Sagemaker inference agent

.github/workflows/pythonbuild.yml

@@ -2,14 +2,13 @@ name: Build
 
 # Schedule runs to run twice a day
 
-
 on:
   push:
     branches:
       - master
   pull_request:
   schedule:
-      - cron: '0 13 * * *' # This schedule runs at 1pm UTC every day
+    - cron: "0 13 * * *" # This schedule runs at 1pm UTC every day
 
 env:
   FLYTE_SDK_LOGGING_LEVEL: 10 # debug
@@ -116,9 +115,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-latest ]
+        os: [ubuntu-latest]
         python-version: ${{fromJson(needs.detect-python-versions.outputs.python-versions)}}
-        pandas: [ "pandas<2.0.0", "pandas>=2.0.0" ]
+        pandas: ["pandas<2.0.0", "pandas>=2.0.0"]
     steps:
       - uses: insightsengineering/disk-space-reclaimer@v1
       - uses: actions/checkout@v4
@@ -257,8 +256,7 @@ jobs:
           - flytekit-async-fsspec
           - flytekit-aws-athena
           - flytekit-aws-batch
-          # TODO: uncomment this when the sagemaker agent is implemented: https://github.com/flyteorg/flyte/issues/4079
-          # - flytekit-aws-sagemaker
+          - flytekit-aws-sagemaker
           - flytekit-bigquery
           - flytekit-dask
           - flytekit-data-fsspec

.pre-commit-config.yaml

@@ -1,11 +1,11 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.6
+    rev: v0.2.2
     hooks:
       # Run the linter.
       - id: ruff
-        args: [--fix, --show-fixes, --show-source]
+        args: [--fix, --show-fixes, --output-format=full]
       # Run the formatter.
       - id: ruff-format
   - repo: https://github.com/pre-commit/pre-commit-hooks

Dockerfile.agent

@@ -13,6 +13,7 @@ RUN pip install --no-cache-dir -U flytekit==$VERSION \
   flytekitplugins-bigquery==$VERSION \
   flytekitplugins-chatgpt==$VERSION \
   flytekitplugins-snowflake==$VERSION \
+  flytekitplugins-awssagemaker==$VERSION \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \
   && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

docs/source/plugins/awssagemaker_inference.rst

@@ -0,0 +1,12 @@
+.. _awssagemaker_inference:
+
+#####################################
+AWS Sagemaker Inference API reference
+#####################################
+
+.. tags:: Integration, MachineLearning, AWS
+
+.. automodule:: flytekitplugins.awssagemaker_inference
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:

docs/source/plugins/index.rst

@@ -30,6 +30,7 @@ Plugin API reference
 * :ref:`Vaex <vaex>` - Vaex API reference
 * :ref:`MLflow <mlflow>` - MLflow API reference
 * :ref:`DuckDB <duckdb>` - DuckDB API reference
+* :ref:`SageMaker Inference <awssagemaker_inference>` - SageMaker Inference API reference
 
 .. toctree::
    :maxdepth: 2
@@ -61,3 +62,4 @@ Plugin API reference
    Vaex <vaex>
    MLflow <mlflow>
    DuckDB <duckdb>
+   SageMaker Inference <awssagemaker_inference>

flytekit/core/base_task.py

@@ -24,7 +24,20 @@
 from abc import abstractmethod
 from base64 import b64encode
 from dataclasses import dataclass
-from typing import Any, Coroutine, Dict, Generic, List, Optional, OrderedDict, Tuple, Type, TypeVar, Union, cast
+from typing import (
+    Any,
+    Coroutine,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    OrderedDict,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from flyteidl.core import artifact_id_pb2 as art_id
 from flyteidl.core import tasks_pb2
@@ -165,9 +178,6 @@ class IgnoreOutputs(Exception):
     """
     This exception should be used to indicate that the outputs generated by this can be safely ignored.
     This is useful in case of distributed training or peer-to-peer parallel algorithms.
-
-    For example look at Sagemaker training, e.g.
-    :py:class:`plugins.awssagemaker.flytekitplugins.awssagemaker.training.SagemakerBuiltinAlgorithmsTask`.
     """
 
     pass
@@ -479,7 +489,10 @@ def __init__(
         self._task_config = task_config
 
         if disable_deck is not None:
-            warnings.warn("disable_deck was deprecated in 1.10.0, please use enable_deck instead", FutureWarning)
+            warnings.warn(
+                "disable_deck was deprecated in 1.10.0, please use enable_deck instead",
+                FutureWarning,
+            )
 
         # Confirm that disable_deck and enable_deck do not contradict each other
         if disable_deck is not None and enable_deck is not None:
@@ -703,7 +716,13 @@ def dispatch_execute(
                         # If executed inside of a workflow being executed locally, then run the coroutine to get the
                         # actual results.
                         return asyncio.run(
-                            self._async_execute(native_inputs, native_outputs, ctx, exec_ctx, new_user_params)
+                            self._async_execute(
+                                native_inputs,
+                                native_outputs,
+                                ctx,
+                                exec_ctx,
+                                new_user_params,
+                            )
                         )
 
                 return self._async_execute(native_inputs, native_outputs, ctx, exec_ctx, new_user_params)
@@ -716,7 +735,10 @@ def dispatch_execute(
             # Short circuit the translation to literal map because what's returned may be a dj spec (or an
             # already-constructed LiteralMap if the dynamic task was a no-op), not python native values
             # dynamic_execute returns a literal map in local execute so this also gets triggered.
-            if isinstance(native_outputs, (_literal_models.LiteralMap, _dynamic_job.DynamicJobSpec)):
+            if isinstance(
+                native_outputs,
+                (_literal_models.LiteralMap, _dynamic_job.DynamicJobSpec),
+            ):
                 return native_outputs
 
             literals_map, native_outputs_as_map = self._output_to_literal_map(native_outputs, exec_ctx)

flytekit/core/condition.py

@@ -428,7 +428,7 @@ def transform_to_comp_expr(expr: ComparisonExpression) -> Tuple[_core_cond.Compa
 
 
 def transform_to_boolexpr(
-    expr: Union[ComparisonExpression, ConjunctionExpression]
+    expr: Union[ComparisonExpression, ConjunctionExpression],
 ) -> Tuple[_core_cond.BooleanExpression, typing.List[Promise]]:
     if isinstance(expr, ConjunctionExpression):
         cexpr, promises = transform_to_conj_expr(expr)

flytekit/core/task.py

@@ -31,8 +31,7 @@ class TaskPlugins(object):
         # Plugin_object_type is a derivative of ``PythonFunctionTask``
 
     Examples of available task plugins include different query-based plugins such as
-    :py:class:`flytekitplugins.athena.task.AthenaTask` and :py:class:`flytekitplugins.hive.task.HiveTask`, ML tools like
-    :py:class:`plugins.awssagemaker.flytekitplugins.awssagemaker.training.SagemakerBuiltinAlgorithmsTask`, kubeflow
+    :py:class:`flytekitplugins.athena.task.AthenaTask` and :py:class:`flytekitplugins.hive.task.HiveTask`, kubeflow
     operators like :py:class:`plugins.kfpytorch.flytekitplugins.kfpytorch.task.PyTorchFunctionTask` and
     :py:class:`plugins.kftensorflow.flytekitplugins.kftensorflow.task.TensorflowFunctionTask`, and generic plugins like
     :py:class:`flytekitplugins.pod.task.PodFunctionTask` which doesn't integrate with third party tools or services.
@@ -103,7 +102,13 @@ def task(
     secret_requests: Optional[List[Secret]] = ...,
     execution_mode: PythonFunctionTask.ExecutionBehavior = ...,
     node_dependency_hints: Optional[
-        Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]]
+        Iterable[
+            Union[
+                PythonFunctionTask,
+                _annotated_launchplan.LaunchPlan,
+                _annotated_workflow.WorkflowBase,
+            ]
+        ]
     ] = ...,
     task_resolver: Optional[TaskResolverMixin] = ...,
     docs: Optional[Documentation] = ...,
@@ -135,7 +140,13 @@ def task(
     secret_requests: Optional[List[Secret]] = ...,
     execution_mode: PythonFunctionTask.ExecutionBehavior = ...,
     node_dependency_hints: Optional[
-        Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]]
+        Iterable[
+            Union[
+                PythonFunctionTask,
+                _annotated_launchplan.LaunchPlan,
+                _annotated_workflow.WorkflowBase,
+            ]
+        ]
     ] = ...,
     task_resolver: Optional[TaskResolverMixin] = ...,
     docs: Optional[Documentation] = ...,
@@ -166,7 +177,13 @@ def task(
     secret_requests: Optional[List[Secret]] = None,
     execution_mode: PythonFunctionTask.ExecutionBehavior = PythonFunctionTask.ExecutionBehavior.DEFAULT,
     node_dependency_hints: Optional[
-        Iterable[Union[PythonFunctionTask, _annotated_launchplan.LaunchPlan, _annotated_workflow.WorkflowBase]]
+        Iterable[
+            Union[
+                PythonFunctionTask,
+                _annotated_launchplan.LaunchPlan,
+                _annotated_workflow.WorkflowBase,
+            ]
+        ]
     ] = None,
     task_resolver: Optional[TaskResolverMixin] = None,
     docs: Optional[Documentation] = None,
@@ -175,7 +192,11 @@ def task(
     pod_template: Optional["PodTemplate"] = None,
     pod_template_name: Optional[str] = None,
     accelerator: Optional[BaseAccelerator] = None,
-) -> Union[Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]], PythonFunctionTask[T], Callable[..., FuncOut]]:
+) -> Union[
+    Callable[[Callable[..., FuncOut]], PythonFunctionTask[T]],
+    PythonFunctionTask[T],
+    Callable[..., FuncOut],
+]:
     """
     This is the core decorator to use for any task type in flytekit.
 
@@ -342,7 +363,13 @@ class ReferenceTask(ReferenceEntity, PythonFunctionTask):  # type: ignore
     """
 
     def __init__(
-        self, project: str, domain: str, name: str, version: str, inputs: Dict[str, type], outputs: Dict[str, Type]
+        self,
+        project: str,
+        domain: str,
+        name: str,
+        version: str,
+        inputs: Dict[str, type],
+        outputs: Dict[str, Type],
     ):
         super().__init__(TaskReference(project, domain, name, version), inputs, outputs)