diff --git a/docs/chip-level-arch.md b/docs/chip-level-arch.md index a0e76f0f5d..8d98821709 100644 --- a/docs/chip-level-arch.md +++ b/docs/chip-level-arch.md @@ -180,6 +180,13 @@ level model (see [hierarchical-level-runtime.md](hierarchical-level-runtime.md)) `CallConfig` is the exception — same type used at every level, with no `Chip*` / unprefixed split (see [task-flow.md](task-flow.md) for details). The unified `Worker(level=N)` factory already routes to the correct backend. +At `level=2` the default `execution_mode="program"` drives `ChipWorker.init` / +`run`. `Worker(level=2, execution_mode="kernel")` routes to the kernel entries +instead: `init(config=...)` to `ChipWorker.kernel_init`, +`kernel_prepare_callable` / `kernel_launch` to `ChipWorker.kernel_*`, and +`close()` to `ChipWorker.finalize`. Its `kernel_mode_supported` property +answers through `ChipWorker.probe_kernel_mode_supported`, before init and without +touching a device. Kernel mode has no L3+ route. When new level-specific types are added (e.g. `ChipCallable`), each concept should follow the same pattern: a `Chip*` concrete type for L2, a prefix-less concrete type for L3+, and optionally a factory function that routes by level. diff --git a/docs/user/reference/python-api.md b/docs/user/reference/python-api.md index 36d7a9f9d2..c344d20462 100644 --- a/docs/user/reference/python-api.md +++ b/docs/user/reference/python-api.md @@ -35,6 +35,7 @@ into `**config` and validated later. The recognized keys: | --- | ---------- | ------- | | `platform` | all | `a2a3`, `a2a3sim`, `a5`, `a5sim` | | `runtime` | all | `tensormap_and_ringbuffer` or `host_build_graph` | +| `execution_mode` | L2 | `"program"` (default) or `"kernel"`; see [Kernel mode](#kernel-mode-l2). `"kernel"` with `level != 2` or a truthy `enable_sdma` raises `ValueError`, as does an unknown value | | `device_id` | L2 | the single chip this worker drives | | `device_ids` | L3+ | one chip child process per entry | | `num_sub_workers` | L3+ | host-side Python callables to fork | @@ -53,8 +54,8 @@ else raises. Remote-worker and remote-memory calls require `level >= 4`. | `unregister(handle_or_slot)` | Releases a registration | | `add_worker(worker) -> int` | Attaches a child worker; returns its id | | `add_remote_worker(spec: RemoteWorkerSpec) -> int` | L4; see the remote-L3 design doc | -| `init(prewarm_config=None)` | Resolves runtime binaries, opens the device, forks children. First place setup errors appear | -| `close()` | Releases the device and reaps children. Put it in a `finally` — a skipped `close()` leaves the device held | +| `init(prewarm_config=None, *, config=None)` | Resolves runtime binaries, opens the device, forks children. First place setup errors appear. `config=` is kernel mode's required context config and is refused in program mode; kernel mode refuses `prewarm_config` | +| `close()` | Releases the device and reaps children. Put it in a `finally` — a skipped `close()` leaves the device held. In `execution_mode="kernel"` it releases only the context; see [Kernel mode](#kernel-mode-l2) | ### Memory @@ -97,6 +98,71 @@ callable is a **Python orchestration function** `f(orch, args, cfg)`, where | `allocate_domain(*, name, workers, window_size, buffers=())` | Context manager returning a handle indexed by domain-local rank | | `alloc_child_tensor(worker_id, shapes, dtype) -> Buffer` | Delegates allocation to the owning Worker; target chip memory is named by the returned handle | +### Kernel mode (L2) + +`Worker(level=2, execution_mode="kernel")` is the entry for a framework that +already owns the device and its streams. The Worker borrows the device that is +current on the calling thread. Each launch enqueues bounded asynchronous work +on the caller's stream and returns without synchronizing. + +| Member | Notes | +| ------ | ----- | +| `init(*, config: CallConfig)` | Kernel init on the already-current `device_id`. `config` is required and validated; it is fixed for the life of the context. A positional argument binds to `prewarm_config` and is refused, so pass `config=CallConfig(...)`. There is no prewarm, no SDMA, and no replay of registrations | +| `kernel_mode_supported -> bool` | Property. `False` when `level != 2`. Otherwise it reports whether the configured platform/runtime build supports kernel mode, whatever the `execution_mode`. The build is probed once, without touching a device, and the result is cached, so it can be asked before `init()` or before choosing a mode. Missing runtime binaries raise rather than answer `False` | +| `kernel_prepare_callable(chip_callable: ChipCallable) -> int` | Uploads the callable and returns the id the runtime minted. It blocks until device-side registration finishes, so a registration failure raises from this call. There is no deduplication: the same callable prepared twice gets two ids. It is unrelated to `register()`. Call it outside ACLGraph capture. The Worker keeps the callable alive until `close()` | +| `kernel_launch(callable_id: int, args: ChipStorageTaskArgs, *, caller_stream: int) -> None` | Enqueues one invocation on `caller_stream`, a nonzero `aclrtStream` address, and returns. Returning means enqueued; device-side errors surface when the caller synchronizes its own stream. No `RunHandle`, no wait. An id this Worker did not mint or a zero stream raises `ValueError`, and any other `args` type raises `TypeError`. A launch that overlaps a prepare, launch or close raises `RuntimeError` at once rather than waiting | +| `close()` | Releases only the resources the context owns. It never synchronizes the caller's stream, resets the device, or finalizes ACL | + +In kernel mode these raise `RuntimeError`: `register` (also before `init()`), +`unregister`, `run`, `submit`, `malloc`, `free`, `copy_to`, `copy_from`, +`create_buffer`, `make_tensor_arg`, `release_buffer`, and `device_memory_info`. +`committed_device_memory` stays available. In program mode, `kernel_prepare_callable` +and `kernel_launch` raise `RuntimeError`. + +The caller's side of the contract: + +- `device_id` is already current on the thread that calls `init()`, and that + thread also calls `close()`. +- `args` tensors are `ChipTensor.make(addr, shapes, dtype, child_memory=True)` + views of caller-owned device memory. The caller keeps `caller_stream` alive + until the work it submitted completes. +- The caller serializes `kernel_prepare_callable`, `kernel_launch` and `close()`. +- Before `close()`, the caller stops launching, synchronizes its own stream, and + destroys every graph that replays the context. Synchronizing, resetting the + device and finalizing ACL after `close()` are also the caller's job. +- `close()` raises `TimeoutError` when an in-flight prepare or launch does not + finish within the rollback grace period, and `RuntimeError` when native + teardown fails. In both cases the context is kept, and a later `close()` + retries. +- Only the process that called `init()` may prepare, launch or close; a forked + child gets `RuntimeError`. Kernel mode has no L3+ form. +- A kernel Worker that is garbage-collected, or still open at interpreter exit, + without `close()` emits a `ResourceWarning`. Its native context is leaked on + purpose rather than finalized on an arbitrary thread. + +```python +from simpler import Worker +from simpler.task_interface import CallConfig, ChipStorageTaskArgs, ChipTensor, DataType + +# The framework has made device 0 current on this thread and owns `stream`, +# the device buffers `x_addr` / `y_addr`, and `chip_callable`. +worker = Worker(level=2, execution_mode="kernel", device_id=0, + platform="a2a3", runtime="tensormap_and_ringbuffer") +if not worker.kernel_mode_supported: + raise RuntimeError("this runtime build has no kernel mode") +worker.init(config=CallConfig()) +try: + cid = worker.kernel_prepare_callable(chip_callable) + args = ChipStorageTaskArgs() + args.add_tensor(ChipTensor.make(x_addr, (n,), DataType.FLOAT32, child_memory=True)) + args.add_tensor(ChipTensor.make(y_addr, (n,), DataType.FLOAT32, child_memory=True)) + args.add_scalar(1.25) + worker.kernel_launch(cid, args, caller_stream=stream) + synchronize(stream) # the framework's own stream synchronize; device errors surface here +finally: + worker.close() +``` + ## Callables and task args ```python @@ -120,7 +186,8 @@ including for callables loaded from cached bytes. The tensor count is Public `Worker` calls use `TaskArgs` containing address-free `Tensor` views at every level. The L2 leaf resolves those views into the internal `ChipStorageTaskArgs` / `ChipTensor` representation; callers do not pass that -internal representation to `Worker.run()`. +internal representation to `Worker.run()`. The exception is kernel mode's +`kernel_launch`, which takes a caller-built `ChipStorageTaskArgs` directly. For L3+ graph construction, `TaskArgs.add_dep(*handles)` adds `WAIT | RETAIN` edges: each consumer waits for its producers and keeps their task-owned diff --git a/docs/zh-cn/kernel-callable-residency.md b/docs/zh-cn/kernel-callable-residency.md index ba0fa7f8d5..2d2670a433 100644 --- a/docs/zh-cn/kernel-callable-residency.md +++ b/docs/zh-cn/kernel-callable-residency.md @@ -41,7 +41,7 @@ context-local `int32_t`，失败写 `-1`。当前上限为 8192。 int32_t callable_id = -1; int rc = simpler_kernel_mode_prepare_callable(ctx, callable, size, &callable_id); if (rc != 0) return rc; -// 在 capture 之外完成 caller 的 warmup 并同步，检查异步准备结果。 +// prepare 已同步 context 自己的 AICPU stream，rc 即注册结果；以下是调用者为 capture 做的 warmup 和同步。 rc = caller_warmup_and_synchronize(caller_stream); // 调用者自己的逻辑 if (rc != 0) return rc; return simpler_kernel_mode_launch(ctx, callable_id, args, caller_stream); diff --git a/docs/zh-cn/kernel-mode-integration-test.md b/docs/zh-cn/kernel-mode-integration-test.md index b2f94d1f37..413b19674c 100644 --- a/docs/zh-cn/kernel-mode-integration-test.md +++ b/docs/zh-cn/kernel-mode-integration-test.md @@ -20,6 +20,13 @@ kernel 模式下，simpler 是一个被调用的库：它借用调用方已经 直接调用 host runtime 动态库，自己扮演调用方，不经过 PyTorch，也不经过 simpler 的 Python Worker。 +经过 Worker 的是另外两个文件，入口都是 `Worker(level=2, execution_mode="kernel")`： + +| 文件 | 硬件 | 覆盖 | +| ---- | ---- | ---- | +| `tests/ut/py/test_worker/test_worker_kernel_mode.py` | 不需要 | 用假的 ChipWorker 检查参数校验、`init(config=...)` 的路由、program 与 kernel 两种模式的接口互斥、prepare/launch/close 之间的串行闸门，以及 teardown 失败后 `close()` 可重试 | +| `tests/ut/py/test_worker/test_worker_kernel_mode_hw.py` | a2a3 真机 | 调用方自己设卡、建 stream，经 `init(config=...)`、`kernel_prepare_callable`、`kernel_launch(..., caller_stream=...)` 和 `close()` 驱动 kernel 模式，再同步自己的 stream 核对结果 | + 被执行的算子是一个 AIV 向量加标量，`y[i] = x[i] + scalar`： | 组成 | 源文件 | 编译产物 | diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index fc8aae0617..82a143f491 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -3623,6 +3623,14 @@ NB_MODULE(_task_interface, m) { "A nonzero context generation, unique and increasing within this host " "process. Generation zero is what the C ABI rejects as invalid." ) + .def_static( + "probe_kernel_mode_supported", &ChipWorker::probe_kernel_mode_supported, nb::arg("host_lib_path"), + nb::arg("sim_context_path"), nb::call_guard(), + "Whether the host runtime at host_lib_path can execute kernel-mode launches. Answers from the " + "runtime build alone on a fresh device context that is destroyed before returning, so it takes no " + "device and needs no init or kernel_init. Raises RuntimeError when the runtime cannot be loaded, " + "lacks a required symbol, or yields no context." + ) .def_prop_ro("device_id", &ChipWorker::device_id) .def_prop_ro("initialized", &ChipWorker::initialized) .def_prop_ro("pipeline_depth", &ChipWorker::pipeline_depth) diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 454865054c..b929898bc6 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1541,6 +1541,34 @@ def kernel_init( with self._lifecycle_lock: self._init_in_progress = False + @staticmethod + def probe_kernel_mode_supported(bins: Any, log_level: int | None = None) -> bool: + """Whether the host runtime in ``bins`` can execute kernel-mode launches. + + Answers before init and independently of every ChipWorker: the native + probe loads the host runtime, asks it on a fresh device context that no + init touches, and destroys that context before returning. The runtime + build alone decides the answer, so no device is taken, no thread is + attached, and the caller needs no current device. + + Args: + bins: same structural type init() takes; only host_path and + sim_context_path are read. + log_level: as for init(). + + Raises: + RuntimeError: the host runtime cannot be loaded, lacks a required + symbol, or yields no device context. + """ + _initialize_host_log(log_level) + sim_context_path = getattr(bins, "sim_context_path", None) + return bool( + _ChipWorker.probe_kernel_mode_supported( + str(bins.host_path), + "" if sim_context_path is None else str(sim_context_path), + ) + ) + @property def kernel_mode_supported(self) -> bool: """Whether the bound runtime can execute kernel-mode launches. @@ -1555,12 +1583,15 @@ def kernel_mode_supported(self) -> bool: def kernel_prepare_callable(self, chip_callable: ChipCallable) -> int: """Register a callable for kernel-mode launches, outside ACLGraph capture. - Returns the ID the runtime minted for it. Registration is pure: the same - callable registered twice takes two distinct, equally valid IDs, and - there is no lookup. It takes no stream — registration enqueues on the - context's own AICPU stream, which every later launch also enqueues on, - so stream FIFO orders registration ahead of each launch. The callable - stays referenced under its ID until finalize() succeeds. + Returns the context-local id simpler minted for it. Registration is + pure: the same callable registered twice takes two distinct, equally + valid ids, and there is no lookup. It takes no stream — registration + enqueues on the context's own AICPU stream, which every later launch + also enqueues on, so stream FIFO orders registration ahead of each + launch. Registration synchronizes that context stream before committing + the callable, so a device-side registration failure raises from this + call; it synchronizes no caller stream and no device. The callable stays + referenced under its id until finalize() succeeds. """ callable_id = int(self._impl.kernel_prepare_callable(chip_callable)) # The registry owns the image for the life of the context: the device diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 790f72f2b4..8182abe844 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -67,6 +67,7 @@ def my_l4_orch(orch, args, config): import json import logging import math +import operator import os import re import shutil @@ -79,6 +80,8 @@ def my_l4_orch(orch, args, config): import threading import time import uuid +import warnings +import weakref from collections.abc import Iterable, Iterator from dataclasses import dataclass, field, replace from multiprocessing import resource_tracker @@ -270,6 +273,7 @@ def my_l4_orch(orch, args, config): CallConfig, ChipCallable, ChipDomainContext, + ChipStorageTaskArgs, ChipWorker, CommBufferSpec, CommDomainHandle, @@ -4623,6 +4627,44 @@ def __len__(self) -> int: return len(self._snapshots) +# ChipWorkers of kernel-mode Workers that were garbage-collected, or still alive at interpreter exit, +# without close(). Entries are never removed; see _pin_unclosed_kernel_chip. +_PINNED_KERNEL_CHIP_WORKERS: list[ChipWorker] = [] + + +class _KernelChipPin: + """The ChipWorker a kernel-mode Worker's GC finalizer pins; ``chip`` is None once close() tore it down.""" + + __slots__ = ("chip",) + + def __init__(self, chip: ChipWorker) -> None: + self.chip: ChipWorker | None = chip + + +def _pin_unclosed_kernel_chip(pin: _KernelChipPin) -> None: + """Finalizer of a kernel-mode Worker collected, or still alive at interpreter exit, without close(). + + ``~ChipWorker`` finalizes its context on whichever thread destroys it, without the caller's + quiescence, stream synchronization and graph destruction that kernel-mode teardown requires. The + ChipWorker is therefore kept alive for the rest of the process instead. The finalizer receives only + the pin, never the Worker, so it does not keep the Worker reachable. + """ + chip = pin.chip + if chip is None: + return + _PINNED_KERNEL_CHIP_WORKERS.append(chip) + # Interpreter finalization clears module globals, this list included, and dropping the list's + # reference there would run ~ChipWorker. This reference is never released, so the object outlives + # finalization and its native destructor never runs. + ctypes.pythonapi.Py_IncRef(ctypes.py_object(chip)) + warnings.warn( + "Worker(execution_mode='kernel') was garbage-collected without close(); its kernel context is " + "intentionally leaked rather than finalized on an arbitrary thread", + ResourceWarning, + stacklevel=2, + ) + + class Worker: """Unified worker for all hierarchy levels. @@ -4632,8 +4674,44 @@ class Worker: level=4+: wraps the C++ Worker composite with Worker(level-1)×N as NEXT_LEVEL children + SubWorker×M. Children are added via add_worker() before init(). + + A level=2 Worker runs in one of two execution modes, fixed at construction by + the ``execution_mode`` config key: + + - ``"program"`` (default): ``init()`` opens the device; callables go through + ``register()`` and run through ``submit()`` / ``run()``. + - ``"kernel"``: ``init(config=...)`` binds the runtime as a kernel-mode + context on a device and stream the caller already owns:: + + worker = Worker(level=2, execution_mode="kernel", device_id=0, + platform="a2a3", runtime="tensormap_and_ringbuffer") + worker.init(config=CallConfig()) # device 0 is already current on this thread + callable_id = worker.kernel_prepare_callable(chip_callable) + worker.kernel_launch(callable_id, args, caller_stream=stream) + ... # the caller synchronizes its own stream + worker.close() + + Preconditions: the thread that calls ``init()`` already has ``device_id`` + current, and the same thread calls ``close()``. The caller serializes + ``kernel_prepare_callable``, ``kernel_launch`` and ``close()``; an + overlapping launch fails immediately. Before ``close()`` the caller stops + launching, synchronizes its own stream, and destroys every graph that + replays the context. ``close()`` never synchronizes the caller's stream, + resets the device, or finalizes ACL. Only the process that called + ``init()`` may drive or close the context. Kernel mode has no L3+ form. + The program-mode APIs (``register`` / ``unregister`` / ``submit`` / + ``run``, ``malloc`` / ``free`` / ``copy_to`` / ``copy_from``, + ``create_buffer`` / ``make_tensor_arg`` / ``release_buffer``, and + ``device_memory_info``) raise ``RuntimeError`` on a kernel-mode Worker. + A kernel-mode Worker garbage-collected without ``close()`` emits a + ``ResourceWarning`` and leaks its context rather than finalizing it on the + collecting thread. """ + # `__init__` sets the instance value; this default makes a Worker built without `__init__` + # (`Worker.__new__`) read as program mode in the mode guards and in close(). + _execution_mode: str = "program" + def __init__( self, level: int, @@ -4714,6 +4792,7 @@ def __init__( self._startup_timeout_s = float(config.get("startup_timeout_s", _STARTUP_TIMEOUT_S)) if not (self._startup_timeout_s > 0 and math.isfinite(self._startup_timeout_s)): raise ValueError("Worker startup_timeout_s must be a positive finite number of seconds") + self._init_execution_mode(level, config) # Per-startup bookkeeping consumed by the rollback path: PIDs the barrier # already reaped (must not be re-SIGKILLed — the PID may be reused) and # PIDs that reached their serve loop (READY → asked to close gracefully @@ -6558,6 +6637,55 @@ def _operation_lease(self, api: str): self._lease_depth[tid] = depth self._hierarchical_start_cv.notify_all() + def _init_execution_mode(self, level: int, config: dict[str, Any]) -> None: + """Validate the ``execution_mode`` config key and set up the kernel-mode state. + + ``"kernel"`` binds a level-2 kernel-mode context in init(); it has no L3+ form and provisions + no SDMA workspace, so both are refused here, before any startup resource exists. + """ + execution_mode = config.get("execution_mode", "program") + if execution_mode not in ("program", "kernel"): + raise ValueError(f"Worker execution_mode must be 'program' or 'kernel', got {execution_mode!r}") + if execution_mode == "kernel" and level != 2: + raise ValueError("execution_mode='kernel' requires level=2") + if execution_mode == "kernel" and config.get("enable_sdma"): + raise ValueError("execution_mode='kernel' does not support enable_sdma") + self._execution_mode = execution_mode + # `_kernel_callables` maps each id kernel_prepare_callable returned to its image; the device + # holds addresses into those images until a native teardown succeeds, so only that teardown + # clears it. `_kernel_gate` linearizes prepare and launch against the native finalize. + # `_kernel_pid` is the process whose kernel_init bound the context, None until one succeeds. + self._kernel_config: CallConfig | None = None + self._kernel_callables: dict[int, ChipCallable] = {} + self._kernel_gate = threading.Lock() + self._kernel_pid: int | None = None + self._kernel_chip_pin: _KernelChipPin | None = None + self._kernel_pin_finalizer: weakref.finalize | None = None + # Cached pre-init kernel_mode_supported answer for this Worker's platform/runtime build. + self._kernel_supported_probe: bool | None = None + self._kernel_supported_lock = threading.Lock() + + def _require_execution_mode(self, api: str, mode: str) -> None: + """Reject an API of the other execution mode; the mode is fixed at construction.""" + if self._execution_mode != mode: + raise RuntimeError( + f"Worker.{api}: requires execution_mode={mode!r}, but this Worker was constructed with " + f"execution_mode={self._execution_mode!r}" + ) + + def _require_kernel_process(self, api: str) -> None: + """Reject a kernel-context operation from a process other than the one whose kernel_init bound it. + + A forked child inherits this Worker object but must not drive or tear down the parent's + context. Before a kernel_init succeeds there is no context to fence. + """ + pid = self._kernel_pid + if pid is not None and os.getpid() != pid: + raise RuntimeError( + f"Worker.{api}: this kernel-mode context belongs to process {pid}; " + f"process {os.getpid()} must not drive or tear it down" + ) + def _invalidate_endpoint_registry(self) -> None: self._endpoint_registry = None self._region_access_service = None @@ -6741,7 +6869,11 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl A post-init dynamic register re-validates eligibility against the frozen topology (``_eligible_target_need``), same as init(). + + Program mode only, before and after init(); a kernel-mode Worker + registers images through ``kernel_prepare_callable``. """ + self._require_execution_mode("register", "program") if isinstance(target, RemoteCallable) and self.level < 4: raise TypeError("Worker.register(RemoteCallable): remote L3 dispatch requires a level >= 4 parent") if self.level == 2 and not isinstance(target, ChipCallable): @@ -6848,6 +6980,123 @@ def _post_start_register_l2(self, reg: _CallableRegistration, target: ChipCallab raise return handle + @property + def kernel_mode_supported(self) -> bool: + """Whether this Worker's platform/runtime build can run kernel-mode launches. + + This is the question a caller asks before choosing ``execution_mode``, so the answer is + independent of the lifecycle and of this Worker's own mode: it is the same before + ``init()``, in program mode, and after ``close()``. A Worker whose level is not 2 has no + kernel mode and answers False. Otherwise the first read loads the runtime build and asks a + fresh device context that touches no device, then caches the answer; a READY kernel-mode + Worker answers True without probing, because kernel_init succeeds only on a supporting + runtime. A missing runtime build raises exactly as ``init()`` would, rather than reading as + unsupported. + """ + if self.level != 2: + return False + with self._hierarchical_start_cv: + if self._execution_mode == "kernel" and self._lifecycle is _Lifecycle.READY: + return True + with self._kernel_supported_lock: + if self._kernel_supported_probe is None: + from simpler_setup.runtime_builder import RuntimeBuilder # noqa: PLC0415 + + binaries = RuntimeBuilder(self._config["platform"]).get_binaries(self._config["runtime"]) + self._kernel_supported_probe = bool(ChipWorker.probe_kernel_mode_supported(binaries)) + return self._kernel_supported_probe + + def kernel_prepare_callable(self, chip_callable: ChipCallable) -> int: + """Register ``chip_callable`` with this kernel-mode context; returns the id the runtime minted. + + Unrelated to ``register()``: there is no digest dedup, no handle, and no pre-init recording, + so the same callable prepared twice takes two distinct, equally valid ids. The call blocks + until the device-side registration is committed and raises its failure; call it outside graph + capture. The Worker keeps every prepared image alive until ``close()`` tears the context down, + because the device holds addresses into it. Requires a READY kernel-mode Worker, called in the + process that initialized it. + """ + self._require_execution_mode("kernel_prepare_callable", "kernel") + if not isinstance(chip_callable, ChipCallable): + raise TypeError( + f"Worker.kernel_prepare_callable: expected a ChipCallable, got {type(chip_callable).__name__}" + ) + self._require_kernel_process("kernel_prepare_callable") + with self._operation_lease("kernel_prepare_callable"), self._kernel_gate: + chip = self._chip_worker + assert chip is not None + callable_id = int(chip.kernel_prepare_callable(chip_callable)) + # Recorded before the id is judged: the runtime may hold the image under an id rejected + # below, and this entry is what keeps the image alive for it. + self._kernel_callables[callable_id] = chip_callable + if callable_id < 0: + raise RuntimeError(f"Worker.kernel_prepare_callable: the runtime minted invalid id {callable_id}") + return callable_id + + def kernel_launch(self, callable_id: int, args: ChipStorageTaskArgs, *, caller_stream: int) -> None: + """Enqueue one kernel-mode invocation of ``callable_id`` on ``caller_stream``. + + Returning means the invocation was enqueued; device execution may still be in flight, and a + device-side failure surfaces on the caller's own synchronize. ``callable_id`` must be an id + this Worker's ``kernel_prepare_callable`` returned, and ``caller_stream`` the caller's + non-null ACL stream for this call. There is no ``RunHandle``, no wait, and no operation + lease: a launch that overlaps another launch, a prepare, or ``close()`` raises + ``RuntimeError`` immediately instead of queueing, because the caller serializes them. + Requires a READY kernel-mode Worker, called in the process that initialized it. + """ + self._require_execution_mode("kernel_launch", "kernel") + stream = int(caller_stream) + if not stream: + raise ValueError("Worker.kernel_launch requires a non-null caller_stream") + if not isinstance(args, ChipStorageTaskArgs): + raise TypeError(f"Worker.kernel_launch: args must be ChipStorageTaskArgs, got {type(args).__name__}") + callable_id = operator.index(callable_id) + self._require_kernel_process("kernel_launch") + if not self._kernel_gate.acquire(blocking=False): + raise RuntimeError( + "Worker.kernel_launch: a kernel prepare/launch/close is in progress; the caller serializes them" + ) + try: + # Read without the lifecycle lock: close() publishes CLOSED before it takes this gate to + # finalize, so a launch holding the gate either observes CLOSED or finishes before finalize. + if self._lifecycle is not _Lifecycle.READY: + raise RuntimeError( + "Worker.kernel_launch: requires an initialized (READY) worker" + ) from self._startup_error + # A negative id stays in `_kernel_callables` only to keep its image alive; prepare + # raised for it, so it is not a launchable id. + if callable_id < 0 or callable_id not in self._kernel_callables: + raise ValueError( + f"Worker.kernel_launch: callable_id {callable_id} was not returned by this Worker's " + "kernel_prepare_callable" + ) + chip = self._chip_worker + assert chip is not None + chip.kernel_launch(callable_id, args, stream) + finally: + self._kernel_gate.release() + + def _arm_kernel_chip_pin(self) -> None: + """Arm the finalizer that leaks this Worker's kernel ChipWorker if the Worker is never closed. + + ``weakref.finalize`` also runs at interpreter exit, so a Worker still alive then is pinned + before module teardown. + """ + assert self._chip_worker is not None + pin = _KernelChipPin(self._chip_worker) + self._kernel_chip_pin = pin + self._kernel_pin_finalizer = weakref.finalize(self, _pin_unclosed_kernel_chip, pin) + + def _disarm_kernel_chip_pin(self) -> None: + """Drop the pin once the kernel context is torn down; nothing is left to leak.""" + pin, finalizer = self._kernel_chip_pin, self._kernel_pin_finalizer + self._kernel_chip_pin = None + self._kernel_pin_finalizer = None + if pin is not None: + pin.chip = None + if finalizer is not None: + finalizer.detach() + def _python_worker_types(self) -> list[WorkerType]: worker_types: list[WorkerType] = [] if self._config.get("num_sub_workers", 0) > 0: @@ -7339,7 +7588,9 @@ def unregister(self, handle_or_slot) -> None: Raises: KeyError: handle was never registered. + RuntimeError: the Worker was constructed with ``execution_mode="kernel"``. """ + self._require_execution_mode("unregister", "program") if self._pre_start_unregister_if_needed(handle_or_slot): return # Every post-start path takes the READY-only lease before touching the @@ -7725,7 +7976,11 @@ def _validate_eligible_targets(self) -> None: ) def init( # noqa: PLR0912, PLR0915 - self, prewarm_config: CallConfig | None = None, *, _startup_deadline: float | None = None + self, + prewarm_config: CallConfig | None = None, + *, + config: CallConfig | None = None, + _startup_deadline: float | None = None, ) -> None: """Initialize the worker and bring its whole subtree to READY. @@ -7740,6 +7995,15 @@ def init( # noqa: PLR0912, PLR0915 ``run`` / ``create_buffer`` / the remote register/memory APIs never trigger startup. + An L2 worker constructed with ``execution_mode="kernel"`` opens no + device: it binds the runtime as a kernel-mode context on a device the + caller already holds. The calling thread must already have + ``device_id`` current (init checks it and never sets it), and that + thread is the one that calls ``close()``. Nothing is prewarmed, no + registration is replayed, and no SDMA workspace is provisioned. A + runtime without kernel mode fails init, which rolls back to FAILED like + any other startup failure. + Args: prewarm_config: Optional CallConfig. When given, its ring sizing (``runtime_env.ring_task_window`` / ``ring_heap`` / @@ -7748,14 +8012,37 @@ def init( # noqa: PLR0912, PLR0915 An L2 worker prewarms here; an L3+ worker prewarms each chip child during hierarchy startup, before it publishes INIT_READY. A no-op for runtimes without a prebuilt arena (host_build_graph). ``None`` - (default) disables prewarm. + (default) disables prewarm. Program mode only. + config: The context-static CallConfig a kernel-mode context binds, + validated before startup. Required with + ``execution_mode="kernel"`` and refused in program mode. It is + keyword-only, so a positional CallConfig binds to + ``prewarm_config``. _startup_deadline: Internal. Absolute ``time.monotonic()`` deadline inherited from a parent's startup epoch so a recursive descendant consumes the parent's remaining budget instead of restarting the timeout. ``None`` starts a fresh epoch. """ - if prewarm_config is not None: - prewarm_config.validate() + if self._execution_mode == "kernel": + if config is None: + raise ValueError( + "Worker.init(): execution_mode='kernel' requires config=CallConfig(...) " + "(a positional argument binds to prewarm_config)" + ) + if prewarm_config is not None: + raise ValueError( + "Worker.init(): execution_mode='kernel' takes no prewarm_config; pass the context-static " + "CallConfig as config=" + ) + config.validate() + else: + if config is not None: + raise ValueError( + "Worker.init(): config= is only accepted with execution_mode='kernel'; " + "program mode takes prewarm_config" + ) + if prewarm_config is not None: + prewarm_config.validate() # Claim the startup epoch atomically: NEW -> INITIALIZING under the # lifecycle lock so a concurrent init / register / close observes one # linear transition and never a half-built Worker. Every level claims the @@ -7783,6 +8070,7 @@ def init( # noqa: PLR0912, PLR0915 # this lock). self._validate_eligible_targets() self._prewarm_config = prewarm_config + self._kernel_config = config self._startup_error = None self._init_owner_thread = threading.current_thread() self._cancel_token = False @@ -7871,6 +8159,17 @@ def _init_level2(self) -> None: builder = RuntimeBuilder(platform) binaries = builder.get_binaries(runtime) + if self._execution_mode == "kernel": + kernel_config = self._kernel_config + assert kernel_config is not None + # Assigned before kernel_init so a raising kernel_init still leaves rollback's + # _finalize_chip a ChipWorker to finalize. + self._chip_worker = ChipWorker() + self._chip_worker.kernel_init(device_id, binaries, kernel_config) + self._kernel_pid = os.getpid() + self._arm_kernel_chip_pin() + return + self._chip_worker = ChipWorker() # The prebuilt runtime-arena is prewarmed inside cw.init for the declared # config's ring sizing (built right after the device comes up), so the @@ -10616,6 +10915,7 @@ def malloc(self, size: int) -> Buffer: allocates child device memory with ``alloc_child_tensor(worker_id, ...)`` instead — a Worker is the only allocator, the Orchestrator never allocates. """ + self._require_execution_mode("malloc", "program") if self.level != 2: raise TypeError("worker.malloc is L2-only; at L3+ use worker.alloc_child_tensor(worker_id, ...)") with self._operation_lease("malloc"): @@ -10677,6 +10977,7 @@ def free(self, handle: Buffer) -> None: The operation lease is re-entrant, so an in-run ``orch.free`` that delegates here nests safely. """ + self._require_execution_mode("free", "program") if self.level != 2 and not self._chip_shms: self._check_chip_worker_id(0) # Lock selection comes from the private registration snapshot. A caller may mutate the @@ -10727,7 +11028,10 @@ def device_memory_info(self, worker_id: int = 0) -> DeviceMemoryInfo: Level 2 queries the in-process chip worker. Level 3 routes by logical *worker_id* to the matching forked chip child. Simulator backends do not synthesize device-wide memory and raise ``NotImplementedError``. + A kernel-mode Worker raises ``RuntimeError``: the native query refuses a + kernel-mode context. """ + self._require_execution_mode("device_memory_info", "program") worker_id = int(worker_id) with self._operation_lease("device_memory_info"): if self.level == 2: @@ -10825,6 +11129,7 @@ def copy_to(self, dst: Buffer, src, *, dst_offset: int = 0, src_offset: int = 0, allocation ``dst`` already names; there is no way to name a sub-range with a handle built at an interior address, because such a handle names no allocation at all. """ + self._require_execution_mode("copy_to", "program") host, src_addr, host_nbytes = self._host_side_of_copy(src, writing=False, api="copy_to") dst_offset, src_offset, nbytes = self._copy_extent( host_nbytes, dst_offset, src_offset, nbytes, host_side="src", api="copy_to" @@ -10885,6 +11190,7 @@ def copy_from( ``nbytes`` defaults to the rest of the host side after ``dst_offset``, so a plain ``copy_from(dst, src)`` still transfers a whole host backing's worth. """ + self._require_execution_mode("copy_from", "program") host, dst_addr, host_nbytes = self._host_side_of_copy(dst, writing=True, api="copy_from") src_offset, dst_offset, nbytes = self._copy_extent( host_nbytes, src_offset, dst_offset, nbytes, host_side="dst", api="copy_from" @@ -10948,6 +11254,7 @@ def create_buffer(self, nbytes: int) -> Buffer: tensor in-process on ``run``. Build a tensor over ``buffer.shm.buf`` with the buffer protocol. Not thread-safe against a concurrent run/create/free on the same Worker. """ + self._require_execution_mode("create_buffer", "program") if self.level < 2: raise TypeError("create_buffer requires a level >= 2 Worker") with self._operation_lease("create_buffer"): @@ -10993,6 +11300,7 @@ def make_tensor_arg(self, tensor, shapes: tuple[int, ...], dtype: int, *, stride it; the ``byte_offset`` this computes is what then separates two views that do not intersect. At L2 (no fork) any host tensor works. ``dtype`` is the ``DataType`` int value. """ + self._require_execution_mode("make_tensor_arg", "program") untyped_storage = getattr(tensor, "untyped_storage", None) if callable(untyped_storage): st = untyped_storage() @@ -11125,6 +11433,7 @@ def release_buffer(self, buffer: Buffer) -> None: The slot is dropped only when it still holds *this* buffer: a buffer_id minted elsewhere can collide with a registry key, and evicting the live entry it names would strand that backing.""" + self._require_execution_mode("release_buffer", "program") if not buffer.closed: # Exclusive, not shared: `shared()` would already exclude admission and so # satisfy the "never mid-callback" argument above, but this keeps the @@ -11204,7 +11513,10 @@ def submit(self, callable, args=None, config=None) -> RunHandle: caller whose first run only completes because a later callback runs would deadlock on a depth-one backend. Completion and cleanup stay attached to each handle. + + Program mode only; a kernel-mode Worker launches through ``kernel_launch``. """ + self._require_execution_mode("submit", "program") try: with self._operation_lease("submit"): result = self._submit_locked(callable, args, config) @@ -11224,6 +11536,7 @@ def run(self, callable, args=None, config=None) -> None: with ``simpler_setup.tools.strace_timing`` (see ``docs/dfx/host-trace.md``). """ + self._require_execution_mode("run", "program") self.submit(callable, args=args, config=config).wait() def _submit_locked(self, callable, args, config) -> RunHandle: @@ -11891,6 +12204,14 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r each resource until its native free succeeds and preserves the child pid/mailbox pair until ``waitpid`` proves the child is gone. - Native teardown runs on the ``init()``-owner thread, being device-bound. + - With ``execution_mode="kernel"`` it releases only the context: it never + synchronizes the caller's stream, resets the device or finalizes ACL. + It raises ``TimeoutError`` when a prepare or launch keeps the kernel + gate past the rollback grace period and ``RuntimeError`` when native + teardown fails; both keep the context for a later ``close()``. A + process other than the one that called ``init()`` is refused. An + unclosed kernel Worker leaks its context with a ``ResourceWarning`` + instead of holding the device through a finalize. """ # close() is a permanent commitment against a resource, not a reversible # attempt: it publishes CLOSED atomically (the sole public admission @@ -12372,31 +12693,51 @@ def _step(fn) -> None: def _finalize_chip() -> None: if self._chip_worker: - # Close the lane before finalizing the worker: a handle the - # caller never waited on still owns device work, and the - # lane drains it here while the device is still up. - # - # The lane rethrows its poison on close. Whether that is - # news depends on who has already seen it: waiting on a - # handle delivers the run's error and retires its entry, so - # a remaining entry is a run whose failure nobody has been - # told about, and only then is close the first report. With - # every run waited, the poison is the error those waits - # already raised, and re-raising it here would turn a - # handled run failure into an unhandled close failure. - impl = getattr(self._chip_worker, "_impl", None) - if impl is not None: - undelivered = bool(self._chip_runs) - try: - impl._close_chip_run_lane() - except Exception: - if undelivered: - raise - with self._registry_lock: - self._chip_runs.clear() - self._chip_run_touched_identities.clear() - self._chip_worker.finalize() - self._chip_worker = None + kernel = self._execution_mode == "kernel" + if kernel: + self._require_kernel_process("close") + # Prepare and launch each hold the gate across one native call. + if not self._kernel_gate.acquire(timeout=_ROLLBACK_GRACEFUL_TIMEOUT_S): + raise TimeoutError( + "Worker.close(): a kernel prepare/launch still holds the kernel gate after " + f"{_ROLLBACK_GRACEFUL_TIMEOUT_S}s; the context is kept, close() again" + ) + try: + # Close the lane before finalizing the worker: a handle the + # caller never waited on still owns device work, and the + # lane drains it here while the device is still up. + # + # The lane rethrows its poison on close. Whether that is + # news depends on who has already seen it: waiting on a + # handle delivers the run's error and retires its entry, so + # a remaining entry is a run whose failure nobody has been + # told about, and only then is close the first report. With + # every run waited, the poison is the error those waits + # already raised, and re-raising it here would turn a + # handled run failure into an unhandled close failure. + impl = getattr(self._chip_worker, "_impl", None) + if impl is not None: + undelivered = bool(self._chip_runs) + try: + impl._close_chip_run_lane() + except Exception: + if undelivered: + raise + with self._registry_lock: + self._chip_runs.clear() + self._chip_run_touched_identities.clear() + # A kernel context whose device teardown fails raises from finalize(), + # which keeps this journal entry, the ChipWorker, its prepared images and + # the GC pin for a later close(). Everything below is reached only once + # the context is actually gone. + self._chip_worker.finalize() + self._chip_worker = None + if kernel: + self._kernel_callables.clear() + self._disarm_kernel_chip_pin() + finally: + if kernel: + self._kernel_gate.release() self._cleanup_journal.add_once("native", "ChipWorker", _finalize_chip) journal_err = self._cleanup_journal.drive({("native", "ChipWorker")}) diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 1ac32cb9d8..b303b729fc 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -455,6 +455,38 @@ uint64_t ChipWorker::next_kernel_context_generation() { return counter.fetch_add(1, std::memory_order_relaxed) + 1; } +bool ChipWorker::probe_kernel_mode_supported(const std::string &host_lib_path, const std::string &sim_context_path) { + if (!sim_context_path.empty()) { + load_sim_context(sim_context_path); + } + + dlerror(); + void *handle = dlopen(host_lib_path.c_str(), RTLD_NOW | RTLD_LOCAL); + if (!handle) { + std::string err = "dlopen failed: "; + const char *msg = dlerror(); + err += msg ? msg : "unknown error"; + throw std::runtime_error(err); + } + DlHandleGuard host_guard(handle); + bind_host_log_state(handle, "host runtime"); + + const auto create_device_context_fn = load_symbol(handle, "create_device_context"); + const auto destroy_device_context_fn = load_symbol(handle, "destroy_device_context"); + const auto kernel_supported_fn = load_symbol(handle, "simpler_kernel_mode_supported"); + + void *ctx = create_device_context_fn(); + if (ctx == nullptr) { + throw std::runtime_error("create_device_context returned null"); + } + // simpler_kernel_mode_supported is a C entry and cannot throw, so nothing + // between create and destroy skips the destroy, and destroy runs before the + // guard unloads the library that implements it. + const int supported = kernel_supported_fn(ctx); + destroy_device_context_fn(ctx); + return supported != 0; +} + void ChipWorker::kernel_init( const std::string &host_lib_path, const std::string &aicpu_path, const std::string &aicore_path, const std::string &dispatcher_path, int device_id, const CallConfig &config, uint64_t context_generation, diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index d797441944..28ff32b274 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -163,6 +163,19 @@ class ChipWorker { /// counter starts at one. static uint64_t next_kernel_context_generation(); + /// Whether the host runtime at `host_lib_path` can execute kernel-mode + /// launches, answered without a worker. Retains the sim context first when + /// `sim_context_path` is non-empty (the same process-wide registry init() + /// uses), loads the runtime, asks simpler_kernel_mode_supported on a fresh + /// device context, destroys that context, and releases its runtime handle. + /// + /// The C ABI requires simpler_kernel_mode_supported to answer from the + /// runtime build alone on a context no init has touched, so this takes no + /// device, attaches no thread, and is independent of init() and + /// kernel_init() on any worker — usable before either. Throws when the + /// runtime cannot be loaded, lacks a required symbol, or yields no context. + static bool probe_kernel_mode_supported(const std::string &host_lib_path, const std::string &sim_context_path); + /// Tear down everything: device resources and runtime library. The worker /// cannot be initialized again afterwards. When a kernel context's device /// teardown fails, this throws ChipWorkerError and keeps the context and the diff --git a/tests/ut/py/test_chip_worker.py b/tests/ut/py/test_chip_worker.py index 23bba8f4d4..0016feac81 100644 --- a/tests/ut/py/test_chip_worker.py +++ b/tests/ut/py/test_chip_worker.py @@ -8,11 +8,14 @@ # ----------------------------------------------------------------------------------------------------------- """Tests for CallConfig and ChipWorker state machine.""" +import ctypes import json import os import shutil import subprocess +import sys import threading +import types from pathlib import Path import pytest @@ -240,15 +243,20 @@ def build( "#include \n" "#include \n" "static int live_contexts = 0;\n" + "static int created_contexts = 0;\n" "struct ContextLeakCheck {\n" " ~ContextLeakCheck() { if (live_contexts != 0) std::abort(); }\n" "};\n" "static ContextLeakCheck context_leak_check;\n" + # A caller that holds the DSO open reads context lifetime through these two. + 'extern "C" int fake_live_contexts() { return live_contexts; }\n' + 'extern "C" int fake_created_contexts() { return created_contexts; }\n' "struct SimplerHostLogState;\n" 'extern "C" int simpler_host_log_bind_state(SimplerHostLogState *) { return 0; }\n' "static std::unordered_set live_handles;\n" "DeviceContextHandle create_device_context() {\n" - " ++live_contexts; auto *ctx = new uint64_t{0}; live_handles.insert(ctx); return ctx;\n" + " ++live_contexts; ++created_contexts; auto *ctx = new uint64_t{0};\n" + " live_handles.insert(ctx); return ctx;\n" "}\n" "void destroy_device_context(DeviceContextHandle ctx) {\n" " --live_contexts; live_handles.erase(ctx); delete static_cast(ctx);\n" @@ -615,6 +623,92 @@ def test_program_teardown_failure_is_not_raised(self, kernel_symbol_runtime): assert not worker.initialized +class TestChipWorkerKernelProbe: + @staticmethod + def _context_counters(runtime): + # ctypes never dlcloses, so the counters outlive the probe's own handle on the DSO. + library = ctypes.CDLL(str(runtime)) + return library.fake_live_contexts, library.fake_created_contexts + + @pytest.mark.parametrize(("supported", "expected"), ((0, False), (1, True))) + def test_probe_reports_runtime_capability(self, kernel_symbol_runtime, supported, expected): + runtime = kernel_symbol_runtime(supported=supported) + assert _ChipWorker.probe_kernel_mode_supported(str(runtime), "") is expected + + @pytest.mark.parametrize("supported", (0, 1)) + def test_probe_destroys_the_context_it_created(self, kernel_symbol_runtime, supported): + runtime = kernel_symbol_runtime(supported=supported) + live, created = self._context_counters(runtime) + created_before = created() + + _ChipWorker.probe_kernel_mode_supported(str(runtime), "") + + assert created() == created_before + 1 + assert live() == 0 + + def test_probe_missing_capability_symbol_raises_before_creating_a_context(self, kernel_symbol_runtime): + runtime = kernel_symbol_runtime(supported=1, missing=("simpler_kernel_mode_supported",)) + live, created = self._context_counters(runtime) + created_before = created() + + with pytest.raises(RuntimeError, match="dlsym failed for 'simpler_kernel_mode_supported'"): + _ChipWorker.probe_kernel_mode_supported(str(runtime), "") + + assert created() == created_before + assert live() == 0 + + def test_probe_nonexistent_library_raises(self): + with pytest.raises(RuntimeError, match="dlopen failed"): + _ChipWorker.probe_kernel_mode_supported("/nonexistent/libfoo.so", "") + + def test_probe_in_fresh_process_needs_no_init(self, kernel_symbol_runtime): + runtime = kernel_symbol_runtime(supported=1) + # The fake DSO aborts when unloaded with a context still live, and the + # probe unloads it before returning, so a leaked context kills this subprocess. + code = ( + "import sys, types\n" + "from simpler.task_interface import ChipWorker\n" + "bins = types.SimpleNamespace(host_path=sys.argv[1], sim_context_path=None)\n" + "print(ChipWorker.probe_kernel_mode_supported(bins))\n" + ) + completed = subprocess.run( + [sys.executable, "-c", code, str(runtime)], capture_output=True, text=True, check=False, timeout=120 + ) + assert completed.returncode == 0, f"{completed.stdout!r} {completed.stderr!r}" + assert completed.stdout.strip().splitlines()[-1] == "True", f"{completed.stdout!r} {completed.stderr!r}" + + def test_public_wrapper_maps_bins_to_the_native_probe(self, monkeypatch): + import simpler.task_interface as task_interface_mod # noqa: PLC0415 + from simpler.task_interface import ChipWorker # noqa: PLC0415 # pyright: ignore[reportAttributeAccessIssue] + + probes = [] + seeded_levels = [] + + class FakeNative: + answers = [1, 0] + + @staticmethod + def probe_kernel_mode_supported(host_lib_path, sim_context_path): + probes.append((host_lib_path, sim_context_path)) + return FakeNative.answers[len(probes) - 1] + + monkeypatch.setattr(task_interface_mod, "_ChipWorker", FakeNative) + monkeypatch.setattr(task_interface_mod, "_initialize_host_log", seeded_levels.append) + + onboard = types.SimpleNamespace(host_path=Path("/rt/libhost_runtime.so"), sim_context_path=None) + sim = types.SimpleNamespace( + host_path=Path("/rt/libhost_runtime.so"), sim_context_path=Path("/rt/libcpu_sim_context.so") + ) + + assert ChipWorker.probe_kernel_mode_supported(onboard, log_level=20) is True + assert ChipWorker.probe_kernel_mode_supported(sim) is False + assert probes == [ + ("/rt/libhost_runtime.so", ""), + ("/rt/libhost_runtime.so", "/rt/libcpu_sim_context.so"), + ] + assert seeded_levels == [20, None] + + class TestChipWorkerStateMachine: def test_initial_state(self): worker = _ChipWorker() diff --git a/tests/ut/py/test_worker/test_worker_kernel_mode.py b/tests/ut/py/test_worker/test_worker_kernel_mode.py new file mode 100644 index 0000000000..c3e813f917 --- /dev/null +++ b/tests/ut/py/test_worker/test_worker_kernel_mode.py @@ -0,0 +1,873 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Device-free UT for ``Worker(level=2, execution_mode="kernel")``. + +``simpler.worker.ChipWorker`` and ``simpler_setup.runtime_builder.RuntimeBuilder`` +are replaced by fakes, so every case runs without an NPU or a built runtime. The +fake kernel ChipWorker records what the Worker forwards to it and fails loudly +on any program-mode entry, which is how a case proves a kernel-mode Worker never +reaches the program path. ``TestSimRuntime`` is the one case that loads a real +runtime build; it skips when the a2a3sim binaries are absent. +""" + +from __future__ import annotations + +import ctypes +import gc +import os +import re +import threading +import warnings +import weakref +from collections.abc import Callable +from typing import Any, cast + +import pytest +import simpler.worker as worker_mod +from simpler.task_interface import CallConfig, ChipStorageTaskArgs, ChipWorkerError +from simpler.worker import Worker + +import simpler_setup.runtime_builder as rb_mod + +from ._harness import SIM_PLATFORM, SIM_RUNTIME, TEST_WALL_BUDGET_S, chip_callable, hard_timeout, install_fake_chip + +_DEVICE_ID = 3 +_STREAM = 0xABC0 +_Lifecycle = worker_mod._Lifecycle + +#: A stand-in for an argument of the wrong type; the refusals under test run before it is used. +_ANY: Any = object() + + +def _run_catch(fn: Callable[[], Any]) -> BaseException | None: + """Run ``fn`` in a thread body, returning None on success or the exception.""" + try: + fn() + return None + except BaseException as e: # noqa: BLE001 + return e + + +class _KernelScript: + """What the fake kernel ChipWorker and runtime builder do; a case sets it before the call it scripts.""" + + def __init__(self) -> None: + self.binaries = object() + self.builder_calls: list[tuple[str, str]] = [] + self.builder_error: BaseException | None = None + self.probe_result = True + self.probe_calls: list[Any] = [] + self.kernel_init_error: BaseException | None = None + # Ids kernel_prepare_callable returns, in order; once exhausted each chip counts up from 0. + self.prepare_ids: list[int] = [] + self.prepare_entered = threading.Event() + self.prepare_release: threading.Event | None = None + self.launch_entered = threading.Event() + self.launch_release: threading.Event | None = None + self.launch_hook: Callable[[], None] | None = None + # Number of finalize() calls that raise ChipWorkerError for a failed device teardown. + self.failed_finalizes = 0 + self.chips: list[_FakeKernelChip] = [] + + +class _FakeKernelImpl: + """The native-handle half of :class:`_FakeKernelChip` (``chip._impl``).""" + + def __init__(self) -> None: + self.initialized = False + self.lane_closes = 0 + + def _close_chip_run_lane(self) -> None: + self.lane_closes += 1 + + def register_callable_from_blob(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached native register_callable_from_blob") + + def run_materialized(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached native run_materialized") + + def _submit_chip_run_direct(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached native _submit_chip_run_direct") + + +class _FakeKernelChip: + """Stand-in for ``simpler.task_interface.ChipWorker`` driven through its kernel-mode surface.""" + + pipeline_depth = 1 + committed_device_memory = 4096 + + def __init__(self, script: _KernelScript) -> None: + self._script = script + self._impl = _FakeKernelImpl() + self._next_id = 0 + self.kernel_init_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.prepared: list[tuple[int, Any]] = [] + self.launches: list[tuple[int, Any, int]] = [] + self.finalize_threads: list[threading.Thread] = [] + # len(finalize_threads) observed by each launch as it returns. + self.finalizes_at_launch_return: list[int] = [] + script.chips.append(self) + + def kernel_init(self, *args: Any, **kwargs: Any) -> None: + self.kernel_init_calls.append((args, kwargs)) + if self._script.kernel_init_error is not None: + raise self._script.kernel_init_error + self._impl.initialized = True + + def kernel_prepare_callable(self, chip_callable_obj: Any) -> int: + script = self._script + script.prepare_entered.set() + if script.prepare_release is not None: + assert script.prepare_release.wait(TEST_WALL_BUDGET_S) + if script.prepare_ids: + callable_id = script.prepare_ids.pop(0) + else: + callable_id = self._next_id + self._next_id += 1 + self.prepared.append((callable_id, chip_callable_obj)) + return callable_id + + def kernel_launch(self, callable_id: int, args: Any, caller_stream: int) -> None: + script = self._script + if script.launch_hook is not None: + script.launch_hook() + script.launch_entered.set() + if script.launch_release is not None: + assert script.launch_release.wait(TEST_WALL_BUDGET_S) + self.launches.append((callable_id, args, caller_stream)) + self.finalizes_at_launch_return.append(len(self.finalize_threads)) + + def finalize(self) -> None: + self.finalize_threads.append(threading.current_thread()) + # A kernel context clears initialized and records the owed teardown before it raises, + # so a retry is a second finalize() rather than a re-init. + self._impl.initialized = False + if self._script.failed_finalizes > 0: + self._script.failed_finalizes -= 1 + raise ChipWorkerError(-77, "ChipWorker::finalize: device teardown failed (-77)") + + def init(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker.init") + + def malloc(self, *_a, **_k) -> int: + raise AssertionError("kernel-mode Worker reached ChipWorker.malloc") + + def free(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker.free") + + def copy_to(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker.copy_to") + + def copy_from(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker.copy_from") + + def device_memory_info(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker.device_memory_info") + + def _register_callable_at_slot(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker._register_callable_at_slot") + + def _unregister_slot(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker._unregister_slot") + + def _run_slot(self, *_a, **_k) -> None: + raise AssertionError("kernel-mode Worker reached ChipWorker._run_slot") + + +class _ObservedGate: + """A kernel gate that reports when a blocking acquire starts; non-blocking attempts are not reported.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self.blocking_acquire = threading.Event() + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if blocking: + self.blocking_acquire.set() + return self._lock.acquire(blocking, timeout) + + def release(self) -> None: + self._lock.release() + + def __enter__(self) -> bool: + return self.acquire() + + def __exit__(self, *_exc_info: Any) -> None: + self.release() + + +class _OtherProcessOs: + """``os`` as a forked child of the test process sees it: another pid, everything else real.""" + + def __init__(self, pid: int) -> None: + self._pid = pid + + def getpid(self) -> int: + return self._pid + + def __getattr__(self, name: str) -> Any: + return getattr(os, name) + + +@pytest.fixture +def script(monkeypatch) -> _KernelScript: + """Bind the fake kernel ChipWorker and runtime builder, and give the GC pin list a per-test instance.""" + kernel_script = _KernelScript() + + class _Chip(_FakeKernelChip): + def __init__(self) -> None: + super().__init__(kernel_script) + + @staticmethod + def probe_kernel_mode_supported(bins: Any, log_level: int | None = None) -> bool: + kernel_script.probe_calls.append(bins) + return kernel_script.probe_result + + class _Builder: + def __init__(self, platform: str, *_a, **_k) -> None: + self._platform = platform + + def get_binaries(self, runtime: str, *_a, **_k) -> Any: + kernel_script.builder_calls.append((self._platform, runtime)) + if kernel_script.builder_error is not None: + raise kernel_script.builder_error + return kernel_script.binaries + + monkeypatch.setattr(worker_mod, "ChipWorker", _Chip) + monkeypatch.setattr(rb_mod, "RuntimeBuilder", _Builder) + monkeypatch.setattr(worker_mod, "_PINNED_KERNEL_CHIP_WORKERS", []) + return kernel_script + + +def _kernel_worker(**config: Any) -> Worker: + return Worker( + level=2, + execution_mode="kernel", + device_id=_DEVICE_ID, + platform=SIM_PLATFORM, + runtime=SIM_RUNTIME, + **config, + ) + + +def _program_worker() -> Worker: + return Worker(level=2, device_id=_DEVICE_ID, platform=SIM_PLATFORM, runtime=SIM_RUNTIME) + + +def _ready_kernel_worker(script: _KernelScript) -> tuple[Worker, _FakeKernelChip]: + worker = _kernel_worker() + worker.init(config=CallConfig()) + assert len(script.chips) == 1 + return worker, script.chips[0] + + +class _RejectingConfig: + def validate(self) -> None: + raise ValueError("injected CallConfig rejection") + + +class TestConstruction: + def test_default_mode_is_program(self, script): + assert _program_worker()._execution_mode == "program" + explicit = Worker(level=2, execution_mode="program", platform=SIM_PLATFORM, runtime=SIM_RUNTIME) + assert explicit._execution_mode == "program" + + def test_unknown_mode_raises(self, script): + with pytest.raises(ValueError, match="execution_mode must be 'program' or 'kernel'"): + Worker(level=2, execution_mode="graph", platform=SIM_PLATFORM, runtime=SIM_RUNTIME) + + def test_kernel_requires_level_2(self, script): + with pytest.raises(ValueError, match=re.escape("execution_mode='kernel' requires level=2")): + Worker(level=3, execution_mode="kernel", num_sub_workers=0) + + def test_kernel_refuses_sdma(self, script): + with pytest.raises(ValueError, match="enable_sdma"): + _kernel_worker(enable_sdma=True) + assert _kernel_worker(enable_sdma=False)._execution_mode == "kernel" + + def test_construction_touches_no_runtime(self, script): + _kernel_worker() + assert script.chips == [] + assert script.builder_calls == [] + assert script.probe_calls == [] + + +class TestInitArguments: + """A refused init() argument spends no startup epoch: the Worker stays NEW and a correct init() follows.""" + + def _assert_untouched(self, worker: Worker, script: _KernelScript) -> None: + assert worker._lifecycle is _Lifecycle.NEW + assert script.chips == [] + assert script.builder_calls == [] + + def test_kernel_mode_requires_config(self, script): + worker = _kernel_worker() + with pytest.raises(ValueError, match=re.escape("config=CallConfig(...)")): + worker.init() + self._assert_untouched(worker, script) + worker.init(config=CallConfig()) + assert worker._lifecycle is _Lifecycle.READY + worker.close() + + def test_positional_config_binds_to_prewarm_config(self, script): + worker = _kernel_worker() + with pytest.raises(ValueError, match="positional argument binds to prewarm_config"): + worker.init(CallConfig()) + self._assert_untouched(worker, script) + worker.close() + + def test_kernel_mode_refuses_prewarm_config(self, script): + worker = _kernel_worker() + with pytest.raises(ValueError, match="takes no prewarm_config"): + worker.init(CallConfig(), config=CallConfig()) + self._assert_untouched(worker, script) + worker.close() + + def test_kernel_config_is_validated_before_startup(self, script): + worker = _kernel_worker() + with pytest.raises(ValueError, match="injected CallConfig rejection"): + worker.init(config=cast(Any, _RejectingConfig())) + self._assert_untouched(worker, script) + worker.close() + + def test_program_mode_refuses_config(self, script): + worker = _program_worker() + with pytest.raises(ValueError, match="only accepted with execution_mode='kernel'"): + worker.init(config=CallConfig()) + self._assert_untouched(worker, script) + worker.close() + + +class TestInitRouting: + def test_init_binds_through_kernel_init_only(self, script): + config = CallConfig() + worker = _kernel_worker() + worker.init(config=config) + try: + assert worker._lifecycle is _Lifecycle.READY + assert script.builder_calls == [(SIM_PLATFORM, SIM_RUNTIME)] + (chip,) = script.chips + assert worker._chip_worker is chip + # No context_generation: ChipWorker.kernel_init mints it. + assert chip.kernel_init_calls == [((_DEVICE_ID, script.binaries, config), {})] + assert chip._impl.initialized + assert worker._callable_registry == {} + assert worker._identity_registry == {} + assert chip.prepared == [] + finally: + worker.close() + + def test_failed_kernel_init_rolls_back_to_failed(self, script): + script.kernel_init_error = RuntimeError("injected kernel_init failure") + worker = _kernel_worker() + with pytest.raises(RuntimeError, match="injected kernel_init failure"): + worker.init(config=CallConfig()) + (chip,) = script.chips + assert worker._lifecycle is _Lifecycle.FAILED + assert chip.finalize_threads == [threading.current_thread()] + assert worker._chip_worker is None + assert worker._kernel_pid is None + assert worker._kernel_chip_pin is None + assert worker._kernel_pin_finalizer is None + + worker.close() + assert len(chip.finalize_threads) == 1 + with pytest.raises(RuntimeError, match="closed"): + worker.init(config=CallConfig()) + + def test_failed_kernel_init_arms_no_gc_pin(self, script): + script.kernel_init_error = RuntimeError("injected kernel_init failure") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + worker = _kernel_worker() + with pytest.raises(RuntimeError, match="injected kernel_init failure"): + worker.init(config=CallConfig()) + # The script holds the raised exception, whose traceback reaches this Worker's init + # frames; without dropping it the Worker stays reachable from the fixture. + script.kernel_init_error = None + ref = weakref.ref(worker) + del worker + gc.collect() + assert ref() is None + assert not [w for w in caught if issubclass(w.category, ResourceWarning)] + assert worker_mod._PINNED_KERNEL_CHIP_WORKERS == [] + + +class TestKernelModeSupported: + @pytest.mark.parametrize("answer", [True, False]) + def test_probe_before_init_reports_runtime_answer(self, script, answer): + script.probe_result = answer + worker = _kernel_worker() + assert worker.kernel_mode_supported is answer + assert script.builder_calls == [(SIM_PLATFORM, SIM_RUNTIME)] + assert script.probe_calls == [script.binaries] + assert script.chips == [] + assert worker._lifecycle is _Lifecycle.NEW + + def test_probe_answer_is_cached(self, script): + worker = _kernel_worker() + assert worker.kernel_mode_supported is True + script.probe_result = False + assert worker.kernel_mode_supported is True + assert len(script.probe_calls) == 1 + + def test_program_mode_worker_probes_the_same_build(self, script): + script.probe_result = True + assert _program_worker().kernel_mode_supported is True + assert script.probe_calls == [script.binaries] + + def test_level_3_is_false_without_probing(self, script): + assert Worker(level=3, num_sub_workers=0).kernel_mode_supported is False + assert script.builder_calls == [] + assert script.probe_calls == [] + + def test_ready_kernel_worker_is_true_without_probing(self, script): + script.probe_result = False + worker, _chip = _ready_kernel_worker(script) + try: + assert worker.kernel_mode_supported is True + assert script.probe_calls == [] + finally: + worker.close() + + def test_missing_binaries_propagate_and_are_not_cached(self, script): + script.builder_error = FileNotFoundError("injected missing runtime build") + worker = _kernel_worker() + with pytest.raises(FileNotFoundError, match="injected missing runtime build"): + _ = worker.kernel_mode_supported + assert script.probe_calls == [] + script.builder_error = None + assert worker.kernel_mode_supported is True + assert len(script.probe_calls) == 1 + + +class TestPrepare: + def test_same_callable_twice_takes_distinct_ids(self, script): + worker, chip = _ready_kernel_worker(script) + try: + target = chip_callable() + first = worker.kernel_prepare_callable(target) + second = worker.kernel_prepare_callable(target) + assert (first, second) == (0, 1) + assert chip.prepared == [(0, target), (1, target)] + assert worker._kernel_callables == {0: target, 1: target} + assert worker._callable_registry == {} + assert worker._identity_registry == {} + assert worker._live_handles == {} + assert worker._active_ops == 0 + finally: + worker.close() + assert worker._kernel_callables == {} + + def test_non_chip_callable_is_refused_before_native(self, script): + worker, chip = _ready_kernel_worker(script) + try: + with pytest.raises(TypeError, match="expected a ChipCallable"): + worker.kernel_prepare_callable(_ANY) + assert chip.prepared == [] + finally: + worker.close() + + def test_refused_before_init(self, script): + worker = _kernel_worker() + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_prepare_callable(chip_callable()) + assert script.chips == [] + worker.close() + + def test_refused_after_close(self, script): + worker, chip = _ready_kernel_worker(script) + worker.close() + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_prepare_callable(chip_callable()) + assert chip.prepared == [] + + def test_negative_id_raises_and_retains_the_callable(self, script): + worker, chip = _ready_kernel_worker(script) + try: + script.prepare_ids = [-1] + target = chip_callable() + with pytest.raises(RuntimeError, match="invalid id -1"): + worker.kernel_prepare_callable(target) + assert chip.prepared == [(-1, target)] + assert worker._kernel_callables == {-1: target} + assert worker._active_ops == 0 + assert worker._kernel_gate.acquire(blocking=False) + worker._kernel_gate.release() + with pytest.raises(ValueError, match="was not returned by this Worker's kernel_prepare_callable"): + worker.kernel_launch(-1, ChipStorageTaskArgs(), caller_stream=_STREAM) + assert chip.launches == [] + assert worker._kernel_callables == {-1: target} + finally: + worker.close() + assert worker._kernel_callables == {} + + +class TestLaunch: + def test_launch_forwards_and_takes_no_lease(self, script): + worker, chip = _ready_kernel_worker(script) + try: + callable_id = worker.kernel_prepare_callable(chip_callable()) + args = ChipStorageTaskArgs() + observed: dict[str, Any] = {} + script.launch_hook = lambda: observed.update( + active_ops=worker._active_ops, lease_depth=dict(worker._lease_depth) + ) + assert worker.kernel_launch(callable_id, args, caller_stream=_STREAM) is None + assert len(chip.launches) == 1 + launched_id, launched_args, launched_stream = chip.launches[0] + assert (launched_id, launched_stream) == (callable_id, _STREAM) + assert launched_args is args + assert observed == {"active_ops": 0, "lease_depth": {}} + assert worker._chip_runs == {} + assert worker._accepted_run_handles == set() + finally: + worker.close() + + def test_bad_arguments_are_refused_without_native_calls(self, script): + worker, chip = _ready_kernel_worker(script) + try: + callable_id = worker.kernel_prepare_callable(chip_callable()) + args = ChipStorageTaskArgs() + with pytest.raises(ValueError, match="non-null caller_stream"): + worker.kernel_launch(callable_id, args, caller_stream=0) + with pytest.raises(TypeError, match="args must be ChipStorageTaskArgs"): + worker.kernel_launch(callable_id, _ANY, caller_stream=_STREAM) + with pytest.raises(ValueError, match="was not returned by this Worker's kernel_prepare_callable"): + worker.kernel_launch(callable_id + 7, args, caller_stream=_STREAM) + with pytest.raises(TypeError): + cast(Any, worker).kernel_launch(callable_id, args, _STREAM) + with pytest.raises(TypeError): + worker.kernel_launch(cast(Any, callable_id + 0.5), args, caller_stream=_STREAM) + assert chip.launches == [] + finally: + worker.close() + + def test_refused_before_init(self, script): + worker = _kernel_worker() + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_launch(0, ChipStorageTaskArgs(), caller_stream=_STREAM) + assert script.chips == [] + worker.close() + + def test_refused_after_close(self, script): + worker, chip = _ready_kernel_worker(script) + callable_id = worker.kernel_prepare_callable(chip_callable()) + worker.close() + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_launch(callable_id, ChipStorageTaskArgs(), caller_stream=_STREAM) + assert chip.launches == [] + + +_PROGRAM_ONLY_CALLS: dict[str, Callable[[Worker], Any]] = { + "register": lambda w: w.register(chip_callable()), + "unregister": lambda w: w.unregister(0), + "submit": lambda w: w.submit(chip_callable()), + "run": lambda w: w.run(chip_callable()), + "malloc": lambda w: w.malloc(16), + "free": lambda w: w.free(_ANY), + "copy_to": lambda w: w.copy_to(_ANY, _ANY), + "copy_from": lambda w: w.copy_from(_ANY, _ANY), + "create_buffer": lambda w: w.create_buffer(16), + "make_tensor_arg": lambda w: w.make_tensor_arg(_ANY, (1,), 0), + "release_buffer": lambda w: w.release_buffer(_ANY), + "device_memory_info": lambda w: w.device_memory_info(), +} + + +class TestModeGuards: + @pytest.mark.parametrize("api", sorted(_PROGRAM_ONLY_CALLS)) + def test_program_only_api_refuses_in_kernel_mode(self, script, api): + worker, chip = _ready_kernel_worker(script) + try: + with pytest.raises(RuntimeError, match=re.escape(f"Worker.{api}: requires execution_mode='program'")): + _PROGRAM_ONLY_CALLS[api](worker) + assert worker._callable_registry == {} + assert chip.prepared == [] + finally: + worker.close() + + def test_register_refuses_in_kernel_mode_at_new(self, script): + worker = _kernel_worker() + with pytest.raises(RuntimeError, match=re.escape("Worker.register: requires execution_mode='program'")): + worker.register(chip_callable()) + assert worker._callable_registry == {} + assert worker._identity_registry == {} + worker.init(config=CallConfig()) + worker.close() + + def test_committed_device_memory_stays_available(self, script): + worker, _chip = _ready_kernel_worker(script) + try: + assert worker.committed_device_memory() == _FakeKernelChip.committed_device_memory + finally: + worker.close() + + def test_kernel_apis_refuse_in_program_mode_at_new(self, script): + worker = _program_worker() + with pytest.raises(RuntimeError, match=re.escape("requires execution_mode='kernel'")): + worker.kernel_prepare_callable(chip_callable()) + with pytest.raises(RuntimeError, match=re.escape("requires execution_mode='kernel'")): + worker.kernel_launch(0, ChipStorageTaskArgs(), caller_stream=0) + worker.close() + + def test_kernel_apis_refuse_on_a_ready_program_worker(self, script, monkeypatch): + install_fake_chip(monkeypatch) + worker = _program_worker() + worker.init() + try: + with pytest.raises(RuntimeError, match=re.escape("requires execution_mode='kernel'")): + worker.kernel_prepare_callable(chip_callable()) + with pytest.raises(RuntimeError, match=re.escape("requires execution_mode='kernel'")): + worker.kernel_launch(0, ChipStorageTaskArgs(), caller_stream=_STREAM) + finally: + worker.close() + + +class TestKernelGate: + def test_launch_during_blocked_prepare_fails_fast(self, script): + with hard_timeout(TEST_WALL_BUDGET_S): + worker, chip = _ready_kernel_worker(script) + try: + callable_id = worker.kernel_prepare_callable(chip_callable()) + args = ChipStorageTaskArgs() + script.prepare_entered.clear() + script.prepare_release = threading.Event() + prepare_result: list[BaseException | None] = [] + preparer = threading.Thread( + target=lambda: prepare_result.append( + _run_catch(lambda: worker.kernel_prepare_callable(chip_callable("second"))) + ) + ) + preparer.start() + try: + assert script.prepare_entered.wait(TEST_WALL_BUDGET_S) + with pytest.raises(RuntimeError, match="a kernel prepare/launch/close is in progress"): + worker.kernel_launch(callable_id, args, caller_stream=_STREAM) + assert chip.launches == [] + finally: + script.prepare_release.set() + preparer.join(TEST_WALL_BUDGET_S) + assert prepare_result == [None] + worker.kernel_launch(callable_id, args, caller_stream=_STREAM) + assert len(chip.launches) == 1 + finally: + worker.close() + + def test_close_during_blocked_launch_finalizes_after_it_returns(self, script): + with hard_timeout(TEST_WALL_BUDGET_S): + worker = _kernel_worker() + gate = _ObservedGate() + worker._kernel_gate = gate + worker.init(config=CallConfig()) + (chip,) = script.chips + callable_id = worker.kernel_prepare_callable(chip_callable()) + gate.blocking_acquire.clear() + script.launch_release = threading.Event() + launch_result: list[BaseException | None] = [] + finalizes_before_release: list[int] = [] + + def release_once_close_waits() -> None: + if gate.blocking_acquire.wait(TEST_WALL_BUDGET_S): + finalizes_before_release.append(len(chip.finalize_threads)) + assert script.launch_release is not None + script.launch_release.set() + + launcher = threading.Thread( + target=lambda: launch_result.append( + _run_catch(lambda: worker.kernel_launch(callable_id, ChipStorageTaskArgs(), caller_stream=_STREAM)) + ) + ) + releaser = threading.Thread(target=release_once_close_waits) + launcher.start() + try: + assert script.launch_entered.wait(TEST_WALL_BUDGET_S) + releaser.start() + worker.close() + finally: + script.launch_release.set() + launcher.join(TEST_WALL_BUDGET_S) + if releaser.ident is not None: + releaser.join(TEST_WALL_BUDGET_S) + assert launch_result == [None] + assert finalizes_before_release == [0] + assert chip.finalizes_at_launch_return == [0] + assert chip.finalize_threads == [threading.current_thread()] + assert worker._chip_worker is None + + def test_close_past_the_gate_budget_keeps_the_context_for_retry(self, script, monkeypatch): + monkeypatch.setattr(worker_mod, "_ROLLBACK_GRACEFUL_TIMEOUT_S", 0.2) + with hard_timeout(TEST_WALL_BUDGET_S): + worker, chip = _ready_kernel_worker(script) + callable_id = worker.kernel_prepare_callable(chip_callable()) + script.launch_release = threading.Event() + launch_result: list[BaseException | None] = [] + launcher = threading.Thread( + target=lambda: launch_result.append( + _run_catch(lambda: worker.kernel_launch(callable_id, ChipStorageTaskArgs(), caller_stream=_STREAM)) + ) + ) + launcher.start() + try: + assert script.launch_entered.wait(TEST_WALL_BUDGET_S) + with pytest.raises(TimeoutError, match=re.escape("close() again")): + worker.close() + assert worker._lifecycle is _Lifecycle.CLOSED + assert worker._chip_worker is chip + assert chip.finalize_threads == [] + assert callable_id in worker._kernel_callables + finally: + script.launch_release.set() + launcher.join(TEST_WALL_BUDGET_S) + assert launch_result == [None] + worker.close() + assert len(chip.finalize_threads) == 1 + assert worker._chip_worker is None + + +class TestClose: + def test_close_finalizes_once_on_the_init_thread(self, script): + worker, chip = _ready_kernel_worker(script) + worker.kernel_prepare_callable(chip_callable()) + finalizer = worker._kernel_pin_finalizer + assert finalizer is not None and finalizer.alive + worker.close() + assert worker._lifecycle is _Lifecycle.CLOSED + assert chip.finalize_threads == [threading.current_thread()] + assert chip._impl.lane_closes == 1 + assert worker._chip_worker is None + assert worker._kernel_callables == {} + assert not finalizer.alive + assert worker._kernel_pin_finalizer is None + worker.close() + assert len(chip.finalize_threads) == 1 + + def test_non_owner_thread_close_raises(self, script): + with hard_timeout(TEST_WALL_BUDGET_S): + worker, chip = _ready_kernel_worker(script) + try: + close_result: list[BaseException | None] = [] + closer = threading.Thread(target=lambda: close_result.append(_run_catch(worker.close))) + closer.start() + closer.join(TEST_WALL_BUDGET_S) + assert len(close_result) == 1 + assert isinstance(close_result[0], RuntimeError) + assert "thread that init()'d it" in str(close_result[0]) + assert worker._lifecycle is _Lifecycle.READY + assert chip.finalize_threads == [] + callable_id = worker.kernel_prepare_callable(chip_callable()) + worker.kernel_launch(callable_id, ChipStorageTaskArgs(), caller_stream=_STREAM) + finally: + worker.close() + assert chip.finalize_threads == [threading.current_thread()] + + def test_failed_teardown_keeps_the_context_for_retry(self, script): + worker, chip = _ready_kernel_worker(script) + target = chip_callable() + callable_id = worker.kernel_prepare_callable(target) + script.failed_finalizes = 1 + + with pytest.raises(ChipWorkerError, match="device teardown failed"): + worker.close() + assert worker._lifecycle is _Lifecycle.CLOSED + assert len(chip.finalize_threads) == 1 + assert worker._chip_worker is chip + assert worker._kernel_callables == {callable_id: target} + finalizer = worker._kernel_pin_finalizer + assert finalizer is not None and finalizer.alive + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_launch(callable_id, ChipStorageTaskArgs(), caller_stream=_STREAM) + with pytest.raises(RuntimeError, match=r"requires an initialized \(READY\) worker"): + worker.kernel_prepare_callable(target) + assert chip.launches == [] + + worker.close() + assert len(chip.finalize_threads) == 2 + assert worker._chip_worker is None + assert worker._kernel_callables == {} + assert not finalizer.alive + worker.close() + assert len(chip.finalize_threads) == 2 + + +class TestGcPin: + @staticmethod + def _abandon_ready_kernel_worker(script: _KernelScript) -> tuple[weakref.ref, _FakeKernelChip]: + worker, chip = _ready_kernel_worker(script) + return weakref.ref(worker), chip + + def test_unclosed_ready_worker_pins_its_chip_instead_of_finalizing(self, script): + with pytest.warns(ResourceWarning, match="garbage-collected without close"): + ref, chip = self._abandon_ready_kernel_worker(script) + gc.collect() + try: + assert ref() is None + assert worker_mod._PINNED_KERNEL_CHIP_WORKERS == [chip] + assert chip.finalize_threads == [] + assert chip._impl.initialized + finally: + # The pin takes a reference it never releases; returning it lets the fake chip, and the + # CallConfig it recorded, go away with this test instead of outliving the process. + ctypes.pythonapi.Py_DecRef(ctypes.py_object(chip)) + + def test_closed_worker_is_not_pinned(self, script): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + worker, chip = _ready_kernel_worker(script) + worker.close() + ref = weakref.ref(worker) + del worker + gc.collect() + assert ref() is None + assert not [w for w in caught if issubclass(w.category, ResourceWarning)] + assert worker_mod._PINNED_KERNEL_CHIP_WORKERS == [] + assert len(chip.finalize_threads) == 1 + + +class TestProcessFence: + def test_other_process_cannot_prepare_launch_or_close(self, script, monkeypatch): + worker, chip = _ready_kernel_worker(script) + callable_id = worker.kernel_prepare_callable(chip_callable()) + args = ChipStorageTaskArgs() + with monkeypatch.context() as patch: + patch.setattr(worker_mod, "os", _OtherProcessOs(os.getpid() + 1)) + with pytest.raises(RuntimeError, match="must not drive or tear it down"): + worker.kernel_prepare_callable(chip_callable()) + with pytest.raises(RuntimeError, match="must not drive or tear it down"): + worker.kernel_launch(callable_id, args, caller_stream=_STREAM) + with pytest.raises(RuntimeError, match="must not drive or tear it down"): + worker.close() + assert len(chip.prepared) == 1 + assert chip.launches == [] + assert chip.finalize_threads == [] + assert worker._chip_worker is chip + + worker.close() + assert chip.finalize_threads == [threading.current_thread()] + assert worker._chip_worker is None + + +class TestSimRuntime: + """The real a2a3sim tensormap_and_ringbuffer build, which has no kernel mode.""" + + def test_sim_runtime_reports_no_kernel_mode_and_init_rolls_back(self): + from simpler_setup.runtime_builder import RuntimeBuilder # noqa: PLC0415 + + try: + RuntimeBuilder(platform=SIM_PLATFORM).get_binaries(SIM_RUNTIME) + except FileNotFoundError as e: + pytest.skip(f"{SIM_PLATFORM} runtime binaries unavailable: {e}") + + worker = Worker(level=2, execution_mode="kernel", platform=SIM_PLATFORM, runtime=SIM_RUNTIME) + assert worker.kernel_mode_supported is False + with pytest.raises(RuntimeError, match="kernel mode"): + worker.init(config=CallConfig()) + assert worker._lifecycle is _Lifecycle.FAILED + assert worker._chip_worker is None + worker.close() diff --git a/tests/ut/py/test_worker/test_worker_kernel_mode_hw.py b/tests/ut/py/test_worker/test_worker_kernel_mode_hw.py new file mode 100644 index 0000000000..98bf42924e --- /dev/null +++ b/tests/ut/py/test_worker/test_worker_kernel_mode_hw.py @@ -0,0 +1,503 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 +"""Hardware UT for ``Worker(level=2, execution_mode="kernel")`` on a2a3. + +The caller side is what kernel mode is: each case initializes ACL, makes the +device current, creates the stream and allocates the device tensors itself, +through ctypes on libascendcl, and lends them to the Worker. The Worker must +never create, reset or destroy any of it, so every case ends by having the +caller synchronize its stream, free its memory, destroy the stream, reset the +device and finalize ACL, and each of those must return 0. + +The executed operator is the AIV vector-add-scalar kernel behind +``kernel_eager_orchestration.cpp`` (``y[i] = x[i] + scalar``), so numeric +checks are exact float32 comparisons after the caller's own synchronize. +No torch is involved. + +Each case runs in a fresh interpreter started with ``python -m`` from the +repository root: kernel_init binds a runtime library into the process and the +ACL device binding is per-thread, so one case's state never reaches the next. + +The ``runtime`` marker is what makes conftest's resource phase dispatch these +cases rather than deselect them, so it is load-bearing rather than descriptive. +The cases that need host_build_graph also load it; every a2a3 build produces it. +""" + +from __future__ import annotations + +import contextlib +import ctypes +import os +import signal +import struct +import subprocess +import sys +from pathlib import Path + +import pytest + +_PROJECT_ROOT = Path(__file__).resolve().parents[4] +_MODULE = "tests.ut.py.test_worker.test_worker_kernel_mode_hw" +_TMR = "tensormap_and_ringbuffer" +_HBG = "host_build_graph" +_CASE_TIMEOUT_S = 300 +_SYNC_TIMEOUT_MS = 60000 +_ACL_MEMCPY_HOST_TO_DEVICE = 1 +_ACL_MEMCPY_DEVICE_TO_HOST = 2 +_COUNT = 128 * 128 +_SENTINEL = -999.0 + +# Runtimes each case loads; a case whose runtime build is missing is skipped before any device work. +_CASE_RUNTIMES = { + "supported_before_init": (_TMR, _HBG), + "eager_numerics": (_TMR,), + "prepare_twice": (_TMR,), + "launch_refusals": (_TMR,), + "second_worker_same_device": (_TMR,), + "hbg_init_refused": (_HBG,), +} + + +# --------------------------------------------------------------------------- +# Caller side: ACL, device, stream and device tensors owned by the test +# --------------------------------------------------------------------------- + + +class _Caller: + """The borrowing side of kernel mode: ACL, one current device, one stream and its device buffers.""" + + def __init__(self, device: int) -> None: + self.device = device + self.acl = self._load_acl() + self._stream = ctypes.c_void_p() + self._allocations: list[ctypes.c_void_p] = [] + self._acl_initialized = False + self._device_bound = False + + @staticmethod + def _load_acl() -> ctypes.CDLL: + lib = None + for name in ("libascendcl.so", "libascendcl.so.1"): + with contextlib.suppress(OSError): + lib = ctypes.CDLL(name) + break + if lib is None: + raise RuntimeError("libascendcl.so is not loadable; source the CANN set_env.sh first") + signatures = { + "aclInit": [ctypes.c_char_p], + "aclFinalize": [], + "aclrtSetDevice": [ctypes.c_int], + "aclrtResetDevice": [ctypes.c_int], + "aclrtCreateStream": [ctypes.POINTER(ctypes.c_void_p)], + "aclrtDestroyStream": [ctypes.c_void_p], + "aclrtSynchronizeStreamWithTimeout": [ctypes.c_void_p, ctypes.c_int32], + "aclrtMalloc": [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t, ctypes.c_int], + "aclrtFree": [ctypes.c_void_p], + "aclrtMemcpy": [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int], + } + for symbol, argtypes in signatures.items(): + function = getattr(lib, symbol) + function.argtypes = argtypes + function.restype = ctypes.c_int + return lib + + def open(self) -> None: + assert self.acl.aclInit(None) == 0 + self._acl_initialized = True + assert self.acl.aclrtSetDevice(self.device) == 0 + self._device_bound = True + assert self.acl.aclrtCreateStream(ctypes.byref(self._stream)) == 0 + assert self._stream.value + + @property + def stream(self) -> int: + return int(self._stream.value or 0) + + def synchronize(self) -> int: + return self.acl.aclrtSynchronizeStreamWithTimeout(self._stream, _SYNC_TIMEOUT_MS) + + def device_buffer(self, values: list[float]) -> int: + """Allocate a float32 device buffer holding ``values``; freed by close().""" + host = (ctypes.c_float * len(values))(*values) + nbytes = ctypes.sizeof(host) + address = ctypes.c_void_p() + assert self.acl.aclrtMalloc(ctypes.byref(address), nbytes, 0) == 0 + self._allocations.append(address) + assert self.acl.aclrtMemcpy(address, nbytes, host, nbytes, _ACL_MEMCPY_HOST_TO_DEVICE) == 0 + return int(address.value or 0) + + def read(self, address: int) -> list[float]: + host = (ctypes.c_float * _COUNT)() + nbytes = ctypes.sizeof(host) + assert self.acl.aclrtMemcpy(host, nbytes, ctypes.c_void_p(address), nbytes, _ACL_MEMCPY_DEVICE_TO_HOST) == 0 + return list(host) + + def close(self) -> dict[str, int]: + """Release everything the caller owns in reverse order; returns the first nonzero status per ACL call.""" + failures: dict[str, int] = {} + + def record(name: str, rc: int) -> None: + if rc != 0: + failures.setdefault(name, rc) + + if self._stream.value: + record("aclrtSynchronizeStreamWithTimeout", self.synchronize()) + for address in reversed(self._allocations): + record("aclrtFree", self.acl.aclrtFree(address)) + self._allocations.clear() + if self._stream.value: + record("aclrtDestroyStream", self.acl.aclrtDestroyStream(self._stream)) + self._stream = ctypes.c_void_p() + if self._device_bound: + record("aclrtResetDevice", self.acl.aclrtResetDevice(self.device)) + self._device_bound = False + if self._acl_initialized: + record("aclFinalize", self.acl.aclFinalize()) + self._acl_initialized = False + return failures + + +@contextlib.contextmanager +def _caller_device(device: int): + """Yield an opened _Caller; on a clean exit every caller-side teardown status must be 0. + + Leaving the block runs the stream synchronize, the frees, aclrtDestroyStream, aclrtResetDevice and + aclFinalize after every Worker in the block has closed, which is the evidence that close() left the + caller's device, stream and memory intact. + """ + caller = _Caller(device) + completed = False + try: + caller.open() + yield caller + completed = True + finally: + failures = caller.close() + if completed: + assert not failures, f"caller-side ACL teardown failed after the Worker closed: {failures}" + elif failures: + print(f"caller-side ACL teardown also failed: {failures}", file=sys.stderr) + + +@contextlib.contextmanager +def _kernel_worker(caller: _Caller, platform: str, runtime: str = _TMR): + """Yield a NEW kernel-mode Worker on the caller's device; it is closed on exit, after a synchronize.""" + from simpler.worker import Worker + + worker = Worker(level=2, execution_mode="kernel", device_id=caller.device, platform=platform, runtime=runtime) + try: + yield worker + finally: + # close() requires the caller's enqueued launches to have drained; a closed Worker's close() is a no-op. + caller.synchronize() + worker.close() + + +# --------------------------------------------------------------------------- +# Callable, config and launch helpers +# --------------------------------------------------------------------------- + + +def _build_eager_callable(platform: str): + """The ChipCallable ``kernel_eager_orchestration(x, y, scalar)`` submitting one AIV add-scalar task.""" + import tempfile + + from simpler.task_interface import ArgDirection, ChipCallable, CoreCallable + + from simpler_setup.elf_parser import extract_text_section + from simpler_setup.kernel_compiler import KernelCompiler + from simpler_setup.pto_isa import ensure_pto_isa_root + + compiler = KernelCompiler(platform) + kernel = _PROJECT_ROOT / "examples" / platform / _TMR / "vector_example/kernels/aiv/kernel_add_scalar.cpp" + orchestration_source = _PROJECT_ROOT / "tests/ut/py/kernel_eager_orchestration.cpp" + with tempfile.TemporaryDirectory(prefix="worker-kernel-eager-") as build_dir: + orchestration = compiler.compile_orchestration(_TMR, str(orchestration_source), build_dir=build_dir) + incore = compiler.compile_incore( + str(kernel), + core_type="aiv", + pto_isa_root=ensure_pto_isa_root(), + extra_include_dirs=compiler.get_orchestration_include_dirs(_TMR), + build_dir=build_dir, + ) + signature = [ArgDirection.IN, ArgDirection.OUT, ArgDirection.SCALAR] + child = CoreCallable.build(signature=signature, binary=extract_text_section(incore)) + return ChipCallable.build( + signature=signature, + func_name="kernel_eager_orchestration", + binary=orchestration, + children=[(0, child)], + ) + + +def _kernel_config(): + from simpler.task_interface import CallConfig + + config = CallConfig() + config.runtime_env.ring_task_window = 64 + config.runtime_env.ring_heap = 1 << 20 + config.runtime_env.ring_dep_pool = 1024 + return config + + +def _input_values(seed: int) -> list[float]: + # Every value and every value + scalar used below is exactly representable in float32. + return [float(i % 127 + seed * 257) for i in range(_COUNT)] + + +def _launch_args(source: int, destination: int, scalar: float | None): + """ChipStorageTaskArgs ``(x, y, scalar)``; ``scalar=None`` omits the scalar the signature requires.""" + from simpler.task_interface import ChipStorageTaskArgs, ChipTensor, DataType + + args = ChipStorageTaskArgs() + args.add_tensor(ChipTensor.make(source, (_COUNT,), DataType.FLOAT32, child_memory=True)) + args.add_tensor(ChipTensor.make(destination, (_COUNT,), DataType.FLOAT32, child_memory=True)) + if scalar is not None: + args.add_scalar(int.from_bytes(struct.pack(" None: + import simpler.worker as worker_module + + assert worker._lifecycle.name == "CLOSED" + assert worker._chip_worker is None + assert worker._kernel_callables == {} + assert not pin_finalizer.alive + assert worker_module._PINNED_KERNEL_CHIP_WORKERS == [] + + +def _init_kernel_worker(worker): + """init(config=...) and return the armed GC-pin finalizer.""" + worker.init(config=_kernel_config()) + assert worker._lifecycle.name == "READY" + pin_finalizer = worker._kernel_pin_finalizer + assert pin_finalizer is not None and pin_finalizer.alive + return pin_finalizer + + +# --------------------------------------------------------------------------- +# Cases (subprocess bodies) +# --------------------------------------------------------------------------- + + +def _case_supported_before_init(platform: str, device: int) -> None: + """kernel_mode_supported answers from the runtime build before init, with no device made current.""" + from simpler.worker import Worker + + for runtime, expected in ((_TMR, True), (_HBG, False)): + kernel = Worker(level=2, execution_mode="kernel", device_id=device, platform=platform, runtime=runtime) + program = Worker(level=2, device_id=device, platform=platform, runtime=runtime) + for worker in (kernel, program): + assert worker.kernel_mode_supported is expected, (runtime, worker._execution_mode) + assert worker.kernel_mode_supported is expected, (runtime, worker._execution_mode) + assert worker._lifecycle.name == "NEW" + assert worker._chip_worker is None + + +def _case_eager_numerics(platform: str, device: int) -> None: + """Prepare once, launch twice with different scalars, close, and leave the caller's device usable.""" + chip = _build_eager_callable(platform) + with _caller_device(device) as caller, _kernel_worker(caller, platform) as worker: + pin_finalizer = _init_kernel_worker(worker) + assert worker.kernel_mode_supported is True + callable_id = worker.kernel_prepare_callable(chip) + with pytest.raises(RuntimeError, match="execution_mode"): + worker.device_memory_info() + committed = worker.committed_device_memory() + assert committed > 0 + + first = _launch_and_check(caller, worker, callable_id, scalar=1.25, seed=0) + source, destination, _ = _launch_and_check(caller, worker, callable_id, scalar=-3.5, seed=1) + # A launch commits no device memory, and the second launch wrote only its own output. + assert worker.committed_device_memory() == committed + assert caller.read(first[1]) == first[2] + + assert caller.synchronize() == 0 + worker.close() + _assert_closed_cleanly(worker, pin_finalizer) + with pytest.raises(RuntimeError, match="READY"): + worker.kernel_launch(callable_id, _launch_args(source, destination, 1.0), caller_stream=caller.stream) + with pytest.raises(RuntimeError, match="READY"): + worker.kernel_prepare_callable(chip) + assert caller.synchronize() == 0 + + +def _case_prepare_twice(platform: str, device: int) -> None: + """The same callable prepared twice takes two distinct ids, and both launch correctly.""" + chip = _build_eager_callable(platform) + with _caller_device(device) as caller, _kernel_worker(caller, platform) as worker: + pin_finalizer = _init_kernel_worker(worker) + first_id = worker.kernel_prepare_callable(chip) + second_id = worker.kernel_prepare_callable(chip) + assert first_id >= 0 and second_id >= 0 + assert first_id != second_id + assert set(worker._kernel_callables) == {first_id, second_id} + + _launch_and_check(caller, worker, second_id, scalar=0.75, seed=2) + _launch_and_check(caller, worker, first_id, scalar=2.5, seed=3) + + assert caller.synchronize() == 0 + worker.close() + _assert_closed_cleanly(worker, pin_finalizer) + + +def _case_launch_refusals(platform: str, device: int) -> None: + """Refused launches enqueue nothing and leave the context launchable.""" + chip = _build_eager_callable(platform) + with _caller_device(device) as caller, _kernel_worker(caller, platform) as worker: + pin_finalizer = _init_kernel_worker(worker) + callable_id = worker.kernel_prepare_callable(chip) + values = _input_values(seed=4) + source = caller.device_buffer(values) + destination = caller.device_buffer([_SENTINEL] * _COUNT) + args = _launch_args(source, destination, scalar=1.25) + + with pytest.raises(ValueError, match="kernel_prepare_callable"): + worker.kernel_launch(callable_id + 1, args, caller_stream=caller.stream) + with pytest.raises(ValueError, match="caller_stream"): + worker.kernel_launch(callable_id, args, caller_stream=0) + # Two tensors and no scalar disagree with the (IN, OUT, SCALAR) signature; the native launch refuses + # them while encoding, before anything is enqueued on the caller's stream. + with pytest.raises(RuntimeError, match="simpler_kernel_mode_launch failed"): + worker.kernel_launch(callable_id, _launch_args(source, destination, None), caller_stream=caller.stream) + assert caller.synchronize() == 0 + assert caller.read(destination) == [_SENTINEL] * _COUNT + + worker.kernel_launch(callable_id, args, caller_stream=caller.stream) + assert caller.synchronize() == 0 + assert caller.read(destination) == [value + 1.25 for value in values] + + worker.close() + _assert_closed_cleanly(worker, pin_finalizer) + + +def _case_second_worker_same_device(platform: str, device: int) -> None: + """A second kernel Worker on a claimed device fails init without disturbing the owner; a successor inits.""" + chip = _build_eager_callable(platform) + with _caller_device(device) as caller: + with _kernel_worker(caller, platform) as owner: + owner_pin = _init_kernel_worker(owner) + owner_id = owner.kernel_prepare_callable(chip) + with _kernel_worker(caller, platform) as rival: + with pytest.raises(RuntimeError): + rival.init(config=_kernel_config()) + assert rival._lifecycle.name == "FAILED" + assert rival._kernel_pin_finalizer is None + rival.close() + assert rival._lifecycle.name == "CLOSED" + assert rival._chip_worker is None + + _launch_and_check(caller, owner, owner_id, scalar=1.25, seed=5) + assert caller.synchronize() == 0 + owner.close() + _assert_closed_cleanly(owner, owner_pin) + + with _kernel_worker(caller, platform) as successor: + successor_pin = _init_kernel_worker(successor) + successor_id = successor.kernel_prepare_callable(chip) + _launch_and_check(caller, successor, successor_id, scalar=-3.5, seed=6) + assert caller.synchronize() == 0 + successor.close() + _assert_closed_cleanly(successor, successor_pin) + + +def _case_hbg_init_refused(platform: str, device: int) -> None: + """host_build_graph has no kernel mode: init fails to FAILED, close() is clean, the caller's stream survives.""" + with _caller_device(device) as caller, _kernel_worker(caller, platform, runtime=_HBG) as worker: + with pytest.raises(RuntimeError, match="kernel mode"): + worker.init(config=_kernel_config()) + assert worker._lifecycle.name == "FAILED" + assert worker._kernel_pin_finalizer is None + assert caller.synchronize() == 0 + worker.close() + assert worker._lifecycle.name == "CLOSED" + assert worker._chip_worker is None + + +_CASES = { + "supported_before_init": _case_supported_before_init, + "eager_numerics": _case_eager_numerics, + "prepare_twice": _case_prepare_twice, + "launch_refusals": _case_launch_refusals, + "second_worker_same_device": _case_second_worker_same_device, + "hbg_init_refused": _case_hbg_init_refused, +} + + +# --------------------------------------------------------------------------- +# pytest side +# --------------------------------------------------------------------------- + + +def _require_prebuilt(platform: str, runtimes: tuple[str, ...]) -> None: + from simpler_setup.runtime_builder import RuntimeBuilder + + for runtime in runtimes: + try: + RuntimeBuilder(platform=platform).get_binaries(runtime) + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +def _run_case_in_subprocess(case: str, platform: str, device: int) -> None: + env = dict(os.environ, PYTHONFAULTHANDLER="1") + proc = subprocess.Popen( + [sys.executable, "-m", _MODULE, case, platform, str(device)], + cwd=str(_PROJECT_ROOT), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + try: + output, _ = proc.communicate(timeout=_CASE_TIMEOUT_S) + except subprocess.TimeoutExpired: + # The case leads its own session, so this also kills the compiler processes it started. + with contextlib.suppress(ProcessLookupError): + os.killpg(proc.pid, signal.SIGKILL) + output, _ = proc.communicate() + pytest.fail(f"case {case} did not exit within {_CASE_TIMEOUT_S}s:\n{output}") + assert proc.returncode == 0, f"case {case} exited with {proc.returncode}:\n{output}" + + +@pytest.mark.requires_hardware +@pytest.mark.platforms(["a2a3"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime("tensormap_and_ringbuffer") +@pytest.mark.parametrize("case", list(_CASE_RUNTIMES)) +def test_worker_kernel_mode_on_caller_device(case, st_platform, st_device_ids): + """Drive one Worker kernel-mode case against a device, stream and memory the case itself owns.""" + assert st_device_ids, "device_count(1) must yield at least one device id" + _require_prebuilt(st_platform, _CASE_RUNTIMES[case]) + _run_case_in_subprocess(case, st_platform, int(st_device_ids[0])) + + +if __name__ == "__main__": + _CASES[sys.argv[1]](sys.argv[2], int(sys.argv[3]))