Configure Kubernetes Plugins

Tags: Kubernetes, Integration, Spark, AWS, GCP, Advanced

This guide provides an overview of setting up the Kubernetes Operator backend plugin in your Flyte deployment.

Spin up a cluster

Flyte binaryFlyte core

Demo clusterHelm chart

PyTorchTensorFlowMPIRaySparkDask

Enable the PyTorch plugin on the demo cluster by adding the following block to ~/.flyte/sandbox/config.yaml:

tasks:
  task-plugins:
    default-for-task-types:
      container: container
      container_array: k8s-array
      sidecar: sidecar
      pytorch: pytorch
    enabled-plugins:
    - container
    - k8s-array
    - sidecar
    - pytorch

Start the demo cluster by running the following command:

flytectl demo start

Note

Add the Flyte chart repo to Helm if you’re installing via the Helm charts.

helm repo add flyteorg https://flyteorg.github.io/flyte

Install the Kubernetes operator

PyTorch/TensorFlow/MPIRaySparkDask

First, install kustomize.

Build and apply the training-operator.

export KUBECONFIG=$KUBECONFIG:~/.kube/config:~/.flyte/k3s/k3s.yaml
kustomize build "https://github.com/kubeflow/training-operator.git/manifests/overlays/standalone?ref=v1.5.0" | kubectl apply -f -

Optional: Using a gang scheduler

To address potential issues with worker pods of distributed training jobs being scheduled at different times due to resource constraints, you can opt for a gang scheduler. This ensures that all worker pods are scheduled simultaneously, reducing the likelihood of job failures caused by timeout errors.

To enable gang scheduling for the Kubeflow training-operator, you can install the Kubernetes scheduler plugins or the Apache YuniKorn scheduler.

1. Install the scheduler plugin or Apache YuniKorn as a second scheduler.

2. Configure the Kubeflow training-operator to use the new scheduler:

   Create a manifest called kustomization.yaml with the following content:

   apiVersion: kustomize.config.k8s.io/v1beta1
   kind: Kustomization
   
   resources:
   - github.com/kubeflow/training-operator/manifests/overlays/standalone
   
   patchesStrategicMerge:
   - patch.yaml
   

   Create a patch file called patch.yaml with the following content:

   apiVersion: apps/v1
   kind: Deployment
   metadata:
     name: training-operator
   spec:
     template:
       spec:
         containers:
         - name: training-operator
           command:
           - /manager
           - --gang-scheduler-name=<scheduler-plugins/yunikorn>
   

   Install the patched kustomization with the following command:

   kustomize build path/to/overlay/directory | kubectl apply -f -
   

   (Only for Apache YuniKorn) To configure gang scheduling with Apache YuniKorn, make sure to set the following annotations in Flyte pod templates:

   • template.metadata.annotations.yunikorn.apache.org/task-group-name

   • template.metadata.annotations.yunikorn.apache.org/task-groups

   • template.metadata.annotations.yunikorn.apache.org/schedulingPolicyParameters

   For more configuration details, refer to the Apache YuniKorn Gang-Scheduling documentation.

3. Use a Flyte pod template with template.spec.schedulerName: scheduler-plugins-scheduler to use the new gang scheduler for your tasks.

   See the Using K8s PodTemplates section for more information on pod templates in Flyte. You can set the scheduler name in the pod template passed to the @task decorator. However, to prevent the two different schedulers from competing for resources, it is recommended to set the scheduler name in the pod template in the flyte namespace which is applied to all tasks. Non distributed training tasks can be scheduled by the gang scheduler as well.

   For more information on pod templates in Flyte, refer to the Using K8s PodTemplates section. You can set the scheduler name in the pod template passed to the @task decorator. However, to avoid resource competition between the two different schedulers, it is recommended to set the scheduler name in the pod template in the flyte namespace, which is applied to all tasks. This allows non-distributed training tasks to be scheduled by the gang scheduler as well.

Specify plugin configuration

PyTorchTensorFlowMPIRaySparkDask

Flyte binaryFlyte core

To specify the plugin when using the Helm chart, edit the relevant YAML file.

tasks:
  task-plugins:
    enabled-plugins:
      - container
      - sidecar
      - k8s-array
      - pytorch
    default-for-task-types:
      - container: container
      - container_array: k8s-array
      - pytorch: pytorch

Upgrade the deployment

Flyte binaryFlyte core

If you are installing Flyte via the Helm chart, run the following command:

Note

There is no need to run helm upgrade for Spark.

helm upgrade <RELEASE_NAME> flyteorg/flyte-binary -n <YOUR_NAMESPACE> --values <YOUR_YAML_FILE>

Replace <RELEASE_NAME> with the name of your release (e.g., flyte-backend), <YOUR_NAMESPACE> with the name of your namespace (e.g., flyte), and <YOUR_YAML_FILE> with the name of your YAML file.

Wait for the upgrade to complete. You can check the status of the deployment pods by running the following command:

kubectl get pods -n --all-namespaces