feat: persist and rehydrate scheduling state across root restarts

[kaiitunnz]

Purpose

The root holds the dispatcher's scheduling state in memory and task events flowed over fire-and-forget pub/sub, so any restart of the root server — a crash, a deploy, an image bump — lost every in-flight workflow. This makes the root restart-survivable: per-task scheduling state is persisted to Redis on each transition and rebuilt on startup, and task lifecycle events flow through a durable, replayable stream consumed from a persisted cursor. The cost is a brief control-plane pause while the server recreates and rehydrates; workers keep running their tasks throughout.

Rolling node image updates — recreate one node at a time, root last, with no full-cluster teardown — fall out as one application of this, via a per-node stack restart primitive. The rollout itself stays operator-driven.

Changes

• Durable scheduler state (task/runtime.py, registries/workflow.py, clients/redis.py) — persist per-task state, deps, and epoch index to Redis on each transition; TaskRuntime.rehydrate() rebuilds the DAG / ready queue / epoch frontiers on startup.
• Replayable task events (supervisor/services/relay_service.py, services/watchdog.py, services/monitoring.py) — task events move to a durable Redis stream consumed from a persisted cursor (advanced per entry); mark_succeeded / mark_failed are idempotent against replay.
• No Redis flush on shutdown (main.py) — state lifetime follows the Redis volume, not the process, so a restart rehydrates instead of starting empty; reset via stack clean. Removes the dead cleanup.py.
• Lifespan hydration (main.py) — rehydrate before the lifespan yields, so readiness is implicit (no traffic until scheduling state is restored).
• Restart hardening — every task transition (mark_started / mark_dispatched / mark_updated, mark_succeeded / mark_failed / mark_cancelled) is guarded so a replayed or conflicting event cannot regress a terminal task or double-count usage; rehydrated in-flight tasks get a heartbeat grace (WORKER_REHYDRATION_GRACE_SEC, default 120s) so a worker still catching up after a restart is not reaped while its pub/sub heartbeats are missing; and the task-event consumer warns when its cursor trails the bounded stream's trim horizon, resumes from the stream start when no cursor is persisted, drops malformed entries, and retries a failing handler up to TASK_EVENT_HANDLER_MAX_ATTEMPTS before dead-lettering it.
• Per-node restart primitive — stack restart [SERVICE ...] (cli/stack/.../stack.py) recreates one or more Compose services in place (--no-deps --force-recreate, optional --pull), leaving Redis up and draining managed workers first (once, via destroy_all_workers()).
• Docs + tests — docs/SERVICE_RESTARTS.md (generic in-place-restart framing), ARCHITECTURE.md, CLI.md, ENV.md; persistence/rehydration, terminal-guard, rehydration-grace, trim-detection, consumer retry/dead-letter, and multi-service stack restart tests.

Design

State is rebuilt, not snapshotted: task IDs are random per parse, so the DAG is reconstructed from persisted per-task records and overlaid with status. Events use a stream + persisted cursor (like the existing log streams) rather than consumer groups; the persist-transition → advance-cursor ordering gives at-least-once delivery, made safe by idempotent handlers and terminal-status guards. In-flight tasks stay assigned on rehydrate — survivors' completions arrive via the stream (no double execution), and departed workers are reclaimed by the watchdog only after the rehydration grace elapses (heartbeats drop during downtime, so a survivor looks briefly stale on resume). Only the server is recreated on a node; Redis stays up so durable state survives.

On a drain, a batch task's in-flight work requeues and re-runs (its executor cancel() is a no-op, so the worker's departure is recovered via WORKER_UNREGISTER); an SSH task — the one executor with real cancellation — is cancelled instead, since a live session cannot resume on another node.

Test Plan

uv run pre-commit run --all-files
uv run pytest tests/server/ tests/shared/ tests/sdk/test_node_client.py tests/cli/test_stack_restart.py

A local single-node e2e (root + one CPU worker, driven outside the repo tree) exercises the stack restart primitive end-to-end — see Test Result.

Test Result

pre-commit run --all-files passed (incl. workspace mypy). Unit tests pass across tests/server/ + tests/shared/ + the SDK/CLI cases above; the full tests/ suite (GPU worker paths) was not run this session.

The single-node e2e drives stack restart through four scenarios, all green (16/16 assertions):

• Queued work survives — a workflow submitted with no worker sits PENDING, survives stack restart server (only the server container is recreated; Redis untouched), and runs to DONE once a worker joins.
• Completed state persists — a DONE workflow is still DONE after a restart (idempotent replay, not re-run).
• In-flight SSH task — an SSH sleep caught at DISPATCHED is cleanly cancelled by the drain (terminal, not orphaned).
• Multi-service restart — restart redis_control server recreates both; a DONE workflow survives (state lives on the Redis volume).

The batch in-flight requeue-and-rerun path is covered by the rehydrate/recovery unit tests and a multi-node manual procedure.

---

Pre-submission Checklist

• I have read the contribution guidelines.
• I have run pre-commit run --all-files and fixed any issues.
• I have added or updated tests covering my changes (if applicable).
• I have verified that uv run pytest tests/ passes locally.
• If I changed shared schemas or proto definitions, I have checked downstream compatibility across Server and Worker.
• If I changed the SDK or CLI, I have verified the affected packages work (uv sync --all-packages --group ci --frozen).
• If this is a breaking change, I have prefixed the PR title with [BREAKING] and described migration steps above.
• I have updated documentation or config examples if user-facing behavior changed.

[timzsu]

@kaiitunnz The purpose of this PR looks sound, but I think the restart contract needs a bit more clarification before we rely on it operationally. It would be helpful to include an e2e pipeline that exercises the intended restart order and demonstrates that queued, dispatched, and in-flight tasks preserve state correctly.

A few design questions I’d like to clarify:

• When rolling workers, do we only restart idle workers, or do we restart all workers regardless of whether they are running tasks? If in-flight workers can be restarted, what is the recovery contract for their active tasks?
• The description says “workers keep running throughout,” but also says the server restart “drains workers first.” Are these referring to different restart modes or stages?
• What happens if a worker becomes unreachable while the root/server is restarting? What if it reconnects after the root has rehydrated, or while that worker’s own node is being updated?
• What is the ordering/atomicity contract for Redis replay? For example, do we persist the task transition before advancing the stream cursor, and are all replayed transitions idempotent if the server crashes between those operations

[kaiitunnz]

@timzsu This is an initial work on decoupling cluster states from the root server so that they can outlive root restarts, so this PR will cover only basic features. New features or advanced restart semantics can be added in future PRs. Below is my reply to your concerns.

• There are three cases for container version updating:

  • Both server and workers: The server and all directly managed workers are torn down and only the server and workers in the worker config file are started up. Other workers must be created and started manually by the operator.
  • Workers only: Individual workers to be updated are torn down and started with "version: " in the worker config.
  • Server only: The server and all directly managed workers are torn down, and only the server and workers in the worker config file are started up. However, the restarted workers will inherit the new version from the server if not manually specified in the worker config file. So, to use the old worker version, the operator must tear down and start the workers with a manually configured version.

  In all cases, the workers are torn down immediately, regardless of their statuses. The in-flight tasks will be considered failed and handled according to the retrying logic in the server, consuming the retry budget. We may make it not consume the retry budget in a follow-up PR.

• On a cluster, there are two kinds of servers: root and worker. To update the versions, we need to restart each individual server, so the workers directly managed by the server must be restarted. The statement that “workers keep running throughout" is from the perspective of the root restart. Previously, when we restart the root, we need to tear down the entire cluster, including all the workers. This PR decouples the root from the cluster states, allowing the worker servers' workers to run independently. Workers directly managed by the root server still need to be torn down. This works because restarting only the server leaves the Redis containers up, so the independent workers' events and heartbeats still reach the root's Redis.

• The normal node restarting logic is to first drain the workers and then restart the server. Intermittently disconnected workers will try to connect to the server's supervisor gRPC endpoint, but once the supervisor has restarted they will not be able to authenticate (as we mint a new token per worker startup, held in the supervisor's in-memory registry) and will be marked stale.

• Task events are appended to a durable Redis stream, and the consumer persists the task transition before advancing the stream cursor. So if the server crashes in between, the event is replayed (at-least-once delivery). mark_succeeded and mark_failed are idempotent — they short-circuit on an already-terminal task. However, mark_started and mark_dispatched are not yet guarded, so a replayed stale event could regress a terminal task and double-count usage. The transition, the workflow-set update, and the cursor are also separate Redis operations (and the stream is on a different Redis instance from the cursor), so there is no cross-operation atomicity; consistency relies on idempotent replay. We will add terminal-status guards to mark_started/mark_dispatched, advance the cursor per entry, and document this contract.

On the e2e pipeline, I have built my own e2e test suites and use them locally for my recent PRs. For this PR, the test covers queued and completed tasks across a server restart. It will be extended to exercise dispatched and in-flight tasks and the intended restart order. I will push the e2e test in this PR, but we have to work on standardizing e2e regression tests in a future PR.

[timzsu]

@kaiitunnz Thanks for the clarification, which looks good to me. May you clarify what the worker's server mean? Is it a RunMesh term?

[kaiitunnz]

@timzsu Worker server is a FlowMesh term. A FlowMesh cluster consists of two kinds of nodes: a root node and a worker node. The root node holds the cluster states and the Redis instances, while the worker nodes only manage the workers. A worker server is just a server in a worker node.

[timzsu]

Some comments, feel free to push back. In addition, EXECUTORS.md is not modified, yet the PR description said it has been changed.

[timzsu]

Two things to discuss. In addition, the PR description contains some changelog-like text (e.g., "falls back to "0-0" (replay from start), not "$" (tail)"), which you may want to clean up.

[timzsu]

One comment. In addition, it seems that a crash during persistence is still not recoverable. Is it acceptable?

[kaiitunnz]

One comment. In addition, it seems that a crash during persistence is still not recoverable. Is it acceptable?

Improving the reliability of TaskRuntime requires significant work. We will address this in a future PR (#72).

[timzsu]

LGTM.