refactor(replay): unify planner runtime boundary [DYN-3850]

components/src/dynamo/mocker/__init__.py

@@ -19,7 +19,6 @@
 
 try:
     from dynamo._core import MockEngineArgs as MockEngineArgs
-    from dynamo._core import PlannerReplayBridge as PlannerReplayBridge
     from dynamo._core import ReasoningConfig as ReasoningConfig
     from dynamo._core import SglangArgs as SglangArgs
     from dynamo._core import TrtllmArgs as TrtllmArgs
@@ -32,7 +31,6 @@
     __all__.extend(
         [
             "MockEngineArgs",
-            "PlannerReplayBridge",
             "ReasoningConfig",
             "SglangArgs",
             "TrtllmArgs",

components/src/dynamo/planner/offline/replay_adapter.py

@@ -1,15 +1,14 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Adapter that drives the planner core from the unified replay loop.
+"""Scaling-policy adapter for the unified offline replay loop.
 
-The Rust offline simulation (``PlannerReplayBridge.run``) owns the drive
-loop and calls back into this adapter once per ``PlannerTick`` event, so
-the adapter is a callback hook rather than an external stepper:
+The public replay entrypoint passes this adapter into the Rust runtime as
+an optional scaling component. Rust owns the drive loop:
 
-    Bridge.run(adapter)                        # Rust owns the loop
+    run_replay(..., scaling_policy=adapter)    # Rust owns the loop
       adapter.initial_tick_ms()      -> first tick time
-      per PlannerTick:
+      per ScalingTick:
         adapter.on_tick(metrics)     -> _build_tick_input() -> TickInput
                                         EngineProtocol.tick() -> PlannerEffects
                                         -> {target_prefill, target_decode, next_tick_ms}
@@ -22,8 +21,7 @@
 ``PlannerEffects.scale_to`` replay contract while using plugin-aware
 observability (Prometheus metrics, audit events, diagnostics).
 
-The simulation steps itself — replay no longer drives the bridge
-externally. Async orchestrator calls (``bootstrap_from_fpms`` / ``tick``)
+Async orchestrator calls (``bootstrap_from_fpms`` / ``tick``)
 run inside a single replay-scoped event loop so callers don't change.
 
 Supports both aggregated and disaggregated topologies. No I/O, no
@@ -35,7 +33,7 @@
 import asyncio
 import logging
 from dataclasses import dataclass, field
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 from dynamo.common.forward_pass_metrics import (
     ForwardPassMetrics,
@@ -85,7 +83,7 @@ class ReplayPlannerReport:
 
 
 def _build_fpm_from_dict(d: dict[str, Any]) -> ForwardPassMetrics:
-    """Convert a bridge FPM snapshot dict into a ForwardPassMetrics struct."""
+    """Convert a replay FPM snapshot dict into a ForwardPassMetrics struct."""
     return ForwardPassMetrics(
         worker_id=str(d["worker_id"]),
         dp_rank=int(d.get("dp_rank", 0)),
@@ -193,33 +191,22 @@ def _weighted(key: str, wa: float, ww: float) -> float:
 
 
 class ReplayPlannerAdapter:
-    """Drives the plugin planner using the PlannerReplayBridge.
-
-    Supports both ``mode="agg"`` and ``mode="disagg"``.
-    """
+    """Context-managed planner scaling policy for offline replay."""
 
     def __init__(
         self,
         planner_config: PlannerConfig,
-        bridge: Any,  # PlannerReplayBridge (Rust pyclass)
+        engine: EngineProtocol,
         capabilities: Optional[WorkerCapabilities] = None,
         warmup_observations: Optional[list[TrafficObservation]] = None,
     ) -> None:
         self._config = planner_config
-        self._bridge = bridge
         self._capabilities = capabilities
         self._is_disagg = planner_config.mode == "disagg"
 
-        self._engine: EngineProtocol
+        self._engine = engine
         self._warmup_observations = list(warmup_observations or [])
         self._orchestrator_bootstrapped = False
-        # Inject a ``VirtualClock`` so plugin scheduler / circuit breaker /
-        # HOLD_LAST cache see trace time, not real wall-clock.
-        self._engine = OrchestratorEngineAdapter(
-            planner_config,
-            capabilities or WorkerCapabilities(),
-            clock=VirtualClock(),
-        )
         # Replay's ``run()`` is synchronous; we own a scoped event loop to
         # drive the async engine calls without forcing callers to use
         # ``asyncio.run``.
@@ -253,24 +240,36 @@ def __init__(
         # installs benchmark FPMs after adapter construction and before the
         # first tick.
 
+    def __enter__(self) -> ReplayPlannerAdapter:
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> Literal[False]:
+        try:
+            self.close()
+        except BaseException:
+            if exc_value is None:
+                raise
+            logger.exception("planner replay cleanup failed")
+        return False
+
     # ------------------------------------------------------------------
     # Sync/async bridging
     # ------------------------------------------------------------------
 
     def _run_sync(self, coro):
         """Run a coroutine on the replay-owned event loop. Used to call
         the orchestrator path's async APIs from replay's sync surface."""
-        assert self._loop is not None, "sync bridge only available on orchestrator path"
+        assert self._loop is not None, "sync execution requires an active replay scope"
         return self._loop.run_until_complete(coro)
 
     def _bootstrap_orchestrator_if_needed(self) -> None:
         if self._orchestrator_bootstrapped:
             return
-        self._run_sync(
-            self._engine.bootstrap_plugins(  # type: ignore[attr-defined]
-                historical_traffic=self._warmup_observations or None
+        bootstrap_plugins = getattr(self._engine, "bootstrap_plugins", None)
+        if bootstrap_plugins is not None:
+            self._run_sync(
+                bootstrap_plugins(historical_traffic=self._warmup_observations or None)
             )
-        )
         self._orchestrator_bootstrapped = True
 
     def install_benchmark_fpms(
@@ -295,10 +294,8 @@ def install_benchmark_fpms(
         )
 
     # ------------------------------------------------------------------
-    # Inverted drive: the Rust ``PlannerReplayBridge.run(self)`` owns the loop and
-    # calls ``initial_tick_ms`` once then ``on_tick`` per ``PlannerTick`` event. The
-    # entrypoint wraps the returned trace_report via ``finalize``. (Replaces the old
-    # Python while-loop that drove ``bridge.advance_to`` + ``bridge.apply_scaling``.)
+    # Rust calls ``initial_tick_ms`` once and ``on_tick`` for each ``ScalingTick``.
+    # The entrypoint wraps the returned trace report via ``finalize``.
     # ------------------------------------------------------------------
 
     def start(self) -> None:
@@ -310,13 +307,13 @@ def start(self) -> None:
         self._total_ticks = 0
 
     def initial_tick_ms(self) -> float:
-        """First tick time in milliseconds (called by the Rust PlannerHook)."""
+        """First tick time in milliseconds."""
         if not self._orchestrator_bootstrapped or not hasattr(self, "_pending_tick"):
             self.start()
         return self._pending_tick.at_s * 1000.0
 
     def on_tick(self, result: dict[str, Any]) -> dict[str, Any]:
-        """Drive one planner tick from the bridge's metrics dict. Returns the scaling
+        """Drive one planner tick from the runtime snapshot. Returns the scaling
         decision (absolute targets, ``None`` = unchanged) + the next tick time in ms
         (``None`` = stop) for the Rust loop to apply and re-arm."""
         tick = self._pending_tick
@@ -328,17 +325,20 @@ def on_tick(self, result: dict[str, Any]) -> dict[str, Any]:
         self._total_ticks += 1
         self._record_diagnostics(tick_input, effects, result, emit_diagnostics)
 
-        # Clear scaling targets once active counts match.
-        active_p = result["active_prefill_count"]
-        active_d = result["active_decode_count"]
+        current_p = result.get(
+            "non_draining_prefill_count", result["active_prefill_count"]
+        )
+        current_d = result.get(
+            "non_draining_decode_count", result["active_decode_count"]
+        )
         if (
             self._scaling_target_prefill is not None
-            and active_p == self._scaling_target_prefill
+            and current_p == self._scaling_target_prefill
         ):
             self._scaling_target_prefill = None
         if (
             self._scaling_target_decode is not None
-            and active_d == self._scaling_target_decode
+            and current_d == self._scaling_target_decode
         ):
             self._scaling_target_decode = None
 
@@ -361,24 +361,18 @@ def on_tick(self, result: dict[str, Any]) -> dict[str, Any]:
         }
 
     def finalize(self, trace_report: dict[str, Any]) -> ReplayPlannerReport:
-        """Assemble the enriched report from accumulated planner state. Called by the
-        entrypoint after the Rust ``run()`` returns the trace_report dict."""
-        try:
-            html_report_path = self._recorder.finalize()
-            return ReplayPlannerReport(
-                trace_report=trace_report,
-                scaling_events=self._scaling_events,
-                diagnostics_log=self._diagnostics_log,
-                total_ticks=self._total_ticks,
-                html_report_path=html_report_path,
-            )
-        finally:
-            self.close()
+        """Assemble the enriched report after successful replay execution."""
+        html_report_path = self._recorder.finalize()
+        return ReplayPlannerReport(
+            trace_report=trace_report,
+            scaling_events=self._scaling_events,
+            diagnostics_log=self._diagnostics_log,
+            total_ticks=self._total_ticks,
+            html_report_path=html_report_path,
+        )
 
     def close(self) -> None:
-        """Shut down the engine and the replay-scoped event loop. Idempotent so it
-        runs cleanly from both ``finalize`` (success) and the entrypoint's error
-        path (when ``bridge.run`` raises before ``finalize`` is reached)."""
+        """Shut down the engine and replay-scoped event loop. Idempotent."""
         loop = self._loop
         if loop is None:
             return
@@ -440,14 +434,18 @@ def _compute_scale_decision(
     ) -> tuple[Optional[int], Optional[int]]:
         """Compute the (prefill, decode) absolute scale targets and record the scaling
         event. Returns ``(None, None)`` for a no-op. The Rust loop applies the targets,
-        so this no longer calls ``bridge.apply_scaling``."""
+        so this method only records the requested transition."""
         scale = effects.scale_to
         if scale is None:
             raise ValueError(
                 "_compute_scale_decision requires effects.scale_to to be set"
             )
-        current_p = result["active_prefill_count"]
-        current_d = result["active_decode_count"]
+        current_p = result.get(
+            "non_draining_prefill_count", result["active_prefill_count"]
+        )
+        current_d = result.get(
+            "non_draining_decode_count", result["active_decode_count"]
+        )
         target_p = scale.num_prefill if scale.num_prefill is not None else current_p
         target_d = scale.num_decode if scale.num_decode is not None else current_d
 
@@ -522,25 +520,27 @@ def _is_easy_mode(self) -> bool:
     def _build_tick_input(
         self, tick: ScheduledTick, result: dict[str, Any]
     ) -> TickInput:
-        """Convert bridge result dict to planner TickInput."""
-        # Keep planner cadence on the scheduled replay clock. The Rust bridge
-        # also advances idle gaps to this timestamp so traffic windows drain
+        """Convert the Rust scaling snapshot to planner ``TickInput``."""
+        # Keep planner cadence on the scheduled replay clock. Rust also
+        # advances idle gaps to this timestamp so traffic windows drain
         # with the same duration the planner sees.
         now_s = tick.at_s
 
         worker_counts = None
         if tick.need_worker_states:
             active_p = result["active_prefill_count"]
             active_d = result["active_decode_count"]
+            current_p = result.get("non_draining_prefill_count", active_p)
+            current_d = result.get("non_draining_decode_count", active_d)
             expected_p = (
                 self._scaling_target_prefill
                 if self._scaling_target_prefill is not None
-                else active_p
+                else current_p
             )
             expected_d = (
                 self._scaling_target_decode
                 if self._scaling_target_decode is not None
-                else active_d
+                else current_d
             )
             worker_counts = WorkerCounts(
                 ready_num_prefill=active_p if self._is_disagg else None,
@@ -550,11 +550,11 @@ def _build_tick_input(
                 prefill_scaling_in_progress=(
                     self._is_disagg
                     and self._scaling_target_prefill is not None
-                    and self._scaling_target_prefill != active_p
+                    and self._scaling_target_prefill != current_p
                 ),
                 decode_scaling_in_progress=(
                     self._scaling_target_decode is not None
-                    and self._scaling_target_decode != active_d
+                    and self._scaling_target_decode != current_d
                 ),
             )
 
@@ -584,7 +584,7 @@ def _build_tick_input(
                 decode=decode_dict,
             )
 
-        # The Rust bridge drains the per-tick traffic window into ``result["traffic"]``;
+        # Rust drains the per-tick traffic window into ``result["traffic"]``;
         # accumulate it so a need_traffic_metrics tick sees the full window since the
         # last consumed one (the planner consumes traffic only on throughput ticks).
         tick_traffic = result.get("traffic")
@@ -632,3 +632,22 @@ def _build_tick_input(
             worker_counts=worker_counts,
             fpm_observations=fpm_observations,
         )
+
+
+def create_replay_planner_adapter(
+    planner_config: PlannerConfig,
+    capabilities: Optional[WorkerCapabilities] = None,
+    warmup_observations: Optional[list[TrafficObservation]] = None,
+) -> ReplayPlannerAdapter:
+    """Create a replay adapter backed by the builtin planner orchestrator."""
+    engine = OrchestratorEngineAdapter(
+        planner_config,
+        capabilities or WorkerCapabilities(),
+        clock=VirtualClock(),
+    )
+    return ReplayPlannerAdapter(
+        planner_config=planner_config,
+        engine=engine,
+        capabilities=capabilities,
+        warmup_observations=warmup_observations,
+    )

components/src/dynamo/planner/tests/offline/test_replay_adapter_fpm.py

@@ -50,7 +50,7 @@ def _agg_config_sla() -> PlannerConfig:
 
 
 def _snap(worker_id: str, wall_time: float, dp_rank: int = 0) -> dict:
-    """A bridge FPM snapshot dict with every key ``_build_fpm_from_dict`` reads."""
+    """A replay FPM snapshot dict with every key ``_build_fpm_from_dict`` reads."""
     return {
         "worker_id": worker_id,
         "dp_rank": dp_rank,
@@ -110,6 +110,33 @@ def _orch_agg_config_sla() -> PlannerConfig:
     return _agg_config_sla()
 
 
+def test_replay_adapter_uses_injected_engine_protocol_and_owns_cleanup():
+    class _Engine:
+        def __init__(self):
+            self.closed = False
+
+        def initial_tick(self, start_s):
+            return ScheduledTick(at_s=start_s + 5.0)
+
+        async def tick(self, scheduled_tick, tick_input):
+            raise AssertionError("tick is not needed by this ownership test")
+
+        async def shutdown(self):
+            self.closed = True
+
+    engine = _Engine()
+    adapter = ReplayPlannerAdapter(
+        PlannerConfig(mode="agg"),
+        capabilities=_agg_caps(),
+        engine=engine,
+    )
+
+    with adapter:
+        assert adapter.initial_tick_ms() == 5_000.0
+
+    assert engine.closed
+
+
 def test_install_benchmark_fpms_installs_regression_on_orchestrator_path():
     """Review #3: the orchestrator replay path must actually install
     regressions. ``ReplayPlannerAdapter.install_benchmark_fpms`` routes to