From cd4464db1a1f77d54a7e39de2b8d0b15090a347a Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Wed, 14 Oct 2026 23:21:43 -0400 Subject: [PATCH 1/5] Fix GPU-info paths to support AMD ROCm and degrade gracefully MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gpu_memory_map(): try nvidia-smi first, fall back to rocm-smi (--showmeminfo vram --csv; free = total - used, bytes -> MB); wrap both in try/except so failures return {0: 4300} not raise. - /gpu endpoint: shutil.which() probe; return 200 + plain text when no SMI tool is found instead of raising FileNotFoundError (HTTP 500). - Add tests/unit/utils/test_generic.py: 7 mocked cases covering both backends, runnable on any CI runner without a GPU. Verified on MI300X (gfx942, ROCm 7.1): gpu_memory_map() -> {0: 65507}. Also corrects the docstring ("memory usage" -> "free memory"; the code has always queried memory.free — documentation fix only). Signed-off-by: Patel, Nilaykumar K Co-authored-by: Anik Chaudhuri Co-authored-by: Vikas C Sajjan Co-authored-by: Soumitra Chatterjee Co-authored-by: Prateek Chokse --- monailabel/endpoints/logs.py | 9 +++- monailabel/utils/others/generic.py | 52 +++++++++++++++----- tests/unit/utils/test_generic.py | 79 ++++++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 14 deletions(-) create mode 100644 tests/unit/utils/test_generic.py diff --git a/monailabel/endpoints/logs.py b/monailabel/endpoints/logs.py index f44e1938b..d249a7c07 100644 --- a/monailabel/endpoints/logs.py +++ b/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,11 @@ 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. + smi = "nvidia-smi" if shutil.which("nvidia-smi") else ("rocm-smi" if shutil.which("rocm-smi") else None) + if smi is None: + 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") diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index a26edec63..0d0e38ca2 100644 --- a/monailabel/utils/others/generic.py +++ b/monailabel/utils/others/generic.py @@ -158,21 +158,49 @@ 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().split("\n")] + return dict(zip(range(len(gpu_memory)), gpu_memory)) + except Exception as e: + logger.info(f"nvidia-smi command didn't work ({e})! - Using default image size [128, 128, 64]") + return {0: 4300} + + 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 = {} + idx = 0 + for line in result.strip().split("\n"): + cols = line.split(",") + # Skip the header and any warning lines rocm-smi prints to stdout. + if not cols or not cols[0].startswith("card"): + continue + try: + total_b, used_b = int(cols[1]), int(cols[2]) + except (IndexError, ValueError): + continue + gpu_memory[idx] = max(0, (total_b - used_b) // (1024 * 1024)) # bytes -> MB free + idx += 1 + if gpu_memory: + return gpu_memory + except Exception as e: + logger.info(f"rocm-smi command didn't work ({e})! - Using default image size [128, 128, 64]") 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(): diff --git a/tests/unit/utils/test_generic.py b/tests/unit/utils/test_generic.py new file mode 100644 index 000000000..062eb77be --- /dev/null +++ b/tests/unit/utils/test_generic.py @@ -0,0 +1,79 @@ +# 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_returns_default(self): + with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("nvidia-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(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_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() From e1370d661bf3aae186cecb2d0ec46a0c120e7914 Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Thu, 15 Oct 2026 01:29:10 -0400 Subject: [PATCH 2/5] Address review comments: fallthrough, card index, planner guard - gpu_memory_map(): nvidia-smi failure now falls through to rocm-smi instead of returning the default immediately. - gpu_memory_map(): parse card index from cols[0] ("cardN" -> N) so valid memory stays associated with the correct device when a row is skipped; preserve a zero entry for devices with unparsable data. - gpu_memory_map(): log when rocm-smi runs but produces no parsable card rows so operators can distinguish subprocess failure from a format change. - planner._get_target_img_size(): guard against nonpositive calculated width (exhausted GPU) before calling log2; log and fall back to the target size rather than raising a math domain error. - Update tests: replace test_nvidia_smi_failure_returns_default with test_nvidia_smi_failure_falls_through_to_rocm; add test_rocm_smi_bad_row_preserves_card_index. Signed-off-by: Patel, Nilaykumar K Co-authored-by: Anik Chaudhuri Co-authored-by: Vikas C Sajjan Co-authored-by: Soumitra Chatterjee Co-authored-by: Prateek Chokse --- monailabel/utils/others/generic.py | 15 +++++++++------ monailabel/utils/others/planner.py | 3 +++ tests/unit/utils/test_generic.py | 28 ++++++++++++++++++++++++---- 3 files changed, 36 insertions(+), 10 deletions(-) diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index 0d0e38ca2..25a91c76d 100644 --- a/monailabel/utils/others/generic.py +++ b/monailabel/utils/others/generic.py @@ -172,8 +172,7 @@ def gpu_memory_map(): gpu_memory = [int(x) for x in result.strip().split("\n")] return dict(zip(range(len(gpu_memory)), gpu_memory)) except Exception as e: - logger.info(f"nvidia-smi command didn't work ({e})! - Using default image size [128, 128, 64]") - return {0: 4300} + 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") @@ -181,20 +180,24 @@ def gpu_memory_map(): # 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 = {} - idx = 0 for line in result.strip().split("\n"): cols = line.split(",") # 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"; preserve a zero entry for unparsable rows. + try: + dev_idx = int(cols[0][4:]) + except ValueError: + continue 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 - gpu_memory[idx] = max(0, (total_b - used_b) // (1024 * 1024)) # bytes -> MB free - idx += 1 + gpu_memory[dev_idx] = 0 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 default image size [128, 128, 64]") return {0: 4300} diff --git a/monailabel/utils/others/planner.py b/monailabel/utils/others/planner.py index 30e4ebe6b..b83beda4d 100644 --- a/monailabel/utils/others/planner.py +++ b/monailabel/utils/others/planner.py @@ -92,6 +92,9 @@ def _get_target_img_size(target_img_size): # 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] diff --git a/tests/unit/utils/test_generic.py b/tests/unit/utils/test_generic.py index 062eb77be..3bd502ec4 100644 --- a/tests/unit/utils/test_generic.py +++ b/tests/unit/utils/test_generic.py @@ -31,11 +31,17 @@ def test_nvidia_smi(self): ): self.assertEqual(gpu_memory_map(), {0: 1000, 1: 2000}) - def test_nvidia_smi_failure_returns_default(self): - with mock.patch("monailabel.utils.others.generic.shutil.which", side_effect=_which("nvidia-smi")), mock.patch( - "monailabel.utils.others.generic.subprocess.check_output", side_effect=OSError("boom") + 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: 4300}) + 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. @@ -70,6 +76,20 @@ def test_rocm_smi_unparsable_output_returns_default(self): ): 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, not 0. + csv = ( + "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=csv + ): + result = gpu_memory_map() + self.assertIn(1, 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}) From 670d2240f7fbede9871c76b595e4b1afdca6b452 Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Thu, 15 Oct 2026 02:02:10 -0400 Subject: [PATCH 3/5] Fix misleading log message in rocm-smi exception handler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When rocm-smi fails, gpu_memory_map() returns {0: 4300}. The planner then uses that estimate to calculate image size — it does not apply [128, 128, 64] directly. The old log message said "Using default image size [128, 128, 64]", which misled operators into thinking a fixed size was being applied rather than a memory-based estimate. Updated to: "Using fallback free-memory estimate 4300 MB for image-size planning". Addresses CodeRabbit review comment on PR #1925. Signed-off-by: Patel, Nilaykumar K --- monailabel/utils/others/generic.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index 25a91c76d..28222c6e5 100644 --- a/monailabel/utils/others/generic.py +++ b/monailabel/utils/others/generic.py @@ -199,7 +199,9 @@ def gpu_memory_map(): 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 default image size [128, 128, 64]") + logger.info( + f"rocm-smi command didn't work ({e})! - Using fallback free-memory estimate 4300 MB for image-size planning" + ) return {0: 4300} logger.info("Neither nvidia-smi nor rocm-smi found! - Using default image size [128, 128, 64]") From d15ff2449fd816b6f61c3e9283279d0b91f40e8b Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Wed, 21 Oct 2026 03:08:32 -0400 Subject: [PATCH 4/5] Address ericspod review comments logs.py: rewrite nested ternary as if/elif/else for clarity. generic.py: use csv.reader + re.findall for rocm-smi CSV parsing; use splitlines() with empty-line filter for nvidia-smi output. Signed-off-by: Patel, Nilaykumar K --- monailabel/endpoints/logs.py | 7 +++++-- monailabel/utils/others/generic.py | 17 ++++++++++------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/monailabel/endpoints/logs.py b/monailabel/endpoints/logs.py index d249a7c07..acb9e4488 100644 --- a/monailabel/endpoints/logs.py +++ b/monailabel/endpoints/logs.py @@ -103,8 +103,11 @@ async def api_get_logs( @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))): # nvidia-smi on NVIDIA, rocm-smi on AMD ROCm. - smi = "nvidia-smi" if shutil.which("nvidia-smi") else ("rocm-smi" if shutil.which("rocm-smi") else None) - if smi is None: + 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") diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index 28222c6e5..c3dfd21b7 100644 --- a/monailabel/utils/others/generic.py +++ b/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 @@ -169,7 +170,9 @@ def gpu_memory_map(): 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().split("\n")] + 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") @@ -180,16 +183,16 @@ def gpu_memory_map(): # 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 = {} - for line in result.strip().split("\n"): - cols = line.split(",") + 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"; preserve a zero entry for unparsable rows. - try: - dev_idx = int(cols[0][4:]) - except ValueError: + # 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 From 5521672cae830361741a5cea3b92ae20880c9c56 Mon Sep 17 00:00:00 2001 From: "Patel, Nilaykumar K" Date: Wed, 21 Oct 2026 05:34:00 -0400 Subject: [PATCH 5/5] Skip unparsable rocm-smi rows; guard planner key-0 access A malformed card row was stored as gpu_memory[dev_idx] = 0, signalling exhausted VRAM to the planner. Unknown memory is not the same as zero free memory. Skip the row instead so it contributes nothing to the dict. planner.py uses gpu_memory_map()[0] unconditionally. If card0's row is absent (e.g. malformed and now skipped, or only card1 parsed), this raises KeyError. Switch to .get(0, 4300) so the planner falls back to the established conservative default when device 0 is not present. Addresses CodeRabbit review comment on generic.py:200. Signed-off-by: Patel, Nilaykumar K --- monailabel/utils/others/generic.py | 2 +- monailabel/utils/others/planner.py | 2 +- tests/unit/utils/test_generic.py | 11 ++++++----- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index c3dfd21b7..756f57d54 100644 --- a/monailabel/utils/others/generic.py +++ b/monailabel/utils/others/generic.py @@ -197,7 +197,7 @@ def gpu_memory_map(): 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): - gpu_memory[dev_idx] = 0 + 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]") diff --git a/monailabel/utils/others/planner.py b/monailabel/utils/others/planner.py index b83beda4d..e21b36503 100644 --- a/monailabel/utils/others/planner.py +++ b/monailabel/utils/others/planner.py @@ -86,7 +86,7 @@ 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 diff --git a/tests/unit/utils/test_generic.py b/tests/unit/utils/test_generic.py index 3bd502ec4..3fcf394a8 100644 --- a/tests/unit/utils/test_generic.py +++ b/tests/unit/utils/test_generic.py @@ -77,18 +77,19 @@ def test_rocm_smi_unparsable_output_returns_default(self): 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, not 0. - csv = ( + # 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=csv + "monailabel.utils.others.generic.subprocess.check_output", return_value=smi_output ): result = gpu_memory_map() - self.assertIn(1, result) - self.assertEqual(result[1], 3072) + 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):