It's not a DAG — it's a state machine
Jul 2026 9 min read
Why I ran a per-customer model-retraining pipeline as a state machine stored in Postgres instead of a workflow tool like Airflow — so it survives restarts and can wait hours on a slow GPU job without tying up a server.
- MLOps
- Postgres
- architecture
Most people reach for a DAG. When I needed to orchestrate the retraining of recommendation models — one model per customer, on a rolling schedule, for a multi-tenant platform — the obvious answer was Airflow or something like it. I ended up building a finite state machine persisted in Postgres instead, and it turned out to be the better fit. This is the story of why, and the handful of techniques that made it hold up in production.
The problem
The job sounds simple: keep every customer’s recommendation model fresh. In practice each retraining run had to:
- Submit work to a remote GPU service and then wait — a single training job could run up to five hours.
- Survive process restarts. Deploys, crashes, and rescheduling happen. A run that was three hours into waiting on a GPU job must not start over.
- Run safely with multiple instances. The orchestrator ran with several replicas for availability, and no two of them could drive the same customer’s pipeline at once.
- Retry intelligently. A failed submission should retry a few times; a job that succeeded remotely but failed to download shouldn’t pay for a full retrain.
A DAG framework can do most of this, but it comes with an assumption I couldn’t satisfy cleanly: a running task holds a worker. Waiting five hours on a remote job by occupying a slot the whole time is exactly the resource profile I wanted to avoid, and layering a “poke-and-defer” sensor on top of a DAG engine felt like fighting the tool. I also didn’t want to introduce a new piece of infrastructure — I already had Postgres, and Postgres turned out to be enough.
The model
I modelled each customer’s retraining run as a state machine. Each state is one row in a table with four fields: the transition function to run, the state to move to on_success, the state to move to on_failure, and a continue_run flag that decides whether the machine advances straight into the next state or is allowed to pause here first:
State = namedtuple("State", ["transition", "on_success", "on_failure", "continue_run"])
states = {
"INITIAL": State(check_readiness, "CHECK_LAST_TRAIN", "DONE", True),
"CHECK_LAST_TRAIN":State(resume_or_fresh, "DOWNLOAD_SAVE", "SUBMIT_JOB", True),
"SUBMIT_JOB": State(submit_train_job, "POLL_JOB", "JOB_FAILED", True),
"POLL_JOB": State(check_job_status, "DOWNLOAD_SAVE", "POLL_JOB", False), # yields
"JOB_FAILED": State(noop, "SUBMIT_JOB", "DONE", True), # retry
"DOWNLOAD_SAVE": State(download_and_save, "DONE", "DONE", True),
}
The shape is a mostly-linear spine with a couple of branches hanging off it:
Nothing exotic there yet. The interesting parts are in how the machine runs and how it remembers where it was.
Durability: the machine lives in Postgres
The state machine doesn’t hold its position in memory — it reads it from the database on every step. Two tables do the work:
- An append-only history table, one row per state visit: the state name, a
run_idfor the whole pass, the visit count, any accumulated errors, and whether the state yielded. - A pointer table with one row per customer, naming the current state.
When the orchestrator boots, it doesn’t restart anything. It looks up each customer’s pointer and resumes from exactly where the last process left off. A deploy in the middle of a five-hour wait costs nothing — the next process picks up polling as if nothing happened. This single decision — state is data in Postgres, not memory in a worker — is what turns “orchestrator” into “durable orchestrator.”
The clever bits
Four techniques carried most of the weight.
1. Keeping the call stack flat
The naive way to advance a state machine is for each transition to call the next one. Chain enough transitions — or add a self-looping poll state — and Python’s recursion limit ends the party. Instead, each step returns the next step instead of calling it directly. A simple loop runs each step it’s handed, one after another — so nothing piles up on the stack:
def trampoline(fn):
result = fn() # run the first step
while callable(result): # did it hand back another step?
result = result() # yes — run that one, then check again
return result # got a plain value instead: we're done
The whole thing rests on one rule that every step follows: to continue, a step returns the next step instead of calling it — unrun, as a zero-argument function like lambda: advance(machine). To finish, it returns a plain value instead. So trampoline just keeps pulling the next step and running it, and stops the moment a step hands back something that isn’t a function.
The payoff is that while loop. Because each step returns to the loop rather than calling the next step from inside itself, a thousand transitions are a thousand quick trips around one loop — not a call stack a thousand frames deep. That’s what lets the pipeline loop through the poll state a hundred times without ever tripping Python’s recursion limit.
2. Letting go of the worker while it waits
This is the piece that made a state machine beat a DAG for me. POLL_JOB is the state that waits for the remote training job to finish, and it’s the only one with continue_run set to False:
"POLL_JOB": State(check_job_status, "DOWNLOAD_SAVE", "POLL_JOB", False), # yields
False means don’t rush on to the next state — pause here if there’s nothing to do yet. So when the machine reaches POLL_JOB, it checks whether the remote job is done. If it isn’t, the machine writes a “come back later” note to the database and stops, letting go of the thread completely. Nothing sleeps, no worker sits waiting, no slot is tied up.
A scheduler tick a few minutes later resumes the machine, which reads the pointer, sees it’s still in POLL_JOB, and checks again. The pipeline spends five hours “in” the poll state while consuming zero resources between checks. Because that pause point is written to the database like everything else, a restart in the middle of the wait changes nothing.
def advance(machine):
state = machine.load_current_state() # from Postgres
if not state.continue_run and not job_ready():
machine.mark_yielded() # persist and bail
return None # trampoline stops; scheduler resumes later
ok, error, params = run_transition(state)
next_name = state.on_success if ok else state.on_failure
machine.save_transition(next_name, error)
return lambda: advance(machine) # hand the next step to the trampoline
3. Capping retries by counting visits
Every state has a maximum number of times it’s allowed to run. SUBMIT_JOB and its JOB_FAILED retry partner are each capped at three; POLL_JOB is capped at a hundred (a hundred polling intervals comfortably covers the remote service’s five-hour ceiling). The same number does two jobs at once: it stops a loop from running forever and it caps how many times we retry. When a state hits its limit, the machine doesn’t throw — it writes a terminal state with the reason “max retries reached” and moves on. Waiting a bit longer between attempts is handled inside the transition functions with ordinary backoff; the visit limit is the outer cap that survives restarts.
4. Turning errors into ordinary transitions
Every transition runs inside a wrapper that normalizes its return value to (ok, error, params) — and catches everything:
def run_transition(state):
try:
return state.transition() # returns (ok, error, params)
except Exception as e:
return (False, f"{type(e).__name__}: {e}", {})
An unhandled exception doesn’t crash the run; it becomes ok = False, which routes the machine down its failure branch and records the error string in the history row. The state machine is the error boundary. There’s no separate try/except scaffolding scattered through the orchestration logic, because the transition table already encodes what “failure” means for every state.
The payoff branch: resume without retraining
My favourite branch is the dashed one in the diagram. The two external steps — submitting/waiting on the GPU job, and downloading the results — aren’t in the same transaction, so there’s a window where the remote job succeeded but the download crashed. CHECK_LAST_TRAIN handles exactly this: if the last job finished successfully and there’s no new data since, it skips straight to DOWNLOAD_SAVE instead of submitting a fresh five-hour job. Recovering from a crash costs a download, not a retrain. Encoding that as a branch in the state table — rather than a special case buried in code — is what made it obvious and testable.
Operating it
Two things made this pleasant to run in production:
- The diagram is generated from the source of truth. A small script reads the same state table the engine runs and renders it to a graph, with the yield points and visit budgets annotated on the nodes. The architecture diagram can’t drift from the code, because it is the code.
- Stuck pipelines are recoverable by hand. Because state is just rows, an operator can repoint a customer stuck in
JOB_FAILEDback to a clean terminal state and let the next cycle start fresh. Durable state cuts both ways — it’s inspectable and editable, not a black box.
The takeaway
Reach for a DAG when the shape is “run these steps in this order.” Reach for a state machine when the shape is “wait, decide, maybe loop, survive a restart, and pick up exactly where you left off.” The moment your pipeline spends most of its wall-clock time waiting on something external, the ability to yield and durably resume matters more than the ability to express a dependency graph — and a few hundred lines over a database you already have will get you there.