refactor(replay): delegate Weka ingestion to AISimulate

.github/workflows/pre-merge.yml

@@ -467,7 +467,6 @@ jobs:
         compile_bench multiturn --bench multiturn_bench
         compile_bench satf --bin request_trace_to_satf
         compile_bench request-trace-to-mooncake --bin request_trace_to_mooncake
-        compile_bench weka-to-agentic-mooncake --bin weka_to_agentic_mooncake
         compile_bench offline-replay --bench offline_replay_bench
         compile_bench mooncake --bench mooncake_bench
         compile_bench active-sequences --bench active_sequences_bench
@@ -477,9 +476,6 @@ jobs:
         cargo test --locked --no-default-features -p dynamo-bench \
           --features request-trace-to-mooncake --test request_trace_to_mooncake
 
-        cargo test --locked --no-default-features -p dynamo-bench \
-          --features weka-to-agentic-mooncake --test weka_to_agentic_mooncake
-
         cargo test --locked -p dynamo-llm --no-default-features \
           --features kv-router-stress --bench kv_router_bench --no-run
 
@@ -551,8 +547,6 @@ jobs:
         lint_bench multiturn --bench multiturn_bench
         lint_bench satf --bin request_trace_to_satf
         lint_bench request-trace-to-mooncake --bin request_trace_to_mooncake
-        lint_bench weka-to-agentic-mooncake --bin weka_to_agentic_mooncake
-        lint_bench weka-to-agentic-mooncake --test weka_to_agentic_mooncake
         lint_bench offline-replay --bench offline_replay_bench
         lint_bench mooncake --bench mooncake_bench
         lint_bench active-sequences --bench active_sequences_bench

Cargo.toml

@@ -55,8 +55,9 @@ repository = "https://github.com/ai-dynamo/dynamo.git"
 keywords = ["llm", "genai", "inference", "nvidia", "distributed"]
 
 [workspace.dependencies]
-# Published separately by https://github.com/ai-dynamo/aisimulate.
-aisimulate-core = { version = "=0.12.0-dev.1" }
+# TODO: replace this public review revision with the matching exact crates.io
+# release once it exports the Weka graph-loading API.
+aisimulate-core = { git = "https://github.com/ai-dynamo/aisimulate.git", rev = "d71203e489541bc98c5f4a616dcec3b225b8fb72" }
 
 # Local crates
 dynamo-runtime = { path = "lib/runtime", version = "1.5.0" }

benchmarks/pyproject.toml

@@ -39,7 +39,9 @@ classifiers = [
 ]
 
 dependencies = [
-    "aisimulate==0.12.0.dev1; python_version >= '3.11' and python_version < '3.14'",
+    # TODO: replace this public review revision with the matching exact PyPI
+    # release once it exports the Weka graph-loading API.
+    "aisimulate @ git+https://github.com/ai-dynamo/aisimulate.git@d71203e489541bc98c5f4a616dcec3b225b8fb72#subdirectory=python/aisimulate ; python_version >= '3.11' and python_version < '3.14'",
     "aiperf==0.12.0",
     "matplotlib",
     "networkx",

components/src/dynamo/planner/core/perf_model/engine_query.py

@@ -218,9 +218,11 @@ def best_available(
         attention_dp_size: int,
     ) -> AicCoreEnginePerfModel:
         if aic_config is None:
-            model = AicForwardPassPerfModel.from_regression(options)
+            model = AicForwardPassPerfModel.from_regression(worker_type, options)
         else:
-            model = AicForwardPassPerfModel.best_available(aic_config, options)
+            model = AicForwardPassPerfModel.best_available(
+                aic_config, worker_type, options
+            )
         return cls(
             model=model,
             worker_type=worker_type,

components/src/dynamo/planner/tests/unit/test_engine_query.py

@@ -93,11 +93,13 @@ def test_best_available_uses_aic_core_wheel_facade(monkeypatch):
 
     class _FakeAicFacade:
         last_config = None
+        last_worker_type = None
         last_options = None
 
         @classmethod
-        def best_available(cls, config, options):
+        def best_available(cls, config, worker_type, options):
             cls.last_config = config
+            cls.last_worker_type = worker_type
             cls.last_options = options
             return sentinel
 
@@ -122,6 +124,7 @@ def best_available(cls, config, options):
     )
 
     assert _FakeAicFacade.last_config is config
+    assert _FakeAicFacade.last_worker_type == "prefill"
     assert _FakeAicFacade.last_options is options
     assert model.diagnostics()["readiness"] == "ready"
 

components/src/dynamo/replay/api.py

@@ -15,6 +15,8 @@
 from dynamo._core import run_mocker_trace_replay as _run_mocker_trace_replay
 from dynamo.replay.report import PlannerReplayDetails, ReplayReport
 
+_AGENTIC_MODEL_PROJECTION_POLICY = "project_to_configured_target"
+
 
 def _planner_replay_adapter():
     """Load Planner replay lazily to break the replay.api/mocker import cycle.
@@ -55,6 +57,7 @@ class _CommonReplayOptions(TypedDict, total=False):
 
 class _TraceReplayOptions(_CommonReplayOptions, total=False):
     agentic_lanes: int | None
+    execution_model: str | None
     trace_block_size: int | None
     trace_format: str
     trace_shared_prefix_ratio: float
@@ -91,15 +94,57 @@ def _materialize_offline_report(
     native,
     *,
     planner: PlannerReplayDetails | None,
+    execution_model: str | None = None,
 ) -> ReplayReport:
+    summary = dict(native.summary)
+    _add_agentic_model_projection(summary, execution_model)
     return ReplayReport(
-        summary=native.summary,
+        summary=summary,
         per_request=native.per_request,
         coverage=native.coverage,
         planner=planner,
     )
 
 
+def _normalize_execution_model(
+    trace_format: str, execution_model: str | None
+) -> str | None:
+    if execution_model is not None:
+        if not isinstance(execution_model, str):
+            raise TypeError("execution_model must be a string or None")
+        execution_model = execution_model.strip()
+        if not execution_model:
+            raise ValueError("execution_model must be non-empty")
+    if (
+        trace_format in {"weka", "agentic_mooncake", "agentic-mooncake"}
+        and execution_model is None
+    ):
+        raise ValueError("agentic execution requires a configured target model")
+    return execution_model
+
+
+def _add_agentic_model_projection(
+    summary: dict[str, Any], execution_model: str | None
+) -> None:
+    graph = summary.get("agentic_graph")
+    if not isinstance(graph, dict):
+        return
+    if execution_model is None:
+        raise ValueError(
+            "agentic execution did not declare its configured target model"
+        )
+    source_models = graph.get("source_models")
+    if not isinstance(source_models, list) or not all(
+        isinstance(model, str) and model for model in source_models
+    ):
+        raise ValueError("agentic graph did not report valid source_models")
+    summary["agentic_model_projection"] = {
+        "policy": _AGENTIC_MODEL_PROJECTION_POLICY,
+        "source_models": source_models,
+        "target_model": execution_model,
+    }
+
+
 @overload
 def run_trace_replay(
     trace_files,
@@ -161,6 +206,7 @@ def run_trace_replay(
     benchmark_granularity=8,
     capture_per_request=False,
     capture_planner_details=True,
+    execution_model=None,
 ) -> ReplayReport | dict[str, Any]:
     """Run trace replay.
 
@@ -171,6 +217,7 @@ def run_trace_replay(
         agentic_lanes is not None and not isinstance(agentic_lanes, int)
     ):
         raise TypeError("agentic_lanes must be an integer or None")
+    execution_model = _normalize_execution_model(trace_format, execution_model)
     trace_files = _normalize_trace_files(trace_files)
     replay_kwargs = {
         "extra_engine_args": extra_engine_args,
@@ -193,6 +240,7 @@ def run_trace_replay(
         "report_jsonl_path": report_jsonl_path,
         "max_sim_time_ms": max_sim_time_ms,
         "model_name": model_name,
+        "execution_model": execution_model,
         "sla_ttft_ms": sla_ttft_ms,
         "sla_itl_ms": sla_itl_ms,
         "sla_e2e_ms": sla_e2e_ms,
@@ -249,15 +297,22 @@ def run_trace_replay(
             return _materialize_offline_report(
                 native,
                 planner=adapter.finalize(native.lifecycle_operations),
+                execution_model=execution_model,
             )
     result = _run_mocker_trace_replay(
         trace_files,
         **replay_kwargs,
         scaling_policy=None,
     )
     if replay_mode == "online":
+        if isinstance(result, dict):
+            _add_agentic_model_projection(result, execution_model)
         return result
-    return _materialize_offline_report(result, planner=None)
+    return _materialize_offline_report(
+        result,
+        planner=None,
+        execution_model=execution_model,
+    )
 
 
 @overload