Union.ai
Flyte
AI

Inside Union: Leases, the Scheduling Engine Behind Flyte 2.0

Ketan Umare

Ketan Umare

Yee Hing Tong

Yee Hing Tong

Sergey Vilgelm

Sergey Vilgelm

Paul Dittamo

Paul Dittamo

The decision path of the Union engine: a leader-elected service that makes every scheduling decision from memory, writes each one down as a durable lease, and recovers from its own restart in seconds. And the cluster worker that holds those leases inside every cluster.

This is part 3 of Inside Union, a five-part series on the engine behind Flyte 2.0. It follows part 1 on the durable AI runtime and part 2 on the replay log, and continues with part 4 on resource-aware scheduling and part 5 on observability.

"A distributed system is one in which the failure of a computer you didn't even know existed can render your own computer unusable."
— Leslie Lamport

Why leases

There are two well-trodden designs here, and we had lived inside one of them. Durable-execution engines keep queues in the control plane. Workers pull steps and the log is safe, but the control plane can't do anything to infrastructure. It can't create, resize, or repair the compute a step needs, because all it owns is the queue. Flyte 1 sat at the other end: a controller running inside the cluster, using etcd as its state store, reconciling every decision locally against Kubernetes. That gives you full command of the infrastructure and binds you to one cluster to get it. Past one cluster you need credentials into every cluster, a Kubernetes client per cluster, and a reconciliation loop that can tell "the pod is gone" from "the API server is slow".

We wanted both halves at once: a control plane that commands infrastructure without living inside it, and one that knows nothing about Kubernetes. It knows about actions, which are immutable descriptions of work, and leases, which are grants to execute one. A lease has a holder, an expiry, and a generation. The worker in each cluster holds leases, does the work, and heartbeats. The tracked lease is what promises the work is durable, and the heartbeats are what keep it live. When heartbeats go missing, the failure is detected and resolved.

Failure is the common case at this scale. Infrastructure fails, and a long run usually hits at least one failure. Leases make every failure look the same: a crashed worker, a network partition, and a deleted cluster all end as a lease that stops heartbeating, expires, and is reassigned. The engine knows why a lease ended and records it on the action. An expiry or an OOM is a system error with its cause, and a raised exception is a user error. That distinction is what lets the retry policy and the UI treat "the node went away" differently from "the code is wrong". Cleanup is itself a lease. And the control plane never holds an inbound connection into a customer's cluster.

This post covers how the engine meets most of the requirements from the first post: survive anything, scale by adding machines, stay multi-tenant without a person in the loop, span clouds, enforce concurrency limits, propagate infrastructure errors with their cause, and keep every action's state visible.

The first post described the split between a write path, the highly sharded action store that absorbs every write from every run, and a decision path. The scheduling engine is that decision path, and we made it deliberately un-distributed at the point of decision. Each shard is owned by one leader-elected process that holds the whole picture in memory: every pending action, every connected worker across every cluster, every queue and quota, every lease and its last heartbeat. One owner means one writer, which means no locks on the decision path, no optimistic concurrency, and no transaction machinery. It also means tenants are isolated by construction. A tenant's work is scheduled by its own owner, so a million-action backfill from one customer can't slow a five-task run from another, and adding capacity to the control plane is adding shards and letting tenants land on them. Concurrency limits, meaning how many runs and actions a tenant or a run may have in flight, are enforced right here at enqueue, where the owner already has the counts. The goal we set at the start, the lowest-latency scheduler we could build that still issues durable promises, comes down to this: decide from memory, then write the decision down before acting on it.

The guarantees, stated up front. A run lease is held by one worker at a time, and only the holder of the current generation can complete it, so assignment is at-most-once. Retries and partitions mean the work itself can run twice, so execution is at-least-once, and tasks should be idempotent or carry the action name into external writes. Every transition is durable before it's acknowledged, with one exception, heartbeat extension, which we'll get to.

Two kinds of lease

What a lease guarantees is easiest to read off its life cycle.

Every action gets exactly one run lease. It starts unassigned, is sent to a worker, and completes when the worker reports a result or the lease expires. The engine then records the attempt durably in the action store and decides whether to finish, retry, or cascade an abort to children. Outcomes are always persisted, and success, failure, timeout, and abort are all handled the same way.

The second kind is the finalize lease, and it exists to guarantee cleanup. It's sticky: routed to the cluster, and where possible the worker, that ran the original attempt, because that's where the pods and the Ray cluster are. Cleanup never waits behind new work, and a user's deadline never kills it. It has its own generous bound.

Generations

Leasing systems are hard to get right because of races, meaning two actors acting on different views of who holds a lease.

The races are concrete. A worker crashes and reconnects within the expiry window, still holding leases from before. A lease expires and is reassigned to worker B while worker A returns from a partition and reports success. The control plane restarts with an empty worker table and a zombie heartbeats for a lease that expired during the restart. The expiration monitor's snapshot is a few milliseconds old and would reap a lease that was just reassigned.

We solve all of these with one rule. Each lease carries a generation, a counter incremented on every assignment and folded into the opaque token the worker receives. Every heartbeat and every terminal report presents the token, and a mismatch is rejected. The reconnecting worker's heartbeats carry the old generation and are refused, so it drops them. Worker A's late success is refused and worker B's result is recorded. The zombie is refused on two counts. And the expiration monitor's transition carries the worker and generation it observed, so the store refuses the stale write.

Generations are 64-bit integers. At a million assignments per second they wrap in about 585,000 years.

A fencing token fences the control plane, not the world. During a partition, worker A may keep running for up to the expiry window while worker B starts the same action. If your task writes to an external system, the write has to be idempotent or carry the action name and generation so the external side can discard the stale one. We document this as a property of the system rather than a bug, because the alternative is a global lock on every side effect.

Heartbeats at scale

A single worker can hold tens of thousands of leases. If each heartbeat were one durable write, a hundred workers renewing five thousand leases each wouldn't fit in any heartbeat interval. Leases expire when they stop heartbeating, and the interval is set at the system level, with tolerance for several missed beats and jitter before expiry.

So heartbeats are batched, a thousand tokens to a request with several in flight, and each request carries the worker's free capacity, which is how the scheduler learns where there's room. On the engine side each batch runs in two phases. The engine handles this cleverly with two phased checks, enabling five thousand extensions to be completed in a handful of milliseconds.

Leases are resolved resiliently especially taking care for the type of workload Union is used for a lot — long running tasks. A rejection means something else entirely: the lease has already moved on to completion processing, and the worker must stop now.

The cluster worker

The other half of every lease is the cluster worker. It runs inside each data-plane cluster, scales horizontally with the work, and is the only thing in the system that touches Kubernetes. It opens an outbound stream to the scheduling engine, receives leases, and executes them through plugins: plain pods, Ray, Spark, Dask, clustered training, Union's reusable containers and warm Ray clusters. It heartbeats everything it holds in the batches described above, and it runs the cleanup phase for every attempt as its own finalize lease. Because it only dials out, a cluster on AWS, GCP, Azure, or a rack in your building looks identical to the control plane, and the control plane holds no credentials into any of them.

The worker also closes the observability loop from the infrastructure side. It watches Kubernetes through informers, so a pod eviction or an unschedulable-pod event reaches the control plane in seconds, attributed to the exact action, rather than as a log line someone correlates later.

And it fails loudly. If availability is compromised, by a partition or by watches that die silently, the worker aborts itself rather than degrading. That avoids split brain and keeps failures visible. The same rule holds in the control plane, where a misconfigured service refuses to serve.

Memory is the truth, the database is the log

Every state transition in the engine is a synchronous, write-through operation: clone the lease, mutate in memory, persist at quorum, and on failure restore the clone. Enqueue, assignment, terminal report, abort, retry, delete, all of them. Only heartbeat extension is best-effort. After startup, nothing reads the database again.

The constraint buys two things. The event-driven, sub-millisecond scheduling pass and the lock-free queue reads the scheduler depends on are only possible because scheduling never waits on I/O. And it removes any need for compare-and-swap in the database. A conditional write in a distributed store is a consensus round, several times the latency of a plain write. We use one exactly once, at startup, when a shard claims leadership of its tenants with a conditional insert and a short TTL that it renews. That's the split-brain guard. It runs off the hot path, it costs tens of milliseconds once, and it's what makes "one owner per shard" true rather than hopeful. The shape is the one multi-Paxos and Raft formalize, a single elected leader making many decisions, each replicated before it takes effect. The store provides the quorum, and the engine's only job is to never act before that write returns.

The same principle runs outward. A worker learns about a new lease by having it pushed down its stream. A parent learns that a child finished the same way. The console learns that an action changed state the same way. Nothing in the engine polls, and nothing waits for a timer to notice a change. A state change streams to whoever is waiting on it the instant it happens, which is why a step in an agent loop can be scheduled in milliseconds rather than on the next poll.

This is also the "we wanted to sleep" promise in practice. Leadership is claimed, not assigned by a person. A shard that dies is re-elected and reloads. Workers reconnect on their own, and the sweep re-drives anything that was mid-flight. Nobody is paged to reconcile state, because there's no state a person could reconcile better than the sweep does.

Recovery on restart is a parallel scan of the shard's partitions, rebuilding the in-memory indexes as rows load, with a budget of half a heartbeat interval. In practice it's a few seconds for tens of thousands of leases and ten to fifteen for a million. Workers aren't reconciled at all. They reconnect, and their next heartbeat either validates or is refused. The system is event-driven with anti-entropy built in, so anything that was mid-transition when the process died is re-driven rather than dropped.

Three timeouts, one enforcer

A task can declare three bounds, and each has a different owner and a different clock.

  • `deadline` — enforced by the engine, before dispatch — anchored to the first enqueue time, immutable across attempts
  • `max_queued_time` — enforced by the engine while queued, then the worker until the pod is running — anchored to this attempt's enqueue time, reset on retry
  • `max_runtime` — enforced by the worker — anchored to this attempt's start time

The scheduling engine enforces all three.

Where it stands

  • 51,889 actions/s dispatched by one shard, 50k actions, p99 97 ms end to end
  • 10–20 ms to durably extend 5,000 leases in one heartbeat batch
  • 3.6 ms p50 dispatch latency at 50k actions with round-robin fairness, against 77.8 ms p99 without
  • 10–14 s to reload a million leases on restart

Under significant load the system stays reliable and fast, and the fairness figure is the one you can feel. A small interactive run sitting next to a 50,000-action fan-out still dispatches in a few milliseconds, at the same overall throughput.

What customers say

"I waited too long to switch to Flyte 2 — it's just so much better. We can scale to 200,000–300,000 pods with the escalation logic baked right in, and the out-of-memory and scheduling headaches I used to fight are simply gone."
— Jay Ganbat, Principal Bioinformatics Engineer, Prima Mente

Next up

Everything above is about whether an action runs and whether its result survives. The next post is about where and when: queues with quotas, resource-aware admission across clusters, gang scheduling, and backfill on predicted runtimes. If you'd rather see the whole path from the user's side first, the life of a run page walks it end to end. And if you'd rather run it than read about it: `pip install flyte` and `flyte start devbox`.

The Inside Union series

  1. The Durable AI Runtime for Flyte 2.0
  2. The Replay Log That Makes Flyte 2.0 Durable
  3. Leases, the Scheduling Engine Behind Flyte 2.0 (this post)
  4. Resource-Aware Scheduling for Flyte 2.0
  5. Observability at a Million Actions
Try the devbox

A free, local sandbox to explore the Union.ai platform.

Chat with an engineer
No items found.