diff --git a/monailabel/endpoints/logs.py b/monailabel/endpoints/logs.py index f44e1938b..acb9e4488 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,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") diff --git a/monailabel/utils/others/generic.py b/monailabel/utils/others/generic.py index a26edec63..756f57d54 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 @@ -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(): diff --git a/monailabel/utils/others/planner.py b/monailabel/utils/others/planner.py index 30e4ebe6b..e21b36503 100644 --- a/monailabel/utils/others/planner.py +++ b/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] diff --git a/tests/unit/utils/test_generic.py b/tests/unit/utils/test_generic.py new file mode 100644 index 000000000..3fcf394a8 --- /dev/null +++ b/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()