swactor/crates/bindings/python/tests/test_bootstrap.py

"""Actor attachment and zero-copy job data-plane guarantees."""

from __future__ import annotations

import asyncio
import ctypes
import json
import os
import struct
import subprocess
import importlib
import sys
import tempfile
from pathlib import Path

import pytest

import swactor

ENV_ARENA = "SWACTOR_ARENA_FD"
_native = importlib.import_module("swactor.swactor")
ENV_ACTOR = "SWACTOR_DATA_PLANE_ACTOR"
ENV_CAPABILITY = "SWACTOR_JOB_CAPABILITY"
ENV_ENDPOINT = "SWACTOR_DATA_PLANE_ENDPOINT"
BOOTSTRAP_ENV = (ENV_ARENA, ENV_ACTOR, ENV_CAPABILITY, ENV_ENDPOINT)
OLD_WAKE_ENV = ("SWACTOR_WAKE_FD", "SWACTOR_HOST_WAKE_FD")
MAGIC = int.from_bytes(b"SWBS", "little")
VERSION = 2
HEADER_LEN = 64
ARENA_SIZE = 4096
PROBE = Path(__file__).parent / "probe.py"


class PyBuffer(ctypes.Structure):
    _fields_ = [
        ("buf", ctypes.c_void_p),
        ("obj", ctypes.c_void_p),
        ("len", ctypes.c_ssize_t),
        ("itemsize", ctypes.c_ssize_t),
        ("readonly", ctypes.c_int),
        ("ndim", ctypes.c_int),
        ("format", ctypes.c_char_p),
        ("shape", ctypes.POINTER(ctypes.c_ssize_t)),
        ("strides", ctypes.POINTER(ctypes.c_ssize_t)),
        ("suboffsets", ctypes.POINTER(ctypes.c_ssize_t)),
        ("internal", ctypes.c_void_p),
    ]


_PY_GET_BUFFER = ctypes.pythonapi.PyObject_GetBuffer
_PY_GET_BUFFER.argtypes = [ctypes.py_object, ctypes.POINTER(PyBuffer), ctypes.c_int]
_PY_GET_BUFFER.restype = ctypes.c_int
_PY_RELEASE_BUFFER = ctypes.pythonapi.PyBuffer_Release
_PY_RELEASE_BUFFER.argtypes = [ctypes.POINTER(PyBuffer)]
_PY_RELEASE_BUFFER.restype = None


@pytest.fixture
def host():
    return _native._test_data_plane_host()


def install_host_env(monkeypatch, host) -> int:
    env = dict(host.env())
    inherited = os.dup(host.arena_fd())
    env[ENV_ARENA] = str(inherited)
    for name in BOOTSTRAP_ENV:
        monkeypatch.setenv(name, env[name])
    for name in OLD_WAKE_ENV:
        monkeypatch.delenv(name, raising=False)
    return inherited


def pack_header(
    *,
    magic: int = MAGIC,
    version: int = VERSION,
    arena_size: int = ARENA_SIZE,
    generation: int = 1,
    control_offset: int = 0,
    control_length: int = 0,
    reserved: tuple[int, int, int] = (0, 0, 0),
) -> bytes:
    return struct.pack(
        "<IHH7Q",
        magic,
        version,
        0,
        arena_size,
        generation,
        control_offset,
        control_length,
        *reserved,
    )


def make_arena(header: bytes, size: int = ARENA_SIZE) -> int:
    fd, path = tempfile.mkstemp(prefix="python-bootstrap-test-")
    os.unlink(path)
    os.ftruncate(fd, size)
    os.pwrite(fd, header, 0)
    return fd


def set_arena(monkeypatch, fd: int) -> None:
    monkeypatch.setenv(ENV_ARENA, str(fd))


def mapping_for_address(address: int) -> str:
    for line in Path("/proc/self/maps").read_text().splitlines():
        extent = line.split(maxsplit=1)[0]
        start, end = (int(value, 16) for value in extent.split("-"))
        if start <= address < end:
            return line
    raise AssertionError(f"no process mapping contains {address:#x}")


def test_run_attaches_before_main_and_maps_blob_buffer_directly(monkeypatch, host):
    inherited = install_host_env(monkeypatch, host)
    observed = {}

    async def main(ctx):
        assert isinstance(ctx, swactor.Context)
        assert isinstance(ctx.data, swactor.DataPlane)
        blob = await ctx.data.read_blob("/models/tiny-linear/weights")
        assert isinstance(blob, swactor.Blob)
        assert blob.length == 24
        assert blob.digest is None
        with blob.map() as mapped:
            assert isinstance(mapped, swactor.BlobView)
            assert struct.unpack_from("<6f", mapped) == pytest.approx(
                (1.5, -2.0, 0.5, 4.0, 0.25, -0.75)
            )
            view = memoryview(mapped)
            assert view.readonly
            assert view.nbytes == 24
            view.release()

            with pytest.raises(BufferError, match="read-only"):
                _PY_GET_BUFFER(mapped, ctypes.byref(PyBuffer()), 1)

            exported = PyBuffer()
            assert _PY_GET_BUFFER(mapped, ctypes.byref(exported), 0) == 0
            try:
                mapping = mapping_for_address(exported.buf)
                assert "data-plane-arena" in mapping
                assert exported.readonly == 1
                assert exported.len == 24
            finally:
                _PY_RELEASE_BUFFER(ctypes.byref(exported))

        public = {name for name in dir(ctx.data) if not name.startswith("_")}
        assert {"read_blob", "write_blob", "read_stream", "write_stream"} <= public
        assert not public & {
            "actor",
            "actor_id",
            "edge",
            "edge_id",
            "offset",
            "peer",
            "ring",
            "socket",
        }
        observed["ran"] = True

    swactor.run(main)
    assert observed == {"ran": True}
    with pytest.raises(OSError):
        os.fstat(inherited)


def test_write_blob_seals_cleanly_and_exception_aborts(monkeypatch, host):
    install_host_env(monkeypatch, host)

    async def main(ctx):
        async with ctx.data.write_blob(
            "/runs/self/results/complete", length=6
        ) as blob:
            assert blob.length == 6
            with blob.map() as mapped:
                struct.pack_into("6s", mapped, 0, b"result")

        class AbortWrite(Exception):
            pass

        try:
            async with ctx.data.write_blob(
                "/runs/self/results/aborted", length=4
            ) as blob:
                with blob.map() as mapped:
                    struct.pack_into("4s", mapped, 0, b"nope")
                raise AbortWrite
        except AbortWrite:
            pass

        complete = await ctx.data.read_blob("/runs/self/results/complete")
        with complete.map() as mapped:
            assert bytes(mapped) == b"result"
        with pytest.raises(swactor.DataPathError, match="not found"):
            await ctx.data.read_blob("/runs/self/results/aborted")

    swactor.run(main)


def test_missing_path_and_authorization_are_typed(monkeypatch, host):
    install_host_env(monkeypatch, host)

    async def main(ctx):
        with pytest.raises(swactor.DataPathError, match="not found"):
            await ctx.data.read_blob("/models/missing")
        with pytest.raises(PermissionError):
            await ctx.data.read_blob("/runs/self/private")

    swactor.run(main)


def test_native_stream_round_trip_has_ordered_eof(monkeypatch, host):
    install_host_env(monkeypatch, host)
    received = []

    async def main(ctx):
        async def receive():
            reader = await ctx.data.read_stream(
                "/runs/self/results/predictions"
            )
            while (chunk := await reader.read()) is not None:
                received.append(chunk)

        receiver = asyncio.create_task(receive())
        async with ctx.data.write_stream(
            "/runs/self/results/predictions"
        ) as stream:
            await stream.write(b"native-")
            await stream.write(b"result")
        await receiver

    swactor.run(main)
    assert b"".join(received) == b"native-result"
    assert "SWACTOR_DATA_PLANE_OUTPUT" not in host.env()


def test_invalid_capability_prevents_main(monkeypatch, host):
    install_host_env(monkeypatch, host)
    monkeypatch.setenv(ENV_CAPABILITY, "00" * 32)
    calls = []

    async def main(_ctx):
        calls.append("ran")

    with pytest.raises(swactor.SessionError, match="CapabilityRejected"):
        swactor.run(main)
    assert calls == []


@pytest.mark.parametrize(
    ("header", "match"),
    [
        (pack_header(magic=MAGIC ^ 0xFF), "magic"),
        (pack_header(version=VERSION + 1), "version"),
        (pack_header(generation=0), "generation"),
        (pack_header(arena_size=ARENA_SIZE + 1), "backing length"),
        (pack_header(control_offset=64, control_length=0), "both offset and length"),
        (pack_header(reserved=(0, 1, 0)), "reserved"),
    ],
)
def test_malformed_arena_fails_before_route_or_main(monkeypatch, host, header, match):
    env = dict(host.env())
    for name, value in env.items():
        monkeypatch.setenv(name, value)
    fd = make_arena(header)
    set_arena(monkeypatch, fd)
    calls = []

    async def main(_ctx):
        calls.append("ran")

    try:
        with pytest.raises(swactor.BootstrapError, match=match):
            swactor.run(main)
        assert calls == []
        with pytest.raises(OSError):
            os.fstat(fd)
    finally:
        try:
            os.close(fd)
        except OSError:
            pass


def test_missing_bootstrap_metadata_prevents_main(monkeypatch):
    for name in (*BOOTSTRAP_ENV, *OLD_WAKE_ENV):
        monkeypatch.delenv(name, raising=False)
    calls = []

    async def main(_ctx):
        calls.append("ran")

    with pytest.raises(swactor.BootstrapError, match=ENV_ARENA):
        swactor.run(main)
    assert calls == []


def test_handoff_uses_one_descriptor_and_no_wake_names(host):
    env = dict(host.env())
    assert set(env) == set(BOOTSTRAP_ENV)
    assert all(name not in env for name in OLD_WAKE_ENV)
    assert env[ENV_ARENA] == str(host.arena_fd())


def test_real_exec_attachment_and_blob_mapping(host):
    env = {
        key: value
        for key, value in os.environ.items()
        if key not in (*BOOTSTRAP_ENV, *OLD_WAKE_ENV)
    }
    env.update(dict(host.env()))
    result = subprocess.run(
        [sys.executable, str(PROBE)],
        env=env,
        pass_fds=(host.arena_fd(),),
        text=True,
        capture_output=True,
        timeout=45,
        check=False,
    )
    assert result.returncode == 0, result.stderr
    facts = dict(
        item.split("=", 1) for item in result.stdout.strip().split() if "=" in item
    )
    assert facts == {
        "HAS_DATA": "1",
        "LENGTH": "24",
        "VALUES": "1.5,-2,0.5,4,0.25,-0.75",
        "ARENA_OPEN": "0",
        "IDENTITY_LEAKS": "0",
    }

@pytest.mark.skipif(
    not Path("/dev/nvidia0").exists(),
    reason="CUDA device is unavailable",
)
def test_real_exec_tinygrad_cuda_scenario(host):
    capture = host.capture_stream("/runs/self/results/inference")
    env = {
        key: value
        for key, value in os.environ.items()
        if key not in (*BOOTSTRAP_ENV, *OLD_WAKE_ENV, "SWACTOR_DATA_PLANE_OUTPUT")
    }
    env.update(dict(host.env()))
    env.update(
        {
            "CUDA_PTX": "1",
            "DEV": "CUDA",
        }
    )
    script = (
        Path(__file__).resolve().parents[4]
        / "apps"
        / "myelin"
        / "jobs"
        / "tiny_linear_inference.py"
    )
    result = subprocess.run(
        [sys.executable, str(script)],
        env=env,
        pass_fds=(host.arena_fd(),),
        text=True,
        capture_output=True,
        timeout=45,
        check=False,
    )
    assert result.returncode == 0, result.stderr
    payload = json.loads(capture.result())
    assert payload["device"].startswith("CUDA")
    assert payload["output"] == pytest.approx([2.75, -8.75])