SageMaker Inference Agent#

This guide provides an overview of how to set up the SageMaker inference agent in your Flyte deployment.

Specify agent configuration#

Edit the relevant YAML file to specify the agent.

kubectl edit configmap flyte-sandbox-config -n flyte

tasks:
  task-plugins:
    enabled-plugins:
      - container
      - sidecar
      - k8s-array
      - agent-service
    default-for-task-types:
      - container: container
      - container_array: k8s-array
      - boto: agent-service
      - sagemaker-endpoint: agent-service

Create a file named values-override.yaml and add the following configuration to it:

configmap:
  enabled_plugins:
    tasks:
      task-plugins:
        enabled-plugins:
          - container
          - sidecar
          - k8s-array
          - agent-service
        default-for-task-types:
          container: container
          sidecar: sidecar
          container_array: k8s-array
          boto: agent-service
          sagemaker-endpoint: agent-service

AWS credentials#

When running the code locally, you can set AWS credentials as environment variables. When running on a production AWS cluster, the IAM role is used by default. Ensure that it has the AmazonSageMakerFullAccess policy attached.

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.

helm upgrade <RELEASE_NAME> flyte/flyte-core -n <YOUR_NAMESPACE> --values values-override.yaml

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

You can refer to the documentation here.