Task abandonment and the typestate lifecycle#

How daisy ensures the run loop terminates even when workers die in pathological ways, and why the per-task counters are protected by a Rust enum rather than runtime guards.

The two terminal outcomes#

Every task ends in exactly one of:

  1. Done — every block reached a terminal counter (success / permanent failure / orphan) through normal processing.

  2. Abandoned — the runner gave up after max_worker_restarts failed attempts. Remaining blocks are accounted as orphaned (they were never executed) so is_done() can become true and the loop can exit.

The whole abandonment subsystem exists to guarantee that is_done() becomes true and stays true even when in-flight messages from dying workers race against the abandonment transition.

Counters#

Each RunningTask (daisy-core/src/task_state.rs) has:

field

meaning

total_block_count

invariant for the whole run

ready_count

in the ready queue, not yet handed out

processing_count

dispatched to a worker, awaiting release_block

completed_count

terminal: success

failed_count

terminal: hit max_retries during real processing

orphaned_count

terminal: never executed (upstream failure or task abandonment)

worker_failure_count

total dirty exits — pure stat, can exceed max_worker_restarts

worker_restart_count

refill spawns the runner has performed, capped at the configured limit

Conservation (Invariant 1): total = ready + processing + completed + failed + orphaned + pending.

Done condition (Invariant 2): is_done() ⟺ ready = 0 ∧ processing = 0 ∧ pending = 0 ⟺ completed + failed + orphaned = total.

Stickiness (Invariant 3): once a task reaches a terminal state, no operation may move it back to non-terminal. This is the property the typestate enforces structurally.

The typestate#

pub enum TaskState {
    Running(RunningTask),
    Done(DoneTask),
    Abandoned(AbandonedTask),
}

Mutation methods (note_acquired, note_completed, note_failed_permanently, note_failed_for_retry, note_ready, note_orphaned, note_worker_died, note_worker_restarted) live on RunningTask. DoneTask and AbandonedTask are read-only snapshots of TaskCounters plus a reason. The compiler enforces that mutating a task requires pattern-matching to extract &mut RunningTask first — terminal variants don’t expose those methods at all.

Scheduler::running_mut(&mut self, task_id) is the single gate every counter writer goes through:

fn running_mut(&mut self, task_id: &str) -> Option<&mut RunningTask> {
    self.task_states.get_mut(task_id)?.as_running_mut()
}

Late events on terminal tasks return None here. The writer drops the message — no need to scatter is_done() checks.

Transitions#

  • TaskState::try_finalize_done() — called after every successful release. If counters balance, transitions Running → Done. Idempotent for terminal states.

  • TaskState::abandon(reason) — consumes the Running variant, accounts remaining blocks as orphaned, replaces self with Abandoned. Returns the count of orphaned blocks for logging. No-op for terminal states.

Both transitions use mem::replace with a placeholder. The HashMap key isn’t changed; only the value variant.

Worker death and restart counting#

Two worker shapes, both folded into the same death/restart cycle:

0-arg spawn_function: the user owns the loop. An exception inside their function unwinds out of PySpawnWorker::spawn, which returns Err. The Rust thread exits dirty.

1-arg process_function: the Rust thread owns the loop. After every client.release_block(), if the block came back with status Failed, the thread exits dirty. This is deliberate: a process-function failure is treated as a worker death, so it bumps worker_failure_count and goes through the same restart-cap accounting as a 0-arg crash. The result is that a buggy process_function is bounded by max_worker_restarts + 1 block failures before the runner gives up — the same invariant you get for 0-arg workers.

In both cases, the dirty exit triggers the same downstream flow:

  1. worker_exit_rx fires on the server’s select!.

  2. check_thread_health joins the thread, decrements allocator.alive, increments worker_failure_count (dirty exits only — it is a reporting stat).

  3. rebalance_workers decides whether to refill, governed by the start budget below.

  4. If the budget is exhausted and alive == 0, abandon_exhausted_tasks transitions the task to Abandoned and BFS-propagates the same transition (with reason UpstreamAbandoned) to all transitively-downstream tasks.

Because each block-failure is a death, in 1-arg mode the budget also bounds total block failures: a buggy process_function will fail at most max_worker_restarts + 1 blocks before the runner gives up.

The worker-start budget#

A task may ever start at most max_workers + max_worker_restarts workers, total, regardless of how or why previous workers exited — dirty, clean, or productive alike (worker_start_count on RunningTask; enforced in rebalance_workers). Starts beyond the first max_workers are accounted as restarts (worker_restart_count), so the abandonment condition and the tqdm ♻=N counter keep their meaning.

Why no exemptions: daisy workers are expected to be long-running — they may hold large models in memory and process many blocks over their lifetime. Recycling workers mid-task (exit after N blocks, rely on respawn) is not a supported pattern; fix the leak, or size max_worker_restarts generously. Exempting “good” exits would also reopen an unbounded-respawn hole: a spawn function whose worker silently fails to start (e.g. subprocess.run(..., check=False) around a command that can’t run on the node) exits cleanly with zero progress and would otherwise respawn forever.

The budget assumes the blocking-spawn contract: a spawn function runs for the lifetime of the worker it starts (sbatch --wait, bsub -K, subprocess.run). alive counts spawn-function threads — daisy has no other window onto external workers. A submit-and-return spawn function makes every submission look like an instant clean death: the budget is consumed by resubmissions while the real jobs (invisible until they connect) pile up, and the abandonment condition (no alive workers + budget exhausted + blocks remaining) can fire while queued jobs are still pending. Counting active per-task client connections toward alive would relax this; it is on the roadmap.

Two practical consequences:

  • Workers that exit because the queue drained consume no further budget — with no ready blocks, rebalance_workers never refills, so over-provisioning max_workers past the block count cannot abandon a completed task.

  • On preemptible or walltime-limited infrastructure, size max_worker_restarts to the number of worker deaths you expect over the run, or rely on done markers and simply start a fresh run to resume.

Race windows the typestate closes#

The motivating problem: when a worker crashes holding a block, three messages travel over independent transports (TCP ReleaseBlock, TCP Disconnect, mpsc WorkerStats) and arrive at the server in any order. Without the typestate gate, the abandonment transition could run while a stale release_block is still in the pipe — the late release would push counters past total_block_count and flip is_done() back to false, hanging the run loop forever.

Four such races, all closed by the gate:

  1. Late ReleaseBlock after direct abandonment — release_block early-returns when as_running_mut() is None.

  2. Downstream ready-block generation after transitive abandonment — queue_ready_block early-returns for non-running destination tasks. Without this, a healthy upstream’s release could repopulate ready_count on an already-abandoned downstream.

  3. Retry-path requeue after abandonment — the retry path goes through release_block and hits the same gate.

  4. Lost-block release after abandonment — bookkeeper-driven lost releases also go through release_block.

All four are now structural — unrepresentable rather than runtime-checked.

Tests#

  • tests/test_worker_restarts.py covers the cycle end-to-end:

    • test_restart_cap_terminates_task_and_marks_remaining_failed — direct abandonment with full restart usage

    • test_task_abandonment_does_not_block_other_tasks — abandonment of one task doesn’t stall an independent peer

    • test_abandoned_upstream_unblocks_downstream — transitive orphan propagation

    • test_abandonment_handles_in_flight_block_release_race — Race 1 stress test (5 iterations)

    • test_block_function_failure_kills_worker_and_drives_abandonment — 1-arg mode hooks into the same cycle

    • test_healthy_upstream_does_not_repopulate_abandoned_downstream — Race 2 regression

  • tests/test_keyboard_interrupt.py — SIGINT during the run aborts cleanly via the same shutdown path.