[cherry-pick releasev0.18.0post1] Codex/release v0.18.0.post1 cherry-pick #2735/#2736/#2555 (#2813)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: Samit <285365963@qq.com>

tests/diffusion/test_diffusion_step_pipeline.py

@@ -443,14 +443,42 @@ class TestWorker:
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
         tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        log_probs = torch.ones(300_000, dtype=torch.float32)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor, trajectory_log_probs=log_probs),
+        )
 
         packed = pack_diffusion_output_shm(output)
         assert isinstance(packed.result.output, dict)
         assert packed.result.output["__tensor_shm__"] is True
+        assert isinstance(packed.result.trajectory_log_probs, dict)
+        assert packed.result.trajectory_log_probs["__tensor_shm__"] is True
 
         unpacked = unpack_diffusion_output_shm(packed)
         assert isinstance(unpacked.result.output, torch.Tensor)
+        assert isinstance(unpacked.result.trajectory_log_probs, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+        torch.testing.assert_close(unpacked.result.trajectory_log_probs, log_probs)
+
+    def test_pack_unpack_runner_output_shm_bfloat16(self):
+        tensor = torch.arange(600_000, dtype=torch.bfloat16)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor),
+        )
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+        assert packed.result.output["torch_dtype"] == str(torch.bfloat16)
+        assert packed.result.output["numpy_dtype"] == "uint16"
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        assert unpacked.result.output.dtype == torch.bfloat16
         torch.testing.assert_close(unpacked.result.output, tensor)
 
 

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_engine():
+    with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
+        engine_instance = MagicMock()
+        mock.make_engine.return_value = engine_instance
+        yield engine_instance
+
+
+@pytest.fixture
+def client(mock_engine):
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_output_type="image",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source="prompt",
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        yield c
+        c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_success(client, mock_engine):
+    mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
+    mock_engine.step.return_value = [mock_result]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "A test prompt", sampling_params)
+
+    # Wait for the task to be processed
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-1"
+    mock_engine.step.assert_called_once()
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == ["A test prompt"]
+    assert request.request_ids == ["req-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_batch_success(client, mock_engine):
+    mock_engine.step.return_value = [
+        OmniRequestOutput.from_diffusion(request_id="req-batch-0", images=[MagicMock()]),
+        OmniRequestOutput.from_diffusion(request_id="req-batch-1", images=[MagicMock()]),
+    ]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    prompts = ["prompt-0", "prompt-1"]
+    await client.add_batch_request_async("req-batch", prompts, sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-batch"
+    assert len(output.images) == 2
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == prompts
+    assert request.request_ids == ["req-batch-0", "req-batch-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_error(client, mock_engine):
+    mock_engine.step.side_effect = RuntimeError("Engine failure")
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-err", "A test prompt", sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output["request_id"] == "req-err"
+    assert output["error"] == "Engine failure"
+    assert output["error_type"] == "RuntimeError"
+
+
+def test_inline_shutdown(client, mock_engine):
+    assert not client._shutting_down
+
+    # Shutting down should cleanly cancel anything queued and close engine
+    client.shutdown()
+
+    assert client._shutting_down
+    mock_engine.close.assert_called_once()