Mirror GPU resource slices to NVIDIA_VISIBLE_DEVICES

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) |

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.
 

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
 

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.<name>.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

docs/user/variables.md

@@ -148,7 +148,7 @@ Available task properties:
 - `task.<name>.nodes[i].ip_address` - IP address of the i-th assigned node
 - `task.<name>.nodes[i].index` - Index of the node within the task's assignment
 - `task.<name>.nodes[i].num_gpus` - Number of GPUs on the node
-- `task.<name>.gpus` - List of GPU indices assigned to the task (from `CUDA_VISIBLE_DEVICES`)
+- `task.<name>.gpus` - List of GPU indices assigned to the task (from `NVIDIA_VISIBLE_DEVICES`)
 - `task.<name>.backend` - Name of the backend used by the task
 - `task.<name>.operator` - Name of the operator used by the task
 

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.

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}"

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]:

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}))
         - >

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`
 

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
+        - decode_server

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