Split automatically detected GPUs into gpus/nvidia and gpus/amd

Author: Kobzol

CHANGELOG.md

@@ -12,6 +12,11 @@
   $ hq worker start --resource="res1=[foo, bar]"
   ```
 
+* HyperQueue now provides built-in support for AMD GPUs. For this reason, the default name of GPU
+resources that are automatically detected on a worker has been changed from `gpus` to `gpus/nvidia`
+for NVIDIA GPUs. AMD GPUs are now autodetected as `gpus/amd`. In the future, we intend to create a way
+to ask for any GPU resource (e.g. `--resource=gpus=2`), regardless of its type.
+
 * AMD GPUs are now automatically detected in workers from the environment variable `ROCR_VISIBLE_DEVICES`
   or `HIP_VISIBLE_DEVICES`.
 

crates/hyperqueue/src/worker/hwdetect.rs

@@ -12,7 +12,7 @@ use tako::format_comma_delimited;
 use tako::internal::has_unique_elements;
 use tako::resources::{
     ResourceDescriptorItem, ResourceDescriptorKind, ResourceIndex, ResourceLabel,
-    GPU_RESOURCE_NAME, MEM_RESOURCE_NAME,
+    AMD_GPU_RESOURCE_NAME, MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME,
 };
 
 pub fn detect_cpus() -> anyhow::Result<ResourceDescriptorKind> {
@@ -57,17 +57,17 @@ pub fn detect_additional_resources(items: &mut Vec<ResourceDescriptorItem>) -> a
     let has_resource =
         |items: &[ResourceDescriptorItem], name: &str| items.iter().any(|x| x.name == name);
 
-    if !has_resource(items, GPU_RESOURCE_NAME) {
-        if let Some(gpus) = detect_gpus_from_env() {
+    if !has_resource(items, NVIDIA_GPU_RESOURCE_NAME) {
+        if let Some(detected) = detect_gpus_from_env() {
             items.push(ResourceDescriptorItem {
-                name: GPU_RESOURCE_NAME.to_string(),
-                kind: gpus,
+                name: detected.resource_name.to_string(),
+                kind: detected.resource,
             });
-        } else if let Ok(count) = read_linux_gpu_count() {
+        } else if let Ok(count) = read_nvidia_linux_gpu_count() {
             if count > 0 {
                 log::info!("Detected {} GPUs from procs", count);
                 items.push(ResourceDescriptorItem {
-                    name: GPU_RESOURCE_NAME.to_string(),
+                    name: NVIDIA_GPU_RESOURCE_NAME.to_string(),
                     kind: ResourceDescriptorKind::simple_indices(count as u32),
                 });
             }
@@ -86,15 +86,20 @@ pub fn detect_additional_resources(items: &mut Vec<ResourceDescriptorItem>) -> a
     Ok(())
 }
 
-pub const GPU_ENV_KEYS: &[&str; 3] = &[
-    "CUDA_VISIBLE_DEVICES",
-    "HIP_VISIBLE_DEVICES",
-    "ROCR_VISIBLE_DEVICES",
+pub const GPU_ENV_KEYS: &[(&str, &str); 3] = &[
+    ("CUDA_VISIBLE_DEVICES", NVIDIA_GPU_RESOURCE_NAME),
+    ("HIP_VISIBLE_DEVICES", AMD_GPU_RESOURCE_NAME),
+    ("ROCR_VISIBLE_DEVICES", AMD_GPU_RESOURCE_NAME),
 ];
 
+struct DetectedGpu {
+    resource_name: &'static str,
+    resource: ResourceDescriptorKind,
+}
+
 /// Tries to detect available GPUs from one of the `GPU_ENV_KEYS` environment variables.
-fn detect_gpus_from_env() -> Option<ResourceDescriptorKind> {
-    GPU_ENV_KEYS.iter().find_map(|env_key| {
+fn detect_gpus_from_env() -> Option<DetectedGpu> {
+    for (env_key, resource_name) in GPU_ENV_KEYS {
         if let Ok(devices_str) = std::env::var(env_key) {
             if let Ok(devices) = parse_comma_separated_values(&devices_str) {
                 log::info!(
@@ -108,15 +113,18 @@ fn detect_gpus_from_env() -> Option<ResourceDescriptorKind> {
 
                 let list =
                     ResourceDescriptorKind::list(devices).expect("List values were not unique");
-                return Some(list);
+                return Some(DetectedGpu {
+                    resource_name,
+                    resource: list,
+                });
             }
         }
-        None
-    })
+    }
+    None
 }
 
 /// Try to find out how many Nvidia GPUs are available on the current node.
-fn read_linux_gpu_count() -> anyhow::Result<usize> {
+fn read_nvidia_linux_gpu_count() -> anyhow::Result<usize> {
     Ok(std::fs::read_dir("/proc/driver/nvidia/gpus")?.count())
 }
 

crates/hyperqueue/src/worker/start.rs

@@ -31,15 +31,15 @@ use crate::common::placeholders::{
 use crate::common::utils::fs::{bytes_to_path, is_implicit_path, path_has_extension};
 use crate::transfer::messages::{PinMode, TaskBody};
 use crate::transfer::stream::ChannelId;
-use crate::worker::hwdetect::GPU_ENV_KEYS;
 use crate::worker::streamer::StreamSender;
 use crate::worker::streamer::StreamerRef;
 use crate::{JobId, JobTaskId};
 use serde::{Deserialize, Serialize};
 use tako::comm::serialize;
 use tako::program::{ProgramDefinition, StdioDef};
 use tako::resources::{
-    Allocation, ResourceAllocation, CPU_RESOURCE_ID, CPU_RESOURCE_NAME, GPU_RESOURCE_NAME,
+    Allocation, ResourceAllocation, AMD_GPU_RESOURCE_NAME, CPU_RESOURCE_ID, CPU_RESOURCE_NAME,
+    NVIDIA_GPU_RESOURCE_NAME,
 };
 
 const MAX_CUSTOM_ERROR_LENGTH: usize = 2048; // 2KiB
@@ -244,16 +244,21 @@ fn insert_resources_into_env(ctx: &LaunchContext, program: &mut ProgramDefinitio
                     );
                 }
             }
-            if resource_name == GPU_RESOURCE_NAME {
-                /* Extra variables for GPUS */
-                for &key in GPU_ENV_KEYS {
-                    program.env.insert(key.into(), labels.clone().into());
-                }
+            if resource_name == NVIDIA_GPU_RESOURCE_NAME {
+                /* Extra variables for Nvidia GPUS */
+                program
+                    .env
+                    .insert("CUDA_VISIBLE_DEVICES".into(), labels.clone().into());
                 program
                     .env
                     .insert("CUDA_DEVICE_ORDER".into(), "PCI_BUS_ID".into());
             }
-
+            if resource_name == AMD_GPU_RESOURCE_NAME {
+                /* Extra variable for AMD GPUS */
+                program
+                    .env
+                    .insert("ROCR_VISIBLE_DEVICES".into(), labels.clone().into());
+            }
             program.env.insert(
                 format!("HQ_RESOURCE_VALUES_{resource_name}").into(),
                 labels.into(),

crates/tako/benches/benchmarks/worker.rs

@@ -13,7 +13,7 @@ use tako::resources::{
     AllocationRequest, ResourceDescriptor, ResourceRequest, ResourceRequestEntry, TimeRequest,
 };
 use tako::resources::{
-    ResourceDescriptorItem, ResourceDescriptorKind, CPU_RESOURCE_NAME, GPU_RESOURCE_NAME,
+    ResourceDescriptorItem, ResourceDescriptorKind, CPU_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME,
 };
 use tako::ItemId;
 use tokio::sync::mpsc::unbounded_channel;
@@ -170,7 +170,7 @@ fn create_resource_queue(num_cpus: u32) -> ResourceWaitQueue {
             kind: ResourceDescriptorKind::simple_indices(num_cpus),
         },
         ResourceDescriptorItem {
-            name: GPU_RESOURCE_NAME.to_string(),
+            name: NVIDIA_GPU_RESOURCE_NAME.to_string(),
             kind: ResourceDescriptorKind::simple_indices(8),
         },
     ]);

crates/tako/src/internal/common/resources/map.rs

@@ -4,7 +4,8 @@ use crate::internal::common::Map;
 pub const CPU_RESOURCE_ID: ResourceId = ResourceId(0);
 
 pub const CPU_RESOURCE_NAME: &str = "cpus";
-pub const GPU_RESOURCE_NAME: &str = "gpus";
+pub const NVIDIA_GPU_RESOURCE_NAME: &str = "gpus/nvidia";
+pub const AMD_GPU_RESOURCE_NAME: &str = "gpus/amd";
 pub const MEM_RESOURCE_NAME: &str = "mem";
 
 #[derive(Debug)]

crates/tako/src/internal/common/resources/mod.rs

@@ -9,7 +9,10 @@ pub use allocation::{Allocation, AllocationValue, ResourceAllocation, ResourceAl
 pub use descriptor::{
     DescriptorError, ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind,
 };
-pub use map::{CPU_RESOURCE_ID, CPU_RESOURCE_NAME, GPU_RESOURCE_NAME, MEM_RESOURCE_NAME};
+pub use map::{
+    AMD_GPU_RESOURCE_NAME, CPU_RESOURCE_ID, CPU_RESOURCE_NAME, MEM_RESOURCE_NAME,
+    NVIDIA_GPU_RESOURCE_NAME,
+};
 pub use request::{
     AllocationRequest, ResourceRequest, ResourceRequestEntries, ResourceRequestEntry, TimeRequest,
 };

crates/tako/src/lib.rs

@@ -31,8 +31,8 @@ pub mod resources {
         Allocation, AllocationRequest, AllocationValue, NumOfNodes, ResourceAllocation,
         ResourceAmount, ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind,
         ResourceIndex, ResourceLabel, ResourceRequest, ResourceRequestEntries,
-        ResourceRequestEntry, TimeRequest, CPU_RESOURCE_ID, CPU_RESOURCE_NAME, GPU_RESOURCE_NAME,
-        MEM_RESOURCE_NAME,
+        ResourceRequestEntry, TimeRequest, AMD_GPU_RESOURCE_NAME, CPU_RESOURCE_ID,
+        CPU_RESOURCE_NAME, MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME,
     };
 
     pub use crate::internal::common::resources::map::ResourceMap;

docs/deployment/allocation.md

@@ -74,13 +74,13 @@ creating a new allocation queue:
     === "PBS"
 
         ```bash
-        $ hq alloc add pbs --time-limit 1h --cpus 4x4 --resource "gpus=range(1-2)" -- -qqprod -AAccount1
+        $ hq alloc add pbs --time-limit 1h --cpus 4x4 --resource "gpus/nvidia=range(1-2)" -- -qqprod -AAccount1
         ```
 
     === "Slurm"
 
         ``` bash
-        $ hq alloc add slurm --time-limit 1h --cpus 4x4 --resource "gpus=range(1-2)" -- --partition=p1
+        $ hq alloc add slurm --time-limit 1h --cpus 4x4 --resource "gpus/nvidia=range(1-2)" -- --partition=p1
         ```
 
     If you do not pass any resources, they will be detected automatically (same as it works with

docs/jobs/resources.md

@@ -90,13 +90,19 @@ The following resources are detected automatically if a resource of a given name
 
 * CPUs are automatically detected as resource named "cpus" (more in [CPU resources](cresources.md)).
 
-* GPUs that are available when a worker is started are automatically detected under the resource
-  name `gpus`. You can use the environment variables `CUDA_VISIBLE_DEVICES`, `HIP_VISIBLE_DEVICES` or
-  `ROCR_VISIBLE_DEVICES` when starting a worker to override the list of available GPUs:
+* GPUs that are available when a worker is started are automatically detected under the following
+  resource names:
+    - NVIDIA GPUs are stored the under resource name `gpus/nvidia`. These GPUs are detected from the
+      environment variable `CUDA_VISIBLE_DEVICES` or from the `procfs` filesystem.
+    - AMD GPUs are stored under the resource name `gpus/amd`. These GPUs are detected from the environment
+      variables `ROCR_VISIBLE_DEVICES` or `HIP_VISIBLE_DEVICES`.
+
+    You can set these environment variables when starting a worker to override the list of available GPUs:
 
-```bash
-$ CUDA_VISIBLE_DEVICES=2,3 hq worker start
-```
+    ```bash
+    $ CUDA_VISIBLE_DEVICES=2,3 hq worker start
+    # The worker will have resource gpus/nvidia=[2,3]
+    ```
 
 * RAM of the node is detected as resource "mem" in bytes. 
 
@@ -111,7 +117,6 @@ The automatic detection of resources can be disabled by argument ``--no-detect-r
 It disables detection of resources other than "cpus";
 if resource "cpus" are not explicitly defined, it will always be detected.
 
-
 ## Resource request
 
 When you submit a job, you can define a **resource requests** with the `--resource` flag:
@@ -133,14 +138,14 @@ task requests.
 
     For example, let's say that a worker has an indexed pool of GPUs:
     ```bash
-    $ hq worker start --resource "gpus=range(1-3)"
+    $ hq worker start --resource "gpus/nvidia=range(1-3)"
     ```
     And we create two jobs, each with a single task. The first job wants 1 GPU, the second one wants
     two GPUs.
 
     ```bash
-    $ hq submit --resource gpus=1 ...
-    $ hq submit --resource gpus=2 ...
+    $ hq submit --resource gpus/nvidia=1 ...
+    $ hq submit --resource gpus/nvidia=2 ...
     ```
 
     Then the first job can be allocated e.g. the GPU `2` and the second job can be allocated the GPUs
@@ -191,16 +196,13 @@ each resource request named `<NAME>`:
 * `HQ_RESOURCE_VALUES_<NAME>` contains the specific resource values allocated for the task as a
 comma-separated list. This variable is only filled for indexed resource pool.
 
-!!! tip
-
-    HQ has a special case for a resource named `gpus`. For that resource, it will also pass the following
-    environment variables to the spawned task:
-
-    * `CUDA_DEVICE_ORDER` set to the value `PCI_BUS_ID`
-    * `CUDA_VISIBLE_DEVICES` set to the same value as `HQ_RESOURCE_VALUES_gpus`
-    * `HIP_VISIBLE_DEVICES` set to the same value as `HQ_RESOURCE_VALUES_gpus`
-    * `ROCR_VISIBLE_DEVICES` set to the same value as `HQ_RESOURCE_VALUES_gpus`
+HQ also sets additional environment variables for various resources with special names:
 
+- For the resource `gpus/nvidia`, HQ will set:
+    - `CUDA_VISIBLE_DEVICES` to the same value as `HQ_RESOURCE_VALUES_gpus/nvidia`
+    - `CUDA_DEVICE_ORDER` to `PCI_BUS_ID`
+- For the resource `gpus/amd`, HQ will set:
+    - `ROCR_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` to the same value as `HQ_RESOURCE_VALUES_gpus/amd`
 
 ## Resource requests and job arrays
 

tests/test_resources.py

@@ -171,44 +171,48 @@ def test_worker_resource_hwdetect_mem(hq_env: HqEnv):
 
 def test_worker_set_gpu_env_for_task(hq_env: HqEnv):
     hq_env.start_server()
-    hq_env.start_worker(args=["--resource", "gpus=[0,1]"])
+    hq_env.start_worker(args=["--resource", "gpus/nvidia=[0,1]"])
     hq_env.command(
         [
             "submit",
             "--resource",
-            "gpus=2",
+            "gpus/nvidia=2",
             "--",
             "bash",
             "-c",
             """
 echo $CUDA_VISIBLE_DEVICES
-echo $HIP_VISIBLE_DEVICES
-echo $ROCR_VISIBLE_DEVICES
 """,
         ]
     )
     wait_for_job_state(hq_env, 1, "FINISHED")
     assert list(
         set(int(v) for v in line.split(","))
         for line in read_file(default_task_output()).splitlines()
-    ) == [{0, 1}, {0, 1}, {0, 1}]
+    ) == [{0, 1}]
 
 
 @pytest.mark.parametrize(
-    "env_key", ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES")
+    "env_and_res",
+    (
+        ("CUDA_VISIBLE_DEVICES", "gpus/nvidia"),
+        ("HIP_VISIBLE_DEVICES", "gpus/amd"),
+        ("ROCR_VISIBLE_DEVICES", "gpus/amd"),
+    ),
 )
-def test_worker_detect_gpus_from_env(hq_env: HqEnv, env_key: str):
+def test_worker_detect_gpus_from_env(hq_env: HqEnv, env_and_res: str):
+    env, resource = env_and_res
     hq_env.start_server()
-    resources = hq_env.command(["worker", "hwdetect"], env={env_key: "1,3"})
-    assert "gpus: [1,3]" in resources
+    resources = hq_env.command(["worker", "hwdetect"], env={env: "1,3"})
+    assert f"{resource}: [1,3]" in resources
 
 
 def test_worker_detect_uuid_gpus_from_env(hq_env: HqEnv):
     hq_env.start_server()
     resources = hq_env.command(
         ["worker", "hwdetect"], env={"CUDA_VISIBLE_DEVICES": "foo,bar"}
     )
-    assert "gpus: [foo,bar]" in resources
+    assert "gpus/nvidia: [foo,bar]" in resources
 
 
 def test_task_info_resources(hq_env: HqEnv):