Add registration for external model configurations (#39452)

test/registered/unit/configs/test_model_config.py

@@ -1,16 +1,23 @@
-"""Unit tests for hybrid attention model configuration."""
+"""Unit tests for model configuration."""
 
 import unittest
+from tempfile import TemporaryDirectory
 from types import SimpleNamespace
+from unittest import mock
 
+from transformers import LlamaConfig
+
+from sglang.srt.arg_groups.overrides import model_config_of
 from sglang.srt.configs.model_config import (
     ModelConfig,
     get_hybrid_layer_ids,
     is_embedding_gemma,
     is_multimodal_model,
+    register_model_config_factory,
     resolve_spec_hidden_size,
 )
 from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig
+from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -112,5 +119,99 @@ class TestDraftModelConfig(CustomTestCase):
         )
 
 
+class TestExternalModelConfig(CustomTestCase):
+    def setUp(self):
+        registry = mock.patch.dict(
+            "sglang.srt.configs.model_config._MODEL_CONFIG_FACTORIES", clear=True
+        )
+        registry.start()
+        self.addCleanup(registry.stop)
+
+    def test_factory_preserves_arguments_and_uses_the_shared_cache(self):
+        class ExternalArgs(ServerArgs):
+            pass
+
+        class DerivedArgs(ExternalArgs):
+            pass
+
+        result = SimpleNamespace(is_hybrid_swa=False)
+        factory = mock.Mock(return_value=result)
+        register_model_config_factory(ExternalArgs, factory)
+        args = DerivedArgs(model_path="dummy")
+        overrides = dict(
+            model_path="draft",
+            model_revision="draft-revision",
+            is_draft_model=True,
+            context_length=512,
+            dtype="bfloat16",
+        )
+        self.assertIs(ModelConfig.from_server_args(args, **overrides), result)
+        factory.assert_called_once_with(args, **overrides)
+
+        factory.reset_mock()
+        self.assertIs(model_config_of(args), result)
+        self.assertIs(model_config_of(args), result)
+        factory.assert_called_once_with(
+            args,
+            model_path=None,
+            model_revision=None,
+            is_draft_model=False,
+            context_length=None,
+        )
+
+    def test_registration_is_idempotent_and_the_nearest_type_wins(self):
+        class ExternalArgs(ServerArgs):
+            pass
+
+        class DerivedArgs(ExternalArgs):
+            pass
+
+        parent_factory = mock.Mock()
+        child_factory = mock.Mock()
+        register_model_config_factory(ExternalArgs, parent_factory)
+        register_model_config_factory(ExternalArgs, parent_factory)
+        with self.assertRaisesRegex(ValueError, "already registered"):
+            register_model_config_factory(ExternalArgs, child_factory)
+        with self.assertRaisesRegex(TypeError, "ServerArgs subclass"):
+            register_model_config_factory(SimpleNamespace, child_factory)
+        register_model_config_factory(DerivedArgs, child_factory)
+        args = DerivedArgs(model_path="dummy")
+        self.assertIs(ModelConfig.from_server_args(args), child_factory.return_value)
+        parent_factory.assert_not_called()
+
+    def test_capabilities_are_available_during_construction(self):
+        class ExternalArgs(ServerArgs):
+            pass
+
+        class ExternalConfig(ModelConfig):
+            def _derive_multimodal_cuda_graph_support(self, enable_multimodal):
+                super()._derive_multimodal_cuda_graph_support(enable_multimodal)
+                self.is_multimodal_breakable_cuda_graph_supported = True
+
+            def _derive_model_shapes(self):
+                assert self.is_multimodal_breakable_cuda_graph_supported
+                super()._derive_model_shapes()
+
+        register_model_config_factory(
+            ExternalArgs,
+            lambda args, **kwargs: ExternalConfig(model_path=args.model_path),
+        )
+        with TemporaryDirectory() as checkpoint:
+            LlamaConfig(
+                architectures=["LlamaForCausalLM"],
+                hidden_size=16,
+                intermediate_size=32,
+                num_attention_heads=2,
+                num_hidden_layers=2,
+                vocab_size=128,
+            ).save_pretrained(checkpoint)
+            ordinary = ModelConfig.from_server_args(ServerArgs(model_path=checkpoint))
+            external = ModelConfig.from_server_args(ExternalArgs(model_path=checkpoint))
+        self.assertIs(type(ordinary), ModelConfig)
+        self.assertFalse(ordinary.is_multimodal_breakable_cuda_graph_supported)
+        self.assertIs(type(external), ExternalConfig)
+        self.assertTrue(external.is_multimodal_breakable_cuda_graph_supported)
+
+
 if __name__ == "__main__":
     unittest.main()