Batch (Kubeflow Train)¶

The JupyterHub notebook interface we used in the previous guide is great for interactive exploration and prototyping, but less appropriate for large scale training jobs that need to run unattended for many hours or days.

In this guide, as before, you'll fine-tune a vision transformer on the Food-101 dataset. This time, however, we'll use the Kubeflow Train Python SDK to run the training via FLAME's Kubeflow Trainer platform to get around the limits of the notebook interface.

Before you start¶

In this guide, we'll work on FLAME using Visual Studio Code, rather than the Jupyter notebook interface. This gives you a file editor and terminal in the cluster, just as though you had connected VSCode to an ordinary Linux server using SSH.

Follow the instructions in the VSCode page to connected to FLAME with VSCode. When launching your JupyterHub server, this time don't choose to add GPU resources. Here our JupyterHub server is just playing host to VSCode for us to edit files. Our actual training will run as it's own independent pod, which we'll request GPUs for separately.

Set up your project¶

First let's create a project folder and install the libraries the training code needs into it.

We'll use uv¹, to create the project and manage its dependencies. In your VSCode terminal (the one connected to the pod), create the project in your home directory:

cd ~
uv init food101
cd food101

(uv init also creates a sample main.py you can ignore or delete.)

Because we're now working in an isolated virtual environment, we need to install PyTorch ourselves rather than inheriting the PyTorch that comes in the image's global Python environment. The --index argument specifies that we want a version of PyTorch built against CUDA 13, which is the default on FLAME. Run the following in a terminal (it will take a while to download):

uv add torch --index https://download.pytorch.org/whl/cu132

Then, install the other dependencies the code will use:

uv add lightning==2.6.5 timm==1.0.27 datasets==2.19.1 pillow==12.2.0 kubeflow==0.4.1

Create the training script¶

Now that the dependencies are installed, create a single .py file containing all the training code. This is mostly the same code from the notebook guide, just gathered into one file rather than spread across notebook cells.

Create train.py inside the project folder you just made — use File → New File and save it there, or run this in the VSCode terminal:

code ~/food101/train.py

Paste the following into train.py and save it:

import torch
import timm
import lightning as L
from lightning.pytorch.loggers import CSVLogger
from torch.utils.data import DataLoader
from datasets import load_dataset

class LitViT(L.LightningModule):
    def __init__(self):
        super().__init__()
        # 101 classes matches the number in our training data
        self.model = timm.create_model("vit_base_patch16_224", pretrained=True, num_classes=101)
        self.criterion = torch.nn.CrossEntropyLoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.criterion(logits, y)
        accuracy = (logits.argmax(dim=1) == y).float().mean()

        # Record loss and accuracy once per epoch. In a batch job these go to
        # the CSV logger configured below, not a live progress bar.
        self.log_dict(
            {"loss": loss, "acc": accuracy},
            on_step=False,
            on_epoch=True,
        )
        return loss

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


def main():
    model = LitViT()

    data_cfg = timm.data.resolve_model_data_config(model.model)
    transform = timm.data.create_transform(**data_cfg, is_training=True)

    # The first call downloads several GB and caches it; later runs reuse the cache.
    train_ds = load_dataset("ethz/food101", split="train")

    def collate(examples):
        images = torch.stack([transform(e["image"].convert("RGB")) for e in examples])
        labels = torch.tensor([e["label"] for e in examples])
        return images, labels

    loader = DataLoader(
        train_ds,
        batch_size=1024,
        num_workers=8,
        collate_fn=collate,
        drop_last=True,
        shuffle=True,
    )

    trainer = L.Trainer(
        precision="bf16-mixed",
        max_epochs=10,
        enable_checkpointing=True,
        default_root_dir="/personal/food101",
        enable_progress_bar=False,
        logger=CSVLogger("/personal/food101", flush_logs_every_n_steps=73),
    )
    trainer.fit(model, loader)

    # Save the trained model
    trainer.save_checkpoint("/personal/food101/final.ckpt")


if __name__ == "__main__":
    main()

/personal is another way to refer to your home directory (see Storage for more details).

This is more-or-less the same code from the notebook version of the guide, just merged into one file rather than spread across seperate cells, and with a few tweaks for running non-interactively (e.g. logging to a CSV file rather than standard out).

Submit your training job¶

Now we'll write a small wrapper script to tell Kubeflow how to run our training. Save the following as submit.py in your workspace:

from kubeflow.trainer import TrainerClient, CustomTrainer

def launch():
    import subprocess
    subprocess.run(
        ["uv", "run", "python", "train.py"],
        cwd="/personal/food101",
        check=True,
    )


client = TrainerClient()

job_name = client.train(
    runtime="torch-gh200",
    trainer=CustomTrainer(
        func=launch,
        resources_per_node={"cpu": 8, "memory": "32Gi", "gpu": 1},
    ),
)

print(f"Submitted TrainJob: {job_name}")

There are two parts here:

1. launch, a small wrapper function that uses uv run to execute our training code. Note that we import subprocess inside the function, contrary to typical Python best practices. This is because kubeflow only runs literally the code inside this function when it executes the training

2. All the TrainerClient stuff at the bottom. This is mostly boilerplate for connecting to the Kubernetes API. The important parts are the arguments to CustomTrainer:

   • runtime="torch-gh200". A "runtime" is what Kubeflow calls a bundle of configuration for a training run: what container image to use, what GPU type you want, what environment variables to set, etc. Here we use "torch-gh200" to run on FLAME's NVIDIA GH200 GPUs².
   • func=launch: The launch function from above is what we want to run.
   • resources_per_node: This tells kubeflow how much memory, CPUs, and GPUs our training needs.

Submit it¶

Run the submission script from your VSCode terminal:

uv run python submit.py

It prints the ID of the job it created:

Submitted TrainJob: cda94e1a31a8 # This is an example, your job name will be different

Confirm it started¶

Now run kubectl get pods. You should see a new pod whose ID is similar to the trainjob ID. e.g. for the example ID above, you might see:

$ kubectl get pods
NAME                                   READY   STATUS              RESTARTS   AGE
cda94e1a31a8-node-0-0-hsw94            0/1     ContainerCreating   0          24s
...

Let's check the logs from this pod (-f means "follow", so we stream the pods logs to the terminal as they arrive):

$ kubectl logs -f cda94e1a31a8-node-0-0-hsw94

You might not see anything for a minute or two while the pod starts, but then you should see it install the dependencies, download the dataset, and begin training.

Press Ctrl + c to exit the log viewer.

Follow training progress¶

train.py writes one row per epoch with the training metrics we defined (loss, accuracy, etc.) to a CSV file under your project folder. You can follow the logs as they come in with:

tail -f ~/food101/lightning_logs/version_0/metrics.csv

Each run creates a new version_N folder (version_0, version_1, …), so you may need to check the latest version and look there instead of version_0 if you've done multiple runs You'll see a row appear per epoch, with the loss falling and accuracy climbing as it trains (exact numbers will vary):

epoch,step,loss,acc
0,72,3.8201,0.1210
1,145,2.5500,0.4300
2,218,1.7100,0.6170
...

When the run finishes, the script writes the trained model to ~/food101/final.ckpt for you to inspect, validate, or use for inference.

Where to go next¶

Some ideas to try next:

• Train a model from scratch. Since a TrainJob can run unattended for many hours or days, you can use it to run much longer trainings, for example, training a model from scratch. Pass pretrained=False to timm.create_model(...) to initialize the ViT randomly and learn Food-101 from the ground up. You'll also need to up the max_epochs in the L.Trainer definition to a much larger number (e.g. a few hundred) to give the training from scratch more time to learn.
• Use more Kubeflow features. The Kubeflow SDK has a lot more useful features you can take advantage of. For example:
  • Data and Model Initializers make GPU resource usage more efficient by pre-downloading the data in pods without GPUs before starting the training pod proper.
  • Hyperparameter tuning automates optimization of hyperparameters (learning rate, model architecture, batch size, etc.)
  • Distributed Training makes it easy to scale your training to multiple nodes.