Files
Andrey VelichkevichandGitHub fbacd9f577 fix(examples): Verify TrainJob Completion (#3331)
* fix(examples): Verify TrainJob Completion

Signed-off-by: Andrey Velichkevich <[email protected]>

* Rename client to TrainerClient()

Signed-off-by: Andrey Velichkevich <[email protected]>

---------

Signed-off-by: Andrey Velichkevich <[email protected]>
2026-03-16 14:40:31 +00:00

29 KiB

PyTorch DDP Fashion MNIST Training Example

This example demonstrates how to train a convolutional neural network (CNN) to classify images using the Fashion MNIST dataset and PyTorch Distributed Data Parallel (DDP).

This notebook walks you through running that example locally, and how to easily scale PyTorch DDP across multiple nodes with Kubeflow TrainJob.

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:

In [ ]:
# !pip install -U kubeflow

Install the PyTorch Dependencies

You also need to install PyTorch and Torchvision to be able to run the example locally:

In [ ]:
!pip install torch==2.10.0
!pip install torchvision==0.22.1

Define the Training Function

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

In [1]:
def train_fashion_mnist(num_epochs=2, batch_size=100, lr=0.1, momentum=0.9):
    import os

    import torch
    import torch.distributed as dist
    import torch.nn.functional as F
    from torch import nn
    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 = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    
    
    # Download FashionMNIST dataset only on local_rank=0 process.
    if local_rank == 0:
        dataset = datasets.FashionMNIST(
            "./data",
            train=True,
            download=True,
            transform=transforms.Compose([transforms.ToTensor()]),
        )
    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=batch_size,
        sampler=DistributedSampler(dataset)
    )

    dist.barrier()
    for epoch in range(1, num_epochs + 1):
        model.train()

        # Iterate over mini-batches from the training set
        for batch_idx, (inputs, labels) in enumerate(train_loader):
            # Copy the data to the GPU device if available
            inputs, labels = inputs.to(device), labels.to(device)
            # Forward pass
            outputs = model(inputs)
            loss = F.nll_loss(outputs, labels)
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if batch_idx % 10 == 0 and dist.get_rank() == 0:
                print(
                    "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                        epoch,
                        batch_idx * len(inputs),
                        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")

    # Finally clean up PyTorch distributed
    dist.destroy_process_group()

Run the Training Locally

We can submit the training function to the local Trainer client to run it in an isolated subprocess.

In [3]:
from kubeflow.trainer import CustomTrainer, TrainerClient, LocalProcessBackendConfig

# Initialize local backend
backend_config = LocalProcessBackendConfig(cleanup_venv=True)
client = TrainerClient(backend_config=backend_config)

# List available runtimes
for runtime in client.list_runtimes():
    if runtime.name == "torch-distributed":
        torch_runtime = runtime
        break

# Submit training job
job_name = client.train(
    trainer=CustomTrainer(
        func=train_fashion_mnist,
        packages_to_install=["torch", "torchvision"],
    ),
    runtime=torch_runtime,
)

# Stream logs
for logline in client.get_job_logs(job_name, follow=True):
    print(logline, end='')
Using Device: cpu, Backend: gloo
Distributed Training for WORLD_SIZE: 1, RANK: 0, LOCAL_RANK: 0
Train Epoch: 1 [0/60000 (0%)]	Loss: 2.300957
Train Epoch: 1 [1000/60000 (2%)]	Loss: 2.068733
Train Epoch: 1 [2000/60000 (3%)]	Loss: 1.136209
Train Epoch: 1 [3000/60000 (5%)]	Loss: 0.897645
Train Epoch: 1 [4000/60000 (7%)]	Loss: 0.868213
Train Epoch: 1 [5000/60000 (8%)]	Loss: 0.943094
Train Epoch: 1 [6000/60000 (10%)]	Loss: 0.785388
Train Epoch: 1 [7000/60000 (12%)]	Loss: 0.772929
Train Epoch: 1 [8000/60000 (13%)]	Loss: 0.612156
Train Epoch: 1 [9000/60000 (15%)]	Loss: 0.783617
Train Epoch: 1 [10000/60000 (17%)]	Loss: 0.608218
Train Epoch: 1 [11000/60000 (18%)]	Loss: 0.595062
Train Epoch: 1 [12000/60000 (20%)]	Loss: 0.612815
Train Epoch: 1 [13000/60000 (22%)]	Loss: 0.523903
Train Epoch: 1 [14000/60000 (23%)]	Loss: 0.673269
Train Epoch: 1 [15000/60000 (25%)]	Loss: 0.675361
Train Epoch: 1 [16000/60000 (27%)]	Loss: 0.591610
Train Epoch: 1 [17000/60000 (28%)]	Loss: 0.515040
Train Epoch: 1 [18000/60000 (30%)]	Loss: 0.535079
Train Epoch: 1 [19000/60000 (32%)]	Loss: 0.721351
Train Epoch: 1 [20000/60000 (33%)]	Loss: 0.413175
Train Epoch: 1 [21000/60000 (35%)]	Loss: 0.407277
Train Epoch: 1 [22000/60000 (37%)]	Loss: 0.356355
Train Epoch: 1 [23000/60000 (38%)]	Loss: 0.685514
Train Epoch: 1 [24000/60000 (40%)]	Loss: 0.439642
Train Epoch: 1 [25000/60000 (42%)]	Loss: 0.399821
Train Epoch: 1 [26000/60000 (43%)]	Loss: 0.556694
Train Epoch: 1 [27000/60000 (45%)]	Loss: 0.326490
Train Epoch: 1 [28000/60000 (47%)]	Loss: 0.401890
Train Epoch: 1 [29000/60000 (48%)]	Loss: 0.449772
Train Epoch: 1 [30000/60000 (50%)]	Loss: 0.427323
Train Epoch: 1 [31000/60000 (52%)]	Loss: 0.400165
Train Epoch: 1 [32000/60000 (53%)]	Loss: 0.334357
Train Epoch: 1 [33000/60000 (55%)]	Loss: 0.298294
Train Epoch: 1 [34000/60000 (57%)]	Loss: 0.551655
Train Epoch: 1 [35000/60000 (58%)]	Loss: 0.456408
Train Epoch: 1 [36000/60000 (60%)]	Loss: 0.409366
Train Epoch: 1 [37000/60000 (62%)]	Loss: 0.449682
Train Epoch: 1 [38000/60000 (63%)]	Loss: 0.433943
Train Epoch: 1 [39000/60000 (65%)]	Loss: 0.403032
Train Epoch: 1 [40000/60000 (67%)]	Loss: 0.427653
Train Epoch: 1 [41000/60000 (68%)]	Loss: 0.596136
Train Epoch: 1 [42000/60000 (70%)]	Loss: 0.451446
Train Epoch: 1 [43000/60000 (72%)]	Loss: 0.415507
Train Epoch: 1 [44000/60000 (73%)]	Loss: 0.431509
Train Epoch: 1 [45000/60000 (75%)]	Loss: 0.269870
Train Epoch: 1 [46000/60000 (77%)]	Loss: 0.603037
Train Epoch: 1 [47000/60000 (78%)]	Loss: 0.361142
Train Epoch: 1 [48000/60000 (80%)]	Loss: 0.533713
Train Epoch: 1 [49000/60000 (82%)]	Loss: 0.334492
Train Epoch: 1 [50000/60000 (83%)]	Loss: 0.360594
Train Epoch: 1 [51000/60000 (85%)]	Loss: 0.479819
Train Epoch: 1 [52000/60000 (87%)]	Loss: 0.408797
Train Epoch: 1 [53000/60000 (88%)]	Loss: 0.402639
Train Epoch: 1 [54000/60000 (90%)]	Loss: 0.291429
Train Epoch: 1 [55000/60000 (92%)]	Loss: 0.308326
Train Epoch: 1 [56000/60000 (93%)]	Loss: 0.503240
Train Epoch: 1 [57000/60000 (95%)]	Loss: 0.449824
Train Epoch: 1 [58000/60000 (97%)]	Loss: 0.339513
Train Epoch: 1 [59000/60000 (98%)]	Loss: 0.260989
Train Epoch: 2 [0/60000 (0%)]	Loss: 0.379833
Train Epoch: 2 [1000/60000 (2%)]	Loss: 0.421474
Train Epoch: 2 [2000/60000 (3%)]	Loss: 0.298147
Train Epoch: 2 [3000/60000 (5%)]	Loss: 0.302185
Train Epoch: 2 [4000/60000 (7%)]	Loss: 0.407515
Train Epoch: 2 [5000/60000 (8%)]	Loss: 0.372201
Train Epoch: 2 [6000/60000 (10%)]	Loss: 0.403835
Train Epoch: 2 [7000/60000 (12%)]	Loss: 0.403082
Train Epoch: 2 [8000/60000 (13%)]	Loss: 0.429058
Train Epoch: 2 [9000/60000 (15%)]	Loss: 0.427616
Train Epoch: 2 [10000/60000 (17%)]	Loss: 0.353761
Train Epoch: 2 [11000/60000 (18%)]	Loss: 0.392035
Train Epoch: 2 [12000/60000 (20%)]	Loss: 0.315083
Train Epoch: 2 [13000/60000 (22%)]	Loss: 0.303653
Train Epoch: 2 [14000/60000 (23%)]	Loss: 0.477725
Train Epoch: 2 [15000/60000 (25%)]	Loss: 0.311283
Train Epoch: 2 [16000/60000 (27%)]	Loss: 0.320934
Train Epoch: 2 [17000/60000 (28%)]	Loss: 0.344658
Train Epoch: 2 [18000/60000 (30%)]	Loss: 0.252863
Train Epoch: 2 [19000/60000 (32%)]	Loss: 0.433338
Train Epoch: 2 [20000/60000 (33%)]	Loss: 0.401821
Train Epoch: 2 [21000/60000 (35%)]	Loss: 0.277403
Train Epoch: 2 [22000/60000 (37%)]	Loss: 0.324770
Train Epoch: 2 [23000/60000 (38%)]	Loss: 0.450259
Train Epoch: 2 [24000/60000 (40%)]	Loss: 0.280156
Train Epoch: 2 [25000/60000 (42%)]	Loss: 0.273340
Train Epoch: 2 [26000/60000 (43%)]	Loss: 0.405919
Train Epoch: 2 [27000/60000 (45%)]	Loss: 0.270347
Train Epoch: 2 [28000/60000 (47%)]	Loss: 0.389775
Train Epoch: 2 [29000/60000 (48%)]	Loss: 0.266520
Train Epoch: 2 [30000/60000 (50%)]	Loss: 0.369550
Train Epoch: 2 [31000/60000 (52%)]	Loss: 0.307309
Train Epoch: 2 [32000/60000 (53%)]	Loss: 0.282144
Train Epoch: 2 [33000/60000 (55%)]	Loss: 0.259714
Train Epoch: 2 [34000/60000 (57%)]	Loss: 0.455187
Train Epoch: 2 [35000/60000 (58%)]	Loss: 0.418436
Train Epoch: 2 [36000/60000 (60%)]	Loss: 0.347754
Train Epoch: 2 [37000/60000 (62%)]	Loss: 0.285930
Train Epoch: 2 [38000/60000 (63%)]	Loss: 0.298169
Train Epoch: 2 [39000/60000 (65%)]	Loss: 0.315216
Train Epoch: 2 [40000/60000 (67%)]	Loss: 0.207120
Train Epoch: 2 [41000/60000 (68%)]	Loss: 0.447531
Train Epoch: 2 [42000/60000 (70%)]	Loss: 0.356669
Train Epoch: 2 [43000/60000 (72%)]	Loss: 0.295995
Train Epoch: 2 [44000/60000 (73%)]	Loss: 0.306820
Train Epoch: 2 [45000/60000 (75%)]	Loss: 0.209821
Train Epoch: 2 [46000/60000 (77%)]	Loss: 0.474326
Train Epoch: 2 [47000/60000 (78%)]	Loss: 0.292024
Train Epoch: 2 [48000/60000 (80%)]	Loss: 0.554571
Train Epoch: 2 [49000/60000 (82%)]	Loss: 0.313117
Train Epoch: 2 [50000/60000 (83%)]	Loss: 0.247606
Train Epoch: 2 [51000/60000 (85%)]	Loss: 0.385026
Train Epoch: 2 [52000/60000 (87%)]	Loss: 0.369550
Train Epoch: 2 [53000/60000 (88%)]	Loss: 0.326454
Train Epoch: 2 [54000/60000 (90%)]	Loss: 0.317689
Train Epoch: 2 [55000/60000 (92%)]	Loss: 0.296784
Train Epoch: 2 [56000/60000 (93%)]	Loss: 0.399627
Train Epoch: 2 [57000/60000 (95%)]	Loss: 0.408119
Train Epoch: 2 [58000/60000 (97%)]	Loss: 0.341115
Train Epoch: 2 [59000/60000 (98%)]	Loss: 0.191851
Training is finished

Scale PyTorch DDP with Kubeflow TrainJob

You can use TrainerClient() from the Kubeflow SDK to communicate with Kubeflow Trainer APIs and scale your training function across multiple PyTorch training nodes.

TrainerClient() verifies that you have required access to the Kubernetes cluster.

Kubeflow Trainer creates a TrainJob resource and automatically sets the appropriate environment variables to set up PyTorch in distributed environment.

In [2]:
from kubeflow.trainer import CustomTrainer, TrainerClient

client = TrainerClient()

List the Training Runtimes

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

Additionally, it might show available accelerator type and number of available resources.

In [3]:
for runtime in client.list_runtimes():
    print(runtime)
    if runtime.name == "torch-distributed":
        torch_runtime = runtime
Runtime(name='deepspeed-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='deepspeed', num_nodes=1, device='Unknown', device_count='1'), pretrained_model=None)
Runtime(name='mlx-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='mlx', num_nodes=1, device='gpu', device_count='1'), pretrained_model=None)
Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None)
Runtime(name='torchtune-llama3.2-1b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, device='gpu', device_count='2'), pretrained_model=None)
Runtime(name='torchtune-llama3.2-3b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, device='gpu', device_count='2.0'), pretrained_model=None)

Run the Distributed TrainJob

Kubeflow TrainJob will train the above model on 3 PyTorch nodes.

In [4]:
job_name = client.train(
    trainer=CustomTrainer(
        func=train_fashion_mnist,
        func_args={"num_epochs": 3, "batch_size": 64, "lr": 0.01, "momentum": 0.9},
        # Set how many PyTorch nodes you want to use for distributed training.
        num_nodes=3,
        # Set the resources for each PyTorch node.
        resources_per_node={
            "cpu": 3,
            "memory": "16Gi",
            # Uncomment this to distribute the TrainJob using GPU nodes.
            # "nvidia.com/gpu": 1,
        },
    ),
    runtime=torch_runtime,
)

Check the TrainJob steps

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

Since the TrainJob performs distributed training across 3 nodes, it generates 3 steps: trainer-node-0 .. trainer-node-2.

You can get the individual status for each of these steps.

In [5]:
# Wait for the running status.
client.wait_for_job_status(name=job_name, status={"Running"})
Out [5]:
TrainJob(name='j819fe212f8e', creation_timestamp=datetime.datetime(2025, 9, 3, 13, 19, 56, tzinfo=TzInfo(UTC)), runtime=Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', num_nodes=1, device='Unknown', device_count='Unknown'), pretrained_model=None), steps=[Step(name='node-0', status='Running', pod_name='j819fe212f8e-node-0-0-nlcqt', device='gpu', device_count='1'), Step(name='node-1', status='Running', pod_name='j819fe212f8e-node-0-1-rjztc', device='gpu', device_count='1'), Step(name='node-2', status='Running', pod_name='j819fe212f8e-node-0-2-l575c', device='gpu', device_count='1')], num_nodes=3, status='Running')
In [6]:
for c in client.get_job(name=job_name).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")
Step: node-0, Status: Running, Devices: gpu x 1

Step: node-1, Status: Running, Devices: gpu x 1

Step: node-2, Status: Running, Devices: gpu x 1

Watch the TrainJob logs

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

Since we run training on 3 GPUs, every PyTorch node uses 60,000/3 = 20,000 images from the dataset.

In [7]:
for logline in client.get_job_logs(job_name, follow=True):
    print(logline)
[W903 13:20:06.250970596 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known).
[W903 13:20:07.936850072 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known).
[W903 13:20:08.873279869 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known).
[W903 13:20:08.673190939 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known).
[W903 13:20:09.683786251 socket.cpp:755] [c10d] The IPv6 network addresses of (j819fe212f8e-node-0-0.j819fe212f8e, 29500) cannot be retrieved (gai error: -2 - Name or service not known).
Using Device: cuda, Backend: nccl
Distributed Training for WORLD_SIZE: 3, RANK: 0, LOCAL_RANK: 0
100%|██████████| 26.4M/26.4M [00:02<00:00, 12.2MB/s]
100%|██████████| 29.5k/29.5k [00:00<00:00, 213kB/s]
100%|██████████| 4.42M/4.42M [00:01<00:00, 3.85MB/s]
100%|██████████| 5.15k/5.15k [00:00<00:00, 43.7MB/s]
/opt/conda/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py:4631: UserWarning: No device id is provided via `init_process_group` or `barrier `. Using the current device set by the user. 
  warnings.warn(  # warn only once
Train Epoch: 1 [0/60000 (0%)]	Loss: 2.297219
Train Epoch: 1 [1000/60000 (5%)]	Loss: 1.803301
Train Epoch: 1 [2000/60000 (10%)]	Loss: 1.663171
Train Epoch: 1 [3000/60000 (15%)]	Loss: 1.462810
Train Epoch: 1 [4000/60000 (20%)]	Loss: 1.270390
Train Epoch: 1 [5000/60000 (25%)]	Loss: 0.977784
Train Epoch: 1 [6000/60000 (30%)]	Loss: 0.791962
Train Epoch: 1 [7000/60000 (35%)]	Loss: 0.626636
Train Epoch: 1 [8000/60000 (40%)]	Loss: 0.613735
Train Epoch: 1 [9000/60000 (45%)]	Loss: 0.590524
Train Epoch: 1 [10000/60000 (50%)]	Loss: 0.529664
Train Epoch: 1 [11000/60000 (55%)]	Loss: 0.448692
Train Epoch: 1 [12000/60000 (60%)]	Loss: 0.508065
Train Epoch: 1 [13000/60000 (65%)]	Loss: 0.519867
Train Epoch: 1 [14000/60000 (70%)]	Loss: 0.485550
Train Epoch: 1 [15000/60000 (75%)]	Loss: 0.377857
Train Epoch: 1 [16000/60000 (80%)]	Loss: 0.456072
Train Epoch: 1 [17000/60000 (85%)]	Loss: 0.563855
Train Epoch: 1 [18000/60000 (90%)]	Loss: 0.426923
Train Epoch: 1 [19000/60000 (95%)]	Loss: 0.379981
Train Epoch: 2 [0/60000 (0%)]	Loss: 0.416116
Train Epoch: 2 [1000/60000 (5%)]	Loss: 0.506344
Train Epoch: 2 [2000/60000 (10%)]	Loss: 0.510181
Train Epoch: 2 [3000/60000 (15%)]	Loss: 0.388432
Train Epoch: 2 [4000/60000 (20%)]	Loss: 0.361425
Train Epoch: 2 [5000/60000 (25%)]	Loss: 0.428971
Train Epoch: 2 [6000/60000 (30%)]	Loss: 0.409087
Train Epoch: 2 [7000/60000 (35%)]	Loss: 0.423781
Train Epoch: 2 [8000/60000 (40%)]	Loss: 0.530143
Train Epoch: 2 [9000/60000 (45%)]	Loss: 0.390188
Train Epoch: 2 [10000/60000 (50%)]	Loss: 0.359034
Train Epoch: 2 [11000/60000 (55%)]	Loss: 0.353203
Train Epoch: 2 [12000/60000 (60%)]	Loss: 0.394956
Train Epoch: 2 [13000/60000 (65%)]	Loss: 0.407927
Train Epoch: 2 [14000/60000 (70%)]	Loss: 0.296730
Train Epoch: 2 [15000/60000 (75%)]	Loss: 0.318915
Train Epoch: 2 [16000/60000 (80%)]	Loss: 0.317562
Train Epoch: 2 [17000/60000 (85%)]	Loss: 0.396108
Train Epoch: 2 [18000/60000 (90%)]	Loss: 0.302800
Train Epoch: 2 [19000/60000 (95%)]	Loss: 0.364170
Training is finished

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 TrainJob is finished, you can delete the resource.

In [ ]:
# client.delete_job(job_name)
In [ ]: