fix(spec): complete CustomSpecAlgo duck-typing interface and guard against drift (#27300)

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.
+# CLI spellings that are not ``SpeculativeAlgorithm`` members but still resolve
-_RESERVED_NAMES = frozenset(
+# to a builtin (e.g. NEXTN -> EAGLE). Reserved alongside the enum members so
-    {"DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM", "NONE"}
+# 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(

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."""