Native session retry support for batch sessions
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 17h 7m
- Merged PRs (30d)
- 358
Description
## Summary
[http://Backend.AI](http://Backend.AI) core lacks automatic retry for failed batch sessions. A `BATCH` session that fails (image pull, transient agent failure, OOM, scheduler timeout, kernel non-zero exit) becomes terminal in `ERROR`, and the user must manually re-create it. Every higher-level orchestrator on top of [http://Backend.AI](http://Backend.AI) ends up re-implementing retry — duplicated logic and inconsistent semantics. This issue proposes adding native, opt-in retry to core.
## Proposal
Introduce a `RetryPolicy` (Pydantic DTO) accepted at session creation, modeled after Apache Airflow's well-understood retry parameters but adapted for [http://Backend.AI](http://Backend.AI) :
- `max_retries: int = 0` (0 disables; preserves current behavior)
- `retry_delay: float = 60.0` (seconds)
- `backoff: "fixed" | "exponential" = "fixed"`, with `backoff_multiplier: float = 2.0`
- `max_retry_delay: float | None = 3600.0` (cap; hard ceiling 24 h)
- `jitter: "none" | "deterministic" | "random" = "deterministic"`, `jitter_ratio: float = 0.25`
- `eligible_causes: set[RetryEligibleCause]` — defaults to all retriable causes (`AGENT_TRANSIENT`, `SCHEDULER_TIMEOUT`, `IMAGE_PULL_FAILURE`, `KERNEL_NONZERO_EXIT`, `OOM_KILLED`, `UNKNOWN`). `USER_CANCELLED` / `VALIDATION_ERROR` / `QUOTA_EXCEEDED` are hardcoded non-retriable and cannot be opted into retry. Setting `max_retries > 0` is the only knob a typical user needs.
**Defaults precedence** (Airflow-style): per-session > project/domain default > cluster default in etcd.
## Data Model
One Alembic migration adds to `sessions`:
- `parent_session_id` (nullable self-FK)
- `retry_count` (int, default 0)
- `max_retries` (int, default 0)
- `retry_policy` (JSONB, nullable)
- `retry_cause` (text, nullable)
No new history table — the chain is a linked list of real `SessionRow`s, each with its own status, kernels, logs, and `status_data`.
## Decision & Dispatch
- New handler in `event_dispatcher/handlers/` subscribes to `session.terminated` / `session.error`.
- `classify_failure(session, status_data) → RetryEligibleCause` central table.
- Eligibility check, atomic `select_for_update`, deterministic `creation_id` for idempotency, then schedule a fresh `SessionService.create_from_params()` with the computed delay (no blocking sleep in the handler).
- Child inherits `retry_policy`; sets `parent_session_id` and `retry_count = parent + 1`.
**No new** `RETRYING` status. Parent goes to `ERROR` as today; child starts `PENDING`. A computed `retry_state` field on the API tells clients `attempt N of M`. Avoids touching the scheduler state machine.
## Surface Area
- REST v2: `POST /sessions` accepts `retry_policy`; `GET /sessions/{id`} returns retry fields; `GET /sessions/{id}/attempts` returns the chain.
- GraphQL v2: mirror in `api/gql/session/types.py` (`parentSession`, `retryCount`, `maxRetries`, `retryPolicy`, `retryCause`, `retryChain` resolver).
- Client SDK v2 + CLI v2 (`./bai session info` shows attempt N/M and links to parent).
- **No retry mutation in v1** — manual retry deferred to v2 once the auto path is stable.
## Observability
- Counters: `bai_session_retry_scheduled_total{cause`}, `bai_session_retry_exhausted_total{cause`}, `bai_session_retry_succeeded_total`.
- Events: `session.retry_scheduled`, `session.retry_exhausted` (consumable by webhook plugin; replaces the need for an Airflow-style `on_retry_callback`).
- Audit log entry per retry dispatch.
## Scope
**In v1:** batch sessions only, automatic retry, opt-in via `max_retries > 0`. Default `0` preserves current behavior across the entire surface.
**Deferred:** manual retry mutation, interactive sessions, OOM-with-resource-bump, DAG-style cross-session orchestration.
## Open Questions
- Quota accounting: do retries count against concurrent-session limits? Likely yes; needs a product call.
- Retry-storm kill switch: an etcd-backed cluster-level disable flag is wise.
- Manual retry in v2: counts toward `max_retries` or independent? Decide before exposing.
## Priority & Estimate
**P1 (Medium-High).** ~3–4 weeks for one engineer. Low blast radius (opt-in, default off).
## Detailed Plan
`docs/investigation/native-session-retry-plan.md` (in-repo working doc) covers the full design including Pydantic schema, backoff formula, work breakdown, and risk analysis.
JIRA Issue: BA-5850
Contributor guide
Assessment
This issue has not been assessed yet.