diff --git a/CMakeLists.txt b/CMakeLists.txt index 842be96b0b..1bb4134c06 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -67,6 +67,8 @@ add_custom_target(build_package_sim if(SKBUILD_MODE) install(DIRECTORY ${CMAKE_SOURCE_DIR}/src/ DESTINATION simpler_setup/_assets/src) + install(DIRECTORY ${CMAKE_SOURCE_DIR}/cmake/ + DESTINATION simpler_setup/_assets/cmake) install(DIRECTORY ${CMAKE_SOURCE_DIR}/build/lib/ DESTINATION simpler_setup/_assets/build/lib OPTIONAL diff --git a/docs/python-packaging.md b/docs/python-packaging.md index 06c094ad6b..8b7ca1da03 100644 --- a/docs/python-packaging.md +++ b/docs/python-packaging.md @@ -31,6 +31,7 @@ simpler_setup/ ← test framework + build/runtime assembly (simpler_s paged_attention.py attention reference (used by multiple paged_attention tests) _assets/ (wheel-only, populated by CMake install) src/ source tree mirror + cmake/ shared modules used by runtime compilation build/lib/ pre-built per-arch/platform/runtime .so/.o _task_interface.*.so nanobind extension at site-packages root @@ -69,10 +70,10 @@ NPU hardware (`a2a3`/`a5` with CANN toolkit). `simpler_setup.environment.PROJECT_ROOT` auto-detects between: -- **Wheel install**: `simpler_setup/_assets/` exists → `PROJECT_ROOT = .../site-packages/simpler_setup/_assets`. The wheel's bundled `_assets/src/` and `_assets/build/lib/` provide everything needed at runtime. +- **Wheel install**: `simpler_setup/_assets/` exists → `PROJECT_ROOT = .../site-packages/simpler_setup/_assets`. The wheel's bundled `_assets/src/`, `_assets/cmake/`, and `_assets/build/lib/` provide everything needed at runtime. - **Source tree / editable install**: `_assets/` doesn't exist → `PROJECT_ROOT = repo root`. Live `src/` and `build/lib/` are used. -Anything that needs to find `src/`, `build/lib/`, or `build/cache/` MUST go through `simpler_setup.environment.PROJECT_ROOT` — never `Path(__file__).parent.parent...`. +Anything that needs to find `src/`, `cmake/`, `build/lib/`, or `build/cache/` MUST go through `simpler_setup.environment.PROJECT_ROOT` — never `Path(__file__).parent.parent...`. ## Import rules diff --git a/examples/workers/README.md b/examples/workers/README.md index 6516dff1df..32cb42b7c0 100644 --- a/examples/workers/README.md +++ b/examples/workers/README.md @@ -38,6 +38,7 @@ workers/ global_tload_mixed_l3/ # Global CommDomain build + cross-machine peer TLOAD on both ranks compute_then_tload_mixed_l3/ # compute round on both L2s, then peer TLOAD through the same domain global_tload_mpirun_l3/ # one mpirun launches an L3 rank per machine; MPI descriptor exchange + vector_add_mpi_direct_l3/ # L4 joins MPI rank 0 and controls both real L3 ranks through direct P2P ``` Why no `tensormap_and_ringbuffer/` layer? Because every example here hard-codes @@ -67,6 +68,11 @@ parent owns a single `mpirun` that launches an L3 rank on each machine example's README for its extra prerequisites (`mpirun` + `mpi4py` on both machines). +`vector_add_mpi_direct_l3` is the direct-MPI alternative: one supervisor +launches L4 as MPI rank 0 and the two real L3 workers as ranks 1 and 2. It also +needs `mpirun` and `mpi4py` on both machines, but task/control traffic is P2P +between L4 and each L3 rather than passing through the PR2 group mailbox. + ### What a new L4 example needs ```text diff --git a/examples/workers/l4/vector_add_mpi_direct_l3/README.md b/examples/workers/l4/vector_add_mpi_direct_l3/README.md new file mode 100644 index 0000000000..6abdb6e235 --- /dev/null +++ b/examples/workers/l4/vector_add_mpi_direct_l3/README.md @@ -0,0 +1,68 @@ +# L4 direct MPI L3 vector add + +This example launches one static MPI world across two machines: rank 0 owns +L4, rank 1 owns a real L3 on the L4 machine, and rank 2 owns a real L3 on the +peer. L4 sends SLR3 task and control frames directly to both L3 ranks, and +each L3 dispatches to its local L2 workers. + +```text +machine A machine B + +MPI rank 0 / L4 -------- direct MPI --------> MPI rank 2 / real L3 + | + +------------- direct MPI -----------> MPI rank 1 / real L3 +``` + +The example validates direct MPI startup and teardown, remote memory copy, +L3-to-L2 vector-add execution, result checking, and Global CommDomain +lifecycle. The vector-add kernel does not perform a peer `TLOAD`. + +## Prerequisites beyond the sibling example + +- `mpirun`/`mpiexec` and `mpi4py` must be installed on **both** machines and + built against the **same** MPI implementation. +- `mpirun` executes one identical command line on both machines, so `--python` + must name an interpreter path valid on BOTH. Point it at a per-machine + launcher script installed at one shared absolute path; each machine's copy + sources CANN, enters that machine's checkout, and execs its `.venv` Python: + +```bash +cat > /tmp/simpler-mpi-python <<'EOF' +#!/usr/bin/env bash +source /usr/local/Ascend/cann/set_env.sh +cd /path/to/this/machines/simpler +exec /path/to/this/machines/simpler/.venv/bin/python "$@" +EOF +chmod +x /tmp/simpler-mpi-python +``` + +## Run on the L4 parent + +`192.0.2.10` / `192.0.2.20` are documentation placeholders. Hosts must be +numeric IPs; the local host is where ranks 0 and 1 run: + +```bash +source .venv/bin/activate +python -m examples.workers.l4.vector_add_mpi_direct_l3.main \ + --local-host 192.0.2.10 --remote-host 192.0.2.20 \ + --python /tmp/simpler-mpi-python \ + --local-devices 0,1 --remote-devices 0,1 \ + --mpirun-path "$(command -v mpiexec)" \ + --launcher-family mpich +``` + +Success requires the process to return status 0 after printing: + +```text +vector_add_mpi_direct_l3 passed +``` + +Any failed vector validation exits non-zero. + +## Running it in CI + +The network1 job's `network1-stage` action writes the per-machine launcher on both +machines at one shared path and exports it as `NETWORK1_MPI_PYTHON`; the +`test_vector_add_mpi_direct_l3.py` wrapper reads it together with `NETWORK1_LOCAL_IP` +and the standard network1 fixtures, and skips when `mpirun`, `mpi4py`, or either +variable is absent. diff --git a/examples/workers/l4/vector_add_mpi_direct_l3/__init__.py b/examples/workers/l4/vector_add_mpi_direct_l3/__init__.py new file mode 100644 index 0000000000..47312d2f6d --- /dev/null +++ b/examples/workers/l4/vector_add_mpi_direct_l3/__init__.py @@ -0,0 +1,9 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""L4-to-L3 direct MPI vector-add example.""" diff --git a/examples/workers/l4/vector_add_mpi_direct_l3/main.py b/examples/workers/l4/vector_add_mpi_direct_l3/main.py new file mode 100755 index 0000000000..bbfff1c4a2 --- /dev/null +++ b/examples/workers/l4/vector_add_mpi_direct_l3/main.py @@ -0,0 +1,287 @@ +#!/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 one L4 MPI rank with local-host and remote-host real L3 MPI ranks.""" + +from __future__ import annotations + +import argparse +import contextlib +import ctypes +import json +import os +import tempfile +from pathlib import Path +from typing import Any + +from simpler.mpi_direct_supervisor import run_supervisor +from simpler.remote_l3_session import get_inner_handle +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CoreCallable, + DataType, + RemoteTensorRef, + TaskArgs, + TensorArgType, +) +from simpler.worker import RemoteCallable + +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 + +CONTROLLER_TARGET = "examples.workers.l4.vector_add_mpi_direct_l3.main:controller" +REMOTE_ORCH_TARGET = "examples.workers.l4.vector_add_mpi_direct_l3.main:remote_l3_vector_orch" +ELEMENTS = 128 * 128 +NBYTES = ELEMENTS * ctypes.sizeof(ctypes.c_float) +FloatArray = ctypes.c_float * ELEMENTS +_REMOTE_KEEPALIVE: list[TaskArgs] = [] + + +def _parse_devices(value: str, *, label: str) -> tuple[int, ...]: + device_ids = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + if not device_ids or any(device_id < 0 for device_id in device_ids): + raise ValueError(f"{label} device list must contain non-negative ids") + if len(set(device_ids)) != len(device_ids): + raise ValueError(f"{label} device list must not contain duplicates") + return device_ids + + +def _digest_from_args(args: TaskArgs) -> bytes: + return b"".join(int(args.scalar(index)).to_bytes(8, "little", signed=False) for index in range(4)) + + +def remote_l3_vector_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + chip_handle = get_inner_handle(_digest_from_args(args).hex()) + chip_args = TaskArgs() + chip_args.add_tensor(args.tensor(0), TensorArgType.INPUT) + chip_args.add_tensor(args.tensor(1), TensorArgType.INPUT) + chip_args.add_tensor(args.tensor(2), TensorArgType.OUTPUT_EXISTING) + _REMOTE_KEEPALIVE[:] = [chip_args] + orch.submit_next_level(chip_handle, chip_args, cfg, worker=0) + + +def _build_chip_callable(platform: str, runtime: str) -> ChipCallable: + root = Path(__file__).resolve().parents[4] + kernels = root / "examples" / "workers" / "l3" / "multi_chip_dispatch" / "kernels" + compiler = KernelCompiler(platform=platform) + kernel = compiler.compile_incore( + source_path=str(kernels / "aiv" / "vector_add_kernel.cpp"), + core_type="aiv", + pto_isa_root=ensure_pto_isa_root(), + extra_include_dirs=compiler.get_orchestration_include_dirs(runtime), + ) + orchestration = compiler.compile_orchestration( + runtime_name=runtime, + source_path=str(kernels / "orchestration" / "vector_add_orch.cpp"), + ) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + func_name="vector_add_orchestration", + binary=orchestration, + children=[ + ( + 0, + CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + binary=extract_text_section(kernel), + ), + ) + ], + ) + + +def _array(value: float) -> Any: + return FloatArray(*([value] * ELEMENTS)) + + +def _task_args(handles, digest: bytes) -> TaskArgs: + args = TaskArgs() + for index, handle in enumerate(handles): + direction = TensorArgType.OUTPUT_EXISTING if index == 2 else TensorArgType.INPUT + args.add_tensor(RemoteTensorRef(handle, shape=(ELEMENTS,), dtype=DataType.FLOAT32), direction) + for offset in range(0, len(digest), 8): + args.add_scalar(int.from_bytes(digest[offset : offset + 8], "little", signed=False)) + return args + + +def controller(context) -> None: + """Rank-0 controller invoked by ``simpler.mpi_direct_runtime``.""" + executors = context.topology.executors + if len(executors) != 2: + raise ValueError("vector_add_mpi_direct_l3 requires exactly two L3 executor ranks") + worker_ids = tuple(spec.worker_id for spec in executors) + worker = context.create_worker( + num_sub_workers=0, + startup_timeout_s=context.topology.startup_timeout_s, + remote_session_timeout_s=context.topology.session_timeout_s, + ) + allocations = [] + keepalive: list[TaskArgs] = [] + try: + first = executors[0] + chip_handle = worker.register(_build_chip_callable(first.platform, first.runtime)) + remote_handle = worker.register(RemoteCallable(REMOTE_ORCH_TARGET), workers=list(worker_ids)) + worker.init() + + digest = bytes(chip_handle.digest) + expected: dict[int, tuple[Any, float]] = {} + task_inputs: dict[int, list] = {} + for index, worker_id in enumerate(worker_ids): + lhs = float(2 + index * 4) + rhs = float(3 + index * 4) + handles = [worker.remote_malloc(worker=worker_id, nbytes=NBYTES) for _ in range(3)] + allocations.extend(handles) + worker.remote_copy_to(handles[0], _array(lhs), NBYTES) + worker.remote_copy_to(handles[1], _array(rhs), NBYTES) + worker.remote_copy_to(handles[2], _array(0.0), NBYTES) + task_inputs[worker_id] = handles + expected[worker_id] = (_array(0.0), lhs + rhs) + + members = tuple( + (spec.worker_id, local_index) for spec in executors for local_index in range(len(spec.device_ids)) + ) + + def parent_orch(orch, _args, cfg): + orch.allocate_global_domain( + name="vector-add-mpi-direct-l3", + members=members, + window_size=1024 * 1024, + ) + local_args = [_task_args(task_inputs[worker_id], digest) for worker_id in worker_ids] + keepalive[:] = local_args + for worker_id, task_args in zip(worker_ids, local_args): + orch.submit_next_level(remote_handle, task_args, cfg, worker=worker_id) + + config = CallConfig() + config.aicpu_thread_num = 2 + worker.run(parent_orch, config=config) + + for worker_id, handles in task_inputs.items(): + output, wanted = expected[worker_id] + worker.remote_copy_from(handles[2], output, NBYTES) + max_diff = max(abs(float(output[index]) - wanted) for index in range(ELEMENTS)) + if max_diff > 1e-5: + raise AssertionError(f"worker {worker_id} vector result mismatch: max_diff={max_diff}") + print("vector_add_mpi_direct_l3 passed") + finally: + keepalive.clear() + for handle in reversed(allocations): + with contextlib.suppress(Exception): + worker.remote_free(handle) + + +def run( # noqa: PLR0913 -- mirrors the two-host CLI surface used by the other L4 examples + *, + local_host: str, + remote_host: str, + python_executable: str, + local_devices: str, + remote_devices: str, + platform: str = "a2a3", + runtime: str = "tensormap_and_ringbuffer", + comm_profile: str = "a3-fabric-v1", + startup_timeout: float = 180.0, + session_timeout: float = 120.0, + mpirun_path: str = "mpirun", + launcher_family: str = "auto", +) -> int: + local_ids = _parse_devices(local_devices, label="local") + remote_ids = _parse_devices(remote_devices, label="remote") + local_global_ranks = tuple(range(len(local_ids))) + remote_global_ranks = tuple(range(len(local_ids), len(local_ids) + len(remote_ids))) + topology = { + "controller_rank": 0, + "controller_host": local_host, + "startup_timeout_s": float(startup_timeout), + "session_timeout_s": float(session_timeout), + "heartbeat_interval_s": 1.0, + "max_pending_frame_bytes": 64 * 1024 * 1024, + "launcher_args": ["-wdir", "/tmp"], + "executor_ranks": [ + { + "rank": 1, + "worker_id": 0, + "host": local_host, + "platform": platform, + "runtime": runtime, + "device_ids": list(local_ids), + "global_device_ranks": list(local_global_ranks), + "num_sub_workers": 0, + "comm_profile": comm_profile, + }, + { + "rank": 2, + "worker_id": 1, + "host": remote_host, + "platform": platform, + "runtime": runtime, + "device_ids": list(remote_ids), + "global_device_ranks": list(remote_global_ranks), + "num_sub_workers": 0, + "comm_profile": comm_profile, + }, + ], + } + fd, topology_path = tempfile.mkstemp(prefix="simpler-mpi-direct-", suffix=".json", text=True) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump(topology, stream, sort_keys=True) + stream.write("\n") + return run_supervisor( + topology_path, + CONTROLLER_TARGET, + mpirun_path=mpirun_path, + launcher_family=launcher_family, + python_executable=python_executable, + ) + finally: + with contextlib.suppress(FileNotFoundError): + os.unlink(topology_path) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--local-host", required=True, help="L4/rank-0 machine address") + parser.add_argument("--remote-host", required=True, help="peer L3 machine address") + parser.add_argument("--python", dest="python_executable", required=True) + parser.add_argument("--local-devices", default="0,1") + parser.add_argument("--remote-devices", default="0,1") + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + parser.add_argument("--comm-profile", default="a3-fabric-v1") + parser.add_argument("--startup-timeout", type=float, default=180.0) + parser.add_argument("--session-timeout", type=float, default=120.0) + parser.add_argument("--mpirun-path", default="mpirun") + parser.add_argument("--launcher-family", choices=("auto", "openmpi", "mpich"), default="auto") + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + return run( + local_host=args.local_host, + remote_host=args.remote_host, + python_executable=args.python_executable, + local_devices=args.local_devices, + remote_devices=args.remote_devices, + platform=args.platform, + runtime=args.runtime, + comm_profile=args.comm_profile, + startup_timeout=args.startup_timeout, + session_timeout=args.session_timeout, + mpirun_path=args.mpirun_path, + launcher_family=args.launcher_family, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/workers/l4/vector_add_mpi_direct_l3/run_parent.sh b/examples/workers/l4/vector_add_mpi_direct_l3/run_parent.sh new file mode 100755 index 0000000000..1d52d9a166 --- /dev/null +++ b/examples/workers/l4/vector_add_mpi_direct_l3/run_parent.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# 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. +# ----------------------------------------------------------------------------------------------------------- +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" + +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_LOCAL_HOST:?set the L4 machine address}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_REMOTE_HOST:?set the peer L3 machine address}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_MPI_PYTHON:?set an absolute launcher path valid on both machines}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_LOCAL_DEVICES:=0,1}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_REMOTE_DEVICES:=0,1}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_PLATFORM:=a2a3}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_RUNTIME:=tensormap_and_ringbuffer}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_STARTUP_TIMEOUT:=180}" +: "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_SESSION_TIMEOUT:=120}" + +cd "${ROOT_DIR}" +if [[ ! -f .venv/bin/activate ]]; then + echo "error: ${ROOT_DIR}/.venv/bin/activate not found; create the virtual environment first" >&2 + exit 1 +fi +# shellcheck source=/dev/null +source .venv/bin/activate + +exec python -m examples.workers.l4.vector_add_mpi_direct_l3.main \ + --local-host "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_LOCAL_HOST}" \ + --remote-host "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_REMOTE_HOST}" \ + --python "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_MPI_PYTHON}" \ + --local-devices "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_LOCAL_DEVICES}" \ + --remote-devices "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_REMOTE_DEVICES}" \ + --platform "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_PLATFORM}" \ + --runtime "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_RUNTIME}" \ + --startup-timeout "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_STARTUP_TIMEOUT}" \ + --session-timeout "${SIMPLER_VECTOR_ADD_MPI_DIRECT_L3_SESSION_TIMEOUT}" \ + --mpirun-path "${MPIRUN:-mpirun}" \ + --launcher-family "${MPI_LAUNCHER_FAMILY:-auto}" diff --git a/examples/workers/l4/vector_add_mpi_direct_l3/test_vector_add_mpi_direct_l3.py b/examples/workers/l4/vector_add_mpi_direct_l3/test_vector_add_mpi_direct_l3.py new file mode 100644 index 0000000000..496f3d4999 --- /dev/null +++ b/examples/workers/l4/vector_add_mpi_direct_l3/test_vector_add_mpi_direct_l3.py @@ -0,0 +1,65 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Network1 ST for the complete L4-to-L3 direct MPI path.""" + +import importlib.util +import os +import shutil + +import pytest + +from simpler_setup import SceneTestLevel, scene_level + +from .main import run + + +def _device_spec(device_ids) -> str: + return ",".join(str(device_id) for device_id in device_ids) + + +def _require_mpi_direct_network1_env() -> tuple[str, str, str]: + mpi_launcher = shutil.which("mpirun") + if mpi_launcher is None: + mpi_launcher = shutil.which("mpiexec") + if mpi_launcher is None: + pytest.skip("mpirun or mpiexec is not on PATH") + assert mpi_launcher is not None + if importlib.util.find_spec("mpi4py") is None: + pytest.skip("mpi4py is not installed") + local_ip = os.environ.get("NETWORK1_LOCAL_IP", "") + if not local_ip: + pytest.skip("NETWORK1_LOCAL_IP is required for the L4/rank-0 startup gate") + mpi_python = os.environ.get("NETWORK1_MPI_PYTHON", "") + if not mpi_python: + pytest.skip("NETWORK1_MPI_PYTHON is required on both MPI hosts") + return mpi_launcher, local_ip, mpi_python + + +@scene_level(SceneTestLevel.NETWORK1) +@pytest.mark.platforms(["a2a3"]) +@pytest.mark.runtime("tensormap_and_ringbuffer") +@pytest.mark.device_count(2) +@pytest.mark.network1_remote_device_count(2) +def test_vector_add_mpi_direct_l3( + st_platform, st_device_ids, st_network1_peer, st_network1_remote_device_ids, st_network1_logs +): + mpirun, local_ip, mpi_python = _require_mpi_direct_network1_env() + remote_host, _daemon_port = st_network1_peer.endpoint.rsplit(":", 1) + rc = run( + local_host=local_ip, + remote_host=remote_host, + python_executable=mpi_python, + local_devices=_device_spec(st_device_ids), + remote_devices=_device_spec(st_network1_remote_device_ids), + platform=st_platform, + startup_timeout=st_network1_peer.session_timeout_s, + session_timeout=st_network1_peer.session_timeout_s, + mpirun_path=mpirun, + ) + assert rc == 0 diff --git a/pyproject.toml b/pyproject.toml index 094eadf000..7c99b07e9b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = ["cloudpickle>=2.2"] [project.scripts] simpler-remote-worker = "simpler.remote_l3_worker:main" simpler-remote-l3-session = "simpler.remote_l3_session:main" +simpler-mpi-direct = "simpler.mpi_direct_supervisor:main" [project.optional-dependencies] # ``torch>=2.3`` is required by ``simpler_setup.torch_interop`` (uses @@ -27,6 +28,7 @@ simpler-remote-l3-session = "simpler.remote_l3_session:main" # torch via a custom index URL and does not rely on this pin; it is here # so ``pip install -e '.[test]'`` resolves a usable version for local devs. test = ["pytest>=6.0", "pytest-timeout>=2.3", "pytest-xdist>=3.0", "PyYAML>=6.0", "torch>=2.3"] +mpi = ["mpi4py>=3.1"] [tool.ruff] line-length = 120 diff --git a/python/bindings/CMakeLists.txt b/python/bindings/CMakeLists.txt index dfb9b249a3..f8a3a77cd3 100644 --- a/python/bindings/CMakeLists.txt +++ b/python/bindings/CMakeLists.txt @@ -24,6 +24,7 @@ set(HIERARCHICAL_SOURCES ${HIERARCHICAL_SRC}/scope.cpp ${HIERARCHICAL_SRC}/remote_wire.cpp ${HIERARCHICAL_SRC}/remote_endpoint.cpp + ${HIERARCHICAL_SRC}/mpi_direct_transport.cpp ${HIERARCHICAL_SRC}/orchestrator.cpp ${HIERARCHICAL_SRC}/worker_manager.cpp ${HIERARCHICAL_SRC}/scheduler.cpp diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 6821055a5a..ab1f042872 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -25,16 +25,19 @@ #pragma once #include +#include #include #include #include #include #include +#include #include #include #include "ring.h" +#include "mpi_direct_transport.h" #include "mpi_group_mailbox.h" #include "orchestrator.h" #include "types.h" @@ -235,6 +238,61 @@ inline void mailbox_store_i32(uint64_t addr, int32_t v) { } inline void bind_worker(nb::module_ &m) { + nb::enum_(m, "_MpiDirectTag") + .value("COMMAND_REQUEST", MpiDirectTag::COMMAND_REQUEST) + .value("COMMAND_REPLY", MpiDirectTag::COMMAND_REPLY) + .value("HEALTH", MpiDirectTag::HEALTH) + .value("LIFECYCLE", MpiDirectTag::LIFECYCLE); + + nb::class_(m, "_MpiDirectTransportHub") + .def(nb::init(), nb::arg("max_pending_frame_bytes")) + .def( + "register_route", &MpiDirectTransportHub::register_route, nb::arg("worker_id"), nb::arg("mpi_rank"), + nb::arg("session_id"), nb::arg("comm_profile") + ) + .def( + "poll_outbound", + [](MpiDirectTransportHub &self, double timeout_s) -> nb::object { + std::optional result; + { + nb::gil_scoped_release release; + result = self.poll_outbound(timeout_s); + } + if (!result.has_value()) return nb::none(); + const auto &frame = result->frame; + return nb::make_tuple( + result->ticket, result->target_rank, static_cast(result->tag), + nb::bytes(reinterpret_cast(frame.data()), frame.size()) + ); + }, + nb::arg("timeout_s") = 0.0 + ) + .def( + "complete_outbound", &MpiDirectTransportHub::complete_outbound, nb::arg("ticket"), + nb::call_guard() + ) + .def( + "deliver", + [](MpiDirectTransportHub &self, int32_t source_rank, int32_t raw_tag, nb::bytes frame) { + if (raw_tag < static_cast(MpiDirectTag::COMMAND_REQUEST) || + raw_tag > static_cast(MpiDirectTag::LIFECYCLE)) { + throw std::invalid_argument("MPI direct tag is outside the fixed transport lanes"); + } + const auto *begin = reinterpret_cast(frame.c_str()); + std::vector native_frame(begin, begin + frame.size()); + { + nb::gil_scoped_release release; + self.deliver(source_rank, static_cast(raw_tag), native_frame); + } + }, + nb::arg("source_rank"), nb::arg("tag"), nb::arg("frame") + ) + .def("fail", &MpiDirectTransportHub::fail, nb::arg("message")) + .def("close", &MpiDirectTransportHub::close) + .def_prop_ro("pending_frame_bytes", &MpiDirectTransportHub::pending_frame_bytes) + .def_prop_ro("terminal", &MpiDirectTransportHub::terminal) + .def_prop_ro("terminal_error", &MpiDirectTransportHub::terminal_error); + // --- WorkerType --- nb::enum_(m, "WorkerType").value("NEXT_LEVEL", WorkerType::NEXT_LEVEL).value("SUB", WorkerType::SUB); @@ -490,13 +548,24 @@ inline void bind_worker(nb::module_ &m) { nb::arg("mpirun_pid"), nb::arg("runtime_timeout_s") = 30.0, "Register one shared-memory MPI group endpoint for each worker id." ) + .def( + "add_remote_l3_mpi", + [](Worker &self, int32_t worker_id, uint64_t session_id, const std::string &transport_name, + const std::shared_ptr &hub, double attach_timeout_s, double runtime_timeout_s) { + nb::gil_scoped_release release; + self.add_remote_l3_mpi(worker_id, session_id, transport_name, hub, attach_timeout_s, runtime_timeout_s); + }, + nb::arg("worker_id"), nb::arg("session_id"), nb::arg("transport_name"), nb::arg("hub"), + nb::arg("attach_timeout_s") = 30.0, nb::arg("runtime_timeout_s") = 30.0, + "Register a directed MPI-backed REMOTE_L3 endpoint after HELLO READY." + ) // Release the GIL while starting the Scheduler thread so another Python // thread can run during it — e.g. a concurrent close() observing // INITIALIZING and failing fast. init/close remain same-thread-only // (enforced by Worker.close()). .def("init", &Worker::init, nb::call_guard(), "Start the Scheduler thread.") - .def("close", &Worker::close, "Stop the Scheduler thread.") + .def("close", &Worker::close, nb::call_guard(), "Stop the Scheduler thread.") .def( "get_orchestrator", &Worker::get_orchestrator, nb::rv_policy::reference_internal, diff --git a/python/simpler/mpi_direct_protocol.py b/python/simpler/mpi_direct_protocol.py new file mode 100644 index 0000000000..cda9ed0435 --- /dev/null +++ b/python/simpler/mpi_direct_protocol.py @@ -0,0 +1,21 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Fixed MPI tag lanes for the direct-MPI transport.""" + +import enum + +MPI_DIRECT_STARTUP_TOKEN_ENV = "SIMPLER_MPI_DIRECT_STARTUP_TOKEN" +MPI_DIRECT_GATE_MAX_BYTES = 64 * 1024 + + +class MpiDirectTag(enum.IntEnum): + COMMAND_REQUEST = 1 + COMMAND_REPLY = 2 + HEALTH = 3 + LIFECYCLE = 4 diff --git a/python/simpler/mpi_direct_runtime.py b/python/simpler/mpi_direct_runtime.py new file mode 100644 index 0000000000..eb9fa84442 --- /dev/null +++ b/python/simpler/mpi_direct_runtime.py @@ -0,0 +1,586 @@ +# 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-role runtime for direct MPI control between L4 and real L3 workers.""" + +from __future__ import annotations + +import argparse +import base64 +import contextlib +import importlib +import json +import os +import runpy +import socket +import struct +import sys +import threading +import time +import traceback +from pathlib import Path +from typing import Any, Callable + +from .mpi_direct_protocol import MPI_DIRECT_GATE_MAX_BYTES, MPI_DIRECT_STARTUP_TOKEN_ENV, MpiDirectTag +from .mpi_direct_topology import MpiDirectExecutorSpec, MpiDirectTopology, load_runtime_manifest_data +from .remote_l3_limits import FRAME_HEADER_BYTES, MAX_FRAME_BYTES + +COMMAND_REQUEST_TAG = int(MpiDirectTag.COMMAND_REQUEST) +COMMAND_REPLY_TAG = int(MpiDirectTag.COMMAND_REPLY) +HEALTH_TAG = int(MpiDirectTag.HEALTH) +LIFECYCLE_TAG = int(MpiDirectTag.LIFECYCLE) +_MPI_POLL_INTERVAL_S = 0.001 +_MPI_GATE_RETRY_INTERVAL_S = 0.01 + + +def _launcher_rank() -> int: + names = ( + "OMPI_COMM_WORLD_RANK", + "PMI_RANK", + "PMIX_RANK", + "MV2_COMM_WORLD_RANK", + "SLURM_PROCID", + ) + values = {name: int(os.environ[name]) for name in names if name in os.environ} + if not values: + raise RuntimeError(f"MPI launcher did not provide a pre-init rank in any of {names}") + if len(set(values.values())) != 1: + raise RuntimeError(f"MPI launcher rank variables disagree: {values}") + return next(iter(values.values())) + + +def _open_fds() -> tuple[int, ...]: + fd_dir = Path("/proc/self/fd") + if not fd_dir.is_dir(): + fd_dir = Path("/dev/fd") + if not fd_dir.is_dir(): + raise RuntimeError("direct MPI executor requires /proc/self/fd or /dev/fd for launcher FD isolation") + return tuple(sorted(int(path.name) for path in fd_dir.iterdir() if path.name.isdigit() and int(path.name) >= 3)) + + +def _import_mpi(): + try: + import mpi4py # noqa: PLC0415 + + mpi4py.rc.initialize = False + mpi4py.rc.finalize = False + mpi4py.rc.thread_level = "serialized" + from mpi4py import MPI # noqa: PLC0415 + except ImportError as exc: + raise RuntimeError("direct MPI runtime requires the optional mpi4py dependency") from exc + return MPI + + +def _init_mpi(MPI, expected_rank: int, expected_world_size: int): + if MPI.Is_initialized(): + raise RuntimeError("MPI was initialized before simpler.mpi_direct_runtime called MPI.Init_thread") + provided = MPI.Init_thread(required=MPI.THREAD_SERIALIZED) + if provided < MPI.THREAD_SERIALIZED: + raise RuntimeError(f"MPI_THREAD_SERIALIZED is required, but MPI provided thread level {provided}") + world = MPI.COMM_WORLD + world.Set_errhandler(MPI.ERRORS_RETURN) + if int(world.Get_rank()) != expected_rank: + raise RuntimeError("pre-init launcher rank does not match MPI_COMM_WORLD rank") + if int(world.Get_size()) != expected_world_size: + raise RuntimeError("MPI_COMM_WORLD size does not match the static topology") + return world + + +def _gate_send(sock: socket.socket, message: dict[str, object]) -> None: + payload = json.dumps(message, separators=(",", ":")).encode("utf-8") + if len(payload) > MPI_DIRECT_GATE_MAX_BYTES: + raise ValueError("MPI startup gate message is too large") + sock.sendall(struct.pack("!I", len(payload)) + payload) + + +def _gate_recv(sock: socket.socket) -> dict[str, object]: + header = b"" + while len(header) < 4: + chunk = sock.recv(4 - len(header)) + if not chunk: + raise ConnectionError("MPI startup gate closed before response") + header += chunk + length = struct.unpack("!I", header)[0] + if length <= 0 or length > MPI_DIRECT_GATE_MAX_BYTES: + raise ValueError("invalid MPI startup gate message length") + payload = bytearray() + while len(payload) < length: + chunk = sock.recv(length - len(payload)) + if not chunk: + raise ConnectionError("MPI startup gate closed while reading response") + payload.extend(chunk) + message = json.loads(payload) + if not isinstance(message, dict): + raise ValueError("MPI startup gate response must be an object") + return message + + +def _pre_mpi_gate( + host: str, + port: int, + token: str, + rank: int, + timeout_s: float, + error: BaseException | None = None, +) -> None: + deadline = time.monotonic() + float(timeout_s) + last_error: BaseException | None = None + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("MPI startup gate connection timed out") from last_error + try: + with socket.create_connection((host, int(port)), timeout=min(remaining, 1.0)) as sock: + sock.settimeout(max(1.0, remaining)) + _gate_send( + sock, + { + "token": token, + "rank": int(rank), + "state": "failed" if error is not None else "ready", + "error": "" if error is None else f"{type(error).__name__}: {error}", + }, + ) + if error is not None: + return + response = _gate_recv(sock) + if response.get("token") != token: + raise RuntimeError("MPI startup gate token mismatch") + state = response.get("state") + if state != "go_mpi": + raise RuntimeError(str(response.get("error") or "MPI startup gate rejected rank")) + return + except (OSError, TimeoutError, ConnectionError) as exc: + last_error = exc + retry_remaining = deadline - time.monotonic() + if retry_remaining > 0: + time.sleep(min(_MPI_GATE_RETRY_INTERVAL_S, retry_remaining)) + continue + + +class _ControllerProgress: + def __init__(self, MPI, world, hub) -> None: + self._MPI = MPI + self._world = world + self._hub = hub + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name="simpler-mpi-direct-progress", daemon=True) + self._error: BaseException | None = None + + def start(self) -> None: + self._thread.start() + + def stop(self, timeout_s: float) -> None: + self._stop.set() + self._thread.join(timeout_s) + if self._thread.is_alive(): + self._hub.fail("MPI progress thread did not drain before shutdown deadline") + raise TimeoutError("MPI progress thread did not stop") + if self._error is not None: + raise RuntimeError("MPI progress failed") from self._error + + def _receive_one(self) -> bool: + status = self._MPI.Status() + message = self._world.improbe( + source=self._MPI.ANY_SOURCE, + tag=self._MPI.ANY_TAG, + status=status, + ) + if message is None: + return False + count = int(status.Get_count(self._MPI.BYTE)) + if count < FRAME_HEADER_BYTES or count > MAX_FRAME_BYTES: + raise RuntimeError(f"inbound MPI frame length {count} is outside the SLR3 bounds") + frame = bytearray(count) + message.Recv([frame, self._MPI.BYTE]) + self._hub.deliver(int(status.Get_source()), int(status.Get_tag()), bytes(frame)) + return True + + def _run(self) -> None: + in_flight: list[tuple[Any, int, bytes]] = [] + try: + while True: + did_work = False + remaining: list[tuple[Any, int, bytes]] = [] + for request, ticket, frame in in_flight: + if request.Test(): + self._hub.complete_outbound(ticket) + did_work = True + else: + remaining.append((request, ticket, frame)) + in_flight = remaining + + outbound = self._hub.poll_outbound(0.0) + if outbound is not None: + ticket, target_rank, tag, frame = outbound + frame = bytes(frame) + request = self._world.Isend( + [frame, self._MPI.BYTE], + dest=int(target_rank), + tag=int(tag), + ) + in_flight.append((request, int(ticket), frame)) + did_work = True + + did_work = self._receive_one() or did_work + if self._stop.is_set() and not in_flight and self._hub.pending_frame_bytes == 0: + return + if not did_work: + time.sleep(_MPI_POLL_INTERVAL_S) + except BaseException as exc: # noqa: BLE001 + self._error = exc + with contextlib.suppress(BaseException): + self._hub.fail(f"MPI progress failure: {type(exc).__name__}: {exc}") + + +class MpiDirectControllerContext: + def __init__(self, MPI, world, topology: MpiDirectTopology, session_id: int) -> None: + from _task_interface import _MpiDirectTransportHub # noqa: PLC0415 + + self.topology = topology + self.session_id = int(session_id) + self.cluster_id = f"mpi-direct-{self.session_id:016x}" + self._hub = _MpiDirectTransportHub(topology.max_pending_frame_bytes) + for spec in topology.executors: + self._hub.register_route(spec.worker_id, spec.rank, self.session_id, spec.comm_profile) + self._progress = _ControllerProgress(MPI, world, self._hub) + self._workers: list[Any] = [] + self._closed = False + self._progress.start() + + def attach(self, worker) -> tuple[int, ...]: + from .worker import _MpiDirectWorkerSpec # noqa: PLC0415 + + if self._closed: + raise RuntimeError("MPI direct controller context is closed") + if worker.level < 4: + raise TypeError("MPI direct controller requires Worker(level>=4)") + if self._workers: + raise RuntimeError("MPI direct controller context supports exactly one attached L4 Worker") + worker._global_cluster_id = self.cluster_id # noqa: SLF001 -- context owns the direct-MPI topology + worker_ids = [] + for executor in self.topology.executors: + worker_id = worker._add_mpi_direct_worker( + _MpiDirectWorkerSpec( + worker_id=executor.worker_id, + mpi_rank=executor.rank, + session_id=self.session_id, + host=executor.host, + comm_profile=executor.comm_profile, + platform=executor.platform, + runtime=executor.runtime, + device_ids=executor.device_ids, + global_device_ranks=executor.global_device_ranks, + hub=self._hub, + attach_timeout_s=self.topology.startup_timeout_s, + runtime_timeout_s=self.topology.session_timeout_s, + ) + ) + worker_ids.append(worker_id) + self._workers.append(worker) + return tuple(worker_ids) + + def create_worker(self, **config): + from .worker import Worker # noqa: PLC0415 + + worker = Worker(level=4, **config) + self.attach(worker) + return worker + + def close(self) -> None: + if self._closed: + return + self._closed = True + close_errors: list[BaseException] = [] + for worker in reversed(self._workers): + try: + worker.close() + except BaseException as exc: # noqa: BLE001 + close_errors.append(exc) + try: + self._progress.stop(self.topology.session_timeout_s) + finally: + self._hub.close() + if close_errors: + raise RuntimeError(f"failed to close {len(close_errors)} attached L4 Worker(s)") from close_errors[0] + + +class _ExecutorFrameSocket: + def __init__(self, MPI, world, spec: MpiDirectExecutorSpec, session_id: int, heartbeat_interval_s: float) -> None: + self._MPI = MPI + self._world = world + self._spec = spec + self._session_id = int(session_id) + self._heartbeat_interval_s = float(heartbeat_interval_s) + self._mpi_mu = threading.Lock() + self._receive_buffer = bytearray() + self._health_stop = threading.Event() + self._health_thread: threading.Thread | None = None + self._health_error: BaseException | None = None + self._health_sequence = 0 + + def _send(self, data: bytes, tag: int) -> None: + with self._mpi_mu: + request = self._world.Isend([data, self._MPI.BYTE], dest=0, tag=tag) + request.Wait() + + def _start_health(self) -> None: + if self._health_thread is not None: + return + self._health_thread = threading.Thread(target=self._health_loop, name="simpler-mpi-direct-health", daemon=True) + self._health_thread.start() + + def _health_loop(self) -> None: + from .remote_l3_protocol import FrameHeader, FrameType, encode_frame # noqa: PLC0415 + + try: + while not self._health_stop.wait(self._heartbeat_interval_s): + self._health_sequence += 1 + frame = encode_frame( + FrameHeader( + FrameType.HEALTH, + self._session_id, + self._spec.worker_id, + self._health_sequence, + ), + b"", + ) + self._send(frame, HEALTH_TAG) + except BaseException as exc: # noqa: BLE001 + self._health_error = exc + + def sendall(self, data: bytes) -> None: + from .remote_l3_protocol import FrameType, decode_frame # noqa: PLC0415 + + raw = bytes(data) + frame = decode_frame(raw) + if frame.header.session_id != self._session_id or frame.header.worker_id != self._spec.worker_id: + raise RuntimeError("executor outbound SLR3 frame identity mismatch") + if frame.header.frame_type in (FrameType.COMPLETION, FrameType.CONTROL_REPLY): + tag = COMMAND_REPLY_TAG + elif frame.header.frame_type == FrameType.HELLO: + tag = LIFECYCLE_TAG + elif frame.header.frame_type == FrameType.HEALTH: + tag = HEALTH_TAG + else: + raise RuntimeError(f"executor cannot send SLR3 frame type {frame.header.frame_type.name}") + self._send(raw, tag) + if frame.header.frame_type == FrameType.HELLO: + self._start_health() + + def _receive_message(self) -> bytes: + from .remote_l3_protocol import FrameType, decode_frame # noqa: PLC0415 + + while True: + if self._health_error is not None: + raise RuntimeError("executor heartbeat failed") from self._health_error + status = self._MPI.Status() + with self._mpi_mu: + message = self._world.improbe(source=0, tag=self._MPI.ANY_TAG, status=status) + if message is not None: + count = int(status.Get_count(self._MPI.BYTE)) + if count < FRAME_HEADER_BYTES or count > MAX_FRAME_BYTES: + raise RuntimeError(f"controller MPI frame length {count} is outside the SLR3 bounds") + frame = bytearray(count) + message.Recv([frame, self._MPI.BYTE]) + if message is None: + time.sleep(_MPI_POLL_INTERVAL_S) + continue + tag = int(status.Get_tag()) + decoded = decode_frame(bytes(frame)) + expected_tag = LIFECYCLE_TAG if decoded.header.frame_type == FrameType.SHUTDOWN else COMMAND_REQUEST_TAG + if tag != expected_tag: + raise RuntimeError("controller MPI tag does not match SLR3 request type") + if decoded.header.session_id != self._session_id or decoded.header.worker_id != self._spec.worker_id: + raise RuntimeError("controller SLR3 frame identity mismatch") + return bytes(frame) + + def recv(self, size: int) -> bytes: + if size < 0: + raise ValueError("recv size must be non-negative") + if not self._receive_buffer: + self._receive_buffer.extend(self._receive_message()) + result = bytes(self._receive_buffer[:size]) + del self._receive_buffer[:size] + return result + + def close(self) -> None: + self._health_stop.set() + if self._health_thread is not None: + self._health_thread.join(self._heartbeat_interval_s + 1.0) + if self._health_thread.is_alive(): + raise RuntimeError("executor heartbeat thread did not stop") + if self._health_error is not None: + raise RuntimeError("executor heartbeat failed") from self._health_error + + +def _load_controller(target: str) -> Callable[[MpiDirectControllerContext], Any]: + path = Path(target) + if path.is_file(): + namespace = runpy.run_path(str(path.resolve()), run_name="simpler_mpi_direct_controller") + callback = namespace.get("main") + else: + if ":" not in target: + raise ValueError("controller must be a Python file or module:callable") + module_name, qualname = target.split(":", 1) + callback = importlib.import_module(module_name) + for part in qualname.split("."): + callback = getattr(callback, part) + if not callable(callback): + raise TypeError("controller target must resolve to a callable main(context)") + return callback + + +def _run_controller(MPI, world, topology: MpiDirectTopology, session_id: int, target: str) -> None: + context = MpiDirectControllerContext(MPI, world, topology, session_id) + try: + _load_controller(target)(context) + finally: + context.close() + + +def _run_executor(MPI, world, topology: MpiDirectTopology, session_id: int, spec, worker) -> None: + from .remote_l3_session import _run_command_loop # noqa: PLC0415 + + channel = _ExecutorFrameSocket(MPI, world, spec, session_id, topology.heartbeat_interval_s) + manifest = { + "session_id": session_id, + "worker_id": spec.worker_id, + "transport": spec.comm_profile, + "comm_profile": spec.comm_profile, + "cluster_id": f"mpi-direct-{session_id:016x}", + "session_timeout_s": topology.session_timeout_s, + "heartbeat_interval_s": topology.heartbeat_interval_s, + "node_rank": spec.rank - 1, + "node_count": len(topology.executors), + "platform": spec.platform, + "runtime": spec.runtime, + "device_ids": list(spec.device_ids), + "num_sub_workers": spec.num_sub_workers, + "global_device_ranks": list(spec.global_device_ranks), + } + try: + _run_command_loop(channel, manifest, worker, {}, {}) # type: ignore[arg-type] + finally: + try: + channel.close() + finally: + worker.close() + + +def run_runtime( # noqa: PLR0912 -- startup gate, role dispatch, and MPI teardown are one ordered lifecycle + topology_path: str | None, + session_id: int | None, + controller: str, + *, + manifest_json: str | None = None, + startup_host: str | None = None, + startup_port: int | None = None, +) -> int: + if manifest_json is not None: + try: + data = json.loads(base64.urlsafe_b64decode(manifest_json.encode("ascii"))) + except (ValueError, UnicodeError, json.JSONDecodeError) as exc: + raise ValueError("invalid inline MPI direct runtime manifest") from exc + topology, manifest_session_id = load_runtime_manifest_data(data) + if session_id is not None and int(session_id) != manifest_session_id: + raise ValueError("session_id does not match inline runtime manifest") + session_id = manifest_session_id + elif topology_path is not None and session_id is not None: + topology = MpiDirectTopology.load(topology_path) + session_id = int(session_id) + else: + raise ValueError("runtime requires either topology+session-id or manifest-json") + session_id = int(session_id) + if session_id == 0: + raise ValueError("session_id must be non-zero") + rank = _launcher_rank() + if rank < 0 or rank >= topology.world_size: + raise ValueError("launcher rank is outside the topology world") + worker = None + startup_error: BaseException | None = None + try: + if rank != 0: + from .worker import Worker # noqa: PLC0415 + + spec = topology.executor_for_rank(rank) + launcher_fds = _open_fds() + worker = Worker( + level=3, + platform=spec.platform, + runtime=spec.runtime, + device_ids=spec.device_ids, + num_sub_workers=spec.num_sub_workers, + comm_profile=spec.comm_profile, + global_device_ranks=spec.global_device_ranks, + startup_timeout_s=topology.startup_timeout_s, + fork_child_close_fds=launcher_fds, + ) + worker.init() + except BaseException as exc: + startup_error = exc + if worker is not None: + with contextlib.suppress(BaseException): + worker.close() + if startup_host is not None or startup_port is not None: + startup_token = os.environ.get(MPI_DIRECT_STARTUP_TOKEN_ENV) + if not (startup_host and startup_port and startup_token): + raise ValueError( + f"startup gate requires host, port, and {MPI_DIRECT_STARTUP_TOKEN_ENV} environment variable" + ) + _pre_mpi_gate(startup_host, int(startup_port), startup_token, rank, topology.startup_timeout_s, startup_error) + if startup_error is not None: + raise startup_error + + MPI = _import_mpi() + world = None + try: + world = _init_mpi(MPI, rank, topology.world_size) + if rank == 0: + _run_controller(MPI, world, topology, session_id, controller) + else: + assert worker is not None + _run_executor(MPI, world, topology, session_id, topology.executor_for_rank(rank), worker) + return 0 + except BaseException: + print( + f"[mpi-direct rank={rank} host={socket.gethostname()}] fatal error after MPI initialization", + file=sys.stderr, + flush=True, + ) + traceback.print_exc(file=sys.stderr) + sys.stderr.flush() + with contextlib.suppress(BaseException): + (world if world is not None else MPI.COMM_WORLD).Abort(1) + raise + finally: + if MPI.Is_initialized() and not MPI.Is_finalized(): + MPI.Finalize() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--topology") + parser.add_argument("--manifest-json") + parser.add_argument("--session-id", type=int) + parser.add_argument("--controller", required=True) + parser.add_argument("--startup-host") + parser.add_argument("--startup-port", type=int) + ns = parser.parse_args(argv) + return run_runtime( + ns.topology, + ns.session_id, + ns.controller, + manifest_json=ns.manifest_json, + startup_host=ns.startup_host, + startup_port=ns.startup_port, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/simpler/mpi_direct_supervisor.py b/python/simpler/mpi_direct_supervisor.py new file mode 100644 index 0000000000..4ab15792e7 --- /dev/null +++ b/python/simpler/mpi_direct_supervisor.py @@ -0,0 +1,397 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""External process owner for a direct-MPI L4 job.""" + +from __future__ import annotations + +import argparse +import base64 +import contextlib +import json +import os +import secrets +import signal +import socket +import struct +import subprocess +import sys +import tempfile +import time +from typing import Any + +from .mpi_direct_protocol import MPI_DIRECT_GATE_MAX_BYTES, MPI_DIRECT_STARTUP_TOKEN_ENV +from .mpi_direct_topology import MpiDirectTopology + +_EXPORTED_ENV_VARS = ( + "PATH", + "LD_LIBRARY_PATH", + "LD_PRELOAD", + "PYTHONPATH", + "VIRTUAL_ENV", + "PYTHONNOUSERSITE", + "PYTHONUNBUFFERED", + "ASCEND_HOME_PATH", + "ASCEND_OPP_PATH", + MPI_DIRECT_STARTUP_TOKEN_ENV, +) + + +class _StartupGateRankFailure(RuntimeError): + pass + + +def _host_slots(topology: MpiDirectTopology) -> tuple[tuple[str, int], ...]: + ordered: list[list[Any]] = [] + for host in topology.hosts: + if ordered and ordered[-1][0] == host: + ordered[-1][1] += 1 + else: + ordered.append([host, 1]) + return tuple((str(host), int(slots)) for host, slots in ordered) + + +def _detect_launcher_family(mpirun_path: str) -> str: + output_parts = [] + for flag in ("--version", "-info"): + try: + result = subprocess.run( + [mpirun_path, flag], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=5.0, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + output_parts.append(f"{flag}: {exc}") + continue + output_parts.append(result.stdout) + lowered = result.stdout.lower() + if "open mpi" in lowered or "openrte" in lowered: + return "openmpi" + if "hydra" in lowered or "mpich" in lowered: + return "mpich" + detail = "\n".join(part.strip() for part in output_parts if part.strip()) + raise RuntimeError( + f"cannot identify MPI launcher {mpirun_path!r}; pass --launcher-family openmpi or mpich" + + (f"\n{detail}" if detail else "") + ) + + +def _mpi_vendor_family(vendor: str) -> str: + lowered = vendor.lower() + if "open mpi" in lowered or "openmpi" in lowered: + return "openmpi" + if "mpich" in lowered or "mvapich" in lowered or "intel(r) mpi" in lowered or "intel mpi" in lowered: + return "mpich" + raise RuntimeError(f"unsupported mpi4py MPI vendor {vendor!r}") + + +def _detect_mpi4py_family(python_executable: str | None = None) -> tuple[str, str]: + script = ( + "import mpi4py; " + "mpi4py.rc.initialize = False; " + "mpi4py.rc.finalize = False; " + "from mpi4py import MPI; " + "print(MPI.get_vendor()[0])" + ) + try: + result = subprocess.run( + [python_executable or sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=10.0, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError(f"cannot inspect mpi4py with {python_executable or sys.executable!r}: {exc}") from exc + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() + raise RuntimeError( + f"cannot import mpi4py with {python_executable or sys.executable!r}" + (f"\n{detail}" if detail else "") + ) + vendor = result.stdout.strip().splitlines()[-1] if result.stdout.strip() else "" + if not vendor: + raise RuntimeError(f"mpi4py did not report its MPI vendor with {python_executable or sys.executable!r}") + return _mpi_vendor_family(vendor), vendor + + +@contextlib.contextmanager +def _launcher_hostfile(topology: MpiDirectTopology, launcher_family: str): + if launcher_family != "mpich": + yield None + return + fd, path = tempfile.mkstemp(prefix="simpler-mpi-hosts-", suffix=".txt", text=True) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + for host, slots in _host_slots(topology): + stream.write(f"{host}:{slots}\n") + yield path + finally: + with contextlib.suppress(FileNotFoundError): + os.unlink(path) + + +def _family_launcher_args(topology: MpiDirectTopology, launcher_family: str, hostfile_path: str | None) -> list[str]: + if launcher_family == "mpich": + if not hostfile_path: + raise ValueError("MPICH/Hydra launch requires a local hostfile") + args = ["-f", hostfile_path] + for name in _EXPORTED_ENV_VARS: + if name in os.environ: + if name == MPI_DIRECT_STARTUP_TOKEN_ENV: + args.extend(("-genvlist", name)) + else: + args.extend(("-genv", name, os.environ[name])) + return args + if launcher_family == "openmpi": + host_spec = ",".join(f"{host}:{slots}" for host, slots in _host_slots(topology)) + args = ["--host", host_spec, "--map-by", "slot", "--bind-to", "none"] + if os.geteuid() == 0: + args.append("--allow-run-as-root") + for name in _EXPORTED_ENV_VARS: + if name in os.environ: + args.extend(("-x", name)) + return args + raise ValueError(f"unsupported MPI launcher family {launcher_family!r}") + + +def _build_command( # noqa: PLR0913 -- launcher construction mirrors the complete MPI CLI surface + topology: MpiDirectTopology, + *, + mpirun_path: str, + topology_path: str | None, + session_id: int, + controller: str, + launcher_family: str, + hostfile_path: str | None = None, + manifest_json: str | None = None, + python_executable: str | None = None, + startup_host: str | None = None, + startup_port: int | None = None, +) -> list[str]: + if manifest_json is None and topology_path is None: + raise ValueError("command requires topology_path or manifest_json") + gate_enabled = startup_host is not None or startup_port is not None + if gate_enabled and not (startup_host and startup_port): + raise ValueError("startup gate requires host and port") + command = [ + mpirun_path, + *_family_launcher_args(topology, launcher_family, hostfile_path), + *topology.launcher_args, + "-np", + str(topology.world_size), + python_executable or sys.executable, + "-m", + "simpler.mpi_direct_runtime", + ] + if manifest_json is not None: + command.extend(("--manifest-json", manifest_json)) + else: + command.extend(("--topology", str(topology_path), "--session-id", str(session_id))) + command.extend(("--controller", controller)) + if gate_enabled: + command.extend( + ( + "--startup-host", + str(startup_host), + "--startup-port", + str(startup_port), + ) + ) + return command + + +def _gate_send(sock: socket.socket, message: dict[str, object]) -> None: + payload = json.dumps(message, separators=(",", ":")).encode("utf-8") + if len(payload) > MPI_DIRECT_GATE_MAX_BYTES: + raise ValueError("MPI startup gate message is too large") + sock.sendall(struct.pack("!I", len(payload)) + payload) + + +def _gate_recv(sock: socket.socket) -> dict[str, object]: + header = bytearray() + while len(header) < 4: + chunk = sock.recv(4 - len(header)) + if not chunk: + raise ConnectionError("MPI startup gate peer closed") + header.extend(chunk) + length = struct.unpack("!I", header)[0] + if length <= 0 or length > MPI_DIRECT_GATE_MAX_BYTES: + raise ValueError("invalid MPI startup gate message length") + payload = bytearray() + while len(payload) < length: + chunk = sock.recv(length - len(payload)) + if not chunk: + raise ConnectionError("MPI startup gate peer closed while reading") + payload.extend(chunk) + message = json.loads(payload) + if not isinstance(message, dict): + raise ValueError("MPI startup gate message must be an object") + return message + + +def _startup_gate( + topology: MpiDirectTopology, + token: str, + listener: socket.socket, + proc: subprocess.Popen[Any], +) -> None: + listener.settimeout(0.25) + deadline = time.monotonic() + topology.startup_timeout_s + peers: dict[int, socket.socket] = {} + try: + while len(peers) < topology.world_size: + if proc.poll() is not None: + raise RuntimeError(f"MPI job exited before startup gate completed (status {proc.returncode})") + if time.monotonic() >= deadline: + raise TimeoutError("MPI startup gate timed out waiting for all ranks") + try: + peer, _address = listener.accept() + except socket.timeout: + continue + try: + peer.settimeout(max(0.5, deadline - time.monotonic())) + message = _gate_recv(peer) + if message.get("token") != token: + raise RuntimeError("MPI startup gate token mismatch") + rank = int(message.get("rank", -1)) # type: ignore[arg-type] + if rank < 0 or rank >= topology.world_size or rank in peers: + raise RuntimeError(f"invalid or duplicate MPI startup rank {rank}") + if message.get("state") == "failed": + raise _StartupGateRankFailure( + f"rank {rank} failed before MPI initialization: {message.get('error', '')}" + ) + if message.get("state") != "ready": + raise RuntimeError(f"rank {rank} sent invalid startup state") + peers[rank] = peer + except _StartupGateRankFailure: + with contextlib.suppress(OSError): + peer.close() + raise + except Exception: + with contextlib.suppress(OSError): + peer.close() + continue + for rank in sorted(peers): + _gate_send(peers[rank], {"token": token, "state": "go_mpi"}) + finally: + for peer in peers.values(): + with contextlib.suppress(OSError): + peer.close() + listener.close() + + +def _terminate_job(proc: subprocess.Popen[Any], timeout_s: float = 5.0) -> None: + if proc.poll() is not None: + proc.wait() + return + with contextlib.suppress(OSError, ProcessLookupError): + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=timeout_s) + return + except subprocess.TimeoutExpired: + pass + with contextlib.suppress(OSError, ProcessLookupError): + os.killpg(proc.pid, signal.SIGKILL) + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=timeout_s) + + +@contextlib.contextmanager +def _startup_token_environment(token: str): + previous = os.environ.get(MPI_DIRECT_STARTUP_TOKEN_ENV) + os.environ[MPI_DIRECT_STARTUP_TOKEN_ENV] = token + try: + yield + finally: + if previous is None: + os.environ.pop(MPI_DIRECT_STARTUP_TOKEN_ENV, None) + else: + os.environ[MPI_DIRECT_STARTUP_TOKEN_ENV] = previous + + +def run_supervisor( + topology_path: str, + controller: str, + *, + mpirun_path: str = "mpirun", + launcher_family: str = "auto", + python_executable: str | None = None, +) -> int: + topology = MpiDirectTopology.load(topology_path) + if launcher_family == "auto": + launcher_family = _detect_launcher_family(mpirun_path) + mpi4py_family, mpi4py_vendor = _detect_mpi4py_family(python_executable) + if mpi4py_family != launcher_family: + raise RuntimeError( + f"MPI launcher family is {launcher_family}, but mpi4py is linked to " + f"{mpi4py_vendor} ({mpi4py_family}); use a matching launcher or rebuild " + "mpi4py with the launcher's MPI compiler" + ) + session_id = secrets.randbits(64) or 1 + manifest = topology.runtime_manifest(session_id) + manifest_json = base64.urlsafe_b64encode(json.dumps(manifest, separators=(",", ":")).encode("utf-8")).decode( + "ascii" + ) + gate_token = secrets.token_urlsafe(32) + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind((topology.controller_host, 0)) + listener.listen(topology.world_size) + gate_host = topology.controller_host + gate_port = int(listener.getsockname()[1]) + proc: subprocess.Popen[Any] | None = None + try: + with _launcher_hostfile(topology, launcher_family) as hostfile_path: + with _startup_token_environment(gate_token): + command = _build_command( + topology, + mpirun_path=mpirun_path, + topology_path=None, + session_id=session_id, + controller=controller, + launcher_family=launcher_family, + hostfile_path=hostfile_path, + manifest_json=manifest_json, + python_executable=python_executable, + startup_host=gate_host, + startup_port=gate_port, + ) + proc = subprocess.Popen(command, start_new_session=True) + _startup_gate(topology, gate_token, listener, proc) + return int(proc.wait()) + except BaseException: + with contextlib.suppress(OSError): + listener.close() + if proc is not None: + _terminate_job(proc) + raise + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--topology", required=True) + parser.add_argument("--controller", required=True, help="Python file with main(context), or module:callable") + parser.add_argument("--mpirun", default="mpirun") + parser.add_argument("--launcher-family", choices=("auto", "openmpi", "mpich"), default="auto") + parser.add_argument("--python", dest="python_executable", help="Python executable available on every MPI host") + ns = parser.parse_args(argv) + return run_supervisor( + ns.topology, + ns.controller, + mpirun_path=ns.mpirun, + launcher_family=ns.launcher_family, + python_executable=ns.python_executable, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/simpler/mpi_direct_topology.py b/python/simpler/mpi_direct_topology.py new file mode 100644 index 0000000000..93d1a1e4b3 --- /dev/null +++ b/python/simpler/mpi_direct_topology.py @@ -0,0 +1,230 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Validated static topology for the direct-MPI L4 control plane.""" + +from __future__ import annotations + +import ipaddress +import json +import math +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from .remote_l3_limits import MAX_FRAME_BYTES + + +def _positive_finite(value: Any, field: str) -> float: + result = float(value) + if not (result > 0.0 and math.isfinite(result)): + raise ValueError(f"{field} must be a positive finite number") + return result + + +def _is_loopback_host(host: str) -> bool: + normalized = host.strip().lower().rstrip(".") + if normalized == "localhost": + return True + try: + return ipaddress.ip_address(normalized).is_loopback + except ValueError: + return False + + +@dataclass(frozen=True) +class MpiDirectExecutorSpec: + rank: int + worker_id: int + host: str + platform: str + runtime: str + device_ids: tuple[int, ...] + num_sub_workers: int + comm_profile: str + global_device_ranks: tuple[int, ...] + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> MpiDirectExecutorSpec: + spec = cls( + rank=int(data["rank"]), + worker_id=int(data["worker_id"]), + host=str(data["host"]), + platform=str(data["platform"]), + runtime=str(data.get("runtime", "tensormap_and_ringbuffer")), + device_ids=tuple(int(item) for item in data.get("device_ids", ())), + num_sub_workers=int(data.get("num_sub_workers", 0)), + comm_profile=str(data.get("comm_profile", "sim")), + global_device_ranks=tuple(int(item) for item in data.get("global_device_ranks", ())), + ) + if spec.rank <= 0: + raise ValueError("executor rank must be positive; rank 0 is the controller") + if spec.worker_id < 0: + raise ValueError("executor worker_id must be non-negative") + if not spec.host or not spec.platform or not spec.runtime or not spec.comm_profile: + raise ValueError("executor host, platform, runtime, and comm_profile must be non-empty") + if spec.num_sub_workers < 0: + raise ValueError("executor num_sub_workers must be non-negative") + if len(set(spec.device_ids)) != len(spec.device_ids) or any(device < 0 for device in spec.device_ids): + raise ValueError("executor device_ids must be unique and non-negative") + if spec.global_device_ranks and len(spec.global_device_ranks) != len(spec.device_ids): + raise ValueError("executor global_device_ranks must match device_ids length") + if len(set(spec.global_device_ranks)) != len(spec.global_device_ranks) or any( + rank < 0 for rank in spec.global_device_ranks + ): + raise ValueError("executor global_device_ranks must be unique and non-negative") + return spec + + def to_dict(self) -> dict[str, Any]: + return { + "rank": self.rank, + "worker_id": self.worker_id, + "host": self.host, + "platform": self.platform, + "runtime": self.runtime, + "device_ids": list(self.device_ids), + "num_sub_workers": self.num_sub_workers, + "comm_profile": self.comm_profile, + "global_device_ranks": list(self.global_device_ranks), + } + + +@dataclass(frozen=True) +class MpiDirectTopology: + controller_rank: int + controller_host: str + executors: tuple[MpiDirectExecutorSpec, ...] + startup_timeout_s: float + session_timeout_s: float + heartbeat_interval_s: float + max_pending_frame_bytes: int + launcher_args: tuple[str, ...] + + @property + def world_size(self) -> int: + return 1 + len(self.executors) + + @property + def hosts(self) -> tuple[str, ...]: + return (self.controller_host,) + tuple(spec.host for spec in self.executors) + + def executor_for_rank(self, rank: int) -> MpiDirectExecutorSpec: + if rank <= 0 or rank >= self.world_size: + raise ValueError(f"MPI rank {rank} is outside executor range") + spec = self.executors[rank - 1] + if spec.rank != rank: + raise RuntimeError("executor table is not rank ordered") + return spec + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> MpiDirectTopology: + parsed = tuple(MpiDirectExecutorSpec.from_dict(item) for item in data.get("executor_ranks", ())) + explicit = {rank for spec in parsed for rank in spec.global_device_ranks} + next_rank = 0 + executors_list = [] + for spec in parsed: + if spec.global_device_ranks: + executors_list.append(spec) + continue + assigned = [] + for _device_id in spec.device_ids: + while next_rank in explicit: + next_rank += 1 + assigned.append(next_rank) + explicit.add(next_rank) + next_rank += 1 + executors_list.append(replace(spec, global_device_ranks=tuple(assigned))) + executors = tuple(executors_list) + topology = cls( + controller_rank=int(data.get("controller_rank", 0)), + controller_host=str(data.get("controller_host", "localhost")), + executors=executors, + startup_timeout_s=_positive_finite(data.get("startup_timeout_s", 60.0), "startup_timeout_s"), + session_timeout_s=_positive_finite(data.get("session_timeout_s", 30.0), "session_timeout_s"), + heartbeat_interval_s=_positive_finite(data.get("heartbeat_interval_s", 1.0), "heartbeat_interval_s"), + max_pending_frame_bytes=int(data.get("max_pending_frame_bytes", 64 * 1024 * 1024)), + launcher_args=tuple(str(arg) for arg in data.get("launcher_args", ())), + ) + topology.validate() + return topology + + @classmethod + def load(cls, path: str) -> MpiDirectTopology: + with Path(path).open("r", encoding="utf-8") as stream: + data = json.load(stream) + if not isinstance(data, dict): + raise ValueError("MPI direct topology root must be an object") + return cls.from_dict(data) + + def validate(self) -> None: + if self.controller_rank != 0: + raise ValueError("PR3 requires controller_rank=0") + if not self.controller_host: + raise ValueError("controller_host must be non-empty") + if not self.executors: + raise ValueError("executor_ranks must contain at least one executor") + if len(set(self.hosts)) > 1 and _is_loopback_host(self.controller_host): + raise ValueError( + f"controller_host={self.controller_host!r} is a loopback address, but the topology spans multiple " + "hosts; set controller_host to an address reachable from all MPI hosts" + ) + closed_hosts: set[str] = set() + previous_host: str | None = None + for host in self.hosts: + if host != previous_host: + if host in closed_hosts: + raise ValueError("topology hosts must be contiguous in rank order") + if previous_host is not None: + closed_hosts.add(previous_host) + previous_host = host + if [spec.rank for spec in self.executors] != list(range(1, self.world_size)): + raise ValueError("executor ranks must be dense and ordered from 1 to world_size-1") + if [spec.worker_id for spec in self.executors] != list(range(len(self.executors))): + raise ValueError("executor worker_ids must be dense and ordered from 0") + all_global_ranks = [rank for spec in self.executors for rank in spec.global_device_ranks] + if len(set(all_global_ranks)) != len(all_global_ranks): + raise ValueError("global_device_ranks must be unique across all executors") + if self.max_pending_frame_bytes < MAX_FRAME_BYTES: + raise ValueError("max_pending_frame_bytes must fit one maximum SLR3 frame") + + def runtime_manifest(self, session_id: int) -> dict[str, Any]: + if int(session_id) == 0: + raise ValueError("session_id must be non-zero") + return { + "version": 1, + "controller_rank": self.controller_rank, + "controller_host": self.controller_host, + "world_size": self.world_size, + "session_id": int(session_id), + "startup_timeout_s": self.startup_timeout_s, + "session_timeout_s": self.session_timeout_s, + "heartbeat_interval_s": self.heartbeat_interval_s, + "max_pending_frame_bytes": self.max_pending_frame_bytes, + "executor_ranks": [spec.to_dict() for spec in self.executors], + } + + +def load_runtime_manifest(path: str) -> tuple[MpiDirectTopology, int]: + with Path(path).open("r", encoding="utf-8") as stream: + data = json.load(stream) + if not isinstance(data, dict) or int(data.get("version", 0)) != 1: + raise ValueError("unsupported MPI direct runtime manifest") + return load_runtime_manifest_data(data) + + +def load_runtime_manifest_data(data: dict[str, Any]) -> tuple[MpiDirectTopology, int]: + if not isinstance(data, dict) or int(data.get("version", 0)) != 1: + raise ValueError("unsupported MPI direct runtime manifest") + topology = MpiDirectTopology.from_dict(data) + declared_world_size = int(data.get("world_size", 0)) + if declared_world_size != topology.world_size: + raise ValueError("runtime manifest world_size does not match executor table") + session_id = int(data.get("session_id", 0)) + if session_id == 0: + raise ValueError("runtime manifest session_id must be non-zero") + return topology, session_id diff --git a/python/simpler/remote_l3_limits.py b/python/simpler/remote_l3_limits.py new file mode 100644 index 0000000000..37dff7aba1 --- /dev/null +++ b/python/simpler/remote_l3_limits.py @@ -0,0 +1,13 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""SLR3 frame limits shared by extension-free launcher code.""" + +FRAME_HEADER_BYTES = 40 +MAX_FRAME_PAYLOAD_BYTES = 16 * 1024 * 1024 +MAX_FRAME_BYTES = FRAME_HEADER_BYTES + MAX_FRAME_PAYLOAD_BYTES diff --git a/python/simpler/remote_l3_protocol.py b/python/simpler/remote_l3_protocol.py index b19b067c95..ef209eb6dd 100644 --- a/python/simpler/remote_l3_protocol.py +++ b/python/simpler/remote_l3_protocol.py @@ -24,13 +24,13 @@ CanonicalIdentity, Tensor, ) +from .remote_l3_limits import FRAME_HEADER_BYTES, MAX_FRAME_PAYLOAD_BYTES from .task_interface import MAX_TENSOR_DIMS, CallConfig, DataType # 3: a TASK's per-argument record is the self-describing wire ``Tensor`` — the embedded # BufferDescriptor plus the strided view. Both ends of a run come from one ``pip install``, # so this constant is a mismatch alarm at the frame header, not a dual-decode selector. PROTOCOL_VERSION = 3 -MAX_FRAME_PAYLOAD_BYTES = 16 * 1024 * 1024 MAX_STRING_BYTES = 1024 MAX_ERROR_BYTES = 4096 MAX_TENSORS = 4096 @@ -48,8 +48,6 @@ REMOTE_BUFFER_ACCESS_WRITE = 1 << 1 REMOTE_BUFFER_ACCESS_READ_WRITE = REMOTE_BUFFER_ACCESS_READ | REMOTE_BUFFER_ACCESS_WRITE CALLABLE_HASH_DIGEST_BYTES = 32 -FRAME_HEADER_BYTES = 40 - # FrameHeader.flags bit: the caller addresses every member of the target's # worker group, not just the worker named in the header. Only group-capable # transports (the MPI mailbox) act on it; point-to-point transports ignore it. diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 90e1a715da..742fc7828e 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -97,6 +97,7 @@ def my_l4_orch(orch, args, config): _l3_child_onboard_region_create, _mailbox_load_i32, _mailbox_store_i32, + _MpiDirectTransportHub, _read_control_copy_request, _set_host_span_level_prefix, _worker_host_mapped_region_ack_cleanup_error, @@ -834,6 +835,22 @@ def __post_init__(self) -> None: # noqa: PLR0912 -- one place validates the pub object.__setattr__(self, "python_executable", str(self.python_executable)) +@dataclass(frozen=True) +class _MpiDirectWorkerSpec: + worker_id: int + mpi_rank: int + session_id: int + host: str + comm_profile: str + platform: str + runtime: str + device_ids: tuple[int, ...] + global_device_ranks: tuple[int, ...] + hub: _MpiDirectTransportHub + attach_timeout_s: float + runtime_timeout_s: float + + @dataclass(frozen=True) class _RemoteSession: worker_id: int @@ -1744,6 +1761,8 @@ def _chip_descriptor_context(worker: Worker) -> tuple[str, str]: contexts.append((str(spec.platform), str(spec.runtime))) for rank in getattr(worker, "_mpi_rank_by_worker_id", {}).values(): contexts.append((str(rank.spec.platform), str(rank.spec.runtime))) + for spec in getattr(worker, "_mpi_direct_worker_specs", []): + contexts.append((str(spec.platform), str(spec.runtime))) if not contexts: return "", "" first = contexts[0] @@ -4375,6 +4394,17 @@ def _on_cancel(_signum, _frame): os._exit(0) +def _close_fork_child_fds(fds) -> None: + for raw_fd in fds: + try: + fd = int(raw_fd) + except (TypeError, ValueError, OverflowError): + continue + if fd >= 3: + with contextlib.suppress(OSError, OverflowError): + os.close(fd) + + # --------------------------------------------------------------------------- # Worker factory # --------------------------------------------------------------------------- @@ -4567,6 +4597,8 @@ def __init__( self._mpi_l3_groups: list[_MpiL3GroupRuntime] = [] self._mpi_worker_ids: list[int] = [] self._mpi_rank_by_worker_id: dict[int, _MpiL3RankRuntime] = {} + self._mpi_direct_worker_specs: list[_MpiDirectWorkerSpec] = [] + self._mpi_direct_worker_ids: list[int] = [] self._next_level_worker_id_count: int = 0 # Fallback ownership for private helpers used outside Worker.submit. # Normal orchestration-owned refs live in RunHandle._resources. @@ -4775,7 +4807,23 @@ def add_mpirun_worker_group(self, spec: MpiL3GroupSpec) -> tuple[int, ...]: return tuple(rank.worker_id for rank in ranks) def _remote_like_worker_ids(self) -> set[int]: - return set(self._remote_worker_ids) | set(self._mpi_worker_ids) + return set(self._remote_worker_ids) | set(self._mpi_worker_ids) | set(self._mpi_direct_worker_ids) + + def _add_mpi_direct_worker(self, spec: _MpiDirectWorkerSpec) -> int: + with self._hierarchical_start_cv: + if self._lifecycle is not _Lifecycle.NEW: + raise RuntimeError("Worker._add_mpi_direct_worker after init") + if self.level < 4: + raise TypeError("direct MPI L3 workers require a level >= 4 parent") + if not isinstance(spec, _MpiDirectWorkerSpec): + raise TypeError("Worker._add_mpi_direct_worker expects an MPI direct worker spec") + expected_id = self._next_level_worker_id_count + if spec.worker_id != expected_id: + raise ValueError("MPI direct worker ids must match dense NEXT_LEVEL allocation order") + worker_id = self._allocate_next_level_worker_id() + self._mpi_direct_worker_specs.append(spec) + self._mpi_direct_worker_ids.append(worker_id) + return worker_id @staticmethod def _parse_remote_endpoint(endpoint: str) -> tuple[str, int]: @@ -4915,7 +4963,9 @@ def _remote_dispatcher_entries_for_worker(self, worker_id: int) -> list[dict[str ) return entries - def _inner_registry_entries_for_spec(self, spec: RemoteWorkerSpec | _MpiL3RankSpec) -> list[dict[str, Any]]: + def _inner_registry_entries_for_spec( + self, spec: RemoteWorkerSpec | _MpiL3RankSpec | _MpiDirectWorkerSpec + ) -> list[dict[str, Any]]: from .remote_l3_protocol import ( # noqa: PLC0415 ChipCallableBlobLocation, RemoteChipCallablePayload, @@ -4996,6 +5046,17 @@ def _resolved_global_nodes(self) -> dict[int, _GlobalNodeRuntime]: True, ) ) + for spec in self._mpi_direct_worker_specs: + configs.append( + ( + int(spec.worker_id), + tuple(spec.device_ids), + spec.platform, + spec.comm_profile, + tuple(spec.global_device_ranks), + True, + ) + ) for worker_id, child in zip(self._next_level_worker_ids, self._next_level_workers): if child.level != 3: continue @@ -5382,6 +5443,78 @@ def _activate_mpirun_worker_groups(self, deadline: float) -> None: if time.monotonic() >= deadline: raise RuntimeError("MPI L3 activation: startup deadline exceeded after attach") + def _activate_mpi_direct_workers(self, deadline: float) -> None: + if not self._mpi_direct_worker_specs: + return + assert self._worker is not None + session_timeout = self._remote_session_timeout_s() + for spec in self._mpi_direct_worker_specs: + remaining = self._remaining_until(deadline, "direct MPI L3 endpoint attach") + self._worker.add_remote_l3_mpi( + spec.worker_id, + spec.session_id, + spec.comm_profile, + spec.hub, + remaining, + min(session_timeout, spec.runtime_timeout_s), + ) + + def _publish_initial_mpi_direct_callables(self) -> None: + if not self._mpi_direct_worker_specs: + return + assert self._worker is not None + direct_worker_ids = {spec.worker_id for spec in self._mpi_direct_worker_specs} + with self._registry_lock: + states = tuple(self._identity_registry.values()) + for state in states: + if state.target_namespace == "REMOTE_TASK_DISPATCHER": + targets = tuple(worker_id for worker_id in state.eligible_worker_ids if worker_id in direct_worker_ids) + if not targets: + continue + registration = _build_callable_registration(self, state.target, workers=list(state.eligible_worker_ids)) + target_registry = "REMOTE_TASK_DISPATCHER" + callable_kind = registration.kind + payloads = {worker_id: registration.payload or b"" for worker_id in targets} + elif state.target_namespace == "LOCAL_CHIP": + targets = tuple(spec.worker_id for spec in self._mpi_direct_worker_specs) + target_registry = "INNER_L3_WORKER" + callable_kind = "CHIP_CALLABLE" + payloads = {} + for spec in self._mpi_direct_worker_specs: + entries = self._inner_registry_entries_for_spec(spec) + entry = next((item for item in entries if item["hashid"] == state.digest.hex()), None) + if entry is None: + raise RuntimeError(f"direct MPI inner chip hashid {state.hashid} was not serialised") + payloads[spec.worker_id] = bytes.fromhex(str(entry["payload_hex"])) + else: + continue + prepared: list[int] = [] + committed: list[int] = [] + try: + for worker_id in targets: + result = self._worker.remote_prepare_register( + worker_id, target_registry, callable_kind, payloads[worker_id], state.digest + ) + if not result.ok: + raise RuntimeError(result.error_message) + prepared.append(worker_id) + for worker_id in targets: + result = self._worker.remote_commit_register( + worker_id, target_registry, callable_kind, state.digest + ) + if not result.ok: + raise RuntimeError(result.error_message) + committed.append(worker_id) + except BaseException: + uncommitted = [worker_id for worker_id in prepared if worker_id not in committed] + for worker_id in uncommitted: + with contextlib.suppress(BaseException): + self._worker.remote_abort_register(worker_id, target_registry, callable_kind, state.digest) + for worker_id in committed: + with contextlib.suppress(BaseException): + self._worker.remote_unregister(worker_id, target_registry, callable_kind, state.digest) + raise + def _require_remote_worker_started(self, worker_id: int) -> None: """Argument + resource gate for the public remote-memory APIs. Admission (READY) is decided by the ``_operation_lease`` these APIs already hold — @@ -6387,6 +6520,11 @@ def _append_endpoint_topology( mpi_path = _format_worker_path(3, parent_path=path, index=int(child_index)) entries.append(_EndpointTopologyEntry(mpi_path, HOST_CPU, mpi_node_identity)) self._append_device_endpoint_topology(entries, mpi_path, rank.spec.device_ids, mpi_node_identity) + for spec in worker._mpi_direct_worker_specs: + mpi_node_identity = _normalize_node_identity(spec.host) + mpi_path = _format_worker_path(3, parent_path=path, index=int(spec.worker_id)) + entries.append(_EndpointTopologyEntry(mpi_path, HOST_CPU, mpi_node_identity)) + self._append_device_endpoint_topology(entries, mpi_path, spec.device_ids, mpi_node_identity) def _append_device_endpoint_topology( self, @@ -6488,7 +6626,7 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl raise TypeError("Worker.register: level 2 only supports ChipCallable targets") reg = _build_callable_registration(self, target, workers=workers) if isinstance(target, RemoteCallable): - if not self._remote_worker_specs and not self._mpi_l3_groups: + if not self._remote_worker_specs and not self._mpi_l3_groups and not self._mpi_direct_worker_specs: raise RuntimeError("Worker.register(RemoteCallable): add at least one remote worker first") remote_worker_ids = self._remote_like_worker_ids() for worker_id in reg.eligible_worker_ids: @@ -7432,6 +7570,8 @@ def has_chip_target(worker: Worker) -> bool: return True if any(rank.spec.device_ids for rank in worker._mpi_rank_by_worker_id.values()): return True + if any(spec.device_ids for spec in worker._mpi_direct_worker_specs): + return True return any(has_chip_target(child) for child in worker._next_level_workers) return None if has_chip_target(self) else "a chip device (device_ids)" @@ -7629,7 +7769,7 @@ def _init_hierarchical(self) -> None: # startup resource (mailbox shm, pre-fork _Worker mmap, child fork, # daemon socket) exists, so an invalid value fails without a # partially-built subtree to roll back. - if self._remote_worker_specs or self._mpi_l3_groups: + if self._remote_worker_specs or self._mpi_l3_groups or self._mpi_direct_worker_specs: self._remote_session_timeout_s() # 1. Allocate sub-worker mailboxes (unified layout, MAILBOX_SIZE each). @@ -7814,6 +7954,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l if pid == 0: buf = self._sub_shms[i].buf assert buf is not None + _close_fork_child_fds(self._config.get("fork_child_close_fds", ())) def _setup(): return _make_local_identity_tables( @@ -7851,6 +7992,7 @@ def _setup(): if self._is_startup_root: with contextlib.suppress(OSError): os.setpgid(0, 0) + _close_fork_child_fds(self._config.get("fork_child_close_fds", ())) # _chip_process_loop publishes INIT_READY/INIT_FAILED itself # (around cw.init + ChipCallable prepare). This guard only # ensures the child exits rather than unwinding into the @@ -7922,6 +8064,7 @@ def _setup(): if pid == 0: buf = self._next_level_shms[idx].buf assert buf is not None + _close_fork_child_fds(self._config.get("fork_child_close_fds", ())) def _setup(inner=inner_worker): # Propagate the fork-constant prewarm sizing and the shared @@ -7972,6 +8115,7 @@ def _setup(inner=inner_worker): # the RemoteL3Endpoint health thread, so both must follow every local # fork. Each remote consumes this process's remaining startup budget. self._activate_mpirun_worker_groups(deadline) + self._activate_mpi_direct_workers(deadline) self._activate_remote_sessions(deadline) # _Worker was constructed in _init_hierarchical (pre-fork) so children @@ -8009,6 +8153,8 @@ def _setup(inner=inner_worker): self._orch = Orchestrator(dw.get_orchestrator(), self) + self._publish_initial_mpi_direct_callables() + # Every ChipCallable in the startup snapshot was already uploaded by its # chip child before that child published INIT_READY (see # _chip_process_loop), and the runtime arena was prewarmed there too — so diff --git a/simpler_setup/environment.py b/simpler_setup/environment.py index bf65e9f055..a247e6c834 100644 --- a/simpler_setup/environment.py +++ b/simpler_setup/environment.py @@ -9,7 +9,7 @@ """Centralized path management. PROJECT_ROOT auto-resolves between two layouts: - - wheel install: simpler_setup/_assets/{src,build/lib} populated by CMakeLists install() + - wheel install: simpler_setup/_assets/{src,cmake,build/lib} populated by CMakeLists install() - source tree / editable: repo root with src/ and build/lib/ in original positions """ diff --git a/src/common/hierarchical/mpi_direct_transport.cpp b/src/common/hierarchical/mpi_direct_transport.cpp new file mode 100644 index 0000000000..7a8d4296ff --- /dev/null +++ b/src/common/hierarchical/mpi_direct_transport.cpp @@ -0,0 +1,378 @@ +/* + * 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 "mpi_direct_transport.h" + +#include +#include +#include + +namespace { + +std::chrono::steady_clock::time_point deadline_from_now(double timeout_s) { + return std::chrono::steady_clock::now() + + std::chrono::duration_cast(std::chrono::duration(timeout_s)); +} + +MpiDirectTag outbound_tag(remote_l3::FrameType type) { + switch (type) { + case remote_l3::FrameType::TASK: + case remote_l3::FrameType::CONTROL: + return MpiDirectTag::COMMAND_REQUEST; + case remote_l3::FrameType::SHUTDOWN: + return MpiDirectTag::LIFECYCLE; + default: + throw std::runtime_error("MpiDirectTransport: invalid outbound SLR3 frame type"); + } +} + +MpiDirectTag inbound_tag(remote_l3::FrameType type) { + switch (type) { + case remote_l3::FrameType::COMPLETION: + case remote_l3::FrameType::CONTROL_REPLY: + return MpiDirectTag::COMMAND_REPLY; + case remote_l3::FrameType::HEALTH: + return MpiDirectTag::HEALTH; + case remote_l3::FrameType::HELLO: + return MpiDirectTag::LIFECYCLE; + default: + throw std::runtime_error("MpiDirectTransportHub: invalid inbound SLR3 frame type"); + } +} + +} // namespace + +MpiDirectTransportHub::MpiDirectTransportHub(size_t max_pending_frame_bytes) : + max_pending_frame_bytes_(max_pending_frame_bytes) { + const size_t max_frame_bytes = remote_l3::FRAME_HEADER_BYTES + remote_l3::MAX_FRAME_PAYLOAD_BYTES; + if (max_pending_frame_bytes_ < max_frame_bytes) { + throw std::invalid_argument("MpiDirectTransportHub: pending byte budget must fit one maximum SLR3 frame"); + } +} + +void MpiDirectTransportHub::register_route( + int32_t worker_id, int32_t mpi_rank, uint64_t session_id, const std::string &comm_profile +) { + if (worker_id < 0) throw std::invalid_argument("MpiDirectTransportHub: worker id must be non-negative"); + if (mpi_rank <= 0) throw std::invalid_argument("MpiDirectTransportHub: executor rank must be positive"); + if (session_id == 0) throw std::invalid_argument("MpiDirectTransportHub: session id must be non-zero"); + if (comm_profile.empty()) throw std::invalid_argument("MpiDirectTransportHub: comm profile must be non-empty"); + + std::lock_guard lk(mu_); + throw_if_terminal_locked(); + if (routes_by_worker_.count(worker_id) != 0 || worker_by_rank_.count(mpi_rank) != 0) { + throw std::invalid_argument("MpiDirectTransportHub: duplicate worker id or MPI rank"); + } + Route route; + route.worker_id = worker_id; + route.mpi_rank = mpi_rank; + route.session_id = session_id; + route.comm_profile = comm_profile; + routes_by_worker_.emplace(worker_id, std::move(route)); + worker_by_rank_.emplace(mpi_rank, worker_id); +} + +void MpiDirectTransportHub::throw_if_terminal_locked() const { + if (!terminal_error_.empty()) throw std::runtime_error("MpiDirectTransportHub: " + terminal_error_); + if (closed_) throw std::runtime_error("MpiDirectTransportHub: closed"); +} + +void MpiDirectTransportHub::throw_if_route_terminal_locked(const Route &route) const { + throw_if_terminal_locked(); + if (!route.terminal_error.empty()) { + throw std::runtime_error( + "MpiDirectTransportHub: route worker " + std::to_string(route.worker_id) + ": " + route.terminal_error + ); + } +} + +void MpiDirectTransportHub::fail_locked(const std::string &message) { + if (terminal_error_.empty()) terminal_error_ = message.empty() ? "terminal failure" : message; + cv_.notify_all(); +} + +void MpiDirectTransportHub::fail(const std::string &message) { + std::lock_guard lk(mu_); + fail_locked(message); +} + +void MpiDirectTransportHub::close() { + std::lock_guard lk(mu_); + closed_ = true; + cv_.notify_all(); +} + +void MpiDirectTransportHub::cancel_route(int32_t worker_id, const std::string &message) { + std::lock_guard lk(mu_); + auto route_it = routes_by_worker_.find(worker_id); + if (route_it == routes_by_worker_.end()) throw std::invalid_argument("MpiDirectTransportHub: unknown worker id"); + if (route_it->second.terminal_error.empty()) { + route_it->second.terminal_error = message.empty() ? "cancelled" : message; + } + cv_.notify_all(); +} + +void MpiDirectTransportHub::enqueue( + int32_t worker_id, MpiDirectTag tag, const std::vector &frame, double timeout_s +) { + if (!(timeout_s > 0.0) || !std::isfinite(timeout_s)) { + throw std::invalid_argument("MpiDirectTransportHub: timeout must be positive and finite"); + } + if (frame.size() > remote_l3::FRAME_HEADER_BYTES + remote_l3::MAX_FRAME_PAYLOAD_BYTES) { + throw std::invalid_argument("MpiDirectTransportHub: frame exceeds maximum"); + } + auto deadline = deadline_from_now(timeout_s); + std::unique_lock lk(mu_); + auto route_it = routes_by_worker_.find(worker_id); + if (route_it == routes_by_worker_.end()) throw std::invalid_argument("MpiDirectTransportHub: unknown worker id"); + while (frame.size() > max_pending_frame_bytes_ - pending_frame_bytes_) { + throw_if_route_terminal_locked(route_it->second); + if (cv_.wait_until(lk, deadline) == std::cv_status::timeout) { + fail_locked("timed out waiting for outbound frame credit"); + throw_if_terminal_locked(); + } + } + throw_if_route_terminal_locked(route_it->second); + MpiDirectOutboundFrame outbound; + outbound.ticket = next_ticket_++; + outbound.target_rank = route_it->second.mpi_rank; + outbound.tag = tag; + outbound.frame = frame; + pending_frame_bytes_ += frame.size(); + outbound_.push_back(std::move(outbound)); + cv_.notify_all(); +} + +std::optional> +MpiDirectTransportHub::poll_inbound(int32_t worker_id, remote_l3::FrameType frame_type, uint64_t sequence) { + std::lock_guard lk(mu_); + auto route_it = routes_by_worker_.find(worker_id); + if (route_it == routes_by_worker_.end()) throw std::invalid_argument("MpiDirectTransportHub: unknown worker id"); + throw_if_route_terminal_locked(route_it->second); + std::deque> &queue = + frame_type == remote_l3::FrameType::HELLO ? route_it->second.lifecycle : route_it->second.replies; + if (queue.empty()) return std::nullopt; + std::vector frame = std::move(queue.front()); + queue.pop_front(); + auto decoded = remote_l3::decode_frame(frame); + if (decoded.header.frame_type != frame_type || decoded.header.sequence != sequence) { + fail_locked("inbound frame type or sequence mismatch"); + throw_if_terminal_locked(); + } + return frame; +} + +std::optional MpiDirectTransportHub::poll_outbound(double timeout_s) { + if (timeout_s < 0.0 || !std::isfinite(timeout_s)) { + throw std::invalid_argument("MpiDirectTransportHub: poll timeout must be finite and non-negative"); + } + std::unique_lock lk(mu_); + if (outbound_.empty() && timeout_s > 0.0 && terminal_error_.empty() && !closed_) { + cv_.wait_for(lk, std::chrono::duration(timeout_s), [this] { + return !outbound_.empty() || !terminal_error_.empty() || closed_; + }); + } + if (outbound_.empty()) { + throw_if_terminal_locked(); + return std::nullopt; + } + MpiDirectOutboundFrame frame = std::move(outbound_.front()); + outbound_.pop_front(); + in_flight_bytes_.emplace(frame.ticket, frame.frame.size()); + return frame; +} + +void MpiDirectTransportHub::complete_outbound(uint64_t ticket) { + std::lock_guard lk(mu_); + auto it = in_flight_bytes_.find(ticket); + if (it == in_flight_bytes_.end()) { + fail_locked("unknown or already completed outbound ticket"); + throw_if_terminal_locked(); + } + pending_frame_bytes_ -= it->second; + in_flight_bytes_.erase(it); + cv_.notify_all(); +} + +void MpiDirectTransportHub::deliver(int32_t source_rank, MpiDirectTag tag, const std::vector &frame) { + remote_l3::DecodedFrame decoded; + try { + decoded = remote_l3::decode_frame(frame); + } catch (const std::exception &e) { + fail(std::string("invalid inbound SLR3 frame: ") + e.what()); + throw; + } + + std::lock_guard lk(mu_); + throw_if_terminal_locked(); + auto worker_it = worker_by_rank_.find(source_rank); + if (worker_it == worker_by_rank_.end()) { + fail_locked("frame arrived from an unknown MPI rank"); + throw_if_terminal_locked(); + } + Route &route = routes_by_worker_.at(worker_it->second); + if (!route.terminal_error.empty()) return; + if (decoded.header.worker_id != route.worker_id || decoded.header.session_id != route.session_id) { + fail_locked("inbound frame identity does not match manifest route"); + throw_if_terminal_locked(); + } + MpiDirectTag expected; + try { + expected = inbound_tag(decoded.header.frame_type); + } catch (const std::exception &e) { + fail_locked(e.what()); + throw_if_terminal_locked(); + } + if (tag != expected) { + fail_locked("inbound MPI tag does not match SLR3 frame type"); + throw_if_terminal_locked(); + } + if (tag == MpiDirectTag::HEALTH) { + if (!decoded.payload.empty()) { + fail_locked("HEALTH frame payload must be empty"); + throw_if_terminal_locked(); + } + route.last_health = std::chrono::steady_clock::now(); + } else if (tag == MpiDirectTag::LIFECYCLE) { + route.lifecycle.push_back(frame); + } else { + route.replies.push_back(frame); + } + cv_.notify_all(); +} + +std::vector MpiDirectTransportHub::wait_inbound( + int32_t worker_id, remote_l3::FrameType frame_type, uint64_t sequence, double timeout_s +) { + auto deadline = deadline_from_now(timeout_s); + std::unique_lock lk(mu_); + auto route_it = routes_by_worker_.find(worker_id); + if (route_it == routes_by_worker_.end()) throw std::invalid_argument("MpiDirectTransportHub: unknown worker id"); + std::deque> &queue = + frame_type == remote_l3::FrameType::HELLO ? route_it->second.lifecycle : route_it->second.replies; + while (queue.empty()) { + throw_if_route_terminal_locked(route_it->second); + if (cv_.wait_until(lk, deadline) == std::cv_status::timeout) { + fail_locked("timed out waiting for inbound frame"); + throw_if_terminal_locked(); + } + } + throw_if_route_terminal_locked(route_it->second); + std::vector frame = std::move(queue.front()); + queue.pop_front(); + auto decoded = remote_l3::decode_frame(frame); + if (decoded.header.frame_type != frame_type || decoded.header.sequence != sequence) { + fail_locked("inbound frame type or sequence mismatch"); + throw_if_terminal_locked(); + } + return frame; +} + +void MpiDirectTransportHub::expect_hello_ready(int32_t worker_id, double timeout_s) { + auto bytes = wait_inbound(worker_id, remote_l3::FrameType::HELLO, 0, timeout_s); + auto frame = remote_l3::decode_frame(bytes); + auto hello = remote_l3::decode_hello(frame.payload.data(), frame.payload.size()); + std::lock_guard lk(mu_); + Route &route = routes_by_worker_.at(worker_id); + throw_if_route_terminal_locked(route); + if (hello.session_id != route.session_id || hello.worker_id != worker_id || + hello.ready_state != remote_l3::ReadyState::READY || hello.comm_profile != route.comm_profile) { + fail_locked("HELLO READY does not match manifest route"); + throw_if_terminal_locked(); + } +} + +size_t MpiDirectTransportHub::pending_frame_bytes() const { + std::lock_guard lk(mu_); + return pending_frame_bytes_; +} + +bool MpiDirectTransportHub::terminal() const { + std::lock_guard lk(mu_); + return !terminal_error_.empty(); +} + +std::string MpiDirectTransportHub::terminal_error() const { + std::lock_guard lk(mu_); + return terminal_error_; +} + +MpiDirectTransport::MpiDirectTransport( + std::shared_ptr hub, int32_t worker_id, double attach_timeout_s, double runtime_timeout_s +) : + hub_(std::move(hub)), + worker_id_(worker_id), + attach_timeout_s_(attach_timeout_s), + runtime_timeout_s_(runtime_timeout_s) { + if (!hub_) throw std::invalid_argument("MpiDirectTransport: null hub"); + if (worker_id_ < 0) throw std::invalid_argument("MpiDirectTransport: worker id must be non-negative"); + if (!(attach_timeout_s_ > 0.0) || !std::isfinite(attach_timeout_s_)) { + throw std::invalid_argument("MpiDirectTransport: attach timeout must be positive and finite"); + } + if (!(runtime_timeout_s_ > 0.0) || !std::isfinite(runtime_timeout_s_)) { + throw std::invalid_argument("MpiDirectTransport: runtime timeout must be positive and finite"); + } +} + +void MpiDirectTransport::expect_hello_ready() { hub_->expect_hello_ready(worker_id_, attach_timeout_s_); } + +void MpiDirectTransport::submit_frame(const std::vector &frame) { + if (closed_.load(std::memory_order_acquire)) throw std::runtime_error("MpiDirectTransport: closed"); + if (progress_active_.load(std::memory_order_acquire)) { + throw std::logic_error("MpiDirectTransport: progress command is active"); + } + auto decoded = remote_l3::decode_frame(frame); + hub_->enqueue(worker_id_, outbound_tag(decoded.header.frame_type), frame, runtime_timeout_s_); +} + +std::vector MpiDirectTransport::wait_for_reply(remote_l3::FrameType frame_type, uint64_t sequence) { + if (closed_.load(std::memory_order_acquire)) throw std::runtime_error("MpiDirectTransport: closed"); + if (progress_active_.load(std::memory_order_acquire)) { + throw std::logic_error("MpiDirectTransport: progress command is active"); + } + return hub_->wait_inbound(worker_id_, frame_type, sequence, runtime_timeout_s_); +} + +void MpiDirectTransport::submit_progress_frame(const std::vector &frame) { + if (closed_.load(std::memory_order_acquire)) throw std::runtime_error("MpiDirectTransport: closed"); + if (progress_active_.load(std::memory_order_acquire)) { + throw std::logic_error("MpiDirectTransport: progress command is already active"); + } + auto decoded = remote_l3::decode_frame(frame); + hub_->enqueue(worker_id_, outbound_tag(decoded.header.frame_type), frame, runtime_timeout_s_); + progress_deadline_ = deadline_from_now(runtime_timeout_s_); + progress_active_.store(true, std::memory_order_release); +} + +bool MpiDirectTransport::poll_progress_reply( + remote_l3::FrameType frame_type, uint64_t sequence, std::vector &reply +) { + if (closed_.load(std::memory_order_acquire)) throw std::runtime_error("MpiDirectTransport: closed"); + if (!progress_active_.load(std::memory_order_acquire)) { + throw std::logic_error("MpiDirectTransport: no progress command is active"); + } + if (std::chrono::steady_clock::now() >= progress_deadline_) { + progress_active_.store(false, std::memory_order_release); + hub_->fail("MpiDirectTransport: progress command timed out"); + throw std::runtime_error("MpiDirectTransport: progress command timed out"); + } + auto result = hub_->poll_inbound(worker_id_, frame_type, sequence); + if (!result.has_value()) return false; + progress_active_.store(false, std::memory_order_release); + reply = std::move(*result); + return true; +} + +void MpiDirectTransport::shutdown() { + closed_.store(true, std::memory_order_release); + progress_active_.store(false, std::memory_order_release); + hub_->cancel_route(worker_id_, "transport shut down"); +} diff --git a/src/common/hierarchical/mpi_direct_transport.h b/src/common/hierarchical/mpi_direct_transport.h new file mode 100644 index 0000000000..f392ced230 --- /dev/null +++ b/src/common/hierarchical/mpi_direct_transport.h @@ -0,0 +1,117 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "remote_endpoint.h" + +enum class MpiDirectTag : int32_t { + COMMAND_REQUEST = 1, + COMMAND_REPLY = 2, + HEALTH = 3, + LIFECYCLE = 4, +}; + +struct MpiDirectOutboundFrame { + uint64_t ticket{0}; + int32_t target_rank{-1}; + MpiDirectTag tag{MpiDirectTag::COMMAND_REQUEST}; + std::vector frame; +}; + +class MpiDirectTransportHub { +public: + explicit MpiDirectTransportHub(size_t max_pending_frame_bytes); + + void register_route(int32_t worker_id, int32_t mpi_rank, uint64_t session_id, const std::string &comm_profile); + std::optional poll_outbound(double timeout_s); + void complete_outbound(uint64_t ticket); + void deliver(int32_t source_rank, MpiDirectTag tag, const std::vector &frame); + void fail(const std::string &message); + void close(); + + size_t pending_frame_bytes() const; + bool terminal() const; + std::string terminal_error() const; + +private: + friend class MpiDirectTransport; + + struct Route { + int32_t worker_id{-1}; + int32_t mpi_rank{-1}; + uint64_t session_id{0}; + std::string comm_profile; + std::deque> replies; + std::deque> lifecycle; + std::chrono::steady_clock::time_point last_health{}; + std::string terminal_error; + }; + + void cancel_route(int32_t worker_id, const std::string &message); + void enqueue(int32_t worker_id, MpiDirectTag tag, const std::vector &frame, double timeout_s); + std::optional> + poll_inbound(int32_t worker_id, remote_l3::FrameType frame_type, uint64_t sequence); + std::vector + wait_inbound(int32_t worker_id, remote_l3::FrameType frame_type, uint64_t sequence, double timeout_s); + void expect_hello_ready(int32_t worker_id, double timeout_s); + void throw_if_terminal_locked() const; + void throw_if_route_terminal_locked(const Route &route) const; + void fail_locked(const std::string &message); + + size_t max_pending_frame_bytes_{0}; + size_t pending_frame_bytes_{0}; + uint64_t next_ticket_{1}; + bool closed_{false}; + std::string terminal_error_; + mutable std::mutex mu_; + std::condition_variable cv_; + std::deque outbound_; + std::unordered_map in_flight_bytes_; + std::unordered_map routes_by_worker_; + std::unordered_map worker_by_rank_; +}; + +class MpiDirectTransport : public RemoteL3Transport { +public: + MpiDirectTransport( + std::shared_ptr hub, int32_t worker_id, double attach_timeout_s, double runtime_timeout_s + ); + + void expect_hello_ready(); + void submit_frame(const std::vector &frame) override; + std::vector wait_for_reply(remote_l3::FrameType frame_type, uint64_t sequence) override; + void submit_progress_frame(const std::vector &frame) override; + bool poll_progress_reply(remote_l3::FrameType frame_type, uint64_t sequence, std::vector &reply) override; + void shutdown() override; + +private: + std::shared_ptr hub_; + int32_t worker_id_{-1}; + double attach_timeout_s_{0.0}; + double runtime_timeout_s_{0.0}; + std::chrono::steady_clock::time_point progress_deadline_{}; + std::atomic progress_active_{false}; + std::atomic closed_{false}; +}; diff --git a/src/common/hierarchical/remote_endpoint.cpp b/src/common/hierarchical/remote_endpoint.cpp index 5002c2ac1c..26cefc5ae7 100644 --- a/src/common/hierarchical/remote_endpoint.cpp +++ b/src/common/hierarchical/remote_endpoint.cpp @@ -1383,7 +1383,12 @@ void RemoteL3Endpoint::request_progress_stop() noexcept { progress_stop_requested_ = true; progress_stop_reason_ = "RemoteL3Endpoint progress stopped"; command_cv_.notify_all(); - transport_->shutdown(); + // Worker::close() stops Scheduler progress before WorkerManager owns + // child shutdown. Keep an idle transport alive across that handoff so + // shutdown_child() can still submit the lifecycle SHUTDOWN frame. An + // in-flight task cannot be drained by normal shutdown and still needs + // transport cancellation to terminalize its progress lane. + if (pending_task_.occupied) transport_->shutdown(); } catch (...) {} } diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 46e08d1fde..77032507ff 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -21,6 +21,7 @@ #include #include "remote_endpoint.h" +#include "mpi_direct_transport.h" // --------------------------------------------------------------------------- // Fork hygiene @@ -161,6 +162,18 @@ void Worker::add_mpi_group_mailbox( } } +void Worker::add_remote_l3_mpi( + int32_t worker_id, uint64_t session_id, const std::string &transport_name, + const std::shared_ptr &hub, double attach_timeout_s, double runtime_timeout_s +) { + if (initialized_) throw std::runtime_error("Worker: add_remote_l3_mpi after init"); + auto transport = std::make_unique(hub, worker_id, attach_timeout_s, runtime_timeout_s); + transport->expect_hello_ready(); + manager_.add_next_level_endpoint( + std::make_unique(worker_id, session_id, transport_name, std::move(transport)) + ); +} + 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 629a4beda8..a60b422fe9 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -91,6 +91,10 @@ class Worker { const std::vector &worker_ids, const std::vector &session_ids, void *mailbox, size_t mailbox_bytes, int mpirun_pid, double runtime_timeout_s ); + void add_remote_l3_mpi( + int32_t worker_id, uint64_t session_id, const std::string &transport_name, + const std::shared_ptr &hub, double attach_timeout_s, 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.cpp b/src/common/hierarchical/worker_manager.cpp index 70be691d45..83e429ccd7 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -1107,6 +1107,14 @@ void WorkerManager::stop_workers() { void WorkerManager::stop() { stop_workers(); + // Stop admission on every endpoint before asking any child to exit. In + // particular, direct-MPI L3 ranks have no out-of-band session owner: their + // command loops leave only after this lifecycle SHUTDOWN reaches them. + // Notify every lane before destroying any endpoint or its transport. + for (auto &wt : next_level_threads_) + wt->shutdown_child(); + for (auto &wt : sub_threads_) + wt->shutdown_child(); next_level_threads_.clear(); sub_threads_.clear(); } diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index cacba117c9..54e8d8d952 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -209,6 +209,7 @@ add_library(hierarchical_objs OBJECT ${HIERARCHICAL_SRC_DIR}/scope.cpp ${HIERARCHICAL_SRC_DIR}/remote_wire.cpp ${HIERARCHICAL_SRC_DIR}/remote_endpoint.cpp + ${HIERARCHICAL_SRC_DIR}/mpi_direct_transport.cpp ${HIERARCHICAL_SRC_DIR}/orchestrator.cpp ${HIERARCHICAL_SRC_DIR}/worker_manager.cpp ${HIERARCHICAL_SRC_DIR}/scheduler.cpp @@ -410,6 +411,7 @@ add_hierarchical_test(test_scheduler hierarchical/test_scheduler.cpp CUSTOM_HOST add_hierarchical_test(test_chip_run_lane hierarchical/test_chip_run_lane.cpp) add_hierarchical_test(test_remote_wire hierarchical/test_remote_wire.cpp) add_hierarchical_test(test_remote_endpoint hierarchical/test_remote_endpoint.cpp) +add_hierarchical_test(test_mpi_direct_transport hierarchical/test_mpi_direct_transport.cpp) add_hierarchical_test(test_pipeline_contract hierarchical/test_pipeline_contract.cpp) # Run-stream pair state machine: publication-aware AICore stream reuse, diff --git a/tests/ut/cpp/hierarchical/test_mpi_direct_transport.cpp b/tests/ut/cpp/hierarchical/test_mpi_direct_transport.cpp new file mode 100644 index 0000000000..3d5582d72b --- /dev/null +++ b/tests/ut/cpp/hierarchical/test_mpi_direct_transport.cpp @@ -0,0 +1,237 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include "mpi_direct_transport.h" +#include "remote_wire.h" +#include "ring.h" +#include "worker_manager.h" + +namespace { + +constexpr int32_t WORKER_ID = 7; +constexpr int32_t MPI_RANK = 3; +constexpr uint64_t SESSION_ID = 0x1234; +constexpr size_t MAX_FRAME_BYTES = 40 + remote_l3::MAX_FRAME_PAYLOAD_BYTES; + +std::vector make_frame(remote_l3::FrameType type, uint64_t sequence, std::vector payload = {}) { + remote_l3::FrameHeader header; + header.frame_type = type; + header.session_id = SESSION_ID; + header.worker_id = WORKER_ID; + header.sequence = sequence; + return remote_l3::encode_frame(header, payload); +} + +std::vector make_hello() { + remote_l3::HelloPayload hello; + hello.session_id = SESSION_ID; + hello.worker_id = WORKER_ID; + hello.protocol_version = remote_l3::PROTOCOL_VERSION; + hello.comm_profile = "sim"; + hello.ready_state = remote_l3::ReadyState::READY; + return make_frame(remote_l3::FrameType::HELLO, 0, remote_l3::encode_hello(hello)); +} + +std::shared_ptr ready_hub() { + auto hub = std::make_shared(2 * MAX_FRAME_BYTES); + hub->register_route(WORKER_ID, MPI_RANK, SESSION_ID, "sim"); + hub->deliver(MPI_RANK, MpiDirectTag::LIFECYCLE, make_hello()); + return hub; +} + +TEST(MpiDirectTransportHub, RoutesRawSlr3FramesWithoutEnvelope) { + auto hub = ready_hub(); + MpiDirectTransport transport(hub, WORKER_ID, 1.0, 1.0); + transport.expect_hello_ready(); + + auto task = make_frame(remote_l3::FrameType::TASK, 11, {1, 2, 3}); + transport.submit_frame(task); + auto outbound = hub->poll_outbound(0.0); + ASSERT_TRUE(outbound.has_value()); + EXPECT_EQ(outbound->target_rank, MPI_RANK); + EXPECT_EQ(outbound->tag, MpiDirectTag::COMMAND_REQUEST); + EXPECT_EQ(outbound->frame, task); + hub->complete_outbound(outbound->ticket); + + auto completion = make_frame(remote_l3::FrameType::COMPLETION, 11, {4, 5}); + hub->deliver(MPI_RANK, MpiDirectTag::COMMAND_REPLY, completion); + EXPECT_EQ(transport.wait_for_reply(remote_l3::FrameType::COMPLETION, 11), completion); + EXPECT_EQ(hub->pending_frame_bytes(), 0U); +} + +TEST(MpiDirectTransport, ProgressApiIsNonBlocking) { + auto hub = ready_hub(); + MpiDirectTransport transport(hub, WORKER_ID, 1.0, 1.0); + transport.expect_hello_ready(); + + auto task = make_frame(remote_l3::FrameType::TASK, 21, {7}); + transport.submit_progress_frame(task); + EXPECT_THROW(transport.submit_frame(make_frame(remote_l3::FrameType::CONTROL, 22)), std::logic_error); + EXPECT_THROW(transport.wait_for_reply(remote_l3::FrameType::CONTROL_REPLY, 22), std::logic_error); + std::vector reply; + EXPECT_FALSE(transport.poll_progress_reply(remote_l3::FrameType::COMPLETION, 21, reply)); + + auto completion = make_frame(remote_l3::FrameType::COMPLETION, 21, {8}); + hub->deliver(MPI_RANK, MpiDirectTag::COMMAND_REPLY, completion); + ASSERT_TRUE(transport.poll_progress_reply(remote_l3::FrameType::COMPLETION, 21, reply)); + EXPECT_EQ(reply, completion); +} + +TEST(MpiDirectTransport, ShutdownWakesOnlyItsRouteWaiter) { + constexpr int32_t OTHER_WORKER_ID = WORKER_ID + 1; + constexpr int32_t OTHER_MPI_RANK = MPI_RANK + 1; + constexpr uint64_t OTHER_SESSION_ID = SESSION_ID + 1; + auto hub = ready_hub(); + hub->register_route(OTHER_WORKER_ID, OTHER_MPI_RANK, OTHER_SESSION_ID, "sim"); + MpiDirectTransport transport(hub, WORKER_ID, 1.0, 5.0); + MpiDirectTransport other(hub, OTHER_WORKER_ID, 1.0, 5.0); + + auto waiter = std::async(std::launch::async, [&] { + return transport.wait_for_reply(remote_l3::FrameType::CONTROL_REPLY, 41); + }); + auto other_waiter = std::async(std::launch::async, [&] { + return other.wait_for_reply(remote_l3::FrameType::CONTROL_REPLY, 42); + }); + EXPECT_EQ(waiter.wait_for(std::chrono::milliseconds(50)), std::future_status::timeout); + EXPECT_EQ(other_waiter.wait_for(std::chrono::milliseconds(50)), std::future_status::timeout); + + transport.shutdown(); + ASSERT_EQ(waiter.wait_for(std::chrono::milliseconds(500)), std::future_status::ready); + EXPECT_THROW((void)waiter.get(), std::runtime_error); + EXPECT_EQ(other_waiter.wait_for(std::chrono::milliseconds(50)), std::future_status::timeout); + EXPECT_FALSE(hub->terminal()); + + EXPECT_NO_THROW( + hub->deliver(MPI_RANK, MpiDirectTag::COMMAND_REPLY, make_frame(remote_l3::FrameType::CONTROL_REPLY, 41)) + ); + + remote_l3::FrameHeader header; + header.frame_type = remote_l3::FrameType::CONTROL_REPLY; + header.session_id = OTHER_SESSION_ID; + header.worker_id = OTHER_WORKER_ID; + header.sequence = 42; + auto other_reply = remote_l3::encode_frame(header, {}); + hub->deliver(OTHER_MPI_RANK, MpiDirectTag::COMMAND_REPLY, other_reply); + ASSERT_EQ(other_waiter.wait_for(std::chrono::milliseconds(500)), std::future_status::ready); + EXPECT_EQ(other_waiter.get(), other_reply); + + header.frame_type = remote_l3::FrameType::CONTROL; + header.sequence = 43; + other.submit_frame(remote_l3::encode_frame(header, {})); + auto outbound = hub->poll_outbound(0.0); + ASSERT_TRUE(outbound.has_value()); + EXPECT_EQ(outbound->target_rank, OTHER_MPI_RANK); + hub->complete_outbound(outbound->ticket); +} + +TEST(MpiDirectTransportHub, RejectsSourceRankMismatchAsTerminal) { + auto hub = std::make_shared(MAX_FRAME_BYTES); + hub->register_route(WORKER_ID, MPI_RANK, SESSION_ID, "sim"); + EXPECT_THROW(hub->deliver(MPI_RANK + 1, MpiDirectTag::LIFECYCLE, make_hello()), std::runtime_error); + EXPECT_TRUE(hub->terminal()); + EXPECT_NE(hub->terminal_error().find("unknown MPI rank"), std::string::npos); +} + +TEST(MpiDirectTransportHub, RejectsTagFrameTypeMismatchAsTerminal) { + auto hub = std::make_shared(MAX_FRAME_BYTES); + hub->register_route(WORKER_ID, MPI_RANK, SESSION_ID, "sim"); + EXPECT_THROW(hub->deliver(MPI_RANK, MpiDirectTag::COMMAND_REPLY, make_hello()), std::runtime_error); + EXPECT_TRUE(hub->terminal()); +} + +TEST(MpiDirectTransportHub, PendingByteCreditIncludesMpiInFlightSend) { + auto hub = std::make_shared(MAX_FRAME_BYTES); + hub->register_route(WORKER_ID, MPI_RANK, SESSION_ID, "sim"); + hub->deliver(MPI_RANK, MpiDirectTag::LIFECYCLE, make_hello()); + MpiDirectTransport transport(hub, WORKER_ID, 1.0, 2.0); + transport.expect_hello_ready(); + std::vector payload(remote_l3::MAX_FRAME_PAYLOAD_BYTES, 0x5a); + auto first = make_frame(remote_l3::FrameType::TASK, 1, payload); + auto second = make_frame(remote_l3::FrameType::TASK, 2, payload); + + transport.submit_frame(first); + auto outbound = hub->poll_outbound(0.0); + ASSERT_TRUE(outbound.has_value()); + + std::atomic submit_started{false}; + std::atomic submit_finished{false}; + std::exception_ptr submit_error; + std::thread submitter([&] { + submit_started.store(true, std::memory_order_release); + try { + transport.submit_frame(second); + submit_finished.store(true, std::memory_order_release); + } catch (...) { + submit_error = std::current_exception(); + } + }); + while (!submit_started.load(std::memory_order_acquire)) + std::this_thread::yield(); + EXPECT_FALSE(submit_finished.load(std::memory_order_acquire)); + hub->complete_outbound(outbound->ticket); + submitter.join(); + ASSERT_EQ(submit_error, nullptr); + EXPECT_TRUE(submit_finished.load(std::memory_order_acquire)); +} + +TEST(MpiDirectTransportHub, HealthFrameIsOutOfBand) { + auto hub = ready_hub(); + MpiDirectTransport transport(hub, WORKER_ID, 1.0, 1.0); + transport.expect_hello_ready(); + hub->deliver(MPI_RANK, MpiDirectTag::HEALTH, make_frame(remote_l3::FrameType::HEALTH, 9)); + + auto completion = make_frame(remote_l3::FrameType::COMPLETION, 10); + hub->deliver(MPI_RANK, MpiDirectTag::COMMAND_REPLY, completion); + EXPECT_EQ(transport.wait_for_reply(remote_l3::FrameType::COMPLETION, 10), completion); +} + +TEST(MpiDirectTransport, WorkerManagerStopSendsLifecycleShutdownAfterProgressStop) { + auto hub = ready_hub(); + auto transport = std::make_unique(hub, WORKER_ID, 1.0, 1.0); + transport->expect_hello_ready(); + + Ring ring; + ring.init(/*heap_bytes=*/0); + WorkerManager manager; + manager.add_next_level_endpoint( + std::make_unique(WORKER_ID, SESSION_ID, "mpi-direct", std::move(transport)) + ); + manager.start(&ring, [](WorkerCompletion) {}, [](WorkerDispatch) {}); + + // Worker::close() stops the Scheduler first. Its final progress pass sees + // the stopped lane before WorkerManager::stop() owns child shutdown. + manager.stop_workers(); + manager.progress(); + manager.stop(); + + auto outbound = hub->poll_outbound(0.0); + ASSERT_TRUE(outbound.has_value()); + EXPECT_EQ(outbound->target_rank, MPI_RANK); + EXPECT_EQ(outbound->tag, MpiDirectTag::LIFECYCLE); + auto frame = remote_l3::decode_frame(outbound->frame); + EXPECT_EQ(frame.header.frame_type, remote_l3::FrameType::SHUTDOWN); + EXPECT_EQ(frame.header.session_id, SESSION_ID); + EXPECT_EQ(frame.header.worker_id, WORKER_ID); + hub->complete_outbound(outbound->ticket); + ring.shutdown(); +} + +} // namespace diff --git a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp index e96a766b9b..c10c72bbb6 100644 --- a/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp +++ b/tests/ut/cpp/hierarchical/test_remote_endpoint.cpp @@ -317,9 +317,14 @@ class FakeRemoteTransport : public RemoteL3Transport { remote_l3::ControlName last_control_name{remote_l3::ControlName::PREPARE_CALLABLE}; remote_l3::RemoteRegistryTarget last_target_registry{remote_l3::RemoteRegistryTarget::REMOTE_TASK_DISPATCHER}; CallableKind last_callable_kind{CallableKind::PYTHON_IMPORT}; + int shutdown_calls{0}; - void submit_frame(const std::vector &frame) override { last_frame = frame; } + void submit_frame(const std::vector &frame) override { + if (shutdown_calls != 0) throw std::runtime_error("FakeRemoteTransport: submit after shutdown"); + last_frame = frame; + } void submit_progress_frame(const std::vector &frame) override { submit_frame(frame); } + void shutdown() override { ++shutdown_calls; } bool poll_progress_reply(remote_l3::FrameType frame_type, uint64_t sequence, std::vector &reply) override { if (progress_polls_before_ready > 0) { @@ -483,6 +488,19 @@ TEST(RemoteEndpoint, ProgressStopReleasesWaitingControl) { ring.shutdown(); } +TEST(RemoteEndpoint, IdleProgressStopDefersTransportShutdownUntilLifecycleFrame) { + auto *transport = new FakeRemoteTransport(); + RemoteL3Endpoint endpoint(3, 99, "fake", std::unique_ptr(transport)); + + endpoint.request_progress_stop(); + EXPECT_EQ(transport->shutdown_calls, 0); + + endpoint.shutdown_child(); + ASSERT_FALSE(transport->last_frame.empty()); + EXPECT_EQ(remote_l3::decode_frame(transport->last_frame).header.frame_type, remote_l3::FrameType::SHUTDOWN); + EXPECT_EQ(transport->shutdown_calls, 1); +} + TEST(RemoteEndpoint, RemoteTaskErrorMapsToTaskFailure) { Ring ring; ring.init(1ULL << 20); diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index 64809fde96..14bcb5283e 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -38,6 +38,7 @@ validate_hashid, ) from simpler.orchestrator import Orchestrator +from simpler.remote_l3_limits import MAX_FRAME_BYTES from simpler.remote_l3_protocol import ( CallableKind, ChipCallableBlobLocation, @@ -2863,6 +2864,56 @@ def remote_abort_register(self, worker_id, *args): assert digest in worker._uncertain_hashids +def test_initial_mpi_direct_register_rollback_aborts_only_uncommitted_workers(): + class FailingCommitWorker: + def __init__(self): + self.aborted = [] + self.unregistered = [] + + def remote_prepare_register(self, worker_id, *args): + return _FakeRemoteControlResult(worker_id) + + def remote_commit_register(self, worker_id, *args): + return _FakeRemoteControlResult(worker_id, ok=worker_id == 0, error_message="commit failed") + + def remote_abort_register(self, worker_id, *args): + self.aborted.append(worker_id) + return _FakeRemoteControlResult(worker_id) + + def remote_unregister(self, worker_id, *args): + self.unregistered.append(worker_id) + return _FakeRemoteControlResult(worker_id) + + worker = Worker(level=4, num_sub_workers=0) + hub = worker_mod._MpiDirectTransportHub(MAX_FRAME_BYTES) + for worker_id in range(2): + worker._add_mpi_direct_worker( + worker_mod._MpiDirectWorkerSpec( + worker_id=worker_id, + mpi_rank=worker_id + 1, + session_id=worker_id + 1, + host="localhost", + comm_profile="sim", + platform="a2a3sim", + runtime="sim", + device_ids=(), + global_device_ranks=(), + hub=hub, + attach_timeout_s=1.0, + runtime_timeout_s=1.0, + ) + ) + worker.register(RemoteCallable(_REMOTE_NOOP_ORCH_TARGET), workers=[0, 1]) + fake = FailingCommitWorker() + worker._worker = fake # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="commit failed"): + worker._publish_initial_mpi_direct_callables() + + assert fake.aborted == [1] + assert fake.unregistered == [0] + + def test_remote_unregister_exception_is_best_effort_and_marks_uncertain(): class FailingUnregisterWorker: def remote_unregister(self, *args): diff --git a/tests/ut/py/test_mpi_direct.py b/tests/ut/py/test_mpi_direct.py new file mode 100644 index 0000000000..34ff6806b6 --- /dev/null +++ b/tests/ut/py/test_mpi_direct.py @@ -0,0 +1,469 @@ +# 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 json +import socket +import threading +from pathlib import Path + +import pytest +import simpler.mpi_direct_runtime as runtime_mod +import simpler.mpi_direct_supervisor as supervisor_mod +from simpler import remote_l3_session +from simpler.mpi_direct_protocol import MPI_DIRECT_STARTUP_TOKEN_ENV, MpiDirectTag +from simpler.mpi_direct_runtime import _pre_mpi_gate +from simpler.mpi_direct_supervisor import ( + _EXPORTED_ENV_VARS, + _build_command, + _family_launcher_args, + _gate_recv, + _gate_send, + _host_slots, + _mpi_vendor_family, + _startup_gate, +) +from simpler.mpi_direct_topology import MpiDirectTopology, load_runtime_manifest, load_runtime_manifest_data + + +def _topology_dict(*, second_host: str = "host-a") -> dict: + return { + "controller_rank": 0, + "controller_host": "host-a", + "startup_timeout_s": 10, + "session_timeout_s": 5, + "heartbeat_interval_s": 1, + "launcher_args": [], + "executor_ranks": [ + { + "rank": 1, + "worker_id": 0, + "host": "host-a", + "platform": "a2a3sim", + "device_ids": [0], + "comm_profile": "sim", + }, + { + "rank": 2, + "worker_id": 1, + "host": second_host, + "platform": "a2a3sim", + "device_ids": [1], + "comm_profile": "sim", + }, + ], + } + + +def test_topology_requires_dense_rank_and_worker_mapping(): + data = _topology_dict() + data["executor_ranks"][1]["worker_id"] = 4 + with pytest.raises(ValueError, match="worker_ids must be dense"): + MpiDirectTopology.from_dict(data) + + data = _topology_dict() + data["executor_ranks"] = list(reversed(data["executor_ranks"])) + with pytest.raises(ValueError, match="ranks must be dense and ordered"): + MpiDirectTopology.from_dict(data) + + +@pytest.mark.parametrize("controller_host", ["localhost", "127.0.0.1", "::1"]) +def test_topology_rejects_loopback_controller_on_multiple_hosts(controller_host): + data = _topology_dict(second_host="host-b") + data["controller_host"] = controller_host + with pytest.raises(ValueError, match="controller_host=.*loopback.*topology spans multiple hosts"): + MpiDirectTopology.from_dict(data) + + +def test_topology_rejects_default_loopback_controller_on_multiple_hosts(): + data = _topology_dict(second_host="host-b") + data.pop("controller_host") + with pytest.raises(ValueError, match="controller_host=.*loopback.*topology spans multiple hosts"): + MpiDirectTopology.from_dict(data) + + +def test_topology_allows_loopback_controller_on_one_host(): + data = _topology_dict(second_host="localhost") + data["controller_host"] = "localhost" + data["executor_ranks"][0]["host"] = "localhost" + MpiDirectTopology.from_dict(data) + + +def test_topology_rejects_noncontiguous_host_order(): + data = _topology_dict(second_host="host-a") + data["executor_ranks"][0]["host"] = "host-b" + with pytest.raises(ValueError, match="hosts must be contiguous"): + MpiDirectTopology.from_dict(data) + + +def test_runtime_manifest_round_trip_binds_one_session(tmp_path: Path): + topology = MpiDirectTopology.from_dict(_topology_dict()) + path = tmp_path / "manifest.json" + path.write_text(json.dumps(topology.runtime_manifest(91)), encoding="utf-8") + loaded, session_id = load_runtime_manifest(str(path)) + assert loaded == topology + assert session_id == 91 + assert loaded.executor_for_rank(2).worker_id == 1 + assert loaded.executor_for_rank(1).global_device_ranks == (0,) + assert loaded.executor_for_rank(2).global_device_ranks == (1,) + + +def test_runtime_manifest_can_be_consumed_without_a_shared_file(): + topology = MpiDirectTopology.from_dict(_topology_dict()) + loaded, session_id = load_runtime_manifest_data(topology.runtime_manifest(92)) + assert loaded == topology + assert session_id == 92 + + +def test_topology_rejects_duplicate_explicit_global_device_ranks(): + data = _topology_dict() + data["executor_ranks"][0]["global_device_ranks"] = [7] + data["executor_ranks"][1]["global_device_ranks"] = [7] + with pytest.raises(ValueError, match="global_device_ranks must be unique"): + MpiDirectTopology.from_dict(data) + + +def test_supervisor_openmpi_command_uses_one_static_world_without_shell(): + topology = MpiDirectTopology.from_dict(_topology_dict(second_host="host-b")) + command = _build_command( + topology, + mpirun_path="/opt/mpi/bin/mpirun", + topology_path="/work/simpler/topology.json", + session_id=91, + controller="case.py", + launcher_family="openmpi", + ) + assert command[:4] == ["/opt/mpi/bin/mpirun", "--host", "host-a:2,host-b:1", "--map-by"] + assert command[-11:-9] == ["-np", "3"] + assert command[-6:] == [ + "--topology", + "/work/simpler/topology.json", + "--session-id", + "91", + "--controller", + "case.py", + ] + + +def test_supervisor_mpich_command_uses_local_hostfile_without_openmpi_flags(monkeypatch): + for name in _EXPORTED_ENV_VARS: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("PYTHONPATH", "/work/simpler") + monkeypatch.setenv("ASCEND_HOME_PATH", "/opt/ascend") + topology = MpiDirectTopology.from_dict(_topology_dict(second_host="host-b")) + assert _host_slots(topology) == (("host-a", 2), ("host-b", 1)) + assert _family_launcher_args(topology, "mpich", "/tmp/hosts") == [ + "-f", + "/tmp/hosts", + "-genv", + "PYTHONPATH", + "/work/simpler", + "-genv", + "ASCEND_HOME_PATH", + "/opt/ascend", + ] + command = _build_command( + topology, + mpirun_path="/opt/mpi/bin/mpiexec", + topology_path="/work/simpler/topology.json", + session_id=91, + controller="case.py", + launcher_family="mpich", + hostfile_path="/tmp/hosts", + ) + assert command[:3] == ["/opt/mpi/bin/mpiexec", "-f", "/tmp/hosts"] + assert command[command.index("-np") : command.index("-np") + 2] == ["-np", "3"] + assert "-x" not in command + assert "--map-by" not in command + + +def test_run_supervisor_keeps_mpich_hostfile_until_job_finishes(tmp_path, monkeypatch): + topology_path = tmp_path / "topology.json" + topology_path.write_text(json.dumps(_topology_dict(second_host="host-b")), encoding="utf-8") + observed: dict[str, str] = {} + + class FakeListener: + def setsockopt(self, *_args): + pass + + def bind(self, _address): + pass + + def listen(self, _backlog): + pass + + def getsockname(self): + return ("host-a", 4567) + + def close(self): + pass + + class FakeProcess: + returncode = None + pid = 1 + + def poll(self): + return None + + def wait(self, *_args, **_kwargs): + assert Path(observed["hostfile"]).is_file() + return 0 + + def popen(command, **_kwargs): + hostfile = command[command.index("-f") + 1] + observed["hostfile"] = hostfile + assert Path(hostfile).is_file() + return FakeProcess() + + def startup_gate(_topology, _token, _listener, _proc): + assert Path(observed["hostfile"]).is_file() + + monkeypatch.setattr(supervisor_mod, "_detect_mpi4py_family", lambda _python: ("mpich", "MPICH")) + monkeypatch.setattr(supervisor_mod.socket, "socket", lambda *_args, **_kwargs: FakeListener()) + monkeypatch.setattr(supervisor_mod.subprocess, "Popen", popen) + monkeypatch.setattr(supervisor_mod, "_startup_gate", startup_gate) + + assert supervisor_mod.run_supervisor(str(topology_path), "controller.py", launcher_family="mpich") == 0 + assert not Path(observed["hostfile"]).exists() + + +def test_supervisor_inline_manifest_command_adds_pre_mpi_gate(): + topology = MpiDirectTopology.from_dict(_topology_dict()) + command = _build_command( + topology, + mpirun_path="mpiexec", + topology_path=None, + session_id=91, + controller="case.py", + launcher_family="mpich", + hostfile_path="/tmp/hosts", + manifest_json="encoded-manifest", + python_executable="/opt/simpler-python", + startup_host="host-a", + startup_port=4567, + ) + assert "/opt/simpler-python" in command + assert "--manifest-json" in command + assert "--topology" not in command + assert command[-4:] == ["--startup-host", "host-a", "--startup-port", "4567"] + assert "token" not in command + + +def test_startup_token_is_exported_without_appearing_in_launcher_args(monkeypatch): + monkeypatch.setenv(MPI_DIRECT_STARTUP_TOKEN_ENV, "secret-token") + topology = MpiDirectTopology.from_dict(_topology_dict(second_host="host-b")) + + mpich_args = _family_launcher_args(topology, "mpich", "/tmp/hosts") + assert ["-genvlist", MPI_DIRECT_STARTUP_TOKEN_ENV] == mpich_args[-2:] + assert "secret-token" not in mpich_args + + openmpi_args = _family_launcher_args(topology, "openmpi", None) + assert ["-x", MPI_DIRECT_STARTUP_TOKEN_ENV] == openmpi_args[-2:] + assert "secret-token" not in openmpi_args + + +def test_pre_mpi_gate_releases_all_ready_ranks(): + topology = MpiDirectTopology.from_dict(_topology_dict(second_host="host-b")) + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen(topology.world_size) + port = int(listener.getsockname()[1]) + errors = [] + + def enter(rank): + try: + _pre_mpi_gate("127.0.0.1", port, "token", rank, 3.0) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=enter, args=(rank,)) for rank in range(topology.world_size)] + for thread in threads: + thread.start() + + class RunningProcess: + returncode = None + + @staticmethod + def poll(): + return None + + _startup_gate(topology, "token", listener, RunningProcess()) # type: ignore[arg-type] + for thread in threads: + thread.join(3.0) + assert not errors + assert all(not thread.is_alive() for thread in threads) + + +def test_pre_mpi_gate_sleeps_between_connection_retries(monkeypatch): + now = [0.0] + sleeps = [] + + def monotonic(): + return now[0] + + def sleep(delay): + sleeps.append(delay) + now[0] = 1.0 + + def refused_connection(*_args, **_kwargs): + raise ConnectionRefusedError("gate is not listening") + + monkeypatch.setattr(runtime_mod.time, "monotonic", monotonic) + monkeypatch.setattr(runtime_mod.time, "sleep", sleep) + monkeypatch.setattr(runtime_mod.socket, "create_connection", refused_connection) + + with pytest.raises(TimeoutError, match="startup gate connection timed out"): + _pre_mpi_gate("controller", 4567, "token", 0, 0.1) + + assert sleeps == [runtime_mod._MPI_GATE_RETRY_INTERVAL_S] + + +def test_startup_gate_ignores_invalid_peer_and_releases_ready_ranks(): + topology = MpiDirectTopology.from_dict(_topology_dict(second_host="host-b")) + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen(topology.world_size) + port = int(listener.getsockname()[1]) + responses = [] + + bad_peer = socket.create_connection(("127.0.0.1", port), timeout=1.0) + _gate_send(bad_peer, {"token": "wrong", "rank": 0, "state": "ready"}) + bad_peer.close() + + def enter(rank): + with socket.create_connection(("127.0.0.1", port), timeout=1.0) as peer: + _gate_send(peer, {"token": "token", "rank": rank, "state": "ready"}) + responses.append(_gate_recv(peer)) + + threads = [threading.Thread(target=enter, args=(rank,)) for rank in range(topology.world_size)] + for thread in threads: + thread.start() + + class RunningProcess: + returncode = None + + @staticmethod + def poll(): + return None + + _startup_gate(topology, "token", listener, RunningProcess()) # type: ignore[arg-type] + for thread in threads: + thread.join(3.0) + + assert responses == [{"token": "token", "state": "go_mpi"}] * topology.world_size + assert all(not thread.is_alive() for thread in threads) + + +def test_executor_cleanup_closes_worker_when_channel_close_fails(monkeypatch): + topology = MpiDirectTopology.from_dict(_topology_dict()) + channel_closed = False + + class FailingChannel: + def close(self): + nonlocal channel_closed + channel_closed = True + raise RuntimeError("heartbeat failed") + + class FakeWorker: + closed = False + + def close(self): + self.closed = True + + channel = FailingChannel() + worker = FakeWorker() + monkeypatch.setattr(runtime_mod, "_ExecutorFrameSocket", lambda *_args: channel) + monkeypatch.setattr(remote_l3_session, "_run_command_loop", lambda *_args: None) + + with pytest.raises(RuntimeError, match="heartbeat failed"): + runtime_mod._run_executor(None, None, topology, 1, topology.executors[0], worker) + + assert channel_closed + assert worker.closed + + +def test_executor_receive_sleeps_between_empty_mpi_probes(monkeypatch): + topology = MpiDirectTopology.from_dict(_topology_dict()) + + class FakeStatus: + pass + + class FakeMPI: + ANY_TAG = -1 + BYTE = object() + + @staticmethod + def Status(): + return FakeStatus() + + class EmptyWorld: + def improbe(self, **_kwargs): + return None + + channel = runtime_mod._ExecutorFrameSocket(FakeMPI, EmptyWorld(), topology.executors[0], 1, 1.0) + sleeps = [] + + def sleep(delay): + sleeps.append(delay) + channel._health_error = RuntimeError("stop probe") + + monkeypatch.setattr(runtime_mod.time, "sleep", sleep) + with pytest.raises(RuntimeError, match="executor heartbeat failed"): + channel._receive_message() + + assert sleeps == [runtime_mod._MPI_POLL_INTERVAL_S] + + +def test_controller_progress_sleeps_when_idle(monkeypatch): + class FakeStatus: + pass + + class FakeMPI: + ANY_SOURCE = -1 + ANY_TAG = -1 + BYTE = object() + + @staticmethod + def Status(): + return FakeStatus() + + class EmptyWorld: + def improbe(self, **_kwargs): + return None + + class EmptyHub: + pending_frame_bytes = 0 + + def poll_outbound(self, _timeout): + return None + + progress = runtime_mod._ControllerProgress(FakeMPI, EmptyWorld(), EmptyHub()) + sleeps = [] + + def sleep(delay): + sleeps.append(delay) + progress._stop.set() + + monkeypatch.setattr(runtime_mod.time, "sleep", sleep) + progress._run() + + assert sleeps == [runtime_mod._MPI_POLL_INTERVAL_S] + + +def test_mpi_vendor_family_matches_supported_launchers(): + assert _mpi_vendor_family("Open MPI") == "openmpi" + assert _mpi_vendor_family("MPICH") == "mpich" + assert _mpi_vendor_family("MVAPICH2") == "mpich" + with pytest.raises(RuntimeError, match="unsupported mpi4py MPI vendor"): + _mpi_vendor_family("unknown") + + +def test_direct_mpi_tags_are_small_fixed_lanes(): + assert tuple(int(tag) for tag in MpiDirectTag) == (1, 2, 3, 4) diff --git a/tests/ut/py/test_worker/test_comm_endpoints.py b/tests/ut/py/test_worker/test_comm_endpoints.py index 44636b9140..54c8efc500 100644 --- a/tests/ut/py/test_worker/test_comm_endpoints.py +++ b/tests/ut/py/test_worker/test_comm_endpoints.py @@ -13,7 +13,7 @@ import pytest from simpler import comm_endpoints as ce from simpler.buffer import BackendKind as BufferBackendKind -from simpler.worker import RemoteWorkerSpec, Worker, _Lifecycle +from simpler.worker import RemoteWorkerSpec, Worker, _Lifecycle, _MpiDirectWorkerSpec def _ready(worker: Worker) -> Worker: @@ -39,6 +39,28 @@ def _l4_with_remote(*specs: RemoteWorkerSpec) -> Worker: return _ready(worker) +def _l4_with_mpi_direct(*, hosts: tuple[str, ...], device_ids_by_rank: tuple[tuple[int, ...], ...]) -> Worker: + worker = Worker(level=4, num_sub_workers=0) + for worker_id, (host, device_ids) in enumerate(zip(hosts, device_ids_by_rank)): + worker._add_mpi_direct_worker( + _MpiDirectWorkerSpec( + worker_id=worker_id, + mpi_rank=worker_id + 1, + session_id=worker_id + 1, + host=host, + comm_profile="sim", + platform="a2a3sim", + runtime="sim", + device_ids=device_ids, + global_device_ranks=(), + hub=object(), # type: ignore[arg-type] + attach_timeout_s=1.0, + runtime_timeout_s=1.0, + ) + ) + return _ready(worker) + + def _record(worker: Worker, path: str, deployment: ce.EndpointDeployment) -> ce.EndpointRecord: return worker._resolve_region_spec([ce.at(path, deployment)], ce.SingleOwner()).members[0] @@ -213,6 +235,24 @@ def test_remote_registry_normalizes_node_identity_by_host_not_remote_status(): assert not registry.same_node(remote_a, remote_c) +def test_mpi_direct_registry_registers_rank_hosts_and_devices(): + worker = _l4_with_mpi_direct( + hosts=("localhost", "10.0.0.7", "10.0.0.7"), + device_ids_by_rank=((12, 13), (14,), (15,)), + ) + registry = worker._get_endpoint_registry() + root = _record(worker, "L4", ce.HOST_CPU) + local_device = _record(worker, "L4/L3[0]/L2[1]", ce.DEVICE_AICORE) + remote_a = _record(worker, "L4/L3[1]", ce.HOST_CPU) + remote_device = _record(worker, "L4/L3[1]/L2[0]", ce.DEVICE_AICORE) + remote_b = _record(worker, "L4/L3[2]", ce.HOST_CPU) + + assert registry.same_node(root, local_device) + assert registry.same_node(remote_a, remote_device) + assert registry.same_node(remote_a, remote_b) + assert not registry.same_node(root, remote_a) + + def test_at_missing_path_reports_path_not_found(): worker = _l3(device_ids=[0]) with pytest.raises(ce.EndpointResolveError) as excinfo: diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 109996d4b2..ed361605f9 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -60,6 +60,7 @@ RunHandle, Worker, _buffer_field_addr, + _close_fork_child_fds, _mailbox_addr, _mailbox_load_i32, _mailbox_store_i32, @@ -86,6 +87,15 @@ def _read_counter(buf) -> int: return struct.unpack_from("i", buf, 0)[0] +def test_close_fork_child_fds_ignores_invalid_entries(monkeypatch): + closed: list[int] = [] + monkeypatch.setattr(worker_mod.os, "close", closed.append) + + _close_fork_child_fds(["invalid", 2, 7, object()]) + + assert closed == [7] + + def _increment_counter(buf) -> None: v = struct.unpack_from("i", buf, 0)[0] struct.pack_into("i", buf, 0, v + 1) diff --git a/tools/verify_packaging.sh b/tools/verify_packaging.sh index 66331aa454..ebccce392d 100755 --- a/tools/verify_packaging.sh +++ b/tools/verify_packaging.sh @@ -59,6 +59,7 @@ from simpler.task_interface import ChipWorker from simpler.orchestrator import Orchestrator from simpler_setup.runtime_builder import RuntimeBuilder from simpler_setup.runtime_compiler import RuntimeCompiler +from simpler_setup.environment import PROJECT_ROOT from simpler_setup.kernel_compiler import KernelCompiler from simpler_setup.elf_parser import extract_text_section from simpler_setup.platform_info import parse_platform, discover_runtimes @@ -74,6 +75,11 @@ for rel in ('pipe_sync.h', os.path.join('common', 'dma_workspace.h')): assert any(os.path.isfile(os.path.join(d, rel)) for d in inc_dirs), \ 'incore helper not shipped: ' + rel + '; include dirs: ' + repr(inc_dirs) print('incore helpers OK:', inc_dirs) +# RuntimeCompiler passes this directory to every host-side CMake configure. +for name in ('host_log_sources.cmake', 'profiling_config.cmake', 'sanitizers.cmake'): + path = PROJECT_ROOT / 'cmake' / name + assert path.is_file(), 'shared CMake module not shipped: ' + str(path) +print('shared CMake modules OK:', PROJECT_ROOT / 'cmake') " ) echo "::endgroup::"