[Log] Wire stat loggers into AsyncOmniEngine to match AsyncLLM (#2551)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

tests/engine/test_async_omni_engine_do_log_stats.py

@@ -0,0 +1,56 @@
+"""Guard tests for AsyncOmniEngine.do_log_stats edge cases.
+
+These are pure-Python tests that bypass __init__ and only exercise the
+no-op branches of do_log_stats, so no stage cores / threads are needed.
+"""
+
+import asyncio
+
+import pytest
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_bare_engine() -> AsyncOmniEngine:
+    # Bypass __init__ so we don't spin up stage cores; we only need the
+    # attributes do_log_stats touches.
+    return AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+
+@pytest.mark.asyncio
+async def test_do_log_stats_noop_when_manager_missing():
+    engine = _make_bare_engine()
+    engine.logger_manager = None
+    engine.orchestrator_loop = None
+    await engine.do_log_stats()  # should silently return
+
+
+@pytest.mark.asyncio
+async def test_do_log_stats_noop_when_loop_missing():
+    engine = _make_bare_engine()
+
+    class _Manager:
+        def log(self) -> None:  # pragma: no cover - must not be called
+            raise AssertionError("log() should not be called without a loop")
+
+    engine.logger_manager = _Manager()
+    engine.orchestrator_loop = None
+    await engine.do_log_stats()
+
+
+@pytest.mark.asyncio
+async def test_do_log_stats_noop_when_loop_not_running():
+    engine = _make_bare_engine()
+
+    class _Manager:
+        def log(self) -> None:  # pragma: no cover - must not be called
+            raise AssertionError("log() should not be called on a stopped loop")
+
+    dead_loop = asyncio.new_event_loop()
+    dead_loop.close()
+
+    engine.logger_manager = _Manager()
+    engine.orchestrator_loop = dead_loop
+    await engine.do_log_stats()

tests/engine/test_async_omni_engine_stage_init.py

@@ -31,6 +31,7 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
     from vllm_omni.platforms import current_omni_platform
 
     engine = object.__new__(AsyncOmniEngine)
+    engine.log_stats = False
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
     engine.num_stages = 1
@@ -280,6 +281,7 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
     )
 
     engine = object.__new__(AsyncOmniEngine)
+    engine.log_stats = False
 
     _stage_client, _out_proc, _vllm_cfg, input_processor = engine._attach_llm_stage(started)
 

tests/engine/test_single_stage_mode.py

@@ -461,6 +461,7 @@ class TestInitializeStagesRouting:
         engine.stage_configs = stage_cfgs
         engine.num_stages = len(stage_cfgs)
         engine.async_chunk = False
+        engine.log_stats = False
         engine.single_stage_mode = single_stage_mode
         engine._single_stage_id_filter = stage_id_filter
         engine._omni_master_address = omni_master_address
@@ -1366,6 +1367,7 @@ class TestLaunchLlmStageSingleStageMode:
     def _build_engine_with_oms(self) -> AsyncOmniEngine:
         engine = object.__new__(AsyncOmniEngine)
         engine.model = "fake-model"
+        engine.log_stats = False
         engine.single_stage_mode = True
         engine._single_stage_id_filter = 0
         engine._llm_stage_launch_lock = threading.Lock()
@@ -1446,6 +1448,7 @@ class TestLaunchLlmStageSingleStageMode:
         """~single_stage_mode → spawn_stage_core + complete_stage_handshake."""
         engine = object.__new__(AsyncOmniEngine)
         engine.model = "fake-model"
+        engine.log_stats = False
         engine.single_stage_mode = False
         engine._omni_master_server = None
         engine._llm_stage_launch_lock = threading.Lock()