unilab.training.cuda_process_sharing

Fail-closed validation for CUDA process-sharing execution modes.

The off-policy learner and collector are separate processes on one rank-local GPU. CUDA MPS is an execution-sharing mode for that required topology: it lets their independent CUDA contexts overlap kernels instead of being arbitrated by coarse driver time-slicing.

This module never starts or stops an MPS daemon. An explicit request is either validated against the host state or rejected before UniLab constructs a learner, collector, or environment.

An unconfigured (null) request still records ambient host MPS state: an environment-level daemon reachable through CUDA_MPS_PIPE_DIRECTORY (or the default pipe directory) multiplexes this rank’s CUDA contexts transparently, so every run must be able to self-report that state. Ambient observation is best-effort and never fails closed; unresolvable facts record None.

Functions

apply_cuda_process_sharing_manifest(...)

Record process-sharing evidence in a producer-owned manifest mapping.

canonicalize_gpu_uuid(value)

Return the case- and dash-insensitive NVIDIA GPU UUID comparison token.

probe_cuda_process_sharing(requested, ...[, ...])

Validate one rank's requested CUDA process-sharing mode.

resolve_mps_visibility_entries(entries, *[, ...])

Return CUDA_VISIBLE_DEVICES entries with numeric indices as UUIDs.

Classes

CudaProcessSharingEvidence

Structured per-rank evidence for a requested mode plus ambient state.
class unilab.training.cuda_process_sharing.CudaProcessSharingEvidence[source]

Bases: object

Structured per-rank evidence for a requested mode plus ambient state.

pipe_directory_env is the raw CUDA_MPS_PIPE_DIRECTORY value (None when unset); pipe_directory/control_pipe/control_pipe_kind/ server_pid describe the effective pipe this rank’s CUDA contexts would use. ambient_server is true only for an unconfigured run whose effective control pipe answered a daemon query — the transparent-MPS case.

Parameters:
configured: str | None
effective: str | None
validated: bool
rank: int = 0
world_size: int = 1
learner_device: str | None = None
collector_device: str | None = None
learner_gpu_uuid: str | None = None
collector_gpu_uuid: str | None = None
pipe_directory_env: str | None = None
pipe_directory: str | None = None
control_pipe: str | None = None
control_pipe_kind: str | None = None
server_pid: int | None = None
ambient_server: bool = False
to_dict()[source]

Return the JSON-shaped diagnostics stored in a runtime manifest.

Return type:

dict[str, Any]

__init__(configured, effective, validated, rank=0, world_size=1, learner_device=None, collector_device=None, learner_gpu_uuid=None, collector_gpu_uuid=None, pipe_directory_env=None, pipe_directory=None, control_pipe=None, control_pipe_kind=None, server_pid=None, ambient_server=False)
Parameters:
unilab.training.cuda_process_sharing.apply_cuda_process_sharing_manifest(runtime_manifest, evidence)[source]

Record process-sharing evidence in a producer-owned manifest mapping.

Parameters:
Return type:

None

unilab.training.cuda_process_sharing.canonicalize_gpu_uuid(value)[source]

Return the case- and dash-insensitive NVIDIA GPU UUID comparison token.

Torch, nvidia-smi, and CUDA_VISIBLE_DEVICES expose the same UUID with different cases, dashes, and an optional GPU- prefix; equality checks must compare this token rather than the raw strings. The persisted/display form stays GPU-<id> (see cuda_mps_cli).

Parameters:

value (str)

Return type:

str

unilab.training.cuda_process_sharing.probe_cuda_process_sharing(requested, learner_device, collector_device, *, backend, rank=0, world_size=1, torch_module=None, run_command=<function run>)[source]

Validate one rank’s requested CUDA process-sharing mode.

Parameters:

  • requested (Any) – Owner value null or mps.

  • learner_device (str) – Rank-local learner CUDA device.

  • collector_device (str | None) – Backend-bound collector CUDA device, when applicable.

  • backend (str) – Configured owner backend identity.

  • rank (int) – Current off-policy DP rank.

  • world_size (int) – Current single-host off-policy DP world size.

  • torch_module (Any | None) – Injectable Torch module for deterministic tests.

  • run_command (Callable[..., Any]) – Injectable subprocess runner for deterministic tests.

Return type:

CudaProcessSharingEvidence

Returns:

Evidence suitable for a per-rank runtime-manifest diagnostic section.

Raises:

ValueError – The first unmet prerequisite for an explicit request.

unilab.training.cuda_process_sharing.resolve_mps_visibility_entries(entries, *, run_command=<function run>)[source]

Return CUDA_VISIBLE_DEVICES entries with numeric indices as UUIDs.

A per-GPU MPS server enumerates only its own device, so a client whose one-entry visibility mask holds a nonzero numeric index overruns that namespace and fails CUDA initialization with “No CUDA GPUs are available”. UUID entries survive that server-side remapping, so an explicit training.cuda_process_sharing=mps launch normalizes visibility before CUDA latches the mask. Non-numeric entries (UUID, MIG, -1) pass through unchanged; the returned UUIDs keep the nvidia-smi display form.

Raises:

ValueError – A numeric entry could not be mapped to a host GPU UUID.

Parameters:
Return type:

tuple[str, ...]