Skip to main content

Module gpus

Module gpus 

Source
Expand description

--gpus flag parsing + single-host cluster envelope synthesis (loopback, one host, N ranks). Used when --gpus is set on a cluster-aware command and no cluster: block is configured in YAML. --gpus flag parsing + single-host cluster envelope synthesis.

The --gpus flag has uniform semantics (“use these GPUs”) but the mechanism depends on the command kind:

  • Cluster-aware commands (cluster: true): N >= 2 GPUs trigger synthesis of a single-host cluster envelope (master=127.0.0.1, lo transport, one host with N ranks) and spawn-per-rank via the existing launcher (see crate::cluster::prepare_cluster_env). The library inside each spawned process reads the envelope from FLODL_INTERNAL_CLUSTER_JSON and uses the same code path as multi-host. N = 1 is degenerate — no synthesis, just runs single-process on that device.

  • Non-cluster commands (test, clippy, etc.): --gpus sets CUDA_VISIBLE_DEVICES on the single child process. No envelope, no spawning. Tests internally manage their own multi-rank coordination (typically via the threaded NcclRankComm pattern in unit tests).

Caller (main.rs) decides which mechanism applies based on whether the resolved command’s cluster: chain enables dispatch.

Enums§

GpusSpec
Parsed --gpus argument value.

Functions§

apply_cuda_visible_devices
Set CUDA_VISIBLE_DEVICES to restrict the spawned process to the given physical CUDA device indices.
count_visible_gpus_via_nvidia_smi
Count visible CUDA devices via nvidia-smi -L.
synthesize_local_cluster
Build a ClusterConfig for single-host loopback from a list of physical CUDA device indices.