diff --git a/docs/components/core.md b/docs/components/core.md index 5172573e..980e5e2c 100644 --- a/docs/components/core.md +++ b/docs/components/core.md @@ -1,7 +1,7 @@ # Core -`cuphoton.core` owns the command-line, context, path, logging, and invariant -framework shared by every cuPhoton component. Component algorithms, datasets, +`cuphoton.core` owns the command-line, context, path, logging, invariant, +and shared Dragon/MPI execution framework. Component algorithms, datasets, scientific validation, and output formatting remain in their owning `cuphoton.*` namespaces. @@ -15,13 +15,18 @@ scientific validation, and output formatting remain in their owning - side-effect-free resolution of component XDG config, state, data, run, and log paths. -The fixed public surface has six groups, 89 domain commands, 86 accepted -command aliases, and 754 declared arguments. Five groups also support a -component-level `version` command, for 94 commands when those built-ins are -included; xDataReader is the exception. +The public surface is pinned by the CLI contract tests. Five groups also +support a component-level `version` command; xDataReader is the exception. Workflow-specific YAML `--config` options belong to each component. +Shared executor rounds measure `batch_wall_sec` through receipt of all worker +completions. Coordinator artifact audits and scientific output merging follow +that interval; `finalization_sec` measures the component's merge separately. +Inputs retained by an adapter are loaded during worker setup. A round summary +records a completed pass, while the root terminal summary is published after +worker shutdown and the executor's lifecycle checks. + ## Public facade Import shared CLI infrastructure from `cuphoton.core.cli`. Its stable facade diff --git a/src/cuphoton/core/_mpi_runtime.py b/src/cuphoton/core/_mpi_runtime.py new file mode 100644 index 00000000..9ae882ba --- /dev/null +++ b/src/cuphoton/core/_mpi_runtime.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Optional MPI startup and launcher contracts for component adapters.""" + +from __future__ import annotations + +import re +import socket +from collections.abc import Callable, Mapping, Sequence +from dataclasses import asdict, dataclass +from importlib.metadata import PackageNotFoundError, version +from typing import Any + +from .bulk import error_payload + +_ALLOCATED_DEVICES = "CUPHOTON_ALLOCATED_CUDA_VISIBLE_DEVICES" +_UNSIGNED = re.compile(r"[0-9]{1,18}") + + +@dataclass(frozen=True) +class _RankContext: + rank: int + local_rank: int + world_size: int + host: str + launcher: str + launch_id: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class _MPIAPI: + MPI: Any + comm: Any + mpi4py_version: str | None + library_version: str | None + + +def _mpi_failure_consensus( + comm: Any, + phase: str, + error: Mapping[str, Any] | None, + *, + root_failure_handler: Callable[[Mapping[str, str]], None] | None = None, +) -> None: + rank = int(comm.Get_rank()) + size = int(comm.Get_size()) + gathered = comm.gather( + {"rank": rank, "error": dict(error) if error else None}, + root=0, + ) + decision = None + if rank == 0: + failures: list[str] = [] + if not isinstance(gathered, Sequence) or len(gathered) != size: + failures.append("invalid consensus result") + else: + for expected_rank, item in enumerate(gathered): + if not isinstance(item, Mapping): + failures.append( + f"rank {expected_rank}: invalid consensus result" + ) + continue + if _integer(item.get("rank")) != expected_rank: + failures.append( + f"rank {expected_rank}: invalid rank identity" + ) + item_error = item.get("error") + if item_error is None: + continue + if not _valid_error(item_error): + failures.append( + f"rank {expected_rank}: invalid error payload" + ) + continue + failures.append( + f"rank {expected_rank}: {item_error['type']}: " + f"{item_error['message']}" + ) + message = ( + f"{phase} failed: " + "; ".join(failures) if failures else None + ) + if message is not None and root_failure_handler is not None: + try: + root_failure_handler( + {"type": "RuntimeError", "message": message} + ) + except Exception as exc: + persistence_error = error_payload(exc) + message += ( + "; failed to persist terminal setup evidence: " + f"{persistence_error['type']}: " + f"{persistence_error['message']}" + ) + decision = {"error": message} + decision = comm.bcast(decision, root=0) + if not isinstance(decision, Mapping): + raise RuntimeError(f"{phase} decision was not a mapping") + if decision.get("error"): + raise RuntimeError(str(decision["error"])) + + +def _valid_error(value: Any) -> bool: + return isinstance(value, Mapping) and all( + isinstance(value.get(field), str) and bool(value[field]) + for field in ("type", "message") + ) + + +def _mpi_context( + mpi: Any, comm: Any, environ: Mapping[str, str] +) -> _RankContext: + rank = _runtime_integer(comm.Get_rank(), "MPI rank") + size = _runtime_integer(comm.Get_size(), "MPI world size") + _validate_topology(rank, size) + context = None + failure = None + try: + try: + shared = comm.Split_type(mpi.COMM_TYPE_SHARED, key=rank) + try: + local_rank = _runtime_integer( + shared.Get_rank(), "MPI local rank" + ) + finally: + shared.Free() + except Exception as exc: + raise RuntimeError("cannot determine MPI local rank") from exc + launcher = _launcher_context(environ) + if launcher is not None and launcher[1:] != (rank, size): + raise RuntimeError("launcher topology disagrees with mpi4py") + launcher_local_rank = _local_rank( + environ, launcher[0] if launcher is not None else None + ) + if ( + launcher_local_rank is not None + and launcher_local_rank != local_rank + ): + raise RuntimeError( + "launcher local rank disagrees with MPI shared communicator" + ) + context = _RankContext( + rank, + local_rank, + size, + socket.gethostname(), + launcher[0] if launcher is not None else "mpi", + ) + except Exception as exc: + failure = error_payload(exc) + if failure is None and context is None: + failure = error_payload( + RuntimeError("rank context validation produced no context") + ) + _mpi_failure_consensus(comm, "MPI rank-context validation", failure) + if context is None: + raise RuntimeError("MPI rank-context consensus returned no context") + return context + + +def _launcher_context( + environ: Mapping[str, str], +) -> tuple[str, int, int] | None: + candidates = ( + ("openmpi", "OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("pmi", "PMI_RANK", "PMI_SIZE"), + ("slurm", "SLURM_PROCID", "SLURM_NTASKS"), + ) + for name, rank_key, size_key in candidates: + present = rank_key in environ, size_key in environ + if all(present): + return ( + name, + _environment_integer(environ, rank_key), + _environment_integer(environ, size_key), + ) + if any(present): + raise RuntimeError( + f"{name} launcher rank and size must appear together" + ) + return None + + +def _local_rank( + environ: Mapping[str, str], launcher: str | None = None +) -> int | None: + keys_by_launcher = { + "openmpi": ("OMPI_COMM_WORLD_LOCAL_RANK",), + "pmi": ("PMI_LOCAL_RANK", "MPI_LOCALRANKID", "PMIX_LOCAL_RANK"), + "slurm": ( + "SLURM_LOCALID", + "PMIX_LOCAL_RANK", + "MPI_LOCALRANKID", + "PMI_LOCAL_RANK", + ), + } + keys = keys_by_launcher.get( + launcher, + ( + "OMPI_COMM_WORLD_LOCAL_RANK", + "PMIX_LOCAL_RANK", + "MPI_LOCALRANKID", + "PMI_LOCAL_RANK", + "SLURM_LOCALID", + ), + ) + values = [ + (key, _environment_integer(environ, key)) + for key in keys + if key in environ + ] + if launcher == "pmi" and not values and "SLURM_LOCALID" in environ: + slurm_keys = ("SLURM_PROCID", "SLURM_NTASKS", "SLURM_LOCALID") + topology_keys = ("PMI_RANK", "PMI_SIZE", *slurm_keys) + if all(key in environ for key in topology_keys): + pmi_topology = ( + _environment_integer(environ, "PMI_RANK"), + _environment_integer(environ, "PMI_SIZE"), + ) + slurm_topology = ( + _environment_integer(environ, "SLURM_PROCID"), + _environment_integer(environ, "SLURM_NTASKS"), + ) + if pmi_topology != slurm_topology: + raise RuntimeError( + "PMI and Slurm launcher rank topology disagree" + ) + values = [ + ( + "SLURM_LOCALID", + _environment_integer(environ, "SLURM_LOCALID"), + ) + ] + if not values: + return None + if any(value != values[0][1] for _, value in values[1:]): + family = launcher or "MPI" + raise RuntimeError(f"{family} launcher local-rank variables disagree") + return values[0][1] + + +def _prebound_visibility(environ: Mapping[str, str]) -> str: + raw_visibility = environ.get("CUDA_VISIBLE_DEVICES") + if raw_visibility is None: + raise RuntimeError( + "CUDA_VISIBLE_DEVICES is unset; bind one device per MPI rank " + "with cuphoton-openmpi-rank-exec or scheduler GPU binding" + ) + tokens = _visible_tokens(raw_visibility) + if len(tokens) != 1: + raise RuntimeError( + "MPI rank must see one CUDA device before mpi4py import; use " + "cuphoton-openmpi-rank-exec or scheduler GPU binding" + ) + allocated_raw = environ.get(_ALLOCATED_DEVICES) + if allocated_raw is None: + return tokens[0] + allocated = _visible_tokens(allocated_raw) + local_rank = _local_rank(environ, "openmpi") + if local_rank is None or local_rank >= len(allocated): + raise RuntimeError( + f"{_ALLOCATED_DEVICES} requires a valid Open MPI local rank; " + "unset it when launching without cuphoton-openmpi-rank-exec" + ) + if allocated[local_rank] != tokens[0]: + raise RuntimeError("pre-bound CUDA device does not match allocation") + return tokens[0] + + +def _visible_tokens(raw: Any) -> tuple[str, ...]: + if not isinstance(raw, str) or not raw.strip() or raw.strip() == "-1": + raise RuntimeError("CUDA_VISIBLE_DEVICES exposes no devices") + tokens = tuple(token.strip() for token in raw.split(",")) + if any(not token or token.startswith("-") for token in tokens) or len( + tokens + ) != len(set(tokens)): + raise RuntimeError("CUDA_VISIBLE_DEVICES is invalid") + return tokens + + +def _environment_integer(environ: Mapping[str, str], key: str) -> int: + raw = environ[key] + if not isinstance(raw, str) or not _UNSIGNED.fullmatch(raw): + raise RuntimeError(f"launcher variable {key} must be unsigned") + return int(raw) + + +def _runtime_integer(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise RuntimeError(f"{field} must be non-negative") + return value + + +def _validate_topology(rank: int, size: int) -> None: + if size <= 0 or rank >= size: + raise RuntimeError("invalid launcher rank topology") + + +def _integer(value: Any) -> int | None: + return ( + value + if isinstance(value, int) and not isinstance(value, bool) + else None + ) + + +def _load_mpi_api() -> _MPIAPI: + try: + from mpi4py import MPI + except (ImportError, OSError, RuntimeError) as exc: + raise RuntimeError( + "MPI collective aggregation requires mpi4py built against the " + "allocation's MPI runtime" + ) from exc + try: + library_version = _normalize_mpi_library_version( + MPI.Get_library_version() + ) + except Exception: # pragma: no cover - implementation-specific + library_version = None + return _MPIAPI( + MPI, + MPI.COMM_WORLD, + _distribution_version("mpi4py"), + library_version, + ) + + +def _normalize_mpi_library_version(value: Any) -> str | None: + if value is None: + return None + return str(value).rstrip("\x00 \t\r\n") + + +def _distribution_version(name: str) -> str | None: + try: + return version(name) + except PackageNotFoundError: + return None diff --git a/src/cuphoton/core/benchmark.py b/src/cuphoton/core/benchmark.py index 7865cdeb..a335736e 100644 --- a/src/cuphoton/core/benchmark.py +++ b/src/cuphoton/core/benchmark.py @@ -149,9 +149,9 @@ def build_benchmark_report( "timing_definitions": { "batch_wall_sec": ( "Coordinator time from before round release through receipt " - "of all worker completions, including ordinary input reads, " - "numerical work, output writes and record publication; " - "excludes subsequent coordinator artifact audits." + "of all worker completions, including any per-round input " + "reads, computation, output writes and record publication. " + "Coordinator artifact audits and finalization are excluded." ), "worker_wall_max_sec": ( "Maximum worker-local elapsed time for the round. Its " diff --git a/src/cuphoton/core/bulk.py b/src/cuphoton/core/bulk.py index e433324b..4ac59abc 100644 --- a/src/cuphoton/core/bulk.py +++ b/src/cuphoton/core/bulk.py @@ -440,3 +440,30 @@ def item_ids_sha256(items: Sequence[WorkItem]) -> str: return hashlib.sha256( "\n".join(item.item_id for item in items).encode() ).hexdigest() + + +def hostnames_match( + requested: str, + actual: str, + *, + allow_loopback_alias: bool = False, +) -> bool: + """Accept exact or equivalent short/FQDN scheduler hostnames.""" + + requested_normalized = requested.rstrip(".").lower() + actual_normalized = actual.rstrip(".").lower() + if allow_loopback_alias and requested_normalized in { + "localhost", + "localhost.localdomain", + }: + # Dragon 0.14.1 reports ``localhost`` for its single-node system + # descriptor even though socket.gethostname() exposes the machine + # hostname inside the launched worker. + return bool(actual_normalized) + if requested_normalized == actual_normalized: + return True + if "." not in requested_normalized: + return actual_normalized.startswith(requested_normalized + ".") + if "." not in actual_normalized: + return requested_normalized.startswith(actual_normalized + ".") + return False diff --git a/src/cuphoton/core/cli/executor.py b/src/cuphoton/core/cli/executor.py new file mode 100644 index 00000000..fc69269b --- /dev/null +++ b/src/cuphoton/core/cli/executor.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Common opt-in distributed execution options for component commands.""" + +from __future__ import annotations + +from ..benchmark import BenchmarkOptions +from .command import CommandError +from .invariants import ( + FloatInvariant, + NonNegativeIntegerInvariant, + PositiveIntegerInvariant, + SetInvariant, +) + + +class ExecutorOptions: + """Keep local command defaults and validate runtime-specific options.""" + + executor = None + max_workers = None + worker_timeout_sec = None + result_timeout_sec = None + rank_setup_timeout_sec = None + warmup_rounds = None + measure_rounds = None + + class ExecutorArg(SetInvariant): + _arg = "--executor" + _help = ( + "Execution runtime: local, dragon, or mpi. [default: %default]" + ) + _set = {"local", "dragon", "mpi"} + _default = "local" + + class MaxWorkersArg(PositiveIntegerInvariant): + _arg = "--max-workers" + _help = "Dragon-only maximum number of GPU workers." + _default = None + + class WorkerTimeoutSecArg(FloatInvariant): + _arg = "--worker-timeout-sec" + _help = "Dragon worker lifetime across all rounds. [default: 3600]" + _default = None + _min = 0.001 + + class ResultTimeoutSecArg(FloatInvariant): + _arg = "--result-timeout-sec" + _help = "Dragon result and artifact grace period. [default: 60]" + _default = None + _min = 0.001 + + class RankSetupTimeoutSecArg(FloatInvariant): + _arg = "--rank-setup-timeout-sec" + _help = "MPI shared-artifact visibility timeout. [default: 600]" + _default = None + _min = 0.001 + + class WarmupRoundsArg(NonNegativeIntegerInvariant): + _arg = "--warmup-rounds" + _help = ( + "Opt into the benchmark round layout with this many warmup " + "passes; outputs are retained. [default when benchmarking: 0]" + ) + _default = None + + class MeasureRoundsArg(PositiveIntegerInvariant): + _arg = "--measure-rounds" + _help = ( + "Opt into the benchmark round layout with this many measured " + "passes in persistent workers. [default when benchmarking: 1]" + ) + _default = None + + def executor_options(self) -> dict: + """Reject ignored flags before loading a numerical runtime.""" + + dragon = { + "max_workers": self.max_workers, + "worker_timeout_sec": self.worker_timeout_sec, + "result_timeout_sec": self.result_timeout_sec, + } + mpi = {"rank_setup_timeout_sec": self.rank_setup_timeout_sec} + rounds = { + "warmup_rounds": self.warmup_rounds, + "measure_rounds": self.measure_rounds, + } + if self.executor == "local": + invalid = {**dragon, **mpi, **rounds} + elif self.executor == "dragon": + invalid = mpi + elif self.executor == "mpi": + invalid = dragon + else: + raise CommandError("executor must be local, dragon, or mpi") + supplied = [ + "--" + name.replace("_", "-") + for name, value in invalid.items() + if value is not None + ] + if supplied: + raise CommandError( + f"{', '.join(supplied)} cannot be used with " + f"--executor {self.executor}" + ) + if self.executor == "local": + return {} + options = { + name: value + for name, value in ( + dragon if self.executor == "dragon" else mpi + ).items() + if value is not None + } + options["benchmark"] = ( + BenchmarkOptions( + warmup_rounds=self.warmup_rounds or 0, + measure_rounds=self.measure_rounds or 1, + ) + if any(value is not None for value in rounds.values()) + else None + ) + return options diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py new file mode 100644 index 00000000..39ed70ac --- /dev/null +++ b/src/cuphoton/core/dragon.py @@ -0,0 +1,1096 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Optional Dragon execution of persistent, independently placed workloads.""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import queue +import socket +import sys +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +from cuphoton.core.benchmark import ( + BenchmarkOptions, + build_benchmark_report, +) +from cuphoton.core.bulk import ( + Placement, + WorkItem, + atomic_write_json, + error_payload, + item_ids_sha256, + json_mapping, + new_run_id, + partition_byte_balanced, + regular_file, + timestamp_utc, + validate_identifier, +) +from cuphoton.core.bulk import ( + hostnames_match as _hostnames_match, +) +from cuphoton.core.execution import ( + ExecutionResult, + WorkloadSpec, + audit_worker_provenance, + execute_worker_round, + factory_reference, + finalize_round, + prepare_run, + resolve_worker_factory, +) + +_LAUNCH_SCHEMA = "cuphoton.core.dragon-launch/v1" +_TEMPLATE_BUDGET_BYTES = 96 * 1024 +_WORKER_POLL_SEC = 0.5 + + +@dataclass(frozen=True) +class _DragonAPI: + System: Any + Node: Any + Policy: Any + ProcessGroup: Any + ProcessTemplate: Any + Queue: Any + + +def discover_gpu_placements( + system_type: Callable[[], Any], + node_type: Callable[[Any], Any], + *, + node_ids: Sequence[Any] | None = None, +) -> tuple[Placement, ...]: + """Enumerate actual Dragon Node.gpus IDs in allocation order.""" + + if node_ids is None: + node_ids = tuple(system_type().nodes) + placements: list[Placement] = [] + seen: set[tuple[str, int]] = set() + for node_id in node_ids: + node = node_type(node_id) + host = str(node.hostname) + for gpu_id in node.gpus or []: + if isinstance(gpu_id, bool) or not isinstance(gpu_id, int): + raise RuntimeError( + f"Dragon node {host!r} reported a non-integer GPU ID" + ) + key = (host, gpu_id) + if key in seen: + raise RuntimeError( + f"duplicate Dragon GPU placement: {host}:{gpu_id}" + ) + seen.add(key) + placements.append( + Placement( + worker_id=len(placements), + host=host, + gpu_id=gpu_id, + ) + ) + return tuple(placements) + + +def _select_gpu_placements( + placements: Sequence[Placement], worker_count: int +) -> tuple[Placement, ...]: + """Select GPUs round-robin across hosts, then assign dense worker IDs.""" + + if isinstance(worker_count, bool) or not isinstance(worker_count, int): + raise TypeError("worker_count must be an integer") + if worker_count <= 0 or worker_count > len(placements): + raise ValueError("worker_count exceeds available Dragon placements") + by_host: dict[str, list[Placement]] = {} + for placement in placements: + by_host.setdefault(placement.host, []).append(placement) + selected: list[Placement] = [] + offset = 0 + while len(selected) < worker_count: + progressed = False + for host_placements in by_host.values(): + if offset >= len(host_placements): + continue + selected.append(host_placements[offset]) + progressed = True + if len(selected) == worker_count: + break + if not progressed: + raise RuntimeError("could not select requested Dragon placements") + offset += 1 + return tuple( + Placement( + worker_id=worker_id, + host=placement.host, + gpu_id=placement.gpu_id, + ) + for worker_id, placement in enumerate(selected) + ) + + +def _singleton_cuda_visibility(expected_gpu_id: int | None = None) -> str: + raw = os.environ.get("CUDA_VISIBLE_DEVICES") + if raw is None: + raise RuntimeError("Dragon worker has no CUDA_VISIBLE_DEVICES") + tokens = [token.strip() for token in raw.split(",")] + if ( + any(not token for token in tokens) + or len(tokens) != 1 + or tokens[0] == "-1" + ): + raise RuntimeError( + f"Dragon worker must see exactly one CUDA device, got {raw!r}" + ) + if expected_gpu_id is not None and tokens[0] != str(expected_gpu_id): + raise RuntimeError( + "Dragon worker GPU placement mismatch: requested " + f"{expected_gpu_id}, got CUDA_VISIBLE_DEVICES={raw!r}" + ) + return tokens[0] + + +def _load_dragon_api() -> _DragonAPI: + try: + from dragon.infrastructure.policy import Policy + from dragon.native.machine import Node, System + from dragon.native.process import ProcessTemplate + from dragon.native.process_group import ProcessGroup + from dragon.native.queue import Queue + except (ImportError, OSError) as exc: + raise RuntimeError( + "The Dragon executor requires the dragonhpc runtime; " + "install it in this Python environment on every node " + "and launch with dragon" + ) from exc + return _DragonAPI( + System=System, + Node=Node, + Policy=Policy, + ProcessGroup=ProcessGroup, + ProcessTemplate=ProcessTemplate, + Queue=Queue, + ) + + +def _distribution_version(name: str) -> str | None: + try: + return version(name) + except PackageNotFoundError: + return None + + +def _stable_gpu_physical_ids( + value: Any, +) -> frozenset[tuple[str, str]] | None: + if ( + not isinstance(value, Mapping) + or "identity_error" not in value + or value["identity_error"] is not None + ): + return None + identities = frozenset( + (field, identifier.strip().casefold()) + for field in ("uuid", "pci_bus_id") + if isinstance((identifier := value.get(field)), str) + and identifier.strip() + ) + return identities or None + + +def _valid_shard_provenance( + value: Any, + *, + placement: Placement, + allow_loopback_alias: bool, + backend: str, +) -> bool: + if not isinstance(value, Mapping): + return False + worker_id = _strict_integer(value.get("worker_id")) + requested_gpu_id = _strict_integer(value.get("requested_gpu_id")) + pid = _strict_integer(value.get("pid")) + requested_host = value.get("requested_host") + hostname = value.get("hostname") + visibility = value.get("cuda_visible_devices") + if ( + worker_id != placement.worker_id + or requested_host != placement.host + or requested_gpu_id != placement.gpu_id + or pid is None + or pid <= 0 + or not isinstance(hostname, str) + or not hostname + or not _hostnames_match( + requested_host, + hostname, + allow_loopback_alias=allow_loopback_alias, + ) + or not isinstance(visibility, str) + or [token.strip() for token in visibility.split(",") if token.strip()] + != [str(placement.gpu_id)] + ): + return False + gpu = value.get("gpu") + if not isinstance(gpu, Mapping) or not gpu: + return False + gpu_backend = gpu.get("backend") + expected_identity_backend = ( + "cupy" if backend in {"cupy", "cutile"} else backend + ) + return ( + isinstance(gpu_backend, str) + and bool(gpu_backend) + and gpu_backend == expected_identity_backend + and _stable_gpu_physical_ids(gpu) is not None + ) + + +def _strict_integer(value: Any) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + + +def run_dragon_work_items( + *, + prepare_workload: Callable[[int], WorkloadSpec], + output_root: Path, + run_id: str | None = None, + max_workers: int | None = None, + result_timeout_sec: float = 60.0, + worker_timeout_sec: float = 3600.0, + benchmark: BenchmarkOptions | None = None, +) -> ExecutionResult: + """Execute a workload with persistent, explicitly placed Dragon workers. + + Worker factories run only after host and singleton CUDA binding checks. + Complete item descriptors live in hashed shared-filesystem launch files; + control queues carry bounded round commands and completion receipts. + + Successful runs explicitly close every worker before joining. On failure, + the coordinator stops the process group; a survivor blocked on its command + queue only runs ``worker.close()`` if Dragon's stop signal unwinds Python. + Forced process termination releases its CUDA resources, but application + cleanup and a CLOSED receipt are not guaranteed. Failed runs retain that + lifecycle error and cannot publish successful benchmark aggregates. + """ + + invocation_start = time.perf_counter() + for name, value in ( + ("result_timeout_sec", result_timeout_sec), + ("worker_timeout_sec", worker_timeout_sec), + ): + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be positive and finite") + if max_workers is not None and ( + isinstance(max_workers, bool) + or not isinstance(max_workers, int) + or max_workers <= 0 + ): + raise ValueError("max_workers must be a positive integer") + if benchmark is not None and not isinstance(benchmark, BenchmarkOptions): + raise TypeError("benchmark must be BenchmarkOptions or None") + spec = prepare_workload(0) + if not isinstance(spec, WorkloadSpec): + raise TypeError("prepare_workload must return WorkloadSpec") + timings = { + "workload_preflight_sec": time.perf_counter() - invocation_start + } + phase_start = time.perf_counter() + api = _load_dragon_api() + node_ids = tuple(api.System().nodes) + available = discover_gpu_placements( + api.System, api.Node, node_ids=node_ids + ) + worker_count = min( + max_workers or len(available), len(available), len(spec.items) + ) + if not worker_count: + raise RuntimeError("Dragon allocation exposes no usable GPUs") + placements = _select_gpu_placements(available, worker_count) + shards = partition_byte_balanced(spec.items, worker_count) + timings["dragon_discovery_sec"] = time.perf_counter() - phase_start + effective_run_id = run_id or new_run_id("dragon-workload") + validate_identifier(effective_run_id, field="run_id") + run_dir = output_root.expanduser().resolve() / effective_run_id + phase_start = time.perf_counter() + prepare_run(run_dir, effective_run_id, spec, "dragon") + for name in ("launch", "startup"): + (run_dir / name).mkdir() + timings["run_artifact_setup_sec"] = time.perf_counter() - phase_start + lifecycle_errors: list[dict[str, Any]] = [] + ready: list[dict[str, Any]] = [] + closed: list[dict[str, Any]] = [] + rounds: list[dict[str, Any]] = [] + terminal_summary: dict[str, Any] = {} + command_queues: list[Any] = [] + results_queue: Any | None = None + group: Any | None = None + started = False + joined = False + exit_status: list[dict[str, int]] = [] + plan = benchmark or BenchmarkOptions() + phase = "process_setup" + phase_start = time.perf_counter() + try: + results_queue = api.Queue(maxsize=2 * worker_count) + group = api.ProcessGroup( + restart=False, + ignore_error_on_exit=False, + walltime=worker_timeout_sec, + ) + for placement, shard in zip(placements, shards): + policy = api.Policy( + placement=api.Policy.Placement.HOST_NAME, + host_name=placement.host, + gpu_affinity=[placement.gpu_id], + ) + commands = api.Queue(maxsize=1, policy=policy) + command_queues.append(commands) + descriptor = { + "schema": _LAUNCH_SCHEMA, + "run_id": effective_run_id, + "run_dir": str(run_dir), + "placement": placement.to_dict(), + "items": [item.to_dict() for item in shard], + "options": dict(spec.options_payload), + "manifest_sha256": spec.manifest_sha256, + "backend": spec.backend, + "worker_factory": factory_reference(spec.worker_factory), + "benchmark": benchmark.to_payload() if benchmark else None, + "allow_loopback_alias": len(node_ids) == 1, + "worker_timeout_sec": worker_timeout_sec, + "result_timeout_sec": result_timeout_sec, + } + descriptor_path = ( + run_dir / "launch" / f"worker-{placement.worker_id:04d}.json" + ) + atomic_write_json(descriptor_path, descriptor, overwrite=False) + encoded = descriptor_path.read_bytes() + context = { + "run_dir": str(run_dir), + "worker_id": placement.worker_id, + "item_count": len(shard), + "weight_bytes": sum(item.weight_bytes for item in shard), + "item_ids_sha256": item_ids_sha256(shard), + "descriptor_bytes": len(encoded), + "result_timeout_sec": result_timeout_sec, + } + template = api.ProcessTemplate( + target=_workload_worker, + args=( + effective_run_id, + str(descriptor_path), + hashlib.sha256(encoded).hexdigest(), + context, + commands, + results_queue, + ), + policy=policy, + ) + if len(template.argdata) > _TEMPLATE_BUDGET_BYTES: + raise ValueError( + "Dragon worker launch arguments exceed 96 KiB" + ) + group.add_process(nproc=1, template=template) + phase = "init" + group.init() + timings["dragon_process_setup_sec"] = ( + time.perf_counter() - phase_start + ) + phase = "start" + phase_start = time.perf_counter() + worker_deadline = time.monotonic() + worker_timeout_sec + started = True + try: + group.start() + finally: + timings["dragon_launch_sec"] = time.perf_counter() - phase_start + phase = "ready" + try: + _collect_messages( + results_queue, + ready, + worker_count=worker_count, + run_id=effective_run_id, + kind="ready", + round_id=None, + deadline=worker_deadline, + group=group, + closed_messages=closed, + ) + provenances = [message["provenance"] for message in ready] + ready_errors = audit_worker_provenance( + provenances, + backend=spec.backend, + expected_worker_count=worker_count, + ) + if ready_errors: + raise ValueError( + f"invalid Dragon READY provenance: {ready_errors}" + ) + for message in ready: + if not _valid_shard_provenance( + message["provenance"], + placement=placements[message["worker_id"]], + allow_loopback_alias=len(node_ids) == 1, + backend=spec.backend, + ): + raise ValueError( + "Dragon READY differs from requested placement" + ) + finally: + timings["readiness_sec"] = time.perf_counter() - invocation_start + for round_spec in plan.rounds(): + phase = f"round:{round_spec.round_id}" + round_dir = ( + run_dir / "rounds" / round_spec.round_id + if benchmark + else run_dir + ) + round_run_id = ( + round_spec.run_id(effective_run_id) + if benchmark + else effective_run_id + ) + if benchmark: + prepare_run(round_dir, round_run_id, spec, "dragon") + messages: list[dict[str, Any]] = [] + round_errors: list[dict[str, Any]] = [] + round_timings: dict[str, float] = {} + batch_start = time.perf_counter() + try: + try: + for commands in command_queues: + remaining = min( + result_timeout_sec, + worker_deadline - time.monotonic(), + ) + if remaining <= 0: + raise TimeoutError( + "Dragon worker deadline expired" + ) + commands.put( + { + "run_id": effective_run_id, + **round_spec.to_payload(), + }, + timeout=remaining, + ) + finally: + round_timings["dispatch_sec"] = ( + time.perf_counter() - batch_start + ) + collection_start = time.perf_counter() + try: + _collect_messages( + results_queue, + messages, + worker_count=worker_count, + run_id=effective_run_id, + kind="round", + round_id=round_spec.round_id, + deadline=worker_deadline, + group=group, + closed_messages=closed, + ) + finally: + round_timings["collection_sec"] = ( + time.perf_counter() - collection_start + ) + except Exception as exc: + round_errors.append(error_payload(exc)) + batch_wall = time.perf_counter() - batch_start + worker_results = [ + message["result"] + for message in messages + if message.get("kind") == "round" + and message.get("run_id") == effective_run_id + and message.get("round_id") == round_spec.round_id + and isinstance(message.get("result"), Mapping) + ] + ready_by_worker = { + message["worker_id"]: message["provenance"] + for message in ready + } + if any( + result.get("provenance") + != ready_by_worker.get(result.get("worker_id")) + for result in worker_results + ): + round_errors.append( + { + "type": "WorkerIdentityChanged", + "message": "worker provenance changed after READY", + } + ) + audit_start = time.perf_counter() + terminal_summary = finalize_round( + round_dir, + round_run_id, + spec, + shards, + worker_results, + artifact_timeout_sec=min(result_timeout_sec, 0.01) + if round_errors + else result_timeout_sec, + ) + round_timings["artifact_audit_sec"] = max( + 0.0, + time.perf_counter() + - audit_start + - terminal_summary["finalization_sec"], + ) + receipt = { + **round_spec.to_payload(), + "status": "success" + if terminal_summary["status"] == "success" + and not round_errors + else "failed", + "batch_wall_sec": batch_wall, + "finalization_sec": terminal_summary["finalization_sec"], + "worker_wall_max_sec": max( + ( + _duration(message.get("worker_wall_sec")) + for message in messages + ), + default=0.0, + ), + "coordinator_timings_sec": round_timings, + "summary_path": str( + (round_dir / "summary.json").relative_to(run_dir) + ), + } + terminal_summary.update( + receipt, + parent_run_id=effective_run_id, + messages=messages, + round_errors=round_errors, + ) + if benchmark: + atomic_write_json( + round_dir / "summary.json", terminal_summary + ) + rounds.append(receipt) + if receipt["status"] != "success": + raise RuntimeError( + f"Dragon round {round_spec.round_id} failed" + ) + phase = "worker_close" + close_start = time.perf_counter() + for commands in command_queues: + remaining = min( + result_timeout_sec, worker_deadline - time.monotonic() + ) + if remaining <= 0: + raise TimeoutError("Dragon worker deadline expired") + commands.put( + {"run_id": effective_run_id, "kind": "close"}, + timeout=remaining, + ) + _collect_messages( + results_queue, + closed, + worker_count=worker_count, + run_id=effective_run_id, + kind="closed", + round_id=None, + deadline=time.monotonic() + result_timeout_sec, + group=group, + ) + timings["worker_close_sec"] = time.perf_counter() - close_start + phase = "join" + phase_start = time.perf_counter() + try: + group.join(timeout=result_timeout_sec) + joined = True + finally: + timings["worker_join_sec"] = time.perf_counter() - phase_start + phase = "unexpected_result" + try: + results_queue.get(timeout=0) + except (queue.Empty, TimeoutError): + pass + else: + raise ValueError("unexpected extra Dragon worker result") + except Exception as exc: + lifecycle_errors.append({"phase": phase, **error_payload(exc)}) + finally: + timings.setdefault( + "dragon_process_setup_sec", time.perf_counter() - phase_start + ) + cleanup_start = time.perf_counter() + if group is not None: + if started and not joined: + try: + group.stop(patience=5.0) + except Exception as exc: + lifecycle_errors.append( + {"phase": "stop_after_failure", **error_payload(exc)} + ) + try: + exit_status = [ + {"puid": int(puid), "exit_code": int(code)} + for puid, code in group.inactive_puids + ] + except Exception as exc: + lifecycle_errors.append( + {"phase": "exit_status", **error_payload(exc)} + ) + try: + group.close(patience=5.0) + except Exception as exc: + lifecycle_errors.append( + {"phase": "group_close", **error_payload(exc)} + ) + cleanup = getattr(group, "_close_no_decorator", None) + if cleanup is not None: + try: + cleanup(patience=5.0) + except Exception as cleanup_exc: + lifecycle_errors.append( + { + "phase": "forced_close", + **error_payload(cleanup_exc), + } + ) + for channel in [*command_queues, results_queue]: + if channel is not None: + try: + channel.close() + except Exception as exc: + lifecycle_errors.append( + {"phase": "queue_close", **error_payload(exc)} + ) + timings["dragon_cleanup_sec"] = time.perf_counter() - cleanup_start + process_audit = { + "ok": len(exit_status) == worker_count + and all(item["exit_code"] == 0 for item in exit_status), + "expected_count": worker_count, + "observed_count": len(exit_status), + "nonzero": [item for item in exit_status if item["exit_code"] != 0], + } + if not process_audit["ok"]: + lifecycle_errors.append( + { + "phase": "exit_status", + "type": "WorkerExitFailure", + "message": "missing or nonzero Dragon worker exit status", + } + ) + report = build_benchmark_report(plan, rounds, errors=lifecycle_errors) + summary = { + **(terminal_summary if benchmark is None else {}), + "schema": "cuphoton.core.dragon-summary/v1", + "executor": "dragon", + "run_id": effective_run_id, + "status": report["status"], + "manifest_sha256": spec.manifest_sha256, + "options": dict(spec.options_payload), + "backend": spec.backend, + "dragonhpc_version": _distribution_version("dragonhpc"), + "worker_timeout_sec": worker_timeout_sec, + "result_timeout_sec": result_timeout_sec, + "worker_count": worker_count, + "allocation_node_count": len(node_ids), + "distinct_host_count": len( + {placement.host for placement in placements} + ), + "placements": [placement.to_dict() for placement in placements], + "command_queue_placement": "consumer", + "coordinator_timings_sec": timings, + "coordinator_wall_sec": time.perf_counter() - invocation_start, + "coordinator_wall_definition": ( + "Function entry through worker cleanup and terminal audits; " + "the final summary.json atomic commit is excluded." + ), + "completed_at_utc": timestamp_utc(), + "ready_messages": ready, + "closed_messages": closed, + "process_exit_status": exit_status, + "process_exit_audit": process_audit, + "lifecycle_errors": lifecycle_errors, + } + if benchmark: + summary["benchmark"] = report + summary_path = run_dir / "summary.json" + atomic_write_json(summary_path, summary) + return ExecutionResult( + "dragon", + effective_run_id, + run_dir, + summary_path, + summary["status"], + summary, + ) + + +def _duration(value: Any) -> float: + return ( + float(value) + if not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(value) + and value >= 0 + else 0.0 + ) + + +def _collect_messages( + results_queue: Any, + messages: list[dict[str, Any]], + *, + worker_count: int, + run_id: str, + kind: str, + round_id: str | None, + deadline: float, + group: Any, + closed_messages: list[dict[str, Any]] | None = None, +) -> None: + seen: dict[int, int] = {} + failed_workers: set[int] = set() + observed_exits: list[tuple[int, int]] | None = None + while len(seen) < worker_count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Dragon {kind} deadline expired") + try: + message = json_mapping( + results_queue.get( + timeout=0 + if observed_exits is not None + else min(remaining, _WORKER_POLL_SEC) + ), + field=f"Dragon {kind} result", + ) + except queue.Empty: + if observed_exits is None: + # Snapshot exits before draining: a worker sends its receipt + # before exiting, possibly just after the timed get expires. + observed_exits = list(group.inactive_puids) + if not observed_exits: + observed_exits = None + continue + missing_exits = [ + (puid, code) + for puid, code in observed_exits + if puid not in seen.values() + ] + if missing_exits: + raise RuntimeError( + f"Dragon worker exited before {kind} completion: " + f"{missing_exits}" + ) + observed_exits = None + continue + worker_id = _strict_integer(message.get("worker_id")) + puid = _strict_integer(message.get("puid")) + if ( + closed_messages is not None + and kind != "closed" + and message.get("kind") == "closed" + ): + if ( + message.get("run_id") != run_id + or message.get("round_id") is not None + or worker_id not in failed_workers + or puid != seen.get(worker_id) + or any( + m.get("worker_id") == worker_id for m in closed_messages + ) + ): + raise ValueError( + "invalid or duplicate Dragon closed identity" + ) + closed_messages.append(message) + continue + messages.append(message) + if ( + message.get("run_id") != run_id + or message.get("kind") != kind + or message.get("round_id") != round_id + or worker_id is None + or worker_id not in range(worker_count) + or worker_id in seen + or puid is None + or puid in seen.values() + ): + raise ValueError(f"invalid or duplicate Dragon {kind} identity") + seen[worker_id] = puid + if message.get("status") != "success": + failed_workers.add(worker_id) + if kind == "round" and message.get("status") == "success": + result = message.get("result") + duration = message.get("worker_wall_sec") + if ( + isinstance(duration, bool) + or not isinstance(duration, (int, float)) + or not math.isfinite(duration) + or duration < 0 + or not isinstance(result, Mapping) + or _strict_integer(result.get("worker_id")) != worker_id + or result.get("status") != "success" + ): + raise ValueError("invalid Dragon round worker result") + if failed_workers: + raise RuntimeError( + f"Dragon workers {sorted(failed_workers)} {kind} failed" + ) + + +def _read_descriptor( + run_id: str, + descriptor_path_raw: str, + digest: str, + context: Mapping[str, Any], +) -> dict[str, Any]: + run_dir = Path(context["run_dir"]) + worker_id = context["worker_id"] + path = Path(descriptor_path_raw) + if ( + not run_dir.is_absolute() + or path != run_dir / "launch" / f"worker-{worker_id:04d}.json" + or not regular_file(path) + ): + raise ValueError("Dragon launch descriptor is not a regular run file") + with path.open("rb") as handle: + encoded = handle.read(context["descriptor_bytes"] + 1) + if ( + len(encoded) != context["descriptor_bytes"] + or hashlib.sha256(encoded).hexdigest() != digest + ): + raise ValueError("Dragon launch descriptor size or SHA-256 differs") + descriptor = json_mapping( + json.loads(encoded), field="Dragon launch descriptor" + ) + if ( + descriptor["schema"] != _LAUNCH_SCHEMA + or descriptor["run_id"] != run_id + or descriptor["run_dir"] != str(run_dir) + or descriptor["placement"]["worker_id"] != worker_id + ): + raise ValueError("Dragon launch descriptor identity differs") + items = tuple( + WorkItem.from_dict(payload) for payload in descriptor["items"] + ) + if ( + len(items) != context["item_count"] + or sum(item.weight_bytes for item in items) != context["weight_bytes"] + or item_ids_sha256(items) != context["item_ids_sha256"] + ): + raise ValueError("Dragon launch descriptor shard identity differs") + return descriptor + + +def _validate_binding( + placement: Placement, *, allow_loopback_alias: bool +) -> tuple[str, str]: + hostname = socket.gethostname() + if not _hostnames_match( + placement.host, hostname, allow_loopback_alias=allow_loopback_alias + ): + raise RuntimeError( + "Dragon worker host differs from requested placement" + ) + visibility = _singleton_cuda_visibility(placement.gpu_id) + premature = [ + name + for name in ("cupy", "numba.cuda", "cuda.tile") + if name in sys.modules + ] + torch = sys.modules.get("torch") + if torch is not None and torch.cuda.is_initialized(): + premature.append("torch CUDA") + if premature: + raise RuntimeError( + "CUDA modules were imported before Dragon worker placement: " + + ", ".join(premature) + ) + return hostname, visibility + + +def _current_puid() -> int: + from dragon.infrastructure.parameters import this_process + + return int(this_process.my_puid) + + +def _workload_worker( + run_id: str, + descriptor_path: str, + digest: str, + context: Mapping[str, Any], + commands: Any, + results_queue: Any, +) -> None: + worker = None + worker_id = context["worker_id"] + puid = _current_puid() + run_dir = Path(context["run_dir"]) + timeout = context["result_timeout_sec"] + ready: dict[str, Any] = { + "kind": "ready", + "puid": puid, + "run_id": run_id, + "round_id": None, + "worker_id": worker_id, + "status": "success", + } + failed = False + failure_error = None + try: + try: + descriptor = _read_descriptor( + run_id, descriptor_path, digest, context + ) + placement = Placement(**descriptor["placement"]) + hostname, visibility = _validate_binding( + placement, + allow_loopback_alias=descriptor["allow_loopback_alias"], + ) + factory = resolve_worker_factory(descriptor["worker_factory"]) + worker = factory(descriptor["options"]) + provenance = { + "worker_id": worker_id, + "requested_host": placement.host, + "requested_gpu_id": placement.gpu_id, + "hostname": hostname, + "pid": os.getpid(), + "cuda_visible_devices": visibility, + "gpu": json_mapping( + worker.gpu_identity, field="worker GPU identity" + ), + } + ready["provenance"] = provenance + except Exception as exc: + ready.update(status="failed", error=error_payload(exc)) + try: + atomic_write_json( + run_dir / "startup" / f"worker-{worker_id:04d}.json", ready + ) + except Exception as exc: + ready.update(status="failed", artifact_error=error_payload(exc)) + results_queue.put(ready, timeout=timeout) + if ready["status"] != "success": + raise RuntimeError("Dragon worker initialization failed") + benchmark = ( + BenchmarkOptions(**descriptor["benchmark"]) + if descriptor["benchmark"] + else None + ) + items = tuple( + WorkItem.from_dict(payload) for payload in descriptor["items"] + ) + for round_spec in (benchmark or BenchmarkOptions()).rounds(): + message: dict[str, Any] = { + "kind": "round", + "puid": puid, + "run_id": run_id, + "round_id": round_spec.round_id, + "worker_id": worker_id, + "status": "success", + } + round_dir = ( + run_dir / "rounds" / round_spec.round_id + if benchmark + else run_dir + ) + try: + command = commands.get( + timeout=descriptor["worker_timeout_sec"] + ) + expected = {"run_id": run_id, **round_spec.to_payload()} + if ( + not isinstance(command, Mapping) + or command != expected + or any( + type(command.get(key)) is not type(value) + for key, value in expected.items() + ) + ): + raise ValueError("unexpected Dragon round command") + start = time.perf_counter() + result = execute_worker_round( + worker, + items=items, + run_id=round_spec.run_id(run_id) if benchmark else run_id, + run_dir=round_dir, + manifest_sha256=descriptor["manifest_sha256"], + worker_id=worker_id, + backend=descriptor["backend"], + provenance=provenance, + ) + message.update( + status=result["status"], + result=result, + worker_wall_sec=time.perf_counter() - start, + ) + except Exception as exc: + message.update(status="failed", error=error_payload(exc)) + try: + atomic_write_json( + round_dir / "errors" / f"worker-{worker_id:04d}.json", + message, + ) + except Exception as write_exc: + message["artifact_error"] = error_payload(write_exc) + results_queue.put(message, timeout=timeout) + if message["status"] != "success": + raise RuntimeError( + f"Dragon worker round {round_spec.round_id} failed" + ) + # Closing produces another result-queue message. Wait until the + # coordinator has collected and audited every final-round receipt. + command = commands.get(timeout=descriptor["worker_timeout_sec"]) + if command != {"run_id": run_id, "kind": "close"}: + raise ValueError("unexpected Dragon worker close command") + except BaseException as exc: + failed = True + failure_error = error_payload(exc) + raise + finally: + closed = { + "kind": "closed", + "puid": puid, + "run_id": run_id, + "round_id": None, + "worker_id": worker_id, + "status": "failed" if failed else "success", + } + if failure_error is not None: + closed["error"] = failure_error + if worker is not None: + try: + worker.close() + except Exception as exc: + closed.update(status="failed", error=error_payload(exc)) + try: + atomic_write_json( + run_dir / "startup" / f"worker-{worker_id:04d}-closed.json", + closed, + ) + except Exception as exc: + closed.update(status="failed", artifact_error=error_payload(exc)) + try: + results_queue.put(closed, timeout=timeout) + except Exception: + if not failed: + raise + if closed["status"] != "success" and not failed: + raise RuntimeError("Dragon worker cleanup failed") diff --git a/src/cuphoton/core/execution.py b/src/cuphoton/core/execution.py new file mode 100644 index 00000000..8c1e1594 --- /dev/null +++ b/src/cuphoton/core/execution.py @@ -0,0 +1,887 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Whole-item execution contracts shared by optional distributed runtimes.""" + +from __future__ import annotations + +import importlib +import json +import math +import os +import stat +import time +from collections import Counter +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Protocol + +from .bulk import ( + WorkItem, + atomic_write_json, + audit_terminal_records, + classify_physical_gpu_pair, + error_payload, + hostnames_match, + item_ids_sha256, + json_mapping, + partition_byte_balanced, + read_json_mapping, + regular_file, + timestamp_utc, + validate_identifier, +) + +ITEM_SCHEMA = "cuphoton.core.execution-item/v1" +WORKER_SCHEMA = "cuphoton.core.execution-worker/v1" +SUMMARY_SCHEMA = "cuphoton.core.execution-summary/v1" + + +class Worker(Protocol): + """Worker-local state constructed after placement and CUDA binding. + + ``gpu_identity`` must contain ``backend`` (``cupy`` for ``cutile``), + ``device_index=0``, ``identity_error=None``, and at least one non-empty + physical ``uuid`` or ``pci_bus_id`` string. + """ + + gpu_identity: Mapping[str, Any] + + def run_item(self, item: WorkItem, output_dir: Path) -> Mapping[str, Any]: + """Execute an item and durably publish ``output_dir/summary.json``. + + Return JSON metadata with ``run_dir`` and ``summary_path`` equal to + those absolute paths, ``backend`` equal to the workload backend, a + non-empty ``device`` string, a ``runtime`` mapping, and mappings + ``timings_sec`` and ``wall_sec`` of names to finite non-negative + seconds. + """ + + def close(self) -> None: + """Complete pending work and release worker-local resources.""" + + +WorkerFactory = Callable[[Mapping[str, Any]], Worker] +RecordValidator = Callable[[Mapping[str, Any], Path], Sequence[str]] +RoundFinalizer = Callable[ + [Path, Sequence[Mapping[str, Any]]], Mapping[str, Any] +] + + +@dataclass(frozen=True) +class WorkloadSpec: + """Preflighted JSON descriptors and worker-local execution callbacks. + + Factories are package-importable. Scientific validators and finalizers run + only on the coordinator; they are never sent to a worker. + """ + + items: Sequence[WorkItem] + options_payload: Mapping[str, Any] + manifest_payload: Mapping[str, Any] + input_identity_payload: Mapping[str, Any] + manifest_sha256: str + backend: str + worker_factory: WorkerFactory + success_record_validator: RecordValidator | None = None + failed_record_validator: RecordValidator | None = None + finalize_round: RoundFinalizer | None = None + + def __post_init__(self) -> None: + items = tuple(self.items) + if not items or not all(isinstance(item, WorkItem) for item in items): + raise ValueError("workload requires non-empty WorkItems") + partition_byte_balanced(items, 1) + object.__setattr__(self, "items", items) + for field in ( + "options_payload", + "manifest_payload", + "input_identity_payload", + ): + object.__setattr__( + self, field, json_mapping(getattr(self, field), field=field) + ) + if ( + not isinstance(self.manifest_sha256, str) + or len(self.manifest_sha256) != 64 + or any(c not in "0123456789abcdef" for c in self.manifest_sha256) + ): + raise ValueError("manifest_sha256 must be a lowercase SHA-256") + if self.backend not in {"cupy", "cutile", "numba-cuda", "torch"}: + raise ValueError("workload requires an explicit GPU backend") + factory_reference(self.worker_factory) + for field in ( + "success_record_validator", + "failed_record_validator", + "finalize_round", + ): + if getattr(self, field) is not None and not callable( + getattr(self, field) + ): + raise TypeError(f"{field} must be callable or None") + + def identity_payload(self) -> dict[str, Any]: + """Return launch data for agreement between ranks and provenance.""" + + return { + "items": [item.to_dict() for item in self.items], + "options": dict(self.options_payload), + "manifest": dict(self.manifest_payload), + "input_identity": dict(self.input_identity_payload), + "manifest_sha256": self.manifest_sha256, + "backend": self.backend, + "worker_factory": factory_reference(self.worker_factory), + } + + +@dataclass(frozen=True) +class ExecutionResult: + """Coordinator handle for a terminal invocation.""" + + executor: str + run_id: str + run_dir: Path + summary_path: Path + status: str + summary: Mapping[str, Any] + + def to_dict(self) -> dict[str, Any]: + return { + "executor": self.executor, + "run_id": self.run_id, + "run_dir": str(self.run_dir), + "summary_path": str(self.summary_path), + "status": self.status, + } + + +def factory_reference(factory: WorkerFactory) -> dict[str, str]: + """Verify an import reference without serializing a live worker object.""" + + reference = { + "module": str(getattr(factory, "__module__", "")), + "qualname": str(getattr(factory, "__qualname__", "")), + } + if resolve_worker_factory(reference) is not factory: + raise ValueError( + "worker factory import resolves to a different callable" + ) + return reference + + +def resolve_worker_factory(reference: Mapping[str, Any]) -> WorkerFactory: + if set(reference) != {"module", "qualname"}: + raise ValueError("worker factory reference has invalid fields") + module = reference["module"] + qualname = reference["qualname"] + if ( + not isinstance(module, str) + or not module + or module == "__main__" + or not isinstance(qualname, str) + or not qualname + or "" in qualname + ): + raise ValueError("worker factory must be package-importable") + result: Any = importlib.import_module(module) + for part in qualname.split("."): + if not part or part.startswith("<"): + raise ValueError("worker factory qualname is invalid") + result = getattr(result, part) + if not callable(result): + raise TypeError("worker factory is not callable") + return result + + +def prepare_run( + run_dir: Path, + run_id: str, + spec: WorkloadSpec, + executor: str, + *, + directory_claimed: bool = False, +) -> None: + """Claim a run directory and publish its immutable input descriptors.""" + + validate_identifier(run_id, field="run_id") + if directory_claimed: + if not _real_directory(run_dir) or any(run_dir.iterdir()): + raise ValueError( + "claimed execution directory must be real and empty" + ) + else: + run_dir.mkdir(parents=True, exist_ok=False) + for name in ("items", "records", "workers"): + (run_dir / name).mkdir() + for name, payload in ( + ("manifest.json", spec.manifest_payload), + ("input-identity.json", spec.input_identity_payload), + ( + "run.json", + { + "schema": "cuphoton.core.execution-run/v1", + "executor": executor, + "run_id": run_id, + "started_at_utc": timestamp_utc(), + **spec.identity_payload(), + }, + ), + ): + atomic_write_json(run_dir / name, payload, overwrite=False) + + +def execute_worker_round( + worker: Worker, + *, + items: Sequence[WorkItem], + run_id: str, + run_dir: Path, + manifest_sha256: str, + worker_id: int, + backend: str, + provenance: Mapping[str, Any], +) -> dict[str, Any]: + """Run a placed worker's shard and retain every terminal item record. + + The caller measures around this function to include the final worker + receipt's durable write. Receipts contain no worker objects or device + values. + """ + + started = time.perf_counter() + started_at = timestamp_utc() + successes = failures = 0 + write_errors: list[dict[str, Any]] = [] + for item in items: + item_started = time.perf_counter() + identity = { + "schema": ITEM_SCHEMA, + "run_id": run_id, + "manifest_sha256": manifest_sha256, + "worker_id": worker_id, + "item_id": item.item_id, + "weight_bytes": item.weight_bytes, + "requested_backend": backend, + "started_at_utc": timestamp_utc(), + } + record = dict(identity) + try: + item_dir = run_dir / "items" / item.item_id + metadata = json_mapping( + worker.run_item(item, item_dir), field="item result metadata" + ) + _validate_item_output(item_dir, metadata) + for field in ("run_dir", "summary_path"): + metadata[field] = str( + Path(metadata[field]).relative_to(run_dir) + ) + record.update(metadata) + record["status"] = "success" + except Exception as exc: + record["status"] = "failed" + record["error"] = error_payload(exc) + if getattr(exc, "__notes__", None): + record["error"]["notes"] = "\n".join(map(str, exc.__notes__)) + record.update(identity) + record.update( + worker_seconds=time.perf_counter() - item_started, + completed_at_utc=timestamp_utc(), + ) + try: + atomic_write_json( + run_dir / "records" / f"{item.item_id}.json", record + ) + except Exception as exc: + write_errors.append( + {"item_id": item.item_id, **error_payload(exc)} + ) + failures += 1 + else: + successes += record["status"] == "success" + failures += record["status"] != "success" + result = { + "schema": WORKER_SCHEMA, + "run_id": run_id, + "manifest_sha256": manifest_sha256, + "worker_id": worker_id, + "assigned_item_ids": [item.item_id for item in items], + "item_ids_sha256": item_ids_sha256(items), + "weight_bytes": sum(item.weight_bytes for item in items), + "success_count": successes, + "failed_count": failures, + "status": "failed" if failures else "success", + "worker_wall_sec": time.perf_counter() - started, + "started_at_utc": started_at, + "completed_at_utc": timestamp_utc(), + "provenance": json_mapping(provenance, field="worker provenance"), + "record_write_errors": write_errors, + "error": None, + } + try: + atomic_write_json( + run_dir / "workers" / f"worker-{worker_id:04d}.json", result + ) + except Exception as exc: + result["status"] = "failed" + result["artifact_error"] = error_payload(exc) + return result + + +def audit_worker_provenance( + provenances: Sequence[Mapping[str, Any]], + *, + backend: str, + expected_worker_count: int, +) -> list[dict[str, Any]]: + """Require one usable, non-overlapping physical GPU per placed worker.""" + + errors: list[dict[str, Any]] = [] + identities: list[tuple[int, str, frozenset[tuple[str, str]]]] = [] + worker_ids = [ + p.get("worker_id") if isinstance(p, Mapping) else None + for p in provenances + ] + if ( + len(provenances) != expected_worker_count + or any(type(value) is not int for value in worker_ids) + or set(worker_ids) != set(range(expected_worker_count)) + ): + errors.append( + {"phase": "provenance", "message": "worker identities differ"} + ) + for index, value in enumerate(provenances): + invalid = [] + if not isinstance(value, Mapping): + errors.append( + { + "phase": "provenance", + "worker_id": index, + "message": "invalid provenance", + } + ) + continue + gpu = value.get("gpu") + host = value.get("hostname") + visibility = value.get("cuda_visible_devices") + if type(value.get("pid")) is not int or value["pid"] <= 0: + invalid.append("pid") + if not isinstance(host, str) or not host: + invalid.append("hostname") + tokens = visibility.split(",") if isinstance(visibility, str) else [] + if ( + len(tokens) != 1 + or not tokens[0].strip() + or tokens[0].strip().startswith("-") + ): + invalid.append("cuda_visible_devices") + ids = frozenset() + expected_backend = "cupy" if backend == "cutile" else backend + if not isinstance(gpu, Mapping): + invalid.append("gpu") + else: + if ( + gpu.get("backend") != expected_backend + or type(gpu.get("device_index")) is not int + or gpu["device_index"] != 0 + or "identity_error" not in gpu + or gpu["identity_error"] is not None + ): + invalid.append("gpu") + ids = frozenset( + (field, identifier.strip().casefold()) + for field in ("uuid", "pci_bus_id") + if isinstance((identifier := gpu.get(field)), str) + and identifier.strip() + ) + if not ids: + invalid.append("physical GPU identity") + if invalid: + errors.append( + { + "phase": "provenance", + "worker_id": value.get("worker_id"), + "message": ", ".join(invalid), + } + ) + else: + identities.append((value["worker_id"], host, ids)) + for index, (worker_id, host, ids) in enumerate(identities): + for other_id, other_host, other_ids in identities[:index]: + same_host = hostnames_match(host, other_host) + relation = classify_physical_gpu_pair( + dict(ids), dict(other_ids), same_host=same_host + ) + if relation != "distinct": + errors.append( + { + "phase": "provenance", + "worker_ids": [other_id, worker_id], + "message": f"{relation} physical GPU identities", + } + ) + return errors + + +def finalize_round( + run_dir: Path, + run_id: str, + spec: WorkloadSpec, + shards: Sequence[Sequence[WorkItem]], + worker_results: Sequence[Mapping[str, Any]], + *, + artifact_timeout_sec: float, +) -> dict[str, Any]: + """Audit evidence, merge outputs, and return an unpublished summary. + + The executor publishes benchmark round summaries separately and commits + the root terminal summary only after worker shutdown and lifecycle audits. + """ + + if ( + not _finite_nonnegative(artifact_timeout_sec) + or artifact_timeout_sec <= 0 + ): + raise ValueError("artifact_timeout_sec must be positive") + errors, records = _read_round_artifacts( + run_dir, shards, worker_results, artifact_timeout_sec + ) + expected_items = [item.item_id for shard in shards for item in shard] + item_audit = audit_terminal_records( + expected_items, + [ + { + **record, + "status": record.get("status") + if isinstance(record.get("status"), str) + else None, + } + for record in records + ], + ) + if not item_audit["ok"] or item_audit["failed_item_ids"]: + errors.append( + {"phase": "items", "message": "terminal item audit failed"} + ) + expected = { + item.item_id: (worker_id, item.weight_bytes) + for worker_id, shard in enumerate(shards) + for item in shard + } + for record in records: + fields = _record_problems(record, run_dir, run_id, spec, expected) + validator = ( + spec.success_record_validator + if record.get("status") == "success" + else spec.failed_record_validator + ) + if validator is not None: + try: + problems = validator(record, run_dir) + if ( + isinstance(problems, (str, bytes)) + or not isinstance(problems, Sequence) + or any( + not isinstance(problem, str) for problem in problems + ) + ): + raise TypeError( + "component validator must return field names" + ) + fields.extend(problems) + except Exception as exc: + fields.append( + f"component validator: {type(exc).__name__}: {exc}" + ) + if fields: + errors.append( + { + "phase": "records", + "item_id": record.get("item_id"), + "message": ", ".join(sorted(set(fields))), + } + ) + errors.extend( + _worker_result_problems(run_id, spec, shards, worker_results, records) + ) + errors.extend( + audit_worker_provenance( + [ + result.get("provenance") + for result in worker_results + if isinstance(result, Mapping) + ], + backend=spec.backend, + expected_worker_count=len(shards), + ) + ) + scientific_result = None + finalization_sec = 0.0 + if not errors and spec.finalize_round is not None: + finalization_start = time.perf_counter() + try: + by_id = {record["item_id"]: record for record in records} + scientific_result = json_mapping( + spec.finalize_round( + run_dir, [by_id[item.item_id] for item in spec.items] + ), + field="component round finalization", + ) + except Exception as exc: + errors.append({"phase": "finalization", **error_payload(exc)}) + finally: + finalization_sec = time.perf_counter() - finalization_start + summary = { + "schema": SUMMARY_SCHEMA, + "run_id": run_id, + "manifest_sha256": spec.manifest_sha256, + "backend": spec.backend, + "status": "failed" if errors else "success", + "completed_at_utc": timestamp_utc(), + "terminal_record_audit": item_audit, + "records": records, + "worker_results": [dict(result) for result in worker_results], + "errors": errors, + "result": scientific_result, + "finalization_sec": finalization_sec, + } + return summary + + +def _read_round_artifacts( + run_dir: Path, + shards: Sequence[Sequence[WorkItem]], + results: Sequence[Mapping[str, Any]], + timeout: float, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + expected = { + f"records/{item.item_id}.json" for shard in shards for item in shard + } | { + f"workers/worker-{worker_id:04d}.json" + for worker_id in range(len(shards)) + } + errors: list[dict[str, Any]] = [] + for result in results: + if not isinstance(result, Mapping): + continue + worker_id = result.get("worker_id") + if type(worker_id) is not int or worker_id not in range(len(shards)): + continue + if result.get("artifact_error") is not None: + expected.discard(f"workers/worker-{worker_id:04d}.json") + errors.append( + { + "phase": "publication", + "worker_id": worker_id, + "error": result["artifact_error"], + } + ) + writes = result.get("record_write_errors") + if isinstance(writes, list): + assigned = {item.item_id for item in shards[worker_id]} + for error in writes: + if ( + isinstance(error, Mapping) + and error.get("item_id") in assigned + ): + expected.discard(f"records/{error['item_id']}.json") + errors.append({"phase": "publication", **dict(error)}) + deadline = time.monotonic() + timeout + mappings: dict[str, dict[str, Any]] = {} + while True: + pending = [] + for label in sorted(expected): + try: + path = run_dir / label + if not regular_file(path): + if os.path.lexists(path): + raise ValueError("artifact is not a regular file") + raise FileNotFoundError(label) + mapping = json_mapping(read_json_mapping(path), field=label) + if ( + label.startswith("records/") + and mapping.get("item_id") != Path(label).stem + ): + raise ValueError( + "terminal item identity differs from record path" + ) + mappings[label] = mapping + if ( + label.startswith("records/") + and mapping.get("status") == "success" + ): + item_id = Path(label).stem + item_dir = run_dir / "items" / item_id + if not _real_directory(item_dir) or not regular_file( + item_dir / "summary.json" + ): + if os.path.lexists(item_dir / "summary.json"): + raise ValueError("item output is not regular") + raise FileNotFoundError( + f"items/{item_id}/summary.json" + ) + except OSError as exc: + pending.append( + { + "phase": "visibility", + "path": label, + **error_payload(exc), + } + ) + except Exception as exc: + errors.append( + { + "phase": "artifacts", + "path": label, + **error_payload(exc), + } + ) + if not pending or errors or time.monotonic() >= deadline: + errors.extend(pending) + break + time.sleep(min(0.05, max(0.0, deadline - time.monotonic()))) + expected_records = { + f"{item.item_id}.json" for shard in shards for item in shard + } + observed_records = { + path.name for path in (run_dir / "records").glob("*.json") + } + if observed_records - expected_records: + errors.append( + {"phase": "artifacts", "message": "unexpected terminal records"} + ) + for result in results: + if ( + not isinstance(result, Mapping) + or type(result.get("worker_id")) is not int + ): + continue + label = f"workers/worker-{result['worker_id']:04d}.json" + if label in mappings and not _same_json(mappings[label], result): + errors.append( + { + "phase": "artifacts", + "path": label, + "message": "durable worker receipt differs from " + "collected receipt", + } + ) + return errors, [ + mapping + for label, mapping in sorted(mappings.items()) + if label.startswith("records/") + ] + + +def _record_problems( + record: Mapping[str, Any], + run_dir: Path, + run_id: str, + spec: WorkloadSpec, + assignments: Mapping[str, tuple[int, int]], +) -> list[str]: + problems = [] + item_id = record.get("item_id") + assignment = ( + assignments.get(item_id) if isinstance(item_id, str) else None + ) + for field, expected in ( + ("schema", ITEM_SCHEMA), + ("run_id", run_id), + ("manifest_sha256", spec.manifest_sha256), + ("requested_backend", spec.backend), + ): + if record.get(field) != expected: + problems.append(field) + if assignment is None: + problems.append("item_id") + else: + for field, value in zip(("worker_id", "weight_bytes"), assignment): + if type(record.get(field)) is not int or record[field] != value: + problems.append(field) + if not _finite_nonnegative(record.get("worker_seconds")): + problems.append("worker_seconds") + if not _valid_timestamps(record): + problems.append("timestamps") + if record.get("status") == "success": + if record.get("backend") != spec.backend: + problems.append("backend") + if record.get("run_dir") != f"items/{item_id}": + problems.append("run_dir") + if record.get("summary_path") != f"items/{item_id}/summary.json": + problems.append("summary_path") + if not isinstance(record.get("device"), str) or not record["device"]: + problems.append("device") + if not isinstance(record.get("runtime"), Mapping): + problems.append("runtime") + for field in ("timings_sec", "wall_sec"): + if not _valid_timings(record.get(field)): + problems.append(field) + if record.get("error") is not None: + problems.append("error") + elif record.get("status") == "failed": + if not _valid_error(record.get("error")): + problems.append("error") + else: + problems.append("status") + return problems + + +def _worker_result_problems( + run_id: str, + spec: WorkloadSpec, + shards: Sequence[Sequence[WorkItem]], + results: Sequence[Mapping[str, Any]], + records: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + errors = [] + identifiers = [ + value.get("worker_id") if isinstance(value, Mapping) else None + for value in results + ] + if ( + len(results) != len(shards) + or any(type(value) is not int for value in identifiers) + or set(identifiers) != set(range(len(shards))) + ): + errors.append( + { + "phase": "workers", + "message": "worker receipt identities differ", + } + ) + for result in results: + if not isinstance(result, Mapping): + errors.append( + {"phase": "workers", "message": "invalid worker receipt"} + ) + continue + worker_id = result.get("worker_id") + if type(worker_id) is not int or worker_id not in range(len(shards)): + continue + shard = shards[worker_id] + assigned_records = [ + r + for r in records + if type(r.get("worker_id")) is int and r["worker_id"] == worker_id + ] + counts = Counter( + r.get("status") if isinstance(r.get("status"), str) else None + for r in assigned_records + ) + writes = result.get("record_write_errors") + valid_writes = isinstance(writes, list) and all( + isinstance(error, Mapping) + and _valid_error(error) + and error.get("item_id") in {item.item_id for item in shard} + for error in writes + ) + failed_count = counts["failed"] + (len(writes) if valid_writes else 0) + checks = { + "schema": result.get("schema") == WORKER_SCHEMA, + "run_id": result.get("run_id") == run_id, + "manifest_sha256": result.get("manifest_sha256") + == spec.manifest_sha256, + "assigned_item_ids": result.get("assigned_item_ids") + == [item.item_id for item in shard], + "item_ids_sha256": result.get("item_ids_sha256") + == item_ids_sha256(shard), + "weight_bytes": type(result.get("weight_bytes")) is int + and result["weight_bytes"] + == sum(item.weight_bytes for item in shard), + "success_count": type(result.get("success_count")) is int + and result["success_count"] == counts["success"], + "failed_count": type(result.get("failed_count")) is int + and result["failed_count"] == failed_count, + "terminal_count": counts["success"] + failed_count == len(shard), + "record_write_errors": valid_writes and not writes, + "artifact_error": result.get("artifact_error") is None, + "error": result.get("error") is None, + "status": result.get("status") == "success" and not failed_count, + "worker_wall_sec": _finite_nonnegative( + result.get("worker_wall_sec") + ), + "provenance_worker_id": isinstance( + result.get("provenance"), Mapping + ) + and type(result["provenance"].get("worker_id")) is int + and result["provenance"]["worker_id"] == worker_id, + "timestamps": _valid_timestamps(result), + } + invalid = [key for key, ok in checks.items() if not ok] + if invalid: + errors.append( + { + "phase": "workers", + "worker_id": worker_id, + "message": ", ".join(invalid), + } + ) + return errors + + +def _validate_item_output( + item_dir: Path, metadata: Mapping[str, Any] +) -> None: + if not _real_directory(item_dir): + raise ValueError("item runner did not create a real output directory") + if Path(metadata.get("run_dir", "")) != item_dir: + raise ValueError("item runner reported a different output directory") + if Path(metadata.get("summary_path", "")) != item_dir / "summary.json": + raise ValueError("item runner reported a different summary path") + if not regular_file(item_dir / "summary.json"): + raise ValueError("item runner did not create a regular summary file") + + +def _real_directory(path: Path) -> bool: + try: + return stat.S_ISDIR(path.lstat().st_mode) + except OSError: + return False + + +def _finite_nonnegative(value: Any) -> bool: + return ( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(value) + and value >= 0 + ) + + +def _valid_timings(value: Any) -> bool: + return isinstance(value, Mapping) and all( + isinstance(key, str) and key and _finite_nonnegative(elapsed) + for key, elapsed in value.items() + ) + + +def _valid_error(value: Any) -> bool: + return isinstance(value, Mapping) and all( + isinstance(value.get(key), str) and value[key] + for key in ("type", "message") + ) + + +def _valid_timestamps(value: Mapping[str, Any]) -> bool: + try: + started = datetime.fromisoformat(value["started_at_utc"]) + completed = datetime.fromisoformat(value["completed_at_utc"]) + return ( + started.utcoffset() is not None + and completed.utcoffset() is not None + and started <= completed + ) + except (KeyError, TypeError, ValueError): + return False + + +def _same_json(left: Any, right: Any) -> bool: + try: + return json.dumps( + left, sort_keys=True, allow_nan=False + ) == json.dumps(right, sort_keys=True, allow_nan=False) + except (TypeError, ValueError): + return False diff --git a/src/cuphoton/core/executors.py b/src/cuphoton/core/executors.py new file mode 100644 index 00000000..3bdc06e8 --- /dev/null +++ b/src/cuphoton/core/executors.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Lazy selection of the distributed runtime for component workloads.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .execution import ExecutionResult, WorkloadSpec + + +def run_workload( + *, + executor: str, + prepare_workload: Callable[[int], WorkloadSpec], + output_root: Path, + run_id: str | None = None, + **options: Any, +) -> ExecutionResult | None: + """Run one component's workload through the selected external launcher.""" + + if executor == "dragon": + from .dragon import run_dragon_work_items as run + elif executor == "mpi": + from .mpi import run_mpi_work_items as run + else: + raise ValueError("distributed executor must be dragon or mpi") + return run( + prepare_workload=prepare_workload, + output_root=output_root, + run_id=run_id, + **options, + ) diff --git a/src/cuphoton/core/mpi.py b/src/cuphoton/core/mpi.py new file mode 100644 index 00000000..a25d409c --- /dev/null +++ b/src/cuphoton/core/mpi.py @@ -0,0 +1,540 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Collective MPI execution of component-neutral, worker-local workloads.""" + +from __future__ import annotations + +import json +import math +import os +import secrets +import sys +import time +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any + +from ._mpi_runtime import ( + _load_mpi_api, + _mpi_context, + _mpi_failure_consensus, + _normalize_mpi_library_version, + _prebound_visibility, +) +from .benchmark import ( + BenchmarkOptions, + BenchmarkRound, + build_benchmark_report, +) +from .bulk import ( + WorkItem, + atomic_write_json, + error_payload, + json_mapping, + new_run_id, + partition_byte_balanced, + read_json_mapping, + timestamp_utc, + validate_identifier, +) +from .execution import ( + ExecutionResult, + Worker, + WorkloadSpec, + audit_worker_provenance, + execute_worker_round, + finalize_round, + prepare_run, + resolve_worker_factory, +) + + +def run_mpi_work_items( + *, + prepare_workload: Callable[[int], WorkloadSpec], + output_root: Path, + run_id: str | None = None, + rank_setup_timeout_sec: float = 600.0, + benchmark: BenchmarkOptions | None = None, + prepare_on_root: bool = False, +) -> ExecutionResult | None: + """Run under an external launcher with one prebound GPU per MPI rank. + + Startup, preflight, worker initialization and close failures are agreed + collectively. Ordinary mode runs once; benchmark mode reuses the + same worker object for every warmup and measured pass. The launcher owns + process failure detection and MPI shutdown. With ``prepare_on_root``, + only rank zero plans the workload; JSON launch descriptors are broadcast + to peers, while scientific validators and finalizers remain on rank zero. + """ + + start = time.perf_counter() + started_at = timestamp_utc() + startup_error = None + visibility = None + resolved_root = None + try: + if benchmark is not None and not isinstance( + benchmark, BenchmarkOptions + ): + raise TypeError("benchmark must be BenchmarkOptions or None") + if ( + isinstance(rank_setup_timeout_sec, bool) + or not isinstance(rank_setup_timeout_sec, (int, float)) + or not math.isfinite(rank_setup_timeout_sec) + or rank_setup_timeout_sec <= 0 + ): + raise ValueError("rank_setup_timeout_sec must be positive") + if run_id is not None: + validate_identifier(run_id, field="run_id") + if not callable(prepare_workload): + raise TypeError("prepare_workload must be callable") + visibility = _prebound_visibility(os.environ) + premature = sorted( + name + for name in ("cupy", "numba.cuda", "cuda.tile") + if name in sys.modules + ) + torch = sys.modules.get("torch") + if torch is not None and torch.cuda.is_initialized(): + premature.append("torch CUDA") + if premature: + raise RuntimeError( + "CUDA was initialized before MPI binding: " + + ", ".join(premature) + ) + resolved_root = Path(output_root).expanduser().resolve(strict=False) + except Exception as exc: + startup_error = error_payload(exc) + api = _load_mpi_api() + comm = api.comm + _mpi_failure_consensus(comm, "MPI rank-startup validation", startup_error) + context = _mpi_context(api.MPI, comm, os.environ) + assert visibility is not None and resolved_root is not None + spec = None + preflight_error = None + try: + if not prepare_on_root or context.rank == 0: + spec = prepare_workload(context.rank) + if not isinstance(spec, WorkloadSpec): + raise TypeError("prepare_workload must return WorkloadSpec") + if context.world_size > len(spec.items): + raise ValueError( + "MPI world size cannot exceed the item count" + ) + except Exception as exc: + preflight_error = error_payload(exc) + _mpi_failure_consensus(comm, "MPI workload preflight", preflight_error) + _agree_payload(comm, {"prepare_on_root": prepare_on_root}) + if prepare_on_root: + payload = comm.bcast( + spec.identity_payload() if context.rank == 0 else None, root=0 + ) + descriptor_error = None + try: + if context.rank != 0: + spec = WorkloadSpec( + items=tuple( + WorkItem.from_dict(item) for item in payload["items"] + ), + options_payload=payload["options"], + manifest_payload=payload["manifest"], + input_identity_payload=payload["input_identity"], + manifest_sha256=payload["manifest_sha256"], + backend=payload["backend"], + worker_factory=resolve_worker_factory( + payload["worker_factory"] + ), + ) + except Exception as exc: + descriptor_error = error_payload(exc) + _mpi_failure_consensus( + comm, "MPI workload descriptor", descriptor_error + ) + assert spec is not None + _agree_payload( + comm, + { + "workload": spec.identity_payload(), + "run_id": run_id, + "output_root": str(resolved_root), + "rank_setup_timeout_sec": rank_setup_timeout_sec, + "benchmark": benchmark.to_payload() if benchmark else None, + "mpi4py_version": api.mpi4py_version, + "mpi_library_version": _normalize_mpi_library_version( + api.library_version + ), + }, + ) + effective_run_id = comm.bcast( + (run_id or new_run_id("mpi-workload")) if context.rank == 0 else None, + root=0, + ) + validate_identifier(effective_run_id, field="run_id") + run_dir = resolved_root / effective_run_id + shards = partition_byte_balanced(spec.items, context.world_size) + worker: Worker | None = None + ownership = {"claimed": False} + rounds: list[dict[str, Any]] = [] + errors: list[dict[str, Any]] = [] + last_summary: dict[str, Any] | None = None + startup_ready_sec = None + phase = "setup" + try: + _prepare_directory( + comm, + context.rank, + run_dir, + effective_run_id, + spec, + rank_setup_timeout_sec, + ownership, + ) + init_error = None + provenance = None + try: + if context.rank == 0: + atomic_write_json( + run_dir / "execution-options.json", + { + "benchmark": benchmark.to_payload() + if benchmark + else None, + "rank_setup_timeout_sec": rank_setup_timeout_sec, + "mpi4py_version": api.mpi4py_version, + "mpi_library_version": api.library_version, + }, + overwrite=False, + ) + worker = spec.worker_factory(spec.options_payload) + if not callable( + getattr(worker, "run_item", None) + ) or not callable(getattr(worker, "close", None)): + raise TypeError("worker requires run_item and close methods") + provenance = { + "worker_id": context.rank, + "hostname": context.host, + "pid": os.getpid(), + "cuda_visible_devices": visibility, + "gpu": json_mapping( + worker.gpu_identity, field="worker GPU identity" + ), + "rank_context": context.to_dict(), + } + except Exception as exc: + init_error = error_payload(exc) + _mpi_failure_consensus(comm, "MPI worker initialization", init_error) + ready = comm.gather(provenance, root=0) + ready_error = None + if context.rank == 0: + try: + problems = audit_worker_provenance( + ready, + backend=spec.backend, + expected_worker_count=context.world_size, + ) + if problems: + raise ValueError(str(problems)) + atomic_write_json( + run_dir / "ready.json", + {"workers": ready}, + overwrite=False, + ) + except Exception as exc: + ready_error = error_payload(exc) + _mpi_failure_consensus(comm, "MPI worker readiness", ready_error) + startup_ready_sec = time.perf_counter() - start + assert worker is not None and provenance is not None + planned_rounds = ( + benchmark.rounds() + if benchmark + else (BenchmarkRound("measure", 0),) + ) + for planned in planned_rounds: + phase = planned.round_id + round_dir = ( + run_dir / "rounds" / planned.round_id + if benchmark + else run_dir + ) + round_run_id = ( + planned.run_id(effective_run_id) + if benchmark + else effective_run_id + ) + if benchmark: + _prepare_directory( + comm, + context.rank, + round_dir, + round_run_id, + spec, + rank_setup_timeout_sec, + ) + round_start = time.perf_counter() + released = comm.bcast( + planned.round_id if context.rank == 0 else None, root=0 + ) + release_sec = time.perf_counter() - round_start + execution_error = None + receipt = None + worker_start = time.perf_counter() + try: + if released != planned.round_id: + raise RuntimeError("MPI round release identity differs") + receipt = execute_worker_round( + worker, + items=shards[context.rank], + run_id=round_run_id, + run_dir=round_dir, + manifest_sha256=spec.manifest_sha256, + worker_id=context.rank, + backend=spec.backend, + provenance=provenance, + ) + except Exception as exc: + execution_error = error_payload(exc) + worker_wall = time.perf_counter() - worker_start + collection_start = time.perf_counter() + gathered = comm.gather( + { + "worker_result": receipt, + "worker_wall_sec": worker_wall, + "error": execution_error, + }, + root=0, + ) + collected_at = time.perf_counter() + decision = None + if context.rank == 0: + try: + execution_errors = [ + {"worker_id": rank, **value["error"]} + for rank, value in enumerate(gathered) + if value["error"] is not None + ] + worker_results = [ + value["worker_result"] + for value in gathered + if value["worker_result"] is not None + ] + last_summary = finalize_round( + round_dir, + round_run_id, + spec, + shards, + worker_results, + artifact_timeout_sec=min(rank_setup_timeout_sec, 0.01) + if execution_errors + else rank_setup_timeout_sec, + ) + if execution_errors: + last_summary["errors"].extend(execution_errors) + last_summary["status"] = "failed" + timing = { + "batch_wall_sec": collected_at - round_start, + "finalization_sec": last_summary["finalization_sec"], + "worker_wall_max_sec": max( + value["worker_wall_sec"] for value in gathered + ), + "coordinator_timings_sec": { + "release_sec": release_sec, + "collection_sec": collected_at - collection_start, + }, + } + last_summary.update(timing) + if benchmark: + atomic_write_json( + round_dir / "summary.json", last_summary + ) + rounds.append( + { + **planned.to_payload(), + "status": last_summary["status"], + "summary_path": str( + (round_dir / "summary.json").relative_to( + run_dir + ) + ), + **timing, + } + ) + decision = { + "status": last_summary["status"], + "error": None, + } + except Exception as exc: + decision = { + "status": "failed", + "error": error_payload(exc), + } + decision = comm.bcast(decision, root=0) + if decision["error"] is not None: + raise RuntimeError(str(decision["error"])) + if decision["status"] != "success": + break + except Exception as exc: + errors.append({"phase": phase, **error_payload(exc)}) + finally: + close_start = time.perf_counter() + close_error = None + if worker is not None: + try: + worker.close() + except Exception as exc: + close_error = error_payload(exc) + closed = comm.gather( + {"worker_id": context.rank, "error": close_error}, root=0 + ) + close_sec = time.perf_counter() - close_start + if context.rank == 0: + errors.extend( + {"phase": "close", **record} + for record in closed + if record["error"] is not None + ) + + result = None + decision = None + if context.rank == 0: + try: + if not ownership["claimed"]: + raise RuntimeError( + "MPI run directory was not claimed: " + + "; ".join( + f"{error['type']}: {error['message']}" + for error in errors + ) + ) + if benchmark: + report = build_benchmark_report( + benchmark, rounds, errors=errors + ) + summary = { + "schema": "cuphoton.core.execution-benchmark-summary/v1", + "status": report["status"], + "benchmark": report, + "errors": errors, + } + else: + summary = dict( + last_summary + or { + "schema": "cuphoton.core.execution-summary/v1", + "status": "failed", + "errors": [], + } + ) + summary["errors"].extend(errors) + if summary["errors"]: + summary["status"] = "failed" + summary.update( + executor="mpi", + run_id=effective_run_id, + manifest_sha256=spec.manifest_sha256, + world_size=context.world_size, + mpi4py_version=api.mpi4py_version, + mpi_library_version=api.library_version, + started_at_utc=started_at, + completed_at_utc=timestamp_utc(), + startup_ready_sec=startup_ready_sec, + close_sec=close_sec, + coordinator_wall_sec=time.perf_counter() - start, + launcher_exit_note=( + "Collective completion does not prove launcher exit." + ), + ) + atomic_write_json(run_dir / "summary.json", summary) + result = ExecutionResult( + "mpi", + effective_run_id, + run_dir, + run_dir / "summary.json", + summary["status"], + summary, + ) + decision = {"status": result.status, "error": None} + except Exception as exc: + decision = {"status": "failed", "error": error_payload(exc)} + decision = comm.bcast(decision, root=0) + if decision["error"] is not None: + raise RuntimeError( + "cannot persist MPI execution aggregate: " + + str(decision["error"]) + ) + if context.rank != 0 and decision["status"] != "success": + raise RuntimeError(f"MPI execution {effective_run_id!r} failed") + return result + + +def _agree_payload(comm: Any, payload: Mapping[str, Any]) -> None: + gathered = comm.gather(payload, root=0) + error = None + if int(comm.Get_rank()) == 0: + try: + identities = { + json.dumps(value, sort_keys=True, allow_nan=False) + for value in gathered + } + if len(gathered) != int(comm.Get_size()) or len(identities) != 1: + raise ValueError( + "MPI workload, run options, or benchmark plans differ" + ) + except Exception as exc: + error = error_payload(exc) + _mpi_failure_consensus(comm, "MPI workload agreement", error) + + +def _prepare_directory( + comm: Any, + rank: int, + path: Path, + run_id: str, + spec: WorkloadSpec, + timeout: float, + ownership: dict[str, bool] | None = None, +) -> None: + state = None + if rank == 0: + try: + path.mkdir(parents=True, exist_ok=False) + if ownership is not None: + ownership["claimed"] = True + prepare_run(path, run_id, spec, "mpi", directory_claimed=True) + nonce = secrets.token_hex(32) + ready = { + "run_id": run_id, + "manifest_sha256": spec.manifest_sha256, + "run_dir": str(path), + "nonce": nonce, + } + atomic_write_json(path / ".ready.json", ready, overwrite=False) + state = {"error": None, "ready": ready} + except Exception as exc: + state = {"error": error_payload(exc), "ready": None} + state = comm.bcast(state, root=0) + if state["error"] is not None: + raise RuntimeError("cannot prepare MPI run: " + str(state["error"])) + shared_error = None + deadline = time.monotonic() + timeout + while True: + try: + ready = read_json_mapping(path / ".ready.json") + if ready != state["ready"]: + raise RuntimeError( + "MPI shared run directory identity differs" + ) + break + except OSError as exc: + if time.monotonic() >= deadline: + shared_error = error_payload(exc) + break + time.sleep(min(0.05, max(0.0, deadline - time.monotonic()))) + except Exception as exc: + shared_error = error_payload(exc) + break + _mpi_failure_consensus( + comm, "MPI shared run-directory validation", shared_error + ) diff --git a/src/cuphoton/xpois/dragon.py b/src/cuphoton/xpois/dragon.py index 608b828f..f3422d33 100644 --- a/src/cuphoton/xpois/dragon.py +++ b/src/cuphoton/xpois/dragon.py @@ -19,7 +19,6 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime -from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any @@ -47,6 +46,18 @@ from cuphoton.core.bulk import ( regular_file as _regular_file, ) +from cuphoton.core.dragon import ( + _distribution_version, + _DragonAPI, + _hostnames_match, + _load_dragon_api, + _select_gpu_placements, + _singleton_cuda_visibility, + _stable_gpu_physical_ids, + _strict_integer, + _valid_shard_provenance, + discover_gpu_placements, +) from .batch import ( BatchFitOptions, @@ -89,16 +100,6 @@ def to_dict(self) -> dict[str, Any]: } -@dataclass(frozen=True) -class _DragonAPI: - System: Any - Node: Any - Policy: Any - ProcessGroup: Any - ProcessTemplate: Any - Queue: Any - - def _resolve_worker_target( reference: Mapping[str, Any], ) -> Callable[..., None]: @@ -1502,78 +1503,6 @@ def _validate_worker_placement( return actual_host, visibility -def discover_gpu_placements( - system_type: Callable[[], Any], - node_type: Callable[[Any], Any], - *, - node_ids: Sequence[Any] | None = None, -) -> tuple[Placement, ...]: - """Enumerate actual Dragon Node.gpus IDs in allocation order.""" - - if node_ids is None: - node_ids = tuple(system_type().nodes) - placements: list[Placement] = [] - seen: set[tuple[str, int]] = set() - for node_id in node_ids: - node = node_type(node_id) - host = str(node.hostname) - for gpu_id in node.gpus or []: - if isinstance(gpu_id, bool) or not isinstance(gpu_id, int): - raise RuntimeError( - f"Dragon node {host!r} reported a non-integer GPU ID" - ) - key = (host, gpu_id) - if key in seen: - raise RuntimeError( - f"duplicate Dragon GPU placement: {host}:{gpu_id}" - ) - seen.add(key) - placements.append( - Placement( - worker_id=len(placements), - host=host, - gpu_id=gpu_id, - ) - ) - return tuple(placements) - - -def _select_gpu_placements( - placements: Sequence[Placement], worker_count: int -) -> tuple[Placement, ...]: - """Select GPUs round-robin across hosts, then assign dense worker IDs.""" - - if isinstance(worker_count, bool) or not isinstance(worker_count, int): - raise TypeError("worker_count must be an integer") - if worker_count <= 0 or worker_count > len(placements): - raise ValueError("worker_count exceeds available Dragon placements") - by_host: dict[str, list[Placement]] = {} - for placement in placements: - by_host.setdefault(placement.host, []).append(placement) - selected: list[Placement] = [] - offset = 0 - while len(selected) < worker_count: - progressed = False - for host_placements in by_host.values(): - if offset >= len(host_placements): - continue - selected.append(host_placements[offset]) - progressed = True - if len(selected) == worker_count: - break - if not progressed: - raise RuntimeError("could not select requested Dragon placements") - offset += 1 - return tuple( - Placement( - worker_id=worker_id, - host=placement.host, - gpu_id=placement.gpu_id, - ) - for worker_id, placement in enumerate(selected) - ) - - def _dragon_shard_worker( run_id: str, run_dir_raw: str, @@ -1824,54 +1753,6 @@ def _validate_success_item_output( raise ValueError("item runner did not create a regular summary file") -def _singleton_cuda_visibility(expected_gpu_id: int | None = None) -> str: - raw = os.environ.get("CUDA_VISIBLE_DEVICES") - if raw is None: - raise RuntimeError("Dragon worker has no CUDA_VISIBLE_DEVICES") - tokens = [token.strip() for token in raw.split(",")] - if ( - any(not token for token in tokens) - or len(tokens) != 1 - or tokens[0] == "-1" - ): - raise RuntimeError( - f"Dragon worker must see exactly one CUDA device, got {raw!r}" - ) - if expected_gpu_id is not None and tokens[0] != str(expected_gpu_id): - raise RuntimeError( - "Dragon worker GPU placement mismatch: requested " - f"{expected_gpu_id}, got CUDA_VISIBLE_DEVICES={raw!r}" - ) - return tokens[0] - - -def _hostnames_match( - requested: str, - actual: str, - *, - allow_loopback_alias: bool = False, -) -> bool: - """Accept exact or equivalent short/FQDN scheduler hostnames.""" - - requested_normalized = requested.rstrip(".").lower() - actual_normalized = actual.rstrip(".").lower() - if allow_loopback_alias and requested_normalized in { - "localhost", - "localhost.localdomain", - }: - # Dragon 0.14.1 reports ``localhost`` for its single-node system - # descriptor even though socket.gethostname() exposes the machine - # hostname inside the launched worker. - return bool(actual_normalized) - if requested_normalized == actual_normalized: - return True - if "." not in requested_normalized: - return actual_normalized.startswith(requested_normalized + ".") - if "." not in actual_normalized: - return requested_normalized.startswith(actual_normalized + ".") - return False - - def _load_terminal_records( run_dir: Path, visible_expected: Sequence[Path] = (), @@ -2452,113 +2333,11 @@ def _parse_timezone_aware_timestamp(value: Any) -> datetime | None: return parsed -def _stable_gpu_physical_ids( - value: Any, -) -> frozenset[tuple[str, str]] | None: - if ( - not isinstance(value, Mapping) - or "identity_error" not in value - or value["identity_error"] is not None - ): - return None - identities = frozenset( - (field, identifier.strip().casefold()) - for field in ("uuid", "pci_bus_id") - if isinstance((identifier := value.get(field)), str) - and identifier.strip() - ) - return identities or None - - -def _valid_shard_provenance( - value: Any, - *, - placement: Placement, - allow_loopback_alias: bool, - backend: str, -) -> bool: - if not isinstance(value, Mapping): - return False - worker_id = _strict_integer(value.get("worker_id")) - requested_gpu_id = _strict_integer(value.get("requested_gpu_id")) - pid = _strict_integer(value.get("pid")) - requested_host = value.get("requested_host") - hostname = value.get("hostname") - visibility = value.get("cuda_visible_devices") - if ( - worker_id != placement.worker_id - or requested_host != placement.host - or requested_gpu_id != placement.gpu_id - or pid is None - or pid <= 0 - or not isinstance(hostname, str) - or not hostname - or not _hostnames_match( - requested_host, - hostname, - allow_loopback_alias=allow_loopback_alias, - ) - or not isinstance(visibility, str) - or [token.strip() for token in visibility.split(",") if token.strip()] - != [str(placement.gpu_id)] - ): - return False - gpu = value.get("gpu") - if not isinstance(gpu, Mapping) or not gpu: - return False - gpu_backend = gpu.get("backend") - expected_identity_backend = ( - "cupy" if backend in {"cupy", "cutile"} else backend - ) - return ( - isinstance(gpu_backend, str) - and bool(gpu_backend) - and gpu_backend == expected_identity_backend - and _stable_gpu_physical_ids(gpu) is not None - ) - - -def _strict_integer(value: Any) -> int | None: - if isinstance(value, bool) or not isinstance(value, int): - return None - return value - - def _shard_result_sort_key(result: Mapping[str, Any]) -> tuple[int, int]: worker_id = _strict_integer(result.get("worker_id")) return (worker_id is None, worker_id if worker_id is not None else 0) -def _load_dragon_api() -> _DragonAPI: - try: - from dragon.infrastructure.policy import Policy - from dragon.native.machine import Node, System - from dragon.native.process import ProcessTemplate - from dragon.native.process_group import ProcessGroup - from dragon.native.queue import Queue - except (ImportError, OSError) as exc: - raise RuntimeError( - "The Dragon executor requires the dragonhpc runtime; " - "install it in this Python environment on every node " - "and launch with dragon" - ) from exc - return _DragonAPI( - System=System, - Node=Node, - Policy=Policy, - ProcessGroup=ProcessGroup, - ProcessTemplate=ProcessTemplate, - Queue=Queue, - ) - - -def _distribution_version(name: str) -> str | None: - try: - return version(name) - except PackageNotFoundError: - return None - - __all__ = [ "DragonBatchResult", "discover_gpu_placements", diff --git a/src/cuphoton/xpois/mpi.py b/src/cuphoton/xpois/mpi.py index ae490f86..7b0ec8be 100644 --- a/src/cuphoton/xpois/mpi.py +++ b/src/cuphoton/xpois/mpi.py @@ -10,7 +10,6 @@ import json import math import os -import re import secrets import socket import stat @@ -18,12 +17,59 @@ import time from collections import Counter from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict, dataclass +from dataclasses import dataclass from datetime import datetime -from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any +from cuphoton.core._mpi_runtime import ( + _MPIAPI as _MPIAPI, +) +from cuphoton.core._mpi_runtime import ( + _distribution_version as _distribution_version, +) +from cuphoton.core._mpi_runtime import ( + _environment_integer as _environment_integer, +) +from cuphoton.core._mpi_runtime import ( + _integer as _integer, +) +from cuphoton.core._mpi_runtime import ( + _launcher_context as _launcher_context, +) +from cuphoton.core._mpi_runtime import ( + _load_mpi_api as _load_mpi_api, +) +from cuphoton.core._mpi_runtime import ( + _local_rank as _local_rank, +) +from cuphoton.core._mpi_runtime import ( + _mpi_context as _mpi_context, +) +from cuphoton.core._mpi_runtime import ( + _mpi_failure_consensus as _mpi_failure_consensus, +) +from cuphoton.core._mpi_runtime import ( + _normalize_mpi_library_version as _normalize_mpi_library_version, +) +from cuphoton.core._mpi_runtime import ( + _prebound_visibility as _prebound_visibility, +) +from cuphoton.core._mpi_runtime import ( + _RankContext as _RankContext, +) +from cuphoton.core._mpi_runtime import ( + _runtime_integer as _runtime_integer, +) +from cuphoton.core._mpi_runtime import ( + _valid_error as _valid_error, +) +from cuphoton.core._mpi_runtime import ( + _validate_topology as _validate_topology, +) +from cuphoton.core._mpi_runtime import ( + _visible_tokens as _visible_tokens, +) from cuphoton.core.benchmark import BenchmarkOptions, build_benchmark_report from cuphoton.core.bulk import ( WorkItem, @@ -60,7 +106,6 @@ _GPU_BACKENDS = frozenset({"cupy", "cutile", "numba-cuda"}) _ALLOCATED_DEVICES = "CUPHOTON_ALLOCATED_CUDA_VISIBLE_DEVICES" _CUDA_MODULES = ("cupy", "numba.cuda", "cuda.tile") -_UNSIGNED = re.compile(r"[0-9]{1,18}") _POLL_INITIAL_SEC = 0.05 _POLL_MAX_SEC = 1.0 @@ -87,27 +132,6 @@ def to_dict(self) -> dict[str, Any]: } -@dataclass(frozen=True) -class _RankContext: - rank: int - local_rank: int - world_size: int - host: str - launcher: str - launch_id: str | None = None - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -@dataclass(frozen=True) -class _MPIAPI: - MPI: Any - comm: Any - mpi4py_version: str | None - library_version: str | None - - class _FileOwnershipError(RuntimeError): """Raised when a rank does not own the requested file-run state.""" @@ -590,70 +614,6 @@ def _run_mpi_benchmark( return result -def _mpi_failure_consensus( - comm: Any, - phase: str, - error: Mapping[str, Any] | None, - *, - root_failure_handler: Callable[[Mapping[str, str]], None] | None = None, -) -> None: - rank = int(comm.Get_rank()) - size = int(comm.Get_size()) - gathered = comm.gather( - {"rank": rank, "error": dict(error) if error else None}, - root=0, - ) - decision = None - if rank == 0: - failures: list[str] = [] - if not isinstance(gathered, Sequence) or len(gathered) != size: - failures.append("invalid consensus result") - else: - for expected_rank, item in enumerate(gathered): - if not isinstance(item, Mapping): - failures.append( - f"rank {expected_rank}: invalid consensus result" - ) - continue - if _integer(item.get("rank")) != expected_rank: - failures.append( - f"rank {expected_rank}: invalid rank identity" - ) - item_error = item.get("error") - if item_error is None: - continue - if not _valid_error(item_error): - failures.append( - f"rank {expected_rank}: invalid error payload" - ) - continue - failures.append( - f"rank {expected_rank}: {item_error['type']}: " - f"{item_error['message']}" - ) - message = ( - f"{phase} failed: " + "; ".join(failures) if failures else None - ) - if message is not None and root_failure_handler is not None: - try: - root_failure_handler( - {"type": "RuntimeError", "message": message} - ) - except Exception as exc: - persistence_error = error_payload(exc) - message += ( - "; failed to persist terminal setup evidence: " - f"{persistence_error['type']}: " - f"{persistence_error['message']}" - ) - decision = {"error": message} - decision = comm.bcast(decision, root=0) - if not isinstance(decision, Mapping): - raise RuntimeError(f"{phase} decision was not a mapping") - if decision.get("error"): - raise RuntimeError(str(decision["error"])) - - def _validate_arguments( aggregation_mode: str, rank_setup_timeout_sec: float, @@ -2397,13 +2357,6 @@ def _parse_timestamp(value: Any) -> datetime | None: return parsed -def _valid_error(value: Any) -> bool: - return isinstance(value, Mapping) and all( - isinstance(value.get(field), str) and bool(value[field]) - for field in ("type", "message") - ) - - def _same_json_value(left: Any, right: Any) -> bool: try: return json.dumps( @@ -3222,57 +3175,6 @@ def _publish_completed_file_ranks( return records, errors -def _mpi_context( - mpi: Any, comm: Any, environ: Mapping[str, str] -) -> _RankContext: - rank = _runtime_integer(comm.Get_rank(), "MPI rank") - size = _runtime_integer(comm.Get_size(), "MPI world size") - _validate_topology(rank, size) - context = None - failure = None - try: - try: - shared = comm.Split_type(mpi.COMM_TYPE_SHARED, key=rank) - try: - local_rank = _runtime_integer( - shared.Get_rank(), "MPI local rank" - ) - finally: - shared.Free() - except Exception as exc: - raise RuntimeError("cannot determine MPI local rank") from exc - launcher = _launcher_context(environ) - if launcher is not None and launcher[1:] != (rank, size): - raise RuntimeError("launcher topology disagrees with mpi4py") - launcher_local_rank = _local_rank( - environ, launcher[0] if launcher is not None else None - ) - if ( - launcher_local_rank is not None - and launcher_local_rank != local_rank - ): - raise RuntimeError( - "launcher local rank disagrees with MPI shared communicator" - ) - context = _RankContext( - rank, - local_rank, - size, - socket.gethostname(), - launcher[0] if launcher is not None else "mpi", - ) - except Exception as exc: - failure = error_payload(exc) - if failure is None and context is None: - failure = error_payload( - RuntimeError("rank context validation produced no context") - ) - _mpi_failure_consensus(comm, "MPI rank-context validation", failure) - if context is None: - raise RuntimeError("MPI rank-context consensus returned no context") - return context - - def _environment_context(environ: Mapping[str, str]) -> _RankContext: launcher = _launcher_context(environ) if launcher is None: @@ -3335,152 +3237,6 @@ def _file_launch_id(environ: Mapping[str, str], launcher: str) -> str: return hashlib.sha256(json.dumps(parts).encode()).hexdigest() -def _launcher_context( - environ: Mapping[str, str], -) -> tuple[str, int, int] | None: - candidates = ( - ("openmpi", "OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), - ("pmi", "PMI_RANK", "PMI_SIZE"), - ("slurm", "SLURM_PROCID", "SLURM_NTASKS"), - ) - for name, rank_key, size_key in candidates: - present = rank_key in environ, size_key in environ - if all(present): - return ( - name, - _environment_integer(environ, rank_key), - _environment_integer(environ, size_key), - ) - if any(present): - raise RuntimeError( - f"{name} launcher rank and size must appear together" - ) - return None - - -def _local_rank( - environ: Mapping[str, str], launcher: str | None = None -) -> int | None: - keys_by_launcher = { - "openmpi": ("OMPI_COMM_WORLD_LOCAL_RANK",), - "pmi": ("PMI_LOCAL_RANK", "MPI_LOCALRANKID", "PMIX_LOCAL_RANK"), - "slurm": ( - "SLURM_LOCALID", - "PMIX_LOCAL_RANK", - "MPI_LOCALRANKID", - "PMI_LOCAL_RANK", - ), - } - keys = keys_by_launcher.get( - launcher, - ( - "OMPI_COMM_WORLD_LOCAL_RANK", - "PMIX_LOCAL_RANK", - "MPI_LOCALRANKID", - "PMI_LOCAL_RANK", - "SLURM_LOCALID", - ), - ) - values = [ - (key, _environment_integer(environ, key)) - for key in keys - if key in environ - ] - if launcher == "pmi" and not values and "SLURM_LOCALID" in environ: - slurm_keys = ("SLURM_PROCID", "SLURM_NTASKS", "SLURM_LOCALID") - topology_keys = ("PMI_RANK", "PMI_SIZE", *slurm_keys) - if all(key in environ for key in topology_keys): - pmi_topology = ( - _environment_integer(environ, "PMI_RANK"), - _environment_integer(environ, "PMI_SIZE"), - ) - slurm_topology = ( - _environment_integer(environ, "SLURM_PROCID"), - _environment_integer(environ, "SLURM_NTASKS"), - ) - if pmi_topology != slurm_topology: - raise RuntimeError( - "PMI and Slurm launcher rank topology disagree" - ) - values = [ - ( - "SLURM_LOCALID", - _environment_integer(environ, "SLURM_LOCALID"), - ) - ] - if not values: - return None - if any(value != values[0][1] for _, value in values[1:]): - family = launcher or "MPI" - raise RuntimeError(f"{family} launcher local-rank variables disagree") - return values[0][1] - - -def _prebound_visibility(environ: Mapping[str, str]) -> str: - raw_visibility = environ.get("CUDA_VISIBLE_DEVICES") - if raw_visibility is None: - raise RuntimeError( - "CUDA_VISIBLE_DEVICES is unset; bind one device per MPI rank " - "with cuphoton-openmpi-rank-exec or scheduler GPU binding" - ) - tokens = _visible_tokens(raw_visibility) - if len(tokens) != 1: - raise RuntimeError( - "MPI rank must see one CUDA device before mpi4py import; use " - "cuphoton-openmpi-rank-exec or scheduler GPU binding" - ) - allocated_raw = environ.get(_ALLOCATED_DEVICES) - if allocated_raw is None: - return tokens[0] - allocated = _visible_tokens(allocated_raw) - local_rank = _local_rank(environ, "openmpi") - if local_rank is None or local_rank >= len(allocated): - raise RuntimeError( - f"{_ALLOCATED_DEVICES} requires a valid Open MPI local rank; " - "unset it when launching without cuphoton-openmpi-rank-exec" - ) - if allocated[local_rank] != tokens[0]: - raise RuntimeError("pre-bound CUDA device does not match allocation") - return tokens[0] - - -def _visible_tokens(raw: Any) -> tuple[str, ...]: - if not isinstance(raw, str) or not raw.strip() or raw.strip() == "-1": - raise RuntimeError("CUDA_VISIBLE_DEVICES exposes no devices") - tokens = tuple(token.strip() for token in raw.split(",")) - if any(not token or token.startswith("-") for token in tokens) or len( - tokens - ) != len(set(tokens)): - raise RuntimeError("CUDA_VISIBLE_DEVICES is invalid") - return tokens - - -def _environment_integer(environ: Mapping[str, str], key: str) -> int: - raw = environ[key] - if not isinstance(raw, str) or not _UNSIGNED.fullmatch(raw): - raise RuntimeError(f"launcher variable {key} must be unsigned") - return int(raw) - - -def _runtime_integer(value: Any, field: str) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value < 0: - raise RuntimeError(f"{field} must be non-negative") - return value - - -def _validate_topology(rank: int, size: int) -> None: - if size <= 0 or rank >= size: - raise RuntimeError("invalid launcher rank topology") - - -def _integer(value: Any) -> int | None: - return ( - value - if isinstance(value, int) and not isinstance(value, bool) - else None - ) - - def _rank_result_sort_key(result: Mapping[str, Any]) -> tuple[bool, int]: rank = _integer(result.get("rank")) return (rank is None, rank if rank is not None else 0) @@ -3852,39 +3608,4 @@ def _aggregate_error( pass -def _load_mpi_api() -> _MPIAPI: - try: - from mpi4py import MPI - except (ImportError, OSError, RuntimeError) as exc: - raise RuntimeError( - "MPI collective aggregation requires mpi4py built against the " - "allocation's MPI runtime" - ) from exc - try: - library_version = _normalize_mpi_library_version( - MPI.Get_library_version() - ) - except Exception: # pragma: no cover - implementation-specific - library_version = None - return _MPIAPI( - MPI, - MPI.COMM_WORLD, - _distribution_version("mpi4py"), - library_version, - ) - - -def _normalize_mpi_library_version(value: Any) -> str | None: - if value is None: - return None - return str(value).rstrip("\x00 \t\r\n") - - -def _distribution_version(name: str) -> str | None: - try: - return version(name) - except PackageNotFoundError: - return None - - __all__ = ["MPIBatchResult", "run_mpi_image_pair_batch"] diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py new file mode 100644 index 00000000..d57239a6 --- /dev/null +++ b/tests/core/test_dragon.py @@ -0,0 +1,833 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Exercise the persistent coordinator with threads and ordinary file IO.""" + +from __future__ import annotations + +import json +import queue +import sys +import threading +import time +from types import SimpleNamespace + +import pytest + +from cuphoton.core import dragon +from cuphoton.core.benchmark import BenchmarkOptions +from cuphoton.core.bulk import WorkItem, atomic_write_json +from cuphoton.core.execution import WorkloadSpec + +_STATE = None + + +class _Worker: + def __init__(self, options): + self.options = options + self.placement = _STATE.local.placement + self.calls = 0 + _STATE.factories.append(self.placement.worker_id) + _STATE.instances.append(self) + self.gpu_identity = { + "backend": options["backend"], + "device_index": 0, + "name": "fake-gpu", + "uuid": f"GPU-{self.placement.worker_id}", + "pci_bus_id": f"0000:{self.placement.gpu_id:02x}:00.0", + "identity_error": None, + } + + def run_item(self, item, output_dir): + self.calls += 1 + _STATE.calls.append( + ( + self.placement.worker_id, + item.item_id, + output_dir, + threading.get_ident(), + ) + ) + if _STATE.item_failure: + raise ValueError("synthetic scientific failure") + output_dir.mkdir() + atomic_write_json( + output_dir / "summary.json", {"value": item.payload["value"]} + ) + return { + "run_dir": str(output_dir), + "summary_path": str(output_dir / "summary.json"), + "requested_backend": self.options["backend"], + "backend": self.options["backend"], + "device": "fake-gpu", + "runtime": {}, + "timings_sec": {"solve": 0.1}, + "wall_sec": {}, + } + + def close(self): + _STATE.closed.append(self.placement.worker_id) + if _STATE.close_failure: + raise RuntimeError("synthetic worker close failure") + + +def _factory(options): + return _Worker(options) + + +def _install_runtime( + monkeypatch, + *, + mutate=None, + item_failure=False, + close_failure=False, + interleave_close=False, +): + global _STATE + state = _STATE = SimpleNamespace( + local=threading.local(), + groups=[], + queues=[], + factories=[], + instances=[], + calls=[], + closed=[], + bindings=[], + item_failure=item_failure, + close_failure=close_failure, + fast_result=threading.Event(), + fast_waiting_for_close=threading.Event(), + release_slow_result=threading.Event(), + ) + + class Queue(queue.Queue): + def __init__(self, maxsize=0, policy=None): + super().__init__(maxsize=maxsize) + self.policy = policy + self.closed = False + state.queues.append(self) + + def put(self, value, block=True, timeout=None): + values = ( + mutate(value) + if mutate is not None and self.policy is None + else [value] + ) + for message in values: + if ( + interleave_close + and self.policy is None + and message["kind"] == "round" + and message["worker_id"] == 1 + ): + assert state.release_slow_result.wait(timeout=5) + super().put(message, block=block, timeout=timeout) + if interleave_close and self.policy is None: + if ( + message["kind"] == "round" + and message["worker_id"] == 0 + ): + state.fast_result.set() + elif ( + message["kind"] == "closed" + and message["worker_id"] == 0 + ): + state.release_slow_result.set() + + def get(self, block=True, timeout=None): + if ( + interleave_close + and self.policy is not None + and self.policy.gpu_affinity == [3] + and state.fast_result.is_set() + ): + state.fast_waiting_for_close.set() + state.release_slow_result.set() + return super().get(block=block, timeout=timeout) + + def close(self): + self.closed = True + + class Policy: + Placement = SimpleNamespace(HOST_NAME="host-name") + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class Template: + def __init__(self, target, args, policy): + self.target, self.args, self.policy = target, args, policy + self.argdata = repr(args).encode() + + class Group: + def __init__(self, **kwargs): + self.options = kwargs + self.templates, self.threads, self.inactive_puids = [], [], [] + self.stopped = self.closed = False + state.groups.append(self) + + def add_process(self, *, nproc, template): + assert nproc == 1 + self.templates.append(template) + + def init(self): + pass + + def start(self): + def invoke(index, template): + try: + state.local.puid = 1000 + index + template.target(*template.args) + except Exception: + self.inactive_puids.append((1000 + index, 1)) + else: + self.inactive_puids.append((1000 + index, 0)) + + for index, template in enumerate(self.templates): + thread = threading.Thread( + target=invoke, args=(index, template) + ) + self.threads.append(thread) + thread.start() + + def join(self, timeout): + deadline = time.monotonic() + timeout + for thread in self.threads: + thread.join(max(0, deadline - time.monotonic())) + if any(thread.is_alive() for thread in self.threads): + raise TimeoutError("group join timed out") + + def stop(self, patience): + self.stopped = True + for template in self.templates: + try: + template.args[4].put_nowait(None) + except queue.Full: + pass + for thread in self.threads: + thread.join(patience) + assert not any(thread.is_alive() for thread in self.threads) + + def close(self, patience): + self.closed = True + + def binding(placement, *, allow_loopback_alias): + assert allow_loopback_alias + state.local.placement = placement + state.bindings.append(placement.worker_id) + return placement.host, str(placement.gpu_id) + + monkeypatch.setattr(dragon, "_validate_binding", binding) + monkeypatch.setattr(dragon, "_current_puid", lambda: state.local.puid) + monkeypatch.setattr( + dragon, + "_load_dragon_api", + lambda: dragon._DragonAPI( + System=lambda: SimpleNamespace(nodes=(1,)), + Node=lambda node_id: SimpleNamespace( + hostname="fake-node.example", gpus=[3, 7] + ), + Policy=Policy, + ProcessGroup=Group, + ProcessTemplate=Template, + Queue=Queue, + ), + ) + return state + + +def _run( + tmp_path, + *, + benchmark=None, + worker_count=2, + worker_timeout=10, + result_timeout=0.2, + large=False, + validator=None, + finalizer=None, +): + spec = WorkloadSpec( + items=( + WorkItem("one", {"value": "x" * 300_000 if large else 1}, 12), + WorkItem("two", {"value": 2}, 6), + ), + options_payload={"backend": "cupy"}, + manifest_payload={"schema": "example"}, + input_identity_payload={"schema": "example-identity"}, + manifest_sha256="0" * 64, + backend="cupy", + worker_factory=_factory, + success_record_validator=validator, + finalize_round=finalizer, + ) + return dragon.run_dragon_work_items( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + run_id="run", + max_workers=worker_count, + result_timeout_sec=result_timeout, + worker_timeout_sec=worker_timeout, + benchmark=benchmark, + ) + + +def test_persistent_workers_use_bounded_descriptors_and_all_rounds( + monkeypatch, tmp_path +): + state = _install_runtime(monkeypatch) + validated, finalized = [], [] + + def validator(record, round_dir): + validated.append((record["item_id"], round_dir.name)) + return () + + def finalizer(round_dir, records): + finalized.append(round_dir.name) + return {"item_ids": [record["item_id"] for record in records]} + + result = _run( + tmp_path, + benchmark=BenchmarkOptions(1, 2), + large=True, + validator=validator, + finalizer=finalizer, + ) + + assert result.status == "success", result.summary + assert sorted(state.factories) == sorted(state.bindings) == [0, 1] + assert sorted(state.closed) == [0, 1] + assert len(state.instances) == 2 + assert all(worker.calls == 3 for worker in state.instances) + assert len(state.groups) == 1 + assert state.groups[0].closed and not state.groups[0].stopped + assert all(channel.closed for channel in state.queues) + assert all( + channel.policy.host_name == "fake-node.example" + for channel in state.queues + if channel.policy + ) + expected = ["warmup-0000", "measure-0000", "measure-0001"] + assert finalized == expected + assert len(validated) == 6 + for worker_id in range(2): + calls = [call for call in state.calls if call[0] == worker_id] + assert [call[2].parent.parent.name for call in calls] == expected + assert len({call[3] for call in calls}) == 1 + for template in state.groups[0].templates: + assert len(template.args) == 6 + assert len(template.argdata) < 2000 + descriptor = json.loads( + (result.run_dir / "launch/worker-0000.json").read_text() + ) + assert len(descriptor["items"][0]["payload"]["value"]) == 300_000 + report = result.summary["benchmark"] + assert report["measured_batch_wall_sec"] is not None + run_ids = set() + for receipt in report["rounds"]: + summary = json.loads( + (result.run_dir / receipt["summary_path"]).read_text() + ) + assert summary["terminal_record_audit"]["ok"] + assert summary["finalization_sec"] == receipt["finalization_sec"] + assert receipt["finalization_sec"] >= 0 + assert summary["result"] == {"item_ids": ["one", "two"]} + run_ids.add(summary["run_id"]) + assert len(run_ids) == 3 + + +def test_ordinary_workload_retains_root_artifacts(monkeypatch, tmp_path): + state = _install_runtime(monkeypatch) + result = _run(tmp_path) + assert result.status == "success", result.summary + assert "benchmark" not in result.summary + assert (result.run_dir / "records/one.json").is_file() + assert not (result.run_dir / "rounds").exists() + assert all(worker.calls == 1 for worker in state.instances) + assert sorted(state.closed) == [0, 1] + + +@pytest.mark.parametrize( + "failure", + [ + "wrong-round", + "duplicate-ready", + "missing-round", + "invalid-duration", + "changed-provenance", + "scientific", + "close", + ], +) +def test_failures_keep_evidence_and_invalidate_aggregates( + monkeypatch, tmp_path, failure +): + def mutate(message): + if failure == "duplicate-ready" and message["kind"] == "ready": + return [message, message] + if message["kind"] == "round": + if failure == "wrong-round": + return [{**message, "round_id": "measure-9999"}] + if failure == "missing-round": + return [] + if failure == "invalid-duration": + return [{**message, "worker_wall_sec": True}] + if failure == "changed-provenance": + result = dict(message["result"]) + result["provenance"] = {**result["provenance"], "pid": 99999} + return [{**message, "result": result}] + return [message] + + state = _install_runtime( + monkeypatch, + mutate=mutate, + item_failure=failure == "scientific", + close_failure=failure == "close", + ) + result = _run( + tmp_path, + benchmark=BenchmarkOptions(0, 1), + worker_count=1, + worker_timeout=0.5 if failure == "missing-round" else 10, + ) + assert result.status == "failed" + assert result.summary["benchmark"]["measured_batch_wall_sec"] is None + assert result.summary["lifecycle_errors"] + assert state.groups[0].stopped and state.groups[0].closed + assert all(channel.closed for channel in state.queues) + assert (result.run_dir / "summary.json").is_file() + if failure == "close": + assert result.summary["benchmark"]["rounds"][0]["status"] == "success" + assert ( + result.summary["closed_messages"][0]["error"]["message"] + == "synthetic worker close failure" + ) + + +def test_descriptor_tampering_fails_before_worker_factory( + monkeypatch, tmp_path +): + state = _install_runtime(monkeypatch) + read_descriptor = dragon._read_descriptor + + def invalid_digest(run_id, descriptor_path, digest, context): + return read_descriptor(run_id, descriptor_path, "f" * 64, context) + + monkeypatch.setattr(dragon, "_read_descriptor", invalid_digest) + result = _run(tmp_path, worker_count=1) + assert result.status == "failed" + assert not state.factories + assert not state.calls + startup = result.summary["ready_messages"][0] + assert startup["status"] == "failed" + assert "SHA-256" in startup["error"]["message"] + assert (result.run_dir / "startup/worker-0000.json").is_file() + assert state.groups[0].closed + + +def test_later_failure_keeps_successful_warmup(monkeypatch, tmp_path): + def corrupt_measured(message): + if ( + message["kind"] == "round" + and message["round_id"] == "measure-0000" + ): + return [{**message, "round_id": "warmup-0000"}] + return [message] + + _install_runtime(monkeypatch, mutate=corrupt_measured) + result = _run(tmp_path, benchmark=BenchmarkOptions(1, 2), worker_count=1) + assert result.status == "failed" + rounds = result.summary["benchmark"]["rounds"] + assert [entry["round_id"] for entry in rounds] == [ + "warmup-0000", + "measure-0000", + ] + assert [entry["status"] for entry in rounds] == ["success", "failed"] + assert (result.run_dir / "rounds/warmup-0000/summary.json").is_file() + assert result.summary["benchmark"]["measured_batch_wall_sec"] is None + + +def test_scientific_validator_failure_suppresses_component_finalizer( + monkeypatch, tmp_path +): + _install_runtime(monkeypatch) + finalized = [] + result = _run( + tmp_path, + worker_count=1, + validator=lambda record, path: ("scientific receipt differs",), + finalizer=lambda path, records: finalized.append(path) or {}, + ) + assert result.status == "failed" + assert not finalized + assert any( + "scientific receipt differs" in error["message"] + for error in result.summary["errors"] + ) + + +def test_fast_worker_waits_for_close_until_all_round_results_arrive( + monkeypatch, tmp_path +): + state = _install_runtime(monkeypatch, interleave_close=True) + result = _run(tmp_path) + + assert result.status == "success", result.summary + assert state.fast_waiting_for_close.is_set() + assert sorted(state.closed) == [0, 1] + assert len(result.summary["messages"]) == 2 + assert len(result.summary["closed_messages"]) == 2 + assert state.groups[0].closed and not state.groups[0].stopped + + +@pytest.mark.parametrize( + "kind, exit_code", + [ + ("ready", -9), + ("round", -9), + ("closed", -9), + ("ready", 0), + ("round", 0), + ("closed", 0), + ], +) +def test_collector_detects_native_exit_on_first_bounded_poll(kind, exit_code): + timeouts = [] + + class EmptyQueue: + def get(self, *, timeout): + timeouts.append(timeout) + raise queue.Empty + + with pytest.raises(RuntimeError, match="worker exited"): + dragon._collect_messages( + EmptyQueue(), + [], + worker_count=1, + run_id="run", + kind=kind, + round_id=None, + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(123, exit_code)]), + ) + assert len(timeouts) == 2 + assert 0 < timeouts[0] <= 1 + assert timeouts[1] == 0 + + +def test_collector_keeps_queued_results_before_reporting_peer_exit(): + receipt = { + "kind": "ready", + "run_id": "run", + "round_id": None, + "worker_id": 0, + "puid": 1000, + "status": "success", + } + + class ResultsQueue: + calls = 0 + + def get(self, *, timeout): + self.calls += 1 + if self.calls == 1: + return receipt + raise queue.Empty + + messages = [] + with pytest.raises(RuntimeError, match="worker exited"): + dragon._collect_messages( + ResultsQueue(), + messages, + worker_count=2, + run_id="run", + kind="ready", + round_id=None, + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(123, -9)]), + ) + assert messages == [receipt] + + +def test_closed_collection_allows_normal_exit_before_delayed_receipt(): + receipt = { + "kind": "closed", + "run_id": "run", + "round_id": None, + "worker_id": 0, + "puid": 1000, + "status": "success", + } + + class DelayedQueue: + calls = 0 + + def get(self, *, timeout): + self.calls += 1 + if self.calls == 1: + raise queue.Empty + return receipt + + channel = DelayedQueue() + messages = [] + dragon._collect_messages( + channel, + messages, + worker_count=1, + run_id="run", + kind="closed", + round_id=None, + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(123, 0)]), + ) + assert channel.calls == 2 + assert messages == [receipt] + + +def test_native_worker_crash_stops_group_without_ready_receipt( + monkeypatch, tmp_path +): + state = _install_runtime(monkeypatch) + + def crash(*args): + raise RuntimeError("synthetic native crash") + + monkeypatch.setattr(dragon, "_workload_worker", crash) + result = _run(tmp_path, worker_count=1) + assert result.status == "failed" + assert not result.summary["ready_messages"] + assert any( + "worker exited before ready" in error["message"] + for error in result.summary["lifecycle_errors"] + ) + assert state.groups[0].stopped and state.groups[0].closed + + +@pytest.mark.parametrize("close_failure", [False, True]) +def test_terminal_summary_is_absent_until_cleanup( + monkeypatch, tmp_path, close_failure +): + state = _install_runtime(monkeypatch, close_failure=close_failure) + close = _Worker.close + group_type = dragon._load_dragon_api().ProcessGroup + join = group_type.join + group_close = group_type.close + observations = [] + + def inspect_close(worker): + observations.append( + ("worker_close", (tmp_path / "run/summary.json").exists()) + ) + return close(worker) + + def inspect_join(group, timeout): + observations.append( + ("join", (tmp_path / "run/summary.json").exists()) + ) + return join(group, timeout) + + def inspect_group_close(group, patience): + observations.append( + ("group_close", (tmp_path / "run/summary.json").exists()) + ) + return group_close(group, patience) + + monkeypatch.setattr(_Worker, "close", inspect_close) + monkeypatch.setattr(group_type, "join", inspect_join) + monkeypatch.setattr(group_type, "close", inspect_group_close) + result = _run(tmp_path) + assert not any(exists for _, exists in observations) + assert [phase for phase, _ in observations].count("worker_close") == 2 + assert ("join", False) in observations or close_failure + assert ("group_close", False) in observations + assert result.status == ("failed" if close_failure else "success") + assert len(result.summary["closed_messages"]) == 2 + assert all(not thread.is_alive() for thread in state.groups[0].threads) + assert json.loads(result.summary_path.read_text()) == result.summary + + +def test_failed_round_collects_slow_peer_after_failed_worker_closes( + monkeypatch, tmp_path +): + failed_closed = threading.Event() + + def observe(message): + if message["kind"] == "closed" and message["worker_id"] == 0: + failed_closed.set() + return [message] + + state = _install_runtime(monkeypatch, mutate=observe) + run_item = _Worker.run_item + + def run_with_failure(worker, item, output_dir): + if worker.placement.worker_id == 0: + raise ValueError("first worker failed") + assert failed_closed.wait(timeout=2) + return run_item(worker, item, output_dir) + + monkeypatch.setattr(_Worker, "run_item", run_with_failure) + result = _run(tmp_path, benchmark=BenchmarkOptions(0, 2)) + assert result.status == "failed" + assert len(result.summary["benchmark"]["rounds"]) == 1 + summary = json.loads( + (result.run_dir / "rounds/measure-0000/summary.json").read_text() + ) + assert len(summary["worker_results"]) == 2 + assert summary["terminal_record_audit"]["ok"] + assert len(summary["terminal_record_audit"]["failed_item_ids"]) == 1 + assert sorted(record["status"] for record in summary["records"]) == [ + "failed", + "success", + ] + assert all(not thread.is_alive() for thread in state.groups[0].threads) + assert not any( + error["phase"] == "stop_after_failure" + for error in result.summary["lifecycle_errors"] + ) + + +def test_collector_drains_receipt_after_exit_snapshot_before_failing(): + receipt = { + "kind": "ready", + "run_id": "run", + "round_id": None, + "worker_id": 0, + "puid": 123, + "status": "failed", + } + timeouts = [] + + class RacingQueue: + def get(self, *, timeout): + timeouts.append(timeout) + if len(timeouts) == 1: + raise queue.Empty + return receipt + + messages = [] + with pytest.raises(RuntimeError, match="ready failed"): + dragon._collect_messages( + RacingQueue(), + messages, + worker_count=1, + run_id="run", + kind="ready", + round_id=None, + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(123, 1)]), + ) + assert messages == [receipt] + assert timeouts[1] == 0 + + +@pytest.mark.parametrize("kind", ["ready", "round", "closed"]) +def test_failed_receipt_does_not_abort_slow_peer_on_reported_exit(kind): + first = { + "kind": kind, + "run_id": "run", + "round_id": None, + "worker_id": 0, + "puid": 1000, + "status": "failed", + } + peer = {**first, "worker_id": 1, "puid": 1001, "status": "failed"} + sequence = iter((first, queue.Empty, queue.Empty, peer)) + + class DelayedQueue: + def get(self, *, timeout): + value = next(sequence) + if value is queue.Empty: + raise queue.Empty + return value + + messages = [] + with pytest.raises(RuntimeError, match=kind + " failed"): + dragon._collect_messages( + DelayedQueue(), + messages, + worker_count=2, + run_id="run", + kind=kind, + round_id=None, + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(1000, 1)]), + ) + assert messages == [first, peer] + + +def test_known_round_failure_uses_short_artifact_timeout( + monkeypatch, tmp_path +): + _install_runtime( + monkeypatch, + mutate=lambda message: ( + [{**message, "status": "failed"}] + if message["kind"] == "round" + else [message] + ), + ) + finalize = dragon.finalize_round + observed = [] + + def inspect_finalize(*args, **kwargs): + observed.append(kwargs["artifact_timeout_sec"]) + return finalize(*args, **kwargs) + + monkeypatch.setattr(dragon, "finalize_round", inspect_finalize) + result = _run(tmp_path, worker_count=1, result_timeout=30) + assert result.status == "failed" + assert observed == [0.01] + + +@pytest.mark.parametrize("reject_terminal_put", [False, True]) +def test_interrupted_worker_keeps_failed_close_evidence( + monkeypatch, tmp_path, reject_terminal_put +): + state = _install_runtime(monkeypatch) + _run(tmp_path, worker_count=1) + args = list(state.groups[0].templates[0].args) + state.local.puid = 1000 + + class InterruptedCommands: + def get(self, **kwargs): + raise KeyboardInterrupt("synthetic interrupt") + + class Results: + def put(self, message, **kwargs): + if reject_terminal_put and message["kind"] == "closed": + raise queue.Full + + args[4:] = [InterruptedCommands(), Results()] + with pytest.raises(KeyboardInterrupt, match="synthetic interrupt"): + dragon._workload_worker(*args) + closed = json.loads( + (tmp_path / "run/startup/worker-0000-closed.json").read_text() + ) + assert closed["status"] == "failed" + assert closed["error"]["type"] == "KeyboardInterrupt" + + +@pytest.mark.parametrize("initialized", [False, True]) +def test_binding_allows_torch_import_without_cuda(monkeypatch, initialized): + for name in ("cupy", "numba.cuda", "cuda.tile"): + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.setitem( + sys.modules, + "torch", + SimpleNamespace( + cuda=SimpleNamespace(is_initialized=lambda: initialized) + ), + ) + monkeypatch.setattr( + dragon.socket, "gethostname", lambda: "worker.example" + ) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3") + placement = dragon.Placement(0, "worker.example", 3) + if initialized: + with pytest.raises(RuntimeError, match="torch CUDA"): + dragon._validate_binding(placement, allow_loopback_alias=False) + else: + assert dragon._validate_binding( + placement, allow_loopback_alias=False + ) == ("worker.example", "3") diff --git a/tests/core/test_execution.py b/tests/core/test_execution.py new file mode 100644 index 00000000..1ef62180 --- /dev/null +++ b/tests/core/test_execution.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +from dataclasses import replace + +import pytest + +from cuphoton.core import execution +from cuphoton.core.bulk import WorkItem, atomic_write_json, read_json_mapping + + +class FakeWorker: + def __init__(self, options): + self.options = options + self.closed = False + self.gpu_identity = { + "backend": options.get("backend", "cupy"), + "device_index": 0, + "uuid": options.get("uuid", "GPU-0"), + "identity_error": None, + } + + def run_item(self, item, output_dir): + if item.payload.get("fail"): + raise ValueError("scientific failure") + output_dir.mkdir() + atomic_write_json(output_dir / "summary.json", item.payload) + return { + "run_dir": str(output_dir), + "summary_path": str(output_dir / "summary.json"), + "backend": self.options.get("backend", "cupy"), + "device": "cuda:0", + "runtime": {}, + "timings_sec": {"work": 0.0}, + "wall_sec": {"item": 0.0}, + } + + def close(self): + self.closed = True + + +def worker_factory(options): + return FakeWorker(options) + + +def workload(**overrides): + values = { + "items": ( + WorkItem("one", {"value": 1}, 1), + WorkItem("two", {"value": 2}, 2), + ), + "options_payload": {}, + "manifest_payload": {"items": ["one", "two"]}, + "input_identity_payload": {}, + "manifest_sha256": "a" * 64, + "backend": "cupy", + "worker_factory": worker_factory, + } + values.update(overrides) + return execution.WorkloadSpec(**values) + + +def provenance(worker_id, gpu): + return { + "worker_id": worker_id, + "hostname": "host", + "pid": os.getpid(), + "cuda_visible_devices": str(worker_id), + "gpu": gpu, + } + + +def stage(tmp_path, spec=None): + spec = spec or workload() + execution.prepare_run(tmp_path, "run", spec, "test") + shards = ((spec.items[1],), (spec.items[0],)) + results = [] + for worker_id, shard in enumerate(shards): + worker = FakeWorker( + {"uuid": f"GPU-{worker_id}", "backend": spec.backend} + ) + results.append( + execution.execute_worker_round( + worker, + items=shard, + run_id="run", + run_dir=tmp_path, + manifest_sha256=spec.manifest_sha256, + worker_id=worker_id, + backend=spec.backend, + provenance=provenance(worker_id, worker.gpu_identity), + ) + ) + return spec, shards, results + + +def test_round_audits_then_finalizes_in_manifest_order(monkeypatch, tmp_path): + calls = [] + clock = [100.0] + + def finalize(run_dir, records): + clock[0] += 2.5 + calls.append([record["item_id"] for record in records]) + return {"ordered": calls[-1]} + + spec, shards, results = stage( + tmp_path / "run", + workload( + items=( + WorkItem("two", {"value": 2}, 2), + WorkItem("one", {"value": 1}, 1), + ), + finalize_round=finalize, + ), + ) + monkeypatch.setattr(execution.time, "perf_counter", lambda: clock[0]) + report = execution.finalize_round( + tmp_path / "run", + "run", + spec, + shards, + results, + artifact_timeout_sec=0.01, + ) + assert report["status"] == "success" + assert report["terminal_record_audit"]["ok"] + assert calls == [["two", "one"]] + assert report["result"] == {"ordered": ["two", "one"]} + assert not (tmp_path / "run" / "summary.json").exists() + assert report["finalization_sec"] == 2.5 + + +@pytest.mark.parametrize( + "tamper", + [ + "record", + "record-path", + "receipt", + "missing", + "extra", + "gpu", + "round", + "callback", + ], +) +def test_round_audit_rejects_invalid_evidence_before_component_merge( + tmp_path, tamper +): + finalized = [] + spec, shards, results = stage( + tmp_path / "run", + workload( + finalize_round=lambda directory, records: ( + finalized.append(True) or {} + ) + ), + ) + run_dir = tmp_path / "run" + if tamper == "record": + path = run_dir / "records" / "one.json" + record = read_json_mapping(path) + record["worker_id"] = 0 + atomic_write_json(path, record) + elif tamper == "record-path": + first = run_dir / "records" / "one.json" + second = run_dir / "records" / "two.json" + first_record = read_json_mapping(first) + atomic_write_json(first, read_json_mapping(second)) + atomic_write_json(second, first_record) + elif tamper == "receipt": + results[0]["success_count"] = 99 + elif tamper == "missing": + (run_dir / "records" / "one.json").unlink() + elif tamper == "extra": + atomic_write_json(run_dir / "records" / "extra.json", {}) + elif tamper == "gpu": + results[1]["provenance"]["gpu"]["uuid"] = "GPU-0" + atomic_write_json( + run_dir / "workers" / "worker-0001.json", results[1] + ) + elif tamper == "round": + results[0]["run_id"] = "other-round" + atomic_write_json( + run_dir / "workers" / "worker-0000.json", results[0] + ) + elif tamper == "callback": + spec = replace( + spec, + success_record_validator=lambda record, directory: [ + "bad science" + ], + ) + report = execution.finalize_round( + run_dir, "run", spec, shards, results, artifact_timeout_sec=0.01 + ) + assert report["status"] == "failed" + assert report["errors"] + assert report["result"] is None + assert not finalized + + +def test_failed_item_and_publication_are_retained(tmp_path, monkeypatch): + spec = workload( + items=(WorkItem("one", {"fail": True}), WorkItem("two", {})) + ) + original = execution.atomic_write_json + + def write(path, payload, **kwargs): + if path.name == "two.json": + raise OSError("record disk failure") + return original(path, payload, **kwargs) + + monkeypatch.setattr(execution, "atomic_write_json", write) + spec, shards, results = stage(tmp_path / "run", spec) + report = execution.finalize_round( + tmp_path / "run", + "run", + spec, + shards, + results, + artifact_timeout_sec=0.01, + ) + assert report["status"] == "failed" + assert "scientific failure" in json.dumps(report) + assert "record disk failure" in json.dumps(report) + + +def test_torch_worker_does_not_require_cupy(tmp_path): + spec, shards, results = stage(tmp_path / "run", workload(backend="torch")) + report = execution.finalize_round( + tmp_path / "run", + "run", + spec, + shards, + results, + artifact_timeout_sec=0.01, + ) + assert report["status"] == "success" + + +def test_factory_reference_is_importable_and_payload_is_json(): + spec = workload() + reference = execution.factory_reference(worker_factory) + assert execution.resolve_worker_factory(reference) is worker_factory + assert ( + json.loads(json.dumps(spec.identity_payload())) + == spec.identity_payload() + ) + with pytest.raises(ValueError, match="package-importable"): + workload(worker_factory=lambda options: FakeWorker(options)) + + +def test_run_directory_is_never_reused(tmp_path): + run_dir = tmp_path / "run" + execution.prepare_run(run_dir, "run", workload(), "test") + with pytest.raises(FileExistsError): + execution.prepare_run(run_dir, "run", workload(), "test") + + +@pytest.mark.parametrize( + "hosts, duplicate", + [ + (("node.alpha.example", "node.beta.example"), False), + (("node", "node.alpha.example"), True), + (("NODE.alpha.example.", "node.alpha.example"), True), + ], +) +def test_provenance_host_equivalence_preserves_domains(hosts, duplicate): + identities = [] + for index, host in enumerate(hosts): + value = provenance( + index, + { + "backend": "cupy", + "device_index": 0, + "identity_error": None, + "pci_bus_id": "0000:01:00.0", + }, + ) + identities.append({**value, "hostname": host}) + errors = execution.audit_worker_provenance( + identities, backend="cupy", expected_worker_count=2 + ) + assert bool(errors) is duplicate diff --git a/tests/core/test_executor_options.py b/tests/core/test_executor_options.py new file mode 100644 index 00000000..2fb17834 --- /dev/null +++ b/tests/core/test_executor_options.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from cuphoton.core.cli import CommandError +from cuphoton.core.cli.executor import ExecutorOptions + + +@pytest.mark.parametrize( + ("executor", "option"), + [ + ("local", "warmup_rounds"), + ("local", "measure_rounds"), + ("local", "max_workers"), + ("local", "rank_setup_timeout_sec"), + ("mpi", "max_workers"), + ("mpi", "worker_timeout_sec"), + ("mpi", "result_timeout_sec"), + ("dragon", "rank_setup_timeout_sec"), + ], +) +def test_rejects_flags_which_selected_executor_cannot_honor(executor, option): + command = ExecutorOptions() + command.executor = executor + setattr(command, option, 1) + with pytest.raises(CommandError, match=option.replace("_", "-")): + command.executor_options() + + +def test_local_default_does_not_enable_distributed_execution(): + command = ExecutorOptions() + command.executor = "local" + assert command.executor_options() == {} + + +@pytest.mark.parametrize("executor", ["dragon", "mpi"]) +def test_round_controls_are_opt_in(executor): + command = ExecutorOptions() + command.executor = executor + assert command.executor_options() == {"benchmark": None} + command.warmup_rounds = 0 + assert command.executor_options()["benchmark"].to_payload() == { + "warmup_rounds": 0, + "measure_rounds": 1, + } diff --git a/tests/core/test_mpi_execution.py b/tests/core/test_mpi_execution.py new file mode 100644 index 00000000..f22a6af5 --- /dev/null +++ b/tests/core/test_mpi_execution.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier, local +from types import SimpleNamespace + +import pytest + +from cuphoton.core import mpi +from cuphoton.core._mpi_runtime import _MPIAPI +from cuphoton.core.benchmark import BenchmarkOptions +from cuphoton.core.bulk import read_json_mapping +from cuphoton.core.execution import ExecutionResult + +from .test_execution import FakeWorker, workload + +_STATE = local() + + +class TrackingWorker(FakeWorker): + def __init__(self, options): + self.rank = _STATE.rank + self.trace = _STATE.trace + self.trace.append(("initialize", self.rank)) + if options["failure"] == "initialize" and self.rank == 1: + raise ValueError("worker initialization failed") + super().__init__({**options, "uuid": f"GPU-{self.rank}"}) + + def run_item(self, item, output_dir): + round_id = output_dir.parent.parent.name + self.trace.append(("execute", self.rank, round_id, id(self))) + if self.options["failure"] == "item" and self.rank == 1: + raise ValueError("item execution failed") + return super().run_item(item, output_dir) + + def close(self): + self.trace.append(("close", self.rank)) + if self.options["failure"] == "close" and self.rank == 1: + raise ValueError("worker close failed") + super().close() + + +def tracking_factory(options): + return TrackingWorker(options) + + +@pytest.mark.parametrize( + "failure", + [None, "preflight", "plan", "initialize", "item", "close", "unexpected"], +) +@pytest.mark.parametrize("prepare_on_root", [False, True]) +def test_two_rank_persistent_lifecycle_and_collective_failures( + monkeypatch, tmp_path, failure, prepare_on_root +): + for name in tuple(os.environ): + if name.startswith( + ( + "OMPI_", + "PMI_", + "PMIX_", + "SLURM_", + "MPI_LOCAL", + "CUPHOTON_ALLOCATED", + ) + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + for name in ("cupy", "numba.cuda", "cuda.tile", "torch"): + monkeypatch.delitem(sys.modules, name, raising=False) + barrier = Barrier(2, timeout=5) + shared = SimpleNamespace(gathered=[None, None], broadcast=None) + trace = [] + + class Comm: + def __init__(self, rank): + self.rank = rank + + def Get_rank(self): + return self.rank + + def Get_size(self): + return 2 + + def Split_type(self, split_type, key): + return SimpleNamespace( + Get_rank=lambda: self.rank, Free=lambda: None + ) + + def gather(self, value, root): + shared.gathered[self.rank] = value + barrier.wait() + result = list(shared.gathered) if self.rank == root else None + barrier.wait() + return result + + def bcast(self, value, root): + if self.rank == root: + shared.broadcast = value + barrier.wait() + result = shared.broadcast + barrier.wait() + return result + + monkeypatch.setattr( + mpi, + "_load_mpi_api", + lambda: _MPIAPI( + SimpleNamespace(COMM_TYPE_SHARED=1), + Comm(_STATE.rank), + "test-mpi4py", + "test-MPI", + ), + ) + if failure == "unexpected": + execute = mpi.execute_worker_round + + def fail_outside_item(*args, **kwargs): + if _STATE.rank == 1: + raise RuntimeError("unexpected worker loop failure") + return execute(*args, **kwargs) + + monkeypatch.setattr(mpi, "execute_worker_round", fail_outside_item) + + def prepare(rank): + trace.append(("preflight", rank)) + if failure == "preflight" and rank == (0 if prepare_on_root else 1): + raise ValueError("bad rank input") + return workload( + worker_factory=tracking_factory, + options_payload={"failure": failure}, + ) + + def run(rank): + _STATE.rank = rank + _STATE.trace = trace + try: + return mpi.run_mpi_work_items( + prepare_workload=prepare, + prepare_on_root=prepare_on_root, + output_root=tmp_path, + run_id="parallel", + rank_setup_timeout_sec=0.01, + benchmark=None + if failure == "plan" and rank == 1 + else BenchmarkOptions(warmup_rounds=1, measure_rounds=2), + ) + except Exception as exc: + return exc + + with ThreadPoolExecutor(max_workers=2) as pool: + root, peer = list(pool.map(run, (0, 1))) + assert sorted(event[1] for event in trace if event[0] == "preflight") == ( + [0] if prepare_on_root else [0, 1] + ) + if failure in {"preflight", "plan"}: + assert isinstance(root, RuntimeError) + assert isinstance(peer, RuntimeError) + assert not (tmp_path / "parallel").exists() + assert not any(event[0] == "initialize" for event in trace) + return + assert isinstance(root, ExecutionResult), repr(root) + assert sorted(event for event in trace if event[0] == "initialize") == [ + ("initialize", 0), + ("initialize", 1), + ] + for rank in (0, 1): + instance_ids = { + event[3] + for event in trace + if event[0] == "execute" and event[1] == rank + } + assert len(instance_ids) <= 1 + closed = sorted(event[1] for event in trace if event[0] == "close") + assert closed == ([0] if failure == "initialize" else [0, 1]) + report = root.summary["benchmark"] + assert read_json_mapping(root.summary_path) == root.summary + if failure is None: + assert root.status == "success" + assert peer is None + assert len(report["rounds"]) == 3 + assert all( + receipt["finalization_sec"] == 0 for receipt in report["rounds"] + ) + assert report["measured_batch_wall_sec"] is not None + assert len([event for event in trace if event[0] == "execute"]) == 6 + else: + assert root.status == "failed" + assert isinstance(peer, RuntimeError) + assert report["measured_batch_wall_sec"] is None + assert report["errors"] + if failure == "close": + assert len(report["rounds"]) == 3 + assert all( + round_result["status"] == "success" + for round_result in report["rounds"] + ) + elif failure == "initialize": + assert report["rounds"] == [] + else: + assert len(report["rounds"]) == 1 + assert report["rounds"][0]["status"] == "failed" + + +def test_ordinary_mode_and_rank_binding_before_worker_factory( + monkeypatch, tmp_path +): + class Comm: + def Get_rank(self): + return 0 + + def Get_size(self): + return 1 + + def Split_type(self, split_type, key): + return SimpleNamespace(Get_rank=lambda: 0, Free=lambda: None) + + def gather(self, value, root): + return [value] + + def bcast(self, value, root): + return value + + for name in tuple(os.environ): + if name.startswith( + ( + "OMPI_", + "PMI_", + "PMIX_", + "SLURM_", + "MPI_LOCAL", + "CUPHOTON_ALLOCATED", + ) + ): + monkeypatch.delenv(name, raising=False) + for name in ("cupy", "numba.cuda", "cuda.tile", "torch"): + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + comm = Comm() + monkeypatch.setattr( + mpi, + "_load_mpi_api", + lambda: _MPIAPI( + SimpleNamespace(COMM_TYPE_SHARED=1), comm, "test", "test" + ), + ) + _STATE.rank = 0 + _STATE.trace = [] + close = TrackingWorker.close + observations = [] + + def inspect_close(worker): + observations.append((tmp_path / "once/summary.json").exists()) + return close(worker) + + monkeypatch.setattr(TrackingWorker, "close", inspect_close) + result = mpi.run_mpi_work_items( + prepare_workload=lambda rank: workload( + worker_factory=tracking_factory, options_payload={"failure": None} + ), + output_root=tmp_path, + run_id="once", + ) + assert result.status == "success" + assert observations == [False] + assert "benchmark" not in result.summary + assert not (result.run_dir / "rounds").exists() + assert ( + len([event for event in _STATE.trace if event[0] == "execute"]) == 2 + ) + assert _STATE.trace[-1] == ("close", 0) + with pytest.raises(RuntimeError, match="FileExistsError"): + mpi.run_mpi_work_items( + prepare_workload=lambda rank: workload(), + output_root=tmp_path, + run_id="once", + ) + assert read_json_mapping(result.summary_path) == result.summary + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + with pytest.raises(RuntimeError, match="one CUDA device"): + mpi.run_mpi_work_items( + prepare_workload=lambda rank: pytest.fail( + "preflight before binding" + ), + output_root=tmp_path, + run_id="unbound", + ) + assert not (tmp_path / "unbound").exists()