anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Support DROID policy server for Cosmos3 OpenPI (#4282)

Browse Source 
Signed-off-by: yuzhud <yuzhudong@gmail.com>
This commit is contained in:
Yuzhu Dong authored and GitHub committed 2026-06-15 21:22:59 +08:00
1 parent ae937320ff
commit e2e79177db
16 files changed
+1397 -154
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_plugin_hooks.py
+33 -1
View File
@@ -10,11 +10,14 @@ This module tests:
- Worker integration: model runner resolved via platform hook
"""
from unittest.mock import patch
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
from vllm_omni.diffusion.registry import (
_DIFFUSION_ACTION_POST_PROCESS_FUNCS,
_DIFFUSION_IR_OP_PRIORITY_FUNCS,
_DIFFUSION_MODELS,
_DIFFUSION_POST_PROCESS_FUNCS,
_DIFFUSION_PRE_PROCESS_FUNCS,
@@ -59,6 +62,8 @@ class TestRegisterDiffusionModel:
original_models = _DIFFUSION_MODELS.copy()
original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
yield
_DIFFUSION_MODELS.clear()
_DIFFUSION_MODELS.update(original_models)
@@ -66,6 +71,10 @@ class TestRegisterDiffusionModel:
_DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
_DIFFUSION_POST_PROCESS_FUNCS.clear()
_DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
_DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
_DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
_DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
_DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
def test_register_new_model(self):
"""Test registering a new diffusion model with pre/post process functions."""
@@ -75,6 +84,8 @@ class TestRegisterDiffusionModel:
class_name="TestPipeline",
pre_process_func_name="test_pre_process",
post_process_func_name="test_post_process",
action_post_process_func_name="test_action_post_process",
ir_op_priority_func_name="test_ir_op_priority",
)
assert "TestPipeline" in _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -84,6 +95,8 @@ class TestRegisterDiffusionModel:
)
assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
class TestWorkerUsesHook:
@@ -108,3 +121,22 @@ class TestWorkerUsesHook:
assert worker.model_runner is mock_runner_instance
mock_platform.get_diffusion_model_runner_cls.assert_called_once()
mock_resolve.assert_called_once_with("custom.path")
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_diffusion_ir_op_priority_func")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
def test_ir_op_priority_hook_receives_platform_default(self, mock_platform, mock_get_hook):
"""Test model IR priority hook merges from the platform default."""
from vllm_omni.diffusion.worker.diffusion_worker import _resolve_ir_op_priority
od_config = SimpleNamespace(model_class_name="TestPipeline")
vllm_config = SimpleNamespace()
default_priority = object()
merged_priority = object()
hook = Mock(return_value=merged_priority)
mock_platform.get_default_ir_op_priority.return_value = default_priority
mock_get_hook.return_value = hook
assert _resolve_ir_op_priority(od_config, vllm_config) is merged_priority
mock_platform.get_default_ir_op_priority.assert_called_once_with(vllm_config)
mock_get_hook.assert_called_once_with(od_config)
hook.assert_called_once_with(default_priority, vllm_config=vllm_config)