sagemaker agent backend setup doc

Signed-off-by: Samhita Alla <[email protected]>

docs/_ext/import_projects.py

@@ -142,7 +142,12 @@ def import_projects(app: Sphinx, config: Config):
         else:
             local_dir = srcdir / import_projects_config.clone_dir / project.dest
             shutil.rmtree(local_dir, ignore_errors=True)
-            repo = Repo.clone_from(project.source, local_dir)
+            if "flytekit" in project.source:
+                repo = Repo.clone_from(project.source, local_dir, branch="sagemaker-agent")
+            elif "flytesnacks" in project.source:
+                repo = Repo.clone_from(project.source, local_dir, branch="add-sagemaker-agent")
+            else:
+                repo = Repo.clone_from(project.source, local_dir)
             show_repo_tags = True
 
         local_docs_path = local_dir / project.docs_path

docs/core_use_cases/machine_learning.md

@@ -108,9 +108,8 @@ There are many ways to extend your workloads:
     [Ray Tune](https://docs.ray.io/en/latest/tune/index.html) for hyperparameter
     optimization, all orchestrated by Flyte as ephemerally-provisioned Ray clusters.
 * - **📦 Ephemeral Cluster Resources**
-  - Use the {ref}`MPI Operator <kf-mpi-op>`, {ref}`Sagemaker <aws-sagemaker>`,
-    {ref}`Kubeflow Tensorflow <kftensorflow-plugin>`, {ref}`Kubeflow Pytorch<kf-pytorch-op>`
-    and {doc}`more <_tags/DistributedComputing>` to do distributed training.
+  - Use the {ref}`MPI Operator <kf-mpi-op>`, {ref}`Kubeflow Tensorflow <kftensorflow-plugin>`,
+    {ref}`Kubeflow Pytorch<kf-pytorch-op>` and {doc}`more <_tags/DistributedComputing>` to do distributed training.
 * - **🔎 Experiment Tracking**
   - Auto-capture training logs with the {py:func}`~flytekitplugins.mlflow.mlflow_autolog`
     decorator, which can be viewed as Flyte Decks with `@task(disable_decks=False)`.

docs/deployment/agents/index.md

@@ -27,6 +27,8 @@ If you are using a managed deployment of Flyte, you will need to contact your de
   - Configuring your Flyte deployment for the MMCloud agent.
 * - {ref}`Sensor Agent <deployment-agent-setup-sensor>`
   - Configuring your Flyte deployment for the sensor agent.
+* - {ref}`SageMaker Inference <deployment-agent-setup-sagemaker-inference>`
+  - Deploy models and create, as well as trigger inference endpoints on SageMaker.
 ```
 
 ```{toctree}
@@ -39,6 +41,7 @@ chatgpt
 databricks
 bigquery
 mmcloud
+sagemaker_inference
 sensor
 snowflake
 ```

docs/deployment/agents/sagemaker_inference.rst

@@ -0,0 +1,126 @@
+.. _deployment-agent-setup-sagemaker-inference:
+
+SageMaker Inference Agent
+=========================
+
+This guide provides an overview of how to set up the SageMaker inference agent in your Flyte deployment.
+
+Specify agent configuration
+---------------------------
+
+.. tabs::
+
+    .. group-tab:: Flyte binary
+
+      Edit the relevant YAML file to specify the agent.
+
+      .. code-block:: bash
+
+        kubectl edit configmap flyte-sandbox-config -n flyte
+
+      .. code-block:: yaml
+        :emphasize-lines: 7,11-12,16-17
+
+        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
+        plugins:
+          agent-service:
+            supportedTaskTypes:
+            - boto
+            - sagemaker-endpoint
+
+    .. group-tab:: Flyte core
+
+      Create a file named ``values-override.yaml`` and add the following configuration to it:
+
+      .. code-block:: yaml
+        :emphasize-lines: 9,14-15,19-20
+
+        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
+            plugins:
+              agent-service:
+                supportedTaskTypes:
+                - boto
+                - sagemaker-endpoint
+
+Add the AWS credentials
+-----------------------
+
+1. Install the flyteagent pod using helm:
+
+.. code-block::
+
+  helm repo add flyteorg https://flyteorg.github.io/flyte
+  helm install flyteagent flyteorg/flyteagent --namespace flyte
+
+2. Get the base64 value of your AWS credentials:
+
+.. code-block::
+  
+  echo -n "<AWS_CREDENTIAL>" | base64
+
+3. Edit the flyteagent secret:
+
+.. code-block:: bash
+  
+  kubectl edit secret flyteagent -n flyte
+
+.. code-block:: yaml
+  :emphasize-lines: 3-5
+
+  apiVersion: v1
+  data:
+    aws-access-key: <BASE64_ENCODED_AWS_ACCESS_KEY>
+    aws-secret-access-key: <BASE64_ENCODED_AWS_SECRET_ACCESS_KEY>
+    aws-session-token: <BASE64_ENCODED_AWS_SESSION_TOKEN>
+  kind: Secret
+
+Upgrade the Flyte Helm release
+------------------------------
+
+.. tabs::
+
+  .. group-tab:: Flyte binary
+
+    .. code-block:: bash
+
+      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.
+
+  .. group-tab:: Flyte core
+
+    .. code-block:: bash
+
+      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 <https://docs.flyte.org/en/latest/flytesnacks/examples/sagemaker_inference_agent/index.html>`__.

docs/flyte_fundamentals/extending_flyte.md

@@ -151,7 +151,7 @@ many more ways to customize Flyte tasks:
     in other languages outside of the `flytekit` SDK language.
 * - {ref}`Backend Plugins <extend-plugin-flyte-backend>`
   - These tasks plugins require implementing a backend plugin to leverage
-    external services like Sagemaker, Snowflake, BigQuery, etc.
+    external services like SageMaker, Snowflake, BigQuery, etc.
 ```
 
 ## What's next?

docs/flyte_fundamentals/optimizing_tasks.md

@@ -242,8 +242,7 @@ at the most granular level of your workflow!
 When this task is executed on a Flyte cluster, it automatically provisions all of
 the resources that you need. In this case, that need is distributed
 training, but Flyte also provides integrations for {ref}`Spark <plugins-spark-k8s>`,
-{ref}`Ray <kube-ray-op>`, {ref}`MPI <kf-mpi-op>`, {ref}`Sagemaker <aws-sagemaker>`,
-{ref}`Snowflake <snowflake_agent>`, and more.
+{ref}`Ray <kube-ray-op>`, {ref}`MPI <kf-mpi-op>`, {ref}`Snowflake <snowflake_agent>`, and more.
 
 Even though Flyte itself is a powerful compute engine and orchestrator for
 data engineering, machine learning, and analytics, perhaps you have existing