vllm_omni.worker.base ¶

Base worker class for vLLM-Omni with device-level GPU memory profiling.

logger `module-attribute` ¶

logger = init_logger(__name__)

OmniGPUWorkerBase ¶

Bases: Worker

Base GPU worker for vLLM-Omni.

This class overrides determine_available_memory() to compute the KV cache budget from device-level profiling data. Concurrent same-device stage initialization is coordinated externally by the parallel-stage-init admission control + SH/EX device locks (see VllmOmniOrchestratorConfig.parallel_stage_init), which keep each measurement quiescent instead of relying on per-process NVML accounting.

It also replaces vLLM's TorchProfilerWrapper with OmniTorchProfilerWrapper for custom trace naming, background gzip, and trace path collection.

profiler `instance-attribute` ¶

profiler = create_omni_profiler(
    profiler_config=profiler_config,
    worker_name=worker_name,
    local_rank=self.local_rank,
)

determine_available_memory ¶

determine_available_memory() -> int

Device-level GPU memory profiling for the KV cache budget.

Algorithm

requested_memory = total_gpu_memory * gpu_memory_utilization (computed in init_device from cache_config)
profiled_usage = weights + peak_activation + non_torch_increase (measured by memory_profiling around profile_run(); non_torch_increase is device-level, so it reflects whatever else is resident on the GPU at profiling time)
available_kv_cache = requested_memory - profiled_usage

Note

Process-scoped NVML estimation was removed in favour of the device-level profiling fallback plus the parallel-stage-init admission + SH/EX device locks (see VllmOmniOrchestratorConfig.parallel_stage_init): the EX lock makes each measurement quiescent, so the device-level number is the correct, conservative budget. The NVML helpers in gpu_memory_utils are retained for diffusion memory reporting.

encoder_loaded ¶

encoder_loaded() -> bool

Check if encoder weights are loaded in the model.

This method is exposed via collective_rpc to check encoder availability for models that support voice cloning with reference audio.

Returns:

Name	Type	Description
`bool`	`bool`	True if encoder weights are available, False otherwise.

handle_sleep_task ¶

handle_sleep_task(task: OmniSleepTask) -> OmniACK

Handle deterministic Sleep command from the main process

handle_wake_task ¶

handle_wake_task(task: OmniWakeTask) -> OmniACK

Handle deterministic Wakeup command from the main process

load_model ¶

load_model(*args, **kwargs)

profile ¶

profile(
    is_start: bool = True, profile_prefix: str | None = None
)

Override to set trace filename before starting the profiler.

Parameters:

Name	Type	Description	Default
`is_start`	`bool`	True to start profiling, False to stop.	`True`
`profile_prefix`	`str \| None`	Optional prefix for trace filename (vLLM compat).	`None`

vLLM's profile() only passes is_start, so we generate a descriptive trace filename here before delegating to the profiler.

sleep ¶

sleep(level: int = 1) -> bool

Put the worker to sleep. Args: level: 1 (Offload weights to CPU), level: 2 (Total Discard).

wake_up ¶

wake_up(tags: list[str] | None = None) -> bool

Physical video memory reloading logic