From 282a9e4223ea6b01d2d69b16c2b3e7c4abfb34eb Mon Sep 17 00:00:00 2001 From: Sachal Malick Date: Mon, 31 Aug 2026 04:51:22 -0400 Subject: [PATCH] feat(planner): add batch scheduling policy and config --- .../dynamo/planner/config/planner_config.py | 344 +++++++++- .../src/dynamo/planner/core/__init__.py | 10 + .../src/dynamo/planner/core/batch_policy.py | 649 ++++++++++++++++++ .../config/test_batch_scheduling_config.py | 306 +++++++++ .../planner/tests/unit/test_batch_policy.py | 603 ++++++++++++++++ 5 files changed, 1899 insertions(+), 13 deletions(-) create mode 100644 components/src/dynamo/planner/core/batch_policy.py create mode 100644 components/src/dynamo/planner/tests/config/test_batch_scheduling_config.py create mode 100644 components/src/dynamo/planner/tests/unit/test_batch_policy.py diff --git a/components/src/dynamo/planner/config/planner_config.py b/components/src/dynamo/planner/config/planner_config.py index eb973563dbdd..19aa0c52e1db 100644 --- a/components/src/dynamo/planner/config/planner_config.py +++ b/components/src/dynamo/planner/config/planner_config.py @@ -17,29 +17,53 @@ import logging import math import os +import re from enum import Enum from pathlib import Path -from typing import Dict, Literal, Optional, Protocol -from urllib.parse import parse_qsl +from typing import TYPE_CHECKING, Dict, Literal, Optional, Protocol +from urllib.parse import parse_qsl, urlsplit import yaml +from dynamo.planner.config.aic_interpolation_spec import AICInterpolationSpec +from dynamo.planner.config.defaults import SLAPlannerDefaults +from dynamo.planner.config.parallelization import PickedParallelConfig +from dynamo.planner.plugins.registry.config import PluginRegistrationConfig +from dynamo.planner.plugins.types import HoldPolicy from pydantic import ( AliasChoices, BaseModel, ConfigDict, Field, + SecretStr, field_validator, model_validator, ) -from dynamo.planner.config.aic_interpolation_spec import AICInterpolationSpec -from dynamo.planner.config.defaults import SLAPlannerDefaults -from dynamo.planner.config.parallelization import PickedParallelConfig -from dynamo.planner.plugins.registry.config import PluginRegistrationConfig -from dynamo.planner.plugins.types import HoldPolicy +if TYPE_CHECKING: + from dynamo.planner.core.batch_policy import BatchSchedulingPolicyConfig logger = logging.getLogger(__name__) +_PROMETHEUS_LABEL_NAME = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + + +def _validate_http_url(value: str, *, field_name: str) -> str: + """Validate a service URL without normalizing its path or trailing slash.""" + + parsed = urlsplit(value) + if parsed.scheme not in ("http", "https") or not parsed.netloc: + raise ValueError(f"{field_name} must be an absolute http(s) URL") + if parsed.username is not None or parsed.password is not None: + raise ValueError( + f"{field_name} must not contain credentials; use a mounted secret" + ) + return value + + +def _batch_redis_url_default() -> SecretStr | None: + value = os.environ.get("DYN_PLANNER_BATCH_REDIS_URL") + return SecretStr(value) if value else None + class MinimumEndpointConfig(Protocol): """Configuration fields used to resolve component endpoint floors.""" @@ -367,6 +391,265 @@ class SchedulingConfig(BaseModel): ) +class BatchGatewaySourceConfig(BaseModel): + """Batch Gateway API settings for native batch-demand observation.""" + + model_config = ConfigDict(extra="forbid") + + base_url: str = Field( + ..., + min_length=1, + description="Base URL of the OpenAI-compatible Batch Gateway API.", + ) + tenant: str = Field( + ..., + min_length=1, + description=( + "Dedicated Batch Gateway tenant observed by this single-pool POC. " + "Sent as the X-MaaS-Username header." + ), + ) + request_timeout_seconds: float = Field( + default=5.0, + gt=0, + allow_inf_nan=False, + description="Per-request timeout for Batch Gateway API calls.", + ) + page_size: int = Field( + default=100, + ge=1, + le=100, + description="Batch list page size accepted by the Gateway API.", + ) + max_pages: int = Field( + default=10_000, + ge=1, + description="Safety bound on Batch Gateway pagination.", + ) + + @field_validator("base_url") + @classmethod + def _validate_base_url(cls, value: str) -> str: + return _validate_http_url(value, field_name="batch gateway base_url") + + +class BatchMetricsConfig(BaseModel): + """Direct OpenMetrics scrape endpoints used by batch-aware planning.""" + + model_config = ConfigDict(extra="forbid") + + frontend_metrics_url: str = Field( + ..., + min_length=1, + description=( + "Direct /metrics URL for the Dynamo frontend serving the shared pool." + ), + ) + dispatcher_metrics_url: str = Field( + ..., + min_length=1, + description="Direct /metrics URL for the llm-d Async dispatcher.", + ) + online_match_labels: dict[str, str] = Field( + default_factory=lambda: {"request_type": "stream"}, + description=( + "Exact frontend metric labels selecting online traffic. The POC " + "uses streaming requests for online traffic and unary requests for " + "batch traffic." + ), + ) + request_timeout_seconds: float = Field( + default=5.0, + gt=0, + allow_inf_nan=False, + description="Per-scrape timeout for both direct OpenMetrics endpoints.", + ) + + @field_validator("frontend_metrics_url") + @classmethod + def _validate_frontend_metrics_url(cls, value: str) -> str: + return _validate_http_url(value, field_name="frontend_metrics_url") + + @field_validator("dispatcher_metrics_url") + @classmethod + def _validate_dispatcher_metrics_url(cls, value: str) -> str: + return _validate_http_url(value, field_name="dispatcher_metrics_url") + + @field_validator("online_match_labels") + @classmethod + def _validate_online_match_labels(cls, value: dict[str, str]) -> dict[str, str]: + if not value: + raise ValueError( + "online_match_labels must explicitly distinguish online traffic" + ) + for label_name, label_value in value.items(): + if not _PROMETHEUS_LABEL_NAME.fullmatch(label_name): + raise ValueError( + f"online_match_labels contains invalid label name {label_name!r}" + ) + if not label_value: + raise ValueError("online_match_labels values must be non-empty strings") + return dict(value) + + +class BatchRedisActuatorConfig(BaseModel): + """Redis connection and key used to publish leased drain limits.""" + + model_config = ConfigDict(extra="forbid") + + url: SecretStr | None = Field( + default_factory=_batch_redis_url_default, + exclude=True, + description=( + "Redis URL loaded from DYN_PLANNER_BATCH_REDIS_URL by default. " + "Excluded from serialized Planner config." + ), + ) + control_key: str = Field( + ..., + min_length=1, + description="Exact llm-d Async Redis hash key for the configured pool.", + ) + connect_timeout_seconds: float = Field( + default=2.0, + gt=0, + allow_inf_nan=False, + description="Redis connection timeout.", + ) + socket_timeout_seconds: float = Field( + default=2.0, + gt=0, + allow_inf_nan=False, + description="Redis command read/write timeout.", + ) + + @field_validator("url") + @classmethod + def _validate_redis_url(cls, value: SecretStr | None) -> SecretStr | None: + if value is None: + return None + raw_url = value.get_secret_value() + parsed = urlsplit(raw_url) + if parsed.scheme not in ("redis", "rediss") or not parsed.netloc: + raise ValueError("batch Redis url must be an absolute redis(s) URL") + return value + + +class BatchPoolConfig(BaseModel): + """Capacity and deadline assumptions for the aggregate batch pool.""" + + model_config = ConfigDict(extra="forbid") + + pool_id: str = Field(..., min_length=1) + work_class: str = Field(..., min_length=1) + safe_rps_per_ready_replica: float = Field( + ..., + gt=0, + allow_inf_nan=False, + ) + cold_start_margin_seconds: float = Field( + default=0.0, + ge=0, + allow_inf_nan=False, + ) + finalization_margin_seconds: float = Field( + default=0.0, + ge=0, + allow_inf_nan=False, + ) + max_observation_age_seconds: float = Field( + default=60.0, + ge=0, + allow_inf_nan=False, + ) + drain_lease_duration_seconds: float = Field( + default=60.0, + gt=0, + allow_inf_nan=False, + description=("Lifetime of each fail-closed llm-d Async drain-limit lease."), + ) + min_replicas: int = Field(default=0, ge=0) + max_replicas: int = Field(..., ge=0) + scale_from_zero_replicas: int = Field( + default=1, + ge=1, + description=( + "Replica floor requested when a fresh, valid active batch job is " + "observed while the pool has zero ready replicas." + ), + ) + max_batch_admission_rps: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + ) + + @model_validator(mode="after") + def _validate_replica_range(self) -> "BatchPoolConfig": + if self.max_replicas < self.min_replicas: + raise ValueError("max_replicas must be >= min_replicas") + if self.scale_from_zero_replicas > self.max_replicas: + raise ValueError("scale_from_zero_replicas must be <= max_replicas") + return self + + +class BatchSchedulingConfig(BaseModel): + """Opt-in native Planner integration for one aggregate batch pool.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + gateway: BatchGatewaySourceConfig | None = None + metrics: BatchMetricsConfig | None = None + redis: BatchRedisActuatorConfig | None = None + pool: BatchPoolConfig | None = None + + @model_validator(mode="after") + def _validate_enabled_dependencies(self) -> "BatchSchedulingConfig": + if not self.enabled: + return self + + missing = [ + field_name + for field_name in ("gateway", "metrics", "redis", "pool") + if getattr(self, field_name) is None + ] + if missing: + raise ValueError( + "batch_scheduling.enabled=True requires: " + ", ".join(missing) + ) + assert self.redis is not None + if self.redis.url is None: + raise ValueError( + "batch_scheduling.enabled=True requires redis.url or " + "DYN_PLANNER_BATCH_REDIS_URL" + ) + return self + + def to_policy_config(self) -> "BatchSchedulingPolicyConfig": + """Translate validated Planner config into the pure policy contract.""" + + if not self.enabled or self.pool is None: + raise ValueError("batch scheduling must be enabled and configured") + + from dynamo.planner.core.batch_policy import BatchSchedulingPolicyConfig + + pool = self.pool + return BatchSchedulingPolicyConfig( + pool_id=pool.pool_id, + work_class=pool.work_class, + safe_rps_per_ready_replica=pool.safe_rps_per_ready_replica, + cold_start_margin_s=pool.cold_start_margin_seconds, + finalization_margin_s=pool.finalization_margin_seconds, + max_observation_age_s=pool.max_observation_age_seconds, + drain_lease_duration_s=pool.drain_lease_duration_seconds, + min_replicas=pool.min_replicas, + max_replicas=pool.max_replicas, + scale_from_zero_replicas=pool.scale_from_zero_replicas, + max_batch_admission_rps=pool.max_batch_admission_rps, + ) + + class PlannerConfig(BaseModel): """Pydantic configuration for the Dynamo Planner. @@ -384,9 +667,9 @@ class PlannerConfig(BaseModel): ), ) - environment: Literal[ - "kubernetes", "virtual", "global-planner" - ] = SLAPlannerDefaults.environment + environment: Literal["kubernetes", "virtual", "global-planner"] = ( + SLAPlannerDefaults.environment + ) namespace: str = Field( default_factory=lambda: os.environ.get("DYN_NAMESPACE", "dynamo"), exclude=True, @@ -610,9 +893,9 @@ def _validate_ca_bundle_path(cls, v: Optional[str]) -> Optional[str]: metric_reporting_prometheus_port: int = Field( default_factory=lambda: int(os.environ.get("PLANNER_PROMETHEUS_PORT", 0)) ) - throughput_metrics_source: Literal[ - "frontend", "router" - ] = SLAPlannerDefaults.throughput_metrics_source + throughput_metrics_source: Literal["frontend", "router"] = ( + SLAPlannerDefaults.throughput_metrics_source + ) model_name: Optional[str] = None @@ -772,6 +1055,13 @@ def _validate_ca_bundle_path(cls, v: Optional[str]) -> Optional[str]: ), ) + batch_scheduling: BatchSchedulingConfig = Field( + default_factory=BatchSchedulingConfig, + description=( + "Opt-in native batch scheduling for one Kubernetes aggregate pool." + ), + ) + plugin_registration: PluginRegistrationConfig = Field( default_factory=PluginRegistrationConfig, description=( @@ -864,6 +1154,34 @@ def _validate_config(self) -> "PlannerConfig": "when mode='agg'; use min_endpoint" ) + if self.batch_scheduling.enabled: + if self.environment != "kubernetes": + raise ValueError( + "batch_scheduling.enabled=True requires environment='kubernetes'" + ) + if self.mode != "agg": + raise ValueError("batch_scheduling.enabled=True requires mode='agg'") + + pool = self.batch_scheduling.pool + assert pool is not None + if pool.max_replicas < self.effective_decode_min_endpoint: + raise ValueError( + "batch_scheduling.pool.max_replicas must be >= the effective " + f"aggregate endpoint minimum ({self.effective_decode_min_endpoint})" + ) + scale_interval_s = self.scheduling.scale_interval_seconds + minimum_lease_s = max( + 2.0 * scale_interval_s, + self.scheduling.tick_max_duration_seconds + scale_interval_s, + ) + if pool.drain_lease_duration_seconds < minimum_lease_s: + raise ValueError( + "batch_scheduling.pool.drain_lease_duration_seconds must be " + "at least max(2 * scheduling.scale_interval_seconds, " + "scheduling.tick_max_duration_seconds + " + f"scheduling.scale_interval_seconds) ({minimum_lease_s}s)" + ) + if self.report_interval_hours is not None: if ( not math.isfinite(self.report_interval_hours) diff --git a/components/src/dynamo/planner/core/__init__.py b/components/src/dynamo/planner/core/__init__.py index 5584e2d9b9a5..e4a2b10bad81 100644 --- a/components/src/dynamo/planner/core/__init__.py +++ b/components/src/dynamo/planner/core/__init__.py @@ -1,6 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from dynamo.planner.core.batch_policy import ( + BatchSchedulingDiagnostics, + BatchSchedulingPlan, + BatchSchedulingPolicyConfig, + plan_batch_schedule, +) from dynamo.planner.core.state_machine import PlannerScalingState from dynamo.planner.core.types import ( BatchDispatcherFeedback, @@ -23,7 +29,10 @@ "BatchDispatcherFeedback", "BatchDrainLimitDecision", "BatchJobDemand", + "BatchSchedulingDiagnostics", "BatchSchedulingObservation", + "BatchSchedulingPlan", + "BatchSchedulingPolicyConfig", "EngineCapabilities", "FpmObservations", "PlannerEffects", @@ -35,4 +44,5 @@ "TrafficObservation", "WorkerCapabilities", "WorkerCounts", + "plan_batch_schedule", ] diff --git a/components/src/dynamo/planner/core/batch_policy.py b/components/src/dynamo/planner/core/batch_policy.py new file mode 100644 index 000000000000..0e53969b00ad --- /dev/null +++ b/components/src/dynamo/planner/core/batch_policy.py @@ -0,0 +1,649 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure batch scheduling policy for one aggregate inference pool. + +The policy performs no observation collection and no actuation. It turns one +``TickInput`` snapshot into two independently consumable recommendations: + +* an absolute replica lower bound for the existing scaling pipeline; and +* a leased maximum batch-admission rate for a dispatcher. + +Deadline pressure is computed with an earliest-deadline-first cumulative demand +bound. Current drain is work-conserving, but is never allowed to consume more +than the capacity left after online offered load. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional + +from dynamo.planner.core.types import ( + BatchDispatcherFeedback, + BatchDrainLimitDecision, + BatchJobDemand, + PoolTrafficDemand, + TickInput, +) + +_NON_DISPATCHABLE_STATUSES = frozenset( + { + "canceled", + "cancelled", + "cancelling", + "completed", + "expired", + "failed", + "finalizing", + } +) + + +@dataclass(frozen=True) +class BatchSchedulingPolicyConfig: + """Static assumptions for the single-pool POC policy. + + ``safe_rps_per_ready_replica`` is calibrated for ``work_class``. The two + margins reserve time before each absolute job deadline. The optional cap is + an independent upper bound on batch admission, even when more pool + headroom is available. + """ + + pool_id: str + work_class: str + safe_rps_per_ready_replica: float + cold_start_margin_s: float + finalization_margin_s: float + max_observation_age_s: float + drain_lease_duration_s: float + min_replicas: int + max_replicas: int + scale_from_zero_replicas: int = 1 + max_batch_admission_rps: Optional[float] = None + + def __post_init__(self) -> None: + if not self.pool_id: + raise ValueError("pool_id must not be empty") + if not self.work_class: + raise ValueError("work_class must not be empty") + _require_positive_finite( + "safe_rps_per_ready_replica", self.safe_rps_per_ready_replica + ) + _require_non_negative_finite("cold_start_margin_s", self.cold_start_margin_s) + _require_non_negative_finite( + "finalization_margin_s", self.finalization_margin_s + ) + _require_non_negative_finite( + "max_observation_age_s", self.max_observation_age_s + ) + _require_positive_finite("drain_lease_duration_s", self.drain_lease_duration_s) + if self.min_replicas < 0: + raise ValueError("min_replicas must be non-negative") + if self.max_replicas < self.min_replicas: + raise ValueError("max_replicas must be >= min_replicas") + if self.scale_from_zero_replicas <= 0: + raise ValueError("scale_from_zero_replicas must be positive") + if self.scale_from_zero_replicas > self.max_replicas: + raise ValueError("scale_from_zero_replicas must be <= max_replicas") + if self.max_batch_admission_rps is not None: + _require_non_negative_finite( + "max_batch_admission_rps", self.max_batch_admission_rps + ) + + +@dataclass(frozen=True) +class BatchSchedulingDiagnostics: + """Auditable intermediate values from one batch policy evaluation.""" + + required_batch_rps: Optional[float] + current_safe_headroom_rps: float + planned_batch_capacity_rps: float + predicted_finish_at_s: Optional[float] + minimum_deadline_slack_s: Optional[float] + infeasible: bool + infeasible_reasons: tuple[str, ...] + safety_paused: bool + active_job_count: int + remaining_requests: int + ready_replicas: Optional[int] + required_replica_floor: Optional[int] + online_offered_rps: Optional[float] + online_observation_age_s: Optional[float] + oldest_job_observation_age_s: Optional[float] + capacity_observation_stale: bool + online_observation_stale: bool + job_observation_stale: bool + dispatcher_feedback_stale: bool + dispatcher_observation_age_s: Optional[float] + actual_dispatch_rps: Optional[float] + applied_max_admission_rps: Optional[float] + + +@dataclass(frozen=True) +class BatchSchedulingPlan: + """Pure policy result; callers merge the floor with other scale opinions.""" + + replica_floor: int + drain_limit: BatchDrainLimitDecision + diagnostics: BatchSchedulingDiagnostics + + +def plan_batch_schedule( + tick_input: TickInput, + config: BatchSchedulingPolicyConfig, + *, + decision_id: str, +) -> BatchSchedulingPlan: + """Compute a deterministic batch plan from one input snapshot. + + ``ready_num_decode`` is the aggregate POC's ready-worker count. A present + worker snapshot is considered current for ``tick_input.now_s`` because the + normal planner adapter gathers it for that tick. Missing or invalid worker + state is treated as stale capacity. + """ + + if not decision_id: + raise ValueError("decision_id must not be empty") + if not math.isfinite(tick_input.now_s): + raise ValueError("tick_input.now_s must be finite") + + now_s = tick_input.now_s + batch = tick_input.batch + ready_replicas, capacity_stale = _ready_replicas(tick_input) + hold_floor = _hold_floor(ready_replicas, config) + + pool_traffic = ( + _latest_pool_traffic(batch.pool_traffic, config.pool_id) if batch else None + ) + online_age_s, online_stale = _observation_freshness( + pool_traffic.observed_at_s if pool_traffic else None, + now_s, + config.max_observation_age_s, + ) + online_offered_rps = pool_traffic.online_offered_rps if pool_traffic else None + if online_offered_rps is not None and ( + not math.isfinite(online_offered_rps) or online_offered_rps < 0 + ): + online_stale = True + + active_jobs = _active_jobs(batch.job_demands, config.pool_id) if batch else [] + oldest_job_age_s: Optional[float] = None + job_stale = False + mismatched_jobs: list[BatchJobDemand] = [] + invalid_jobs: list[BatchJobDemand] = [] + for job in active_jobs: + age_s, stale = _observation_freshness( + job.observed_at_s, + now_s, + config.max_observation_age_s, + ) + if age_s is not None: + oldest_job_age_s = ( + age_s if oldest_job_age_s is None else max(oldest_job_age_s, age_s) + ) + job_stale = job_stale or stale + if job.work_class != config.work_class: + mismatched_jobs.append(job) + if job.remaining_requests < 0 or ( + job.deadline_at_s is not None and not math.isfinite(job.deadline_at_s) + ): + invalid_jobs.append(job) + + dispatcher_feedback = ( + _latest_dispatcher_feedback(batch.dispatcher_feedback, config.pool_id) + if batch + else None + ) + dispatcher_age_s, dispatcher_stale = _observation_freshness( + dispatcher_feedback.observed_at_s if dispatcher_feedback else None, + now_s, + config.max_observation_age_s, + ) + + safety_reasons: list[str] = [] + if batch is None: + safety_reasons.append("batch_observation_missing") + if capacity_stale: + safety_reasons.append("capacity_observation_missing_or_invalid") + if pool_traffic is None: + safety_reasons.append("online_traffic_observation_missing") + elif online_stale: + safety_reasons.append("online_traffic_observation_stale_or_invalid") + if job_stale: + safety_reasons.append("active_job_observation_stale_or_invalid") + if invalid_jobs: + safety_reasons.extend( + f"active_job_invalid:{job.job_id}" for job in invalid_jobs + ) + if mismatched_jobs: + safety_reasons.extend( + f"work_class_mismatch:{job.job_id}:{job.work_class}" + for job in mismatched_jobs + ) + + # A fresh durable job snapshot is enough to break the scale-from-zero + # observation loop: frontend traffic may remain unavailable until a worker + # is serving, but admission remains fail-closed until that sample exists. + # Stale or malformed demand must never trigger capacity. + scale_from_zero_floor = ( + config.scale_from_zero_replicas + if ready_replicas == 0 + and active_jobs + and not job_stale + and not invalid_jobs + and not mismatched_jobs + else 0 + ) + + if safety_reasons: + return _safety_pause_plan( + now_s=now_s, + decision_id=decision_id, + config=config, + replica_floor=max(hold_floor, scale_from_zero_floor), + required_replica_floor=( + scale_from_zero_floor if scale_from_zero_floor > 0 else None + ), + reasons=tuple(safety_reasons), + active_jobs=active_jobs, + ready_replicas=ready_replicas, + online_offered_rps=online_offered_rps, + online_age_s=online_age_s, + oldest_job_age_s=oldest_job_age_s, + capacity_stale=capacity_stale, + online_stale=online_stale, + job_stale=job_stale, + dispatcher_feedback=dispatcher_feedback, + dispatcher_age_s=dispatcher_age_s, + dispatcher_stale=dispatcher_stale, + ) + + # These values are proven present by the safety gate above. + assert ready_replicas is not None + assert online_offered_rps is not None + + required_batch_rps, demand_bounds, infeasible_reasons = _edf_demand_bound( + active_jobs, + now_s, + config, + ) + required_replica_floor: Optional[int] + if math.isfinite(required_batch_rps): + required_replica_floor = max( + config.min_replicas, + scale_from_zero_floor, + math.ceil( + (online_offered_rps + required_batch_rps) + / config.safe_rps_per_ready_replica + ), + ) + replica_floor = min(required_replica_floor, config.max_replicas) + if required_replica_floor > config.max_replicas: + infeasible_reasons.append("required_replica_floor_exceeds_max_replicas") + else: + required_replica_floor = None + replica_floor = config.max_replicas + + if ( + config.max_batch_admission_rps is not None + and required_batch_rps > config.max_batch_admission_rps + ): + infeasible_reasons.append("batch_cap_below_required_rps") + + current_safe_headroom_rps = max( + 0.0, + ready_replicas * config.safe_rps_per_ready_replica - online_offered_rps, + ) + candidate_drain_cap_rps = current_safe_headroom_rps if active_jobs else 0.0 + if config.max_batch_admission_rps is not None: + candidate_drain_cap_rps = min( + candidate_drain_cap_rps, + config.max_batch_admission_rps, + ) + + # Scaling and drain decisions are applied independently downstream, so the + # advertised replica floor must be able to sustain the advertised drain. + # Without this coupling, a best-effort job could consume today's headroom + # while another scaling opinion simultaneously reduced the fleet to the + # deadline-only floor. + if candidate_drain_cap_rps > 0: + drain_sustaining_floor = math.ceil( + (online_offered_rps + candidate_drain_cap_rps) + / config.safe_rps_per_ready_replica + ) + replica_floor = min( + config.max_replicas, + max(replica_floor, drain_sustaining_floor), + ) + + # Predict against the advertised lower bound, not today's possibly larger + # fleet. This keeps the deadline/slack claim valid if another scaling + # opinion elects to scale down to exactly this floor. + planned_replicas = replica_floor + planned_batch_capacity_rps = max( + 0.0, + planned_replicas * config.safe_rps_per_ready_replica - online_offered_rps, + ) + if config.max_batch_admission_rps is not None: + planned_batch_capacity_rps = min( + planned_batch_capacity_rps, + config.max_batch_admission_rps, + ) + + # A configured replica ceiling can prevent the floor from preserving all + # current headroom. Clip the drain in that case so the joint decision still + # satisfies online + batch <= floor * safe per-replica capacity. + drain_cap_rps = min( + candidate_drain_cap_rps, + planned_batch_capacity_rps, + ) + + predicted_finish_at_s, minimum_slack_s = _finish_diagnostics( + active_jobs=active_jobs, + demand_bounds=demand_bounds, + now_s=now_s, + planned_batch_capacity_rps=planned_batch_capacity_rps, + scaling_up=replica_floor > ready_replicas, + config=config, + ) + if demand_bounds and planned_batch_capacity_rps <= 0: + infeasible_reasons.append("no_planned_batch_capacity") + elif minimum_slack_s is not None and minimum_slack_s < -1e-9: + infeasible_reasons.append("negative_deadline_slack") + + infeasible_reasons = list(dict.fromkeys(infeasible_reasons)) + diagnostics = BatchSchedulingDiagnostics( + required_batch_rps=required_batch_rps, + current_safe_headroom_rps=current_safe_headroom_rps, + planned_batch_capacity_rps=planned_batch_capacity_rps, + predicted_finish_at_s=predicted_finish_at_s, + minimum_deadline_slack_s=minimum_slack_s, + infeasible=bool(infeasible_reasons), + infeasible_reasons=tuple(infeasible_reasons), + safety_paused=False, + active_job_count=len(active_jobs), + remaining_requests=sum(job.remaining_requests for job in active_jobs), + ready_replicas=ready_replicas, + required_replica_floor=required_replica_floor, + online_offered_rps=online_offered_rps, + online_observation_age_s=online_age_s, + oldest_job_observation_age_s=oldest_job_age_s, + capacity_observation_stale=False, + online_observation_stale=False, + job_observation_stale=False, + dispatcher_feedback_stale=dispatcher_stale, + dispatcher_observation_age_s=dispatcher_age_s, + actual_dispatch_rps=( + dispatcher_feedback.actual_dispatch_rps if dispatcher_feedback else None + ), + applied_max_admission_rps=( + dispatcher_feedback.applied_max_admission_rps + if dispatcher_feedback + else None + ), + ) + return BatchSchedulingPlan( + replica_floor=replica_floor, + drain_limit=BatchDrainLimitDecision( + pool_id=config.pool_id, + max_admission_rps=drain_cap_rps, + valid_until_s=now_s + config.drain_lease_duration_s, + decision_id=decision_id, + ), + diagnostics=diagnostics, + ) + + +def _require_positive_finite(name: str, value: float) -> None: + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be positive and finite") + + +def _require_non_negative_finite(name: str, value: float) -> None: + if not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be non-negative and finite") + + +def _ready_replicas(tick_input: TickInput) -> tuple[Optional[int], bool]: + worker_counts = tick_input.worker_counts + if worker_counts is None or worker_counts.ready_num_decode is None: + return None, True + ready_replicas = worker_counts.ready_num_decode + if ready_replicas < 0: + return None, True + return ready_replicas, False + + +def _hold_floor( + ready_replicas: Optional[int], config: BatchSchedulingPolicyConfig +) -> int: + if ready_replicas is None: + return config.max_replicas + return min(config.max_replicas, max(config.min_replicas, ready_replicas)) + + +def _observation_freshness( + observed_at_s: Optional[float], + now_s: float, + max_age_s: float, +) -> tuple[Optional[float], bool]: + if observed_at_s is None or not math.isfinite(observed_at_s): + return None, True + age_s = now_s - observed_at_s + return age_s, age_s < 0 or age_s > max_age_s + + +def _latest_pool_traffic( + observations: list[PoolTrafficDemand], pool_id: str +) -> Optional[PoolTrafficDemand]: + matches = [ + observation for observation in observations if observation.pool_id == pool_id + ] + if not matches: + return None + return max( + matches, + key=lambda observation: ( + _sortable_timestamp(observation.observed_at_s), + observation.online_offered_rps, + ), + ) + + +def _latest_dispatcher_feedback( + observations: list[BatchDispatcherFeedback], pool_id: str +) -> Optional[BatchDispatcherFeedback]: + matches = [ + observation for observation in observations if observation.pool_id == pool_id + ] + if not matches: + return None + return max( + matches, + key=lambda observation: ( + _sortable_timestamp(observation.observed_at_s), + observation.queued_requests, + observation.inflight_requests, + observation.actual_dispatch_rps, + ), + ) + + +def _active_jobs( + observations: list[BatchJobDemand], pool_id: str +) -> list[BatchJobDemand]: + latest_by_id: dict[str, BatchJobDemand] = {} + for job in observations: + if job.pool_id != pool_id: + continue + previous = latest_by_id.get(job.job_id) + if previous is None or _job_order_key(job) > _job_order_key(previous): + latest_by_id[job.job_id] = job + return sorted( + ( + job + for job in latest_by_id.values() + if job.remaining_requests != 0 + and job.status.strip().lower() not in _NON_DISPATCHABLE_STATUSES + ), + key=lambda job: job.job_id, + ) + + +def _job_order_key(job: BatchJobDemand) -> tuple[float, str, int, int, int, float, str]: + return ( + _sortable_timestamp(job.observed_at_s), + job.status, + job.total_requests, + job.completed_requests, + job.failed_requests, + ( + _sortable_timestamp(job.deadline_at_s) + if job.deadline_at_s is not None + else math.inf + ), + job.work_class, + ) + + +def _sortable_timestamp(timestamp_s: float) -> float: + return timestamp_s if math.isfinite(timestamp_s) else math.inf + + +def _edf_demand_bound( + active_jobs: list[BatchJobDemand], + now_s: float, + config: BatchSchedulingPolicyConfig, +) -> tuple[float, list[tuple[float, int]], list[str]]: + deadline_jobs = sorted( + (job for job in active_jobs if job.deadline_at_s is not None), + key=lambda job: (job.deadline_at_s, job.job_id), + ) + cumulative_requests = 0 + required_batch_rps = 0.0 + demand_bounds: list[tuple[float, int]] = [] + infeasible_reasons: list[str] = [] + for job in deadline_jobs: + assert job.deadline_at_s is not None + cumulative_requests += job.remaining_requests + demand_bounds.append((job.deadline_at_s, cumulative_requests)) + available_s = ( + job.deadline_at_s + - now_s + - config.cold_start_margin_s + - config.finalization_margin_s + ) + if available_s <= 0: + required_batch_rps = math.inf + infeasible_reasons.append(f"deadline_window_exhausted:{job.job_id}") + continue + required_batch_rps = max( + required_batch_rps, + cumulative_requests / available_s, + ) + return required_batch_rps, demand_bounds, infeasible_reasons + + +def _finish_diagnostics( + *, + active_jobs: list[BatchJobDemand], + demand_bounds: list[tuple[float, int]], + now_s: float, + planned_batch_capacity_rps: float, + scaling_up: bool, + config: BatchSchedulingPolicyConfig, +) -> tuple[Optional[float], Optional[float]]: + remaining_requests = sum(job.remaining_requests for job in active_jobs) + if remaining_requests == 0: + return now_s, None + if planned_batch_capacity_rps <= 0: + return None, -math.inf if demand_bounds else None + + start_at_s = now_s + (config.cold_start_margin_s if scaling_up else 0.0) + predicted_finish_at_s = ( + start_at_s + + remaining_requests / planned_batch_capacity_rps + + config.finalization_margin_s + ) + slacks = [ + deadline_at_s + - ( + start_at_s + + cumulative_requests / planned_batch_capacity_rps + + config.finalization_margin_s + ) + for deadline_at_s, cumulative_requests in demand_bounds + ] + return predicted_finish_at_s, min(slacks) if slacks else None + + +def _safety_pause_plan( + *, + now_s: float, + decision_id: str, + config: BatchSchedulingPolicyConfig, + replica_floor: int, + required_replica_floor: Optional[int], + reasons: tuple[str, ...], + active_jobs: list[BatchJobDemand], + ready_replicas: Optional[int], + online_offered_rps: Optional[float], + online_age_s: Optional[float], + oldest_job_age_s: Optional[float], + capacity_stale: bool, + online_stale: bool, + job_stale: bool, + dispatcher_feedback: Optional[BatchDispatcherFeedback], + dispatcher_age_s: Optional[float], + dispatcher_stale: bool, +) -> BatchSchedulingPlan: + diagnostics = BatchSchedulingDiagnostics( + required_batch_rps=None, + current_safe_headroom_rps=0.0, + planned_batch_capacity_rps=0.0, + predicted_finish_at_s=None, + minimum_deadline_slack_s=None, + infeasible=True, + infeasible_reasons=reasons, + safety_paused=True, + active_job_count=len(active_jobs), + remaining_requests=sum(max(0, job.remaining_requests) for job in active_jobs), + ready_replicas=ready_replicas, + required_replica_floor=required_replica_floor, + online_offered_rps=online_offered_rps, + online_observation_age_s=online_age_s, + oldest_job_observation_age_s=oldest_job_age_s, + capacity_observation_stale=capacity_stale, + online_observation_stale=online_stale, + job_observation_stale=job_stale, + dispatcher_feedback_stale=dispatcher_stale, + dispatcher_observation_age_s=dispatcher_age_s, + actual_dispatch_rps=( + dispatcher_feedback.actual_dispatch_rps if dispatcher_feedback else None + ), + applied_max_admission_rps=( + dispatcher_feedback.applied_max_admission_rps + if dispatcher_feedback + else None + ), + ) + return BatchSchedulingPlan( + replica_floor=replica_floor, + drain_limit=BatchDrainLimitDecision( + pool_id=config.pool_id, + max_admission_rps=0.0, + valid_until_s=now_s + config.drain_lease_duration_s, + decision_id=decision_id, + ), + diagnostics=diagnostics, + ) + + +__all__ = [ + "BatchSchedulingDiagnostics", + "BatchSchedulingPlan", + "BatchSchedulingPolicyConfig", + "plan_batch_schedule", +] diff --git a/components/src/dynamo/planner/tests/config/test_batch_scheduling_config.py b/components/src/dynamo/planner/tests/config/test_batch_scheduling_config.py new file mode 100644 index 000000000000..4720a54c45e1 --- /dev/null +++ b/components/src/dynamo/planner/tests/config/test_batch_scheduling_config.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Focused validation tests for native batch scheduling configuration.""" + +from __future__ import annotations + +from copy import deepcopy + +import pytest +from dynamo.planner.config.planner_config import PlannerConfig +from pydantic import ValidationError + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +@pytest.fixture(autouse=True) +def _clear_batch_redis_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DYN_PLANNER_BATCH_REDIS_URL", raising=False) + + +def _batch_config() -> dict[str, object]: + return { + "enabled": True, + "gateway": { + "base_url": "http://batch-gateway-apiserver:8000", + "tenant": "planner-poc", + }, + "metrics": { + "frontend_metrics_url": "http://frontend:8000/metrics", + "dispatcher_metrics_url": "http://llm-d-async:9090/metrics", + "online_match_labels": {"request_type": "stream"}, + }, + "redis": { + "url": "redis://batch-gateway-valkey:6379/0", + "control_key": "llm-d-async:drain-limit:dynamo-batch", + }, + "pool": { + "pool_id": "dynamo-batch", + "work_class": "gsm8k-128", + "safe_rps_per_ready_replica": 10.0, + "cold_start_margin_seconds": 15.0, + "finalization_margin_seconds": 5.0, + "max_observation_age_seconds": 20.0, + "drain_lease_duration_seconds": 60.0, + "min_replicas": 1, + "max_replicas": 8, + "scale_from_zero_replicas": 2, + "max_batch_admission_rps": 25.0, + }, + } + + +def _planner_config( + batch_config: dict[str, object] | None = None, + **overrides: object, +) -> PlannerConfig: + raw: dict[str, object] = { + "namespace": "test-ns", + "environment": "kubernetes", + "mode": "agg", + "batch_scheduling": _batch_config() if batch_config is None else batch_config, + } + raw.update(overrides) + return PlannerConfig.model_validate(raw) + + +def test_batch_scheduling_is_disabled_and_empty_by_default() -> None: + config = PlannerConfig(namespace="test-ns") + + assert config.batch_scheduling.enabled is False + assert config.batch_scheduling.gateway is None + assert config.batch_scheduling.metrics is None + assert config.batch_scheduling.redis is None + assert config.batch_scheduling.pool is None + + +def test_enabled_batch_config_maps_exactly_to_policy_config() -> None: + config = _planner_config() + + policy = config.batch_scheduling.to_policy_config() + + assert policy.pool_id == "dynamo-batch" + assert policy.work_class == "gsm8k-128" + assert policy.safe_rps_per_ready_replica == 10.0 + assert policy.cold_start_margin_s == 15.0 + assert policy.finalization_margin_s == 5.0 + assert policy.max_observation_age_s == 20.0 + assert policy.drain_lease_duration_s == 60.0 + assert policy.min_replicas == 1 + assert policy.max_replicas == 8 + assert policy.scale_from_zero_replicas == 2 + assert policy.max_batch_admission_rps == 25.0 + + +def test_batch_pool_scale_from_zero_defaults_to_one_replica() -> None: + batch = _batch_config() + pool = batch["pool"] + assert isinstance(pool, dict) + pool.pop("scale_from_zero_replicas") + + config = _planner_config(batch) + + assert config.batch_scheduling.pool is not None + assert config.batch_scheduling.pool.scale_from_zero_replicas == 1 + assert config.batch_scheduling.to_policy_config().scale_from_zero_replicas == 1 + + +@pytest.mark.parametrize("scale_from_zero_replicas", [0, 9]) +def test_batch_pool_rejects_invalid_scale_from_zero_floor( + scale_from_zero_replicas: int, +) -> None: + batch = _batch_config() + pool = batch["pool"] + assert isinstance(pool, dict) + pool["scale_from_zero_replicas"] = scale_from_zero_replicas + + with pytest.raises(ValidationError, match="scale_from_zero_replicas"): + _planner_config(batch) + + +def test_direct_metrics_config_has_safe_stream_selector_default() -> None: + batch = _batch_config() + metrics = batch["metrics"] + assert isinstance(metrics, dict) + metrics.pop("online_match_labels") + + config = _planner_config(batch) + + assert config.batch_scheduling.metrics is not None + assert config.batch_scheduling.metrics.online_match_labels == { + "request_type": "stream" + } + + +def test_redis_url_loads_from_env_and_is_excluded_from_serialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis_url = "redis://batch-gateway-valkey:6379/0" + monkeypatch.setenv("DYN_PLANNER_BATCH_REDIS_URL", redis_url) + batch = _batch_config() + redis_config = batch["redis"] + assert isinstance(redis_config, dict) + redis_config.pop("url") + + config = _planner_config(batch) + + assert config.batch_scheduling.redis is not None + assert config.batch_scheduling.redis.url is not None + assert config.batch_scheduling.redis.url.get_secret_value() == redis_url + dumped = config.model_dump(mode="json") + dumped_batch = dumped["batch_scheduling"] + assert isinstance(dumped_batch, dict) + dumped_redis = dumped_batch["redis"] + assert isinstance(dumped_redis, dict) + assert "url" not in dumped_redis + assert redis_url not in config.model_dump_json() + + +@pytest.mark.parametrize("missing", ["gateway", "metrics", "redis", "pool"]) +def test_enabled_batch_config_requires_every_subtree(missing: str) -> None: + batch = _batch_config() + batch.pop(missing) + + with pytest.raises(ValidationError, match=missing): + _planner_config(batch) + + +def test_enabled_batch_config_requires_redis_url_or_env() -> None: + batch = _batch_config() + redis_config = batch["redis"] + assert isinstance(redis_config, dict) + redis_config.pop("url") + + with pytest.raises(ValidationError, match="DYN_PLANNER_BATCH_REDIS_URL"): + _planner_config(batch) + + +def test_enabled_batch_scheduling_is_kubernetes_only() -> None: + with pytest.raises(ValidationError, match="environment='kubernetes'"): + _planner_config(environment="virtual") + + +def test_enabled_batch_scheduling_is_aggregate_only() -> None: + with pytest.raises(ValidationError, match="mode='agg'"): + _planner_config(mode="disagg") + + +def test_batch_pool_ceiling_covers_effective_aggregate_minimum() -> None: + batch = _batch_config() + pool = batch["pool"] + assert isinstance(pool, dict) + pool["max_replicas"] = 2 + + with pytest.raises(ValidationError, match="effective aggregate endpoint minimum"): + _planner_config(batch, min_endpoint=3) + + +@pytest.mark.parametrize( + ("tick_max_duration", "scale_interval", "minimum_lease"), + [ + (20.0, 5.0, 25.0), + (2.0, 5.0, 10.0), + ], +) +def test_drain_lease_covers_tick_deadline_and_two_tick_intervals( + tick_max_duration: float, + scale_interval: float, + minimum_lease: float, +) -> None: + batch = _batch_config() + pool = batch["pool"] + assert isinstance(pool, dict) + pool["drain_lease_duration_seconds"] = minimum_lease - 0.001 + + with pytest.raises(ValidationError, match=f"{minimum_lease}s"): + _planner_config( + batch, + scheduling={ + "tick_max_duration_seconds": tick_max_duration, + "scale_interval_seconds": scale_interval, + }, + ) + + pool["drain_lease_duration_seconds"] = minimum_lease + config = _planner_config( + batch, + scheduling={ + "tick_max_duration_seconds": tick_max_duration, + "scale_interval_seconds": scale_interval, + }, + ) + assert config.batch_scheduling.pool is not None + assert config.batch_scheduling.pool.drain_lease_duration_seconds == minimum_lease + + +def test_default_sixty_second_lease_is_valid_for_five_second_ticks() -> None: + config = _planner_config() + + assert config.scheduling.scale_interval_seconds == 5.0 + assert config.scheduling.tick_max_duration_seconds == 30.0 + assert config.batch_scheduling.pool is not None + assert config.batch_scheduling.pool.drain_lease_duration_seconds == 60.0 + + +@pytest.mark.parametrize("subtree", ["batch", "gateway", "metrics", "redis", "pool"]) +def test_batch_config_rejects_unknown_nested_fields(subtree: str) -> None: + batch = deepcopy(_batch_config()) + if subtree == "batch": + target = batch + else: + target = batch[subtree] + assert isinstance(target, dict) + target["unknown_field"] = "unsafe" + + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + _planner_config(batch) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("frontend_metrics_url", "prometheus.monitoring:9090/metrics"), + ("dispatcher_metrics_url", "ftp://llm-d-async/metrics"), + ], +) +def test_metrics_urls_must_be_absolute_http_urls(field: str, value: str) -> None: + batch = _batch_config() + metrics = batch["metrics"] + assert isinstance(metrics, dict) + metrics[field] = value + + with pytest.raises(ValidationError, match=field): + _planner_config(batch) + + +@pytest.mark.parametrize( + "labels", + [ + {}, + {"not-a-label": "stream"}, + {"request_type": ""}, + ], +) +def test_online_metric_selector_must_be_explicit_and_valid( + labels: dict[str, str], +) -> None: + batch = _batch_config() + metrics = batch["metrics"] + assert isinstance(metrics, dict) + metrics["online_match_labels"] = labels + + with pytest.raises(ValidationError, match="online_match_labels"): + _planner_config(batch) + + +def test_policy_conversion_rejects_disabled_config() -> None: + config = PlannerConfig(namespace="test-ns") + + with pytest.raises(ValueError, match="must be enabled"): + config.batch_scheduling.to_policy_config() diff --git a/components/src/dynamo/planner/tests/unit/test_batch_policy.py b/components/src/dynamo/planner/tests/unit/test_batch_policy.py new file mode 100644 index 000000000000..dafa8face98e --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_batch_policy.py @@ -0,0 +1,603 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic tests for the pure single-pool batch scheduling policy.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +from dynamo.planner.core.batch_policy import ( + BatchSchedulingPolicyConfig, + plan_batch_schedule, +) +from dynamo.planner.core.types import ( + BatchDispatcherFeedback, + BatchJobDemand, + BatchSchedulingObservation, + PoolTrafficDemand, + TickInput, + WorkerCounts, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +NOW_S = 1_700_000_000.0 + + +def _config(**overrides: object) -> BatchSchedulingPolicyConfig: + config = BatchSchedulingPolicyConfig( + pool_id="pool-a", + work_class="chat-8k", + safe_rps_per_ready_replica=10.0, + cold_start_margin_s=0.0, + finalization_margin_s=0.0, + max_observation_age_s=30.0, + drain_lease_duration_s=15.0, + min_replicas=1, + max_replicas=20, + scale_from_zero_replicas=1, + ) + return replace(config, **overrides) + + +def _job( + job_id: str, + *, + remaining_requests: int, + deadline_offset_s: float | None, + observed_offset_s: float = 0.0, + work_class: str = "chat-8k", + status: str = "in_progress", +) -> BatchJobDemand: + total_requests = max(1_000, remaining_requests) + completed_requests = total_requests - remaining_requests + return BatchJobDemand( + observed_at_s=NOW_S + observed_offset_s, + pool_id="pool-a", + job_id=job_id, + status=status, + total_requests=total_requests, + completed_requests=completed_requests, + failed_requests=0, + deadline_at_s=( + NOW_S + deadline_offset_s if deadline_offset_s is not None else None + ), + work_class=work_class, + ) + + +def _tick_input( + *, + ready_replicas: int | None = 10, + online_offered_rps: float | None = 90.0, + online_observed_offset_s: float = 0.0, + jobs: list[BatchJobDemand] | None = None, + dispatcher_feedback: list[BatchDispatcherFeedback] | None = None, +) -> TickInput: + worker_counts = ( + WorkerCounts(ready_num_decode=ready_replicas) + if ready_replicas is not None + else None + ) + return TickInput( + now_s=NOW_S, + worker_counts=worker_counts, + batch=BatchSchedulingObservation( + job_demands=jobs or [], + pool_traffic=( + [ + PoolTrafficDemand( + observed_at_s=NOW_S + online_observed_offset_s, + pool_id="pool-a", + online_offered_rps=online_offered_rps, + ) + ] + if online_offered_rps is not None + else [] + ), + dispatcher_feedback=dispatcher_feedback or [], + ), + ) + + +def _plan( + tick_input: TickInput, + config: BatchSchedulingPolicyConfig | None = None, +): + return plan_batch_schedule( + tick_input, + config or _config(), + decision_id="decision-123", + ) + + +def test_work_conserving_headroom_drives_leased_drain_cap(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=600, + deadline_offset_s=100.0, + ) + ] + ) + ) + + assert plan.replica_floor == 10 + assert plan.drain_limit.pool_id == "pool-a" + assert plan.drain_limit.max_admission_rps == 10.0 + assert plan.drain_limit.valid_until_s == NOW_S + 15.0 + assert plan.drain_limit.decision_id == "decision-123" + assert plan.diagnostics.required_batch_rps == 6.0 + assert plan.diagnostics.current_safe_headroom_rps == 10.0 + assert plan.diagnostics.predicted_finish_at_s == NOW_S + 60.0 + assert plan.diagnostics.minimum_deadline_slack_s == 40.0 + assert plan.diagnostics.infeasible is False + + +def test_optional_batch_cap_limits_drain_and_predicted_finish(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=600, + deadline_offset_s=100.0, + ) + ] + ), + _config(max_batch_admission_rps=8.0), + ) + + assert plan.drain_limit.max_admission_rps == 8.0 + assert plan.diagnostics.planned_batch_capacity_rps == 8.0 + assert plan.diagnostics.predicted_finish_at_s == NOW_S + 75.0 + assert plan.diagnostics.minimum_deadline_slack_s == 25.0 + assert plan.diagnostics.infeasible is False + + +def test_edf_uses_cumulative_work_across_multiple_deadlines(): + plan = _plan( + _tick_input( + ready_replicas=5, + online_offered_rps=20.0, + jobs=[ + _job("job-a", remaining_requests=100, deadline_offset_s=10.0), + _job("job-b", remaining_requests=300, deadline_offset_s=20.0), + ], + ) + ) + + # Individual rates would be 10 and 15 RPS. EDF cumulative demand at the + # second deadline is (100 + 300) / 20 = 20 RPS, which is the binding bound. + assert plan.diagnostics.required_batch_rps == 20.0 + assert plan.diagnostics.required_replica_floor == 4 + # Current headroom permits a 30 RPS drain, so the coupled floor remains at + # five replicas instead of allowing an unsafe simultaneous scale-down. + assert plan.replica_floor == 5 + assert plan.drain_limit.max_admission_rps == 30.0 + assert plan.diagnostics.current_safe_headroom_rps == 30.0 + assert plan.diagnostics.planned_batch_capacity_rps == 30.0 + assert plan.diagnostics.predicted_finish_at_s == pytest.approx(NOW_S + 400.0 / 30.0) + assert plan.diagnostics.minimum_deadline_slack_s == pytest.approx(20.0 / 3.0) + assert plan.diagnostics.infeasible is False + + +def test_deadline_required_rps_scales_the_absolute_replica_floor(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=2_000, + deadline_offset_s=100.0, + ) + ] + ) + ) + + assert plan.diagnostics.required_batch_rps == 20.0 + assert plan.diagnostics.required_replica_floor == 11 + assert plan.replica_floor == 11 + # Drain remains bounded by *current* safe headroom until replica 11 is ready. + assert plan.drain_limit.max_admission_rps == 10.0 + assert plan.diagnostics.planned_batch_capacity_rps == 20.0 + + +def test_cold_start_and_finalization_margins_reduce_deadline_window(): + plan = _plan( + _tick_input( + ready_replicas=1, + online_offered_rps=0.0, + jobs=[ + _job( + "job-a", + remaining_requests=800, + deadline_offset_s=100.0, + ) + ], + ), + _config(cold_start_margin_s=10.0, finalization_margin_s=10.0), + ) + + assert plan.diagnostics.required_batch_rps == 10.0 + assert plan.replica_floor == 1 + assert plan.diagnostics.predicted_finish_at_s == NOW_S + 90.0 + assert plan.diagnostics.minimum_deadline_slack_s == 10.0 + + +def test_replica_ceiling_surfaces_infeasible_deadline_but_drains_safely(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=2_000, + deadline_offset_s=100.0, + ) + ] + ), + _config(max_replicas=10), + ) + + assert plan.replica_floor == 10 + assert plan.drain_limit.max_admission_rps == 10.0 + assert plan.diagnostics.infeasible is True + assert "required_replica_floor_exceeds_max_replicas" in ( + plan.diagnostics.infeasible_reasons + ) + assert "negative_deadline_slack" in plan.diagnostics.infeasible_reasons + assert plan.diagnostics.minimum_deadline_slack_s == -100.0 + + +def test_batch_cap_below_required_rate_is_diagnosed_infeasible(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=1_000, + deadline_offset_s=100.0, + ) + ] + ), + _config(max_batch_admission_rps=5.0), + ) + + assert plan.drain_limit.max_admission_rps == 5.0 + assert plan.diagnostics.infeasible is True + assert "batch_cap_below_required_rps" in plan.diagnostics.infeasible_reasons + + +def test_zero_batch_cap_is_an_explicit_pause_with_fresh_inputs(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=100, + deadline_offset_s=None, + ) + ] + ), + _config(max_batch_admission_rps=0.0), + ) + + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.required_batch_rps == 0.0 + assert plan.diagnostics.safety_paused is False + + +def test_exhausted_deadline_window_uses_max_floor_and_reports_infinite_rate(): + plan = _plan( + _tick_input( + jobs=[ + _job( + "job-a", + remaining_requests=100, + deadline_offset_s=10.0, + ) + ] + ), + _config(cold_start_margin_s=5.0, finalization_margin_s=5.0), + ) + + assert plan.replica_floor == 20 + assert plan.diagnostics.required_batch_rps == float("inf") + assert "deadline_window_exhausted:job-a" in (plan.diagnostics.infeasible_reasons) + + +def test_stale_online_traffic_pauses_drain_and_holds_ready_floor(): + plan = _plan( + _tick_input( + ready_replicas=7, + online_observed_offset_s=-31.0, + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=100.0)], + ) + ) + + assert plan.replica_floor == 7 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.safety_paused is True + assert plan.diagnostics.required_batch_rps is None + assert plan.diagnostics.online_observation_stale is True + assert "online_traffic_observation_stale_or_invalid" in ( + plan.diagnostics.infeasible_reasons + ) + + +def test_missing_capacity_pauses_drain_and_uses_max_floor(): + plan = _plan( + _tick_input( + ready_replicas=None, + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=100.0)], + ) + ) + + assert plan.replica_floor == 20 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.capacity_observation_stale is True + assert "capacity_observation_missing_or_invalid" in ( + plan.diagnostics.infeasible_reasons + ) + + +def test_missing_batch_snapshot_pauses_and_holds_ready_floor(): + plan = _plan( + TickInput( + now_s=NOW_S, + worker_counts=WorkerCounts(ready_num_decode=7), + batch=None, + ) + ) + + assert plan.replica_floor == 7 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.required_batch_rps is None + assert "batch_observation_missing" in plan.diagnostics.infeasible_reasons + + +def test_fresh_active_job_bootstraps_zero_pool_while_missing_traffic_pauses_drain(): + plan = _plan( + _tick_input( + ready_replicas=0, + online_offered_rps=None, + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)], + ), + _config(min_replicas=0, max_replicas=4), + ) + + assert plan.replica_floor == 1 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.safety_paused is True + assert plan.diagnostics.required_replica_floor == 1 + assert plan.diagnostics.active_job_count == 1 + assert plan.diagnostics.remaining_requests == 100 + assert "online_traffic_observation_missing" in (plan.diagnostics.infeasible_reasons) + + +def test_explicit_zero_traffic_is_fresh_and_still_bootstraps_zero_pool(): + plan = _plan( + _tick_input( + ready_replicas=0, + online_offered_rps=0.0, + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)], + ), + _config(min_replicas=0, max_replicas=4), + ) + + assert plan.replica_floor == 1 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.safety_paused is False + assert plan.diagnostics.required_batch_rps == 0.0 + assert plan.diagnostics.required_replica_floor == 1 + assert plan.diagnostics.online_offered_rps == 0.0 + + +def test_missing_traffic_without_active_jobs_does_not_bootstrap_zero_pool(): + plan = _plan( + _tick_input( + ready_replicas=0, + online_offered_rps=None, + jobs=[], + ), + _config(min_replicas=0, max_replicas=4), + ) + + assert plan.replica_floor == 0 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.required_replica_floor is None + + +def test_stale_active_job_does_not_bootstrap_zero_pool(): + plan = _plan( + _tick_input( + ready_replicas=0, + online_offered_rps=None, + jobs=[ + _job( + "job-a", + remaining_requests=100, + deadline_offset_s=None, + observed_offset_s=-31.0, + ) + ], + ), + _config(min_replicas=0, max_replicas=4), + ) + + assert plan.replica_floor == 0 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.required_replica_floor is None + assert plan.diagnostics.job_observation_stale is True + + +def test_missing_traffic_holds_nonzero_ready_pool_without_extra_bootstrap(): + plan = _plan( + _tick_input( + ready_replicas=2, + online_offered_rps=None, + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)], + ), + _config(min_replicas=0, max_replicas=4), + ) + + assert plan.replica_floor == 2 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.required_replica_floor is None + + +def test_stale_active_job_pauses_drain_and_holds_ready_floor(): + plan = _plan( + _tick_input( + ready_replicas=7, + jobs=[ + _job( + "job-a", + remaining_requests=100, + deadline_offset_s=100.0, + observed_offset_s=-31.0, + ) + ], + ) + ) + + assert plan.replica_floor == 7 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.job_observation_stale is True + + +def test_work_class_mismatch_pauses_and_holds_clipped_ready_floor(): + plan = _plan( + _tick_input( + ready_replicas=25, + jobs=[ + _job( + "job-a", + remaining_requests=100, + deadline_offset_s=100.0, + work_class="embeddings", + ) + ], + ) + ) + + assert plan.replica_floor == 20 + assert plan.drain_limit.max_admission_rps == 0.0 + assert plan.diagnostics.infeasible is True + assert "work_class_mismatch:job-a:embeddings" in ( + plan.diagnostics.infeasible_reasons + ) + + +def test_latest_raw_job_counters_win_and_terminal_jobs_do_not_add_demand(): + old = _job( + "job-a", + remaining_requests=900, + deadline_offset_s=100.0, + observed_offset_s=-10.0, + ) + latest = _job( + "job-a", + remaining_requests=300, + deadline_offset_s=100.0, + ) + terminal = _job( + "job-b", + remaining_requests=500, + deadline_offset_s=50.0, + status="completed", + ) + + plan = _plan(_tick_input(jobs=[old, latest, terminal])) + + assert plan.diagnostics.active_job_count == 1 + assert plan.diagnostics.remaining_requests == 300 + assert plan.diagnostics.required_batch_rps == 3.0 + + +def test_dispatcher_feedback_is_diagnostic_and_freshness_is_explicit(): + feedback = BatchDispatcherFeedback( + observed_at_s=NOW_S, + pool_id="pool-a", + observation_window_s=10.0, + queued_requests=100, + inflight_requests=5, + actual_dispatch_rps=7.5, + applied_max_admission_rps=8.0, + ) + plan = _plan( + _tick_input( + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)], + dispatcher_feedback=[feedback], + ) + ) + + assert plan.diagnostics.dispatcher_feedback_stale is False + assert plan.diagnostics.dispatcher_observation_age_s == 0.0 + assert plan.diagnostics.actual_dispatch_rps == 7.5 + assert plan.diagnostics.applied_max_admission_rps == 8.0 + + +def test_best_effort_drain_keeps_a_floor_that_can_sustain_it(): + plan = _plan( + _tick_input( + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)] + ) + ) + + assert plan.diagnostics.required_batch_rps == 0.0 + assert plan.drain_limit.max_admission_rps == 10.0 + assert plan.diagnostics.required_replica_floor == 9 + assert plan.replica_floor == 10 + assert plan.drain_limit.max_admission_rps <= (plan.replica_floor * 10.0 - 90.0) + assert plan.diagnostics.minimum_deadline_slack_s is None + + +def test_drain_is_clipped_when_replica_ceiling_cannot_sustain_headroom(): + plan = _plan( + _tick_input( + jobs=[_job("job-a", remaining_requests=100, deadline_offset_s=None)] + ), + _config(max_replicas=9), + ) + + assert plan.diagnostics.current_safe_headroom_rps == 10.0 + assert plan.diagnostics.required_replica_floor == 9 + assert plan.replica_floor == 9 + assert plan.diagnostics.planned_batch_capacity_rps == 0.0 + assert plan.drain_limit.max_admission_rps == 0.0 + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"safe_rps_per_ready_replica": 0.0}, "positive and finite"), + ({"cold_start_margin_s": -1.0}, "non-negative and finite"), + ({"drain_lease_duration_s": 0.0}, "positive and finite"), + ({"min_replicas": -1}, "min_replicas"), + ({"min_replicas": 2, "max_replicas": 1}, "max_replicas"), + ({"scale_from_zero_replicas": 0}, "scale_from_zero_replicas"), + ( + {"scale_from_zero_replicas": 21}, + "scale_from_zero_replicas must be <= max_replicas", + ), + ({"max_batch_admission_rps": -1.0}, "non-negative and finite"), + ], +) +def test_config_rejects_unsafe_values(overrides: dict[str, object], message: str): + with pytest.raises(ValueError, match=message): + _config(**overrides) + + +def test_decision_id_is_required(): + with pytest.raises(ValueError, match="decision_id"): + plan_batch_schedule(_tick_input(), _config(), decision_id="")