From 4a25f7f0e46c953c2fddb602f97369b8f6704806 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 08:48:11 -0700 Subject: [PATCH 1/3] Add configurable GPU sharing for pipeline workers Allow independent image pairs to share a GPU through Dragon processes or local process and thread workers. Expose worker counts, MPS service selection and free-threading checks through the CLI and Python API. Preserve explicit placement, retained outputs and worker cleanup across repeated rounds. Signed-off-by: Trent Nelson --- docs/components/xscan.md | 109 +++- docs/dragon.md | 7 + src/cuphoton/core/cli/executor.py | 50 +- src/cuphoton/core/dragon.py | 148 +++++- src/cuphoton/core/execution.py | 43 +- src/cuphoton/core/mps.py | 70 +++ src/cuphoton/xscan/commands.py | 20 +- src/cuphoton/xscan/pipeline_executor.py | 32 +- src/cuphoton/xscan/sharing.py | 638 ++++++++++++++++++++++++ tests/core/test_cli_contract.py | 6 +- tests/core/test_dragon.py | 220 +++++++- tests/core/test_execution.py | 57 +++ tests/core/test_executor_options.py | 11 +- tests/core/test_mps.py | 61 +++ tests/xscan/_sharing_support.py | 108 ++++ tests/xscan/test_pipeline_executor.py | 139 +++++- tests/xscan/test_sharing.py | 216 ++++++++ 17 files changed, 1904 insertions(+), 31 deletions(-) create mode 100644 src/cuphoton/core/mps.py create mode 100644 src/cuphoton/xscan/sharing.py create mode 100644 tests/core/test_mps.py create mode 100644 tests/xscan/_sharing_support.py create mode 100644 tests/xscan/test_sharing.py diff --git a/docs/components/xscan.md b/docs/components/xscan.md index 3ebbc065..3d93053e 100644 --- a/docs/components/xscan.md +++ b/docs/components/xscan.md @@ -561,8 +561,8 @@ scientific evidence; Dragon saves each result in `items//summary.json` and checks it again in the coordinator. The evidence decoder returns the 22 named arrays for comparison. -The `run-pipeline` command accepts the same descriptors with either Dragon or -MPI. Save the configuration and items above as a manifest: +The `run-pipeline` command accepts the same descriptors with Dragon, MPI, +or local GPU workers. Save the configuration and items above as a manifest: ```python import json @@ -616,6 +616,111 @@ decoding and reader stream completion; `input_h2d` measures NPY upload calls. Timings are host elapsed times, so distinguish model initialization, first-item compilation and warmed context reuse when comparing runs. +### Share a GPU between image pairs + +`--workers-per-gpu N` allows independent image pairs to use the same GPU. +The default is one worker per GPU. Each worker retains its own model and +pipeline context, so additional workers consume additional device and host +memory. Start with a small count and measure throughput and memory on the +intended images; increasing the count can also slow a saturated GPU. + +For a local run, expose exactly one GPU before starting Python and select +`--executor processes` or `--executor threads`: + +```bash +CUDA_VISIBLE_DEVICES=0 cuphoton xscan run-pipeline \ + --executor processes --workers-per-gpu 4 \ + --manifest pipeline.json --output-dir runs --name processes-4 \ + --worker-timeout-sec 600 --warmup-rounds 1 --measure-rounds 3 + +CUDA_VISIBLE_DEVICES=0 cuphoton xscan run-pipeline \ + --executor threads --workers-per-gpu 4 \ + --manifest pipeline.json --output-dir runs --name threads-4 \ + --worker-timeout-sec 600 --warmup-rounds 1 --measure-rounds 3 +``` + +These options call the same Python entry point: + +```python +from pathlib import Path +from cuphoton.xscan.pipeline_executor import run_pipeline_manifest + +result = run_pipeline_manifest( + executor="processes", + manifest_path=Path("pipeline.json"), + output_root=Path("runs"), + workers_per_gpu=4, + worker_timeout_sec=600, +) +``` + +Process mode creates persistent spawned processes. Thread mode keeps workers +in one supervised process, with a separate model, CuPy memory pool and pair +of CUDA streams per thread. Initialization and first use are serialized; +use a warmup round to exclude that cost from measured rounds. Outputs from +every planned round are retained. The worker count is capped by the number +of items. A timeout or worker failure fails the run and stops its owned +processes; consult the terminal `summary.json` before using timing results. + +Thread mode accepts prepared NPY inputs. It works with ordinary Python and +can also use a compatible free-threaded interpreter. To require the GIL to +remain disabled, add `--require-free-threaded` (Python: +`require_free_threaded=True`) to the threads command. This checks the running +interpreter; it does not disable the GIL. Extensions may re-enable it on +import, and forcing it off does not establish extension thread safety. See +[Python's free-threading guide](https://docs.python.org/3/howto/free-threading-python.html). +Use process mode for FITS inputs. Free-threaded execution requires validation +of the installed scientific libraries; it is not a general compatibility +claim for native extensions. + +With the shared Dragon executor, the same count applies to every discovered +GPU. `--max-workers` caps the total number of processes. Dragon uses distinct +GPUs before assigning additional workers to each one: + +```bash +dragon cuphoton xscan run-pipeline --executor dragon \ + --workers-per-gpu 4 --max-workers 16 \ + --manifest pipeline.json --output-dir runs --name dragon-shared +``` + +The shared Dragon Python API accepts `workers_per_gpu` as well. Physical GPU +checks require workers assigned to the same GPU to agree on its identity, +and still reject collisions between different requested GPUs. MPI retains +one rank per physical GPU and rejects the sharing options. + +### Select an existing MPS service + +Process workers can run under CUDA's ordinary process scheduling or an +externally managed Multi-Process Service (MPS). MPS allows CUDA work from +different processes to overlap on the GPU. To require a connection to an +existing MPS v2 service, add an absolute pipe directory: + +```bash +CUDA_VISIBLE_DEVICES=GPU-your-allocated-device-uuid \ + cuphoton xscan run-pipeline --executor processes --workers-per-gpu 4 \ + --mps-pipe-directory /tmp/my-job-mps \ + --manifest pipeline.json --output-dir runs --name mps-4 +``` + +The Python keyword is `mps_pipe_directory`. It is also accepted by the shared +Dragon executor. Every worker must appear in that service's client list +after CUDA initialization, or startup fails. The service and clients must +use a common PID namespace for this check. On multiple hosts use the same +absolute node-local directory on each host, such as a job-specific `/tmp` +path. Dragon resolves its assigned physical device to a UUID before CUDA +initialization and checks the reported UUID, avoiding MPS ordinal remapping. +This requires `nvidia-smi` on each worker host. Record the runtime's +configuration when comparing runs. Without this option, process workers +inherit the environment's CUDA +and MPS settings without asserting an MPS connection. + +cuPhoton leaves service lifetime, device visibility and resource limits to +the launcher or administrator. It does not start or stop MPS or change GPU +compute mode. The explicit connection check uses the +[MPS v2 control interface](https://docs.nvidia.com/deploy/mps/latest/mpsv2-interface.html). +Thread mode rejects `--mps-pipe-directory`; its threads share a process and +CUDA context without needing an interprocess service. + ## Review workflow ```bash diff --git a/docs/dragon.md b/docs/dragon.md index 2fe64a3d..fbf95688 100644 --- a/docs/dragon.md +++ b/docs/dragon.md @@ -31,6 +31,13 @@ pipeline keeps intermediate arrays inside each GPU worker while it processes an image pair. Each GPU handles complete items; these paths do not split one image across GPUs. +The shared executor accepts `--workers-per-gpu` to place several independent +workers on each GPU; its default is one. `--max-workers` remains a total +process cap. See [GPU sharing](components/xscan.md#share-a-gpu-between-image-pairs) +for CLI and Python examples, local process/thread modes, and external MPS +connection checks. This option applies to the shared component executor; +the standalone XPOIS batch API retains its existing placement behavior. + ### Launch descriptors Dragon workers load immutable, hash-checked launch descriptors from the shared diff --git a/src/cuphoton/core/cli/executor.py b/src/cuphoton/core/cli/executor.py index 38cdb04e..647d2dda 100644 --- a/src/cuphoton/core/cli/executor.py +++ b/src/cuphoton/core/cli/executor.py @@ -6,6 +6,8 @@ from __future__ import annotations +from collections.abc import Mapping + from ..benchmark import BenchmarkOptions from .command import CommandError from .invariants import ( @@ -13,6 +15,7 @@ NonNegativeIntegerInvariant, PositiveIntegerInvariant, SetInvariant, + StringInvariant, ) @@ -26,6 +29,8 @@ class ExecutorOptions: rank_setup_timeout_sec: float | None = None warmup_rounds: int | None = None measure_rounds: int | None = None + workers_per_gpu: int | None = None + mps_pipe_directory: str | None = None class ExecutorArg(SetInvariant): _arg = "--executor" @@ -42,10 +47,26 @@ class MaxWorkersArg(PositiveIntegerInvariant): class WorkerTimeoutSecArg(FloatInvariant): _arg = "--worker-timeout-sec" - _help = "Dragon worker lifetime across all rounds. [default: 3600]" + _help = "Dragon or local GPU worker lifetime. [default: 3600]" _default = None _min = 0.001 + class WorkersPerGpuArg(PositiveIntegerInvariant): + _arg = "--workers-per-gpu" + _help = ( + "Maximum workers sharing each GPU with Dragon or a local " + "pipeline executor. [default: 1]" + ) + _default = None + + class MpsPipeDirectoryArg(StringInvariant): + _arg = "--mps-pipe-directory" + _help = ( + "Existing MPS v2 pipe directory for Dragon or process workers; " + "require every worker to connect to that service." + ) + _default = None + class ResultTimeoutSecArg(FloatInvariant): _arg = "--result-timeout-sec" _help = "Dragon result and artifact grace period. [default: 60]" @@ -77,22 +98,32 @@ class MeasureRoundsArg(PositiveIntegerInvariant): def executor_options(self) -> dict: """Reject ignored flags before loading a numerical runtime.""" - dragon = { + dragon_only = { "max_workers": self.max_workers, - "worker_timeout_sec": self.worker_timeout_sec, "result_timeout_sec": self.result_timeout_sec, } + sharing = { + "workers_per_gpu": self.workers_per_gpu, + "worker_timeout_sec": self.worker_timeout_sec, + } + mps = {"mps_pipe_directory": self.mps_pipe_directory} + dragon = {**dragon_only, **sharing, **mps} mpi = {"rank_setup_timeout_sec": self.rank_setup_timeout_sec} rounds = { "warmup_rounds": self.warmup_rounds, "measure_rounds": self.measure_rounds, } + invalid: Mapping[str, str | int | float | None] if self.executor == "local": invalid = {**dragon, **mpi, **rounds} elif self.executor == "dragon": invalid = mpi elif self.executor == "mpi": invalid = dragon + elif self.executor == "processes": + invalid = {**dragon_only, **mpi} + elif self.executor == "threads": + invalid = {**dragon_only, **mpi, **mps} else: raise CommandError("executor must be local, dragon, or mpi") supplied = [ @@ -107,11 +138,16 @@ def executor_options(self) -> dict: ) if self.executor == "local": return {} - options: dict[str, int | float | BenchmarkOptions | None] = { + by_executor: dict[str, Mapping[str, str | int | float | None]] = { + "dragon": dragon, + "mpi": mpi, + "processes": {**sharing, **mps}, + "threads": sharing, + } + runtime_options = by_executor[self.executor] + options: dict[str, str | int | float | BenchmarkOptions | None] = { name: value - for name, value in ( - dragon if self.executor == "dragon" else mpi - ).items() + for name, value in runtime_options.items() if value is not None } options["benchmark"] = ( diff --git a/src/cuphoton/core/dragon.py b/src/cuphoton/core/dragon.py index 11b8fbdb..c825e92a 100644 --- a/src/cuphoton/core/dragon.py +++ b/src/cuphoton/core/dragon.py @@ -12,6 +12,7 @@ import os import queue import socket +import subprocess import sys import time from collections.abc import Callable, Mapping, Sequence @@ -19,6 +20,7 @@ from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any +from uuid import UUID from cuphoton.core.benchmark import ( BenchmarkOptions, @@ -103,31 +105,40 @@ def discover_gpu_placements( def _select_gpu_placements( - placements: Sequence[Placement], worker_count: int + placements: Sequence[Placement], + worker_count: int, + *, + workers_per_gpu: int = 1, ) -> tuple[Placement, ...]: - """Select GPUs round-robin across hosts, then assign dense worker IDs.""" + """Spread workers across hosts and GPUs before reusing GPU slots.""" 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): + if type(workers_per_gpu) is not int or workers_per_gpu <= 0: + raise ValueError("workers_per_gpu must be a positive integer") + if worker_count <= 0 or worker_count > len(placements) * workers_per_gpu: 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: + unique_count = min(worker_count, len(placements)) + while len(selected) < unique_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: + if len(selected) == unique_count: break if not progressed: raise RuntimeError("could not select requested Dragon placements") offset += 1 + selected = [ + selected[index % unique_count] for index in range(worker_count) + ] return tuple( Placement( worker_id=worker_id, @@ -208,12 +219,58 @@ def _stable_gpu_physical_ids( return identities or None +def _normalized_gpu_uuid(value: Any) -> str | None: + if not isinstance(value, str): + return None + try: + return UUID(value.removeprefix("GPU-")).hex + except ValueError: + return None + + +def _mps_gpu_uuid(gpu_id: int) -> str: + """Resolve one Dragon/NVML ordinal without initializing CUDA. + + Dragon's NVIDIA discovery reports nvidia-smi device indices. MPS may + remap these ordinals, so bind the corresponding UUID in the client. + """ + + try: + result = subprocess.run( + [ + "nvidia-smi", + f"--id={gpu_id}", + "--query-gpu=index,uuid", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=10, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError("cannot resolve Dragon GPU UUID for MPS") from exc + fields = [field.strip() for field in result.stdout.strip().split(",")] + if ( + result.returncode != 0 + or len(fields) != 2 + or fields[0] != str(gpu_id) + or not fields[1].startswith("GPU-") + or _normalized_gpu_uuid(fields[1]) is None + ): + raise RuntimeError( + "nvidia-smi did not identify the requested Dragon GPU for MPS" + ) + return fields[1] + + def _valid_shard_provenance( value: Any, *, placement: Placement, allow_loopback_alias: bool, backend: str, + mps_pipe_directory: str | None = None, ) -> bool: if not isinstance(value, Mapping): return False @@ -223,6 +280,25 @@ def _valid_shard_provenance( requested_host = value.get("requested_host") hostname = value.get("hostname") visibility = value.get("cuda_visible_devices") + expected_visibility = str(placement.gpu_id) + requested_uuid = None + if mps_pipe_directory is not None: + raw_uuid = value.get("requested_gpu_uuid") + if not isinstance(raw_uuid, str): + return False + expected_visibility = raw_uuid + requested_uuid = _normalized_gpu_uuid(expected_visibility) + if requested_uuid is None: + return False + mps = value.get("mps") + if ( + not isinstance(mps, Mapping) + or mps.get("pipe_directory") != mps_pipe_directory + or _strict_integer(mps.get("client_pid")) != pid + or (server_pid := _strict_integer(mps.get("server_pid"))) is None + or server_pid <= 0 + ): + return False if ( worker_id != placement.worker_id or requested_host != placement.host @@ -238,12 +314,20 @@ def _valid_shard_provenance( ) or not isinstance(visibility, str) or [token.strip() for token in visibility.split(",") if token.strip()] - != [str(placement.gpu_id)] + != [expected_visibility] + or ( + mps_pipe_directory is not None + and value.get("mps_pipe_directory") != mps_pipe_directory + ) ): return False gpu = value.get("gpu") if not isinstance(gpu, Mapping) or not gpu: return False + if requested_uuid is not None and ( + _normalized_gpu_uuid(gpu.get("uuid")) != requested_uuid + ): + return False gpu_backend = gpu.get("backend") expected_identity_backend = ( "cupy" if backend in {"cupy", "cutile", "native"} else backend @@ -268,6 +352,8 @@ def run_dragon_work_items( output_root: Path, run_id: str | None = None, max_workers: int | None = None, + workers_per_gpu: int = 1, + mps_pipe_directory: str | None = None, result_timeout_sec: float = 60.0, worker_timeout_sec: float = 3600.0, benchmark: BenchmarkOptions | None = None, @@ -278,6 +364,11 @@ def run_dragon_work_items( Complete item descriptors live in hashed shared-filesystem launch files; control queues carry bounded round commands and completion receipts. + ``workers_per_gpu`` permits that many independent worker processes on + each allocated GPU; ``max_workers`` caps the total process count. An + explicit MPS pipe requires every worker to connect to an externally + managed MPS server. This executor never starts or stops that server. + 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. @@ -287,6 +378,13 @@ def run_dragon_work_items( """ invocation_start = time.perf_counter() + if type(workers_per_gpu) is not int or workers_per_gpu <= 0: + raise ValueError("workers_per_gpu must be a positive integer") + if mps_pipe_directory is not None and ( + not isinstance(mps_pipe_directory, str) + or not Path(mps_pipe_directory).is_absolute() + ): + raise ValueError("mps_pipe_directory must be an absolute path") for name, value in ( ("result_timeout_sec", result_timeout_sec), ("worker_timeout_sec", worker_timeout_sec), @@ -319,11 +417,22 @@ def run_dragon_work_items( api.System, api.Node, node_ids=node_ids ) worker_count = min( - max_workers or len(available), len(available), len(spec.items) + max_workers or len(available) * workers_per_gpu, + len(available) * workers_per_gpu, + len(spec.items), ) if not worker_count: raise RuntimeError("Dragon allocation exposes no usable GPUs") - placements = _select_gpu_placements(available, worker_count) + placements = _select_gpu_placements( + available, worker_count, workers_per_gpu=workers_per_gpu + ) + group_ids: dict[tuple[str, int], int] = {} + gpu_groups = tuple( + group_ids.setdefault( + (placement.host, placement.gpu_id), len(group_ids) + ) + for placement in placements + ) 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") @@ -377,6 +486,7 @@ def run_dragon_work_items( "allow_loopback_alias": len(node_ids) == 1, "worker_timeout_sec": worker_timeout_sec, "result_timeout_sec": result_timeout_sec, + "mps_pipe_directory": mps_pipe_directory, } descriptor_path = ( run_dir / "launch" / f"worker-{placement.worker_id:04d}.json" @@ -444,6 +554,7 @@ def run_dragon_work_items( provenances, backend=spec.backend, expected_worker_count=worker_count, + gpu_groups=gpu_groups, ) if ready_errors: raise ValueError( @@ -455,6 +566,7 @@ def run_dragon_work_items( placement=placements[message["worker_id"]], allow_loopback_alias=len(node_ids) == 1, backend=spec.backend, + mps_pipe_directory=mps_pipe_directory, ): raise ValueError( "Dragon READY differs from requested placement" @@ -554,6 +666,7 @@ def run_dragon_work_items( artifact_timeout_sec=min(result_timeout_sec, 0.01) if round_errors else result_timeout_sec, + gpu_groups=gpu_groups, ) round_timings["artifact_audit_sec"] = max( 0.0, @@ -712,6 +825,9 @@ def run_dragon_work_items( "worker_timeout_sec": worker_timeout_sec, "result_timeout_sec": result_timeout_sec, "worker_count": worker_count, + "workers_per_gpu": workers_per_gpu, + "gpu_count": len(group_ids), + "mps_pipe_directory": mps_pipe_directory, "allocation_node_count": len(node_ids), "distinct_host_count": len( {placement.host for placement in placements} @@ -969,8 +1085,20 @@ def _workload_worker( placement, allow_loopback_alias=descriptor["allow_loopback_alias"], ) + mps_pipe_directory = descriptor.get("mps_pipe_directory") + requested_gpu_uuid = None + if mps_pipe_directory is not None: + os.environ["CUDA_MPS_PIPE_DIRECTORY"] = mps_pipe_directory + requested_gpu_uuid = _mps_gpu_uuid(placement.gpu_id) + os.environ["CUDA_VISIBLE_DEVICES"] = requested_gpu_uuid + visibility = requested_gpu_uuid factory = resolve_worker_factory(descriptor["worker_factory"]) worker = factory(descriptor["options"]) + mps_receipt = None + if mps_pipe_directory is not None: + from .mps import require_mps_client + + mps_receipt = require_mps_client(os.getpid()) provenance = { "worker_id": worker_id, "requested_host": placement.host, @@ -982,6 +1110,10 @@ def _workload_worker( worker.gpu_identity, field="worker GPU identity" ), } + if mps_pipe_directory is not None: + provenance["mps_pipe_directory"] = mps_pipe_directory + provenance["requested_gpu_uuid"] = requested_gpu_uuid + provenance["mps"] = mps_receipt ready["provenance"] = provenance except Exception as exc: ready.update(status="failed", error=error_payload(exc)) diff --git a/src/cuphoton/core/execution.py b/src/cuphoton/core/execution.py index d8a001bc..1d75c911 100644 --- a/src/cuphoton/core/execution.py +++ b/src/cuphoton/core/execution.py @@ -346,8 +346,30 @@ def audit_worker_provenance( *, backend: str, expected_worker_count: int, + gpu_groups: Sequence[int] | None = None, ) -> list[dict[str, Any]]: - """Require one usable, non-overlapping physical GPU per placed worker.""" + """Check physical devices against the coordinator's worker placement. + + Workers in the same requested group must share a physical GPU; different + groups must have distinct devices. Omitted groups require one GPU per + worker, preserving the default distributed-execution contract. + """ + + if gpu_groups is None: + groups = tuple(range(expected_worker_count)) + else: + if ( + not isinstance(gpu_groups, Sequence) + or isinstance(gpu_groups, (str, bytes)) + or len(gpu_groups) != expected_worker_count + or any( + type(group) is not int or group < 0 for group in gpu_groups + ) + ): + raise ValueError( + "gpu_groups must contain one non-negative integer per worker" + ) + groups = tuple(gpu_groups) errors: list[dict[str, Any]] = [] identities: list[tuple[int, str, frozenset[tuple[str, str]]]] = [] @@ -377,6 +399,11 @@ def audit_worker_provenance( gpu = value.get("gpu") host = value.get("hostname") visibility = value.get("cuda_visible_devices") + if ( + type(value.get("worker_id")) is not int + or not 0 <= value["worker_id"] < expected_worker_count + ): + invalid.append("worker_id") if type(value.get("pid")) is not int or value["pid"] <= 0: invalid.append("pid") if not isinstance(host, str) or not host: @@ -428,12 +455,20 @@ def audit_worker_provenance( relation = classify_physical_gpu_pair( dict(ids), dict(other_ids), same_host=same_host ) - if relation != "distinct": + expected = ( + "duplicate" + if groups[worker_id] == groups[other_id] + else "distinct" + ) + if relation != expected: errors.append( { "phase": "provenance", "worker_ids": [other_id, worker_id], - "message": f"{relation} physical GPU identities", + "message": ( + f"{relation} physical GPU identities; " + f"expected {expected} for requested GPU groups" + ), } ) return errors @@ -447,6 +482,7 @@ def finalize_round( worker_results: Sequence[Mapping[str, Any]], *, artifact_timeout_sec: float, + gpu_groups: Sequence[int] | None = None, ) -> dict[str, Any]: """Audit evidence, merge outputs, and return an unpublished summary. @@ -529,6 +565,7 @@ def finalize_round( ], backend=spec.backend, expected_worker_count=len(shards), + gpu_groups=gpu_groups, ) ) scientific_result = None diff --git a/src/cuphoton/core/mps.py b/src/cuphoton/core/mps.py new file mode 100644 index 00000000..2638bbe1 --- /dev/null +++ b/src/cuphoton/core/mps.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Read-only connection checks for an externally managed MPS v2 service.""" + +from __future__ import annotations + +import os +import subprocess +import time +from pathlib import Path + + +def require_mps_client(pid: int) -> dict[str, str | int]: + """Confirm an initialized CUDA worker appears in the selected service. + + The caller sets ``CUDA_MPS_PIPE_DIRECTORY`` before initializing CUDA. + Client and service must report PIDs in the same namespace. This function + never starts a daemon, terminates a client, or changes GPU policy. + """ + + if type(pid) is not int or pid <= 0: + raise ValueError("MPS client PID must be a positive integer") + directory = os.environ.get("CUDA_MPS_PIPE_DIRECTORY") + if not directory or not Path(directory).is_dir(): + raise RuntimeError("selected MPS pipe directory does not exist") + deadline = time.monotonic() + 10.0 + + def query(command: str) -> list[int]: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError("MPS connection verification timed out") + try: + result = subprocess.run( + ["nvidia-cuda-mps-control"], + input=command + "\n", + text=True, + capture_output=True, + timeout=remaining, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + "cannot query selected MPS v2 service" + ) from exc + if result.returncode != 0: + raise RuntimeError( + "cannot query selected MPS v2 service: " + + result.stderr.strip() + ) + tokens = result.stdout.split() + if any( + not token.isascii() or not token.isdecimal() for token in tokens + ): + raise RuntimeError( + "unexpected MPS v2 PID list: " + result.stdout.strip() + ) + return [int(token) for token in tokens] + + for server_pid in query("get_server_list"): + if server_pid > 0 and pid in query(f"get_client_list {server_pid}"): + return { + "pipe_directory": directory, + "server_pid": server_pid, + "client_pid": pid, + } + raise RuntimeError( + f"worker PID {pid} is not connected to the selected MPS v2 service" + ) diff --git a/src/cuphoton/xscan/commands.py b/src/cuphoton/xscan/commands.py index 9d5b5f15..4fb945c7 100644 --- a/src/cuphoton/xscan/commands.py +++ b/src/cuphoton/xscan/commands.py @@ -947,15 +947,25 @@ class RunPipelineCommand(FitsReaderOptions, ExecutorOptions, XScanCommand): manifest: str | None = None output_dir: str | None = None run_name = None + require_free_threaded = None class ExecutorArg(ExecutorOptions.ExecutorArg): - _set = {"dragon", "mpi"} + _set = {"dragon", "mpi", "processes", "threads"} _default = None _mandatory = True _help = ( - "Distributed runtime: dragon or mpi, with its matching launcher." + "GPU runtime: dragon or mpi with its launcher, or local " + "processes or threads on one visible GPU." ) + class RequireFreeThreadedArg(BoolInvariant): + _arg = "--require-free-threaded" + _help = ( + "Require the threads executor to keep the Python GIL disabled. " + "The interpreter and its extensions must support free threading." + ) + _default = False + class ManifestArg(PathSpecInvariant): _arg = "--manifest" _mandatory = True @@ -975,6 +985,12 @@ def run(self) -> None: assert self.manifest is not None assert self.output_dir is not None options = self.executor_options() + if self.require_free_threaded: + if self.executor != "threads": + raise CommandError( + "--require-free-threaded requires --executor threads" + ) + options["require_free_threaded"] = True from .pipeline_executor import run_pipeline_manifest result = self._call( diff --git a/src/cuphoton/xscan/pipeline_executor.py b/src/cuphoton/xscan/pipeline_executor.py index 0a13bd61..bff31e75 100644 --- a/src/cuphoton/xscan/pipeline_executor.py +++ b/src/cuphoton/xscan/pipeline_executor.py @@ -14,7 +14,11 @@ from cuphoton.core.bulk import WorkItem, collect_gpu_identity, json_mapping -from .device_pipeline import DevicePipelineConfig, DevicePipelineItem +from .device_pipeline import ( + DevicePipelineConfig, + DevicePipelineItem, + NpyArrayDescriptor, +) from .dragon_pipeline import ( DRAGON_DEVICE_PIPELINE_OPTIONS_SCHEMA, _failed_record_problems, @@ -223,8 +227,34 @@ def prepare(rank: int): config, items = load_pipeline_manifest( manifest_path, fits_reader=fits_reader, xdr_options=xdr_options ) + if executor == "threads" and any( + descriptor is not None + and not isinstance(descriptor, NpyArrayDescriptor) + for item in items + for descriptor in ( + item.reference, + item.target, + item.variance, + item.fit_mask, + ) + ): + raise ValueError( + "threads requires prepared NPY inputs; use processes, " + "Dragon or MPI for FITS inputs" + ) return prepare_pipeline_workload(items, config, rank=rank) + if executor in {"processes", "threads"}: + from .sharing import run_shared_pipeline + + return run_shared_pipeline( + executor=executor, + prepare_workload=prepare, + output_root=output_root, + run_id=run_id, + **options, + ) + return run_workload( executor=executor, prepare_workload=prepare, diff --git a/src/cuphoton/xscan/sharing.py b/src/cuphoton/xscan/sharing.py new file mode 100644 index 00000000..0726a255 --- /dev/null +++ b/src/cuphoton/xscan/sharing.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Persistent pipeline workers sharing one explicitly bound GPU.""" + +from __future__ import annotations + +import importlib +import math +import multiprocessing +import os +import socket +import sys +import threading +import time +from collections.abc import Callable, Mapping, Sequence +from contextlib import nullcontext +from multiprocessing.connection import Connection, wait +from pathlib import Path +from typing import Any + +from cuphoton.core.benchmark import ( + BenchmarkOptions, + BenchmarkRound, + build_benchmark_report, +) +from cuphoton.core.bulk import ( + WorkItem, + atomic_write_json, + error_payload, + json_mapping, + new_run_id, + partition_byte_balanced, + read_json_mapping, + timestamp_utc, + validate_identifier, +) +from cuphoton.core.execution import ( + ExecutionResult, + Worker, + WorkloadSpec, + audit_worker_provenance, + execute_worker_round, + finalize_round, + prepare_run, + resolve_worker_factory, +) + + +def _gil_enabled() -> bool: + return bool(getattr(sys, "_is_gil_enabled", lambda: True)()) + + +def _check_gil(required: bool) -> bool: + enabled = _gil_enabled() + if required and enabled: + raise RuntimeError( + "free-threaded execution requires the GIL to remain disabled " + "after importing the pipeline dependencies" + ) + return enabled + + +def _load_cuda() -> Any: + # Torch establishes CUDA library loading before CuPy touches the driver. + torch = importlib.import_module("torch") + cp = importlib.import_module("cupy") + if ( + int(cp.cuda.runtime.getDeviceCount()) != 1 + or int(torch.cuda.device_count()) != 1 + ): + raise RuntimeError("GPU sharing requires exactly one visible GPU") + torch.cuda.set_device(0) + cp.cuda.Device(0).use() + return cp + + +class _FirstUseWorker: + """Serialize each thread's first item without creating hidden work.""" + + def __init__(self, worker: Worker, lock: Any, barrier: Any) -> None: + self.worker = worker + self.lock = lock + self.barrier = barrier + self.first = True + + def run_item(self, item: WorkItem, output_dir: Path) -> Mapping[str, Any]: + if self.first: + self.first = False + try: + with self.lock: + result = self.worker.run_item(item, output_dir) + except BaseException: + self.barrier.abort() + raise + self.barrier.wait() + return result + return self.worker.run_item(item, output_dir) + + @property + def gpu_identity(self) -> Mapping[str, Any]: + return self.worker.gpu_identity + + def close(self) -> None: + self.worker.close() + + +def _worker_loop( + connection: Connection, + payload: Mapping[str, Any], + worker_id: int, + lock: Any = None, + barrier: Any = None, +) -> None: + worker: Worker | None = None + pool = None + close_error = None + stage = "ready" + try: + with lock if lock is not None else nullcontext(): + cp = _load_cuda() + pool = cp.cuda.MemoryPool() if lock is not None else None + allocator = ( + cp.cuda.using_allocator(pool.malloc) + if pool is not None + else nullcontext() + ) + allocator.__enter__() + factory = resolve_worker_factory(payload["worker_factory"]) + worker = factory(payload["options"]) + if not callable( + getattr(worker, "run_item", None) + ) or not callable(getattr(worker, "close", None)): + raise TypeError("worker requires run_item and close methods") + gil_enabled = _check_gil(payload["require_free_threaded"]) + mps = None + if payload["mps_pipe_directory"] is not None: + from cuphoton.core.mps import require_mps_client + + mps = require_mps_client(os.getpid()) + provenance = { + "worker_id": worker_id, + "hostname": socket.gethostname(), + "pid": os.getpid(), + "thread_id": threading.get_native_id(), + "cuda_visible_devices": os.environ["CUDA_VISIBLE_DEVICES"], + "gpu": json_mapping( + worker.gpu_identity, field="worker GPU identity" + ), + "gil_enabled": gil_enabled, + "mps_pipe_directory": os.environ.get( + "CUDA_MPS_PIPE_DIRECTORY" + ), + "mps": mps, + } + if lock is not None: + worker = _FirstUseWorker(worker, lock, barrier) + connection.send({"phase": "ready", "provenance": provenance}) + items = tuple( + WorkItem.from_dict(item) for item in payload["shards"][worker_id] + ) + while True: + command = connection.recv() + stage = command["phase"] + if stage == "close": + break + if stage != "round": + raise ValueError("unknown shared-worker command") + _check_gil(payload["require_free_threaded"]) + started = time.perf_counter() + execute_worker_round( + worker, + items=items, + run_id=command["run_id"], + run_dir=Path(command["run_dir"]), + manifest_sha256=payload["manifest_sha256"], + worker_id=worker_id, + backend=payload["backend"], + provenance=provenance, + ) + _check_gil(payload["require_free_threaded"]) + connection.send( + { + "phase": "round", + "worker_wall_sec": time.perf_counter() - started, + } + ) + except BaseException as exc: + if barrier is not None: + barrier.abort() + try: + connection.send({"phase": stage, "error": error_payload(exc)}) + except (BrokenPipeError, EOFError, OSError): + pass + finally: + try: + if worker is not None: + worker.close() + worker = None + if pool is not None: + pool.free_all_blocks() + except BaseException as exc: + close_error = error_payload(exc) + finally: + if "allocator" in locals(): + try: + allocator.__exit__(None, None, None) + except BaseException as exc: + close_error = error_payload(exc) + try: + connection.send({"phase": "closed", "error": close_error}) + except (BrokenPipeError, EOFError, OSError): + pass + connection.close() + + +def _process_entry( + connections: Sequence[Connection], + payload: Mapping[str, Any], + worker_ids: Sequence[int], + threaded: bool, +) -> None: + if payload["mps_pipe_directory"] is not None: + os.environ["CUDA_MPS_PIPE_DIRECTORY"] = payload["mps_pipe_directory"] + if threaded: + lock = threading.Lock() + barrier = threading.Barrier(len(worker_ids)) + threads = [ + threading.Thread( + target=_worker_loop, + args=(connection, payload, worker_id, lock, barrier), + name=f"cuphoton-pipeline-{worker_id}", + ) + for connection, worker_id in zip( + connections, worker_ids, strict=True + ) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + else: + _worker_loop(connections[0], payload, worker_ids[0]) + + +def _collect( + connections: Sequence[Connection], + owners: Sequence[Any], + phase: str, + timeout: float, +) -> list[dict[str, Any]]: + deadline = time.monotonic() + timeout + pending = dict(enumerate(connections)) + received: dict[int, dict[str, Any]] = {} + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"shared workers timed out during {phase}") + ready = wait(list(pending.values()), timeout=min(remaining, 0.1)) + for worker_id, connection in list(pending.items()): + if connection not in ready: + if owners[worker_id].exitcode is not None: + # Drain a last message sent immediately before exit. + if not connection.poll(): + raise RuntimeError( + f"worker {worker_id} exited during {phase}" + ) + else: + continue + message = connection.recv() + if message.get("phase") != phase: + raise RuntimeError( + f"worker {worker_id} sent an unexpected {phase} receipt" + ) + received[worker_id] = message + del pending[worker_id] + return [received[index] for index in range(len(connections))] + + +def _stop_processes( + processes: Sequence[Any], timeout: float +) -> list[dict[str, Any]]: + errors = [] + deadline = time.monotonic() + timeout + for process in processes: + process.join(timeout=max(0.0, deadline - time.monotonic())) + for process in processes: + if process.is_alive(): + errors.append( + { + "phase": "shutdown", + "message": "worker required termination", + } + ) + process.terminate() + for process in processes: + process.join(timeout=2) + if process.is_alive(): + process.kill() + process.join(timeout=2) + if process.exitcode != 0: + errors.append( + { + "phase": "shutdown", + "pid": process.pid, + "exitcode": process.exitcode, + } + ) + return errors + + +def run_shared_pipeline( + *, + prepare_workload: Callable[[int], WorkloadSpec], + output_root: Path, + run_id: str | None = None, + executor: str, + workers_per_gpu: int = 1, + worker_timeout_sec: float = 3600, + benchmark: BenchmarkOptions | None = None, + require_free_threaded: bool = False, + mps_pipe_directory: str | None = None, +) -> ExecutionResult: + """Run whole pipeline items on one GPU with persistent local workers. + + The parent never initializes CUDA. Process workers may attach to an + externally managed MPS server. Thread workers each own a model, streams, + and CuPy allocator inside one supervised child process. Initialization + and each thread's first item are serialized; that first item remains in + the first round's outputs and elapsed time. One deadline bounds startup, + all rounds and close. A stuck child is terminated before return. + """ + + if executor not in {"processes", "threads"}: + raise ValueError("sharing executor must be processes or threads") + if type(workers_per_gpu) is not int or workers_per_gpu < 1: + raise ValueError("workers_per_gpu must be a positive integer") + if ( + isinstance(worker_timeout_sec, bool) + or not isinstance(worker_timeout_sec, (int, float)) + or not math.isfinite(worker_timeout_sec) + or worker_timeout_sec <= 0 + ): + raise ValueError("worker_timeout_sec must be positive") + if type(require_free_threaded) is not bool: + raise TypeError("require_free_threaded must be a boolean") + if require_free_threaded and executor != "threads": + raise ValueError( + "require_free_threaded requires the threads executor" + ) + if mps_pipe_directory is not None: + if executor != "processes": + raise ValueError( + "mps_pipe_directory requires the processes executor" + ) + if not isinstance(mps_pipe_directory, str) or not mps_pipe_directory: + raise ValueError("mps_pipe_directory must be a nonempty path") + mps_pipe_directory = str( + Path(mps_pipe_directory).expanduser().resolve() + ) + if benchmark is not None and not isinstance(benchmark, BenchmarkOptions): + raise TypeError("benchmark must be BenchmarkOptions or None") + visibility = os.environ.get("CUDA_VISIBLE_DEVICES", "") + tokens = visibility.split(",") + if ( + len(tokens) != 1 + or not tokens[0].strip() + or tokens[0].strip().startswith("-") + ): + raise ValueError("set CUDA_VISIBLE_DEVICES to exactly one GPU") + effective_run_id = run_id or new_run_id("shared-pipeline") + validate_identifier(effective_run_id, field="run_id") + started = time.perf_counter() + started_at = timestamp_utc() + spec = prepare_workload(0) + if not isinstance(spec, WorkloadSpec): + raise TypeError("prepare_workload must return WorkloadSpec") + worker_count = min(workers_per_gpu, len(spec.items)) + shards = partition_byte_balanced(spec.items, worker_count) + run_dir = Path(output_root).expanduser().resolve() / effective_run_id + prepare_run(run_dir, effective_run_id, spec, executor) + options = { + "workers_per_gpu": workers_per_gpu, + "worker_count": worker_count, + "worker_timeout_sec": worker_timeout_sec, + "benchmark": benchmark.to_payload() if benchmark else None, + "require_free_threaded": require_free_threaded, + "mps_pipe_directory": mps_pipe_directory, + "serialized_first_item": executor == "threads", + } + atomic_write_json(run_dir / "execution-options.json", options) + payload = { + **spec.identity_payload(), + "shards": [[item.to_dict() for item in shard] for shard in shards], + "require_free_threaded": require_free_threaded, + "mps_pipe_directory": mps_pipe_directory, + } + context = multiprocessing.get_context("spawn") + pairs = [context.Pipe() for _ in shards] + parents = [pair[0] for pair in pairs] + children = [pair[1] for pair in pairs] + processes: list[Any] = [] + owners: list[Any] = [] + errors: list[dict[str, Any]] = [] + rounds: list[dict[str, Any]] = [] + last_summary = None + startup_ready_sec = None + phase = "startup" + deadline = time.monotonic() + worker_timeout_sec + try: + groups = ( + [list(range(worker_count))] + if executor == "threads" + else [[index] for index in range(worker_count)] + ) + for group in groups: + process = context.Process( + target=_process_entry, + args=( + [children[index] for index in group], + payload, + group, + executor == "threads", + ), + name=f"cuphoton-{executor}-{group[0]}", + ) + process.start() + processes.append(process) + owners.extend([process] * len(group)) + for connection in children: + connection.close() + ready = _collect( + parents, owners, "ready", deadline - time.monotonic() + ) + for index, receipt in enumerate(ready): + if receipt.get("error") is not None: + raise RuntimeError( + f"worker {index} initialization failed: " + f"{receipt['error']}" + ) + provenance = receipt["provenance"] + if ( + provenance.get("worker_id") != index + or provenance.get("pid") != owners[index].pid + or type(provenance.get("thread_id")) is not int + or provenance["thread_id"] <= 0 + ): + raise RuntimeError( + "worker readiness differs from process ownership" + ) + if ( + executor == "threads" + and len({receipt["provenance"]["thread_id"] for receipt in ready}) + != worker_count + ): + raise RuntimeError( + "thread workers require distinct native threads" + ) + problems = audit_worker_provenance( + [receipt["provenance"] for receipt in ready], + backend=spec.backend, + expected_worker_count=worker_count, + gpu_groups=[0] * worker_count, + ) + if problems: + raise RuntimeError(str(problems)) + atomic_write_json(run_dir / "ready.json", {"workers": ready}) + startup_ready_sec = time.perf_counter() - started + for planned in ( + benchmark.rounds() + if benchmark + else (BenchmarkRound("measure", 0),) + ): + phase = planned.round_id + round_dir = ( + run_dir / "rounds" / planned.round_id + if benchmark + else run_dir + ) + round_id = ( + planned.run_id(effective_run_id) + if benchmark + else effective_run_id + ) + if benchmark: + prepare_run(round_dir, round_id, spec, executor) + round_started = time.perf_counter() + for connection in parents: + connection.send( + { + "phase": "round", + "run_id": round_id, + "run_dir": str(round_dir), + } + ) + results = _collect( + parents, owners, "round", deadline - time.monotonic() + ) + collected = time.perf_counter() + failures = [ + {"phase": phase, "worker_id": index, **result["error"]} + for index, result in enumerate(results) + if result.get("error") is not None + ] + worker_results = [ + read_json_mapping( + round_dir / "workers" / f"worker-{index:04d}.json" + ) + for index, result in enumerate(results) + if result.get("error") is None + ] + for worker_result in worker_results: + worker_id = worker_result.get("worker_id") + if ( + type(worker_id) is not int + or not 0 <= worker_id < worker_count + or worker_result.get("provenance") + != ready[worker_id]["provenance"] + ): + raise RuntimeError( + "worker provenance changed after readiness" + ) + last_summary = finalize_round( + round_dir, + round_id, + spec, + shards, + worker_results, + artifact_timeout_sec=0.01, + gpu_groups=[0] * worker_count, + ) + last_summary["errors"].extend(failures) + if failures: + last_summary["status"] = "failed" + timing = { + "batch_wall_sec": collected - round_started, + "finalization_sec": last_summary["finalization_sec"], + "worker_wall_max_sec": max( + ( + result.get("worker_wall_sec", 0.0) + for result in results + ), + default=0.0, + ), + } + last_summary.update(timing) + if benchmark: + atomic_write_json(round_dir / "summary.json", last_summary) + rounds.append( + { + **planned.to_payload(), + "status": last_summary["status"], + "summary_path": str( + (round_dir / "summary.json").relative_to(run_dir) + ), + **timing, + } + ) + if last_summary["status"] != "success": + break + except BaseException as exc: + errors.append({"phase": phase, **error_payload(exc)}) + if not isinstance(exc, Exception): + raise + finally: + close_started = time.perf_counter() + try: + for connection in parents: + connection.send({"phase": "close"}) + closed = _collect( + parents, + owners, + "closed", + min(deadline - time.monotonic(), 1.0) + if errors + else deadline - time.monotonic(), + ) + errors.extend( + {"phase": "close", "worker_id": index, **receipt["error"]} + for index, receipt in enumerate(closed) + if receipt.get("error") is not None + ) + except Exception as exc: + errors.append({"phase": "close", **error_payload(exc)}) + errors.extend( + _stop_processes( + processes, + 0.0 + if errors + else max(0.0, min(deadline - time.monotonic(), 5.0)), + ) + ) + for connection in (*parents, *children): + connection.close() + close_sec = time.perf_counter() - close_started + 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=executor, + run_id=effective_run_id, + manifest_sha256=spec.manifest_sha256, + worker_count=worker_count, + workers_per_gpu=workers_per_gpu, + physical_gpu_count=1, + 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() - started, + ) + atomic_write_json(run_dir / "summary.json", summary) + return ExecutionResult( + executor, + effective_run_id, + run_dir, + run_dir / "summary.json", + summary["status"], + summary, + ) diff --git a/tests/core/test_cli_contract.py b/tests/core/test_cli_contract.py index 45d41b04..ad4fdec1 100644 --- a/tests/core/test_cli_contract.py +++ b/tests/core/test_cli_contract.py @@ -71,9 +71,9 @@ def test_public_command_surface_counts_are_exact() -> None: assert per_group == [ ("xdr", 1, 0, 17, 0), - ("xfit", 3, 3, 36, 1), + ("xfit", 3, 3, 38, 1), ("xpois", 7, 7, 162, 1), - ("xscan", 44, 44, 224, 1), + ("xscan", 44, 44, 229, 1), ("xrep", 6, 6, 123, 1), ("xray", 33, 31, 392, 1), ] @@ -81,7 +81,7 @@ def test_public_command_surface_counts_are_exact() -> None: assert sum(item[1] for item in per_group) == 94 assert sum(item[1] + item[4] for item in per_group) == 99 assert sum(item[2] for item in per_group) == 91 - assert sum(item[3] for item in per_group) == 954 + assert sum(item[3] for item in per_group) == 961 def test_public_registry_order_and_component_derivations() -> None: diff --git a/tests/core/test_dragon.py b/tests/core/test_dragon.py index 34cecebf..9aa812f9 100644 --- a/tests/core/test_dragon.py +++ b/tests/core/test_dragon.py @@ -37,7 +37,9 @@ def __init__(self, options): "backend": options["backend"], "device_index": 0, "name": "fake-gpu", - "uuid": f"GPU-{self.placement.worker_id}", + "uuid": ( + f"GPU-00000000-0000-0000-0000-{self.placement.gpu_id:012x}" + ), "pci_bus_id": f"0000:{self.placement.gpu_id:02x}:00.0", "identity_error": None, } @@ -251,9 +253,13 @@ def _run( large=False, validator=None, finalizer=None, + workers_per_gpu=1, + items=None, + mps_pipe_directory=None, ): spec = WorkloadSpec( - items=( + items=items + or ( WorkItem("one", {"value": "x" * 300_000 if large else 1}, 12), WorkItem("two", {"value": 2}, 6), ), @@ -274,9 +280,219 @@ def _run( result_timeout_sec=result_timeout, worker_timeout_sec=worker_timeout, benchmark=benchmark, + workers_per_gpu=workers_per_gpu, + mps_pipe_directory=mps_pipe_directory, ) +def test_shared_gpu_workers_remain_persistent_and_keep_disjoint_shards( + monkeypatch, tmp_path +): + state = _install_runtime(monkeypatch) + result = _run( + tmp_path, + worker_count=None, + workers_per_gpu=2, + items=tuple( + WorkItem(f"item-{index}", {"value": index}) for index in range(5) + ), + benchmark=BenchmarkOptions(0, 2), + ) + assert result.status == "success", result.summary + assert result.summary["worker_count"] == 4 + assert result.summary["gpu_count"] == 2 + assert result.summary["workers_per_gpu"] == 2 + assert [value["gpu_id"] for value in result.summary["placements"]] == [ + 3, + 7, + 3, + 7, + ] + assert sorted(state.factories) == sorted(state.closed) == [0, 1, 2, 3] + assert len(state.calls) == 10 + for round_info in result.summary["benchmark"]["rounds"]: + summary = json.loads( + (result.run_dir / round_info["summary_path"]).read_text() + ) + assert summary["terminal_record_audit"]["ok"] + assert len(summary["records"]) == 5 + + +def test_shared_gpu_max_workers_is_still_total_process_cap( + monkeypatch, tmp_path +): + _install_runtime(monkeypatch) + result = _run( + tmp_path, + worker_count=3, + workers_per_gpu=4, + items=tuple( + WorkItem(f"item-{index}", {"value": index}) for index in range(5) + ), + ) + assert result.status == "success", result.summary + assert result.summary["worker_count"] == 3 + assert [value["gpu_id"] for value in result.summary["placements"]] == [ + 3, + 7, + 3, + ] + + +def test_shared_placements_visit_each_host_and_gpu_before_reusing(): + available = tuple( + dragon.Placement(index, host, gpu) + for index, (host, gpu) in enumerate((("a", 3), ("a", 7), ("b", 4))) + ) + selected = dragon._select_gpu_placements(available, 5, workers_per_gpu=2) + assert [(value.host, value.gpu_id) for value in selected] == [ + ("a", 3), + ("b", 4), + ("a", 7), + ("a", 3), + ("b", 4), + ] + + +@pytest.mark.parametrize("workers_per_gpu", [0, -1, True, 1.5]) +def test_invalid_sharing_count_fails_before_workload_preparation( + workers_per_gpu, +): + def unexpected(rank): + pytest.fail( + "invalid sharing count must fail before workload preparation" + ) + + with pytest.raises(ValueError, match="workers_per_gpu"): + dragon.run_dragon_work_items( + prepare_workload=unexpected, + output_root=None, + workers_per_gpu=workers_per_gpu, + ) + + +@pytest.mark.parametrize("connected", [True, False]) +def test_requested_mps_is_checked_after_worker_cuda_initialization( + monkeypatch, tmp_path, connected +): + state = _install_runtime(monkeypatch) + pipe_directory = str(tmp_path / "mps") + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "previous-directory") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "previous-device") + monkeypatch.setattr( + dragon, + "_mps_gpu_uuid", + lambda gpu_id: f"GPU-00000000-0000-0000-0000-{gpu_id:012x}", + ) + checked = [] + resolve = dragon.resolve_worker_factory + + def resolve_factory(reference): + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == pipe_directory + assert os.environ["CUDA_VISIBLE_DEVICES"].startswith("GPU-") + return resolve(reference) + + monkeypatch.setattr(dragon, "resolve_worker_factory", resolve_factory) + + def require_client(pid): + assert state.factories + assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == pipe_directory + checked.append(pid) + if not connected: + raise RuntimeError("worker did not connect to MPS") + return { + "pipe_directory": pipe_directory, + "server_pid": 456, + "client_pid": pid, + } + + monkeypatch.setitem( + sys.modules, + "cuphoton.core.mps", + SimpleNamespace(require_mps_client=require_client), + ) + result = _run(tmp_path, mps_pipe_directory=pipe_directory) + assert len(checked) == 2 + assert result.status == ("success" if connected else "failed") + assert sorted(state.closed) == [0, 1] + assert result.summary["mps_pipe_directory"] == pipe_directory + if connected: + assert all( + message["provenance"]["mps_pipe_directory"] == pipe_directory + and message["provenance"]["mps"] + == { + "pipe_directory": pipe_directory, + "server_pid": 456, + "client_pid": message["provenance"]["pid"], + } + for message in result.summary["ready_messages"] + ) + + +@pytest.mark.parametrize("defect", [None, "gpu", "client", "server", "pipe"]) +def test_mps_provenance_binds_actual_device_and_verified_client(defect): + requested_uuid = "GPU-00000000-0000-0000-0000-000000000003" + actual_uuid = "00000000000000000000000000000003" + value = { + "worker_id": 0, + "pid": 123, + "requested_host": "worker.example", + "hostname": "worker.example", + "requested_gpu_id": 3, + "requested_gpu_uuid": requested_uuid, + "cuda_visible_devices": requested_uuid, + "mps_pipe_directory": "/tmp/mps", + "mps": { + "pipe_directory": "/tmp/other" + if defect == "pipe" + else "/tmp/mps", + "server_pid": True if defect == "server" else 456, + "client_pid": 124 if defect == "client" else 123, + }, + "gpu": { + "backend": "cupy", + "device_index": 0, + "identity_error": None, + "uuid": "00000000000000000000000000000000" + if defect == "gpu" + else actual_uuid, + }, + } + assert dragon._valid_shard_provenance( + value, + placement=dragon.Placement(0, "worker.example", 3), + allow_loopback_alias=False, + backend="cupy", + mps_pipe_directory="/tmp/mps", + ) is (defect is None) + + +@pytest.mark.parametrize("reported_index", [3, 0]) +def test_mps_uuid_lookup_queries_only_the_requested_dragon_ordinal( + monkeypatch, reported_index +): + expected = "GPU-00000000-0000-0000-0000-000000000003" + + def query(argv, **options): + assert argv == [ + "nvidia-smi", + "--id=3", + "--query-gpu=index,uuid", + "--format=csv,noheader,nounits", + ] + assert options["timeout"] == 10 + return SimpleNamespace( + returncode=0, stdout=f"{reported_index}, {expected}\n" + ) + + monkeypatch.setattr(dragon.subprocess, "run", query) + if reported_index == 3: + assert dragon._mps_gpu_uuid(3) == expected + else: + with pytest.raises(RuntimeError, match="requested Dragon GPU"): + dragon._mps_gpu_uuid(3) + + def test_persistent_workers_use_bounded_descriptors_and_all_rounds( monkeypatch, tmp_path ): diff --git a/tests/core/test_execution.py b/tests/core/test_execution.py index 1ef62180..597af5cc 100644 --- a/tests/core/test_execution.py +++ b/tests/core/test_execution.py @@ -287,3 +287,60 @@ def test_provenance_host_equivalence_preserves_domains(hosts, duplicate): identities, backend="cupy", expected_worker_count=2 ) assert bool(errors) is duplicate + + +@pytest.mark.parametrize( + ("gpu_ids", "groups", "passes"), + [ + ((0, 1, 0, 1), (0, 1, 0, 1), True), + ((0, 1, 0, 1), None, False), + ((0, 1, 2, 1), (0, 1, 0, 1), False), + ((0, 0, 0, 0), (0, 1, 0, 1), False), + ], +) +def test_provenance_requires_exact_requested_gpu_sharing( + gpu_ids, groups, passes +): + values = [ + provenance(index, FakeWorker({"uuid": f"GPU-{gpu}"}).gpu_identity) + for index, gpu in enumerate(gpu_ids) + ] + # READY messages need not arrive in worker order. + errors = execution.audit_worker_provenance( + values[::-1], + backend="cupy", + expected_worker_count=4, + gpu_groups=groups, + ) + assert (not errors) is passes + + +@pytest.mark.parametrize("groups", [(0,), (0, True), (0, -1), (0, "1"), "01"]) +def test_provenance_rejects_invalid_gpu_groups(groups): + with pytest.raises(ValueError, match="gpu_groups"): + execution.audit_worker_provenance( + [], backend="cupy", expected_worker_count=2, gpu_groups=groups + ) + + +def test_shared_gpu_round_still_audits_and_finalizes_all_items(tmp_path): + run_dir = tmp_path / "run" + spec, shards, results = stage( + run_dir, + workload( + finalize_round=lambda path, records: {"count": len(records)} + ), + ) + results[1]["provenance"]["gpu"]["uuid"] = "GPU-0" + atomic_write_json(run_dir / "workers" / "worker-0001.json", results[1]) + report = execution.finalize_round( + run_dir, + "run", + spec, + shards, + results, + artifact_timeout_sec=0.01, + gpu_groups=(0, 0), + ) + assert report["status"] == "success", report["errors"] + assert report["result"] == {"count": 2} diff --git a/tests/core/test_executor_options.py b/tests/core/test_executor_options.py index 2fb17834..3111cb96 100644 --- a/tests/core/test_executor_options.py +++ b/tests/core/test_executor_options.py @@ -21,6 +21,13 @@ ("mpi", "worker_timeout_sec"), ("mpi", "result_timeout_sec"), ("dragon", "rank_setup_timeout_sec"), + ("local", "workers_per_gpu"), + ("local", "mps_pipe_directory"), + ("mpi", "workers_per_gpu"), + ("mpi", "mps_pipe_directory"), + ("processes", "max_workers"), + ("threads", "result_timeout_sec"), + ("threads", "mps_pipe_directory"), ], ) def test_rejects_flags_which_selected_executor_cannot_honor(executor, option): @@ -37,7 +44,9 @@ def test_local_default_does_not_enable_distributed_execution(): assert command.executor_options() == {} -@pytest.mark.parametrize("executor", ["dragon", "mpi"]) +@pytest.mark.parametrize( + "executor", ["dragon", "mpi", "processes", "threads"] +) def test_round_controls_are_opt_in(executor): command = ExecutorOptions() command.executor = executor diff --git a/tests/core/test_mps.py b/tests/core/test_mps.py new file mode 100644 index 00000000..04076018 --- /dev/null +++ b/tests/core/test_mps.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import subprocess +from types import SimpleNamespace + +import pytest + +from cuphoton.core.mps import require_mps_client + + +def test_mps_verification_requires_actual_client_membership( + tmp_path, monkeypatch +): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", str(tmp_path)) + commands = [] + + def control(argv, **kwargs): + assert argv == ["nvidia-cuda-mps-control"] + assert 0 < kwargs["timeout"] <= 10 + command = kwargs["input"] + commands.append(command) + output = { + "get_server_list\n": "41\n42\n", + "get_client_list 41\n": "51\n", + "get_client_list 42\n": "52\n53\n", + }[command] + return SimpleNamespace(returncode=0, stdout=output, stderr="") + + monkeypatch.setattr(subprocess, "run", control) + assert require_mps_client(53) == { + "pipe_directory": str(tmp_path), + "server_pid": 42, + "client_pid": 53, + } + assert len(commands) == 3 + with pytest.raises(RuntimeError, match="not connected"): + require_mps_client(99) + + +@pytest.mark.parametrize("failure", ["unavailable", "malformed", "timeout"]) +def test_mps_query_failure_is_not_silent_fallback( + tmp_path, monkeypatch, failure +): + monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", str(tmp_path)) + + def control(*args, **kwargs): + if failure == "timeout": + raise subprocess.TimeoutExpired(args[0], kwargs["timeout"]) + return SimpleNamespace( + returncode=1 if failure == "unavailable" else 0, + stdout="Server error", + stderr="No daemon", + ) + + monkeypatch.setattr(subprocess, "run", control) + with pytest.raises(RuntimeError, match="MPS"): + require_mps_client(53) diff --git a/tests/xscan/_sharing_support.py b/tests/xscan/_sharing_support.py new file mode 100644 index 00000000..68b1c08e --- /dev/null +++ b/tests/xscan/_sharing_support.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Importable CPU worker and CUDA stubs for real spawn lifecycle tests.""" + +from __future__ import annotations + +import os +import sys +import threading +import time +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace + +from cuphoton.core.bulk import atomic_write_json + + +class Worker: + def __init__(self, options): + self.options = options + self.owner = threading.get_ident() + self.identifier = f"{os.getpid()}-{threading.get_native_id()}" + self.trace = Path(options["trace"]) / f"{self.identifier}.txt" + self.record("initialize") + if options.get("failure") == "initialize": + raise ValueError("initialization failed") + self.gpu_identity = { + "backend": "cupy", + "device_index": 0, + "uuid": "GPU-shared", + "identity_error": None, + } + + def record(self, message): + assert self.owner == threading.get_ident() + with self.trace.open("a") as stream: + stream.write(message + "\n") + + def run_item(self, item, output_dir): + self.record(f"item:{item.item_id}") + time.sleep(self.options.get("item_delay", 0.0)) + if self.options.get("failure") == "hang": + time.sleep(30) + if self.options.get("failure") == "exit": + os._exit(7) + if self.options.get("failure") == "item": + raise ValueError("item failed") + 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": "cupy", + "device": "cuda:0", + "runtime": {"instance": self.identifier}, + "timings_sec": {"work": 0.0}, + "wall_sec": {"item": 0.0}, + } + + def close(self): + self.record("close") + if self.options.get("failure") == "close": + raise ValueError("close failed") + + +def factory(options): + return Worker(options) + + +def process_entry(connections, payload, worker_ids, threaded): + from cuphoton.xscan import sharing + + count = payload["options"].get("device_count", 1) + pool = SimpleNamespace( + malloc=lambda size: None, free_all_blocks=lambda: None + ) + cp = SimpleNamespace( + cuda=SimpleNamespace( + runtime=SimpleNamespace(getDeviceCount=lambda: count), + Device=lambda index: SimpleNamespace(use=lambda: None), + MemoryPool=lambda: pool, + using_allocator=lambda allocator: nullcontext(), + ) + ) + torch = SimpleNamespace( + cuda=SimpleNamespace( + device_count=lambda: count, + set_device=lambda index: None, + ) + ) + + def require_mps_client(pid): + if payload["options"].get("failure") == "mps": + raise RuntimeError("MPS client verification failed") + return { + "client_pid": pid, + "server_pid": 41, + "pipe_directory": os.environ["CUDA_MPS_PIPE_DIRECTORY"], + } + + sys.modules["cuphoton.core.mps"] = SimpleNamespace( + require_mps_client=require_mps_client + ) + sys.modules["cupy"] = cp + sys.modules["torch"] = torch + sharing._process_entry(connections, payload, worker_ids, threaded) diff --git a/tests/xscan/test_pipeline_executor.py b/tests/xscan/test_pipeline_executor.py index 99a2bc5c..b1b7dd0a 100644 --- a/tests/xscan/test_pipeline_executor.py +++ b/tests/xscan/test_pipeline_executor.py @@ -188,7 +188,9 @@ def test_preflight_rejects_changed_input_on_root(tmp_path): pipeline.prepare_pipeline_workload((item,), config) -@pytest.mark.parametrize("executor", ["dragon", "mpi"]) +@pytest.mark.parametrize( + "executor", ["dragon", "mpi", "processes", "threads"] +) def test_entrypoint_prepares_inside_runtime_callback( tmp_path, monkeypatch, executor ): @@ -202,7 +204,12 @@ def run(**kwargs): assert spec.options_payload["configuration"] == config.to_payload() return "result" - monkeypatch.setattr("cuphoton.core.executors.run_workload", run) + monkeypatch.setattr( + "cuphoton.xscan.sharing.run_shared_pipeline" + if executor in {"processes", "threads"} + else "cuphoton.core.executors.run_workload", + run, + ) assert ( pipeline.run_pipeline_manifest( executor=executor, @@ -379,3 +386,131 @@ def test_pipeline_command_rejects_wrong_executor_limits_before_runtime( ) != 0 ) + + +@pytest.mark.parametrize( + ("executor", "extra", "expected"), + [ + ("dragon", ["--workers-per-gpu", "4"], {"workers_per_gpu": 4}), + ( + "processes", + [ + "--workers-per-gpu", + "2", + "--mps-pipe-directory", + "/tmp/test-mps", + ], + {"workers_per_gpu": 2, "mps_pipe_directory": "/tmp/test-mps"}, + ), + ( + "threads", + ["--workers-per-gpu", "3", "--require-free-threaded"], + {"workers_per_gpu": 3, "require_free_threaded": True}, + ), + ], +) +def test_cli_selects_sharing_logic( + tmp_path, monkeypatch, executor, extra, expected +): + calls = [] + monkeypatch.setattr( + pipeline, + "run_pipeline_manifest", + lambda **kwargs: calls.append(kwargs), + ) + assert ( + run_component( + "xscan", + [ + "run-pipeline", + "--executor", + executor, + "--manifest", + str(tmp_path / "manifest.json"), + "--output-dir", + str(tmp_path / "runs"), + "--worker-timeout-sec", + "30", + *extra, + ], + ) + == 0 + ) + assert len(calls) == 1 + assert calls[0]["executor"] == executor + assert calls[0]["worker_timeout_sec"] == 30.0 + for name, value in expected.items(): + assert calls[0][name] == value + + +@pytest.mark.parametrize( + ("executor", "flags"), + [ + ("mpi", ["--workers-per-gpu", "2"]), + ("processes", ["--require-free-threaded"]), + ("threads", ["--mps-pipe-directory", "/tmp/test-mps"]), + ("threads", ["--workers-per-gpu", "0"]), + ], +) +def test_cli_rejects_invalid_sharing_before_launch( + tmp_path, monkeypatch, executor, flags +): + monkeypatch.setattr( + pipeline, + "run_pipeline_manifest", + lambda **kwargs: pytest.fail("invalid sharing flags reached runtime"), + ) + assert ( + run_component( + "xscan", + [ + "run-pipeline", + "--executor", + executor, + "--manifest", + str(tmp_path / "manifest.json"), + "--output-dir", + str(tmp_path / "runs"), + *flags, + ], + ) + != 0 + ) + + +def test_threads_reject_fits_before_scientific_preflight( + tmp_path, monkeypatch +): + path, config, item = _manifest(tmp_path) + monkeypatch.setattr( + pipeline, + "load_pipeline_manifest", + lambda *args, **kwargs: ( + config, + ( + SimpleNamespace( + reference=object(), + target=item.target, + variance=None, + fit_mask=None, + ), + ), + ), + ) + monkeypatch.setattr( + pipeline, + "prepare_pipeline_workload", + lambda *args, **kwargs: pytest.fail( + "FITS reached threaded preflight" + ), + ) + monkeypatch.setattr( + "cuphoton.xscan.sharing.run_shared_pipeline", + lambda **kwargs: kwargs["prepare_workload"](0), + ) + with pytest.raises(ValueError, match="prepared NPY"): + pipeline.run_pipeline_manifest( + executor="threads", + manifest_path=path, + output_root=tmp_path / "runs", + ) diff --git a/tests/xscan/test_sharing.py b/tests/xscan/test_sharing.py new file mode 100644 index 00000000..60570237 --- /dev/null +++ b/tests/xscan/test_sharing.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +import multiprocessing +import sys +import time +from pathlib import Path + +import pytest + +from cuphoton.core.benchmark import BenchmarkOptions +from cuphoton.core.bulk import WorkItem, read_json_mapping +from cuphoton.core.execution import WorkloadSpec +from cuphoton.xscan import sharing + + +@pytest.fixture +def cpu_runtime(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).parent)) + support = importlib.import_module("_sharing_support") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-shared") + monkeypatch.setattr(sharing, "_process_entry", support.process_entry) + return support + + +def workload(tmp_path, cpu_runtime, **options): + trace = tmp_path / "trace" + trace.mkdir() + return WorkloadSpec( + items=tuple( + WorkItem(str(index), {"value": index}, 1) for index in range(4) + ), + options_payload={"trace": str(trace), **options}, + manifest_payload={"test": True}, + input_identity_payload={}, + manifest_sha256="a" * 64, + backend="cupy", + worker_factory=cpu_runtime.factory, + ) + + +@pytest.mark.parametrize("executor", ["processes", "threads"]) +def test_spawn_reuses_workers_and_audits_every_round( + tmp_path, cpu_runtime, executor +): + spec = workload(tmp_path, cpu_runtime) + modules_before = set(sys.modules) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + run_id="persistent", + executor=executor, + workers_per_gpu=2, + worker_timeout_sec=10, + benchmark=BenchmarkOptions(warmup_rounds=1, measure_rounds=2), + ) + assert result.status == "success", result.summary + assert result.summary["worker_count"] == 2 + traces = [ + path.read_text().splitlines() + for path in (tmp_path / "trace").iterdir() + ] + assert len(traces) == 2 + assert all( + trace[0] == "initialize" and trace[-1] == "close" for trace in traces + ) + assert all(len(trace) == 8 for trace in traces) + for planned in BenchmarkOptions(1, 2).rounds(): + directory = result.run_dir / "rounds" / planned.round_id + summary = read_json_mapping(directory / "summary.json") + assert summary["terminal_record_audit"]["ok"] + assert len(list((directory / "records").glob("*.json"))) == 4 + assert {"cupy", "torch"} & (set(sys.modules) - modules_before) == set() + assert not multiprocessing.active_children() + + +@pytest.mark.parametrize("executor", ["processes", "threads"]) +@pytest.mark.parametrize( + "failure", ["initialize", "item", "close", "hang", "exit"] +) +def test_spawn_failures_do_not_report_success_or_leave_children( + tmp_path, cpu_runtime, executor, failure +): + spec = workload(tmp_path, cpu_runtime, failure=failure) + started = time.monotonic() + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor=executor, + workers_per_gpu=2, + worker_timeout_sec=0.7 if failure == "hang" else 10, + ) + assert result.status == "failed" + assert read_json_mapping(result.summary_path)["status"] == "failed" + assert time.monotonic() - started < 12 + assert not multiprocessing.active_children() + + +def test_actual_device_count_and_free_threading_are_verified_after_imports( + tmp_path, cpu_runtime +): + spec = workload(tmp_path, cpu_runtime, device_count=2) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + run_id="two-gpus", + executor="processes", + worker_timeout_sec=10, + ) + assert result.status == "failed" + assert "exactly one visible GPU" in str(result.summary) + if sharing._gil_enabled(): + spec = WorkloadSpec( + **{ + **spec.__dict__, + "options_payload": { + **spec.options_payload, + "device_count": 1, + }, + } + ) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + run_id="gil-enabled", + executor="threads", + require_free_threaded=True, + worker_timeout_sec=10, + ) + assert result.status == "failed" + assert "GIL to remain disabled" in str(result.summary) + + +def test_worker_count_is_capped_by_item_count(tmp_path, cpu_runtime): + spec = workload(tmp_path, cpu_runtime) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="threads", + workers_per_gpu=8, + worker_timeout_sec=10, + ) + assert result.status == "success", result.summary + assert result.summary["worker_count"] == 4 + assert result.summary["workers_per_gpu"] == 8 + + +@pytest.mark.parametrize( + "options, message", + [ + ({"workers_per_gpu": 0}, "positive integer"), + ({"worker_timeout_sec": float("nan")}, "must be positive"), + ({"require_free_threaded": True}, "threads executor"), + ( + {"executor": "threads", "mps_pipe_directory": "/tmp/mps"}, + "processes executor", + ), + ], +) +def test_invalid_options_do_not_prepare_workload(tmp_path, options, message): + def prepare(rank): + pytest.fail("invalid launch called prepare_workload") + + with pytest.raises(ValueError, match=message): + sharing.run_shared_pipeline( + prepare_workload=prepare, + output_root=tmp_path, + **{"executor": "processes", **options}, + ) + + +@pytest.mark.parametrize("failure", [None, "mps"]) +def test_explicit_mps_requires_verified_client_receipt( + tmp_path, cpu_runtime, failure +): + spec = workload(tmp_path, cpu_runtime, failure=failure) + pipe = str(tmp_path / "mps") + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="processes", + mps_pipe_directory=pipe, + worker_timeout_sec=10, + ) + if failure: + assert result.status == "failed" + assert "MPS client verification failed" in str(result.summary) + else: + assert result.status == "success", result.summary + provenance = read_json_mapping(result.run_dir / "ready.json")[ + "workers" + ][0]["provenance"] + assert provenance["mps_pipe_directory"] == pipe + assert provenance["mps"]["server_pid"] == 41 + assert provenance["mps"]["client_pid"] == provenance["pid"] + + +def test_timeout_bounds_all_rounds_together(tmp_path, cpu_runtime): + spec = workload(tmp_path, cpu_runtime, item_delay=0.15) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="processes", + worker_timeout_sec=3.0, + benchmark=BenchmarkOptions(measure_rounds=8), + ) + assert result.status == "failed" + report = result.summary["benchmark"] + assert 1 <= len(report["rounds"]) < 8 + assert report["measured_batch_wall_sec"] is None + assert "timed out" in str(result.summary) + assert not multiprocessing.active_children() From 93b0d1bb2bd7aed2d8dab374def463b85380999f Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 09:28:11 -0700 Subject: [PATCH 2/3] Preserve worker results and cleanup after sharing failures Keep successful first-item results when another thread breaks the barrier. Close each started worker independently, drain pending receipts, and allow bounded graceful shutdown after failures. Cover isolated worker and partial-start failures, and clarify the sharing options in the CLI and component documentation. Signed-off-by: Trent Nelson --- docs/cli.md | 10 +-- docs/components/xfit.md | 6 ++ docs/dragon.md | 5 +- src/cuphoton/core/cli/executor.py | 2 +- src/cuphoton/xscan/sharing.py | 73 +++++++++++++++------- tests/xscan/_sharing_support.py | 18 +++++- tests/xscan/test_sharing.py | 100 ++++++++++++++++++++++++++++++ 7 files changed, 185 insertions(+), 29 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 2854bfd7..50cd125e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -121,10 +121,12 @@ standalone raw-comparison and Alard--Lupton review servers are available as command-specific contracts. See [xScan](components/xscan.md). `infer-real-bogus --executor dragon|mpi` distributes complete minibatches; -`run-pipeline --executor dragon|mpi` distributes complete image pairs through -xPois, xFit and xScan. Optional round flags retain the workers across passes. -These commands require a shared input/output filesystem and the matching -runtime launcher. +`run-pipeline --executor dragon|mpi|processes|threads` distributes complete +image pairs through xPois, xFit and xScan. Optional round flags retain the +workers across passes. Dragon and MPI require a shared input/output +filesystem and the matching runtime launcher. The local `processes` and +`threads` pipeline executors share one visible GPU; see +[GPU sharing](components/xscan.md#share-a-gpu-between-image-pairs). `infer-real-bogus` consumes packaged datasets. Direct FITS candidate scoring uses the CUDA Python API `cuphoton.xscan.fits_inference.predict_fits`. diff --git a/docs/components/xfit.md b/docs/components/xfit.md index 5fcc8c96..181a2589 100644 --- a/docs/components/xfit.md +++ b/docs/components/xfit.md @@ -355,6 +355,12 @@ IDs and input order are restored in the merged artifacts. The task count must be at least the MPI rank count. Dragon uses the smaller of the requested worker count and task count. +With Dragon, `--workers-per-gpu` allows multiple workers per GPU and +`--mps-pipe-directory` requires each worker to connect to an existing MPS v2 +service. `--max-workers` caps the total worker count. See +[GPU sharing](xscan.md#share-a-gpu-between-image-pairs) for service setup and +connection requirements. + The same executor commands accept a FITS candidate manifest. Planning checks headers, candidate bounds and source hashes without decoding image pixels. Each bound worker reads the candidate region and retains its stamps during diff --git a/docs/dragon.md b/docs/dragon.md index fbf95688..d0c18295 100644 --- a/docs/dragon.md +++ b/docs/dragon.md @@ -35,8 +35,9 @@ The shared executor accepts `--workers-per-gpu` to place several independent workers on each GPU; its default is one. `--max-workers` remains a total process cap. See [GPU sharing](components/xscan.md#share-a-gpu-between-image-pairs) for CLI and Python examples, local process/thread modes, and external MPS -connection checks. This option applies to the shared component executor; -the standalone XPOIS batch API retains its existing placement behavior. +connection checks. These sharing and MPS options apply to `xfit fit-dipoles`, +`xscan infer-real-bogus`, and `xscan run-pipeline` with `--executor dragon`. +The standalone XPOIS batch API retains its existing placement behavior. ### Launch descriptors diff --git a/src/cuphoton/core/cli/executor.py b/src/cuphoton/core/cli/executor.py index 647d2dda..fc1d04b9 100644 --- a/src/cuphoton/core/cli/executor.py +++ b/src/cuphoton/core/cli/executor.py @@ -125,7 +125,7 @@ def executor_options(self) -> dict: elif self.executor == "threads": invalid = {**dragon_only, **mpi, **mps} else: - raise CommandError("executor must be local, dragon, or mpi") + raise CommandError(f"unsupported executor: {self.executor}") supplied = [ "--" + name.replace("_", "-") for name, value in invalid.items() diff --git a/src/cuphoton/xscan/sharing.py b/src/cuphoton/xscan/sharing.py index 0726a255..c5e17b0a 100644 --- a/src/cuphoton/xscan/sharing.py +++ b/src/cuphoton/xscan/sharing.py @@ -94,7 +94,11 @@ def run_item(self, item: WorkItem, output_dir: Path) -> Mapping[str, Any]: except BaseException: self.barrier.abort() raise - self.barrier.wait() + try: + self.barrier.wait() + except threading.BrokenBarrierError: + # Another thread's failure must not discard this result. + pass return result return self.worker.run_item(item, output_dir) @@ -270,6 +274,12 @@ def _collect( continue message = connection.recv() if message.get("phase") != phase: + if phase == "closed" and message.get("phase") in { + "ready", + "round", + }: + # A failed collection can leave peer receipts queued. + continue raise RuntimeError( f"worker {worker_id} sent an unexpected {phase} receipt" ) @@ -565,30 +575,51 @@ def run_shared_pipeline( raise finally: close_started = time.perf_counter() - try: - for connection in parents: + for connection in children: + connection.close() + active_parents = parents[: len(owners)] + close_deadline = ( + min(deadline, time.monotonic() + 2.0) if errors else deadline + ) + for connection in active_parents: + try: connection.send({"phase": "close"}) - closed = _collect( - parents, - owners, - "closed", - min(deadline - time.monotonic(), 1.0) - if errors - else deadline - time.monotonic(), - ) - errors.extend( - {"phase": "close", "worker_id": index, **receipt["error"]} - for index, receipt in enumerate(closed) - if receipt.get("error") is not None - ) - except Exception as exc: - errors.append({"phase": "close", **error_payload(exc)}) + except (BrokenPipeError, EOFError, OSError): + # A failed worker may already have sent its closed receipt. + pass + for index, (connection, owner) in enumerate( + zip(active_parents, owners, strict=True) + ): + try: + receipt = _collect( + [connection], + [owner], + "closed", + close_deadline - time.monotonic(), + )[0] + if receipt.get("error") is not None: + errors.append( + { + "phase": "close", + "worker_id": index, + **receipt["error"], + } + ) + except Exception as exc: + errors.append( + { + "phase": "close", + "worker_id": index, + **error_payload(exc), + } + ) errors.extend( _stop_processes( processes, - 0.0 - if errors - else max(0.0, min(deadline - time.monotonic(), 5.0)), + max( + 0.0, + min(deadline - time.monotonic(), 2.0 if errors else 5.0), + ), ) ) for connection in (*parents, *children): diff --git a/tests/xscan/_sharing_support.py b/tests/xscan/_sharing_support.py index 68b1c08e..e22ae31d 100644 --- a/tests/xscan/_sharing_support.py +++ b/tests/xscan/_sharing_support.py @@ -45,7 +45,9 @@ def run_item(self, item, output_dir): time.sleep(30) if self.options.get("failure") == "exit": os._exit(7) - if self.options.get("failure") == "item": + if self.options.get("failure") == "item" or ( + self.options.get("fail_item") == item.item_id + ): raise ValueError("item failed") output_dir.mkdir() atomic_write_json(output_dir / "summary.json", item.payload) @@ -72,6 +74,20 @@ def factory(options): def process_entry(connections, payload, worker_ids, threaded): from cuphoton.xscan import sharing + worker_loop = sharing._worker_loop + + def run_worker(connection, worker_payload, worker_id, *args): + options = dict(worker_payload["options"]) + if options.get("failure_worker_id", worker_id) != worker_id: + options.pop("failure", None) + worker_loop( + connection, + {**worker_payload, "options": options}, + worker_id, + *args, + ) + + sharing._worker_loop = run_worker count = payload["options"].get("device_count", 1) pool = SimpleNamespace( malloc=lambda size: None, free_all_blocks=lambda: None diff --git a/tests/xscan/test_sharing.py b/tests/xscan/test_sharing.py index 60570237..478ebf72 100644 --- a/tests/xscan/test_sharing.py +++ b/tests/xscan/test_sharing.py @@ -9,6 +9,7 @@ import sys import time from pathlib import Path +from types import SimpleNamespace import pytest @@ -135,6 +136,105 @@ def test_actual_device_count_and_free_threading_are_verified_after_imports( assert "GIL to remain disabled" in str(result.summary) +def test_threads_keep_successful_items_when_a_first_item_fails( + tmp_path, cpu_runtime +): + spec = workload(tmp_path, cpu_runtime, fail_item="0") + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="threads", + workers_per_gpu=2, + worker_timeout_sec=10, + ) + assert result.status == "failed" + records = [ + read_json_mapping(path) + for path in (result.run_dir / "records").glob("*.json") + ] + assert len(records) == len(spec.items) + assert { + record["item_id"] + for record in records + if record["status"] == "failed" + } == {"0"} + assert result.summary["terminal_record_audit"]["ok"] + assert "BrokenBarrierError" not in str(result.summary) + assert not multiprocessing.active_children() + + +@pytest.mark.parametrize("executor", ["processes", "threads"]) +def test_initialization_failure_still_closes_healthy_workers( + tmp_path, cpu_runtime, executor +): + spec = workload( + tmp_path, cpu_runtime, failure="initialize", failure_worker_id=0 + ) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor=executor, + workers_per_gpu=2, + worker_timeout_sec=10, + ) + assert result.status == "failed" + assert "initialization failed" in str(result.summary) + traces = sorted( + path.read_text().splitlines() + for path in (tmp_path / "trace").iterdir() + ) + assert traces == [["initialize"], ["initialize", "close"]] + assert all( + error["phase"] == "startup" for error in result.summary["errors"] + ) + assert not multiprocessing.active_children() + + +@pytest.mark.parametrize("failed_start", [0, 1]) +def test_process_start_failure_preserves_error_and_closes_started_workers( + tmp_path, cpu_runtime, monkeypatch, failed_start +): + spec = workload(tmp_path, cpu_runtime) + context = multiprocessing.get_context("spawn") + created = 0 + + def process(**kwargs): + nonlocal created + child = context.Process(**kwargs) + if created == failed_start: + + def fail_start(): + raise OSError("process could not start") + + monkeypatch.setattr(child, "start", fail_start) + created += 1 + return child + + monkeypatch.setattr( + sharing.multiprocessing, + "get_context", + lambda method: SimpleNamespace(Pipe=context.Pipe, Process=process), + ) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="processes", + workers_per_gpu=2, + worker_timeout_sec=10, + ) + assert result.status == "failed" + assert "process could not start" in str(result.summary) + assert all( + error["phase"] == "startup" for error in result.summary["errors"] + ) + traces = [ + path.read_text().splitlines() + for path in (tmp_path / "trace").iterdir() + ] + assert traces == [["initialize", "close"]] * failed_start + assert not multiprocessing.active_children() + + def test_worker_count_is_capped_by_item_count(tmp_path, cpu_runtime): spec = workload(tmp_path, cpu_runtime) result = sharing.run_shared_pipeline( From 04a2e6107274162b36b95f4c32a0dd70b44bf002 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Fri, 9 Oct 2026 12:13:20 -0700 Subject: [PATCH 3/3] Test per-thread allocator ownership during worker cleanup Signed-off-by: Trent Nelson --- tests/xscan/_sharing_support.py | 50 +++++++++++++++++++++++++++++---- tests/xscan/test_sharing.py | 22 +++++++++++++++ 2 files changed, 66 insertions(+), 6 deletions(-) diff --git a/tests/xscan/_sharing_support.py b/tests/xscan/_sharing_support.py index e22ae31d..27f59f83 100644 --- a/tests/xscan/_sharing_support.py +++ b/tests/xscan/_sharing_support.py @@ -10,12 +10,14 @@ import sys import threading import time -from contextlib import nullcontext +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from cuphoton.core.bulk import atomic_write_json +_allocator_state = threading.local() + class Worker: def __init__(self, options): @@ -23,6 +25,7 @@ def __init__(self, options): self.owner = threading.get_ident() self.identifier = f"{os.getpid()}-{threading.get_native_id()}" self.trace = Path(options["trace"]) / f"{self.identifier}.txt" + self.pool = getattr(_allocator_state, "pool", None) self.record("initialize") if options.get("failure") == "initialize": raise ValueError("initialization failed") @@ -35,6 +38,10 @@ def __init__(self, options): def record(self, message): assert self.owner == threading.get_ident() + if self.options.get("pool_trace"): + assert self.pool is not None + assert _allocator_state.pool is self.pool + self.pool.record(message) with self.trace.open("a") as stream: stream.write(message + "\n") @@ -89,15 +96,46 @@ def run_worker(connection, worker_payload, worker_id, *args): sharing._worker_loop = run_worker count = payload["options"].get("device_count", 1) - pool = SimpleNamespace( - malloc=lambda size: None, free_all_blocks=lambda: None - ) + + class RecordingPool: + def __init__(self): + self.owner = threading.get_ident() + self.trace = payload["options"].get("pool_trace") + self.record("create") + + def record(self, message): + assert self.owner == threading.get_ident() + if self.trace: + path = Path(self.trace) / f"{id(self)}.txt" + with path.open("a") as stream: + stream.write(message + "\n") + + def malloc(self, size): + assert _allocator_state.pool is self + + def free_all_blocks(self): + assert _allocator_state.pool is self + self.record("release") + + @contextmanager + def using_allocator(allocator): + assert getattr(_allocator_state, "pool", None) is None + pool = allocator.__self__ + _allocator_state.pool = pool + pool.record("enter") + try: + yield + finally: + assert _allocator_state.pool is pool + del _allocator_state.pool + pool.record("exit") + cp = SimpleNamespace( cuda=SimpleNamespace( runtime=SimpleNamespace(getDeviceCount=lambda: count), Device=lambda index: SimpleNamespace(use=lambda: None), - MemoryPool=lambda: pool, - using_allocator=lambda allocator: nullcontext(), + MemoryPool=RecordingPool, + using_allocator=using_allocator, ) ) torch = SimpleNamespace( diff --git a/tests/xscan/test_sharing.py b/tests/xscan/test_sharing.py index 478ebf72..fed55e30 100644 --- a/tests/xscan/test_sharing.py +++ b/tests/xscan/test_sharing.py @@ -79,6 +79,28 @@ def test_spawn_reuses_workers_and_audits_every_round( assert not multiprocessing.active_children() +def test_threads_own_pools_through_close_and_release(tmp_path, cpu_runtime): + pools = tmp_path / "pools" + pools.mkdir() + spec = workload(tmp_path, cpu_runtime, pool_trace=str(pools)) + result = sharing.run_shared_pipeline( + prepare_workload=lambda rank: spec, + output_root=tmp_path, + executor="threads", + workers_per_gpu=2, + worker_timeout_sec=10, + ) + assert result.status == "success", result.summary + traces = [path.read_text().splitlines() for path in pools.iterdir()] + assert len(traces) == 2 + for trace in traces: + assert trace[:3] == ["create", "enter", "initialize"] + assert len(trace[3:-3]) == 2 + assert all(event.startswith("item:") for event in trace[3:-3]) + assert trace[-3:] == ["close", "release", "exit"] + assert not multiprocessing.active_children() + + @pytest.mark.parametrize("executor", ["processes", "threads"]) @pytest.mark.parametrize( "failure", ["initialize", "item", "close", "hang", "exit"]