Sagemaker Plugin#

This guide provides an overview of setting up Sagemaker in your Flyte deployment.

Note

The Sagemaker plugin requires Flyte deployment in the AWS cloud; it is not compatible with demo/GCP/Azure.

Set up AWS Flyte cluster#

  • Ensure you have a functional Flyte cluster running in AWS.

  • Verify that your AWS role is set up correctly for SageMaker.

  • Install the AWS SageMaker k8s operator in your Kubernetes cluster.

  • Confirm that you have the correct kubeconfig and have selected the appropriate Kubernetes context.

  • Verify the presence of the correct Flytectl configuration at ~/.flyte/config.yaml.

Specify the plugin configuration#

Edit the relevant YAML file to specify the plugin.

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

Ensure that the propeller has the correct service account for Sagemaker.

Upgrade the Flyte Helm release#

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.