From 7430c56b20f8845362e00af448f68a7d0a162c65 Mon Sep 17 00:00:00 2001 From: hirakunaramuka2 <295779535+hirakunaramuka2@users.noreply.github.com> Date: Wed, 24 Jun 2026 17:45:01 +0800 Subject: [PATCH] [Misc] Use logger instead of print() in utils/common.py (#29004) Co-authored-by: Claude Fable 5 --- python/sglang/srt/utils/common.py | 46 ++++++----- test/registered/unit/utils/test_common.py | 99 ++++++++++++++++++++++- 2 files changed, 126 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 9227c0a4a..ea9e9b252 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -566,9 +566,11 @@ def get_available_gpu_memory( assert gpu_id < num_gpus if torch.cuda.current_device() != gpu_id: - print( - f"WARNING: current device is not {gpu_id}, but {torch.cuda.current_device()}, ", - "which may cause useless memory allocation for torch CUDA context.", + logger.warning( + "current device is not %s, but %s, which may cause useless " + "memory allocation for torch CUDA context.", + gpu_id, + torch.cuda.current_device(), ) if empty_cache: @@ -588,9 +590,11 @@ def get_available_gpu_memory( assert gpu_id < num_gpus if torch.xpu.current_device() != gpu_id: - print( - f"WARNING: current device is not {gpu_id}, but {torch.xpu.current_device()}, ", - "which may cause useless memory allocation for torch XPU context.", + logger.warning( + "current device is not %s, but %s, which may cause useless " + "memory allocation for torch XPU context.", + gpu_id, + torch.xpu.current_device(), ) if empty_cache: @@ -604,9 +608,11 @@ def get_available_gpu_memory( assert gpu_id < num_gpus if torch.hpu.current_device() != gpu_id: - print( - f"WARNING: current device is not {gpu_id}, but {torch.hpu.current_device()}, ", - "which may cause useless memory allocation for torch HPU context.", + logger.warning( + "current device is not %s, but %s, which may cause useless " + "memory allocation for torch HPU context.", + gpu_id, + torch.hpu.current_device(), ) free_gpu_memory, total_gpu_memory = torch.hpu.mem_get_info() @@ -621,9 +627,11 @@ def get_available_gpu_memory( assert gpu_id < num_gpus if torch.npu.current_device() != gpu_id: - print( - f"WARNING: current device is not {gpu_id}, but {torch.npu.current_device()}, ", - "which may cause useless memory allocation for torch NPU context.", + logger.warning( + "current device is not %s, but %s, which may cause useless " + "memory allocation for torch NPU context.", + gpu_id, + torch.npu.current_device(), ) if empty_cache: empty_device_cache(torch.npu) @@ -642,9 +650,11 @@ def get_available_gpu_memory( assert gpu_id < num_gpus if torch.musa.current_device() != gpu_id: - print( - f"WARNING: current device is not {gpu_id}, but {torch.musa.current_device()}, ", - "which may cause useless memory allocation for torch MUSA context.", + logger.warning( + "current device is not %s, but %s, which may cause useless " + "memory allocation for torch MUSA context.", + gpu_id, + torch.musa.current_device(), ) if empty_cache: empty_device_cache(torch.musa) @@ -1539,7 +1549,7 @@ def delete_directory(dirpath): # This will remove the directory and all its contents shutil.rmtree(dirpath) except OSError as e: - print(f"Warning: {dirpath} : {e.strerror}") + logger.warning("Failed to delete directory %s: %s", dirpath, e.strerror) # Temporary directory for prometheus multiprocess mode @@ -3848,7 +3858,7 @@ def get_nvidia_driver_version() -> tuple: @lru_cache(maxsize=1) -def get_nvidia_driver_version_str() -> str: +def get_nvidia_driver_version_str() -> str | None: """Return the NVIDIA driver version string, e.g. '595.58.03'. Returns None on failure.""" try: @@ -3930,7 +3940,7 @@ def get_device_sm_nvidia_smi(): except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e: # Handle cases where nvidia-smi isn't available or output is unexpected - print(f"Error getting compute capability: {e}") + logger.error("Error getting compute capability: %s", e) return (0, 0) # Default/fallback value diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py index b38e1e7f7..f38047601 100644 --- a/test/registered/unit/utils/test_common.py +++ b/test/registered/unit/utils/test_common.py @@ -3,7 +3,11 @@ from array import array import torch -from sglang.srt.utils.common import flatten_arrays_to_int64_tensor +from sglang.srt.utils.common import ( + flatten_arrays_to_int64_tensor, + get_device_sm_nvidia_smi, + get_nvidia_driver_version_str, +) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -46,5 +50,98 @@ class TestFlattenArraysToInt64Tensor(CustomTestCase): self._check(parts, [10, 20, 30, 100, 200, 1000]) +class TestNvidiaDriverVersionStr(CustomTestCase): + """`get_nvidia_driver_version_str` is typed as `str | None`: it returns + `None` when nvidia-smi is missing, fails, or emits an empty string. These + tests exercise both the success and the None-return paths by monkey- + patching `subprocess.run`, so they don't require a GPU. The function is + `@lru_cache`d, so the cache is cleared around each test to make the patch + observable. + """ + + def setUp(self): + get_nvidia_driver_version_str.cache_clear() + + def tearDown(self): + get_nvidia_driver_version_str.cache_clear() + + def test_returns_version_string(self): + import subprocess + + class _R: + stdout = "595.58.03\n" + + original = subprocess.run + subprocess.run = lambda *a, **k: _R() + try: + self.assertEqual(get_nvidia_driver_version_str(), "595.58.03") + finally: + subprocess.run = original + + def test_returns_none_on_empty_output(self): + import subprocess + + class _R: + stdout = "\n" + + original = subprocess.run + subprocess.run = lambda *a, **k: _R() + try: + self.assertIsNone(get_nvidia_driver_version_str()) + finally: + subprocess.run = original + + def test_returns_none_on_called_process_error(self): + import subprocess + + original = subprocess.run + + def boom(*a, **k): + raise subprocess.CalledProcessError(1, "nvidia-smi") + + subprocess.run = boom + try: + self.assertIsNone(get_nvidia_driver_version_str()) + finally: + subprocess.run = original + + def test_returns_none_on_file_not_found(self): + import subprocess + + original = subprocess.run + + def boom(*a, **k): + raise FileNotFoundError("nvidia-smi") + + subprocess.run = boom + try: + self.assertIsNone(get_nvidia_driver_version_str()) + finally: + subprocess.run = original + + +class TestGetDeviceSmNvidiaSmi(CustomTestCase): + """`get_device_sm_nvidia_smi` parses nvidia-smi output into a (major, + minor) tuple and falls back to (0, 0) -- logging via `logger.error` -- + when nvidia-smi fails. The success path needs a GPU; the fallback path is + covered here by forcing a failure and asserting the (0, 0) return. The + fallback path needs no GPU, so this test runs on CPU. + """ + + def test_fallback_on_failure_returns_zero_zero(self): + import subprocess + + original = subprocess.run + + def boom(*a, **k): + raise subprocess.CalledProcessError(1, "nvidia-smi") + + subprocess.run = boom + try: + self.assertEqual(get_device_sm_nvidia_smi(), (0, 0)) + finally: + subprocess.run = original + + if __name__ == "__main__": unittest.main()