Add assertions to reject training data generated under the sglang default weight version

miles/ray/rollout/rollout_executor.py

@@ -34,6 +34,7 @@
 from miles.utils.misc import load_function
 from miles.utils.timer import timer
 from miles.utils.tracking_utils.tracking import init_tracking
+from miles.utils.weight_version import assert_samples_weight_version_sane
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 logging.getLogger("httpcore").setLevel(logging.WARNING)
@@ -224,6 +225,7 @@ async def _get_rollout_data(self, rollout_id):
             data, metadata = postprocess_rollout_data(
                 self.args, data, train_parallel_config=self.train_parallel_config
             )
+            assert_samples_weight_version_sane(self.args, samples=data)
             if RolloutDataInjectionUtil.should_inject(self.args, rollout_id):
                 generated_data = data
                 data, metadata = RolloutDataInjectionUtil.load(self.args, rollout_id=rollout_id)

miles/utils/weight_version.py

@@ -0,0 +1,30 @@
+import re
+from argparse import Namespace
+from typing import TYPE_CHECKING
+
+from miles.utils.lora import is_lora_enabled
+
+if TYPE_CHECKING:
+    from miles.utils.types import Sample
+
+SGLANG_DEFAULT_WEIGHT_VERSION = "default"
+
+_NUMERIC_VERSION_PATTERN = re.compile(r"[0-9]+")
+
+
+def assert_samples_weight_version_sane(args: Namespace, samples: list["Sample"]) -> None:
+    if args.debug_rollout_only or args.debug_skip_weight_update or is_lora_enabled(args):
+        return
+
+    for sample in samples:
+        for span in sample.all_weight_version_spans:
+            assert span.version != SGLANG_DEFAULT_WEIGHT_VERSION, (
+                f"sample index={sample.index} tokens [{span.abs_start}, {span.abs_end}) were generated under "
+                f"weight version {SGLANG_DEFAULT_WEIGHT_VERSION!r}, the sglang placeholder for an engine whose "
+                f"weights were never updated; training data must only come from engines that received a weight update"
+            )
+            assert _NUMERIC_VERSION_PATTERN.fullmatch(span.version), (
+                f"sample index={sample.index} tokens [{span.abs_start}, {span.abs_end}) carry weight version "
+                f"{span.version!r}, which is not the numeric version miles stamps on weight updates; "
+                f"the engine serving this sample got its weights from somewhere miles does not know about"
+            )

tests/fast/ray/rollout/conftest.py

@@ -64,6 +64,7 @@ def make_args(**overrides: Any) -> Namespace:
         # placement / colocation
         debug_train_only=False,
         debug_rollout_only=False,
+        debug_skip_weight_update=False,
         colocate=False,
         actor_num_nodes=1,
         actor_num_gpus_per_node=8,

tests/fast/ray/rollout/real_ray/test_rollout_executor.py

@@ -14,6 +14,7 @@
     RolloutFnTrainInput,
     RolloutFnTrainOutput,
 )
+from miles.utils.types import WeightVersionSpan, WeightVersionsPerCall
 
 
 @pytest.fixture
@@ -164,6 +165,24 @@ def fake_rollout_fn(input):
             assert "loss_masks" in partition
             assert len(partition["tokens"]) == 4
 
+    async def test_rejects_samples_generated_under_the_default_weight_version(self, ray_local_mode, patch_low_level):
+        """A batch carrying the sglang never-updated version must fail get(), not reach training."""
+        args = _make_test_args()
+        args.global_batch_size = 8
+
+        executor = _make_executor(args)
+        executor.set_train_parallel_config({"dp_size": 2})
+
+        samples = make_samples_grouped(n_groups=2, group_size=4)
+        samples[0].weight_versions = [
+            WeightVersionsPerCall(spans=[WeightVersionSpan(version="default", abs_start=0, abs_end=1)])
+        ]
+
+        executor.generate_rollout = lambda input: RolloutFnTrainOutput(samples=[samples], metrics=None)
+
+        with pytest.raises(AssertionError, match="never updated"):
+            await executor.get(rollout_id=42)
+
     async def test_does_not_touch_the_inference_side(self, ray_local_mode, patch_low_level):
         """The controller is a driver-side object the executor cannot reach, so generate must not need it."""
         args = _make_test_args()

tests/fast/utils/test_weight_version.py

@@ -0,0 +1,81 @@
+from argparse import Namespace
+
+import pytest
+
+from miles.utils.types import Sample, WeightVersionSpan, WeightVersionsPerCall
+from miles.utils.weight_version import SGLANG_DEFAULT_WEIGHT_VERSION, assert_samples_weight_version_sane
+
+SGLANG_LITERAL = "default"
+
+
+def _make_args(**overrides: object) -> Namespace:
+    args = Namespace(debug_rollout_only=False, debug_skip_weight_update=False, lora_rank=0, lora_adapter_path=None)
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+def _make_sample(versions: list[str], index: int = 0) -> Sample:
+    calls = [
+        WeightVersionsPerCall(spans=[WeightVersionSpan(version=version, abs_start=i, abs_end=i + 1)])
+        for i, version in enumerate(versions)
+    ]
+    return Sample(
+        index=index,
+        tokens=list(range(len(versions))),
+        response_length=len(versions),
+        weight_versions=calls,
+    )
+
+
+class TestAssertSamplesWeightVersionSane:
+    def test_the_guarded_literal_is_the_sglang_wire_value(self):
+        """sglang sends the literal 'default'; a drifted constant would let real default spans through."""
+        assert SGLANG_DEFAULT_WEIGHT_VERSION == SGLANG_LITERAL
+
+    def test_an_updated_engine_version_passes(self):
+        """Tokens generated after a weight update carry a real version and must be accepted."""
+        assert_samples_weight_version_sane(_make_args(), samples=[_make_sample(["3"])])
+
+    def test_multiple_numeric_weight_versions_pass(self):
+        """A trajectory spanning several weight updates carries mixed numeric versions and must be accepted."""
+        assert_samples_weight_version_sane(_make_args(), samples=[_make_sample(["1", "2", "10"])])
+
+    def test_the_sglang_default_version_fails(self):
+        """Tokens generated by a never-updated engine must be rejected, not silently trained on."""
+        with pytest.raises(AssertionError, match="never updated"):
+            assert_samples_weight_version_sane(_make_args(), samples=[_make_sample([SGLANG_LITERAL])])
+
+    def test_a_default_prefix_before_an_update_fails(self):
+        """A request whose first tokens predate the first weight update must be caught despite later real versions."""
+        with pytest.raises(AssertionError, match="never updated"):
+            assert_samples_weight_version_sane(_make_args(), samples=[_make_sample([SGLANG_LITERAL, "1"])])
+
+    def test_a_later_sample_with_a_trailing_foreign_version_fails(self):
+        """Every sample and every span must be checked, not only the first sample or the first span."""
+        samples = [_make_sample(["1", "2"], index=0), _make_sample(["3", SGLANG_LITERAL], index=7)]
+        with pytest.raises(AssertionError, match=r"index=7 tokens \[1, 2\) were generated under weight version"):
+            assert_samples_weight_version_sane(_make_args(), samples=samples)
+
+    @pytest.mark.parametrize("version", ["mock-v0", "v3", "3.0", "-1", "", "3 "])
+    def test_a_non_numeric_version_fails(self, version: str):
+        """miles stamps bare numeric versions, so anything else means the weights came from outside this run."""
+        with pytest.raises(AssertionError, match="not the numeric version"):
+            assert_samples_weight_version_sane(_make_args(), samples=[_make_sample([version])])
+
+    def test_a_sample_without_spans_passes(self):
+        """Samples with no recorded weight versions carry nothing to validate."""
+        assert_samples_weight_version_sane(_make_args(), samples=[_make_sample([])])
+
+    @pytest.mark.parametrize(
+        "overrides",
+        [
+            {"debug_rollout_only": True},
+            {"debug_skip_weight_update": True},
+            {"lora_rank": 8},
+            {"lora_adapter_path": "/adapters/foo"},
+        ],
+    )
+    def test_modes_that_never_push_weights_are_exempt(self, overrides: dict[str, object]):
+        """Debug modes and LoRA legitimately serve un-updated weights, so the default version is allowed."""
+        assert_samples_weight_version_sane(_make_args(**overrides), samples=[_make_sample([SGLANG_LITERAL])])