Kubeflow PyTorch

flytekitplugins-kfpytorch

FlytekitML Trainingpytorchkubeflowdistributed-trainingdeep-learning

This plugin uses the Kubeflow PyTorch Operator and provides an extremely simplified interface for executing distributed training using various PyTorch backends.

Install

pip install flytekitplugins-kfpytorch

Quick Start(example, may need adjustment)

See full examples

pip install flytekitplugins-kfpytorch

from flytekit import task, workflow
from flytekitplugins.kfpytorch import CleanPodPolicy, Elastic, Master, PyTorch

config = CleanPodPolicy(...)

@task
def my_task() -> None:
    ...

@workflow
def my_workflow() -> None:
    my_task()

Available Imports (7)

enumCleanPodPolicy

CleanPodPolicy describes how to deal with pods when the job is finished.

extends Enum — enumeration of predefined options

from flytekitplugins.kfpytorch import CleanPodPolicy

configElastic

Configuration for [`torch elastic training`](https://pytorch.org/docs/stable/elastic/run.html).

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.kfpytorch import Elastic

configMaster

Configuration for master replica group.

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.kfpytorch import Master

configPyTorch

Configuration for an executable [`PyTorch Job`](https://github.com/kubeflow/pytorch-operator).

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.kfpytorch import PyTorch

enumRestartPolicy

RestartPolicy describes how the replicas should be restarted.

extends Enum — enumeration of predefined options

from flytekitplugins.kfpytorch import RestartPolicy

configRunPolicy

RunPolicy describes some policy to apply to the execution of a kubeflow job.

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.kfpytorch import RunPolicy

configWorker

Configuration type for Kubeflow PyTorch.

extends dataclass — configuration or data structure for plugin setup

from flytekitplugins.kfpytorch import Worker

Dependencies

cloudpicklekubernetes

Related Plugins

Kubeflow TensorFlow

ML Training

This plugin uses the Kubeflow TensorFlow Operator and provides an extremely simplified interface for executing distributed training using various TensorFlow backends.

pytorch

ML Training

Union can execute PyTorch distributed training jobs natively on a Kubernetes Cluster, which manages the lifecycle of worker pods, rendezvous coordination, spin-up, and tear down. It leverages the open-sourced TorchElastic (torch.distributed.elastic) launcher and the Kubeflow PyTorch Operator, enabling fault-tolerant and elastic training across multiple nodes.