diff --git a/docs/user/intro.md b/docs/user/intro.md index 2ecd180..b65072d 100644 --- a/docs/user/intro.md +++ b/docs/user/intro.md @@ -196,7 +196,7 @@ This user guide reflects actual code behavior. Not all planned features may be a | Named inputs (paths, images, etc.) | [Artifacts](./artifacts.md) | | Compute backends (local, Slurm) | [Backends](./backends.md) | | Task launch methods (bash, srun, containers) | [Operators](./operators.md) | -| Node/GPU placement, CUDA_VISIBLE_DEVICES | [Resources](./resources.md) | +| Node/GPU placement, NVIDIA_VISIBLE_DEVICES | [Resources](./resources.md) | | Parallel/sequential replicas, sweeps | [Replicas](./replicas.md) | | Composable configs, sweeps, missable tasks | [Modular Workflows](./modular-workflows.md) | | Readiness/failure gates for services | [Probes](./probes.md) | diff --git a/docs/user/quick-reference.md b/docs/user/quick-reference.md index 12508cf..06e5e31 100644 --- a/docs/user/quick-reference.md +++ b/docs/user/quick-reference.md @@ -192,7 +192,7 @@ For detailed explanations and examples, see [Configuration](./configuration.md). |-------|----------|------|---------|-------------| | `nodes.indices` | | list[int / expr] | `null` | Specific node indices (e.g. `[0]`). | | `nodes.count` | | int / expr | `null` | Number of nodes. | -| `gpus.count` | Yes | int / expr | — | Number of GPUs (sets `CUDA_VISIBLE_DEVICES`). | +| `gpus.count` | Yes | int / expr | — | Number of GPUs (sets `NVIDIA_VISIBLE_DEVICES`). | ## Task Replicas @@ -273,7 +273,7 @@ These are automatically set by sflow and available in every task script. | `SFLOW_REPLICA_INDEX` | Zero-based replica index (`0`, `1`, `2`, ...). | | `SFLOW_TASK_ASSIGNED_NODE_NAMES` | Comma-separated hostnames of nodes assigned to this task. | | `SFLOW_TASK_ASSIGNED_NODE_IPS` | Comma-separated IP addresses of nodes assigned to this task. | -| `CUDA_VISIBLE_DEVICES` | Comma-separated GPU indices allocated to this task (set when `resources.gpus.count` is used). | +| `NVIDIA_VISIBLE_DEVICES` | Comma-separated GPU indices allocated to this task (set when `resources.gpus.count` is used). | In addition, all resolved `variables` and `artifacts` paths are injected as environment variables accessible via `${VAR_NAME}` in scripts. diff --git a/docs/user/replicas.md b/docs/user/replicas.md index 57c0762..0fd3e24 100644 --- a/docs/user/replicas.md +++ b/docs/user/replicas.md @@ -161,7 +161,8 @@ If `CONCURRENCY` has `domain: [8, 16, 32, 64]` (4 values), this works. If the do ## GPU Allocation with Replicas -When combining replicas with GPU resources, sflow automatically slices `CUDA_VISIBLE_DEVICES` for each replica: +When combining replicas with GPU resources, sflow automatically slices +`NVIDIA_VISIBLE_DEVICES` for each replica: ```yaml backends: @@ -179,14 +180,14 @@ workflow: gpus: count: 2 script: - - echo "My GPUs: $CUDA_VISIBLE_DEVICES" + - echo "My GPUs: $NVIDIA_VISIBLE_DEVICES" ``` Each replica gets 2 GPUs: -- `worker_0`: CUDA_VISIBLE_DEVICES=0,1 -- `worker_1`: CUDA_VISIBLE_DEVICES=2,3 -- `worker_2`: CUDA_VISIBLE_DEVICES=4,5 -- `worker_3`: CUDA_VISIBLE_DEVICES=6,7 +- `worker_0`: `NVIDIA_VISIBLE_DEVICES=0,1` +- `worker_1`: `NVIDIA_VISIBLE_DEVICES=2,3` +- `worker_2`: `NVIDIA_VISIBLE_DEVICES=4,5` +- `worker_3`: `NVIDIA_VISIBLE_DEVICES=6,7` ## Real-World Examples diff --git a/docs/user/resources.md b/docs/user/resources.md index a595e1d..2f828e1 100644 --- a/docs/user/resources.md +++ b/docs/user/resources.md @@ -5,14 +5,18 @@ sidebar_position: 7 `resources` lets you constrain where a task runs (which nodes) and how many GPUs it should get. -## GPUs: `CUDA_VISIBLE_DEVICES` slicing (Slurm) +## GPUs: `NVIDIA_VISIBLE_DEVICES` slicing (Slurm) GPU resource example: Key idea: - Set `backends..gpus_per_node` so sflow can **pack and slice** GPU indices per task/replica. -- Set `task.resources.gpus.count` to request GPUs for that task. +- Set `task.resources.gpus.count` to request GPUs for that task. SFlow exports + the resulting slice as `NVIDIA_VISIBLE_DEVICES` before launching `srun`, so + NVIDIA container runtimes see the intended GPU subset. SFlow does not set + `CUDA_VISIBLE_DEVICES`; task scripts may set it themselves if a framework + specifically needs it. Minimal example: @@ -37,7 +41,7 @@ backends: nodes: ${{ variables.SLURM_NODES }} workflow: - name: slurm_gpu_cuda_visible + name: slurm_gpu_nvidia_visible tasks: - name: t2 replicas: @@ -47,7 +51,7 @@ workflow: gpus: count: 2 script: - - echo "replica=$SFLOW_REPLICA_INDEX CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" + - echo "replica=$SFLOW_REPLICA_INDEX NVIDIA_VISIBLE_DEVICES=$NVIDIA_VISIBLE_DEVICES" ``` ## Nodes: pin tasks to the same node diff --git a/docs/user/variables.md b/docs/user/variables.md index 2bcfd9e..3fa0b1f 100644 --- a/docs/user/variables.md +++ b/docs/user/variables.md @@ -148,7 +148,7 @@ Available task properties: - `task..nodes[i].ip_address` - IP address of the i-th assigned node - `task..nodes[i].index` - Index of the node within the task's assignment - `task..nodes[i].num_gpus` - Number of GPUs on the node -- `task..gpus` - List of GPU indices assigned to the task (from `CUDA_VISIBLE_DEVICES`) +- `task..gpus` - List of GPU indices assigned to the task (from `NVIDIA_VISIBLE_DEVICES`) - `task..backend` - Name of the backend used by the task - `task..operator` - Name of the operator used by the task diff --git a/src/sflow/app/assembly.py b/src/sflow/app/assembly.py index e60ce5b..b2f3c9f 100644 --- a/src/sflow/app/assembly.py +++ b/src/sflow/app/assembly.py @@ -72,12 +72,12 @@ def _build_task_info( } ) - # Parse GPU indices from CUDA_VISIBLE_DEVICES env var + # Parse GPU indices from NVIDIA_VISIBLE_DEVICES env var. gpus: list[int] = [] - cuda_visible = task.envs.get("CUDA_VISIBLE_DEVICES") - if cuda_visible: + nvidia_visible = task.envs.get("NVIDIA_VISIBLE_DEVICES") + if nvidia_visible: try: - gpus = [int(g.strip()) for g in cuda_visible.split(",") if g.strip()] + gpus = [int(g.strip()) for g in nvidia_visible.split(",") if g.strip()] except ValueError: gpus = [] @@ -1303,10 +1303,10 @@ def _assigned_nodelist( # # This enables global GPU sharing across workflow tasks. Example: # - Node0 has 4 GPUs - # - Task A requests 2 GPUs -> gets CUDA_VISIBLE_DEVICES=0,1 on node0 - # - Task B requests 2 GPUs -> gets CUDA_VISIBLE_DEVICES=2,3 on node0 + # - Task A requests 2 GPUs -> gets NVIDIA_VISIBLE_DEVICES=0,1 on node0 + # - Task B requests 2 GPUs -> gets NVIDIA_VISIBLE_DEVICES=2,3 on node0 # - # We use the same planning-time cursor as CUDA_VISIBLE_DEVICES assignment (`gpu_next`) + # We use the same planning-time cursor as NVIDIA_VISIBLE_DEVICES assignment (`gpu_next`) # so the chosen node reflects already-reserved GPU slices from earlier tasks/replicas. if gpus_count_raw is not None and runtime_backend.allocation is not None: gpus_needed = _resolve_int( @@ -1332,7 +1332,7 @@ def _assigned_nodelist( cursor_key = (runtime_backend.name, node.name) start = gpu_next.get(cursor_key, 0) if start + gpus_needed <= cap: - # Pin to a single node; CUDA_VISIBLE_DEVICES will allocate a slice later. + # Pin to a single node; NVIDIA_VISIBLE_DEVICES will allocate a slice later. return [node.name], True # If nodes are not explicitly requested but GPUs are, we can infer a minimum node count @@ -1364,7 +1364,7 @@ def _assigned_nodelist( # don't place a multi-node task onto nodes whose GPUs were already reserved by # earlier tasks/replicas. # - # Additionally, since CUDA_VISIBLE_DEVICES is computed as a uniform slice across + # Additionally, since NVIDIA_VISIBLE_DEVICES is computed as a uniform slice across # the assigned nodes (per-node env), we require the selected nodes to share the # same planning cursor (gpu_next) value. alloc_nodes_map = {n.name: n for n in alloc_nodes_by_name} @@ -1416,7 +1416,7 @@ def _assigned_nodelist( return [n.name for n in alloc_nodes], False - def _cuda_visible_devices( + def _nvidia_visible_devices( *, task_name: str, base_task_name: str, @@ -1496,13 +1496,13 @@ def _backend_gpu_state_summary() -> str: f"Consider increasing backend nodes or reducing concurrent GPU requests." ) - # CUDA_VISIBLE_DEVICES is evaluated per-node, so indices must be local to that node. + # NVIDIA_VISIBLE_DEVICES is evaluated per-node, so indices must be local to that node. slice_str = ",".join(str(i) for i in range(start, start + count)) gpu_next[cursor_key] = start + count return slice_str # Multi-node GPU request: interpret `count` as a total GPU request across assigned nodes, - # and expose an even-ish per-node slice via CUDA_VISIBLE_DEVICES (since env is per-node). + # and expose an even-ish per-node slice via NVIDIA_VISIBLE_DEVICES (since env is per-node). if runtime_backend.allocation and assigned_nodes and len(assigned_nodes) > 1: alloc_nodes_by_name = {n.name: n for n in runtime_backend.allocation.nodes} caps: list[int] = [] @@ -1703,7 +1703,7 @@ def _make_replica_name( if backend is None: raise ValueError(f"Task '{t_conf.name}' references unknown backend") - # Apply resources -> runtime node subset + CUDA_VISIBLE_DEVICES env + # Apply resources -> runtime node subset + NVIDIA_VISIBLE_DEVICES env nodes_indices_raw = None nodes_count_raw = None nodes_exclude_raw = None @@ -1857,7 +1857,7 @@ def _mount_key(mount: str) -> tuple[str, str] | None: ) if nodes_inferred_from_gpus: nodes_are_pinned = True - cuda_visible = _cuda_visible_devices( + nvidia_visible = _nvidia_visible_devices( task_name=node_name, base_task_name=base, runtime_backend=backend, @@ -1994,8 +1994,8 @@ def _mount_key(mount: str) -> tuple[str, str] | None: if apath is not None: task.envs.setdefault(aname, str(apath)) task.envs.update(replica_envs.get(node_name, {})) - if cuda_visible is not None: - task.envs["CUDA_VISIBLE_DEVICES"] = cuda_visible + if nvidia_visible is not None: + task.envs["NVIDIA_VISIBLE_DEVICES"] = nvidia_visible task_graph.dag.add_node(node_name, task) # If the task is replicated sequentially, enforce replica order by chaining edges. diff --git a/src/sflow/app/sflow.py b/src/sflow/app/sflow.py index ae42ccc..7ea4c4f 100644 --- a/src/sflow/app/sflow.py +++ b/src/sflow/app/sflow.py @@ -44,18 +44,18 @@ def extract_container_mounts_from_extra_args(extra_args: list[str]) -> list[str] return mounts -def parse_cuda_visible_devices(cuda_visible: str | None) -> list[int]: +def parse_gpu_visible_devices(visible_devices: str | None) -> list[int]: """ - Parse CUDA_VISIBLE_DEVICES into a list of GPU indices. + Parse a comma-separated GPU visibility string into a list of GPU indices. Supports comma-separated indices and simple ranges like ``0-3``. Non-numeric tokens are ignored. """ - if not cuda_visible: + if not visible_devices: return [] indices: list[int] = [] - for part in str(cuda_visible).split(","): + for part in str(visible_devices).split(","): token = part.strip() if not token: continue @@ -132,8 +132,8 @@ def _shorten(value: str, max_len: int = 18) -> str: if not assigned_nodes and alloc.nodes: assigned_nodes = [node.name for node in alloc.nodes] - gpu_indices = parse_cuda_visible_devices( - getattr(task, "envs", {}).get("CUDA_VISIBLE_DEVICES") + gpu_indices = parse_gpu_visible_devices( + getattr(task, "envs", {}).get("NVIDIA_VISIBLE_DEVICES") ) for node_name in assigned_nodes: @@ -889,7 +889,7 @@ def _validate_container_mounts( op_type_str = getattr(op_conf, "type", None) or "unknown" nodelist = getattr(op_conf, "nodelist", None) or [] - cuda_visible = t.envs.get("CUDA_VISIBLE_DEVICES") + nvidia_visible = t.envs.get("NVIDIA_VISIBLE_DEVICES") task_out_dir = t.envs.get("SFLOW_TASK_OUTPUT_DIR") retry = getattr(t, "retries", None) retry_str = ( @@ -908,9 +908,9 @@ def _validate_container_mounts( f" ├─ depends_on: {list(deps) if deps else '[]'}" ) _logger.info(f" ├─ nodelist: {nodelist}") - if cuda_visible: + if nvidia_visible: _logger.info( - f" ├─ CUDA_VISIBLE_DEVICES: {cuda_visible}" + f" ├─ NVIDIA_VISIBLE_DEVICES: {nvidia_visible}" ) _logger.info(f" ├─ task_output_dir: {task_out_dir}") _logger.info(f" ├─ retries: {retry_str}") @@ -1234,12 +1234,12 @@ class VisualizeResult: ) def _gpu_label(task) -> str: - cuda = task.envs.get("CUDA_VISIBLE_DEVICES") - if not cuda: + nvidia_visible = task.envs.get("NVIDIA_VISIBLE_DEVICES") + if not nvidia_visible: return "" # "0,1,2" -> 3 try: - n = len([x for x in cuda.split(",") if x.strip() != ""]) + n = len([x for x in nvidia_visible.split(",") if x.strip() != ""]) except Exception: n = 0 return f" gpus={n}" diff --git a/src/sflow/core/compute_node.py b/src/sflow/core/compute_node.py index 1848ae3..3c385ef 100644 --- a/src/sflow/core/compute_node.py +++ b/src/sflow/core/compute_node.py @@ -10,7 +10,7 @@ class ComputeNode: name: str ip_address: str index: int - # GPU count available on this node (if known). Used for CUDA_VISIBLE_DEVICES packing/validation. + # GPU count available on this node (if known). Used for NVIDIA_VISIBLE_DEVICES packing/validation. num_gpus: int | None = None def to_dict(self) -> dict[str, Any]: diff --git a/src/sflow/skills/writing-sflow-yaml/examples.md b/src/sflow/skills/writing-sflow-yaml/examples.md index 4cac592..45276a5 100644 --- a/src/sflow/skills/writing-sflow-yaml/examples.md +++ b/src/sflow/skills/writing-sflow-yaml/examples.md @@ -281,7 +281,7 @@ workflow: policy: "parallel" script: - set -x - - export FIRST_CUDA_DEVICE=$(echo ${CUDA_VISIBLE_DEVICES} | cut -d',' -f1) + - export FIRST_CUDA_DEVICE=$(echo ${NVIDIA_VISIBLE_DEVICES} | cut -d',' -f1) - export VLLM_NIXL_SIDE_CHANNEL_PORT=$((5557 + ${FIRST_CUDA_DEVICE})) - export DYN_SYSTEM_PORT=$((8082 + ${FIRST_CUDA_DEVICE})) - > @@ -325,7 +325,7 @@ workflow: policy: "parallel" script: - set -x - - export FIRST_CUDA_DEVICE=$(echo ${CUDA_VISIBLE_DEVICES} | cut -d',' -f1) + - export FIRST_CUDA_DEVICE=$(echo ${NVIDIA_VISIBLE_DEVICES} | cut -d',' -f1) - export VLLM_NIXL_SIDE_CHANNEL_PORT=$((5557 + ${FIRST_CUDA_DEVICE})) - export DYN_SYSTEM_PORT=$((8082 + ${FIRST_CUDA_DEVICE})) - > diff --git a/src/sflow/skills/writing-sflow-yaml/schema-reference.md b/src/sflow/skills/writing-sflow-yaml/schema-reference.md index d74b644..7fd4633 100644 --- a/src/sflow/skills/writing-sflow-yaml/schema-reference.md +++ b/src/sflow/skills/writing-sflow-yaml/schema-reference.md @@ -326,14 +326,14 @@ workflow: ```yaml resources: gpus: - count: 4 # CUDA_VISIBLE_DEVICES slicing + count: 4 # NVIDIA_VISIBLE_DEVICES slicing nodes: indices: [0] # pin to specific node indices count: 2 # OR request N nodes (mutually exclusive with indices) ``` GPU allocation: -- Each task/replica gets `count` GPUs via `CUDA_VISIBLE_DEVICES` +- Each task/replica gets `count` GPUs via `NVIDIA_VISIBLE_DEVICES` - GPUs are sliced sequentially across replicas on each node - Total GPUs across all tasks must not exceed `nodes * gpus_per_node` diff --git a/tests/e2e_tests/test_task_context.yaml b/tests/e2e_tests/test_task_context.yaml index 859b98a..07248db 100644 --- a/tests/e2e_tests/test_task_context.yaml +++ b/tests/e2e_tests/test_task_context.yaml @@ -64,7 +64,7 @@ workflow: gpus: count: ${{ variables.PREFILL_GPUS_PER_SERVER }} script: - - echo "Starting prefill server ${SFLOW_REPLICA_INDEX} on nodes ${SFLOW_TASK_ASSIGNED_NODE_NAMES} with GPU ${CUDA_VISIBLE_DEVICES}" + - echo "Starting prefill server ${SFLOW_REPLICA_INDEX} on nodes ${SFLOW_TASK_ASSIGNED_NODE_NAMES} with GPU ${NVIDIA_VISIBLE_DEVICES}" - echo "Assigned IPs - ${SFLOW_TASK_ASSIGNED_NODE_IPS}" - name: decode_server @@ -76,7 +76,7 @@ workflow: count: ${{ variables.DECODE_GPUS_PER_SERVER }} script: # Use SFLOW_TASK_ASSIGNED_* env vars to access this task's own assigned nodes - - echo "Starting decode server ${SFLOW_REPLICA_INDEX} on nodes ${SFLOW_TASK_ASSIGNED_NODE_NAMES} with GPU ${CUDA_VISIBLE_DEVICES}" + - echo "Starting decode server ${SFLOW_REPLICA_INDEX} on nodes ${SFLOW_TASK_ASSIGNED_NODE_NAMES} with GPU ${NVIDIA_VISIBLE_DEVICES}" - echo "Assigned IPs - ${SFLOW_TASK_ASSIGNED_NODE_IPS}" - name: client @@ -135,4 +135,4 @@ workflow: - echo "Client Task Operator - ${{ task.client.operator }}" depends_on: - prefill_server - - decode_server \ No newline at end of file + - decode_server diff --git a/tests/integration/guide/sflow_slurm_gpu_cuda_visible.yaml b/tests/integration/guide/sflow_slurm_gpu_cuda_visible.yaml index 177a526..a7b6954 100644 --- a/tests/integration/guide/sflow_slurm_gpu_cuda_visible.yaml +++ b/tests/integration/guide/sflow_slurm_gpu_cuda_visible.yaml @@ -17,7 +17,7 @@ backends: - name: slurm_cluster type: slurm default: true - # Used by sflow to pack/slice CUDA_VISIBLE_DEVICES + # Used by sflow to pack/slice NVIDIA_VISIBLE_DEVICES gpus_per_node: ${{ variables.GPUS_PER_NODE }} account: ${{ variables.SLURM_ACCOUNT }} partition: ${{ variables.SLURM_PARTITION }} @@ -31,13 +31,13 @@ operators: container_mount_home: false workflow: - name: slurm_gpu_cuda_visible + name: slurm_gpu_nvidia_visible tasks: # Expectation with SLURM_NODES=2 and GPUS_PER_NODE=4: # - sflow will pin base tasks to distinct nodes by default (t2 -> node0, t4 -> node1). - # - t2_0: CUDA_VISIBLE_DEVICES=0,1 - # - t2_1: CUDA_VISIBLE_DEVICES=2,3 + # - t2_0: NVIDIA_VISIBLE_DEVICES=0,1 + # - t2_1: NVIDIA_VISIBLE_DEVICES=2,3 - name: t2 replicas: count: 2 @@ -46,7 +46,7 @@ workflow: gpus: count: 2 script: - - echo "task=t2 replica=$SFLOW_REPLICA_INDEX CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" + - echo "task=t2 replica=$SFLOW_REPLICA_INDEX NVIDIA_VISIBLE_DEVICES=$NVIDIA_VISIBLE_DEVICES" - nvidia-smi -L || true # Expectation with GPUS_PER_NODE=4: @@ -59,5 +59,5 @@ workflow: gpus: count: 1 script: - - echo "task=t4 replica=$SFLOW_REPLICA_INDEX CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" + - echo "task=t4 replica=$SFLOW_REPLICA_INDEX NVIDIA_VISIBLE_DEVICES=$NVIDIA_VISIBLE_DEVICES" - nvidia-smi -L || true diff --git a/tests/unit/test_app_assembly_build_task_graph.py b/tests/unit/test_app_assembly_build_task_graph.py index 8f819d5..fe84acc 100644 --- a/tests/unit/test_app_assembly_build_task_graph.py +++ b/tests/unit/test_app_assembly_build_task_graph.py @@ -606,7 +606,7 @@ def test_build_task_graph_resources_nodes_count_compact_allocation_for_parallel_ assert t11.operator.config.nodelist == ["n3", "n4"] -def test_build_task_graph_resources_gpus_count_sets_cuda_visible_devices_with_offset(): +def test_build_task_graph_resources_gpus_count_sets_nvidia_visible_devices_with_offset(): state = _state() state.backends = {"local": _FakeBackend("local", allocation=None)} state.default_backend = state.backends["local"] @@ -627,8 +627,12 @@ def test_build_task_graph_resources_gpus_count_sets_cuda_visible_devices_with_of ) tg = build_task_graph(config, state) - assert tg.get_task("t1_0").envs["CUDA_VISIBLE_DEVICES"] == "0,1" - assert tg.get_task("t1_1").envs["CUDA_VISIBLE_DEVICES"] == "2,3" + t10 = tg.get_task("t1_0") + t11 = tg.get_task("t1_1") + assert t10.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1" + assert t11.envs["NVIDIA_VISIBLE_DEVICES"] == "2,3" + assert "CUDA_VISIBLE_DEVICES" not in t10.envs + assert "CUDA_VISIBLE_DEVICES" not in t11.envs def test_build_task_graph_resources_gpus_respects_compute_node_num_gpus(): @@ -661,8 +665,12 @@ def test_build_task_graph_resources_gpus_respects_compute_node_num_gpus(): ), ) tg = build_task_graph(ok, state) - assert tg.get_task("t1_0").envs["CUDA_VISIBLE_DEVICES"] == "0,1" - assert tg.get_task("t1_1").envs["CUDA_VISIBLE_DEVICES"] == "2,3" + t10 = tg.get_task("t1_0") + t11 = tg.get_task("t1_1") + assert t10.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1" + assert t11.envs["NVIDIA_VISIBLE_DEVICES"] == "2,3" + assert "CUDA_VISIBLE_DEVICES" not in t10.envs + assert "CUDA_VISIBLE_DEVICES" not in t11.envs bad = SflowConfig( version="0.1", @@ -761,7 +769,7 @@ def test_build_task_graph_resources_gpus_count_auto_expands_nodes_when_total_exc assert t1.operator.config.nodelist == ["n1", "n2"] assert t1.operator.config.nodes == 2 # Multi-node request exposes a per-node slice (env is evaluated per node). - assert t1.envs["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + assert t1.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1,2,3" def test_build_task_graph_resources_gpus_count_requires_multiple_of_per_node_capacity_for_auto_expand(): @@ -856,12 +864,12 @@ def test_build_task_graph_resources_with_multiple_tasks_pins_distinct_nodes_and_ assert t22.operator.config.nodelist == ["n2"] assert t23.operator.config.nodelist == ["n2"] - assert t10.envs["CUDA_VISIBLE_DEVICES"] == "0,1" - assert t11.envs["CUDA_VISIBLE_DEVICES"] == "2,3" - assert t20.envs["CUDA_VISIBLE_DEVICES"] == "0" - assert t21.envs["CUDA_VISIBLE_DEVICES"] == "1" - assert t22.envs["CUDA_VISIBLE_DEVICES"] == "2" - assert t23.envs["CUDA_VISIBLE_DEVICES"] == "3" + assert t10.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1" + assert t11.envs["NVIDIA_VISIBLE_DEVICES"] == "2,3" + assert t20.envs["NVIDIA_VISIBLE_DEVICES"] == "0" + assert t21.envs["NVIDIA_VISIBLE_DEVICES"] == "1" + assert t22.envs["NVIDIA_VISIBLE_DEVICES"] == "2" + assert t23.envs["NVIDIA_VISIBLE_DEVICES"] == "3" def test_build_task_graph_gpu_packing_allows_multiple_base_tasks_to_share_remaining_gpus_on_same_node(): @@ -918,17 +926,17 @@ def test_build_task_graph_gpu_packing_allows_multiple_base_tasks_to_share_remain # t1 consumes 2 GPUs on n1 -> leaves 2 GPUs (2,3) available on n1. assert t1.operator.config.nodelist == ["n1"] - assert t1.envs["CUDA_VISIBLE_DEVICES"] == "0,1" + assert t1.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1" # t2 replicas should use the remaining GPUs on n1. assert t20.operator.config.nodelist == ["n1"] assert t21.operator.config.nodelist == ["n1"] - assert t20.envs["CUDA_VISIBLE_DEVICES"] == "2" - assert t21.envs["CUDA_VISIBLE_DEVICES"] == "3" + assert t20.envs["NVIDIA_VISIBLE_DEVICES"] == "2" + assert t21.envs["NVIDIA_VISIBLE_DEVICES"] == "3" # t3 requests a full node (4 GPUs) -> should land on n2. assert t3.operator.config.nodelist == ["n2"] - assert t3.envs["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + assert t3.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1,2,3" def test_build_task_graph_multi_node_gpu_request_reserves_gpus_so_later_tasks_use_other_nodes(): @@ -981,12 +989,12 @@ def test_build_task_graph_multi_node_gpu_request_reserves_gpus_so_later_tasks_us d1 = tg.get_task("decode_1") assert prefill.operator.config.nodelist == ["n1", "n2"] - assert prefill.envs["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + assert prefill.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1,2,3" assert d0.operator.config.nodelist == ["n3"] assert d1.operator.config.nodelist == ["n3"] - assert d0.envs["CUDA_VISIBLE_DEVICES"] == "0,1" - assert d1.envs["CUDA_VISIBLE_DEVICES"] == "2,3" + assert d0.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1" + assert d1.envs["NVIDIA_VISIBLE_DEVICES"] == "2,3" def test_build_task_graph_multi_node_decode_avoids_node_fully_consumed_by_sequential_prefill_replicas(): @@ -1039,15 +1047,15 @@ def test_build_task_graph_multi_node_decode_avoids_node_fully_consumed_by_sequen assert tg.get_task("prefill_1").operator.config.nodelist == ["n1"] assert tg.get_task("prefill_2").operator.config.nodelist == ["n1"] assert tg.get_task("prefill_3").operator.config.nodelist == ["n1"] - assert tg.get_task("prefill_0").envs["CUDA_VISIBLE_DEVICES"] == "0" - assert tg.get_task("prefill_1").envs["CUDA_VISIBLE_DEVICES"] == "1" - assert tg.get_task("prefill_2").envs["CUDA_VISIBLE_DEVICES"] == "2" - assert tg.get_task("prefill_3").envs["CUDA_VISIBLE_DEVICES"] == "3" + assert tg.get_task("prefill_0").envs["NVIDIA_VISIBLE_DEVICES"] == "0" + assert tg.get_task("prefill_1").envs["NVIDIA_VISIBLE_DEVICES"] == "1" + assert tg.get_task("prefill_2").envs["NVIDIA_VISIBLE_DEVICES"] == "2" + assert tg.get_task("prefill_3").envs["NVIDIA_VISIBLE_DEVICES"] == "3" # Decode should avoid n1 and use n2+n3. decode = tg.get_task("decode") assert decode.operator.config.nodelist == ["n2", "n3"] - assert decode.envs["CUDA_VISIBLE_DEVICES"] == "0,1,2,3" + assert decode.envs["NVIDIA_VISIBLE_DEVICES"] == "0,1,2,3" def test_build_task_graph_replica_tasks_have_independent_container_mounts(tmp_path): @@ -1587,7 +1595,7 @@ def test_build_task_graph_resources_nodes_exclude_with_gpus(): t1 = tg.get_task("t1") # Pool after exclude: only n2. GPU packing should place task on n2. assert t1.operator.config.nodelist == ["n2"] - assert t1.envs.get("CUDA_VISIBLE_DEVICES") == "0,1" + assert t1.envs.get("NVIDIA_VISIBLE_DEVICES") == "0,1" def test_build_task_graph_resources_nodes_exclude_all_raises(): diff --git a/tests/unit/test_app_sflow_dry_run.py b/tests/unit/test_app_sflow_dry_run.py index a0b6376..45dbeaf 100644 --- a/tests/unit/test_app_sflow_dry_run.py +++ b/tests/unit/test_app_sflow_dry_run.py @@ -12,7 +12,7 @@ from sflow.app.sflow import ( build_allocation_map_lines, extract_container_mounts_from_extra_args, - parse_cuda_visible_devices, + parse_gpu_visible_devices, ) from sflow.core.backend import Allocation from sflow.core.compute_node import ComputeNode @@ -94,18 +94,18 @@ def test_container_mounts_with_rw_mode(self): assert result == ["/host:/container:rw", "/data:/data:ro"] -class TestParseCudaVisibleDevices: +class TestParseGpuVisibleDevices: def test_none_returns_empty(self): - assert parse_cuda_visible_devices(None) == [] + assert parse_gpu_visible_devices(None) == [] def test_comma_separated_indices(self): - assert parse_cuda_visible_devices("0,2,3") == [0, 2, 3] + assert parse_gpu_visible_devices("0,2,3") == [0, 2, 3] def test_range_syntax(self): - assert parse_cuda_visible_devices("0-3") == [0, 1, 2, 3] + assert parse_gpu_visible_devices("0-3") == [0, 1, 2, 3] def test_mixed_tokens_ignores_invalid(self): - assert parse_cuda_visible_devices("0,abc,2-3") == [0, 2, 3] + assert parse_gpu_visible_devices("0,abc,2-3") == [0, 2, 3] class TestBuildAllocationMapLines: @@ -124,14 +124,14 @@ def test_builds_node_and_gpu_chart(self): name="prefill_0", backend_name="slurm_cluster", assigned_nodes=["n1"], - envs={"CUDA_VISIBLE_DEVICES": "0,1"}, + envs={"NVIDIA_VISIBLE_DEVICES": "0,1"}, operator=None, ), SimpleNamespace( name="decode_0", backend_name="slurm_cluster", assigned_nodes=["n1"], - envs={"CUDA_VISIBLE_DEVICES": "2"}, + envs={"NVIDIA_VISIBLE_DEVICES": "2"}, operator=None, ), SimpleNamespace( @@ -145,7 +145,7 @@ def test_builds_node_and_gpu_chart(self): name="decode_1", backend_name="slurm_cluster", assigned_nodes=["n2"], - envs={"CUDA_VISIBLE_DEVICES": "0,1,2,3"}, + envs={"NVIDIA_VISIBLE_DEVICES": "0,1,2,3"}, operator=None, ), ]