Add Pythonic SageMaker training and processing

flytekit/extend/backend/utils.py

@@ -37,7 +37,12 @@ def is_terminal_phase(phase: TaskExecution.Phase) -> bool:
     """
     Return true if the phase is terminal.
     """
-    return phase in [TaskExecution.SUCCEEDED, TaskExecution.ABORTED, TaskExecution.FAILED]
+    return phase in [
+        TaskExecution.SUCCEEDED,
+        TaskExecution.ABORTED,
+        TaskExecution.FAILED,
+        TaskExecution.RETRYABLE_FAILED,
+    ]
 
 
 def get_connector_secret(secret_key: str) -> str:

plugins/README.md

@@ -6,7 +6,7 @@ All the Flytekit plugins maintained by the core team are added here. It is not n
 
 | Plugin                       | Installation                                          | Description                                                                                                                | Version                                                                                                                                                      | Type          |
 | ---------------------------- | ----------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------- |
-| AWS SageMaker                | `bash pip install flytekitplugins-awssagemaker`       | Deploy SageMaker models and manage inference endpoints with ease.                                                          | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-awssagemaker.svg)](https://pypi.python.org/pypi/flytekitplugins-awssagemaker/)             | Flytekit-only |
+| AWS SageMaker                | `bash pip install flytekitplugins-awssagemaker`       | Run SageMaker training, processing, tuning, transform, recommendation, and deployment workloads.                           | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-awssagemaker.svg)](https://pypi.python.org/pypi/flytekitplugins-awssagemaker/)             | Flytekit-only |
 | dask                         | `bash pip install flytekitplugins-dask `              | Installs SDK to author dask jobs that can be executed natively on Kubernetes using the Flyte backend plugin                | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-dask.svg)](https://pypi.python.org/pypi/flytekitplugins-dask/)                             | Backend       |
 | Hive Queries                 | `bash pip install flytekitplugins-hive `              | Installs SDK to author Hive Queries that can be executed on a configured hive backend using Flyte backend plugin           | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-hive.svg)](https://pypi.python.org/pypi/flytekitplugins-hive/)                             | Backend       |
 | K8s distributed PyTorch Jobs | `bash pip install flytekitplugins-kfpytorch `         | Installs SDK to author Distributed pyTorch Jobs in python using Kubeflow PyTorch Operator                                  | [![PyPI version fury.io](https://badge.fury.io/py/flytekitplugins-kfpytorch.svg)](https://pypi.python.org/pypi/flytekitplugins-kfpytorch/)                   | Backend       |

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/__init__.py

@@ -0,0 +1,27 @@
+"""
+.. currentmodule:: flytekitplugins.awssagemaker_batch_transform
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   SageMakerTransformJobConnector
+   SageMakerTransformJobTask
+   SageMakerStopTransformJobTask
+   SageMakerDescribeTransformJobTask
+"""
+
+from .connector import SageMakerTransformJobConnector, SageMakerTransformJobMetadata
+from .task import (
+    SageMakerDescribeTransformJobTask,
+    SageMakerStopTransformJobTask,
+    SageMakerTransformJobTask,
+)
+
+__all__ = [
+    "SageMakerTransformJobConnector",
+    "SageMakerTransformJobMetadata",
+    "SageMakerTransformJobTask",
+    "SageMakerStopTransformJobTask",
+    "SageMakerDescribeTransformJobTask",
+]

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/connector.py

@@ -0,0 +1,157 @@
+"""SageMaker batch-transform connector.
+
+Mirrors the training-job connector. Targets ``CreateTransformJob`` /
+``DescribeTransformJob`` / ``StopTransformJob``. Surfaces the predictions
+``S3OutputPath`` so downstream Flyte tasks can read scores written by SageMaker
+without any extra plumbing.
+
+Note: ``TransformJobStatus`` has no ``Deleting`` state and there is no
+``SecondaryStatus`` — running phase has no live progress signal beyond the job
+being in flight.
+"""
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import cloudpickle
+from flyteidl.core.execution_pb2 import TaskExecution
+from flytekitplugins.awssagemaker_inference.boto3_mixin import (
+    Boto3ConnectorMixin,
+    CustomException,
+)
+
+from flytekit.extend.backend.base_connector import (
+    AsyncConnectorBase,
+    ConnectorRegistry,
+    Resource,
+    ResourceMeta,
+)
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+
+@dataclass
+class SageMakerTransformJobMetadata(ResourceMeta):
+    config: Dict[str, Any]
+    region: Optional[str] = None
+    inputs: Optional[LiteralMap] = None
+
+    def encode(self) -> bytes:
+        return cloudpickle.dumps(self)
+
+    @classmethod
+    def decode(cls, data: bytes) -> "SageMakerTransformJobMetadata":
+        return cloudpickle.loads(data)
+
+
+_STATE_MAP = {
+    "InProgress": TaskExecution.RUNNING,
+    "Stopping": TaskExecution.RUNNING,
+    "Completed": TaskExecution.SUCCEEDED,
+    "Failed": TaskExecution.FAILED,
+    "Stopped": TaskExecution.FAILED,
+}
+
+
+def _isoformat(value: Any) -> Any:
+    if isinstance(value, datetime):
+        return value.isoformat()
+    return value
+
+
+def _build_outputs(describe_response: Dict[str, Any]) -> Dict[str, Any]:
+    """Project describe_transform_job down to a stable, downstream-friendly dict."""
+    transform_output = describe_response.get("TransformOutput") or {}
+    return {
+        "TransformJobArn": describe_response.get("TransformJobArn"),
+        "TransformJobName": describe_response.get("TransformJobName"),
+        "ModelName": describe_response.get("ModelName"),
+        "TransformOutput": {"S3OutputPath": transform_output.get("S3OutputPath")},
+        "TransformStartTime": _isoformat(describe_response.get("TransformStartTime")),
+        "TransformEndTime": _isoformat(describe_response.get("TransformEndTime")),
+    }
+
+
+class SageMakerTransformJobConnector(Boto3ConnectorMixin, AsyncConnectorBase):
+    """Long-running connector for SageMaker batch-transform jobs."""
+
+    name = "SageMaker Transform Job Connector"
+
+    def __init__(self):
+        super().__init__(
+            service="sagemaker",
+            task_type_name="sagemaker-transform-job",
+            metadata_type=SageMakerTransformJobMetadata,
+        )
+
+    async def create(
+        self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs
+    ) -> SageMakerTransformJobMetadata:
+        custom = task_template.custom
+        config = custom.get("config")
+        region = custom.get("region")
+
+        try:
+            await self._call(
+                method="create_transform_job",
+                config=config,
+                inputs=inputs,
+                region=region,
+            )
+        except CustomException as e:
+            original_exception = e.original_exception
+            error_code = original_exception.response["Error"]["Code"]
+            error_message = original_exception.response["Error"]["Message"]
+
+            if e.idempotence_token and (
+                error_code == "ResourceInUse"
+                or (error_code == "ValidationException" and "Cannot create already existing" in error_message)
+            ):
+                return SageMakerTransformJobMetadata(config=config, region=region, inputs=inputs)
+            raise e
+
+        return SageMakerTransformJobMetadata(config=config, region=region, inputs=inputs)
+
+    async def get(self, resource_meta: SageMakerTransformJobMetadata, **kwargs) -> Resource:
+        describe_response, _ = await self._call(
+            method="describe_transform_job",
+            config={"TransformJobName": resource_meta.config.get("TransformJobName")},
+            inputs=resource_meta.inputs,
+            region=resource_meta.region,
+        )
+
+        current_state = describe_response.get("TransformJobStatus")
+        flyte_phase = _STATE_MAP.get(current_state, TaskExecution.RUNNING)
+
+        message: Optional[str] = None
+        if current_state in ("Failed", "Stopped"):
+            message = describe_response.get("FailureReason")
+
+        outputs: Optional[Dict[str, Any]] = None
+        if current_state == "Completed":
+            outputs = {"result": _build_outputs(describe_response)}
+
+        return Resource(phase=flyte_phase, outputs=outputs, message=message)
+
+    async def delete(self, resource_meta: SageMakerTransformJobMetadata, **kwargs):
+        try:
+            await self._call(
+                method="stop_transform_job",
+                config={"TransformJobName": resource_meta.config.get("TransformJobName")},
+                region=resource_meta.region,
+                inputs=resource_meta.inputs,
+            )
+        except CustomException as e:
+            original_exception = e.original_exception
+            error_code = original_exception.response["Error"]["Code"]
+            error_message = original_exception.response["Error"]["Message"]
+
+            if error_code == "ResourceNotFound" or (
+                error_code == "ValidationException" and "non-running" in error_message
+            ):
+                return
+            raise e
+
+
+ConnectorRegistry.register(SageMakerTransformJobConnector())

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/task.py

@@ -0,0 +1,101 @@
+"""User-facing tasks for SageMaker batch-transform jobs."""
+
+from typing import Any, Dict, Optional, Type
+
+from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask
+
+from flytekit import kwtypes
+from flytekit.configuration import SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.core.interface import Interface
+from flytekit.extend.backend.base_connector import AsyncConnectorExecutorMixin
+
+
+class SageMakerTransformJobTask(AsyncConnectorExecutorMixin, PythonTask):
+    """Run a SageMaker batch-transform job and emit the predictions ``S3OutputPath``.
+
+    Outputs a single ``result: dict`` literal containing ``TransformJobArn``,
+    ``TransformJobName``, ``ModelName``, ``TransformOutput.S3OutputPath`` (the S3
+    prefix where SageMaker wrote one ``<input>.out`` per input object — feed this
+    into a downstream Flyte task to consume the predictions), ``TransformStartTime``
+    and ``TransformEndTime``.
+
+    Set ``DataProcessing.JoinSource: "Input"`` in the config for tabular predictive
+    workloads so each output line carries the original input fields alongside the
+    prediction (otherwise rows have no key to join back).
+
+    ``name`` identifies the Flyte task. ``config`` is the boto3
+    ``create_transform_job`` request and may contain ``{inputs.X}`` and
+    ``{idempotence_token}`` placeholders. ``region`` selects the AWS region, and
+    ``inputs`` maps input placeholders to Flyte types.
+    """
+
+    _TASK_TYPE = "sagemaker-transform-job"
+
+    def __init__(
+        self,
+        name: str,
+        config: Dict[str, Any],
+        region: Optional[str] = None,
+        inputs: Optional[Dict[str, Type]] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            task_type=self._TASK_TYPE,
+            interface=Interface(inputs=inputs, outputs=kwtypes(result=dict)),
+            **kwargs,
+        )
+        self._config = config
+        self._region = region
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {"config": self._config, "region": self._region}
+
+
+class SageMakerStopTransformJobTask(BotoTask):
+    """Sync helper task that stops a running SageMaker transform job by name."""
+
+    def __init__(
+        self,
+        name: str,
+        config: Dict[str, Any],
+        region: Optional[str] = None,
+        inputs: Optional[Dict[str, Type]] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            task_config=BotoConfig(
+                service="sagemaker",
+                method="stop_transform_job",
+                config=config,
+                region=region,
+            ),
+            inputs=inputs,
+            **kwargs,
+        )
+
+
+class SageMakerDescribeTransformJobTask(BotoTask):
+    """Sync helper task that returns the full ``describe_transform_job`` response."""
+
+    def __init__(
+        self,
+        name: str,
+        config: Dict[str, Any],
+        region: Optional[str] = None,
+        inputs: Optional[Dict[str, Type]] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            task_config=BotoConfig(
+                service="sagemaker",
+                method="describe_transform_job",
+                config=config,
+                region=region,
+            ),
+            inputs=inputs,
+            **kwargs,
+        )

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/__init__.py

@@ -0,0 +1,30 @@
+"""
+.. currentmodule:: flytekitplugins.awssagemaker_hyperparameter_tuning
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   SageMakerHyperParameterTuningJobConnector
+   SageMakerHyperParameterTuningJobTask
+   SageMakerStopHyperParameterTuningJobTask
+   SageMakerDescribeHyperParameterTuningJobTask
+"""
+
+from .connector import (
+    SageMakerHyperParameterTuningJobConnector,
+    SageMakerHyperParameterTuningJobMetadata,
+)
+from .task import (
+    SageMakerDescribeHyperParameterTuningJobTask,
+    SageMakerHyperParameterTuningJobTask,
+    SageMakerStopHyperParameterTuningJobTask,
+)
+
+__all__ = [
+    "SageMakerHyperParameterTuningJobConnector",
+    "SageMakerHyperParameterTuningJobMetadata",
+    "SageMakerHyperParameterTuningJobTask",
+    "SageMakerStopHyperParameterTuningJobTask",
+    "SageMakerDescribeHyperParameterTuningJobTask",
+]