Clustered task environment

from __future__ import annotations

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

# The torch wheel from PyPI bundles CUDA + NCCL, so the same image runs on CPU
# and GPU nodes. `flyte` itself provides the `clustered` runtime entrypoint that
# bootstraps torchrun inside each pod — no extra runtime library is needed.
image = flyte.Image.from_debian_base().with_pip_packages("torch", "numpy")

resources = (
    flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1")
    if USE_GPU
    else flyte.Resources(cpu=(1, 2), memory=("1Gi", "2Gi"))
)

env = ClusteredTaskEnvironment(
    name="ddp_env",
    image=image,
    resources=resources,
    replicas=REPLICAS,  # number of pods (nodes) in the JobSet
    nproc_per_node=NPROC_PER_NODE,  # processes (one per GPU) per pod
    runtime=TorchRun(rdzv_backend="static", max_restarts=0),
    failure_policy=ClusterFailurePolicy(max_restarts=1),
)

@env.task
async def train_ddp(steps: int = 50, lr: float = 0.05) -> float:
    """Run DDP training across the cluster and return the final (rank-0) loss."""
    import torch
    import torch.distributed as dist
    import torch.nn as nn
    from torch.nn.parallel import DistributedDataParallel as DDP

    ctx = flyte.ctx()

    # Bind this rank to its local GPU BEFORE init_process_group so NCCL binds
    # the right device.
    if _BACKEND == "nccl" and torch.cuda.is_available():
        torch.cuda.set_device(ctx.local_rank or 0)
        device = torch.device(f"cuda:{ctx.local_rank or 0}")
    else:
        device = torch.device("cpu")

    # torchrun has already populated RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT.
    dist.init_process_group(backend=_BACKEND)
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    print(
        f"[rank {rank}/{world_size}] device={device} "
        f"node_rank={ctx.node_rank} nnodes={ctx.nnodes} master_addr={ctx.master_addr}",
        flush=True,
    )

    # Tiny model the workers train cooperatively: learn y = x · [1,1,1,1].
    torch.manual_seed(0)
    model = nn.Linear(4, 1).to(device)
    ddp = DDP(model, device_ids=[device.index] if device.type == "cuda" else None)
    opt = torch.optim.SGD(ddp.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    # Each rank trains on its own shard of synthetic data.
    g = torch.Generator().manual_seed(rank)
    x = torch.randn(64, 4, generator=g).to(device)
    y = x.sum(dim=1, keepdim=True)

    last_loss = 0.0
    for step in range(steps):
        opt.zero_grad()
        loss = loss_fn(ddp(x), y)
        loss.backward()
        opt.step()
        last_loss = float(loss.detach())
        if rank == 0 and step % 10 == 0:
            print(f"[rank 0] step {step:3d}  loss {last_loss:.5f}", flush=True)

    dist.barrier()
    dist.destroy_process_group()
    print(f"[rank {rank}] done — final loss {last_loss:.5f}", flush=True)
    return last_loss

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(train_ddp, steps=50)
    print("Run URL:", run.url)
    run.wait()
    print("Final phase:", run.phase)

ctx = flyte.ctx()
torch.cuda.set_device(ctx.local_rank or 0)
torch.distributed.init_process_group(backend="nccl")

# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.5.18",
#    "torch",
# ]
# main = "train_fsdp"
# params = "steps=30"
# ///
"""
Fully Sharded Data Parallel (FSDP) training on a ClusteredTaskEnvironment.

FSDP shards a model's parameters across ranks, so each GPU holds only a slice —
the way you train models too big to fit on one device. The cluster shape is
identical to DDP (one JobSet, N rank-uniform pods); only the in-task wrapping
differs. This example also shows the checkpoint-to-durable-storage pattern that
FSDP requires: pod-local disk is wiped on restart, so model state must be
persisted through the task's checkpoint store.
"""

from __future__ import annotations

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

USE_GPU = True  # FSDP targets GPUs; CPU/gloo here is only a wiring smoke
REPLICAS = 2
NPROC_PER_NODE = 1

_BACKEND = "nccl" if USE_GPU else "gloo"

image = flyte.Image.from_debian_base().with_pip_packages("torch")

resources = (
    flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1")
    if USE_GPU
    else flyte.Resources(cpu=(1, 2), memory=("2Gi", "4Gi"))
)

env = ClusteredTaskEnvironment(
    name="fsdp_env",
    image=image,
    resources=resources,
    replicas=REPLICAS,
    nproc_per_node=NPROC_PER_NODE,
    runtime=TorchRun(rdzv_backend="static", max_restarts=0),
    failure_policy=ClusterFailurePolicy(max_restarts=1),
)

@env.task
async def train_fsdp(steps: int = 30, lr: float = 0.01) -> float:
    """Train a small transformer under FSDP and checkpoint the gathered state."""
    import io

    import torch
    import torch.distributed as dist
    import torch.nn as nn
    from torch.distributed.fsdp import FullStateDictConfig, StateDictType
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    ctx = flyte.ctx()

    if _BACKEND == "nccl" and torch.cuda.is_available():
        torch.cuda.set_device(ctx.local_rank or 0)
        device = torch.device(f"cuda:{ctx.local_rank or 0}")
    else:
        device = torch.device("cpu")
    dist.init_process_group(backend=_BACKEND)
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    print(f"[rank {rank}/{world_size}] device={device} nnodes={ctx.nnodes}", flush=True)

    # A small transformer encoder — each FSDP unit's parameters are sharded across ranks.
    torch.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(32, 64),
        nn.TransformerEncoderLayer(d_model=64, nhead=4, dim_feedforward=128, batch_first=True),
        nn.Linear(64, 1),
    ).to(device)
    fsdp_model = FSDP(model, device_id=device.index if device.type == "cuda" else None)
    opt = torch.optim.AdamW(fsdp_model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    # Synthetic per-rank data: learn to sum a sequence's features.
    g = torch.Generator().manual_seed(rank)
    x = torch.randn(16, 8, 32, generator=g).to(device)  # (batch, seq, features)
    y = x.mean(dim=(1, 2), keepdim=True).squeeze(-1)[:, :1]

    last_loss = 0.0
    for step in range(steps):
        opt.zero_grad()
        out = fsdp_model(x).mean(dim=1)  # pool over sequence -> (batch, 1)
        loss = loss_fn(out, y)
        loss.backward()
        opt.step()
        last_loss = float(loss.detach())
        if rank == 0 and step % 10 == 0:
            print(f"[rank 0] step {step:3d}  loss {last_loss:.5f}", flush=True)

    # --- Checkpoint: gather a FULL state dict onto rank-0 and persist it. -----
    # For models too large to gather, switch to StateDictType.SHARDED_STATE_DICT
    # and have EVERY rank save its own shard under a rank-suffixed key instead.
    save_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_cfg):
        full_state = fsdp_model.state_dict()  # populated only on rank-0

    cp = ctx.checkpoint
    if rank == 0 and cp is not None:
        buf = io.BytesIO()
        torch.save(full_state, buf)
        await cp.save(buf.getvalue())
        print(f"[rank 0] saved FSDP checkpoint to {cp.path}", flush=True)

    dist.barrier()
    dist.destroy_process_group()
    print(f"[rank {rank}] done — final loss {last_loss:.5f}", flush=True)
    return last_loss

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(train_fsdp, steps=30)
    print("Run URL:", run.url)
    run.wait()
    print("Final phase:", run.phase)

# /// script
# requires-python = "==3.13"
# dependencies = [
#    "flyte>=2.5.18",
#    "torch",
#    "torchvision",
#    "lightning",
# ]
# main = "train_lightning"
# params = "max_steps=50"
# ///
"""
PyTorch Lightning training on a ClusteredTaskEnvironment.

Lightning has its own multi-process launching machinery, but it also rides the
torchrun contract: we let the clustered runtime start one process per rank and
configure Lightning's `Trainer` to attach to the EXISTING process group rather
than spawning its own. The key is `strategy="ddp"` (not a `*_spawn` strategy)
with `devices`/`num_nodes` taken from the clustered context — Lightning then
reads RANK / LOCAL_RANK / WORLD_SIZE / MASTER_ADDR from the environment.
"""

from __future__ import annotations

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

USE_GPU = True
REPLICAS = 2  # nodes
NPROC_PER_NODE = 1  # processes (GPUs) per node

image = flyte.Image.from_debian_base().with_pip_packages("torch", "torchvision", "lightning")

resources = (
    flyte.Resources(cpu=(2, 4), memory=("4Gi", "8Gi"), gpu="L4:1")
    if USE_GPU
    else flyte.Resources(cpu=(2, 4), memory=("2Gi", "4Gi"))
)

env = ClusteredTaskEnvironment(
    name="lightning_env",
    image=image,
    resources=resources,
    replicas=REPLICAS,
    nproc_per_node=NPROC_PER_NODE,
    runtime=TorchRun(rdzv_backend="static", max_restarts=0),
    failure_policy=ClusterFailurePolicy(max_restarts=1),
)

@env.task
async def train_lightning(max_steps: int = 50) -> float:
    """Train a tiny MNIST autoencoder with Lightning DDP over the JobSet's ranks."""
    import lightning as L
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import DataLoader, TensorDataset

    ctx = flyte.ctx()
    print(
        f"[rank {ctx.rank}/{ctx.world_size}] node_rank={ctx.node_rank} "
        f"nnodes={ctx.nnodes} local_rank={ctx.local_rank}",
        flush=True,
    )

    class AutoEncoder(L.LightningModule):
        def __init__(self):
            super().__init__()
            self.encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
            self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

        def training_step(self, batch, _):
            x, _y = batch
            x = x.view(x.size(0), -1)
            z = self.encoder(x)
            loss = F.mse_loss(self.decoder(z), x)
            self.log("train_loss", loss, prog_bar=True)
            return loss

        def configure_optimizers(self):
            return torch.optim.Adam(self.parameters(), lr=1e-3)

    # Synthetic MNIST-shaped data so the example needs no dataset download; swap
    # in torchvision MNIST for the real thing. Lightning's DistributedSampler
    # shards this across ranks automatically.
    torch.manual_seed(0)
    images = torch.rand(512, 1, 28, 28)
    labels = torch.zeros(512, dtype=torch.long)
    loader = DataLoader(TensorDataset(images, labels), batch_size=32, shuffle=True)

    # Attach to the torchrun-provided process group: ddp (not ddp_spawn), with
    # devices per node and num_nodes from the clustered context.
    trainer = L.Trainer(
        accelerator="gpu" if USE_GPU else "cpu",
        devices=NPROC_PER_NODE,
        num_nodes=ctx.nnodes or REPLICAS,
        strategy="ddp",
        max_steps=max_steps,
        enable_checkpointing=False,
        logger=False,
    )
    model = AutoEncoder()
    trainer.fit(model, loader)

    final_loss = float(trainer.callback_metrics.get("train_loss", torch.tensor(0.0)))
    print(f"[rank {ctx.rank}] done — final loss {final_loss:.5f}", flush=True)
    return final_loss

if __name__ == "__main__":
    flyte.init_from_config()
    run = flyte.run(train_lightning, max_steps=50)
    print("Run URL:", run.url)
    run.wait()
    print("Final phase:", run.phase)