Add managed SageMaker job connectors

Expand the existing SageMaker plugin with async lifecycle support for training, processing, batch transform, hyperparameter tuning, and Inference Recommender jobs. Include typed task wrappers, stable projected outputs, idempotent retries, public documentation, and comprehensive unit tests without requiring Propeller changes.

Signed-off-by: Rohit Sharma <rohitrsh@gmail.com>

Author: rohitrsh

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/README.md

@@ -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/__init__.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/connector.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_batch_transform/task.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",
+]