diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index 4e96deff1b8b..a1a2178182d0 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -24,6 +24,7 @@ from dynamo.common.configuration.arg_group import ArgGroup from dynamo.common.configuration.config_base import ConfigBase from dynamo.common.configuration.utils import ( + ArgEnv, Deprecated, add_argument, add_negatable_bool_argument, @@ -297,9 +298,17 @@ def kv_router_kwargs(self) -> dict: class KvRouterArgGroup(ArgGroup): - """CLI arguments for the shared KvRouterConfig parameters.""" + """CLI arguments for the shared KvRouterConfig parameters. + + Args: + env: Another process's env (see ``ArgEnv``). + """ + + def __init__(self, *, env: Optional[ArgEnv] = None) -> None: + self.env = env def add_arguments(self, parser) -> None: + environ = os.environ if self.env is None else self.env.environ g = parser.add_argument_group("KV Router Options") add_negatable_bool_argument( @@ -320,6 +329,7 @@ def add_arguments(self, parser) -> None: "and prefill-token load tracking, and disables remote/shared cache indexers. " "The builtin policy does not request cache inputs in this mode." ), + env=self.env, ) add_argument( g, @@ -334,6 +344,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="overlap_score_credit", + env=self.env, ) add_argument( g, @@ -349,6 +360,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="overlap_score_credit_decay", + env=self.env, ) add_argument( g, @@ -356,7 +368,7 @@ def add_arguments(self, parser) -> None: obsolete_flag="--kv-overlap-score-weight", env_var=( "DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT" - if "DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT" in os.environ + if "DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT" in environ else "DYN_OVERLAP_SCORE_WEIGHT" ), dest="overlap_score_weight", @@ -364,6 +376,7 @@ def add_arguments(self, parser) -> None: deprecated=_policy_parameter("prefill_load_scale"), default=None, help=argparse.SUPPRESS, + env=self.env, ) add_argument( g, @@ -377,6 +390,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="prefill_load_scale", + env=self.env, ) add_argument( g, @@ -392,6 +406,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="decode_active_request_weight", + env=self.env, ) add_argument( g, @@ -406,6 +421,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="host_cache_hit_weight", + env=self.env, ) add_argument( g, @@ -419,6 +435,7 @@ def add_arguments(self, parser) -> None: ), arg_type=float, dest="disk_cache_hit_weight", + env=self.env, ) add_argument( g, @@ -431,6 +448,7 @@ def add_arguments(self, parser) -> None: "Higher values promote more randomness, and 0 falls back to deterministic." ), arg_type=float, + env=self.env, ) add_negatable_bool_argument( g, @@ -444,6 +462,7 @@ def add_arguments(self, parser) -> None: ), dest="use_kv_events", obsolete_flag="--kv-events", + env=self.env, ) add_negatable_bool_argument( g, @@ -454,6 +473,7 @@ def add_arguments(self, parser) -> None: "KV Router: Enable best-effort active-sequence synchronization through " "the Runtime event plane." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -466,6 +486,7 @@ def add_arguments(self, parser) -> None: "By default, active blocks are tracked for load balancing." ), obsolete_flag="--track-active-blocks", + env=self.env, ) add_negatable_bool_argument( g, @@ -480,6 +501,7 @@ def add_arguments(self, parser) -> None: "suffix; shared prompt blocks retain full weight." ), obsolete_flag="--track-output-blocks", + env=self.env, ) add_negatable_bool_argument( g, @@ -493,6 +515,7 @@ def add_arguments(self, parser) -> None: "(when KV cache reuse is not expected)." ), obsolete_flag="--assume-kv-reuse", + env=self.env, ) add_negatable_bool_argument( g, @@ -505,6 +528,7 @@ def add_arguments(self, parser) -> None: "Use --no-router-track-prefill-tokens to ignore prompt tokens in router " "prefill-token load, queue pressure, and active_prefill_tokens metrics." ), + env=self.env, ) add_argument( g, @@ -513,6 +537,7 @@ def add_arguments(self, parser) -> None: default="public-xxh3-v1", choices=["public-xxh3-v1", "keyed-xxh3-v1"], help="KV Router: Hash function for router-derived active-sequence identities.", + env=self.env, ) add_argument( g, @@ -520,6 +545,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_ROUTER_TRACKING_KEY_FILE", default=None, help="KV Router: File containing the 32-byte provider tracking key.", + env=self.env, ) add_argument( g, @@ -527,6 +553,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_ROUTER_TRACKING_KEY_ID", default=None, help="KV Router: Provider-managed tracking-key epoch identifier.", + env=self.env, ) add_argument( g, @@ -540,6 +567,7 @@ def add_arguments(self, parser) -> None: "'none' keeps static prompt load accounting. " "'ais' decays the oldest active prefill request using AISimulate-predicted duration." ), + env=self.env, ) add_argument( g, @@ -551,6 +579,7 @@ def add_arguments(self, parser) -> None: "Only used when --no-router-kv-events is set." ), arg_type=float, + env=self.env, ) add_argument( g, @@ -564,6 +593,7 @@ def add_arguments(self, parser) -> None: "the existing time-based behavior; 'lru' uses each worker rank's " "advertised KV capacity. Side indexers remain TTL-only." ), + env=self.env, ) add_argument( g, @@ -588,6 +618,7 @@ def add_arguments(self, parser) -> None: "explicitly for predictable semantics, or use a smaller threshold." ), arg_type=nullable_float, + env=self.env, ) add_argument( g, @@ -602,6 +633,7 @@ def add_arguments(self, parser) -> None: "router_queue_threshold is set." ), arg_type=str, + env=self.env, ) add_negatable_bool_argument( g, @@ -616,6 +648,7 @@ def add_arguments(self, parser) -> None: "for each request." ), dest="conditional_disagg_enabled", + env=self.env, ) add_argument( g, @@ -631,6 +664,7 @@ def add_arguments(self, parser) -> None: arg_type=_conditional_disagg_config_arg, dest="conditional_disagg_config", metavar="JSON", + env=self.env, ) add_argument( g, @@ -647,6 +681,7 @@ def add_arguments(self, parser) -> None: "Use 'default' for Dynamo's built-in worker selector." ), arg_type=str, + env=self.env, ) add_argument( g, @@ -663,6 +698,7 @@ def add_arguments(self, parser) -> None: "Use 'default' for Dynamo's built-in worker selector." ), arg_type=str, + env=self.env, ) add_argument( g, @@ -675,6 +711,7 @@ def add_arguments(self, parser) -> None: "approximate routing when --no-router-kv-events is set." ), arg_type=int, + env=self.env, ) add_argument( g, @@ -692,6 +729,7 @@ def add_arguments(self, parser) -> None: ), arg_type=str, choices=["fcfs", "wspt"], + env=self.env, ) add_negatable_bool_argument( g, @@ -703,6 +741,7 @@ def add_arguments(self, parser) -> None: "component via the request plane instead of maintaining a local radix tree." ), dest="use_remote_indexer", + env=self.env, ) add_negatable_bool_argument( g, @@ -715,6 +754,7 @@ def add_arguments(self, parser) -> None: "index for routing decisions yet." ), dest="enable_session_prefix_index", + env=self.env, ) add_argument( g, @@ -730,6 +770,7 @@ def add_arguments(self, parser) -> None: "The default policy uses 0.5 when shared cache is enabled." ), arg_type=float, + env=self.env, ) add_argument( g, @@ -744,6 +785,7 @@ def add_arguments(self, parser) -> None: ), arg_type=str, choices=["none", "hicache"], + env=self.env, ) add_argument( g, @@ -756,4 +798,5 @@ def add_arguments(self, parser) -> None: "Independent of --router-ttl-secs, which covers pure approximate mode." ), arg_type=float, + env=self.env, ) diff --git a/components/src/dynamo/common/configuration/groups/router_args.py b/components/src/dynamo/common/configuration/groups/router_args.py index 728c76ea05fa..243fc36153e0 100644 --- a/components/src/dynamo/common/configuration/groups/router_args.py +++ b/components/src/dynamo/common/configuration/groups/router_args.py @@ -24,7 +24,12 @@ KvRouterArgGroup, KvRouterConfigBase, ) -from dynamo.common.configuration.utils import add_argument, nullable_float, nullable_int +from dynamo.common.configuration.utils import ( + ArgEnv, + add_argument, + nullable_float, + nullable_int, +) if TYPE_CHECKING: from dynamo.llm import RouterConfig @@ -64,7 +69,14 @@ ) -class _IgnoredAdmissionControlAction(argparse.Action): +class _IgnoredAction(argparse.Action): + """Store nothing, so the namespace never carries the value.""" + + def __call__(self, parser, namespace, values, option_string=None): + return None + + +class _IgnoredAdmissionControlAction(_IgnoredAction): """Warn and store nothing, so the namespace never carries the value.""" def __call__(self, parser, namespace, values, option_string=None): @@ -165,6 +177,8 @@ class RouterArgGroup(ArgGroup): consumes. ``--router-min-initial-workers`` gates frontend startup and is not carried on the model card, so a worker registering it would ship a flag that does nothing. + env: Another process's env (see ``ArgEnv``); the removed and + deprecated flags then parse without their notices. """ def __init__( @@ -172,15 +186,18 @@ def __init__( *, default_router_mode: Optional[str], include_frontend_only: bool, + env: Optional[ArgEnv] = None, ) -> None: self.default_router_mode = default_router_mode self.include_frontend_only = include_frontend_only + self.env = env def add_arguments(self, parser) -> None: - if "DYN_ADMISSION_CONTROL" in os.environ: - logger.warning(_ADMISSION_CONTROL_REMOVAL_WARNING) - if "DYN_ENFORCE_DISAGG" in os.environ: - logger.warning(_ENFORCE_DISAGG_DEPRECATION, "DYN_ENFORCE_DISAGG") + if self.env is None: + if "DYN_ADMISSION_CONTROL" in os.environ: + logger.warning(_ADMISSION_CONTROL_REMOVAL_WARNING) + if "DYN_ENFORCE_DISAGG" in os.environ: + logger.warning(_ENFORCE_DISAGG_DEPRECATION, "DYN_ENFORCE_DISAGG") g = parser.add_argument_group("Router Options") @@ -195,7 +212,11 @@ def add_arguments(self, parser) -> None: g.add_argument( "--admission-control", choices=("token-capacity", "none"), - action=_IgnoredAdmissionControlAction, + action=( + _IgnoredAdmissionControlAction + if self.env is None + else _IgnoredAction + ), default=argparse.SUPPRESS, help=argparse.SUPPRESS, ) @@ -210,7 +231,12 @@ def add_arguments(self, parser) -> None: "readiness are determined from registered worker types." ), arg_type=None, - action=_DeprecatedEnforceDisaggAction, + action=( + _DeprecatedEnforceDisaggAction + if self.env is None + else argparse.BooleanOptionalAction + ), + env=self.env, ) add_argument( g, @@ -225,6 +251,7 @@ def add_arguments(self, parser) -> None: ), arg_type=int, dest="min_initial_workers", + env=self.env, ) add_argument( @@ -249,6 +276,7 @@ def add_arguments(self, parser) -> None: "least-loaded", "device-aware-weighted", ], + env=self.env, ) add_argument( g, @@ -263,6 +291,7 @@ def add_arguments(self, parser) -> None: ), arg_type=int, dest="session_affinity_ttl_secs", + env=self.env, ) add_argument( g, @@ -276,6 +305,7 @@ def add_arguments(self, parser) -> None: ), choices=("hard", "soft"), dest="session_affinity_mode", + env=self.env, ) add_argument( g, @@ -288,6 +318,7 @@ def add_arguments(self, parser) -> None: "Unset by default; pass 'None' to disable it." ), arg_type=nullable_float, + env=self.env, ) add_argument( g, @@ -303,6 +334,7 @@ def add_arguments(self, parser) -> None: "with --active-prefill-tokens-threshold-frac." ), arg_type=nullable_int, + env=self.env, ) add_argument( g, @@ -316,6 +348,7 @@ def add_arguments(self, parser) -> None: "OR logic with --active-prefill-tokens-threshold." ), arg_type=nullable_float, + env=self.env, ) @@ -344,16 +377,19 @@ class WorkerRouterConfig(RouterConfigBase, KvRouterConfigBase): enforce_disagg: bool = False -def add_worker_router_arguments(parser: argparse.ArgumentParser) -> None: +def add_worker_router_arguments( + parser: argparse.ArgumentParser, *, env: Optional[ArgEnv] = None +) -> None: """Register the worker-side router flags. ``--router-mode`` defaults to ``None``: a worker that omits it advertises - nothing and inherits the frontend's configuration. + nothing and inherits the frontend's configuration. ``env`` is another + process's env (see ``ArgEnv``). """ - RouterArgGroup(default_router_mode=None, include_frontend_only=False).add_arguments( - parser - ) - KvRouterArgGroup().add_arguments(parser) + RouterArgGroup( + default_router_mode=None, include_frontend_only=False, env=env + ).add_arguments(parser) + KvRouterArgGroup(env=env).add_arguments(parser) def parse_worker_router_config( diff --git a/components/src/dynamo/common/configuration/utils.py b/components/src/dynamo/common/configuration/utils.py index e3e090b83fe0..8a441be1d47f 100644 --- a/components/src/dynamo/common/configuration/utils.py +++ b/components/src/dynamo/common/configuration/utils.py @@ -7,7 +7,8 @@ import os import re import warnings -from dataclasses import dataclass +from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Any, Callable, Optional, TypeVar, Union T = TypeVar("T") @@ -30,6 +31,23 @@ def warn(self, name: str) -> None: warnings.warn(self.message(name), FutureWarning, stacklevel=3) +@dataclass +class ArgEnv: + """The env of another process, for registering that process's arguments. + + Argument groups read their env defaults from ``os.environ`` and warn + about deprecated or removed settings. A process that parses another + process's configuration (the planner reads its frontend's from the + DynamoGraphDeployment spec) passes an ``ArgEnv`` instead: env defaults + come from ``environ`` only, deprecation and removal notices are skipped + because they belong to the process that runs with that env, and + ``env_dests`` records the env var -> dest of every argument registered. + """ + + environ: Mapping[str, str] + env_dests: dict[str, str] = field(default_factory=dict) + + def parse_bool(value: str) -> bool: """Parse Dynamo's truthy and falsy configuration values.""" normalized = value.strip().lower() @@ -64,6 +82,8 @@ def env_or_default( env_var: str, default: T, value_type: Optional[Union[type, Callable[..., Any]]] = None, + *, + env: Optional[ArgEnv] = None, ) -> T: """ Get value from environment variable or return default. @@ -76,11 +96,12 @@ def env_or_default( value_type: If provided, use this type to convert the env value. If None, the type is taken from type(default). Use value_type when default is None but you still want the env value coerced (e.g. env_or_default("DYN_FOO", None, value_type=int)). + env: Read ``env.environ`` instead of ``os.environ``. Returns: Environment variable value (type-converted) or default """ - value = os.environ.get(env_var) + value = (os.environ if env is None else env.environ).get(env_var) if value is None: return default @@ -129,6 +150,7 @@ def add_argument( arg_type: Optional[Union[type, Callable[..., Any]]] = str, env_value_type: Optional[Union[type, Callable[..., Any]]] = None, deprecated: Optional[Deprecated] = None, + env: Optional[ArgEnv] = None, **kwargs: Any, ) -> None: """ @@ -147,6 +169,7 @@ def add_argument( arg_type: Type for the argument (default: str) env_value_type: Optional parser used only for the environment value deprecated: Warn on explicit CLI/env use and show the migration in help. + env: Another process's env (see ``ArgEnv``); skips ``deprecated``. """ arg_dest = _get_dest_name(flag_name, kwargs.get("dest")) value_type_for_env = env_value_type @@ -154,7 +177,12 @@ def add_argument( value_type_for_env = arg_type if isinstance(default, list) and (arg_type is None or arg_type is str): value_type_for_env = None - default_with_env = env_or_default(env_var, default, value_type=value_type_for_env) + default_with_env = env_or_default( + env_var, default, value_type=value_type_for_env, env=env + ) + if env is not None: + env.env_dests[env_var] = arg_dest + deprecated = None names = [flag_name] @@ -210,6 +238,7 @@ def add_negatable_bool_argument( obsolete_flag: Optional[str] = None, env_value_type: Optional[Callable[..., bool]] = None, deprecated: Optional[Deprecated] = None, + env: Optional[ArgEnv] = None, ) -> None: """ Add negatable boolean flag (--foo / --no-foo). @@ -223,6 +252,7 @@ def add_negatable_bool_argument( dest: Optional destination name for the parsed value obsolete_flag: Optional obsolete/legacy flag (for help msg only, must start with '--') env_value_type: Optional strict parser for the environment value + env: Another process's env (see ``ArgEnv``) """ add_argument( parser, @@ -236,6 +266,7 @@ def add_negatable_bool_argument( env_value_type=env_value_type, action=argparse.BooleanOptionalAction, deprecated=deprecated, + env=env, ) diff --git a/components/src/dynamo/common/tests/configuration/test_utils.py b/components/src/dynamo/common/tests/configuration/test_utils.py index 5c25dbf2da24..32ec739efca5 100644 --- a/components/src/dynamo/common/tests/configuration/test_utils.py +++ b/components/src/dynamo/common/tests/configuration/test_utils.py @@ -8,6 +8,7 @@ import pytest from dynamo.common.configuration.utils import ( + ArgEnv, Deprecated, add_argument, add_negatable_bool_argument, @@ -391,6 +392,35 @@ def test_deprecated_argument_env_and_cli_precedence(monkeypatch): assert parser.parse_args(["--old", "4"]).old == 4 +def test_arg_env_replaces_os_environ_and_skips_deprecation(monkeypatch, recwarn): + monkeypatch.setenv("TEST_OLD", "3") + env = ArgEnv(environ={"TEST_FLAG": "false", "TEST_OLD": "5"}) + parser = argparse.ArgumentParser() + add_argument( + parser, + flag_name="--old", + env_var="TEST_OLD", + default=1, + help="Old", + arg_type=int, + deprecated=Deprecated("--new", remove_in="v2.0"), + env=env, + ) + add_negatable_bool_argument( + parser, + flag_name="--flag", + env_var="TEST_FLAG", + default=True, + help="Flag", + env=env, + ) + assert vars(parser.parse_args([])) == {"old": 5, "flag": False} + assert parser.parse_args(["--old", "4"]).old == 4 + assert env.env_dests == {"TEST_OLD": "old", "TEST_FLAG": "flag"} + assert env_or_default("TEST_OLD", 0, env=ArgEnv(environ={})) == 0 + assert not recwarn + + def test_deprecated_hidden_argument_is_quiet_unless_used(monkeypatch, recwarn): monkeypatch.delenv("TEST_OLD", raising=False) parser = argparse.ArgumentParser() diff --git a/components/src/dynamo/frontend/frontend_args.py b/components/src/dynamo/frontend/frontend_args.py index 67dd92cc28c6..b60180a78fb6 100644 --- a/components/src/dynamo/frontend/frontend_args.py +++ b/components/src/dynamo/frontend/frontend_args.py @@ -4,6 +4,7 @@ import argparse import os import pathlib +from collections.abc import Mapping from typing import Any, Dict, Optional from dynamo.common.config_dump import register_encoder @@ -22,6 +23,7 @@ RouterConfigBase, ) from dynamo.common.configuration.utils import ( + ArgEnv, add_argument, add_negatable_bool_argument, env_or_default, @@ -368,7 +370,17 @@ def _preprocess_for_encode_config(config: FrontendConfig) -> Dict[str, Any]: class FrontendArgGroup(ArgGroup): - """Frontend configuration parameters.""" + """Frontend configuration parameters. + + Args: + env: Another process's env (see ``ArgEnv``). Nothing then reads the + local filesystem: ``--model-path`` is not checked, and the + AISimulate perf-model arguments, whose values name files, are not + registered. + """ + + def __init__(self, *, env: Optional[ArgEnv] = None) -> None: + self.env = env def add_arguments(self, parser) -> None: parser.add_argument( @@ -383,7 +395,7 @@ def add_arguments(self, parser) -> None: "--interactive", dest="interactive", action=argparse.BooleanOptionalAction, - default=env_or_default("DYN_INTERACTIVE", False), + default=env_or_default("DYN_INTERACTIVE", False, env=self.env), help="Interactive text chat.\nenv var: DYN_INTERACTIVE", ) @@ -396,6 +408,7 @@ def add_arguments(self, parser) -> None: "Dynamo namespace for model discovery scoping. Use for exact namespace matching. " "If --namespace-prefix is also specified, prefix takes precedence." ), + env=self.env, ) add_argument( @@ -405,6 +418,7 @@ def add_arguments(self, parser) -> None: default=None, help="KV cache block size (u32).", arg_type=int, + env=self.env, ) add_argument( @@ -413,6 +427,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_HTTP_HOST", default="0.0.0.0", help="HTTP host for the engine (str).", + env=self.env, ) add_argument( g, @@ -421,6 +436,7 @@ def add_arguments(self, parser) -> None: default=8000, help="HTTP port for the engine (u16).", arg_type=int, + env=self.env, ) add_argument( g, @@ -428,6 +444,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_BATCH_GATEWAY_URL", default=None, help="Internal Batch Gateway URL for proxying /v1/files and /v1/batches.", + env=self.env, ) add_negatable_bool_argument( g, @@ -436,6 +453,7 @@ def add_arguments(self, parser) -> None: default=False, help="Serve this frontend's local KV indexers over the request plane.", dest="serve_indexer", + env=self.env, ) add_argument( g, @@ -444,6 +462,7 @@ def add_arguments(self, parser) -> None: default=None, help="TLS certificate path, PEM format.", arg_type=pathlib.Path, + env=self.env, ) add_argument( g, @@ -452,6 +471,7 @@ def add_arguments(self, parser) -> None: default=None, help="TLS certificate key path, PEM format.", arg_type=pathlib.Path, + env=self.env, ) add_argument( g, @@ -460,6 +480,7 @@ def add_arguments(self, parser) -> None: default=None, help="Client CA certificate path for mutual TLS, PEM format.", arg_type=pathlib.Path, + env=self.env, ) add_argument( @@ -468,6 +489,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_TCP_TLS_CERT_PATH", default=None, help="Path to PEM certificate for the TCP server.", + env=self.env, ) add_argument( @@ -476,6 +498,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_TCP_TLS_KEY_PATH", default=None, help="Path to PEM private key for the TCP server certificate.", + env=self.env, ) add_argument( @@ -484,6 +507,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_TCP_TLS_CA_CERT_PATH", default=None, help="Path to PEM CA certificate used to verify the TCP peer's certificate.", + env=self.env, ) add_argument( @@ -492,6 +516,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_TCP_TLS_CLIENT_CERT_PATH", default=None, help="Path to PEM client certificate presented to the TCP server for mTLS.", + env=self.env, ) add_argument( @@ -500,6 +525,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_TCP_TLS_CLIENT_KEY_PATH", default=None, help="Path to PEM private key for the TCP client certificate (mTLS).", + env=self.env, ) add_argument( @@ -509,6 +535,7 @@ def add_arguments(self, parser) -> None: default=None, help="Path to PEM CA certificate the TCP server uses to verify client " "certificates. When set, clients must present a trusted certificate (mTLS enforced).", + env=self.env, ) add_argument( @@ -517,6 +544,7 @@ def add_arguments(self, parser) -> None: env_var="NATS_TLS_CA_CERT_PATH", default=None, help="Path to PEM CA certificate for verifying the NATS server.", + env=self.env, ) add_negatable_bool_argument( @@ -525,6 +553,7 @@ def add_arguments(self, parser) -> None: env_var="NATS_TLS_INSECURE", default=False, help="Disable NATS TLS certificate verification. For local development only.", + env=self.env, ) add_argument( @@ -533,6 +562,7 @@ def add_arguments(self, parser) -> None: env_var="NATS_TLS_CLIENT_CERT_PATH", default=None, help="Path to PEM client certificate presented to the NATS server for mTLS.", + env=self.env, ) add_argument( @@ -541,16 +571,20 @@ def add_arguments(self, parser) -> None: env_var="NATS_TLS_CLIENT_KEY_PATH", default=None, help="Path to PEM private key for the NATS client certificate (mTLS).", + env=self.env, ) # Router options (shared with dynamo.router) RouterArgGroup( - default_router_mode="round-robin", include_frontend_only=True + default_router_mode="round-robin", + include_frontend_only=True, + env=self.env, ).add_arguments(parser) # KV router options (shared with dynamo.router) - KvRouterArgGroup().add_arguments(parser) - AisPerfArgGroup().add_arguments(parser) + KvRouterArgGroup(env=self.env).add_arguments(parser) + if self.env is None: + AisPerfArgGroup().add_arguments(parser) add_argument( g, @@ -562,6 +596,7 @@ def add_arguments(self, parser) -> None: "namespaces starting with this prefix (e.g., 'ns' matches 'ns', 'ns-abc123', " "'ns-def456'). Takes precedence over --namespace if both are specified." ), + env=self.env, ) add_argument( @@ -577,6 +612,7 @@ def add_arguments(self, parser) -> None: "session; unfinished audio is replayed to a compatible warm worker." ), arg_type=int, + env=self.env, ) add_argument( @@ -591,6 +627,7 @@ def add_arguments(self, parser) -> None: "Default: no limit." ), arg_type=int, + env=self.env, ) add_argument( @@ -600,6 +637,7 @@ def add_arguments(self, parser) -> None: default=None, help="Model name as a string (e.g., 'Llama-3.2-1B-Instruct')", arg_type=validate_model_name, + env=self.env, ) add_argument( g, @@ -607,7 +645,8 @@ def add_arguments(self, parser) -> None: env_var="DYN_MODEL_PATH", default=None, help="Path to model directory on disk (e.g., /tmp/model_cache/llama3.2_1B/)", - arg_type=validate_model_path, + arg_type=validate_model_path if self.env is None else str, + env=self.env, ) add_argument( g, @@ -618,6 +657,7 @@ def add_arguments(self, parser) -> None: "Prefix for Dynamo frontend metrics. If unset, uses DYN_METRICS_PREFIX env var " "or 'dynamo_frontend'." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -625,6 +665,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_KSERVE_GRPC_SERVER", default=False, help="Start KServe gRPC server.", + env=self.env, ) add_argument( g, @@ -636,6 +677,7 @@ def add_arguments(self, parser) -> None: "Defaults to 8788." ), arg_type=int, + env=self.env, ) add_argument( @@ -644,6 +686,7 @@ def add_arguments(self, parser) -> None: env_var="DYN_DUMP_CONFIG_TO", default=None, help="Dump config to the specified file path.", + env=self.env, ) add_argument( @@ -659,6 +702,7 @@ def add_arguments(self, parser) -> None: "path. May be repeated. DYN_FRONTEND_ROUTE_EXTENSIONS accepts " "whitespace-separated values." ), + env=self.env, ) add_argument( @@ -672,6 +716,7 @@ def add_arguments(self, parser) -> None: "File uses root dir from env var DYN_FILE_KV or defaults to $TMPDIR/dynamo_store_kv." ), choices=["kubernetes", "etcd", "file", "mem"], + env=self.env, ) add_argument( g, @@ -683,6 +728,7 @@ def add_arguments(self, parser) -> None: "'tcp' is fastest [nats|tcp]" ), choices=["nats", "tcp"], + env=self.env, ) add_argument( g, @@ -691,6 +737,7 @@ def add_arguments(self, parser) -> None: default="tcp", help="Select the response transport. Frontend and workers must match.", choices=["tcp", "quic", "velo"], + env=self.env, ) add_argument( g, @@ -701,6 +748,7 @@ def add_arguments(self, parser) -> None: "defaults to 'zmq' for all discovery backends. Set to 'nats' to use a " "NATS-based event plane.", choices=["nats", "zmq"], + env=self.env, ) add_negatable_bool_argument( g, @@ -711,6 +759,7 @@ def add_arguments(self, parser) -> None: "[EXPERIMENTAL] Enable Anthropic Messages API endpoint (/v1/messages). " "This feature is experimental and may change." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -721,6 +770,7 @@ def add_arguments(self, parser) -> None: "[EXPERIMENTAL] Enable typed decision APIs (/v1/decisions and " "/v1/systemone) for qualified aggregate SGLang models." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -731,6 +781,7 @@ def add_arguments(self, parser) -> None: "Strip the Claude Code billing preamble (x-anthropic-billing-header) " "from the system prompt. Saves tokens and improves prompt caching." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -743,6 +794,7 @@ def add_arguments(self, parser) -> None: "for each complete tool call before finish_reason arrives. " "Can be combined with --enable-streaming-reasoning-dispatch." ), + env=self.env, ) add_negatable_bool_argument( g, @@ -755,6 +807,7 @@ def add_arguments(self, parser) -> None: "with the complete reasoning block once thinking ends. " "Can be combined with --enable-streaming-tool-dispatch." ), + env=self.env, ) add_argument( g, @@ -765,6 +818,7 @@ def add_arguments(self, parser) -> None: "OpenAI-compatible response field used for emitted reasoning content." ), choices=["reasoning_content", "reasoning"], + env=self.env, ) # NOTE: This flag also exists in DynamoRuntimeArgGroup (runtime_args.py). # Both definitions are needed: runtime_args controls the Rust-native @@ -781,6 +835,7 @@ def add_arguments(self, parser) -> None: "tool_choice='none'. Prevents models from generating raw XML " "tool calls in the content field." ), + env=self.env, ) add_argument( g, @@ -795,6 +850,7 @@ def add_arguments(self, parser) -> None: "parsing, and reasoning parsing." ), choices=["dynamo", "vllm", "sglang"], + env=self.env, ) add_negatable_bool_argument( @@ -808,13 +864,16 @@ def add_arguments(self, parser) -> None: "Logs per-function timing, request concurrency, and hot-path section durations. " "Supported with '--dyn-chat-processor vllm' and '--dyn-chat-processor sglang'." ), + env=self.env, ) g.add_argument( "--dyn-preprocess-workers", dest="preprocess_workers", type=int, - default=env_or_default("DYN_PREPROCESS_WORKERS", None, value_type=int), + default=env_or_default( + "DYN_PREPROCESS_WORKERS", None, value_type=int, env=self.env + ), help=( "[EXPERIMENTAL] SGLang preprocessing worker processes per model. " "When unset, selects up to 2 workers per model from visible CPU capacity, " @@ -839,6 +898,7 @@ def add_arguments(self, parser) -> None: "Has no effect on TikToken models." ), choices=["default", "fastokens", "basetenkenizer"], + env=self.env, ) add_negatable_bool_argument( @@ -856,6 +916,7 @@ def add_arguments(self, parser) -> None: "continues running." ), env_value_type=parse_bool, + env=self.env, ) add_negatable_bool_argument( @@ -867,4 +928,28 @@ def add_arguments(self, parser) -> None: "Trust remote code when loading the tokenizer. Required for models " "that ship custom tokenizer code (e.g. Qwen, Falcon)." ), - ) + env=self.env, + ) + + +def build_frontend_parser( + env: Mapping[str, str], + parser: Optional[argparse.ArgumentParser] = None, +) -> tuple[argparse.ArgumentParser, dict[str, str]]: + """The argument parser of a frontend that runs with ``env``. + + For a process that parses another frontend's command line (the planner + reads one from the DynamoGraphDeployment spec): env defaults come from + ``env`` only, never ``os.environ``, and nothing logs, warns or reads the + local filesystem (see ``FrontendArgGroup``). Registers the arguments on + ``parser`` when given, else on a new ``ArgumentParser(add_help=False)``. + Returns the parser and the env var -> dest map of every argument with an + env var. Raises what the frontend raises at startup on invalid env + (``ValueError``, ``TypeError`` or ``argparse.ArgumentTypeError``). + """ + + arg_env = ArgEnv(environ=dict(env)) + if parser is None: + parser = argparse.ArgumentParser(add_help=False) + FrontendArgGroup(env=arg_env).add_arguments(parser) + return parser, arg_env.env_dests diff --git a/components/src/dynamo/frontend/tests/test_frontend_parser_builder.py b/components/src/dynamo/frontend/tests/test_frontend_parser_builder.py new file mode 100644 index 000000000000..5fc49e153b01 --- /dev/null +++ b/components/src/dynamo/frontend/tests/test_frontend_parser_builder.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``build_frontend_parser``: another frontend's parser, against an explicit env.""" + +import argparse +import logging +import os +import sys + +import pytest + +from dynamo.common.configuration.utils import env_or_default +from dynamo.frontend.frontend_args import build_frontend_parser + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + + +def test_reads_only_the_given_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DYN_ROUTER_MODE", "kv") + monkeypatch.setenv("DYN_ROUTER_TEMPERATURE", "0.9") + before = dict(os.environ) + + parser, env_dests = build_frontend_parser({"DYN_ROUTER_MODE": "random"}) + namespace = parser.parse_args([]) + + assert namespace.router_mode == "random" + assert namespace.router_temperature == 0.0 + assert env_dests["DYN_ROUTER_MODE"] == "router_mode" + assert env_dests["DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT"] == "overlap_score_credit" + assert dict(os.environ) == before + + +def test_no_registered_env_var_is_read_from_os_environ( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # An argument registered without ``env=`` would take its default from + # the parsing process's os.environ. Record every env var the + # registrations look up, wherever ``env_or_default`` was imported. + registered: set[str] = set() + + def recording_env_or_default(env_var, default, value_type=None, *, env=None): + registered.add(env_var) + return env_or_default(env_var, default, value_type, env=env) + + for name, module in list(sys.modules.items()): + if name.startswith("dynamo.") and ( + vars(module).get("env_or_default") is env_or_default + ): + monkeypatch.setattr(module, "env_or_default", recording_env_or_default) + build_frontend_parser({}) + assert "DYN_ROUTER_MODE" in registered + for name in registered: + monkeypatch.delenv(name, raising=False) + defaults = vars(build_frontend_parser({})[0].parse_args([])) + + # Two sentinels, so that a leaked env var changes at least one default. + for sentinel in ("0", "1"): + for name in registered: + monkeypatch.setenv(name, sentinel) + parser, _ = build_frontend_parser({}) + assert vars(parser.parse_args([])) == defaults, sentinel + + +def test_parses_without_notices_or_local_files( + caplog: pytest.LogCaptureFixture, recwarn: pytest.WarningsRecorder +) -> None: + env = {"DYN_ADMISSION_CONTROL": "none", "DYN_ENFORCE_DISAGG": "1"} + parser = argparse.ArgumentParser(add_help=False) + with caplog.at_level(logging.WARNING): + assert build_frontend_parser(env, parser)[0] is parser + namespace, unknown = parser.parse_known_args( + [ + "--router-temperature", + "0.5", + "--admission-control", + "none", + "--enforce-disagg", + "--model-path", + "/not/a/local/dir", + "--ais-perf-config", + "/not/a/local/file.yaml", + ] + ) + + assert namespace.router_temperature == 0.5 + assert namespace.model_path == "/not/a/local/dir" + assert not hasattr(namespace, "admission_control") + # The AISimulate perf-model flags name local files and are not registered. + assert unknown == ["--ais-perf-config", "/not/a/local/file.yaml"] + assert not caplog.records + assert not recwarn diff --git a/components/src/dynamo/planner/__main__.py b/components/src/dynamo/planner/__main__.py index e2f495deb106..7c8e24807527 100644 --- a/components/src/dynamo/planner/__main__.py +++ b/components/src/dynamo/planner/__main__.py @@ -16,6 +16,9 @@ import argparse import asyncio import logging +import signal +from collections.abc import Awaitable +from typing import Any, Optional from pydantic import BaseModel @@ -25,6 +28,9 @@ logger = logging.getLogger(__name__) +# How long a stopped planner waits for its endpoint to stop serving. +_SERVE_STOP_TIMEOUT_S = 10.0 + class RequestType(BaseModel): text: str @@ -43,13 +49,87 @@ async def generate(request: RequestType): # once the handler is registered with the transport. Running it concurrently # with planner.run() ensures the system server (/health, /live) reports the # planner as ready only after _async_init() has completed. - await asyncio.gather( - generate_endpoint.serve_endpoint( - generate, # type: ignore[arg-type] - health_check_payload={"text": "health"}, - ), - planner.run(), + serve = generate_endpoint.serve_endpoint( + generate, # type: ignore[arg-type] + health_check_payload={"text": "health"}, ) + if not planner.may_hold_retune_lease: + await asyncio.gather(serve, planner.run()) + return + await _run_until_sigterm(runtime, serve, planner.run(), planner.retune_lease_found) + + +async def _run_until_sigterm( + runtime: Any, + serve: Awaitable[Any], + run: Awaitable[None], + lease_found: asyncio.Event, +) -> None: + """Serve and run the planner; once ``lease_found`` is set, SIGTERM stops + the planner cleanly. + + Live retune, and the clean-up of a removed one, hold a leadership Lease, + and live retune may have a sweep in flight. Without a handler SIGTERM + kills the process before ``run()``'s ``finally`` stops the sweep and + releases the Lease, so the next leader would wait out the Lease. The + handler is installed when the planner may hold the Lease: at once with + live retune, or when the clean-up check finds a Lease. It cancels only + the run task; once its shutdown finished, the runtime is shut down to + stop serving the endpoint and the call returns normally. + """ + + loop = asyncio.get_running_loop() + run_task = asyncio.ensure_future(run) + serve_task = asyncio.ensure_future(serve) + terminating = False + installed = False + + def on_sigterm() -> None: + nonlocal terminating + if terminating: + return # a second SIGTERM must not interrupt the shutdown + terminating = True + logger.info("SIGTERM received; stopping the planner") + run_task.cancel() + + def install() -> None: + nonlocal installed + try: + loop.add_signal_handler(signal.SIGTERM, on_sigterm) + installed = True + except (NotImplementedError, RuntimeError, ValueError) as exc: + # No signal handlers on this platform or off the main thread. + logger.warning("Cannot install a planner SIGTERM handler: %s", exc) + + async def install_once_found() -> None: + await lease_found.wait() + install() + + installer: Optional[asyncio.Future[None]] = None + if lease_found.is_set(): + install() + else: + installer = asyncio.ensure_future(install_once_found()) + try: + await asyncio.gather(serve_task, run_task) + except asyncio.CancelledError: + if not (terminating and run_task.done() and run_task.cancelled()): + raise + logger.info("Planner stopped; shutting down the runtime") + runtime.shutdown() + done, _ = await asyncio.wait({serve_task}, timeout=_SERVE_STOP_TIMEOUT_S) + if not done: + serve_task.cancel() + elif not serve_task.cancelled() and serve_task.exception() is not None: + logger.warning( + "Planner endpoint stopped with an error: %s", serve_task.exception() + ) + finally: + if installer is not None and not installer.done(): + installer.cancel() + await asyncio.wait([installer]) + if installed: + loop.remove_signal_handler(signal.SIGTERM) async def init_planner(runtime: DistributedRuntime, config: PlannerConfig): diff --git a/components/src/dynamo/planner/config/planner_config.py b/components/src/dynamo/planner/config/planner_config.py index c2b2903af8c5..784b198460f7 100644 --- a/components/src/dynamo/planner/config/planner_config.py +++ b/components/src/dynamo/planner/config/planner_config.py @@ -23,7 +23,7 @@ from dataclasses import asdict from enum import Enum from functools import lru_cache -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Annotated, Any, Dict, Literal, Optional, Protocol, cast from urllib.parse import parse_qsl, urlsplit @@ -847,6 +847,463 @@ def replay_kwargs( ) +class DgdrRunConfig(BaseModel): + """How each DGDR v1beta2 Run of a live retune sweep evaluates its points.""" + + model_config = ConfigDict(extra="forbid") + + parallelism: int = Field( + default=8, + ge=1, + description=( + "AISimulate optimizer.parallelism of each Run. The Run Job's CPU " + "and memory may lower the effective value." + ), + ) + candidate_timeout_seconds: float = Field( + default=1800.0, + gt=0, + allow_inf_nan=False, + description="AISimulate optimizer.candidate_timeout_seconds of each Run.", + ) + max_shards: int = Field( + default=3, + ge=1, + le=16, + description=( + "Most grid shards (Requests) per window. A shard is a Cartesian " + "sub-grid of at most 100 points; off-grid points and the CRN probe " + "are extra single-point Requests and do not count." + ), + ) + + +# Kubernetes request timeout of the live retune supervisor +# (``SupervisorSettings``). +LIVE_RETUNE_K8S_TIMEOUT_SECONDS = 5.0 + + +def live_retune_lease_margin_seconds(k8s_timeout_s: float) -> float: + """How long before the leadership Lease could expire local leadership + ends.""" + + return 3.0 * k8s_timeout_s + + +def live_retune_io_timeout_seconds(k8s_timeout_s: float) -> float: + """I/O budget of one blocking Kubernetes call, such as a whole renewal + (a GET and a replace): a backstop, the client applies the request + timeout too.""" + + return 2.0 * k8s_timeout_s + 1.0 + + +def live_retune_min_lease_seconds(k8s_timeout_s: float) -> float: + """Shortest leadership Lease the live supervisor can keep renewing. + + Local leadership ends the lease margin before the Lease could expire, + counted from the send of the renewal that granted it. That renewal may + use a whole I/O budget, and the next one starts at least one I/O budget + before local leadership ends, so two budgets must fit. Lease rounds are + single renewals: an initialization renews concurrently instead of + delaying the next one. + """ + + margin_s = live_retune_lease_margin_seconds(k8s_timeout_s) + return margin_s + 2.0 * live_retune_io_timeout_seconds(k8s_timeout_s) + + +class DgdrExecutorConfig(BaseModel): + """Live retune through DGDR v1beta2 (``retune.executor: dgdr``). + + The planner cuts each window from the frontend request traces on a shared + RWX claim, creates DGDR v1beta2 Requests that replay it, scores their + Candidates and applies the winner to the frontends' routers. + """ + + model_config = ConfigDict(extra="forbid") + + phase_seconds: float = Field( + default=0.0, + ge=0, + allow_inf_nan=False, + description=( + "Window phase: windows are [phase + k*interval, phase + (k+1)*interval) " + "in epoch seconds, so they survive planner restarts." + ), + ) + claim_name: str = Field( + min_length=1, + description="RWX PersistentVolumeClaim, in the DGD namespace, holding traces.", + ) + mount_path: str = Field( + default="/trace", + min_length=1, + description="Where the planner container mounts claim_name.", + ) + segments_dir: str = Field( + default="segments", + min_length=1, + description="Claim-relative directory of frontend trace segments: ///.", + ) + windows_dir: str = Field( + default="windows", + min_length=1, + description=( + "Claim-relative directory of window files: " + "//-.jsonl.gz, where " + " is a digest of the window's start, end and cutoff." + ), + ) + model: str = Field( + min_length=1, + description="HF repo id the workers serve (worker --model); AISimulate resolves it.", + ) + sku: str = Field(min_length=1, description="spec.hardware.gpu.sku of each Request.") + engine: dict[str, Any] = Field( + description=( + "AISimulate public engine section with every leaf a scalar " + "(mode, backend, backend_version, context_length, " + "workers..parallelism.preset, scheduler, kv_cache). Replicas " + "come from the DGD spec at each launch and must not be set here." + ), + ) + cut_budget_seconds: float = Field( + default=900.0, + gt=0, + allow_inf_nan=False, + description="Deadline of the subprocess that cuts a window file.", + ) + job_startup_seconds: float = Field( + default=300.0, + ge=0, + allow_inf_nan=False, + description="Run Job start-up time reserved when deciding whether a shard can still finish.", + ) + run: DgdrRunConfig = Field(default_factory=DgdrRunConfig) + keep_requests: int = Field( + default=5, + ge=0, + description="Windows whose terminal Requests are kept for audit.", + ) + max_window_rows: int = Field( + default=5_000_000, + ge=1, + description="A window with more requests fails (window_too_large).", + ) + promote_last_advisory: bool = Field( + default=False, + description=( + "When switching from advisory to active, apply the last advisory " + "recommendation if it is recent and its fingerprint still matches." + ), + ) + lease_duration_seconds: float = Field( + default=60.0, + allow_inf_nan=False, + # Not ge=: it would reject before _lease_fits_two_renewals explains. + json_schema_extra={ + "minimum": live_retune_min_lease_seconds(LIVE_RETUNE_K8S_TIMEOUT_SECONDS) + }, + description=( + "Leadership Lease duration. Renewed every third of it, or sooner " + "so that a renewal that uses its whole I/O budget ends before " + "local leadership does (the duration minus the lease margin). " + "At least 37 s: the 15 s lease margin plus two 11 s renewal I/O " + "budgets, with the 5 s Kubernetes request timeout." + ), + ) + actuator_interval_seconds: float = Field( + default=15.0, + gt=0, + allow_inf_nan=False, + description="Period of the router actuator and reconciler rounds.", + ) + + @field_validator("lease_duration_seconds") + @classmethod + def _lease_fits_two_renewals(cls, value: float) -> float: + timeout_s = LIVE_RETUNE_K8S_TIMEOUT_SECONDS + minimum_s = live_retune_min_lease_seconds(timeout_s) + if value < minimum_s: + raise ValueError( + "retune.dgdr.lease_duration_seconds must be at least the lease " + "margin + 2 x the I/O budget of one renewal " + f"({live_retune_lease_margin_seconds(timeout_s):g}s + 2 x " + f"{live_retune_io_timeout_seconds(timeout_s):g}s = {minimum_s:g}s " + f"with the {timeout_s:g}s Kubernetes request timeout)" + ) + return value + + +# Live retune (``retune.executor: dgdr``) limits. A sweep must end this long +# before the next window's launch: harvest, the decision write and clean-up +# run in that margin. +DGDR_TIMEOUT_MARGIN_SECONDS = 600.0 +# DGDR v1beta2 ``spec.recommendation.maxCandidates`` maximum: one Request (a +# grid shard) carries at most this many router points. +DGDR_MAX_SHARD_POINTS = 100 +# Router knobs a DGDR v1beta2 Request can search, by public name +# (``contracts.DGDR_ROUTER_KNOBS``); the Request's router section has no +# place for the retune extension knobs. +_DGDR_ROUTER_SECTION_KEYS = frozenset( + { + "overlap_score_credit", + "prefill_load_scale", + "temperature", + "policy", + "prefill_load_model", + } +) +_DGDR_ENGINE_MODES = {"agg": "aggregated", "disagg": "disaggregated"} +_DGDR_ENGINE_ROLES = {"agg": ("aggregated",), "disagg": ("prefill", "decode")} +# ``parallelism.preset`` entry keys: AISimulate requires all of these plus +# ``replicas``, which the executor sets from the DGD spec at each launch. +_DGDR_PRESET_KEYS = ("tensor", "pipeline", "attention_data", "moe_tensor", "moe_expert") +# The pinned AISimulate recommend runs prefill_context only at 1, and its +# preset entry schema has no decode_context. When the pin moves, update this +# tuple together with _validate_dgdr_parallelism (its decode_context branch, +# the "optionally prefill_context: 1" message and the prefill_context == 1 +# check), test_pinned_aisimulate_rejects_decode_context and the preset rules +# in docs/fern/pages/kubernetes/auto-deployment/planner-live-retune.md. +_DGDR_PRESET_OPTIONAL_KEYS = ("prefill_context",) +# Scheduler knobs AISimulate searches over a default domain when omitted. +_DGDR_SCHEDULER_KNOBS = ("max_batched_tokens", "max_sequences") + + +def dgdr_rounds(points: int, parallelism: int) -> int: + """Evaluation rounds of a DGDR Run that scores ``points`` candidates. + + The AISimulate optimizer evaluates the first candidate alone, then + ``parallelism`` candidates at a time. + """ + + if points < 1: + raise ValueError(f"a DGDR shard has at least one point, got {points}") + if parallelism < 1: + raise ValueError(f"parallelism must be at least 1, got {parallelism}") + return 1 + math.ceil((points - 1) / parallelism) + + +def dgdr_need_seconds( + points: int, dgdr: DgdrExecutorConfig, *, needs_cut: bool +) -> float: + """Wall-clock seconds one DGDR shard of ``points`` router points needs. + + ``cut_budget`` (when the window file must still be cut) + Run Job + start-up + ``dgdr_rounds(points, run.parallelism)`` rounds of + ``candidate_timeout``. ``parallelism`` is the configured value; Job + resources may lower it, so this is optimistic. + """ + + rounds = dgdr_rounds(points, dgdr.run.parallelism) + cut = dgdr.cut_budget_seconds if needs_cut else 0.0 + return cut + dgdr.job_startup_seconds + rounds * dgdr.run.candidate_timeout_seconds + + +def balanced_shard_parts( + sizes: Sequence[int], max_shards: int, *, parallelism: int +) -> tuple[int, ...]: + """How many chunks to cut each router grid dimension into for DGDR sharding. + + ``sizes`` are the value counts of the grid dimensions. Cutting dimension + ``i`` into ``parts[i]`` contiguous chunks whose sizes differ by at most + one gives ``prod(parts)`` Cartesian shards, the largest with + ``largest_shard_points(sizes, parts)`` points. Of the parts with + ``prod(parts) <= max_shards``, returns the one that minimizes, in order: + + 1. whether the largest shard exceeds ``DGDR_MAX_SHARD_POINTS``, and when + every split does, its points (so the caller reports the smallest); + 2. the largest shard's wall time, its ``dgdr_rounds`` with + ``parallelism``; + 3. the shard count, the largest shard's points, then the parts. + + P13 calls it with ``parallelism=run.parallelism``; the live executor + must plan each window's shards with the same call, so that both see the + same largest shard. The default 3 x 9 x 4 grid with ``max_shards=3`` and + ``parallelism=8`` gives ``(1, 3, 1)``, three shards of 36 points. + """ + + if max_shards < 1: + raise ValueError(f"max_shards must be at least 1, got {max_shards}") + if any(size < 1 for size in sizes): + raise ValueError(f"every grid dimension needs a value, got sizes {sizes}") + if parallelism < 1: + raise ValueError(f"parallelism must be at least 1, got {parallelism}") + + def splits(index: int, budget: int): + if index == len(sizes): + yield () + return + for part in range(1, min(sizes[index], budget) + 1): + for rest in splits(index + 1, budget // part): + yield (part, *rest) + + def rank(parts: tuple[int, ...]) -> tuple[Any, ...]: + largest = largest_shard_points(sizes, parts) + if largest > DGDR_MAX_SHARD_POINTS: + return (True, largest, math.prod(parts), parts) + rounds = dgdr_rounds(largest, parallelism) + return (False, rounds, math.prod(parts), largest, parts) + + return min(splits(0, max_shards), key=rank) + + +def largest_shard_points(sizes: Sequence[int], parts: Sequence[int]) -> int: + """Points in the largest shard of ``balanced_shard_parts``' split.""" + + return math.prod( + math.ceil(size / part) for size, part in zip(sizes, parts, strict=True) + ) + + +def _validate_dgdr_engine( + engine: Mapping[str, Any], planner_mode: str, backend: str +) -> None: + """P5: ``retune.dgdr.engine`` pins one engine with scalars only. + + The section is the Request's ``inputOverrides.engine`` merge patch, so + any AISimulate search domain left in it (a ``choices`` or ``range`` + mapping, a parallelism preset list or a default domain) would make the + Run search engine settings next to the router knobs. A written + ``backend`` must be the deployed one (``PlannerConfig.backend``), which + every Request names and which fills an omitted ``backend``. + """ + + name = "retune.dgdr.engine" + _check_dgdr_engine_scalars(engine, (name,)) + written = engine.get("backend") + if written is not None and written != backend: + raise ValueError( + f"{name}.backend {written!r} differs from the deployed backend " + f"{backend!r} (P5); remove it or set it to {backend!r}" + ) + expected_mode = _DGDR_ENGINE_MODES[planner_mode] + if engine.get("mode") != expected_mode: + raise ValueError( + f"{name}.mode must be {expected_mode!r} for planner mode " + f"{planner_mode!r} (P5), got {engine.get('mode')!r}; an omitted mode " + "searches both aggregated and disaggregated engines" + ) + for key, owner in (("model", "model"), ("hardware", "sku")): + if key in engine: + raise ValueError( + f"{name}.{key} would override the Request's typed field (P5); " + f"set retune.dgdr.{owner} instead" + ) + workers = engine.get("workers") + if not isinstance(workers, Mapping): + raise ValueError( + f"{name}.workers must be a mapping that pins every worker role of " + f"planner mode {planner_mode!r} (P5)" + ) + used = _DGDR_ENGINE_ROLES[planner_mode] + for role, worker in workers.items(): + if role not in used and worker is not None: + raise ValueError( + f"{name}.workers.{role} is not a worker role of planner mode " + f"{planner_mode!r}; write it as null or remove it (P5)" + ) + for role in used: + path = f"{name}.workers.{role}" + worker = workers.get(role) + if not isinstance(worker, Mapping): + raise ValueError( + f"{path} is required for planner mode {planner_mode!r} (P5)" + ) + _validate_dgdr_parallelism(worker.get("parallelism"), f"{path}.parallelism") + scheduler = worker.get("scheduler") + scheduler = scheduler if isinstance(scheduler, Mapping) else {} + for knob in _DGDR_SCHEDULER_KNOBS: + if not _is_positive_int(scheduler.get(knob)): + raise ValueError( + f"{path}.scheduler.{knob} must be a positive integer (P5); " + "omitted, AISimulate searches its default domain" + ) + + +def _validate_dgdr_parallelism(parallelism: Any, path: str) -> None: + preset = parallelism.get("preset") if isinstance(parallelism, Mapping) else None + if not ( + isinstance(preset, list) and len(preset) == 1 and isinstance(preset[0], Mapping) + ): + raise ValueError( + f"{path}.preset must be a list of exactly one mapping with " + f"{', '.join(_DGDR_PRESET_KEYS)} (P5); got {preset!r}" + ) + others = sorted(set(parallelism) - {"preset"}) + if others: + raise ValueError( + f"{path} sets {others} next to preset; pin the parallelism in the " + "single preset entry (P5)" + ) + entry = preset[0] + if "decode_context" in entry: + raise ValueError( + f"{path}.preset[0].decode_context is not accepted (P5); the pinned " + "AISimulate preset schema has no decode_context" + ) + missing = [key for key in _DGDR_PRESET_KEYS if key not in entry] + unknown = sorted( + set(entry) - set(_DGDR_PRESET_KEYS) - set(_DGDR_PRESET_OPTIONAL_KEYS) + ) + if missing or unknown: + raise ValueError( + f"{path}.preset[0] must set exactly {', '.join(_DGDR_PRESET_KEYS)} " + "and optionally prefill_context: 1 " + f"(P5); missing={missing}, unknown={unknown}" + ) + for key, value in entry.items(): + if not _is_positive_int(value): + raise ValueError( + f"{path}.preset[0].{key} must be a positive integer (P5), got {value!r}" + ) + if entry.get("prefill_context", 1) != 1: + raise ValueError( + f"{path}.preset[0].prefill_context must be 1 or omitted (P5); " + "AISimulate recommend rejects context parallelism above 1, " + f"got {entry['prefill_context']!r}" + ) + + +def _check_dgdr_engine_scalars(value: Any, path: tuple[str, ...]) -> None: + dotted = ".".join(path) + if isinstance(value, Mapping): + domain = sorted({"choices", "range"} & set(value)) + if domain: + raise ValueError( + f"{dotted} is a search domain ({domain[0]}); every " + "retune.dgdr.engine value must be a scalar (P5)" + ) + for key, item in value.items(): + if key == "replicas": + raise ValueError( + f"{dotted}.replicas must not be set: the DGD spec provides " + "replicas at each launch (P5)" + ) + _check_dgdr_engine_scalars(item, (*path, str(key))) + elif isinstance(value, list): + # Only workers..parallelism.preset is a list; its shape is + # checked per used role. + if not ( + len(path) == 5 + and path[1] == "workers" + and path[3:] == ("parallelism", "preset") + ): + raise ValueError( + f"{dotted} is a list; every retune.dgdr.engine value except " + "workers..parallelism.preset must be a scalar (P5)" + ) + for index, item in enumerate(value): + _check_dgdr_engine_scalars(item, (*path[:-1], f"{path[-1]}[{index}]")) + elif value is not None and not isinstance(value, (str, int, float)): + raise ValueError(f"{dotted} must be a scalar, got {type(value).__name__} (P5)") + + +def _is_positive_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + class RetuneConfig(BaseModel): """Periodic replay-based retune of router and planner runtime parameters. @@ -981,6 +1438,17 @@ class RetuneConfig(BaseModel): "AISimulate default domain; null does not search the router." ), ) + executor: Literal["local", "dgdr"] = Field( + default="local", + description=( + "Where sweeps run: local worker processes (replay) or DGDR v1beta2 " + "Requests (live; requires dgdr)." + ), + ) + dgdr: DgdrExecutorConfig | None = Field( + default=None, + description="Live DGDR executor settings; required iff executor is dgdr.", + ) planner: Literal["default"] | dict[str, Any] | None = Field( default=None, description=( @@ -1019,6 +1487,12 @@ def _validate_domain_syntax( parse_retune_section(section, value) # type: ignore[call-overload] return value + @model_validator(mode="after") + def _validate_executor(self) -> "RetuneConfig": + if (self.executor == "dgdr") != (self.dgdr is not None): + raise ValueError("retune.dgdr is required iff retune.executor is 'dgdr'") + return self + @model_validator(mode="after") def _validate_timing(self) -> "RetuneConfig": timeout = self.effective_timeout_seconds @@ -1456,8 +1930,8 @@ def _validate_ca_bundle_path(cls, v: Optional[str]) -> Optional[str]: ) # Advisory mode: compute, log and report scaling decisions without - # executing them, live and in planner-in-the-loop replay; replay records - # retune decisions (replay-only) the same way. + # executing them, live and in planner-in-the-loop replay; retune records + # a decision that would apply as advisory and applies nothing. advisory: bool = SLAPlannerDefaults.advisory # --- Power-aware budget (read-only caps; DGD-owned) --- @@ -1565,9 +2039,11 @@ def _validate_ca_bundle_path(cls, v: Optional[str]) -> Optional[str]: default=None, description=( "Opt-in periodic replay-based retune of router and planner runtime " - "parameters. Replay-only in this release. Setting both " - "enable_load_scaling and enable_throughput_scaling explicitly to " - "false with retune set runs retune-only with fixed replicas." + "parameters. executor 'local' runs in planner-in-the-loop replay; " + "a live Kubernetes planner requires executor 'dgdr' " + "(experimental). Setting both enable_load_scaling and " + "enable_throughput_scaling explicitly to false with retune set " + "runs retune-only with fixed replicas." ), ) @@ -1952,6 +2428,9 @@ def _validate_retune(self, retune: RetuneConfig) -> None: raise ValueError( f"retune requires mode='agg' or mode='disagg', got {self.mode!r}" ) + if retune.executor == "dgdr": + # Before the domain expansion, whose errors assume replay. + self._validate_live_retune(retune) if self.scaling_enabled() and retune.objective not in RETUNE_PER_GPU_OBJECTIVES: raise ValueError( f"retune.objective {retune.objective!r} does not divide by " @@ -1996,6 +2475,174 @@ def _validate_retune(self, retune: RetuneConfig) -> None: "omitted AISimulate knob searches its default domain, so pin " "knobs with a scalar to shrink the grid" ) + if retune.executor == "dgdr": + self._validate_live_retune_budget(retune, router_domain) + + def _validate_live_retune(self, retune: RetuneConfig) -> None: + """Configuration premises P1, P2, P3 and P5 of live retune. + + A live planner sweeps through DGDR v1beta2 Requests in the Kubernetes + namespace of its DGD, retunes the frontend router only (Scenario 1), + and decides when a sweep returns; the frontend-side premises are + checked at runtime against the DGD. + """ + + dgdr = retune.dgdr + if dgdr is None: # RetuneConfig requires it with executor dgdr + raise ValueError("retune.dgdr is required iff retune.executor is 'dgdr'") + if self.environment != "kubernetes": + raise ValueError( + "retune.executor 'dgdr' requires environment='kubernetes' (P1), " + f"got {self.environment!r}" + ) + latency = retune.decision_latency_seconds + if latency == "measured": + if retune.decision_latency_scale != 1.0: + raise ValueError( + "retune.executor 'dgdr' decides when the sweep returns, so " + "decision_latency_seconds: measured requires " + "decision_latency_scale 1 (P1), got " + f"{retune.decision_latency_scale}" + ) + elif latency != 0: + raise ValueError( + "retune.executor 'dgdr' decides when the sweep returns, so " + "retune.decision_latency_seconds must be 0 or 'measured' (P1), " + f"got {latency}" + ) + if retune.timeout_seconds is None: + raise ValueError( + "retune.executor 'dgdr' requires an explicit " + "retune.timeout_seconds (P1): the sweep deadline, window cut " + "included" + ) + limit = retune.interval_seconds - DGDR_TIMEOUT_MARGIN_SECONDS + if retune.timeout_seconds > limit: + raise ValueError( + f"retune.timeout_seconds must be <= interval_seconds - " + f"{DGDR_TIMEOUT_MARGIN_SECONDS:g} = {limit:g} with executor " + f"'dgdr' (P1), got {retune.timeout_seconds:g}" + ) + if self.scaling_enabled(): + raise ValueError( + "retune.executor 'dgdr' supports retune-only mode (P2): set " + "enable_load_scaling and enable_throughput_scaling both to false" + ) + if retune.planner is not None: + raise ValueError( + "retune.executor 'dgdr' retunes the router only (P2); remove " + "retune.planner" + ) + self._validate_live_retune_router(retune) + _validate_dgdr_engine(dgdr.engine, self.mode, self.backend) + # Claim paths: the planner, the frontends and the Run Jobs mount the + # same claim, and clean-up deletes under these directories. A DGDR + # trace path rejects empty, '.' and '..' segments. + dirs: dict[str, list[str]] = {} + for name in ("segments_dir", "windows_dir"): + value = getattr(dgdr, name) + dirs[name] = value.split("/") + if any(segment in ("", ".", "..") for segment in dirs[name]): + raise ValueError( + f"retune.dgdr.{name} must be a relative directory inside " + "the claim: '/'-separated names without empty, '.' or " + f"'..' segments, got {value!r}" + ) + if not PurePosixPath(dgdr.mount_path).is_absolute(): + raise ValueError( + f"retune.dgdr.mount_path must be absolute, got {dgdr.mount_path!r}" + ) + segments, windows = dirs["segments_dir"], dirs["windows_dir"] + shorter = min(len(segments), len(windows)) + if segments[:shorter] == windows[:shorter]: + raise ValueError( + "retune.dgdr.segments_dir and retune.dgdr.windows_dir must not " + f"contain one another, got {dgdr.segments_dir!r} and " + f"{dgdr.windows_dir!r}" + ) + + @staticmethod + def _validate_live_retune_router(retune: RetuneConfig) -> None: + """P3: the router domain fits the DGDR v1beta2 router section. + + A Request searches ``overlap_score_credit``, ``prefill_load_scale`` + and ``temperature`` of a ``kv_router`` with ``prefill_load_model: + none``; ``RouterRetuneConfig`` already pins the policy. + """ + + spec = retune.router + if not isinstance(spec, Mapping): + return + extra = sorted(set(spec) - _DGDR_ROUTER_SECTION_KEYS) + if extra: + raise ValueError( + f"retune.router {extra} cannot be searched with executor 'dgdr' " + "(P3): a DGDR Request searches only overlap_score_credit, " + "prefill_load_scale and temperature" + ) + if "prefill_load_model" not in spec: + return + # Imported on use: the adapter schema imports this module. + from aisimulate.config.common import Choices + + from dynamo.planner.retune.domains import parse_retune_section + + model = parse_retune_section("router", spec).prefill_load_model + if "type" not in model.model_fields_set: + return + written = ( + list(model.type.choices) + if isinstance(model.type, Choices) + else [model.type] + ) + if written != ["none"]: + raise ValueError( + "retune.router.prefill_load_model.type must be 'none' with " + f"executor 'dgdr' (P3), got {written}" + ) + + @staticmethod + def _validate_live_retune_budget( + retune: RetuneConfig, router_domain: Mapping[str, Sequence[float]] + ) -> None: + """P13: the largest grid shard can finish within the sweep deadline. + + The baseline can fall in any grid shard, and a window whose baseline + shard cannot finish fails with ``launch_too_late``, so the largest + shard must fit. The domain is the widest one: a zero startup credit + only narrows it at runtime. + """ + + dgdr = retune.dgdr + if dgdr is None or retune.timeout_seconds is None: + # _validate_live_retune rejects both first. + raise ValueError("retune.executor 'dgdr' requires dgdr and timeout_seconds") + sizes = [len(values) for values in router_domain.values()] + parts = balanced_shard_parts( + sizes, dgdr.run.max_shards, parallelism=dgdr.run.parallelism + ) + points = largest_shard_points(sizes, parts) + if points > DGDR_MAX_SHARD_POINTS: + raise ValueError( + f"the retune router grid of {math.prod(sizes)} points cannot be " + f"split into at most retune.dgdr.run.max_shards=" + f"{dgdr.run.max_shards} Cartesian shards of at most " + f"{DGDR_MAX_SHARD_POINTS} points (P13); the largest would have " + f"{points}. Raise max_shards or pin router knobs with a scalar" + ) + need = dgdr_need_seconds(points, dgdr, needs_cut=True) + if need > retune.timeout_seconds: + rounds = dgdr_rounds(points, dgdr.run.parallelism) + raise ValueError( + f"the largest retune grid shard ({points} points) needs " + f"{need:g}s: cut_budget_seconds {dgdr.cut_budget_seconds:g} + " + f"job_startup_seconds {dgdr.job_startup_seconds:g} + {rounds} " + f"rounds x candidate_timeout_seconds " + f"{dgdr.run.candidate_timeout_seconds:g}, more than " + f"retune.timeout_seconds={retune.timeout_seconds:g} (P13), so " + "every window would fail launch_too_late. Raise timeout_seconds " + "or run.parallelism, add run.max_shards, or shrink the grid" + ) def _validate_retune_planner_grid( self, diff --git a/components/src/dynamo/planner/core/base.py b/components/src/dynamo/planner/core/base.py index 593703c7909d..258b7dbbc861 100644 --- a/components/src/dynamo/planner/core/base.py +++ b/components/src/dynamo/planner/core/base.py @@ -12,17 +12,19 @@ from __future__ import annotations import asyncio +import importlib import logging import time from collections.abc import AsyncIterator, Iterable from contextlib import asynccontextmanager -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional, Protocol import aiohttp.web from prometheus_client import start_http_server from dynamo.planner.config.defaults import SubComponentType, TargetReplica from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.connectors.kubernetes import KubernetesConnector from dynamo.planner.control_api import ( _MinimumEndpointUnavailableError, _MinimumEndpointValidationError, @@ -53,18 +55,93 @@ from dynamo.planner.offline.trace_data import extract_traffic_observations_from_trace from dynamo.planner.reconfiguration import ( PlannerReconfiguration, + live_retune_sla_errors, runtime_candidate_errors, runtime_configuration_errors, ) +from dynamo.planner.retune.live.contracts import KubeError, KubeNotFound +from dynamo.planner.retune.live.kube import KubernetesRetuneKube +from dynamo.planner.retune.live.state import find_retune_leases from dynamo.runtime import DistributedRuntime if TYPE_CHECKING: + from dynamo.planner.core.types import RetuneInputs from dynamo.planner.monitoring.worker_info import WorkerInfo from dynamo.planner.plugins.builtins.local_planner import BuiltinLoadPredict + from dynamo.planner.plugins.orchestrator.engine_adapter import RetuneHooks logger = logging.getLogger(__name__) + +def _live_retune_bootstrap() -> Any: + """The live retune stack, loaded only by a planner that runs it. + + It pulls in the DGDR sweep executor and the replay stack behind it; + a planner with retune config loads it at startup, one without only + when a removed live retune left its Leases behind. + """ + + return importlib.import_module("dynamo.planner.retune.live.bootstrap") + + +def _log_legacy_retune_failure(task: asyncio.Task[None]) -> None: + if not task.cancelled() and task.exception() is not None: + logger.error( + "Live retune legacy clean-up check failed", exc_info=task.exception() + ) + + +class LiveRetuneHandle(Protocol): + """What ``retune.live.bootstrap.build_live_retune`` returns to the base. + + ``retune_inputs`` is ``None`` while the frontend's startup router config + does not satisfy the retune premises (the engine then starts with retune + disabled). ``start`` runs after the engine is bootstrapped and may call + back ``replace_live_retune``; ``on_effects`` sees every tick's effects + (it must be cheap and never block); ``stop`` runs before the engine shuts + down: it aborts the sweep, waits for the sweep thread and releases + leadership. + """ + + def retune_inputs(self) -> Optional[RetuneInputs]: + ... + + def hooks(self) -> RetuneHooks: + ... + + async def start(self, host: NativePlannerBase) -> None: + ... + + async def stop(self) -> None: + ... + + def on_effects(self, effects: PlannerEffects) -> None: + ... + + +class RetuneCleanupHandle(Protocol): + """What ``retune.live.bootstrap.build_legacy_cleanup`` returns. + + A planner without retune config whose DGD still has the Leases of a + removed live retune rolls back the routers that retune deployed, then + deletes its Requests and Leases; ``start`` begins that clean-up and + ``stop`` releases leadership. + """ + + async def start(self, host: NativePlannerBase) -> None: + ... + + async def stop(self) -> None: + ... + + _CONFIG_LOCK_TIMEOUT_SECONDS = 10.0 +# The background check for a removed live retune reads two Leases, then, +# only if one exists, the DGD; each Kubernetes call is bounded by this. +_LEGACY_RETUNE_CALL_TIMEOUT_SECONDS = 10.0 +# The check runs once per start; a transient API error is retried once after +# this long (403 and 404 are not). +_LEGACY_RETUNE_RETRY_SECONDS = 15.0 # A tick waits briefly for the common fast submission path, then leaves the # task running. It is deliberately not cancelled: a remote request may have # been accepted even when its response has not arrived, so retrying could @@ -131,8 +208,13 @@ def __init__( config: PlannerConfig, environment: PlannerEnvironment, ) -> None: - if config.retune is not None: - raise ValueError("planner.retune is replay-only in this release") + if config.retune is not None and config.retune.executor == "local": + # Local sweeps replay in this process; a live planner retunes + # through DGDR v1beta2 Requests only. + raise ValueError( + "planner.retune is replay-only with executor 'local'; " + "use executor: dgdr" + ) self.runtime = runtime self.config = config self.environment = environment @@ -174,6 +256,13 @@ def __init__( self._environment_initialized = False self._engine: Optional[EngineProtocol] = None self._last_worker_counts: Optional[WorkerCounts] = None + # Live retune (``retune.executor: dgdr``): built in _async_init. + self._live_retune: Optional[LiveRetuneHandle] = None + # Without retune config, a Kubernetes planner checks in the background + # whether a removed live retune left its Leases, and cleans up after it. + self._legacy_retune_task: Optional[asyncio.Task[None]] = None + self._legacy_retune_cleanup: Optional[RetuneCleanupHandle] = None + self._retune_lease_found = asyncio.Event() async def _async_init(self) -> None: # Shutdown is safe for a partially initialized environment and is @@ -183,9 +272,19 @@ async def _async_init(self) -> None: await self.environment.initialize() self._validate_min_endpoint_budgets_at_startup() + if self.config.retune is not None: + # Before the first _ensure_engine (regression bootstrap builds + # the engine): the engine takes the live retune inputs/hooks. + self._live_retune = await self._build_live_retune() + self._retune_lease_found.set() + await self._bootstrap_regression() await self._bootstrap_engine_plugins_if_needed() + if self._live_retune is not None: + await self._live_retune.start(self) + self._start_legacy_retune_check() + if self.config.enable_runtime_reconfiguration: if self.runtime is None: raise ValueError("system-port reconfiguration requires a runtime") @@ -221,6 +320,124 @@ async def _async_init(self) -> None: await self._shutdown_runtime() raise + @property + def may_hold_retune_lease(self) -> bool: + """Whether live retune, or the clean-up of a removed one, may hold + the retune leadership Lease (known once ``_async_init`` returned).""" + + return self._live_retune is not None or self._legacy_retune_task is not None + + @property + def retune_lease_found(self) -> asyncio.Event: + """Set once this planner may hold a retune Lease: at start with live + retune configured, or when the clean-up check of a removed one finds + a Lease. ``__main__`` installs its SIGTERM handler then.""" + + return self._retune_lease_found + + async def _build_live_retune(self) -> LiveRetuneHandle: + return await _live_retune_bootstrap().build_live_retune(self) + + def _kubernetes_connector(self) -> Optional[KubernetesConnector]: + # Only some environments have a deployment controller. + controller = getattr(self.environment, "controller", None) + return controller if isinstance(controller, KubernetesConnector) else None + + def _start_legacy_retune_check(self) -> None: + """Look for a removed live retune in the background. + + Only a Kubernetes planner without retune config. Startup never waits + for it, and it never fails the planner. + """ + + connector = self._kubernetes_connector() + if self.config.retune is not None or connector is None: + return + task = asyncio.create_task( + self._legacy_retune_check(connector), name="planner-retune-legacy-check" + ) + task.add_done_callback(_log_legacy_retune_failure) + self._legacy_retune_task = task + + async def _legacy_retune_check(self, connector: KubernetesConnector) -> None: + """Clean up after a live retune whose config was removed. + + Reads the DGD's two retune Leases; only when one exists does it load + the live retune stack, which rolls back the routers that retune + deployed and then deletes its Requests and Leases. Without Lease RBAC + (the default chart) the read answers 403, logged at debug. Any other + API error but a 404 is retried once. + """ + + kube = KubernetesRetuneKube.from_kube_api(connector.kube_api) + cleanup: Optional[RetuneCleanupHandle] = None + for retry in (True, False): + try: + cleanup = await self._find_legacy_retune( + kube, connector.parent_dgd_name + ) + break + except KubeError as error: + # ``status`` is set by the Kubernetes client wrapper only. + if getattr(error, "status", None) == 403: + logger.debug("Live retune legacy clean-up check skipped: %s", error) + return + if not retry or isinstance(error, KubeNotFound): + logger.info("Live retune legacy clean-up skipped: %s", error) + return + logger.info( + "Live retune legacy clean-up check failed; retrying in %.0f s: %s", + _LEGACY_RETUNE_RETRY_SECONDS, + error, + ) + await asyncio.sleep(_LEGACY_RETUNE_RETRY_SECONDS) + if cleanup is None: + return + self._legacy_retune_cleanup = cleanup + await cleanup.start(self) + + async def _find_legacy_retune( + self, kube: KubernetesRetuneKube, dgd_name: str + ) -> Optional[RetuneCleanupHandle]: + """The clean-up of a removed live retune; ``None`` without its Leases.""" + + leases = await asyncio.to_thread( + find_retune_leases, + kube, + dgd_name, + timeout_s=_LEGACY_RETUNE_CALL_TIMEOUT_SECONDS, + ) + if not leases: + return None + # The clean-up takes the leadership Lease: stop cleanly on SIGTERM. + self._retune_lease_found.set() + return await _live_retune_bootstrap().build_legacy_cleanup( + self, kube, timeout_s=_LEGACY_RETUNE_CALL_TIMEOUT_SECONDS + ) + + async def replace_live_retune( + self, + inputs: Optional[RetuneInputs], + *, + generation: int, + ) -> None: + """Swap the engine's live retune controller (supervisor callback). + + Serialized with ticks by the config lock, so no tick runs a + half-replaced controller. ``inputs=None`` disables retune; a + construction error propagates to the caller with retune disabled. + """ + + from dynamo.planner.plugins.orchestrator.engine_adapter import ( + OrchestratorEngineAdapter, + ) + + async with self._config_lock: + engine = self._engine + if not isinstance(engine, OrchestratorEngineAdapter): + raise RuntimeError("live retune requires the orchestrator engine") + engine.replace_retune(inputs, generation=generation, now_s=time.time()) + async def _shutdown_runtime(self) -> None: """Release initialized planner resources after normal or failed startup.""" @@ -249,6 +466,29 @@ async def _shutdown_runtime(self) -> None: except Exception: logger.exception("Failed to stop planner live dashboard") + # Before the engine: the supervisor aborts the retune sweep, waits + # for the engine's sweep thread and then releases leadership. + live_retune = self._live_retune + self._live_retune = None + if live_retune is not None: + try: + await live_retune.stop() + except Exception: + logger.exception("Failed to stop planner live retune") + + legacy_retune_task = self._legacy_retune_task + self._legacy_retune_task = None + if legacy_retune_task is not None: + legacy_retune_task.cancel() + await asyncio.wait([legacy_retune_task]) + legacy_retune_cleanup = self._legacy_retune_cleanup + self._legacy_retune_cleanup = None + if legacy_retune_cleanup is not None: + try: + await legacy_retune_cleanup.stop() + except Exception: + logger.exception("Failed to stop the live retune legacy clean-up") + engine = self._engine self._engine = None if engine is not None: @@ -344,6 +584,8 @@ async def get_min_endpoints(self) -> dict[str, object]: def validate_runtime_candidate(self, candidate: PlannerConfig) -> None: errors = runtime_candidate_errors(candidate, self._build_worker_capabilities()) + if self._live_retune is not None: + errors.extend(live_retune_sla_errors(self.config, candidate)) if errors: raise ValueError("; ".join(errors)) self._reconfigurable_engine() @@ -467,6 +709,7 @@ def _ensure_engine(self) -> EngineProtocol: OrchestratorEngineAdapter, ) + live_retune = self._live_retune self._engine = OrchestratorEngineAdapter( self.config, self._build_worker_capabilities(), @@ -475,6 +718,10 @@ def _ensure_engine(self) -> EngineProtocol: require_prefill=self.require_prefill, require_decode=self.require_decode, ), + retune_inputs=( + None if live_retune is None else live_retune.retune_inputs() + ), + retune_hooks=None if live_retune is None else live_retune.hooks(), ) return self._engine @@ -1034,6 +1281,14 @@ async def _run_one_tick( async with self._config_lock: config_generation = self._config_generation effects = await engine.tick(tick, tick_input) + if self._live_retune is not None: + # Retune decisions go to the live supervisor's outbox, which + # persists them before actuating routers; _apply_effects never + # sees router parameters. + try: + self._live_retune.on_effects(effects) + except Exception: + logger.exception("Live retune failed to take the tick's effects") await self._submit_effects(effects, config_generation) emit_diagnostics = self._should_emit_tick_diagnostics(tick, effects) if emit_diagnostics: diff --git a/components/src/dynamo/planner/core/types.py b/components/src/dynamo/planner/core/types.py index 428d6712ef65..117b14a1c033 100644 --- a/components/src/dynamo/planner/core/types.py +++ b/components/src/dynamo/planner/core/types.py @@ -11,6 +11,7 @@ from __future__ import annotations +import copy import json import math from dataclasses import dataclass, field @@ -342,7 +343,9 @@ def __hash__(self) -> int: class RetuneDecision: """Audit record of one retune decision. - ``decision_id`` is ``retune--w`` and is the live reconfiguration + ``decision_id`` is ``retune--w`` in replay and + ``retune-w`` live (windows anchored at an epoch phase, so + the id survives planner restarts); it is the live reconfiguration ``request_id``. Times are planner-clock seconds: epoch seconds live, virtual seconds in replay. ``as_of_s`` is the censoring cutoff (the scheduled launch time). ``chosen`` equals ``previous`` unless the outcome @@ -355,7 +358,10 @@ class RetuneDecision: ``(params.key(), score)`` pairs, best first. ``sweep_wall_s`` is the sweep's wall-clock time and ``decision_latency_s`` the delay the decision waited for after launch (at least one tick; ``decided_at_s`` is the tick - that observed it). + that observed it). ``failure_reason`` is a stable machine-readable reason + next to the free-text ``reason`` (for example ``replicas_changed``, or + the live sweep executor's), and ``details`` carries JSON-shaped + executor facts; both are empty in replay. """ decision_id: str @@ -379,6 +385,8 @@ class RetuneDecision: decision_latency_s: Optional[float] = None top: list[tuple[str, float]] = field(default_factory=list) reason: Optional[str] = None + failure_reason: Optional[str] = None + details: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: _require_non_empty_string("decision_id", self.decision_id) @@ -394,6 +402,11 @@ def __post_init__(self) -> None: raise ValueError("previous and chosen must be RetuneParams") if self.predictor_warmup is not None: object.__setattr__(self, "predictor_warmup", tuple(self.predictor_warmup)) + if self.failure_reason is not None: + _require_non_empty_string("failure_reason", self.failure_reason) + # A private deep copy: the record must not change with the sweep + # result, whose details may nest mutable values. + object.__setattr__(self, "details", copy.deepcopy(self.details)) def to_dict(self) -> dict[str, Any]: """JSON-shaped record; ``predictor_warmup`` becomes its length.""" @@ -426,6 +439,8 @@ def to_dict(self) -> dict[str, Any]: "decision_latency_s": self.decision_latency_s, "top": [[key, score] for key, score in self.top], "reason": self.reason, + "failure_reason": self.failure_reason, + "details": json.loads(json.dumps(self.details)), } diff --git a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py index 2c5fa1f5d625..f39f06b1d479 100644 --- a/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py +++ b/components/src/dynamo/planner/plugins/orchestrator/engine_adapter.py @@ -62,8 +62,9 @@ from __future__ import annotations import logging -from collections.abc import Iterable -from dataclasses import replace +from collections.abc import Callable, Iterable +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Optional, Sequence from dynamo.common.forward_pass_metrics import encode as _encode_fpm_record @@ -74,6 +75,7 @@ import grpc.aio from dynamo.planner.retune.controller import RetuneController + from dynamo.planner.retune.live.contracts import HarvestAllowed, LaunchGate from dynamo.planner.core.batch_policy import plan_batch_schedule from dynamo.planner.core.budget import ( @@ -136,6 +138,26 @@ log = logging.getLogger(__name__) +@dataclass(frozen=True) +class RetuneHooks: + """Live retune hooks a planner host hands every retune controller. + + ``sweep_fn(request, *, cancel, abort)`` runs one sweep (the DGDR + executor); ``launch_gate`` and ``harvest_allowed`` are the live + supervisor's predicates (see ``retune.controller``); ``phase_s`` is the + epoch phase the window schedule is anchored at. Replay passes no hooks. + """ + + sweep_fn: Callable[..., Any] + launch_gate: LaunchGate + harvest_allowed: HarvestAllowed + phase_s: float = 0.0 + + +def _no_sweep() -> None: + """Queued behind the live sweeps by ``retune_sweeps_done``.""" + + class OrchestratorEngineAdapter: """``EngineProtocol``-compatible wrapper around the builtin plugin chain. @@ -155,6 +177,13 @@ class OrchestratorEngineAdapter: ``RetuneController`` built from ``retune_inputs``: ``initial_tick`` anchors its window schedule, every ``tick`` runs it after the plugin pipeline, and ``shutdown`` closes it. + + A live host also passes ``retune_hooks``. Live, missing + ``retune_inputs`` (the frontend's startup router config does not satisfy + the retune premises) or a controller that fails construction leaves the + adapter in a disabled state, retune configured but no controller, + instead of failing the planner; ``replace_retune`` later installs or + removes a controller without rebuilding the engine. """ def __init__( @@ -165,6 +194,7 @@ def __init__( observe_plugin: Optional[EnvironmentObserver] = None, clock: Optional[Clock] = None, retune_inputs: Optional[RetuneInputs] = None, + retune_hooks: Optional[RetuneHooks] = None, ) -> None: self._config = config self._capabilities = capabilities @@ -335,15 +365,46 @@ def _factory(plugin_id, endpoint, *, in_process_instance=None): self._plugins_bootstrapped = False self._in_process_plugins_loaded = False self._register_builtin_plugins() - self._retune = self._build_retune_controller(retune_inputs) + self._retune_hooks = retune_hooks + self._retune: Optional[RetuneController] = None + # Live: the one sweep thread every controller generation submits to, + # so a rebuilt controller's sweep queues behind a dropped one and two + # generations never sweep at once. + self._retune_executor: Optional[ThreadPoolExecutor] = None + if retune_hooks is None or self._config.retune is None: + self._retune = self._build_retune_controller(retune_inputs) + elif retune_inputs is None: + log.warning( + "planner.retune is configured but its live premises do not " + "hold; retune stays disabled until they do" + ) + else: + try: + self._retune = self._build_retune_controller( + retune_inputs, generation=0 + ) + except ValueError as exc: + # A live planner never exits on a retune premise: the + # supervisor keeps running (to roll routers back) and may + # install a controller later through replace_retune. + log.error( + "planner.retune controller construction failed; retune " + "stays disabled: %s", + exc, + ) def _build_retune_controller( - self, retune_inputs: Optional[RetuneInputs] + self, + retune_inputs: Optional[RetuneInputs], + *, + generation: int = 0, ) -> Optional[RetuneController]: """Build the retune controller iff ``config.retune`` is set. The controller validates the retune setup in its constructor, so an - invalid one fails here, at engine construction. + invalid one fails here, at engine construction. With live hooks the + controller gets them, ``generation`` and the adapter's shared sweep + executor; without (replay) it is built exactly as before. """ if self._config.retune is None: return None @@ -358,14 +419,100 @@ def _build_retune_controller( # this module. from dynamo.planner.retune.controller import RetuneController + hooks = self._retune_hooks + if hooks is None: + return RetuneController( + self._config, + retune_inputs, + self._capabilities, + self._clock, + self._scale_interval, + ) + if self._retune_executor is None: + self._retune_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="planner-retune" + ) return RetuneController( self._config, retune_inputs, self._capabilities, self._clock, self._scale_interval, + sweep_fn=hooks.sweep_fn, + launch_gate=hooks.launch_gate, + harvest_allowed=hooks.harvest_allowed, + generation=generation, + executor=self._retune_executor, ) + @property + def retune_controller(self) -> Optional[RetuneController]: + """The installed retune controller; ``None`` without retune or when + a live retune is disabled.""" + + return self._retune + + def replace_retune( + self, + retune_inputs: Optional[RetuneInputs], + *, + generation: int, + now_s: float, + ) -> None: + """Swap the live retune controller without rebuilding the engine. + + The old controller's sweep is aborted (``drop_pending``) and the + controller closed. With ``retune_inputs`` a new controller of + ``generation`` is built with the same hooks and anchored at ``now_s`` + and the hooks' phase; with ``None`` the adapter stays disabled. Every + generation submits to the adapter's one sweep thread, so the new + controller's first sweep queues behind the old one while that still + runs. A construction error is raised after the old controller is + gone, leaving the adapter disabled. Callers serialize this with + ticks (the planner's config lock). + """ + hooks = self._retune_hooks + if self._config.retune is None or hooks is None: + raise RuntimeError("replace_retune requires a live retune configuration") + old = self._retune + self._retune = None + if old is not None: + try: + old.drop_pending() + except Exception: + log.exception("retune controller drop_pending failed") + try: + old.close() + except Exception: + log.exception("retune controller close failed") + if retune_inputs is None: + log.info( + "retune controller generation %d: disabled (premises do not hold)", + generation, + ) + return + controller = self._build_retune_controller(retune_inputs, generation=generation) + if controller is None: + raise RuntimeError("replace_retune requires a live retune configuration") + controller.start(now_s, hooks.phase_s) + self._retune = controller + log.info("retune controller generation %d installed", generation) + + def retune_sweeps_done(self) -> Optional[Future[None]]: + """A future done once every live sweep submitted so far returned. + + The sweep thread runs one sweep at a time in submission order, so a + no-op queued behind them completes after the last one. ``None`` + without live retune or once the engine is shut down. + """ + executor = self._retune_executor + if executor is None: + return None + try: + return executor.submit(_no_sweep) + except RuntimeError: # shut down + return None + @property def plugins_bootstrapped(self) -> bool: return self._plugins_bootstrapped @@ -610,8 +757,12 @@ def initial_tick(self, start_s: float) -> ScheduledTick: if plugin.is_builtin and plugin.last_call_at == float("-inf"): plugin.registered_at = self._last_tick_monotonic if self._retune is not None: - # Anchors the retune window schedule (s0) at the planner start. - self._retune.start(start_s) + if self._retune_hooks is None: + # Anchors the retune window schedule (s0) at the planner start. + self._retune.start(start_s) + else: + # Live: anchored at the configured epoch phase instead. + self._retune.start(start_s, self._retune_hooks.phase_s) return self._compute_next_scheduled_tick() def apply_runtime_configuration( @@ -722,9 +873,13 @@ async def tick( decision_id=ctx.decision_id, ) worker_counts = tick_input.worker_counts or WorkerCounts() - # A retune-only planner keeps the deployment's replicas: no - # GPU-budget, min-endpoint, or batch-floor repair. - retune_only = self._retune is not None and not self._config.scaling_enabled() + # A retune-only planner (retune with both scaling modes off) keeps + # the deployment's replicas: no GPU-budget, min-endpoint, or + # batch-floor repair. Decided by config, so a disabled live retune + # (no controller) keeps them too. + retune_only = ( + self._config.retune is not None and not self._config.scaling_enabled() + ) scale_to = ( None if retune_only @@ -849,6 +1004,8 @@ async def shutdown(self) -> None: type(exc).__name__, exc, ) + if self._retune_executor is not None: + self._retune_executor.shutdown(wait=False, cancel_futures=True) await self._orchestrator.shutdown() # ------------------------------------------------------------------ diff --git a/components/src/dynamo/planner/reconfiguration.py b/components/src/dynamo/planner/reconfiguration.py index fe7827b36c1f..a8874a0d243f 100644 --- a/components/src/dynamo/planner/reconfiguration.py +++ b/components/src/dynamo/planner/reconfiguration.py @@ -244,6 +244,40 @@ def runtime_candidate_errors( ) +def live_retune_sla_errors( + config: PlannerConfig, candidate: PlannerConfig +) -> list[str]: + """Return why ``candidate`` cannot replace ``config`` under live retune. + + The live retune controller and the supervisor's state fingerprint read + the scoring SLA only when they are built: at startup, or at a rebuild for + a spec, P4 or DGD change. A runtime change of an inherited ``ttft_ms`` or + ``itl_ms`` would leave the running controller scoring with the old SLA + until such a rebuild, which would then apply the new SLA and, through its + new fingerprint, reset the persisted retune state. + """ + + if config.retune is None or candidate.retune is None: + return [] + current = config.retune.scoring_sla(config.ttft_ms, config.itl_ms) + proposed = candidate.retune.scoring_sla(candidate.ttft_ms, candidate.itl_ms) + changed = [ + field + for field, old, new in zip(("ttft_ms", "itl_ms"), current, proposed) + if old != new + ] + if not changed: + return [] + fields = " and ".join(changed) + overrides = " and ".join(f"retune.sla_{field}" for field in changed) + return [ + f"cannot change {fields} at runtime while retune.executor is 'dgdr': " + f"live retune inherits its scoring SLA from {fields} ({overrides} unset); " + f"restart the planner with the new {fields}, or restart it with " + f"{overrides} set to keep {fields} changeable at runtime" + ] + + def prepare_predictor( config: PlannerConfig, capabilities: WorkerCapabilities, diff --git a/components/src/dynamo/planner/retune/controller.py b/components/src/dynamo/planner/retune/controller.py index 4385d6ca3708..1f0b263e85df 100644 --- a/components/src/dynamo/planner/retune/controller.py +++ b/components/src/dynamo/planner/retune/controller.py @@ -23,6 +23,26 @@ known once the sweep is done. In replay the old parameters therefore serve traffic for the modeled sweep time; live, ``measured`` with scale 1 decides on the first tick after the sweep finishes. + +A live host (``retune.live``) passes hooks that replay never sets, so replay +behaviour is unchanged: + +- ``start(start_s, phase_s)`` anchors ``s0`` at an epoch phase instead of + the planner start, so windows and ``decision_id = retune-w`` + survive restarts and leader changes. +- ``launch_gate(k, window)`` is asked before a due window is launched or + recorded ``busy``: ``not_leader``/``not_ready`` leave the window due (it + is retried on a later tick), ``already_decided`` consumes it silently. +- ``harvest_allowed(generation)`` is asked before a finished sweep becomes + a decision; a refused or abandoned (``NON_DECISION_FAILURES``) sweep + records nothing and rewinds the window so it can be relaunched or adopted. +- ``drop_pending``/``rewind`` let the host abandon an in-flight sweep + without consuming its window; ``restore_desired`` re-installs persisted + router values. +- ``executor`` is the live host's single sweep thread, shared by every + controller generation it builds, so two generations never sweep at once: + a launch queues behind a dropped sweep that still runs. Harvest waits for + the queued sweep's future, so the queue delay needs no bookkeeping. """ from __future__ import annotations @@ -33,7 +53,7 @@ import math import threading from collections import deque -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from typing import Any @@ -63,6 +83,12 @@ grid_points, required_history_seconds, ) +from dynamo.planner.retune.live.contracts import ( + NON_DECISION_FAILURES, + GateResult, + HarvestAllowed, + LaunchGate, +) from dynamo.planner.retune.sweep import ( SweepRequest, SweepResult, @@ -76,6 +102,11 @@ _TOP_SCORES = 5 _EPSILON = 1e-9 +# Gate reasons that leave a due window due, to be retried on a later tick. +_GATE_RETRY_REASONS = frozenset({"not_leader", "not_ready"}) + +# ``sweep_fn(request, *, cancel)`` in replay; live hooks add ``abort``. +SweepFn = Callable[..., SweepResult] @dataclass(frozen=True) @@ -105,6 +136,8 @@ class _PendingSweep: # Both None while a ``measured`` sweep is still running. decision_latency_s: float | None ready_at_s: float | None + # Set by the supervisor to stop this sweep (lost leadership, rebuild). + abort: threading.Event | None = None class RetuneController: @@ -118,6 +151,15 @@ class RetuneController: from it. Under ``config.advisory`` a decision that would apply is recorded as ``advisory`` and leaves the desired state at the deployed point. + + The keyword-only hooks are for a live host and default to replay + behaviour: ``sweep_fn`` replaces ``run_sweep`` (and must accept + ``abort=`` when ``launch_gate`` is set), ``launch_gate`` and + ``harvest_allowed`` are the supervisor's predicates (module docstring), + and ``generation`` identifies this controller among the ones a live host + builds over its lifetime; every sweep request carries it. ``executor`` + is the host's single-thread sweep executor, shared across generations; + without one (replay) the controller owns its sweep thread. """ def __init__( @@ -127,12 +169,27 @@ def __init__( capabilities: WorkerCapabilities, clock: Clock, scale_interval_s: float, + *, + sweep_fn: SweepFn | None = None, + launch_gate: LaunchGate | None = None, + harvest_allowed: HarvestAllowed | None = None, + generation: int = 0, + executor: ThreadPoolExecutor | None = None, ) -> None: retune = config.retune if retune is None: raise ValueError("RetuneController requires config.retune") if not (math.isfinite(scale_interval_s) and scale_interval_s > 0): raise ValueError("scale_interval_s must be positive and finite") + if launch_gate is not None and sweep_fn is None: + # run_sweep takes no abort event; a gated launch passes one. + raise ValueError("RetuneController launch_gate requires a sweep_fn") + if isinstance(generation, bool) or not isinstance(generation, int): + raise ValueError("generation must be an integer") + self._sweep_fn = sweep_fn + self._launch_gate = launch_gate + self._harvest_allowed = harvest_allowed + self._generation = generation self._config = config self._retune = retune self._inputs = inputs @@ -156,6 +213,7 @@ def __init__( ) self._router_maps: dict[str, dict[str, dict[str, float]]] = {} router_domain = self._router_domain(retune) + self._router_domain_values = router_domain self._router_keys = tuple(sorted(router_domain)) self._router_searched = bool(router_domain) @@ -193,12 +251,15 @@ def __init__( self._startup_point = self.current_point() self._s0 = 0.0 + # Live: s0 is an epoch phase and decision ids carry the window end. + self._phase_anchored = False self._started = False self._last_k = -1 self._history: deque[tuple[float, dict[str, int]]] = deque() self._pending: _PendingSweep | None = None self._cancel = threading.Event() - self._executor = ThreadPoolExecutor( + self._owns_executor = executor is None + self._executor = executor or ThreadPoolExecutor( max_workers=1, thread_name_prefix="planner-retune" ) logger.info( @@ -226,10 +287,23 @@ def __init__( # ------------------------------------------------------------------ API - def start(self, start_s: float) -> None: - """Anchor the window schedule at the planner's first tick time.""" + def start(self, start_s: float, phase_s: float | None = None) -> None: + """Anchor the window schedule. + + Replay anchors ``s0`` at the planner's first tick time ``start_s``. + A live host passes ``phase_s``: ``s0`` is then that epoch phase, so + every planner process and controller generation cuts the same + windows, and decision ids become ``retune-w``. + """ - self._s0 = float(start_s) + if phase_s is None: + self._s0 = float(start_s) + self._phase_anchored = False + else: + if not math.isfinite(phase_s): + raise ValueError("phase_s must be finite") + self._s0 = float(phase_s) + self._phase_anchored = True self._started = True self._last_k = -1 @@ -259,19 +333,33 @@ async def on_tick( # must land in the tick that launches the next window. if now_s + _EPSILON * self._retune.interval_seconds >= pending.ready_at_s: self._pending = None - decisions.append(self._harvest(pending, now_s)) + harvested = self._harvest_or_rewind(pending, now_s) + if harvested is None: + audit.append(f"retune:dropped w={pending.request.window_index}") + else: + decisions.append(harvested) due = self._due_window(now_s) if due > self._last_k: - self._last_k = due - if self._pending is not None: - decisions.append(self._skipped(due, now_s, "busy", "sweep in flight")) + # The gate runs before both branches: a window this process may + # not work on is neither launched nor recorded busy. + gate = self._gate(due) + if gate is not None and gate.reason in _GATE_RETRY_REASONS: + pass + elif gate is not None and gate.reason == "already_decided": + self._last_k = due else: - launched = await self._launch(due, now_s) - if launched is None: - audit.append(f"retune:launch w={due}") + self._last_k = due + if self._pending is not None: + decisions.append( + self._skipped(due, now_s, "busy", "sweep in flight") + ) else: - decisions.append(launched) + launched = await self._launch(due, now_s, gate) + if launched is None: + audit.append(f"retune:launch w={due}") + else: + decisions.append(launched) for decision in decisions: audit.append(f"retune:{decision.outcome} w={decision.window_index}") @@ -283,10 +371,89 @@ async def on_tick( ) def close(self) -> None: - """Kill the sweep workers and release the sweep thread.""" + """Kill the sweep workers and release the sweep thread. + + A shared (host-owned) executor stays up: only this controller's + sweep is cancelled, or told to stop if it already runs. + """ self._cancel.set() - self._executor.shutdown(wait=False, cancel_futures=True) + if self._owns_executor: + self._executor.shutdown(wait=False, cancel_futures=True) + elif self._pending is not None: + self._pending.future.cancel() + + def drop_pending(self) -> None: + """Abandon the in-flight sweep without deciding its window. + + Sets the sweep's abort event, cancels it if it has not started, and + forgets it, so a following ``rewind`` takes effect and the window is + relaunched (or adopted by another leader) instead of recorded busy. + A sweep that is already running keeps the single sweep thread until + it sees the abort; a later launch queues behind it. + """ + + pending = self._pending + if pending is None: + return + self._pending = None + if pending.abort is not None: + pending.abort.set() + pending.future.cancel() + logger.info( + "retune dropped the in-flight sweep of window %d", + pending.request.window_index, + ) + + def rewind(self, last_decided_k: int) -> None: + """Make every window after ``last_decided_k`` due again. + + A no-op while a sweep is pending (``drop_pending`` first): its + window is consumed until it is harvested or dropped. + """ + + if self._pending is None: + self._last_k = min(self._last_k, int(last_decided_k)) + + def restore_desired(self, flat: Mapping[str, float]) -> None: + """Re-install persisted desired router values (live re-election). + + ``flat`` maps live router knob names to values and must cover every + searched key; extra keys are ignored. The startup point is not + touched: it stays a candidate of every sweep. + """ + + try: + values = {key: flat[key] for key in self._router_keys} + except KeyError as error: + raise ValueError( + f"restore_desired needs every searched router key " + f"{list(self._router_keys)}; missing {error.args[0]!r}" + ) from None + # Validates finite, non-negative values and a resolvable config, so + # the per-tick desired map can never fail afterwards. + router = RetuneParams(router=values).router + if self._router_searched: + self._router_map(router) + self._desired_router = router + + @property + def generation(self) -> int: + """The generation the host built this controller with.""" + + return self._generation + + @property + def router_keys(self) -> tuple[str, ...]: + """Searched router knobs, by live (PATCH) name, sorted.""" + + return self._router_keys + + @property + def router_domain(self) -> dict[str, list[float]]: + """The searched router domain: live knob name -> candidate values.""" + + return {key: list(values) for key, values in self._router_domain_values.items()} def current_point(self) -> RetuneParams: """The point in effect: desired router values and live planner values. @@ -308,63 +475,11 @@ def current_point(self) -> RetuneParams: def _router_domain(self, retune: RetuneConfig) -> dict[str, list[float]]: if retune.router is None: return {} - engine = self._inputs.engine - startup = self._startup_router_json - # Only the builtin default policy is modelled identically. - if engine.router_mode != "kv_router": + if self._inputs.engine.router_mode != "kv_router": raise ValueError("retune.router requires engine router_mode='kv_router'") - if startup.get("router_policy_config") is not None: - raise ValueError( - "retune.router requires the builtin default router policy " - "(router_config.router_policy_config must be unset)" - ) - if startup.get("overlap_score_weight") is not None: - raise ValueError( - "retune.router cannot be used with the deprecated " - "overlap_score_weight, which overrides prefill_load_scale" - ) - if startup.get("conditional_disagg_enabled"): - raise ValueError( - "retune.router does not support conditional disaggregation: " - "replay and live resolve decode credit differently" - ) - - startup_maps = self._router_map({}) - # A worker type that starts with no overlap credit can never get it. - zero_credit = { - worker_type - for worker_type, parameters in startup_maps.items() - if worker_type != "decode" and parameters[OVERLAP_SCORE_CREDIT] == 0 - } - # The schema keeps prefill_load_model.type at its startup value; the - # serde form of the AIS load model is "aic". - startup_load_model = json.loads( - KvRouterConfig.from_json(json.dumps(startup)).to_json() - )["router_prefill_load_model"] - domain = expand_router_domain( - retune, - startup_load_model=startup_load_model, - zero_startup_credit=bool(zero_credit), + return startup_router_domain( + retune, self._startup_router_json, self._worker_types ) - # The schema admits only PolicyParameters names, with the bounds - # KvRouterConfig checks; resolving every candidate also guards - # against the two drifting apart. - for point in grid_points(domain, {}): - try: - maps = self._router_map(point.router) - except ValueError as error: - raise ValueError( - f"retune.router candidate {point.router} is not a valid " - f"router configuration: {error}" - ) from None - for worker_type in zero_credit: - if maps[worker_type][OVERLAP_SCORE_CREDIT] != 0: - raise ValueError( - f"retune.router candidate {point.router} gives the " - f"{worker_type} router overlap_score_credit > 0, but its " - "startup credit is 0 (a live router rejects that update)" - ) - return domain def _feasible_grid( self, @@ -491,13 +606,127 @@ def _ms(self, time_s: float) -> float: return inputs.trace_origin_ms + time_s * 1000.0 * inputs.arrival_speedup_ratio def _decision_id(self, k: int) -> str: + if self._phase_anchored: + _, end_s, _ = self._window_times(k) + return f"retune-w{int(round(self._ms(end_s)))}" return f"retune-{int(self._s0)}-w{k}" - async def _launch(self, k: int, now_s: float) -> RetuneDecision | None: + def _gate(self, k: int) -> GateResult | None: + """Ask the live host whether window ``k`` may launch now. + + ``None`` without a gate (replay). A gate that raises or answers an + unknown reason holds the window like ``not_ready``: the window stays + due and the gate is asked again on the next tick. + """ + + if self._launch_gate is None: + return None + try: + gate = self._launch_gate(k, self._window_times(k)) + except Exception: + logger.exception("retune launch gate failed for window %d", k) + return GateResult(reason="not_ready") + if gate.reason is not None and gate.reason not in ( + *_GATE_RETRY_REASONS, + "already_decided", + ): + logger.warning( + "retune launch gate returned unknown reason %r for window %d; " + "holding the window", + gate.reason, + k, + ) + return GateResult(reason="not_ready") + return gate + + def _harvest_or_rewind( + self, pending: _PendingSweep, now_s: float + ) -> RetuneDecision | None: + """Turn a finished sweep into a decision, or abandon its window. + + The window is abandoned (no decision, ``_last_k`` rewound so it is + due again) when the host refuses the harvest, for instance because + this process lost leadership or rebuilt its controller, or when the + sweep reports a ``NON_DECISION_FAILURES`` reason. + """ + + k = pending.request.window_index + allowed = True + if self._harvest_allowed is not None: + try: + allowed = bool(self._harvest_allowed(pending.request.generation)) + except Exception: + logger.exception("retune harvest_allowed failed for window %d", k) + allowed = False + if not allowed: + logger.info("retune discards the sweep of window %d: harvest refused", k) + self.rewind(k - 1) + return None + result = self._sweep_result(pending) + if result.failure_reason in NON_DECISION_FAILURES: + logger.info( + "retune abandons window %d: sweep stopped (%s)", + k, + result.failure_reason, + ) + self.rewind(k - 1) + return None + return self._decide(pending, result, now_s) + + def _replicas_for_launch( + self, start_s: float, end_s: float, spec_replicas: Mapping[str, int] | None + ) -> tuple[dict[str, int] | None, bool, str | None]: + """``(replicas, verified, mismatch)`` for a window launch. + + Without ``spec_replicas`` (replay), the ready counts at window start, + verified. With them (live), every ready count recorded over the + window, including the one in effect at its start, must equal the + spec; a mismatch is reported. The spec replicas are used either way, + unverified when no recorded count covers the window start (the + window began before this controller started). + """ + + if spec_replicas is None: + return self._replicas_at(start_s), True, None + spec = {key: int(value) for key, value in spec_replicas.items()} + covering: dict[str, int] | None = None + observed: list[dict[str, int]] = [] + for at_s, counts in self._history: + if at_s <= start_s + _EPSILON: + covering = counts + elif at_s < end_s - _EPSILON: + observed.append(counts) + if covering is not None: + observed.insert(0, covering) + for counts in observed: + if counts != spec: + return spec, covering is not None, json.dumps(counts, sort_keys=True) + return spec, covering is not None, None + + async def _launch( + self, k: int, now_s: float, gate: GateResult | None = None + ) -> RetuneDecision | None: """Submit the sweep of window ``k``; a decision only if it cannot run.""" start_s, end_s, as_of_s = self._window_times(k) - replicas = self._replicas_at(start_s) + spec_replicas = gate.spec_replicas if gate is not None else None + replicas, verified, mismatch = self._replicas_for_launch( + start_s, end_s, spec_replicas + ) + if mismatch is not None: + return self._skipped( + k, + now_s, + "failed", + f"ready worker counts {mismatch} differ from the DGD spec " + f"replicas {json.dumps(spec_replicas, sort_keys=True)} during " + "the window", + failure_reason="replicas_changed", + details={ + "spec_replicas": dict(spec_replicas or {}), + "observed_replicas": json.loads(mismatch), + }, + ) if replicas is None: return self._skipped( k, now_s, "failed", "no ready worker counts recorded at window start" @@ -535,9 +764,12 @@ async def _launch(self, k: int, now_s: float) -> RetuneDecision | None: sla_itl_ms=self._sla_itl_ms, max_parallel=self._retune.max_parallel, timeout_s=self._retune.effective_timeout_seconds, + # 0 in replay, the field's default. + generation=self._generation, + **self._live_request_fields(as_of_s, spec_replicas, verified), ) - # Looked up at call time so tests can substitute the sweep function. - future = self._executor.submit(run_sweep, request, cancel=self._cancel) + abort = gate.abort if gate is not None else None + future = self._submit(request, abort, gated=gate is not None) if isinstance(self._clock, VirtualClock): # Virtual time is frozen while the sweep runs; wait for it here. await asyncio.wait([asyncio.wrap_future(future)]) @@ -555,17 +787,64 @@ async def _launch(self, k: int, now_s: float) -> RetuneDecision | None: launched_at_s=now_s, decision_latency_s=latency, ready_at_s=None if latency is None else now_s + latency, + abort=abort, ) logger.info( - "retune launched window %d [%s, %s) with %d candidates, replicas %s", + "retune launched window %d [%s, %s) with %d candidates, replicas %s%s", k, start_s, end_s, len(candidates), replicas, + "" + if verified + else " (unverified: the window began before this controller)", ) return None + def _live_request_fields( + self, + as_of_s: float, + spec_replicas: Mapping[str, int] | None, + verified: bool, + ) -> dict[str, Any]: + """The ``SweepRequest`` fields only a gated (live) launch sets.""" + + if self._launch_gate is None: + return {} + return { + "router_domain": { + key: tuple(float(value) for value in values) + for key, values in self._router_domain_values.items() + }, + "deadline_s": as_of_s + self._retune.effective_timeout_seconds, + "spec_replicas": None if spec_replicas is None else dict(spec_replicas), + "replicas_verified": verified, + } + + def _submit( + self, + request: SweepRequest, + abort: threading.Event | None, + *, + gated: bool, + ) -> Future[SweepResult]: + """Submit the sweep to the single sweep thread. + + ``run_sweep`` is looked up at call time so tests can substitute it. + A gated (live) launch also passes the gate's abort event. While the + sweep thread still runs a dropped sweep, of this or an earlier + controller generation, the launch queues behind it; the decision + waits for its future either way. + """ + + sweep_fn: SweepFn = self._sweep_fn if self._sweep_fn is not None else run_sweep + if not gated: + return self._executor.submit(sweep_fn, request, cancel=self._cancel) + return self._executor.submit( + sweep_fn, request, cancel=self._cancel, abort=abort + ) + def _decision_latency_s(self, future: Future[SweepResult]) -> float: """Delay from launch to decision; ``measured`` needs a done sweep.""" latency = self._retune.decision_latency_seconds @@ -581,15 +860,16 @@ def _decision_latency_s(self, future: Future[SweepResult]) -> float: ) return max(latency, self._scale_interval_s) - def _harvest(self, pending: _PendingSweep, now_s: float) -> RetuneDecision: + @staticmethod + def _sweep_result(pending: _PendingSweep) -> SweepResult: + """The finished sweep's result; a raised exception is ``failed``.""" + error = pending.future.exception() if error is not None: - result = SweepResult( + return SweepResult( status="failed", error=f"{type(error).__name__}: {error}" ) - else: - result = pending.future.result() - return self._decide(pending, result, now_s) + return pending.future.result() def _decide( self, pending: _PendingSweep, result: SweepResult, now_s: float @@ -621,6 +901,8 @@ def record(outcome: RetuneOutcome, **overrides: Any) -> RetuneDecision: "predictor_presweep": result.predictor_presweep, "sweep_wall_s": result.wall_s, "decision_latency_s": pending.decision_latency_s, + "failure_reason": result.failure_reason, + "details": result.details, } values.update(overrides) return RetuneDecision(**values) @@ -743,7 +1025,14 @@ def _improves(self, score: float, baseline_score: float) -> bool: return gain > margin def _skipped( - self, k: int, now_s: float, outcome: RetuneOutcome, reason: str + self, + k: int, + now_s: float, + outcome: RetuneOutcome, + reason: str, + *, + failure_reason: str | None = None, + details: Mapping[str, Any] | None = None, ) -> RetuneDecision: start_s, end_s, as_of_s = self._window_times(k) current = self.current_point() @@ -763,6 +1052,8 @@ def _skipped( baseline_score=None, best_score=None, reason=reason, + failure_reason=failure_reason, + details=dict(details or {}), ) # ------------------------------------------------------- router helpers @@ -804,7 +1095,7 @@ def _config_value(self, key: str) -> Any: def _log_decision(decision: RetuneDecision) -> None: logger.info( "retune decision %s: %s (baseline=%s best=%s requests=%d " - "chosen=%s planner_changes=%s reason=%s)", + "chosen=%s planner_changes=%s reason=%s failure_reason=%s)", decision.decision_id, decision.outcome, decision.baseline_score, @@ -817,6 +1108,7 @@ def _log_decision(decision: RetuneDecision) -> None: ), decision.planner_changes, decision.reason, + decision.failure_reason, ) @@ -846,6 +1138,78 @@ def _check_inner_planner_isolation(config: PlannerConfig) -> None: ) +def startup_router_domain( + retune: RetuneConfig, + startup_router_json: Mapping[str, Any], + worker_types: Sequence[str], +) -> dict[str, list[float]]: + """The ``retune.router`` domain on a startup ``KvRouterConfig`` JSON. + + The startup-dependent router checks of the controller constructor, which + live retune's premise P4 runs on the frontend startup it parses, so that + a startup P4 accepts builds a controller. ``worker_types`` are the router + worker types of the planner mode. Raises ``ValueError`` when the startup + cannot be retuned or a candidate does not resolve on it. + """ + + # Only the builtin default policy is modelled identically. + if startup_router_json.get("router_policy_config") is not None: + raise ValueError( + "retune.router requires the builtin default router policy " + "(router_config.router_policy_config must be unset)" + ) + if startup_router_json.get("overlap_score_weight") is not None: + raise ValueError( + "retune.router cannot be used with the deprecated " + "overlap_score_weight, which overrides prefill_load_scale" + ) + if startup_router_json.get("conditional_disagg_enabled"): + raise ValueError( + "retune.router does not support conditional disaggregation: " + "replay and live resolve decode credit differently" + ) + + def router_maps(router: Mapping[str, float]) -> dict[str, dict[str, float]]: + config = KvRouterConfig.from_json(json.dumps({**startup_router_json, **router})) + return {wt: dict(config.policy_parameters(wt)) for wt in worker_types} + + # A worker type that starts with no overlap credit can never get it. + zero_credit = { + worker_type + for worker_type, parameters in router_maps({}).items() + if worker_type != "decode" and parameters[OVERLAP_SCORE_CREDIT] == 0 + } + # The schema keeps prefill_load_model.type at its startup value; the + # serde form of the AIS load model is "aic". + startup_load_model = json.loads( + KvRouterConfig.from_json(json.dumps(dict(startup_router_json))).to_json() + )["router_prefill_load_model"] + domain = expand_router_domain( + retune, + startup_load_model=startup_load_model, + zero_startup_credit=bool(zero_credit), + ) + # The schema admits only PolicyParameters names, with the bounds + # KvRouterConfig checks; resolving every candidate also guards against + # the two drifting apart. + for point in grid_points(domain, {}): + try: + maps = router_maps(point.router) + except ValueError as error: + raise ValueError( + f"retune.router candidate {point.router} is not a valid " + f"router configuration: {error}" + ) from None + for worker_type in zero_credit: + if maps[worker_type][OVERLAP_SCORE_CREDIT] != 0: + raise ValueError( + f"retune.router candidate {point.router} gives the " + f"{worker_type} router overlap_score_credit > 0, but its " + "startup credit is 0 (a live router rejects that update)" + ) + return domain + + def _candidate_error( config: PlannerConfig, updates: dict[str, Any], diff --git a/components/src/dynamo/planner/retune/live/__init__.py b/components/src/dynamo/planner/retune/live/__init__.py new file mode 100644 index 000000000000..452ce8294de1 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Live retune: run the retune sweep through DGDR v1beta2 and apply the winner. + +The replay retune loop (``RetuneController`` and ``_decide``) is unchanged; +this package supplies its live pieces: the DGDR sweep executor, the +supervisor that owns leadership and persisted state, the router actuator, +the reconciler and the trace window cutter. ``contracts`` holds the types +they share. +""" diff --git a/components/src/dynamo/planner/retune/live/actuator.py b/components/src/dynamo/planner/retune/live/actuator.py new file mode 100644 index 000000000000..154e1c0dbce4 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/actuator.py @@ -0,0 +1,886 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Apply the retune router parameters to every frontend router. + +``RouterActuator.reconcile`` runs one round: list the DGD's frontend pods, +``GET /router_policy`` on each, keep the routers of the DGD's dynamo +namespace, latch each router's initial configuration, check it against the +parsed startup maps (P14), write what the target needs and report whether +every router has converged. + +Wire format (``lib/llm/src/http/service/router_policy.rs``):: + + GET /router_policy -> {"routers": [{router_id, model, namespace, + component, endpoint, worker_type, + configuration: {revision, parameters} | null}]} + PATCH /router_policy {router_id, expected_revision, parameters: {...}} + {router_id, expected_revision, reset: true} + -> one router object as above + +A PATCH replaces only the keys it names; ``reset: true`` restores every +parameter to the router's startup values. Both bump the router's revision +(``lib/kv-router/src/plugins/worker_selection/runtime.rs``), and GET never +returns the startup values, so the actuator latches them itself: the +configuration of a router seen at revision 0, or returned by a reset, is that +router's initial. + +Modes (``ActuationTarget.mode``): + +- ``apply``: PATCH explicit values for every one of the 8 policy parameters + that differs from ``desired_by_wt[worker_type]``. A router whose initial is + unknown (revision > 0, never latched) is reset once first to learn it; a + router whose initial differs from the startup map (P14) is never written, + so planner-parsed values cannot leak into keys the sweep did not search. + A router is PATCHed only once the ``initials`` passed in (the persisted + ones) hold its initial: a router latched in this round is held, reported + in ``new_initials`` and not converged, so a process stopped before the + state write never leaves a router off an initial nobody recorded. +- ``rollback``: reset only routers whose configuration differs from their + latched initial. Never by revision: a reset itself bumps the revision. +- ``observe``: no writes; latch initials and check parity only. + +``ActuationTarget.learn_unknown=False`` (advisory) never resets a router to +learn its initial: such a router is left alone, and rollback does not wait +for it. + +Targeted pods: Ready frontends (plus NotReady ones with ``include_not_ready``) +that have an IP, are not terminating and are not in phase Succeeded/Failed +(an evicted pod can keep its IP but serves nothing). + +Convergence: at least ``max(1, expected_frontends)`` Ready frontends, every +Ready frontend answered GET and has a router with a runtime configuration +for every worker type of the planner mode, and every router seen has its 8 +parameters equal to the target (``apply``: the desired map; +``rollback``/``observe``: the router's latched initial). A NotReady pod +that answers GET is checked like the others; one that does not answer is +ignored (a restarted frontend comes back at revision 0, its initial). No +routers at all is not converged. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Optional + +import aiohttp + +from dynamo.planner.retune.live.contracts import ( + POD_LABEL_COMPONENT_TYPE, + POD_LABEL_DGD_NAME, + POLICY_PARAMETER_KEYS, + ActuationReport, + ActuationTarget, + DgdRef, + PodInfo, + PolicyCallResult, + RetuneKube, + RetuneMetricsSink, + RouterInfo, + RouterPolicyClient, +) + +logger = logging.getLogger(__name__) + +ROUTER_POLICY_PATH = "/router_policy" +# Container port of the frontend when the pod has no port named ``http`` +# (deploy/operator/internal/consts/consts.go DynamoServicePort). +DEFAULT_FRONTEND_PORT = 8000 +# Consecutive rounds of GET 404 on a frontend after which its admin API is +# reported disabled (P12). +ADMIN_API_DISABLED_ROUNDS = 3 +# Largest error text kept from a failed call. +_MAX_ERROR_CHARS = 512 +# Relative/absolute tolerance of parameter comparisons. Values round-trip +# exactly through serde_json and Python floats; the tolerance only absorbs +# formatting noise from hand-written targets. +_REL_TOL = 1e-9 +_ABS_TOL = 1e-12 +# Statuses after which the pod is re-read once and the write retried. +_RETRY_AFTER_REGET = frozenset({404, 409, 503}) +# The frontend bounds each router's control call by CONTROL_TIMEOUT +# (lib/llm/src/http/service/router_policy.rs): GET then answers 200 without +# the routers that timed out, and PATCH answers 503. The client's total +# timeout must exceed it, or the client gives up first and turns those +# answers into status 0 (a whole pod lost for one stuck router). +ROUTER_CONTROL_TIMEOUT_S = 5.0 +DEFAULT_CLIENT_TIMEOUT_S = ROUTER_CONTROL_TIMEOUT_S + 3.0 +# Pod phases after which no container runs again. +_FINISHED_POD_PHASES = frozenset({"Succeeded", "Failed"}) +# (pod uid, restarts): one run of a frontend; its routers live only as long. +_PodInstance = tuple[str, Optional[int]] + + +# ----------------------------------------------------------------- client + + +class AiohttpRouterPolicyClient: + """``RouterPolicyClient`` over a dedicated ``aiohttp.ClientSession``. + + The session is created on first use inside the running event loop and + closed by ``close()``. Every call returns a ``PolicyCallResult`` and never + raises for HTTP, network or decoding failures: ``status`` is the HTTP + status, or 0 when no well-formed response arrived (connection error, + timeout, a pod without an IP, or a body that is not the documented JSON). + ``timeout_s`` bounds each call and must stay above the server's + ``ROUTER_CONTROL_TIMEOUT_S``. + """ + + def __init__( + self, + timeout_s: float = DEFAULT_CLIENT_TIMEOUT_S, + scheme: str = "http", + *, + default_port: int = DEFAULT_FRONTEND_PORT, + max_connections: int = 32, + ) -> None: + if not (math.isfinite(timeout_s) and timeout_s > 0): + raise ValueError("timeout_s must be positive and finite") + if scheme not in ("http", "https"): + raise ValueError(f"scheme must be 'http' or 'https', got {scheme!r}") + self._timeout = aiohttp.ClientTimeout(total=float(timeout_s)) + self._scheme = scheme + self._default_port = int(default_port) + self._max_connections = int(max_connections) + self._session: Optional[aiohttp.ClientSession] = None + + def url(self, pod: PodInfo) -> Optional[str]: + """The pod's ``/router_policy`` URL, or None without a pod IP.""" + + if not pod.ip: + return None + host = f"[{pod.ip}]" if ":" in pod.ip else pod.ip + port = pod.http_port if pod.http_port is not None else self._default_port + return f"{self._scheme}://{host}:{port}{ROUTER_POLICY_PATH}" + + async def get(self, pod: PodInfo) -> PolicyCallResult: + return await self._call(pod, "GET", None, _parse_list) + + async def patch( + self, + pod: PodInfo, + router_id: str, + expected_revision: int, + parameters: Mapping[str, float], + ) -> PolicyCallResult: + if not parameters: + # The server rejects an empty patch with 400; do not send one. + return PolicyCallResult(status=0, error="empty parameters") + body = { + "router_id": router_id, + "expected_revision": int(expected_revision), + "parameters": {key: float(value) for key, value in parameters.items()}, + } + return await self._call(pod, "PATCH", body, _parse_single) + + async def reset( + self, pod: PodInfo, router_id: str, expected_revision: int + ) -> PolicyCallResult: + body = { + "router_id": router_id, + "expected_revision": int(expected_revision), + "reset": True, + } + return await self._call(pod, "PATCH", body, _parse_single) + + async def close(self) -> None: + session, self._session = self._session, None + if session is not None and not session.closed: + await session.close() + + async def _call( + self, + pod: PodInfo, + method: str, + body: Optional[dict[str, Any]], + parse: Any, + ) -> PolicyCallResult: + url = self.url(pod) + if url is None: + return PolicyCallResult(status=0, error=f"pod {pod.name} has no IP") + try: + session = self._ensure_session() + async with session.request(method, url, json=body) as response: + raw = await response.read() + status = response.status + except asyncio.CancelledError: + raise + except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as exc: + return PolicyCallResult( + status=0, error=_truncate(f"{method} {url}: {exc!r}") + ) + if status != 200: + return PolicyCallResult(status=status, error=_error_text(raw)) + try: + return parse(json.loads(raw)) + except (ValueError, TypeError, KeyError) as exc: + return PolicyCallResult( + status=0, error=_truncate(f"malformed {method} response: {exc}") + ) + + def _ensure_session(self) -> aiohttp.ClientSession: + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession( + timeout=self._timeout, + connector=aiohttp.TCPConnector(limit=self._max_connections), + trust_env=False, + ) + return self._session + + +def _parse_router(item: Mapping[str, Any]) -> RouterInfo: + configuration = item.get("configuration") + parameters: Optional[dict[str, float]] = None + revision = 0 + if configuration is not None: + revision = int(configuration["revision"]) + parameters = { + str(key): float(value) + for key, value in dict(configuration["parameters"]).items() + if isinstance(value, (int, float)) and not isinstance(value, bool) + } + model = item.get("model") + return RouterInfo( + router_id=str(item["router_id"]), + namespace=str(item["namespace"]), + worker_type=str(item["worker_type"]), + revision=revision, + configuration=parameters, + model=None if model is None else str(model), + ) + + +def _parse_list(payload: Any) -> PolicyCallResult: + routers = payload["routers"] + if not isinstance(routers, list): + raise TypeError("routers is not a list") + return PolicyCallResult( + status=200, routers=tuple(_parse_router(item) for item in routers) + ) + + +def _parse_single(payload: Any) -> PolicyCallResult: + return PolicyCallResult(status=200, router=_parse_router(payload)) + + +def _error_text(raw: bytes) -> str: + text = raw.decode("utf-8", errors="replace") + try: + payload = json.loads(text) + except ValueError: + return _truncate(text) + if isinstance(payload, dict) and isinstance(payload.get("error"), str): + return _truncate(payload["error"]) + return _truncate(text) + + +def _truncate(text: str) -> str: + return text if len(text) <= _MAX_ERROR_CHARS else text[:_MAX_ERROR_CHARS] + "..." + + +# --------------------------------------------------------------- actuator + + +@dataclass +class _Round: + """Mutable bookkeeping of one ``reconcile`` call.""" + + target: ActuationTarget + initials: Mapping[str, Mapping[str, float]] + errors: dict[str, int] = field(default_factory=dict) + + def error(self, kind: str) -> None: + self.errors[kind] = self.errors.get(kind, 0) + 1 + + +@dataclass +class _PodOutcome: + pod: PodInfo + ok: bool # the last GET of this pod succeeded + get_status: Optional[int] = None + routers: tuple[RouterInfo, ...] = () # in-scope routers, final state + + +class RouterActuator: + """Converges the frontends' router policies to an ``ActuationTarget``. + + ``startup_maps`` is the parsed startup ``policy_parameters(wt)`` per + worker type of the planner mode (``aggregated``, or ``prefill`` and + ``decode``); its keys are the worker types the actuator manages. The + actuator keeps a local cache of the initials it latched, so a write in + the same round, or a state write that has not landed yet, never makes a + router look like it has an unknown initial. ``reconcile`` never raises + (except ``asyncio.CancelledError``). + """ + + def __init__( + self, + kube: RetuneKube, + client: RouterPolicyClient, + dgd: DgdRef, + startup_maps: dict[str, dict[str, float]], + metrics: RetuneMetricsSink, + max_concurrency: int = 8, + *, + list_timeout_s: float = 10.0, + admin_disabled_rounds: int = ADMIN_API_DISABLED_ROUNDS, + ) -> None: + if max_concurrency < 1: + raise ValueError("max_concurrency must be >= 1") + if admin_disabled_rounds < 1: + raise ValueError("admin_disabled_rounds must be >= 1") + self._kube = kube + self._client = client + self._dgd = dgd + self._metrics = metrics + self._max_concurrency = int(max_concurrency) + self._list_timeout_s = float(list_timeout_s) + self._admin_disabled_rounds = int(admin_disabled_rounds) + # router_id -> initial configuration latched by this process. + self._latched: dict[str, dict[str, float]] = {} + # router_id -> the pod instances it may run on: the one it was last + # seen on, or for an initial not seen by this process, every + # frontend listed when it was first missing. + self._router_hosts: dict[str, frozenset[_PodInstance]] = {} + # router_id -> signature of a write the router rejected with 422; + # the same write is not retried until the target changes. + self._rejected: dict[str, str] = {} + # pod uid -> consecutive rounds whose GET returned 404. + self._get_404_rounds: dict[str, int] = {} + self._last_pods: Optional[tuple[PodInfo, ...]] = None + self._startup_maps: dict[str, dict[str, float]] = {} + self.set_startup_maps(startup_maps) + + # ------------------------------------------------------------------ API + + @property + def label_selector(self) -> str: + return ( + f"{POD_LABEL_DGD_NAME}={self._dgd.name}," + f"{POD_LABEL_COMPONENT_TYPE}=frontend" + ) + + @property + def worker_types(self) -> tuple[str, ...]: + return tuple(self._startup_maps) + + @property + def last_pods(self) -> Optional[tuple[PodInfo, ...]]: + """Frontend pods of the last round's listing; None if it failed.""" + + return self._last_pods + + def set_startup_maps(self, startup_maps: Mapping[str, Mapping[str, float]]) -> None: + """Replace the parsed startup maps (after a frontend spec change).""" + + if not startup_maps: + raise ValueError("startup_maps must name at least one worker type") + self._startup_maps = { + worker_type: {key: float(value) for key, value in parameters.items()} + for worker_type, parameters in startup_maps.items() + } + self._rejected.clear() + + async def close(self) -> None: + """Close the policy client when it has a ``close`` coroutine (the + ``RouterPolicyClient`` protocol does not require one).""" + + close = getattr(self._client, "close", None) + if close is not None: + await close() + + async def reconcile( + self, + target: ActuationTarget, + *, + initials: Mapping[str, Mapping[str, float]], + expected_frontends: int, + ) -> ActuationReport: + """Run one round; see the module docstring. Never raises.""" + + try: + report = await self._reconcile(target, initials, expected_frontends) + except asyncio.CancelledError: + raise + except Exception: # never raise: the supervisor retries next round + logger.exception("retune router actuator round failed") + report = ActuationReport( + converged=False, + seen_router_ids=frozenset(initials) | frozenset(self._latched), + errors={"internal": 1}, + ) + self._metrics.router_converged(report.converged) + return report + + # ------------------------------------------------------------- the round + + async def _reconcile( + self, + target: ActuationTarget, + initials: Mapping[str, Mapping[str, float]], + expected_frontends: int, + ) -> ActuationReport: + state = _Round(target=target, initials=initials) + target_error = self._target_error(target) + if target_error is not None: + logger.error("retune router actuator: invalid target: %s", target_error) + state.error("invalid_target") + # Latch and check parity without writing; never report converged. + state.target = ActuationTarget(mode="observe") + + try: + pods = await asyncio.to_thread( + self._kube.list_pods, + self.label_selector, + timeout_s=self._list_timeout_s, + ) + except Exception as exc: # KubeError or a client bug: retry next round + logger.warning( + "retune router actuator: listing frontend pods failed: %s", exc + ) + self._last_pods = None + state.error("list_pods") + return ActuationReport( + converged=False, + seen_router_ids=frozenset(initials) | frozenset(self._latched), + errors=dict(state.errors), + ) + self._last_pods = tuple(pods) + + live = [pod for pod in pods if _serving(pod)] + ready = [pod for pod in live if pod.ready and pod.ip] + targeted = [ + pod + for pod in live + if pod.ip and (pod.ready or state.target.include_not_ready) + ] + + semaphore = asyncio.Semaphore(self._max_concurrency) + + async def run(pod: PodInfo) -> _PodOutcome: + async with semaphore: + return await self._process_pod(pod, state) + + outcomes = list(await asyncio.gather(*(run(pod) for pod in targeted))) + + self._track_admin_api(pods, outcomes) + converged = target_error is None and self._converged( + state, ready, outcomes, max(1, int(expected_frontends)) + ) + routers_present = self._routers_present( + ready, outcomes, max(1, int(expected_frontends)) + ) + parity_ok, mismatches = self._parity(state, outcomes) + seen = self._seen_router_ids(state, pods, outcomes) + self._prune(seen) + new_initials = { + router_id: dict(configuration) + for router_id, configuration in self._latched.items() + if router_id in seen + and not _same_map(initials.get(router_id), configuration) + } + admin_api_disabled = any( + count >= self._admin_disabled_rounds + for count in self._get_404_rounds.values() + ) + if mismatches: + logger.warning( + "retune router actuator: P14 parity mismatch: %s", "; ".join(mismatches) + ) + return ActuationReport( + converged=converged, + new_initials=new_initials, + seen_router_ids=frozenset(seen), + parity_ok=parity_ok, + parity_mismatches=tuple(mismatches), + admin_api_disabled=admin_api_disabled, + errors=dict(state.errors), + routers_present=routers_present, + ) + + def _routers_present( + self, + ready: Sequence[PodInfo], + outcomes: Sequence[_PodOutcome], + required_ready: int, + ) -> bool: + """P12: enough Ready frontends, each GET answered with a configured + router for every worker type of the planner mode.""" + + if len(ready) < required_ready: + return False + ready_names = {pod.name for pod in ready} + answered = 0 + for outcome in outcomes: + if outcome.pod.name not in ready_names: + continue + if not outcome.ok: + return False + covered = { + router.worker_type + for router in outcome.routers + if router.configuration is not None + } + if covered != set(self._startup_maps): + return False + answered += 1 + return answered >= required_ready + + def _target_error(self, target: ActuationTarget) -> Optional[str]: + if target.mode not in ("apply", "rollback", "observe"): + return f"unknown mode {target.mode!r}" + if target.mode != "apply": + return None + desired = target.desired_by_wt + if desired is None: + return "apply needs desired_by_wt" + for worker_type in self._startup_maps: + parameters = desired.get(worker_type) + if parameters is None: + return f"desired_by_wt has no {worker_type!r} map" + for key in POLICY_PARAMETER_KEYS: + value = parameters.get(key) + if not _valid_parameter(value): + return f"desired_by_wt[{worker_type!r}][{key!r}] is {value!r}" + return None + + # ------------------------------------------------------------- per pod + + async def _process_pod(self, pod: PodInfo, state: _Round) -> _PodOutcome: + result = await self._client.get(pod) + if result.status != 200: + state.error(f"get:{_code(result.status)}") + logger.debug( + "retune router actuator: GET %s -> %s %s", + pod.name, + result.status, + result.error, + ) + return _PodOutcome(pod=pod, ok=False, get_status=result.status) + routers = self._in_scope(pod, result.routers, state) + retried = False + while True: + final: list[RouterInfo] = [] + retry = False + for router in routers: + current, needs_reget = await self._drive(pod, router, state) + final.append(current) + retry = retry or needs_reget + if not retry or retried: + return _PodOutcome( + pod=pod, ok=True, get_status=200, routers=tuple(final) + ) + # 409, 503 or PATCH 404: re-read the pod once and retry. + retried = True + result = await self._client.get(pod) + if result.status != 200: + state.error(f"get:{_code(result.status)}") + return _PodOutcome(pod=pod, ok=False, get_status=result.status) + routers = self._in_scope(pod, result.routers, state) + + def _in_scope( + self, pod: PodInfo, routers: Sequence[RouterInfo], state: _Round + ) -> list[RouterInfo]: + """Routers of the DGD's dynamo namespace and the planner's worker + types; latch the initials of routers at revision 0.""" + + scoped = [] + for router in routers: + if router.namespace != self._dgd.dynamo_namespace: + continue + if router.worker_type not in self._startup_maps: + continue + self._router_hosts[router.router_id] = frozenset({_instance(pod)}) + if router.configuration is None: + scoped.append(router) + continue + if router.revision == 0 and self._initial(router.router_id, state) is None: + self._latch(router.router_id, router.configuration) + scoped.append(router) + return scoped + + def _initial(self, router_id: str, state: _Round) -> Optional[dict[str, float]]: + latched = self._latched.get(router_id) + if latched is not None: + return latched + persisted = state.initials.get(router_id) + return None if persisted is None else dict(persisted) + + def _latch(self, router_id: str, configuration: Mapping[str, float]) -> None: + self._latched[router_id] = { + key: float(value) for key, value in configuration.items() + } + + async def _drive( + self, pod: PodInfo, router: RouterInfo, state: _Round + ) -> tuple[RouterInfo, bool]: + """Write what ``router`` needs; return its final state and whether + the pod should be re-read and the router retried.""" + + current = router + # At most a learning reset followed by one PATCH. + for _ in range(2): + action = self._plan(current, state) + if action is None: + return current, False + kind, parameters = action + signature = self._signature(kind, current.worker_type, state) + if self._rejected.get(current.router_id) == signature: + return current, False + if kind == "reset": + result = await self._client.reset( + pod, current.router_id, current.revision + ) + else: + result = await self._client.patch( + pod, current.router_id, current.revision, parameters + ) + if result.status != 200 or result.router is None: + code = _code(result.status) + state.error(f"{kind}:{code}") + self._metrics.patch_error(code) + logger.warning( + "retune router actuator: %s %s router %s rev %d -> %s %s", + kind, + pod.name, + current.router_id, + current.revision, + result.status, + result.error, + ) + if result.status == 422: + self._rejected[current.router_id] = signature + return current, result.status in _RETRY_AFTER_REGET + self._rejected.pop(current.router_id, None) + current = result.router + if kind == "reset" and current.configuration is not None: + # A reset restores every parameter to the router's startup + # values: that response is the router's initial. + self._latch(current.router_id, current.configuration) + return current, False + + def _plan( + self, router: RouterInfo, state: _Round + ) -> Optional[tuple[str, dict[str, float]]]: + """``("reset", {})``, ``("patch", diff)`` or None for one router.""" + + mode = state.target.mode + configuration = router.configuration + if mode == "observe" or configuration is None: + return None + initial = self._initial(router.router_id, state) + if initial is None and not state.target.learn_unknown: + # Never latched, and learning by reset is off (advisory): the + # planner never wrote it, so it is not the planner's to touch. + return None + if mode == "rollback": + if initial is None or not _same_map(configuration, initial): + return ("reset", {}) + return None + # apply + if initial is None: + # Revision > 0 and never latched: learn the initial first. + return ("reset", {}) + if not _same_map(state.initials.get(router.router_id), initial): + # Latched but not persisted yet: hold until the state records it. + return None + if not _same_map(initial, self._startup_maps[router.worker_type]): + # P14 failed for this router; writing explicit values would put + # planner-parsed values into keys the sweep never searched. + return None + desired = state.target.desired_by_wt[router.worker_type] # type: ignore[index] + diff = { + key: float(desired[key]) + for key in POLICY_PARAMETER_KEYS + if not _close(configuration.get(key), desired[key]) + } + return ("patch", diff) if diff else None + + @staticmethod + def _signature(kind: str, worker_type: str, state: _Round) -> str: + """Identity of a write for the 422 memory: mode, kind and target.""" + + desired = (state.target.desired_by_wt or {}).get(worker_type) + return json.dumps( + [state.target.mode, kind, desired], sort_keys=True, default=float + ) + + # ------------------------------------------------------------ verdicts + + def _left_alone(self, router: RouterInfo, state: _Round) -> bool: + """A rollback without learning skips routers with no known initial.""" + + return ( + state.target.mode == "rollback" + and not state.target.learn_unknown + and self._initial(router.router_id, state) is None + ) + + def _target_map( + self, router: RouterInfo, state: _Round + ) -> Optional[Mapping[str, float]]: + if state.target.mode == "apply": + return state.target.desired_by_wt[router.worker_type] # type: ignore[index] + return self._initial(router.router_id, state) + + def _converged( + self, + state: _Round, + ready: Sequence[PodInfo], + outcomes: Sequence[_PodOutcome], + required_ready: int, + ) -> bool: + if len(ready) < required_ready: + return False + configured = 0 + for outcome in outcomes: + if not outcome.ok: + if outcome.pod.ready: + return False + # An unreachable NotReady pod (crash loop, still starting) is + # ignored: it comes back at revision 0, i.e. at its initial. + continue + covered: set[str] = set() + for router in outcome.routers: + if router.configuration is None: + continue + covered.add(router.worker_type) + configured += 1 + if self._left_alone(router, state): + continue + expected = self._target_map(router, state) + if expected is None or not _same_map(router.configuration, expected): + return False + if outcome.pod.ready and covered != set(self._startup_maps): + return False + return configured > 0 + + def _parity( + self, state: _Round, outcomes: Sequence[_PodOutcome] + ) -> tuple[Optional[bool], list[str]]: + mismatches: list[str] = [] + known = 0 + unknown = 0 + for outcome in outcomes: + for router in outcome.routers: + if router.configuration is None: + continue + initial = self._initial(router.router_id, state) + if initial is None: + unknown += 1 + continue + known += 1 + startup = self._startup_maps[router.worker_type] + differing = [ + f"{key} initial={initial.get(key)} startup={startup.get(key)}" + for key in POLICY_PARAMETER_KEYS + if not _close(initial.get(key), startup.get(key)) + ] + if differing: + mismatches.append( + f"{outcome.pod.name}/{router.router_id}/{router.worker_type}: " + + ", ".join(differing) + ) + if mismatches: + return False, mismatches + if known == 0 or unknown: + return None, mismatches + return True, mismatches + + def _seen_router_ids( + self, + state: _Round, + pods: Sequence[PodInfo], + outcomes: Sequence[_PodOutcome], + ) -> set[str]: + """Routers observed this round plus known ones not provably gone. + + A GET answers without the routers whose control call timed out or + failed (``router_policy.rs``), so a router missing from its pod's + answer may still be there. A known router is gone only once every + pod instance it may run on is gone: no longer listed, finished + (Succeeded/Failed), or restarted (a restarted frontend registers new + routers). + """ + + observed = { + router.router_id for outcome in outcomes for router in outcome.routers + } + present = frozenset( + _instance(pod) for pod in pods if pod.phase not in _FINISHED_POD_PHASES + ) + seen = set(observed) + for router_id in set(state.initials) | set(self._latched): + if router_id in observed: + continue + # A persisted initial this process has not seen on any pod: its + # router runs on a frontend listed now (a pod never comes back), + # or on none. + hosts = self._router_hosts.setdefault(router_id, present) + if hosts & present: + seen.add(router_id) + return seen + + def _prune(self, seen: set[str]) -> None: + for cache in (self._latched, self._router_hosts, self._rejected): + for router_id in [key for key in cache if key not in seen]: + del cache[router_id] + + def _track_admin_api( + self, pods: Sequence[PodInfo], outcomes: Sequence[_PodOutcome] + ) -> None: + live_uids = {pod.uid for pod in pods} + for uid in [uid for uid in self._get_404_rounds if uid not in live_uids]: + del self._get_404_rounds[uid] + for outcome in outcomes: + uid = outcome.pod.uid + if outcome.get_status == 404: + self._get_404_rounds[uid] = self._get_404_rounds.get(uid, 0) + 1 + elif outcome.get_status is not None: + self._get_404_rounds.pop(uid, None) + + +# ---------------------------------------------------------------- helpers + + +def _instance(pod: PodInfo) -> _PodInstance: + """One run of a pod's frontend: its uid and container restart count.""" + + return (pod.uid, pod.restarts) + + +def _serving(pod: PodInfo) -> bool: + """Not terminating and not finished: its frontend may still answer.""" + + return not pod.terminating and pod.phase not in _FINISHED_POD_PHASES + + +def _valid_parameter(value: Any) -> bool: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return math.isfinite(value) and value >= 0 + + +def _code(status: int) -> str: + return str(status) if status else "network" + + +def _close(left: Optional[float], right: Optional[float]) -> bool: + if left is None or right is None: + return False + return math.isclose(float(left), float(right), rel_tol=_REL_TOL, abs_tol=_ABS_TOL) + + +def _same_map( + left: Optional[Mapping[str, float]], right: Optional[Mapping[str, float]] +) -> bool: + """Equal on the 8 policy parameters (a missing key is never equal).""" + + if left is None or right is None: + return False + return all(_close(left.get(key), right.get(key)) for key in POLICY_PARAMETER_KEYS) diff --git a/components/src/dynamo/planner/retune/live/bootstrap.py b/components/src/dynamo/planner/retune/live/bootstrap.py new file mode 100644 index 000000000000..274a73123f18 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/bootstrap.py @@ -0,0 +1,580 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Build the live retune supervisor of a Kubernetes planner. + +``build_live_retune(host)`` runs in ``NativePlannerBase._async_init`` after +``environment.initialize()`` and before the engine is built. It reads the +planner's DGD, parses the frontend's startup router configuration with a +pure function of the DGD spec, evaluates the spec premises, and returns a +``LiveRetuneSupervisor`` whose ``retune_inputs()`` the engine's first +controller is built from (``None`` while P4 fails: the adapter then starts +with retune disabled) and whose ``hooks()`` every controller gets. + +This module imports the whole live retune stack (Kubernetes client, router +actuator, reconciler, DGDR executor, window cutter). +``NativePlannerBase`` imports it only when the planner runs live retune or +must clean up after a removed one; ``build_supervisor`` takes the +collaborators injectable for tests. + +``build_legacy_cleanup(host, kube)`` is the clean-up for a planner whose +retune was removed. ``NativePlannerBase`` calls it in the background, only +after finding that a previous live retune left its Leases behind; it returns +a supervisor in decommission mode that rolls back what that retune deployed, +then deletes its Requests and Leases. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import posixpath +import socket +import uuid +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Optional + +from dynamo.llm import KvRouterConfig +from dynamo.planner.config.defaults import SubComponentType +from dynamo.planner.config.planner_config import ( + DgdrExecutorConfig, + PlannerConfig, + RetuneConfig, + RetuneEngineSpec, +) +from dynamo.planner.core.types import RetuneInputs +from dynamo.planner.errors import ComponentError +from dynamo.planner.monitoring.dgd_services import get_component_from_type_or_name +from dynamo.planner.retune.domains import expand_router_domain +from dynamo.planner.retune.live import frontend_config, window_cut +from dynamo.planner.retune.live.actuator import ( + AiohttpRouterPolicyClient, + RouterActuator, +) +from dynamo.planner.retune.live.contracts import ( + DgdRef, + FrontendStartup, + KubeAlreadyExists, + KubeInvalid, + PremiseReport, + RetuneKube, + RetuneMetricsSink, + WindowCutter, +) +from dynamo.planner.retune.live.dgdr import DgdrContext, Shard, build_request +from dynamo.planner.retune.live.kube import KubernetesRetuneKube +from dynamo.planner.retune.live.metrics import RetuneMetrics +from dynamo.planner.retune.live.reconciler import Reconciler +from dynamo.planner.retune.live.state import fingerprint +from dynamo.planner.retune.live.supervisor import ( + ActuatorFactory, + LiveRetuneSupervisor, + SpecEvaluation, + SupervisorSettings, + worker_types_for, +) + +logger = logging.getLogger(__name__) + +# Shard id and window of the P11 canary Request (dry run only, never stored). +CANARY_SHARD_ID = "canary" +CANARY_WINDOW_END_MS = 0 + + +# ------------------------------------------------------------- resolver + + +class KvRouterResolver: + """``RouterResolver`` over the parsed startup ``KvRouterConfig`` JSON. + + Resolves exactly as the retune controller does: the startup value of a + searched key is its serde value, or its aggregated resolution when the + startup leaves it unset; a point's per worker type parameters are + ``policy_parameters(wt)`` of ``startup ⊕ point``. + """ + + def __init__( + self, router_config: Mapping[str, Any], worker_types: Sequence[str] + ) -> None: + self._router_config = dict(router_config) + self._worker_types = tuple(worker_types) + + def startup_flat(self, keys: Sequence[str]) -> dict[str, float]: + startup = KvRouterConfig.from_json(json.dumps(self._router_config)) + serialized = json.loads(startup.to_json()) + resolved = startup.policy_parameters("aggregated") + values: dict[str, float] = {} + for key in keys: + value = serialized.get(key) + values[key] = float(resolved[key] if value is None else value) + return values + + def by_wt(self, flat: Mapping[str, float]) -> dict[str, dict[str, float]]: + config = KvRouterConfig.from_json(json.dumps({**self._router_config, **flat})) + return { + worker_type: { + key: float(value) + for key, value in dict(config.policy_parameters(worker_type)).items() + } + for worker_type in self._worker_types + } + + +# ------------------------------------------------------------ DGD reads + + +def worker_spec_replicas( + dgd: Mapping[str, Any], planner_mode: str +) -> Optional[dict[str, int]]: + """DGD spec replicas per worker type of the planner mode. + + ``None`` when a worker component cannot be resolved or has no replicas. + An omitted ``replicas`` is the workload default of 1. + """ + + roles = ( + (("prefill", SubComponentType.PREFILL), ("decode", SubComponentType.DECODE)) + if planner_mode == "disagg" + else (("aggregated", SubComponentType.DECODE),) + ) + replicas: dict[str, int] = {} + for worker_type, role in roles: + try: + service = get_component_from_type_or_name(dict(dgd), role) + except ComponentError as error: + logger.warning( + "retune: cannot resolve the %s workers: %s", worker_type, error + ) + return None + raw = service.service.get("replicas") + try: + count = 1 if raw is None else int(raw) + except (TypeError, ValueError): + return None + if count < 1: + return None + replicas[worker_type] = count + return replicas + + +def build_retune_inputs( + config: PlannerConfig, + startup: FrontendStartup, + *, + dgd_name: str, + model_name: Optional[str], +) -> RetuneInputs: + """Live ``RetuneInputs``: the parsed startup, traces on the claim. + + The engine pins only the router (the DGDR Run's engine comes from + ``retune.dgdr.engine``); ``prefill/decode_engine_args`` are empty maps in + disagg so the controller's engine-shape check holds. Live trace time is + epoch milliseconds (origin 0, speedup 1). + """ + + _, dgdr = _dgdr_retune(config) + disagg = config.mode == "disagg" + engine = RetuneEngineSpec( + router_mode=startup.router_mode, # type: ignore[arg-type] + router_config=dict(startup.router_config), + model_name=model_name or None, + prefill_engine_args={} if disagg else None, + decode_engine_args={} if disagg else None, + extra_engine_args=None, + ) + trace = posixpath.join(dgdr.mount_path, dgdr.segments_dir.strip("/"), dgd_name) + return RetuneInputs( + engine=engine, + trace_paths=(trace,), + trace_origin_ms=0.0, + arrival_speedup_ratio=1.0, + ) + + +def evaluate_spec( + dgd: Mapping[str, Any], + *, + config: PlannerConfig, + k8s_namespace: str, + dynamo_namespace: str, + frontend: Any = None, +) -> SpecEvaluation: + """Parse the frontend and evaluate the spec premises of one DGD read. + + ``frontend`` is the ``frontend_config`` module (injectable for tests). + """ + + if frontend is None: + frontend = frontend_config + metadata = dgd.get("metadata") or {} + name = str(metadata.get("name") or "") + uid = metadata.get("uid") + if not name or not uid: + raise ValueError("the DGD has no metadata.name or metadata.uid") + ref = DgdRef( + name=name, + namespace=str(metadata.get("namespace") or k8s_namespace), + uid=str(uid), + dynamo_namespace=dynamo_namespace, + ) + parsed = frontend.parse_frontend(dgd, planner_mode=config.mode) + startup = parsed if isinstance(parsed, FrontendStartup) else None + report: PremiseReport = frontend.evaluate_premises(dgd, planner_config=config) + inputs = None + if report.p4_ok and startup is not None: + try: + inputs = build_retune_inputs( + config, + startup, + dgd_name=name, + model_name=frontend.served_model_name(dgd), + ) + except ValueError as error: + logger.error("retune: building the live retune inputs failed: %s", error) + report = PremiseReport( + p4_ok=False, + spec_ok=False, + reasons=(*report.reasons, f"P4: {error}"), + warnings=report.warnings, + startup=report.startup, + ) + generation = metadata.get("generation") + return SpecEvaluation( + report=report, + startup=startup, + inputs=inputs, + spec_replicas=worker_spec_replicas(dgd, config.mode), + dgd=ref, + generation=generation if isinstance(generation, int) else None, + ) + + +# --------------------------------------------------------- DGDR pieces + + +def dgdr_context_factory( + config: PlannerConfig, cut_window: Optional[WindowCutter] = None +) -> Callable[[LiveRetuneSupervisor], Any]: + """Bind a ``dgdr.DgdrContext`` to the supervisor's current DGD.""" + + def factory(supervisor: LiveRetuneSupervisor) -> DgdrContext: + retune, dgdr = _dgdr_retune(config) + return DgdrContext( + dgd=supervisor.dgd, + cfg=dgdr, + retune=retune, + backend=str(config.backend), + kube=supervisor.kube, + registry=supervisor.registry, + cut_window=cut_window or window_cut.cut_window, + is_leader=supervisor.is_leader, + ) + + return factory + + +def dgdr_canary( + config: PlannerConfig, *, timeout_s: float = 10.0 +) -> Callable[[LiveRetuneSupervisor], bool]: + """P11: v1beta2 accepts ``workload.trace`` (a Strict dry-run create). + + The canary is the single-point Request of the startup router point that + a real window would create, so a CRD or webhook that drops or rejects any + field the executor relies on fails it. ``KubeInvalid`` is a failed + premise; other API errors propagate (the supervisor counts them as + failing and retries next round). + """ + + def canary(supervisor: LiveRetuneSupervisor) -> bool: + retune, dgdr = _dgdr_retune(config) + replicas = supervisor.evaluation.spec_replicas + point = supervisor.startup_flat() + if not replicas or not point: + return False + sla_ttft_ms, sla_itl_ms = retune.scoring_sla(config.ttft_ms, config.itl_ms) + shard = Shard( + id=CANARY_SHARD_ID, + kind="point", + choices={key: (value,) for key, value in point.items()}, + points=(dict(point),), + ) + body = build_request( + dgd=supervisor.dgd, + cfg=dgdr, + retune=retune, + sla_ttft_ms=sla_ttft_ms, + sla_itl_ms=sla_itl_ms, + window_start_ms=CANARY_WINDOW_END_MS - 1, + window_end_ms=CANARY_WINDOW_END_MS, + as_of_ms=CANARY_WINDOW_END_MS, + shard=shard, + spec_replicas=replicas, + backend=str(config.backend), + attempt=0, + ) + try: + supervisor.kube.create_request(body, dry_run=True, timeout_s=timeout_s) + except KubeAlreadyExists: + return True + except KubeInvalid as error: + logger.error( + "retune P11: the DGDR v1beta2 API rejects the Request: %s", error + ) + return False + return True + + return canary + + +def default_actuator_factory( + kube: RetuneKube, metrics: RetuneMetricsSink +) -> ActuatorFactory: + def factory( + dgd: DgdRef, startup_maps: dict[str, dict[str, float]] + ) -> RouterActuator: + return RouterActuator( + kube, AiohttpRouterPolicyClient(), dgd, startup_maps, metrics + ) + + return factory + + +def planner_identity() -> str: + """``/``: the Lease holder identity of this process. + + The nonce makes a restarted container a new holder, so it never renews + the Lease of its previous incarnation. + """ + + pod = ( + os.environ.get("POD_NAME") + or os.environ.get("DYN_POD_NAME") + or os.environ.get("HOSTNAME") + or socket.gethostname() + ) + return f"{pod}/{uuid.uuid4().hex[:12]}" + + +def searched_router_keys(config: PlannerConfig) -> tuple[str, ...]: + """Live names of the searched router keys (independent of the startup).""" + + retune = config.retune + if retune is None: + return () + return tuple(sorted(expand_router_domain(retune))) + + +def _dgdr_retune(config: PlannerConfig) -> tuple[RetuneConfig, DgdrExecutorConfig]: + """``config.retune`` and its DGDR executor section, which must be set.""" + + retune = config.retune + if retune is None or retune.dgdr is None: + raise ValueError("live retune requires retune.executor 'dgdr'") + return retune, retune.dgdr + + +def _runtime_namespace(host: Any) -> str: + """The routers' dynamo namespace, else the planner's configured one.""" + + value = host.environment.runtime_namespace() + return str(value) if value else str(host.config.namespace) + + +# --------------------------------------------------------------- build + + +async def build_supervisor( + config: PlannerConfig, + *, + kube: RetuneKube, + dgd_name: str, + namespace_fn: Callable[[], str], + metrics: Optional[RetuneMetricsSink] = None, + identity: Optional[str] = None, + frontend: Any = None, + actuator_factory: Optional[ActuatorFactory] = None, + reconciler: Any = None, + canary: Optional[Callable[[LiveRetuneSupervisor], bool]] = None, + context_factory: Optional[Callable[[LiveRetuneSupervisor], Any]] = None, + sweep_runner: Optional[Callable[..., Any]] = None, + **supervisor_kwargs: Any, +) -> LiveRetuneSupervisor: + """A supervisor for ``dgd_name`` with the production collaborators by default.""" + + retune = config.retune + if retune is None or retune.executor != "dgdr" or retune.dgdr is None: + raise ValueError( + "live retune requires retune.executor 'dgdr' (retune.executor " + f"is {None if retune is None else retune.executor!r})" + ) + dgdr = retune.dgdr + if metrics is None: + metrics = RetuneMetrics() + dgd_obj = await asyncio.to_thread(kube.get_dgd, dgd_name) + + def evaluate(obj: dict) -> SpecEvaluation: + return evaluate_spec( + obj, + config=config, + k8s_namespace=kube.namespace, + dynamo_namespace=namespace_fn(), + frontend=frontend, + ) + + evaluation = evaluate(dgd_obj) + if reconciler is None: + reconciler = Reconciler( + kube, evaluation.dgd, dgdr, retune, dgdr.mount_path, metrics + ) + worker_types = worker_types_for(config.mode) + settings = SupervisorSettings.from_config( + config, router_keys=searched_router_keys(config) + ) + supervisor = LiveRetuneSupervisor( + settings=settings, + kube=kube, + evaluation=evaluation, + identity=identity or planner_identity(), + premise_evaluator=evaluate, + actuator_factory=actuator_factory or default_actuator_factory(kube, metrics), + fingerprint_fn=lambda startup: fingerprint(config, startup), + resolver_factory=lambda startup: ( + None + if startup is None + else KvRouterResolver(startup.router_config, worker_types) + ), + reconciler=reconciler, + metrics=metrics, + canary=canary or dgdr_canary(config), + sweep_runner=sweep_runner, + context_factory=context_factory or dgdr_context_factory(config), + **supervisor_kwargs, + ) + report = evaluation.report + logger.info( + "live retune for DGD %s/%s (uid %s, dynamo namespace %s) as %s: " + "P4 %s, spec_ok %s, spec replicas %s, advisory %s", + evaluation.dgd.namespace, + evaluation.dgd.name, + evaluation.dgd.uid, + evaluation.dgd.dynamo_namespace, + supervisor.identity, + "ok" if report.p4_ok else "failed", + report.spec_ok, + evaluation.spec_replicas, + settings.advisory, + ) + return supervisor + + +async def build_live_retune(host: Any) -> LiveRetuneSupervisor: + """The supervisor of a ``NativePlannerBase`` with ``retune.executor: dgdr``. + + ``host.environment.controller`` is the planner's ``KubernetesConnector``; + its ``kube_api`` client and ``parent_dgd_name`` name the DGD. + """ + + config = host.config + retune = config.retune + if retune is None or retune.executor != "dgdr": + raise ValueError( + "build_live_retune requires retune.executor 'dgdr'; executor " + f"{None if retune is None else retune.executor!r} is replay-only" + ) + connector = host.environment.controller + kube = KubernetesRetuneKube.from_kube_api(connector.kube_api) + return await build_supervisor( + config, + kube=kube, + dgd_name=connector.parent_dgd_name, + namespace_fn=lambda: _runtime_namespace(host), + ) + + +async def build_legacy_cleanup( + host: Any, kube: RetuneKube, *, timeout_s: float = 10.0 +) -> LiveRetuneSupervisor: + """A decommission supervisor for the Leases a removed retune left. + + ``NativePlannerBase`` calls it once it found a Lease of the DGD. Run it + like the live supervisor (``start``/``stop``): it rolls back what the + old retune deployed, then deletes the DGD's retune Requests and both + Leases and stops itself (``decommissioned``). Reading the DGD may raise + ``KubeError``. + """ + + connector = host.environment.controller + dgd_name = connector.parent_dgd_name + dgd_obj = await asyncio.to_thread(kube.get_dgd, dgd_name, timeout_s=timeout_s) + logger.info( + "retune legacy clean-up: DGD %s has Leases of a removed live retune; " + "rolling back what it deployed, then deleting its Requests and Leases", + dgd_name, + ) + return build_decommission_supervisor( + host.config, + kube=kube, + dgd=dgd_obj, + dynamo_namespace=_runtime_namespace(host), + ) + + +def build_decommission_supervisor( + config: PlannerConfig, + *, + kube: RetuneKube, + dgd: Mapping[str, Any], + dynamo_namespace: str, + metrics: Optional[RetuneMetricsSink] = None, + actuator_factory: Optional[ActuatorFactory] = None, + frontend: Any = None, + identity: Optional[str] = None, + **supervisor_kwargs: Any, +) -> LiveRetuneSupervisor: + """A supervisor that only rolls back and deletes (no retune config).""" + + if frontend is None: + frontend = frontend_config + if metrics is None: + metrics = RetuneMetrics() + metadata = dgd.get("metadata") or {} + ref = DgdRef( + name=str(metadata.get("name")), + namespace=str(metadata.get("namespace") or kube.namespace), + uid=str(metadata.get("uid")), + dynamo_namespace=dynamo_namespace, + ) + parsed = frontend.parse_frontend(dgd, planner_mode=config.mode) + startup = parsed if isinstance(parsed, FrontendStartup) else None + + def evaluate(obj: dict) -> SpecEvaluation: + return SpecEvaluation( + report=PremiseReport(p4_ok=False, spec_ok=True), + startup=startup, + inputs=None, + spec_replicas=None, + dgd=ref, + ) + + evaluation = evaluate(dict(dgd)) + settings = SupervisorSettings( + interval_s=3600.0, + grace_s=0.0, + timeout_s=3000.0, + advisory=True, + worker_types=worker_types_for(config.mode), + decommission=True, + ) + return LiveRetuneSupervisor( + settings=settings, + kube=kube, + evaluation=evaluation, + identity=identity or planner_identity(), + premise_evaluator=evaluate, + actuator_factory=actuator_factory or default_actuator_factory(kube, metrics), + fingerprint_fn=lambda _: "decommission", + metrics=metrics, + **supervisor_kwargs, + ) diff --git a/components/src/dynamo/planner/retune/live/contracts.py b/components/src/dynamo/planner/retune/live/contracts.py new file mode 100644 index 000000000000..50c743cd65ea --- /dev/null +++ b/components/src/dynamo/planner/retune/live/contracts.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Types shared by the live retune modules. + +Every module in ``dynamo.planner.retune.live`` depends on this one and on +nothing else in the package except through the protocols defined here, so +each can be built and tested with fakes. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any, Literal, Optional, Protocol + +# The eight builtin default-policy parameters a frontend router exposes on +# GET/PATCH /router_policy, by their live (PATCH) names. +POLICY_PARAMETER_KEYS: tuple[str, ...] = ( + "decode_active_request_weight", + "disk_cache_hit_weight", + "host_cache_hit_weight", + "overlap_score_credit", + "overlap_score_credit_decay", + "prefill_load_scale", + "router_temperature", + "shared_cache_multiplier", +) + +# Router knobs a DGDR v1beta2 Request can search, AISimulate public name -> +# live name. ``dynamo.router.simulation.config.ROUTER_PUBLIC_TO_LIVE`` is the +# source of truth; this mirrors it for modules that must not import the +# router simulation stack. +DGDR_ROUTER_KNOBS: dict[str, str] = { + "overlap_score_credit": "overlap_score_credit", + "prefill_load_scale": "prefill_load_scale", + "temperature": "router_temperature", +} + +LABEL_PREFIX = "nvidia.com/planner-retune-" +LABEL_DGD = LABEL_PREFIX + "dgd" +LABEL_DGD_UID = LABEL_PREFIX + "dgd-uid" +LABEL_WINDOW = LABEL_PREFIX + "window" +LABEL_SHARD = LABEL_PREFIX + "shard" +LABEL_SPEC = LABEL_PREFIX + "spec" +LABEL_ATTEMPT = LABEL_PREFIX + "attempt" + +# Pod labels the operator sets (deploy/operator/internal/consts/consts.go). +POD_LABEL_DGD_NAME = "nvidia.com/dynamo-graph-deployment-name" +POD_LABEL_COMPONENT_TYPE = "nvidia.com/dynamo-component-type" + +DGDR_GROUP = "nvidia.com" +DGDR_VERSION = "v1beta2" +DGDR_REQUESTS = "dynamographdeploymentrequests" +DGDR_RUNS = "dynamographdeploymentruns" +DGDR_CANDIDATES = "dynamographdeploymentcandidates" +RUN_UID_LABEL = "nvidia.com/dgdr-run-uid" + +# SweepResult.failure_reason values that mean "this process stopped working +# on the window", not "the window failed": no decision is recorded and the +# controller rewinds so the window can be relaunched or adopted by the next +# leader. +NON_DECISION_FAILURES = frozenset({"shutdown", "lost_leadership"}) + +GateReason = Literal["not_leader", "not_ready", "already_decided"] + + +# --------------------------------------------------------------- controller + + +@dataclass(frozen=True) +class GateResult: + """What ``launch_gate`` tells the controller about a due window. + + ``reason`` is ``None`` to let the window launch (or record ``busy``). + ``not_leader``/``not_ready`` keep ``_last_k`` so the window is retried on + a later tick; ``already_decided`` advances it without a decision. + ``spec_replicas`` (worker type -> DGD spec replicas) and ``abort`` (set + by the supervisor when this process stops leading) are only meaningful + when ``reason`` is ``None``. + """ + + reason: Optional[GateReason] = None + spec_replicas: Optional[dict[str, int]] = None + abort: Optional[threading.Event] = None + + +# (k, (window_start_s, window_end_s, as_of_s)) -> GateResult +LaunchGate = Callable[[int, tuple[float, float, float]], GateResult] +# controller generation of the sweep's controller -> may it be harvested? +HarvestAllowed = Callable[[int], bool] + + +# ------------------------------------------------------------ kubernetes + + +class KubeError(Exception): + """A Kubernetes call failed (network, timeout, 5xx, unexpected status).""" + + +class KubeNotFound(KubeError): + pass + + +class KubeConflict(KubeError): + """409 on update (stale resourceVersion).""" + + +class KubeAlreadyExists(KubeError): + """409 on create (the name is taken).""" + + +class KubeInvalid(KubeError): + """422 / 400: the API server rejected the object (e.g. Strict pruning).""" + + +@dataclass(frozen=True) +class DgdRef: + """The DynamoGraphDeployment this planner serves.""" + + name: str + namespace: str # Kubernetes namespace + uid: str + dynamo_namespace: str # the routers' ``namespace`` on GET /router_policy + + +@dataclass(frozen=True) +class PodInfo: + name: str + uid: str + ip: Optional[str] + ready: bool + terminating: bool + labels: Mapping[str, str] = field(default_factory=dict) + # Container port named ``http`` of the frontend container, if any. + http_port: Optional[int] = None + # ``status.phase`` (Pending, Running, Succeeded, Failed, Unknown). + phase: Optional[str] = None + # ``restartCount`` of the ``main`` container, if reported. A restarted + # frontend registers new routers (new router ids). + restarts: Optional[int] = None + + +class RetuneKube(Protocol): + """Synchronous Kubernetes access for live retune. + + Every method takes an optional ``timeout_s`` (connect+read, applied as + the kubernetes client's ``_request_timeout`` tuple of ints) and raises + the ``Kube*`` errors above instead of client exceptions. Callers on the + event loop wrap calls in ``asyncio.to_thread``. + """ + + namespace: str + + def get_dgd(self, name: str, *, timeout_s: float = 10.0) -> dict: + ... + + def list_pods( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[PodInfo]: + ... + + def create_request( + self, + body: dict, + *, + dry_run: bool = False, + timeout_s: float = 10.0, + ) -> dict: + """Create with ``field_validation='Strict'`` (and ``dry_run='All'``).""" + ... + + def get_request(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + ... + + def list_requests( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + ... + + def delete_request( + self, + name: str, + *, + propagation: Literal["Background", "Foreground"] = "Background", + timeout_s: float = 10.0, + ) -> None: + """Missing objects are not an error.""" + ... + + def get_run(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + ... + + def list_candidates( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + ... + + def get_lease(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + ... + + def create_lease(self, body: dict, *, timeout_s: float = 10.0) -> dict: + ... + + def replace_lease(self, name: str, body: dict, *, timeout_s: float = 10.0) -> dict: + """Update with ``body['metadata']['resourceVersion']``; 409 -> KubeConflict.""" + ... + + def delete_lease( + self, + name: str, + *, + resource_version: Optional[str] = None, + timeout_s: float = 10.0, + ) -> None: + """Missing Leases are not an error. With ``resource_version`` the + delete is conditional on it; a changed Lease -> KubeConflict.""" + ... + + +# ------------------------------------------------------- frontend config + + +@dataclass(frozen=True) +class FrontendStartup: + """The frontend router configuration parsed from the DGD spec. + + ``router_config`` is the startup ``KvRouterConfig`` JSON (only keys the + frontend sets); ``startup_maps`` is ``policy_parameters(wt)`` for each + worker type of the planner mode. ``fingerprint_source`` is a canonical + JSON string of everything parsed that the state fingerprint covers. + """ + + router_mode: str # "kv_router" or "round_robin" + router_config: dict[str, Any] + startup_maps: dict[str, dict[str, float]] + frontend_component: str # DGD component name of the frontend + replicas: int + replica_sync: bool + fingerprint_source: str + + +@dataclass(frozen=True) +class PremiseReport: + """Spec-class premise evaluation (P4, P7-P10; P14 is the actuator's). + + ``p4_ok`` alone decides whether a controller can be built; ``spec_ok`` + is the conjunction of every spec premise. ``reasons`` explain each + failed premise as ``"P: ..."``; ``warnings`` carry warn-only ones. + """ + + p4_ok: bool + spec_ok: bool + reasons: tuple[str, ...] = () + warnings: tuple[str, ...] = () + startup: Optional[FrontendStartup] = None + + +# ---------------------------------------------------------- router policy + + +@dataclass(frozen=True) +class RouterInfo: + """One router on GET /router_policy.""" + + router_id: str + namespace: str + worker_type: str + revision: int + # The 8 PolicyParameters; None for a custom policy without runtime config. + configuration: Optional[dict[str, float]] + model: Optional[str] = None + + +@dataclass(frozen=True) +class PolicyCallResult: + """Outcome of one GET or PATCH. ``status`` is the HTTP status (0 on a + network error or timeout).""" + + status: int + routers: tuple[RouterInfo, ...] = () # GET + router: Optional[RouterInfo] = None # PATCH/reset response, when 200 + error: Optional[str] = None + + +class RouterPolicyClient(Protocol): + async def get(self, pod: PodInfo) -> PolicyCallResult: + ... + + async def patch( + self, + pod: PodInfo, + router_id: str, + expected_revision: int, + parameters: Mapping[str, float], + ) -> PolicyCallResult: + ... + + async def reset( + self, pod: PodInfo, router_id: str, expected_revision: int + ) -> PolicyCallResult: + """PATCH ``{router_id, expected_revision, reset: true}``.""" + ... + + +ActuationMode = Literal["apply", "rollback", "observe"] + + +@dataclass(frozen=True) +class ActuationTarget: + """What one actuator round converges to. + + ``apply``: PATCH explicit values so every router's 8 parameters equal + ``desired_by_wt[worker_type]``. ``rollback``: reset every router whose + configuration differs from its latched initial (never by revision). + ``observe``: no writes; latch initials and check parity only. + ``include_not_ready`` also targets NotReady, non-terminating pods. + ``learn_unknown`` lets the actuator reset a router whose initial was + never latched (revision > 0) once to learn it; that is active mode only. + With False (advisory) such a router is never written: rollback + leaves it alone and does not wait for it to converge. + """ + + mode: ActuationMode + desired_by_wt: Optional[dict[str, dict[str, float]]] = None + include_not_ready: bool = False + learn_unknown: bool = True + + +@dataclass(frozen=True) +class ActuationReport: + converged: bool + # router_id -> configuration latched this round (revision 0 or reset). + new_initials: dict[str, dict[str, float]] = field(default_factory=dict) + # router_ids seen this round, or known and not provably gone; the + # initials of every other router are pruned. + seen_router_ids: frozenset[str] = frozenset() + # P14: every known initial equals startup_maps[worker_type]. + parity_ok: Optional[bool] = None + parity_mismatches: tuple[str, ...] = () + admin_api_disabled: bool = False + errors: dict[str, int] = field(default_factory=dict) # by status code / kind + # P12: enough Ready frontends, each with a configured router for every + # worker type, whatever their values. Gates launches; convergence does not. + routers_present: Optional[bool] = None + + +# ----------------------------------------------------------- reconciler + + +class SweepRegistry(Protocol): + """Thread-safe link between a running sweep and the reconciler. + + Protection is reference-counted per window: every ``protect`` call + takes one reference and needs one matching ``release``. A window, with + the union of every Request name given to it, stays protected while any + reference is held, so two sweeps of one window (a rebuilt controller + relaunching it while the old sweep winds down) cannot drop each + other's protection. + """ + + def protect(self, window_end_ms: int, request_names: Iterable[str]) -> None: + """Mark a window file and Requests as in use by this process (one + reference).""" + ... + + def release(self, window_end_ms: int) -> None: + """Drop one reference taken by ``protect``.""" + ... + + def enqueue_delete(self, request_names: Iterable[str]) -> None: + ... + + +# ---------------------------------------------------------------- window + + +@dataclass(frozen=True) +class WindowCutResult: + path: Optional[str] # None for an empty window (no file written) + num_requests: int + trace_block_size: Optional[int] + row_counts: dict[str, int] + dropped_members: int = 0 + dropped_bytes: int = 0 + reused: bool = False # the final file already existed (EEXIST) + + +class SupportsIsSet(Protocol): + """A cancellation flag: a ``threading.Event`` or any object with ``is_set()``.""" + + def is_set(self) -> bool: + ... + + +class WindowCutter(Protocol): + def __call__( + self, + segments_root: str, + start_ms: int, + end_ms: int, + as_of_ms: int, + out_path: str, + *, + max_rows: int, + deadline_s: float, + cancel: Optional[SupportsIsSet] = None, + ) -> WindowCutResult: + ... + + +# --------------------------------------------------------------- metrics + + +class RetuneMetricsSink(Protocol): + """``dynamo_planner_retune_*`` metrics; every method is cheap and safe.""" + + def decision(self, outcome: str) -> None: + ... + + def failure(self, reason: str) -> None: + ... + + def skipped(self, reason: str) -> None: + ... + + def supervisor_state( + self, + *, + leader: bool, + ready: bool, + initialized: bool, + rollback_pending: bool, + generation: int, + ) -> None: + ... + + def scores( + self, + *, + baseline: Optional[float], + best: Optional[float], + sweep_seconds: Optional[float], + missing_candidates: int, + ) -> None: + ... + + def router_converged(self, converged: bool) -> None: + ... + + def patch_error(self, code: str) -> None: + ... + + def trace_sink_missing(self, pod: str, missing: bool) -> None: + ... + + def dropped_members(self, count: int) -> None: + ... + + +class NullMetrics: + """A ``RetuneMetricsSink`` that records nothing (tests, disabled metrics).""" + + def decision(self, outcome: str) -> None: + pass + + def failure(self, reason: str) -> None: + pass + + def skipped(self, reason: str) -> None: + pass + + def supervisor_state(self, **_: Any) -> None: + pass + + def scores(self, **_: Any) -> None: + pass + + def router_converged(self, converged: bool) -> None: + pass + + def patch_error(self, code: str) -> None: + pass + + def trace_sink_missing(self, pod: str, missing: bool) -> None: + pass + + def dropped_members(self, count: int) -> None: + pass diff --git a/components/src/dynamo/planner/retune/live/dgdr.py b/components/src/dynamo/planner/retune/live/dgdr.py new file mode 100644 index 000000000000..232f58c0c37f --- /dev/null +++ b/components/src/dynamo/planner/retune/live/dgdr.py @@ -0,0 +1,1927 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Live retune sweep executor: one window through DGDR v1beta2 Requests. + +``run_dgdr_sweep`` is the live ``sweep_fn`` of ``RetuneController``. It +runs on the controller's single sweep thread and never raises: + +1. Shard the router grid (``plan_shards``): balanced Cartesian sub-grids of + at most 100 points (a Request's ``maxCandidates`` bound), one + single-point shard for an off-grid baseline or startup point, and a CRN + probe that evaluates one fixed grid point a second time. +2. Adopt the Requests a previous leader created for this window, matched by + ``(dgd-uid, window, shard, spec hash)``; others of the window are queued + for deletion. The spec names the window file, whose name holds the + trace selection (``window_claim_path``), so a Request of another window + start or cutoff with the same end is never adopted. A window whose + every shard is adopted as Succeeded is harvested without cutting it + again. An adopted attempt whose Run failed or ran another input is + replaced by attempt+1 while the shard still fits before the deadline. +3. Cut the window file (a subprocess, ``context.cut_window``), then create + one Request per remaining shard with Strict field validation and check + that the API server kept every submitted field. +4. Wait for each Request's Run, check that ``Run.status.aisimulate.input`` + is what the Request asked for, and read the Run's Candidates. +5. Map each Candidate's router adapter parameters back to live names, + require identical non-router parameters across the window, score each + with ``score_summary`` on ``retune.objective``, and check the CRN probe. + +The executor never deletes anything itself: Requests to remove go to the +``SweepRegistry`` deletion queue, which the reconciler drains while this +process leads. Shutdown (``cancel``) and loss of leadership (``abort``) +return with ``failure_reason`` ``shutdown``/``lost_leadership`` and leave +every Request in place for the next leader to adopt. Shutdown ends the wait +between polls at once; ``abort`` is checked when it ends. Every mutating +call first checks ``abort`` and ``context.is_leader()``. + +A shard that holds the baseline, or the probe, is critical: when it fails +the window fails. Any other failed shard only leaves its points unscored. +""" + +from __future__ import annotations + +import copy +import functools +import hashlib +import itertools +import json +import logging +import math +import os +import re +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Literal, NoReturn, Optional + +from dynamo.planner.config.planner_config import ( + DGDR_MAX_SHARD_POINTS, + DgdrExecutorConfig, + RetuneConfig, + balanced_shard_parts, + dgdr_need_seconds, + largest_shard_points, +) +from dynamo.planner.core.types import RetuneParams +from dynamo.planner.retune.live.contracts import ( + DGDR_GROUP, + DGDR_ROUTER_KNOBS, + DGDR_VERSION, + LABEL_ATTEMPT, + LABEL_DGD, + LABEL_DGD_UID, + LABEL_PREFIX, + LABEL_SHARD, + LABEL_SPEC, + LABEL_WINDOW, + RUN_UID_LABEL, + DgdRef, + KubeAlreadyExists, + KubeError, + KubeInvalid, + KubeNotFound, + RetuneKube, + SweepRegistry, + WindowCutResult, + WindowCutter, +) +from dynamo.planner.retune.live.reconciler import trace_pvc_source +from dynamo.planner.retune.sweep import ( + SweepRequest, + SweepResult, + SweepStatus, + score_summary, +) + +logger = logging.getLogger(__name__) + +ShardKind = Literal["grid", "point", "probe"] + +# Seconds a new Request may go without an Accepted condition before the +# window fails as operator_not_reconciling. +ACCEPT_TIMEOUT_S = 60.0 +# Relative score difference the CRN probe tolerates. +CRN_REL_TOLERANCE = 1e-9 +# Annotation holding the cut window's row counts, so a window adopted as +# Succeeded reports them without cutting the window again. +ROWS_ANNOTATION = LABEL_PREFIX + "rows" +# Operator constants (deploy/operator at #17483). +REQUEST_GENERATION_ANNOTATION = "nvidia.com/dgdr-request-generation" +TRACE_MOUNT_PATH = "/var/run/aisimulate/trace" +ROUTER_ADAPTER = "dynamo.router" +_ROLES = ("aggregated", "prefill", "decode") +_TEMPERATURE = DGDR_ROUTER_KNOBS["temperature"] +# spec.objective.metric -> AISimulate optimization.target (operator lowering). +_METRIC_TARGETS = { + "Throughput": "throughput", + "ThroughputPerGPU": "throughput_per_gpu", + "ThroughputPerUser": "throughput_per_user", + "TTFT": "ttft", + "E2ELatency": "e2e_latency", + "Goodput": "goodput", + "GoodputPerGPU": "goodput_per_gpu", +} +# Statuses a retry cannot fix (bad request, auth, missing API). +_PERMANENT_STATUSES = frozenset({400, 401, 403, 404, 405, 422}) +# How the API server words a rejection by an admission webhook. +_ADMISSION_DENIED = re.compile(r'admission webhook "[^"]*" denied the request') +_K8S_TIMEOUT_S = 10.0 +_SNAP_TOLERANCE = 1e-9 +_MIN_K8S_TIMEOUT_S = 1.0 +_ABSENT = object() + + +# ------------------------------------------------------------------ shards + + +@dataclass(frozen=True) +class Shard: + """One Request's slice of the window's candidates. + + ``choices`` maps each live router knob to the values the shard spans; + ``points`` is their Cartesian product (live names). A ``point`` or + ``probe`` shard has one value per knob. + """ + + id: str + kind: ShardKind + choices: dict[str, tuple[float, ...]] + points: tuple[dict[str, float], ...] + + @property + def size(self) -> int: + return len(self.points) + + def contains(self, point: Mapping[str, float]) -> bool: + return all(point.get(key) in values for key, values in self.choices.items()) + + +def plan_shards( + router_domain: Mapping[str, Sequence[float]], + *, + baseline: Mapping[str, float], + startup: Mapping[str, float], + max_shards: int, + parallelism: int, + probe: bool = True, +) -> list[Shard]: + """Split ``router_domain`` (live knob -> values) into Request shards. + + The grid becomes at most ``max_shards`` Cartesian sub-grids of at most + ``DGDR_MAX_SHARD_POINTS`` points: each knob's sorted values are cut into + contiguous, near-equal chunks by ``balanced_shard_parts``, the split + the configuration check (P13) uses. It minimizes the largest shard's + evaluation rounds at the Run ``parallelism`` (``dgdr_rounds``), then + the shard count, so a shard that would not shorten the window is not + added. + ``baseline`` and ``startup`` (live knob -> value) each get a + single-point shard when off the grid. + With ``probe``, a ``probe`` shard repeats one fixed grid point (the + middle value of each knob, a positive temperature when the grid has + one) to check common random numbers across Runs. Raises ``ValueError`` + when no split fits or a point does not cover the domain's knobs. + """ + + if not router_domain: + raise ValueError("router domain is empty; there is nothing to shard") + if max_shards < 1: + raise ValueError(f"max_shards must be at least 1, got {max_shards}") + keys = tuple(sorted(router_domain)) + axes: dict[str, tuple[float, ...]] = {} + for key in keys: + values = tuple(sorted({float(value) for value in router_domain[key]})) + if not values: + raise ValueError(f"router domain {key!r} has no values") + if not all(math.isfinite(value) for value in values): + raise ValueError(f"router domain {key!r} has a non-finite value") + axes[key] = values + for name, point in (("baseline", baseline), ("startup", startup)): + if set(point) != set(keys): + raise ValueError( + f"{name} router point keys {sorted(point)} do not match the " + f"router domain {list(keys)}" + ) + + sizes = [len(axes[key]) for key in keys] + parts = balanced_shard_parts(sizes, max_shards, parallelism=parallelism) + largest = largest_shard_points(sizes, parts) + if largest > DGDR_MAX_SHARD_POINTS: + raise ValueError( + f"the router grid of {math.prod(sizes)} points cannot be split into " + f"at most {max_shards} Cartesian shards of at most " + f"{DGDR_MAX_SHARD_POINTS} points (the smallest largest shard has " + f"{largest}); raise retune.dgdr.run.max_shards or shrink the domain" + ) + + shards: list[Shard] = [] + chunked = [_chunks(axes[key], part) for key, part in zip(keys, parts)] + for index, combo in enumerate(itertools.product(*chunked)): + choices = dict(zip(keys, combo)) + shards.append( + Shard( + id=f"grid-{index}", + kind="grid", + choices=choices, + points=_product(choices), + ) + ) + + def in_grid(point: Mapping[str, float]) -> bool: + return all(float(point[key]) in axes[key] for key in keys) + + def single(shard_id: str, kind: ShardKind, point: Mapping[str, float]) -> Shard: + values = {key: float(point[key]) for key in keys} + return Shard( + id=shard_id, + kind=kind, + choices={key: (value,) for key, value in values.items()}, + points=(values,), + ) + + baseline_point = {key: float(baseline[key]) for key in keys} + startup_point = {key: float(startup[key]) for key in keys} + if not in_grid(baseline_point): + shards.append(single("point-baseline", "point", baseline_point)) + if not in_grid(startup_point) and startup_point != baseline_point: + shards.append(single("point-startup", "point", startup_point)) + if probe: + shards.append(single("probe", "probe", _probe_point(axes))) + return shards + + +def _chunks(values: tuple[float, ...], count: int) -> list[tuple[float, ...]]: + """``values`` cut into ``count`` contiguous chunks whose sizes differ by <= 1.""" + + base, extra = divmod(len(values), count) + chunks: list[tuple[float, ...]] = [] + start = 0 + for index in range(count): + end = start + base + (1 if index < extra else 0) + chunks.append(values[start:end]) + start = end + return chunks + + +def _product(choices: Mapping[str, tuple[float, ...]]) -> tuple[dict[str, float], ...]: + keys = sorted(choices) + return tuple( + dict(zip(keys, values)) + for values in itertools.product(*(choices[key] for key in keys)) + ) + + +def _probe_point(axes: Mapping[str, tuple[float, ...]]) -> dict[str, float]: + """The fixed CRN probe point: each knob's middle value; a positive + temperature when the grid has one, so the probe exercises sampling.""" + + point: dict[str, float] = {} + for key, values in axes.items(): + pool = values + if key == _TEMPERATURE: + positive = tuple(value for value in values if value > 0) + pool = positive or values + point[key] = pool[len(pool) // 2] + return point + + +# ----------------------------------------------------------------- Request + + +def spec_hash(spec: Mapping[str, Any]) -> str: + """First 16 hex digits of the SHA-256 of the canonical JSON of ``spec``.""" + + canonical = json.dumps(spec, sort_keys=True, separators=(",", ":"), allow_nan=False) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16] + + +def window_selection(window_start_ms: int, window_end_ms: int, as_of_ms: int) -> str: + """First 12 hex digits of the SHA-256 of ``||`` (epoch ms).""" + + key = f"{int(window_start_ms)}|{int(window_end_ms)}|{int(as_of_ms)}" + return hashlib.sha256(key.encode("utf-8")).hexdigest()[:12] + + +def window_claim_path( + cfg: DgdrExecutorConfig, + dgd: DgdRef, + *, + window_start_ms: int, + window_end_ms: int, + as_of_ms: int, +) -> str: + """Claim-relative path of a window file: + ``//-.jsonl.gz``. + + The Request spec names this path, so the selection (``[start, end)`` + as of ``as_of``) is part of the spec hash and the Request name: a + window cut or a Request of another selection with the same end is + never reused. The reconciler reads the window end from the leading + digits. + """ + + selection = window_selection(window_start_ms, window_end_ms, as_of_ms) + return ( + f"{cfg.windows_dir.strip('/')}/{dgd.name}/" + f"{int(window_end_ms)}-{selection}.jsonl.gz" + ) + + +def request_name( + dgd: DgdRef, window_end_ms: int, shard_id: str, digest: str, attempt: int +) -> str: + """``-rt-``.""" + + key = f"{dgd.uid}|{int(window_end_ms)}|{shard_id}|{digest}|{attempt}" + suffix = hashlib.sha256(key.encode("utf-8")).hexdigest()[:10] + prefix = dgd.name[:40].rstrip("-.") or "dgd" + return f"{prefix}-rt-{suffix}" + + +def _label_value(value: str) -> str: + """``value`` when it is a valid label value, else a hashed 63-char form.""" + + if len(value) <= 63 and re.fullmatch( + r"([A-Za-z0-9]([-A-Za-z0-9_.]*[A-Za-z0-9])?)?", value + ): + return value + digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:10] + head = re.sub(r"[^-A-Za-z0-9_.]", "-", value)[:52].strip("-_.") + return f"{head}-{digest}" if head else digest + + +def build_request( + *, + dgd: DgdRef, + cfg: DgdrExecutorConfig, + retune: RetuneConfig, + sla_ttft_ms: float, + sla_itl_ms: float, + window_start_ms: int, + window_end_ms: int, + as_of_ms: int, + shard: Shard, + spec_replicas: Mapping[str, int], + backend: str, + attempt: int, +) -> dict: + """The DGDR v1beta2 Request that replays one shard of a window. + + The trace is the window file of ``[window_start_ms, window_end_ms)`` as + of ``as_of_ms`` (``window_claim_path``). The engine is ``cfg.engine`` + with every leaf a scalar and each used role's + ``parallelism.preset[0].replicas`` set from ``spec_replicas`` + (unused roles are deleted with ``null``); the GPU budget is the sum of + replicas times GPUs per worker (tensor x pipeline x attention_data). + Router knobs use their AISimulate public names: ``{choices: [...]}`` + for a grid shard, a scalar for a point or probe shard. ``cfg.engine`` + is the one the configuration check (P5) accepted. Raises + ``ValueError`` for ``spec_replicas`` or a shard the Request cannot + express. + """ + + engine = _pinned_engine(cfg.engine, spec_replicas, backend) + size = shard.size + if not 1 <= size <= DGDR_MAX_SHARD_POINTS: + raise ValueError( + f"shard {shard.id} has {size} points; " + f"1..{DGDR_MAX_SHARD_POINTS} fit a Request" + ) + spec: dict[str, Any] = { + "model": {"name": cfg.model}, + "backends": [backend], + "hardware": { + "gpu": {"sku": cfg.sku, "budget": _gpu_budget(engine, spec_replicas)} + }, + # The trace format is left to its CRD default, Dynamo. + "workload": { + "trace": { + "source": { + "persistentVolumeClaim": { + "claimName": cfg.claim_name, + "subPath": window_claim_path( + cfg, + dgd, + window_start_ms=window_start_ms, + window_end_ms=window_end_ms, + as_of_ms=as_of_ms, + ), + } + }, + } + }, + "objective": { + "mode": "Optimize", + "metric": retune.objective, + "sla": {"ttftMs": float(sla_ttft_ms), "itlMs": float(sla_itl_ms)}, + }, + "recommendation": {"maxCandidates": size}, + "aisimulate": { + "inputOverrides": { + "engine": engine, + "router": _router_section(shard), + "optimizer": { + "algorithm": "random", + "max_trials": 2 * size, + "parallelism": cfg.run.parallelism, + "candidate_timeout_seconds": cfg.run.candidate_timeout_seconds, + }, + "optimization": {"strict_sla": False}, + # The planner reads only each Candidate's parameters and + # score, never its DGD. The direct generator keeps retune + # Runs independent of the AIC generator, whose rendering can + # reject a Candidate (ai-dynamo/dynamo#17581: no explicit + # block size) and so fail the Run and the window. + "dgdr_run": {"generator": "direct"}, + } + }, + } + digest = spec_hash(spec) + return { + "apiVersion": f"{DGDR_GROUP}/{DGDR_VERSION}", + "kind": "DynamoGraphDeploymentRequest", + "metadata": { + "name": request_name(dgd, window_end_ms, shard.id, digest, attempt), + "namespace": dgd.namespace, + "labels": { + LABEL_DGD: _label_value(dgd.name), + LABEL_DGD_UID: dgd.uid, + LABEL_WINDOW: str(int(window_end_ms)), + LABEL_SHARD: shard.id, + LABEL_SPEC: digest, + LABEL_ATTEMPT: str(attempt), + }, + "ownerReferences": [ + { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "name": dgd.name, + "uid": dgd.uid, + "blockOwnerDeletion": False, + } + ], + }, + "spec": spec, + } + + +def _pinned_engine( + engine: Mapping[str, Any], spec_replicas: Mapping[str, int], backend: str +) -> dict[str, Any]: + """``engine`` with the deployed backend and the DGD spec replicas. + + The configuration check (P5) has already pinned ``engine``: its backend, + if written, is ``backend``, and each worker role of the planner mode has + one parallelism preset entry of positive integers. ``spec_replicas`` + comes from the live DGD, so its roles and counts are checked here. + """ + + roles = set(spec_replicas) + if roles not in ({"aggregated"}, {"prefill", "decode"}): + raise ValueError( + "spec_replicas must name 'aggregated', or 'prefill' and 'decode'; " + f"got {sorted(roles)}" + ) + pinned = copy.deepcopy(dict(engine)) + pinned["backend"] = backend + workers = pinned["workers"] + for role in _ROLES: + if role not in roles: + workers[role] = None + continue + replicas = int(spec_replicas[role]) + if replicas < 1: + raise ValueError( + f"{role} has {replicas} spec replicas; at least 1 is needed" + ) + workers[role]["parallelism"]["preset"][0]["replicas"] = replicas + return pinned + + +def _gpu_budget(engine: Mapping[str, Any], spec_replicas: Mapping[str, int]) -> int: + budget = 0 + for role, replicas in spec_replicas.items(): + entry = engine["workers"][role]["parallelism"]["preset"][0] + per_worker = entry["tensor"] * entry["pipeline"] * entry["attention_data"] + budget += int(replicas) * per_worker + return budget + + +def _router_section(shard: Shard) -> dict[str, Any]: + live_knobs = set(DGDR_ROUTER_KNOBS.values()) + if set(shard.choices) != live_knobs: + raise ValueError( + f"shard {shard.id} router knobs {sorted(shard.choices)} must be exactly " + f"{sorted(live_knobs)}, the knobs a DGDR Request can express" + ) + section: dict[str, Any] = { + "policy": "kv_router", + "prefill_load_model": {"type": "none"}, + } + for public, live in DGDR_ROUTER_KNOBS.items(): + values = [float(value) for value in shard.choices[live]] + section[public] = {"choices": values} if shard.kind == "grid" else values[0] + return section + + +# --------------------------------------------------------- input read-back + + +def _get(obj: Any, path: Sequence[str], default: Any = None) -> Any: + for key in path: + if not isinstance(obj, Mapping) or key not in obj: + return default + obj = obj[key] + return obj + + +def _same_json(left: Any, right: Any) -> bool: + """JSON equality of leaves: numbers by value (2 == 2.0), booleans never + equal numbers.""" + + if isinstance(left, bool) or isinstance(right, bool): + return isinstance(left, bool) and isinstance(right, bool) and left == right + if isinstance(left, (int, float)) and isinstance(right, (int, float)): + return float(left) == float(right) + return type(left) is type(right) and left == right + + +def missing_fields(stored: Any, submitted: Any, path: str = "") -> list[str]: + """Submitted leaves the stored object lacks or changed (pruned or rejected). + + ``stored`` may hold more (admission defaults); a submitted ``null`` only + requires the stored value to be absent or ``null``. + """ + + prefix = (path,) if path else () + return [ + ".".join(where) or "." + for where, _, _ in _differences(stored, submitted, prefix) + ] + + +def _differences( + stored: Any, submitted: Any, path: tuple[str, ...] +) -> list[tuple[tuple[str, ...], Any, Any]]: + """``(path, submitted, stored)`` of every submitted value ``stored`` lacks. + + A mapping may hold more keys than submitted, at any depth, including + the mappings of a list of the submitted length. A submitted ``null`` + requires the stored value to be absent or ``null``; ``_ABSENT`` stands + for an absent stored value. + """ + + if submitted is None: + if stored is None or stored is _ABSENT: + return [] + return [(path, None, stored)] + if isinstance(submitted, Mapping): + if not isinstance(stored, Mapping): + return [(path, submitted, stored)] + found: list[tuple[tuple[str, ...], Any, Any]] = [] + for key, value in submitted.items(): + found += _differences(stored.get(key, _ABSENT), value, (*path, str(key))) + return found + if ( + isinstance(submitted, list) + and isinstance(stored, list) + and len(stored) == len(submitted) + ): + found = [] + for index, (left, right) in enumerate(zip(stored, submitted)): + found += _differences(left, right, (*path, str(index))) + return found + return [] if _same_json(stored, submitted) else [(path, submitted, stored)] + + +def _overlay(target: Any, patch: Any) -> Any: + """``patch`` merged over ``target`` like a merge patch, but a ``null`` in + ``patch`` stays (as "must be unset") instead of deleting the key.""" + + if not isinstance(patch, Mapping): + return copy.deepcopy(patch) + result = dict(target) if isinstance(target, Mapping) else {} + for key, value in patch.items(): + result[key] = _overlay(result.get(key), value) + return result + + +def input_mismatches( + spec: Mapping[str, Any], effective: Mapping[str, Any] +) -> list[str]: + """How ``Run.status.aisimulate.input`` differs from what ``spec`` asks for. + + The expected input is what the operator lowers ``spec``'s typed fields + to (model, backends, GPU SKU and budget, objective, SLA, the trace file + on its read-only mount, ``maxCandidates``) with ``inputOverrides`` + merged over them. ``inputOverrides`` wins: nothing restores the typed + trace source after it. The check is the create path's subset check + (``missing_fields``), so the input may carry more, such as a new + operator default, but must hold every expected value; an expected + ``null`` (a worker role the overrides delete, an SLA or a replay + speedup the Request leaves unset) must be absent or ``null``. + """ + + _, sub_path = trace_pvc_source(spec) + mounted = TRACE_MOUNT_PATH + ( + "/trace.jsonl.gz" if (sub_path or "").endswith(".gz") else "/trace.jsonl" + ) + backends = list(spec.get("backends") or []) + sla = _get(spec, ("objective", "sla"), {}) or {} + typed = { + "engine": { + "model": _get(spec, ("model", "name")), + "backend": {"choices": backends} if backends else None, + "hardware": _get(spec, ("hardware", "gpu", "sku")), + }, + "optimization": { + "target": _METRIC_TARGETS.get(_get(spec, ("objective", "metric"), "")), + "constraints": { + "max_candidate_gpus": _get(spec, ("hardware", "gpu", "budget")) + }, + }, + "evaluation": { + "sla": { + "ttft_ms": sla.get("ttftMs"), + "itl_ms": sla.get("itlMs"), + "e2e_ms": sla.get("e2eMs"), + } + }, + "traffic": { + "source": {"type": "trace", "format": "dynamo", "paths": [mounted]}, + "load": { + "type": "trace_timestamps", + "speedup": _get( + spec, ("workload", "trace", "replay", "arrivalSpeedupRatio") + ), + }, + }, + "dgdr_run": {"max_candidates": _get(spec, ("recommendation", "maxCandidates"))}, + } + overrides = _get(spec, ("aisimulate", "inputOverrides"), {}) or {} + # Normalize like the API server (JSON numbers, no tuples). + expected = json.loads(json.dumps(_overlay(typed, overrides))) + actual = json.loads(json.dumps(effective)) + return [ + f"{'.'.join(path)}: expected " + f"{'unset' if want is None else repr(want)}, got " + f"{'unset' if got is _ABSENT else repr(got)}" + for path, want, got in _differences(actual, expected, ()) + ] + + +# ---------------------------------------------------------------- executor + + +@dataclass(frozen=True) +class DgdrContext: + """Everything ``run_dgdr_sweep`` needs besides the ``SweepRequest``. + + ``backend`` is the deployed backend (``vllm``, ``sglang``, ``trtllm``). + ``is_leader`` is checked before every mutating call. ``now`` is the + wall clock (``deadline_s`` is epoch seconds). Between polls the sweep + waits on ``cancel``, so shutdown ends the wait at once; ``sleep``, when + set, replaces that wait (a test's simulated clock). + """ + + dgd: DgdRef + cfg: DgdrExecutorConfig + retune: RetuneConfig + backend: str + kube: RetuneKube + registry: SweepRegistry + cut_window: WindowCutter + is_leader: Callable[[], bool] + now: Callable[[], float] = time.time + sleep: Optional[Callable[[float], None]] = None + poll_interval_s: float = 5.0 + + +def run_dgdr_sweep( + request: SweepRequest, + *, + cancel: threading.Event, + abort: Optional[threading.Event], + context: DgdrContext, +) -> SweepResult: + """Sweep one window through DGDR v1beta2 Requests; never raises. + + Scores are keyed by ``candidate.key()`` for every candidate of + ``request`` (live router names, no planner part); a candidate no Run + published scores ``None``. See the module docstring for the protocol. + """ + + started = context.now() + try: + return _DgdrSweep(request, cancel, abort, context, started).run() + except Exception as error: # defensive: the executor must never raise + logger.exception("retune DGDR sweep of window %s failed", request.window_index) + return SweepResult( + status="failed", + wall_s=max(0.0, context.now() - started), + error=f"{type(error).__name__}: {error}", + failure_reason="internal_error", + details={}, + ) + + +class _Interrupted(Exception): + """Shutdown or loss of leadership: return without a decision.""" + + def __init__(self, reason: str) -> None: + super().__init__(reason) + self.reason = reason + + +class _WindowFailed(Exception): + """The window ends without scores (``failed``, ``timeout`` or ``empty``).""" + + def __init__( + self, status: SweepStatus, reason: Optional[str], message: str + ) -> None: + super().__init__(message) + self.status = status + self.reason = reason + self.message = message + + +class _Stop: + """Shutdown (``cancel``) or loss of leadership (``abort``) as the one + ``is_set()`` flag the window cutter polls.""" + + def __init__( + self, cancel: threading.Event, abort: Optional[threading.Event] + ) -> None: + self._events = tuple(event for event in (cancel, abort) if event is not None) + + def is_set(self) -> bool: + return any(event.is_set() for event in self._events) + + +ShardState = Literal["create", "unknown", "wait", "done", "missing"] + + +@dataclass +class _ShardRun: + """Progress of one shard's Request within this sweep.""" + + shard: Shard + critical: bool + spec: dict + digest: str + attempt: int = 1 + state: ShardState = "create" + name: Optional[str] = None + adopted: bool = False + # This sweep sent a create for the current attempt that may have landed; + # False for an attempt a previous leader created. + created: bool = False + reason: Optional[str] = None + message: Optional[str] = None + # Why the shard waits: the Request's Accepted reason and message while + # it has no Run (RunNameConflict, RunCreating, CancelingSupersededRun), + # then the open Run's Completed reason and message (JobRunning, + # ContainerWaiting, TraceSourceMissing); None once the Run ends. + run_status: Optional[str] = None + run_name: Optional[str] = None + run_uid: Optional[str] = None + checked_run_uid: Optional[str] = None + unaccepted_since: Optional[float] = None + pending_since: Optional[float] = None + run_pending: bool = False + rows: Optional[dict] = None + published: int = 0 + + def require_name(self) -> str: + """The current attempt's Request name. + + It is set once a create is sent or a Request is adopted, so a shard in + the ``unknown`` or ``wait`` state, where it is read, always has it. + """ + + if self.name is None: + raise RuntimeError( + f"shard {self.shard.id} has no Request name in state {self.state}" + ) + return self.name + + def status(self) -> Optional[str]: + """Why the open Request or Run waits, else why the shard ended or + its last attempt was replaced.""" + + return self.run_status or self.message + + def summary(self) -> dict[str, Any]: + return { + "kind": self.shard.kind, + "points": self.shard.size, + "critical": self.critical, + "request": self.name, + "attempt": self.attempt, + "adopted": self.adopted, + "run": self.run_name, + "state": self.state, + "reason": self.reason, + "message": self.status(), + "published": self.published, + "run_pending": self.run_pending, + } + + +@dataclass +class _Scored: + score: Optional[float] + entry: dict[str, Any] = field(default_factory=dict) + + +class _DgdrSweep: + def __init__( + self, + request: SweepRequest, + cancel: threading.Event, + abort: Optional[threading.Event], + context: DgdrContext, + started: float, + ) -> None: + self.request = request + self.cancel = cancel + self.abort = abort + self.stop = _Stop(cancel, abort) + self.ctx = context + self.cfg = context.cfg + self.started = started + self.deadline = ( + float(request.deadline_s) + if request.deadline_s is not None + else started + request.timeout_s + ) + # Rounded once: the window file's path, the Request specs and the + # cut must agree on the selection. + self.window_start_ms = int(round(request.window_start_ms)) + self.window_end_ms = int(round(request.window_end_ms)) + self.as_of_ms = int(round(request.as_of_ms)) + # Registry references this sweep took; each is released at the end. + self._protections = 0 + self.claim_path = window_claim_path( + self.cfg, + context.dgd, + window_start_ms=self.window_start_ms, + window_end_ms=self.window_end_ms, + as_of_ms=self.as_of_ms, + ) + self.local_path = os.path.join(self.cfg.mount_path, self.claim_path) + self.shards: dict[str, _ShardRun] = {} + self.cut_done = False + self.num_requests = 0 + self.row_counts: dict[str, int] = {} + self.details: dict[str, Any] = { + "window": {"path": self.claim_path, "cut": False}, + "deadline_s": self.deadline, + "generation": request.generation, + "replicas_verified": request.replicas_verified, + } + self.k8s_errors = 0 + self.runtime: dict[str, Any] = {} + + # ------------------------------------------------------------ driver + + def run(self) -> SweepResult: + try: + self._prepare() + self._check_interrupt() + self._adopt() + self._create_all() + self._poll() + return self._harvest() + except _Interrupted as stop: + logger.info( + "retune window %s sweep stopped (%s); its Requests stay for the " + "next leader", + self.request.window_index, + stop.reason, + ) + return self._result( + "failed", + failure_reason=stop.reason, + error=f"sweep stopped: {stop.reason}", + ) + except _WindowFailed as failed: + self._abandon_in_flight() + log = logger.info if failed.status == "empty" else logger.warning + log( + "retune window %s %s (%s): %s", + self.request.window_index, + failed.status, + failed.reason, + failed.message, + ) + return self._result( + failed.status, + failure_reason=failed.reason, + error=failed.message if failed.status != "empty" else None, + ) + finally: + # One release per protect: the registry counts references per + # window, and another generation's sweep may hold its own. + try: + for _ in range(self._protections): + self.ctx.registry.release(self.window_end_ms) + except Exception: # pragma: no cover - registry is in-process + logger.exception("retune registry release failed") + self._protections = 0 + + def _protect(self, names: list[str]) -> None: + self.ctx.registry.protect(self.window_end_ms, names) + self._protections += 1 + + def _result( + self, + status: SweepStatus, + *, + failure_reason: Optional[str] = None, + error: Optional[str] = None, + scores: Optional[dict[str, Optional[float]]] = None, + metrics: Optional[dict[str, dict[str, Any]]] = None, + ) -> SweepResult: + details = dict(self.details) + details["shards"] = {key: run.summary() for key, run in self.shards.items()} + details["k8s_errors"] = self.k8s_errors + details.update(self.runtime) + return SweepResult( + status=status, + num_requests=self.num_requests, + row_counts=dict(self.row_counts), + scores=scores or {}, + metrics=metrics or {}, + wall_s=max(0.0, self._now() - self.started), + error=error, + failure_reason=failure_reason, + details=details, + ) + + # ----------------------------------------------------------- helpers + + def _now(self) -> float: + return self.ctx.now() + + def _timeout(self) -> float: + return max(_MIN_K8S_TIMEOUT_S, min(_K8S_TIMEOUT_S, self.deadline - self._now())) + + def _pause(self, seconds: float) -> None: + """Wait between polls. Shutdown ends the wait at once; loss of + leadership is seen when it ends, and every mutating call checks + ``is_leader()`` first anyway.""" + + if self.ctx.sleep is not None: + self.ctx.sleep(seconds) + else: + self.cancel.wait(seconds) + self._check_interrupt() + + def _check_interrupt(self) -> None: + if self.cancel.is_set(): + raise _Interrupted("shutdown") + if self.abort is not None and self.abort.is_set(): + raise _Interrupted("lost_leadership") + + def _check_mutation(self) -> None: + self._check_interrupt() + if not self.ctx.is_leader(): + raise _Interrupted("lost_leadership") + + def _fail( + self, reason: Optional[str], message: str, status: SweepStatus = "failed" + ) -> NoReturn: + raise _WindowFailed(status, reason, message) + + def _retrying(self, what: str, call: Callable[[], Any]) -> Any: + """``call`` until it succeeds; a permanent error or the deadline fails.""" + + while True: + self._check_interrupt() + try: + return call() + except KubeError as error: + self.k8s_errors += 1 + if _permanent(error): + self._fail("k8s_error", f"{what}: {error}") + if self._now() + self.ctx.poll_interval_s >= self.deadline: + self._fail("k8s_error", f"{what}: {error}", "timeout") + logger.warning("retune %s failed, retrying: %s", what, error) + self._pause(self.ctx.poll_interval_s) + + def _read_failed(self, what: str, error: KubeError) -> None: + """A GET failed: a permanent error fails the window; else poll again.""" + + self.k8s_errors += 1 + if _permanent(error): + self._fail("k8s_error", f"{what}: {error}") + logger.warning("retune %s failed: %s", what, error) + + def _enqueue_delete(self, names: Iterable[Optional[str]]) -> None: + present = [name for name in names if name] + if present: + self.ctx.registry.enqueue_delete(present) + + def _abandon_in_flight(self) -> None: + """Queue this window's unfinished Requests for deletion.""" + + self._enqueue_delete( + run.name + for run in self.shards.values() + if run.state in ("wait", "unknown") and run.name + ) + + # ------------------------------------------------------------- setup + + def _prepare(self) -> None: + request = self.request + domain = request.router_domain + if not domain: + self._fail("invalid_request", "the sweep request has no router domain") + if any(candidate.planner for candidate in request.candidates): + self._fail( + "invalid_request", + "a DGDR sweep searches router parameters only (Scenario 1)", + ) + keys = set(domain) + startup = self._startup_point(domain) + try: + shards = plan_shards( + domain, + baseline=request.baseline.router, + startup=startup, + max_shards=self.cfg.run.max_shards, + parallelism=self.cfg.run.parallelism, + ) + except (KeyError, ValueError) as error: + self._fail("invalid_domain", str(error)) + replicas = request.spec_replicas or request.replicas + self.details["replicas"] = dict(replicas) + probe = next((shard for shard in shards if shard.kind == "probe"), None) + for shard in shards: + try: + body = build_request( + dgd=self.ctx.dgd, + cfg=self.cfg, + retune=self.ctx.retune, + sla_ttft_ms=request.sla_ttft_ms, + sla_itl_ms=request.sla_itl_ms, + window_start_ms=self.window_start_ms, + window_end_ms=self.window_end_ms, + as_of_ms=self.as_of_ms, + shard=shard, + spec_replicas=replicas, + backend=self.ctx.backend, + attempt=1, + ) + except (KeyError, TypeError, ValueError) as error: + self._fail("invalid_engine", f"shard {shard.id}: {error}") + # The baseline must be scored, and the probe's grid twin is what + # the probe is compared with. + critical = ( + shard.kind == "probe" + or shard.contains(request.baseline.router) + or (probe is not None and shard.contains(probe.points[0])) + ) + self.shards[shard.id] = _ShardRun( + shard=shard, + critical=critical, + spec=body["spec"], + digest=body["metadata"]["labels"][LABEL_SPEC], + ) + uncovered = [ + candidate.router + for candidate in request.candidates + if set(candidate.router) != keys + or not any( + run.shard.kind != "probe" and run.shard.contains(candidate.router) + for run in self.shards.values() + ) + ] + if uncovered: + self.details["uncovered_candidates"] = uncovered[:10] + + def _startup_point(self, domain: Mapping[str, Sequence[float]]) -> dict[str, float]: + """The off-grid candidate other than the baseline, else the baseline. + + The controller's candidates are the grid plus the startup point and + the baseline, so an off-grid candidate is one of those two. + """ + + baseline = self.request.baseline + for candidate in self.request.candidates: + router = candidate.router + if candidate == baseline or set(router) != set(domain): + continue + if not all(float(router[key]) in domain[key] for key in domain): + return dict(router) + return dict(baseline.router) + + def _build(self, run: _ShardRun) -> dict: + body = build_request( + dgd=self.ctx.dgd, + cfg=self.cfg, + retune=self.ctx.retune, + sla_ttft_ms=self.request.sla_ttft_ms, + sla_itl_ms=self.request.sla_itl_ms, + window_start_ms=self.window_start_ms, + window_end_ms=self.window_end_ms, + as_of_ms=self.as_of_ms, + shard=run.shard, + spec_replicas=self.details["replicas"], + backend=self.ctx.backend, + attempt=run.attempt, + ) + rows = {"num_requests": self.num_requests, "row_counts": self.row_counts} + body["metadata"]["annotations"] = { + ROWS_ANNOTATION: json.dumps(rows, sort_keys=True, separators=(",", ":")) + } + return body + + # ------------------------------------------------------------- adopt + + def _adopt(self) -> None: + """Adopt the latest attempt of each shard; queue the rest for deletion. + + Each adopted attempt is observed once here: a Succeeded one needs no + cut, and a failed one is replaced while there is time (``_run_failed``). + """ + + selector = ( + f"{LABEL_DGD_UID}={self.ctx.dgd.uid},{LABEL_WINDOW}={self.window_end_ms}" + ) + existing = self._retrying( + "list window Requests", + lambda: self.ctx.kube.list_requests(selector, timeout_s=self._timeout()), + ) + stale: list[Optional[str]] = [] + latest: dict[str, dict] = {} + for obj in existing: + metadata = obj.get("metadata") or {} + labels = metadata.get("labels") or {} + name = metadata.get("name") + run = self.shards.get(labels.get(LABEL_SHARD, "")) + if run is None or labels.get(LABEL_SPEC) != run.digest: + stale.append(name) + continue + attempt = _int(labels.get(LABEL_ATTEMPT), 0) + if metadata.get("deletionTimestamp"): + # Being deleted: never adopt it, and never reuse its attempt. + run.attempt = max(run.attempt, attempt + 1) + continue + current = latest.get(run.shard.id) + if current is None or attempt > _int( + current["metadata"]["labels"].get(LABEL_ATTEMPT), 0 + ): + if current is not None: + stale.append(current["metadata"]["name"]) + latest[run.shard.id] = obj + else: + stale.append(name) + self._enqueue_delete(stale) + for shard_id, obj in latest.items(): + run = self.shards[shard_id] + metadata = obj["metadata"] + run.name = metadata["name"] + run.attempt = max( + run.attempt, _int(metadata["labels"].get(LABEL_ATTEMPT), 1) + ) + run.adopted = True + run.state = "wait" + run.rows = _rows(metadata.get("annotations")) + logger.info( + "retune window %s adopts Request %s for shard %s", + self.request.window_index, + run.name, + shard_id, + ) + if latest: + self._protect([obj["metadata"]["name"] for obj in latest.values()]) + # Classify the adopted Requests: a window whose shards all + # succeeded is harvested without cutting it again. + for run in self.shards.values(): + if run.state == "wait": + self._observe(run) + + # ------------------------------------------------------------ create + + def _create_all(self) -> None: + pending = [run for run in self._ordered() if run.state == "create"] + if not pending: + rows = next( + (run.rows for run in self.shards.values() if run.rows is not None), + None, + ) + if rows is not None: + self.num_requests = _int(rows.get("num_requests"), 0) + self.row_counts = dict(rows.get("row_counts") or {}) + return + # Gate every shard before the cut, so a critical shard that cannot + # finish fails the window without spending the cut budget. + for run in pending: + self._too_late(run) + for run in pending: + if run.state == "create": + self._create(run) + + def _ordered(self) -> list[_ShardRun]: + """Critical shards first, so they get the time that is left.""" + + return sorted(self.shards.values(), key=lambda run: not run.critical) + + def _too_late(self, run: _ShardRun) -> bool: + """Apply the launch gate; True when the shard was given up.""" + + need = dgdr_need_seconds(run.shard.size, self.cfg, needs_cut=not self.cut_done) + left = self.deadline - self._now() + if left >= need: + return False + message = ( + f"shard {run.shard.id} needs {need:.0f}s but {max(left, 0.0):.0f}s are " + "left before the deadline" + ) + if run.critical: + self._fail("launch_too_late", message, "timeout") + run.state = "missing" + run.reason = "launch_too_late" + run.message = message + logger.warning("retune window %s: %s", self.request.window_index, message) + return True + + def _create(self, run: _ShardRun) -> None: + self._check_mutation() + if self._too_late(run): + return + if not self.cut_done: + self._cut() + if self._too_late(run): + return + self._check_mutation() + body = self._build(run) + name = body["metadata"]["name"] + run.name = name + run.rows = None + self._protect([name]) + try: + created = self.ctx.kube.create_request(body, timeout_s=self._timeout()) + except KubeAlreadyExists: + self._adopt_existing(run) + return + except KubeError as error: + self.k8s_errors += 1 + if _admission_denied(error): + self._fail("request_rejected", f"Request {name} rejected: {error}") + if isinstance(error, KubeInvalid): + self._fail("schema_unsupported", f"Request {name} rejected: {error}") + if _permanent(error): + self._fail("k8s_error", f"create Request {name}: {error}") + # The create may have landed; look it up on the next poll. + run.created = True + run.state = "unknown" + run.message = f"create: {error}" + logger.warning("retune create of Request %s failed: %s", name, error) + return + run.created = True + dropped = missing_fields(created.get("spec"), body["spec"], "spec") + labels = (created.get("metadata") or {}).get("labels") or {} + dropped += [ + f"metadata.labels.{key}" + for key, value in body["metadata"]["labels"].items() + if labels.get(key) != value + ] + if dropped: + self._enqueue_delete([name]) + run.state = "missing" + self._fail( + "schema_unsupported", + f"the API server dropped or changed {sorted(dropped)} of Request {name}", + ) + run.state = "wait" + run.unaccepted_since = self._now() + logger.info( + "retune window %s created Request %s (shard %s, %d points, attempt %d)", + self.request.window_index, + name, + run.shard.id, + run.shard.size, + run.attempt, + ) + + def _adopt_existing(self, run: _ShardRun) -> None: + """The name exists (an earlier create landed): adopt it if it is ours.""" + + existing = self._retrying( + f"get Request {run.name}", + lambda: self.ctx.kube.get_request( + run.require_name(), timeout_s=self._timeout() + ), + ) + if existing is None: + # Deleted in between; create it again on the next poll. + run.state = "create" + return + labels = (existing.get("metadata") or {}).get("labels") or {} + if ( + labels.get(LABEL_SPEC) == run.digest + and labels.get(LABEL_DGD_UID) == self.ctx.dgd.uid + and labels.get(LABEL_SHARD) == run.shard.id + ): + run.state = "wait" + run.adopted = True + run.unaccepted_since = self._now() + return + self._shard_failed( + run, "request_conflict", f"Request {run.name} exists with other labels" + ) + + # ------------------------------------------------------------ window + + def _cut(self) -> None: + """Cut the window file; an empty window ends the sweep as ``empty``.""" + + self._check_interrupt() + segments_root = ( + self.request.trace_paths[0] + if self.request.trace_paths + else os.path.join( + self.cfg.mount_path, self.cfg.segments_dir, self.ctx.dgd.name + ) + ) + budget = min(self.cfg.cut_budget_seconds, self.deadline - self._now()) + if budget <= 0: + self._fail("launch_too_late", "no time left to cut the window", "timeout") + # WindowCutter takes the deadline as absolute wall-clock seconds. + cut_deadline_s = self._now() + budget + try: + os.makedirs(os.path.dirname(self.local_path), exist_ok=True) + result: WindowCutResult = self.ctx.cut_window( + segments_root, + self.window_start_ms, + self.window_end_ms, + self.as_of_ms, + self.local_path, + max_rows=self.cfg.max_window_rows, + deadline_s=cut_deadline_s, + cancel=self.stop, + ) + except Exception as error: + # A cut cancelled by shutdown or abort reports which one it was. + self._check_interrupt() + reason = getattr(error, "reason", None) + if reason not in ( + "window_cut_timeout", + "window_too_large", + "trace_corrupt", + "window_cut_failed", + ): + reason = "window_cut_failed" + self._fail(reason, f"window cut: {type(error).__name__}: {error}") + self.cut_done = True + self.num_requests = int(result.num_requests) + self.row_counts = dict(result.row_counts) + self.details["window"].update( + cut=True, + num_requests=result.num_requests, + reused=result.reused, + dropped_members=result.dropped_members, + dropped_bytes=result.dropped_bytes, + trace_block_size=result.trace_block_size, + ) + if result.path is None or result.num_requests == 0: + self._fail(None, "empty window", "empty") + + # -------------------------------------------------------------- wait + + def _poll(self) -> None: + while True: + self._check_interrupt() + active = [ + run + for run in self._ordered() + if run.state in ("create", "unknown", "wait") + ] + if not active: + return + if self._now() >= self.deadline: + pending = [run.shard.id for run in active if run.run_pending] + last = [ + f"{run.shard.id}: {run.status()}" + for run in sorted(active, key=lambda run: run.shard.id) + if run.status() + ] + self._fail( + "run_pending" if pending else "timeout", + "the deadline passed with shards " + f"{sorted(run.shard.id for run in active)} unfinished" + + (f" (Runs pending: {sorted(pending)})" if pending else "") + + (f"; {'; '.join(last)}" if last else ""), + "timeout", + ) + for run in active: + self._check_interrupt() + if run.state == "create": + self._create(run) + elif run.state == "unknown": + self._resolve_unknown(run) + else: + self._observe(run) + if any( + run.state in ("create", "unknown", "wait") + for run in self.shards.values() + ): + left = self.deadline - self._now() + self._pause(max(0.0, min(self.ctx.poll_interval_s, left))) + + def _resolve_unknown(self, run: _ShardRun) -> None: + """After a create whose outcome is unknown: adopt it or create again.""" + + try: + existing = self.ctx.kube.get_request( + run.require_name(), timeout_s=self._timeout() + ) + except KubeError as error: + self._read_failed(f"get Request {run.name}", error) + return + if existing is None: + run.state = "create" + return + self._adopt_existing(run) + + def _shard_failed(self, run: _ShardRun, reason: str, message: str) -> None: + """A shard ended without scores; its Request (if any) is kept for audit. + + A critical shard fails the window; any other leaves its points unscored. + """ + + run.state = "missing" + run.reason = reason + run.message = message + if run.critical: + self._fail(reason, f"shard {run.shard.id}: {message}") + logger.warning( + "retune window %s shard %s %s: %s; its points are unscored", + self.request.window_index, + run.shard.id, + reason, + message, + ) + + def _observe(self, run: _ShardRun) -> None: + """One look at a shard's Request and Run. + + The Request must be Accepted for its current generation (within + ``ACCEPT_TIMEOUT_S``); its Run's input is read back once per Run + UID; a Succeeded Run makes the shard ``done``, a failed one goes to + ``_run_failed``, and a live one is waited on. + """ + + kube = self.ctx.kube + try: + obj = kube.get_request(run.require_name(), timeout_s=self._timeout()) + except KubeError as error: + self._read_failed(f"get Request {run.name}", error) + return + if obj is None or (obj.get("metadata") or {}).get("deletionTimestamp"): + self._shard_failed(run, "request_deleted", f"Request {run.name} is gone") + return + if run.rows is None: + run.rows = _rows((obj.get("metadata") or {}).get("annotations")) + generation = _int((obj.get("metadata") or {}).get("generation"), 1) + accepted = _condition(obj, "Accepted") + if accepted is None or _int(accepted.get("observedGeneration"), 0) < generation: + if run.unaccepted_since is None: + run.unaccepted_since = self._now() + elif self._now() - run.unaccepted_since >= ACCEPT_TIMEOUT_S: + self._fail( + "operator_not_reconciling", + f"Request {run.name} has no Accepted condition after " + f"{ACCEPT_TIMEOUT_S:.0f}s", + ) + return + run.unaccepted_since = None + if accepted.get("status") != "True": + if accepted.get("reason") == "InvalidSpec": + self._fail( + "invalid_spec", + f"Request {run.name}: {accepted.get('message', '')}", + ) + # RunCreating, RunNameConflict, CancelingSupersededRun: wait, and + # keep why, so a deadline that passes here reports it. + run.run_status = ( + f"Request {run.name} {accepted.get('reason')}: " + f"{accepted.get('message', '')}" + ) + return + try: + found = self._find_run(obj) + except KubeError as error: + self._read_failed(f"get the Run of Request {run.name}", error) + return + if found is None: + run.run_status = None + return # created, not yet visible + run.run_name = found["metadata"]["name"] + run_uid = found["metadata"].get("uid") + status = found.get("status") or {} + completed = _condition(found, "Completed") or {} + terminal = completed.get("status") == "True" + succeeded = terminal and completed.get("reason") == "SearchSucceeded" + outcome = ( + f"Run {run.run_name} {completed.get('reason')}: " + f"{completed.get('message', '')}" + ) + run.run_status = outcome if not terminal and completed.get("reason") else None + if not terminal: + self._track_pending(run, status) + if run.checked_run_uid != run_uid: + effective = _get(status, ("aisimulate", "input")) + if not isinstance(effective, Mapping): + if not terminal: + return # the input is recorded together with jobRef + if not succeeded: + self._run_failed(run, outcome) + return + self._input_mismatch(run, ["status.aisimulate.input is missing"]) + return + if not self._check_input(run, effective): + return + run.checked_run_uid = run_uid + if succeeded: + run.state = "done" + run.run_uid = run_uid + elif terminal: + self._run_failed(run, outcome) + + def _track_pending(self, run: _ShardRun, status: Mapping[str, Any]) -> None: + """Flag a live Run that has had no Job for a long time (``run_pending``). + + The operator writes ``status.aisimulate.input`` and ``jobRef`` in one + status patch once the Job exists; until then the Run has no status, + or only a ``Completed=False`` condition that says why it waits (such + as ``TraceSourceMissing``), so this runs before the input is read + back. + """ + + if status.get("jobRef") or status.get("startTime"): + run.pending_since = None + run.run_pending = False + return + now = self._now() + if run.pending_since is None: + run.pending_since = now + elif now - run.pending_since >= max( + self.cfg.job_startup_seconds, self.ctx.poll_interval_s + ): + if not run.run_pending: + logger.warning( + "retune Run %s has no Job after %.0fs (run_pending): %s", + run.run_name, + now - run.pending_since, + run.run_status or "no status yet", + ) + run.run_pending = True + + def _run_failed(self, run: _ShardRun, outcome: str) -> None: + """A Run ended without success. + + An attempt a previous leader created is replaced while there is time, + since only a live or Succeeded attempt is worth adopting; otherwise + the shard fails. + """ + + if run.adopted and not run.created: + name = run.name + if self._replace_attempt(run, outcome): + logger.warning( + "retune window %s replaces adopted Request %s of shard %s: %s", + self.request.window_index, + name, + run.shard.id, + outcome, + ) + return + self._shard_failed(run, "run_failed", outcome) + + def _find_run(self, obj: Mapping[str, Any]) -> Optional[dict]: + """The Run of the Request's current generation, by the operator's refs.""" + + metadata = obj.get("metadata") or {} + status = obj.get("status") or {} + generation = _int(metadata.get("generation"), 1) + names: list[str] = [] + active = (status.get("activeRunRef") or {}).get("name") + if active: + names.append(active) + names += [ + ref.get("name") + for ref in status.get("recentRunRefs") or () + if ref.get("name") + ] + names.append(f"{metadata.get('name')}-{generation}") + uid = metadata.get("uid") + for name in dict.fromkeys(names): + found = self.ctx.kube.get_run(name, timeout_s=self._timeout()) + if found is None: + continue + run_meta = found.get("metadata") or {} + annotations = run_meta.get("annotations") or {} + owners = run_meta.get("ownerReferences") or () + if annotations.get(REQUEST_GENERATION_ANNOTATION) == str( + generation + ) and any(owner.get("uid") == uid for owner in owners): + return found + return None + + def _check_input(self, run: _ShardRun, effective: Mapping[str, Any]) -> bool: + """Read back the Run's lowered input; False when the shard moved on.""" + + dgdr_run = effective.get("dgdr_run") or {} + image = dgdr_run.get("runtime_image") + version = dgdr_run.get("runtime_version_override") + if image or version: + # For audit: the runtime that replayed the window. + self.runtime = {"runtime_image": image, "runtime_version": version} + problems = input_mismatches(run.spec, effective) + if problems: + self._input_mismatch(run, problems) + return False + return True + + def _input_mismatch(self, run: _ShardRun, problems: list[str]) -> None: + """An adopted attempt is replaced while there is time; else it fails.""" + + message = f"Run {run.run_name} input differs: {'; '.join(problems[:5])}" + logger.warning("retune window %s: %s", self.request.window_index, message) + if run.adopted and self._replace_attempt(run, message): + return + self._enqueue_delete([run.name]) + run.state = "missing" + if run.critical: + self._fail("input_mismatch", message) + run.reason = "input_mismatch" + run.message = message + + def _replace_attempt(self, run: _ShardRun, message: str) -> bool: + """Queue the attempt for deletion and create attempt+1, if it fits. + + False (nothing changed) when ``D_k - now < need(shard)``. + """ + + need = dgdr_need_seconds(run.shard.size, self.cfg, needs_cut=not self.cut_done) + if self.deadline - self._now() < need: + return False + self._enqueue_delete([run.name]) + run.attempt += 1 + run.state = "create" + run.adopted = False + run.created = False + run.name = None + run.run_name = None + run.run_status = None + run.checked_run_uid = None + run.unaccepted_since = None + run.pending_since = None + run.run_pending = False + run.message = message + return True + + # ----------------------------------------------------------- harvest + + def _harvest(self) -> SweepResult: + request = self.request + domain_keys = set(request.router_domain or {}) + results: dict[str, _Scored] = {} + probe_results: dict[str, _Scored] = {} + engine_signature: Optional[str] = None + signature_source: Optional[str] = None + unmapped: list[str] = [] + for run in self._ordered(): + if run.state != "done": + continue + selector = f"{RUN_UID_LABEL}={run.run_uid}" + candidates = self._retrying( + f"list Candidates of Run {run.run_name}", + functools.partial(self._list_candidates, selector), + ) + run.published = len(candidates) + for candidate in sorted( + candidates, key=lambda obj: (obj.get("metadata") or {}).get("name", "") + ): + name = (candidate.get("metadata") or {}).get("name", "") + parameters = _get(candidate, ("spec", "parameters")) + try: + point, signature = _map_parameters(parameters, run.shard) + except _Unmapped as error: + if error.fatal: + self._fail("engine_not_pinned", f"Candidate {name}: {error}") + unmapped.append(f"{name}: {error}") + continue + if engine_signature is None: + engine_signature, signature_source = signature, name + elif signature != engine_signature: + self._fail( + "engine_not_pinned", + f"Candidates {signature_source} and {name} differ in " + "non-router parameters", + ) + key = RetuneParams(router=point).key() + target = probe_results if run.shard.kind == "probe" else results + if key in target: + continue + metrics = _get(candidate, ("status", "metrics"), {}) or {} + entry: dict[str, Any] = { + "params": RetuneParams(router=point).to_dict(), + "shard": run.shard.id, + "candidate": name, + } + try: + score, info = score_summary( + metrics, request.window_s, request.objective + ) + except ValueError as error: + entry["error"] = str(error) + target[key] = _Scored(None, entry) + continue + entry.update(info) + target[key] = _Scored(score, entry) + if unmapped: + self.details["unmapped_candidates"] = unmapped[:10] + self._check_probe(results, probe_results) + + scores: dict[str, Optional[float]] = {} + metrics_out: dict[str, dict[str, Any]] = {} + missing = 0 + for candidate in request.candidates: + key = candidate.key() + found = results.get(key) + if found is not None and set(candidate.router) == domain_keys: + scores[key] = found.score + metrics_out[key] = found.entry + else: + scores[key] = None + metrics_out[key] = { + "params": candidate.to_dict(), + "error": self._missing_reason(candidate.router), + } + if scores[key] is None: + missing += 1 + self.details["missing_candidates"] = missing + logger.info( + "retune window %s harvested %d of %d candidates from %d Requests", + request.window_index, + len(scores) - missing, + len(scores), + sum(1 for run in self.shards.values() if run.state == "done"), + ) + return self._result("ok", scores=scores, metrics=metrics_out) + + def _list_candidates(self, selector: str) -> list[dict]: + return self.ctx.kube.list_candidates(selector, timeout_s=self._timeout()) + + def _missing_reason(self, point: Mapping[str, float]) -> str: + for run in self.shards.values(): + if run.shard.kind != "probe" and run.shard.contains(point): + if run.state == "missing": + return f"shard {run.shard.id} {run.reason}: {run.message}" + return f"not published by Run {run.run_name} (infeasible or failed)" + return "not covered by any shard" + + def _check_probe( + self, results: Mapping[str, _Scored], probe_results: Mapping[str, _Scored] + ) -> None: + probe = self.shards.get("probe") + if probe is None: + return + point = probe.shard.points[0] + key = RetuneParams(router=point).key() + grid = results.get(key) + repeat = probe_results.get(key) + grid_score = None if grid is None else grid.score + probe_score = None if repeat is None else repeat.score + report: dict[str, Any] = { + "point": point, + "grid_score": grid_score, + "probe_score": probe_score, + } + self.details["crn_probe"] = report + if grid_score is None or probe_score is None: + self._fail( + "crn_mismatch", + f"CRN probe point {point} is unscored (grid {grid_score}, " + f"probe {probe_score})", + ) + scale = max(abs(grid_score), abs(probe_score)) + error = 0.0 if scale == 0 else abs(grid_score - probe_score) / scale + report["relative_error"] = error + if not error <= CRN_REL_TOLERANCE: + self._fail( + "crn_mismatch", + f"CRN probe point {point} scored {probe_score!r}, its grid " + f"evaluation {grid_score!r} (relative error {error:.3g})", + ) + + +class _Unmapped(Exception): + """A Candidate that maps to no point; ``fatal`` when the engine drifted.""" + + def __init__(self, message: str, *, fatal: bool = False) -> None: + super().__init__(message) + self.fatal = fatal + + +def _map_parameters(parameters: Any, shard: Shard) -> tuple[dict[str, float], str]: + """A Candidate's live router point and its non-router signature. + + The signature is the canonical JSON of the parameters without the three + searched knobs; every Candidate of a window must share it. Raises + ``_Unmapped`` for parameters that name no point of ``shard``, fatally + for a router other than the pinned ``kv_router`` without a load model. + """ + + if not isinstance(parameters, Mapping): + raise _Unmapped("spec.parameters is not an object") + adapters = parameters.get("adapters") + if not isinstance(adapters, Mapping): + raise _Unmapped("spec.parameters.adapters is missing") + router = adapters.get(ROUTER_ADAPTER) + if not isinstance(router, Mapping): + raise _Unmapped(f"spec.parameters.adapters[{ROUTER_ADAPTER!r}] is missing") + load_model = router.get("prefill_load_model") + load_type = ( + load_model.get("type") if isinstance(load_model, Mapping) else load_model + ) + if router.get("policy") != "kv_router" or load_type != "none": + raise _Unmapped( + f"router policy {router.get('policy')!r} with prefill load model " + f"{load_type!r}; the Request pinned kv_router and none", + fatal=True, + ) + point: dict[str, float] = {} + for public, live in DGDR_ROUTER_KNOBS.items(): + value = router.get(public) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise _Unmapped(f"router knob {public} is {value!r}") + snapped = _snap(float(value), shard.choices.get(live, ())) + if snapped is None: + raise _Unmapped( + f"router knob {public}={value!r} is not one of the shard's values" + ) + point[live] = snapped + rest = dict(parameters) + rest["adapters"] = { + **adapters, + ROUTER_ADAPTER: { + key: value for key, value in router.items() if key not in DGDR_ROUTER_KNOBS + }, + } + signature = json.dumps(rest, sort_keys=True, separators=(",", ":"), default=str) + return point, signature + + +def _snap(value: float, values: Sequence[float]) -> Optional[float]: + """The domain value ``value`` denotes, allowing for float round-trips.""" + + for candidate in values: + if abs(value - candidate) <= _SNAP_TOLERANCE * max(1.0, abs(candidate)): + return candidate + return None + + +def _condition(obj: Mapping[str, Any], kind: str) -> Optional[dict]: + for condition in _get(obj, ("status", "conditions"), ()) or (): + if isinstance(condition, Mapping) and condition.get("type") == kind: + return dict(condition) + return None + + +def _permanent(error: KubeError) -> bool: + """A Kubernetes error a retry cannot fix (bad request, auth, missing API).""" + + return getattr(error, "status", None) in _PERMANENT_STATUSES or isinstance( + error, (KubeInvalid, KubeNotFound) + ) + + +def _admission_denied(error: KubeError) -> bool: + """A permanent rejection whose message comes from an admission webhook.""" + + return _permanent(error) and _ADMISSION_DENIED.search(str(error)) is not None + + +def _int(value: Any, default: int) -> int: + try: + return int(value) + except (TypeError, ValueError): + return default + + +def _rows(annotations: Any) -> Optional[dict]: + if not isinstance(annotations, Mapping) or ROWS_ANNOTATION not in annotations: + return None + try: + rows = json.loads(annotations[ROWS_ANNOTATION]) + except (TypeError, ValueError): + return None + return rows if isinstance(rows, dict) else None diff --git a/components/src/dynamo/planner/retune/live/frontend_config.py b/components/src/dynamo/planner/retune/live/frontend_config.py new file mode 100644 index 000000000000..fba1b3204ed9 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/frontend_config.py @@ -0,0 +1,2405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Frontend router startup parsed from the DGD spec, and the spec premises. + +Live retune replays each window in DGDR with AISimulate's router defaults +plus the searched knobs, while the live frontends start from whatever their +argv and env say. The two agree only when the frontend leaves every +unsearched routing field at its default (premise P4), so the planner reads +the frontend's startup from the DynamoGraphDeployment spec, never from a +running pod: + +- ``find_frontend`` picks the one standalone frontend component. +- ``parse_frontend`` rebuilds the frontend's effective argv (the operator's + ``command``/``args`` defaulting) and env (operator-injected, graph ``env`` + and the main container's, with Kubernetes ``$(VAR)`` expansion), then + parses them with the frontend's own parser (``build_frontend_parser``), + which reads env defaults from that env alone, so the planner process's + environment is neither read nor modified. A shell-wrapped command, a + ``--router-policy-config`` file, router env the planner cannot read and + an argument whose unreadable ``$(VAR)`` could set a router flag make the + frontend unparseable. +- Env set through ``valueFrom`` is unreadable. A ConfigMap, Secret or + annotation value could also be or create an option ("dangerous"), as + could a ``$(VAR)`` the env leaves undefined that an ``envFrom`` source + may define (its ``prefix`` plus a ConfigMap or Secret key). Any other + ``fieldRef`` (the operator's ``POD_*`` among them), a + ``resourceFieldRef`` and the operator env of ``_operator_inert_env`` + (``DYN_NAMESPACE`` among them) never start with ``-`` or hold + whitespace, so they stay unreadable values but never stand for an option + ("inert"). ``envFrom`` sources are otherwise assumed to set no router + env; ``evaluate_premises`` warns about them, and P14 compares the + routers' policy parameters with the parsed startup. +- ``evaluate_premises`` checks P4 and P7-P10 and warns on P6 and P7; P14 + is the router actuator's. P4 also parses every worker component's router + flags and env the same way (``_worker_router_errors``), each role's pod + template for a component that has only those, and skips a component + with a native ``runtime`` sidecar, whose main runs the engine: a worker + that resolves a router mode advertises its own router configuration, + which replaces the frontend's for its worker set, and one where + dangerous env could stand for an option fails too. ``p4_ok`` alone + decides whether a retune controller can be built. + +Every function is pure (a dict in, a value out, no Kubernetes calls) and +never reads the planner's own environment. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import posixpath +import re +import shlex +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Optional + +from dynamo.common.configuration.groups.kv_router_args import ( + _KV_ROUTER_FIELDS, + KvRouterConfigBase, +) +from dynamo.common.configuration.groups.router_args import ( + RouterConfigBase, + add_worker_router_arguments, +) +from dynamo.common.configuration.utils import ArgEnv +from dynamo.frontend.frontend_args import build_frontend_parser +from dynamo.llm import KvRouterConfig +from dynamo.planner.retune.controller import startup_router_domain +from dynamo.planner.retune.domains import expand_router_domain +from dynamo.planner.retune.live.contracts import FrontendStartup, PremiseReport + +if TYPE_CHECKING: + from dynamo.planner.config.planner_config import PlannerConfig, RetuneConfig + +logger = logging.getLogger(__name__) + +FRONTEND_TYPE = "frontend" +PLANNER_TYPE = "planner" +EPP_TYPE = "epp" +# Component types that run inference workers; v1beta1 ``type`` is optional. +WORKER_TYPES = frozenset({"worker", "prefill", "decode", ""}) + +MAIN_CONTAINER = "main" +# The operator's frontend invocation (component_frontend.go GetBaseContainer). +DEFAULT_FRONTEND_COMMAND = ("python3",) +DEFAULT_FRONTEND_ARGS = ("-m", "dynamo.frontend") +FRONTEND_MODULE = "dynamo.frontend" +# The operator's worker invocation (component_common.go getCommonContainer), +# kept by graph.go's merge under an empty authored command. +DEFAULT_WORKER_COMMAND = ("/bin/sh", "-c") +# consts.go RuntimeContainerName: an init container of this name makes the +# component a native sidecar, whose main runs the engine. +RUNTIME_CONTAINER = "runtime" +# runtime_status.go ResolveComponentRuntimeStatus, in its priority order. +_MODEL_FLAGS = ("--served-model-name", "--model-name", "--model", "--model-path") +# Worker options known to take at least one value, so an unreadable value +# right after one is a value: were it itself an option, the listed one would +# lack its value and the worker would fail at startup. The worker's router +# parser knows only the router flags, not the arity of engine or runtime +# flags (--enforce-eager takes none, --connector a variable number), so only +# these exempt a value. Checked arities: vLLM --model (one) and +# --served-model-name (nargs "+", whose first value this is); +# SGLang --model-path and --served-model-name (one); TensorRT-LLM --model, +# --model-path and --served-model-name (one); the mocker's --model-path and +# --model-name (one); --speculative-config (one JSON value). Before adding +# an option, check its nargs and action in every backend that accepts it. +_WORKER_VALUE_OPTIONS = frozenset({"--speculative-config", *_MODEL_FLAGS}) +# Left-over frontend options known to take exactly one value, in the sglang +# chat processor and in vLLM's. Other left-over options go to the chat +# processor too, where some take none (vLLM's --enable-auto-tool-choice). +_FRONTEND_VALUE_OPTIONS = frozenset( + {"--chat-template", "--reasoning-parser", "--tool-call-parser"} +) + +# Env the operator injects into every component (component_common.go +# getCommonContainer); user env of the same name overrides it. +ENV_PARENT_DGD_NAME = "DYN_PARENT_DGD_K8S_NAME" +ENV_PARENT_DGD_NAMESPACE = "DYN_PARENT_DGD_K8S_NAMESPACE" +_OPERATOR_FIELD_REF_ENV = { + "POD_NAME": "metadata.name", + "POD_NAMESPACE": "metadata.namespace", + "POD_UID": "metadata.uid", +} +# Operator env whose value the planner does not compute but which never +# starts with "-" or holds whitespace: getCommonContainer's Dynamo namespace +# and component type; on a worker, component_worker.go GetBaseContainer's +# system server, health check, NIXL telemetry and forward-pass metric +# settings; on a multinode component, graph.go addMultinodeTopologyEnvVars' +# leader hostname and rank. Container env wins over envFrom, so no envFrom +# source defines them. Names the operator sets only in some configurations +# (NATS_SERVER, ETCD_ENDPOINTS, DYN_DISCOVERY_BACKEND, CONTAINER_NAME, the +# TLS paths) are not modeled. +_OPERATOR_INERT_ENV = ("DYN_COMPONENT", "DYN_NAMESPACE") +_OPERATOR_WORKER_INERT_ENV = ( + "DYN_FORWARDPASS_METRIC_PORT", + "DYN_HEALTH_CHECK_ENABLED", + "DYN_SYSTEM_ENABLED", + "DYN_SYSTEM_PORT", + "DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS", + "NIXL_TELEMETRY_ENABLE", + "NIXL_TELEMETRY_EXPORTER", + "NIXL_TELEMETRY_PROMETHEUS_PORT", +) +_OPERATOR_MULTINODE_INERT_ENV = ("DYNAMO_LEADER_ADDRESS", "DYNAMO_RANK") + +# The frontend's request trace recipe (P10). +ENV_POD_NAME = "DYN_POD_NAME" +ENV_TRACE = "DYN_REQUEST_TRACE" +ENV_TRACE_SINKS = "DYN_REQUEST_TRACE_SINKS" +ENV_TRACE_FORMAT = "DYN_REQUEST_TRACE_FILE_FORMAT" +ENV_TRACE_RECORDS = "DYN_REQUEST_TRACE_RECORDS" +ENV_TRACE_PATH = "DYN_REQUEST_TRACE_FILE_PATH" +# lib/truthy is_truthy; lib/llm/src/request_trace/config.rs parse_file_format. +_TRUTHY = frozenset({"1", "true", "on", "yes"}) +_GZ_FORMATS = frozenset({"jsonl_gz", "jsonl.gz"}) + +ENV_WORKER_SELECTION_POLICY = "DYN_ROUTER_WORKER_SELECTION_POLICY" +ENV_LOAD_AWARE = "DYN_ROUTER_LOAD_AWARE" +# A worker that resolves a router mode (router_args.add_worker_router_arguments) +# advertises a RouterConfig in its model card. +ENV_ROUTER_MODE = "DYN_ROUTER_MODE" +_ROUTER_MODE_DEST = "router_mode" +_ROUTER_ENV_PREFIX = "DYN_ROUTER_" +# Router env the Rust router reads directly (no frontend flag) that tunes +# transport or startup only, never which worker is picked. +_TRANSPORT_ROUTER_ENV = frozenset( + { + "DYN_ROUTER_ACTIVE_SEQUENCE_DIRECT_ZMQ", + "DYN_ROUTER_ACTIVE_SEQUENCE_ZMQ_RCVHWM", + "DYN_ROUTER_MODEL_CARD_WAIT_SECS", + "DYN_ROUTER_ZMQ_ENDPOINTS_PER_SUB", + } +) +# The deprecated overlap_score_weight: KvRouterArgGroup reads the first name +# when it is set and the second otherwise. +_OVERLAP_WEIGHT_ENV = ("DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT", "DYN_OVERLAP_SCORE_WEIGHT") + +# RouterConfig (not KvRouterConfig) fields that change routing or admission: +# session affinity pins requests and the busy thresholds reject them. The +# DGDR replay models neither. +_ROUTER_ADMISSION_FIELDS = ( + "session_affinity_ttl_secs", + "active_decode_blocks_threshold", + "active_prefill_tokens_threshold", + "active_prefill_tokens_threshold_frac", +) +# KvRouterConfig fields serde skips (to_json drops them and from_json +# rejects them), so they are checked on the parsed flags. +_SERDE_SKIPPED_DEFAULTS: dict[str, Any] = { + "router_approximate_cache_policy": "ttl", + "router_prefill_policy": None, + "router_decode_policy": None, +} +# KvRouterConfig JSON keys a frontend may set to a non-default value: they +# do not change which worker the router picks, or another premise covers +# them. +TOLERATED_ROUTER_FIELDS = frozenset( + { + "router_replica_sync", # P8 + "router_event_threads", # indexer thread count + # Hashing of router-derived active-sequence identities. + "router_tracking_hash", + "router_tracking_key_file", + "router_tracking_key_id", + "serve_indexer", # serves this router's indexer to other routers + "enable_session_prefix_index", # tracked, not consumed by routing + "skip_initial_worker_wait", # startup only + } +) +# Fields with their own P4 message. +_LOAD_MODEL = "router_prefill_load_model" +_CONDITIONAL_DISAGG = "conditional_disagg_enabled" + +_TLS_DESTS = ("tls_cert_path", "tls_key_path") +# Parsed destinations whose value the startup depends on. +_RELEVANT_DESTS = frozenset( + { + *_KV_ROUTER_FIELDS, + *_ROUTER_ADMISSION_FIELDS, + *_TLS_DESTS, + "router_mode", + "session_affinity_mode", + "load_aware", + "conditional_disagg_config", + } +) +_SHELLS = frozenset({"sh", "bash", "dash", "zsh", "ash", "ksh"}) +# Shell option letters that take the next argument (-o pipefail, -O extglob), +# and bash's long options that do. +_SHELL_OPERAND_LETTERS = frozenset("oO") +_SHELL_LONG_OPTIONS_WITH_VALUE = frozenset({"--init-file", "--rcfile"}) +# Python options that may precede ``-m``; the second set takes a value. +_PYTHON_FLAGS = frozenset({"-u", "-B", "-E", "-I", "-O", "-OO", "-P", "-q", "-s", "-S"}) +_PYTHON_FLAGS_WITH_VALUE = frozenset({"-X", "-W"}) +# Python's -m, alone or closing a cluster of flags (-um), with its module +# attached or in the next word. +_PYTHON_MODULE_OPTION = re.compile(r"-[A-Za-z]*?m(.*)") +# Marks an argument part that references env the planner cannot read, as +# "\x00": _OPAQUE where the value could be or create an +# option, _INERT where it cannot. +_UNREADABLE = "\x00unreadable" +_OPAQUE = _UNREADABLE + ":" +_INERT = _UNREADABLE + "-inert:" +# A marker, also as argparse quotes it with repr. +_MARKER = re.compile(r"(?:\x00|\\x00)unreadable(?:-inert)?:(.*?)(?:\x00|\\x00|$)") +# The keys a ConfigMap or Secret may hold, so what an envFrom source defines. +_ENV_FROM_KEY = re.compile(r"[-._a-zA-Z0-9]+") + +# Shell syntax the worker script scan reads. +_SHELL_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +_SHELL_ASSIGNMENT = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\+?=") +_SHELL_BLANKS = frozenset(" \t") +_SHELL_OPERATORS = frozenset("\n;&|()<>") +# Reserved words that may precede a command's name. +_SHELL_RESERVED = frozenset( + {"!", "{", "do", "elif", "else", "if", "then", "time", "until", "while"} +) +_SHELL_DECLARATIONS = frozenset({"declare", "export", "local", "readonly", "typeset"}) +# Commands that parse a quoted word again or turn it into env, so that one +# word can still reach the worker's options or env. +_SHELL_REPARSING = frozenset({".", "env", "eval", "export", "source"}) +# Commands that run their first operand as the command. +_SHELL_LAUNCHERS = frozenset({"builtin", "command", "exec", "nohup"}) + + +class FrontendParseError(ValueError): + """The frontend's router configuration cannot be derived from the spec.""" + + +# ------------------------------------------------------------------ API + + +def worker_types(planner_mode: str) -> tuple[str, ...]: + """Router worker types of a planner mode, as the retune controller uses them.""" + + return ("prefill", "decode") if planner_mode == "disagg" else ("aggregated",) + + +def canonical_json(value: Any) -> str: + """Deterministic JSON: sorted keys, no whitespace.""" + + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def find_frontend(dgd: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: + """Return the name and spec of the DGD's one standalone frontend component. + + Reads the v1beta1 ``spec.components`` list (the planner reads every DGD + as v1beta1). Raises ``FrontendParseError`` ("P7: ...") unless there + is exactly one standalone frontend, naming an EPP or ``frontendSidecar`` + layout when that is why. + """ + + components = _components(dgd) + frontends = [c for c in components if _type(c) == FRONTEND_TYPE] + if len(frontends) == 1: + return str(frontends[0].get("name", "")), dict(frontends[0]) + if frontends: + names = ", ".join(sorted(str(c.get("name")) for c in frontends)) + raise FrontendParseError( + f"P7: the DGD has {len(frontends)} frontend components ({names}); " + "live retune needs exactly one" + ) + sidecars = _sidecar_components(components) + if sidecars: + raise FrontendParseError( + f"P7: the frontend runs as a frontendSidecar of {', '.join(sidecars)}; " + "live retune needs a standalone frontend component" + ) + if any(_type(c) == EPP_TYPE for c in components): + raise FrontendParseError( + "P7: the DGD routes through an EPP component and has no standalone " + "frontend component" + ) + raise FrontendParseError("P7: the DGD has no frontend component") + + +def served_model_names( + dgd: Mapping[str, Any], +) -> tuple[dict[str, str], tuple[str, ...]]: + """Served model name of each worker component, and the undetermined ones. + + Uses the operator's ``status.components[].servedModelName`` while + the status is current, else the first of ``--served-model-name``, + ``--model-name``, ``--model`` and ``--model-path`` in the main + container's command and args (the operator's order; of each role's pod + template, for a component with role templates only), after Kubernetes + ``$(VAR)`` expansion against the component's env, skipping values that + reference env the planner cannot read. + """ + + projected = _projected_status(dgd) + models: dict[str, str] = {} + unknown: list[str] = [] + for component in _components(dgd): + if _type(component) not in WORKER_TYPES: + continue + name = str(component.get("name", "")) + model = (projected.get(name) or {}).get("servedModelName") or _flag_model( + _component_model_flags(dgd, component) + ) + if model: + models[name] = str(model) + else: + unknown.append(name) + return models, tuple(sorted(unknown)) + + +def served_model_name(dgd: Mapping[str, Any]) -> Optional[str]: + """The DGD's served model when exactly one name is known, else ``None``.""" + + names = set(served_model_names(dgd)[0].values()) + return next(iter(names)) if len(names) == 1 else None + + +def parse_frontend( + dgd: Mapping[str, Any], *, planner_mode: str +) -> FrontendStartup | str: + """Parse the frontend's startup router configuration from the DGD spec. + + Returns the startup, or a ``"P4: ..."`` / ``"P7: ..."`` reason when it + cannot be derived. ``router_config`` holds only the keys whose serde + value differs from ``KvRouterConfig()``; ``startup_maps`` resolves + ``policy_parameters`` for each worker type of ``planner_mode``. A + frontend that parses but violates other P4 conditions (a non-KV router + mode, a non-default field) still returns its startup; see + ``evaluate_premises``. + """ + + analysis = _analyze(dgd, planner_mode) + if analysis.startup is None: + return analysis.error or "P4: the frontend configuration cannot be parsed" + return analysis.startup + + +def evaluate_premises( + dgd: Mapping[str, Any], *, planner_config: PlannerConfig +) -> PremiseReport: + """Evaluate the spec premises P4 and P7-P10, and the P6 and P7 warnings. + + Parses the frontend as ``parse_frontend`` does (cached, so cheap) and + reports its startup; P4 also covers what a ``FrontendStartup`` cannot + carry (serde-skipped flags, session affinity, unread router env) and + fails for a worker component that advertises its own router + configuration. A malformed DGD never raises; every problem becomes a + reason. Raises ``ValueError`` only when ``planner_config.retune.dgdr`` is + unset. + """ + + retune = planner_config.retune + if retune is None or retune.dgdr is None: + raise ValueError("evaluate_premises requires planner_config.retune.dgdr") + dgdr = retune.dgdr + mode = str(planner_config.mode) + analysis = _analyze(dgd, mode) + startup = analysis.startup + components = _components(dgd) + p4: list[str] = [] + others: list[str] = [] + warnings: list[str] = [] + + # P4: a parsed KV router whose unsearched fields are the defaults and + # whose startup passes the retune controller's constructor checks. + if startup is None: + p4.append(analysis.error or "P4: the frontend configuration cannot be parsed") + else: + p4.extend(analysis.p4_violations) + p4.extend(_startup_errors(startup, retune, mode)) + # A worker's advertised router configuration replaces the frontend's for + # its worker set (watcher.rs effective_router_config), whatever the + # frontend's flags say. + p4.extend(_worker_router_errors(dgd, components, warnings)) + + # P7: one served model, one standalone frontend, no EPP, no TLS. + models, unknown = served_model_names(dgd) + distinct = sorted(set(models.values())) + if len(distinct) > 1: + others.append( + f"P7: the DGD serves {len(distinct)} models ({', '.join(distinct)}); " + "live retune needs one" + ) + if unknown: + warnings.append( + f"P7: cannot determine the served model of {', '.join(unknown)}; " + "live retune needs the DGD to serve one model" + ) + # Every DGDR Run simulates retune.dgdr.model; an alias given with + # --served-model-name can hide the HF id, so a miss only warns. + names = _worker_model_names(dgd, components) + if names and dgdr.model not in names: + warnings.append( + f"P7: retune.dgdr.model {dgdr.model!r} is not among the workers' " + f"--model/--model-path/--served-model-name values " + f"({', '.join(sorted(names))}); every DGDR Run simulates " + "retune.dgdr.model, so make sure it is the model the DGD serves" + ) + if analysis.component is None: + if analysis.error is not None and analysis.error.startswith("P7"): + others.append(analysis.error) + else: + if any(_type(c) == EPP_TYPE for c in components): + others.append("P7: the DGD routes through an EPP component (GAIE)") + sidecars = _sidecar_components(components) + if sidecars: + others.append(f"P7: {', '.join(sidecars)} run a frontendSidecar") + if analysis.tls: + others.append( + "P7: the frontend serves TLS (--tls-cert-path / DYN_TLS_CERT_PATH)" + ) + + if analysis.env is not None: + for source in analysis.env.env_from: + if _env_from_may_set_router_env(source): + warnings.append( + f"P4: the frontend container uses envFrom {_describe_env_from(source)}; " + "the planner cannot read it and assumes it sets no router env " + "(P14 still checks the routers' policy parameters)" + ) + + # P8: one frontend replica, or synchronized router replicas. + if startup is not None and startup.replicas > 1 and not startup.replica_sync: + others.append( + f"P8: the frontend has {startup.replicas} replicas without " + "router_replica_sync; AISimulate models a single router" + ) + + # P9: worker replicas only change through the DGD spec, and one planner. + for component in components: + if _type(component) in WORKER_TYPES and _scaling_adapter(component): + others.append( + f"P9: worker component {component.get('name')} has a " + "scalingAdapter; its replicas could change during a window" + ) + others.extend(_planner_replica_errors(components, warnings)) + + # P10: the frontend records request_end traces onto the claim. + if analysis.component is not None and analysis.env is not None: + others.extend( + _trace_recipe_errors( + analysis, + dgd_name=_dgd_name(dgd), + claim_name=dgdr.claim_name, + segments_dir=dgdr.segments_dir, + warnings=warnings, + ) + ) + + # P6 (warn only): AISimulate resolves the model by its HF id. + warning = _model_warning(dgdr.model) + if warning is not None: + warnings.append(warning) + + reasons = tuple(_unique([*p4, *others])) + return PremiseReport( + p4_ok=not p4, + spec_ok=not reasons, + reasons=reasons, + warnings=tuple(_unique(warnings)), + startup=startup, + ) + + +# --------------------------------------------------------- frontend facts + + +@dataclass(frozen=True) +class _Env: + """A component main container's env as far as the planner can know it. + + ``entries`` is the merged raw env keyed by name (operator-injected, then + graph ``env``, then the container's; later wins). ``values`` holds each + name whose value the planner can compute, after ``$(VAR)`` expansion; + ``opaque`` the names it cannot (``valueFrom``, or a reference to + unreadable env), and ``inert`` those of them whose value can never be or + create an option. + """ + + entries: Mapping[str, Mapping[str, Any]] + values: Mapping[str, str] + opaque: frozenset[str] + inert: frozenset[str] + env_from: tuple[Mapping[str, Any], ...] + + @property + def env_from_prefixes(self) -> tuple[str, ...]: + return _env_from_prefixes(self.env_from) + + def expand(self, value: str) -> str: + return _expand( + value, + self.values, + self.opaque, + inert=self.inert, + env_from_prefixes=self.env_from_prefixes, + ) + + +@dataclass(frozen=True) +class _Analysis: + component: Optional[str] = None + spec: Optional[Mapping[str, Any]] = None + env: Optional[_Env] = None + startup: Optional[FrontendStartup] = None + error: Optional[str] = None + # P4 violations a FrontendStartup cannot show. + p4_violations: tuple[str, ...] = () + tls: bool = False + + +@dataclass(frozen=True) +class _Parsed: + """The cacheable outcome of parsing one frontend argv and env.""" + + error: Optional[str] = None + router_mode: str = "" + router_config: Optional[Mapping[str, Any]] = None + startup_maps: Optional[Mapping[str, Mapping[str, float]]] = None + replica_sync: bool = False + p4_violations: tuple[str, ...] = () + tls: bool = False + + +def _analyze(dgd: Mapping[str, Any], planner_mode: str) -> _Analysis: + try: + name, spec = find_frontend(dgd) + except FrontendParseError as error: + return _Analysis(error=str(error)) + env = _component_env(dgd, spec) + try: + argv = _effective_argv(spec) + except FrontendParseError as error: + return _Analysis(component=name, spec=spec, env=env, error=str(error)) + parsed = _parse_cached( + canonical_json( + { + "argv": list(argv), + "values": dict(env.values), + "opaque": sorted(env.opaque), + "inert": sorted(env.inert), + "env_from_prefixes": list(env.env_from_prefixes), + } + ), + worker_types(planner_mode), + ) + if parsed.error is not None: + return _Analysis(component=name, spec=spec, env=env, error=parsed.error) + assert parsed.router_config is not None and parsed.startup_maps is not None + router_config = dict(parsed.router_config) + startup_maps = {wt: dict(params) for wt, params in parsed.startup_maps.items()} + startup = FrontendStartup( + router_mode=parsed.router_mode, + router_config=router_config, + startup_maps=startup_maps, + frontend_component=name, + replicas=_replicas(spec), + replica_sync=parsed.replica_sync, + # Replicas stay out: scaling the frontend is P8's concern and must not + # reset the persisted state. + fingerprint_source=canonical_json( + { + "router_mode": parsed.router_mode, + "router_config": router_config, + "startup_maps": startup_maps, + } + ), + ) + return _Analysis( + component=name, + spec=spec, + env=env, + startup=startup, + p4_violations=parsed.p4_violations, + tls=parsed.tls, + ) + + +@lru_cache(maxsize=32) +def _parse_cached(key: str, wts: tuple[str, ...]) -> _Parsed: + """Parse once per distinct frontend (revalidation re-parses every round).""" + + inputs = json.loads(key) + try: + return _parse( + inputs["argv"], + inputs["values"], + frozenset(inputs["opaque"]), + wts, + inert=frozenset(inputs["inert"]), + env_from_prefixes=tuple(inputs["env_from_prefixes"]), + ) + except FrontendParseError as error: + return _Parsed(error=str(error)) + except Exception as error: # noqa: BLE001 - fail closed, never crash the planner + logger.exception("live retune could not parse the frontend configuration") + return _Parsed( + error=( + "P4: the planner failed to parse the frontend configuration " + f"({type(error).__name__}: {error})" + ) + ) + + +def _parse( + argv: Sequence[str], + values: Mapping[str, str], + opaque: frozenset[str], + wts: tuple[str, ...], + *, + inert: frozenset[str] = frozenset(), + env_from_prefixes: tuple[str, ...] = (), +) -> _Parsed: + args = _frontend_args(argv) + parser, registered = _frontend_parser(values) + + def relevant(name: str) -> bool: + if name in registered: + return registered[name] in _RELEVANT_DESTS + return name in _OVERLAP_WEIGHT_ENV or ( + name.startswith(_ROUTER_ENV_PREFIX) and name not in _TRANSPORT_ROUTER_ENV + ) + + # Router env the planner cannot read makes the startup unknowable. + hidden = sorted(name for name in opaque if relevant(name)) + if hidden: + raise FrontendParseError( + f"P4: frontend router env {', '.join(hidden)} is set through valueFrom " + "(or references such a value, or a $(VAR) an envFrom source may " + "define); the planner cannot read it, so set it with a literal value" + ) + + # An unreadable $(VAR) becomes a marker the parser carries as a value: a + # non-router flag keeps it unconverted, a router flag fails the parse, + # and a dangerous one left over where a flag could stand fails below. + expanded = [ + _expand(arg, values, opaque, inert=inert, env_from_prefixes=env_from_prefixes) + for arg in args + ] + namespace, unknown = parser.parse_known_args(expanded) + hidden_flag = _opaque_flag_position(unknown, value_options=_FRONTEND_VALUE_OPTIONS) + if hidden_flag is not None: + original = dict(zip(expanded, args)).get(hidden_flag, hidden_flag) + raise FrontendParseError( + f"P4: frontend arg {original!r} references env the planner cannot " + "read and is neither attached to an option nor the value of a " + "frontend flag or of one known to take a value " + f"({', '.join(sorted(_FRONTEND_VALUE_OPTIONS))}), so it could set a " + "router flag" + ) + unresolved = sorted( + dest + for dest, value in vars(namespace).items() + if dest in _RELEVANT_DESTS and _UNREADABLE in str(value) + ) + if unresolved: + raise FrontendParseError( + f"P4: frontend args set {', '.join(unresolved)} from env the planner " + "cannot read" + ) + + view = _RouterView.from_cli_args(namespace) + # FrontendConfig.validate: the load-aware preset implies KV routing. + if view.load_aware: + view.router_mode = "kv" + if view.router_policy_config is not None: + raise FrontendParseError( + "P4: the frontend sets --router-policy-config " + f"({view.router_policy_config}); the planner cannot read that file, " + "and live retune needs the builtin default policy" + ) + try: + view.apply_load_aware_preset() + view.apply_conditional_disagg_config() + kwargs = {name: getattr(view, name) for name in _KV_ROUTER_FIELDS} + config = KvRouterConfig(**kwargs) + except (ValueError, TypeError) as error: + raise FrontendParseError( + f"P4: the frontend's router configuration is invalid: {error}" + ) from None + + full = json.loads(config.to_json()) + defaults = json.loads(KvRouterConfig().to_json()) + router_config = { + key: value + for key, value in full.items() + if key not in defaults or defaults[key] != value + } + try: + maps = _policy_maps(canonical_json(router_config), wts) + except ValueError as error: + raise FrontendParseError( + f"P4: cannot resolve the frontend's startup policy parameters: {error}" + ) from None + + violations: list[str] = [] + if kwargs["overlap_score_weight"] is not None: + violations.append( + "P4: the frontend sets the deprecated overlap_score_weight, which " + "overrides prefill_load_scale" + ) + for name, default in _SERDE_SKIPPED_DEFAULTS.items(): + if kwargs[name] != default: + violations.append( + f"P4: the frontend sets {name}={kwargs[name]!r} (default " + f"{default!r}), which live retune cannot model" + ) + if view.load_aware: + violations.append( + f"P4: the frontend enables the load-aware preset (--load-aware / " + f"{ENV_LOAD_AWARE})" + ) + for name in _ROUTER_ADMISSION_FIELDS: + value = getattr(view, name) + if value is not None: + violations.append( + f"P4: the frontend sets {name}={value!r}; live retune cannot " + "model session affinity or busy-worker rejection" + ) + for name in sorted(values): + if ( + not name.startswith(_ROUTER_ENV_PREFIX) + or name in registered + or name in _TRANSPORT_ROUTER_ENV + ): + continue + if name == ENV_WORKER_SELECTION_POLICY: + violations.append( + f"P4: the frontend sets {name}, which selects a worker-selection " + "policy instance instead of the builtin default policy" + ) + else: + violations.append( + f"P4: the frontend sets {name}, a router setting live retune " + "cannot model" + ) + + return _Parsed( + router_mode=_router_mode(str(view.router_mode)), + router_config=router_config, + startup_maps=maps, + replica_sync=bool(full.get("router_replica_sync", False)), + p4_violations=tuple(violations), + tls=any(getattr(view, dest, None) for dest in _TLS_DESTS), + ) + + +class _RouterView(RouterConfigBase, KvRouterConfigBase): + """The router fields of ``FrontendConfig``, without its other checks.""" + + tls_cert_path: Optional[str] = None + tls_key_path: Optional[str] = None + + +@lru_cache(maxsize=256) +def _policy_maps( + router_config_json: str, wts: tuple[str, ...] +) -> dict[str, dict[str, float]]: + config = KvRouterConfig.from_json(router_config_json) + return { + wt: {key: float(value) for key, value in config.policy_parameters(wt).items()} + for wt in wts + } + + +def _router_mode(cli_mode: str) -> str: + """``--router-mode`` in the ``RetuneEngineSpec.router_mode`` spelling.""" + + return "kv_router" if cli_mode == "kv" else cli_mode.replace("-", "_") + + +# ----------------------------------------------------------- argv and env + + +def _effective_argv(spec: Mapping[str, Any]) -> tuple[str, ...]: + """The frontend container's argv after the operator's defaulting. + + graph.go: an authored frontend ``command`` drops the default args; with + none, the command stays ``python3`` and the args default to + ``-m dynamo.frontend`` unless authored. + """ + + if spec.get("roles"): + raise FrontendParseError( + "P4: the frontend component uses roles; live retune needs a " + "single-pod frontend" + ) + main = _main_container(spec) + command = [str(part) for part in main.get("command") or []] + args = [str(part) for part in main.get("args") or []] + if command: + return tuple(command + args) + return (*DEFAULT_FRONTEND_COMMAND, *(args or DEFAULT_FRONTEND_ARGS)) + + +def _frontend_args(argv: Sequence[str]) -> list[str]: + """The arguments after ``python -m dynamo.frontend``; anything else fails.""" + + if not argv: + raise FrontendParseError("P4: the frontend container has an empty command") + program = posixpath.basename(argv[0]) + if program in _SHELLS or (len(argv) > 1 and argv[1] == "-c"): + raise FrontendParseError( + f"P4: the frontend command is shell-wrapped ({' '.join(argv[:2])} ...); " + "write it as command: [python3] and args: [-m, dynamo.frontend, ...]" + ) + if not program.startswith("python"): + raise FrontendParseError( + f"P4: the frontend command starts with {argv[0]!r}, not a python " + f"interpreter running -m {FRONTEND_MODULE}" + ) + index = 1 + while index < len(argv): + token = argv[index] + if token in _PYTHON_FLAGS: + index += 1 + continue + if token in _PYTHON_FLAGS_WITH_VALUE: + index += 2 + continue + if token == "-m" and index + 1 < len(argv): + module, rest = argv[index + 1], index + 2 + elif token.startswith("-m") and len(token) > 2: + module, rest = token[2:], index + 1 + else: + break + if module == FRONTEND_MODULE: + return list(argv[rest:]) + break + raise FrontendParseError( + f"P4: the frontend command {' '.join(argv[:4])!r} does not run " + f"python -m {FRONTEND_MODULE}" + ) + + +def _component_env(dgd: Mapping[str, Any], spec: Mapping[str, Any]) -> _Env: + """Merge a component's env the way the operator renders it. + + The operator injects ``DYN_PARENT_DGD_K8S_*``, ``POD_*`` and the inert + names of ``_operator_inert_env``, except into a native-sidecar main, + which gets no base container; MergeEnvs overlays graph ``env`` and the + main container's ``env`` and sorts the result by name; Kubernetes + expands ``$(VAR)`` against earlier names and every ``envFrom`` key. + """ + + metadata = dgd.get("metadata") or {} + entries: dict[str, Mapping[str, Any]] = {} + operator_inert: set[str] = set() + if not _native_sidecar(spec): + entries[ENV_PARENT_DGD_NAME] = { + "name": ENV_PARENT_DGD_NAME, + "value": str(metadata.get("name", "")), + } + entries[ENV_PARENT_DGD_NAMESPACE] = { + "name": ENV_PARENT_DGD_NAMESPACE, + "value": str(metadata.get("namespace", "")), + } + for name, path in _OPERATOR_FIELD_REF_ENV.items(): + entries[name] = { + "name": name, + "valueFrom": {"fieldRef": {"fieldPath": path}}, + } + operator_inert.update(_operator_inert_env(spec)) + main = _main_container(spec) + for entry in [*_graph_env(dgd), *(main.get("env") or [])]: + if isinstance(entry, Mapping) and entry.get("name"): + entries[str(entry["name"])] = entry + operator_inert.discard(str(entry["name"])) + env_from = tuple( + source for source in main.get("envFrom") or [] if isinstance(source, Mapping) + ) + prefixes = _env_from_prefixes(env_from) + + values: dict[str, str] = {} + opaque: set[str] = set() + inert: set[str] = set() + for name in sorted({*entries, *operator_inert}): + if name in operator_inert: + opaque.add(name) + inert.add(name) + continue + entry = entries[name] + source = entry.get("valueFrom") + if source is not None: + opaque.add(name) + if _inert_source(source): + inert.add(name) + continue + expanded = _expand( + str(entry.get("value") or ""), + values, + opaque, + inert=inert, + env_from_prefixes=prefixes, + ) + if _UNREADABLE not in expanded: + values[name] = expanded + continue + opaque.add(name) + # Literal text around inert values may still make an option. + if _OPAQUE not in expanded and not ( + expanded.startswith("-") or any(char.isspace() for char in expanded) + ): + inert.add(name) + return _Env( + entries=entries, + values=values, + opaque=frozenset(opaque), + inert=frozenset(inert), + env_from=env_from, + ) + + +def _operator_inert_env(spec: Mapping[str, Any]) -> list[str]: + """The ``_OPERATOR_*INERT_ENV`` names the operator sets on a component. + + The worker names need an explicit worker type (component_common.go + ComponentDefaultsFactory gives an untyped component the common base + container only), the multinode ones a ``multinode`` component, whose + ``nodeCount`` defaults to 2. + """ + + names = list(_OPERATOR_INERT_ENV) + if _type(spec) and _type(spec) in WORKER_TYPES: + names.extend(_OPERATOR_WORKER_INERT_ENV) + if isinstance(spec.get("multinode"), Mapping): + names.extend(_OPERATOR_MULTINODE_INERT_ENV) + return names + + +def _inert_source(source: Any) -> bool: + """Whether a ``valueFrom`` value can never be or create an option. + + A fieldRef other than an annotation (a name, namespace, uid, label, IP, + node or service account name) and a resourceFieldRef (a quantity) never + start with ``-`` or hold whitespace; a ConfigMap or Secret key and an + annotation can hold anything. + """ + + if not isinstance(source, Mapping): + return False + if isinstance(source.get("resourceFieldRef"), Mapping): + return True + field_ref = source.get("fieldRef") + if not isinstance(field_ref, Mapping): + return False + return not str(field_ref.get("fieldPath") or "").startswith("metadata.annotations") + + +def _env_from_prefixes(env_from: Iterable[Mapping[str, Any]]) -> tuple[str, ...]: + return tuple(sorted({str(source.get("prefix") or "") for source in env_from})) + + +def _env_from_may_define(name: str, prefixes: Iterable[str]) -> bool: + """Whether an ``envFrom`` source with one of ``prefixes`` could define ``name``.""" + + return any( + name.startswith(prefix) and _ENV_FROM_KEY.fullmatch(name[len(prefix) :]) + for prefix in prefixes + ) + + +def _expand( + value: str, + values: Mapping[str, str], + opaque: Iterable[str], + *, + inert: Iterable[str] = (), + env_from_prefixes: Iterable[str] = (), +) -> str: + """Kubernetes ``$(VAR)`` expansion, marking what the planner cannot read. + + ``$$`` escapes a ``$``. A reference to an ``opaque`` name, or to an + undefined one an ``envFrom`` source may define, is replaced by a marker: + ``_INERT`` for an ``inert`` name, ``_OPAQUE`` otherwise. Any other + undefined reference stays as written. + """ + + hidden, harmless = set(opaque), set(inert) + prefixes = tuple(env_from_prefixes) + out: list[str] = [] + index = 0 + while index < len(value): + if value.startswith("$$", index): + out.append("$") + index += 2 + continue + if value.startswith("$(", index): + close = value.find(")", index + 2) + if close != -1: + name = value[index + 2 : close] + if name in harmless: + out.append(f"{_INERT}$({name})\x00") + elif name in hidden or ( + name not in values and _env_from_may_define(name, prefixes) + ): + out.append(f"{_OPAQUE}$({name})\x00") + else: + out.append(values.get(name, value[index : close + 1])) + index = close + 1 + continue + out.append(value[index]) + index += 1 + return "".join(out) + + +def _opaque_flag_position( + unknown: Sequence[str], *, value_options: frozenset[str] +) -> Optional[str]: + """A left-over argument whose dangerous env could stand for a flag. + + ``unknown`` is what the frontend parser, or a worker's router parser, + leaves over in exec form. The parser consumes the value of every flag it + knows, so a dangerous argument it leaves over sits where a flag could, + unless it is the value of a left-over option, and a left-over option + whose name is unreadable always could. Inert env never stands for a + flag, and ``--option=$(VAR)`` is the option's value whatever it holds. + + A reference expands to exactly one argument, so right after one of + ``value_options``, the left-over options known to take a value, it is + that value: one that starts with ``-`` makes that option's own parser + fail at startup. After any other left-over option it could be parsed as + a flag: the frontend hands its left-over options to its chat processor + and a worker to its engine, and both have flags that take no value + (vLLM's ``--enable-auto-tool-choice``, ``--enforce-eager``) or + optional values (``--connector``). + """ + + previous: Optional[str] = None + for token in unknown: + if token.startswith("-"): + if _UNREADABLE in token.split("=", 1)[0]: + return token + elif token.startswith(_OPAQUE) and previous not in value_options: + return token + previous = token + return None + + +def _env_from_note( + reference: Optional[str], opaque: frozenset[str], prefixes: Sequence[str] +) -> str: + """Says why a reference to env (``$(NAME)``, ``$NAME``) is unreadable. + + Empty unless the name is not opaque env and an envFrom source may + define it. A script variable's reference, which names its source + (``$NAME(=)``), never gets the note. + """ + + name = (reference or "").strip("$()") + if name in opaque or not _env_from_may_define(name, prefixes): + return "" + return " (the container's env leaves it to an envFrom source)" + + +def _opaque_reference(text: str) -> Optional[str]: + """The reference (``$(NAME)``, ``$NAME``) of the first marker in ``text``. + + argparse quotes a rejected value with ``repr``, so the marker may appear + escaped. + """ + + match = _MARKER.search(text) + return None if match is None else match.group(1) + + +def _env_from_may_set_router_env(source: Mapping[str, Any]) -> bool: + """Whether an envFrom source could define a ``DYN_*`` name.""" + + prefix = str(source.get("prefix") or "") + return "DYN_".startswith(prefix) or prefix.startswith("DYN_") + + +def _describe_env_from(source: Mapping[str, Any]) -> str: + for kind in ("configMapRef", "secretRef"): + reference = source.get(kind) + if isinstance(reference, Mapping): + return f"{kind} {reference.get('name')!r}" + return canonical_json(dict(source)) + + +# ---------------------------------------------------- frontend arg parser + + +def _parse_error_message(message: str) -> str: + reference = _opaque_reference(message) + if reference is None: + return f"P4: cannot parse the frontend args: {message.strip()}" + return ( + f"P4: a frontend router arg references {reference}, which the planner " + "cannot read" + ) + + +def _worker_parse_error_message(message: str) -> str: + reference = _opaque_reference(message) + if reference is None: + return f"has router flags the planner cannot parse ({message.strip()})" + return ( + f"has a router flag that references {reference}, which the planner " + "cannot read" + ) + + +class _ArgsParser(argparse.ArgumentParser): + """An ``ArgumentParser`` that raises instead of printing and exiting. + + A value carrying an unreadable marker skips type conversion and the + choices check unless its destination is in ``relevant``, so an + unreadable value of a flag that does not matter here parses, while one + of a flag that does fails with a message naming the reference. + """ + + def __init__( + self, + *, + relevant: frozenset[str] = _RELEVANT_DESTS, + describe: Callable[[str], str] = _parse_error_message, + ) -> None: + super().__init__(add_help=False) + self._relevant = relevant + self._describe = describe + + def error(self, message: str): # type: ignore[override] + raise FrontendParseError(self._describe(message)) + + def exit(self, status: int = 0, message: Optional[str] = None): # type: ignore[override] + raise FrontendParseError( + self._describe(message or f"argument parsing exited with {status}") + ) + + def _opaque_irrelevant(self, action: argparse.Action, value: Any) -> bool: + return ( + isinstance(value, str) + and _UNREADABLE in value + and action.dest not in self._relevant + ) + + def _get_value(self, action, arg_string): # type: ignore[no-untyped-def] + if self._opaque_irrelevant(action, arg_string): + return arg_string + return super()._get_value(action, arg_string) + + def _check_value(self, action, value): # type: ignore[no-untyped-def] + if self._opaque_irrelevant(action, value): + return None + return super()._check_value(action, value) + + +def _frontend_parser( + env: Mapping[str, str], +) -> tuple[argparse.ArgumentParser, dict[str, str]]: + """The frontend's argument parser with every env default read from ``env``. + + ``build_frontend_parser`` registers the frontend's arguments without + reading the planner's env or filesystem and without the frontend's + notices. Returns the parser and the env name -> dest map of everything + it registered. + """ + + try: + return build_frontend_parser(env, _ArgsParser()) + except (ValueError, TypeError, argparse.ArgumentTypeError) as error: + # The frontend itself fails to start on such env. + raise FrontendParseError(f"P4: invalid frontend env: {error}") from None + + +def _worker_router_parser( + env: Mapping[str, str], +) -> tuple[argparse.ArgumentParser, dict[str, str]]: + """A worker's router-advertisement parser with env defaults from ``env``. + + Registers ``router_args.add_worker_router_arguments``, which every + backend (vLLM, SGLang, TensorRT-LLM, mocker) uses, against ``env``. Only + ``router_mode`` is relevant: it decides whether the worker advertises a + router configuration at all. + """ + + arg_env = ArgEnv(environ=dict(env)) + parser = _ArgsParser( + relevant=frozenset({_ROUTER_MODE_DEST}), + describe=_worker_parse_error_message, + ) + try: + add_worker_router_arguments(parser, env=arg_env) + except (ValueError, TypeError, argparse.ArgumentTypeError) as error: + # The worker itself fails to start on such env. + raise FrontendParseError(f"has invalid router env ({error})") from None + return parser, arg_env.env_dests + + +# --------------------------------------------------------- premise checks + + +def _startup_errors( + startup: FrontendStartup, retune: RetuneConfig, planner_mode: str +) -> list[str]: + """P4 on a parsed startup: KV mode, defaults, the controller's checks.""" + + if startup.router_mode != "kv_router": + return [ + f"P4: the frontend routes with {startup.router_mode!r}; live retune " + "needs --router-mode kv" + ] + searched = set(expand_router_domain(retune)) + config = startup.router_config + errors: list[str] = [] + load_model = config.get(_LOAD_MODEL, "none") + if load_model != "none": + errors.append( + f"P4: the frontend's {_LOAD_MODEL} is {load_model!r}; live retune " + "needs 'none'" + ) + if config.get(_CONDITIONAL_DISAGG): + errors.append( + "P4: the frontend enables conditional disaggregation, which replay " + "and live routers resolve differently" + ) + for key in sorted(config): + if ( + key in searched + or key in TOLERATED_ROUTER_FIELDS + or key in (_LOAD_MODEL, _CONDITIONAL_DISAGG) + ): + continue + errors.append( + f"P4: the frontend sets {key}={config[key]!r}; only the searched " + f"router knobs ({', '.join(sorted(searched))}) may differ from " + "their defaults" + ) + if errors: + return errors + # The controller constructor's own checks, so that P4 passes only for a + # startup a controller can be built on. + try: + startup_router_domain(retune, config, worker_types(planner_mode)) + except ValueError as error: + return [f"P4: retune.router does not fit the frontend startup: {error}"] + return [] + + +def _worker_router_errors( + dgd: Mapping[str, Any], + components: Sequence[Mapping[str, Any]], + warnings: list[str], +) -> list[str]: + """P4 on the workers: none advertises its own router configuration. + + A worker that resolves a router mode (``--router-mode`` or + ``DYN_ROUTER_MODE``, from its container or the graph env) puts a + ``RouterConfig`` in its model card, and the frontend routes that worker + set with it instead of its own flags; replay would model the frontend's. + A worker whose mode the planner cannot rule out fails too: one where + dangerous env could stand for an option. In exec form that is a + reference neither attached to its option (``--option=$(VAR)``) nor + after one of ``_WORKER_VALUE_OPTIONS``; in shell form any dangerous + ``$(VAR)``, which Kubernetes substitutes before the shell parses, and + the script references ``_shell_expand`` marks. + + A component with role pod templates and no component ``podTemplate`` + is checked role by role, each role's template as its pods run it. A + component with a native ``runtime`` sidecar is skipped: its main runs + the engine with its image entrypoint, and the sidecar registers no + router configuration. + """ + + errors: list[str] = [] + env_from: list[str] = [] + for component in components: + if _type(component) not in WORKER_TYPES or _native_sidecar(component): + continue + name = str(component.get("name", "")) + roles = [ + (role, spec, _component_env(dgd, spec)) + for role, spec in _role_specs(component) + ] + for role, spec, env in roles: + argv, shell = _worker_argv(spec, env) + detail = _worker_router_cached( + canonical_json( + { + "argv": argv, + "shell": shell, + "values": dict(env.values), + "opaque": sorted(env.opaque), + "env_from_prefixes": list(env.env_from_prefixes), + } + ) + ) + if detail is not None: + label = name if role is None else f"{name} (role {role})" + errors.append( + f"P4: worker component {label} {detail}; a worker's router " + "configuration replaces the frontend's for its workers, so set " + f"--router-mode and {ENV_ROUTER_MODE} on the frontend only (a " + "graph-level env entry reaches every worker)" + ) + break + if any( + _env_from_may_set_router_env(source) + for _, _, env in roles + for source in env.env_from + ): + env_from.append(name) + if env_from: + warnings.append( + f"P4: worker components {', '.join(sorted(env_from))} use envFrom; " + "the planner cannot read it and assumes it sets no router env (a " + f"{ENV_ROUTER_MODE} there would make the worker advertise its own " + "router configuration)" + ) + return errors + + +@lru_cache(maxsize=64) +def _worker_router_cached(key: str) -> Optional[str]: + """Why one worker advertises a router configuration, or ``None``.""" + + inputs = json.loads(key) + try: + return _worker_router_mode( + inputs["argv"], + bool(inputs["shell"]), + inputs["values"], + frozenset(inputs["opaque"]), + tuple(inputs["env_from_prefixes"]), + ) + except FrontendParseError as error: + return str(error) + except Exception as error: # noqa: BLE001 - fail closed, never crash the planner + logger.exception("live retune could not parse a worker's router flags") + return ( + "could not be checked for a router configuration " + f"({type(error).__name__}: {error})" + ) + + +def _worker_router_mode( + argv: Sequence[str], + shell: bool, + values: Mapping[str, str], + opaque: frozenset[str], + env_from_prefixes: Sequence[str], +) -> Optional[str]: + parser, registered = _worker_router_parser(values) + mode_env = sorted( + name for name, dest in registered.items() if dest == _ROUTER_MODE_DEST + ) + hidden = [name for name in mode_env if name in opaque] + if hidden: + return ( + f"sets {', '.join(hidden)} through valueFrom (or a reference to such " + "a value, or a $(VAR) an envFrom source may define), so the planner " + "cannot tell whether it advertises a router configuration" + ) + # A shell script can also export the env inline. + inline = sorted( + {token.split("=", 1)[0] for token in argv if shell and "=" in token} + & set(mode_env) + ) + if inline: + return f"sets {', '.join(inline)} in its shell command" + namespace, remainder = parser.parse_known_args(list(argv)) + mode = getattr(namespace, _ROUTER_MODE_DEST, None) + if mode is not None: + return ( + f"advertises its own router configuration (router mode {mode!r} from " + f"--router-mode or {ENV_ROUTER_MODE})" + ) + if shell: + # Kubernetes substitutes $(VAR) into the script before the shell + # parses it, so a dangerous one can become several words anywhere; + # _shell_expand marks only the shell's own references that can. + opaque_word = next((word for word in argv if _OPAQUE in word), None) + if opaque_word is None: + return None + reference = _opaque_reference(opaque_word) or "" + note = _env_from_note(reference, opaque, env_from_prefixes) + if reference.startswith("$("): + return ( + f"has a shell script that references {reference}, which " + "Kubernetes substitutes into the script before the shell parses " + "it, so its value could become an option wherever it stands, and " + f"the planner cannot read it{note}, so it could select a router " + "mode; give the value literally or use the shell's own reference " + f'attached to its option (--option="${reference[2:-1]}")' + ) + return ( + f"has a shell script that references {_describe_reference(reference)} " + f"where it could become an option, and the planner cannot read it{note}, " + "so it could select a router mode; give the value literally, quote the " + 'reference ("$VAR") outside the worker\'s arguments, or attach it to ' + 'its option (--option="$VAR")' + ) + hidden_flag = _opaque_flag_position(remainder, value_options=_WORKER_VALUE_OPTIONS) + if hidden_flag is None: + return None + reference = _opaque_reference(hidden_flag) or "" + note = _env_from_note(reference, opaque, env_from_prefixes) + if hidden_flag.startswith("-"): + return ( + f"has an option whose name references {reference}, which the planner " + f"cannot read{note}, so it could select a router mode; spell the " + "option's name literally" + ) + return ( + f"has an argument that references {reference}, which the planner cannot " + f"read{note} and which is neither attached to an option nor the value of " + "one known to take a value " + f"({', '.join(sorted(_WORKER_VALUE_OPTIONS))}), so it could select a " + "router mode; give the value literally or attach it to its option " + "(--option=$(VAR))" + ) + + +def _describe_reference(reference: str) -> str: + """A shell reference as a message shows it, with its source when it has one. + + ``_ShellScan`` writes a script variable assigned from dangerous env as + ``$NAME(=)``, since a marker holds no whitespace. + """ + + match = re.fullmatch(r"(\$[A-Za-z_][A-Za-z0-9_]*)\(=(.*)\)", reference) + if match is None: + return reference + return f"{match.group(1)} (assigned from {match.group(2)})" + + +def _worker_argv(component: Mapping[str, Any], env: _Env) -> tuple[list[str], bool]: + """A worker's argv after ``$(VAR)`` expansion, and whether a shell runs it. + + An empty command is the operator's ``/bin/sh -c``. A shell runs the + command when it starts with one, or when a shell's ``-c`` follows a + launcher (``env bash -c``, ``tini -- sh -c``); then its script's own + references are resolved (``_shell_expand``) and every part is split + into words, so the flags of the commands it runs are seen. Unreadable + ``$(VAR)`` references stay marked. A standalone ``--`` is dropped: a + launcher consumes it, while the router parser would read every later + flag as a positional. + """ + + main = _main_container(component) + command = list(main.get("command") or []) or list(DEFAULT_WORKER_COMMAND) + parts = [env.expand(str(part)) for part in [*command, *(main.get("args") or [])]] + script = _shell_script_index(parts) + shell = script is not None or ( + bool(parts) + and (posixpath.basename(parts[0]) in _SHELLS or parts[1:2] == ["-c"]) + ) + words = parts + if shell: + words = [] + for index, part in enumerate(parts): + if index == script: + part = _shell_expand(part, env) + try: + words.extend(shlex.split(part)) + except ValueError: + words.extend(part.split()) + return [word for word in words if word != "--"], shell + + +def _native_sidecar(component: Mapping[str, Any]) -> bool: + """Whether the operator runs the component's Dynamo runtime as a native sidecar. + + v1beta1_helpers.go GetComponentDynamoSidecar: an init container named + ``runtime`` in the pod template, or in the leader role's. + """ + + templates = [component] + for role in component.get("roles") or []: + if isinstance(role, Mapping) and role.get("name") == "leader": + templates.append(role) + return any( + isinstance(container, Mapping) and container.get("name") == RUNTIME_CONTAINER + for template in templates + for container in _pod_spec(template).get("initContainers") or [] + ) + + +def _role_specs( + component: Mapping[str, Any], +) -> list[tuple[Optional[str], Mapping[str, Any]]]: + """The component as each of its pods runs it, with the role's name. + + v1beta1_helpers.go EffectiveComponentForRole: a component ``podTemplate``, + or no role ``podTemplate``, serves every pod (role ``None`` here); + otherwise each role with a ``podTemplate`` runs that template. The + roles stay in each copy, so ``_native_sidecar`` still reads the + leader's. + """ + + roles = [ + role + for role in component.get("roles") or [] + if isinstance(role, Mapping) and isinstance(role.get("podTemplate"), Mapping) + ] + if isinstance(component.get("podTemplate"), Mapping) or not roles: + return [(None, component)] + return [ + (str(role.get("name", "")), {**component, "podTemplate": role["podTemplate"]}) + for role in roles + ] + + +def _shell_script_index(parts: Sequence[str]) -> Optional[int]: + """The index of the script a shell's ``-c`` runs in ``parts``, if any. + + It is ``len(parts)`` when ``parts`` end before the script. + """ + + for index, part in enumerate(parts): + if posixpath.basename(part) in _SHELLS: + script = _shell_options_end(parts, index + 1) + if script is not None: + return script + return None + + +def _shell_options_end(parts: Sequence[str], index: int) -> Optional[int]: + """Where the options a shell reads from ``index`` on end, if they include ``-c``. + + The shell reads ``-`` and ``+`` clusters up to ``--``, ``-`` or a word + that is none; each ``o`` or ``O`` in a cluster takes the next argument + (``-euo pipefail``). With ``-c`` the script is the first word after them. + """ + + command = False + while index < len(parts): + option = parts[index] + if option[:1] not in ("-", "+"): + break + index += 1 + if option in ("-", "--"): + break + if option.startswith("--"): + if option in _SHELL_LONG_OPTIONS_WITH_VALUE: + index += 1 + continue + command |= option.startswith("-") and "c" in option + index += sum(letter in _SHELL_OPERAND_LETTERS for letter in option[1:]) + return index if command else None + + +def _shell_expand(script: str, env: _Env) -> str: + """A worker script with the shell's own parameter references resolved. + + Runs on the text the shell parses, after Kubernetes ``$(VAR)`` + substitution, and keeps it apart from each ``$NAME``, ``${NAME}`` and + ``${NAMEword}`` outside single quotes, comments and heredocs: + + - literal env becomes its value as the shell passes it on (split into + words when unquoted), so the router parser sees an option it holds; + - dangerous env (a ConfigMap, Secret or annotation ``valueFrom``, or a + name the container's env leaves to an ``envFrom`` source, which with + an unprefixed source includes image env such as ``PATH``) becomes + ``_OPAQUE`` where its value could be or create an option of the + worker: unquoted, which the shell splits into words; anywhere in an + argument of ``eval`` or of a shell given ``-c`` (also behind a + launcher), which parse it again; or quoted as a whole argument of the + command that runs the worker (``-m dynamo.``) or of one that turns + it into code or env (source, env, export, a shell, named by path or + behind exec, nohup, command or builtin), unless it follows one of + ``_WORKER_VALUE_OPTIONS``. A redirection target other than a + here-string names a file and never counts. An assignment neither + splits nor passes it on, but leaves its variable dangerous, and a + reference to that variable names the env it was assigned from; + - other unreadable env becomes ``_INERT``; + - script variables, special parameters and image env that no + ``envFrom`` source may define stay as written; + - a default or alternate word is scanned in their place wherever the + shell may use it. + + The scan does not follow control flow, so an assignment may not run: + it leaves literal env and explicit dangerous env of its name, and a + dangerous earlier assignment, in place. Command substitution, backticks + and arithmetic are neither evaluated nor rejected, and neither are + scripts the worker sources or runs, literal text that ``eval`` or a + nested shell parses again, text that a shell or ``xargs`` reads from a + heredoc, a here-string or a pipe, or a script ``su -c`` runs; a script + variable counts as dangerous only when the script assigns it from + dangerous env directly (not through an array, positional parameters, a + function argument or a command's output). + """ + + return _ShellScan(script, env).run() + + +class _ShellScan: + """One pass of ``_shell_expand`` over a script.""" + + def __init__(self, script: str, env: _Env) -> None: + self.script = script + self.env = env + self.out: list[str] = [] + self.index = 0 + self.quoted = False + # The current word as far as the scan knows it: an expansion it does + # not evaluate adds nothing, since it may be empty. + self.word = "" + self.in_word = False + self.assignment = False + self.dangerous = False + # The known words of the current simple command, and its name. + self.words: list[str] = [] + self.command: Optional[str] = None + self.command_dangerous = False + self.redirect = False + self.here_string = False + self.heredocs: list[tuple[str, bool]] = [] + # Script variables, and whether each was assigned dangerous env. + self.variables: dict[str, bool] = {} + # The first dangerous reference in the current word, and the one + # each dangerous script variable was assigned from. + self.source: Optional[str] = None + self.sources: dict[str, str] = {} + + def run(self) -> str: + while self.index < len(self.script): + # A default or alternate word is spliced into the script. + script = self.script + char = script[self.index] + if char == "\x00": + end = script.find("\x00", self.index + 1) + end = len(script) if end == -1 else end + 1 + if script.startswith(_OPAQUE, self.index): + self.dangerous = True + self.source = self.source or _opaque_reference( + script[self.index : end] + ) + self.emit(end, "") + elif not self.quoted and char in _SHELL_BLANKS: + self.end_word() + self.copy(self.index + 1) + elif not self.quoted and char in _SHELL_OPERATORS: + self.operator() + elif not self.quoted and char == "#" and not self.in_word: + end = script.find("\n", self.index) + self.copy(len(script) if end == -1 else end) + elif char == "\\": + self.escape() + elif char == "'" and not self.quoted: + end = script.find("'", self.index + 1) + end = len(script) if end == -1 else end + 1 + self.emit(end, script[self.index + 1 : end - 1]) + elif char == '"': + self.emit(self.index + 1, "") + self.quoted = not self.quoted + elif char == "`": + self.emit(_backtick_end(script, self.index + 1), "") + elif char == "$": + self.dollar() + else: + self.emit(self.index + 1, char) + self.end_word() + self.end_command() + return "".join(self.out) + + def copy(self, end: int) -> None: + self.out.append(self.script[self.index : end]) + self.index = end + + def emit(self, end: int, known: str, text: Optional[str] = None) -> None: + """Add the script up to ``end`` (or ``text`` for it) to the current word.""" + + if not self.in_word: + self.in_word = True + self.assignment = self.command is None and bool( + _SHELL_ASSIGNMENT.match(self.script, self.index) + ) + self.out.append(self.script[self.index : end] if text is None else text) + self.word += known + self.index = end + + def end_word(self) -> None: + if not self.in_word: + return + word = self.word + self.command_dangerous |= self.dangerous + if self.redirect: + self.redirect = self.here_string = False + else: + match = _SHELL_ASSIGNMENT.match(word) + if match and (self.assignment or self.command in _SHELL_DECLARATIONS): + self.assign(match.group(1), self.dangerous) + elif _SHELL_NAME.fullmatch(word) and ( + self.command == "read" + or (self.command == "for" and self.words == ["for"]) + ): + self.assign(word, False) + self.words.append(word) + if ( + self.command is None + and not self.assignment + and word not in _SHELL_RESERVED + ): + self.command = word + self.word, self.in_word, self.assignment, self.dangerous = ( + "", + False, + False, + False, + ) + self.source = None + + def end_command(self) -> None: + if self.command == "for" and len(self.words) > 1 and self.command_dangerous: + self.variables[self.words[1]] = True + self.words, self.command = [], None + self.command_dangerous = self.redirect = self.here_string = False + + def assign(self, name: str, dangerous: bool) -> None: + # The scan does not follow control flow, so the assignment may not + # run: it never clears explicit dangerous env or an earlier + # dangerous assignment. + env = self.env + self.variables[name] = ( + dangerous + or self.variables.get(name, False) + or (name in env.opaque and name not in env.inert) + ) + if dangerous and self.source is not None: + self.sources.setdefault(name, self.source) + + def operator(self) -> None: + script, index = self.script, self.index + self.end_word() + if script.startswith("<<", index) and not script.startswith("<<<", index): + heredoc, end = _shell_heredoc(script, index) + if heredoc is not None: + self.heredocs.append(heredoc) + self.copy(end) + return + if script[index] in "<>" or script.startswith("&>", index): + # A redirection, whose target is never an argument of the + # command, unless a here-string feeds it the target's text. + end = index + 1 + while end < len(script) and script[end] in "<>&|": + end += 1 + self.here_string = script.startswith("<<<", index) + self.copy(end) + self.redirect = True + return + self.end_command() + self.copy(index + 1) + if script[index] == "\n": + for delimiter, strip in self.heredocs: + self.copy(_shell_heredoc_end(script, self.index, delimiter, strip)) + self.heredocs.clear() + + def escape(self) -> None: + index = self.index + following = self.script[index + 1 : index + 2] + if following == "\n": + # A line continuation joins the lines. + self.index += 2 + elif self.quoted and following not in ('"', "\\", "$", "`"): + self.emit(index + 1, "\\") + else: + self.emit(index + 2, following) + + def dollar(self) -> None: + script, index = self.script, self.index + following = script[index + 1 : index + 2] + name = _SHELL_NAME.match(script, index + 1) + if following == "(": + # Command substitution or arithmetic, not evaluated. + self.emit(_shell_substitution_end(script, index + 2), "") + elif following == "{": + self.braced() + elif name is not None: + self.reference(name.group(0), name.end()) + elif following and following in "@*#?$!-0123456789": + self.emit(index + 2, "") + else: + self.emit(index + 1, "$") + + def braced(self) -> None: + script, index = self.script, self.index + close = _shell_brace_end(script, index + 2) + if close == -1: + self.emit(len(script), "") + return + body = script[index + 2 : close] + name = _SHELL_NAME.match(body) + if name is None: + # A length, an indirection or a positional parameter. + self.emit(close + 1, "") + return + operator = body[name.end() :] + word_operator = operator.lstrip(":")[:1] + if word_operator in ("-", "=", "+"): + word = operator[operator.index(word_operator) + 1 :] + kind = self.kind(name.group(0)) + value = self.env.values.get(name.group(0), "") + known_set = kind == "readable" and ( + bool(value) or not operator.startswith(":") + ) + # The shell expands the word where it replaces the value: an + # alternate one where the value may be set, a default one where + # it may be unset and is not unreadable env. + if word_operator == "+": + if known_set or kind != "readable": + self.script = script[:index] + word + script[close + 1 :] + else: + self.index = close + 1 + return + if not known_set and kind not in ("dangerous", "inert"): + self.script = script[:index] + word + script[close + 1 :] + return + self.reference(name.group(0), close + 1, operator) + + def reference(self, name: str, end: int, operator: str = "") -> None: + kind = self.kind(name) + if kind == "readable": + value = self.env.values[name] + # The value, unless a default word replaces an empty one or an + # operator derives another value from it. + if operator[:1] in ("", "-", "=", "?") or ( + operator[:2] in (":-", ":=", ":?") and value + ): + self.literal(value, end) + else: + self.emit(end, "") + elif kind == "unknown": + self.emit(end, "") + else: + dangerous = kind == "dangerous" + self.dangerous |= dangerous + # A marker holds no whitespace: _worker_argv splits the script. + shown = f"${name}" + if self.variables.get(name) and name in self.sources: + shown += f"(={self.sources[name]})" + if dangerous: + self.source = self.source or shown + marker = _OPAQUE if dangerous and self.option_position() else _INERT + self.emit(end, "", f"{marker}{shown}\x00") + + def kind(self, name: str) -> str: + env = self.env + assigned = self.variables.get(name) + if assigned: + return "dangerous" + if name in env.values: + # The script's own assignment may not run either. + return "readable" + if assigned is not None: + return "unknown" + if name in env.inert: + return "inert" + if name in env.opaque or _env_from_may_define(name, env.env_from_prefixes): + return "dangerous" + return "unknown" + + def literal(self, value: str, end: int) -> None: + if self.quoted: + self.emit(end, value, value.replace("\\", "\\\\").replace('"', '\\"')) + elif self.in_word and self.assignment: + self.emit(end, value, shlex.quote(value)) + else: + # The shell splits an unquoted value into words. + pieces = value.split() + for position, piece in enumerate(pieces): + if position or value[:1].isspace(): + self.end_word() + self.out.append(" ") + self.emit(self.index, piece, shlex.quote(piece)) + if value[-1:].isspace(): + self.end_word() + self.out.append(" ") + self.index = end + + def option_position(self) -> bool: + """Whether a dangerous value here could be or create an option of the worker.""" + + if self.in_word and self.assignment: + return False + if self.redirect and not self.here_string: + # A redirection target names a file or descriptor. + return False + if not self.quoted or self.reparses(): + return True + lead = self.word if self.in_word else "" + if lead: + # Attached to its option, or inside one's name. + return lead.startswith("-") and "=" not in lead + if self.words and self.words[-1] in _WORKER_VALUE_OPTIONS: + return False + if self.command is None: + # The command's name: one word, a program to run. + return False + launched = self.launched() + if launched in _SHELL_REPARSING or launched in _SHELLS: + return True + words = self.words + for word, following in zip(words, [*words[1:], ""]): + module = _PYTHON_MODULE_OPTION.fullmatch(word) + if module is not None and (module.group(1) or following).startswith( + "dynamo." + ): + return True + return False + + def launched(self) -> str: + """The current command's name by basename, past ``_SHELL_LAUNCHERS``.""" + + if self.command is None: + return "" + for word in self.words[self.words.index(self.command) :]: + name = posixpath.basename(word) + if not word.startswith("-") and name not in _SHELL_LAUNCHERS: + return name + return "" + + def reparses(self) -> bool: + """Whether the current command parses its arguments again as a script. + + ``eval`` does, also behind ``_SHELL_LAUNCHERS``, and so does a shell + given ``-c``, also behind a launcher (``exec``, ``nohup``, ``env``), + whose positional parameters its script may expand unquoted. + """ + + return self.launched() == "eval" or _shell_script_index(self.words) is not None + + +def _backtick_end(script: str, index: int) -> int: + """The index after the backtick closing one opened before ``index``.""" + + while index < len(script): + if script[index] == "\\": + index += 2 + elif script[index] == "`": + return index + 1 + else: + index += 1 + return len(script) + + +def _shell_brace_end(script: str, index: int) -> int: + """The index of the ``}`` closing a ``${`` whose body starts at ``index``, or -1.""" + + depth = 0 + while index < len(script): + if script[index] == "\\": + index += 2 + continue + if script.startswith("${", index): + depth += 1 + index += 2 + continue + if script[index] == "}": + if depth == 0: + return index + depth -= 1 + index += 1 + return -1 + + +def _shell_substitution_end(script: str, index: int) -> int: + """The index after the ``)`` closing a ``$(`` whose body starts at ``index``.""" + + depth = 0 + quote = "" + heredocs: list[tuple[str, bool]] = [] + word_start = True + while index < len(script): + char = script[index] + if quote == "'": + quote = "" if char == "'" else quote + index += 1 + continue + if char == "\\": + index += 2 + word_start = False + continue + if char == "`": + index = _backtick_end(script, index + 1) + continue + if script.startswith("$(", index): + index = _shell_substitution_end(script, index + 2) + continue + if quote == '"': + quote = "" if char == '"' else quote + index += 1 + continue + if char in "'\"": + quote = char + elif char == "#" and word_start: + end = script.find("\n", index) + index = len(script) if end == -1 else end + continue + elif script.startswith("<<", index) and not script.startswith("<<<", index): + heredoc, index = _shell_heredoc(script, index) + if heredoc is not None: + heredocs.append(heredoc) + continue + elif char == "\n" and heredocs: + index += 1 + for delimiter, strip in heredocs: + index = _shell_heredoc_end(script, index, delimiter, strip) + heredocs.clear() + word_start = True + continue + elif char == "(": + depth += 1 + elif char == ")": + if depth == 0: + return index + 1 + depth -= 1 + word_start = char in _SHELL_BLANKS or char in _SHELL_OPERATORS + index += 1 + return len(script) + + +def _shell_heredoc(script: str, index: int) -> tuple[Optional[tuple[str, bool]], int]: + """The heredoc ``<<`` at ``index`` opens, and the index after its delimiter. + + The heredoc is its delimiter and whether ``<<-`` strips leading tabs. + """ + + index += 2 + strip = script.startswith("-", index) + if strip: + index += 1 + while index < len(script) and script[index] in _SHELL_BLANKS: + index += 1 + start = index + while ( + index < len(script) + and script[index] not in _SHELL_BLANKS + and script[index] not in _SHELL_OPERATORS + ): + index += 1 + delimiter = re.sub(r"[\"'\\]", "", script[start:index]) + return ((delimiter, strip) if delimiter else None), index + + +def _shell_heredoc_end(script: str, index: int, delimiter: str, strip: bool) -> int: + """The index after a heredoc body starting at ``index`` and its delimiter line.""" + + while index < len(script): + end = script.find("\n", index) + end = len(script) if end == -1 else end + line = script[index:end] + index = end + 1 + if (line.lstrip("\t") if strip else line) == delimiter: + break + return min(index, len(script)) + + +def _planner_replica_errors( + components: Sequence[Mapping[str, Any]], warnings: list[str] +) -> list[str]: + planners = [c for c in components if _type(c) == PLANNER_TYPE] + if not planners: + warnings.append( + "P9: the DGD has no planner component; its replicas cannot be checked" + ) + return [ + f"P9: planner component {c.get('name')} has {_replicas(c)} replicas; " + "live retune needs 1" + for c in planners + if _replicas(c) != 1 + ] + + +def _trace_recipe_errors( + analysis: _Analysis, + *, + dgd_name: str, + claim_name: str, + segments_dir: str, + warnings: list[str], +) -> list[str]: + """P10: the frontend writes request_end segments onto the claim.""" + + assert analysis.spec is not None and analysis.env is not None + entries = analysis.env.entries + errors: list[str] = [] + + def literal(name: str) -> Optional[str]: + entry = entries.get(name) + if entry is None: + return None + if entry.get("valueFrom") is not None: + errors.append(f"P10: frontend env {name} must be a literal value") + return None + return str(entry.get("value") or "") + + trace = literal(ENV_TRACE) + if trace is None or trace.strip().lower() not in _TRUTHY: + errors.append(f"P10: the frontend must set {ENV_TRACE}=1") + + sinks = _csv(literal(ENV_TRACE_SINKS)) + if "file" not in sinks or {"jsonl", "jsonl_gz"} & set(sinks): + errors.append(f"P10: the frontend must set {ENV_TRACE_SINKS}=file") + elif set(sinks) != {"file"}: + extra = ", ".join(sorted(set(sinks) - {"file"})) + warnings.append(f"P10: {ENV_TRACE_SINKS} also enables {extra}") + + file_format = literal(ENV_TRACE_FORMAT) + if file_format is not None and file_format.strip().lower() not in _GZ_FORMATS: + errors.append(f"P10: {ENV_TRACE_FORMAT} must be jsonl_gz or unset") + + if set(_csv(literal(ENV_TRACE_RECORDS))) != {"request_end"}: + errors.append(f"P10: the frontend must set {ENV_TRACE_RECORDS}=request_end") + + pod_name = entries.get(ENV_POD_NAME) or {} + field_ref = (pod_name.get("valueFrom") or {}).get("fieldRef") or {} + if field_ref.get("fieldPath") != "metadata.name": + errors.append( + f"P10: the frontend must set {ENV_POD_NAME} from fieldRef metadata.name" + ) + parent = entries.get(ENV_PARENT_DGD_NAME) or {} + if parent.get("valueFrom") is not None or parent.get("value") != dgd_name: + errors.append( + f"P10: the frontend overrides {ENV_PARENT_DGD_NAME}; it must stay the " + f"DGD name {dgd_name!r}" + ) + + mounts, mount_errors = _claim_mounts(analysis.spec, claim_name) + errors.extend(mount_errors) + path = literal(ENV_TRACE_PATH) + if path is None: + errors.append(f"P10: the frontend must set {ENV_TRACE_PATH}") + elif mounts: + prefixes = [ + _join_dir(mount, segments_dir, dgd_dir, f"$({ENV_POD_NAME})") + for mount in mounts + for dgd_dir in (f"$({ENV_PARENT_DGD_NAME})", dgd_name) + ] + # Directly in the pod's directory: the reconciler prunes only the + # segments at its top level. + if not any( + path.startswith(prefix) + and len(path) > len(prefix) + and "/" not in path[len(prefix) :] + for prefix in prefixes + ): + errors.append( + f"P10: {ENV_TRACE_PATH}={path!r} must be a file prefix directly " + f"under {prefixes[0]!r} (no further '/')" + ) + return errors + + +def _claim_mounts( + spec: Mapping[str, Any], claim_name: str +) -> tuple[list[str], list[str]]: + """Writable main-container mount paths of the claim's root, and errors.""" + + pod_spec = _pod_spec(spec) + volumes = { + str(volume.get("name")) + for volume in pod_spec.get("volumes") or [] + if isinstance(volume, Mapping) + and (volume.get("persistentVolumeClaim") or {}).get("claimName") == claim_name + } + if not volumes: + return [], [f"P10: the frontend pod has no volume for claim {claim_name!r}"] + candidates = [ + mount + for mount in _main_container(spec).get("volumeMounts") or [] + if isinstance(mount, Mapping) and mount.get("name") in volumes + ] + if not candidates: + return [], [ + f"P10: the frontend main container does not mount claim {claim_name!r}" + ] + mounts = [ + str(mount["mountPath"]) + for mount in candidates + if mount.get("mountPath") + and not mount.get("readOnly") + and not mount.get("subPath") + and not mount.get("subPathExpr") + ] + if not mounts: + return [], [ + f"P10: the frontend mounts claim {claim_name!r} read-only or through a " + "subPath; mount the claim root writable" + ] + return mounts, [] + + +def _join_dir(mount: str, *parts: str) -> str: + return "/".join([mount.rstrip("/"), *(part for part in parts if part)]) + "/" + + +@lru_cache(maxsize=8) +def _model_warning(model: str) -> Optional[str]: + """P6 for ``model``; cached, so a broken install logs once per model.""" + + try: + # Optional dependency of a warn-only premise: without it, P6 is not + # checked. + from aisimulate_core.sdk.common import DefaultHFModels + except ImportError: + return None + except Exception as error: # noqa: BLE001 - evaluate_premises never raises + logger.exception("live retune could not load AISimulate's default HF models") + return ( + f"P6: cannot check retune.dgdr.model {model!r} against AISimulate's " + f"default HF models ({type(error).__name__}: {error}); each DGDR Run " + "must be able to download it" + ) + if model in DefaultHFModels: + return None + return ( + f"P6: retune.dgdr.model {model!r} is not an AISimulate default HF model; " + "each DGDR Run must be able to download it" + ) + + +# ------------------------------------------------------------ DGD helpers + + +def _components(dgd: Mapping[str, Any]) -> list[Mapping[str, Any]]: + components = (dgd.get("spec") or {}).get("components") + if isinstance(components, list): + return [c for c in components if isinstance(c, Mapping)] + return [] + + +def _graph_env(dgd: Mapping[str, Any]) -> list[Any]: + return list((dgd.get("spec") or {}).get("env") or []) + + +def _pod_spec(spec: Mapping[str, Any]) -> Mapping[str, Any]: + template = spec.get("podTemplate") + pod_spec = template.get("spec") if isinstance(template, Mapping) else None + return pod_spec if isinstance(pod_spec, Mapping) else {} + + +def _main_container(spec: Mapping[str, Any]) -> Mapping[str, Any]: + for container in _pod_spec(spec).get("containers") or []: + if isinstance(container, Mapping) and container.get("name") == MAIN_CONTAINER: + return container + return {} + + +def _sidecar_components(components: Sequence[Mapping[str, Any]]) -> list[str]: + return sorted(str(c.get("name")) for c in components if c.get("frontendSidecar")) + + +def _type(component: Mapping[str, Any]) -> str: + return str(component.get("type") or "") + + +def _replicas(component: Mapping[str, Any]) -> int: + """Spec replicas; an omitted value is the workload default of 1.""" + + replicas = component.get("replicas") + try: + return 1 if replicas is None else int(replicas) + except (TypeError, ValueError): + return 1 + + +def _scaling_adapter(component: Mapping[str, Any]) -> bool: + return component.get("scalingAdapter") is not None + + +def _dgd_name(dgd: Mapping[str, Any]) -> str: + return str((dgd.get("metadata") or {}).get("name", "")) + + +def _projected_status(dgd: Mapping[str, Any]) -> Mapping[str, Any]: + """``status.components`` while the status is current, else empty.""" + + status = dgd.get("status") or {} + generation = (dgd.get("metadata") or {}).get("generation") + current = generation is not None and generation == status.get("observedGeneration") + return (status.get("components") or {}) if current else {} + + +def _model_flags( + component: Mapping[str, Any], env: _Env +) -> list[tuple[str, Optional[str]]]: + """Each model flag in the main container's command and args, in order. + + Values are ``$(VAR)``-expanded against the component's env; a value + that still references env (or looks like a flag) is ``None``. + """ + + main = _main_container(component) + tokens: list[str] = [] + for part in [*(main.get("command") or []), *(main.get("args") or [])]: + text = env.expand(str(part)) + try: + tokens.extend(shlex.split(text)) + except ValueError: + tokens.append(text) + found: list[tuple[str, Optional[str]]] = [] + for index, token in enumerate(tokens): + for flag in _MODEL_FLAGS: + if token == flag and index + 1 < len(tokens): + value = tokens[index + 1] + elif token.startswith(flag + "="): + value = token[len(flag) + 1 :] + else: + continue + resolved = ( + value + and "$" not in value + and _UNREADABLE not in value + and not value.startswith("-") + ) + found.append((flag, value if resolved else None)) + return found + + +def _component_model_flags( + dgd: Mapping[str, Any], component: Mapping[str, Any] +) -> list[tuple[str, Optional[str]]]: + """``_model_flags`` of every pod template the component runs (``_role_specs``).""" + + return [ + flag + for _, spec in _role_specs(component) + for flag in _model_flags(spec, _component_env(dgd, spec)) + ] + + +def _flag_model(flags: Sequence[tuple[str, Optional[str]]]) -> Optional[str]: + """The operator's served model: the first flag of ``_MODEL_FLAGS`` given.""" + + for flag in _MODEL_FLAGS: + for name, value in flags: + if name == flag: + return value + return None + + +def _worker_model_names( + dgd: Mapping[str, Any], components: Sequence[Mapping[str, Any]] +) -> set[str]: + """Every model name or path the workers give, and their served names.""" + + projected = _projected_status(dgd) + names: set[str] = set() + for component in components: + if _type(component) not in WORKER_TYPES: + continue + served = (projected.get(str(component.get("name", ""))) or {}).get( + "servedModelName" + ) + if served: + names.add(str(served)) + for _, value in _component_model_flags(dgd, component): + if value: + names.update(part for part in re.split(r"[\s,]+", value) if part) + return names + + +def _csv(value: Optional[str]) -> list[str]: + if not value: + return [] + return [part.strip().lower() for part in value.split(",") if part.strip()] + + +def _unique(items: Iterable[str]) -> list[str]: + return list(dict.fromkeys(items)) diff --git a/components/src/dynamo/planner/retune/live/kube.py b/components/src/dynamo/planner/retune/live/kube.py new file mode 100644 index 000000000000..3746f699b1b0 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/kube.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``contracts.RetuneKube`` over the official ``kubernetes`` client. + +Every call is synchronous (callers on the event loop use +``asyncio.to_thread``), carries ``_request_timeout=(connect, read)`` and +raises the ``contracts.Kube*`` errors instead of client exceptions. The +error carries the HTTP status as ``error.status`` (``None`` for a network +error or timeout), so a caller can tell 403 from the rest. + +Custom resources (the DGD at v1beta1, DGDR v1beta2 Requests, Runs and +Candidates) come back as the API server's JSON. Leases are read without the +client's model deserialization (``_preload_content=False``) so a Lease dict +round-trips unchanged: deserializing turns ``renewTime`` into a ``datetime`` +whose ``isoformat`` drops a zero fraction, which a later replace would send +back in a form the ``MicroTime`` parser rejects. +""" + +from __future__ import annotations + +import json +import logging +import math +from collections.abc import Callable, Mapping +from typing import Any, Literal, Optional, TypeVar + +from kubernetes import client +from kubernetes.client import ApiException + +from dynamo.planner.retune.live.contracts import ( + DGDR_CANDIDATES, + DGDR_GROUP, + DGDR_REQUESTS, + DGDR_RUNS, + DGDR_VERSION, + KubeAlreadyExists, + KubeConflict, + KubeError, + KubeInvalid, + KubeNotFound, + PodInfo, +) + +logger = logging.getLogger(__name__) + +DGD_VERSION = "v1beta1" +DGD_PLURAL = "dynamographdeployments" +# The frontend's container and HTTP port names (deploy/operator/internal/consts). +MAIN_CONTAINER_NAME = "main" +HTTP_PORT_NAME = "http" +FIELD_MANAGER = "dynamo-planner-retune" + +_T = TypeVar("_T") +_Operation = Literal["read", "create", "update", "delete"] + + +def request_timeout(timeout_s: float) -> tuple[int, int]: + """The client's ``(connect, read)`` timeout for ``timeout_s``, at least 1 s each.""" + + if not math.isfinite(timeout_s) or timeout_s <= 0: + seconds = 1 + else: + seconds = max(1, int(math.ceil(timeout_s))) + return (seconds, seconds) + + +def _kube_error(error: Exception, operation: _Operation, what: str) -> KubeError: + """Map a client exception to the contracts error for ``operation``.""" + + if isinstance(error, ApiException): + status = error.status + detail = _api_message(error) + message = f"{operation} {what}: HTTP {status}: {detail}" + mapped: KubeError + if status == 404: + mapped = KubeNotFound(message) + elif status == 409 and operation == "create": + mapped = KubeAlreadyExists(message) + elif status == 409: + mapped = KubeConflict(message) + elif status in (400, 422): + mapped = KubeInvalid(message) + else: + mapped = KubeError(message) + mapped.status = status # type: ignore[attr-defined] + return mapped + mapped = KubeError(f"{operation} {what}: {type(error).__name__}: {error}") + mapped.status = None # type: ignore[attr-defined] + return mapped + + +def _api_message(error: ApiException) -> str: + """The ``Status.message`` of an API error body, else its reason.""" + + body = error.body + if isinstance(body, bytes): + body = body.decode("utf-8", errors="replace") + if isinstance(body, str) and body: + try: + parsed = json.loads(body) + except ValueError: + return body[:500] + if isinstance(parsed, dict) and parsed.get("message"): + return str(parsed["message"])[:500] + return str(error.reason or "") + + +def _pod_info(pod: Any) -> PodInfo: + """``PodInfo`` of a ``V1Pod``.""" + + metadata = pod.metadata + status = pod.status + ready = any( + condition.type == "Ready" and condition.status == "True" + for condition in (status.conditions if status is not None else None) or () + ) + return PodInfo( + name=metadata.name, + uid=metadata.uid, + ip=status.pod_ip if status is not None else None, + ready=ready, + terminating=metadata.deletion_timestamp is not None, + labels=dict(metadata.labels or {}), + http_port=_http_port(pod.spec), + phase=status.phase if status is not None else None, + restarts=_main_restarts(status), + ) + + +def _main_restarts(status: Any) -> Optional[int]: + """``restartCount`` of the main container, if its status is reported.""" + + for container in (status.container_statuses if status is not None else None) or (): + if container.name == MAIN_CONTAINER_NAME: + return int(container.restart_count) + return None + + +def _http_port(spec: Any) -> Optional[int]: + """The ``http`` container port of the main container, else of any container.""" + + containers = list((spec.containers if spec is not None else None) or ()) + containers.sort(key=lambda container: container.name != MAIN_CONTAINER_NAME) + for container in containers: + for port in container.ports or (): + if port.name == HTTP_PORT_NAME: + return int(port.container_port) + return None + + +def _decode(response: Any) -> dict: + """The JSON body of a ``_preload_content=False`` response.""" + + data = response.data + if isinstance(data, bytes): + data = data.decode("utf-8") + return json.loads(data) if data else {} + + +class KubernetesRetuneKube: + """``contracts.RetuneKube`` for one namespace over ``api_client``. + + The client must carry a loaded configuration (in-cluster or kubeconfig): + a bare ``ApiClient()`` targets localhost. ``from_kube_api`` reuses the + planner connector's client, which ``KubernetesAPI`` configured. + """ + + def __init__(self, namespace: str, api_client: client.ApiClient): + if not namespace: + raise ValueError("namespace must be set") + self.namespace = namespace + self.api_client = api_client + self._custom = client.CustomObjectsApi(self.api_client) + self._core = client.CoreV1Api(self.api_client) + self._coordination = client.CoordinationV1Api(self.api_client) + + @classmethod + def from_kube_api(cls, kube_api: Any) -> KubernetesRetuneKube: + """Share a ``KubernetesAPI``'s client and namespace.""" + + return cls( + kube_api.current_namespace, + api_client=kube_api.custom_api.api_client, + ) + + # ------------------------------------------------------------ helpers + + def _call(self, operation: _Operation, what: str, function: Callable[[], _T]) -> _T: + try: + return function() + except KubeError: + raise + except Exception as error: # ApiException, urllib3 and socket errors + raise _kube_error(error, operation, what) from error + + def _get_custom( + self, version: str, plural: str, name: str, timeout_s: float + ) -> Optional[dict]: + try: + return self._call( + "read", + f"{plural}/{name}", + lambda: self._custom.get_namespaced_custom_object( + group=DGDR_GROUP, + version=version, + namespace=self.namespace, + plural=plural, + name=name, + _request_timeout=request_timeout(timeout_s), + ), + ) + except KubeNotFound: + return None + + def _list_custom( + self, version: str, plural: str, label_selector: str, timeout_s: float + ) -> list[dict]: + result = self._call( + "read", + f"{plural}?labelSelector={label_selector}", + lambda: self._custom.list_namespaced_custom_object( + group=DGDR_GROUP, + version=version, + namespace=self.namespace, + plural=plural, + label_selector=label_selector, + _request_timeout=request_timeout(timeout_s), + ), + ) + return list(result.get("items") or []) + + # ----------------------------------------------------------------- DGD + + def get_dgd(self, name: str, *, timeout_s: float = 10.0) -> dict: + found = self._get_custom(DGD_VERSION, DGD_PLURAL, name, timeout_s) + if found is None: + error = KubeNotFound( + f"DynamoGraphDeployment {self.namespace}/{name} not found" + ) + error.status = 404 # type: ignore[attr-defined] + raise error + return found + + # ---------------------------------------------------------------- pods + + def list_pods( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[PodInfo]: + result = self._call( + "read", + f"pods?labelSelector={label_selector}", + lambda: self._core.list_namespaced_pod( + namespace=self.namespace, + label_selector=label_selector, + _request_timeout=request_timeout(timeout_s), + ), + ) + return [_pod_info(pod) for pod in result.items or ()] + + # ---------------------------------------------------- DGDR v1beta2 API + + def create_request( + self, + body: dict, + *, + dry_run: bool = False, + timeout_s: float = 10.0, + ) -> dict: + kwargs: dict[str, Any] = { + "field_validation": "Strict", + "field_manager": FIELD_MANAGER, + "_request_timeout": request_timeout(timeout_s), + } + if dry_run: + kwargs["dry_run"] = "All" + name = (body.get("metadata") or {}).get("name", "") + return self._call( + "create", + f"{DGDR_REQUESTS}/{name}", + lambda: self._custom.create_namespaced_custom_object( + group=DGDR_GROUP, + version=DGDR_VERSION, + namespace=self.namespace, + plural=DGDR_REQUESTS, + body=body, + **kwargs, + ), + ) + + def get_request(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + return self._get_custom(DGDR_VERSION, DGDR_REQUESTS, name, timeout_s) + + def list_requests( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + return self._list_custom(DGDR_VERSION, DGDR_REQUESTS, label_selector, timeout_s) + + def delete_request( + self, + name: str, + *, + propagation: Literal["Background", "Foreground"] = "Background", + timeout_s: float = 10.0, + ) -> None: + try: + self._call( + "delete", + f"{DGDR_REQUESTS}/{name}", + lambda: self._custom.delete_namespaced_custom_object( + group=DGDR_GROUP, + version=DGDR_VERSION, + namespace=self.namespace, + plural=DGDR_REQUESTS, + name=name, + propagation_policy=propagation, + _request_timeout=request_timeout(timeout_s), + ), + ) + except KubeNotFound: + pass + + def get_run(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + return self._get_custom(DGDR_VERSION, DGDR_RUNS, name, timeout_s) + + def list_candidates( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + return self._list_custom( + DGDR_VERSION, DGDR_CANDIDATES, label_selector, timeout_s + ) + + # -------------------------------------------------------------- Leases + + def get_lease(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + try: + return self._call( + "read", + f"leases/{name}", + lambda: _decode( + self._coordination.read_namespaced_lease( + name=name, + namespace=self.namespace, + _preload_content=False, + _request_timeout=request_timeout(timeout_s), + ) + ), + ) + except KubeNotFound: + return None + + def create_lease(self, body: dict, *, timeout_s: float = 10.0) -> dict: + name = (body.get("metadata") or {}).get("name", "") + return self._call( + "create", + f"leases/{name}", + lambda: _decode( + self._coordination.create_namespaced_lease( + namespace=self.namespace, + body=_lease_body(body), + field_manager=FIELD_MANAGER, + _preload_content=False, + _request_timeout=request_timeout(timeout_s), + ) + ), + ) + + def replace_lease(self, name: str, body: dict, *, timeout_s: float = 10.0) -> dict: + return self._call( + "update", + f"leases/{name}", + lambda: _decode( + self._coordination.replace_namespaced_lease( + name=name, + namespace=self.namespace, + body=_lease_body(body), + field_manager=FIELD_MANAGER, + _preload_content=False, + _request_timeout=request_timeout(timeout_s), + ) + ), + ) + + def delete_lease( + self, + name: str, + *, + resource_version: Optional[str] = None, + timeout_s: float = 10.0, + ) -> None: + """Missing Leases are not an error. ``resource_version`` makes the + delete conditional (a 409 -> ``KubeConflict`` when it changed).""" + + options = None + if resource_version is not None: + options = client.V1DeleteOptions( + preconditions=client.V1Preconditions(resource_version=resource_version) + ) + try: + self._call( + "delete", + f"leases/{name}", + lambda: self._coordination.delete_namespaced_lease( + name=name, + namespace=self.namespace, + body=options, + _request_timeout=request_timeout(timeout_s), + ), + ) + except KubeNotFound: + pass + + +def _lease_body(body: Mapping[str, Any]) -> dict: + """A Lease body with its type fields set (the API server requires them).""" + + lease = dict(body) + lease.setdefault("apiVersion", "coordination.k8s.io/v1") + lease.setdefault("kind", "Lease") + return lease diff --git a/components/src/dynamo/planner/retune/live/metrics.py b/components/src/dynamo/planner/retune/live/metrics.py new file mode 100644 index 000000000000..b586425f4efa --- /dev/null +++ b/components/src/dynamo/planner/retune/live/metrics.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``dynamo_planner_retune_*`` Prometheus metrics of live retune. + +``RetuneMetrics`` implements ``contracts.RetuneMetricsSink``: decision, +failure and skip counters, the supervisor's state, the scores of the latest +decision, router convergence, PATCH errors and trace health (the table in +the live retune guide's Metrics section). +Every method is cheap and never raises, so callers on the event loop, in the +sweep thread or in the actuator need no guard. + +The planner serves ``prometheus_client.REGISTRY`` on +``metric_reporting_prometheus_port``, so the default registry is the one to +use. The collectors are created once per registry and shared: building a +second ``RetuneMetrics`` on the same registry (a controller rebuild, a test) +does not fail with a duplicate registration. +""" + +from __future__ import annotations + +import logging +import math +import threading +import weakref +from typing import Optional + +from prometheus_client import REGISTRY, CollectorRegistry, Counter, Gauge + +from dynamo.planner.core.types import RETUNE_OUTCOMES + +logger = logging.getLogger(__name__) + +PREFIX = "dynamo_planner_retune" +# Gate reasons (``contracts.GateReason``) and decision outcomes are known up +# front; their series start at 0 so ``increase()`` sees the first event. +_SKIP_REASONS = ("not_leader", "not_ready", "already_decided") +# Live outcomes recorded in the state history besides the controller's. +_LIVE_OUTCOMES = ("promoted", "rolled_back") + + +class _Collectors: + """The metric objects of one registry.""" + + def __init__(self, registry: CollectorRegistry) -> None: + self.decisions_total = Counter( + f"{PREFIX}_decisions_total", + "Retune decisions, counted once the state write persisted them.", + labelnames=["outcome"], + registry=registry, + ) + self.failures_total = Counter( + f"{PREFIX}_failures_total", + "Failed retune windows and supervisor failures, by reason.", + labelnames=["reason"], + registry=registry, + ) + self.skipped_total = Counter( + f"{PREFIX}_skipped_total", + "Due windows the launch gate held back, once per window and reason.", + labelnames=["reason"], + registry=registry, + ) + self.leader = Gauge( + f"{PREFIX}_leader", + "1 while this planner holds the retune leadership Lease.", + registry=registry, + ) + self.ready = Gauge( + f"{PREFIX}_ready", + "1 while every spec and health premise holds, so windows launch.", + registry=registry, + ) + self.initialized = Gauge( + f"{PREFIX}_initialized", + "1 once the leader finished initializing from the persisted state.", + registry=registry, + ) + self.rollback_pending = Gauge( + f"{PREFIX}_rollback_pending", + "1 while routers are being reset to their startup parameters.", + registry=registry, + ) + self.controller_generation = Gauge( + f"{PREFIX}_controller_generation", + "Generation of the installed retune controller; it grows on each rebuild.", + registry=registry, + ) + self.baseline_score = Gauge( + f"{PREFIX}_baseline_score", + "Objective score of the current parameters on the latest decided " + "window (NaN when not scored).", + registry=registry, + ) + self.best_score = Gauge( + f"{PREFIX}_best_score", + "Best candidate score on the latest decided window (NaN when none).", + registry=registry, + ) + self.sweep_seconds = Gauge( + f"{PREFIX}_sweep_seconds", + "Wall-clock duration of the latest sweep (NaN when it did not run).", + registry=registry, + ) + self.missing_candidates = Gauge( + f"{PREFIX}_missing_candidates", + "Candidates of the latest sweep without a score.", + registry=registry, + ) + self.router_converged = Gauge( + f"{PREFIX}_router_converged", + "1 when every frontend router holds the target policy parameters.", + registry=registry, + ) + self.router_patch_errors_total = Counter( + f"{PREFIX}_router_patch_errors_total", + "Failed /router_policy PATCH or reset calls, by HTTP status or kind.", + labelnames=["code"], + registry=registry, + ) + self.trace_sink_missing = Gauge( + f"{PREFIX}_trace_sink_missing", + "1 for a Ready frontend pod without its own trace segment directory.", + labelnames=["pod"], + registry=registry, + ) + self.trace_dropped_members_total = Counter( + f"{PREFIX}_trace_dropped_members_total", + "Trace gzip members dropped by window cuts for a bad CRC, size or " + "truncation.", + registry=registry, + ) + for outcome in (*RETUNE_OUTCOMES, *_LIVE_OUTCOMES): + self.decisions_total.labels(outcome=outcome) + for reason in _SKIP_REASONS: + self.skipped_total.labels(reason=reason) + for gauge in (self.baseline_score, self.best_score, self.sweep_seconds): + gauge.set(math.nan) + + +_LOCK = threading.Lock() +_BY_REGISTRY: weakref.WeakKeyDictionary[ + CollectorRegistry, _Collectors +] = weakref.WeakKeyDictionary() + + +def _collectors(registry: CollectorRegistry) -> _Collectors: + with _LOCK: + collectors = _BY_REGISTRY.get(registry) + if collectors is None: + collectors = _Collectors(registry) + _BY_REGISTRY[registry] = collectors + return collectors + + +def _label(value: object) -> str: + text = str(value) + return text if text else "unknown" + + +def _number(value: Optional[float]) -> float: + return math.nan if value is None else float(value) + + +class RetuneMetrics: + """``RetuneMetricsSink`` over Prometheus collectors on ``registry``. + + ``registry`` defaults to ``prometheus_client.REGISTRY``, which the + planner's metrics server exposes. Instances on the same registry share + their collectors. + """ + + def __init__(self, registry: CollectorRegistry | None = None) -> None: + self._m = _collectors(REGISTRY if registry is None else registry) + + def decision(self, outcome: str) -> None: + try: + self._m.decisions_total.labels(outcome=_label(outcome)).inc() + except Exception: + logger.debug("retune decision metric failed", exc_info=True) + + def failure(self, reason: str) -> None: + try: + self._m.failures_total.labels(reason=_label(reason)).inc() + except Exception: + logger.debug("retune failure metric failed", exc_info=True) + + def skipped(self, reason: str) -> None: + try: + self._m.skipped_total.labels(reason=_label(reason)).inc() + except Exception: + logger.debug("retune skipped metric failed", exc_info=True) + + def supervisor_state( + self, + *, + leader: bool, + ready: bool, + initialized: bool, + rollback_pending: bool, + generation: int, + ) -> None: + try: + self._m.leader.set(1 if leader else 0) + self._m.ready.set(1 if ready else 0) + self._m.initialized.set(1 if initialized else 0) + self._m.rollback_pending.set(1 if rollback_pending else 0) + self._m.controller_generation.set(generation) + except Exception: + logger.debug("retune supervisor metrics failed", exc_info=True) + + def scores( + self, + *, + baseline: Optional[float], + best: Optional[float], + sweep_seconds: Optional[float], + missing_candidates: int, + ) -> None: + try: + self._m.baseline_score.set(_number(baseline)) + self._m.best_score.set(_number(best)) + self._m.sweep_seconds.set(_number(sweep_seconds)) + self._m.missing_candidates.set(missing_candidates) + except Exception: + logger.debug("retune score metrics failed", exc_info=True) + + def router_converged(self, converged: bool) -> None: + try: + self._m.router_converged.set(1 if converged else 0) + except Exception: + logger.debug("retune router_converged metric failed", exc_info=True) + + def patch_error(self, code: str) -> None: + try: + self._m.router_patch_errors_total.labels(code=_label(code)).inc() + except Exception: + logger.debug("retune patch error metric failed", exc_info=True) + + def trace_sink_missing(self, pod: str, missing: bool) -> None: + """Set the pod's series to 1, or drop it once the sink is back. + + Dropping instead of writing 0 keeps no series for pods that are gone. + """ + + try: + if missing: + self._m.trace_sink_missing.labels(pod=_label(pod)).set(1) + else: + try: + self._m.trace_sink_missing.remove(_label(pod)) + except KeyError: + pass + except Exception: + logger.debug("retune trace_sink_missing metric failed", exc_info=True) + + def dropped_members(self, count: int) -> None: + try: + if count > 0: + self._m.trace_dropped_members_total.inc(count) + except Exception: + logger.debug("retune dropped members metric failed", exc_info=True) diff --git a/components/src/dynamo/planner/retune/live/reconciler.py b/components/src/dynamo/planner/retune/live/reconciler.py new file mode 100644 index 000000000000..a3ce986b1b10 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/reconciler.py @@ -0,0 +1,932 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Garbage collection of live retune files and DGDR Requests. + +The supervisor runs ``Reconciler.run_once`` only while this process leads and +is initialized. It deletes, with ``fs_root`` the planner's mount of the trace +claim, ``T``/``G`` the retune interval/grace and ``timeout`` the sweep +timeout:: + + ////.NNNNNN.jsonl.gz frontend segments + ///-.jsonl.gz window files + ///.tmp-- cut temp files + +- Segments, only when the frontend pod set is known and nonempty: never the + highest-numbered segment of a live pod; other segments of a live pod, and + the whole directory of a pod that no longer exists, once every file is + older than ``T + G + timeout + VISIBILITY_MARGIN_S``. +- Window files, never while protected: ``window_end_ms`` in + ``[last_window_end_ms, due_window_end_ms]`` with ``now < D_k``, or in use + by this process's sweep (``protect``). Otherwise deleted when at least one + Request references the file and every referencing Request's current Run + has ``Completed=True`` and no Job pod left, or when no Request references + it and it is older than ``T + timeout``. +- Temp files older than ``cut_budget_seconds + TEMP_MARGIN_S``, unless + their window is in use by this process's sweep. +- Requests of this DGD name, never one this process's sweep protects: a + terminal Request outside the newest ``keep_requests`` windows, a + non-terminal Request whose window deadline has passed, a Request in the + delete queue, and a Request of another DGD uid (the DGD was recreated). + +``Reconciler`` is also the ``SweepRegistry`` the DGDR sweep uses, from the +controller's executor thread, to protect its window and Requests and to +queue deletions. +""" + +from __future__ import annotations + +import logging +import os +import re +import shutil +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, Optional + +from dynamo.planner.config.planner_config import DgdrExecutorConfig, RetuneConfig +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD, + LABEL_DGD_UID, + LABEL_WINDOW, + DgdRef, + KubeError, + PodInfo, + RetuneKube, + RetuneMetricsSink, +) +from dynamo.planner.retune.live.window_cut import ( + SEGMENT_SUFFIX, + TEMP_NAME_RE, + TEMP_PREFIX, +) + +logger = logging.getLogger(__name__) + +# Allowance for RWX attribute caching and the frontend's flush interval on +# top of T + G + timeout before a segment counts as unused. +VISIBILITY_MARGIN_S = 600.0 +# Allowance on top of cut_budget_seconds before a temp file counts as +# abandoned (the cut subprocess is killed at its deadline). +TEMP_MARGIN_S = 300.0 +KUBE_TIMEOUT_S = 10.0 + +# Label the Job controller puts on a Job's pods (batchv1.JobNameLabel). +JOB_NAME_LABEL = "batch.kubernetes.io/job-name" +_FINISHED_POD_PHASES = frozenset({"Succeeded", "Failed"}) +# Annotation the DGDR v1beta2 Request controller puts on each Run: the +# Request generation it was created for. +RUN_GENERATION_ANNOTATION = "nvidia.com/dgdr-request-generation" +RUN_COMPLETED_CONDITION = "Completed" +REQUEST_ACCEPTED_CONDITION = "Accepted" +REQUEST_INVALID_SPEC_REASON = "InvalidSpec" + +# A window file: [-].jsonl.gz +_WINDOW_RE = re.compile(r"^(\d+)(?:-[0-9a-f]+)?" + re.escape(SEGMENT_SUFFIX) + r"$") +# A Rust gzip JSONL sink segment: ..jsonl.gz +_SEGMENT_RE = re.compile(r"^(.*)\.(\d{6,})" + re.escape(SEGMENT_SUFFIX) + r"$") + + +def trace_pvc_source(spec: Any) -> tuple[Optional[str], Optional[str]]: + """``(claimName, subPath)`` of a DGDR spec's trace claim. + + The fields are ``workload.trace.source.persistentVolumeClaim.claimName`` + and ``.subPath`` (the claim-relative file the Run replays); each is + ``None`` when absent or not a string. + """ + + pvc = _dig(spec, "workload", "trace", "source", "persistentVolumeClaim") + if not isinstance(pvc, Mapping): + return None, None + claim, sub_path = pvc.get("claimName"), pvc.get("subPath") + return ( + claim if isinstance(claim, str) else None, + sub_path if isinstance(sub_path, str) else None, + ) + + +@dataclass(frozen=True) +class ReconcileReport: + """What one ``run_once`` round did.""" + + deleted_segments: int = 0 + deleted_pod_dirs: tuple[str, ...] = () + deleted_windows: tuple[int, ...] = () + deleted_temp_files: int = 0 + deleted_requests: tuple[str, ...] = () + # Ready frontend pods without a segment directory (trace_sink_missing). + sink_missing: tuple[str, ...] = () + errors: dict[str, int] = field(default_factory=dict) + # The round ended early: ``should_continue`` said no, or the DGD + # identity changed (``set_dgd``) while it ran. + stopped: bool = False + + +@dataclass +class _RequestView: + name: str + uid: str + dgd_uid: Optional[str] + window_end_ms: Optional[int] + trace_file: Optional[str] # window file name if it reads one of ours + obj: dict[str, Any] + + +@dataclass +class _Protection: + """One window's protection: references held and the names protected.""" + + count: int = 0 + names: set[str] = field(default_factory=set) + + +@dataclass +class _Round: + """Mutable bookkeeping of one ``run_once`` call. Counts accumulate as + deletions happen, so a phase that fails half way still reports them.""" + + now: float + dgd: DgdRef + segments_root: str + windows_root: str + should_continue: Optional[Callable[[], bool]] + errors: dict[str, int] = field(default_factory=dict) + stopped: bool = False + deleted_segments: int = 0 + deleted_dirs: list[str] = field(default_factory=list) + deleted_windows: list[int] = field(default_factory=list) + deleted_temp: int = 0 + deleted_requests: list[str] = field(default_factory=list) + sink_missing: list[str] = field(default_factory=list) + + def error(self, kind: str) -> None: + self.errors[kind] = self.errors.get(kind, 0) + 1 + + +class Reconciler: + """Deletes what live retune no longer needs; see the module docstring.""" + + def __init__( + self, + kube: RetuneKube, + dgd: DgdRef, + cfg: DgdrExecutorConfig, + retune: RetuneConfig, + fs_root: str, + metrics: RetuneMetricsSink, + now: Callable[[], float] = time.time, + ) -> None: + self._kube = kube + self._cfg = cfg + self._fs_root = fs_root + self._metrics = metrics + self._now = now + interval_s = float(retune.interval_seconds) + sweep_timeout_s = float(retune.effective_timeout_seconds) + self._segment_max_age_s = ( + interval_s + + float(retune.grace_seconds) + + sweep_timeout_s + + VISIBILITY_MARGIN_S + ) + self._window_max_age_s = interval_s + sweep_timeout_s + self._temp_max_age_s = float(cfg.cut_budget_seconds) + TEMP_MARGIN_S + + self._lock = threading.Lock() + # Guarded by _lock (set_dgd may run while a round is on its thread). + self._dgd = dgd + self._segments_root = self._root(cfg.segments_dir, dgd) + self._windows_root = self._root(cfg.windows_dir, dgd) + # window_end_ms -> references and Request names of running sweeps. + self._protected: dict[int, _Protection] = {} + self._delete_queue: set[str] = set() + # (request uid, generation) of Requests known terminal / finished; + # neither can revert for one generation. + self._terminal: set[tuple[str, int]] = set() + self._finished: set[tuple[str, int]] = set() + # Owned by the round (one runs at a time): the pods whose + # trace_sink_missing gauge is set, those last reported missing, and + # the (Request uid, label) pairs already reported missing. + self._sink_reported: set[str] = set() + self._sink_missing: set[str] = set() + self._unlabelled: set[tuple[str, str]] = set() + + def _root(self, directory: str, dgd: DgdRef) -> str: + return os.path.join(self._fs_root, directory.strip("/"), dgd.name) + + # ------------------------------------------------------------- identity + + def set_dgd(self, dgd: DgdRef) -> None: + """Follow the DGD identity (it was recreated under the same name). + + Requests labelled with the new uid become this DGD's; those of the + old uid are then deleted as ``other_dgd_uid``. The Request-state + caches are cleared. A round running on another thread ends at its + next delete, since it classified Requests by the old identity. + """ + + with self._lock: + if dgd != self._dgd: + logger.info( + "retune reconciler: DGD %s uid %s -> %s %s", + self._dgd.name, + self._dgd.uid, + dgd.name, + dgd.uid, + ) + self._dgd = dgd + self._segments_root = self._root(self._cfg.segments_dir, dgd) + self._windows_root = self._root(self._cfg.windows_dir, dgd) + self._terminal.clear() + self._finished.clear() + + @property + def dgd(self) -> DgdRef: + with self._lock: + return self._dgd + + # --------------------------------------------------------- SweepRegistry + + def protect(self, window_end_ms: int, request_names: Iterable[str]) -> None: + """Take one reference on the window and add ``request_names`` to + its protected names; each call needs a matching ``release``.""" + + names = set(request_names) + with self._lock: + entry = self._protected.setdefault(int(window_end_ms), _Protection()) + entry.count += 1 + entry.names.update(names) + + def release(self, window_end_ms: int) -> None: + """Drop one reference; the window and its names are unprotected + once no reference is left. Releasing an unknown window is a no-op.""" + + with self._lock: + window = int(window_end_ms) + entry = self._protected.get(window) + if entry is None: + return + entry.count -= 1 + if entry.count <= 0: + del self._protected[window] + + def enqueue_delete(self, request_names: Iterable[str]) -> None: + names = set(request_names) + with self._lock: + self._delete_queue.update(names) + + # ---------------------------------------------------------- inspection + + @property + def segments_root(self) -> str: + with self._lock: + return self._segments_root + + @property + def windows_root(self) -> str: + with self._lock: + return self._windows_root + + def protected_windows(self) -> frozenset[int]: + with self._lock: + return frozenset(self._protected) + + def protected_requests(self) -> frozenset[str]: + with self._lock: + return frozenset( + name for entry in self._protected.values() for name in entry.names + ) + + def pending_deletes(self) -> frozenset[str]: + with self._lock: + return frozenset(self._delete_queue) + + # ---------------------------------------------------------------- round + + def run_once( + self, + *, + last_window_end_ms: Optional[int], + due_window_end_ms: Optional[int], + deadline_for: Callable[[int], float], + frontend_pods: Optional[Sequence[PodInfo]], + should_continue: Optional[Callable[[], bool]] = None, + ) -> ReconcileReport: + """One GC round. ``deadline_for(window_end_ms)`` is that window's + sweep deadline ``D_k`` in epoch seconds; ``frontend_pods`` is the + actuator's listing (None when it failed). ``should_continue`` is + called before every delete, unlink or Kubernetes read; once it + returns False (or raises) the round deletes and reads nothing more + and returns. It blocks on the filesystem and the Kubernetes API: + call it with ``asyncio.to_thread``.""" + + now = float(self._now()) + with self._lock: + state = _Round( + now=now, + dgd=self._dgd, + segments_root=self._segments_root, + windows_root=self._windows_root, + should_continue=should_continue, + ) + + def in_window_range(window_end_ms: int) -> bool: + """Undecided (from the last decided window to the due one) and + before its deadline D_k: a sweep may still cut or adopt it.""" + + if last_window_end_ms is not None and window_end_ms < last_window_end_ms: + return False + if due_window_end_ms is not None and window_end_ms > due_window_end_ms: + return False + return now < deadline_for(window_end_ms) + + try: + self._collect_segments(state, frontend_pods) + except Exception: # keep the sink check, Request and window GC running + logger.exception("retune reconciler: segment GC failed") + state.error("segments") + if not state.stopped: + try: + state.sink_missing = self._report_sinks(state, frontend_pods) + except Exception: + logger.exception("retune reconciler: trace sink check failed") + state.error("sinks") + + requests: Optional[list[_RequestView]] = None + if self._proceed(state): + try: + requests = self._list_requests(state.dgd) + except KubeError as exc: + logger.warning("retune reconciler: listing Requests failed: %s", exc) + state.error("list_requests") + except Exception: + logger.exception("retune reconciler: listing Requests failed") + state.error("list_requests") + try: + self._collect_requests(state, requests, in_window_range, deadline_for) + except Exception: + logger.exception("retune reconciler: Request GC failed") + state.error("requests") + + if not state.stopped: + try: + remaining = ( + None + if requests is None + else [ + view + for view in requests + if view.name not in state.deleted_requests + ] + ) + self._collect_windows(state, remaining, in_window_range) + except Exception: + logger.exception("retune reconciler: window GC failed") + state.error("windows") + + report = ReconcileReport( + deleted_segments=state.deleted_segments, + deleted_pod_dirs=tuple(sorted(state.deleted_dirs)), + deleted_windows=tuple(sorted(state.deleted_windows)), + deleted_temp_files=state.deleted_temp, + deleted_requests=tuple(sorted(state.deleted_requests)), + sink_missing=tuple(sorted(state.sink_missing)), + errors=state.errors, + stopped=state.stopped, + ) + if ( + report.deleted_segments + or report.deleted_pod_dirs + or report.deleted_windows + or report.deleted_temp_files + or report.deleted_requests + ): + logger.info( + "retune reconciler: deleted %d segments, pod dirs %s, windows %s, " + "%d temp files, Requests %s", + report.deleted_segments, + list(report.deleted_pod_dirs), + list(report.deleted_windows), + report.deleted_temp_files, + list(report.deleted_requests), + ) + return report + + def _proceed(self, state: _Round) -> bool: + """Whether the round may delete or read one more thing; asked right + before every delete, unlink or Kubernetes read. A no is final for + the round.""" + + if state.stopped: + return False + with self._lock: + current = self._dgd + if current != state.dgd: + logger.info( + "retune reconciler: the DGD identity changed during the round; " + "stopping it" + ) + state.stopped = True + return False + if state.should_continue is None: + return True + try: + proceed = bool(state.should_continue()) + except Exception: + logger.exception("retune reconciler: should_continue failed") + proceed = False + if not proceed: + logger.info("retune reconciler: told to stop; ending the round early") + state.stopped = True + return proceed + + # ------------------------------------------------------------- segments + + def _collect_segments( + self, state: _Round, frontend_pods: Optional[Sequence[PodInfo]] + ) -> None: + if not frontend_pods: + # Unknown or empty pod set: every directory would look orphaned. + return + live = {pod.name for pod in frontend_pods} + cutoff = state.now - self._segment_max_age_s + for entry in _scandir(state.segments_root): + if state.stopped: + return + # One pod directory that cannot be read (another uid, umask 077) + # must not starve the pods after it in scandir order. + try: + if not entry.is_dir(follow_symlinks=False): + continue + if entry.name in live: + self._collect_pod_segments(state, entry.path, cutoff) + continue + if _newest_mtime(entry.path) >= cutoff: + continue + if not self._proceed(state): + return + shutil.rmtree(entry.path) + state.deleted_dirs.append(entry.name) + except FileNotFoundError: + continue + except OSError as exc: + logger.warning( + "retune reconciler: segment GC of %s failed: %s", entry.path, exc + ) + state.error("segments") + + def _collect_pod_segments( + self, state: _Round, directory: str, cutoff: float + ) -> None: + # Top level only: the frontend premise check puts the sink's file + # prefix directly under the pod directory, so there is no + # subdirectory to walk. + segments: dict[str, list[tuple[int, os.DirEntry[str]]]] = {} + for entry in _scandir(directory): + match = _SEGMENT_RE.match(entry.name) + if match is None or not entry.is_file(follow_symlinks=False): + continue + segments.setdefault(match.group(1), []).append((int(match.group(2)), entry)) + for entries in segments.values(): + entries.sort(key=lambda item: item[0]) + # The highest-numbered segment is the one the frontend appends to. + for _, entry in entries[:-1]: + try: + if entry.stat(follow_symlinks=False).st_mtime >= cutoff: + continue + if not self._proceed(state): + return + os.unlink(entry.path) + state.deleted_segments += 1 + except FileNotFoundError: + continue + except OSError as exc: + logger.warning( + "retune reconciler: removing %s failed: %s", entry.path, exc + ) + state.error("segments") + + def _report_sinks( + self, state: _Round, frontend_pods: Optional[Sequence[PodInfo]] + ) -> list[str]: + if frontend_pods is None: + return [] + missing: list[str] = [] + current: set[str] = set() + for pod in frontend_pods: + if not pod.ready or pod.terminating: + continue + current.add(pod.name) + absent = not os.path.isdir(os.path.join(state.segments_root, pod.name)) + if absent: + missing.append(pod.name) + self._metrics.trace_sink_missing(pod.name, absent) + for name in self._sink_reported - current: + self._metrics.trace_sink_missing(name, False) + self._sink_reported = current + # The gauge tracks the condition; the log marks where it starts. + newly = sorted(set(missing) - self._sink_missing) + self._sink_missing = set(missing) + if newly: + logger.warning( + "retune reconciler: Ready frontends without a trace directory " + "under %s: %s", + state.segments_root, + newly, + ) + return missing + + # -------------------------------------------------------------- windows + + def _collect_windows( + self, + state: _Round, + requests: Optional[list[_RequestView]], + in_window_range: Callable[[int], bool], + ) -> None: + protected = self.protected_windows() + references: dict[str, list[_RequestView]] = {} + for view in requests or (): + if view.trace_file is not None: + references.setdefault(view.trace_file, []).append(view) + + for entry in _scandir(state.windows_root): + if state.stopped: + return + try: + self._collect_window_entry( + state, entry, requests, references, protected, in_window_range + ) + except FileNotFoundError: + continue + except OSError as exc: + logger.warning( + "retune reconciler: window GC of %s failed: %s", entry.path, exc + ) + state.error("windows") + + def _collect_window_entry( + self, + state: _Round, + entry: os.DirEntry[str], + requests: Optional[list[_RequestView]], + references: Mapping[str, list[_RequestView]], + protected: frozenset[int], + in_window_range: Callable[[int], bool], + ) -> None: + if not entry.is_file(follow_symlinks=False): + return + if entry.name.startswith(TEMP_PREFIX): + temp = TEMP_NAME_RE.match(entry.name) + window = int(temp.group(1)) if temp is not None else None + # A temp file of a window this process's sweep is cutting + # belongs to the running cut subprocess. + if window is not None and window in self.protected_windows(): + return + if self._older_than(entry, state.now, self._temp_max_age_s): + if self._unlink(state, entry.path): + state.deleted_temp += 1 + return + match = _WINDOW_RE.match(entry.name) + if match is None or requests is None: + # Unknown file, or references unknown: keep it. + return + window = int(match.group(1)) + if window in protected or in_window_range(window): + return + referrers = references.get(entry.name, []) + if referrers: + if not all(self._finished_with(state, view) for view in referrers): + return + elif not self._older_than(entry, state.now, self._window_max_age_s): + return + # The sweep may have protected it since the snapshot above. + if window in self.protected_windows(): + return + if self._unlink(state, entry.path): + state.deleted_windows.append(window) + + # ------------------------------------------------------------- requests + + def _list_requests(self, dgd: DgdRef) -> list[_RequestView]: + items = self._kube.list_requests( + f"{LABEL_DGD}={dgd.name}", timeout_s=KUBE_TIMEOUT_S + ) + views = [] + prefix = self._cfg.windows_dir.strip("/") + "/" + dgd.name + "/" + for obj in items: + metadata = obj.get("metadata") or {} + labels = metadata.get("labels") or {} + claim, path = trace_pvc_source(obj.get("spec")) + trace_file = None + if ( + claim == self._cfg.claim_name + and path is not None + and path.startswith(prefix) + and "/" not in path[len(prefix) :] + ): + trace_file = path[len(prefix) :] + views.append( + _RequestView( + name=str(metadata.get("name", "")), + uid=str(metadata.get("uid", "")), + dgd_uid=labels.get(LABEL_DGD_UID), + window_end_ms=_int_or_none(labels.get(LABEL_WINDOW)), + trace_file=trace_file, + obj=obj, + ) + ) + return [view for view in views if view.name] + + def _collect_requests( + self, + state: _Round, + requests: Optional[list[_RequestView]], + in_window_range: Callable[[int], bool], + deadline_for: Callable[[int], float], + ) -> None: + now = state.now + dgd_uid = state.dgd.uid + reasons: dict[str, str] = {name: "queued" for name in self.pending_deletes()} + if requests is not None: + listed = {view.uid for view in requests} + with self._lock: + for cache in (self._terminal, self._finished): + cache.difference_update( + [key for key in cache if key[0] not in listed] + ) + self._unlabelled.difference_update( + [key for key in self._unlabelled if key[0] not in listed] + ) + ours = [view for view in requests if view.dgd_uid == dgd_uid] + for view in requests: + if view.dgd_uid is None: + self._report_unlabelled(view, LABEL_DGD_UID) + elif view.dgd_uid != dgd_uid: + # The DGD was deleted and recreated under the same name. + reasons.setdefault(view.name, "other_dgd_uid") + # Each uncached verdict reads Runs; _is_terminal asks + # should_continue first, so a stopped round reads nothing more. + states = {view.name: self._is_terminal(state, view) for view in ours} + terminal = {name for name, value in states.items() if value is True} + kept_windows = sorted( + { + view.window_end_ms + for view in ours + if view.name in terminal and view.window_end_ms is not None + }, + reverse=True, + )[: self._cfg.keep_requests] + for view in ours: + window = view.window_end_ms + if window is None: + self._report_unlabelled(view, LABEL_WINDOW) + continue + value = states[view.name] + if value is None: + continue # Run state unknown this round + if value: + if window not in kept_windows and not in_window_range(window): + reasons.setdefault(view.name, "terminal") + elif now >= deadline_for(window): + # Not terminal and past D_k: no sweep can adopt it anymore. + reasons.setdefault(view.name, "deadline_passed") + + for name in sorted(reasons): + if name in self.protected_requests(): + continue + if not self._proceed(state): + return + try: + self._kube.delete_request(name, timeout_s=KUBE_TIMEOUT_S) + except KubeError as exc: + logger.warning( + "retune reconciler: deleting Request %s (%s) failed: %s", + name, + reasons[name], + exc, + ) + state.error("delete_request") + continue + logger.info( + "retune reconciler: deleted Request %s (%s)", name, reasons[name] + ) + state.deleted_requests.append(name) + with self._lock: + self._delete_queue.discard(name) + + def _report_unlabelled(self, view: _RequestView, label: str) -> None: + """Warn once per Request about a missing (or unparsable) label.""" + + key = (view.uid, label) + if key not in self._unlabelled: + self._unlabelled.add(key) + logger.warning( + "retune reconciler: Request %s has no %s label", view.name, label + ) + + # ------------------------------------------------- Request / Run state + + def _cached(self, cache: set[tuple[str, int]], key: tuple[str, int]) -> bool: + with self._lock: + return key in cache + + def _remember( + self, state: _Round, cache: set[tuple[str, int]], key: tuple[str, int] + ) -> None: + """Cache a verdict unless ``set_dgd`` cleared the caches meanwhile + (the caches are cleared in place, never replaced).""" + + with self._lock: + if self._dgd == state.dgd: + cache.add(key) + + def _is_terminal(self, state: _Round, view: _RequestView) -> Optional[bool]: + """Whether the Request's current-generation Run completed, or the + operator rejected the current generation (InvalidSpec); None when + the Run could not be read.""" + + generation = _int_or_none(_dig(view.obj, "metadata", "generation")) or 0 + key = (view.uid, generation) + if self._cached(self._terminal, key): + return True + if not self._proceed(state): + return None + try: + run = self._current_run(view, generation) + except KubeError as exc: + logger.warning( + "retune reconciler: reading the Run of Request %s failed: %s", + view.name, + exc, + ) + state.error("get_run") + return None + terminal = ( + _condition_true(run, RUN_COMPLETED_CONDITION) + if run is not None + else _invalid_spec(view.obj, generation) + ) + if terminal: + self._remember(state, self._terminal, key) + return terminal + + def _finished_with(self, state: _Round, view: _RequestView) -> bool: + """The Request no longer reads its window file: terminal, no active + Run, and no pod of the current Run's Job left.""" + + generation = _int_or_none(_dig(view.obj, "metadata", "generation")) or 0 + key = (view.uid, generation) + if self._cached(self._finished, key): + return True + if _dig(view.obj, "status", "activeRunRef", "name"): + return False + if not self._proceed(state): + return False + try: + run = self._current_run(view, generation) + if run is None: + finished = _invalid_spec(view.obj, generation) + elif not _condition_true(run, RUN_COMPLETED_CONDITION): + finished = False + else: + job = _dig(run, "status", "jobRef", "name") + # A Succeeded or Failed pod no longer reads the trace; one + # with no phase reported is treated as still running. + finished = job is None or not any( + pod.phase not in _FINISHED_POD_PHASES + for pod in self._kube.list_pods( + f"{JOB_NAME_LABEL}={job}", timeout_s=KUBE_TIMEOUT_S + ) + ) + except KubeError as exc: + logger.warning( + "retune reconciler: reading Run state of Request %s failed: %s", + view.name, + exc, + ) + return False + if finished: + self._remember(state, self._finished, key) + return finished + + def _current_run( + self, view: _RequestView, generation: int + ) -> Optional[dict[str, Any]]: + """The Run created for the Request's current generation, if any. + + The Run's generation annotation identifies it, whatever its name; + the Request status names the active and the retained terminal Runs. + """ + + status = view.obj.get("status") or {} + names = [] + active = (status.get("activeRunRef") or {}).get("name") + if active: + names.append(active) + names.extend( + ref.get("name") + for ref in status.get("recentRunRefs") or () + if ref.get("name") + ) + for name in names: + run = self._kube.get_run(name, timeout_s=KUBE_TIMEOUT_S) + if run is None: + continue + annotations = _dig(run, "metadata", "annotations") or {} + if _int_or_none(annotations.get(RUN_GENERATION_ANNOTATION)) == generation: + return run + return None + + # ---------------------------------------------------------- filesystem + + @staticmethod + def _older_than(entry: os.DirEntry[str], now: float, age_s: float) -> bool: + try: + return entry.stat(follow_symlinks=False).st_mtime < now - age_s + except FileNotFoundError: + return False + + def _unlink(self, state: _Round, path: str) -> bool: + if not self._proceed(state): + return False + try: + os.unlink(path) + return True + except FileNotFoundError: + return False + except OSError as exc: + logger.warning("retune reconciler: removing %s failed: %s", path, exc) + state.error("windows") + return False + + +# ---------------------------------------------------------------- helpers + + +def _scandir(path: str) -> list[os.DirEntry[str]]: + """Entries of ``path``; none when it is missing. Other ``OSError``s + (e.g. ``PermissionError``) propagate to the caller's per-entry handler.""" + + try: + with os.scandir(path) as entries: + return list(entries) + except (FileNotFoundError, NotADirectoryError): + return [] + + +def _newest_mtime(path: str) -> float: + """Newest mtime of ``path`` and everything under it.""" + + newest = os.lstat(path).st_mtime + for root, dirs, files in os.walk(path): + for name in dirs + files: + try: + newest = max(newest, os.lstat(os.path.join(root, name)).st_mtime) + except FileNotFoundError: + continue + return newest + + +def _dig(obj: Any, *keys: str) -> Any: + for key in keys: + if not isinstance(obj, Mapping): + return None + obj = obj.get(key) + return obj + + +def _int_or_none(value: Any) -> Optional[int]: + if value is None or isinstance(value, bool): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _conditions(obj: Optional[Mapping[str, Any]]) -> list[Mapping[str, Any]]: + conditions = _dig(obj, "status", "conditions") + return [c for c in conditions or () if isinstance(c, Mapping)] + + +def _condition_true(obj: Optional[Mapping[str, Any]], kind: str) -> bool: + return any( + c.get("type") == kind and c.get("status") == "True" for c in _conditions(obj) + ) + + +def _invalid_spec(request: Mapping[str, Any], generation: int) -> bool: + """The operator rejected this generation; only a new one changes that.""" + + for condition in _conditions(request): + if ( + condition.get("type") == REQUEST_ACCEPTED_CONDITION + and condition.get("status") == "False" + and condition.get("reason") == REQUEST_INVALID_SPEC_REASON + and _int_or_none(condition.get("observedGeneration")) == generation + ): + return True + return False diff --git a/components/src/dynamo/planner/retune/live/state.py b/components/src/dynamo/planner/retune/live/state.py new file mode 100644 index 000000000000..c3bbdef4f898 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/state.py @@ -0,0 +1,1043 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Persisted live retune state and the leadership Lease. + +Two ``coordination.k8s.io/v1`` Leases per DGD, both owned by the DGD so the +garbage collector removes them with it: + +- ``-planner-retune``: leadership. ``LeaseLeadership`` renews it with + get -> check holder -> update(resourceVersion), takes over an empty holder + at once and a foreign holder only after observing it unchanged (holder and + resourceVersion) for a whole lease duration on the local monotonic clock. +- ``-planner-retune-state``: the state snapshot (``RetuneState``) as + JSON in two annotations. Every write replaces the snapshot and is + conditional on the resourceVersion read; the supervisor additionally + requires ``state.leader_term`` to equal its own term for every write but + the claim. + +The write-rule functions below return a new ``RetuneState`` and never touch +Kubernetes; ``StateStore`` reads and writes the snapshot. Everything here is +synchronous: callers on the event loop run it in a thread. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import logging +import math +import threading +import time +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, Literal, Optional + +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD, + LABEL_DGD_UID, + DgdRef, + FrontendStartup, + KubeAlreadyExists, + KubeConflict, + KubeError, + RetuneKube, +) + +if TYPE_CHECKING: + from dynamo.planner.config.planner_config import PlannerConfig + +logger = logging.getLogger(__name__) + +STATE_ANNOTATION = "nvidia.com/planner-retune-state" +HISTORY_ANNOTATION = "nvidia.com/planner-retune-history" +HISTORY_LIMIT = 50 +STATE_FORMAT_VERSION = 1 + +DGD_API_VERSION = "nvidia.com/v1beta1" +DGD_KIND = "DynamoGraphDeployment" +LEASE_API_VERSION = "coordination.k8s.io/v1" + +RetuneMode = Literal["active", "advisory"] +_MODES = ("active", "advisory") +_REL_TOL = 1e-9 +_ABS_TOL = 1e-12 + + +def leadership_lease_name(dgd_name: str) -> str: + return f"{dgd_name}-planner-retune" + + +def state_lease_name(dgd_name: str) -> str: + return f"{dgd_name}-planner-retune-state" + + +def find_retune_leases( + kube: RetuneKube, dgd_name: str, *, timeout_s: float +) -> tuple[str, ...]: + """Names of the DGD's live retune Leases that exist. + + A planner without retune config calls it to find what a removed live + retune left behind. Raises ``KubeError`` (403 without Lease RBAC). + """ + + return tuple( + name + for name in (leadership_lease_name(dgd_name), state_lease_name(dgd_name)) + if kube.get_lease(name, timeout_s=timeout_s) is not None + ) + + +def owner_reference(dgd: DgdRef) -> dict[str, Any]: + """ownerReference to the DGD; the GC deletes the Leases with it.""" + + return { + "apiVersion": DGD_API_VERSION, + "kind": DGD_KIND, + "name": dgd.name, + "uid": dgd.uid, + "blockOwnerDeletion": False, + } + + +def leader_term(lease_transitions: int, identity: str, term_seq: int) -> str: + """``//``. + + ``term_seq`` is bumped on every entry into initialization, so a write of + an earlier term (another process, or this one before re-initializing) + always fails the ``leader_term`` precondition. + """ + + return f"{int(lease_transitions)}/{identity}/{int(term_seq)}" + + +def micro_time(epoch_s: float) -> str: + """A Kubernetes ``MicroTime`` (RFC 3339 with microseconds, UTC).""" + + stamp = datetime.fromtimestamp(epoch_s, tz=timezone.utc) + return stamp.strftime("%Y-%m-%dT%H:%M:%S.%fZ") + + +def canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def maps_equal( + left: Optional[Mapping[str, Any]], right: Optional[Mapping[str, Any]] +) -> bool: + """Deep equality of parameter maps, floats within a tiny tolerance.""" + + if left is None or right is None: + return left is None and right is None + if set(left) != set(right): + return False + for key, value in left.items(): + other = right[key] + if isinstance(value, Mapping) or isinstance(other, Mapping): + if not ( + isinstance(value, Mapping) + and isinstance(other, Mapping) + and maps_equal(value, other) + ): + return False + elif isinstance(value, (int, float)) and isinstance(other, (int, float)): + if not math.isclose( + float(value), float(other), rel_tol=_REL_TOL, abs_tol=_ABS_TOL + ): + return False + elif value != other: + return False + return True + + +def fingerprint(config: PlannerConfig, startup: Optional[FrontendStartup]) -> str: + """What the deployed point was chosen under; a change resets it. + + Covers the router search domain as configured (its expansion is a pure + function of it and the startup), the parsed frontend startup, the pinned + DGDR engine section with its model, SKU and backend, the planner mode, + the objective and the scoring SLA. ``advisory`` is not part of it: the + state's ``mode`` records it. + """ + + retune = config.retune + if retune is None: + raise ValueError("fingerprint requires config.retune") + dgdr = retune.dgdr + sla_ttft_ms, sla_itl_ms = retune.scoring_sla(config.ttft_ms, config.itl_ms) + payload = { + "version": STATE_FORMAT_VERSION, + "router": retune.router, + "planner": retune.planner, + "objective": retune.objective, + "sla": [float(sla_ttft_ms), float(sla_itl_ms)], + "mode": config.mode, + "backend": config.backend, + "engine": None if dgdr is None else dgdr.engine, + "model": None if dgdr is None else dgdr.model, + "sku": None if dgdr is None else dgdr.sku, + "startup": None if startup is None else startup.fingerprint_source, + } + digest = hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest() + return digest[:16] + + +# ------------------------------------------------------------------ state + + +class StateFormatError(ValueError): + """The state annotations do not hold a readable snapshot. + + The snapshot is treated as corrupt: the next claim write replaces it. + """ + + +class StateVersionError(Exception): + """The snapshot has a format version this planner does not know. + + Not a ``StateFormatError``: a newer planner's state is never treated as + corrupt and overwritten. Initialization refuses to proceed until a + planner that reads the version leads, or the state Lease is deleted. + """ + + +def _finite_number(value: Any) -> bool: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return math.isfinite(value) + + +def _float_map(value: Any, what: str) -> dict[str, float]: + if not isinstance(value, Mapping): + raise StateFormatError(f"{what} must be an object") + out: dict[str, float] = {} + for key, item in value.items(): + if not _finite_number(item): + raise StateFormatError(f"{what}.{key} must be a finite number") + out[str(key)] = float(item) + return out + + +def _writable_entry(item: Any) -> bool: + """A history entry that ``canonical_json`` can write back: a JSON object + without NaN or infinities (``json.loads`` accepts both).""" + + if not isinstance(item, Mapping): + return False + try: + canonical_json(item) + except ValueError: + return False + return True + + +def _nested_map(value: Any, what: str) -> dict[str, dict[str, float]]: + if not isinstance(value, Mapping): + raise StateFormatError(f"{what} must be an object") + return {str(key): _float_map(item, f"{what}.{key}") for key, item in value.items()} + + +@dataclass(frozen=True) +class RetuneState: + """The persisted live retune state of one DGD. + + ``mode is None`` is "no state": only the claim (``dgd_uid`` and + ``leader_term``) has been written. ``by_wt`` is the deployed per worker + type 8-parameter map, ``flat`` the controller's router point (live + names), ``initials`` the latched initial configuration of each router. + ``history`` is newest last, at most ``HISTORY_LIMIT`` entries. + """ + + dgd_uid: str + leader_term: str = "" + fingerprint: Optional[str] = None + mode: Optional[RetuneMode] = None + rollback_pending: bool = False + deploy_seq: int = 0 + flat: dict[str, float] = field(default_factory=dict) + by_wt: dict[str, dict[str, float]] = field(default_factory=dict) + initials: dict[str, dict[str, float]] = field(default_factory=dict) + last_window_end_ms: Optional[int] = None + updated_ms: int = 0 + history: tuple[dict[str, Any], ...] = () + + @property + def has_deployment(self) -> bool: + """Whether an initialization has recorded what is deployed.""" + + return self.mode is not None + + def to_annotations(self) -> dict[str, str]: + body = { + "version": STATE_FORMAT_VERSION, + "dgd_uid": self.dgd_uid, + "leader_term": self.leader_term, + "fingerprint": self.fingerprint, + "mode": self.mode, + "rollback_pending": self.rollback_pending, + "deploy_seq": self.deploy_seq, + "flat": self.flat, + "by_wt": self.by_wt, + "initials": self.initials, + "last_window_end_ms": self.last_window_end_ms, + "updated_ms": self.updated_ms, + } + return { + STATE_ANNOTATION: canonical_json(body), + HISTORY_ANNOTATION: canonical_json(list(self.history)), + } + + @classmethod + def from_annotations( + cls, annotations: Optional[Mapping[str, str]] + ) -> Optional[RetuneState]: + """Parse the snapshot; ``None`` when the lease holds none. + + Raises ``StateVersionError`` for a snapshot of another format version + (checked first, so nothing else of it is judged) and + ``StateFormatError`` for a snapshot that cannot be read. + """ + + raw = (annotations or {}).get(STATE_ANNOTATION) + if raw is None: + return None + try: + body = json.loads(raw) + except ValueError as error: + raise StateFormatError(f"state annotation is not JSON: {error}") from None + if not isinstance(body, Mapping): + raise StateFormatError("state annotation must be a JSON object") + version = body.get("version") + if isinstance(version, bool) or not isinstance(version, int): + raise StateFormatError("state.version must be an integer") + if version != STATE_FORMAT_VERSION: + raise StateVersionError( + f"state format version {version} is not {STATE_FORMAT_VERSION}" + ) + dgd_uid = body.get("dgd_uid") + if not isinstance(dgd_uid, str) or not dgd_uid: + raise StateFormatError("state.dgd_uid must be a non-empty string") + mode = body.get("mode") + if mode is not None and mode not in _MODES: + raise StateFormatError(f"state.mode must be one of {_MODES} or null") + fingerprint = body.get("fingerprint") + if fingerprint is not None and not isinstance(fingerprint, str): + raise StateFormatError("state.fingerprint must be a string or null") + deploy_seq = body.get("deploy_seq", 0) + if isinstance(deploy_seq, bool) or not isinstance(deploy_seq, int): + raise StateFormatError("state.deploy_seq must be an integer") + last = body.get("last_window_end_ms") + if last is not None and not _finite_number(last): + raise StateFormatError("state.last_window_end_ms must be a finite number") + updated = body.get("updated_ms") + if updated is not None and not _finite_number(updated): + raise StateFormatError("state.updated_ms must be a finite number") + history_raw = (annotations or {}).get(HISTORY_ANNOTATION) + history: list[dict[str, Any]] = [] + if history_raw: + try: + parsed = json.loads(history_raw) + except ValueError: + # History is an audit trail; a damaged one never blocks. + logger.warning("retune state: unreadable history annotation dropped") + parsed = [] + if isinstance(parsed, list): + # A NaN in one entry would make every later write raise. + history = [dict(item) for item in parsed if _writable_entry(item)] + if len(history) < len(parsed): + logger.warning( + "retune state: %d malformed history entries dropped", + len(parsed) - len(history), + ) + return cls( + dgd_uid=dgd_uid, + leader_term=str(body.get("leader_term") or ""), + fingerprint=fingerprint, + mode=mode, + rollback_pending=bool(body.get("rollback_pending", False)), + deploy_seq=int(deploy_seq), + flat=_float_map(body.get("flat") or {}, "state.flat"), + by_wt=_nested_map(body.get("by_wt") or {}, "state.by_wt"), + initials=_nested_map(body.get("initials") or {}, "state.initials"), + last_window_end_ms=None if last is None else int(last), + updated_ms=0 if updated is None else int(updated), + history=tuple(history[-HISTORY_LIMIT:]), + ) + + +def deployed_off( + state: Optional[RetuneState], + startup_by_wt: Optional[Mapping[str, Mapping[str, float]]], +) -> bool: + """Whether the routers may be off their startup values. + + ``rollback_pending ∨ by_wt ≠ startup``, false without a state, with one + refinement: only an ``active`` record can have deployed anything. An + ``advisory`` record without ``rollback_pending`` was written by a planner + that only observed (or finished rolling back), so its ``by_wt`` is the + startup it saw, not a deployment; after a frontend startup change it + must not make an advisory planner reset routers it never wrote. With an + unknown startup (the frontend could not be parsed) an active record + counts as deployed: only the routers' latched initials can tell. + """ + + if state is None or not state.has_deployment: + return False + if state.rollback_pending: + return True + if state.mode != "active": + return False + if startup_by_wt is None: + return True + return not maps_equal(state.by_wt, startup_by_wt) + + +# ------------------------------------------------------------ write rules + + +def claim(state: Optional[RetuneState], *, dgd_uid: str, term: str) -> RetuneState: + """Claim write: only ``leader_term`` changes (a new record without state).""" + + if state is None: + return RetuneState(dgd_uid=dgd_uid, leader_term=term) + return dataclasses.replace(state, leader_term=term) + + +def _replace_deployment( + state: RetuneState, + *, + fingerprint: str, + mode: RetuneMode, + flat: Mapping[str, float], + by_wt: Mapping[str, Mapping[str, float]], + rollback_pending: bool, +) -> RetuneState: + return dataclasses.replace( + state, + fingerprint=fingerprint, + mode=mode, + flat={key: float(value) for key, value in flat.items()}, + by_wt={wt: dict(params) for wt, params in by_wt.items()}, + rollback_pending=bool(rollback_pending), + deploy_seq=state.deploy_seq + 1, + ) + + +def deploy( + state: RetuneState, + *, + fingerprint: str, + mode: RetuneMode, + flat: Mapping[str, float], + by_wt: Mapping[str, Mapping[str, float]], +) -> RetuneState: + """applied / promote: replace the deployment, clear rollback_pending.""" + + return _replace_deployment( + state, + fingerprint=fingerprint, + mode=mode, + flat=flat, + by_wt=by_wt, + rollback_pending=False, + ) + + +def rollback_start( + state: RetuneState, + *, + fingerprint: str, + mode: RetuneMode, + flat: Mapping[str, float], + by_wt: Mapping[str, Mapping[str, float]], +) -> RetuneState: + """Rollback start (advisory, or spec_ok false): startup, pending set.""" + + return _replace_deployment( + state, + fingerprint=fingerprint, + mode=mode, + flat=flat, + by_wt=by_wt, + rollback_pending=True, + ) + + +def reset( + state: RetuneState, + *, + fingerprint: str, + mode: RetuneMode, + flat: Mapping[str, float], + by_wt: Mapping[str, Mapping[str, float]], + rollback_pending: bool, +) -> RetuneState: + """Reset (fingerprint change, no state, mode switch): startup, + ``rollback_pending = deployed_off``.""" + + return _replace_deployment( + state, + fingerprint=fingerprint, + mode=mode, + flat=flat, + by_wt=by_wt, + rollback_pending=rollback_pending, + ) + + +def rollback_complete( + state: RetuneState, + *, + expected_deploy_seq: int, + entry: Optional[Mapping[str, Any]] = None, +) -> Optional[RetuneState]: + """Flip only ``rollback_pending`` to false. + + ``None`` (discard, never retry) when the deployment changed since the + rollback started (``deploy_seq``) or the rollback is already complete. + ``entry`` (a ``rolled_back`` record) is appended to the history. + """ + + if not state.rollback_pending or state.deploy_seq != expected_deploy_seq: + return None + new = dataclasses.replace(state, rollback_pending=False) + if entry is not None: + new = append_history(new, entry) + return new + + +def merge_initials( + state: RetuneState, + latched: Mapping[str, Mapping[str, float]], + *, + keep: Optional[Iterable[str]] = None, +) -> RetuneState: + """Merge newly latched initials; drop router ids not in ``keep``.""" + + initials = {rid: dict(cfg) for rid, cfg in state.initials.items()} + for router_id, configuration in latched.items(): + initials[str(router_id)] = {k: float(v) for k, v in configuration.items()} + if keep is not None: + kept = set(keep) + initials = {rid: cfg for rid, cfg in initials.items() if rid in kept} + return dataclasses.replace(state, initials=initials) + + +def append_history(state: RetuneState, entry: Mapping[str, Any]) -> RetuneState: + """Append one entry, deduplicated by ``decision_id``; keep the newest 50.""" + + decision_id = entry.get("decision_id") + history = [ + item + for item in state.history + if decision_id is None or item.get("decision_id") != decision_id + ] + history.append(dict(entry)) + return dataclasses.replace(state, history=tuple(history[-HISTORY_LIMIT:])) + + +def has_decision(state: RetuneState, decision_id: str) -> bool: + return any(item.get("decision_id") == decision_id for item in state.history) + + +def record_decision( + state: RetuneState, entry: Mapping[str, Any], window_end_ms: Optional[int] +) -> RetuneState: + """Every decision: history entry, ``last_window_end_ms = max(old, end)``.""" + + new = append_history(state, entry) + if window_end_ms is None: + return new + last = state.last_window_end_ms + return dataclasses.replace( + new, + last_window_end_ms=int(window_end_ms) + if last is None + else max(last, int(window_end_ms)), + ) + + +# ----------------------------------------------------------------- store + + +@dataclass(frozen=True) +class StateSnapshot: + """One read of the state Lease. ``lease is None``: it does not exist.""" + + state: Optional[RetuneState] + resource_version: Optional[str] + lease: Optional[dict[str, Any]] + # The annotations held a snapshot that could not be parsed. + corrupt: bool = False + + +class StateStore: + """Reads and conditionally writes ``-planner-retune-state``.""" + + def __init__( + self, + kube: RetuneKube, + dgd: DgdRef, + *, + timeout_s: float = 5.0, + clock: Callable[[], float] = time.time, + ) -> None: + self._kube = kube + self._dgd = dgd + self._timeout_s = float(timeout_s) + self._clock = clock + self.name = state_lease_name(dgd.name) + + @property + def dgd(self) -> DgdRef: + return self._dgd + + def read(self) -> StateSnapshot: + """Read the snapshot; an unreadable one comes back ``corrupt``. + + A snapshot of another format version raises ``StateVersionError``. + """ + + lease = self._kube.get_lease(self.name, timeout_s=self._timeout_s) + if lease is None: + return StateSnapshot(state=None, resource_version=None, lease=None) + metadata = lease.get("metadata") or {} + rv = metadata.get("resourceVersion") + try: + state = RetuneState.from_annotations(metadata.get("annotations")) + except StateFormatError as error: + logger.error("retune state lease %s is unreadable: %s", self.name, error) + return StateSnapshot( + state=None, resource_version=rv, lease=lease, corrupt=True + ) + return StateSnapshot(state=state, resource_version=rv, lease=lease) + + def write(self, state: RetuneState, snapshot: StateSnapshot) -> StateSnapshot: + """Replace the snapshot read as ``snapshot`` (create it if absent). + + A changed resourceVersion, or a lease created by someone else in the + meantime, raises ``KubeConflict``; other failures raise ``KubeError`` + (the write may still have landed: re-read before deciding). + """ + + state = dataclasses.replace(state, updated_ms=int(round(self._clock() * 1000))) + annotations = state.to_annotations() + if snapshot.lease is None: + body = self._new_body(annotations) + try: + created = self._kube.create_lease(body, timeout_s=self._timeout_s) + except KubeAlreadyExists as error: + raise KubeConflict(str(error)) from error + return StateSnapshot( + state=state, + resource_version=(created.get("metadata") or {}).get("resourceVersion"), + lease=created, + ) + body = json.loads(json.dumps(snapshot.lease)) + metadata = body.setdefault("metadata", {}) + metadata["resourceVersion"] = snapshot.resource_version + merged = dict(metadata.get("annotations") or {}) + merged.update(annotations) + metadata["annotations"] = merged + metadata["ownerReferences"] = [owner_reference(self._dgd)] + labels = dict(metadata.get("labels") or {}) + labels.update({LABEL_DGD: self._dgd.name, LABEL_DGD_UID: self._dgd.uid}) + metadata["labels"] = labels + replaced = self._kube.replace_lease(self.name, body, timeout_s=self._timeout_s) + return StateSnapshot( + state=state, + resource_version=(replaced.get("metadata") or {}).get("resourceVersion"), + lease=replaced, + ) + + def delete(self, snapshot: StateSnapshot) -> None: + """Delete the state Lease read or written as ``snapshot``. + + Conditional on its resourceVersion: a Lease changed since (another + term claimed it) raises ``KubeConflict`` and stays. + """ + + self._kube.delete_lease( + self.name, + resource_version=snapshot.resource_version, + timeout_s=self._timeout_s, + ) + + def _new_body(self, annotations: Mapping[str, str]) -> dict[str, Any]: + return { + "apiVersion": LEASE_API_VERSION, + "kind": "Lease", + "metadata": { + "name": self.name, + "namespace": self._kube.namespace, + "labels": {LABEL_DGD: self._dgd.name, LABEL_DGD_UID: self._dgd.uid}, + "annotations": dict(annotations), + "ownerReferences": [owner_reference(self._dgd)], + }, + "spec": {}, + } + + +# ------------------------------------------------------------ leadership + + +LeaseStep = Literal["acquired", "renewed", "standby", "lost", "error"] + + +class LeaseLeadership: + """Leadership over ``-planner-retune``. + + ``step`` runs one renewal/acquisition attempt and returns what happened: + + - ``acquired``: this process now leads (a new lease, an empty holder, an + expired foreign holder, or its own lease after it had lost track). + - ``renewed``: it led and still does. + - ``lost``: it led and another holder now owns the lease. + - ``standby``: another holder leads, or ``release`` has run. + - ``error``: the API call failed; leadership, if any, lasts until + ``leader_until``. + + ``is_leader`` holds from a successful write until ``leader_until`` = + send time + lease duration - ``margin_s``, all on the local monotonic + clock: a contender takes over only after watching the holder unchanged + for a whole lease duration, which starts after the send. ``margin_s`` + must be >= 0 and less than the lease duration. The caller must start + each renewal early enough for it to end before ``leader_until`` + (``renewal_wait_s``): a renewal whose read of the Lease returns at or + after ``leader_until`` reports ``acquired`` (a new term), never + ``renewed``. + """ + + def __init__( + self, + kube: RetuneKube, + dgd: DgdRef, + *, + identity: str, + lease_duration_s: float = 60.0, + margin_s: Optional[float] = None, + timeout_s: float = 5.0, + monotonic: Callable[[], float] = time.monotonic, + clock: Callable[[], float] = time.time, + ) -> None: + if not identity: + raise ValueError("identity must be non-empty") + if not (math.isfinite(lease_duration_s) and lease_duration_s > 0): + raise ValueError("lease_duration_s must be positive and finite") + self._kube = kube + self._dgd = dgd + self.identity = identity + self.name = leadership_lease_name(dgd.name) + self.lease_duration_s = float(lease_duration_s) + self._duration_int = max(1, int(math.ceil(lease_duration_s))) + if margin_s is None: + margin_s = 3.0 * timeout_s + if not 0.0 <= margin_s < self.lease_duration_s: + raise ValueError("margin_s must be >= 0 and less than lease_duration_s") + self.margin_s = float(margin_s) + self._timeout_s = float(timeout_s) + self._monotonic = monotonic + self._clock = clock + self._lock = threading.Lock() + self._held = False + self._leader_until = -math.inf + self._acquisitions = 0 + # Set by ``release``: a step that reads the Lease after it never leads. + self._released = False + self.transitions = 0 + # resourceVersion of this process's last successful write. + self.resource_version: Optional[str] = None + # (holderIdentity, resourceVersion) of a foreign lease, and when this + # process first saw it unchanged. + self._observed: Optional[tuple[str, Optional[str]]] = None + self._observed_since = 0.0 + + # ------------------------------------------------------------- queries + + def is_leader(self) -> bool: + with self._lock: + return self._held and self._monotonic() < self._leader_until + + @property + def held(self) -> bool: + with self._lock: + return self._held + + @property + def leader_until(self) -> float: + with self._lock: + return self._leader_until + + @property + def acquisitions(self) -> int: + """How many steps have reported ``acquired``.""" + + with self._lock: + return self._acquisitions + + def leads_since(self, acquisitions: int) -> bool: + """``is_leader`` with no ``acquired`` step since ``acquisitions``. + + Checked atomically with the ``leader_until`` a re-acquisition sets, + so a step still in flight on another thread cannot hide its lapse. + """ + + with self._lock: + return ( + self._held + and self._monotonic() < self._leader_until + and self._acquisitions == acquisitions + ) + + def renewal_wait_s( + self, *, interval_s: float, budget_s: float, floor_s: float = 0.0 + ) -> float: + """Time until the next ``step``. + + ``interval_s``, but while leading no later than ``budget_s`` before + ``leader_until`` (at once if that has passed), so that a renewal + taking up to ``budget_s`` ends while this process still leads. + Never less than ``floor_s``, which takes precedence. + """ + + now = self._monotonic() + with self._lock: + leading = self._held and now < self._leader_until + latest_s = self._leader_until - budget_s - now + return max(floor_s, min(interval_s, latest_s) if leading else interval_s) + + def set_owner(self, dgd: DgdRef) -> None: + """The DGD the Lease's ownerReference names from now on.""" + + self._dgd = dgd + + def drop(self) -> None: + """Forget leadership locally (no API call).""" + + with self._lock: + self._held = False + self._leader_until = -math.inf + + # ---------------------------------------------------------------- step + + def step(self) -> LeaseStep: + t_send = self._monotonic() + try: + lease = self._kube.get_lease(self.name, timeout_s=self._timeout_s) + except KubeError as error: + logger.warning( + "retune leadership: reading lease %s failed: %s", self.name, error + ) + return "error" + with self._lock: + if self._released: + # A step still running on another thread when release ran. + return "standby" + if lease is None: + return self._create(t_send) + + spec = lease.get("spec") or {} + metadata = lease.get("metadata") or {} + holder = spec.get("holderIdentity") or "" + rv = metadata.get("resourceVersion") + transitions = _int(spec.get("leaseTransitions")) + + if holder == self.identity: + # Expired locally counts as lost: renewing it starts a new term. + was_held = self.is_leader() + body = self._body(lease, transitions=transitions, acquire=not was_held) + try: + updated = self._kube.replace_lease( + self.name, body, timeout_s=self._timeout_s + ) + except KubeConflict: + # Changed since the read: re-read on the next step. + logger.info("retune leadership: renewal of %s conflicted", self.name) + return "error" + except KubeError as error: + logger.warning( + "retune leadership: renewing %s failed: %s", self.name, error + ) + return "error" + new_holder = (updated.get("spec") or {}).get("holderIdentity") or "" + if new_holder != self.identity: + return self._lose() + self._lead(t_send, transitions, updated, acquired=not was_held) + return "renewed" if was_held else "acquired" + + lost = self.held + if lost: + self.drop() + logger.warning( + "retune leadership: lease %s is now held by %r", + self.name, + holder or None, + ) + if not holder: + # Released by its holder: take it over at once. + return self._take_over(lease, t_send, transitions, lost) + observed = (holder, rv) + if self._observed != observed: + self._observed = observed + self._observed_since = t_send + return "lost" if lost else "standby" + duration = _positive(spec.get("leaseDurationSeconds")) or self.lease_duration_s + if t_send - self._observed_since < duration: + return "lost" if lost else "standby" + logger.info( + "retune leadership: holder %r of %s unchanged for %.0fs; taking over", + holder, + self.name, + t_send - self._observed_since, + ) + return self._take_over(lease, t_send, transitions, lost) + + def release(self) -> bool: + """Clear the holder if it is still this process (conditional update). + + Final: a ``step`` whose read of the Lease returns after this started + reports ``standby`` and writes nothing. A step that read the Lease + earlier can still write; this clears that write only if it lands + before this method's own read. + """ + + with self._lock: + self._released = True + held = self.held + self.drop() + if not held: + return False + try: + lease = self._kube.get_lease(self.name, timeout_s=self._timeout_s) + if lease is None: + return False + spec = lease.get("spec") or {} + if (spec.get("holderIdentity") or "") != self.identity: + return False + body = json.loads(json.dumps(lease)) + body_spec = body.setdefault("spec", {}) + body_spec["holderIdentity"] = None + body_spec["renewTime"] = micro_time(self._clock()) + self._kube.replace_lease(self.name, body, timeout_s=self._timeout_s) + logger.info("retune leadership: released lease %s", self.name) + return True + except KubeError as error: + logger.warning( + "retune leadership: releasing %s failed: %s", self.name, error + ) + return False + + # ------------------------------------------------------------ helpers + + def _lead( + self, + t_send: float, + transitions: int, + written: dict[str, Any], + *, + acquired: bool, + ) -> None: + with self._lock: + self._held = True + self._leader_until = t_send + self.lease_duration_s - self.margin_s + if acquired: + self._acquisitions += 1 + self.transitions = transitions + self.resource_version = (written.get("metadata") or {}).get("resourceVersion") + self._observed = None + + def _lose(self) -> LeaseStep: + lost = self.held + self.drop() + return "lost" if lost else "standby" + + def _create(self, t_send: float) -> LeaseStep: + now = micro_time(self._clock()) + body = { + "apiVersion": LEASE_API_VERSION, + "kind": "Lease", + "metadata": { + "name": self.name, + "namespace": self._kube.namespace, + "labels": {LABEL_DGD: self._dgd.name, LABEL_DGD_UID: self._dgd.uid}, + "ownerReferences": [owner_reference(self._dgd)], + }, + "spec": { + "holderIdentity": self.identity, + "leaseDurationSeconds": self._duration_int, + "acquireTime": now, + "renewTime": now, + "leaseTransitions": 0, + }, + } + try: + created = self._kube.create_lease(body, timeout_s=self._timeout_s) + except (KubeAlreadyExists, KubeConflict): + return self._lose() + except KubeError as error: + logger.warning( + "retune leadership: creating %s failed: %s", self.name, error + ) + return "error" + self._lead(t_send, 0, created, acquired=True) + logger.info( + "retune leadership: acquired new lease %s as %s", self.name, self.identity + ) + return "acquired" + + def _take_over( + self, lease: dict[str, Any], t_send: float, transitions: int, lost: bool + ) -> LeaseStep: + body = self._body(lease, transitions=transitions + 1, acquire=True) + try: + written = self._kube.replace_lease( + self.name, body, timeout_s=self._timeout_s + ) + except KubeConflict: + return "lost" if lost else "standby" + except KubeError as error: + logger.warning( + "retune leadership: taking over %s failed: %s", self.name, error + ) + return "error" + self._lead(t_send, transitions + 1, written, acquired=True) + logger.info( + "retune leadership: acquired lease %s as %s (transitions=%d)", + self.name, + self.identity, + transitions + 1, + ) + return "acquired" + + def _body( + self, lease: dict[str, Any], *, transitions: int, acquire: bool + ) -> dict[str, Any]: + body = json.loads(json.dumps(lease)) + metadata = body.setdefault("metadata", {}) + metadata["ownerReferences"] = [owner_reference(self._dgd)] + spec = body.setdefault("spec", {}) + now = micro_time(self._clock()) + spec["holderIdentity"] = self.identity + spec["leaseDurationSeconds"] = self._duration_int + spec["renewTime"] = now + spec["leaseTransitions"] = int(transitions) + if acquire or not spec.get("acquireTime"): + spec["acquireTime"] = now + return body + + +def _int(value: Any) -> int: + if isinstance(value, bool): + return 0 + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + +def _positive(value: Any) -> Optional[float]: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) if value > 0 else None diff --git a/components/src/dynamo/planner/retune/live/supervisor.py b/components/src/dynamo/planner/retune/live/supervisor.py new file mode 100644 index 000000000000..d69ece31fbaa --- /dev/null +++ b/components/src/dynamo/planner/retune/live/supervisor.py @@ -0,0 +1,2248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``LiveRetuneSupervisor``: leadership, persisted state and actuation. + +The retune controller decides exactly as in replay; the supervisor decides +when it may launch and harvest, persists every decision before anything is +applied, and owns the routers' parameters. + +Hooks the controller calls on the planner's event loop: + +- ``launch_gate(k, window)`` before both the launch and the busy branch: + ``not_leader``/``not_ready`` keep the window due, ``already_decided`` skips + a window the state already records, otherwise it hands out the DGD spec + replicas and a fresh abort event. +- ``harvest_allowed(generation)``: leader, initialized and the sweep's + controller generation is current. +- ``sweep_fn(request, *, cancel, abort)`` on the sweep thread: binds a + ``DgdrContext`` and runs ``dgdr.run_dgdr_sweep``; never raises. + +The host calls ``on_effects`` after every engine tick; it enqueues the tick's +decisions in the outbox bound to (term_seq, generation, fingerprint). + +Background tasks (all catch every exception; every Kubernetes call runs in a +thread with a timeout): + +- lease: renew or acquire ``-planner-retune`` every third of the lease + duration; while leading, early enough that a renewal using its whole I/O + budget still ends before local leadership does. Becoming leader, + re-acquiring after losing track, or a re-init request enters + initialization (claim the state with ``term_seq + 1``, then apply the + first matching row of ``plan_initialization``); the Lease is still renewed + on the same schedule while it runs, and a renewal that reports a lapse + aborts it. Losing leadership runs the local loss. +- control: re-read the DGD and re-validate the spec premises (rebuilding the + controller through ``host.replace_live_retune`` when the startup, + ``spec_ok`` or P4 change), then, as an initialized leader, the P11 canary + and one actuator round, and it starts one reconciler round. +- reconciler: one round at a time in its own task, never awaited by the + control loop, within ``reconcile_budget_s``. A round still running (a hung + filesystem or slow API calls) skips the next ones; the reconciler stops + before its next delete once this process is no longer the initialized + leader of the round's epoch or the budget has passed. +- outbox: write each decision, conditional on the state's ``leader_term`` + being this process's term and on the resourceVersion read; N consecutive + failures re-initialize. + +Local leadership loss, a failed outbox, and a controller rebuild share one +path: clear ``initialized``, stop the outbox writer and the in-flight +actuator/reconciler round, clear the outbox, abort the sweep and drop it +(``drop_pending``), and rewind the controller to the last persisted window, +so a due window can be relaunched or adopted before its deadline. + +The actuation target follows the state: ``rollback`` while +``rollback_pending`` (an advisory planner resets only routers whose initial +it latched); ``observe`` (never write) in advisory; ``rollback`` when a spec +premise fails; otherwise ``apply`` the state's ``by_wt``. Like a decision, a +router's latched initial is persisted before the router is moved off it, so +a process stopped at any point leaves only routers whose initial the state +records. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import dataclasses +import functools +import inspect +import logging +import math +import threading +import time +from collections import deque +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from concurrent.futures import Future +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, Union + +from dynamo.planner.config.planner_config import ( + LIVE_RETUNE_K8S_TIMEOUT_SECONDS, + live_retune_io_timeout_seconds, + live_retune_lease_margin_seconds, + live_retune_min_lease_seconds, +) +from dynamo.planner.plugins.orchestrator.engine_adapter import RetuneHooks +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD_UID, + ActuationReport, + ActuationTarget, + DgdRef, + FrontendStartup, + GateReason, + GateResult, + KubeConflict, + KubeError, + KubeNotFound, + NullMetrics, + PremiseReport, + RetuneKube, + RetuneMetricsSink, + SweepRegistry, +) +from dynamo.planner.retune.live.dgdr import run_dgdr_sweep +from dynamo.planner.retune.live.state import ( + LeaseLeadership, + LeaseStep, + RetuneMode, + RetuneState, + StateSnapshot, + StateStore, + StateVersionError, + append_history, + claim, + deploy, + deployed_off, + has_decision, + maps_equal, + merge_initials, + record_decision, + reset, + rollback_complete, + rollback_start, +) +from dynamo.planner.retune.sweep import SweepRequest, SweepResult + +if TYPE_CHECKING: + from dynamo.planner.config.planner_config import PlannerConfig + from dynamo.planner.core.types import PlannerEffects, RetuneDecision, RetuneInputs + +logger = logging.getLogger(__name__) + +_EPSILON = 1e-9 +# Claim-write attempts on resourceVersion conflicts before giving up the +# initialization (the next lease round retries with term_seq + 1). +_CLAIM_ATTEMPTS = 3 +# Conflict retries of one state write. +_WRITE_ATTEMPTS = 3 +# Remembered (k, reason) pairs of skipped launches (the metric counts each +# pair once). +_SKIP_MEMORY = 64 + +BlockingRunner = Callable[[Callable[[], Any]], Awaitable[Any]] + + +async def _to_thread(call: Callable[[], Any]) -> Any: + return await asyncio.to_thread(call) + + +# ---------------------------------------------------------------- types + + +@dataclass(frozen=True) +class SpecEvaluation: + """One evaluation of the DGD spec (frontend parse and spec premises). + + ``inputs`` is the live ``RetuneInputs``, ``None`` iff P4 fails. + ``spec_replicas`` maps worker type to the DGD spec replicas (``None`` + when the workers cannot be resolved). ``dgd`` carries the DGD's current + uid and the routers' dynamo namespace. + """ + + report: PremiseReport + startup: Optional[FrontendStartup] + inputs: Optional[RetuneInputs] + spec_replicas: Optional[dict[str, int]] + dgd: DgdRef + generation: Optional[int] = None + + +# DGD object -> evaluation; may raise (the round is then skipped). +PremiseEvaluator = Callable[[dict], SpecEvaluation] + + +class RouterResolver(Protocol): + """Resolves router points against the parsed startup ``KvRouterConfig``.""" + + def startup_flat(self, keys: Sequence[str]) -> dict[str, float]: + """The startup value of each searched key (the controller's start).""" + ... + + def by_wt(self, flat: Mapping[str, float]) -> dict[str, dict[str, float]]: + """``policy_parameters(wt)`` of ``startup ⊕ flat`` per worker type.""" + ... + + +ResolverFactory = Callable[[Optional[FrontendStartup]], Optional[RouterResolver]] + + +class ActuatorLike(Protocol): + async def reconcile( + self, + target: ActuationTarget, + *, + initials: Mapping[str, Mapping[str, float]], + expected_frontends: int, + ) -> ActuationReport: + ... + + def set_startup_maps(self, startup_maps: Mapping[str, Mapping[str, float]]) -> None: + ... + + +# (DgdRef, startup maps per worker type) -> actuator +ActuatorFactory = Callable[[DgdRef, dict[str, dict[str, float]]], ActuatorLike] +SweepRunner = Callable[..., SweepResult] + + +class LiveRetuneHost(Protocol): + """What the supervisor needs from ``NativePlannerBase``. + + ``replace_live_retune`` takes the planner's config lock and calls the + engine adapter's ``replace_retune``. The supervisor reaches the current + controller through ``host._engine.retune_controller`` and serializes its + calls on it with ``host._config_lock``; ``host._engine.retune_sweeps_done`` + tells it when the engine's sweep thread is idle. + """ + + async def replace_live_retune( + self, inputs: Optional[RetuneInputs], *, generation: int + ) -> None: + ... + + +@dataclass(frozen=True) +class SupervisorSettings: + """Timing and policy of one supervisor (from ``PlannerConfig``).""" + + interval_s: float + grace_s: float + timeout_s: float + phase_s: float = 0.0 + advisory: bool = False + promote_last_advisory: bool = False + worker_types: tuple[str, ...] = ("aggregated",) + # Searched router keys (live names); the controller's router_keys. + router_keys: tuple[str, ...] = () + lease_duration_s: float = 60.0 + actuator_interval_s: float = 15.0 + k8s_timeout_s: float = LIVE_RETUNE_K8S_TIMEOUT_SECONDS + write_failure_limit: int = 3 + canary_interval_s: float = 600.0 + stop_timeout_s: float = 10.0 + # How long one reconciler round may run before the reconciler is told to + # stop (a blocked filesystem call cannot be interrupted; later rounds + # wait for it instead of piling up threads). + reconcile_budget_s: float = 120.0 + # Legacy clean-up of a removed retune: roll back what it deployed, then + # delete the Requests and both Leases. + decommission: bool = False + + def __post_init__(self) -> None: + for name in ( + "interval_s", + "lease_duration_s", + "actuator_interval_s", + "reconcile_budget_s", + "k8s_timeout_s", + ): + value = getattr(self, name) + if not (math.isfinite(value) and value > 0): + raise ValueError(f"{name} must be positive and finite") + if self.write_failure_limit < 1: + raise ValueError("write_failure_limit must be >= 1") + minimum_s = live_retune_min_lease_seconds(self.k8s_timeout_s) + if self.lease_duration_s < minimum_s: + raise ValueError( + f"lease_duration_s {self.lease_duration_s:g}s is too short for " + f"k8s_timeout_s {self.k8s_timeout_s:g}s: it must be at least the " + f"lease margin ({self.lease_margin_s:g}s) + 2 x the I/O budget " + f"of one renewal ({self.io_timeout_s:g}s) = {minimum_s:g}s" + ) + + @property + def renew_interval_s(self) -> float: + return self.lease_duration_s / 3.0 + + @property + def lease_margin_s(self) -> float: + """How long before the Lease could expire local leadership ends.""" + + return live_retune_lease_margin_seconds(self.k8s_timeout_s) + + @property + def io_timeout_s(self) -> float: + """Backstop on one blocking call; the client applies its own.""" + + return live_retune_io_timeout_seconds(self.k8s_timeout_s) + + @classmethod + def from_config( + cls, config: PlannerConfig, *, router_keys: Sequence[str] = () + ) -> SupervisorSettings: + retune = config.retune + if retune is None or retune.dgdr is None: + raise ValueError("live retune requires retune.executor 'dgdr'") + dgdr = retune.dgdr + return cls( + interval_s=float(retune.interval_seconds), + grace_s=float(retune.grace_seconds), + timeout_s=float(retune.effective_timeout_seconds), + phase_s=float(dgdr.phase_seconds), + advisory=bool(config.advisory), + promote_last_advisory=bool(dgdr.promote_last_advisory), + worker_types=worker_types_for(config.mode), + router_keys=tuple(sorted(router_keys)), + lease_duration_s=float(dgdr.lease_duration_seconds), + actuator_interval_s=float(dgdr.actuator_interval_seconds), + ) + + +def worker_types_for(planner_mode: str) -> tuple[str, ...]: + return ("prefill", "decode") if planner_mode == "disagg" else ("aggregated",) + + +def due_window_end_ms( + now_s: float, *, phase_s: float, interval_s: float, grace_s: float +) -> int: + """End of the latest window whose launch time is at or before ``now_s``. + + Same arithmetic as ``RetuneController._due_window`` with ``s0 = phase``. + """ + + k = math.floor((now_s - phase_s - grace_s) / interval_s + _EPSILON) - 1 + return int(round((phase_s + (k + 1) * interval_s) * 1000.0)) + + +def current_window_end_ms(now_s: float, *, phase_s: float, interval_s: float) -> int: + """End of the window in progress at ``now_s`` (it contains ``now_s``). + + It ends after the due window (``due_window_end_ms``) and after every + window that ended but is not due yet (within the grace). + """ + + k = math.floor((now_s - phase_s) / interval_s + _EPSILON) + return int(round((phase_s + (k + 1) * interval_s) * 1000.0)) + + +def window_index(window_end_ms: int, *, phase_s: float, interval_s: float) -> int: + """The controller's ``k`` of the window ending at ``window_end_ms``.""" + + return int(round((window_end_ms / 1000.0 - phase_s) / interval_s)) - 1 + + +# ---------------------------------------------------------- initialization + +InitRow = Literal[ + "spec_not_ok", + "advisory_rollback", + "advisory_observe", + "active_restore", + "active_reset", + "promote", +] + + +@dataclass(frozen=True) +class InitPlan: + """What initialization writes and restores. + + ``state`` is the snapshot to persist (``write`` says whether it differs + from the claimed one); ``restore_flat`` is passed to the controller's + ``restore_desired`` (``None``: leave the controller alone). + """ + + row: InitRow + state: RetuneState + write: bool + restore_flat: Optional[dict[str, float]] + + +def plan_initialization( + state: RetuneState, + *, + spec_ok: bool, + advisory: bool, + fingerprint: str, + startup_flat: Optional[Mapping[str, float]], + startup_by_wt: Optional[Mapping[str, Mapping[str, float]]], + router_keys: Sequence[str], + promote: bool, + now_ms: int, + interval_ms: int, + first_window_end_ms: int, +) -> InitPlan: + """What a new term writes and restores; the first matching row wins. + + 1. ``spec_not_ok``: a spec premise fails. Start a rollback if the routers + may be off their startup (``deployed_off``), else record the startup. + 2. ``advisory_rollback``/``advisory_observe``: advisory. Roll back what + an active planner deployed, else only observe. + 3. ``active_restore``: an active, current (same mode and fingerprint), + complete deployment without a pending rollback is restored as is. + 4. ``promote``: with ``promote_last_advisory``, the last advisory + recommendation, if recent and of the same fingerprint, is deployed. + 5. ``active_reset``: otherwise deploy the startup point, with a rollback + pending when the routers may be off it. + + ``state`` is the snapshot after the claim write. A state without a + deployment record is "no state". The first initialization of a DGD also + persists ``last_window_end_ms = first_window_end_ms`` (the end of the + window in progress now): neither the due window nor the window whose + trace began mid-way (tracing enabled with the planner) is ever swept, + even after a restart. + """ + + mode: RetuneMode = "advisory" if advisory else "active" + off = deployed_off(state, startup_by_wt) + flat0 = {key: float(value) for key, value in (startup_flat or {}).items()} + by0 = {wt: dict(params) for wt, params in (startup_by_wt or {}).items()} + stale = ( + not state.has_deployment + or state.mode != mode + or state.fingerprint != fingerprint + ) + restore: Optional[dict[str, float]] = None + + if not spec_ok: + row: InitRow = "spec_not_ok" + if off: + new = rollback_start( + state, fingerprint=fingerprint, mode=mode, flat=flat0, by_wt=by0 + ) + elif stale: + new = reset( + state, + fingerprint=fingerprint, + mode=mode, + flat=flat0, + by_wt=by0, + rollback_pending=False, + ) + else: + new = state + elif advisory: + if off: + row = "advisory_rollback" + new = rollback_start( + state, fingerprint=fingerprint, mode=mode, flat=flat0, by_wt=by0 + ) + else: + row = "advisory_observe" + new = ( + reset( + state, + fingerprint=fingerprint, + mode=mode, + flat=flat0, + by_wt=by0, + rollback_pending=False, + ) + if stale + else state + ) + elif ( + not stale + and not state.rollback_pending + and set(state.flat) == set(router_keys) + and state.by_wt + ): + row = "active_restore" + new = state + restore = dict(state.flat) + else: + entry = ( + _promotable(state, fingerprint, router_keys, now_ms, interval_ms) + if promote + else None + ) + if entry is not None: + row = "promote" + flat = {key: float(value) for key, value in entry["flat"].items()} + by_wt = {wt: dict(params) for wt, params in entry["by_wt"].items()} + new = deploy( + state, fingerprint=fingerprint, mode="active", flat=flat, by_wt=by_wt + ) + new = append_history( + new, + { + "decision_id": f"promote-{entry.get('decision_id')}", + "outcome": "promoted", + "window_end_ms": entry.get("window_end_ms"), + "fingerprint": fingerprint, + "flat": flat, + "by_wt": by_wt, + "deploy_seq": new.deploy_seq, + "at_ms": now_ms, + }, + ) + restore = flat + else: + row = "active_reset" + new = reset( + state, + fingerprint=fingerprint, + mode="active", + flat=flat0, + by_wt=by0, + rollback_pending=off, + ) + restore = flat0 + if new.last_window_end_ms is None: + new = dataclasses.replace(new, last_window_end_ms=int(first_window_end_ms)) + return InitPlan(row=row, state=new, write=new != state, restore_flat=restore) + + +def _promotable( + state: RetuneState, + fingerprint: str, + router_keys: Sequence[str], + now_ms: int, + interval_ms: int, +) -> Optional[Mapping[str, Any]]: + """The last advisory recommendation, if it may be promoted. + + It must be the newest history entry, an ``advisory`` outcome of the same + fingerprint, at most one interval old, and cover every searched key. + """ + + if not state.history: + return None + entry = state.history[-1] + if entry.get("outcome") != "advisory" or entry.get("fingerprint") != fingerprint: + return None + window_end_ms = entry.get("window_end_ms") + if ( + not isinstance(window_end_ms, (int, float)) + or now_ms - window_end_ms > interval_ms + ): + return None + flat = entry.get("flat") + by_wt = entry.get("by_wt") + if not isinstance(flat, Mapping) or set(flat) != set(router_keys): + return None + if not isinstance(by_wt, Mapping) or not by_wt: + return None + return entry + + +# ------------------------------------------------------------------ outbox + + +@dataclass(frozen=True) +class OutboxEntry: + """One decision waiting to be persisted, bound to the term that made it.""" + + decision: RetuneDecision + router_parameters: Optional[dict[str, dict[str, float]]] + term_seq: int + generation: int + fingerprint: str + epoch: int + + +WriteOutcome = Literal["ok", "already", "discarded", "lost_term", "error"] +_ALREADY = object() +Mutation = Callable[[RetuneState], Union[RetuneState, None, object]] + + +class _NullRegistry: + def protect(self, window_end_ms: int, request_names: Any) -> None: + pass + + def release(self, window_end_ms: int) -> None: + pass + + def enqueue_delete(self, request_names: Any) -> None: + pass + + +# -------------------------------------------------------------- supervisor + + +class LiveRetuneSupervisor: + """Leadership, state, outbox, actuation and clean-up of live retune. + + Collaborators are injected so tests never need a cluster: ``kube`` (a + ``RetuneKube``), ``premise_evaluator`` (DGD -> ``SpecEvaluation``), + ``actuator_factory`` (a ``RouterActuator``), ``reconciler`` (a + ``Reconciler``, also the sweep registry), ``canary`` (P11), + ``sweep_runner``/``context_factory`` (``dgdr.run_dgdr_sweep`` and its + ``DgdrContext``), the clocks, ``sleep`` and ``blocking`` (how a blocking + call runs off the event loop). + """ + + def __init__( + self, + *, + settings: SupervisorSettings, + kube: RetuneKube, + evaluation: SpecEvaluation, + identity: str, + premise_evaluator: PremiseEvaluator, + actuator_factory: ActuatorFactory, + fingerprint_fn: Callable[[Optional[FrontendStartup]], str], + resolver_factory: Optional[ResolverFactory] = None, + reconciler: Optional[Any] = None, + metrics: Optional[RetuneMetricsSink] = None, + canary: Optional[Callable[[LiveRetuneSupervisor], bool]] = None, + sweep_runner: Optional[SweepRunner] = None, + context_factory: Optional[Callable[[LiveRetuneSupervisor], Any]] = None, + clock: Callable[[], float] = time.time, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[Any]] = asyncio.sleep, + blocking: Optional[BlockingRunner] = None, + leadership: Optional[LeaseLeadership] = None, + store: Optional[StateStore] = None, + ) -> None: + self._settings = settings + self._kube = kube + self._identity = identity + self._premise_evaluator = premise_evaluator + self._actuator_factory = actuator_factory + self._fingerprint_fn = fingerprint_fn + self._resolver_factory = resolver_factory + self._reconciler = reconciler + self._metrics: RetuneMetricsSink = ( + metrics if metrics is not None else NullMetrics() + ) + self._canary = canary + self._sweep_runner = sweep_runner + self._context_factory = context_factory + self._clock = clock + self._monotonic = monotonic + self._sleep = sleep + self._blocking: BlockingRunner = ( + blocking if blocking is not None else _to_thread + ) + + self._dgd = evaluation.dgd + self._leadership = leadership or LeaseLeadership( + kube, + evaluation.dgd, + identity=identity, + lease_duration_s=settings.lease_duration_s, + margin_s=settings.lease_margin_s, + timeout_s=settings.k8s_timeout_s, + monotonic=monotonic, + clock=clock, + ) + self._store = store or StateStore( + kube, evaluation.dgd, timeout_s=settings.k8s_timeout_s, clock=clock + ) + + # Spec evaluation and what is derived from it. + self._evaluation = evaluation + self._p14_ok: Optional[bool] = None + self._spec_ok = bool(evaluation.report.spec_ok) + self._fingerprint = fingerprint_fn(evaluation.startup) + self._resolver = self._make_resolver(evaluation.startup) + self._actuator = actuator_factory( + evaluation.dgd, self._actuator_maps(evaluation.startup) + ) + self._log_report(evaluation.report) + + # Leadership, term and the cached state snapshot. + self._term_seq = 0 + self._term: Optional[str] = None + self._snapshot: Optional[StateSnapshot] = None + self._initialized = False + self._reinit_requested = False + # Bumped by every local loss: results of work started before it are + # discarded. + self._epoch = 0 + + # Controller. + self._host: Optional[Any] = None + self._generation = 0 + self._controller_ok = evaluation.inputs is not None + + # Outbox. + self._outbox: deque[OutboxEntry] = deque() + self._write_failures = 0 + + # Health (P11, P12) and actuation. + self._canary_ok: Optional[bool] = None + self._canary_at: Optional[float] = None + self._last_report: Optional[ActuationReport] = None + # (deploy_seq, monotonic since) of a converged rollback. + self._rollback_since: Optional[tuple[int, float]] = None + + # Skipped launches already counted. + self._skips_seen: deque[tuple[int, str]] = deque(maxlen=_SKIP_MEMORY) + + # Tasks and locks. + self._init_lock = asyncio.Lock() + self._write_lock = asyncio.Lock() + # The lease loop and an initialization both step the leadership + # Lease; their GET/replace pairs must not race on its resourceVersion. + self._lease_step_lock = asyncio.Lock() + self._tasks: list[asyncio.Task[None]] = [] + self._outbox_task: Optional[asyncio.Task[None]] = None + self._round_task: Optional[asyncio.Task[None]] = None + # The reconciler round's task (it waits up to the budget) and its + # work (the reconciler call itself, done only once it has returned). + self._reconcile_task: Optional[asyncio.Task[None]] = None + self._reconcile_work: Optional[asyncio.Future[Any]] = None + self._lease_wake = asyncio.Event() + self._outbox_wake = asyncio.Event() + self._background = False + self._stopping = False + self._stopped = False + self.decommissioned = False + + # ------------------------------------------------------------ queries + + @property + def identity(self) -> str: + return self._identity + + @property + def generation(self) -> int: + return self._generation + + @property + def term(self) -> Optional[str]: + return self._term + + @property + def term_seq(self) -> int: + return self._term_seq + + @property + def initialized(self) -> bool: + return self._initialized + + @property + def state(self) -> Optional[RetuneState]: + return None if self._snapshot is None else self._snapshot.state + + @property + def spec_ok(self) -> bool: + return self._spec_ok + + @property + def fingerprint(self) -> str: + return self._fingerprint + + @property + def evaluation(self) -> SpecEvaluation: + return self._evaluation + + @property + def dgd(self) -> DgdRef: + return self._dgd + + @property + def kube(self) -> RetuneKube: + return self._kube + + @property + def registry(self) -> SweepRegistry: + reconciler = self._reconciler + if reconciler is not None and hasattr(reconciler, "protect"): + return reconciler + return _NullRegistry() + + @property + def actuator(self) -> ActuatorLike: + return self._actuator + + @property + def outbox_size(self) -> int: + return len(self._outbox) + + @property + def last_report(self) -> Optional[ActuationReport]: + return self._last_report + + @property + def reconcile_task(self) -> Optional[asyncio.Task[None]]: + """The latest reconciler round's task (done once the round returned + or its budget passed); tests await it after ``control_round``.""" + + return self._reconcile_task + + def is_leader(self) -> bool: + """Thread-safe; ``False`` once ``leader_until`` has passed.""" + + return self._leadership.is_leader() + + def ready(self) -> bool: + """Whether a due window may launch (spec, controller and health).""" + + return ( + self._spec_ok + and self._controller_ok + and self._evaluation.spec_replicas is not None + and self._canary_ok is True + and self._p12_ok() + and not self._settings.decommission + ) + + def retune_inputs(self) -> Optional[RetuneInputs]: + """Inputs for the engine's first controller; ``None`` if P4 fails.""" + + return self._evaluation.inputs + + def hooks(self) -> RetuneHooks: + """``RetuneHooks`` for ``OrchestratorEngineAdapter``.""" + + return RetuneHooks( + sweep_fn=self.sweep_fn, + launch_gate=self.launch_gate, + harvest_allowed=self.harvest_allowed, + phase_s=self._settings.phase_s, + ) + + def target(self) -> Optional[ActuationTarget]: + """What the actuator converges to now (``None``: nothing).""" + + state = self.state + if not self._initialized or state is None: + return None + if state.rollback_pending: + # Only an active planner resets a router it never latched + # (revision > 0) to learn its initial; an advisory planner leaves + # it alone. + return ActuationTarget( + mode="rollback", + include_not_ready=True, + learn_unknown=not ( + self._settings.advisory or self._settings.decommission + ), + ) + if self._settings.advisory or self._settings.decommission: + return ActuationTarget(mode="observe") + if not self._spec_ok: + return ActuationTarget(mode="rollback", include_not_ready=True) + if not state.by_wt: + return ActuationTarget(mode="observe") + return ActuationTarget( + mode="apply", + desired_by_wt={wt: dict(params) for wt, params in state.by_wt.items()}, + ) + + # ----------------------------------------------------- controller hooks + + def launch_gate(self, k: int, window: tuple[float, float, float]) -> GateResult: + """Whether window ``k`` may launch now; see the module docstring.""" + + reason: Optional[GateReason] + _, end_s, _ = window + state = self.state + if not self.is_leader(): + reason = "not_leader" + elif not self._initialized or state is None or not self.ready(): + reason = "not_ready" + elif ( + state.last_window_end_ms is not None + and int(round(end_s * 1000.0)) <= state.last_window_end_ms + ): + reason = "already_decided" + else: + reason = None + if reason is not None: + if (k, reason) not in self._skips_seen: + self._skips_seen.append((k, reason)) + self._metrics.skipped(reason) + logger.info("retune gate: window %d %s", k, reason) + return GateResult(reason=reason) + # Set by the controller's drop_pending (local loss, rebuild, stop). + return GateResult( + reason=None, + spec_replicas=dict(self._evaluation.spec_replicas or {}), + abort=threading.Event(), + ) + + def harvest_allowed(self, generation: int) -> bool: + """Read-only: leader, initialized, and the sweep's generation is current.""" + + return ( + self.is_leader() + and self._initialized + and int(generation) == self._generation + ) + + def sweep_fn( + self, + request: SweepRequest, + *, + cancel: threading.Event, + abort: Optional[threading.Event] = None, + ) -> SweepResult: + """Run one DGDR sweep on the controller's sweep thread; never raises.""" + + try: + runner = self._sweep_runner or run_dgdr_sweep + if self._context_factory is None: + raise RuntimeError("no DGDR sweep context factory is configured") + context = self._context_factory(self) + return runner(request, cancel=cancel, abort=abort, context=context) + except Exception as error: # the sweep contract: never raise + logger.exception("retune sweep of window %s failed", request.window_index) + return SweepResult( + status="failed", + error=f"{type(error).__name__}: {error}", + failure_reason="sweep_error", + ) + + def on_effects(self, effects: PlannerEffects) -> None: + """Enqueue the tick's decisions; never raises into the tick.""" + + decisions = list(effects.retune_decisions or ()) + if not decisions: + return + if not (self.is_leader() and self._initialized): + logger.warning( + "retune: dropping %d decision(s) %s: not an initialized leader", + len(decisions), + [decision.decision_id for decision in decisions], + ) + self._rewind_now() + return + router_parameters = effects.router_parameters + for decision in decisions: + self._outbox.append( + OutboxEntry( + decision=decision, + router_parameters=( + None + if router_parameters is None + else {wt: dict(p) for wt, p in router_parameters.items()} + ), + term_seq=self._term_seq, + generation=self._generation, + fingerprint=self._fingerprint, + epoch=self._epoch, + ) + ) + self._outbox_wake.set() + + # ------------------------------------------------------------ lifecycle + + async def start(self, host: Any, *, background: bool = True) -> None: + """Bind the planner host and start the background tasks. + + With ``background=False`` (tests) nothing runs on its own: drive + ``lease_round``, ``control_round`` and ``outbox_round`` directly. + ``control_round`` still starts the reconciler round as a task (it is + never waited for); await ``reconcile_task`` to finish it. + """ + + self._host = host + controller = self._controller() + if controller is not None: + self._generation = int(controller.generation) + self._controller_ok = controller is not None + self._background = background + self._publish() + if not background: + return + self._tasks = [ + asyncio.create_task(self._lease_loop(), name="planner-retune-lease"), + asyncio.create_task(self._control_loop(), name="planner-retune-control"), + ] + + async def stop(self) -> None: + """Abort the sweep, wait for it briefly and release the Lease. + + Runs before the engine shuts down. The local loss drops the + controller's sweep (``drop_pending`` sets its abort event); then the + engine's sweep thread is waited for, at most ``stop_timeout_s``, so + no sweep of this process creates a Request after the next leader + took over. + """ + + if self._stopped: + return + self._stopping = True + tasks = list(self._tasks) + self._tasks = [] + for task in tasks: + task.cancel() + if tasks: + await asyncio.wait(tasks) + try: + async with self._init_lock: + await self._local_loss_locked("shutdown") + except Exception: + logger.exception("retune supervisor: shutdown clean-up failed") + done = self._sweeps_done() + if done is not None: + try: + await asyncio.wait_for( + asyncio.wrap_future(done), timeout=self._settings.stop_timeout_s + ) + except asyncio.TimeoutError: + logger.warning( + "retune supervisor: a sweep is still running at shutdown" + ) + try: + await self._io(self._leadership.release) + except Exception: + logger.exception("retune supervisor: releasing the lease failed") + await self._close_actuator(self._actuator) + self._stopped = True + self._publish() + + # ------------------------------------------------------------- rounds + + async def lease_round(self) -> Optional[LeaseStep]: + """One renewal/acquisition; initialize or run the local loss. + + Returns the renewal/acquisition step, ``None`` once stopping. + """ + + if self._stopping: + return None + step = await self._lease_step() + try: + if self.is_leader(): + if ( + step == "acquired" + or self._reinit_requested + or not self._initialized + ): + async with self._init_lock: + if self.is_leader() and not self._stopping: + await self._initialize_locked() + elif self._initialized or self._term is not None: + async with self._init_lock: + if self._term is not None: + await self._local_loss_locked(f"not leader ({step})") + self._term = None + except Exception: + logger.exception("retune lease round: initialization failed") + self._publish() + return step + + async def control_round(self) -> None: + """Re-validate the spec; as an initialized leader, actuate and start + a reconciler round (never waited for).""" + + if self._stopping: + return + try: + await self.revalidate() + except Exception: + logger.exception("retune revalidation failed") + if self._initialized and self.is_leader(): + task = asyncio.create_task(self._leader_round()) + self._round_task = task + await asyncio.wait([task]) + if self._round_task is task: + self._round_task = None + if not task.cancelled() and task.exception() is not None: + logger.error("retune actuation round failed", exc_info=task.exception()) + self._publish() + + async def outbox_round(self) -> None: + """Write the queued decisions in order; stop at the first failure.""" + + while self._outbox and not self._stopping: + entry = self._outbox[0] + if not self._entry_current(entry): + self._outbox.popleft() + logger.info( + "retune outbox: dropped %s of an earlier term/generation", + entry.decision.decision_id, + ) + continue + outcome = await self._write_decision(entry) + if outcome in ("ok", "already", "discarded"): + if self._outbox and self._outbox[0] is entry: + self._outbox.popleft() + self._write_failures = 0 + continue + if outcome == "lost_term": + self._request_reinit("the state's leader_term is not ours") + return + # "error": counted as failures_total{reason="state_write"} by + # _write_state. + self._write_failures += 1 + logger.warning( + "retune outbox: writing %s failed (%d/%d)", + entry.decision.decision_id, + self._write_failures, + self._settings.write_failure_limit, + ) + if self._write_failures >= self._settings.write_failure_limit: + self._request_reinit("consecutive state write failures") + return + + async def revalidate(self) -> None: + """Re-read the DGD and rebuild on a spec-class change.""" + + try: + dgd_obj = await self._io(self._kube.get_dgd, self._dgd.name) + except KubeNotFound: + logger.warning("retune: DGD %s not found", self._dgd.name) + return + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune: reading DGD %s failed: %s", self._dgd.name, error) + return + evaluation = self._premise_evaluator(dgd_obj) + await self.apply_evaluation(evaluation) + + async def apply_evaluation(self, evaluation: SpecEvaluation) -> None: + """Adopt an evaluation; rebuild when the startup, ``spec_ok``, P4 or + the DGD identity changed.""" + + old = self._evaluation + startup_changed = _startup_key(old.startup) != _startup_key(evaluation.startup) + p14 = None if startup_changed else self._p14_ok + spec_ok = bool(evaluation.report.spec_ok) and p14 is not False + reasons = [] + if startup_changed: + reasons.append("startup changed") + if spec_ok != self._spec_ok: + reasons.append(f"spec_ok {self._spec_ok} -> {spec_ok}") + if bool(old.report.p4_ok) != bool(evaluation.report.p4_ok): + reasons.append(f"P4 {old.report.p4_ok} -> {evaluation.report.p4_ok}") + if (old.dgd.name, old.dgd.namespace, old.dgd.uid) != ( + evaluation.dgd.name, + evaluation.dgd.namespace, + evaluation.dgd.uid, + ): + reasons.append(f"DGD uid {old.dgd.uid} -> {evaluation.dgd.uid}") + if reasons: + await self.rebuild(evaluation, "; ".join(reasons)) + return + # Replicas, warnings or the routers' dynamo namespace (a worker + # rollout) change without a rebuild. + self._evaluation = evaluation + if old.dgd != evaluation.dgd: + logger.info( + "retune: routers' dynamo namespace %s -> %s", + old.dgd.dynamo_namespace, + evaluation.dgd.dynamo_namespace, + ) + await self._swap_actuator(evaluation) + + async def rebuild(self, evaluation: SpecEvaluation, reason: str) -> None: + """Local loss, generation + 1, ``replace_live_retune``, re-initialize.""" + + async with self._init_lock: + logger.info("retune: rebuilding the controller: %s", reason) + await self._local_loss_locked(f"rebuild: {reason}") + await self._install_evaluation(evaluation) + self._generation += 1 + host = self._host + if host is not None: + try: + await host.replace_live_retune( + evaluation.inputs, generation=self._generation + ) + self._controller_ok = evaluation.inputs is not None + except Exception: + logger.exception( + "retune: installing controller generation %d failed; not ready", + self._generation, + ) + self._controller_ok = False + else: + self._controller_ok = evaluation.inputs is not None + if self.is_leader() and not self._stopping: + await self._initialize_locked() + self._publish() + + # ------------------------------------------------- init and local loss + + async def _initialize_locked(self) -> bool: + """Claim a new term and initialize under ``_init_lock``. + + Confirms leadership, claims the state (only ``leader_term``), writes + what ``plan_initialization`` decides, restores the controller and + rewinds it to the last persisted window. ``True`` once initialized. + The leadership Lease is renewed meanwhile; a renewal that reports a + lapse (``acquired``) aborts the initialization. + """ + + await self._local_loss_locked("initialize") + # Confirm on the API server right before the claim: a holder change + # this process has not seen yet (it believes it leads until + # leader_until) aborts the initialization instead of claiming. + started = self._monotonic() + step = await self._lease_step() + if not self.is_leader(): + return False + acquisitions = self._leadership.acquisitions + keepalive = asyncio.create_task( + self._renew_while_initializing(step, started), + name="planner-retune-init-lease", + ) + try: + return await self._claim_locked(acquisitions) + finally: + await self._end_renewals(keepalive) + + async def _claim_locked(self, acquisitions: int) -> bool: + """``_initialize_locked`` after the confirming renewal; gives up once + a local loss or a re-acquisition since ``acquisitions`` made the new + term stale.""" + + self._term_seq += 1 + self._reinit_requested = False + term = f"{self._leadership.transitions}/{self._identity}/{self._term_seq}" + self._term = term + epoch = self._epoch + settings = self._settings + + def current() -> bool: + return epoch == self._epoch and self._leadership.leads_since(acquisitions) + + async with self._write_lock: + snapshot = await self._read_own_snapshot() + if snapshot is None: + return False + # Claim: only leader_term, conditional on the resourceVersion. + for attempt in range(_CLAIM_ATTEMPTS): + if not current(): + return False + claimed = claim(snapshot.state, dgd_uid=self._dgd.uid, term=term) + try: + snapshot = await self._io(self._store.write, claimed, snapshot) + break + except KubeConflict: + snapshot = await self._read_own_snapshot() + if snapshot is None: + return False + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune init: claim write failed: %s", error) + self._metrics.failure("state_write") + return False + else: + logger.warning("retune init: claim write kept conflicting") + self._metrics.failure("state_write") + return False + assert snapshot.state is not None + now_s = self._clock() + plan = plan_initialization( + snapshot.state, + spec_ok=self._spec_ok, + advisory=settings.advisory or settings.decommission, + fingerprint=self._fingerprint, + startup_flat=self._startup_flat(), + startup_by_wt=self._startup_by_wt(), + router_keys=self._router_keys(), + promote=settings.promote_last_advisory, + now_ms=int(round(now_s * 1000.0)), + interval_ms=int(round(settings.interval_s * 1000.0)), + # First enable: skip the due window and the window in + # progress, whose trace may have begun mid-way. + first_window_end_ms=current_window_end_ms( + now_s, phase_s=settings.phase_s, interval_s=settings.interval_s + ), + ) + if plan.write: + if not current(): + return False + try: + snapshot = await self._io(self._store.write, plan.state, snapshot) + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune init: %s write failed: %s", plan.row, error) + self._metrics.failure("state_write") + return False + self._snapshot = snapshot + logger.info( + "retune init: term %s row %s (mode=%s rollback_pending=%s deploy_seq=%s " + "last_window_end_ms=%s spec_ok=%s)", + term, + plan.row, + plan.state.mode, + plan.state.rollback_pending, + plan.state.deploy_seq, + plan.state.last_window_end_ms, + self._spec_ok, + ) + controller = self._controller() + async with self._host_lock(): + if plan.restore_flat is not None and controller is not None: + try: + controller.restore_desired(plan.restore_flat) + except Exception: + logger.exception("retune init: restore_desired failed") + return False + if controller is not None: + self._rewind(controller) + if not current(): + return False + self._initialized = True + self._rollback_since = None + if self._background: + self._outbox_task = asyncio.create_task( + self._outbox_loop(), name="planner-retune-outbox" + ) + self._outbox_wake.set() + return True + + async def _read_own_snapshot(self) -> Optional[StateSnapshot]: + """Read the state; drop a snapshot of another DGD uid. + + A recreated DGD of the same name finds the old uid's state: it is + deleted with that uid's Requests. A snapshot of an unknown format + version is never overwritten: the read fails (``None``). + """ + + try: + snapshot = await self._io(self._store.read) + if snapshot.state is not None and snapshot.state.dgd_uid != self._dgd.uid: + logger.warning( + "retune init: state belongs to DGD uid %s, not %s; recreating it", + snapshot.state.dgd_uid, + self._dgd.uid, + ) + await self._io(self._store.delete, snapshot) + await self._delete_foreign_requests(snapshot.state.dgd_uid) + snapshot = StateSnapshot(state=None, resource_version=None, lease=None) + elif snapshot.corrupt: + # Replaced by the claim (conditional on its resourceVersion). + snapshot = dataclasses.replace(snapshot, state=None) + return snapshot + except StateVersionError as error: + logger.error( + "retune init: the state Lease %s was written by another planner " + "version (%s); not initializing until a planner that reads it " + "leads or the Lease is deleted", + self._store.name, + error, + ) + self._metrics.failure("state_version") + return None + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune init: reading the state failed: %s", error) + return None + + async def _delete_foreign_requests(self, dgd_uid: str) -> None: + try: + requests = await self._io( + self._kube.list_requests, f"{LABEL_DGD_UID}={dgd_uid}" + ) + except (KubeError, asyncio.TimeoutError) as error: + logger.warning( + "retune: listing Requests of DGD uid %s failed: %s", dgd_uid, error + ) + return + names = [(request.get("metadata") or {}).get("name") for request in requests] + self.registry.enqueue_delete([name for name in names if name]) + + async def _local_loss_locked(self, reason: str) -> None: + """Local leadership loss under ``_init_lock``. + + Clears ``initialized``, stops the outbox writer and the in-flight + actuator/reconciler round, drops the queued decisions and the sweep, + and rewinds the controller to the last persisted window. + """ + + was = self._initialized or bool(self._outbox) or self._term is not None + self._initialized = False + self._epoch += 1 + self._rollback_since = None + # Health is re-established by the next term's first actuator round. + self._last_report = None + await self._cancel(self._outbox_task) + self._outbox_task = None + await self._cancel(self._round_task) + self._round_task = None + # The reconciler stops before its next delete (its should_continue + # checks the epoch); _reconcile_work stays pending until it returns. + await self._cancel(self._reconcile_task) + self._reconcile_task = None + dropped = len(self._outbox) + self._outbox.clear() + self._write_failures = 0 + controller = self._controller() + if controller is not None: + async with self._host_lock(): + try: + controller.drop_pending() + except Exception: + logger.exception("retune: drop_pending failed") + self._rewind(controller) + if was: + logger.info( + "retune: local leadership loss (%s); dropped %d queued decision(s)", + reason, + dropped, + ) + + def _request_reinit(self, reason: str) -> None: + logger.warning("retune: re-initializing: %s", reason) + self._initialized = False + self._reinit_requested = True + self._lease_wake.set() + + # ------------------------------------------------------ state writes + + async def _write_state(self, mutate: Mutation, what: str) -> WriteOutcome: + """Conditional write (term and resourceVersion) of ``mutate(state)``. + + ``mutate`` returns the new state, ``None`` to discard, or + ``_ALREADY`` when the server state already has it. A timeout is + resolved by re-reading: a write that landed is a success. Every + failed write (``error``) counts as ``failures_total{state_write}``. + """ + + outcome = await self._conditional_write(mutate, what) + if outcome == "error": + self._metrics.failure("state_write") + return outcome + + async def _conditional_write(self, mutate: Mutation, what: str) -> WriteOutcome: + epoch = self._epoch + async with self._write_lock: + snapshot = self._snapshot + for attempt in range(_WRITE_ATTEMPTS): + if epoch != self._epoch or not self.is_leader(): + return "discarded" + if snapshot is None or snapshot.state is None: + snapshot = await self._reread() + if snapshot is None or snapshot.state is None: + return "error" + state = snapshot.state + if state.leader_term != self._term: + return "lost_term" + new = mutate(state) + if new is _ALREADY: + self._snapshot = snapshot + return "already" + if new is None: + self._snapshot = snapshot + return "discarded" + assert isinstance(new, RetuneState) + try: + written = await self._io(self._store.write, new, snapshot) + except KubeConflict: + snapshot = await self._reread() + continue + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune state write (%s) failed: %s", what, error) + snapshot = await self._reread() + if ( + snapshot is not None + and snapshot.state is not None + and snapshot.state.leader_term == self._term + and mutate(snapshot.state) is _ALREADY + ): + self._snapshot = snapshot + return "ok" + return "error" + if epoch != self._epoch: + return "discarded" + self._snapshot = written + return "ok" + return "error" + + async def _reread(self) -> Optional[StateSnapshot]: + try: + snapshot = await self._io(self._store.read) + except (KubeError, asyncio.TimeoutError, StateVersionError) as error: + logger.warning("retune: re-reading the state failed: %s", error) + return None + if snapshot.state is not None and snapshot.state.dgd_uid != self._dgd.uid: + return StateSnapshot(state=None, resource_version=None, lease=None) + self._snapshot = snapshot + return snapshot + + async def _write_decision(self, entry: OutboxEntry) -> WriteOutcome: + decision = entry.decision + history = self._history_entry(entry) + window_end_ms = history["window_end_ms"] + applied = decision.outcome == "applied" + by_wt = entry.router_parameters + + def mutate(state: RetuneState) -> Union[RetuneState, None, object]: + if not self._entry_current(entry): + return None + if has_decision(state, decision.decision_id): + return _ALREADY + new = record_decision(state, history, window_end_ms) + if applied: + if not by_wt: + logger.error( + "retune: applied decision %s has no router parameters", + decision.decision_id, + ) + return new + new = deploy( + new, + fingerprint=entry.fingerprint, + mode="active", + flat=dict(decision.chosen.router), + by_wt=by_wt, + ) + return new + + outcome = await self._write_state(mutate, f"decision {decision.decision_id}") + if outcome in ("ok", "already"): + self._count_decision(decision) + logger.info( + "retune: persisted decision %s (%s)%s", + decision.decision_id, + decision.outcome, + " - deploying" if applied else "", + ) + return outcome + + def _count_decision(self, decision: RetuneDecision) -> None: + self._metrics.decision(decision.outcome) + if decision.outcome in ("failed", "timeout"): + self._metrics.failure(decision.failure_reason or decision.outcome) + details = decision.details + # The DGDR executor reports the window cut under details["window"]. + window = details.get("window") + dropped = ( + window.get("dropped_members") + if isinstance(window, Mapping) + else details.get("dropped_members") + ) + if isinstance(dropped, int) and dropped > 0: + self._metrics.dropped_members(dropped) + if decision.launched_at_s is not None: + missing = details.get("missing_candidates", 0) + self._metrics.scores( + baseline=_finite(decision.baseline_score), + best=_finite(decision.best_score), + sweep_seconds=_finite(decision.sweep_wall_s), + missing_candidates=int(missing) if isinstance(missing, int) else 0, + ) + + def _history_entry(self, entry: OutboxEntry) -> dict[str, Any]: + decision = entry.decision + flat = {key: float(value) for key, value in decision.chosen.router.items()} + item: dict[str, Any] = { + "decision_id": decision.decision_id, + "window_end_ms": int(round(decision.window_end_s * 1000.0)), + "outcome": decision.outcome, + "failure_reason": decision.failure_reason, + "reason": decision.reason, + "fingerprint": entry.fingerprint, + "baseline": _finite(decision.baseline_score), + "best": _finite(decision.best_score), + "flat": flat, + "decided_ms": int(round(decision.decided_at_s * 1000.0)), + } + if decision.outcome == "applied" and entry.router_parameters is not None: + item["by_wt"] = entry.router_parameters + elif decision.outcome == "advisory": + item["by_wt"] = self._resolve_by_wt(flat) + details = decision.details + image = { + key: details[key] + for key in ("runtime_image", "runtime_version", "runtime_version_override") + if isinstance(details.get(key), str) + } + if image: + item["image"] = image + return item + + def _entry_current(self, entry: OutboxEntry) -> bool: + return ( + self._initialized + and self.is_leader() + and entry.epoch == self._epoch + and entry.term_seq == self._term_seq + and entry.generation == self._generation + and entry.fingerprint == self._fingerprint + ) + + # ----------------------------------------------------------- actuation + + async def _leader_round(self) -> None: + epoch = self._epoch + if self._settings.decommission: + await self._actuate() + await self._finish_decommission(epoch) + return + await self._check_canary() + await self._actuate() + # Off the control loop's critical path: revalidation and actuation + # never wait for the reconciler. + self._start_reconcile(epoch) + + async def _actuate(self) -> None: + target = self.target() + if target is None: + return + epoch = self._epoch + report = await self._reconcile_routers(target) + if report is None: + return + outcome = await self._persist_initials(report) + if ( + outcome == "ok" + and target.mode == "apply" + and report.new_initials + and not report.converged + and epoch == self._epoch + ): + # The actuator moves a router off its initial only once the state + # records that initial: a process stopped between the write and + # the state write would leave a router whose initial nobody + # knows, which an advisory or clean-up rollback never resets. + # The initials learned this round are persisted now, so the + # writes held for them go out in this round, not the next. + report = await self._reconcile_routers(target) + if report is None: + return + await self._persist_initials(report) + await self._track_rollback(target, report) + + async def _reconcile_routers( + self, target: ActuationTarget + ) -> Optional[ActuationReport]: + """One actuator pass against the persisted initials; ``None`` when + there is no state or a local loss overtook it.""" + + state = self.state + if state is None: + return None + epoch = self._epoch + startup = self._evaluation.startup + report = await self._actuator.reconcile( + target, + initials=state.initials, + expected_frontends=max(1, startup.replicas if startup is not None else 1), + ) + if epoch != self._epoch: + return None + self._last_report = report + if startup is not None and report.parity_ok is not None: + if report.parity_ok != self._p14_ok: + logger.info( + "retune: P14 parity %s -> %s", self._p14_ok, report.parity_ok + ) + self._p14_ok = report.parity_ok + return report + + async def _persist_initials( + self, report: ActuationReport + ) -> Optional[WriteOutcome]: + """Merge the newly latched initials into the state and drop those of + routers that are gone; ``None`` when nothing changes.""" + + state = self.state + if state is None: + return None + merged = merge_initials(state, report.new_initials, keep=report.seen_router_ids) + if maps_equal(merged.initials, state.initials): + return None + latched = dict(report.new_initials) + keep = frozenset(report.seen_router_ids) + + def mutate(current: RetuneState) -> Union[RetuneState, None, object]: + new = merge_initials(current, latched, keep=keep) + return _ALREADY if maps_equal(new.initials, current.initials) else new + + outcome = await self._write_state(mutate, "initials") + if outcome == "lost_term": + self._request_reinit("the state's leader_term is not ours") + return outcome + + async def _track_rollback( + self, target: ActuationTarget, report: ActuationReport + ) -> None: + """Write "rollback complete" after a lease duration of convergence.""" + + state = self.state + if state is None or target.mode != "rollback" or not state.rollback_pending: + self._rollback_since = None + return + if not report.converged: + self._rollback_since = None + return + now = self._monotonic() + since = self._rollback_since + if since is None or since[0] != state.deploy_seq: + self._rollback_since = (state.deploy_seq, now) + return + if now - since[1] < self._settings.lease_duration_s: + return + expected = since[0] + entry = { + "decision_id": f"rolled_back-{expected}", + "outcome": "rolled_back", + "deploy_seq": expected, + "fingerprint": self._fingerprint, + "at_ms": int(round(self._clock() * 1000.0)), + } + + def mutate(current: RetuneState) -> Union[RetuneState, None, object]: + return rollback_complete(current, expected_deploy_seq=expected, entry=entry) + + outcome = await self._write_state(mutate, "rollback complete") + self._rollback_since = None + if outcome == "ok": + logger.info("retune: rollback complete (deploy_seq %d)", expected) + elif outcome == "lost_term": + self._request_reinit("the state's leader_term is not ours") + + async def _check_canary(self) -> None: + """P11: a dry-run Strict Request with ``workload.trace``.""" + + canary = self._canary + if canary is None: + self._canary_ok = True + return + now = self._monotonic() + if ( + self._canary_ok + and self._canary_at is not None + and now - self._canary_at < self._settings.canary_interval_s + ): + return + try: + ok = bool(await self._io(canary, self)) + except Exception as error: + logger.warning("retune canary (P11) failed: %s", error) + ok = False + if ok != self._canary_ok: + logger.info("retune canary (P11): %s", "ok" if ok else "failing") + self._canary_ok = ok + self._canary_at = now + + def _reconcile_allowed(self, epoch: int) -> bool: + """Only the initialized leader of ``epoch`` deletes anything. + + Called from the reconciler's thread too (``is_leader`` is + thread-safe, the rest are single reads). + """ + + return ( + not self._stopping + and self._initialized + and epoch == self._epoch + and self.is_leader() + ) + + def reconcile_running(self) -> bool: + """Whether a reconciler round (its task or its call) is still running.""" + + task, work = self._reconcile_task, self._reconcile_work + return (task is not None and not task.done()) or ( + work is not None and not work.done() + ) + + def _start_reconcile(self, epoch: int) -> None: + """Start one reconciler round in its own task, unless one still runs.""" + + if self._reconciler is None or not self._reconcile_allowed(epoch): + return + if self.reconcile_running(): + logger.debug("retune reconciler: the previous round is still running") + return + self._reconcile_task = asyncio.create_task( + self._reconcile(epoch), name="planner-retune-reconcile" + ) + + async def _reconcile(self, epoch: int) -> None: + """One reconciler round, waited for at most ``reconcile_budget_s``. + + The reconciler calls ``should_continue`` before every delete or + unlink and stops once it is ``False``: this process is no longer the + initialized leader of ``epoch``, or the budget has passed. A blocking + call cannot be cancelled, so ``_reconcile_work`` stays pending until + it returns and no other round starts meanwhile. + """ + + reconciler = self._reconciler + if reconciler is None or not self._reconcile_allowed(epoch): + return + state = self.state + settings = self._settings + expired = threading.Event() + + def should_continue() -> bool: + return not expired.is_set() and self._reconcile_allowed(epoch) + + call = functools.partial( + reconciler.run_once, + last_window_end_ms=None if state is None else state.last_window_end_ms, + due_window_end_ms=due_window_end_ms( + self._clock(), + phase_s=settings.phase_s, + interval_s=settings.interval_s, + grace_s=settings.grace_s, + ), + deadline_for=self._deadline_for, + frontend_pods=getattr(self._actuator, "last_pods", None), + should_continue=should_continue, + ) + cancellable = inspect.iscoroutinefunction(reconciler.run_once) + work: asyncio.Future[Any] = asyncio.ensure_future( + call() if cancellable else self._blocking(call) + ) + work.add_done_callback(_log_reconcile_failure) + self._reconcile_work = work + try: + await asyncio.wait_for( + asyncio.shield(work), timeout=settings.reconcile_budget_s + ) + except asyncio.TimeoutError: + logger.warning( + "retune reconciler: the round exceeded its %.0fs budget; it " + "stops before its next delete, later rounds wait for it", + settings.reconcile_budget_s, + ) + except asyncio.CancelledError: + raise + except Exception: + pass # logged by _log_reconcile_failure + finally: + if not work.done(): + expired.set() + if cancellable: + work.cancel() + + def _deadline_for(self, window_end_ms: int) -> float: + """``D_k = L_k + timeout`` in epoch seconds.""" + + settings = self._settings + return window_end_ms / 1000.0 + settings.grace_s + settings.timeout_s + + async def _finish_decommission(self, epoch: int) -> None: + """Once nothing is left to roll back, delete Requests and Leases. + + Only as the initialized leader of ``epoch``, checked again before + every delete: a lost term (a failed state write of this round) or + lost leadership stops it, so it never deletes under another leader. + The local clock is not the only fence: each Lease delete is + conditional on the resourceVersion this process last read or wrote, + so it fails once another term has claimed the state or taken over + the leadership Lease. + """ + + snapshot = self._snapshot + if ( + self.decommissioned + or snapshot is None + or snapshot.state is None + or snapshot.state.rollback_pending + ): + return + if not self._reconcile_allowed(epoch): + return + dgd = self._dgd + try: + requests = await self._io( + self._kube.list_requests, f"{LABEL_DGD_UID}={dgd.uid}" + ) + for request in requests: + name = (request.get("metadata") or {}).get("name") + if not name: + continue + if not self._reconcile_allowed(epoch): + return + await self._io(self._kube.delete_request, name) + if not self._reconcile_allowed(epoch): + return + await self._io(self._store.delete, snapshot) + except KubeConflict: + logger.warning( + "retune clean-up: the state changed since this term read it; " + "not deleting it" + ) + self._request_reinit("the state changed under the clean-up") + return + except (KubeError, asyncio.TimeoutError) as error: + logger.warning("retune clean-up: deleting Requests/state failed: %s", error) + return + if not self._reconcile_allowed(epoch): + return + try: + await self._io( + self._kube.delete_lease, + self._leadership.name, + resource_version=self._leadership.resource_version, + ) + except (KubeError, asyncio.TimeoutError) as error: + # A KubeConflict: the Lease changed since this process wrote it. + # It stays for the DGD's garbage collection or a later clean-up. + logger.warning("retune clean-up: deleting the lease failed: %s", error) + self._leadership.drop() + self._initialized = False + self.decommissioned = True + self._stopping = True + logger.info("retune clean-up of DGD %s finished", dgd.name) + + # ------------------------------------------------------------- helpers + + async def _install_evaluation(self, evaluation: SpecEvaluation) -> None: + old = self._evaluation + startup_changed = _startup_key(old.startup) != _startup_key(evaluation.startup) + if startup_changed: + self._p14_ok = None + self._evaluation = evaluation + self._spec_ok = bool(evaluation.report.spec_ok) and self._p14_ok is not False + self._fingerprint = self._fingerprint_fn(evaluation.startup) + if startup_changed: + self._resolver = self._make_resolver(evaluation.startup) + if old.dgd != evaluation.dgd: + if old.dgd.uid != evaluation.dgd.uid: + # Before the re-initialization deletes the old uid's + # Requests: the reconciler must treat the new uid's Requests + # as its own and the old ones as another DGD's. + self._set_reconciler_dgd(evaluation.dgd) + self._store = StateStore( + self._kube, + evaluation.dgd, + timeout_s=self._settings.k8s_timeout_s, + clock=self._clock, + ) + self._leadership.set_owner(evaluation.dgd) + await self._swap_actuator(evaluation) + elif startup_changed: + self._actuator.set_startup_maps(self._actuator_maps(evaluation.startup)) + self._last_report = None + self._log_report(evaluation.report) + + def _set_reconciler_dgd(self, dgd: DgdRef) -> None: + setter = getattr(self._reconciler, "set_dgd", None) + if setter is None: + return + try: + setter(dgd) + except Exception: + logger.exception( + "retune: moving the reconciler to DGD uid %s failed", dgd.uid + ) + + async def _swap_actuator(self, evaluation: SpecEvaluation) -> None: + """A new actuator for a new DGD identity or dynamo namespace.""" + + self._dgd = evaluation.dgd + previous = self._actuator + self._actuator = self._actuator_factory( + evaluation.dgd, self._actuator_maps(evaluation.startup) + ) + self._last_report = None + await self._close_actuator(previous) + + def _actuator_maps( + self, startup: Optional[FrontendStartup] + ) -> dict[str, dict[str, float]]: + if startup is not None and startup.startup_maps: + return {wt: dict(params) for wt, params in startup.startup_maps.items()} + # Unknown startup: rollback needs only the latched initials. + return {wt: {} for wt in self._settings.worker_types} + + def _make_resolver( + self, startup: Optional[FrontendStartup] + ) -> Optional[RouterResolver]: + if startup is None or self._resolver_factory is None: + return None + try: + return self._resolver_factory(startup) + except Exception: + logger.exception("retune: building the router resolver failed") + return None + + def _router_keys(self) -> tuple[str, ...]: + controller = self._controller() + keys = controller.router_keys if controller is not None else None + return tuple(keys) if keys else self._settings.router_keys + + def _startup_by_wt(self) -> Optional[dict[str, dict[str, float]]]: + startup = self._evaluation.startup + if startup is None: + return None + return {wt: dict(params) for wt, params in startup.startup_maps.items()} + + def startup_flat(self) -> Optional[dict[str, float]]: + """The controller's startup router point (live names), if known.""" + + return self._startup_flat() + + def _startup_flat(self) -> Optional[dict[str, float]]: + resolver = self._resolver + keys = self._router_keys() + if resolver is None or not keys: + return None + try: + return resolver.startup_flat(keys) + except Exception: + logger.exception("retune: resolving the startup router point failed") + return None + + def _resolve_by_wt( + self, flat: Mapping[str, float] + ) -> Optional[dict[str, dict[str, float]]]: + resolver = self._resolver + if resolver is None: + return None + try: + return resolver.by_wt(flat) + except Exception: + logger.exception("retune: resolving router point %s failed", dict(flat)) + return None + + def _p12_ok(self) -> bool: + """P12 from the last actuator round (routers present, admin API on).""" + + report = self._last_report + if report is None or report.admin_api_disabled: + return False + present = report.routers_present + return bool(report.converged if present is None else present) + + def _last_decided_k(self) -> Optional[int]: + state = self.state + if state is None or state.last_window_end_ms is None: + return None + return window_index( + state.last_window_end_ms, + phase_s=self._settings.phase_s, + interval_s=self._settings.interval_s, + ) + + def _rewind(self, controller: Any) -> None: + last = self._last_decided_k() + try: + controller.rewind(-1 if last is None else last) + except Exception: + logger.exception("retune: rewind failed") + + def _rewind_now(self) -> None: + controller = self._controller() + if controller is not None: + self._rewind(controller) + + def _engine(self) -> Any: + """The host's engine adapter; ``None`` before ``start`` or without one.""" + + return getattr(self._host, "_engine", None) + + def _controller(self) -> Any: + engine = self._engine() + return None if engine is None else engine.retune_controller + + def _sweeps_done(self) -> Optional[Future[None]]: + engine = self._engine() + return None if engine is None else engine.retune_sweeps_done() + + @contextlib.asynccontextmanager + async def _host_lock(self) -> AsyncIterator[None]: + lock = getattr(self._host, "_config_lock", None) + if lock is None: + yield + return + async with lock: + yield + + async def _io(self, function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + call = functools.partial(function, *args, **kwargs) + return await asyncio.wait_for( + self._blocking(call), timeout=self._settings.io_timeout_s + ) + + @staticmethod + async def _cancel(task: Optional[asyncio.Task[Any]]) -> None: + if task is None or task.done() or task is asyncio.current_task(): + return + task.cancel() + await asyncio.wait([task]) + + @staticmethod + async def _close_actuator(actuator: Any) -> None: + client = getattr(actuator, "_client", None) + close = getattr(client, "close", None) + if close is None: + return + try: + result = close() + if inspect.isawaitable(result): + await result + except Exception: + logger.debug( + "retune: closing the router policy client failed", exc_info=True + ) + + def _log_report(self, report: PremiseReport) -> None: + for reason in report.reasons: + logger.warning("retune premise failed: %s", reason) + for warning in report.warnings: + logger.warning("retune premise warning: %s", warning) + + def _publish(self) -> None: + state = self.state + try: + self._metrics.supervisor_state( + leader=self.is_leader(), + ready=self.ready(), + initialized=self._initialized, + rollback_pending=bool(state is not None and state.rollback_pending), + generation=self._generation, + ) + except Exception: + logger.debug("retune: publishing supervisor metrics failed", exc_info=True) + + # ---------------------------------------------------------------- loops + + async def _lease_loop(self) -> None: + while not self._stopping: + started = self._monotonic() + step: Optional[LeaseStep] = "error" + try: + step = await self.lease_round() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("retune lease loop iteration failed") + wait_s = self._lease_wait_s(failed=step == "error", started=started) + await self._wait(self._lease_wake, wait_s) + + def _lease_wait_s(self, *, failed: bool, started: float) -> float: + """Time until the next renewal after an attempt begun at ``started``. + + While leading, after a successful attempt the next starts no later + than one I/O budget before ``leader_until`` (at once if that has + passed), so that it ends while this process still leads. After a + failed attempt the next starts no sooner than one request timeout + after the failed one started, even within one I/O budget of + ``leader_until`` or past it: a fast-failing API server is not polled + in a busy loop, and a retry after a slow failure is not pushed back + further. + """ + + settings = self._settings + floor_s = 0.0 + if failed: + elapsed_s = self._monotonic() - started + floor_s = max(0.0, settings.k8s_timeout_s - elapsed_s) + return self._leadership.renewal_wait_s( + interval_s=settings.renew_interval_s, + budget_s=settings.io_timeout_s, + floor_s=floor_s, + ) + + async def _lease_step(self) -> LeaseStep: + """One ``LeaseLeadership.step``; ``error`` if it raised.""" + + async with self._lease_step_lock: + try: + return await self._io(self._leadership.step) + except Exception as error: + logger.warning("retune lease renewal failed: %s", error) + return "error" + + async def _renew_while_initializing(self, step: LeaseStep, started: float) -> None: + """Renew on the lease loop's schedule while an initialization runs. + + The lease loop waits for the initialization (or for ``_init_lock`` + during a rebuild), which may outlast the local lifetime left after + its confirming renewal. ``step`` and ``started`` are that renewal's. + A lapse (``acquired``) requests a re-initialization; the leadership + re-acquired is still renewed until the stale initialization ends. + """ + + while True: + await self._sleep( + self._lease_wait_s(failed=step == "error", started=started) + ) + if self._stopping: + return + started = self._monotonic() + step = await self._lease_step() + if step == "acquired": + self._request_reinit("leadership lapsed during initialization") + elif step not in ("renewed", "error"): + return + + async def _end_renewals(self, keepalive: asyncio.Task[None]) -> None: + """Cancel ``_renew_while_initializing`` and wait for it to end. + + A renewal in flight is waited for first: cancelled, its thread would + go on and race the next step on the resourceVersion. If this task is + cancelled meanwhile (``stop``), the renewal task is cancelled without + waiting, as the lease loop's own step is: the renewal's thread may + still finish its step, and once ``stop`` has released the + ``LeaseLeadership``, a step that reads the Lease after that is a + no-op. + """ + + try: + async with self._lease_step_lock: + await self._cancel(keepalive) + except asyncio.CancelledError: + keepalive.cancel() + await asyncio.wait([keepalive]) + raise + + async def _control_loop(self) -> None: + while not self._stopping: + try: + await self.control_round() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("retune control loop iteration failed") + await self._wait(None, self._settings.actuator_interval_s) + + async def _outbox_loop(self) -> None: + while not self._stopping: + await self._outbox_wake.wait() + self._outbox_wake.clear() + try: + await self.outbox_round() + except asyncio.CancelledError: + raise + except Exception: + logger.exception("retune outbox round failed") + if self._reinit_requested: + return + if self._outbox: + # A failed write: retry after a short pause. + await self._sleep(min(5.0, self._settings.k8s_timeout_s)) + self._outbox_wake.set() + + async def _wait(self, event: Optional[asyncio.Event], timeout_s: float) -> None: + """Sleep ``timeout_s`` (the injectable sleep) or until ``event`` is set.""" + + sleeper = asyncio.ensure_future(self._sleep(timeout_s)) + waiters: set[asyncio.Future[Any]] = {sleeper} + if event is not None: + waiters.add(asyncio.ensure_future(event.wait())) + try: + await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED) + finally: + for waiter in waiters: + waiter.cancel() + await asyncio.wait(waiters) + if event is not None: + event.clear() + + +def _log_reconcile_failure(work: asyncio.Future[Any]) -> None: + """Done callback of a reconciler call (it may outlive its round's task).""" + + if work.cancelled(): + return + error = work.exception() + if error is not None: + logger.error("retune reconciler round failed", exc_info=error) + + +def _startup_key(startup: Optional[FrontendStartup]) -> Optional[str]: + return None if startup is None else startup.fingerprint_source + + +def _finite(value: Optional[float]) -> Optional[float]: + if value is None: + return None + try: + number = float(value) + except (TypeError, ValueError): + return None + return number if math.isfinite(number) else None diff --git a/components/src/dynamo/planner/retune/live/window_cut.py b/components/src/dynamo/planner/retune/live/window_cut.py new file mode 100644 index 000000000000..a44006201487 --- /dev/null +++ b/components/src/dynamo/planner/retune/live/window_cut.py @@ -0,0 +1,1406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cut a retune window file from the frontends' live request-trace segments. + +The frontends append ``dynamo.request.trace.v1`` rows to rotating +``//req.NNNNNN.jsonl.gz`` segments on a shared RWX claim, +one gzip member per flush (``lib/llm/src/telemetry/jsonl_gz.rs``). +``cut_window`` turns the ``request_end`` rows that arrived in +``[start_ms, end_ms)`` and completed by ``as_of_ms`` into the bare-row window +file a DGDR Run replays. Selection, censoring, de-duplication, the single +block-size rule and ``row_counts`` are those of +``retune.window.read_trace_window`` over the segments the cut reads, and the +rows are the same bytes; only the output is gzip-compressed and the memory +is bounded. The cut runs in a child process, reads the segments in two +passes (``_Selection``, ``_write_rows``) and publishes the file atomically +(``_publish``). The file's gzip header records its selection (``start_ms``, +``end_ms``, ``as_of_ms`` in whole ms) and counts; a file already at the +output path is reused only for the same selection. + +A gzip member that fails validation is dropped whole. Runs of bad bytes +above ``MAX_DROPPED_MEMBER_RATIO`` of the members, or bad bytes above that +ratio of the gzip bytes, fail the cut (``trace_corrupt``). A run that +reaches the end of a segment is a writer's in-flight flush, or the last +flush of a writer that crashed, and is counted but does not count toward +either ratio: after an in-place container restart the sink continues in a +new segment, so that tail can end any segment. +""" + +from __future__ import annotations + +import errno +import hashlib +import json +import logging +import math +import multiprocessing +import os +import re +import secrets +import stat +import struct +import time +import zlib +from array import array +from collections import OrderedDict +from collections.abc import Callable, Generator, Iterable, Mapping, Sequence +from dataclasses import dataclass +from multiprocessing.connection import Connection, wait +from operator import attrgetter +from pathlib import Path +from typing import Any, BinaryIO, Literal, Optional, Protocol, Union, get_args + +import numpy as np + +from dynamo.planner.retune.live.contracts import SupportsIsSet, WindowCutResult +from dynamo.planner.retune.window import ( + _dumps, + _new_counts, + _parse_request_end, + _request_end_event, + _RequestRow, + _window_record, +) + +logger = logging.getLogger(__name__) + +StrPath = Union[str, "os.PathLike[str]"] + + +class _ByteSink(Protocol): + def write(self, data: bytes, /) -> Any: + ... + + +WindowCutFailure = Literal[ + "window_cut_timeout", # the deadline passed; the child was killed + "window_too_large", # more than max_rows de-duplicated rows + "trace_corrupt", # too many bad gzip members, or a malformed row + "window_cut_failed", # I/O error, mixed block sizes, child died, ... + "cancelled", # the caller's cancel event was set; the child was killed +] +WINDOW_CUT_FAILURES: tuple[str, ...] = get_args(WindowCutFailure) + +# Runs of bad bytes (other than segment tails) over all members, or their +# bytes over all gzip bytes, above which the cut fails as trace_corrupt. A run +# counts once however many members it wiped out (a zero-filled hole); its +# bytes measure how much of the trace it took. +MAX_DROPPED_MEMBER_RATIO = 0.01 +# A row is appended when its request completes, after it arrived, so a +# segment last modified before start_ms holds no row of the window and is not +# read. The margin covers clock skew between the frontends, which stamp +# arrivals, and the file server, which stamps modification times. +SEGMENT_MTIME_MARGIN_S = 300.0 +# Window files are written fast rather than small: on hash-chain rows level +# 1 is ~5x faster than level 6 for ~8% more bytes, and the cut budget pays +# for compression. +OUTPUT_COMPRESSLEVEL = 1 +# Temporary window files: ``.tmp--`` next to out_path. +TEMP_PREFIX = ".tmp-" +TEMP_NAME_RE = re.compile(r"^" + re.escape(TEMP_PREFIX) + r"(\d+)-[0-9A-Za-z]+$") +SEGMENT_SUFFIX = ".jsonl.gz" + +_GZIP_MAGIC = b"\x1f\x8b\x08" # ID1, ID2, CM=deflate +_GZIP_WBITS = 16 + zlib.MAX_WBITS # one gzip member, header and trailer +# The header of a window file: FLG=FCOMMENT, MTIME=0 (equal cuts are equal +# bytes), XFL=0, OS=unknown. The zero-terminated comment that follows holds +# the cut's selection and counts. +_WINDOW_HEADER = _GZIP_MAGIC + b"\x10" + bytes(4) + b"\x00\xff" +_WINDOW_HEADER_READ = 1 << 16 +_READ_CHUNK = 1 << 20 +_DECODE_STEP = 1 << 16 +_COMPACT_AT = 1 << 20 # consumed reader bytes before the buffer is compacted +# A trace file the reconciler deleted under the walk (ESTALE: on NFS). +_VANISHED = frozenset({errno.ENOENT, errno.ESTALE}) +# Pass-1 columns may hold this fraction of max_rows (at least the minimum) +# beyond max_rows before duplicates are removed. +_DEDUP_SLACK_RATIO = 0.25 +_DEDUP_SLACK_MIN = 1024 +_PASS2_CACHE_BYTES = 256 << 20 +_PASS2_BLOCK_ROWS = 1 << 16 +_LINE_OVERHEAD = 64 # bytes object and dict slot per cached line +_OPEN_FILES = 32 +_WRITE_BUFFER = 1 << 20 +_FLAG_AGENT_CONTEXT = 1 +_FLAG_NO_TOTAL_TIME = 2 +_POLL_INTERVAL_S = 0.25 +_EXIT_GRACE_S = 1.0 +_JOIN_TIMEOUT_S = 5.0 +_CHILD_NICENESS = 10 +_MAX_DROP_EXAMPLES = 5 + + +class WindowCutError(Exception): + """A window cut failed; ``reason`` is one of ``WINDOW_CUT_FAILURES``.""" + + def __init__(self, reason: WindowCutFailure, message: str) -> None: + if reason not in WINDOW_CUT_FAILURES: + raise ValueError(f"unknown window cut failure reason {reason!r}") + super().__init__(message) + self.reason: WindowCutFailure = reason + self.message = message + + def __reduce__(self) -> tuple[Any, ...]: + return (type(self), (self.reason, self.message)) + + +def temp_file_name(window_end_ms: float, nonce: Optional[str] = None) -> str: + """The name of a window's temporary file (``.tmp--``).""" + + if nonce is None: + nonce = secrets.token_hex(8) + elif not re.fullmatch(r"[0-9A-Za-z]+", nonce): + raise ValueError(f"nonce must be alphanumeric, got {nonce!r}") + return f"{TEMP_PREFIX}{int(round(window_end_ms))}-{nonce}" + + +# ------------------------------------------------------------- entry points + + +def cut_window( + segments_root: StrPath, + start_ms: float, + end_ms: float, + as_of_ms: float, + out_path: StrPath, + *, + max_rows: int, + deadline_s: float, + cancel: Optional[SupportsIsSet] = None, +) -> WindowCutResult: + """Cut a window file in a spawned child process (``contracts.WindowCutter``). + + ``deadline_s`` is absolute wall-clock time (``time.time()`` seconds), e.g. + ``min(now + cut_budget_seconds, D_k)``. The child is killed at the + deadline (``window_cut_timeout``) or as soon as ``cancel`` is set + (``cancelled``; any object with ``is_set()`` works, so a caller can pass + one that watches several events). Every failure raises ``WindowCutError``; + invalid arguments raise ``ValueError`` before a child is started. + """ + + _validate(start_ms, end_ms, as_of_ms, out_path, max_rows) + if not math.isfinite(deadline_s): + raise ValueError("deadline_s must be finite") + deadline = time.monotonic() + (deadline_s - time.time()) + out_path = os.fspath(out_path) + temp_path = os.path.join( + os.path.dirname(os.path.abspath(out_path)), temp_file_name(end_ms) + ) + started = time.monotonic() + result = _run_in_child( + cut_window_inprocess, + (os.fspath(segments_root), start_ms, end_ms, as_of_ms, out_path), + {"max_rows": max_rows, "temp_path": temp_path}, + deadline=deadline, + cancel=cancel, + cleanup=(temp_path,), + name=f"retune-window-cut-{int(round(end_ms))}", + ) + logger.info( + "retune window [%s, %s) as of %s: %d requests in %s (%s) in %.1fs; " + "dropped %d gzip members (%d bytes)", + start_ms, + end_ms, + as_of_ms, + result.num_requests, + result.path, + "reused" if result.reused else "written", + time.monotonic() - started, + result.dropped_members, + result.dropped_bytes, + ) + return result + + +def cut_window_inprocess( + segments_root: StrPath, + start_ms: float, + end_ms: float, + as_of_ms: float, + out_path: StrPath, + *, + max_rows: int, + temp_path: Optional[StrPath] = None, +) -> WindowCutResult: + """Cut a window file in this process; the body of ``cut_window``'s child. + + An empty window writes no file (``path=None``). When ``out_path`` already + exists, no segment is read: if its header records the same selection + (``start_ms``, ``end_ms`` and ``as_of_ms`` rounded to whole ms), the + result carries the counts in its header (``reused=True``, nothing + dropped); a file of another selection, or one that records none, fails + as ``window_cut_failed``. ``temp_path`` names the temporary file (it + must be in ``out_path``'s directory); by default a fresh + ``temp_file_name`` there. Failures raise ``WindowCutError``. + """ + + _validate(start_ms, end_ms, as_of_ms, out_path, max_rows) + out_path = os.fspath(out_path) + directory = os.path.dirname(os.path.abspath(out_path)) + if temp_path is None: + temp_path = os.path.join(directory, temp_file_name(end_ms)) + temp_path = os.fspath(temp_path) + if os.path.dirname(os.path.abspath(temp_path)) != directory: + # os.link needs one filesystem; the reconciler looks for it there. + raise ValueError(f"temp_path {temp_path!r} is not next to {out_path!r}") + try: + if os.path.lexists(out_path): + # Another process (an earlier leader, an adopted Request's cut) + # already published this window; DGDR Requests name it by path. + return _published( + out_path, start_ms=start_ms, end_ms=end_ms, as_of_ms=as_of_ms + ) + files = _segment_files(os.fspath(segments_root)) + selection = _Selection(start_ms, end_ms, as_of_ms, max_rows) + items = _trace_items(files, start_ms / 1000.0 - SEGMENT_MTIME_MARGIN_S) + try: + selection.scan(items) + finally: + items.close() + selection.finish() + result = WindowCutResult( + path=out_path if selection.num_rows else None, + num_requests=selection.num_rows, + trace_block_size=selection.block_size, + row_counts=dict(selection.counts), + dropped_members=selection.dropped_members, + dropped_bytes=selection.dropped_bytes, + ) + if result.path is None: + return result + created = _publish( + out_path, + temp_path, + _counts_comment( + result, start_ms=start_ms, end_ms=end_ms, as_of_ms=as_of_ms + ), + lambda sink: _write_rows(selection, files, sink), + ) + if created: + return result + # Another process published it first: report the counts it holds. + return _published(out_path, start_ms=start_ms, end_ms=end_ms, as_of_ms=as_of_ms) + except OSError as error: + raise WindowCutError( + "window_cut_failed", f"{type(error).__name__}: {error}" + ) from error + + +def _validate( + start_ms: float, + end_ms: float, + as_of_ms: float, + out_path: StrPath, + max_rows: int, +) -> None: + for name, value in ( + ("start_ms", start_ms), + ("end_ms", end_ms), + ("as_of_ms", as_of_ms), + ): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number, got {value!r}") + if not math.isfinite(value): + raise ValueError(f"{name} must be finite") + if not start_ms < end_ms: + raise ValueError("start_ms must be before end_ms") + if as_of_ms < end_ms: + raise ValueError("as_of_ms must not be before end_ms") + if isinstance(max_rows, bool) or not isinstance(max_rows, int) or max_rows < 1: + raise ValueError(f"max_rows must be a positive integer, got {max_rows!r}") + if not os.fspath(out_path).endswith(".gz"): + raise ValueError(f"window file must be gzip (.gz): {os.fspath(out_path)!r}") + + +# ------------------------------------------------------------ child process + + +def _run_in_child( + function: Callable[..., Any], + args: Sequence[Any], + kwargs: Mapping[str, Any], + *, + deadline: float, + cancel: Optional[SupportsIsSet], + cleanup: Sequence[str] = (), + name: str = "retune-window-cut", +) -> Any: + """``function(*args, **kwargs)`` in a ``spawn`` child; its return value. + + ``deadline`` is ``time.monotonic()`` time. The child is killed at the + deadline, on cancel, or once it has answered; ``cleanup`` paths are + removed unless it answered with a result (a killed child leaves its + temporary file behind). A ``WindowCutError`` raised by ``function`` is + re-raised with its reason; any other failure, the child's death included, + is ``window_cut_failed``. + """ + + _check_interrupted(deadline, cancel) + context = multiprocessing.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + process = context.Process( + target=_child_main, + args=(sender, function, tuple(args), dict(kwargs)), + name=name, + daemon=True, + ) + outcome: Optional[tuple[Any, ...]] = None + try: + try: + process.start() + except Exception as error: + raise WindowCutError( + "window_cut_failed", f"cannot start the window cut process: {error}" + ) from error + finally: + sender.close() + while outcome is None: + _check_interrupted(deadline, cancel) + timeout = max(0.0, min(deadline - time.monotonic(), _POLL_INTERVAL_S)) + if not wait([receiver, process.sentinel], timeout): + continue + try: + outcome = receiver.recv() + except (EOFError, OSError): + process.join(_JOIN_TIMEOUT_S) + raise WindowCutError( + "window_cut_failed", + f"window cut process exited with code {process.exitcode} " + "before returning a result", + ) from None + finally: + _stop(process, grace_s=_EXIT_GRACE_S if outcome is not None else 0.0) + receiver.close() + if outcome is None or outcome[0] != "ok": + for path in cleanup: + _unlink_quietly(path) + assert outcome is not None # the loop only ends with one + if outcome[0] == "ok": + return outcome[1] + _, reason, message = outcome + raise WindowCutError(reason, message) + + +def _check_interrupted(deadline: float, cancel: Optional[SupportsIsSet]) -> None: + if cancel is not None and cancel.is_set(): + raise WindowCutError("cancelled", "window cut cancelled") + if time.monotonic() >= deadline: + raise WindowCutError("window_cut_timeout", "window cut exceeded its deadline") + + +def _stop(process: multiprocessing.process.BaseProcess, *, grace_s: float) -> None: + """Let ``process`` exit for ``grace_s``, then kill it and reap it. + + A child blocked in uninterruptible I/O (a hard-mounted RWX volume) does + not die on SIGKILL until the I/O returns; it is left behind rather than + holding the sweep thread. SIGKILL takes effect before it returns to user + space, so it cannot publish a window afterwards. + """ + + if process.pid is None: + return + if grace_s > 0: + process.join(grace_s) + if process.is_alive(): + process.kill() + process.join(_JOIN_TIMEOUT_S) + if process.is_alive(): + logger.error( + "retune window cut process %s (pid %s) did not exit after SIGKILL; " + "leaving it behind", + process.name, + process.pid, + ) + + +def _child_main( + conn: Connection, + function: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> None: + """Child entry point: answer ``("ok", value)`` or ``("error", reason, message)``.""" + + # The cut shares the planner pod's CPU; the planner's own loop comes first. + try: + os.nice(_CHILD_NICENESS) + except OSError: + pass + try: + outcome: tuple[Any, ...] = ("ok", function(*args, **kwargs)) + except WindowCutError as error: + outcome = ("error", error.reason, error.message) + except BaseException as error: # noqa: BLE001 - every failure is reported + outcome = ("error", "window_cut_failed", f"{type(error).__name__}: {error}") + try: + conn.send(outcome) + except Exception as error: + conn.send(("error", "window_cut_failed", f"unsendable result: {error}")) + finally: + conn.close() + + +# ------------------------------------------------------------ gzip members + + +@dataclass(frozen=True) +class _Member: + """A gzip member that decoded and matched its CRC32 and ISIZE.""" + + offset: int + end: int + payload: bytes + + +@dataclass(frozen=True) +class _Garbage: + """A maximal run of bytes ``[offset, end)`` that holds no valid member.""" + + offset: int + end: int + at_eof: bool + + +class _GzipMemberReader: + """Iterate a gzip file as validated members and runs of bad bytes. + + Every byte from ``offset`` to the end of the file is in exactly one item, + in file order. A member is yielded only after zlib reached its end and + its trailer's CRC32 and ISIZE match the decompressed bytes (zlib checks + them too; the explicit check does not depend on how zlib was built). + Anything else - a member cut short by the end of the file, a corrupt + one, zeros - starts a ``_Garbage`` run that ends at the next gzip header + that decodes, or at the end of the file. Memory holds one member's + compressed and decompressed bytes plus a read chunk. + """ + + def __init__(self, source: BinaryIO, offset: int = 0) -> None: + if offset: + source.seek(offset) + self._source = source + self._buf = bytearray() + self._pos = 0 # current position in _buf + self._base = offset # file offset of _buf[0] + self._eof = False + self._pending: Optional[_Member] = None + + def __iter__(self) -> _GzipMemberReader: + return self + + def __next__(self) -> Union[_Member, _Garbage]: + if self._pending is not None: + pending, self._pending = self._pending, None + return pending + if self._pos >= len(self._buf) and not self._fill(): + raise StopIteration + start = self._offset() + member = self._decode() + if member is not None: + return member + while self._skip_to_magic(): + member = self._decode() + if member is not None: + self._pending = member + return _Garbage(start, member.offset, at_eof=False) + return _Garbage(start, self._offset(), at_eof=True) + + def _offset(self) -> int: + return self._base + self._pos + + def _fill(self) -> bool: + if self._eof: + return False + chunk = self._source.read(_READ_CHUNK) + if not chunk: + self._eof = True + return False + self._buf += chunk + return True + + def _consume(self, count: int) -> None: + self._pos += count + if self._pos >= _COMPACT_AT and 2 * self._pos >= len(self._buf): + del self._buf[: self._pos] + self._base += self._pos + self._pos = 0 + + def _decode(self) -> Optional[_Member]: + """Decode and consume the member at the current position. + + Returns ``None``, consuming nothing, when no valid member starts here. + """ + + decoder = zlib.decompressobj(wbits=_GZIP_WBITS) + parts: list[bytes] = [] + index = self._pos + while True: + if index >= len(self._buf) and not self._fill(): + return None # cut short by the end of the file + piece = self._buf[index : index + _DECODE_STEP] + try: + parts.append(decoder.decompress(piece)) + except zlib.error: + return None + if decoder.eof: + end = index + len(piece) - len(decoder.unused_data) + payload = b"".join(parts) + if not _trailer_matches(self._buf[end - 8 : end], payload): + return None + member = _Member(self._offset(), self._base + end, payload) + self._consume(end - self._pos) + return member + index += len(piece) + + def _skip_to_magic(self) -> bool: + """Consume at least one byte, then up to the next gzip magic. + + ``False`` (everything consumed) when no magic is left in the file. + """ + + self._consume(1) + keep = len(_GZIP_MAGIC) - 1 # a magic may straddle two reads + while True: + found = self._buf.find(_GZIP_MAGIC, self._pos) + if found >= 0: + self._consume(found - self._pos) + return True + self._consume(max(0, len(self._buf) - keep - self._pos)) + if not self._fill(): + self._consume(len(self._buf) - self._pos) + return False + + +def _trailer_matches(trailer: Union[bytes, bytearray], payload: bytes) -> bool: + crc, size = struct.unpack(" Optional[bytes]: + """The payload of ``data`` when it is exactly one valid gzip member.""" + + decoder = zlib.decompressobj(wbits=_GZIP_WBITS) + try: + payload = decoder.decompress(data) + except zlib.error: + return None + if not decoder.eof or decoder.unused_data or len(data) < 8: + return None + return payload if _trailer_matches(data[-8:], payload) else None + + +# ---------------------------------------------------------------- pass 1 + + +def _segment_files(root: str) -> list[str]: + """Trace segments under ``root``, in replay-reader order. + + The files are the sorted ``*.jsonl.gz`` regular files under ``root``, in + the order ``resolve_trace_files([root])`` reads them; plain ``*.jsonl`` + files are ignored (logged), as the reconciler ignores them. A directory + or file under ``root`` that vanishes during the walk (ENOENT, or ESTALE + on NFS: the reconciler deletes a dead pod's segments) is skipped; any + other error listing or stat-ing one (EACCES, EIO, ETIMEDOUT on a soft + mount, ...) fails the cut as ``window_cut_failed`` rather than silently + losing a pod's rows, as does any error on ``root`` itself. + """ + + try: + root_mode = os.stat(root).st_mode + except OSError as error: + if error.errno in _VANISHED: + raise WindowCutError( + "window_cut_failed", f"trace segments root {root} does not exist" + ) from error + raise _unreadable(root, error) from error + if stat.S_ISDIR(root_mode): + + def walk_error(error: OSError) -> None: + if error.filename != root and error.errno in _VANISHED: + logger.info( + "retune window cut: trace directory %s vanished (%s); " + "skipping it", + error.filename, + error, + ) + return + raise _unreadable(error.filename, error) from error + + found: list[Path] = [] + plain: list[str] = [] + for directory, _, names in os.walk(root, onerror=walk_error): + for name in names: + candidate = os.path.join(directory, name) + if name.endswith(".jsonl"): + plain.append(candidate) + elif name.endswith(SEGMENT_SUFFIX) and _is_regular_file(candidate): + found.append(Path(candidate)) + if plain: + logger.warning( + "retune window cut: ignoring %d plain JSONL files under %s (the " + "live trace sink must write %s segments), e.g. %s", + len(plain), + root, + SEGMENT_SUFFIX, + min(plain), + ) + return [str(path) for path in sorted(found)] + if root.endswith(SEGMENT_SUFFIX): + return [root] + raise WindowCutError( + "window_cut_failed", f"trace segment {root} is not a {SEGMENT_SUFFIX} file" + ) + + +def _is_regular_file(path: str) -> bool: + """Whether ``path`` is a regular file (following symlinks, like ``is_file``). + + ``False`` when it vanished; any other stat error fails the cut. + """ + + try: + return stat.S_ISREG(os.stat(path).st_mode) + except OSError as error: + if error.errno in _VANISHED: + return False + raise _unreadable(path, error) from error + + +def _unreadable(path: Any, error: OSError) -> WindowCutError: + return WindowCutError( + "window_cut_failed", + f"cannot read trace segments at {path}: {type(error).__name__}: {error}", + ) + + +# (file index, path, item); an OSError item means the file vanished. +_TraceItem = tuple[int, str, Union[_Member, _Garbage, OSError]] + + +def _trace_items( + files: Sequence[str], not_before_s: float +) -> Generator[_TraceItem, None, None]: + """Every member and bad byte run of ``files``, in file order. + + A file last modified before ``not_before_s`` (epoch seconds) is skipped. + Its modification time is read once it is open: opening a file on NFS + revalidates its cached attributes (close-to-open consistency). A file + that vanishes (ENOENT, or ESTALE on NFS) before or while it is read + yields its error and the walk goes on; other I/O errors raise. + """ + + for file_index, path in enumerate(files): + try: + with open(path, "rb") as source: + if os.fstat(source.fileno()).st_mtime < not_before_s: + logger.debug( + "retune window cut: %s was last written before the " + "window; skipping it", + path, + ) + continue + for item in _GzipMemberReader(source): + yield file_index, path, item + except OSError as error: + if error.errno not in _VANISHED: + raise + yield file_index, path, error + + +class _Selection: + """Pass 1: the selected rows of every trace file read, as columns. + + Row ``i`` is ``arrival[i]``, ``completion[i]``, ``rid[16*i:16*i+16]`` + (BLAKE2b of ``request_id``), line ``line[i]`` of member + ``members[member[i]]``, ``block[i]`` and ``flags[i]``. Rows are in file + order, the order the replay reader sees them, so de-duplication keeps the + same (first) row it does. A row the replay reader would reject, in a + member whose CRC is good, fails the cut as ``trace_corrupt``. + """ + + def __init__( + self, start_ms: float, end_ms: float, as_of_ms: float, max_rows: int + ) -> None: + self.start_ms = start_ms + self.end_ms = end_ms + self.as_of_ms = as_of_ms + self.max_rows = max_rows + self._dedup_at = max_rows + max( + _DEDUP_SLACK_MIN, int(max_rows * _DEDUP_SLACK_RATIO) + ) + self.counts = _new_counts() + # (file index, offset, end) of each member with a selected row. + self.members: list[tuple[int, int, int]] = [] + self.arrival = array("q") + self.completion = array("q") + self.rid = bytearray() + self.member = array("I") + self.line = array("I") + self.block = array("q") + self.flags = array("B") + self.duplicates = 0 + self.valid_members = 0 + # Runs of bad bytes (``_Garbage``), each counted once however many + # members it held, and their bytes. + self.dropped_members = 0 + self.dropped_bytes = 0 + # The runs, and their bytes, that count toward MAX_DROPPED_MEMBER_RATIO. + self.corrupt_members = 0 + self.corrupt_bytes = 0 + # Bytes of every gzip member and run of bad bytes. + self.gzip_bytes = 0 + self.block_size: Optional[int] = None + self._drop_examples: list[str] = [] + + @property + def num_rows(self) -> int: + return len(self.arrival) + + def digest(self, row: int) -> bytes: + return bytes(self.rid[16 * row : 16 * row + 16]) + + def scan(self, items: Iterable[_TraceItem]) -> None: + """Consume ``_trace_items`` (in file order).""" + + for file_index, path, item in items: + if isinstance(item, _Member): + self.valid_members += 1 + self.gzip_bytes += item.end - item.offset + self._scan_payload( + item.payload, file_index, path, item.offset, item.end + ) + elif isinstance(item, _Garbage): + self._dropped(path, item) + else: + # The reconciler deletes only segments far older than any + # window being cut; a row taken from this one fails pass 2. + logger.warning( + "retune window cut: trace file %s vanished (%s); skipping it", + path, + item, + ) + + def finish(self) -> None: + """De-duplicate, then apply the corruption, size and block-size rules.""" + + self._dedup() + self.counts["duplicate"] = self.duplicates + total = self.valid_members + self.dropped_members + # A run counts as one member: its bytes catch a hole that wiped out + # many members. + if ( + self.corrupt_members > MAX_DROPPED_MEMBER_RATIO * total + or self.corrupt_bytes > MAX_DROPPED_MEMBER_RATIO * self.gzip_bytes + ): + raise WindowCutError( + "trace_corrupt", + f"{self.corrupt_members} of {total} gzip members of the trace " + f"failed validation ({self.corrupt_bytes} of {self.gzip_bytes} " + f"bytes; limit {MAX_DROPPED_MEMBER_RATIO:.0%} of either); " + f"dropped byte runs: {', '.join(self._drop_examples)}", + ) + if self.num_rows > self.max_rows: + raise _too_large(self.max_rows) + if not self.num_rows: + return + blocks = np.unique(np.frombuffer(self.block, dtype=self.block.typecode)) + if len(blocks) > 1: + raise WindowCutError( + "window_cut_failed", + "trace window mixes trace_block_size values " + f"{blocks.tolist()}; replay needs one block size", + ) + self.block_size = int(blocks[0]) + flags: np.ndarray = np.frombuffer(self.flags, dtype=self.flags.typecode) + self.counts["agent_context_stripped"] = int( + np.count_nonzero(flags & _FLAG_AGENT_CONTEXT) + ) + self.counts["censor_unavailable"] = int( + np.count_nonzero(flags & _FLAG_NO_TOTAL_TIME) + ) + + def _dropped(self, path: str, item: _Garbage) -> None: + size = item.end - item.offset + self.dropped_members += 1 + self.dropped_bytes += size + self.gzip_bytes += size + if item.at_eof: + # The flush its writer is appending, or the last one of a writer + # that crashed (an in-place restart continues in a new segment). + logger.debug( + "retune window cut: %d trailing bytes of %s hold no complete " + "gzip member (in-flight or interrupted write)", + size, + path, + ) + return + self.corrupt_members += 1 + self.corrupt_bytes += size + if len(self._drop_examples) < _MAX_DROP_EXAMPLES: + self._drop_examples.append(f"{path}@{item.offset}+{size}") + logger.warning( + "retune window cut: dropped %d bytes at %s offset %d that hold no " + "valid gzip member", + size, + path, + item.offset, + ) + + def _scan_payload( + self, + payload: bytes, + file_index: int, + path: str, + offset: int, + end: int, + ) -> None: + lines = payload.split(b"\n") + if not lines[-1]: + lines.pop() + member_id = -1 + start_ms, end_ms, as_of_ms = self.start_ms, self.end_ms, self.as_of_ms + counts = self.counts + for index, raw in enumerate(lines): + if not raw.strip(): + continue + location = f"{path}@{offset}:{index + 1}" + event = _event(raw, location, "trace_corrupt") + if event is None: + continue + row = _request_row(event, location, "trace_corrupt") + if row is None: + counts["missing_output_tokens"] += 1 + continue + if not start_ms <= row.arrival_ms < end_ms: + continue + if row.completion_ms > as_of_ms: + counts["censored"] += 1 + continue + if member_id < 0: + member_id = len(self.members) + self.members.append((file_index, offset, end)) + self._append(row, member_id, index, location) + + def _append( + self, row: _RequestRow, member_id: int, index: int, location: str + ) -> None: + try: + self.arrival.append(row.arrival_ms) + self.completion.append(row.completion_ms) + self.block.append(row.trace_block_size) + except OverflowError as error: + raise WindowCutError( + "trace_corrupt", f"trace row value out of range at {location}" + ) from error + self.rid += _digest(row.request_id) + self.member.append(member_id) + self.line.append(index) + self.flags.append( + (_FLAG_AGENT_CONTEXT if row.had_agent_context else 0) + | (_FLAG_NO_TOTAL_TIME if row.total_time_ms is None else 0) + ) + if len(self.arrival) > self._dedup_at: + # Duplicates only shrink the window, and rows already kept stay + # kept: past max_rows after de-duplication, the window is too large. + self._dedup() + if len(self.arrival) > self.max_rows: + raise _too_large(self.max_rows) + + def _dedup(self) -> None: + count = len(self.arrival) + if count < 2: + return + keep = _first_occurrences(self.rid, count) + kept = int(np.count_nonzero(keep)) + if kept == count: + return + self.duplicates += count - kept + self.arrival = _compress(self.arrival, keep) + self.completion = _compress(self.completion, keep) + self.member = _compress(self.member, keep) + self.line = _compress(self.line, keep) + self.block = _compress(self.block, keep) + self.flags = _compress(self.flags, keep) + self.rid = bytearray( + np.frombuffer(self.rid, dtype=np.uint8).reshape(count, 16)[keep].tobytes() + ) + + +def _first_occurrences(rid: bytearray, count: int) -> np.ndarray: + """Mask of the rows whose 128-bit digest has not occurred before them.""" + + words: np.ndarray = np.frombuffer(rid, dtype=np.uint64).reshape(count, 2) + high, low = words[:, 0], words[:, 1] + order = np.lexsort((np.arange(count), low, high)) + high, low = high[order], low[order] + repeat: np.ndarray = np.zeros(count, dtype=bool) + repeat[1:] = (high[1:] == high[:-1]) & (low[1:] == low[:-1]) + keep: np.ndarray = np.ones(count, dtype=bool) + keep[order[repeat]] = False + return keep + + +def _compress(column: array, keep: np.ndarray) -> array: + kept = array(column.typecode) + kept.frombytes(np.frombuffer(column, dtype=column.typecode)[keep].tobytes()) + return kept + + +def _too_large(max_rows: int) -> WindowCutError: + return WindowCutError( + "window_too_large", + f"the window has more than {max_rows} requests (max_window_rows)", + ) + + +def _digest(request_id: str) -> bytes: + # UTF-8 with surrogatepass is injective on str, lone surrogates included. + return hashlib.blake2b( + request_id.encode("utf-8", "surrogatepass"), digest_size=16 + ).digest() + + +def _event( + raw: bytes, location: str, failure: WindowCutFailure +) -> Optional[dict[str, Any]]: + """The ``request_end`` event of a trace line, as the replay reader reads it.""" + + try: + record = json.loads(raw) + except ValueError as error: # JSONDecodeError, UnicodeDecodeError + raise WindowCutError(failure, f"invalid JSON at {location}: {error}") from error + if not isinstance(record, dict): + raise WindowCutError(failure, f"trace row at {location} must be a JSON object") + try: + return _request_end_event(record, location) + except ValueError as error: + raise WindowCutError(failure, str(error)) from error + + +def _request_row( + event: dict[str, Any], location: str, failure: WindowCutFailure +) -> Optional[_RequestRow]: + try: + return _parse_request_end(event, location, need_hashes=True) + except ValueError as error: + raise WindowCutError(failure, str(error)) from error + + +# ---------------------------------------------------------------- pass 2 + + +def _publish( + out_path: str, + temp_path: str, + comment: bytes, + write: Callable[[_ByteSink], None], +) -> bool: + """Write ``out_path`` atomically through ``temp_path``. + + The file is one gzip member with ``comment`` in its header and what + ``write`` writes as its data. Returns ``False`` when ``out_path`` already + existed: it is kept and the new file is discarded. ``temp_path`` is + removed in every case. + """ + + directory = os.path.dirname(os.path.abspath(out_path)) + os.makedirs(directory, exist_ok=True) + try: + descriptor = os.open( + temp_path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_CLOEXEC", 0), + 0o600, + ) + with open(descriptor, "wb") as raw: + sink = _GzipWriter(raw, comment) + write(sink) + sink.close() + raw.flush() + os.fsync(raw.fileno()) + # The Run's sweeper reads the window as uid 1000 with supplemental + # group 65532 and no fsGroup, so it must be readable by others: 0600 + # would fail every Run of a Planner that runs as another uid. + os.chmod(temp_path, 0o644) + try: + os.link(temp_path, out_path) + created = True + except FileExistsError: + created = False + finally: + _unlink_quietly(temp_path) + _fsync_directory(directory) + return created + + +def _fsync_directory(directory: str) -> None: + try: + descriptor = os.open(directory, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + except OSError as error: + logger.debug("retune window cut: cannot open %s to fsync: %s", directory, error) + return + try: + os.fsync(descriptor) + except OSError as error: + # Some network filesystems reject fsync on a directory; their LINK + # is synchronous at the server anyway. + logger.debug("retune window cut: fsync of %s failed: %s", directory, error) + finally: + os.close(descriptor) + + +def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except FileNotFoundError: + pass + except OSError as error: + logger.warning("retune window cut: cannot remove %s: %s", path, error) + + +class _GzipWriter: + """One gzip member: ``_WINDOW_HEADER``, the zero-terminated ``comment``, + the deflated bytes of every ``write`` and the CRC32/ISIZE trailer. + + ``gzip.GzipFile`` cannot write a header comment. + """ + + def __init__(self, raw: BinaryIO, comment: bytes) -> None: + raw.write(_WINDOW_HEADER + comment + b"\x00") + self._raw = raw + self._deflate = zlib.compressobj( + OUTPUT_COMPRESSLEVEL, zlib.DEFLATED, -zlib.MAX_WBITS + ) + self._crc = 0 + self._size = 0 + + def write(self, data: bytes) -> None: + self._crc = zlib.crc32(data, self._crc) + self._size += len(data) + self._raw.write(self._deflate.compress(data)) + + def close(self) -> None: + self._raw.write(self._deflate.flush()) + self._raw.write(struct.pack(" dict[str, int]: + return { + "start_ms": int(round(start_ms)), + "end_ms": int(round(end_ms)), + "as_of_ms": int(round(as_of_ms)), + } + + +def _describe_selection(bounds: Mapping[str, Any]) -> str: + return f"[{bounds['start_ms']}, {bounds['end_ms']}) as of {bounds['as_of_ms']}" + + +def _counts_comment( + result: WindowCutResult, *, start_ms: float, end_ms: float, as_of_ms: float +) -> bytes: + """The header comment of a window file: the selection and the counts of + its cut, as ASCII JSON (which never holds the NUL that ends a gzip + comment).""" + + counts = { + "selection": _selection_bounds(start_ms, end_ms, as_of_ms), + "num_requests": result.num_requests, + "trace_block_size": result.trace_block_size, + "row_counts": result.row_counts, + } + return json.dumps(counts, sort_keys=True, separators=(",", ":")).encode("ascii") + + +def _published( + path: str, *, start_ms: float, end_ms: float, as_of_ms: float +) -> WindowCutResult: + """The result of the window file already at ``path``, from its header. + + Nothing was dropped by this call: it read no segment. A file without + the counts of a cut in its header, or cut for another selection than + ``[start_ms, end_ms)`` as of ``as_of_ms`` (whole ms), or recording no + selection, fails as ``window_cut_failed``. + """ + + with open(path, "rb") as source: + header = source.read(_WINDOW_HEADER_READ) + start = len(_WINDOW_HEADER) + end = header.find(b"\x00", start) + try: + if not header.startswith(_WINDOW_HEADER) or end < 0: + raise ValueError("no header comment") + counts = json.loads(header[start:end]) + num_requests = counts["num_requests"] + block_size = counts["trace_block_size"] + row_counts = counts["row_counts"] + if not ( + _is_count(num_requests) + and num_requests > 0 + and _is_count(block_size) + and isinstance(row_counts, dict) + and all(_is_count(value) for value in row_counts.values()) + ): + raise ValueError(f"malformed counts {counts!r}") + except (ValueError, TypeError, KeyError) as error: + raise WindowCutError( + "window_cut_failed", + f"window file {path} exists but its gzip header holds no cut " + f"counts: {error}", + ) from error + wanted = _selection_bounds(start_ms, end_ms, as_of_ms) + found = counts.get("selection") + if found != wanted: + if not (isinstance(found, dict) and set(found) == set(wanted)): + raise WindowCutError( + "window_cut_failed", + f"window file {path} exists but its gzip header holds no " + "selection, so it is not reused", + ) + raise WindowCutError( + "window_cut_failed", + f"window file {path} was cut for {_describe_selection(found)}, " + f"not {_describe_selection(wanted)}", + ) + return WindowCutResult( + path=path, + num_requests=num_requests, + trace_block_size=block_size, + row_counts=row_counts, + reused=True, + ) + + +def _is_count(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + + +def _write_rows(selection: _Selection, paths: Sequence[str], sink: _ByteSink) -> None: + """Pass 2: write the kept rows to ``sink`` in window-file order. + + The order is ``read_trace_window``'s: (arrival, completion, request_id). + The columns order rows by the first two; rows with an equal pair are + collected and ordered by their ``request_id`` before they are written. + Only the members that hold kept rows are read again; one that no longer + matches pass 1 fails the cut (the segment changed). + """ + + count = selection.num_rows + arrival: np.ndarray = np.frombuffer( + selection.arrival, dtype=selection.arrival.typecode + ) + completion: np.ndarray = np.frombuffer( + selection.completion, dtype=selection.completion.typecode + ) + members: np.ndarray = np.frombuffer( + selection.member, dtype=selection.member.typecode + ) + lines: np.ndarray = np.frombuffer(selection.line, dtype=selection.line.typecode) + order = np.lexsort((completion, arrival)) + position: np.ndarray = np.empty(count, dtype=np.int64) + position[order] = np.arange(count, dtype=np.int64) + # Each member's rows, in the order they are written. + by_member = np.lexsort((position, members)) + bounds = np.asarray( + np.searchsorted(members[by_member], np.arange(len(selection.members) + 1)) + ) + cache = _MemberLines( + paths, + selection.members, + bounds, + position[by_member], + lines[by_member], + capacity=_PASS2_CACHE_BYTES, + ) + del position, by_member + group: list[_RequestRow] = [] + group_key: Optional[tuple[int, int]] = None + writer = _LineWriter(sink) + try: + for first in range(0, count, _PASS2_BLOCK_ROWS): + block = order[first : first + _PASS2_BLOCK_ROWS] + for row, arrival_ms, completion_ms, member_id, line_index in zip( + block.tolist(), + arrival[block].tolist(), + completion[block].tolist(), + members[block].tolist(), + lines[block].tolist(), + ): + if (arrival_ms, completion_ms) != group_key: + _write_group(group, writer) + group = [] + group_key = (arrival_ms, completion_ms) + raw = cache.take(member_id, line_index) + location = cache.location(member_id, line_index) + event = _event(raw, location, "window_cut_failed") + parsed = ( + None + if event is None + else _request_row(event, location, "window_cut_failed") + ) + if ( + parsed is None + or parsed.arrival_ms != arrival_ms + or parsed.completion_ms != completion_ms + or _digest(parsed.request_id) != selection.digest(row) + ): + raise _changed(f"the row at {location} differs") + group.append(parsed) + _write_group(group, writer) + writer.flush() + finally: + cache.close() + + +def _write_group(group: list[_RequestRow], writer: _LineWriter) -> None: + if len(group) > 1: + group.sort(key=attrgetter("request_id")) + for row in group: + writer.write(_dumps(_window_record(row))) + + +def _changed(detail: str) -> WindowCutError: + return WindowCutError( + "window_cut_failed", f"trace segment changed between passes: {detail}" + ) + + +class _LineWriter: + """Joins window lines into writes of about ``_WRITE_BUFFER`` bytes.""" + + def __init__(self, sink: _ByteSink) -> None: + self._sink = sink + self._parts: list[str] = [] + self._size = 0 + + def write(self, line: str) -> None: + self._parts.append(line) + self._parts.append("\n") + self._size += len(line) + 1 + if self._size >= _WRITE_BUFFER: + self.flush() + + def flush(self) -> None: + if self._parts: + self._sink.write("".join(self._parts).encode("utf-8")) + self._parts = [] + self._size = 0 + + +class _MemberLines: + """The lines pass 2 still has to write, decoded member by member. + + Rows are taken in output order. A member is read and decoded when one of + its rows is first needed and keeps only its lines still to be written; it + is dropped when its last one is taken. Above ``capacity`` bytes, the + members whose next row is written last are evicted (Belady) and decoded + again when needed, so memory stays bounded whatever the request durations. + + ``bounds[m]:bounds[m+1]`` indexes member ``m``'s rows in ``positions`` + (their output positions, ascending) and ``lines`` (their line indices). + """ + + def __init__( + self, + paths: Sequence[str], + members: Sequence[tuple[int, int, int]], + bounds: np.ndarray, + positions: np.ndarray, + lines: np.ndarray, + *, + capacity: int, + ) -> None: + self._paths = paths + self._members = members + self._cursor: np.ndarray = bounds[:-1].astype(np.int64) + self._ends = bounds[1:] + self._positions = positions + self._lines = lines + self._capacity = capacity + self._loaded: dict[int, dict[int, bytes]] = {} + self._sizes: dict[int, int] = {} + self._bytes = 0 + self._files: OrderedDict[int, BinaryIO] = OrderedDict() + self.loads = 0 + + def location(self, member_id: int, line_index: int) -> str: + file_index, offset, _ = self._members[member_id] + return f"{self._paths[file_index]}@{offset}:{line_index + 1}" + + def take(self, member_id: int, line_index: int) -> bytes: + lines = self._loaded.get(member_id) + if lines is None: + lines = self._load(member_id) + raw = lines.pop(line_index, None) + if raw is None: + raise WindowCutError( + "window_cut_failed", + f"window cut lost track of the row at " + f"{self.location(member_id, line_index)}", + ) + self._cursor[member_id] += 1 + size = len(raw) + _LINE_OVERHEAD + self._sizes[member_id] -= size + self._bytes -= size + if not lines: + del self._loaded[member_id] + del self._sizes[member_id] + elif self._bytes > self._capacity: + self._evict() + return raw + + def close(self) -> None: + for source in self._files.values(): + source.close() + self._files.clear() + + def _load(self, member_id: int) -> dict[int, bytes]: + file_index, offset, end = self._members[member_id] + data = self._read(file_index, offset, end) + payload = _decode_exact_member(data) + if payload is None: + raise _changed( + f"the gzip member at {self._paths[file_index]}@{offset} " + "no longer decodes" + ) + pieces = payload.split(b"\n") + wanted = self._lines[self._cursor[member_id] : self._ends[member_id]].tolist() + if max(wanted) >= len(pieces): + raise _changed( + f"the member at {self._paths[file_index]}@{offset} has fewer lines" + ) + lines = {index: pieces[index] for index in wanted} + size = sum(map(len, lines.values())) + _LINE_OVERHEAD * len(lines) + self._loaded[member_id] = lines + self._sizes[member_id] = size + self._bytes += size + self.loads += 1 + return lines + + def _read(self, file_index: int, offset: int, end: int) -> bytes: + source = self._files.pop(file_index, None) + if source is None: + if len(self._files) >= _OPEN_FILES: + _, oldest = self._files.popitem(last=False) + oldest.close() + source = open(self._paths[file_index], "rb") + self._files[file_index] = source + source.seek(offset) + data = source.read(end - offset) + if len(data) != end - offset: + raise _changed(f"{self._paths[file_index]} is shorter than when scanned") + return data + + def _evict(self) -> None: + target = self._capacity * 3 // 4 + furthest_first = sorted( + self._loaded, + key=lambda member_id: int(self._positions[self._cursor[member_id]]), + reverse=True, + ) + for member_id in furthest_first: + if self._bytes <= target: + return + del self._loaded[member_id] + self._bytes -= self._sizes.pop(member_id) diff --git a/components/src/dynamo/planner/retune/sweep.py b/components/src/dynamo/planner/retune/sweep.py index 1dc508dac62f..d8cf1521f201 100644 --- a/components/src/dynamo/planner/retune/sweep.py +++ b/components/src/dynamo/planner/retune/sweep.py @@ -171,6 +171,17 @@ class SweepRequest: ``retune.domains`` expands ``retune.planner.load_predictor``; its winner's predictor fields are fixed into every non-baseline candidate. ``objective`` is the metric ``score_summary`` scores each replay on. + + The remaining fields are set only by a live controller (``launch_gate`` + set) and keep their defaults in replay, where ``run_sweep`` ignores them: + ``router_domain`` is the controller's router domain (live knob name -> + values), ``deadline_s`` the absolute epoch deadline ``D_k = L_k + + timeout``, ``spec_replicas`` the DGD spec replicas by worker type the + gate reported, ``replicas_verified`` whether the recorded ready counts + cover the whole window (``False`` when the window began before the + controller started), and ``generation`` the launching controller's + generation. Abort events are passed to the sweep function, never stored + here, so the request stays picklable. """ window_index: int @@ -192,6 +203,11 @@ class SweepRequest: max_parallel: int timeout_s: float history_paths: tuple[str, ...] = () + router_domain: dict[str, tuple[float, ...]] | None = None + deadline_s: float | None = None + spec_replicas: dict[str, int] | None = None + replicas_verified: bool = True + generation: int = 0 def __post_init__(self) -> None: keys = [candidate.key() for candidate in self.candidates] @@ -207,6 +223,8 @@ def __post_init__(self) -> None: raise ValueError("max_parallel must be at least 1") if not self.timeout_s > 0: raise ValueError("timeout_s must be positive") + if self.deadline_s is not None and not math.isfinite(self.deadline_s): + raise ValueError("deadline_s must be finite") if self.predictor_candidates: labels = [label for label, _ in self.predictor_candidates] if len(set(labels)) != len(labels): @@ -238,7 +256,11 @@ class SweepResult: prepares it on ``warmup_observations`` as a host does before a swap; ``predictor_warmup_error`` is that preparation's error, if any. ``row_counts`` holds the window reader's counts and, when the history - was read, ``history_requests``. + was read, ``history_requests``. ``failure_reason`` is a stable, + machine-readable reason a live executor sets next to the free-text + ``error`` (``retune.live.contracts.NON_DECISION_FAILURES`` mean the + window was abandoned, not failed); ``details`` carries JSON-shaped + executor facts for the decision record. ``run_sweep`` sets neither. """ status: SweepStatus @@ -251,6 +273,8 @@ class SweepResult: predictor_warmup_error: str | None = None wall_s: float = 0.0 error: str | None = None + failure_reason: str | None = None + details: dict[str, Any] = field(default_factory=dict) def with_predictor_fields( diff --git a/components/src/dynamo/planner/tests/integration/test_live_retune_e2e.py b/components/src/dynamo/planner/tests/integration/test_live_retune_e2e.py new file mode 100644 index 000000000000..e601b19bf81a --- /dev/null +++ b/components/src/dynamo/planner/tests/integration/test_live_retune_e2e.py @@ -0,0 +1,1210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the retune imports. + +"""Live retune end to end: the real modules, faked only at the edges. + +Real: ``bootstrap.build_supervisor`` (frontend parse, spec premises, the P11 +canary), ``LiveRetuneSupervisor`` with its Lease leadership and state store, +``OrchestratorEngineAdapter`` with its ``RetuneController``, +``dgdr.run_dgdr_sweep``, the window cutter, ``RouterActuator`` and +``Reconciler``. + +Fake: the Kubernetes API (``FakeRetuneKube``, with ``FakeDgdrOperator`` +driving each DGDR v1beta2 Request to a Run and its Candidates), the +frontends' ``/router_policy`` (a revision-checked PATCH that merges, and a +reset to the startup values), the trace claim (a temporary directory of real ``jsonl.gz`` segments) and the +clocks. The test advances the clocks in 15 s steps; each step runs every +live planner the way ``NativePlannerBase`` and the supervisor's background +tasks do: lease round, the due planner tick (its decisions go to the +outbox), outbox round, control round (revalidation, canary, actuator, +reconciler). A launched sweep runs on the controller's sweep thread and the +test waits for it before the next step, so every run is deterministic. The +DGDR operator advances one phase per poll of the sweep. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import gzip +import json +import os +import time +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Optional + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from dynamo.llm import KvRouterConfig +from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.core.types import RetuneDecision, TickInput, WorkerCounts +from dynamo.planner.plugins.clock import WallClock +from dynamo.planner.plugins.orchestrator.engine_adapter import OrchestratorEngineAdapter +from dynamo.planner.retune.live import bootstrap +from dynamo.planner.retune.live.actuator import RouterActuator +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD_UID, + LABEL_SHARD, + LABEL_WINDOW, + DgdRef, + NullMetrics, + PodInfo, + PolicyCallResult, + RouterInfo, +) +from dynamo.planner.retune.live.reconciler import Reconciler +from dynamo.planner.retune.live.state import ( + RetuneState, + leadership_lease_name, + state_lease_name, +) +from dynamo.planner.retune.live.supervisor import LiveRetuneSupervisor +from dynamo.planner.retune.live.window_cut import cut_window, cut_window_inprocess +from dynamo.planner.tests.live_retune_fakes import ( + FakeDgdrOperator, + FakeRetuneKube, + default_metrics, +) +from dynamo.planner.tests.retune_stubs import caps + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.integration, + pytest.mark.planner, +] + +INTERVAL = 3600.0 +GRACE = 120.0 +TIMEOUT = 3000.0 +SCALE_INTERVAL = 60.0 +# Lease renewal (a third of the lease) and the actuator cadence both fit it. +STEP = 15.0 +LEASE = 60.0 +# Window n is [T0 + (n - 1) * INTERVAL, T0 + n * INTERVAL), launched GRACE +# after it ends. T0 is far ahead of the real clock: files the cutter writes +# (real mtimes) look old to the reconciler, and a real deadline derived from +# the simulated one is always in the future. +T0 = 3_600_000_000.0 +DGD_NAME = "graph" +K8S_NAMESPACE = "retune-e2e" +DGD_UID = "dgd-uid-1" +DYNAMO_NAMESPACE = "retune-e2e-graph" +CLAIM = "request-traces" +MODEL = "Qwen/Qwen3-0.6B" +FRONTENDS = ("graph-frontend-a", "graph-frontend-b") +WORKER_REPLICAS = 2 +REQUESTS_PER_POD = 12 +ROWS_PER_WINDOW = REQUESTS_PER_POD * len(FRONTENDS) +TRACE_SCHEMA = "dynamo.request.trace.v1" +TRACE_FILE_PATH = "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req" +CREDIT = "overlap_score_credit" +LOAD = "prefill_load_scale" +TEMP = "router_temperature" +STARTUP = {CREDIT: 1.0, LOAD: 1.0, TEMP: 0.0} +# The frontend runs ``--router-mode kv --router-replica-sync``. +FRONTEND_ROUTER_CONFIG = {"router_replica_sync": True} + + +def end_s(n: int) -> float: + return T0 + n * INTERVAL + + +def end_ms(n: int) -> int: + return int(round(end_s(n) * 1000.0)) + + +def launch_s(n: int) -> float: + return end_s(n) + GRACE + + +def decision_id(n: int) -> str: + return f"retune-w{end_ms(n)}" + + +def by_wt(point: Mapping[str, float]) -> dict[str, dict[str, float]]: + """``policy_parameters`` of the frontend's startup config with ``point``.""" + + config = KvRouterConfig.from_json(json.dumps({**FRONTEND_ROUTER_CONFIG, **point})) + return { + "aggregated": { + key: float(value) + for key, value in dict(config.policy_parameters("aggregated")).items() + } + } + + +def startup_map() -> dict[str, float]: + return by_wt({})["aggregated"] + + +# ------------------------------------------------------------- DGD, config + + +def dgd_object() -> dict[str, Any]: + """A DGD that satisfies every spec premise (P4, P6-P10).""" + + env = [ + { + "name": "DYN_POD_NAME", + "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}, + }, + {"name": "DYN_REQUEST_TRACE", "value": "1"}, + {"name": "DYN_REQUEST_TRACE_RECORDS", "value": "request_end"}, + {"name": "DYN_REQUEST_TRACE_SINKS", "value": "file"}, + {"name": "DYN_REQUEST_TRACE_FILE_FORMAT", "value": "jsonl_gz"}, + {"name": "DYN_REQUEST_TRACE_FILE_PATH", "value": TRACE_FILE_PATH}, + ] + frontend = { + "name": "Frontend", + "type": "frontend", + "replicas": len(FRONTENDS), + "podTemplate": { + "spec": { + "containers": [ + { + "name": "main", + "image": "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.8.0", + "command": ["python3"], + "args": [ + "-m", + "dynamo.frontend", + "--router-mode", + "kv", + "--router-replica-sync", + ], + "env": env, + "volumeMounts": [{"name": "traces", "mountPath": "/trace"}], + } + ], + "volumes": [ + {"name": "traces", "persistentVolumeClaim": {"claimName": CLAIM}} + ], + } + }, + } + worker = { + "name": "worker", + "type": "worker", + "replicas": WORKER_REPLICAS, + "podTemplate": { + "spec": { + "containers": [ + { + "name": "main", + "command": ["python3", "-m", "dynamo.vllm"], + "args": [f"--model={MODEL}"], + } + ] + } + }, + } + planner = { + "name": "Planner", + "type": "planner", + "replicas": 1, + "podTemplate": {"spec": {"containers": [{"name": "main"}]}}, + } + return { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": DGD_NAME, "namespace": K8S_NAMESPACE, "uid": DGD_UID}, + "spec": {"components": [frontend, worker, planner]}, + } + + +def planner_config(mount_path: Path, *, advisory: bool) -> PlannerConfig: + return PlannerConfig( + mode="agg", + environment="kubernetes", + advisory=advisory, + enable_load_scaling=False, + enable_throughput_scaling=False, + load_adjustment_interval_seconds=int(SCALE_INTERVAL), + scheduling={"scale_interval_seconds": SCALE_INTERVAL}, + retune={ + "interval_seconds": INTERVAL, + "grace_seconds": GRACE, + "timeout_seconds": TIMEOUT, + "router": { + "overlap_score_credit": {"choices": [0.5, 1.0]}, + "prefill_load_scale": 1.0, + "temperature": {"choices": [0.0, 0.2]}, + }, + "executor": "dgdr", + "dgdr": { + "claim_name": CLAIM, + "mount_path": str(mount_path), + "model": MODEL, + "sku": "h200_sxm", + "engine": { + "mode": "aggregated", + "backend": "vllm", + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + { + "tensor": 1, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, + } + ] + }, + "scheduler": { + "max_batched_tokens": 8192, + "max_sequences": 256, + }, + }, + "prefill": None, + "decode": None, + }, + }, + "run": { + "parallelism": 8, + "candidate_timeout_seconds": 600, + "max_shards": 3, + }, + "cut_budget_seconds": 300, + "job_startup_seconds": 60, + "keep_requests": 1, + "lease_duration_seconds": LEASE, + "actuator_interval_seconds": STEP, + }, + }, + ) + + +# ------------------------------------------------------------------ fakes + + +class SimClock: + """Wall and monotonic time, advanced only by the test.""" + + def __init__(self, wall: float) -> None: + self.wall = wall + self.mono = 1_000.0 + + def time(self) -> float: + return self.wall + + def monotonic(self) -> float: + return self.mono + + def advance(self, seconds: float) -> None: + self.wall += seconds + self.mono += seconds + + +async def inline(call): + return call() + + +async def no_sleep(_seconds: float) -> None: + return None + + +class RecordingMetrics(NullMetrics): + def __init__(self) -> None: + self.events: list[tuple[str, Any]] = [] + + def decision(self, outcome: str) -> None: + self.events.append(("decision", outcome)) + + def failure(self, reason: str) -> None: + self.events.append(("failure", reason)) + + def skipped(self, reason: str) -> None: + self.events.append(("skipped", reason)) + + def patch_error(self, code: str) -> None: + self.events.append(("patch_error", code)) + + def of(self, kind: str) -> list[Any]: + return [value for event, value in self.events if event == kind] + + +@dataclasses.dataclass +class FakeRouter: + """One frontend-local router: startup values, current values, revision.""" + + router_id: str + initial: dict[str, float] + parameters: dict[str, float] + revision: int = 0 + worker_type: str = "aggregated" + + def info(self) -> RouterInfo: + return RouterInfo( + router_id=self.router_id, + namespace=DYNAMO_NAMESPACE, + worker_type=self.worker_type, + revision=self.revision, + configuration=dict(self.parameters), + model=MODEL, + ) + + +@dataclasses.dataclass(frozen=True) +class PolicyWrite: + """One PATCH or reset that reached a frontend, and the persisted retune + state at that moment (to check persist-before-actuate).""" + + writer: str + kind: str + pod: str + router_id: str + expected_revision: int + parameters: Optional[dict[str, Any]] + status: int + state: Optional[RetuneState] + + +class FakeFrontends: + """The DGD's frontend pods and their routers.""" + + def __init__(self, kube: FakeRetuneKube) -> None: + self.kube = kube + self.routers: dict[str, list[FakeRouter]] = {} + self.writes: list[PolicyWrite] = [] + + def start(self, pod: str) -> None: + self.kube.add_pod(pod, dgd=DGD_NAME, component_type="frontend", ready=True) + self._register(pod) + + def restart(self, pod: str) -> None: + """A new frontend process in the pod (a container restart).""" + + self.kube.restart_pod(pod) + self._register(pod) + + def _register(self, pod: str) -> None: + """A new router id at revision 0 with the startup values + (``RouterPolicyRegistry::register``).""" + + startup = startup_map() + self.routers[pod] = [ + FakeRouter( + router_id=str(uuid.uuid4()), + initial=dict(startup), + parameters=dict(startup), + ) + ] + + def router(self, pod: str) -> FakeRouter: + (router,) = self.routers[pod] + return router + + def router_ids(self) -> set[str]: + return {self.router(pod).router_id for pod in FRONTENDS} + + def persisted_state(self) -> Optional[RetuneState]: + lease = self.kube.leases.get(state_lease_name(DGD_NAME)) + if lease is None: + return None + return RetuneState.from_annotations(lease["metadata"].get("annotations")) + + +class FakeRouterPolicyClient: + """``RouterPolicyClient`` over ``FakeFrontends``: a stale + ``expected_revision`` is 409; a PATCH replaces only the keys it names, a + reset restores every startup value, and both bump the revision.""" + + def __init__(self, frontends: FakeFrontends, writer: str) -> None: + self.frontends = frontends + self.writer = writer + self.gets = 0 + + async def get(self, pod: PodInfo) -> PolicyCallResult: + self.gets += 1 + routers = self.frontends.routers.get(pod.name) + if routers is None: + return PolicyCallResult(status=0, error="connection refused") + return PolicyCallResult( + status=200, routers=tuple(router.info() for router in routers) + ) + + async def patch( + self, + pod: PodInfo, + router_id: str, + expected_revision: int, + parameters: Mapping[str, float], + ) -> PolicyCallResult: + return self._update( + pod, "patch", router_id, expected_revision, dict(parameters) + ) + + async def reset( + self, pod: PodInfo, router_id: str, expected_revision: int + ) -> PolicyCallResult: + return self._update(pod, "reset", router_id, expected_revision, None) + + def _update( + self, + pod: PodInfo, + kind: str, + router_id: str, + expected_revision: int, + parameters: Optional[dict[str, Any]], + ) -> PolicyCallResult: + status, router, error = self._apply( + pod.name, router_id, expected_revision, parameters + ) + self.frontends.writes.append( + PolicyWrite( + writer=self.writer, + kind=kind, + pod=pod.name, + router_id=router_id, + expected_revision=expected_revision, + parameters=None if parameters is None else dict(parameters), + status=status, + state=self.frontends.persisted_state(), + ) + ) + if status != 200 or router is None: + return PolicyCallResult(status=status, error=error) + return PolicyCallResult(status=200, router=router.info()) + + def _apply( + self, + pod: str, + router_id: str, + expected_revision: int, + parameters: Optional[dict[str, Any]], + ) -> tuple[int, Optional[FakeRouter], Optional[str]]: + (router,) = [r for r in self.frontends.routers[pod] if r.router_id == router_id] + if expected_revision != router.revision: + return 409, None, "policy revision conflict" + if parameters is None: + router.parameters = dict(router.initial) + else: + router.parameters = {**router.parameters, **parameters} + router.revision += 1 + return 200, router, None + + +class TraceClaim: + """The RWX trace claim: ``/segments///req.NNNNNN.jsonl.gz``.""" + + def __init__(self, mount: Path) -> None: + self.mount = mount + self.segments = mount / "segments" / DGD_NAME + self.windows = mount / "windows" / DGD_NAME + self._index = 0 + for pod in FRONTENDS: + (self.segments / pod).mkdir(parents=True) + + def write_window(self, n: int, now_s: float) -> None: + """A new segment per frontend with window ``n``'s request_end records, + one gzip member per record as the frontend's gzip sink flushes.""" + + self._index += 1 + start = end_ms(n - 1) + span = end_ms(n) - start + for p, pod in enumerate(FRONTENDS): + path = self.segments / pod / f"req.{self._index:06d}.jsonl.gz" + with open(path, "wb") as sink: + for i in range(REQUESTS_PER_POD): + received = start + (i + 1) * span // (REQUESTS_PER_POD + 1) + p + total = 400 + 25 * i + record = { + "schema": TRACE_SCHEMA, + "event_type": "request_end", + "event_time_unix_ms": received + total, + "request": { + "request_id": f"w{n}-{pod}-{i}", + "model": MODEL, + "request_received_ms": received, + "total_time_ms": float(total), + "output_tokens": 16 + i, + "replay": { + "trace_block_size": 64, + "input_length": 256 + 8 * i, + "input_sequence_hashes": [n, p, i, 11 * i + 3], + }, + }, + } + envelope = {"timestamp": received + total, "event": record} + sink.write(gzip.compress(json.dumps(envelope).encode() + b"\n")) + os.utime(path, (now_s, now_s)) + + def segment_files(self, pod: str) -> list[str]: + return sorted(os.listdir(self.segments / pod)) + + def window_files(self) -> set[int]: + if not self.windows.is_dir(): + return set() + return { + int(name.split(".", 1)[0].split("-", 1)[0]) + for name in os.listdir(self.windows) + if name.endswith(".jsonl.gz") + } + + +# ---------------------------------------------------------------- planners + + +class PlannerHost: + """The slice of ``NativePlannerBase`` the supervisor calls back into.""" + + def __init__(self, engine: OrchestratorEngineAdapter, clock: SimClock) -> None: + self._engine = engine + self._config_lock = asyncio.Lock() + self._clock = clock + + async def replace_live_retune(self, inputs, *, generation): + async with self._config_lock: + self._engine.replace_retune( + inputs, generation=generation, now_s=self._clock.time() + ) + + +class Planner: + """One planner process: its supervisor, engine, actuator and reconciler.""" + + def __init__(self, world: World, name: str, *, advisory: bool) -> None: + self.world = world + self.name = name + self.identity = f"{name}/{uuid.uuid4().hex[:12]}" + self.config = planner_config(world.claim.mount, advisory=advisory) + self.metrics = RecordingMetrics() + self.client = FakeRouterPolicyClient(world.frontends, name) + self.decisions: list[RetuneDecision] = [] + self.router_parameters: Optional[dict[str, dict[str, float]]] = None + self.supervisor: LiveRetuneSupervisor + self.engine: OrchestratorEngineAdapter + self.host: PlannerHost + + async def boot(self) -> None: + """``NativePlannerBase._async_init`` then ``run``'s first tick.""" + + world = self.world + clock = world.clock + retune = self.config.retune + assert retune is not None and retune.dgdr is not None + reconciler = Reconciler( + world.kube, + world.dgd_ref, + retune.dgdr, + retune, + retune.dgdr.mount_path, + self.metrics, + now=clock.time, + ) + contexts = bootstrap.dgdr_context_factory(self.config, cut_window=world.cutter) + + def context_factory(supervisor: LiveRetuneSupervisor) -> Any: + # The executor's clock is the simulated one; each poll of the + # sweep lets the DGDR operator advance one phase. + return dataclasses.replace( + contexts(supervisor), + now=clock.time, + sleep=world.operator_sleep, + poll_interval_s=5.0, + ) + + self.supervisor = await bootstrap.build_supervisor( + self.config, + kube=world.kube, + dgd_name=DGD_NAME, + namespace_fn=lambda: DYNAMO_NAMESPACE, + metrics=self.metrics, + identity=self.identity, + actuator_factory=lambda dgd, maps: RouterActuator( + world.kube, self.client, dgd, maps, self.metrics + ), + reconciler=reconciler, + context_factory=context_factory, + clock=clock.time, + monotonic=clock.monotonic, + sleep=no_sleep, + blocking=inline, + ) + self.engine = OrchestratorEngineAdapter( + self.config, + caps(), + clock=WallClock(), + retune_inputs=self.supervisor.retune_inputs(), + retune_hooks=self.supervisor.hooks(), + ) + self.host = PlannerHost(self.engine, clock) + await self.supervisor.start(self.host, background=False) + self.next_tick = self.engine.initial_tick(clock.time()) + + async def step(self) -> None: + supervisor = self.supervisor + await supervisor.lease_round() + if self.world.clock.time() >= self.next_tick.at_s: + await self.tick() + await supervisor.outbox_round() + await supervisor.control_round() + # The reconciler runs as its own task; finish it to keep each step + # deterministic. + if supervisor.reconcile_task is not None: + await supervisor.reconcile_task + + async def tick(self) -> None: + """``NativePlannerBase._run_one_tick`` without the environment.""" + + counts = WorkerCounts( + ready_num_decode=WORKER_REPLICAS, expected_num_decode=WORKER_REPLICAS + ) + async with self.host._config_lock: + effects = await self.engine.tick( + self.next_tick, + TickInput(now_s=self.world.clock.time(), worker_counts=counts), + ) + self.supervisor.on_effects(effects) + assert effects.next_tick is not None + self.next_tick = effects.next_tick + self.decisions.extend(effects.retune_decisions) + self.router_parameters = effects.router_parameters + # Live, the sweep runs on the controller's thread while ticks go on; + # finishing it here keeps the interleaving deterministic. + controller = self.engine.retune_controller + pending = None if controller is None else controller._pending + if pending is not None and not pending.future.done(): + await asyncio.wait_for(asyncio.wrap_future(pending.future), timeout=60) + + async def stop(self) -> None: + """SIGTERM: ``_shutdown_runtime`` stops the supervisor (it aborts the + sweep and waits for the sweep thread), then the engine.""" + + self.world.planners.remove(self) + await self.supervisor.stop() + await self.engine.shutdown() + + def crash(self) -> None: + """SIGKILL: nothing of this process runs again; its Lease stays.""" + + self.world.planners.remove(self) + self.world.crashed.append(self) + + def decision(self, n: int) -> RetuneDecision: + (found,) = [d for d in self.decisions if d.decision_id == decision_id(n)] + return found + + +class World: + """The cluster, the frontends, the trace claim and the clocks.""" + + def __init__(self, tmp_path: Path, *, spawn_cutter: bool = False) -> None: + # Boot inside window 0: first enable skips it (its trace starts + # mid-way), so window 1 is the first one recorded in full. + self.clock = SimClock(T0 - 600.0) + self.kube = FakeRetuneKube(K8S_NAMESPACE, now=self.clock.time) + self.kube.put_dgd(dgd_object()) + self.dgd_ref = DgdRef(DGD_NAME, K8S_NAMESPACE, DGD_UID, DYNAMO_NAMESPACE) + self.frontends = FakeFrontends(self.kube) + for pod in FRONTENDS: + self.frontends.start(pod) + self.claim = TraceClaim(tmp_path / "claim") + # window end ms -> the router point (live names) the replay favours. + self.preferences: dict[int, Mapping[str, float]] = {} + self.operator = FakeDgdrOperator(self.kube, evaluate=self._evaluate) + self.spawn_cutter = spawn_cutter + self.cuts: list[tuple[int, int, bool]] = [] + self.polls = 0 + self.planners: list[Planner] = [] + self.crashed: list[Planner] = [] + self._next_trace_window = 1 + + # ----------------------------------------------------------- edges + + def prefer(self, n: int, point: Mapping[str, float]) -> None: + self.preferences[end_ms(n)] = dict(point) + + def _evaluate(self, request: dict, point: Mapping[str, float]) -> dict: + """Replay summary of a router point (AISimulate public names): the + window's preferred point scores best; deterministic, so the CRN + probe always matches its grid twin.""" + + window = int(request["metadata"]["labels"][LABEL_WINDOW]) + preferred = self.preferences.get(window, STARTUP) + distance = abs(point["overlap_score_credit"] - preferred[CREDIT]) + abs( + point["temperature"] - preferred[TEMP] + ) + metrics = default_metrics(point) + metrics["goodput_output_throughput_tok_s"] = 1000.0 - 200.0 * distance + return metrics + + def operator_sleep(self, _seconds: float) -> None: + self.polls += 1 + if self.polls > 10_000: + raise AssertionError("the DGDR sweep never finished polling") + self.operator.step() + + def cutter( + self, + segments_root, + start_ms, + end_ms_, + as_of_ms, + out_path, + *, + max_rows, + deadline_s, + cancel=None, + ): + if self.spawn_cutter: + # The child runs on the real clock; keep the remaining budget. + real_deadline = time.time() + (deadline_s - self.clock.time()) + result = cut_window( + segments_root, + start_ms, + end_ms_, + as_of_ms, + out_path, + max_rows=max_rows, + deadline_s=real_deadline, + cancel=cancel, + ) + else: + result = cut_window_inprocess( + segments_root, start_ms, end_ms_, as_of_ms, out_path, max_rows=max_rows + ) + self.cuts.append((int(end_ms_), result.num_requests, result.reused)) + return result + + # ------------------------------------------------------------ driving + + async def boot(self, name: str, *, advisory: bool = False) -> Planner: + planner = Planner(self, name, advisory=advisory) + await planner.boot() + self.planners.append(planner) + return planner + + async def run_until(self, until_s: float) -> None: + while self.clock.wall < until_s: + self.clock.advance(min(STEP, until_s - self.clock.wall)) + while self.clock.wall >= end_s(self._next_trace_window): + self.claim.write_window(self._next_trace_window, self.clock.wall) + self._next_trace_window += 1 + for planner in list(self.planners): + await planner.step() + + async def close(self) -> None: + for planner in list(self.planners): + await planner.stop() + for planner in self.crashed: + await planner.engine.shutdown() + + # ---------------------------------------------------------- inspection + + def state(self) -> RetuneState: + state = self.frontends.persisted_state() + assert state is not None + return state + + def history(self, decision: str) -> dict[str, Any]: + (entry,) = [e for e in self.state().history if e["decision_id"] == decision] + return entry + + def lease_spec(self) -> dict[str, Any]: + return self.kube.leases[leadership_lease_name(DGD_NAME)]["spec"] + + def requests_of(self, n: int) -> dict[str, dict]: + return { + name: obj + for name, obj in self.kube.requests.items() + if obj["metadata"]["labels"][LABEL_WINDOW] == str(end_ms(n)) + } + + def request_windows(self) -> set[int]: + return { + int(obj["metadata"]["labels"][LABEL_WINDOW]) + for obj in self.kube.requests.values() + } + + def created(self, n: Optional[int] = None) -> list[str]: + """Names of the stored (non dry-run) Request creates, in order.""" + + return [ + call.target["metadata"]["name"] + for call in self.kube.calls_of("create_request") + if not call.kwargs["dry_run"] + and ( + n is None + or call.target["metadata"]["labels"][LABEL_WINDOW] == str(end_ms(n)) + ) + ] + + def canaries(self) -> int: + return sum( + 1 for call in self.kube.calls_of("create_request") if call.kwargs["dry_run"] + ) + + def deleted(self) -> set[str]: + return {name for name, _ in self.kube.deleted_requests} + + def configurations(self) -> dict[str, tuple[int, dict[str, float]]]: + return { + pod: ( + self.frontends.router(pod).revision, + self.frontends.router(pod).parameters, + ) + for pod in FRONTENDS + } + + +def converged_to(world: World, point: Mapping[str, float], revision: int) -> None: + expected = by_wt(point)["aggregated"] + assert world.configurations() == {pod: (revision, expected) for pod in FRONTENDS} + + +@pytest.fixture +async def make_world(tmp_path): + worlds: list[World] = [] + + def make(**kwargs: Any) -> World: + world = World(tmp_path, **kwargs) + worlds.append(world) + return world + + yield make + for world in worlds: + await world.close() + + +# ------------------------------------------------------------------ tests + + +@pytest.mark.timeout(180) +async def test_advisory_records_recommendations_and_never_writes(make_world): + world = make_world(spawn_cutter=True) + recommended_1 = {CREDIT: 0.5, LOAD: 1.0, TEMP: 0.0} + recommended_2 = {CREDIT: 0.5, LOAD: 1.0, TEMP: 0.2} + world.prefer(1, recommended_1) + world.prefer(2, recommended_2) + planner = await world.boot("planner-a", advisory=True) + + await world.run_until(launch_s(1) - STEP) + supervisor = planner.supervisor + assert supervisor.is_leader() and supervisor.initialized and supervisor.ready() + state = world.state() + # First enable: an advisory record of the startup point; the window in + # progress at start-up (its trace began mid-way) counts as decided. + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "advisory", + False, + 1, + ) + assert state.flat == STARTUP and state.by_wt == by_wt(STARTUP) + assert state.last_window_end_ms == end_ms(0) + assert state.initials == { + router_id: startup_map() for router_id in world.frontends.router_ids() + } + assert planner.decisions == [] and world.created() == [] + assert world.canaries() >= 1 and world.kube.requests == {} + + # Window 1: swept through DGDR, recorded as a recommendation only. + await world.run_until(launch_s(1) + SCALE_INTERVAL + 2 * STEP) + decision = planner.decision(1) + assert decision.outcome == "advisory", decision + assert decision.previous.router == STARTUP + assert decision.chosen.router == recommended_1 + assert decision.num_requests == ROWS_PER_WINDOW + assert decision.failure_reason is None + assert decision.best_score > decision.baseline_score + details = decision.details + assert details["window"]["cut"] is True + assert details["window"]["num_requests"] == ROWS_PER_WINDOW + assert details["window"]["dropped_members"] == 0 + assert details["missing_candidates"] == 0 + assert details["crn_probe"]["relative_error"] == 0.0 + assert set(details["shards"]) == {"grid-0", "probe"} + assert world.cuts == [(end_ms(1), ROWS_PER_WINDOW, False)] + + state = world.state() + assert state.mode == "advisory" and state.flat == STARTUP + assert state.by_wt == by_wt(STARTUP) and state.deploy_seq == 1 + assert state.last_window_end_ms == end_ms(1) + entry = world.history(decision_id(1)) + assert entry["outcome"] == "advisory" + assert entry["flat"] == recommended_1 + assert entry["by_wt"] == by_wt(recommended_1) + assert supervisor.target().mode == "observe" + assert supervisor.last_report.converged + assert planner.router_parameters == by_wt(STARTUP) + + shards = { + obj["metadata"]["labels"][LABEL_SHARD] for obj in world.requests_of(1).values() + } + assert shards == {"grid-0", "probe"} + assert all( + obj["metadata"]["labels"][LABEL_DGD_UID] == DGD_UID + for obj in world.kube.requests.values() + ) + + # Window 2: another recommendation; window 1's Requests are cleaned up + # once window 2's are terminal (keep_requests: 1). + window_1_requests = set(world.requests_of(1)) + await world.run_until(launch_s(2) + SCALE_INTERVAL + 2 * STEP) + decision = planner.decision(2) + assert decision.outcome == "advisory" + assert decision.previous.router == STARTUP + assert decision.chosen.router == recommended_2 + assert world.history(decision_id(2))["by_wt"] == by_wt(recommended_2) + assert world.request_windows() == {end_ms(2)} + assert window_1_requests <= world.deleted() + assert world.claim.window_files() == {end_ms(2)} + + # Advisory never writes a router. + assert world.frontends.writes == [] + converged_to(world, STARTUP, revision=0) + assert planner.metrics.of("decision") == ["advisory", "advisory"] + assert planner.metrics.of("patch_error") == [] + + +@pytest.mark.timeout(180) +async def test_apply_restart_handover_and_advisory_rollback(make_world): + world = make_world() + applied_1 = {CREDIT: 0.5, LOAD: 1.0, TEMP: 0.0} + applied_3 = {CREDIT: 0.5, LOAD: 1.0, TEMP: 0.2} + world.prefer(1, applied_1) + world.prefer(2, applied_1) + world.prefer(3, applied_3) + world.prefer(4, applied_1) + pods = sorted(FRONTENDS) + + # ---------------------------------------- active: apply a better point + p1 = await world.boot("planner-a") + await world.run_until(launch_s(1) - STEP) + state = world.state() + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "active", + False, + 1, + ) + assert state.flat == STARTUP and state.last_window_end_ms == end_ms(0) + assert state.initials == { + router_id: startup_map() for router_id in world.frontends.router_ids() + } + assert world.frontends.writes == [] and p1.decisions == [] + + await world.run_until(launch_s(1) + SCALE_INTERVAL + 2 * STEP) + decision = p1.decision(1) + assert decision.outcome == "applied", decision + assert decision.previous.router == STARTUP + assert decision.chosen.router == applied_1 + assert decision.num_requests == ROWS_PER_WINDOW + assert decision.best_score > decision.baseline_score + assert decision.details["replicas"] == {"aggregated": WORKER_REPLICAS} + # The planner started before window 1, so its replicas were observed. + assert decision.details["replicas_verified"] is True + assert decision.details["crn_probe"]["relative_error"] == 0.0 + + state = world.state() + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "active", + False, + 2, + ) + assert state.flat == applied_1 and state.by_wt == by_wt(applied_1) + assert state.last_window_end_ms == end_ms(1) + assert world.history(decision_id(1))["by_wt"] == by_wt(applied_1) + + # The routers converge with explicit values, only for the changed key, + # and only after the decision is persisted. + writes = world.frontends.writes + assert sorted( + (w.pod, w.kind, w.expected_revision, w.parameters, w.status) for w in writes + ) == [(pod, "patch", 0, {CREDIT: 0.5}, 200) for pod in pods] + assert all(w.writer == "planner-a" for w in writes) + assert all(w.state.flat == applied_1 and w.state.deploy_seq == 2 for w in writes) + converged_to(world, applied_1, revision=1) + assert p1.supervisor.last_report.converged + assert p1.router_parameters == by_wt(applied_1) + window_1_requests = set(world.requests_of(1)) + assert { + obj["metadata"]["labels"][LABEL_SHARD] for obj in world.requests_of(1).values() + } == {"grid-0", "probe"} + assert sorted(world.created()) == sorted(window_1_requests) + + # A frontend restarts: its new router (revision 0, startup values) is + # latched and re-applied with explicit values; the persisted initials + # follow the new router id (the final rollback depends on them). + restarted = FRONTENDS[1] + gone = world.frontends.router(restarted).router_id + world.frontends.restart(restarted) + writes_before = len(world.frontends.writes) + await world.run_until(launch_s(1) + SCALE_INTERVAL + 4 * STEP) + reapplied = world.frontends.writes[writes_before:] + assert [ + (w.pod, w.kind, w.expected_revision, w.parameters, w.status) for w in reapplied + ] == [(restarted, "patch", 0, {CREDIT: 0.5}, 200)] + converged_to(world, applied_1, revision=1) + state = world.state() + assert gone not in state.initials + assert state.initials == { + router_id: startup_map() for router_id in world.frontends.router_ids() + } + assert state.deploy_seq == 2 + + # -------------------------------- restart: restore, never roll back + restart_at = launch_s(1) + 600.0 + await world.run_until(restart_at) + await p1.stop() + assert not world.lease_spec().get("holderIdentity") # released + writes_before = len(world.frontends.writes) + p2 = await world.boot("planner-b") + await world.run_until(restart_at + 5 * STEP) + assert p2.supervisor.is_leader() and p2.supervisor.initialized + assert world.lease_spec()["holderIdentity"] == p2.identity + assert world.lease_spec()["leaseTransitions"] == 1 + state = world.state() + assert state.leader_term.startswith(f"1/{p2.identity}/") + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "active", + False, + 2, + ) + assert state.flat == applied_1 and state.by_wt == by_wt(applied_1) + controller = p2.engine.retune_controller + assert controller is not None and controller.current_point().router == applied_1 + assert p2.router_parameters == by_wt(applied_1) + assert len(world.frontends.writes) == writes_before + converged_to(world, applied_1, revision=1) + + # The decided window is not swept again. + await world.run_until(launch_s(2) - STEP) + assert p2.decisions == [] + assert sorted(world.created()) == sorted(window_1_requests) + assert [cut[0] for cut in world.cuts] == [end_ms(1)] + + # Window 2: the restarted planner keeps the applied point. + await world.run_until(launch_s(2) + SCALE_INTERVAL + 2 * STEP) + kept = p2.decision(2) + assert kept.outcome == "kept", kept + assert kept.previous.router == applied_1 and kept.chosen.router == applied_1 + assert kept.best_score == kept.baseline_score + state = world.state() + assert state.deploy_seq == 2 and state.last_window_end_ms == end_ms(2) + assert len(world.frontends.writes) == writes_before + # Window 1's Requests and window file are gone (keep_requests: 1). + assert world.request_windows() == {end_ms(2)} + assert window_1_requests <= world.deleted() + assert world.claim.window_files() == {end_ms(2)} + window_2_requests = set(world.requests_of(2)) + + # ---------------------------- leadership handover (holder crashes) + p3 = await world.boot("planner-c") + await world.run_until(launch_s(3)) + # The leader launched window 3 (its Requests all succeeded); the + # standby never launched anything. + assert not p3.supervisor.is_leader() and not p3.supervisor.initialized + assert p3.decisions == [] and "not_leader" in p3.metrics.of("skipped") + window_3_requests = set(world.requests_of(3)) + assert window_3_requests and set(world.created(3)) == window_3_requests + assert [d.decision_id for d in p2.decisions] == [decision_id(2)] + transitions = world.lease_spec()["leaseTransitions"] + p2.crash() + await world.run_until(launch_s(3) + LEASE + 2 * SCALE_INTERVAL + 2 * STEP) + assert p3.supervisor.is_leader() and p3.supervisor.initialized + assert world.lease_spec()["holderIdentity"] == p3.identity + assert world.lease_spec()["leaseTransitions"] == transitions + 1 + + decision = p3.decision(3) + assert decision.outcome == "applied", decision + assert decision.previous.router == applied_1 + assert decision.chosen.router == applied_3 + assert decision.num_requests == ROWS_PER_WINDOW + # The new leader adopted the crashed leader's Requests: no new create, + # no second cut. + assert all(shard["adopted"] for shard in decision.details["shards"].values()) + assert decision.details["window"]["cut"] is False + assert set(world.created(3)) == window_3_requests + assert [cut[0] for cut in world.cuts] == [end_ms(1), end_ms(2), end_ms(3)] + + state = world.state() + assert state.leader_term.startswith(f"{transitions + 1}/{p3.identity}/") + assert (state.mode, state.deploy_seq) == ("active", 3) + assert state.flat == applied_3 and state.by_wt == by_wt(applied_3) + new_writes = world.frontends.writes[writes_before:] + assert sorted( + (w.pod, w.kind, w.expected_revision, w.parameters, w.status) for w in new_writes + ) == [(pod, "patch", 1, {TEMP: 0.2}, 200) for pod in pods] + assert all( + w.writer == "planner-c" and w.state.flat == applied_3 for w in new_writes + ) + converged_to(world, applied_3, revision=2) + + # ----------------------- switch to advisory: two-phase rollback + switch_at = launch_s(3) + 600.0 + await world.run_until(switch_at) + await p3.stop() + writes_before = len(world.frontends.writes) + p4 = await world.boot("planner-d", advisory=True) + await world.run_until(switch_at + 2 * STEP) + state = world.state() + # Phase 1: rollback_pending persisted, then the routers are reset. + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "advisory", + True, + 4, + ) + assert state.flat == STARTUP and state.by_wt == by_wt(STARTUP) + resets = world.frontends.writes[writes_before:] + assert sorted( + (w.pod, w.kind, w.expected_revision, w.parameters, w.status) for w in resets + ) == [(pod, "reset", 2, None, 200) for pod in pods] + assert all( + w.writer == "planner-d" + and w.state.rollback_pending + and w.state.mode == "advisory" + and w.state.deploy_seq == 4 + for w in resets + ) + converged_to(world, STARTUP, revision=3) + assert p4.supervisor.target().mode == "rollback" + + # Phase 2: converged for a lease duration -> rollback complete. + await world.run_until(switch_at + 2 * STEP + LEASE + 3 * STEP) + state = world.state() + assert (state.mode, state.rollback_pending, state.deploy_seq) == ( + "advisory", + False, + 4, + ) + last = state.history[-1] + assert (last["outcome"], last["deploy_seq"]) == ("rolled_back", 4) + assert p4.supervisor.target().mode == "observe" + assert len(world.frontends.writes) == writes_before + len(pods) + + # Window 4 under advisory: a recommendation, no write. + await world.run_until(launch_s(4) + SCALE_INTERVAL + 2 * STEP) + decision = p4.decision(4) + assert decision.outcome == "advisory", decision + assert decision.previous.router == STARTUP + assert decision.chosen.router == applied_1 + state = world.state() + assert state.flat == STARTUP and state.last_window_end_ms == end_ms(4) + assert len(world.frontends.writes) == writes_before + len(pods) + converged_to(world, STARTUP, revision=3) + + # Clean-up across the four windows: only the newest window's Requests + # and window file remain, and old segments are gone (the newest of each + # pod is never deleted). + assert world.request_windows() == {end_ms(4)} + assert window_1_requests | window_2_requests | window_3_requests <= world.deleted() + assert world.claim.window_files() == {end_ms(4)} + for pod in FRONTENDS: + assert world.claim.segment_files(pod) == [ + "req.000003.jsonl.gz", + "req.000004.jsonl.gz", + ] + for planner in (p1, p2, p3, p4): + assert planner.metrics.of("patch_error") == [] diff --git a/components/src/dynamo/planner/tests/live_retune_fakes.py b/components/src/dynamo/planner/tests/live_retune_fakes.py new file mode 100644 index 000000000000..a026b23e3019 --- /dev/null +++ b/components/src/dynamo/planner/tests/live_retune_fakes.py @@ -0,0 +1,1144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""In-memory fakes of the Kubernetes side of live retune. + +``FakeRetuneKube`` implements ``contracts.RetuneKube`` over plain dicts with +the API-server behaviour the live modules depend on: ``resourceVersion`` +compare-and-swap on Lease replace, ``AlreadyExists`` on create, Strict field +validation and silent pruning on Request create, label selectors, cascading +Request deletion, and injectable faults (including writes that time out but +land). ``FakeDgdrOperator`` drives the Requests in a ``FakeRetuneKube`` +through the DGDR v1beta2 lifecycle (Accepted -> Run -> Completed with +Candidates) the way the operator at #17483 does, lowering each Request into +``Run.status.aisimulate.input`` with the same merge-patch rules. The +operator's golden +``deploy/operator/internal/aisimulate/testdata/trace_request_lowering.golden.json`` +is the oracle for that lowering; ``test_live_retune_dgdr`` checks +``lower_request_input`` against it. + +This is a plain module, not a conftest: tests import what they need. +""" + +from __future__ import annotations + +import copy +import functools +import itertools +import json +import re +import threading +import time +import uuid +from collections import deque +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field, replace +from typing import Any, Literal, Optional + +from dynamo.planner.retune.live.contracts import ( + DGDR_ROUTER_KNOBS, + POD_LABEL_COMPONENT_TYPE, + POD_LABEL_DGD_NAME, + RUN_UID_LABEL, + KubeAlreadyExists, + KubeConflict, + KubeError, + KubeInvalid, + KubeNotFound, + PodInfo, +) + +# Operator constants (deploy/operator/internal/controller at #17483). +REQUEST_GENERATION_ANNOTATION = "nvidia.com/dgdr-request-generation" +CANDIDATE_ID_LABEL = "nvidia.com/dgdr-candidate-id" +TRACE_MOUNT_PATH = "/var/run/aisimulate/trace" +DEFAULT_RUNTIME_IMAGE = "nvcr.io/nvidia/ai-dynamo/dynamo-planner:1.8.0" +DEFAULT_DYNAMO_VERSION = "1.8.0" +ROUTER_ADAPTER = "dynamo.router" +# spec.workload.trace.format: the CRD default, and each API value's +# AISimulate traffic.source.format spelling. +DEFAULT_TRACE_FORMAT = "Dynamo" +_TRACE_FORMATS = {"Dynamo": "dynamo"} + +_METRIC_TARGETS = { + "Throughput": "throughput", + "ThroughputPerGPU": "throughput_per_gpu", + "ThroughputPerUser": "throughput_per_user", + "TTFT": "ttft", + "E2ELatency": "e2e_latency", + "Goodput": "goodput", + "GoodputPerGPU": "goodput_per_gpu", +} + +KubeMethod = Literal[ + "get_dgd", + "list_pods", + "create_request", + "get_request", + "list_requests", + "delete_request", + "get_run", + "list_candidates", + "get_lease", + "create_lease", + "replace_lease", + "delete_lease", +] + + +# ------------------------------------------------------------------ helpers + + +def merge_patch(target: Any, patch: Any) -> Any: + """RFC 7386 JSON merge patch; returns a new value, never mutates.""" + + if not isinstance(patch, Mapping): + return copy.deepcopy(patch) + result = dict(target) if isinstance(target, Mapping) else {} + for key, value in patch.items(): + if value is None: + result.pop(key, None) + else: + result[key] = merge_patch(result.get(key), value) + return result + + +def nested_get(obj: Any, *path: str, default: Any = None) -> Any: + for key in path: + if not isinstance(obj, Mapping) or key not in obj: + return default + obj = obj[key] + return obj + + +def _nested_set(obj: dict, path: tuple[str, ...], value: Any) -> None: + for key in path[:-1]: + child = obj.get(key) + if not isinstance(child, dict): + child = {} + obj[key] = child + obj = child + obj[path[-1]] = value + + +def _nested_pop(obj: Any, path: tuple[str, ...]) -> bool: + for key in path[:-1]: + if not isinstance(obj, dict) or key not in obj: + return False + obj = obj[key] + if isinstance(obj, dict) and path[-1] in obj: + del obj[path[-1]] + return True + return False + + +def _split(path: str) -> tuple[str, ...]: + return tuple(path.split(".")) + + +def rfc3339(epoch_s: float) -> str: + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(epoch_s)) + + +_SELECTOR_TERM = re.compile( + r"^\s*(!?)\s*([A-Za-z0-9._/-]+)\s*(?:(==|=|!=)\s*(\S*))?\s*$" +) + + +def match_labels(labels: Mapping[str, str] | None, selector: str) -> bool: + """Equality-based label selector: ``k=v``, ``k==v``, ``k!=v``, ``k``, ``!k``.""" + + labels = labels or {} + if not selector.strip(): + return True + for term in selector.split(","): + found = _SELECTOR_TERM.match(term) + if found is None: + raise ValueError(f"unsupported label selector term {term!r}") + negate, key, op, value = found.groups() + if op is None: + present = key in labels + if present == bool(negate): + return False + elif negate: + raise ValueError(f"unsupported label selector term {term!r}") + elif op in ("=", "=="): + if labels.get(key) != value: + return False + elif labels.get(key) == value: # != + return False + return True + + +def lower_request_input( + request: Mapping[str, Any], + *, + run_name: str, + runtime_image: str = DEFAULT_RUNTIME_IMAGE, + dynamo_version: str = DEFAULT_DYNAMO_VERSION, + gpus_per_node: int = 8, +) -> dict[str, Any]: + """``Run.status.aisimulate.input`` for a Request, as the operator lowers it. + + Mirrors ``aisimulate.LowerWithUnderlay`` and ``dgdrRunInput`` at #17483: + the public role underlay, the ``dgdr_run`` static underlay, the typed + underlay, the derived defaults, ``inputOverrides`` as the final merge + patch, then the operator-owned ``dgdr_run`` fields. ``inputOverrides`` + wins over the typed trace source: nothing restores it afterwards. For + the golden's ``request_spec`` it returns the golden's ``run_input``. + """ + + spec = request.get("spec") or {} + max_candidates = int(nested_get(spec, "recommendation", "maxCandidates", default=1)) + underlay: dict[str, Any] = { + "engine": {"workers": {"aggregated": {}, "prefill": {}, "decode": {}}} + } + underlay = merge_patch( + underlay, {"dgdr_run": {"name": run_name, "max_candidates": max_candidates}} + ) + typed: dict[str, Any] = {} + model = nested_get(spec, "model", "name") + if model is not None: + _nested_set(typed, ("engine", "model"), model) + if spec.get("backends"): + _nested_set(typed, ("engine", "backend", "choices"), list(spec["backends"])) + sku = nested_get(spec, "hardware", "gpu", "sku") + if sku: + _nested_set(typed, ("engine", "hardware"), sku) + budget = nested_get(spec, "hardware", "gpu", "budget") + if budget is not None: + _nested_set( + typed, ("optimization", "constraints", "max_candidate_gpus"), int(budget) + ) + trace = nested_get(spec, "workload", "trace") + if trace is not None: + sub_path = nested_get( + trace, "source", "persistentVolumeClaim", "subPath", default="" + ) + mounted = TRACE_MOUNT_PATH + ( + "/trace.jsonl.gz" if sub_path.endswith(".gz") else "/trace.jsonl" + ) + api_format = trace.get("format", DEFAULT_TRACE_FORMAT) + if api_format not in _TRACE_FORMATS: + raise ValueError(f"unsupported trace format {api_format!r}") + trace_source = { + "type": "trace", + "format": _TRACE_FORMATS[api_format], + "paths": [mounted], + } + _nested_set(typed, ("traffic", "source"), trace_source) + _nested_set(typed, ("traffic", "load", "type"), "trace_timestamps") + speedup = nested_get(trace, "replay", "arrivalSpeedupRatio") + if speedup is not None: + _nested_set(typed, ("traffic", "load", "speedup"), speedup) + objective = spec.get("objective") or {} + if objective.get("metric") is not None: + _nested_set( + typed, ("optimization", "target"), _METRIC_TARGETS[objective["metric"]] + ) + sla = objective.get("sla") or {} + for api_key, input_key in ( + ("ttftMs", "ttft_ms"), + ("itlMs", "itl_ms"), + ("e2eMs", "e2e_ms"), + ): + if sla.get(api_key) is not None: + _nested_set(typed, ("evaluation", "sla", input_key), sla[api_key]) + underlay = merge_patch(underlay, typed) + if "traffic" in underlay: + source_type = nested_get(underlay, "traffic", "source", "type") + if source_type != "trace" and nested_get(underlay, "traffic", "stop") is None: + _nested_set(underlay, ("traffic", "stop", "requests"), 1000.0) + if nested_get(underlay, "optimizer", "max_trials") is None: + _nested_set(underlay, ("optimizer", "max_trials"), 320.0) + if nested_get(underlay, "evaluation", "sla") and ( + nested_get(underlay, "optimization", "strict_sla") is None + ): + _nested_set(underlay, ("optimization", "strict_sla"), True) + overrides = nested_get(spec, "aisimulate", "inputOverrides", default={}) or {} + effective = merge_patch(underlay, overrides) + reserved = { + "name": run_name, + "max_candidates": max_candidates, + "snapshot_dir": "/var/run/aisimulate/output", + "namespace": nested_get(request, "metadata", "namespace"), + "runtime_image": runtime_image, + "runtime_version_override": dynamo_version, + "derive_runtime_images": True, + "num_gpus_per_node": gpus_per_node, + } + for key, value in reserved.items(): + _nested_set(effective, ("dgdr_run", key), value) + # Round-trip through JSON like the API server does. + return json.loads(json.dumps(effective)) + + +def router_points(effective_input: Mapping[str, Any]) -> list[dict[str, float]]: + """The router points (public knob names) an input's router section spans.""" + + router = effective_input.get("router") or {} + axes: list[tuple[str, list[float]]] = [] + for public in DGDR_ROUTER_KNOBS: + value = router.get(public) + if isinstance(value, Mapping): + axes.append((public, [float(v) for v in value["choices"]])) + elif value is not None: + axes.append((public, [float(value)])) + return [ + dict(zip([name for name, _ in axes], values)) + for values in itertools.product(*(values for _, values in axes)) + ] + + +def engine_parameters(effective_input: Mapping[str, Any]) -> dict[str, Any]: + """The non-router parameters every Candidate of a pinned engine shares.""" + + engine = effective_input.get("engine") or {} + workers = engine.get("workers") or {} + params: dict[str, Any] = { + "backend": engine.get("backend"), + "deployment_mode": "disagg" if engine.get("mode") == "disaggregated" else "agg", + "model": engine.get("model"), + "hardware": engine.get("hardware"), + } + for role, worker in sorted(workers.items()): + if isinstance(worker, Mapping): + preset = nested_get(worker, "parallelism", "preset", default=[{}]) + params[f"{role}_parallelism"] = dict(preset[0]) if preset else {} + return json.loads(json.dumps(params)) + + +def candidate_parameters( + effective_input: Mapping[str, Any], point_public: Mapping[str, float] +) -> dict[str, Any]: + """``Candidate.spec.parameters`` of one evaluated router point.""" + + router = effective_input.get("router") or {} + return { + **engine_parameters(effective_input), + "adapters": { + ROUTER_ADAPTER: { + "policy": router.get("policy", "kv_router"), + "prefill_load_model": dict( + router.get("prefill_load_model") or {"type": "none"} + ), + **{name: float(value) for name, value in point_public.items()}, + } + }, + } + + +def default_metrics(point_public: Mapping[str, float]) -> dict[str, float]: + """Deterministic replay summary metrics of a router point (for tests).""" + + credit = point_public.get("overlap_score_credit", 1.0) + scale = point_public.get("prefill_load_scale", 1.0) + temperature = point_public.get("temperature", 0.0) + goodput = 1000.0 + 100.0 * credit - 3.0 * abs(scale - 2.0) - 10.0 * temperature + return { + "goodput_output_throughput_tok_s": goodput, + "output_throughput_tok_s": goodput * 1.1, + "gpu_hours": 2.0, + "duration_ms": 3_600_000.0, + "mean_ttft_ms": 100.0 + scale, + "num_ttft_samples": 1000.0, + "mean_e2e_latency_ms": 1000.0 + scale, + "num_e2e_latency_samples": 1000.0, + "mean_output_token_throughput_per_user": 50.0 + credit, + } + + +# ------------------------------------------------------------------- faults + + +@dataclass +class Fault: + """One injected failure of a ``FakeRetuneKube`` method. + + ``landed=True`` performs the call and then raises (a write that timed + out but reached the API server). ``match`` limits the fault to calls + whose first positional argument (a name, selector or body) satisfies it. + """ + + error: Exception + times: int = 1 + landed: bool = False + match: Optional[Callable[[Any], bool]] = None + + +@dataclass +class KubeCall: + method: str + target: Any + kwargs: dict[str, Any] = field(default_factory=dict) + + +# --------------------------------------------------------------------- kube + + +class FakeRetuneKube: + """In-memory ``contracts.RetuneKube`` (thread-safe). + + Stores DGDs, pods (``PodInfo``), Requests, Runs, Candidates and Leases + as plain dicts. Objects come back as deep copies, so callers can never + alias the store. ``now`` stamps ``creationTimestamp``. + + Request create honours ``field_validation='Strict'``: a path listed in + ``unknown_fields`` (dotted, e.g. ``spec.workload.trace``) is rejected + with ``KubeInvalid``; a path in ``pruned_fields`` is dropped silently + (an API server that ignores the validation directive). Like the CRD + schema, create defaults ``spec.workload.trace.format`` to ``Dynamo``. + ``reject_create`` rejects every create with ``KubeInvalid``. + ``on_request_created`` hooks run after each stored (non dry-run) create. + """ + + def __init__( + self, + namespace: str = "test-ns", + *, + now: Callable[[], float] = time.time, + ) -> None: + self.namespace = namespace + self.now = now + self.dgds: dict[str, dict] = {} + self.pods: dict[str, PodInfo] = {} + self.requests: dict[str, dict] = {} + self.runs: dict[str, dict] = {} + # Each Run's lowered input, recorded in its status by ``create_job``. + self.run_inputs: dict[str, dict] = {} + self.candidates: dict[str, dict] = {} + self.leases: dict[str, dict] = {} + self.calls: list[KubeCall] = [] + self.deleted_requests: list[tuple[str, str]] = [] + self.unknown_fields: set[str] = set() + self.pruned_fields: set[str] = set() + self.reject_create: Optional[str] = None + self.on_request_created: list[Callable[[FakeRetuneKube, dict], None]] = [] + self.cascade_deletes = True + self._faults: dict[str, deque[Fault]] = {} + self._rv = itertools.count(1) + self._lock = threading.RLock() + + # ------------------------------------------------------------ faults + + def fail( + self, + method: KubeMethod, + error: Exception | None = None, + *, + times: int = 1, + landed: bool = False, + match: Optional[Callable[[Any], bool]] = None, + ) -> None: + """Make the next ``times`` matching calls of ``method`` raise ``error``.""" + + with self._lock: + self._faults.setdefault(method, deque()).append( + Fault( + error=error or KubeError(f"injected {method} failure"), + times=times, + landed=landed, + match=match, + ) + ) + + def clear_faults(self) -> None: + with self._lock: + self._faults.clear() + + def calls_of(self, method: str) -> list[KubeCall]: + with self._lock: + return [call for call in self.calls if call.method == method] + + def _call( + self, + method: str, + target: Any, + kwargs: dict[str, Any], + operation: Callable[[], Any], + ) -> Any: + with self._lock: + self.calls.append(KubeCall(method, copy.deepcopy(target), dict(kwargs))) + fault = self._take_fault(method, target) + if fault is not None and not fault.landed: + raise fault.error + result = operation() + if fault is not None: + raise fault.error + return result + + def _take_fault(self, method: str, target: Any) -> Optional[Fault]: + queue = self._faults.get(method) + if not queue: + return None + for fault in list(queue): + if fault.match is None or fault.match(target): + fault.times -= 1 + if fault.times <= 0: + queue.remove(fault) + return fault + return None + + def _stamp(self, obj: dict, *, kind: str, api_version: str) -> dict: + metadata = obj.setdefault("metadata", {}) + obj.setdefault("apiVersion", api_version) + obj.setdefault("kind", kind) + metadata.setdefault("namespace", self.namespace) + metadata.setdefault("uid", str(uuid.uuid4())) + metadata.setdefault("creationTimestamp", rfc3339(self.now())) + metadata.setdefault("generation", 1) + metadata["resourceVersion"] = str(next(self._rv)) + return obj + + def bump(self, obj: dict) -> None: + """Give a stored object a new ``resourceVersion``.""" + + obj["metadata"]["resourceVersion"] = str(next(self._rv)) + + # --------------------------------------------------------------- DGD + + def put_dgd(self, dgd: dict) -> dict: + with self._lock: + stored = copy.deepcopy(dgd) + self._stamp( + stored, kind="DynamoGraphDeployment", api_version="nvidia.com/v1beta1" + ) + self.dgds[stored["metadata"]["name"]] = stored + return copy.deepcopy(stored) + + def get_dgd(self, name: str, *, timeout_s: float = 10.0) -> dict: + def operation() -> dict: + if name not in self.dgds: + raise KubeNotFound(f"DynamoGraphDeployment {name!r} not found") + return copy.deepcopy(self.dgds[name]) + + return self._call("get_dgd", name, {"timeout_s": timeout_s}, operation) + + # -------------------------------------------------------------- pods + + def add_pod( + self, + name: str, + *, + dgd: str | None = None, + component_type: str = "frontend", + ready: bool = True, + terminating: bool = False, + ip: str | None = None, + http_port: int | None = 8000, + labels: Mapping[str, str] | None = None, + restarts: int | None = 0, + ) -> PodInfo: + all_labels = dict(labels or {}) + if dgd is not None: + all_labels.setdefault(POD_LABEL_DGD_NAME, dgd) + if component_type: + all_labels.setdefault(POD_LABEL_COMPONENT_TYPE, component_type) + pod = PodInfo( + name=name, + uid=str(uuid.uuid4()), + ip=ip if ip is not None else f"10.0.0.{len(self.pods) + 1}", + ready=ready, + terminating=terminating, + labels=all_labels, + http_port=http_port, + restarts=restarts, + ) + with self._lock: + self.pods[name] = pod + return pod + + def restart_pod(self, name: str) -> PodInfo: + """The pod's ``main`` container restarted.""" + + with self._lock: + pod = self.pods[name] + pod = replace(pod, restarts=(pod.restarts or 0) + 1) + self.pods[name] = pod + return pod + + def remove_pod(self, name: str) -> None: + with self._lock: + self.pods.pop(name, None) + + def list_pods( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[PodInfo]: + def operation() -> list[PodInfo]: + return [ + pod + for pod in self.pods.values() + if match_labels(pod.labels, label_selector) + ] + + return self._call( + "list_pods", label_selector, {"timeout_s": timeout_s}, operation + ) + + # ----------------------------------------------------------- Requests + + def create_request( + self, + body: dict, + *, + dry_run: bool = False, + timeout_s: float = 10.0, + ) -> dict: + def operation() -> dict: + if self.reject_create is not None: + raise KubeInvalid(self.reject_create) + stored = copy.deepcopy(body) + for path in sorted(self.unknown_fields): + if nested_get(stored, *_split(path)) is not None: + raise KubeInvalid(f'strict decoding error: unknown field "{path}"') + for path in self.pruned_fields: + _nested_pop(stored, _split(path)) + trace = nested_get(stored, "spec", "workload", "trace") + if isinstance(trace, dict): + trace.setdefault("format", DEFAULT_TRACE_FORMAT) + name = nested_get(stored, "metadata", "name") + if not name: + raise KubeInvalid("metadata.name is required") + if name in self.requests: + raise KubeAlreadyExists( + f"dynamographdeploymentrequests {name!r} already exists" + ) + self._stamp( + stored, + kind="DynamoGraphDeploymentRequest", + api_version="nvidia.com/v1beta2", + ) + if dry_run: + return stored + self.requests[name] = stored + for hook in list(self.on_request_created): + hook(self, copy.deepcopy(stored)) + return copy.deepcopy(stored) + + return self._call( + "create_request", + body, + {"dry_run": dry_run, "timeout_s": timeout_s}, + operation, + ) + + def get_request(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + def operation() -> Optional[dict]: + found = self.requests.get(name) + return None if found is None else copy.deepcopy(found) + + return self._call("get_request", name, {"timeout_s": timeout_s}, operation) + + def list_requests( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + def operation() -> list[dict]: + return [ + copy.deepcopy(obj) + for obj in self.requests.values() + if match_labels(obj["metadata"].get("labels"), label_selector) + ] + + return self._call( + "list_requests", label_selector, {"timeout_s": timeout_s}, operation + ) + + def delete_request( + self, + name: str, + *, + propagation: Literal["Background", "Foreground"] = "Background", + timeout_s: float = 10.0, + ) -> None: + def operation() -> None: + request = self.requests.pop(name, None) + if request is None: + return + self.deleted_requests.append((name, propagation)) + if not self.cascade_deletes: + return + request_uid = request["metadata"]["uid"] + for run_name, run in list(self.runs.items()): + owners = run["metadata"].get("ownerReferences") or [] + if any(owner.get("uid") == request_uid for owner in owners): + self._delete_run(run_name) + + return self._call( + "delete_request", + name, + {"propagation": propagation, "timeout_s": timeout_s}, + operation, + ) + + # --------------------------------------------------- Runs, Candidates + + def _delete_run(self, run_name: str) -> None: + run = self.runs.pop(run_name, None) + self.run_inputs.pop(run_name, None) + if run is None: + return + run_uid = run["metadata"]["uid"] + for candidate_name, candidate in list(self.candidates.items()): + labels = candidate["metadata"].get("labels") or {} + if labels.get(RUN_UID_LABEL) == run_uid: + del self.candidates[candidate_name] + + def get_run(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + def operation() -> Optional[dict]: + found = self.runs.get(name) + return None if found is None else copy.deepcopy(found) + + return self._call("get_run", name, {"timeout_s": timeout_s}, operation) + + def list_candidates( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[dict]: + def operation() -> list[dict]: + return [ + copy.deepcopy(obj) + for obj in self.candidates.values() + if match_labels(obj["metadata"].get("labels"), label_selector) + ] + + return self._call( + "list_candidates", label_selector, {"timeout_s": timeout_s}, operation + ) + + # ------------------------------------------------------------- Leases + + def get_lease(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + def operation() -> Optional[dict]: + found = self.leases.get(name) + return None if found is None else copy.deepcopy(found) + + return self._call("get_lease", name, {"timeout_s": timeout_s}, operation) + + def create_lease(self, body: dict, *, timeout_s: float = 10.0) -> dict: + def operation() -> dict: + stored = copy.deepcopy(body) + name = nested_get(stored, "metadata", "name") + if not name: + raise KubeInvalid("metadata.name is required") + if name in self.leases: + raise KubeAlreadyExists(f"leases {name!r} already exists") + stored["metadata"].pop("resourceVersion", None) + self._stamp(stored, kind="Lease", api_version="coordination.k8s.io/v1") + self.leases[name] = stored + return copy.deepcopy(stored) + + return self._call("create_lease", body, {"timeout_s": timeout_s}, operation) + + def replace_lease(self, name: str, body: dict, *, timeout_s: float = 10.0) -> dict: + def operation() -> dict: + current = self.leases.get(name) + if current is None: + raise KubeNotFound(f"leases {name!r} not found") + sent = nested_get(body, "metadata", "resourceVersion") + if sent != current["metadata"]["resourceVersion"]: + raise KubeConflict( + f"leases {name!r}: resourceVersion {sent!r} is stale " + f"(current {current['metadata']['resourceVersion']!r})" + ) + stored = copy.deepcopy(body) + metadata = stored.setdefault("metadata", {}) + for key in ("uid", "creationTimestamp", "namespace", "generation"): + metadata[key] = current["metadata"].get(key) + metadata["name"] = name + stored.setdefault("apiVersion", current.get("apiVersion")) + stored.setdefault("kind", current.get("kind")) + self.bump(stored) + self.leases[name] = stored + return copy.deepcopy(stored) + + return self._call( + "replace_lease", (name, body), {"timeout_s": timeout_s}, operation + ) + + def delete_lease( + self, + name: str, + *, + resource_version: Optional[str] = None, + timeout_s: float = 10.0, + ) -> None: + def operation() -> None: + current = self.leases.get(name) + if current is None: + return + found = current["metadata"]["resourceVersion"] + if resource_version is not None and resource_version != found: + raise KubeConflict( + f"leases {name!r}: precondition resourceVersion " + f"{resource_version!r} is stale (current {found!r})" + ) + del self.leases[name] + + return self._call( + "delete_lease", + name, + {"resource_version": resource_version, "timeout_s": timeout_s}, + operation, + ) + + # ------------------------------------------------- operator simulation + + def set_request_condition( + self, + name: str, + *, + status: str, + reason: str, + message: str = "", + observed_generation: int | None = None, + ) -> None: + """Set the Request's ``Accepted`` condition.""" + + with self._lock: + request = self.requests[name] + generation = request["metadata"].get("generation", 1) + conditions = [ + condition + for condition in request.setdefault("status", {}).get("conditions", []) + if condition.get("type") != "Accepted" + ] + conditions.append( + { + "type": "Accepted", + "status": status, + "reason": reason, + "message": message, + "observedGeneration": ( + generation + if observed_generation is None + else observed_generation + ), + "lastTransitionTime": rfc3339(self.now()), + } + ) + request["status"]["conditions"] = conditions + self.bump(request) + + def accept_request( + self, + name: str, + *, + runtime_image: str = DEFAULT_RUNTIME_IMAGE, + dynamo_version: str = DEFAULT_DYNAMO_VERSION, + mutate_input: Optional[Callable[[dict], None]] = None, + ) -> dict: + """Create the Request's Run (``-``) and mark it Accepted. + + As with the operator, the new Run has no status yet: ``create_job`` + later records its lowered input (kept in ``run_inputs`` until then) + together with ``jobRef``. + """ + + with self._lock: + request = self.requests[name] + generation = int(request["metadata"].get("generation", 1)) + run_name = f"{name}-{generation}" + effective = lower_request_input( + request, + run_name=run_name, + runtime_image=runtime_image, + dynamo_version=dynamo_version, + ) + if mutate_input is not None: + mutate_input(effective) + run: dict[str, Any] = { + "metadata": { + "name": run_name, + "annotations": {REQUEST_GENERATION_ANNOTATION: str(generation)}, + "ownerReferences": [ + { + "apiVersion": "nvidia.com/v1beta2", + "kind": "DynamoGraphDeploymentRequest", + "name": name, + "uid": request["metadata"]["uid"], + "controller": True, + "blockOwnerDeletion": True, + } + ], + }, + "spec": copy.deepcopy( + {k: v for k, v in request.get("spec", {}).items() if k != "v1beta1"} + ), + "status": {}, + } + self._stamp( + run, kind="DynamoGraphDeploymentRun", api_version="nvidia.com/v1beta2" + ) + self.runs[run_name] = run + self.run_inputs[run_name] = effective + status = request.setdefault("status", {}) + status["observedGeneration"] = generation + status["activeRunRef"] = {"name": run_name} + self.set_request_condition( + name, + status="True", + reason="RunCreated", + message=f"Run {run_name!r} exists for generation {generation}", + ) + return copy.deepcopy(run) + + def create_job(self, run_name: str, *, record_input: bool = True) -> None: + """Create the Run's Job, as the operator's ``reconcileRunStatus`` does. + + One status write records ``status.aisimulate.input`` and ``jobRef`` + with ``Completed=False`` (``JobPending``). ``record_input=False`` + leaves the input out (an operator that does not record it). + """ + + with self._lock: + run = self.runs[run_name] + status = run.setdefault("status", {}) + if record_input: + status["aisimulate"] = { + "input": copy.deepcopy(self.run_inputs[run_name]) + } + status["jobRef"] = {"name": f"{run_name}-job", "uid": str(uuid.uuid4())} + status["phase"] = "Pending" + self._set_completed( + run, "False", "JobPending", "AISimulate Job is waiting to start" + ) + + def start_run(self, run_name: str, *, record_input: bool = True) -> None: + """Start the Run's Job (creating it first, see ``create_job``).""" + + with self._lock: + run = self.runs[run_name] + if not run["status"].get("jobRef"): + self.create_job(run_name, record_input=record_input) + status = run["status"] + status["startTime"] = rfc3339(self.now()) + status["phase"] = "Running" + self._set_completed(run, "False", "JobRunning", "AISimulate Job is running") + + def set_run_completed( + self, run_name: str, *, status: str, reason: str, message: str = "" + ) -> None: + """Set the Run's ``Completed`` condition, e.g. ``False`` with + ``ContainerWaiting`` while a Job's container cannot start.""" + + with self._lock: + run = self.runs[run_name] + run.setdefault("status", {}) + self._set_completed(run, status, reason, message) + + def add_candidate( + self, + run_name: str, + parameters: Mapping[str, Any], + metrics: Mapping[str, Any], + ) -> dict: + with self._lock: + run = self.runs[run_name] + canonical = json.dumps(parameters, sort_keys=True, separators=(",", ":")) + candidate_id = ( + "evaluated-point-" + uuid.uuid5(uuid.NAMESPACE_OID, canonical).hex[:12] + ) + candidate = { + "metadata": { + "name": f"{run_name}-{candidate_id}", + "labels": { + RUN_UID_LABEL: run["metadata"]["uid"], + CANDIDATE_ID_LABEL: candidate_id, + }, + "ownerReferences": [ + { + "apiVersion": "nvidia.com/v1beta2", + "kind": "DynamoGraphDeploymentRun", + "name": run_name, + "uid": run["metadata"]["uid"], + "controller": True, + } + ], + }, + "spec": {"parameters": json.loads(canonical)}, + "status": {"metrics": json.loads(json.dumps(dict(metrics)))}, + } + self._stamp( + candidate, + kind="DynamoGraphDeploymentCandidate", + api_version="nvidia.com/v1beta2", + ) + self.candidates[candidate["metadata"]["name"]] = candidate + return copy.deepcopy(candidate) + + def complete_run( + self, + run_name: str, + *, + reason: str = "SearchSucceeded", + message: str = "", + candidates: Iterable[tuple[Mapping[str, Any], Mapping[str, Any]]] = (), + ) -> None: + """Publish ``candidates`` (parameters, metrics) and complete the Run.""" + + with self._lock: + run = self.runs[run_name] + refs = [ + { + "name": self.add_candidate(run_name, params, metrics)["metadata"][ + "name" + ] + } + for params, metrics in candidates + ] + status = run["status"] + status["candidateRefs"] = refs + status["phase"] = "Succeeded" if reason == "SearchSucceeded" else "Failed" + status["completionTime"] = rfc3339(self.now()) + self._set_completed(run, "True", reason, message) + owner = (run["metadata"].get("ownerReferences") or [{}])[0] + request = self.requests.get(owner.get("name", "")) + if request is not None: + request_status = request.setdefault("status", {}) + active = request_status.get("activeRunRef") or {} + if active.get("name") == run_name: + request_status.pop("activeRunRef", None) + recent = [ + ref + for ref in request_status.get("recentRunRefs", []) + if ref.get("name") != run_name + ] + request_status["recentRunRefs"] = [{"name": run_name}, *recent] + self.bump(request) + + def _set_completed(self, run: dict, status: str, reason: str, message: str) -> None: + conditions = [ + condition + for condition in run["status"].get("conditions", []) + if condition.get("type") != "Completed" + ] + conditions.append( + { + "type": "Completed", + "status": status, + "reason": reason, + "message": message, + "lastTransitionTime": rfc3339(self.now()), + } + ) + run["status"]["conditions"] = conditions + self.bump(run) + + def run_of(self, request_name: str) -> Optional[dict]: + """The stored Run a Request owns, if any (the newest).""" + + with self._lock: + request = self.requests.get(request_name) + if request is None: + return None + owned = [ + run + for run in self.runs.values() + if any( + owner.get("uid") == request["metadata"]["uid"] + for owner in run["metadata"].get("ownerReferences") or [] + ) + ] + return copy.deepcopy(owned[-1]) if owned else None + + +Evaluator = Callable[[dict, Mapping[str, float]], Optional[Mapping[str, Any]]] + + +class FakeDgdrOperator: + """Drives every Request of a ``FakeRetuneKube`` one phase per ``step()``. + + Phases: accept (a Run with no status, as before its Job exists), start + (the lowered input, jobRef and startTime; the operator records the + input together with jobRef), complete (one Candidate per router point + the input spans, metrics from ``evaluate(request, point_public)``; + ``None`` leaves the point unpublished, as AISimulate drops an + infeasible point). ``hold`` keeps + named Requests in their current phase; ``fail_runs`` maps a Request + name to the ``Completed`` reason its Run ends with. ``mutate_input`` + edits a Run's lowered input and ``mutate_parameters`` a Candidate's + parameters (both get the Request name), to simulate version skew. + ``max_phase`` stops every Request at ``accepted`` or ``started``. + """ + + def __init__( + self, + kube: FakeRetuneKube, + *, + evaluate: Optional[Evaluator] = None, + runtime_image: str = DEFAULT_RUNTIME_IMAGE, + dynamo_version: str = DEFAULT_DYNAMO_VERSION, + mutate_input: Optional[Callable[[str, dict], None]] = None, + mutate_parameters: Optional[Callable[[str, dict], None]] = None, + ) -> None: + self.kube = kube + self.evaluate: Evaluator = evaluate or ( + lambda _request, point: default_metrics(point) + ) + self.runtime_image = runtime_image + self.dynamo_version = dynamo_version + self.mutate_input = mutate_input + self.mutate_parameters = mutate_parameters + self.max_phase: Optional[Literal["accepted", "started"]] = None + self.hold: set[str] = set() + self.fail_runs: dict[str, str] = {} + self.steps = 0 + + def step(self) -> int: + """Advance every Request one phase; returns how many moved.""" + + moved = 0 + with self.kube._lock: + self.steps += 1 + for name in list(self.kube.requests): + if name in self.hold: + continue + run = self.kube.run_of(name) + if run is None: + hook = self.mutate_input + self.kube.accept_request( + name, + runtime_image=self.runtime_image, + dynamo_version=self.dynamo_version, + mutate_input=( + None if hook is None else functools.partial(hook, name) + ), + ) + moved += 1 + continue + run_name = run["metadata"]["name"] + completed = next( + ( + condition + for condition in run["status"].get("conditions", []) + if condition.get("type") == "Completed" + ), + None, + ) + if completed is not None and completed.get("status") == "True": + continue + if self.max_phase == "accepted": + continue + if not run["status"].get("startTime"): + self.kube.start_run(run_name) + moved += 1 + continue + if self.max_phase == "started": + continue + reason = self.fail_runs.get(name) + if reason is not None: + self.kube.complete_run(run_name, reason=reason, message="injected") + moved += 1 + continue + effective = run["status"]["aisimulate"]["input"] + request = self.kube.requests[name] + published = [] + for point in router_points(effective): + metrics = self.evaluate(copy.deepcopy(request), point) + if metrics is not None: + parameters = candidate_parameters(effective, point) + if self.mutate_parameters is not None: + self.mutate_parameters(name, parameters) + published.append((parameters, metrics)) + self.kube.complete_run(run_name, candidates=published) + moved += 1 + return moved diff --git a/components/src/dynamo/planner/tests/retune_stubs.py b/components/src/dynamo/planner/tests/retune_stubs.py index aedfa4518e78..d9d137c369d8 100644 --- a/components/src/dynamo/planner/tests/retune_stubs.py +++ b/components/src/dynamo/planner/tests/retune_stubs.py @@ -50,21 +50,43 @@ class StubController: """Records the host's calls; returns one queued tick result per tick.""" instances: list[StubController] = [] - - def __init__(self, config, inputs, capabilities, clock, scale_interval_s): + # Raised by the next construction, then cleared (live rebuild tests). + construct_error: Exception | None = None + + def __init__(self, config, inputs, capabilities, clock, scale_interval_s, **hooks): + error = StubController.construct_error + if error is not None: + StubController.construct_error = None + raise error self.args = (config, inputs, capabilities, clock, scale_interval_s) + # Live hooks (sweep_fn, launch_gate, harvest_allowed, generation, + # executor); empty when the host built the controller the replay way. + self.hooks = hooks self.config = config self.started: list[float] = [] + # phase_s of each start(); None when the host passed none. + self.phases: list[float | None] = [] self.ticks: list[tuple[float, Any]] = [] # Searched planner values as the controller saw them on each tick. self.seen_planner_values: list[dict[str, Any]] = [] self.results: list[StubTickResult] = [] self.closed = 0 self.close_error: Exception | None = None + # Ordered lifecycle calls: "start", "drop_pending", "close". + self.calls: list[str] = [] StubController.instances.append(self) - def start(self, start_s: float) -> None: + @property + def generation(self) -> int: + return self.hooks.get("generation", 0) + + def start(self, start_s: float, phase_s: float | None = None) -> None: self.started.append(start_s) + self.phases.append(phase_s) + self.calls.append("start") + + def drop_pending(self) -> None: + self.calls.append("drop_pending") async def on_tick(self, now_s, worker_counts): self.ticks.append((now_s, worker_counts)) @@ -78,6 +100,7 @@ async def on_tick(self, now_s, worker_counts): def close(self) -> None: self.closed += 1 + self.calls.append("close") if self.close_error is not None: raise self.close_error @@ -89,6 +112,7 @@ def install_stub_controller(monkeypatch) -> type[StubController]: module.RetuneController = StubController # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, CONTROLLER_MODULE, module) StubController.instances = [] + StubController.construct_error = None return StubController diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_actuator.py b/components/src/dynamo/planner/tests/unit/test_live_retune_actuator.py new file mode 100644 index 000000000000..4a1237b96b3b --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_actuator.py @@ -0,0 +1,837 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the live retune router actuator and its HTTP client. + +``FakePolicyClient`` only records the calls and answers them: a successful +write sets the router to what the actuator asked for (the PATCHed keys, or +its initial after a reset) and bumps the revision. Error statuses are +injected by the tests; the server's own rules (``router_policy.rs``) are not +modelled here. +""" + +from __future__ import annotations + +import json +import re +import socket +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Optional + +import pytest +from aiohttp import test_utils, web + +from dynamo.planner.retune.live.actuator import ( + DEFAULT_CLIENT_TIMEOUT_S, + ROUTER_CONTROL_TIMEOUT_S, + AiohttpRouterPolicyClient, + RouterActuator, +) +from dynamo.planner.retune.live.contracts import ( + POD_LABEL_COMPONENT_TYPE, + POD_LABEL_DGD_NAME, + ActuationTarget, + DgdRef, + KubeError, + PodInfo, + PolicyCallResult, + RouterInfo, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +DGD = DgdRef(name="dgd", namespace="ns", uid="dgd-uid", dynamo_namespace="dyn-ns") + +# KvRouterConfig defaults resolved by policy_parameters("aggregated"). +STARTUP = { + "decode_active_request_weight": 0.0, + "disk_cache_hit_weight": 0.25, + "host_cache_hit_weight": 0.75, + "overlap_score_credit": 1.0, + "overlap_score_credit_decay": 0.0, + "prefill_load_scale": 1.0, + "router_temperature": 0.0, + "shared_cache_multiplier": 0.0, +} +CHOSEN = { + **STARTUP, + "overlap_score_credit": 0.5, + "prefill_load_scale": 2.0, + "router_temperature": 0.2, +} +# policy_parameters("decode") of the defaults: plain decode has no cache input. +DECODE_STARTUP = { + **STARTUP, + "overlap_score_credit": 0.0, + "disk_cache_hit_weight": 0.0, + "host_cache_hit_weight": 0.0, +} + + +# ------------------------------------------------------------------ fakes + + +@dataclass +class FakeRouter: + router_id: str + initial: dict[str, float] + worker_type: str = "aggregated" + namespace: str = DGD.dynamo_namespace + revision: int = 0 + current: Optional[dict[str, float]] = None + runtime_config: bool = True # False: a custom policy, configuration null + + def __post_init__(self) -> None: + if self.current is None: + self.current = dict(self.initial) + + def info(self) -> RouterInfo: + return RouterInfo( + router_id=self.router_id, + namespace=self.namespace, + worker_type=self.worker_type, + revision=self.revision if self.runtime_config else 0, + configuration=dict(self.current) if self.runtime_config else None, + model="model", + ) + + +@dataclass +class FakeFrontend: + routers: dict[str, FakeRouter] = field(default_factory=dict) + # Statuses returned by the next GETs before the real answer. + get_failures: list[int] = field(default_factory=list) + # Routers GET leaves out (their control call timed out). + timed_out: set[str] = field(default_factory=set) + + def add(self, router: FakeRouter) -> FakeRouter: + self.routers[router.router_id] = router + return router + + +class FakePolicyClient: + """In-memory frontends keyed by pod name; records every call.""" + + def __init__(self) -> None: + self.frontends: dict[str, FakeFrontend] = {} + self.calls: list[tuple[Any, ...]] = [] + # Statuses returned by the next writes (any router) instead of applying. + self.write_failures: list[int] = [] + # Called before a write is answered: (op, pod name, router_id). + self.before_write: Optional[Callable[[str, str, str], None]] = None + + def frontend(self, pod: str) -> FakeFrontend: + return self.frontends.setdefault(pod, FakeFrontend()) + + def writes(self) -> list[tuple[Any, ...]]: + return [call for call in self.calls if call[0] in ("patch", "reset")] + + async def get(self, pod: PodInfo) -> PolicyCallResult: + self.calls.append(("get", pod.name)) + frontend = self.frontends.get(pod.name) + if frontend is None: + return PolicyCallResult(status=0, error="connection refused") + if frontend.get_failures: + return PolicyCallResult(status=frontend.get_failures.pop(0)) + return PolicyCallResult( + status=200, + routers=tuple( + router.info() + for router in frontend.routers.values() + if router.router_id not in frontend.timed_out + ), + ) + + async def patch( + self, + pod: PodInfo, + router_id: str, + expected_revision: int, + parameters: Mapping[str, float], + ) -> PolicyCallResult: + self.calls.append( + ("patch", pod.name, router_id, expected_revision, dict(parameters)) + ) + return self._answer("patch", pod, router_id, parameters) + + async def reset( + self, pod: PodInfo, router_id: str, expected_revision: int + ) -> PolicyCallResult: + self.calls.append(("reset", pod.name, router_id, expected_revision)) + return self._answer("reset", pod, router_id, None) + + def _answer( + self, + op: str, + pod: PodInfo, + router_id: str, + parameters: Optional[Mapping[str, float]], + ) -> PolicyCallResult: + if self.before_write is not None: + self.before_write(op, pod.name, router_id) + if self.write_failures: + return PolicyCallResult(status=self.write_failures.pop(0), error="injected") + router = self.frontend(pod.name).routers.get(router_id) + if router is None: + return PolicyCallResult(status=404, error="router instance not found") + if parameters is None: + router.current = dict(router.initial) + else: + router.current = {**(router.current or {}), **parameters} + router.revision += 1 + return PolicyCallResult(status=200, router=router.info()) + + +class FakeKube: + namespace = DGD.namespace + + def __init__(self, pods: Optional[list[PodInfo]] = None) -> None: + self.pods: list[PodInfo] = list(pods or []) + self.selectors: list[str] = [] + self.fail: Optional[Exception] = None + + def list_pods( + self, label_selector: str, *, timeout_s: float = 10.0 + ) -> list[PodInfo]: + self.selectors.append(label_selector) + if self.fail is not None: + raise self.fail + return list(self.pods) + + +class RecordingMetrics: + def __init__(self) -> None: + self.converged: list[bool] = [] + self.patch_errors: list[str] = [] + + def router_converged(self, converged: bool) -> None: + self.converged.append(converged) + + def patch_error(self, code: str) -> None: + self.patch_errors.append(code) + + def __getattr__(self, name: str) -> Callable[..., None]: + return lambda *args, **kwargs: None + + +def pod( + name: str, + *, + ready: bool = True, + terminating: bool = False, + ip: Optional[str] = "10.0.0.1", + phase: Optional[str] = None, + restarts: int = 0, +) -> PodInfo: + return PodInfo( + name=name, + uid=f"{name}-uid", + ip=ip, + ready=ready, + terminating=terminating, + labels={POD_LABEL_DGD_NAME: DGD.name, POD_LABEL_COMPONENT_TYPE: "frontend"}, + http_port=8000, + phase=phase, + restarts=restarts, + ) + + +def make_actuator( + pods: list[PodInfo], + *, + startup_maps: Optional[dict[str, dict[str, float]]] = None, +) -> tuple[RouterActuator, FakeKube, FakePolicyClient, RecordingMetrics]: + kube = FakeKube(pods) + client = FakePolicyClient() + metrics = RecordingMetrics() + actuator = RouterActuator( + kube, client, DGD, startup_maps or {"aggregated": dict(STARTUP)}, metrics + ) + return actuator, kube, client, metrics + + +def apply(desired: Mapping[str, float], wt: str = "aggregated") -> ActuationTarget: + return ActuationTarget(mode="apply", desired_by_wt={wt: dict(desired)}) + + +OBSERVE = ActuationTarget(mode="observe") +ROLLBACK = ActuationTarget(mode="rollback", include_not_ready=True) + + +# ---------------------------------------------------------------- latching + + +async def test_observe_latches_initials_and_never_writes() -> None: + actuator, kube, client, metrics = make_actuator([pod("fe-0")]) + client.frontend("fe-0").add(FakeRouter("r1", STARTUP)) + + report = await actuator.reconcile(OBSERVE, initials={}, expected_frontends=1) + + assert kube.selectors == [ + f"{POD_LABEL_DGD_NAME}=dgd,{POD_LABEL_COMPONENT_TYPE}=frontend" + ] + assert client.writes() == [] + assert report.converged + assert report.new_initials == {"r1": STARTUP} + assert report.seen_router_ids == frozenset({"r1"}) + assert report.parity_ok is True + assert report.parity_mismatches == () + assert metrics.converged == [True] + + # Persisted initials are no longer reported as new. + report = await actuator.reconcile( + OBSERVE, initials={"r1": STARTUP}, expected_frontends=1 + ) + assert report.new_initials == {} + assert report.converged + + +# ------------------------------------------------------------------- apply + + +async def test_apply_holds_a_router_until_its_initial_is_persisted() -> None: + # A process stopped between a PATCH and the state write of the router's + # initial would leave a router no rollback can restore: an advisory or + # clean-up planner never resets a router whose initial it does not know. + actuator, _, client, _ = make_actuator([pod("fe-0")]) + router = client.frontend("fe-0").add(FakeRouter("r1", STARTUP)) + + for _ in range(2): # the state write has not landed yet + report = await actuator.reconcile( + apply(CHOSEN), initials={}, expected_frontends=1 + ) + assert client.writes() == [] + assert report.new_initials == {"r1": STARTUP} + assert not report.converged + + report = await actuator.reconcile( + apply(CHOSEN), initials={"r1": STARTUP}, expected_frontends=1 + ) + assert [call[0] for call in client.writes()] == ["patch"] + assert router.current == CHOSEN + assert report.converged and report.new_initials == {} + + +async def test_apply_patches_only_differing_keys_with_explicit_values() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + router = client.frontend("fe-0").add(FakeRouter("r1", STARTUP)) + initials = {"r1": dict(STARTUP)} + + report = await actuator.reconcile( + apply(CHOSEN), initials=initials, expected_frontends=1 + ) + + assert client.writes() == [ + ( + "patch", + "fe-0", + "r1", + 0, + { + "overlap_score_credit": 0.5, + "prefill_load_scale": 2.0, + "router_temperature": 0.2, + }, + ) + ] + assert router.current == CHOSEN + assert report.converged + assert report.parity_ok is True + + # Converged: the next rounds write nothing. + for _ in range(2): + report = await actuator.reconcile( + apply(CHOSEN), initials=initials, expected_frontends=1 + ) + assert report.converged + assert len(client.writes()) == 1 + + +async def test_apply_learns_an_unknown_initial_with_a_reset_first() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + tampered = {**STARTUP, "router_temperature": 0.9} + router = client.frontend("fe-0").add( + FakeRouter("r1", STARTUP, revision=2, current=tampered) + ) + + # A reset restores the router's initial and reports it; the PATCH waits + # for the state to record it. + report = await actuator.reconcile(apply(CHOSEN), initials={}, expected_frontends=1) + assert client.writes() == [("reset", "fe-0", "r1", 2)] + assert router.current == STARTUP + assert report.new_initials == {"r1": STARTUP} + assert report.parity_ok is True + assert not report.converged + + report = await actuator.reconcile( + apply(CHOSEN), initials=report.new_initials, expected_frontends=1 + ) + assert [call[0] for call in client.writes()] == ["reset", "patch"] + assert router.current == CHOSEN + assert report.converged + + +async def test_parity_mismatch_blocks_apply_and_never_loops() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0"), pod("fe-1", ip="10.0.0.2")]) + odd = {**STARTUP, "overlap_score_credit_decay": 0.5} + client.frontend("fe-0").add(FakeRouter("good", STARTUP)) + client.frontend("fe-1").add(FakeRouter("odd", odd)) + initials = {"good": dict(STARTUP), "odd": dict(odd)} + + for _ in range(3): + report = await actuator.reconcile( + apply(CHOSEN), initials=initials, expected_frontends=2 + ) + assert report.parity_ok is False + assert len(report.parity_mismatches) == 1 + assert report.parity_mismatches[0].startswith("fe-1/odd/aggregated:") + assert "overlap_score_credit_decay initial=0.5 startup=0.0" in ( + report.parity_mismatches[0] + ) + assert not report.converged + # Only the router whose initial matches startup was written, once. + assert [call[2] for call in client.writes()] == ["good"] + assert client.frontend("fe-1").routers["odd"].current == odd + + +async def test_disagg_needs_both_worker_types() -> None: + maps = {"prefill": dict(STARTUP), "decode": dict(DECODE_STARTUP)} + actuator, _, client, _ = make_actuator([pod("fe-0")], startup_maps=maps) + frontend = client.frontend("fe-0") + prefill = frontend.add(FakeRouter("p", STARTUP, worker_type="prefill")) + target = ActuationTarget( + mode="apply", + desired_by_wt={ + "prefill": dict(CHOSEN), + "decode": {**DECODE_STARTUP, "router_temperature": 0.2}, + }, + ) + initials = {"p": dict(STARTUP), "d": dict(DECODE_STARTUP)} + + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert not report.converged # no decode router yet + assert prefill.current == CHOSEN + + decode = frontend.add(FakeRouter("d", DECODE_STARTUP, worker_type="decode")) + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert decode.current == {**DECODE_STARTUP, "router_temperature": 0.2} + assert report.converged + assert report.parity_ok is True + + +async def test_only_the_dgd_namespace_and_planner_worker_types_are_managed() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + frontend = client.frontend("fe-0") + ours = frontend.add(FakeRouter("ours", STARTUP)) + foreign = frontend.add(FakeRouter("foreign", STARTUP, namespace="other-ns")) + encode = frontend.add(FakeRouter("encode", STARTUP, worker_type="encode")) + + report = await actuator.reconcile( + apply(CHOSEN), initials={"ours": dict(STARTUP)}, expected_frontends=1 + ) + + assert {call[2] for call in client.writes()} == {"ours"} + assert ours.current == CHOSEN + assert foreign.current == STARTUP and encode.current == STARTUP + assert report.seen_router_ids == frozenset({"ours"}) + assert report.converged + + +# ---------------------------------------------------------------- rollback + + +async def test_rollback_resets_only_deviating_routers_and_never_by_revision() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + frontend = client.frontend("fe-0") + # At its initial although its revision is not 0 (patched, then reset). + at_initial = frontend.add(FakeRouter("a", STARTUP, revision=3)) + deviating = frontend.add(FakeRouter("b", STARTUP, revision=1, current=dict(CHOSEN))) + initials = {"a": dict(STARTUP), "b": dict(STARTUP)} + + report = await actuator.reconcile(ROLLBACK, initials=initials, expected_frontends=1) + + assert client.writes() == [("reset", "fe-0", "b", 1)] + assert deviating.current == STARTUP + assert at_initial.revision == 3 + assert report.converged + + # The reset bumped b's revision; comparing configurations, not + # revisions, means no further resets. + for _ in range(3): + report = await actuator.reconcile( + ROLLBACK, initials=initials, expected_frontends=1 + ) + assert report.converged + assert len(client.writes()) == 1 + + +async def test_rollback_without_learning_leaves_unknown_routers_alone() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + frontend = client.frontend("fe-0") + # Changed by someone else before any planner latched it. + foreign = frontend.add( + FakeRouter("foreign", STARTUP, revision=4, current=dict(CHOSEN)) + ) + ours = frontend.add(FakeRouter("ours", STARTUP, revision=1, current=dict(CHOSEN))) + target = ActuationTarget( + mode="rollback", include_not_ready=True, learn_unknown=False + ) + + for _ in range(2): + report = await actuator.reconcile( + target, initials={"ours": STARTUP}, expected_frontends=1 + ) + assert report.converged + assert client.writes() == [("reset", "fe-0", "ours", 1)] + assert ours.current == STARTUP + assert foreign.current == CHOSEN and foreign.revision == 4 + assert report.new_initials == {} + assert report.parity_ok is None # the foreign router's initial is unknown + + # The default (active) target learns it with one reset. + report = await actuator.reconcile( + ROLLBACK, initials={"ours": STARTUP}, expected_frontends=1 + ) + assert client.writes()[-1] == ("reset", "fe-0", "foreign", 4) + assert foreign.current == STARTUP + assert report.converged + + +async def test_rollback_targets_not_ready_pods_and_waits_only_for_ready_ones() -> None: + pods = [ + pod("ready", ip="10.0.0.1"), + pod("not-ready", ready=False, ip="10.0.0.2"), + pod("crashing", ready=False, ip="10.0.0.3"), + pod("pending", ready=False, ip=None), + pod("leaving", terminating=True, ip="10.0.0.4"), + ] + actuator, kube, client, _ = make_actuator(pods) + initials = {} + for name in ("ready", "not-ready", "leaving"): + client.frontend(name).add( + FakeRouter(f"r-{name}", STARTUP, revision=1, current=dict(CHOSEN)) + ) + initials[f"r-{name}"] = dict(STARTUP) + # CrashLoopBackOff: an unreachable NotReady pod never blocks convergence. + client.frontend("crashing").get_failures = [0] + + report = await actuator.reconcile(ROLLBACK, initials=initials, expected_frontends=1) + assert sorted(call[1] for call in client.writes()) == ["not-ready", "ready"] + assert ("get", "leaving") not in client.calls + assert ("get", "pending") not in client.calls + assert report.converged and report.errors == {"get:network": 1} + + # An unreachable Ready pod does block it; apply never targets NotReady pods. + kube.pods = [pod("ready"), pod("crashing", ip="10.0.0.3")] + client.frontend("crashing").get_failures = [0] + client.calls.clear() + report = await actuator.reconcile( + apply(CHOSEN), initials=initials, expected_frontends=1 + ) + assert not report.converged + kube.pods = pods + client.calls.clear() + await actuator.reconcile(apply(CHOSEN), initials=initials, expected_frontends=1) + assert {call[1] for call in client.calls} == {"ready"} + + +# ----------------------------------------------------------- status codes + + +async def test_conflict_regets_and_retries_once() -> None: + actuator, _, client, metrics = make_actuator([pod("fe-0")]) + router = client.frontend("fe-0").add(FakeRouter("r1", STARTUP)) + client.write_failures = [409] + + def someone_else_writes(op: str, pod_name: str, router_id: str) -> None: + if client.write_failures: # the write that answers 409 + router.current = {**STARTUP, "router_temperature": 0.7} + router.revision += 1 + + client.before_write = someone_else_writes + + report = await actuator.reconcile( + apply(CHOSEN), initials={"r1": STARTUP}, expected_frontends=1 + ) + + patches = [call for call in client.writes() if call[0] == "patch"] + assert [call[3] for call in patches] == [0, 1] + # The retry diff comes from the fresh GET (temperature 0.7 -> 0.2). + assert patches[1][4]["router_temperature"] == 0.2 + assert router.current == CHOSEN + assert report.errors == {"patch:409": 1} + assert metrics.patch_errors == ["409"] + assert report.converged + + +# ------------------------------------------------------------- convergence + + +async def test_convergence_and_presence_need_the_expected_frontends() -> None: + actuator, _, client, _ = make_actuator([pod("fe-0")]) + frontend = client.frontend("fe-0") + frontend.add(FakeRouter("custom", STARTUP, runtime_config=False)) + + # No configured router: neither converged nor present (P12). + report = await actuator.reconcile(apply(CHOSEN), initials={}, expected_frontends=1) + assert not report.converged and report.routers_present is False + assert report.parity_ok is None and client.writes() == [] + + # Present even though the router does not match the apply target. + frontend.add(FakeRouter("r1", {**STARTUP, "router_temperature": 0.9})) + report = await actuator.reconcile(OBSERVE, initials={}, expected_frontends=1) + assert report.converged and report.routers_present is True + report = await actuator.reconcile(OBSERVE, initials={}, expected_frontends=2) + assert not report.converged and report.routers_present is False + + +# ------------------------------------------------------------ gone routers + + +async def test_seen_router_ids_drop_only_routers_that_are_provably_gone() -> None: + pods = [pod("fe-0"), pod("fe-1", ip="10.0.0.2")] + actuator, kube, client, _ = make_actuator(pods) + frontend0 = client.frontend("fe-0") + frontend0.add(FakeRouter("a", STARTUP)) + client.frontend("fe-1").add(FakeRouter("b", STARTUP)) + report = await actuator.reconcile(OBSERVE, initials={}, expected_frontends=2) + initials = dict(report.new_initials) + assert report.seen_router_ids == frozenset({"a", "b"}) + + # GET leaves out a router whose control call timed out, and a pod that + # does not answer shows none: both are kept, as on a list failure. + frontend0.timed_out.add("a") + client.frontend("fe-1").get_failures = [0] + report = await actuator.reconcile(OBSERVE, initials=initials, expected_frontends=2) + assert report.seen_router_ids == frozenset({"a", "b"}) + kube.fail = KubeError("apiserver down") + report = await actuator.reconcile(OBSERVE, initials=initials, expected_frontends=2) + assert report.errors == {"list_pods": 1} and actuator.last_pods is None + assert report.seen_router_ids == frozenset({"a", "b"}) + kube.fail = None + + # fe-0's frontend restarted with a new router; fe-1 was evicted (the pod + # object stays, with its IP, in a finished phase and is never asked). + del frontend0.routers["a"] + frontend0.timed_out.clear() + frontend0.add(FakeRouter("a2", STARTUP)) + kube.pods = [ + pod("fe-0", restarts=1), + pod("fe-1", ready=False, ip="10.0.0.2", phase="Failed"), + ] + client.calls.clear() + report = await actuator.reconcile(OBSERVE, initials=initials, expected_frontends=1) + assert ("get", "fe-1") not in client.calls + assert report.seen_router_ids == frozenset({"a2"}) + + +async def test_a_router_missing_from_get_is_still_rolled_back() -> None: + # A clean-up (advisory, learn_unknown off) planner restarts with the + # initial of a retuned router persisted. Its first GET leaves the router + # out (control call timed out): the initial must survive, or the router + # comes back at revision > 0, is left alone, and stays retuned. + target = ActuationTarget( + mode="rollback", include_not_ready=True, learn_unknown=False + ) + actuator, kube, client, _ = make_actuator([pod("fe-0")]) + frontend = client.frontend("fe-0") + router = frontend.add(FakeRouter("r1", STARTUP, revision=1, current=dict(CHOSEN))) + frontend.timed_out.add("r1") + initials = {"r1": dict(STARTUP)} + + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert report.seen_router_ids == frozenset({"r1"}) + assert not report.converged and client.writes() == [] + + frontend.timed_out.clear() + initials = {k: v for k, v in initials.items() if k in report.seen_router_ids} + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert client.writes() == [("reset", "fe-0", "r1", 1)] + assert router.current == STARTUP and report.converged + + # Once every frontend listed at the first miss is gone, so is the router. + actuator, kube, client, _ = make_actuator([pod("fe-0")]) + client.frontend("fe-0").add(FakeRouter("r2", STARTUP)) + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert report.seen_router_ids == frozenset({"r1", "r2"}) + kube.pods = [pod("fe-new", ip="10.0.0.3")] + report = await actuator.reconcile(target, initials=initials, expected_frontends=1) + assert report.seen_router_ids == frozenset() + + +# ------------------------------------------------------------- HTTP client + + +class _FrontendApp: + """A tiny aiohttp server speaking the /router_policy wire format.""" + + def __init__(self) -> None: + self.requests: list[tuple[str, Any]] = [] + self.patch_status = 200 + self.list_body: Any = { + "routers": [ + { + "router_id": "6f1c", + "model": "Qwen/Qwen3-0.6B", + "namespace": "dyn-ns", + "component": "backend", + "endpoint": "generate", + "worker_type": "aggregated", + "configuration": {"revision": 3, "parameters": dict(STARTUP)}, + }, + { + "router_id": "custom", + "model": "Qwen/Qwen3-0.6B", + "namespace": "dyn-ns", + "component": "backend", + "endpoint": "generate", + "worker_type": "decode", + "configuration": None, + }, + ] + } + + def app(self) -> web.Application: + app = web.Application() + app.router.add_get("/router_policy", self.list) + app.router.add_patch("/router_policy", self.update) + return app + + async def list(self, request: web.Request) -> web.Response: + self.requests.append(("GET", dict(request.query))) + if isinstance(self.list_body, str): + return web.Response(text=self.list_body, content_type="application/json") + return web.json_response(self.list_body) + + async def update(self, request: web.Request) -> web.Response: + body = await request.json() + self.requests.append(("PATCH", body)) + if self.patch_status != 200: + return web.json_response( + {"error": "policy revision conflict: expected 1, current 2"}, + status=self.patch_status, + ) + parameters = dict(STARTUP) + parameters.update(body.get("parameters") or {}) + return web.json_response( + { + "router_id": body["router_id"], + "model": "Qwen/Qwen3-0.6B", + "namespace": "dyn-ns", + "component": "backend", + "endpoint": "generate", + "worker_type": "aggregated", + "configuration": { + "revision": body["expected_revision"] + 1, + "parameters": parameters, + }, + } + ) + + +@pytest.fixture +async def frontend_server(): + frontend = _FrontendApp() + server = test_utils.TestServer(frontend.app(), host="127.0.0.1") + await server.start_server() + client = AiohttpRouterPolicyClient(timeout_s=2.0) + target = PodInfo( + name="fe-0", + uid="u", + ip="127.0.0.1", + ready=True, + terminating=False, + http_port=server.port, + ) + try: + yield frontend, client, target + finally: + await client.close() + await server.close() + + +async def test_client_speaks_the_router_policy_wire_format(frontend_server) -> None: + frontend, client, target = frontend_server + + result = await client.get(target) + assert result.status == 200 and result.error is None + first, custom = result.routers + assert first == RouterInfo( + router_id="6f1c", + namespace="dyn-ns", + worker_type="aggregated", + revision=3, + configuration=STARTUP, + model="Qwen/Qwen3-0.6B", + ) + assert custom.configuration is None and custom.revision == 0 + + patched = await client.patch(target, "6f1c", 3, {"overlap_score_credit": 0}) + reset = await client.reset(target, "6f1c", 7) + assert frontend.requests == [ + ("GET", {}), # no ?model= + ( + "PATCH", + { + "router_id": "6f1c", + "expected_revision": 3, + "parameters": {"overlap_score_credit": 0.0}, + }, + ), + ("PATCH", {"router_id": "6f1c", "expected_revision": 7, "reset": True}), + ] + assert patched.router is not None and patched.router.revision == 4 + assert patched.router.configuration["overlap_score_credit"] == 0.0 + assert reset.status == 200 and reset.router.revision == 8 + + +async def test_client_errors_are_statuses(frontend_server) -> None: + frontend, client, target = frontend_server + frontend.patch_status = 409 + result = await client.patch(target, "6f1c", 1, {"router_temperature": 0.1}) + assert (result.status, result.router) == (409, None) + assert result.error == "policy revision conflict: expected 1, current 2" + + for body in ("not json", json.dumps({"routers": "x"}), json.dumps({"other": []})): + frontend.list_body = body + result = await client.get(target) + assert result.status == 0 + assert result.error and result.error.startswith("malformed GET response") + + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + closed_port = sock.getsockname()[1] + result = await client.get(replace(target, http_port=closed_port)) + assert result.status == 0 and result.error + + +def test_client_timeout_outlasts_the_server_control_timeout() -> None: + # A stuck router makes the server wait CONTROL_TIMEOUT before answering + # GET 200 without it (or PATCH 503); the client must still be listening. + assert DEFAULT_CLIENT_TIMEOUT_S == 8.0 + assert DEFAULT_CLIENT_TIMEOUT_S > ROUTER_CONTROL_TIMEOUT_S + assert AiohttpRouterPolicyClient()._timeout.total == DEFAULT_CLIENT_TIMEOUT_S + + source = next( + ( + parent / "lib/llm/src/http/service/router_policy.rs" + for parent in Path(__file__).resolve().parents + if (parent / "lib/llm/src/http/service/router_policy.rs").is_file() + ), + None, + ) + if source is None: + pytest.skip("router_policy.rs is not in this checkout") + match = re.search( + r"CONTROL_TIMEOUT: Duration = Duration::from_secs\((\d+)\)", + source.read_text(), + ) + assert match is not None + assert float(match.group(1)) == ROUTER_CONTROL_TIMEOUT_S diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_adapter.py b/components/src/dynamo/planner/tests/unit/test_live_retune_adapter.py new file mode 100644 index 000000000000..34cbd95a6f71 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_adapter.py @@ -0,0 +1,836 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Live retune wiring of the engine adapter, the planner base and ``__main__``. + +The retune controller is the recording stub of ``retune_stubs`` (one test +builds the real one), and the live supervisor is a fake ``LiveRetuneHandle``, +so these tests cover only the host contract: the disabled state, +``replace_retune``, the base's supervisor lifecycle, its runtime +configuration API and the SIGTERM path. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import signal +import threading +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import patch + +import pytest + +from dynamo.planner import __main__ as planner_main +from dynamo.planner.config.planner_config import PlannerConfig, RetuneConfig +from dynamo.planner.core.base import NativePlannerBase +from dynamo.planner.core.types import ( + PlannerEffects, + RetuneDecision, + RetuneInputs, + RetuneParams, + ScheduledTick, + TickInput, + WorkerCounts, +) +from dynamo.planner.environment.state import DeploymentState +from dynamo.planner.plugins.clock import WallClock +from dynamo.planner.plugins.orchestrator.engine_adapter import ( + OrchestratorEngineAdapter, + RetuneHooks, +) +from dynamo.planner.reconfiguration import ReconfigureRequest, candidate_config +from dynamo.planner.retune.live.contracts import GateResult +from dynamo.planner.tests.retune_stubs import ( + ROUTER_MAP, + StubTickResult, + caps, + install_stub_controller, + retune_inputs, + retune_only_config, +) +from dynamo.planner.tests.unit.test_live_retune_state import live_config + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +PHASE = 120.0 + + +def _sweep(request, *, cancel, abort): + raise AssertionError("the stub controller never sweeps") + + +def _hooks(**overrides) -> RetuneHooks: + values = { + "sweep_fn": _sweep, + "launch_gate": lambda _k, _window: GateResult(), + "harvest_allowed": lambda _generation: True, + "phase_s": PHASE, + } + values.update(overrides) + return RetuneHooks(**values) + + +@pytest.fixture +def stub_controller(monkeypatch): + return install_stub_controller(monkeypatch) + + +def _adapter(*, inputs=None, hooks=None, clock=None, **overrides): + return OrchestratorEngineAdapter( + retune_only_config(**overrides), + caps(), + clock=clock or WallClock(), + retune_inputs=inputs, + retune_hooks=hooks if hooks is not None else _hooks(), + ) + + +def _decision() -> RetuneDecision: + point = RetuneParams(router={"router_temperature": 0.5}) + return RetuneDecision( + decision_id="retune-w60000", + outcome="kept", + window_index=0, + window_start_s=0.0, + window_end_s=60.0, + as_of_s=120.0, + launched_at_s=120.0, + decided_at_s=125.0, + num_requests=3, + num_candidates=2, + previous=point, + chosen=point, + baseline_score=1.0, + best_score=1.0, + ) + + +# --------------------------------------------------------------------------- +# Engine adapter: construction and the disabled state +# --------------------------------------------------------------------------- + + +def test_live_controller_gets_the_hooks_and_the_phase(stub_controller): + hooks = _hooks() + engine = _adapter(inputs=retune_inputs(), hooks=hooks) + + (controller,) = stub_controller.instances + assert engine.retune_controller is controller + executor = controller.hooks.pop("executor") + assert isinstance(executor, ThreadPoolExecutor) + assert controller.hooks == { + "sweep_fn": hooks.sweep_fn, + "launch_gate": hooks.launch_gate, + "harvest_allowed": hooks.harvest_allowed, + "generation": 0, + } + engine.initial_tick(1_000.0) + assert (controller.started, controller.phases) == ([1_000.0], [PHASE]) + + +def test_live_without_inputs_or_a_valid_controller_starts_disabled( + stub_controller, caplog +): + with caplog.at_level(logging.WARNING): + engine = _adapter(inputs=None) + assert engine.retune_controller is None and stub_controller.instances == [] + assert "retune stays disabled" in caplog.text + engine.initial_tick(0.0) # no controller to start + + stub_controller.construct_error = ValueError("retune.router requires kv_router") + with caplog.at_level(logging.ERROR): + engine = _adapter(inputs=retune_inputs()) + assert engine.retune_controller is None + assert "requires kv_router" in caplog.text + + +async def test_disabled_retune_only_planner_never_reconciles_replicas( + stub_controller, +): + # retune_only follows the config, not the controller: a disabled live + # retune keeps the deployment's replicas (here below min_endpoint) and + # stays quiet. + engine = _adapter(inputs=None, optimization_target="sla", min_endpoint=2) + counts = WorkerCounts(ready_num_decode=1, expected_num_decode=1) + + scheduled = engine.initial_tick(0.0) + for _ in range(10): + assert not scheduled.run_load_scaling + assert not scheduled.need_worker_fpm + effects = await engine.tick( + scheduled, TickInput(now_s=scheduled.at_s, worker_counts=counts) + ) + assert effects.scale_to is None + assert effects.router_parameters is None + assert effects.retune_decisions == [] + assert effects.diagnostics.audit_events == [] + assert effects.next_tick is not None + scheduled = effects.next_tick + + +# --------------------------------------------------------------------------- +# Engine adapter: replace_retune +# --------------------------------------------------------------------------- + + +async def test_replace_retune_swaps_the_controller(stub_controller): + engine = _adapter(inputs=retune_inputs()) + engine.initial_tick(0.0) + (old,) = stub_controller.instances + inputs = RetuneInputs(engine=retune_inputs().engine, trace_paths=("/trace/b",)) + + engine.replace_retune(inputs, generation=4, now_s=500.0) + + old_, new = stub_controller.instances + assert old_ is old + # Abort the old sweep first, then close the old controller. + assert old.calls == ["start", "drop_pending", "close"] + assert new.args[1] is inputs + assert new.generation == 4 + # Every generation submits to the adapter's one sweep thread. + assert new.hooks["executor"] is old.hooks["executor"] + assert new.hooks["launch_gate"] is old.hooks["launch_gate"] + assert new.calls == ["start"] + assert (new.started, new.phases) == ([500.0], [PHASE]) + assert engine.retune_controller is new + + # The next tick runs the new controller only. + new.results.append( + StubTickResult(decisions=[_decision()], router_parameters=ROUTER_MAP) + ) + scheduled = ScheduledTick(at_s=505.0, at_monotonic_s=505.0) + effects = await engine.tick( + scheduled, + TickInput(now_s=505.0, worker_counts=WorkerCounts(ready_num_decode=1)), + ) + assert old.ticks == [] and len(new.ticks) == 1 + assert effects.retune_decisions == [_decision()] + assert effects.router_parameters == ROUTER_MAP + + await engine.shutdown() + assert (old.closed, new.closed) == (1, 1) + + +def test_replace_retune_with_no_inputs_disables_retune(stub_controller): + engine = _adapter(inputs=retune_inputs()) + (old,) = stub_controller.instances + + engine.replace_retune(None, generation=2, now_s=10.0) + + assert engine.retune_controller is None + assert old.calls == ["drop_pending", "close"] + assert len(stub_controller.instances) == 1 + + +@pytest.mark.timeout(10) +async def test_rebuilt_real_controllers_never_sweep_at_once(): + """Two generations of one window never run their sweeps concurrently. + + Generation 0's sweep of window k is dropped but, like a DGDR sweep that + polls, only returns a while later; generation 1 relaunches window k in + the meantime. Its sweep must start only after generation 0's returned, + and it is decided once it is done. + """ + + pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", + ) + from dynamo.planner.retune.sweep import SweepResult + + lock = threading.Lock() + running: list[int] = [] + overlaps: list[list[int]] = [] + order: list[tuple[int, int]] = [] + gen0_started = threading.Event() + gen1_started = threading.Event() + release = threading.Event() + + def sweep(request, *, cancel, abort): + with lock: + running.append(request.generation) + order.append((request.generation, request.window_index)) + if len(running) > 1: + overlaps.append(list(running)) + try: + if request.generation == 0: + gen0_started.set() + # Sees the abort late, as a DGDR sweep at its next poll. + release.wait(30) + return SweepResult(status="failed", failure_reason="lost_leadership") + gen1_started.set() + return SweepResult( + status="ok", + num_requests=1, + scores={candidate.key(): 1.0 for candidate in request.candidates}, + ) + finally: + with lock: + running.remove(request.generation) + + engine = _adapter( + inputs=retune_inputs(), + hooks=_hooks( + sweep_fn=sweep, + launch_gate=lambda _k, _window: GateResult( + spec_replicas={"aggregated": 2}, abort=threading.Event() + ), + phase_s=0.0, + ), + ) + counts = WorkerCounts(ready_num_decode=2, expected_num_decode=2) + # On a window boundary: W_k = [now - 120, now - 60) launches at now and + # stays the only due window until now + 60. + now = 1_800_000_000.0 + try: + scheduled = engine.initial_tick(now) + await engine.tick(scheduled, TickInput(now_s=now, worker_counts=counts)) + assert gen0_started.wait(30) + + engine.replace_retune(retune_inputs(), generation=1, now_s=now) + await engine.tick(scheduled, TickInput(now_s=now + 1.0, worker_counts=counts)) + controller = engine.retune_controller + assert controller is not None and controller.generation == 1 + relaunch = controller._pending.future + # Queued behind generation 0's sweep, which still runs. + assert not gen1_started.wait(0.2) + + release.set() + relaunch.result(timeout=30) + effects = await engine.tick( + scheduled, TickInput(now_s=now + 30.0, worker_counts=counts) + ) + finally: + release.set() + await engine.shutdown() + + assert overlaps == [] + ((_, k0), (_, k1)) = order + assert [g for g, _ in order] == [0, 1] and k0 == k1 + (decision,) = effects.retune_decisions + assert decision.window_index == k0 + + +# --------------------------------------------------------------------------- +# Planner base wiring +# --------------------------------------------------------------------------- + + +class _Environment: + def __init__(self) -> None: + self.state = DeploymentState() + self.state.decode.num_gpus = 1 + self.calls: list[str] = [] + + def deployment_state(self) -> DeploymentState: + return self.state + + async def initialize(self) -> None: + self.calls.append("environment.initialize") + + async def refresh(self) -> None: + pass + + async def shutdown(self) -> None: + self.calls.append("environment.shutdown") + + +class FakeLiveRetune: + """A ``LiveRetuneHandle`` that records the base's calls.""" + + def __init__(self, calls, *, inputs=None, hooks=None): + self.calls = calls + self.inputs = inputs + self._hooks = hooks or _hooks() + self.effects: list[PlannerEffects] = [] + self.host = None + self.on_effects_error: Exception | None = None + # The host's engine when stop ran (it must still be up). + self.engine_at_stop = None + + def retune_inputs(self): + self.calls.append("live.retune_inputs") + return self.inputs + + def hooks(self): + return self._hooks + + async def start(self, host): + self.calls.append("live.start") + self.host = host + + async def stop(self): + self.calls.append("live.stop") + if self.host is not None: + self.engine_at_stop = self.host._engine + + def on_effects(self, effects): + self.effects.append(effects) + if self.on_effects_error is not None: + raise self.on_effects_error + + +def _live_config(executor: str = "dgdr") -> PlannerConfig: + config = PlannerConfig( + namespace="test-ns", + mode="agg", + enable_load_scaling=False, + enable_throughput_scaling=False, + metric_reporting_prometheus_port=0, + live_dashboard_port=0, + control_api_port=0, + max_gpu_budget=8, + ) + retune = { + "interval_seconds": 21600, + "grace_seconds": 120, + "timeout_seconds": 20400, + "router": {"overlap_score_credit": {"choices": [0.0, 1.0]}}, + "executor": executor, + } + if executor == "dgdr": + retune["dgdr"] = { + "claim_name": "request-traces", + "model": "Qwen/Qwen3-0.6B", + "sku": "h200_sxm", + "engine": {"mode": "aggregated", "backend": "vllm"}, + } + # Assigned after construction: the base wiring is under test here, not + # the live retune config validation of PlannerConfig. + config.retune = RetuneConfig.model_validate(retune) + return config + + +def _base(config: PlannerConfig, environment: _Environment) -> NativePlannerBase: + with patch("dynamo.planner.core.base.PlannerPrometheusMetrics"): + return NativePlannerBase(None, config, environment) + + +def _planner(executor="dgdr", **live_kwargs): + environment = _Environment() + planner = _base(_live_config(executor), environment) + live = FakeLiveRetune(environment.calls, **live_kwargs) + + async def build(): + environment.calls.append("live.build") + return live + + planner._build_live_retune = build # type: ignore[method-assign] + return planner, environment, live + + +def test_live_planner_rejects_the_local_executor(): + with pytest.raises( + ValueError, match="replay-only with executor 'local'; use executor: dgdr" + ): + _base(_live_config("local"), _Environment()) + + +async def test_async_init_builds_the_engine_with_live_retune_and_starts_it( + stub_controller, +): + planner, environment, live = _planner(inputs=retune_inputs()) + + async def bootstrap_regression(): + environment.calls.append("bootstrap_regression") + planner._ensure_engine() + + planner._bootstrap_regression = bootstrap_regression # type: ignore[method-assign] + await planner._async_init() + try: + assert environment.calls[:5] == [ + "environment.initialize", + "live.build", + "bootstrap_regression", + "live.retune_inputs", + "live.start", + ] + assert live.host is planner + engine = planner._engine + assert isinstance(engine, OrchestratorEngineAdapter) + (controller,) = stub_controller.instances + assert engine.retune_controller is controller + assert controller.hooks["launch_gate"] is live.hooks().launch_gate + finally: + await planner._shutdown_runtime() + + # The supervisor stops while the engine is up (it waits for the engine's + # sweep thread), then the engine closes the controller. + assert live.engine_at_stop is engine + assert controller.closed == 1 + assert environment.calls[-2:] == ["live.stop", "environment.shutdown"] + assert planner._live_retune is None + + +async def test_failed_init_still_stops_live_retune(stub_controller): + planner, environment, live = _planner(inputs=retune_inputs()) + + async def failing_start(_host): + environment.calls.append("live.start") + raise RuntimeError("lease API unavailable") + + live.start = failing_start # type: ignore[method-assign] + with pytest.raises(RuntimeError, match="lease API unavailable"): + await planner._async_init() + + assert environment.calls[-2:] == ["live.stop", "environment.shutdown"] + + +class _FakeEngine: + def __init__(self, effects): + self.effects = effects + + async def observe(self, tick, now_s): + return TickInput(now_s=now_s) + + async def tick(self, tick, tick_input): + return self.effects + + +async def test_each_tick_hands_its_effects_to_live_retune(caplog): + planner, _, live = _planner() + planner._live_retune = live + effects = PlannerEffects( + next_tick=ScheduledTick(at_s=10.0), + router_parameters=ROUTER_MAP, + retune_decisions=[_decision()], + ) + applied = [] + + async def apply_effects(applied_effects): + applied.append(applied_effects) + + planner._apply_effects = apply_effects # type: ignore[method-assign] + engine = _FakeEngine(effects) + + next_tick = await planner._run_one_tick(engine, ScheduledTick(at_s=5.0)) + assert next_tick.at_s == 10.0 + assert live.effects == [effects] + assert applied == [effects] + + # A failing outbox never fails the tick or blocks the effects. + live.on_effects_error = RuntimeError("outbox full") + with caplog.at_level(logging.ERROR): + await planner._run_one_tick(engine, ScheduledTick(at_s=5.0)) + assert "Live retune failed to take the tick's effects" in caplog.text + assert len(applied) == 2 + + +@pytest.mark.timeout(10) +async def test_replace_live_retune_holds_the_config_lock(stub_controller): + planner, _, _ = _planner(inputs=None) + await planner._async_init() + try: + await planner._config_lock.acquire() + task = asyncio.create_task( + planner.replace_live_retune(retune_inputs(), generation=2) + ) + for _ in range(5): # let the task run until it waits for the lock + await asyncio.sleep(0) + assert not task.done() and stub_controller.instances == [] + planner._config_lock.release() + await asyncio.wait_for(task, timeout=5) + + (controller,) = stub_controller.instances + assert planner._engine.retune_controller is controller # type: ignore[union-attr] + assert controller.generation == 2 + assert controller.phases == [PHASE] + finally: + await planner._shutdown_runtime() + + +# --------------------------------------------------------------------------- +# Runtime reconfiguration +# --------------------------------------------------------------------------- + + +def _sla_planner(**retune) -> NativePlannerBase: + """A planner whose validated dgdr config scores with SLA 500/50 ms by + default, with its live retune supervisor installed.""" + + config = live_config( + ttft_ms=500.0, + itl_ms=50.0, + retune=retune, + metric_reporting_prometheus_port=0, + live_dashboard_port=0, + control_api_port=0, + ) + planner = _base(config, _Environment()) + planner._live_retune = FakeLiveRetune([]) + return planner + + +async def _patch(planner: NativePlannerBase, changes: dict) -> dict: + """PATCH ``changes`` through the runtime API; an accepted update is + awaited and answered with its operation.""" + + reconfiguration = planner.reconfiguration + request_id = f"update-{planner._config_generation}" + body = ReconfigureRequest( + instance_id=reconfiguration.instance_id, + expected_revision=planner._config_generation, + request_id=request_id, + config=changes, + ).model_dump() + response = await reconfiguration.handle_request( + {"method": "PATCH", "path": "config", "body": body} + ) + if response["status"] != 202: + return response + await reconfiguration._task + return {"status": 202, "body": reconfiguration.get_operation(request_id)} + + +@pytest.mark.parametrize( + ("retune", "changes", "inherited"), + [ + ({}, {"ttft_ms": 250.0}, "ttft_ms"), + ({}, {"itl_ms": 25.0}, "itl_ms"), + ({}, {"ttft_ms": 250.0, "itl_ms": 25.0}, "ttft_ms and itl_ms"), + ({}, {"ttft_ms": 250.0, "max_gpu_budget": 16}, "ttft_ms"), + ({"sla_ttft_ms": 400.0}, {"itl_ms": 25.0}, "itl_ms"), + ], +) +async def test_live_retune_rejects_a_runtime_change_of_its_inherited_sla( + retune, changes, inherited +): + planner = _sla_planner(**retune) + # The update passes every other runtime check: only live retune stops it. + candidate = candidate_config(planner.config, changes) + assert {field: getattr(candidate, field) for field in changes} == changes + before = await planner.reconfiguration.get_config() + + response = await _patch(planner, changes) + + assert response["status"] == 422 + overrides = " and ".join(f"retune.sla_{f}" for f in inherited.split(" and ")) + assert response["body"]["error"] == ( + f"cannot change {inherited} at runtime while retune.executor is 'dgdr': " + f"live retune inherits its scoring SLA from {inherited} ({overrides} unset); " + f"restart the planner with the new {inherited}, or restart it with " + f"{overrides} set to keep {inherited} changeable at runtime" + ) + assert await planner.reconfiguration.get_config() == before + assert (planner.config.ttft_ms, planner.config.itl_ms) == (500.0, 50.0) + await planner._shutdown_runtime() + + +@pytest.mark.parametrize( + ("retune", "changes"), + [ + ( + {"sla_ttft_ms": 400.0, "sla_itl_ms": 40.0}, + {"ttft_ms": 250.0, "itl_ms": 25.0}, + ), + ({"sla_ttft_ms": 400.0}, {"ttft_ms": 250.0}), + ({}, {"ttft_ms": 500.0, "max_gpu_budget": 16}), + ], +) +async def test_live_retune_accepts_runtime_updates_that_keep_its_scoring_sla( + retune, changes +): + planner = _sla_planner(**retune) + config = planner.config + assert config.retune is not None + scoring_sla = config.retune.scoring_sla(config.ttft_ms, config.itl_ms) + + response = await _patch(planner, changes) + + assert response["body"]["state"] == "applied" + assert {field: getattr(config, field) for field in changes} == changes + assert config.retune.scoring_sla(config.ttft_ms, config.itl_ms) == scoring_sla + await planner._shutdown_runtime() + + +async def test_planner_without_live_retune_changes_its_sla_at_runtime(): + config = PlannerConfig( + namespace="test-ns", + mode="agg", + metric_reporting_prometheus_port=0, + live_dashboard_port=0, + control_api_port=0, + ) + planner = _base(config, _Environment()) + + response = await _patch(planner, {"ttft_ms": 250.0, "itl_ms": 25.0}) + + assert response["body"]["state"] == "applied" + assert (planner.config.ttft_ms, planner.config.itl_ms) == (250.0, 25.0) + await planner._shutdown_runtime() + + +# --------------------------------------------------------------------------- +# SIGTERM +# --------------------------------------------------------------------------- + + +class _Runtime: + def __init__(self) -> None: + self.stopped = asyncio.Event() + self.shutdowns = 0 + + def shutdown(self) -> None: + self.shutdowns += 1 + self.stopped.set() + + +def _found() -> asyncio.Event: + found = asyncio.Event() + found.set() + return found + + +def _sigterm_installed() -> bool: + loop = asyncio.get_running_loop() + return signal.SIGTERM in loop._signal_handlers # type: ignore[attr-defined] + + +async def test_sigterm_cancels_the_planner_run_and_stops_serving(): + runtime = _Runtime() + events: list[str] = [] + running = asyncio.Event() + + async def serve(): + await runtime.stopped.wait() + events.append("serve stopped") + + async def run(): + running.set() + try: + await asyncio.sleep(3600) + finally: + # run()'s finally: the planner shutdown that releases the Lease. + await asyncio.sleep(0) + events.append("run finally") + + main = asyncio.create_task( + planner_main._run_until_sigterm(runtime, serve(), run(), _found()) + ) + await asyncio.wait_for(running.wait(), timeout=5) + assert _sigterm_installed() + os.kill(os.getpid(), signal.SIGTERM) + os.kill(os.getpid(), signal.SIGTERM) # a second one changes nothing + await asyncio.wait_for(main, timeout=5) + + assert events == ["run finally", "serve stopped"] + assert runtime.shutdowns == 1 + assert not _sigterm_installed() + + +async def test_sigterm_is_handled_only_once_a_retune_lease_is_found(): + # The clean-up check of a removed live retune runs in the background; + # a planner whose check finds no Lease keeps the default SIGTERM. + runtime = _Runtime() + found = asyncio.Event() + running = asyncio.Event() + + async def serve(): + await runtime.stopped.wait() + + async def run(): + running.set() + await asyncio.sleep(3600) + + main = asyncio.create_task( + planner_main._run_until_sigterm(runtime, serve(), run(), found) + ) + await asyncio.wait_for(running.wait(), timeout=5) + assert not _sigterm_installed() + found.set() + for _ in range(3): + await asyncio.sleep(0) + assert _sigterm_installed() + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.wait_for(main, timeout=5) + assert runtime.shutdowns == 1 and not _sigterm_installed() + + +async def test_run_failure_and_outer_cancellation_propagate(): + runtime = _Runtime() + + async def serve(): + await runtime.stopped.wait() + + async def failing(): + raise RuntimeError("tick failed") + + with pytest.raises(RuntimeError, match="tick failed"): + await planner_main._run_until_sigterm(runtime, serve(), failing(), _found()) + assert runtime.shutdowns == 0 and not _sigterm_installed() + + started = asyncio.Event() + + async def run(): + started.set() + await asyncio.sleep(3600) + + # Never found: the waiting installer is cancelled with the call. + main = asyncio.create_task( + planner_main._run_until_sigterm(runtime, serve(), run(), asyncio.Event()) + ) + await asyncio.wait_for(started.wait(), timeout=5) + main.cancel() + with pytest.raises(asyncio.CancelledError): + await main + assert runtime.shutdowns == 0 + runtime.shutdown() # let the serve tasks finish + + +async def test_start_planner_runs_the_sigterm_path_only_when_a_lease_may_be_held( + monkeypatch, +): + # With live retune, or the clean-up check of a removed one (every + # Kubernetes planner without retune config), the planner may hold the + # retune leadership Lease. + seen: list[asyncio.Event] = [] + + class _Planner: + def __init__(self, may_hold: bool) -> None: + self.may_hold_retune_lease = may_hold + self.retune_lease_found = asyncio.Event() + + async def _async_init(self): + pass + + async def run(self): + return None + + class _Endpoint: + async def serve_endpoint(self, *_args, **_kwargs): + return None + + class _StartRuntime: + def endpoint(self, _name): + return _Endpoint() + + async def fake_run_until_sigterm(runtime, serve, run, lease_found): + seen.append(lease_found) + await asyncio.gather(serve, run) + + planners = iter([_Planner(False), _Planner(True)]) + monkeypatch.setattr(planner_main, "construct_planner", lambda **_: next(planners)) + monkeypatch.setattr(planner_main, "_run_until_sigterm", fake_run_until_sigterm) + + await planner_main.start_planner(_StartRuntime(), PlannerConfig(mode="agg")) + assert seen == [] + await planner_main.start_planner(_StartRuntime(), PlannerConfig(mode="agg")) + assert len(seen) == 1 and not seen[0].is_set() + + +async def test_live_retune_finds_its_lease_at_start(stub_controller): + planner, _, _ = _planner(inputs=None) + assert not planner.may_hold_retune_lease + await planner._async_init() + try: + assert planner.may_hold_retune_lease + assert planner.retune_lease_found.is_set() + finally: + await planner._shutdown_runtime() diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_config.py b/components/src/dynamo/planner/tests/unit/test_live_retune_config.py new file mode 100644 index 000000000000..ef2cbddf43d4 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_config.py @@ -0,0 +1,719 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the retune imports. + +"""Configuration premises of live retune (``retune.executor: dgdr``). + +P1 (Kubernetes, decision latency, explicit timeout within the interval), P2 +(retune-only, router only), P3 (the DGDR router knobs), P5 (a pinned engine +section) and P13 (the largest grid shard fits the sweep deadline), plus the +shard-split and need-seconds helpers the executor shares with validation. +""" + +from __future__ import annotations + +import copy +import math + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from aisimulate.config.engine import ParallelismRecommendationConfig +from pydantic import ValidationError + +from dynamo.planner.config.planner_config import ( + DGDR_MAX_SHARD_POINTS, + LIVE_RETUNE_K8S_TIMEOUT_SECONDS, + DgdrExecutorConfig, + PlannerConfig, + balanced_shard_parts, + dgdr_need_seconds, + dgdr_rounds, + largest_shard_points, + live_retune_min_lease_seconds, +) +from dynamo.planner.retune.domains import expand_router_domain +from dynamo.planner.retune.live import dgdr + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +PRESET = { + "tensor": 2, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, +} +SCHEDULER = {"max_batched_tokens": 8192, "max_sequences": 256} + + +def _worker(**overrides): + return { + "parallelism": {"preset": [dict(PRESET)]}, + "scheduler": dict(SCHEDULER), + "kv_cache": {"block_size": 64, "capacity": {"memory_fraction": 0.9}}, + **overrides, + } + + +ENGINE_AGG = { + "mode": "aggregated", + "backend": "vllm", + "backend_version": "0.11.0", + "context_length": 32768, + "workers": {"aggregated": _worker(), "prefill": None, "decode": None}, +} +ENGINE_DISAGG = { + "mode": "disaggregated", + "backend": "vllm", + "backend_version": {"vllm": "0.11.0"}, + "context_length": "max", + "workers": {"prefill": _worker(), "decode": _worker()}, +} +# The live retune defaults: T = 6 h, timeout = T - 20 min. +INTERVAL = 21600.0 +TIMEOUT = 20400.0 +# The default router domain is 3 x 9 x 4 = 108 points; with max_shards 3 the +# largest shard has 36: 900 + 300 + (1 + ceil(35 / 8)) x 1800 = 12000 s. +DEFAULT_NEED = 12000.0 + + +def _dgdr(**overrides): + return { + "claim_name": "request-traces", + "model": "Qwen/Qwen3-32B", + "sku": "h200_sxm", + "engine": copy.deepcopy(ENGINE_AGG), + **overrides, + } + + +def _retune(**overrides): + return { + "interval_seconds": INTERVAL, + "timeout_seconds": TIMEOUT, + "router": "default", + "executor": "dgdr", + "dgdr": _dgdr(), + **overrides, + } + + +def _config(retune=None, **overrides) -> PlannerConfig: + values = { + "namespace": "test-ns", + "environment": "kubernetes", + "mode": "agg", + "enable_load_scaling": False, + "enable_throughput_scaling": False, + "retune": _retune() if retune is None else retune, + **overrides, + } + return PlannerConfig(**values) + + +def _rejects(match: str, retune=None, **overrides) -> None: + with pytest.raises(ValidationError, match=match): + _config(retune, **overrides) + + +def _engine(mode: str = "agg") -> dict: + return copy.deepcopy(ENGINE_AGG if mode == "agg" else ENGINE_DISAGG) + + +# ------------------------------------------------------------------ valid + + +@pytest.mark.parametrize("mode", ["agg", "disagg"]) +def test_live_retune_config_accepts_the_design_example(mode): + retune = _retune(dgdr=_dgdr(engine=_engine(mode))) + config = _config(retune, mode=mode) + + assert config.retune is not None + assert config.retune.executor == "dgdr" + assert config.retune.dgdr is not None + assert config.retune.dgdr.engine == _engine(mode) + # Validation never rewrites the sections it checks. + assert config.retune.router == "default" + + +def test_live_retune_accepts_a_measured_decision_latency_at_scale_one(): + config = _config(_retune(decision_latency_seconds="measured")) + assert config.retune is not None + assert config.retune.decision_latency_seconds == "measured" + + +def test_live_retune_accepts_the_largest_timeout_the_interval_allows(): + config = _config(_retune(timeout_seconds=INTERVAL - 600)) + assert config.retune is not None + + +def test_live_retune_accepts_the_three_dgdr_router_knobs_and_none_load_model(): + router = { + "overlap_score_credit": {"choices": [0.0, 1.0]}, + "prefill_load_scale": {"range": {"min": 0.5, "max": 2.0, "step": 0.5}}, + "temperature": 0.0, + "policy": "kv_router", + "prefill_load_model": {"type": "none"}, + } + assert _config(_retune(router=router)).retune is not None + router["prefill_load_model"] = {"type": {"choices": ["none"]}} + assert _config(_retune(router=router)).retune is not None + + +def test_replay_retune_keeps_its_rules(): + # executor local: none of the live premises apply. + config = _config( + {"interval_seconds": 600, "router": "default"}, + environment="virtual", + ) + assert config.retune is not None + assert config.retune.executor == "local" + + +# --------------------------------------------------------------------- P1 + + +def test_p1_requires_kubernetes(): + _rejects(r"environment='kubernetes' \(P1\)", environment="virtual") + + +@pytest.mark.parametrize("latency", [30.0, 0.5]) +def test_p1_rejects_a_fixed_nonzero_decision_latency(latency): + _rejects( + r"decision_latency_seconds must be 0 or 'measured' \(P1\)", + _retune(decision_latency_seconds=latency), + ) + + +def test_p1_rejects_a_scaled_measured_decision_latency(): + # Valid for RetuneConfig (2 x 10000 <= 21600), rejected by P1 before the + # time budget is checked. + _rejects( + r"decision_latency_scale 1 \(P1\)", + _retune( + decision_latency_seconds="measured", + decision_latency_scale=2.0, + timeout_seconds=10000.0, + ), + ) + + +def test_p1_requires_an_explicit_timeout(): + retune = _retune() + del retune["timeout_seconds"] + _rejects(r"explicit retune.timeout_seconds \(P1\)", retune) + + +def test_p1_rejects_a_timeout_within_ten_minutes_of_the_interval(): + _rejects( + r"timeout_seconds must be <= interval_seconds - 600 = 21000 .*\(P1\)", + _retune(timeout_seconds=INTERVAL - 599), + ) + + +# --------------------------------------------------------------------- P2 + + +@pytest.mark.parametrize( + "scaling", + [ + {"enable_load_scaling": True, "enable_throughput_scaling": False}, + { + "enable_load_scaling": False, + "enable_throughput_scaling": True, + "optimization_target": "sla", + }, + ], +) +def test_p2_requires_retune_only(scaling): + _rejects(r"retune-only mode \(P2\)", **scaling) + + +def test_p2_requires_both_scaling_flags_explicitly(): + # Without both flags written, easy mode keeps load scaling on. + with pytest.raises(ValidationError, match=r"retune-only mode \(P2\)"): + PlannerConfig( + namespace="test-ns", + environment="kubernetes", + mode="agg", + enable_throughput_scaling=False, + retune=_retune(), + ) + + +def test_p2_rejects_a_planner_domain(): + _rejects( + r"retunes the router only \(P2\)", + _retune(planner={"load_sensitivity": {"preset": ["default"]}}), + ) + + +# --------------------------------------------------------------------- P3 + + +@pytest.mark.parametrize( + "knob", ["overlap_score_credit_decay", "decode_active_request_weight"] +) +def test_p3_rejects_retune_extension_knobs(knob): + _rejects( + rf"retune.router \['{knob}'\] cannot be searched with executor 'dgdr' \(P3\)", + _retune(router={knob: {"choices": [0.0, 1.0]}}), + ) + + +@pytest.mark.parametrize( + "load_model", + [{"type": "ais"}, {"type": {"choices": ["none", "ais"]}}], +) +def test_p3_requires_the_none_prefill_load_model(load_model): + _rejects( + r"prefill_load_model.type must be 'none' with executor 'dgdr' \(P3\)", + _retune(router={"temperature": 0.0, "prefill_load_model": load_model}), + ) + + +def test_p3_keeps_the_schema_errors_for_other_policies(): + # RouterRetuneConfig already pins the policy; its error stays first. + _rejects( + r"retune.router.policy is fixed at startup", + _retune(router={"policy": "round_robin"}), + ) + + +# --------------------------------------------------------------------- P5 + + +def _rejects_engine(match: str, engine: dict, mode: str = "agg") -> None: + _rejects(match, _retune(dgdr=_dgdr(engine=engine)), mode=mode) + + +@pytest.mark.parametrize( + ("path", "value", "match"), + [ + ( + ("workers", "aggregated", "scheduler", "max_sequences"), + {"choices": [128, 256]}, + r"engine.workers.aggregated.scheduler.max_sequences is a search " + r"domain \(choices\); .* \(P5\)", + ), + ( + ("workers", "aggregated", "kv_cache", "capacity", "memory_fraction"), + {"range": {"min": 0.8, "max": 0.9}}, + r"memory_fraction is a search domain \(range\)", + ), + ( + ("backend",), + {"choices": ["vllm", "sglang"]}, + r"engine.backend is a search domain", + ), + ( + ("workers", "aggregated", "timing", "systems_paths"), + ["/opt/systems"], + r"engine.workers.aggregated.timing.systems_paths is a list; .* \(P5\)", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "replicas"), + 2, + r"engine.workers.aggregated.parallelism.preset\[0\].replicas must " + r"not be set: the DGD spec provides replicas at each launch \(P5\)", + ), + ( + ("workers", "aggregated", "parallelism", "replicas"), + 2, + r"parallelism.replicas must not be set: the DGD spec", + ), + ( + ("workers", "aggregated", "parallelism", "preset"), + [dict(PRESET), dict(PRESET, tensor=4)], + r"parallelism.preset must be a list of exactly one mapping .* \(P5\)", + ), + ( + ("workers", "aggregated", "parallelism", "preset"), + "default", + r"parallelism.preset must be a list of exactly one mapping", + ), + ( + ("workers", "aggregated", "parallelism", "tensor"), + 2, + r"parallelism sets \['tensor'\] next to preset", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "tensor"), + True, + r"preset\[0\].tensor must be a positive integer \(P5\)", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "expert"), + 1, + r"preset\[0\] must set exactly .*moe_expert and optionally " + r"prefill_context: 1 \(P5\); missing=\[\], unknown=\['expert'\]", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "decode_context"), + 1, + r"preset\[0\].decode_context is not accepted \(P5\); the pinned " + r"AISimulate preset schema has no decode_context", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "prefill_context"), + 2, + r"preset\[0\].prefill_context must be 1 or omitted \(P5\); AISimulate " + r"recommend rejects context parallelism above 1, got 2", + ), + ( + ("workers", "aggregated", "parallelism", "preset", 0, "prefill_context"), + 0, + r"preset\[0\].prefill_context must be a positive integer \(P5\), got 0", + ), + ( + ("workers", "aggregated", "scheduler", "max_batched_tokens"), + None, + r"scheduler.max_batched_tokens must be a positive integer \(P5\); " + r"omitted, AISimulate searches its default domain", + ), + ( + ("mode",), + "disaggregated", + r"engine.mode must be 'aggregated' for planner mode 'agg' \(P5\)", + ), + ( + ("model",), + "Qwen/Qwen3-32B", + r"engine.model would override the Request's typed field \(P5\); " + r"set retune.dgdr.model", + ), + ( + ("hardware",), + "h200_sxm", + r"set retune.dgdr.sku instead", + ), + ( + ("workers", "decode"), + _worker(), + r"engine.workers.decode is not a worker role of planner mode 'agg'", + ), + ( + ("workers", "encoder"), + {"tensor": 1}, + r"engine.workers.encoder is not a worker role", + ), + ], +) +def test_p5_rejects_an_unpinned_engine(path, value, match): + engine = _engine() + target = engine + for key in path[:-1]: + target = target.setdefault(key, {}) if isinstance(key, str) else target[key] + target[path[-1]] = value + _rejects_engine(match, engine) + + +def test_p5_rejects_a_missing_preset_knob(): + engine = _engine() + del engine["workers"]["aggregated"]["parallelism"]["preset"][0]["moe_expert"] + _rejects_engine(r"missing=\['moe_expert'\]", engine) + + +def test_p5_requires_mode(): + engine = _engine() + del engine["mode"] + _rejects_engine(r"an omitted mode searches both", engine) + + +def test_p5_requires_every_role_of_the_planner_mode(): + engine = _engine("disagg") + del engine["workers"]["decode"] + _rejects_engine( + r"engine.workers.decode is required for planner mode 'disagg' \(P5\)", + engine, + mode="disagg", + ) + + +def test_p5_requires_the_disaggregated_engine_mode_in_disagg(): + engine = _engine("disagg") + engine["mode"] = "aggregated" + _rejects_engine( + r"engine.mode must be 'disaggregated' for planner mode 'disagg'", + engine, + mode="disagg", + ) + + +def test_p5_engine_backend_must_be_the_deployed_backend(): + engine = _engine() + engine["backend"] = "sglang" + # PlannerConfig.backend defaults to vllm; every Run would fail its build. + _rejects_engine( + r"retune.dgdr.engine.backend 'sglang' differs from the deployed " + r"backend 'vllm' \(P5\)", + engine, + ) + config = _config(_retune(dgdr=_dgdr(engine=engine)), backend="sglang") + assert config.retune is not None and config.retune.dgdr is not None + assert config.retune.dgdr.engine["backend"] == "sglang" + _rejects( + r"retune.dgdr.engine.backend 'vllm' differs from the deployed backend " + r"'trtllm' \(P5\)", + _retune(dgdr=_dgdr(engine=_engine())), + backend="trtllm", + ) + + +def test_p5_accepts_an_omitted_engine_backend(): + # The Request fills it with the deployed backend. + engine = _engine() + del engine["backend"] + for backend in ("vllm", "sglang"): + config = _config( + _retune(dgdr=_dgdr(engine=copy.deepcopy(engine))), backend=backend + ) + assert config.retune is not None + + +def test_live_retune_accepts_nested_claim_directories(): + dgdr = _dgdr(segments_dir="retune/segments", windows_dir="retune/windows") + assert _config(_retune(dgdr=dgdr)).retune is not None + + +def test_p5_accepts_unused_roles_absent_or_null(): + engine = _engine() + del engine["workers"]["prefill"] + engine["workers"]["encoder"] = None + assert _config(_retune(dgdr=_dgdr(engine=engine))).retune is not None + + +def test_p5_accepts_prefill_context_one(): + engine = _engine() + engine["workers"]["aggregated"]["parallelism"]["preset"][0]["prefill_context"] = 1 + assert _config(_retune(dgdr=_dgdr(engine=engine))).retune is not None + + +def test_p5_rejects_decode_context_on_a_disaggregated_role(): + engine = _engine("disagg") + engine["workers"]["decode"]["parallelism"]["preset"][0]["decode_context"] = 1 + _rejects_engine( + r"engine.workers.decode.parallelism.preset\[0\].decode_context is not " + r"accepted \(P5\)", + engine, + mode="disagg", + ) + + +def _spec_replicas(mode: str) -> dict[str, int]: + return {"aggregated": 2} if mode == "agg" else {"prefill": 1, "decode": 3} + + +def _validate_with_pinned_aisimulate(engine: dict, mode: str) -> None: + spec_replicas = _spec_replicas(mode) + pinned = dgdr._pinned_engine(engine, spec_replicas, "vllm") + for role in spec_replicas: + ParallelismRecommendationConfig.model_validate( + pinned["workers"][role]["parallelism"] + ) + + +@pytest.mark.parametrize("mode", ["agg", "disagg"]) +@pytest.mark.parametrize("prefill_context", [None, 1]) +def test_p5_presets_validate_against_the_pinned_aisimulate(mode, prefill_context): + engine = _engine(mode) + if prefill_context is not None: + for role in _spec_replicas(mode): + preset = engine["workers"][role]["parallelism"]["preset"] + preset[0]["prefill_context"] = prefill_context + config = _config(_retune(dgdr=_dgdr(engine=engine)), mode=mode) + assert config.retune is not None and config.retune.dgdr is not None + + _validate_with_pinned_aisimulate(config.retune.dgdr.engine, mode) + + +def test_pinned_aisimulate_rejects_decode_context(): + # Why P5 rejects decode_context; once a pin bump makes this fail, P5 may + # admit it. + engine = _engine() + engine["workers"]["aggregated"]["parallelism"]["preset"][0]["decode_context"] = 1 + with pytest.raises(ValidationError, match=r"decode_context\s+Extra inputs"): + _validate_with_pinned_aisimulate(engine, "agg") + + +# -------------------------------------------------------------- claim paths + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"segments_dir": "/trace/segments"}, r"segments_dir must be a relative"), + ({"windows_dir": "../windows"}, r"windows_dir must be a relative"), + ({"windows_dir": "."}, r"windows_dir must be a relative"), + ({"windows_dir": "a/./windows"}, r"windows_dir must be a relative"), + ({"windows_dir": "windows/"}, r"windows_dir must be a relative"), + ({"segments_dir": "a//segments"}, r"segments_dir must be a relative"), + ({"mount_path": "trace"}, r"mount_path must be absolute"), + ({"windows_dir": "segments"}, r"must not contain one another"), + ({"windows_dir": "segments/windows"}, r"must not contain one another"), + ], +) +def test_live_retune_rejects_unsafe_claim_paths(overrides, match): + _rejects(match, _retune(dgdr=_dgdr(**overrides))) + + +# --------------------------------------------------------- leadership Lease + + +def test_live_retune_lease_fits_two_renewals_before_its_margin(): + minimum = live_retune_min_lease_seconds(LIVE_RETUNE_K8S_TIMEOUT_SECONDS) + # 3 x 5 s margin + 2 x (2 x 5 s + 1 s) I/O budget. + assert minimum == 37.0 + for lease in (15.0, minimum - 0.001): + _rejects( + r"retune.dgdr.lease_duration_seconds must be at least the lease " + r"margin \+ 2 x the I/O budget of one renewal \(15s \+ 2 x 11s = 37s " + r"with the 5s Kubernetes request timeout\)", + _retune(dgdr=_dgdr(lease_duration_seconds=lease)), + ) + config = _config(_retune(dgdr=_dgdr(lease_duration_seconds=minimum))) + assert config.retune is not None and config.retune.dgdr is not None + assert config.retune.dgdr.lease_duration_seconds == minimum + schema = DgdrExecutorConfig.model_json_schema()["properties"] + assert schema["lease_duration_seconds"]["minimum"] == minimum + assert "At least 37 s" in schema["lease_duration_seconds"]["description"] + + +# -------------------------------------------------------------------- P13 + + +def test_p13_accepts_a_shard_that_exactly_fits(): + assert _config(_retune(timeout_seconds=DEFAULT_NEED)).retune is not None + + +def test_p13_rejects_a_largest_shard_that_cannot_finish(): + _rejects( + r"largest retune grid shard \(36 points\) needs 12000s: " + r"cut_budget_seconds 900 \+ job_startup_seconds 300 \+ 6 rounds x " + r"candidate_timeout_seconds 1800, more than " + r"retune.timeout_seconds=11999 \(P13\)", + _retune(timeout_seconds=DEFAULT_NEED - 1), + ) + + +def test_p13_counts_the_largest_shard_of_fewer_max_shards(): + # max_shards 2: 3 x 9 x 2 = 54 points -> 1 + ceil(53 / 8) = 8 rounds. + dgdr = _dgdr(run={"max_shards": 2}) + need = 900 + 300 + 8 * 1800 + assert _config(_retune(dgdr=dgdr, timeout_seconds=need)).retune is not None + _rejects(r"\(54 points\)", _retune(dgdr=dgdr, timeout_seconds=need - 1)) + + +def test_p13_rejects_a_grid_that_needs_more_shards(): + _rejects( + r"router grid of 108 points cannot be split into at most " + r"retune.dgdr.run.max_shards=1 Cartesian shards of at most 100 points " + r"\(P13\); the largest would have 108", + _retune(dgdr=_dgdr(run={"max_shards": 1})), + ) + + +def test_p13_uses_the_configured_run_parallelism_and_timeouts(): + dgdr = _dgdr( + cut_budget_seconds=60, + job_startup_seconds=0, + run={"parallelism": 36, "candidate_timeout_seconds": 100}, + ) + # 60 + 0 + (1 + ceil(35 / 36)) x 100 = 260. + assert _config(_retune(dgdr=dgdr, timeout_seconds=260)).retune is not None + _rejects(r"needs 260s", _retune(dgdr=dgdr, timeout_seconds=259)) + + +# --------------------------------------------------------------- helpers + + +def test_balanced_shard_parts_minimizes_the_rounds_exhaustively(): + for sizes in ((3, 9, 4), (5, 7), (2, 3, 5), (12,), (1, 1, 30), (250,)): + for max_shards in range(1, 17): + best = min( + largest_shard_points(sizes, (a, b, c)[: len(sizes)]) + for a in range(1, max_shards + 1) + for b in range(1, max_shards + 1) + for c in range(1, max_shards + 1) + if math.prod((a, b, c)[: len(sizes)]) <= max_shards + and all(p <= n for p, n in zip((a, b, c)[: len(sizes)], sizes)) + ) + # With the Run parallelism, the split takes the fewest rounds, so + # its largest shard fits and finishes exactly when the smallest + # largest shard does: P13 and the executor agree. + for parallelism in (1, 8, 36): + parts = balanced_shard_parts(sizes, max_shards, parallelism=parallelism) + largest = largest_shard_points(sizes, parts) + if best > DGDR_MAX_SHARD_POINTS: + assert largest == best + else: + assert largest <= DGDR_MAX_SHARD_POINTS + assert dgdr_rounds(largest, parallelism) == dgdr_rounds( + best, parallelism + ) + + +@pytest.mark.parametrize( + ("sizes", "max_shards", "parallelism", "parts"), + [ + # A shard is added only when it saves an evaluation round; the first + # candidate runs alone, so single points save one. + ((2, 2, 2), 3, 8, (1, 1, 1)), + ((2, 2, 2), 16, 8, (2, 2, 2)), + ((2, 2, 2), 3, 1, (1, 1, 2)), + ((3, 9, 4), 3, 8, (1, 3, 1)), + ], +) +def test_balanced_shard_parts_with_parallelism(sizes, max_shards, parallelism, parts): + assert balanced_shard_parts(sizes, max_shards, parallelism=parallelism) == parts + + +@pytest.mark.parametrize( + ("sizes", "max_shards", "parallelism"), + [((3, 0), 3, 8), ((3,), 0, 8), ((3,), 3, 0)], +) +def test_balanced_shard_parts_rejects_bad_input(sizes, max_shards, parallelism): + with pytest.raises(ValueError): + balanced_shard_parts(sizes, max_shards, parallelism=parallelism) + + +def test_default_router_grid_splits_into_three_shards_of_36(): + retune = _config().retune + assert retune is not None + sizes = [len(values) for values in expand_router_domain(retune).values()] + assert sizes == [3, 9, 4] + assert retune.dgdr is not None + parts = balanced_shard_parts(sizes, 3, parallelism=retune.dgdr.run.parallelism) + assert largest_shard_points(sizes, parts) == 36 <= DGDR_MAX_SHARD_POINTS + + +@pytest.mark.parametrize( + ("points", "needs_cut", "expected"), + [ + (1, True, 900 + 300 + 1 * 1800), + (1, False, 300 + 1 * 1800), + (9, True, 900 + 300 + 2 * 1800), + (10, True, 900 + 300 + 3 * 1800), + (36, False, 300 + 6 * 1800), + ], +) +def test_dgdr_need_seconds(points, needs_cut, expected): + dgdr = DgdrExecutorConfig.model_validate(_dgdr()) + assert dgdr_need_seconds(points, dgdr, needs_cut=needs_cut) == expected + + +def test_dgdr_need_seconds_rejects_an_empty_shard(): + dgdr = DgdrExecutorConfig.model_validate(_dgdr()) + with pytest.raises(ValueError): + dgdr_need_seconds(0, dgdr, needs_cut=True) diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_controller_hooks.py b/components/src/dynamo/planner/tests/unit/test_live_retune_controller_hooks.py new file mode 100644 index 000000000000..a0f52a33e496 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_controller_hooks.py @@ -0,0 +1,599 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the controller imports the sweep. + +"""Live retune hooks of ``RetuneController`` with a scripted sweep function. + +Covers what a live host adds on top of the replay loop: epoch-phase +anchoring, the launch gate, harvest gating and abandoned sweeps, +``drop_pending``/``rewind``, the DGD spec replicas check, ``restore_desired`` +and the structured failure reason of a decision. The sweep thread shared +across generations is tested with the engine adapter. +""" + +from __future__ import annotations + +import json +import threading +from collections import deque + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from dynamo.llm import KvRouterConfig +from dynamo.planner.config.planner_config import PlannerConfig, RetuneEngineSpec +from dynamo.planner.core.types import ( + RetuneInputs, + RetuneParams, + WorkerCapabilities, + WorkerCounts, +) +from dynamo.planner.plugins.clock import VirtualClock, WallClock +from dynamo.planner.retune.controller import RetuneController +from dynamo.planner.retune.live.contracts import NON_DECISION_FAILURES, GateResult +from dynamo.planner.retune.sweep import SweepResult + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +INTERVAL = 60.0 +GRACE = 5.0 +STEP = 5.0 +# An epoch-like planner time on a window boundary (phase 0): W_K = [T0-60, T0). +T0 = 1_800_000_000.0 +K = int(T0 // INTERVAL) - 1 +L_K = T0 + GRACE +ROUTER = { + "overlap_score_credit": {"choices": [0.0, 1.0]}, + "prefill_load_scale": 1.0, + "temperature": {"choices": [0.0, 0.5]}, +} +ROUTER_DOMAIN = { + "overlap_score_credit": (0.0, 1.0), + "prefill_load_scale": (1.0,), + "router_temperature": (0.0, 0.5), +} +SPEC = {"aggregated": 2} + + +def _point(credit, temperature): + return RetuneParams( + router={ + "overlap_score_credit": credit, + "prefill_load_scale": 1.0, + "router_temperature": temperature, + } + ) + + +STARTUP = _point(1.0, 0.0) +BEST = _point(0.0, 0.5) + + +class FakeSweeps: + """Scripted sweep function: records each call, returns queued results.""" + + def __init__(self): + self.requests = [] + self.cancels = [] + self.aborts = [] + self.results = deque() + self.default = None + + def __call__(self, request, *, cancel=None, abort=None): + self.requests.append(request) + self.cancels.append(cancel) + self.aborts.append(abort) + result = self.results.popleft() if self.results else self.default + return result(request) if callable(result) else result + + +class Gate: + """Scripted ``launch_gate``: a queue of results, then ``default``.""" + + def __init__(self, *results, default=None): + self.results = deque(results) + self.default = default if default is not None else GateResult() + self.calls = [] + + def __call__(self, k, window): + self.calls.append((k, window)) + result = self.results.popleft() if self.results else self.default + if isinstance(result, Exception): + raise result + return result + + +class Harvest: + """Scripted ``harvest_allowed``.""" + + def __init__(self, *answers, default=True): + self.answers = deque(answers) + self.default = default + self.calls = [] + + def __call__(self, generation): + self.calls.append(generation) + return self.answers.popleft() if self.answers else self.default + + +def _scored(score_of, *, wall_s=0.5, **extra): + def build(request): + return SweepResult( + status="ok", + num_requests=10, + scores={c.key(): score_of(c) for c in request.candidates}, + wall_s=wall_s, + **extra, + ) + + return build + + +def _flat(score=1.0, **extra): + return _scored(lambda _c: score, **extra) + + +def _config(**retune): + return PlannerConfig( + namespace="test-ns", + mode="agg", + enable_load_scaling=False, + enable_throughput_scaling=False, + retune={ + "interval_seconds": INTERVAL, + "grace_seconds": GRACE, + "router": ROUTER, + **retune, + }, + ) + + +def _inputs(origin_ms=0.0): + return RetuneInputs( + engine=RetuneEngineSpec( + router_mode="kv_router", + router_config=json.loads(KvRouterConfig().to_json()), + ), + trace_paths=("/trace/segments/dgd",), + trace_origin_ms=origin_ms, + ) + + +def _live( + sweeps, + *, + gate=None, + harvest=None, + clock=None, + generation=3, + start=T0, + phase=0.0, + config=None, + executor=None, +): + controller = RetuneController( + config or _config(), + _inputs(), + WorkerCapabilities(), + clock or VirtualClock(), + STEP, + sweep_fn=sweeps, + launch_gate=( + gate if gate is not None else Gate(default=GateResult(spec_replicas=SPEC)) + ), + harvest_allowed=harvest, + generation=generation, + executor=executor, + ) + controller.start(start, phase) + return controller + + +def _counts(decode=2): + return WorkerCounts(ready_num_decode=decode) + + +async def _tick(controller, now, decode=2): + return await controller.on_tick(now, _counts(decode)) + + +def _window_ms(k): + return k * INTERVAL * 1000.0, (k + 1) * INTERVAL * 1000.0 + + +# --------------------------------------------------------------------------- +# Phase anchoring +# --------------------------------------------------------------------------- + + +async def test_phase_anchoring_gives_restart_stable_windows_and_ids(): + sweeps = FakeSweeps() + sweeps.default = _flat() + # Two processes started at different times share the window schedule. + first = _live(sweeps, start=T0 - 7.0) + second = _live(sweeps, start=T0 + 3.0) + try: + for controller in (first, second): + await _tick(controller, L_K) + result = await _tick(controller, L_K + STEP) + (decision,) = result.decisions + start_ms, end_ms = _window_ms(K) + assert decision.window_index == K + assert decision.decision_id == f"retune-w{int(end_ms)}" + assert (decision.window_start_s, decision.window_end_s) == ( + T0 - INTERVAL, + T0, + ) + first_request, second_request = sweeps.requests + assert (first_request.window_start_ms, first_request.window_end_ms) == ( + start_ms, + end_ms, + ) + assert first_request.as_of_ms == (T0 + GRACE) * 1000.0 + assert (second_request.window_start_ms, second_request.window_end_ms) == ( + start_ms, + end_ms, + ) + finally: + first.close() + second.close() + + +# --------------------------------------------------------------------------- +# Launch gate +# --------------------------------------------------------------------------- + + +async def test_gated_launch_passes_abort_and_live_request_fields(): + sweeps = FakeSweeps() + sweeps.default = _flat() + abort = threading.Event() + gate = Gate(GateResult(spec_replicas=SPEC, abort=abort)) + controller = _live(sweeps, gate=gate, generation=4) + try: + result = await _tick(controller, L_K) + finally: + controller.close() + + assert result.audit_events == [f"retune:launch w={K}"] + assert gate.calls == [(K, (T0 - INTERVAL, T0, T0 + GRACE))] + (request,) = sweeps.requests + assert sweeps.aborts == [abort] + assert sweeps.cancels == [controller._cancel] + assert request.router_domain == ROUTER_DOMAIN + # D_k = L_k + timeout (timeout defaults to the interval). + assert request.deadline_s == L_K + INTERVAL + assert request.spec_replicas == SPEC + assert request.replicas == SPEC + # Started at the window end: no recorded count covers the window start. + assert request.replicas_verified is False + assert request.generation == 4 + assert request.trace_paths == ("/trace/segments/dgd",) + + +@pytest.mark.parametrize("reason", ["not_leader", "not_ready"]) +async def test_retry_reasons_hold_the_window_until_it_may_launch(reason): + sweeps = FakeSweeps() + sweeps.default = _flat() + gate = Gate( + GateResult(reason=reason), + GateResult(reason=reason), + default=GateResult(spec_replicas=SPEC), + ) + controller = _live(sweeps, gate=gate) + try: + held = [await _tick(controller, L_K + i * STEP) for i in range(2)] + launched = await _tick(controller, L_K + 2 * STEP) + decided = await _tick(controller, L_K + 3 * STEP) + finally: + controller.close() + + for result in held: + assert result.decisions == [] and result.audit_events == [] + assert [k for k, _ in gate.calls] == [K, K, K] + assert launched.audit_events == [f"retune:launch w={K}"] + assert [d.window_index for d in decided.decisions] == [K] + assert decided.decisions[0].launched_at_s == L_K + 2 * STEP + assert len(sweeps.requests) == 1 + + +async def test_already_decided_consumes_the_window_without_a_decision(): + sweeps = FakeSweeps() + sweeps.default = _flat() + gate = Gate( + GateResult(reason="already_decided"), default=GateResult(spec_replicas=SPEC) + ) + controller = _live(sweeps, gate=gate) + try: + consumed = await _tick(controller, L_K) + quiet = await _tick(controller, L_K + STEP) + next_window = await _tick(controller, L_K + INTERVAL) + finally: + controller.close() + + assert consumed.decisions == [] and consumed.audit_events == [] + assert quiet.decisions == [] and quiet.audit_events == [] + # Asked once for K (consumed), then for K + 1 at its launch time. + assert [k for k, _ in gate.calls] == [K, K + 1] + assert next_window.audit_events == [f"retune:launch w={K + 1}"] + assert [r.window_index for r in sweeps.requests] == [K + 1] + + +async def test_gate_runs_before_the_busy_branch(): + release = threading.Event() + sweeps = FakeSweeps() + + def blocking(request): + release.wait(30) + return _flat()(request) + + sweeps.default = blocking + gate = Gate( + GateResult(spec_replicas=SPEC), + GateResult(reason="not_leader"), + default=GateResult(spec_replicas=SPEC), + ) + controller = _live(sweeps, gate=gate, clock=WallClock()) + try: + await _tick(controller, L_K) + held = await _tick(controller, L_K + INTERVAL) + busy = await _tick(controller, L_K + INTERVAL + STEP) + finally: + release.set() + controller.close() + + # A window this process may not work on is not recorded busy. + assert held.decisions == [] + assert [(d.outcome, d.window_index) for d in busy.decisions] == [("busy", K + 1)] + assert busy.decisions[0].failure_reason is None + assert len(sweeps.requests) == 1 + + +# --------------------------------------------------------------------------- +# Harvest gating and abandoned sweeps +# --------------------------------------------------------------------------- + + +async def test_refused_harvest_drops_the_result_and_relaunches_the_window(): + sweeps = FakeSweeps() + sweeps.default = _flat() + harvest = Harvest(False) + controller = _live(sweeps, harvest=harvest, generation=5) + try: + await _tick(controller, L_K) + refused = await _tick(controller, L_K + STEP) + decided = await _tick(controller, L_K + 2 * STEP) + finally: + controller.close() + + assert harvest.calls == [5, 5] + # No decision; the window is due again and relaunches in the same tick. + assert refused.decisions == [] + assert refused.audit_events == [f"retune:dropped w={K}", f"retune:launch w={K}"] + assert [r.window_index for r in sweeps.requests] == [K, K] + assert [(d.outcome, d.window_index) for d in decided.decisions] == [("kept", K)] + + +@pytest.mark.parametrize("reason", sorted(NON_DECISION_FAILURES)) +async def test_non_decision_failures_record_nothing_and_rewind(reason): + sweeps = FakeSweeps() + sweeps.results.append( + SweepResult(status="failed", error="stopped", failure_reason=reason) + ) + sweeps.default = _flat() + controller = _live(sweeps) + try: + await _tick(controller, L_K) + abandoned = await _tick(controller, L_K + STEP) + decided = await _tick(controller, L_K + 2 * STEP) + finally: + controller.close() + + assert abandoned.decisions == [] + assert abandoned.audit_events == [f"retune:dropped w={K}", f"retune:launch w={K}"] + assert [(d.outcome, d.window_index) for d in decided.decisions] == [("kept", K)] + + +async def test_abandoned_window_is_not_relaunched_once_its_successor_is_due(): + sweeps = FakeSweeps() + sweeps.results.append(SweepResult(status="failed", failure_reason="shutdown")) + sweeps.default = _flat() + controller = _live(sweeps, clock=WallClock()) + try: + await _tick(controller, L_K) + controller._pending.future.result(timeout=30) + # Harvested only after the next window is due: K + 1 launches. + result = await _tick(controller, L_K + INTERVAL) + finally: + controller.close() + + assert result.audit_events == [f"retune:dropped w={K}", f"retune:launch w={K + 1}"] + + +# --------------------------------------------------------------------------- +# drop_pending and rewind +# --------------------------------------------------------------------------- + + +async def test_drop_pending_and_rewind_relaunch_instead_of_busy(): + release = threading.Event() + sweeps = FakeSweeps() + + def blocking(request): + release.wait(30) + return _flat()(request) + + sweeps.results.append(blocking) + sweeps.default = _flat() + abort = threading.Event() + gate = Gate( + GateResult(spec_replicas=SPEC, abort=abort), + default=GateResult(spec_replicas=SPEC), + ) + controller = _live(sweeps, gate=gate, clock=WallClock()) + try: + await _tick(controller, L_K) + dropped = controller._pending.future + # A pending sweep keeps its window: rewind alone does nothing. + controller.rewind(K - 1) + assert controller._last_k == K + + controller.drop_pending() + assert abort.is_set() + assert controller._pending is None + controller.rewind(K - 1) + relaunch = await _tick(controller, L_K + STEP) + release.set() + dropped.result(timeout=30) + # The relaunch queued behind the dropped sweep; it is decided once + # its own sweep is done. + controller._pending.future.result(timeout=30) + decided = await _tick(controller, L_K + 3 * STEP) + finally: + release.set() + controller.close() + + assert relaunch.decisions == [] + assert relaunch.audit_events == [f"retune:launch w={K}"] + assert [r.window_index for r in sweeps.requests] == [K, K] + assert [(d.outcome, d.window_index) for d in decided.decisions] == [("kept", K)] + + +# --------------------------------------------------------------------------- +# DGD spec replicas +# --------------------------------------------------------------------------- + + +async def _run_window(controller, counts_at, *, start, until): + """Tick from ``start`` to ``until``; return window K's decisions. + + Earlier windows are due from the first tick too; only K is asserted. + """ + + decisions = [] + now = start + while now <= until + 1e-9: + result = await controller.on_tick(now, _counts(counts_at(now))) + decisions.extend((now, d) for d in result.decisions if d.window_index == K) + now += STEP + return decisions + + +def _requests_for_k(sweeps): + return [request for request in sweeps.requests if request.window_index == K] + + +async def test_a_ready_count_change_in_the_window_fails_it(): + sweeps = FakeSweeps() + sweeps.default = _flat() + start = T0 - INTERVAL - 2 * STEP + + def counts(now): + # One worker restarts for a tick in the middle of the window. + return 1 if now == T0 - INTERVAL / 2 else 2 + + controller = _live(sweeps, start=start) + try: + ((now, decision),) = await _run_window( + controller, counts, start=start, until=L_K + ) + finally: + controller.close() + + assert now == L_K + assert (decision.outcome, decision.failure_reason) == ("failed", "replicas_changed") + assert decision.details == { + "spec_replicas": SPEC, + "observed_replicas": {"aggregated": 1}, + } + assert decision.decision_id == f"retune-w{int(T0 * 1000)}" + assert decision.launched_at_s is None + assert _requests_for_k(sweeps) == [] + + +@pytest.mark.parametrize( + "counts, failed", + [ + # 3 ready before the window, 2 from the tick at its start: clean. + (lambda now: 3 if now < T0 - INTERVAL else 2, False), + # The tick at the window start still reports 3: that count was in + # effect when the window began, so it fails. + (lambda now: 3 if now <= T0 - INTERVAL else 2, True), + ], + ids=["changed-before-start", "changed-after-start"], +) +async def test_the_count_in_effect_at_window_start_must_match_too(counts, failed): + sweeps = FakeSweeps() + sweeps.default = _flat() + start = T0 - INTERVAL - 2 * STEP + controller = _live(sweeps, start=start) + try: + decisions = await _run_window(controller, counts, start=start, until=L_K) + finally: + controller.close() + + if failed: + assert [(d.outcome, d.failure_reason) for _, d in decisions] == [ + ("failed", "replicas_changed") + ] + assert _requests_for_k(sweeps) == [] + else: + assert decisions == [] + (request,) = _requests_for_k(sweeps) + assert request.replicas_verified is True + + +async def test_window_started_before_the_controller_launches_unverified(): + sweeps = FakeSweeps() + sweeps.default = _flat() + start = T0 - INTERVAL / 2 + controller = _live(sweeps, start=start) + try: + decisions = await _run_window( + controller, lambda _now: 2, start=start, until=L_K + ) + finally: + controller.close() + + assert decisions == [] + (request,) = _requests_for_k(sweeps) + assert request.replicas_verified is False + assert request.replicas == SPEC + + +# --------------------------------------------------------------------------- +# restore_desired +# --------------------------------------------------------------------------- + + +async def test_restore_desired_sets_the_deployed_point_but_not_startup(): + sweeps = FakeSweeps() + sweeps.default = _flat() + controller = _live(sweeps) + try: + controller.restore_desired( + {**BEST.router, "decode_active_request_weight": 4.0} # extra: ignored + ) + assert controller.current_point() == BEST + result = await _tick(controller, L_K) + finally: + controller.close() + + expected = KvRouterConfig.from_json( + json.dumps({**json.loads(KvRouterConfig().to_json()), **BEST.router}) + ).policy_parameters("aggregated") + assert result.router_parameters == {"aggregated": dict(expected)} + (request,) = sweeps.requests + # The restored point is the baseline; startup stays a candidate. + assert request.baseline == BEST + assert STARTUP in request.candidates diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_dgdr.py b/components/src/dynamo/planner/tests/unit/test_live_retune_dgdr.py new file mode 100644 index 000000000000..f6fcf3c1f8f1 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_dgdr.py @@ -0,0 +1,2118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the executor imports the sweep. + +"""Unit tests for the live retune DGDR executor (``retune.live.dgdr``). + +Each sweep runs against ``FakeRetuneKube`` with ``FakeDgdrOperator`` +advancing every Request one phase per ``sleep`` of a fake clock, so the +Request -> Run -> Candidate lifecycle, adoption, failures and timeouts are +all deterministic. +""" + +from __future__ import annotations + +import dataclasses +import gzip +import hashlib +import itertools +import json +import math +import threading +import time +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any, Optional + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from dynamo.planner.config.planner_config import ( + DgdrExecutorConfig, + RetuneConfig, + RetuneEngineSpec, + dgdr_need_seconds, +) +from dynamo.planner.core.types import RetuneParams +from dynamo.planner.retune.live import dgdr +from dynamo.planner.retune.live.contracts import ( + LABEL_ATTEMPT, + LABEL_DGD_UID, + LABEL_SHARD, + LABEL_SPEC, + LABEL_WINDOW, + NON_DECISION_FAILURES, + DgdRef, + KubeAlreadyExists, + KubeError, + KubeInvalid, + KubeNotFound, + WindowCutResult, +) +from dynamo.planner.retune.live.dgdr import ( + DgdrContext, + Shard, + build_request, + input_mismatches, + missing_fields, + plan_shards, + request_name, + run_dgdr_sweep, + spec_hash, + window_claim_path, + window_selection, +) +from dynamo.planner.retune.live.reconciler import trace_pvc_source +from dynamo.planner.retune.sweep import SweepRequest, score_summary +from dynamo.planner.tests.live_retune_fakes import ( + FakeDgdrOperator, + FakeRetuneKube, + default_metrics, + lower_request_input, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +CREDIT = "overlap_score_credit" +SCALE = "prefill_load_scale" +TEMP = "router_temperature" +LIVE_TO_PUBLIC = {CREDIT: CREDIT, SCALE: SCALE, TEMP: "temperature"} +DEFAULT_DOMAIN = { + CREDIT: (0.0, 0.5, 1.0), + SCALE: (0.0, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0), + TEMP: (0.0, 0.2, 0.5, 1.0), +} +SMALL_DOMAIN = {CREDIT: (0.5, 1.0), SCALE: (1.0, 2.0), TEMP: (0.0, 0.5)} +WINDOW_END_MS = 1_793_000_000_000 +INTERVAL_S = 3600.0 +GRACE_S = 120.0 +TIMEOUT_S = 3000.0 +LAUNCH_S = WINDOW_END_MS / 1000.0 + GRACE_S +WINDOW_START_MS = WINDOW_END_MS - int(INTERVAL_S * 1000) +AS_OF_MS = WINDOW_END_MS + int(GRACE_S * 1000) +WINDOW_FILE = ( + f"windows/llm/{WINDOW_END_MS}-" + f"{window_selection(WINDOW_START_MS, WINDOW_END_MS, AS_OF_MS)}.jsonl.gz" +) +POLL_S = 5.0 +DGD = DgdRef(name="llm", namespace="test-ns", uid="dgd-uid-1", dynamo_namespace="dyn") +AGG_ENGINE = { + "mode": "aggregated", + "backend": "vllm", + "context_length": 32768, + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + { + "tensor": 2, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, + } + ] + } + }, + "prefill": None, + "decode": None, + }, +} +DISAGG_ENGINE = { + "mode": "disaggregated", + "backend": "vllm", + "workers": { + "aggregated": None, + "prefill": { + "parallelism": { + "preset": [ + { + "tensor": 1, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, + } + ] + } + }, + "decode": { + "parallelism": { + "preset": [ + { + "tensor": 2, + "pipeline": 1, + "attention_data": 2, + "moe_tensor": 1, + "moe_expert": 1, + } + ] + } + }, + }, +} + + +# ------------------------------------------------------------------ helpers + + +def make_cfg(**overrides: Any) -> DgdrExecutorConfig: + values: dict[str, Any] = { + "claim_name": "request-traces", + "model": "Qwen/Qwen3-32B", + "sku": "h200_sxm", + "engine": AGG_ENGINE, + "cut_budget_seconds": 100.0, + "job_startup_seconds": 60.0, + "run": {"parallelism": 8, "candidate_timeout_seconds": 300.0, "max_shards": 3}, + } + values.update(overrides) + return DgdrExecutorConfig.model_validate(values) + + +def make_retune(cfg: DgdrExecutorConfig, **overrides: Any) -> RetuneConfig: + values: dict[str, Any] = { + "interval_seconds": INTERVAL_S, + "grace_seconds": GRACE_S, + "timeout_seconds": TIMEOUT_S, + "router": "default", + "executor": "dgdr", + "dgdr": cfg, + "objective": "GoodputPerGPU", + } + values.update(overrides) + return RetuneConfig.model_validate(values) + + +def grid(domain: Mapping[str, tuple[float, ...]]) -> list[dict[str, float]]: + keys = sorted(domain) + return [ + dict(zip(keys, values)) + for values in itertools.product(*(domain[key] for key in keys)) + ] + + +def make_request( + domain: Mapping[str, tuple[float, ...]] = SMALL_DOMAIN, + *, + baseline: Optional[Mapping[str, float]] = None, + startup: Optional[Mapping[str, float]] = None, + deadline_s: Optional[float] = LAUNCH_S + TIMEOUT_S, + spec_replicas: Optional[dict[str, int]] = None, + objective: str = "GoodputPerGPU", + router_domain: Any = "same", + planner: Optional[dict[str, Any]] = None, + trace_paths: tuple[str, ...] = ("/trace/segments/llm",), +) -> SweepRequest: + points = grid(domain) + baseline = dict(baseline or points[-1]) + startup = dict(startup or baseline) + candidates = {RetuneParams(router=point) for point in points} + candidates.add(RetuneParams(router=startup)) + base = RetuneParams(router=baseline, planner=planner or {}) + candidates.add(base) + replicas = spec_replicas or {"aggregated": 3} + disagg = "prefill" in replicas + return SweepRequest( + window_index=7, + window_start_ms=WINDOW_END_MS - INTERVAL_S * 1000.0, + window_end_ms=float(WINDOW_END_MS), + history_start_ms=WINDOW_END_MS - INTERVAL_S * 1000.0, + as_of_ms=WINDOW_END_MS + GRACE_S * 1000.0, + trace_paths=trace_paths, + arrival_speedup_ratio=1.0, + engine=RetuneEngineSpec( + router_mode="kv_router", + router_config={}, + prefill_engine_args={} if disagg else None, + decode_engine_args={} if disagg else None, + ), + candidates=tuple(sorted(candidates, key=RetuneParams.key)), + baseline=base, + replicas=dict(replicas), + base_planner=None, + predictor_candidates=None, + objective=objective, # type: ignore[arg-type] + sla_ttft_ms=500.0, + sla_itl_ms=30.0, + max_parallel=1, + timeout_s=TIMEOUT_S, + router_domain=( + {key: tuple(values) for key, values in domain.items()} + if router_domain == "same" + else router_domain + ), + deadline_s=deadline_s, + spec_replicas=dict(replicas), + generation=3, + ) + + +def public(point: Mapping[str, float]) -> dict[str, float]: + return {LIVE_TO_PUBLIC[key]: value for key, value in point.items()} + + +def expected_score(point: Mapping[str, float], request: SweepRequest) -> float: + score, _ = score_summary( + default_metrics(public(point)), request.window_s, request.objective + ) + return score + + +class FakeClock: + """Wall clock for the executor; ``sleep`` advances it and runs hooks.""" + + def __init__(self, start: float = LAUNCH_S) -> None: + self.now = start + self.sleeps: list[float] = [] + self.hooks: list[Callable[[], None]] = [] + + def time(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += max(seconds, 0.001) + for hook in list(self.hooks): + hook() + + +class FakeRegistry: + def __init__(self) -> None: + self.protected: dict[int, set[str]] = {} + self.released: list[int] = [] + self.references: dict[int, int] = {} + self.deleted: list[str] = [] + self.lock = threading.Lock() + + def protect(self, window_end_ms: int, request_names) -> None: + with self.lock: + self.protected.setdefault(window_end_ms, set()).update(request_names) + self.references[window_end_ms] = self.references.get(window_end_ms, 0) + 1 + + def release(self, window_end_ms: int) -> None: + with self.lock: + self.released.append(window_end_ms) + self.references[window_end_ms] = self.references.get(window_end_ms, 0) - 1 + + def balanced(self, window_end_ms: int) -> bool: + """Every reference the sweep took on the window was released.""" + return ( + window_end_ms in self.released + and self.references.get(window_end_ms, 0) == 0 + ) + + def enqueue_delete(self, request_names) -> None: + with self.lock: + self.deleted.extend(request_names) + + +class FakeCutter: + """Records cuts; returns ``num_requests`` rows or raises ``error``.""" + + def __init__( + self, + *, + num_requests: int = 1234, + error: Optional[Exception] = None, + on_cut: Optional[Callable[[], None]] = None, + ) -> None: + self.num_requests = num_requests + self.error = error + self.on_cut = on_cut + self.calls: list[dict[str, Any]] = [] + + def __call__( + self, + segments_root, + start_ms, + end_ms, + as_of_ms, + out_path, + *, + max_rows, + deadline_s, + cancel=None, + ) -> WindowCutResult: + self.calls.append( + { + "segments_root": segments_root, + "start_ms": start_ms, + "end_ms": end_ms, + "as_of_ms": as_of_ms, + "out_path": out_path, + "max_rows": max_rows, + "deadline_s": deadline_s, + "cancel": cancel, + } + ) + if self.on_cut is not None: + self.on_cut() + if self.error is not None: + raise self.error + if self.num_requests == 0: + return WindowCutResult( + path=None, num_requests=0, trace_block_size=None, row_counts={} + ) + return WindowCutResult( + path=out_path, + num_requests=self.num_requests, + trace_block_size=64, + row_counts={"selected": self.num_requests}, + dropped_members=1, + dropped_bytes=10, + ) + + +class CutError(Exception): + def __init__(self, reason: str) -> None: + super().__init__(reason) + self.reason = reason + + +class Harness: + """A fake cluster, operator, clock, registry and cutter for one window.""" + + def __init__( + self, + tmp_path, + *, + cfg: Optional[DgdrExecutorConfig] = None, + operator: bool = True, + **retune_overrides: Any, + ) -> None: + self.clock = FakeClock() + self.kube = FakeRetuneKube("test-ns", now=self.clock.time) + self.cfg = cfg or make_cfg(mount_path=str(tmp_path / "trace")) + self.retune = make_retune(self.cfg, **retune_overrides) + self.registry = FakeRegistry() + self.cutter = FakeCutter() + self.leader = True + self.cancel = threading.Event() + self.abort = threading.Event() + self.operator = FakeDgdrOperator(self.kube) + if operator: + self.clock.hooks.append(self.operator.step) + + def context(self, **overrides: Any) -> DgdrContext: + values: dict[str, Any] = { + "dgd": DGD, + "cfg": self.cfg, + "retune": self.retune, + "backend": "vllm", + "kube": self.kube, + "registry": self.registry, + "cut_window": self.cutter, + "is_leader": lambda: self.leader, + "now": self.clock.time, + "sleep": self.clock.sleep, + "poll_interval_s": POLL_S, + } + values.update(overrides) + return DgdrContext(**values) + + def sweep(self, request: SweepRequest, **overrides: Any): + return run_dgdr_sweep( + request, + cancel=self.cancel, + abort=self.abort, + context=self.context(**overrides), + ) + + def created(self) -> list[dict]: + return [ + call.target + for call in self.kube.calls_of("create_request") + if not call.kwargs.get("dry_run") + ] + + def request_of_shard(self, shard_id: str) -> dict: + found = [ + obj + for obj in self.kube.requests.values() + if obj["metadata"]["labels"][LABEL_SHARD] == shard_id + ] + assert len(found) == 1, (shard_id, [o["metadata"]["name"] for o in found]) + return found[0] + + +def assert_scores(result, request: SweepRequest, missing=()) -> None: + missing_keys = {RetuneParams(router=point).key() for point in missing} + assert set(result.scores) == {candidate.key() for candidate in request.candidates} + for candidate in request.candidates: + key = candidate.key() + if key in missing_keys: + assert result.scores[key] is None, key + assert "error" in result.metrics[key] + else: + assert result.scores[key] == pytest.approx( + expected_score(candidate.router, request), rel=1e-12 + ), key + assert result.metrics[key]["params"] == candidate.to_dict() + + +def kube_error(cls: type[KubeError], status: Optional[int], message: str) -> KubeError: + """A ``Kube*`` error as ``kube.py`` raises it, with its HTTP ``status``.""" + + error = cls(f"HTTP {status}: {message}") + error.status = status # type: ignore[attr-defined] + return error + + +def lose_leadership_after_start(h: Harness, request: SweepRequest) -> None: + """A first leader creates the window's Requests, the operator starts their + Runs, and the leader loses leadership; the next sweep is another leader.""" + + h.operator.max_phase = "started" + h.clock.hooks.append(lambda: h.abort.set() if len(h.clock.sleeps) >= 2 else None) + assert h.sweep(request).failure_reason == "lost_leadership" + h.abort = threading.Event() + h.operator.max_phase = None + h.clock.hooks = [h.operator.step] + + +# ------------------------------------------------------------------ shards + + +def test_plan_shards_default_grid_is_three_balanced_shards_and_a_probe(): + points = grid(DEFAULT_DOMAIN) + baseline = points[0] + shards = plan_shards( + DEFAULT_DOMAIN, + baseline=baseline, + startup=baseline, + max_shards=3, + parallelism=8, + ) + grid_shards = [shard for shard in shards if shard.kind == "grid"] + assert [shard.size for shard in grid_shards] == [36, 36, 36] + assert [shard.id for shard in shards] == ["grid-0", "grid-1", "grid-2", "probe"] + covered = [point for shard in grid_shards for point in shard.points] + assert sorted(map(str, covered)) == sorted(map(str, points)) + for shard in grid_shards: + # Each shard is the Cartesian product of its choices. + assert shard.size == math.prod(len(v) for v in shard.choices.values()) + probe = shards[-1] + # The middle value of each knob, a positive temperature. + assert probe.points == ({CREDIT: 0.5, SCALE: 2.0, TEMP: 0.5},) + assert probe.choices == {CREDIT: (0.5,), SCALE: (2.0,), TEMP: (0.5,)} + + +def test_plan_shards_adds_off_grid_baseline_and_startup_points(): + off_baseline = {CREDIT: 0.75, SCALE: 1.0, TEMP: 0.0} + off_startup = {CREDIT: 1.0, SCALE: 3.0, TEMP: 0.0} + shards = plan_shards( + SMALL_DOMAIN, + baseline=off_baseline, + startup=off_startup, + max_shards=3, + parallelism=8, + ) + by_id = {shard.id: shard for shard in shards} + assert set(by_id) == {"grid-0", "point-baseline", "point-startup", "probe"} + assert by_id["point-baseline"].points == (off_baseline,) + assert by_id["point-baseline"].kind == "point" + assert by_id["point-startup"].points == (off_startup,) + # The same off-grid point once. + shards = plan_shards( + SMALL_DOMAIN, + baseline=off_baseline, + startup=off_baseline, + max_shards=3, + parallelism=8, + ) + assert [shard.id for shard in shards] == ["grid-0", "point-baseline", "probe"] + # An on-grid startup needs no shard; no probe when disabled. + shards = plan_shards( + SMALL_DOMAIN, + baseline=grid(SMALL_DOMAIN)[0], + startup=grid(SMALL_DOMAIN)[1], + max_shards=3, + probe=False, + parallelism=8, + ) + assert [shard.id for shard in shards] == ["grid-0"] + + +def test_plan_shards_splits_uneven_dimension_and_rejects_impossible_grids(): + domain = {CREDIT: tuple(float(v) for v in range(250))} + zero = {CREDIT: 0.0} + shards = plan_shards( + domain, baseline=zero, startup=zero, max_shards=3, parallelism=8, probe=False + ) + assert [shard.size for shard in shards] == [84, 83, 83] + assert shards[0].choices[CREDIT][0] == 0.0 + assert shards[-1].choices[CREDIT][-1] == 249.0 + with pytest.raises(ValueError, match="cannot be split"): + plan_shards(domain, baseline=zero, startup=zero, max_shards=2, parallelism=8) + with pytest.raises(ValueError, match="cannot be split"): + plan_shards( + DEFAULT_DOMAIN, + baseline=grid(DEFAULT_DOMAIN)[0], + startup=grid(DEFAULT_DOMAIN)[0], + max_shards=1, + parallelism=8, + ) + point = grid(SMALL_DOMAIN)[0] + + def sizes(**kwargs): + shards = plan_shards( + SMALL_DOMAIN, baseline=point, startup=point, probe=False, **kwargs + ) + return [shard.size for shard in shards] + + # A shard is added only when it saves an evaluation round (AISimulate + # runs the first candidate alone, so single points save one). + assert sizes(max_shards=3, parallelism=8) == [8] + assert sizes(max_shards=16, parallelism=8) == [1] * 8 + assert sizes(max_shards=3, parallelism=1) == [4, 4] + # The default grid is three shards of 36. + default = plan_shards( + DEFAULT_DOMAIN, + baseline=grid(DEFAULT_DOMAIN)[0], + startup=grid(DEFAULT_DOMAIN)[0], + max_shards=3, + parallelism=8, + ) + assert [shard.size for shard in default] == [36, 36, 36, 1] + + +def test_plan_shards_validates_inputs_and_probe_falls_back_to_zero_temperature(): + point = {CREDIT: 0.5, SCALE: 1.0} + with pytest.raises(ValueError, match="do not match"): + plan_shards( + SMALL_DOMAIN, baseline=point, startup=point, max_shards=3, parallelism=8 + ) + with pytest.raises(ValueError, match="empty"): + plan_shards({}, baseline={}, startup={}, max_shards=3, parallelism=8) + domain = {CREDIT: (0.0, 1.0), SCALE: (1.0,), TEMP: (0.0,)} + shards = plan_shards( + domain, + baseline=grid(domain)[0], + startup=grid(domain)[0], + max_shards=3, + parallelism=8, + ) + assert shards[-1].points == ({CREDIT: 1.0, SCALE: 1.0, TEMP: 0.0},) + + +# ---------------------------------------------------------------- Request + + +def test_spec_hash_is_canonical(): + assert spec_hash({"a": 1, "b": [1, 2]}) == spec_hash({"b": [1, 2], "a": 1}) + assert spec_hash({"a": 1}) != spec_hash({"a": 2}) + assert len(spec_hash({})) == 16 + int(spec_hash({}), 16) + + +def _shards(domain=SMALL_DOMAIN, **kwargs): + point = grid(domain)[0] + return plan_shards( + domain, baseline=point, startup=point, max_shards=3, parallelism=8, **kwargs + ) + + +def test_build_request_aggregated_grid_shard(): + cfg = make_cfg() + retune = make_retune(cfg) + shard = _shards()[0] + body = build_request( + dgd=DGD, + cfg=cfg, + retune=retune, + sla_ttft_ms=500, + sla_itl_ms=30, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + shard=shard, + spec_replicas={"aggregated": 3}, + backend="vllm", + attempt=1, + ) + spec = body["spec"] + assert body["apiVersion"] == "nvidia.com/v1beta2" + assert body["kind"] == "DynamoGraphDeploymentRequest" + assert spec["model"] == {"name": "Qwen/Qwen3-32B"} + assert spec["backends"] == ["vllm"] + # 3 replicas x (tensor 2 x pipeline 1 x attention_data 1). + assert spec["hardware"] == {"gpu": {"sku": "h200_sxm", "budget": 6}} + # The trace format is left to the CRD default. + assert spec["workload"] == { + "trace": { + "source": { + "persistentVolumeClaim": { + "claimName": "request-traces", + "subPath": WINDOW_FILE, + } + }, + } + } + assert spec["objective"] == { + "mode": "Optimize", + "metric": "GoodputPerGPU", + "sla": {"ttftMs": 500.0, "itlMs": 30.0}, + } + assert spec["recommendation"] == {"maxCandidates": 8} + overrides = spec["aisimulate"]["inputOverrides"] + workers = overrides["engine"]["workers"] + assert workers["aggregated"]["parallelism"]["preset"][0]["replicas"] == 3 + assert workers["prefill"] is None and workers["decode"] is None + assert overrides["engine"]["backend"] == "vllm" + assert overrides["router"] == { + "policy": "kv_router", + "prefill_load_model": {"type": "none"}, + "overlap_score_credit": {"choices": [0.5, 1.0]}, + "prefill_load_scale": {"choices": [1.0, 2.0]}, + "temperature": {"choices": [0.0, 0.5]}, + } + assert overrides["optimizer"] == { + "algorithm": "random", + "max_trials": 16, + "parallelism": 8, + "candidate_timeout_seconds": 300.0, + } + assert overrides["optimization"] == {"strict_sla": False} + assert overrides["dgdr_run"] == {"generator": "direct"} + # The config's engine is not mutated. + assert ( + "replicas" + not in AGG_ENGINE["workers"]["aggregated"]["parallelism"]["preset"][0] + ) + + metadata = body["metadata"] + digest = spec_hash(spec) + assert metadata["labels"] == { + "nvidia.com/planner-retune-dgd": "llm", + LABEL_DGD_UID: "dgd-uid-1", + LABEL_WINDOW: str(WINDOW_END_MS), + LABEL_SHARD: "grid-0", + LABEL_SPEC: digest, + LABEL_ATTEMPT: "1", + } + assert metadata["ownerReferences"] == [ + { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "name": "llm", + "uid": "dgd-uid-1", + "blockOwnerDeletion": False, + } + ] + assert metadata["name"] == request_name(DGD, WINDOW_END_MS, "grid-0", digest, 1) + assert metadata["name"].startswith("llm-rt-") and len(metadata["name"]) == 17 + assert request_name(DGD, WINDOW_END_MS, "grid-0", digest, 2) != metadata["name"] + + +def test_build_request_point_shard_uses_scalars_and_disagg_budget(): + cfg = make_cfg(engine=DISAGG_ENGINE) + retune = make_retune(cfg, objective="TTFT") + probe = _shards()[-1] + body = build_request( + dgd=DGD, + cfg=cfg, + retune=retune, + sla_ttft_ms=200, + sla_itl_ms=20, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + shard=probe, + spec_replicas={"prefill": 2, "decode": 4}, + backend="vllm", + attempt=3, + ) + spec = body["spec"] + # 2 x 1 + 4 x (2 x 1 x 2). + assert spec["hardware"]["gpu"]["budget"] == 18 + assert spec["objective"]["metric"] == "TTFT" + overrides = spec["aisimulate"]["inputOverrides"] + workers = overrides["engine"]["workers"] + assert workers["aggregated"] is None + assert workers["prefill"]["parallelism"]["preset"][0]["replicas"] == 2 + assert workers["decode"]["parallelism"]["preset"][0]["replicas"] == 4 + assert overrides["router"]["temperature"] == 0.5 + assert overrides["router"]["overlap_score_credit"] == 1.0 + assert overrides["router"]["prefill_load_scale"] == 2.0 + assert overrides["optimizer"]["max_trials"] == 2 + assert spec["recommendation"]["maxCandidates"] == 1 + assert body["metadata"]["labels"][LABEL_ATTEMPT] == "3" + + +@pytest.mark.parametrize( + "replicas, match", + [({"prefill": 1}, "spec_replicas"), ({"aggregated": 0}, "at least 1")], +) +def test_build_request_rejects_spec_replicas_it_cannot_pin(replicas, match): + cfg = make_cfg() + with pytest.raises(ValueError, match=match): + build_request( + dgd=DGD, + cfg=cfg, + retune=make_retune(cfg), + sla_ttft_ms=1, + sla_itl_ms=1, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + shard=_shards()[0], + spec_replicas=replicas, + backend="vllm", + attempt=1, + ) + + +def test_build_request_requires_the_three_dgdr_knobs(): + cfg = make_cfg() + shard = Shard( + id="grid-0", kind="grid", choices={CREDIT: (1.0,)}, points=({CREDIT: 1.0},) + ) + with pytest.raises(ValueError, match="exactly"): + build_request( + dgd=DGD, + cfg=cfg, + retune=make_retune(cfg), + sla_ttft_ms=1, + sla_itl_ms=1, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + shard=shard, + spec_replicas={"aggregated": 1}, + backend="vllm", + attempt=1, + ) + + +def _body( + shard_index=0, *, window_start_ms=WINDOW_START_MS, as_of_ms=AS_OF_MS, **kwargs +): + cfg = make_cfg(**kwargs) + return build_request( + dgd=DGD, + cfg=cfg, + retune=make_retune(cfg), + sla_ttft_ms=500, + sla_itl_ms=30, + window_start_ms=window_start_ms, + window_end_ms=WINDOW_END_MS, + as_of_ms=as_of_ms, + shard=_shards()[shard_index], + spec_replicas={"aggregated": 3}, + backend="vllm", + attempt=1, + ) + + +def test_window_file_is_named_by_its_end_and_selection_digest(): + key = f"{WINDOW_START_MS}|{WINDOW_END_MS}|{AS_OF_MS}".encode() + digest = hashlib.sha256(key).hexdigest()[:12] + assert window_selection(WINDOW_START_MS, WINDOW_END_MS, AS_OF_MS) == digest + path = window_claim_path( + make_cfg(windows_dir="retune/windows"), + DGD, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + ) + assert path == f"retune/windows/llm/{WINDOW_END_MS}-{digest}.jsonl.gz" + + +def _identity(**selection: int) -> tuple[Optional[str], str, str]: + """The window file, spec hash and name of a grid shard's Request.""" + + body = _body(**selection) + _, path = trace_pvc_source(body["spec"]) + return path, body["metadata"]["labels"][LABEL_SPEC], body["metadata"]["name"] + + +@pytest.mark.parametrize( + "selection", + [ + {"window_start_ms": WINDOW_END_MS - 1_800_000}, + {"as_of_ms": WINDOW_END_MS + 300_000}, + ], + ids=["start", "as_of"], +) +def test_request_identity_follows_the_trace_selection(selection): + same = _identity() + assert _identity() == same + assert same[0] == WINDOW_FILE + path, digest, name = _identity(**selection) + assert path is not None and path.startswith(f"windows/llm/{WINDOW_END_MS}-") + assert path != same[0] + assert digest != same[1] + assert name != same[2] + + +def test_input_read_back_accepts_the_operator_lowering_and_finds_drift(): + body = _body() + body["metadata"]["namespace"] = "test-ns" + effective = lower_request_input(body, run_name="r-1") + assert input_mismatches(body["spec"], effective) == [] + + drifted = lower_request_input(body, run_name="r-1") + drifted["engine"]["workers"]["aggregated"]["parallelism"]["preset"][0][ + "replicas" + ] = 1 + # A role the overrides delete must stay deleted. + drifted["engine"]["workers"]["prefill"] = {} + drifted["router"]["temperature"] = {"choices": [0.0]} + drifted["evaluation"]["sla"]["e2e_ms"] = 10.0 + drifted["traffic"]["load"]["speedup"] = 2.0 + drifted["traffic"]["source"]["paths"] = ["/elsewhere/trace.jsonl.gz"] + drifted["optimization"]["strict_sla"] = True + problems = input_mismatches(body["spec"], drifted) + assert sorted(problem.split(":")[0] for problem in problems) == [ + "engine.workers.aggregated.parallelism.preset.0.replicas", + "engine.workers.prefill", + "evaluation.sla.e2e_ms", + "optimization.strict_sla", + "router.temperature.choices", + "traffic.load.speedup", + "traffic.source.paths.0", + ] + assert "engine.workers.prefill: expected unset, got {}" in problems + + stripped = lower_request_input(body, run_name="r-1") + del stripped["dgdr_run"]["generator"] + del stripped["optimization"]["constraints"] + problems = input_mismatches(body["spec"], stripped) + assert "dgdr_run.generator: expected 'direct', got unset" in problems + assert any(p.startswith("optimization.constraints:") for p in problems) + + +def test_input_read_back_allows_operator_defaults_and_follows_overrides(): + body = _body() + effective = lower_request_input(body, run_name="r-1") + # Fields the Request does not set, such as a new operator default. + effective["engine"]["kv_cache_dtype"] = "auto" + preset = effective["engine"]["workers"]["aggregated"]["parallelism"]["preset"] + preset[0]["prefill_context"] = 1 + effective["router"]["queue_threshold"] = 0.5 + effective["optimizer"]["seed"] = 0 + effective["traffic"]["source"]["compression"] = "auto" + assert input_mismatches(body["spec"], effective) == [] + + # inputOverrides win over the typed trace source. + overrides = body["spec"]["aisimulate"]["inputOverrides"] + overrides["traffic"] = {"source": {"paths": ["/data/other.jsonl"]}} + effective = lower_request_input(body, run_name="r-1") + assert effective["traffic"]["source"]["paths"] == ["/data/other.jsonl"] + assert input_mismatches(body["spec"], effective) == [] + + +# Written by the operator's TestDGDRV1beta2TraceRequestLoweringGolden +# (deploy/operator/internal/controller): request_spec is the Request this +# executor builds for _body(), run_input what dgdrRunInput lowers it to. +LOWERING_GOLDEN = Path( + "deploy/operator/internal/aisimulate/testdata/trace_request_lowering.golden.json" +) + + +def _lowering_golden() -> dict[str, Any]: + for parent in Path(__file__).resolve().parents: + candidate = parent / LOWERING_GOLDEN + if candidate.is_file(): + return json.loads(candidate.read_text(encoding="utf-8")) + raise AssertionError(f"{LOWERING_GOLDEN} not found above {Path(__file__)}") + + +def test_request_and_run_input_match_the_operator_golden(): + golden = _lowering_golden() + body = _body() + assert json.loads(json.dumps(body["spec"])) == golden["request_spec"] + run_input = golden["run_input"] + assert input_mismatches(body["spec"], run_input) == [] + # The fake operator lowers the Request like the real one. + request = {"metadata": {"namespace": "test-ns"}, "spec": golden["request_spec"]} + assert ( + lower_request_input( + request, + run_name=run_input["dgdr_run"]["name"], + runtime_image=run_input["dgdr_run"]["runtime_image"], + dynamo_version=run_input["dgdr_run"]["runtime_version_override"], + gpus_per_node=run_input["dgdr_run"]["num_gpus_per_node"], + ) + == run_input + ) + + restored = json.loads(json.dumps(run_input)) + restored["engine"]["workers"]["prefill"] = {} + assert input_mismatches(body["spec"], restored) == [ + "engine.workers.prefill: expected unset, got {}" + ] + moved = json.loads(json.dumps(run_input)) + moved["traffic"]["source"]["paths"] = ["/var/run/aisimulate/trace/trace.jsonl"] + assert input_mismatches(body["spec"], moved) == [ + "traffic.source.paths.0: expected '/var/run/aisimulate/trace/trace.jsonl.gz'," + " got '/var/run/aisimulate/trace/trace.jsonl'" + ] + # An uncompressed window is mounted as trace.jsonl. + uncompressed = json.loads(json.dumps(body["spec"])) + claim = uncompressed["workload"]["trace"]["source"]["persistentVolumeClaim"] + claim["subPath"] = claim["subPath"].removesuffix(".gz") + assert input_mismatches(uncompressed, run_input) == [ + "traffic.source.paths.0: expected '/var/run/aisimulate/trace/trace.jsonl'," + " got '/var/run/aisimulate/trace/trace.jsonl.gz'" + ] + + +def test_missing_fields_reports_pruned_leaves_and_allows_defaults(): + submitted = {"a": {"b": 1, "c": [1, 2], "gone": None}, "d": True} + stored = {"a": {"b": 1.0, "c": [1, 2], "extra": 5}, "d": True, "default": 1} + assert missing_fields(stored, submitted) == [] + assert missing_fields({"a": {"b": 1}}, submitted, "spec") == ["spec.a.c", "spec.d"] + assert missing_fields( + {"a": {"b": 1, "c": [1, 2], "gone": 3}, "d": 1}, submitted + ) == [ + "a.gone", + "d", + ] + # The mappings of a list may hold defaults too; the list keeps its length. + assert missing_fields({"p": [{"a": 1, "b": 2}]}, {"p": [{"a": 1}]}) == [] + assert missing_fields({"p": [{"a": 2}]}, {"p": [{"a": 1}]}) == ["p.0.a"] + assert missing_fields({"p": [{"a": 1}, {"a": 2}]}, {"p": [{"a": 1}]}) == ["p"] + + +# ------------------------------------------------------------ the sweep + + +def test_fresh_window_creates_cuts_harvests_and_scores(tmp_path): + h = Harness(tmp_path) + request = make_request() + result = h.sweep(request) + + assert result.status == "ok", (result.error, result.failure_reason) + assert result.failure_reason is None + assert result.num_requests == 1234 + assert result.row_counts == {"selected": 1234} + assert_scores(result, request) + # One grid shard (8 points) and the probe. + created = h.created() + assert [obj["metadata"]["labels"][LABEL_SHARD] for obj in created] == [ + "grid-0", + "probe", + ] + assert all( + call.kwargs["timeout_s"] <= 10.0 for call in h.kube.calls_of("create_request") + ) + # The window was cut once, into the claim path the Requests name. + (cut,) = h.cutter.calls + assert cut["out_path"] == str(tmp_path / "trace" / WINDOW_FILE) + assert cut["segments_root"] == "/trace/segments/llm" + assert cut["start_ms"] == WINDOW_END_MS - 3_600_000 + assert cut["end_ms"] == WINDOW_END_MS + assert cut["as_of_ms"] == WINDOW_END_MS + 120_000 + # Absolute wall-clock seconds: now plus the cut budget. + assert cut["deadline_s"] == LAUNCH_S + 100.0 + assert cut["max_rows"] == h.cfg.max_window_rows + assert (tmp_path / "trace" / "windows" / "llm").is_dir() + # Requests were protected, nothing deleted, the window released. + assert h.registry.protected[WINDOW_END_MS] == { + o["metadata"]["name"] for o in created + } + assert h.registry.deleted == [] + assert h.registry.balanced(WINDOW_END_MS) + details = result.details + assert details["crn_probe"]["relative_error"] == 0.0 + assert details["runtime_image"].endswith(":1.8.0") + assert details["runtime_version"] == "1.8.0" + assert details["missing_candidates"] == 0 + assert details["shards"]["grid-0"]["critical"] is True + assert details["window"]["dropped_members"] == 1 + assert details["generation"] == 3 + assert result.wall_s > 0 + + +def test_default_grid_with_off_grid_startup_runs_five_requests(tmp_path): + h = Harness(tmp_path) + startup = {CREDIT: 1.0, SCALE: 1.5, TEMP: 0.0} + request = make_request(DEFAULT_DOMAIN, baseline=startup, startup=startup) + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + shard_ids = sorted(obj["metadata"]["labels"][LABEL_SHARD] for obj in h.created()) + assert shard_ids == ["grid-0", "grid-1", "grid-2", "point-baseline", "probe"] + # The probe point's grid twin is critical; the baseline shard is too. + critical = {k for k, v in result.details["shards"].items() if v["critical"]} + assert critical == {"grid-1", "point-baseline", "probe"} + + +def test_disaggregated_window(tmp_path): + h = Harness(tmp_path, cfg=make_cfg(engine=DISAGG_ENGINE, mount_path=str(tmp_path))) + request = make_request(spec_replicas={"prefill": 1, "decode": 2}) + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + spec = h.created()[0]["spec"] + assert spec["hardware"]["gpu"]["budget"] == 1 + 2 * 4 + + +def test_window_adopted_as_succeeded_is_not_cut_again(tmp_path): + h = Harness(tmp_path) + request = make_request() + first = h.sweep(request) + assert first.status == "ok" + creates = len(h.kube.calls_of("create_request")) + + # A new leader (fresh cutter) sweeps the same window. + h.cutter = FakeCutter(num_requests=1) + second = h.sweep(request) + assert second.status == "ok", second.error + assert h.cutter.calls == [] + assert len(h.kube.calls_of("create_request")) == creates + assert second.scores == first.scores + # Row counts come from the Request annotation of the cut. + assert second.num_requests == 1234 + assert second.row_counts == {"selected": 1234} + assert all(shard["adopted"] for shard in second.details["shards"].values()) + + +@pytest.mark.parametrize( + "change", + [ + # Restarted with half the interval: the window ends at the same time + # but starts later. + { + "window_index": 15, + "window_start_ms": WINDOW_END_MS - INTERVAL_S * 500.0, + "history_start_ms": WINDOW_END_MS - INTERVAL_S * 500.0, + }, + # Restarted with a longer grace: the same window, a later cutoff. + {"as_of_ms": WINDOW_END_MS + 2 * GRACE_S * 1000.0}, + ], + ids=["interval", "grace"], +) +def test_changed_selection_at_a_shared_end_is_not_adopted(tmp_path, change): + h = Harness(tmp_path) + first = h.sweep(make_request()) + assert first.status == "ok", first.error + old = dict(h.kube.requests) + (first_cut,) = h.cutter.calls + + h.cutter = FakeCutter() + request = dataclasses.replace(make_request(), **change) + second = h.sweep(request) + + assert second.status == "ok", second.error + assert len(h.cutter.calls) == 1, "the new selection was not cut" + (cut,) = h.cutter.calls + assert (cut["start_ms"], cut["end_ms"], cut["as_of_ms"]) == ( + round(request.window_start_ms), + WINDOW_END_MS, + round(request.as_of_ms), + ) + assert cut["out_path"] != first_cut["out_path"] + assert not any(shard["adopted"] for shard in second.details["shards"].values()) + assert sorted(h.registry.deleted) == sorted(old) + new = [obj for name, obj in h.kube.requests.items() if name not in old] + assert len(new) == len(old) + new_files = {trace_pvc_source(obj["spec"])[1] for obj in new} + old_files = {trace_pvc_source(obj["spec"])[1] for obj in old.values()} + assert {str(tmp_path / "trace" / path) for path in new_files} == {cut["out_path"]} + assert new_files.isdisjoint(old_files) + new_specs = {obj["metadata"]["labels"][LABEL_SPEC] for obj in new} + old_specs = {obj["metadata"]["labels"][LABEL_SPEC] for obj in old.values()} + assert new_specs.isdisjoint(old_specs) + assert_scores(second, request) + + +def test_lost_leadership_leaves_requests_for_the_next_leader(tmp_path): + h = Harness(tmp_path) + h.operator.max_phase = "started" + h.clock.hooks.append(lambda: h.abort.set() if len(h.clock.sleeps) >= 2 else None) + request = make_request() + result = h.sweep(request) + assert result.status == "failed" + assert result.failure_reason == "lost_leadership" + assert result.failure_reason in NON_DECISION_FAILURES + assert h.registry.deleted == [] + assert h.registry.balanced(WINDOW_END_MS) + names = set(h.kube.requests) + assert len(names) == 2 + + # The next leader adopts the in-flight Requests without cutting. + h.abort = threading.Event() + h.operator.max_phase = None + h.clock.hooks = [h.operator.step] + h.cutter = FakeCutter() + second = h.sweep(request) + assert second.status == "ok", second.error + assert set(h.kube.requests) == names + assert h.cutter.calls == [] + assert_scores(second, request) + + +def test_shutdown_returns_at_once_without_deleting(tmp_path): + h = Harness(tmp_path) + h.operator.max_phase = "accepted" + h.clock.hooks.append(h.cancel.set) + result = h.sweep(make_request()) + assert result.failure_reason == "shutdown" + assert result.status == "failed" + assert h.registry.deleted == [] + assert len(h.kube.requests) == 2 + + +def test_shutdown_ends_the_wait_between_polls(tmp_path): + # Wall-clock time and the sweep's own wait: the poll interval is far + # longer than the test, so only shutdown can end the wait. + h = Harness(tmp_path, operator=False) + request = make_request(deadline_s=time.time() + TIMEOUT_S) + timer = threading.Timer(0.2, h.cancel.set) + timer.start() + started = time.monotonic() + try: + result = h.sweep(request, now=time.time, sleep=None, poll_interval_s=600.0) + finally: + timer.cancel() + assert result.failure_reason == "shutdown" + assert time.monotonic() - started < 60.0 + assert len(h.kube.requests) == 2 + + +def test_not_leader_creates_nothing(tmp_path): + h = Harness(tmp_path) + h.leader = False + result = h.sweep(make_request()) + assert result.failure_reason == "lost_leadership" + assert h.created() == [] + assert h.cutter.calls == [] + + +def test_abort_during_cut_reports_lost_leadership(tmp_path): + h = Harness(tmp_path) + h.cutter = FakeCutter(on_cut=h.abort.set, error=CutError("cancelled")) + result = h.sweep(make_request()) + assert result.failure_reason == "lost_leadership" + assert h.created() == [] + # The cutter's cancel event reflects abort. + assert h.cutter.calls[0]["cancel"].is_set() + + +def test_stale_spec_and_older_attempts_are_queued_for_deletion(tmp_path): + h = Harness(tmp_path) + request = make_request() + h.sweep(request) + grid_request = h.request_of_shard("grid-0") + name = grid_request["metadata"]["name"] + + # A stale shard spec and an older attempt of the same spec. + stale = { + "metadata": { + "name": "llm-rt-stale", + "labels": { + **grid_request["metadata"]["labels"], + LABEL_SPEC: "0" * 16, + }, + }, + "spec": {}, + } + older = { + "metadata": { + "name": "llm-rt-older", + "labels": {**grid_request["metadata"]["labels"], LABEL_ATTEMPT: "0"}, + }, + "spec": {}, + } + unknown_shard = { + "metadata": { + "name": "llm-rt-unknown", + "labels": {**grid_request["metadata"]["labels"], LABEL_SHARD: "grid-9"}, + }, + "spec": {}, + } + for obj in (stale, older, unknown_shard): + h.kube.create_request(obj) + other_window = { + "metadata": { + "name": "llm-rt-other", + "labels": {**grid_request["metadata"]["labels"], LABEL_WINDOW: "1"}, + }, + "spec": {}, + } + h.kube.create_request(other_window) + h.registry.deleted.clear() + result = h.sweep(request) + assert result.status == "ok", result.error + assert sorted(h.registry.deleted) == [ + "llm-rt-older", + "llm-rt-stale", + "llm-rt-unknown", + ] + assert result.details["shards"]["grid-0"]["request"] == name + + +def test_create_that_lands_but_times_out_is_adopted(tmp_path): + h = Harness(tmp_path) + h.kube.fail("create_request", KubeError("timeout"), landed=True) + request = make_request() + result = h.sweep(request) + assert result.status == "ok", result.error + assert len(h.kube.requests) == 2 + assert result.details["k8s_errors"] == 1 + assert_scores(result, request) + + +def test_create_already_exists_with_same_spec_is_adopted(tmp_path): + h = Harness(tmp_path) + h.kube.fail("create_request", KubeAlreadyExists("exists"), landed=True) + result = h.sweep(make_request()) + assert result.status == "ok", result.error + assert len(h.kube.requests) == 2 + + +def test_create_that_fails_is_retried(tmp_path): + h = Harness(tmp_path) + h.kube.fail("create_request", KubeError("503"), times=2) + result = h.sweep(make_request()) + assert result.status == "ok", result.error + assert len(h.kube.requests) == 2 + assert result.details["k8s_errors"] == 2 + + +@pytest.mark.parametrize( + "error", + [ + kube_error(KubeError, 403, "requests is forbidden: cannot create"), + kube_error(KubeNotFound, 404, "the server could not find the resource"), + kube_error(KubeError, 405, "method not allowed"), + KubeNotFound("no such API"), + ], + ids=["403", "404", "405", "not-found"], +) +def test_permanent_create_error_fails_fast(tmp_path, error): + h = Harness(tmp_path) + h.kube.fail("create_request", error, times=1000) + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + assert result.status == "failed" + assert result.failure_reason == "k8s_error" + assert str(error) in result.error + # No create -> unknown -> GET -> create loop until the launch gate trips. + assert len(h.created()) == 1 + assert h.clock.sleeps == [] + assert result.details["k8s_errors"] == 1 + assert h.kube.requests == {} + + +@pytest.mark.parametrize( + "cls, status", [(KubeInvalid, 422), (KubeInvalid, 400), (KubeError, 403)] +) +def test_admission_webhook_rejection_is_request_rejected(tmp_path, cls, status): + h = Harness(tmp_path) + message = ( + 'admission webhook "vdynamographdeploymentrequest-v1beta2.kb.io" denied ' + 'the request: DynamoGraphDeploymentRequest.nvidia.com "x" is invalid: ' + 'spec.hardware.gpu.sku: Invalid value: "h200_sxm": no node offers it' + ) + h.kube.fail("create_request", kube_error(cls, status, message)) + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "request_rejected" + assert "no node offers it" in result.error + assert len(h.created()) == 1 + + +def test_webhook_unreachable_is_retried(tmp_path): + h = Harness(tmp_path) + error = kube_error( + KubeError, + 500, + 'Internal error occurred: failed calling webhook "x.kb.io": connection refused', + ) + h.kube.fail("create_request", error, times=2) + request = make_request() + result = h.sweep(request) + assert result.status == "ok", result.error + assert result.details["k8s_errors"] == 2 + assert_scores(result, request) + + +@pytest.mark.parametrize( + "error", + [kube_error(KubeError, 403, "requests is forbidden: cannot get"), KubeNotFound()], + ids=["403", "not-found"], +) +def test_permanent_get_error_fails_fast(tmp_path, error): + h = Harness(tmp_path) + # The create's outcome is unknown, so the next poll looks the Request up. + h.kube.fail("create_request", KubeError("timeout")) + h.kube.fail("get_request", error, times=1000) + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + assert result.status == "failed" + assert result.failure_reason == "k8s_error" + assert "get Request" in result.error + assert h.clock.now <= LAUNCH_S + POLL_S + + +def test_transient_get_errors_are_retried(tmp_path): + h = Harness(tmp_path) + h.kube.fail("get_request", KubeError("timeout"), times=3) + h.kube.fail("get_run", kube_error(KubeError, 503, "unavailable"), times=2) + request = make_request() + result = h.sweep(request) + assert result.status == "ok", result.error + assert result.details["k8s_errors"] == 5 + assert_scores(result, request) + + +def test_name_taken_by_a_foreign_request_fails_a_critical_shard(tmp_path): + h = Harness(tmp_path) + request = make_request() + # Pre-create a foreign object under the grid shard's name. + body = build_request( + dgd=DGD, + cfg=h.cfg, + retune=h.retune, + sla_ttft_ms=500, + sla_itl_ms=30, + window_start_ms=WINDOW_START_MS, + window_end_ms=WINDOW_END_MS, + as_of_ms=AS_OF_MS, + shard=_shards()[0], + spec_replicas={"aggregated": 3}, + backend="vllm", + attempt=1, + ) + h.kube.requests[body["metadata"]["name"]] = { + "metadata": {"name": body["metadata"]["name"], "labels": {"foreign": "yes"}}, + "spec": {}, + } + result = h.sweep(request) + assert result.status == "failed" + assert result.failure_reason == "request_conflict" + assert body["metadata"]["name"] not in h.registry.deleted + + +def test_strict_rejection_is_schema_unsupported(tmp_path): + h = Harness(tmp_path) + h.kube.unknown_fields = {"spec.workload.trace"} + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "schema_unsupported" + assert h.kube.requests == {} + + +def test_silently_pruned_field_is_schema_unsupported(tmp_path): + h = Harness(tmp_path) + h.kube.pruned_fields = {"spec.workload.trace.source"} + result = h.sweep(make_request()) + assert result.failure_reason == "schema_unsupported" + assert "spec.workload.trace.source" in result.error + # The pruned Request is queued for deletion. + assert h.registry.deleted == list(h.kube.requests) + + +def test_invalid_spec_fails_the_window_and_abandons_in_flight(tmp_path): + h = Harness(tmp_path, operator=False) + + def reject(kube, created): + if created["metadata"]["labels"][LABEL_SHARD] == "probe": + kube.set_request_condition( + created["metadata"]["name"], + status="False", + reason="InvalidSpec", + message="bad engine", + ) + + h.kube.on_request_created.append(reject) + result = h.sweep(make_request()) + assert result.failure_reason == "invalid_spec" + assert "bad engine" in result.error + assert sorted(h.registry.deleted) == sorted(h.kube.requests) + + +def test_operator_not_reconciling_after_sixty_seconds(tmp_path): + h = Harness(tmp_path, operator=False) + result = h.sweep(make_request()) + assert result.failure_reason == "operator_not_reconciling" + assert h.clock.now - LAUNCH_S >= dgdr.ACCEPT_TIMEOUT_S + assert h.clock.now - LAUNCH_S < dgdr.ACCEPT_TIMEOUT_S + 2 * POLL_S + + +def test_run_name_conflict_keeps_waiting(tmp_path): + h = Harness(tmp_path, operator=False) + phases = {"steps": 0} + + def operator(): + phases["steps"] += 1 + for name in list(h.kube.requests): + if phases["steps"] < 30: + h.kube.set_request_condition( + name, status="False", reason="RunNameConflict" + ) + if phases["steps"] >= 30: + h.operator.step() + + h.clock.hooks.append(operator) + request = make_request() + result = h.sweep(request) + assert result.status == "ok", result.error + assert h.clock.now - LAUNCH_S > dgdr.ACCEPT_TIMEOUT_S + assert_scores(result, request) + + +def test_a_run_name_conflict_until_the_deadline_reports_why(tmp_path): + # A recreated Request whose old Run still waits for garbage collection. + h = Harness(tmp_path, operator=False) + message = ( + 'Run "llm-rt-0123456789-1" exists and is not the Run of generation 1 ' + "of this Request; trying again once it is gone" + ) + + def operator(): + for name in list(h.kube.requests): + h.kube.set_request_condition( + name, status="False", reason="RunNameConflict", message=message + ) + + h.clock.hooks.append(operator) + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + assert result.status == "timeout" + assert result.failure_reason == "timeout" + for shard_id in ("grid-0", "probe"): + shard = result.details["shards"][shard_id] + assert shard["state"] == "wait" + assert shard["message"] == ( + f"Request {shard['request']} RunNameConflict: {message}" + ) + assert f"{shard_id}: {shard['message']}" in result.error + + +def test_run_pending_until_the_deadline_times_out(tmp_path): + h = Harness(tmp_path) + # The operator never creates the Jobs. Like the real operator, it then + # writes no Run status at all: the input arrives together with jobRef. + h.operator.max_phase = "accepted" + request = make_request(deadline_s=LAUNCH_S + 900.0) + result = h.sweep(request) + assert [h.kube.run_of(name)["status"] for name in h.kube.requests] == [{}, {}] + assert result.status == "timeout" + assert result.failure_reason == "run_pending" + assert "Runs pending" in result.error + assert h.clock.now >= LAUNCH_S + 900.0 + assert sorted(h.registry.deleted) == sorted(h.kube.requests) + assert all(shard["run_pending"] for shard in result.details["shards"].values()) + + +def test_run_pending_clears_once_the_job_exists(tmp_path, caplog): + h = Harness(tmp_path) + h.operator.max_phase = "accepted" + + def create_jobs_late(): + if h.clock.now - LAUNCH_S >= 3 * h.cfg.job_startup_seconds: + h.operator.max_phase = None + + h.clock.hooks.insert(0, create_jobs_late) + request = make_request() + with caplog.at_level("WARNING", logger=dgdr.__name__): + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + assert "(run_pending)" in caplog.text + assert not any(shard["run_pending"] for shard in result.details["shards"].values()) + + +def test_job_created_but_not_started_is_not_run_pending(tmp_path): + h = Harness(tmp_path, operator=False) + + def create_jobs(): + for name in list(h.kube.requests): + if h.kube.run_of(name) is None: + h.kube.create_job(h.kube.accept_request(name)["metadata"]["name"]) + + h.clock.hooks.append(create_jobs) + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + for name in h.kube.requests: + status = h.kube.run_of(name)["status"] + assert status["aisimulate"]["input"] and status["jobRef"] + assert "startTime" not in status + assert result.status == "timeout" + assert result.failure_reason == "timeout" + assert not any(shard["run_pending"] for shard in result.details["shards"].values()) + + +@pytest.mark.parametrize( + "job, reason, message, failure", + [ + (True, "ContainerWaiting", "sweeper: ErrImagePull", "timeout"), + ( + False, + "TraceSourceMissing", + 'trace PersistentVolumeClaim "request-traces" does not exist', + "run_pending", + ), + ], +) +def test_a_waiting_run_reports_why_at_the_deadline( + tmp_path, job, reason, message, failure +): + h = Harness(tmp_path, operator=False) + + def operator(): + for name in list(h.kube.requests): + if h.kube.run_of(name) is None: + run = h.kube.accept_request(name)["metadata"]["name"] + if job: + h.kube.start_run(run) + h.kube.set_run_completed( + run, status="False", reason=reason, message=message + ) + + h.clock.hooks.append(operator) + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + assert result.status == "timeout" + assert result.failure_reason == failure + for shard_id in ("grid-0", "probe"): + shard = result.details["shards"][shard_id] + assert shard["state"] == "wait" + assert shard["message"].endswith(f"{reason}: {message}") + assert f"{shard_id}: {shard['message']}" in result.error + + +def test_running_job_past_the_deadline_times_out(tmp_path): + h = Harness(tmp_path) + h.operator.max_phase = "started" + result = h.sweep(make_request(deadline_s=LAUNCH_S + 900.0)) + assert result.status == "timeout" + assert result.failure_reason == "timeout" + assert sorted(h.registry.deleted) == sorted(h.kube.requests) + + +def test_launch_too_late_for_a_critical_shard(tmp_path): + h = Harness(tmp_path) + # The 8-point grid shard needs 100 + 60 + 2 * 300 = 760 s. + result = h.sweep(make_request(deadline_s=LAUNCH_S + 700.0)) + assert result.status == "timeout" + assert result.failure_reason == "launch_too_late" + assert h.cutter.calls == [] + assert h.created() == [] + + +def test_launch_too_late_for_an_ordinary_shard_leaves_its_points_missing(tmp_path): + cfg = make_cfg( + mount_path=str(tmp_path), + run={"parallelism": 1, "candidate_timeout_seconds": 300.0, "max_shards": 3}, + ) + h = Harness(tmp_path, cfg=cfg) + domain = { + CREDIT: tuple(float(v) for v in range(10)), + SCALE: (1.0,), + TEMP: (0.0, 0.5), + } + baseline = {CREDIT: 2.5, SCALE: 1.0, TEMP: 0.0} # off the grid + shards = plan_shards( + domain, baseline=baseline, startup=baseline, max_shards=3, parallelism=1 + ) + sizes = {shard.id: shard.size for shard in shards} + # Ten credits in chunks of 4, 3 and 3; the probe (credit 5) is in grid-1. + assert sizes == { + "grid-0": 8, + "grid-1": 6, + "grid-2": 6, + "point-baseline": 1, + "probe": 1, + } + assert shards[-1].points[0][CREDIT] == 5.0 + need6 = dgdr_need_seconds(shards[1].size, cfg, needs_cut=True) + need8 = dgdr_need_seconds(shards[0].size, cfg, needs_cut=True) + assert need6 < need8 + request = make_request(domain, baseline=baseline, deadline_s=LAUNCH_S + need6 + 1) + result = h.sweep(request) + assert result.status == "ok", result.error + assert result.details["shards"]["grid-0"]["reason"] == "launch_too_late" + assert result.details["shards"]["grid-0"]["request"] is None + assert result.details["shards"]["grid-2"]["state"] == "done" + assert_scores(result, request, missing=shards[0].points) + assert ( + "launch_too_late" + in result.metrics[RetuneParams(router=shards[0].points[0]).key()]["error"] + ) + + +def test_failed_run_of_an_ordinary_shard_leaves_its_points_missing(tmp_path): + h = Harness(tmp_path) + startup = {CREDIT: 1.0, SCALE: 1.5, TEMP: 0.0} + request = make_request(baseline=grid(SMALL_DOMAIN)[0], startup=startup) + + def fail_startup(kube, created): + if created["metadata"]["labels"][LABEL_SHARD] == "point-startup": + h.operator.fail_runs[created["metadata"]["name"]] = "SearchFailed" + + h.kube.on_request_created.append(fail_startup) + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request, missing=[startup]) + shard = result.details["shards"]["point-startup"] + assert shard["reason"] == "run_failed" + assert "SearchFailed" in shard["message"] + # A terminal Request is kept for audit. + assert h.registry.deleted == [] + + +def test_failed_run_of_the_baseline_shard_fails_the_window(tmp_path): + h = Harness(tmp_path) + + def fail_grid(kube, created): + if created["metadata"]["labels"][LABEL_SHARD] == "grid-0": + h.operator.fail_runs[created["metadata"]["name"]] = "JobFailed" + + h.kube.on_request_created.append(fail_grid) + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "run_failed" + assert "JobFailed" in result.error + probe = h.request_of_shard("probe")["metadata"]["name"] + grid0 = h.request_of_shard("grid-0")["metadata"]["name"] + # The probe may already be done (kept) or in flight (deleted); the failed + # baseline shard is kept for audit. + assert grid0 not in h.registry.deleted + assert set(h.registry.deleted) <= {probe} + + +def test_input_mismatch_of_a_created_critical_shard_fails(tmp_path): + h = Harness(tmp_path) + + def skew(name, effective): + effective["engine"]["context_length"] = 1024 + + h.operator.mutate_input = skew + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "input_mismatch" + assert "engine" in result.error + assert set(h.registry.deleted) == set(h.kube.requests) + + +def test_adopted_attempt_with_input_mismatch_is_replaced(tmp_path): + h = Harness(tmp_path) + request = make_request() + skewed: set[str] = set() + + def skew(name, effective): + if not skewed or name in skewed: + skewed.add(name) + effective["router"]["prefill_load_model"] = {"type": "ais"} + + h.operator.mutate_input = skew + # The first leader creates the Requests and loses leadership before the + # operator accepts them. + h.clock.hooks = [lambda: h.abort.set() if len(h.clock.sleeps) >= 2 else None] + assert h.sweep(request).failure_reason == "lost_leadership" + + h.abort = threading.Event() + h.clock.hooks = [h.operator.step] + result = h.sweep(request) + (bad,) = skewed + assert result.status == "ok", result.error + assert_scores(result, request) + assert bad in h.registry.deleted + replaced = [s for s in result.details["shards"].values() if s["attempt"] == 2] + assert len(replaced) == 1 + # The replacement needed the window file, so the window was cut. + assert len(h.cutter.calls) == 2 + + +def test_adopted_attempt_whose_run_failed_is_replaced(tmp_path): + h = Harness(tmp_path) + request = make_request() + lose_leadership_after_start(h, request) + old = h.request_of_shard("grid-0")["metadata"]["name"] + probe = h.request_of_shard("probe")["metadata"]["name"] + # While no leader runs, the baseline shard's Job fails. + h.kube.complete_run( + h.kube.run_of(old)["metadata"]["name"], reason="JobFailed", message="evicted" + ) + creates = len(h.created()) + + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + shard = result.details["shards"]["grid-0"] + assert shard["attempt"] == 2 + assert not shard["adopted"] + assert shard["request"] not in (old, None) + assert "JobFailed" in shard["message"] + assert result.details["shards"]["probe"]["request"] == probe + assert h.registry.deleted == [old] + assert len(h.created()) == creates + 1 + # The new attempt needed the window file, so the window was cut again. + assert len(h.cutter.calls) == 2 + + +def test_adopted_failed_run_without_time_for_a_new_attempt_fails(tmp_path): + h = Harness(tmp_path) + # The grid shard needs 100 + 60 + 2 * 300 = 760 s with a cut. + request = make_request(deadline_s=LAUNCH_S + 800.0) + lose_leadership_after_start(h, request) + old = h.request_of_shard("grid-0")["metadata"]["name"] + h.kube.complete_run( + h.kube.run_of(old)["metadata"]["name"], reason="JobFailed", message="evicted" + ) + h.clock.now = LAUNCH_S + 100.0 # 700 s are left + creates = len(h.created()) + + result = h.sweep(request) + assert result.status == "failed" + assert result.failure_reason == "run_failed" + assert "JobFailed" in result.error + assert len(h.created()) == creates + assert result.details["shards"]["grid-0"]["attempt"] == 1 + # The failed attempt is kept for audit. + assert old not in h.registry.deleted + + +def test_adopted_failed_run_of_an_ordinary_shard_is_replaced(tmp_path): + h = Harness(tmp_path) + startup = {CREDIT: 1.0, SCALE: 1.5, TEMP: 0.0} + request = make_request(baseline=grid(SMALL_DOMAIN)[0], startup=startup) + lose_leadership_after_start(h, request) + old = h.request_of_shard("point-startup")["metadata"]["name"] + h.kube.complete_run( + h.kube.run_of(old)["metadata"]["name"], + reason="TraceSourceUnavailable", + message="claim is gone", + ) + + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + assert result.details["shards"]["point-startup"]["attempt"] == 2 + assert h.registry.deleted == [old] + + +def test_failed_run_of_a_request_this_sweep_created_is_not_replaced(tmp_path): + h = Harness(tmp_path) + + def grid0(body) -> bool: + return body["metadata"]["labels"][LABEL_SHARD] == "grid-0" + + # The create lands but reports an error, so the sweep adopts its own + # Request by name; its failed Run is not retried. + h.kube.fail("create_request", KubeError("timeout"), landed=True, match=grid0) + + def fail_grid(kube, created): + if grid0(created): + h.operator.fail_runs[created["metadata"]["name"]] = "JobFailed" + + h.kube.on_request_created.append(fail_grid) + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "run_failed" + shard = result.details["shards"]["grid-0"] + assert shard["adopted"] and shard["attempt"] == 1 + assert len([b for b in h.created() if grid0(b)]) == 1 + + +def test_input_mismatch_of_an_ordinary_shard_leaves_points_missing(tmp_path): + h = Harness(tmp_path) + startup = {CREDIT: 1.0, SCALE: 1.5, TEMP: 0.0} + request = make_request(baseline=grid(SMALL_DOMAIN)[0], startup=startup) + + def skew(name, effective): + if h.kube.requests[name]["metadata"]["labels"][LABEL_SHARD] == "point-startup": + effective["optimizer"]["parallelism"] = 1 + + h.operator.mutate_input = skew + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request, missing=[startup]) + assert result.details["shards"]["point-startup"]["reason"] == "input_mismatch" + + +def test_missing_run_input_on_a_succeeded_run_is_a_mismatch(tmp_path): + h = Harness(tmp_path, operator=False) + + def operator(): + for name in list(h.kube.requests): + if h.kube.run_of(name) is None: + h.kube.accept_request(name) + run = h.kube.run_of(name)["metadata"]["name"] + h.kube.start_run(run, record_input=False) + h.kube.complete_run(run) + + h.clock.hooks.append(operator) + result = h.sweep(make_request()) + assert result.failure_reason == "input_mismatch" + assert "status.aisimulate.input is missing" in result.error + + +def test_crn_probe_mismatch_fails_the_window(tmp_path): + def evaluate(request, point): + metrics = default_metrics(point) + if request["metadata"]["labels"][LABEL_SHARD] == "probe": + metrics["goodput_output_throughput_tok_s"] *= 1.0 + 1e-6 + return metrics + + h = Harness(tmp_path) + h.operator.evaluate = evaluate + result = h.sweep(make_request()) + assert result.failure_reason == "crn_mismatch" + assert result.details["crn_probe"]["relative_error"] == pytest.approx( + 1e-6, rel=1e-3 + ) + + +def test_crn_probe_within_tolerance_passes(tmp_path): + def evaluate(request, point): + metrics = default_metrics(point) + if request["metadata"]["labels"][LABEL_SHARD] == "probe": + metrics["goodput_output_throughput_tok_s"] *= 1.0 + 1e-13 + return metrics + + h = Harness(tmp_path) + h.operator.evaluate = evaluate + result = h.sweep(make_request()) + assert result.status == "ok", result.error + + +def test_unscored_probe_is_a_crn_mismatch(tmp_path): + probe_point = {CREDIT: 1.0, SCALE: 2.0, TEMP: 0.5} + + def evaluate(request, point): + if request["metadata"]["labels"][LABEL_SHARD] == "probe": + return None + return default_metrics(point) + + h = Harness(tmp_path) + h.operator.evaluate = evaluate + result = h.sweep(make_request()) + assert result.failure_reason == "crn_mismatch" + assert result.details["crn_probe"]["point"] == probe_point + + +def test_unpublished_and_unscorable_points_are_none(tmp_path): + request = make_request() + dropped = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.0} + broken = {CREDIT: 0.5, SCALE: 2.0, TEMP: 0.0} + + def evaluate(request_obj, point): + if point == public(dropped): + return None + metrics = default_metrics(point) + if point == public(broken): + del metrics["gpu_hours"] + return metrics + + h = Harness(tmp_path) + h.operator.evaluate = evaluate + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request, missing=[dropped, broken]) + assert "gpu_hours" in result.metrics[RetuneParams(router=broken).key()]["error"] + assert ( + "not published" in result.metrics[RetuneParams(router=dropped).key()]["error"] + ) + assert result.details["missing_candidates"] == 2 + + +def test_scores_follow_the_retune_objective(tmp_path): + h = Harness(tmp_path, objective="TTFT") + request = make_request(objective="TTFT") + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request) + assert h.created()[0]["spec"]["objective"]["metric"] == "TTFT" + + +def test_engine_drift_between_candidates_is_engine_not_pinned(tmp_path): + h = Harness(tmp_path) + + def drift(name, parameters): + if h.kube.requests[name]["metadata"]["labels"][LABEL_SHARD] == "probe": + parameters["aggregated_parallelism"]["tensor"] = 4 + + h.operator.mutate_parameters = drift + result = h.sweep(make_request()) + assert result.failure_reason == "engine_not_pinned" + + +def test_router_policy_drift_is_engine_not_pinned(tmp_path): + h = Harness(tmp_path) + + def drift(name, parameters): + parameters["adapters"]["dynamo.router"]["prefill_load_model"] = {"type": "ais"} + + h.operator.mutate_parameters = drift + result = h.sweep(make_request()) + assert result.failure_reason == "engine_not_pinned" + + +def test_candidate_with_an_unknown_point_is_ignored(tmp_path): + h = Harness(tmp_path) + request = make_request() + odd = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.0} + + def drift(name, parameters): + router = parameters["adapters"]["dynamo.router"] + if ( + router["overlap_score_credit"], + router["prefill_load_scale"], + router["temperature"], + ) == (0.5, 1.0, 0.0): + router["prefill_load_scale"] = 7.0 + + h.operator.mutate_parameters = drift + result = h.sweep(request) + assert result.status == "ok", result.error + assert_scores(result, request, missing=[odd]) + assert "not one of the shard's values" in result.details["unmapped_candidates"][0] + + +def test_unscored_baseline_is_reported_for_the_controller(tmp_path): + request = make_request(baseline=grid(SMALL_DOMAIN)[0]) + baseline = request.baseline.router + + def evaluate(request_obj, point): + return None if point == public(baseline) else default_metrics(point) + + h = Harness(tmp_path) + h.operator.evaluate = evaluate + result = h.sweep(request) + # Status ok: _decide records the window failed for its missing baseline. + assert result.status == "ok" + assert result.scores[request.baseline.key()] is None + + +def test_empty_window_creates_nothing(tmp_path): + h = Harness(tmp_path) + h.cutter = FakeCutter(num_requests=0) + result = h.sweep(make_request()) + assert result.status == "empty" + assert result.failure_reason is None + assert result.error is None + assert h.created() == [] + + +@pytest.mark.parametrize( + "error, reason", + [ + (CutError("window_too_large"), "window_too_large"), + (CutError("trace_corrupt"), "trace_corrupt"), + (CutError("window_cut_timeout"), "window_cut_timeout"), + (CutError("something_else"), "window_cut_failed"), + (OSError("ENOENT"), "window_cut_failed"), + ], +) +def test_cut_failures_fail_the_window(tmp_path, error, reason): + h = Harness(tmp_path) + h.cutter = FakeCutter(error=error) + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == reason + assert h.created() == [] + + +def test_transient_list_failures_are_retried(tmp_path): + h = Harness(tmp_path) + h.kube.fail("list_requests", KubeError("timeout"), times=2) + h.kube.fail("list_candidates", KubeError("timeout")) + request = make_request() + result = h.sweep(request) + assert result.status == "ok", result.error + assert result.details["k8s_errors"] == 3 + assert_scores(result, request) + + +def test_permanent_list_failure_fails_the_window(tmp_path): + h = Harness(tmp_path) + error = KubeError("forbidden") + error.status = 403 # type: ignore[attr-defined] + h.kube.fail("list_requests", error) + result = h.sweep(make_request()) + assert result.failure_reason == "k8s_error" + + +def test_unexpected_exception_never_escapes(tmp_path): + h = Harness(tmp_path) + h.kube.fail("list_requests", RuntimeError("boom")) + result = h.sweep(make_request()) + assert result.status == "failed" + assert result.failure_reason == "internal_error" + assert "boom" in result.error + + +@pytest.mark.parametrize( + "kwargs, reason", + [ + ({"router_domain": None}, "invalid_request"), + ({"planner": {"load_scaling_down_sensitivity": 50}}, "invalid_request"), + ( + {"router_domain": {CREDIT: tuple(float(v) for v in range(500))}}, + "invalid_domain", + ), + ], +) +def test_invalid_requests_fail_cleanly(tmp_path, kwargs, reason): + h = Harness(tmp_path) + if "router_domain" in kwargs and kwargs["router_domain"] is not None: + # A domain whose keys do not match the candidates' router keys. + request = make_request(router_domain=kwargs["router_domain"]) + else: + request = make_request(**kwargs) + result = h.sweep(request) + assert result.status == "failed" + assert result.failure_reason == reason + assert h.created() == [] + + +def test_invalid_engine_fails_cleanly(tmp_path): + h = Harness(tmp_path) + result = h.sweep(make_request(spec_replicas={"aggregated": 0})) + assert result.failure_reason == "invalid_engine" + assert h.created() == [] + + +def test_deadline_defaults_to_the_request_timeout(tmp_path): + h = Harness(tmp_path) + h.operator.max_phase = "started" + result = h.sweep(make_request(deadline_s=None)) + assert result.status == "timeout" + assert result.details["deadline_s"] == LAUNCH_S + TIMEOUT_S + assert h.clock.now >= LAUNCH_S + TIMEOUT_S + + +def test_sweep_with_the_real_window_cutter(tmp_path): + """The executor and ``window_cut.cut_window`` agree on paths and deadline.""" + + window_cut = pytest.importorskip("dynamo.planner.retune.live.window_cut") + from dynamo.planner.tests.retune_sweep_helpers import trace_rows + + pod_dir = tmp_path / "trace" / "segments" / "llm" / "frontend-0" + pod_dir.mkdir(parents=True) + arrivals = [WINDOW_END_MS - 1000 * i for i in range(1, 6)] + rows = trace_rows(arrivals) + trace_rows([WINDOW_END_MS + 5], prefix="late") + with gzip.open(pod_dir / "req.000001.jsonl.gz", "wt", encoding="utf-8") as out: + out.writelines(json.dumps(row) + "\n" for row in rows) + + h = Harness(tmp_path) + clock = h.clock + + def real_cut(*args, deadline_s, **kwargs): + # The fake clock runs in the past; keep the same budget in real time. + return window_cut.cut_window( + *args, deadline_s=time.time() + (deadline_s - clock.now), **kwargs + ) + + request = make_request(trace_paths=(str(tmp_path / "trace" / "segments" / "llm"),)) + result = h.sweep(request, cut_window=real_cut) + assert result.status == "ok", result.error + assert result.num_requests == 5 + window_file = tmp_path / "trace" / WINDOW_FILE + with gzip.open(window_file, "rt", encoding="utf-8") as lines: + assert len(lines.readlines()) == 5 + _, claim_path = trace_pvc_source(h.created()[0]["spec"]) + assert tmp_path / "trace" / claim_path == window_file + assert_scores(result, request) diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_frontend_config.py b/components/src/dynamo/planner/tests/unit/test_live_retune_frontend_config.py new file mode 100644 index 000000000000..544339bc85fb --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_frontend_config.py @@ -0,0 +1,2449 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the retune imports. + +"""Frontend startup parsing and the spec premises of live retune. + +Covers ``find_frontend``, ``parse_frontend`` (the operator's argv +defaulting, env merging and ``$(VAR)`` expansion, the frontend's own parser +built against an explicit env, a golden comparison with the real frontend +parser, disaggregated maps) and +``evaluate_premises`` (P4, P6-P10) on table-driven DGD fixtures, including +P4 on workers that advertise their own router configuration (compared with +the real worker parser) and the P7 check of ``retune.dgdr.model``. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import sys +import types +import warnings +from typing import Any + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from dynamo.common.configuration.groups.kv_router_args import _KV_ROUTER_FIELDS +from dynamo.frontend.frontend_args import FrontendArgGroup, FrontendConfig +from dynamo.llm import KvRouterConfig +from dynamo.planner.config.planner_config import PlannerConfig +from dynamo.planner.retune.live import frontend_config as fc +from dynamo.planner.retune.live.contracts import FrontendStartup, PremiseReport + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +DGD = "llm" +NAMESPACE = "team-a" +CLAIM = "request-traces" +MODEL = "Qwen/Qwen3-32B" +TRACE_PATH = "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req" +KV_ARGS = ("-m", "dynamo.frontend", "--router-mode", "kv") +DEFAULT_AGG_MAP = { + "decode_active_request_weight": 0.0, + "disk_cache_hit_weight": 0.25, + "host_cache_hit_weight": 0.75, + "overlap_score_credit": 1.0, + "overlap_score_credit_decay": 0.0, + "prefill_load_scale": 1.0, + "router_temperature": 0.0, + "shared_cache_multiplier": 0.0, +} + +PRESET = { + "tensor": 2, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, +} +ENGINE_WORKER = { + "parallelism": {"preset": [PRESET]}, + "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + "kv_cache": {"block_size": 64, "capacity": {"memory_fraction": 0.9}}, +} +ENGINES = { + "agg": { + "mode": "aggregated", + "backend": "vllm", + "backend_version": "0.11.0", + "context_length": 32768, + "workers": {"aggregated": ENGINE_WORKER, "prefill": None, "decode": None}, + }, + "disagg": { + "mode": "disaggregated", + "backend": "vllm", + "backend_version": "0.11.0", + "context_length": 32768, + "workers": {"prefill": ENGINE_WORKER, "decode": ENGINE_WORKER}, + }, +} + + +# ------------------------------------------------------------- fixtures + + +def trace_env(**overrides: Any) -> list[dict[str, Any]]: + """The P10 trace recipe env; an override of None drops the entry.""" + + entries: dict[str, dict[str, Any]] = { + "DYN_POD_NAME": { + "name": "DYN_POD_NAME", + "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}, + }, + "DYN_REQUEST_TRACE": {"name": "DYN_REQUEST_TRACE", "value": "1"}, + "DYN_REQUEST_TRACE_RECORDS": { + "name": "DYN_REQUEST_TRACE_RECORDS", + "value": "request_end", + }, + "DYN_REQUEST_TRACE_SINKS": {"name": "DYN_REQUEST_TRACE_SINKS", "value": "file"}, + "DYN_REQUEST_TRACE_FILE_FORMAT": { + "name": "DYN_REQUEST_TRACE_FILE_FORMAT", + "value": "jsonl_gz", + }, + "DYN_REQUEST_TRACE_FILE_PATH": { + "name": "DYN_REQUEST_TRACE_FILE_PATH", + "value": TRACE_PATH, + }, + "DYN_REQUEST_TRACE_FILE_ROLL_BYTES": { + "name": "DYN_REQUEST_TRACE_FILE_ROLL_BYTES", + "value": "67108864", + }, + "DYN_REQUEST_TRACE_CAPACITY": { + "name": "DYN_REQUEST_TRACE_CAPACITY", + "value": "8192", + }, + } + for name, entry in overrides.items(): + if entry is None: + entries.pop(name, None) + elif isinstance(entry, dict): + entries[name] = {"name": name, **entry} + else: + entries[name] = {"name": name, "value": entry} + return list(entries.values()) + + +def frontend( + *, + args: tuple[str, ...] | list[str] | None = KV_ARGS, + command: tuple[str, ...] | list[str] | None = ("python3",), + env: list[dict[str, Any]] | None = None, + replicas: int | None = 1, + mounts: list[dict[str, Any]] | None = None, + volumes: list[dict[str, Any]] | None = None, + name: str = "Frontend", + **extra: Any, +) -> dict[str, Any]: + main: dict[str, Any] = { + "name": "main", + "image": "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.3.0", + "env": trace_env() if env is None else env, + "volumeMounts": ( + [{"name": "traces", "mountPath": "/trace"}] if mounts is None else mounts + ), + } + if command is not None: + main["command"] = list(command) + if args is not None: + main["args"] = list(args) + component: dict[str, Any] = { + "name": name, + "type": "frontend", + "podTemplate": { + "spec": { + "containers": [main], + "volumes": ( + [{"name": "traces", "persistentVolumeClaim": {"claimName": CLAIM}}] + if volumes is None + else volumes + ), + } + }, + **extra, + } + if replicas is not None: + component["replicas"] = replicas + return component + + +def worker( + name: str = "agg", *, type_: str = "worker", model: str = MODEL, **extra: Any +) -> dict[str, Any]: + return { + "name": name, + "type": type_, + "replicas": 2, + "podTemplate": { + "spec": { + "containers": [ + { + "name": "main", + "command": ["python3", "-m", "dynamo.vllm"], + "args": [f"--model={model}", "--tensor-parallel-size=2"], + } + ] + } + }, + **extra, + } + + +def planner(replicas: int = 1) -> dict[str, Any]: + return { + "name": "Planner", + "type": "planner", + "replicas": replicas, + "podTemplate": {"spec": {"containers": [{"name": "main"}]}}, + } + + +def dgd(*components: dict[str, Any], **spec: Any) -> dict[str, Any]: + if not components: + components = (frontend(), worker(), planner()) + return { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "metadata": {"name": DGD, "namespace": NAMESPACE, "uid": "u-1"}, + "spec": {"components": list(components), **spec}, + } + + +def planner_config( + mode: str = "agg", router: Any = "default", **dgdr: Any +) -> PlannerConfig: + """A validated live retune config (P1-P3, P5 and P13 hold).""" + + return PlannerConfig( + namespace="test-ns", + environment="kubernetes", + mode=mode, + enable_load_scaling=False, + enable_throughput_scaling=False, + retune={ + "interval_seconds": 21600.0, + "timeout_seconds": 20400.0, + "router": router, + "executor": "dgdr", + "dgdr": { + "claim_name": CLAIM, + "model": MODEL, + "sku": "h200_sxm", + "engine": copy.deepcopy(ENGINES[mode]), + **dgdr, + }, + }, + ) + + +def parse(deployment: dict[str, Any], mode: str = "agg") -> FrontendStartup: + startup = fc.parse_frontend(deployment, planner_mode=mode) + assert isinstance(startup, FrontendStartup), startup + return startup + + +def parse_error(deployment: dict[str, Any], mode: str = "agg") -> str: + result = fc.parse_frontend(deployment, planner_mode=mode) + assert isinstance(result, str), result + return result + + +def evaluate( + deployment: dict[str, Any], *, config: PlannerConfig | None = None +) -> PremiseReport: + return fc.evaluate_premises(deployment, planner_config=config or planner_config()) + + +def reasons_matching(report: PremiseReport, prefix: str) -> list[str]: + return [reason for reason in report.reasons if reason.startswith(prefix)] + + +def with_frontend(**kwargs: Any) -> dict[str, Any]: + return dgd(frontend(**kwargs), worker(), planner()) + + +# --------------------------------------------------------- find_frontend + + +def test_find_frontend_v1beta1(): + name, spec = fc.find_frontend(dgd()) + assert name == "Frontend" + assert spec["type"] == "frontend" + + +@pytest.mark.parametrize( + ("components", "match"), + [ + ((worker(), planner()), "no frontend component"), + ( + (frontend(), frontend(name="Frontend2"), worker()), + "2 frontend components", + ), + ((worker(frontendSidecar="sidecar-frontend"),), "frontendSidecar of agg"), + (({"name": "Epp", "type": "epp"}, worker()), "EPP component"), + ], +) +def test_find_frontend_rejects(components, match): + with pytest.raises(fc.FrontendParseError, match=match): + fc.find_frontend(dgd(*components)) + assert parse_error(dgd(*components)).startswith("P7:") + + +# ----------------------------------------------------- argv and parsing + + +def test_parse_recipe_frontend(): + startup = parse(dgd()) + assert startup.router_mode == "kv_router" + assert startup.router_config == {} + assert startup.startup_maps == {"aggregated": DEFAULT_AGG_MAP} + assert startup.frontend_component == "Frontend" + assert startup.replicas == 1 + assert startup.replica_sync is False + assert json.loads(startup.fingerprint_source) == { + "router_mode": "kv_router", + "router_config": {}, + "startup_maps": {"aggregated": DEFAULT_AGG_MAP}, + } + assert startup.fingerprint_source == fc.canonical_json( + json.loads(startup.fingerprint_source) + ) + + +@pytest.mark.parametrize( + ("command", "args", "mode"), + [ + # Operator default: python3 -m dynamo.frontend (round-robin). + (None, None, "round_robin"), + # Authored args keep the default python3 command. + (None, KV_ARGS, "kv_router"), + (("python3",), KV_ARGS, "kv_router"), + (("python3", "-m", "dynamo.frontend"), ("--router-mode", "kv"), "kv_router"), + ( + ("/usr/bin/python3", "-u", "-X", "dev", "-m", "dynamo.frontend"), + ("--router-mode=kv",), + "kv_router", + ), + (("python", "-mdynamo.frontend", "--router-mode", "kv"), None, "kv_router"), + (("python3",), ("-m", "dynamo.frontend", "--router-mode", "random"), "random"), + ( + ("python3",), + ("-m", "dynamo.frontend", "--router-mode", "power-of-two"), + "power_of_two", + ), + ], +) +def test_effective_argv(command, args, mode): + startup = parse(with_frontend(command=command, args=args)) + assert startup.router_mode == mode + + +def test_authored_command_drops_default_args(): + # graph.go: an authored frontend command clears the default args, so the + # command alone must name the module. + error = parse_error(with_frontend(command=("python3",), args=())) + assert "does not run python -m dynamo.frontend" in error + + +@pytest.mark.parametrize( + "command,args", + [ + (("/bin/sh", "-c"), ("python3 -m dynamo.frontend --router-mode kv",)), + (("bash", "-lc", "exec python3 -m dynamo.frontend --router-mode kv"), None), + (("python3", "-c", "import runpy"), None), + ], +) +def test_shell_wrapped_command_is_unparseable(command, args): + error = parse_error(with_frontend(command=command, args=args)) + assert error.startswith("P4:") + assert "shell-wrapped" in error + + +@pytest.mark.parametrize( + "command,args,match", + [ + (("dynamo-frontend",), ("--router-mode", "kv"), "not a python interpreter"), + (("python3",), ("-m", "dynamo.vllm"), "does not run python -m dynamo.frontend"), + (("python3",), ("serve.py",), "does not run python -m dynamo.frontend"), + ], +) +def test_non_frontend_command_is_unparseable(command, args, match): + assert match in parse_error(with_frontend(command=command, args=args)) + + +def test_frontend_roles_are_unparseable(): + error = parse_error(dgd(frontend(roles=[{"name": "leader"}]), worker())) + assert "uses roles" in error + + +def test_env_sets_router_mode_and_knobs(): + env = trace_env( + DYN_ROUTER_MODE="kv", + DYN_ROUTER_TEMPERATURE="0.5", + DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT="0.25", + ) + startup = parse(with_frontend(args=("-m", "dynamo.frontend"), env=env)) + assert startup.router_mode == "kv_router" + assert startup.router_config == { + "router_temperature": 0.5, + "overlap_score_credit": 0.25, + } + assert startup.startup_maps["aggregated"]["router_temperature"] == 0.5 + assert startup.startup_maps["aggregated"]["overlap_score_credit"] == 0.25 + + +def test_cli_wins_over_env_and_container_env_over_graph_env(): + env = trace_env(DYN_ROUTER_TEMPERATURE="0.5", DYN_ROUTER_PREFILL_LOAD_SCALE="2") + deployment = dgd( + frontend(args=(*KV_ARGS, "--router-temperature", "0.2"), env=env), + worker(), + env=[ + {"name": "DYN_ROUTER_PREFILL_LOAD_SCALE", "value": "3"}, + {"name": "DYN_ROUTER_KV_OVERLAP_SCORE_CREDIT", "value": "0.5"}, + ], + ) + assert parse(deployment).router_config == { + "router_temperature": 0.2, + "prefill_load_scale": 2.0, + "overlap_score_credit": 0.5, + } + + +def test_planner_process_env_is_ignored(monkeypatch): + monkeypatch.setenv("DYN_ROUTER_MODE", "kv") + monkeypatch.setenv("DYN_ROUTER_TEMPERATURE", "0.9") + monkeypatch.setenv("DYN_ROUTER_REPLICA_SYNC", "1") + before = dict(os.environ) + startup = parse(with_frontend(args=("-m", "dynamo.frontend"))) + assert startup.router_mode == "round_robin" + assert startup.router_config == {} + assert startup.replica_sync is False + assert dict(os.environ) == before + + +def test_abbreviated_flag_matches_frontend_parser(): + # The frontend parses with argparse abbreviations enabled. + startup = parse(with_frontend(args=(*KV_ARGS, "--router-temp", "0.4"))) + assert startup.router_config == {"router_temperature": 0.4} + + +def test_unknown_frontend_args_are_ignored(): + args = ( + *KV_ARGS, + "--http-port", + "8000", + "--kv-cache-block-size", + "64", + "--trust-remote-code", + "--some-vllm-processor-flag", + "x", + ) + assert parse(with_frontend(args=args)).router_config == {} + + +def test_legacy_and_removed_flags_parse_quietly(caplog): + args = ( + *KV_ARGS, + "--no-kv-events", + "--admission-control", + "none", + "--enforce-disagg", + ) + startup = parse(with_frontend(args=args)) + assert startup.router_config == {"use_kv_events": False} + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_replica_sync_and_disagg_maps(): + startup = parse( + with_frontend( + args=(*KV_ARGS, "--router-replica-sync", "--router-temperature", "0.3") + ), + mode="disagg", + ) + assert startup.replica_sync is True + assert startup.router_config == { + "router_replica_sync": True, + "router_temperature": 0.3, + } + assert set(startup.startup_maps) == {"prefill", "decode"} + prefill, decode = startup.startup_maps["prefill"], startup.startup_maps["decode"] + assert prefill["overlap_score_credit"] == 1.0 + assert prefill["host_cache_hit_weight"] == 0.75 + # Plain decode routers report zero cache credits. + assert decode["overlap_score_credit"] == 0.0 + assert decode["host_cache_hit_weight"] == 0.0 + assert decode["disk_cache_hit_weight"] == 0.0 + assert decode["router_temperature"] == 0.3 + + +def test_zero_credit_links_cache_weights(): + startup = parse( + with_frontend(args=(*KV_ARGS, "--router-kv-overlap-score-credit", "0")) + ) + agg = startup.startup_maps["aggregated"] + assert agg["overlap_score_credit"] == 0.0 + assert agg["host_cache_hit_weight"] == 0.0 + assert agg["disk_cache_hit_weight"] == 0.0 + + +def test_load_aware_implies_kv_mode(): + startup = parse(with_frontend(args=("-m", "dynamo.frontend", "--load-aware"))) + assert startup.router_mode == "kv_router" + assert startup.router_config["use_kv_events"] is False + assert startup.router_config["overlap_score_credit"] == 0.0 + + +def test_fingerprint_ignores_replicas_but_not_router_config(): + base = parse(dgd(frontend(replicas=1), worker())) + scaled = parse(dgd(frontend(replicas=3), worker())) + tuned = parse( + dgd(frontend(args=(*KV_ARGS, "--router-temperature", "0.1")), worker()) + ) + assert scaled.replicas == 3 + assert base.fingerprint_source == scaled.fingerprint_source + assert base.fingerprint_source != tuned.fingerprint_source + + +def test_omitted_replicas_default_to_one(): + assert parse(dgd(frontend(replicas=None), worker())).replicas == 1 + + +# -------------------------------------------- env references and sources + + +def test_arg_expands_literal_env(): + env = trace_env(TEMP="0.3") + args = (*KV_ARGS, "--router-temperature", "$(TEMP)") + assert parse(with_frontend(args=args, env=env)).router_config == { + "router_temperature": 0.3 + } + + +def test_escaped_reference_stays_literal(): + # $$(TEMP) is the literal "$(TEMP)", which is not a float. + env = trace_env(TEMP="0.3") + args = (*KV_ARGS, "--router-temperature", "$$(TEMP)") + error = parse_error(with_frontend(args=args, env=env)) + assert "cannot parse the frontend args" in error + + +def test_arg_referencing_value_from_is_unparseable(): + env = trace_env(TEMP={"valueFrom": {"configMapKeyRef": {"name": "c", "key": "t"}}}) + args = (*KV_ARGS, "--router-temperature", "$(TEMP)") + error = parse_error(with_frontend(args=args, env=env)) + assert "$(TEMP)" in error and "cannot read" in error + + +def test_opaque_string_router_arg_is_unparseable(): + env = trace_env(MODE={"valueFrom": {"secretKeyRef": {"name": "s", "key": "k"}}}) + args = ("-m", "dynamo.frontend", "--router-tracking-key-id", "id-$(MODE)") + error = parse_error(with_frontend(args=args, env=env)) + assert "router_tracking_key_id" in error + + +def test_arg_that_is_an_opaque_reference_is_unparseable(): + env = trace_env(EXTRA={"valueFrom": {"secretKeyRef": {"name": "s", "key": "k"}}}) + error = parse_error(with_frontend(args=(*KV_ARGS, "$(EXTRA)"), env=env)) + assert "$(EXTRA)" in error + + +def test_opaque_reference_in_a_non_router_arg_is_fine(): + args = (*KV_ARGS, "--metrics-prefix", "fe-$(POD_NAME)") + assert parse(with_frontend(args=args)).router_config == {} + + +OPAQUE_PORT = {"valueFrom": {"configMapKeyRef": {"name": "c", "key": "port"}}} + + +@pytest.mark.parametrize( + "args", + [ + # Typed non-router flags keep the unreadable value unconverted. + (*KV_ARGS, "--http-port", "$(HTTP_PORT)"), + (*KV_ARGS, "--http-port=$(HTTP_PORT)"), + (*KV_ARGS, "--grpc-metrics-port", "$(POD_UID)"), + # A non-router choices flag. + (*KV_ARGS, "--dyn-chat-processor", "$(HTTP_PORT)"), + # The value of a flag the frontend hands to its chat processor. + ( + *KV_ARGS, + "--dyn-chat-processor", + "vllm", + "--tool-call-parser", + "$(HTTP_PORT)", + ), + (*KV_ARGS, "--dyn-chat-processor", "vllm", "--tool-call-parser=$(HTTP_PORT)"), + # Attached to an option that takes no value, it is the option's value. + ( + *KV_ARGS, + "--dyn-chat-processor", + "vllm", + "--enable-auto-tool-choice", + "--some-vllm-opt=$(HTTP_PORT)", + ), + ], +) +def test_opaque_reference_in_a_typed_non_router_arg_is_fine(args): + startup = parse(with_frontend(args=args, env=trace_env(HTTP_PORT=OPAQUE_PORT))) + assert startup.router_mode == "kv_router" + assert startup.router_config == {} + + +@pytest.mark.parametrize( + ("args", "match"), + [ + # Router flags whose value is unreadable. + ( + ("-m", "dynamo.frontend", "--router-mode", "$(HTTP_PORT)"), + "router arg references $(HTTP_PORT)", + ), + ( + (*KV_ARGS, "--router-temp", "$(HTTP_PORT)"), + "router arg references $(HTTP_PORT)", + ), + ( + (*KV_ARGS, "--router-ttl-secs=$(HTTP_PORT)"), + "router arg references $(HTTP_PORT)", + ), + # Left over where a flag could stand. + ( + (*KV_ARGS, "--router-kv-events", "$(HTTP_PORT)"), + "could set a router flag", + ), + ((*KV_ARGS, "--$(HTTP_PORT)"), "could set a router flag"), + ((*KV_ARGS, "--$(HTTP_PORT)=1"), "could set a router flag"), + ( + (*KV_ARGS, "--tool-call-parser", "--$(HTTP_PORT)", "$(HTTP_PORT)"), + "could set a router flag", + ), + # The chat processor's flag takes no value, so the frontend would + # parse what follows it as its own flag. + ( + ( + *KV_ARGS, + "--dyn-chat-processor", + "vllm", + "--enable-auto-tool-choice", + "$(HTTP_PORT)", + ), + "could set a router flag", + ), + ( + ( + *KV_ARGS, + "--dyn-chat-processor", + "vllm", + "--tool-call-parser", + "hermes", + "--enable-auto-tool-choice", + "--trust-remote-code", + "$(HTTP_PORT)", + ), + "could set a router flag", + ), + # Inert env is still unreadable as a router value or an option name. + ((*KV_ARGS, "--router-temp", "$(POD_UID)"), "router arg references $(POD_UID)"), + ((*KV_ARGS, "--$(POD_NAME)"), "could set a router flag"), + ], +) +def test_opaque_reference_that_could_set_a_router_flag_is_unparseable(args, match): + error = parse_error(with_frontend(args=args, env=trace_env(HTTP_PORT=OPAQUE_PORT))) + assert error.startswith("P4:") + assert match in error + + +@pytest.mark.parametrize( + "args", + [ + (*KV_ARGS, "$(POD_NAME)"), + (*KV_ARGS, "--router-kv-events", "$(POD_UID)"), + (*KV_ARGS, "$(NODE)"), + ], +) +def test_inert_reference_where_a_flag_could_stand_is_fine(args): + # A pod name, uid or node name never starts with "-". + env = trace_env(NODE={"valueFrom": {"fieldRef": {"fieldPath": "spec.nodeName"}}}) + assert parse(with_frontend(args=args, env=env)).router_config == {} + + +@pytest.mark.parametrize( + ("args", "prefix", "match"), + [ + ((*KV_ARGS, "$(ROUTER_ARG)"), None, "could set a router flag"), + ( + (*KV_ARGS, "--router-temperature", "$(T)"), + None, + "router arg references $(T)", + ), + ((*KV_ARGS, "$(APP_ROUTER_ARG)"), "APP_", "could set a router flag"), + # A prefix that cannot produce the name, and command substitution. + ((*KV_ARGS, "$(ROUTER_ARG)"), "APP_", None), + ((*KV_ARGS, "--metrics-prefix", "$(seq 1 3)"), None, None), + ], +) +def test_env_from_may_define_an_undefined_reference(args, prefix, match): + deployment = with_frontend(args=args) + source: dict[str, Any] = {"configMapRef": {"name": "router-options"}} + if prefix is not None: + source["prefix"] = prefix + deployment["spec"]["components"][0]["podTemplate"]["spec"]["containers"][0][ + "envFrom" + ] = [source] + result = fc.parse_frontend(deployment, planner_mode="agg") + if match is None: + assert isinstance(result, FrontendStartup), result + else: + assert isinstance(result, str) and match in result, result + + +def test_router_env_from_value_from_is_unparseable(): + env = trace_env( + DYN_ROUTER_TEMPERATURE={ + "valueFrom": {"configMapKeyRef": {"name": "c", "key": "t"}} + } + ) + error = parse_error(with_frontend(env=env)) + assert "DYN_ROUTER_TEMPERATURE" in error and "valueFrom" in error + + +def test_router_env_referencing_value_from_is_unparseable(): + env = trace_env( + A_TEMP={"valueFrom": {"configMapKeyRef": {"name": "c", "key": "t"}}}, + DYN_ROUTER_TEMPERATURE="$(A_TEMP)", + ) + assert "DYN_ROUTER_TEMPERATURE" in parse_error(with_frontend(env=env)) + + +def test_router_env_references_earlier_literal(): + env = trace_env(A_TEMP="0.7", DYN_ROUTER_TEMPERATURE="$(A_TEMP)") + assert parse(with_frontend(env=env)).router_config == {"router_temperature": 0.7} + + +def test_non_router_value_from_is_fine(): + env = trace_env( + HF_TOKEN={"valueFrom": {"secretKeyRef": {"name": "hf", "key": "t"}}} + ) + assert parse(with_frontend(env=env)).router_mode == "kv_router" + + +@pytest.mark.parametrize( + ("source", "warned"), + [ + ({"secretRef": {"name": "hf-token-secret"}}, True), + ({"configMapRef": {"name": "router"}, "prefix": "DYN_"}, True), + ({"configMapRef": {"name": "router"}, "prefix": "D"}, True), + ({"secretRef": {"name": "hf-token-secret"}, "prefix": "HF_"}, False), + ], +) +def test_env_from_is_parsed_and_warned_when_it_may_set_router_env(source, warned): + deployment = with_frontend() + deployment["spec"]["components"][0]["podTemplate"]["spec"]["containers"][0][ + "envFrom" + ] = [source] + assert isinstance( + fc.parse_frontend(deployment, planner_mode="agg"), FrontendStartup + ) + report = evaluate(deployment) + envfrom_warnings = [w for w in report.warnings if "envFrom" in w] + assert bool(envfrom_warnings) == warned + assert not reasons_matching(report, "P4") + + +def test_router_policy_config_is_unparseable(): + error = parse_error( + with_frontend(args=(*KV_ARGS, "--router-policy-config", "/etc/policy.yaml")) + ) + assert "--router-policy-config" in error + + +@pytest.mark.parametrize( + "args,env,match", + [ + ((*KV_ARGS, "--router-kv-overlap-score-credit", "-1"), {}, "invalid"), + ((*KV_ARGS, "--router-temperature", "hot"), {}, "cannot parse"), + ((*KV_ARGS,), {"DYN_ROUTER_TEMPERATURE": "hot"}, "invalid frontend env"), + ( + (*KV_ARGS, "--router-approximate-cache-policy", "lru"), + {}, + "invalid", + ), + ], +) +def test_invalid_router_config_is_unparseable(args, env, match): + error = parse_error(with_frontend(args=args, env=trace_env(**env))) + assert error.startswith("P4:") + assert match in error + + +# ------------------------------------------------------ parser fidelity + + +GOLDEN_CASES = [ + ((*KV_ARGS,), {}), + ((*KV_ARGS, "--router-temperature", "0.5", "--router-replica-sync"), {}), + ( + ("-m", "dynamo.frontend"), + {"DYN_ROUTER_MODE": "kv", "DYN_ROUTER_USE_KV_EVENTS": "0"}, + ), + ((*KV_ARGS, "--no-kv-events", "--router-ttl-secs", "30"), {}), + ( + (*KV_ARGS, "--router-kv-overlap-score-credit", "0"), + {"DYN_ROUTER_TEMPERATURE": "0.1"}, + ), + (("-m", "dynamo.frontend", "--load-aware"), {}), + ((*KV_ARGS,), {"DYN_OVERLAP_SCORE_WEIGHT": "2"}), + ( + (*KV_ARGS,), + {"DYN_ROUTER_KV_OVERLAP_SCORE_WEIGHT": "0", "DYN_OVERLAP_SCORE_WEIGHT": "2"}, + ), + ((*KV_ARGS, "--router-temp", "0.25", "--router-event-threads", "8"), {}), + ( + (*KV_ARGS, "--router-tracking-hash", "keyed-xxh3-v1"), + {"DYN_ROUTER_TRACKING_KEY_FILE": "/k", "DYN_ROUTER_TRACKING_KEY_ID": "e1"}, + ), +] + + +@pytest.mark.parametrize(("args", "env"), GOLDEN_CASES) +def test_matches_real_frontend_parser(monkeypatch, args, env): + """The frontend's own parser, run with its env in os.environ, agrees.""" + + for name in list(os.environ): + if name.startswith("DYN_"): + monkeypatch.delenv(name) + for name, value in env.items(): + monkeypatch.setenv(name, value) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + parser = argparse.ArgumentParser() + FrontendArgGroup().add_arguments(parser) + namespace, _ = parser.parse_known_args(list(args[2:])) + config = FrontendConfig.from_cli_args(namespace) + if config.load_aware: + config.router_mode = "kv" + config.apply_load_aware_preset() + config.apply_conditional_disagg_config() + expected = json.loads( + KvRouterConfig(**{f: getattr(config, f) for f in _KV_ROUTER_FIELDS}).to_json() + ) + + for name in env: + monkeypatch.delenv(name) + startup = parse(with_frontend(args=args, env=trace_env(**env))) + actual = json.loads( + KvRouterConfig.from_json(json.dumps(startup.router_config)).to_json() + ) + assert actual == expected + assert startup.router_mode == fc._router_mode(config.router_mode) + + +def test_router_fields_are_all_covered(): + """Every frontend router field is in the JSON or checked on the flags.""" + + serialized = set( + json.loads( + KvRouterConfig( + router_tracking_hash="keyed-xxh3-v1", + router_tracking_key_file="/k", + router_tracking_key_id="e", + decode_active_request_weight=1.0, + enable_session_prefix_index=True, + shared_cache_multiplier=0.5, + conditional_disagg_enabled=True, + conditional_disagg_policy="isl_or_load", + conditional_disagg_eff_isl_threshold=1, + conditional_disagg_eff_isl_ratio_threshold=0.5, + conditional_disagg_prefill_busy_threshold=0.5, + conditional_disagg_decode_busy_threshold=0.5, + ).to_json() + ) + ) + covered = serialized | set(fc._SERDE_SKIPPED_DEFAULTS) | {"overlap_score_weight"} + assert set(_KV_ROUTER_FIELDS) <= covered + + +# --------------------------------------------------------- P4 premises + + +def test_recipe_dgd_satisfies_every_premise(): + report = evaluate(dgd()) + assert report.p4_ok and report.spec_ok, report.reasons + assert report.reasons == () + assert report.startup is not None + assert not [w for w in report.warnings if w.startswith("P6")] + + +def test_parse_error_fails_p4_only(): + deployment = with_frontend( + command=("/bin/sh", "-c"), args=("python3 -m dynamo.frontend",) + ) + report = evaluate(deployment) + assert not report.p4_ok and not report.spec_ok + assert len(report.reasons) == 1 and "shell-wrapped" in report.reasons[0] + assert report.startup is None + + +@pytest.mark.parametrize( + ("args", "env", "match"), + [ + (("-m", "dynamo.frontend"), {}, "routes with 'round_robin'"), + ( + ("-m", "dynamo.frontend", "--router-mode", "random"), + {}, + "routes with 'random'", + ), + ((*KV_ARGS, "--no-router-kv-events"), {}, "use_kv_events=False"), + ( + (*KV_ARGS, "--router-host-cache-hit-weight", "0.5"), + {}, + "host_cache_hit_weight=0.5", + ), + ( + (*KV_ARGS, "--router-queue-threshold", "0.5"), + {}, + "router_queue_threshold=0.5", + ), + ( + (*KV_ARGS, "--router-kv-overlap-score-credit-decay", "1"), + {}, + "overlap_score_credit_decay", + ), + ( + (*KV_ARGS, "--router-decode-active-request-weight", "1"), + {}, + "decode_active_request_weight", + ), + ( + (*KV_ARGS, "--router-prefill-load-model", "ais"), + {}, + "router_prefill_load_model is 'aic'", + ), + ( + (*KV_ARGS, "--router-conditional-disagg"), + {}, + "conditional disaggregation", + ), + ( + ( + *KV_ARGS, + "--no-router-kv-events", + "--router-approximate-cache-policy", + "lru", + ), + {}, + "router_approximate_cache_policy='lru'", + ), + ((*KV_ARGS, "--router-prefill-policy", "default"), {}, "router_prefill_policy"), + ((*KV_ARGS, "--router-decode-policy", "default"), {}, "router_decode_policy"), + ((*KV_ARGS,), {"DYN_OVERLAP_SCORE_WEIGHT": "2"}, "overlap_score_weight"), + ( + ("-m", "dynamo.frontend"), + {"DYN_ROUTER_LOAD_AWARE": "1"}, + "load-aware preset", + ), + ( + (*KV_ARGS,), + {"DYN_ROUTER_WORKER_SELECTION_POLICY": "mine"}, + "worker-selection", + ), + ( + (*KV_ARGS,), + {"DYN_ROUTER_POSITIONAL_SEARCH_MODE": "x"}, + "DYN_ROUTER_POSITIONAL_SEARCH_MODE", + ), + ( + (*KV_ARGS,), + {"DYN_ROUTER_SESSION_AFFINITY_TTL_SECS": "60"}, + "session_affinity_ttl_secs=60", + ), + ( + (*KV_ARGS, "--active-decode-blocks-threshold", "0.9"), + {}, + "active_decode_blocks_threshold", + ), + ( + (*KV_ARGS, "--active-prefill-tokens-threshold", "100"), + {}, + "active_prefill_tokens_threshold", + ), + ], +) +def test_p4_rejects(args, env, match): + report = evaluate(with_frontend(args=args, env=trace_env(**env))) + assert not report.p4_ok and not report.spec_ok + assert any( + match in reason for reason in reasons_matching(report, "P4:") + ), report.reasons + + +@pytest.mark.parametrize( + ("args", "env"), + [ + # Searched knobs may start anywhere. + ( + ( + *KV_ARGS, + "--router-temperature", + "0.5", + "--router-prefill-load-scale", + "2", + ), + {}, + ), + ((*KV_ARGS, "--router-kv-overlap-score-credit", "0.5"), {}), + # Tolerated fields. + ((*KV_ARGS, "--router-event-threads", "8", "--router-replica-sync"), {}), + ((*KV_ARGS, "--enable-session-prefix-index", "--serve-indexer"), {}), + ( + (*KV_ARGS, "--router-tracking-hash", "keyed-xxh3-v1"), + {"DYN_ROUTER_TRACKING_KEY_FILE": "/k", "DYN_ROUTER_TRACKING_KEY_ID": "e"}, + ), + # Frontend-only and transport settings. + ( + (*KV_ARGS,), + { + "DYN_ROUTER_MIN_INITIAL_WORKERS": "2", + "DYN_ROUTER_ZMQ_ENDPOINTS_PER_SUB": "4", + }, + ), + ((*KV_ARGS, "--router-session-affinity-mode", "soft"), {}), + ], +) +def test_p4_accepts(args, env): + report = evaluate(with_frontend(args=args, env=trace_env(**env))) + assert report.p4_ok, report.reasons + + +def test_p4_unsearched_knob_when_router_domain_is_pinned(): + # A scalar still searches (pins) the knob: it stays in the domain. + config = planner_config( + router={ + "overlap_score_credit": 1.0, + "prefill_load_scale": 1.0, + "temperature": 0.0, + } + ) + report = evaluate( + with_frontend(args=(*KV_ARGS, "--router-temperature", "0.5")), config=config + ) + assert report.p4_ok, report.reasons + + +def test_p4_zero_startup_credit_rejects_positive_candidates(): + zero = (*KV_ARGS, "--router-kv-overlap-score-credit", "0") + # The default domain collapses the omitted credit to [0.0]. + assert evaluate(with_frontend(args=zero)).p4_ok + config = planner_config( + router={"overlap_score_credit": {"choices": [0.0, 0.5]}, "temperature": 0.0} + ) + report = evaluate(with_frontend(args=zero), config=config) + assert not report.p4_ok + assert "overlap_score_credit > 0" in report.reasons[0] + + +def test_p4_disagg_decode_credit_zero_is_not_a_zero_credit_startup(): + # Plain decode routers always report credit 0; only prefill counts. + config = planner_config(mode="disagg") + report = evaluate( + dgd( + frontend(), + worker("prefill", type_="prefill"), + worker("decode", type_="decode"), + planner(), + ), + config=config, + ) + assert report.p4_ok, report.reasons + assert report.startup is not None + assert set(report.startup.startup_maps) == {"prefill", "decode"} + + +# ------------------------------------------- P4: worker router advertisement + + +SECRET = {"valueFrom": {"secretKeyRef": {"name": "s", "key": "k"}}} +# A Secret that may hold any key. +ENV_FROM = [{"secretRef": {"name": "worker-launch-options"}}] +# A worker launch in a shell script. +WORKER = f"python3 -m dynamo.vllm --model {MODEL}" + + +def router_worker( + *args: str, + env: list[dict[str, Any]] | None = None, + env_from: list[dict[str, Any]] | None = None, + name: str = "agg", + type_: str = "worker", +) -> dict[str, Any]: + component = worker(name, type_=type_) + main = component["podTemplate"]["spec"]["containers"][0] + main["args"] = [*main["args"], *args] + if env is not None: + main["env"] = env + if env_from is not None: + main["envFrom"] = env_from + return component + + +def shell_worker( + script: str, + *, + command: tuple[str, ...] | None = ("/bin/sh", "-c"), + env: list[dict[str, Any]] | None = None, + env_from: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """A worker running ``script``; ``command=None`` omits the command.""" + + component = worker() + main = component["podTemplate"]["spec"]["containers"][0] + main["args"] = [script] + if command is None: + del main["command"] + else: + main["command"] = list(command) + if env is not None: + main["env"] = env + if env_from is not None: + main["envFrom"] = env_from + return component + + +def field_ref_env(name: str, path: str) -> dict[str, Any]: + return {"name": name, "valueFrom": {"fieldRef": {"fieldPath": path}}} + + +def router_arg_env() -> list[dict[str, Any]]: + """ROUTER_ARG from a ConfigMap the planner cannot read.""" + + return [ + { + "name": "ROUTER_ARG", + "valueFrom": { + "configMapKeyRef": { + "name": "worker-launch-options", + "key": "router-arg", + } + }, + } + ] + + +def minimax_prefill_worker(*router_args: str) -> dict[str, Any]: + """The PrefillWorker of recipes/minimax-m3/vllm/disagg-gb200-agentic.""" + + return { + "name": "PrefillWorker", + "type": "prefill", + "replicas": 3, + "podTemplate": { + "spec": { + "containers": [ + { + "name": "main", + "command": ["python3"], + "args": [ + "-m", + "dynamo.vllm", + *router_args, + "--model", + "$(MODEL_NAME)", + "--served-model-name", + "$(SERVED_MODEL_NAME)", + "--tensor-parallel-size", + "4", + "--speculative-config", + "$(SPECULATIVE_CONFIG)", + "--disaggregation-mode", + "prefill", + "--kv-transfer-config", + "$(KV_TRANSFER_CONFIG)", + "--kv-events-config", + '{"enable_kv_cache_events":true,"publisher":"zmq"}', + ], + "env": [ + { + "name": "KV_TRANSFER_CONFIG", + "value": '{"kv_connector":"NixlConnector",' + '"kv_role":"kv_both"}', + }, + { + "name": "SPECULATIVE_CONFIG", + "valueFrom": { + "configMapKeyRef": { + "name": "minimax-m3-config", + "key": "speculative-config", + } + }, + }, + {"name": "MODEL_NAME", "value": "nvidia/MiniMax-M3-NVFP4"}, + { + "name": "SERVED_MODEL_NAME", + "value": "nvidia/MiniMax-M3-NVFP4", + }, + ], + } + ] + } + }, + } + + +@pytest.mark.parametrize( + ("args", "env", "match"), + [ + (("--router-mode", "kv"), None, "advertises its own router configuration"), + (("--router-mode=round-robin",), None, "router mode 'round-robin'"), + (("--router-m", "kv", "--no-router-kv-events"), None, "router mode 'kv'"), + ((), [{"name": "DYN_ROUTER_MODE", "value": "kv"}], "router mode 'kv'"), + ((), [{"name": "DYN_ROUTER_MODE", "value": ""}], "router mode ''"), + ( + (), + [{"name": "DYN_ROUTER_MODE", **SECRET}], + "DYN_ROUTER_MODE through valueFrom", + ), + ( + (), + [ + {"name": "A_MODE", **SECRET}, + {"name": "DYN_ROUTER_MODE", "value": "$(A_MODE)"}, + ], + "DYN_ROUTER_MODE through valueFrom", + ), + ( + ("--router-mode", "$(A_MODE)"), + [{"name": "A_MODE", **SECRET}], + "router flag that references $(A_MODE)", + ), + ( + ("--router-mode", "$(A_MODE)"), + [{"name": "A_MODE", "value": "kv"}], + "router mode 'kv'", + ), + (("--router-mode", "fastest"), None, "cannot parse"), + ( + (), + [{"name": "DYN_ROUTER_TEMPERATURE", "value": "hot"}], + "invalid router env", + ), + # An unreadable argument where an option could stand, e.g. a + # ConfigMap value "--router-mode=kv". + ( + ("--no-router-kv-events", "$(ROUTER_ARG)"), + router_arg_env(), + "has an argument that references $(ROUTER_ARG)", + ), + ( + ("$(ROUTER_ARG)",), + router_arg_env(), + "has an argument that references $(ROUTER_ARG)", + ), + ( + ("--$(A_MODE)=kv",), + [{"name": "A_MODE", **SECRET}], + "has an option whose name references $(A_MODE)", + ), + ( + ("--router-$(A_MODE)", "kv"), + [{"name": "A_MODE", **SECRET}], + "has an option whose name references $(A_MODE)", + ), + # Inert env could still spell an option's name. + ( + ("--x-$(MY_POD)=1",), + [field_ref_env("MY_POD", "metadata.name")], + "has an option whose name references $(MY_POD)", + ), + ( + ("--speculative-config", "$(A_SECRET)", "$(A_SECRET)"), + [{"name": "A_SECRET", **SECRET}], + "has an argument that references $(A_SECRET)", + ), + # Only the first value of --served-model-name (nargs "+") is exempt: + # a later one starting with "-" is parsed as a flag. + ( + ("--served-model-name", "x", "$(ROUTER_ARG)"), + router_arg_env(), + "has an argument that references $(ROUTER_ARG)", + ), + # After an engine flag that takes no value or optional values, the + # worker would read "--router-mode=kv" as a flag. + ( + ("--enforce-eager", "$(ROUTER_ARG)"), + router_arg_env(), + "has an argument that references $(ROUTER_ARG)", + ), + ( + ("--connector", "$(ROUTER_ARG)"), + router_arg_env(), + "has an argument that references $(ROUTER_ARG)", + ), + ( + ("$(NOTE)",), + [field_ref_env("NOTE", "metadata.annotations['router']")], + "has an argument that references $(NOTE)", + ), + ], +) +def test_p4_rejects_a_worker_that_advertises_a_router_config(args, env, match): + report = evaluate(dgd(frontend(), router_worker(*args, env=env), planner())) + assert not report.p4_ok and not report.spec_ok + assert report.startup is not None # the frontend itself parses + rejected = reasons_matching(report, "P4: worker component agg ") + assert any(match in reason for reason in rejected), report.reasons + assert all("on the frontend only" in reason for reason in rejected) + + +@pytest.mark.parametrize( + ("args", "env"), + [ + # Router flags and env without a mode advertise nothing. + (("--no-router-kv-events", "--router-temperature", "0.5"), None), + ((), [{"name": "DYN_ROUTER_TEMPERATURE", "value": "0.5"}]), + ((), [{"name": "DYN_ROUTER_MODE_HINT", "value": "kv"}]), + # Unreadable env as the value of a non-mode router flag, of an + # option known to take one, or attached to its option. + (("--router-temperature", "$(A_SECRET)"), [{"name": "A_SECRET", **SECRET}]), + (("--speculative-config", "$(A_SECRET)"), [{"name": "A_SECRET", **SECRET}]), + (("--speculative-config", "$(ROUTER_ARG)"), router_arg_env()), + (("--model", "$(ROUTER_ARG)"), router_arg_env()), + (("--model-path", "$(ROUTER_ARG)"), router_arg_env()), + (("--served-model-name", "$(ROUTER_ARG)"), router_arg_env()), + (("--some-opt=$(ROUTER_ARG)",), router_arg_env()), + # Inert env never stands for an option. + (("$(POD_NAME)",), None), + (("--enforce-eager", "$(MY_POD)"), [field_ref_env("MY_POD", "metadata.name")]), + ], +) +def test_p4_accepts_workers_that_advertise_nothing(args, env): + report = evaluate(dgd(frontend(), router_worker(*args, env=env), planner())) + assert report.p4_ok and report.spec_ok, report.reasons + + +@pytest.mark.parametrize( + ("args", "env_from", "match"), + [ + # A Secret or ConfigMap key may define what env leaves undefined. + ( + ("$(ROUTER_ARG)",), + ENV_FROM, + "references $(ROUTER_ARG), which the planner cannot read (the " + "container's env leaves it to an envFrom source)", + ), + ( + ("$(APP_ROUTER_ARG)",), + [{**ENV_FROM[0], "prefix": "APP_"}], + "references $(APP_ROUTER_ARG)", + ), + # A prefix that cannot produce the name, the worker's own env, and + # the value of an option known to take one. + (("$(ROUTER_ARG)",), [{**ENV_FROM[0], "prefix": "APP_"}], None), + (("--model", "$(MODEL_ID)"), ENV_FROM, None), + (("--model", "$(MODEL_NAME)"), ENV_FROM, None), + # The operator sets these on the container, whose env wins over + # envFrom, and none of them holds an option. + (("--draft-endpoint", "$(DYN_NAMESPACE)/backend/generate"), ENV_FROM, None), + (("--x", "$(DYN_COMPONENT)", "$(DYN_SYSTEM_PORT)"), ENV_FROM, None), + # The operator sets this one only for some discovery configurations. + (("$(NATS_SERVER)",), ENV_FROM, "references $(NATS_SERVER)"), + ], +) +def test_p4_worker_reference_an_env_from_source_may_define(args, env_from, match): + component = router_worker( + *args, env=[{"name": "MODEL_ID", "value": MODEL}], env_from=env_from + ) + report = evaluate(dgd(frontend(), component, planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert any(match in reason for reason in rejected), report.reasons + + +@pytest.mark.parametrize( + ("command", "script", "match"), + [ + (("/bin/sh", "-c"), "$(ROUTER_ARG)", "references $(ROUTER_ARG)"), + (None, "$(ROUTER_ARG)", "references $(ROUTER_ARG)"), + # Image env too, which a shell splits unless quoted; a script's own + # assignment overrides what such a source may define. + ( + ("/bin/sh", "-c"), + f"export LD_LIBRARY_PATH=/opt/lib:$LD_LIBRARY_PATH; exec {WORKER}", + "references $LD_LIBRARY_PATH where it could become an option, and " + "the planner cannot read it (the container's env leaves it to an " + "envFrom source)", + ), + ( + ("/bin/sh", "-c"), + f'export LD_LIBRARY_PATH="/opt/lib:$LD_LIBRARY_PATH"; exec {WORKER}', + None, + ), + (("/bin/sh", "-c"), f"ARGS=; exec {WORKER} $ARGS", None), + # A redirection target is a file, never an argument. + (("/bin/sh", "-c"), f'exec {WORKER} > "$HOME/worker.log" 2>&1', None), + # Operator env the container's env defines. + (None, f"echo $DYN_NAMESPACE $DYN_COMPONENT; exec {WORKER}", None), + (None, f"exec {WORKER} --endpoint dyn://$DYN_NAMESPACE.backend.generate", None), + (None, f"echo $DYN_SYSTEM_PORT; exec {WORKER}", None), + ], +) +def test_p4_worker_script_reference_an_env_from_source_may_define( + command, script, match +): + component = shell_worker( + script, + command=command, + env=[{"name": "MODEL_ID", "value": MODEL}], + env_from=ENV_FROM, + ) + report = evaluate(dgd(frontend(), component, planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert any(match in reason for reason in rejected), report.reasons + + +@pytest.mark.parametrize( + "env_from", [ENV_FROM, [{**ENV_FROM[0], "prefix": "APP_"}]], ids=["bare", "APP_"] +) +def test_p4_names_the_env_a_dangerous_script_variable_holds(env_from): + component = shell_worker( + f"X=$ROUTER_ARG; exec {WORKER} $X", env=router_arg_env(), env_from=env_from + ) + report = evaluate(dgd(frontend(), component, planner())) + (reason,) = reasons_matching(report, "P4: worker component agg ") + assert "references $X (assigned from $ROUTER_ARG) where" in reason + assert "envFrom source" not in reason + + +@pytest.mark.parametrize( + ("env", "match"), + [ + (None, None), + # User env of an operator name overrides the operator's. + ( + [ + { + "name": "DYN_NAMESPACE", + "valueFrom": {"configMapKeyRef": {"name": "c", "key": "ns"}}, + } + ], + "references $(DYN_NAMESPACE)", + ), + ], +) +def test_p4_user_env_overrides_operator_env(env, match): + component = router_worker( + "--draft-endpoint", "$(DYN_NAMESPACE)/backend/generate", env=env + ) + report = evaluate(dgd(frontend(), component, planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert any(match in reason for reason in rejected), report.reasons + + +def test_p4_untyped_component_gets_no_worker_operator_env(): + # ComponentDefaultsFactory gives an untyped component the common base + # container only, so an envFrom source may define DYN_SYSTEM_PORT. + component = router_worker("$(DYN_SYSTEM_PORT)", env_from=ENV_FROM, type_="") + report = evaluate(dgd(frontend(), component, planner())) + assert any( + "references $(DYN_SYSTEM_PORT)" in reason + for reason in reasons_matching(report, "P4: worker component agg ") + ), report.reasons + + +def test_p4_checks_an_untyped_worker(): + report = evaluate( + dgd( + frontend(), + router_worker("--router-mode", "kv", name="w", type_=""), + planner(), + ) + ) + assert any( + "router mode 'kv'" in reason + for reason in reasons_matching(report, "P4: worker component w ") + ), report.reasons + + +def test_p4_rejects_graph_env_router_mode_reaching_the_workers(): + # recipes/solar-open2-250b: a graph-level DYN_ROUTER_MODE reaches the + # frontend and every worker. + deployment = dgd( + frontend(args=("-m", "dynamo.frontend")), + worker("prefill", type_="prefill"), + worker("decode", type_="decode"), + planner(), + env=[{"name": "DYN_ROUTER_MODE", "value": "kv"}], + ) + report = evaluate(deployment, config=planner_config("disagg")) + assert report.startup is not None and report.startup.router_mode == "kv_router" + assert not report.p4_ok + assert [r.split(" advertises")[0] for r in reasons_matching(report, "P4:")] == [ + "P4: worker component prefill", + "P4: worker component decode", + ] + + +def test_p4_checks_the_real_recipe_prefill_worker(): + config = planner_config("disagg", model="nvidia/MiniMax-M3-NVFP4") + advertising = dgd( + frontend(), + minimax_prefill_worker( + "--router-mode", "kv", "--no-router-track-output-blocks" + ), + worker("DecodeWorker", type_="decode", model="nvidia/MiniMax-M3-NVFP4"), + planner(), + ) + report = evaluate(advertising, config=config) + assert not report.p4_ok + assert reasons_matching( + report, + "P4: worker component PrefillWorker advertises its own router " + "configuration (router mode 'kv'", + ) + + # Without the router flags the worker inherits the frontend's routing; + # its unreadable --speculative-config value does not matter, and its + # $(MODEL_NAME) / $(SERVED_MODEL_NAME) expand. + inheriting = dgd( + frontend(), + minimax_prefill_worker(), + worker("DecodeWorker", type_="decode", model="nvidia/MiniMax-M3-NVFP4"), + planner(), + ) + report = evaluate(inheriting, config=config) + assert report.p4_ok and report.spec_ok, report.reasons + assert not [w for w in report.warnings if w.startswith("P7:")], report.warnings + assert fc.served_model_name(inheriting) == "nvidia/MiniMax-M3-NVFP4" + + +@pytest.mark.parametrize( + "command", + # The operator's /bin/sh -c applies to an omitted or empty command; a + # launcher may start the shell. + [ + ("/bin/sh", "-c"), + None, + (), + ("/usr/bin/env", "bash", "-c"), + ("tini", "--", "/bin/sh", "-c"), + ("dumb-init", "--", "bash", "-c"), + ("/bin/bash", "-o", "pipefail", "-c"), + ("bash", "-euo", "pipefail", "-c"), + ], + ids=[ + "explicit", + "omitted", + "empty", + "env-bash", + "tini-sh", + "dumb-init-bash", + "bash-o-pipefail", + "bash-euo-pipefail", + ], +) +@pytest.mark.parametrize( + ("script", "match"), + [ + (f"python3 -m dynamo.vllm --model {MODEL} --router-mode kv", "router mode"), + ( + f"python3 -m dynamo.vllm --model {MODEL} --router-mode kv " + "--no-router-kv-events", + "router mode 'kv'", + ), + ( + f"export DYN_ROUTER_MODE=kv && exec python3 -m dynamo.vllm --model {MODEL}", + "sets DYN_ROUTER_MODE in its shell command", + ), + ( + f"DYN_ROUTER_MODE=kv python3 -m dynamo.vllm --model {MODEL}", + "sets DYN_ROUTER_MODE in its shell command", + ), + ('python3 -m dynamo.vllm --router-mode "$MODE"', "cannot parse"), + (f"python3 -m dynamo.vllm --model {MODEL} --no-router-kv-events", None), + ], +) +def test_p4_shell_wrapped_worker(command, script, match): + report = evaluate(dgd(frontend(), shell_worker(script, command=command), planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert not report.p4_ok + assert any(match in reason for reason in rejected), report.reasons + + +@pytest.mark.parametrize( + "command", + # A shell option's operand (pipefail) is no script. + [ + ("/bin/sh", "-c"), + ("bash", "-lc"), + None, + ("/bin/bash", "-o", "pipefail", "-c"), + ("bash", "-euo", "pipefail", "-c"), + ], + ids=["sh", "bash-lc", "omitted", "bash-o-pipefail", "bash-euo-pipefail"], +) +@pytest.mark.parametrize( + ("script", "match"), + [ + ("$(WORKER_CMD)", "has a shell script that references $(WORKER_CMD)"), + (f"{WORKER} $(ROUTER_ARG)", "has a shell script that references $(ROUTER_ARG)"), + # Kubernetes substitutes $(VAR) before the shell parses quotes, so + # even an option's value could carry --router-mode. + ( + "python3 -m dynamo.vllm --speculative-config $(CFG)", + "has a shell script that references $(CFG)", + ), + ( + f"{WORKER} --router-temperature $(CFG)", + "has a shell script that references $(CFG)", + ), + ("python3 -m dynamo.vllm --model $(MODEL_NAME)", None), + # The shell splits an unquoted variable into words. + (f"{WORKER} $ROUTER_ARG", "has a shell script that references $ROUTER_ARG"), + (f"{WORKER} ${{ROUTER_ARG}}", "has a shell script that references $ROUTER_ARG"), + ( + "python3 -m dynamo.vllm --speculative-config $SPECULATIVE_CONFIG", + "has a shell script that references $SPECULATIVE_CONFIG", + ), + ( + f"X=$ROUTER_ARG; exec {WORKER} $X", + "has a shell script that references $X (assigned from $ROUTER_ARG)", + ), + # An assignment may not run, so it clears neither env nor an + # earlier dangerous value. + ( + 'if [ -z "$ROUTER_ARG" ]; then ROUTER_ARG=x; fi; ' + f"exec {WORKER} $ROUTER_ARG", + "has a shell script that references $ROUTER_ARG", + ), + ( + f'[ -n "$ROUTER_ARG" ] || ROUTER_ARG=""; exec {WORKER} $ROUTER_ARG', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'X=$ROUTER_ARG; [ -n "$X" ] || X=; exec {WORKER} $X', + "has a shell script that references $X", + ), + ( + f'[ -n "$OVERRIDE" ] && EXTRA="$OVERRIDE"; exec {WORKER} $EXTRA', + "router mode 'kv'", + ), + # A quoted one stays one word: an option's value, unless it stands + # alone as an argument of the worker. + (f'{WORKER} "$ROUTER_ARG"', "has a shell script that references $ROUTER_ARG"), + ( + f'exec env A=1 {WORKER} "${{ROUTER_ARG:-}}"', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'exec python3 -um dynamo.vllm --model {MODEL} "$ROUTER_ARG"', + "has a shell script that references $ROUTER_ARG", + ), + # eval and a shell's -c parse their arguments again, also behind a + # launcher and within a larger word. + ('eval "$ROUTER_ARG"', "has a shell script that references $ROUTER_ARG"), + ( + f'eval "{WORKER} $ROUTER_ARG"', + "has a shell script that references $ROUTER_ARG", + ), + ('exec sh -c "$ROUTER_ARG"', "has a shell script that references $ROUTER_ARG"), + ( + f'exec sh -c "{WORKER} $ROUTER_ARG"', + "has a shell script that references $ROUTER_ARG", + ), + ( + f"exec bash -c 'exec {WORKER} $@' bash \"$ROUTER_ARG\"", + "has a shell script that references $ROUTER_ARG", + ), + # source, env, export and a shell turn one quoted word into code or + # env, also by path and behind exec, nohup, command or builtin. + ( + f'env "$ROUTER_ARG" {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'exec env "$ROUTER_ARG" {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'exec env A=1 "$ROUTER_ARG" {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'/usr/bin/env "$ROUTER_ARG" {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'nohup env "$ROUTER_ARG" {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + f'command export "$ROUTER_ARG"; exec {WORKER}', + "has a shell script that references $ROUTER_ARG", + ), + ( + 'command eval "$ROUTER_ARG"', + "has a shell script that references $ROUTER_ARG", + ), + (f'cp -r /cfg/. "$ROUTER_ARG"; exec {WORKER}', None), + # A redirection target names a file, unless a here-string feeds + # its text to the command. + (f'exec {WORKER} > "$CFG" 2>&1', None), + (f"exec {WORKER} 2>$CFG", None), + ( + f'exec {WORKER} > "$CFG" "$ROUTER_ARG"', + "has a shell script that references $ROUTER_ARG", + ), + ("exec bash <<< $ROUTER_ARG", "has a shell script that references $ROUTER_ARG"), + # The value of an option known to take one. + ('exec python3 -m dynamo.vllm --model "$ROUTER_ARG"', None), + (f'{WORKER} --speculative-config="${{SPECULATIVE_CONFIG}}"', None), + (f'{WORKER} --speculative-config "$SPECULATIVE_CONFIG"', None), + (f'{WORKER} "--some-opt=$ROUTER_ARG"', None), + (f'[ -n "$ROUTER_ARG" ] && echo "$ROUTER_ARG"; exec {WORKER}', None), + (f'export LD_LIBRARY_PATH="/opt/lib:$ROUTER_ARG"; exec {WORKER}', None), + # Literal env reached through the shell is read, and so are the + # default and alternate words the shell may use. + (f"{WORKER} $EXTRA", "router mode 'kv'"), + (f'{WORKER} --router-mode "${{MODE:-random}}"', "router mode 'kv'"), + (f"{WORKER} ${{MODE:+--router-mode=kv}}", "router mode 'kv'"), + ( + f"{WORKER} ${{UNSET:-$ROUTER_ARG}}", + "has a shell script that references $ROUTER_ARG", + ), + # Inert env never stands for an option. + (f'{WORKER} $POD_NAME "$MY_POD" $(MY_POD) $(POD_UID)', None), + (f'[ "$POD_INDEX" = "0" ] && exec {WORKER}', None), + (f"{WORKER} $NOTE", "has a shell script that references $NOTE"), + # Neither evaluated nor rejected: command substitution and script + # variables built from it. + (f"{WORKER} $(cat /etc/router-args)", None), + (f"ARGS=`cat /etc/router-args`; {WORKER} $ARGS", None), + ], +) +def test_p4_shell_wrapped_worker_with_unreadable_env(command, script, match): + env = [ + {"name": name, "valueFrom": {"configMapKeyRef": {"name": "c", "key": "k"}}} + for name in ("CFG", "ROUTER_ARG", "SPECULATIVE_CONFIG", "WORKER_CMD") + ] + env += [ + {"name": "MODEL_NAME", "value": MODEL}, + {"name": "EXTRA", "value": "--router-mode kv"}, + {"name": "MODE", "value": "kv"}, + field_ref_env("MY_POD", "metadata.name"), + field_ref_env("POD_INDEX", "metadata.labels['grove.io/podclique-pod-index']"), + field_ref_env("NOTE", "metadata.annotations['router']"), + ] + report = evaluate( + dgd(frontend(), shell_worker(script, command=command, env=env), planner()) + ) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert not report.p4_ok + assert any(match in reason for reason in rejected), report.reasons + assert all("on the frontend only" in reason for reason in rejected) + + +@pytest.mark.parametrize( + ("command", "args", "match"), + [ + # A launcher drops its "--", so the worker still sees later flags. + ( + ("tini", "--", "python3", "-m", "dynamo.vllm"), + ("--router-mode", "kv"), + "router mode 'kv'", + ), + ( + ("dumb-init", "--"), + ("python3", "-m", "dynamo.vllm", "--router-mode", "kv"), + "router mode 'kv'", + ), + ( + ("/bin/sh", "-c"), + ("exec tini -- python3 -m dynamo.vllm --router-mode kv",), + "router mode 'kv'", + ), + (("tini", "--", "python3", "-m", "dynamo.vllm"), ("--model", MODEL), None), + ], +) +def test_p4_worker_behind_a_launcher(command, args, match): + component = worker() + component["podTemplate"]["spec"]["containers"][0].update( + command=list(command), args=list(args) + ) + report = evaluate(dgd(frontend(), component, planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if match is None: + assert report.p4_ok, report.reasons + else: + assert any(match in reason for reason in rejected), report.reasons + + +MODEL_FROM_CONFIG_MAP = [ + {"name": "MODEL", "valueFrom": {"configMapKeyRef": {"name": "c", "key": "model"}}} +] +RUNTIME_SIDECAR = {"name": "runtime", "restartPolicy": "Always"} + + +@pytest.mark.parametrize( + ("command", "args"), + [ + (None, ["--model-path", "$(MODEL)", "--tp", "2"]), + (["python3", "-m", "tensorrt_llm.commands.serve"], ["$(MODEL)"]), + (None, ["--model-path=$(MODEL)", "--router-mode", "kv"]), + ], +) +@pytest.mark.parametrize("leader_only", [False, True], ids=["component", "leader"]) +def test_p4_skips_native_sidecar_mains(command, args, leader_only): + # The main runs the engine with its image entrypoint, and the runtime + # sidecar registers no router configuration. + component = worker() + spec = component["podTemplate"]["spec"] + main = spec["containers"][0] + main.update(args=args, env=MODEL_FROM_CONFIG_MAP) + if command is None: + del main["command"] + else: + main["command"] = command + if leader_only: + component["multinode"] = {"nodeCount": 2} + component["roles"] = [ + { + "name": "leader", + "podTemplate": {"spec": {"initContainers": [RUNTIME_SIDECAR]}}, + } + ] + else: + spec["initContainers"] = [RUNTIME_SIDECAR] + report = evaluate(dgd(frontend(), component, planner())) + assert report.p4_ok, report.reasons + + +def test_p4_reads_an_omitted_command_without_a_native_sidecar_as_a_script(): + component = worker() + main = component["podTemplate"]["spec"]["containers"][0] + del main["command"] + main.update(args=["--model-path=$(MODEL)", "--tp", "2"], env=MODEL_FROM_CONFIG_MAP) + report = evaluate(dgd(frontend(), component, planner())) + (reason,) = reasons_matching(report, "P4: worker component agg ") + # Quoting or attaching never helps a $(VAR): Kubernetes substitutes it + # before the shell parses the script. + assert "has a shell script that references $(MODEL), which Kubernetes" in reason + assert '(--option="$MODEL")' in reason + assert "quote the reference" not in reason + + +def role_worker( + leader: dict[str, Any], worker_: dict[str, Any], **extra: Any +) -> dict[str, Any]: + """A multinode worker whose launches are only in its role pod templates.""" + + roles = [] + for name, main in (("leader", leader), ("worker", worker_)): + container = {"name": "main", "command": ["python3", "-m", "dynamo.vllm"]} + container.update(main) + roles.append( + {"name": name, "podTemplate": {"spec": {"containers": [container]}}} + ) + return { + "name": "agg", + "type": "worker", + "replicas": 1, + "multinode": {"nodeCount": 2}, + "roles": roles, + **extra, + } + + +ROLE_ARGS = ["--model", MODEL, "--tensor-parallel-size", "16"] + + +@pytest.mark.parametrize( + ("leader", "worker_", "match"), + [ + ( + {"args": [*ROLE_ARGS, "--router-mode", "kv"]}, + {"args": ROLE_ARGS}, + "agg (role leader) advertises its own router configuration", + ), + ( + {"args": ROLE_ARGS}, + {"args": ROLE_ARGS, "env": [{"name": "DYN_ROUTER_MODE", "value": "kv"}]}, + "agg (role worker) advertises its own router configuration", + ), + ( + {"args": [*ROLE_ARGS, "$(ROUTER_ARG)"], "env": router_arg_env()}, + {"args": ROLE_ARGS}, + "agg (role leader) has an argument that references $(ROUTER_ARG)", + ), + ({"args": ROLE_ARGS}, {"args": ROLE_ARGS}, None), + # The operator sets the rank, so envFrom cannot. + ( + {"args": ROLE_ARGS, "envFrom": ENV_FROM}, + { + "args": [*ROLE_ARGS, "--node-rank", "$(DYNAMO_RANK)"], + "envFrom": ENV_FROM, + }, + None, + ), + ], +) +def test_p4_checks_each_role_pod_template(leader, worker_, match): + report = evaluate(dgd(frontend(), role_worker(leader, worker_), planner())) + if match is None: + assert report.p4_ok, report.reasons + assert fc.served_model_names(dgd(role_worker(leader, worker_)))[0] == { + "agg": MODEL + } + else: + assert any( + match in reason + for reason in reasons_matching(report, "P4: worker component ") + ), report.reasons + + +def test_worker_check_ignores_the_planner_env(monkeypatch): + monkeypatch.setenv("DYN_ROUTER_MODE", "kv") + fc._worker_router_cached.cache_clear() + try: + report = evaluate(dgd()) + finally: + fc._worker_router_cached.cache_clear() + assert report.p4_ok, report.reasons + assert os.environ["DYN_ROUTER_MODE"] == "kv" + + +WORKER_GOLDEN_CASES = [ + ((), {}), + (("--router-mode", "kv"), {}), + (("--router-m", "random"), {}), + (("--router-mode=power-of-two",), {}), + ((), {"DYN_ROUTER_MODE": "kv"}), + (("--router-mode", "least-loaded"), {"DYN_ROUTER_MODE": "kv"}), + (("--no-router-kv-events", "--router-temperature", "0.5"), {}), + ((), {"DYN_ROUTER_TEMPERATURE": "0.5"}), + # What an unreadable "$(ROUTER_ARG)" after --no-router-kv-events can be. + (("--no-router-kv-events", "--router-mode=kv"), {}), +] + + +@pytest.mark.parametrize(("args", "env"), WORKER_GOLDEN_CASES) +def test_worker_router_mode_matches_the_real_worker_parser(monkeypatch, args, env): + """Backends parse their left-over argv with parse_worker_router_config.""" + + from dynamo.common.configuration.groups.router_args import ( + parse_worker_router_config, + ) + + for name in list(os.environ): + if name.startswith("DYN_"): + monkeypatch.delenv(name) + for name, value in env.items(): + monkeypatch.setenv(name, value) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + config, _ = parse_worker_router_config([f"--model={MODEL}", *args]) + for name in env: + monkeypatch.delenv(name) + + worker_env = [{"name": name, "value": value} for name, value in env.items()] + report = evaluate(dgd(frontend(), router_worker(*args, env=worker_env), planner())) + rejected = reasons_matching(report, "P4: worker component agg ") + if config.router_mode is None: + assert rejected == [], rejected + else: + assert any( + f"router mode {config.router_mode!r}" in reason for reason in rejected + ), rejected + + +def test_worker_env_from_is_warned_once(): + deployment = dgd(frontend(), worker("a"), worker("b"), planner()) + for component in deployment["spec"]["components"]: + component["podTemplate"]["spec"]["containers"][0]["envFrom"] = [ + {"secretRef": {"name": "hf-token-secret"}} + ] + report = evaluate(deployment) + assert report.p4_ok and report.spec_ok, report.reasons + env_from = sorted(w for w in report.warnings if "envFrom" in w) + assert env_from == [ + "P4: the frontend container uses envFrom secretRef 'hf-token-secret'; " + "the planner cannot read it and assumes it sets no router env (P14 " + "still checks the routers' policy parameters)", + "P4: worker components a, b use envFrom; the planner cannot read it and " + "assumes it sets no router env (a DYN_ROUTER_MODE there would make the " + "worker advertise its own router configuration)", + ] + + +# ---------------------------------------------------------- P6-P10 + + +def test_p6_warns_on_unknown_model(): + report = evaluate(dgd(), config=planner_config(model="acme/private-model")) + assert report.spec_ok + assert any(w.startswith("P6:") for w in report.warnings) + + +@pytest.fixture +def uncached_model_warning(): + fc._model_warning.cache_clear() + yield + fc._model_warning.cache_clear() + + +@pytest.mark.usefixtures("uncached_model_warning") +def test_p6_without_aisimulate_is_not_checked(monkeypatch): + monkeypatch.setitem(sys.modules, "aisimulate_core.sdk.common", None) + assert fc._model_warning("acme/private-model") is None + + +@pytest.mark.usefixtures("uncached_model_warning") +def test_p6_broken_aisimulate_warns(monkeypatch, caplog): + def broken(name: str) -> Any: + raise RuntimeError("broken install") + + module = types.ModuleType("aisimulate_core.sdk.common") + module.__getattr__ = broken # type: ignore[method-assign] + monkeypatch.setitem(sys.modules, "aisimulate_core.sdk.common", module) + warning = fc._model_warning(MODEL) + assert warning is not None and warning.startswith("P6: cannot check") + assert "RuntimeError: broken install" in warning + assert any(r.levelname == "ERROR" for r in caplog.records) + + +def test_p7_served_models(): + two = dgd(frontend(), worker("a"), worker("b", model="meta/llama"), planner()) + report = evaluate(two) + assert reasons_matching(report, "P7:") == [ + f"P7: the DGD serves 2 models ({MODEL}, meta/llama); live retune needs one" + ] + assert fc.served_model_name(two) is None + assert fc.served_model_name(dgd()) == MODEL + + +def test_p7_served_model_prefers_current_status_and_warns_when_unknown(): + unknown = worker("b") + unknown["podTemplate"]["spec"]["containers"][0]["args"] = ["--model", "$(MODEL)"] + deployment = dgd(frontend(), worker("a"), unknown, planner()) + report = evaluate(deployment) + assert report.spec_ok + assert any("served model of b" in w for w in report.warnings) + + deployment["metadata"]["generation"] = 3 + deployment["status"] = { + "observedGeneration": 3, + "components": {"b": {"servedModelName": "other/model"}}, + } + assert fc.served_model_names(deployment) == ({"a": MODEL, "b": "other/model"}, ()) + deployment["status"]["observedGeneration"] = 2 + assert fc.served_model_names(deployment)[1] == ("b",) + + +def test_p7_shell_wrapped_worker_model(): + shell = worker("a") + shell["podTemplate"]["spec"]["containers"][0].update( + command=["/bin/sh", "-c"], + args=[ + f"python3 -m dynamo.sglang --model-path {MODEL} --served-model-name qwen" + ], + ) + assert fc.served_model_names(dgd(frontend(), shell)) == ({"a": "qwen"}, ()) + + +def p7_model_warnings(report: PremiseReport) -> list[str]: + return [w for w in report.warnings if w.startswith("P7: retune.dgdr.model")] + + +def test_p7_warns_when_dgdr_model_is_not_what_the_workers_serve(): + deployment = dgd(frontend(), worker(model="meta-llama/Llama-3.1-8B"), planner()) + report = evaluate(deployment) + assert report.p4_ok and report.spec_ok, report.reasons + assert p7_model_warnings(report) == [ + f"P7: retune.dgdr.model {MODEL!r} is not among the workers' " + "--model/--model-path/--served-model-name values " + "(meta-llama/Llama-3.1-8B); every DGDR Run simulates retune.dgdr.model, " + "so make sure it is the model the DGD serves" + ] + assert not p7_model_warnings(evaluate(dgd())) + + +@pytest.mark.parametrize( + "args", + [ + # An alias: the HF id is still among the values. + ["--model", "/models/qwen", "--served-model-name", MODEL], + ["--model-path", MODEL, "--served-model-name", "qwen"], + ["--model", "/models/qwen", f"--served-model-name=qwen,{MODEL}"], + # Expanded against the worker's env. + ["--model", "$(MODEL_ID)"], + ], +) +def test_p7_dgdr_model_matches_any_worker_model_value(args): + component = worker() + component["podTemplate"]["spec"]["containers"][0].update( + args=args, env=[{"name": "MODEL_ID", "value": MODEL}] + ) + report = evaluate(dgd(frontend(), component, planner())) + assert not p7_model_warnings(report), report.warnings + + +def test_p7_dgdr_model_counts_the_current_served_model_status(): + component = worker() + component["podTemplate"]["spec"]["containers"][0]["args"] = ["--model", "$(X)"] + deployment = dgd(frontend(), component, planner()) + deployment["metadata"]["generation"] = 1 + deployment["status"] = { + "observedGeneration": 1, + "components": {"agg": {"servedModelName": "other/model"}}, + } + assert p7_model_warnings(evaluate(deployment)) + deployment["status"]["components"]["agg"]["servedModelName"] = MODEL + assert not p7_model_warnings(evaluate(deployment)) + # Nothing known: only the undetermined-model warning. + del deployment["status"] + report = evaluate(deployment) + assert not p7_model_warnings(report) + assert any("served model of agg" in w for w in report.warnings) + + +@pytest.mark.parametrize( + ("extra", "match"), + [ + ({"name": "Epp", "type": "epp"}, "EPP component"), + ( + worker("side", frontendSidecar="sidecar-frontend"), + "side run a frontendSidecar", + ), + ], +) +def test_p7_gaie_layouts_with_a_frontend(extra, match): + report = evaluate(dgd(frontend(), worker(), planner(), extra)) + assert report.p4_ok + assert any(match in r for r in reasons_matching(report, "P7:")), report.reasons + + +def test_p7_no_frontend(): + report = evaluate(dgd(worker(frontendSidecar="sidecar-frontend"), planner())) + assert not report.p4_ok + assert all(r.startswith("P7:") for r in report.reasons) + assert len(report.reasons) == 1 + + +@pytest.mark.parametrize( + ("args", "env"), + [ + ((*KV_ARGS, "--tls-cert-path", "/c.pem", "--tls-key-path", "/k.pem"), {}), + ((*KV_ARGS,), {"DYN_TLS_CERT_PATH": "/c.pem", "DYN_TLS_KEY_PATH": "/k.pem"}), + ], +) +def test_p7_tls(args, env): + report = evaluate(with_frontend(args=args, env=trace_env(**env))) + assert report.p4_ok + assert reasons_matching(report, "P7:") == [ + "P7: the frontend serves TLS (--tls-cert-path / DYN_TLS_CERT_PATH)" + ] + + +@pytest.mark.parametrize( + ("replicas", "args", "ok"), + [ + (1, KV_ARGS, True), + (0, KV_ARGS, True), + (2, KV_ARGS, False), + (2, (*KV_ARGS, "--router-replica-sync"), True), + ], +) +def test_p8_frontend_replicas(replicas, args, ok): + report = evaluate(dgd(frontend(replicas=replicas, args=args), worker(), planner())) + assert bool(reasons_matching(report, "P8:")) is not ok + assert report.p4_ok + + +def test_p9_scaling_adapter_and_planner_replicas(): + adapter = dgd(frontend(), worker(scalingAdapter={}), planner()) + assert reasons_matching(evaluate(adapter), "P9:") == [ + "P9: worker component agg has a scalingAdapter; its replicas could change during a window" + ] + two_planners = dgd(frontend(), worker(), planner(replicas=2)) + assert "2 replicas" in reasons_matching(evaluate(two_planners), "P9:")[0] + no_planner = evaluate(dgd(frontend(), worker())) + assert no_planner.spec_ok + assert any(w.startswith("P9:") for w in no_planner.warnings) + + +P10_CASES = [ + ({"DYN_REQUEST_TRACE": "0"}, None, "DYN_REQUEST_TRACE=1"), + ({"DYN_REQUEST_TRACE": "true"}, None, None), + ({"DYN_REQUEST_TRACE_SINKS": "stderr"}, None, "DYN_REQUEST_TRACE_SINKS=file"), + ( + {"DYN_REQUEST_TRACE_SINKS": "file,jsonl_gz"}, + None, + "DYN_REQUEST_TRACE_SINKS=file", + ), + ({"DYN_REQUEST_TRACE_SINKS": "file,stderr"}, None, None), + ({"DYN_REQUEST_TRACE_FILE_FORMAT": None}, None, None), + ({"DYN_REQUEST_TRACE_FILE_FORMAT": "jsonl"}, None, "DYN_REQUEST_TRACE_FILE_FORMAT"), + ( + {"DYN_REQUEST_TRACE_RECORDS": None}, + None, + "DYN_REQUEST_TRACE_RECORDS=request_end", + ), + ( + {"DYN_REQUEST_TRACE_RECORDS": "request_end,tool"}, + None, + "DYN_REQUEST_TRACE_RECORDS", + ), + ({"DYN_REQUEST_TRACE_RECORDS": " Request_End "}, None, None), + ( + {"DYN_REQUEST_TRACE_FILE_PATH": None}, + None, + "must set DYN_REQUEST_TRACE_FILE_PATH", + ), + ( + {"DYN_REQUEST_TRACE_FILE_PATH": f"/trace/segments/{DGD}/$(DYN_POD_NAME)/req"}, + None, + None, + ), + # The reconciler prunes only the top level of the pod's directory. + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/a/req" + }, + None, + "no further '/'", + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req-$(DYN_POD_NAME)" + }, + None, + None, + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/" + }, + None, + "file prefix", + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(POD_NAME)/req" + }, + None, + "file prefix", + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/trace/other/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req" + }, + None, + "file prefix", + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": "/data/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req" + }, + None, + "file prefix", + ), + ( + { + "DYN_REQUEST_TRACE_FILE_PATH": { + "valueFrom": {"configMapKeyRef": {"name": "c", "key": "p"}} + } + }, + None, + "must be a literal value", + ), + ( + {"DYN_POD_NAME": {"valueFrom": {"fieldRef": {"fieldPath": "metadata.uid"}}}}, + None, + "DYN_POD_NAME from fieldRef", + ), + ({"DYN_PARENT_DGD_K8S_NAME": "other"}, None, "overrides DYN_PARENT_DGD_K8S_NAME"), + ( + {}, + {"volumes": [{"name": "traces", "persistentVolumeClaim": {"claimName": "x"}}]}, + "no volume for claim", + ), + ({}, {"mounts": []}, "does not mount claim"), + ( + {}, + {"mounts": [{"name": "traces", "mountPath": "/trace", "readOnly": True}]}, + "read-only or through a subPath", + ), + ({}, {"mounts": [{"name": "traces", "mountPath": "/trace/"}]}, None), + ( + {}, + { + "mounts": [ + {"name": "traces", "mountPath": "/ro", "readOnly": True}, + {"name": "traces", "mountPath": "/trace"}, + ] + }, + None, + ), +] + + +@pytest.mark.parametrize(("env", "frontend_kwargs", "match"), P10_CASES) +def test_p10_trace_recipe(env, frontend_kwargs, match): + deployment = dgd( + frontend(env=trace_env(**env), **(frontend_kwargs or {})), worker(), planner() + ) + report = evaluate(deployment) + assert report.p4_ok, report.reasons + p10 = reasons_matching(report, "P10:") + if match is None: + assert p10 == [], p10 + assert report.spec_ok + else: + assert any(match in r for r in p10), p10 + assert not report.spec_ok + + +def test_p10_custom_segments_dir_and_graph_env(): + path = "/trace/rt/segs/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req" + deployment = dgd( + frontend(env=trace_env(DYN_REQUEST_TRACE_FILE_PATH=None)), + worker(), + planner(), + env=[{"name": "DYN_REQUEST_TRACE_FILE_PATH", "value": path}], + ) + report = evaluate(deployment, config=planner_config(segments_dir="rt/segs")) + assert report.spec_ok, report.reasons + + +def test_p10_extra_sinks_warn(): + report = evaluate( + with_frontend(env=trace_env(DYN_REQUEST_TRACE_SINKS="file,stderr")) + ) + assert report.spec_ok + assert any("also enables stderr" in w for w in report.warnings) + + +def test_p10_is_checked_when_parsing_fails(): + deployment = with_frontend( + command=("/bin/sh", "-c"), + args=("python3 -m dynamo.frontend",), + env=trace_env(DYN_REQUEST_TRACE=None), + ) + report = evaluate(deployment) + assert reasons_matching(report, "P4:") + assert reasons_matching(report, "P10:") + + +def test_evaluate_requires_dgdr(): + config = PlannerConfig(namespace="test-ns", environment="kubernetes", mode="agg") + with pytest.raises(ValueError, match="retune.dgdr"): + fc.evaluate_premises(dgd(), planner_config=config) + + +def test_malformed_dgd_never_raises(): + odd_frontend = { + "name": "f", + "type": "frontend", + "replicas": "many", + "podTemplate": "not-a-template", + } + for deployment in ( + {}, + {"spec": None}, + {"spec": {"components": [None, 3]}}, + {"spec": {"components": [odd_frontend]}}, + {"spec": {"components": [frontend(volumes=["x"], mounts=[None])]}}, + ): + report = evaluate(deployment) + assert not report.spec_ok + + +def test_unexpected_parser_failure_fails_closed(monkeypatch, caplog): + def broken(env): + raise RuntimeError("frontend arg groups changed") + + monkeypatch.setattr(fc, "_frontend_parser", broken) + fc._parse_cached.cache_clear() + try: + error = parse_error(dgd()) + finally: + fc._parse_cached.cache_clear() + assert error.startswith("P4: the planner failed to parse") + assert "RuntimeError: frontend arg groups changed" in error + assert any(r.levelname == "ERROR" for r in caplog.records) + + +# ------------------------------------------------------------- helpers + + +def test_worker_types_and_canonical_json(): + assert fc.worker_types("disagg") == ("prefill", "decode") + assert fc.worker_types("agg") == ("aggregated",) + assert fc.canonical_json({"b": 1, "a": [1.5]}) == '{"a":[1.5],"b":1}' + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("plain", "plain"), + ("$(A)/x", "1/x"), + ("$$(A)", "$(A)"), + ("$(MISSING)", "$(MISSING)"), + ("$(A", "$(A"), + ("$", "$"), + ("a$$b", "a$b"), + ("x-$(SECRET)", "x-{opaque}$(SECRET)\x00"), + ("$(POD)", "{inert}$(POD)\x00"), + # Undefined names an envFrom source with prefix APP_ may define. + ("$(APP_ARG)", "{opaque}$(APP_ARG)\x00"), + ("$(APP_a.b-c)", "{opaque}$(APP_a.b-c)\x00"), + # The container's env wins over envFrom. + ("$(APP_DEFINED)", "2"), + ("$(ARG)", "$(ARG)"), + ("$(APP_ seq 1 60)", "$(APP_ seq 1 60)"), + ], +) +def test_kubernetes_expansion(value, expected): + expanded = fc._expand( + value, + {"A": "1", "APP_DEFINED": "2"}, + {"SECRET", "POD"}, + inert={"POD"}, + env_from_prefixes=("APP_",), + ) + assert expanded == expected.format(opaque=fc._OPAQUE, inert=fc._INERT) + + +def test_parse_is_cached_per_spec(): + fc._parse_cached.cache_clear() + parse(dgd()) + parse(dgd()) + info = fc._parse_cached.cache_info() + assert info.hits >= 1 and info.misses == 1 + parse(with_frontend(args=(*KV_ARGS, "--router-temperature", "0.5"))) + assert fc._parse_cached.cache_info().misses == 2 diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_kube.py b/components/src/dynamo/planner/tests/unit/test_live_retune_kube.py new file mode 100644 index 000000000000..c85fb01dba1b --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_kube.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for ``KubernetesRetuneKube``. + +The tests run the real ``kubernetes`` API classes and replace only +the REST transport, so they check the URLs, query parameters, timeouts and +bodies the client library actually sends, and the error mapping. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any, Optional + +import pytest +from kubernetes import client +from kubernetes.client import ApiException + +from dynamo.planner.retune.live.contracts import ( + KubeAlreadyExists, + KubeConflict, + KubeError, + KubeInvalid, + KubeNotFound, + PodInfo, + RetuneKube, +) +from dynamo.planner.retune.live.kube import KubernetesRetuneKube, request_timeout + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +NS = "test-ns" + + +class FakeResponse: + def __init__(self, body: Any, status: int = 200) -> None: + self.status = status + self.reason = "OK" + self.data = body if isinstance(body, (str, bytes)) else json.dumps(body) + + def getheaders(self) -> dict: + return {"content-type": "application/json"} + + def getheader(self, name: str, default: Optional[str] = None) -> Optional[str]: + return self.getheaders().get(name.lower(), default) + + +class Transport: + """Stands in for ``RESTClientObject.request``: records and replies.""" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + self.replies: list[Any] = [] + + def reply(self, value: Any) -> None: + self.replies.append(value) + + def __call__( + self, + method, + url, + query_params=None, + headers=None, + post_params=None, + body=None, + _preload_content=True, + _request_timeout=None, + ): + self.calls.append( + { + "method": method, + "url": url, + "query": dict(query_params or []), + "body": body, + "preload": _preload_content, + "timeout": _request_timeout, + } + ) + value = self.replies.pop(0) if self.replies else {} + if isinstance(value, Exception): + raise value + if isinstance(value, FakeResponse): + return value + return FakeResponse(value) + + +def api_error(status: int, message: str = "nope") -> ApiException: + error = ApiException(status=status, reason="Error") + error.body = json.dumps({"kind": "Status", "message": message}) + return error + + +@pytest.fixture +def kube(): + api_client = client.ApiClient(client.Configuration()) + transport = Transport() + api_client.rest_client.request = transport + retune_kube = KubernetesRetuneKube(NS, api_client=api_client) + yield retune_kube, transport + api_client.close() + + +def path(call: dict) -> str: + return call["url"].split("://", 1)[-1].split("/", 1)[-1] + + +# ---------------------------------------------------------------- client + + +def test_request_timeout_is_a_tuple_of_whole_seconds(): + assert request_timeout(10.0) == (10, 10) + assert request_timeout(2.5) == (3, 3) + assert request_timeout(0.2) == (1, 1) + assert request_timeout(0.0) == (1, 1) + assert request_timeout(float("nan")) == (1, 1) + + +def test_satisfies_the_protocol(kube): + retune_kube, _ = kube + assert isinstance(retune_kube, KubernetesRetuneKube) + _: RetuneKube = retune_kube + assert retune_kube.namespace == NS + with pytest.raises(ValueError): + KubernetesRetuneKube("", api_client=retune_kube.api_client) + + +def test_get_dgd(kube): + retune_kube, transport = kube + dgd = {"metadata": {"name": "llm", "uid": "u1"}, "spec": {"services": {}}} + transport.reply(dgd) + assert retune_kube.get_dgd("llm", timeout_s=4) == dgd + (call,) = transport.calls + assert call["method"] == "GET" + assert ( + path(call) + == f"apis/nvidia.com/v1beta1/namespaces/{NS}/dynamographdeployments/llm" + ) + assert call["timeout"] == (4, 4) + + transport.reply(api_error(404)) + with pytest.raises(KubeNotFound) as raised: + retune_kube.get_dgd("gone") + assert raised.value.status == 404 + + +def test_create_request_is_strict_and_dry_run_on_request(kube): + retune_kube, transport = kube + body = {"metadata": {"name": "r1"}, "spec": {"model": {"name": "m"}}} + transport.reply({**body, "metadata": {"name": "r1", "uid": "x"}}) + created = retune_kube.create_request(body, timeout_s=2.5) + assert created["metadata"]["uid"] == "x" + call = transport.calls[-1] + assert call["method"] == "POST" + assert ( + path(call) + == f"apis/nvidia.com/v1beta2/namespaces/{NS}/dynamographdeploymentrequests" + ) + assert call["query"]["fieldValidation"] == "Strict" + assert call["query"]["fieldManager"] == "dynamo-planner-retune" + assert "dryRun" not in call["query"] + assert call["body"] == body + assert call["timeout"] == (3, 3) + + transport.reply(body) + retune_kube.create_request(body, dry_run=True) + assert transport.calls[-1]["query"]["dryRun"] == "All" + assert transport.calls[-1]["query"]["fieldValidation"] == "Strict" + + +@pytest.mark.parametrize( + "status, error_type", + [ + (409, KubeAlreadyExists), + (422, KubeInvalid), + (400, KubeInvalid), + (403, KubeError), + (500, KubeError), + ], +) +def test_create_request_maps_api_errors(kube, status, error_type): + retune_kube, transport = kube + transport.reply(api_error(status, "strict decoding error: unknown field")) + with pytest.raises(error_type) as raised: + retune_kube.create_request({"metadata": {"name": "r1"}}) + assert raised.value.status == status + assert f"HTTP {status}" in str(raised.value) + assert "unknown field" in str(raised.value) + if error_type is KubeError: + assert type(raised.value) is KubeError + + +def test_network_errors_become_kube_errors(kube): + retune_kube, transport = kube + # Imported here: the marker-report hook collects tests without urllib3. + from urllib3.exceptions import ReadTimeoutError + + transport.reply(ReadTimeoutError(None, "url", "timed out")) + with pytest.raises(KubeError) as raised: + retune_kube.get_request("r1") + assert raised.value.status is None + assert "ReadTimeoutError" in str(raised.value) + + +def test_requests_runs_and_candidates(kube): + retune_kube, transport = kube + transport.reply({"metadata": {"name": "r1"}}) + assert retune_kube.get_request("r1") == {"metadata": {"name": "r1"}} + transport.reply(api_error(404)) + assert retune_kube.get_request("missing") is None + + transport.reply({"items": [{"metadata": {"name": "r1"}}]}) + assert retune_kube.list_requests("a=b,c=d") == [{"metadata": {"name": "r1"}}] + call = transport.calls[-1] + assert ( + path(call) + == f"apis/nvidia.com/v1beta2/namespaces/{NS}/dynamographdeploymentrequests" + ) + assert call["query"]["labelSelector"] == "a=b,c=d" + + transport.reply({"metadata": {"name": "r1-1"}}) + assert retune_kube.get_run("r1-1") == {"metadata": {"name": "r1-1"}} + assert path(transport.calls[-1]).endswith("dynamographdeploymentruns/r1-1") + transport.reply(api_error(404)) + assert retune_kube.get_run("gone") is None + + transport.reply({"items": []}) + assert retune_kube.list_candidates("nvidia.com/dgdr-run-uid=u") == [] + call = transport.calls[-1] + assert path(call).endswith("dynamographdeploymentcandidates") + assert call["query"]["labelSelector"] == "nvidia.com/dgdr-run-uid=u" + + transport.reply({}) + retune_kube.delete_request("r1", propagation="Foreground", timeout_s=7) + call = transport.calls[-1] + assert call["method"] == "DELETE" + assert path(call).endswith("dynamographdeploymentrequests/r1") + assert call["query"]["propagationPolicy"] == "Foreground" + assert call["timeout"] == (7, 7) + transport.reply(api_error(404)) + retune_kube.delete_request("gone") # missing is not an error + transport.reply(api_error(500)) + with pytest.raises(KubeError): + retune_kube.delete_request("r1") + + +def _pod( + name: str, + *, + ready: bool = True, + deleting: bool = False, + containers: Optional[list] = None, + ip: Optional[str] = "10.0.0.5", +) -> dict: + return { + "metadata": { + "name": name, + "uid": f"uid-{name}", + "labels": {"nvidia.com/dynamo-component-type": "frontend"}, + **({"deletionTimestamp": "2026-10-28T00:00:00Z"} if deleting else {}), + }, + "spec": { + "containers": containers + if containers is not None + else [ + {"name": "sidecar", "ports": [{"name": "http", "containerPort": 9999}]}, + { + "name": "main", + "ports": [ + {"name": "metrics", "containerPort": 9090}, + {"name": "http", "containerPort": 8000}, + ], + }, + ] + }, + "status": { + "podIP": ip, + "conditions": [ + {"type": "PodScheduled", "status": "True"}, + {"type": "Ready", "status": "True" if ready else "False"}, + ], + "containerStatuses": [ + {"name": "sidecar", "restartCount": 7, **_CONTAINER_STATUS}, + {"name": "main", "restartCount": 2, **_CONTAINER_STATUS}, + ], + }, + } + + +# The fields V1ContainerStatus requires besides name and restartCount. +_CONTAINER_STATUS = {"image": "img", "imageID": "", "ready": True} + + +def test_list_pods_builds_pod_info(kube): + retune_kube, transport = kube + transport.reply( + { + "kind": "PodList", + "apiVersion": "v1", + "metadata": {}, + "items": [ + _pod("a"), + _pod("b", ready=False, deleting=True, ip=None), + _pod( + "c", + containers=[ + { + "name": "other", + "ports": [{"name": "http", "containerPort": 8080}], + } + ], + ), + _pod("d", containers=[{"name": "main"}]), + ], + } + ) + pods = retune_kube.list_pods("x=y", timeout_s=3) + call = transport.calls[-1] + assert path(call) == f"api/v1/namespaces/{NS}/pods" + assert call["query"]["labelSelector"] == "x=y" + assert call["timeout"] == (3, 3) + assert pods[0] == PodInfo( + name="a", + uid="uid-a", + ip="10.0.0.5", + ready=True, + terminating=False, + labels={"nvidia.com/dynamo-component-type": "frontend"}, + http_port=8000, # the main container's port wins + restarts=2, # the main container's + ) + assert (pods[1].ready, pods[1].terminating, pods[1].ip) == (False, True, None) + assert pods[2].http_port == 8080 + assert pods[3].http_port is None + + +def test_leases_round_trip_raw_json(kube): + retune_kube, transport = kube + lease = { + "apiVersion": "coordination.k8s.io/v1", + "kind": "Lease", + "metadata": {"name": "llm-planner-retune", "resourceVersion": "5"}, + "spec": { + "holderIdentity": "pod/abc", + "leaseDurationSeconds": 60, + "renewTime": "2026-10-28T12:00:00.000000Z", + }, + } + transport.reply(FakeResponse(json.dumps(lease).encode())) + # MicroTime strings come back verbatim, not re-serialized datetimes. + assert retune_kube.get_lease("llm-planner-retune") == lease + call = transport.calls[-1] + assert call["preload"] is False + assert ( + path(call) + == f"apis/coordination.k8s.io/v1/namespaces/{NS}/leases/llm-planner-retune" + ) + transport.reply(api_error(404)) + assert retune_kube.get_lease("gone") is None + + body = {"metadata": {"name": "l2"}, "spec": {"holderIdentity": "me"}} + transport.reply({**body, "metadata": {"name": "l2", "resourceVersion": "1"}}) + created = retune_kube.create_lease(body) + assert created["metadata"]["resourceVersion"] == "1" + call = transport.calls[-1] + assert call["method"] == "POST" + assert call["body"]["apiVersion"] == "coordination.k8s.io/v1" + assert call["body"]["kind"] == "Lease" + assert "apiVersion" not in body # the caller's dict is not mutated + transport.reply(api_error(409)) + with pytest.raises(KubeAlreadyExists): + retune_kube.create_lease(body) + + transport.reply(lease) + retune_kube.replace_lease("llm-planner-retune", lease, timeout_s=2) + call = transport.calls[-1] + assert call["method"] == "PUT" + assert call["body"]["metadata"]["resourceVersion"] == "5" + assert call["timeout"] == (2, 2) + transport.reply(api_error(409, "the object has been modified")) + with pytest.raises(KubeConflict) as raised: + retune_kube.replace_lease("llm-planner-retune", lease) + assert raised.value.status == 409 + transport.reply(api_error(404)) + with pytest.raises(KubeNotFound): + retune_kube.replace_lease("gone", lease) + + transport.reply({"kind": "Status", "status": "Success"}) + retune_kube.delete_lease("l2") + assert transport.calls[-1]["method"] == "DELETE" + assert transport.calls[-1]["body"] is None + transport.reply(api_error(404)) + retune_kube.delete_lease("l2") + # Conditional on a resourceVersion: a changed Lease is a conflict. + transport.reply(api_error(409, "Precondition failed: ResourceVersion")) + with pytest.raises(KubeConflict): + retune_kube.delete_lease("l2", resource_version="5") + assert transport.calls[-1]["body"]["preconditions"] == {"resourceVersion": "5"} + + +def test_from_kube_api_shares_the_connector_client(): + api_client = client.ApiClient(client.Configuration()) + kube_api = SimpleNamespace( + current_namespace="ns2", + custom_api=client.CustomObjectsApi(api_client), + ) + retune_kube = KubernetesRetuneKube.from_kube_api(kube_api) + assert retune_kube.namespace == "ns2" + assert retune_kube.api_client is api_client + api_client.close() diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_metrics.py b/components/src/dynamo/planner/tests/unit/test_live_retune_metrics.py new file mode 100644 index 000000000000..9d08e1c81a80 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_metrics.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``dynamo_planner_retune_*`` metrics: the names the live retune guide +lists, the series that start at zero, the ``trace_sink_missing`` series a +recovered pod drops, and collectors shared per registry.""" + +from __future__ import annotations + +import math + +import pytest +from prometheus_client import REGISTRY, CollectorRegistry, generate_latest + +from dynamo.planner.core.types import RETUNE_OUTCOMES +from dynamo.planner.retune.live.metrics import PREFIX, RetuneMetrics + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + + +@pytest.fixture +def registry() -> CollectorRegistry: + return CollectorRegistry() + + +def _value(registry: CollectorRegistry, name: str, **labels: str) -> float | None: + return registry.get_sample_value(f"{PREFIX}_{name}", labels) + + +def test_exposition_uses_the_documented_names(registry): + metrics = RetuneMetrics(registry) + metrics.trace_sink_missing("frontend-0", True) + text = generate_latest(registry).decode() + + for name in ( + "decisions_total", + "failures_total", + "skipped_total", + "leader", + "ready", + "initialized", + "rollback_pending", + "controller_generation", + "baseline_score", + "best_score", + "sweep_seconds", + "missing_candidates", + "router_converged", + "router_patch_errors_total", + "trace_sink_missing", + "trace_dropped_members_total", + ): + assert f"\n# HELP {PREFIX}_{name} " in f"\n{text}", name + + +def test_known_outcomes_and_skip_reasons_start_at_zero(registry): + RetuneMetrics(registry) + + for outcome in (*RETUNE_OUTCOMES, "promoted", "rolled_back"): + assert _value(registry, "decisions_total", outcome=outcome) == 0 + for reason in ("not_leader", "not_ready", "already_decided"): + assert _value(registry, "skipped_total", reason=reason) == 0 + # Gauges of the latest decision read NaN until one is scored. + for name in ("baseline_score", "best_score", "sweep_seconds"): + value = _value(registry, name) + assert value is not None and math.isnan(value) + + +def test_trace_sink_missing_drops_the_series_once_the_sink_is_back(registry): + metrics = RetuneMetrics(registry) + + metrics.trace_sink_missing("frontend-0", True) + metrics.trace_sink_missing("frontend-1", True) + metrics.trace_sink_missing("frontend-0", False) + assert _value(registry, "trace_sink_missing", pod="frontend-0") is None + assert _value(registry, "trace_sink_missing", pod="frontend-1") == 1 + # A pod that never had a series is not an error. + metrics.trace_sink_missing("frontend-9", False) + + +def test_instances_share_the_collectors_of_their_registry(registry): + first = RetuneMetrics(registry) + second = RetuneMetrics(registry) + first.decision("applied") + second.decision("applied") + assert _value(registry, "decisions_total", outcome="applied") == 2 + + # The planner serves prometheus_client.REGISTRY; a controller rebuild or + # a second test must not hit "Duplicated timeseries". + before = ( + REGISTRY.get_sample_value(f"{PREFIX}_failures_total", {"reason": "probe"}) + or 0.0 + ) + RetuneMetrics().failure("probe") + RetuneMetrics().failure("probe") + assert ( + REGISTRY.get_sample_value(f"{PREFIX}_failures_total", {"reason": "probe"}) + == before + 2 + ) diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_reconciler.py b/components/src/dynamo/planner/tests/unit/test_live_retune_reconciler.py new file mode 100644 index 000000000000..58988605bdfd --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_reconciler.py @@ -0,0 +1,1094 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the live retune reconciler (segment, window, temp-file and +DGDR Request garbage collection) and its ``SweepRegistry``.""" + +from __future__ import annotations + +import logging +import os +import shutil +import stat +from collections.abc import Callable +from pathlib import Path +from typing import Any, Optional + +import pytest + +from dynamo.planner.config.planner_config import DgdrExecutorConfig, RetuneConfig +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD, + LABEL_DGD_UID, + LABEL_WINDOW, + DgdRef, + KubeError, + PodInfo, +) +from dynamo.planner.retune.live.reconciler import ( + JOB_NAME_LABEL, + RUN_GENERATION_ANNOTATION, + TEMP_MARGIN_S, + VISIBILITY_MARGIN_S, + Reconciler, +) + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +DGD = DgdRef(name="dgd", namespace="ns", uid="uid-now", dynamo_namespace="dyn-ns") +NOW = 1_800_000_000.0 +INTERVAL = 3600.0 +GRACE = 120.0 +TIMEOUT = 1800.0 +SEGMENT_AGE = INTERVAL + GRACE + TIMEOUT + VISIBILITY_MARGIN_S +WINDOW_AGE = INTERVAL + TIMEOUT +TEMP_AGE = 900.0 + TEMP_MARGIN_S # cut_budget_seconds + margin +OLD = 100_000.0 # older than every threshold + + +def window_ms(k: int) -> int: + """End of the k-th window before NOW's current one, in epoch ms.""" + + return int((NOW - k * INTERVAL) * 1000) + + +def deadline_for(window_end_ms: int) -> float: + return window_end_ms / 1000 + GRACE + TIMEOUT + + +# ------------------------------------------------------------------ fakes + + +class FakeKube: + namespace = DGD.namespace + + def __init__(self) -> None: + self.requests: dict[str, dict[str, Any]] = {} + self.runs: dict[str, dict[str, Any]] = {} + self.job_pods: dict[str, list[PodInfo]] = {} + self.deleted: list[str] = [] + self.selectors: list[str] = [] + self.get_run_calls: list[str] = [] + self.fail_list = False + self.fail_get_run = False + self.fail_delete: set[str] = set() + self.on_delete: Optional[Callable[[str], None]] = None + + def list_requests(self, label_selector: str, *, timeout_s: float = 10.0): + self.selectors.append(label_selector) + if self.fail_list: + raise KubeError("list failed") + return [dict(obj) for obj in self.requests.values()] + + def delete_request( + self, name: str, *, propagation: str = "Background", timeout_s: float = 10.0 + ): + if name in self.fail_delete: + raise KubeError("delete failed") + if self.on_delete is not None: + self.on_delete(name) + self.deleted.append(name) + self.requests.pop(name, None) + + def get_run(self, name: str, *, timeout_s: float = 10.0): + self.get_run_calls.append(name) + if self.fail_get_run: + raise KubeError("get failed") + return self.runs.get(name) + + def list_pods(self, label_selector: str, *, timeout_s: float = 10.0): + key, _, value = label_selector.partition("=") + assert key == JOB_NAME_LABEL + return list(self.job_pods.get(value, [])) + + # -- builders + + def add_request( + self, + name: str, + window: int, + *, + dgd_uid: Optional[str] = DGD.uid, + state: str = "completed", # completed | running | creating | invalid + job_pods: int = 0, + job_pod_phase: Optional[str] = None, + claim: str = "traces", + sub_path: Optional[str] = None, + generation: int = 1, + ) -> None: + labels = {LABEL_DGD: DGD.name, LABEL_WINDOW: str(window)} + if dgd_uid is not None: + labels[LABEL_DGD_UID] = dgd_uid + status: dict[str, Any] = {"conditions": []} + run_name = f"{name}-{generation}" + job = f"{run_name}-job" + if state in ("completed", "running"): + completed = state == "completed" + self.runs[run_name] = { + "metadata": { + "name": run_name, + "annotations": {RUN_GENERATION_ANNOTATION: str(generation)}, + }, + "status": { + "jobRef": {"name": job}, + "conditions": [ + { + "type": "Completed", + "status": "True" if completed else "False", + "reason": "SearchSucceeded" if completed else "JobRunning", + } + ], + }, + } + if completed: + status["recentRunRefs"] = [{"name": run_name}] + else: + status["activeRunRef"] = {"name": run_name} + self.job_pods[job] = [ + PodInfo( + name=f"{job}-{i}", + uid=str(i), + ip=None, + ready=False, + terminating=False, + phase=job_pod_phase, + ) + for i in range(job_pods) + ] + accepted = ("True", "RunCreated") + elif state == "invalid": + accepted = ("False", "InvalidSpec") + else: + accepted = ("False", "RunCreating") + status["activeRunRef"] = {"name": run_name} + status["conditions"].append( + { + "type": "Accepted", + "status": accepted[0], + "reason": accepted[1], + "observedGeneration": generation, + } + ) + self.requests[name] = { + "metadata": { + "name": name, + "uid": f"{name}-uid", + "generation": generation, + "labels": labels, + }, + "spec": { + "workload": { + "trace": { + "source": { + "persistentVolumeClaim": { + "claimName": claim, + "subPath": sub_path or f"windows/dgd/{window}.jsonl.gz", + } + }, + } + } + }, + "status": status, + } + + def finish_job(self, name: str) -> None: + for run in self.runs.values(): + if run["metadata"]["name"].startswith(name): + self.job_pods[run["status"]["jobRef"]["name"]] = [] + + +class RecordingMetrics: + def __init__(self) -> None: + self.sinks: list[tuple[str, bool]] = [] + self.fail_sinks = False + + def trace_sink_missing(self, pod: str, missing: bool) -> None: + if self.fail_sinks: + raise RuntimeError("metrics backend failed") + self.sinks.append((pod, missing)) + + def __getattr__(self, name: str): + return lambda *args, **kwargs: None + + +def make( + tmp_path: Path, *, keep_requests: int = 2 +) -> tuple[Reconciler, FakeKube, RecordingMetrics]: + cfg = DgdrExecutorConfig( + claim_name="traces", + model="Qwen/Qwen3-0.6B", + sku="h200_sxm", + engine={}, + cut_budget_seconds=900, + keep_requests=keep_requests, + ) + retune = RetuneConfig( + interval_seconds=INTERVAL, + grace_seconds=GRACE, + timeout_seconds=TIMEOUT, + executor="dgdr", + dgdr=cfg, + ) + kube = FakeKube() + metrics = RecordingMetrics() + reconciler = Reconciler( + kube, DGD, cfg, retune, str(tmp_path), metrics, now=lambda: NOW + ) + return reconciler, kube, metrics + + +def touch(path: Path, age_s: float, content: bytes = b"x") -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + os.utime(path, (NOW - age_s, NOW - age_s)) + return path + + +def age_dir(path: Path, age_s: float) -> None: + os.utime(path, (NOW - age_s, NOW - age_s)) + + +def frontend(name: str, *, ready: bool = True, terminating: bool = False) -> PodInfo: + return PodInfo( + name=name, + uid=f"{name}-uid", + ip="10.0.0.1", + ready=ready, + terminating=terminating, + ) + + +def run_once( + reconciler: Reconciler, + *, + last: Optional[int] = None, + due: Optional[int] = None, + pods: Optional[list[PodInfo]] = None, + should_continue: Optional[Callable[[], bool]] = None, +): + return reconciler.run_once( + last_window_end_ms=last, + due_window_end_ms=due, + deadline_for=deadline_for, + frontend_pods=pods, + should_continue=should_continue, + ) + + +# ----------------------------------------------------------------- segments + + +def test_segments_of_a_live_pod_keep_the_newest_and_recent_ones(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + pod_dir = Path(reconciler.segments_root) / "fe-0" + old0 = touch(pod_dir / "req.000000.jsonl.gz", OLD) + recent = touch(pod_dir / "req.000001.jsonl.gz", SEGMENT_AGE - 60) + old2 = touch(pod_dir / "req.000002.jsonl.gz", OLD) + newest = touch(pod_dir / "req.000010.jsonl.gz", OLD) # numeric, not lexical + other = touch(pod_dir / "notes.txt", OLD) + + report = run_once(reconciler, pods=[frontend("fe-0")]) + + assert report.deleted_segments == 2 + assert not old0.exists() and not old2.exists() + assert recent.exists() and newest.exists() and other.exists() + assert report.errors == {} + + +def test_directory_of_a_gone_pod_goes_only_when_everything_is_old( + tmp_path: Path, +) -> None: + reconciler, _, _ = make(tmp_path) + root = Path(reconciler.segments_root) + touch(root / "gone-old" / "req.000000.jsonl.gz", OLD) + touch(root / "gone-old" / "req.000001.jsonl.gz", SEGMENT_AGE + 1) + age_dir(root / "gone-old", OLD) + touch(root / "gone-recent" / "req.000000.jsonl.gz", OLD) + touch(root / "gone-recent" / "req.000001.jsonl.gz", 10) + age_dir(root / "gone-recent", OLD) + touch(root / "fe-0" / "req.000000.jsonl.gz", 10) + + report = run_once(reconciler, pods=[frontend("fe-0")]) + + assert report.deleted_pod_dirs == ("gone-old",) + assert not (root / "gone-old").exists() + assert (root / "gone-recent").exists() + assert (root / "fe-0").exists() + + +@pytest.mark.parametrize("pods", [None, []]) +def test_no_segment_is_deleted_without_a_known_pod_set(tmp_path: Path, pods) -> None: + reconciler, _, _ = make(tmp_path) + root = Path(reconciler.segments_root) + seg = touch(root / "fe-0" / "req.000000.jsonl.gz", OLD) + touch(root / "fe-0" / "req.000001.jsonl.gz", OLD) + age_dir(root / "fe-0", OLD) + + report = run_once(reconciler, pods=pods) + + assert report.deleted_segments == 0 and report.deleted_pod_dirs == () + assert seg.exists() + + +def test_trace_sink_missing_is_reported_and_cleared(tmp_path: Path) -> None: + reconciler, _, metrics = make(tmp_path) + root = Path(reconciler.segments_root) + touch(root / "fe-0" / "req.000000.jsonl.gz", 0) + pods = [ + frontend("fe-0"), + frontend("fe-1"), + frontend("starting", ready=False), + frontend("leaving", terminating=True), + ] + + report = run_once(reconciler, pods=pods) + assert report.sink_missing == ("fe-1",) + assert sorted(metrics.sinks) == [("fe-0", False), ("fe-1", True)] + + metrics.sinks.clear() + (root / "fe-1").mkdir() + report = run_once(reconciler, pods=pods) + assert report.sink_missing == () + assert ("fe-1", False) in metrics.sinks + + # A pod that leaves has its gauge cleared. + metrics.sinks.clear() + run_once(reconciler, pods=[frontend("fe-0")]) + assert sorted(metrics.sinks) == [("fe-0", False), ("fe-1", False)] + + # An unknown pod set reports nothing. + metrics.sinks.clear() + assert run_once(reconciler, pods=None).sink_missing == () + assert metrics.sinks == [] + + +@pytest.mark.skipif( + hasattr(os, "geteuid") and os.geteuid() == 0, + reason="root reads a directory whatever its mode", +) +def test_an_unreadable_pod_directory_does_not_stop_the_others(tmp_path: Path) -> None: + reconciler, kube, metrics = make(tmp_path, keep_requests=0) + root = Path(reconciler.segments_root) + names = [f"fe-{i}" for i in range(6)] + for name in names: + touch(root / name / "req.000000.jsonl.gz", OLD) + touch(root / name / "req.000001.jsonl.gz", OLD) + touch(root / name / "req.000002.jsonl.gz", OLD) + kube.add_request("rt-old", window_ms(6), state="completed") + # Created by a frontend running as another uid, or with umask 077. + locked = root / "fe-3" + locked.chmod(0) + try: + report = run_once( + reconciler, due=window_ms(0), pods=[frontend(name) for name in names] + ) + finally: + locked.chmod(stat.S_IRWXU) + + # Every other pod lost its two old segments, whatever the scandir order. + assert report.deleted_segments == 2 * (len(names) - 1) + assert len(list(locked.iterdir())) == 3 + assert report.errors == {"segments": 1} + # The sink check and the Request GC still ran. + assert sorted(metrics.sinks) == [(name, False) for name in names] + assert report.deleted_requests == ("rt-old",) + + +def test_a_failing_sink_check_keeps_the_segment_counts(tmp_path: Path) -> None: + reconciler, _, metrics = make(tmp_path) + pod_dir = Path(reconciler.segments_root) / "fe-0" + touch(pod_dir / "req.000000.jsonl.gz", OLD) + touch(pod_dir / "req.000001.jsonl.gz", OLD) + metrics.fail_sinks = True + + report = run_once(reconciler, pods=[frontend("fe-0")]) + + assert report.deleted_segments == 1 + assert report.errors == {"sinks": 1} + + +def test_only_the_top_level_segments_of_a_live_pod_are_collected( + tmp_path: Path, +) -> None: + # The frontend premise check requires the sink's file prefix directly + # under the pod directory, so the reconciler does not walk + # subdirectories of a live pod. + reconciler, _, _ = make(tmp_path) + pod_dir = Path(reconciler.segments_root) / "fe-0" + nested = [touch(pod_dir / "sub" / f"req.00000{i}.jsonl.gz", OLD) for i in range(3)] + top = [touch(pod_dir / f"req.00000{i}.jsonl.gz", OLD) for i in range(3)] + + report = run_once(reconciler, pods=[frontend("fe-0")]) + + assert report.deleted_segments == 2 + assert [path.exists() for path in top] == [False, False, True] + assert all(path.exists() for path in nested) + + +# ------------------------------------------------------------------ windows + + +def window_path(reconciler: Reconciler, window: int) -> Path: + return Path(reconciler.windows_root) / f"{window}.jsonl.gz" + + +def selection_path(reconciler: Reconciler, window: int, digest: str) -> Path: + return Path(reconciler.windows_root) / f"{window}-{digest}.jsonl.gz" + + +def test_protected_windows_are_never_deleted(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + last, due = window_ms(2), int((NOW - 60) * 1000) + assert NOW >= deadline_for(last) and NOW < deadline_for(due) + last_path = touch(window_path(reconciler, last), OLD) + due_path = touch(window_path(reconciler, due), OLD) + stale = window_ms(5) + stale_path = touch(window_path(reconciler, stale), OLD) + registry_window = window_ms(7) + registry_path = touch(window_path(reconciler, registry_window), OLD) + reconciler.protect(registry_window, []) + + report = run_once(reconciler, last=last, due=due, pods=[frontend("fe-0")]) + + assert due_path.exists() # undecided and before its deadline + assert not last_path.exists() # in the range, but past its deadline + assert registry_path.exists() # this process's sweep uses it + assert not stale_path.exists() + assert report.deleted_windows == tuple(sorted((last, stale))) + + reconciler.release(registry_window) + report = run_once(reconciler, last=last, due=due) + assert report.deleted_windows == (registry_window,) + assert due_path.exists() + + +def test_range_protection_uses_the_window_deadline(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + # Two undecided windows, both before their deadline. + young = int((NOW - 60) * 1000) + older = int((NOW - 1000) * 1000) + paths = [touch(window_path(reconciler, w), OLD) for w in (young, older)] + + run_once(reconciler, last=older, due=young) + assert all(path.exists() for path in paths) + + # Without a persisted last window, every window before the due one with + # a live deadline is protected too. + run_once(reconciler, last=None, due=young) + assert all(path.exists() for path in paths) + + +def test_referenced_window_goes_once_every_referrer_is_finished(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + # A window outside the protected range whose deadline has not passed, + # so only the referrer rule decides. + window = int((NOW - 60) * 1000) + path = touch(window_path(reconciler, window), 10) # young: no fallback + kube.add_request("rt-a", window, state="completed") + kube.add_request("rt-b", window, state="running", job_pods=1) + + def round_(): + return run_once(reconciler, last=window_ms(0), due=window_ms(0)) + + round_() + assert path.exists() # rt-b's Run is still active + + # Completed, but the Job pod is still there. + kube.add_request("rt-b", window, state="completed", job_pods=1) + round_() + assert path.exists() + + kube.finish_job("rt-b") + report = round_() + assert report.deleted_windows == (window,) + assert not path.exists() + assert kube.deleted == [] # the Requests stay (keep_requests) + + +@pytest.mark.parametrize( + ("phase", "deleted"), + [("Succeeded", True), ("Failed", True), ("Running", False), (None, False)], +) +def test_finished_job_pods_of_a_completed_run_release_its_window( + tmp_path: Path, phase: Optional[str], deleted: bool +) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + window = int((NOW - 60) * 1000) + path = touch(window_path(reconciler, window), 10) # young: no fallback + kube.add_request("rt-a", window, state="completed", job_pods=1, job_pod_phase=phase) + + report = run_once(reconciler, last=window_ms(0), due=window_ms(0)) + assert (report.deleted_windows == (window,)) == deleted + assert path.exists() != deleted + + +def test_invalid_spec_request_does_not_hold_its_window(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + window = window_ms(4) + path = touch(window_path(reconciler, window), 10) + kube.add_request("rt-a", window, state="invalid") + + report = run_once(reconciler, due=window_ms(0)) + assert report.deleted_windows == (window,) + assert not path.exists() + + +def test_request_waiting_for_its_run_holds_its_window(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + window = window_ms(4) + path = touch(window_path(reconciler, window), OLD) + kube.add_request("rt-a", window, state="creating") + assert path.exists() + # The deadline has passed, so the Request itself is deleted; the window + # file then has no referrer and is old. + report = run_once(reconciler, due=window_ms(0)) + assert kube.deleted == ["rt-a"] + assert report.deleted_windows == (window,) + assert not path.exists() + + reconciler2, kube2, _ = make(tmp_path / "other", keep_requests=10) + young_window = int((NOW - 60) * 1000) + young_path = touch(window_path(reconciler2, young_window), OLD) + kube2.add_request("rt-b", young_window, state="creating") + report = run_once(reconciler2, due=None, last=None) + assert young_path.exists() + assert kube2.deleted == [] + + +def test_unreferenced_window_needs_the_fallback_age(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + young = touch(window_path(reconciler, window_ms(5)), WINDOW_AGE - 60) + old = touch(window_path(reconciler, window_ms(6)), WINDOW_AGE + 60) + unknown = touch(Path(reconciler.windows_root) / "README", OLD) + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_windows == (window_ms(6),) + assert young.exists() and not old.exists() and unknown.exists() + + +def test_windows_named_by_their_selection_are_protected_by_their_end( + tmp_path: Path, +) -> None: + reconciler, _, _ = make(tmp_path) + last, due = window_ms(2), int((NOW - 60) * 1000) + due_path = touch(selection_path(reconciler, due, "0123456789ab"), OLD) + stale_path = touch(selection_path(reconciler, window_ms(5), "0123456789ab"), OLD) + legacy_path = touch(window_path(reconciler, window_ms(6)), OLD) + + report = run_once(reconciler, last=last, due=due, pods=[frontend("fe-0")]) + + assert due_path.exists() # undecided and before its deadline + assert not stale_path.exists() and not legacy_path.exists() + assert report.deleted_windows == (window_ms(6), window_ms(5)) + + +def test_a_window_named_by_its_selection_goes_once_its_referrer_is_finished( + tmp_path: Path, +) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + # Outside the protected range and before its deadline, so only its + # referrer decides. + window = int((NOW - 60) * 1000) + path = touch(selection_path(reconciler, window, "0123456789ab"), 10) + other_selection = touch(selection_path(reconciler, window, "ba9876543210"), 10) + kube.add_request( + "rt-a", + window, + state="running", + job_pods=1, + sub_path=f"windows/dgd/{path.name}", + ) + + def round_(): + return run_once(reconciler, last=window_ms(0), due=window_ms(0)) + + round_() + assert path.exists() and other_selection.exists() + + kube.add_request( + "rt-a", window, state="completed", sub_path=f"windows/dgd/{path.name}" + ) + report = round_() + assert report.deleted_windows == (window,) + assert not path.exists() + assert other_selection.exists() # unreferenced but young + + +def test_windows_are_kept_when_requests_cannot_be_listed(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path) + path = touch(window_path(reconciler, window_ms(6)), OLD) + temp = touch(Path(reconciler.windows_root) / f".tmp-{window_ms(6)}-abc", OLD) + kube.fail_list = True + + report = run_once(reconciler, due=window_ms(0)) + + assert path.exists() + assert not temp.exists() # temp files do not depend on references + assert report.errors == {"list_requests": 1} + + +def test_a_request_of_another_claim_does_not_reference_the_window( + tmp_path: Path, +) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + window = window_ms(6) + path = touch(window_path(reconciler, window), OLD) + kube.add_request("rt-a", window, state="running", claim="other-claim") + kube.add_request( + "rt-b", + window, + state="running", + sub_path=f"windows/other-dgd/{window}.jsonl.gz", + ) + # Both Requests are past their deadline and get deleted anyway; the + # window had no referrer and is old. + report = run_once(reconciler, due=window_ms(0)) + assert report.deleted_windows == (window,) + assert not path.exists() + + +def test_temp_files(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + root = Path(reconciler.windows_root) + protected = window_ms(1) + reconciler.protect(protected, []) + old = touch(root / f".tmp-{window_ms(3)}-a1", TEMP_AGE + 1) + young = touch(root / f".tmp-{window_ms(3)}-b2", TEMP_AGE - 1) + in_use = touch(root / f".tmp-{protected}-c3", OLD) + odd = touch(root / ".tmp-garbage", TEMP_AGE + 1) + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_temp_files == 2 + assert not old.exists() and not odd.exists() + assert young.exists() and in_use.exists() + + +# ----------------------------------------------------------------- requests + + +def test_terminal_requests_outside_the_newest_kept_windows_go(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=2) + windows = [window_ms(k) for k in (5, 4, 3)] + for i, window in enumerate(windows): + kube.add_request(f"rt-{i}-a", window, state="completed") + kube.add_request(f"rt-{i}-b", window, state="invalid") + + report = run_once(reconciler, due=window_ms(0)) + + assert kube.selectors == [f"{LABEL_DGD}=dgd"] + assert report.deleted_requests == ("rt-0-a", "rt-0-b") + assert set(kube.requests) == {"rt-1-a", "rt-1-b", "rt-2-a", "rt-2-b"} + + +def test_keep_zero_still_keeps_requests_of_an_undecided_window(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=0) + due = int((NOW - 60) * 1000) + kube.add_request("rt-due", due, state="completed") + kube.add_request("rt-old", window_ms(4), state="completed") + + report = run_once(reconciler, last=window_ms(4), due=due) + + # rt-due's window is undecided and before its deadline: the next sweep + # adopts its Succeeded Request instead of recreating it. + assert report.deleted_requests == ("rt-old",) + + +def test_non_terminal_request_goes_after_its_deadline(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + current = int((NOW - 60) * 1000) + kube.add_request("rt-live", current, state="running", job_pods=1) + kube.add_request("rt-late", window_ms(2), state="running", job_pods=1) + + report = run_once(reconciler, due=current) + + assert report.deleted_requests == ("rt-late",) + + +def test_requests_of_a_recreated_dgd_are_deleted(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + current = int((NOW - 60) * 1000) + kube.add_request("rt-old-uid", current, dgd_uid="uid-before", state="running") + kube.add_request("rt-no-uid", current, dgd_uid=None, state="running") + kube.add_request("rt-ours", current, state="running") + + report = run_once(reconciler, due=current) + + assert report.deleted_requests == ("rt-old-uid",) + + +def test_delete_queue(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + current = int((NOW - 60) * 1000) + kube.add_request("rt-a", current, state="running") + kube.add_request("rt-b", current, state="running") + reconciler.enqueue_delete(["rt-a", "rt-b", "rt-missing"]) + kube.fail_delete = {"rt-b"} + + report = run_once(reconciler, due=current) + + assert report.deleted_requests == ("rt-a", "rt-missing") + assert report.errors == {"delete_request": 1} + assert reconciler.pending_deletes() == frozenset({"rt-b"}) + + kube.fail_delete = set() + report = run_once(reconciler, due=current) + assert report.deleted_requests == ("rt-b",) + assert reconciler.pending_deletes() == frozenset() + + +def test_the_delete_queue_is_processed_when_listing_fails(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path) + reconciler.enqueue_delete(["rt-a"]) + kube.fail_list = True + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_requests == ("rt-a",) + + +def test_protected_requests_are_never_deleted(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=0) + old = window_ms(6) + kube.add_request("rt-terminal", old, state="completed") + kube.add_request("rt-late", old, state="running") + kube.add_request("rt-other-uid", old, dgd_uid="uid-before", state="running") + reconciler.protect(old, ["rt-terminal", "rt-late", "rt-other-uid"]) + reconciler.enqueue_delete(["rt-late"]) + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_requests == () + assert kube.deleted == [] + assert reconciler.pending_deletes() == frozenset({"rt-late"}) + assert reconciler.protected_requests() == frozenset( + {"rt-terminal", "rt-late", "rt-other-uid"} + ) + + reconciler.release(old) + report = run_once(reconciler, due=window_ms(0)) + assert report.deleted_requests == ("rt-late", "rt-other-uid", "rt-terminal") + + +def test_unreadable_run_state_keeps_the_request(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=0) + kube.add_request("rt-a", window_ms(6), state="completed") + window = touch(window_path(reconciler, window_ms(6)), 10) + kube.fail_get_run = True + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_requests == () + assert report.deleted_windows == () + assert window.exists() + assert report.errors == {"get_run": 1} + + +def test_terminal_and_finished_verdicts_are_cached(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + window = window_ms(6) + kube.add_request("rt-a", window, state="completed") + touch(window_path(reconciler, window), 10) + + run_once(reconciler, due=window_ms(0)) + calls = len(kube.get_run_calls) + assert calls > 0 + touch(window_path(reconciler, window), 10) # recreated: still finished + report = run_once(reconciler, due=window_ms(0)) + assert len(kube.get_run_calls) == calls + assert report.deleted_windows == (window,) + + +def test_deleted_requests_stop_referencing_their_window(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=0) + old_window, young_window = window_ms(6), window_ms(7) + old = touch(window_path(reconciler, old_window), WINDOW_AGE + 60) + young = touch(window_path(reconciler, young_window), WINDOW_AGE - 60) + kube.add_request("rt-old", old_window, state="running") + kube.add_request("rt-young", young_window, state="running") + + report = run_once(reconciler, due=window_ms(0)) + + assert report.deleted_requests == ("rt-old", "rt-young") + assert report.deleted_windows == (old_window,) + assert not old.exists() and young.exists() + + +# ------------------------------------------------------------ early stop + + +def one_of_each(tmp_path: Path) -> tuple[Reconciler, FakeKube]: + """Six deletable things: 2 segments, a pod dir, a temp file, a Request + and the window file it referenced.""" + + reconciler, kube, _ = make(tmp_path, keep_requests=0) + root = Path(reconciler.segments_root) + for i in range(3): + touch(root / "fe-0" / f"req.00000{i}.jsonl.gz", OLD) + touch(root / "gone" / "req.000000.jsonl.gz", OLD) + age_dir(root / "gone", OLD) + touch(Path(reconciler.windows_root) / f".tmp-{window_ms(9)}-a1", OLD) + touch(window_path(reconciler, window_ms(6)), OLD) + kube.add_request("rt-a", window_ms(6), state="completed") + return reconciler, kube + + +def deletions(report) -> int: + return ( + report.deleted_segments + + len(report.deleted_pod_dirs) + + len(report.deleted_windows) + + report.deleted_temp_files + + len(report.deleted_requests) + ) + + +def test_should_continue_is_asked_before_every_delete_and_read( + tmp_path: Path, +) -> None: + reconciler, _ = one_of_each(tmp_path) + calls: list[int] = [] + + def should_continue() -> bool: + calls.append(1) + return True + + report = run_once( + reconciler, + due=window_ms(0), + pods=[frontend("fe-0")], + should_continue=should_continue, + ) + + assert deletions(report) == 6 + # Six deletes, the Request listing and rt-a's Run read. + assert len(calls) == 8 + assert not report.stopped + + +@pytest.mark.parametrize(("allowed", "run_reads"), [(2, 1), (4, 3)]) +def test_a_stopped_round_reads_no_more_run_state( + tmp_path: Path, allowed: int, run_reads: int +) -> None: + # Each read can wait for its timeout: once the budget has passed or + # leadership is lost, neither the Request verdicts (allowed=2: the + # listing, then one Run read) nor the window referrers (allowed=4) + # read anything more. + reconciler, kube, _ = make(tmp_path, keep_requests=10) + for k in (4, 5, 6): + kube.add_request(f"rt-{k}", window_ms(k), state="completed") + touch(window_path(reconciler, window_ms(k)), 10) + answers = [True] * allowed + + def should_continue() -> bool: + return answers.pop(0) if answers else False + + report = run_once(reconciler, due=window_ms(0), should_continue=should_continue) + + assert report.stopped + assert len(kube.get_run_calls) == run_reads + assert kube.deleted == [] and report.deleted_windows == () + + +def test_a_no_from_should_continue_ends_the_round(tmp_path: Path) -> None: + reconciler, kube = one_of_each(tmp_path) + answers = [True, True] + + def should_continue() -> bool: + return answers.pop(0) if answers else False + + report = run_once( + reconciler, + due=window_ms(0), + pods=[frontend("fe-0")], + should_continue=should_continue, + ) + + # Two segment deletions, then nothing: no Request was even listed. + assert deletions(report) == 2 + assert report.stopped + assert kube.selectors == [] and kube.deleted == [] + assert window_path(reconciler, window_ms(6)).exists() + + report = run_once(reconciler, due=window_ms(0), pods=[frontend("fe-0")]) + assert deletions(report) == 4 + assert not report.stopped + + +def test_a_failing_should_continue_ends_the_round(tmp_path: Path) -> None: + reconciler, kube = one_of_each(tmp_path) + + def should_continue() -> bool: + raise RuntimeError("lease state unknown") + + report = run_once( + reconciler, + due=window_ms(0), + pods=[frontend("fe-0")], + should_continue=should_continue, + ) + + assert deletions(report) == 0 + assert report.stopped + assert kube.deleted == [] + + +def test_persistent_conditions_are_logged_when_they_start( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + current = int((NOW - 60) * 1000) + kube.add_request("rt-no-uid", current, dgd_uid=None, state="running") + kube.add_request("rt-no-window", current, state="running") + del kube.requests["rt-no-window"]["metadata"]["labels"][LABEL_WINDOW] + pods = [frontend("fe-0")] + sink = Path(reconciler.segments_root) / "fe-0" + + def warnings() -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING + ] + + with caplog.at_level(logging.WARNING, logger=Reconciler.__module__): + for _ in range(3): + run_once(reconciler, due=current, pods=pods) + assert len(warnings()) == 3 # fe-0's sink, one per missing label + sink.mkdir(parents=True) + run_once(reconciler, due=current, pods=pods) + shutil.rmtree(sink) # missing again: a new occurrence + run_once(reconciler, due=current, pods=pods) + run_once(reconciler, due=current, pods=pods) + + messages = warnings() + assert len(messages) == 4 + assert sum("without a trace directory" in m for m in messages) == 2 + assert sum(LABEL_DGD_UID in m for m in messages) == 1 + assert sum(LABEL_WINDOW in m for m in messages) == 1 + + +# --------------------------------------------------------------- set_dgd + + +RECREATED = DgdRef( + name=DGD.name, namespace=DGD.namespace, uid="uid-new", dynamo_namespace="dyn-ns" +) + + +def test_set_dgd_follows_a_recreated_dgd(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=2) + current = int((NOW - 60) * 1000) + # The new DGD's Requests: one kept as a recent terminal window, one + # still running before its deadline. + kube.add_request("new-w1-done", window_ms(1), dgd_uid="uid-new") + kube.add_request("new-w0-running", current, dgd_uid="uid-new", state="running") + kube.add_request("old-w2", window_ms(2), dgd_uid=DGD.uid, state="running") + + reconciler.set_dgd(RECREATED) + report = run_once(reconciler, due=current) + + assert reconciler.dgd == RECREATED + assert report.deleted_requests == ("old-w2",) + assert set(kube.requests) == {"new-w1-done", "new-w0-running"} + + +def test_set_dgd_clears_the_request_state_caches(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=10) + kube.add_request("rt-a", window_ms(3), state="completed") + + run_once(reconciler, due=window_ms(0)) + calls = len(kube.get_run_calls) + assert calls > 0 + run_once(reconciler, due=window_ms(0)) + assert len(kube.get_run_calls) == calls # cached terminal verdict + + reconciler.set_dgd(DGD) + run_once(reconciler, due=window_ms(0)) + assert len(kube.get_run_calls) > calls + + +def test_set_dgd_during_a_round_stops_it(tmp_path: Path) -> None: + reconciler, kube, _ = make(tmp_path) + reconciler.enqueue_delete(["rt-a", "rt-b", "rt-c"]) + kube.on_delete = lambda name: reconciler.set_dgd(RECREATED) + + report = run_once(reconciler, due=window_ms(0)) + + # It classified Requests by the old uid: nothing after the change. + assert report.deleted_requests == ("rt-a",) + assert report.stopped + assert reconciler.pending_deletes() == frozenset({"rt-b", "rt-c"}) + + +def test_set_dgd_with_a_new_name_moves_the_roots(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + renamed = DgdRef(name="other", namespace="ns", uid="u2", dynamo_namespace="d") + + reconciler.set_dgd(renamed) + + assert reconciler.segments_root == str(tmp_path / "segments" / "other") + assert reconciler.windows_root == str(tmp_path / "windows" / "other") + + +# ------------------------------------------------------------------ registry + + +def test_registry_is_reference_counted_per_window(tmp_path: Path) -> None: + reconciler, _, _ = make(tmp_path) + reconciler.protect(5, []) + reconciler.protect(5, ["rt-a"]) + reconciler.protect(5, ["rt-b"]) + reconciler.protect(6, ["rt-c"]) + assert reconciler.protected_windows() == frozenset({5, 6}) + assert reconciler.protected_requests() == frozenset({"rt-a", "rt-b", "rt-c"}) + + # Every protect call holds one reference; the names stay the union of + # everything given to the window while any reference is left. + reconciler.release(5) + reconciler.release(5) + assert reconciler.protected_windows() == frozenset({5, 6}) + assert reconciler.protected_requests() == frozenset({"rt-a", "rt-b", "rt-c"}) + reconciler.release(5) + assert reconciler.protected_windows() == frozenset({6}) + assert reconciler.protected_requests() == frozenset({"rt-c"}) + + reconciler.release(5) # already released: no-op, never negative + reconciler.release(99) # unknown: no-op + assert reconciler.protected_windows() == frozenset({6}) + reconciler.protect(5, ["rt-d"]) + assert reconciler.protected_requests() == frozenset({"rt-c", "rt-d"}) + reconciler.release(5) + assert reconciler.protected_windows() == frozenset({6}) + assert reconciler.protected_requests() == frozenset({"rt-c"}) + + +def test_an_old_sweeps_release_keeps_a_relaunched_sweeps_protection( + tmp_path: Path, +) -> None: + reconciler, kube, _ = make(tmp_path, keep_requests=0) + window = window_ms(6) + path = touch(window_path(reconciler, window), OLD) + kube.add_request("rt-old", window, state="running") + kube.add_request("rt-new", window, state="running") + # A rebuilt controller relaunches the window while the aborted sweep + # of the previous generation is still winding down. + reconciler.protect(window, ["rt-old"]) + reconciler.protect(window, ["rt-new"]) + reconciler.release(window) # the aborted sweep's finally + + report = run_once(reconciler, due=window_ms(0)) + assert report.deleted_requests == () + assert report.deleted_windows == () + assert path.exists() + + reconciler.release(window) # the relaunched sweep ends + report = run_once(reconciler, due=window_ms(0)) + assert report.deleted_requests == ("rt-new", "rt-old") + assert report.deleted_windows == (window,) diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_state.py b/components/src/dynamo/planner/tests/unit/test_live_retune_state.py new file mode 100644 index 000000000000..74f7e4c6a260 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_state.py @@ -0,0 +1,755 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the live retune state snapshot, write rules and Lease leadership.""" + +from __future__ import annotations + +import dataclasses +import json +import math +from collections.abc import Callable +from typing import Optional + +import pytest + +from dynamo.planner.config.planner_config import ( + LIVE_RETUNE_K8S_TIMEOUT_SECONDS, + PlannerConfig, + live_retune_io_timeout_seconds, + live_retune_lease_margin_seconds, + live_retune_min_lease_seconds, +) +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD, + LABEL_DGD_UID, + DgdRef, + FrontendStartup, + KubeConflict, + KubeError, +) +from dynamo.planner.retune.live.state import ( + HISTORY_ANNOTATION, + HISTORY_LIMIT, + STATE_ANNOTATION, + STATE_FORMAT_VERSION, + LeaseLeadership, + RetuneState, + StateFormatError, + StateSnapshot, + StateStore, + StateVersionError, + append_history, + claim, + deploy, + deployed_off, + find_retune_leases, + fingerprint, + has_decision, + leader_term, + leadership_lease_name, + maps_equal, + merge_initials, + micro_time, + owner_reference, + record_decision, + reset, + rollback_complete, + rollback_start, + state_lease_name, +) +from dynamo.planner.tests.live_retune_fakes import FakeRetuneKube + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +DGD = DgdRef(name="graph", namespace="test-ns", uid="uid-1", dynamo_namespace="dyn") +STARTUP = {"aggregated": {"overlap_score_credit": 1.0, "router_temperature": 0.0}} +DEPLOYED = {"aggregated": {"overlap_score_credit": 0.5, "router_temperature": 0.2}} + + +class FakeClock: + """Shared wall and monotonic clocks a test advances by hand.""" + + def __init__(self, wall: float = 1_800_000_000.0, mono: float = 1000.0) -> None: + self.wall = wall + self.mono = mono + + def time(self) -> float: + return self.wall + + def monotonic(self) -> float: + return self.mono + + def advance(self, seconds: float) -> None: + self.wall += seconds + self.mono += seconds + + +def live_config(**overrides) -> PlannerConfig: + """A planner config that passes the live retune configuration premises.""" + + retune = { + "interval_seconds": 21600, + "grace_seconds": 120, + "timeout_seconds": 20400, + "router": { + "overlap_score_credit": {"choices": [0.5, 1.0]}, + "prefill_load_scale": 1.0, + "temperature": {"choices": [0.0, 0.2]}, + }, + "executor": "dgdr", + "dgdr": { + "claim_name": "request-traces", + "model": "Qwen/Qwen3-0.6B", + "sku": "h200_sxm", + "engine": { + "mode": "aggregated", + "backend": "vllm", + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + { + "tensor": 1, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, + } + ] + }, + "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + }, + "prefill": None, + "decode": None, + }, + }, + }, + } + retune.update(overrides.pop("retune", {})) + return PlannerConfig( + mode="agg", + environment="kubernetes", + enable_load_scaling=False, + enable_throughput_scaling=False, + retune=retune, + **overrides, + ) + + +def startup(source: str = "a") -> FrontendStartup: + return FrontendStartup( + router_mode="kv_router", + router_config={}, + startup_maps=STARTUP, + frontend_component="Frontend", + replicas=1, + replica_sync=False, + fingerprint_source=source, + ) + + +def deployed_state(**changes) -> RetuneState: + state = RetuneState( + dgd_uid="uid-1", + leader_term="1/me/1", + fingerprint="fp", + mode="active", + deploy_seq=3, + flat={"overlap_score_credit": 0.5}, + by_wt=DEPLOYED, + initials={"r1": dict(STARTUP["aggregated"])}, + last_window_end_ms=1_000, + ) + return dataclasses.replace(state, **changes) + + +# --------------------------------------------------------------- helpers + + +def test_names_term_owner_and_time() -> None: + assert leadership_lease_name("graph") == "graph-planner-retune" + assert state_lease_name("graph") == "graph-planner-retune-state" + assert leader_term(2, "pod/abc", 7) == "2/pod/abc/7" + assert owner_reference(DGD) == { + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeployment", + "name": "graph", + "uid": "uid-1", + "blockOwnerDeletion": False, + } + assert micro_time(0.5) == "1970-01-01T00:00:00.500000Z" + + +def test_maps_equal() -> None: + assert maps_equal({"a": 1, "b": {"c": 0.1 + 0.2}}, {"a": 1.0, "b": {"c": 0.3}}) + assert not maps_equal({"a": 1}, {"a": 1, "b": 2}) + assert not maps_equal({"a": {"c": 1.0}}, {"a": 1.0}) + assert not maps_equal({"a": 1.0}, {"a": 1.1}) + assert not maps_equal(None, {}) + assert maps_equal(None, None) + + +def test_fingerprint_tracks_what_the_point_was_chosen_under() -> None: + config = live_config() + base = fingerprint(config, startup("a")) + assert base == fingerprint(live_config(), startup("a")) + assert len(base) == 16 + # The parsed frontend startup is covered. + assert fingerprint(config, startup("b")) != base + assert fingerprint(config, None) != base + # So are the objective and the scoring SLA ... + assert ( + fingerprint(live_config(retune={"objective": "Throughput"}), startup()) != base + ) + assert fingerprint(live_config(retune={"sla_ttft_ms": 99.0}), startup()) != base + # ... but not advisory, which the state's mode records. + assert fingerprint(live_config(advisory=True), startup()) == base + + +# ----------------------------------------------------------------- state + + +def test_state_round_trips_through_annotations() -> None: + state = append_history(deployed_state(rollback_pending=True), {"decision_id": "d1"}) + annotations = state.to_annotations() + assert set(annotations) == {STATE_ANNOTATION, HISTORY_ANNOTATION} + body = json.loads(annotations[STATE_ANNOTATION]) + assert body["leader_term"] == "1/me/1" and body["deploy_seq"] == 3 + assert RetuneState.from_annotations(annotations) == state + + +def _body(**fields) -> str: + return json.dumps({"version": STATE_FORMAT_VERSION, **fields}) + + +def test_state_parsing_edge_cases() -> None: + assert RetuneState.from_annotations(None) is None + assert RetuneState.from_annotations({"other": "x"}) is None + claimed = RetuneState.from_annotations( + {STATE_ANNOTATION: _body(dgd_uid="u", leader_term="t")} + ) + assert claimed == RetuneState(dgd_uid="u", leader_term="t") + assert not claimed.has_deployment + for bad in ( + "{", + "[]", + json.dumps({"dgd_uid": "u"}), # no version + _body(leader_term="t"), + _body(dgd_uid="u", mode="bogus"), + _body(dgd_uid="u", flat={"a": "x"}), + _body(dgd_uid="u", deploy_seq="3"), + # json.loads accepts NaN and Infinity; they are not numbers here. + _body(dgd_uid="u", flat={"a": math.nan}), + _body(dgd_uid="u", last_window_end_ms=math.nan), + _body(dgd_uid="u", last_window_end_ms=math.inf), + _body(dgd_uid="u", updated_ms="x"), + _body(dgd_uid="u", updated_ms=[1]), + _body(dgd_uid="u", fingerprint=math.nan), + ): + with pytest.raises(StateFormatError): + RetuneState.from_annotations({STATE_ANNOTATION: bad}) + # A damaged history never blocks the state. + parsed = RetuneState.from_annotations( + {STATE_ANNOTATION: _body(dgd_uid="u"), HISTORY_ANNOTATION: "not json"} + ) + assert parsed is not None and parsed.history == () + + +def test_unwritable_history_entries_are_dropped_on_read() -> None: + # json.loads accepts NaN and Infinity, canonical_json does not: one such + # entry kept would make every later state write raise. + history = [ + {"decision_id": "ok", "best": 1.5}, + {"decision_id": "nan", "best": math.nan}, + {"decision_id": "inf", "by_wt": {"aggregated": {"t": math.inf}}}, + "not an object", + ] + parsed = RetuneState.from_annotations( + { + STATE_ANNOTATION: _body(dgd_uid="u"), + HISTORY_ANNOTATION: json.dumps(history), + } + ) + assert parsed is not None + assert [entry["decision_id"] for entry in parsed.history] == ["ok"] + written = append_history(parsed, {"decision_id": "next"}).to_annotations() + assert json.loads(written[HISTORY_ANNOTATION])[0] == history[0] + + +def test_another_format_version_is_rejected_not_corrupt() -> None: + # Checked first: a newer planner's state is never judged by this + # planner's field rules (which would make it corrupt, then overwritten). + newer = json.dumps( + {"version": STATE_FORMAT_VERSION + 1, "dgd_uid": "u", "flat": {"a": "x"}} + ) + with pytest.raises(StateVersionError): + RetuneState.from_annotations({STATE_ANNOTATION: newer}) + assert not issubclass(StateVersionError, StateFormatError) + + kube = FakeRetuneKube() + store = StateStore(kube, DGD) + kube.create_lease( + { + "metadata": {"name": store.name, "annotations": {STATE_ANNOTATION: newer}}, + "spec": {}, + } + ) + with pytest.raises(StateVersionError): + store.read() + + +def test_find_retune_leases() -> None: + kube = FakeRetuneKube() + assert find_retune_leases(kube, "graph", timeout_s=1.0) == () + kube.create_lease({"metadata": {"name": state_lease_name("graph")}, "spec": {}}) + assert find_retune_leases(kube, "graph", timeout_s=1.0) == ( + state_lease_name("graph"), + ) + + +def test_deployed_off() -> None: + assert not deployed_off(None, STARTUP) + assert not deployed_off(RetuneState(dgd_uid="u"), STARTUP) + assert deployed_off(deployed_state(), STARTUP) + assert not deployed_off(deployed_state(by_wt=STARTUP), STARTUP) + assert deployed_off(deployed_state(by_wt=STARTUP, rollback_pending=True), STARTUP) + # Unknown startup: an active record may have deployed something. + assert deployed_off(deployed_state(by_wt=STARTUP), None) + # An advisory record only observed: its by_wt (an older startup) is no + # deployment, unless a rollback is still pending. + assert not deployed_off(deployed_state(mode="advisory"), STARTUP) + assert not deployed_off(deployed_state(mode="advisory", by_wt=STARTUP), None) + assert deployed_off(deployed_state(mode="advisory", rollback_pending=True), STARTUP) + + +# ----------------------------------------------------------- write rules + + +def test_claim_changes_only_the_term() -> None: + fresh = claim(None, dgd_uid="uid-1", term="t2") + assert fresh == RetuneState(dgd_uid="uid-1", leader_term="t2") + state = deployed_state() + claimed = claim(state, dgd_uid="ignored", term="t9") + assert claimed == dataclasses.replace(state, leader_term="t9") + + +def test_deploy_rollback_and_reset_replace_the_deployment() -> None: + state = deployed_state(rollback_pending=True) + applied = deploy( + state, + fingerprint="fp2", + mode="active", + flat={"k": 1.0}, + by_wt={"w": {"k": 1.0}}, + ) + assert (applied.fingerprint, applied.mode, applied.flat, applied.by_wt) == ( + "fp2", + "active", + {"k": 1.0}, + {"w": {"k": 1.0}}, + ) + assert applied.rollback_pending is False and applied.deploy_seq == 4 + assert applied.initials == state.initials and applied.history == state.history + + started = rollback_start( + deployed_state(), fingerprint="fp", mode="advisory", flat={}, by_wt=STARTUP + ) + assert started.rollback_pending and started.deploy_seq == 4 + assert started.by_wt == STARTUP and started.mode == "advisory" + + for pending in (True, False): + reset_state = reset( + deployed_state(), + fingerprint="fp3", + mode="active", + flat={}, + by_wt=STARTUP, + rollback_pending=pending, + ) + assert reset_state.rollback_pending is pending + assert reset_state.deploy_seq == 4 and reset_state.fingerprint == "fp3" + + +def test_rollback_complete_is_conditional_on_deploy_seq() -> None: + pending = deployed_state(rollback_pending=True, by_wt=STARTUP) + entry = {"decision_id": "rolled_back-3", "outcome": "rolled_back"} + done = rollback_complete(pending, expected_deploy_seq=3, entry=entry) + assert done is not None + # Only the flag flips (plus the history record). + assert dataclasses.replace(done, history=pending.history) == dataclasses.replace( + pending, rollback_pending=False + ) + assert done.history[-1] == entry + # An applied (deploy_seq moved) or an earlier completion: discard. + assert rollback_complete(pending, expected_deploy_seq=2) is None + assert rollback_complete(done, expected_deploy_seq=3) is None + + +def test_merge_initials_merges_and_prunes() -> None: + state = deployed_state(initials={"r1": {"a": 1.0}, "r2": {"a": 2.0}}) + merged = merge_initials(state, {"r3": {"a": 3}}, keep={"r1", "r3"}) + assert merged.initials == {"r1": {"a": 1.0}, "r3": {"a": 3.0}} + assert merge_initials(state, {}).initials == state.initials + + +def test_history_is_deduplicated_and_bounded() -> None: + state = deployed_state(last_window_end_ms=None) + for index in range(HISTORY_LIMIT + 5): + state = record_decision( + state, {"decision_id": f"d{index}", "outcome": "kept"}, 1000 + index + ) + assert len(state.history) == HISTORY_LIMIT + assert state.history[-1]["decision_id"] == f"d{HISTORY_LIMIT + 4}" + assert state.last_window_end_ms == 1000 + HISTORY_LIMIT + 4 + # Same id again: replaced in place at the end; the window never moves back. + again = record_decision(state, {"decision_id": "d10", "outcome": "failed"}, 5) + assert again.last_window_end_ms == state.last_window_end_ms + assert [item["decision_id"] for item in again.history].count("d10") == 1 + assert again.history[-1] == {"decision_id": "d10", "outcome": "failed"} + assert has_decision(again, "d10") and not has_decision(again, "nope") + assert record_decision(state, {"decision_id": "x"}, None).last_window_end_ms == ( + state.last_window_end_ms + ) + + +# ----------------------------------------------------------------- store + + +def test_store_creates_reads_and_conditionally_replaces() -> None: + clock = FakeClock() + kube = FakeRetuneKube(now=clock.time) + store = StateStore(kube, DGD, clock=clock.time) + empty = store.read() + assert empty == StateSnapshot(state=None, resource_version=None, lease=None) + + created = store.write(RetuneState(dgd_uid="uid-1", leader_term="t1"), empty) + lease = kube.leases[store.name] + assert lease["metadata"]["ownerReferences"] == [owner_reference(DGD)] + assert lease["metadata"]["labels"] == {LABEL_DGD: "graph", LABEL_DGD_UID: "uid-1"} + assert created.state.updated_ms == int(clock.wall * 1000) + + read = store.read() + assert ( + read.state == created.state + and read.resource_version == created.resource_version + ) + + updated = store.write(dataclasses.replace(read.state, leader_term="t2"), read) + assert store.read().state.leader_term == "t2" + # The first snapshot is stale now. + with pytest.raises(KubeConflict): + store.write(dataclasses.replace(read.state, leader_term="t3"), read) + # A create that loses the race is a conflict too. + with pytest.raises(KubeConflict): + store.write(RetuneState(dgd_uid="uid-1"), empty) + assert store.read().state.leader_term == "t2" + assert updated.resource_version != read.resource_version + + # The delete is conditional too: a stale snapshot leaves the Lease. + with pytest.raises(KubeConflict): + store.delete(read) + store.delete(store.read()) + assert store.read().lease is None + + +def test_store_reports_a_corrupt_snapshot() -> None: + kube = FakeRetuneKube() + store = StateStore(kube, DGD) + kube.create_lease( + { + "metadata": { + "name": store.name, + "annotations": {STATE_ANNOTATION: "{broken"}, + }, + "spec": {}, + } + ) + snapshot = store.read() + assert snapshot.corrupt and snapshot.state is None and snapshot.lease is not None + # The claim overwrites it, conditional on its resourceVersion. + written = store.write(RetuneState(dgd_uid="uid-1", leader_term="t"), snapshot) + assert written.state.leader_term == "t" and not store.read().corrupt + + +# ------------------------------------------------------------- leadership + + +def leadership(kube, clock, identity="a/1", **kwargs) -> LeaseLeadership: + return LeaseLeadership( + kube, + DGD, + identity=identity, + lease_duration_s=60.0, + margin_s=15.0, + monotonic=clock.monotonic, + clock=clock.time, + **kwargs, + ) + + +def holder(kube) -> str | None: + return kube.leases[leadership_lease_name("graph")]["spec"].get("holderIdentity") + + +def test_leadership_create_renew_and_expire_locally() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock) + assert not a.is_leader() + assert a.step() == "acquired" + assert a.is_leader() and a.transitions == 0 + lease = kube.leases[a.name] + assert lease["spec"]["leaseDurationSeconds"] == 60 + assert lease["metadata"]["ownerReferences"] == [owner_reference(DGD)] + assert a.leader_until == clock.mono + 45.0 + + clock.advance(20) + assert a.step() == "renewed" + assert a.leader_until == clock.mono + 45.0 + # What a conditional delete of the Lease needs. + assert a.resource_version == kube.leases[a.name]["metadata"]["resourceVersion"] + # Renewal failures keep leadership only until leader_until. + kube.fail("get_lease", times=10) + clock.advance(20) + assert a.step() == "error" and a.is_leader() + clock.advance(25) + assert a.step() == "error" and not a.is_leader() + kube.clear_faults() + # Its own lease again after losing track: acquired (a new term). + assert a.step() == "acquired" and a.is_leader() + + +def test_leadership_takes_over_only_after_an_unchanged_lease_duration() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + b = leadership(kube, clock, identity="b/1") + assert a.step() == "acquired" + assert b.step() == "standby" + # a keeps renewing: b never takes over. + for _ in range(6): + clock.advance(20) + assert a.step() == "renewed" + assert b.step() == "standby" + # a dies (SIGKILL) after its last renewal, which b saw: b takes over + # once it has watched the lease unchanged for 60s. + clock.advance(30) + assert b.step() == "standby" + clock.advance(29) + assert b.step() == "standby" + assert not a.is_leader() # a's local leadership expired long before + clock.advance(1) + assert b.step() == "acquired" + assert b.is_leader() and b.transitions == 1 and holder(kube) == "b/1" + # a comes back: it has lost, then stands by. + assert a.step() == "lost" and not a.is_leader() + assert a.step() == "standby" + + +def test_leadership_never_overlaps() -> None: + """a stops renewing at an arbitrary point; b never leads while a does.""" + + for stop_after in range(0, 60, 7): + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + b = leadership(kube, clock, identity="b/1") + a.step() + b.step() + for second in range(0, 400): + clock.advance(1) + if second < stop_after and second % 20 == 0: + a.step() + if second % 7 == 0: + b.step() + assert not (a.is_leader() and b.is_leader()) + assert b.is_leader() + + +def test_leadership_lost_to_another_holder() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + assert a.step() == "acquired" + lease = kube.leases[a.name] + lease["spec"]["holderIdentity"] = "intruder" + kube.bump(lease) + assert a.step() == "lost" + assert not a.is_leader() + + +def test_empty_holder_is_taken_over_at_once_and_release_is_conditional() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + b = leadership(kube, clock, identity="b/1") + assert a.step() == "acquired" + assert b.step() == "standby" + assert a.release() + assert holder(kube) is None and not a.is_leader() + # Released: b takes over without waiting a lease duration. + assert b.step() == "acquired" and b.transitions == 1 + # a no longer holds it: release is a no-op that never clears b. + assert not a.release() + assert holder(kube) == "b/1" + + +def test_a_step_after_release_never_leads_again() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + assert a.step() == "acquired" + assert a.release() and holder(kube) is None + # A step still running when the Lease was released stands by instead of + # taking the empty holder over, or creating the Lease again. + assert a.step() == "standby" + assert holder(kube) is None and not a.is_leader() + del kube.leases[a.name] + assert a.step() == "standby" + assert a.name not in kube.leases and not a.is_leader() + + +def test_leadership_conflicts_and_errors() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = leadership(kube, clock, identity="a/1") + # Creating the lease races another creator. + kube.fail("create_lease", KubeConflict("raced")) + assert a.step() == "standby" + kube.fail("create_lease", KubeError("timeout")) + assert a.step() == "error" + assert a.step() == "acquired" + clock.advance(20) + kube.fail("replace_lease", KubeConflict("stale")) + assert a.step() == "error" and a.is_leader() + kube.fail("replace_lease", KubeError("timeout")) + assert a.step() == "error" and a.is_leader() + assert a.step() == "renewed" + # Releasing tolerates API errors. + kube.fail("get_lease", KubeError("down")) + assert not a.release() and not a.is_leader() + + +def test_leadership_validates_arguments() -> None: + kube = FakeRetuneKube() + clock = FakeClock() + with pytest.raises(ValueError): + LeaseLeadership(kube, DGD, identity="") + with pytest.raises(ValueError): + LeaseLeadership(kube, DGD, identity="a", lease_duration_s=0) + # No time left to lead in; the default margin is three request timeouts. + for margin_s in (None, 15.0, 100.0, -1.0, math.nan): + with pytest.raises(ValueError, match="margin_s"): + LeaseLeadership( + kube, DGD, identity="a", lease_duration_s=15.0, margin_s=margin_s + ) + assert leadership(kube, clock).margin_s == 15.0 + short = LeaseLeadership( + kube, DGD, identity="a", lease_duration_s=15.0, margin_s=10.0 + ) + assert short.margin_s == 10.0 + short.set_owner(dataclasses.replace(DGD, uid="uid-2")) + assert short.step() == "acquired" + owner = kube.leases[short.name]["metadata"]["ownerReferences"][0] + assert owner["uid"] == "uid-2" + + +# The shortest lease the live supervisor accepts, with its request timeout. +K8S_TIMEOUT = LIVE_RETUNE_K8S_TIMEOUT_SECONDS +MIN_LEASE = live_retune_min_lease_seconds(K8S_TIMEOUT) + + +def minimum_leadership(kube, clock) -> LeaseLeadership: + return LeaseLeadership( + kube, + DGD, + identity="a/1", + lease_duration_s=MIN_LEASE, + margin_s=live_retune_lease_margin_seconds(K8S_TIMEOUT), + timeout_s=K8S_TIMEOUT, + monotonic=clock.monotonic, + clock=clock.time, + ) + + +INTERVAL = MIN_LEASE / 3.0 +BUDGET = live_retune_io_timeout_seconds(K8S_TIMEOUT) + + +def renewal_wait(a: LeaseLeadership) -> float: + """The supervisor's lease loop wait after a successful renewal.""" + + return a.renewal_wait_s(interval_s=INTERVAL, budget_s=BUDGET) + + +def test_leadership_renews_the_minimum_lease_in_one_term_until_it_lapses() -> None: + clock = FakeClock() + kube = FakeRetuneKube() + a = minimum_leadership(kube, clock) + assert renewal_wait(a) == INTERVAL + assert a.step() == "acquired" and a.acquisitions == 1 + acquired = kube.leases[a.name]["spec"]["acquireTime"] + # One I/O budget before leader_until, sooner than the renew interval. + margin = live_retune_lease_margin_seconds(K8S_TIMEOUT) + assert renewal_wait(a) == MIN_LEASE - margin - BUDGET < INTERVAL + for _ in range(10): + clock.advance(renewal_wait(a)) + assert a.leads_since(1) + assert a.step() == "renewed" + spec = kube.leases[a.name]["spec"] + assert a.transitions == 0 and spec["leaseTransitions"] == 0 + assert spec["acquireTime"] == acquired + # A failed attempt's floor wins over the cap, even past leader_until. + clock.advance(a.leader_until - clock.mono - 1.0) + assert a.renewal_wait_s(interval_s=INTERVAL, budget_s=BUDGET, floor_s=5.0) == 5.0 + clock.advance(5.0) + assert not a.is_leader() and renewal_wait(a) == INTERVAL + assert a.step() == "acquired" and a.acquisitions == 2 + assert a.is_leader() and not a.leads_since(1) + + +class SlowLeaseKube(FakeRetuneKube): + """Every Lease call takes ``latency_s`` of ``clock``. A replace calls + ``probe`` when its response is due, before the caller gets it.""" + + def __init__(self, clock: FakeClock, latency_s: float) -> None: + super().__init__(now=clock.time) + self.clock = clock + self.latency_s = latency_s + self.probe: Optional[Callable[[], bool]] = None + self.probes: list[bool] = [] + + def get_lease(self, name: str, *, timeout_s: float = 10.0) -> Optional[dict]: + self.clock.advance(self.latency_s) + return super().get_lease(name, timeout_s=timeout_s) + + def create_lease(self, body: dict, *, timeout_s: float = 10.0) -> dict: + self.clock.advance(self.latency_s) + return super().create_lease(body, timeout_s=timeout_s) + + def replace_lease(self, name: str, body: dict, *, timeout_s: float = 10.0) -> dict: + self.clock.advance(self.latency_s) + if self.probe is not None: + self.probes.append(self.probe()) + return super().replace_lease(name, body, timeout_s=timeout_s) + + +def test_leadership_never_lapses_while_slow_renewals_are_in_flight() -> None: + # The GET and the replace together take just under the I/O budget that + # bounds one renewal (the supervisor's backstop around step()); the + # client's (connect, read) timeout alone does not cap a call at K8S_TIMEOUT. + clock = FakeClock() + kube = SlowLeaseKube(clock, latency_s=BUDGET / 2 - 0.01) + a = minimum_leadership(kube, clock) + kube.probe = a.is_leader + assert a.step() == "acquired" + for _ in range(10): + clock.advance(renewal_wait(a)) + assert a.step() == "renewed" + assert kube.probes == [True] * 10 and a.transitions == 0 diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_supervisor.py b/components/src/dynamo/planner/tests/unit/test_live_retune_supervisor.py new file mode 100644 index 000000000000..e9864dddfd15 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_supervisor.py @@ -0,0 +1,3119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402 +# Optional-dependency preflight must run before the supervisor imports the sweep. + +"""Unit tests for ``LiveRetuneSupervisor`` and its bootstrap. + +The supervisor drives the real engine adapter, ``RetuneController`` and +``RouterActuator``; sweeps run on the adapter's sweep thread, in +``StubSweeps``, until the test finishes them. Kubernetes is +``FakeRetuneKube`` (shared Leases), the routers are ``FakeFrontend`` (PATCH +merges and bumps the revision, reset restores the initial) behind +``FrontendPolicyClient``, and clocks are advanced by hand. Kubernetes calls +run inline. +""" + +from __future__ import annotations + +import asyncio +import copy +import dataclasses +import itertools +import json +import logging +import math +import threading +from types import SimpleNamespace +from typing import Any, Optional + +import pytest + +pytest.importorskip( + "aisimulate.sweeper", + reason="AI Simulate is an optional Dynamo simulation dependency", +) + +from dynamo.planner.config.planner_config import ( + LIVE_RETUNE_K8S_TIMEOUT_SECONDS, + PlannerConfig, + RetuneEngineSpec, + live_retune_min_lease_seconds, +) +from dynamo.planner.core.types import ( + EngineCapabilities, + RetuneDecision, + RetuneInputs, + RetuneParams, + WorkerCapabilities, + WorkerCounts, +) +from dynamo.planner.plugins.clock import WallClock +from dynamo.planner.plugins.orchestrator.engine_adapter import ( + OrchestratorEngineAdapter, + RetuneHooks, +) +from dynamo.planner.retune import controller as controller_module +from dynamo.planner.retune.controller import RetuneController +from dynamo.planner.retune.live import bootstrap +from dynamo.planner.retune.live.actuator import RouterActuator +from dynamo.planner.retune.live.contracts import ( + LABEL_DGD_UID, + ActuationReport, + ActuationTarget, + DgdRef, + FrontendStartup, + KubeConflict, + KubeError, + NullMetrics, + PolicyCallResult, + PremiseReport, + RouterInfo, +) +from dynamo.planner.retune.live.state import ( + STATE_ANNOTATION, + RetuneState, + StateStore, + append_history, + leadership_lease_name, + state_lease_name, +) +from dynamo.planner.retune.live.supervisor import ( + LiveRetuneSupervisor, + SpecEvaluation, + SupervisorSettings, + _to_thread, + current_window_end_ms, + due_window_end_ms, + plan_initialization, + window_index, +) +from dynamo.planner.retune.sweep import SweepResult +from dynamo.planner.tests.live_retune_fakes import FakeRetuneKube +from dynamo.planner.tests.unit.test_live_retune_state import live_config + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +INTERVAL = 600.0 +GRACE = 30.0 +TIMEOUT = 540.0 +LEASE = 60.0 +START_WALL = 1_800_000_100.0 +DGD_NAME = "graph" +STATE_NAME = state_lease_name(DGD_NAME) +LEADER_NAME = leadership_lease_name(DGD_NAME) +CREDIT = "overlap_score_credit" +SCALE = "prefill_load_scale" +TEMP = "router_temperature" +KEYS = (CREDIT, SCALE, TEMP) +# The searched router domain of live_config(), by AISimulate knob name. +ROUTER = { + "overlap_score_credit": {"choices": [0.5, 1.0]}, + "prefill_load_scale": 1.0, + "temperature": {"choices": [0.0, 0.2]}, +} +# The same domain by live knob name. +DOMAIN = {CREDIT: [0.5, 1.0], SCALE: [1.0], TEMP: [0.0, 0.2]} +# Two frontend startups: the KvRouterConfig JSON the frontend parses to. +ROUTER_CONFIG = {"a": {CREDIT: 1.0}, "b": {CREDIT: 0.5}} +SPEC_COUNTS = WorkerCounts(ready_num_decode=2) +# The due window at START_WALL. +DUE0 = math.floor((START_WALL - GRACE) / INTERVAL) - 1 +# The window in progress at START_WALL (the one after DUE0). First enable +# records it as decided: neither it (its trace may have begun mid-way) nor +# the due window is swept. +K0 = math.floor(START_WALL / INTERVAL) +assert K0 == DUE0 + 1 + + +def resolver(variant: str) -> bootstrap.KvRouterResolver: + return bootstrap.KvRouterResolver(ROUTER_CONFIG[variant], ("aggregated",)) + + +def by_wt(variant: str, flat: dict[str, float]) -> dict[str, dict[str, float]]: + return resolver(variant).by_wt(flat) + + +def startup_flat(variant: str) -> dict[str, float]: + return resolver(variant).startup_flat(KEYS) + + +# policy_parameters("aggregated") of each startup: its routers' 8 parameters. +_DEFAULT_POLICY = { + "decode_active_request_weight": 0.0, + "disk_cache_hit_weight": 0.25, + "host_cache_hit_weight": 0.75, + "overlap_score_credit": 1.0, + "overlap_score_credit_decay": 0.0, + "prefill_load_scale": 1.0, + "router_temperature": 0.0, + "shared_cache_multiplier": 0.0, +} +VARIANTS = { + variant: {"aggregated": {**_DEFAULT_POLICY, **config}} + for variant, config in ROUTER_CONFIG.items() +} + + +def window(k: int) -> tuple[float, float, float]: + start = k * INTERVAL + return start, start + INTERVAL, start + INTERVAL + GRACE + + +def end_ms(k: int) -> int: + return int(round(window(k)[1] * 1000)) + + +# ------------------------------------------------------------------ fakes + + +class Clock: + def __init__(self) -> None: + self.wall = START_WALL + self.mono = 10_000.0 + + def time(self) -> float: + return self.wall + + def monotonic(self) -> float: + return self.mono + + def advance(self, seconds: float) -> None: + self.wall += seconds + self.mono += seconds + + +class AdapterClock(WallClock): + """The engine adapter's clock: monotonic time is the test's.""" + + def __init__(self, clock: Clock) -> None: + self._clock = clock + + def now(self) -> float: + return self._clock.time() + + def monotonic(self) -> float: + return self._clock.monotonic() + + +async def inline(call): + return call() + + +async def frozen_sleep(_seconds: float) -> None: + """The test's clock moves only when the test moves it: no sleep ends.""" + + await asyncio.get_running_loop().create_future() + + +class RecordingMetrics(NullMetrics): + def __init__(self) -> None: + self.events: list[tuple[str, Any]] = [] + self.state: dict[str, Any] = {} + + def decision(self, outcome: str) -> None: + self.events.append(("decision", outcome)) + + def failure(self, reason: str) -> None: + self.events.append(("failure", reason)) + + def skipped(self, reason: str) -> None: + self.events.append(("skipped", reason)) + + def supervisor_state(self, **kwargs: Any) -> None: + self.state = kwargs + + def scores(self, **kwargs: Any) -> None: + self.events.append(("scores", kwargs)) + + def dropped_members(self, count: int) -> None: + self.events.append(("dropped_members", count)) + + def of(self, kind: str) -> list[Any]: + return [value for event, value in self.events if event == kind] + + +@dataclasses.dataclass +class Router: + router_id: str + worker_type: str + revision: int + config: dict[str, float] + initial: dict[str, float] + manual: bool = False + + +FRONTEND_POD = "frontend-0" + + +class FakeFrontend: + """The DGD's frontend pod and its routers, with the runtime.rs + semantics: PATCH merges and bumps the revision, reset restores the + initial and bumps the revision. A restart is a new frontend process in + the same pod (a container restart): new routers at revision 0.""" + + def __init__(self, kube: FakeRetuneKube, maps: dict[str, dict[str, float]]): + self._ids = itertools.count() + self.kube = kube + self.down = False + # (writer, advisory, kind, router_id) + self.writes: list[tuple[str, bool, str, str]] = [] + self.routers: dict[str, Router] = {} + kube.add_pod(FRONTEND_POD, dgd=DGD_NAME) + self._start(maps) + + def _start(self, maps: dict[str, dict[str, float]]) -> None: + self.routers = {} + for worker_type, params in maps.items(): + router_id = f"r{next(self._ids)}" + self.routers[router_id] = Router( + router_id, worker_type, 0, dict(params), dict(params) + ) + + def restart(self, maps: dict[str, dict[str, float]]) -> None: + self.kube.restart_pod(FRONTEND_POD) + self._start(maps) + + def write(self, router: Router, kind: str, params, writer, advisory) -> None: + if kind == "reset": + router.config = dict(router.initial) + else: + router.config.update(params) + router.revision += 1 + self.writes.append((writer, advisory, kind, router.router_id)) + + def manual_patch(self, key: str, value: float) -> None: + for router in self.routers.values(): + router.config[key] = value + router.revision += 1 + + +class FrontendPolicyClient: + """``RouterPolicyClient`` of the real ``RouterActuator`` over the + world's ``FakeFrontend``: revision compare-and-swap, then the write. + ``writer`` and ``advisory`` label the writes it makes.""" + + def __init__(self, world: World, writer: str, advisory: bool) -> None: + self.world = world + self.writer = writer + self.advisory = advisory + + def _info(self, router: Router) -> RouterInfo: + return RouterInfo( + router_id=router.router_id, + namespace=self.world.spec.dynamo_namespace, + worker_type=router.worker_type, + revision=router.revision, + configuration=dict(router.config), + model="model", + ) + + async def get(self, pod) -> PolicyCallResult: + frontend = self.world.frontend + if frontend.down: + return PolicyCallResult(status=0, error="connection refused") + routers = tuple(self._info(r) for r in frontend.routers.values()) + return PolicyCallResult(status=200, routers=routers) + + async def patch(self, pod, router_id, expected_revision, parameters): + return self._write(router_id, expected_revision, "patch", dict(parameters)) + + async def reset(self, pod, router_id, expected_revision): + return self._write(router_id, expected_revision, "reset", None) + + def _write(self, router_id, expected_revision, kind, params) -> PolicyCallResult: + frontend = self.world.frontend + router = frontend.routers.get(router_id) + if router is None: + return PolicyCallResult(status=404, error="router instance not found") + if router.revision != expected_revision: + return PolicyCallResult(status=409, error="policy revision conflict") + frontend.write(router, kind, params, self.writer, self.advisory) + return PolicyCallResult(status=200, router=self._info(router)) + + +def real_actuators(world: World, writer: str, advisory: bool = False): + """The ``actuator_factory`` of a process: ``RouterActuator`` over the + world's frontend.""" + + return lambda dgd, maps: RouterActuator( + world.kube, + FrontendPolicyClient(world, writer, advisory), + dgd, + maps, + NullMetrics(), + ) + + +class StubSweeps: + """The ``sweep_runner`` of every process: a sweep runs on its adapter's + sweep thread until the test finishes it, or until its abort or cancel + event is set. Results are made on the test's thread, so a seeded run + is deterministic.""" + + @dataclasses.dataclass + class Running: + cancel: threading.Event + abort: Optional[threading.Event] + done: threading.Event = dataclasses.field(default_factory=threading.Event) + result: Optional[SweepResult] = None + + def stopping(self) -> bool: + return self.cancel.is_set() or ( + self.abort is not None and self.abort.is_set() + ) + + def __init__(self) -> None: + self._cond = threading.Condition() + self._running: dict[int, StubSweeps.Running] = {} + self._closed = False + self.requests: list[Any] = [] + + def __call__(self, request, *, cancel, abort, context): + running = StubSweeps.Running(cancel, abort) + with self._cond: + self._running[id(request)] = running + self.requests.append(request) + self._cond.notify_all() + try: + while not running.done.wait(0.002): + if self._closed or cancel.is_set(): + return SweepResult(status="failed", failure_reason="shutdown") + if abort is not None and abort.is_set(): + return SweepResult( + status="failed", failure_reason="lost_leadership" + ) + assert running.result is not None + return running.result + finally: + with self._cond: + del self._running[id(request)] + self._cond.notify_all() + + def finish(self, controller: RetuneController, result_for) -> bool: + """Finish the controller's pending sweep with ``result_for(request)`` + and wait until it returned; ``False`` when none is pending.""" + + pending = controller._pending + if pending is None or pending.future.done(): + return False + with self._cond: + assert self._cond.wait_for( + lambda: id(pending.request) in self._running or pending.future.done(), + timeout=10, + ) + running = self._running.get(id(pending.request)) + if running is None: + return False + running.result = result_for(pending.request) + running.done.set() + pending.future.result(timeout=10) + return True + + def quiesce(self, controllers) -> None: + """Wait until every stopped sweep returned and every pending sweep + runs, so no sweep starts after the test's clock moved on.""" + + def quiet() -> bool: + if any(running.stopping() for running in self._running.values()): + return False + for controller in controllers: + pending = controller._pending + if ( + pending is not None + and not pending.future.done() + and id(pending.request) not in self._running + ): + return False + return True + + with self._cond: + assert self._cond.wait_for(quiet, timeout=10) + + def close(self) -> None: + with self._cond: + self._closed = True + + +def failed(request) -> SweepResult: + return SweepResult(status="failed", error="injected", failure_reason="crn_mismatch") + + +def scored(best=None): + """A sweep result whose ``best`` candidate (else none) beats the rest.""" + + def result_for(request) -> SweepResult: + return SweepResult( + status="ok", + num_requests=10, + scores={ + candidate.key(): (2.0 if candidate.router == best else 1.0) + for candidate in request.candidates + }, + wall_s=1.0, + ) + + return result_for + + +class TracedController(RetuneController): + """The real controller, remembering the windows the supervisor dropped.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.dropped: list[int] = [] + + def drop_pending(self) -> None: + if self._pending is not None: + self.dropped.append(self._pending.request.window_index) + super().drop_pending() + + +@pytest.fixture(autouse=True) +def traced_controllers(monkeypatch): + # The engine adapter looks the controller class up when it builds one. + monkeypatch.setattr(controller_module, "RetuneController", TracedController) + yield + while WORLDS: + WORLDS.pop().close() + + +def controller_config(*, advisory: bool = False) -> PlannerConfig: + """The real controller's config, with this file's windows. + + Built with the replay executor: the live config's timeout must leave + 600 s before the next window, which 600 s windows cannot, and the + controller never reads the executor. + """ + + return PlannerConfig( + namespace="test-ns", + mode="agg", + advisory=advisory, + enable_load_scaling=False, + enable_throughput_scaling=False, + retune={ + "interval_seconds": INTERVAL, + "grace_seconds": GRACE, + "timeout_seconds": TIMEOUT, + "router": ROUTER, + }, + ) + + +def retune_inputs(variant: str) -> RetuneInputs: + return RetuneInputs( + engine=RetuneEngineSpec( + router_mode="kv_router", router_config=dict(ROUTER_CONFIG[variant]) + ), + trace_paths=("/trace/segments/graph",), + ) + + +def variant_of(inputs: RetuneInputs) -> str: + (variant,) = [ + v for v, c in ROUTER_CONFIG.items() if c == inputs.engine.router_config + ] + return variant + + +def make_decision( + k: int, + outcome: str, + previous: Optional[RetuneParams] = None, + *, + chosen: Optional[dict[str, float]] = None, + failure_reason: Optional[str] = None, + details: Optional[dict[str, Any]] = None, +) -> RetuneDecision: + previous = previous or RetuneParams(router=startup_flat("a")) + start_s, end_s, as_of_s = window(k) + return RetuneDecision( + decision_id=f"retune-w{end_ms(k)}", + outcome=outcome, + window_index=k, + window_start_s=start_s, + window_end_s=end_s, + as_of_s=as_of_s, + launched_at_s=None if outcome == "busy" else as_of_s, + decided_at_s=as_of_s + 60.0, + num_requests=10, + num_candidates=4, + previous=previous, + chosen=previous if chosen is None else RetuneParams(router=chosen), + baseline_score=1.0, + best_score=1.5, + sweep_wall_s=50.0, + failure_reason=failure_reason, + details=details or {}, + ) + + +class FakeHost: + """``NativePlannerBase`` as the supervisor sees it: the config lock and + the real engine adapter, whose controllers share one sweep thread.""" + + def __init__( + self, + world: World, + hooks: RetuneHooks, + inputs: Optional[RetuneInputs], + *, + advisory: bool, + ) -> None: + self.clock = world.clock + self._config_lock = asyncio.Lock() + self._engine = OrchestratorEngineAdapter( + controller_config(advisory=advisory), + WorkerCapabilities( + decode=EngineCapabilities( + num_gpu=1, max_num_batched_tokens=2048, max_kv_tokens=16384 + ) + ), + clock=AdapterClock(world.clock), + retune_inputs=inputs, + retune_hooks=hooks, + ) + if self.controller is not None: + # What the adapter's initial_tick does. + self.controller.start(world.clock.time(), hooks.phase_s) + # The frontend startup the current controller was built for. + self.variant = "a" if inputs is None else variant_of(inputs) + self.replacements: list[tuple[Any, int]] = [] + self.fail_replace = False + + @property + def controller(self) -> Optional[TracedController]: + return self._engine.retune_controller + + async def replace_live_retune(self, inputs, *, generation): + async with self._config_lock: + self.replacements.append((inputs, generation)) + now = self.clock.time() + if self.fail_replace: + self._engine.replace_retune(None, generation=generation, now_s=now) + raise ValueError("controller construction failed") + self._engine.replace_retune(inputs, generation=generation, now_s=now) + if inputs is not None: + self.variant = variant_of(inputs) + + def close(self) -> None: + if self.controller is not None: + self.controller.close() + executor = self._engine._retune_executor + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) + + +class Spec: + """The DGD spec as the premise evaluator sees it (tests mutate it).""" + + def __init__(self) -> None: + self.variant = "a" + self.spec_ok = True + self.p4_ok = True + self.replicas: Optional[dict[str, int]] = {"aggregated": 2} + self.uid = "uid-1" + self.dynamo_namespace = "dyn" + self.frontend_replicas = 1 + + def startup(self) -> FrontendStartup: + return FrontendStartup( + router_mode="kv_router", + router_config=dict(ROUTER_CONFIG[self.variant]), + startup_maps=copy.deepcopy(VARIANTS[self.variant]), + frontend_component="Frontend", + replicas=self.frontend_replicas, + replica_sync=False, + fingerprint_source=self.variant, + ) + + def evaluation(self, _dgd: Optional[dict] = None) -> SpecEvaluation: + reasons = [] if self.p4_ok else ["P4: router mode round-robin"] + if not self.spec_ok: + reasons.append("P8: two frontends without replica sync") + return SpecEvaluation( + report=PremiseReport( + p4_ok=self.p4_ok, + spec_ok=self.p4_ok and self.spec_ok, + reasons=tuple(reasons), + ), + startup=self.startup(), + inputs=retune_inputs(self.variant) if self.p4_ok else None, + spec_replicas=None if self.replicas is None else dict(self.replicas), + dgd=DgdRef(DGD_NAME, "test-ns", self.uid, self.dynamo_namespace), + ) + + +# Worlds of the running test, closed after it (sweep threads end). +WORLDS: list[World] = [] + + +class World: + """Shared kube, frontend, spec, clock and sweeps of one test.""" + + def __init__(self, kube: Optional[FakeRetuneKube] = None) -> None: + self.clock = Clock() + self.kube = kube or FakeRetuneKube(now=self.clock.time) + self.spec = Spec() + self.frontend = FakeFrontend(self.kube, VARIANTS["a"]) + self.sweeps = StubSweeps() + self.hosts: list[FakeHost] = [] + self.kube.put_dgd( + {"metadata": {"name": DGD_NAME, "uid": "uid-1"}, "spec": {"components": []}} + ) + WORLDS.append(self) + + def shared_state(self) -> Optional[RetuneState]: + lease = self.kube.leases.get(STATE_NAME) + if lease is None: + return None + return RetuneState.from_annotations(lease["metadata"].get("annotations")) + + def holder(self) -> Optional[str]: + lease = self.kube.leases.get(LEADER_NAME) + return None if lease is None else lease["spec"].get("holderIdentity") + + def close(self) -> None: + self.sweeps.close() + for host in self.hosts: + host.close() + + +@dataclasses.dataclass +class Process: + name: str + sup: LiveRetuneSupervisor + host: FakeHost + metrics: RecordingMetrics + advisory: bool + kube: Any + alive: bool = True + + @property + def controller(self) -> Optional[TracedController]: + return self.host.controller + + +def settings_for( + *, advisory: bool = False, promote: bool = False, **changes: Any +) -> SupervisorSettings: + values: dict[str, Any] = dict( + interval_s=INTERVAL, + grace_s=GRACE, + timeout_s=TIMEOUT, + advisory=advisory, + promote_last_advisory=promote, + router_keys=KEYS, + lease_duration_s=LEASE, + actuator_interval_s=15.0, + k8s_timeout_s=5.0, + write_failure_limit=3, + ) + values.update(changes) + return SupervisorSettings(**values) + + +def make_process( + world: World, + name: str = "p1", + *, + advisory: bool = False, + promote: bool = False, + kube: Any = None, + canary=None, + reconciler=None, + sweep_runner=None, + context_factory=lambda supervisor: None, + settings: Optional[SupervisorSettings] = None, + blocking=None, + actuator_factory=None, +) -> Process: + kube = kube or world.kube + metrics = RecordingMetrics() + evaluation = world.spec.evaluation() + settings = settings or settings_for(advisory=advisory, promote=promote) + sup = LiveRetuneSupervisor( + settings=settings, + kube=kube, + evaluation=evaluation, + identity=f"{name}/{len(world.kube.calls)}", + premise_evaluator=world.spec.evaluation, + actuator_factory=actuator_factory or real_actuators(world, name, advisory), + fingerprint_fn=lambda startup: f"fp-{None if startup is None else startup.fingerprint_source}", + resolver_factory=lambda startup: bootstrap.KvRouterResolver( + startup.router_config, ("aggregated",) + ), + reconciler=reconciler, + metrics=metrics, + canary=canary, + sweep_runner=sweep_runner or world.sweeps, + context_factory=context_factory, + clock=world.clock.time, + monotonic=world.clock.monotonic, + sleep=frozen_sleep, + blocking=blocking or inline, + ) + host = FakeHost(world, sup.hooks(), evaluation.inputs, advisory=advisory) + world.hosts.append(host) + return Process(name, sup, host, metrics, advisory, kube) + + +async def started(world: World, name: str = "p1", **kwargs) -> Process: + process = make_process(world, name, **kwargs) + await process.sup.start(process.host, background=False) + return process + + +async def leading(world: World, name: str = "p1", **kwargs) -> Process: + """A started process that leads, is initialized and is ready.""" + + process = await started(world, name, **kwargs) + await process.sup.lease_round() + assert process.sup.is_leader() and process.sup.initialized + await process.sup.control_round() + return process + + +async def settle(processes: list[Process], world: World, rounds: int = 6) -> None: + for _ in range(rounds): + world.clock.advance(20) + for process in processes: + if process.alive: + await process.sup.lease_round() + await process.sup.outbox_round() + await process.sup.control_round() + + +def enqueue(process: Process, *decisions: RetuneDecision) -> None: + """Hand the supervisor a tick's decisions with the controller's routers.""" + + controller = process.controller + desired = ( + controller.current_point().router + if controller is not None + else startup_flat("a") + ) + process.sup.on_effects( + SimpleNamespace( + retune_decisions=list(decisions), + router_parameters=by_wt(process.host.variant, desired), + ) + ) + + +def frontend_configs(world: World) -> list[dict[str, float]]: + return [router.config for router in world.frontend.routers.values()] + + +# ------------------------------------------------------- pure functions + + +def claimed(**changes) -> RetuneState: + base = RetuneState(dgd_uid="uid-1", leader_term="1/p/2") + return dataclasses.replace(base, **changes) + + +def deployed(**changes) -> RetuneState: + values = dict( + fingerprint="fp", + mode="active", + deploy_seq=4, + flat={CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2}, + by_wt=by_wt("a", {CREDIT: 0.5, TEMP: 0.2}), + last_window_end_ms=777, + ) + values.update(changes) + return claimed(**values) + + +def plan(state: RetuneState, **overrides): + kwargs = dict( + spec_ok=True, + advisory=False, + fingerprint="fp", + startup_flat=startup_flat("a"), + startup_by_wt=VARIANTS["a"], + router_keys=KEYS, + promote=False, + now_ms=10_000_000, + interval_ms=int(INTERVAL * 1000), + first_window_end_ms=9_999, + ) + kwargs.update(overrides) + return plan_initialization(state, **kwargs) + + +def test_init_table_first_enable_persists_the_due_window() -> None: + result = plan(claimed()) + assert result.row == "active_reset" and result.write + assert result.state.mode == "active" and not result.state.rollback_pending + assert result.state.by_wt == VARIANTS["a"] and result.state.flat == startup_flat( + "a" + ) + assert result.state.last_window_end_ms == 9_999 + assert result.state.deploy_seq == 1 + assert result.restore_flat == startup_flat("a") + + +def test_init_table_spec_not_ok_is_the_first_row() -> None: + # Deployed X: roll back (startup, pending), controller untouched. + result = plan(deployed(), spec_ok=False) + assert result.row == "spec_not_ok" and result.restore_flat is None + assert result.state.rollback_pending and result.state.by_wt == VARIANTS["a"] + assert result.state.deploy_seq == 5 + # Nothing deployed and the record current: nothing to write. + clean = deployed(by_wt=VARIANTS["a"], flat=startup_flat("a")) + assert not plan(clean, spec_ok=False).write + # Even an advisory planner takes this row (with its own mode). + advisory = plan(deployed(), spec_ok=False, advisory=True) + assert advisory.row == "spec_not_ok" and advisory.state.mode == "advisory" + assert advisory.state.rollback_pending + # Unknown startup and an active record: roll back to the latched initials. + unknown = plan( + deployed(by_wt=VARIANTS["a"]), + spec_ok=False, + startup_flat=None, + startup_by_wt=None, + ) + assert unknown.state.rollback_pending and unknown.state.by_wt == {} + + +def test_init_table_advisory_rows() -> None: + rollback = plan(deployed(), advisory=True) + assert rollback.row == "advisory_rollback" and rollback.restore_flat is None + assert rollback.state.mode == "advisory" and rollback.state.rollback_pending + assert rollback.state.by_wt == VARIANTS["a"] + + observe = plan(deployed(by_wt=VARIANTS["a"], flat=startup_flat("a")), advisory=True) + assert observe.row == "advisory_observe" and observe.write + assert observe.state.mode == "advisory" and not observe.state.rollback_pending + # Already advisory and current: no write. + again = plan(observe.state, advisory=True) + assert again.row == "advisory_observe" and not again.write + # An advisory record of an older startup is not a deployment. + older = plan(observe.state, advisory=True, startup_by_wt=VARIANTS["b"]) + assert older.row == "advisory_observe" and not older.state.rollback_pending + + +def test_init_table_active_rows() -> None: + restore = plan(deployed()) + assert restore.row == "active_restore" and not restore.write + assert restore.restore_flat == {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + + for state, pending in ( + (deployed(fingerprint="old"), True), + (deployed(mode="advisory"), False), + (deployed(rollback_pending=True), True), + (deployed(flat={CREDIT: 0.5}), True), + (deployed(fingerprint="old", by_wt=VARIANTS["a"]), False), + ): + result = plan(state) + assert result.row == "active_reset", state + assert result.state.rollback_pending is pending, state + assert result.state.by_wt == VARIANTS["a"] and result.state.mode == "active" + assert result.restore_flat == startup_flat("a") + assert result.state.last_window_end_ms == 777 + + +def test_init_table_promotes_a_recent_matching_advisory() -> None: + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + entry = { + "decision_id": "retune-w9000000", + "outcome": "advisory", + "window_end_ms": 9_900_000, + "fingerprint": "fp", + "flat": chosen, + "by_wt": by_wt("a", chosen), + } + state = append_history( + deployed(mode="advisory", by_wt=VARIANTS["a"], flat=startup_flat("a")), entry + ) + result = plan(state, promote=True) + assert result.row == "promote" and result.restore_flat == chosen + assert result.state.mode == "active" and result.state.by_wt == by_wt("a", chosen) + assert not result.state.rollback_pending and result.state.deploy_seq == 5 + assert result.state.history[-1]["outcome"] == "promoted" + # Not promoted: disabled, stale fingerprint, too old, or not the latest. + assert plan(state).row == "active_reset" + assert plan(state, promote=True, fingerprint="new").row == "active_reset" + assert plan(state, promote=True, now_ms=9_900_000 + 600_001).row == "active_reset" + later = append_history(state, {"decision_id": "x", "outcome": "kept"}) + assert plan(later, promote=True).row == "active_reset" + + +def test_window_arithmetic_matches_the_controller() -> None: + now = START_WALL + end = due_window_end_ms(now, phase_s=0.0, interval_s=INTERVAL, grace_s=GRACE) + assert end == end_ms(DUE0) + assert window_index(end, phase_s=0.0, interval_s=INTERVAL) == DUE0 + # Exactly at L_k the window is due. + at_launch = window(DUE0 + 1)[2] + assert due_window_end_ms( + at_launch, phase_s=0.0, interval_s=INTERVAL, grace_s=GRACE + ) == end_ms(DUE0 + 1) + phased = due_window_end_ms(1000.0, phase_s=100.0, interval_s=300.0, grace_s=0.0) + assert phased == 1_000_000 + assert window_index(phased, phase_s=100.0, interval_s=300.0) == 2 + + +def test_current_window_is_after_the_due_and_the_not_yet_due_windows() -> None: + def current(now_s: float, phase_s: float = 0.0) -> int: + return current_window_end_ms(now_s, phase_s=phase_s, interval_s=INTERVAL) + + # Mid-window: the window in progress, one after the due window. + assert current(START_WALL) == end_ms(K0) == end_ms(DUE0 + 1) + # Within the grace after a window ended: that window is not due yet and + # the one in progress is two after the due window; both are skipped. + start = window(K0)[0] + within_grace = start + GRACE / 2 + due = due_window_end_ms( + within_grace, phase_s=0.0, interval_s=INTERVAL, grace_s=GRACE + ) + assert due == end_ms(K0 - 2) and current(within_grace) == end_ms(K0) + # A window boundary belongs to the window starting there. + assert current(start) == end_ms(K0) + # Phased: windows [100 + k T, 100 + (k + 1) T); 1000 is in [700, 1300). + assert current(1000.0, phase_s=100.0) == 1_300_000 + + +# ----------------------------------------------------------- gate/harvest + + +async def test_gate_reasons_and_the_skip_metric() -> None: + world = World() + process = await started(world) + sup = process.sup + assert sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "not_leader" + assert sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "not_leader" + await sup.lease_round() + # Initialized but no health yet (no actuator round). + assert sup.initialized + assert sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "not_ready" + await sup.control_round() + assert sup.ready() + # First enable recorded the window in progress at start as decided: it + # and the due window before it are skipped. + assert world.shared_state().last_window_end_ms == end_ms(K0) + assert sup.launch_gate(DUE0, window(DUE0)).reason == "already_decided" + assert sup.launch_gate(K0, window(K0)).reason == "already_decided" + gate = sup.launch_gate(K0 + 1, window(K0 + 1)) + assert gate.reason is None and gate.spec_replicas == {"aggregated": 2} + assert isinstance(gate.abort, threading.Event) and not gate.abort.is_set() + # Each (k, reason) counts once. + assert process.metrics.of("skipped") == [ + "not_leader", + "not_ready", + "already_decided", + "already_decided", + ] + # Health and spec gate only launches. + world.spec.replicas = None + await sup.control_round() + assert sup.launch_gate(K0 + 2, window(K0 + 2)).reason == "not_ready" + + +async def test_harvest_allowed_needs_leader_initialized_and_generation() -> None: + world = World() + process = await leading(world) + sup = process.sup + assert sup.harvest_allowed(0) + assert not sup.harvest_allowed(1) + world.spec.spec_ok = False + await sup.control_round() # rebuild: generation 1 + assert sup.generation == 1 + assert not sup.harvest_allowed(0) and sup.harvest_allowed(1) + world.clock.advance(LEASE) # leadership lapses without renewal + assert not sup.harvest_allowed(1) + + +# ---------------------------------------------------------------- outbox + + +async def test_applied_decision_is_persisted_then_deployed() -> None: + world = World() + process = await leading(world) + sup, controller = process.sup, process.controller + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + controller.restore_desired(chosen) + enqueue(process, make_decision(K0 + 1, "applied", chosen=chosen)) + assert sup.outbox_size == 1 + assert process.metrics.of("decision") == [] # counted only once persisted + await sup.outbox_round() + state = world.shared_state() + assert sup.outbox_size == 0 + assert state.flat == chosen and state.by_wt == by_wt("a", chosen) + assert state.mode == "active" and state.deploy_seq == 2 + assert state.last_window_end_ms == end_ms(K0 + 1) + entry = state.history[-1] + assert entry["decision_id"] == f"retune-w{end_ms(K0 + 1)}" + assert entry["outcome"] == "applied" and entry["by_wt"] == by_wt("a", chosen) + assert entry["fingerprint"] == "fp-a" and entry["baseline"] == 1.0 + assert process.metrics.of("decision") == ["applied"] + assert sup.target() == ActuationTarget( + mode="apply", desired_by_wt=by_wt("a", chosen) + ) + await sup.control_round() + assert frontend_configs(world) == [by_wt("a", chosen)["aggregated"]] + # The next window launches; the decided one does not. + assert sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "already_decided" + + +async def test_advisory_decision_records_the_recommendation_only() -> None: + world = World() + process = await leading(world, advisory=True) + sup = process.sup + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + enqueue( + process, + make_decision( + K0 + 1, + "advisory", + chosen=chosen, + details={"runtime_image": "img:1", "missing_candidates": 2}, + ), + ) + await sup.outbox_round() + state = world.shared_state() + assert state.mode == "advisory" and state.by_wt == VARIANTS["a"] + entry = state.history[-1] + assert entry["outcome"] == "advisory" and entry["flat"] == chosen + assert entry["by_wt"] == by_wt("a", chosen) + assert entry["image"] == {"runtime_image": "img:1"} + scores = process.metrics.of("scores")[-1] + assert scores["missing_candidates"] == 2 and scores["best"] == 1.5 + assert sup.target() == ActuationTarget(mode="observe") + world.frontend.manual_patch(TEMP, 0.7) + await settle([process], world) + assert world.frontend.writes == [] + + +async def test_failed_decision_counts_its_reason() -> None: + world = World() + process = await leading(world) + enqueue( + process, + make_decision( + K0 + 1, + "failed", + failure_reason="crn_mismatch", + details={"window": {"dropped_members": 3, "cut": True}}, + ), + ) + await process.sup.outbox_round() + assert process.metrics.of("failure") == ["crn_mismatch"] + assert process.metrics.of("dropped_members") == [3] + assert world.shared_state().history[-1]["failure_reason"] == "crn_mismatch" + + +async def test_outbox_timeout_that_landed_and_conflict() -> None: + world = World() + process = await leading(world) + sup = process.sup + # A write that timed out but landed: the re-read finds it. + world.kube.fail( + "replace_lease", + KubeError("timeout"), + landed=True, + match=lambda target: target[0] == STATE_NAME, + ) + enqueue(process, make_decision(K0 + 1, "kept")) + await sup.outbox_round() + ids = [item["decision_id"] for item in world.shared_state().history] + assert ids.count(f"retune-w{end_ms(K0 + 1)}") == 1 + assert sup.outbox_size == 0 and process.metrics.of("decision") == ["kept"] + # A stale cached resourceVersion: conflict, re-read, retry. + world.kube.bump(world.kube.leases[STATE_NAME]) + enqueue(process, make_decision(K0 + 2, "kept")) + await sup.outbox_round() + ids = [item["decision_id"] for item in world.shared_state().history] + assert ids[-2:] == [f"retune-w{end_ms(K0 + 1)}", f"retune-w{end_ms(K0 + 2)}"] + assert sup.outbox_size == 0 and process.metrics.of("decision") == ["kept", "kept"] + + +async def test_consecutive_write_failures_reinitialize_and_rewind() -> None: + world = World() + process = await leading(world) + sup, controller = process.sup, process.controller + term_seq = sup.term_seq + # Window K0 + 2 launches (K0 + 1 is skipped: K0 + 2 is due already). + world.clock.wall = window(K0 + 2)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + assert controller._pending.request.window_index == K0 + 2 + world.kube.fail( + "replace_lease", + KubeError("timeout"), + times=3, + match=lambda target: target[0] == STATE_NAME, + ) + enqueue(process, make_decision(K0 + 1, "kept")) + for _ in range(3): + await sup.outbox_round() + assert not sup.initialized # re-init requested + assert process.metrics.of("failure") == ["state_write"] * 3 + await sup.lease_round() + assert sup.initialized and sup.term_seq == term_seq + 1 + assert sup.outbox_size == 0 + # The sweep was dropped and the controller rewound to the last decided + # window, so both windows can still run. + assert controller._pending is None and controller.dropped == [K0 + 2] + assert controller._last_k == K0 + assert world.shared_state().leader_term == sup.term + + +async def test_every_failed_state_write_counts_as_state_write() -> None: + world = World() + process = await started(world) + sup = process.sup + + def state_lease(target) -> bool: + return target[0] == STATE_NAME + + def failures() -> list[str]: + return process.metrics.of("failure") + + # First enable: the claim creates the state, the plan write fails. + world.kube.fail("replace_lease", KubeError("quota"), match=state_lease) + await sup.lease_round() + assert not sup.initialized and failures() == ["state_write"] + await sup.lease_round() + assert sup.initialized + # The claim write of a re-initialization. + world.kube.fail("replace_lease", KubeError("forbidden"), match=state_lease) + sup._request_reinit("test") + await sup.lease_round() + assert not sup.initialized and failures() == ["state_write"] * 2 + await sup.lease_round() + assert sup.initialized + # The initials of a restarted frontend (retried next round). + world.frontend.restart(VARIANTS["a"]) + world.kube.fail("replace_lease", KubeError("too large"), match=state_lease) + await sup.control_round() + assert failures() == ["state_write"] * 3 + await sup.control_round() + assert set(world.shared_state().initials) == set(world.frontend.routers) + # A lost term is not a failed write. + store = StateStore(world.kube, sup.dgd) + snapshot = store.read() + store.write(dataclasses.replace(snapshot.state, leader_term="9/x/1"), snapshot) + enqueue(process, make_decision(K0 + 1, "kept")) + await sup.outbox_round() + assert not sup.initialized and failures() == ["state_write"] * 3 + + +async def test_decisions_are_dropped_and_rewound_without_leadership() -> None: + world = World() + process = await started(world) + controller = process.controller + controller._last_k = K0 + 1 + enqueue(process, make_decision(K0 + 1, "kept")) + assert process.sup.outbox_size == 0 + assert controller._last_k == -1 # no state: every window is due again + + +async def test_a_foreign_term_forces_reinitialization() -> None: + world = World() + process = await leading(world) + sup = process.sup + store = StateStore(world.kube, sup.dgd) + snapshot = store.read() + store.write(dataclasses.replace(snapshot.state, leader_term="9/x/1"), snapshot) + enqueue(process, make_decision(K0 + 1, "kept")) + await sup.outbox_round() + assert not sup.initialized and sup.outbox_size == 1 + await sup.lease_round() + assert sup.initialized and world.shared_state().leader_term == sup.term + assert sup.outbox_size == 0 + + +async def test_a_late_write_of_an_earlier_term_fails() -> None: + world = World() + process = await leading(world) + sup = process.sup + store = StateStore(world.kube, sup.dgd) + stale = store.read() + # Same-process re-initialization (term_seq + 1) moves the term on. + sup._request_reinit("test") + await sup.lease_round() + assert world.shared_state().leader_term == sup.term != stale.state.leader_term + with pytest.raises(KubeConflict): + store.write(append_history(stale.state, {"decision_id": "late"}), stale) + + +def steal_lease(world: World, identity: str) -> None: + """The leadership Lease changes hands without expiring (an admin edit, + or a contender whose clock runs fast).""" + + lease = world.kube.leases[LEADER_NAME] + lease["spec"]["holderIdentity"] = identity + lease["spec"]["leaseTransitions"] = lease["spec"].get("leaseTransitions", 0) + 1 + world.kube.bump(lease) + + +async def test_a_stolen_lease_rejects_the_old_leaders_writes() -> None: + """Two processes believe they lead; the state term precondition keeps + the old one from writing once the new one has claimed.""" + + world = World() + old = await leading(world) + new = await started(world, "p2") + steal_lease(world, new.sup.identity) + await new.sup.lease_round() + assert new.sup.initialized and old.sup.is_leader() # both believe + enqueue(old, make_decision(K0 + 1, "kept")) + await old.sup.outbox_round() + state = world.shared_state() + assert state.leader_term == new.sup.term + assert all( + item["decision_id"] != f"retune-w{end_ms(K0 + 1)}" for item in state.history + ) + assert not old.sup.initialized + await old.sup.lease_round() # sees the new holder: local loss + assert not old.sup.is_leader() and old.sup.term is None + + +async def test_a_stolen_lease_aborts_the_old_leaders_claim() -> None: + world = World() + old = await leading(world) + term = world.shared_state().leader_term + steal_lease(world, "intruder/0") + assert old.sup.is_leader() # not seen yet + world.spec.spec_ok = False + await old.sup.control_round() # a rebuild re-initializes: no claim + assert world.shared_state().leader_term == term + assert not old.sup.is_leader() and not old.sup.initialized + + +# ----------------------------------------------------- initialization + + +async def test_restart_restores_the_persisted_point() -> None: + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.stop() + assert world.holder() is None # released + second = await leading(world, "p2") + assert second.controller.current_point().router == chosen + # The persisted window is decided: the gate skips it, the next launches. + assert second.sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "already_decided" + assert second.sup.launch_gate(K0 + 2, window(K0 + 2)).reason is None + state = world.shared_state() + assert state.leader_term == second.sup.term and state.deploy_seq == 2 + assert second.sup.target().desired_by_wt == by_wt("a", chosen) + + +async def test_fingerprint_change_resets_with_a_two_phase_rollback() -> None: + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.control_round() + assert frontend_configs(world) == [by_wt("a", chosen)["aggregated"]] + # The frontend args change (new startup): restart under it. + await first.sup.stop() + world.spec.variant = "b" + world.frontend.restart(VARIANTS["b"]) + world.frontend.manual_patch(TEMP, 0.9) + second = await leading(world, "p2") + state = world.shared_state() + assert state.rollback_pending and state.by_wt == VARIANTS["b"] + assert state.fingerprint == "fp-b" and state.deploy_seq == 3 + assert second.controller.current_point().router == startup_flat("b") + assert second.sup.target().mode == "rollback" + # Reset once, then converged: completes after a lease duration. + resets = len(world.frontend.writes) + await second.sup.control_round() + assert len(world.frontend.writes) == resets + assert world.shared_state().rollback_pending + await settle([second], world, rounds=2) + assert world.shared_state().rollback_pending # 40s: not yet + await settle([second], world, rounds=1) + state = world.shared_state() + assert not state.rollback_pending and state.deploy_seq == 3 + assert state.history[-1]["outcome"] == "rolled_back" + assert second.sup.target().desired_by_wt == VARIANTS["b"] + await settle([second], world, rounds=3) + assert len(world.frontend.writes) == resets # no further writes + + +async def test_applied_during_a_pending_rollback_wins() -> None: + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.stop() + world.spec.variant = "b" + world.frontend.restart(VARIANTS["b"]) + second = await leading(world, "p2") + assert world.shared_state().rollback_pending + seq = world.shared_state().deploy_seq + # An applied decision lands while the rollback is pending. + applied = {CREDIT: 1.0, SCALE: 1.0, TEMP: 0.2} + second.controller.restore_desired(applied) + enqueue(second, make_decision(K0 + 2, "applied", chosen=applied)) + await second.sup.outbox_round() + state = world.shared_state() + assert not state.rollback_pending and state.deploy_seq == seq + 1 + assert second.sup.target().mode == "apply" + # A rollback-complete write racing it is discarded, never retried. + stale_seq = seq + from dynamo.planner.retune.live.state import rollback_complete + + assert rollback_complete(state, expected_deploy_seq=stale_seq) is None + + +@pytest.mark.parametrize("racing", ["deploy", "rollback_start"]) +async def test_a_stale_rollback_complete_is_discarded(racing: str) -> None: + """Right before the rollback-complete write of deploy_seq S, a write of + the same term moves deploy_seq on (an applied, or a new rollback): the + completion of S is discarded, never retried, never applied to S+1.""" + + from dynamo.planner.retune.live.state import deploy, rollback_start + + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.stop() + world.spec.variant = "b" + world.frontend.restart(VARIANTS["b"]) + second = await leading(world, "p2") + sup = second.sup + seq = world.shared_state().deploy_seq + assert world.shared_state().rollback_pending + store = sup._store + original = store.write + raced = [] + + def racing_write(state, snapshot): + if not raced: + raced.append(1) + current = store.read() + rule = deploy if racing == "deploy" else rollback_start + original( + rule( + current.state, + fingerprint="fp-b", + mode="active", + flat=startup_flat("b"), + by_wt=VARIANTS["b"], + ), + current, + ) + return original(state, snapshot) + + store.write = racing_write + for _ in range(6): + if raced: + break + await settle([second], world, rounds=1) + assert raced + state = world.shared_state() + assert state.deploy_seq == seq + 1 + assert state.rollback_pending is (racing == "rollback_start") + assert all(item.get("outcome") != "rolled_back" for item in state.history) + # The new rollback completes on its own deploy_seq. + if racing == "rollback_start": + await settle([second], world, rounds=4) + state = world.shared_state() + assert not state.rollback_pending + assert state.history[-1]["decision_id"] == f"rolled_back-{seq + 1}" + + +async def test_advisory_switch_rolls_back_only_what_was_deployed() -> None: + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.control_round() + await first.sup.stop() + second = await leading(world, "p2", advisory=True) + state = world.shared_state() + assert state.mode == "advisory" and state.rollback_pending + # Advisory never restores a deployed point into the controller. + assert second.controller.current_point().router == startup_flat("a") + await settle([second], world, rounds=5) + state = world.shared_state() + assert not state.rollback_pending + assert frontend_configs(world) == [VARIANTS["a"]["aggregated"]] + planner_writes = [w for w in world.frontend.writes if w[0] == "p2"] + assert planner_writes and all(w[2] == "reset" for w in planner_writes) + # Now advisory only observes: a manual change is left alone. + world.frontend.manual_patch(TEMP, 0.9) + before = len(world.frontend.writes) + await settle([second], world, rounds=3) + assert len(world.frontend.writes) == before + + +async def test_advisory_rollback_leaves_unlatched_routers_alone() -> None: + """A router changed at revision > 0 before anyone latched it is reset to + learn its initial only in active mode.""" + + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.control_round() + target = first.sup.target() + assert target.mode == "apply" and target.learn_unknown + await first.sup.stop() + # The frontend restarts and an operator patches its router before any + # planner latched it. + world.frontend.restart(VARIANTS["a"]) + world.frontend.manual_patch(TEMP, 0.7) + (manual,) = world.frontend.routers.values() + second = await leading(world, "p2", advisory=True) + target = second.sup.target() + assert target.mode == "rollback" and not target.learn_unknown + await settle([second], world, rounds=5) + # Never written, and the rollback still completes without it. + assert not [w for w in world.frontend.writes if w[0] == "p2"] + assert manual.config[TEMP] == 0.7 and manual.revision == 1 + assert not world.shared_state().rollback_pending + await second.sup.stop() + + +def test_rollback_targets_learn_unknown_routers_only_when_active() -> None: + world = World() + for advisory, decommission, learn in ( + (False, False, True), + (True, False, False), + (True, True, False), + ): + process = make_process( + world, + settings=settings_for(advisory=advisory, decommission=decommission), + ) + sup = process.sup + sup._initialized = True + sup._snapshot = SimpleNamespace(state=claimed(rollback_pending=True)) + target = sup.target() + assert target.mode == "rollback" and target.include_not_ready + assert target.learn_unknown is learn + # A failed spec premise rolls back an active planner's own deployment. + sup = make_process(world).sup + sup._initialized = True + sup._snapshot = SimpleNamespace(state=claimed()) + sup._spec_ok = False + target = sup.target() + assert target.mode == "rollback" and target.learn_unknown + + +async def test_promote_applies_the_last_advisory_recommendation() -> None: + world = World() + first = await leading(world, advisory=True) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + enqueue(first, make_decision(K0 + 1, "advisory", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.stop() + world.clock.advance(100) + second = await leading(world, "p2", promote=True) + state = world.shared_state() + assert state.mode == "active" and state.by_wt == by_wt("a", chosen) + assert state.history[-1]["outcome"] == "promoted" + assert second.controller.current_point().router == chosen + assert frontend_configs(world) == [by_wt("a", chosen)["aggregated"]] + + +async def test_a_state_of_another_dgd_uid_is_recreated() -> None: + world = World() + store = StateStore(world.kube, DgdRef(DGD_NAME, "test-ns", "old-uid", "dyn")) + snapshot = store.read() + store.write( + RetuneState( + dgd_uid="old-uid", + leader_term="0/x/1", + fingerprint="fp-a", + mode="active", + last_window_end_ms=123, + ), + snapshot, + ) + world.kube.create_request( + {"metadata": {"name": "old-req", "labels": {LABEL_DGD_UID: "old-uid"}}} + ) + reconciler = SimpleNamespace( + deleted=[], + protect=lambda *a: None, + release=lambda *a: None, + run_once=lambda **kwargs: None, + ) + reconciler.enqueue_delete = lambda names: reconciler.deleted.extend(names) + process = await leading(world, reconciler=reconciler) + state = world.shared_state() + assert state.dgd_uid == "uid-1" and state.last_window_end_ms == end_ms(K0) + owner = world.kube.leases[STATE_NAME]["metadata"]["ownerReferences"][0] + assert owner["uid"] == "uid-1" + assert reconciler.deleted == ["old-req"] + assert process.sup.initialized + + +# ----------------------------------------------------------- revalidation + + +async def test_spec_flip_rebuilds_rolls_back_and_recovers() -> None: + world = World() + process = await leading(world) + sup, host = process.sup, process.host + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + process.controller.restore_desired(chosen) + enqueue(process, make_decision(K0 + 1, "applied", chosen=chosen)) + await sup.outbox_round() + await sup.control_round() + old_controller = process.controller + # P8 breaks (a second frontend without replica sync). + world.spec.spec_ok = False + await sup.control_round() + assert host.replacements[-1][1] == 1 and sup.generation == 1 + assert process.controller is not old_controller + assert old_controller._cancel.is_set() # closed + state = world.shared_state() + assert state.rollback_pending and state.by_wt == VARIANTS["a"] + assert sup.target().mode == "rollback" and not sup.ready() + assert frontend_configs(world) == [VARIANTS["a"]["aggregated"]] + assert sup.launch_gate(K0 + 2, window(K0 + 2)).reason == "not_ready" + # Fixed: rebuilt again, back to ready. + world.spec.spec_ok = True + await settle([process], world, rounds=5) + assert sup.generation == 2 and sup.ready() + assert not world.shared_state().rollback_pending + assert sup.launch_gate(K0 + 2, window(K0 + 2)).reason is None + + +async def test_p4_failure_disables_until_it_recovers() -> None: + world = World() + world.spec.p4_ok = False + process = await leading(world) + sup, host = process.sup, process.host + assert sup.retune_inputs() is None and process.controller is None + assert not sup.ready() + world.spec.p4_ok = True + await sup.control_round() + assert variant_of(host.replacements[-1][0]) == "a" + assert process.controller is not None and sup.generation == 1 + await sup.control_round() + assert sup.ready() + + +async def test_controller_construction_failure_keeps_not_ready() -> None: + world = World() + process = await leading(world) + process.host.fail_replace = True + world.spec.variant = "b" + world.frontend.restart(VARIANTS["b"]) + await process.sup.control_round() # no exception escapes + assert process.controller is None and not process.sup.ready() + assert process.sup.initialized # still leads, rolls back, persists + + +async def test_p14_parity_failure_is_a_spec_failure() -> None: + world = World() + process = await leading(world) + sup = process.sup + # A frontend restarts with startup values the planner did not parse. + world.frontend.restart(VARIANTS["b"]) + await sup.control_round() + assert sup._p14_ok is False + await sup.control_round() # revalidation sees spec_ok flip + assert not sup.spec_ok and sup.generation == 1 + assert sup.target().mode == "rollback" + world.frontend.restart(VARIANTS["a"]) + await settle([process], world, rounds=3) + assert sup.spec_ok and sup.ready() + + +async def test_namespace_change_swaps_the_actuator_only() -> None: + world = World() + process = await leading(world) + old = process.sup.actuator + world.spec.dynamo_namespace = "dyn-v2" + await process.sup.control_round() + assert process.sup.actuator is not old + assert process.sup.dgd.dynamo_namespace == "dyn-v2" + assert process.sup.generation == 0 and process.host.replacements == [] + + +async def test_health_premises_only_gate_launches() -> None: + world = World() + verdict = {"ok": False} + process = await leading(world, canary=lambda sup: verdict["ok"]) + sup = process.sup + rv = world.kube.leases[STATE_NAME]["metadata"]["resourceVersion"] + assert not sup.ready() + world.frontend.down = True + await sup.control_round() + world.frontend.down = False + verdict["ok"] = True + assert world.kube.leases[STATE_NAME]["metadata"]["resourceVersion"] == rv + await sup.control_round() + assert sup.ready() + + +async def test_an_actuator_round_overtaken_by_a_local_loss_is_discarded() -> None: + world = World() + process = await leading(world) + sup = process.sup + report_before = sup.last_report + rv = world.kube.leases[STATE_NAME]["metadata"]["resourceVersion"] + + class OvertakenActuator(RouterActuator): + async def reconcile(self, target, *, initials, expected_frontends): + sup._epoch += 1 # a local loss ran while this round awaited I/O + return ActuationReport( + converged=False, + new_initials={"stale-router": {CREDIT: 9.0}}, + seen_router_ids=frozenset({"stale-router"}), + parity_ok=False, + ) + + sup._actuator = OvertakenActuator( + world.kube, + FrontendPolicyClient(world, "p1", False), + sup.dgd, + VARIANTS["a"], + NullMetrics(), + ) + await sup.control_round() + assert sup.last_report is report_before and sup._p14_ok is not False + assert world.kube.leases[STATE_NAME]["metadata"]["resourceVersion"] == rv + + +async def test_initials_are_persisted_and_pruned() -> None: + world = World() + process = await leading(world) + first_ids = set(world.frontend.routers) + assert set(world.shared_state().initials) == first_ids + world.frontend.restart(VARIANTS["a"]) + await process.sup.control_round() + assert set(world.shared_state().initials) == set(world.frontend.routers) + assert not first_ids & set(world.shared_state().initials) + + +# ------------------------------------------------------------ reconciler + + +class FakeReconciler: + """Records each round's ``should_continue``; a round can hang (a + hard-mounted NFS call) until ``unblock`` is set.""" + + def __init__(self, *, hang: bool = False) -> None: + self.rounds: list[Any] = [] + self.answers_after: list[bool] = [] + self.entered = threading.Event() + self.unblock = threading.Event() + if not hang: + self.unblock.set() + self.events: list[tuple[str, Any]] = [] + + def protect(self, window_end_ms, names) -> None: + pass + + def release(self, window_end_ms) -> None: + pass + + def enqueue_delete(self, names) -> None: + self.events.append(("enqueue_delete", sorted(names))) + + def set_dgd(self, dgd: DgdRef) -> None: + self.events.append(("set_dgd", dgd.uid)) + + def run_once( + self, + *, + last_window_end_ms, + due_window_end_ms, + deadline_for, + frontend_pods, + should_continue=None, + ): + self.rounds.append(should_continue) + self.entered.set() + assert self.unblock.wait(10), "never unblocked" + self.answers_after.append(should_continue()) + + +def reconciler_on_a_thread(reconciler: FakeReconciler): + """Run only the reconciler's round on a worker thread (as live).""" + + async def blocking(call): + if getattr(call, "func", None) == reconciler.run_once: + return await asyncio.to_thread(call) + return call() + + return blocking + + +async def reconciled(sup: LiveRetuneSupervisor) -> None: + task = sup.reconcile_task + if task is not None: + await task + + +async def test_a_hung_reconciler_never_blocks_revalidation_or_actuation() -> None: + world = World() + reconciler = FakeReconciler(hang=True) + process = await leading( + world, + reconciler=reconciler, + blocking=reconciler_on_a_thread(reconciler), + settings=settings_for(reconcile_budget_s=0.05), + ) + sup = process.sup + try: + assert await asyncio.to_thread(reconciler.entered.wait, 5) + # The round runs in its own task; the control loop goes on and never + # starts a second round while the first one's call still runs. + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + process.controller.restore_desired(chosen) + enqueue(process, make_decision(K0 + 1, "applied", chosen=chosen)) + await sup.outbox_round() + await asyncio.wait_for(sup.control_round(), 5) + assert frontend_configs(world) == [by_wt("a", chosen)["aggregated"]] + assert len(reconciler.rounds) == 1 and sup.reconcile_running() + # Past its budget the round's task ends and the reconciler is told + # to stop before its next delete; its call still runs. + await asyncio.wait_for(reconciled(sup), 5) + assert sup.reconcile_running() and reconciler.rounds[0]() is False + # A spec break is still revalidated and rolled back meanwhile. + world.spec.spec_ok = False + await asyncio.wait_for(sup.control_round(), 5) + await asyncio.wait_for(sup.control_round(), 5) + assert not sup.spec_ok and sup.target().mode == "rollback" + assert frontend_configs(world) == [VARIANTS["a"]["aggregated"]] + assert len(reconciler.rounds) == 1 + finally: + reconciler.unblock.set() + work = sup._reconcile_work + assert work is not None + await asyncio.wait_for(work, 5) + assert reconciler.answers_after == [False] + # Once the call returned, the next control round starts a new round. + assert not sup.reconcile_running() + await sup.control_round() + await reconciled(sup) + await asyncio.wait_for(sup._reconcile_work, 5) + assert len(reconciler.rounds) == 2 and reconciler.answers_after[-1] is True + await sup.stop() + + +async def test_reconciler_runs_only_for_the_leader_of_its_epoch() -> None: + world = World() + reconciler = FakeReconciler() + process = await leading(world, reconciler=reconciler) + sup = process.sup + await reconciled(sup) + assert len(reconciler.rounds) == 1 and reconciler.rounds[0]() is True + # Another process took the term; the actuator round finds out while + # persisting new initials: no reconciler round starts after it. + store = StateStore(world.kube, sup.dgd) + snapshot = store.read() + store.write(dataclasses.replace(snapshot.state, leader_term="9/x/1"), snapshot) + world.frontend.restart(VARIANTS["a"]) + await sup.control_round() + await reconciled(sup) + assert not sup.initialized + assert len(reconciler.rounds) == 1 and reconciler.rounds[0]() is False + # Re-initialized (a new epoch): the old round's should_continue stays + # False, the new round's is True until leadership lapses. + await sup.lease_round() + await sup.control_round() + await reconciled(sup) + assert len(reconciler.rounds) == 2 + assert reconciler.rounds[0]() is False and reconciler.rounds[1]() is True + world.clock.advance(LEASE) + assert not sup.is_leader() and reconciler.rounds[1]() is False + await sup.control_round() + assert len(reconciler.rounds) == 2 + + +async def test_recreated_dgd_moves_the_reconciler_first() -> None: + world = World() + reconciler = FakeReconciler() + process = await leading(world, reconciler=reconciler) + world.kube.create_request( + {"metadata": {"name": "old-req", "labels": {LABEL_DGD_UID: "uid-1"}}} + ) + world.spec.uid = "uid-2" # deleted and recreated under the same name + await process.sup.control_round() + assert process.sup.dgd.uid == "uid-2" and process.sup.initialized + assert world.shared_state().dgd_uid == "uid-2" + # The reconciler follows the new uid before the old uid's Requests are + # queued for deletion, so the new DGD's Requests stay its own. + assert reconciler.events == [("set_dgd", "uid-2"), ("enqueue_delete", ["old-req"])] + # A namespace-only change keeps the reconciler's identity. + world.spec.dynamo_namespace = "dyn-2" + await process.sup.control_round() + assert reconciler.events[-1] == ("enqueue_delete", ["old-req"]) + + +# ----------------------------------------------------------- sweeps, stop + + +async def test_sweep_fn_binds_the_context_and_never_raises() -> None: + world = World() + seen = {} + + def runner(request, *, cancel, abort, context): + seen.update(request=request, cancel=cancel, abort=abort, context=context) + return "result" + + process = await started( + world, sweep_runner=runner, context_factory=lambda sup: ("ctx", sup.dgd) + ) + cancel, abort = threading.Event(), threading.Event() + request = SimpleNamespace(window_index=5) + assert process.sup.sweep_fn(request, cancel=cancel, abort=abort) == "result" + assert seen["context"] == ("ctx", process.sup.dgd) and seen["abort"] is abort + + def broken(*args, **kwargs): + raise RuntimeError("boom") + + process.sup._sweep_runner = broken + result = process.sup.sweep_fn(request, cancel=cancel, abort=abort) + assert result.status == "failed" and result.failure_reason == "sweep_error" + assert "boom" in result.error + no_context = await started(world, "p2", sweep_runner=runner, context_factory=None) + result = no_context.sup.sweep_fn(request, cancel=cancel, abort=abort) + assert result.status == "failed" + + +async def test_stop_aborts_the_sweep_and_waits_for_the_sweep_thread() -> None: + world = World() + process = await leading(world) + sup, controller = process.sup, process.controller + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + pending = controller._pending + world.sweeps.quiesce([controller]) # running on the sweep thread + abort = pending.abort + + await sup.stop() + + # drop_pending set the sweep's abort; stop waited for the sweep thread. + assert abort.is_set() and pending.future.done() + assert pending.future.result().failure_reason == "lost_leadership" + assert world.holder() is None and not sup.is_leader() + await sup.stop() # idempotent + + +async def test_stop_gives_up_on_a_sweep_that_never_returns(caplog) -> None: + world = World() + release = threading.Event() + + def stuck(request, *, cancel, abort, context): + release.wait(10) # ignores its abort + return SweepResult(status="failed") + + process = await leading( + world, sweep_runner=stuck, settings=settings_for(stop_timeout_s=0.05) + ) + sup, controller = process.sup, process.controller + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + try: + with caplog.at_level(logging.WARNING): + await sup.stop() + assert "a sweep is still running at shutdown" in caplog.text + assert world.holder() is None + finally: + release.set() + + +async def test_background_loops_run_and_stop() -> None: + world = World() + process = make_process(world) + sup = process.sup + sup._sleep = lambda s: asyncio.sleep(0.001) + await sup.start(process.host) + for _ in range(500): + await asyncio.sleep(0.001) + if sup.initialized and sup.ready(): + break + assert sup.initialized and sup.ready() + enqueue(process, make_decision(K0 + 1, "kept")) + for _ in range(500): + await asyncio.sleep(0.001) + if sup.outbox_size == 0: + break + assert world.shared_state().history[-1]["outcome"] == "kept" + await sup.stop() + assert world.holder() is None + assert process.metrics.state["leader"] is False + + +# ------------------------------------------------------------ lease timing + +MIN_LEASE = live_retune_min_lease_seconds(LIVE_RETUNE_K8S_TIMEOUT_SECONDS) + + +def test_settings_require_a_lease_that_fits_two_renewals() -> None: + for lease in (15.0, MIN_LEASE - 0.001): + with pytest.raises( + ValueError, + match=r"k8s_timeout_s 5s: it must be at least the lease margin " + r"\(15s\) \+ 2 x the I/O budget of one renewal \(11s\) = 37s", + ): + settings_for(lease_duration_s=lease) + for timeout in (0.0, -1.0, math.inf, math.nan): + with pytest.raises(ValueError, match="k8s_timeout_s must be positive"): + settings_for(k8s_timeout_s=timeout) + settings = settings_for(lease_duration_s=MIN_LEASE) + assert settings.lease_margin_s == 15.0 and settings.io_timeout_s == 11.0 + # The legacy clean-up supervisor's settings (the defaults). + SupervisorSettings(interval_s=3600.0, grace_s=0.0, timeout_s=3000.0) + dgdr = live_config().retune.dgdr.model_dump() + dgdr["lease_duration_seconds"] = MIN_LEASE + config = live_config(retune={"dgdr": dgdr}) + assert SupervisorSettings.from_config(config).lease_duration_s == MIN_LEASE + + +class SlowLeases: + """``world.kube``, but each read or write of a Lease named in + ``latency_s`` first advances the test's clock by its latency.""" + + def __init__(self, world: World) -> None: + self._kube = world.kube + self._clock = world.clock + self.latency_s: dict[str, float] = {} + + def __getattr__(self, name: str) -> Any: + return getattr(self._kube, name) + + def _wait(self, lease: str) -> None: + self._clock.advance(self.latency_s.get(lease, 0.0)) + + def get_lease(self, name: str, **kwargs: Any) -> Optional[dict]: + self._wait(name) + return self._kube.get_lease(name, **kwargs) + + def create_lease(self, body: dict, **kwargs: Any) -> dict: + self._wait(body["metadata"]["name"]) + return self._kube.create_lease(body, **kwargs) + + def replace_lease(self, name: str, body: dict, **kwargs: Any) -> dict: + self._wait(name) + return self._kube.replace_lease(name, body, **kwargs) + + +class LeaseLoop: + """The supervisor's own lease loop, one round at a time: each of its + waits advances the test's clock by the time the loop asked for (a + re-init request does not cut it short). + + Every other sleep (an initialization's renewals) ends once Kubernetes + calls have advanced the clock past it: each call runs inline, then lets + the sleepers it made due run before its caller goes on. + """ + + def __init__(self, sup: LiveRetuneSupervisor, clock: Clock) -> None: + self._sup = sup + self._clock = clock + self._asleep = asyncio.Event() + self._wake = asyncio.Event() + self._task: Optional[asyncio.Task[None]] = None + self._sleepers: list[tuple[float, asyncio.Future[None]]] = [] + self.waits: list[float] = [] + + async def __aenter__(self) -> LeaseLoop: + self._sup._wait = self._wait # type: ignore[method-assign] + self._sup._sleep = self._sleep + self._sup._blocking = self._run + self._task = asyncio.create_task(self._sup._lease_loop()) + await self._round_done() + return self + + async def __aexit__(self, *exc_info: Any) -> None: + assert self._task is not None + self._task.cancel() + await asyncio.wait([self._task]) + + async def _wait(self, _event: Optional[asyncio.Event], seconds: float) -> None: + self.waits.append(seconds) + self._asleep.set() + await self._wake.wait() + self._wake.clear() + self._clock.advance(seconds) + + async def _sleep(self, seconds: float) -> None: + if seconds <= 0: + return + future = asyncio.get_running_loop().create_future() + self._sleepers.append((self._clock.monotonic() + seconds, future)) + await future + + async def _run(self, call): + try: + return call() + finally: + now = self._clock.monotonic() + for deadline, future in self._sleepers: + if deadline <= now and not future.done(): + future.set_result(None) + self._sleepers = [ + (deadline, future) + for deadline, future in self._sleepers + if not future.done() + ] + await asyncio.sleep(0) + + async def _round_done(self) -> None: + await asyncio.wait_for(self._asleep.wait(), timeout=10) + self._asleep.clear() + + async def next_round(self) -> None: + """Let the loop wait as asked, then run its next round.""" + + self._wake.set() + await self._round_done() + + +async def test_renewals_after_a_slow_initialization_keep_the_term() -> None: + world = World() + kube = SlowLeases(world) + process = await started(world, kube=kube) + sup, controller = process.sup, process.controller + # The state read, claim and plan write after the confirming renewal take + # 27 s, each within its 11 s I/O budget: more than the local lifetime + # (45 s) minus the renew interval (20 s). + kube.latency_s[STATE_NAME] = 9.0 + async with LeaseLoop(sup, world.clock) as loop: + assert sup.initialized + kube.latency_s.clear() + term, term_seq = sup.term, sup.term_seq + await sup.control_round() + assert sup.ready() + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + pending = controller._pending + assert pending is not None + for _ in range(6): + await loop.next_round() + assert sup.is_leader() and sup.initialized + assert (sup.term, sup.term_seq) == (term, term_seq) + assert controller._pending is pending and not pending.abort.is_set() + assert controller.dropped == [] + + +async def test_minimum_lease_renewals_keep_the_term_and_the_sweep() -> None: + world = World() + kube = SlowLeases(world) + settings = settings_for(lease_duration_s=MIN_LEASE) + process = await leading(world, kube=kube, settings=settings) + sup, controller = process.sup, process.controller + term, term_seq = sup.term, sup.term_seq + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + pending = controller._pending + assert pending is not None + # Every leadership Lease call takes its whole request timeout. + kube.latency_s[LEADER_NAME] = settings.k8s_timeout_s + async with LeaseLoop(sup, world.clock) as loop: + for _ in range(8): + await loop.next_round() + assert sup.is_leader() and sup.initialized + assert (sup.term, sup.term_seq) == (term, term_seq) + assert world.kube.leases[LEADER_NAME]["spec"]["leaseTransitions"] == 0 + assert controller._pending is pending and not pending.abort.is_set() + assert controller.dropped == [] + + +@pytest.mark.parametrize("latency_s", [4.0, 5.0]) +async def test_minimum_lease_initialization_with_slow_leases_keeps_the_term( + latency_s: float, +) -> None: + world = World() + kube = SlowLeases(world) + settings = settings_for(lease_duration_s=MIN_LEASE) + process = await started(world, kube=kube, settings=settings) + sup, controller = process.sup, process.controller + # Without renewals in between, local leadership would end soon after the + # state read, claim and plan write that follow the confirming renewal + # (4 s per Lease call: before the next renewal is sent), or before they + # end (5 s, the whole request timeout). + kube.latency_s[LEADER_NAME] = latency_s + kube.latency_s[STATE_NAME] = latency_s + async with LeaseLoop(sup, world.clock) as loop: + assert sup.is_leader() and sup.initialized + term, term_seq = sup.term, sup.term_seq + del kube.latency_s[STATE_NAME] + await sup.control_round() + assert sup.ready() + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + pending = controller._pending + assert pending is not None + for _ in range(8): + await loop.next_round() + assert sup.is_leader() and sup.initialized + assert (sup.term, sup.term_seq) == (term, term_seq) + assert controller._pending is pending and not pending.abort.is_set() + assert controller.dropped == [] + + +async def test_a_lapse_during_initialization_aborts_it_and_reinitializes_once() -> None: + world = World() + kube = SlowLeases(world) + settings = settings_for(lease_duration_s=MIN_LEASE) + process = await started(world, kube=kube, settings=settings) + sup = process.sup + kube.latency_s[LEADER_NAME] = 1.0 + kube.latency_s[STATE_NAME] = 9.0 + # From the new term on (set before its claim), reading the leadership + # Lease fails until local leadership has ended: the next renewal that + # succeeds re-acquires while the initialization of that term still runs. + world.kube.fail( + "get_lease", + KubeError("unavailable"), + times=100, + match=lambda target: target == LEADER_NAME + and sup.term is not None + and sup.is_leader(), + ) + async with LeaseLoop(sup, world.clock) as loop: + assert sup.is_leader() and not sup.initialized + assert sup.term_seq == 1 + world.kube.clear_faults() + kube.latency_s.clear() + for _ in range(3): + await loop.next_round() + assert sup.is_leader() and sup.initialized + assert sup.term_seq == 2 + assert world.shared_state().leader_term == sup.term + + +async def test_a_cancelled_initialization_leaves_no_renewal_behind() -> None: + world = World() + process = await started(world) + sup = process.sup + steps = 0 + in_flight = asyncio.Event() + gate, returned = threading.Event(), threading.Event() + late: list[str] = [] + + def gated_step(call): + gate.wait(10) + try: + late.append(call()) + finally: + returned.set() + + async def gated(call): + # The third leadership step, after the acquiring and the confirming + # ones, is the initialization's first renewal: it runs on a thread, + # as the production runner's steps do, and blocks there until the + # gate opens. The initialization goes on only once it is in flight. + nonlocal steps + if getattr(call, "func", None) == sup._leadership.step: + steps += 1 + if steps == 3: + in_flight.set() + return await _to_thread(lambda: gated_step(call)) + elif steps >= 2: + await in_flight.wait() + return call() + + sup._blocking = gated + sup._sleep = lambda _seconds: asyncio.sleep(0) + task = asyncio.create_task(sup.lease_round()) + for _ in range(1000): + if sup.initialized: + break + await asyncio.sleep(0) + # Initialized; the round waits for the renewal in flight. + assert sup.initialized and in_flight.is_set() and not task.done() + task.cancel() + await asyncio.wait([task]) + pending = [ + other + for other in asyncio.all_tasks() + if other.get_name() == "planner-retune-init-lease" + ] + assert pending == [] + await sup.stop() + assert world.holder() is None + # The cancelled renewal's thread goes on and reads the Lease after stop + # released it: it neither takes it over nor leads. + gate.set() + assert await asyncio.to_thread(returned.wait, 10) + assert late == ["standby"] + assert world.holder() is None and not sup._leadership.is_leader() + + +async def test_a_lapse_reported_after_initialization_ends_the_term() -> None: + world = World() + process = await started(world) + sup = process.sup + steps = 0 + in_flight, gate = asyncio.Event(), asyncio.Event() + + async def gated(call): + # The third leadership step, the initialization's first renewal, + # stays in flight until the initialization has finished and returns + # after local leadership ended: it reports a lapse (acquired). + nonlocal steps + if getattr(call, "func", None) == sup._leadership.step: + steps += 1 + if steps == 3: + in_flight.set() + await gate.wait() + world.clock.advance(LEASE) + elif steps >= 2: + await in_flight.wait() + return call() + + sup._blocking = gated + sup._sleep = lambda _seconds: asyncio.sleep(0) + task = asyncio.create_task(sup.lease_round()) + for _ in range(1000): + if sup.initialized: + break + await asyncio.sleep(0) + assert sup.initialized and in_flight.is_set() + term_seq = sup.term_seq + # The round waits for the renewal in flight. + for _ in range(100): + await asyncio.sleep(0) + assert not task.done() + gate.set() + assert await task == "acquired" + lease = world.kube.leases[LEADER_NAME] + assert sup._leadership.resource_version == lease["metadata"]["resourceVersion"] + # The lapse ends the term the initialization started. + assert sup._leadership.acquisitions == 2 and not sup.initialized + sup._blocking = inline + await sup.lease_round() + assert sup.initialized and sup.term_seq == term_seq + 1 + assert world.shared_state().leader_term == sup.term + + +async def test_leadership_steps_of_a_rebuild_and_the_lease_loop_never_overlap() -> None: + world = World() + process = await leading(world) + sup = process.sup + active = overlaps = cancelled = 0 + + async def counting(call): + # Every leadership step yields to the event loop while in flight. + nonlocal active, overlaps, cancelled + if getattr(call, "func", None) != sup._leadership.step: + return call() + active += 1 + if active > 1: + overlaps += 1 + try: + for _ in range(5): + await asyncio.sleep(0) + return call() + except asyncio.CancelledError: + cancelled += 1 + raise + finally: + active -= 1 + + sup._blocking = counting + # P8 breaks: the rebuild re-initializes next to a lease round. + world.spec.spec_ok = False + world.clock.advance(20) + await asyncio.gather(sup.control_round(), sup.lease_round()) + assert sup.generation == 1 and sup.is_leader() and sup.initialized + assert overlaps == 0 and cancelled == 0 + + +async def test_failed_renewals_are_retried_a_request_timeout_apart() -> None: + world = World() + process = await leading(world) + sup = process.sup + world.kube.fail( + "get_lease", + KubeError("unavailable"), + times=100, + match=lambda target: target == LEADER_NAME, + ) + async with LeaseLoop(sup, world.clock) as loop: + for _ in range(5): + await loop.next_round() + # The renew interval, then one I/O budget before leader_until (45 s after + # the last renewal), then one request timeout apart until it passed. + assert loop.waits == [20.0, 14.0, 5.0, 5.0, 5.0, 20.0] + assert sup.term is None and not sup.initialized + + +async def test_a_retry_after_a_slow_failure_keeps_the_minimum_lease() -> None: + world = World() + kube = SlowLeases(world) + settings = settings_for(lease_duration_s=MIN_LEASE) + process = await leading(world, kube=kube, settings=settings) + sup, controller = process.sup, process.controller + term, term_seq = sup.term, sup.term_seq + world.clock.wall = window(K0 + 1)[2] + await controller.on_tick(world.clock.time(), SPEC_COUNTS) + pending = controller._pending + assert pending is not None + kube.latency_s[LEADER_NAME] = 2.0 + async with LeaseLoop(sup, world.clock) as loop: + # The renewal one I/O budget before leader_until fails after its + # whole request timeout; the retry goes at once and still renews. + kube.latency_s[LEADER_NAME] = settings.k8s_timeout_s + world.kube.fail( + "get_lease", + KubeError("unavailable"), + match=lambda target: target == LEADER_NAME, + ) + await loop.next_round() + kube.latency_s[LEADER_NAME] = 2.0 + await loop.next_round() + assert sup.is_leader() and sup.initialized + assert (sup.term, sup.term_seq) == (term, term_seq) + # leader_until is 22 s after each renewal's send; a renewal takes 4 s. + assert loop.waits == [7.0, 0.0, 7.0] + assert controller._pending is pending and not pending.abort.is_set() + + +async def test_hooks_are_the_adapter_hooks() -> None: + world = World() + process = await started(world) + hooks = process.sup.hooks() + assert isinstance(hooks, RetuneHooks) + assert hooks.launch_gate == process.sup.launch_gate + assert hooks.harvest_allowed == process.sup.harvest_allowed + assert hooks.sweep_fn == process.sup.sweep_fn and hooks.phase_s == 0.0 + + +def decommission_supervisor(world: World) -> LiveRetuneSupervisor: + """The clean-up supervisor of a planner whose retune was removed.""" + + frontend = SimpleNamespace( + parse_frontend=lambda dgd, planner_mode: world.spec.startup() + ) + return bootstrap.build_decommission_supervisor( + live_config(), + kube=world.kube, + dgd=world.kube.get_dgd(DGD_NAME), + dynamo_namespace="dyn", + metrics=NullMetrics(), + actuator_factory=real_actuators(world, "gc", advisory=True), + frontend=frontend, + clock=world.clock.time, + monotonic=world.clock.monotonic, + sleep=frozen_sleep, + blocking=inline, + ) + + +async def test_decommission_rolls_back_then_deletes_everything() -> None: + world = World() + first = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first.controller.restore_desired(chosen) + enqueue(first, make_decision(K0 + 1, "applied", chosen=chosen)) + await first.sup.outbox_round() + await first.sup.control_round() + await first.sup.stop() + world.kube.create_request( + {"metadata": {"name": "rt-1", "labels": {LABEL_DGD_UID: "uid-1"}}} + ) + sup = decommission_supervisor(world) + await sup.start(SimpleNamespace(), background=False) + for _ in range(6): + if sup.decommissioned: + break + await sup.lease_round() + await sup.control_round() + world.clock.advance(30) + assert sup.decommissioned + assert frontend_configs(world) == [VARIANTS["a"]["aggregated"]] + assert world.kube.requests == {} and world.kube.leases == {} + assert sup.launch_gate(K0 + 1, window(K0 + 1)).reason == "not_leader" + await sup.stop() + + +async def decommissioning(world: World) -> LiveRetuneSupervisor: + """An initialized clean-up supervisor with nothing to roll back.""" + + first = await leading(world, advisory=True) # nothing deployed + await first.sup.stop() + world.kube.create_request( + {"metadata": {"name": "rt-1", "labels": {LABEL_DGD_UID: "uid-1"}}} + ) + sup = decommission_supervisor(world) + await sup.start(SimpleNamespace(), background=False) + await sup.lease_round() + assert sup.initialized and not world.shared_state().rollback_pending + return sup + + +async def test_decommission_never_deletes_after_losing_its_term() -> None: + world = World() + sup = await decommissioning(world) + # Another replica claims the state; this round's initials write (a + # restarted frontend) finds out before the clean-up would delete. + store = StateStore(world.kube, sup.dgd) + snapshot = store.read() + store.write(dataclasses.replace(snapshot.state, leader_term="9/x/1"), snapshot) + world.frontend.restart(VARIANTS["a"]) + await sup.control_round() + assert not sup.decommissioned and not sup.initialized + assert set(world.kube.requests) == {"rt-1"} + assert set(world.kube.leases) == {STATE_NAME, LEADER_NAME} + await sup.stop() + + +async def test_decommission_deletes_only_the_leases_it_last_saw() -> None: + # The local lease clock still says leader and no write of the round + # notices the other term: only the delete preconditions stop it. + world = World() + sup = await decommissioning(world) + store = StateStore(world.kube, sup.dgd) + snapshot = store.read() + store.write(dataclasses.replace(snapshot.state, leader_term="9/x/1"), snapshot) + await sup.control_round() + assert not sup.decommissioned and not sup.initialized + assert set(world.kube.leases) == {STATE_NAME, LEADER_NAME} + assert world.shared_state().leader_term == "9/x/1" + await sup.stop() + + world = World() + sup = await decommissioning(world) + steal_lease(world, "intruder/0") + await sup.control_round() + assert sup.decommissioned and set(world.kube.leases) == {LEADER_NAME} + assert world.holder() == "intruder/0" + await sup.stop() + + +# -------------------------------------------------------------- bootstrap + + +def dgd_object(mode: str = "agg", **worker) -> dict: + components = [{"name": "Frontend", "type": "frontend", "replicas": 1}] + if mode == "disagg": + components += [ + {"name": "prefill", "type": "prefill", "replicas": 2}, + {"name": "decode", "type": "decode", **worker}, + ] + else: + components.append({"name": "worker", "type": "worker", **worker}) + return { + "metadata": { + "name": DGD_NAME, + "namespace": "test-ns", + "uid": "uid-1", + "generation": 3, + }, + "spec": {"components": components}, + } + + +def test_worker_spec_replicas() -> None: + assert bootstrap.worker_spec_replicas(dgd_object(replicas=3), "agg") == { + "aggregated": 3 + } + # Omitted replicas are the workload default of 1. + assert bootstrap.worker_spec_replicas(dgd_object(), "agg") == {"aggregated": 1} + assert bootstrap.worker_spec_replicas(dgd_object(replicas=0), "agg") is None + assert bootstrap.worker_spec_replicas( + dgd_object("disagg", replicas=4), "disagg" + ) == {"prefill": 2, "decode": 4} + assert bootstrap.worker_spec_replicas({"spec": {"components": []}}, "agg") is None + + +def test_build_retune_inputs_agg_and_disagg() -> None: + config = live_config() + startup = Spec().startup() + inputs = bootstrap.build_retune_inputs( + config, startup, dgd_name="graph", model_name="m" + ) + assert inputs.trace_paths == ("/trace/segments/graph",) + assert (inputs.trace_origin_ms, inputs.arrival_speedup_ratio) == (0.0, 1.0) + assert inputs.engine.router_mode == "kv_router" and inputs.engine.model_name == "m" + assert not inputs.engine.is_disagg and inputs.engine.extra_engine_args is None + disagg = bootstrap.build_retune_inputs( + config.model_copy(update={"mode": "disagg"}), + startup, + dgd_name="graph", + model_name=None, + ) + assert disagg.engine.is_disagg and disagg.engine.prefill_engine_args == {} + + +def test_kv_router_resolver_matches_the_router_config() -> None: + llm = pytest.importorskip("dynamo.llm") + assert VARIANTS == {variant: by_wt(variant, {}) for variant in ROUTER_CONFIG} + resolver = bootstrap.KvRouterResolver({}, ("prefill", "decode")) + flat = resolver.startup_flat(KEYS) + default = llm.KvRouterConfig.from_json("{}") + serialized = json.loads(default.to_json()) + resolved = dict(default.policy_parameters("aggregated")) + assert flat == { + key: float(resolved[key] if serialized.get(key) is None else serialized[key]) + for key in KEYS + } + resolved = resolver.by_wt({CREDIT: 0.5}) + config = llm.KvRouterConfig.from_json(json.dumps({CREDIT: 0.5})) + assert resolved == { + wt: {k: float(v) for k, v in dict(config.policy_parameters(wt)).items()} + for wt in ("prefill", "decode") + } + + +class FakeFrontendModule: + def __init__(self, startup: Optional[FrontendStartup], p4_ok=True, spec_ok=True): + self.startup = startup + self.p4_ok = p4_ok + self.spec_ok = spec_ok + + def parse_frontend(self, dgd, *, planner_mode): + return self.startup if self.startup is not None else "P4: shell-wrapped" + + def evaluate_premises(self, dgd, *, planner_config): + parsed = self.startup is not None + return PremiseReport( + p4_ok=self.p4_ok and parsed, + spec_ok=self.spec_ok and self.p4_ok and parsed, + reasons=() if parsed else ("P4: shell-wrapped",), + startup=self.startup, + ) + + def served_model_name(self, dgd): + return "served" + + +def test_evaluate_spec() -> None: + config = live_config() + frontend = FakeFrontendModule(Spec().startup()) + evaluation = bootstrap.evaluate_spec( + dgd_object(replicas=2), + config=config, + k8s_namespace="ns", + dynamo_namespace="dyn", + frontend=frontend, + ) + assert evaluation.dgd == DgdRef(DGD_NAME, "test-ns", "uid-1", "dyn") + assert evaluation.spec_replicas == {"aggregated": 2} and evaluation.generation == 3 + assert evaluation.inputs.engine.model_name == "served" + assert evaluation.report.p4_ok and evaluation.startup is not None + + unparsed = FakeFrontendModule(None) + evaluation = bootstrap.evaluate_spec( + dgd_object(), + config=config, + k8s_namespace="ns", + dynamo_namespace="dyn", + frontend=unparsed, + ) + assert evaluation.inputs is None and evaluation.startup is None + assert evaluation.report.reasons == ("P4: shell-wrapped",) + with pytest.raises(ValueError): + bootstrap.evaluate_spec( + {"metadata": {"name": "x"}}, + config=config, + k8s_namespace="ns", + dynamo_namespace="dyn", + frontend=frontend, + ) + + +async def test_build_supervisor_wires_the_collaborators() -> None: + world = World() + world.kube.put_dgd(dgd_object(replicas=2)) + frontend = FakeFrontendModule(Spec().startup()) + reconciler = SimpleNamespace( + protect=lambda *a: None, + release=lambda *a: None, + enqueue_delete=lambda *a: None, + run_once=lambda **kwargs: None, + ) + sup = await bootstrap.build_supervisor( + live_config(), + kube=world.kube, + dgd_name=DGD_NAME, + namespace_fn=lambda: "dyn", + metrics=NullMetrics(), + identity="pod/1", + frontend=frontend, + actuator_factory=real_actuators(world, "x"), + reconciler=reconciler, + canary=lambda s: True, + ) + assert sup.identity == "pod/1" and sup.registry is reconciler + assert sup.retune_inputs() is not None + assert sup.evaluation.spec_replicas == {"aggregated": 2} + assert sup._settings.router_keys == KEYS + assert sup._settings.interval_s == 21600 and sup._settings.lease_duration_s == 60 + assert len(sup.fingerprint) == 16 + + +async def test_build_live_retune_requires_the_dgdr_executor() -> None: + host = SimpleNamespace( + config=live_config( + retune={"executor": "local", "dgdr": None, "timeout_seconds": None} + ) + ) + with pytest.raises(ValueError, match="replay-only"): + await bootstrap.build_live_retune(host) + + +async def test_build_live_retune_uses_the_connector(monkeypatch) -> None: + world = World() + world.kube.put_dgd(dgd_object(replicas=2)) + from dynamo.planner.retune.live import kube as kube_module + + monkeypatch.setattr( + kube_module.KubernetesRetuneKube, + "from_kube_api", + classmethod(lambda cls, api: world.kube), + ) + captured = {} + + async def fake_build(config, **kwargs): + captured.update(kwargs, config=config) + return "supervisor" + + monkeypatch.setattr(bootstrap, "build_supervisor", fake_build) + host = SimpleNamespace( + config=live_config(), + environment=SimpleNamespace( + controller=SimpleNamespace(kube_api=object(), parent_dgd_name=DGD_NAME), + runtime_namespace=lambda: "dyn-hash", + ), + ) + assert await bootstrap.build_live_retune(host) == "supervisor" + assert captured["kube"] is world.kube and captured["dgd_name"] == DGD_NAME + assert captured["namespace_fn"]() == "dyn-hash" + + +async def test_build_legacy_cleanup(monkeypatch) -> None: + world = World() + world.kube.put_dgd(dgd_object(replicas=2)) + monkeypatch.setattr( + bootstrap.frontend_config, + "parse_frontend", + lambda dgd, planner_mode: "P4: no frontend", + ) + monkeypatch.setattr(bootstrap, "RetuneMetrics", NullMetrics) + host = SimpleNamespace( + config=live_config().model_copy(update={"retune": None}), + environment=SimpleNamespace( + controller=SimpleNamespace(parent_dgd_name=DGD_NAME), + runtime_namespace=lambda: "", + ), + ) + sup = await bootstrap.build_legacy_cleanup(host, world.kube) + assert sup._settings.decommission and sup.retune_inputs() is None + # The planner's configured namespace when the runtime has none. + assert sup.dgd == DgdRef(DGD_NAME, "test-ns", "uid-1", host.config.namespace) + world.kube.fail("get_dgd", KubeError("forbidden")) + with pytest.raises(KubeError): + await bootstrap.build_legacy_cleanup(host, world.kube) + + +class PlannerEnvironment: + """The ``PlannerEnvironment`` surface ``NativePlannerBase`` startup uses.""" + + def __init__(self, controller: Any) -> None: + from dynamo.planner.environment.state import DeploymentState + + self.controller = controller + self.state = DeploymentState() + self.state.decode.num_gpus = 1 + self.calls: list[str] = [] + + def deployment_state(self): + return self.state + + async def initialize(self) -> None: + self.calls.append("environment.initialize") + + async def refresh(self) -> None: + pass + + async def shutdown(self) -> None: + self.calls.append("environment.shutdown") + + +class FakeCleanup: + def __init__(self, calls: list[str]): + self.calls = calls + self.host = None + + async def start(self, host) -> None: + self.calls.append("cleanup.start") + self.host = host + + async def stop(self) -> None: + self.calls.append("cleanup.stop") + + +def forbidden() -> KubeError: + error = KubeError("read leases/graph-planner-retune: HTTP 403: forbidden") + error.status = 403 # type: ignore[attr-defined] + return error + + +def planner_without_retune(monkeypatch, controller: Any, *, leases=(), build=None): + """A ``NativePlannerBase`` without retune config; the engine bootstrap + is a no-op. ``leases`` is what the Lease check finds (or raises), a list + of those for consecutive checks, and ``build(host, kube)`` builds the + clean-up.""" + + from unittest.mock import patch + + from dynamo.planner.core import base as base_module + + config = live_config( + metric_reporting_prometheus_port=0, live_dashboard_port=0, control_api_port=0 + ).model_copy(update={"retune": None}) + environment = PlannerEnvironment(controller) + kube = object() + + outcomes = list(leases) if isinstance(leases, list) else None + + def find_retune_leases(found_kube, dgd_name, *, timeout_s): + assert found_kube is kube and dgd_name == DGD_NAME + environment.calls.append("leases.check") + found = leases if outcomes is None else outcomes.pop(0) + if isinstance(found, Exception): + raise found + return tuple(found) + + async def build_legacy_cleanup(host, found_kube, *, timeout_s): + assert found_kube is kube + environment.calls.append("cleanup.build") + return await build(host) + + monkeypatch.setattr(base_module, "find_retune_leases", find_retune_leases) + monkeypatch.setattr( + base_module.KubernetesRetuneKube, + "from_kube_api", + classmethod(lambda cls, api: kube), + ) + monkeypatch.setattr(bootstrap, "build_legacy_cleanup", build_legacy_cleanup) + with patch("dynamo.planner.core.base.PlannerPrometheusMetrics"): + planner = base_module.NativePlannerBase(None, config, environment) + + async def bootstrap_engine(): + environment.calls.append("engine") + + async def nothing(): + pass + + planner._bootstrap_regression = bootstrap_engine # type: ignore[method-assign] + planner._bootstrap_engine_plugins_if_needed = nothing # type: ignore[method-assign] + return planner, environment + + +def kubernetes_connector() -> Any: + from unittest.mock import MagicMock + + from dynamo.planner.connectors.kubernetes import KubernetesConnector + + connector = MagicMock(spec=KubernetesConnector) + connector.kube_api = object() + connector.parent_dgd_name = DGD_NAME + return connector + + +async def test_a_planner_without_retune_cleans_up_in_the_background( + monkeypatch, +) -> None: + cleanups: list[FakeCleanup] = [] + + async def build(host): + cleanups.append(FakeCleanup(environment.calls)) + return cleanups[-1] + + planner, environment = planner_without_retune( + monkeypatch, kubernetes_connector(), leases=(STATE_NAME,), build=build + ) + await planner._async_init() + try: + # Startup never waits for the check. + assert environment.calls == ["environment.initialize", "engine"] + assert planner.may_hold_retune_lease + assert not planner.retune_lease_found.is_set() + await planner._legacy_retune_task + # The clean-up holds the leadership Lease: SIGTERM must stop it. + assert planner.retune_lease_found.is_set() + assert environment.calls[2:] == [ + "leases.check", + "cleanup.build", + "cleanup.start", + ] + assert cleanups[0].host is planner and planner._live_retune is None + finally: + await planner._shutdown_runtime() + assert environment.calls[-2:] == ["cleanup.stop", "environment.shutdown"] + assert planner._legacy_retune_cleanup is None + + +@pytest.mark.parametrize("outcome", ["nothing_left", "forbidden", "not_found", "bug"]) +async def test_the_legacy_check_never_fails_the_planner( + monkeypatch, caplog, outcome +) -> None: + from dynamo.planner.core import base as base_module + from dynamo.planner.retune.live.contracts import KubeNotFound + + leases: Any = (LEADER_NAME,) + if outcome == "nothing_left": + leases = () + elif outcome == "forbidden": + leases = forbidden() + + async def build(host): + if outcome == "not_found": + raise KubeNotFound("read dynamographdeployments/graph: HTTP 404") + raise RuntimeError("a bug") + + if outcome in ("nothing_left", "forbidden"): + # The live retune stack is never loaded. + monkeypatch.setattr(base_module, "_live_retune_bootstrap", None) + planner, environment = planner_without_retune( + monkeypatch, kubernetes_connector(), leases=leases, build=build + ) + with caplog.at_level(logging.DEBUG, logger="dynamo.planner.core.base"): + await planner._async_init() + try: + await asyncio.wait([planner._legacy_retune_task]) + assert planner._legacy_retune_cleanup is None + assert environment.calls.count("leases.check") == 1 # no retry + # The default SIGTERM stays unless a Lease was found. + assert planner.retune_lease_found.is_set() is ( + outcome in ("not_found", "bug") + ) + records = [r for r in caplog.records if "legacy clean-up" in r.message] + levels = {r.levelname for r in records} + if outcome == "nothing_left": + assert records == [] + elif outcome == "forbidden": + assert levels == {"DEBUG"} # the default chart: no Lease RBAC + elif outcome == "not_found": + assert levels == {"INFO"} and "skipped" in caplog.text + else: + assert levels == {"ERROR"} and "a bug" in caplog.text + finally: + await planner._shutdown_runtime() + assert "cleanup.stop" not in environment.calls + + +async def test_the_legacy_check_retries_a_transient_error_once(monkeypatch) -> None: + from dynamo.planner.core import base as base_module + + monkeypatch.setattr(base_module, "_LEGACY_RETUNE_RETRY_SECONDS", 0.0) + cleanups: list[FakeCleanup] = [] + + async def build(host): + cleanups.append(FakeCleanup(environment.calls)) + return cleanups[-1] + + for leases, found in ( + ([KubeError("timeout"), (STATE_NAME,)], True), + ([KubeError("timeout"), KubeError("timeout")], False), + ): + planner, environment = planner_without_retune( + monkeypatch, kubernetes_connector(), leases=leases, build=build + ) + await planner._async_init() + try: + await planner._legacy_retune_task + assert environment.calls.count("leases.check") == 2 + assert (planner._legacy_retune_cleanup is not None) is found + assert planner.retune_lease_found.is_set() is found + finally: + await planner._shutdown_runtime() + + +async def test_the_legacy_check_is_kubernetes_only(monkeypatch) -> None: + for controller in (None, SimpleNamespace(kube_api=object())): + planner, environment = planner_without_retune(monkeypatch, controller) + await planner._async_init() + try: + assert planner._legacy_retune_task is None + assert not planner.may_hold_retune_lease + finally: + await planner._shutdown_runtime() + assert "leases.check" not in environment.calls + + +async def test_dgdr_canary_and_context() -> None: + world = World() + world.kube.put_dgd(dgd_object(replicas=2)) + config = live_config() + process = await leading(world) + canary = bootstrap.dgdr_canary(config) + assert canary(process.sup) is True + created = world.kube.calls_of("create_request")[-1] + assert created.kwargs["dry_run"] is True and world.kube.requests == {} + world.kube.reject_create = ( + "strict decoding error: unknown field spec.workload.trace" + ) + assert canary(process.sup) is False + world.spec.replicas = None + await process.sup.control_round() + assert canary(process.sup) is False + + context = bootstrap.dgdr_context_factory(config, cut_window="cutter")(process.sup) + assert context.dgd == process.sup.dgd and context.kube is world.kube + assert context.cut_window == "cutter" and context.backend == config.backend + assert context.is_leader() is True + + +# ------------------------------------------- initials before router writes + + +async def test_a_stop_before_the_initials_write_leaves_no_router_retuned() -> None: + """A router seen at revision 0 is latched and would be PATCHed in the same + round. Stopped (SIGTERM) before the state records its initial, the next + planner rolls back with an advisory target, which never resets a router + whose initial it does not know: so no router may be PATCHed before its + initial is persisted.""" + + world = World() + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + first = await started(world) + store = StateStore(world.kube, first.sup.dgd) + await first.sup.lease_round() + snapshot = store.read() + # An active deployment is persisted, and the routers restart before the + # actuator has latched them (revision 0, initials unknown). + store.write( + dataclasses.replace( + snapshot.state, flat=chosen, by_wt=by_wt("a", chosen), deploy_seq=2 + ), + snapshot, + ) + first.sup._snapshot = store.read() + world.kube.fail( + "replace_lease", + KubeError("timeout"), + match=lambda target: target[0] == STATE_NAME, + ) + await first.sup.control_round() + assert world.frontend.writes == [] # held: the initial is not persisted + assert world.shared_state().initials == {} + await first.sup.stop() + + second = await leading(world, "p2", advisory=True) + await settle([second], world, rounds=4) + assert frontend_configs(world) == [VARIANTS["a"]["aggregated"]] + assert not world.shared_state().rollback_pending + assert world.frontend.writes == [] + + +async def test_initials_are_persisted_before_the_routers_are_patched() -> None: + world = World() + process = await leading(world) + chosen = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + process.controller.restore_desired(chosen) + enqueue(process, make_decision(K0 + 1, "applied", chosen=chosen)) + await process.sup.outbox_round() + world.frontend.restart(VARIANTS["a"]) # new routers at revision 0 + (router_id,) = world.frontend.routers + patched_with: list[dict] = [] + original = world.frontend.write + + def write(router, kind, *args): + if kind == "patch": + patched_with.append(dict(world.shared_state().initials)) + original(router, kind, *args) + + world.frontend.write = write # type: ignore[method-assign] + # One round: latch, persist the initial, then PATCH. + await process.sup.control_round() + assert patched_with and all(router_id in seen for seen in patched_with) + assert frontend_configs(world) == [by_wt("a", chosen)["aggregated"]] + assert process.sup.last_report.converged + + +async def test_a_state_of_an_unknown_format_version_is_never_overwritten() -> None: + world = World() + newer = json.dumps({"version": 99, "dgd_uid": "uid-1", "anything": "new"}) + world.kube.create_lease( + { + "metadata": {"name": STATE_NAME, "annotations": {STATE_ANNOTATION: newer}}, + "spec": {}, + } + ) + process = await started(world) + for _ in range(2): + await process.sup.lease_round() + assert process.sup.is_leader() and not process.sup.initialized + lease = world.kube.leases[STATE_NAME] + assert lease["metadata"]["annotations"][STATE_ANNOTATION] == newer + assert process.metrics.of("failure") == ["state_version", "state_version"] + + +# ------------------------------------------- with the real retune controller + + +async def test_real_controller_launches_harvests_and_persists() -> None: + """The supervisor's hooks drive ``RetuneController`` end to end: gate, + launch with the spec replicas and the abort event, sweep_fn, harvest + gating, on_effects, outbox and the persisted deployment.""" + + world = World() + process = await leading(world) + sup, controller = process.sup, process.controller + assert sup.ready() and controller.router_keys == KEYS + best = {CREDIT: 0.5, SCALE: 1.0, TEMP: 0.2} + + async def tick(now_s): + world.clock.wall = now_s + result = await controller.on_tick(now_s, SPEC_COUNTS) + sup.on_effects( + SimpleNamespace( + retune_decisions=result.decisions, + router_parameters=result.router_parameters, + ) + ) + await sup.outbox_round() + return result + + # First enable recorded the window in progress at start as decided: + # neither it (its trace began mid-way) nor the due window launches. + await tick(START_WALL) + await tick(window(K0)[2]) + assert controller._pending is None and world.sweeps.requests == [] + # The next window launches at its L_k with the spec replicas. + await tick(window(K0 + 1)[2]) + pending = controller._pending + assert pending is not None and pending.abort is not None + assert world.sweeps.finish(controller, scored(best)) + (request,) = world.sweeps.requests + assert request.window_index == K0 + 1 and request.generation == 0 + assert request.spec_replicas == {"aggregated": 2} and request.replicas_verified + # Harvested on the first tick at least a scale interval after the launch: + # applied, persisted, then deployed. + result = await tick(window(K0 + 1)[2] + 5.0) + (decided,) = result.decisions + assert decided.outcome == "applied" + assert decided.decision_id == f"retune-w{end_ms(K0 + 1)}" + state = world.shared_state() + assert state.flat == best and state.deploy_seq == 2 + assert state.last_window_end_ms == end_ms(K0 + 1) + assert state.by_wt == by_wt("a", best) + await sup.control_round() + assert frontend_configs(world) == [by_wt("a", best)["aggregated"]] + # A re-initialization restores the persisted point into the controller. + controller.restore_desired(startup_flat("a")) + sup._request_reinit("test") + await sup.lease_round() + assert controller.current_point().router == best diff --git a/components/src/dynamo/planner/tests/unit/test_live_retune_window_cut.py b/components/src/dynamo/planner/tests/unit/test_live_retune_window_cut.py new file mode 100644 index 000000000000..c686a1b42be9 --- /dev/null +++ b/components/src/dynamo/planner/tests/unit/test_live_retune_window_cut.py @@ -0,0 +1,1244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the live retune window cutter. + +The golden tests compare the cut window, decompressed, byte for byte with +``read_trace_window`` over the same clean segments. A few tests spawn the +real child process; they carry a timeout. +""" + +import errno +import gzip +import io +import json +import os +import pickle +import random +import shutil +import stat +import threading +import time +from dataclasses import replace + +import pytest + +from dynamo.planner.retune.live import window_cut +from dynamo.planner.retune.live.contracts import WindowCutResult +from dynamo.planner.retune.live.window_cut import ( + MAX_DROPPED_MEMBER_RATIO, + SEGMENT_MTIME_MARGIN_S, + TEMP_NAME_RE, + WINDOW_CUT_FAILURES, + WindowCutError, + _decode_exact_member, + _Garbage, + _GzipMemberReader, + _Member, + _run_in_child, + _segment_files, + cut_window, + cut_window_inprocess, + temp_file_name, +) +from dynamo.planner.retune.window import ROW_COUNT_KEYS, read_trace_window + +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + +SCHEMA = "dynamo.request.trace.v1" +START, END, AS_OF = 5_000, 15_000, 15_500 +_UNSET = object() + + +# ------------------------------------------------------------------ helpers + + +def _row( + request_id, + received_ms=_UNSET, + *, + total_ms=10.0, + event_ms=None, + isl=128, + osl=8, + block_size=64, + model="m", + agent_context=None, + output_tokens=_UNSET, +): + request = {"request_id": request_id} + if model is not None: + request["model"] = model + if received_ms is not _UNSET and received_ms is not None: + request["request_received_ms"] = received_ms + if total_ms is not None: + request["total_time_ms"] = total_ms + if output_tokens is not None: + request["output_tokens"] = osl if output_tokens is _UNSET else output_tokens + request["replay"] = { + "trace_block_size": block_size, + "input_length": isl, + "input_sequence_hashes": [len(request_id), 7, 11_400_714_819_323_198_485], + } + if event_ms is None: + base = received_ms if isinstance(received_ms, int) else 0 + event_ms = base + (round(total_ms) if total_ms is not None else 0) + record = { + "schema": SCHEMA, + "event_type": "request_end", + "event_time_unix_ms": event_ms, + "request": request, + } + if agent_context is not None: + record["agent_context"] = agent_context + return record + + +def _envelope(record, timestamp=0): + # The frontend's gzip sink wraps every record like this. + return {"timestamp": timestamp, "event": record} + + +def _lines(records): + return b"".join(json.dumps(record).encode("utf-8") + b"\n" for record in records) + + +def _member(records): + return gzip.compress(_lines(records)) + + +def _write(path, *chunks): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"".join(chunks)) + return path + + +def _out(tmp_path, name="w.jsonl.gz"): + return tmp_path / "windows" / "dgd" / name + + +def _cut(root, tmp_path, start=START, end=END, as_of=AS_OF, *, max_rows=10**6): + out = _out(tmp_path) + return cut_window_inprocess(root, start, end, as_of, out, max_rows=max_rows), out + + +def _golden(root, tmp_path, start=START, end=END, as_of=AS_OF): + out = tmp_path / "golden" / "window.jsonl" + window = read_trace_window([str(root)], start, end, as_of, out) + return window, out.read_bytes() + + +def _assert_golden(result, out, window, golden_bytes): + assert isinstance(result, WindowCutResult) + assert result.num_requests == window.num_requests + assert result.trace_block_size == window.trace_block_size + assert result.row_counts == window.row_counts + assert set(result.row_counts) == set(ROW_COUNT_KEYS) + if window.num_requests: + assert result.path == str(out) + assert gzip.decompress(out.read_bytes()) == golden_bytes + else: + assert result.path is None + assert not out.exists() + + +def _temp_files(directory): + if not directory.exists(): + return [] + return [name for name in os.listdir(directory) if name.startswith(".tmp-")] + + +def _random_trace(root, seed, *, pods=3, segments=3, members=5): + """Segments with every row shape ``read_trace_window`` handles. + + Arrivals and durations are coarse so (arrival, completion) pairs tie + across files; some request IDs repeat (first in file order wins), some + are not ASCII. + """ + + rng = random.Random(seed) + issued = [] + serial = 0 + for pod in range(pods): + for segment in range(segments): + chunks = [] + for _ in range(members): + records = [] + for _ in range(rng.randint(0, 14)): + roll = rng.random() + if roll < 0.04: + records.append( + { + "schema": SCHEMA, + "event_type": rng.choice(["tool_start", "tool_end"]), + "event_time_unix_ms": rng.randrange(20_000), + "agent_context": {"session_id": "s"}, + } + ) + continue + if issued and roll < 0.16: + request_id = rng.choice(issued) + else: + serial += 1 + prefix = rng.choice(["req-", "req-", "réq-", "請求-", "Req-"]) + request_id = f"{prefix}{serial % 97}-{serial}" + issued.append(request_id) + received = rng.randrange(0, 400) * 50 + total = rng.choice([10.0, 10.5, 49.5, 99.4, 600.0, 1_200.25]) + kwargs = {} + if rng.random() < 0.08: + kwargs["received_ms"] = None + kwargs["event_ms"] = received + round(total) + else: + kwargs["received_ms"] = received + if rng.random() < 0.05: + kwargs["total_ms"] = None + kwargs["event_ms"] = received + rng.randrange(0, 900) + else: + kwargs["total_ms"] = total + if rng.random() < 0.03: + kwargs["output_tokens"] = None + if rng.random() < 0.05: + kwargs["agent_context"] = {"session_id": f"s{serial}"} + if rng.random() < 0.1: + kwargs["model"] = None + received_arg = kwargs.pop("received_ms") + record = _row( + request_id, + received_arg, + isl=rng.randrange(1, 4096), + osl=rng.randrange(1, 512), + **kwargs, + ) + if rng.random() < 0.5: + record = _envelope(record, timestamp=rng.randrange(10**6)) + records.append(record) + payload = _lines(records) + if rng.random() < 0.1: + payload += b"\n \n" # blank lines are skipped + chunks.append(payload) + _write( + root / f"frontend-{pod}" / f"req.{segment:06d}.jsonl.gz", + *(gzip.compress(chunk) for chunk in chunks), + ) + return root + + +def _items(data): + return list(_GzipMemberReader(io.BytesIO(data))) + + +# ----------------------------------------------------------- gzip members + + +@pytest.fixture(params=["default", "tiny", "small"]) +def reader_chunks(request, monkeypatch): + """Run reader tests with reads and decode steps that split headers.""" + + sizes = {"default": None, "tiny": (7, 5, 16), "small": (64, 3, 128)}[request.param] + if sizes is not None: + read_chunk, decode_step, compact_at = sizes + monkeypatch.setattr(window_cut, "_READ_CHUNK", read_chunk) + monkeypatch.setattr(window_cut, "_DECODE_STEP", decode_step) + monkeypatch.setattr(window_cut, "_COMPACT_AT", compact_at) + return request.param + + +def _spans(*members): + offsets = [0] + for member in members: + offsets.append(offsets[-1] + len(member)) + return offsets + + +def test_clean_multi_member_file(reader_chunks): + payloads = [b'{"a":1}\n', b"", b'{"b":2}\n{"c":3}\n', b"x" * 5000 + b"\n"] + members = [gzip.compress(payload) for payload in payloads] + offsets = _spans(*members) + + items = _items(b"".join(members)) + + assert items == [ + _Member(offsets[i], offsets[i + 1], payloads[i]) for i in range(len(payloads)) + ] + + +@pytest.mark.parametrize("keep", ["header", "middle", "trailer"]) +def test_truncated_last_member_is_dropped(reader_chunks, keep): + first, last = gzip.compress(b"a\n" * 50), gzip.compress(b"b\n" * 50) + cut = {"header": 5, "middle": len(last) // 2, "trailer": len(last) - 3}[keep] + + items = _items(first + last[:cut]) + + assert items == [ + _Member(0, len(first), b"a\n" * 50), + _Garbage(len(first), len(first) + cut, at_eof=True), + ] + + +def test_partial_member_followed_by_complete_members(reader_chunks): + first = gzip.compress(b"first\n" * 40) + full = gzip.compress(b"".join(b"partial %d\n" % i for i in range(400))) + partial = full[: len(full) // 2] + third = gzip.compress(b"third\n") + fourth = gzip.compress(b"fourth\n") + offsets = _spans(first, partial, third, fourth) + + items = _items(first + partial + third + fourth) + + assert items == [ + _Member(offsets[0], offsets[1], b"first\n" * 40), + _Garbage(offsets[1], offsets[2], at_eof=False), + _Member(offsets[2], offsets[3], b"third\n"), + _Member(offsets[3], offsets[4], b"fourth\n"), + ] + + +@pytest.mark.parametrize("where", ["crc", "isize", "data"]) +def test_flipped_member_is_dropped_and_reader_resyncs(reader_chunks, where): + first = gzip.compress(b"first\n") + middle = bytearray(gzip.compress(b"middle row\n" * 30)) + index = {"crc": len(middle) - 8, "isize": len(middle) - 2, "data": 14}[where] + middle[index] ^= 0x40 + last = gzip.compress(b"last\n") + offsets = _spans(first, middle, last) + + items = _items(first + bytes(middle) + last) + + assert items == [ + _Member(offsets[0], offsets[1], b"first\n"), + _Garbage(offsets[1], offsets[2], at_eof=False), + _Member(offsets[2], offsets[3], b"last\n"), + ] + + +def test_zero_filled_hole_between_members(reader_chunks): + first, last = gzip.compress(b"first\n"), gzip.compress(b"last\n") + hole = bytes(4096) + + items = _items(first + hole + last) + + assert items == [ + _Member(0, len(first), b"first\n"), + _Garbage(len(first), len(first) + len(hole), at_eof=False), + _Member(len(first) + len(hole), len(first) + len(hole) + len(last), b"last\n"), + ] + + +def test_false_gzip_magic_inside_garbage_is_one_run(reader_chunks): + first, last = gzip.compress(b"first\n"), gzip.compress(b"last\n") + junk = b"\x1f\x8b\x08\x00junk" * 20 + b"\x1f\x8b" + + items = _items(first + junk + last) + + assert items == [ + _Member(0, len(first), b"first\n"), + _Garbage(len(first), len(first) + len(junk), at_eof=False), + _Member(len(first) + len(junk), len(first) + len(junk) + len(last), b"last\n"), + ] + + +def test_garbage_only_and_empty_files(reader_chunks): + assert _items(b"") == [] + assert _items(b"plain text, not gzip\n") == [_Garbage(0, 21, at_eof=True)] + member = gzip.compress(b"a\n") + assert _items(member + bytes(10)) == [ + _Member(0, len(member), b"a\n"), + _Garbage(len(member), len(member) + 10, at_eof=True), + ] + + +def test_reader_starts_at_an_offset(): + first, second = gzip.compress(b"first\n"), gzip.compress(b"second\n") + source = io.BytesIO(first + second) + + items = list(_GzipMemberReader(source, len(first))) + + assert items == [_Member(len(first), len(first) + len(second), b"second\n")] + + +def test_decode_exact_member(): + member = gzip.compress(b"row\n") + assert _decode_exact_member(member) == b"row\n" + assert _decode_exact_member(member + b"x") is None + assert _decode_exact_member(member[:-1]) is None + flipped = bytearray(member) + flipped[-6] ^= 1 + assert _decode_exact_member(bytes(flipped)) is None + assert _decode_exact_member(b"") is None + + +# ------------------------------------------------------------ golden cuts + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_cut_matches_read_trace_window(tmp_path, seed): + root = _random_trace(tmp_path / "segments" / "dgd", seed) + window, golden_bytes = _golden(root, tmp_path) + assert window.num_requests > 50 + assert window.row_counts["duplicate"] > 0 + assert window.row_counts["censored"] > 0 + + result, out = _cut(root, tmp_path) + + _assert_golden(result, out, window, golden_bytes) + assert result.dropped_members == 0 + assert result.dropped_bytes == 0 + assert not result.reused + assert stat.S_IMODE(out.stat().st_mode) == 0o644 + assert _temp_files(out.parent) == [] + + +def test_plain_jsonl_files_are_ignored(tmp_path): + # The reconciler never collects them, and the live sink writes gzip. + root = _random_trace(tmp_path / "segments" / "dgd", 3) + window, golden_bytes = _golden(root, tmp_path) + for pod in ("frontend-0", "frontend-9"): + _write(root / pod / "req.000000.jsonl", _lines([_row("plain", START + 1)])) + + result, out = _cut(root, tmp_path) + + _assert_golden(result, out, window, golden_bytes) + + +def test_cut_matches_golden_under_eviction_and_small_reads(tmp_path, monkeypatch): + root = _random_trace(tmp_path / "segments" / "dgd", 4, members=8) + window, golden_bytes = _golden(root, tmp_path) + monkeypatch.setattr(window_cut, "_PASS2_CACHE_BYTES", 1) + monkeypatch.setattr(window_cut, "_PASS2_BLOCK_ROWS", 7) + monkeypatch.setattr(window_cut, "_READ_CHUNK", 97) + monkeypatch.setattr(window_cut, "_DECODE_STEP", 13) + monkeypatch.setattr(window_cut, "_COMPACT_AT", 64) + monkeypatch.setattr(window_cut, "_WRITE_BUFFER", 10) + monkeypatch.setattr(window_cut, "_OPEN_FILES", 2) + loads = [] + original_load = window_cut._MemberLines._load + + def counting_load(self, member_id): + loads.append(member_id) + return original_load(self, member_id) + + monkeypatch.setattr(window_cut._MemberLines, "_load", counting_load) + + result, out = _cut(root, tmp_path) + + _assert_golden(result, out, window, golden_bytes) + # Every member was evicted after each row and decoded again. + assert len(loads) == window.num_requests + assert len(set(loads)) < len(loads) + + +def test_dedup_compaction_keeps_first_rows_and_bounds_rows(tmp_path, monkeypatch): + rows = [_row(f"r{i}", START + 10 * i) for i in range(10)] + # Many repeats of the same IDs, in later files: the first row wins. + repeats = [_row(f"r{i % 10}", START + 5 + i, isl=999) for i in range(40)] + root = tmp_path / "segments" + _write(root / "a" / "req.000000.jsonl.gz", _member(rows)) + _write(root / "b" / "req.000000.jsonl.gz", *(_member([row]) for row in repeats)) + window, golden_bytes = _golden(root, tmp_path) + assert window.num_requests == 10 + monkeypatch.setattr(window_cut, "_DEDUP_SLACK_MIN", 1) + monkeypatch.setattr(window_cut, "_DEDUP_SLACK_RATIO", 0.0) + compactions = [] + original_dedup = window_cut._Selection._dedup + + def counting_dedup(self): + compactions.append(self.num_rows) + original_dedup(self) + + monkeypatch.setattr(window_cut._Selection, "_dedup", counting_dedup) + + result, out = _cut(root, tmp_path, max_rows=10) + + _assert_golden(result, out, window, golden_bytes) + assert result.row_counts["duplicate"] == 40 + assert len(compactions) > 2 # during pass 1, not only at the end + + +def test_too_many_rows_is_window_too_large(tmp_path, monkeypatch): + root = tmp_path / "segments" + rows = [_row(f"r{i}", START + i) for i in range(30)] + _write(root / "a" / "req.000000.jsonl.gz", _member(rows)) + monkeypatch.setattr(window_cut, "_DEDUP_SLACK_MIN", 1) + monkeypatch.setattr(window_cut, "_DEDUP_SLACK_RATIO", 0.0) + + result, out = _cut(root, tmp_path, max_rows=30) + assert result.num_requests == 30 + out.unlink() + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path, max_rows=29) + assert caught.value.reason == "window_too_large" + + +def test_empty_window_writes_no_file(tmp_path): + root = tmp_path / "segments" + _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("early", START - 100), _row("late", END + 1)]), + _member([_row("no-output", START + 1, output_tokens=None)]), + ) + window, golden_bytes = _golden(root, tmp_path) + + result, out = _cut(root, tmp_path) + + assert golden_bytes == b"" + _assert_golden(result, out, window, golden_bytes) + assert result.row_counts["missing_output_tokens"] == 1 + assert not out.parent.exists() + + +def test_empty_segments_root_is_an_empty_window(tmp_path): + root = tmp_path / "segments" + (root / "frontend-0").mkdir(parents=True) + + result, out = _cut(root, tmp_path) + + assert result.path is None + assert result.num_requests == 0 + assert result.trace_block_size is None + + +def test_missing_segments_root_fails(tmp_path): + with pytest.raises(WindowCutError) as caught: + _cut(tmp_path / "nowhere", tmp_path) + assert caught.value.reason == "window_cut_failed" + assert "does not exist" in str(caught.value) + + +def test_mixed_block_sizes_fail(tmp_path): + root = tmp_path / "segments" + _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("a", START, block_size=64), _row("b", START, block_size=16)]), + ) + with pytest.raises(ValueError, match="mixes trace_block_size"): + _golden(root, tmp_path) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + assert caught.value.reason == "window_cut_failed" + assert "mixes trace_block_size values [16, 64]" in str(caught.value) + + +@pytest.mark.parametrize( + "line, message", + [ + (b"{not json", "invalid JSON"), + (b"[1, 2]", "must be a JSON object"), + (b'{"schema": "other", "event_type": "request_end"}', "expected"), + ( + json.dumps( + { + "schema": SCHEMA, + "event_type": "request_end", + "event_time_unix_ms": 1, + "request": { + "request_id": "x", + "output_tokens": 1, + "replay": {"input_length": 1, "trace_block_size": 64}, + }, + } + ).encode(), + "input_sequence_hashes", + ), + ], +) +def test_malformed_row_in_a_valid_member_is_trace_corrupt(tmp_path, line, message): + # The member's CRC is good: the writer wrote this, as replay would see it. + root = tmp_path / "segments" + _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("good", START)]), + gzip.compress(line + b"\n"), + ) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + assert caught.value.reason == "trace_corrupt" + assert message in str(caught.value) + + +def test_out_of_range_value_is_trace_corrupt(tmp_path): + root = tmp_path / "segments" + _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("huge", START, block_size=2**70)]), + ) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + assert caught.value.reason == "trace_corrupt" + + +# --------------------------------------------------- dropped gzip members + + +def _tiny_members(prefix, count, first_ms=START): + return [ + _member([_envelope(_row(f"{prefix}{i}", first_ms + i), timestamp=i)]) + for i in range(count) + ] + + +def _cut_and_twin(tmp_path, build): + """Cut ``build(bad=True)`` and the golden of ``build(bad=False)``. + + ``build(root, bad)`` writes the segments with or without the bad bytes; + the rows of a dropped member never reach the window. + """ + + clean = tmp_path / "clean" + build(clean, bad=False) + window, golden_bytes = _golden(clean, tmp_path) + damaged = tmp_path / "damaged" + build(damaged, bad=True) + result, out = _cut(damaged, tmp_path) + _assert_golden(result, out, window, golden_bytes) + return result + + +@pytest.mark.parametrize("truncated", ["req.000001", "req.000000"]) +def test_truncated_tail_of_a_segment_is_not_corruption(tmp_path, truncated): + # req.000001's tail is the flush its writer is appending. req.000000's is + # the last flush of a writer that crashed: after an in-place restart the + # sink continues in req.000001. + lost = _member([_row("lost", START + 900)]) + + def build(root, bad): + for name, prefix in (("req.000000", "old"), ("req.000001", "new")): + tail = [lost[: len(lost) // 2]] if bad and name == truncated else [] + _write( + root / "a" / f"{name}.jsonl.gz", + *_tiny_members(prefix, 3), + *tail, + ) + + result = _cut_and_twin(tmp_path, build) + + # 1 of 7 members is far above the ratio, but it is a writer's tail. + assert result.dropped_members == 1 + assert result.dropped_bytes == len(lost) // 2 + + +@pytest.mark.parametrize("damage", ["partial", "crc", "hole"]) +def test_damaged_members_within_the_ratio_are_dropped(tmp_path, damage): + # 1 bad member among 100 is exactly the limit: the cut goes on. Its bytes + # are at most one member's, below the limit too. + count = int(1 / MAX_DROPPED_MEMBER_RATIO) - 1 + bad_row = _member([_row("bad", START + 50)]) + if damage == "partial": + damaged = bad_row[: len(bad_row) - 12] + elif damage == "crc": + flipped = bytearray(bad_row) + flipped[-7] ^= 0x10 + damaged = bytes(flipped) + else: + damaged = bytes(len(bad_row)) + assert len(damaged) <= min(map(len, _tiny_members("m", count))) + + def build(root, bad): + members = _tiny_members("m", count) + if bad: + members.insert(count // 2, damaged) + _write(root / "a" / "req.000000.jsonl.gz", *members) + + result = _cut_and_twin(tmp_path, build) + + assert result.dropped_members == 1 + assert result.dropped_bytes == len(damaged) + + +def test_damaged_members_above_the_ratio_are_trace_corrupt(tmp_path): + root = tmp_path / "segments" + members = _tiny_members("m", 98) + members.insert(10, bytes(64)) + members.insert(50, bytes(64)) + _write(root / "a" / "req.000000.jsonl.gz", *members) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + assert caught.value.reason == "trace_corrupt" + assert "2 of 100 gzip members" in str(caught.value) + + +def test_hole_spanning_many_members_is_trace_corrupt(tmp_path): + # A zero-filled hole over 20 of 200 members is one run of bad bytes: 1 of + # 181 members is within the ratio, but 10% of the bytes is not. + root = tmp_path / "segments" + members = _tiny_members("m", 200) + hole = sum(map(len, members[50:70])) + data = b"".join([*members[:50], bytes(hole), *members[70:]]) + _write(root / "a" / "req.000000.jsonl.gz", data) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + + assert caught.value.reason == "trace_corrupt" + message = str(caught.value) + assert "1 of 181 gzip members" in message + assert f"({hole} of " in message + assert f"req.000000.jsonl.gz@{sum(map(len, members[:50]))}+{hole}" in message + + +def test_hole_within_the_byte_ratio_is_dropped(tmp_path): + # A hole over 1 of 200 members: 1 run, under 1% of the bytes. + members = _tiny_members("m", 200) + + def build(root, bad): + kept = [*members[:50], *members[51:]] + if bad: + kept.insert(50, bytes(len(members[50]))) + _write(root / "a" / "req.000000.jsonl.gz", *kept) + + result = _cut_and_twin(tmp_path, build) + + assert result.dropped_members == 1 + assert result.dropped_bytes == len(members[50]) + + +def test_segment_files_order(tmp_path): + root = tmp_path / "segments" + for name in ( + "b/req.000999.jsonl.gz", + "b/req.1000000.jsonl.gz", + "a/req.000000.jsonl.gz", + "a/req.000001.jsonl", # plain JSONL: ignored + "a/notes.txt", + "c/single.jsonl.gz", + ): + _write(root / name, b"") + + files = _segment_files(str(root)) + + assert files == [ + str(root / "a/req.000000.jsonl.gz"), + str(root / "b/req.000999.jsonl.gz"), + str(root / "b/req.1000000.jsonl.gz"), + str(root / "c/single.jsonl.gz"), + ] + + +def test_a_segments_root_that_is_not_a_gzip_segment_fails(tmp_path): + plain = _write(tmp_path / "req.000000.jsonl", _lines([_row("a", START)])) + + with pytest.raises(WindowCutError) as caught: + _cut(plain, tmp_path) + assert caught.value.reason == "window_cut_failed" + single = _write(tmp_path / "req.000000.jsonl.gz", _member([_row("a", START)])) + assert _cut(single, tmp_path)[0].num_requests == 1 + + +def test_segments_last_written_before_the_window_are_not_read(tmp_path): + # A row is written when its request completes, after it arrived: a + # segment last modified before the window started holds none of its + # rows. This one does (it cannot in production) to show it is not read, + # and its hole would fail the cut as trace_corrupt if it were. + start, end = 1_700_000_000_000, 1_700_000_010_000 + root = tmp_path / "segments" + old = _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("old-0", start + 1)]), + bytes(64), + _member([_row("old-1", start + 2)]), + ) + _write(root / "a" / "req.000001.jsonl.gz", _member([_row("new", start + 3)])) + out = _out(tmp_path) + + def cut(): + return cut_window_inprocess(root, start, end, end + 500, out, max_rows=10) + + skipped = start / 1000 - SEGMENT_MTIME_MARGIN_S - 1 + os.utime(old, (skipped, skipped)) + result = cut() + assert result.num_requests == 1 + assert result.dropped_members == 0 + out.unlink() + + within_margin = start / 1000 - SEGMENT_MTIME_MARGIN_S + 1 + os.utime(old, (within_margin, within_margin)) + with pytest.raises(WindowCutError) as caught: + cut() + assert caught.value.reason == "trace_corrupt" + + +def test_vanished_file_is_skipped(tmp_path, monkeypatch): + root = tmp_path / "segments" + _write(root / "a" / "req.000000.jsonl.gz", _member([_row("a", START)])) + listed = _segment_files(str(root)) + gone = str(root / "a" / "req.000001.jsonl.gz") + monkeypatch.setattr(window_cut, "_segment_files", lambda _: [*listed, gone]) + + result, _ = _cut(root, tmp_path) + + assert result.num_requests == 1 + + +def _two_pods(root): + for pod in ("a", "b"): + _write( + root / pod / "req.000000.jsonl.gz", + _member([_row(f"{pod}-{i}", START + i) for i in range(5)]), + ) + return root + + +def _skip_as_root(): + if hasattr(os, "geteuid") and os.geteuid() == 0: + pytest.skip("root ignores directory permissions") + + +@pytest.mark.parametrize("mode", [0o000, 0o600]) +def test_unreadable_pod_directory_fails(tmp_path, mode): + # 0o000 cannot be listed; 0o600 can, but its files cannot be stat-ed. + # Either way that pod's rows would be missing from the window. + _skip_as_root() + root = _two_pods(tmp_path / "segments") + pod = root / "b" + pod.chmod(mode) + try: + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + finally: + pod.chmod(0o755) + + assert caught.value.reason == "window_cut_failed" + assert str(pod) in str(caught.value) + assert "PermissionError" in str(caught.value) + + +def test_unreadable_segments_root_fails(tmp_path): + # Not an empty window: the cut cannot see any segment. + _skip_as_root() + root = _two_pods(tmp_path / "segments") + root.chmod(0o000) + try: + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + finally: + root.chmod(0o755) + + assert caught.value.reason == "window_cut_failed" + assert str(root) in str(caught.value) + + +@pytest.mark.parametrize("code", [errno.ENOENT, errno.EACCES]) +def test_segment_walk_skips_only_vanished_entries(tmp_path, monkeypatch, code): + root = _two_pods(tmp_path / "segments") + pod = root / "b" + real = os.scandir + + def failing(path, *args, **kwargs): + if os.fspath(path) == str(pod): + raise OSError(code, os.strerror(code), os.fspath(path)) + return real(path, *args, **kwargs) + + monkeypatch.setattr(os, "scandir", failing) + + if code == errno.ENOENT: + # The reconciler deleted pod b's segments under the walk. + assert _segment_files(str(root)) == [str(root / "a" / "req.000000.jsonl.gz")] + return + with pytest.raises(WindowCutError) as caught: + _segment_files(str(root)) + assert caught.value.reason == "window_cut_failed" + assert str(pod) in str(caught.value) + assert os.strerror(code) in str(caught.value) + + +# -------------------------------------------------------- atomic output + + +def _reused(result, path): + return replace( + result, path=str(path), reused=True, dropped_members=0, dropped_bytes=0 + ) + + +def _header_comment(path): + data = path.read_bytes() + return json.loads(data[10 : data.index(b"\x00", 10)]) + + +def test_published_window_is_reused_without_reading_the_trace(tmp_path): + root = _random_trace(tmp_path / "segments", 5) + first, out = _cut(root, tmp_path) + published = out.read_bytes() + assert _header_comment(out)["selection"] == { + "start_ms": START, + "end_ms": END, + "as_of_ms": AS_OF, + } + # A later leader adopting the window must not fail it on segments that + # changed, broke or were collected since it was published. + shutil.rmtree(root) + + result, _ = _cut(root, tmp_path) + + assert result == _reused(first, out) + assert out.read_bytes() == published + assert _temp_files(out.parent) == [] + + +@pytest.mark.parametrize( + "selection", + [{"start": START + 1000}, {"as_of": AS_OF + 100}], + ids=["start", "as_of"], +) +def test_published_window_of_another_selection_is_not_reused(tmp_path, selection): + root = _random_trace(tmp_path / "segments", 5) + _, out = _cut(root, tmp_path) + published = out.read_bytes() + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path, **selection) + + assert caught.value.reason == "window_cut_failed" + assert f"was cut for [{START}, {END}) as of {AS_OF}, not [" in str(caught.value) + assert out.read_bytes() == published + assert _temp_files(out.parent) == [] + + +def test_published_window_without_a_selection_is_not_reused(tmp_path): + root = _random_trace(tmp_path / "segments", 5) + _, out = _cut(root, tmp_path) + data = out.read_bytes() + end = data.index(b"\x00", 10) + counts = _header_comment(out) + del counts["selection"] + out.write_bytes(data[:10] + json.dumps(counts).encode("ascii") + data[end:]) + + with pytest.raises(WindowCutError) as caught: + _cut(root, tmp_path) + assert caught.value.reason == "window_cut_failed" + assert "holds no selection" in str(caught.value) + + +def test_a_window_file_without_cut_counts_fails(tmp_path): + root = _random_trace(tmp_path / "segments", 5) + out = _out(tmp_path) + _write(out, gzip.compress(_lines([_row("a", START)]))) + + with pytest.raises(WindowCutError) as caught: + cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10**6) + assert caught.value.reason == "window_cut_failed" + assert "holds no cut counts" in str(caught.value) + + +def _lose_link_race_to(monkeypatch, winner_out): + """Make the cut's link find ``winner_out`` already published at its target.""" + + real_link = os.link + linked = [] + + def racing_link(source, target): + linked.append(os.path.basename(source)) + real_link(winner_out, target) # another process published first + return real_link(source, target) + + monkeypatch.setattr(os, "link", racing_link) + return linked + + +def test_link_race_reports_the_winners_counts(tmp_path, monkeypatch): + root = _random_trace(tmp_path / "segments", 6) + winner_out = _out(tmp_path, "winner.jsonl.gz") + winner = cut_window_inprocess( + _random_trace(tmp_path / "other", 7), + START, + END, + AS_OF, + winner_out, + max_rows=10**6, + ) + assert winner.num_requests != _golden(root, tmp_path)[0].num_requests + out = _out(tmp_path) + linked = _lose_link_race_to(monkeypatch, winner_out) + + result = cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10**6) + + assert result == _reused(winner, out) + assert out.read_bytes() == winner_out.read_bytes() + assert linked[0].startswith(f".tmp-{END}-") + assert _temp_files(out.parent) == [] + + +def test_link_race_won_by_another_selection_fails(tmp_path, monkeypatch): + root = _random_trace(tmp_path / "segments", 6) + winner_out = _out(tmp_path, "winner.jsonl.gz") + cut_window_inprocess(root, START + 1000, END, AS_OF, winner_out, max_rows=10**6) + out = _out(tmp_path) + linked = _lose_link_race_to(monkeypatch, winner_out) + + with pytest.raises(WindowCutError) as caught: + cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10**6) + + assert caught.value.reason == "window_cut_failed" + assert f"was cut for [{START + 1000}, {END}) as of {AS_OF}" in str(caught.value) + assert linked and out.read_bytes() == winner_out.read_bytes() + assert _temp_files(out.parent) == [] + + +def test_link_enoent_is_window_cut_failed(tmp_path, monkeypatch): + root = _random_trace(tmp_path / "segments", 7) + out = _out(tmp_path) + + def vanished_link(source, target): + raise FileNotFoundError(2, "No such file or directory", source) + + monkeypatch.setattr(os, "link", vanished_link) + + with pytest.raises(WindowCutError) as caught: + cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10**6) + assert caught.value.reason == "window_cut_failed" + assert not out.exists() + assert _temp_files(out.parent) == [] + + +def test_output_is_deterministic_and_0644_under_any_umask(tmp_path): + root = _random_trace(tmp_path / "segments", 8) + first, second = _out(tmp_path, "a.jsonl.gz"), _out(tmp_path, "b.jsonl.gz") + previous = os.umask(0o077) + try: + cut_window_inprocess(root, START, END, AS_OF, first, max_rows=10**6) + cut_window_inprocess(root, START, END, AS_OF, second, max_rows=10**6) + finally: + os.umask(previous) + + assert first.read_bytes() == second.read_bytes() + assert stat.S_IMODE(first.stat().st_mode) == 0o644 + with gzip.open(first, "rb") as source: + assert source.read(1) + # A comment (the counts) but no file name or mtime in the gzip header. + header = first.read_bytes()[:10] + assert header[3] == 0x10 + assert header[4:8] == b"\x00\x00\x00\x00" + + +@pytest.mark.parametrize("change", ["rewritten", "truncated", "deleted"]) +def test_segment_changed_between_passes_fails(tmp_path, monkeypatch, change): + root = tmp_path / "segments" + segment = _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("a", START), _row("b", START + 1)]), + _member([_row("c", START + 2)]), + ) + original_write_rows = window_cut._write_rows + + def mutate_then_write(selection, paths, sink): + if change == "rewritten": + segment.write_bytes( + _member([_row("x", START), _row("y", START + 1)]) + + _member([_row("z", START + 2)]) + ) + elif change == "truncated": + segment.write_bytes(segment.read_bytes()[:40]) + else: + segment.unlink() + return original_write_rows(selection, paths, sink) + + monkeypatch.setattr(window_cut, "_write_rows", mutate_then_write) + out = _out(tmp_path) + + with pytest.raises(WindowCutError) as caught: + cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10) + assert caught.value.reason == "window_cut_failed" + assert not out.exists() + assert _temp_files(out.parent) == [] + + +def test_write_failure_is_window_cut_failed_and_leaves_no_files(tmp_path, monkeypatch): + root = _random_trace(tmp_path / "segments", 10) + out = _out(tmp_path) + + def failing_write(selection, paths, sink): + raise OSError(28, "No space left on device") + + monkeypatch.setattr(window_cut, "_write_rows", failing_write) + + with pytest.raises(WindowCutError) as caught: + cut_window_inprocess(root, START, END, AS_OF, out, max_rows=10**6) + assert caught.value.reason == "window_cut_failed" + assert "No space left" in str(caught.value) + assert not out.exists() + assert _temp_files(out.parent) == [] + + +# ------------------------------------------------- arguments and helpers + + +@pytest.mark.parametrize( + "start, end, as_of, out, max_rows", + [ + (10, 10, 20, "w.jsonl.gz", 1), + (10, 20, 19, "w.jsonl.gz", 1), + (10, 20, 20, "w.jsonl", 1), + (10, 20, 20, "w.jsonl.gz", 0), + (10, 20, 20, "w.jsonl.gz", True), + (float("nan"), 20, 20, "w.jsonl.gz", 1), + ("10", 20, 20, "w.jsonl.gz", 1), + ], +) +def test_invalid_arguments_raise_value_error( + tmp_path, start, end, as_of, out, max_rows +): + with pytest.raises(ValueError): + cut_window_inprocess( + tmp_path, start, end, as_of, tmp_path / out, max_rows=max_rows + ) + with pytest.raises(ValueError): + cut_window( + tmp_path, + start, + end, + as_of, + tmp_path / out, + max_rows=max_rows, + deadline_s=time.time() + 60, + ) + + +def test_temp_path_must_be_next_to_out_path(tmp_path): + with pytest.raises(ValueError, match="not next to"): + cut_window_inprocess( + tmp_path, + START, + END, + AS_OF, + tmp_path / "a" / "w.jsonl.gz", + max_rows=1, + temp_path=tmp_path / "b" / temp_file_name(END), + ) + + +def test_temp_file_name(): + assert temp_file_name(1_700_000_000_000.0, "ab12") == ".tmp-1700000000000-ab12" + assert temp_file_name(5) != temp_file_name(5) + # The reconciler finds a temp file's window with TEMP_NAME_RE. + assert TEMP_NAME_RE.match(temp_file_name(5)).group(1) == "5" + with pytest.raises(ValueError): + temp_file_name(5, "../x") + + +def test_window_cut_error_reason_and_pickle(): + error = WindowCutError("trace_corrupt", "bad member") + copy = pickle.loads(pickle.dumps(error)) + assert (copy.reason, str(copy)) == ("trace_corrupt", "bad member") + assert set(WINDOW_CUT_FAILURES) == { + "window_cut_timeout", + "window_too_large", + "trace_corrupt", + "window_cut_failed", + "cancelled", + } + with pytest.raises(ValueError): + WindowCutError("nope", "x") + + +# ------------------------------------------------------- child process + + +def test_interrupted_before_spawn_does_not_start_a_child(tmp_path, monkeypatch): + def no_spawn(*_): + raise AssertionError("a child was started") + + monkeypatch.setattr(window_cut.multiprocessing, "get_context", no_spawn) + out = _out(tmp_path) + + with pytest.raises(WindowCutError) as caught: + cut_window( + tmp_path, START, END, AS_OF, out, max_rows=1, deadline_s=time.time() - 1 + ) + assert caught.value.reason == "window_cut_timeout" + + cancel = threading.Event() + cancel.set() + with pytest.raises(WindowCutError) as caught: + cut_window( + tmp_path, + START, + END, + AS_OF, + out, + max_rows=1, + deadline_s=time.time() + 60, + cancel=cancel, + ) + assert caught.value.reason == "cancelled" + + +@pytest.mark.timeout(120) +def test_cut_window_in_child_matches_golden(tmp_path): + root = _random_trace(tmp_path / "segments" / "dgd", 9) + window, golden_bytes = _golden(root, tmp_path) + out = _out(tmp_path) + + result = cut_window( + root, START, END, AS_OF, out, max_rows=10**6, deadline_s=time.time() + 100 + ) + + _assert_golden(result, out, window, golden_bytes) + assert _temp_files(out.parent) == [] + + +@pytest.mark.timeout(120) +def test_cut_window_in_child_reports_the_failure_reason(tmp_path): + root = tmp_path / "segments" + _write( + root / "a" / "req.000000.jsonl.gz", + _member([_row("a", START), _row("b", START + 1)]), + ) + + with pytest.raises(WindowCutError) as caught: + cut_window( + root, + START, + END, + AS_OF, + _out(tmp_path), + max_rows=1, + deadline_s=time.time() + 100, + ) + assert caught.value.reason == "window_too_large" + assert "more than 1 requests" in str(caught.value) + + +@pytest.mark.timeout(120) +def test_child_that_dies_is_window_cut_failed(): + with pytest.raises(WindowCutError) as caught: + _run_in_child(os._exit, (3,), {}, deadline=time.monotonic() + 100, cancel=None) + assert caught.value.reason == "window_cut_failed" + assert "exited with code 3" in str(caught.value) + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("interrupt", ["cancel", "deadline"]) +def test_interrupt_kills_child_and_removes_its_temp_file(tmp_path, interrupt): + temp = tmp_path / temp_file_name(END) + temp.write_bytes(b"partial window") + cancel = threading.Event() + timer = threading.Timer(0.5, cancel.set) + deadline = time.monotonic() + (100 if interrupt == "cancel" else 0.5) + started = time.monotonic() + timer.start() + try: + with pytest.raises(WindowCutError) as caught: + _run_in_child( + time.sleep, + (600,), + {}, + deadline=deadline, + cancel=cancel if interrupt == "cancel" else None, + cleanup=(str(temp),), + ) + finally: + timer.cancel() + timer.join() + expected = "cancelled" if interrupt == "cancel" else "window_cut_timeout" + assert caught.value.reason == expected + assert time.monotonic() - started < 20 + assert not temp.exists() diff --git a/components/src/dynamo/replay/api.py b/components/src/dynamo/replay/api.py index 0ac237c70b99..e1fab3db88f9 100644 --- a/components/src/dynamo/replay/api.py +++ b/components/src/dynamo/replay/api.py @@ -97,7 +97,7 @@ class _TraceReplayOptions(_CommonReplayOptions, total=False): trace_num_prefix_groups: int report_jsonl_path: str | os.PathLike[str] | None max_sim_time_ms: float | None - # Internal to retune sweeps; not part of the public API. + # Internal; not part of the public API. _planner_warmup_observations: "Sequence[TrafficObservation] | None" _deterministic: bool @@ -155,6 +155,14 @@ def _retune_inputs( AISimulate setting. """ + retune = config.retune + if retune is not None and retune.executor != "local": + # Replay runs local sweeps only; executor 'dgdr' needs a live + # Planner, so do not silently replay it as a local sweep. + raise ValueError( + f"planner.retune executor {retune.executor!r} is not supported in " + "replay; replay retune runs executor 'local'" + ) if trace_format != "dynamo": raise ValueError( "planner.retune replays a Dynamo request trace: it requires " @@ -420,8 +428,9 @@ def run_trace_replay( # Internal: a retune sweep's inner planner is warmed with the history it # already bucketed instead of its load_predictor_warmup_trace. _planner_warmup_observations=None, - # Internal: a retune sweep replays the baseline and every candidate with - # the canonical router seed and ordinal request ids (common random numbers). + # Internal: replay with the canonical router seed and ordinal request ids + # (common random numbers). Retune sweeps and every trace replay through + # the Sweeper's Dynamo runner set it. _deterministic=False, ) -> ReplayReport | dict[str, Any]: """Run trace replay. diff --git a/components/src/dynamo/replay/simulation.py b/components/src/dynamo/replay/simulation.py index 5b1a5bd49936..251e033199b4 100644 --- a/components/src/dynamo/replay/simulation.py +++ b/components/src/dynamo/replay/simulation.py @@ -409,6 +409,9 @@ def _run_trace( # embedded value instead of imposing the synthetic-workload default. if trace_format not in {"weka", "dynamo"} and trace_block_size is None: trace_block_size = self.trace_block_size + # Trace replays use ordinal request ids and the canonical router seed, + # so a candidate's score depends only on its configuration and the + # trace. if deployment.deployment_mode == "agg": return run_trace_replay( trace_files=trace_files, @@ -423,6 +426,7 @@ def _run_trace( execution_model=execution_model, extra_engine_args=self._engine_args(deployment, "aggregated"), num_workers=deployment.num_workers, + _deterministic=True, **common, ) return run_trace_replay( @@ -440,6 +444,7 @@ def _run_trace( decode_engine_args=self._engine_args(deployment, "decode"), num_prefill_workers=deployment.num_prefill_workers, num_decode_workers=deployment.num_decode_workers, + _deterministic=True, **common, ) diff --git a/components/src/dynamo/replay/tests/test_retune_replay.py b/components/src/dynamo/replay/tests/test_retune_replay.py index f902a8566173..791292404e9a 100644 --- a/components/src/dynamo/replay/tests/test_retune_replay.py +++ b/components/src/dynamo/replay/tests/test_retune_replay.py @@ -578,6 +578,38 @@ def test_retune_rejects_an_engine_shape_that_differs_from_the_planner_mode( ) +@pytest.mark.unit +def test_retune_rejects_the_dgdr_executor(tmp_path): + # Replay runs local sweeps only; executor 'dgdr' needs a live Planner. + preset = dict.fromkeys( + ("tensor", "pipeline", "attention_data", "moe_tensor", "moe_expert"), 1 + ) + worker = { + "parallelism": {"preset": [preset]}, + "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + } + engine = {"mode": "aggregated", "workers": {"aggregated": worker}} + dgdr = {"claim_name": "traces", "model": "m", "sku": "h200_sxm", "engine": engine} + config = PlannerConfig.model_validate( + _guard_config() + | { + "retune": { + "interval_seconds": 21600, + "timeout_seconds": 20400, + "router": "default", + "executor": "dgdr", + "dgdr": dgdr, + } + } + ) + trace = _write_trace( + tmp_path / "trace.jsonl", count=4, gap_ms=100, blocks=1, output_tokens=2 + ) + + with pytest.raises(ValueError, match="executor 'dgdr' is not supported in replay"): + api._retune_inputs(config, **_inputs_kwargs(trace)) + + @pytest.mark.unit def test_synthetic_replay_rejects_retune(): with pytest.raises( diff --git a/components/src/dynamo/replay/tests/test_simulation.py b/components/src/dynamo/replay/tests/test_simulation.py index 4c7c7b3b7cea..23725992c6d5 100644 --- a/components/src/dynamo/replay/tests/test_simulation.py +++ b/components/src/dynamo/replay/tests/test_simulation.py @@ -97,7 +97,35 @@ def _agg_deployment() -> BackendDeploymentSpec: ) -def test_trace_runner_preserves_current_replay_arguments(monkeypatch) -> None: +def _disagg_deployment() -> BackendDeploymentSpec: + return BackendDeploymentSpec( + deployment_mode="disagg", + backend="vllm", + backend_version="0.11.0", + prefill_engine_args={ + "worker_type": "prefill", + "timing_model": {"type": "polynomial"}, + }, + decode_engine_args={ + "worker_type": "decode", + "timing_model": {"type": "polynomial"}, + }, + num_prefill_workers=1, + num_decode_workers=2, + ) + + +@pytest.mark.parametrize( + ("deployment", "workers"), + [ + (_agg_deployment(), {"num_workers": 3}), + (_disagg_deployment(), {"num_prefill_workers": 1, "num_decode_workers": 2}), + ], + ids=["agg", "disagg"], +) +def test_trace_runner_preserves_current_replay_arguments( + monkeypatch, deployment, workers +) -> None: seen = {} def fake_run_trace_replay(**kwargs): @@ -113,7 +141,7 @@ def fake_run_trace_replay(**kwargs): monkeypatch.setattr(simulation, "KvRouterConfig", _FakeRouterConfig) monkeypatch.setattr(simulation, "run_trace_replay", fake_run_trace_replay) spec = ReplaySpec( - backend_deployment=_agg_deployment(), + backend_deployment=deployment, workload={ "trace_path": "tiny.jsonl", "trace_format": "dynamo", @@ -159,7 +187,8 @@ def fake_run_trace_replay(**kwargs): assert seen["trace_files"] == "tiny.jsonl" assert seen["trace_format"] == "dynamo" - assert seen["num_workers"] == 3 + assert {name: seen[name] for name in workers} == workers + assert seen["_deterministic"] is True assert seen["router_mode"] == "kv_router" assert seen["planner_config"] == {"mode": "agg"} assert seen["arrival_speedup_ratio"] == 2.0 diff --git a/deploy/helm/charts/platform/README.md b/deploy/helm/charts/platform/README.md index c7e1c680d98e..4cd9b75912d3 100644 --- a/deploy/helm/charts/platform/README.md +++ b/deploy/helm/charts/platform/README.md @@ -686,6 +686,7 @@ Kubernetes: `>=1.30.0-0` | dynamo-operator.upgradePreflight.checks.upgradePath.unsupportedUpgradeOverride | string | `""` | Target minor (N.M) allowed for an otherwise unsupported forward upgrade. The value is ignored for other target minors and never bypasses other checks. | | dynamo-operator.upgradePreflight.checks.upgradePath.unsupportedDowngradeOverride | string | `""` | Target minor (N.M) allowed for an otherwise unsupported downgrade. The value is ignored for other target minors and never bypasses other checks. | | dynamo-operator.dgdrV1Beta2.enabled | bool | `false` | EXPERIMENTAL, DEVELOPMENT CLUSTERS ONLY: Serve the unfinished DGDR v1beta2 Request API with conversion and install the Run and Candidate APIs. Run and Candidate objects are persisted in etcd; incompatible schema changes may make them unreadable or block future CRD upgrades. When true, the operator reconciles v1beta2 Requests into Runs and executes each Run as an AISimulate planner Job with a publisher sidecar that persists deployment Candidates. | +| dynamo-operator.dgdrV1Beta2.plannerLiveRetune | bool | `false` | EXPERIMENTAL, DEVELOPMENT CLUSTERS ONLY: With dgdrV1Beta2.enabled, grant every Planner ServiceAccount what Planner live retune (retune.executor: dgdr) needs: create, read and delete DGDR v1beta2 Requests, read their Runs and Candidates, and read, create, update and delete any Lease in the namespace. | | dynamo-operator.natsAddr | string | `""` | NATS server address for operator communication. When empty, the operator uses bundled NATS only if global.nats.install=true; otherwise NATS is not configured. Format: `nats://hostname:4222` | | dynamo-operator.etcdAddr | string | `""` | etcd server address for an external etcd instance. Only needed when using external etcd without the bundled subchart. Format: `http://hostname:2379` or `https://hostname:2379` | | dynamo-operator.modelExpressURL | string | `""` | URL for the Model Express server if not deployed by this helm chart. This is ignored if Model Express server is installed by this helm chart (global.model-express.enabled is true). | diff --git a/deploy/helm/charts/platform/components/operator/files/role.yaml b/deploy/helm/charts/platform/components/operator/files/role.yaml index f667b372eac5..97258d1d9b3f 100644 --- a/deploy/helm/charts/platform/components/operator/files/role.yaml +++ b/deploy/helm/charts/platform/components/operator/files/role.yaml @@ -37,6 +37,7 @@ rules: - delete - get - list + - patch - watch - apiGroups: - "" diff --git a/deploy/helm/charts/platform/components/operator/templates/planner.yaml b/deploy/helm/charts/platform/components/operator/templates/planner.yaml index 2c9ef2bb73d0..60a67ba8874c 100644 --- a/deploy/helm/charts/platform/components/operator/templates/planner.yaml +++ b/deploy/helm/charts/platform/components/operator/templates/planner.yaml @@ -57,6 +57,24 @@ rules: - apiGroups: [""] resources: ["pods"] verbs: ["list"] +{{- if and .Values.dgdrV1Beta2.enabled .Values.dgdrV1Beta2.plannerLiveRetune }} +# EXPERIMENTAL live retune (planner retune.executor: dgdr): sweep through DGDR +# v1beta2 Requests, read their Runs and list their Candidates, and hold the +# leadership and state Leases (read, create, replace and delete; never patch). +# Lease names are per DGD, so the Lease rule covers every Lease in the namespace. +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentrequests"] + verbs: ["create", "get", "list", "delete"] +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentruns"] + verbs: ["get"] +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentcandidates"] + verbs: ["list"] +- apiGroups: ["coordination.k8s.io"] + resources: ["leases"] + verbs: ["get", "create", "update", "delete"] +{{- end }} --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding @@ -78,6 +96,10 @@ roleRef: # Even though this is a ClusterRole, the operator binds it via a namespace-scoped # RoleBinding (internal/rbac/manager.go EnsureServiceAccountWithRBAC), so all # permissions — including pods/list — are effective only in the DGD's namespace. +# RBAC escalation prevention lets the operator bind only permissions it holds, +# so every rule here must also be in the operator's role (with +# dgdrV1Beta2.plannerLiveRetune, the DGDR v1beta2 Run and Candidate rules of +# the operator's v1beta2 integration). --- apiVersion: rbac.authorization.k8s.io/v1 kind: ClusterRole @@ -107,4 +129,22 @@ rules: - apiGroups: [""] resources: ["pods"] verbs: ["list"] +{{- if and .Values.dgdrV1Beta2.enabled .Values.dgdrV1Beta2.plannerLiveRetune }} +# EXPERIMENTAL live retune (planner retune.executor: dgdr): sweep through DGDR +# v1beta2 Requests, read their Runs and list their Candidates, and hold the +# leadership and state Leases (read, create, replace and delete; never patch). +# Lease names are per DGD, so the Lease rule covers every Lease in the namespace. +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentrequests"] + verbs: ["create", "get", "list", "delete"] +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentruns"] + verbs: ["get"] +- apiGroups: ["nvidia.com"] + resources: ["dynamographdeploymentcandidates"] + verbs: ["list"] +- apiGroups: ["coordination.k8s.io"] + resources: ["leases"] + verbs: ["get", "create", "update", "delete"] +{{- end }} {{- end }} diff --git a/deploy/helm/charts/platform/components/operator/values.yaml b/deploy/helm/charts/platform/components/operator/values.yaml index 269173f55baa..c8f6056082f7 100644 --- a/deploy/helm/charts/platform/components/operator/values.yaml +++ b/deploy/helm/charts/platform/components/operator/values.yaml @@ -40,6 +40,11 @@ dgdrV1Beta2: # Requests into Runs and executes each Run as an AISimulate planner Job with # a publisher sidecar that persists deployment Candidates. enabled: false + # -- EXPERIMENTAL, DEVELOPMENT CLUSTERS ONLY: With dgdrV1Beta2.enabled, grant every + # Planner ServiceAccount what Planner live retune (retune.executor: dgdr) needs: + # create, read and delete DGDR v1beta2 Requests, read their Runs and Candidates, + # and read, create, update and delete any Lease in the namespace. + plannerLiveRetune: false # Environment variables to pass to operator Deployment. env: [] diff --git a/deploy/helm/charts/platform/tests/planner_rbac_test.yaml b/deploy/helm/charts/platform/tests/planner_rbac_test.yaml index 389d2332c571..c1ba096a9eb4 100644 --- a/deploy/helm/charts/platform/tests/planner_rbac_test.yaml +++ b/deploy/helm/charts/platform/tests/planner_rbac_test.yaml @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -suite: planner RBAC — Pod lifecycle observation +suite: planner RBAC — Pod lifecycle observation and live retune templates: - charts/dynamo-operator/templates/planner.yaml @@ -114,3 +114,157 @@ tests: - get - patch documentIndex: 1 + + # ── EXPERIMENTAL live retune (dgdrV1Beta2.plannerLiveRetune) ──────────────── + # The base role has 6 rules; live retune adds exactly 4. Counting the rules + # also catches an extra verb or a split rule. + + - it: omits live retune rules from the ClusterRole by default + asserts: + - lengthEqual: + path: rules + count: 6 + documentIndex: 0 + + - it: omits live retune rules from the ClusterRole with dgdrV1Beta2 alone + set: + dynamo-operator.dgdrV1Beta2.enabled: true + asserts: + - lengthEqual: + path: rules + count: 6 + documentIndex: 0 + + - it: omits live retune rules from the ClusterRole with plannerLiveRetune alone + set: + dynamo-operator.dgdrV1Beta2.plannerLiveRetune: true + asserts: + - lengthEqual: + path: rules + count: 6 + documentIndex: 0 + + - it: grants DGDR v1beta2 and Lease access in the ClusterRole with plannerLiveRetune + set: + dynamo-operator.dgdrV1Beta2.enabled: true + dynamo-operator.dgdrV1Beta2.plannerLiveRetune: true + asserts: + - lengthEqual: + path: rules + count: 10 + documentIndex: 0 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentrequests + verbs: + - create + - get + - list + - delete + documentIndex: 0 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentruns + verbs: + - get + documentIndex: 0 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentcandidates + verbs: + - list + documentIndex: 0 + - contains: + path: rules + content: + apiGroups: + - coordination.k8s.io + resources: + - leases + verbs: + - get + - create + - update + - delete + documentIndex: 0 + + - it: omits live retune rules from the Role with dgdrV1Beta2 alone in namespace-restricted mode + set: + dynamo-operator.namespaceRestriction.enabled: true + dynamo-operator.upgradeCRD: false + dynamo-operator.dgdrV1Beta2.enabled: true + asserts: + - lengthEqual: + path: rules + count: 6 + documentIndex: 1 + + - it: grants DGDR v1beta2 and Lease access in the Role with plannerLiveRetune in namespace-restricted mode + set: + dynamo-operator.namespaceRestriction.enabled: true + dynamo-operator.upgradeCRD: false + dynamo-operator.dgdrV1Beta2.enabled: true + dynamo-operator.dgdrV1Beta2.plannerLiveRetune: true + asserts: + - lengthEqual: + path: rules + count: 10 + documentIndex: 1 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentrequests + verbs: + - create + - get + - list + - delete + documentIndex: 1 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentruns + verbs: + - get + documentIndex: 1 + - contains: + path: rules + content: + apiGroups: + - nvidia.com + resources: + - dynamographdeploymentcandidates + verbs: + - list + documentIndex: 1 + - contains: + path: rules + content: + apiGroups: + - coordination.k8s.io + resources: + - leases + verbs: + - get + - create + - update + - delete + documentIndex: 1 diff --git a/deploy/helm/charts/platform/values.yaml b/deploy/helm/charts/platform/values.yaml index cd2913a97994..b6c2d8b4a1aa 100644 --- a/deploy/helm/charts/platform/values.yaml +++ b/deploy/helm/charts/platform/values.yaml @@ -131,6 +131,11 @@ dynamo-operator: # reconciles v1beta2 Requests into Runs and executes each Run as an AISimulate # planner Job with a publisher sidecar that persists deployment Candidates. enabled: false + # -- EXPERIMENTAL, DEVELOPMENT CLUSTERS ONLY: With dgdrV1Beta2.enabled, grant every + # Planner ServiceAccount what Planner live retune (retune.executor: dgdr) needs: + # create, read and delete DGDR v1beta2 Requests, read their Runs and Candidates, + # and read, create, update and delete any Lease in the namespace. + plannerLiveRetune: false # Environment variables to pass to operator Deployment. env: [] diff --git a/deploy/operator/.golangci-kal.yml b/deploy/operator/.golangci-kal.yml new file mode 100644 index 000000000000..c197eb455f36 --- /dev/null +++ b/deploy/operator/.golangci-kal.yml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kubernetes API conventions for the v1beta2 API types, checked by kube-api-linter +# (sigs.k8s.io/kube-api-linter). It runs through the golangci-lint build that the +# kube-api-linter module ships, so `make lint-api` uses this configuration alone. + +version: "2" + +run: + timeout: 10m + +linters: + default: none + enable: + - kubeapilinter + settings: + custom: + kubeapilinter: + type: module + description: kube-api-linter checks Kubernetes API conventions. + settings: + linters: + disable: + - "*" + enable: + # Every serialized field carries exactly one of +optional or +required. + - optionalorrequired + exclusions: + generated: strict + rules: + - linters: + - kubeapilinter + path-except: ^api/v1beta2/ + +issues: + max-issues-per-linter: 0 + max-same-issues: 0 diff --git a/deploy/operator/Dockerfile b/deploy/operator/Dockerfile index b903b4c20838..419dac24b0fe 100644 --- a/deploy/operator/Dockerfile +++ b/deploy/operator/Dockerfile @@ -35,6 +35,11 @@ RUN go install github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.15.0 # Run linting RUN golangci-lint run --timeout=5m +# Check the v1beta2 API types against Kubernetes API conventions. Keep the version in sync with +# KUBE_API_LINTER_VERSION in the Makefile. +RUN go install sigs.k8s.io/kube-api-linter/cmd/golangci-lint-kube-api-linter@v0.0.0-20261012124120-79b557aa8c8d && \ + golangci-lint-kube-api-linter run --timeout=5m -c .golangci-kal.yml ./api/v1beta2/... + # Test stage FROM base AS tester diff --git a/deploy/operator/Makefile b/deploy/operator/Makefile index 7af4283fdb08..01256da0747b 100644 --- a/deploy/operator/Makefile +++ b/deploy/operator/Makefile @@ -385,13 +385,17 @@ integration: envtest setup-clustertest ## Run envtests and clustertests, creatin $(MAKE) clustertest .PHONY: lint -lint: golangci-lint ## Run golangci-lint linter & yamllint +lint: golangci-lint lint-api ## Run golangci-lint linter, the Kubernetes API linter & yamllint $(GOLANGCI_LINT) run .PHONY: lint-fix lint-fix: golangci-lint ## Run golangci-lint linter and perform fixes $(GOLANGCI_LINT) run --fix +.PHONY: lint-api +lint-api: kube-api-linter ## Check the v1beta2 API types against Kubernetes API conventions + $(KUBE_API_LINTER) run -c .golangci-kal.yml ./api/v1beta2/... + ##@ Build .PHONY: build @@ -579,6 +583,8 @@ YQ_VERSION ?= v4.45.4 CONTROLLER_TOOLS_VERSION ?= v0.17.3 ENVTEST_VERSION ?= release-0.19 GOLANGCI_LINT_VERSION ?= v2.15.0 +# kube-api-linter publishes no tags; this pseudo-version is its 2026-10-12 main. +KUBE_API_LINTER_VERSION ?= v0.0.0-20261012124120-79b557aa8c8d ## Tool Binaries KUBECTL ?= $(LOCALBIN)/kubectl-$(KUBECTL_VERSION) @@ -588,6 +594,7 @@ YQ ?= $(LOCALBIN)/yq-$(YQ_VERSION) CONTROLLER_GEN ?= $(LOCALBIN)/controller-gen-$(CONTROLLER_TOOLS_VERSION) ENVTEST ?= $(LOCALBIN)/setup-envtest-$(ENVTEST_VERSION) GOLANGCI_LINT = $(LOCALBIN)/golangci-lint-$(GOLANGCI_LINT_VERSION) +KUBE_API_LINTER = $(LOCALBIN)/golangci-lint-kube-api-linter-$(KUBE_API_LINTER_VERSION) ## ---------- Binary download targets ---------- @@ -651,6 +658,11 @@ golangci-lint: $(GOLANGCI_LINT) ## Download golangci-lint locally if necessary. $(GOLANGCI_LINT): $(LOCALBIN) $(call go-install-tool,$(GOLANGCI_LINT),github.com/golangci/golangci-lint/v2/cmd/golangci-lint,${GOLANGCI_LINT_VERSION}) +.PHONY: kube-api-linter +kube-api-linter: $(KUBE_API_LINTER) ## Build the golangci-lint that includes kube-api-linter locally if necessary. +$(KUBE_API_LINTER): $(LOCALBIN) + $(call go-install-tool,$(KUBE_API_LINTER),sigs.k8s.io/kube-api-linter/cmd/golangci-lint-kube-api-linter,${KUBE_API_LINTER_VERSION}) + # go-install-tool will 'go install' any package with custom target and name of binary, if it doesn't exist # $1 - target path with name of binary (ideally with version) # $2 - package url which can be installed diff --git a/deploy/operator/api/roundtrip_fuzz_test.go b/deploy/operator/api/roundtrip_fuzz_test.go index 6fc6158d09eb..9672d002b73b 100644 --- a/deploy/operator/api/roundtrip_fuzz_test.go +++ b/deploy/operator/api/roundtrip_fuzz_test.go @@ -563,7 +563,31 @@ func fuzzBeta2Hardware(hardware *v1beta2.HardwareSpec, c randfill.Continue) { } func fuzzBeta2Workload(workload *v1beta2.WorkloadSpec, c randfill.Continue) { - if workload == nil || workload.Static == nil { + if workload == nil { + return + } + + // A trace is exclusive with static and names one admitted claim-relative file. When both are + // filled, keep one of them at random so static and trace are equally likely. + if workload.Trace != nil && workload.Static != nil { + if c.Bool() { + workload.Static = nil + } else { + workload.Trace = nil + } + } + if workload.Trace != nil { + workload.Trace.Format = v1beta2.TraceFormatDynamo + workload.Trace.Source.PersistentVolumeClaim = &v1beta2.TracePVCSource{ + ClaimName: fmt.Sprintf("traces-%d", c.Uint32()), + SubPath: fmt.Sprintf("retune/window-%d.jsonl.gz", c.Uint32()), + } + if workload.Trace.Replay != nil && workload.Trace.Replay.ArrivalSpeedupRatio != nil { + workload.Trace.Replay.ArrivalSpeedupRatio = oneOfPtr(c, 0.5, 1.0, 4.0) + } + return + } + if workload.Static == nil { return } diff --git a/deploy/operator/api/v1beta2/dynamographdeploymentcandidate_types.go b/deploy/operator/api/v1beta2/dynamographdeploymentcandidate_types.go index 9dd1fe007845..62c083350a11 100644 --- a/deploy/operator/api/v1beta2/dynamographdeploymentcandidate_types.go +++ b/deploy/operator/api/v1beta2/dynamographdeploymentcandidate_types.go @@ -40,6 +40,7 @@ type DynamoGraphDeploymentCandidateSpec struct { // Parameters are the resolved evaluation parameters for this point. // +kubebuilder:pruning:PreserveUnknownFields // +kubebuilder:validation:Type=object + // +required Parameters runtime.RawExtension `json:"parameters"` } @@ -52,10 +53,14 @@ type DynamoGraphDeploymentCandidateSpec struct { // DynamoGraphDeploymentCandidate is one bounded, user-visible search result. type DynamoGraphDeploymentCandidate struct { - metav1.TypeMeta `json:",inline"` + metav1.TypeMeta `json:",inline"` + // +optional metav1.ObjectMeta `json:"metadata,omitempty"` - Spec DynamoGraphDeploymentCandidateSpec `json:"spec"` + // +required + Spec DynamoGraphDeploymentCandidateSpec `json:"spec"` + + // +optional Status DynamoGraphDeploymentCandidateStatus `json:"status,omitempty"` } diff --git a/deploy/operator/api/v1beta2/dynamographdeploymentrequest_types.go b/deploy/operator/api/v1beta2/dynamographdeploymentrequest_types.go index 25536741ca97..2685f6b3e65a 100644 --- a/deploy/operator/api/v1beta2/dynamographdeploymentrequest_types.go +++ b/deploy/operator/api/v1beta2/dynamographdeploymentrequest_types.go @@ -28,6 +28,7 @@ import ( type ModelSpec struct { // Name identifies the model in the syntax accepted by the selected backend. // +kubebuilder:validation:MinLength=1 + // +required Name string `json:"name"` } @@ -91,12 +92,14 @@ type KVCapacityFractionRange struct { // +kubebuilder:validation:ExclusiveMinimum=true // +kubebuilder:validation:Minimum=0 // +kubebuilder:validation:Maximum=1 + // +required Min float64 `json:"min"` // Max is the upper end of the searched fraction range. // +kubebuilder:validation:ExclusiveMinimum=true // +kubebuilder:validation:Minimum=0 // +kubebuilder:validation:Maximum=1 + // +required Max float64 `json:"max"` } @@ -154,13 +157,99 @@ type StaticWorkloadSpec struct { Load *StaticLoadSpec `json:"load,omitempty"` } -// WorkloadSpec defines the synthetic traffic model used for every candidate evaluation. +// TraceFormat identifies the decoder of a recorded request trace. Lowering translates it to +// AISimulate's native traffic.source.format. +// +kubebuilder:validation:Enum=Dynamo +type TraceFormat string + +const ( + // TraceFormatDynamo is the Dynamo request trace (dynamo.request.trace.v1) that the frontend + // records, as JSON Lines or gzip-compressed JSON Lines. It lowers to AISimulate's dynamo format. + TraceFormatDynamo TraceFormat = "Dynamo" +) + +// TracePVCSource names one trace file on a PersistentVolumeClaim in the Request namespace. +type TracePVCSource struct { + // ClaimName is the PersistentVolumeClaim, in the Request namespace, that holds the trace. + // The Run Job mounts only the trace file from it, read-only, and only into the AISimulate + // container. Each Run protects the claim with the finalizer `nvidia.com/dgdr-run-` + // until the Run completes or, for a Run deleted before it completes, until no Job of the Run + // is left. Garbage collection deletes the Job of a deleted Run before the Job's pod, so the + // finalizer of a deleted Run can be released while that pod still stops; Kubernetes + // PersistentVolumeClaim protection keeps the claim until no pod uses it. Deleting the claim + // cancels the Runs that read it. + // +required + // +kubebuilder:validation:MinLength=1 + // +kubebuilder:validation:MaxLength=253 + // +kubebuilder:validation:Pattern=`^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$` + ClaimName string `json:"claimName"` + + // SubPath names one trace file relative to the root of the claim; the Run Job mounts it as + // the subPath of the claim volume. It is neither a directory nor a glob. It must not start or + // end with '/' and must not contain empty, '.' or '..' segments or a NUL character, which no + // file name can hold. The file must be complete before the Request is created and must never + // be rewritten: every Run of the Request reads it when its Job starts. + // +required + // +kubebuilder:validation:MinLength=1 + // +kubebuilder:validation:MaxLength=1024 + // +kubebuilder:validation:XValidation:rule="!self.startsWith('/') && !self.endsWith('/') && !self.contains('//') && !self.matches('(^|/)[.][.]?(/|$)')",message="subPath must be relative to the claim root without empty, '.' or '..' segments" + // +kubebuilder:validation:XValidation:rule="!self.contains('\\u0000')",message="subPath must not contain a NUL character" + SubPath string `json:"subPath"` +} + +// TraceSourceSpec identifies exactly one trace artifact. A PersistentVolumeClaim is the only +// source in this release. +// +kubebuilder:validation:XValidation:rule="has(self.persistentVolumeClaim)",message="persistentVolumeClaim is required" +type TraceSourceSpec struct { + // PersistentVolumeClaim reads the trace from a PersistentVolumeClaim in the Request + // namespace. + // +optional + PersistentVolumeClaim *TracePVCSource `json:"persistentVolumeClaim,omitempty"` +} + +// TraceReplaySpec configures how AISimulate replays a decoded trace. +type TraceReplaySpec struct { + // ArrivalSpeedupRatio multiplies the recorded request-arrival rate. 1.0 preserves it. + // When omitted, the recorded arrival times are replayed unchanged. + // +optional + // +kubebuilder:validation:ExclusiveMinimum=true + // +kubebuilder:validation:Minimum=0 + ArrivalSpeedupRatio *float64 `json:"arrivalSpeedupRatio,omitempty"` +} + +// TraceWorkloadSpec replays a recorded request trace open loop at its recorded arrival times. +type TraceWorkloadSpec struct { + // Format identifies the trace decoder. Defaults to Dynamo. + // +optional + // +kubebuilder:default=Dynamo + Format TraceFormat `json:"format,omitempty"` + + // Source identifies the trace artifact. + // +required + Source TraceSourceSpec `json:"source"` + + // Replay configures how the decoded trace is replayed. + // +optional + Replay *TraceReplaySpec `json:"replay,omitempty"` +} + +// WorkloadSpec defines the traffic model used for every candidate evaluation. static and +// trace are mutually exclusive. +// +kubebuilder:validation:XValidation:rule="!(has(self.static) && has(self.trace))",message="static and trace are mutually exclusive" type WorkloadSpec struct { - // Static defines AISimulate synthetic traffic. When omitted, admission supplies the - // complete static default. An explicit empty object opts out; any nonempty object is - // completed with omitted inputTokens, outputTokens and load defaults. + // Static defines AISimulate synthetic traffic. When omitted without a trace, admission + // supplies the complete static default. An explicit empty object opts out; any nonempty + // object is completed with omitted inputTokens, outputTokens and load defaults. // +optional Static *StaticWorkloadSpec `json:"static,omitempty"` + + // Trace replays a recorded request trace from a PersistentVolumeClaim in the Request + // namespace. Admission supplies no static default when it is set. The Run Job mounts only the + // trace file, read-only through a subPath of the claim, at /var/run/aisimulate/trace/trace.jsonl, + // or trace.jsonl.gz when subPath ends in .gz. The trace lowers to the default traffic.source + // and traffic.load, which spec.aisimulate.inputOverrides can replace or delete. + // +optional + Trace *TraceWorkloadSpec `json:"trace,omitempty"` } // ObjectiveMode selects the optimization mode. Pareto is reserved for a future @@ -239,9 +328,18 @@ type AISimulateSpec struct { // runtime_image, runtime_version_override, derive_runtime_images, // num_gpus_per_node, and namespace are operator-owned. Request admission // rejects them, and the Run controller enforces its values while lowering the - // immutable snapshot. The operator runs the AISimulate version from its - // configured planner image; Requests and Runs cannot select a different image - // or version. + // immutable snapshot. A typed spec.workload.trace is an underlay like the other + // typed fields: it supplies the default traffic.source (type trace, format + // dynamo and the mounted trace file in paths) and the recorded-arrival + // traffic.load. Overrides can replace those values, including + // traffic.source.type, format, and paths, and null deletes the named field or + // parent object; the Run controller does not restore it. The typed + // PersistentVolumeClaim source alone decides the Run Job's trace mount and its + // claim dependency, so overriding a native path adds no volume mount. + // AISimulate validates the resulting input when the Run's Job starts; admission + // does not validate it against the AISimulate schema yet. The operator runs + // the AISimulate version from its configured planner image; Requests and Runs + // cannot select a different image or version. // +optional // +kubebuilder:pruning:PreserveUnknownFields // +kubebuilder:validation:Type=object @@ -279,12 +377,14 @@ type NativeSearchInput struct { // +optional Hardware *HardwareSpec `json:"hardware,omitempty"` - // Workload defines the synthetic traffic model used for every candidate evaluation. + // Workload defines the traffic model used for every candidate evaluation: synthetic static + // traffic or a recorded trace. // When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; // an explicit empty workload receives the same shape. workload.static: {} opts out of all - // traffic defaults. Any other static object is completed with its missing defaults. + // traffic defaults. Any other static object is completed with its missing defaults, and a + // trace workload receives no static default. // spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over - // every lowered traffic field it names. + // every lowered traffic field it names, including the fields a trace lowers. // The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction // search. Lowering independently defaults optimizer.max_trials to 320. // +optional @@ -374,10 +474,14 @@ type DynamoGraphDeploymentRequestStatus struct { // DynamoGraphDeploymentRequest is the Schema for simulation-backed deployment searches. type DynamoGraphDeploymentRequest struct { - metav1.TypeMeta `json:",inline"` + metav1.TypeMeta `json:",inline"` + // +optional metav1.ObjectMeta `json:"metadata,omitempty"` - Spec DynamoGraphDeploymentRequestSpec `json:"spec"` + // +required + Spec DynamoGraphDeploymentRequestSpec `json:"spec"` + + // +optional Status DynamoGraphDeploymentRequestStatus `json:"status,omitempty"` } diff --git a/deploy/operator/api/v1beta2/dynamographdeploymentrun_types.go b/deploy/operator/api/v1beta2/dynamographdeploymentrun_types.go index affb00815de3..fb2cf2dde1bf 100644 --- a/deploy/operator/api/v1beta2/dynamographdeploymentrun_types.go +++ b/deploy/operator/api/v1beta2/dynamographdeploymentrun_types.go @@ -27,7 +27,7 @@ import ( // DynamoGraphDeploymentRunSpec is an immutable snapshot of an accepted DGDR spec. // Request admission resolves its structural defaults before the Request controller copies it; // Request admission defaulting does not apply to Runs. -// +kubebuilder:validation:XValidation:rule="has(self.model) && has(self.backends) && size(self.backends) > 0 && has(self.hardware) && ((!has(self.hardware.gpu) && !has(self.hardware.numGpusPerNode)) || (has(self.hardware.gpu) && has(self.hardware.gpu.sku) && has(self.hardware.gpu.budget) && has(self.hardware.numGpusPerNode))) && has(self.workload) && has(self.workload.static) && ((!has(self.workload.static.inputTokens) && !has(self.workload.static.outputTokens) && !has(self.workload.static.load)) || (has(self.workload.static.inputTokens) && has(self.workload.static.outputTokens) && has(self.workload.static.load))) && has(self.objective) && has(self.objective.mode) && has(self.objective.metric) && has(self.objective.sla) && has(self.recommendation) && has(self.recommendation.maxCandidates) && self.recommendation.maxCandidates > 0",message="spec must be a fully resolved Request admission snapshot" +// +kubebuilder:validation:XValidation:rule="has(self.model) && has(self.backends) && size(self.backends) > 0 && has(self.hardware) && ((!has(self.hardware.gpu) && !has(self.hardware.numGpusPerNode)) || (has(self.hardware.gpu) && has(self.hardware.gpu.sku) && has(self.hardware.gpu.budget) && has(self.hardware.numGpusPerNode))) && has(self.workload) && (has(self.workload.trace) || (has(self.workload.static) && ((!has(self.workload.static.inputTokens) && !has(self.workload.static.outputTokens) && !has(self.workload.static.load)) || (has(self.workload.static.inputTokens) && has(self.workload.static.outputTokens) && has(self.workload.static.load))))) && has(self.objective) && has(self.objective.mode) && has(self.objective.metric) && has(self.objective.sla) && has(self.recommendation) && has(self.recommendation.maxCandidates) && self.recommendation.maxCandidates > 0",message="spec must be a fully resolved Request admission snapshot" // // The Run reconciler lowers this immutable typed snapshot plus inputOverrides exactly once and // records the resulting public recommendation input in status.aisimulate.input. @@ -49,6 +49,7 @@ const ( // ResourceReference identifies a namespaced resource by name and UID. type ResourceReference struct { // Name is the resource name. + // +required Name string `json:"name"` // UID is the resource UID. @@ -59,51 +60,64 @@ type ResourceReference struct { // RoundProgress reports completed and planned optimizer rounds. type RoundProgress struct { // Completed is the number of finished rounds. + // +required Completed int32 `json:"completed"` // Total is the number of planned rounds. + // +required Total int32 `json:"total"` } // BranchProgress reports progress for one independently optimized AISimulate branch. type BranchProgress struct { // Name is the executor-defined stable branch identifier. + // +required Name string `json:"name"` // Rounds reports optimizer progress for this branch. + // +required Rounds RoundProgress `json:"rounds"` } // EvaluationProgress reports unique candidate outcomes without creating one resource per trial. type EvaluationProgress struct { // Scheduled is the number of unique evaluations accepted for execution. + // +optional Scheduled int64 `json:"scheduled,omitempty"` // Running is the number of evaluations currently executing. + // +optional Running int64 `json:"running,omitempty"` // Feasible is the number of completed evaluations that satisfy all constraints. + // +optional Feasible int64 `json:"feasible,omitempty"` // Infeasible is the number of completed evaluations that violate at least one constraint. + // +optional Infeasible int64 `json:"infeasible,omitempty"` // Failed is the number of evaluations that ended in an execution error. + // +optional Failed int64 `json:"failed,omitempty"` // Unsupported is the number of configurations rejected as unsupported before evaluation. + // +optional Unsupported int64 `json:"unsupported,omitempty"` // CacheHits is the number of evaluations satisfied from a prior result. + // +optional CacheHits int64 `json:"cacheHits,omitempty"` } // CandidateProgress reports bounded Kubernetes publication, not all internal trials. type CandidateProgress struct { // ParetoFront is the number of feasible points on the current Pareto front. + // +optional ParetoFront int64 `json:"paretoFront,omitempty"` // Published is the number of Candidate resources currently referenced by this run. + // +optional Published int32 `json:"published,omitempty"` } @@ -189,10 +203,14 @@ type DynamoGraphDeploymentRunStatus struct { // DynamoGraphDeploymentRun is one immutable accepted DGDR request and its execution status. type DynamoGraphDeploymentRun struct { - metav1.TypeMeta `json:",inline"` + metav1.TypeMeta `json:",inline"` + // +optional metav1.ObjectMeta `json:"metadata,omitempty"` - Spec DynamoGraphDeploymentRunSpec `json:"spec"` + // +required + Spec DynamoGraphDeploymentRunSpec `json:"spec"` + + // +optional Status DynamoGraphDeploymentRunStatus `json:"status,omitempty"` } diff --git a/deploy/operator/api/v1beta2/zz_generated.deepcopy.go b/deploy/operator/api/v1beta2/zz_generated.deepcopy.go index 899c5121053e..8cf585efa0eb 100644 --- a/deploy/operator/api/v1beta2/zz_generated.deepcopy.go +++ b/deploy/operator/api/v1beta2/zz_generated.deepcopy.go @@ -854,6 +854,82 @@ func (in *StaticWorkloadSpec) DeepCopy() *StaticWorkloadSpec { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TracePVCSource) DeepCopyInto(out *TracePVCSource) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TracePVCSource. +func (in *TracePVCSource) DeepCopy() *TracePVCSource { + if in == nil { + return nil + } + out := new(TracePVCSource) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TraceReplaySpec) DeepCopyInto(out *TraceReplaySpec) { + *out = *in + if in.ArrivalSpeedupRatio != nil { + in, out := &in.ArrivalSpeedupRatio, &out.ArrivalSpeedupRatio + *out = new(float64) + **out = **in + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TraceReplaySpec. +func (in *TraceReplaySpec) DeepCopy() *TraceReplaySpec { + if in == nil { + return nil + } + out := new(TraceReplaySpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TraceSourceSpec) DeepCopyInto(out *TraceSourceSpec) { + *out = *in + if in.PersistentVolumeClaim != nil { + in, out := &in.PersistentVolumeClaim, &out.PersistentVolumeClaim + *out = new(TracePVCSource) + **out = **in + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TraceSourceSpec. +func (in *TraceSourceSpec) DeepCopy() *TraceSourceSpec { + if in == nil { + return nil + } + out := new(TraceSourceSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *TraceWorkloadSpec) DeepCopyInto(out *TraceWorkloadSpec) { + *out = *in + in.Source.DeepCopyInto(&out.Source) + if in.Replay != nil { + in, out := &in.Replay, &out.Replay + *out = new(TraceReplaySpec) + (*in).DeepCopyInto(*out) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new TraceWorkloadSpec. +func (in *TraceWorkloadSpec) DeepCopy() *TraceWorkloadSpec { + if in == nil { + return nil + } + out := new(TraceWorkloadSpec) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *WorkloadSpec) DeepCopyInto(out *WorkloadSpec) { *out = *in @@ -862,6 +938,11 @@ func (in *WorkloadSpec) DeepCopyInto(out *WorkloadSpec) { *out = new(StaticWorkloadSpec) (*in).DeepCopyInto(*out) } + if in.Trace != nil { + in, out := &in.Trace, &out.Trace + *out = new(TraceWorkloadSpec) + (*in).DeepCopyInto(*out) + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new WorkloadSpec. diff --git a/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentrequests.yaml b/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentrequests.yaml index 0be519ccd537..6fa1c6d8edd5 100644 --- a/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentrequests.yaml +++ b/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentrequests.yaml @@ -9617,9 +9617,18 @@ spec: runtime_image, runtime_version_override, derive_runtime_images, num_gpus_per_node, and namespace are operator-owned. Request admission rejects them, and the Run controller enforces its values while lowering the - immutable snapshot. The operator runs the AISimulate version from its - configured planner image; Requests and Runs cannot select a different image - or version. + immutable snapshot. A typed spec.workload.trace is an underlay like the other + typed fields: it supplies the default traffic.source (type trace, format + dynamo and the mounted trace file in paths) and the recorded-arrival + traffic.load. Overrides can replace those values, including + traffic.source.type, format, and paths, and null deletes the named field or + parent object; the Run controller does not restore it. The typed + PersistentVolumeClaim source alone decides the Run Job's trace mount and its + claim dependency, so overriding a native path adds no volume mount. + AISimulate validates the resulting input when the Run's Job starts; admission + does not validate it against the AISimulate schema yet. The operator runs + the AISimulate version from its configured planner image; Requests and Runs + cannot select a different image or version. type: object x-kubernetes-preserve-unknown-fields: true type: object @@ -9787,20 +9796,22 @@ spec: x-kubernetes-preserve-unknown-fields: true workload: description: |- - Workload defines the synthetic traffic model used for every candidate evaluation. + Workload defines the traffic model used for every candidate evaluation: synthetic static + traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: {} opts out of all - traffic defaults. Any other static object is completed with its missing defaults. + traffic defaults. Any other static object is completed with its missing defaults, and a + trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over - every lowered traffic field it names. + every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. properties: static: description: |- - Static defines AISimulate synthetic traffic. When omitted, admission supplies the - complete static default. An explicit empty object opts out; any nonempty object is - completed with omitted inputTokens, outputTokens and load defaults. + Static defines AISimulate synthetic traffic. When omitted without a trace, admission + supplies the complete static default. An explicit empty object opts out; any nonempty + object is completed with omitted inputTokens, outputTokens and load defaults. properties: inputTokens: description: InputTokens is the synthetic request input length in tokens. @@ -9869,7 +9880,83 @@ spec: minimum: 1 type: integer type: object + trace: + description: |- + Trace replays a recorded request trace from a PersistentVolumeClaim in the Request + namespace. Admission supplies no static default when it is set. The Run Job mounts only the + trace file, read-only through a subPath of the claim, at /var/run/aisimulate/trace/trace.jsonl, + or trace.jsonl.gz when subPath ends in .gz. The trace lowers to the default traffic.source + and traffic.load, which spec.aisimulate.inputOverrides can replace or delete. + properties: + format: + default: Dynamo + description: Format identifies the trace decoder. Defaults to Dynamo. + enum: + - Dynamo + type: string + replay: + description: Replay configures how the decoded trace is replayed. + properties: + arrivalSpeedupRatio: + description: |- + ArrivalSpeedupRatio multiplies the recorded request-arrival rate. 1.0 preserves it. + When omitted, the recorded arrival times are replayed unchanged. + exclusiveMinimum: true + minimum: 0 + type: number + type: object + source: + description: Source identifies the trace artifact. + properties: + persistentVolumeClaim: + description: |- + PersistentVolumeClaim reads the trace from a PersistentVolumeClaim in the Request + namespace. + properties: + claimName: + description: |- + ClaimName is the PersistentVolumeClaim, in the Request namespace, that holds the trace. + The Run Job mounts only the trace file from it, read-only, and only into the AISimulate + container. Each Run protects the claim with the finalizer `nvidia.com/dgdr-run-` + until the Run completes or, for a Run deleted before it completes, until no Job of the Run + is left. Garbage collection deletes the Job of a deleted Run before the Job's pod, so the + finalizer of a deleted Run can be released while that pod still stops; Kubernetes + PersistentVolumeClaim protection keeps the claim until no pod uses it. Deleting the claim + cancels the Runs that read it. + maxLength: 253 + minLength: 1 + pattern: ^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$ + type: string + subPath: + description: |- + SubPath names one trace file relative to the root of the claim; the Run Job mounts it as + the subPath of the claim volume. It is neither a directory nor a glob. It must not start or + end with '/' and must not contain empty, '.' or '..' segments or a NUL character, which no + file name can hold. The file must be complete before the Request is created and must never + be rewritten: every Run of the Request reads it when its Job starts. + maxLength: 1024 + minLength: 1 + type: string + x-kubernetes-validations: + - message: subPath must be relative to the claim root without empty, '.' or '..' segments + rule: '!self.startsWith(''/'') && !self.endsWith(''/'') && !self.contains(''//'') && !self.matches(''(^|/)[.][.]?(/|$)'')' + - message: subPath must not contain a NUL character + rule: '!self.contains(''\u0000'')' + required: + - claimName + - subPath + type: object + type: object + x-kubernetes-validations: + - message: persistentVolumeClaim is required + rule: has(self.persistentVolumeClaim) + required: + - source + type: object type: object + x-kubernetes-validations: + - message: static and trace are mutually exclusive + rule: '!(has(self.static) && has(self.trace))' type: object x-kubernetes-validations: - message: exactly one of native v1beta2 fields or v1beta1 must be specified diff --git a/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentruns.yaml b/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentruns.yaml index 5110346080b9..7c4942a688a6 100644 --- a/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentruns.yaml +++ b/deploy/operator/config/crd/bases/nvidia.com_dynamographdeploymentruns.yaml @@ -77,9 +77,18 @@ spec: runtime_image, runtime_version_override, derive_runtime_images, num_gpus_per_node, and namespace are operator-owned. Request admission rejects them, and the Run controller enforces its values while lowering the - immutable snapshot. The operator runs the AISimulate version from its - configured planner image; Requests and Runs cannot select a different image - or version. + immutable snapshot. A typed spec.workload.trace is an underlay like the other + typed fields: it supplies the default traffic.source (type trace, format + dynamo and the mounted trace file in paths) and the recorded-arrival + traffic.load. Overrides can replace those values, including + traffic.source.type, format, and paths, and null deletes the named field or + parent object; the Run controller does not restore it. The typed + PersistentVolumeClaim source alone decides the Run Job's trace mount and its + claim dependency, so overriding a native path adds no volume mount. + AISimulate validates the resulting input when the Run's Job starts; admission + does not validate it against the AISimulate schema yet. The operator runs + the AISimulate version from its configured planner image; Requests and Runs + cannot select a different image or version. type: object x-kubernetes-preserve-unknown-fields: true type: object @@ -239,20 +248,22 @@ spec: type: object workload: description: |- - Workload defines the synthetic traffic model used for every candidate evaluation. + Workload defines the traffic model used for every candidate evaluation: synthetic static + traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: {} opts out of all - traffic defaults. Any other static object is completed with its missing defaults. + traffic defaults. Any other static object is completed with its missing defaults, and a + trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over - every lowered traffic field it names. + every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. properties: static: description: |- - Static defines AISimulate synthetic traffic. When omitted, admission supplies the - complete static default. An explicit empty object opts out; any nonempty object is - completed with omitted inputTokens, outputTokens and load defaults. + Static defines AISimulate synthetic traffic. When omitted without a trace, admission + supplies the complete static default. An explicit empty object opts out; any nonempty + object is completed with omitted inputTokens, outputTokens and load defaults. properties: inputTokens: description: InputTokens is the synthetic request input length in tokens. @@ -321,11 +332,87 @@ spec: minimum: 1 type: integer type: object + trace: + description: |- + Trace replays a recorded request trace from a PersistentVolumeClaim in the Request + namespace. Admission supplies no static default when it is set. The Run Job mounts only the + trace file, read-only through a subPath of the claim, at /var/run/aisimulate/trace/trace.jsonl, + or trace.jsonl.gz when subPath ends in .gz. The trace lowers to the default traffic.source + and traffic.load, which spec.aisimulate.inputOverrides can replace or delete. + properties: + format: + default: Dynamo + description: Format identifies the trace decoder. Defaults to Dynamo. + enum: + - Dynamo + type: string + replay: + description: Replay configures how the decoded trace is replayed. + properties: + arrivalSpeedupRatio: + description: |- + ArrivalSpeedupRatio multiplies the recorded request-arrival rate. 1.0 preserves it. + When omitted, the recorded arrival times are replayed unchanged. + exclusiveMinimum: true + minimum: 0 + type: number + type: object + source: + description: Source identifies the trace artifact. + properties: + persistentVolumeClaim: + description: |- + PersistentVolumeClaim reads the trace from a PersistentVolumeClaim in the Request + namespace. + properties: + claimName: + description: |- + ClaimName is the PersistentVolumeClaim, in the Request namespace, that holds the trace. + The Run Job mounts only the trace file from it, read-only, and only into the AISimulate + container. Each Run protects the claim with the finalizer `nvidia.com/dgdr-run-` + until the Run completes or, for a Run deleted before it completes, until no Job of the Run + is left. Garbage collection deletes the Job of a deleted Run before the Job's pod, so the + finalizer of a deleted Run can be released while that pod still stops; Kubernetes + PersistentVolumeClaim protection keeps the claim until no pod uses it. Deleting the claim + cancels the Runs that read it. + maxLength: 253 + minLength: 1 + pattern: ^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$ + type: string + subPath: + description: |- + SubPath names one trace file relative to the root of the claim; the Run Job mounts it as + the subPath of the claim volume. It is neither a directory nor a glob. It must not start or + end with '/' and must not contain empty, '.' or '..' segments or a NUL character, which no + file name can hold. The file must be complete before the Request is created and must never + be rewritten: every Run of the Request reads it when its Job starts. + maxLength: 1024 + minLength: 1 + type: string + x-kubernetes-validations: + - message: subPath must be relative to the claim root without empty, '.' or '..' segments + rule: '!self.startsWith(''/'') && !self.endsWith(''/'') && !self.contains(''//'') && !self.matches(''(^|/)[.][.]?(/|$)'')' + - message: subPath must not contain a NUL character + rule: '!self.contains(''\u0000'')' + required: + - claimName + - subPath + type: object + type: object + x-kubernetes-validations: + - message: persistentVolumeClaim is required + rule: has(self.persistentVolumeClaim) + required: + - source + type: object type: object + x-kubernetes-validations: + - message: static and trace are mutually exclusive + rule: '!(has(self.static) && has(self.trace))' type: object x-kubernetes-validations: - message: spec must be a fully resolved Request admission snapshot - rule: has(self.model) && has(self.backends) && size(self.backends) > 0 && has(self.hardware) && ((!has(self.hardware.gpu) && !has(self.hardware.numGpusPerNode)) || (has(self.hardware.gpu) && has(self.hardware.gpu.sku) && has(self.hardware.gpu.budget) && has(self.hardware.numGpusPerNode))) && has(self.workload) && has(self.workload.static) && ((!has(self.workload.static.inputTokens) && !has(self.workload.static.outputTokens) && !has(self.workload.static.load)) || (has(self.workload.static.inputTokens) && has(self.workload.static.outputTokens) && has(self.workload.static.load))) && has(self.objective) && has(self.objective.mode) && has(self.objective.metric) && has(self.objective.sla) && has(self.recommendation) && has(self.recommendation.maxCandidates) && self.recommendation.maxCandidates > 0 + rule: has(self.model) && has(self.backends) && size(self.backends) > 0 && has(self.hardware) && ((!has(self.hardware.gpu) && !has(self.hardware.numGpusPerNode)) || (has(self.hardware.gpu) && has(self.hardware.gpu.sku) && has(self.hardware.gpu.budget) && has(self.hardware.numGpusPerNode))) && has(self.workload) && (has(self.workload.trace) || (has(self.workload.static) && ((!has(self.workload.static.inputTokens) && !has(self.workload.static.outputTokens) && !has(self.workload.static.load)) || (has(self.workload.static.inputTokens) && has(self.workload.static.outputTokens) && has(self.workload.static.load))))) && has(self.objective) && has(self.objective.mode) && has(self.objective.metric) && has(self.objective.sla) && has(self.recommendation) && has(self.recommendation.maxCandidates) && self.recommendation.maxCandidates > 0 status: description: DynamoGraphDeploymentRunStatus represents one search execution's observed state. properties: diff --git a/deploy/operator/internal/aisimulate/defaults.go b/deploy/operator/internal/aisimulate/defaults.go index d26bbd54fea5..ba2be457abbc 100644 --- a/deploy/operator/internal/aisimulate/defaults.go +++ b/deploy/operator/internal/aisimulate/defaults.go @@ -57,7 +57,9 @@ func applyDefaults(underlay map[string]any) error { b := builder{root: underlay} if has(underlay, trafficPath...) { - if !has(underlay, trafficStopRequests...) { + // A trace accepts only max_virtual_time_seconds; a request count would be rejected. + sourceType, _, _ := unstructured.NestedString(underlay, trafficSourceType...) + if sourceType != "trace" && !has(underlay, trafficStopRequests...) { b.set(float64(DefaultRequestCount), trafficStopRequests) } } diff --git a/deploy/operator/internal/aisimulate/doc.go b/deploy/operator/internal/aisimulate/doc.go index a41ef4bee3d7..63f8a3fbe335 100644 --- a/deploy/operator/internal/aisimulate/doc.go +++ b/deploy/operator/internal/aisimulate/doc.go @@ -27,12 +27,14 @@ // Lower starts with the public engine-role underlay required for AISimulate's default // aggregated-and-disaggregated search. It then translates the typed spec (model, backends, // hardware, workload, objective) into that input. The derived defaults that no API field -// carries are added next: the request count, a bounded trial budget for a -// KV-load range search, and optimization.strict_sla when an SLA bound is set. Last, -// spec.aisimulate.inputOverrides is applied on top as a JSON Merge Patch (RFC 7386): upstream -// objects merge recursively, arrays and scalars replace the value, and null deletes it. The -// input overrides are authoritative: they win on every path they name, including the derived -// defaults, and are applied once, last. +// carries are added next: the request count (not for a trace, which accepts only a virtual-time +// stop), a bounded trial budget for a KV-load range search, and optimization.strict_sla when an +// SLA bound is set. Last, spec.aisimulate.inputOverrides is applied on top as a JSON Merge +// Patch (RFC 7386): upstream objects merge recursively, arrays and scalars replace the value, +// and null deletes it. The input overrides are authoritative: they win on every path they name, +// including the derived defaults and the traffic a typed trace lowers, and are applied once, last. +// Only the operator-owned dgdr_run fields are exempt: Request admission rejects them and the Run +// controller restores them after this lowering. // // The document stays unstructured (map[string]any) because its schema belongs to the // AISimulate release bundled in the operator's Planner image, not to this package. Reads and writes go diff --git a/deploy/operator/internal/aisimulate/input.go b/deploy/operator/internal/aisimulate/input.go index 54d42bc1c679..ad9a6e41fef4 100644 --- a/deploy/operator/internal/aisimulate/input.go +++ b/deploy/operator/internal/aisimulate/input.go @@ -161,6 +161,8 @@ func Underlay(spec *v1beta2.NativeSearchInput) (map[string]any, error) { } } + b.setTraceWorkload(spec.Workload) + if objective := spec.Objective; objective != nil { if objective.Mode == v1beta2.ObjectiveModePareto { if len(objective.Metrics) > 0 && !isDefaultParetoMetrics(objective.Metrics) { @@ -191,6 +193,25 @@ func Underlay(spec *v1beta2.NativeSearchInput) (map[string]any, error) { return b.root, b.err } +// setTraceWorkload lowers workload.trace: it replays the recorded arrival times, scaled by the +// speedup only when one is set. workload may be nil; a workload without a trace lowers nothing. +func (b *builder) setTraceWorkload(workload *v1beta2.WorkloadSpec) { + if workload == nil || workload.Trace == nil { + return + } + trace := workload.Trace + source, err := traceSource(trace) + if err != nil { + b.err = err + return + } + b.set(source, trafficSource) + b.set("trace_timestamps", trafficLoadType) + if replay := trace.Replay; replay != nil && replay.ArrivalSpeedupRatio != nil { + b.set(*replay.ArrivalSpeedupRatio, trafficLoadSpeedup) + } +} + func hasStaticWorkloadFields(static *v1beta2.StaticWorkloadSpec) bool { return static.InputTokens != nil || static.OutputTokens != nil || static.Load != nil } diff --git a/deploy/operator/internal/aisimulate/paths.go b/deploy/operator/internal/aisimulate/paths.go index 25ef533ecd6c..82a5b4ab5b08 100644 --- a/deploy/operator/internal/aisimulate/paths.go +++ b/deploy/operator/internal/aisimulate/paths.go @@ -37,6 +37,7 @@ var ( optimizerMaxTrials = []string{"optimizer", "max_trials"} trafficPath = []string{"traffic"} + trafficSource = []string{"traffic", "source"} trafficSourceType = []string{"traffic", "source", "type"} trafficSourceInputTokens = []string{"traffic", "source", "input_tokens"} trafficSourceOutputTokens = []string{"traffic", "source", "output_tokens"} @@ -44,6 +45,7 @@ var ( trafficLoadConcurrency = []string{"traffic", "load", "concurrency"} trafficLoadRequestsPerSecond = []string{"traffic", "load", "requests_per_second"} trafficLoadFraction = []string{"traffic", "load", "fraction"} + trafficLoadSpeedup = []string{"traffic", "load", "speedup"} trafficStopRequests = []string{"traffic", "stop", "requests"} ) diff --git a/deploy/operator/internal/aisimulate/testdata/trace_request_lowering.golden.json b/deploy/operator/internal/aisimulate/testdata/trace_request_lowering.golden.json new file mode 100644 index 000000000000..59560e0daec5 --- /dev/null +++ b/deploy/operator/internal/aisimulate/testdata/trace_request_lowering.golden.json @@ -0,0 +1,191 @@ +{ + "request_spec": { + "model": { + "name": "Qwen/Qwen3-32B" + }, + "backends": [ + "vllm" + ], + "hardware": { + "gpu": { + "sku": "h200_sxm", + "budget": 6 + } + }, + "workload": { + "trace": { + "source": { + "persistentVolumeClaim": { + "claimName": "request-traces", + "subPath": "windows/llm/1793000000000-ba081802aba3.jsonl.gz" + } + } + } + }, + "objective": { + "mode": "Optimize", + "metric": "GoodputPerGPU", + "sla": { + "ttftMs": 500.0, + "itlMs": 30.0 + } + }, + "recommendation": { + "maxCandidates": 8 + }, + "aisimulate": { + "inputOverrides": { + "engine": { + "mode": "aggregated", + "backend": "vllm", + "context_length": 32768, + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + { + "tensor": 2, + "pipeline": 1, + "attention_data": 1, + "moe_tensor": 1, + "moe_expert": 1, + "replicas": 3 + } + ] + } + }, + "prefill": null, + "decode": null + } + }, + "router": { + "policy": "kv_router", + "prefill_load_model": { + "type": "none" + }, + "overlap_score_credit": { + "choices": [ + 0.5, + 1.0 + ] + }, + "prefill_load_scale": { + "choices": [ + 1.0, + 2.0 + ] + }, + "temperature": { + "choices": [ + 0.0, + 0.5 + ] + } + }, + "optimizer": { + "algorithm": "random", + "max_trials": 16, + "parallelism": 8, + "candidate_timeout_seconds": 300.0 + }, + "optimization": { + "strict_sla": false + }, + "dgdr_run": { + "generator": "direct" + } + } + } + }, + "run_input": { + "dgdr_run": { + "derive_runtime_images": true, + "generator": "direct", + "max_candidates": 8, + "name": "llm-rt-golden-1", + "namespace": "test-ns", + "num_gpus_per_node": 8, + "runtime_image": "example.com/dynamo-planner:1.2.3", + "runtime_version_override": "1.2.3", + "snapshot_dir": "/var/run/aisimulate/output" + }, + "engine": { + "backend": "vllm", + "context_length": 32768, + "hardware": "h200_sxm", + "mode": "aggregated", + "model": "Qwen/Qwen3-32B", + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + { + "attention_data": 1, + "moe_expert": 1, + "moe_tensor": 1, + "pipeline": 1, + "replicas": 3, + "tensor": 2 + } + ] + } + } + } + }, + "evaluation": { + "sla": { + "itl_ms": 30, + "ttft_ms": 500 + } + }, + "optimization": { + "constraints": { + "max_candidate_gpus": 6 + }, + "strict_sla": false, + "target": "goodput_per_gpu" + }, + "optimizer": { + "algorithm": "random", + "candidate_timeout_seconds": 300, + "max_trials": 16, + "parallelism": 8 + }, + "router": { + "overlap_score_credit": { + "choices": [ + 0.5, + 1 + ] + }, + "policy": "kv_router", + "prefill_load_model": { + "type": "none" + }, + "prefill_load_scale": { + "choices": [ + 1, + 2 + ] + }, + "temperature": { + "choices": [ + 0, + 0.5 + ] + } + }, + "traffic": { + "load": { + "type": "trace_timestamps" + }, + "source": { + "format": "dynamo", + "paths": [ + "/var/run/aisimulate/trace/trace.jsonl.gz" + ], + "type": "trace" + } + } + } +} diff --git a/deploy/operator/internal/aisimulate/trace.go b/deploy/operator/internal/aisimulate/trace.go new file mode 100644 index 000000000000..0a98d2808bd2 --- /dev/null +++ b/deploy/operator/internal/aisimulate/trace.go @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package aisimulate + +import ( + "fmt" + "strings" + + v1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" +) + +// TraceMountPath is the directory of the Run Job's sweeper container that holds the trace file. +const TraceMountPath = "/var/run/aisimulate/trace" + +// TracePath is where the Run Job mounts the trace file at subPath of the claim: one fixed file +// under TraceMountPath. The Job mounts only that file, through a subPath of the claim, so the +// kubelet resolves subPath and refuses a symlink that resolves outside the claim. The name keeps +// the ".gz" suffix, by which AISimulate recognizes gzip-compressed JSON Lines. +func TracePath(subPath string) string { + if strings.HasSuffix(subPath, ".gz") { + return TraceMountPath + "/trace.jsonl.gz" + } + return TraceMountPath + "/trace.jsonl" +} + +// traceSource returns the traffic.source object a typed trace lowers to: a trace in AISimulate's +// native format read from its one claim file under TraceMountPath. trace must not be nil. A trace +// without a persistentVolumeClaim source has no paths; the CRD schema rejects that shape. +func traceSource(trace *v1beta2.TraceWorkloadSpec) (map[string]any, error) { + format, err := nativeTraceFormat(trace.Format) + if err != nil { + return nil, err + } + source := map[string]any{"type": "trace", "format": format} + if claim := trace.Source.PersistentVolumeClaim; claim != nil { + source["paths"] = []any{TracePath(claim.SubPath)} + } + return source, nil +} + +// nativeTraceFormat maps each API trace format to AISimulate's traffic.source.format spelling. +// The CRD schema defaults an omitted format, so an empty format is an invalid input here. +func nativeTraceFormat(format v1beta2.TraceFormat) (string, error) { + switch format { + case v1beta2.TraceFormatDynamo: + return "dynamo", nil + default: + return "", fmt.Errorf("spec.workload.trace.format: unsupported trace format %q", format) + } +} diff --git a/deploy/operator/internal/aisimulate/trace_test.go b/deploy/operator/internal/aisimulate/trace_test.go new file mode 100644 index 000000000000..2183dcdd193e --- /dev/null +++ b/deploy/operator/internal/aisimulate/trace_test.go @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package aisimulate + +import ( + "testing" + + v1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" +) + +func withTrace(spec *v1beta2.DynamoGraphDeploymentRequestSpec, speedup *float64) *v1beta2.DynamoGraphDeploymentRequestSpec { + spec.Workload = &v1beta2.WorkloadSpec{Trace: &v1beta2.TraceWorkloadSpec{ + Format: v1beta2.TraceFormatDynamo, + Source: v1beta2.TraceSourceSpec{PersistentVolumeClaim: &v1beta2.TracePVCSource{ClaimName: "request-traces", SubPath: "retune/w-1.jsonl.gz"}}, + }} + if speedup != nil { + spec.Workload.Trace.Replay = &v1beta2.TraceReplaySpec{ArrivalSpeedupRatio: speedup} + } + return spec +} + +func TestLowerTraceWorkload(t *testing.T) { + tests := []struct { + name string + spec *v1beta2.DynamoGraphDeploymentRequestSpec + wantTraffic string + }{ + { + name: "recorded arrival times in the native dynamo format without a request stop", + spec: withTrace(defaultedSpec(), nil), + wantTraffic: `{"load":{"type":"trace_timestamps"},` + + `"source":{"format":"dynamo","paths":["/var/run/aisimulate/trace/trace.jsonl.gz"],"type":"trace"}}`, + }, + { + name: "an arrival speedup lowers to traffic.load.speedup", + spec: withTrace(defaultedSpec(), new(2.5)), + wantTraffic: `{"load":{"speedup":2.5,"type":"trace_timestamps"},` + + `"source":{"format":"dynamo","paths":["/var/run/aisimulate/trace/trace.jsonl.gz"],"type":"trace"}}`, + }, + { + name: "a stop override still applies", + spec: withInputOverrides(withTrace(defaultedSpec(), nil), `{"traffic":{"stop":{"max_virtual_time_seconds":600}}}`), + wantTraffic: `{"load":{"type":"trace_timestamps"},` + + `"source":{"format":"dynamo","paths":["/var/run/aisimulate/trace/trace.jsonl.gz"],"type":"trace"},` + + `"stop":{"max_virtual_time_seconds":600}}`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + traffic := nested(t, lower(t, tt.spec), "traffic").(map[string]any) + if got := encode(t, traffic); got != tt.wantTraffic { + t.Fatalf("traffic\n got: %s\nwant: %s", got, tt.wantTraffic) + } + }) + } +} + +func TestLowerTraceWorkloadRejectsAFormatWithoutANativeSpelling(t *testing.T) { + for _, format := range []v1beta2.TraceFormat{"", "dynamo", "mooncake"} { + t.Run(string(format), func(t *testing.T) { + spec := withTrace(defaultedSpec(), nil) + spec.Workload.Trace.Format = format + if _, err := Lower(&spec.NativeSearchInput); err == nil { + t.Fatalf("Lower accepted trace format %q", format) + } + }) + } +} + +func TestTracePathIsOneFixedFileUnderTheMount(t *testing.T) { + for subPath, want := range map[string]string{ + "retune/w-1.jsonl.gz": "/var/run/aisimulate/trace/trace.jsonl.gz", + "retune/w-1.jsonl": "/var/run/aisimulate/trace/trace.jsonl", + } { + if got := TracePath(subPath); got != want { + t.Errorf("TracePath(%q) = %q, want %q", subPath, got, want) + } + } +} diff --git a/deploy/operator/internal/controller/clusterenv_test.go b/deploy/operator/internal/controller/clusterenv_test.go index a9333464cceb..fb7d3dc3ec02 100644 --- a/deploy/operator/internal/controller/clusterenv_test.go +++ b/deploy/operator/internal/controller/clusterenv_test.go @@ -47,6 +47,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" logf "sigs.k8s.io/controller-runtime/pkg/log" "sigs.k8s.io/controller-runtime/pkg/log/zap" + "sigs.k8s.io/yaml" ) var clusterTestEnv *clusterenv.Env @@ -367,6 +368,26 @@ func TestClusterDynamoGraphDeploymentRequestV1beta2ProducesCandidates(t *testing t.Log("Allow the AISimulate Job Pod to run while preventing unrelated ReplicaSet workloads") env.BlockReplicaSets() + t.Log("Write a trace scenario's trace to the claim file its Request names before applying it") + traceFixture := filepath.Join(scenarioDir, "trace.jsonl") + var trace clusterTestTraceFile + traceBefore := "" + if _, err := os.Stat(traceFixture); err == nil { + scenario := &nvidiacomv1beta2.DynamoGraphDeploymentRequest{} + data, err := os.ReadFile(input) + require.NoError(t, err) + require.NoError(t, yaml.Unmarshal(data, scenario)) + require.NotNil(t, scenario.Spec.Workload, "trace scenario input has no workload") + require.NotNil(t, scenario.Spec.Workload.Trace, "trace scenario input has no trace workload") + require.NotNil(t, scenario.Spec.Workload.Trace.Source.PersistentVolumeClaim, "trace scenario input has no trace claim") + trace = clusterTestTraceFile{ + claim: scenario.Spec.Workload.Trace.Source.PersistentVolumeClaim.ClaimName, + subPath: scenario.Spec.Workload.Trace.Source.PersistentVolumeClaim.SubPath, + } + createClusterTestTraceClaim(t, env, trace) + traceBefore = writeClusterTestTrace(t, env, trace, traceFixture, plannerImage) + } + t.Log("Apply the native v1beta2 Request through production admission") objects := golden.ApplyManifests(t, filepath.Join(scenarioDir, "input.yaml"), env.Client(), env.Namespace()) if len(objects) != 1 || objects[0].GetKind() != "DynamoGraphDeploymentRequest" { @@ -486,6 +507,12 @@ func TestClusterDynamoGraphDeploymentRequestV1beta2ProducesCandidates(t *testing t.Log("Match the published Candidate API contract against the reviewed golden") golden.EventuallyMatchManifests(t, env.Client(), env.Namespace(), filepath.Join(scenarioDir, "output.yaml")) + + if traceBefore != "" { + t.Log("The Run replayed the mounted trace file and left its contents, owner and mode unchanged") + requireClusterTestTraceInput(t, &completedRun, trace) + require.Equal(t, traceBefore, statClusterTestTrace(t, env, trace, "trace-after", plannerImage)) + } }) } } diff --git a/deploy/operator/internal/controller/dgdr_v1beta2_trace_clustertest_test.go b/deploy/operator/internal/controller/dgdr_v1beta2_trace_clustertest_test.go new file mode 100644 index 000000000000..c0b787b31872 --- /dev/null +++ b/deploy/operator/internal/controller/dgdr_v1beta2_trace_clustertest_test.go @@ -0,0 +1,449 @@ +//go:build clustertest + +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "context" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + "testing" + "time" + + semver "github.com/Masterminds/semver/v3" + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/aisimulate" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" + "github.com/ai-dynamo/dynamo/deploy/operator/internal/testing/clusterenv" + "github.com/stretchr/testify/require" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/meta" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/watch" + "k8s.io/utils/ptr" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/yaml" +) + +// The writer gives a trace file an owner the Job's users do not have and the trace group, so the +// sweeper can read it only through the supplemental group of the Run Job's pod. +const ( + clusterTestTraceUID = 4321 + clusterTestTraceGID = dgdrRunTraceGroup + clusterTestTraceMode = "640" +) + +// clusterTestCreateContainerConfigError is the kubelet's waiting reason for a container whose +// configuration it cannot apply, such as a subPath that resolves outside its volume. +const clusterTestCreateContainerConfigError = "CreateContainerConfigError" + +// clusterTestTraceFile is the claim and the claim-relative subPath a trace Request reads. +type clusterTestTraceFile struct { + claim, subPath string +} + +// clusterTestTraceWriter gzips the fixture to the claim file and gives it the requested owner and mode. +const clusterTestTraceWriter = ` +import gzip, os, shutil, sys +source, target, uid, gid, mode = sys.argv[1:] +os.makedirs(os.path.dirname(target), exist_ok=True) +with open(source, "rb") as src, gzip.GzipFile(target, "wb", mtime=0) as dst: + shutil.copyfileobj(src, dst) +os.chown(target, int(uid), int(gid)) +os.chmod(target, int(mode, 8)) +` + +// clusterTestTraceLinker places a symlink at the claim file that points outside the claim. +const clusterTestTraceLinker = ` +import os, sys +link, target = sys.argv[1:] +os.makedirs(os.path.dirname(link), exist_ok=True) +os.symlink(target, link) +` + +// clusterTestTraceStat reports the claim file's owner, mode and digest as its termination message. +const clusterTestTraceStat = ` +import hashlib, os, sys +st = os.lstat(sys.argv[1]) +with open(sys.argv[1], "rb") as f: + digest = hashlib.sha256(f.read()).hexdigest() +with open("/dev/termination-log", "w") as log: + log.write(f"{st.st_uid} {st.st_gid} {st.st_mode & 0o7777:o} {digest}") +` + +// createClusterTestTraceClaim creates the static filesystem claim of trace in env's namespace. Kind +// installs no dynamic provisioner the test may rely on, so the claim binds a PersistentVolume of its +// own on one node, which is deleted with the test. +func createClusterTestTraceClaim(t *testing.T, env *clusterenv.TestEnv, trace clusterTestTraceFile) { + t.Helper() + ctx := t.Context() + + t.Log("Pin the trace volume to one node, so every pod that mounts it sees the same files") + var nodes corev1.NodeList + require.NoError(t, env.Client().List(ctx, &nodes)) + require.NotEmpty(t, nodes.Items, "the cluster has no nodes") + hostname := nodes.Items[0].Labels[corev1.LabelHostname] + require.NotEmpty(t, hostname, "node %q has no %s label", nodes.Items[0].Name, corev1.LabelHostname) + + t.Log("Create the PersistentVolume and the claim bound to it") + volume := &corev1.PersistentVolume{ + Name: "dgdr-trace-" + env.Namespace(), + Spec: corev1.PersistentVolumeSpec{ + Capacity: corev1.ResourceList{corev1.ResourceStorage: resource.MustParse("16Mi")}, + AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce}, + PersistentVolumeReclaimPolicy: corev1.PersistentVolumeReclaimRetain, + ClaimRef: &corev1.ObjectReference{Namespace: env.Namespace(), Name: trace.claim}, + PersistentVolumeSource: corev1.PersistentVolumeSource{HostPath: &corev1.HostPathVolumeSource{ + Path: "/var/local/dynamo-clustertest/" + env.Namespace(), + Type: ptr.To(corev1.HostPathDirectoryOrCreate), + }}, + NodeAffinity: &corev1.VolumeNodeAffinity{Required: &corev1.NodeSelector{NodeSelectorTerms: []corev1.NodeSelectorTerm{{ + MatchExpressions: []corev1.NodeSelectorRequirement{{ + Key: corev1.LabelHostname, Operator: corev1.NodeSelectorOpIn, Values: []string{hostname}, + }}, + }}}}, + }, + } + require.NoError(t, env.Client().Create(ctx, volume)) + t.Cleanup(func() { + if err := env.Client().Delete(context.Background(), volume); client.IgnoreNotFound(err) != nil { + t.Errorf("delete trace PersistentVolume %q: %v", volume.Name, err) + } + }) + claim := &corev1.PersistentVolumeClaim{ + Name: trace.claim, Namespace: env.Namespace(), + Spec: corev1.PersistentVolumeClaimSpec{ + AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce}, + VolumeMode: ptr.To(corev1.PersistentVolumeFilesystem), + StorageClassName: new(""), + VolumeName: volume.Name, + Resources: corev1.VolumeResourceRequirements{ + Requests: corev1.ResourceList{corev1.ResourceStorage: resource.MustParse("16Mi")}, + }, + }, + } + require.NoError(t, env.Client().Create(ctx, claim)) +} + +// runClusterTestTraceJob runs script with args in image, with claim mounted at /claim (read-only +// unless writable) and the files of configMap, if named, at /fixture. It waits for the Job to +// succeed and returns its termination message. +func runClusterTestTraceJob( + t *testing.T, + env *clusterenv.TestEnv, + name, image, claim, configMap string, + writable bool, + script string, + args ...string, +) string { + t.Helper() + ctx := t.Context() + + t.Logf("Run the %s Job as root with the trace claim mounted", name) + pod := corev1.PodSpec{ + RestartPolicy: corev1.RestartPolicyNever, + SecurityContext: &corev1.PodSecurityContext{RunAsUser: ptr.To[int64](0), RunAsGroup: ptr.To[int64](0)}, + Containers: []corev1.Container{{ + Name: "trace", + Image: image, + ImagePullPolicy: corev1.PullIfNotPresent, + Command: append([]string{"python3", "-c", script}, args...), + VolumeMounts: []corev1.VolumeMount{{Name: "claim", MountPath: "/claim", ReadOnly: !writable}}, + }}, + Volumes: []corev1.Volume{{ + Name: "claim", + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: claim, ReadOnly: !writable}, + }}, + } + if configMap != "" { + pod.Containers[0].VolumeMounts = append(pod.Containers[0].VolumeMounts, corev1.VolumeMount{Name: "fixture", MountPath: "/fixture", ReadOnly: true}) + pod.Volumes = append(pod.Volumes, corev1.Volume{ + Name: "fixture", + ConfigMap: &corev1.ConfigMapVolumeSource{Name: configMap}, + }) + } + job := &batchv1.Job{ + Name: name, Namespace: env.Namespace(), + Spec: batchv1.JobSpec{BackoffLimit: ptr.To[int32](0), Template: corev1.PodTemplateSpec{Spec: pod}}, + } + require.NoError(t, env.Client().Create(ctx, job)) + + t.Logf("Wait for the %s Job to succeed", name) + require.Eventually(t, func() bool { + require.NoError(t, env.Client().Get(ctx, client.ObjectKeyFromObject(job), job)) + require.Zero(t, job.Status.Failed, "%s Job failed: %+v", name, job.Status.Conditions) + return job.Status.Succeeded == 1 + }, 5*time.Minute, time.Second) + + t.Logf("Read the %s Job's termination message", name) + var pods corev1.PodList + require.NoError(t, env.Client().List(ctx, &pods, client.InNamespace(env.Namespace()), + client.MatchingLabels{batchv1.JobNameLabel: name})) + require.Len(t, pods.Items, 1) + require.Len(t, pods.Items[0].Status.ContainerStatuses, 1) + terminated := pods.Items[0].Status.ContainerStatuses[0].State.Terminated + require.NotNil(t, terminated, "%s pod did not terminate", name) + return terminated.Message +} + +// writeClusterTestTrace gzips the fixture at fixture into trace and returns the file's owner, mode +// and digest as the claim holds it. +func writeClusterTestTrace(t *testing.T, env *clusterenv.TestEnv, trace clusterTestTraceFile, fixture, image string) string { + t.Helper() + ctx := t.Context() + + t.Log("Carry the checked-in trace fixture into the cluster") + content, err := os.ReadFile(fixture) + require.NoError(t, err) + require.NoError(t, env.Client().Create(ctx, &corev1.ConfigMap{ + Name: "trace-fixture", Namespace: env.Namespace(), + Data: map[string]string{"trace.jsonl": string(content)}, + })) + + runClusterTestTraceJob(t, env, "trace-writer", image, trace.claim, "trace-fixture", true, clusterTestTraceWriter, + "/fixture/trace.jsonl", filepath.Join("/claim", trace.subPath), + fmt.Sprint(clusterTestTraceUID), fmt.Sprint(clusterTestTraceGID), clusterTestTraceMode) + return statClusterTestTrace(t, env, trace, "trace-before", image) +} + +// statClusterTestTrace returns the owner, mode and digest of trace as its claim holds it, read by a +// Job named name. +func statClusterTestTrace(t *testing.T, env *clusterenv.TestEnv, trace clusterTestTraceFile, name, image string) string { + t.Helper() + stat := runClusterTestTraceJob(t, env, name, image, trace.claim, "", false, clusterTestTraceStat, + filepath.Join("/claim", trace.subPath)) + require.Regexp(t, fmt.Sprintf(`^%d %d %s [0-9a-f]{64}$`, clusterTestTraceUID, clusterTestTraceGID, clusterTestTraceMode), stat) + return stat +} + +// requireClusterTestTraceInput checks that run replayed trace through its fixed mount path. +func requireClusterTestTraceInput(t *testing.T, run *nvidiacomv1beta2.DynamoGraphDeploymentRun, trace clusterTestTraceFile) { + t.Helper() + require.NotNil(t, run.Status.AISimulate, "completed Run has no AISimulate status") + var input struct { + Traffic struct { + Source map[string]any `json:"source"` + Load map[string]any `json:"load"` + } `json:"traffic"` + } + require.NoError(t, yaml.Unmarshal(run.Status.AISimulate.Input.Raw, &input)) + require.Equal(t, map[string]any{ + "type": "trace", "format": "dynamo", "paths": []any{aisimulate.TracePath(trace.subPath)}, + }, input.Traffic.Source) + require.Equal(t, "trace_timestamps", input.Traffic.Load["type"]) +} + +func TestClusterDynamoGraphDeploymentRequestV1beta2RefusesATraceSymlinkThatLeavesItsClaim(t *testing.T) { + ctx := t.Context() + plannerImage := clusterenv.RequireEnv(t, "DYNAMO_CLUSTERTEST_PROFILER_IMAGE") + operatorImage := clusterenv.RequireEnv(t, "DYNAMO_CLUSTERTEST_OPERATOR_IMAGE") + + t.Log("Create an isolated namespace in the explicitly unlocked cluster") + env := clusterTestEnv.RunT(t) + env.BlockReplicaSets() + + t.Log("Place a symlink at the trace path that points at a file outside the claim") + trace := clusterTestTraceFile{claim: "linked-traces", subPath: "retune/escape.jsonl.gz"} + createClusterTestTraceClaim(t, env, trace) + runClusterTestTraceJob(t, env, "trace-linker", plannerImage, trace.claim, "", true, clusterTestTraceLinker, + filepath.Join("/claim", trace.subPath), "/etc/passwd") + + t.Log("Apply a Request that reads the linked trace path through production admission") + request := &nvidiacomv1beta2.DynamoGraphDeploymentRequest{ + Name: "aisimulate-trace-symlink", Namespace: env.Namespace(), + Spec: nvidiacomv1beta2.DynamoGraphDeploymentRequestSpec{ + DynamoGraphDeploymentRequestNativeSpec: nvidiacomv1beta2.DynamoGraphDeploymentRequestNativeSpec{ + NativeSearchInput: nvidiacomv1beta2.NativeSearchInput{ + Model: &nvidiacomv1beta2.ModelSpec{Name: "Qwen/Qwen3-32B"}, + Backends: []nvidiacomv1beta2.Backend{nvidiacomv1beta2.BackendVLLM}, + Hardware: &nvidiacomv1beta2.HardwareSpec{GPU: &nvidiacomv1beta2.GPUHardwareSpec{ + SKU: new("h200_sxm"), Budget: new(int32(2)), + }}, + Workload: &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Format: nvidiacomv1beta2.TraceFormatDynamo, + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: trace.claim, SubPath: trace.subPath, + }}, + }}, + }, + }, + }, + } + require.NoError(t, env.Client().Create(ctx, request)) + + t.Log("Watch every Run transition before starting the production controllers") + watchClient, err := client.NewWithWatch(env.RESTConfig(), client.Options{Scheme: newClusterTestScheme()}) + require.NoError(t, err) + runWatch, err := watchClient.Watch(ctx, &nvidiacomv1beta2.DynamoGraphDeploymentRunList{}, client.InNamespace(env.Namespace())) + require.NoError(t, err) + defer runWatch.Stop() + startClusterTestDGDRV1beta2Controllers(t, env, plannerImage, operatorImage) + + // An image pull or other waiting state can come first and recover; only the configuration error + // shows that the kubelet tried the trace mount and refused it. + t.Log("The kubelet refuses the subPath, so the sweeper waits with CreateContainerConfigError") + var sweeperPod types.UID + require.Eventually(t, func() bool { + var pods corev1.PodList + require.NoError(t, env.Client().List(ctx, &pods, client.InNamespace(env.Namespace()))) + for i := range pods.Items { + for _, status := range pods.Items[i].Status.ContainerStatuses { + if status.Name == dgdrRunSweeperContainer && status.State.Waiting != nil && + status.State.Waiting.Reason == clusterTestCreateContainerConfigError { + sweeperPod = pods.Items[i].UID + return true + } + } + } + return false + }, 5*time.Minute, time.Second, "the sweeper did not wait with %s", clusterTestCreateContainerConfigError) + + t.Log("The open Run reports that configuration error of the sweeper") + var waiting nvidiacomv1beta2.DynamoGraphDeploymentRun + awaitClusterTestRun(t, runWatch, "the Run did not report its sweeper's configuration error", func(run *nvidiacomv1beta2.DynamoGraphDeploymentRun) bool { + waiting = *run.DeepCopy() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + return condition != nil && condition.Reason == dgdrRunReasonContainerWaiting && strings.HasPrefix(condition.Message, + fmt.Sprintf("container %q is waiting: %s: ", dgdrRunSweeperContainer, clusterTestCreateContainerConfigError)) + }) + condition := meta.FindStatusCondition(waiting.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + require.Equal(t, metav1.ConditionFalse, condition.Status) + require.NotNil(t, waiting.Status.JobRef) + require.Nil(t, waiting.Status.CompletionTime) + + t.Log("The sweeper never started: its pod reports the kubelet's waiting state, which the Run carries verbatim") + var pods corev1.PodList + require.NoError(t, env.Client().List(ctx, &pods, client.InNamespace(env.Namespace()), + client.MatchingLabels{batchv1.JobNameLabel: waiting.Status.JobRef.Name})) + require.Len(t, pods.Items, 1) + require.Equal(t, sweeperPod, pods.Items[0].UID) + var sweeper *corev1.ContainerStatus + for i := range pods.Items[0].Status.ContainerStatuses { + if pods.Items[0].Status.ContainerStatuses[i].Name == dgdrRunSweeperContainer { + sweeper = &pods.Items[0].Status.ContainerStatuses[i] + } + } + require.NotNil(t, sweeper) + require.NotNil(t, sweeper.State.Waiting) + require.Equal(t, clusterTestCreateContainerConfigError, sweeper.State.Waiting.Reason) + require.Nil(t, sweeper.State.Running) + require.Nil(t, sweeper.LastTerminationState.Terminated) + require.Zero(t, sweeper.RestartCount) + require.Contains(t, condition.Message, fmt.Sprintf("container %q is waiting: %s: ", dgdrRunSweeperContainer, sweeper.State.Waiting.Reason)) + + t.Log("The waiting container alone neither ends the Run nor makes the operator delete its Job") + require.Never(t, func() bool { + job := &batchv1.Job{} + err := env.Client().Get(ctx, client.ObjectKey{Namespace: env.Namespace(), Name: waiting.Status.JobRef.Name}, job) + run := &nvidiacomv1beta2.DynamoGraphDeploymentRun{} + require.NoError(t, env.Client().Get(ctx, client.ObjectKeyFromObject(&waiting), run)) + return err != nil || job.DeletionTimestamp != nil || + meta.IsStatusConditionTrue(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + }, 20*time.Second, time.Second) + + t.Log("The Run's finalizer protects the claim; deleting the claim cancels the Run") + claim := &corev1.PersistentVolumeClaim{} + require.NoError(t, env.Client().Get(ctx, client.ObjectKey{Namespace: env.Namespace(), Name: trace.claim}, claim)) + require.Contains(t, claim.Finalizers, dgdrRunTraceClaimFinalizer(waiting.UID)) + require.NoError(t, env.Client().Delete(ctx, claim)) + var canceled nvidiacomv1beta2.DynamoGraphDeploymentRun + awaitClusterTestRun(t, runWatch, "the Run was not canceled with its claim", func(run *nvidiacomv1beta2.DynamoGraphDeploymentRun) bool { + canceled = *run.DeepCopy() + return meta.IsStatusConditionTrue(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + }) + + t.Log("The Job and its pod were gone before the Run failed with TraceSourceDeleted") + var jobs batchv1.JobList + require.NoError(t, env.Client().List(ctx, &jobs, client.InNamespace(env.Namespace()))) + for i := range jobs.Items { + require.False(t, metav1.IsControlledBy(&jobs.Items[i], &canceled), "Job %q of the canceled Run still exists", jobs.Items[i].Name) + } + require.NoError(t, env.Client().List(ctx, &pods, client.InNamespace(env.Namespace()), + client.MatchingLabels{batchv1.JobNameLabel: waiting.Status.JobRef.Name})) + require.Empty(t, pods.Items) + condition = meta.FindStatusCondition(canceled.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + require.Equal(t, nvidiacomv1beta2.DGDRRunPhaseFailed, canceled.Status.Phase) + require.Equal(t, dgdrRunReasonTraceSourceDeleted, condition.Reason) + + t.Log("No Candidate was published, and the completed Run releases the claim") + var candidates nvidiacomv1beta2.DynamoGraphDeploymentCandidateList + require.NoError(t, env.Client().List(ctx, &candidates, client.InNamespace(env.Namespace()))) + require.Empty(t, candidates.Items) + require.Eventually(t, func() bool { + stored := &corev1.PersistentVolumeClaim{} + err := env.Client().Get(ctx, client.ObjectKeyFromObject(claim), stored) + if apierrors.IsNotFound(err) { + return true + } + require.NoError(t, err) + return !slices.Contains(stored.Finalizers, dgdrRunTraceClaimFinalizer(waiting.UID)) + }, time.Minute, time.Second, "the canceled Run did not release its claim") +} + +// awaitClusterTestRun reads Run events from runWatch until done accepts one, failing with message +// after five minutes. No Run may ever succeed: none of them may read a trace outside its claim. +func awaitClusterTestRun( + t *testing.T, + runWatch watch.Interface, + message string, + done func(*nvidiacomv1beta2.DynamoGraphDeploymentRun) bool, +) { + t.Helper() + deadline := time.NewTimer(5 * time.Minute) + defer deadline.Stop() + for { + select { + case event, ok := <-runWatch.ResultChan(): + require.True(t, ok, "Run watch closed: %s", message) + require.NotEqual(t, watch.Error, event.Type, "Run watch failed") + run, ok := event.Object.(*nvidiacomv1beta2.DynamoGraphDeploymentRun) + require.Truef(t, ok, "Run watch returned %T", event.Object) + require.NotEqual(t, nvidiacomv1beta2.DGDRRunPhaseSucceeded, run.Status.Phase, "Run read a trace outside its claim") + if done(run) { + return + } + case <-deadline.C: + t.Fatal(message) + } + } +} + +// startClusterTestDGDRV1beta2Controllers starts the production Request and Run controllers for env. +func startClusterTestDGDRV1beta2Controllers(t *testing.T, env *clusterenv.TestEnv, plannerImage, operatorImage string) { + t.Helper() + operatorVersion, err := clusterTestOperatorVersion() + require.NoError(t, err) + parsedOperatorVersion, err := semver.NewVersion(operatorVersion) + require.NoError(t, err) + dynamoVersion := fmt.Sprintf("%d.%d.%d", parsedOperatorVersion.Major(), parsedOperatorVersion.Minor(), parsedOperatorVersion.Patch()) + + t.Log("Start the production Request, Run, AISimulate Job, and publisher chain") + operatorConfig := clusterTestRestrictedConfig(env.Namespace()) + runtimeConfig := &commoncontroller.RuntimeConfig{Gate: features.Gates{}} + env.StartManager(func(mgr ctrl.Manager) error { + setupOptions := SetupOptions{Config: operatorConfig, RuntimeConfig: runtimeConfig} + if err := SetupDynamoGraphDeploymentRequestV1beta2(mgr, setupOptions); err != nil { + return err + } + return SetupDynamoGraphDeploymentRunV1beta2(mgr, DynamoGraphDeploymentRunSetupOptions{ + SetupOptions: setupOptions, + PlannerImage: plannerImage, + DynamoVersion: dynamoVersion, + OperatorImage: operatorImage, + OperatorImagePullPolicy: corev1.PullIfNotPresent, + }) + }) +} diff --git a/deploy/operator/internal/controller/dgdr_v1beta2_trace_golden_envtest_test.go b/deploy/operator/internal/controller/dgdr_v1beta2_trace_golden_envtest_test.go new file mode 100644 index 000000000000..ceaa601ab181 --- /dev/null +++ b/deploy/operator/internal/controller/dgdr_v1beta2_trace_golden_envtest_test.go @@ -0,0 +1,161 @@ +//go:build !clustertest + +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "bytes" + "encoding/json" + "os" + "testing" + + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" + "github.com/pmezard/go-difflib/difflib" + "sigs.k8s.io/controller-runtime/pkg/client" +) + +// traceRequestLoweringGoldenPath is read by the live-retune planner's tests, which check that the +// planner builds request_spec and accepts run_input as the effective input of its Run. +const traceRequestLoweringGoldenPath = "../aisimulate/testdata/trace_request_lowering.golden.json" + +// liveRetuneShardRequestSpec is the spec of one grid shard the live-retune planner submits: its +// build_request for an aggregated vLLM deployment with 3 replicas of tensor parallelism 2, the window +// [1792996400000, 1793000000000) ms as of 1793000120000 ms, whose selection digest ends the subPath, +// and a 2x2x2 router grid. It omits workload.trace.format, which the CRD schema defaults. Change it +// together with the planner and regenerate the golden file. +const liveRetuneShardRequestSpec = `{ + "model": {"name": "Qwen/Qwen3-32B"}, + "backends": ["vllm"], + "hardware": {"gpu": {"sku": "h200_sxm", "budget": 6}}, + "workload": { + "trace": { + "source": { + "persistentVolumeClaim": { + "claimName": "request-traces", + "subPath": "windows/llm/1793000000000-ba081802aba3.jsonl.gz" + } + } + } + }, + "objective": { + "mode": "Optimize", + "metric": "GoodputPerGPU", + "sla": {"ttftMs": 500.0, "itlMs": 30.0} + }, + "recommendation": {"maxCandidates": 8}, + "aisimulate": { + "inputOverrides": { + "engine": { + "mode": "aggregated", + "backend": "vllm", + "context_length": 32768, + "workers": { + "aggregated": { + "parallelism": { + "preset": [ + {"tensor": 2, "pipeline": 1, "attention_data": 1, "moe_tensor": 1, "moe_expert": 1, "replicas": 3} + ] + } + }, + "prefill": null, + "decode": null + } + }, + "router": { + "policy": "kv_router", + "prefill_load_model": {"type": "none"}, + "overlap_score_credit": {"choices": [0.5, 1.0]}, + "prefill_load_scale": {"choices": [1.0, 2.0]}, + "temperature": {"choices": [0.0, 0.5]} + }, + "optimizer": { + "algorithm": "random", + "max_trials": 16, + "parallelism": 8, + "candidate_timeout_seconds": 300.0 + }, + "optimization": {"strict_sla": false}, + "dgdr_run": {"generator": "direct"} + } + } +}` + +// The live-retune planner submits a trace Request and checks the input its Run executes. Both +// sides share one golden file, so a change to the trace API, admission or lowering that the planner +// does not follow fails here or in the planner's tests. A mismatch writes the result next to the +// golden file as .new; to regenerate the golden file, run the test and move the .new file over it: +// +// go test ./internal/controller -run '^TestDGDRV1beta2TraceRequestLoweringGolden$' +// mv internal/aisimulate/testdata/trace_request_lowering.golden.json.new \ +// internal/aisimulate/testdata/trace_request_lowering.golden.json +func TestDGDRV1beta2TraceRequestLoweringGolden(t *testing.T) { + env := sharedEnv.ForTest(t) + c := env.Client() + ctx := t.Context() + + t.Log("Decode the planner's shard Request; every field must belong to the v1beta2 API") + var spec nvidiacomv1beta2.DynamoGraphDeploymentRequestSpec + decoder := json.NewDecoder(bytes.NewReader([]byte(liveRetuneShardRequestSpec))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&spec); err != nil { + t.Fatalf("decode the planner's Request spec: %v", err) + } + + t.Log("Create the Request through admission and the CRD schema, which complete it") + request := &nvidiacomv1beta2.DynamoGraphDeploymentRequest{Name: "llm-rt-golden", Namespace: env.Namespace(), Spec: spec} + if err := c.Create(ctx, request); err != nil { + t.Fatalf("create the Request: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(request), request); err != nil { + t.Fatalf("get the Request: %v", err) + } + + t.Log("Freeze the admitted Request into its Run as the Request controller does, and store the Run") + run, err := (&DynamoGraphDeploymentRequestV1beta2Reconciler{Client: c}).buildRun(request) + if err != nil { + t.Fatalf("build the Run: %v", err) + } + if err := c.Create(ctx, run); err != nil { + t.Fatalf("create the Run: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(run), run); err != nil { + t.Fatalf("get the Run: %v", err) + } + + t.Log("Lower the stored Run as the Run controller does, in a fixed namespace so the golden file is stable") + run.Namespace = "test-ns" + input, err := dgdrRunInput(run, "example.com/dynamo-planner:1.2.3", "1.2.3") + if err != nil { + t.Fatalf("lower the Run: %v", err) + } + got, err := json.MarshalIndent(struct { + RequestSpec json.RawMessage `json:"request_spec"` + RunInput map[string]any `json:"run_input"` + }{RequestSpec: json.RawMessage(liveRetuneShardRequestSpec), RunInput: input}, "", " ") + if err != nil { + t.Fatalf("encode the golden file: %v", err) + } + got = append(got, '\n') + + t.Log("Compare with the golden file; a mismatch writes the result next to it as .new") + want, err := os.ReadFile(traceRequestLoweringGoldenPath) + if err != nil { + t.Fatalf("read the golden file: %v", err) + } + if !bytes.Equal(got, want) { + newPath := traceRequestLoweringGoldenPath + ".new" + if err := os.WriteFile(newPath, got, 0o644); err != nil { + t.Fatalf("write %s: %v", newPath, err) + } + diff, _ := difflib.GetUnifiedDiffString(difflib.UnifiedDiff{ + A: difflib.SplitLines(string(want)), B: difflib.SplitLines(string(got)), + FromFile: traceRequestLoweringGoldenPath, ToFile: newPath, Context: 3, + }) + t.Fatalf("%s is stale; review %s, move it over the golden file and update the planner if needed:\n%s", + traceRequestLoweringGoldenPath, newPath, diff) + } +} diff --git a/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller.go b/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller.go index 4cd79565f8e1..3fce0c179186 100644 --- a/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller.go +++ b/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller.go @@ -207,6 +207,13 @@ func (r *DynamoGraphDeploymentRequestV1beta2Reconciler) reconcileMissingRun( } marker := meta.FindStatusCondition(request.Status.Conditions, nvidiacomv1beta2.DGDRRequestConditionAccepted) + if request.Status.ActiveRunRef == nil && marker != nil && marker.ObservedGeneration == request.Generation && + marker.Reason == dgdrReasonInvalidSpec { + // The Run CRD rejected this generation's Run; only a new generation can change that. + condition := dgdrV1beta2Condition(request, metav1.ConditionFalse, dgdrReasonInvalidSpec, marker.Message) + return &condition, ctrl.Result{}, nil + } + markerMatches := request.Status.ActiveRunRef != nil && request.Status.ActiveRunRef.Name == wantName && marker != nil && marker.ObservedGeneration == request.Generation if markerMatches && (marker.Reason == dgdrReasonRunCreating || marker.Reason == dgdrReasonRunCreated) { @@ -298,6 +305,17 @@ func (r *DynamoGraphDeploymentRequestV1beta2Reconciler) createReservedRun( original *nvidiacomv1beta2.DynamoGraphDeploymentRequest, ) (ctrl.Result, error) { err := r.Create(ctx, run) + + // The Run CRD rejects the frozen spec, such as a trace that reached the Request without + // admission. Retrying the same spec cannot succeed, so the reservation is released and the + // Request reports the invalid spec until a new generation. + if apierrors.IsInvalid(err) { + log.FromContext(ctx).Info("Run CRD rejects the Request spec", "reason", err.Error()) + request.Status.ActiveRunRef = nil + meta.SetStatusCondition(&request.Status.Conditions, + dgdrV1beta2Condition(request, metav1.ConditionFalse, dgdrReasonInvalidSpec, err.Error())) + return ctrl.Result{}, r.patchStatus(ctx, request, original) + } if err != nil && !apierrors.IsAlreadyExists(err) { return ctrl.Result{}, fmt.Errorf("create Run %q: %w", run.Name, err) } diff --git a/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller_test.go b/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller_test.go index 22c4b313ac45..8c4153d0d754 100644 --- a/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller_test.go +++ b/deploy/operator/internal/controller/dynamographdeploymentrequest_v1beta2_controller_test.go @@ -21,6 +21,7 @@ import ( "context" "errors" "fmt" + "slices" "strings" "testing" @@ -32,7 +33,9 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/validation/field" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" @@ -740,3 +743,66 @@ func TestV1beta2RequestControllerPreservesInputOverridesForRunLowering(t *testin t.Fatalf("Run inputOverrides changed with the Request: %s", got) } } + +func TestV1beta2RequestControllerReportsARunTheCRDRejectsInsteadOfRetryingIt(t *testing.T) { + invalid := apierrors.NewInvalid( + schema.GroupKind{Group: nvidiacomv1beta2.GroupVersion.Group, Kind: "DynamoGraphDeploymentRun"}, "search-1", + field.ErrorList{field.Invalid(field.NewPath("spec", "workload", "trace", "source", "persistentVolumeClaim", "subPath"), "../escape.jsonl", + "subPath must be relative to the claim root without empty, '.' or '..' segments")}, + ) + tests := []struct { + name string + createErr error + wantRuns []string + wantReason string + wantMessage string + }{ + { + name: "a Run the API server accepts is created", + wantRuns: []string{"search-1"}, + wantReason: dgdrReasonRunCreated, + }, + { + name: "a Run the CRD rejects is reported once as an invalid spec", + createErr: invalid, + wantReason: dgdrReasonInvalidSpec, + wantMessage: invalid.Error(), + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + creates := 0 + f := newDGDRV1beta2FixtureWithInterceptor(t, interceptor.Funcs{ + Create: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.CreateOption) error { + if _, ok := obj.(*nvidiacomv1beta2.DynamoGraphDeploymentRun); ok { + creates++ + if tt.createErr != nil { + return tt.createErr + } + } + return c.Create(ctx, obj, opts...) + }, + }, dgdrV1beta2Request(true)) + + t.Log("Reconcile twice: an invalid spec is a condition for this generation, not a create to retry") + f.reconcile() + f.reconcile() + if names := f.runNames(); !slices.Equal(names, tt.wantRuns) { + t.Fatalf("runs = %v, want %v", names, tt.wantRuns) + } + if creates != 1 { + t.Fatalf("Run creates = %d, want 1", creates) + } + if request := f.request(); tt.createErr != nil && request.Status.ActiveRunRef != nil { + t.Fatalf("activeRunRef = %+v, want the reservation released", request.Status.ActiveRunRef) + } + condition := dgdrV1beta2AcceptedCondition(t, f.request()) + if condition.Reason != tt.wantReason { + t.Fatalf("Accepted = %+v, want reason %s", condition, tt.wantReason) + } + if tt.wantMessage != "" && condition.Message != tt.wantMessage { + t.Fatalf("Accepted message = %q, want %q", condition.Message, tt.wantMessage) + } + }) + } +} diff --git a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller.go b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller.go index f9b10fb3b046..fcb829431ec2 100644 --- a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller.go +++ b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller.go @@ -34,15 +34,19 @@ import ( "k8s.io/apimachinery/pkg/types" "k8s.io/utils/ptr" ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/handler" "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/predicate" "sigs.k8s.io/controller-runtime/pkg/reconcile" "sigs.k8s.io/yaml" ) const ( dgdrRunJobControllerUIDIndex = ".metadata.dgdrRunControllerUID" + dgdrRunTraceClaimIndex = ".spec.workload.trace.source.persistentVolumeClaim.claimName" dgdrRunMaxJobNameLength = 63 dgdrRunInputKey = "input.yaml" @@ -63,10 +67,13 @@ const ( dgdrRunPublisherContainer = "publisher" dgdrRunInputVolume = "input" dgdrRunSnapshotVolume = "snapshot" + dgdrRunTraceVolume = "trace" dgdrRunServiceAccountVolume = "publisher-kube-api" dgdrRunSnapshotPathAnnotation = "nvidia.com/dgdr-snapshot-path" dgdrRunResourceComponent = "aisimulate" dgdrRunServiceAccountTokenTTL = int64(3600) + dgdrRunTraceGroup = int64(65532) + dgdrRunLegacyFSGroup = int64(65532) dgdrRunReasonPending = "JobPending" dgdrRunReasonRunning = "JobRunning" dgdrRunReasonSucceeded = "SearchSucceeded" @@ -121,6 +128,7 @@ type DynamoGraphDeploymentRunV1beta2Reconciler struct { // +kubebuilder:rbac:groups=nvidia.com,resources=dynamographdeploymentcandidates/status,verbs=get;update;patch // +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;delete // +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch +// +kubebuilder:rbac:groups="",resources=persistentvolumeclaims,verbs=get;list;watch;patch // +kubebuilder:rbac:groups="",resources=configmaps;serviceaccounts,verbs=get;list;watch;create;update;patch // +kubebuilder:rbac:groups=rbac.authorization.k8s.io,resources=roles;rolebindings,verbs=get;list;watch;create;update;patch @@ -132,9 +140,14 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) Reconcile( if err := r.Get(ctx, req.NamespacedName, run); err != nil { return ctrl.Result{}, client.IgnoreNotFound(err) } - if run.DeletionTimestamp != nil || meta.IsStatusConditionTrue(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) { + if run.DeletionTimestamp != nil { + // The trace claim reconciler releases the claim once the Jobs of the deleted Run are gone. return ctrl.Result{}, nil } + if meta.IsStatusConditionTrue(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) { + return r.retireCompletedRun(ctx, run) + } + if condition := dgdrRunPendingPermanentFailure(run); condition != nil { return r.finishPermanentFailure(ctx, run, condition.Reason, errors.New(condition.Message)) } @@ -157,7 +170,8 @@ func dgdrRunPendingPermanentFailure(run *nvidiacomv1beta2.DynamoGraphDeploymentR return nil } switch condition.Reason { - case dgdrRunReasonInvalidExecution, dgdrRunReasonInvalidResource, commoncontroller.EventReasonOwnershipConflict: + case dgdrRunReasonInvalidExecution, dgdrRunReasonInvalidResource, commoncontroller.EventReasonOwnershipConflict, + dgdrRunReasonTraceSourceUnavailable: return condition default: return nil @@ -178,8 +192,10 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) beginPermanentFailure( Message: cause.Error(), ObservedGeneration: run.Generation, }) + // Like every Run status write, the marker carries the resourceVersion the cache showed, so a + // stale Run cannot overwrite a newer status, such as a terminal cancellation. if !apiequality.Semantic.DeepEqual(original.Status, run.Status) { - if err := r.Status().Patch(ctx, run, client.MergeFrom(original)); err != nil { + if err := r.Status().Patch(ctx, run, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { return ctrl.Result{}, &dgdrRunRetryableError{cause: fmt.Errorf("record permanent failure for Run %q: %w", run.Name, err)} } } @@ -198,16 +214,13 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) finishPermanentFailure( if err != nil { return ctrl.Result{}, err } - stopped, err := r.stopOwnedJobs(ctx, run, jobs, "after permanent failure") + completed, err := r.stopAndCompleteRun(ctx, run, jobs, "after permanent failure", reason, cause.Error()) if err != nil { return ctrl.Result{}, err } - if !stopped { + if !completed { return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil } - if err := r.completeRun(ctx, run, reason, cause.Error()); err != nil { - return ctrl.Result{}, err - } return ctrl.Result{}, reconcile.TerminalError(cause) } @@ -229,8 +242,44 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) reconcileRun( job = &jobs[0] } if supersedingGeneration, cancelRequested := dgdrRunSupersedingGeneration(run); cancelRequested { - return r.cancelSupersededRun(ctx, run, job, supersedingGeneration) + completed, err := r.stopAndCompleteRun(ctx, run, jobs, "for supersession", dgdrRunReasonSuperseded, + fmt.Sprintf("Run was superseded by Request generation %s", supersedingGeneration)) + if err != nil || completed { + return ctrl.Result{}, err + } + return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil } + + // The trace claim decides whether the Run may execute, before anything is created for it. + var claim *corev1.PersistentVolumeClaim + if source := dgdrRunTraceClaim(run); source != nil { + claim, err = r.traceClaim(ctx, run, source) + if err != nil { + return ctrl.Result{}, err + } + + // Deleting the claim cancels the Run whatever its pod has started, unless its Job already + // finished on its own: the Job's own outcome then decides the Run. A Job that is being deleted + // is cancelled even when it finished, because the Job controller fails a Job whose pod the + // deletion removes, and a started cancellation must end as one. A claim that is gone after the Run + // created its Job was protected by the Run, so it was deleted, too. Before the first Job, a + // missing claim keeps the Run open, and only a claim no Job could ever mount fails it. + executed := job != nil || run.Status.JobRef != nil + switch { + case (claim != nil && claim.DeletionTimestamp != nil) || (claim == nil && executed): + if job == nil || !dgdrRunJobFinishedOnItsOwn(job) { + return r.cancelForDeletedTraceClaim(ctx, run, jobs, source) + } + case claim == nil: + return ctrl.Result{}, r.reportMissingTraceClaim(ctx, run, source) + case !executed: + if unusable := dgdrRunUnusableTraceClaim(claim); unusable != "" { + return r.beginPermanentFailure(ctx, run, dgdrRunReasonTraceSourceUnavailable, fmt.Errorf( + "trace PersistentVolumeClaim %q in namespace %q %s", source.ClaimName, run.Namespace, unusable)) + } + } + } + resourceName := dgdrRunResourceName(run.Name) if err := r.reconcilePublisherAccess(ctx, run, resourceName); err != nil { return ctrl.Result{}, err @@ -243,6 +292,13 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) reconcileRun( return ctrl.Result{}, err } + // The Run protects its trace claim before it creates the Job and while the Job runs. claim is + // nil here for a Run without a trace, or for one whose claim is gone after its Job finished. + if claim != nil { + if err := r.protectTraceClaim(ctx, run, claim, job); err != nil { + return ctrl.Result{}, err + } + } if job == nil { job = r.buildJob(run, resourceName, execution) if err := controllerutil.SetControllerReference(run, job, r.Scheme()); err != nil { @@ -271,30 +327,31 @@ func dgdrRunSupersedingGeneration(run *nvidiacomv1beta2.DynamoGraphDeploymentRun return cancelGenerationValue, err == nil && cancelGeneration > runGeneration } -func (r *DynamoGraphDeploymentRunV1beta2Reconciler) cancelSupersededRun( +// stopAndCompleteRun stops every execution of run and, once the API server confirms that none is +// left, completes run as Failed with reason and message. It reports whether run was completed; +// otherwise the caller retries after the stop. jobs are the Jobs of run the cache shows. run must +// not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) stopAndCompleteRun( ctx context.Context, run *nvidiacomv1beta2.DynamoGraphDeploymentRun, - job *batchv1.Job, - supersedingGeneration string, -) (ctrl.Result, error) { - var cached []batchv1.Job - if job != nil { - cached = []batchv1.Job{*job} - } - stopped, err := r.stopOwnedJobs(ctx, run, cached, "for supersession") - if err != nil { - return ctrl.Result{}, err + jobs []batchv1.Job, + action, reason, message string, +) (bool, error) { + stopped, err := r.stopOwnedJobs(ctx, run, jobs, action) + if err != nil || !stopped { + return false, err } - if !stopped { - return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil + if err := r.completeRun(ctx, run, reason, message); err != nil { + return false, err } - - return ctrl.Result{}, r.completeRun(ctx, run, dgdrRunReasonSuperseded, - fmt.Sprintf("Run was superseded by Request generation %s", supersedingGeneration)) + return true, nil } // stopOwnedJobs stops every execution visible through either the informer cache or the API -// server. Completion is safe only after the uncached reader also confirms absence. +// server. Completion is safe only after the uncached reader also confirms absence. Each deletion is +// fenced by the UID and resourceVersion of the Job as it was read, so a Job that changed since, such +// as one that finished meanwhile, is not stopped on a stale decision: the API server reports a +// Conflict, and the next pass decides again from the changed Job. func (r *DynamoGraphDeploymentRunV1beta2Reconciler) stopOwnedJobs( ctx context.Context, run *nvidiacomv1beta2.DynamoGraphDeploymentRun, @@ -316,7 +373,9 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) stopOwnedJobs( if !jobs[i].DeletionTimestamp.IsZero() { continue } - if err := r.Delete(ctx, &jobs[i], client.PropagationPolicy(metav1.DeletePropagationForeground)); err != nil && !apierrors.IsNotFound(err) { + err := r.Delete(ctx, &jobs[i], client.PropagationPolicy(metav1.DeletePropagationForeground), + client.Preconditions{UID: &jobs[i].UID, ResourceVersion: &jobs[i].ResourceVersion}) + if err != nil && !apierrors.IsNotFound(err) && !apierrors.IsConflict(err) { return false, fmt.Errorf("stop Job %q %s of Run %q: %w", jobs[i].Name, action, run.Name, err) } } @@ -342,7 +401,9 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) completeRun( if apiequality.Semantic.DeepEqual(original.Status, run.Status) { return nil } - if err := r.Status().Patch(ctx, run, client.MergeFrom(original)); err != nil { + // The write carries the resourceVersion the cache showed, so a stale Run conflicts instead of + // overwriting a newer status. + if err := r.Status().Patch(ctx, run, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { return &dgdrRunRetryableError{cause: fmt.Errorf("complete Run %q: %w", run.Name, err)} } return nil @@ -377,6 +438,334 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) ownedJobs( return owned, nil } +const ( + dgdrRunReasonTraceSourceUnavailable = "TraceSourceUnavailable" + dgdrRunReasonTraceSourceMissing = "TraceSourceMissing" + dgdrRunReasonTraceSourceDeleted = "TraceSourceDeleted" + dgdrRunReasonContainerWaiting = "ContainerWaiting" + // dgdrRunTraceClaimFinalizerPrefix and the Run UID name the finalizer by which one Run protects + // its trace claim: nvidia.com/dgdr-run-. + dgdrRunTraceClaimFinalizerPrefix = "nvidia.com/dgdr-run-" +) + +// dgdrRunTraceClaimFinalizer returns the finalizer by which the Run with uid protects its trace +// claim. Each Run holds its own, so Runs that share a claim never release each other's protection. +func dgdrRunTraceClaimFinalizer(uid types.UID) string { + return dgdrRunTraceClaimFinalizerPrefix + string(uid) +} + +// traceClaim returns the claim source names, as the cache shows it, or nil when the cache shows no +// such claim. A claim that has not appeared yet keeps the Run pending, and the claim's events wake +// the Run, so the cache answers. run and source must not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) traceClaim( + ctx context.Context, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + source *nvidiacomv1beta2.TracePVCSource, +) (*corev1.PersistentVolumeClaim, error) { + claim := &corev1.PersistentVolumeClaim{} + err := r.Get(ctx, types.NamespacedName{Namespace: run.Namespace, Name: source.ClaimName}, claim) + if apierrors.IsNotFound(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("get trace PersistentVolumeClaim %q in namespace %q: %w", source.ClaimName, run.Namespace, err) + } + return claim, nil +} + +// dgdrRunUnusableTraceClaim returns why a Job could never mount claim read-only as a filesystem, or +// "" when it can. Only structured fields decide: a raw block claim, or a claim that lost its +// volume. A Pending claim is usable: a WaitForFirstConsumer claim binds only once its first pod is +// scheduled. Access modes are not checked: whether another pod holds a ReadWriteOnce or +// ReadWriteOncePod claim is a scheduling question. claim must not be nil. +func dgdrRunUnusableTraceClaim(claim *corev1.PersistentVolumeClaim) string { + switch { + case claim.Spec.VolumeMode != nil && *claim.Spec.VolumeMode == corev1.PersistentVolumeBlock: + return "is a raw block volume, not a filesystem" + case claim.Status.Phase == corev1.ClaimLost: + return "has lost its PersistentVolume" + default: + return "" + } +} + +// reportMissingTraceClaim keeps run open while its trace claim does not exist. The claim's creation +// wakes the Run, which then starts. run and source must not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) reportMissingTraceClaim( + ctx context.Context, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + source *nvidiacomv1beta2.TracePVCSource, +) error { + original := run.DeepCopy() + run.Status.Phase = nvidiacomv1beta2.DGDRRunPhasePending + meta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: nvidiacomv1beta2.DGDRRunConditionCompleted, + Status: metav1.ConditionFalse, + Reason: dgdrRunReasonTraceSourceMissing, + Message: fmt.Sprintf("trace PersistentVolumeClaim %q does not exist in namespace %q; the Run starts once it is created", + source.ClaimName, run.Namespace), + ObservedGeneration: run.Generation, + }) + if apiequality.Semantic.DeepEqual(original.Status, run.Status) { + return nil + } + if err := r.Status().Patch(ctx, run, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { + return &dgdrRunRetryableError{cause: fmt.Errorf("report the missing trace claim of Run %q: %w", run.Name, err)} + } + return nil +} + +// protectTraceClaim adds the finalizer of run to claim, which the cache shows not being deleted. +// Before the Run creates its Job, job is nil and the finalizer is patched even when the cache shows +// it already: the patch carries the resourceVersion the cache showed, so the API server fails it +// when the cached claim is stale, such as when the cache has not shown the claim's deletion yet, and +// a successful patch confirms that the claim is not being deleted. The API server also refuses a new +// finalizer on a claim being deleted. Every failure is retryable: the refused finalizer is reported +// as Invalid, which would otherwise fail the Run permanently, and the next pass reads the claim again +// and cancels the Run if the claim is being deleted. While job runs, a finalizer removed by hand is +// restored; a finished job needs no protection. claim is updated to the stored object. run and claim +// must not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) protectTraceClaim( + ctx context.Context, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + claim *corev1.PersistentVolumeClaim, + job *batchv1.Job, +) error { + if job != nil && (dgdrRunJobFinished(job) || controllerutil.ContainsFinalizer(claim, dgdrRunTraceClaimFinalizer(run.UID))) { + return nil + } + original := claim.DeepCopy() + controllerutil.AddFinalizer(claim, dgdrRunTraceClaimFinalizer(run.UID)) + if err := r.Patch(ctx, claim, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { + return &dgdrRunRetryableError{cause: fmt.Errorf("protect trace PersistentVolumeClaim %q for Run %q: %w", claim.Name, run.Name, err)} + } + return nil +} + +// cancelForDeletedTraceClaim stops the execution of run because its trace claim is being deleted or +// is gone, and then fails run with TraceSourceDeleted. jobs are the Jobs of run the cache shows. +// When the cache shows none, a Job the API server already holds is stopped as well, unless it +// finished on its own: its own outcome decides the Run once the cache shows it. The Run keeps its finalizer on +// the claim until it completes. run and source must not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) cancelForDeletedTraceClaim( + ctx context.Context, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + jobs []batchv1.Job, + source *nvidiacomv1beta2.TracePVCSource, +) (ctrl.Result, error) { + message := fmt.Sprintf("trace PersistentVolumeClaim %q in namespace %q was deleted, which cancels the Run", + source.ClaimName, run.Namespace) + if len(jobs) == 0 { + live, err := r.uncachedOwnedJobs(ctx, run) + if err != nil { + return ctrl.Result{}, err + } + for i := range live { + if dgdrRunJobFinishedOnItsOwn(&live[i]) { + return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil + } + } + + // The API server has just confirmed that no Job is left, so the Run completes without + // listing the Jobs again. + if len(live) == 0 { + return ctrl.Result{}, r.completeRun(ctx, run, dgdrRunReasonTraceSourceDeleted, message) + } + jobs = live + } + + completed, err := r.stopAndCompleteRun(ctx, run, jobs, "for its deleted trace claim", dgdrRunReasonTraceSourceDeleted, message) + if err != nil || completed { + return ctrl.Result{}, err + } + return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil +} + +// retireCompletedRun keeps a completed Run free of live executions and then releases its trace +// claim. Completed=True means that no execution of the Run is left, but a stale pass can still +// create a Job after the Run completed; such an unfinished Job is stopped. Releasing the claim is an +// absence decision the cache cannot make alone, because the Job cache can miss a Job that was just +// created, so the API server confirms that no unfinished Job of the Run is left first. run must not +// be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) retireCompletedRun( + ctx context.Context, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, +) (ctrl.Result, error) { + jobs, err := r.ownedJobs(ctx, run) + if err != nil { + return ctrl.Result{}, err + } + unfinished := make([]batchv1.Job, 0, len(jobs)) + for i := range jobs { + if !dgdrRunJobFinished(&jobs[i]) { + unfinished = append(unfinished, jobs[i]) + } + } + if len(unfinished) > 0 { + if _, err := r.stopOwnedJobs(ctx, run, unfinished, "of a completed Run"); err != nil { + return ctrl.Result{}, err + } + return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil + } + + // Only a trace claim that still carries the Run's finalizer needs the uncached confirmation. + source := dgdrRunTraceClaim(run) + if source == nil { + return ctrl.Result{}, nil + } + claim, err := r.traceClaim(ctx, run, source) + if err != nil || claim == nil || !controllerutil.ContainsFinalizer(claim, dgdrRunTraceClaimFinalizer(run.UID)) { + return ctrl.Result{}, err + } + live, err := r.uncachedOwnedJobs(ctx, run) + if err != nil { + return ctrl.Result{}, err + } + for i := range live { + if !dgdrRunJobFinished(&live[i]) { + return ctrl.Result{RequeueAfter: dgdrRunExecutionCacheRetry}, nil + } + } + + original := claim.DeepCopy() + controllerutil.RemoveFinalizer(claim, dgdrRunTraceClaimFinalizer(run.UID)) + if err := r.Patch(ctx, claim, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { + return ctrl.Result{}, fmt.Errorf("release trace PersistentVolumeClaim %q of Run %q: %w", claim.Name, run.Name, err) + } + return ctrl.Result{}, nil +} + +// dgdrRunTraceClaim returns the claim the Run's trace is read from, or nil when it replays no trace. +// run must not be nil. +func dgdrRunTraceClaim(run *nvidiacomv1beta2.DynamoGraphDeploymentRun) *nvidiacomv1beta2.TracePVCSource { + if run.Spec.Workload == nil || run.Spec.Workload.Trace == nil { + return nil + } + return run.Spec.Workload.Trace.Source.PersistentVolumeClaim +} + +// dgdrRunTraceClaimNames indexes a Run by the claim its trace is read from. +func dgdrRunTraceClaimNames(obj client.Object) []string { + run, ok := obj.(*nvidiacomv1beta2.DynamoGraphDeploymentRun) + if !ok { + return nil + } + if source := dgdrRunTraceClaim(run); source != nil { + return []string{source.ClaimName} + } + return nil +} + +// runsForTraceClaim maps a claim event to the Runs in the claim's namespace that read their trace +// from it and that the claim can still affect: a Run that has not completed, or a completed Run +// whose finalizer the claim still carries and that has yet to release it. Many Runs can share one +// claim, so waking every completed Run on each finalizer change would reconcile them all again +// for nothing. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) runsForTraceClaim(ctx context.Context, claim client.Object) []reconcile.Request { + runs := &nvidiacomv1beta2.DynamoGraphDeploymentRunList{} + if err := r.List(ctx, runs, client.InNamespace(claim.GetNamespace()), + client.MatchingFields{dgdrRunTraceClaimIndex: claim.GetName()}); err != nil { + log.FromContext(ctx).Error(err, "Could not list the Runs that read a trace claim", "claim", claim.GetName()) + return nil + } + requests := make([]reconcile.Request, 0, len(runs.Items)) + for i := range runs.Items { + run := &runs.Items[i] + if meta.IsStatusConditionTrue(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) && + !controllerutil.ContainsFinalizer(claim, dgdrRunTraceClaimFinalizer(run.UID)) { + continue + } + requests = append(requests, reconcile.Request{NamespacedName: client.ObjectKeyFromObject(run)}) + } + return requests +} + +// dgdrRunJobFinished reports whether job has a terminal JobComplete or JobFailed condition. The Job +// controller adds either only after the Job's pod terminated, so a finished Job no longer uses the +// trace claim and its outcome is final. job must not be nil. +func dgdrRunJobFinished(job *batchv1.Job) bool { + for _, conditionType := range []batchv1.JobConditionType{batchv1.JobComplete, batchv1.JobFailed} { + if condition := dgdrRunJobCondition(job, conditionType); condition != nil && condition.Status == corev1.ConditionTrue { + return true + } + } + return false +} + +// dgdrRunJobFinishedOnItsOwn reports whether job finished while nothing was deleting it. A Job +// that is being deleted can still be marked failed, because the Job controller counts the pods the +// deletion removes as failed, so its outcome would only echo the deletion. job must not be nil. +func dgdrRunJobFinishedOnItsOwn(job *batchv1.Job) bool { + return job.DeletionTimestamp.IsZero() && dgdrRunJobFinished(job) +} + +// jobPods returns the pods job controls, as the cache shows them. job must not be nil. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) jobPods(ctx context.Context, job *batchv1.Job) ([]corev1.Pod, error) { + list := &corev1.PodList{} + if err := r.List(ctx, list, client.InNamespace(job.Namespace), client.MatchingLabels{batchv1.JobNameLabel: job.Name}); err != nil { + return nil, fmt.Errorf("list pods of Job %q: %w", job.Name, err) + } + pods := make([]corev1.Pod, 0, len(list.Items)) + for i := range list.Items { + if metav1.IsControlledBy(&list.Items[i], job) { + pods = append(pods, list.Items[i]) + } + } + return pods, nil +} + +// dgdrRunWaitingContainer returns the first sweeper, then publisher, container of pods that waits +// for a reason other than its ordinary start, such as a mount, configuration or image error. The +// result is only a diagnostic: a waiting container may start once its cause is fixed, so it never +// ends the Run. Only the container name and its waiting reason and message are read, which the +// operator's pod cache keeps. +func dgdrRunWaitingContainer(pods []corev1.Pod) (string, corev1.ContainerStateWaiting, bool) { + for _, name := range []string{dgdrRunSweeperContainer, dgdrRunPublisherContainer} { + for i := range pods { + for _, status := range pods[i].Status.ContainerStatuses { + waiting := status.State.Waiting + if status.Name != name || waiting == nil { + continue + } + switch waiting.Reason { + case "", "ContainerCreating", "PodInitializing": + continue + } + return name, *waiting, true + } + } + } + return "", corev1.ContainerStateWaiting{}, false +} + +// dgdrRunJobPod reports whether obj is a pod of a Job named like a Run's AISimulate Job. +func dgdrRunJobPod(obj client.Object) bool { + owner := metav1.GetControllerOf(obj) + return owner != nil && owner.Kind == "Job" && owner.APIVersion == batchv1.SchemeGroupVersion.String() && + strings.HasSuffix(owner.Name, "-"+dgdrRunResourceComponent) +} + +// runForJobPod maps a pod event to the Run that controls the pod's Job. +func (r *DynamoGraphDeploymentRunV1beta2Reconciler) runForJobPod(ctx context.Context, pod client.Object) []reconcile.Request { + owner := metav1.GetControllerOf(pod) + if owner == nil { + return nil + } + job := &batchv1.Job{} + if err := r.Get(ctx, types.NamespacedName{Namespace: pod.GetNamespace(), Name: owner.Name}, job); err != nil { + if !apierrors.IsNotFound(err) { + log.FromContext(ctx).Error(err, "Could not get the Job of a pod", "pod", pod.GetName(), "job", owner.Name) + } + return nil + } + run := metav1.GetControllerOf(job) + if job.UID != owner.UID || run == nil || + run.Kind != "DynamoGraphDeploymentRun" || run.APIVersion != nvidiacomv1beta2.GroupVersion.String() { + return nil + } + return []reconcile.Request{{Namespace: job.Namespace, Name: run.Name}} +} + func (r *DynamoGraphDeploymentRunV1beta2Reconciler) uncachedOwnedJobs( ctx context.Context, run *nvidiacomv1beta2.DynamoGraphDeploymentRun, @@ -794,7 +1183,7 @@ func validateDGDRRunJobExecution( if len(pod.InitContainers) != 0 { return errors.New("init containers are not allowed") } - if err := validateDGDRRunPodSecurityContext(pod.SecurityContext); err != nil { + if err := validateDGDRRunPodSecurityContext(pod.SecurityContext, run); err != nil { return err } if pod.ServiceAccountName != inputConfigMapName || pod.AutomountServiceAccountToken == nil || *pod.AutomountServiceAccountToken { @@ -809,7 +1198,7 @@ func validateDGDRRunJobExecution( if !found { return errors.New("planner image identities differ") } - if err := validateDGDRRunPlannerContainer(&planner, execution); err != nil { + if err := validateDGDRRunPlannerContainer(&planner, run, execution); err != nil { return err } publisher, found := dgdrRunContainer(pod.Containers, dgdrRunPublisherContainer) @@ -819,7 +1208,7 @@ func validateDGDRRunJobExecution( if err := validateDGDRRunPublisherContainer(&publisher, run, execution); err != nil { return err } - if err := validateDGDRRunVolumes(pod.Volumes, inputConfigMapName); err != nil { + if err := validateDGDRRunVolumes(pod.Volumes, run, inputConfigMapName); err != nil { return err } if !apiequality.Semantic.DeepEqual(pod.ImagePullSecrets, execution.ImagePullSecrets) { @@ -828,7 +1217,11 @@ func validateDGDRRunJobExecution( return nil } -func validateDGDRRunPlannerContainer(planner *corev1.Container, execution dgdrRunExecution) error { +func validateDGDRRunPlannerContainer( + planner *corev1.Container, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + execution dgdrRunExecution, +) error { if planner.Image != execution.PlannerImage || planner.ImagePullPolicy != execution.PlannerPullPolicy { return errors.New("planner image identities differ") } @@ -845,7 +1238,7 @@ func validateDGDRRunPlannerContainer(planner *corev1.Container, execution dgdrRu if planner.Lifecycle != nil { return errors.New("planner lifecycle differs") } - if !apiequality.Semantic.DeepEqual(planner.VolumeMounts, dgdrRunSweeperVolumeMounts(execution.SnapshotPath)) { + if !apiequality.Semantic.DeepEqual(planner.VolumeMounts, dgdrRunSweeperVolumeMounts(run, execution.SnapshotPath)) { return errors.New("planner volume mounts differ") } return nil @@ -905,10 +1298,20 @@ func validateDGDRRunJobSpec(spec *batchv1.JobSpec) error { return nil } -func validateDGDRRunPodSecurityContext(security *corev1.PodSecurityContext) error { - if security == nil || security.FSGroup == nil || *security.FSGroup != 65532 { +func validateDGDRRunPodSecurityContext(security *corev1.PodSecurityContext, run *nvidiacomv1beta2.DynamoGraphDeploymentRun) error { + if security == nil { + return errors.New("pod fsGroup differs") + } + + // The pod requests no fsGroup. Earlier operators set fsGroup 65532 on every Run Job; such a Job + // mounts no trace claim, so a Run without a trace keeps accepting it and an operator upgrade does + // not fail its in-flight sweep. + if security.FSGroup != nil && (dgdrRunTraceClaim(run) != nil || *security.FSGroup != dgdrRunLegacyFSGroup) { return errors.New("pod fsGroup differs") } + if !apiequality.Semantic.DeepEqual(security.SupplementalGroups, dgdrRunPodSecurityContext(run).SupplementalGroups) { + return errors.New("pod supplemental groups differ") + } if security.RunAsNonRoot == nil || !*security.RunAsNonRoot { return errors.New("pod runAsNonRoot differs") } @@ -967,8 +1370,12 @@ func validateDGDRRunPublisherEnv(env []corev1.EnvVar) error { return nil } -func validateDGDRRunVolumes(volumes []corev1.Volume, inputConfigMapName string) error { - if len(volumes) != 3 { +func validateDGDRRunVolumes( + volumes []corev1.Volume, + run *nvidiacomv1beta2.DynamoGraphDeploymentRun, + inputConfigMapName string, +) error { + if len(volumes) != len(dgdrRunVolumes(run, inputConfigMapName)) { return errors.New("input, snapshot or publisher credential volumes differ") } input, found := dgdrRunVolume(volumes, dgdrRunInputVolume) @@ -986,6 +1393,14 @@ func validateDGDRRunVolumes(volumes []corev1.Volume, inputConfigMapName string) !apiequality.Semantic.DeepEqual(publisher.Projected.Sources, dgdrRunPublisherKubeAPIVolume().Projected.Sources) { return errors.New("input, snapshot or publisher credential volumes differ") } + + // A Run that replays a trace mounts exactly its claim, read-only. + if source := dgdrRunTraceClaim(run); source != nil { + trace, found := dgdrRunVolume(volumes, dgdrRunTraceVolume) + if !found || !apiequality.Semantic.DeepEqual(trace, dgdrRunTraceClaimVolume(source)) { + return errors.New("trace volume differs") + } + } return nil } @@ -1056,14 +1471,8 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) buildJob( ServiceAccountName: name, AutomountServiceAccountToken: new(false), RestartPolicy: corev1.RestartPolicyNever, - SecurityContext: &corev1.PodSecurityContext{ - FSGroup: ptr.To[int64](65532), - RunAsNonRoot: new(true), - SeccompProfile: &corev1.SeccompProfile{ - Type: corev1.SeccompProfileTypeRuntimeDefault, - }, - }, - ImagePullSecrets: append([]corev1.LocalObjectReference(nil), execution.ImagePullSecrets...), + SecurityContext: dgdrRunPodSecurityContext(run), + ImagePullSecrets: append([]corev1.LocalObjectReference(nil), execution.ImagePullSecrets...), Containers: []corev1.Container{ { Name: dgdrRunSweeperContainer, @@ -1072,7 +1481,7 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) buildJob( Command: dgdrRunSweeperCommand(), Args: dgdrRunSweeperArgs(execution.SnapshotPath), SecurityContext: sweeperSecurityContext, - VolumeMounts: dgdrRunSweeperVolumeMounts(execution.SnapshotPath), + VolumeMounts: dgdrRunSweeperVolumeMounts(run, execution.SnapshotPath), }, { Name: dgdrRunPublisherContainer, @@ -1085,7 +1494,7 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) buildJob( VolumeMounts: dgdrRunPublisherVolumeMounts(execution.SnapshotPath), }, }, - Volumes: dgdrRunVolumes(name), + Volumes: dgdrRunVolumes(run, name), }, }, }, @@ -1127,11 +1536,37 @@ func dgdrRunPublisherEnv() []corev1.EnvVar { } } -func dgdrRunSweeperVolumeMounts(snapshotPath string) []corev1.VolumeMount { - return []corev1.VolumeMount{ +// dgdrRunPodSecurityContext returns the pod security context of run's Job. The pod sets no +// fsGroup, for which kubelet and CSI drivers may chown and chmod a volume recursively; a Run that +// replays a trace gets the trace group as a supplemental group instead, which changes no volume. +// run must not be nil. +func dgdrRunPodSecurityContext(run *nvidiacomv1beta2.DynamoGraphDeploymentRun) *corev1.PodSecurityContext { + security := &corev1.PodSecurityContext{ + RunAsNonRoot: new(true), + SeccompProfile: &corev1.SeccompProfile{ + Type: corev1.SeccompProfileTypeRuntimeDefault, + }, + } + if dgdrRunTraceClaim(run) != nil { + security.SupplementalGroups = []int64{dgdrRunTraceGroup} + } + return security +} + +// dgdrRunSweeperVolumeMounts returns the sweeper's mounts. Only AISimulate reads the trace: it +// mounts the one trace file read-only through a subPath, which the kubelet resolves inside the +// claim, refusing a symlink that resolves outside it. run must not be nil. +func dgdrRunSweeperVolumeMounts(run *nvidiacomv1beta2.DynamoGraphDeploymentRun, snapshotPath string) []corev1.VolumeMount { + mounts := []corev1.VolumeMount{ {Name: dgdrRunInputVolume, MountPath: dgdrRunInputMountPath, ReadOnly: true}, {Name: dgdrRunSnapshotVolume, MountPath: snapshotPath}, } + if source := dgdrRunTraceClaim(run); source != nil { + mounts = append(mounts, corev1.VolumeMount{ + Name: dgdrRunTraceVolume, MountPath: aisimulate.TracePath(source.SubPath), SubPath: source.SubPath, ReadOnly: true, + }) + } + return mounts } func dgdrRunPublisherVolumeMounts(snapshotPath string) []corev1.VolumeMount { @@ -1141,12 +1576,30 @@ func dgdrRunPublisherVolumeMounts(snapshotPath string) []corev1.VolumeMount { } } -func dgdrRunVolumes(inputConfigMapName string) []corev1.Volume { - return []corev1.Volume{ +// dgdrRunVolumes returns the volumes of run's Job, with the trace claim mounted read-only when run +// replays a trace. run must not be nil. +func dgdrRunVolumes(run *nvidiacomv1beta2.DynamoGraphDeploymentRun, inputConfigMapName string) []corev1.Volume { + volumes := []corev1.Volume{ {Name: dgdrRunInputVolume, ConfigMap: &corev1.ConfigMapVolumeSource{Name: inputConfigMapName}}, {Name: dgdrRunSnapshotVolume, EmptyDir: &corev1.EmptyDirVolumeSource{}}, dgdrRunPublisherKubeAPIVolume(), } + if source := dgdrRunTraceClaim(run); source != nil { + volumes = append(volumes, dgdrRunTraceClaimVolume(source)) + } + return volumes +} + +// dgdrRunTraceClaimVolume returns the read-only volume of the claim source names. source must not +// be nil. +func dgdrRunTraceClaimVolume(source *nvidiacomv1beta2.TracePVCSource) corev1.Volume { + return corev1.Volume{ + Name: dgdrRunTraceVolume, + PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ + ClaimName: source.ClaimName, + ReadOnly: true, + }, + } } func dgdrRunContainerSecurityContext() *corev1.SecurityContext { @@ -1213,6 +1666,20 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) reconcileRunStatus( reason, message = dgdrRunReasonRunning, dgdrRunMessageRunning } + // A container of the unfinished Job that waits on an error is reported verbatim. It keeps the + // Run open: the container may start once its cause is fixed, and only the Job's conditions or an + // explicit cancellation end the Run. + if completed == metav1.ConditionFalse { + pods, err := r.jobPods(ctx, job) + if err != nil { + return err + } + if container, waiting, found := dgdrRunWaitingContainer(pods); found { + reason = dgdrRunReasonContainerWaiting + message = fmt.Sprintf("container %q is waiting: %s: %s", container, waiting.Reason, waiting.Message) + } + } + meta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ Type: nvidiacomv1beta2.DGDRRunConditionCompleted, Status: completed, @@ -1220,10 +1687,13 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) reconcileRunStatus( Message: message, ObservedGeneration: run.Generation, }) + + // A stale Run must not overwrite a newer status, such as a terminal cancellation, so the write + // carries the resourceVersion the cache showed. if apiequality.Semantic.DeepEqual(original.Status, run.Status) { return nil } - if err := r.Status().Patch(ctx, run, client.MergeFrom(original)); err != nil { + if err := r.Status().Patch(ctx, run, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { return &dgdrRunRetryableError{cause: fmt.Errorf("update Run %q status: %w", run.Name, err)} } return nil @@ -1355,6 +1825,11 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) SetupWithManager(mgr ctrl.Ma ); err != nil { return fmt.Errorf("index Jobs by Run controller UID: %w", err) } + if err := mgr.GetFieldIndexer().IndexField( + context.Background(), &nvidiacomv1beta2.DynamoGraphDeploymentRun{}, dgdrRunTraceClaimIndex, dgdrRunTraceClaimNames, + ); err != nil { + return fmt.Errorf("index Runs by trace claim: %w", err) + } return ctrl.NewControllerManagedBy(mgr). Named("dynamographdeploymentrun-v1beta2"). For(&nvidiacomv1beta2.DynamoGraphDeploymentRun{}). @@ -1363,6 +1838,9 @@ func (r *DynamoGraphDeploymentRunV1beta2Reconciler) SetupWithManager(mgr ctrl.Ma Owns(&corev1.ServiceAccount{}). Owns(&rbacv1.Role{}). Owns(&rbacv1.RoleBinding{}). + Watches(&corev1.PersistentVolumeClaim{}, handler.EnqueueRequestsFromMapFunc(r.runsForTraceClaim)). + Watches(&corev1.Pod{}, handler.EnqueueRequestsFromMapFunc(r.runForJobPod), + builder.WithPredicates(predicate.NewPredicateFuncs(dgdrRunJobPod))). WithEventFilter(commoncontroller.EphemeralDeploymentEventFilter(r.Config, r.RuntimeConfig)). Complete(r) } diff --git a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_envtest_test.go b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_envtest_test.go index c80688bbec30..d44552c82a05 100644 --- a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_envtest_test.go +++ b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_envtest_test.go @@ -8,15 +8,19 @@ package controller import ( + "context" "errors" "testing" + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" corev1 "k8s.io/api/core/v1" rbacv1 "k8s.io/api/rbac/v1" apiequality "k8s.io/apimachinery/pkg/api/equality" apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" "sigs.k8s.io/controller-runtime/pkg/envtest" ) @@ -117,3 +121,71 @@ func roleAllowsRunFinalizerUpdate(rules []rbacv1.PolicyRule, runName string) boo } return false } + +func TestDGDRRunTraceClaimProtectionIsFencedByTheAPIServer(t *testing.T) { + env := sharedEnv.ForTest(t) + admin := env.Client() + ctx := t.Context() + reconciler := &DynamoGraphDeploymentRunV1beta2Reconciler{Client: admin} + first := &nvidiacomv1beta2.DynamoGraphDeploymentRun{Name: "first", UID: "first-uid"} + second := &nvidiacomv1beta2.DynamoGraphDeploymentRun{Name: "second", UID: "second-uid"} + + t.Log("Create a trace claim that another finalizer keeps while it is deleted") + claim := &corev1.PersistentVolumeClaim{ + Name: "request-traces", Namespace: env.Namespace(), Finalizers: []string{"example.com/hold"}, + Spec: corev1.PersistentVolumeClaimSpec{ + AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteMany}, + Resources: corev1.VolumeResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceStorage: resource.MustParse("1Mi")}}, + }, + } + if err := admin.Create(ctx, claim); err != nil { + t.Fatalf("create claim: %v", err) + } + t.Cleanup(func() { + stored := &corev1.PersistentVolumeClaim{} + if err := admin.Get(context.Background(), client.ObjectKeyFromObject(claim), stored); err != nil { + return + } + stored.Finalizers = nil + if err := admin.Update(context.Background(), stored); client.IgnoreNotFound(err) != nil { + t.Errorf("release claim: %v", err) + } + }) + stale := claim.DeepCopy() + + t.Log("The first Run protects the claim") + if err := reconciler.protectTraceClaim(ctx, first, claim, nil); err != nil { + t.Fatalf("protect claim: %v", err) + } + if !controllerutil.ContainsFinalizer(claim, dgdrRunTraceClaimFinalizer(first.UID)) { + t.Fatalf("finalizers = %v, want the first Run's", claim.Finalizers) + } + protected := claim.DeepCopy() + + t.Log("A protection read from a stale claim conflicts and is retried") + err := reconciler.protectTraceClaim(ctx, second, stale, nil) + if !apierrors.IsConflict(err) || dgdrRunIsPermanentError(err) { + t.Fatalf("stale protection error = %v, want a retryable conflict", err) + } + + t.Log("Delete the claim") + if err := admin.Delete(ctx, claim); err != nil { + t.Fatalf("delete claim: %v", err) + } + + t.Log("A protected claim whose deletion the cache has not shown fails the fence before a Job is created") + err = reconciler.protectTraceClaim(ctx, first, protected, nil) + if !apierrors.IsConflict(err) || dgdrRunIsPermanentError(err) { + t.Fatalf("fence error for a stale protected claim = %v, want a retryable conflict", err) + } + + t.Log("The API server refuses a new finalizer on the claim being deleted, and the refusal is retried") + deleting := &corev1.PersistentVolumeClaim{} + if err := admin.Get(ctx, client.ObjectKeyFromObject(claim), deleting); err != nil { + t.Fatalf("get claim: %v", err) + } + err = reconciler.protectTraceClaim(ctx, second, deleting, nil) + if !apierrors.IsInvalid(err) || dgdrRunIsPermanentError(err) { + t.Fatalf("protection error for a claim being deleted = %v, want a retryable Invalid", err) + } +} diff --git a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_test.go b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_test.go index f97d199f9fe6..a761ed6a5090 100644 --- a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_test.go +++ b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_controller_test.go @@ -9,7 +9,9 @@ import ( "context" "encoding/json" "errors" + "fmt" "reflect" + "slices" "strings" "testing" "time" @@ -27,6 +29,7 @@ import ( "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/validation" "k8s.io/apimachinery/pkg/util/validation/field" clientgoscheme "k8s.io/client-go/kubernetes/scheme" "k8s.io/utils/ptr" @@ -34,6 +37,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/controller-runtime/pkg/client/interceptor" + "sigs.k8s.io/controller-runtime/pkg/event" "sigs.k8s.io/controller-runtime/pkg/reconcile" "sigs.k8s.io/yaml" ) @@ -95,6 +99,7 @@ func newDGDRRunFixtureWithInterceptor( WithScheme(scheme). WithObjects(objects...). WithIndex(&batchv1.Job{}, dgdrRunJobControllerUIDIndex, dgdrRunJobControllerUIDs). + WithIndex(&nvidiacomv1beta2.DynamoGraphDeploymentRun{}, dgdrRunTraceClaimIndex, dgdrRunTraceClaimNames). WithStatusSubresource(&nvidiacomv1beta2.DynamoGraphDeploymentRun{}, &batchv1.Job{}, &corev1.Pod{}). WithInterceptorFuncs(funcs). Build() @@ -856,6 +861,20 @@ func TestValidateDGDRRunJobExecutionRejectsImmutableBindingSkew(t *testing.T) { }, wantErr: "pod seccomp profile differs", }, + { + name: "pod fsGroup", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.SecurityContext.FSGroup = ptr.To[int64](1000) + }, + wantErr: "pod fsGroup differs", + }, + { + name: "pod supplemental group", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.SecurityContext.SupplementalGroups = []int64{65532} + }, + wantErr: "pod supplemental groups differ", + }, { name: "planner privilege escalation", mutate: func(job *batchv1.Job) { @@ -1030,6 +1049,105 @@ func TestValidateDGDRRunJobExecutionRejectsImmutableBindingSkew(t *testing.T) { t.Fatalf("validation error = %v, want API defaults accepted", err) } }) + + t.Run("a Job from an operator that set fsGroup 65532 remains compatible", func(t *testing.T) { + legacy := job.DeepCopy() + legacy.Spec.Template.Spec.SecurityContext.FSGroup = ptr.To[int64](65532) + if err := validateDGDRRunJobExecution(legacy, run, configMap.Name, execution); err != nil { + t.Fatalf("validation error = %v, want the legacy fsGroup accepted for a Run without a trace", err) + } + }) +} + +func TestValidateDGDRRunJobExecutionBindsTheTraceMount(t *testing.T) { + run := dgdrRunWithTraceForTest(t, map[string]any{}) + f := newDGDRRunFixture(t, run, traceClaimForTest()) + f.reconcile() + job := f.job() + configMap := &corev1.ConfigMap{} + if err := f.client.Get(t.Context(), client.ObjectKeyFromObject(job), configMap); err != nil { + t.Fatalf("get execution ConfigMap: %v", err) + } + execution, err := dgdrRunExecutionFromConfigMap(configMap) + if err != nil { + t.Fatalf("decode execution: %v", err) + } + + t.Log("The Job built for a trace Run matches its frozen execution") + if err := validateDGDRRunJobExecution(job, run, configMap.Name, execution); err != nil { + t.Fatalf("validation error = %v, want the built trace Job accepted", err) + } + + t.Log("Any change to the trace claim, its mount or its group is skew") + tests := []struct { + name string + mutate func(*batchv1.Job) + wantErr string + }{ + { + name: "missing trace group", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.SecurityContext.SupplementalGroups = nil + }, + wantErr: "pod supplemental groups differ", + }, + { + name: "fsGroup in place of the trace group", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.SecurityContext.FSGroup = ptr.To[int64](65532) + }, + wantErr: "pod fsGroup differs", + }, + { + name: "writable trace claim", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.Volumes[3].PersistentVolumeClaim.ReadOnly = false + }, + wantErr: "trace volume differs", + }, + { + name: "other trace claim", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.Volumes[3].PersistentVolumeClaim.ClaimName = "other-traces" + }, + wantErr: "trace volume differs", + }, + { + name: "missing trace claim", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.Volumes = job.Spec.Template.Spec.Volumes[:3] + }, + wantErr: "input, snapshot or publisher credential volumes differ", + }, + { + name: "whole claim mounted", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.Containers[0].VolumeMounts[2].SubPath = "" + }, + wantErr: "planner volume mounts differ", + }, + { + name: "trace exposed to the publisher", + mutate: func(job *batchv1.Job) { + job.Spec.Template.Spec.Containers[1].VolumeMounts = append( + job.Spec.Template.Spec.Containers[1].VolumeMounts, + job.Spec.Template.Spec.Containers[0].VolumeMounts[2], + ) + }, + wantErr: "publisher volume mounts differ", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mutated := job.DeepCopy() + tt.mutate(mutated) + err := validateDGDRRunJobExecution(mutated, run, configMap.Name, execution) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("validation error = %v, want %q", err, tt.wantErr) + } + }) + } } func TestV1beta2RunControllerStopsAllJobsBeforePermanentCompletion(t *testing.T) { @@ -1254,6 +1372,9 @@ func assertDGDRRunJobIsolation( !hasMountAt(publisher.VolumeMounts, dgdrRunSnapshotVolume, dgdrRunSnapshotPath) { t.Fatal("sweeper and publisher do not share the frozen snapshot directory") } + if hasMount(publisher.VolumeMounts, dgdrRunTraceVolume) { + t.Fatal("publisher unexpectedly mounts the trace claim") + } } func assertDGDRRunPublisherIdentity( @@ -1460,6 +1581,17 @@ func assertDGDRRunPendingStatus(t *testing.T, stored *nvidiacomv1beta2.DynamoGra } } +func assertDGDRRunRunningStatus(t *testing.T, stored *nvidiacomv1beta2.DynamoGraphDeploymentRun, jobName string) { + t.Helper() + if stored.Status.Phase != nvidiacomv1beta2.DGDRRunPhaseRunning || stored.Status.JobRef == nil || stored.Status.JobRef.Name != jobName { + t.Fatalf("Run status = %+v, want Running and JobRef %q", stored.Status, jobName) + } + condition := meta.FindStatusCondition(stored.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if condition == nil || condition.Status != metav1.ConditionFalse || condition.Reason != dgdrRunReasonRunning { + t.Fatalf("Completed condition = %+v, want False/%s", condition, dgdrRunReasonRunning) + } +} + func TestV1beta2RunControllerOwnsOnlyLifecycleStatus(t *testing.T) { f := newDGDRRunFixture(t, dgdrRunForTest(t, map[string]any{})) f.reconcile() @@ -1675,3 +1807,1421 @@ func hasMountAt(mounts []corev1.VolumeMount, name, path string) bool { } return false } + +func dgdrRunWithTraceForTest(t *testing.T, inputOverrides map[string]any) *nvidiacomv1beta2.DynamoGraphDeploymentRun { + t.Helper() + run := dgdrRunForTest(t, inputOverrides) + run.Spec.Workload = &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Format: nvidiacomv1beta2.TraceFormatDynamo, + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: "request-traces", SubPath: "retune/w-1.jsonl.gz", + }}, + }} + return run +} + +func traceClaimForTest() *corev1.PersistentVolumeClaim { + return &corev1.PersistentVolumeClaim{ + Name: "request-traces", Namespace: "default", + Spec: corev1.PersistentVolumeClaimSpec{ + AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteMany}, + VolumeMode: ptr.To(corev1.PersistentVolumeFilesystem), + }, + Status: corev1.PersistentVolumeClaimStatus{Phase: corev1.ClaimBound}, + } +} + +func TestV1beta2RunControllerMountsTheTraceFileReadOnlyIntoTheSweeperOnly(t *testing.T) { + run := dgdrRunWithTraceForTest(t, map[string]any{}) + f := newDGDRRunFixture(t, run, traceClaimForTest()) + f.reconcile() + + t.Log("The pod mounts the claim read-only as the trace volume") + job := f.job() + var trace *corev1.Volume + for i := range job.Spec.Template.Spec.Volumes { + if job.Spec.Template.Spec.Volumes[i].Name == dgdrRunTraceVolume { + trace = &job.Spec.Template.Spec.Volumes[i] + } + } + if trace == nil || trace.PersistentVolumeClaim == nil || + trace.PersistentVolumeClaim.ClaimName != "request-traces" || !trace.PersistentVolumeClaim.ReadOnly { + t.Fatalf("trace volume = %#v, want the read-only request-traces claim", trace) + } + + t.Log("Only the sweeper mounts the one trace file, read-only through a subPath at the fixed path") + sweeper := containerNamed(t, job.Spec.Template.Spec.Containers, dgdrRunSweeperContainer) + want := corev1.VolumeMount{ + Name: dgdrRunTraceVolume, MountPath: "/var/run/aisimulate/trace/trace.jsonl.gz", + SubPath: "retune/w-1.jsonl.gz", ReadOnly: true, + } + var traceMounts []corev1.VolumeMount + for _, mount := range sweeper.VolumeMounts { + if mount.Name == dgdrRunTraceVolume { + traceMounts = append(traceMounts, mount) + } + } + if !reflect.DeepEqual(traceMounts, []corev1.VolumeMount{want}) { + t.Fatalf("sweeper trace mounts = %#v, want exactly %#v", traceMounts, want) + } + assertDGDRRunJobIsolation(t, job, run, f.reconciler.PlannerImage, f.reconciler.OperatorImage) + + t.Log("The pod grants the trace group as a supplemental group and requests no fsGroup") + security := job.Spec.Template.Spec.SecurityContext + if security == nil || security.FSGroup != nil || !reflect.DeepEqual(security.SupplementalGroups, []int64{65532}) { + t.Fatalf("pod securityContext = %#v, want supplemental group 65532 and no fsGroup", security) + } + + t.Log("The frozen input reads the mounted trace file") + configMap := &corev1.ConfigMap{} + if err := f.client.Get(t.Context(), client.ObjectKeyFromObject(job), configMap); err != nil { + t.Fatalf("get ConfigMap: %v", err) + } + var input map[string]any + if err := yaml.Unmarshal([]byte(configMap.Data[dgdrRunInputKey]), &input); err != nil { + t.Fatalf("decode input: %v", err) + } + source := input["traffic"].(map[string]any)["source"].(map[string]any) + if !reflect.DeepEqual(source["paths"], []any{"/var/run/aisimulate/trace/trace.jsonl.gz"}) { + t.Fatalf("traffic.source.paths = %v", source["paths"]) + } +} + +func TestV1beta2RunControllerRequestsNoPodGroupsWithoutATrace(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunForTest(t, map[string]any{})) + f.reconcile() + + t.Log("A Run without a trace requests neither an fsGroup nor a supplemental group") + security := f.job().Spec.Template.Spec.SecurityContext + if security == nil || security.FSGroup != nil || len(security.SupplementalGroups) != 0 { + t.Fatalf("pod securityContext = %#v, want no fsGroup and no supplemental groups", security) + } +} + +func TestDGDRRunInputLetsOverridesReplaceOrDeleteTheTypedTraceSource(t *testing.T) { + typedSource := map[string]any{"type": "trace", "format": "dynamo", "paths": []any{"/var/run/aisimulate/trace/trace.jsonl.gz"}} + tests := []struct { + name string + overrides map[string]any + wantTraffic any + wantFound bool + }{ + { + name: "no override keeps the typed underlay", + overrides: map[string]any{}, + wantTraffic: map[string]any{"source": typedSource, "load": map[string]any{"type": "trace_timestamps"}}, + wantFound: true, + }, + { + name: "replaced source fields and load are kept", + overrides: map[string]any{"traffic": map[string]any{ + "source": map[string]any{"type": "synthetic", "format": "mooncake", "paths": []any{"/etc/passwd"}}, + "load": map[string]any{"speedup": 4}, + }}, + wantTraffic: map[string]any{ + "source": map[string]any{"type": "synthetic", "format": "mooncake", "paths": []any{"/etc/passwd"}}, + "load": map[string]any{"type": "trace_timestamps", "speedup": json.Number("4")}, + }, + wantFound: true, + }, + { + name: "a null source stays deleted", + overrides: map[string]any{"traffic": map[string]any{"source": nil}}, + wantTraffic: map[string]any{"load": map[string]any{"type": "trace_timestamps"}}, + wantFound: true, + }, + { + name: "a null traffic stays deleted", + overrides: map[string]any{"traffic": nil}, + }, + { + name: "a scalar traffic.source replaces the object", + overrides: map[string]any{"traffic": map[string]any{"source": "invalid"}}, + wantTraffic: map[string]any{"source": "invalid", "load": map[string]any{"type": "trace_timestamps"}}, + wantFound: true, + }, + { + name: "a scalar traffic replaces the object", + overrides: map[string]any{"traffic": "invalid"}, + wantTraffic: "invalid", + wantFound: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Log("Lower a trace Run whose overrides name the traffic the trace lowers") + input, err := dgdrRunInput(dgdrRunWithTraceForTest(t, tt.overrides), "example.com/dynamo-planner:1.2.3", "1.2.3") + if err != nil { + t.Fatalf("lower input: %v", err) + } + + t.Log("The overrides win, and what they delete is not restored") + traffic, found := input["traffic"] + if found != tt.wantFound || !reflect.DeepEqual(traffic, tt.wantTraffic) { + t.Fatalf("traffic = %#v (found %v), want %#v (found %v)", traffic, found, tt.wantTraffic, tt.wantFound) + } + }) + } +} + +func TestV1beta2RunControllerFailsARunWhoseTraceClaimCannotBeMounted(t *testing.T) { + tests := []struct { + name string + mutate func(*corev1.PersistentVolumeClaim) + reason string + }{ + { + name: "raw block claim", + mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.VolumeMode = ptr.To(corev1.PersistentVolumeBlock) + }, + reason: `trace PersistentVolumeClaim "request-traces" in namespace "default" is a raw block volume`, + }, + { + name: "lost claim", + mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Status.Phase = corev1.ClaimLost + }, + reason: `trace PersistentVolumeClaim "request-traces" in namespace "default" has lost its PersistentVolume`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + tt.mutate(claim) + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + request := ctrl.Request{Namespace: "default", Name: "search-1"} + + t.Log("The first pass records the failure marker, the second confirms no Job exists and fails the Run") + result, err := f.reconciler.Reconcile(t.Context(), request) + if err != nil || result.RequeueAfter != dgdrRunExecutionCacheRetry { + t.Fatalf("marker result/error = %+v/%v, want a durable marker and a retry", result, err) + } + if _, err := f.reconciler.Reconcile(t.Context(), request); !errors.Is(err, reconcile.TerminalError(nil)) { + t.Fatalf("reconcile error = %v, want a terminal failure", err) + } + + t.Log("The Run is Failed and Completed with the claim named in the message") + run := f.run() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if run.Status.Phase != nvidiacomv1beta2.DGDRRunPhaseFailed || condition == nil || + condition.Status != metav1.ConditionTrue || condition.Reason != dgdrRunReasonTraceSourceUnavailable { + t.Fatalf("status = %+v, want Failed and Completed=True/%s", run.Status, dgdrRunReasonTraceSourceUnavailable) + } + if !strings.HasPrefix(condition.Message, tt.reason) || run.Status.CompletionTime == nil || run.Status.JobRef != nil { + t.Fatalf("condition message = %q, completionTime = %v, jobRef = %v; want message prefix %q", + condition.Message, run.Status.CompletionTime, run.Status.JobRef, tt.reason) + } + if condition.ObservedGeneration != run.Generation { + t.Fatalf("condition observedGeneration = %d, want %d", condition.ObservedGeneration, run.Generation) + } + + t.Log("Nothing was created for the Run, and the claim was never protected") + assertNoDGDRRunResources(t, f.client) + if stored, found := f.claim(); !found || len(stored.Finalizers) != 0 { + t.Fatalf("claim = %+v (found %t), want it without finalizers", stored, found) + } + }) + } +} + +func TestV1beta2RunControllerKeepsARunOpenUntilItsTraceClaimExists(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{})) + + t.Log("Without the claim the Run stays open with a diagnostic, and nothing is created for it") + f.reconcile() + f.reconcile() + run := f.run() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if run.Status.Phase != nvidiacomv1beta2.DGDRRunPhasePending || condition == nil || + condition.Status != metav1.ConditionFalse || condition.Reason != dgdrRunReasonTraceSourceMissing { + t.Fatalf("status = %+v, want Pending and Completed=False/%s", run.Status, dgdrRunReasonTraceSourceMissing) + } + want := `trace PersistentVolumeClaim "request-traces" does not exist in namespace "default"; the Run starts once it is created` + if condition.Message != want || run.Status.CompletionTime != nil { + t.Fatalf("condition message = %q, completionTime = %v; want %q and no completion", condition.Message, run.Status.CompletionTime, want) + } + assertNoDGDRRunResources(t, f.client) + + t.Log("The claim's creation wakes the Run, which protects the claim and starts") + if err := f.client.Create(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("create claim: %v", err) + } + if got := f.reconciler.runsForTraceClaim(t.Context(), traceClaimForTest()); len(got) != 1 || got[0].Name != "search-1" { + t.Fatalf("claim event requests = %v, want the Run", got) + } + f.reconcile() + assertDGDRRunPendingStatus(t, f.run(), f.job().Name) + assertDGDRRunTraceClaimFinalizers(t, f, dgdrRunTraceClaimFinalizer("run-uid")) +} + +func TestV1beta2RunControllerCreatesNothingForATraceClaimBeingDeleted(t *testing.T) { + claim := traceClaimForTest() + claim.DeletionTimestamp = new(metav1.Now()) + claim.Finalizers = []string{"kubernetes.io/pvc-protection"} + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + + t.Log("Deleting the claim before the first Job cancels the Run like a deletion after it") + f.reconcile() + f.reconcile() + assertDGDRRunCanceledForTrace(t, f.run(), "") + + t.Log("Nothing was created, and no finalizer was added to the claim being deleted") + assertNoDGDRRunResources(t, f.client) + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection") +} + +func assertNoDGDRRunResources(t *testing.T, c client.Client) { + t.Helper() + for _, list := range []client.ObjectList{ + &batchv1.JobList{}, &corev1.ConfigMapList{}, &corev1.ServiceAccountList{}, + &rbacv1.RoleList{}, &rbacv1.RoleBindingList{}, + } { + if err := c.List(t.Context(), list, client.InNamespace("default")); err != nil { + t.Fatalf("list %T: %v", list, err) + } + if n := meta.LenList(list); n != 0 { + t.Fatalf("%T has %d items, want none", list, n) + } + } +} + +// claim returns the stored trace claim and whether it exists. +func (f *dgdrRunFixture) claim() (*corev1.PersistentVolumeClaim, bool) { + f.t.Helper() + claim := &corev1.PersistentVolumeClaim{} + err := f.client.Get(f.t.Context(), types.NamespacedName{Namespace: "default", Name: "request-traces"}, claim) + if apierrors.IsNotFound(err) { + return nil, false + } + if err != nil { + f.t.Fatalf("get claim: %v", err) + } + return claim, true +} + +// assertDGDRRunTraceClaimFinalizers checks that the stored trace claim carries exactly want. +func assertDGDRRunTraceClaimFinalizers(t *testing.T, f *dgdrRunFixture, want ...string) { + t.Helper() + claim, found := f.claim() + if !found { + t.Fatalf("claim is gone, want finalizers %v", want) + } + if !slices.Equal(claim.Finalizers, want) { + t.Fatalf("claim finalizers = %v, want %v", claim.Finalizers, want) + } +} + +func TestDGDRRunTraceClaimFinalizerIsAQualifiedName(t *testing.T) { + t.Log("A Run UID completes the finalizer to a qualified name within the length limits") + finalizer := dgdrRunTraceClaimFinalizer("0f8fad5b-d9cb-469f-a165-70867728950e") + if finalizer != "nvidia.com/dgdr-run-0f8fad5b-d9cb-469f-a165-70867728950e" { + t.Fatalf("finalizer = %q", finalizer) + } + if errs := validation.IsQualifiedName(finalizer); len(errs) != 0 { + t.Fatalf("finalizer %q is not a qualified name: %v", finalizer, errs) + } +} + +func TestV1beta2RunControllerAcceptsAnUnboundOrSharedTraceClaim(t *testing.T) { + tests := []struct { + name string + mutate func(*corev1.PersistentVolumeClaim) + }{ + {name: "WaitForFirstConsumer claim still Pending", mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Status.Phase = corev1.ClaimPending + }}, + {name: "ReadWriteOnce claim", mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.AccessModes = []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce} + }}, + {name: "ReadWriteOncePod claim", mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.AccessModes = []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOncePod} + }}, + {name: "claim without a volume mode", mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.VolumeMode = nil + }}, + {name: "claim protected by Kubernetes and another controller", mutate: func(claim *corev1.PersistentVolumeClaim) { + claim.Finalizers = []string{"kubernetes.io/pvc-protection", "example.com/hold"} + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + tt.mutate(claim) + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + f.reconcile() + assertDGDRRunPendingStatus(t, f.run(), f.job().Name) + + t.Log("The claim keeps its finalizers and gains the Run's") + assertDGDRRunTraceClaimFinalizers(t, f, append(claim.Finalizers, dgdrRunTraceClaimFinalizer("run-uid"))...) + }) + } +} + +func TestV1beta2RunControllerRetriesATransientTraceClaimReadError(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if _, ok := obj.(*corev1.PersistentVolumeClaim); ok { + return errors.New("cache unavailable") + } + return c.Get(ctx, key, obj, opts...) + }, + }) + + t.Log("A read error other than NotFound is returned for a retry") + _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}) + if err == nil || !strings.Contains(err.Error(), "cache unavailable") { + t.Fatalf("reconcile error = %v, want the claim read error", err) + } + + t.Log("The Run stays open, and nothing was created for it") + if run := f.run(); run.Status.Phase != "" || len(run.Status.Conditions) != 0 { + t.Fatalf("status = %+v, want untouched", run.Status) + } + assertNoDGDRRunResources(t, f.client) +} + +func TestV1beta2RunControllerProtectsTheTraceClaimBeforeCreatingTheJob(t *testing.T) { + tests := []struct { + name string + // patchErr, when set, fails the claim patch that protects the claim. + patchErr error + }{ + {name: "the protection is persisted"}, + {name: "the claim patch fails", patchErr: errors.New("API server unavailable")}, + { + name: "the claim patch conflicts with a newer claim", + patchErr: apierrors.NewConflict(schema.GroupResource{Resource: "persistentvolumeclaims"}, "request-traces", + errors.New("the object has been modified")), + }, + { + name: "the API server refuses a new finalizer on a claim being deleted", + patchErr: apierrors.NewInvalid(schema.GroupKind{Kind: "PersistentVolumeClaim"}, "request-traces", field.ErrorList{ + field.Forbidden(field.NewPath("metadata", "finalizers"), "no new finalizers can be added if the object is being deleted"), + }), + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var calls []string + f := newDGDRRunFixtureWithInterceptor(t, interceptor.Funcs{ + Patch: func(ctx context.Context, c client.WithWatch, obj client.Object, patch client.Patch, opts ...client.PatchOption) error { + if _, ok := obj.(*corev1.PersistentVolumeClaim); ok { + calls = append(calls, "protect claim") + if tt.patchErr != nil { + return tt.patchErr + } + } + return c.Patch(ctx, obj, patch, opts...) + }, + Create: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.CreateOption) error { + if _, ok := obj.(*batchv1.Job); ok { + calls = append(calls, "create Job") + } + return c.Create(ctx, obj, opts...) + }, + }, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + + t.Log("Reconcile the Run once") + _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{Namespace: "default", Name: "search-1"}) + + if tt.patchErr == nil { + t.Log("The claim is protected before the Job is created") + if err != nil { + t.Fatalf("reconcile: %v", err) + } + if want := []string{"protect claim", "create Job"}; !reflect.DeepEqual(calls, want) { + t.Fatalf("calls = %v, want %v", calls, want) + } + assertDGDRRunTraceClaimFinalizers(t, f, dgdrRunTraceClaimFinalizer("run-uid")) + return + } + + t.Log("A failed protection is retried: no Job is created and the Run is not failed") + if err == nil || dgdrRunIsPermanentError(err) { + t.Fatalf("reconcile error = %v, want a retryable error", err) + } + if want := []string{"protect claim"}; !reflect.DeepEqual(calls, want) { + t.Fatalf("calls = %v, want %v", calls, want) + } + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want none without protection", len(jobs)) + } + if run := f.run(); len(run.Status.Conditions) != 0 { + t.Fatalf("status = %+v, want untouched", run.Status) + } + }) + } +} + +func TestV1beta2RunControllerCancelsARunWhoseTraceClaimIsDeleted(t *testing.T) { + tests := []struct { + name string + status batchv1.JobStatus + // removeClaim removes the claim completely, as when its finalizers were removed by hand. + removeClaim bool + // failDuringDeletion marks the Job failed while it is being deleted, as the Job controller + // does once the deletion removed its pod. + failDuringDeletion bool + }{ + { + name: "the Job runs", + status: batchv1.JobStatus{Active: 1, Ready: new(int32(1))}, + }, + { + name: "the Job met its success criteria", + status: batchv1.JobStatus{Conditions: []batchv1.JobCondition{{Type: batchv1.JobSuccessCriteriaMet, Status: corev1.ConditionTrue}}}, + }, + { + name: "the Job reached its failure target", + status: batchv1.JobStatus{Conditions: []batchv1.JobCondition{{Type: batchv1.JobFailureTarget, Status: corev1.ConditionTrue}}}, + }, + { + name: "the claim is gone after its finalizers were removed by hand", + status: batchv1.JobStatus{Active: 1, Ready: new(int32(1))}, + removeClaim: true, + }, + { + name: "the Job is marked failed while it is being deleted", + status: batchv1.JobStatus{Active: 1, Ready: new(int32(1))}, + failDuringDeletion: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + claim.Finalizers = []string{"kubernetes.io/pvc-protection"} + var propagation []metav1.DeletionPropagation + f := newDGDRRunFixtureWithInterceptor(t, interceptor.Funcs{ + Delete: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.DeleteOption) error { + if _, ok := obj.(*batchv1.Job); ok { + options := &client.DeleteOptions{} + options.ApplyOptions(opts) + propagation = append(propagation, ptr.Deref(options.PropagationPolicy, "")) + } + return c.Delete(ctx, obj, opts...) + }, + }, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + f.reconcile() + finalizer := dgdrRunTraceClaimFinalizer("run-uid") + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection", finalizer) + + t.Log("The Job reports the row's execution state; a test finalizer stands in for its pods' cleanup") + job := f.job() + job.Finalizers = []string{"example.com/pods-terminating"} + if err := f.client.Update(t.Context(), job); err != nil { + t.Fatalf("hold Job deletion: %v", err) + } + job.Status = tt.status + job.Status.StartTime = new(metav1.Now()) + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("update Job status: %v", err) + } + f.reconcile() + + t.Log("Delete the claim") + if err := f.client.Delete(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("delete claim: %v", err) + } + if tt.removeClaim { + stored, _ := f.claim() + stored.Finalizers = nil + if err := f.client.Update(t.Context(), stored); err != nil { + t.Fatalf("remove claim finalizers: %v", err) + } + } + + t.Log("The Run stops its Job in the foreground and stays open while the Job is deleted") + f.reconcile() + f.reconcile() + if stopped := f.job(); stopped.UID != job.UID || stopped.DeletionTimestamp == nil { + t.Fatalf("Job = %+v, want the Job %q being deleted", stopped.ObjectMeta, job.UID) + } + if !reflect.DeepEqual(propagation, []metav1.DeletionPropagation{metav1.DeletePropagationForeground}) { + t.Fatalf("Job deletions = %v, want one foreground deletion", propagation) + } + if meta.IsStatusConditionTrue(f.run().Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) { + t.Fatalf("status = %+v, want the Run open until its Job is gone", f.run().Status) + } + if !tt.removeClaim { + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection", finalizer) + } + + if tt.failDuringDeletion { + t.Log("A Job marked failed during its deletion is still cancelled, not taken as its own outcome") + deleting := f.job() + deleting.Status.Conditions = []batchv1.JobCondition{ + {Type: batchv1.JobFailureTarget, Status: corev1.ConditionTrue}, + {Type: batchv1.JobFailed, Status: corev1.ConditionTrue}, + } + if err := f.client.Status().Update(t.Context(), deleting); err != nil { + t.Fatalf("fail the deleting Job: %v", err) + } + f.reconcile() + if meta.IsStatusConditionTrue(f.run().Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) { + t.Fatalf("status = %+v, want the Run open until its Job is gone", f.run().Status) + } + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection", finalizer) + } + + t.Log("Once the Job and its pods are gone, the Run fails with TraceSourceDeleted and keeps the claim until then") + job = f.job() + job.Finalizers = nil + if err := f.client.Update(t.Context(), job); err != nil { + t.Fatalf("finish Job deletion: %v", err) + } + f.reconcile() + assertDGDRRunCanceledForTrace(t, f.run(), job.Name) + if !tt.removeClaim { + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection", finalizer) + } + + t.Log("The completed Run releases the claim and never recreates its Job") + f.reconcile() + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want none", len(jobs)) + } + if tt.removeClaim { + if _, found := f.claim(); found { + t.Fatal("claim exists, want it gone") + } + return + } + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection") + }) + } +} + +// createDGDRRunJobPodForTest stores the pod the Job controller would create for job, bound to node +// unless node is empty, with status. +func createDGDRRunJobPodForTest(t *testing.T, c client.Client, job *batchv1.Job, node string, status corev1.PodStatus) { + t.Helper() + pod := &corev1.Pod{ + Name: job.Name + "-pod", Namespace: job.Namespace, + Labels: map[string]string{batchv1.JobNameLabel: job.Name}, + OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(job, batchv1.SchemeGroupVersion.WithKind("Job"))}, + Spec: corev1.PodSpec{NodeName: node}, + } + if err := c.Create(t.Context(), pod); err != nil { + t.Fatalf("create Job pod: %v", err) + } + pod.Status = status + if err := c.Status().Update(t.Context(), pod); err != nil { + t.Fatalf("update Job pod status: %v", err) + } +} + +// assertDGDRRunCanceledForTrace checks that run failed because its trace claim was deleted, with the +// jobRef of jobName, or none when jobName is empty. +func assertDGDRRunCanceledForTrace(t *testing.T, run *nvidiacomv1beta2.DynamoGraphDeploymentRun, jobName string) { + t.Helper() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if run.Status.Phase != nvidiacomv1beta2.DGDRRunPhaseFailed || condition == nil || + condition.Status != metav1.ConditionTrue || condition.Reason != dgdrRunReasonTraceSourceDeleted { + t.Fatalf("status = %+v, want Failed and Completed=True/%s", run.Status, dgdrRunReasonTraceSourceDeleted) + } + want := `trace PersistentVolumeClaim "request-traces" in namespace "default" was deleted, which cancels the Run` + if condition.Message != want || run.Status.CompletionTime == nil { + t.Fatalf("condition message = %q, completionTime = %v; want %q and a completion time", + condition.Message, run.Status.CompletionTime, want) + } + if (jobName == "" && run.Status.JobRef != nil) || (jobName != "" && (run.Status.JobRef == nil || run.Status.JobRef.Name != jobName)) { + t.Fatalf("jobRef = %+v, want %q", run.Status.JobRef, jobName) + } +} + +func TestV1beta2RunControllerLetsAFinishedJobDecideAfterItsTraceClaimIsDeleted(t *testing.T) { + tests := []struct { + name string + condition batchv1.JobConditionType + wantPhase nvidiacomv1beta2.DGDRRunPhase + wantReason string + // hideJob hides the Job from the cached List for the first pass after the deletion. + hideJob bool + // staleJob shows the Job unfinished in the cached List for the first pass after the deletion. + staleJob bool + }{ + {name: "the Job completed", condition: batchv1.JobComplete, + wantPhase: nvidiacomv1beta2.DGDRRunPhaseSucceeded, wantReason: dgdrRunReasonSucceeded}, + {name: "the Job failed", condition: batchv1.JobFailed, + wantPhase: nvidiacomv1beta2.DGDRRunPhaseFailed, wantReason: dgdrRunReasonJobFailed}, + {name: "the Job completed but the cache does not show it yet", condition: batchv1.JobComplete, + wantPhase: nvidiacomv1beta2.DGDRRunPhaseSucceeded, wantReason: dgdrRunReasonSucceeded, hideJob: true}, + {name: "the Job completed but the cache still shows it running", condition: batchv1.JobComplete, + wantPhase: nvidiacomv1beta2.DGDRRunPhaseSucceeded, wantReason: dgdrRunReasonSucceeded, staleJob: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + claim.Finalizers = []string{"kubernetes.io/pvc-protection"} + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + f.reconcile() + + t.Log("The Job finishes, and then the claim is deleted") + job := f.job() + unfinished := job.DeepCopy() + job.Status.StartTime = new(metav1.Now()) + job.Status.Conditions = []batchv1.JobCondition{{Type: tt.condition, Status: corev1.ConditionTrue}} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("finish Job: %v", err) + } + if err := f.client.Delete(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("delete claim: %v", err) + } + + if tt.hideJob { + t.Log("A pass that finds the finished Job only in the API server waits for the cache") + cached := f.reconciler.Client + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok { + return nil + } + return c.List(ctx, list, opts...) + }, + }) + result, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}) + if err != nil || result.RequeueAfter != dgdrRunExecutionCacheRetry { + t.Fatalf("result/error = %+v/%v, want a retry", result, err) + } + f.reconciler.Client = cached + } + + if tt.staleJob { + t.Log("A pass that sees the Job unfinished in a stale cache cannot stop the finished Job") + cached := f.reconciler.Client + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if jobs, ok := list.(*batchv1.JobList); ok { + jobs.Items = []batchv1.Job{*unfinished.DeepCopy()} + return nil + } + return c.List(ctx, list, opts...) + }, + }) + result, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}) + if err != nil || result.RequeueAfter != dgdrRunExecutionCacheRetry { + t.Fatalf("result/error = %+v/%v, want a retry", result, err) + } + f.reconciler.Client = cached + } + + t.Log("The finished Job is kept and decides the Run") + f.reconcile() + if kept := f.job(); kept.UID != job.UID || kept.DeletionTimestamp != nil { + t.Fatalf("Job = %+v, want the finished Job %q kept", kept.ObjectMeta, job.UID) + } + run := f.run() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if run.Status.Phase != tt.wantPhase || condition == nil || condition.Status != metav1.ConditionTrue || + condition.Reason != tt.wantReason { + t.Fatalf("status = %+v, want %s and Completed=True/%s", run.Status, tt.wantPhase, tt.wantReason) + } + + t.Log("The completed Run releases the claim") + f.reconcile() + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection") + }) + } +} + +func TestV1beta2RunControllerReleasesTheTraceClaimOnlyAfterTheRunCompletes(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + f.reconcile() + finalizer := dgdrRunTraceClaimFinalizer("run-uid") + + t.Log("The running Run keeps its protection") + job := f.job() + job.Status = batchv1.JobStatus{Active: 1, Ready: new(int32(1)), StartTime: new(metav1.Now())} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("start Job: %v", err) + } + f.reconcile() + assertDGDRRunRunningStatus(t, f.run(), job.Name) + assertDGDRRunTraceClaimFinalizers(t, f, finalizer) + + t.Log("The Job completes; the pass that records the outcome keeps the protection") + job = f.job() + job.Status.Active, job.Status.Ready = 0, new(int32(0)) + job.Status.Conditions = []batchv1.JobCondition{{Type: batchv1.JobComplete, Status: corev1.ConditionTrue}} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("complete Job: %v", err) + } + f.reconcile() + if f.run().Status.Phase != nvidiacomv1beta2.DGDRRunPhaseSucceeded { + t.Fatalf("status = %+v, want Succeeded", f.run().Status) + } + assertDGDRRunTraceClaimFinalizers(t, f, finalizer) + + t.Log("The completed Run releases the claim and keeps its finished Job") + f.reconcile() + assertDGDRRunTraceClaimFinalizers(t, f) + if kept := f.job(); kept.UID != job.UID { + t.Fatalf("Job UID = %q, want the finished Job %q kept", kept.UID, job.UID) + } +} + +func TestV1beta2RunControllerStopsAnUnfinishedJobOfACompletedRunBeforeReleasingItsClaim(t *testing.T) { + tests := []struct { + name string + // hideJob hides the Job from the cached List until the last step. + hideJob bool + }{ + {name: "the cache shows the Job"}, + {name: "only the API server shows the Job", hideJob: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + f.reconcile() + job := f.job() + finalizer := dgdrRunTraceClaimFinalizer("run-uid") + + t.Log("The Run is completed while its Job still runs, as after a stale pass created the Job") + run := f.run() + meta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: nvidiacomv1beta2.DGDRRunConditionCompleted, Status: metav1.ConditionTrue, + Reason: dgdrRunReasonSuperseded, Message: "Run was superseded by Request generation 2", + }) + run.Status.Phase = nvidiacomv1beta2.DGDRRunPhaseFailed + if err := f.client.Status().Update(t.Context(), run); err != nil { + t.Fatalf("complete Run: %v", err) + } + + if tt.hideJob { + t.Log("While the cache misses the Job, the API server keeps the claim protected") + cached := f.reconciler.Client + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok { + return nil + } + return c.List(ctx, list, opts...) + }, + }) + result, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(run)}) + if err != nil || result.RequeueAfter != dgdrRunExecutionCacheRetry { + t.Fatalf("result/error = %+v/%v, want a retry", result, err) + } + assertDGDRRunTraceClaimFinalizers(t, f, finalizer) + f.reconciler.Client = cached + } + + t.Log("The completed Run stops the Job before it releases the claim") + f.reconcile() + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want the unfinished Job %q stopped", len(jobs), job.UID) + } + assertDGDRRunTraceClaimFinalizers(t, f, finalizer) + f.reconcile() + assertDGDRRunTraceClaimFinalizers(t, f) + if condition := meta.FindStatusCondition(f.run().Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted); condition == nil || + condition.Reason != dgdrRunReasonSuperseded { + t.Fatalf("Completed condition = %+v, want the recorded outcome kept", condition) + } + }) + } +} + +func TestV1beta2RunControllerKeepsAnotherRunsProtectionOfASharedTraceClaim(t *testing.T) { + first := dgdrRunWithTraceForTest(t, map[string]any{}) + second := dgdrRunWithTraceForTest(t, map[string]any{}) + second.Name, second.UID = "search-2", "run-uid-2" + f := newDGDRRunFixture(t, first, second, traceClaimForTest()) + + t.Log("Each Run protects the shared claim with its own finalizer") + f.reconcile() + if _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(second)}); err != nil { + t.Fatalf("reconcile second Run: %v", err) + } + assertDGDRRunTraceClaimFinalizers(t, f, dgdrRunTraceClaimFinalizer("run-uid"), dgdrRunTraceClaimFinalizer("run-uid-2")) + + t.Log("The first Run completes and releases only its own protection") + var job *batchv1.Job + jobs := f.jobs() + for i := range jobs { + if metav1.IsControlledBy(&jobs[i], first) { + job = &jobs[i] + } + } + if job == nil { + t.Fatalf("Jobs = %d, want one of the first Run", len(jobs)) + } + job.Status.StartTime = new(metav1.Now()) + job.Status.Conditions = []batchv1.JobCondition{{Type: batchv1.JobComplete, Status: corev1.ConditionTrue}} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("complete first Job: %v", err) + } + f.reconcile() + f.reconcile() + assertDGDRRunTraceClaimFinalizers(t, f, dgdrRunTraceClaimFinalizer("run-uid-2")) +} + +func TestV1beta2RunControllerNeverRecreatesACanceledJobWhileItsOutcomeIsUnrecorded(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + f.reconcile() + f.reconcile() + jobName := f.job().Name + assertDGDRRunPendingStatus(t, f.run(), jobName) + + t.Log("Delete the claim; the first pass stops the Job") + if err := f.client.Delete(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("delete claim: %v", err) + } + f.reconcile() + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want the Job stopped", len(jobs)) + } + + t.Log("Fail the status patch that records the cancellation") + cached := f.reconciler.Client + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + SubResourcePatch: func(context.Context, client.Client, string, client.Object, client.Patch, ...client.SubResourcePatchOption) error { + return errors.New("injected status patch failure") + }, + }) + if _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}); err == nil { + t.Fatal("reconcile error = nil, want the injected status patch failure") + } + + t.Log("The claim is still being deleted, so no pass creates another Job") + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want none", len(jobs)) + } + assertDGDRRunPendingStatus(t, f.run(), jobName) + + t.Log("The next pass records the cancellation") + f.reconciler.Client = cached + f.reconcile() + assertDGDRRunCanceledForTrace(t, f.run(), jobName) + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want none", len(jobs)) + } +} + +func TestV1beta2RunControllerDoesNotLetAStaleRunUndoATraceCancellation(t *testing.T) { + tests := []struct { + name string + // recreateClaim replaces the deleted claim with a new one under the same name before the stale + // pass, so the stale pass protects it and creates a Job for the completed Run. + recreateClaim bool + }{ + {name: "the claim is still being deleted"}, + {name: "the claim was recreated under the same name", recreateClaim: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + claim.Finalizers = []string{"kubernetes.io/pvc-protection"} + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + f.reconcile() + job := f.job() + + t.Log("The Job runs while the Run is Running") + job.Status = batchv1.JobStatus{Active: 1, Ready: new(int32(1)), StartTime: new(metav1.Now())} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("update Job status: %v", err) + } + f.reconcile() + stale := f.run() + assertDGDRRunRunningStatus(t, stale, job.Name) + + t.Log("The claim is deleted, the Job is stopped and the cancellation is recorded") + if err := f.client.Delete(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("delete claim: %v", err) + } + f.reconcile() + f.reconcile() + canceled := f.run() + assertDGDRRunCanceledForTrace(t, canceled, job.Name) + if tt.recreateClaim { + f.reconcile() + if stored, found := f.claim(); found { + stored.Finalizers = nil + if err := f.client.Update(t.Context(), stored); err != nil { + t.Fatalf("finish claim deletion: %v", err) + } + } + if err := f.client.Create(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("recreate claim: %v", err) + } + } + + t.Log("A pass that still reads the Running Run conflicts instead of overwriting the cancellation") + cached := f.reconciler.Client + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + Get: func(ctx context.Context, c client.WithWatch, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error { + if run, ok := obj.(*nvidiacomv1beta2.DynamoGraphDeploymentRun); ok { + stale.DeepCopyInto(run) + return nil + } + return c.Get(ctx, key, obj, opts...) + }, + }) + _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(stale)}) + if !apierrors.IsConflict(err) { + t.Fatalf("reconcile error = %v, want a status conflict", err) + } + f.reconciler.Client = cached + if got := f.run(); !reflect.DeepEqual(got.Status, canceled.Status) { + t.Fatalf("status = %+v, want the recorded cancellation %+v", got.Status, canceled.Status) + } + wantStaleJobs := 0 + if tt.recreateClaim { + wantStaleJobs = 1 + } + if jobs := f.jobs(); len(jobs) != wantStaleJobs { + t.Fatalf("Jobs = %d after the stale pass, want %d", len(jobs), wantStaleJobs) + } + + t.Log("The completed Run stops any Job the stale pass created, then releases the claim") + f.reconcile() + f.reconcile() + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want none left", len(jobs)) + } + if tt.recreateClaim { + assertDGDRRunTraceClaimFinalizers(t, f) + } else { + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection") + } + if got := f.run(); !reflect.DeepEqual(got.Status, canceled.Status) { + t.Fatalf("status = %+v, want the recorded cancellation %+v", got.Status, canceled.Status) + } + }) + } +} + +func TestV1beta2RunControllerReportsAWaitingContainerWithoutEndingTheRun(t *testing.T) { + tests := []struct { + name string + container string + reason string + message string + // wantReason is the Run's Completed reason; a diagnostic carries the container's state verbatim. + wantReason string + }{ + { + name: "a sweeper whose trace mount cannot be prepared", + container: dgdrRunSweeperContainer, + reason: "CreateContainerConfigError", + message: `failed to prepare subPath for volumeMount "trace" of container "sweeper"`, + wantReason: dgdrRunReasonContainerWaiting, + }, + { + name: "the publisher", + container: dgdrRunPublisherContainer, + reason: "CreateContainerConfigError", + message: `failed to prepare subPath for volumeMount "trace" of container "publisher"`, + wantReason: dgdrRunReasonContainerWaiting, + }, + { + name: "a sweeper image that cannot be pulled", + container: dgdrRunSweeperContainer, + reason: "ErrImagePull", + message: "failed to pull image", + wantReason: dgdrRunReasonContainerWaiting, + }, + { + name: "a container that is being created", + container: dgdrRunSweeperContainer, + reason: "ContainerCreating", + wantReason: dgdrRunReasonRunning, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var deleted bool + f := newDGDRRunFixtureWithInterceptor(t, interceptor.Funcs{ + Delete: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.DeleteOption) error { + deleted = true + return c.Delete(ctx, obj, opts...) + }, + }, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + f.reconcile() + job := f.job() + + t.Log("The scheduled pod's container waits with the row's reason") + job.Status = batchv1.JobStatus{Active: 1, Ready: new(int32(0)), StartTime: new(metav1.Now())} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("update Job status: %v", err) + } + createDGDRRunJobPodForTest(t, f.client, job, "node-a", corev1.PodStatus{ + Phase: corev1.PodPending, + ContainerStatuses: []corev1.ContainerStatus{{ + Name: tt.container, + State: corev1.ContainerState{Waiting: &corev1.ContainerStateWaiting{Reason: tt.reason, Message: tt.message}}, + }}, + }) + + t.Log("Reconcile twice: the Run stays open with the diagnostic and its Job is kept") + f.reconcile() + f.reconcile() + if kept := f.job(); kept.UID != job.UID || kept.DeletionTimestamp != nil || deleted { + t.Fatalf("Job = %+v (deleted %t), want the Job %q kept", kept.ObjectMeta, deleted, job.UID) + } + run := f.run() + condition := meta.FindStatusCondition(run.Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted) + if run.Status.Phase != nvidiacomv1beta2.DGDRRunPhaseRunning || run.Status.CompletionTime != nil || condition == nil || + condition.Status != metav1.ConditionFalse || condition.Reason != tt.wantReason { + t.Fatalf("status = %+v, want Running and Completed=False/%s", run.Status, tt.wantReason) + } + if want := fmt.Sprintf("container %q is waiting: %s: %s", tt.container, tt.reason, tt.message); tt.wantReason == dgdrRunReasonContainerWaiting && + condition.Message != want { + t.Fatalf("condition message = %q, want %q", condition.Message, want) + } + + t.Log("Once the container starts, the Run reports the running Job again") + pod := &corev1.Pod{} + if err := f.client.Get(t.Context(), types.NamespacedName{Namespace: "default", Name: job.Name + "-pod"}, pod); err != nil { + t.Fatalf("get pod: %v", err) + } + pod.Status.ContainerStatuses[0].State = corev1.ContainerState{Running: &corev1.ContainerStateRunning{}} + if err := f.client.Status().Update(t.Context(), pod); err != nil { + t.Fatalf("start container: %v", err) + } + f.reconcile() + assertDGDRRunRunningStatus(t, f.run(), job.Name) + }) + } +} + +func TestV1beta2RunControllerWakesTheRunsThatReadAChangedTraceClaim(t *testing.T) { + reader := dgdrRunWithTraceForTest(t, map[string]any{}) + otherClaim := dgdrRunWithTraceForTest(t, map[string]any{}) + otherClaim.Name = "search-2" + otherClaim.Spec.Workload.Trace.Source.PersistentVolumeClaim.ClaimName = "other-traces" + otherNamespace := dgdrRunWithTraceForTest(t, map[string]any{}) + otherNamespace.Namespace = "elsewhere" + withoutTrace := dgdrRunForTest(t, map[string]any{}) + withoutTrace.Name = "search-3" + completedHolding := dgdrRunWithTraceForTest(t, map[string]any{}) + completedHolding.Name, completedHolding.UID = "search-4", "run-uid-4" + completedReleased := dgdrRunWithTraceForTest(t, map[string]any{}) + completedReleased.Name, completedReleased.UID = "search-5", "run-uid-5" + for _, run := range []*nvidiacomv1beta2.DynamoGraphDeploymentRun{completedHolding, completedReleased} { + meta.SetStatusCondition(&run.Status.Conditions, metav1.Condition{ + Type: nvidiacomv1beta2.DGDRRunConditionCompleted, Status: metav1.ConditionTrue, Reason: dgdrRunReasonSucceeded, + }) + } + f := newDGDRRunFixture(t, reader, otherClaim, otherNamespace, withoutTrace, completedHolding, completedReleased) + + t.Log("An event on request-traces in default wakes the open Run that reads it and the completed Run whose finalizer it still carries") + claim := traceClaimForTest() + claim.Finalizers = []string{dgdrRunTraceClaimFinalizer(completedHolding.UID)} + got := f.reconciler.runsForTraceClaim(t.Context(), claim) + want := []reconcile.Request{ + {NamespacedName: client.ObjectKeyFromObject(reader)}, + {NamespacedName: client.ObjectKeyFromObject(completedHolding)}, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("requests = %v, want %v", got, want) + } +} + +func TestV1beta2RunControllerWakesTheRunOfAJobPod(t *testing.T) { + run := dgdrRunWithTraceForTest(t, map[string]any{}) + f := newDGDRRunFixture(t, run, traceClaimForTest()) + f.reconcile() + job := f.job() + otherJob := &batchv1.Job{Name: "other-aisimulate", Namespace: "default", UID: "other-job"} + if err := f.client.Create(t.Context(), otherJob); err != nil { + t.Fatalf("create Job without a Run: %v", err) + } + tests := []struct { + name string + owner metav1.OwnerReference + wantWatch bool + want []reconcile.Request + }{ + { + name: "a pod of the Run's Job", + owner: *metav1.NewControllerRef(job, batchv1.SchemeGroupVersion.WithKind("Job")), + wantWatch: true, + want: []reconcile.Request{{NamespacedName: client.ObjectKeyFromObject(f.run())}}, + }, + { + name: "a pod of a Job no Run controls", + owner: *metav1.NewControllerRef(otherJob, batchv1.SchemeGroupVersion.WithKind("Job")), + wantWatch: true, + }, + { + name: "a pod of a Job that no longer exists", + owner: metav1.OwnerReference{APIVersion: "batch/v1", Kind: "Job", Name: "gone-aisimulate", UID: "gone", Controller: new(true)}, + wantWatch: true, + }, + { + name: "a pod of a Job with the same name but another UID", + owner: metav1.OwnerReference{APIVersion: "batch/v1", Kind: "Job", Name: job.Name, UID: "replaced", Controller: new(true)}, + wantWatch: true, + }, + { + name: "a pod of a Job named unlike a Run's Job", + owner: metav1.OwnerReference{APIVersion: "batch/v1", Kind: "Job", Name: "nightly", UID: "nightly", Controller: new(true)}, + }, + { + name: "a pod of a ReplicaSet", + owner: metav1.OwnerReference{APIVersion: "apps/v1", Kind: "ReplicaSet", Name: "web-aisimulate", UID: "rs", Controller: new(true)}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pod := &corev1.Pod{Name: "pod", Namespace: "default", OwnerReferences: []metav1.OwnerReference{tt.owner}} + + t.Log("Only pods of Jobs named like a Run's Job pass the watch filter") + if got := dgdrRunJobPod(pod); got != tt.wantWatch { + t.Fatalf("dgdrRunJobPod = %v, want %v", got, tt.wantWatch) + } + + t.Log("Only a pod of the Job a Run controls wakes that Run") + if got := f.reconciler.runForJobPod(t.Context(), pod); !reflect.DeepEqual(got, tt.want) { + t.Fatalf("requests = %v, want %v", got, tt.want) + } + }) + } +} + +func TestV1beta2RunControllerCancelsAJobWhoseJobRefWasNeverRecorded(t *testing.T) { + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), traceClaimForTest()) + cached := f.reconciler.Client + + t.Log("Create the Job but fail the status patch that would record its jobRef") + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + SubResourcePatch: func(context.Context, client.Client, string, client.Object, client.Patch, ...client.SubResourcePatchOption) error { + return errors.New("injected status patch failure") + }, + }) + if _, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}); err == nil { + t.Fatal("reconcile error = nil, want the injected status patch failure") + } + job := f.job() + if f.run().Status.JobRef != nil { + t.Fatalf("jobRef = %+v, want none after the failed status patch", f.run().Status.JobRef) + } + + t.Log("The Job's pod becomes ready") + job.Status = batchv1.JobStatus{Active: 1, Ready: new(int32(1)), StartTime: new(metav1.Now())} + if err := f.client.Status().Update(t.Context(), job); err != nil { + t.Fatalf("update Job status: %v", err) + } + + t.Log("Delete the claim and hide the Job from the cached List; the API server still shows the Job, which is stopped") + if err := f.client.Delete(t.Context(), traceClaimForTest()); err != nil { + t.Fatalf("delete claim: %v", err) + } + f.reconciler.Client = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok { + return nil + } + return c.List(ctx, list, opts...) + }, + }) + result, err := f.reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(f.run())}) + if err != nil || result.RequeueAfter != dgdrRunExecutionCacheRetry { + t.Fatalf("result/error = %+v/%v, want the stop and a retry", result, err) + } + if jobs := f.jobs(); len(jobs) != 0 { + t.Fatalf("Jobs = %d, want the Job stopped", len(jobs)) + } + if condition := meta.FindStatusCondition(f.run().Status.Conditions, nvidiacomv1beta2.DGDRRunConditionCompleted); condition != nil { + t.Fatalf("Completed condition = %+v, want none until the stop is confirmed", condition) + } + + t.Log("Once the API server confirms the Job is gone, the Run is canceled without a second uncached Job list") + apiJobLists := 0 + f.reconciler.APIReader = interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok { + apiJobLists++ + } + return c.List(ctx, list, opts...) + }, + }) + f.reconciler.Client = cached + f.reconcile() + assertDGDRRunCanceledForTrace(t, f.run(), "") + if apiJobLists != 1 { + t.Fatalf("uncached Job lists = %d, want 1", apiJobLists) + } +} + +// newDGDRRunTraceClaimFixture returns a trace claim reconciler over a fake client holding objects, +// whose cached Job List the returned hide flag can make empty, and the count of the Job Lists its +// uncached reader serves. +func newDGDRRunTraceClaimFixture( + t *testing.T, + objects ...client.Object, +) (*DynamoGraphDeploymentRunTraceClaimReconciler, client.Client, *bool, *int) { + t.Helper() + f := newDGDRRunFixture(t, objects...) + hide := new(bool) + cached := interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok && *hide { + return nil + } + return c.List(ctx, list, opts...) + }, + }) + apiJobLists := new(int) + uncached := interceptor.NewClient(f.client.(client.WithWatch), interceptor.Funcs{ + List: func(ctx context.Context, c client.WithWatch, list client.ObjectList, opts ...client.ListOption) error { + if _, ok := list.(*batchv1.JobList); ok { + *apiJobLists++ + } + return c.List(ctx, list, opts...) + }, + }) + return &DynamoGraphDeploymentRunTraceClaimReconciler{Client: cached, APIReader: uncached}, f.client, hide, apiJobLists +} + +func TestDGDRRunTraceClaimReconcilerReleasesOnlyRunsThatCanNoLongerExecute(t *testing.T) { + const otherFinalizer = "kubernetes.io/pvc-protection" + live := dgdrRunWithTraceForTest(t, map[string]any{}) + live.Name, live.UID = "live", "live-uid" + deleting := dgdrRunWithTraceForTest(t, map[string]any{}) + deleting.Name, deleting.UID = "deleting", "deleting-uid" + deleting.DeletionTimestamp = new(metav1.Now()) + deleting.Finalizers = []string{metav1.FinalizerDeleteDependents} + tests := []struct { + name string + // runs exist besides the claim; jobFor lists the UIDs that control a Job. + runs []*nvidiacomv1beta2.DynamoGraphDeploymentRun + jobFor []types.UID + // hideJobs makes the cache miss every Job, as a lagging informer would. + hideJobs bool + uids []types.UID + claimDeleted bool + want []string + }{ + {name: "a Run that exists", runs: []*nvidiacomv1beta2.DynamoGraphDeploymentRun{live}, + uids: []types.UID{"live-uid"}, want: []string{dgdrRunTraceClaimFinalizer("live-uid")}}, + {name: "a deleted Run whose Job remains", jobFor: []types.UID{"gone-uid"}, + uids: []types.UID{"gone-uid"}, want: []string{dgdrRunTraceClaimFinalizer("gone-uid")}}, + {name: "a deleted Run whose Job only the API server shows", jobFor: []types.UID{"gone-uid"}, hideJobs: true, + uids: []types.UID{"gone-uid"}, want: []string{dgdrRunTraceClaimFinalizer("gone-uid")}}, + {name: "a deleted Run whose Job is gone", uids: []types.UID{"gone-uid"}}, + {name: "two deleted Runs whose Jobs are gone", uids: []types.UID{"gone-uid", "other-gone-uid"}}, + {name: "a Run being deleted whose Job remains", runs: []*nvidiacomv1beta2.DynamoGraphDeploymentRun{deleting}, + jobFor: []types.UID{"deleting-uid"}, uids: []types.UID{"deleting-uid"}, + want: []string{dgdrRunTraceClaimFinalizer("deleting-uid")}}, + {name: "a Run being deleted whose Job is gone", runs: []*nvidiacomv1beta2.DynamoGraphDeploymentRun{deleting}, + uids: []types.UID{"deleting-uid"}}, + {name: "a deleted Run and a Run that exists share the claim", runs: []*nvidiacomv1beta2.DynamoGraphDeploymentRun{live}, + uids: []types.UID{"gone-uid", "live-uid"}, want: []string{dgdrRunTraceClaimFinalizer("live-uid")}}, + {name: "a claim being deleted after the operator restarted", uids: []types.UID{"gone-uid"}, claimDeleted: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + claim := traceClaimForTest() + claim.Finalizers = []string{otherFinalizer} + for _, uid := range tt.uids { + claim.Finalizers = append(claim.Finalizers, dgdrRunTraceClaimFinalizer(uid)) + } + if tt.claimDeleted { + claim.DeletionTimestamp = new(metav1.Now()) + } + objects := []client.Object{claim} + for _, run := range tt.runs { + objects = append(objects, run.DeepCopy()) + } + for i, uid := range tt.jobFor { + owner := &nvidiacomv1beta2.DynamoGraphDeploymentRun{Name: string(uid), UID: uid} + objects = append(objects, &batchv1.Job{ + Name: fmt.Sprintf("job-%d-aisimulate", i), Namespace: "default", + OwnerReferences: []metav1.OwnerReference{ + *metav1.NewControllerRef(owner, nvidiacomv1beta2.GroupVersion.WithKind("DynamoGraphDeploymentRun")), + }, + }) + } + reconciler, c, hide, apiJobLists := newDGDRRunTraceClaimFixture(t, objects...) + *hide = tt.hideJobs + + t.Log("Reconcile the claim") + if _, err := reconciler.Reconcile(t.Context(), ctrl.Request{NamespacedName: client.ObjectKeyFromObject(claim)}); err != nil { + t.Fatalf("reconcile: %v", err) + } + + t.Log("Only the finalizers of Runs that can still execute remain, beside foreign ones") + stored := &corev1.PersistentVolumeClaim{} + if err := c.Get(t.Context(), client.ObjectKeyFromObject(claim), stored); err != nil { + t.Fatalf("get claim: %v", err) + } + want := append([]string{otherFinalizer}, tt.want...) + if !reflect.DeepEqual(stored.Finalizers, want) { + t.Fatalf("finalizers = %v, want %v", stored.Finalizers, want) + } + + t.Log("One uncached Job list at most confirms every stale finalizer") + if *apiJobLists > 1 { + t.Fatalf("uncached Job lists = %d, want at most 1", *apiJobLists) + } + }) + } +} + +func TestDGDRRunTraceClaimReconcilerReleasesTheClaimOfADeletedRun(t *testing.T) { + claim := traceClaimForTest() + claim.Finalizers = []string{"kubernetes.io/pvc-protection"} + f := newDGDRRunFixture(t, dgdrRunWithTraceForTest(t, map[string]any{}), claim) + claimReconciler := &DynamoGraphDeploymentRunTraceClaimReconciler{Client: f.client, APIReader: f.client} + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(claim)} + f.reconcile() + job := f.job() + + t.Log("The Run is deleted while its Job still runs; the Job keeps the claim protected") + if err := f.client.Delete(t.Context(), f.run()); err != nil { + t.Fatalf("delete Run: %v", err) + } + if _, err := claimReconciler.Reconcile(t.Context(), request); err != nil { + t.Fatalf("reconcile claim: %v", err) + } + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection", dgdrRunTraceClaimFinalizer("run-uid")) + + t.Log("The Job's deletion event maps to the claim, which is then released") + if got := dgdrRunTraceClaimOfJob(t.Context(), job); !reflect.DeepEqual(got, []reconcile.Request{request}) { + t.Fatalf("Job event requests = %v, want %v", got, []reconcile.Request{request}) + } + if err := f.client.Delete(t.Context(), job); err != nil { + t.Fatalf("delete Job: %v", err) + } + if _, err := claimReconciler.Reconcile(t.Context(), request); err != nil { + t.Fatalf("reconcile claim: %v", err) + } + assertDGDRRunTraceClaimFinalizers(t, f, "kubernetes.io/pvc-protection") +} + +func TestDGDRRunTraceClaimReconcilerWatchesOnlyWhatReleasesAClaim(t *testing.T) { + run := dgdrRunWithTraceForTest(t, map[string]any{}) + withoutTrace := dgdrRunForTest(t, map[string]any{}) + claimRequest := []reconcile.Request{{Namespace: "default", Name: "request-traces"}} + unowned := &batchv1.Job{Name: "nightly", Namespace: "default", Spec: batchv1.JobSpec{Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{Volumes: []corev1.Volume{dgdrRunTraceClaimVolume(run.Spec.Workload.Trace.Source.PersistentVolumeClaim)}}, + }}} + + t.Log("Claims pass the filter only while a Run protects them") + protected := traceClaimForTest() + protected.Finalizers = []string{"kubernetes.io/pvc-protection", dgdrRunTraceClaimFinalizer("run-uid")} + if !dgdrRunProtectsTraceClaim(protected) || dgdrRunProtectsTraceClaim(traceClaimForTest()) { + t.Fatal("only a claim with a Run finalizer passes the filter") + } + + t.Log("Run and Job deletions map to the claim the Run reads") + if got := dgdrRunTraceClaimOfRun(t.Context(), run); !reflect.DeepEqual(got, claimRequest) { + t.Fatalf("Run requests = %v, want %v", got, claimRequest) + } + if got := dgdrRunTraceClaimOfRun(t.Context(), withoutTrace); len(got) != 0 { + t.Fatalf("Run without a trace requests = %v, want none", got) + } + if got := dgdrRunTraceClaimOfJob(t.Context(), unowned); len(got) != 0 { + t.Fatalf("Job no Run controls requests = %v, want none", got) + } + + t.Log("Only deletions pass the Run and Job event filter") + if dgdrRunDeleteEvents.Create(event.CreateEvent{Object: run}) || dgdrRunDeleteEvents.Update(event.UpdateEvent{ObjectOld: run, ObjectNew: run}) || + !dgdrRunDeleteEvents.Delete(event.DeleteEvent{Object: run}) || dgdrRunDeleteEvents.Generic(event.GenericEvent{Object: run}) { + t.Fatal("the event filter must pass deletions only") + } +} diff --git a/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_trace_claim_controller.go b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_trace_claim_controller.go new file mode 100644 index 000000000000..884bd23a1857 --- /dev/null +++ b/deploy/operator/internal/controller/dynamographdeploymentrun_v1beta2_trace_claim_controller.go @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package controller + +import ( + "context" + "fmt" + "slices" + "strings" + + configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1" + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" + commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common" + batchv1 "k8s.io/api/batch/v1" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/predicate" + "sigs.k8s.io/controller-runtime/pkg/reconcile" +) + +// DynamoGraphDeploymentRunTraceClaimReconciler releases the finalizers by which Runs protect their +// trace claims when the Run reconciler cannot: a Run has no finalizer of its own, so it can be +// deleted, by its Request or with its namespace, before it released its claim. The claim keeps a +// Run's finalizer while that Run exists and is not being deleted, or while a Job the Run controls +// exists. The Run reconciler releases the finalizer of a completed Run itself. +// +// For a deleted Run, no Job does not mean no pod: garbage collection deletes the Run's Job in the +// background, before the Job's pod has stopped, so the finalizer can be released while that pod +// still has the trace file mounted. Kubernetes PersistentVolumeClaim protection, not this +// finalizer, keeps the claim until the pod is gone. +type DynamoGraphDeploymentRunTraceClaimReconciler struct { + client.Client + APIReader client.Reader + Config *configv1alpha1.OperatorConfiguration + RuntimeConfig *commoncontroller.RuntimeConfig +} + +func (r *DynamoGraphDeploymentRunTraceClaimReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + claim := &corev1.PersistentVolumeClaim{} + if err := r.Get(ctx, req.NamespacedName, claim); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + + // The finalizers of Runs that the cache shows gone and without a Job. + stale, err := r.staleTraceClaimFinalizers(ctx, claim) + if err != nil || len(stale) == 0 { + return ctrl.Result{}, err + } + + // The Job cache can miss a Job a Run created just before it was deleted, and releasing the claim + // cannot be undone, so one uncached list of the namespace's Jobs confirms every stale finalizer. + live := &batchv1.JobList{} + if err := r.APIReader.List(ctx, live, client.InNamespace(claim.Namespace)); err != nil { + return ctrl.Result{}, fmt.Errorf("list the Jobs of the Runs that read trace PersistentVolumeClaim %q from the API: %w", + claim.Name, err) + } + for i := range live.Items { + for _, uid := range dgdrRunJobControllerUIDs(&live.Items[i]) { + delete(stale, dgdrRunTraceClaimFinalizer(types.UID(uid))) + } + } + + // Drop the finalizer of every Run that can no longer execute against the claim. + original := claim.DeepCopy() + for finalizer := range stale { + controllerutil.RemoveFinalizer(claim, finalizer) + } + if len(claim.Finalizers) == len(original.Finalizers) { + return ctrl.Result{}, nil + } + if err := r.Patch(ctx, claim, client.MergeFromWithOptions(original, client.MergeFromWithOptimisticLock{})); err != nil { + return ctrl.Result{}, fmt.Errorf("release trace PersistentVolumeClaim %q: %w", claim.Name, err) + } + return ctrl.Result{}, nil +} + +// staleTraceClaimFinalizers returns the Run finalizers on claim that the cache no longer justifies: +// the Run is gone or being deleted, and the cache shows no Job it controls. A Run never appears after +// its finalizer, so the cache answers for the Run; only the Jobs need the API server's confirmation, +// which the caller makes once for all of them. claim must not be nil. +func (r *DynamoGraphDeploymentRunTraceClaimReconciler) staleTraceClaimFinalizers( + ctx context.Context, + claim *corev1.PersistentVolumeClaim, +) (map[string]struct{}, error) { + runs := &nvidiacomv1beta2.DynamoGraphDeploymentRunList{} + if err := r.List(ctx, runs, client.InNamespace(claim.Namespace), + client.MatchingFields{dgdrRunTraceClaimIndex: claim.Name}); err != nil { + return nil, fmt.Errorf("list the Runs that read trace PersistentVolumeClaim %q: %w", claim.Name, err) + } + active := make(map[string]struct{}, len(runs.Items)) + for i := range runs.Items { + if runs.Items[i].DeletionTimestamp == nil { + active[dgdrRunTraceClaimFinalizer(runs.Items[i].UID)] = struct{}{} + } + } + + stale := map[string]struct{}{} + for _, finalizer := range claim.Finalizers { + uid, ok := strings.CutPrefix(finalizer, dgdrRunTraceClaimFinalizerPrefix) + if _, held := active[finalizer]; !ok || held { + continue + } + + // A Job of the Run as the cache shows it. + jobs := &batchv1.JobList{} + if err := r.List(ctx, jobs, client.InNamespace(claim.Namespace), + client.MatchingFields{dgdrRunJobControllerUIDIndex: uid}); err != nil { + return nil, fmt.Errorf("list the Jobs of Run %q: %w", uid, err) + } + if len(jobs.Items) == 0 { + stale[finalizer] = struct{}{} + } + } + return stale, nil +} + +// dgdrRunProtectsTraceClaim reports whether obj carries the finalizer of a Run. +func dgdrRunProtectsTraceClaim(obj client.Object) bool { + return slices.ContainsFunc(obj.GetFinalizers(), func(finalizer string) bool { + return strings.HasPrefix(finalizer, dgdrRunTraceClaimFinalizerPrefix) + }) +} + +// dgdrRunTraceClaimOfRun maps a Run event to the trace claim the Run reads, if any. +func dgdrRunTraceClaimOfRun(_ context.Context, obj client.Object) []reconcile.Request { + run, ok := obj.(*nvidiacomv1beta2.DynamoGraphDeploymentRun) + if !ok { + return nil + } + source := dgdrRunTraceClaim(run) + if source == nil { + return nil + } + return []reconcile.Request{{Namespace: run.Namespace, Name: source.ClaimName}} +} + +// dgdrRunTraceClaimOfJob maps an event of a Job a Run controls to the trace claim the Job mounts, if +// any. +func dgdrRunTraceClaimOfJob(_ context.Context, obj client.Object) []reconcile.Request { + job, ok := obj.(*batchv1.Job) + if !ok || len(dgdrRunJobControllerUIDs(job)) == 0 { + return nil + } + trace, found := dgdrRunVolume(job.Spec.Template.Spec.Volumes, dgdrRunTraceVolume) + if !found || trace.PersistentVolumeClaim == nil { + return nil + } + return []reconcile.Request{{Namespace: job.Namespace, Name: trace.PersistentVolumeClaim.ClaimName}} +} + +// dgdrRunDeleteEvents passes only deletions: a Run or Job releases a claim only by going away, and +// claim events cover every other state, including the initial list after a restart. +var dgdrRunDeleteEvents = predicate.Funcs{ + CreateFunc: func(event.CreateEvent) bool { return false }, + UpdateFunc: func(event.UpdateEvent) bool { return false }, + DeleteFunc: func(event.DeleteEvent) bool { return true }, + GenericFunc: func(event.GenericEvent) bool { return false }, +} + +// SetupWithManager registers the reconciler. It relies on the Run and Job indexes that the Run +// reconciler registers, so it is set up after that reconciler. +func (r *DynamoGraphDeploymentRunTraceClaimReconciler) SetupWithManager(mgr ctrl.Manager) error { + if r.APIReader == nil { + r.APIReader = mgr.GetAPIReader() + } + return ctrl.NewControllerManagedBy(mgr). + Named("dynamographdeploymentrun-v1beta2-trace-claim"). + For(&corev1.PersistentVolumeClaim{}, builder.WithPredicates(predicate.NewPredicateFuncs(dgdrRunProtectsTraceClaim))). + Watches(&nvidiacomv1beta2.DynamoGraphDeploymentRun{}, handler.EnqueueRequestsFromMapFunc(dgdrRunTraceClaimOfRun), + builder.WithPredicates(dgdrRunDeleteEvents)). + Watches(&batchv1.Job{}, handler.EnqueueRequestsFromMapFunc(dgdrRunTraceClaimOfJob), + builder.WithPredicates(dgdrRunDeleteEvents)). + WithEventFilter(commoncontroller.EphemeralDeploymentEventFilter(r.Config, r.RuntimeConfig)). + Complete(r) +} diff --git a/deploy/operator/internal/controller/pod_cache_contract_test.go b/deploy/operator/internal/controller/pod_cache_contract_test.go index 91f4c9422fd7..6cdea36e8d6a 100644 --- a/deploy/operator/internal/controller/pod_cache_contract_test.go +++ b/deploy/operator/internal/controller/pod_cache_contract_test.go @@ -138,6 +138,46 @@ func TestProjectedPodSupportsControllerContract(t *testing.T) { assert.Contains(t, detail, "profiling failed") }) + t.Run("DGDR Run diagnostics", func(t *testing.T) { + t.Log("Project a pod of a Run's Job whose sweeper waits on a mount error") + job := &batchv1.Job{Name: "search-1-aisimulate", Namespace: pod.Namespace, UID: types.UID("job-uid")} + runPod := podcache.Project(&corev1.Pod{ + Name: "search-1-aisimulate-abcde", + Namespace: pod.Namespace, + Labels: map[string]string{batchv1.JobNameLabel: job.Name}, + OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(job, batchv1.SchemeGroupVersion.WithKind("Job"))}, + Spec: corev1.PodSpec{ + NodeName: "node-a", + Containers: []corev1.Container{{Name: dgdrRunSweeperContainer}, {Name: dgdrRunPublisherContainer}}, + }, + Status: corev1.PodStatus{ + Phase: corev1.PodPending, + ContainerStatuses: []corev1.ContainerStatus{ + {Name: dgdrRunPublisherContainer, State: corev1.ContainerState{Running: &corev1.ContainerStateRunning{}}}, + {Name: dgdrRunSweeperContainer, State: corev1.ContainerState{Waiting: &corev1.ContainerStateWaiting{ + Reason: "CreateContainerConfigError", Message: `failed to prepare subPath for volumeMount "trace"`, + }}}, + }, + }, + }) + + t.Log("The Run controller finds the pod of its Job and reports the waiting sweeper verbatim") + scheme := runtime.NewScheme() + require.NoError(t, corev1.AddToScheme(scheme)) + reconciler := &DynamoGraphDeploymentRunV1beta2Reconciler{ + Client: fake.NewClientBuilder().WithScheme(scheme).WithObjects(runPod).Build(), + } + pods, err := reconciler.jobPods(context.Background(), job) + require.NoError(t, err) + require.Len(t, pods, 1) + container, waiting, found := dgdrRunWaitingContainer(pods) + assert.True(t, found) + assert.Equal(t, dgdrRunSweeperContainer, container) + assert.Equal(t, corev1.ContainerStateWaiting{ + Reason: "CreateContainerConfigError", Message: `failed to prepare subPath for volumeMount "trace"`, + }, waiting) + }) + t.Run("Recreate", func(t *testing.T) { // PR #11909 lists old worker Pods by these labels and treats only // Failed/Succeeded phases (or cache disappearance) as terminated. diff --git a/deploy/operator/internal/controller/setup.go b/deploy/operator/internal/controller/setup.go index 56355c42473d..d11d48583225 100644 --- a/deploy/operator/internal/controller/setup.go +++ b/deploy/operator/internal/controller/setup.go @@ -156,7 +156,8 @@ func SetupDynamoGraphDeploymentRequestV1beta2(mgr ctrl.Manager, opts SetupOption } // SetupDynamoGraphDeploymentRunV1beta2 registers the controller that executes immutable -// v1beta2 Runs as AISimulate Jobs and observes their lifecycle. +// v1beta2 Runs as AISimulate Jobs and observes their lifecycle, and the controller that releases +// the trace claims of deleted Runs. func SetupDynamoGraphDeploymentRunV1beta2(mgr ctrl.Manager, opts DynamoGraphDeploymentRunSetupOptions) error { if err := (&DynamoGraphDeploymentRunV1beta2Reconciler{ Client: mgr.GetClient(), @@ -170,6 +171,13 @@ func SetupDynamoGraphDeploymentRunV1beta2(mgr ctrl.Manager, opts DynamoGraphDepl }).SetupWithManager(mgr); err != nil { return fmt.Errorf("unable to create DynamoGraphDeploymentRun v1beta2 controller: %w", err) } + if err := (&DynamoGraphDeploymentRunTraceClaimReconciler{ + Client: mgr.GetClient(), + Config: opts.Config, + RuntimeConfig: opts.RuntimeConfig, + }).SetupWithManager(mgr); err != nil { + return fmt.Errorf("unable to create DynamoGraphDeploymentRun v1beta2 trace claim controller: %w", err) + } return nil } diff --git a/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/input.yaml b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/input.yaml new file mode 100644 index 000000000000..5fe13fc54714 --- /dev/null +++ b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/input.yaml @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# The test writes trace.jsonl, gzip-compressed, to retune/w-1.jsonl.gz on the request-traces claim +# before it applies this Request. +apiVersion: nvidia.com/v1beta2 +kind: DynamoGraphDeploymentRequest +metadata: + name: aisimulate-vllm-trace +spec: + model: + name: Qwen/Qwen3-32B + backends: + - vllm + workload: + trace: + # format defaults to Dynamo. + source: + persistentVolumeClaim: + claimName: request-traces + subPath: retune/w-1.jsonl.gz + replay: + arrivalSpeedupRatio: 2.0 + objective: + mode: Optimize + metric: ThroughputPerGPU + sla: {} + aisimulate: + inputOverrides: + engine: + mode: aggregated + context_length: 32768 + workers: + # TODO(#17581): Remove this explicit value after AISimulate preserves an unset runtime block size. + aggregated: + kv_cache: + block_size: 64 + prefill: null + decode: null + optimizer: + algorithm: random + max_trials: 3 + parallelism: 1 + candidate_timeout_seconds: 120 + dgdr_run: + snapshot_interval_seconds: 0 + recommendation: + maxCandidates: 1 diff --git a/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/output.yaml b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/output.yaml new file mode 100644 index 000000000000..81ce298dedb0 --- /dev/null +++ b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/output.yaml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +apiVersion: nvidia.com/v1beta2 +kind: DynamoGraphDeploymentCandidate +metadata: + $strict: false + name: $glob:aisimulate-vllm-trace-1-*-evaluated-point-* + labels: + $strict: false + nvidia.com/dgdr-candidate-id: $pattern:^evaluated-point-[0-9a-f]{12}$ + nvidia.com/dgdr-run-uid: $ignore +spec: + $strict: false + parameters: + $strict: false + backend: vllm +status: + metrics: + $strict: false + score: $ignore + usedGpus: $ignore diff --git a/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/trace.jsonl b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/trace.jsonl new file mode 100644 index 000000000000..e7b9ff7a1b9f --- /dev/null +++ b/deploy/operator/internal/controller/testdata/dgdr-v1beta2/vllm-trace/trace.jsonl @@ -0,0 +1,40 @@ +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000001147, "request": {"request_id": "r0", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000000247, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 256, "input_sequence_hashes": [3000, 3001, 1000000000, 1000000001]}, "output_tokens": 49}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000001458, "request": {"request_id": "r1", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000000558, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 640, "input_sequence_hashes": [3000, 3001, 3002, 3003, 3004, 1000000100, 1000000101, 1000000102, 1000000103, 1000000104]}, "output_tokens": 54}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000001752, "request": {"request_id": "r2", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000000852, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 2005, 1000000200, 1000000201, 1000000202, 1000000203, 1000000204, 1000000205, 1000000206]}, "output_tokens": 43}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000002060, "request": {"request_id": "r3", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000001160, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 512, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1000000300, 1000000301, 1000000302, 1000000303]}, "output_tokens": 33}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000002158, "request": {"request_id": "r4", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000001258, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 768, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 2005, 1000000400, 1000000401, 1000000402, 1000000403, 1000000404, 1000000405]}, "output_tokens": 93}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000002283, "request": {"request_id": "r5", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000001383, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [2000, 2001, 1000000500, 1000000501, 1000000502]}, "output_tokens": 25}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000002683, "request": {"request_id": "r6", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000001783, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 704, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 1000000600, 1000000601, 1000000602, 1000000603, 1000000604, 1000000605]}, "output_tokens": 87}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000002784, "request": {"request_id": "r7", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000001884, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 640, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 1000000700, 1000000701, 1000000702, 1000000703, 1000000704]}, "output_tokens": 56}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000003146, "request": {"request_id": "r8", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000002246, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 768, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1000000800, 1000000801, 1000000802, 1000000803, 1000000804, 1000000805]}, "output_tokens": 77}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000003422, "request": {"request_id": "r9", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000002522, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 256, "input_sequence_hashes": [2000, 2001, 1000000900, 1000000901]}, "output_tokens": 86}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000003479, "request": {"request_id": "r10", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000002579, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 960, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 6, 1000001000, 1000001001, 1000001002, 1000001003, 1000001004, 1000001005, 1000001006, 1000001007]}, "output_tokens": 67}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000003871, "request": {"request_id": "r11", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000002971, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 1000001100, 1000001101, 1000001102, 1000001103, 1000001104, 1000001105, 1000001106]}, "output_tokens": 79}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000004091, "request": {"request_id": "r12", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000003191, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 960, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1000001200, 1000001201, 1000001202, 1000001203, 1000001204, 1000001205, 1000001206, 1000001207]}, "output_tokens": 57}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000004173, "request": {"request_id": "r13", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000003273, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1000001300, 1000001301, 1000001302, 1000001303, 1000001304, 1000001305, 1000001306]}, "output_tokens": 44}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000004345, "request": {"request_id": "r14", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000003445, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 1024, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1000001400, 1000001401, 1000001402, 1000001403, 1000001404, 1000001405, 1000001406, 1000001407]}, "output_tokens": 85}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000004624, "request": {"request_id": "r15", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000003724, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [0, 1, 1000001500, 1000001501, 1000001502]}, "output_tokens": 56}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000004934, "request": {"request_id": "r16", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000004034, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [3000, 3001, 1000001600, 1000001601, 1000001602]}, "output_tokens": 54}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000005266, "request": {"request_id": "r17", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000004366, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 960, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 2005, 2006, 1000001700, 1000001701, 1000001702, 1000001703, 1000001704, 1000001705, 1000001706, 1000001707]}, "output_tokens": 31}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000005596, "request": {"request_id": "r18", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000004696, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 768, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 2005, 1000001800, 1000001801, 1000001802, 1000001803, 1000001804, 1000001805]}, "output_tokens": 42}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000005954, "request": {"request_id": "r19", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000005054, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 704, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 1000001900, 1000001901, 1000001902, 1000001903, 1000001904, 1000001905]}, "output_tokens": 27}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000006309, "request": {"request_id": "r20", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000005409, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 576, "input_sequence_hashes": [3000, 3001, 3002, 3003, 1000002000, 1000002001, 1000002002, 1000002003, 1000002004]}, "output_tokens": 89}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000006482, "request": {"request_id": "r21", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000005582, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 384, "input_sequence_hashes": [2000, 2001, 2002, 1000002100, 1000002101, 1000002102]}, "output_tokens": 40}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000006627, "request": {"request_id": "r22", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000005727, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 1000002200, 1000002201, 1000002202, 1000002203, 1000002204, 1000002205, 1000002206]}, "output_tokens": 49}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000006920, "request": {"request_id": "r23", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000006020, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [0, 1, 1000002300, 1000002301, 1000002302]}, "output_tokens": 32}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000007046, "request": {"request_id": "r24", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000006146, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [0, 1, 1000002400, 1000002401, 1000002402]}, "output_tokens": 85}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000007445, "request": {"request_id": "r25", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000006545, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 960, "input_sequence_hashes": [3000, 3001, 3002, 3003, 3004, 3005, 3006, 1000002500, 1000002501, 1000002502, 1000002503, 1000002504, 1000002505, 1000002506, 1000002507]}, "output_tokens": 83}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000007636, "request": {"request_id": "r26", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000006736, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 448, "input_sequence_hashes": [1000, 1001, 1002, 1000002600, 1000002601, 1000002602, 1000002603]}, "output_tokens": 91}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000007900, "request": {"request_id": "r27", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000007000, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 704, "input_sequence_hashes": [2000, 2001, 2002, 2003, 2004, 1000002700, 1000002701, 1000002702, 1000002703, 1000002704, 1000002705]}, "output_tokens": 79}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000008288, "request": {"request_id": "r28", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000007388, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [2000, 2001, 1000002800, 1000002801, 1000002802]}, "output_tokens": 57}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000008651, "request": {"request_id": "r29", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000007751, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 704, "input_sequence_hashes": [0, 1, 2, 3, 4, 1000002900, 1000002901, 1000002902, 1000002903, 1000002904, 1000002905]}, "output_tokens": 91}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009023, "request": {"request_id": "r30", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000008123, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 448, "input_sequence_hashes": [2000, 2001, 2002, 1000003000, 1000003001, 1000003002, 1000003003]}, "output_tokens": 47}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009081, "request": {"request_id": "r31", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000008181, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [2000, 2001, 1000003100, 1000003101, 1000003102]}, "output_tokens": 44}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009321, "request": {"request_id": "r32", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000008421, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 576, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1000003200, 1000003201, 1000003202, 1000003203, 1000003204]}, "output_tokens": 70}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009402, "request": {"request_id": "r33", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000008502, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 1024, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 6, 7, 1000003300, 1000003301, 1000003302, 1000003303, 1000003304, 1000003305, 1000003306, 1000003307]}, "output_tokens": 34}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009564, "request": {"request_id": "r34", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000008664, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 1000003400, 1000003401, 1000003402, 1000003403, 1000003404, 1000003405, 1000003406]}, "output_tokens": 84}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000009922, "request": {"request_id": "r35", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000009022, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 256, "input_sequence_hashes": [0, 1, 1000003500, 1000003501]}, "output_tokens": 31}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000010297, "request": {"request_id": "r36", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000009397, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1000003600, 1000003601, 1000003602, 1000003603, 1000003604, 1000003605, 1000003606]}, "output_tokens": 89}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000010408, "request": {"request_id": "r37", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000009508, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 320, "input_sequence_hashes": [3000, 3001, 1000003700, 1000003701, 1000003702]}, "output_tokens": 63}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000010517, "request": {"request_id": "r38", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000009617, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 832, "input_sequence_hashes": [0, 1, 2, 3, 4, 5, 1000003800, 1000003801, 1000003802, 1000003803, 1000003804, 1000003805, 1000003806]}, "output_tokens": 18}} +{"schema": "dynamo.request.trace.v1", "event_type": "request_end", "event_time_unix_ms": 1761000010666, "request": {"request_id": "r39", "model": "Qwen/Qwen3-32B", "request_received_ms": 1761000009766, "total_time_ms": 900.0, "replay": {"trace_block_size": 64, "input_length": 960, "input_sequence_hashes": [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1000003900, 1000003901, 1000003902, 1000003903, 1000003904, 1000003905, 1000003906, 1000003907]}, "output_tokens": 31}} diff --git a/deploy/operator/internal/webhook/defaulting/dynamographdeploymentrequest_v1beta2.go b/deploy/operator/internal/webhook/defaulting/dynamographdeploymentrequest_v1beta2.go index 612b4fb9eb71..6abb2fb1f707 100644 --- a/deploy/operator/internal/webhook/defaulting/dynamographdeploymentrequest_v1beta2.go +++ b/deploy/operator/internal/webhook/defaulting/dynamographdeploymentrequest_v1beta2.go @@ -54,8 +54,9 @@ const ( // objective.mode, the scalar metric, objective.sla, recommendation.maxCandidates, // workload.static and discovered cluster hardware. workload.static: {} opts out of the static // workload; a nonempty static object receives any missing inputTokens, outputTokens and load -// defaults. objective.sla: {} opts out of the latency defaults. hardware: {} opts out of -// cluster discovery and typed hardware defaults; a nonempty object is completed. +// defaults, and workload.trace suppresses the static default. objective.sla: {} opts out of the +// latency defaults. hardware: {} opts out of cluster discovery and typed hardware defaults; a +// nonempty object is completed. // spec.aisimulate.inputOverrides never changes a default. It overrides the defaulted typed // values later, when the input is lowered. @@ -211,6 +212,11 @@ func defaultGPUHardwareSpecV1beta2(gpu *nvidiacomv1beta2.GPUHardwareSpec, defaul } func defaultWorkloadSpecV1beta2(workload *nvidiacomv1beta2.WorkloadSpec) { + // A trace is the complete traffic model and receives no static default. The CRD schema, not + // admission, defaults its format. + if workload.Trace != nil { + return + } if workload.Static == nil { workload.Static = &nvidiacomv1beta2.StaticWorkloadSpec{} } else if workload.Static.InputTokens == nil && workload.Static.OutputTokens == nil && workload.Static.Load == nil { diff --git a/deploy/operator/internal/webhook/defaulting/suite_envtest_test.go b/deploy/operator/internal/webhook/defaulting/suite_envtest_test.go index 3001d912142d..8d92e344de86 100644 --- a/deploy/operator/internal/webhook/defaulting/suite_envtest_test.go +++ b/deploy/operator/internal/webhook/defaulting/suite_envtest_test.go @@ -367,6 +367,19 @@ func TestDynamoGraphDeploymentRequestV1beta2DefaultingAdmission(t *testing.T) { discovered: discoveredH200, check: expectV1beta2ExplicitValues, }, + { + name: "v1beta2 trace workload receives the format default and no static default", + request: v1beta2DefaultingRequest(func(spec *nvidiacomv1beta2.DynamoGraphDeploymentRequestSpec) { + spec.Workload = &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: "request-traces", SubPath: "retune/w-1.jsonl.gz", + }}, + }} + }), + discovered: discoveredH200, + check: expectV1beta2TraceWithoutStaticDefault, + idempotent: true, + }, { name: "v1beta2 explicit node width wins while admission discovers a missing budget", request: v1beta2DefaultingRequest(func(spec *nvidiacomv1beta2.DynamoGraphDeploymentRequestSpec) { @@ -570,6 +583,23 @@ func expectV1beta2ExplicitValues(t *testing.T, request *nvidiacomv1beta2.DynamoG expectV1beta2H200x8Hardware(t, request) } +func expectV1beta2TraceWithoutStaticDefault(t *testing.T, request *nvidiacomv1beta2.DynamoGraphDeploymentRequest) { + t.Helper() + want := &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Format: nvidiacomv1beta2.TraceFormatDynamo, + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: "request-traces", SubPath: "retune/w-1.jsonl.gz", + }}, + }} + if !reflect.DeepEqual(request.Spec.Workload, want) { + t.Fatalf("workload = %#v, want the trace alone", request.Spec.Workload) + } + if request.Spec.Objective == nil || request.Spec.Recommendation == nil { + t.Fatalf("objective = %#v, recommendation = %#v, want the other defaults", request.Spec.Objective, request.Spec.Recommendation) + } + expectV1beta2H200x8Hardware(t, request) +} + func expectV1beta2MetadataOnlyUpdate(t *testing.T, request *nvidiacomv1beta2.DynamoGraphDeploymentRequest) { t.Helper() expectV1beta2H200x8Hardware(t, request) diff --git a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_carrier_envtest_test.go b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_carrier_envtest_test.go index 7f431b046ad3..5759d4c685dc 100644 --- a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_carrier_envtest_test.go +++ b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_carrier_envtest_test.go @@ -24,10 +24,12 @@ import ( "testing" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" "k8s.io/client-go/dynamic" + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" "github.com/ai-dynamo/dynamo/deploy/operator/internal/features" ) @@ -45,9 +47,34 @@ func TestDynamoGraphDeploymentRequestCarrierObjectsStayWritable(t *testing.T) { t.Run("native v1beta2 written through v1alpha1", func(t *testing.T) { carrierRoundTrip(t, beta2DGDRForAdmission(setBeta2AdmissionHardware), alphaDGDRForAdmission(nil), "v1beta2") }) + t.Run("native v1beta2 trace written through v1beta1", func(t *testing.T) { + foreign := carrierRoundTrip(t, beta2DGDRForAdmission(setBeta2AdmissionTrace), betaDGDRForAdmission(nil), "v1beta2") + + t.Log("The v1beta1 carrier holds the complete trace with its defaulted format and no static default") + trace, found, err := unstructured.NestedMap(foreign.Object, "spec", "v1beta2", "workload", "trace") + if err != nil || !found { + t.Fatalf("spec.v1beta2.workload.trace found=%v err=%v", found, err) + } + subPath, _, _ := unstructured.NestedString(trace, "source", "persistentVolumeClaim", "subPath") + if trace["format"] != "Dynamo" || subPath != "retune/w-1.jsonl.gz" { + t.Fatalf("spec.v1beta2.workload.trace = %#v", trace) + } + if _, found, _ := unstructured.NestedFieldNoCopy(foreign.Object, "spec", "v1beta2", "workload", "static"); found { + t.Fatal("the carrier gained a static workload beside the trace") + } + }) +} + +func setBeta2AdmissionTrace(request *nvidiacomv1beta2.DynamoGraphDeploymentRequest) { + setBeta2AdmissionHardware(request) + request.Spec.Workload = &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: "request-traces", SubPath: "retune/w-1.jsonl.gz", + }}, + }} } -func carrierRoundTrip(t *testing.T, ownerFixture, otherFixture runtime.Object, carrierKey string) { +func carrierRoundTrip(t *testing.T, ownerFixture, otherFixture runtime.Object, carrierKey string) *unstructured.Unstructured { t.Helper() env := admissionEnv.ForTest(t) admissionGate.set(features.Gates{GPUDiscovery: true}) @@ -119,4 +146,5 @@ func carrierRoundTrip(t *testing.T, ownerFixture, otherFixture runtime.Object, c if updated.GetGeneration() != created.GetGeneration() { t.Fatalf("generation = %d, want %d: an unchanged update rewrote the spec", updated.GetGeneration(), created.GetGeneration()) } + return foreign } diff --git a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_v1beta2.go b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_v1beta2.go index 8f8c0e14eab9..e5b29131ef65 100644 --- a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_v1beta2.go +++ b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_v1beta2.go @@ -60,6 +60,9 @@ func (v *dynamoGraphDeploymentRequestValidation) validateDynamoGraphDeploymentRe if spec.Hardware != nil { allErrs = append(allErrs, v.validateHardwareSpecV1beta2(spec.Hardware, fldPath.Child("hardware"))...) } + if spec.Workload != nil { + allErrs = append(allErrs, validateWorkloadSpec(spec.Workload, fldPath.Child("workload"))...) + } if spec.AISimulate == nil { return allErrs } diff --git a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_validation_envtest_test.go b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_validation_envtest_test.go index 2a7d18b5d322..1019d45083de 100644 --- a/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_validation_envtest_test.go +++ b/deploy/operator/internal/webhook/validation/dynamographdeploymentrequest_validation_envtest_test.go @@ -751,6 +751,276 @@ func TestDynamoGraphDeploymentRequestValidator_Validate(t *testing.T) { discovered: discoveredH200, wantCELErr: `spec.workload.static.load: Invalid value: exactly one of requestsPerSecond, concurrency, or kvCapacityFraction must be specified`, }, + // v1beta2 recorded-trace workload: the shared schema rules, the format default and overrides. + { + name: "v1beta2 trace workload is kept and receives no static default", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("retune/w-1.jsonl.gz")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + check: func(t *testing.T, result *unstructured.Unstructured) { + if _, found := v1beta2Field(t, result, "spec", "workload", "static"); found { + t.Fatal("workload.static was defaulted beside a trace") + } + if got := v1beta2String(t, result, "spec", "workload", "trace", "source", "persistentVolumeClaim", "subPath"); got != "retune/w-1.jsonl.gz" { + t.Fatalf("workload.trace.source.persistentVolumeClaim.subPath = %q", got) + } + if got := v1beta2String(t, result, "spec", "workload", "trace", "format"); got != "Dynamo" { + t.Fatalf("workload.trace.format = %q, want Dynamo", got) + } + }, + }, + { + name: "v1beta2 trace subPath that leaves the claim is rejected by CEL", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("retune/../../w.jsonl")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must be relative to the claim root without empty, '.' or '..' segments`, + }, + { + name: "v1beta2 absolute trace subPath is rejected by CEL", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("/w.jsonl")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must be relative to the claim root without empty, '.' or '..' segments`, + }, + { + name: "v1beta2 trace subPath with an empty segment is rejected by CEL", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("retune//w.jsonl")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must be relative to the claim root without empty, '.' or '..' segments`, + }, + { + name: "v1beta2 trace subPath with a NUL character is rejected by CEL", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("retune/w\x00.jsonl")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must not contain a NUL character`, + }, + { + name: "v1beta2 trace claim name must be a DNS subdomain", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + trace := v1beta2TraceForAdmission("w.jsonl") + trace["source"] = map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "Bad_Claim", "subPath": "w.jsonl"}} + spec["workload"] = map[string]any{"trace": trace} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantSchemaErr: `spec.workload.trace.source.persistentVolumeClaim.claimName: Invalid value: "Bad_Claim": spec.workload.trace.source.persistentVolumeClaim.claimName in body should match '^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$'`, + }, + { + name: "v1beta2 omitted trace format defaults to Dynamo", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + trace := v1beta2TraceForAdmission("w.jsonl") + delete(trace, "format") + spec["workload"] = map[string]any{"trace": trace} + }), + gpuDiscovery: true, + discovered: discoveredH200, + check: func(t *testing.T, result *unstructured.Unstructured) { + if got := v1beta2String(t, result, "spec", "workload", "trace", "format"); got != "Dynamo" { + t.Fatalf("workload.trace.format = %q, want the Dynamo default", got) + } + }, + }, + { + name: "v1beta2 trace format admits only dynamo", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + trace := v1beta2TraceForAdmission("w.jsonl") + trace["format"] = "mooncake" + spec["workload"] = map[string]any{"trace": trace} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantSchemaErr: `spec.workload.trace.format: Unsupported value: "mooncake": supported values: "Dynamo"`, + }, + { + name: "v1beta2 trace source requires a persistentVolumeClaim", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + trace := v1beta2TraceForAdmission("w.jsonl") + trace["source"] = map[string]any{} + spec["workload"] = map[string]any{"trace": trace} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload.trace.source: Invalid value: persistentVolumeClaim is required`, + }, + { + name: "v1beta2 static and trace are mutually exclusive", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"static": map[string]any{}, "trace": v1beta2TraceForAdmission("w.jsonl")} + }), + gpuDiscovery: true, + discovered: discoveredH200, + wantCELErr: `spec.workload: Invalid value: static and trace are mutually exclusive`, + }, + { + name: "v1beta2 input overrides may restate the trace source", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("w.jsonl.gz")} + setV1beta2InputOverrides(spec, map[string]any{"traffic": map[string]any{ + "source": map[string]any{ + "type": "trace", "format": "dynamo", "paths": []any{"/var/run/aisimulate/trace/trace.jsonl.gz"}, + }, + }}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + }, + { + name: "v1beta2 input overrides may replace the trace source", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("w.jsonl")} + setV1beta2InputOverrides(spec, map[string]any{"traffic": map[string]any{ + "source": map[string]any{"type": "synthetic", "format": "mooncake", "paths": []any{"/etc/passwd"}}, + "load": map[string]any{"speedup": 2.0}, + }}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + check: func(t *testing.T, result *unstructured.Unstructured) { + if got := v1beta2String(t, result, "spec", "aisimulate", "inputOverrides", "traffic", "source", "format"); got != "mooncake" { + t.Fatalf("inputOverrides.traffic.source.format = %q, want the override kept", got) + } + }, + }, + // A v1beta2 trace in a v1beta1 carrier stored without admission: CEL ratchets the unchanged + // trace on the next v1beta2 write, so the structural webhook recursion rejects it. + { + name: "v1beta2 carrier trace subPath with a NUL character", + oldRequest: v1beta1CarrierForAdmission(map[string]any{"trace": v1beta2TraceForAdmission("retune/w\x00.jsonl")}), + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("retune/w\x00.jsonl")} + }), + seedWithoutWebhook: true, + gpuDiscovery: true, + discovered: discoveredH200, + wantWebhook: []string{ + "spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: \"retune/w\\x00.jsonl\": subPath must not contain a NUL character", + }, + // CEL ratchets the unchanged trace and reports what it tolerated as a warning. + wantWarnings: []string{ + `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must not contain a NUL character`, + }, + }, + { + name: "v1beta2 carrier trace subPath longer than its limit", + oldRequest: v1beta1CarrierForAdmission(map[string]any{"trace": v1beta2TraceForAdmission(strings.Repeat("a", 1025))}), + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission(strings.Repeat("a", 1025))} + }), + seedWithoutWebhook: true, + gpuDiscovery: true, + discovered: discoveredH200, + wantWebhook: []string{ + "spec.workload.trace.source.persistentVolumeClaim.subPath: Too long: may not be more than 1024 characters", + }, + }, + { + name: "v1beta2 carrier trace without a format receives the default on the next write", + oldRequest: v1beta1CarrierForAdmission(map[string]any{"trace": map[string]any{"source": map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "request-traces", "subPath": "w.jsonl"}}}}), + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": map[string]any{"source": map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "request-traces", "subPath": "w.jsonl"}}}} + }), + seedWithoutWebhook: true, + gpuDiscovery: true, + discovered: discoveredH200, + check: func(t *testing.T, result *unstructured.Unstructured) { + if got := v1beta2String(t, result, "spec", "workload", "trace", "format"); got != "Dynamo" { + t.Fatalf("workload.trace.format = %q, want the Dynamo default", got) + } + }, + }, + { + name: "v1beta2 carrier trace source without a persistentVolumeClaim", + oldRequest: v1beta1CarrierForAdmission(map[string]any{"trace": map[string]any{"format": "Dynamo", "source": map[string]any{}}}), + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": map[string]any{"format": "Dynamo", "source": map[string]any{}}} + }), + seedWithoutWebhook: true, + gpuDiscovery: true, + discovered: discoveredH200, + wantWebhook: []string{ + `spec.workload.trace.source.persistentVolumeClaim: Required value`, + }, + // CEL ratchets the unchanged trace and reports what it tolerated as a warning. + wantWarnings: []string{ + `spec.workload.trace.source: Invalid value: "object": persistentVolumeClaim is required`, + }, + }, + { + name: "v1beta2 carrier trace errors are reported together in field order", + oldRequest: v1beta1CarrierForAdmission(map[string]any{"static": map[string]any{}, "trace": map[string]any{"format": "mooncake", "source": map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "Bad_Claim", "subPath": "../w.jsonl"}}, "replay": map[string]any{"arrivalSpeedupRatio": 0.0}}}), + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"static": map[string]any{}, "trace": map[string]any{"format": "mooncake", "source": map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "Bad_Claim", "subPath": "../w.jsonl"}}, "replay": map[string]any{"arrivalSpeedupRatio": 0.0}}} + }), + seedWithoutWebhook: true, + gpuDiscovery: true, + discovered: discoveredH200, + wantWebhook: []string{ + `spec.workload: Forbidden: static and trace are mutually exclusive`, + `spec.workload.trace.source.persistentVolumeClaim.claimName: Invalid value: "Bad_Claim": a lowercase RFC 1123 subdomain must consist of lower case alphanumeric characters, '-' or '.', and must start and end with an alphanumeric character (e.g. 'example.com', regex used for validation is '[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*')`, + `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "../w.jsonl": subPath must be relative to the claim root without empty, '.' or '..' segments`, + `spec.workload.trace.replay.arrivalSpeedupRatio: Invalid value: 0: must be greater than 0`, + `spec: Invalid value: "": cannot be lowered to the public AISimulate recommendation input: spec.workload.trace.format: unsupported trace format "mooncake"`, + }, + // CEL ratchets the unchanged trace and reports what it tolerated as a warning. + wantWarnings: []string{ + `spec.workload: Invalid value: "object": static and trace are mutually exclusive`, + `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must be relative to the claim root without empty, '.' or '..' segments`, + }, + }, + { + name: "v1beta2 a trace accepts load, source and stop overrides", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("w.jsonl")} + setV1beta2InputOverrides(spec, map[string]any{"traffic": map[string]any{ + "load": map[string]any{"speedup": 2.0}, + "source": map[string]any{"block_size": 64}, + "stop": map[string]any{"requests": 8}, + }}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + }, + { + name: "v1beta2 source overrides stay allowed without a trace", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + setV1beta2InputOverrides(spec, map[string]any{"traffic": map[string]any{ + "source": map[string]any{"type": "trace", "format": "dynamo", "paths": []any{"/x"}}, + }}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + }, + { + name: "v1beta2 a null traffic override may delete the traffic a trace lowers", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("w.jsonl")} + setV1beta2InputOverrides(spec, map[string]any{"traffic": nil}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + }, + { + name: "v1beta2 a null source override may delete the source a trace lowers", + request: v1beta2DGDRForAdmission(func(spec map[string]any) { + spec["workload"] = map[string]any{"trace": v1beta2TraceForAdmission("w.jsonl")} + setV1beta2InputOverrides(spec, map[string]any{"traffic": map[string]any{"source": nil}}) + }), + gpuDiscovery: true, + discovered: discoveredH200, + }, { name: "v1beta2 update to duplicate backends is rejected", oldRequest: v1beta2DGDRForAdmission(nil), @@ -883,6 +1153,58 @@ func TestDynamoGraphDeploymentRunSpecIsImmutable(t *testing.T) { } } +// A Run has no admission webhook, so the trace rules it shares with the Request must hold in its +// own CRD schema: a directly created Run cannot name a file outside its claim, and its omitted +// format receives the same default. +func TestDynamoGraphDeploymentRunTraceIsValidatedByTheCRD(t *testing.T) { + tests := []struct { + name string + subPath string + withStatic bool + wantCELErr string + }{ + {name: "a valid trace is stored with the default format", subPath: "retune/w-1.jsonl.gz"}, + { + name: "a subPath that leaves the claim is rejected", + subPath: "../w.jsonl", + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must be relative to the claim root without empty, '.' or '..' segments`, + }, + { + name: "a subPath with a NUL character is rejected", + subPath: "retune/w\x00.jsonl", + wantCELErr: `spec.workload.trace.source.persistentVolumeClaim.subPath: Invalid value: "string": subPath must not contain a NUL character`, + }, + { + name: "static beside a trace is rejected", + subPath: "retune/w-1.jsonl.gz", + withStatic: true, + wantCELErr: `spec.workload: Invalid value: static and trace are mutually exclusive`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + run := v1beta2RunForAdmission(v1beta2AdmissionModel, nil) + run.Spec.Workload = &nvidiacomv1beta2.WorkloadSpec{Trace: &nvidiacomv1beta2.TraceWorkloadSpec{ + Source: nvidiacomv1beta2.TraceSourceSpec{PersistentVolumeClaim: &nvidiacomv1beta2.TracePVCSource{ + ClaimName: "request-traces", SubPath: tt.subPath, + }}, + }} + if tt.withStatic { + run.Spec.Workload.Static = &nvidiacomv1beta2.StaticWorkloadSpec{} + } + actual := runAdmissionTest(t, admissionTestCase{object: run, wantCELError: tt.wantCELErr}) + if tt.wantCELErr == "" { + if got := v1beta2String(t, actual, "spec", "workload", "trace", "source", "persistentVolumeClaim", "subPath"); got != "retune/w-1.jsonl.gz" { + t.Fatalf("stored workload.trace.source.persistentVolumeClaim.subPath = %q, want the trace kept", got) + } + if got := v1beta2String(t, actual, "spec", "workload", "trace", "format"); got != "Dynamo" { + t.Fatalf("stored workload.trace.format = %q, want the Dynamo default", got) + } + } + }) + } +} + func TestDynamoGraphDeploymentRequestStatusRepresentationMatchesSpec(t *testing.T) { tests := []struct { name string @@ -1127,6 +1449,28 @@ func v1beta2DGDRForAdmission(mutate func(spec map[string]any)) *unstructured.Uns }} } +// v1beta1CarrierForAdmission is a v1beta1 Request carrying a v1beta2 spec with workload under the +// name v1beta2DGDRForAdmission uses. Only conversion writes such a carrier, so seed it without the +// webhook. +func v1beta1CarrierForAdmission(workload map[string]any) *unstructured.Unstructured { + return &unstructured.Unstructured{Object: map[string]any{ + "apiVersion": "nvidia.com/v1beta1", + "kind": "DynamoGraphDeploymentRequest", + "metadata": map[string]any{"name": "test-v1beta2-dgdr", "namespace": "default"}, + "spec": map[string]any{"v1beta2": map[string]any{ + "model": map[string]any{"name": v1beta2AdmissionModel}, + "workload": workload, + }}, + }} +} + +func v1beta2TraceForAdmission(subPath string) map[string]any { + return map[string]any{ + "format": "Dynamo", + "source": map[string]any{"persistentVolumeClaim": map[string]any{"claimName": "request-traces", "subPath": subPath}}, + } +} + func v1beta2Field(t *testing.T, obj *unstructured.Unstructured, path ...string) (any, bool) { t.Helper() value, found, err := unstructured.NestedFieldNoCopy(obj.Object, path...) diff --git a/deploy/operator/internal/webhook/validation/shared_helpers.go b/deploy/operator/internal/webhook/validation/shared_helpers.go index c8ac790efd58..6b03744b25f2 100644 --- a/deploy/operator/internal/webhook/validation/shared_helpers.go +++ b/deploy/operator/internal/webhook/validation/shared_helpers.go @@ -414,6 +414,20 @@ func groveNativeMinimumPath(component *nvidiacomv1beta1.DynamoComponentDeploymen return fldPath.Child("providerOverride", "value", "spec", "minAvailable") } +// traceClaimPathHasOnlyNamedSegments reports whether subPath is relative to the claim root and +// every segment names an entry: none is empty, "." or "..". +func traceClaimPathHasOnlyNamedSegments(subPath string) bool { + if subPath == "" { + return false + } + for segment := range strings.SplitSeq(subPath, "/") { + if segment == "" || segment == "." || segment == ".." { + return false + } + } + return true +} + // hasRoleSidecarV1Alpha1 resolves component-wide mode from the authored leader. // spec must not be nil. func hasRoleSidecarV1Alpha1(spec *nvidiacomv1alpha1.DynamoComponentDeploymentSharedSpec) bool { diff --git a/deploy/operator/internal/webhook/validation/shared_v1beta2.go b/deploy/operator/internal/webhook/validation/shared_v1beta2.go new file mode 100644 index 000000000000..58205e9cd47d --- /dev/null +++ b/deploy/operator/internal/webhook/validation/shared_v1beta2.go @@ -0,0 +1,97 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package validation + +import ( + "strings" + "unicode/utf8" + + nvidiacomv1beta2 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1beta2" + "k8s.io/apimachinery/pkg/util/validation" + "k8s.io/apimachinery/pkg/util/validation/field" +) + +// maxTraceClaimSubPathLength is the CRD maxLength of +// spec.workload.trace.source.persistentVolumeClaim.subPath. +const maxTraceClaimSubPathLength = 1024 + +// validateWorkloadSpec validates workload, a type that Requests and Runs share. The CRD schema and +// CEL enforce the same rules on every v1beta2 write; this recursion reaches a Request whose v1beta1 +// carrier was stored without admission and that CEL ratcheting then lets through. workload and +// fldPath must not be nil. +func validateWorkloadSpec(workload *nvidiacomv1beta2.WorkloadSpec, fldPath *field.Path) field.ErrorList { + allErrs := field.ErrorList{} + if workload.Trace == nil { + return allErrs + } + if workload.Static != nil { + allErrs = append(allErrs, field.Forbidden(fldPath, "static and trace are mutually exclusive")) + } + return append(allErrs, validateTraceWorkloadSpec(workload.Trace, fldPath.Child("trace"))...) +} + +// validateTraceWorkloadSpec validates trace. Its format is not checked here: the CRD schema +// defaults and enumerates it, and lowering rejects a format it cannot translate. trace and fldPath +// must not be nil. +func validateTraceWorkloadSpec(trace *nvidiacomv1beta2.TraceWorkloadSpec, fldPath *field.Path) field.ErrorList { + allErrs := validateTraceSourceSpec(&trace.Source, fldPath.Child("source")) + if trace.Replay != nil { + allErrs = append(allErrs, validateTraceReplaySpec(trace.Replay, fldPath.Child("replay"))...) + } + return allErrs +} + +// validateTraceSourceSpec validates source. source and fldPath must not be nil. +func validateTraceSourceSpec(source *nvidiacomv1beta2.TraceSourceSpec, fldPath *field.Path) field.ErrorList { + if source.PersistentVolumeClaim == nil { + return field.ErrorList{field.Required(fldPath.Child("persistentVolumeClaim"), "")} + } + return validateTracePVCSource(source.PersistentVolumeClaim, fldPath.Child("persistentVolumeClaim")) +} + +// validateTracePVCSource validates claim: a PersistentVolumeClaim name and one nameable file inside +// it. claim and fldPath must not be nil. +func validateTracePVCSource(claim *nvidiacomv1beta2.TracePVCSource, fldPath *field.Path) field.ErrorList { + allErrs := field.ErrorList{} + for _, msg := range validation.IsDNS1123Subdomain(claim.ClaimName) { + allErrs = append(allErrs, field.Invalid(fldPath.Child("claimName"), claim.ClaimName, msg)) + } + + // The subPath names one file below the claim root: relative, with only named segments. + subPath := fldPath.Child("subPath") + if utf8.RuneCountInString(claim.SubPath) > maxTraceClaimSubPathLength { + allErrs = append(allErrs, field.TooLongCharacters(subPath, claim.SubPath, maxTraceClaimSubPathLength)) + } + if !traceClaimPathHasOnlyNamedSegments(claim.SubPath) { + allErrs = append(allErrs, field.Invalid(subPath, claim.SubPath, + "subPath must be relative to the claim root without empty, '.' or '..' segments")) + } + if strings.ContainsRune(claim.SubPath, 0) { + allErrs = append(allErrs, field.Invalid(subPath, claim.SubPath, "subPath must not contain a NUL character")) + } + return allErrs +} + +// validateTraceReplaySpec validates replay. replay and fldPath must not be nil. +func validateTraceReplaySpec(replay *nvidiacomv1beta2.TraceReplaySpec, fldPath *field.Path) field.ErrorList { + if replay.ArrivalSpeedupRatio != nil && *replay.ArrivalSpeedupRatio <= 0 { + return field.ErrorList{field.Invalid(fldPath.Child("arrivalSpeedupRatio"), *replay.ArrivalSpeedupRatio, + "must be greater than 0")} + } + return nil +} diff --git a/docs/fern/index.yml b/docs/fern/index.yml index 8e310e8f86ea..cda0bf76fb4c 100644 --- a/docs/fern/index.yml +++ b/docs/fern/index.yml @@ -224,6 +224,8 @@ navigation: path: pages/kubernetes/auto-deployment/dynamo-profiler.md - page: Dynamo Planner path: pages/kubernetes/auto-deployment/dynamo-planner.mdx + - page: Planner Live Retune + path: pages/kubernetes/auto-deployment/planner-live-retune.md - section: Operations icon: screwdriver-wrench diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx index 4c8340ee3ec2..67ad1ec95ecc 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx @@ -471,16 +471,18 @@ searches from the point it applied. To predict what an active Planner does, repl ## Simulate Retune (Experimental) -**Experimental.** Retune runs only in offline simulation in this release. A live Planner rejects a -configuration that sets `retune` at startup. +**Experimental.** This section covers replay retune, `retune.executor: local` (the default), which +runs only in offline simulation; a live Planner rejects it at startup. To retune a live deployment +on Kubernetes, which uses `executor: dgdr` and is for development clusters only, see +[Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). Retune makes the Planner retune its router and Planner runtime parameters from the traffic of the last window. For the concept and the two scenarios, see [Periodic Retune](planner-guide.md#periodic-retune-experimental). For windows, scoring, and decision -timing, see [Planner Design](planner-design.md#retune-experimental-simulation-only). For every +timing, see [Planner Design](planner-design.md#retune-experimental). For every `retune` field, its default, and the search-domain syntax, see the -[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only). +[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental). ### Requirements @@ -507,7 +509,7 @@ With `dynamo.replay.run_trace_replay`: `agent_context` (an agentic Dynamo trace). Retune windows drop tool rows. `execution_model` is accepted, because it applies only to agentic requests. - `router_mode="kv_router"` and a router configuration that meets the - [`retune.router` requirements](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only) + [`retune.router` requirements](../../../../reference/components/planner-configuration.mdx#retune-experimental) when `retune.router` is set. - With throughput scaling, retune reads `load_predictor_warmup_trace` as recorded history. It must then be a Dynamo request trace that ends at or before the first replayed request, on the same @@ -712,7 +714,7 @@ decision lands. Choose the value by what the run is for: - A fixed number of seconds gives the same decision ticks on every run. Use it to compare configurations, and set it to the sweep time you expect on the host that will run sweeps. - `measured` charges the sweep's own wall-clock time on the replay host, as the - [`decision_latency_seconds` field](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only) + [`decision_latency_seconds` field](../../../../reference/components/planner-configuration.mdx#retune-experimental) defines it. Decision ticks then depend on the speed and load of the replay host, so two runs of the same config can switch at different ticks and report different results. Set `decision_latency_scale` above 1 to model a slower sweep host, or below 1 for a faster one. Above @@ -909,7 +911,9 @@ Scenario 2 at about 4.5 GB with `max_parallel: 8`. ### Limitations -- **Simulation only.** A live Planner rejects `retune` at startup. +- **Simulation only.** A live Planner rejects `executor: local` at startup. Live retune uses + `executor: dgdr`; see + [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). - **Open-loop Dynamo request traces only.** See [Requirements](#requirements). - **Cold window replays.** Each candidate replay starts with an empty KV cache and the replica counts at the window start, so prefix-cache reuse from before the window does not carry over. diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-design.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-design.md index 03713cd4314f..f7871aae831d 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-design.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-design.md @@ -289,18 +289,26 @@ When both modes are enabled, throughput-based scaling (longer interval) sets a l In aggregated mode (`mode: agg`), engines handle both prefill and decode via chunked prefill. The planner maintains both TTFT and ITL regression models but uses per-worker time-averaged metrics (not instantaneous) for regression training to smooth out chunked prefill noise. Scale up if either prefill or decode signals overload; scale down only if both signal underload. -## Retune (Experimental, Simulation Only) + + +## Retune (Experimental) > [!WARNING] -> **Experimental.** Retune runs only in planner-in-the-loop offline replay in this release. A live -> Planner rejects a configuration that sets `retune` at startup. +> **Experimental.** With `retune.executor: local`, the default, retune runs only in +> planner-in-the-loop offline replay, and a live Planner rejects it at startup. With +> `executor: dgdr`, a live Planner on Kubernetes retunes three KV router knobs with fixed replicas +> through DGDR v1beta2 Requests; it is for development clusters only. See +> [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). Retune periodically replays the last window of recorded traffic against a grid of router and Planner parameters and applies the best point. The [Planner Guide](planner-guide.md#periodic-retune-experimental) covers the concept, the -[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only) -the fields, and [Simulate Retune](aisimulate-planner-replay.mdx#simulate-retune-experimental) how to -run it. This section describes the mechanics. +[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental) +the fields, [Simulate Retune](aisimulate-planner-replay.mdx#simulate-retune-experimental) how to +run it in replay, and +[Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md) +how to run it on a live deployment. This section describes the mechanics of replay retune, +executor `local`. ### Engine Integration @@ -326,10 +334,12 @@ Retune lives in the Planner engine (`OrchestratorEngineAdapter`), not in a plugi type, and `PlannerEffects.retune_decisions`, which carry the `planner_changes` for the host to apply. Router parameters are level-triggered desired state: the host applies the difference. - In retune-only mode, `retune` set with both scaling flags explicitly `false`, the adapter returns - no `scale_to`. The [`retune` field](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only) + no `scale_to`. The [`retune` field](../../../../reference/components/planner-configuration.mdx#retune-experimental) lists what else that mode turns off. -- The live runtime base, `NativePlannerBase`, raises `planner.retune is replay-only in this release` - at construction. +- The live runtime base, `NativePlannerBase`, raises + `planner.retune is replay-only with executor 'local'; use executor: dgdr` at construction when + `retune.executor` is `local`. With `dgdr`, it starts the live retune supervisor of + [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md) instead. ### Windows @@ -480,7 +490,7 @@ With `D` = `decision_latency_seconds` and `Δ` = `scheduling.scale_interval_seco launched at tick time `t` is ready at `t + max(D, Δ)`, and the decision lands on the first tick at or after that time at which the sweep is done, never in the launch tick. With `measured`, `D` is the modeled sweep time that the -[`decision_latency_seconds` field](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only) +[`decision_latency_seconds` field](../../../../reference/components/planner-configuration.mdx#retune-experimental) defines. Under a `VirtualClock`, the replay host, the controller waits for the sweep inside the launch tick, @@ -544,15 +554,14 @@ request trace, and its error names the AISimulate setting. 1. **Adjustment interval vs scaling latency**: If a plugin interval is shorter than the time to scale, later ticks may observe an in-progress transition and hold rather than stacking new replica changes. 2. **Average-based prediction**: Throughput-based scaling uses average ISL/OSL, which may not represent bimodal or heavy-tailed distributions well. 3. **Local Planner scope**: Each local Planner instance directly manages one DGD. Use GlobalPlanner for shared GPU budgets and coordination across multiple DGDs. -4. **Retune is experimental and simulation-only**: See the retune [limitations](aisimulate-planner-replay.mdx#limitations). +4. **Retune is experimental**: executor `local` runs only in simulation; see the retune [limitations](aisimulate-planner-replay.mdx#limitations). Executor `dgdr` runs on a live Planner with the limitations in [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md#limitations). ## Future Work - Distribution-aware interpolation (beyond mean ISL/OSL) - Adaptive adjustment interval based on observed scaling latency -- Retune on a live Planner: reading the frontend request traces as they are written, patching - `/router_policy` on the frontends, and restoring the previous point when a deployment rejects an - applied change +- Live retune of Planner parameters, with load or throughput scaling, and of the router knobs + that a DGDR v1beta2 Request cannot express ## File Map @@ -573,6 +582,7 @@ request trace, and its error names the AISimulate setting. | `components/src/dynamo/planner/retune/sweep.py` | Sweep worker processes, candidate replays, scoring, and the load-predictor pre-sweep. | | `components/src/dynamo/planner/retune/window.py` | Request-trace window and history reader for sweeps. | | `components/src/dynamo/planner/retune/domains.py` | Expansion of `retune.router` and `retune.planner` into the candidate grid. | +| `components/src/dynamo/planner/retune/live/` | Live retune (`executor: dgdr`): deployment rules, window cuts, DGDR v1beta2 sweeps, leadership and state Leases, router actuation, and clean-up. | | `components/src/dynamo/planner/offline/replay_adapter.py` | Replay host: tick records, scaling events, router parameter updates, and retune actuation. | | `components/src/dynamo/planner/simulation/` | AISimulate Planner adapter: the `planner` section of `aisimulate predict` and `recommend`. | | `components/src/dynamo/router/simulation/config.py` | AISimulate router section schemas, including `retune.router`. | diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md index f29eea09ce23..58a34a71f76f 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md @@ -69,7 +69,7 @@ spec: advisory: true ``` -Advisory mode is suggestion-only. The Planner computes recommended replica counts, logs them, exports them as diagnostics, and shows them in HTML reports. The recommendations are not applied as scaling decisions: the Planner does not execute scaling actions or change the deployment. Planner-in-the-loop replay honors it too: the replay keeps the simulated deployment fixed, reports the recommendations, and records [retune](#periodic-retune-experimental) decisions, which run only in simulation, without applying them; see [Replay in Advisory Mode](aisimulate-planner-replay.mdx#replay-in-advisory-mode). +Advisory mode is suggestion-only. The Planner computes recommended replica counts, logs them, exports them as diagnostics, and shows them in HTML reports. The recommendations are not applied as scaling decisions: the Planner does not execute scaling actions or change the deployment. Planner-in-the-loop replay honors it too: the replay keeps the simulated deployment fixed, reports the recommendations, and records [retune](#periodic-retune-experimental) decisions without applying them; see [Replay in Advisory Mode](aisimulate-planner-replay.mdx#replay-in-advisory-mode). ### Optimization Target @@ -328,9 +328,12 @@ spec: ## Periodic Retune (Experimental) > [!WARNING] -> **Experimental.** In this release, retune runs only in offline simulation: -> `aisimulate predict --stack dynamo` and `dynamo.replay.run_trace_replay`. A live Planner rejects a -> configuration that sets `retune` at startup. +> **Experimental.** `retune.executor` selects where retune runs. With `local`, the default, it runs +> only in offline simulation: `aisimulate predict --stack dynamo` and +> `dynamo.replay.run_trace_replay`, and a live Planner rejects it at startup. With `dgdr`, a live +> Planner on Kubernetes retunes three KV router knobs with fixed replicas through DGDR v1beta2 +> Requests; it is for development clusters only. See +> [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). Retune lets the Planner tune its own runtime parameters from the traffic it served. Every `retune.interval_seconds`, it takes the window of the request trace that just ended and replays it @@ -375,8 +378,11 @@ A candidate that changes decisions is judged on one replay of one window, so This minimal fixed-replica configuration searches the 108-point default router domain every six hours. It is a `PlannerConfig` for the `planner_config` argument of `dynamo.replay.run_trace_replay` (see [Run Retune from Python](aisimulate-planner-replay.mdx#run-retune-from-python)), not a DGDR -`spec.features.planner` fragment: a live Planner rejects it. In `aisimulate predict`, the `planner` -section takes `retune` and the two scaling flags, and `mode` comes from `engine.mode`; see +`spec.features.planner` fragment: a live Planner rejects it, because it uses the default +`executor: local`. For a live Planner, see +[Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). +In `aisimulate predict`, the `planner` section takes `retune` and the two scaling flags, and `mode` +comes from `engine.mode`; see [Scenario 1](aisimulate-planner-replay.mdx#scenario-1-fixed-replicas-router-only). ```yaml @@ -389,11 +395,12 @@ retune: ``` To run retune and read its decisions, see -[Simulate Retune](aisimulate-planner-replay.mdx#simulate-retune-experimental). For windows, -candidates, scoring, and decision timing, see -[Planner Design](planner-design.md#retune-experimental-simulation-only). For every field and the +[Simulate Retune](aisimulate-planner-replay.mdx#simulate-retune-experimental), and to run it on a +live deployment, see [Planner Live Retune](../../../../kubernetes/auto-deployment/planner-live-retune.md). +For windows, candidates, scoring, and decision timing, see +[Planner Design](planner-design.md#retune-experimental). For every field and the search-domain syntax, see the -[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental-simulation-only). +[`retune` reference](../../../../reference/components/planner-configuration.mdx#retune-experimental). ## Integration with Profiler diff --git a/docs/fern/pages/kubernetes/auto-deployment/planner-live-retune.md b/docs/fern/pages/kubernetes/auto-deployment/planner-live-retune.md new file mode 100644 index 000000000000..117db53ca8d0 --- /dev/null +++ b/docs/fern/pages/kubernetes/auto-deployment/planner-live-retune.md @@ -0,0 +1,408 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Planner Live Retune +subtitle: Let the Planner tune the KV router of a running DynamoGraphDeployment from its own recorded traffic, through DGDR v1beta2 search runs. +--- + +> [!CAUTION] +> **Experimental, development clusters only.** Live retune builds on DGDR v1beta2, which the +> operator serves only with the experimental `dgdrV1Beta2.enabled` Helm value, and it writes +> router parameters through the frontend admin API, which has no authentication. Do not enable it +> on a cluster that serves production traffic. + +Live retune is the Kubernetes form of [Periodic Retune](../../developer-guide/knowledge-base/modular-components/planner/planner-guide.md#periodic-retune-experimental). +Every `retune.interval_seconds`, the Planner cuts the window of request traces that just ended, +creates DynamoGraphDeploymentRequests (DGDR) v1beta2 that replay the window for a grid of KV +router parameters, scores the candidates the Runs report on `retune.objective`, and keeps the +current parameters unless a candidate beats them by more than `retune.min_improvement_percentage`. +Outside advisory mode it then sets the winner on every frontend router with +`PATCH /router_policy`. + +Live retune searches the three router knobs a DGDR Request can express: +`overlap_score_credit`, `prefill_load_scale` and `temperature`. Replicas stay fixed: the Planner +runs in retune-only mode and never scales. + +## How it works + +1. **Record.** Each frontend writes a `request_end` record per request to its own directory on a + shared ReadWriteMany claim (`///`). +2. **Cut a window.** At `L_k = window end + grace_seconds`, the Planner merges the records of the + window into one canonical trace file, + `//-.jsonl.gz`, in a subprocess bounded by + `cut_budget_seconds`. `` is a digest of the window's start, end and cutoff + (`window end + grace_seconds`), so a Planner restarted with another `interval_seconds` or + `grace_seconds` cuts its own file and creates its own Requests instead of reusing those of a + window with the same end. +3. **Search.** The Planner splits the router grid into at most `run.max_shards` Cartesian shards of + at most 100 points and creates one DGDR v1beta2 Request per shard. Each Request replays the + window file on the pinned `retune.dgdr.engine` with the replica counts of the DGD spec. One more + single-point Request replays a grid point again as a common-random-numbers probe: its score + must match the grid's within a relative error of 1e-9. +4. **Decide.** The Planner reads each Run's Candidates, scores them on `retune.objective`, and + decides with the same rules as replay retune. A decision is persisted in a Lease before it + reaches any router. +5. **Apply.** The Planner sets all eight router policy parameters of every frontend router to + the decided values, and keeps them there: a manual change to them is reverted. + +Windows are aligned to `retune.dgdr.phase_seconds + k × interval_seconds` in epoch time, so a +restarted Planner continues the same windows. Leadership and state live in two Leases owned by +the DGD, `-planner-retune` and `-planner-retune-state`. The first window after retune is +first enabled is skipped, since its trace is incomplete. + +Every DGDR trace Run replays with common random numbers: ordinal request ids and a fixed router +seed. A Run's score therefore does not depend on the other candidates or the Run's parallelism. +This applies to all DGDR trace Runs, not only those the Planner creates. + +## Prerequisites + +- **Operator with DGDR v1beta2 traces.** Install the platform chart with + `dynamo-operator.dgdrV1Beta2.enabled=true` and an operator that accepts `spec.workload.trace`. + Also set `dynamo-operator.dgdrV1Beta2.plannerLiveRetune=true` to grant the Planner + ServiceAccount the access it needs (see [RBAC](#rbac)). + In a cluster-wide install, the operator must also hold the DGDR v1beta2 Run and Candidate + permissions of [ai-dynamo/dynamo#17308](https://github.com/ai-dynamo/dynamo/pull/17308); see + [RBAC](#rbac). +- **Run Job scheduling.** The namespace needs a LimitRange for the Run Jobs, and the operator + must discover GPUs, through DCGM or GPU discovery with labeled nodes. See + [Auto Deploy with DGDR](auto-deploy-with-dgdr.md#prerequisites). +- **A public model.** The Run Job has no Hugging Face credentials and no model mount, so + `retune.dgdr.model` must be a public Hugging Face repository id that AISimulate can resolve. +- **A ReadWriteMany claim** in the DGD namespace, dedicated to traces. Create `segments/` and + `windows/` at its root, writable by the user the frontend and Planner containers run as. The + Planner writes each window file with mode `0644`, so the Run Job, which reads it as UID 1000 + with supplemental group 65532, can read it whatever user the Planner runs as. +- **A KV-router frontend with default routing fields**, and workers that do not set their own + router mode or pass a value the Planner cannot read where an option could stand. See + [Deployment rules](#deployment-rules). +- **One Planner replica.** + +## Configure the frontend + +Add the trace environment and the claim to the frontend container of the DGD. The planner checks +these values against `retune.dgdr` at startup and whenever the DGD changes. + +```yaml +spec: + components: + - name: Frontend + type: frontend + replicas: 1 + podTemplate: + spec: + containers: + - name: main + command: [python3] + args: [-m, dynamo.frontend, --router-mode, kv] + env: + - name: DYN_POD_NAME + valueFrom: {fieldRef: {fieldPath: metadata.name}} + - {name: DYN_REQUEST_TRACE, value: "1"} + - {name: DYN_REQUEST_TRACE_RECORDS, value: request_end} + - {name: DYN_REQUEST_TRACE_SINKS, value: file} + - {name: DYN_REQUEST_TRACE_FILE_FORMAT, value: jsonl_gz} + - name: DYN_REQUEST_TRACE_FILE_PATH + value: /trace/segments/$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/req + - {name: DYN_REQUEST_TRACE_FILE_ROLL_BYTES, value: "67108864"} + - {name: DYN_REQUEST_TRACE_CAPACITY, value: "8192"} + volumeMounts: + - {name: request-traces, mountPath: /trace} + volumes: + - name: request-traces + persistentVolumeClaim: {claimName: request-traces} +``` + +- `DYN_REQUEST_TRACE_FILE_PATH` is a file prefix directly under + `//$(DYN_PARENT_DGD_K8S_NAME)/$(DYN_POD_NAME)/`, such as `req`: the + part after `$(DYN_POD_NAME)/` must not contain `/`. `` is where the frontend mounts + `retune.dgdr.claim_name`. The operator injects `DYN_PARENT_DGD_K8S_NAME`. +- Records are dropped without a count for `prompt_embeds`, multimodal, `n > 1` and `best_of > 1` + requests, and when the writer lags or a batch write fails. +- The trace holds unkeyed xxh3 hash chains of the prompts. Anyone who can create pods in the + namespace can mount the claim and read them. + +See [Agent Tracing](../../use-cases/agents/agent-tracing.md) for the `request_end` record. + +## Configure the Planner + +Mount the same claim in the Planner container at `retune.dgdr.mount_path`, and pass a +configuration like the one below with `--config`: + +```yaml +environment: kubernetes +backend: vllm +mode: agg +advisory: true # start advisory; see Roll out +enable_load_scaling: false # retune-only: both flags false +enable_throughput_scaling: false +retune: + interval_seconds: 21600 # T: window length and decision cadence + grace_seconds: 120 # covers p99 request time + 1 s flush + claim visibility + timeout_seconds: 20400 # required; at most T - 600; includes the window cut + decision_latency_seconds: 0 + objective: GoodputPerGPU + sla_ttft_ms: 500 + sla_itl_ms: 30 + router: default # 3 x 9 x 4 = 108 points + executor: dgdr + dgdr: + phase_seconds: 0 + claim_name: request-traces + mount_path: /trace + segments_dir: segments + windows_dir: windows + model: Qwen/Qwen3-32B + sku: h200_sxm + engine: + mode: aggregated + backend: vllm + context_length: 32768 + workers: + aggregated: + parallelism: + preset: + - {tensor: 2, pipeline: 1, attention_data: 1, moe_tensor: 1, moe_expert: 1} + scheduler: {max_batched_tokens: 8192, max_sequences: 256} + prefill: null + decode: null + cut_budget_seconds: 900 + job_startup_seconds: 300 + run: + parallelism: 8 + candidate_timeout_seconds: 1800 + max_shards: 3 + keep_requests: 5 + max_window_rows: 5000000 + promote_last_advisory: false +``` + +`retune.dgdr` fields: + +| Field | Default | Meaning | +|---|---|---| +| `claim_name` | required | ReadWriteMany claim, in the DGD namespace, that holds the traces | +| `mount_path` | `/trace` | Where the Planner container mounts the claim; absolute | +| `segments_dir`, `windows_dir` | `segments`, `windows` | Claim-relative directories of the frontend segments and the window files; neither may contain the other | +| `phase_seconds` | `0` | Window phase in epoch seconds | +| `model` | required | Hugging Face repository id the workers serve; the Requests' `spec.model.name` | +| `sku` | required | The Requests' `spec.hardware.gpu.sku` | +| `engine` | required | The pinned AISimulate engine section; see [Engine section](#engine-section) | +| `cut_budget_seconds` | `900` | Deadline of the window cut subprocess | +| `job_startup_seconds` | `300` | Run Job start-up time reserved when deciding whether a shard can still finish | +| `run.parallelism` | `8` | AISimulate `optimizer.parallelism` of each Run; Job resources may lower it | +| `run.candidate_timeout_seconds` | `1800` | AISimulate `optimizer.candidate_timeout_seconds` of each Run | +| `run.max_shards` | `3` | Most grid shards (Requests) per window; the probe and off-grid points do not count | +| `keep_requests` | `5` | Windows whose finished Requests are kept for audit | +| `max_window_rows` | `5000000` | A larger window fails with `window_too_large` | +| `promote_last_advisory` | `false` | On switching to active, apply the last advisory recommendation if it is at most one interval old and the configuration is unchanged | +| `lease_duration_seconds` | `60` | Leadership Lease duration, renewed every third of it or sooner; at least `37`, so two renewals of up to 11 s each fit before the Planner stops leading, 15 s before the Lease expires | +| `actuator_interval_seconds` | `15` | Period of the router and clean-up rounds | + +For the other `retune` fields and the limits that `executor: dgdr` adds to them, see the +[Retune field reference](../../reference/components/planner-configuration.mdx#retune-experimental). + +Live retune takes its scoring SLA when it starts: `retune.sla_ttft_ms` and `retune.sla_itl_ms`, +or the Planner's `ttft_ms` and `itl_ms` where they are unset. An update of `ttft_ms` or `itl_ms` +through the +[runtime reconfiguration API](../../reference/components/planner-configuration.mdx#system-port-reconfiguration-api) +that would change that SLA fails with HTTP `422`. To change the SLA, restart the Planner with the +new values; to keep `ttft_ms` and `itl_ms` changeable at runtime, set both `retune` SLA fields. + +### Engine section + +`retune.dgdr.engine` becomes `spec.aisimulate.inputOverrides.engine` of every Request, so every +value in it must be a scalar. A `{choices: ...}` or `{range: ...}` domain, or an omitted knob that +AISimulate searches by default, would make the Run search engine settings as well. + +- `mode` is `aggregated` for `mode: agg` and `disaggregated` for `mode: disagg`. +- Each worker role of the mode (`aggregated`, or `prefill` and `decode`) sets + `parallelism.preset` to a list of exactly one entry with `tensor`, `pipeline`, `attention_data`, + `moe_tensor` and `moe_expert`, and sets `scheduler.max_batched_tokens` and + `scheduler.max_sequences`. +- The preset entry may also set `prefill_context`, but only to `1`, and must not set + `decode_context`: the pinned AISimulate recommendation does not run context parallelism. +- Do not set `replicas` anywhere: the Planner fills them in from the DGD spec at each launch and + sets `spec.hardware.gpu.budget` to their GPU total. +- Write roles the mode does not use as `null`, or leave them out. +- Set the model and GPU with `retune.dgdr.model` and `retune.dgdr.sku`, not in the engine. +- Leave `backend_version` out to let the pinned AISimulate pick its version for + `retune.dgdr.sku`. If you set it, it must be a version AISimulate accepts for that GPU; + otherwise every Run fails with `UnlistedQueryVersionError`. + +### Time budget + +Each shard must finish before the window's deadline, `L_k + timeout_seconds`. A shard of `n` +points needs + +```text +need(n) = cut_budget_seconds + job_startup_seconds + + (1 + ceil((n - 1) / run.parallelism)) x run.candidate_timeout_seconds +``` + +AISimulate evaluates the first candidate alone, hence the `1 +`. The Planner rejects a +configuration whose largest grid shard needs more than `timeout_seconds`, because the baseline +can fall in any shard and a window without a baseline score fails. With the defaults, the 108-point +grid splits into 3 shards of 36 points and needs 900 + 300 + 6 × 1800 = 12000 seconds. Measure +replay time and memory on your trace before relying on these defaults. + +## Deployment rules + +The Planner checks three kinds of rules. None of them stops a running Planner except a +configuration error at startup. + +**Configuration**, checked when the Planner starts: + +| Rule | Requirement | +|---|---| +| P1 | `environment: kubernetes`; `decision_latency_seconds` is `0`, or `measured` with `decision_latency_scale: 1`; `timeout_seconds` is set and at most `interval_seconds - 600` | +| P2 | `enable_load_scaling` and `enable_throughput_scaling` are both `false`; `retune.planner` is not set | +| P3 | `retune.router` writes only `overlap_score_credit`, `prefill_load_scale`, `temperature`, `policy: kv_router` and `prefill_load_model: {type: none}` | +| P5 | `retune.dgdr.engine` is all scalars, as in [Engine section](#engine-section) | +| P13 | The largest grid shard has at most 100 points and fits the [time budget](#time-budget) | + +**Deployment**, checked at startup, when the DGD changes and before every router round. While one +fails, the Planner launches no windows, resets the routers to their startup parameters, and +retries: + +| Rule | Requirement | +|---|---| +| P4 | The frontend runs `--router-mode kv` with every routing field the Request cannot express at its default, `kv_event_lag_ms` unset or 0, no `DYN_ROUTER_WORKER_SELECTION_POLICY`, a `none` prefill load model, router env given as plain values (no `valueFrom`), and a command that is not wrapped in a shell. The Planner cannot read `envFrom` sources: it assumes they set no router env, only warns about a source whose `prefix` could produce a `DYN_` name, and treats a `$(VAR)` that only an `envFrom` source could define as unreadable. Keep router env out of `envFrom`: P14 catches only a changed policy parameter. No worker advertises its own router configuration: worker components set neither `--router-mode` nor `DYN_ROUTER_MODE`, directly or through the graph-level `spec.env`, because a worker's advertised configuration replaces the frontend's for its workers. The Planner reads a worker's command as the operator runs it: an omitted `command` is `/bin/sh -c` with the first of `args` as the script, and a shell started through `env`, `tini --` or `dumb-init --` is a script too. For a worker whose launches are in role pod templates (`roles[].podTemplate`, no component `podTemplate`), it reads each role's main container and env. A worker with a native `runtime` sidecar is not checked: its main runs the engine, and the sidecar advertises no router configuration. A value the Planner cannot read and that could hold an option (a ConfigMap, Secret or annotation `valueFrom`, or a `$(VAR)` that the container's `env` leaves to an `envFrom` source) must not stand where a worker option could: attach it to its option (`--option=$(VAR)`) or pass it as the separate value right after `--model`, `--model-path`, `--served-model-name`, `--model-name` or `--speculative-config`. A frontend argument follows the same rule, with `--tool-call-parser`, `--reasoning-parser` and `--chat-template` as the options whose separate value it may be. In a worker shell script any such `$(VAR)` fails, since Kubernetes substitutes it before the shell parses the script; use the shell's own `$VAR` instead. A `$VAR` or `${VAR}` of such a value fails when it is unquoted, stands alone as an argument of the worker or of `env`, `export`, `source` or a shell (also by path and behind `exec`, `nohup`, `command` or `builtin`), or sits in text that `eval` or a shell given `-c` parses again; quote it outside the worker's arguments and attach it to its option among them (`--option="$VAR"`). A redirection target, such as `> "$LOG"`, is fine. The Planner does not follow the script's conditions, so the script's own assignment to a ConfigMap, Secret or annotation variable does not clear it. With an `envFrom` source that has no `prefix`, every name the container's `env` leaves undefined counts as such a value, image env such as `PATH` and `LD_LIBRARY_PATH` included, so quote those references too (`export LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH"`). Names the operator sets on the container do not count: `DYN_NAMESPACE`, `DYN_COMPONENT`, `DYN_PARENT_DGD_K8S_*` and `POD_*`, a worker's `DYN_SYSTEM_*`, `DYN_HEALTH_CHECK_ENABLED`, `NIXL_TELEMETRY_*` and `DYN_FORWARDPASS_METRIC_PORT`, and a multinode component's `DYNAMO_RANK` and `DYNAMO_LEADER_ADDRESS`. Other `fieldRef` values, such as the pod name, may stand anywhere except inside an option's name (`--x-$(POD_NAME)`) or as the value of a frontend router flag. The Planner does not evaluate command substitution, backticks, the scripts a worker runs, text that a shell or `xargs` reads from a heredoc, a here-string or a pipe, or a script that `su -c` runs, so keep router flags and such values out of them | +| P7 | The DGD serves one model; the frontend is not an EPP, not a `frontendSidecar`, and has no TLS | +| P8 | One frontend replica, or `--router-replica-sync` | +| P9 | No `scalingAdapter` on worker components; one Planner replica | +| P10 | The frontend trace environment and claim mount match [Configure the frontend](#configure-the-frontend), with the trace file prefix directly under `$(DYN_POD_NAME)/` | +| P14 | Every router's startup parameters equal the ones the Planner parsed from the frontend | + +**Health**, checked every round. While one fails, windows do not launch; nothing is reset: + +| Rule | Requirement | +|---|---| +| P11 | The API server accepts a dry-run Request with `spec.workload.trace` | +| P12 | At least `max(1, replicas)` frontends are Ready, each with a router per worker type, and `/router_policy` answers | + +The Planner only warns when AISimulate does not list the model among its default models (P6). + +## Roll out + +1. Prepare the claim and its directories, then add the frontend trace configuration and the + Planner mount. +2. Start with `advisory: true`. In advisory mode the Planner records the decision it would make + and changes no router value it has not set itself. Watch `dynamo_planner_retune_ready` and the + decision history: + + ```bash + kubectl get lease -planner-retune-state -n \ + -o jsonpath='{.metadata.annotations.nvidia\.com/planner-retune-history}' + ``` + +3. Set `advisory: false`. With `retune.dgdr.promote_last_advisory: true`, the Planner applies the + last advisory recommendation right away if it is at most one interval old and the + configuration has not changed since. + +## Disable + +1. Set `advisory: true` and wait until `dynamo_planner_retune_router_converged` is 1 and + `dynamo_planner_retune_rollback_pending` is 0. The routers are then back at their startup + parameters. The mode stored in the state alone does not show that. +2. Remove `retune` from the Planner configuration. +3. Remove the trace environment from the frontend. +4. Delete the Planner's Requests and both Leases. Select the Requests by the DGD's + `metadata.uid`: the `nvidia.com/planner-retune-dgd` label holds a hash instead of the name + when the name is longer than 63 characters. + + ```bash + kubectl delete dynamographdeploymentrequests.nvidia.com -n \ + -l nvidia.com/planner-retune-dgd-uid=$(kubectl get dgd -n -o jsonpath='{.metadata.uid}') + kubectl delete lease -n -planner-retune -planner-retune-state + ``` + + While the chart still grants the live retune RBAC, a Planner without `retune` does this on + its own: it looks for the Leases in the background after it starts (once more after a + transient API error), resets any router the removed retune changed, then deletes the + Requests and both Leases. + +5. Delete `/` and `/` from the claim. + +Before upgrading the operator or its CRDs, switch to advisory or disable retune, then delete the +retune Requests. + +## Security + +- **Frontend admin API.** `/router_policy` is served on the frontend HTTP port, next to the + OpenAI API, without authentication, and is on by default. An Ingress that forwards `/` exposes + it. Deny the admin paths at L7 and allow in-cluster access only from the ingress controller and + the Planner. Do not set `DYN_DISABLE_FRONTEND_ADMIN_API` on the frontend: the Planner reports a + frontend whose `/router_policy` answers 404 as not ready. +- **Ownership.** Outside advisory mode the Planner owns the eight router policy parameters and + reverts manual changes. +- **Trace data.** See the privacy note in [Configure the frontend](#configure-the-frontend). + +### RBAC + +With `dynamo-operator.dgdrV1Beta2.enabled=true` and +`dynamo-operator.dgdrV1Beta2.plannerLiveRetune=true`, the platform chart adds these rules to the +Planner Role (namespace-restricted installs) and ClusterRole (cluster-wide installs). Both values +default to `false`; `dgdrV1Beta2.enabled` alone grants the Planner nothing new. + +| API group | Resources | Verbs | +|---|---|---| +| `nvidia.com` | `dynamographdeploymentrequests` | create, get, list, delete | +| `nvidia.com` | `dynamographdeploymentruns` | get | +| `nvidia.com` | `dynamographdeploymentcandidates` | list | +| `coordination.k8s.io` | `leases` | get, create, update, delete | + +The chart has one Planner role for the whole install, so Planners without `retune` get these +rules too. The Planner also lists pods with its existing permission. These rules are namespace-wide: a +Role or a RoleBinding of the ClusterRole has no `resourceNames`, because the Planner creates its +Leases and Requests under per-DGD names. As a result, any workload that runs as the Planner +ServiceAccount can: + +- create DGDR Requests, which start Run Jobs, in its namespace; +- read, update or delete any Lease in its namespace, not only the two retune Leases. In a + namespace-restricted install, the Planner Role lives in the release namespace, which by default + also holds the operator's leader-election Lease, so such a workload can disrupt operator leader + election. + +In a cluster-wide install, the operator binds the Planner ClusterRole into each DGD namespace with +a RoleBinding, and Kubernetes lets it bind only permissions it holds itself. Set +`dynamo-operator.dgdrV1Beta2.plannerLiveRetune=true` there only with an operator whose own role +grants `get` on `dynamographdeploymentruns` and `list` on `dynamographdeploymentcandidates`, +which the DGDR v1beta2 operator change in +[ai-dynamo/dynamo#17308](https://github.com/ai-dynamo/dynamo/pull/17308) adds. +Without it, the API server rejects the Planner RoleBinding in every namespace that does not have +one yet, and the DGDs there fail to reconcile their Planner RBAC, with or without live retune. + +## Metrics + +The Planner exports these on `metric_reporting_prometheus_port`: + +| Metric | Type | Meaning | +|---|---|---| +| `dynamo_planner_retune_decisions_total{outcome}` | counter | Decisions, counted after the state write persisted them | +| `dynamo_planner_retune_failures_total{reason}` | counter | Failed windows by reason, such as `launch_too_late`, `replicas_changed`, `window_cut_failed`, `trace_corrupt`, `crn_mismatch` or `k8s_error` | +| `dynamo_planner_retune_skipped_total{reason}` | counter | Due windows held back: `not_leader`, `not_ready` or `already_decided` | +| `dynamo_planner_retune_leader`, `_ready`, `_initialized`, `_rollback_pending` | gauge | Supervisor state, 0 or 1 | +| `dynamo_planner_retune_controller_generation` | gauge | Generation of the retune controller; it grows when a DGD change rebuilds it | +| `dynamo_planner_retune_baseline_score`, `_best_score` | gauge | Scores of the latest decision; NaN when not scored | +| `dynamo_planner_retune_sweep_seconds`, `_missing_candidates` | gauge | Duration of the latest sweep and its candidates without a score | +| `dynamo_planner_retune_router_converged` | gauge | 1 when every frontend router holds the target parameters | +| `dynamo_planner_retune_router_patch_errors_total{code}` | counter | Failed router updates by HTTP status or error kind | +| `dynamo_planner_retune_trace_sink_missing{pod}` | gauge | 1 for a Ready frontend without its own trace directory | +| `dynamo_planner_retune_trace_dropped_members_total` | counter | Corrupt or truncated trace gzip members skipped by window cuts | + +## Limitations + +| Not supported | Reason | +|---|---| +| Retune with load or throughput scaling, or of Planner parameters | A DGDR Run has no base Planner configuration, warm-up or predictor pre-sweep | +| Frontends with non-default routing fields beyond the three knobs, or `kv_event_lag_ms` other than 0 | The DGDR router section expresses only the three knobs, and the Run does not pass a lag | +| Workers that set their own router mode (`--router-mode` or `DYN_ROUTER_MODE`) | A worker's advertised router configuration replaces the frontend's for its workers, and the Run replays only the frontend's | +| Worker arguments that pass a ConfigMap, Secret, annotation or `envFrom` value where an option could stand, including an unquoted `$VAR` of one in a worker shell script or one that `eval` or a shell given `-c` parses again (P4 fails) | The Planner cannot read the value, so it cannot rule out a router mode the worker would advertise | +| Router flags that a worker produces through shell command substitution, backticks, script variables, positional parameters, literal text that `eval` or a nested shell parses again, or a script file, and ConfigMap, Secret, annotation or `envFrom` values in text that a shell or `xargs` reads from a heredoc, a here-string or a pipe, or that `su -c` runs (P4 does not see them) | The Planner reads the worker's spec and does not run its script | +| Raising `overlap_score_credit` on a frontend that starts it at 0: an omitted credit then searches only 0, and a written positive value fails P4 | A live router rejects a positive credit after starting at 0 | +| Private, fine-tuned or local-path models | The Run Job has no Hugging Face credentials or model mount | +| Several frontends without router replica sync | AISimulate simulates one router | +| EPP or Gateway API Inference Extension frontends, `frontendSidecar`, TLS frontends | No `/router_policy` there, different addressing, or certificates | +| Run Job resources, deadline and tolerations | Fixed by the DGDR Run Job | +| Time-based segment rotation and a segment count limit | The frontend trace writer rotates by size only | +| A separate or authenticated admin listener | Not available in the frontend yet | diff --git a/docs/fern/pages/reference/components/planner-configuration.mdx b/docs/fern/pages/reference/components/planner-configuration.mdx index b5a9ef4ee116..5d9652107c89 100644 --- a/docs/fern/pages/reference/components/planner-configuration.mdx +++ b/docs/fern/pages/reference/components/planner-configuration.mdx @@ -381,7 +381,7 @@ These fields default from environment variables and are excluded from serialized ## Advisory mode - Suggestion-only mode. The Planner computes, logs, exports, and reports recommended replica counts and [retune](#retune-experimental-simulation-only) decisions without executing any of them or changing the deployment. Use it to evaluate a new configuration or validate SLA targets against production traffic. The runtime configuration APIs, such as `PATCH /v1/min-endpoints`, still change the Planner's own settings. + Suggestion-only mode. The Planner computes, logs, exports, and reports recommended replica counts and [retune](#retune-experimental) decisions without executing any of them or changing the deployment. Use it to evaluate a new configuration or validate SLA targets against production traffic. The runtime configuration APIs, such as `PATCH /v1/min-endpoints`, still change the Planner's own settings. Planner-in-the-loop replay honors it too: an advisory replay keeps the simulated deployment at its starting replicas and router parameters, reports the recommendations in `planner.recommended_scaling_events`, and records a retune decision that would apply with outcome `advisory`. See [Replay in Advisory Mode](../../developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx#replay-in-advisory-mode). @@ -476,6 +476,11 @@ Fields retain their normal mode-specific meaning. For example, the explicit queu thresholds are used by `optimization_target: load`. Predictor replacement requires `optimization_target: sla` with throughput scaling enabled. Scheduling intervals, AIS performance models, deployment topology, and other fields are rejected by this API. +With [live retune](../../kubernetes/auto-deployment/planner-live-retune.md) (`retune.executor: dgdr`), +an update of `ttft_ms` or `itl_ms` that changes the retune scoring SLA, because `retune.sla_ttft_ms` +or `retune.sla_itl_ms` is unset, is rejected with HTTP `422`. To change the SLA, restart the +Planner with the new values; to keep `ttft_ms` and `itl_ms` changeable at runtime, start it with +both `retune` SLA fields set. Use the instance ID and revision returned by GET when submitting an update: @@ -627,12 +632,17 @@ The Planner runs the builtin plugin pipeline by default. `scheduling` controls p -## Retune (Experimental, Simulation Only) + -`retune` makes the Planner periodically replay the last window of recorded traffic against a grid of router and Planner parameters and apply the best point by `objective`, SLA goodput per GPU by default. This section is the field contract. For the concept, see [Periodic Retune](../../developer-guide/knowledge-base/modular-components/planner/planner-guide.md#periodic-retune-experimental); for windows, scoring, and decision timing, see [Planner Design](../../developer-guide/knowledge-base/modular-components/planner/planner-design.md#retune-experimental-simulation-only); to run it, see [Simulate Retune](../../developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx#simulate-retune-experimental). +## Retune (Experimental) + +`retune` makes the Planner periodically replay the last window of recorded traffic against a grid of router and Planner parameters and apply the best point by `objective`, SLA goodput per GPU by default. This section is the field contract. For the concept, see [Periodic Retune](../../developer-guide/knowledge-base/modular-components/planner/planner-guide.md#periodic-retune-experimental); for windows, scoring, and decision timing, see [Planner Design](../../developer-guide/knowledge-base/modular-components/planner/planner-design.md#retune-experimental); to run it in replay, see [Simulate Retune](../../developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx#simulate-retune-experimental); to run it on a live deployment, see [Planner Live Retune](../../kubernetes/auto-deployment/planner-live-retune.md). -**Experimental.** In this release, retune runs only inside planner-in-the-loop offline replay: `dynamo.replay.run_trace_replay` with `trace_format="dynamo"`, or `aisimulate predict --stack dynamo` on a Dynamo request trace. A live Planner rejects a config that sets `retune` at startup. +**Experimental.** Where retune runs depends on `executor`: + +- `local`, the default, runs only inside planner-in-the-loop offline replay: `dynamo.replay.run_trace_replay` with `trace_format="dynamo"`, or `aisimulate predict --stack dynamo` on a Dynamo request trace. A live Planner rejects it at startup. +- `dgdr` runs on a live Planner on Kubernetes, through DGDR v1beta2 Requests. It retunes three KV router knobs with fixed replicas and is meant for development clusters only; see [Planner Live Retune](../../kubernetes/auto-deployment/planner-live-retune.md). @@ -656,13 +666,13 @@ The Planner runs the builtin plugin pipeline by default. `scheduling` controls p Trace history before each window that warms the load predictors of candidate planners and of a newly selected predictor. With throughput scaling, replay reads this history from the replayed trace and, before its first request, from `load_predictor_warmup_trace`, which must then be a Dynamo request trace that ends at or before the first replayed request; without throughput scaling no history is read. `null` resolves to the minimum the predictors need: `throughput_adjustment_interval_seconds` × max(`prophet_window_size`, `kalman_min_points`) over the current predictor and the load-predictor candidates, or 0 without throughput scaling. An explicit value must be 0 or greater, and a value smaller than that minimum is rejected. - Delay from sweep launch to decision. A number is a fixed delay, 0 or greater. `measured` uses the sweep's own wall-clock time, capped at `timeout_seconds`, times `decision_latency_scale`; it depends on the replay host, so it is not reproducible across hosts or runs. The decision lands no earlier than one Planner tick after launch. See [Decision Timing](../../developer-guide/knowledge-base/modular-components/planner/planner-design.md#decision-timing) for the mechanics and [Decision Latency](../../developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx#decision-latency) for choosing a value. + Delay from sweep launch to decision. A number is a fixed delay, 0 or greater. `measured` uses the sweep's own wall-clock time, capped at `timeout_seconds`, times `decision_latency_scale`; it depends on the replay host, so it is not reproducible across hosts or runs. The decision lands no earlier than one Planner tick after launch. See [Decision Timing](../../developer-guide/knowledge-base/modular-components/planner/planner-design.md#decision-timing) for the mechanics and [Decision Latency](../../developer-guide/knowledge-base/modular-components/planner/aisimulate-planner-replay.mdx#decision-latency) for choosing a value. With executor `dgdr`, the decision lands when the sweep returns, so this must be 0, or `measured` with `decision_latency_scale: 1`. Multiplier on the measured sweep time, for example to model a production sweep host that is slower or faster than the replay host. Must be greater than 0. A value other than 1 requires `decision_latency_seconds: measured`. Above 1, set `timeout_seconds` ≤ `interval_seconds` / scale (the default, `interval_seconds`, is rejected); the modeled deadline is then scale × `timeout_seconds`. - Wall-clock deadline for one sweep; its workers are killed after it and the decision is `timeout`. `null` means `interval_seconds`. Must be greater than 0. Requires `timeout_seconds` ≤ `interval_seconds`. A numeric `decision_latency_seconds` must also be ≤ `timeout_seconds`; with `measured`, `decision_latency_scale` × `timeout_seconds` must also be ≤ `interval_seconds`. This keeps at most one sweep in flight. + Wall-clock deadline for one sweep; its workers are killed after it and the decision is `timeout`. `null` means `interval_seconds`. Must be greater than 0. Requires `timeout_seconds` ≤ `interval_seconds`. A numeric `decision_latency_seconds` must also be ≤ `timeout_seconds`; with `measured`, `decision_latency_scale` × `timeout_seconds` must also be ≤ `interval_seconds`. This keeps at most one sweep in flight. With executor `dgdr`, it is required, must be ≤ `interval_seconds` − 600, and includes the window cut. Metric each candidate replay is scored on, written as the DGDR v1beta2 `spec.objective.metric` value. `TTFT` and `E2ELatency` are minimized; the others are maximized. See [Objectives](#objectives) for what each one scores. With load or throughput scaling enabled, only `GoodputPerGPU` and `ThroughputPerGPU` are accepted; retune-only mode accepts every value. @@ -682,7 +692,7 @@ The Planner runs the builtin plugin pipeline by default. `scheduling` controls p Upper bound on the expanded grid size (the product of all domain sizes), on the values of each knob, and separately on the number of load-predictor candidates. Must be at least 1. A grid that exceeds it is rejected with the size of each dimension. For the candidate set a sweep replays, see [Candidates and the Sweep](../../developer-guide/knowledge-base/modular-components/planner/planner-design.md#candidates-and-the-sweep). - Number of sweep worker processes. Must be at least 1. Each worker replays one candidate at a time, so memory grows with this value and the window size. + Number of sweep worker processes of executor `local`. Must be at least 1. Each worker replays one candidate at a time, so memory grows with this value and the window size. Executor `dgdr` ignores it and uses `dgdr.run.parallelism`. Router search domain, written like the `router` section of `aisimulate recommend --stack dynamo` (see [Search domains](#search-domains)). A knob omitted from the map searches its AISimulate default domain; write a scalar to pin it. `default` searches the whole AISimulate default domain. `null` (or no `router` field) keeps every router parameter at its startup value; an empty map `{}` is rejected. A searched router requires `router_mode="kv_router"` (`router.policy: kv_router` in `aisimulate predict`) with the default worker-selection policy, no `overlap_score_weight`, and conditional disaggregation disabled. @@ -690,6 +700,77 @@ The Planner runs the builtin plugin pipeline by default. `scheduling` controls p Planner search domain, written like the `planner` section of `aisimulate recommend --stack dynamo` and limited to what a running Planner can change (see [Search domains](#search-domains)). `default` searches the parts of the AISimulate default domain that apply to this Planner. `null` (or no `planner` field) keeps every Planner field unchanged; an empty map `{}` is rejected. A written knob or group requires load or throughput scaling, and the scaling mode that uses it: for example, `load_sensitivity` needs load scaling, and `load_predictor` needs `optimization_target: sla` with throughput scaling. A searched field must have a value in the Planner config. `ttft_ms` and `itl_ms` are never searched. + + Where sweeps run. `local` replays the candidates in worker processes of the Planner and runs only inside planner-in-the-loop replay; a live Planner rejects it at startup with `planner.retune is replay-only with executor 'local'; use executor: dgdr`. `dgdr` runs each sweep as DGDR v1beta2 Requests from a live Planner on Kubernetes and requires the `dgdr` section below. **Experimental**, development clusters only: see [Planner Live Retune](../../kubernetes/auto-deployment/planner-live-retune.md). + + With `dgdr`, the Planner also rejects at startup a config that breaks a [configuration rule](../../kubernetes/auto-deployment/planner-live-retune.md#deployment-rules) of live retune: it requires `environment: kubernetes`, retune-only mode with no `planner` domain, a `router` domain that writes only `overlap_score_credit`, `prefill_load_scale`, `temperature`, `policy: kv_router` and `prefill_load_model: {type: none}`, the `timeout_seconds` and `decision_latency_seconds` limits above, a `dgdr.engine` that follows the rules below, and a largest grid shard of at most 100 points that fits the [time budget](../../kubernetes/auto-deployment/planner-live-retune.md#time-budget). + + Allowed values: local dgdr + + + Settings of executor `dgdr`. Required when `executor` is `dgdr`, and rejected otherwise. Rejects unknown fields. [Planner Live Retune](../../kubernetes/auto-deployment/planner-live-retune.md) shows a complete example. + + + + ReadWriteMany PersistentVolumeClaim, in the DGD namespace, that holds the frontend traces and the window files. The frontends and the Planner mount it. + + + Where the Planner container mounts `claim_name`. Must be absolute. + + + Claim-relative directory of the frontend trace segments, written as `///`. `/`-separated names without empty, `.` or `..` segments. `segments_dir` and `windows_dir` must not contain one another. + + + Claim-relative directory of the window files, written as `//-.jsonl.gz`, where `` is a digest of the window's start, end and cutoff. Same rules as `segments_dir`. + + + Window phase. Window `k` is `[phase_seconds + k × interval_seconds, phase_seconds + (k+1) × interval_seconds)` in epoch seconds, so a restarted Planner continues the same windows. Must be 0 or greater. + + + Hugging Face repository id the workers serve, written as each Request's `spec.model.name`. The Run Job has no Hugging Face credentials or model mount, so it must be a public repository that AISimulate can resolve. + + + GPU SKU, written as each Request's `spec.hardware.gpu.sku`. + + + The pinned AISimulate `engine` section, written as each Request's `spec.aisimulate.inputOverrides.engine`. Every value must be a scalar, each worker role of `mode` sets one `parallelism.preset` entry and both scheduler knobs, and `replicas` must not be set: the Planner fills them in from the DGD spec at each launch. See [Engine section](../../kubernetes/auto-deployment/planner-live-retune.md#engine-section). + + + Deadline of the subprocess that cuts a window file from the segments. Must be greater than 0. + + + Run Job start-up time reserved when deciding whether a shard can still finish before the window's deadline. Must be 0 or greater. + + + How each DGDR Run evaluates its points. Rejects unknown fields. + + + + AISimulate `optimizer.parallelism` of each Run. Must be at least 1. The Run Job's CPU and memory may lower the effective value. + + + AISimulate `optimizer.candidate_timeout_seconds` of each Run. Must be greater than 0. + + + Most grid shards, one Request each, per window. A shard is a Cartesian sub-grid of at most 100 points. The common-random-numbers probe and off-grid points are extra single-point Requests and do not count. Must be in [1, 16]. + + + + Windows whose finished Requests are kept for audit; older ones are deleted. Must be 0 or greater. + + + A window with more requests fails with `window_too_large`. Must be at least 1. + + + On switching from `advisory: true` to `false`, apply the last advisory recommendation if it is at most one interval old and the configuration is unchanged. + + + Duration of the leadership Lease, renewed every third of it or sooner. Must be at least 37: the Planner stops acting as leader 15 seconds (three 5-second Kubernetes request timeouts) before the Lease expires, and two renewals of up to 11 seconds each must fit before then. + + + Period of the router and clean-up rounds. Must be greater than 0. + + ### Objectives @@ -954,11 +1035,11 @@ The Planner enforces these cross-field rules at startup and rejects a config tha - `fpm_sample_bucket_size` must be a perfect square. - `global_planner_namespace` is required when `environment` is `global-planner`. - Under `optimization_target: load`, prefill modes require `prefill_scale_up_queue_tokens` > `prefill_scale_down_queue_tokens`, and decode modes require `decode_scale_up_kv_rate` > `decode_scale_down_kv_rate`. -- At least one scaling mode must be enabled. Under any `optimization_target` other than `sla`, the Planner forces load scaling on and throughput scaling off and ignores `ttft_ms` / `itl_ms`. The only exception is retune-only mode: `retune` set with both `enable_load_scaling` and `enable_throughput_scaling` explicitly `false` (see [Retune](#retune-experimental-simulation-only)). +- At least one scaling mode must be enabled. Under any `optimization_target` other than `sla`, the Planner forces load scaling on and throughput scaling off and ignores `ttft_ms` / `itl_ms`. The only exception is retune-only mode: `retune` set with both `enable_load_scaling` and `enable_throughput_scaling` explicitly `false` (see [Retune](#retune-experimental)). - When `ais_perf_model` is set, its `roles` must cover the deployment: `prefill` and `decode` for `disagg`, the matching role for `prefill`/`decode`, or `aggregated` for `agg`. Role keys must agree with `worker_type`. Unknown canonical fields and the retired `aic_perf_model` input are rejected. - `scheduling.scale_interval_seconds` must evenly divide `load_adjustment_interval_seconds` and, when throughput scaling is enabled, `throughput_adjustment_interval_seconds`. - When both scaling modes are enabled, `load_adjustment_interval_seconds` must be shorter than `throughput_adjustment_interval_seconds`. -- `retune` has its own cross-field rules, listed with its fields in [Retune](#retune-experimental-simulation-only). A live Planner rejects any config that sets `retune`. +- `retune` has its own cross-field rules, listed with its fields in [Retune](#retune-experimental). A live Planner rejects a `retune` with executor `local`; executor `dgdr` adds the [configuration rules](../../kubernetes/auto-deployment/planner-live-retune.md#deployment-rules) of live retune. ## Related pages diff --git a/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md b/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md index 580284f7fea2..b36b21d6769d 100644 --- a/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md +++ b/docs/fern/pages/reference/kubernetes-api/additional-resources/api-reference-k8s.md @@ -3321,7 +3321,7 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `inputOverrides` _[RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg)_ | InputOverrides is an RFC 7386 JSON Merge Patch over the input lowered from
the typed fields. It uses the upstream AISimulate JSON contract, including
snake_case names. Values here take precedence, and null deletes a lowered
field. The dgdr_run member must remain an object, and its operator-owned
transport, identity, and runtime fields name, max_candidates, snapshot_dir,
runtime_image, runtime_version_override, derive_runtime_images,
num_gpus_per_node, and namespace are operator-owned. Request admission
rejects them, and the Run controller enforces its values while lowering the
immutable snapshot. The operator runs the AISimulate version from its
configured planner image; Requests and Runs cannot select a different image
or version. | | Type: object
Optional: \{\}
| +| `inputOverrides` _[RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg)_ | InputOverrides is an RFC 7386 JSON Merge Patch over the input lowered from
the typed fields. It uses the upstream AISimulate JSON contract, including
snake_case names. Values here take precedence, and null deletes a lowered
field. The dgdr_run member must remain an object, and its operator-owned
transport, identity, and runtime fields name, max_candidates, snapshot_dir,
runtime_image, runtime_version_override, derive_runtime_images,
num_gpus_per_node, and namespace are operator-owned. Request admission
rejects them, and the Run controller enforces its values while lowering the
immutable snapshot. A typed spec.workload.trace is an underlay like the other
typed fields: it supplies the default traffic.source (type trace, format
dynamo and the mounted trace file in paths) and the recorded-arrival
traffic.load. Overrides can replace those values, including
traffic.source.type, format, and paths, and null deletes the named field or
parent object; the Run controller does not restore it. The typed
PersistentVolumeClaim source alone decides the Run Job's trace mount and its
claim dependency, so overriding a native path adds no volume mount.
AISimulate validates the resulting input when the Run's Job starts; admission
does not validate it against the AISimulate schema yet. The operator runs
the AISimulate version from its configured planner image; Requests and Runs
cannot select a different image or version. | | Type: object
Optional: \{\}
| #### Backend @@ -3359,8 +3359,8 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `name` _string_ | Name is the executor-defined stable branch identifier. | | | -| `rounds` _[RoundProgress](#roundprogress)_ | Rounds reports optimizer progress for this branch. | | | +| `name` _string_ | Name is the executor-defined stable branch identifier. | | Required: \{\}
| +| `rounds` _[RoundProgress](#roundprogress)_ | Rounds reports optimizer progress for this branch. | | Required: \{\}
| #### CandidateProgress @@ -3376,8 +3376,8 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `paretoFront` _integer_ | ParetoFront is the number of feasible points on the current Pareto front. | | | -| `published` _integer_ | Published is the number of Candidate resources currently referenced by this run. | | | +| `paretoFront` _integer_ | ParetoFront is the number of feasible points on the current Pareto front. | | Optional: \{\}
| +| `published` _integer_ | Published is the number of Candidate resources currently referenced by this run. | | Optional: \{\}
| #### DGDRRunPhase @@ -3432,9 +3432,9 @@ DynamoGraphDeploymentCandidate is one bounded, user-visible search result. | --- | --- | --- | --- | | `apiVersion` _string_ | `nvidia.com/v1beta2` | | | | `kind` _string_ | `DynamoGraphDeploymentCandidate` | | | -| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | | -| `spec` _[DynamoGraphDeploymentCandidateSpec](#dynamographdeploymentcandidatespec)_ | | | | -| `status` _[DynamoGraphDeploymentCandidateStatus](#dynamographdeploymentcandidatestatus)_ | | | | +| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | Optional: \{\}
| +| `spec` _[DynamoGraphDeploymentCandidateSpec](#dynamographdeploymentcandidatespec)_ | | | Required: \{\}
| +| `status` _[DynamoGraphDeploymentCandidateStatus](#dynamographdeploymentcandidatestatus)_ | | | Optional: \{\}
| #### DynamoGraphDeploymentCandidateSpec @@ -3461,7 +3461,7 @@ _Appears in:_ | `restart` _[Restart](#restart)_ | restart specifies the restart policy for the graph deployment. | | Optional: \{\}
| | `topologyConstraint` _[SpecTopologyConstraint](#spectopologyconstraint)_ | topologyConstraint is the deployment-level topology constraint. When
set, `spec.topologyConstraint.clusterTopologyName` names the ClusterTopology
CR to use. `spec.topologyConstraint.packDomain` is optional at this
level and can be omitted when only components carry constraints.
Components without their own `topologyConstraint` inherit from this value. | | Optional: \{\}
| | `experimental` _[DynamoGraphDeploymentExperimentalSpec](#dynamographdeploymentexperimentalspec)_ | experimental groups graph-level preview features whose API shape and
behavior may change in breaking ways between v1beta1 releases. | | Optional: \{\}
| -| `parameters` _[RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg)_ | Parameters are the resolved evaluation parameters for this point. | | Type: object
| +| `parameters` _[RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg)_ | Parameters are the resolved evaluation parameters for this point. | | Type: object
Required: \{\}
| #### DynamoGraphDeploymentCandidateStatus @@ -3494,9 +3494,9 @@ DynamoGraphDeploymentRequest is the Schema for simulation-backed deployment sear | --- | --- | --- | --- | | `apiVersion` _string_ | `nvidia.com/v1beta2` | | | | `kind` _string_ | `DynamoGraphDeploymentRequest` | | | -| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | | -| `spec` _[DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec)_ | | | | -| `status` _[DynamoGraphDeploymentRequestStatus](#v1beta2-dynamographdeploymentrequeststatus)_ | | | | +| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | Optional: \{\}
| +| `spec` _[DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec)_ | | | Required: \{\}
| +| `status` _[DynamoGraphDeploymentRequestStatus](#v1beta2-dynamographdeploymentrequeststatus)_ | | | Optional: \{\}
| #### v1beta2 DynamoGraphDeploymentRequestNativeSpec @@ -3515,7 +3515,7 @@ _Appears in:_ | `model` _[ModelSpec](#modelspec)_ | Model identifies the model whose deployment configurations are evaluated. | | Optional: \{\}
| | `backends` _[Backend](#backend) array_ | Backends lists one or more inference backends searched by AISimulate and used by generated candidates. | | Enum: [vllm sglang trtllm]
MinItems: 1
Optional: \{\}
| | `hardware` _[HardwareSpec](#v1beta2-hardwarespec)_ | Hardware bounds the accelerator configurations evaluated by the search.
When omitted, admission discovers SKU, budget and node width. An exactly empty object opts
out of operator hardware discovery and typed defaults. Admission completes every nonempty
object with any missing hardware values; explicit values remain authoritative. | | Optional: \{\}
| -| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the synthetic traffic model used for every candidate evaluation.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| +| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the traffic model used for every candidate evaluation: synthetic static
traffic or a recorded trace.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults, and a
trace workload receives no static default.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names, including the fields a trace lowers.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| | `objective` _[ObjectiveSpec](#objectivespec)_ | Objective defines scalar optimization. When omitted,
admission supplies the object. Its mode defaults to Optimize, its scalar metric defaults
to ThroughputPerGPU, and its absent SLA defaults to TTFT 2000 ms and ITL 30 ms. | | Optional: \{\}
| | `aisimulate` _[AISimulateSpec](#aisimulatespec)_ | AISimulate supplies public recommendation input to the operator's AISimulate execution bundle. | | Optional: \{\}
| | `recommendation` _[RecommendationSpec](#recommendationspec)_ | Recommendation controls bounded projection into DGDC resources. | | Optional: \{\}
| @@ -3557,7 +3557,7 @@ _Appears in:_ | `model` _[ModelSpec](#modelspec)_ | Model identifies the model whose deployment configurations are evaluated. | | Optional: \{\}
| | `backends` _[Backend](#backend) array_ | Backends lists one or more inference backends searched by AISimulate and used by generated candidates. | | Enum: [vllm sglang trtllm]
MinItems: 1
Optional: \{\}
| | `hardware` _[HardwareSpec](#v1beta2-hardwarespec)_ | Hardware bounds the accelerator configurations evaluated by the search.
When omitted, admission discovers SKU, budget and node width. An exactly empty object opts
out of operator hardware discovery and typed defaults. Admission completes every nonempty
object with any missing hardware values; explicit values remain authoritative. | | Optional: \{\}
| -| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the synthetic traffic model used for every candidate evaluation.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| +| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the traffic model used for every candidate evaluation: synthetic static
traffic or a recorded trace.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults, and a
trace workload receives no static default.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names, including the fields a trace lowers.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| | `objective` _[ObjectiveSpec](#objectivespec)_ | Objective defines scalar optimization. When omitted,
admission supplies the object. Its mode defaults to Optimize, its scalar metric defaults
to ThroughputPerGPU, and its absent SLA defaults to TTFT 2000 ms and ITL 30 ms. | | Optional: \{\}
| | `aisimulate` _[AISimulateSpec](#aisimulatespec)_ | AISimulate supplies public recommendation input to the operator's AISimulate execution bundle. | | Optional: \{\}
| | `recommendation` _[RecommendationSpec](#recommendationspec)_ | Recommendation controls bounded projection into DGDC resources. | | Optional: \{\}
| @@ -3598,9 +3598,9 @@ DynamoGraphDeploymentRun is one immutable accepted DGDR request and its executio | --- | --- | --- | --- | | `apiVersion` _string_ | `nvidia.com/v1beta2` | | | | `kind` _string_ | `DynamoGraphDeploymentRun` | | | -| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | | -| `spec` _[DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec)_ | | | | -| `status` _[DynamoGraphDeploymentRunStatus](#dynamographdeploymentrunstatus)_ | | | | +| `metadata` _[ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta)_ | Refer to Kubernetes API documentation for fields of `metadata`. | | Optional: \{\}
| +| `spec` _[DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec)_ | | | Required: \{\}
| +| `status` _[DynamoGraphDeploymentRunStatus](#dynamographdeploymentrunstatus)_ | | | Optional: \{\}
| #### DynamoGraphDeploymentRunSpec @@ -3624,7 +3624,7 @@ _Appears in:_ | `model` _[ModelSpec](#modelspec)_ | Model identifies the model whose deployment configurations are evaluated. | | Optional: \{\}
| | `backends` _[Backend](#backend) array_ | Backends lists one or more inference backends searched by AISimulate and used by generated candidates. | | Enum: [vllm sglang trtllm]
MinItems: 1
Optional: \{\}
| | `hardware` _[HardwareSpec](#v1beta2-hardwarespec)_ | Hardware bounds the accelerator configurations evaluated by the search.
When omitted, admission discovers SKU, budget and node width. An exactly empty object opts
out of operator hardware discovery and typed defaults. Admission completes every nonempty
object with any missing hardware values; explicit values remain authoritative. | | Optional: \{\}
| -| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the synthetic traffic model used for every candidate evaluation.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| +| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the traffic model used for every candidate evaluation: synthetic static
traffic or a recorded trace.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults, and a
trace workload receives no static default.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names, including the fields a trace lowers.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| | `objective` _[ObjectiveSpec](#objectivespec)_ | Objective defines scalar optimization. When omitted,
admission supplies the object. Its mode defaults to Optimize, its scalar metric defaults
to ThroughputPerGPU, and its absent SLA defaults to TTFT 2000 ms and ITL 30 ms. | | Optional: \{\}
| | `aisimulate` _[AISimulateSpec](#aisimulatespec)_ | AISimulate supplies public recommendation input to the operator's AISimulate execution bundle. | | Optional: \{\}
| | `recommendation` _[RecommendationSpec](#recommendationspec)_ | Recommendation controls bounded projection into DGDC resources. | | Optional: \{\}
| @@ -3667,13 +3667,13 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `scheduled` _integer_ | Scheduled is the number of unique evaluations accepted for execution. | | | -| `running` _integer_ | Running is the number of evaluations currently executing. | | | -| `feasible` _integer_ | Feasible is the number of completed evaluations that satisfy all constraints. | | | -| `infeasible` _integer_ | Infeasible is the number of completed evaluations that violate at least one constraint. | | | -| `failed` _integer_ | Failed is the number of evaluations that ended in an execution error. | | | -| `unsupported` _integer_ | Unsupported is the number of configurations rejected as unsupported before evaluation. | | | -| `cacheHits` _integer_ | CacheHits is the number of evaluations satisfied from a prior result. | | | +| `scheduled` _integer_ | Scheduled is the number of unique evaluations accepted for execution. | | Optional: \{\}
| +| `running` _integer_ | Running is the number of evaluations currently executing. | | Optional: \{\}
| +| `feasible` _integer_ | Feasible is the number of completed evaluations that satisfy all constraints. | | Optional: \{\}
| +| `infeasible` _integer_ | Infeasible is the number of completed evaluations that violate at least one constraint. | | Optional: \{\}
| +| `failed` _integer_ | Failed is the number of evaluations that ended in an execution error. | | Optional: \{\}
| +| `unsupported` _integer_ | Unsupported is the number of configurations rejected as unsupported before evaluation. | | Optional: \{\}
| +| `cacheHits` _integer_ | CacheHits is the number of evaluations satisfied from a prior result. | | Optional: \{\}
| #### GPUHardwareSpec @@ -3727,8 +3727,8 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `min` _float_ | Min is the lower end of the searched fraction range. | | ExclusiveMinimum: true
Maximum: 1
Minimum: 0
| -| `max` _float_ | Max is the upper end of the searched fraction range. | | ExclusiveMinimum: true
Maximum: 1
Minimum: 0
| +| `min` _float_ | Min is the lower end of the searched fraction range. | | ExclusiveMinimum: true
Maximum: 1
Minimum: 0
Required: \{\}
| +| `max` _float_ | Max is the upper end of the searched fraction range. | | ExclusiveMinimum: true
Maximum: 1
Minimum: 0
Required: \{\}
| #### KVCapacityFractionSpec @@ -3764,7 +3764,7 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `name` _string_ | Name identifies the model in the syntax accepted by the selected backend. | | MinLength: 1
| +| `name` _string_ | Name identifies the model in the syntax accepted by the selected backend. | | MinLength: 1
Required: \{\}
| #### NativeSearchInput @@ -3786,7 +3786,7 @@ _Appears in:_ | `model` _[ModelSpec](#modelspec)_ | Model identifies the model whose deployment configurations are evaluated. | | Optional: \{\}
| | `backends` _[Backend](#backend) array_ | Backends lists one or more inference backends searched by AISimulate and used by generated candidates. | | Enum: [vllm sglang trtllm]
MinItems: 1
Optional: \{\}
| | `hardware` _[HardwareSpec](#v1beta2-hardwarespec)_ | Hardware bounds the accelerator configurations evaluated by the search.
When omitted, admission discovers SKU, budget and node width. An exactly empty object opts
out of operator hardware discovery and typed defaults. Admission completes every nonempty
object with any missing hardware values; explicit values remain authoritative. | | Optional: \{\}
| -| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the synthetic traffic model used for every candidate evaluation.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| +| `workload` _[WorkloadSpec](#v1beta2-workloadspec)_ | Workload defines the traffic model used for every candidate evaluation: synthetic static
traffic or a recorded trace.
When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10;
an explicit empty workload receives the same shape. workload.static: \{\} opts out of all
traffic defaults. Any other static object is completed with its missing defaults, and a
trace workload receives no static default.
spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over
every lowered traffic field it names, including the fields a trace lowers.
The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction
search. Lowering independently defaults optimizer.max_trials to 320. | | Optional: \{\}
| | `objective` _[ObjectiveSpec](#objectivespec)_ | Objective defines scalar optimization. When omitted,
admission supplies the object. Its mode defaults to Optimize, its scalar metric defaults
to ThroughputPerGPU, and its absent SLA defaults to TTFT 2000 ms and ITL 30 ms. | | Optional: \{\}
| | `aisimulate` _[AISimulateSpec](#aisimulatespec)_ | AISimulate supplies public recommendation input to the operator's AISimulate execution bundle. | | Optional: \{\}
| | `recommendation` _[RecommendationSpec](#recommendationspec)_ | Recommendation controls bounded projection into DGDC resources. | | Optional: \{\}
| @@ -3898,7 +3898,7 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `name` _string_ | Name is the resource name. | | | +| `name` _string_ | Name is the resource name. | | Required: \{\}
| | `uid` _[UID](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#uid-types-pkg)_ | UID is the resource UID. | | Optional: \{\}
| @@ -3915,8 +3915,8 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `completed` _integer_ | Completed is the number of finished rounds. | | | -| `total` _integer_ | Total is the number of planned rounds. | | | +| `completed` _integer_ | Completed is the number of finished rounds. | | Required: \{\}
| +| `total` _integer_ | Total is the number of planned rounds. | | Required: \{\}
| #### StaticLoadSpec @@ -3956,11 +3956,98 @@ _Appears in:_ | `load` _[StaticLoadSpec](#staticloadspec)_ | Load selects the synthetic load model. When omitted from a nonempty static
workload, admission defaults it to concurrency 10. | | Optional: \{\}
| +#### TraceFormat + +_Underlying type:_ _string_ + +TraceFormat identifies the decoder of a recorded request trace. Lowering translates it to +AISimulate's native traffic.source.format. + +_Validation:_ +- Enum: [Dynamo] + +_Appears in:_ +- [TraceWorkloadSpec](#traceworkloadspec) + +| Field | Description | +| --- | --- | +| `Dynamo` | TraceFormatDynamo is the Dynamo request trace (dynamo.request.trace.v1) that the frontend
records, as JSON Lines or gzip-compressed JSON Lines. It lowers to AISimulate's dynamo format.
| + + +#### TracePVCSource + + + +TracePVCSource names one trace file on a PersistentVolumeClaim in the Request namespace. + + + +_Appears in:_ +- [TraceSourceSpec](#tracesourcespec) + +| Field | Description | Default | Validation | +| --- | --- | --- | --- | +| `claimName` _string_ | ClaimName is the PersistentVolumeClaim, in the Request namespace, that holds the trace.
The Run Job mounts only the trace file from it, read-only, and only into the AISimulate
container. Each Run protects the claim with the finalizer `nvidia.com/dgdr-run-`
until the Run completes or, for a Run deleted before it completes, until no Job of the Run
is left. Garbage collection deletes the Job of a deleted Run before the Job's pod, so the
finalizer of a deleted Run can be released while that pod still stops; Kubernetes
PersistentVolumeClaim protection keeps the claim until no pod uses it. Deleting the claim
cancels the Runs that read it. | | MaxLength: 253
MinLength: 1
Pattern: `^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$`
Required: \{\}
| +| `subPath` _string_ | SubPath names one trace file relative to the root of the claim; the Run Job mounts it as
the subPath of the claim volume. It is neither a directory nor a glob. It must not start or
end with '/' and must not contain empty, '.' or '..' segments or a NUL character, which no
file name can hold. The file must be complete before the Request is created and must never
be rewritten: every Run of the Request reads it when its Job starts. | | MaxLength: 1024
MinLength: 1
Required: \{\}
| + + +#### TraceReplaySpec + + + +TraceReplaySpec configures how AISimulate replays a decoded trace. + + + +_Appears in:_ +- [TraceWorkloadSpec](#traceworkloadspec) + +| Field | Description | Default | Validation | +| --- | --- | --- | --- | +| `arrivalSpeedupRatio` _float_ | ArrivalSpeedupRatio multiplies the recorded request-arrival rate. 1.0 preserves it.
When omitted, the recorded arrival times are replayed unchanged. | | ExclusiveMinimum: true
Minimum: 0
Optional: \{\}
| + + +#### TraceSourceSpec + + + +TraceSourceSpec identifies exactly one trace artifact. A PersistentVolumeClaim is the only +source in this release. + + + +_Appears in:_ +- [TraceWorkloadSpec](#traceworkloadspec) + +| Field | Description | Default | Validation | +| --- | --- | --- | --- | +| `persistentVolumeClaim` _[TracePVCSource](#tracepvcsource)_ | PersistentVolumeClaim reads the trace from a PersistentVolumeClaim in the Request
namespace. | | Optional: \{\}
| + + +#### TraceWorkloadSpec + + + +TraceWorkloadSpec replays a recorded request trace open loop at its recorded arrival times. + + + +_Appears in:_ +- [WorkloadSpec](#v1beta2-workloadspec) + +| Field | Description | Default | Validation | +| --- | --- | --- | --- | +| `format` _[TraceFormat](#traceformat)_ | Format identifies the trace decoder. Defaults to Dynamo. | Dynamo | Enum: [Dynamo]
Optional: \{\}
| +| `source` _[TraceSourceSpec](#tracesourcespec)_ | Source identifies the trace artifact. | | Required: \{\}
| +| `replay` _[TraceReplaySpec](#tracereplayspec)_ | Replay configures how the decoded trace is replayed. | | Optional: \{\}
| + + #### v1beta2 WorkloadSpec -WorkloadSpec defines the synthetic traffic model used for every candidate evaluation. +WorkloadSpec defines the traffic model used for every candidate evaluation. static and +trace are mutually exclusive. @@ -3972,7 +4059,8 @@ _Appears in:_ | Field | Description | Default | Validation | | --- | --- | --- | --- | -| `static` _[StaticWorkloadSpec](#staticworkloadspec)_ | Static defines AISimulate synthetic traffic. When omitted, admission supplies the
complete static default. An explicit empty object opts out; any nonempty object is
completed with omitted inputTokens, outputTokens and load defaults. | | Optional: \{\}
| +| `static` _[StaticWorkloadSpec](#staticworkloadspec)_ | Static defines AISimulate synthetic traffic. When omitted without a trace, admission
supplies the complete static default. An explicit empty object opts out; any nonempty
object is completed with omitted inputTokens, outputTokens and load defaults. | | Optional: \{\}
| +| `trace` _[TraceWorkloadSpec](#traceworkloadspec)_ | Trace replays a recorded request trace from a PersistentVolumeClaim in the Request
namespace. Admission supplies no static default when it is set. The Run Job mounts only the
trace file, read-only through a subPath of the claim, at /var/run/aisimulate/trace/trace.jsonl,
or trace.jsonl.gz when subPath ends in .gz. The trace lowers to the default traffic.source
and traffic.load, which spec.aisimulate.inputOverrides can replace or delete. | | Optional: \{\}
| diff --git a/docs/fern/pages/reference/kubernetes-api/dynamo-graph-deployment-request.mdx b/docs/fern/pages/reference/kubernetes-api/dynamo-graph-deployment-request.mdx index 0e12d9aa2644..f70be532557b 100644 --- a/docs/fern/pages/reference/kubernetes-api/dynamo-graph-deployment-request.mdx +++ b/docs/fern/pages/reference/kubernetes-api/dynamo-graph-deployment-request.mdx @@ -1114,6 +1114,165 @@ kubectl get dgdr my-model -n $NAMESPACE \ - The relationship is tracked via labels: `dgdr.nvidia.com/name` and `dgdr.nvidia.com/namespace`. - Additional resources (planner ConfigMaps) are created in the same namespace and labeled with `dgdr.nvidia.com/name`. +## Replay a recorded trace (v1beta2) + + +The `nvidia.com/v1beta2` API is experimental and disabled by default. See the +[installation guide](../../kubernetes/installation/install-dynamo.md) before you enable it. + + +A native `v1beta2` request can evaluate every candidate against a recorded request trace instead +of synthetic traffic. Set `spec.workload.trace` instead of `spec.workload.static`; the two are +mutually exclusive, and admission adds no static default when `trace` is set. To switch an existing +request from static traffic to a trace, set `spec.workload.static` to `null` in the same update. + +```yaml +apiVersion: nvidia.com/v1beta2 +kind: DynamoGraphDeploymentRequest +metadata: + name: qwen3-trace-replay +spec: + model: + name: Qwen/Qwen3-0.6B + workload: + trace: + # format defaults to Dynamo + source: + persistentVolumeClaim: + claimName: request-traces # PersistentVolumeClaim in this namespace + subPath: windows/2026-10-27T10-00.jsonl.gz + replay: + arrivalSpeedupRatio: 1.0 # optional; 2.0 replays arrivals twice as fast +``` + +Before you create the request: + +- Create the PersistentVolumeClaim in the request's namespace. Use a filesystem volume, not + `volumeMode: Block`. When the process that writes the trace keeps running, use `ReadWriteMany`. + With `ReadWriteOnce`, a Job scheduled to another node than the writer waits in + `ContainerCreating`, and with `ReadWriteOncePod`, a Job waits in `Pending` until the writer's pod + is gone. The Run reports `Running` while its Job waits. +- Finish writing the file. `subPath` names one file relative to the claim root; directories and globs + are not expanded. Every Run of the request reads the file when its Job starts, so do not rewrite + it afterwards. To replay several rolled `jsonl_gz` segments, concatenate them into one `.jsonl.gz` + file. A relative symlink to another file in the claim works. An absolute symlink does not, + because the kubelet resolves it on the node rather than inside the claim, and neither does a + symlink that leaves the claim. +- Make the file readable by others (for example, mode `0644` in `0755` directories) or by group + `65532`, which the Job's pod has as a supplemental group. The pod sets no `fsGroup`, so neither + the kubelet nor a CSI driver changes the ownership or permissions of anything on the claim. + +The operator mounts only this file, read-only and only into the AISimulate container, through a +`subPath` of the claim at `/var/run/aisimulate/trace/trace.jsonl.gz`, or at +`/var/run/aisimulate/trace/trace.jsonl` when `subPath` does not end in `.gz`. The kubelet resolves +`subPath` inside the claim and refuses to mount a file outside it. If the file does not exist when +the Job starts, the kubelet cannot mount it either, and the AISimulate container waits; see +[Trace claim lifecycle](#trace-claim-lifecycle). On a volume type whose node directory the kubelet +can write, such as `hostPath`, the kubelet instead creates an empty directory at `subPath` on the +claim, and the Run fails when AISimulate cannot read it. + +`format` defaults to `Dynamo`, the only supported format. It reads the Dynamo request trace +(`dynamo.request.trace.v1`) written by the frontend, as JSON Lines or as gzip-compressed JSON Lines +when the file name ends in `.gz`. See the [Request Trace Reference](../observability/request-traces.mdx). + +The operator lowers the trace to the AISimulate input `traffic.source` (`type: trace`, +`format: dynamo`, and `paths` with the mounted file) and `traffic.load.type: trace_timestamps`, +with `traffic.load.speedup` when `arrivalSpeedupRatio` is set. It adds no `traffic.stop.requests`. +These values are an underlay, like every value lowered from a typed field: +`spec.aisimulate.inputOverrides` is applied last as a JSON Merge Patch and can replace any of them, +including `traffic.source.type`, `traffic.source.format`, and `traffic.source.paths`. A `null` +override deletes the named field or parent object, and the operator does not restore it. For +example, this override deletes the lowered source: + +```yaml +spec: + aisimulate: + inputOverrides: + traffic: + source: null +``` + +The typed `persistentVolumeClaim` source alone decides which claim the Run's Job mounts and which +claim the Run protects. Overriding `traffic.source.paths` adds no volume mount. AISimulate validates +the resulting input when the Run's Job starts, and an input it rejects fails the Run; admission does +not validate it against the AISimulate schema yet. + +### Trace claim lifecycle + +The operator ties the claim's lifetime to the Runs that read it. It decides from structured state +only: the claim's `deletionTimestamp`, `volumeMode`, and phase, the Job's conditions, and the +request's `spec`. It copies container status messages into the Run but never interprets them. + +- **Protection.** Before a Run creates its Job, the operator adds the finalizer + `nvidia.com/dgdr-run-` to the claim, where `` is the Run's `metadata.uid`. Every + Run that reads the claim adds its own finalizer, so Runs of several requests can share a claim. The + operator removes a Run's finalizer once the Run is complete and none of its Jobs is still running, + or once the Run and its Jobs are deleted, for example when the request retires old Runs or is + deleted. Kubernetes garbage collection deletes the Job of a deleted Run before the Job's pod has + stopped, so the finalizer of a Run deleted before it completed can be removed while that pod still + mounts the file; `kubernetes.io/pvc-protection` keeps the claim until the pod is gone. +- **Deleting the claim cancels its Runs.** When the claim has a `deletionTimestamp`, the operator + stops the Job of every Run that reads it through foreground deletion, whether or not the Job's pod + has started. Once the Job and its pod are gone, the Run fails with reason `TraceSourceDeleted`, and + the operator then removes the Run's finalizer. A Job that already completed or failed before the + claim's deletion, and is not itself being deleted, is not stopped; its own outcome decides the Run. + Once the operator has started stopping a Job, the Run fails with `TraceSourceDeleted` even if the + Job is marked failed while its pod is removed. A Run whose claim is already being deleted before it + creates its Job fails with `TraceSourceDeleted` without creating anything. +- **A missing claim keeps the Run open.** When the claim does not exist, the Run stays `Pending` + with `Completed=False` and reason `TraceSourceMissing`, and starts once you create the claim. +- **A claim no Job can mount fails the Run before it starts.** Before the Run creates its Job, a + `volumeMode: Block` claim or a claim in phase `Lost` fails the Run with reason + `TraceSourceUnavailable`. After the Job exists, the operator does not check the claim's volume mode + or phase again; the Job's pod waits until the volume is usable, the Job fails, or you cancel the + Run. +- **A waiting container is a diagnostic, not a failure.** While a container of the Run's Job waits + on an error, such as a trace file the kubelet cannot mount (a missing file, an absolute symlink, + or a symlink that leaves the claim) or an image it cannot pull, the Run stays open with + `Completed=False`, reason `ContainerWaiting`, and the message + `container "" is waiting: : `, where `` and `` are copied + from the container's status. The operator does not delete the Job, and there is no timeout. The + kubelet keeps retrying the container, so fixing the cause, for example writing the missing file, + lets the same Run continue. Waiting in `ContainerCreating` or `PodInitializing` is not reported: + volume errors that keep a pod there, such as a `ReadWriteOnce` claim attached to another node, + appear only in the pod's events. + +To stop a Run that cannot continue, cancel it explicitly: change the request's `spec` (the Run +fails with reason `Superseded`), delete the request, or delete the claim (`TraceSourceDeleted`). A +Pending claim, such as one that waits for its first consumer, is accepted. The request starts one +Run per `metadata.generation`, so after a Run failed, change the request's `spec` to start a new +Run. + +```bash +# Active and recent Runs of the request +kubectl get dynamographdeploymentrequests.v1beta2.nvidia.com qwen3-trace-replay -n "$NAMESPACE" \ + -o jsonpath='{.status.activeRunRef}{"\n"}{.status.recentRunRefs}{"\n"}' + +# Why a Run failed or waits +kubectl get dynamographdeploymentrun qwen3-trace-replay-1 -n "$NAMESPACE" \ + -o jsonpath='{.status.conditions}' +``` + + +Kubernetes keeps a deleted claim `Terminating` while any pod that mounts it exists +(`kubernetes.io/pvc-protection`), including the completed pods of finished Jobs. The operator keeps +the Runs, Jobs, and pods of a request's five most recent finished Runs, so a deleted trace claim can +stay `Terminating` after every Run released it, until those Runs are retired or the request is +deleted. + + +If the operator is uninstalled, or the `v1beta2` API is disabled, while Runs are in flight, their +finalizers stay on the claim and keep it from being deleted. List them, then remove only the +`nvidia.com/dgdr-run-` finalizers (requires `jq`): + +```bash +kubectl get pvc request-traces -n "$NAMESPACE" -o jsonpath='{.metadata.finalizers}{"\n"}' + +kubectl get pvc request-traces -n "$NAMESPACE" -o json \ + | jq '.metadata.finalizers |= map(select(startswith("nvidia.com/dgdr-run-") | not))' \ + | kubectl replace -f - +``` + ## Known issues - **PCIe profiler data not yet available.** See the PCIe callout under [SKU format](#sku-format). diff --git a/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx b/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx index 9fb73f668e27..fe8bd8ae4813 100644 --- a/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx +++ b/docs/fern/pages/reference/kubernetes-api/full-api-reference.mdx @@ -10,12 +10,12 @@ subtitle: Full Kubernetes CRD + operator-config API reference, generated from th > [!WARNING] > **Auto-generated from source.** This page is regenerated from the Dynamo operator CRDs by `docs/fern/scripts/gen_kubernetes_api.py` and covers **every** type across all three API packages: the deprecated `nvidia.com/v1alpha1` surface, the supported `nvidia.com/v1beta1` surface, and the operator's own `operator.config.dynamo.nvidia.com/v1alpha1` configuration. The trimmed [DGD](dynamo-graph-deployment.mdx), [DGDR](dynamo-graph-deployment-request.mdx), and [DCD](dynamo-component-deployment.mdx) references cover only user-facing `v1beta1` fields. To edit the surface, change the Go types under `deploy/operator/api/` and let CI regenerate this page. -Dynamo publishes 4 Kubernetes API packages with 214 typed CRD and config sections. +Dynamo publishes 4 Kubernetes API packages with 219 typed CRD and config sections. 74 types 76 types -35 types +40 types 29 types @@ -4634,7 +4634,7 @@ planner image. **Appears in:** [DynamoGraphDeploymentRequestNativeSpec](#v1beta2-dynamographdeploymentrequestnativespec), [DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec), [DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec), [NativeSearchInput](#nativesearchinput) -InputOverrides is an RFC 7386 JSON Merge Patch over the input lowered from the typed fields. It uses the upstream AISimulate JSON contract, including snake_case names. Values here take precedence, and null deletes a lowered field. The dgdr_run member must remain an object, and its operator-owned transport, identity, and runtime fields name, max_candidates, snapshot_dir, runtime_image, runtime_version_override, derive_runtime_images, num_gpus_per_node, and namespace are operator-owned. Request admission rejects them, and the Run controller enforces its values while lowering the immutable snapshot. The operator runs the AISimulate version from its configured planner image; Requests and Runs cannot select a different image or version. +InputOverrides is an RFC 7386 JSON Merge Patch over the input lowered from the typed fields. It uses the upstream AISimulate JSON contract, including snake_case names. Values here take precedence, and null deletes a lowered field. The dgdr_run member must remain an object, and its operator-owned transport, identity, and runtime fields name, max_candidates, snapshot_dir, runtime_image, runtime_version_override, derive_runtime_images, num_gpus_per_node, and namespace are operator-owned. Request admission rejects them, and the Run controller enforces its values while lowering the immutable snapshot. A typed spec.workload.trace is an underlay like the other typed fields: it supplies the default traffic.source (type trace, format dynamo and the mounted trace file in paths) and the recorded-arrival traffic.load. Overrides can replace those values, including traffic.source.type, format, and paths, and null deletes the named field or parent object; the Run controller does not restore it. The typed PersistentVolumeClaim source alone decides the Run Job's trace mount and its claim dependency, so overriding a native path adds no volume mount. AISimulate validates the resulting input when the Run's Job starts; admission does not validate it against the AISimulate schema yet. The operator runs the AISimulate version from its configured planner image; Requests and Runs cannot select a different image or version. See [RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg). **Validation:** Type: object Optional: \{\} @@ -4663,13 +4663,15 @@ BranchProgress reports progress for one independently optimized AISimulate branc **Kind:** `type` **Appears in:** [DGDRRunProgress](#dgdrrunprogress) - + Name is the executor-defined stable branch identifier. +**Validation:** Required: \{\} - + Rounds reports optimizer progress for this branch. See [RoundProgress](#roundprogress). +**Validation:** Required: \{\} @@ -4682,10 +4684,12 @@ CandidateProgress reports bounded Kubernetes publication, not all internal trial ParetoFront is the number of feasible points on the current Pareto front. +**Validation:** Optional: \{\} Published is the number of Candidate resources currently referenced by this run. +**Validation:** Optional: \{\} @@ -4749,14 +4753,17 @@ DynamoGraphDeploymentCandidate is one bounded, user-visible search result. Refer to Kubernetes API documentation for fields of `metadata`. See [ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta). +**Validation:** Optional: \{\} - + See [DynamoGraphDeploymentCandidateSpec](#dynamographdeploymentcandidatespec). +**Validation:** Required: \{\} See [DynamoGraphDeploymentCandidateStatus](#dynamographdeploymentcandidatestatus). +**Validation:** Optional: \{\} @@ -4824,10 +4831,10 @@ See [DynamoGraphDeploymentExperimentalSpec](#dynamographdeploymentexperimentalsp **Validation:** Optional: \{\} - + Parameters are the resolved evaluation parameters for this point. See [RawExtension](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#rawextension-runtime-pkg). -**Validation:** Type: object +**Validation:** Type: object Required: \{\} @@ -4862,14 +4869,17 @@ DynamoGraphDeploymentRequest is the Schema for simulation-backed deployment sear Refer to Kubernetes API documentation for fields of `metadata`. See [ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta). +**Validation:** Optional: \{\} - + See [DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec). +**Validation:** Required: \{\} See [DynamoGraphDeploymentRequestStatus](#v1beta2-dynamographdeploymentrequeststatus). +**Validation:** Optional: \{\} @@ -4899,7 +4909,7 @@ See [HardwareSpec](#v1beta2-hardwarespec). -Workload defines the synthetic traffic model used for every candidate evaluation. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. +Workload defines the traffic model used for every candidate evaluation: synthetic static traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults, and a trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. See [WorkloadSpec](#v1beta2-workloadspec). **Validation:** Optional: \{\} @@ -4981,7 +4991,7 @@ See [HardwareSpec](#v1beta2-hardwarespec). -Workload defines the synthetic traffic model used for every candidate evaluation. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. +Workload defines the traffic model used for every candidate evaluation: synthetic static traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults, and a trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. See [WorkloadSpec](#v1beta2-workloadspec). **Validation:** Optional: \{\} @@ -5065,14 +5075,17 @@ DynamoGraphDeploymentRun is one immutable accepted DGDR request and its executio Refer to Kubernetes API documentation for fields of `metadata`. See [ObjectMeta](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.28/#objectmeta-v1-meta). +**Validation:** Optional: \{\} - + See [DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec). +**Validation:** Required: \{\} See [DynamoGraphDeploymentRunStatus](#dynamographdeploymentrunstatus). +**Validation:** Optional: \{\} @@ -5107,7 +5120,7 @@ See [HardwareSpec](#v1beta2-hardwarespec). -Workload defines the synthetic traffic model used for every candidate evaluation. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. +Workload defines the traffic model used for every candidate evaluation: synthetic static traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults, and a trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. See [WorkloadSpec](#v1beta2-workloadspec). **Validation:** Optional: \{\} @@ -5202,30 +5215,37 @@ EvaluationProgress reports unique candidate outcomes without creating one resour Scheduled is the number of unique evaluations accepted for execution. +**Validation:** Optional: \{\} Running is the number of evaluations currently executing. +**Validation:** Optional: \{\} Feasible is the number of completed evaluations that satisfy all constraints. +**Validation:** Optional: \{\} Infeasible is the number of completed evaluations that violate at least one constraint. +**Validation:** Optional: \{\} Failed is the number of evaluations that ended in an execution error. +**Validation:** Optional: \{\} Unsupported is the number of configurations rejected as unsupported before evaluation. +**Validation:** Optional: \{\} CacheHits is the number of evaluations satisfied from a prior result. +**Validation:** Optional: \{\} @@ -5274,14 +5294,14 @@ KVCapacityFractionRange defines the searched fraction of a candidate's KV capaci **Kind:** `type` **Appears in:** [KVCapacityFractionSpec](#kvcapacityfractionspec) - + Min is the lower end of the searched fraction range. -**Validation:** ExclusiveMinimum: true Maximum: 1 Minimum: 0 +**Validation:** ExclusiveMinimum: true Maximum: 1 Minimum: 0 Required: \{\} - + Max is the upper end of the searched fraction range. -**Validation:** ExclusiveMinimum: true Maximum: 1 Minimum: 0 +**Validation:** ExclusiveMinimum: true Maximum: 1 Minimum: 0 Required: \{\} @@ -5311,9 +5331,9 @@ ModelSpec identifies the model whose deployment configurations are evaluated. **Kind:** `type` **Appears in:** [DynamoGraphDeploymentRequestNativeSpec](#v1beta2-dynamographdeploymentrequestnativespec), [DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec), [DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec), [NativeSearchInput](#nativesearchinput) - + Name identifies the model in the syntax accepted by the selected backend. -**Validation:** MinLength: 1 +**Validation:** MinLength: 1 Required: \{\} @@ -5344,7 +5364,7 @@ See [HardwareSpec](#v1beta2-hardwarespec). -Workload defines the synthetic traffic model used for every candidate evaluation. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. +Workload defines the traffic model used for every candidate evaluation: synthetic static traffic or a recorded trace. When omitted, admission supplies 1024 input tokens, 128 output tokens and concurrency 10; an explicit empty workload receives the same shape. workload.static: \{\} opts out of all traffic defaults. Any other static object is completed with its missing defaults, and a trace workload receives no static default. spec.aisimulate.inputOverrides is applied as an RFC 7386 JSON Merge Patch and wins over every lowered traffic field it names, including the fields a trace lowers. The default evaluates a fixed concurrency, so it does not create a KV-capacity-fraction search. Lowering independently defaults optimizer.max_trials to 320. See [WorkloadSpec](#v1beta2-workloadspec). **Validation:** Optional: \{\} @@ -5470,8 +5490,9 @@ ResourceReference identifies a namespaced resource by name and UID. **Kind:** `type` **Appears in:** [DynamoGraphDeploymentRunStatus](#dynamographdeploymentrunstatus) - + Name is the resource name. +**Validation:** Required: \{\} @@ -5488,12 +5509,14 @@ RoundProgress reports completed and planned optimizer rounds. **Kind:** `type` **Appears in:** [BranchProgress](#branchprogress) - + Completed is the number of finished rounds. +**Validation:** Required: \{\} - + Total is the number of planned rounds. +**Validation:** Required: \{\} @@ -5546,18 +5569,112 @@ See [StaticLoadSpec](#staticloadspec). + +TraceFormat identifies the decoder of a recorded request trace. Lowering translates it to +AISimulate's native traffic.source.format. + + +**Kind:** `enum` +**Underlying type:** `string` +**Validation:** Enum: [Dynamo] +**Appears in:** [TraceWorkloadSpec](#traceworkloadspec) + +**Allowed values** + +- Dynamo TraceFormatDynamo is the Dynamo request trace (dynamo.request.trace.v1) that the frontend records, as JSON Lines or gzip-compressed JSON Lines. It lowers to AISimulate's dynamo format. + + + +TracePVCSource names one trace file on a PersistentVolumeClaim in the Request namespace. + + +**Kind:** `type` +**Appears in:** [TraceSourceSpec](#tracesourcespec) + + +ClaimName is the PersistentVolumeClaim, in the Request namespace, that holds the trace. The Run Job mounts only the trace file from it, read-only, and only into the AISimulate container. Each Run protects the claim with the finalizer `nvidia.com/dgdr-run-` until the Run completes or, for a Run deleted before it completes, until no Job of the Run is left. Garbage collection deletes the Job of a deleted Run before the Job's pod, so the finalizer of a deleted Run can be released while that pod still stops; Kubernetes PersistentVolumeClaim protection keeps the claim until no pod uses it. Deleting the claim cancels the Runs that read it. +**Validation:** MaxLength: 253 MinLength: 1 Pattern: `^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$` Required: \{\} + + + +SubPath names one trace file relative to the root of the claim; the Run Job mounts it as the subPath of the claim volume. It is neither a directory nor a glob. It must not start or end with '/' and must not contain empty, '.' or '..' segments or a NUL character, which no file name can hold. The file must be complete before the Request is created and must never be rewritten: every Run of the Request reads it when its Job starts. +**Validation:** MaxLength: 1024 MinLength: 1 Required: \{\} + + + + +TraceReplaySpec configures how AISimulate replays a decoded trace. + + +**Kind:** `type` +**Appears in:** [TraceWorkloadSpec](#traceworkloadspec) + + +ArrivalSpeedupRatio multiplies the recorded request-arrival rate. 1.0 preserves it. When omitted, the recorded arrival times are replayed unchanged. +**Validation:** ExclusiveMinimum: true Minimum: 0 Optional: \{\} + + + + +TraceSourceSpec identifies exactly one trace artifact. A PersistentVolumeClaim is the only +source in this release. + + +**Kind:** `type` +**Appears in:** [TraceWorkloadSpec](#traceworkloadspec) + + +PersistentVolumeClaim reads the trace from a PersistentVolumeClaim in the Request namespace. +See [TracePVCSource](#tracepvcsource). +**Validation:** Optional: \{\} + + + + +TraceWorkloadSpec replays a recorded request trace open loop at its recorded arrival times. + + +**Kind:** `type` +**Appears in:** [WorkloadSpec](#v1beta2-workloadspec) + + +Format identifies the trace decoder. Defaults to Dynamo. +See [TraceFormat](#traceformat). +**Validation:** Enum: [Dynamo] Optional: \{\} + + + +Source identifies the trace artifact. +See [TraceSourceSpec](#tracesourcespec). +**Validation:** Required: \{\} + + + +Replay configures how the decoded trace is replayed. +See [TraceReplaySpec](#tracereplayspec). +**Validation:** Optional: \{\} + + + -WorkloadSpec defines the synthetic traffic model used for every candidate evaluation. +WorkloadSpec defines the traffic model used for every candidate evaluation. static and +trace are mutually exclusive. **Kind:** `type` **Appears in:** [DynamoGraphDeploymentRequestNativeSpec](#v1beta2-dynamographdeploymentrequestnativespec), [DynamoGraphDeploymentRequestSpec](#v1beta2-dynamographdeploymentrequestspec), [DynamoGraphDeploymentRunSpec](#dynamographdeploymentrunspec), [NativeSearchInput](#nativesearchinput) -Static defines AISimulate synthetic traffic. When omitted, admission supplies the complete static default. An explicit empty object opts out; any nonempty object is completed with omitted inputTokens, outputTokens and load defaults. +Static defines AISimulate synthetic traffic. When omitted without a trace, admission supplies the complete static default. An explicit empty object opts out; any nonempty object is completed with omitted inputTokens, outputTokens and load defaults. See [StaticWorkloadSpec](#staticworkloadspec). **Validation:** Optional: \{\} + + +Trace replays a recorded request trace from a PersistentVolumeClaim in the Request namespace. Admission supplies no static default when it is set. The Run Job mounts only the trace file, read-only through a subPath of the claim, at /var/run/aisimulate/trace/trace.jsonl, or trace.jsonl.gz when subPath ends in .gz. The trace lowers to the default traffic.source and traffic.load, which spec.aisimulate.inputOverrides can replace or delete. +See [TraceWorkloadSpec](#traceworkloadspec). +**Validation:** Optional: \{\} +
diff --git a/docs/fern/scripts/tests/test_gen_kubernetes_api.py b/docs/fern/scripts/tests/test_gen_kubernetes_api.py index 787166cc6d6d..2b480cf15ee6 100644 --- a/docs/fern/scripts/tests/test_gen_kubernetes_api.py +++ b/docs/fern/scripts/tests/test_gen_kubernetes_api.py @@ -51,7 +51,7 @@ # owned by github.com/ai-dynamo/snapshot. "nvidia.com/v1alpha1": 74, "nvidia.com/v1beta1": 76, - "nvidia.com/v1beta2": 35, + "nvidia.com/v1beta2": 40, "operator.config.dynamo.nvidia.com/v1alpha1": 29, } EXPECTED_OPERATOR_DEFAULT_SECTIONS = ( diff --git a/docs/fern/translations/zh-CN/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md b/docs/fern/translations/zh-CN/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md index 8a5b477a8673..42729c33194d 100644 --- a/docs/fern/translations/zh-CN/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md +++ b/docs/fern/translations/zh-CN/pages/developer-guide/knowledge-base/modular-components/planner/planner-guide.md @@ -60,7 +60,7 @@ spec: advisory: true ``` -advisory 模式仅提供建议。Planner 会计算建议副本数、记录日志、将其导出为诊断信息，并显示在 HTML 报告中。这些建议不会作为扩缩容决策应用：Planner 不会执行扩缩容操作，也不会更改部署。planner-in-the-loop replay 也遵循此设置：模拟部署保持不变，只报告建议，retune 决策（retune 目前仅在离线仿真中运行）也只记录、不应用。 +advisory 模式仅提供建议。Planner 会计算建议副本数、记录日志、将其导出为诊断信息，并显示在 HTML 报告中。这些建议不会作为扩缩容决策应用：Planner 不会执行扩缩容操作，也不会更改部署。planner-in-the-loop replay 也遵循此设置：模拟部署保持不变，只报告建议，retune 决策（`executor: local` 的 retune 仅在离线仿真中运行；线上部署使用 `executor: dgdr`，目前为实验功能）也只记录、不应用。 ### 优化目标 diff --git a/examples/deployments/GKE/sglang/disagg.yaml b/examples/deployments/GKE/sglang/disagg.yaml index dda3390cddac..6cacd6c2acb2 100644 --- a/examples/deployments/GKE/sglang/disagg.yaml +++ b/examples/deployments/GKE/sglang/disagg.yaml @@ -20,8 +20,8 @@ spec: containers: - args: - | - export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:$LD_LIBRARY_PATH - export PATH=$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64 + export LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" + export PATH="$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64" /sbin/ldconfig nvidia-smi exec python3 -m dynamo.sglang --model-path Qwen/Qwen3-0.6B --served-model-name Qwen/Qwen3-0.6B --page-size 16 --tp 1 --trust-remote-code --skip-tokenizer-init --disaggregation-mode decode --disaggregation-transfer-backend nixl --disaggregation-bootstrap-port "12345" --host "0.0.0.0" @@ -45,8 +45,8 @@ spec: containers: - args: - |- - export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:$LD_LIBRARY_PATH - export PATH=$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64 + export LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" + export PATH="$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64" /sbin/ldconfig nvidia-smi exec python3 -m dynamo.sglang --model-path Qwen/Qwen3-0.6B --served-model-name Qwen/Qwen3-0.6B --page-size 16 --tp 1 --trust-remote-code --skip-tokenizer-init --disaggregation-mode prefill --disaggregation-transfer-backend nixl --disaggregation-bootstrap-port "12345" --host "0.0.0.0" diff --git a/examples/deployments/GKE/vllm/disagg.yaml b/examples/deployments/GKE/vllm/disagg.yaml index 206ff36cb908..8538fe401a54 100644 --- a/examples/deployments/GKE/vllm/disagg.yaml +++ b/examples/deployments/GKE/vllm/disagg.yaml @@ -21,8 +21,8 @@ spec: containers: - args: - | - export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:$LD_LIBRARY_PATH - export PATH=$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64 + export LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" + export PATH="$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64" /sbin/ldconfig exec python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --disaggregation-mode decode command: @@ -45,8 +45,8 @@ spec: containers: - args: - | - export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:$LD_LIBRARY_PATH - export PATH=$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64 + export LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" + export PATH="$PATH:/usr/local/nvidia/bin:/usr/local/nvidia/lib64" /sbin/ldconfig exec python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --disaggregation-mode prefill command: diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index 12a59d76ad99..067320487ffa 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -117,6 +117,7 @@ "aisimulate.config", "aisimulate.config.cli", "aisimulate.config.common", + "aisimulate.config.engine", "aisimulate.config_adapter", "aisimulate.generator.api", "aisimulate.generator.enumerate",