From c748e48f8521c329fde756fc0d6d023a20de1970 Mon Sep 17 00:00:00 2001 From: puddingfjz <2811443837@qq.com> Date: Thu, 16 Jul 2026 23:47:13 -0700 Subject: [PATCH] Fix: expose host buffers to L3 subworkers - Broadcast host-buffer map and unmap controls to subworkers - Rewrite parent host addresses before SubWorker task decoding - Add focused regression coverage for post-fork visibility --- docs/comm-domain.md | 19 ++- python/simpler/worker.py | 144 ++++++++++-------- .../test_host_buffer_registration.py | 138 ++++++++++++++++- 3 files changed, 222 insertions(+), 79 deletions(-) diff --git a/docs/comm-domain.md b/docs/comm-domain.md index 4b2517e373..e1f43b9048 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -150,20 +150,19 @@ with orch.allocate_domain(...) as handle: ## 6. Host tensor visibility for `worker.run` -A host tensor passed to `worker.run(...)` / `orch.submit_next_level(...)` is -ultimately dereferenced from the forked chip child, not the parent, so its -memory must be reachable there. Zero-copy requires born-shared memory (a virtual -address is not portable across the fork, so the child can only reach the same -physical pages via a MAP_SHARED backing established at allocation time), which is -why the buffer is worker-allocated rather than a user tensor. Two sources are -legal: +A host tensor passed to `worker.run(...)` / `orch.submit_next_level(...)` / +`orch.submit_sub(...)` is ultimately dereferenced from a forked local L3 child, +not the parent, so its memory must be backed by pages mapped into that child. +Fork-inherited MAP_SHARED mappings retain their virtual address, while post-fork +worker-allocated buffers may map at a different address and have their pointers +rewritten before decoding. Two sources are legal: | Source | How | Why it works | | ------ | --- | ------------ | -| **fork-inherited** | `tensor.share_memory_()` **before the chip children are forked** (i.e. before the first `Worker.run()`) | the child inherits the MAP_SHARED page at fork | -| **worker-allocated post-fork** | `worker.create_host_buffer(nbytes)` after the chips exist | born-shared memory attached into every child, **zero-copy** | +| **fork-inherited** | `tensor.share_memory_()` **before the local L3 children are forked** (i.e. before the first `Worker.run()`) | the child inherits the MAP_SHARED page at fork | +| **worker-allocated post-fork** | `worker.create_host_buffer(nbytes)` after the children exist | born-shared memory attached into every local child, **zero-copy** | -The chip children are forked lazily on the **first** `run()`. A host tensor +The local L3 children are forked lazily on the **first** `run()`. A host tensor created after that — the natural dynamic-shape serving pattern — is invisible to the children unless it lives in a `create_host_buffer` buffer: diff --git a/python/simpler/worker.py b/python/simpler/worker.py index f580ea9d54..d77c1e4d6a 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -232,7 +232,7 @@ def my_l4_orch(orch, args, config): _CTRL_PY_UNREGISTER = 11 _CTRL_PY_IMPORT_REGISTER = 12 # Host-buffer registration. MAP_HOST maps a named host-buffer shm -# into every chip child *post-fork* and keeps it mapped so later runs can copy +# into every local L3 child *post-fork* and keeps it mapped so later runs can copy # through it; UNMAP_HOST drops one. The child also records the parent VA range # the shm stands in for, so the per-task blob's host pointers (raw parent VAs) # can be rewritten to the child's own mapping before the runtime dereferences @@ -387,7 +387,7 @@ class _RemoteSession: class _HostBufEntry: """Parent-side record for a born-shared post-fork host buffer. - The worker owns ``shm`` — a named buffer the chip children attach and + The worker owns ``shm`` — a named buffer the local L3 children attach and read/write through. The user builds a tensor over it (via the buffer protocol on :class:`HostBuffer`), so the buffer *is* the shm: ``data_ptr == shm_base`` and no per-run copy is needed (the child reads and writes the same @@ -407,7 +407,7 @@ class HostBuffer: """Handle for a worker-allocated, born-shared host buffer (zero-copy). Returned by ``Worker.create_host_buffer``. ``buffer`` is a ``memoryview`` - over shared memory already attached into every chip child; wrap it with + over shared memory already attached into every local L3 child; wrap it with ``torch.frombuffer`` / ``np.frombuffer`` to get a real tensor whose writes land directly in the child-visible pages — no per-run copy. ``token`` / ``data_ptr`` / ``nbytes`` identify the mapping; pass this handle back to @@ -934,46 +934,62 @@ def _sub_worker_loop( rethrows it as ``std::runtime_error``. """ state_addr = _buffer_field_addr(buf, _OFF_STATE) - while True: - state = _mailbox_load_i32(state_addr) - if state == _TASK_READY: - digest = _read_task_digest(buf) - cid = identity_table.get(digest) - fn = registry.get(int(cid)) if cid is not None else None - code = 0 - msg = "" - if fn is None: - code = 1 - msg = f"sub_worker: callable hash {_format_digest(digest)} not registered" - else: + host_buf_table: dict[int, tuple[SharedMemory, int, int, int]] = {} + host_buf_ranges: list[tuple[int, int, int]] = [] + try: + while True: + state = _mailbox_load_i32(state_addr) + if state == _TASK_READY: + digest = _read_task_digest(buf) + cid = identity_table.get(digest) + fn = registry.get(int(cid)) if cid is not None else None + code = 0 + msg = "" + if fn is None: + code = 1 + msg = f"sub_worker: callable hash {_format_digest(digest)} not registered" + else: + try: + if host_buf_ranges: + _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) + args = _read_args_from_mailbox(buf) + fn(args) + except Exception as e: # noqa: BLE001 + code = 1 + msg = _format_exc("sub_worker", e) + _write_error(buf, code, msg) + _mailbox_store_i32(state_addr, _TASK_DONE) + elif state == _CONTROL_REQUEST: + sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] + code = 0 + msg = "" try: - args = _read_args_from_mailbox(buf) - fn(args) + if sub_cmd == _CTRL_MAP_HOST: + _handle_ctrl_map_host(buf, host_buf_table, host_buf_ranges) + elif sub_cmd == _CTRL_UNMAP_HOST: + _handle_ctrl_unmap_host(buf, host_buf_table, host_buf_ranges) + else: + _handle_py_callable_control( + buf, + registry, + identity_table, + identity_refs, + int(sub_cmd), + context="sub_worker", + ) except Exception as e: # noqa: BLE001 code = 1 - msg = _format_exc("sub_worker", e) - _write_error(buf, code, msg) - _mailbox_store_i32(state_addr, _TASK_DONE) - elif state == _CONTROL_REQUEST: - sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] - code = 0 - msg = "" + msg = _format_exc("sub_worker control", e) + _write_error(buf, code, msg) + _mailbox_store_i32(state_addr, _CONTROL_DONE) + elif state == _SHUTDOWN: + break + finally: + for host_shm, _lo, _hi, _base in host_buf_table.values(): try: - _handle_py_callable_control( - buf, - registry, - identity_table, - identity_refs, - int(sub_cmd), - context="sub_worker", - ) - except Exception as e: # noqa: BLE001 - code = 1 - msg = _format_exc("sub_worker control", e) - _write_error(buf, code, msg) - _mailbox_store_i32(state_addr, _CONTROL_DONE) - elif state == _SHUTDOWN: - break + host_shm.close() + except Exception: # noqa: BLE001 + pass def _read_shm_name(buf, offset: int) -> str: @@ -1796,7 +1812,7 @@ def __init__( # Post-fork zero-copy host buffers (``create_host_buffer``). Keyed by the # born-shared shm's mapped base (== the buffer's data_ptr); each entry maps - # a named shm into every chip child so memory created after the children + # a named shm into every local L3 child so memory created after the children # were forked is still reachable by a later run — with no per-run copy. self._host_buf_registry: dict[int, _HostBufEntry] = {} # Immutable read snapshot for the lock-free per-submit lookup @@ -4204,10 +4220,10 @@ def copy_from(self, dst: int, src: int, size: int, worker_id: int = 0) -> None: # ------------------------------------------------------------------ def create_host_buffer(self, nbytes: int) -> HostBuffer: - """Allocate a born-shared host buffer, attached into every chip child, + """Allocate a born-shared host buffer, attached into every local L3 child, that a later ``run()`` reads/writes with **no per-run copy**. - L3 chip children are forked lazily on the first ``run()``; memory created + Local L3 children are forked lazily on the first ``run()``; memory created afterwards is not in their address space. This hands you memory that is *born* in a shm already attached into every child, so there is nothing to copy: the child reads and writes the same physical pages the parent sees. @@ -4223,7 +4239,7 @@ def create_host_buffer(self, nbytes: int) -> HostBuffer: worker.free_host_buffer(buf) # drop the tensor first simpler stays framework-free: torch/numpy appear only on the user's side - (``frombuffer``). Blocks until every chip child has attached the buffer; + (``frombuffer``). Blocks until every local L3 child has attached the buffer; not thread-safe against a concurrent ``run`` / ``create`` / ``free`` on the same Worker — drive them from one thread, as the L3 worker is otherwise. @@ -4270,17 +4286,10 @@ def create_host_buffer(self, nbytes: int) -> HostBuffer: self._rebuild_host_buf_snapshot() payload = _HOST_BUF_MAP_HEADER.pack(token, data_ptr, nbytes) + shm.name.encode("utf-8") - results = self._worker.broadcast_control_all( - WorkerType.NEXT_LEVEL, - int(_CTRL_MAP_HOST), - payload, - None, - timeout_s=self._py_control_timeout_s, - ) - errors = self._control_errors(list(results)) + errors = self._broadcast_host_control(_CTRL_MAP_HOST, payload) if errors: raise RuntimeError( - f"create_host_buffer: MAP_HOST failed on {len(errors)} chip children; first error: {errors[0]}" + f"create_host_buffer: MAP_HOST failed on {len(errors)} local L3 children; first error: {errors[0]}" ) except BaseException: # Roll back on any failure — a staging error before the map, a partial @@ -4315,7 +4324,7 @@ def create_host_buffer(self, nbytes: int) -> HostBuffer: def free_host_buffer(self, handle: HostBuffer) -> None: """Release a born-shared buffer created by ``create_host_buffer``. - Unmaps it from every chip child and frees the parent shm. Drop every + Unmaps it from every local L3 child and frees the parent shm. Drop every tensor / ``memoryview`` you built over ``handle.buffer`` *first*: a live view keeps the shm's pages exported, so ``close()`` cannot release them and the buffer only warns (and is reclaimed once the last view is gone). @@ -4344,7 +4353,7 @@ def free_host_buffer(self, handle: HostBuffer) -> None: if errors: sys.stderr.write( f"[worker pid={os.getpid()}] WARN: free_host_buffer token={entry.token} " - f"failed on {len(errors)} chip children; first error: {errors[0]}\n" + f"failed on {len(errors)} local L3 children; first error: {errors[0]}\n" ) sys.stderr.flush() @@ -4396,17 +4405,24 @@ def _release_all_host_buffers(self) -> None: self._close_host_shm(entry) def _broadcast_host_unmap(self, token: int) -> list[str]: - """Broadcast _CTRL_UNMAP_HOST for ``token`` to every chip child.""" + """Broadcast _CTRL_UNMAP_HOST for ``token`` to every local L3 child.""" + return self._broadcast_host_control(_CTRL_UNMAP_HOST, _HOST_BUF_UNMAP.pack(token)) + + def _broadcast_host_control(self, sub_cmd: int, payload: bytes) -> list[str]: if self._worker is None: return [] - results = self._worker.broadcast_control_all( - WorkerType.NEXT_LEVEL, - int(_CTRL_UNMAP_HOST), - _HOST_BUF_UNMAP.pack(token), - None, - timeout_s=self._py_control_timeout_s, - ) - return self._control_errors(list(results)) + results = [] + for worker_type in (WorkerType.NEXT_LEVEL, WorkerType.SUB): + results.extend( + self._worker.broadcast_control_all( + worker_type, + int(sub_cmd), + payload, + None, + timeout_s=self._py_control_timeout_s, + ) + ) + return self._control_errors(results) def _stage_host_buffers_for_chip_submit(self, args: Any) -> None: """Validate the host tensors of one chip submit before dispatch. @@ -4610,7 +4626,7 @@ def close(self) -> None: # noqa: PLR0912 -- parallel teardown for _worker + sub sys.stderr.flush() # Release any host buffers the user never unregistered. Must run while - # the chip mailboxes are still usable (before _worker.close()). + # the local L3 child mailboxes are still usable (before _worker.close()). self._release_all_host_buffers() if self.level == 2: diff --git a/tests/ut/py/test_worker/test_host_buffer_registration.py b/tests/ut/py/test_worker/test_host_buffer_registration.py index 283d2820a7..9dba23bfb8 100644 --- a/tests/ut/py/test_worker/test_host_buffer_registration.py +++ b/tests/ut/py/test_worker/test_host_buffer_registration.py @@ -6,24 +6,28 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""Host-side staging for post-fork zero-copy host buffers. +"""Host-side handling for post-fork zero-copy host buffers. ``submit_next_level`` calls ``_stage_host_buffers_for_chip_submit`` before any dispatch. A ``create_host_buffer`` buffer is born-shared — its bytes already live in the child-visible shm and the child writes results back into the same pages — so staging copies in neither direction; it only validates that each -in-range view fits inside its buffer. These are pure host-side unit tests: they -inject the staging state directly and never fork, compile a kernel, or touch a -device. The end-to-end round-trip lives in the a2a3sim scene test. +in-range view fits inside its buffer. The SubWorker tests also exercise +MAP/UNMAP control handling and parent-to-child address rewriting directly +against the child loop. These tests never compile a kernel or touch a device; +the end-to-end chip round-trip lives in the a2a3sim scene test. """ from __future__ import annotations import ctypes +import struct +from multiprocessing.shared_memory import SharedMemory import pytest +import simpler.worker as worker_mod import torch -from simpler.task_interface import TaskArgs, TensorArgType +from simpler.task_interface import MAILBOX_SIZE, TaskArgs, TensorArgType, WorkerType from simpler.worker import Worker, _HostBufEntry from simpler_setup import make_tensor_arg @@ -107,3 +111,127 @@ def test_subview_inside_buffer_is_accepted(self): w = _staging_worker(_entry_for(parent, shm_buf)) w._stage_host_buffers_for_chip_submit(_arg(parent[64:128], TensorArgType.INPUT)) assert list(shm_buf[:4]) == [9.0, 9.0, 9.0, 9.0] + + +def test_host_buffer_control_reaches_l3_sub_workers(): + calls = [] + + class FakeWorker: + def broadcast_control_all(self, worker_type, sub_cmd, payload, digest, *, timeout_s): + calls.append((worker_type, sub_cmd, bytes(payload), digest, timeout_s)) + return [] + + worker = Worker(level=3) + worker._initialized = True + worker._hierarchical_started = True + worker._chip_shms = [None] + worker._worker = FakeWorker() + worker._start_hierarchical = lambda: None + + handle = worker.create_host_buffer(8) + try: + assert [(call[0], call[1]) for call in calls] == [ + (WorkerType.NEXT_LEVEL, worker_mod._CTRL_MAP_HOST), + (WorkerType.SUB, worker_mod._CTRL_MAP_HOST), + ] + finally: + handle.buffer.release() + worker.free_host_buffer(handle) + + assert [(call[0], call[1]) for call in calls[2:]] == [ + (WorkerType.NEXT_LEVEL, worker_mod._CTRL_UNMAP_HOST), + (WorkerType.SUB, worker_mod._CTRL_UNMAP_HOST), + ] + + +def test_l3_sub_worker_maps_rewrites_and_unmaps_host_buffer(monkeypatch): + mailbox = SharedMemory(create=True, size=MAILBOX_SIZE) + host = SharedMemory(create=True, size=8) + staged = None + mailbox_buf = mailbox.buf + host_buf = host.buf + assert mailbox_buf is not None + assert host_buf is not None + + try: + struct.pack_into("