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/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_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..8333f146ae --- /dev/null +++ b/python/simpler/mpi_group_mailbox.py @@ -0,0 +1,479 @@ +# 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) + + +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[:] = b"\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._buffer[MAILBOX_REQUEST_OFFSET : MAILBOX_REQUEST_OFFSET + len(data)] = 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 = bytes(self._buffer[MAILBOX_REQUEST_OFFSET : 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._buffer[MAILBOX_RESPONSE_OFFSET : MAILBOX_RESPONSE_OFFSET + len(data)] = 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._buffer[MAILBOX_ERROR_OFFSET : MAILBOX_ERROR_OFFSET + len(data)] = 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 = bytes(self._buffer[MAILBOX_RESPONSE_OFFSET : 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 = bytes(self._buffer[MAILBOX_ERROR_OFFSET : 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._buffer[MAILBOX_ERROR_OFFSET : MAILBOX_ERROR_OFFSET + len(data)] = 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 bytes(self._buffer[MAILBOX_ERROR_OFFSET : MAILBOX_ERROR_OFFSET + error_bytes]).decode( + "utf-8", errors="replace" + ) + + def overwrite_request_payload_for_test(self, data: bytes) -> None: + value = bytes(data) + self._buffer[MAILBOX_REQUEST_OFFSET : MAILBOX_REQUEST_OFFSET + len(value)] = 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 buffer + + 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..afb67c04b9 100644 --- a/python/simpler/mpi_l3_session.py +++ b/python/simpler/mpi_l3_session.py @@ -6,24 +6,20 @@ # 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 import argparse 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 +30,34 @@ 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 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 @@ -162,24 +180,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 +281,294 @@ 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 + exit_code = 0 + try: + while True: + if rank == 0: + assert mailbox is not None + while mailbox.request_state not in ( + MailboxRequestState.REQUEST_READY, + MailboxRequestState.SHUTDOWN_READY, + ): + if mailbox.group_state is MailboxGroupState.TERMINAL: + break + 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 BaseException: + 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 = _payload_for_rank(request, rank) + try: + 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 +580,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/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/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..9c40f69d66 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,420 @@ 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_); + 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"); + } + } +} + +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; + 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; + 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 +1041,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 +1104,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/task.md b/task.md new file mode 100644 index 0000000000..914adbb3c7 --- /dev/null +++ b/task.md @@ -0,0 +1,60 @@ +# Task Record + +## Current Task + +- Summary: Implement mailbox and MPI dispatch for MPI worker groups on PR #1623. +- Status: completed +- Local Branch: `skx/mpi-mailbox-broadcast` +- Delivery Branch: `origin/use-mpi-and-remove-socket` +- Validation Hosts: servers 37 and 35 (external handoff required) +- Baseline PR Head: `d3d17b21593e642c6c5f3956ab163589cdb2bbe1` +- Latest Tested Tree: final local branch tree; final commit SHA is reported at handoff +- Last Updated: 2026-08-01 + +## Work Completed + +- Read the repository instructions, always-on rules, and relevant workflows. +- Refreshed `upstream/pr-1623` to GitHub head `d3d17b21` and created an + independent branch without rebasing or modifying `add-mpi-run`. +- Used this worktree's project `.venv`. +- Added a versioned named shared-memory mailbox protocol with explicit READY, + accepted, done, failed, shutdown, terminal, sequence, target, and payload + semantics. +- Added the C++ mailbox transport and endpoint. A complete MPI group submission + becomes one per-rank request; directed and subset submissions remain directed. +- Replaced the MPI group's Simpler command/health TCP activation with a rank-0 + mailbox and a single-threaded MPI dispatcher using separate dispatch and + Global CommDomain communicators. +- Added ranked result aggregation, terminal timeout behavior, direct `mpirun` + process-group monitoring, and mailbox/manifest cleanup. +- Preserved the ordinary non-MPI Remote L3 socket path. +- Added unit tests, a device-free real-endpoint smoke, and the A3 2x2 compute + plus global TLOAD smoke. + +## Validation + +- Local Windows named-mailbox unit tests: + `10 passed in 0.56s`. +- Staged-file hooks passed: headers, English-only, large files, EOF, + whitespace, markdownlint, Ruff, formatting, and Pyright. +- Clang-tidy is unavailable on this Windows host because the repository setup + imports POSIX-only modules. +- The tests that import the Linux C++ extension cannot be collected on this + Windows host. +- Single-host real `mpirun`, two-host MPI, and A3/NPU smoke are **NOT + VALIDATED** here. Per user direction, `myserver` is not a validation host. +- The official MPI and A3 validation must run on servers 37 and 35 using the + separately delivered agent prompt. No result may be called passing until + its log is returned. + +## Protocol Conclusions + +1. MPI group operations never fall back to Simpler TCP. A mailbox/MPI failure + terminates the group. +2. Full worker-id sets in `submit_next_level_group` use one `PER_RANK` + mailbox request. One worker ID means directed execution; a subset remains + an ordered set of directed requests. +3. Remote controls remain uniquely numbered 1 through 17. Local control 18 + remains reserved for committed-device-memory handling. +4. Bare host pointers, child-memory pointers without a transferable sidecar, + and other unsafe remote addresses are rejected before dispatch. 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..aefe69af60 100644 --- a/tests/ut/py/test_global_comm_domain.py +++ b/tests/ut/py/test_global_comm_domain.py @@ -286,8 +286,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 +394,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: 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..81c6005113 --- /dev/null +++ b/tests/ut/py/test_mpi_group_mailbox.py @@ -0,0 +1,211 @@ +# 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 pytest +from simpler.mpi_group_mailbox import ( + MailboxGroupState, + MailboxOpcode, + MailboxRequestState, + MailboxTarget, + MpiGroupError, + MpiGroupMailbox, + MpiRankError, + open_rank_mailbox, +) + + +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/tools/a3_l4_tcp_smoke/README.md b/tools/a3_l4_tcp_smoke/README.md new file mode 100644 index 0000000000..20cd8bf069 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/README.md @@ -0,0 +1,27 @@ +# 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 120.9.10.37 \ + --host-35 120.9.10.35 \ + --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..aff6015727 --- /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", default="120.9.10.37") + parser.add_argument("--host-35", default="120.9.10.35") + parser.add_argument("--roce-37", default="10.30.2.1,10.30.2.2") + parser.add_argument("--roce-35", default="10.30.0.1,10.30.0.2") + 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())