From e6b7015b0de6bbed1b6eda6a1e85e14c3a764f0d Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:23:29 -0700 Subject: [PATCH 01/10] Add shared whole-item worker execution contracts Signed-off-by: Trent Nelson --- src/cuphoton/core/execution.py | 862 +++++++++++++++++++++++++++++++++ tests/core/test_execution.py | 236 +++++++++ 2 files changed, 1098 insertions(+) create mode 100644 src/cuphoton/core/execution.py create mode 100644 tests/core/test_execution.py diff --git a/src/cuphoton/core/execution.py b/src/cuphoton/core/execution.py new file mode 100644 index 0000000..d6c3b39 --- /dev/null +++ b/src/cuphoton/core/execution.py @@ -0,0 +1,862 @@ +# 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, + 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: Mapping[str, Any] + + def run_item(self, item: WorkItem, output_dir: Path) -> Mapping[str, Any]: + """Execute one item and durably publish its ordinary output files.""" + + 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 = ( + host.rstrip(".").casefold().split(".")[0] + == other_host.rstrip(".").casefold().split(".")[0] + ) + 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 durable execution evidence, then run a component's merge step.""" + + 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 + if not errors and spec.finalize_round is not None: + 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)}) + 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, + } + atomic_write_json(run_dir / "summary.json", summary) + 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) + 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/tests/core/test_execution.py b/tests/core/test_execution.py new file mode 100644 index 0000000..81f92be --- /dev/null +++ b/tests/core/test_execution.py @@ -0,0 +1,236 @@ +# 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(tmp_path): + calls = [] + + def finalize(run_dir, records): + calls.append([record["item_id"] for record in records]) + return {"ordered": calls[-1]} + + spec, shards, results = stage( + tmp_path / "run", workload(finalize_round=finalize) + ) + 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 == [["one", "two"]] + assert report["result"] == {"ordered": ["one", "two"]} + assert read_json_mapping(tmp_path / "run" / "summary.json") == report + + +@pytest.mark.parametrize( + "tamper", + ["record", "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 == "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") From 8f9d9ae57ee3414e63c1f2034e4b0956619f3c90 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:25:39 -0700 Subject: [PATCH 02/10] Run persistent workloads through the shared Dragon executor Initialize workload factories after placement and retain their state across rounds. Use compact launch descriptors and shared artifact audits, then require worker cleanup and process exits before accepting results. Signed-off-by: Trent Nelson --- src/cuphoton/core/dragon.py | 1008 ++++++++++++++++++++++++++++++++++ src/cuphoton/xpois/dragon.py | 245 +-------- tests/core/test_dragon.py | 424 ++++++++++++++ 3 files changed, 1444 insertions(+), 233 deletions(-) create mode 100644 src/cuphoton/core/dragon.py create mode 100644 tests/core/test_dragon.py diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py new file mode 100644 index 0000000..5f79874 --- /dev/null +++ b/src/cuphoton/core/dragon.py @@ -0,0 +1,1008 @@ +# 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.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 + + +@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 _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_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. + """ + + 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=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, + ) + 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, + ) + 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=result_timeout_sec, + ) + round_timings["artifact_audit_sec"] = ( + time.perf_counter() - audit_start + ) + receipt = { + **round_spec.to_payload(), + "status": "success" + if terminal_summary["status"] == "success" + and not round_errors + else "failed", + "batch_wall_sec": batch_wall, + "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, + ) + 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" + _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, + ) + 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, +) -> None: + seen: set[int] = set() + 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=remaining), + field=f"Dragon {kind} result", + ) + except queue.Empty as exc: + raise TimeoutError(f"Dragon {kind} deadline expired") from exc + messages.append(message) + worker_id = _strict_integer(message.get("worker_id")) + 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 + ): + raise ValueError(f"invalid or duplicate Dragon {kind} identity") + seen.add(worker_id) + if message.get("status") != "success": + raise RuntimeError(f"Dragon worker {worker_id} {kind} failed") + if kind == "round": + 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") + + +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", "torch") + if name in sys.modules + ] + if premature: + raise RuntimeError( + "CUDA modules were imported before Dragon worker placement: " + + ", ".join(premature) + ) + return hostname, visibility + + +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"] + run_dir = Path(context["run_dir"]) + timeout = context["result_timeout_sec"] + ready: dict[str, Any] = { + "kind": "ready", + "run_id": run_id, + "round_id": None, + "worker_id": worker_id, + "status": "success", + } + failed = False + 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", + "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" + ) + except Exception: + failed = True + raise + finally: + closed = { + "kind": "closed", + "run_id": run_id, + "round_id": None, + "worker_id": worker_id, + "status": "failed" if failed else "success", + } + 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)) + results_queue.put(closed, timeout=timeout) + if closed["status"] != "success" and not failed: + raise RuntimeError("Dragon worker cleanup failed") diff --git a/src/cuphoton/xpois/dragon.py b/src/cuphoton/xpois/dragon.py index 608b828..f3422d3 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/tests/core/test_dragon.py b/tests/core/test_dragon.py new file mode 100644 index 0000000..b6f0bf3 --- /dev/null +++ b/tests/core/test_dragon.py @@ -0,0 +1,424 @@ +# 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 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 +): + global _STATE + state = _STATE = SimpleNamespace( + local=threading.local(), + groups=[], + queues=[], + factories=[], + instances=[], + calls=[], + closed=[], + bindings=[], + item_failure=item_failure, + close_failure=close_failure, + ) + + 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: + super().put(message, 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: + 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, + "_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, + 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=0.2, + 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["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"] + ) From ac1b1913e5166813d8f230fabe2e1986e41a1d63 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:29:54 -0700 Subject: [PATCH 03/10] Release Dragon worker cleanup after the final round audit Signed-off-by: Trent Nelson --- src/cuphoton/core/dragon.py | 17 ++++++++++++ tests/core/test_dragon.py | 53 ++++++++++++++++++++++++++++++++++++- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index 5f79874..3b41bcc 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -595,6 +595,17 @@ def run_dragon_work_items( 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, @@ -604,6 +615,7 @@ def run_dragon_work_items( round_id=None, deadline=time.monotonic() + result_timeout_sec, ) + timings["worker_close_sec"] = time.perf_counter() - close_start phase = "join" phase_start = time.perf_counter() try: @@ -980,6 +992,11 @@ def _workload_worker( 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 Exception: failed = True raise diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index b6f0bf3..5325d1d 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -76,7 +76,12 @@ def _factory(options): def _install_runtime( - monkeypatch, *, mutate=None, item_failure=False, close_failure=False + monkeypatch, + *, + mutate=None, + item_failure=False, + close_failure=False, + interleave_close=False, ): global _STATE state = _STATE = SimpleNamespace( @@ -90,6 +95,9 @@ def _install_runtime( 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): @@ -106,7 +114,36 @@ def put(self, value, block=True, timeout=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 @@ -422,3 +459,17 @@ def test_scientific_validator_failure_suppresses_component_finalizer( "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 From a3d9e55201a07b6703962151cd7506a2c29053c2 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:30:45 -0700 Subject: [PATCH 04/10] Share executor selection and round options across commands Signed-off-by: Trent Nelson --- src/cuphoton/core/cli/executor.py | 117 ++++++++++++++++++++++++++++ src/cuphoton/core/executors.py | 38 +++++++++ tests/core/test_executor_options.py | 49 ++++++++++++ 3 files changed, 204 insertions(+) create mode 100644 src/cuphoton/core/cli/executor.py create mode 100644 src/cuphoton/core/executors.py create mode 100644 tests/core/test_executor_options.py diff --git a/src/cuphoton/core/cli/executor.py b/src/cuphoton/core/cli/executor.py new file mode 100644 index 0000000..92110ea --- /dev/null +++ b/src/cuphoton/core/cli/executor.py @@ -0,0 +1,117 @@ +# 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. [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 = "Persistent-worker warmup passes, with outputs retained." + _default = None + + class MeasureRoundsArg(PositiveIntegerInvariant): + _arg = "--measure-rounds" + _help = "Measured passes in persistent workers. [default: 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/executors.py b/src/cuphoton/core/executors.py new file mode 100644 index 0000000..3bdc06e --- /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/tests/core/test_executor_options.py b/tests/core/test_executor_options.py new file mode 100644 index 0000000..2fb1783 --- /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, + } From d1da8527a26cdef7f0bef9d8acfb7eeb8467d4fd Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:31:26 -0700 Subject: [PATCH 05/10] Execute component workloads in persistent MPI ranks Signed-off-by: Trent Nelson --- src/cuphoton/core/_mpi_runtime.py | 343 ++++++++++++++++++++ src/cuphoton/core/mpi.py | 498 ++++++++++++++++++++++++++++++ src/cuphoton/xpois/mpi.py | 377 +++------------------- tests/core/test_mpi_execution.py | 269 ++++++++++++++++ 4 files changed, 1159 insertions(+), 328 deletions(-) create mode 100644 src/cuphoton/core/_mpi_runtime.py create mode 100644 src/cuphoton/core/mpi.py create mode 100644 tests/core/test_mpi_execution.py diff --git a/src/cuphoton/core/_mpi_runtime.py b/src/cuphoton/core/_mpi_runtime.py new file mode 100644 index 0000000..9ae882b --- /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/mpi.py b/src/cuphoton/core/mpi.py new file mode 100644 index 0000000..35bceec --- /dev/null +++ b/src/cuphoton/core/mpi.py @@ -0,0 +1,498 @@ +# 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 ( + 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, +) + + +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, +) -> 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. + """ + + 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: + 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) + 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, + "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) + 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") + 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/mpi.py b/src/cuphoton/xpois/mpi.py index ae490f8..7b0ec8b 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_mpi_execution.py b/tests/core/test_mpi_execution.py new file mode 100644 index 0000000..2bd4c3d --- /dev/null +++ b/tests/core/test_mpi_execution.py @@ -0,0 +1,269 @@ +# 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"], +) +def test_two_rank_persistent_lifecycle_and_collective_failures( + monkeypatch, tmp_path, failure +): + 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 == 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, + 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))) + 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 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 = [] + 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 "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) + 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() From ec329c0b0ad03ae64938f1b6717c82324e221e0b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:34:32 -0700 Subject: [PATCH 06/10] Bind terminal item identities to their record paths Signed-off-by: Trent Nelson --- src/cuphoton/core/execution.py | 7 +++++++ tests/core/test_execution.py | 17 ++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/cuphoton/core/execution.py b/src/cuphoton/core/execution.py index d6c3b39..a85ae14 100644 --- a/src/cuphoton/core/execution.py +++ b/src/cuphoton/core/execution.py @@ -582,6 +582,13 @@ def _read_round_artifacts( 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/") diff --git a/tests/core/test_execution.py b/tests/core/test_execution.py index 81f92be..454ef5d 100644 --- a/tests/core/test_execution.py +++ b/tests/core/test_execution.py @@ -126,7 +126,16 @@ def finalize(run_dir, records): @pytest.mark.parametrize( "tamper", - ["record", "receipt", "missing", "extra", "gpu", "round", "callback"], + [ + "record", + "record-path", + "receipt", + "missing", + "extra", + "gpu", + "round", + "callback", + ], ) def test_round_audit_rejects_invalid_evidence_before_component_merge( tmp_path, tamper @@ -146,6 +155,12 @@ def test_round_audit_rejects_invalid_evidence_before_component_merge( 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": From 43e56dab7c8aa80cc316f32754da890ee00e1d93 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Wed, 23 Sep 2026 21:36:35 -0700 Subject: [PATCH 07/10] Detect Dragon worker exits while collecting results Signed-off-by: Trent Nelson --- src/cuphoton/core/dragon.py | 21 ++++++- tests/core/test_dragon.py | 110 ++++++++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 3 deletions(-) diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index 3b41bcc..9a884f3 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -50,6 +50,7 @@ _LAUNCH_SCHEMA = "cuphoton.core.dragon-launch/v1" _TEMPLATE_BUDGET_BYTES = 96 * 1024 +_WORKER_POLL_SEC = 0.5 @dataclass(frozen=True) @@ -447,6 +448,7 @@ def run_dragon_work_items( kind="ready", round_id=None, deadline=worker_deadline, + group=group, ) provenances = [message["provenance"] for message in ready] ready_errors = audit_worker_provenance( @@ -520,6 +522,7 @@ def run_dragon_work_items( kind="round", round_id=round_spec.round_id, deadline=worker_deadline, + group=group, ) finally: round_timings["collection_sec"] = ( @@ -614,6 +617,7 @@ def run_dragon_work_items( kind="closed", round_id=None, deadline=time.monotonic() + result_timeout_sec, + group=group, ) timings["worker_close_sec"] = time.perf_counter() - close_start phase = "join" @@ -762,6 +766,7 @@ def _collect_messages( kind: str, round_id: str | None, deadline: float, + group: Any, ) -> None: seen: set[int] = set() while len(seen) < worker_count: @@ -770,11 +775,21 @@ def _collect_messages( raise TimeoutError(f"Dragon {kind} deadline expired") try: message = json_mapping( - results_queue.get(timeout=remaining), + results_queue.get(timeout=min(remaining, _WORKER_POLL_SEC)), field=f"Dragon {kind} result", ) - except queue.Empty as exc: - raise TimeoutError(f"Dragon {kind} deadline expired") from exc + except queue.Empty: + failed_exits = [ + (puid, code) + for puid, code in group.inactive_puids + if code != 0 + ] + if failed_exits: + raise RuntimeError( + f"Dragon worker exited before {kind} completion: " + f"{failed_exits}" + ) + continue messages.append(message) worker_id = _strict_integer(message.get("worker_id")) if ( diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 5325d1d..887db2c 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -473,3 +473,113 @@ def test_fast_worker_waits_for_close_until_all_round_results_arrive( 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", ["ready", "round", "closed"]) +def test_collector_detects_native_exit_on_first_bounded_poll(kind): + 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() + 3600, + group=SimpleNamespace(inactive_puids=[(123, -9)]), + ) + assert len(timeouts) == 1 + assert 0 < timeouts[0] <= 1 + + +def test_collector_keeps_queued_results_before_reporting_peer_exit(): + receipt = { + "kind": "ready", + "run_id": "run", + "round_id": None, + "worker_id": 0, + "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() + 3600, + 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, + "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() + 3600, + 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 From 2984ab99382860d05be2ce3382980bb2ba91954b Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 14:13:50 -0700 Subject: [PATCH 08/10] Publish executor success only after worker shutdown Signed-off-by: Trent Nelson --- docs/components/core.md | 17 ++- src/cuphoton/core/benchmark.py | 6 +- src/cuphoton/core/bulk.py | 27 ++++ src/cuphoton/core/cli/executor.py | 14 +- src/cuphoton/core/dragon.py | 136 +++++++++++------ src/cuphoton/core/execution.py | 34 ++++- src/cuphoton/core/mpi.py | 16 +- tests/core/test_dragon.py | 238 +++++++++++++++++++++++++++++- tests/core/test_execution.py | 48 +++++- tests/core/test_mpi_execution.py | 19 +++ 10 files changed, 473 insertions(+), 82 deletions(-) diff --git a/docs/components/core.md b/docs/components/core.md index 5172573..980e5e2 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/benchmark.py b/src/cuphoton/core/benchmark.py index 7865cde..a335736 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 e433324..4ac59ab 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 index 92110ea..fc69269 100644 --- a/src/cuphoton/core/cli/executor.py +++ b/src/cuphoton/core/cli/executor.py @@ -29,7 +29,9 @@ class ExecutorOptions: class ExecutorArg(SetInvariant): _arg = "--executor" - _help = "Execution runtime. [default: %default]" + _help = ( + "Execution runtime: local, dragon, or mpi. [default: %default]" + ) _set = {"local", "dragon", "mpi"} _default = "local" @@ -58,12 +60,18 @@ class RankSetupTimeoutSecArg(FloatInvariant): class WarmupRoundsArg(NonNegativeIntegerInvariant): _arg = "--warmup-rounds" - _help = "Persistent-worker warmup passes, with outputs retained." + _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 = "Measured passes in persistent workers. [default: 1]" + _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: diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index 9a884f3..4ac5afd 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -37,6 +37,9 @@ timestamp_utc, validate_identifier, ) +from cuphoton.core.bulk import ( + hostnames_match as _hostnames_match, +) from cuphoton.core.execution import ( ExecutionResult, WorkloadSpec, @@ -156,33 +159,6 @@ def _singleton_cuda_visibility(expected_gpu_id: int | None = None) -> str: 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_dragon_api() -> _DragonAPI: try: from dragon.infrastructure.policy import Policy @@ -365,7 +341,7 @@ def run_dragon_work_items( phase = "process_setup" phase_start = time.perf_counter() try: - results_queue = api.Queue(maxsize=worker_count) + results_queue = api.Queue(maxsize=2 * worker_count) group = api.ProcessGroup( restart=False, ignore_error_on_exit=False, @@ -449,6 +425,7 @@ def run_dragon_work_items( round_id=None, deadline=worker_deadline, group=group, + closed_messages=closed, ) provenances = [message["provenance"] for message in ready] ready_errors = audit_worker_provenance( @@ -523,6 +500,7 @@ def run_dragon_work_items( round_id=round_spec.round_id, deadline=worker_deadline, group=group, + closed_messages=closed, ) finally: round_timings["collection_sec"] = ( @@ -561,10 +539,15 @@ def run_dragon_work_items( spec, shards, worker_results, - artifact_timeout_sec=result_timeout_sec, + artifact_timeout_sec=min(result_timeout_sec, 0.01) + if round_errors + else result_timeout_sec, ) - round_timings["artifact_audit_sec"] = ( - time.perf_counter() - audit_start + round_timings["artifact_audit_sec"] = max( + 0.0, + time.perf_counter() + - audit_start + - terminal_summary["finalization_sec"], ) receipt = { **round_spec.to_payload(), @@ -573,6 +556,7 @@ def run_dragon_work_items( 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")) @@ -591,7 +575,10 @@ def run_dragon_work_items( messages=messages, round_errors=round_errors, ) - atomic_write_json(round_dir / "summary.json", terminal_summary) + if benchmark: + atomic_write_json( + round_dir / "summary.json", terminal_summary + ) rounds.append(receipt) if receipt["status"] != "success": raise RuntimeError( @@ -767,31 +754,66 @@ def _collect_messages( round_id: str | None, deadline: float, group: Any, + closed_messages: list[dict[str, Any]] | None = None, ) -> None: - seen: set[int] = set() + 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=min(remaining, _WORKER_POLL_SEC)), + 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: - failed_exits = [ + 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 group.inactive_puids - if code != 0 + for puid, code in observed_exits + if puid not in seen.values() ] - if failed_exits: + if missing_exits: raise RuntimeError( f"Dragon worker exited before {kind} completion: " - f"{failed_exits}" + f"{missing_exits}" ) + observed_exits = None continue - messages.append(message) 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 @@ -799,12 +821,14 @@ def _collect_messages( 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.add(worker_id) + seen[worker_id] = puid if message.get("status") != "success": - raise RuntimeError(f"Dragon worker {worker_id} {kind} failed") - if kind == "round": + failed_workers.add(worker_id) + if kind == "round" and message.get("status") == "success": result = message.get("result") duration = message.get("worker_wall_sec") if ( @@ -817,6 +841,10 @@ def _collect_messages( 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( @@ -887,6 +915,12 @@ def _validate_binding( 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, @@ -897,16 +931,19 @@ def _workload_worker( ) -> 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( @@ -953,6 +990,7 @@ def _workload_worker( 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, @@ -1012,17 +1050,21 @@ def _workload_worker( command = commands.get(timeout=descriptor["worker_timeout_sec"]) if command != {"run_id": run_id, "kind": "close"}: raise ValueError("unexpected Dragon worker close command") - except Exception: + 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() @@ -1035,6 +1077,10 @@ def _workload_worker( ) except Exception as exc: closed.update(status="failed", artifact_error=error_payload(exc)) - results_queue.put(closed, timeout=timeout) + 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 index a85ae14..8c1e159 100644 --- a/src/cuphoton/core/execution.py +++ b/src/cuphoton/core/execution.py @@ -25,6 +25,7 @@ audit_terminal_records, classify_physical_gpu_pair, error_payload, + hostnames_match, item_ids_sha256, json_mapping, partition_byte_balanced, @@ -40,12 +41,24 @@ class Worker(Protocol): - """Worker-local state constructed after placement and CUDA binding.""" + """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 one item and durably publish its ordinary output files.""" + """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.""" @@ -397,10 +410,7 @@ def audit_worker_provenance( 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 = ( - host.rstrip(".").casefold().split(".")[0] - == other_host.rstrip(".").casefold().split(".")[0] - ) + same_host = hostnames_match(host, other_host) relation = classify_physical_gpu_pair( dict(ids), dict(other_ids), same_host=same_host ) @@ -424,7 +434,11 @@ def finalize_round( *, artifact_timeout_sec: float, ) -> dict[str, Any]: - """Audit durable execution evidence, then run a component's merge step.""" + """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) @@ -504,7 +518,9 @@ def finalize_round( ) ) 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( @@ -515,6 +531,8 @@ def finalize_round( ) 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, @@ -527,8 +545,8 @@ def finalize_round( "worker_results": [dict(result) for result in worker_results], "errors": errors, "result": scientific_result, + "finalization_sec": finalization_sec, } - atomic_write_json(run_dir / "summary.json", summary) return summary diff --git a/src/cuphoton/core/mpi.py b/src/cuphoton/core/mpi.py index 35bceec..a9a76dc 100644 --- a/src/cuphoton/core/mpi.py +++ b/src/cuphoton/core/mpi.py @@ -301,6 +301,7 @@ def run_mpi_work_items( 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 ), @@ -310,9 +311,10 @@ def run_mpi_work_items( }, } last_summary.update(timing) - atomic_write_json( - round_dir / "summary.json", last_summary - ) + if benchmark: + atomic_write_json( + round_dir / "summary.json", last_summary + ) rounds.append( { **planned.to_payload(), @@ -365,7 +367,13 @@ def run_mpi_work_items( if context.rank == 0: try: if not ownership["claimed"]: - raise RuntimeError("MPI run directory was not 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 diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 887db2c..4687c6f 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -176,6 +176,7 @@ def init(self): 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)) @@ -217,6 +218,7 @@ def binding(placement, *, allow_loopback_alias): 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", @@ -240,6 +242,7 @@ def _run( benchmark=None, worker_count=2, worker_timeout=10, + result_timeout=0.2, large=False, validator=None, finalizer=None, @@ -263,7 +266,7 @@ def _run( output_root=tmp_path, run_id="run", max_workers=worker_count, - result_timeout_sec=0.2, + result_timeout_sec=result_timeout, worker_timeout_sec=worker_timeout, benchmark=benchmark, ) @@ -326,6 +329,8 @@ def finalizer(round_dir, records): (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 @@ -475,8 +480,18 @@ def test_fast_worker_waits_for_close_until_all_round_results_arrive( assert state.groups[0].closed and not state.groups[0].stopped -@pytest.mark.parametrize("kind", ["ready", "round", "closed"]) -def test_collector_detects_native_exit_on_first_bounded_poll(kind): +@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: @@ -492,11 +507,12 @@ def get(self, *, timeout): run_id="run", kind=kind, round_id=None, - deadline=time.monotonic() + 3600, - group=SimpleNamespace(inactive_puids=[(123, -9)]), + deadline=time.monotonic() + 2, + group=SimpleNamespace(inactive_puids=[(123, exit_code)]), ) - assert len(timeouts) == 1 + assert len(timeouts) == 2 assert 0 < timeouts[0] <= 1 + assert timeouts[1] == 0 def test_collector_keeps_queued_results_before_reporting_peer_exit(): @@ -505,6 +521,7 @@ def test_collector_keeps_queued_results_before_reporting_peer_exit(): "run_id": "run", "round_id": None, "worker_id": 0, + "puid": 1000, "status": "success", } @@ -526,7 +543,7 @@ def get(self, *, timeout): run_id="run", kind="ready", round_id=None, - deadline=time.monotonic() + 3600, + deadline=time.monotonic() + 2, group=SimpleNamespace(inactive_puids=[(123, -9)]), ) assert messages == [receipt] @@ -538,6 +555,7 @@ def test_closed_collection_allows_normal_exit_before_delayed_receipt(): "run_id": "run", "round_id": None, "worker_id": 0, + "puid": 1000, "status": "success", } @@ -559,7 +577,7 @@ def get(self, *, timeout): run_id="run", kind="closed", round_id=None, - deadline=time.monotonic() + 3600, + deadline=time.monotonic() + 2, group=SimpleNamespace(inactive_puids=[(123, 0)]), ) assert channel.calls == 2 @@ -583,3 +601,207 @@ def crash(*args): 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: ( + [] 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, worker_timeout=0.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" diff --git a/tests/core/test_execution.py b/tests/core/test_execution.py index 454ef5d..1ef6218 100644 --- a/tests/core/test_execution.py +++ b/tests/core/test_execution.py @@ -99,16 +99,26 @@ def stage(tmp_path, spec=None): return spec, shards, results -def test_round_audits_then_finalizes_in_manifest_order(tmp_path): +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(finalize_round=finalize) + 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", @@ -119,9 +129,10 @@ def finalize(run_dir, records): ) assert report["status"] == "success" assert report["terminal_record_audit"]["ok"] - assert calls == [["one", "two"]] - assert report["result"] == {"ordered": ["one", "two"]} - assert read_json_mapping(tmp_path / "run" / "summary.json") == report + 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( @@ -249,3 +260,30 @@ def test_run_directory_is_never_reused(tmp_path): 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_mpi_execution.py b/tests/core/test_mpi_execution.py index 2bd4c3d..2d62846 100644 --- a/tests/core/test_mpi_execution.py +++ b/tests/core/test_mpi_execution.py @@ -179,6 +179,9 @@ def run(rank): 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: @@ -243,6 +246,14 @@ def bcast(self, value, root): ) _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} @@ -251,12 +262,20 @@ def bcast(self, value, root): 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( From defb890586611187ac41b88b3e92d5481da6bd23 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 14:22:07 -0700 Subject: [PATCH 09/10] Make the failed-round artifact timeout test deterministic Signed-off-by: Trent Nelson --- tests/core/test_dragon.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 4687c6f..1d4977b 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -761,7 +761,9 @@ def test_known_round_failure_uses_short_artifact_timeout( _install_runtime( monkeypatch, mutate=lambda message: ( - [] if message["kind"] == "round" else [message] + [{**message, "status": "failed"}] + if message["kind"] == "round" + else [message] ), ) finalize = dragon.finalize_round @@ -772,9 +774,7 @@ def inspect_finalize(*args, **kwargs): return finalize(*args, **kwargs) monkeypatch.setattr(dragon, "finalize_round", inspect_finalize) - result = _run( - tmp_path, worker_count=1, worker_timeout=0.1, result_timeout=30 - ) + result = _run(tmp_path, worker_count=1, result_timeout=30) assert result.status == "failed" assert observed == [0.01] From 7cbb908b1544eb19bc5d624a47ff4276745dd1a1 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Thu, 24 Sep 2026 17:42:53 -0700 Subject: [PATCH 10/10] Align worker startup and allow MPI coordinator-only planning Signed-off-by: Trent Nelson --- src/cuphoton/core/dragon.py | 12 ++++++++- src/cuphoton/core/mpi.py | 46 +++++++++++++++++++++++++++----- tests/core/test_dragon.py | 26 ++++++++++++++++++ tests/core/test_mpi_execution.py | 9 +++++-- 4 files changed, 84 insertions(+), 9 deletions(-) diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index 4ac5afd..39ed70a 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -276,6 +276,13 @@ def run_dragon_work_items( 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() @@ -904,9 +911,12 @@ def _validate_binding( visibility = _singleton_cuda_visibility(placement.gpu_id) premature = [ name - for name in ("cupy", "numba.cuda", "cuda.tile", "torch") + 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: " diff --git a/src/cuphoton/core/mpi.py b/src/cuphoton/core/mpi.py index a9a76dc..a25d409 100644 --- a/src/cuphoton/core/mpi.py +++ b/src/cuphoton/core/mpi.py @@ -29,6 +29,7 @@ build_benchmark_report, ) from .bulk import ( + WorkItem, atomic_write_json, error_payload, json_mapping, @@ -46,6 +47,7 @@ execute_worker_round, finalize_round, prepare_run, + resolve_worker_factory, ) @@ -56,13 +58,16 @@ def run_mpi_work_items( 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. + 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() @@ -111,14 +116,43 @@ def run_mpi_work_items( spec = None preflight_error = None try: - 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") + 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, diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 1d4977b..d57239a 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -8,6 +8,7 @@ import json import queue +import sys import threading import time from types import SimpleNamespace @@ -805,3 +806,28 @@ def put(self, message, **kwargs): ) 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_mpi_execution.py b/tests/core/test_mpi_execution.py index 2d62846..f22a6af 100644 --- a/tests/core/test_mpi_execution.py +++ b/tests/core/test_mpi_execution.py @@ -54,8 +54,9 @@ def tracking_factory(options): "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 + monkeypatch, tmp_path, failure, prepare_on_root ): for name in tuple(os.environ): if name.startswith( @@ -128,7 +129,7 @@ def fail_outside_item(*args, **kwargs): def prepare(rank): trace.append(("preflight", rank)) - if failure == "preflight" and rank == 1: + if failure == "preflight" and rank == (0 if prepare_on_root else 1): raise ValueError("bad rank input") return workload( worker_factory=tracking_factory, @@ -141,6 +142,7 @@ def run(rank): 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, @@ -153,6 +155,9 @@ def run(rank): 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)