Orchestrate PyTorch, Ray and pure Python to train models, run RL, make agents durable and serve inference in your own cloud.
Find us in the Community Expo · Concourse Level
Self-serve signup at union.ai/signup
Union adds durability and resource aware scheduling to your existing PyTorch.
On the programWhere's My Training Job?
· Demo Theater, Tue Oct 20, 2:15 PM
Out-of-memory kills and preemptions show up as Python exceptions such as OOMError. Your program can catch one and retry with more memory. You set retries, checkpoints and spot recovery on the task.
On the programMaking Enterprise Agentic Inference Production-Ready
Branch, loop and fan out based on what the model just returned. Each awaited call is an action with its own retry, cache and record, so a long agent or RL run picks up where it stopped.
On the programFaster LLM Serving Startup with fastsafetensors
Reusable Containers keep replicas live across actions, so model weights load once. DynamicBatcher keeps the GPU full. Serving runs in the same program as training and shares its lineage.
This is the code behind the run log.
import asynciofrom datetime import timedeltaimport flytefrom flyte.io import File vllm_image = flyte.Image.from_debian_base().with_pip_packages("vllm", "unionai-reuse>=0.1.10")torch_image = flyte.Image.from_debian_base().with_pip_packages("torch") # rollouts: replicas stay live across actionsrollout_env = flyte.TaskEnvironment( name="rollouts", image=vllm_image, resources=flyte.Resources(gpu="L4:1", memory="32Gi"), reusable=flyte.ReusePolicy(replicas=4, concurrency=1, scaledown_ttl=timedelta(minutes=10)),)learner_env = flyte.TaskEnvironment( name="learner", image=torch_image, resources=flyte.Resources(gpu="H100:8", memory="320Gi"),)env = flyte.TaskEnvironment(name="rl-loop", depends_on=[rollout_env, learner_env]) @rollout_env.task(retries=3)async def rollout(policy: File, prompt: str) -> dict: ... @learner_env.taskasync def update(policy: File, batch: list[dict]) -> tuple[File, float]: ... @env.taskasync def rl_loop(policy: File, prompts: list[str], target: float) -> File: while True: # plain Python control flow batch = await asyncio.gather(*[rollout(policy, p) for p in prompts]) try: policy, reward = await update(policy, batch) except flyte.errors.OOMError: # out-of-memory raised as an error policy, reward = await update.override( resources=flyte.Resources(gpu="H100:8", memory="640Gi"))(policy, batch) if reward >= target: # loop length set at runtime return policyTrimmed for the page: task bodies are elided.
Rollouts run on L4s and the learner runs on 8×H100. Each environment declares its own image and hardware.
gather256 prompts become 256 actions spread over 4 live replicas. The batch size can change every loop.
Catch the error and override memory for that one call. The run keeps going.
Tasks are plain async Python functions. Build locally, then run the same program remotely and securely in your own cloud.
Qwen2.5-7B on vLLM, one L4. Reusable Containers compared with a fresh pod per call.
Read the benchmark →Coding agents writing Flyte 2 compared with Flyte 1, across 48 trials on a real cluster.
Read the paper →Remote Image Builder and Cold Boot Optimization removed the image-pull tax across 8,000-node pipelines in 4 AWS regions.
Customer result · name withheldEach framework is a plugin on a task environment. Union provisions the nodes, runs the job, and records it in the same lineage as the rest of the run.
Set Elastic(nnodes, nproc_per_node) on an environment, and the task runs as a multi-node PyTorch job. The pretraining tutorial trains a GPT-2-style model with PyTorch Lightning and FSDP on H200s. It streams MDS shards across ranks, checkpoints to S3, resumes from the last checkpoint, and pushes loss curves to a live Flyte report.
import flyte
from flyteplugins.pytorch.task import Elastic
NUM_NODES = 2
DEVICES_PER_NODE = 8
image = flyte.Image.from_debian_base().with_pip_packages("torch", "lightning", "flyteplugins-pytorch")
train_env = flyte.TaskEnvironment(
name="distributed_llm_training_h200",
image=image,
resources=flyte.Resources(
cpu=64, memory="512Gi", gpu=f"H200:{DEVICES_PER_NODE}",
shm="16Gi", # shared memory for NCCL
),
plugin_config=Elastic(nnodes=NUM_NODES, nproc_per_node=DEVICES_PER_NODE),
cache="auto",
)Attach a RayJobConfig to an environment, and Union brings up the head and worker nodes for the task. To reuse a cluster, point later tasks at a running one with address, and they connect without waiting for a new cluster to start. Autoscaling, the Ray runtime env and the cleanup TTL are all set in the same config.
import flyte
from flyteplugins.ray.task import RayJobConfig, WorkerNodeConfig
image = flyte.Image.from_debian_base().with_pip_packages("ray[default]", "flyteplugins-ray")
ray_env = flyte.TaskEnvironment(
name="ray_cluster",
image=image,
plugin_config=RayJobConfig(
worker_node_config=[WorkerNodeConfig(group_name="ray-group", replicas=2)],
address="ray://existing-cluster:10001", # reuse a running cluster
),
)Spark and Dask work the same way. Running pipelines on your own Kubernetes? Bring one to a session and we'll port it with you. Union runs in your cloud account, and your data stays there.
A 20-minute working session with a Union engineer.