Merge pull request #2324 from transformerlab/add/runpod_amd

Add Runpod support for AMD GPUs and images

Author: dadmobile

api/test/compute_providers/test_runpod_image.py

@@ -56,3 +56,25 @@ def test_default_image_when_no_override():
     config = ClusterConfig(accelerators="A100:1")
     provider.launch_cluster("clust-1", config)
     assert captured["pod_data"]["imageName"] == "runpod/pytorch:1.0.3-cu1281-torch290-ubuntu2204"
+
+
+def test_amd_gpu_defaults_to_rocm_image():
+    provider = _provider()
+    captured = _capture_pod_data(provider)
+    # AMD/ROCm card resolves to a ROCm image, not the default CUDA image.
+    provider._map_gpu_type_to_runpod = lambda _accel: "AMD Instinct MI300X OAM"
+    config = ClusterConfig(accelerators="MI300X:1")
+    provider.launch_cluster("clust-1", config)
+    assert captured["pod_data"]["imageName"] == "rocm/pytorch:latest"
+
+
+def test_amd_gpu_image_name_override_still_honored():
+    provider = _provider()
+    captured = _capture_pod_data(provider)
+    provider._map_gpu_type_to_runpod = lambda _accel: "AMD Instinct MI300X OAM"
+    config = ClusterConfig(
+        accelerators="MI300X:1",
+        provider_config={"image_name": "my-org/rocm-custom:latest"},
+    )
+    provider.launch_cluster("clust-1", config)
+    assert captured["pod_data"]["imageName"] == "my-org/rocm-custom:latest"

api/transformerlab/compute_providers/runpod.py

@@ -68,8 +68,33 @@
     "RTX5000": "NVIDIA RTX 5000 Ada Generation",
     "RTX4000": "NVIDIA RTX 4000 Ada Generation",
     "RTX2000": "NVIDIA RTX 2000 Ada Generation",
+    # AMD Instinct (ROCm). Launching these requires a ROCm container image, not
+    # the default CUDA image — see _is_amd_gpu_id and the image selection in
+    # launch_cluster().
+    "MI300X": "AMD Instinct MI300X OAM",
 }
 
+# Runpod GPU type IDs for AMD/ROCm cards. Used to pick a ROCm container image at
+# launch time instead of the default CUDA image.
+_RUNPOD_AMD_GPU_IDS: frozenset = frozenset(
+    {
+        "AMD Instinct MI300X OAM",
+    }
+)
+
+# Default ROCm container image for AMD pods. Overridable per-launch via the
+# provider config `image_name` (or legacy `template_id`).
+_RUNPOD_DEFAULT_ROCM_IMAGE = "rocm/pytorch:latest"
+# Default CUDA container image for NVIDIA pods.
+_RUNPOD_DEFAULT_CUDA_IMAGE = "runpod/pytorch:1.0.3-cu1281-torch290-ubuntu2204"
+
+
+def _is_amd_gpu_id(gpu_type_id: Optional[str]) -> bool:
+    """Return True if the resolved Runpod GPU type ID is an AMD/ROCm card."""
+    if not gpu_type_id:
+        return False
+    return gpu_type_id in _RUNPOD_AMD_GPU_IDS or "AMD" in gpu_type_id
+
 
 async def fetch_runpod_provider_logs(
     provider_instance: ComputeProvider,
@@ -360,8 +385,12 @@ def launch_cluster(self, cluster_name: str, config: ClusterConfig) -> Dict[str,
                 except ValueError:
                     gpu_count = 1
 
-            # Use GPU-enabled image
-            default_image = "runpod/pytorch:1.0.3-cu1281-torch290-ubuntu2204"
+            # Use a GPU-enabled image. AMD/ROCm cards need a ROCm image; the
+            # default CUDA image will not run on them.
+            if _is_amd_gpu_id(gpu_type_id):
+                default_image = _RUNPOD_DEFAULT_ROCM_IMAGE
+            else:
+                default_image = _RUNPOD_DEFAULT_CUDA_IMAGE
         else:
             # CPU pod
             compute_type = "CPU"