feat(config): add v6.1 ruleset and pin RNG seeds from submission ruleset (#389)

* feat(config): add v6.1 ruleset and pin RNG seeds from submission ruleset

Adds the MLPerf Inference v6.1 RoundRuleset (schedule/sample_index seeds from
loadgen/mlperf.conf; per-model latency targets unchanged from v5.1) and makes
it CURRENT. The registry now registers every known round by version, so both
v5.1 and v6.1 are resolvable.

When a config carries submission_ref, _resolve_and_validate now pins the
runtime RNG seeds from the selected ruleset during construction — before the
config is dumped to the report dir, so the persisted config.yaml,
RuntimeSettings, and report all reflect the seeds that actually ran. The
ruleset value wins over any user-supplied seed (mirrors LoadGen locking core
seeds from user.conf). Lenient by design: an unregistered ruleset leaves the
config unchanged, so placeholder/non-submission configs are unaffected.

Duration/latency overrides and the runtime-path integration (surfacing ruleset
latency targets in the report) are deferred to a follow-up.

---------

Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>

Author: arekay-nv

src/inference_endpoint/config/ruleset_registry.py

@@ -28,6 +28,7 @@
 
 from typing import TYPE_CHECKING
 
+from .rulesets.mlcommons.rules import ALL_ROUNDS as mlcommons_rounds
 from .rulesets.mlcommons.rules import CURRENT as mlcommons_current
 from .rulesets.mlcommons.rules import EDGE_CURRENT as mlcommons_edge_current
 
@@ -44,7 +45,16 @@ def register_ruleset(name: str, ruleset: BenchmarkSuiteRuleset) -> None:
     Args:
         name: Unique name/version identifier (e.g., "mlperf-inference-v5.1")
         ruleset: The ruleset instance to register
+
+    Raises:
+        ValueError: If a ruleset is already registered under ``name`` (prevents
+            an accidental overwrite via this function). The MLCommons
+            auto-registration path uses ``setdefault`` and does not hit this
+            guard; round-version uniqueness is enforced by
+            ``test_round_versions_are_unique``.
     """
+    if name in _RULESET_REGISTRY:
+        raise ValueError(f"Ruleset {name!r} is already registered")
     _RULESET_REGISTRY[name] = ruleset
 
 
@@ -79,16 +89,22 @@ def list_rulesets() -> list[str]:
 
 # Auto-register MLCommons rulesets
 def _auto_register_mlcommons():
-    """Auto-register MLCommons rulesets."""
-    # Register with version-specific name
-    register_ruleset(f"mlperf-inference-{mlcommons_current.version}", mlcommons_current)
-    # Also register as "mlcommons-current" for convenience
-    register_ruleset("mlcommons-current", mlcommons_current)
-    # Edge-Agentic (BFCL v4) ruleset
-    register_ruleset(
+    """Auto-register MLCommons rulesets (idempotent).
+
+    Uses ``setdefault`` so re-running (e.g. a module re-import) is a no-op
+    rather than tripping ``register_ruleset``'s duplicate guard. Uniqueness of
+    round versions is enforced separately by ``test_round_versions_are_unique``.
+    """
+    # Register every known round under its version-specific name
+    for ruleset in mlcommons_rounds:
+        _RULESET_REGISTRY.setdefault(f"mlperf-inference-{ruleset.version}", ruleset)
+    # Also register the current round as "mlcommons-current" for convenience
+    _RULESET_REGISTRY.setdefault("mlcommons-current", mlcommons_current)
+    # Edge-agentic ruleset: by version and as the edge "current".
+    _RULESET_REGISTRY.setdefault(
         f"mlperf-{mlcommons_edge_current.version}", mlcommons_edge_current
     )  # -> "mlperf-edge-v0.1"
-    register_ruleset("mlperf-edge-current", mlcommons_edge_current)
+    _RULESET_REGISTRY.setdefault("mlperf-edge-current", mlcommons_edge_current)
 
 
 # Auto-register on import

src/inference_endpoint/config/rulesets/mlcommons/rules.py

@@ -19,6 +19,7 @@
 https://github.com/mlcommons/inference_policies/blob/master/inference_rules.adoc
 """
 
+import copy
 import random
 from dataclasses import dataclass, field
 from enum import Enum
@@ -272,7 +273,30 @@ def apply_user_config(
 )
 
 
-CURRENT = _v5_1
+# v6.1 per-model latency targets are unchanged from v5.1 (they match the
+# current mlperf.conf); only the round RNG seeds rotate. Seeds are the
+# schedule_rng_seed / sample_index_rng_seed values from loadgen/mlperf.conf,
+# pinned to a specific upstream commit for traceability:
+# https://github.com/mlcommons/inference/blob/10f823448fd38bb739e52690efe8191c3a55412b/loadgen/mlperf.conf#L42-L43
+# qsl_rng_seed is intentionally not modeled: it only selects the load order of
+# the first-N working set, which the sample-index shuffle already covers.
+_v6_1 = RoundRuleset(
+    version="v6.1",
+    scheduler_rng_seed=16159082839903944936,
+    sample_index_rng_seed=2747215439041700203,
+    benchmark_rulesets={
+        # Keep the model-singleton keys; deep-copy only the per-model rules so
+        # v6.1's mutable leaves (e.g. reported_metrics) can't leak into v5.1.
+        model: copy.deepcopy(rules)
+        for model, rules in _v5_1.benchmark_rulesets.items()
+    },
+)
+
+
+# All known rounds, registered by version in the ruleset registry.
+ALL_ROUNDS = [_v5_1, _v6_1]
+
+CURRENT = _v6_1
 
 
 # --- Edge-Agentic (BFCL v4) ruleset ---

src/inference_endpoint/config/schema.py

@@ -22,6 +22,7 @@
 
 from __future__ import annotations
 
+import logging
 from collections import Counter
 from enum import Enum
 from pathlib import Path
@@ -50,6 +51,8 @@
 from .ruleset_base import BenchmarkSuiteRuleset
 from .utils import parse_dataset_string, resolve_env_vars
 
+logger = logging.getLogger(__name__)
+
 
 class SystemDefaults(BaseModel):
     DEFAULT_TIMEOUT: ClassVar[float] = 300.0
@@ -1029,8 +1032,92 @@ def _resolve_and_validate(self) -> Self:
                 f"got '{lp.type}'"
             )
 
+        # Pin RNG seeds from the submission ruleset. Done last so the values
+        # are baked into the config before any consumer reads them — the config
+        # dump to the report dir, RuntimeSettings.from_config, and the report
+        # seeds block all see the pinned values.
+        self._apply_ruleset_seed_overrides()
+
         return self
 
+    def _apply_ruleset_seed_overrides(self) -> None:
+        """Override runtime + warmup RNG seeds from the selected submission ruleset.
+
+        MLPerf rounds pin the RNG seeds; this mirrors LoadGen locking the core
+        seeds from ``user.conf`` (a submitter cannot substitute their own).
+        If ``submission_ref`` is unset, the config is left unchanged. If it
+        names an unregistered ruleset, a ``type=SUBMISSION`` config errors (a
+        submission cannot silently fall back to default seeds), while any other
+        type is left unchanged so non-submission/placeholder configs still work.
+
+        The warmup phase is reseeded from the sample-index (dataloader) seed so
+        its sample order derives from the same pinned seed as the perf phase.
+        Only the seed *value* is propagated — each phase builds its own
+        ``random.Random`` downstream, so the RNG object is never shared.
+        """
+        if self.submission_ref is None:
+            return
+        try:
+            ruleset = self.submission_ref.get_ruleset_instance()
+        except KeyError as e:
+            if self.type == TestType.SUBMISSION:
+                raise ValueError(
+                    f"submission_ref.ruleset {self.submission_ref.ruleset!r} is not "
+                    "registered; a submission must pin official RNG seeds and cannot "
+                    "fall back to defaults."
+                ) from e
+            logger.warning(
+                "submission_ref.ruleset %r is not registered; skipping ruleset "
+                "seed overrides.",
+                self.submission_ref.ruleset,
+            )
+            return
+
+        # A ruleset used as a submission_ref must pin both seeds. ``None`` means
+        # "unseeded" in the general ruleset contract (ruleset_base.py), but an
+        # unseeded submission is incoherent and would silently diverge from the
+        # random.Random(None) path in RoundRuleset.apply_user_config. Reject it.
+        if ruleset.scheduler_rng_seed is None or ruleset.sample_index_rng_seed is None:
+            raise ValueError(
+                f"submission_ref.ruleset {self.submission_ref.ruleset!r} leaves an "
+                "RNG seed unset; a pinned ruleset must define both the scheduler "
+                "and sample-index seeds."
+            )
+
+        # Rebuild through model_validate (not model_copy(update=)): with
+        # extra='forbid' this validates seed *values* and rejects renamed/unknown
+        # fields. model_copy(update=) writes straight into __dict__, so a
+        # wrong-typed (e.g. str) or renamed seed would slip through unchecked.
+        runtime = self.settings.runtime
+        new_runtime = type(runtime).model_validate(
+            {
+                **runtime.model_dump(),
+                "scheduler_random_seed": ruleset.scheduler_rng_seed,
+                "dataloader_random_seed": ruleset.sample_index_rng_seed,
+            }
+        )
+        warmup = self.settings.warmup
+        new_warmup = type(warmup).model_validate(
+            {
+                **warmup.model_dump(),
+                "warmup_random_seed": ruleset.sample_index_rng_seed,
+            }
+        )
+        object.__setattr__(
+            self,
+            "settings",
+            self.settings.model_copy(
+                update={"runtime": new_runtime, "warmup": new_warmup}
+            ),
+        )
+        logger.debug(
+            "Pinned RNG seeds from ruleset %r: scheduler=%s sample_index=%s "
+            "(warmup reseeded from sample_index)",
+            self.submission_ref.ruleset,
+            ruleset.scheduler_rng_seed,
+            ruleset.sample_index_rng_seed,
+        )
+
     @model_validator(mode="after")
     def _propagate_client_api_type(self) -> Self:
         """Sync client.api_type from endpoint_config.api_type at construction.

src/inference_endpoint/config/templates/submission_template.yaml

@@ -7,7 +7,7 @@ benchmark_mode: "offline" # Required for submission: "offline" or "online"
 # Submission reference (model + ruleset for compliance checks)
 submission_ref:
   model: "llama-2-70b"
-  ruleset: "mlperf-inference-v6.0"
+  ruleset: "mlperf-inference-v6.1"
 
 model_params:
   name: "meta-llama/Llama-2-70B-Chat-HF"

tests/unit/commands/test_benchmark.py

@@ -382,7 +382,7 @@ def test_from_config_submission_defaults_to_both(self, mock_run, tmp_path):
   - path: "test.jsonl"
 submission_ref:
   model: "test-model"
-  ruleset: "test"
+  ruleset: "mlperf-inference-v6.1"
 """
         config_file = tmp_path / "sub.yaml"
         config_file.write_text(yaml_content)

tests/unit/commands/test_util_commands.py

@@ -81,12 +81,12 @@ def test_submission_ref_logging(self, tmp_path, caplog):
   - path: "test.jsonl"
 submission_ref:
   model: "llama"
-  ruleset: "mlperf-v6"
+  ruleset: "mlperf-inference-v6.1"
 """)
         with caplog.at_level("INFO"):
             execute_validate(config_file)
         assert "Submission:" in caplog.text
-        assert "mlperf-v6" in caplog.text
+        assert "mlperf-inference-v6.1" in caplog.text
 
     @pytest.mark.unit
     def test_invalid_yaml(self, tmp_path):

tests/unit/config/rulesets/mlcommons/test_rules.py

@@ -15,9 +15,14 @@
 
 import pytest
 from inference_endpoint import metrics
-from inference_endpoint.config.ruleset_registry import get_ruleset, list_rulesets
+from inference_endpoint.config.ruleset_registry import (
+    get_ruleset,
+    list_rulesets,
+    register_ruleset,
+)
 from inference_endpoint.config.rulesets.mlcommons import datasets, models
 from inference_endpoint.config.rulesets.mlcommons.rules import (
+    ALL_ROUNDS,
     CURRENT,
     EDGE_CURRENT,
     OptimizationPriority,
@@ -116,6 +121,62 @@ def test_apply_user_config_min_sample_count_override():
 
 
 @pytest.mark.unit
+def test_current_round_is_v6_1():
+    assert CURRENT.version == "v6.1"
+
+
+@pytest.mark.unit
+def test_v6_1_official_seeds():
+    """Seeds are the schedule/sample_index values from loadgen/mlperf.conf."""
+    v6_1 = get_ruleset("mlperf-inference-v6.1")
+    assert v6_1.scheduler_rng_seed == 16159082839903944936
+    assert v6_1.sample_index_rng_seed == 2747215439041700203
+
+
+@pytest.mark.unit
+def test_all_rounds_registered_by_version():
+    names = list_rulesets()
+    for ruleset in ALL_ROUNDS:
+        assert f"mlperf-inference-{ruleset.version}" in names
+    # Both the prior and current round remain resolvable.
+    assert "mlperf-inference-v5.1" in names
+    assert "mlperf-inference-v6.1" in names
+    assert get_ruleset("mlcommons-current") is CURRENT
+
+
+@pytest.mark.unit
+def test_round_versions_are_unique():
+    """Each registered round must own a distinct version string, else the
+    registry would silently overwrite one round's entry with another's."""
+    assert len(ALL_ROUNDS) == len({r.version for r in ALL_ROUNDS})
+
+
+@pytest.mark.unit
+def test_register_ruleset_rejects_duplicate():
+    """Re-registering an existing name raises rather than clobbering it."""
+    with pytest.raises(ValueError):
+        register_ruleset("mlperf-inference-v6.1", CURRENT)
+
+
+@pytest.mark.unit
+def test_v6_1_latency_targets_match_v5_1():
+    """Only the round seeds rotate; per-model targets are identical."""
+    v5_1 = get_ruleset("mlperf-inference-v5.1")
+    v6_1 = get_ruleset("mlperf-inference-v6.1")
+    assert v6_1.benchmark_rulesets == v5_1.benchmark_rulesets
+    # v6.1 holds an independent copy, not a shared reference to v5.1's dict.
+    assert v6_1.benchmark_rulesets is not v5_1.benchmark_rulesets
+    # Anchor on a concrete target so the equality above verifies real structure.
+    assert (
+        v6_1.benchmark_rulesets[models.Llama3_1_405b][
+            OptimizationPriority.THROUGHPUT
+        ].max_tpot_latency_ms
+        == 175
+    )
+    assert v6_1.scheduler_rng_seed != v5_1.scheduler_rng_seed
+    assert v6_1.sample_index_rng_seed != v5_1.sample_index_rng_seed
+
+
 def test_edge_ruleset_registered():
     # Resolvable by version-specific name and the "current" alias.
     assert get_ruleset("mlperf-edge-v0.1") is EDGE_CURRENT