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 — terminal-status guards on mark_started / mark_dispatched / mark_updated so a replayed stale event can't regress a terminal task or double-count usage; heartbeat grace for rehydrated in-flight tasks (WORKER_REHYDRATION_GRACE_SEC, default 120s) so a worker catching up after a restart isn't reaped (heartbeats are pub/sub and drop during downtime); trim-horizon detection logs (instead of silently dropping) when the cursor falls behind a bounded stream; the event consumer survives a poison event (logs and skips) instead of the handler exception killing the thread.
• Per-node restart primitive — stack restart [SERVICE] (cli/stack/.../stack.py) recreates one Compose service in place (--no-deps --force-recreate, optional --pull), leaving Redis up and draining managed workers first; NodeClient.drain_workers() (sdk/stack/.../node_client.py) exposes the drain over destroy_all_workers().
• echo delay_sec (worker/executors/echo_executor.py) — holds a task in flight for testing dispatch/recovery paths.
• Docs + tests — docs/ROLLING_UPDATES.md, ARCHITECTURE.md, CLI.md, ENV.md, EXECUTORS.md; persistence/rehydration, terminal-guard, rehydration-grace, trim-detection, consumer-resilience, echo-delay, drain_workers, and 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.

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 tests/worker/test_echo_delay.py
bash tmp/e2e_tests/rolling_node_restart/run.sh

Test Result

pre-commit run --all-files all passed (incl. workspace mypy). Unit tests: 444 passed (tests/server/ + tests/shared/ + the SDK/CLI/echo cases above); the full tests/ suite (GPU worker paths) was not run this session. E2E on a single root + CPU worker: 14/14 — queued and completed workflows survive a stack restart server, only the server container is recreated (Redis untouched), rehydrated work runs to DONE, and an in-flight task caught at DISPATCHED survives the restart, is requeued (attempts 0→1), and runs to DONE once a worker is back.

---

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.