TensorFlow Distributed#

Tags: Integration, DistributedComputing, MachineLearning, KubernetesOperator, Advanced

TensorFlow operator is useful to natively run distributed TensorFlow training jobs on Flyte. It leverages the Kubeflow training operator.

Install the plugin#

To install the Kubeflow TensorFlow plugin, run the following command:

pip install flytekitplugins-kftensorflow

To enable the plugin in the backend, follow instructions outlined in the Configure Kubernetes Plugins guide.

Run the example on the Flyte cluster#

To run the provided example on the Flyte cluster, use the following command:

pyflyte run --remote \
  https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/kftensorflow_plugin/kftensorflow_plugin/tf_mnist.py \
  mnist_tensorflow_workflow