[Misc] Use logger instead of print() in utils/common.py (#29004)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
hirakunaramuka2
2026-06-24 09:45:01 +00:00
committed by GitHub
co-authored by Claude Fable 5
parent 10e0bcd622
commit 7430c56b20
2 changed files with 126 additions and 19 deletions
+28 -18
View File
@@ -566,9 +566,11 @@ def get_available_gpu_memory(
assert gpu_id < num_gpus assert gpu_id < num_gpus
if torch.cuda.current_device() != gpu_id: if torch.cuda.current_device() != gpu_id:
print( logger.warning(
f"WARNING: current device is not {gpu_id}, but {torch.cuda.current_device()}, ", "current device is not %s, but %s, which may cause useless "
"which may cause useless memory allocation for torch CUDA context.", "memory allocation for torch CUDA context.",
gpu_id,
torch.cuda.current_device(),
) )
if empty_cache: if empty_cache:
@@ -588,9 +590,11 @@ def get_available_gpu_memory(
assert gpu_id < num_gpus assert gpu_id < num_gpus
if torch.xpu.current_device() != gpu_id: if torch.xpu.current_device() != gpu_id:
print( logger.warning(
f"WARNING: current device is not {gpu_id}, but {torch.xpu.current_device()}, ", "current device is not %s, but %s, which may cause useless "
"which may cause useless memory allocation for torch XPU context.", "memory allocation for torch XPU context.",
gpu_id,
torch.xpu.current_device(),
) )
if empty_cache: if empty_cache:
@@ -604,9 +608,11 @@ def get_available_gpu_memory(
assert gpu_id < num_gpus assert gpu_id < num_gpus
if torch.hpu.current_device() != gpu_id: if torch.hpu.current_device() != gpu_id:
print( logger.warning(
f"WARNING: current device is not {gpu_id}, but {torch.hpu.current_device()}, ", "current device is not %s, but %s, which may cause useless "
"which may cause useless memory allocation for torch HPU context.", "memory allocation for torch HPU context.",
gpu_id,
torch.hpu.current_device(),
) )
free_gpu_memory, total_gpu_memory = torch.hpu.mem_get_info() 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 assert gpu_id < num_gpus
if torch.npu.current_device() != gpu_id: if torch.npu.current_device() != gpu_id:
print( logger.warning(
f"WARNING: current device is not {gpu_id}, but {torch.npu.current_device()}, ", "current device is not %s, but %s, which may cause useless "
"which may cause useless memory allocation for torch NPU context.", "memory allocation for torch NPU context.",
gpu_id,
torch.npu.current_device(),
) )
if empty_cache: if empty_cache:
empty_device_cache(torch.npu) empty_device_cache(torch.npu)
@@ -642,9 +650,11 @@ def get_available_gpu_memory(
assert gpu_id < num_gpus assert gpu_id < num_gpus
if torch.musa.current_device() != gpu_id: if torch.musa.current_device() != gpu_id:
print( logger.warning(
f"WARNING: current device is not {gpu_id}, but {torch.musa.current_device()}, ", "current device is not %s, but %s, which may cause useless "
"which may cause useless memory allocation for torch MUSA context.", "memory allocation for torch MUSA context.",
gpu_id,
torch.musa.current_device(),
) )
if empty_cache: if empty_cache:
empty_device_cache(torch.musa) empty_device_cache(torch.musa)
@@ -1539,7 +1549,7 @@ def delete_directory(dirpath):
# This will remove the directory and all its contents # This will remove the directory and all its contents
shutil.rmtree(dirpath) shutil.rmtree(dirpath)
except OSError as e: 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 # Temporary directory for prometheus multiprocess mode
@@ -3848,7 +3858,7 @@ def get_nvidia_driver_version() -> tuple:
@lru_cache(maxsize=1) @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'. """Return the NVIDIA driver version string, e.g. '595.58.03'.
Returns None on failure.""" Returns None on failure."""
try: try:
@@ -3930,7 +3940,7 @@ def get_device_sm_nvidia_smi():
except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e: except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
# Handle cases where nvidia-smi isn't available or output is unexpected # 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 return (0, 0) # Default/fallback value
+98 -1
View File
@@ -3,7 +3,11 @@ from array import array
import torch 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.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -46,5 +50,98 @@ class TestFlattenArraysToInt64Tensor(CustomTestCase):
self._check(parts, [10, 20, 30, 100, 200, 1000]) 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__": if __name__ == "__main__":
unittest.main() unittest.main()