Integrate BackgroundResources into Omni and AsyncOmni

Signed-off-by: Daniel Huang <daniel1.huang@intel.com>

Author: pi314ever

vllm_omni/entrypoints/async_omni.py

@@ -3,7 +3,6 @@
 import asyncio
 import copy
 import time
-import weakref
 from collections.abc import AsyncGenerator, Iterable, Sequence
 from typing import Any
 
@@ -18,12 +17,11 @@
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length, try_send_via_connector
-from vllm_omni.distributed.ray_utils.utils import try_close_ray
 from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
+from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
@@ -38,34 +36,6 @@
 logger = init_logger(__name__)
 
 
-def _weak_close_cleanup_async(stage_list, stage_in_queues, stage_out_queues, ray_pg, output_handler, zmq_ctx=None):
-    """Weak reference cleanup function for AsyncOmni instances."""
-    if stage_list:
-        for q in stage_in_queues:
-            try:
-                q.put_nowait(SHUTDOWN_TASK)
-            except Exception as e:
-                logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
-            close_fn = getattr(q, "close", None)
-            if callable(close_fn):
-                close_fn()
-        for q in stage_out_queues:
-            close_fn = getattr(q, "close", None)
-            if callable(close_fn):
-                close_fn()
-        for stage in stage_list:
-            try:
-                stage.stop_stage_worker()
-            except Exception as e:
-                logger.warning(f"Failed to stop stage worker: {e}")
-    try_close_ray(ray_pg)
-    # Cancel output handler
-    if output_handler is not None:
-        output_handler.cancel()
-    if zmq_ctx is not None:
-        zmq_ctx.term()
-
-
 class AsyncOmni(OmniBase):
     """Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models.
 
@@ -107,22 +77,9 @@ def __init__(self, model: str, **kwargs: dict[str, Any]) -> None:
 
         # Request state tracking
         self.request_states: dict[str, ClientRequestState] = {}
-        self.output_handler: asyncio.Task | None = None
 
         super().__init__(model, **kwargs)
 
-        # Register weak reference cleanup (called on garbage collection)
-        self._weak_finalizer = weakref.finalize(
-            self,
-            _weak_close_cleanup_async,
-            self.stage_list,
-            self._stage_in_queues,
-            self._stage_out_queues,
-            self._ray_pg,
-            self.output_handler,
-            self._zmq_ctx,
-        )
-
     def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
         """Create default diffusion stage configuration."""
         # TODO: here is different from the Omni class. We should merge the two in the future.
@@ -216,7 +173,7 @@ def _process_stage_ready(self, stage: OmniStage, stage_id: int, result: dict[str
     def _wait_for_stages_ready(self, timeout: int = 120) -> None:
         """Wait for all stages to report readiness."""
         super()._wait_for_stages_ready(timeout)
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             if stage.vllm_config is not None and stage.tokenizer is not None:
                 try:
                     vllm_config = stage.vllm_config
@@ -305,11 +262,13 @@ async def generate(
             if sampling_params_list is None:
                 sampling_params_list = self.default_sampling_params_list
 
-            if len(sampling_params_list) != len(self.stage_list):
-                raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
+            if len(sampling_params_list) != len(self.resources.stage_list):
+                raise ValueError(
+                    f"Expected {len(self.resources.stage_list)} sampling params, got {len(sampling_params_list)}"
+                )
 
             # Orchestrator keeps stage objects for input derivation
-            num_stages = len(self.stage_list)
+            num_stages = len(self.resources.stage_list)
             # Track per-request start time for end-to-end timing
             _req_start_ts: dict[int, float] = {}
             _wall_start_ts: float = time.time()
@@ -318,7 +277,7 @@ async def generate(
             # Determine the final stage for E2E stats (highest stage_id with
             # final_output=True; fallback to last stage)
             final_stage_id_for_e2e = get_final_stage_id_for_e2e(
-                output_modalities, self.output_modalities, self.stage_list
+                output_modalities, self.output_modalities, self.resources.stage_list
             )
 
             # Metrics/aggregation helper
@@ -337,7 +296,7 @@ async def generate(
                 "engine_inputs": prompt,
                 "sampling_params": sp0,
             }
-            self.stage_list[0].submit(task)
+            self.resources.stage_list[0].submit(task)
             metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
             _req_start_ts[request_id] = time.time()
             logger.info(
@@ -399,7 +358,7 @@ async def _process_async_results(
         all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
         submit_flag = True
         while not all(all_stages_finished.values()):
-            for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
+            for stage_id, stage in enumerate(self.resources.stage_list[: final_stage_id_for_e2e + 1]):
                 if all_stages_finished[stage_id]:
                     continue
                 try:
@@ -420,13 +379,13 @@ async def _process_async_results(
                     next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
                     engine_input["prompt_token_ids"] = [0] * next_prompt_len
                     engine_input["multi_modal_data"] = engine_input["mm_processor_kwargs"] = None
-                    for i in range(1, len(self.stage_list)):
+                    for i in range(1, len(self.resources.stage_list)):
                         task = {
                             "request_id": request_id,
                             "engine_inputs": engine_input,
                             "sampling_params": sampling_params_list[i],
                         }
-                        self.stage_list[i].submit(task)
+                        self.resources.stage_list[i].submit(task)
                         metrics.stage_first_ts[i] = time.time()
                 all_stages_finished[stage_id] = finished
 
@@ -461,10 +420,10 @@ async def _process_sequential_results(
             # Forward to next stage if there is one
             next_stage_id = stage_id + 1
             if next_stage_id <= final_stage_id_for_e2e:
-                next_stage: OmniStage = self.stage_list[next_stage_id]
+                next_stage: OmniStage = self.resources.stage_list[next_stage_id]
                 # Derive inputs for the next stage, record postprocess time
                 with metrics.stage_postprocess_timer(stage_id, request_id):
-                    next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
+                    next_inputs = next_stage.process_engine_inputs(self.resources.stage_list, prompt)
                 sp_next: SamplingParams = sampling_params_list[next_stage_id]
 
                 # Check if we have a connector for this edge
@@ -481,7 +440,7 @@ async def _process_sequential_results(
                         next_inputs=next_inputs,
                         sampling_params=sp_next,
                         original_prompt=prompt,
-                        next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                        next_stage_queue_submit_fn=self.resources.stage_list[next_stage_id].submit,
                         metrics=metrics,
                     )
 
@@ -574,10 +533,10 @@ def _process_single_result(
         return engine_outputs, finished, output_to_yield
 
     def _run_output_handler(self) -> None:
-        if self.output_handler is not None:
+        if self.resources.output_handler is not None:
             return
 
-        stage_list = self.stage_list
+        stage_list = self.resources.stage_list
         request_states = self.request_states
 
         async def output_handler():
@@ -623,14 +582,14 @@ async def output_handler():
                     else:
                         await req_state.queue.put(error_msg)
                     error_msg = {"request_id": req_state.request_id, "error": str(e)}
-                self.output_handler = None  # Make possible for restart
+                self.resources.output_handler = None  # Make possible for restart
 
-        self.output_handler = asyncio.create_task(output_handler())
+        self.resources.output_handler = asyncio.create_task(output_handler())
 
     @property
     def is_running(self) -> bool:
         # Is None before the loop is started.
-        return len(self._stage_in_queues) > 0
+        return len(self.resources._stage_in_queues) > 0
 
     @property
     def is_stopped(self) -> bool:
@@ -654,20 +613,20 @@ def dead_error(self) -> BaseException:
 
     async def abort(self, request_id: str | Iterable[str]) -> None:
         abort_task = {"type": OmniStageTaskType.ABORT, "request_id": request_id}
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             stage.submit(abort_task)
         return None
 
     async def get_vllm_config(self) -> VllmConfig:
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             if stage.is_comprehension:
                 # Use the vllm_config received from worker process
                 if stage.vllm_config is not None:
                     return stage.vllm_config
         return None
 
     async def get_model_config(self) -> OmniModelConfig:
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             if stage.is_comprehension:
                 # Use the vllm_config received from worker process
                 if stage.vllm_config is not None:
@@ -678,13 +637,13 @@ async def get_input_preprocessor(self) -> InputPreprocessor:
         return None
 
     async def get_tokenizer(self) -> TokenizerLike:
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             if stage.is_comprehension:
                 return stage.tokenizer
         return None
 
     async def is_tracing_enabled(self) -> bool:
-        for stage in self.stage_list:
+        for stage in self.resources.stage_list:
             if stage.is_comprehension:
                 return stage.is_tracing_enabled
         return False