fix(spec): complete CustomSpecAlgo duck-typing interface and guard against drift (#27300)
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user