Multi platform Plugin (#21388)
Co-authored-by: root <root@tjzj-inf-sci-k8s-bzz2-0183.tjzj.baidu.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai> Co-authored-by: Alex Nails <alexj.nails@gmail.com> Co-authored-by: root <root@tjzj-inf-sci-k8s-bzz2-0000.tjzj.baidu.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Baidu-AIAK
co-authored by
root
Alex Nails
Alex Nails
root
Mick
root
Alex Nails
Alex Nails
root
Mick
parent
ebcc2b3eec
commit
7ca3566130
@@ -0,0 +1,448 @@
|
||||
"""
|
||||
Unit tests for the hook registry system.
|
||||
|
||||
Covers: basic hooks (AROUND/BEFORE/AFTER/REPLACE), descriptor preservation
|
||||
(classmethod/staticmethod), hook ordering, cross-target conflict detection,
|
||||
patch propagation, and edge cases.
|
||||
|
||||
Run: python -m pytest test/registered/unit/plugins/test_hook_registry.py -v
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
import uuid
|
||||
|
||||
from sglang.srt.plugins.hook_registry import HookRegistry, HookType, plugin_hook
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers: synthetic module creation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SYNTH_MODULE_PREFIX = "_synth_hook_test_"
|
||||
|
||||
|
||||
def _make_module(**attrs):
|
||||
"""Create a throwaway module registered in sys.modules."""
|
||||
name = f"{_SYNTH_MODULE_PREFIX}{uuid.uuid4().hex[:8]}"
|
||||
mod = types.ModuleType(name)
|
||||
for k, v in attrs.items():
|
||||
setattr(mod, k, v)
|
||||
sys.modules[name] = mod
|
||||
return mod, name
|
||||
|
||||
|
||||
def _cleanup_synth_modules():
|
||||
"""Remove all synthetic modules from sys.modules."""
|
||||
to_del = [k for k in sys.modules if k.startswith(_SYNTH_MODULE_PREFIX)]
|
||||
for k in to_del:
|
||||
del sys.modules[k]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Base class for hook tests (shared setUp/tearDown)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _HookTestCase(CustomTestCase):
|
||||
"""Base class that resets HookRegistry and cleans up synth modules."""
|
||||
|
||||
def setUp(self):
|
||||
HookRegistry.reset()
|
||||
_cleanup_synth_modules()
|
||||
|
||||
def tearDown(self):
|
||||
HookRegistry.reset()
|
||||
_cleanup_synth_modules()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestBasicHooks
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestBasicHooks(_HookTestCase):
|
||||
"""AROUND / BEFORE / AFTER / REPLACE on plain functions, class REPLACE,
|
||||
and the @plugin_hook decorator."""
|
||||
|
||||
def test_around_function(self):
|
||||
def orig(x):
|
||||
return x * 2
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def add_one(original_fn, x):
|
||||
return original_fn(x) + 1
|
||||
|
||||
HookRegistry.register(f"{name}.orig", add_one, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 7) # 3*2 + 1
|
||||
|
||||
def test_before_modifies_args(self):
|
||||
"""BEFORE hook returns (args, kwargs) to modify arguments."""
|
||||
|
||||
def orig(x, y=0):
|
||||
return x + y
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def double_x(x, y=0):
|
||||
return (x * 2,), {"y": y + 1}
|
||||
|
||||
HookRegistry.register(f"{name}.orig", double_x, HookType.BEFORE)
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 7) # x=3*2=6, y=0+1=1, 6+1=7
|
||||
|
||||
def test_before_returning_none(self):
|
||||
"""BEFORE hook returning None leaves arguments unchanged."""
|
||||
|
||||
def orig(x):
|
||||
return x * 2
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def before_noop(x):
|
||||
return None # leave args unchanged
|
||||
|
||||
HookRegistry.register(f"{name}.orig", before_noop, HookType.BEFORE)
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 6) # args unchanged
|
||||
|
||||
def test_after_function(self):
|
||||
def orig(x):
|
||||
return x * 2
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def add_ten(result, x):
|
||||
return result + 10
|
||||
|
||||
HookRegistry.register(f"{name}.orig", add_ten, HookType.AFTER)
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 16) # 3*2 + 10
|
||||
|
||||
def test_replace_function(self):
|
||||
def orig(x):
|
||||
return x * 2
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def replacement(x):
|
||||
return x * 100
|
||||
|
||||
HookRegistry.register(f"{name}.orig", replacement, HookType.REPLACE)
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 300)
|
||||
|
||||
def test_class_replace(self):
|
||||
class Original:
|
||||
def greet(self):
|
||||
return "original"
|
||||
|
||||
mod, name = _make_module(Original=Original)
|
||||
|
||||
class Replacement(Original):
|
||||
def greet(self):
|
||||
return "replaced"
|
||||
|
||||
HookRegistry.register(f"{name}.Original", Replacement, HookType.REPLACE)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
self.assertIs(mod.Original, Replacement)
|
||||
self.assertIsInstance(mod.Original(), Replacement)
|
||||
self.assertEqual(mod.Original().greet(), "replaced")
|
||||
|
||||
def test_plugin_hook_decorator(self):
|
||||
def orig(x):
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
@plugin_hook(f"{name}.orig", type=HookType.REPLACE)
|
||||
def my_replace(x):
|
||||
return x + 42
|
||||
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(0), 42)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestDescriptorPreservation (Bug B regression tests)
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestDescriptorPreservation(_HookTestCase):
|
||||
"""Hooks on classmethod/staticmethod must preserve descriptor semantics."""
|
||||
|
||||
def _make_cls_module(self):
|
||||
class MyClass:
|
||||
@classmethod
|
||||
def cm(cls, x):
|
||||
return ("cm", cls.__name__, x)
|
||||
|
||||
@staticmethod
|
||||
def sm(x):
|
||||
return ("sm", x)
|
||||
|
||||
mod, name = _make_module(MyClass=MyClass)
|
||||
return mod, name, MyClass
|
||||
|
||||
def test_around_classmethod(self):
|
||||
mod, name, MyClass = self._make_cls_module()
|
||||
|
||||
def add_tag(original_fn, cls, x):
|
||||
return original_fn(cls, x) + ("around",)
|
||||
|
||||
HookRegistry.register(f"{name}.MyClass.cm", add_tag, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
result = mod.MyClass.cm(1)
|
||||
self.assertEqual(result, ("cm", "MyClass", 1, "around"))
|
||||
|
||||
def test_replace_classmethod(self):
|
||||
mod, name, MyClass = self._make_cls_module()
|
||||
|
||||
def new_cm(cls, x):
|
||||
return ("replaced_cm", cls.__name__, x)
|
||||
|
||||
HookRegistry.register(f"{name}.MyClass.cm", new_cm, HookType.REPLACE)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
result = mod.MyClass.cm(1)
|
||||
self.assertEqual(result, ("replaced_cm", "MyClass", 1))
|
||||
|
||||
def test_around_staticmethod(self):
|
||||
mod, name, MyClass = self._make_cls_module()
|
||||
|
||||
def wrap_sm(original_fn, x):
|
||||
return original_fn(x) + ("around",)
|
||||
|
||||
HookRegistry.register(f"{name}.MyClass.sm", wrap_sm, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
result = mod.MyClass.sm(1)
|
||||
self.assertEqual(result, ("sm", 1, "around"))
|
||||
|
||||
def test_replace_staticmethod(self):
|
||||
mod, name, MyClass = self._make_cls_module()
|
||||
|
||||
def new_sm(x):
|
||||
return ("replaced_sm", x)
|
||||
|
||||
HookRegistry.register(f"{name}.MyClass.sm", new_sm, HookType.REPLACE)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
result = mod.MyClass.sm(1)
|
||||
self.assertEqual(result, ("replaced_sm", 1))
|
||||
|
||||
def test_classmethod_subclass_cls(self):
|
||||
mod, name, MyClass = self._make_cls_module()
|
||||
|
||||
def add_tag(original_fn, cls, x):
|
||||
return original_fn(cls, x) + ("around",)
|
||||
|
||||
HookRegistry.register(f"{name}.MyClass.cm", add_tag, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
class Sub(mod.MyClass):
|
||||
pass
|
||||
|
||||
result = Sub.cm(1)
|
||||
self.assertEqual(result, ("cm", "Sub", 1, "around"))
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestHookOrdering
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestHookOrdering(_HookTestCase):
|
||||
"""Verify REPLACE is applied first, then other hooks wrap it."""
|
||||
|
||||
def test_replace_then_around(self):
|
||||
def orig(x):
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def repl(x):
|
||||
return x * 10
|
||||
|
||||
def add_one(original_fn, x):
|
||||
return original_fn(x) + 1
|
||||
|
||||
HookRegistry.register(f"{name}.orig", repl, HookType.REPLACE)
|
||||
HookRegistry.register(f"{name}.orig", add_one, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
# REPLACE first: x*10, then AROUND: +1 => 31
|
||||
self.assertEqual(mod.orig(3), 31)
|
||||
|
||||
def test_replace_before_after(self):
|
||||
def orig(x):
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def repl(x):
|
||||
return x * 10
|
||||
|
||||
def double_arg(x):
|
||||
return (x * 2,), {}
|
||||
|
||||
def add_hundred(result, x):
|
||||
return result + 100
|
||||
|
||||
HookRegistry.register(f"{name}.orig", repl, HookType.REPLACE)
|
||||
HookRegistry.register(f"{name}.orig", double_arg, HookType.BEFORE)
|
||||
HookRegistry.register(f"{name}.orig", add_hundred, HookType.AFTER)
|
||||
HookRegistry.apply_hooks()
|
||||
# BEFORE doubles x: 3*2=6 → REPLACE: 6*10=60 → AFTER: 60+100=160
|
||||
self.assertEqual(mod.orig(3), 160)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestCrossTargetConflict
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestCrossTargetConflict(_HookTestCase):
|
||||
"""Verify warning for class REPLACE + method REPLACE combo."""
|
||||
|
||||
def test_class_replace_then_method_replace_warns(self):
|
||||
class Original:
|
||||
def foo(self):
|
||||
return "orig"
|
||||
|
||||
mod, name = _make_module(Original=Original)
|
||||
|
||||
class Replacement(Original):
|
||||
def foo(self):
|
||||
return "class_replaced"
|
||||
|
||||
HookRegistry.register(f"{name}.Original", Replacement, HookType.REPLACE)
|
||||
|
||||
def method_repl(self):
|
||||
return "method_replaced"
|
||||
|
||||
HookRegistry.register(f"{name}.Original.foo", method_repl, HookType.REPLACE)
|
||||
|
||||
with self.assertLogs("sglang.srt.plugins.hook_registry", level="WARNING") as cm:
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
self.assertTrue(any("will override" in msg for msg in cm.output))
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestPatchPropagation
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestPatchPropagation(_HookTestCase):
|
||||
"""Verify that patches propagate to other modules that imported the target."""
|
||||
|
||||
def test_same_reference_propagates(self):
|
||||
def orig(x):
|
||||
return x * 2
|
||||
|
||||
source_mod, source_name = _make_module(orig=orig)
|
||||
importer_mod, _ = _make_module(orig=orig) # same reference
|
||||
|
||||
def add_one(fn, x):
|
||||
return fn(x) + 1
|
||||
|
||||
HookRegistry.register(f"{source_name}.orig", add_one, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
self.assertEqual(source_mod.orig(3), 7)
|
||||
self.assertEqual(importer_mod.orig(3), 7)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestEdgeCases
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestEdgeCases(_HookTestCase):
|
||||
"""Reset, type validation, multi-AROUND onion, idempotent apply."""
|
||||
|
||||
def test_reset(self):
|
||||
def orig(x):
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def noop(fn, x):
|
||||
return fn(x)
|
||||
|
||||
HookRegistry.register(f"{name}.orig", noop, HookType.AROUND)
|
||||
HookRegistry.reset()
|
||||
|
||||
HookRegistry.apply_hooks()
|
||||
self.assertEqual(mod.orig(3), 3)
|
||||
|
||||
def test_register_class_with_wrong_type(self):
|
||||
class BadHook:
|
||||
pass
|
||||
|
||||
for ht in (HookType.BEFORE, HookType.AFTER, HookType.AROUND):
|
||||
with self.assertRaises(TypeError):
|
||||
HookRegistry.register("some.target", BadHook, ht)
|
||||
|
||||
def test_multi_around_onion(self):
|
||||
call_order = []
|
||||
|
||||
def orig(x):
|
||||
call_order.append("orig")
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def around1(fn, x):
|
||||
call_order.append("a1_before")
|
||||
result = fn(x)
|
||||
call_order.append("a1_after")
|
||||
return result + 1
|
||||
|
||||
def around2(fn, x):
|
||||
call_order.append("a2_before")
|
||||
result = fn(x)
|
||||
call_order.append("a2_after")
|
||||
return result + 10
|
||||
|
||||
HookRegistry.register(f"{name}.orig", around1, HookType.AROUND)
|
||||
HookRegistry.register(f"{name}.orig", around2, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
|
||||
result = mod.orig(0)
|
||||
self.assertEqual(result, 11)
|
||||
self.assertEqual(
|
||||
call_order, ["a2_before", "a1_before", "orig", "a1_after", "a2_after"]
|
||||
)
|
||||
|
||||
def test_apply_idempotent(self):
|
||||
call_count = [0]
|
||||
|
||||
def orig(x):
|
||||
return x
|
||||
|
||||
mod, name = _make_module(orig=orig)
|
||||
|
||||
def counter(fn, x):
|
||||
call_count[0] += 1
|
||||
return fn(x)
|
||||
|
||||
HookRegistry.register(f"{name}.orig", counter, HookType.AROUND)
|
||||
HookRegistry.apply_hooks()
|
||||
HookRegistry.apply_hooks() # second apply should be no-op
|
||||
|
||||
mod.orig(1)
|
||||
self.assertEqual(call_count[0], 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
@@ -0,0 +1,187 @@
|
||||
"""
|
||||
Unit tests for the plugin loading flow.
|
||||
|
||||
Covers: idempotency, apply_hooks invocation, exception resilience,
|
||||
SGLANG_PLUGINS whitelist, SGLANG_PLATFORM exclusion logic,
|
||||
and _current_plugin_source context var reset.
|
||||
|
||||
Run: python -m pytest test/registered/unit/plugins/test_load_plugins.py -v
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.plugins import (
|
||||
_current_plugin_source,
|
||||
_get_excluded_dists,
|
||||
load_plugins,
|
||||
load_plugins_by_group,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
def _make_ep(name, dist_name=None, load_fn=None):
|
||||
"""Create a mock entry point."""
|
||||
ep = MagicMock()
|
||||
ep.name = name
|
||||
ep.value = f"fake_module:{name}"
|
||||
ep.dist = MagicMock()
|
||||
ep.dist.name = dist_name or f"{name}-dist"
|
||||
if load_fn is not None:
|
||||
ep.load.return_value = load_fn
|
||||
else:
|
||||
ep.load.return_value = MagicMock()
|
||||
return ep
|
||||
|
||||
|
||||
def _reset_plugins_loaded():
|
||||
"""Reset the _plugins_loaded flag so load_plugins() can run again."""
|
||||
import sglang.srt.plugins as plugins_mod
|
||||
|
||||
plugins_mod._plugins_loaded = False
|
||||
|
||||
|
||||
class TestLoadPlugins(CustomTestCase):
|
||||
"""Tests for load_plugins() and related helpers."""
|
||||
|
||||
def setUp(self):
|
||||
_reset_plugins_loaded()
|
||||
|
||||
def tearDown(self):
|
||||
_reset_plugins_loaded()
|
||||
|
||||
@patch("sglang.srt.plugins.HookRegistry")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
@patch("sglang.srt.plugins.entry_points", return_value=[])
|
||||
def test_load_plugins_idempotent_and_calls_apply(
|
||||
self, mock_eps, mock_envs, mock_registry
|
||||
):
|
||||
"""Second call is a no-op; first call invokes apply_hooks."""
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
mock_envs.SGLANG_PLUGINS.get.return_value = ""
|
||||
|
||||
load_plugins()
|
||||
self.assertEqual(mock_registry.apply_hooks.call_count, 1)
|
||||
|
||||
load_plugins() # should be skipped
|
||||
self.assertEqual(mock_registry.apply_hooks.call_count, 1)
|
||||
|
||||
@patch("sglang.srt.plugins.HookRegistry")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
@patch("sglang.srt.plugins.entry_points")
|
||||
def test_plugin_exception_does_not_crash(self, mock_eps, mock_envs, mock_registry):
|
||||
"""A failing plugin should not prevent others from loading."""
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
mock_envs.SGLANG_PLUGINS.get.return_value = ""
|
||||
|
||||
def bad_plugin():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
good_call_log = []
|
||||
|
||||
def good_plugin():
|
||||
good_call_log.append("ok")
|
||||
|
||||
eps = [
|
||||
_make_ep("bad", load_fn=bad_plugin),
|
||||
_make_ep("good", load_fn=good_plugin),
|
||||
]
|
||||
mock_eps.return_value = eps
|
||||
|
||||
with self.assertLogs("sglang.srt.plugins", level="ERROR") as cm:
|
||||
load_plugins()
|
||||
|
||||
self.assertTrue(any("boom" in msg for msg in cm.output))
|
||||
self.assertEqual(good_call_log, ["ok"])
|
||||
mock_registry.apply_hooks.assert_called_once()
|
||||
|
||||
@patch("sglang.srt.plugins.entry_points")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
def test_sglang_plugins_whitelist(self, mock_envs, mock_eps):
|
||||
"""Only plugins named in SGLANG_PLUGINS should be loaded."""
|
||||
mock_envs.SGLANG_PLUGINS.get.return_value = "alpha,gamma"
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
|
||||
alpha_fn = MagicMock()
|
||||
beta_fn = MagicMock()
|
||||
gamma_fn = MagicMock()
|
||||
|
||||
eps = [
|
||||
_make_ep("alpha", load_fn=alpha_fn),
|
||||
_make_ep("beta", load_fn=beta_fn),
|
||||
_make_ep("gamma", load_fn=gamma_fn),
|
||||
]
|
||||
mock_eps.return_value = eps
|
||||
|
||||
result = load_plugins_by_group("test.group")
|
||||
self.assertIn("alpha", result)
|
||||
self.assertNotIn("beta", result)
|
||||
self.assertIn("gamma", result)
|
||||
|
||||
@patch("sglang.srt.plugins.entry_points")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
def test_excluded_dists(self, mock_envs, mock_eps):
|
||||
"""SGLANG_PLATFORM excludes other platform dists; empty when unset."""
|
||||
# Case 1: no env set → empty
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
self.assertEqual(_get_excluded_dists(), set())
|
||||
|
||||
# Case 2: env set → exclude other dists
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = "kunlun"
|
||||
ep_kunlun = _make_ep("kunlun", dist_name="kunlun-pkg")
|
||||
ep_other = _make_ep("other_hw", dist_name="other-pkg")
|
||||
mock_eps.return_value = [ep_kunlun, ep_other]
|
||||
|
||||
excluded = _get_excluded_dists()
|
||||
self.assertNotIn("kunlun-pkg", excluded)
|
||||
self.assertIn("other-pkg", excluded)
|
||||
|
||||
@patch("sglang.srt.plugins.HookRegistry")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
@patch("sglang.srt.plugins.entry_points")
|
||||
def test_current_plugin_source_set_during_and_reset_after(
|
||||
self, mock_eps, mock_envs, mock_registry
|
||||
):
|
||||
"""_current_plugin_source is set during plugin execution, reset after."""
|
||||
sources_seen = []
|
||||
|
||||
def spy_plugin():
|
||||
sources_seen.append(_current_plugin_source.get())
|
||||
|
||||
mock_eps.return_value = [_make_ep("spy", load_fn=spy_plugin)]
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
mock_envs.SGLANG_PLUGINS.get.return_value = ""
|
||||
|
||||
load_plugins()
|
||||
# During execution: source was set (not None)
|
||||
self.assertEqual(len(sources_seen), 1)
|
||||
self.assertIsNotNone(sources_seen[0])
|
||||
self.assertEqual(sources_seen[0].plugin_name, "spy")
|
||||
# After execution: source is back to None
|
||||
self.assertIsNone(_current_plugin_source.get())
|
||||
|
||||
@patch("sglang.srt.plugins.HookRegistry")
|
||||
@patch("sglang.srt.plugins.envs")
|
||||
@patch("sglang.srt.plugins.entry_points")
|
||||
def test_current_plugin_source_reset_after_exception(
|
||||
self, mock_eps, mock_envs, mock_registry
|
||||
):
|
||||
"""_current_plugin_source is reset to None even when a plugin raises."""
|
||||
mock_envs.SGLANG_PLATFORM.get.return_value = ""
|
||||
mock_envs.SGLANG_PLUGINS.get.return_value = ""
|
||||
|
||||
def bad_plugin():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
mock_eps.return_value = [_make_ep("bad", load_fn=bad_plugin)]
|
||||
|
||||
load_plugins()
|
||||
self.assertIsNone(_current_plugin_source.get())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user