Fix /gpu endpoint and gpu_memory_map() to support AMD ROCm and degrade gracefully

monailabel/endpoints/logs.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 
 import os
+import shutil
 import subprocess
 from collections import deque
 from typing import Optional
@@ -99,7 +100,14 @@ async def api_get_logs(
     return get_logs(os.path.join(settings.MONAI_LABEL_APP_DIR, "logs", "app.log"), lines, html, text, refresh)
 
 
-@router.get("/gpu", summary=f"{RBAC_ADMIN}Get GPU Info (nvidia-smi)")
+@router.get("/gpu", summary=f"{RBAC_ADMIN}Get GPU Info (nvidia-smi / rocm-smi)")
 async def gpu_info(user: User = Depends(RBAC(settings.MONAI_LABEL_AUTH_ROLE_ADMIN))):
-    response = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE).stdout.decode("utf-8")
+    # nvidia-smi on NVIDIA, rocm-smi on AMD ROCm.
+    if shutil.which("nvidia-smi"):
+        smi = "nvidia-smi"
+    elif shutil.which("rocm-smi"):
+        smi = "rocm-smi"
+    else:
+        return Response(content="No GPU SMI tool found (nvidia-smi / rocm-smi)", media_type="text/plain")
+    response = subprocess.run([smi], stdout=subprocess.PIPE).stdout.decode("utf-8")
     return Response(content=response, media_type="text/plain")

monailabel/utils/others/generic.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import csv
 import distutils.util
 import hashlib
 import json
@@ -158,21 +159,56 @@ def gpu_memory_map():
     -------
     usage: dict
         Keys are device ids as integers.
-        Values are memory usage as integers in MB.
+        Values are free memory as integers in MB.
     """
-    logger.info("Using nvidia-smi command")
-    if shutil.which("nvidia-smi") is None:
-        logger.info("nvidia-smi command didn't work! - Using default image size [128, 128, 64]")
+    # Prefer nvidia-smi; fall back to rocm-smi on AMD ROCm hosts. Both report
+    # per-device free VRAM. On any failure keep the original safe default so the
+    # caller's image-size auto-tuning degrades gracefully instead of raising.
+    if shutil.which("nvidia-smi") is not None:
+        logger.info("Using nvidia-smi command")
+        try:
+            result = subprocess.check_output(
+                ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,nounits,noheader"], encoding="utf-8"
+            )
+            gpu_memory = [int(x) for x in result.strip().splitlines() if x.strip()]
+            if not gpu_memory:
+                raise ValueError("nvidia-smi returned no GPU memory information!")
+            return dict(zip(range(len(gpu_memory)), gpu_memory))
+        except Exception as e:
+            logger.info(f"nvidia-smi command didn't work ({e})! - Trying rocm-smi")
+
+    if shutil.which("rocm-smi") is not None:
+        logger.info("Using rocm-smi command")
+        try:
+            # rocm-smi reports total + used VRAM in bytes per card; free = total - used.
+            result = subprocess.check_output(["rocm-smi", "--showmeminfo", "vram", "--csv"], encoding="utf-8")
+            gpu_memory = {}
+            reader = csv.reader(result.strip().splitlines())
+            for cols in reader:
+                # Skip the header and any warning lines rocm-smi prints to stdout.
+                if not cols or not cols[0].startswith("card"):
+                    continue
+                # Parse device index from "cardN" (e.g. "card0" -> 0).
+                idx_match = re.findall(r"card([0-9]+)", cols[0])
+                if not idx_match:
+                    continue
+                dev_idx = int(idx_match[0])
+                try:
+                    total_b, used_b = int(cols[1]), int(cols[2])
+                    gpu_memory[dev_idx] = (total_b - used_b) // (1024 * 1024)  # bytes -> MB free
+                except (IndexError, ValueError):
+                    continue  # skip unparsable row; unknown != exhausted
+            if gpu_memory:
+                return gpu_memory
+            logger.info("rocm-smi output had no parsable card rows! - Using default image size [128, 128, 64]")
+        except Exception as e:
+            logger.info(
+                f"rocm-smi command didn't work ({e})! - Using fallback free-memory estimate 4300 MB for image-size planning"
+            )
         return {0: 4300}
 
-    result = subprocess.check_output(
-        ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,nounits,noheader"], encoding="utf-8"
-    )
-
-    # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split("\n")]
-    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory))
-    return gpu_memory_map
+    logger.info("Neither nvidia-smi nor rocm-smi found! - Using default image size [128, 128, 64]")
+    return {0: 4300}
 
 
 def gpu_count():

monailabel/utils/others/planner.py

@@ -86,12 +86,15 @@ def _get_target_img_size(target_img_size):
         # This should return an image according to the free gpu memory available
         # Equation obtained from curve fitting using table:
         # https://tinyurl.com/tableGPUMemory
-        gpu_mem = gpu_memory_map()[0]
+        gpu_mem = gpu_memory_map().get(0, 4300)
         # Get a number in base 2 close to the mean depth
         depth_base_2 = int(2 ** np.ceil(np.log2(target_img_size[2])))
         # Get the maximum width according available GPU memory
         # This equation roughly estimates the image size that fits in the available GPU memory using DynUNet
         width = (gpu_mem - 2000) / (0.5 * depth_base_2)
+        if width <= 0:
+            logger.info(f"Insufficient free GPU memory ({gpu_mem} MB); falling back to target size {target_img_size}")
+            return [int(target_img_size[0]), int(target_img_size[1]), depth_base_2]
         width_base_2 = int(2 ** np.round(np.log2(width)))
         if width_base_2 < np.maximum(target_img_size[0], target_img_size[1]):
             return [width_base_2, width_base_2, depth_base_2]

tests/unit/utils/test_generic.py

@@ -0,0 +1,100 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from unittest import mock
+
+from monailabel.utils.others.generic import gpu_memory_map
+
+GIB = 1024 * 1024 * 1024
+
+
+def _which(available):
+    def f(cmd):
+        return f"/usr/bin/{cmd}" if cmd == available else None
+
+    return f
+
+
+class TestGpuMemoryMap(unittest.TestCase):
+    def test_nvidia_smi(self):
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("nvidia-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", return_value="1000\n2000\n"
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 1000, 1: 2000})
+
+    def test_nvidia_smi_failure_falls_through_to_rocm(self):
+        # A failing nvidia-smi must not block the rocm-smi fallback.
+        csv = f"device,VRAM Total Memory (B),VRAM Total Used Memory (B)\ncard0,{2 * GIB},{1 * GIB}\n"
+
+        def _which_both(cmd):
+            return f"/usr/bin/{cmd}" if cmd in ("nvidia-smi", "rocm-smi") else None
+
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which_both), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", side_effect=[OSError("boom"), csv]
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 1024})
+
+    def test_rocm_smi(self):
+        # rocm-smi reports total/used VRAM in bytes; free is derived as total - used.
+        csv = f"device,VRAM Total Memory (B),VRAM Total Used Memory (B)\ncard0,{2 * GIB},{1 * GIB}\n"
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("rocm-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", return_value=csv
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 1024})
+
+    def test_rocm_smi_multi_card_skips_warning_lines(self):
+        # rocm-smi prints warnings to stdout ahead of the CSV header; both must be skipped.
+        csv = (
+            "WARNING: AMD GPU device(s) is/are in a low-power state.\n"
+            "device,VRAM Total Memory (B),VRAM Total Used Memory (B)\n"
+            f"card0,{2 * GIB},{1 * GIB}\n"
+            f"card1,{4 * GIB},{1 * GIB}\n"
+        )
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("rocm-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", return_value=csv
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 1024, 1: 3072})
+
+    def test_rocm_smi_failure_returns_default(self):
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("rocm-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", side_effect=OSError("boom")
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 4300})
+
+    def test_rocm_smi_unparsable_output_returns_default(self):
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("rocm-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", return_value="no cards here\n"
+        ):
+            self.assertEqual(gpu_memory_map(), {0: 4300})
+
+    def test_rocm_smi_bad_row_preserves_card_index(self):
+        # card0 has malformed columns; card1 is valid. card1's memory must map to key 1,
+        # and card0 must be absent entirely (not stored as zero free memory).
+        smi_output = (
+            "device,VRAM Total Memory (B),VRAM Total Used Memory (B)\n"
+            "card0,bad,data\n"
+            f"card1,{4 * GIB},{1 * GIB}\n"
+        )
+        with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("rocm-smi")), mock.patch(
+            "monailabel.utils.others.generic.subprocess.check_output", return_value=smi_output
+        ):
+            result = gpu_memory_map()
+            self.assertNotIn(0, result)
+            self.assertEqual(result, {1: 3072})
+
+    def test_no_smi_tool_returns_default(self):
+        with mock.patch("monailabel.utils.others.generic.shutil.which", return_value=None):
+            self.assertEqual(gpu_memory_map(), {0: 4300})
+
+
+if __name__ == "__main__":
+    unittest.main()