Sagemaker Plugin Setup#

This guide gives an overview of how to set up Sagemaker in your Flyte deployment. Sagemaker plugin needs Flyte deployment in AWS cloud; sandbox/GCP/Azure wouldn’t work.

  1. Setup the AWS Flyte cluster

  1. Create a file named values-override.yaml and add the following config to it. Please make sure that the propeller has the correct service account for Sagemaker.

    # -- Tasks specific configuration [structure](
      # -- Plugins configuration, [structure](
        # -- [Enabled Plugins](
        # plugins
          - container
          - sidecar
          - k8s-array
          - sagemaker_training
          - sagemaker_hyperparameter_tuning
          container: container
          sidecar: sidecar
          container_array: k8s-array
  1. Upgrade the Flyte Helm release.

helm upgrade -n flyte -f values-override.yaml flyteorg/flyte-core
  1. Register the Sagemaker plugin example.

flytectl register files --archive -p flytesnacks -d development
  1. Launch an execution

  • Navigate to Flyte Console’s UI (e.g. sandbox) and find the workflow.

  • Click on Launch to open up the launch form.

  • Submit the form.

  • Retrieve an execution form in the form of a YAML file:

    flytectl get launchplan --config ~/.flyte/flytectl.yaml --project flytesnacks --domain development sagemaker_training.sagemaker_custom_training.mnist_trainer --latest --execFile exec_spec.yaml
  • Launch! 🚀

    flytectl --config ~/.flyte/flytectl.yaml create execution -p <project> -d <domain> --execFile ~/exec_spec.yaml