anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[BUGFIX] Delete the CUDA context in the stage process. (#661)

Browse Source 
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: Chenguang Zheng <645327136@qq.com>
This commit is contained in:
Chenguang Zheng authored and GitHub committed 2026-01-06 22:22:04 +08:00
1 parent 14c04d0fd4
commit aaba243832
2 files changed
+1 -58
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/test_stage_utils.py
-1
View File
@@ -55,4 +55,3 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
set_stage_devices(stage_id=0, devices="0,1")
assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
assert call_log and call_log[0] == 0 # current device set after remap
vllm_omni/entrypoints/stage_utils.py
+1 -57
View File
@@ -64,35 +64,6 @@ def set_stage_devices(
env_var = get_device_control_env_var()
# Select device-specific torch functions
if device_type == "npu":
try:
import torch.npu # type: ignore[import-untyped]
except ImportError:
logger.debug("[Stage-%s] torch.npu not available, skipping NPU device setup", stage_id)
return
is_available_fn = torch.npu.is_available
set_device_fn = torch.npu.set_device
device_count_fn = torch.npu.device_count
get_device_properties_fn = torch.npu.get_device_properties
mem_get_info_fn = torch.npu.mem_get_info
get_device_name_fn = torch.npu.get_device_name
device_type_label = "NPU"
elif device_type == "cuda":
import torch
is_available_fn = torch.cuda.is_available
set_device_fn = torch.cuda.set_device
device_count_fn = torch.cuda.device_count
get_device_properties_fn = torch.cuda.get_device_properties
mem_get_info_fn = torch.cuda.mem_get_info
get_device_name_fn = torch.cuda.get_device_name
device_type_label = "CUDA"
else:
logger.debug("[Stage-%s] Unsupported device type: %s", stage_id, device_type)
return
try:
selected_physical: int | None = None
logical_idx: int | None = None
@@ -130,7 +101,7 @@ def set_stage_devices(
selected_physical,
)
except Exception as e:
logger.debug("[Stage-%s] Failed to parse first %s device: %s", stage_id, device_type_label, e)
logger.debug("[Stage-%s] Failed to parse first %s device: %s", stage_id, device_type, e)
selected_physical = None
elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
logical_idx = max(0, int(devices))
@@ -159,33 +130,6 @@ def set_stage_devices(
selected_physical = int(str(devices))
os.environ[env_var] = str(selected_physical)
logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
try:
import torch
if is_available_fn():
try:
set_device_fn(0)
except Exception as e:
logger.debug(
"[Stage-%s] %s set_device(0) failed: %s", stage_id, device_type_label, e, exc_info=True
)
num = device_count_fn()
info = []
for i in range(num):
total = get_device_properties_fn(i).total_memory
free, _ = mem_get_info_fn(i)
info.append(
{
"idx": i,
"name": get_device_name_fn(i),
"total": int(total),
"free": int(free),
}
)
logger.debug("[Stage-%s] %s devices visible=%s info=%s", stage_id, device_type_label, num, info)
except Exception as e:
logger.debug("[Stage-%s] Failed to query %s devices: %s", stage_id, device_type_label, e, exc_info=True)
except Exception as e:
logger.warning("Failed to interpret devices for stage %s: %s", stage_id, e)