From d3d17b21593e642c6c5f3956ab163589cdb2bbe1 Mon Sep 17 00:00:00 2001 From: sunkaixuan2018 Date: Sat, 1 Aug 2026 17:41:42 +0800 Subject: [PATCH 1/2] Add: support mpirun-launched multi-host L3 workers Add static MPI-launched L3 groups with rank-0 manifest broadcast and ready coordination, and extend Global CommDomains across mixed local, TCP-remote, and MPI-backed L3 nodes. Harden the control path with validated control names, bounded MPI collectives and L4 fanout, rollback-safe peer mappings, isolated process cleanup, fixed-port session binding compatibility, diagnostics, documentation, and regression coverage. Co-authored-by: Leaf-Salix <2503954024@qq.com> Co-authored-by: xl <21039015+xl1123@users.noreply.github.com> --- docs/comm-domain.md | 65 + docs/remote-l3-worker-design.md | 16 +- .../implementation-record.md | 16 +- docs/remote-l3-worker-design/protocol.md | 25 +- python/bindings/CMakeLists.txt | 1 + python/bindings/task_interface.cpp | 30 + python/bindings/worker_bind.h | 41 + python/simpler/global_comm_domain.py | 572 +++++ python/simpler/mpi_l3_session.py | 257 ++ python/simpler/orchestrator.py | 88 + python/simpler/remote_l3_protocol.py | 2 + python/simpler/remote_l3_session.py | 257 +- python/simpler/remote_l3_worker.py | 19 + python/simpler/task_interface.py | 139 ++ python/simpler/worker.py | 2096 ++++++++++++++++- src/a2a3/platform/onboard/host/comm_hccl.cpp | 200 ++ src/a5/platform/onboard/host/comm_hccl.cpp | 16 + src/common/hierarchical/remote_endpoint.cpp | 7 + src/common/hierarchical/remote_endpoint.h | 2 + src/common/hierarchical/remote_wire.cpp | 6 +- src/common/hierarchical/remote_wire.h | 4 + src/common/hierarchical/worker.h | 10 + src/common/hierarchical/worker_manager.cpp | 46 + src/common/hierarchical/worker_manager.h | 10 + src/common/platform_comm/comm.h | 79 + src/common/platform_comm/comm_sim.cpp | 202 ++ src/common/worker/chip_worker.cpp | 82 + src/common/worker/chip_worker.h | 16 + tests/ut/cpp/CMakeLists.txt | 1 + tests/ut/py/test_callable_identity.py | 80 + tests/ut/py/test_global_comm_domain.py | 1024 ++++++++ .../py/test_worker/test_startup_readiness.py | 49 +- 32 files changed, 5389 insertions(+), 69 deletions(-) create mode 100644 python/simpler/global_comm_domain.py create mode 100644 python/simpler/mpi_l3_session.py create mode 100644 tests/ut/py/test_global_comm_domain.py 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 From a3483130e2ccfddeb5260cef8fc3657c44b289f3 Mon Sep 17 00:00:00 2001 From: sunkaixuan2018 Date: Sat, 1 Aug 2026 18:09:17 +0800 Subject: [PATCH 2/2] Add: dispatch MPI worker groups through mailbox - Route MPI group task and control traffic through a named rank-0 mailbox - Distribute per-rank payloads and collect ranked results with MPI collectives - Preserve the TCP transport for ordinary non-MPI Remote L3 workers - Add protocol, lifecycle, timeout, compatibility, and smoke coverage --- docs/mpi-l3-mailbox.md | 132 +++++ docs/remote-l3-worker-design.md | 5 + .../implementation-record.md | 4 +- mkdocs.yml | 1 + python/bindings/worker_bind.h | 14 + python/simpler/global_comm_domain.py | 5 +- python/simpler/global_comm_smoke.py | 167 ++++++ python/simpler/mpi_group_mailbox.py | 507 +++++++++++++++++ python/simpler/mpi_group_smoke.py | 29 + python/simpler/mpi_l3_session.py | 516 +++++++++++++++--- python/simpler/orchestrator.py | 4 + python/simpler/task_interface.py | 2 +- python/simpler/worker.py | 467 ++++++---------- src/a2a3/platform/onboard/host/comm_hccl.cpp | 9 +- src/common/hierarchical/mpi_group_mailbox.h | 76 +++ src/common/hierarchical/remote_endpoint.cpp | 445 ++++++++++++++- src/common/hierarchical/remote_endpoint.h | 82 ++- src/common/hierarchical/worker.cpp | 20 + src/common/hierarchical/worker.h | 5 + src/common/hierarchical/worker_manager.h | 1 + src/common/platform_comm/comm_sim.cpp | 5 + .../cpp/hierarchical/test_remote_endpoint.cpp | 94 ++++ tests/ut/py/test_global_comm_domain.py | 85 +-- tests/ut/py/test_mpi_group_mailbox.py | 227 ++++++++ tests/ut/py/test_mpi_l3_group.py | 144 +++++ tests/ut/py/test_remote_l3_lifecycle.py | 10 +- tools/a3_l4_tcp_smoke/README.md | 29 + .../kernels/aiv/global_tload_kernel.cpp | 86 +++ .../kernels/aiv/local_add_kernel.cpp | 61 +++ .../orchestration/global_tload_orch.cpp | 35 ++ .../kernels/orchestration/local_add_orch.cpp | 34 ++ .../mpirun_compute_then_tload_2x2_smoke.py | 326 +++++++++++ tools/mpi_group_mailbox_smoke.py | 222 ++++++++ tools/mpi_l3_group_smoke.py | 143 +++++ 34 files changed, 3580 insertions(+), 412 deletions(-) create mode 100644 docs/mpi-l3-mailbox.md create mode 100644 python/simpler/global_comm_smoke.py create mode 100644 python/simpler/mpi_group_mailbox.py create mode 100644 python/simpler/mpi_group_smoke.py create mode 100644 src/common/hierarchical/mpi_group_mailbox.h create mode 100644 tests/ut/py/test_mpi_group_mailbox.py create mode 100644 tests/ut/py/test_mpi_l3_group.py create mode 100644 tools/a3_l4_tcp_smoke/README.md create mode 100644 tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp create mode 100644 tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py create mode 100644 tools/mpi_group_mailbox_smoke.py create mode 100644 tools/mpi_l3_group_smoke.py diff --git a/docs/mpi-l3-mailbox.md b/docs/mpi-l3-mailbox.md new file mode 100644 index 0000000000..bf45d62855 --- /dev/null +++ b/docs/mpi-l3-mailbox.md @@ -0,0 +1,132 @@ +# MPI L3 group mailbox protocol + +An MPI L3 group has one L4-owned, named shared-memory mailbox. Only local MPI +rank 0 opens that mailbox. All ranks participate in the same ordered +`dispatch_comm` collectives, while Global CommDomain descriptor exchange uses +a separate `domain_comm`. + +```text +L4 / MpiGroupMailboxEndpoint + | + | named SharedMemory (one request lane) + v +local MPI rank 0 / L3 + | + | dispatch_comm Bcast + Gather + v +all MPI ranks / L3 -> each rank's local L2 workers +``` + +MPI groups never start or connect Simpler command/health TCP sockets. Ordinary +`RemoteWorkerSpec` workers still use `RemoteL3SocketTransport` and keep their +existing command and health lanes. + +## Startup and shutdown + +1. L4 creates the mailbox and writes its name, protocol version, size, and + world size into the group manifest. +2. L4 starts `mpirun` as a new process group and monitors that direct child. +3. Each rank creates and initializes its own L3 `Worker`. +4. All ranks complete a readiness `allgather`. +5. Rank 0 reopens the mailbox by name and publishes `READY`. Other ranks never + map it. +6. L4 attaches every stable MPI worker id to the same + `MpiGroupMailboxChannel`; it does not create `RemoteL3Endpoint` sockets. +7. Shutdown is a mailbox `SHUTDOWN` request, followed by one MPI broadcast. + Each rank closes its inner worker, communicators are freed, `mpirun` exits, + and L4 unlinks the mailbox and manifest directory. + +If startup, a collective, the mailbox, or `mpirun` fails, the group becomes +terminal. Runtime timeout also kills the complete `mpirun` process group. +There is no TCP fallback. + +## Envelope and state + +Protocol version 1 has a fixed 256-byte header and two 16 MiB payload regions. +The header contains: + +- magic `SMPIBOX\0` +- protocol version and layout size +- MPI world size +- group state: `INITIALIZING`, `READY`, `TERMINAL`, or `CLOSED` +- request state: `IDLE`, `REQUEST_READY`, `TASK_ACCEPTED`, `TASK_DONE`, + `TASK_FAILED`, `SHUTDOWN_READY`, or `SHUTDOWN_DONE` +- monotonic mailbox `sequence_id` +- opcode: `TASK`, `CONTROL`, `PING`, or `SHUTDOWN` +- target: `GROUP`, `RANK`, or `PER_RANK` +- target rank, payload count, and byte lengths + +Rank 0 copies the complete request to private memory before publishing +`TASK_ACCEPTED`. It publishes `TASK_DONE` only after gathering every rank's +status. Duplicate or decreasing sequence ids make the group terminal. + +Every gathered error contains `rank`, `error_type`, and `message`. Any target +rank failure fails the group operation. A broken command processor or +collective is terminal; an ordinary task/control application error is returned +to L4 and the communicator may be reused. + +## Target and API semantics + +- `orch.submit_next_level(..., worker=id)` remains a directed rank operation. + Every MPI rank receives the envelope in collective order, but only the + selected rank executes it. +- `orch.submit_next_level_group(args_list, workers=...)` remains one DAG node. + When `workers` is the complete MPI group, C++ batches all members into one + `PER_RANK` mailbox request. Rank `workers[i]` uses `args_list[i]`. +- A subset group remains supported as ordered directed requests. It is not + silently widened to the complete MPI group. +- Group-wide controls use one `GROUP` mailbox request. + +The existing remote task codec is reused. It serializes scalar values, tensor +metadata, inline host payloads, and `RemoteTensorRef` descriptors. Bare host or +child virtual addresses without a valid remote sidecar are rejected before +execution; a pointer value is never forwarded as if it were meaningful on +another rank. `PYTHON_SERIALIZED` callable payloads remain unsupported by the +underlying Remote L3 protocol; `PYTHON_IMPORT` and inline `CHIP_CALLABLE` +registration are supported. + +## Remote protocol audit and MPI mapping + +The wire `FrameType` values remain unchanged: + +| Existing frame | MPI mailbox mapping | +| -------------- | ------------------- | +| `HELLO` / ready | rank-local initialization, readiness `allgather`, then rank 0 publishes mailbox `READY` | +| `TASK` | `TASK`; directed `RANK`, or one full-group `PER_RANK` vector | +| `CONTROL` / `CONTROL_REPLY` | `CONTROL`; directed except the group-wide controls below | +| `COMPLETION` | gathered per-rank status; selected/per-rank replies returned to L4 | +| `HEALTH` | `PING` to `GROUP`, gathered before success | +| `SHUTDOWN` | `SHUTDOWN` to `GROUP`, gathered before `SHUTDOWN_DONE` | + +All existing remote controls use the mailbox path: + +| Number | Control | MPI target | +| -----: | ------- | ---------- | +| 1 | `UNREGISTER_CALLABLE` | directed rank | +| 2 | `PREPARE_REGISTER_CALLABLE` | directed rank | +| 3 | `COMMIT_REGISTER_CALLABLE` | directed rank | +| 4 | `ABORT_REGISTER_CALLABLE` | directed rank | +| 5 | `PREPARE_CALLABLE` | directed rank | +| 6 | `ALLOC_REMOTE_BUFFER` | directed rank | +| 7 | `FREE_REMOTE_BUFFER` | directed rank | +| 8 | `COPY_TO_REMOTE` | directed rank | +| 9 | `COPY_FROM_REMOTE` | directed rank | +| 10 | `EXPORT_BUFFER` | directed rank | +| 11 | `IMPORT_BUFFER` | directed rank | +| 12 | `RELEASE_IMPORT` | directed rank | +| 13 | `COMM_INIT` | directed rank | +| 14 | `ALLOC_DOMAIN` prepare/import/commit/abort | one group request; descriptor work uses `domain_comm` | +| 15 | `RELEASE_DOMAIN` | one group request | +| 16 | `COPY_TO_DOMAIN` | directed rank | +| 17 | `COPY_FROM_DOMAIN` | directed rank | + +Remote control number 18 is intentionally not assigned. The local hierarchical +protocol keeps number 18 for committed-device-memory control. + +## Threading + +Only the main dispatcher thread calls MPI. The existing command processor runs +on a rank-local thread over an in-memory, socket-shaped queue. Global +CommDomain operations cross back to the dispatcher through a queue and +`threading.Event`, so they use `domain_comm` on the MPI-owning thread. This +design does not require `MPI_THREAD_MULTIPLE`. diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index 1324911b62..0024ea129f 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -9,6 +9,7 @@ Detailed protocol, buffer, transport, and rollout notes live in: - [protocol.md](remote-l3-worker-design/protocol.md) - [buffers-and-transports.md](remote-l3-worker-design/buffers-and-transports.md) +- [MPI L3 group mailbox](mpi-l3-mailbox.md) - [implementation-plan.md](remote-l3-worker-design/implementation-plan.md) - [pr-split-and-audit-plan.md][split-audit-plan] @@ -63,6 +64,10 @@ Implemented: - Socket-backed simulation remote sessions via `simpler-remote-worker` and `simpler-remote-l3-session`, including `HELLO READY`, TASK/COMPLETION, CONTROL/CONTROL_REPLY, SHUTDOWN, and an independent health lane. +- MPI L3 groups use one rank-0 named shared-memory mailbox plus ordered MPI + collectives for task, control, health, Global CommDomain, error, and shutdown + handling. They do not create Simpler command or health TCP sockets; ordinary + non-MPI Remote L3 sessions retain the socket transport. - Simulation remote buffer allocation, copy, export, import, release-import, imported-handle scheduling eligibility, and deferred owner free. - Registry-scope-aware remote callable manifest/control install for dispatcher diff --git a/docs/remote-l3-worker-design/implementation-record.md b/docs/remote-l3-worker-design/implementation-record.md index 9d1a554535..af5f4a49cd 100644 --- a/docs/remote-l3-worker-design/implementation-record.md +++ b/docs/remote-l3-worker-design/implementation-record.md @@ -92,8 +92,8 @@ It is updated as each documented feature is completed and verified. 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. + released or the Worker closes. The sim shm and `a3-fabric-v1` profiles 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. diff --git a/mkdocs.yml b/mkdocs.yml index c9ac8b3c72..bf1ba4a8a4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -129,6 +129,7 @@ nav: - L3-L2 Message Queue: l3-l2-message-queue.md - Directed NEXT_LEVEL Scheduling: directed-next-level-scheduling.md - Remote L3 Worker Design: remote-l3-worker-design.md + - MPI L3 group mailbox: mpi-l3-mailbox.md - Profiling and DFX: - Overview: dfx/README.md - Profiling Framework: dfx/profiling-framework.md diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index f3b1b968ed..8336183618 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -445,6 +445,20 @@ inline void bind_worker(nb::module_ &m) { nb::arg("health_host"), nb::arg("health_port"), nb::arg("attach_timeout_s") = 30.0, nb::arg("runtime_timeout_s") = 30.0, "Register a REMOTE_L3 endpoint after the session reports HELLO READY." ) + .def( + "add_mpi_group_mailbox", + [](Worker &self, const std::vector &worker_ids, const std::vector &session_ids, + uint64_t mailbox_ptr, size_t mailbox_bytes, int mpirun_pid, double runtime_timeout_s) { + nb::gil_scoped_release release; + self.add_mpi_group_mailbox( + worker_ids, session_ids, reinterpret_cast(mailbox_ptr), mailbox_bytes, mpirun_pid, + runtime_timeout_s + ); + }, + nb::arg("worker_ids"), nb::arg("session_ids"), nb::arg("mailbox_ptr"), nb::arg("mailbox_bytes"), + nb::arg("mpirun_pid"), nb::arg("runtime_timeout_s") = 30.0, + "Register one shared-memory MPI group endpoint for each worker id." + ) // Release the GIL while starting the Scheduler thread so another Python // thread can run during it — e.g. a concurrent close() observing diff --git a/python/simpler/global_comm_domain.py b/python/simpler/global_comm_domain.py index 039d63bd71..177ac5902a 100644 --- a/python/simpler/global_comm_domain.py +++ b/python/simpler/global_comm_domain.py @@ -16,6 +16,7 @@ GLOBAL_DOMAIN_VERSION = 1 GLOBAL_DOMAIN_MAX_RANKS = 64 +GLOBAL_DOMAIN_MAX_BUFFERS = 64 GLOBAL_DOMAIN_HANDLE_BYTES = 256 GLOBAL_DOMAIN_DESCRIPTOR = struct.Struct(" bytes: 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(command.buffers) > GLOBAL_DOMAIN_MAX_BUFFERS: + raise ValueError("global domain command buffer count exceeds maximum") 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): @@ -450,7 +453,7 @@ def decode_domain_command(data: bytes) -> GlobalDomainCommand: 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: + if buffer_count > GLOBAL_DOMAIN_MAX_BUFFERS: 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() diff --git a/python/simpler/global_comm_smoke.py b/python/simpler/global_comm_smoke.py new file mode 100644 index 0000000000..47e1d0cd78 --- /dev/null +++ b/python/simpler/global_comm_smoke.py @@ -0,0 +1,167 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Rank-local L3 callbacks used by the A3 MPI compute + TLOAD smoke.""" + +from __future__ import annotations + +from .task_interface import CallConfig, DataType, TaskArgs, Tensor, TensorArgType + +_SMOKE_COUNT = 256 + + +def _digest_from_scalars(args: TaskArgs, start: int) -> bytes: + return b"".join(int(args.scalar(start + i)).to_bytes(8, "little") for i in range(4)) + + +def remote_compute_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Submit one local vector-add task from an MPI L3 rank to its selected L2.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("remote compute task expects domain_id, local_worker_id, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_id = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + context = orch.get_global_domain(domain_id)[local_worker_id] + + chip_args = TaskArgs() + for buffer_name in ("lhs", "rhs"): + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs[buffer_name], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=local_worker_id) + + +def remote_rank_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Submit the peer-TLOAD kernel from an MPI L3 rank to its selected L2.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("global TLOAD remote task expects domain_id, local_worker_id, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_id = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + context = orch.get_global_domain(domain_id)[local_worker_id] + + chip_args = TaskArgs() + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["result"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + chip_args.add_scalar(context.domain_size) + chip_args.add_scalar(context.device_ctx) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=local_worker_id) + + +def remote_compute_group_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Drive every L2 owned by one MPI L3 rank with distinct domain buffers.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("MPI compute group task expects domain_id, local worker count, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_count = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + domain = orch.get_global_domain(domain_id) + group_args = [] + workers = [] + for local_worker_id in range(local_worker_count): + context = domain[local_worker_id] + chip_args = TaskArgs() + for buffer_name in ("lhs", "rhs"): + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs[buffer_name], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + group_args.append(chip_args) + workers.append(local_worker_id) + orch.submit_next_level_group(chip_handle, group_args, cfg, workers=workers) + + +def remote_rank_group_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Drive peer TLOAD on every L2 owned by one MPI L3 rank.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("MPI TLOAD group task expects domain_id, local worker count, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_count = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + domain = orch.get_global_domain(domain_id) + group_args = [] + workers = [] + for local_worker_id in range(local_worker_count): + context = domain[local_worker_id] + chip_args = TaskArgs() + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["result"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + chip_args.add_scalar(context.domain_size) + chip_args.add_scalar(context.device_ctx) + group_args.append(chip_args) + workers.append(local_worker_id) + orch.submit_next_level_group(chip_handle, group_args, cfg, workers=workers) diff --git a/python/simpler/mpi_group_mailbox.py b/python/simpler/mpi_group_mailbox.py new file mode 100644 index 0000000000..bb3b0fc99f --- /dev/null +++ b/python/simpler/mpi_group_mailbox.py @@ -0,0 +1,507 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Named shared-memory protocol used between L4 and an MPI L3 group.""" + +from __future__ import annotations + +import ctypes +import enum +import json +import struct +from dataclasses import asdict, dataclass +from multiprocessing.shared_memory import SharedMemory +from typing import Any + +MAILBOX_MAGIC = b"SMPIBOX\0" +MAILBOX_PROTOCOL_VERSION = 1 +MAILBOX_HEADER_BYTES = 256 +MAILBOX_PAYLOAD_BYTES = 16 * 1024 * 1024 +MAILBOX_ERROR_BYTES = 64 * 1024 +MAILBOX_REQUEST_OFFSET = MAILBOX_HEADER_BYTES +MAILBOX_RESPONSE_OFFSET = MAILBOX_REQUEST_OFFSET + MAILBOX_PAYLOAD_BYTES +MAILBOX_ERROR_OFFSET = MAILBOX_RESPONSE_OFFSET + MAILBOX_PAYLOAD_BYTES +MAILBOX_SIZE = MAILBOX_ERROR_OFFSET + MAILBOX_ERROR_BYTES + +_OFF_MAGIC = 0 +_OFF_PROTOCOL_VERSION = 8 +_OFF_HEADER_BYTES = 12 +_OFF_MAILBOX_BYTES = 16 +_OFF_WORLD_SIZE = 24 +_OFF_GROUP_STATE = 28 +_OFF_REQUEST_STATE = 32 +_OFF_SEQUENCE_ID = 40 +_OFF_OPCODE = 48 +_OFF_TARGET = 52 +_OFF_TARGET_RANK = 56 +_OFF_REQUEST_COUNT = 60 +_OFF_REQUEST_BYTES = 64 +_OFF_RESPONSE_COUNT = 68 +_OFF_RESPONSE_BYTES = 72 +_OFF_ERROR_BYTES = 76 + + +class MailboxGroupState(enum.IntEnum): + INITIALIZING = 0 + READY = 1 + TERMINAL = 2 + CLOSED = 3 + + +class MailboxRequestState(enum.IntEnum): + IDLE = 0 + REQUEST_READY = 1 + TASK_ACCEPTED = 2 + TASK_DONE = 3 + TASK_FAILED = 4 + SHUTDOWN_READY = 5 + SHUTDOWN_DONE = 6 + + +class MailboxOpcode(enum.IntEnum): + TASK = 1 + CONTROL = 2 + PING = 3 + SHUTDOWN = 4 + + +class MailboxTarget(enum.IntEnum): + GROUP = 1 + RANK = 2 + PER_RANK = 3 + + +@dataclass(frozen=True) +class MpiRankError: + rank: int + error_type: str + message: str + + +@dataclass(frozen=True) +class MailboxRequest: + sequence_id: int + opcode: MailboxOpcode + target: MailboxTarget + target_rank: int + payloads: tuple[bytes, ...] + + +@dataclass(frozen=True) +class MailboxResult: + sequence_id: int + payloads: tuple[bytes, ...] + + +class MpiGroupError(RuntimeError): + """A mailbox or MPI group operation failed.""" + + +def _encode_payloads(payloads: tuple[bytes, ...]) -> bytes: + values = tuple(bytes(payload) for payload in payloads) + prefix_bytes = 4 + 4 * len(values) + payload_bytes = sum(len(payload) for payload in values) + if prefix_bytes + payload_bytes > MAILBOX_PAYLOAD_BYTES: + raise ValueError("MPI group mailbox payload vector exceeds capacity") + out = bytearray(struct.pack(" tuple[bytes, ...]: + if len(data) < 4: + raise MpiGroupError("MPI group mailbox payload vector is truncated") + (count,) = struct.unpack_from(" len(data): + raise MpiGroupError("MPI group mailbox payload lengths are truncated") + lengths = struct.unpack_from(f"<{count}I", data, 4) if count else () + offset = prefix_bytes + payloads: list[bytes] = [] + for length in lengths: + if offset > len(data) or length > len(data) - offset: + raise MpiGroupError("MPI group mailbox payload entry is truncated") + payloads.append(bytes(data[offset : offset + length])) + offset += length + if offset != len(data): + raise MpiGroupError("MPI group mailbox payload vector has trailing bytes") + return tuple(payloads) + + +def _as_byte_view(buffer: memoryview) -> memoryview: + """Normalize platform-specific shared-memory formats to a byte view.""" + return buffer.cast("B") + + +def _buffer_address(buffer: memoryview) -> int: + return ctypes.addressof(ctypes.c_char.from_buffer(buffer)) + + +class MpiGroupMailbox: + """Owner or reopened view of one MPI group mailbox.""" + + def __init__(self, shm: SharedMemory, *, owner: bool) -> None: + self._shm = shm + self._owner = bool(owner) + self._closed = False + self._validate_header() + + @classmethod + def create(cls, *, world_size: int) -> MpiGroupMailbox: + if int(world_size) <= 0: + raise ValueError("MPI group mailbox world_size must be positive") + shm = SharedMemory(create=True, size=MAILBOX_SIZE) + buffer = shm.buf + if buffer is None: + shm.close() + shm.unlink() + raise MpiGroupError("MPI group mailbox mapping returned no buffer") + buffer = _as_byte_view(buffer) + ctypes.memset(_buffer_address(buffer), 0, MAILBOX_SIZE) + struct.pack_into( + "<8sIIQ", buffer, _OFF_MAGIC, MAILBOX_MAGIC, MAILBOX_PROTOCOL_VERSION, MAILBOX_HEADER_BYTES, MAILBOX_SIZE + ) + struct.pack_into(" MpiGroupMailbox: + try: + shm = SharedMemory(name=str(name), create=False, track=False) + except TypeError: # Python < 3.13 has no per-instance resource-tracker switch. + shm = SharedMemory(name=str(name), create=False) + # A reopened rank-0 view is not the owner. Older Python versions + # register every SharedMemory view and otherwise try to unlink it + # when rank 0 exits, racing the L4 owner and emitting leak warnings. + from multiprocessing import resource_tracker # noqa: PLC0415 + + resource_tracker.unregister(shm._name, "shared_memory") # noqa: SLF001 + return cls(shm, owner=False) + + @property + def name(self) -> str: + return self._shm.name + + @property + def address(self) -> int: + self._require_open() + return ctypes.addressof(ctypes.c_char.from_buffer(self._buffer)) + + @property + def world_size(self) -> int: + return self._read_u32(_OFF_WORLD_SIZE) + + @property + def group_state(self) -> MailboxGroupState: + return MailboxGroupState(self._load_i32(_OFF_GROUP_STATE)) + + @property + def request_state(self) -> MailboxRequestState: + return MailboxRequestState(self._load_i32(_OFF_REQUEST_STATE)) + + def manifest(self) -> dict[str, Any]: + return { + "name": self.name, + "protocol_version": MAILBOX_PROTOCOL_VERSION, + "mailbox_bytes": MAILBOX_SIZE, + "world_size": self.world_size, + } + + def publish_ready(self) -> None: + if self.group_state is not MailboxGroupState.INITIALIZING: + raise MpiGroupError("MPI group mailbox READY can only be published from INITIALIZING") + self._store_i32(_OFF_GROUP_STATE, int(MailboxGroupState.READY)) + + def publish_closed(self) -> None: + if self.group_state is MailboxGroupState.TERMINAL: + return + self._store_i32(_OFF_GROUP_STATE, int(MailboxGroupState.CLOSED)) + + def write_request( + self, + *, + sequence_id: int, + opcode: MailboxOpcode, + target: MailboxTarget, + target_rank: int, + payloads: tuple[bytes, ...], + ) -> None: + self._require_ready() + if self.request_state is not MailboxRequestState.IDLE: + raise MpiGroupError(f"MPI group mailbox is busy in state {self.request_state.name}") + sequence_id = int(sequence_id) + if sequence_id <= 0: + raise ValueError("MPI group mailbox sequence_id must be positive") + opcode = MailboxOpcode(opcode) + target = MailboxTarget(target) + target_rank = int(target_rank) + payloads = tuple(bytes(payload) for payload in payloads) + if target is MailboxTarget.RANK: + if target_rank < 0 or target_rank >= self.world_size: + raise ValueError("MPI group mailbox target rank is outside the group") + if len(payloads) != 1: + raise ValueError("rank-targeted MPI group request requires one payload") + elif target is MailboxTarget.GROUP: + if target_rank != -1 or len(payloads) != 1: + raise ValueError("group-targeted MPI request requires target_rank=-1 and one payload") + elif target is MailboxTarget.PER_RANK: + if target_rank != -1 or len(payloads) != self.world_size: + raise ValueError("per-rank MPI request requires one payload for every rank") + data = _encode_payloads(payloads) + self._write_bytes(MAILBOX_REQUEST_OFFSET, data) + self._write_u64(_OFF_SEQUENCE_ID, sequence_id) + self._write_u32(_OFF_OPCODE, int(opcode)) + self._write_u32(_OFF_TARGET, int(target)) + self._write_i32(_OFF_TARGET_RANK, target_rank) + self._write_u32(_OFF_REQUEST_COUNT, len(payloads)) + self._write_u32(_OFF_REQUEST_BYTES, len(data)) + self._write_u32(_OFF_RESPONSE_COUNT, 0) + self._write_u32(_OFF_RESPONSE_BYTES, 0) + self._write_u32(_OFF_ERROR_BYTES, 0) + ready_state = ( + MailboxRequestState.SHUTDOWN_READY + if opcode is MailboxOpcode.SHUTDOWN + else MailboxRequestState.REQUEST_READY + ) + self._store_i32(_OFF_REQUEST_STATE, int(ready_state)) + + def accept_request(self, *, last_sequence_id: int) -> MailboxRequest: + state = self.request_state + if state not in (MailboxRequestState.REQUEST_READY, MailboxRequestState.SHUTDOWN_READY): + raise MpiGroupError(f"MPI group mailbox has no request to accept (state={state.name})") + sequence_id = self._read_u64(_OFF_SEQUENCE_ID) + request_bytes = self._read_u32(_OFF_REQUEST_BYTES) + request_count = self._read_u32(_OFF_REQUEST_COUNT) + if request_bytes > MAILBOX_PAYLOAD_BYTES: + self.mark_terminal("request payload exceeds mailbox capacity") + raise MpiGroupError("MPI group mailbox request payload exceeds capacity") + data = self._read_bytes(MAILBOX_REQUEST_OFFSET, request_bytes) + try: + if sequence_id <= int(last_sequence_id): + raise MpiGroupError( + f"MPI group mailbox sequence_id {sequence_id} is not newer than {int(last_sequence_id)}" + ) + request = MailboxRequest( + sequence_id=sequence_id, + opcode=MailboxOpcode(self._read_u32(_OFF_OPCODE)), + target=MailboxTarget(self._read_u32(_OFF_TARGET)), + target_rank=self._read_i32(_OFF_TARGET_RANK), + payloads=_decode_payloads(data, request_count), + ) + except BaseException as exc: + self.mark_terminal(str(exc)) + raise + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.TASK_ACCEPTED)) + return request + + def complete_request(self, *, sequence_id: int, payloads: tuple[bytes, ...]) -> None: + self._validate_active_sequence(sequence_id) + values = tuple(bytes(payload) for payload in payloads) + data = _encode_payloads(values) + self._write_bytes(MAILBOX_RESPONSE_OFFSET, data) + self._write_u32(_OFF_RESPONSE_COUNT, len(values)) + self._write_u32(_OFF_RESPONSE_BYTES, len(data)) + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.TASK_DONE)) + + def complete_shutdown(self, *, sequence_id: int) -> None: + self._validate_active_sequence(sequence_id) + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.SHUTDOWN_DONE)) + + def fail_request( + self, + *, + sequence_id: int, + errors: tuple[MpiRankError, ...], + terminal: bool, + ) -> None: + self._validate_active_sequence(sequence_id) + if not errors: + raise ValueError("MPI group mailbox failure requires at least one rank error") + data = json.dumps([asdict(error) for error in errors], sort_keys=True).encode("utf-8") + if len(data) > MAILBOX_ERROR_BYTES: + data = data[: MAILBOX_ERROR_BYTES - 1] + self._write_bytes(MAILBOX_ERROR_OFFSET, data) + self._write_u32(_OFF_ERROR_BYTES, len(data)) + if terminal: + self._store_i32(_OFF_GROUP_STATE, int(MailboxGroupState.TERMINAL)) + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.TASK_FAILED)) + + def read_result(self, *, sequence_id: int) -> MailboxResult: + if self._read_u64(_OFF_SEQUENCE_ID) != int(sequence_id): + raise MpiGroupError("MPI group mailbox result sequence does not match the request") + state = self.request_state + if state is MailboxRequestState.TASK_DONE: + response_bytes = self._read_u32(_OFF_RESPONSE_BYTES) + response_count = self._read_u32(_OFF_RESPONSE_COUNT) + if response_bytes > MAILBOX_PAYLOAD_BYTES: + self.mark_terminal("response payload exceeds mailbox capacity") + raise MpiGroupError("MPI group mailbox response payload exceeds capacity") + data = self._read_bytes(MAILBOX_RESPONSE_OFFSET, response_bytes) + result = MailboxResult(int(sequence_id), _decode_payloads(data, response_count)) + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.IDLE)) + return result + if state is MailboxRequestState.TASK_FAILED: + error_bytes = min(self._read_u32(_OFF_ERROR_BYTES), MAILBOX_ERROR_BYTES) + raw = self._read_bytes(MAILBOX_ERROR_OFFSET, error_bytes) + try: + entries = json.loads(raw.decode("utf-8")) + message = "; ".join( + f"rank {int(entry['rank'])}: {entry['error_type']}: {entry['message']}" for entry in entries + ) + except BaseException: + message = raw.decode("utf-8", errors="replace") or "MPI group request failed" + if self.group_state is MailboxGroupState.READY: + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.IDLE)) + raise MpiGroupError(message) + raise MpiGroupError(f"MPI group mailbox result is not ready (state={state.name})") + + def mark_terminal(self, reason: str) -> None: + data = str(reason).encode("utf-8")[:MAILBOX_ERROR_BYTES] + self._write_bytes(MAILBOX_ERROR_OFFSET, data) + self._write_u32(_OFF_ERROR_BYTES, len(data)) + self._store_i32(_OFF_GROUP_STATE, int(MailboxGroupState.TERMINAL)) + self._store_i32(_OFF_REQUEST_STATE, int(MailboxRequestState.TASK_FAILED)) + + def terminal_reason(self) -> str: + error_bytes = min(self._read_u32(_OFF_ERROR_BYTES), MAILBOX_ERROR_BYTES) + return self._read_bytes(MAILBOX_ERROR_OFFSET, error_bytes).decode("utf-8", errors="replace") + + def overwrite_request_payload_for_test(self, data: bytes) -> None: + value = bytes(data) + self._write_bytes(MAILBOX_REQUEST_OFFSET, value) + + def close(self, *, unlink: bool = False) -> None: + if self._closed: + return + if unlink and not self._owner: + raise RuntimeError("only the MPI group mailbox owner may unlink it") + self._shm.close() + if unlink: + try: + self._shm.unlink() + except FileNotFoundError: + # Python < 3.13 may let a non-owner process resource tracker + # unlink the name first; the owner mapping still closes here. + pass + self._closed = True + + def _validate_header(self) -> None: + if len(self._buffer) < MAILBOX_SIZE: + raise MpiGroupError("MPI group mailbox is smaller than the protocol layout") + magic, version, header_bytes, mailbox_bytes = struct.unpack_from("<8sIIQ", self._buffer, _OFF_MAGIC) + if magic != MAILBOX_MAGIC: + raise MpiGroupError("MPI group mailbox magic does not match") + if version != MAILBOX_PROTOCOL_VERSION: + raise MpiGroupError(f"MPI group mailbox protocol version {version} is not supported") + if header_bytes != MAILBOX_HEADER_BYTES or mailbox_bytes != MAILBOX_SIZE: + raise MpiGroupError("MPI group mailbox layout does not match the protocol") + if self._read_u32(_OFF_WORLD_SIZE) == 0: + raise MpiGroupError("MPI group mailbox world_size must be positive") + + def _validate_active_sequence(self, sequence_id: int) -> None: + if self.request_state is not MailboxRequestState.TASK_ACCEPTED: + raise MpiGroupError("MPI group mailbox request has not been accepted") + if self._read_u64(_OFF_SEQUENCE_ID) != int(sequence_id): + raise MpiGroupError("MPI group mailbox active sequence does not match") + + def _require_ready(self) -> None: + state = self.group_state + if state is MailboxGroupState.TERMINAL: + reason = self.terminal_reason() or "terminal failure" + raise MpiGroupError(f"MPI group mailbox is terminal: {reason}") + if state is not MailboxGroupState.READY: + raise MpiGroupError(f"MPI group mailbox is not ready (state={state.name})") + + def _require_open(self) -> None: + if self._closed: + raise RuntimeError("MPI group mailbox is closed") + + @property + def _buffer(self) -> memoryview: + buffer = self._shm.buf + if buffer is None: + raise RuntimeError("MPI group mailbox mapping has no buffer") + return _as_byte_view(buffer) + + def _write_bytes(self, offset: int, data: bytes) -> None: + self._require_open() + offset = int(offset) + buffer = self._buffer + if offset < 0 or offset > len(buffer) or len(data) > len(buffer) - offset: + raise ValueError("MPI group mailbox write is outside the mapping") + if data: + ctypes.memmove(_buffer_address(buffer) + offset, data, len(data)) + + def _read_bytes(self, offset: int, length: int) -> bytes: + self._require_open() + offset = int(offset) + length = int(length) + buffer = self._buffer + if offset < 0 or length < 0 or offset > len(buffer) or length > len(buffer) - offset: + raise ValueError("MPI group mailbox read is outside the mapping") + if length == 0: + return b"" + return ctypes.string_at(_buffer_address(buffer) + offset, length) + + def _load_i32(self, offset: int) -> int: + self._require_open() + try: + from .task_interface import _mailbox_load_i32 # noqa: PLC0415 + except (ImportError, AttributeError): + return self._read_i32(offset) + return int(_mailbox_load_i32(self.address + offset)) + + def _store_i32(self, offset: int, value: int) -> None: + self._require_open() + try: + from .task_interface import _mailbox_store_i32 # noqa: PLC0415 + except (ImportError, AttributeError): + self._write_i32(offset, value) + return + _mailbox_store_i32(self.address + offset, int(value)) + + def _read_i32(self, offset: int) -> int: + return int(struct.unpack_from(" None: + struct.pack_into(" int: + return int(struct.unpack_from(" None: + struct.pack_into(" int: + return int(struct.unpack_from(" None: + struct.pack_into(" MpiGroupMailbox | None: + """Open the L4 mailbox on rank 0; all other ranks stay MPI-only.""" + + if int(rank) != 0: + return None + if int(manifest.get("protocol_version", -1)) != MAILBOX_PROTOCOL_VERSION: + raise MpiGroupError("MPI group manifest mailbox protocol version does not match") + if int(manifest.get("mailbox_bytes", -1)) != MAILBOX_SIZE: + raise MpiGroupError("MPI group manifest mailbox size does not match") + mailbox = MpiGroupMailbox.open(name=str(manifest["name"])) + if mailbox.world_size != int(manifest.get("world_size", -1)): + mailbox.close() + raise MpiGroupError("MPI group manifest world_size does not match the mailbox") + return mailbox diff --git a/python/simpler/mpi_group_smoke.py b/python/simpler/mpi_group_smoke.py new file mode 100644 index 0000000000..3f1768d685 --- /dev/null +++ b/python/simpler/mpi_group_smoke.py @@ -0,0 +1,29 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Device-free RemoteCallable callbacks for the MPI L3 integration smoke.""" + +from __future__ import annotations + +import json +import os + +from .task_interface import CallConfig, TaskArgs + + +def record_rank_value(_orch, args: TaskArgs, _config: CallConfig) -> None: + if args.scalar_count() != 1: + raise ValueError("MPI group smoke callback expects one scalar") + rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) + value = int(args.scalar(0)) + if value == 0xFFFF: + raise ValueError(f"injected callback failure on MPI rank {rank}") + output_dir = os.environ["SIMPLER_MPI_SMOKE_DIR"] + path = os.path.join(output_dir, f"rank-{rank}.json") + with open(path, "w", encoding="utf-8") as output_file: + json.dump({"rank": rank, "value": value}, output_file, sort_keys=True) diff --git a/python/simpler/mpi_l3_session.py b/python/simpler/mpi_l3_session.py index 82581f314d..44ad73b5a9 100644 --- a/python/simpler/mpi_l3_session.py +++ b/python/simpler/mpi_l3_session.py @@ -6,12 +6,7 @@ # 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. -""" +"""MPI L3 group runner driven by one rank-0 named shared-memory mailbox.""" from __future__ import annotations @@ -19,11 +14,13 @@ import json import math import os +import queue import signal -import socket import struct import sys +import threading import time +from dataclasses import dataclass from typing import Any from .global_comm_domain import ( @@ -34,12 +31,48 @@ encode_descriptor_table, validate_descriptor_table, ) -from .remote_l3_session import _format_remote_error, run_session +from .mpi_group_mailbox import ( + MailboxGroupState, + MailboxOpcode, + MailboxRequest, + MailboxRequestState, + MailboxTarget, + MpiGroupMailbox, + MpiRankError, + open_rank_mailbox, +) +from .remote_l3_protocol import ( + FrameHeader, + FrameType, + decode_frame, + encode_frame, +) +from .remote_l3_session import ( + _format_remote_error, + _install_manifest_dispatcher_registry, + _install_manifest_inner_registry, + _run_command_loop, + _startup_deadline, +) from .worker import Worker +_POLL_WAIT = threading.Event() + + +def _cooperative_poll_backoff(spins: int) -> int: + """Bound a hot poll with a scheduler yield, without sleep-based polling.""" + spins += 1 + if spins < 64: + return spins + if hasattr(os, "sched_yield"): + os.sched_yield() + else: + _POLL_WAIT.wait(0.0002) + return 0 + class MpiGlobalDomainExchange: - """Descriptor exchange hook for a full mpirun-launched L3 group.""" + """Run Global CommDomain collectives on the dispatcher thread.""" def __init__(self, comm: Any, *, group_worker_ids: tuple[int, ...], timeout_s: float) -> None: self._comm = comm @@ -51,19 +84,12 @@ def __init__(self, comm: Any, *, group_worker_ids: tuple[int, ...], timeout_s: f 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") + # The mpi4py version used by the deployment has no non-blocking pickle + # convenience named ``iallgather``. ``Iallgather`` is the buffer API, + # while lowercase ``allgather`` pickles Python objects and returns a + # list. The outer L4 mailbox timeout remains the group-level watchdog + # and terminates the complete mpirun process group if this call stalls. + return list(self._comm.allgather(payload)) def prepare_import(self, command: GlobalDomainCommand, inner_worker: Worker, worker_id: int) -> bytes | None: if command.phase is not GlobalDomainPhase.PREPARE_EXPORT: @@ -162,24 +188,87 @@ def release_local() -> None: 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) +@dataclass +class _DomainRequest: + exchange: MpiGlobalDomainExchange + command: GlobalDomainCommand + inner_worker: Worker + worker_id: int + done: threading.Event + result: bytes | None = None + error: BaseException | None = None + def execute(self) -> None: + try: + self.result = self.exchange.prepare_import(self.command, self.inner_worker, self.worker_id) + except BaseException as exc: # noqa: BLE001 + self.error = exc + finally: + self.done.set() -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(" None: + self._events = events + self._exchange = exchange -def _raise_keyboard_interrupt(_signum, _frame): - raise KeyboardInterrupt + def prepare_import(self, command: GlobalDomainCommand, inner_worker: Worker, worker_id: int) -> bytes | None: + request = _DomainRequest( + self._exchange, + command, + inner_worker, + int(worker_id), + threading.Event(), + ) + self._events.put(("domain", request)) + request.done.wait() + if request.error is not None: + raise request.error + return request.result + + +class _InMemoryCommandConnection: + """Socket-shaped blocking stream backed by queues, with no OS socket.""" + + def __init__(self) -> None: + self._incoming: queue.Queue[bytes] = queue.Queue() + self._events: queue.Queue[tuple[str, Any]] = queue.Queue() + self._recv_buffer = bytearray() + + @property + def events(self) -> queue.Queue[tuple[str, Any]]: + return self._events + + def recv(self, nbytes: int) -> bytes: + while not self._recv_buffer: + self._recv_buffer.extend(self._incoming.get()) + count = min(int(nbytes), len(self._recv_buffer)) + data = bytes(self._recv_buffer[:count]) + del self._recv_buffer[:count] + return data + + def sendall(self, data: bytes) -> None: + self._events.put(("reply", bytes(data))) + + def feed(self, frame: bytes) -> None: + self._incoming.put(bytes(frame)) + + def next_reply(self) -> bytes: + while True: + kind, value = self._events.get() + if kind == "domain": + value.execute() + continue + if kind == "processor_error": + error_type, message = value + raise RuntimeError(f"{error_type}: {message}") + if kind != "reply": + raise RuntimeError(f"MPI command processor emitted unknown event {kind!r}") + return bytes(value) + + def exchange(self, frame: bytes) -> bytes: + self.feed(frame) + return self.next_reply() def _load_group_manifest_from_rank0(comm: Any, manifest_path: str) -> dict[str, Any]: @@ -200,6 +289,309 @@ def _load_group_manifest_from_rank0(comm: Any, manifest_path: str) -> dict[str, return value +def _rewrite_frame_identity( + frame_bytes: bytes, + manifest: dict[str, Any], + *, + sequence: int | None = None, +) -> bytes: + frame = decode_frame(frame_bytes) + header = FrameHeader( + frame_type=frame.header.frame_type, + session_id=int(manifest["session_id"]), + worker_id=int(manifest["worker_id"]), + sequence=frame.header.sequence if sequence is None else int(sequence), + flags=frame.header.flags, + ) + return encode_frame(header, frame.payload) + + +def _reply_error(reply_bytes: bytes) -> tuple[str, str] | None: + frame = decode_frame(reply_bytes) + if frame.header.frame_type is FrameType.COMPLETION: + if len(frame.payload) < 16: + return ("ProtocolError", "truncated completion reply") + _sequence, error_code, message_bytes = struct.unpack_from(" len(frame.payload) - offset: + return ("ProtocolError", "reply error message is truncated") + if error_code == 0: + return None + message = frame.payload[offset : offset + message_bytes].decode("utf-8", errors="replace") + return ("RemoteOperationError", message) + + +def _payload_for_rank(request: MailboxRequest, rank: int) -> bytes | None: + if request.target is MailboxTarget.GROUP: + return request.payloads[0] + if request.target is MailboxTarget.RANK: + return request.payloads[0] if rank == request.target_rank else None + if request.target is MailboxTarget.PER_RANK: + return request.payloads[rank] + raise ValueError(f"unsupported MPI mailbox target {request.target}") + + +def _shutdown_frame(manifest: dict[str, Any]) -> bytes: + return encode_frame( + FrameHeader( + frame_type=FrameType.SHUTDOWN, + session_id=int(manifest["session_id"]), + worker_id=int(manifest["worker_id"]), + sequence=0, + ), + b"", + ) + + +def _select_mailbox_replies( + request: MailboxRequest, + gathered: list[tuple[int, bool, bytes, MpiRankError | None]], + worker_ids: tuple[int, ...], +) -> tuple[bytes, ...]: + def _for_request(reply: bytes, request_frame: bytes) -> bytes: + if not reply: + return b"" + decoded_reply = decode_frame(reply) + decoded_request = decode_frame(request_frame) + reply_payload = bytearray(decoded_reply.payload) + if decoded_reply.header.frame_type in (FrameType.COMPLETION, FrameType.CONTROL_REPLY): + if len(reply_payload) < 8: + raise RuntimeError("MPI rank reply is truncated before its sequence field") + struct.pack_into(" int: + rank = int(dispatch_comm.Get_rank()) + worker_ids = tuple(int(value) for value in group_manifest["worker_ids"]) + mailbox: MpiGroupMailbox | None = None + inner_worker = Worker( + level=3, + platform=str(manifest["platform"]), + runtime=str(manifest.get("runtime", "tensormap_and_ringbuffer")), + device_ids=tuple(int(value) for value in manifest.get("device_ids", ())), + num_sub_workers=int(manifest.get("num_sub_workers", 0)), + heap_ring_size=int(manifest["heap_ring_size"]) if manifest.get("heap_ring_size") is not None else None, + ) + connection = _InMemoryCommandConnection() + processor_thread: threading.Thread | None = None + startup_ok = True + startup_error = "" + + try: + dispatch_registry = _install_manifest_dispatcher_registry(manifest) + inner_handles = _install_manifest_inner_registry(manifest, inner_worker) + inner_worker.init(_startup_deadline=_startup_deadline(manifest)) + exchange = MpiGlobalDomainExchange( + domain_comm, + group_worker_ids=worker_ids, + timeout_s=float(manifest["session_timeout_s"]), + ) + bridge = _DomainBridge(connection.events, exchange) + + def _processor() -> None: + try: + _run_command_loop( + connection, # type: ignore[arg-type] + manifest, + inner_worker, + inner_handles, + dispatch_registry, + bridge.prepare_import, + ) + except BaseException as exc: # noqa: BLE001 + connection.events.put(("processor_error", (type(exc).__name__, str(exc)))) + + processor_thread = threading.Thread(target=_processor, name=f"simpler-mpi-command-rank-{rank}") + processor_thread.start() + hello = decode_frame(connection.next_reply()) + if hello.header.frame_type is not FrameType.HELLO: + raise RuntimeError("MPI command processor did not publish HELLO") + except BaseException as exc: # noqa: BLE001 + startup_ok = False + startup_error = _format_remote_error(f"MPI rank {rank} startup", exc) + + readiness = dispatch_comm.allgather((rank, startup_ok, startup_error)) + mailbox_ready = True + mailbox_error = "" + if rank == 0: + try: + mailbox = open_rank_mailbox(group_manifest["mailbox"], rank=rank) + assert mailbox is not None + failures = [(ready_rank, error) for ready_rank, ok, error in readiness if not ok] + if failures: + mailbox_error = "; ".join(f"rank {ready_rank}: {error}" for ready_rank, error in failures) + mailbox.mark_terminal(mailbox_error) + mailbox_ready = False + else: + mailbox.publish_ready() + except BaseException as exc: # noqa: BLE001 + mailbox_ready = False + mailbox_error = _format_remote_error("MPI rank 0 mailbox startup", exc) + if mailbox is not None: + mailbox.mark_terminal(mailbox_error) + mailbox_ready, mailbox_error = dispatch_comm.bcast((mailbox_ready, mailbox_error), root=0) + if not mailbox_ready: + if processor_thread is not None: + connection.feed(_shutdown_frame(manifest)) + processor_thread.join(timeout=1.0) + inner_worker.close() + if mailbox is not None: + mailbox.close() + return 1 + + last_sequence_id = 0 + local_command_sequence = 0 + session_timeout_s = float(manifest["session_timeout_s"]) + exit_code = 0 + try: + while True: + if rank == 0: + assert mailbox is not None + request_wait_deadline = time.monotonic() + session_timeout_s + spins = 0 + while True: + request_state = mailbox.request_state + if request_state in ( + MailboxRequestState.REQUEST_READY, + MailboxRequestState.SHUTDOWN_READY, + ): + break + if mailbox.group_state is MailboxGroupState.TERMINAL: + break + now = time.monotonic() + if request_state is MailboxRequestState.IDLE: + request_wait_deadline = now + session_timeout_s + elif now >= request_wait_deadline: + mailbox.mark_terminal(f"MPI group mailbox request lane stalled in state {request_state.name}") + break + spins = _cooperative_poll_backoff(spins) + if mailbox.group_state is MailboxGroupState.TERMINAL: + request = None + else: + try: + request = mailbox.accept_request(last_sequence_id=last_sequence_id) + last_sequence_id = request.sequence_id + except Exception as exc: + mailbox.mark_terminal(_format_remote_error("MPI rank 0 mailbox request", exc)) + request = None + else: + request = None + request = dispatch_comm.bcast(request, root=0) + if request is None: + exit_code = 1 + break + + local_reply = b"" + local_error: MpiRankError | None = None + payload = None + try: + payload = _payload_for_rank(request, rank) + if request.opcode is MailboxOpcode.PING: + local_reply = b"" + elif request.opcode is MailboxOpcode.SHUTDOWN: + assert payload is not None + connection.feed(_rewrite_frame_identity(payload, manifest)) + elif payload is not None: + local_command_sequence += 1 + local_reply = connection.exchange( + _rewrite_frame_identity(payload, manifest, sequence=local_command_sequence) + ) + except BaseException as exc: # noqa: BLE001 + local_error = MpiRankError(rank, type(exc).__name__, str(exc)) + + gathered = dispatch_comm.gather((rank, payload is not None, local_reply, local_error), root=0) + terminal_dispatch_failure = False + if rank == 0: + assert mailbox is not None and gathered is not None + transport_errors = tuple(item[3] for item in gathered if item[3] is not None) + application_errors: list[MpiRankError] = [] + for result_rank, executed, reply, _error in gathered: + if not executed or not reply: + continue + decoded_error = _reply_error(reply) + if decoded_error is not None: + error_type, message = decoded_error + application_errors.append(MpiRankError(int(result_rank), error_type, message)) + if transport_errors: + mailbox.fail_request( + sequence_id=request.sequence_id, + errors=transport_errors, + terminal=True, + ) + exit_code = 1 + terminal_dispatch_failure = True + elif application_errors: + mailbox.fail_request( + sequence_id=request.sequence_id, + errors=tuple(application_errors), + terminal=False, + ) + elif request.opcode is MailboxOpcode.SHUTDOWN: + mailbox.complete_shutdown(sequence_id=request.sequence_id) + mailbox.publish_closed() + else: + replies = _select_mailbox_replies(request, gathered, worker_ids) + mailbox.complete_request(sequence_id=request.sequence_id, payloads=replies) + terminal_dispatch_failure = dispatch_comm.bcast(terminal_dispatch_failure, root=0) + if terminal_dispatch_failure: + break + if request.opcode is MailboxOpcode.SHUTDOWN: + break + finally: + if processor_thread is not None: + if processor_thread.is_alive(): + connection.feed(_shutdown_frame(manifest)) + processor_thread.join(timeout=5.0) + try: + inner_worker.close() + finally: + if mailbox is not None: + mailbox.close() + return exit_code + + +def _raise_keyboard_interrupt(_signum, _frame): + raise KeyboardInterrupt + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser() parser.add_argument("--group-manifest", required=True) @@ -211,46 +603,30 @@ def main(argv: list[str] | None = None) -> int: 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) + world = MPI.COMM_WORLD + group_manifest = _load_group_manifest_from_rank0(world, 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: + world_size = int(world.Get_size()) + rank = int(world.Get_rank()) + if len(rank_manifests) != world_size or int(group_manifest.get("world_size", -1)) != 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: + if len(group_manifest.get("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, - ) + dispatch_comm = world.Dup() + domain_comm = world.Dup() + try: + return _run_group_session( + dispatch_comm=dispatch_comm, + domain_comm=domain_comm, + group_manifest=group_manifest, + manifest=dict(rank_manifests[rank]), + ) + finally: + domain_comm.Free() + dispatch_comm.Free() if __name__ == "__main__": diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 08d78114f8..8f90057340 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -247,6 +247,10 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli ``workers`` contains the exact stable NEXT_LEVEL worker id for each member. For L3 chip dispatch, these are the existing chip worker ids. + When ``workers`` is the complete worker-id set of one MPI L3 group, + ``args_list`` becomes one per-rank mailbox request: MPI rank + ``workers[i]`` executes only ``args_list[i]``. A single worker id remains + directed and is never silently widened to the whole group. """ cfg = config if config is not None else CallConfig() worker_ids = [_require_next_level_worker_id(value, argument="workers entries") for value in workers] diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index bef445bdf3..89e5878df8 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1105,8 +1105,8 @@ def __repr__(self) -> str: def release(self) -> None: if self._released: return - self._release_fn(self) self._released = True + self._release_fn(self) def __enter__(self) -> GlobalCommDomainHandle: return self diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 9795bf2cf3..d370dd20d6 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -516,18 +516,25 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class MpiL3GroupSpec: - """Describes a group of L3 workers launched by one parent-owned ``mpirun``.""" + """Describes L3 workers launched by one parent-owned ``mpirun``. + + ``command_port_base``, ``health_port_base``, ``session_listen_hosts``, + ``connect_hosts``, ``allow_wildcard_session_bind``, ``ready_host``, and + ``ready_port`` are accepted only for source compatibility with PR #1623. + MPI groups ignore them and use the local named mailbox plus MPI collectives. + Non-MPI ``RemoteWorkerSpec`` continues to use its TCP fields. + """ 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, ...], ...] = () + command_port_base: int | None = None + health_port_base: int | None = None session_listen_hosts: tuple[str, ...] = () connect_hosts: tuple[str, ...] = () allow_wildcard_session_bind: bool = False @@ -573,28 +580,25 @@ def __post_init__(self) -> None: # noqa: PLR0912 -- one place validates the pub ) 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): + session_listen_hosts = tuple(str(host) for host in self.session_listen_hosts) + connect_hosts = tuple(str(host) for host in self.connect_hosts) + if session_listen_hosts and len(session_listen_hosts) != len(hosts): raise ValueError("MpiL3GroupSpec.session_listen_hosts must match hosts length") - if len(connect_hosts) != len(hosts): + if connect_hosts and 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") + command_port_base = None if self.command_port_base is None else int(self.command_port_base) + health_port_base = None if self.health_port_base is None else int(self.health_port_base) + for name, value in (("command_port_base", command_port_base), ("health_port_base", health_port_base)): + if value is not None and (value <= 0 or value + len(hosts) - 1 > 65535): + raise ValueError(f"MpiL3GroupSpec.{name} deprecated port range 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") + raise ValueError("MpiL3GroupSpec.transport must be 'sim'") 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"): @@ -636,16 +640,24 @@ class _RemoteSession: pid: int +@dataclass(frozen=True) +class _MpiL3RankSpec: + platform: str + runtime: str + device_ids: tuple[int, ...] + num_sub_workers: int + transport: str + comm_profile: str + global_device_ranks: tuple[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 + session_id: int + spec: _MpiL3RankSpec @dataclass @@ -656,6 +668,9 @@ class _MpiL3GroupRuntime: process: subprocess.Popen[Any] | None = None manifest_path: str | None = None ready_dir: str | None = None + mailbox: Any | None = None + monitor_thread: threading.Thread | None = None + closing: bool = False @dataclass @@ -3151,10 +3166,9 @@ def add_remote_worker(self, spec: RemoteWorkerSpec) -> int: 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. + A single named mailbox is created for the group during ``init()``. + The returned ids remain exact NEXT_LEVEL targets, but all of them route + through that group mailbox and the MPI collective dispatcher. """ with self._hierarchical_start_cv: if self._lifecycle is not _Lifecycle.NEW: @@ -3163,33 +3177,12 @@ def add_mpirun_worker_group(self, spec: MpiL3GroupSpec) -> tuple[int, ...]: 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): + for rank in range(len(spec.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}", + rank_spec = _MpiL3RankSpec( platform=spec.platform, runtime=spec.runtime, device_ids=spec.device_ids_by_rank[rank], @@ -3197,18 +3190,13 @@ def add_mpirun_worker_group(self, spec: MpiL3GroupSpec) -> tuple[int, ...]: 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, + session_id=self._new_remote_session_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) @@ -3358,7 +3346,7 @@ 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]]: + def _inner_registry_entries_for_spec(self, spec: RemoteWorkerSpec | _MpiL3RankSpec) -> list[dict[str, Any]]: from .remote_l3_protocol import ( # noqa: PLC0415 ChipCallableBlobLocation, RemoteChipCallablePayload, @@ -3561,6 +3549,37 @@ def _build_remote_manifest( "feature_flags": [], } + def _build_mpi_rank_manifest( + self, + *, + rank: _MpiL3RankRuntime, + startup_remaining_s: float, + ) -> dict[str, Any]: + spec = rank.spec + runtime = self._resolved_global_nodes()[int(rank.worker_id)] + return { + "session_id": int(rank.session_id), + "parent_worker_level": int(self.level), + "remote_worker_level": 3, + "worker_id": int(rank.worker_id), + "platform": spec.platform, + "runtime": spec.runtime, + "device_ids": list(spec.device_ids), + "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": runtime.node_rank, + "node_count": runtime.node_count, + "global_device_ranks": list(runtime.global_device_ranks), + "session_timeout_s": self._remote_session_timeout_s(), + "startup_remaining_s": float(startup_remaining_s), + "remote_task_dispatcher": self._remote_dispatcher_entries_for_worker(rank.worker_id), + "inner_l3_worker": self._inner_registry_entries_for_spec(spec), + "feature_flags": ["mpi-group-mailbox-v1"], + } + def _open_remote_session( self, *, spec: RemoteWorkerSpec, worker_id: int, session_id: int, deadline: float ) -> _RemoteSession: @@ -3619,99 +3638,30 @@ def _mpirun_args_select_hosts(args: tuple[str, ...]) -> bool: 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, + def _close_mpirun_group( # noqa: PLR0912 -- cleanup reports every independent process/resource failure 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]: + *, + timeout_s: float, + ) -> list[str]: failures: list[str] = [] + group.closing = True proc = group.process if proc is not None: try: if proc.poll() is None: try: - proc.terminate() + os.killpg(proc.pid, signal.SIGTERM) + except ProcessLookupError: + pass except BaseException as exc: # noqa: BLE001 failures.append(f"terminate: {exc}") try: proc.wait(timeout=timeout_s) except subprocess.TimeoutExpired: try: - proc.kill() + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass except BaseException as exc: # noqa: BLE001 failures.append(f"kill: {exc}") try: @@ -3722,6 +3672,16 @@ def _close_mpirun_group(group: _MpiL3GroupRuntime, *, timeout_s: float) -> list[ failures.append(f"wait after terminate: {exc}") finally: group.process = None + if group.monitor_thread is not None: + group.monitor_thread.join(timeout=timeout_s) + group.monitor_thread = None + if group.mailbox is not None: + try: + group.mailbox.close(unlink=True) + except BaseException as exc: # noqa: BLE001 + failures.append(f"close mailbox: {exc}") + finally: + group.mailbox = None try: if group.ready_dir is not None: shutil.rmtree(group.ready_dir) @@ -3749,38 +3709,43 @@ def _close_mpirun_groups( if not suppress_errors: raise RuntimeError(failures[0]) + @staticmethod + def _monitor_mpirun_group(group: _MpiL3GroupRuntime, proc: subprocess.Popen[Any]) -> None: + returncode = proc.wait() + if group.closing or group.mailbox is None: + return + with contextlib.suppress(BaseException): + group.mailbox.mark_terminal(f"mpirun exited unexpectedly with status {returncode}") + + @staticmethod + def _mark_mpirun_groups_closing(groups: list[_MpiL3GroupRuntime]) -> None: + for group in groups: + group.closing = True + def _activate_mpirun_worker_groups(self, deadline: float) -> None: if not self._mpi_l3_groups: return + from .mpi_group_mailbox import MAILBOX_SIZE as MPI_MAILBOX_SIZE # noqa: PLC0415 + from .mpi_group_mailbox import MailboxGroupState, MpiGroupMailbox # noqa: PLC0415 + 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-") + ready_dir = tempfile.mkdtemp(prefix="simpler-mpirun-manifest-") 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]) + mailbox = MpiGroupMailbox.create(world_size=len(group.ranks)) + group.mailbox = mailbox 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, + manifest = self._build_mpi_rank_manifest( + rank=rank, 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), @@ -3789,23 +3754,11 @@ def _activate_mpirun_worker_groups(self, deadline: float) -> None: } ) 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, + "mailbox": mailbox.manifest(), "rank_manifests": rank_manifests, } with open(manifest_path, "w", encoding="utf-8") as f: @@ -3825,33 +3778,31 @@ def _activate_mpirun_worker_groups(self, deadline: float) -> None: 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, - ) + group.process = subprocess.Popen(cmd, start_new_session=True) + while mailbox.group_state is MailboxGroupState.INITIALIZING: + if group.process.poll() is not None: + raise RuntimeError( + f"MPI L3 group {group.group_id} exited before READY (status {group.process.returncode})" + ) + self._remaining_until(deadline, "MPI L3 mailbox READY") + time.sleep(_STARTUP_POLL_INTERVAL_S) + if mailbox.group_state is not MailboxGroupState.READY: + raise RuntimeError(f"MPI L3 group {group.group_id} failed before READY: {mailbox.terminal_reason()}") + self._worker.add_mpi_group_mailbox( + [rank.worker_id for rank in group.ranks], + [rank.session_id for rank in group.ranks], + mailbox.address, + MPI_MAILBOX_SIZE, + group.process.pid, + session_timeout, + ) + group.monitor_thread = threading.Thread( + target=self._monitor_mpirun_group, + args=(group, group.process), + daemon=True, + name=f"simpler-mpirun-monitor-{group.group_id[:8]}", + ) + group.monitor_thread.start() if time.monotonic() >= deadline: raise RuntimeError("MPI L3 activation: startup deadline exceeded after attach") @@ -6085,6 +6036,7 @@ def _cleanup_partial_init(self) -> None: remote_sessions = list(self._remote_sessions) if self._worker is not None: try: + self._mark_mpirun_groups_closing(self._mpi_l3_groups) self._worker.close() except BaseException: # noqa: BLE001 pass @@ -7088,67 +7040,6 @@ def _global_domain_control(self, worker_id: int, control_name: int, payload: byt 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: @@ -7157,23 +7048,17 @@ def _mpi_group_for_involved_nodes(self, involved_nodes: tuple[int, ...]) -> _Mpi 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 _mpi_group_control( + self, + group: _MpiL3GroupRuntime, + control_name: int, + payload: bytes, + ) -> bytes: + if not group.ranks: + raise RuntimeError(f"MPI L3 group {group.group_id} has no ranks") + # One L4 request enters the rank-0 mailbox. The endpoint marks group-wide + # controls as GROUP, so every MPI rank receives the same envelope. + return self._global_domain_control(group.ranks[0].worker_id, control_name, payload) def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prepare/import/commit rollback stay ordered self, @@ -7274,6 +7159,7 @@ def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prep ) prepared_nodes: list[int] = [] + mpi_group: _MpiL3GroupRuntime | None = None try: for node_worker_id in involved_nodes: node = nodes[node_worker_id] @@ -7299,17 +7185,12 @@ def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prep 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, + reply = self._mpi_group_control( + mpi_group, 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, ) + descriptors = decode_descriptor_table(reply) else: descriptor_by_rank: dict[int, GlobalDomainDescriptor] = {} for node_worker_id in involved_nodes: @@ -7340,11 +7221,10 @@ def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prep descriptors=descriptors, ) if mpi_group is not None: - self._global_domain_control_many( - involved_nodes, + self._mpi_group_control( + mpi_group, ControlName.ALLOC_DOMAIN, encode_domain_command(commit_command), - mpi_group=mpi_group, ) else: import_command = GlobalDomainCommand( @@ -7381,13 +7261,21 @@ def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prep members=domain_members_tuple, buffers=global_buffers, ) - for node_worker_id in prepared_nodes: + if mpi_group is not None and prepared_nodes: with contextlib.suppress(BaseException): - self._global_domain_control( - node_worker_id, + self._mpi_group_control( + mpi_group, ControlName.ALLOC_DOMAIN, encode_domain_command(abort_command), ) + else: + 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( @@ -7451,11 +7339,19 @@ def _release_global_domain_claimed(self, handle: GlobalCommDomainHandle) -> None 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): + involved_nodes = tuple(dict.fromkeys(member.node_worker_id for member in handle.members)) + mpi_group = self._mpi_group_for_involved_nodes(involved_nodes) + if mpi_group is not None: try: - self._global_domain_control(node_worker_id, ControlName.RELEASE_DOMAIN, command) + self._mpi_group_control(mpi_group, ControlName.RELEASE_DOMAIN, command) except BaseException as exc: # noqa: BLE001 errors.append(exc) + else: + for node_worker_id in involved_nodes: + 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: @@ -8781,6 +8677,7 @@ def _finalize_chip() -> None: def _close_worker() -> None: if self._worker: + self._mark_mpirun_groups_closing(self._mpi_l3_groups) self._worker.close() self._worker = None self._orch = None diff --git a/src/a2a3/platform/onboard/host/comm_hccl.cpp b/src/a2a3/platform/onboard/host/comm_hccl.cpp index 6765e105d1..20fbd2f45f 100644 --- a/src/a2a3/platform/onboard/host/comm_hccl.cpp +++ b/src/a2a3/platform/onboard/host/comm_hccl.cpp @@ -95,6 +95,7 @@ static_assert( sizeof(aclrtMemFabricHandle) <= COMM_GLOBAL_DOMAIN_HANDLE_BYTES, "Fabric handle exceeds global descriptor" ); static std::unordered_map> global_domain_allocations; +static std::mutex global_domain_allocations_mutex; struct CommHandle_ { int rank; @@ -1546,7 +1547,11 @@ extern "C" int comm_global_domain_prepare( ) 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) { + local_window_base_out == nullptr) { + return -1; + } + std::lock_guard lock(global_domain_allocations_mutex); + if (global_domain_allocations.count(domain_id) != 0) { return -1; } @@ -1602,6 +1607,7 @@ extern "C" int comm_global_domain_prepare( extern "C" int comm_global_domain_import( uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out ) try { + std::lock_guard lock(global_domain_allocations_mutex); auto it = global_domain_allocations.find(domain_id); if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { return -1; @@ -1690,6 +1696,7 @@ extern "C" int comm_global_domain_import( } extern "C" int comm_global_domain_release(uint64_t domain_id) try { + std::lock_guard lock(global_domain_allocations_mutex); auto it = global_domain_allocations.find(domain_id); if (it == global_domain_allocations.end()) { return 0; diff --git a/src/common/hierarchical/mpi_group_mailbox.h b/src/common/hierarchical/mpi_group_mailbox.h new file mode 100644 index 0000000000..434a9d2c94 --- /dev/null +++ b/src/common/hierarchical/mpi_group_mailbox.h @@ -0,0 +1,76 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include + +namespace mpi_group_mailbox { + +inline constexpr uint8_t MAGIC[8] = {'S', 'M', 'P', 'I', 'B', 'O', 'X', '\0'}; +inline constexpr uint32_t PROTOCOL_VERSION = 1; +inline constexpr size_t HEADER_BYTES = 256; +inline constexpr size_t PAYLOAD_BYTES = 16U * 1024U * 1024U; +inline constexpr size_t ERROR_BYTES = 64U * 1024U; +inline constexpr size_t REQUEST_OFFSET = HEADER_BYTES; +inline constexpr size_t RESPONSE_OFFSET = REQUEST_OFFSET + PAYLOAD_BYTES; +inline constexpr size_t ERROR_OFFSET = RESPONSE_OFFSET + PAYLOAD_BYTES; +inline constexpr size_t MAILBOX_BYTES = ERROR_OFFSET + ERROR_BYTES; + +inline constexpr size_t OFF_MAGIC = 0; +inline constexpr size_t OFF_PROTOCOL_VERSION = 8; +inline constexpr size_t OFF_HEADER_BYTES = 12; +inline constexpr size_t OFF_MAILBOX_BYTES = 16; +inline constexpr size_t OFF_WORLD_SIZE = 24; +inline constexpr size_t OFF_GROUP_STATE = 28; +inline constexpr size_t OFF_REQUEST_STATE = 32; +inline constexpr size_t OFF_SEQUENCE_ID = 40; +inline constexpr size_t OFF_OPCODE = 48; +inline constexpr size_t OFF_TARGET = 52; +inline constexpr size_t OFF_TARGET_RANK = 56; +inline constexpr size_t OFF_REQUEST_COUNT = 60; +inline constexpr size_t OFF_REQUEST_BYTES = 64; +inline constexpr size_t OFF_RESPONSE_COUNT = 68; +inline constexpr size_t OFF_RESPONSE_BYTES = 72; +inline constexpr size_t OFF_ERROR_BYTES = 76; + +enum class GroupState : int32_t { + INITIALIZING = 0, + READY = 1, + TERMINAL = 2, + CLOSED = 3, +}; + +enum class RequestState : int32_t { + IDLE = 0, + REQUEST_READY = 1, + TASK_ACCEPTED = 2, + TASK_DONE = 3, + TASK_FAILED = 4, + SHUTDOWN_READY = 5, + SHUTDOWN_DONE = 6, +}; + +enum class Opcode : uint32_t { + TASK = 1, + CONTROL = 2, + PING = 3, + SHUTDOWN = 4, +}; + +enum class Target : uint32_t { + GROUP = 1, + RANK = 2, + PER_RANK = 3, +}; + +} // namespace mpi_group_mailbox diff --git a/src/common/hierarchical/remote_endpoint.cpp b/src/common/hierarchical/remote_endpoint.cpp index 2c0eac74cb..1da106da14 100644 --- a/src/common/hierarchical/remote_endpoint.cpp +++ b/src/common/hierarchical/remote_endpoint.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -355,6 +356,13 @@ void validate_owner_buffer_handle(const RemoteBufferHandle &handle, size_t reque } // namespace +std::vector +RemoteL3Transport::exchange_group_task(const std::vector &frame, uint64_t, int32_t, int32_t) { + const auto decoded = remote_l3::decode_frame(frame); + submit_frame(frame); + return wait_for_reply(remote_l3::FrameType::COMPLETION, decoded.header.sequence); +} + RemoteL3SocketTransport::RemoteL3SocketTransport( std::string host, uint16_t port, std::string health_host, uint16_t health_port, double attach_timeout_s, double runtime_timeout_s @@ -612,15 +620,427 @@ std::vector RemoteL3SocketTransport::wait_for_reply(remote_l3::FrameTyp void RemoteL3SocketTransport::shutdown() { close_socket(); } +MpiGroupMailboxChannel::MpiGroupMailboxChannel( + void *mailbox, size_t mailbox_bytes, int32_t world_size, int mpirun_pid, double runtime_timeout_s +) : + mailbox_(static_cast(mailbox)), + mailbox_bytes_(mailbox_bytes), + world_size_(world_size), + mpirun_pid_(mpirun_pid), + runtime_timeout_s_(runtime_timeout_s) { + using namespace mpi_group_mailbox; + if (mailbox_ == nullptr) throw std::invalid_argument("MpiGroupMailboxChannel: null mailbox"); + if (mailbox_bytes_ < MAILBOX_BYTES) throw std::invalid_argument("MpiGroupMailboxChannel: mailbox is too small"); + if (world_size_ <= 0) throw std::invalid_argument("MpiGroupMailboxChannel: world_size must be positive"); + if (!(runtime_timeout_s_ > 0.0)) { + throw std::invalid_argument("MpiGroupMailboxChannel: runtime_timeout_s must be positive"); + } + if (std::memcmp(mailbox_ + OFF_MAGIC, MAGIC, sizeof(MAGIC)) != 0) { + throw std::invalid_argument("MpiGroupMailboxChannel: mailbox magic mismatch"); + } + if (read_u32(OFF_PROTOCOL_VERSION) != PROTOCOL_VERSION || read_u32(OFF_HEADER_BYTES) != HEADER_BYTES) { + throw std::invalid_argument("MpiGroupMailboxChannel: mailbox protocol mismatch"); + } + if (read_u64(OFF_MAILBOX_BYTES) != MAILBOX_BYTES || + read_u32(OFF_WORLD_SIZE) != static_cast(world_size_)) { + throw std::invalid_argument("MpiGroupMailboxChannel: mailbox layout or world size mismatch"); + } +} + +int32_t MpiGroupMailboxChannel::load_i32(size_t offset) const { + int32_t value = 0; + __atomic_load(reinterpret_cast(mailbox_ + offset), &value, __ATOMIC_ACQUIRE); + return value; +} + +void MpiGroupMailboxChannel::store_i32(size_t offset, int32_t value) { + __atomic_store(reinterpret_cast(mailbox_ + offset), &value, __ATOMIC_RELEASE); +} + +uint32_t MpiGroupMailboxChannel::read_u32(size_t offset) const { + uint32_t value = 0; + std::memcpy(&value, mailbox_ + offset, sizeof(value)); + return value; +} + +uint64_t MpiGroupMailboxChannel::read_u64(size_t offset) const { + uint64_t value = 0; + std::memcpy(&value, mailbox_ + offset, sizeof(value)); + return value; +} + +void MpiGroupMailboxChannel::write_u32(size_t offset, uint32_t value) { + std::memcpy(mailbox_ + offset, &value, sizeof(value)); +} + +void MpiGroupMailboxChannel::write_u64(size_t offset, uint64_t value) { + std::memcpy(mailbox_ + offset, &value, sizeof(value)); +} + +std::string MpiGroupMailboxChannel::terminal_reason() const { + using namespace mpi_group_mailbox; + const size_t size = std::min(static_cast(read_u32(OFF_ERROR_BYTES)), ERROR_BYTES); + return std::string(reinterpret_cast(mailbox_ + ERROR_OFFSET), size); +} + +void MpiGroupMailboxChannel::mark_terminal(const std::string &reason) { + using namespace mpi_group_mailbox; + const size_t size = std::min(reason.size(), ERROR_BYTES); + if (size > 0) std::memcpy(mailbox_ + ERROR_OFFSET, reason.data(), size); + write_u32(OFF_ERROR_BYTES, static_cast(size)); + store_i32(OFF_GROUP_STATE, static_cast(GroupState::TERMINAL)); + store_i32(OFF_REQUEST_STATE, static_cast(RequestState::TASK_FAILED)); +} + +void MpiGroupMailboxChannel::kill_mpirun_group() const { + if (mpirun_pid_ <= 0) return; + (void)::kill(-mpirun_pid_, SIGKILL); + (void)::kill(mpirun_pid_, SIGKILL); +} + +bool MpiGroupMailboxChannel::terminal() const { + return load_i32(mpi_group_mailbox::OFF_GROUP_STATE) == + static_cast(mpi_group_mailbox::GroupState::TERMINAL); +} + +std::vector> MpiGroupMailboxChannel::run_exchange( + const std::vector> &frames, mpi_group_mailbox::Opcode opcode, mpi_group_mailbox::Target target, + int32_t target_rank +) { + using namespace mpi_group_mailbox; + if (frames.empty()) throw std::invalid_argument("MpiGroupMailboxChannel: request requires a payload"); + if (target == Target::PER_RANK && frames.size() != static_cast(world_size_)) { + throw std::invalid_argument("MpiGroupMailboxChannel: per-rank request must include every rank"); + } + if (target != Target::PER_RANK && frames.size() != 1) { + throw std::invalid_argument("MpiGroupMailboxChannel: rank/group request requires one payload"); + } + size_t encoded_request_bytes = 4 + 4 * frames.size(); + for (const auto &frame : frames) { + if (frame.size() > UINT32_MAX || frame.size() > PAYLOAD_BYTES || + encoded_request_bytes > PAYLOAD_BYTES - frame.size()) { + throw std::invalid_argument("MpiGroupMailboxChannel: request frame vector exceeds mailbox capacity"); + } + encoded_request_bytes += frame.size(); + } + if (target == Target::RANK && (target_rank < 0 || target_rank >= world_size_)) { + throw std::invalid_argument("MpiGroupMailboxChannel: target rank is outside the group"); + } + const auto group_state = static_cast(load_i32(OFF_GROUP_STATE)); + if (group_state == GroupState::TERMINAL) { + throw std::runtime_error("MpiGroupMailboxChannel: group is terminal: " + terminal_reason()); + } + if (group_state != GroupState::READY) { + throw std::runtime_error("MpiGroupMailboxChannel: group is not ready"); + } + if (static_cast(load_i32(OFF_REQUEST_STATE)) != RequestState::IDLE) { + throw std::runtime_error("MpiGroupMailboxChannel: request lane is not idle"); + } + + const uint64_t sequence = next_sequence_++; + const uint32_t count = static_cast(frames.size()); + std::memcpy(mailbox_ + REQUEST_OFFSET, &count, sizeof(count)); + size_t request_offset = REQUEST_OFFSET + 4; + for (const auto &frame : frames) { + const uint32_t frame_size = static_cast(frame.size()); + std::memcpy(mailbox_ + request_offset, &frame_size, sizeof(frame_size)); + request_offset += 4; + } + for (const auto &frame : frames) { + if (!frame.empty()) std::memcpy(mailbox_ + request_offset, frame.data(), frame.size()); + request_offset += frame.size(); + } + write_u64(OFF_SEQUENCE_ID, sequence); + write_u32(OFF_OPCODE, static_cast(opcode)); + write_u32(OFF_TARGET, static_cast(target)); + std::memcpy(mailbox_ + OFF_TARGET_RANK, &target_rank, sizeof(target_rank)); + write_u32(OFF_REQUEST_COUNT, count); + write_u32(OFF_REQUEST_BYTES, static_cast(encoded_request_bytes)); + write_u32(OFF_RESPONSE_COUNT, 0); + write_u32(OFF_RESPONSE_BYTES, 0); + write_u32(OFF_ERROR_BYTES, 0); + const RequestState ready_state = + opcode == Opcode::SHUTDOWN ? RequestState::SHUTDOWN_READY : RequestState::REQUEST_READY; + store_i32(OFF_REQUEST_STATE, static_cast(ready_state)); + + const Deadline deadline = deadline_from_now(runtime_timeout_s_); + uint32_t poll_spins = 0; + while (true) { + const auto state = static_cast(load_i32(OFF_REQUEST_STATE)); + if (state == RequestState::TASK_DONE) { + const uint32_t response_count = read_u32(OFF_RESPONSE_COUNT); + const size_t response_bytes = read_u32(OFF_RESPONSE_BYTES); + const uint32_t expected_count = target == Target::PER_RANK ? static_cast(world_size_) : 1U; + const size_t prefix_bytes = 4 + 4 * static_cast(response_count); + if (response_count != expected_count || response_bytes < prefix_bytes || response_bytes > PAYLOAD_BYTES) { + mark_terminal("MPI group mailbox returned an invalid response vector"); + kill_mpirun_group(); + throw std::runtime_error("MpiGroupMailboxChannel: invalid response vector"); + } + uint32_t encoded_count = 0; + std::memcpy(&encoded_count, mailbox_ + RESPONSE_OFFSET, sizeof(encoded_count)); + if (encoded_count != response_count) { + mark_terminal("MPI group mailbox response vector length mismatch"); + kill_mpirun_group(); + throw std::runtime_error("MpiGroupMailboxChannel: response vector length mismatch"); + } + std::vector payload_sizes(response_count); + size_t total_payload_bytes = 0; + for (uint32_t i = 0; i < response_count; ++i) { + std::memcpy( + &payload_sizes[i], mailbox_ + RESPONSE_OFFSET + 4 + 4 * static_cast(i), + sizeof(payload_sizes[i]) + ); + total_payload_bytes += payload_sizes[i]; + } + if (prefix_bytes + total_payload_bytes != response_bytes) { + mark_terminal("MPI group mailbox response vector length mismatch"); + kill_mpirun_group(); + throw std::runtime_error("MpiGroupMailboxChannel: response vector length mismatch"); + } + std::vector> responses; + responses.reserve(response_count); + size_t response_offset = RESPONSE_OFFSET + prefix_bytes; + for (uint32_t payload_size : payload_sizes) { + std::vector response(payload_size); + if (payload_size > 0) { + std::memcpy(response.data(), mailbox_ + response_offset, payload_size); + } + response_offset += payload_size; + responses.push_back(std::move(response)); + } + store_i32(OFF_REQUEST_STATE, static_cast(RequestState::IDLE)); + return responses; + } + if (state == RequestState::SHUTDOWN_DONE) { + store_i32(OFF_REQUEST_STATE, static_cast(RequestState::IDLE)); + return {{}}; + } + if (state == RequestState::TASK_FAILED || terminal()) { + const std::string reason = terminal_reason(); + if (!terminal()) store_i32(OFF_REQUEST_STATE, static_cast(RequestState::IDLE)); + throw std::runtime_error( + "MpiGroupMailboxChannel: MPI group request failed" + (reason.empty() ? std::string() : ": " + reason) + ); + } + if (std::chrono::steady_clock::now() >= deadline) { + mark_terminal("MPI group mailbox request timed out at sequence " + std::to_string(sequence)); + kill_mpirun_group(); + throw std::runtime_error("MpiGroupMailboxChannel: request timed out"); + } + if (++poll_spins >= 64) { + std::this_thread::yield(); + poll_spins = 0; + } + } +} + +std::vector MpiGroupMailboxChannel::exchange(const std::vector &frame, int32_t target_rank) { + std::lock_guard lock(lane_mu_); + const auto decoded = remote_l3::decode_frame(frame); + mpi_group_mailbox::Opcode opcode = mpi_group_mailbox::Opcode::CONTROL; + mpi_group_mailbox::Target target = mpi_group_mailbox::Target::RANK; + if (decoded.header.frame_type == remote_l3::FrameType::TASK) { + opcode = mpi_group_mailbox::Opcode::TASK; + } else if (decoded.header.frame_type == remote_l3::FrameType::HEALTH) { + opcode = mpi_group_mailbox::Opcode::PING; + target = mpi_group_mailbox::Target::GROUP; + target_rank = -1; + } else if (decoded.header.frame_type == remote_l3::FrameType::CONTROL) { + const auto control = remote_l3::decode_control(decoded.payload.data(), decoded.payload.size()); + if (control.control_name == remote_l3::ControlName::ALLOC_DOMAIN || + control.control_name == remote_l3::ControlName::RELEASE_DOMAIN) { + target = mpi_group_mailbox::Target::GROUP; + target_rank = -1; + } + } else { + throw std::runtime_error("MpiGroupMailboxChannel: unsupported request frame type"); + } + auto responses = run_exchange({frame}, opcode, target, target_rank); + return std::move(responses.front()); +} + +std::vector MpiGroupMailboxChannel::exchange_group_task( + const std::vector &frame, int32_t target_rank, uint64_t task_slot, int32_t group_size +) { + if (group_size != world_size_) { + // Existing submit_next_level_group permits subsets. Keep that behavior + // as ordered rank-targeted requests; only a full MPI group is batched + // into one PER_RANK mailbox envelope. + return exchange(frame, target_rank); + } + if (target_rank < 0 || target_rank >= world_size_) { + throw std::invalid_argument("MpiGroupMailboxChannel: group task target rank is outside the group"); + } + + const Deadline deadline = deadline_from_now(runtime_timeout_s_); + bool leader = false; + std::vector> frames; + { + std::unique_lock lock(group_mu_); + while (group_active_ && group_task_slot_ != task_slot) { + if (group_cv_.wait_until(lock, deadline) == std::cv_status::timeout) { + lock.unlock(); + mark_terminal("MPI group task batching timed out waiting for the prior task"); + kill_mpirun_group(); + throw std::runtime_error("MpiGroupMailboxChannel: group task batching timed out"); + } + } + if (!group_active_) { + group_active_ = true; + group_done_ = false; + group_task_slot_ = task_slot; + group_arrived_ = 0; + group_departed_ = 0; + group_frames_.assign(static_cast(world_size_), {}); + group_replies_.clear(); + group_error_ = nullptr; + } + auto &rank_frame = group_frames_[static_cast(target_rank)]; + if (!rank_frame.empty()) { + throw std::runtime_error("MpiGroupMailboxChannel: duplicate rank in one MPI group task"); + } + rank_frame = frame; + ++group_arrived_; + if (group_arrived_ == world_size_) { + leader = true; + frames = group_frames_; + } else { + while (!group_done_) { + if (group_cv_.wait_until(lock, deadline) == std::cv_status::timeout) { + group_error_ = std::make_exception_ptr(std::runtime_error("MPI group task batching timed out")); + group_done_ = true; + ++group_departed_; + lock.unlock(); + mark_terminal("MPI group task batching timed out waiting for all rank payloads"); + kill_mpirun_group(); + group_cv_.notify_all(); + throw std::runtime_error("MpiGroupMailboxChannel: group task batching timed out"); + } + } + } + } + + if (leader) { + try { + std::lock_guard lane_lock(lane_mu_); + auto replies = + run_exchange(frames, mpi_group_mailbox::Opcode::TASK, mpi_group_mailbox::Target::PER_RANK, -1); + std::lock_guard group_lock(group_mu_); + group_replies_ = std::move(replies); + group_done_ = true; + } catch (...) { + std::lock_guard group_lock(group_mu_); + group_error_ = std::current_exception(); + group_done_ = true; + } + group_cv_.notify_all(); + } + + std::unique_lock lock(group_mu_); + while (!group_done_) { + if (group_cv_.wait_until(lock, deadline) == std::cv_status::timeout) { + group_error_ = std::make_exception_ptr(std::runtime_error("MPI group task dispatch timed out")); + group_done_ = true; + ++group_departed_; + lock.unlock(); + mark_terminal("MPI group task dispatch timed out"); + kill_mpirun_group(); + group_cv_.notify_all(); + throw std::runtime_error("MpiGroupMailboxChannel: group task dispatch timed out"); + } + } + std::exception_ptr error = group_error_; + std::vector reply; + if (error == nullptr) { + if (group_replies_.size() != static_cast(world_size_)) { + error = std::make_exception_ptr( + std::runtime_error("MpiGroupMailboxChannel: MPI group task reply count mismatch") + ); + } else { + reply = group_replies_[static_cast(target_rank)]; + } + } + ++group_departed_; + if (group_departed_ == world_size_) { + group_active_ = false; + group_done_ = false; + group_frames_.clear(); + group_replies_.clear(); + group_error_ = nullptr; + group_cv_.notify_all(); + } + lock.unlock(); + if (error != nullptr) std::rethrow_exception(error); + return reply; +} + +void MpiGroupMailboxChannel::shutdown(const std::vector &frame) { + std::lock_guard lock(lane_mu_); + if (shutdown_sent_ || terminal()) return; + shutdown_sent_ = true; + (void)run_exchange({frame}, mpi_group_mailbox::Opcode::SHUTDOWN, mpi_group_mailbox::Target::GROUP, -1); +} + +MpiGroupMailboxTransport::MpiGroupMailboxTransport( + std::shared_ptr channel, int32_t target_rank +) : + channel_(std::move(channel)), + target_rank_(target_rank) { + if (!channel_) throw std::invalid_argument("MpiGroupMailboxTransport: null channel"); + if (target_rank_ < 0) throw std::invalid_argument("MpiGroupMailboxTransport: target rank must be non-negative"); +} + +void MpiGroupMailboxTransport::submit_frame(const std::vector &frame) { + if (pending_) throw std::runtime_error("MpiGroupMailboxTransport: a request is already pending"); + pending_frame_ = frame; + pending_ = true; +} + +std::vector MpiGroupMailboxTransport::wait_for_reply(remote_l3::FrameType frame_type, uint64_t sequence) { + if (!pending_) throw std::runtime_error("MpiGroupMailboxTransport: no request is pending"); + auto frame = std::move(pending_frame_); + pending_frame_.clear(); + pending_ = false; + auto reply = channel_->exchange(frame, target_rank_); + auto decoded = remote_l3::decode_frame(reply); + if (decoded.header.frame_type != frame_type || decoded.header.sequence != sequence) { + throw std::runtime_error("MpiGroupMailboxTransport: reply frame type or sequence mismatch"); + } + return reply; +} + +std::vector MpiGroupMailboxTransport::exchange_group_task( + const std::vector &frame, uint64_t task_slot, int32_t, int32_t group_size +) { + const auto request = remote_l3::decode_frame(frame); + auto reply = channel_->exchange_group_task(frame, target_rank_, task_slot, group_size); + const auto decoded = remote_l3::decode_frame(reply); + if (decoded.header.frame_type != remote_l3::FrameType::COMPLETION || + decoded.header.sequence != request.header.sequence) { + throw std::runtime_error("MpiGroupMailboxTransport: group task reply frame type or sequence mismatch"); + } + return reply; +} + +void MpiGroupMailboxTransport::shutdown() { + if (!pending_) return; + auto frame = std::move(pending_frame_); + pending_frame_.clear(); + pending_ = false; + channel_->shutdown(frame); +} + RemoteL3Endpoint::RemoteL3Endpoint( - int32_t worker_id, uint64_t session_id, std::string transport_name, std::unique_ptr transport + int32_t worker_id, uint64_t session_id, std::string transport_name, std::unique_ptr transport, + WorkerEndpointKind endpoint_kind ) : session_id_(session_id), transport_(std::move(transport)) { if (worker_id < 0) throw std::invalid_argument("RemoteL3Endpoint: worker_id must be non-negative"); if (session_id == 0) throw std::invalid_argument("RemoteL3Endpoint: session_id must be non-zero"); if (!transport_) throw std::invalid_argument("RemoteL3Endpoint: null transport"); - caps_.kind = WorkerEndpointKind::REMOTE_L3; + caps_.kind = endpoint_kind; caps_.worker_id = worker_id; caps_.remote = true; caps_.supports_task_dispatch = true; @@ -628,6 +1048,14 @@ RemoteL3Endpoint::RemoteL3Endpoint( caps_.transport = std::move(transport_name); } +MpiGroupMailboxEndpoint::MpiGroupMailboxEndpoint( + int32_t worker_id, uint64_t session_id, int32_t rank, std::shared_ptr channel +) : + RemoteL3Endpoint( + worker_id, session_id, "mpi-group-mailbox", + std::make_unique(std::move(channel), rank), WorkerEndpointKind::MPI_GROUP_MAILBOX + ) {} + remote_l3::TaskPayloadWire RemoteL3Endpoint::build_task_payload(const TaskSlotState &slot, int32_t group_index) const { remote_l3::TaskPayloadWire payload; payload.callable_digest = slot.callable.digest; @@ -683,9 +1111,16 @@ WorkerCompletion RemoteL3Endpoint::run(Ring *ring, const WorkerDispatch &dispatc header.session_id = session_id_; header.worker_id = caps_.worker_id; header.sequence = sequence; - transport_->submit_frame(remote_l3::encode_frame(header, payload)); - - auto reply_bytes = transport_->wait_for_reply(remote_l3::FrameType::COMPLETION, sequence); + auto frame = remote_l3::encode_frame(header, payload); + std::vector reply_bytes; + if (slot.is_group() && transport_->supports_group_batch()) { + reply_bytes = transport_->exchange_group_task( + frame, static_cast(dispatch.task_slot), dispatch.group_index, slot.group_size() + ); + } else { + transport_->submit_frame(frame); + reply_bytes = transport_->wait_for_reply(remote_l3::FrameType::COMPLETION, sequence); + } auto reply = remote_l3::decode_frame(reply_bytes); if (reply.header.frame_type != remote_l3::FrameType::COMPLETION) { throw std::runtime_error("RemoteL3Endpoint::run: expected COMPLETION reply"); diff --git a/src/common/hierarchical/remote_endpoint.h b/src/common/hierarchical/remote_endpoint.h index aec470d24c..9836e73659 100644 --- a/src/common/hierarchical/remote_endpoint.h +++ b/src/common/hierarchical/remote_endpoint.h @@ -15,10 +15,13 @@ #include #include #include +#include +#include #include #include #include +#include "mpi_group_mailbox.h" #include "remote_wire.h" #include "worker_manager.h" @@ -27,6 +30,9 @@ class RemoteL3Transport { virtual ~RemoteL3Transport() = default; virtual void submit_frame(const std::vector &frame) = 0; virtual std::vector wait_for_reply(remote_l3::FrameType frame_type, uint64_t sequence) = 0; + virtual bool supports_group_batch() const { return false; } + virtual std::vector + exchange_group_task(const std::vector &frame, uint64_t task_slot, int32_t group_index, int32_t group_size); virtual void shutdown() {} }; @@ -74,10 +80,77 @@ class RemoteL3SocketTransport : public RemoteL3Transport { std::vector read_frame(std::chrono::steady_clock::time_point deadline); }; +class MpiGroupMailboxChannel { +public: + MpiGroupMailboxChannel( + void *mailbox, size_t mailbox_bytes, int32_t world_size, int mpirun_pid, double runtime_timeout_s + ); + + std::vector exchange(const std::vector &frame, int32_t target_rank); + std::vector + exchange_group_task(const std::vector &frame, int32_t target_rank, uint64_t task_slot, int32_t group_size); + void shutdown(const std::vector &frame); + bool terminal() const; + +private: + uint8_t *mailbox_{nullptr}; + size_t mailbox_bytes_{0}; + int32_t world_size_{0}; + int mpirun_pid_{-1}; + double runtime_timeout_s_{30.0}; + uint64_t next_sequence_{1}; + bool shutdown_sent_{false}; + mutable std::mutex lane_mu_; + std::mutex group_mu_; + std::condition_variable group_cv_; + bool group_active_{false}; + bool group_done_{false}; + uint64_t group_task_slot_{0}; + int32_t group_arrived_{0}; + int32_t group_departed_{0}; + std::vector> group_frames_; + std::vector> group_replies_; + std::exception_ptr group_error_; + + int32_t load_i32(size_t offset) const; + void store_i32(size_t offset, int32_t value); + uint32_t read_u32(size_t offset) const; + uint64_t read_u64(size_t offset) const; + void write_u32(size_t offset, uint32_t value); + void write_u64(size_t offset, uint64_t value); + std::string terminal_reason() const; + void mark_terminal(const std::string &reason); + void kill_mpirun_group() const; + std::vector> run_exchange( + const std::vector> &frames, mpi_group_mailbox::Opcode opcode, + mpi_group_mailbox::Target target, int32_t target_rank + ); +}; + +class MpiGroupMailboxTransport : public RemoteL3Transport { +public: + MpiGroupMailboxTransport(std::shared_ptr channel, int32_t target_rank); + + void submit_frame(const std::vector &frame) override; + std::vector wait_for_reply(remote_l3::FrameType frame_type, uint64_t sequence) override; + bool supports_group_batch() const override { return true; } + std::vector exchange_group_task( + const std::vector &frame, uint64_t task_slot, int32_t group_index, int32_t group_size + ) override; + void shutdown() override; + +private: + std::shared_ptr channel_; + int32_t target_rank_{-1}; + std::vector pending_frame_; + bool pending_{false}; +}; + class RemoteL3Endpoint : public WorkerEndpoint { public: RemoteL3Endpoint( - int32_t worker_id, uint64_t session_id, std::string transport_name, std::unique_ptr transport + int32_t worker_id, uint64_t session_id, std::string transport_name, + std::unique_ptr transport, WorkerEndpointKind endpoint_kind = WorkerEndpointKind::REMOTE_L3 ); const WorkerEndpointCaps &caps() const override { return caps_; } @@ -124,3 +197,10 @@ class RemoteL3Endpoint : public WorkerEndpoint { remote_l3::ControlReplyPayload run_control(remote_l3::ControlName control_name, const std::vector &command_bytes); }; + +class MpiGroupMailboxEndpoint final : public RemoteL3Endpoint { +public: + MpiGroupMailboxEndpoint( + int32_t worker_id, uint64_t session_id, int32_t rank, std::shared_ptr channel + ); +}; diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 90efac32a3..a771d242d4 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -93,6 +93,26 @@ void Worker::add_remote_l3_socket( ); } +void Worker::add_mpi_group_mailbox( + const std::vector &worker_ids, const std::vector &session_ids, void *mailbox, + size_t mailbox_bytes, int mpirun_pid, double runtime_timeout_s +) { + if (initialized_) throw std::runtime_error("Worker: add_mpi_group_mailbox after init"); + if (worker_ids.empty() || worker_ids.size() != session_ids.size()) { + throw std::invalid_argument("Worker: MPI group worker_ids and session_ids must have the same non-zero size"); + } + auto channel = std::make_shared( + mailbox, mailbox_bytes, static_cast(worker_ids.size()), mpirun_pid, runtime_timeout_s + ); + for (size_t rank = 0; rank < worker_ids.size(); ++rank) { + manager_.add_next_level_endpoint( + std::make_unique( + worker_ids[rank], session_ids[rank], static_cast(rank), channel + ) + ); + } +} + void Worker::init() { if (initialized_) throw std::runtime_error("Worker: already initialized"); diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index 27531b7e89..5cc9785c91 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -40,6 +40,7 @@ #pragma once +#include #include #include #include @@ -86,6 +87,10 @@ class Worker { uint16_t port, const std::string &health_host, uint16_t health_port, double attach_timeout_s, double runtime_timeout_s ); + void add_mpi_group_mailbox( + const std::vector &worker_ids, const std::vector &session_ids, void *mailbox, + size_t mailbox_bytes, int mpirun_pid, double runtime_timeout_s + ); // Start the scheduler thread. Must be called AFTER the parent has forked // any child workers — init() spins up threads in the parent that would diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 4c14eef357..4aa6e23b2a 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -195,6 +195,7 @@ struct WorkerDispatch; enum class WorkerEndpointKind : int32_t { LOCAL_MAILBOX = 0, REMOTE_L3 = 1, + MPI_GROUP_MAILBOX = 2, }; struct WorkerEndpointCaps { diff --git a/src/common/platform_comm/comm_sim.cpp b/src/common/platform_comm/comm_sim.cpp index d6b08b38ed..cd65a404fb 100644 --- a/src/common/platform_comm/comm_sim.cpp +++ b/src/common/platform_comm/comm_sim.cpp @@ -45,6 +45,7 @@ #include #include #include +#include #include #include #include @@ -217,6 +218,7 @@ struct GlobalDomainAllocation { static_assert(sizeof(CommGlobalDomainDescriptor) == 288, "global domain descriptor ABI changed"); static std::unordered_map> global_domain_allocations; +static std::mutex global_domain_allocations_mutex; struct CommHandle_ { int rank; @@ -747,6 +749,7 @@ extern "C" int comm_global_domain_prepare( local_window_base_out == nullptr) { return -1; } + std::lock_guard lock(global_domain_allocations_mutex); if (global_domain_allocations.count(domain_id) != 0) { return -1; } @@ -806,6 +809,7 @@ extern "C" int comm_global_domain_prepare( extern "C" int comm_global_domain_import( uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out ) try { + std::lock_guard lock(global_domain_allocations_mutex); auto it = global_domain_allocations.find(domain_id); if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { return -1; @@ -873,6 +877,7 @@ extern "C" int comm_global_domain_import( } extern "C" int comm_global_domain_release(uint64_t domain_id) try { + std::lock_guard lock(global_domain_allocations_mutex); auto it = global_domain_allocations.find(domain_id); if (it == global_domain_allocations.end()) { return 0; diff --git a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp index a7ef07b7de..013acaa70c 100644 --- a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp +++ b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp @@ -624,3 +624,97 @@ TEST(RemoteEndpoint, BareHostPointerWithoutSidecarIsEndpointFailure) { EXPECT_TRUE(transport->last_frame.empty()); ring.shutdown(); } + +namespace { + +std::vector ready_mpi_mailbox(int32_t world_size) { + using namespace mpi_group_mailbox; + std::vector mailbox(MAILBOX_BYTES, 0); + std::memcpy(mailbox.data() + OFF_MAGIC, MAGIC, sizeof(MAGIC)); + const uint32_t version = PROTOCOL_VERSION; + const uint32_t header_bytes = HEADER_BYTES; + const uint64_t mailbox_bytes = MAILBOX_BYTES; + const uint32_t size = static_cast(world_size); + const int32_t ready = static_cast(GroupState::READY); + const int32_t idle = static_cast(RequestState::IDLE); + std::memcpy(mailbox.data() + OFF_PROTOCOL_VERSION, &version, sizeof(version)); + std::memcpy(mailbox.data() + OFF_HEADER_BYTES, &header_bytes, sizeof(header_bytes)); + std::memcpy(mailbox.data() + OFF_MAILBOX_BYTES, &mailbox_bytes, sizeof(mailbox_bytes)); + std::memcpy(mailbox.data() + OFF_WORLD_SIZE, &size, sizeof(size)); + std::memcpy(mailbox.data() + OFF_GROUP_STATE, &ready, sizeof(ready)); + std::memcpy(mailbox.data() + OFF_REQUEST_STATE, &idle, sizeof(idle)); + return mailbox; +} + +int32_t mailbox_state(const std::vector &mailbox, size_t offset) { + int32_t value = 0; + __atomic_load(reinterpret_cast(mailbox.data() + offset), &value, __ATOMIC_ACQUIRE); + return value; +} + +void set_mailbox_state(std::vector &mailbox, size_t offset, int32_t value) { + __atomic_store(reinterpret_cast(mailbox.data() + offset), &value, __ATOMIC_RELEASE); +} + +void respond_with_payloads(std::vector &mailbox, const std::vector> &payloads) { + using namespace mpi_group_mailbox; + while (mailbox_state(mailbox, OFF_REQUEST_STATE) != static_cast(RequestState::REQUEST_READY)) {} + const uint32_t count = static_cast(payloads.size()); + std::memcpy(mailbox.data() + RESPONSE_OFFSET, &count, sizeof(count)); + size_t offset = RESPONSE_OFFSET + 4; + size_t response_bytes = 4 + 4 * payloads.size(); + for (const auto &payload : payloads) { + const uint32_t size = static_cast(payload.size()); + std::memcpy(mailbox.data() + offset, &size, sizeof(size)); + offset += 4; + response_bytes += payload.size(); + } + for (const auto &payload : payloads) { + if (!payload.empty()) std::memcpy(mailbox.data() + offset, payload.data(), payload.size()); + offset += payload.size(); + } + std::memcpy(mailbox.data() + OFF_RESPONSE_COUNT, &count, sizeof(count)); + const uint32_t encoded_bytes = static_cast(response_bytes); + std::memcpy(mailbox.data() + OFF_RESPONSE_BYTES, &encoded_bytes, sizeof(encoded_bytes)); + set_mailbox_state(mailbox, OFF_REQUEST_STATE, static_cast(mpi_group_mailbox::RequestState::TASK_DONE)); +} + +} // namespace + +TEST(MpiGroupMailboxChannel, FullGroupTaskUsesOnePerRankEnvelope) { + using namespace mpi_group_mailbox; + auto mailbox = ready_mpi_mailbox(2); + MpiGroupMailboxChannel channel(mailbox.data(), mailbox.size(), 2, -1, 2.0); + std::vector reply0; + std::vector reply1; + std::thread rank0([&]() { + reply0 = channel.exchange_group_task({0x10}, 0, 7, 2); + }); + std::thread rank1([&]() { + reply1 = channel.exchange_group_task({0x20}, 1, 7, 2); + }); + + while (mailbox_state(mailbox, OFF_REQUEST_STATE) != static_cast(RequestState::REQUEST_READY)) {} + uint32_t target = 0; + uint32_t request_count = 0; + std::memcpy(&target, mailbox.data() + OFF_TARGET, sizeof(target)); + std::memcpy(&request_count, mailbox.data() + OFF_REQUEST_COUNT, sizeof(request_count)); + EXPECT_EQ(target, static_cast(Target::PER_RANK)); + EXPECT_EQ(request_count, 2U); + respond_with_payloads(mailbox, {{0xA0}, {0xA1}}); + + rank0.join(); + rank1.join(); + EXPECT_EQ(reply0, std::vector({0xA0})); + EXPECT_EQ(reply1, std::vector({0xA1})); + EXPECT_EQ(mailbox_state(mailbox, OFF_REQUEST_STATE), static_cast(RequestState::IDLE)); +} + +TEST(MpiGroupMailboxChannel, TimeoutMakesGroupTerminal) { + using namespace mpi_group_mailbox; + auto mailbox = ready_mpi_mailbox(1); + MpiGroupMailboxChannel channel(mailbox.data(), mailbox.size(), 1, -1, 0.01); + EXPECT_THROW(channel.exchange_group_task({0x10}, 0, 9, 1), std::runtime_error); + EXPECT_TRUE(channel.terminal()); + EXPECT_EQ(mailbox_state(mailbox, OFF_GROUP_STATE), static_cast(GroupState::TERMINAL)); +} diff --git a/tests/ut/py/test_global_comm_domain.py b/tests/ut/py/test_global_comm_domain.py index 80be90692d..ddd4ca457c 100644 --- a/tests/ut/py/test_global_comm_domain.py +++ b/tests/ut/py/test_global_comm_domain.py @@ -25,6 +25,7 @@ CTRL_GLOBAL_DOMAIN_PREPARE, CTRL_GLOBAL_DOMAIN_RELEASE, GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + GLOBAL_DOMAIN_MAX_BUFFERS, GLOBAL_DOMAIN_PROFILE_IDS, GLOBAL_DOMAIN_VERSION, GlobalCommInitCommand, @@ -174,6 +175,22 @@ def test_global_domain_wire_round_trips_topology_and_descriptor_table(): assert GLOBAL_DOMAIN_DESCRIPTOR_BYTES == 288 +def test_global_domain_encode_rejects_too_many_buffers(): + command = GlobalDomainCommand( + phase=GlobalDomainPhase.PREPARE_EXPORT, + domain_id=11, + generation=1, + name="tp", + profile="sim", + window_size=GLOBAL_DOMAIN_MAX_BUFFERS + 1, + members=_members(), + buffers=tuple(GlobalDomainBuffer(f"payload-{index}", 1) for index in range(GLOBAL_DOMAIN_MAX_BUFFERS + 1)), + ) + + with pytest.raises(ValueError, match="buffer count exceeds maximum"): + encode_domain_command(command) + + 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 @@ -286,8 +303,6 @@ def _mpi_static_worker(): 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,)), @@ -396,9 +411,15 @@ def control(worker_id, control_name, payload): 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.PREPARE_EXPORT] == 1 + assert counts[GlobalDomainPhase.COMMIT] == 1 assert counts[GlobalDomainPhase.IMPORT] == 0 + group_phases = [ + worker_id + for phase, worker_id in calls + if phase in (GlobalDomainPhase.PREPARE_EXPORT, GlobalDomainPhase.COMMIT) + ] + assert group_phases == [node_ids[0], node_ids[0]] assert worker._live_global_domains["mpi-static"] is handle assert resources.live_global_domains["mpi-static"] is handle finally: @@ -527,7 +548,7 @@ def test_global_domain_descriptor_table_rejects_different_mapping_sizes(): validate_descriptor_table(tuple(descriptors), rank_count=2, profile="sim") -def test_global_domain_release_retries_after_callback_failure(): +def test_global_domain_release_stays_released_after_callback_failure(): from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 attempts = 0 @@ -535,8 +556,7 @@ def test_global_domain_release_retries_after_callback_failure(): def release_fn(_handle): nonlocal attempts attempts += 1 - if attempts == 1: - raise RuntimeError("transient release failure") + raise RuntimeError("release failure") handle = GlobalCommDomainHandle( name="retry", @@ -549,13 +569,12 @@ def release_fn(_handle): _release_fn=release_fn, ) - with pytest.raises(RuntimeError, match="transient release failure"): + with pytest.raises(RuntimeError, match="release failure"): handle.release() - assert not handle.released - handle.release() assert handle.released - assert attempts == 2 + handle.release() + assert attempts == 1 def test_global_domain_handle_repr_reports_lifecycle_state(): @@ -579,61 +598,45 @@ def test_global_domain_handle_repr_reports_lifecycle_state(): assert repr(handle) == "GlobalCommDomainHandle(name='repr', members=0, freed)" -def test_mpi_global_domain_collective_timeout_releases_local_state(monkeypatch): +def test_mpi_global_domain_collective_uses_pickle_allgather(): from simpler.mpi_l3_session import MpiGlobalDomainExchange # noqa: PLC0415 - class _PendingRequest: - @staticmethod - def test(): - return False, None - class _Comm: - aborted = False + payloads = [] @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") + def allgather(self, payload): + self.payloads.append(payload) + return [payload, b"peer"] - 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)) + gathered = exchange._allgather( + b"payload", + operation="prepare", + on_timeout=lambda: pytest.fail("blocking allgather must use the outer group watchdog"), + ) - assert releases == [True] - assert comm.aborted + assert gathered == [b"payload", b"peer"] + assert comm.payloads == [b"payload"] 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) + def allgather(payload): + return [payload] class _InnerWorker: released = False diff --git a/tests/ut/py/test_mpi_group_mailbox.py b/tests/ut/py/test_mpi_group_mailbox.py new file mode 100644 index 0000000000..5a09050bc2 --- /dev/null +++ b/tests/ut/py/test_mpi_group_mailbox.py @@ -0,0 +1,227 @@ +# 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 ctypes + +import pytest +from simpler.mpi_group_mailbox import ( + MailboxGroupState, + MailboxOpcode, + MailboxRequestState, + MailboxTarget, + MpiGroupError, + MpiGroupMailbox, + MpiRankError, + _as_byte_view, + open_rank_mailbox, +) + + +def test_shared_memory_view_is_normalized_to_unsigned_bytes(): + raw = (ctypes.c_ubyte * 4)() + original_view = memoryview(raw) + byte_view = _as_byte_view(original_view) + + byte_view[:] = b"test" + + assert byte_view is not original_view + assert byte_view.format == "B" + assert byte_view.ndim == 1 + assert bytes(raw) == b"test" + + +def test_named_mailbox_reopens_by_manifest_name_and_only_rank0_opens(): + owner = MpiGroupMailbox.create(world_size=2) + reopened = None + try: + manifest = owner.manifest() + assert open_rank_mailbox(manifest, rank=1) is None + + reopened = open_rank_mailbox(manifest, rank=0) + assert reopened is not None + assert reopened.name == owner.name + assert reopened.world_size == 2 + assert reopened.group_state is MailboxGroupState.INITIALIZING + + reopened.publish_ready() + assert owner.group_state is MailboxGroupState.READY + finally: + if reopened is not None: + reopened.close() + owner.close(unlink=True) + + +@pytest.mark.parametrize( + ("target", "target_rank", "payloads"), + [ + (MailboxTarget.GROUP, -1, (b"broadcast",)), + (MailboxTarget.RANK, 1, (b"rank-1",)), + (MailboxTarget.PER_RANK, -1, (b"rank-0", b"rank-1")), + ], +) +def test_request_envelope_round_trips_all_target_types(target, target_rank, payloads): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.TASK, + target=target, + target_rank=target_rank, + payloads=payloads, + ) + + request = mailbox.accept_request(last_sequence_id=0) + assert request.sequence_id == 1 + assert request.opcode is MailboxOpcode.TASK + assert request.target is target + assert request.target_rank == target_rank + assert request.payloads == payloads + assert mailbox.request_state is MailboxRequestState.TASK_ACCEPTED + + mailbox.complete_request(sequence_id=1, payloads=(b"ok",)) + result = mailbox.read_result(sequence_id=1) + assert result.payloads == (b"ok",) + assert mailbox.request_state is MailboxRequestState.IDLE + finally: + mailbox.close(unlink=True) + + +def test_accept_copies_payload_before_publishing_task_accepted(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.TASK, + target=MailboxTarget.RANK, + target_rank=0, + payloads=(b"immutable-copy",), + ) + request = mailbox.accept_request(last_sequence_id=0) + assert mailbox.request_state is MailboxRequestState.TASK_ACCEPTED + + mailbox.overwrite_request_payload_for_test(b"changed-after-accept") + assert request.payloads == (b"immutable-copy",) + finally: + mailbox.close(unlink=True) + + +def test_duplicate_sequence_marks_group_terminal(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=7, + opcode=MailboxOpcode.PING, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"",), + ) + with pytest.raises(MpiGroupError, match="sequence_id 7 is not newer than 7"): + mailbox.accept_request(last_sequence_id=7) + assert mailbox.group_state is MailboxGroupState.TERMINAL + finally: + mailbox.close(unlink=True) + + +def test_rank_failure_is_reported_with_rank_and_terminal_is_reusable_only_when_requested(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.TASK, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"task",), + ) + mailbox.accept_request(last_sequence_id=0) + mailbox.fail_request( + sequence_id=1, + errors=(MpiRankError(rank=1, error_type="ValueError", message="rank one failed"),), + terminal=False, + ) + with pytest.raises(MpiGroupError, match=r"rank 1.*ValueError.*rank one failed"): + mailbox.read_result(sequence_id=1) + assert mailbox.group_state is MailboxGroupState.READY + + mailbox.write_request( + sequence_id=2, + opcode=MailboxOpcode.PING, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"",), + ) + finally: + mailbox.close(unlink=True) + + +def test_terminal_group_rejects_new_requests(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.mark_terminal("collective timed out") + with pytest.raises(MpiGroupError, match="collective timed out"): + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.PING, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"",), + ) + finally: + mailbox.close(unlink=True) + + +def test_rank0_failure_is_reported_and_marks_terminal(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.CONTROL, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"control",), + ) + mailbox.accept_request(last_sequence_id=0) + mailbox.fail_request( + sequence_id=1, + errors=(MpiRankError(rank=0, error_type="RuntimeError", message="rank zero exited"),), + terminal=True, + ) + with pytest.raises(MpiGroupError, match=r"rank 0.*rank zero exited"): + mailbox.read_result(sequence_id=1) + assert mailbox.group_state is MailboxGroupState.TERMINAL + finally: + mailbox.close(unlink=True) + + +def test_shutdown_has_distinct_ready_and_done_states(): + mailbox = MpiGroupMailbox.create(world_size=2) + try: + mailbox.publish_ready() + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.SHUTDOWN, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"shutdown-frame",), + ) + assert mailbox.request_state is MailboxRequestState.SHUTDOWN_READY + request = mailbox.accept_request(last_sequence_id=0) + assert request.opcode is MailboxOpcode.SHUTDOWN + assert mailbox.request_state is MailboxRequestState.TASK_ACCEPTED + mailbox.complete_shutdown(sequence_id=1) + assert mailbox.request_state is MailboxRequestState.SHUTDOWN_DONE + finally: + mailbox.close(unlink=True) diff --git a/tests/ut/py/test_mpi_l3_group.py b/tests/ut/py/test_mpi_l3_group.py new file mode 100644 index 0000000000..d494132b61 --- /dev/null +++ b/tests/ut/py/test_mpi_l3_group.py @@ -0,0 +1,144 @@ +# 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 L3 group activation and transport-isolation tests.""" + +from __future__ import annotations + +import json +from typing import Any, cast +from unittest.mock import MagicMock + +import simpler.mpi_group_mailbox as mailbox_mod +import simpler.worker as worker_mod +from simpler.mpi_group_mailbox import MAILBOX_SIZE, MailboxGroupState +from simpler.remote_l3_protocol import ControlName +from simpler.worker import MpiL3GroupSpec, Worker + + +class _ReadyMailbox: + def __init__(self, world_size: int) -> None: + self.world_size = int(world_size) + self.group_state = MailboxGroupState.READY + self.address = 0x12340000 + self.closed = False + self.terminal_messages: list[str] = [] + + def manifest(self): + return { + "name": "pytest-mpi-mailbox", + "protocol_version": 1, + "mailbox_bytes": MAILBOX_SIZE, + "world_size": self.world_size, + } + + def terminal_reason(self): + return self.terminal_messages[-1] if self.terminal_messages else "" + + def mark_terminal(self, message): + self.group_state = MailboxGroupState.TERMINAL + self.terminal_messages.append(str(message)) + + def close(self, *, unlink=False): + assert unlink + self.closed = True + + +class _FakeProcess: + pid = 4242 + returncode = None + + def poll(self): + return self.returncode + + +class _InertThread: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + def start(self): + return None + + def join(self, timeout=None): + return None + + +def _mpi_worker() -> Worker: + worker = Worker(level=4, num_sub_workers=0) + worker.add_mpirun_worker_group( + MpiL3GroupSpec( + hosts=("127.0.0.1", "127.0.0.1"), + platform="a2a3sim", + device_ids_by_rank=((0,), (1,)), + global_device_ranks_by_rank=((0,), (1,)), + ) + ) + return worker + + +def test_mpi_activation_attaches_only_mailbox_endpoint(monkeypatch): + worker = _mpi_worker() + mailbox = _ReadyMailbox(world_size=2) + native_worker = MagicMock() + worker._worker = native_worker + monkeypatch.setattr(mailbox_mod.MpiGroupMailbox, "create", lambda **_kwargs: mailbox) + monkeypatch.setattr(worker_mod.subprocess, "Popen", lambda *_args, **_kwargs: _FakeProcess()) + monkeypatch.setattr(worker_mod.threading, "Thread", _InertThread) + try: + worker._activate_mpirun_worker_groups(worker_mod.time.monotonic() + 10.0) + native_worker.add_mpi_group_mailbox.assert_called_once() + native_worker.add_remote_l3_socket.assert_not_called() + group = worker._mpi_l3_groups[0] + assert group.manifest_path is not None + with open(group.manifest_path, encoding="utf-8") as manifest_file: + manifest = json.load(manifest_file) + assert manifest["mailbox"]["name"] == "pytest-mpi-mailbox" + assert "command_port" not in json.dumps(manifest) + assert "health_port" not in json.dumps(manifest) + assert "listen_host" not in json.dumps(manifest) + assert "connect_host" not in json.dumps(manifest) + finally: + group = worker._mpi_l3_groups[0] + group.process = None + group.monitor_thread = None + worker._worker = None + worker.close() + assert mailbox.closed + + +def test_unexpected_mpirun_exit_marks_group_terminal(): + worker = _mpi_worker() + mailbox = _ReadyMailbox(world_size=2) + group = worker._mpi_l3_groups[0] + group.mailbox = mailbox + process = _FakeProcess() + process.returncode = 17 + process.wait = lambda: 17 + worker._monitor_mpirun_group(group, cast(Any, process)) + try: + assert mailbox.group_state is MailboxGroupState.TERMINAL + assert "status 17" in mailbox.terminal_reason() + finally: + group.process = None + group.mailbox = None + worker.close() + + +def test_mpi_spec_tcp_fields_are_optional_and_control_18_stays_reserved(): + spec = MpiL3GroupSpec( + hosts=("127.0.0.1", "127.0.0.1"), + platform="a2a3sim", + device_ids_by_rank=((0,), (1,)), + ) + assert spec.command_port_base is None + assert spec.health_port_base is None + assert spec.session_listen_hosts == () + assert spec.connect_hosts == () + assert worker_mod._CTRL_COMMITTED_DEVICE_MEMORY == 18 + assert 18 not in {int(control) for control in ControlName} diff --git a/tests/ut/py/test_remote_l3_lifecycle.py b/tests/ut/py/test_remote_l3_lifecycle.py index acd644cb3d..e06e206847 100644 --- a/tests/ut/py/test_remote_l3_lifecycle.py +++ b/tests/ut/py/test_remote_l3_lifecycle.py @@ -268,7 +268,7 @@ def close(self): monkeypatch.setattr(remote_l3_session, "Worker", FakeWorker) monkeypatch.setattr(remote_l3_session, "_install_manifest_dispatcher_registry", lambda manifest: {}) monkeypatch.setattr(remote_l3_session, "_install_manifest_inner_registry", lambda manifest, worker: {}) - monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda host: sockets.pop(0)) + monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda _host, _port=0: sockets.pop(0)) monkeypatch.setattr(remote_l3_session, "_health_loop", lambda *args: None) try: @@ -343,7 +343,7 @@ def close(self): monkeypatch.setattr(remote_l3_session, "Worker", FakeWorker) monkeypatch.setattr(remote_l3_session, "_install_manifest_dispatcher_registry", lambda manifest: {}) monkeypatch.setattr(remote_l3_session, "_install_manifest_inner_registry", lambda manifest, worker: {}) - monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda host: sockets.pop(0)) + monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda _host, _port=0: sockets.pop(0)) monkeypatch.setattr(remote_l3_session, "_health_loop", lambda *args: None) monkeypatch.setattr(remote_l3_session, "_run_command_loop", lambda *args, **kwargs: None) @@ -403,7 +403,7 @@ def close(self): monkeypatch.setattr(remote_l3_session, "Worker", FakeWorker) monkeypatch.setattr(remote_l3_session, "_install_manifest_dispatcher_registry", lambda manifest: {}) monkeypatch.setattr(remote_l3_session, "_install_manifest_inner_registry", lambda manifest, worker: {}) - monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda host: sockets.pop(0)) + monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda _host, _port=0: sockets.pop(0)) monkeypatch.setattr(remote_l3_session, "_health_loop", lambda *args: None) try: @@ -464,7 +464,7 @@ def close(self): monkeypatch.setattr(remote_l3_session, "Worker", FakeWorker) monkeypatch.setattr(remote_l3_session, "_install_manifest_dispatcher_registry", slow_dispatch_registry) monkeypatch.setattr(remote_l3_session, "_install_manifest_inner_registry", lambda manifest, worker: {}) - monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda host: sockets.pop(0)) + monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda _host, _port=0: sockets.pop(0)) monkeypatch.setattr(remote_l3_session, "_health_loop", lambda *args: None) try: @@ -524,7 +524,7 @@ def close(self): monkeypatch.setattr(remote_l3_session, "Worker", FakeWorker) monkeypatch.setattr(remote_l3_session, "_install_manifest_dispatcher_registry", slow_dispatch_registry) monkeypatch.setattr(remote_l3_session, "_install_manifest_inner_registry", lambda manifest, worker: {}) - monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda host: sockets.pop(0)) + monkeypatch.setattr(remote_l3_session, "_bind_listener", lambda _host, _port=0: sockets.pop(0)) monkeypatch.setattr(remote_l3_session, "_health_loop", lambda *args: None) try: diff --git a/tools/a3_l4_tcp_smoke/README.md b/tools/a3_l4_tcp_smoke/README.md new file mode 100644 index 0000000000..fc4d12e6eb --- /dev/null +++ b/tools/a3_l4_tcp_smoke/README.md @@ -0,0 +1,29 @@ +# A3 L4 MPI mailbox smoke + +`mpirun_compute_then_tload_2x2_smoke.py` validates the current two-server +topology: + +- server 37: L4 plus MPI rank 0/L3, managing two L2 devices; +- server 35: MPI rank 1/L3, managing two L2 devices; +- four global device ranks perform local addition, then peer TLOAD. + +The L4 process creates one named shared-memory mailbox. Full-group submissions +use `submit_next_level_group`, so one `PER_RANK` request carries a distinct +payload for each MPI rank. No Simpler command or health TCP endpoint is +created. OpenMPI, SSH, or the MPI implementation may still use TCP internally. + +Example from server 37: + +```bash +python tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py \ + --host-37 "$RANK0_SSH_HOST" \ + --host-35 "$RANK1_SSH_HOST" \ + --roce-37 "$RANK0_ROCE_INTERFACES" \ + --roce-35 "$RANK1_ROCE_INTERFACES" \ + --devices-37 0,1 \ + --devices-35 0,1 +``` + +Success requires all four compute and TLOAD `max_diff` values to stay within +their tolerances. Shutdown also checks that `mpirun`, the named shared-memory +object, and the temporary manifest are gone. diff --git a/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp b/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp new file mode 100644 index 0000000000..231a75a6a1 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp @@ -0,0 +1,86 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "platform_comm/comm_context.h" +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static constexpr size_t kCount = 256; +static constexpr int kMaxRanks = 16; + +template +AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *local_ptr, int peer_rank) { + uint64_t local_base = ctx->windowsIn[ctx->rankId]; + uint64_t offset = reinterpret_cast(local_ptr) - local_base; + return reinterpret_cast<__gm__ T *>(ctx->windowsIn[peer_rank] + offset); +} + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + int rank_count = static_cast(args[2]); + __gm__ CommContext *comm_ctx = reinterpret_cast<__gm__ CommContext *>(args[3]); + + if (rank_count <= 0 || rank_count > kMaxRanks) { + pipe_barrier(PIPE_ALL); + return; + } + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset; + + using Shape = pto::Shape; + using Stride = pto::Stride; + using Global = pto::GlobalTensor; + using Tile = pto::Tile; + + Shape shape(1, 1, 1, 1, kCount); + Stride stride(kCount, kCount, kCount, kCount, 1); + Tile accumulator(1, kCount); + Tile peer_tile(1, kCount); + TASSIGN(accumulator, 0x0); + TASSIGN(peer_tile, 0x10000); + + Global local_input(input, shape, stride); + TLOAD(accumulator, local_input); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + int my_rank = static_cast(comm_ctx->rankId); + for (int peer = 0; peer < rank_count; ++peer) { + if (peer == my_rank) continue; + __gm__ float *peer_input = CommRemotePtr(comm_ctx, input, peer); + Global peer_global(peer_input, shape, stride); + TLOAD(peer_tile, peer_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TADD(accumulator, accumulator, peer_tile); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + } + + Global result_global(result, shape, stride); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(result_global, accumulator); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pipe_barrier(PIPE_ALL); +} diff --git a/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp b/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp new file mode 100644 index 0000000000..9d50f4f4a1 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp @@ -0,0 +1,61 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static constexpr size_t kCount = 256; + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *lhs_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *rhs_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *lhs = reinterpret_cast<__gm__ float *>(lhs_tensor->buffer.addr) + lhs_tensor->start_offset; + __gm__ float *rhs = reinterpret_cast<__gm__ float *>(rhs_tensor->buffer.addr) + rhs_tensor->start_offset; + __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset; + + using Shape = pto::Shape; + using Stride = pto::Stride; + using Global = pto::GlobalTensor; + using Tile = pto::Tile; + + Shape shape(1, 1, 1, 1, kCount); + Stride stride(kCount, kCount, kCount, kCount, 1); + Tile lhs_tile(1, kCount); + Tile rhs_tile(1, kCount); + Tile result_tile(1, kCount); + TASSIGN(lhs_tile, 0x0); + TASSIGN(rhs_tile, 0x10000); + TASSIGN(result_tile, 0x20000); + + Global lhs_global(lhs, shape, stride); + Global rhs_global(rhs, shape, stride); + Global result_global(result, shape, stride); + TLOAD(lhs_tile, lhs_global); + TLOAD(rhs_tile, rhs_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(result_tile, lhs_tile, rhs_tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(result_global, result_tile); + pipe_barrier(PIPE_ALL); +} diff --git a/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp b/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp new file mode 100644 index 0000000000..fe9c2da937 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp @@ -0,0 +1,35 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +global_tload_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 4, + }; +} + +__attribute__((visibility("default"))) void global_tload_orchestration(const L2TaskArgs &orch_args) { + L0TaskArgs params; + params.add_input(orch_args.tensor(0).ref()); + params.add_output(orch_args.tensor(1).ref()); + params.add_scalar(orch_args.scalar(0)); + params.add_scalar(orch_args.scalar(1)); + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp b/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp new file mode 100644 index 0000000000..1c0a2d81e9 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp @@ -0,0 +1,34 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +local_add_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +__attribute__((visibility("default"))) void local_add_orchestration(const L2TaskArgs &orch_args) { + L0TaskArgs params; + params.add_input(orch_args.tensor(0).ref()); + params.add_input(orch_args.tensor(1).ref()); + params.add_output(orch_args.tensor(2).ref()); + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py b/tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py new file mode 100644 index 0000000000..87a7d19229 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py @@ -0,0 +1,326 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Run the mailbox->MPI two-L3/four-L2 A3 compute + peer-TLOAD smoke on servers 37 and 35.""" + +from __future__ import annotations + +import argparse +import os +import struct +import sys +from multiprocessing.shared_memory import SharedMemory + +from simpler.task_interface import ArgDirection, CallConfig, ChipCallable, CommBufferSpec, CoreCallable, TaskArgs +from simpler.worker import MpiL3GroupSpec, RemoteCallable, Worker + +from simpler_setup.elf_parser import extract_text_section +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.pto_isa import ensure_pto_isa_root + +HERE = os.path.dirname(os.path.abspath(__file__)) +LOCAL_ADD_AIV = os.path.join(HERE, "kernels", "aiv", "local_add_kernel.cpp") +LOCAL_ADD_ORCH = os.path.join(HERE, "kernels", "orchestration", "local_add_orch.cpp") +GLOBAL_TLOAD_AIV = os.path.join(HERE, "kernels", "aiv", "global_tload_kernel.cpp") +GLOBAL_TLOAD_ORCH = os.path.join(HERE, "kernels", "orchestration", "global_tload_orch.cpp") +COUNT = 256 +FLOAT_BYTES = 4 +WINDOW_SIZE = 4096 + + +def _csv_ints(raw: str) -> tuple[int, ...]: + values = tuple(int(item.strip()) for item in raw.split(",") if item.strip()) + if len(values) != 2: + raise ValueError("this smoke requires exactly two devices on each server") + return values + + +def _csv_strings(raw: str) -> tuple[str, ...]: + values = tuple(item.strip() for item in raw.split(",") if item.strip()) + if len(values) != 2: + raise ValueError("this smoke requires exactly two RoCE addresses on each server") + return values + + +def _compile_aiv(compiler: KernelCompiler, platform: str, runtime: str, source: str) -> bytes: + include_dirs = compiler.get_orchestration_include_dirs(runtime) + kernel_include_dirs = list(include_dirs) + [str(compiler.project_root / "src" / "common")] + binary = compiler.compile_incore( + source_path=source, + core_type="aiv", + pto_isa_root=ensure_pto_isa_root(), + extra_include_dirs=kernel_include_dirs, + ) + return binary if platform.endswith("sim") else extract_text_section(binary) + + +def _build_chip_callable( + platform: str, + runtime: str, + kernel_source: str, + orchestration_source: str, + signature: list[ArgDirection], + func_name: str, + config_name: str, +) -> ChipCallable: + compiler = KernelCompiler(platform=platform) + core = CoreCallable.build( + signature=signature, + binary=_compile_aiv(compiler, platform, runtime, kernel_source), + ) + return ChipCallable.build( + signature=signature, + func_name=func_name, + config_name=config_name, + binary=compiler.compile_orchestration(runtime_name=runtime, source_path=orchestration_source), + children=[(0, core)], + ) + + +def _digest_scalars(digest: bytes) -> tuple[int, ...]: + if len(digest) != 32: + raise ValueError("callable digest must be 32 bytes") + return tuple(int.from_bytes(digest[offset : offset + 8], "little") for offset in range(0, 32, 8)) + + +def _rank_args(domain_id: int, local_worker_count: int, digest: bytes) -> TaskArgs: + args = TaskArgs() + args.add_scalar(domain_id) + args.add_scalar(local_worker_count) + for value in _digest_scalars(digest): + args.add_scalar(value) + return args + + +def _lhs(rank: int) -> tuple[float, ...]: + return tuple(float(rank * 100 + index) for index in range(COUNT)) + + +def _rhs(rank: int) -> tuple[float, ...]: + return tuple(float(rank * 10 + 2 * index) for index in range(COUNT)) + + +def _compute_expected(rank: int) -> tuple[float, ...]: + return tuple(a + b for a, b in zip(_lhs(rank), _rhs(rank), strict=True)) + + +def _tload_expected(rank_count: int) -> tuple[float, ...]: + by_rank = tuple(_compute_expected(rank) for rank in range(rank_count)) + return tuple(sum(values[index] for values in by_rank) for index in range(COUNT)) + + +def _unpack(raw: bytes) -> tuple[float, ...]: + return tuple(float(value) for value in struct.unpack(f"<{COUNT}f", raw)) + + +def _max_diff(actual: tuple[float, ...], expected: tuple[float, ...]) -> float: + return max(abs(a - b) for a, b in zip(actual, expected, strict=True)) + + +def _assert_mailbox_unlinked(name: str) -> None: + try: + reopened = SharedMemory(name=name, create=False) + except FileNotFoundError: + return + reopened.close() + raise AssertionError(f"MPI group shared-memory object {name!r} remained after shutdown") + + +def run(args: argparse.Namespace) -> int: + devices37 = _csv_ints(args.devices_37) + devices35 = _csv_ints(args.devices_35) + roce37 = _csv_strings(args.roce_37) + roce35 = _csv_strings(args.roce_35) + print(f"[mpi-mailbox-2x2] rank0/L3 server37={args.host_37}, devices={devices37}, roce={roce37}") + print(f"[mpi-mailbox-2x2] rank1/L3 server35={args.host_35}, devices={devices35}, roce={roce35}") + + compute_callable = _build_chip_callable( + args.platform, + args.runtime, + LOCAL_ADD_AIV, + LOCAL_ADD_ORCH, + [ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + "local_add_orchestration", + "local_add_orchestration_config", + ) + tload_callable = _build_chip_callable( + args.platform, + args.runtime, + GLOBAL_TLOAD_AIV, + GLOBAL_TLOAD_ORCH, + [ArgDirection.IN, ArgDirection.OUT], + "global_tload_orchestration", + "global_tload_orchestration_config", + ) + mpirun_args = tuple(args.mpirun_arg) if args.mpirun_arg else ("--map-by", "ppr:1:node") + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=args.timeout_s) + node37, node35 = worker.add_mpirun_worker_group( + MpiL3GroupSpec( + hosts=(f"{args.host_37}:1", f"{args.host_35}:1"), + platform=args.platform, + device_ids_by_rank=(devices37, devices35), + runtime=args.runtime, + comm_profile="a3-fabric-v1", + global_device_ranks_by_rank=((0, 1), (2, 3)), + mpirun_path=args.mpirun_path, + mpirun_args=mpirun_args, + python_executable=args.python_executable, + ) + ) + compute_handle = worker.register(compute_callable) + tload_handle = worker.register(tload_callable) + rank_compute_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_compute_group_orch"), + workers=[node37, node35], + ) + rank_tload_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_rank_group_orch"), + workers=[node37, node35], + ) + targets = ((node37, 0), (node37, 1), (node35, 0), (node35, 1)) + captured: dict[str, object] = {} + compute_results: list[tuple[float, ...]] = [] + tload_results: list[tuple[float, ...]] = [] + mailbox_name = "" + manifest_path = "" + mpirun_process = None + try: + worker.init() + group = worker._mpi_l3_groups[0] + assert group.mailbox is not None + mailbox_name = group.mailbox.name + manifest_path = group.manifest_path + mpirun_process = group.process + if worker._remote_sessions: + raise AssertionError("MPI group unexpectedly created a Simpler TCP Remote L3 session") + print(f"[mpi-mailbox-2x2] mailbox={mailbox_name}; Simpler TCP sessions=0") + + def compute_phase(orch, _args, cfg): + domain = orch.allocate_global_domain( + name="a3-mpi-mailbox-2x2-compute-tload", + members=targets, + window_size=WINDOW_SIZE, + buffers=( + CommBufferSpec("lhs", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("rhs", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("input", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("result", "float32", COUNT, COUNT * FLOAT_BYTES), + ), + retain_after_run=True, + ) + for global_rank in range(4): + orch.copy_to_global_domain( + domain, + global_rank, + struct.pack(f"<{COUNT}f", *_lhs(global_rank)), + buffer="lhs", + ) + orch.copy_to_global_domain( + domain, + global_rank, + struct.pack(f"<{COUNT}f", *_rhs(global_rank)), + buffer="rhs", + ) + orch.submit_next_level_group( + rank_compute_handle, + [ + _rank_args(domain.domain_id, 2, compute_handle.digest), + _rank_args(domain.domain_id, 2, compute_handle.digest), + ], + cfg, + workers=[node37, node35], + ) + captured["domain"] = domain + + worker.run(compute_phase, args=None, config=CallConfig()) + domain = captured["domain"] + + def tload_phase(orch, _args, cfg): + for global_rank in range(4): + compute_results.append( + _unpack( + orch.copy_from_global_domain( + domain, + global_rank, + COUNT * FLOAT_BYTES, + buffer="input", + ) + ) + ) + orch.submit_next_level_group( + rank_tload_handle, + [ + _rank_args(domain.domain_id, 2, tload_handle.digest), + _rank_args(domain.domain_id, 2, tload_handle.digest), + ], + cfg, + workers=[node37, node35], + ) + + worker.run(tload_phase, args=None, config=CallConfig()) + + def verify_phase(orch, _args, _cfg): + try: + for global_rank in range(4): + tload_results.append( + _unpack( + orch.copy_from_global_domain( + domain, + global_rank, + COUNT * FLOAT_BYTES, + buffer="result", + ) + ) + ) + finally: + domain.release() + + worker.run(verify_phase, args=None, config=CallConfig()) + for global_rank, observed in enumerate(compute_results): + diff = _max_diff(observed, _compute_expected(global_rank)) + print(f"[mpi-mailbox-2x2] compute global_rank={global_rank} max_diff={diff:.3e}") + if diff > 1e-5: + raise AssertionError(f"compute mismatch on global rank {global_rank}: max_diff={diff}") + expected = _tload_expected(4) + for global_rank, observed in enumerate(tload_results): + diff = _max_diff(observed, expected) + print(f"[mpi-mailbox-2x2] TLOAD global_rank={global_rank} max_diff={diff:.3e}") + if diff > 1e-3: + raise AssertionError(f"TLOAD mismatch on global rank {global_rank}: max_diff={diff}") + print("[mpi-mailbox-2x2] PASS: one per-rank mailbox task drove both MPI ranks and all four L2 devices") + return 0 + finally: + worker.close() + if mpirun_process is not None and mpirun_process.poll() is None: + raise AssertionError("mpirun remained alive after Worker.close()") + if mailbox_name: + _assert_mailbox_unlinked(mailbox_name) + if manifest_path and os.path.exists(manifest_path): + raise AssertionError(f"MPI group manifest remained after shutdown: {manifest_path}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host-37", required=True) + parser.add_argument("--host-35", required=True) + parser.add_argument("--roce-37", required=True) + parser.add_argument("--roce-35", required=True) + parser.add_argument("--devices-37", default="0,1") + parser.add_argument("--devices-35", default="0,1") + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + parser.add_argument("--timeout-s", type=float, default=180.0) + parser.add_argument("--mpirun-path", default="mpirun") + parser.add_argument("--mpirun-arg", action="append", default=[]) + parser.add_argument("--python-executable", default=sys.executable) + return run(parser.parse_args()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/mpi_group_mailbox_smoke.py b/tools/mpi_group_mailbox_smoke.py new file mode 100644 index 0000000000..996c901904 --- /dev/null +++ b/tools/mpi_group_mailbox_smoke.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Device-free two-rank integration smoke for the named mailbox and MPI dispatch loop.""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import time + +from simpler.mpi_group_mailbox import ( + MailboxGroupState, + MailboxOpcode, + MailboxRequestState, + MailboxTarget, + MpiGroupError, + MpiGroupMailbox, + MpiRankError, + open_rank_mailbox, +) + + +def _wait_until(predicate, *, deadline: float, label: str) -> None: + while not predicate(): + if time.monotonic() >= deadline: + raise TimeoutError(f"timed out waiting for {label}") + + +def _worker(manifest_path: str) -> int: # noqa: PLR0912 -- mirrors the ordered rank dispatcher state machine + from mpi4py import MPI # noqa: PLC0415 + + comm = MPI.COMM_WORLD.Dup() + rank = int(comm.Get_rank()) + if int(comm.Get_size()) != 2: + raise RuntimeError("MPI mailbox smoke requires exactly two ranks") + if rank == 0: + with open(manifest_path, encoding="utf-8") as manifest_file: + manifest = json.load(manifest_file) + else: + manifest = None + manifest = comm.bcast(manifest, root=0) + mailbox = None + try: + ready = comm.allgather(rank) + if ready != [0, 1]: + raise RuntimeError(f"unexpected readiness result: {ready}") + mailbox = open_rank_mailbox(manifest, rank=rank) + if rank == 0: + assert mailbox is not None + mailbox.publish_ready() + last_sequence = 0 + while True: + if rank == 0: + assert mailbox is not None + _wait_until( + lambda: mailbox.request_state + in (MailboxRequestState.REQUEST_READY, MailboxRequestState.SHUTDOWN_READY), + deadline=time.monotonic() + 10.0, + label="mailbox request", + ) + request = mailbox.accept_request(last_sequence_id=last_sequence) + last_sequence = request.sequence_id + else: + request = None + request = comm.bcast(request, root=0) + if request.opcode is MailboxOpcode.SHUTDOWN: + gathered = comm.gather((rank, True, b"", None), root=0) + if rank == 0: + assert gathered is not None + assert mailbox is not None + mailbox.complete_shutdown(sequence_id=request.sequence_id) + mailbox.publish_closed() + break + + payload = request.payloads[rank] if request.target is MailboxTarget.PER_RANK else request.payloads[0] + local_error = None + local_result = b"" + try: + command = json.loads(payload.decode()) + if command.get("fail_rank") == rank: + raise ValueError(f"injected failure on rank {rank}") + local_result = json.dumps( + { + "rank": rank, + "input": int(command["value"]), + "result": int(command["value"]) * (rank + 2), + }, + sort_keys=True, + ).encode() + except BaseException as exc: # noqa: BLE001 + local_error = MpiRankError(rank, type(exc).__name__, str(exc)) + gathered = comm.gather((rank, True, local_result, local_error), root=0) + if rank == 0: + assert mailbox is not None + errors = tuple(item[3] for item in gathered if item[3] is not None) + if errors: + mailbox.fail_request(sequence_id=request.sequence_id, errors=errors, terminal=False) + else: + mailbox.complete_request( + sequence_id=request.sequence_id, + payloads=tuple(item[2] for item in gathered), + ) + return 0 + finally: + if mailbox is not None: + mailbox.close() + comm.Free() + + +def _parent(mpirun: str, python: str) -> int: + mailbox = MpiGroupMailbox.create(world_size=2) + process = None + with tempfile.TemporaryDirectory(prefix="simpler-mpi-mailbox-smoke-") as temp_dir: + manifest_path = os.path.join(temp_dir, "mailbox.json") + with open(manifest_path, "w", encoding="utf-8") as manifest_file: + json.dump(mailbox.manifest(), manifest_file) + cmd = [mpirun, "-np", "2", python, os.path.abspath(__file__), "--worker", manifest_path] + try: + process = subprocess.Popen(cmd, start_new_session=True) + deadline = time.monotonic() + 20.0 + _wait_until( + lambda: mailbox.group_state is not MailboxGroupState.INITIALIZING or process.poll() is not None, + deadline=deadline, + label="MPI mailbox READY", + ) + if process.poll() is not None: + raise RuntimeError(f"mpirun exited during startup with status {process.returncode}") + if mailbox.group_state is not MailboxGroupState.READY: + raise RuntimeError(f"MPI mailbox failed during startup: {mailbox.terminal_reason()}") + + mailbox.write_request( + sequence_id=1, + opcode=MailboxOpcode.TASK, + target=MailboxTarget.PER_RANK, + target_rank=-1, + payloads=(b'{"value": 7}', b'{"value": 11}'), + ) + _wait_until( + lambda: mailbox.request_state in (MailboxRequestState.TASK_DONE, MailboxRequestState.TASK_FAILED), + deadline=time.monotonic() + 10.0, + label="per-rank task", + ) + result = mailbox.read_result(sequence_id=1) + decoded = tuple(json.loads(payload.decode()) for payload in result.payloads) + assert decoded == ( + {"input": 7, "rank": 0, "result": 14}, + {"input": 11, "rank": 1, "result": 33}, + ) + print(f"[mpi-mailbox-local] per-rank results: {decoded}") + + mailbox.write_request( + sequence_id=2, + opcode=MailboxOpcode.TASK, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b'{"value": 5, "fail_rank": 1}',), + ) + _wait_until( + lambda: mailbox.request_state in (MailboxRequestState.TASK_DONE, MailboxRequestState.TASK_FAILED), + deadline=time.monotonic() + 10.0, + label="rank failure", + ) + try: + mailbox.read_result(sequence_id=2) + except MpiGroupError as exc: + if "rank 1" not in str(exc): + raise + print(f"[mpi-mailbox-local] expected rank failure: {exc}") + else: + raise AssertionError("rank-1 failure unexpectedly succeeded") + + mailbox.write_request( + sequence_id=3, + opcode=MailboxOpcode.SHUTDOWN, + target=MailboxTarget.GROUP, + target_rank=-1, + payloads=(b"shutdown",), + ) + _wait_until( + lambda: mailbox.request_state is MailboxRequestState.SHUTDOWN_DONE, + deadline=time.monotonic() + 10.0, + label="shutdown", + ) + process.wait(timeout=10.0) + if process.returncode != 0: + raise RuntimeError(f"mpirun returned status {process.returncode}") + print("[mpi-mailbox-local] PASS: READY, one per-rank request, ranked failure, and shutdown") + return 0 + finally: + if process is not None and process.poll() is None: + if os.name == "posix": + os.killpg(process.pid, 9) + else: + process.kill() + process.wait(timeout=5.0) + mailbox.close(unlink=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mpirun", default="mpirun") + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--worker") + args = parser.parse_args() + if args.worker: + return _worker(args.worker) + return _parent(args.mpirun, args.python) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/mpi_l3_group_smoke.py b/tools/mpi_l3_group_smoke.py new file mode 100644 index 0000000000..72189bff58 --- /dev/null +++ b/tools/mpi_l3_group_smoke.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Device-free integration smoke for the real L4->mailbox->MPI L3 command path.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import tempfile +from multiprocessing.shared_memory import SharedMemory + +from simpler.task_interface import CallConfig, TaskArgs +from simpler.worker import MpiL3GroupSpec, RemoteCallable, Worker + + +def _args(value: int) -> TaskArgs: + args = TaskArgs() + args.add_scalar(value) + return args + + +def _assert_unlinked(name: str) -> None: + try: + reopened = SharedMemory(name=name, create=False) + except FileNotFoundError: + return + reopened.close() + raise AssertionError(f"MPI L3 mailbox {name!r} remained after shutdown") + + +def run(args: argparse.Namespace) -> int: # noqa: PLR0915 -- integration lifecycle stays visible in one function + with tempfile.TemporaryDirectory(prefix="simpler-mpi-l3-group-smoke-") as output_dir: + os.environ["SIMPLER_MPI_SMOKE_DIR"] = output_dir + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=args.timeout_s) + worker_ids = worker.add_mpirun_worker_group( + MpiL3GroupSpec( + hosts=("localhost", "localhost"), + platform="a2a3sim", + device_ids_by_rank=((0,), (1,)), + num_sub_workers_by_rank=(0, 0), + global_device_ranks_by_rank=((0,), (1,)), + mpirun_path=args.mpirun, + mpirun_args=( + "--oversubscribe", + "--mca", + "btl", + "self,vader,tcp", + "--mca", + "pml", + "ob1", + ), + python_executable=args.python, + ) + ) + callback = worker.register( + RemoteCallable("simpler.mpi_group_smoke:record_rank_value"), + workers=list(worker_ids), + ) + mailbox_name = "" + manifest_path = "" + process = None + try: + worker.init() + group = worker._mpi_l3_groups[0] + assert group.mailbox is not None + assert group.manifest_path is not None + mailbox_name = group.mailbox.name + manifest_path = group.manifest_path + process = group.process + if worker._remote_sessions: + raise AssertionError("MPI L3 group unexpectedly created a Simpler TCP session") + with open(manifest_path, encoding="utf-8") as manifest_file: + manifest_text = manifest_file.read() + for forbidden in ("command_port", "health_port", "listen_host", "connect_host"): + if forbidden in manifest_text: + raise AssertionError(f"MPI group manifest unexpectedly contains {forbidden}") + + def submit_values(orch, _run_args, cfg): + orch.submit_next_level_group( + callback, + [_args(101), _args(202)], + cfg, + workers=list(worker_ids), + ) + + worker.run(submit_values, args=None, config=CallConfig()) + observed = [] + for rank in range(2): + with open(os.path.join(output_dir, f"rank-{rank}.json"), encoding="utf-8") as output_file: + observed.append(json.load(output_file)) + if observed != [{"rank": 0, "value": 101}, {"rank": 1, "value": 202}]: + raise AssertionError(f"unexpected per-rank callback results: {observed}") + print(f"[mpi-l3-group-local] per-rank results: {observed}") + + def submit_failure(orch, _run_args, cfg): + orch.submit_next_level_group( + callback, + [_args(303), _args(0xFFFF)], + cfg, + workers=list(worker_ids), + ) + + try: + worker.run(submit_failure, args=None, config=CallConfig()) + except RuntimeError as exc: + if "rank 1" not in str(exc): + raise + print(f"[mpi-l3-group-local] expected rank failure: {exc}") + else: + raise AssertionError("rank-1 callback failure unexpectedly succeeded") + + worker.run(submit_values, args=None, config=CallConfig()) + print("[mpi-l3-group-local] PASS: real endpoint batching, ranked failure, reuse, and no Simpler TCP") + return 0 + finally: + worker.close() + if process is not None and process.poll() is None: + raise AssertionError("mpirun remained alive after MPI L3 group shutdown") + if mailbox_name: + _assert_unlinked(mailbox_name) + if manifest_path and os.path.exists(manifest_path): + raise AssertionError(f"MPI L3 group manifest remained after shutdown: {manifest_path}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mpirun", default="mpirun") + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--timeout-s", type=float, default=30.0) + return run(parser.parse_args()) + + +if __name__ == "__main__": + sys.exit(main())