From c54bda300ab2d9128eded374b802a4779302b9ff Mon Sep 17 00:00:00 2001 From: HAI Date: Mon, 2 Dec 2024 00:15:45 -0800 Subject: [PATCH] Use rocminfo instead of rocm-smi for more OS/WSL support (#2310) --- python/sglang/srt/utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/utils.py b/python/sglang/srt/utils.py index fef4c58a5..d0bb767d7 100644 --- a/python/sglang/srt/utils.py +++ b/python/sglang/srt/utils.py @@ -903,7 +903,9 @@ def get_amdgpu_memory_capacity(): try: # Run rocm-smi and capture the output result = subprocess.run( - ["rocm-smi --showmeminfo vram | grep 'Total Memory' | awk '{print $NF}'"], + [ + "rocminfo | grep 'gfx94' -A 100 | grep 'Pool 1' -A 5 | grep 'Size:' | awk '{print $2}'" + ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, @@ -914,9 +916,8 @@ def get_amdgpu_memory_capacity(): # Parse the output to extract memory values in MiB memory_values = [ - float(mem) / 1024 / 1024 + float(mem.split("(")[0].strip()) / 1024 for mem in result.stdout.strip().split("\n") - if re.match(r"^\d+(\.\d+)?$", mem.strip()) ] if not memory_values: