From 800ddbb18d71519ea39fe20141793258dfaf5b0a Mon Sep 17 00:00:00 2001 From: superxf <1208713646@qq.com> Date: Thu, 3 Sep 2026 17:13:11 +0800 Subject: [PATCH] support register_chip_control_extension --- python/simpler/__init__.py | 9 ++++- python/simpler/worker.py | 82 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/python/simpler/__init__.py b/python/simpler/__init__.py index 995d77f267..3bb7ff58eb 100644 --- a/python/simpler/__init__.py +++ b/python/simpler/__init__.py @@ -37,6 +37,7 @@ __all__ = [ "DEFAULT_THRESHOLD", "Worker", + "register_chip_control_extension", "NUL", "TIMING", "get_current_config", @@ -46,7 +47,13 @@ ] # name -> (module, attribute). Resolved by __getattr__ on first access. -_LAZY_ATTRS = {"Worker": (f"{__name__}.worker", "Worker")} +_LAZY_ATTRS = { + "Worker": (f"{__name__}.worker", "Worker"), + "register_chip_control_extension": ( + f"{__name__}.worker", + "register_chip_control_extension", + ), +} _LAZY_SUBMODULES = ("comm_endpoints", "task_interface") diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 76949d2f0e..a200ce2f46 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -630,6 +630,24 @@ def _require_copy_span(extent: int, offset: int, nbytes: int, *, side: str, api: _LOCAL_GLOBAL_CONTROL_HEADER = struct.Struct(" None: + """Register a trusted Python control handler inherited by chip children.""" + if not isinstance(name, str) or not name or len(name.encode("utf-8")) > 255: + raise ValueError("chip control extension name must contain 1 to 255 UTF-8 bytes") + if not callable(handler): + raise TypeError("chip control extension handler must be callable") + existing = _chip_control_extensions.get(name) + if existing is not None and existing is not handler: + raise ValueError(f"chip control extension {name!r} is already registered") + _chip_control_extensions[name] = handler + # Layout of the CTRL_COMM_INIT request shm. _COMM_INIT_HEADER = struct.Struct(" str: return raw[: nul if nul >= 0 else _CTRL_SHM_NAME_BYTES].decode("utf-8", "replace") +def _read_ctrl_staged_payload(buf: memoryview) -> bytes: + payload_size = int(struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0]) + shm_name = _read_ctrl_staged_shm_name(buf) + if payload_size <= 0 or not shm_name: + raise ValueError("chip control extension payload must not be empty") + shm = SharedMemory(name=shm_name) + shm_buf = cast(memoryview, shm.buf) + try: + if payload_size > shm.size: + raise ValueError(f"chip control extension payload size {payload_size} exceeds shm size {shm.size}") + return bytes(shm_buf[:payload_size]) + finally: + shm_buf.release() + shm.close() + + +def _handle_chip_control_extension(cw: ChipWorker, buf: memoryview, device_id: int) -> None: + envelope = _read_ctrl_staged_payload(buf) + if len(envelope) < _CHIP_EXTENSION_HEADER.size: + raise ValueError("chip control extension envelope is truncated") + (name_size,) = _CHIP_EXTENSION_HEADER.unpack_from(envelope) + name_end = _CHIP_EXTENSION_HEADER.size + name_size + if name_size == 0 or name_end > len(envelope): + raise ValueError("chip control extension name is invalid") + name = envelope[_CHIP_EXTENSION_HEADER.size : name_end].decode("utf-8") + handler = _chip_control_extensions.get(name) + if handler is None: + raise KeyError(f"chip control extension {name!r} is not registered") + error = handler(cw, envelope[name_end:], device_id) + if error: + raise RuntimeError(str(error)) + + def _allocate_local_slot(registry: dict[int, Any]) -> int: for i in range(MAX_REGISTERED_CALLABLE_IDS): if i not in registry: @@ -3006,6 +3057,8 @@ def handle_control( # noqa: PLR0912, PLR0915 -- one branch per control sub-comm elif sub_cmd == _CTRL_DEVICE_MEMORY_INFO: info = cw.device_memory_info() _DEVICE_MEMORY_INFO.pack_into(buf, _CTRL_OFF_RESULT, info.free_bytes, info.total_bytes) + elif sub_cmd == _CTRL_CHIP_EXTENSION: + _handle_chip_control_extension(cw, buf, device_id) elif sub_cmd == _CTRL_IMPORT_RELEASE: import_registry.unregister(_unpack_identity_wire(_read_control_digest(buf))) elif sub_cmd == CTRL_GLOBAL_DOMAIN_PREPARE: @@ -10584,6 +10637,35 @@ def _copy_extent( _require_copy_span(host_nbytes, host_offset, nbytes, side=host_side, api=api) return device_offset, host_offset, nbytes + def run_chip_control_extension( + self, + name: str, + payload: bytes, + *, + timeout_s: float | None = None, + ) -> None: + """Broadcast one named control payload to every local chip child.""" + with self._operation_lease("run_chip_control_extension"): + if self.level < 3 or self._worker is None: + raise TypeError("chip control extensions require an initialized level >= 3 Worker") + if name not in _chip_control_extensions: + raise KeyError(f"chip control extension {name!r} is not registered in the parent") + name_bytes = name.encode("utf-8") + envelope = _CHIP_EXTENSION_HEADER.pack(len(name_bytes)) + name_bytes + bytes(payload) + with self._device_control_admission("run_chip_control_extension"): + results = self._worker.broadcast_control_all( + WorkerType.NEXT_LEVEL, + _CTRL_CHIP_EXTENSION, + envelope, + None, + timeout_s=timeout_s, + ) + errors = self._control_errors(list(results)) + if errors: + raise RuntimeError( + f"chip control extension {name!r} failed on {len(errors)} child workers; first error: {errors[0]}" + ) + @staticmethod def _require_device_end(handle: Buffer, *, api: str) -> None: """Reject a host handle on the device end by what it is, before the identity is looked up.