From Slurm to Flyte

A guide to moving batch and ML workloads off a Slurm cluster onto Flyte 2.

The compute code mostly survives the move. The training loop, the preprocessing script, and the eval harness come over unchanged. What changes is everything around them: where the environment comes from, how a job gets its inputs, how one job triggers the next, and how you ask for capacity.

Three shifts account for most of the work:

• module load and conda environments become container images. You declare the environment once, in Python, instead of relying on what happens to be installed on the login node.
• The shared filesystem becomes explicit inputs and outputs. A script that assumes /scratch exists on every node needs to take its data as arguments instead.
• sbatch chains become function calls. Dependency flags, sentinel files, and the cron entries that glue them together become ordinary Python.

The rest of this guide maps Slurm constructs onto their Flyte equivalents.

The mapping at a glance

The job script becomes a task

A representative Slurm job:

#!/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 in Flyte:

from datetime import timedelta

import flyte

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:
    ...

The #SBATCH block splits in two. Anything that describes the environment (image, resources) belongs on the flyte.TaskEnvironment, which is shared by every task declared against it. Anything that describes this job (retries, timeout, caching) belongs on the @env.task decorator. Both can be overridden per invocation with task.override(...), which has no Slurm equivalent: the same task can run with 1 GPU in one call and 8 in the next.

The function signature also does more work than a job script’s argv. lr: float and -> flyte.io.File are the task’s interface, and Flyte records the value of every input and output of every run, so a checkpoint six months old can be traced back to the arguments that produced it.

Docs: TaskEnvironment · Resources · Overrides

module load becomes an image

This is usually the slowest part of a Slurm migration, and the one worth doing carefully. On a Slurm cluster the environment is ambient: modules, a conda env on NFS, whatever the sysadmin installed. In Flyte the environment is part of the task definition.

You do not need to write a Dockerfile. flyte.Image builds one for you from Python:

image = (
    flyte.Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("git")
    .with_pip_packages("torch", "transformers", "datasets")
)

If you already keep your dependencies in a pyproject.toml, point at it directly with .with_uv_project("pyproject.toml"). If your organization already publishes a blessed CUDA image, pass it by reference instead and skip the build entirely:

env = flyte.TaskEnvironment(
    name="training",
    image="registry.example.com/ml-platform/cuda-torch:2026.04.01",
)

Two practical notes for people coming from modules:

• Different tasks can use different images. There is no single cluster-wide Python environment to keep everyone happy. A preprocessing task on a slim CPU image and a training task on a CUDA image are part of the same run.
• Images are content-hashed. Rerunning with an unchanged spec reuses the previous build instead of rebuilding.

Where that build runs depends on your deployment.

Docs: Container images

Job arrays become fan-out

A Slurm job array indexes into work with $SLURM_ARRAY_TASK_ID and leaves collection of the results to you, usually as files in a shared directory plus a script that reads them back.

#SBATCH --array=0-999%50
python process.py --shard $SLURM_ARRAY_TASK_ID

In Flyte the fan-out is a call, and the results come back as return values:

@env.task
async def process(shard: int) -> int: ...


@env.task
async def main(n_shards: int = 1000) -> int:
    counts = flyte.map(process, range(n_shards), concurrency=50)
    return sum(c for c in counts if not isinstance(c, Exception))

concurrency=50 is the equivalent of the %50 throttle: at most 50 shards run at once, and the rest wait. flyte.map yields results in input order and returns an exception object in place of a result for shards that failed, so a partial failure does not cost you the whole array.

When the items are not uniform, or you want to fan out across different tasks, use asyncio.gather instead:

results = await asyncio.gather(*(process(s) for s in shards), return_exceptions=True)

Docs: Mapping over inputs · Fanout · Controlling parallelism

Job dependencies become ordinary Python

--dependency=afterok:$JOBID handles a linear chain. Anything past that (a fan-out that joins, a branch on the result of an earlier step, a retry of just one stage) tends to become a driver bash script, some sentinel files, and a wiki page describing the arrangement.

In Flyte, a pipeline is a task that calls other tasks. There is no workflow DSL and no graph to compile:

@env.task
async def main(ds: str) -> Report:
    raw = await ingest(ds)
    clean = await filter_rows(raw)
    shards = await asyncio.gather(*(tokenize(clean, i) for i in range(8)))
    model = await train(shards)
    return await evaluate(model)

Because the driver runs at execution time as normal Python, control flow is normal Python too. Branching is if. Early exit is return. Failure handling is try/except/finally, and specific failure modes are catchable by type:

import flyte.errors


@env.task
async def main(ds: str) -> int:
    try:
        return await transform(ds)
    except flyte.errors.OOMError:
        return await transform.override(resources=flyte.Resources(memory="64Gi"))(ds)

Scheduled submission moves from a cron entry on the login node to a trigger on the task itself:

@env.task(triggers=flyte.Trigger("nightly", flyte.Cron("0 2 * * *")))
async def nightly_eval() -> Report: ...

Docs: Triggers · Error handling

--requeue becomes retries, spot handling, and checkpoints

--requeue restarts the script from the top and leaves the rest to you. Flyte splits the problem into pieces that can be configured separately.

Retries are declarative, and count only against failures your code is responsible for:

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

Spot capacity is a flag. interruptible=True schedules the task on spot or preemptible instances. Preemptions are recorded as system failures rather than task failures, so they do not consume the retry budget, and the last attempt falls back to on-demand so a task cannot loop forever on reclaimed capacity.

Checkpoints make the retry cheap. flyte.ctx().checkpoint writes to object storage rather than a shared filesystem, so the next attempt resumes on whatever node it lands on:

@env.task(retries=5)
async def train(steps: int) -> flyte.io.File:
    ckpt = flyte.ctx().checkpoint
    start = 0
    if (prev := await ckpt.load()) is not None:
        start = load_state(prev)

    for step in range(start, steps):
        ...
        if step % 100 == 0:
            await ckpt.save(state_path)

Task-level caching covers the other half of the problem. A cached task with unchanged inputs is skipped on re-execution, so rerunning a twelve-hour pipeline after fixing step nine starts at step nine instead of step one.

Docs: Retries and timeouts · Interruptible tasks · Intra-task checkpoints

Multi-node jobs become clustered tasks

--nodes=4 --ntasks-per-node=8 maps onto a ClusteredTaskEnvironment, which launches all replicas together as a single Kubernetes JobSet with torchrun handling rendezvous:

import flyte
from flyte.clustered import ClusteredTaskEnvironment, ClusterFailurePolicy, TorchRun

env = ClusteredTaskEnvironment(
    name="pretrain",
    image=image,
    resources=flyte.Resources(cpu=16, memory="64Gi", gpu="H100:8", shm="auto"),
    replicas=4,           # nodes
    nproc_per_node=8,     # processes per node, so world size is 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")
    ...

Training code that already runs under srun with torchrun needs no changes: RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT are populated in each worker as usual. The same values are available from flyte.ctx() (rank, local_rank, node_rank, nnodes, world_size, master_addr) if you would rather read them from Python. Only rank 0 uploads the task’s outputs.

ClusterFailurePolicy distinguishes between two things Slurm treats alike. restart_on_host_maintenance=True restarts the job when the underlying node is preempted or drained for maintenance, without spending the max_restarts budget you set aside for actual crashes.

Docs: Clustered task environments

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


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

The shared filesystem becomes explicit data

This mapping has no one-line translation. Doing it properly is what buys you lineage, caching, and reproducible reruns; working around it costs you all three.

On Slurm, your home directory and /scratch are visible from the login node and every compute node, so a job reads and writes paths and the filesystem does the rest. There is no equivalent guarantee in Flyte. Data moves because a task takes it as an argument or returns it.

For most cases, the change is mechanical. Values that fit in a return type (numbers, strings, dataclasses, Pydantic models) travel as return values. Anything larger travels as flyte.io.File or flyte.io.Dir, which are typed references to object storage:

from flyte.io import Dir, File


@env.task
async def tokenize(raw: Dir) -> File:
    out = File.new_remote()
    async with out.open("wb") as f:
        ...
    return out


@env.task
async def train(tokens: File) -> File:
    local = await tokens.download()
    ...

A File passes between tasks the way an int does. The upload on write and the download on read are handled for you, and the object supports streaming, so reading a range out of a 500 GB file does not require pulling the whole thing to local disk first.

Workloads that genuinely need a filesystem view have an escape hatch. If you already run a parallel filesystem such as FSx for Lustre or a shared NFS export, mount it into tasks with a pod template. Pod templates give you the full Kubernetes pod spec: volumes, node selectors, tolerations, service accounts, and sidecars.

Docs: Files and directories · Pod templates

Watching and debugging jobs

squeue and sacct map onto the CLI and the UI:

flyte get run                # recent runs and their phases
flyte get run <run-name>     # actions within a run
flyte get logs <run-name>    # streaming logs

Runs are also visible in the UI with per-task logs, GPU and memory utilization, and inputs and outputs for every action.

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

Staging the migration

Migrations tend to go badly when the first workload moved is the most expensive one. A sequence that works:

1. Start with pipeline-shaped work. Data processing, evaluation, hyperparameter sweeps, batch inference. These gain the most from typed inputs and outputs, caching, and retries, and they are the workloads where the sbatch glue was worst. Nothing expensive is at risk while you settle the image and data-access questions.
2. Move single-node training next. By this point the image is validated and the data is coming in as arguments, so what you pick up is reproducible environments, spot with automatic fallback, checkpoint recovery, and run metadata.
3. Move multi-node training last. It is the most performance-sensitive workload in the stack, and by the time you get to it, everything underneath it has already been exercised.

Existing binaries do not have to be rewritten to come along. A container task runs an arbitrary image with typed inputs and outputs, whatever language the tool is written in, which covers the bioinformatics binaries and vendor CLIs that tend to be wrapped in srun today.

For the inner loop, flyte run --local executes the same code in your local Python process with no cluster involved, which is the closest thing to iterating on the login node before submitting.

Docs: Running locally · Container tasks

Scheduler features without an equivalent

Three things Slurm’s scheduler does are not available, and if your workload depends on them, plan around them explicitly:

• Gang admission. A clustered task’s replicas are launched together and torchrun waits for the full group before training starts, but there is no scheduler-level all-or-nothing admission.
• Topology-aware placement. There is no way to request workers on the same rack or switch.
• Preemption. Nothing evicts running low-priority work to make room for high-priority work the way a Slurm QOS can. Priority affects the order work starts in, not what happens to work that has already started.

These are gaps in the Kubernetes batch ecosystem rather than anything specific to Flyte.