[Serve][RFC] Auto-rollback for `serve deploy`

[abrarsheikh]

1. How rollouts work today

• The CLI/REST entry PUT /api/serve/applications/ and serve deploy both end at a single controller method, ServeController.apply_config (controller.py:1102).
• apply_config overwrites a single KV blob CONFIG_CHECKPOINT_KEY = "serve-app-config-checkpoint" and calls ApplicationStateManager.apply_app_configs(...), which per-app does either an in-place override (same code_version) or a fresh build_app task (application_state.py:1299).
• Each app independently transitions DEPLOYING -> RUNNING | DEPLOY_FAILED | UNHEALTHY via ApplicationState._determine_app_status (:760). DEPLOY_FAILED is set when any underlying deployment is DEPLOY_FAILED, which the deployment FSM in common.py (handle_transition) triggers on:
  • build_app task error (import/runtime env)
  • replica startup retries reaching min(max_constructor_retry_count, target_num_replicas * MAX_PER_REPLICA_RETRY_COUNT)
  • health-check failure while still UPDATING
• There is no previous-version retention today. The KV blob holds only the current target.

2. State machine introduced by this feature

stateDiagram-v2
    [*] --> Idle
    Idle --> Watching: apply_config new submitted
    Watching --> Promoted: all apps RUNNING
    Watching --> RollingBack: any app DEPLOY_FAILED
    Watching --> Cancelled: user submits another config
    RollingBack --> RolledBack: rollback reaches RUNNING
    RollingBack --> RollbackFailed: rollback also DEPLOY_FAILED
    Promoted --> Idle: promote pending to last_good
    RolledBack --> Idle
    RollbackFailed --> Idle: alert and stop, no auto retry
    Cancelled --> Watching: new submission becomes pending

Loading

3. Important considerations

• Trigger is ApplicationStatus.DEPLOY_FAILED only. No progress timeout, no UNHEALTHY-driven rollback. This avoids false positives on slow-starting workloads and post-rollout regressions.
• Scope is per-submission. Even if only one app in a multi-app config fails, all apps revert together to the last-good ServeDeploySchema. This matches the atomic semantics of apply_config.
• Declarative-only. Imperative serve.run / client.deploy_applications is out of scope for v1; the feature keys off the KV checkpoint that only declarative deploys produce.
• "Last good" definition = the most recent ServeDeploySchema where every app reached RUNNING after apply_config. Promotion happens only on success, so we never roll back to a config that itself was failing.
• First deploy has no previous good. If the very first apply_config fails, there is nothing to roll back to. Behavior: leave the partial deploy in DEPLOY_FAILED and emit a clear status message ("no previous successful config to roll back to"). Optionally tear down the failed apps if --rollback-on-failure is set; document the trade-off.
• Concurrent submissions. If a user pushes a new config while we are watching or rolling back, the new submission wins (cancels the watch / rollback). Reconciliation already serializes through the controller actor, so no extra locking is needed beyond a pending_submission_id.
• Avoid rollback ping-pong. If the last-good config also ends in DEPLOY_FAILED after rollback (e.g., environment regressed), do not auto-rollback again. Mark ROLLBACK_FAILED, log clearly, surface in serve status. Operator must intervene.
• Persistence and recovery. The watcher state must survive controller crashes. We will extend CONFIG_CHECKPOINT_KEY with an additional dict so _recover_state_from_checkpoint can resume the watch (and re-check whether to roll back) after a restart.
• Side effects of revert. Rolling back re-runs apply_config(last_good). Things that get reverted: code versions, replica counts, runtime envs, route prefixes, target_capacity, target_capacity_direction. Apps that the failed config newly added are deleted (since they are not in last-good). Apps the failed config deleted will be re-created from scratch — this resurrects deployments the user may have intentionally deleted, an important UX caveat to document.
• HTTP/gRPC/proxy options in ServeDeploySchema are not currently part of apply_config's persisted-config flow (they go through serve_head.py start-time). Rollback semantics deliberately scoped to applications + target_capacity to avoid restarting proxies.
• Imperative apps coexisting with declarative apps. apply_app_configs only deletes apps with api_type == DECLARATIVE. Auto-rollback re-applies a ServeDeploySchema, so it inherits the same scope: imperative apps (deployed via serve.run) are untouched.
• Idempotency. DeploymentState.deploy() already short-circuits on no-op; re-applying an unchanged app costs nothing. Rollback that touches only the failed app is essentially free for the others.
• Detection latency. ApplicationState.update() runs on each control loop tick, so DEPLOY_FAILED is observable within one or two control loops of the failure. The watcher polls inside the same loop, so no extra timer is required.
• Status surface. New fields on ApplicationDetails and ServeInstanceDetails:
  • rollout_status: WATCHING | NONE | ROLLING_BACK | ROLLED_BACK | ROLLBACK_FAILED
  • rolled_back_from_deployment_time, last_good_deployment_time
  • serve status prints these alongside the existing ApplicationStatus.
• Observability. Emit a controller log Auto-rollback triggered: app '<X>' is DEPLOY_FAILED, reverting to config from <ts>. Add metrics: serve_auto_rollback_triggered_total, serve_auto_rollback_succeeded_total, serve_auto_rollback_failed_total.
• Opt-in vs opt-out. Recommendation: opt-in via a new field on ServeDeploySchema (rollout_strategy.auto_rollback: bool = False and serve deploy --rollback-on-failure), to avoid surprising existing users. The value is persisted with the config so it survives controller restarts.
• Tests. Need coverage for: build-task failure rollback, replica-startup failure rollback, mid-rollout health failure rollback, controller restart during watch, controller restart during rollback, multi-app partial failure, rollback ping-pong (rollback also fails), auto_rollback=False keeps current behavior.

4. High-level approach

4a. Persist the previous good config

Today CONFIG_CHECKPOINT_KEY stores (deployment_time, target_capacity, target_capacity_direction, config_dict) (controller.py:1143). Replace with a forward-compatible dict:

{
    "version": 2,
    "current": {
        "deployment_time": float,
        "target_capacity": float | None,
        "target_capacity_direction": TargetCapacityDirection | None,
        "config_dict": Dict[str, app_config_dict],
        "auto_rollback_enabled": bool,
    },
    "last_good": { ... same shape ... } | None,
    "rollout_state": "WATCHING" | "NONE" | "ROLLING_BACK" | "ROLLBACK_FAILED",
    "rollout_started_at": float,
}

_read_config_checkpoint (controller.py:778) gets a v1 -> v2 migration shim. _recover_state_from_checkpoint (controller.py:760) restores the watcher state and decides whether to resume rollback.

4b. Add a RolloutSupervisor inside the controller

A small object owned by ServeController, ticked from the main control loop right after application_state_manager.update() (controller.py:585). Responsibilities:

• Track the current "pending submission" (app names, started_at, auto_rollback_enabled).
• On each tick, read app statuses via ApplicationStateManager.list_app_statuses(...).
  • If every app in pending.config_dict is RUNNING -> promote: last_good := pending, clear pending, persist KV, emit metric.
  • If any app is DEPLOY_FAILED and auto_rollback_enabled:
    • If last_good is None -> mark rollout_state = "NONE"; surface "no prior good config" message; stop watching.
    • Else -> set rollout_state = "ROLLING_BACK", persist KV, call self._apply_config_internal(last_good_config, is_rollback=True). After the rollback's reconciliation is itself watched: if it reaches RUNNING set ROLLED_BACK -> NONE and keep last_good unchanged; if it ends in DEPLOY_FAILED set ROLLBACK_FAILED and stop.
• New submission supersedes anything in flight: clear any WATCHING / ROLLING_BACK state and start a fresh watch.

4c. Wire opt-in flag

• Add an optional RolloutStrategySchema to ServeDeploySchema (schema.py:974):

  rollout_strategy: Optional[RolloutStrategySchema]
      auto_rollback: bool = False
  

• Add --rollback-on-failure / --no-rollback-on-failure to serve deploy in scripts.py (:341); the flag mutates the schema before ServeSubmissionClient(...).deploy_applications(...).

• Forward-compat preserved by extra="allow" already on ServeDeploySchema.

4d. Refactor apply_config to feed the supervisor

Split current apply_config into:

• public apply_config(config, deployment_time=0) (unchanged signature) -> calls internal helper with is_rollback=False and arms the supervisor.
• private _apply_config_internal(config, deployment_time, *, is_rollback: bool) -> does today's work (persist KV, call apply_app_configs, save_checkpoint). When is_rollback=True, the supervisor sets rollout_state="ROLLING_BACK" instead of WATCHING and does not arm a second auto-rollback (prevents ping-pong even if last_good itself fails).

4e. Surface state to users

• Extend ApplicationDetails / ServeInstanceDetails in schema.py with the new rollout fields.
• Update ServeController.get_serve_instance_details (:1258) and scripts.py:_get_status (:711) to render them.
• Emit info logs at supervisor transitions; add Prometheus counters via existing metrics_pusher infrastructure used by DeploymentStateManager.

[vaishdho1]

A few ideas around implementation:

1. Checkpoint Migration: Extract a Shared _parse_checkpoint helper

   The current checkpoint (CONFIG_CHECKPOINT_KEY) is a pickled 4-tuple: (deployment_time, target_capacity, target_capacity_direction, config_checkpoints_dict). It is unpacked independently in two places:

   • _read_config_checkpoint, which is used during recovery and by apply_config to compute target_capacity_direction.
   • get_app_configs, which is used by get_serve_instance_details to populate deployed_app_config in the status response.

   Both currently call pickle.loads(checkpoint) and destructure the result as a 4-tuple. The RFC's v2 checkpoint format, which is a versioned dict with current, last_good, and rollout_state, would break both paths unless there is migration logic.

   We could use a single _parse_checkpoint(raw_bytes) helper that both methods call. It checks isinstance(data, tuple) vs. isinstance(data, dict), normalizes to the v2 structure, and returns a typed dataclass. This avoids duplicating migration logic across the two read sites and makes the format easy to evolve further.

2. Overlap with serve deploy --merge

   Both features interact with the same code paths:

   • Checkpoint format: --merge needs to read existing per-app configs from the checkpoint to merge new apps in without deleting unmentioned ones. The v2 checkpoint format stores config_dict per app, which --merge would also read. Whichever lands first should design the format to accommodate both.

   • apply_app_configs deletion logic: Today this deletes any declarative app not in the new config. --merge would need to skip this deletion for unmentioned apps. Auto-rollback re-calls apply_config(last_good), which goes through the same deletion logic. If --merge changes this behavior, for example via a flag on apply_app_configs, rollback must use the full-replacement mode so apps added by the failed config are properly cleaned up.

   • apply_config split: --merge might also need a hook before _apply_config_internal to merge the new apps into the existing checkpoint. The refactor proposed for auto-rollback, splitting apply_config into public/private methods, is a natural extension point for --merge.

   Suggestion: Land the checkpoint v2 migration and apply_config split as a standalone preparatory PR. Both features benefit from it, this can unblock parallel development.

[abrarsheikh]

@vaishdho1 I don't think a migration is warranted here because when serve is deployed on a new cluster, it does not rely on the old checkpoint. Serve relies on checkpoints only between controller restarts on the same cluster.

When a new version of serve is deployed fresh using kuberay, we provision a new cluster and old checkpoint is not used at all. Let me know if that makes sense. In short, backward compatibility of the checkpoint data schema is not a design consideration.

[vaishdho1]

Got it. Thank you for the clarification

[abrarsheikh]

additional implementation considerations

1. We should change the replica to replica routing such that we only consider downstream replicas running the same version as the upstream replica
2. In the current implementation the way we handle rolling update is by shutting down 20% replicas first before starting 20% replicas on new version, that should change to start 20% replicas first before shutting down old replicas

[vaishdho1]

2. In the current implementation the way we handle rolling update is by shutting down 20% replicas first before starting 20% replicas on new version, that should change to start 20% replicas first before shutting down old replicas

The second point would mean that if there are already 100 replicas running we need to start say 20 more before stopping the old 20. Resource reservations and cluster capacity should be taken into consideration for this right, which is currently not an issue since we are already shutting down the replicas before starting

[abrarsheikh]

The second point would mean that if there are already 100 replicas running we need to start say 20 more before stopping the old 20. Resource reservations and cluster capacity should be taken into consideration for this right, which is currently not an issue since we are already shutting down the replicas before starting

correct, maybe we need a flag like --allow-surge, to indicate that we should over-provision first

[vaishdho1]

I'd like to pick this up. This will span multiple PRs:

• Checkpoint format change and apply_config refactor (also help serve deploy --merge)
• RolloutSupervisor core and persistence
• Schema, CLI flag, and status surface
• Metrics and tests