diff --git a/tests/worker/test_process_gpu_memory.py b/tests/worker/test_process_gpu_memory.py new file mode 100644 index 000000000..d8e0b969d --- /dev/null +++ b/tests/worker/test_process_gpu_memory.py @@ -0,0 +1,154 @@ +"""Tests for process-scoped GPU memory accounting.""" + +import os +from unittest import mock + +import pytest + + +class TestParseCudaVisibleDevices: + def test_empty(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {}, clear=True): + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + assert parse_cuda_visible_devices() == [] + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + assert parse_cuda_visible_devices() == [] + + def test_integer_indices(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}): + assert parse_cuda_visible_devices() == [2, 3, 5] + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}): + assert parse_cuda_visible_devices() == [0] + + def test_uuids(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + uuid1 = "GPU-12345678-1234-1234-1234-123456789abc" + uuid2 = "GPU-87654321-4321-4321-4321-cba987654321" + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}): + assert parse_cuda_visible_devices() == [uuid1, uuid2] + + def test_mig_ids(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0" + mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0" + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}): + assert parse_cuda_visible_devices() == [mig1, mig2] + + def test_spaces(self): + from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices + + with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}): + assert parse_cuda_visible_devices() == [2, 3, 5] + + +class TestGetProcessGpuMemory: + @pytest.mark.skipif(not os.path.exists("/dev/nvidia0"), reason="No GPU") + def test_returns_memory_for_current_process(self): + import torch + + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + device = torch.device("cuda:0") + tensor = torch.zeros(1000, 1000, device=device) + + memory = get_process_gpu_memory(0) + assert memory >= 0 + + del tensor + torch.cuda.empty_cache() + + def test_raises_on_invalid_device(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1), + ): + with pytest.raises(RuntimeError, match="Invalid GPU device"): + get_process_gpu_memory(5) + + def test_returns_zero_when_process_not_found(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), + ): + memory = get_process_gpu_memory(0) + assert memory == 0 + + def test_uses_uuid_when_provided(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + uuid = "GPU-12345678-1234-1234-1234-123456789abc" + mock_handle = mock.MagicMock() + + with ( + mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid, + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]), + ): + memory = get_process_gpu_memory(0) + assert memory == 0 + mock_by_uuid.assert_called_once_with(uuid) + + def test_raises_on_invalid_uuid(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + uuid = "GPU-invalid-uuid" + + with ( + mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")), + ): + with pytest.raises(RuntimeError, match="Failed to get NVML handle"): + get_process_gpu_memory(0) + + def test_returns_none_on_nvml_init_failure(self): + from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + ): + result = get_process_gpu_memory(0) + assert result is None + + +class TestIsProcessScopedMemoryAvailable: + def test_returns_true_when_nvml_works(self): + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + with ( + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"), + mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"), + ): + assert is_process_scoped_memory_available() is True + + def test_returns_false_when_nvml_fails(self): + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + with mock.patch( + "vllm_omni.worker.gpu_memory_utils.nvmlInit", + side_effect=Exception("NVML unavailable"), + ): + assert is_process_scoped_memory_available() is False diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index a2070f89d..4935c43cc 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -16,6 +16,7 @@ import sys import time import traceback from collections.abc import Sequence +from contextlib import contextmanager from dataclasses import fields from typing import Any, Literal, cast @@ -55,6 +56,144 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +@contextmanager +def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int = 300): + """Acquire device locks for sequential init if NVML is unavailable. + + If process-scoped memory tracking is available (NVML works), stages can + safely initialize concurrently — each measures only its own GPU memory. + Otherwise, fall back to file-based locks to serialize initialization. + """ + from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available + + if is_process_scoped_memory_available(): + logger.debug("NVML process-scoped memory available — concurrent init is safe, skipping locks") + yield + return + + logger.debug("NVML unavailable — using sequential init locks") + + from vllm_omni.platforms import current_omni_platform + + # Get all parallel sizes from engine_args or parallel_config (defaults to 1) + if "parallel_config" in engine_args: + parallel_config = engine_args["parallel_config"] + tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) + pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) + data_parallel_size = parallel_config.get("data_parallel_size", 1) + prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) + sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) + cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + else: + tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) + pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) + data_parallel_size = engine_args.get("data_parallel_size", 1) + prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) + sequence_parallel_size = 1 + cfg_parallel_size = 1 + + num_devices_per_stage = ( + tensor_parallel_size + * pipeline_parallel_size + * data_parallel_size + * prefill_context_parallel_size + * sequence_parallel_size + * cfg_parallel_size + ) + + # Get physical device IDs from device control env var + device_control_env = current_omni_platform.device_control_env_var + visible_devices_str = os.environ.get(device_control_env) + physical_devices = [] + + if visible_devices_str: + try: + physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] + except (ValueError, IndexError): + pass + + if not physical_devices: + num_devices = current_omni_platform.get_device_count() + physical_devices = list(range(num_devices)) + + num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) + devices_to_lock = sorted(physical_devices[:num_devices_to_lock]) + + logger.debug( + "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", + tensor_parallel_size, + pipeline_parallel_size, + data_parallel_size, + prefill_context_parallel_size, + sequence_parallel_size, + cfg_parallel_size, + num_devices_to_lock, + devices_to_lock, + ) + + # Acquire exclusive locks for all devices using fcntl.flock + wait_start = time.time() + acquired_lock_fds = [] + + for device_id in devices_to_lock: + lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" + lock_acquired = False + + while not lock_acquired: + try: + lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644) + + try: + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + os.ftruncate(lock_fd, 0) + os.write(lock_fd, f"{os.getpid()}\n".encode()) + os.fsync(lock_fd) + lock_acquired = True + acquired_lock_fds.append(lock_fd) + logger.debug("Acquired exclusive lock for device %s", device_id) + except BlockingIOError: + os.close(lock_fd) + + if time.time() - wait_start > stage_init_timeout: + logger.warning( + "Timeout waiting for device %s initialization lock, proceeding anyway", + device_id, + ) + break + + time.sleep(0.1) + except OSError as e: + logger.debug( + "Failed to acquire lock for device %s: %s, continuing anyway", + device_id, + e, + ) + try: + os.close(lock_fd) + except (OSError, NameError): + pass + break + + # Set FD_CLOEXEC to prevent child processes from inheriting locks + for lock_fd in acquired_lock_fds: + try: + flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) + fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + except (OSError, ValueError): + pass + + try: + yield + finally: + for lock_fd in acquired_lock_fds: + try: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + logger.debug("Released initialization lock (fd=%s)", lock_fd) + except (OSError, ValueError): + pass + + def _resolve_worker_cls(engine_args: dict[str, Any]) -> None: worker_type = engine_args.pop("worker_type", None) if not worker_type: @@ -541,183 +680,32 @@ def _stage_worker( except Exception as e: logger.warning("Device setup failed: %s", e) - # Sequential initialization on the same device to avoid memory calculation errors - # when multiple instances start simultaneously - # For TP/PP/DP/SP, we need to lock ALL devices that will be used by this stage - lock_files = [] - try: - # Get all parallel sizes from engine_args or parallel_config (defaults to 1) - if "parallel_config" in engine_args: - parallel_config = engine_args["parallel_config"] - tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) - pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) - data_parallel_size = parallel_config.get("data_parallel_size", 1) - prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) - sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) - cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) + # Use sequential init locks only when NVML is unavailable + with _sequential_init_lock(engine_args, stage_init_timeout): + # Init engine based on stage_type + logger.debug( + "[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()) + ) + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id + if stage_type == "diffusion": + engine_args.pop("model_stage", None) + engine_args.pop("model", None) + stage_engine = OmniDiffusion( + model=model, + stage_id=stage_id, + engine_input_source=stage_payload.get("engine_input_source", []), + **engine_args, + ) else: - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) - sequence_parallel_size = 1 # not use in omni model - cfg_parallel_size = 1 # not used in omni model - - # Calculate total number of devices needed for this stage - # For a single stage worker: - # - TP: splits model across devices (always needed) - # - PP: splits layers across pipeline stages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # - SP: sequence parallelism, typically uses TP devices - # - CFG: Classifier-Free Guidance parallelism for diffusion models - # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size - * pipeline_parallel_size - * data_parallel_size - * prefill_context_parallel_size - * sequence_parallel_size - * cfg_parallel_size - ) - - # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) - # After set_stage_devices, this env var is set to physical device(s) - device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = _os.environ.get(device_control_env) - physical_devices = [] - - if visible_devices_str: - try: - physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] - except (ValueError, IndexError): - pass - - if not physical_devices: - # Fallback: use logical device count if device control env var not set - num_devices = current_omni_platform.get_device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - - logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - sequence_parallel_size, - cfg_parallel_size, - num_devices_to_lock, - devices_to_lock, - ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, proceeding anyway", - device_id, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", - device_id, - e, - ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass - break - - lock_files = acquired_lock_fds - - # Set FD_CLOEXEC on all lock file descriptors to prevent child processes - # (e.g., EngineCore) from inheriting them, which would cause deadlock - for lock_fd in acquired_lock_fds: - try: - flags = fcntl.fcntl(lock_fd, fcntl.F_GETFD) - fcntl.fcntl(lock_fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) - except (OSError, ValueError): - pass - except Exception as e: - logger.debug( - "[Stage-%s] Failed to set up sequential initialization lock: %s", - stage_id, - e, - ) - - # Init engine based on stage_type - logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys())) - if engine_args.get("async_chunk", False): - logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) - stage_connector_spec = {} - for v in connectors_config.values(): - stage_connector_spec = dict(v.get("spec", {})) - break - engine_args["stage_connector_spec"] = stage_connector_spec - engine_args["stage_id"] = stage_id - if stage_type == "diffusion": - engine_args.pop("model_stage", None) - engine_args.pop("model", None) - stage_engine = OmniDiffusion( - model=model, - stage_id=stage_id, - engine_input_source=stage_payload.get("engine_input_source", []), - **engine_args, - ) - else: - # Default to LLM engine - stage_engine = OmniLLM(model=model, **engine_args) - - # Release all locks AFTER engine initialization completes - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass - lock_files = [] # Clear after release + # Default to LLM engine + stage_engine = OmniLLM(model=model, **engine_args) logger.debug("Engine initialized") # Initialize OmniConnectors if configured @@ -1086,155 +1074,23 @@ async def _stage_worker_async( return connectors = built_connectors - # Sequential initialization on the same device to avoid memory calculation errors - # when multiple instances start simultaneously - # For TP/PP/DP/PCP, we need to lock ALL devices that will be used by this stage - lock_files = [] - try: - # Get all parallel sizes from engine_args or parallel_config (defaults to 1) - if "parallel_config" in engine_args: - parallel_config = engine_args["parallel_config"] - tensor_parallel_size = parallel_config.get("tensor_parallel_size", 1) - pipeline_parallel_size = parallel_config.get("pipeline_parallel_size", 1) - data_parallel_size = parallel_config.get("data_parallel_size", 1) - prefill_context_parallel_size = parallel_config.get("prefill_context_parallel_size", 1) - sequence_parallel_size = parallel_config.get("sequence_parallel_size", 1) - cfg_parallel_size = parallel_config.get("cfg_parallel_size", 1) - else: - tensor_parallel_size = engine_args.get("tensor_parallel_size", 1) - pipeline_parallel_size = engine_args.get("pipeline_parallel_size", 1) - data_parallel_size = engine_args.get("data_parallel_size", 1) - prefill_context_parallel_size = engine_args.get("prefill_context_parallel_size", 1) - sequence_parallel_size = 1 # not use in omni model - cfg_parallel_size = 1 # not used in omni model - - # Calculate total number of devices needed for this stage - # For a single stage worker: - # - TP: splits model across devices (always needed) - # - PP: splits layers across pipeline stages, but each stage uses TP devices - # - DP: replicates model, but each replica uses TP devices - # - PCP: context parallelism, typically uses TP devices - # - SP: sequence parallelism, typically uses TP devices - # - CFG: Classifier-Free Guidance parallelism for diffusion models - # The number of devices per stage is determined by TP * PP * DP * PCP * SP * CFG size - # (PP/DP/PCP are higher-level parallelism that don't add devices per stage) - num_devices_per_stage = ( - tensor_parallel_size - * pipeline_parallel_size - * data_parallel_size - * prefill_context_parallel_size - * sequence_parallel_size - * cfg_parallel_size - ) - - # Get physical device IDs from device control env var (e.g., CUDA_VISIBLE_DEVICES) - # After set_stage_devices, this env var is set to physical device(s) - device_control_env = current_omni_platform.device_control_env_var - visible_devices_str = _os.environ.get(device_control_env) - physical_devices = [] - - if visible_devices_str: - try: - physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()] - except (ValueError, IndexError): - pass - - if not physical_devices: - # Fallback: use logical device count if device control env var not set - num_devices = current_omni_platform.get_device_count() - physical_devices = list(range(num_devices)) - - # Determine which devices will be used (min of devices per stage and available devices) - num_devices_to_lock = min(num_devices_per_stage, len(physical_devices)) - devices_to_lock = physical_devices[:num_devices_to_lock] - - # Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order) - devices_to_lock = sorted(devices_to_lock) - + # Use sequential init locks only when NVML is unavailable + with _sequential_init_lock(engine_args, stage_init_timeout): + # Init engine based on stage_type logger.debug( - "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s", - tensor_parallel_size, - pipeline_parallel_size, - data_parallel_size, - prefill_context_parallel_size, - sequence_parallel_size, - cfg_parallel_size, - num_devices_to_lock, - devices_to_lock, + "[Stage-%s] Initializing %s engine with args keys=%s", + stage_id, + stage_type, + list(engine_args.keys()), ) - - # Acquire exclusive locks for all devices using fcntl.flock - # Locks are automatically released when process dies - wait_start = _time.time() - acquired_lock_fds = [] # Store file descriptors to keep locks alive - - for device_id in devices_to_lock: - lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock" - lock_acquired = False - - while not lock_acquired: - try: - # Open or create the lock file - lock_fd = _os.open(lock_file, _os.O_CREAT | _os.O_RDWR, 0o644) - - # Try to acquire exclusive lock (non-blocking first) - try: - fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - # Successfully acquired lock - write PID - _os.ftruncate(lock_fd, 0) # Clear file - _os.write(lock_fd, f"{_os.getpid()}\n".encode()) - _os.fsync(lock_fd) # Ensure written to disk - lock_acquired = True - acquired_lock_fds.append(lock_fd) - logger.debug("Acquired exclusive lock for device %s", device_id) - except BlockingIOError: - # Lock is held by another process - _os.close(lock_fd) - - # Check if we've been waiting too long - if _time.time() - wait_start > stage_init_timeout: - logger.warning( - "Timeout waiting for device %s initialization lock, proceeding anyway with timeout %s", - device_id, - stage_init_timeout, - ) - break - - # Wait a bit before retrying - _time.sleep(0.1) - except OSError as e: - # Other error - log and continue without lock - logger.debug( - "Failed to acquire lock for device %s: %s, continuing anyway", - device_id, - e, - ) - try: - _os.close(lock_fd) - except (OSError, NameError): - pass - break - - lock_files = acquired_lock_fds - except Exception as e: - logger.debug("Failed to set up sequential initialization lock: %s", e) - - # Init engine based on stage_type - logger.debug( - "[Stage-%s] Initializing %s engine with args keys=%s", - stage_id, - stage_type, - list(engine_args.keys()), - ) - if engine_args.get("async_chunk", False): - logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) - stage_connector_spec = {} - for v in connectors_config.values(): - stage_connector_spec = dict(v.get("spec", {})) - break - engine_args["stage_connector_spec"] = stage_connector_spec - engine_args["stage_id"] = stage_id - try: + if engine_args.get("async_chunk", False): + logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id) + stage_connector_spec = {} + for v in connectors_config.values(): + stage_connector_spec = dict(v.get("spec", {})) + break + engine_args["stage_connector_spec"] = stage_connector_spec + engine_args["stage_id"] = stage_id if stage_type == "diffusion": # For diffusion, we need to extract diffusion-specific config od_config = _build_od_config(engine_args, model) @@ -1261,17 +1117,6 @@ async def _stage_worker_async( usage_context=usage_context, engine_args=omni_engine_args, ) - finally: - # Release all locks by closing file descriptors - # Locks are automatically released when file descriptors are closed - # or when process dies - for lock_fd in lock_files: - try: - fcntl.flock(lock_fd, fcntl.LOCK_UN) - _os.close(lock_fd) - logger.debug("Released initialization lock (fd=%s)", lock_fd) - except (OSError, ValueError): - pass omni_stage.set_async_engine(stage_engine) if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats: diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py new file mode 100644 index 000000000..2e7d07858 --- /dev/null +++ b/vllm_omni/worker/base.py @@ -0,0 +1,102 @@ +"""Base worker class for vLLM-Omni with process-scoped GPU memory accounting.""" + +from __future__ import annotations + +import os + +import torch +from vllm.logger import init_logger +from vllm.utils.mem_utils import format_gib, memory_profiling +from vllm.v1.worker.gpu_worker import Worker as GPUWorker + +from vllm_omni.worker.gpu_memory_utils import ( + get_process_gpu_memory, +) + +logger = init_logger(__name__) + + +class OmniGPUWorkerBase(GPUWorker): + """Base GPU worker for vLLM-Omni with process-scoped memory accounting. + + This class overrides determine_available_memory() to use per-process GPU + memory tracking via pynvml, allowing multiple stages to initialize + concurrently on the same GPU without memory accounting interference. + """ + + @torch.inference_mode() + def determine_available_memory(self) -> int: + """Process-scoped GPU memory profiling for concurrent stage initialization. + + Algorithm: + 1. requested_memory = total_gpu_memory * gpu_memory_utilization + (computed in init_device from cache_config) + + 2. process_memory = memory used by THIS process only (via pynvml) + - Uses nvmlDeviceGetComputeRunningProcesses to get per-PID memory + - Supports CUDA_VISIBLE_DEVICES with indices, UUIDs, or MIG IDs + + 3. available_kv_cache = requested_memory - process_memory + + Fallback: + If NVML is unavailable, falls back to profiling data: + available = requested - (weights + activations + non_torch) + """ + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes: + self.model_runner.profile_run() + logger.info( + "Using explicit kv_cache_memory_bytes: %s GiB", + format_gib(kv_cache_memory_bytes), + ) + return kv_cache_memory_bytes + + with memory_profiling( + self.init_snapshot, + weights_memory=int(self.model_runner.model_memory_usage), + ) as profile_result: + self.model_runner.profile_run() + + self.non_torch_memory = profile_result.non_torch_increase + self.peak_activation_memory = profile_result.torch_peak_increase + + process_memory = get_process_gpu_memory(self.local_rank) + + if process_memory is not None: + # NVML available: use per-process memory + self.available_kv_cache_memory_bytes = max(0, self.requested_memory - process_memory) + logger.debug( + "Process-scoped memory (PID %d, GPU %d): requested=%s, used=%s, available=%s", + os.getpid(), + self.local_rank, + format_gib(self.requested_memory), + format_gib(process_memory), + format_gib(self.available_kv_cache_memory_bytes), + ) + logger.info_once( + "Available KV cache memory: %s GiB (process-scoped)", + format_gib(self.available_kv_cache_memory_bytes), + scope="local", + ) + else: + # NVML unavailable: use profiling data as conservative fallback + profiled_usage = ( + int(self.model_runner.model_memory_usage) + + profile_result.torch_peak_increase + + profile_result.non_torch_increase + ) + self.available_kv_cache_memory_bytes = max(0, self.requested_memory - profiled_usage) + logger.debug( + "Profiling fallback (PID %d, GPU %d): requested=%s, profiled=%s, available=%s", + os.getpid(), + self.local_rank, + format_gib(self.requested_memory), + format_gib(profiled_usage), + format_gib(self.available_kv_cache_memory_bytes), + ) + logger.info_once( + "Available KV cache memory: %s GiB (profiling fallback)", + format_gib(self.available_kv_cache_memory_bytes), + scope="local", + ) + + return int(self.available_kv_cache_memory_bytes) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 98b5afdda..30edba643 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -7,18 +7,18 @@ from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUARWorker(OmniWorkerMixin, GPUWorker): +class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU worker for autoregressive omni model stages. Extends the base GPUWorker to initialize and manage autoregressive diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 91f90d6b4..8ba4efa3d 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -7,18 +7,18 @@ from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats -from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) -class GPUGenerationWorker(OmniWorkerMixin, GPUWorker): +class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU Worker for Generation model (non-autoregressive waveform generation). Usage in stage config: diff --git a/vllm_omni/worker/gpu_memory_utils.py b/vllm_omni/worker/gpu_memory_utils.py new file mode 100644 index 000000000..d005c5e58 --- /dev/null +++ b/vllm_omni/worker/gpu_memory_utils.py @@ -0,0 +1,124 @@ +"""NVML-based per-process GPU memory utilities. + +Shared across worker types (OmniGPUWorkerBase, DiffusionWorker, etc.) +for process-scoped GPU memory accounting. +""" + +from __future__ import annotations + +import os + +from vllm.logger import init_logger +from vllm.third_party.pynvml import ( + nvmlDeviceGetComputeRunningProcesses, + nvmlDeviceGetHandleByIndex, + nvmlInit, + nvmlShutdown, +) + +logger = init_logger(__name__) + + +def is_process_scoped_memory_available() -> bool: + """Check if NVML process-scoped memory tracking is available. + + When True, concurrent stage initialization is safe because each + process can accurately measure its own GPU memory via NVML. + When False, sequential initialization (file locks) is needed. + """ + try: + nvmlInit() + nvmlShutdown() + return True + except Exception: + return False + + +def parse_cuda_visible_devices() -> list[str | int]: + """Parse CUDA_VISIBLE_DEVICES into a list of device identifiers. + + Returns list of integers (physical indices) or strings (UUIDs/MIG IDs). + """ + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible_devices: + return [] + + result: list[str | int] = [] + for item in visible_devices.split(","): + item = item.strip() + if not item: + continue + try: + result.append(int(item)) + except ValueError: + # UUID (GPU-xxx) or MIG ID (MIG-xxx) + result.append(item) + return result + + +def get_device_handle(device_id: str | int): + """Get NVML device handle by index or UUID.""" + if isinstance(device_id, int): + return nvmlDeviceGetHandleByIndex(device_id) + else: + from vllm.third_party.pynvml import nvmlDeviceGetHandleByUUID + + return nvmlDeviceGetHandleByUUID(device_id) + + +def get_process_gpu_memory(local_rank: int) -> int | None: + """Get GPU memory used by current process via pynvml. + + Supports CUDA_VISIBLE_DEVICES with integer indices, UUIDs, or MIG IDs. + + Returns: + Memory in bytes used by this process, or None if NVML unavailable. + + Raises: + RuntimeError: If device validation fails (invalid index or UUID). + """ + from vllm.third_party.pynvml import nvmlDeviceGetCount + + my_pid = os.getpid() + visible_devices = parse_cuda_visible_devices() + + try: + nvmlInit() + except Exception as e: + logger.warning("NVML init failed, will use profiling fallback: %s", e) + return None + + try: + if visible_devices and local_rank < len(visible_devices): + device_id = visible_devices[local_rank] + try: + handle = get_device_handle(device_id) + except Exception as e: + raise RuntimeError( + f"Failed to get NVML handle for device '{device_id}' (local_rank={local_rank}). " + f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + ) from e + else: + # No CUDA_VISIBLE_DEVICES or local_rank out of range: use index directly + device_count = nvmlDeviceGetCount() + if local_rank >= device_count: + raise RuntimeError( + f"Invalid GPU device {local_rank}. Only {device_count} GPU(s) available. " + f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + ) + handle = nvmlDeviceGetHandleByIndex(local_rank) + + for proc in nvmlDeviceGetComputeRunningProcesses(handle): + if proc.pid == my_pid: + return proc.usedGpuMemory + return 0 + except RuntimeError: + raise + except Exception as e: + logger.warning("NVML query failed, will use profiling fallback: %s", e) + return None + finally: + try: + nvmlShutdown() + except Exception: + pass