diff --git a/components/src/dynamo/replay/api.py b/components/src/dynamo/replay/api.py index 0ac237c70b99..baf1a285538a 100644 --- a/components/src/dynamo/replay/api.py +++ b/components/src/dynamo/replay/api.py @@ -97,7 +97,7 @@ class _TraceReplayOptions(_CommonReplayOptions, total=False): trace_num_prefix_groups: int report_jsonl_path: str | os.PathLike[str] | None max_sim_time_ms: float | None - # Internal to retune sweeps; not part of the public API. + # Internal; not part of the public API. _planner_warmup_observations: "Sequence[TrafficObservation] | None" _deterministic: bool @@ -420,8 +420,9 @@ def run_trace_replay( # Internal: a retune sweep's inner planner is warmed with the history it # already bucketed instead of its load_predictor_warmup_trace. _planner_warmup_observations=None, - # Internal: a retune sweep replays the baseline and every candidate with - # the canonical router seed and ordinal request ids (common random numbers). + # Internal: replay with the canonical router seed and ordinal request ids + # (common random numbers). Retune sweeps and every trace replay through + # the Sweeper's Dynamo runner set it. _deterministic=False, ) -> ReplayReport | dict[str, Any]: """Run trace replay. diff --git a/components/src/dynamo/replay/simulation.py b/components/src/dynamo/replay/simulation.py index 5b1a5bd49936..251e033199b4 100644 --- a/components/src/dynamo/replay/simulation.py +++ b/components/src/dynamo/replay/simulation.py @@ -409,6 +409,9 @@ def _run_trace( # embedded value instead of imposing the synthetic-workload default. if trace_format not in {"weka", "dynamo"} and trace_block_size is None: trace_block_size = self.trace_block_size + # Trace replays use ordinal request ids and the canonical router seed, + # so a candidate's score depends only on its configuration and the + # trace. if deployment.deployment_mode == "agg": return run_trace_replay( trace_files=trace_files, @@ -423,6 +426,7 @@ def _run_trace( execution_model=execution_model, extra_engine_args=self._engine_args(deployment, "aggregated"), num_workers=deployment.num_workers, + _deterministic=True, **common, ) return run_trace_replay( @@ -440,6 +444,7 @@ def _run_trace( decode_engine_args=self._engine_args(deployment, "decode"), num_prefill_workers=deployment.num_prefill_workers, num_decode_workers=deployment.num_decode_workers, + _deterministic=True, **common, ) diff --git a/components/src/dynamo/replay/tests/test_simulation.py b/components/src/dynamo/replay/tests/test_simulation.py index 4c7c7b3b7cea..23725992c6d5 100644 --- a/components/src/dynamo/replay/tests/test_simulation.py +++ b/components/src/dynamo/replay/tests/test_simulation.py @@ -97,7 +97,35 @@ def _agg_deployment() -> BackendDeploymentSpec: ) -def test_trace_runner_preserves_current_replay_arguments(monkeypatch) -> None: +def _disagg_deployment() -> BackendDeploymentSpec: + return BackendDeploymentSpec( + deployment_mode="disagg", + backend="vllm", + backend_version="0.11.0", + prefill_engine_args={ + "worker_type": "prefill", + "timing_model": {"type": "polynomial"}, + }, + decode_engine_args={ + "worker_type": "decode", + "timing_model": {"type": "polynomial"}, + }, + num_prefill_workers=1, + num_decode_workers=2, + ) + + +@pytest.mark.parametrize( + ("deployment", "workers"), + [ + (_agg_deployment(), {"num_workers": 3}), + (_disagg_deployment(), {"num_prefill_workers": 1, "num_decode_workers": 2}), + ], + ids=["agg", "disagg"], +) +def test_trace_runner_preserves_current_replay_arguments( + monkeypatch, deployment, workers +) -> None: seen = {} def fake_run_trace_replay(**kwargs): @@ -113,7 +141,7 @@ def fake_run_trace_replay(**kwargs): monkeypatch.setattr(simulation, "KvRouterConfig", _FakeRouterConfig) monkeypatch.setattr(simulation, "run_trace_replay", fake_run_trace_replay) spec = ReplaySpec( - backend_deployment=_agg_deployment(), + backend_deployment=deployment, workload={ "trace_path": "tiny.jsonl", "trace_format": "dynamo", @@ -159,7 +187,8 @@ def fake_run_trace_replay(**kwargs): assert seen["trace_files"] == "tiny.jsonl" assert seen["trace_format"] == "dynamo" - assert seen["num_workers"] == 3 + assert {name: seen[name] for name in workers} == workers + assert seen["_deterministic"] is True assert seen["router_mode"] == "kv_router" assert seen["planner_config"] == {"mode": "agg"} assert seen["arrival_speedup_ratio"] == 2.0