Union.ai
Flyte
Compare

A Modern Alternative to Slurm for ML Workloads

Samhita Alla

Samhita Alla

What one of computing's most successful schedulers got right, why ML teams are leaving it anyway, and what your workflows look like on the other side.

Slurm has been scheduling the world's largest computers for more than twenty years. It came out of Lawrence Livermore National Laboratory in 2002, and today it schedules over 65% of the TOP500 supercomputers, ships as the default scheduler on nearly every GPU neocloud, and is the first tool most researchers touch when they get cluster access in grad school. If your team trains models on more than a handful of GPUs, someone on it can write an `sbatch` script from muscle memory.

So when teams tell us they're migrating off Slurm, both halves of that sentence deserequeuerve to be taken seriously. They chose Slurm for real reasons and they're leaving for real reasons. This post is for two kinds of readers. The first is running ML workloads on a Slurm cluster and wondering whether the growing pile of bash glue around it is normal (it is) and permanent (it isn't). The second has already decided to move to Kubernetes and is discovering that vanilla Kubernetes answers a different question than the one Slurm answered. What follows is the case as we'd make it to a colleague: what Slurm does well, why ML outgrew it, why raw Kubernetes disappoints, and what the same work looks like on Union.

What Slurm got right

Any honest story about migrating away from Slurm has to start with why people adopted it in the first place.

Take scheduling. When Slurm grants a 64-GPU job its allocation, all 64 GPUs start together. Distributed training can't tolerate partial starts; you can't begin with 63 GPUs and hope the last one shows up. Slurm treats the whole job, not the individual process, as the unit of scheduling, and that one design decision explains much of its staying power in large-scale training. It also knows your network. Tightly coupled jobs land on nodes that share a switch or an InfiniBand island, which is where a meaningful fraction of training throughput comes from at scale.

Slurm also solved the sociology of shared machines. Decades of academic clusters produced hierarchical fair-share, QOS tiers, per-account limits and accounting good enough to settle arguments about who used what. The control plane is light, scheduling throughput is high, and an allocation once granted, is yours. No evictions and no noisy neighbors on your nodes.

And researchers genuinely like using it. A job is a bash script with some `#SBATCH` pragmas at the top. `squeue` tells you where things stand. Your home directory sits on a shared filesystem, so the code and data you see on the login node are exactly what your job sees on every compute node. There is very little to learn and almost nothing to package.

That last point matters more than it looks. Most of the friction in leaving Slurm comes from giving up that shared-filesystem, ssh-and-edit way of working, not from the scheduler itself. Keep it in mind; we'll come back to it.

Why teams leave anyway

The short version is that Slurm schedules jobs and ML work stopped looking like jobs.

A modern training effort is a pipeline. Data ingestion, filtering, tokenization, training, evaluation, quantization, deployment and monitoring all have different resource shapes and are wired together with dependencies and retries, and Slurm sees none of that structure. Job dependencies exist (`--dependency=afterok:...`), but anything beyond a linear chain turns into bash scripts, sentinel files, cron entries and a wiki page explaining the whole arrangement to new hires. Reinforcement learning makes it worse, since a GRPO or PPO loop interleaves rollouts, reward scoring and gradient updates in one tight cycle. That's a workflow, not a job.

The cluster stopped being fixed too. Slurm assumes a static pool of near-identical, always-on nodes, which described a national lab machine perfectly and describes almost no ML fleet today. GPU capacity now arrives as a reservation in one cloud, spot capacity in another, and a neocloud contract in a third region. Elasticity, heterogeneous instance types and scale-to-zero are foreign to a scheduler built for machines that never turn off.

Then there's everything that isn't a batch job at all. The model you trained has to be served, and inference endpoints, eval dashboards, labeling tools and agent runtimes are long-running services, a concept Slurm simply doesn't have. Teams end up running Slurm for training and a separate Kubernetes cluster for serving, with two operational stacks, two access-control systems, and a wall between them exactly where the train-eval-deploy loop needs to flow.

Software environments round out the list. `module load cuda/12.1` plus a conda environment on NFS works until two users need conflicting versions or until you need to reproduce last quarter's result. Containers solved this everywhere else in infrastructure, and while Slurm can run them through extensions like Pyxis and enroot, the ecosystem's center of gravity, from PyTorch base images to inference servers, assumes a container-native platform.

Why Kubernetes isn’t enough

None of this is a failure of Slurm's engineering as Slurm was built for a machine and ML teams now operate a supply chain.

So why not just switch to Kubernetes? That's where things get tricky. Vanilla Kubernetes doesn't come with many of the scheduling guarantees and HPC features that Slurm users are used to.

Kubernetes was designed in 2014 to keep stateless services alive and it is excellent at that. Batch scheduling was bolted on later. Out of the box there is no job queue, no fair-share, no priority-based ordering worth the name and no gang scheduling. A 64-pod training job can get 60 pods placed and sit deadlocked while the remaining four wait for capacity that the first 60 are now blocking. The default scheduler places pods one at a time, optimizing for availability rather than batch throughput. And the interface it offers your researchers is YAML. A five-line `sbatch` script becomes a couple hundred lines of manifests, and the person who used to type `squeue` now needs to understand pods, nodeSelectors, tolerations and why their job is `Pending` with no further comment. Nebius wrote a good deeper treatment of this design gap.

The ecosystem knows all this, which is why the gap is slowly closing. Kueue adds quotas, queueing and all-or-nothing admission. Volcano and YuniKorn offer batch schedulers. JobSet gives multi-node jobs a first-class API and Dynamic Resource Allocation finally gives the scheduler structured knowledge of GPUs and other devices, the kind of hardware awareness Slurm has had for decades. The convergence is running in both directions, too. CoreWeave ships SUNK, which is Slurm running on Kubernetes. SchedMD, the company behind Slurm, maintains Slinky, its own Slurm-on-Kubernetes project, and was acquired by NVIDIA last year. The industry has effectively settled the substrate question in favor of Kubernetes. What nobody has settled is the interface question, because researchers should not have to become Kubernetes operators to train a model.

That gap between substrate and interface is where Union sits.

What the same work looks like on Union

Union is an AI runtime built on Kubernetes and powered by Flyte. It runs in your own cloud account or on clusters you already operate, including neoclouds like CoreWeave, Crusoe and Nebius, and any conformant Kubernetes cluster. Union keeps Kubernetes as the substrate and returns the things you actually liked about Slurm: a simple mental model, queues, batch semantics for multi-node jobs, while adding what Slurm was never going to grow: pipelines, elasticity, serving, reproducibility. Here's the mapping concept by concept:

Slurm Union
`sbatch train.sh` `flyte run train.py main`
`#SBATCH --gres=gpu:a100:8` `flyte.Resources(gpu="A100:8")`
`#SBATCH --cpus-per-task=16 --mem=64G` `flyte.Resources(cpu=16, memory="64Gi")`
`#SBATCH --array=0-999` `flyte.map(step, range(1000))`
`#SBATCH --nodes=4 --ntasks-per-node=8` `ClusteredTaskEnvironment(replicas=4, nproc_per_node=8)`
`#SBATCH --partition=gpu --qos=high` `queue="gpu-high"`
`#SBATCH --requeue` `retries=3`, plus `interruptible=True` for spot
`#SBATCH --time=04:00:00` `timeout=timedelta(hours=4)`
`#SBATCH --begin=...` / cron `triggers=flyte.Trigger(...)` with `flyte.Cron`
`module load cuda && source venv/bin/activate` `flyte.Image.from_debian_base().with_pip_packages(...)`
`$SLURM_PROCID`, `$SLURM_NNODES` `flyte.ctx()` rank, world size (torchrun-compatible)
`squeue`, `sacct` `flyte get run`, `flyte get queue --watch`, the UI
`srun --pty bash`, `ssh node042` `--debug`: VS Code in the pod or SSH into the task

The job script becomes a function

A representative Slurm job:

Copied to clipboard!
#!/bin/bash
#SBATCH --job-name=train
#SBATCH --partition=gpu
#SBATCH --gres=gpu:a100:8
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --time=04:00:00
#SBATCH --requeue
module load cuda/12.1
source ~/venvs/train/bin/activate
srun python train.py --lr 3e-4

The same job on Union:

Copied to clipboard!
import flyte
from datetime import timedelta

env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages("torch"),
    resources=flyte.Resources(cpu=16, memory="64Gi", gpu="A100:8"),
)

@env.task(retries=3, timeout=timedelta(hours=4))
async def train(lr: float = 3e-4) -> flyte.io.File:
    ...

What changed is more than syntax. The task has typed inputs and outputs, so composing tasks is just calling functions, and every input and output of every run is recorded. There is no workflow DSL to learn. A pipeline is a task that calls other tasks, and control flow is ordinary Python. A hyperparameter sweep that would be a job array plus a results-collection script on Slurm becomes:

Copied to clipboard!
@env.task
async def sweep(configs: list[TrainConfig]) -> TrainConfig:
    results = await asyncio.gather(*(train(c) for c in configs))
    return max(zip(results, configs))[1]

For bounded fan-out at larger scale, `flyte.map(train, configs, concurrency=200)` runs a worker pool and streams results back in order. Runs scale to tens of thousands of tasks; a single run supports up to 50,000 actions.

The `module load` line is a good place to start because it's where one of the biggest misconceptions about moving from Slurm to Kubernetes shows up: that you'll suddenly have to start writing Dockerfiles. In practice, you don't. `flyte.Image` lets you define your environment in Python instead. Start with a maintained base image, add your Python or system dependencies, or even an entire `uv` project, and Union builds and caches the image remotely the first time you run it. You don't need Docker installed locally or have to think about container registries. And because every environment is content-hashed, it's always clear exactly which dependency versions produced a given checkpoint, instead of having to piece it together later.

Multi-node training without the YAML

Multi-node training is where Slurm users are most skeptical of Kubernetes-based platforms, understandably, so here is the complete setup:

Copied to clipboard!
import flyte
from flyte.clustered import ClusteredTaskEnvironment, TorchRun, ClusterFailurePolicy

env = ClusteredTaskEnvironment(
    name="pretrain",
    image=image,
    resources=flyte.Resources(cpu=16, memory="64Gi", gpu="H100:8", shm="auto"),
    replicas=4,
    nproc_per_node=8, # world size = 32
    runtime=TorchRun(rdzv_backend="c10d"),
    failure_policy=ClusterFailurePolicy(max_restarts=2, restart_on_host_maintenance=True),
)

@env.task
async def pretrain(steps: int) -> flyte.io.File:
    import torch.distributed as dist
    dist.init_process_group(backend="nccl")
    ...

A clustered task launches all replicas as a single Kubernetes JobSet, with headless networking between pods and torchrun handling rendezvous, so your training code sees exactly the `RANK`, `WORLD_SIZE` and `MASTER_ADDR` it expects. Rank, node rank and world size are also available directly from `flyte.ctx()` and only rank 0 uploads outputs.

The failure policy is where things start to go beyond `--requeue`. With `restart_on_host_maintenance=True`, your job automatically restarts if the underlying node is preempted or taken down for maintenance without eating into the restart budget you've set aside for actual failures. That's an important distinction: the cloud taking away your node isn't the same as your code crashing. On shared cloud GPU fleets, that difference comes up all the time, and it's something Slurm was never really designed to account for.

The same philosophy extends to distributed compute. Instead of provisioning and managing long-lived clusters, you can spin up Ray, Spark or Dask as ephemeral, per-task clusters that are created when the task starts and torn down automatically when it finishes.

Two caveats are worth keeping in mind. First, clustered tasks are still new and primarily designed for `torchrun` workloads. If you're running MPI applications rather than PyTorch, Ray or Spark, Union isn't targeting that use case just yet.

Second, some of the scheduling guarantees you'd expect from Slurm are still a work in progress. JobSet launches the workers together, and `torchrun` waits until the full group has assembled before training begins, but scheduler-level gang scheduling and topology-aware placement (for example, keeping workers on the same rack) are on the roadmap.

Partitions and QOS become queues

Union queues are named scheduling lanes that map to your compute and they answer the question every shared cluster fights about, which is how forty people share a fixed number of GPUs without either chaos or idle hardware.

Source: https://www.union.ai/blog-post/introducing-queues-and-cluster-controls-durable-workloads-under-contention

A queue routes work to a cluster pool and puts limits on it. You can cap how many runs and how many individual tasks execute at once, cap the backlog so that submissions past a certain depth get back-pressure instead of queueing forever, set priority levels that control ordering, and pick a fairness policy that decides how work from different projects interleaves when they share a lane. Targeting a queue is one parameter, settable at the environment, task, invocation or trigger level:

Copied to clipboard!
@env.task(queue="research-h100")
async def evaluate(model: flyte.io.File) -> EvalReport: ...

await train.override(queue="prod-high")(cfg)   # per-invocation

`flyte get queue research-h100 --watch` streams live occupancy, which is the closest thing to `squeue` muscle-memory you'll find here. Queues are in beta and the current priority levels order work rather than preempt it. Preemption, where high-priority work evicts running low-priority work the way Slurm QOS can, is in development.

How INESC TEC keeps shared GPUs fair and busy

INESC TEC is a research institute in Porto with more than 30 years of work spanning AI, robotics, energy and bioengineering. Slurm is already part of the HPC landscape there. Some research groups run their own dedicated Slurm clusters, deployments that predate the institute's platform work. At the institute level, however, the infrastructure team deployed Union on INESC TEC's local HPC and GPU infrastructure. This gives research groups across the institute access to a larger, shared pool of GPUs and a consistent way to define and run their AI workflows.

The challenge was a familiar one: many research groups sharing a finite pool of GPUs, with no reliable way to predict who would need what or when. The team needed to allocate capacity fairly across groups, let individual projects draw from their group's quota, and make sure idle GPUs didn't sit unused. Slurm's fair-share scheduler can approximate this, but it's relatively coarse-grained, and tuning it well is practically a discipline of its own.

At the time, Union didn't yet have queues, so the team built a two-layer quota model on top of Kueue, with everything declared and versioned alongside the rest of their infrastructure. Researchers don't have to think about any of this. They write Python, submit workflows, and the platform handles quotas and resource allocation behind the scenes.

The platform also enforces quotas at the project level. If a request exceeds the available quota, it fails immediately with a clear error rather than sitting in `Pending` indefinitely.

Looking ahead, the team is exploring how to extend this model to larger HPC and supercomputing environments: running Kubernetes-based execution environments on top of those resources and exposing them through Union. Researchers would continue working at the workflow level, while the infrastructure layer determines where the underlying compute resources actually reside.

Requeue becomes real fault tolerance

Slurm's `--requeue` simply restarts your script from the beginning and leaves the rest up to you. Union breaks the problem into separate pieces with each one handled by the layer that's best suited for it.

Retries with backoff are declarative, either a simple `retries=5` or a full `RetryStrategy` with exponential backoff. Spot capacity is a boolean. `interruptible=True` schedules the task on spot or preemptible instances, preemptions are tracked as system failures that don't consume your retry budget, and the final attempt automatically falls back to on-demand so a job can't preemption-loop forever. Checkpoints save to object storage through a small API, so a retried attempt resumes from the last checkpoint on any node with no shared filesystem required; there are examples for HuggingFace Trainer and PyTorch Lightning resume-from-last.

One layer finer, `@flyte.trace` makes individual function calls inside a task durable. A traced call whose result was already recorded is skipped on retry and returns the stored result, so in a pipeline that interleaves cheap steps with expensive ones, recovery replays the cheap parts and skips the expensive ones. And task-level caching keyed on code and inputs means re-running a twelve-hour pipeline after fixing step nine starts at step nine.

The login node becomes a warm pool

Kubernetes' per-task cold start, scheduling a pod, pulling an image, importing torch, is real overhead and it's the real reason Slurm feels fast: your allocation is already running. Union's answer is reusable containers:

Copied to clipboard!
env = flyte.TaskEnvironment(
    name="rollouts",
    resources=flyte.Resources(gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=(2, 10), concurrency=4, idle_ttl=300),
    image=flyte.Image.from_debian_base().with_pip_packages("unionai-reuse", "vllm"),
)

A reuse policy keeps a pool of warm containers that persist across task invocations, autoscale between bounds, and hold in-memory state, so a model loaded once serves thousands of subsequent tasks with scheduling measured in milliseconds rather than pod-startup seconds. This is the pattern behind the RL setups mentioned earlier, where a warm GPU pool runs rollouts continuously while ephemeral tasks handle training steps and reward scoring. There is one trade-off to understand. Each replica is a long-lived Python process, so in-memory state survives across invocations on that replica (though not across restarts). That persistence is the point but it means global state deserves the same care you'd give it in any long-running server.

The shared filesystem becomes explicit data

This is where the mental model really changes. There's no shared `/home` directory that every job can rely on. That's often the first assumption people bring from Slurm and it's one that doesn't hold up on Kubernetes. In Union, data movement is explicit. It's part of the workflow, not something that happens in the background.

Tasks pass `flyte.io.File` and `flyte.io.Dir` objects, which are typed references to object storage. They stream, so you can read a range of a 500 GB file without downloading it. They upload and download without any S3 boilerplate in your code. And because every task's inputs and outputs are recorded, you get lineage for free; six months later you can trace exactly which data produced which checkpoint, which no NFS home directory will ever tell you.

For workloads that genuinely want a filesystem view, Union Volumes provide a durable, FUSE-mounted filesystem backed by object storage, with versioned commits and copy-on-write forks for parallel branches of work (one writer at a time while mounted). Teams keeping a parallel filesystem like FSx for Lustre can mount it into tasks through standard CSI drivers via pod templates, which are also the general escape hatch to full Kubernetes pod specs, node selectors, tolerations and sidecars included, when you need them.

SSH to the node becomes a debug session

On Slurm, when a job misbehaves, you get on the node. `srun --pty bash` or a straight `ssh` to the machine, then `nvidia-smi`, `htop`, a look at the filesystem, and you know what's actually happening. Losing that is many researchers' single biggest fear about a managed platform, so it's worth saying plainly: you don't lose it.

Union gives you three ways into a running task. From the UI, one click opens a live debugging session on any action. Launching a run with `debug=True` (or `--debug` from the CLI) starts a browser-based VS Code server inside the task pod, where you can set breakpoints and step through your code on the same hardware, against the same data, images and dependencies the run uses. And when you want your own terminal, you can SSH directly into the task pod (currently in beta) or attach VS Code Remote-SSH from your laptop. If anything, the scope improves on Slurm: you land inside the failing task's exact container rather than on a shared node whose modules and paths may not match the job, and there's no login node in the middle that everyone else is competing for.

What Slurm can’t give you

Everything so far is about parity: familiar Slurm workflows but on an elastic Kubernetes backend. That's important because it makes the transition easier. But it's not why teams stick around.

The real value comes from things Slurm was never designed to do. Training and serving live on the same platform. The checkpoint produced by a clustered training job can be deployed as a vLLM application from the same codebase, with model weights streamed directly from object storage into GPU memory instead of first downloading them to local disk. Serving apps autoscale with traffic, including all the way down to zero when they're idle. Training, evaluation and deployment become parts of the same workflow instead of separate systems owned by different teams.

The economics change too. Cluster pools can span multiple Kubernetes clusters, cloud providers and regions behind a single control plane. GPU node groups scale up when demand increases and scale back down when they're no longer needed, so you're not paying for idle capacity. Spot instances also become much more practical for training because interruptions are handled by the platform's retry and checkpointing mechanisms rather than by custom logic in every training script.

You also get observability without having to assemble it yourself. Live logs, GPU metrics, resource utilization and HTML reports all come out of the box. You can keep an eye on loss curves while training is still running, and when someone asks why the GPU bill was so high, the data is already there instead of buried across half a dozen tools.

Real-time logs in the Union UI
Monitor resource utilization in the Union UI
Render interactive HTML reports in the Union UI

The enterprise layer is someone else's job now. SSO against any OIDC or SAML IdP, role-based access control down to project scope, audit logs on every API call, and a two-plane architecture in which your code, data, images and secrets never leave your account.

The control plane sees orchestration metadata and pointers, nothing else, and the data plane only ever dials out:

  • there is no inbound port to defend
  • the Kubernetes API never faces the public internet

Secrets are write-only, so the platform can inject them into tasks but no API will ever return one.

Union is SOC 2 Type II certified, and for organizations whose policy says no third-party network may ever reach their data, a familiar posture in healthcare, finance and defense, a sovereign deployment option locks the data plane inside the corporate network entirely, beyond the reach of even Union's own employees.

And the dev loop survives the move. `flyte run --local` executes the same code in-process on a laptop, and the same task submits to the cluster from a Python script, a CI job or a Jupyter notebook. And when something misbehaves remotely, the debug session from earlier, browser VS Code or SSH straight into the task pod, is one click away.

What a migration actually involves

We won't pretend the migration is effortless. There are a few mindset shifts along the way and it's better to know about them upfront.

The first is moving from modules to images. The image builder means you don't have to write Dockerfiles, but you do have to declare your environment once instead of relying on whatever happens to be installed on the login node. For most teams, this is the slowest part of the migration. It's also the point where your environments become reproducible instead of accidental.

The second is making data access explicit. Scripts that assume `/scratch` exists everywhere need to be rewritten to accept files and directories as inputs. It's usually a straightforward change, but it's also what unlocks lineage, caching and reproducibility.

The third is treating jobs as functions instead of scripts. The `sbatch` chains, sentinel files and cron jobs that glue workflows together become ordinary Python code. In most cases, that means less orchestration code, not more. And if you have existing binaries you don't want to rewrite, you don't have to as they can still run as container tasks regardless of the language they're written in, with typed inputs and outputs handled automatically.

You also don't have to do the mechanical translation by hand. We've packaged the mapping in this post as a migration skill for AI coding agents. Point it at an `sbatch` script and it'll generate the Flyte 2 equivalent while flagging the parts only you can decide, such as queue names, GPU types or cluster-specific settings. Treat the output as a first draft to review rather than a finished migration, but it can eliminate most of the repetitive work.

Where to start

The migrations we've seen go most smoothly when teams start with workloads where Union is clearly a better fit than Slurm: pipeline-shaped work that benefits from composition, lineage, and failure recovery (data processing, evaluation, hyperparameter sweeps, batch inference, RL rollouts and reward scoring and offline inference). These workloads immediately gain from typed inputs and outputs, caching, retries, checkpointing and queueing, without asking teams to bet their most expensive training jobs on a new platform.

From there, most teams move single-node training onto Union to pick up the same operational benefits: reproducible environments, spot/on-demand fallback, checkpoint recovery, rich run metadata, and a single control plane for training, evaluation and deployment.

Multi-node training usually comes last because it's the most performance-sensitive workload in the stack. By that point, teams have already validated their images, data access, observability and production workflows, making the final step far more predictable.

The gaps we’re closing

A fair gap analysis comes down to scheduler primitives, not scale. There are three things Slurm's scheduler does that vanilla Kubernetes historically hasn't: all-or-nothing admission of multi-node jobs, topology-aware placement against the network fabric, and preemptive fair-share under long-tuned accounting policies. These are Kubernetes ecosystem gaps rather than anything specific to Union, and the ecosystem is moving on all three. JobSet gives multi-node jobs a first-class grouped API. Dynamic Resource Allocation, GA since v1.34, gives the scheduler structured knowledge of GPUs and other devices. And gang scheduling is now in the upstream kube-scheduler itself: workload-aware scheduling shipped as alpha in v1.35 and grew group-level preemption in v1.36, so all-or-nothing placement and group preemption are on their way to being things Kubernetes just does.

Our platform engineering is pointed at exactly this list. Gang admission and topology awareness for clustered tasks, preemption for queues, and first-class DRA support in the SDK are all in flight. If your training depends on hand-tuned fabric placement or preemptive fair-share right now, talk to us about where these sit on the roadmap.

Deciding

If your Slurm cluster runs jobs fine but your team spends its time between jobs, moving data, wiring dependencies, rebuilding environments, standing up serving infrastructure next door, then the scheduler was never the bottleneck and a better-scheduled version of the same architecture won't fix it. If you're mid-migration to Kubernetes, the substrate question is settled and the open question is whether your researchers get handed kubectl or something better. Both paths end at the same test: take one real workload, a sweep, an eval pipeline, a fine-tune on spot instances with checkpoint resume, run it through the migration skill, and run the result end to end on a Flyte devbox on your machine for free. You'll learn more from that week than from any comparison post, including this one.

Union runs in your cloud account or on your own clusters. If you're planning a move off Slurm or partway through one, we'd like to talk: https://www.union.ai/get-started.

Try the devbox

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

Chat with an engineer
No items found.