diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index e438cbbc0..c0f7d1897 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -48,6 +48,9 @@ class CustomSpecAlgo: def __repr__(self) -> str: return f"CustomSpecAlgo({self.name!r})" + def is_some(self) -> bool: + return True + def is_none(self) -> bool: return False @@ -60,6 +63,9 @@ class CustomSpecAlgo: def is_eagle3(self) -> bool: return False + def is_frozen_kv_mtp(self) -> bool: + return False + def is_dflash(self) -> bool: return False @@ -104,10 +110,58 @@ class CustomSpecAlgo: _REGISTRY: Dict[str, CustomSpecAlgo] = {} -# Builtin enum members + the NEXTN alias; plugins cannot shadow these. -_RESERVED_NAMES = frozenset( - {"DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM", "NONE"} -) +# CLI spellings that are not ``SpeculativeAlgorithm`` members but still resolve +# to a builtin (e.g. NEXTN -> EAGLE). Reserved alongside the enum members so +# plugins cannot shadow them. +_RESERVED_ALIASES = frozenset({"NEXTN"}) + + +def _reserved_names() -> frozenset: + """Names plugins cannot register under: every ``SpeculativeAlgorithm`` + member plus ``_RESERVED_ALIASES``. + + Derived from the enum (lazily, to avoid a circular import — ``spec_info`` + imports this module) so any new builtin is reserved automatically without + editing a second list. + """ + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + return frozenset(algo.name for algo in SpeculativeAlgorithm) | _RESERVED_ALIASES + + +def _assert_custom_spec_algo_conforms(spec_class: Type[CustomSpecAlgo]) -> None: + """Fail fast if ``spec_class`` drifts from the ``SpeculativeAlgorithm`` + duck-typing contract. + + ``from_string`` returns either type and callers dispatch on the shared + ``is_*()`` / ``supports_*()`` interface without isinstance checks, so every + such method on the enum must also exist on the registered spec class — + otherwise a plugin-registered algo hits ``AttributeError`` at a call site + (this is how ``is_some`` / ``is_frozen_kv_mtp`` silently went missing). New + predicates are covered automatically; no second list to maintain. + + Called from ``register_algorithm`` rather than at import time because + ``spec_info`` imports this module, so ``SpeculativeAlgorithm`` does not yet + exist while this module is loading; at registration time it is fully + defined. + """ + # NOTE: use ``vars()`` not ``dir()`` for the enum — ``EnumMeta.__dir__`` + # hides instance methods, so ``dir(SpeculativeAlgorithm)`` would yield an + # empty interface and turn this guard into a silent no-op. + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + interface = { + name + for name in vars(SpeculativeAlgorithm) + if name.startswith(("is_", "supports_")) + } + missing = sorted(interface - set(dir(spec_class))) + if missing: + raise TypeError( + f"{spec_class.__name__} is missing duck-typed methods from " + f"SpeculativeAlgorithm: {missing}. Add them to {spec_class.__name__} " + "so plugin-registered algorithms stay dispatchable." + ) def register_algorithm( @@ -123,12 +177,13 @@ def register_algorithm( ``is_*()`` / ``supports_*()`` / ``create_worker`` method. """ upper = name.upper() - if upper in _RESERVED_NAMES: + if upper in _reserved_names(): raise ValueError( f"'{upper}' is a reserved speculative algorithm name; cannot be re-registered." ) if upper in _REGISTRY: raise ValueError(f"Speculative algorithm '{upper}' already registered.") + _assert_custom_spec_algo_conforms(spec_class) def decorator(factory: WorkerFactory) -> WorkerFactory: _REGISTRY[upper] = spec_class( diff --git a/test/registered/unit/spec/test_spec_registry.py b/test/registered/unit/spec/test_spec_registry.py index baeb9d632..b6dc07baa 100644 --- a/test/registered/unit/spec/test_spec_registry.py +++ b/test/registered/unit/spec/test_spec_registry.py @@ -6,8 +6,9 @@ from unittest.mock import MagicMock from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_registry import ( _REGISTRY, - _RESERVED_NAMES, CustomSpecAlgo, + _assert_custom_spec_algo_conforms, + _reserved_names, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -104,10 +105,21 @@ class TestRegister(_RegistryIsolated): return MagicMock def test_reserved_name_raises(self): - for reserved in _RESERVED_NAMES: + for reserved in _reserved_names(): with self.assertRaisesRegex(ValueError, "reserved"): SpeculativeAlgorithm.register(reserved) + def test_reserved_names_cover_all_enum_members(self): + # Reserved names are derived from the enum, so every builtin (including + # FROZEN_KV_MTP, which a hand-maintained list had omitted) is reserved. + for member in SpeculativeAlgorithm: + self.assertIn(member.name, _reserved_names()) + self.assertIn("NEXTN", _reserved_names()) # CLI alias + + def test_reserved_name_is_case_insensitive(self): + with self.assertRaisesRegex(ValueError, "reserved"): + SpeculativeAlgorithm.register("frozen_kv_mtp") + def test_register_is_case_insensitive_on_collision(self): @SpeculativeAlgorithm.register("MY_FOO") def _factory(server_args): @@ -136,10 +148,19 @@ class TestCustomSpecAlgoInterface(_RegistryIsolated): self.assertFalse(self.algo.is_none()) self.assertFalse(self.algo.is_eagle()) self.assertFalse(self.algo.is_eagle3()) + self.assertFalse(self.algo.is_frozen_kv_mtp()) self.assertFalse(self.algo.is_dflash()) self.assertFalse(self.algo.is_standalone()) self.assertFalse(self.algo.is_ngram()) self.assertTrue(self.algo.is_speculative()) + # A registered plugin is never NONE -> is_some() mirrors the enum. + self.assertTrue(self.algo.is_some()) + + def test_is_some_matches_enum_semantics(self): + # is_some() is called on spec algos in overlap_utils.py; a CustomSpecAlgo + # must answer it the same way the enum does (True iff not NONE). + self.assertEqual(self.algo.is_some(), not self.algo.is_none()) + self.assertEqual(SpeculativeAlgorithm.EAGLE.is_some(), self.algo.is_some()) def test_supports_spec_v2_follows_supports_overlap(self): # Plugin registered with supports_overlap=False -> not spec_v2. @@ -216,6 +237,60 @@ class TestSubclassOverride(_RegistryIsolated): self.assertEqual(algo.create_worker(MagicMock()), "custom-dispatched") +class TestConformanceGuard(_RegistryIsolated): + """register_algorithm rejects spec classes that drift from the enum's + is_*() / supports_*() duck-typing interface.""" + + def test_base_custom_spec_algo_conforms(self): + # The shipped base class must satisfy its own contract. + _assert_custom_spec_algo_conforms(CustomSpecAlgo) + + def test_conforming_subclass_passes(self): + class Good(CustomSpecAlgo): + def is_eagle(self) -> bool: + return True + + _assert_custom_spec_algo_conforms(Good) # does not raise + + @staticmethod + def _spec_class_missing(method: str) -> type: + # Build a class exposing the full enum interface except `method`. A real + # subclass can't be "missing" an inherited method, so the failure mode + # the guard catches is the base class itself dropping one — simulated + # here with a standalone class. + interface = { + name + for name in vars(SpeculativeAlgorithm) + if name.startswith(("is_", "supports_")) + } + body = {m: (lambda self: False) for m in interface if m != method} + return type("Broken", (), body) + + def test_missing_predicate_raises(self): + Broken = self._spec_class_missing("is_some") + with self.assertRaisesRegex(TypeError, "is_some"): + _assert_custom_spec_algo_conforms(Broken) + + def test_register_rejects_nonconforming_spec_class(self): + Broken = self._spec_class_missing("is_some") + with self.assertRaisesRegex(TypeError, "missing duck-typed methods"): + + @SpeculativeAlgorithm.register("MY_BROKEN", spec_class=Broken) + def _factory(server_args): + return MagicMock + + def test_enum_interface_subset_of_custom_spec_algo(self): + # Every is_*/supports_* method on the enum exists on CustomSpecAlgo. + # vars() (not dir()) because EnumMeta.__dir__ hides instance methods. + interface = { + name + for name in vars(SpeculativeAlgorithm) + if name.startswith(("is_", "supports_")) + } + self.assertTrue(interface) # guard against an empty (no-op) interface + self.assertEqual(interface - set(dir(CustomSpecAlgo)), set()) + + class TestCrossTypeIdentity(_RegistryIsolated): """A plugin algo and a builtin enum value must never compare equal."""