fix(replay): keep timing and capacity on the selected database version (cherry-pick #291)

python/aisimulate/src/aisimulate/aic.py

@@ -31,6 +31,16 @@ def materialize_aic_num_gpu_blocks(raw: dict[str, Any]) -> dict[str, Any]:
     """Return engine arguments with rank-local AIC KV capacity materialized."""
 
     lowered = dict(raw)
+    timing = lowered.get("timing_model")
+    timing_systems_path = None
+    if (
+        isinstance(timing, dict)
+        and timing.get("type") == "external"
+        and timing.get("provider") == "aic"
+        and isinstance(timing.get("config"), dict)
+    ):
+        timing_systems_path = timing["config"].get("systems_path")
+
     attention_dp = lowered.get("aic_attention_dp_size")
     dp = attention_dp or 1
     configured_dp = lowered.get("dp_size") or 1
@@ -57,6 +67,11 @@ def materialize_aic_num_gpu_blocks(raw: dict[str, Any]) -> dict[str, Any]:
     if not model:
         raise ValueError("AIC KV cache capacity estimation requires aic_model_path in engine args")
 
+    capacity_systems_path = lowered.get("systems_path")
+    if timing_systems_path is not None:
+        # Keep capacity on the same database root as the authored AIC timing.
+        capacity_systems_path = timing_systems_path
+
     lowered["num_gpu_blocks"] = estimate_num_gpu_blocks(
         backend_name=backend,
         system=lowered.get("aic_system") or _DEFAULT_AIC_SYSTEM,
@@ -74,7 +89,11 @@ def materialize_aic_num_gpu_blocks(raw: dict[str, Any]) -> dict[str, Any]:
         gpu_memory_utilization=lowered.get("gpu_memory_utilization"),
         mem_fraction_static=lowered.get("mem_fraction_static"),
         free_gpu_memory_fraction=lowered.get("free_gpu_memory_fraction"),
-        backend_version=lowered.get("aic_backend_version"),
+        backend_version=(
+            lowered.get("aic_backend_version")
+            if lowered.get("aic_backend_version") is not None
+            else lowered.get("backend_version")
+        ),
         pp_size=(lowered.get("aic_pp_size") if lowered.get("aic_pp_size") is not None else 1),
         moe_tp_size=lowered.get("aic_moe_tp_size"),
         moe_ep_size=lowered.get("aic_moe_ep_size"),
@@ -84,7 +103,7 @@ def materialize_aic_num_gpu_blocks(raw: dict[str, Any]) -> dict[str, Any]:
         fmha_dtype=lowered.get("aic_fmha_dtype"),
         kv_cache_dtype=lowered.get("aic_kv_cache_dtype"),
         comm_dtype=lowered.get("aic_comm_dtype"),
-        systems_path=lowered.get("systems_path"),
+        systems_path=capacity_systems_path,
     )
     return lowered
 
@@ -129,6 +148,15 @@ def estimate_num_gpu_blocks(
         estimate_num_gpu_blocks as aic_estimate_num_gpu_blocks,
     )
 
+    if backend_version is None:
+        from aiconfigurator_core.sdk.perf_database import get_latest_database_version
+
+        # Use the maintained current slot (or the latest version in a custom
+        # legacy root), matching the database used for this capacity estimate.
+        backend_version = get_latest_database_version(system, backend_name, systems_paths=systems_path)
+        if backend_version is None:
+            raise ValueError(f"no perf database for system={system!r}, backend={backend_name!r}")
+
     if backend_name == "trtllm":
         memory_fraction_kind = "of_free"
         memory_fraction_value = (
@@ -148,9 +176,7 @@ def estimate_num_gpu_blocks(
             model_path,
             system,
             backend_name,
-            backend_version=(
-                backend_version if backend_version is not None else DEFAULT_BACKEND_VERSIONS[backend_name]
-            ),
+            backend_version=backend_version,
             scheduler_block_size=block_size,
             max_num_tokens=max_num_batched_tokens,
             max_batch_size=max_num_sequences,

python/aisimulate/src/aisimulate/runner.py

@@ -468,6 +468,42 @@ def _materialize_engine_role(
     num_gpu_blocks_is_explicit = False
     if "rank" not in role_config:
         num_gpu_blocks_is_explicit = role_config.get("num_gpu_blocks") is not None
+        for name in ("aic_backend_version", "backend_version"):
+            if role_config.get(name) is None:
+                role_config.pop(name, None)
+        # Propagate authored versions before capacity preflight. Otherwise a
+        # deployment pin reaches timing only after capacity has used current.
+        timing = role_config.get("timing_model")
+        timing_config = None
+        if (
+            isinstance(timing, dict)
+            and timing.get("type") == "external"
+            and timing.get("provider") == "aic"
+            and isinstance(timing.get("config"), dict)
+        ):
+            timing_config = timing["config"]
+        version = next(
+            (
+                value
+                for value in (
+                    role_config.get("aic_backend_version"),
+                    role_config.get("backend_version"),
+                    timing_config.get("backend_version") if timing_config is not None else None,
+                    deployment_backend_version or None,
+                )
+                if value is not None
+            ),
+            None,
+        )
+        if version is not None:
+            if (
+                not num_gpu_blocks_is_explicit
+                and role_config.get("aic_backend") is not None
+                and not any(name in role_config for name in ("aic_backend_version", "backend_version"))
+            ):
+                role_config["aic_backend_version"] = version
+            if timing_config is not None and timing_config.get("backend_version") is None:
+                role_config["timing_model"] = {**timing, "config": {**timing_config, "backend_version": version}}
         role_config = materialize_aic_num_gpu_blocks(role_config)
         capacity_materialized = role_config.get("num_gpu_blocks") is not None
     for name in ("engine_type", "aic_backend"):

tests/test_aic.py

@@ -15,8 +15,17 @@
 ]
 
 
+@pytest.mark.parametrize(
+    "version_fields",
+    [
+        {"backend_version": "test-version"},
+        {"aic_backend_version": "test-version"},
+        {"aic_backend_version": None, "backend_version": "test-version"},
+    ],
+)
 def test_materializer_sets_rank_local_capacity_without_forwarding_nextn(
     monkeypatch,
+    version_fields,
 ) -> None:
     calls = []
 
@@ -29,6 +38,7 @@ def estimate(**kwargs):
         {
             "engine_type": "vllm",
             "aic_backend": "vllm",
+            **version_fields,
             "aic_system": "h200_sxm",
             "aic_model_path": "test-model",
             "aic_attention_dp_size": 2,
@@ -43,6 +53,7 @@ def estimate(**kwargs):
     assert lowered["num_gpu_blocks"] == 46000
     assert lowered["dp_size"] == 2
     assert calls[0]["attention_dp_size"] == 2
+    assert calls[0]["backend_version"] == "test-version"
     assert calls[0]["pp_size"] == 3
     assert calls[0]["systems_path"] == "/tmp/custom-systems.yaml"
     assert calls[0]["max_num_sequences"] == 37
@@ -58,9 +69,16 @@ def estimate(*args, **kwargs):
         calls.append((args, kwargs))
         return 123
 
-    from aiconfigurator_core.sdk import memory
+    from aiconfigurator_core.sdk import memory, perf_database
 
     monkeypatch.setattr(memory, "estimate_num_gpu_blocks", estimate)
+    version_calls = []
+
+    def latest(system, backend, *, systems_paths):
+        version_calls.append((system, backend, systems_paths))
+        return "test-current"
+
+    monkeypatch.setattr(perf_database, "get_latest_database_version", latest)
     blocks = aic.estimate_num_gpu_blocks(
         backend_name="vllm",
         system="h200_sxm",
@@ -78,7 +96,8 @@ def estimate(*args, **kwargs):
     assert blocks == 123
     args, kwargs = calls[0]
     assert args == ("test-model", "h200_sxm", "vllm")
-    assert kwargs["backend_version"] == "0.19.0"
+    assert kwargs["backend_version"] == "test-current"
+    assert version_calls == [("h200_sxm", "vllm", "/tmp/custom-systems.yaml")]
     assert kwargs["memory_fraction_kind"] == "of_total"
     assert kwargs["memory_fraction_value"] == 0.9
     assert kwargs["pp_size"] == 3
@@ -100,6 +119,11 @@ def test_explicit_capacity_is_preserved_without_estimation(monkeypatch) -> None:
         "aic_backend": "vllm",
         "aic_model_path": "test-model",
         "num_gpu_blocks": 17,
+        "timing_model": {
+            "type": "external",
+            "provider": "aic",
+            "config": {"systems_paths": ["/path/that/does/not/exist"]},
+        },
     }
     assert aic.materialize_aic_num_gpu_blocks(raw) == raw
 
@@ -133,3 +157,94 @@ def estimate(**kwargs):
     assert calls[0]["gpu_memory_utilization"] == 0.0
     assert calls[0]["mem_fraction_static"] == 0.0
     assert calls[0]["free_gpu_memory_fraction"] == 0.0
+
+
+@pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"])
+@pytest.mark.parametrize("timing", ["default", "fixed", "polynomial"])
+def test_prediction_omitted_version_uses_current_database_for_capacity(monkeypatch, backend, timing):
+    from aiconfigurator_core.sdk import perf_database
+    from aisimulate.compiler import prediction_to_replay_spec
+    from aisimulate.config import CorePredictionConfig
+
+    expected = perf_database.get_version_slots("h200_sxm", backend)["current"]
+    get_database = perf_database.get_database
+    versions = []
+
+    def recorded_database(*args, **kwargs):
+        database = get_database(*args, **kwargs)
+        versions.append(database.version)
+        return database
+
+    monkeypatch.setattr(perf_database, "get_database", recorded_database)
+    timing_config = {"type": timing}
+    if timing == "fixed":
+        timing_config.update(prefill_ms=1, decode_ms=1)
+    config = CorePredictionConfig.model_validate(
+        {
+            "engine": {
+                "model": "Qwen/Qwen3-32B-FP8",
+                "hardware": "h200_sxm",
+                "backend": backend,
+                "context_length": 4096,
+                "workers": {"aggregated": {"parallelism": {"tensor": 2}, "timing": timing_config}},
+            }
+        }
+    )
+    deployment = prediction_to_replay_spec(config).backend_deployment
+    args = aic.materialize_aic_num_gpu_blocks(deployment.agg_engine_args)
+    assert args["num_gpu_blocks"] > 0
+    assert versions and set(versions) == {expected}
+    assert config.engine.backend_version is None
+    if timing != "default":
+        assert args["timing_model"]["type"] == timing
+
+
+@pytest.mark.parametrize("version", ["current", "next", "0.24.0", "0.19.0"])
+def test_capacity_preserves_explicit_backend_version(monkeypatch, version):
+    from aiconfigurator_core.sdk import memory, perf_database
+
+    monkeypatch.setattr(
+        perf_database,
+        "get_latest_database_version",
+        lambda *_args, **_kwargs: pytest.fail("explicit versions must not be replaced"),
+    )
+    versions = []
+
+    def estimate(*args, **kwargs):
+        versions.append(kwargs["backend_version"])
+        return 128
+
+    monkeypatch.setattr(memory, "estimate_num_gpu_blocks", estimate)
+    assert (
+        aic.estimate_num_gpu_blocks(
+            backend_name="vllm",
+            backend_version=version,
+            system="h200_sxm",
+            model_path="test-model",
+            tp_size=1,
+            block_size=64,
+            max_num_batched_tokens=4096,
+        )
+        == 128
+    )
+    assert versions == [version]
+
+
+def test_capacity_without_a_database_fails_before_estimation(monkeypatch):
+    from aiconfigurator_core.sdk import memory, perf_database
+
+    monkeypatch.setattr(perf_database, "get_latest_database_version", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(
+        memory,
+        "estimate_num_gpu_blocks",
+        lambda *_args, **_kwargs: pytest.fail("missing versions must fail before estimation"),
+    )
+    with pytest.raises(ValueError, match="no perf database.*unsupported-test-gpu.*vllm"):
+        aic.estimate_num_gpu_blocks(
+            backend_name="vllm",
+            system="unsupported-test-gpu",
+            model_path="test-model",
+            tp_size=1,
+            block_size=64,
+            max_num_batched_tokens=4096,
+        )