[Bugfix][M9] Restore model-neutral PD coordinator metadata

docs/unified_execution.md

@@ -157,6 +157,20 @@ and optional stream overlap remain execution choices behind this interface.
 | Preemption snapshots | Engine-owned PreemptionManager, invoked through scheduler callbacks |
 | Transport staging and remote-access lifetime | Connector; local Request/KV changes are coordinated through StageEngine |
 
+**Model configuration.** Deployment entry points construct a shared `ModelConfig`
+from the resolved checkpoint JSON, an existing concrete config, or `module.config`
+before spawning workers. It retains the concrete model config and delegates common
+request limits, including derived properties. StageClient and the PD coordinator
+retain it; spawned workers receive copies. Checkpoint workers load weights using
+the concrete config without reparsing JSON; config-only inputs construct models,
+and module inputs retain their supplied weights. Models and StageEngine continue
+to use the concrete config. Configuration is treated as read-only after initialization.
+Worker `ready` messages report runtime capacity and transport information, not a
+serialized model config. PD still verifies worker-computed weight/KV fingerprints
+before connecting peers; static configuration does not prove runtime compatibility.
+PD validates model-family support before spawning workers and explicitly rejects
+unsupported families.
+
 **Connector and deployment.** Worker calls `connector.progress()` independently
 of engine progress. Connector readiness/backpressure reaches Scheduler through
 the engine's readiness predicate. Import/export and safe release also coordinate

tests/test_auto_model.py

@@ -12,12 +12,85 @@
 from tests.helpers import tiny_nanbeige_config, tiny_ouro_config
 from vllm_rlt.models import (
     AutoModelForCausalLM,
+    ModelConfig,
     NanbeigeForCausalLM,
     OuroForCausalLM,
+    prepare_model,
     resolve_local_model_path,
 )
 
 
+@pytest.mark.parametrize("family", ["ouro", "nanbeige", "huginn"])
+def test_entry_config_preserves_properties_and_checkpoint_weights(tmp_path, monkeypatch, family):
+    from tests.test_huginn import tiny_huginn_config
+    from vllm_rlt.models import HuginnForCausalLM
+
+    config_factory, model_cls = {
+        "ouro": (tiny_ouro_config, OuroForCausalLM),
+        "nanbeige": (tiny_nanbeige_config, NanbeigeForCausalLM),
+        "huginn": (tiny_huginn_config, HuginnForCausalLM),
+    }[family]
+    concrete = config_factory()
+    original = model_cls(concrete)
+    (tmp_path / "config.json").write_text(json.dumps(concrete.to_dict()))
+    save_file(
+        {key: value.clone() for key, value in original.state_dict().items()},
+        tmp_path / "model.safetensors",
+    )
+    source, revision, config = prepare_model(tmp_path)
+    assert config.total_ut_steps == concrete.total_ut_steps
+    assert config.max_position_embeddings == concrete.max_position_embeddings
+    assert config.vocab_size == concrete.vocab_size
+    assert ModelConfig(concrete).model_specific_config is concrete
+    assert prepare_model(original)[2].model_specific_config is original.config
+    with pytest.raises(ValueError, match="do not match"):
+        prepare_model(original, model_config=ModelConfig(tiny_ouro_config(vocab_size=99)))
+
+    # Neither the Auto factory nor the family loader may reconstruct configuration.
+    (tmp_path / "config.json").unlink()
+    resolver = Mock(side_effect=AssertionError("configuration resolved twice"))
+    monkeypatch.setattr("vllm_rlt.models.resolve_model_config", resolver)
+    loaded = AutoModelForCausalLM.from_pretrained(
+        source, revision=revision, model_config=config, dtype=torch.float32
+    )
+    assert loaded.config is config.model_specific_config
+    for key, value in original.state_dict().items():
+        torch.testing.assert_close(loaded.state_dict()[key], value, rtol=0, atol=0)
+    resolver.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "model_cls,config_factory",
+    [
+        (OuroForCausalLM, tiny_ouro_config),
+        (NanbeigeForCausalLM, tiny_nanbeige_config),
+    ],
+)
+def test_stage_checkpoint_uses_entry_config_after_spawn(tmp_path, model_cls, config_factory):
+    from vllm_rlt.config import CacheConfig
+    from vllm_rlt.engine.llm_engine import StageClient
+
+    concrete = config_factory()
+    model = model_cls(concrete)
+    save_file(model.state_dict(), tmp_path / "model.safetensors")
+    # No config.json exists: the child must use the supplied, pickled configuration.
+    config = ModelConfig(concrete)
+    client = StageClient(
+        tmp_path,
+        model_config=config,
+        dtype="float32",
+        attention_backend="torch",
+        cache_config=CacheConfig(64, 2),
+    )
+    try:
+        assert client.model_config is config
+        assert client.ready
+        assert "model" not in client.info and "model_type" not in client.info
+        assert client.stats()["requests"] == 0
+    finally:
+        client.close()
+
+
 @pytest.fixture
 def confirm_download(monkeypatch):
     monkeypatch.setattr("sys.stdin.isatty", lambda: True)

tests/test_pd.py

@@ -105,6 +105,77 @@ def test_pd_configuration_validation(options):
         PDConfig(**options)
 
 
+def test_pd_rejects_huginn_before_spawning(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm_rlt.models import HuginnConfig
+    from vllm_rlt.pd.engine import PDEngine
+
+    clients = Mock()
+    monkeypatch.setattr("vllm_rlt.pd.engine.StageClient", clients)
+    with pytest.raises(ValueError, match="PD does not support Huginn"):
+        PDEngine(HuginnConfig())
+    clients.assert_not_called()
+
+
+def test_pd_nanbeige_startup_uses_entry_config_for_request_validation(monkeypatch):
+    from unittest.mock import Mock
+
+    from tests.helpers import tiny_nanbeige_config
+    from vllm_rlt.pd.engine import PDEngine
+
+    config = tiny_nanbeige_config()
+
+    def ready_client(model, *, stage_config, **options):
+        # Replace only the process boundary. The coordinator still constructs,
+        # connects peers and validates requests without model metadata in ready.
+        assert model is config
+        assert options["model_config"].model_specific_config is config
+        return Mock(
+            ready=True,
+            connected=True,
+            stopped=True,
+            info=dict(
+                kind="ready",
+                stage_id=stage_config.stage_id,
+                role=stage_config.role,
+                fingerprint="same-model-and-cache",
+                info={"agent": stage_config.stage_id},
+                num_blocks=64,
+                block_size=16,
+                depths=model.total_ut_steps,
+            ),
+        )
+
+    clients = Mock(side_effect=ready_client)
+    monkeypatch.setattr("vllm_rlt.pd.engine.StageClient", clients)
+    with PDEngine(config, attention_backend="triton") as engine:
+        assert clients.call_count == 2
+        assert engine.model_config.model_specific_config is config
+        assert all(
+            call.kwargs["model_config"] is engine.model_config for call in clients.call_args_list
+        )
+        # Accept the upper vocabulary/context boundaries at Nanbeige's depth.
+        params = SamplingParams(
+            max_tokens=config.max_position_embeddings, max_loops=config.total_ut_steps
+        )
+        engine.add_request("valid", [config.vocab_size - 1], params)
+        assert engine.has_request("valid")
+        with pytest.raises(ValueError, match="valid model token IDs"):
+            engine.add_request("bad_token", [config.vocab_size], params)
+        with pytest.raises(ValueError, match="loop bounds exceed model depth"):
+            engine.add_request(
+                "bad_depth", [1], SamplingParams(max_loops=config.total_ut_steps + 1)
+            )
+        with pytest.raises(ValueError, match="exceeds context capacity"):
+            engine.add_request(
+                "bad_context", [1], SamplingParams(max_tokens=config.max_position_embeddings + 1)
+            )
+        assert engine.request_ids() == ("valid",)
+    assert engine.closed
+    assert not engine.has_unfinished_requests()
+
+
 def drain(engine):
     outputs = {}
     deadline = time.monotonic() + 90

tests/test_stage_engine.py

@@ -254,12 +254,14 @@ def test_spawned_mixed_worker_matches_local_without_connector(asynchronous):
     import multiprocessing
 
     from vllm_rlt.engine.worker import worker_main
+    from vllm_rlt.models import ModelConfig
 
     local = make_stage_engine(mixed=True, asynchronous=asynchronous)
     options = dict(
         seed=123,
         dtype="float32",
         revision=None,
+        model_config=ModelConfig(tiny_ouro_config()),
         engine=dict(
             cache_config=CacheConfig(128, 2),
             scheduler_config=SchedulerConfig(
@@ -296,6 +298,7 @@ def receive():
         ready = receive()
         assert ready["kind"] == "ready"
         assert ready["role"] == "prefill_decode"
+        assert "model" not in ready and "model_type" not in ready
         assert "info" not in ready  # No registered transport or PD handshake.
         params = SamplingParams(max_tokens=3, ignore_eos=True)
         for request_id, tokens in (("a", [1, 2, 3, 4, 5]), ("b", [2, 3])):
@@ -561,6 +564,9 @@ def test_module_weight_source_preserves_explicit_deployment(
     _, device, weights, settings = launch.call_args.args
     assert device == expected_device
     assert settings["dtype"] == expected_dtype
+    assert not client.ready
+    assert client.model_config.model_specific_config is model.config
+    assert settings["model_config"] is client.model_config
     assert weights is not model
     assert next(weights.parameters()).device.type == "cpu"
     client.closed = True

vllm_rlt/engine/llm_engine.py

@@ -4,14 +4,14 @@
 import time
 from collections import deque
 from dataclasses import replace
-from types import SimpleNamespace
 from uuid import uuid4
 
 import torch
 
 from vllm_rlt.engine.stage_config import StageConfig
 from vllm_rlt.engine.stage_engine import StageEngine
 from vllm_rlt.engine.worker import start_stage_worker
+from vllm_rlt.models import prepare_model
 
 LLMEngine = StageEngine
 
@@ -30,6 +30,7 @@ def __init__(
         device=None,
         dtype=None,
         revision=None,
+        model_config=None,
         seed=0,
         startup_timeout=120,
         shutdown_timeout=30,
@@ -39,6 +40,9 @@ def __init__(
         wait_for_ready=True,
         **engine_options,
     ):
+        model, revision, self.model_config = prepare_model(
+            model, revision=revision, model_config=model_config
+        )
         # An explicitly supplied model is a weight source, never a second runner.
         # Checkpoint paths load directly inside the worker without a host copy.
         if isinstance(model, torch.nn.Module):
@@ -69,7 +73,13 @@ def __init__(
             self.stage_config,
             device,
             model,
-            dict(seed=seed, dtype=dtype, revision=revision, engine=engine_options),
+            dict(
+                seed=seed,
+                dtype=dtype,
+                revision=revision,
+                engine=engine_options,
+                model_config=self.model_config,
+            ),
             connector_factory=connector_factory,
             max_control_messages=max_control_messages,
         )
@@ -131,9 +141,6 @@ def _receive(self, *, timeout=0):
             raise error(f"stage worker failed:\n{self.failure}")
         if kind == "ready":
             self.info, self.ready = message, True
-            self.model_config = SimpleNamespace(
-                **message["model"], model_type=message["model_type"]
-            )
         elif kind == "connected":
             self.connected = True
         elif kind == "reply":

vllm_rlt/engine/worker.py

@@ -9,7 +9,7 @@
 import time
 import traceback
 from collections.abc import Callable
-from dataclasses import asdict, is_dataclass
+from dataclasses import is_dataclass
 
 import torch
 
@@ -43,14 +43,21 @@ def __init__(
         self.channel = channel
         self.max_control_messages = max_control_messages
         dtype = getattr(torch, options["dtype"])
+        model_config = options["model_config"]
         if isinstance(model_source, torch.nn.Module):
             self.model = model_source.to(device=device, dtype=dtype)
         elif is_dataclass(model_source):
-            model_cls = MODEL_MAPPING[model_source.to_dict()["model_type"]]
-            self.model = model_cls(model_source).to(device=device, dtype=dtype)
+            model_cls = MODEL_MAPPING[model_config.model_type]
+            self.model = model_cls(model_config.model_specific_config).to(
+                device=device, dtype=dtype
+            )
         else:
             self.model = AutoModelForCausalLM.from_pretrained(
-                model_source, revision=options["revision"], device=str(device), dtype=dtype
+                model_source,
+                revision=options["revision"],
+                device=str(device),
+                dtype=dtype,
+                model_config=model_config,
             )
         self.engine = StageEngine(self.model, stage_config=stage_config, **options["engine"])
         self.connector = (
@@ -66,8 +73,6 @@ def __init__(
         info = dict(
             stage_id=stage_config.stage_id,
             role=stage_config.role,
-            model=asdict(self.model.config),
-            model_type=self.model.config.to_dict()["model_type"],
             num_blocks=cache.num_blocks,
             block_size=cache.block_size,
             depths=cache.storage_depths,

vllm_rlt/entrypoints/llm.py

@@ -5,7 +5,7 @@
 from vllm_rlt.engine.llm_engine import StageClient
 from vllm_rlt.models import (
     load_tokenizer,
-    resolve_model_config,
+    prepare_model,
     resolve_model_source,
 )
 from vllm_rlt.sampling_params import SamplingParams
@@ -29,15 +29,17 @@ def __init__(
         execution_config=None,
         speculative_config=None,
     ):
-        self._tokenizer_source = None
-        if isinstance(model, str):
-            tokenizer_source = resolve_model_source(model)
-            model, revision, _ = resolve_model_config(model, revision=revision)
-            if tokenizer is None:
-                self._tokenizer_source = (tokenizer_source, revision)
+        tokenizer_source = resolve_model_source(model) if isinstance(model, str) else None
+        model, revision, model_config = prepare_model(model, revision=revision)
+        self._tokenizer_source = (
+            (tokenizer_source, revision)
+            if tokenizer is None and tokenizer_source is not None
+            else None
+        )
         self.tokenizer = tokenizer
         self.engine = StageClient(
             model,
+            model_config=model_config,
             device=device,
             dtype=dtype,
             revision=revision,

vllm_rlt/entrypoints/pd_serve.py

@@ -11,7 +11,7 @@
     profile_config_from_args,
     runtime_configs,
 )
-from vllm_rlt.models import load_tokenizer, resolve_model_config, resolve_model_source
+from vllm_rlt.models import ModelConfig, load_tokenizer, resolve_model_config, resolve_model_source
 from vllm_rlt.pd.config import PDConfig
 from vllm_rlt.pd.engine import PDEngine
 
@@ -37,6 +37,7 @@ def load_engine(args):
     cache = options.pop("cache_config")
     engine = PDEngine(
         source,
+        model_config=ModelConfig.from_dict(config),
         revision=revision,
         dtype=args.dtype,
         pd_config=PDConfig(

vllm_rlt/entrypoints/serve.py

@@ -15,6 +15,7 @@
     runtime_configs,
 )
 from vllm_rlt.models import (
+    ModelConfig,
     load_tokenizer,
     resolve_model_config,
     resolve_model_source,
@@ -42,6 +43,7 @@ def load_engine(args):
             raise ValueError("serving requires the Ouro byte-level tokenizer")
     engine = StageClient(
         source,
+        model_config=ModelConfig.from_dict(config),
         revision=revision,
         device=args.device,
         dtype=args.dtype,