From 1035c24a6e117ebcee99d717cc56ffa7c7e72082 Mon Sep 17 00:00:00 2001 From: Sachal Malick Date: Mon, 31 Aug 2026 05:01:14 -0400 Subject: [PATCH] feat(planner): schedule batch work on native ticks --- components/src/dynamo/planner/core/base.py | 37 +- .../planner/plugins/builtins/observe.py | 34 +- .../plugins/orchestrator/engine_adapter.py | 267 ++++++++++++- .../orchestrator/test_engine_adapter.py | 357 +++++++++++++++++- .../test_environment_observe_plugin.py | 66 +++- .../planner/tests/unit/test_tick_lifecycle.py | 148 +++++++- 6 files changed, 886 insertions(+), 23 deletions(-) diff --git a/components/src/dynamo/planner/core/base.py b/components/src/dynamo/planner/core/base.py index 8dfe06cfb117..446f64a3e452 100644 --- a/components/src/dynamo/planner/core/base.py +++ b/components/src/dynamo/planner/core/base.py @@ -564,6 +564,7 @@ async def _gather_tick_input(self, tick: ScheduledTick) -> TickInput: traffic = None worker_counts = None fpm_obs = None + batch = None if tick.need_traffic_metrics: if tick.use_full_traffic_metrics: @@ -576,11 +577,22 @@ async def _gather_tick_input(self, tick: ScheduledTick) -> TickInput: worker_counts = await self._collect_worker_counts() if tick.need_worker_fpm: fpm_obs = self._collect_fpm() + if tick.need_batch_scheduling: + try: + batch = await self.environment.collect_batch_scheduling() + except asyncio.CancelledError: + raise + except Exception: + logger.exception( + "Batch scheduling observation failed; continuing fail-closed" + ) + now = time.time() return TickInput( now_s=now, traffic=traffic, worker_counts=worker_counts, fpm_observations=fpm_obs, + batch=batch, ) async def _observe_tick( @@ -803,6 +815,7 @@ def _should_emit_tick_diagnostics( tick.run_load_scaling or tick.run_throughput_scaling or effects.scale_to is not None + or bool(effects.batch_drain_limits) or bool(diag.audit_events) or bool(diag.short_circuit_reason) ) @@ -822,7 +835,29 @@ async def _run_one_tick( # lock is released before connector rollouts that may take minutes. async with self._config_lock: effects = await engine.tick(tick, tick_input) - await self._apply_effects(effects) + batch_actuation_succeeded = True + if effects.batch_drain_limits and not self.config.advisory: + try: + # Apply the leased admission decision before any replica + # mutation. If Redis is unavailable, retain the current fleet; + # the previous lease expires fail-closed and the next tick + # retries without terminating the Planner loop. + await self.environment.apply_batch_drain_limits( + effects.batch_drain_limits + ) + except asyncio.CancelledError: + raise + except Exception as exc: + batch_actuation_succeeded = False + logger.exception( + "Batch drain-limit actuation failed; skipping replica " + "actuation for this tick" + ) + effects.diagnostics.audit_events.append( + f"batch_drain_actuation_failed:{type(exc).__name__}" + ) + if batch_actuation_succeeded: + await self._apply_effects(effects) emit_diagnostics = self._should_emit_tick_diagnostics(tick, effects) if emit_diagnostics: self._report_diagnostics(tick, effects.diagnostics) diff --git a/components/src/dynamo/planner/plugins/builtins/observe.py b/components/src/dynamo/planner/plugins/builtins/observe.py index 2db4c7fd8f39..b43bf009ec86 100644 --- a/components/src/dynamo/planner/plugins/builtins/observe.py +++ b/components/src/dynamo/planner/plugins/builtins/observe.py @@ -5,12 +5,18 @@ from __future__ import annotations +import asyncio +import logging +import time +from collections.abc import Callable from dataclasses import dataclass from typing import Protocol from dynamo.planner.core.types import ScheduledTick, TickInput, WorkerCounts from dynamo.planner.environment.interface import PlannerEnvironment +logger = logging.getLogger(__name__) + @dataclass class ObserveStageRequest: @@ -47,16 +53,19 @@ def __init__( *, require_prefill: bool, require_decode: bool, + clock: Callable[[], float] = time.time, ) -> None: self.environment = environment self.require_prefill = require_prefill self.require_decode = require_decode + self._clock = clock async def Observe(self, request: ObserveStageRequest) -> ObserveStageResponse: tick = request.scheduled_tick traffic = None worker_counts = None fpm_observations = None + batch = None if tick.need_traffic_metrics: if tick.use_full_traffic_metrics: @@ -69,13 +78,36 @@ async def Observe(self, request: ObserveStageRequest) -> ObserveStageResponse: worker_counts = self._collect_worker_counts() if tick.need_worker_fpm: fpm_observations = self.environment.collect_fpm() + if tick.need_batch_scheduling: + try: + batch = await self.environment.collect_batch_scheduling() + except asyncio.CancelledError: + raise + except Exception: + # Batch policy treats a missing snapshot as an explicit safety + # pause. Keep unrelated worker/traffic observations available + # and let the native tick renew a zero-rate lease. + logger.exception( + "Batch scheduling observation failed; continuing fail-closed" + ) + + # Batch source timestamps are sampled during collection. Anchor policy + # freshness and lease expiry after that I/O so a valid observation can + # never appear to come from the future and the full lease remains for + # the decision/actuation phases. + tick_now_s = ( + max(request.now_s, self._clock()) + if tick.need_batch_scheduling + else request.now_s + ) return ObserveStageResponse( tick_input=TickInput( - now_s=request.now_s, + now_s=tick_now_s, traffic=traffic, worker_counts=worker_counts, fpm_observations=fpm_observations, + batch=batch, ) ) diff --git a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py index 94c4a7d58fd4..b03d47b3765f 100644 --- a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py +++ b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py @@ -34,7 +34,9 @@ Extracts ``traffic`` into ``TrafficMetrics``, ``worker_counts`` (counts + scaling-in-progress flags) into ``WorkerState``, and per-engine FPM observations into ``FpmData`` (msgspec/msgpack- - encoded, keyed by ``"/"``) on ``ObservationData``. + encoded, keyed by ``"/"``), plus generic batch + scheduling observations into ``BatchSchedulingData`` on + ``ObservationData``. External plugins declaring ``needs=["observations.fpm"]`` receive the FPM map; an empty/absent submap means "no FPM this tick". 3. **FPM regression observation**: @@ -68,6 +70,7 @@ if TYPE_CHECKING: import grpc.aio +from dynamo.planner.core.batch_policy import plan_batch_schedule from dynamo.planner.core.budget import ( apply_power_budget, proportional_clamp_pair, @@ -75,6 +78,8 @@ ) from dynamo.planner.core.state_machine import PlannerScalingState from dynamo.planner.core.types import ( + BatchDrainLimitDecision, + BatchSchedulingObservation, FpmObservations, PlannerEffects, ScalingDecision, @@ -103,12 +108,22 @@ from dynamo.planner.plugins.scheduler import PluginScheduler from dynamo.planner.plugins.transport.config import make_transport_for_endpoint from dynamo.planner.plugins.types import ( + BatchDispatcherFeedback as BatchDispatcherFeedbackData, +) +from dynamo.planner.plugins.types import ( + BatchJobDemand as BatchJobDemandData, +) +from dynamo.planner.plugins.types import ( + BatchSchedulingData, FpmData, ObservationData, PipelineContext, TrafficMetrics, WorkerState, ) +from dynamo.planner.plugins.types import ( + PoolTrafficDemand as PoolTrafficDemandData, +) log = logging.getLogger(__name__) @@ -175,6 +190,14 @@ def __init__( self._last_tick_monotonic: float = 0.0 self._last_load_loop_monotonic: float = 0.0 self._last_throughput_loop_monotonic: float = 0.0 + + # Batch scheduling is a native planner decision made on every pipeline + # tick when enabled, independent of whether any plugin is due. + # ``getattr`` plus the strict bool check keeps legacy tests/config + # mocks (which predate the config subtree) disabled instead of + # accidentally enabling the path through a truthy ``MagicMock``. + batch_config = getattr(config, "batch_scheduling", None) + self._batch_scheduling_enabled = getattr(batch_config, "enabled", False) is True # Emit one rollout-hold warning per continuous mid-rollout stretch; # reset when the deployment is stable again so the next rollout warns. self._power_rollout_hold_warned: bool = False @@ -580,6 +603,22 @@ async def tick( # 3. Build PipelineContext + baseline and drive the orchestrator. ctx = self._tick_input_to_context(tick_input) + batch_plan = None + if getattr(scheduled_tick, "need_batch_scheduling", False): + batch_config = getattr(self._config, "batch_scheduling", None) + to_policy_config = getattr(batch_config, "to_policy_config", None) + if not callable(to_policy_config): + raise RuntimeError( + "batch scheduling tick is due but " + "batch_scheduling.to_policy_config() is unavailable" + ) + # Compute exactly once. Both the replica floor and drain lease + # below must be projections of the same input snapshot/decision. + batch_plan = plan_batch_schedule( + tick_input, + to_policy_config(), + decision_id=ctx.decision_id, + ) baseline = self._baseline_from_worker_counts(tick_input.worker_counts) outcome = await self._orchestrator.tick( ctx, @@ -588,8 +627,16 @@ async def tick( ) # 4. Project PipelineOutcome onto PlannerEffects. + worker_counts = tick_input.worker_counts or WorkerCounts() scale_to = self._project_scale_to( - outcome, tick_input.worker_counts or WorkerCounts() + outcome, + worker_counts, + batch_replica_floor=( + batch_plan.replica_floor if batch_plan is not None else None + ), + ) + batch_drain_limits, batch_audit_event = self._project_batch_drain_limits( + outcome, batch_plan, scale_to, worker_counts ) # 5. Populate diagnostics from the shared scaling state. Consumed by @@ -628,11 +675,14 @@ async def tick( diagnostics.execute_action = outcome.execute_action diagnostics.short_circuit_reason = outcome.short_circuit_reason diagnostics.audit_events = list(outcome.audit_events) + if batch_audit_event is not None: + diagnostics.audit_events.append(batch_audit_event) return PlannerEffects( scale_to=scale_to, next_tick=self._compute_next_scheduled_tick(), diagnostics=diagnostics, + batch_drain_limits=batch_drain_limits, ) async def observe(self, scheduled_tick: ScheduledTick, now_s: float) -> TickInput: @@ -706,6 +756,7 @@ def _compute_next_scheduled_tick(self) -> ScheduledTick: at_monotonic, ) ) + batch_loop_due = self._batch_scheduling_enabled # Lazy traffic pull: only when some currently-registered, # currently-due plugin actually consumes @@ -772,6 +823,7 @@ def due_consumers(path: str): at_monotonic_s=at_monotonic, run_load_scaling=load_loop_due, run_throughput_scaling=throughput_loop_due, + need_batch_scheduling=batch_loop_due, need_worker_states=True, need_worker_fpm=bool(fpm_consumers_due) or internal_fpm_due, need_traffic_metrics=need_traffic, @@ -840,10 +892,62 @@ def _tick_input_to_context(self, ti: TickInput) -> PipelineContext: # and could not implement load-based decisions through # the public PipelineContext API. fpm = self._encode_fpm(ti.fpm_observations) + batch = self._encode_batch(ti.batch) return PipelineContext( request_id=f"tick-{ti.now_s}", decision_id=f"d-{ti.now_s}", - observations=ObservationData(traffic=traffic, fpm=fpm, workers=workers), + observations=ObservationData( + traffic=traffic, + fpm=fpm, + workers=workers, + batch=batch, + ), + ) + + @staticmethod + def _encode_batch( + obs: Optional[BatchSchedulingObservation], + ) -> Optional[BatchSchedulingData]: + """Project generic core batch observations onto the plugin contract.""" + + if obs is None: + return None + return BatchSchedulingData( + job_demands=[ + BatchJobDemandData( + observed_at_s=job.observed_at_s, + pool_id=job.pool_id, + job_id=job.job_id, + status=job.status, + total_requests=job.total_requests, + completed_requests=job.completed_requests, + failed_requests=job.failed_requests, + deadline_at_s=job.deadline_at_s, + work_class=job.work_class, + remaining_requests=job.remaining_requests, + ) + for job in obs.job_demands + ], + pool_traffic=[ + PoolTrafficDemandData( + observed_at_s=pool.observed_at_s, + pool_id=pool.pool_id, + online_offered_rps=pool.online_offered_rps, + ) + for pool in obs.pool_traffic + ], + dispatcher_feedback=[ + BatchDispatcherFeedbackData( + observed_at_s=feedback.observed_at_s, + pool_id=feedback.pool_id, + observation_window_s=feedback.observation_window_s, + queued_requests=feedback.queued_requests, + inflight_requests=feedback.inflight_requests, + actual_dispatch_rps=feedback.actual_dispatch_rps, + applied_max_admission_rps=feedback.applied_max_admission_rps, + ) + for feedback in obs.dispatcher_feedback + ], ) @staticmethod @@ -909,7 +1013,109 @@ def _baseline_from_worker_counts( out[ComponentKey(sub_component_type="decode")] = counts.ready_num_decode return out - def _project_scale_to(self, outcome, worker_counts: WorkerCounts): + @staticmethod + def _project_batch_drain_limits( + outcome, batch_plan, scale_to, worker_counts: WorkerCounts + ): + """Project one batch plan to a leased dispatcher decision. + + A rejected/timed-out plugin pipeline must not leave a positive prior + lease live until expiry. The policy was still evaluated on this tick, + so reuse its newly computed lease identity/expiry while forcing + admission to zero. ``skip_no_targets`` is not a rejection: it means + the plugin pipeline had no scaling opinion, while the native batch + decision remains valid. + """ + + if batch_plan is None: + return [], None + + drain_limit = batch_plan.drain_limit + effective_decode = ( + scale_to.num_decode + if scale_to is not None and scale_to.num_decode is not None + else worker_counts.ready_num_decode + ) + if outcome.execute_action in ("skip_short_circuit", "skip_tick_timeout"): + drain_limit = BatchDrainLimitDecision( + pool_id=drain_limit.pool_id, + max_admission_rps=0.0, + valid_until_s=drain_limit.valid_until_s, + decision_id=drain_limit.decision_id, + ) + audit_event = ( + "batch_drain_limit_safety_pause:" + f"pipeline_action={outcome.execute_action}:" + f"pool_id={drain_limit.pool_id}:" + f"decision_id={drain_limit.decision_id}" + ) + log.warning( + "Batch scheduling safety pause: pipeline_action=%s " + "pool_id=%s decision_id=%s valid_until_s=%s", + outcome.execute_action, + drain_limit.pool_id, + drain_limit.decision_id, + drain_limit.valid_until_s, + ) + elif ( + effective_decode is not None and effective_decode < batch_plan.replica_floor + ): + # GPU/power caps are the final scaling boundary and are allowed to + # make the requested floor infeasible. Never pair that reduced + # capacity target with a positive drain lease. + drain_limit = BatchDrainLimitDecision( + pool_id=drain_limit.pool_id, + max_admission_rps=0.0, + valid_until_s=drain_limit.valid_until_s, + decision_id=drain_limit.decision_id, + ) + audit_event = ( + "batch_floor_blocked_by_final_budget:" + f"pool_id={drain_limit.pool_id}:" + f"replica_floor={batch_plan.replica_floor}:" + f"effective_decode={effective_decode}:" + f"decision_id={drain_limit.decision_id}" + ) + log.warning( + "Batch scheduling drain paused because final budget blocked " + "the replica floor: pool_id=%s replica_floor=%s " + "effective_decode=%s decision_id=%s valid_until_s=%s", + drain_limit.pool_id, + batch_plan.replica_floor, + effective_decode, + drain_limit.decision_id, + drain_limit.valid_until_s, + ) + else: + audit_event = ( + "batch_drain_limit_decision:" + f"pipeline_action={outcome.execute_action}:" + f"pool_id={drain_limit.pool_id}:" + f"max_admission_rps={drain_limit.max_admission_rps}:" + f"replica_floor={batch_plan.replica_floor}:" + f"decision_id={drain_limit.decision_id}" + ) + log.info( + "Batch scheduling decision: pipeline_action=%s pool_id=%s " + "replica_floor=%s max_admission_rps=%s decision_id=%s " + "valid_until_s=%s", + outcome.execute_action, + drain_limit.pool_id, + batch_plan.replica_floor, + drain_limit.max_admission_rps, + drain_limit.decision_id, + drain_limit.valid_until_s, + ) + + return [drain_limit], audit_event + + def _project_scale_to( + self, + outcome, + worker_counts: WorkerCounts, + *, + batch_replica_floor: Optional[int] = None, + ): """Project the pipeline outcome onto ``PlannerEffects.scale_to`` with planner "no change -> None" detection. @@ -918,12 +1124,27 @@ def _project_scale_to(self, outcome, worker_counts: WorkerCounts): actually targeted before that merge, so the power path can charge baseline peers without treating them as adjustable targets. """ - if outcome.execute_action != "apply" or outcome.final_proposal is None: + if batch_replica_floor is not None and ( + isinstance(batch_replica_floor, bool) + or not isinstance(batch_replica_floor, int) + or batch_replica_floor < 0 + ): + raise ValueError("batch_replica_floor must be a non-negative integer") + + pipeline_has_proposal = ( + outcome.execute_action == "apply" and outcome.final_proposal is not None + ) + batch_floor_allowed = batch_replica_floor is not None and ( + outcome.execute_action in ("apply", "skip_no_targets") + ) + if not pipeline_has_proposal and not batch_floor_allowed: return None - by_comp = { - t.sub_component_type: t.replicas for t in outcome.final_proposal.targets - } + by_comp = ( + {t.sub_component_type: t.replicas for t in outcome.final_proposal.targets} + if pipeline_has_proposal + else {} + ) num_p = by_comp.get("prefill") num_d = by_comp.get("decode") @@ -961,14 +1182,30 @@ def _role_scaling( and current_d is not None and current_d < decode_min_endpoint ) - floor_reconcile = mode == "disagg" and ( - prefill_floor_needed or decode_floor_needed - ) if prefill_floor_needed: num_p = max(num_p or 0, prefill_min_endpoint) if decode_floor_needed: num_d = max(num_d or 0, decode_min_endpoint) + # The batch policy floor is a native, dynamic final invariant. It can + # raise an explicit plugin target or synthesize a decode target when + # the plugin pipeline has no opinion. Do not echo the ready baseline + # when it already satisfies the floor: that remains a true no-op. + batch_floor_needed = False + if batch_floor_allowed and mode in ("disagg", "decode", "agg"): + batch_floor_baseline = num_d if num_d is not None else current_d + if batch_floor_baseline is None: + if batch_replica_floor > 0: + num_d = batch_replica_floor + batch_floor_needed = True + elif batch_floor_baseline < batch_replica_floor: + num_d = batch_replica_floor + batch_floor_needed = True + + floor_reconcile = mode == "disagg" and ( + prefill_floor_needed or decode_floor_needed or batch_floor_needed + ) + prefill_key = ComponentKey(sub_component_type="prefill") decode_key = ComponentKey(sub_component_type="decode") prefill_proposed = prefill_key in outcome.proposed_components @@ -985,7 +1222,12 @@ def _role_scaling( and not floor_reconcile ): num_p = None - if not decode_proposed and not decode_floor_needed and not floor_reconcile: + if ( + not decode_proposed + and not decode_floor_needed + and not batch_floor_needed + and not floor_reconcile + ): num_d = None if num_p is None and num_d is None: return None @@ -1013,6 +1255,7 @@ def _role_scaling( and num_d == current_d and not decode_proposed and not decode_floor_needed + and not batch_floor_needed ): num_d = None diff --git a/components/src/dynamo/planner/tests/plugins/orchestrator/test_engine_adapter.py b/components/src/dynamo/planner/tests/plugins/orchestrator/test_engine_adapter.py index d46dded2fa12..402954bebb57 100644 --- a/components/src/dynamo/planner/tests/plugins/orchestrator/test_engine_adapter.py +++ b/components/src/dynamo/planner/tests/plugins/orchestrator/test_engine_adapter.py @@ -22,12 +22,19 @@ from __future__ import annotations -import pytest +from types import SimpleNamespace +import dynamo.planner.plugins.orchestrator.engine_adapter as engine_adapter_module +import pytest from dynamo.planner.config.planner_config import PlannerConfig from dynamo.planner.core.types import ( + BatchDispatcherFeedback, + BatchDrainLimitDecision, + BatchJobDemand, + BatchSchedulingObservation, EngineCapabilities, FpmObservations, + PoolTrafficDemand, ScheduledTick, TickInput, TrafficObservation, @@ -69,6 +76,40 @@ def _agg_config_throughput_on() -> PlannerConfig: ) +def _enable_batch_for_test( + config: PlannerConfig, +): + """Attach the minimal batch config seam without coupling these tests to + environment/Redis source validation owned by PlannerConfig tests. + """ + + policy_config = object() + batch_config = SimpleNamespace( + enabled=True, + to_policy_config=lambda: policy_config, + ) + object.__setattr__(config, "batch_scheduling", batch_config) + return policy_config + + +def _batch_plan( + *, + replica_floor: int = 4, + max_admission_rps: float = 3.5, + valid_until_s: float = 65.0, + decision_id: str = "d-5.0", +): + return SimpleNamespace( + replica_floor=replica_floor, + drain_limit=BatchDrainLimitDecision( + pool_id="pool-a", + max_admission_rps=max_admission_rps, + valid_until_s=valid_until_s, + decision_id=decision_id, + ), + ) + + def _agg_config_custom_intervals( *, load_interval: float, throughput_interval: float ) -> PlannerConfig: @@ -384,6 +425,48 @@ def test_project_scale_to_non_apply_action_returns_none(): assert adapter._project_scale_to(outcome, wc) is None +def test_project_scale_to_merges_batch_floor_after_plugin_proposal(): + adapter = OrchestratorEngineAdapter(_agg_config_throughput_on(), _caps()) + wc = WorkerCounts(ready_num_decode=5) + outcome = _apply_outcome([ComponentTarget(sub_component_type="decode", replicas=2)]) + + decision = adapter._project_scale_to(outcome, wc, batch_replica_floor=4) + + assert decision is not None + assert decision.num_prefill is None + assert decision.num_decode == 4 + + +def test_project_scale_to_synthesizes_batch_floor_on_skip_no_targets(): + adapter = OrchestratorEngineAdapter(_agg_config_throughput_on(), _caps()) + outcome = PipelineOutcome(execute_action="skip_no_targets", final_proposal=None) + + decision = adapter._project_scale_to( + outcome, + WorkerCounts(ready_num_decode=2), + batch_replica_floor=4, + ) + + assert decision is not None + assert decision.num_prefill is None + assert decision.num_decode == 4 + + +@pytest.mark.parametrize("action", ["skip_short_circuit", "skip_tick_timeout"]) +def test_project_scale_to_never_scales_from_batch_floor_on_pipeline_failure(action): + adapter = OrchestratorEngineAdapter(_agg_config_throughput_on(), _caps()) + outcome = PipelineOutcome(execute_action=action, final_proposal=None) + + assert ( + adapter._project_scale_to( + outcome, + WorkerCounts(ready_num_decode=2), + batch_replica_floor=10, + ) + is None + ) + + def test_project_scale_to_applies_final_gpu_budget_to_external_proposal(): cfg = PlannerConfig( mode="disagg", @@ -612,6 +695,56 @@ def test_tick_input_to_context_maps_observations_and_fpm(): assert back.worker_id == "w1" +def test_tick_input_to_context_maps_batch_observations(): + adapter = OrchestratorEngineAdapter(_agg_config_throughput_on(), _caps()) + ti = TickInput( + now_s=1_700_000_010.0, + batch=BatchSchedulingObservation( + job_demands=[ + BatchJobDemand( + observed_at_s=1_700_000_000.0, + pool_id="pool-a", + job_id="batch-123", + status="in_progress", + total_requests=1_000, + completed_requests=275, + failed_requests=25, + deadline_at_s=1_700_003_600.0, + work_class="chat-8k", + ) + ], + pool_traffic=[ + PoolTrafficDemand( + observed_at_s=1_700_000_001.0, + pool_id="pool-a", + online_offered_rps=90.0, + ) + ], + dispatcher_feedback=[ + BatchDispatcherFeedback( + observed_at_s=1_700_000_002.0, + pool_id="pool-a", + observation_window_s=30.0, + queued_requests=700, + inflight_requests=20, + actual_dispatch_rps=9.5, + applied_max_admission_rps=10.0, + ) + ], + ), + ) + + ctx = adapter._tick_input_to_context(ti) + assert ctx.observations.batch is not None + assert ctx.observations.batch.job_demands[0].remaining_requests == 700 + assert ctx.observations.batch.job_demands[0].deadline_at_s == 1_700_003_600.0 + assert ctx.observations.batch.pool_traffic[0].online_offered_rps == 90.0 + assert ctx.observations.batch.dispatcher_feedback[0].queued_requests == 700 + assert ( + ctx.observations.batch.dispatcher_feedback[0].applied_max_admission_rps == 10.0 + ) + + def test_initial_tick_with_throughput_scaling_enabled_does_not_attribute_error(): """``initial_tick`` used to read the non-existent ``throughput_adjustment_interval`` attribute (canonical name has a @@ -636,6 +769,216 @@ def test_initial_tick_with_throughput_scaling_enabled_does_not_attribute_error() assert tick.run_load_scaling or tick.run_throughput_scaling +def test_batch_scheduling_runs_on_every_pipeline_tick_when_enabled(): + config = _agg_config_throughput_on() + config.scheduling.scale_interval_seconds = 5.0 + _enable_batch_for_test(config) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + first = adapter.initial_tick(start_s=0.0) + assert first.at_s == pytest.approx(5.0) + assert first.need_batch_scheduling is True + + adapter._last_tick_s = first.at_s + adapter._last_tick_monotonic = first.at_monotonic_s + second = adapter._compute_next_scheduled_tick() + assert second.at_s == pytest.approx(10.0) + assert second.need_batch_scheduling is True + + +@pytest.mark.asyncio +async def test_tick_computes_one_batch_plan_and_projects_both_outputs(monkeypatch): + config = _agg_config_throughput_on() + config.scheduling.scale_interval_seconds = 5.0 + policy_config = _enable_batch_for_test(config) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + calls = [] + + def fake_plan(tick_input, passed_config, *, decision_id): + calls.append((tick_input, passed_config, decision_id)) + return _batch_plan( + replica_floor=4, + max_admission_rps=3.5, + valid_until_s=tick_input.now_s + 60.0, + decision_id=decision_id, + ) + + monkeypatch.setattr(engine_adapter_module, "plan_batch_schedule", fake_plan) + + async def fake_orchestrator_tick(ctx, baseline, *, tick_now=None): + return PipelineOutcome(execute_action="skip_no_targets", final_proposal=None) + + adapter._orchestrator.tick = fake_orchestrator_tick # type: ignore[method-assign] + + scheduled = adapter.initial_tick(start_s=0.0) + assert scheduled.need_batch_scheduling is True + tick_input = TickInput( + now_s=scheduled.at_s, + worker_counts=WorkerCounts(ready_num_decode=2), + ) + + effects = await adapter.tick(scheduled, tick_input) + + assert calls == [(tick_input, policy_config, "d-5.0")] + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 4 + assert effects.batch_drain_limits == [ + BatchDrainLimitDecision( + pool_id="pool-a", + max_admission_rps=3.5, + valid_until_s=65.0, + decision_id="d-5.0", + ) + ] + assert effects.next_tick is not None + assert effects.next_tick.need_batch_scheduling is True + assert any( + event.startswith("batch_drain_limit_decision:") + for event in effects.diagnostics.audit_events + ) + + +@pytest.mark.asyncio +async def test_tick_pauses_batch_drain_when_final_budget_blocks_floor(monkeypatch): + config = _agg_config_throughput_on() + config.scheduling.scale_interval_seconds = 5.0 + config.min_gpu_budget = -1 + config.max_gpu_budget = 2 + _enable_batch_for_test(config) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + def fake_plan(tick_input, passed_config, *, decision_id): + return _batch_plan( + replica_floor=4, + max_admission_rps=3.5, + valid_until_s=tick_input.now_s + 60.0, + decision_id=decision_id, + ) + + monkeypatch.setattr(engine_adapter_module, "plan_batch_schedule", fake_plan) + + async def fake_orchestrator_tick(ctx, baseline, *, tick_now=None): + return PipelineOutcome(execute_action="skip_no_targets", final_proposal=None) + + adapter._orchestrator.tick = fake_orchestrator_tick # type: ignore[method-assign] + scheduled = adapter.initial_tick(start_s=0.0) + + effects = await adapter.tick( + scheduled, + TickInput( + now_s=scheduled.at_s, + worker_counts=WorkerCounts(ready_num_decode=1), + ), + ) + + assert effects.scale_to is not None + assert effects.scale_to.num_decode == 2 + assert effects.batch_drain_limits[0].max_admission_rps == 0.0 + assert any( + event.startswith("batch_floor_blocked_by_final_budget:") + for event in effects.diagnostics.audit_events + ) + + +@pytest.mark.asyncio +async def test_tick_pauses_batch_drain_when_budget_clamps_floor_to_current_noop( + monkeypatch, +): + config = _agg_config_throughput_on() + config.scheduling.scale_interval_seconds = 5.0 + config.min_gpu_budget = -1 + config.max_gpu_budget = 1 + _enable_batch_for_test(config) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + def fake_plan(tick_input, passed_config, *, decision_id): + return _batch_plan( + replica_floor=2, + max_admission_rps=3.5, + valid_until_s=tick_input.now_s + 60.0, + decision_id=decision_id, + ) + + monkeypatch.setattr(engine_adapter_module, "plan_batch_schedule", fake_plan) + + async def fake_orchestrator_tick(ctx, baseline, *, tick_now=None): + return PipelineOutcome(execute_action="skip_no_targets", final_proposal=None) + + adapter._orchestrator.tick = fake_orchestrator_tick # type: ignore[method-assign] + scheduled = adapter.initial_tick(start_s=0.0) + + effects = await adapter.tick( + scheduled, + TickInput( + now_s=scheduled.at_s, + worker_counts=WorkerCounts(ready_num_decode=1), + ), + ) + + # Final GPU clamping returned the current ready count, so scale projection + # correctly suppressed it as a no-op. Drain still fails closed because the + # effective count remains below the plan's floor. + assert effects.scale_to is None + assert effects.batch_drain_limits[0].max_admission_rps == 0.0 + assert any( + "replica_floor=2:effective_decode=1" in event + for event in effects.diagnostics.audit_events + if event.startswith("batch_floor_blocked_by_final_budget:") + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["skip_short_circuit", "skip_tick_timeout"]) +async def test_tick_emits_fresh_batch_pause_on_pipeline_failure(monkeypatch, action): + config = _agg_config_throughput_on() + config.scheduling.scale_interval_seconds = 5.0 + _enable_batch_for_test(config) + adapter = OrchestratorEngineAdapter(config, _caps(), clock=VirtualClock()) + + calls = [] + + def fake_plan(tick_input, passed_config, *, decision_id): + calls.append(decision_id) + return _batch_plan( + replica_floor=10, + max_admission_rps=7.0, + valid_until_s=tick_input.now_s + 60.0, + decision_id=decision_id, + ) + + monkeypatch.setattr(engine_adapter_module, "plan_batch_schedule", fake_plan) + + async def fake_orchestrator_tick(ctx, baseline, *, tick_now=None): + return PipelineOutcome(execute_action=action, final_proposal=None) + + adapter._orchestrator.tick = fake_orchestrator_tick # type: ignore[method-assign] + scheduled = adapter.initial_tick(start_s=0.0) + + effects = await adapter.tick( + scheduled, + TickInput( + now_s=scheduled.at_s, + worker_counts=WorkerCounts(ready_num_decode=2), + ), + ) + + assert calls == ["d-5.0"] + assert effects.scale_to is None + assert effects.batch_drain_limits == [ + BatchDrainLimitDecision( + pool_id="pool-a", + max_admission_rps=0.0, + valid_until_s=65.0, + decision_id="d-5.0", + ) + ] + assert any( + event.startswith(f"batch_drain_limit_safety_pause:pipeline_action={action}:") + for event in effects.diagnostics.audit_events + ) + + def test_orchestrator_path_honours_configured_protocol_version_range(): """``planner.plugin_registration.protocol_version_min/max`` must flow into the orchestrator-path registry server (previously dropped @@ -1376,15 +1719,15 @@ def _make_traffic_plugin(plugin_id: str, needs: list[str]) -> RegisteredPlugin: "p_sibling", ["observations.traffic_legacy"] ) sched = adapter._compute_next_scheduled_tick() - assert ( - sched.need_traffic_metrics is False - ), "prefix match must not over-fire on sibling field 'observations.traffic_legacy'" + assert sched.need_traffic_metrics is False, ( + "prefix match must not over-fire on sibling field 'observations.traffic_legacy'" + ) # --- Case 4: completely unrelated needs must NOT match --- adapter._orchestrator._registry._plugins["p_other"] = _make_traffic_plugin( "p_other", ["observations.fpm", "predictions"] ) sched = adapter._compute_next_scheduled_tick() - assert ( - sched.need_traffic_metrics is False - ), "unrelated needs must not trigger the traffic pull" + assert sched.need_traffic_metrics is False, ( + "unrelated needs must not trigger the traffic pull" + ) diff --git a/components/src/dynamo/planner/tests/plugins/orchestrator/test_environment_observe_plugin.py b/components/src/dynamo/planner/tests/plugins/orchestrator/test_environment_observe_plugin.py index 3585365137e9..fbb49416c728 100644 --- a/components/src/dynamo/planner/tests/plugins/orchestrator/test_environment_observe_plugin.py +++ b/components/src/dynamo/planner/tests/plugins/orchestrator/test_environment_observe_plugin.py @@ -3,7 +3,12 @@ import pytest -from dynamo.planner.core.types import FpmObservations, ScheduledTick, TrafficObservation +from dynamo.planner.core.types import ( + BatchSchedulingObservation, + FpmObservations, + ScheduledTick, + TrafficObservation, +) from dynamo.planner.environment.state import DeploymentState from dynamo.planner.plugins.builtins.observe import ( EnvironmentObservePlugin, @@ -24,6 +29,8 @@ def __init__(self) -> None: self.full_traffic_calls = 0 self.kv_hit_calls: list[float] = [] self.fpm_calls = 0 + self.batch_calls = 0 + self.batch_error: Exception | None = None def deployment_state(self) -> DeploymentState: return self.state @@ -46,6 +53,12 @@ def collect_fpm(self): self.fpm_calls += 1 return FpmObservations(prefill={}, decode={}) + async def collect_batch_scheduling(self): + self.batch_calls += 1 + if self.batch_error is not None: + raise self.batch_error + return BatchSchedulingObservation() + async def test_environment_observe_plugin_collects_requested_observations(): env = _FakeEnvironment() @@ -106,3 +119,54 @@ async def test_environment_observe_plugin_collects_load_only_traffic_window(): assert env.full_traffic_calls == 0 assert env.kv_hit_calls == [5] assert env.fpm_calls == 0 + + +async def test_environment_observe_plugin_collects_batch_and_anchors_after_io(): + env = _FakeEnvironment() + plugin = EnvironmentObservePlugin( + env, + require_prefill=False, + require_decode=True, + clock=lambda: 500.0, + ) + + response = await plugin.Observe( + ObserveStageRequest( + scheduled_tick=ScheduledTick( + at_s=123, + need_batch_scheduling=True, + ), + now_s=456, + ) + ) + + assert response.tick_input.now_s == 500.0 + assert response.tick_input.batch == BatchSchedulingObservation() + assert env.batch_calls == 1 + + +async def test_environment_observe_plugin_batch_failure_is_fail_closed(): + env = _FakeEnvironment() + env.batch_error = RuntimeError("telemetry unavailable") + plugin = EnvironmentObservePlugin( + env, + require_prefill=False, + require_decode=True, + clock=lambda: 500.0, + ) + + response = await plugin.Observe( + ObserveStageRequest( + scheduled_tick=ScheduledTick( + at_s=123, + need_worker_states=True, + need_batch_scheduling=True, + ), + now_s=456, + ) + ) + + assert response.tick_input.now_s == 500.0 + assert response.tick_input.batch is None + assert response.tick_input.worker_counts is not None + assert env.batch_calls == 1 diff --git a/components/src/dynamo/planner/tests/unit/test_tick_lifecycle.py b/components/src/dynamo/planner/tests/unit/test_tick_lifecycle.py index f7999ee11971..a2876fb82080 100644 --- a/components/src/dynamo/planner/tests/unit/test_tick_lifecycle.py +++ b/components/src/dynamo/planner/tests/unit/test_tick_lifecycle.py @@ -10,7 +10,13 @@ from dynamo.planner.config.defaults import SubComponentType from dynamo.planner.config.planner_config import PlannerConfig from dynamo.planner.core.adapters import AggPlanner -from dynamo.planner.core.types import PlannerEffects, ScalingDecision, ScheduledTick +from dynamo.planner.core.types import ( + BatchDrainLimitDecision, + PlannerEffects, + ScalingDecision, + ScheduledTick, + TickInput, +) from dynamo.planner.environment.state import DeploymentState from dynamo.planner.monitoring.traffic_metrics import Metrics from dynamo.planner.monitoring.worker_info import WorkerInfo @@ -124,3 +130,143 @@ async def _apply_effects(self, effects): else: assert applied_targets == [] assert events == expected_events + + +@pytest.mark.asyncio +@pytest.mark.parametrize("advisory", [False, True]) +async def test_batch_drain_is_applied_before_scaling_and_respects_advisory(advisory): + events = [] + state = DeploymentState() + state.decode.info = WorkerInfo(k8s_name="decode-worker") + state.decode.replicas.active = 1 + + environment = MagicMock() + environment.deployment_state.return_value = state + environment.metrics_state.return_value = Metrics() + environment.refresh = AsyncMock(return_value=state) + environment.apply_batch_drain_limits = AsyncMock( + side_effect=lambda _decisions: events.append("drain") + ) + environment.apply_scaling = AsyncMock( + side_effect=lambda _targets, blocking=False: events.append("scale") + ) + + next_tick = ScheduledTick(at_s=20.0) + drain = BatchDrainLimitDecision( + pool_id="pool", + max_admission_rps=5.0, + valid_until_s=100.0, + decision_id="decision", + ) + + class Engine: + async def tick(self, _scheduled_tick, _tick_input): + return PlannerEffects( + scale_to=ScalingDecision(num_decode=2), + next_tick=next_tick, + batch_drain_limits=[drain], + ) + + # Keep this test focused on effect ordering; observation composition is + # covered by the EnvironmentObservePlugin tests. + engine = Engine() + + async def observe(_scheduled_tick, now_s): + return TickInput( + now_s=now_s, + worker_counts=EnvironmentObservePlugin( + environment, + require_prefill=False, + require_decode=True, + )._collect_worker_counts(), + ) + + engine.observe = observe + + config = PlannerConfig( + mode="agg", + advisory=advisory, + namespace="test-namespace", + metric_reporting_prometheus_port=0, + live_dashboard_port=0, + report_interval_hours=None, + ) + with patch( + "dynamo.planner.core.base.PlannerPrometheusMetrics", + return_value=MagicMock(), + ): + planner = AggPlanner(None, config, environment) + + await planner._run_one_tick( + engine, + ScheduledTick(at_s=10.0, need_worker_states=True), + ) + + if advisory: + assert events == [] + else: + assert events == ["drain", "scale"] + + +@pytest.mark.asyncio +async def test_batch_actuation_failure_skips_scaling_without_stopping_tick(): + state = DeploymentState() + state.decode.info = WorkerInfo(k8s_name="decode-worker") + state.decode.replicas.active = 1 + environment = MagicMock() + environment.deployment_state.return_value = state + environment.metrics_state.return_value = Metrics() + environment.refresh = AsyncMock(return_value=state) + environment.apply_batch_drain_limits = AsyncMock( + side_effect=RuntimeError("redis unavailable") + ) + environment.apply_scaling = AsyncMock() + + next_tick = ScheduledTick(at_s=20.0) + + class Engine: + async def observe(self, _scheduled_tick, now_s): + return TickInput( + now_s=now_s, + worker_counts=EnvironmentObservePlugin( + environment, + require_prefill=False, + require_decode=True, + )._collect_worker_counts(), + ) + + async def tick(self, _scheduled_tick, _tick_input): + return PlannerEffects( + scale_to=ScalingDecision(num_decode=2), + next_tick=next_tick, + batch_drain_limits=[ + BatchDrainLimitDecision( + pool_id="pool", + max_admission_rps=5.0, + valid_until_s=100.0, + decision_id="decision", + ) + ], + ) + + config = PlannerConfig( + mode="agg", + namespace="test-namespace", + metric_reporting_prometheus_port=0, + live_dashboard_port=0, + report_interval_hours=None, + ) + with patch( + "dynamo.planner.core.base.PlannerPrometheusMetrics", + return_value=MagicMock(), + ): + planner = AggPlanner(None, config, environment) + + engine = Engine() + result = await planner._run_one_tick( + engine, + ScheduledTick(at_s=10.0, need_worker_states=True), + ) + + assert result is next_tick + environment.apply_scaling.assert_not_awaited()