trainer/examples/local/local-container-mnist.ipynb

Kubeflow Trainer: Container Backend (Single-Node) Training

This notebook demonstrates how to run single-node training using the Container Backend with Docker or Podman.

Container Backend

• Container Runtime: Docker or Podman required
• Use Case: Testing container workflows, simulating production environments
• Prerequisites: Python 3.9+ and Docker Desktop/Engine OR Podman

This example trains a CNN on the classic MNIST handwritten digit dataset using PyTorch.

Install the Kubeflow SDK

You need to install the Kubeflow SDK with container backend support:

In [ ]:

# Uncomment to install
# %pip install -U kubeflow[docker]  # For Docker
# %pip install -U kubeflow[podman]  # For Podman

Define the Training Function

The first step is to create a function to train CNN model using MNIST data.

In [ ]:

def train_mnist():
    import os
    import torch
    import torch.nn.functional as F
    from torch import nn, optim
    import torch.distributed as dist
    from torch.utils.data import DataLoader, DistributedSampler
    from torchvision import datasets, transforms

    # Define the PyTorch CNN model to be trained
    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(1, 20, 5, 1)
            self.conv2 = nn.Conv2d(20, 50, 5, 1)
            self.fc1 = nn.Linear(4 * 4 * 50, 500)
            self.fc2 = nn.Linear(500, 10)

        def forward(self, x):
            x = F.relu(self.conv1(x))
            x = F.max_pool2d(x, 2, 2)
            x = F.relu(self.conv2(x))
            x = F.max_pool2d(x, 2, 2)
            x = x.view(-1, 4 * 4 * 50)
            x = F.relu(self.fc1(x))
            x = self.fc2(x)
            return F.log_softmax(x, dim=1)

    # Use NCCL if a GPU is available, otherwise use Gloo as communication backend.
    device, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")
    print(f"Using Device: {device}, Backend: {backend}")

    # Setup PyTorch distributed.
    local_rank = int(os.getenv("LOCAL_RANK", 0))
    dist.init_process_group(backend=backend)
    print(
        "Distributed Training for WORLD_SIZE: {}, RANK: {}, LOCAL_RANK: {}".format(
            dist.get_world_size(),
            dist.get_rank(),
            local_rank,
        )
    )

    # Create the model and load it into the device.
    device = torch.device(f"{device}:{local_rank}")
    model=Net().to(device)
    if torch.cuda.is_available():
        model=torch.compile(model)
    model = nn.parallel.DistributedDataParallel(model)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

    # Download FashionMNIST dataset only on global rank 0 process.
    if dist.get_rank() == 0:
        datasets.FashionMNIST(
            "./data",
            train=True,
            download=True,
        )
    dist.barrier()
    dataset = datasets.FashionMNIST(
        "./data",
        train=True,
        download=False,
        transform=transforms.Compose([transforms.ToTensor()]),
    )

    # Shard the dataset across workers.
    train_loader = DataLoader(
        dataset,
        batch_size=100,
        sampler=DistributedSampler(dataset)
    )

    dist.barrier()
    for epoch in range(1, 3):
        model.train()
        
        # Iterate over mini-batches from the training set
        for batch_idx, (data, target) in enumerate(train_loader):
            # Forward pass
            data, target = data.to(device), target.to(device)
            outputs = model(data)
            loss = F.nll_loss(outputs, target)
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if batch_idx % 100 == 0:
                print(
                    "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                        epoch,
                        batch_idx * len(data),
                        len(train_loader.dataset),
                        100.0 * batch_idx / len(train_loader),
                        loss.item(),
                    )
                )

    # Wait for the distributed training to complete
    dist.barrier()
    if dist.get_rank() == 0:
        print("Training is finished")
        torch.save(model.state_dict(), "mnist_cnn.pt")

    # Finally clean up PyTorch distributed
    dist.destroy_process_group()

Configure Container Backend

The container backend automatically detects and uses either Docker or Podman:

In [ ]:

from kubeflow.trainer import TrainerClient, ContainerBackendConfig
import os

# Auto-detects Docker or Podman
backend_config = ContainerBackendConfig()

# Optional: Force specific runtime
# backend_config = ContainerBackendConfig(runtime="docker")  # Force Docker
# backend_config = ContainerBackendConfig(runtime="podman")  # Force Podman

# Optional: For Colima on macOS
# backend_config = ContainerBackendConfig(
#     container_host=f"unix://{os.path.expanduser('~')}/.colima/default/docker.sock"
# )

# Optional: For Podman Machine on macOS
# backend_config = ContainerBackendConfig(
#     runtime="podman",
#     container_host="unix:///run/user/1000/podman/podman.sock"
# )

Initialize Client

Initialize the TrainerClient with the Container Backend:

In [1]:

client = TrainerClient(backend_config=backend_config)

List the Training Runtimes

You can get the list of available Training Runtimes to start your TrainJob.

In [ ]:

for runtime in client.list_runtimes():
    print(runtime)
    if runtime.name == "torch-distributed":
        torch_runtime = runtime

Run the TrainJob

Submit the training job to the Container Backend (single container):

In [1]:

from kubeflow.trainer import CustomTrainer

job_name = client.train(
    trainer=CustomTrainer(
        func=train_mnist,
        packages_to_install=["torchvision"],
        num_nodes=2,
    ),
    runtime=torch_runtime,
)

Check the TrainJob Status

You can check the status of the TrainJob that's created.

In [1]:

job = client.get_job(job_name)
print("Job: {}, Status: {}".format(job.name, job.status))

Job: p6072ef2ca48, Status: Running

Watch the TrainJob Logs

We can use the get_job_logs() API to get the TrainJob logs.

In [ ]:

for logline in client.get_job_logs(job_name, follow=True):
    print(logline, end='')

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:

client.wait_for_job_status(name=job_name, timeout=20)

Delete the TrainJob

When the TrainJob is finished, you can delete the resource.

In [ ]:

client.delete_job(job_name)