diff --git a/docs/comm-domain.md b/docs/comm-domain.md index e0b3f8bfc4..6ebc7c92a8 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -48,6 +48,71 @@ allocation: if `sum(b.nbytes) > window_size`, `allocate_domain` raises Kernels read peer windows through `device_ctx` (which holds every rank's window base, local + imported peer); `buffer_ptrs[name]` is the local slice. +### Global CommDomain across local and remote L3 nodes + +An L4 worker can build the same `CommContext` shape across any combination of +forked local L3 workers (`add_worker`) and TCP-connected L3 workers +(`add_remote_worker`) without `mpirun`: + +```python +with orch.allocate_global_domain( + name="tp", + members=[(node0_worker_id, 0), (node1_worker_id, 0)], + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], +) as domain: + ... +``` + +Each member is `(l3_worker_id, local_l2_worker_id)`. The order defines dense +domain ranks. A remote node reads `comm_profile` and `global_device_ranks` +from `RemoteWorkerSpec`; a local L3 reads the same fields from its `Worker` +configuration. All participating nodes must use the same profile. + +An MPI-launched group registered with `add_mpirun_worker_group` uses the same +member contract. Rank 0 writes the group manifest before launch, every MPI rank +must publish READY before the L4 parent exposes the returned worker ids, and +the `MpiL3GroupSpec.hosts` order defines node ranks. Global CommDomain members +must include the complete returned group; their order still defines dense +domain ranks. + +Global CommDomain capability follows the backend that the node actually +loads: a platform ending in `sim` supports the `sim` profile, and a real +`a2a3` platform supports `a3-fabric-v1`. Real A5 and any other +platform/profile combination currently reject allocation before `PREPARE`. +Each local or remote L3 repeats the same check during `COMM_INIT`, so an +unsupported backend never advertises a usable descriptor capability. + +The control flow is: + +1. L4 sends `COMM_INIT` with cluster, node, global-device, and domain-rank + identities. +2. Each L3 asks its participating L2 children to create a local window and + export a transport descriptor. +3. L4 validates and assembles one complete rank-ordered descriptor table. +4. L4 returns that table to every L3, which forwards it to each L2 for import. +5. L4 commits only after all imports succeed. Any earlier failure sends + `ABORT` and releases every prepared local window. + +The descriptor reports the backend's actual mapped size. A3 Fabric may align +the requested size to its VMM granularity; bounds checks against the mapped +window use the returned mapping size, while named-buffer offsets and limits use +each buffer's `nbytes`. The L4 `GlobalCommDomainHandle` exposes only topology +and buffer metadata. L3-local `ChipDomainContext` objects retain device context +and pointers for kernel submission. Remote orchestration code calls +`orch.get_global_domain(domain_id)` to obtain only its committed L3-local +contexts. + +`copy_to_global_domain` and `copy_from_global_domain` provide bounded +control-plane staging and smoke checks. Normal communication still runs in +L2 kernels through the imported `CommContext`. + +By default a live Global CommDomain is swept after the current `Worker.run` +drains. Set `retain_after_run=True` when a communication kernel writes results +into the window and a second L4 run must inspect them. The later run should +call `domain.release()` after copying the results; `Worker.close()` is the +final safety net. + --- ## 2. Lifetime model diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index 607a6f9bc4..1324911b62 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -67,11 +67,18 @@ Implemented: imported-handle scheduling eligibility, and deferred owner free. - Registry-scope-aware remote callable manifest/control install for dispatcher `PYTHON_IMPORT`, inner `PYTHON_IMPORT`, and inner inline `CHIP_CALLABLE`. + Pre-init `ChipCallable` registrations on an L4 worker are serialized into + each remote session manifest and installed on that L3's L2 children. +- A no-`mpirun` A3 TCP smoke that keeps a Global CommDomain across two L4 + runs, executes peer `TLOAD` from each remote L2, and verifies the reduced + values before release. +- Two-server hardware validation covers L4-brokered peer `TLOAD`, one L2 + compute followed by cross-machine communication, and two-NPU-per-node + remote L3 group compute. Still pending: - A2 RoCE, A3 HCCS, and A5 UB HCOMM profiles. -- Remote `CommDomain` allocation/import and hardware-gated validation. - Negotiated `PYTHON_SERIALIZED` remote callable payloads and staged `CHIP_CALLABLE` blob adapters. @@ -449,10 +456,9 @@ Session execution rules: the current one-`WorkerThread`-per-child local scheduling model and keeps ordering, buffer lifetime, and callable visibility simple. - State-changing CONTROL frames such as register, unregister, buffer free, - copy, export/import, and import release serialize with TASK execution on the - ordered command lane. They are not applied concurrently with a running TASK - on the same endpoint. Future Remote CommDomain controls follow the same - ordering rule when they enter scope. + copy, export/import, import release, and Global CommDomain transactions + serialize with TASK execution on the ordered command lane. They are not + applied concurrently with a running TASK on the same endpoint. - Bulk data movement may use a separate data plane, but the state change that makes staged bytes, callable payloads, or imported handles visible is ordered by the command lane. diff --git a/docs/remote-l3-worker-design/implementation-record.md b/docs/remote-l3-worker-design/implementation-record.md index 9958df3816..9d1a554535 100644 --- a/docs/remote-l3-worker-design/implementation-record.md +++ b/docs/remote-l3-worker-design/implementation-record.md @@ -16,7 +16,7 @@ It is updated as each documented feature is completed and verified. | 5 | Versioned remote frame codec | In progress | TASK/COMPLETION/CONTROL_REPLY/HELLO/CONTROL/HEALTH exist; core fuzz/bounds coverage is present, with more exhaustive corpus testing still possible. | | 6 | Remote callable registry | In progress | Dispatcher `PYTHON_IMPORT`, inner manifest/control `PYTHON_IMPORT`, and inner manifest/control inline `CHIP_CALLABLE` are implemented; serialized payloads and staged chip blobs remain negotiated extensions. | | 7 | Fork-safe simulation session runner | In progress | Daemon/session bootstrap and HELLO READY barrier are implemented for sim transport. | -| 8 | Remote control-plane parity | In progress | Registry, alloc/free/copy, export/import/release-import controls are implemented for sim; Remote CommDomain controls are reserved/unsupported. | +| 8 | Remote control-plane parity | In progress | Registry, remote buffers, and Global CommDomain prepare/import/commit/release/copy controls are implemented. | | 9 | Remote buffer registry | In progress | Sim owner/imported buffers, TASK materialization, public memory API, opaque handles, slot/import-ref capture, and deferred free/release-import are implemented. | | 10 | A2 RoCE HCOMM profile | Pending | Hardware-gated profile. | | 11 | A3 HCCS HCOMM profile | Pending | Hardware-gated profile. | @@ -88,6 +88,18 @@ It is updated as each documented feature is completed and verified. Imports use shared-memory backed mappings in the session runner, imported handles remain opaque on the parent, and owner frees wait for live imports and slot refs to drain. +- Added L4-brokered Global CommDomain setup without MPI. L2 export + descriptors are collected by L3, assembled by L4, returned to every L3/L2 + for import, and released after the L4 DAG drain by default. Domains created + with `retain_after_run=True` remain live for a later run until explicitly + released or the Worker closes. Sim shm and A3 Fabric V2 use the same + descriptor ABI. +- Added startup-manifest delivery for pre-registered inner `CHIP_CALLABLE` + payloads, allowing remote sessions to resolve installed chip callables + before task dispatch. +- Remote buffers use L3-owned child-visible host buffers whenever the L3 has + forked chip children, while childless sim sessions keep the shared-memory + fallback. - Documented the v1 remote registry target/kind matrix, inner `INNER_L3_WORKER` visibility rules, remote `CHIP_CALLABLE` staged/inline payload contract, partial-register cleanup outcomes, and health-expiry @@ -95,6 +107,8 @@ It is updated as each documented feature is completed and verified. ## Verification +- Global CommDomain codec/validation tests and the Linux two-daemon sim + transaction test live in `tests/ut/py/test_global_comm_domain.py`. - Python focused sidecar/callable tests: `tests/ut/py/test_task_interface.py tests/ut/py/test_callable_identity.py` passed with `145 passed`. diff --git a/docs/remote-l3-worker-design/protocol.md b/docs/remote-l3-worker-design/protocol.md index d4e3c24da9..0c4c2b183c 100644 --- a/docs/remote-l3-worker-design/protocol.md +++ b/docs/remote-l3-worker-design/protocol.md @@ -296,15 +296,30 @@ Required remote controls: - `IMPORT_BUFFER` - `RELEASE_IMPORT` -Reserved future controls for Remote CommDomain: +Required Global CommDomain controls: - `COMM_INIT` - `ALLOC_DOMAIN` - `RELEASE_DOMAIN` - -The first Remote L3 task-dispatch cut rejects the reserved domain controls -with an unsupported-control reply. They become required only when Remote -CommDomain enters scope. +- `COPY_TO_DOMAIN` +- `COPY_FROM_DOMAIN` + +`COMM_INIT` validates the cluster id, node identity, communication profile, +global device ranks, and dense domain-rank table. `ALLOC_DOMAIN` is a +transaction with `PREPARE_EXPORT`, `IMPORT`, `COMMIT`, and `ABORT` phases. +Each L2 exports its local transport descriptor during prepare. L4 assembles +the complete rank-ordered table and sends it to every L3; each L3 forwards it +to its L2 children for import. No domain becomes visible to a remote task +before every node acknowledges `COMMIT`. + +`RELEASE_DOMAIN` is idempotent. The handle becomes released when the caller +requests release; physical backend teardown runs after the owning L4 DAG +drains. An allocation marked `retain_after_run` may remain live for a later L4 +run that reads kernel results. Explicit release or session shutdown requests +the same fence-ordered teardown. +`COPY_TO_DOMAIN` and `COPY_FROM_DOMAIN` are bounded smoke/control data +operations for a committed local window. They do not replace kernel data +movement through `CommContext`. The register-family controls are registry-scope-aware. `PREPARE_REGISTER_CALLABLE` carries: diff --git a/python/bindings/CMakeLists.txt b/python/bindings/CMakeLists.txt index 75d24dff82..f1df0c248b 100644 --- a/python/bindings/CMakeLists.txt +++ b/python/bindings/CMakeLists.txt @@ -61,6 +61,7 @@ target_include_directories(_task_interface PRIVATE ${CMAKE_SOURCE_DIR}/src/common/task_interface ${CMAKE_SOURCE_DIR}/src/common/worker ${CMAKE_SOURCE_DIR}/src/common/hierarchical + ${CMAKE_SOURCE_DIR}/src/common/platform/include ${CMAKE_SOURCE_DIR}/src/common/platform/include/common ${CMAKE_SOURCE_DIR}/src/common/platform/include/host ${CMAKE_CURRENT_SOURCE_DIR} diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index b89f523d23..0c56dc6451 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1585,6 +1585,36 @@ NB_MODULE(_task_interface, m) { nb::arg("allocation_id"), nb::arg("rank_count"), nb::arg("domain_rank"), "Pair to comm_alloc_domain_windows: collectively release the per-rank pool." ) + .def( + "comm_global_domain_prepare", + [](ChipWorker &self, uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, + uint32_t profile) { + auto [descriptor, local_window_base, actual_window_size] = + self.comm_global_domain_prepare(domain_id, domain_rank, rank_count, window_size, profile); + return nb::make_tuple( + nb::bytes(reinterpret_cast(descriptor.data()), descriptor.size()), local_window_base, + actual_window_size + ); + }, + nb::arg("domain_id"), nb::arg("domain_rank"), nb::arg("rank_count"), nb::arg("window_size"), + nb::arg("profile"), "Create a Global CommDomain local window and return its transport descriptor." + ) + .def( + "comm_global_domain_import", + [](ChipWorker &self, uint64_t domain_id, nb::bytes descriptors) { + std::vector descriptor_bytes( + reinterpret_cast(descriptors.c_str()), + reinterpret_cast(descriptors.c_str()) + descriptors.size() + ); + return self.comm_global_domain_import(domain_id, descriptor_bytes); + }, + nb::arg("domain_id"), nb::arg("descriptors"), + "Import a rank-ordered Global CommDomain descriptor table and return the device context." + ) + .def( + "comm_global_domain_release", &ChipWorker::comm_global_domain_release, nb::arg("domain_id"), + "Release a prepared or imported Global CommDomain." + ) .def("comm_barrier", &ChipWorker::comm_barrier, nb::arg("comm_handle"), "Synchronize all ranks.") .def( "comm_destroy", &ChipWorker::comm_destroy, nb::arg("comm_handle"), diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 8b5a844412..f3b1b968ed 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -701,6 +701,28 @@ inline void bind_worker(nb::module_ &m) { nb::arg("importer_worker_id"), nb::arg("owner_worker_id"), nb::arg("buffer_id"), nb::arg("generation"), nb::arg("import_id"), "Release an imported remote buffer mapping." ) + .def( + "remote_domain_control", + [](Worker &self, int worker_id, uint32_t control_name, nb::bytes command) { + if (!remote_l3::valid_control_name(control_name)) { + throw nb::value_error("control_name is not supported"); + } + std::vector command_bytes( + reinterpret_cast(command.c_str()), + reinterpret_cast(command.c_str()) + command.size() + ); + std::vector result; + { + nb::gil_scoped_release release; + result = self.control_remote_domain( + worker_id, static_cast(control_name), command_bytes + ); + } + return nb::bytes(reinterpret_cast(result.data()), result.size()); + }, + nb::arg("worker_id"), nb::arg("control_name"), nb::arg("command"), + "Send one Global CommDomain control to a remote L3 endpoint." + ) .def( "broadcast_unregister_all", [](Worker &self, nb::object digest) { @@ -742,6 +764,25 @@ inline void bind_worker(nb::module_ &m) { "If payload is a Python buffer, C++ stages it in POSIX shm and writes the shm name " "into the mailbox. Returns per-child ControlResult entries." ) + .def( + "control_payload", + [](Worker &self, WorkerType worker_type, int worker_id, uint64_t sub_cmd, nb::object payload, + nb::object timeout_s) { + std::string payload_bytes = buffer_to_string(payload, "payload"); + double timeout_val = timeout_s.is_none() ? -1.0 : nb::cast(timeout_s); + std::vector result; + { + nb::gil_scoped_release release; + result = self.control_payload( + worker_type, worker_id, sub_cmd, payload_bytes.data(), payload_bytes.size(), timeout_val + ); + } + return nb::bytes(reinterpret_cast(result.data()), result.size()); + }, + nb::arg("worker_type"), nb::arg("worker_id"), nb::arg("sub_cmd"), nb::arg("payload"), + nb::arg("timeout_s") = nb::none(), + "Drive one local worker control with a mutable staged payload and return its final bytes." + ) .def( "control_alloc_domain", &Worker::control_alloc_domain, nb::arg("worker_id"), nb::arg("request_shm_name"), nb::arg("reply_shm_name"), nb::call_guard(), diff --git a/python/simpler/global_comm_domain.py b/python/simpler/global_comm_domain.py new file mode 100644 index 0000000000..039d63bd71 --- /dev/null +++ b/python/simpler/global_comm_domain.py @@ -0,0 +1,572 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Versioned codecs and data models for L4-brokered Global CommDomains.""" + +from __future__ import annotations + +import enum +import struct +from dataclasses import dataclass + +GLOBAL_DOMAIN_VERSION = 1 +GLOBAL_DOMAIN_MAX_RANKS = 64 +GLOBAL_DOMAIN_HANDLE_BYTES = 256 +GLOBAL_DOMAIN_DESCRIPTOR = struct.Struct(" bytes: + _validate_descriptor(self) + handle = bytes(self.handle) + return GLOBAL_DOMAIN_DESCRIPTOR.pack( + int(self.version), + int(self.profile_id), + int(self.domain_rank), + int(self.rank_count), + int(self.mapping_size), + len(handle), + 0, + handle + b"\x00" * (GLOBAL_DOMAIN_HANDLE_BYTES - len(handle)), + ) + + @classmethod + def decode(cls, data: bytes) -> GlobalDomainDescriptor: + if len(data) != GLOBAL_DOMAIN_DESCRIPTOR_BYTES: + raise ValueError("global domain descriptor size mismatch") + version, profile_id, domain_rank, rank_count, mapping_size, handle_size, reserved, handle = ( + GLOBAL_DOMAIN_DESCRIPTOR.unpack(data) + ) + if reserved != 0: + raise ValueError("global domain descriptor reserved field must be zero") + if handle_size > GLOBAL_DOMAIN_HANDLE_BYTES: + raise ValueError("global domain descriptor handle size exceeds maximum") + descriptor = cls( + version=int(version), + profile_id=int(profile_id), + domain_rank=int(domain_rank), + rank_count=int(rank_count), + mapping_size=int(mapping_size), + handle=bytes(handle[:handle_size]), + ) + _validate_descriptor(descriptor) + return descriptor + + +@dataclass(frozen=True) +class GlobalCommInitCommand: + cluster_id: str + topology_hash: str + profile: str + node_rank: int + node_count: int + members: tuple[GlobalDomainMember, ...] + + +@dataclass(frozen=True) +class GlobalCommInitResult: + profile: str + max_ranks: int + descriptor_bytes: int + local_device_count: int + + +def resolve_global_comm_capability(*, platform: str, profile: str, local_device_count: int) -> GlobalCommInitResult: + """Return the capability implemented by the selected platform backend.""" + platform = str(platform) + profile = str(profile) + supported = (platform.endswith("sim") and profile == GLOBAL_DOMAIN_PROFILE_SIM) or ( + platform.startswith("a2a3") and not platform.endswith("sim") and profile == GLOBAL_DOMAIN_PROFILE_A3_FABRIC + ) + if not supported: + raise ValueError(f"Global CommDomain is not supported by platform {platform!r} with comm_profile {profile!r}") + if local_device_count <= 0: + raise ValueError("Global CommDomain capability requires at least one local device") + return GlobalCommInitResult( + profile=profile, + max_ranks=GLOBAL_DOMAIN_MAX_RANKS, + descriptor_bytes=GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + local_device_count=int(local_device_count), + ) + + +@dataclass(frozen=True) +class GlobalDomainCommand: + phase: GlobalDomainPhase + domain_id: int + generation: int + name: str + profile: str + window_size: int + members: tuple[GlobalDomainMember, ...] + buffers: tuple[GlobalDomainBuffer, ...] + descriptors: tuple[GlobalDomainDescriptor, ...] = () + + +@dataclass(frozen=True) +class GlobalDomainReleaseCommand: + domain_id: int + generation: int + + +@dataclass(frozen=True) +class GlobalDomainCopyCommand: + domain_id: int + generation: int + domain_rank: int + offset: int + nbytes: int + data: bytes = b"" + + +class _Reader: + def __init__(self, data: bytes) -> None: + self._data = data + self._offset = 0 + + def _take(self, size: int, field: str) -> bytes: + if size < 0 or self._offset > len(self._data) or size > len(self._data) - self._offset: + raise ValueError(f"global domain wire truncated {field}") + result = self._data[self._offset : self._offset + size] + self._offset += size + return result + + def u32(self) -> int: + return int(struct.unpack(" int: + return int(struct.unpack(" int: + return int(struct.unpack(" str: + size = self.u32() + if size > GLOBAL_DOMAIN_MAX_STRING_BYTES: + raise ValueError(f"global domain wire {field} exceeds maximum") + return self._take(size, field).decode("utf-8") + + def blob(self, maximum: int, field: str) -> bytes: + size = self.u32() + if size > maximum: + raise ValueError(f"global domain wire {field} exceeds maximum") + return self._take(size, field) + + def fixed(self, size: int, field: str) -> bytes: + return self._take(size, field) + + def done(self, field: str) -> None: + if self._offset != len(self._data): + raise ValueError(f"global domain wire trailing bytes after {field}") + + +def _put_string(out: bytearray, value: str, field: str) -> None: + encoded = str(value).encode("utf-8") + if len(encoded) > GLOBAL_DOMAIN_MAX_STRING_BYTES: + raise ValueError(f"global domain wire {field} exceeds maximum") + out.extend(struct.pack(" None: + encoded = bytes(value) + if len(encoded) > maximum: + raise ValueError(f"global domain wire {field} exceeds maximum") + out.extend(struct.pack(" None: + if descriptor.version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global domain descriptor version mismatch") + if descriptor.profile_id not in GLOBAL_DOMAIN_PROFILE_IDS.values(): + raise ValueError("global domain descriptor profile is unknown") + if descriptor.rank_count <= 0 or descriptor.rank_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor rank_count is invalid") + if descriptor.domain_rank < 0 or descriptor.domain_rank >= descriptor.rank_count: + raise ValueError("global domain descriptor domain_rank is invalid") + if descriptor.mapping_size <= 0: + raise ValueError("global domain descriptor mapping_size must be positive") + if not descriptor.handle or len(descriptor.handle) > GLOBAL_DOMAIN_HANDLE_BYTES: + raise ValueError("global domain descriptor handle size is invalid") + + +def validate_member_table(members: tuple[GlobalDomainMember, ...]) -> None: + if not members or len(members) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain members must contain between 1 and 64 devices") + ranks = [member.domain_rank for member in members] + if ranks != list(range(len(members))): + raise ValueError("global domain members must be in dense domain-rank order") + devices = [(member.node_worker_id, member.local_worker_id) for member in members] + if len(set(devices)) != len(devices): + raise ValueError("global domain members contain duplicate node/local devices") + if any(node < 0 or local < 0 for node, local in devices): + raise ValueError("global domain member node/local ids must be non-negative") + global_ranks = [member.global_device_rank for member in members] + if len(set(global_ranks)) != len(global_ranks) or any(rank < 0 for rank in global_ranks): + raise ValueError("global domain members require unique non-negative global device ranks") + + +def validate_descriptor_table( + descriptors: tuple[GlobalDomainDescriptor, ...], *, rank_count: int, profile: str +) -> None: + if profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {profile!r}") + if len(descriptors) != rank_count: + raise ValueError("global domain descriptor table is incomplete") + expected_profile = GLOBAL_DOMAIN_PROFILE_IDS[profile] + ranks: set[int] = set() + mapping_size: int | None = None + for descriptor in descriptors: + _validate_descriptor(descriptor) + if descriptor.profile_id != expected_profile or descriptor.rank_count != rank_count: + raise ValueError("global domain descriptor profile or rank_count mismatch") + if descriptor.domain_rank in ranks: + raise ValueError("global domain descriptor table contains a duplicate rank") + ranks.add(descriptor.domain_rank) + if mapping_size is None: + mapping_size = descriptor.mapping_size + elif mapping_size != descriptor.mapping_size: + raise ValueError("global domain descriptor mapping sizes differ") + if ranks != set(range(rank_count)): + raise ValueError("global domain descriptor table has missing ranks") + + +def _put_member(out: bytearray, member: GlobalDomainMember) -> None: + out.extend( + struct.pack( + " GlobalDomainMember: + return GlobalDomainMember( + node_worker_id=reader.i32(), + local_worker_id=reader.u32(), + global_device_rank=reader.u32(), + domain_rank=reader.u32(), + ) + + +def encode_comm_init(command: GlobalCommInitCommand) -> bytes: + validate_member_table(command.members) + if not command.cluster_id or not command.topology_hash: + raise ValueError("global comm init cluster_id and topology_hash must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if command.node_rank < 0 or command.node_count <= 0 or command.node_rank >= command.node_count: + raise ValueError("global comm init node identity is invalid") + out = bytearray(struct.pack(" GlobalCommInitCommand: + reader = _Reader(data) + version = reader.u32() + if version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global comm init version mismatch") + node_rank = reader.u32() + node_count = reader.u32() + cluster_id = reader.string("cluster_id") + topology_hash = reader.string("topology_hash") + profile = reader.string("profile") + member_count = reader.u32() + if member_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global comm init member count exceeds maximum") + members = tuple(_read_member(reader) for _ in range(member_count)) + reader.done("COMM_INIT") + command = GlobalCommInitCommand(cluster_id, topology_hash, profile, node_rank, node_count, members) + validate_member_table(command.members) + if not command.cluster_id or not command.topology_hash: + raise ValueError("global comm init cluster_id and topology_hash must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if node_count <= 0 or node_rank >= node_count: + raise ValueError("global comm init node identity is invalid") + return command + + +def encode_comm_init_result(result: GlobalCommInitResult) -> bytes: + out = bytearray( + struct.pack( + " GlobalCommInitResult: + reader = _Reader(data) + max_ranks = reader.u32() + descriptor_bytes = reader.u32() + local_device_count = reader.u32() + profile = reader.string("profile") + reader.done("COMM_INIT result") + result = GlobalCommInitResult( + profile=profile, + max_ranks=max_ranks, + descriptor_bytes=descriptor_bytes, + local_device_count=local_device_count, + ) + return result + + +def encode_domain_command(command: GlobalDomainCommand) -> bytes: + validate_member_table(command.members) + if command.domain_id == 0 or command.generation == 0 or command.window_size <= 0: + raise ValueError("global domain command identity and window_size must be positive") + if not command.name: + raise ValueError("global domain command name must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if len({buffer.name for buffer in command.buffers}) != len(command.buffers): + raise ValueError("global domain command contains duplicate buffer names") + if any(not buffer.name or buffer.nbytes <= 0 for buffer in command.buffers): + raise ValueError("global domain buffers require a name and positive size") + if sum(buffer.nbytes for buffer in command.buffers) > command.window_size: + raise ValueError("global domain buffers exceed the requested window") + if command.descriptors: + validate_descriptor_table(command.descriptors, rank_count=len(command.members), profile=command.profile) + if command.phase in (GlobalDomainPhase.IMPORT, GlobalDomainPhase.COMMIT): + if len(command.descriptors) != len(command.members): + raise ValueError("global domain IMPORT/COMMIT requires a complete descriptor table") + elif command.descriptors: + raise ValueError("global domain PREPARE/ABORT must not carry descriptors") + + out = bytearray( + struct.pack( + " GlobalDomainCommand: + reader = _Reader(data) + version = reader.u32() + if version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global domain command version mismatch") + try: + phase = GlobalDomainPhase(reader.u32()) + except ValueError as exc: + raise ValueError("global domain command phase is unknown") from exc + domain_id = reader.u64() + generation = reader.u64() + window_size = reader.u64() + name = reader.string("name") + profile = reader.string("profile") + member_count = reader.u32() + if member_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command member count exceeds maximum") + members = tuple(_read_member(reader) for _ in range(member_count)) + buffer_count = reader.u32() + if buffer_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command buffer count exceeds maximum") + buffers = tuple(GlobalDomainBuffer(reader.string("buffer.name"), reader.u64()) for _ in range(buffer_count)) + descriptor_count = reader.u32() + if descriptor_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command descriptor count exceeds maximum") + descriptors = tuple( + GlobalDomainDescriptor.decode(reader.fixed(GLOBAL_DOMAIN_DESCRIPTOR_BYTES, "descriptor")) + for _ in range(descriptor_count) + ) + reader.done("ALLOC_DOMAIN") + command = GlobalDomainCommand( + phase=phase, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=window_size, + members=members, + buffers=buffers, + descriptors=descriptors, + ) + encode_domain_command(command) + return command + + +def encode_descriptor_table(descriptors: tuple[GlobalDomainDescriptor, ...]) -> bytes: + if len(descriptors) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor table exceeds maximum") + return struct.pack(" tuple[GlobalDomainDescriptor, ...]: + reader = _Reader(data) + count = reader.u32() + if count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor table exceeds maximum") + descriptors = tuple( + GlobalDomainDescriptor.decode(reader.fixed(GLOBAL_DOMAIN_DESCRIPTOR_BYTES, "descriptor")) for _ in range(count) + ) + reader.done("descriptor table") + return descriptors + + +def encode_release_command(command: GlobalDomainReleaseCommand) -> bytes: + if command.domain_id == 0 or command.generation == 0: + raise ValueError("global domain release identity must be positive") + return struct.pack(" GlobalDomainReleaseCommand: + if len(data) != struct.calcsize(" bytes: + if ( + command.domain_id == 0 + or command.generation == 0 + or command.domain_rank < 0 + or command.offset < 0 + or command.nbytes <= 0 + or command.nbytes > GLOBAL_DOMAIN_MAX_COPY_BYTES + ): + raise ValueError("global domain copy fields are invalid") + if include_data and len(command.data) != command.nbytes: + raise ValueError("global domain copy payload size mismatch") + if not include_data and command.data: + raise ValueError("global domain copy-from request must not contain data") + out = bytearray( + struct.pack( + " GlobalDomainCopyCommand: + header_size = struct.calcsize(" bytes: + out = bytearray() + _put_blob(out, data, GLOBAL_DOMAIN_MAX_COPY_BYTES, "copy result") + return bytes(out) + + +def decode_copy_result(data: bytes) -> bytes: + reader = _Reader(data) + result = reader.blob(GLOBAL_DOMAIN_MAX_COPY_BYTES, "copy result") + reader.done("copy result") + return result diff --git a/python/simpler/mpi_l3_session.py b/python/simpler/mpi_l3_session.py new file mode 100644 index 0000000000..82581f314d --- /dev/null +++ b/python/simpler/mpi_l3_session.py @@ -0,0 +1,257 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""MPI-launched Remote L3 session runner. + +Each mpirun rank runs one ordinary Remote L3 session. The L4 parent still owns +the task/control lane over the existing RemoteL3 socket transport; MPI is used +only when the full mpirun group participates in one Global CommDomain prepare. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import signal +import socket +import struct +import sys +import time +from typing import Any + +from .global_comm_domain import ( + GlobalDomainCommand, + GlobalDomainPhase, + GlobalDomainReleaseCommand, + decode_descriptor_table, + encode_descriptor_table, + validate_descriptor_table, +) +from .remote_l3_session import _format_remote_error, run_session +from .worker import Worker + + +class MpiGlobalDomainExchange: + """Descriptor exchange hook for a full mpirun-launched L3 group.""" + + def __init__(self, comm: Any, *, group_worker_ids: tuple[int, ...], timeout_s: float) -> None: + self._comm = comm + self._rank = int(comm.Get_rank()) + self._group_worker_ids = tuple(int(worker_id) for worker_id in group_worker_ids) + self._group_worker_id_set = set(self._group_worker_ids) + self._timeout_s = float(timeout_s) + if not (self._timeout_s > 0 and math.isfinite(self._timeout_s)): + raise ValueError("MPI Global CommDomain timeout must be a positive finite number") + + def _allgather(self, payload: Any, *, operation: str, on_timeout) -> list[Any]: + request = self._comm.iallgather(payload) + deadline = time.monotonic() + self._timeout_s + while True: + complete, gathered = request.test() + if complete: + return list(gathered) + if time.monotonic() >= deadline: + on_timeout() + try: + self._comm.Abort(1) + except BaseException as exc: # noqa: BLE001 + raise TimeoutError(f"MPI Global CommDomain {operation} timed out") from exc + raise TimeoutError(f"MPI Global CommDomain {operation} timed out") + + def prepare_import(self, command: GlobalDomainCommand, inner_worker: Worker, worker_id: int) -> bytes | None: + if command.phase is not GlobalDomainPhase.PREPARE_EXPORT: + return None + if {int(member.node_worker_id) for member in command.members} != self._group_worker_id_set: + return None + + ok = True + error_message = "" + local_payload = b"" + release_command = GlobalDomainReleaseCommand(command.domain_id, command.generation) + + def release_local() -> None: + inner_worker._release_global_domain_node( # noqa: SLF001 + release_command, + suppress_errors=True, + ) + + try: + descriptors = inner_worker._prepare_global_domain_node(command, int(worker_id)) # noqa: SLF001 + local_payload = encode_descriptor_table(descriptors) + except BaseException as exc: # noqa: BLE001 + release_local() + ok = False + error_message = _format_remote_error( + f"mpi global domain prepare rank={self._rank} worker_id={worker_id}", + exc, + ) + + try: + gathered = self._allgather( + (self._rank, ok, error_message, local_payload), + operation="prepare", + on_timeout=release_local, + ) + except BaseException: + release_local() + raise + errors = [(rank, message) for rank, rank_ok, message, _payload in gathered if not rank_ok] + if errors: + if ok: + release_local() + rank, message = errors[0] + raise RuntimeError(f"MPI Global CommDomain prepare failed on rank {rank}: {message}") + + try: + descriptor_by_rank = {} + for _rank, _ok, _message, payload in sorted(gathered, key=lambda item: int(item[0])): + for descriptor in decode_descriptor_table(bytes(payload)): + if descriptor.domain_rank in descriptor_by_rank: + raise RuntimeError("MPI Global CommDomain exchange returned a duplicate domain rank") + descriptor_by_rank[descriptor.domain_rank] = descriptor + descriptors = tuple(descriptor_by_rank[rank] for rank in range(len(command.members))) + validate_descriptor_table(descriptors, rank_count=len(command.members), profile=command.profile) + except BaseException: + release_local() + raise + + import_command = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=command.domain_id, + generation=command.generation, + name=command.name, + profile=command.profile, + window_size=command.window_size, + members=command.members, + buffers=command.buffers, + descriptors=descriptors, + ) + import_ok = True + import_error = "" + try: + inner_worker._import_global_domain_node(import_command, int(worker_id)) # noqa: SLF001 + except BaseException as exc: # noqa: BLE001 + release_local() + import_ok = False + import_error = _format_remote_error( + f"mpi global domain import rank={self._rank} worker_id={worker_id}", + exc, + ) + + try: + statuses = self._allgather( + (self._rank, import_ok, import_error), + operation="import", + on_timeout=release_local, + ) + except BaseException: + release_local() + raise + import_errors = [(rank, message) for rank, rank_ok, message in statuses if not rank_ok] + if import_errors: + release_local() + rank, message = import_errors[0] + raise RuntimeError(f"MPI Global CommDomain import failed on rank {rank}: {message}") + return encode_descriptor_table(descriptors) + + +def _write_ready_file(ready_dir: str, rank: int, payload: dict[str, Any]) -> None: + os.makedirs(ready_dir, exist_ok=True) + final_path = os.path.join(ready_dir, f"rank-{int(rank)}.json") + tmp_path = f"{final_path}.tmp.{os.getpid()}" + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(payload, f, sort_keys=True) + f.write("\n") + os.replace(tmp_path, final_path) + + +def _send_ready_tcp(host: str, port: int, payload: dict[str, Any]) -> None: + data = json.dumps(payload, sort_keys=True).encode("utf-8") + with socket.create_connection((host, int(port)), timeout=10.0) as sock: + sock.sendall(struct.pack(" dict[str, Any]: + rank = int(comm.Get_rank()) + if rank == 0: + try: + with open(manifest_path, encoding="utf-8") as f: + payload = (True, json.load(f)) + except BaseException as exc: # noqa: BLE001 + payload = (False, f"{type(exc).__name__}: {exc}") + else: + payload = None + ok, value = comm.bcast(payload, root=0) + if not ok: + raise RuntimeError(f"MPI L3 rank0 failed to read group manifest {manifest_path!r}: {value}") + if not isinstance(value, dict): + raise ValueError("MPI L3 group manifest broadcast returned a non-object payload") + return value + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--group-manifest", required=True) + ns = parser.parse_args(argv) + + signal.signal(signal.SIGTERM, _raise_keyboard_interrupt) + try: + from mpi4py import MPI # noqa: PLC0415 + except ImportError as exc: + raise RuntimeError("simpler.mpi_l3_session requires mpi4py inside the mpirun Python environment") from exc + + comm = MPI.COMM_WORLD + rank = int(comm.Get_rank()) + world_size = int(comm.Get_size()) + + group_manifest = _load_group_manifest_from_rank0(comm, ns.group_manifest) + rank_manifests = group_manifest.get("rank_manifests") + if not isinstance(rank_manifests, list): + raise ValueError("MPI L3 group manifest requires a rank_manifests list") + if len(rank_manifests) != world_size: + raise ValueError("MPI L3 group manifest world size does not match MPI_COMM_WORLD") + manifest = dict(rank_manifests[rank]) + group_worker_ids = tuple(int(x) for x in group_manifest.get("worker_ids", ())) + if len(group_worker_ids) != world_size: + raise ValueError("MPI L3 group manifest worker_ids must match MPI_COMM_WORLD size") + + ready_dir = str(group_manifest["ready_dir"]) + ready_host = str(group_manifest.get("ready_host") or "") + ready_port = int(group_manifest.get("ready_port") or 0) + ready_token = str(group_manifest.get("ready_token") or "") + + def ready_writer(payload: dict[str, Any]) -> None: + payload = dict(payload) + payload["mpi_rank"] = rank + if ready_host: + payload["ready_token"] = ready_token + _send_ready_tcp(ready_host, ready_port, payload) + else: + _write_ready_file(ready_dir, rank, payload) + + exchange = MpiGlobalDomainExchange( + comm, + group_worker_ids=group_worker_ids, + timeout_s=float(manifest["session_timeout_s"]), + ) + return run_session( + manifest, + None, + ready_writer=ready_writer, + global_domain_prepare_import=exchange.prepare_import, + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index f0e71acebf..08d78114f8 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -48,6 +48,8 @@ def my_orch(orch, args, cfg): CommBufferSpec, CommDomainHandle, DataType, + GlobalCommDomainHandle, + GlobalCommDomainView, RemoteAddressSpace, TaskArgs, Tensor, @@ -414,6 +416,92 @@ def release_domain(self, handle: CommDomainHandle) -> None: """Collective release. Equivalent to ``handle.release()``.""" handle.release() + def allocate_global_domain( + self, + *, + name: str, + members: Sequence[tuple[int, int]], + window_size: int, + buffers: Sequence[CommBufferSpec] = (), + retain_after_run: bool = False, + ) -> GlobalCommDomainHandle: + """Create a CommDomain across local and/or remote L3 nodes without MPI. + + Each member is ``(l3_worker_id, local_l2_worker_id)``. The L3 worker + may have been registered by ``Worker.add_worker`` or + ``Worker.add_remote_worker``. L4 collects every L2 export descriptor, + sends the complete rank-ordered table back to every L3, and commits + only after all L2 imports succeed. + ``retain_after_run=True`` keeps the domain live after the current DAG + drains so a later run can inspect communication results; explicit + release or ``Worker.close()`` still tears it down. + """ + if self._worker is None: + raise RuntimeError("allocate_global_domain requires an Orchestrator bound to a Worker") + return self._worker._allocate_global_domain( + name=str(name), + members=tuple((int(node), int(local)) for node, local in members), + window_size=int(window_size), + buffers=list(buffers), + retain_after_run=bool(retain_after_run), + ) + + def release_global_domain(self, handle: GlobalCommDomainHandle) -> None: + """Request fence-ordered release of an L4-owned Global CommDomain.""" + handle.release() + + def get_global_domain(self, domain_id: int) -> GlobalCommDomainView: + """Return the committed L3-local view for a domain created by L4.""" + if self._worker is None: + raise RuntimeError("get_global_domain requires an Orchestrator bound to a Worker") + return self._worker._get_global_domain(int(domain_id)) + + @staticmethod + def _global_copy_range(handle: GlobalCommDomainHandle, *, buffer: str | None, offset: int, nbytes: int) -> int: + absolute = int(offset) + if absolute < 0 or nbytes <= 0: + raise ValueError("Global CommDomain copy offset must be non-negative and size must be positive") + limit = handle.mapping_size + if buffer is not None: + buffer_offset, buffer_nbytes = handle.buffer_range(str(buffer)) + if absolute > buffer_nbytes or nbytes > buffer_nbytes - absolute: + raise ValueError(f"Global CommDomain copy exceeds buffer {buffer!r}") + absolute += buffer_offset + elif absolute > limit or nbytes > limit - absolute: + raise ValueError("Global CommDomain copy exceeds the mapped window") + return absolute + + def copy_to_global_domain( + self, + handle: GlobalCommDomainHandle, + domain_rank: int, + data: bytes, + *, + buffer: str | None = None, + offset: int = 0, + ) -> None: + """Copy bytes into one rank's mapped window or named buffer.""" + payload = bytes(data) + absolute = self._global_copy_range(handle, buffer=buffer, offset=int(offset), nbytes=len(payload)) + if self._worker is None: + raise RuntimeError("copy_to_global_domain requires an Orchestrator bound to a Worker") + self._worker._copy_to_global_domain(handle, int(domain_rank), payload, absolute) + + def copy_from_global_domain( + self, + handle: GlobalCommDomainHandle, + domain_rank: int, + nbytes: int, + *, + buffer: str | None = None, + offset: int = 0, + ) -> bytes: + """Copy bytes from one rank's mapped window or named buffer.""" + absolute = self._global_copy_range(handle, buffer=buffer, offset=int(offset), nbytes=int(nbytes)) + if self._worker is None: + raise RuntimeError("copy_from_global_domain requires an Orchestrator bound to a Worker") + return self._worker._copy_from_global_domain(handle, int(domain_rank), int(nbytes), absolute) + def create_l3_l2_region(self, *, worker_id: int, payload_bytes: int, counter_bytes: int): """Create an L3-L2 communication region on one NEXT_LEVEL chip worker.""" if self._worker is None: diff --git a/python/simpler/remote_l3_protocol.py b/python/simpler/remote_l3_protocol.py index 24ca8ab402..f1263b80f3 100644 --- a/python/simpler/remote_l3_protocol.py +++ b/python/simpler/remote_l3_protocol.py @@ -64,6 +64,8 @@ class ControlName(enum.IntEnum): COMM_INIT = 13 ALLOC_DOMAIN = 14 RELEASE_DOMAIN = 15 + COPY_TO_DOMAIN = 16 + COPY_FROM_DOMAIN = 17 class RemoteRegistryTarget(enum.IntEnum): diff --git a/python/simpler/remote_l3_session.py b/python/simpler/remote_l3_session.py index c62a613779..b784a64633 100644 --- a/python/simpler/remote_l3_session.py +++ b/python/simpler/remote_l3_session.py @@ -42,6 +42,19 @@ parse_python_import_target, validate_hashid, ) +from .global_comm_domain import ( + GlobalDomainCommand, + GlobalDomainPhase, + GlobalDomainReleaseCommand, + decode_comm_init, + decode_copy_command, + decode_domain_command, + decode_release_command, + encode_comm_init_result, + encode_copy_result, + encode_descriptor_table, + resolve_global_comm_capability, +) from .remote_l3_protocol import ( PROTOCOL_VERSION, CallableKind, @@ -75,7 +88,7 @@ send_frame, ) from .task_interface import ChipCallable, TaskArgs, Tensor -from .worker import Worker +from .worker import Worker, _NoHostBufferChildrenError sys.modules.setdefault("simpler.remote_l3_session", sys.modules[__name__]) @@ -111,6 +124,7 @@ class _RemoteBufferEntry: nbytes: int generation: int address_space: RemoteAddressSpace + owner: Worker | None = None offset: int = 0 released: bool = False @@ -120,16 +134,27 @@ def addr(self) -> int: buf = self.data.buf assert buf is not None return ctypes.addressof(ctypes.c_char.from_buffer(buf)) + if hasattr(self.data, "data_ptr"): + return int(self.data.data_ptr) return ctypes.addressof(self.data) @property def shm_name(self) -> str: - if not isinstance(self.data, shared_memory.SharedMemory): - raise ValueError("remote buffer is not backed by SharedMemory") - return self.data.name + if isinstance(self.data, shared_memory.SharedMemory): + return self.data.name + if hasattr(self.data, "shm_name"): + return str(self.data.shm_name) + raise ValueError("remote buffer is not backed by SharedMemory") def close(self, *, unlink: bool = False) -> None: if not isinstance(self.data, shared_memory.SharedMemory): + if self.owner is not None: + owner, self.owner = self.owner, None + # HostBuffer itself owns a memoryview. Release it before asking + # Worker to close the backing shm so teardown is prompt and does + # not report a false live-view warning. + self.data.buffer.release() + owner.free_host_buffer(self.data) return self.data.close() if unlink: @@ -149,10 +174,10 @@ def _load_import_target(target: str) -> Callable[..., Any]: return obj -def _bind_listener(host: str) -> socket.socket: +def _bind_listener(host: str, port: int = 0) -> socket.socket: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind((host, 0)) + sock.bind((host, int(port))) sock.listen(1) return sock @@ -461,6 +486,20 @@ def _control_reply( send_frame(conn, FrameHeader(FrameType.CONTROL_REPLY, session_id, worker_id, sequence), payload) +def _control_result_reply( + conn: socket.socket, + manifest: dict[str, Any], + sequence: int, + control_name: ControlName, + version: int, + result: bytes, +) -> None: + session_id = int(manifest["session_id"]) + worker_id = int(manifest["worker_id"]) + payload = encode_control_reply(sequence, control_name, version, 0, "", bytes(result)) + send_frame(conn, FrameHeader(FrameType.CONTROL_REPLY, session_id, worker_id, sequence), payload) + + def _copy_command_header(data: bytes) -> tuple[int, int, int, int, int, bytes]: if len(data) < 36: raise ValueError("remote buffer copy command is truncated") @@ -533,6 +572,7 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 inner_worker: Worker, manifest_inner_handles: dict[tuple[CallableKind, bytes], CallableHandle] | None = None, manifest_dispatch_registry: dict[bytes, Callable[..., Any]] | None = None, + global_domain_prepare_import: Callable[[GlobalDomainCommand, Worker, int], bytes | None] | None = None, ) -> None: session_id = int(manifest["session_id"]) worker_id = int(manifest["worker_id"]) @@ -544,12 +584,13 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 next_export_id = 1 next_import_id = 1 buffers: dict[tuple[int, ...], _RemoteBufferEntry] = {} + global_comm_inits: dict[str, Any] = {} hello = HelloPayload( session_id=session_id, worker_id=worker_id, protocol_version=PROTOCOL_VERSION, - comm_profile=str(manifest["transport"]), + comm_profile=str(manifest.get("comm_profile", manifest["transport"])), feature_flags=0, ready_state=ReadyState.READY, ) @@ -656,9 +697,24 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 buffer_id = next_buffer_id next_buffer_id += 1 generation = 1 - buf = shared_memory.SharedMemory(create=True, size=int(nbytes)) + try: + buf = inner_worker.create_host_buffer(int(nbytes)) + entry = _RemoteBufferEntry( + buf, + int(nbytes), + generation, + RemoteAddressSpace.REMOTE_DEVICE, + owner=inner_worker, + ) + except _NoHostBufferChildrenError: + buf = shared_memory.SharedMemory(create=True, size=int(nbytes)) + entry = _RemoteBufferEntry( + buf, + int(nbytes), + generation, + RemoteAddressSpace.REMOTE_DEVICE, + ) key = _buffer_key(buffer_id, generation) - entry = _RemoteBufferEntry(buf, int(nbytes), generation, RemoteAddressSpace.REMOTE_DEVICE) buffers[key] = entry remote_addr = entry.addr result = struct.pack( @@ -846,19 +902,143 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 _control_reply( conn, manifest, header.sequence, control.control_name, control.control_version, 0, "" ) - elif control.control_name in ( - ControlName.COMM_INIT, - ControlName.ALLOC_DOMAIN, - ControlName.RELEASE_DOMAIN, - ): + elif control.control_name == ControlName.COMM_INIT: + command = decode_comm_init(control.command_bytes) + manifest_profile = str(manifest.get("comm_profile", manifest["transport"])) + if command.cluster_id != str(manifest.get("cluster_id", "")): + raise ValueError("COMM_INIT cluster_id does not match the session manifest") + if command.profile != manifest_profile: + raise ValueError("COMM_INIT profile does not match the session manifest") + if command.node_rank != int(manifest.get("node_rank", 0)) or command.node_count != int( + manifest.get("node_count", 1) + ): + raise ValueError("COMM_INIT node identity does not match the session manifest") + global_ranks = tuple(int(rank) for rank in manifest.get("global_device_ranks", ())) + local_members = tuple( + member for member in command.members if member.node_worker_id == worker_id + ) + if not local_members: + raise ValueError("COMM_INIT topology has no local members") + for member in local_members: + if member.local_worker_id < 0 or member.local_worker_id >= len(global_ranks): + raise ValueError("COMM_INIT local worker id exceeds the session device list") + if member.global_device_rank != global_ranks[member.local_worker_id]: + raise ValueError("COMM_INIT global device rank does not match the manifest") + prior = global_comm_inits.get(command.topology_hash) + if prior is not None and prior != command: + raise ValueError("COMM_INIT topology hash conflicts with an earlier command") + result = resolve_global_comm_capability( + platform=str(manifest["platform"]), + profile=manifest_profile, + local_device_count=len(global_ranks), + ) + global_comm_inits[command.topology_hash] = command + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + encode_comm_init_result(result), + ) + elif control.control_name == ControlName.ALLOC_DOMAIN: + command = decode_domain_command(control.command_bytes) + if not any( + init.profile == command.profile and init.members == command.members + for init in global_comm_inits.values() + ): + raise RuntimeError("ALLOC_DOMAIN requires a matching COMM_INIT topology") + if command.phase is GlobalDomainPhase.PREPARE_EXPORT: + if command.descriptors: + raise ValueError("PREPARE_EXPORT must not carry descriptors") + result = ( + global_domain_prepare_import(command, inner_worker, worker_id) + if global_domain_prepare_import is not None + else None + ) + if result is None: + descriptors = inner_worker._prepare_global_domain_node(command, worker_id) # noqa: SLF001 + result = encode_descriptor_table(descriptors) + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + result, + ) + elif command.phase is GlobalDomainPhase.IMPORT: + inner_worker._import_global_domain_node(command, worker_id) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif command.phase is GlobalDomainPhase.COMMIT: + inner_worker._commit_global_domain_node(command) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif command.phase is GlobalDomainPhase.ABORT: + inner_worker._release_global_domain_node( # noqa: SLF001 + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + else: + raise ValueError("ALLOC_DOMAIN phase is not supported") + elif control.control_name == ControlName.RELEASE_DOMAIN: + command = decode_release_command(control.command_bytes) + inner_worker._release_global_domain_node(command) # noqa: SLF001 _control_reply( conn, manifest, header.sequence, control.control_name, control.control_version, - 1, - f"unsupported reserved remote domain control {control.control_name.name}", + 0, + "", + ) + elif control.control_name == ControlName.COPY_TO_DOMAIN: + command = decode_copy_command(control.command_bytes, include_data=True) + inner_worker._copy_global_domain_node(command, copy_to_device=True) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif control.control_name == ControlName.COPY_FROM_DOMAIN: + command = decode_copy_command(control.command_bytes, include_data=False) + result = inner_worker._copy_global_domain_node(command, copy_to_device=False) # noqa: SLF001 + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + encode_copy_result(result), ) else: _control_reply( @@ -924,7 +1104,13 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 _INNER_HANDLES.clear() -def run_session(manifest: dict[str, Any], ready_fd: int) -> int: +def run_session( + manifest: dict[str, Any], + ready_fd: int | None, + *, + ready_writer: Callable[[dict[str, Any]], None] | None = None, + global_domain_prepare_import: Callable[[GlobalDomainCommand, Worker, int], bytes | None] | None = None, +) -> int: inner_worker = Worker( level=3, platform=str(manifest["platform"]), @@ -937,6 +1123,16 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: health_sock: socket.socket | None = None stop_health = threading.Event() health_thread: threading.Thread | None = None + ready_sent = False + + def _publish_ready(payload: dict[str, Any]) -> None: + nonlocal ready_sent + if ready_writer is not None: + ready_writer(payload) + elif ready_fd is not None: + _send_ready(ready_fd, payload) + ready_sent = True + try: # Validate the runtime command timeout wire value up front (rejects a # malformed session_timeout_s); the command lane itself idle-waits blocking. @@ -962,8 +1158,10 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: inner_worker.init(_startup_deadline=startup_deadline) listen_host = str(manifest.get("listen_host", "127.0.0.1")) - command_sock = _bind_listener(listen_host) - health_sock = _bind_listener(listen_host) + command_port = int(manifest.get("command_port", 0) or 0) + health_port = int(manifest.get("health_port", 0) or 0) + command_sock = _bind_listener(listen_host) if command_port == 0 else _bind_listener(listen_host, command_port) + health_sock = _bind_listener(listen_host) if health_port == 0 else _bind_listener(listen_host, health_port) health_thread = threading.Thread( target=_health_loop, args=(health_sock, stop_health, int(manifest["session_id"]), int(manifest["worker_id"])), @@ -973,8 +1171,7 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: command_port = int(command_sock.getsockname()[1]) health_port = int(health_sock.getsockname()[1]) - _send_ready( - ready_fd, + _publish_ready( { "ok": True, "command_host": str(manifest.get("connect_host", listen_host)), @@ -1001,13 +1198,21 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: # parent closes the command socket (read_frame sees EOF). conn.settimeout(None) with conn: - _run_command_loop(conn, manifest, inner_worker, manifest_inner_handles, manifest_dispatch_registry) + _run_command_loop( + conn, + manifest, + inner_worker, + manifest_inner_handles, + manifest_dispatch_registry, + global_domain_prepare_import, + ) return 0 except BaseException as exc: # noqa: BLE001 - try: - _send_ready(ready_fd, {"ok": False, "error": _format_remote_error("remote session startup", exc)}) - except OSError: - pass + if not ready_sent: + try: + _publish_ready({"ok": False, "error": _format_remote_error("remote session startup", exc)}) + except OSError: + pass return 1 finally: stop_health.set() diff --git a/python/simpler/remote_l3_worker.py b/python/simpler/remote_l3_worker.py index 3561ee95c4..fbb5a8a11e 100644 --- a/python/simpler/remote_l3_worker.py +++ b/python/simpler/remote_l3_worker.py @@ -26,6 +26,8 @@ import time from typing import Any +from .global_comm_domain import GLOBAL_DOMAIN_PROFILE_A3_FABRIC, GLOBAL_DOMAIN_PROFILE_IDS + def _read_exact(sock: socket.socket, n: int) -> bytes: data = bytearray() @@ -64,6 +66,23 @@ def _validate_manifest(manifest: dict[str, Any]) -> None: raise ValueError("manifest platform must be non-empty") if str(manifest["transport"]) != "sim": raise ValueError("only sim transport is accepted by simpler-remote-worker") + comm_profile = str(manifest.get("comm_profile", manifest["transport"])) + if comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError("manifest comm_profile is not supported") + if comm_profile == GLOBAL_DOMAIN_PROFILE_A3_FABRIC and not str(manifest["platform"]).startswith("a2a3"): + raise ValueError("manifest a3-fabric-v1 comm_profile requires an a2a3 platform") + if comm_profile == GLOBAL_DOMAIN_PROFILE_A3_FABRIC and str(manifest["platform"]).endswith("sim"): + raise ValueError("manifest a3-fabric-v1 comm_profile requires real A3 devices") + node_rank = int(manifest.get("node_rank", 0)) + node_count = int(manifest.get("node_count", 1)) + if node_count <= 0 or node_rank < 0 or node_rank >= node_count: + raise ValueError("manifest node identity is invalid") + device_ids = [int(device_id) for device_id in manifest.get("device_ids", [])] + global_device_ranks = [int(rank) for rank in manifest.get("global_device_ranks", range(len(device_ids)))] + if len(global_device_ranks) != len(device_ids): + raise ValueError("manifest global_device_ranks must match device_ids length") + if any(rank < 0 for rank in global_device_ranks) or len(set(global_device_ranks)) != len(global_device_ranks): + raise ValueError("manifest global_device_ranks must be unique and non-negative") def _session_timeout_s(manifest: dict[str, Any]) -> float: diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 332b436eb9..bef445bdf3 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -130,6 +130,8 @@ def _assert_bindings_match_source_tree() -> None: _assert_bindings_match_source_tree() +from .global_comm_domain import GlobalDomainBuffer, GlobalDomainMember # noqa: E402 + __all__ = [ "DataType", "get_element_size", @@ -163,6 +165,8 @@ def _assert_bindings_match_source_tree() -> None: "CommBufferSpec", "ChipDomainContext", "CommDomainHandle", + "GlobalCommDomainHandle", + "GlobalCommDomainView", ] COMM_MAX_RANK_NUM = 64 @@ -1017,6 +1021,141 @@ def __repr__(self) -> str: return f"CommDomainHandle(name={self.name!r}, workers={self.workers}, {state})" +class GlobalCommDomainHandle: + """L4-owned handle for one CommDomain spanning local and/or remote L3 nodes. + + The handle contains stable topology and buffer offsets only. Device + addresses remain in the L3/L2 process that imported the transport handles. + """ + + __slots__ = ( + "_freed", + "_release_fn", + "_released", + "buffers", + "domain_id", + "generation", + "mapping_size", + "members", + "name", + "retain_after_run", + ) + + def __init__( + self, + *, + name: str, + members: tuple[GlobalDomainMember, ...], + buffers: tuple[GlobalDomainBuffer, ...], + domain_id: int, + generation: int, + mapping_size: int, + retain_after_run: bool, + _release_fn, + ) -> None: + self.name = str(name) + self.members = tuple(members) + self.buffers = tuple(buffers) + self.domain_id = int(domain_id) + self.generation = int(generation) + self.mapping_size = int(mapping_size) + self.retain_after_run = bool(retain_after_run) + self._release_fn = _release_fn + self._released = False + self._freed = False + + def member(self, domain_rank: int) -> GlobalDomainMember: + if self._released: + raise RuntimeError(f"GlobalCommDomainHandle({self.name!r}) is already released") + rank = int(domain_rank) + if rank < 0 or rank >= len(self.members): + raise IndexError(f"global domain rank {rank} is out of range") + member = self.members[rank] + if member.domain_rank != rank: + raise RuntimeError("global domain member table is not rank ordered") + return member + + def buffer_range(self, name: str) -> tuple[int, int]: + if self._released: + raise RuntimeError(f"GlobalCommDomainHandle({self.name!r}) is already released") + offset = 0 + for buffer in self.buffers: + if buffer.name == name: + return offset, buffer.nbytes + offset += buffer.nbytes + raise KeyError(f"global domain {self.name!r} has no buffer {name!r}") + + @property + def released(self) -> bool: + return self._released + + @property + def freed(self) -> bool: + return self._freed + + def __repr__(self) -> str: + if self._freed: + state = "freed" + elif self._released: + state = "released" + else: + state = "live" + return f"GlobalCommDomainHandle(name={self.name!r}, members={len(self.members)}, {state})" + + def release(self) -> None: + if self._released: + return + self._release_fn(self) + self._released = True + + def __enter__(self) -> GlobalCommDomainHandle: + return self + + def __exit__(self, *_): + self.release() + + +class GlobalCommDomainView: + """L3-local imported view exposed to remote orchestration callables.""" + + __slots__ = ( + "_committed", + "contexts", + "domain_id", + "generation", + "mapping_size", + "members", + "name", + ) + + def __init__( + self, + *, + name: str, + members: tuple[GlobalDomainMember, ...], + contexts: dict[int, ChipDomainContext], + domain_id: int, + generation: int, + mapping_size: int, + ) -> None: + self.name = str(name) + self.members = tuple(members) + self.contexts = dict(contexts) + self.domain_id = int(domain_id) + self.generation = int(generation) + self.mapping_size = int(mapping_size) + self._committed = False + + def __getitem__(self, local_worker_id: int) -> ChipDomainContext: + if not self._committed: + raise RuntimeError(f"GlobalCommDomainView({self.name!r}) is not committed") + return self.contexts[int(local_worker_id)] + + @property + def committed(self) -> bool: + return self._committed + + # Process-wide RTLD_GLOBAL preload registry. host_runtime.so resolves its # undefined HostLogger / unified_log_* (and, on sim, sim_context_*) symbols # against these globals, so they must be loaded — exactly once — before any diff --git a/python/simpler/worker.py b/python/simpler/worker.py index d3fd1f1bd3..9795bf2cf3 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -63,15 +63,19 @@ def my_l4_orch(orch, args, config): import contextlib import ctypes import enum +import hashlib import importlib import json import math import os import re +import shutil import signal import socket import struct +import subprocess import sys +import tempfile import threading import time import uuid @@ -107,6 +111,51 @@ def my_l4_orch(orch, args, config): parse_python_callable_payload, parse_python_import_target, ) +from .global_comm_domain import ( + CTRL_GLOBAL_DOMAIN_COPY_FROM, + CTRL_GLOBAL_DOMAIN_COPY_TO, + CTRL_GLOBAL_DOMAIN_IMPORT, + CTRL_GLOBAL_DOMAIN_PREPARE, + CTRL_GLOBAL_DOMAIN_RELEASE, + GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + GLOBAL_DOMAIN_MAX_COPY_BYTES, + GLOBAL_DOMAIN_MAX_RANKS, + GLOBAL_DOMAIN_MAX_STRING_BYTES, + GLOBAL_DOMAIN_PROFILE_IDS, + GLOBAL_DOMAIN_VERSION, + LOCAL_COPY_REPLY, + LOCAL_COPY_REQUEST, + LOCAL_DOMAIN_MAGIC, + LOCAL_IMPORT_REPLY, + LOCAL_IMPORT_REQUEST, + LOCAL_PREPARE_REPLY, + LOCAL_PREPARE_REQUEST, + LOCAL_RELEASE_REQUEST, + GlobalCommInitCommand, + GlobalDomainBuffer, + GlobalDomainCommand, + GlobalDomainCopyCommand, + GlobalDomainDescriptor, + GlobalDomainMember, + GlobalDomainPhase, + GlobalDomainReleaseCommand, + decode_comm_init, + decode_comm_init_result, + decode_copy_command, + decode_copy_result, + decode_descriptor_table, + decode_domain_command, + decode_release_command, + encode_comm_init, + encode_comm_init_result, + encode_copy_command, + encode_copy_result, + encode_descriptor_table, + encode_domain_command, + encode_release_command, + resolve_global_comm_capability, + validate_descriptor_table, +) from .l3_l2_orch_comm import ( _CTRL_SHM_TOKEN_BYTES, _REGION_CREATE_REPLY, @@ -136,6 +185,8 @@ def my_l4_orch(orch, args, config): ChipWorker, CommBufferSpec, CommDomainHandle, + GlobalCommDomainHandle, + GlobalCommDomainView, RemoteAddressSpace, RemoteBufferExport, RemoteBufferHandle, @@ -228,6 +279,9 @@ def my_l4_orch(orch, args, config): # to close gracefully (so it unlinks the nested mailbox shms only it knows the # names of) before being SIGKILLed. This bounds that graceful wait. _ROLLBACK_GRACEFUL_TIMEOUT_S = 10.0 +# SIGKILL should make a child waitable promptly, but use a separate bounded +# window so normal close can confirm the reap without ever blocking forever. +_FORCED_REAP_TIMEOUT_S = 1.0 # Bounded re-check interval for a close() joiner waiting on an in-flight # _CloseAttempt. A joiner normally wakes immediately on the completing thread's # notify_all(); the timeout is a backstop so that if that notify is skipped (an @@ -313,6 +367,12 @@ def my_l4_orch(orch, args, config): _CTRL_L3_L2_REGION_CREATE = 16 _CTRL_L3_L2_REGION_RELEASE = 17 _CTRL_COMMITTED_DEVICE_MEMORY = 18 +# L4-to-local-L3 envelope for the Global CommDomain control protocol. The +# enclosed command uses remote_l3_protocol.ControlName. Values 19-23 belong +# to the local L3-to-L2 Global CommDomain controls. +_CTRL_GLOBAL_DOMAIN_NODE = 24 +_LOCAL_GLOBAL_CONTROL_HEADER = struct.Struct(" None: object.__setattr__(self, "platform", str(self.platform)) object.__setattr__(self, "runtime", str(self.runtime)) object.__setattr__(self, "transport", str(self.transport)) + object.__setattr__(self, "comm_profile", str(self.comm_profile)) object.__setattr__( self, "session_listen_host", @@ -431,9 +494,135 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "allow_wildcard_session_bind", bool(self.allow_wildcard_session_bind)) object.__setattr__(self, "device_ids", tuple(int(x) for x in self.device_ids)) + object.__setattr__(self, "global_device_ranks", tuple(int(x) for x in self.global_device_ranks)) object.__setattr__(self, "num_sub_workers", int(self.num_sub_workers)) if self.num_sub_workers < 0: raise ValueError("RemoteWorkerSpec.num_sub_workers must be non-negative") + if self.transport != "sim": + raise ValueError("RemoteWorkerSpec.transport must be 'sim' for the TCP daemon control plane") + if self.comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"RemoteWorkerSpec.comm_profile {self.comm_profile!r} is not supported") + if self.comm_profile == "a3-fabric-v1" and not self.platform.startswith("a2a3"): + raise ValueError("RemoteWorkerSpec.comm_profile 'a3-fabric-v1' requires an a2a3 platform") + if self.comm_profile == "a3-fabric-v1" and self.platform.endswith("sim"): + raise ValueError("RemoteWorkerSpec.comm_profile 'a3-fabric-v1' requires real A3 devices") + if self.global_device_ranks and len(self.global_device_ranks) != len(self.device_ids): + raise ValueError("RemoteWorkerSpec.global_device_ranks must match device_ids length") + if any(rank < 0 for rank in self.global_device_ranks) or len(set(self.global_device_ranks)) != len( + self.global_device_ranks + ): + raise ValueError("RemoteWorkerSpec.global_device_ranks must be unique and non-negative") + + +@dataclass(frozen=True) +class MpiL3GroupSpec: + """Describes a group of L3 workers launched by one parent-owned ``mpirun``.""" + + hosts: tuple[str, ...] + platform: str + command_port_base: int + health_port_base: int + device_ids_by_rank: tuple[tuple[int, ...], ...] + runtime: str = "tensormap_and_ringbuffer" + num_sub_workers_by_rank: tuple[int, ...] = () + transport: str = "sim" + comm_profile: str = "sim" + global_device_ranks_by_rank: tuple[tuple[int, ...], ...] = () + session_listen_hosts: tuple[str, ...] = () + connect_hosts: tuple[str, ...] = () + allow_wildcard_session_bind: bool = False + ready_host: str = "" + ready_port: int = 0 + mpirun_path: str = "mpirun" + mpirun_args: tuple[str, ...] = () + python_executable: str = field(default_factory=lambda: sys.executable) + + def __post_init__(self) -> None: # noqa: PLR0912 -- one place validates the public mpirun rank contract + hosts = tuple(str(host) for host in self.hosts) + if not hosts: + raise ValueError("MpiL3GroupSpec.hosts must be non-empty") + if not self.platform: + raise ValueError("MpiL3GroupSpec.platform must be non-empty") + device_ids_by_rank = tuple(tuple(int(device_id) for device_id in rank) for rank in self.device_ids_by_rank) + if len(device_ids_by_rank) != len(hosts): + raise ValueError("MpiL3GroupSpec.device_ids_by_rank must match hosts length") + if any(not rank for rank in device_ids_by_rank): + raise ValueError("MpiL3GroupSpec.device_ids_by_rank entries must be non-empty") + num_sub_workers_by_rank = ( + tuple(0 for _ in hosts) + if not self.num_sub_workers_by_rank + else tuple(int(count) for count in self.num_sub_workers_by_rank) + ) + if len(num_sub_workers_by_rank) != len(hosts): + raise ValueError("MpiL3GroupSpec.num_sub_workers_by_rank must match hosts length") + if any(count < 0 for count in num_sub_workers_by_rank): + raise ValueError("MpiL3GroupSpec.num_sub_workers_by_rank entries must be non-negative") + global_device_ranks_by_rank = ( + tuple(() for _ in hosts) + if not self.global_device_ranks_by_rank + else tuple(tuple(int(rank) for rank in ranks) for ranks in self.global_device_ranks_by_rank) + ) + if len(global_device_ranks_by_rank) != len(hosts): + raise ValueError("MpiL3GroupSpec.global_device_ranks_by_rank must match hosts length") + rank_pairs = zip(global_device_ranks_by_rank, device_ids_by_rank, strict=True) + for rank_index, rank_pair in enumerate(rank_pairs): + global_ranks, device_ids = rank_pair + if global_ranks and len(global_ranks) != len(device_ids): + raise ValueError( + f"MpiL3GroupSpec.global_device_ranks_by_rank[{rank_index}] must match device_ids length" + ) + if any(rank < 0 for rank in global_ranks) or len(set(global_ranks)) != len(global_ranks): + raise ValueError("MpiL3GroupSpec.global_device_ranks_by_rank entries must be unique non-negative ranks") + session_listen_hosts = ( + hosts if not self.session_listen_hosts else tuple(str(host) for host in self.session_listen_hosts) + ) + connect_hosts = hosts if not self.connect_hosts else tuple(str(host) for host in self.connect_hosts) + if len(session_listen_hosts) != len(hosts): + raise ValueError("MpiL3GroupSpec.session_listen_hosts must match hosts length") + if len(connect_hosts) != len(hosts): + raise ValueError("MpiL3GroupSpec.connect_hosts must match hosts length") + if any(not host for host in session_listen_hosts) or any(not host for host in connect_hosts): + raise ValueError("MpiL3GroupSpec host fields must be non-empty") + command_port_base = int(self.command_port_base) + health_port_base = int(self.health_port_base) + last_command_port = command_port_base + len(hosts) - 1 + last_health_port = health_port_base + len(hosts) - 1 + if command_port_base <= 0 or health_port_base <= 0 or last_command_port > 65535 or last_health_port > 65535: + raise ValueError("MpiL3GroupSpec command/health port ranges must be within 1..65535") + ready_host = str(self.ready_host) + ready_port = int(self.ready_port) + if ready_port < 0 or ready_port > 65535: + raise ValueError("MpiL3GroupSpec.ready_port must be within 0..65535") + if self.transport != "sim": + raise ValueError("MpiL3GroupSpec.transport must be 'sim' for the TCP control plane") + if self.comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"MpiL3GroupSpec.comm_profile {self.comm_profile!r} is not supported") + if self.comm_profile == "a3-fabric-v1" and not self.platform.startswith("a2a3"): + raise ValueError("MpiL3GroupSpec.comm_profile 'a3-fabric-v1' requires an a2a3 platform") + if self.comm_profile == "a3-fabric-v1" and self.platform.endswith("sim"): + raise ValueError("MpiL3GroupSpec.comm_profile 'a3-fabric-v1' requires real A3 devices") + if not self.mpirun_path: + raise ValueError("MpiL3GroupSpec.mpirun_path must be non-empty") + if not self.python_executable: + raise ValueError("MpiL3GroupSpec.python_executable must be non-empty") + object.__setattr__(self, "hosts", hosts) + object.__setattr__(self, "platform", str(self.platform)) + object.__setattr__(self, "runtime", str(self.runtime)) + object.__setattr__(self, "transport", str(self.transport)) + object.__setattr__(self, "comm_profile", str(self.comm_profile)) + object.__setattr__(self, "device_ids_by_rank", device_ids_by_rank) + object.__setattr__(self, "num_sub_workers_by_rank", num_sub_workers_by_rank) + object.__setattr__(self, "global_device_ranks_by_rank", global_device_ranks_by_rank) + object.__setattr__(self, "session_listen_hosts", session_listen_hosts) + object.__setattr__(self, "connect_hosts", connect_hosts) + object.__setattr__(self, "command_port_base", command_port_base) + object.__setattr__(self, "health_port_base", health_port_base) + object.__setattr__(self, "allow_wildcard_session_bind", bool(self.allow_wildcard_session_bind)) + object.__setattr__(self, "ready_host", ready_host) + object.__setattr__(self, "ready_port", ready_port) + object.__setattr__(self, "mpirun_path", str(self.mpirun_path)) + object.__setattr__(self, "mpirun_args", tuple(str(arg) for arg in self.mpirun_args)) + object.__setattr__(self, "python_executable", str(self.python_executable)) @dataclass(frozen=True) @@ -447,6 +636,53 @@ class _RemoteSession: pid: int +@dataclass(frozen=True) +class _MpiL3RankRuntime: + group_id: str + rank: int + worker_id: int + spec: RemoteWorkerSpec + command_host: str + command_port: int + health_host: str + health_port: int + + +@dataclass +class _MpiL3GroupRuntime: + group_id: str + spec: MpiL3GroupSpec + ranks: tuple[_MpiL3RankRuntime, ...] + process: subprocess.Popen[Any] | None = None + manifest_path: str | None = None + ready_dir: str | None = None + + +@dataclass +class _GlobalNodeDomainState: + command: GlobalDomainCommand + prepared_domain_ranks: set[int] = field(default_factory=set) + descriptors: dict[int, GlobalDomainDescriptor] = field(default_factory=dict) + local_window_bases: dict[int, int] = field(default_factory=dict) + mapping_sizes: dict[int, int] = field(default_factory=dict) + contexts: dict[int, ChipDomainContext] = field(default_factory=dict) + view: GlobalCommDomainView | None = None + phase: GlobalDomainPhase = GlobalDomainPhase.PREPARE_EXPORT + + +@dataclass(frozen=True) +class _GlobalNodeRuntime: + worker_id: int + device_ids: tuple[int, ...] + platform: str + comm_profile: str + global_device_ranks: tuple[int, ...] + node_rank: int + node_count: int + cluster_id: str + is_remote: bool + + _IdentitySnapshotEntry = tuple[bytes, Any, int, str, str] @@ -530,6 +766,11 @@ class HostBuffer: data_ptr: int nbytes: int buffer: memoryview + shm_name: str + + +class _NoHostBufferChildrenError(RuntimeError): + """The Worker has no process child that can attach a host buffer.""" def _rewrite_blob_host_addrs(buf: memoryview, blob_off: int, ranges: list[tuple[int, int, int]]) -> None: @@ -729,14 +970,17 @@ def _pack_py_callable_payload(target) -> bytes: def _chip_descriptor_context(worker: Worker) -> tuple[str, str]: platform = str(worker._config.get("platform", "")) runtime = str(worker._config.get("runtime", "")) - if platform or runtime: - return platform, runtime - contexts: list[tuple[str, str]] = [] + if platform or runtime: + contexts.append((platform, runtime)) for child in getattr(worker, "_next_level_workers", []): child_context = _chip_descriptor_context(child) if child_context != ("", ""): contexts.append(child_context) + for spec in getattr(worker, "_remote_worker_specs", []): + contexts.append((str(spec.platform), str(spec.runtime))) + for rank in getattr(worker, "_mpi_rank_by_worker_id", {}).values(): + contexts.append((str(rank.spec.platform), str(rank.spec.runtime))) if not contexts: return "", "" first = contexts[0] @@ -1295,6 +1539,25 @@ class _L2HostL3L2RegionStore: next_region_id: int = 1 +@dataclass +class _L2GlobalDomain: + domain_id: int + generation: int + domain_rank: int + rank_count: int + descriptor: GlobalDomainDescriptor + local_window_base: int + mapping_size: int + requested_window_size: int + device_ctx: int = 0 + descriptor_table: bytes = b"" + + +@dataclass +class _L2GlobalDomainStore: + domains: dict[int, _L2GlobalDomain] = field(default_factory=dict) + + @dataclass(frozen=True) class _L2HostL3L2RegionReplyMeta: payload_base: int @@ -1456,6 +1719,233 @@ def _sweep_l2_host_l3_l2_regions(store: _L2HostL3L2RegionStore) -> None: pass +def _open_global_domain_payload(buf: memoryview) -> tuple[SharedMemory, memoryview, int]: + payload_size = int(struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0]) + if payload_size <= 0: + raise RuntimeError("Global CommDomain control payload must be non-empty") + staged = SharedMemory(name=_read_ctrl_staged_shm_name(buf)) + staged_buf = cast(memoryview, staged.buf) + if payload_size > staged.size: + staged_buf.release() + staged.close() + raise RuntimeError("Global CommDomain control payload exceeds staged shm") + return staged, staged_buf, payload_size + + +def _validate_local_global_header( + magic: bytes, version: int, domain_id: int, generation: int, *, operation: str +) -> None: + if magic != LOCAL_DOMAIN_MAGIC or version != GLOBAL_DOMAIN_VERSION: + raise RuntimeError(f"{operation}: local protocol magic or version mismatch") + if domain_id == 0 or generation == 0: + raise RuntimeError(f"{operation}: domain identity must be positive") + + +def _handle_ctrl_global_domain_prepare(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < max(LOCAL_PREPARE_REQUEST.size, LOCAL_PREPARE_REPLY.size + GLOBAL_DOMAIN_DESCRIPTOR_BYTES): + raise RuntimeError("Global CommDomain prepare payload is too small") + fields = LOCAL_PREPARE_REQUEST.unpack_from(payload, 0) + magic, version, domain_id, generation, domain_rank, rank_count, profile_id, window_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="prepare") + if rank_count <= 0 or rank_count > GLOBAL_DOMAIN_MAX_RANKS or domain_rank >= rank_count: + raise RuntimeError("Global CommDomain prepare rank identity is invalid") + if profile_id not in GLOBAL_DOMAIN_PROFILE_IDS.values() or window_size <= 0: + raise RuntimeError("Global CommDomain prepare profile or window size is invalid") + prior = store.domains.get(int(domain_id)) + if prior is not None: + if ( + prior.generation != generation + or prior.domain_rank != domain_rank + or prior.rank_count != rank_count + or prior.descriptor.profile_id != profile_id + or prior.requested_window_size != window_size + ): + raise RuntimeError("Global CommDomain prepare conflicts with a live domain") + descriptor = prior.descriptor + local_base = prior.local_window_base + mapping_size = prior.mapping_size + else: + descriptor_bytes, local_base, mapping_size = cw._impl.comm_global_domain_prepare( + int(domain_id), + int(domain_rank), + int(rank_count), + int(window_size), + int(profile_id), + ) + try: + descriptor = GlobalDomainDescriptor.decode(bytes(descriptor_bytes)) + except BaseException: + with contextlib.suppress(BaseException): + cw._impl.comm_global_domain_release(int(domain_id)) + raise + if ( + descriptor.domain_rank != domain_rank + or descriptor.rank_count != rank_count + or descriptor.profile_id != profile_id + or descriptor.mapping_size != mapping_size + or descriptor.mapping_size < window_size + ): + cw._impl.comm_global_domain_release(int(domain_id)) + raise RuntimeError("Global CommDomain backend returned an inconsistent descriptor") + store.domains[int(domain_id)] = _L2GlobalDomain( + domain_id=int(domain_id), + generation=int(generation), + domain_rank=int(domain_rank), + rank_count=int(rank_count), + descriptor=descriptor, + local_window_base=int(local_base), + mapping_size=int(mapping_size), + requested_window_size=int(window_size), + ) + LOCAL_PREPARE_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + int(local_base), + int(mapping_size), + ) + start = LOCAL_PREPARE_REPLY.size + payload[start : start + GLOBAL_DOMAIN_DESCRIPTOR_BYTES] = descriptor.encode() + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_import(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_IMPORT_REQUEST.size: + raise RuntimeError("Global CommDomain import payload is truncated") + magic, version, domain_id, generation, descriptor_count = LOCAL_IMPORT_REQUEST.unpack_from(payload, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="import") + expected_size = LOCAL_IMPORT_REQUEST.size + int(descriptor_count) * GLOBAL_DOMAIN_DESCRIPTOR_BYTES + if descriptor_count <= 0 or descriptor_count > GLOBAL_DOMAIN_MAX_RANKS or expected_size > payload_size: + raise RuntimeError("Global CommDomain import descriptor table size is invalid") + entry = store.domains.get(int(domain_id)) + if entry is None or entry.generation != generation: + raise RuntimeError("Global CommDomain import requires a matching prepared domain") + descriptor_bytes = bytes(payload[LOCAL_IMPORT_REQUEST.size : expected_size]) + descriptors = tuple( + GlobalDomainDescriptor.decode(descriptor_bytes[offset : offset + GLOBAL_DOMAIN_DESCRIPTOR_BYTES]) + for offset in range(0, len(descriptor_bytes), GLOBAL_DOMAIN_DESCRIPTOR_BYTES) + ) + profile = next( + name for name, profile_id in GLOBAL_DOMAIN_PROFILE_IDS.items() if profile_id == entry.descriptor.profile_id + ) + validate_descriptor_table(descriptors, rank_count=entry.rank_count, profile=profile) + if descriptors[entry.domain_rank] != entry.descriptor: + raise RuntimeError("Global CommDomain import table does not contain the local exported descriptor") + if entry.descriptor_table and entry.descriptor_table != descriptor_bytes: + raise RuntimeError("Global CommDomain repeated import carries a different descriptor table") + if entry.device_ctx == 0: + entry.device_ctx = int(cw._impl.comm_global_domain_import(int(domain_id), descriptor_bytes)) + if entry.device_ctx == 0: + raise RuntimeError("Global CommDomain backend returned a zero device context") + entry.descriptor_table = descriptor_bytes + if payload_size < LOCAL_IMPORT_REPLY.size: + raise RuntimeError("Global CommDomain import reply capacity is too small") + LOCAL_IMPORT_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + entry.device_ctx, + entry.local_window_base, + entry.mapping_size, + ) + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_release(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_RELEASE_REQUEST.size: + raise RuntimeError("Global CommDomain release payload is truncated") + magic, version, domain_id, generation = LOCAL_RELEASE_REQUEST.unpack_from(payload, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="release") + entry = store.domains.get(int(domain_id)) + if entry is not None and entry.generation != generation: + raise RuntimeError("Global CommDomain release generation mismatch") + if entry is not None: + cw._impl.comm_global_domain_release(int(domain_id)) + store.domains.pop(int(domain_id), None) + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_copy( + cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore, *, copy_to_device: bool +) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_COPY_REQUEST.size: + raise RuntimeError("Global CommDomain copy payload is truncated") + magic, version, domain_id, generation, offset, nbytes = LOCAL_COPY_REQUEST.unpack_from(payload, 0) + operation = "copy-to" if copy_to_device else "copy-from" + _validate_local_global_header(magic, version, domain_id, generation, operation=operation) + entry = store.domains.get(int(domain_id)) + if entry is None or entry.generation != generation or entry.device_ctx == 0: + raise RuntimeError(f"Global CommDomain {operation} requires an imported live domain") + if nbytes <= 0 or nbytes > GLOBAL_DOMAIN_MAX_COPY_BYTES: + raise RuntimeError(f"Global CommDomain {operation} size is invalid") + if offset > entry.mapping_size or nbytes > entry.mapping_size - offset: + raise RuntimeError(f"Global CommDomain {operation} range exceeds the local window") + if copy_to_device: + data_offset = LOCAL_COPY_REQUEST.size + if data_offset + nbytes > payload_size: + raise RuntimeError("Global CommDomain copy-to data is truncated") + exported = ctypes.c_char.from_buffer(payload, data_offset) + try: + cw.copy_to(entry.local_window_base + int(offset), ctypes.addressof(exported), int(nbytes)) + finally: + del exported + else: + data_offset = LOCAL_COPY_REPLY.size + if data_offset + nbytes > payload_size: + raise RuntimeError("Global CommDomain copy-from reply capacity is too small") + exported = ctypes.c_char.from_buffer(payload, data_offset) + try: + cw.copy_from(ctypes.addressof(exported), entry.local_window_base + int(offset), int(nbytes)) + finally: + del exported + LOCAL_COPY_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + int(nbytes), + ) + finally: + payload.release() + staged.close() + + +def _sweep_l2_global_domains(cw: ChipWorker, store: _L2GlobalDomainStore) -> None: + first_error: Exception | None = None + for domain_id in list(store.domains): + store.domains.pop(domain_id, None) + try: + cw._impl.comm_global_domain_release(int(domain_id)) + except Exception as exc: # noqa: BLE001 + if first_error is None: + first_error = exc + if first_error is not None: + sys.stderr.write(f"[worker pid={os.getpid()}] Global CommDomain sweep release failed: {first_error}\n") + sys.stderr.flush() + + def _handle_ctrl_release_domain(cw: ChipWorker, buf: memoryview) -> None: """CTRL_RELEASE_DOMAIN handler — collective free for one allocation.""" request_shm_name = _read_shm_name(buf, _OFF_ARGS) @@ -1527,6 +2017,7 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de """ prepared = prepared if prepared is not None else set() l3_l2_region_store = _L2HostL3L2RegionStore() + global_domain_store = _L2GlobalDomainStore() # Post-fork host buffers mapped into this child. `host_buf_table` # owns the mmap per token (for unmap + teardown); `host_buf_ranges` is the # parent-VA → child-VA translation table the per-task blob rewrite consults, @@ -1585,7 +2076,9 @@ def handle_task() -> tuple[int, str]: code, msg = on_task_done_success() return code, msg - def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one branch per control sub-command + def handle_control( # noqa: PLR0912, PLR0915 -- one branch per control sub-command + sub_cmd: int, + ) -> tuple[int, str]: code = 0 msg = "" try: @@ -1682,6 +2175,26 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra _handle_ctrl_l3_l2_region_release(buf, l3_l2_region_store) elif sub_cmd == _CTRL_COMMITTED_DEVICE_MEMORY: struct.pack_into("Q", buf, _CTRL_OFF_RESULT, cw.committed_device_memory) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_PREPARE: + _handle_ctrl_global_domain_prepare(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_IMPORT: + _handle_ctrl_global_domain_import(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_RELEASE: + _handle_ctrl_global_domain_release(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_COPY_TO: + _handle_ctrl_global_domain_copy( + cw, + buf, + global_domain_store, + copy_to_device=True, + ) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_COPY_FROM: + _handle_ctrl_global_domain_copy( + cw, + buf, + global_domain_store, + copy_to_device=False, + ) else: raise RuntimeError(f"unknown control sub-command {int(sub_cmd)}") except Exception as e: # noqa: BLE001 @@ -1696,6 +2209,7 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra try: _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: + _sweep_l2_global_domains(cw, global_domain_store) _sweep_l2_host_l3_l2_regions(l3_l2_region_store) for host_shm, _lo, _hi, _base in host_buf_table.values(): try: @@ -1822,12 +2336,91 @@ def _read_config_from_mailbox(buf: memoryview) -> CallConfig: return cfg +def _run_local_global_domain_control( # noqa: PLR0912 -- one ordered dispatcher for the Global CommDomain protocol + inner_worker: Worker, + runtime: _GlobalNodeRuntime, + comm_inits: dict[str, GlobalCommInitCommand], + control_name: int, + request: bytes, +) -> bytes: + """Execute one Global CommDomain command inside an add_worker L3 child.""" + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + control = ControlName(control_name) + if control is ControlName.COMM_INIT: + command = decode_comm_init(request) + if command.cluster_id != runtime.cluster_id: + raise ValueError("COMM_INIT cluster_id does not match the local L3 topology") + if command.profile != runtime.comm_profile: + raise ValueError("COMM_INIT profile does not match the local L3 topology") + if command.node_rank != runtime.node_rank or command.node_count != runtime.node_count: + raise ValueError("COMM_INIT node identity does not match the local L3 topology") + local_members = tuple(member for member in command.members if member.node_worker_id == runtime.worker_id) + if not local_members: + raise ValueError("COMM_INIT topology has no local members") + for member in local_members: + if member.local_worker_id < 0 or member.local_worker_id >= len(runtime.global_device_ranks): + raise ValueError("COMM_INIT local worker id exceeds the local L3 device list") + if member.global_device_rank != runtime.global_device_ranks[member.local_worker_id]: + raise ValueError("COMM_INIT global device rank does not match the local L3 topology") + prior = comm_inits.get(command.topology_hash) + if prior is not None and prior != command: + raise ValueError("COMM_INIT topology hash conflicts with an earlier command") + capability = resolve_global_comm_capability( + platform=runtime.platform, + profile=runtime.comm_profile, + local_device_count=len(runtime.global_device_ranks), + ) + comm_inits[command.topology_hash] = command + return encode_comm_init_result(capability) + + if control is ControlName.ALLOC_DOMAIN: + command = decode_domain_command(request) + if not any(init.profile == command.profile and init.members == command.members for init in comm_inits.values()): + raise RuntimeError("ALLOC_DOMAIN requires a matching COMM_INIT topology") + if command.phase is GlobalDomainPhase.PREPARE_EXPORT: + if command.descriptors: + raise ValueError("PREPARE_EXPORT must not carry descriptors") + return encode_descriptor_table(inner_worker._prepare_global_domain_node(command, runtime.worker_id)) + if command.phase is GlobalDomainPhase.IMPORT: + inner_worker._import_global_domain_node(command, runtime.worker_id) + return b"" + if command.phase is GlobalDomainPhase.COMMIT: + inner_worker._commit_global_domain_node(command) + return b"" + if command.phase is GlobalDomainPhase.ABORT: + inner_worker._release_global_domain_node( + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + return b"" + raise ValueError("ALLOC_DOMAIN phase is not supported") + + if control is ControlName.RELEASE_DOMAIN: + inner_worker._release_global_domain_node(decode_release_command(request)) + return b"" + if control is ControlName.COPY_TO_DOMAIN: + inner_worker._copy_global_domain_node( + decode_copy_command(request, include_data=True), + copy_to_device=True, + ) + return b"" + if control is ControlName.COPY_FROM_DOMAIN: + result = inner_worker._copy_global_domain_node( + decode_copy_command(request, include_data=False), + copy_to_device=False, + ) + return encode_copy_result(result) + raise ValueError(f"unsupported local Global CommDomain control {int(control)}") + + def _child_worker_loop( buf: memoryview, registry: dict[int, Any], identity_table: dict[bytes, int], identity_refs: dict[bytes, int], inner_worker: Worker, + global_node: _GlobalNodeRuntime | None = None, ) -> None: """Runs in forked child process. Any-level Worker as child of its parent. @@ -1839,6 +2432,7 @@ def _child_worker_loop( into the inner Worker (see docs section 7). """ state_addr = _buffer_field_addr(buf, _OFF_STATE) + global_comm_inits: dict[str, GlobalCommInitCommand] = {} def handle_task() -> tuple[int, str]: digest = _read_task_digest(buf) @@ -1889,6 +2483,41 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: sub_cmd, context=f"child_worker level={inner_worker.level}", ) + elif sub_cmd == _CTRL_GLOBAL_DOMAIN_NODE: + if global_node is None: + raise RuntimeError("Global CommDomain control requires a local L3 child") + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < _LOCAL_GLOBAL_CONTROL_HEADER.size: + raise RuntimeError("local Global CommDomain control payload is truncated") + control_name, request_size, response_size = _LOCAL_GLOBAL_CONTROL_HEADER.unpack_from(payload, 0) + capacity = payload_size - _LOCAL_GLOBAL_CONTROL_HEADER.size + if request_size > capacity: + raise RuntimeError("local Global CommDomain request exceeds staged payload") + if response_size != 0: + raise RuntimeError("local Global CommDomain request contains a response") + start = _LOCAL_GLOBAL_CONTROL_HEADER.size + request = bytes(payload[start : start + request_size]) + response = _run_local_global_domain_control( + inner_worker, + global_node, + global_comm_inits, + int(control_name), + request, + ) + if len(response) > capacity: + raise RuntimeError("local Global CommDomain response exceeds staged payload") + payload[start : start + len(response)] = response + _LOCAL_GLOBAL_CONTROL_HEADER.pack_into( + payload, + 0, + int(control_name), + int(request_size), + len(response), + ) + finally: + payload.release() + staged.close() else: raise RuntimeError(f"unknown control sub-command {sub_cmd}") except Exception as e: # noqa: BLE001 @@ -1970,6 +2599,8 @@ class _RunResources: remote_slot_refs: list[RemoteBufferHandle] = field(default_factory=list) live_domains: dict[str, CommDomainHandle] = field(default_factory=dict) pending_release_domains: list[CommDomainHandle] = field(default_factory=list) + live_global_domains: dict[str, GlobalCommDomainHandle] = field(default_factory=dict) + pending_release_global_domains: list[GlobalCommDomainHandle] = field(default_factory=list) l3_l2_regions: list[Any] = field(default_factory=list) l3_l2_orch_comm_host_buffers: dict[int, int] = field(default_factory=dict) # True once the owning run's fence has claimed the domains above. A release @@ -2381,6 +3012,9 @@ def __init__( self._remote_worker_specs: list[RemoteWorkerSpec] = [] self._remote_worker_ids: list[int] = [] self._remote_sessions: list[_RemoteSession] = [] + self._mpi_l3_groups: list[_MpiL3GroupRuntime] = [] + self._mpi_worker_ids: list[int] = [] + self._mpi_rank_by_worker_id: dict[int, _MpiL3RankRuntime] = {} self._next_level_worker_id_count: int = 0 # Fallback ownership for private helpers used outside Worker.submit. # Normal orchestration-owned refs live in RunHandle._resources. @@ -2392,6 +3026,10 @@ def __init__( # among live handles). ``orch.allocate_domain`` adds entries here; # ``release()`` removes them and queues a deferred backend free. self._live_domains: dict[str, CommDomainHandle] = {} + self._live_global_domains: dict[str, GlobalCommDomainHandle] = {} + self._failed_global_domain_releases: dict[int, GlobalCommDomainHandle] = {} + self._global_node_domains: dict[int, _GlobalNodeDomainState] = {} + self._global_cluster_id = uuid.uuid4().hex # Monotonic per-Worker counter; mixed into IPC barrier filenames so # two concurrent allocations don't share a marker file. Wraps after # 2^64 allocations — far beyond any realistic Worker lifetime. @@ -2405,6 +3043,8 @@ def __init__( # down under an in-flight release. self._domain_free_mu = threading.Lock() self._domain_free_results: dict[int, BaseException | None] = {} + self._global_domain_free_mu = threading.Lock() + self._global_domain_free_results: dict[int, BaseException | None] = {} self._alloc_id_lock = threading.Lock() # Base HCCL/sim communicator is built lazily on the first # ``orch.allocate_domain`` call (see ``_ensure_comm_base``). We @@ -2508,6 +3148,78 @@ def add_remote_worker(self, spec: RemoteWorkerSpec) -> int: self._remote_worker_ids.append(worker_id) return worker_id + def add_mpirun_worker_group(self, spec: MpiL3GroupSpec) -> tuple[int, ...]: + """Register L3 workers that will be launched by one parent-owned ``mpirun``. + + The returned ids are ordinary NEXT_LEVEL worker ids for dispatch and + Global CommDomain membership. They must be used as a complete set for + the v1 MPI descriptor-exchange path; partial groups fall back to the + existing L4 descriptor broker. + """ + with self._hierarchical_start_cv: + if self._lifecycle is not _Lifecycle.NEW: + raise RuntimeError("Worker.add_mpirun_worker_group after init") + if self.level < 4: + raise TypeError("Worker.add_mpirun_worker_group: MPI L3 groups require a level >= 4 parent") + if not isinstance(spec, MpiL3GroupSpec): + raise TypeError("Worker.add_mpirun_worker_group expects a MpiL3GroupSpec") + for host in spec.connect_hosts: + self._validate_numeric_endpoint_host(host) + if spec.ready_host: + self._validate_numeric_endpoint_host(spec.ready_host) + seen_listeners: set[tuple[str, int]] = set() + for rank, listen_host in enumerate(spec.session_listen_hosts): + if self._is_wildcard_session_host(listen_host): + if not spec.allow_wildcard_session_bind: + raise ValueError( + "MpiL3GroupSpec wildcard session bind requires allow_wildcard_session_bind=True" + ) + else: + self._validate_numeric_endpoint_host(listen_host) + for port in (spec.command_port_base + rank, spec.health_port_base + rank): + key = (listen_host, int(port)) + if key in seen_listeners: + raise ValueError("MpiL3GroupSpec command/health ports overlap on the same listen host") + seen_listeners.add(key) + + group_id = uuid.uuid4().hex + ranks: list[_MpiL3RankRuntime] = [] + for rank, connect_host in enumerate(spec.connect_hosts): + worker_id = self._allocate_next_level_worker_id() + command_port = spec.command_port_base + rank + health_port = spec.health_port_base + rank + rank_spec = RemoteWorkerSpec( + endpoint=f"{connect_host}:{command_port}", + platform=spec.platform, + runtime=spec.runtime, + device_ids=spec.device_ids_by_rank[rank], + num_sub_workers=spec.num_sub_workers_by_rank[rank], + transport=spec.transport, + comm_profile=spec.comm_profile, + global_device_ranks=spec.global_device_ranks_by_rank[rank], + session_listen_host=spec.session_listen_hosts[rank], + allow_wildcard_session_bind=spec.allow_wildcard_session_bind, + ) + runtime = _MpiL3RankRuntime( + group_id=group_id, + rank=rank, + worker_id=worker_id, + spec=rank_spec, + command_host=connect_host, + command_port=command_port, + health_host=connect_host, + health_port=health_port, + ) + ranks.append(runtime) + self._mpi_worker_ids.append(worker_id) + self._mpi_rank_by_worker_id[worker_id] = runtime + group = _MpiL3GroupRuntime(group_id=group_id, spec=spec, ranks=tuple(ranks)) + self._mpi_l3_groups.append(group) + return tuple(rank.worker_id for rank in ranks) + + def _remote_like_worker_ids(self) -> set[int]: + return set(self._remote_worker_ids) | set(self._mpi_worker_ids) + @staticmethod def _parse_remote_endpoint(endpoint: str) -> tuple[str, int]: if endpoint.count(":") != 1: @@ -2646,6 +3358,165 @@ def _remote_dispatcher_entries_for_worker(self, worker_id: int) -> list[dict[str ) return entries + def _inner_registry_entries_for_spec(self, spec: RemoteWorkerSpec) -> list[dict[str, Any]]: + from .remote_l3_protocol import ( # noqa: PLC0415 + ChipCallableBlobLocation, + RemoteChipCallablePayload, + encode_remote_chip_callable_payload, + ) + + entries: list[dict[str, Any]] = [] + with self._registry_lock: + states = list(self._identity_registry.values()) + for state in states: + if state.target_namespace != "LOCAL_CHIP": + continue + if not isinstance(state.target, ChipCallable): + raise RuntimeError(f"inner chip hashid {state.hashid} does not carry a ChipCallable target") + descriptor = build_chip_callable_descriptor( + target=state.target, + platform=spec.platform, + runtime=spec.runtime, + ) + if descriptor != state.descriptor: + raise RuntimeError(f"inner chip hashid {state.hashid} was registered for a different platform/runtime") + blob = ctypes.string_at(int(state.target.buffer_ptr()), int(state.target.buffer_size())) + payload = encode_remote_chip_callable_payload( + RemoteChipCallablePayload( + descriptor_bytes=descriptor, + blob_location=ChipCallableBlobLocation.INLINE_BLOB, + blob_size=len(blob), + blob_sha256=hashlib.sha256(blob).digest(), + inline_blob=blob, + staged_blob_token=b"", + ) + ) + entries.append( + { + "hashid": state.digest.hex(), + "kind": "CHIP_CALLABLE", + "target_registry": "INNER_L3_WORKER", + "payload_version": 1, + "payload_hex": payload.hex(), + } + ) + return entries + + @staticmethod + def _validate_global_node_config( + *, + label: str, + platform: str, + device_ids: tuple[int, ...], + comm_profile: str, + global_device_ranks: tuple[int, ...], + ) -> None: + if comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"{label} comm_profile {comm_profile!r} is not supported") + if comm_profile == "a3-fabric-v1" and (not platform.startswith("a2a3") or platform.endswith("sim")): + raise ValueError(f"{label} comm_profile 'a3-fabric-v1' requires real A3 devices") + if global_device_ranks and len(global_device_ranks) != len(device_ids): + raise ValueError(f"{label} global_device_ranks must match device_ids length") + if any(rank < 0 for rank in global_device_ranks) or len(set(global_device_ranks)) != len(global_device_ranks): + raise ValueError(f"{label} global_device_ranks must be unique and non-negative") + + def _resolved_global_nodes(self) -> dict[int, _GlobalNodeRuntime]: + configs: list[tuple[int, tuple[int, ...], str, str, tuple[int, ...], bool]] = [] + for worker_id, spec in zip(self._remote_worker_ids, self._remote_worker_specs, strict=True): + configs.append( + ( + int(worker_id), + tuple(spec.device_ids), + spec.platform, + spec.comm_profile, + tuple(spec.global_device_ranks), + True, + ) + ) + for worker_id in self._mpi_worker_ids: + rank = self._mpi_rank_by_worker_id[int(worker_id)] + configs.append( + ( + int(worker_id), + tuple(rank.spec.device_ids), + rank.spec.platform, + rank.spec.comm_profile, + tuple(rank.spec.global_device_ranks), + True, + ) + ) + for worker_id, child in zip(self._next_level_worker_ids, self._next_level_workers, strict=True): + if child.level != 3: + continue + device_ids = tuple(int(device_id) for device_id in child._config.get("device_ids", ())) + platform = str(child._config.get("platform", "")) + comm_profile = str(child._config.get("comm_profile", "sim")) + global_device_ranks = tuple(int(rank) for rank in child._config.get("global_device_ranks", ())) + self._validate_global_node_config( + label=f"local L3 worker {worker_id}", + platform=platform, + device_ids=device_ids, + comm_profile=comm_profile, + global_device_ranks=global_device_ranks, + ) + configs.append( + ( + int(worker_id), + device_ids, + platform, + comm_profile, + global_device_ranks, + False, + ) + ) + configs.sort(key=lambda item: item[0]) + + explicit_ranks: set[int] = set() + for worker_id, _device_ids, _platform, _profile, ranks, _is_remote in configs: + overlap = explicit_ranks.intersection(ranks) + if overlap: + raise ValueError( + f"Global CommDomain worker {worker_id} duplicates global_device_ranks {sorted(overlap)}" + ) + explicit_ranks.update(ranks) + + used = set(explicit_ranks) + next_rank = 0 + resolved: dict[int, _GlobalNodeRuntime] = {} + node_count = len(configs) + for node_rank, (worker_id, device_ids, platform, profile, ranks, is_remote) in enumerate(configs): + self._validate_global_node_config( + label=f"{'remote' if is_remote else 'local'} L3 worker {worker_id}", + platform=platform, + device_ids=device_ids, + comm_profile=profile, + global_device_ranks=ranks, + ) + if not ranks: + assigned: list[int] = [] + for _device_id in device_ids: + while next_rank in used: + next_rank += 1 + assigned.append(next_rank) + used.add(next_rank) + next_rank += 1 + ranks = tuple(assigned) + resolved[worker_id] = _GlobalNodeRuntime( + worker_id=worker_id, + device_ids=device_ids, + platform=platform, + comm_profile=profile, + global_device_ranks=ranks, + node_rank=node_rank, + node_count=node_count, + cluster_id=self._global_cluster_id, + is_remote=is_remote, + ) + return resolved + + def _resolved_global_device_ranks(self) -> dict[int, tuple[int, ...]]: + return {worker_id: runtime.global_device_ranks for worker_id, runtime in self._resolved_global_nodes().items()} + def _build_remote_manifest( self, *, spec: RemoteWorkerSpec, worker_id: int, session_id: int, startup_remaining_s: float ) -> dict[str, Any]: @@ -2653,6 +3524,15 @@ def _build_remote_manifest( listen_host = spec.session_listen_host or ("127.0.0.1" if daemon_host == "localhost" else daemon_host) if self._is_wildcard_session_host(listen_host) and not spec.allow_wildcard_session_bind: raise ValueError("RemoteWorkerSpec wildcard session bind requires allow_wildcard_session_bind=True") + if worker_id in self._remote_like_worker_ids(): + runtime = self._resolved_global_nodes()[int(worker_id)] + node_rank = runtime.node_rank + node_count = runtime.node_count + global_device_ranks = runtime.global_device_ranks + else: + node_rank = 0 + node_count = 1 + global_device_ranks = spec.global_device_ranks or tuple(range(len(spec.device_ids))) return { "session_id": int(session_id), "parent_worker_level": int(self.level), @@ -2664,6 +3544,11 @@ def _build_remote_manifest( "num_sub_workers": int(spec.num_sub_workers), "heap_ring_size": self._config.get("remote_heap_ring_size", None), "transport": spec.transport, + "comm_profile": spec.comm_profile, + "cluster_id": self._global_cluster_id, + "node_rank": node_rank, + "node_count": node_count, + "global_device_ranks": list(global_device_ranks), # session_timeout_s bounds the runtime command socket; startup_remaining_s # bounds this session's slice of the single root startup budget. They are # distinct: the remote must not spend runtime-command time as startup time. @@ -2672,7 +3557,7 @@ def _build_remote_manifest( "listen_host": listen_host, "connect_host": daemon_host, "remote_task_dispatcher": self._remote_dispatcher_entries_for_worker(worker_id), - "inner_l3_worker": [], + "inner_l3_worker": self._inner_registry_entries_for_spec(spec), "feature_flags": [], } @@ -2723,6 +3608,253 @@ def _close_remote_sessions(self, sessions: list[_RemoteSession]) -> None: for session in reversed(sessions): self._close_remote_session(session) + @staticmethod + def _new_remote_session_id() -> int: + session_id = uuid.uuid4().int & ((1 << 63) - 1) + return session_id if session_id != 0 else 1 + + @staticmethod + def _mpirun_args_select_hosts(args: tuple[str, ...]) -> bool: + host_args = {"--host", "-host", "-H", "--hostfile", "-hostfile", "--machinefile", "-machinefile"} + return any(arg in host_args or arg.startswith("--host=") or arg.startswith("--hostfile=") for arg in args) + + @staticmethod + def _mpi_ready_path(ready_dir: str, rank: int) -> str: + return os.path.join(ready_dir, f"rank-{int(rank)}.json") + + def _wait_mpi_ready_files(self, group: _MpiL3GroupRuntime, deadline: float) -> dict[int, dict[str, Any]]: + if group.ready_dir is None: + raise RuntimeError("MPI L3 group has no ready directory") + pending = {rank.rank for rank in group.ranks} + ready: dict[int, dict[str, Any]] = {} + while pending: + if group.process is not None and group.process.poll() is not None: + raise RuntimeError( + f"MPI L3 group {group.group_id} exited before ranks became ready " + f"(status {group.process.returncode})" + ) + for rank in list(pending): + path = self._mpi_ready_path(group.ready_dir, rank) + try: + with open(path, encoding="utf-8") as f: + payload = json.load(f) + except FileNotFoundError: + continue + if not payload.get("ok", False): + raise RuntimeError(f"MPI L3 rank {rank} startup failed: {payload.get('error')}") + ready[rank] = payload + pending.remove(rank) + if pending: + remaining = self._remaining_until(deadline, "MPI L3 group ready") + time.sleep(min(0.05, remaining)) + return ready + + @staticmethod + def _open_mpi_ready_listener(host: str, port: int) -> socket.socket: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((host, int(port))) + sock.listen() + return sock + except BaseException: + sock.close() + raise + + def _wait_mpi_ready_tcp( + self, + group: _MpiL3GroupRuntime, + ready_sock: socket.socket, + ready_token: str, + deadline: float, + ) -> dict[int, dict[str, Any]]: + pending = {rank.rank for rank in group.ranks} + ready: dict[int, dict[str, Any]] = {} + while pending: + if group.process is not None and group.process.poll() is not None: + raise RuntimeError( + f"MPI L3 group {group.group_id} exited before ranks became ready " + f"(status {group.process.returncode})" + ) + ready_sock.settimeout(min(0.2, self._remaining_until(deadline, "MPI L3 TCP ready"))) + try: + conn, _addr = ready_sock.accept() + except TimeoutError: + continue + except socket.timeout: + continue + with conn: + payload = self._recv_remote_daemon_json(conn, deadline) + rank_value = int(payload.get("mpi_rank", -1)) + if payload.get("ready_token") != ready_token: + raise RuntimeError("MPI L3 rank published ready with an unexpected token") + if rank_value not in pending: + raise RuntimeError(f"MPI L3 rank published duplicate or unexpected ready rank={rank_value}") + if not payload.get("ok", False): + raise RuntimeError(f"MPI L3 rank {rank_value} startup failed: {payload.get('error')}") + ready[rank_value] = payload + pending.remove(rank_value) + return ready + + @staticmethod + def _close_mpirun_group(group: _MpiL3GroupRuntime, *, timeout_s: float) -> list[str]: + failures: list[str] = [] + proc = group.process + if proc is not None: + try: + if proc.poll() is None: + try: + proc.terminate() + except BaseException as exc: # noqa: BLE001 + failures.append(f"terminate: {exc}") + try: + proc.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + try: + proc.kill() + except BaseException as exc: # noqa: BLE001 + failures.append(f"kill: {exc}") + try: + proc.wait(timeout=timeout_s) + except BaseException as exc: # noqa: BLE001 + failures.append(f"wait after kill: {exc}") + except BaseException as exc: # noqa: BLE001 + failures.append(f"wait after terminate: {exc}") + finally: + group.process = None + try: + if group.ready_dir is not None: + shutil.rmtree(group.ready_dir) + except FileNotFoundError: + pass + except BaseException as exc: # noqa: BLE001 + failures.append(f"remove ready directory: {exc}") + finally: + group.ready_dir = None + group.manifest_path = None + return [f"MPI L3 group {group.group_id} cleanup {failure}" for failure in failures] + + def _close_mpirun_groups( + self, + *, + timeout_s: float = _ROLLBACK_GRACEFUL_TIMEOUT_S, + suppress_errors: bool = False, + ) -> None: + failures: list[str] = [] + for group in reversed(self._mpi_l3_groups): + failures.extend(self._close_mpirun_group(group, timeout_s=timeout_s)) + if failures: + sys.stderr.write("\n".join(f"[worker pid={os.getpid()}] WARN: {failure}" for failure in failures) + "\n") + sys.stderr.flush() + if not suppress_errors: + raise RuntimeError(failures[0]) + + def _activate_mpirun_worker_groups(self, deadline: float) -> None: + if not self._mpi_l3_groups: + return + session_timeout = self._remote_session_timeout_s() + assert self._worker is not None + for group in self._mpi_l3_groups: + ready_dir = tempfile.mkdtemp(prefix="simpler-mpirun-ready-") + manifest_path = os.path.join(ready_dir, "group.json") + group.ready_dir = ready_dir + group.manifest_path = manifest_path + ready_sock: socket.socket | None = None + ready_token = uuid.uuid4().hex + ready_host = group.spec.ready_host + ready_port = 0 + if ready_host: + ready_sock = self._open_mpi_ready_listener(ready_host, group.spec.ready_port) + ready_port = int(ready_sock.getsockname()[1]) + rank_manifests: list[dict[str, Any]] = [] + sessions: dict[int, _RemoteSession] = {} + startup_remaining_s = self._remaining_until(deadline, "MPI L3 group manifest") + for rank in group.ranks: + session_id = self._new_remote_session_id() + manifest = self._build_remote_manifest( + spec=rank.spec, + worker_id=rank.worker_id, + session_id=session_id, + startup_remaining_s=startup_remaining_s, + ) + manifest.update( + { + "command_port": rank.command_port, + "health_port": rank.health_port, + "mpi_group_id": group.group_id, + "mpi_rank": rank.rank, + "mpi_world_size": len(group.ranks), + "mpi_group_worker_ids": [item.worker_id for item in group.ranks], + "mpi_global_domain_exchange": True, + } + ) + rank_manifests.append(manifest) + sessions[rank.rank] = _RemoteSession( + worker_id=rank.worker_id, + session_id=session_id, + command_host=rank.command_host, + command_port=rank.command_port, + health_host=rank.health_host, + health_port=rank.health_port, + pid=0, + ) + group_manifest = { + "group_id": group.group_id, + "world_size": len(group.ranks), + "worker_ids": [rank.worker_id for rank in group.ranks], + "ready_dir": ready_dir, + "ready_host": ready_host, + "ready_port": ready_port, + "ready_token": ready_token, + "rank_manifests": rank_manifests, + } + with open(manifest_path, "w", encoding="utf-8") as f: + json.dump(group_manifest, f, sort_keys=True) + f.write("\n") + + cmd = [group.spec.mpirun_path, "-np", str(len(group.ranks))] + if not self._mpirun_args_select_hosts(group.spec.mpirun_args): + cmd.extend(["--host", ",".join(group.spec.hosts)]) + cmd.extend(group.spec.mpirun_args) + cmd.extend( + [ + group.spec.python_executable, + "-m", + "simpler.mpi_l3_session", + "--group-manifest", + manifest_path, + ] + ) + group.process = subprocess.Popen(cmd) + try: + if ready_sock is None: + ready = self._wait_mpi_ready_files(group, deadline) + else: + ready = self._wait_mpi_ready_tcp(group, ready_sock, ready_token, deadline) + finally: + if ready_sock is not None: + ready_sock.close() + for rank in group.ranks: + payload = ready[rank.rank] + if int(payload["command_port"]) != rank.command_port or int(payload["health_port"]) != rank.health_port: + raise RuntimeError(f"MPI L3 rank {rank.rank} published unexpected command/health ports") + session = sessions[rank.rank] + self._remote_sessions.append(session) + remaining = self._remaining_until(deadline, "MPI L3 endpoint attach") + self._worker.add_remote_l3_socket( + session.worker_id, + session.session_id, + rank.spec.comm_profile, + str(payload.get("command_host", rank.command_host)), + int(payload["command_port"]), + str(payload.get("health_host", rank.health_host)), + int(payload["health_port"]), + remaining, + session_timeout, + ) + if time.monotonic() >= deadline: + raise RuntimeError("MPI L3 activation: startup deadline exceeded after attach") + def _require_remote_worker_started(self, worker_id: int) -> None: """Argument + resource gate for the public remote-memory APIs. Admission (READY) is decided by the ``_operation_lease`` these APIs already hold — @@ -2732,8 +3864,10 @@ def _require_remote_worker_started(self, worker_id: int) -> None: instead of spuriously failing.""" if self.level < 4: raise TypeError("remote memory APIs require a level >= 4 parent Worker") - if int(worker_id) not in set(self._remote_worker_ids): - raise ValueError("remote memory APIs require a remote worker id returned by add_remote_worker") + if int(worker_id) not in self._remote_like_worker_ids(): + raise ValueError( + "remote memory APIs require a remote worker id returned by add_remote_worker or add_mpirun_worker_group" + ) if self._worker is None: raise RuntimeError("remote memory APIs require a started hierarchical Worker") @@ -2745,8 +3879,10 @@ def _require_remote_transport(self, worker_id: int) -> None: public lifecycle, so teardown keeps its capability without re-opening public admission. Public entrypoints validate READY separately via ``_require_remote_worker_started``.""" - if int(worker_id) not in set(self._remote_worker_ids): - raise ValueError("remote memory APIs require a remote worker id returned by add_remote_worker") + if int(worker_id) not in self._remote_like_worker_ids(): + raise ValueError( + "remote memory APIs require a remote worker id returned by add_remote_worker or add_mpirun_worker_group" + ) if self._worker is None: raise RuntimeError("remote memory APIs require a started hierarchical Worker") @@ -3376,14 +4512,14 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl raise TypeError("Worker.register: level 2 only supports ChipCallable targets") reg = _build_callable_registration(self, target, workers=workers) if isinstance(target, RemoteCallable): - if not self._remote_worker_specs: + remote_worker_ids = self._remote_like_worker_ids() + if not remote_worker_ids: raise RuntimeError("Worker.register(RemoteCallable): add at least one remote worker first") - remote_worker_ids = set(self._remote_worker_ids) for worker_id in reg.eligible_worker_ids: if worker_id not in remote_worker_ids: raise ValueError( "Worker.register(RemoteCallable): workers must name remote worker ids returned by " - "add_remote_worker" + "add_remote_worker or add_mpirun_worker_group" ) # Linearize against the startup epoch exactly like the local path: a # register that races an in-progress init() waits for it, then a @@ -4194,11 +5330,21 @@ def _eligible_target_need(self, namespace: str | None, eligible_worker_ids) -> s has_python_child = self._config.get("num_sub_workers", 0) > 0 or bool(self._next_level_workers) return None if has_python_child else "a SUB or next-level child" if namespace == "LOCAL_CHIP": - return None if bool(self._config.get("device_ids")) else "a chip device (device_ids)" + + def has_chip_target(worker: Worker) -> bool: + if worker._config.get("device_ids"): + return True + if any(spec.device_ids for spec in worker._remote_worker_specs): + return True + if any(rank.spec.device_ids for rank in worker._mpi_rank_by_worker_id.values()): + return True + return any(has_chip_target(child) for child in worker._next_level_workers) + + return None if has_chip_target(self) else "a local or remote chip device" if namespace == "REMOTE_TASK_DISPATCHER": - has_remote_workers = set(self._remote_worker_ids) + has_remote_workers = self._remote_like_worker_ids() ok = bool(has_remote_workers) and set(eligible_worker_ids) <= has_remote_workers - return None if ok else "its named remote worker(s) (add_remote_worker)" + return None if ok else "its named remote worker(s) (add_remote_worker/add_mpirun_worker_group)" return None def _validate_eligible_targets(self) -> None: @@ -4376,7 +5522,7 @@ def _init_hierarchical(self) -> None: # startup resource (mailbox shm, pre-fork _Worker mmap, child fork, # daemon socket) exists, so an invalid value fails without a # partially-built subtree to roll back. - if self._remote_worker_specs: + if self._remote_worker_specs or self._mpi_l3_groups: self._remote_session_timeout_s() # 1. Allocate sub-worker mailboxes (unified layout, MAILBOX_SIZE each). @@ -4468,7 +5614,7 @@ def _activate_remote_sessions(self, deadline: float) -> None: self._worker.add_remote_l3_socket( session.worker_id, session.session_id, - spec.transport, + spec.comm_profile, session.command_host, session.command_port, session.health_host, @@ -4492,6 +5638,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) deadline = self._startup_deadline + global_nodes = self._resolved_global_nodes() if self.level >= 4 else {} # Freeze the startup registry snapshot. init() already holds the epoch in # the INITIALIZING state, so a concurrent register/unregister is blocked @@ -4605,6 +5752,8 @@ def _setup(): # L3 child → L3's chip/sub grandchildren) and INIT_READY propagates up # only after the whole subtree is ready. for idx, inner_worker in enumerate(self._next_level_workers): + worker_id = self._next_level_worker_ids[idx] + global_node = global_nodes.get(worker_id) pid = os.fork() if pid == 0: buf = self._next_level_shms[idx].buf @@ -4636,7 +5785,12 @@ def _setup(inner=inner_worker): buf, f"next_level worker {idx}", _setup, - lambda tables, b=buf, inner=inner_worker: _child_worker_loop(b, *tables, inner), + lambda tables, b=buf, inner=inner_worker, node=global_node: _child_worker_loop( + b, + *tables, + inner, + node, + ), make_group_leader=self._is_startup_root, ) else: @@ -4653,6 +5807,7 @@ def _setup(inner=inner_worker): # L3 sessions: opening starts the remote subtree and registering spawns # the RemoteL3Endpoint health thread, so both must follow every local # fork. Each remote consumes this process's remaining startup budget. + self._activate_mpirun_worker_groups(deadline) self._activate_remote_sessions(deadline) # _Worker was constructed in _init_hierarchical (pre-fork) so children @@ -4934,6 +6089,7 @@ def _cleanup_partial_init(self) -> None: except BaseException: # noqa: BLE001 pass self._close_remote_sessions(remote_sessions) + self._close_mpirun_groups(suppress_errors=True) if self._chip_worker is not None: try: self._chip_worker.finalize() @@ -5405,8 +6561,13 @@ def _retire_run_domains(self, resources: _RunResources) -> None: resources.retired = True stragglers = list(resources.pending_release_domains) resources.pending_release_domains.clear() + global_stragglers = list(resources.pending_release_global_domains) + resources.pending_release_global_domains.clear() + resources.live_global_domains.clear() for handle in stragglers: self._free_domain_after_fence(handle) + for handle in global_stragglers: + self._free_global_domain_after_fence(handle) def _free_domain_after_fence(self, handle: CommDomainHandle) -> None: """Back-end free for a handle whose owning run has retired. @@ -5578,6 +6739,804 @@ def dispatch(chip_idx: int) -> None: f"{len(errors)}/{len(workers)} chips; first error chip={first[0]}: {first[1]}" ) + @staticmethod + def _global_domain_command_identity(command: GlobalDomainCommand) -> tuple[Any, ...]: + return ( + command.domain_id, + command.generation, + command.name, + command.profile, + command.window_size, + command.members, + command.buffers, + ) + + @staticmethod + def _global_domain_provenance_id(domain_id: int) -> int: + # Local CommDomain allocation ids are positive. Keep remote Global + # CommDomains in a disjoint namespace while reusing the same exact-pointer + # provenance table that protects child-memory task submissions. + return -int(domain_id) + + def _global_local_members( + self, command: GlobalDomainCommand, node_worker_id: int + ) -> tuple[GlobalDomainMember, ...]: + members = tuple(member for member in command.members if member.node_worker_id == node_worker_id) + if not members: + raise ValueError(f"Global CommDomain has no members on node worker {node_worker_id}") + local_count = len(self._config.get("device_ids", [])) + for member in members: + if member.local_worker_id < 0 or member.local_worker_id >= local_count: + raise ValueError( + f"Global CommDomain local worker {member.local_worker_id} is outside [0, {local_count})" + ) + return members + + def _prepare_global_domain_node( + self, command: GlobalDomainCommand, node_worker_id: int + ) -> tuple[GlobalDomainDescriptor, ...]: + if self.level != 3 or self._worker is None: + raise RuntimeError("Global CommDomain node prepare requires a ready L3 Worker") + prior = self._global_node_domains.get(command.domain_id) + if prior is not None: + if self._global_domain_command_identity(prior.command) != self._global_domain_command_identity(command): + raise RuntimeError("Global CommDomain prepare conflicts with a live domain") + return tuple(prior.descriptors[rank] for rank in sorted(prior.descriptors)) + + state = _GlobalNodeDomainState(command=command) + self._global_node_domains[command.domain_id] = state + local_members = self._global_local_members(command, node_worker_id) + capacity = max(LOCAL_PREPARE_REQUEST.size, LOCAL_PREPARE_REPLY.size + GLOBAL_DOMAIN_DESCRIPTOR_BYTES) + try: + for member in local_members: + payload = bytearray(capacity) + LOCAL_PREPARE_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + member.domain_rank, + len(command.members), + GLOBAL_DOMAIN_PROFILE_IDS[command.profile], + command.window_size, + ) + state.prepared_domain_ranks.add(member.domain_rank) + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_PREPARE, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + fields = LOCAL_PREPARE_REPLY.unpack_from(reply, 0) + magic, version, domain_id, generation, local_base, mapping_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="prepare reply") + if domain_id != command.domain_id or generation != command.generation: + raise RuntimeError("Global CommDomain prepare reply identity mismatch") + start = LOCAL_PREPARE_REPLY.size + descriptor = GlobalDomainDescriptor.decode(reply[start : start + GLOBAL_DOMAIN_DESCRIPTOR_BYTES]) + if descriptor.domain_rank != member.domain_rank: + raise RuntimeError("Global CommDomain prepare reply rank mismatch") + state.descriptors[member.domain_rank] = descriptor + state.local_window_bases[member.local_worker_id] = int(local_base) + state.mapping_sizes[member.local_worker_id] = int(mapping_size) + return tuple(state.descriptors[rank] for rank in sorted(state.descriptors)) + except BaseException: + self._release_global_domain_node( + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + raise + + def _import_global_domain_node(self, command: GlobalDomainCommand, node_worker_id: int) -> None: + if self.level != 3 or self._worker is None: + raise RuntimeError("Global CommDomain node import requires a ready L3 Worker") + state = self._global_node_domains.get(command.domain_id) + if state is None or state.command.generation != command.generation: + raise RuntimeError("Global CommDomain import requires a matching prepared domain") + if self._global_domain_command_identity(state.command) != self._global_domain_command_identity(command): + raise RuntimeError("Global CommDomain import command conflicts with prepare") + validate_descriptor_table( + command.descriptors, + rank_count=len(command.members), + profile=command.profile, + ) + local_members = self._global_local_members(command, node_worker_id) + descriptor_bytes = b"".join(descriptor.encode() for descriptor in command.descriptors) + request_size = LOCAL_IMPORT_REQUEST.size + len(descriptor_bytes) + capacity = max(request_size, LOCAL_IMPORT_REPLY.size) + for member in local_members: + payload = bytearray(capacity) + LOCAL_IMPORT_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + len(command.descriptors), + ) + payload[LOCAL_IMPORT_REQUEST.size : request_size] = descriptor_bytes + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_IMPORT, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + fields = LOCAL_IMPORT_REPLY.unpack_from(reply, 0) + magic, version, domain_id, generation, device_ctx, local_base, mapping_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="import reply") + if domain_id != command.domain_id or generation != command.generation: + raise RuntimeError("Global CommDomain import reply identity mismatch") + if mapping_size != command.descriptors[member.domain_rank].mapping_size: + raise RuntimeError("Global CommDomain import reply mapping size mismatch") + offset = 0 + buffer_ptrs: dict[str, int] = {} + for buffer in command.buffers: + buffer_ptrs[buffer.name] = int(local_base) + offset + offset += buffer.nbytes + state.contexts[member.local_worker_id] = ChipDomainContext( + name=command.name, + domain_rank=member.domain_rank, + domain_size=len(command.members), + device_ctx=int(device_ctx), + local_window_base=int(local_base), + actual_window_size=int(mapping_size), + buffer_ptrs=buffer_ptrs, + ) + provenance_id = self._global_domain_provenance_id(command.domain_id) + with self._child_prov_lock: + self._child_prov_record_domain( + member.local_worker_id, + int(local_base), + provenance_id, + int(mapping_size), + ) + for buffer in command.buffers: + self._child_prov_record_domain( + member.local_worker_id, + buffer_ptrs[buffer.name], + provenance_id, + int(buffer.nbytes), + ) + state.command = command + state.phase = GlobalDomainPhase.IMPORT + state.view = GlobalCommDomainView( + name=command.name, + members=command.members, + contexts=state.contexts, + domain_id=command.domain_id, + generation=command.generation, + mapping_size=command.descriptors[0].mapping_size, + ) + + def _commit_global_domain_node(self, command: GlobalDomainCommand) -> None: + state = self._global_node_domains.get(command.domain_id) + if state is None or state.command.generation != command.generation: + raise RuntimeError("Global CommDomain commit requires a matching imported domain") + if state.phase is not GlobalDomainPhase.IMPORT or state.view is None: + raise RuntimeError("Global CommDomain commit requires IMPORT completion") + if ( + self._global_domain_command_identity(state.command) != self._global_domain_command_identity(command) + or state.command.descriptors != command.descriptors + ): + raise RuntimeError("Global CommDomain commit command conflicts with IMPORT") + state.phase = GlobalDomainPhase.COMMIT + state.view._committed = True # noqa: SLF001 -- session owns the transaction + + def _release_global_domain_node( + self, command: GlobalDomainReleaseCommand, *, suppress_errors: bool = False + ) -> None: + state = self._global_node_domains.get(command.domain_id) + if state is None: + return + if state.command.generation != command.generation: + raise RuntimeError("Global CommDomain release generation mismatch") + with self._child_prov_lock: + self._child_prov_drop_domain(self._global_domain_provenance_id(command.domain_id)) + if self._worker is None: + return + errors: list[BaseException] = [] + local_members = tuple( + member for member in state.command.members if member.domain_rank in state.prepared_domain_ranks + ) + for member in local_members: + payload = bytearray(LOCAL_RELEASE_REQUEST.size) + LOCAL_RELEASE_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + ) + try: + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_RELEASE, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if not errors: + self._global_node_domains.pop(command.domain_id, None) + if errors and not suppress_errors: + raise RuntimeError(f"Global CommDomain node release failed: {errors[0]}") from errors[0] + + def _release_all_global_domain_nodes(self) -> None: + for state in list(self._global_node_domains.values())[::-1]: + try: + self._release_global_domain_node( + GlobalDomainReleaseCommand(state.command.domain_id, state.command.generation) + ) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._release_all_global_domain_nodes: domain_id={state.command.domain_id} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _get_global_domain(self, domain_id: int) -> GlobalCommDomainView: + state = self._global_node_domains.get(int(domain_id)) + if state is None or state.phase is not GlobalDomainPhase.COMMIT or state.view is None: + raise KeyError(f"Global CommDomain {domain_id} is not committed on this L3 node") + return state.view + + def _copy_global_domain_node(self, command: GlobalDomainCopyCommand, *, copy_to_device: bool) -> bytes: + state = self._global_node_domains.get(command.domain_id) + if ( + state is None + or state.command.generation != command.generation + or state.phase is not GlobalDomainPhase.COMMIT + ): + raise RuntimeError("Global CommDomain copy requires a committed live domain") + if command.domain_rank >= len(state.command.members): + raise ValueError("Global CommDomain copy rank is out of range") + member = state.command.members[command.domain_rank] + if member.local_worker_id not in state.contexts: + raise RuntimeError("Global CommDomain copy rank is not local to this L3 node") + request_size = LOCAL_COPY_REQUEST.size + (command.nbytes if copy_to_device else 0) + reply_size = LOCAL_COPY_REPLY.size + (command.nbytes if not copy_to_device else 0) + payload = bytearray(max(request_size, reply_size)) + LOCAL_COPY_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + command.offset, + command.nbytes, + ) + if copy_to_device: + payload[LOCAL_COPY_REQUEST.size : request_size] = command.data + assert self._worker is not None + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_COPY_TO if copy_to_device else CTRL_GLOBAL_DOMAIN_COPY_FROM, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + magic, version, domain_id, generation, nbytes = LOCAL_COPY_REPLY.unpack_from(reply, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="copy reply") + if domain_id != command.domain_id or generation != command.generation or nbytes != command.nbytes: + raise RuntimeError("Global CommDomain copy reply mismatch") + if copy_to_device: + return b"" + return reply[LOCAL_COPY_REPLY.size : LOCAL_COPY_REPLY.size + command.nbytes] + + @staticmethod + def _local_global_domain_response_capacity(control_name: int, payload: bytes) -> int: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + control = ControlName(control_name) + if control is ControlName.COMM_INIT: + return struct.calcsize(" bytes: + if self._worker is None: + raise RuntimeError("Global CommDomain control requires a ready hierarchical Worker") + if worker_id not in self._next_level_worker_ids: + raise ValueError(f"Global CommDomain worker {worker_id} is not a local L3 worker") + response_capacity = self._local_global_domain_response_capacity(control_name, payload) + capacity = max(len(payload), response_capacity) + staged = bytearray(_LOCAL_GLOBAL_CONTROL_HEADER.size + capacity) + _LOCAL_GLOBAL_CONTROL_HEADER.pack_into(staged, 0, int(control_name), len(payload), 0) + start = _LOCAL_GLOBAL_CONTROL_HEADER.size + staged[start : start + len(payload)] = payload + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + int(worker_id), + _CTRL_GLOBAL_DOMAIN_NODE, + staged, + _PY_CONTROL_TIMEOUT_S, + ) + ) + reply_control, reply_request_size, response_size = _LOCAL_GLOBAL_CONTROL_HEADER.unpack_from(reply, 0) + if reply_control != int(control_name) or reply_request_size != len(payload) or response_size > capacity: + raise RuntimeError("local Global CommDomain control reply is invalid") + return reply[start : start + response_size] + + def _global_domain_control(self, worker_id: int, control_name: int, payload: bytes) -> bytes: + if self._worker is None: + raise RuntimeError("Global CommDomain control requires a ready hierarchical Worker") + if worker_id in self._remote_like_worker_ids(): + return bytes(self._worker.remote_domain_control(int(worker_id), int(control_name), bytes(payload))) + if worker_id in self._next_level_worker_ids: + return self._local_global_domain_control(worker_id, control_name, payload) + raise ValueError(f"Global CommDomain worker {worker_id} is not a registered L3 worker") + + def _global_domain_control_many( + self, + worker_ids: tuple[int, ...], + control_name: int, + payload: bytes, + *, + mpi_group: _MpiL3GroupRuntime, + ) -> dict[int, bytes]: + replies: dict[int, bytes] = {} + errors: list[tuple[int, BaseException]] = [] + completed = threading.Condition() + + def _send(worker_id: int) -> None: + try: + reply = self._global_domain_control(worker_id, control_name, payload) + except BaseException as exc: # noqa: BLE001 + with completed: + errors.append((worker_id, exc)) + completed.notify_all() + return + with completed: + replies[worker_id] = reply + completed.notify_all() + + threads = [ + (int(worker_id), threading.Thread(target=_send, args=(int(worker_id),), daemon=True)) + for worker_id in worker_ids + ] + for _worker_id, thread in threads: + thread.start() + deadline = time.monotonic() + self._py_control_timeout_s + with completed: + while len(replies) + len(errors) < len(worker_ids) and not errors: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + completed.wait(timeout=remaining) + + pending = [worker_id for worker_id, thread in threads if thread.is_alive()] + if pending: + cleanup_failures = self._close_mpirun_group( + mpi_group, + timeout_s=min(1.0, self._py_control_timeout_s), + ) + if cleanup_failures: + sys.stderr.write( + "\n".join(f"[worker pid={os.getpid()}] WARN: {failure}" for failure in cleanup_failures) + "\n" + ) + sys.stderr.flush() + for _worker_id, thread in threads: + thread.join(timeout=1.0) + if errors: + worker_id, exc = errors[0] + raise RuntimeError(f"Global CommDomain control fanout failed on worker {worker_id}: {exc}") from exc + if pending: + raise TimeoutError(f"Global CommDomain control fanout timed out on workers {pending}") + missing = sorted(set(worker_ids) - replies.keys()) + if missing: + raise RuntimeError(f"Global CommDomain control fanout returned no reply for workers {missing}") + return replies + + def _mpi_group_for_involved_nodes(self, involved_nodes: tuple[int, ...]) -> _MpiL3GroupRuntime | None: + involved = set(int(worker_id) for worker_id in involved_nodes) + for group in self._mpi_l3_groups: + group_workers = {rank.worker_id for rank in group.ranks} + if involved == group_workers: + return group + return None + + @staticmethod + def _descriptor_table_from_mpi_replies( + replies: dict[int, bytes], + *, + rank_count: int, + profile: str, + ) -> tuple[GlobalDomainDescriptor, ...]: + tables = tuple(decode_descriptor_table(payload) for payload in replies.values()) + if not tables: + raise RuntimeError("MPI Global CommDomain prepare returned no descriptor tables") + first = tables[0] + validate_descriptor_table(first, rank_count=rank_count, profile=profile) + for table in tables[1:]: + validate_descriptor_table(table, rank_count=rank_count, profile=profile) + if table != first: + raise RuntimeError("MPI Global CommDomain prepare returned inconsistent descriptor tables") + return first + + def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prepare/import/commit rollback stay ordered + self, + *, + name: str, + members: tuple[tuple[int, int], ...], + window_size: int, + buffers: list[CommBufferSpec], + retain_after_run: bool, + ) -> GlobalCommDomainHandle: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + if self.level < 4 or self._worker is None: + raise RuntimeError("allocate_global_domain requires a ready L4+ Worker") + resources = self._building_run_resources + if resources is None: + raise RuntimeError("allocate_global_domain is only valid while a run's graph is being built") + if not name: + raise ValueError("allocate_global_domain: name must be non-empty") + if name in self._live_global_domains: + raise ValueError(f"allocate_global_domain: domain {name!r} is already live") + if not members or len(members) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("allocate_global_domain: members must contain between 1 and 64 devices") + if len(set(members)) != len(members): + raise ValueError("allocate_global_domain: members contain duplicate node/local devices") + if window_size <= 0: + raise ValueError("allocate_global_domain: window_size must be positive") + if len({buffer.name for buffer in buffers}) != len(buffers): + raise ValueError("allocate_global_domain: buffer names must be unique") + if any(not buffer.name or int(buffer.nbytes) <= 0 for buffer in buffers): + raise ValueError("allocate_global_domain: buffers require a name and positive nbytes") + if sum(int(buffer.nbytes) for buffer in buffers) > window_size: + raise ValueError("allocate_global_domain: buffers exceed window_size") + + nodes = self._resolved_global_nodes() + profiles: set[str] = set() + domain_members: list[GlobalDomainMember] = [] + for domain_rank, (node_worker_id, local_worker_id) in enumerate(members): + node = nodes.get(int(node_worker_id)) + if node is None: + raise ValueError(f"allocate_global_domain: worker {node_worker_id} is not a registered L3") + if local_worker_id < 0 or local_worker_id >= len(node.device_ids): + raise ValueError( + f"allocate_global_domain: local worker {local_worker_id} is outside " + f"worker {node_worker_id}'s device list" + ) + profiles.add(node.comm_profile) + domain_members.append( + GlobalDomainMember( + node_worker_id=int(node_worker_id), + local_worker_id=int(local_worker_id), + global_device_rank=node.global_device_ranks[int(local_worker_id)], + domain_rank=domain_rank, + ) + ) + if len(profiles) != 1: + raise ValueError("allocate_global_domain: all participating nodes must use the same comm_profile") + profile = next(iter(profiles)) + global_buffers = tuple(GlobalDomainBuffer(buffer.name, int(buffer.nbytes)) for buffer in buffers) + domain_members_tuple = tuple(domain_members) + involved_nodes = tuple(dict.fromkeys(member.node_worker_id for member in domain_members_tuple)) + for node_worker_id in involved_nodes: + node = nodes[node_worker_id] + resolve_global_comm_capability( + platform=node.platform, + profile=node.comm_profile, + local_device_count=len(node.device_ids), + ) + topology_bytes = repr( + ( + self._global_cluster_id, + profile, + tuple( + ( + member.node_worker_id, + member.local_worker_id, + member.global_device_rank, + member.domain_rank, + ) + for member in domain_members_tuple + ), + ) + ).encode() + topology_hash = hashlib.sha256(topology_bytes).hexdigest() + with self._alloc_id_lock: + self._next_alloc_id += 1 + domain_id = self._next_alloc_id + generation = 1 + base_command = GlobalDomainCommand( + phase=GlobalDomainPhase.PREPARE_EXPORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + ) + + prepared_nodes: list[int] = [] + try: + for node_worker_id in involved_nodes: + node = nodes[node_worker_id] + init = GlobalCommInitCommand( + cluster_id=self._global_cluster_id, + topology_hash=topology_hash, + profile=profile, + node_rank=node.node_rank, + node_count=node.node_count, + members=domain_members_tuple, + ) + result = decode_comm_init_result( + self._global_domain_control(node_worker_id, ControlName.COMM_INIT, encode_comm_init(init)) + ) + if ( + result.profile != profile + or result.max_ranks < len(domain_members_tuple) + or result.descriptor_bytes != GLOBAL_DOMAIN_DESCRIPTOR_BYTES + or result.local_device_count != len(node.device_ids) + ): + raise RuntimeError(f"Global CommDomain COMM_INIT capability mismatch on node {node_worker_id}") + + mpi_group = self._mpi_group_for_involved_nodes(involved_nodes) + if mpi_group is not None: + prepared_nodes.extend(involved_nodes) + replies = self._global_domain_control_many( + involved_nodes, + ControlName.ALLOC_DOMAIN, + encode_domain_command(base_command), + mpi_group=mpi_group, + ) + descriptors = self._descriptor_table_from_mpi_replies( + replies, + rank_count=len(domain_members_tuple), + profile=profile, + ) + else: + descriptor_by_rank: dict[int, GlobalDomainDescriptor] = {} + for node_worker_id in involved_nodes: + prepared_nodes.append(node_worker_id) + reply = self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(base_command), + ) + for descriptor in decode_descriptor_table(reply): + if descriptor.domain_rank in descriptor_by_rank: + raise RuntimeError("Global CommDomain prepare returned a duplicate rank") + descriptor_by_rank[descriptor.domain_rank] = descriptor + descriptors = tuple(descriptor_by_rank[rank] for rank in range(len(domain_members_tuple))) + validate_descriptor_table(descriptors, rank_count=len(domain_members_tuple), profile=profile) + if descriptors[0].mapping_size < window_size: + raise RuntimeError("Global CommDomain backend mapped less than the requested window size") + + commit_command = GlobalDomainCommand( + phase=GlobalDomainPhase.COMMIT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + descriptors=descriptors, + ) + if mpi_group is not None: + self._global_domain_control_many( + involved_nodes, + ControlName.ALLOC_DOMAIN, + encode_domain_command(commit_command), + mpi_group=mpi_group, + ) + else: + import_command = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + descriptors=descriptors, + ) + for node_worker_id in involved_nodes: + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(import_command), + ) + for node_worker_id in involved_nodes: + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(commit_command), + ) + except BaseException: + abort_command = GlobalDomainCommand( + phase=GlobalDomainPhase.ABORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + ) + for node_worker_id in prepared_nodes: + with contextlib.suppress(BaseException): + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(abort_command), + ) + raise + + handle = GlobalCommDomainHandle( + name=name, + members=domain_members_tuple, + buffers=global_buffers, + domain_id=domain_id, + generation=generation, + mapping_size=descriptors[0].mapping_size, + retain_after_run=retain_after_run, + _release_fn=self._release_global_domain_handle, + ) + self._live_global_domains[name] = handle + resources.live_global_domains[name] = handle + return handle + + def _release_global_domain_handle(self, handle: GlobalCommDomainHandle) -> None: + if self._worker is None: + return + resources = self._building_run_resources + if resources is not None: + with resources.domain_lock: + if resources.live_global_domains.get(handle.name) is handle: + resources.live_global_domains.pop(handle.name) + if self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + if not resources.retired: + resources.pending_release_global_domains.append(handle) + return + elif self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + self._free_global_domain_after_fence(handle) + + def _free_global_domain_after_fence(self, handle: GlobalCommDomainHandle) -> None: + if handle.freed: + return + try: + self._release_global_domain_now(handle) + handle._freed = True # noqa: SLF001 -- runtime owns this transition + self._failed_global_domain_releases.pop(handle.domain_id, None) + except Exception: + self._failed_global_domain_releases[handle.domain_id] = handle + raise + + def _release_global_domain_now(self, handle: GlobalCommDomainHandle) -> None: + with self._global_domain_free_mu: + if handle.domain_id in self._global_domain_free_results: + failure = self._global_domain_free_results[handle.domain_id] + if failure is not None: + raise failure + return + try: + self._release_global_domain_claimed(handle) + except BaseException as exc: + self._global_domain_free_results[handle.domain_id] = exc + raise + self._global_domain_free_results[handle.domain_id] = None + + def _release_global_domain_claimed(self, handle: GlobalCommDomainHandle) -> None: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + command = encode_release_command(GlobalDomainReleaseCommand(handle.domain_id, handle.generation)) + errors: list[BaseException] = [] + for node_worker_id in dict.fromkeys(member.node_worker_id for member in handle.members): + try: + self._global_domain_control(node_worker_id, ControlName.RELEASE_DOMAIN, command) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if errors: + raise RuntimeError(f"Global CommDomain release failed: {errors[0]}") from errors[0] + if self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + + def _execute_pending_global_domain_releases(self, resources: _RunResources) -> None: + pending = list(resources.pending_release_global_domains) + resources.pending_release_global_domains.clear() + for handle in pending: + try: + self._free_global_domain_after_fence(handle) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._execute_pending_global_domain_releases: {handle.name!r} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _release_all_live_global_domains( + self, + resources: _RunResources | None = None, + *, + include_retained: bool = True, + ) -> None: + live_domains = self._live_global_domains if resources is None else resources.live_global_domains + for handle in list(live_domains.values())[::-1]: + if handle.retain_after_run and not include_retained: + continue + try: + handle._released = True # noqa: SLF001 -- runtime owns this transition + self._free_global_domain_after_fence(handle) + if live_domains.get(handle.name) is handle: + live_domains.pop(handle.name) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._release_all_live_global_domains: {handle.name!r} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _copy_to_global_domain( + self, handle: GlobalCommDomainHandle, domain_rank: int, data: bytes, offset: int + ) -> None: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + payload = bytes(data) + member = handle.member(domain_rank) + command = GlobalDomainCopyCommand( + domain_id=handle.domain_id, + generation=handle.generation, + domain_rank=int(domain_rank), + offset=int(offset), + nbytes=len(payload), + data=payload, + ) + self._global_domain_control( + member.node_worker_id, + ControlName.COPY_TO_DOMAIN, + encode_copy_command(command, include_data=True), + ) + + def _copy_from_global_domain( + self, handle: GlobalCommDomainHandle, domain_rank: int, nbytes: int, offset: int + ) -> bytes: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + member = handle.member(domain_rank) + command = GlobalDomainCopyCommand( + domain_id=handle.domain_id, + generation=handle.generation, + domain_rank=int(domain_rank), + offset=int(offset), + nbytes=int(nbytes), + ) + reply = self._global_domain_control( + member.node_worker_id, + ControlName.COPY_FROM_DOMAIN, + encode_copy_command(command, include_data=False), + ) + return decode_copy_result(reply) + def _release_all_live_domains(self, resources: _RunResources | None = None) -> None: """Best-effort release of every still-live domain handle (LIFO). @@ -5750,7 +7709,7 @@ def _require_local_next_level_target(self, worker_id: int, *, api: str) -> None: C++ target check only rejects unregistered ids, so a registered remote worker slips through — this guards that hole. """ - if worker_id in set(self._remote_worker_ids): + if worker_id in self._remote_like_worker_ids(): raise ValueError( f"orch.{api}: worker {worker_id} is a remote NEXT_LEVEL worker; a local callable " f"must target a local child (remote workers only run RemoteCallable dispatches)" @@ -5891,7 +7850,7 @@ def _create_host_buffer_locked(self, nbytes: int) -> HostBuffer: # and sub alike, via _broadcast_host_control). Only a truly childless L3 # has nowhere to attach it. if not self._chip_shms and not self._sub_shms: - raise RuntimeError( + raise _NoHostBufferChildrenError( "create_host_buffer requires at least one forked chip or sub child (this Worker has none)" ) assert self._worker is not None @@ -5961,7 +7920,7 @@ def _create_host_buffer_locked(self, nbytes: int) -> HostBuffer: buf_view = shm.buf assert buf_view is not None - return HostBuffer(token=token, data_ptr=data_ptr, nbytes=nbytes, buffer=buf_view) + return HostBuffer(token=token, data_ptr=data_ptr, nbytes=nbytes, buffer=buf_view, shm_name=shm.name) def free_host_buffer(self, handle: HostBuffer) -> None: """Release a born-shared buffer created by ``create_host_buffer``. @@ -6288,6 +8247,14 @@ def _step(fn) -> None: _step(lambda: self._cleanup_l3_l2_regions(resources)) finally: resources.l3_l2_orch_comm_host_buffers.clear() + _step(lambda: self._execute_pending_global_domain_releases(resources)) + if resources.live_global_domains: + _step( + lambda: self._release_all_live_global_domains( + resources, + include_retained=False, + ) + ) _step(lambda: self._execute_pending_domain_releases(resources)) if resources.live_domains: _step(lambda: self._release_all_live_domains(resources)) @@ -6369,8 +8336,11 @@ def _has_live_resources(self) -> bool: self._has_native_tree() or bool(self._sub_pids or self._chip_pids or self._next_level_pids) or bool(self._sub_shms or self._chip_shms or self._next_level_shms) + or any(group.process is not None or group.ready_dir is not None for group in self._mpi_l3_groups) or bool(self._live_l3_l2_regions) or bool(self._live_domains) + or bool(self._live_global_domains or self._failed_global_domain_releases) + or bool(self._global_node_domains) or bool(self._host_buf_registry) or bool(self._pending_remote_buffer_frees or self._pending_remote_import_releases) ) @@ -6387,10 +8357,19 @@ def _describe_live_resources(self) -> str: n_shms = len(self._sub_shms) + len(self._chip_shms) + len(self._next_level_shms) if n_shms: parts.append(f"{n_shms} child shm(s)") + n_mpi = sum(1 for group in self._mpi_l3_groups if group.process is not None or group.ready_dir is not None) + if n_mpi: + parts.append(f"{n_mpi} mpirun group(s)") if self._live_l3_l2_regions: parts.append(f"{len(self._live_l3_l2_regions)} L3-L2 region(s)") if self._live_domains: parts.append(f"{len(self._live_domains)} comm domain(s)") + if self._live_global_domains or self._failed_global_domain_releases: + live_global_ids = {handle.domain_id for handle in self._live_global_domains.values()} + live_global_ids.update(self._failed_global_domain_releases) + parts.append(f"{len(live_global_ids)} global comm domain(s)") + if self._global_node_domains: + parts.append(f"{len(self._global_node_domains)} imported global comm domain(s)") if self._host_buf_registry: parts.append(f"{len(self._host_buf_registry)} host buffer(s)") n_remote = len(self._pending_remote_buffer_frees) + len(self._pending_remote_import_releases) @@ -6624,8 +8603,12 @@ def _broadcast_child_shutdown(shms: list[SharedMemory]) -> None: raise errors[0] @staticmethod - def _reap_child_groups( # noqa: PLR0912 -- interleaved reap across groups / bounded poll / conditional shm-free - groups: list[tuple[list[SharedMemory], list[int]]], deadline: float + def _reap_child_groups( # noqa: PLR0912, PLR0915 -- interleaved graceful/hard reap across groups + groups: list[tuple[list[SharedMemory], list[int]]], + deadline: float, + *, + kill_survivors: bool = False, + kill_process_groups: bool = False, ) -> None: """Reap + free every child across ALL groups within one shared deadline. @@ -6634,11 +8617,12 @@ def _reap_child_groups( # noqa: PLR0912 -- interleaved reap across groups / bou every group each round, so a child wedged in one group never starves the reap of healthy children in another (the serial-per-group variant let the first stuck group burn the whole budget and left later groups as - one-poll survivors). ``pids[i]`` pairs with ``shms[i]``; a shm is freed + one-poll survivors). When ``kill_survivors`` is true, children that + exhaust the graceful budget are SIGKILLed and given one additional + bounded reap window. ``pids[i]`` pairs with ``shms[i]``; a shm is freed ONLY once its pid is reaped (freeing a live child's mailbox is a - use-after-free), so a survivor keeps BOTH. Teardown is terminal — a - survivor LEAKS and is reported as an error so close() never returns - success while a child is alive; an abnormal exit (signal / non-zero code) + use-after-free), so an unreaped survivor keeps BOTH and is reported as + an error. An unexpected signal / non-zero exit during the graceful phase is likewise reported. The first error is raised after every child is attempted. """ @@ -6674,6 +8658,45 @@ def _reap_child_groups( # noqa: PLR0912 -- interleaved reap across groups / bou time.sleep(_STARTUP_POLL_INTERVAL_S) else: break + if pending and kill_survivors: + for g, i in pending: + _shms, pids = groups[g] + if kill_process_groups: + with contextlib.suppress(ProcessLookupError, OSError): + os.killpg(pids[i], signal.SIGKILL) + try: + os.kill(pids[i], signal.SIGKILL) + except ProcessLookupError: + reaped.add((g, i)) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + # A SIGKILLed child should become waitable immediately, but keep + # this poll bounded too: an uninterruptible (D-state) child must + # still surface as a survivor rather than pin close() forever. + kill_deadline = time.monotonic() + _FORCED_REAP_TIMEOUT_S + pending = [item for item in pending if item not in reaped] + while pending: + still = [] + for g, i in pending: + _shms, pids = groups[g] + try: + wpid, _status = os.waitpid(pids[i], os.WNOHANG) + except ChildProcessError: + reaped.add((g, i)) + continue + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + continue + if wpid != 0: + reaped.add((g, i)) + else: + still.append((g, i)) + pending = still + if pending and time.monotonic() <= kill_deadline: + time.sleep(_STARTUP_POLL_INTERVAL_S) + else: + break survivors: list[int] = [] for g, (shms, pids) in enumerate(groups): n = min(len(shms), len(pids)) @@ -6731,6 +8754,10 @@ def _step(fn) -> None: # C++ scheduler: once `dw.close()` runs the chip mailboxes are unusable # and we can no longer drive CTRL_RELEASE_DOMAIN. _step(self._cleanup_l3_l2_regions) + if self._live_global_domains: + _step(self._release_all_live_global_domains) + if self._global_node_domains: + _step(self._release_all_global_domain_nodes) if self._live_domains: _step(self._release_all_live_domains) _step(self._clear_child_prov) @@ -6759,6 +8786,7 @@ def _close_worker() -> None: self._orch = None _step(_close_worker) + _step(self._close_mpirun_groups) # Two-phase child shutdown: broadcast SHUTDOWN to EVERY group first, # then reap all groups together within the shared deadline. Sending # SHUTDOWN per-group-then-reap (serial) let a stuck child in the first @@ -6773,10 +8801,18 @@ def _close_worker() -> None: _step(lambda shms=shms: self._broadcast_child_shutdown(shms)) # Grace starts NOW, once SHUTDOWN is delivered to every group — not at # teardown entry — so the (blocking) pre-child cleanup above cannot - # eat it. Reap removes reclaimed pids/shms in place; a surviving child - # is left in place and reported as an error (terminal, not retried). + # eat it. Reap gives a graceful survivor a hard-kill backstop, then + # removes reclaimed pids/shms in place. A child still unreaped after + # both bounded phases is left in place and reported as an error. reap_deadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S - _step(lambda: self._reap_child_groups(groups, reap_deadline)) + _step( + lambda: self._reap_child_groups( + groups, + reap_deadline, + kill_survivors=True, + kill_process_groups=self._is_startup_root, + ) + ) _step(self._close_l3_l2_orch_comm) # Drop next-level worker refs only once their pids/shms are reclaimed. if not self._next_level_pids and not self._next_level_shms: diff --git a/src/a2a3/platform/onboard/host/comm_hccl.cpp b/src/a2a3/platform/onboard/host/comm_hccl.cpp index 8ce02540ff..6765e105d1 100644 --- a/src/a2a3/platform/onboard/host/comm_hccl.cpp +++ b/src/a2a3/platform/onboard/host/comm_hccl.cpp @@ -84,11 +84,18 @@ struct DomainAllocation { int rank = 0; int nranks = 0; + bool release_failed = false; VmmWindow local_window; std::vector peer_windows; CommContext *device_ctx = nullptr; // aclrtMalloc'd CommContext mirror }; +static_assert(sizeof(CommGlobalDomainDescriptor) == 288, "global domain descriptor ABI changed"); +static_assert( + sizeof(aclrtMemFabricHandle) <= COMM_GLOBAL_DOMAIN_HANDLE_BYTES, "Fabric handle exceeds global descriptor" +); +static std::unordered_map> global_domain_allocations; + struct CommHandle_ { int rank; int nranks; @@ -1533,6 +1540,199 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +) try { + if (domain_id == 0 || rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || + window_size == 0 || profile != COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC || descriptor_out == nullptr || + local_window_base_out == nullptr || global_domain_allocations.count(domain_id) != 0) { + return -1; + } + + int32_t device_id = -1; + if (aclrtGetDevice(&device_id) != ACL_SUCCESS) { + return -1; + } + auto allocation = std::make_unique(); + aclError status = create_local_fabric_window(device_id, window_size, &allocation->local_window); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] create local Fabric window -> %d", domain_rank, static_cast(status)); + return -1; + } + + aclrtMemFabricHandle fabric_handle{}; + status = export_fabric_window(allocation->local_window, &fabric_handle); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] export Fabric handle -> %d", domain_rank, static_cast(status)); + release_domain_windows(allocation.get()); + return -1; + } + status = + aclrtMemset(allocation->local_window.base, allocation->local_window.size, 0, allocation->local_window.size); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] zero local Fabric window -> %d", domain_rank, static_cast(status)); + release_domain_windows(allocation.get()); + return -1; + } + + allocation->rank = static_cast(domain_rank); + allocation->nranks = static_cast(rank_count); + CommGlobalDomainDescriptor descriptor{}; + descriptor.version = COMM_GLOBAL_DOMAIN_VERSION; + descriptor.profile = COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC; + descriptor.domain_rank = domain_rank; + descriptor.rank_count = rank_count; + descriptor.mapping_size = allocation->local_window.size; + descriptor.handle_size = sizeof(fabric_handle); + std::memcpy(descriptor.handle, &fabric_handle, sizeof(fabric_handle)); + + *descriptor_out = descriptor; + *local_window_base_out = reinterpret_cast(allocation->local_window.base); + global_domain_allocations.emplace(domain_id, std::move(allocation)); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] prepare exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] prepare unknown exception"); + return -1; +} + +extern "C" int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { + return -1; + } + auto &allocation = it->second; + if (descriptor_count != static_cast(allocation->nranks) || allocation->device_ctx != nullptr) { + return -1; + } + aclError cleanup_status = release_vmm_windows(&allocation->peer_windows); + if (cleanup_status != ACL_SUCCESS) { + LOG_ERROR( + "[global domain rank %d] clear stale peer windows -> %d", allocation->rank, static_cast(cleanup_status) + ); + return -1; + } + + std::vector rank_order(descriptor_count, nullptr); + for (size_t i = 0; i < descriptor_count; ++i) { + const auto &descriptor = descriptors[i]; + if (descriptor.version != COMM_GLOBAL_DOMAIN_VERSION || + descriptor.profile != COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC || + descriptor.rank_count != static_cast(allocation->nranks) || + descriptor.domain_rank >= static_cast(allocation->nranks) || + descriptor.mapping_size != allocation->local_window.size || + descriptor.handle_size != sizeof(aclrtMemFabricHandle) || rank_order[descriptor.domain_rank] != nullptr) { + return -1; + } + rank_order[descriptor.domain_rank] = &descriptor; + } + + int32_t device_id = -1; + if (aclrtGetDevice(&device_id) != ACL_SUCCESS) { + return -1; + } + CommContext ctx{}; + ctx.rankId = static_cast(allocation->rank); + ctx.rankNum = static_cast(allocation->nranks); + ctx.winSize = allocation->local_window.size; + std::vector peer_windows; + peer_windows.reserve(descriptor_count - 1); + for (uint32_t rank = 0; rank < static_cast(allocation->nranks); ++rank) { + const auto *descriptor = rank_order[rank]; + if (descriptor == nullptr) { + return -1; + } + uint64_t window_addr = reinterpret_cast(allocation->local_window.base); + if (rank != static_cast(allocation->rank)) { + aclrtMemFabricHandle fabric_handle{}; + std::memcpy(&fabric_handle, descriptor->handle, sizeof(fabric_handle)); + VmmWindow peer_window; + aclError status = import_fabric_window(device_id, fabric_handle, descriptor->mapping_size, &peer_window); + if (status != ACL_SUCCESS) { + LOG_ERROR( + "[global domain rank %d] import peer rank %u -> %d", allocation->rank, rank, + static_cast(status) + ); + return -1; + } + window_addr = reinterpret_cast(peer_window.base); + peer_windows.push_back(std::move(peer_window)); + } + ctx.windowsIn[rank] = window_addr; + ctx.windowsOut[rank] = window_addr; + } + + void *device_ctx = nullptr; + aclError status = aclrtMalloc(&device_ctx, sizeof(CommContext), ACL_MEM_MALLOC_HUGE_FIRST); + if (status != ACL_SUCCESS) { + return -1; + } + status = aclrtMemcpy(device_ctx, sizeof(CommContext), &ctx, sizeof(CommContext), ACL_MEMCPY_HOST_TO_DEVICE); + if (status != ACL_SUCCESS) { + aclrtFree(device_ctx); + return -1; + } + allocation->peer_windows = std::move(peer_windows); + allocation->device_ctx = static_cast(device_ctx); + *device_ctx_out = reinterpret_cast(device_ctx); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] import exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] import unknown exception"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t domain_id) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end()) { + return 0; + } + auto &allocation = it->second; + if (allocation->release_failed) { + LOG_ERROR( + "[global domain] domain %llu previously failed a partial release", + static_cast(domain_id) + ); + return -1; + } + int rc = 0; + if (allocation->device_ctx != nullptr) { + if (aclrtFree(allocation->device_ctx) != ACL_SUCCESS) { + rc = -1; + } + allocation->device_ctx = nullptr; + } + if (release_domain_windows(allocation.get()) != ACL_SUCCESS) { + rc = -1; + } + if (rc != 0) { + // The release helpers are destructive: some fields may already be + // cleared before a later ACL operation fails. Keep a terminal failure + // record so a repeated call cannot falsely report success. + allocation->release_failed = true; + LOG_ERROR( + "[global domain] domain %llu entered terminal partial-release state", + static_cast(domain_id) + ); + return rc; + } + global_domain_allocations.erase(it); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] release exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] release unknown exception"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (!h) return -1; diff --git a/src/a5/platform/onboard/host/comm_hccl.cpp b/src/a5/platform/onboard/host/comm_hccl.cpp index b00517539c..e11817c67d 100644 --- a/src/a5/platform/onboard/host/comm_hccl.cpp +++ b/src/a5/platform/onboard/host/comm_hccl.cpp @@ -1396,6 +1396,22 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int +comm_global_domain_prepare(uint64_t, uint32_t, uint32_t, size_t, uint32_t, CommGlobalDomainDescriptor *, uint64_t *) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + +extern "C" int comm_global_domain_import(uint64_t, const CommGlobalDomainDescriptor *, size_t, uint64_t *) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (!h) return -1; diff --git a/src/common/hierarchical/remote_endpoint.cpp b/src/common/hierarchical/remote_endpoint.cpp index fb96e3e6f4..2c0eac74cb 100644 --- a/src/common/hierarchical/remote_endpoint.cpp +++ b/src/common/hierarchical/remote_endpoint.cpp @@ -934,6 +934,13 @@ void RemoteL3Endpoint::control_remote_release_import(const RemoteBufferHandle &h run_control(remote_l3::ControlName::RELEASE_IMPORT, remote_l3::encode_release_import_request(request)); } +std::vector RemoteL3Endpoint::control_remote_domain( + remote_l3::ControlName control_name, const std::vector &command_bytes +) { + auto reply = run_control(control_name, command_bytes); + return std::move(reply.result_bytes); +} + void RemoteL3Endpoint::shutdown_child() { if (!transport_) return; try { diff --git a/src/common/hierarchical/remote_endpoint.h b/src/common/hierarchical/remote_endpoint.h index 39fc86c2c8..aec470d24c 100644 --- a/src/common/hierarchical/remote_endpoint.h +++ b/src/common/hierarchical/remote_endpoint.h @@ -110,6 +110,8 @@ class RemoteL3Endpoint : public WorkerEndpoint { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ) override; void control_remote_release_import(const RemoteBufferHandle &handle) override; + std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes) override; private: WorkerEndpointCaps caps_; diff --git a/src/common/hierarchical/remote_wire.cpp b/src/common/hierarchical/remote_wire.cpp index 8e3dbd820e..fe47ee4311 100644 --- a/src/common/hierarchical/remote_wire.cpp +++ b/src/common/hierarchical/remote_wire.cpp @@ -130,7 +130,7 @@ bool valid_frame_type(uint32_t v) { return false; } -bool valid_control_name(uint32_t v) { +bool valid_control_name_impl(uint32_t v) { switch (static_cast(v)) { case ControlName::UNREGISTER_CALLABLE: case ControlName::PREPARE_REGISTER_CALLABLE: @@ -147,6 +147,8 @@ bool valid_control_name(uint32_t v) { case ControlName::COMM_INIT: case ControlName::ALLOC_DOMAIN: case ControlName::RELEASE_DOMAIN: + case ControlName::COPY_TO_DOMAIN: + case ControlName::COPY_FROM_DOMAIN: return true; } return false; @@ -233,6 +235,8 @@ void validate_desc_against_inline_payload(const RemoteTensorDesc &desc, size_t i } // namespace +bool valid_control_name(uint32_t value) { return valid_control_name_impl(value); } + std::vector encode_frame(const FrameHeader &header, const std::vector &payload) { ensure(payload.size() <= MAX_FRAME_PAYLOAD_BYTES, "remote_wire: frame payload exceeds maximum"); ensure(header.flags == 0, "remote_wire: frame flags are reserved in v1"); diff --git a/src/common/hierarchical/remote_wire.h b/src/common/hierarchical/remote_wire.h index 720ccedba9..8d7e4ebc06 100644 --- a/src/common/hierarchical/remote_wire.h +++ b/src/common/hierarchical/remote_wire.h @@ -63,8 +63,12 @@ enum class ControlName : uint32_t { COMM_INIT = 13, ALLOC_DOMAIN = 14, RELEASE_DOMAIN = 15, + COPY_TO_DOMAIN = 16, + COPY_FROM_DOMAIN = 17, }; +bool valid_control_name(uint32_t value); + enum class ReadyState : uint32_t { NOT_READY = 0, READY = 1, diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index f8e2e82b09..27531b7e89 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -128,6 +128,11 @@ class Worker { control_digest_only(WorkerType type, int worker_id, uint64_t sub_cmd, const uint8_t *digest, double timeout_s) { return manager_.control_digest_only(type, worker_id, sub_cmd, digest, timeout_s); } + std::vector control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s + ) { + return manager_.control_payload(type, worker_id, sub_cmd, payload, payload_size, timeout_s); + } ControlResult remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -175,6 +180,11 @@ class Worker { return manager_.control_remote_import(importer_worker_id, export_desc, requested_access_flags); } void remote_release_import(const RemoteBufferHandle &handle) { manager_.control_remote_release_import(handle); } + std::vector control_remote_domain( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes + ) { + return manager_.control_remote_domain(worker_id, control_name, command_bytes); + } // Broadcast CTRL_REGISTER / CTRL_UNREGISTER for a ChipCallable digest to // every NEXT_LEVEL child in parallel. `blob_ptr`/`blob_size` describe diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index ff34e4b6eb..29f0988a9d 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -127,6 +127,9 @@ RemoteBufferHandle WorkerEndpoint::control_remote_import(int32_t, const RemoteBu void WorkerEndpoint::control_remote_release_import(const RemoteBufferHandle &) { throw_unsupported_control("control_remote_release_import"); } +std::vector WorkerEndpoint::control_remote_domain(remote_l3::ControlName, const std::vector &) { + throw_unsupported_control("control_remote_domain"); +} void WorkerEndpoint::control_generic(uint64_t, const char *, size_t, double, const uint8_t *) { throw_unsupported_control("control_generic"); } @@ -960,6 +963,12 @@ void WorkerThread::control_remote_release_import(const RemoteBufferHandle &handl endpoint_->control_remote_release_import(handle); } +std::vector +WorkerThread::control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes) { + if (!endpoint_) throw std::runtime_error("control_remote_domain: null endpoint"); + return endpoint_->control_remote_domain(control_name, command_bytes); +} + void WorkerThread::control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ) { @@ -1173,6 +1182,24 @@ ControlResult WorkerManager::control_digest_only( return result; } +std::vector WorkerManager::control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s +) { + if (payload == nullptr || payload_size == 0) { + throw std::runtime_error("control_payload: payload must be non-empty"); + } + WorkerThread *wt = get_worker_by_id(type, worker_id); + if (wt == nullptr) { + throw std::runtime_error("control_payload: invalid worker_id " + std::to_string(worker_id)); + } + std::string shm_name = make_shm_name(); + PosixShmHolder shm(shm_name, payload_size); + std::memcpy(shm.addr(), payload, payload_size); + wt->control_generic(sub_cmd, shm_name.c_str(), payload_size, timeout_s, nullptr); + auto *begin = static_cast(shm.addr()); + return {begin, begin + payload_size}; +} + ControlResult WorkerManager::control_remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -1315,6 +1342,25 @@ void WorkerManager::control_remote_release_import(const RemoteBufferHandle &hand wt->control_remote_release_import(handle); } +std::vector WorkerManager::control_remote_domain( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes +) { + WorkerThread *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (wt == nullptr) { + throw std::runtime_error("control_remote_domain: invalid worker_id " + std::to_string(worker_id)); + } + switch (control_name) { + case remote_l3::ControlName::COMM_INIT: + case remote_l3::ControlName::ALLOC_DOMAIN: + case remote_l3::ControlName::RELEASE_DOMAIN: + case remote_l3::ControlName::COPY_TO_DOMAIN: + case remote_l3::ControlName::COPY_FROM_DOMAIN: + return wt->control_remote_domain(control_name, command_bytes); + default: + throw std::runtime_error("control_remote_domain: control name is not a domain operation"); + } +} + std::vector WorkerManager::broadcast_register_all(const void *blob_ptr, size_t blob_size, const uint8_t *digest) { std::vector results; diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 622bcbf0ac..4c14eef357 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -252,6 +252,8 @@ class WorkerEndpoint { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); virtual void control_remote_release_import(const RemoteBufferHandle &handle); + virtual std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes); virtual void control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ); @@ -452,6 +454,8 @@ class WorkerThread { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); void control_remote_release_import(const RemoteBufferHandle &handle); + std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes); void control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ); @@ -538,6 +542,9 @@ class WorkerManager { void control_l3_l2_region_release(int worker_id, uint64_t region_id); ControlResult control_digest_only(WorkerType type, int worker_id, uint64_t sub_cmd, const uint8_t *digest, double timeout_s); + std::vector control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s + ); ControlResult control_remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -566,6 +573,9 @@ class WorkerManager { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); void control_remote_release_import(const RemoteBufferHandle &handle); + std::vector control_remote_domain( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes + ); // Broadcast CTRL_REGISTER for `digest` to every NEXT_LEVEL worker in // parallel. Stages `blob_size` bytes from `blob_ptr` into a per-call diff --git a/src/common/platform_comm/comm.h b/src/common/platform_comm/comm.h index b43a60675d..11da9b9114 100644 --- a/src/common/platform_comm/comm.h +++ b/src/common/platform_comm/comm.h @@ -37,6 +37,49 @@ extern "C" { typedef struct CommHandle_ *CommHandle; +#define COMM_GLOBAL_DOMAIN_VERSION 1U +#define COMM_GLOBAL_DOMAIN_HANDLE_BYTES 256U +#define COMM_GLOBAL_DOMAIN_DESCRIPTOR_BYTES 288U + +typedef enum CommGlobalDomainProfile { + COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM = 1, + COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC = 2, +} CommGlobalDomainProfile; + +typedef struct CommGlobalDomainDescriptor { + uint32_t version; + uint32_t profile; + uint32_t domain_rank; + uint32_t rank_count; + uint64_t mapping_size; + uint32_t handle_size; + uint32_t reserved; + uint8_t handle[COMM_GLOBAL_DOMAIN_HANDLE_BYTES]; +} CommGlobalDomainDescriptor; + +#ifdef __cplusplus +static_assert( + sizeof(CommGlobalDomainDescriptor) == COMM_GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + "CommGlobalDomainDescriptor wire size drift" +); +static_assert(offsetof(CommGlobalDomainDescriptor, version) == 0, "CommGlobalDomainDescriptor.version offset drift"); +static_assert(offsetof(CommGlobalDomainDescriptor, profile) == 4, "CommGlobalDomainDescriptor.profile offset drift"); +static_assert( + offsetof(CommGlobalDomainDescriptor, domain_rank) == 8, "CommGlobalDomainDescriptor.domain_rank offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, rank_count) == 12, "CommGlobalDomainDescriptor.rank_count offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, mapping_size) == 16, "CommGlobalDomainDescriptor.mapping_size offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, handle_size) == 24, "CommGlobalDomainDescriptor.handle_size offset drift" +); +static_assert(offsetof(CommGlobalDomainDescriptor, reserved) == 28, "CommGlobalDomainDescriptor.reserved offset drift"); +static_assert(offsetof(CommGlobalDomainDescriptor, handle) == 32, "CommGlobalDomainDescriptor.handle offset drift"); +#endif + /** Bit mask of DmaWorkspaceKind values this platform can provision. */ uint32_t dma_workspace_supported_mask(void); @@ -223,6 +266,42 @@ int comm_alloc_domain_windows( */ int comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_count, uint32_t domain_rank); +/** + * Create one local window for a Global CommDomain and export its transport + * descriptor. This operation has no + * HCCL or filesystem bootstrap dependency. + * + * The caller relays the descriptor through its control plane. It later + * passes + * the complete rank-ordered table to comm_global_domain_import(). A live + * domain_id is unique within one + * ChipWorker process. + */ +int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +); + +/** + * Import every peer descriptor and publish a device CommContext. + * + * descriptors must contain exactly one entry + * for every dense domain rank. + * The returned context and all imported mappings remain live until + * + * comm_global_domain_release(). + */ +int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +); + +/** + * Release a prepared or imported Global CommDomain. This is local teardown. + * The L4 owner is responsible for + * draining all rank tasks before fanout. + */ +int comm_global_domain_release(uint64_t domain_id); + /** * Synchronize all ranks. * diff --git a/src/common/platform_comm/comm_sim.cpp b/src/common/platform_comm/comm_sim.cpp index c57ebe2229..d6b08b38ed 100644 --- a/src/common/platform_comm/comm_sim.cpp +++ b/src/common/platform_comm/comm_sim.cpp @@ -50,6 +50,7 @@ #include #include #include +#include #include namespace { @@ -165,6 +166,58 @@ struct DomainAllocation { std::unique_ptr host_ctx; // device_ctx points here on sim }; +struct GlobalPeerMapping { + void *base = nullptr; + size_t size = 0; + + GlobalPeerMapping() = default; + GlobalPeerMapping(void *mapping_base, size_t mapping_size) : + base(mapping_base), + size(mapping_size) {} + ~GlobalPeerMapping() { + if (base != nullptr) { + munmap(base, size); + } + } + GlobalPeerMapping(const GlobalPeerMapping &) = delete; + GlobalPeerMapping &operator=(const GlobalPeerMapping &) = delete; + GlobalPeerMapping(GlobalPeerMapping &&other) noexcept : + base(std::exchange(other.base, nullptr)), + size(std::exchange(other.size, 0)) {} + GlobalPeerMapping &operator=(GlobalPeerMapping &&other) noexcept { + if (this != &other) { + if (base != nullptr) { + munmap(base, size); + } + base = std::exchange(other.base, nullptr); + size = std::exchange(other.size, 0); + } + return *this; + } +}; + +struct GlobalDomainAllocation { + ~GlobalDomainAllocation() { + if (local_base != nullptr) { + munmap(local_base, mapping_size); + } + if (!shm_name.empty()) { + shm_unlink(shm_name.c_str()); + } + } + + uint32_t rank = 0; + uint32_t nranks = 0; + std::string shm_name; + void *local_base = nullptr; + size_t mapping_size = 0; + std::vector peer_mappings; + std::unique_ptr host_ctx; +}; + +static_assert(sizeof(CommGlobalDomainDescriptor) == 288, "global domain descriptor ABI changed"); +static std::unordered_map> global_domain_allocations; + struct CommHandle_ { int rank; int nranks; @@ -685,6 +738,155 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +) try { + if (domain_id == 0 || rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || + window_size == 0 || profile != COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM || descriptor_out == nullptr || + local_window_base_out == nullptr) { + return -1; + } + if (global_domain_allocations.count(domain_id) != 0) { + return -1; + } + + std::string identity = + std::to_string(static_cast(domain_id)) + ":" + std::to_string(domain_rank); + std::string shm_name = make_shm_name(static_cast(getpid()), hash_id(identity.c_str())); + if (shm_name.empty() || shm_name.size() > COMM_GLOBAL_DOMAIN_HANDLE_BYTES) { + return -1; + } + + int fd = shm_open(shm_name.c_str(), O_CREAT | O_EXCL | O_RDWR, 0600); + if (fd < 0) { + return -1; + } + if (ftruncate(fd, static_cast(window_size)) != 0) { + close(fd); + shm_unlink(shm_name.c_str()); + return -1; + } + void *base = mmap(nullptr, window_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + close(fd); + if (base == MAP_FAILED) { + shm_unlink(shm_name.c_str()); + return -1; + } + std::memset(base, 0, window_size); + + auto allocation = std::make_unique(); + allocation->rank = domain_rank; + allocation->nranks = rank_count; + allocation->shm_name = shm_name; + allocation->local_base = base; + allocation->mapping_size = window_size; + + CommGlobalDomainDescriptor descriptor{}; + descriptor.version = COMM_GLOBAL_DOMAIN_VERSION; + descriptor.profile = COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM; + descriptor.domain_rank = domain_rank; + descriptor.rank_count = rank_count; + descriptor.mapping_size = window_size; + descriptor.handle_size = static_cast(shm_name.size()); + std::memcpy(descriptor.handle, shm_name.data(), shm_name.size()); + + *descriptor_out = descriptor; + *local_window_base_out = reinterpret_cast(base); + global_domain_allocations.emplace(domain_id, std::move(allocation)); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_prepare: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_prepare: unknown exception\n"); + return -1; +} + +extern "C" int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { + return -1; + } + auto &allocation = it->second; + if (descriptor_count != allocation->nranks || allocation->host_ctx != nullptr) { + return -1; + } + allocation->peer_mappings.clear(); + + std::vector rank_order(descriptor_count, nullptr); + for (size_t i = 0; i < descriptor_count; ++i) { + const auto &descriptor = descriptors[i]; + if (descriptor.version != COMM_GLOBAL_DOMAIN_VERSION || + descriptor.profile != COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM || descriptor.rank_count != allocation->nranks || + descriptor.domain_rank >= allocation->nranks || descriptor.mapping_size != allocation->mapping_size || + descriptor.handle_size == 0 || descriptor.handle_size > COMM_GLOBAL_DOMAIN_HANDLE_BYTES || + rank_order[descriptor.domain_rank] != nullptr) { + return -1; + } + rank_order[descriptor.domain_rank] = &descriptor; + } + + auto ctx = std::make_unique(); + ctx->rankId = allocation->rank; + ctx->rankNum = allocation->nranks; + ctx->winSize = allocation->mapping_size; + std::vector peer_mappings; + peer_mappings.reserve(allocation->nranks - 1); + for (uint32_t rank = 0; rank < allocation->nranks; ++rank) { + const auto *descriptor = rank_order[rank]; + if (descriptor == nullptr) { + return -1; + } + void *base = allocation->local_base; + if (rank != allocation->rank) { + std::string peer_name( + reinterpret_cast(descriptor->handle), static_cast(descriptor->handle_size) + ); + int fd = shm_open(peer_name.c_str(), O_RDWR, 0600); + if (fd < 0) { + return -1; + } + base = mmap(nullptr, allocation->mapping_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + close(fd); + if (base == MAP_FAILED) { + return -1; + } + peer_mappings.emplace_back(base, allocation->mapping_size); + } + ctx->windowsIn[rank] = reinterpret_cast(base); + ctx->windowsOut[rank] = reinterpret_cast(base); + } + + allocation->peer_mappings = std::move(peer_mappings); + *device_ctx_out = reinterpret_cast(ctx.get()); + allocation->host_ctx = std::move(ctx); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_import: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_import: unknown exception\n"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t domain_id) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end()) { + return 0; + } + global_domain_allocations.erase(it); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_release: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_release: unknown exception\n"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (h == nullptr) return -1; diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 8cc111168e..7507c53dfc 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -150,6 +150,9 @@ void ChipWorker::init( comm_alloc_domain_windows_fn_ = load_symbol(handle, "comm_alloc_domain_windows"); comm_release_domain_windows_fn_ = load_symbol(handle, "comm_release_domain_windows"); + comm_global_domain_prepare_fn_ = load_symbol(handle, "comm_global_domain_prepare"); + comm_global_domain_import_fn_ = load_symbol(handle, "comm_global_domain_import"); + comm_global_domain_release_fn_ = load_symbol(handle, "comm_global_domain_release"); comm_barrier_fn_ = load_symbol(handle, "comm_barrier"); comm_destroy_fn_ = load_symbol(handle, "comm_destroy"); } catch (...) { @@ -263,6 +266,9 @@ void ChipWorker::init( comm_get_window_size_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); @@ -310,6 +316,9 @@ void ChipWorker::init( comm_derive_context_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); @@ -340,6 +349,15 @@ void ChipWorker::init( } void ChipWorker::finalize() { + // Global domains are independent of the legacy communicator sessions. + // Release them while the host runtime and device context are still alive. + if (comm_global_domain_release_fn_ != nullptr) { + for (uint64_t domain_id : global_domain_ids_) { + comm_global_domain_release_fn_(domain_id); + } + } + global_domain_ids_.clear(); + // Defensive: if the user never called comm_destroy, reclaim all owned // communicator handles and streams before tearing down the device context. clear_comm_sessions(); @@ -386,6 +404,9 @@ void ChipWorker::finalize() { comm_derive_context_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); @@ -825,6 +846,67 @@ void ChipWorker::comm_release_domain_windows( } } +std::tuple, uint64_t, size_t> ChipWorker::comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile +) { + if (comm_global_domain_prepare_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_prepare is not supported by this runtime"); + } + auto [tracked, inserted] = global_domain_ids_.insert(domain_id); + if (!inserted) { + throw std::runtime_error("comm_global_domain_prepare received a duplicate domain_id"); + } + CommGlobalDomainDescriptor descriptor{}; + uint64_t local_window_base = 0; + int rc = comm_global_domain_prepare_fn_( + domain_id, domain_rank, rank_count, window_size, profile, &descriptor, &local_window_base + ); + if (rc != 0) { + global_domain_ids_.erase(tracked); + throw std::runtime_error("comm_global_domain_prepare failed with code " + std::to_string(rc)); + } + if (local_window_base == 0 || descriptor.mapping_size == 0) { + comm_global_domain_release_fn_(domain_id); + global_domain_ids_.erase(domain_id); + throw std::runtime_error("comm_global_domain_prepare returned an invalid window"); + } + const auto *begin = reinterpret_cast(&descriptor); + std::vector descriptor_bytes(begin, begin + sizeof(descriptor)); + return {std::move(descriptor_bytes), local_window_base, static_cast(descriptor.mapping_size)}; +} + +uint64_t ChipWorker::comm_global_domain_import(uint64_t domain_id, const std::vector &descriptors) { + if (comm_global_domain_import_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_import is not supported by this runtime"); + } + if (descriptors.empty() || descriptors.size() % sizeof(CommGlobalDomainDescriptor) != 0) { + throw std::runtime_error("comm_global_domain_import descriptor table size is invalid"); + } + uint64_t device_ctx = 0; + int rc = comm_global_domain_import_fn_( + domain_id, reinterpret_cast(descriptors.data()), + descriptors.size() / sizeof(CommGlobalDomainDescriptor), &device_ctx + ); + if (rc != 0) { + throw std::runtime_error("comm_global_domain_import failed with code " + std::to_string(rc)); + } + if (device_ctx == 0) { + throw std::runtime_error("comm_global_domain_import returned a null device context"); + } + return device_ctx; +} + +void ChipWorker::comm_global_domain_release(uint64_t domain_id) { + if (comm_global_domain_release_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_release is not supported by this runtime"); + } + int rc = comm_global_domain_release_fn_(domain_id); + if (rc != 0) { + throw std::runtime_error("comm_global_domain_release failed with code " + std::to_string(rc)); + } + global_domain_ids_.erase(domain_id); +} + void ChipWorker::comm_barrier(uint64_t comm_handle) { int rc = comm_barrier_fn_(reinterpret_cast(comm_handle)); if (rc != 0) { diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 2a31935641..02df3a837b 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -14,9 +14,12 @@ #include #include +#include #include +#include #include +#include "../platform_comm/comm.h" #include "../task_interface/call_config.h" #include "../task_interface/task_args.h" #include "pipeline_slot_pool.h" @@ -155,6 +158,11 @@ class ChipWorker { /// inside the backend's per-allocation record). void comm_release_domain_windows(uint64_t comm_handle, uint64_t allocation_id, size_t rank_count, uint32_t domain_rank); + std::tuple, uint64_t, size_t> comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile + ); + uint64_t comm_global_domain_import(uint64_t domain_id, const std::vector &descriptors); + void comm_global_domain_release(uint64_t domain_id); void comm_barrier(uint64_t comm_handle); void comm_destroy(uint64_t comm_handle); void comm_destroy_all(); @@ -216,6 +224,10 @@ class ChipWorker { using CommAllocDomainWindowsFn = int (*)(void *, uint64_t, const uint32_t *, size_t, uint32_t, size_t, uint64_t *, uint64_t *); using CommReleaseDomainWindowsFn = int (*)(void *, uint64_t, size_t, uint32_t); + using CommGlobalDomainPrepareFn = + int (*)(uint64_t, uint32_t, uint32_t, size_t, uint32_t, CommGlobalDomainDescriptor *, uint64_t *); + using CommGlobalDomainImportFn = int (*)(uint64_t, const CommGlobalDomainDescriptor *, size_t, uint64_t *); + using CommGlobalDomainReleaseFn = int (*)(uint64_t); using CommBarrierFn = int (*)(void *); using CommDestroyFn = int (*)(void *); @@ -269,11 +281,15 @@ class ChipWorker { CommDeriveContextFn comm_derive_context_fn_ = nullptr; CommAllocDomainWindowsFn comm_alloc_domain_windows_fn_ = nullptr; CommReleaseDomainWindowsFn comm_release_domain_windows_fn_ = nullptr; + CommGlobalDomainPrepareFn comm_global_domain_prepare_fn_ = nullptr; + CommGlobalDomainImportFn comm_global_domain_import_fn_ = nullptr; + CommGlobalDomainReleaseFn comm_global_domain_release_fn_ = nullptr; CommBarrierFn comm_barrier_fn_ = nullptr; CommDestroyFn comm_destroy_fn_ = nullptr; void *device_ctx_ = nullptr; std::vector comm_sessions_; std::unordered_map comm_session_index_; + std::unordered_set global_domain_ids_; uint64_t base_comm_handle_ = 0; // Slot 0 with no generation bookkeeping: an unleased run is a caller that diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 4e97288242..fb27b1e8bc 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -197,6 +197,7 @@ add_library(hierarchical_objs OBJECT ) target_include_directories(hierarchical_objs PUBLIC ${HIERARCHICAL_SRC_DIR} + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/include ${CMAKE_SOURCE_DIR}/../../../src/common/task_interface ${WORKER_SRC_DIR} ) diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index e8ce434496..ad847532e3 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -42,10 +42,14 @@ encode_register_callable_command, encode_remote_chip_callable_payload, ) +from simpler.remote_l3_protocol import ( + RemoteAddressSpace as WireRemoteAddressSpace, +) from simpler.remote_l3_session import ( _install_manifest_dispatcher_registry, _install_manifest_inner_registry, _prepare_register_callable, + _RemoteBufferEntry, _unpublish_inner_handle, get_inner_handle, ) @@ -115,6 +119,44 @@ def _remote_sum_u8_orch(orch, args, cfg): dst_data[0] = sum(int(src_data[i]) for i in range(src.nbytes())) & 0xFF +def test_remote_buffer_entry_releases_child_visible_host_buffer(): + raw = bytearray(8) + view = memoryview(raw) + + class FakeHostBuffer: + data_ptr = ctypes.addressof(ctypes.c_char.from_buffer(raw)) + buffer = view + shm_name = "host-buffer-test" + + class FakeOwner: + def __init__(self): + self.freed = [] + + def free_host_buffer(self, handle): + self.freed.append(handle) + + data = FakeHostBuffer() + owner = FakeOwner() + entry = _RemoteBufferEntry( + data=data, + nbytes=len(raw), + generation=1, + address_space=WireRemoteAddressSpace.REMOTE_DEVICE, + owner=cast(Worker, owner), + ) + + assert entry.addr == data.data_ptr + assert entry.shm_name == data.shm_name + entry.close() + assert owner.freed == [data] + assert entry.owner is None + with pytest.raises(ValueError, match="released memoryview"): + _ = view[0] + + # Session finalization may defensively revisit an already closed entry. + entry.close() + + class _FakeRemoteControlResult: def __init__(self, worker_id: int, ok: bool = True, error_message: str = ""): self.worker_type = "NEXT_LEVEL" @@ -628,6 +670,44 @@ def test_remote_session_manifest_uses_endpoint_host_as_default_bind(): worker.close() +def test_remote_manifest_carries_pre_registered_inner_chip_callable(): + worker = Worker(level=4, num_sub_workers=0) + chip = ChipCallable.build(signature=[], func_name="x", binary=b"\x01", children=[]) + try: + worker_id = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint="127.0.0.1:19073", + platform="a2a3sim", + device_ids=(0,), + ) + ) + handle = worker.register(chip) + manifest = worker._build_remote_manifest( + spec=worker._remote_worker_specs[0], + worker_id=worker_id, + session_id=1, + startup_remaining_s=30.0, + ) + + assert len(manifest["inner_l3_worker"]) == 1 + entry = manifest["inner_l3_worker"][0] + assert entry["hashid"] == handle.digest.hex() + command = encode_register_callable_command( + RemoteRegistryTarget.INNER_L3_WORKER, + CallableKind.CHIP_CALLABLE, + handle.digest, + 1, + bytes.fromhex(entry["payload_hex"]), + ) + digest, kind, registry, target = _prepare_register_callable(command, manifest) + assert digest == handle.digest + assert kind is CallableKind.CHIP_CALLABLE + assert registry is RemoteRegistryTarget.INNER_L3_WORKER + assert isinstance(target, ChipCallable) + finally: + worker.close() + + def test_remote_session_manifest_requires_wildcard_bind_opt_in(): worker = Worker(level=4, num_sub_workers=0) try: diff --git a/tests/ut/py/test_global_comm_domain.py b/tests/ut/py/test_global_comm_domain.py new file mode 100644 index 0000000000..80be90692d --- /dev/null +++ b/tests/ut/py/test_global_comm_domain.py @@ -0,0 +1,1024 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import time +from collections import Counter +from typing import cast + +import pytest +from simpler.global_comm_domain import ( + CTRL_GLOBAL_DOMAIN_COPY_FROM, + CTRL_GLOBAL_DOMAIN_COPY_TO, + CTRL_GLOBAL_DOMAIN_IMPORT, + CTRL_GLOBAL_DOMAIN_PREPARE, + CTRL_GLOBAL_DOMAIN_RELEASE, + GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + GLOBAL_DOMAIN_PROFILE_IDS, + GLOBAL_DOMAIN_VERSION, + GlobalCommInitCommand, + GlobalDomainBuffer, + GlobalDomainCommand, + GlobalDomainDescriptor, + GlobalDomainMember, + GlobalDomainPhase, + decode_comm_init, + decode_descriptor_table, + decode_domain_command, + encode_comm_init, + encode_comm_init_result, + encode_descriptor_table, + encode_domain_command, + resolve_global_comm_capability, + validate_descriptor_table, +) + + +def test_global_domain_control_ids_do_not_overlap_worker_controls(): + from simpler.worker import _CTRL_COMMITTED_DEVICE_MEMORY, _CTRL_GLOBAL_DOMAIN_NODE # noqa: PLC0415 + + control_ids = ( + _CTRL_COMMITTED_DEVICE_MEMORY, + CTRL_GLOBAL_DOMAIN_PREPARE, + CTRL_GLOBAL_DOMAIN_IMPORT, + CTRL_GLOBAL_DOMAIN_RELEASE, + CTRL_GLOBAL_DOMAIN_COPY_TO, + CTRL_GLOBAL_DOMAIN_COPY_FROM, + _CTRL_GLOBAL_DOMAIN_NODE, + ) + + assert len(control_ids) == len(set(control_ids)) + + +def _members() -> tuple[GlobalDomainMember, ...]: + return ( + GlobalDomainMember(0, 0, 3, 0), + GlobalDomainMember(1, 0, 7, 1), + ) + + +def _descriptors() -> tuple[GlobalDomainDescriptor, ...]: + return tuple( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS["sim"], + domain_rank=rank, + rank_count=2, + mapping_size=4096, + handle=f"/simpler-test-{rank}".encode(), + ) + for rank in range(2) + ) + + +@pytest.mark.parametrize( + ("platform", "profile"), + ( + ("a2a3sim", "sim"), + ("a5sim", "sim"), + ("a2a3", "a3-fabric-v1"), + ), +) +def test_global_comm_capability_reports_only_implemented_backends(platform, profile): + result = resolve_global_comm_capability(platform=platform, profile=profile, local_device_count=2) + + assert result.profile == profile + assert result.max_ranks == 64 + assert result.descriptor_bytes == GLOBAL_DOMAIN_DESCRIPTOR_BYTES + assert result.local_device_count == 2 + + +@pytest.mark.parametrize( + ("platform", "profile"), + ( + ("a2a3", "sim"), + ("a2a3sim", "a3-fabric-v1"), + ("a5", "sim"), + ("a5", "a3-fabric-v1"), + ), +) +def test_global_comm_capability_rejects_unimplemented_backends(platform, profile): + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + resolve_global_comm_capability(platform=platform, profile=profile, local_device_count=2) + + +def test_local_l3_comm_init_rejects_unsupported_capability_without_caching_topology(): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + from simpler.worker import Worker, _GlobalNodeRuntime, _run_local_global_domain_control # noqa: PLC0415 + + inner_worker = Worker(level=3, num_sub_workers=0) + runtime = _GlobalNodeRuntime( + worker_id=0, + device_ids=(0,), + platform="a5", + comm_profile="sim", + global_device_ranks=(0,), + node_rank=0, + node_count=1, + cluster_id="cluster", + is_remote=False, + ) + comm_inits = {} + command = GlobalCommInitCommand( + cluster_id="cluster", + topology_hash="topology", + profile="sim", + node_rank=0, + node_count=1, + members=(GlobalDomainMember(0, 0, 0, 0),), + ) + + try: + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + _run_local_global_domain_control( + inner_worker, + runtime, + comm_inits, + ControlName.COMM_INIT, + encode_comm_init(command), + ) + + assert comm_inits == {} + finally: + inner_worker.close() + + +def test_global_domain_wire_round_trips_topology_and_descriptor_table(): + init = GlobalCommInitCommand("cluster", "topology", "sim", 0, 2, _members()) + command = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=11, + generation=1, + name="tp", + profile="sim", + window_size=2048, + members=_members(), + buffers=(GlobalDomainBuffer("payload", 128),), + descriptors=_descriptors(), + ) + + assert decode_comm_init(encode_comm_init(init)) == init + assert decode_domain_command(encode_domain_command(command)) == command + assert decode_descriptor_table(encode_descriptor_table(_descriptors())) == _descriptors() + assert GLOBAL_DOMAIN_DESCRIPTOR_BYTES == 288 + + +def test_global_domain_node_import_records_window_and_buffer_extents(): + from simpler.global_comm_domain import LOCAL_DOMAIN_MAGIC, LOCAL_IMPORT_REPLY # noqa: PLC0415 + from simpler.worker import Worker, _GlobalNodeDomainState # noqa: PLC0415 + + domain_id = 41 + generation = 3 + node_worker_id = 7 + local_worker_id = 0 + local_base = 0x100000 + mapping_size = 4096 + member = GlobalDomainMember(node_worker_id, local_worker_id, 0, 0) + buffers = ( + GlobalDomainBuffer("first", 256), + GlobalDomainBuffer("second", 512), + ) + prepared = GlobalDomainCommand( + phase=GlobalDomainPhase.PREPARE_EXPORT, + domain_id=domain_id, + generation=generation, + name="mpi-import", + profile="sim", + window_size=mapping_size, + members=(member,), + buffers=buffers, + ) + imported = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=domain_id, + generation=generation, + name=prepared.name, + profile=prepared.profile, + window_size=prepared.window_size, + members=prepared.members, + buffers=prepared.buffers, + descriptors=( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS["sim"], + domain_rank=0, + rank_count=1, + mapping_size=mapping_size, + handle=b"/mpi-import", + ), + ), + ) + + class _ControlStub: + def control_payload(self, _worker_type, worker_id, sub_cmd, _payload, _timeout): + assert worker_id == local_worker_id + assert sub_cmd == CTRL_GLOBAL_DOMAIN_IMPORT + reply = bytearray(LOCAL_IMPORT_REPLY.size) + LOCAL_IMPORT_REPLY.pack_into( + reply, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + domain_id, + generation, + 0x55, + local_base, + mapping_size, + ) + return bytes(reply) + + worker = Worker(level=3, device_ids=(0,), num_sub_workers=0) + worker._worker = _ControlStub() + worker._global_node_domains[domain_id] = _GlobalNodeDomainState(command=prepared) + try: + worker._import_global_domain_node(imported, node_worker_id) + + provenance_id = worker._global_domain_provenance_id(domain_id) + window_entry = worker._child_alloc_prov[(local_worker_id, local_base)] + second_buffer_entry = worker._child_alloc_prov[(local_worker_id, local_base + buffers[0].nbytes)] + assert window_entry.domain_allocation_ids[provenance_id] == mapping_size + assert second_buffer_entry.domain_allocation_ids[provenance_id] == buffers[1].nbytes + finally: + worker._worker = None + worker._global_node_domains.clear() + worker._child_alloc_prov.clear() + worker.close() + + +def _failure_injection_worker(*, platform: str = "a2a3sim", profile: str = "sim"): + from simpler.worker import RemoteWorkerSpec, Worker, _RunResources # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{19073 + index}", + platform=platform, + device_ids=(0,), + comm_profile=profile, + global_device_ranks=(index,), + ) + ) + for index in range(2) + ) + resources = _RunResources() + worker._worker = object() + worker._building_run_resources = resources + return worker, resources, node_ids + + +def _mpi_static_worker(): + from simpler.worker import MpiL3GroupSpec, Worker, _RunResources # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + node_ids = worker.add_mpirun_worker_group( + MpiL3GroupSpec( + hosts=("127.0.0.1", "127.0.0.1"), + platform="a2a3sim", + command_port_base=21073, + health_port_base=22073, + device_ids_by_rank=((0,), (0,)), + comm_profile="sim", + global_device_ranks_by_rank=((0,), (1,)), + ) + ) + resources = _RunResources() + worker._worker = object() + worker._building_run_resources = resources + return worker, resources, node_ids + + +def _install_global_domain_failure_injector(monkeypatch, worker, *, fail_phase, fail_node): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + + calls = [] + + def control(worker_id, control_name, payload): + control_name = ControlName(control_name) + if control_name is ControlName.COMM_INIT: + init = decode_comm_init(payload) + calls.append(("COMM_INIT", worker_id)) + return encode_comm_init_result( + resolve_global_comm_capability( + platform="a2a3sim", + profile=init.profile, + local_device_count=1, + ) + ) + + assert control_name is ControlName.ALLOC_DOMAIN + command = decode_domain_command(payload) + calls.append((command.phase, worker_id)) + if command.phase is fail_phase and worker_id == fail_node: + raise RuntimeError(f"injected {command.phase.name} failure") + if command.phase is not GlobalDomainPhase.PREPARE_EXPORT: + return b"" + descriptors = tuple( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS[command.profile], + domain_rank=member.domain_rank, + rank_count=len(command.members), + mapping_size=4096, + handle=f"/injected-{member.domain_rank}".encode(), + ) + for member in command.members + if member.node_worker_id == worker_id + ) + return encode_descriptor_table(descriptors) + + monkeypatch.setattr(worker, "_global_domain_control", control) + return calls + + +def test_mpirun_group_global_domain_uses_mpi_prepare_commit_without_l4_import(monkeypatch): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _mpi_static_worker() + calls = [] + + def control(worker_id, control_name, payload): + control_name = ControlName(control_name) + if control_name is ControlName.COMM_INIT: + init = decode_comm_init(payload) + calls.append(("COMM_INIT", worker_id)) + return encode_comm_init_result( + resolve_global_comm_capability( + platform="a2a3sim", + profile=init.profile, + local_device_count=1, + ) + ) + assert control_name is ControlName.ALLOC_DOMAIN + command = decode_domain_command(payload) + calls.append((command.phase, worker_id)) + if command.phase is GlobalDomainPhase.PREPARE_EXPORT: + descriptors = tuple( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS[command.profile], + domain_rank=member.domain_rank, + rank_count=len(command.members), + mapping_size=4096, + handle=f"/mpi-prepared-{member.domain_rank}".encode(), + ) + for member in command.members + ) + return encode_descriptor_table(descriptors) + if command.phase is GlobalDomainPhase.IMPORT: + raise RuntimeError("L4 broker IMPORT should not run for a full mpirun group") + return b"" + + monkeypatch.setattr(worker, "_global_domain_control", control) + try: + handle = worker._allocate_global_domain( + name="mpi-static", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + assert handle.mapping_size == 4096 + assert handle.members[0].global_device_rank == 0 + assert handle.members[1].global_device_rank == 1 + counts = Counter(phase for phase, _worker_id in calls) + assert counts["COMM_INIT"] == 2 + assert counts[GlobalDomainPhase.PREPARE_EXPORT] == 2 + assert counts[GlobalDomainPhase.COMMIT] == 2 + assert counts[GlobalDomainPhase.IMPORT] == 0 + assert worker._live_global_domains["mpi-static"] is handle + assert resources.live_global_domains["mpi-static"] is handle + finally: + _close_failure_injection_worker(worker, resources) + + +def _close_failure_injection_worker(worker, resources): + worker._building_run_resources = None + worker._live_global_domains.clear() + resources.live_global_domains.clear() + worker._worker = None + worker.close() + + +@pytest.mark.parametrize( + "fail_phase", + ( + GlobalDomainPhase.PREPARE_EXPORT, + GlobalDomainPhase.IMPORT, + GlobalDomainPhase.COMMIT, + ), +) +def test_global_domain_transaction_aborts_all_prepared_nodes_after_phase_failure(monkeypatch, fail_phase): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker() + calls = _install_global_domain_failure_injector( + monkeypatch, + worker, + fail_phase=fail_phase, + fail_node=node_ids[1], + ) + try: + with pytest.raises(RuntimeError, match=f"injected {fail_phase.name} failure"): + worker._allocate_global_domain( + name="failure-injection", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + abort_nodes = [node_id for phase, node_id in calls if phase is GlobalDomainPhase.ABORT] + assert abort_nodes == list(node_ids) + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_global_domain_abort_failure_preserves_primary_error_and_continues_cleanup(monkeypatch): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker() + calls = _install_global_domain_failure_injector( + monkeypatch, + worker, + fail_phase=GlobalDomainPhase.IMPORT, + fail_node=node_ids[1], + ) + original_control = worker._global_domain_control + + def fail_first_abort(worker_id, control_name, payload): + if ControlName(control_name) is ControlName.ALLOC_DOMAIN: + command = decode_domain_command(payload) + if command.phase is GlobalDomainPhase.ABORT and worker_id == node_ids[0]: + calls.append((command.phase, worker_id)) + raise RuntimeError("injected ABORT failure") + return original_control(worker_id, control_name, payload) + + monkeypatch.setattr(worker, "_global_domain_control", fail_first_abort) + try: + with pytest.raises(RuntimeError, match="injected IMPORT failure"): + worker._allocate_global_domain( + name="abort-failure-injection", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + abort_nodes = [node_id for phase, node_id in calls if phase is GlobalDomainPhase.ABORT] + assert abort_nodes == list(node_ids) + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_allocate_global_domain_rejects_unsupported_capability_before_control(monkeypatch): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker(platform="a5", profile="sim") + calls = [] + monkeypatch.setattr(worker, "_global_domain_control", lambda *args: calls.append(args)) + try: + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + worker._allocate_global_domain( + name="unsupported", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + assert calls == [] + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_global_domain_descriptor_table_rejects_different_mapping_sizes(): + descriptors = list(_descriptors()) + descriptors[1] = GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS["sim"], + domain_rank=1, + rank_count=2, + mapping_size=8192, + handle=b"/simpler-test-1", + ) + + with pytest.raises(ValueError, match="mapping sizes differ"): + validate_descriptor_table(tuple(descriptors), rank_count=2, profile="sim") + + +def test_global_domain_release_retries_after_callback_failure(): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + + attempts = 0 + + def release_fn(_handle): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("transient release failure") + + handle = GlobalCommDomainHandle( + name="retry", + members=(), + buffers=(), + domain_id=17, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=release_fn, + ) + + with pytest.raises(RuntimeError, match="transient release failure"): + handle.release() + + assert not handle.released + handle.release() + assert handle.released + assert attempts == 2 + + +def test_global_domain_handle_repr_reports_lifecycle_state(): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + + handle = GlobalCommDomainHandle( + name="repr", + members=(), + buffers=(), + domain_id=18, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=lambda _handle: None, + ) + + assert repr(handle) == "GlobalCommDomainHandle(name='repr', members=0, live)" + handle.release() + assert repr(handle) == "GlobalCommDomainHandle(name='repr', members=0, released)" + handle._freed = True + assert repr(handle) == "GlobalCommDomainHandle(name='repr', members=0, freed)" + + +def test_mpi_global_domain_collective_timeout_releases_local_state(monkeypatch): + from simpler.mpi_l3_session import MpiGlobalDomainExchange # noqa: PLC0415 + + class _PendingRequest: + @staticmethod + def test(): + return False, None + + class _Comm: + aborted = False + + @staticmethod + def Get_rank(): + return 0 + + @staticmethod + def iallgather(_payload): + return _PendingRequest() + + def Abort(self, _error_code): + self.aborted = True + raise RuntimeError("fake MPI abort") + + now = iter((0.0, 2.0)) + monkeypatch.setattr("simpler.mpi_l3_session.time.monotonic", lambda: next(now)) + comm = _Comm() + exchange = MpiGlobalDomainExchange(comm, group_worker_ids=(7,), timeout_s=1.0) + releases = [] + + with pytest.raises(TimeoutError, match="prepare timed out"): + exchange._allgather(b"payload", operation="prepare", on_timeout=lambda: releases.append(True)) + + assert releases == [True] + assert comm.aborted + + +def test_mpi_global_domain_prepare_failure_releases_before_collective(): + from simpler.mpi_l3_session import MpiGlobalDomainExchange # noqa: PLC0415 + from simpler.worker import Worker # noqa: PLC0415 + + class _CompletedRequest: + def __init__(self, payload): + self._payload = payload + + def test(self): + return True, [self._payload] + + class _Comm: + @staticmethod + def Get_rank(): + return 0 + + @staticmethod + def iallgather(payload): + return _CompletedRequest(payload) + + class _InnerWorker: + released = False + + @staticmethod + def _prepare_global_domain_node(_command, _worker_id): + raise RuntimeError("injected prepare failure") + + def _release_global_domain_node(self, _command, *, suppress_errors): + assert suppress_errors + self.released = True + + command = GlobalDomainCommand( + phase=GlobalDomainPhase.PREPARE_EXPORT, + domain_id=20, + generation=1, + name="mpi-failure", + profile="sim", + window_size=4096, + members=(GlobalDomainMember(7, 0, 0, 0),), + buffers=(), + ) + inner_worker = _InnerWorker() + exchange = MpiGlobalDomainExchange(_Comm(), group_worker_ids=(7,), timeout_s=1.0) + + with pytest.raises(RuntimeError, match="prepare failed on rank 0"): + exchange.prepare_import(command, cast(Worker, inner_worker), 7) + + assert inner_worker.released + + +def test_mpirun_group_cleanup_continues_after_one_process_wait_fails(): + class _Process: + def __init__(self, *, fail_wait): + self.fail_wait = fail_wait + self.waited = False + + @staticmethod + def poll(): + return 0 + + def wait(self, *, timeout): + assert timeout == 0.1 + self.waited = True + if self.fail_wait: + raise RuntimeError("injected wait failure") + return 0 + + worker, resources, _node_ids = _mpi_static_worker() + group = worker._mpi_l3_groups[0] + first_process = _Process(fail_wait=True) + second_process = _Process(fail_wait=False) + first = type(group)( + group_id="first", + spec=group.spec, + ranks=group.ranks, + process=cast(subprocess.Popen, first_process), + ) + second = type(group)( + group_id="second", + spec=group.spec, + ranks=group.ranks, + process=cast(subprocess.Popen, second_process), + ) + worker._mpi_l3_groups[:] = [first, second] + try: + with pytest.raises(RuntimeError, match="first cleanup wait after terminate"): + worker._close_mpirun_groups(timeout_s=0.1) + + assert first_process.waited + assert second_process.waited + assert first.process is None + assert second.process is None + finally: + worker._mpi_l3_groups.clear() + _close_failure_injection_worker(worker, resources) + + +def test_old_global_domain_release_does_not_remove_same_name_replacement(): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + from simpler.worker import Worker, _RunResources # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + resources = _RunResources() + + def make_handle(domain_id: int) -> GlobalCommDomainHandle: + return GlobalCommDomainHandle( + name="reuse", + members=(), + buffers=(), + domain_id=domain_id, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=worker._release_global_domain_handle, + ) + + first = make_handle(17) + second = make_handle(18) + worker._worker = object() + worker._building_run_resources = resources + worker._live_global_domains[first.name] = first + resources.live_global_domains[first.name] = first + try: + first.release() + worker._live_global_domains[second.name] = second + resources.live_global_domains[second.name] = second + + worker._execute_pending_global_domain_releases(resources) + + assert first.freed + assert worker._live_global_domains[second.name] is second + assert resources.live_global_domains[second.name] is second + finally: + worker._building_run_resources = None + worker._live_global_domains.clear() + resources.live_global_domains.clear() + worker._worker = None + worker.close() + + +def test_global_domain_backend_release_failure_is_terminal(monkeypatch): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + from simpler.worker import Worker # noqa: PLC0415 + + attempts = 0 + worker = Worker(level=4, num_sub_workers=0) + worker._worker = object() + + def fail_control(_worker_id, _control_name, _payload): + nonlocal attempts + attempts += 1 + raise RuntimeError("partial backend release") + + monkeypatch.setattr(worker, "_global_domain_control", fail_control) + handle = GlobalCommDomainHandle( + name="terminal", + members=(_members()[0],), + buffers=(), + domain_id=19, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=worker._release_global_domain_handle, + ) + try: + with pytest.raises(RuntimeError, match="partial backend release"): + worker._free_global_domain_after_fence(handle) + with pytest.raises(RuntimeError, match="partial backend release"): + worker._free_global_domain_after_fence(handle) + + assert attempts == 1 + assert not handle.freed + assert worker._failed_global_domain_releases[handle.domain_id] is handle + finally: + worker._failed_global_domain_releases.clear() + worker._worker = None + worker.close() + + +def test_childless_host_buffer_uses_dedicated_exception(): + from simpler.worker import Worker, _NoHostBufferChildrenError # noqa: PLC0415 + + worker = Worker(level=3, num_sub_workers=0) + try: + with pytest.raises(_NoHostBufferChildrenError, match="at least one forked chip or sub child"): + worker._create_host_buffer_locked(64) + finally: + worker.close() + + +def _free_tcp_ports(count: int) -> tuple[int, ...]: + sockets = [] + try: + for _ in range(count): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + sockets.append(sock) + return tuple(int(sock.getsockname()[1]) for sock in sockets) + finally: + for sock in sockets: + sock.close() + + +def _wait_for_tcp_ports(ports: tuple[int, ...], timeout_s: float = 5.0) -> None: + pending = set(ports) + deadline = time.monotonic() + timeout_s + while pending: + for port in tuple(pending): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"remote L3 daemons did not become ready on ports {sorted(pending)}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=min(0.1, remaining)): + pending.remove(port) + except OSError: + pass + if pending: + time.sleep(0.01) + + +def _stop_daemons(daemons) -> str: + diagnostics = [] + for index, daemon in enumerate(daemons): + if daemon.poll() is None: + daemon.terminate() + try: + stdout, stderr = daemon.communicate(timeout=5) + except subprocess.TimeoutExpired: + daemon.kill() + stdout, stderr = daemon.communicate(timeout=5) + diagnostics.append(f"daemon[{index}] stdout:\n{stdout}\ndaemon[{index}] stderr:\n{stderr}") + return "\n".join(diagnostics) + + +@pytest.mark.skipif(os.name == "nt", reason="hierarchical workers require fork") +def test_two_remote_daemons_build_and_copy_global_domain_without_mpirun(): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + from simpler.worker import RemoteWorkerSpec, Worker # noqa: PLC0415 + + ports = _free_tcp_ports(2) + daemons = [ + subprocess.Popen( + [ + sys.executable, + "-m", + "simpler.remote_l3_worker", + "--host", + "127.0.0.1", + "--port", + str(port), + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for port in ports + ] + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=20) + captured: dict[str, object] = {} + try: + try: + _wait_for_tcp_ports(ports) + except TimeoutError as exc: + diagnostics = _stop_daemons(daemons) + daemons.clear() + raise TimeoutError(f"{exc}\n{diagnostics}") from exc + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{port}", + platform="a2a3sim", + device_ids=(0,), + comm_profile="sim", + ) + ) + for port in ports + ) + worker.init() + + def parent_orch(orch, _args, _cfg): + domain = orch.allocate_global_domain( + name="tcp-global", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=(CommBufferSpec("payload", "uint8", 64, 64),), + retain_after_run=True, + ) + orch.copy_to_global_domain(domain, 0, b"node-zero", buffer="payload") + orch.copy_to_global_domain(domain, 1, b"node-one", buffer="payload") + captured["ranks"] = tuple(member.global_device_rank for member in domain.members) + captured["handle"] = domain + + worker.run(parent_orch) + assert not captured["handle"].freed + + def read_orch(orch, _args, _cfg): + domain = captured["handle"] + try: + captured["rank0"] = orch.copy_from_global_domain(domain, 0, len(b"node-zero"), buffer="payload") + captured["rank1"] = orch.copy_from_global_domain(domain, 1, len(b"node-one"), buffer="payload") + finally: + domain.release() + + worker.run(read_orch) + assert captured["rank0"] == b"node-zero" + assert captured["rank1"] == b"node-one" + assert captured["ranks"] == (0, 1) + assert captured["handle"].freed + finally: + worker.close() + _stop_daemons(daemons) + + +@pytest.mark.skipif(os.name == "nt", reason="hierarchical workers require fork") +def test_local_and_remote_l3_build_and_copy_global_domain_without_mpirun(): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + from simpler.worker import RemoteWorkerSpec, Worker # noqa: PLC0415 + + (port,) = _free_tcp_ports(1) + daemon = subprocess.Popen( + [ + sys.executable, + "-m", + "simpler.remote_l3_worker", + "--host", + "127.0.0.1", + "--port", + str(port), + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=20) + captured: dict[str, object] = {} + try: + try: + _wait_for_tcp_ports((port,)) + except TimeoutError as exc: + diagnostics = _stop_daemons([daemon]) + daemon = None + raise TimeoutError(f"{exc}\n{diagnostics}") from exc + local_node_id = worker.add_worker( + Worker( + level=3, + device_ids=[0], + num_sub_workers=0, + platform="a2a3sim", + runtime="tensormap_and_ringbuffer", + comm_profile="sim", + global_device_ranks=(0,), + ) + ) + remote_node_id = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{port}", + platform="a2a3sim", + device_ids=(0,), + comm_profile="sim", + global_device_ranks=(1,), + ) + ) + worker.init() + + def build_orch(orch, _args, _cfg): + domain = orch.allocate_global_domain( + name="mixed-global", + members=((local_node_id, 0), (remote_node_id, 0)), + window_size=4096, + buffers=(CommBufferSpec("payload", "uint8", 64, 64),), + retain_after_run=True, + ) + orch.copy_to_global_domain(domain, 0, b"local-l3", buffer="payload") + orch.copy_to_global_domain(domain, 1, b"remote-l3", buffer="payload") + captured["ranks"] = tuple(member.global_device_rank for member in domain.members) + captured["domain"] = domain + + worker.run(build_orch) + domain = captured["domain"] + assert not domain.freed + + def read_orch(orch, _args, _cfg): + try: + captured["local"] = orch.copy_from_global_domain( + domain, + 0, + len(b"local-l3"), + buffer="payload", + ) + captured["remote"] = orch.copy_from_global_domain( + domain, + 1, + len(b"remote-l3"), + buffer="payload", + ) + finally: + domain.release() + + worker.run(read_orch) + assert captured["local"] == b"local-l3" + assert captured["remote"] == b"remote-l3" + assert captured["ranks"] == (0, 1) + assert domain.freed + finally: + worker.close() + if daemon is not None: + _stop_daemons([daemon]) diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index 82be5ed568..6cbbaa821d 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -941,8 +941,10 @@ def test_reap_deadline_starts_after_shutdown_broadcast(self, monkeypatch): monkeypatch.setattr(Worker, "_release_all_host_buffers", lambda self: time.sleep(0.6)) captured: dict = {} - def capture_reap(groups, deadline): + def capture_reap(groups, deadline, *, kill_survivors=False, kill_process_groups=False): captured["remaining"] = deadline - time.monotonic() + captured["kill_survivors"] = kill_survivors + captured["kill_process_groups"] = kill_process_groups monkeypatch.setattr(Worker, "_reap_child_groups", staticmethod(capture_reap)) with _hard_timeout(_TEST_WALL_BUDGET_S): @@ -950,6 +952,8 @@ def capture_reap(groups, deadline): # The reap got ~full grace (1.0s), not 1.0 - 0.6 left over from a deadline # fixed at teardown entry. assert captured["remaining"] > 0.7 + assert captured["kill_survivors"] is True + assert captured["kill_process_groups"] is True def test_reap_child_groups_stuck_child_no_starvation(self, monkeypatch): # A stuck child in one group must not starve the reap of healthy children @@ -994,6 +998,49 @@ def fake_waitpid(pid, _flags): # ...but the healthy children were reaped, freed, and removed. assert chip_pids == [] and chip_shms == [] and next_pids == [] and next_shms == [] + def test_reap_child_groups_force_kills_and_reaps_survivor(self, monkeypatch): + # A READY-tree close first gives children the full graceful budget, then + # force-kills only the survivors. Once waitpid confirms the forced exit, + # the child's mailbox can be safely freed and close need not leak/fail. + import simpler.worker as worker_mod # noqa: PLC0415 + + class _FakeShm: + def __init__(self): + self.closed = False + self.unlinked = False + + def close(self): + self.closed = True + + def unlink(self): + self.unlinked = True + + pid = 90004 + killed: list[tuple[int, int]] = [] + + def fake_kill(target, sig): + killed.append((target, sig)) + + def fake_waitpid(target, _flags): + if killed: + return (target, signal.SIGKILL) + return (0, 0) + + monkeypatch.setattr(worker_mod.os, "kill", fake_kill) + monkeypatch.setattr(worker_mod.os, "waitpid", fake_waitpid) + shm = _FakeShm() + shms, pids = [shm], [pid] + groups = [(shms, pids)] + with _hard_timeout(_TEST_WALL_BUDGET_S): + Worker._reap_child_groups( + groups, # type: ignore[arg-type] + time.monotonic() - 1.0, + kill_survivors=True, + ) + assert killed == [(pid, signal.SIGKILL)] + assert pids == [] and shms == [] + assert shm.closed and shm.unlinked + def test_non_owner_close_of_ready_raises(self, monkeypatch): # A READY worker holds same-thread-only native objects, so a close() from # a thread other than the init owner is rejected before touching the