How I stopped losing 3 hours of GPU work every time a job crashed

#python #opensource #devops #testing

We had a pipeline that could run for hours. GPU time is expensive. A crash at 95%
was the same as a crash at 0% — except we were angrier.

I got tired of it. Not in the "we need a framework" way. In the "this is stupid" way.

What we had before

A long script:

research → script → render (forever) → assemble → done

If the process died in render, we started over. Every time.
People say "use a job queue". We had workers. We still lost time, because
progress lived only in RAM.

The fix (boring on purpose)

After every stage, write down that you finished it.

research   [done]
render     [done]   ← crash here? next run skips research+render
assemble   [pending]

Then:

run() executes pending stages
resume() skips done, retries failed
state is written atomically (tmp file → rename), so a kill mid-write doesn't corrupt everything
workers send a heartbeat; watchdog notices when they go silent

That's it. No magic scheduler. Just refusing to forget finished work.

Why not "just checkpoint tensors"?

We don't need mid-layer optimizer state for every job. Most of our cost is
stage-level work: research, render, package. Stage checkpoints cover 95% of the pain
with 5% of the complexity.

Trade-off on purpose: a stage still re-runs from its start. That's acceptable.
Losing the whole pipeline is not.

What broke along the way

Non-atomic state writes. Kill during save → corrupted JSON → resume lied to us. Fixed with write-to-temp + os.replace.
Retries that duplicated side effects. Stage contract now says: re-entry must be safe (idempotent), or the stage owns its own local temp dir.
"Busy" hung workers. Looked alive in the process list. Heartbeat file + staleness check.

Numbers

Long jobs: tens of minutes to hours
Before: crash = full re-run
After: crash = re-run from last finished stage
Tests: 12 behavioral contracts (resume, fail-fast, atomic state, stale heartbeat)

Practical checklist (copy/paste)

[ ] Checkpoint after each stage, not "sometime later"
[ ] Atomic write (tmp + rename). Kill mid-write must not corrupt state
[ ] resume() skips done, retries failed
[ ] Stages should be idempotent on re-entry
[ ] Heartbeat file + staleness check for hung workers
[ ] Fail-fast on poison stages
[ ] Test: crash after stage 1 → resume does stage 2 only

Rule of thumb

If re-running the job costs more than 5 minutes — checkpoint it.
If it costs GPU-hours — checkpoint it today.

Takeaways

Checkpoint I/O is cheap. Re-running 3 hours of GPU is not.
Atomic writes matter more than fancy progress bars.
Resume is a feature, not an ops afterthought.
Heartbeat ≠ "it works". It's only "it's alive".
Fail-fast on poison stages beats "keep going and hope".

Open-sourced here: checkpointed-task-pipeline

Related reliability pieces:

Top comments (1)

Pierre-Laurent Medori • Oct 16

One thing I hit when testing a crash between two stages: state.json stores the status, but not what the stage returned. data only goes into the in-memory context, so after a crash in assemble, resume skips render, assemble runs with ctx["frames"] == None, and then gets marked done. The status survived the crash; the output didn't. Persisting data next to the status (same atomic write) would close it, and test_context_flows_between_stages could assert got == 42 across a resume.

Same lesson one level up, for LLM subagents and workers. Their RAM is the context window, and compaction or a timeout at 95% is your crash. What works for us:

a handoff file after each step, not just in the final message: what's done, what it produced, where it lives, what's next
"done" means the result is written down; an agent saying "done" is a heartbeat, not proof
side effects (send, publish, deploy) logged with their IDs, so a resumed worker doesn't fire them twice

The payoff is tokens. Each worker can run its window to the end instead of wrapping up early to protect its state, and the next one spends its budget on new work, not on re-reading and re-deriving what was already found. Re-reading is the agent version of re-running 3 GPU hours.