From 7f9c53478b86ba31ad9492fea89c63dde3e47af9 Mon Sep 17 00:00:00 2001 From: Rohit Sharma Date: Tue, 4 Aug 2026 11:22:57 +0100 Subject: [PATCH] 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 --- plugins/README.md | 2 +- plugins/flytekit-aws-sagemaker/README.md | 528 ++++++++++++++++- .../awssagemaker_batch_transform/__init__.py | 27 + .../awssagemaker_batch_transform/connector.py | 157 +++++ .../awssagemaker_batch_transform/task.py | 101 ++++ .../__init__.py | 30 + .../connector.py | 246 ++++++++ .../task.py | 121 ++++ .../awssagemaker_inference/boto3_connector.py | 17 +- .../awssagemaker_inference/boto3_mixin.py | 23 +- .../__init__.py | 30 + .../connector.py | 232 ++++++++ .../task.py | 104 ++++ .../awssagemaker_processing/__init__.py | 27 + .../awssagemaker_processing/connector.py | 186 ++++++ .../awssagemaker_processing/task.py | 112 ++++ .../awssagemaker_training/__init__.py | 27 + .../awssagemaker_training/connector.py | 186 ++++++ .../awssagemaker_training/task.py | 108 ++++ plugins/flytekit-aws-sagemaker/setup.py | 25 +- .../tests/test_batch_transform_connector.py | 262 ++++++++ .../tests/test_batch_transform_task.py | 70 +++ .../tests/test_boto3_mixin.py | 96 +++ .../test_hyperparameter_tuning_connector.py | 560 ++++++++++++++++++ .../tests/test_hyperparameter_tuning_task.py | 95 +++ .../test_inference_recommender_connector.py | 346 +++++++++++ .../tests/test_inference_recommender_task.py | 72 +++ .../tests/test_processing_connector.py | 398 +++++++++++++ .../tests/test_processing_task.py | 87 +++ .../tests/test_training_connector.py | 370 ++++++++++++ .../tests/test_training_task.py | 74 +++ 31 files changed, 4698 insertions(+), 21 deletions(-) create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/__init__.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/connector.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/task.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/__init__.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/connector.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/task.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/__init__.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/connector.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/task.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/__init__.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/connector.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/task.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/__init__.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/connector.py create mode 100644 plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/task.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_batch_transform_connector.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_batch_transform_task.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_connector.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_task.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_connector.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_task.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_processing_connector.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_processing_task.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_training_connector.py create mode 100644 plugins/flytekit-aws-sagemaker/tests/test_training_task.py diff --git a/plugins/README.md b/plugins/README.md index acc7eec4d9..4ef41aaefd 100644 --- a/plugins/README.md +++ b/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 | diff --git a/plugins/flytekit-aws-sagemaker/README.md b/plugins/flytekit-aws-sagemaker/README.md index dd9e447eaa..18481191c2 100644 --- a/plugins/flytekit-aws-sagemaker/README.md +++ b/plugins/flytekit-aws-sagemaker/README.md @@ -1,6 +1,8 @@ # AWS SageMaker Plugin -The plugin currently features a SageMaker deployment connector. +The plugin features connectors for SageMaker deployment, model training, +processing, hyperparameter tuning, batch inference (a.k.a. batch transform), +and inference recommendations. ## Inference @@ -70,3 +72,527 @@ def model_deployment_workflow( instance_type="ml.m4.xlarge", ) ``` + +## Training + +`SageMakerTrainingJobTask` runs a `CreateTrainingJob` and waits for it to reach a +terminal state. The describe-poll loop runs server-side via the connector; no +Flyte worker holds a session open for the training duration. While running, the +task surfaces SageMaker's `SecondaryStatus` (`Starting`, `Downloading`, +`Training`, `Uploading`, …) as the live message. On success it emits a single +`result: dict` literal with: + +- `TrainingJobArn`, `TrainingJobName` +- `ModelArtifacts.S3ModelArtifacts` — the S3 URI of the trained `model.tar.gz` +- `OutputDataConfig.S3OutputPath` — sibling location for checkpoints / TensorBoard +- `FinalMetricDataList` — last value of every metric defined in `MetricDefinitions` +- `BillableTimeInSeconds`, `TrainingTimeInSeconds` + +```python +from flytekitplugins.awssagemaker_training import SageMakerTrainingJobTask +from flytekit import kwtypes, workflow + +training = SageMakerTrainingJobTask( + name="train-xgboost", + config={ + "TrainingJobName": "xgb-{idempotence_token}", + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + "MetricDefinitions": [ + {"Name": "validation:auc", "Regex": "auc=([0-9\\.]+)"}, + ], + }, + "RoleArn": "{inputs.execution_role_arn}", + "InputDataConfig": [ + { + "ChannelName": "train", + "DataSource": { + "S3DataSource": { + "S3DataType": "S3Prefix", + "S3Uri": "{inputs.train_data}", + "S3DataDistributionType": "FullyReplicated", + } + }, + } + ], + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "ResourceConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + region="", + images={"training_image": ""}, + inputs=kwtypes(execution_role_arn=str, train_data=str, output_prefix=str), +) +``` + +A training job writes `model.tar.gz` to S3 but does **not** create a SageMaker +`Model` entity. Chain a `SageMakerModelTask` downstream, feeding it +`result["ModelArtifacts"]["S3ModelArtifacts"]` as `PrimaryContainer.ModelDataUrl`, +to deploy the trained artefact via an endpoint or a batch-transform job. + +Registering the artifact with SageMaker Model Registry is a separate +`CreateModelPackage` operation and is not performed by this task. + +Inputs are S3-resident. To use a Glue/Athena-backed dataset, either pass the +underlying S3 location of the Glue table directly, or stage query results to S3 +with an upstream Flyte task and pass that S3 URI in. + +## Processing + +`SageMakerProcessingJobTask` runs a `CreateProcessingJob` and waits for it to +reach a terminal state, using the same server-side describe-poll loop as the +training task. Processing jobs cover the steps that bookend training — feature +engineering / data cleaning (pre-training), and model evaluation, batch scoring +with custom pre/post-processing, or SageMaker Clarify bias & explainability +(post-training) — on managed SageMaker infra under the same execution role. + +Unlike training, the container image lives at `AppSpecification.ImageUri`. +Inputs are commonly S3-resident, while `ProcessingOutputConfig` can write to S3 +or SageMaker Feature Store. Processing jobs expose no `SecondaryStatus`, so the +live message is empty while running; on failure the task surfaces +`FailureReason` (falling back to `ExitMessage`). On success it emits a single +`result: dict` literal with: + +- `ProcessingJobArn`, `ProcessingJobName` +- `Outputs` — a list projected from `ProcessingOutputConfig.Outputs`. Each + item contains `OutputName` plus either `S3Uri` for an S3 destination or + `FeatureGroupName` for a Feature Store destination. +- `ExitMessage`, `ProcessingStartTime`, `ProcessingEndTime` + +```python +from flytekitplugins.awssagemaker_processing import SageMakerProcessingJobTask +from flytekit import kwtypes + +preprocess = SageMakerProcessingJobTask( + name="preprocess-features", + config={ + "ProcessingJobName": "prep-{idempotence_token}", + "AppSpecification": { + "ImageUri": "{images.processing_image}", + "ContainerEntrypoint": ["python3", "/opt/ml/processing/preprocess.py"], + }, + "RoleArn": "{inputs.execution_role_arn}", + "ProcessingInputs": [ + { + "InputName": "raw", + "S3Input": { + "S3Uri": "{inputs.raw_data}", + "LocalPath": "/opt/ml/processing/input", + "S3DataType": "S3Prefix", + "S3InputMode": "File", + }, + } + ], + "ProcessingOutputConfig": { + "Outputs": [ + { + "OutputName": "train", + "S3Output": { + "S3Uri": "{inputs.output_prefix}", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + }, + } + ] + }, + "ProcessingResources": { + "ClusterConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + } + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + region="", + images={"processing_image": ""}, + inputs=kwtypes(execution_role_arn=str, raw_data=str, output_prefix=str), +) +``` + +Chain it before a `SageMakerTrainingJobTask` (feed an output's `S3Uri` in as the +training `InputDataConfig` S3 URI) or after one for evaluation. The +`SageMakerStopProcessingJobTask` / `SageMakerDescribeProcessingJobTask` helpers +mirror their training-job counterparts. + +## Hyperparameter Tuning + +`SageMakerHyperParameterTuningJobTask` runs `CreateHyperParameterTuningJob` and +waits for it to reach a terminal state. The polling loop is identical in shape +to `SageMakerTrainingJobTask`, but each trial is a child training job — so while +running, the task's message field surfaces a compact trial counter +(`"3 Completed / 1 InProgress / 0 Failed trials"`) instead of a single job's +`SecondaryStatus`. + +On completion the task emits a single `result: dict` literal with: + +- `HyperParameterTuningJobArn`, `HyperParameterTuningJobName` +- `BestTrainingJob` — the winning trial. Contains `TrainingJobName`, + `TrainingJobArn`, `TunedHyperParameters`, `ObjectiveStatus`, + `FinalHyperParameterTuningJobObjectiveMetric.{MetricName, Value}` and — + crucially — `ModelArtifacts.S3ModelArtifacts`. SageMaker's + `DescribeHyperParameterTuningJob` response does *not* include the trained + model URI; the connector resolves it via a single follow-up + `describe_training_job` call so this output chains directly into + `SageMakerModelTask`. +- `ModelArtifacts.S3ModelArtifacts` — top-level convenience copy of the best + trial's model URI so the result dict is **shape-compatible with + `SageMakerTrainingJobTask`'s output**. Any downstream task that reads + `result["ModelArtifacts"]["S3ModelArtifacts"]` works against either task + unchanged. +- `TrainingJobStatusCounters` — `Completed` / `InProgress` / `RetryableError` + / `NonRetryableError` / `Stopped` counts across all trials. +- `ObjectiveStatusCounters` — `Succeeded` / `Pending` / `Failed`. Note these + count objective-metric *evaluation*, not trial completion. A trial can + Complete but fail to emit the configured objective metric, in which case it + lands in `ObjectiveStatusCounters.Failed`. + +```python +from flytekitplugins.awssagemaker_hyperparameter_tuning import ( + SageMakerHyperParameterTuningJobTask, +) +from flytekit import kwtypes + +tuning = SageMakerHyperParameterTuningJobTask( + name="tune-xgboost", + config={ + "HyperParameterTuningJobName": "xgb-tune-{idempotence_token}", + "HyperParameterTuningJobConfig": { + "Strategy": "Bayesian", # Bayesian | Random | Hyperband | Grid + "HyperParameterTuningJobObjective": { + "Type": "Minimize", + "MetricName": "validation:rmse", + }, + "ResourceLimits": { + "MaxNumberOfTrainingJobs": 20, + "MaxParallelTrainingJobs": 4, + }, + "ParameterRanges": { + "ContinuousParameterRanges": [ + {"Name": "eta", "MinValue": "0.01", "MaxValue": "0.5", + "ScalingType": "Logarithmic"}, + ], + "IntegerParameterRanges": [ + {"Name": "max_depth", "MinValue": "3", "MaxValue": "9"}, + {"Name": "num_round", "MinValue": "10", "MaxValue": "200"}, + ], + }, + "TrainingJobEarlyStoppingType": "Auto", + }, + "TrainingJobDefinition": { + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + }, + "RoleArn": "{inputs.execution_role_arn}", + "StaticHyperParameters": {"objective": "reg:squarederror"}, + "InputDataConfig": [ + {"ChannelName": "train", "DataSource": {...}}, + {"ChannelName": "validation", "DataSource": {...}}, # required to emit validation:rmse + ], + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "ResourceConfig": {"InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 30}, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + }, + region="", + images={"training_image": ""}, + inputs=kwtypes(execution_role_arn=str, output_prefix=str), +) +``` + +A few important notes: + +- **Cost.** A tuning job's total cost is roughly `MaxNumberOfTrainingJobs × + per-trial instance-hours`. Start with a small `ResourceLimits` (4-8 trials) + while iterating on ranges, then scale up. +- **Objective metric must be emitted.** For SageMaker built-in algorithms + (XGBoost, BlazingText, etc.) the supported metric names are predefined + (`validation:rmse`, `validation:auc`, …) and require the corresponding + channel (e.g. a `validation` `InputDataConfig` channel) — set up the + channels accordingly. For custom containers, define + `AlgorithmSpecification.MetricDefinitions` with a regex that matches your + container's stdout/stderr so SageMaker can scrape the metric out. +- **`Strategy: "Hyperband"`** only works with iterative algorithms that emit + intermediate objective values, since Hyperband prunes weak trials early. + Default Bayesian is the safest pick if you're not sure. + +To chain HPO directly into the rest of the pipeline, just consume +`result["ModelArtifacts"]["S3ModelArtifacts"]` the same way you would with a +training-job result: + +```python +@workflow +def tune_and_deploy() -> dict: + hpo_result = tuning(...) + model_result, _ = model_task( + model_data=hpo_result["ModelArtifacts"]["S3ModelArtifacts"] + ) + return model_result +``` + +These tasks can be composed as HPO → model → Inference Recommender → batch +transform, passing the best model artifact and recommended instance type through +normal Flyte task outputs. + +Helper sync tasks: `SageMakerStopHyperParameterTuningJobTask` and +`SageMakerDescribeHyperParameterTuningJobTask` follow the stop/describe pattern +from the training and batch-transform connectors. + +## Batch Transform (Batch Inference) + +`SageMakerTransformJobTask` runs `CreateTransformJob` for offline scoring of a +dataset stored on S3 against an existing SageMaker `Model`. SageMaker writes one +`.out` per input object under `TransformOutput.S3OutputPath`. The task +emits a `result: dict` containing `TransformJobArn`, `TransformJobName`, +`ModelName`, `TransformOutput.S3OutputPath`, `TransformStartTime` and +`TransformEndTime`. + +For tabular predictive workloads, set `DataProcessing.JoinSource: "Input"` so +each output line carries the original input columns alongside the prediction — +otherwise the predictions have no key to join back to the source rows. + +```python +from flytekitplugins.awssagemaker_batch_transform import SageMakerTransformJobTask +from flytekit import kwtypes + +batch_score = SageMakerTransformJobTask( + name="batch-score", + config={ + "TransformJobName": "score-{idempotence_token}", + "ModelName": "{inputs.model_name}", + "TransformInput": { + "DataSource": { + "S3DataSource": { + "S3DataType": "S3Prefix", + "S3Uri": "{inputs.input_data}", + } + }, + "ContentType": "text/csv", + "SplitType": "Line", + }, + "TransformOutput": { + "S3OutputPath": "{inputs.output_prefix}", + "AssembleWith": "Line", + }, + "TransformResources": {"InstanceType": "ml.m5.xlarge", "InstanceCount": 1}, + "BatchStrategy": "MultiRecord", + "DataProcessing": {"JoinSource": "Input"}, + }, + region="", + inputs=kwtypes(model_name=str, input_data=str, output_prefix=str), +) +``` + +`ModelName` must reference an existing SageMaker `Model` — typically created +upstream by a `SageMakerModelTask` consuming a training job's +`S3ModelArtifacts` output. + +## Inference Recommender + +`SageMakerInferenceRecommenderJobTask` runs `CreateInferenceRecommendationsJob` +and waits for it to reach a terminal state. SageMaker benchmarks the model +across several real or candidate instance types and returns a ranked list of +`InferenceRecommendations`. The task emits a single `result: dict` containing: + +- `JobArn`, `JobName`, `JobType` (`Default` or `Advanced`) +- `InferenceRecommendations` — ranked list. Each entry has + `EndpointConfiguration.InstanceType`, `InitialInstanceCount`, optional + `ServerlessConfig`, plus `Metrics` (`CostPerHour`, `CostPerInference`, + `MaxInvocations`, `ModelLatency`, `CpuUtilization`, `MemoryUtilization`, + `ModelSetupTime`) and `ModelConfiguration`. +- `EndpointPerformances` — populated for `Default` jobs that benchmark existing + endpoints supplied through `InputConfig.Endpoints`. +- `CompletionTime` + +Two input modes are supported by SageMaker: + +- `ModelPackageVersionArn` — point at a versioned entry in a Model Package Group. +- `ModelName` + `ContainerConfig` — point at a bare `SageMaker.Model` plus a + payload archive and framework hint. Easier to chain after a fresh + `SageMakerTrainingJobTask`/`SageMakerModelTask` because no model-package + registration is required. + +`ContainerConfig.PayloadConfig.SamplePayloadUrl` must be an S3 URL to a single +`.tar.gz` archive containing the sample request body the Recommender will use +when benchmarking. `SupportedInstanceTypes` constrains the sweep to a fixed +list (omit it for a full sweep of the framework's supported instances). + +#### Default vs Advanced jobs — what fields each accepts + +The `CreateInferenceRecommendationsJob` boto3 API exposes a lot of bounds +fields under both `InputConfig` and the top level regardless of `JobType`, but +**AWS only accepts most of them when `JobType="Advanced"`** and returns +`ValidationException` if you set them on a `Default` job. Default jobs are +essentially fire-and-forget for ~45 minutes; the only effective bound is +`ContainerConfig.SupportedInstanceTypes`. + +| Field | `Default` | `Advanced` | +|---|---|---| +| `InputConfig.ModelPackageVersionArn` *or* `ModelName` + `ContainerConfig` | required | required | +| `InputConfig.ContainerConfig.SupportedInstanceTypes` | **the only bound** | optional | +| `InputConfig.JobDurationInSeconds` | rejected | required | +| `InputConfig.TrafficPattern` | rejected | required | +| `InputConfig.ResourceLimit` | rejected | required | +| `InputConfig.EndpointConfigurations` | rejected | required | +| top-level `StoppingConditions` | rejected | optional | +| `OutputConfig` | optional | optional | + +```python +from flytekitplugins.awssagemaker_inference_recommender import ( + SageMakerInferenceRecommenderJobTask, +) +from flytekit import kwtypes + +# Default job — fire-and-forget instance recommendation across the listed +# SupportedInstanceTypes. No StoppingConditions / JobDurationInSeconds. +recommend = SageMakerInferenceRecommenderJobTask( + name="recommend-instance", + config={ + "JobName": "rec-{idempotence_token}", + "JobType": "Default", + "RoleArn": "{inputs.execution_role_arn}", + "InputConfig": { + "ModelName": "{inputs.model_name}", + "ContainerConfig": { + "Domain": "MACHINE_LEARNING", + "Task": "OTHER", + "Framework": "XGBOOST", + "FrameworkVersion": "1.7", + "PayloadConfig": { + "SamplePayloadUrl": "{inputs.payload_url}", + "SupportedContentTypes": ["text/csv"], + }, + "SupportedInstanceTypes": [ + "ml.m5.large", + "ml.m5.xlarge", + "ml.c5.large", + "ml.c5.xlarge", + ], + }, + }, + }, + region="", + inputs=kwtypes(execution_role_arn=str, model_name=str, payload_url=str), +) +``` + +For an `Advanced` load test, the same task class accepts the full set of +fields the Default config rejects: + +```python +from flytekit import kwtypes + +recommend_advanced = SageMakerInferenceRecommenderJobTask( + name="recommend-load-test", + config={ + "JobName": "rec-adv-{idempotence_token}", + "JobType": "Advanced", + "RoleArn": "{inputs.execution_role_arn}", + "InputConfig": { + "ModelName": "{inputs.model_name}", + "ContainerConfig": {...}, # same as above + "JobDurationInSeconds": 7200, # Advanced-only + "TrafficPattern": { # Advanced-only + "TrafficType": "PHASES", + "Phases": [ + {"InitialNumberOfUsers": 1, "SpawnRate": 1, "DurationInSeconds": 120}, + ], + }, + "ResourceLimit": { # Advanced-only + "MaxNumberOfTests": 10, + "MaxParallelOfTests": 2, + }, + "EndpointConfigurations": [ # Advanced-only + {"InstanceType": "ml.m5.xlarge"}, + {"InstanceType": "ml.c5.xlarge"}, + ], + }, + "StoppingConditions": { # top-level, Advanced-only + "MaxInvocations": 500, + "ModelLatencyThresholds": [ + {"Percentile": "P95", "ValueInMilliseconds": 500}, + ], + }, + }, + region="", + inputs=kwtypes(execution_role_arn=str, model_name=str, payload_url=str), +) +``` + +### End-to-end: train, recommend, then batch transform on the recommended instance + +The recommender's `result["InferenceRecommendations"][0]["EndpointConfiguration"]["InstanceType"]` +is a stable scalar — pull it out in a small `@task` and feed it directly to the +next SageMaker task as a Flyte Promise. The boto3 mixin substitutes `{inputs.X}` +placeholders into the config at runtime, so the recommended instance type lands +in `TransformResources.InstanceType` (or `ProductionVariants[*].InstanceType`) +without any extra plumbing. + +```python +from flytekit import kwtypes, task, workflow +from flytekitplugins.awssagemaker_batch_transform import SageMakerTransformJobTask +from flytekitplugins.awssagemaker_inference import SageMakerModelTask +from flytekitplugins.awssagemaker_inference_recommender import ( + SageMakerInferenceRecommenderJobTask, +) +from flytekitplugins.awssagemaker_training import SageMakerTrainingJobTask + + +@task +def top_instance_type(recommender_result: dict) -> str: + """Pick the cheapest-meets-SLA instance the Recommender returned.""" + return recommender_result["InferenceRecommendations"][0]["EndpointConfiguration"]["InstanceType"] + + +training = SageMakerTrainingJobTask(...) # see Training section +model = SageMakerModelTask(...) # wraps S3ModelArtifacts as a Model +recommend = SageMakerInferenceRecommenderJobTask(...) # see snippet above +batch_score = SageMakerTransformJobTask( + name="batch-score-recommended", + config={ + "TransformJobName": "score-{idempotence_token}", + "ModelName": "{inputs.model_name}", + "TransformInput": {...}, + "TransformOutput": {"S3OutputPath": "{inputs.output_prefix}"}, + "TransformResources": { + # Recommender's pick flows in here via the Flyte Promise wired up below. + "InstanceType": "{inputs.instance_type}", + "InstanceCount": 1, + }, + }, + region="", + inputs=kwtypes(model_name=str, instance_type=str, output_prefix=str), +) + + +@workflow +def train_recommend_transform() -> dict: + train_result = training(...) + model_result, _ = model( + model_data=train_result["ModelArtifacts"]["S3ModelArtifacts"] + ) + model_name = model_result["ModelArn"].rsplit("/", 1)[-1] + + rec_result = recommend(model_name=model_name) + instance_type = top_instance_type(recommender_result=rec_result) + + return batch_score( + model_name=model_name, + instance_type=instance_type, + output_prefix="s3:///predictions/", + ) +``` + +The same pattern composes into an end-to-end training → model → recommender → +batch-transform workflow using normal Flyte task outputs. + +Helper sync tasks: `SageMakerStopInferenceRecommenderJobTask` and +`SageMakerDescribeInferenceRecommenderJobTask` mirror the stop/describe pattern +from the training and batch-transform connectors and are useful for inspecting +historical recommender runs from a Flyte workflow. diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/__init__.py new file mode 100644 index 0000000000..ce8de52cb0 --- /dev/null +++ b/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", +] diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/connector.py new file mode 100644 index 0000000000..fc06223032 --- /dev/null +++ b/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()) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/task.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_batch_transform/task.py new file mode 100644 index 0000000000..e181381655 --- /dev/null +++ b/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 ``.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, + ) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/__init__.py new file mode 100644 index 0000000000..723e6c5001 --- /dev/null +++ b/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", +] diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/connector.py new file mode 100644 index 0000000000..b1480e6180 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/connector.py @@ -0,0 +1,246 @@ +"""SageMaker hyperparameter-tuning-job connector. + +Mirrors the training-job connector: same long-running async lifecycle +(create -> describe-poll -> stop) but targets ``CreateHyperParameterTuningJob``. +On completion, surfaces ``BestTrainingJob`` plus the trained +``S3ModelArtifacts`` (looked up via a single follow-up +``describe_training_job`` call, since ``DescribeHyperParameterTuningJob`` does +not include it) so the result chains straight into ``SageMakerModelTask`` +without any extra workflow plumbing. +""" + +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 SageMakerHyperParameterTuningJobMetadata(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) -> "SageMakerHyperParameterTuningJobMetadata": + return cloudpickle.loads(data) + + +# HyperParameterTuningJobStatus -> Flyte phase. +# - Stopping is "still in flight" (SageMaker is asking each child training job to stop +# gracefully) so we report Running while the tear-down happens. +# - Stopped / Deleting / DeleteFailed are terminal admin states; treat as failure. +# - Failed covers both genuine job failure and warm-start parent failure. +_STATE_MAP = { + "InProgress": TaskExecution.RUNNING, + "Stopping": TaskExecution.RUNNING, + "Completed": TaskExecution.SUCCEEDED, + "Failed": TaskExecution.FAILED, + "Stopped": TaskExecution.FAILED, + "Deleting": TaskExecution.FAILED, + "DeleteFailed": TaskExecution.FAILED, +} + + +def _isoformat(value: Any) -> Any: + if isinstance(value, datetime): + return value.isoformat() + return value + + +def _project_best_training_job(best: Dict[str, Any], s3_model_artifacts: Optional[str]) -> Dict[str, Any]: + """Trim BestTrainingJob to the fields downstream tasks key off of.""" + objective_metric = best.get("FinalHyperParameterTuningJobObjectiveMetric") or {} + return { + "TrainingJobName": best.get("TrainingJobName"), + "TrainingJobArn": best.get("TrainingJobArn"), + "TrainingJobStatus": best.get("TrainingJobStatus"), + "ObjectiveStatus": best.get("ObjectiveStatus"), + "FinalHyperParameterTuningJobObjectiveMetric": { + "MetricName": objective_metric.get("MetricName"), + "Value": objective_metric.get("Value"), + }, + "TunedHyperParameters": dict(best.get("TunedHyperParameters") or {}), + # Carry the same nested shape Training emits so downstream tasks that + # consume train_result["ModelArtifacts"]["S3ModelArtifacts"] work + # unchanged against an HPO result. + "ModelArtifacts": {"S3ModelArtifacts": s3_model_artifacts}, + "TrainingStartTime": _isoformat(best.get("TrainingStartTime")), + "TrainingEndTime": _isoformat(best.get("TrainingEndTime")), + } + + +def _build_outputs(describe_response: Dict[str, Any], best_s3_model_artifacts: Optional[str]) -> Dict[str, Any]: + """Project DescribeHyperParameterTuningJob into a stable, downstream-friendly dict.""" + best = describe_response.get("BestTrainingJob") or {} + counters = describe_response.get("TrainingJobStatusCounters") or {} + obj_counters = describe_response.get("ObjectiveStatusCounters") or {} + + return { + "HyperParameterTuningJobArn": describe_response.get("HyperParameterTuningJobArn"), + "HyperParameterTuningJobName": describe_response.get("HyperParameterTuningJobName"), + "BestTrainingJob": _project_best_training_job(best, best_s3_model_artifacts), + # Same nested shape as ModelArtifacts above — promotes BestTrainingJob's + # artifacts to the top level so `result["ModelArtifacts"]["S3ModelArtifacts"]` + # is symmetric with the plain training-job task's output. + "ModelArtifacts": {"S3ModelArtifacts": best_s3_model_artifacts}, + "TrainingJobStatusCounters": { + "Completed": counters.get("Completed"), + "InProgress": counters.get("InProgress"), + "RetryableError": counters.get("RetryableError"), + "NonRetryableError": counters.get("NonRetryableError"), + "Stopped": counters.get("Stopped"), + }, + "ObjectiveStatusCounters": { + "Succeeded": obj_counters.get("Succeeded"), + "Pending": obj_counters.get("Pending"), + "Failed": obj_counters.get("Failed"), + }, + } + + +def _running_message(describe_response: Dict[str, Any]) -> Optional[str]: + """Compact 'N completed / M in-progress / K failed' status line for the Flyte UI.""" + counters = describe_response.get("TrainingJobStatusCounters") or {} + completed = counters.get("Completed") or 0 + in_progress = counters.get("InProgress") or 0 + failed = (counters.get("RetryableError") or 0) + (counters.get("NonRetryableError") or 0) + return f"{completed} Completed / {in_progress} InProgress / {failed} Failed trials" + + +class SageMakerHyperParameterTuningJobConnector(Boto3ConnectorMixin, AsyncConnectorBase): + """Long-running connector for SageMaker hyperparameter-tuning jobs.""" + + name = "SageMaker Hyperparameter Tuning Job Connector" + + def __init__(self): + super().__init__( + service="sagemaker", + task_type_name="sagemaker-hyperparameter-tuning-job", + metadata_type=SageMakerHyperParameterTuningJobMetadata, + ) + + async def create( + self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs + ) -> SageMakerHyperParameterTuningJobMetadata: + custom = task_template.custom + config = custom.get("config") + region = custom.get("region") + images = custom.get("images") + + try: + await self._call( + method="create_hyper_parameter_tuning_job", + config=config, + images=images, + 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"] + + # Idempotent re-runs: SageMaker rejects duplicate tuning job names. Treat as already-running. + if e.idempotence_token and ( + error_code == "ResourceInUse" + or (error_code == "ValidationException" and "Cannot create already existing" in error_message) + ): + return SageMakerHyperParameterTuningJobMetadata(config=config, region=region, inputs=inputs) + raise e + + return SageMakerHyperParameterTuningJobMetadata(config=config, region=region, inputs=inputs) + + async def _best_training_job_artifacts( + self, + describe_response: Dict[str, Any], + resource_meta: SageMakerHyperParameterTuningJobMetadata, + ) -> Optional[str]: + """Resolve BestTrainingJob -> S3ModelArtifacts via one extra describe call. + + ``DescribeHyperParameterTuningJob`` returns the best job's name and + tuned hyperparameters but NOT its ``ModelArtifacts`` — that lives on + ``DescribeTrainingJob``. We do the follow-up here so the connector's + ``result`` dict is self-contained: downstream tasks can chain on + ``result["ModelArtifacts"]["S3ModelArtifacts"]`` exactly like they do + for the plain training-job task. + """ + best = describe_response.get("BestTrainingJob") or {} + best_name = best.get("TrainingJobName") + if not best_name: + return None + training_describe, _ = await self._call( + method="describe_training_job", + config={"TrainingJobName": best_name}, + inputs=resource_meta.inputs, + region=resource_meta.region, + ) + return (training_describe.get("ModelArtifacts") or {}).get("S3ModelArtifacts") + + async def get(self, resource_meta: SageMakerHyperParameterTuningJobMetadata, **kwargs) -> Resource: + describe_response, _ = await self._call( + method="describe_hyper_parameter_tuning_job", + config={"HyperParameterTuningJobName": resource_meta.config.get("HyperParameterTuningJobName")}, + inputs=resource_meta.inputs, + region=resource_meta.region, + ) + + current_state = describe_response.get("HyperParameterTuningJobStatus") + flyte_phase = _STATE_MAP.get(current_state, TaskExecution.RUNNING) + + # While running we surface the trial counters so users see live progress + # ("3 Completed / 1 InProgress / 0 Failed trials"). On terminal failure + # FailureReason is the most useful single line. + message: Optional[str] = None + if current_state == "InProgress": + message = _running_message(describe_response) + elif current_state in ("Failed", "Stopped", "Deleting", "DeleteFailed"): + message = describe_response.get("FailureReason") or _running_message(describe_response) + + outputs: Optional[Dict[str, Any]] = None + if current_state == "Completed": + s3_model_artifacts = await self._best_training_job_artifacts(describe_response, resource_meta) + outputs = {"result": _build_outputs(describe_response, s3_model_artifacts)} + + return Resource(phase=flyte_phase, outputs=outputs, message=message) + + async def delete(self, resource_meta: SageMakerHyperParameterTuningJobMetadata, **kwargs): + try: + await self._call( + method="stop_hyper_parameter_tuning_job", + config={"HyperParameterTuningJobName": resource_meta.config.get("HyperParameterTuningJobName")}, + 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"] + + # Same swallow-on-already-terminal behaviour as the training connector. + if error_code == "ResourceNotFound" or ( + error_code == "ValidationException" and "non-running" in error_message + ): + return + raise e + + +ConnectorRegistry.register(SageMakerHyperParameterTuningJobConnector()) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/task.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/task.py new file mode 100644 index 0000000000..f29d300384 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_hyperparameter_tuning/task.py @@ -0,0 +1,121 @@ +"""User-facing tasks for SageMaker hyperparameter-tuning jobs.""" + +from typing import Any, Dict, Optional, Type, Union + +from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask + +from flytekit import ImageSpec, 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 +from flytekit.image_spec.image_spec import ImageBuildEngine + + +class SageMakerHyperParameterTuningJobTask(AsyncConnectorExecutorMixin, PythonTask): + """Run a SageMaker hyperparameter-tuning job and emit the best trial's artefacts. + + Outputs a single ``result: dict`` literal containing: + + - ``HyperParameterTuningJobArn``, ``HyperParameterTuningJobName`` + - ``BestTrainingJob`` — the trial SageMaker picked: ``TrainingJobName``, + ``TrainingJobArn``, ``TunedHyperParameters``, + ``FinalHyperParameterTuningJobObjectiveMetric.{MetricName, Value}``, + ``ObjectiveStatus``, plus the trial's ``ModelArtifacts.S3ModelArtifacts`` + resolved via a follow-up ``describe_training_job`` call (so this output + chains directly into ``SageMakerModelTask`` the same way the plain + ``SageMakerTrainingJobTask`` does) + - ``ModelArtifacts.S3ModelArtifacts`` — top-level convenience copy of the + best trial's model URI so workflows can consume it symmetrically with + training-job results + - ``TrainingJobStatusCounters`` — how many trials Completed / InProgress / + RetryableError / NonRetryableError / Stopped + - ``ObjectiveStatusCounters`` — Succeeded / Pending / Failed at the + objective-metric layer (a trial can Complete but Fail to emit the + objective metric — that lands in ``ObjectiveStatusCounters.Failed``) + + ``name`` identifies the Flyte task. ``config`` is the boto3 + ``create_hyper_parameter_tuning_job`` request and may contain + ``{inputs.X}``, ``{images.X}``, and ``{idempotence_token}`` placeholders. + ``region`` selects the AWS region. ``images`` maps trial-image placeholders + to image URIs or ``ImageSpec`` objects, and ``inputs`` maps input + placeholders to Flyte types. + """ + + _TASK_TYPE = "sagemaker-hyperparameter-tuning-job" + + def __init__( + self, + name: str, + config: Dict[str, Any], + region: Optional[str] = None, + images: Optional[Dict[str, Union[str, ImageSpec]]] = 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 + self._images = images + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + images = self._images + if images is not None: + for key, image in images.items(): + if isinstance(image, ImageSpec): + ImageBuildEngine.build(image) + images[key] = image.image_name() + return {"config": self._config, "region": self._region, "images": images} + + +class SageMakerStopHyperParameterTuningJobTask(BotoTask): + """Sync helper task that stops a running SageMaker hyperparameter-tuning 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_hyper_parameter_tuning_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) + + +class SageMakerDescribeHyperParameterTuningJobTask(BotoTask): + """Sync helper task that returns the full ``describe_hyper_parameter_tuning_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_hyper_parameter_tuning_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_connector.py index 48ff965381..7f084bf6a3 100644 --- a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_connector.py +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_connector.py @@ -19,19 +19,6 @@ from .boto3_mixin import Boto3ConnectorMixin, CustomException -# https://github.com/flyteorg/flyte/issues/4505 -def convert_floats_with_no_fraction_to_ints(data): - if isinstance(data, dict): - for key, value in data.items(): - data[key] = convert_floats_with_no_fraction_to_ints(value) - elif isinstance(data, list): - for i, item in enumerate(data): - data[i] = convert_floats_with_no_fraction_to_ints(item) - elif isinstance(data, float) and data.is_integer(): - return int(data) - return data - - class BotoConnector(SyncConnectorBase): """A general purpose boto3 connector that can be used to call any boto3 method.""" @@ -50,9 +37,7 @@ async def do( custom = task_template.custom service = custom.get("service") - raw_config = custom.get("config") - convert_floats_with_no_fraction_to_ints(raw_config) - config = raw_config + config = custom.get("config") region = custom.get("region") method = custom.get("method") images = custom.get("images") diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py index 4228b49c5e..ae61b59b25 100644 --- a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py @@ -21,11 +21,27 @@ def sorted_dict_str(d): if isinstance(d, dict): return "{" + ", ".join(f"{sorted_dict_str(k)}: {sorted_dict_str(v)}" for k, v in sorted(d.items())) + "}" elif isinstance(d, list): - return "[" + ", ".join(sorted_dict_str(i) for i in sorted(d, key=lambda x: str(x))) + "]" + # Dictionary order is irrelevant to a request, but list order can be + # semantically significant (for example, ContainerArguments). + return "[" + ", ".join(sorted_dict_str(i) for i in d) + "]" else: return str(d) +# https://github.com/flyteorg/flyte/issues/4505 +def convert_floats_with_no_fraction_to_ints(data): + """Recursively rewrite whole-number floats to ints so boto3 doesn't reject integer fields.""" + if isinstance(data, dict): + for key, value in data.items(): + data[key] = convert_floats_with_no_fraction_to_ints(value) + elif isinstance(data, list): + for i, item in enumerate(data): + data[i] = convert_floats_with_no_fraction_to_ints(item) + elif isinstance(data, float) and data.is_integer(): + return int(data) + return data + + account_id_map = { "us-east-1": "785573368785", "us-east-2": "007439368137", @@ -120,6 +136,11 @@ async def _call( updated_config = format_dict(self._service, config, args) + # boto3 rejects whole-number floats for integer-typed fields (e.g. InstanceCount, + # MaxRuntimeInSeconds). Normalize before hashing so semantically equivalent + # integer values produce the same idempotence token. + updated_config = convert_floats_with_no_fraction_to_ints(updated_config) + hash = "" if "idempotence_token" in str(updated_config): # compute hash of the config diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/__init__.py new file mode 100644 index 0000000000..ffb951e97a --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/__init__.py @@ -0,0 +1,30 @@ +""" +.. currentmodule:: flytekitplugins.awssagemaker_inference_recommender + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + SageMakerInferenceRecommenderJobConnector + SageMakerInferenceRecommenderJobTask + SageMakerStopInferenceRecommenderJobTask + SageMakerDescribeInferenceRecommenderJobTask +""" + +from .connector import ( + SageMakerInferenceRecommenderJobConnector, + SageMakerInferenceRecommenderJobMetadata, +) +from .task import ( + SageMakerDescribeInferenceRecommenderJobTask, + SageMakerInferenceRecommenderJobTask, + SageMakerStopInferenceRecommenderJobTask, +) + +__all__ = [ + "SageMakerInferenceRecommenderJobConnector", + "SageMakerInferenceRecommenderJobMetadata", + "SageMakerInferenceRecommenderJobTask", + "SageMakerStopInferenceRecommenderJobTask", + "SageMakerDescribeInferenceRecommenderJobTask", +] diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/connector.py new file mode 100644 index 0000000000..684fdd03c2 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/connector.py @@ -0,0 +1,232 @@ +"""SageMaker Inference Recommender job connector. + +Mirrors the training-job / batch-transform connectors. Targets +``CreateInferenceRecommendationsJob`` / ``DescribeInferenceRecommendationsJob`` / +``StopInferenceRecommendationsJob``. Surfaces the ranked +``InferenceRecommendations`` list (Default jobs) and the +``EndpointPerformances`` list (Default jobs targeting existing endpoints) so +downstream Flyte tasks can pick an instance type / endpoint config without +re-querying SageMaker. + +Note: ``InferenceRecommendationsJob`` ``Status`` values are ALL_CAPS +(``PENDING`` / ``IN_PROGRESS`` / ``COMPLETED`` / ``FAILED`` / ``STOPPING`` / +``STOPPED`` / ``DELETING`` / ``DELETED``), unlike training/transform jobs which +use PascalCase. +""" + +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, 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 SageMakerInferenceRecommenderJobMetadata(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) -> "SageMakerInferenceRecommenderJobMetadata": + return cloudpickle.loads(data) + + +# Status values per boto3 reference. PENDING and IN_PROGRESS keep the job +# in flight; STOPPING is still a running tear-down. STOPPED covers both +# user-stop and timeout. DELETING/DELETED are admin states - treat as failure +# so we don't silently surface partial recommendations. +_STATE_MAP = { + "PENDING": TaskExecution.RUNNING, + "IN_PROGRESS": TaskExecution.RUNNING, + "STOPPING": TaskExecution.RUNNING, + "COMPLETED": TaskExecution.SUCCEEDED, + "FAILED": TaskExecution.FAILED, + "STOPPED": TaskExecution.FAILED, + "DELETING": TaskExecution.FAILED, + "DELETED": TaskExecution.FAILED, +} + + +def _isoformat(value: Any) -> Any: + if isinstance(value, datetime): + return value.isoformat() + return value + + +def _project_recommendation(rec: Dict[str, Any]) -> Dict[str, Any]: + """Trim a single InferenceRecommendations entry to the fields users actually pick on.""" + metrics = rec.get("Metrics") or {} + endpoint_config = rec.get("EndpointConfiguration") or {} + model_config = rec.get("ModelConfiguration") or {} + serverless_config = endpoint_config.get("ServerlessConfig") or {} + + return { + "RecommendationId": rec.get("RecommendationId"), + "Metrics": { + "CostPerHour": metrics.get("CostPerHour"), + "CostPerInference": metrics.get("CostPerInference"), + "MaxInvocations": metrics.get("MaxInvocations"), + "ModelLatency": metrics.get("ModelLatency"), + "CpuUtilization": metrics.get("CpuUtilization"), + "MemoryUtilization": metrics.get("MemoryUtilization"), + "ModelSetupTime": metrics.get("ModelSetupTime"), + }, + "EndpointConfiguration": { + "EndpointName": endpoint_config.get("EndpointName"), + "VariantName": endpoint_config.get("VariantName"), + "InstanceType": endpoint_config.get("InstanceType"), + "InitialInstanceCount": endpoint_config.get("InitialInstanceCount"), + "ServerlessConfig": { + "MemorySizeInMB": serverless_config.get("MemorySizeInMB"), + "MaxConcurrency": serverless_config.get("MaxConcurrency"), + "ProvisionedConcurrency": serverless_config.get("ProvisionedConcurrency"), + } + if serverless_config + else None, + }, + "ModelConfiguration": { + "InferenceSpecificationName": model_config.get("InferenceSpecificationName"), + "CompilationJobName": model_config.get("CompilationJobName"), + }, + "InvocationStartTime": _isoformat(rec.get("InvocationStartTime")), + "InvocationEndTime": _isoformat(rec.get("InvocationEndTime")), + } + + +def _project_endpoint_performance(perf: Dict[str, Any]) -> Dict[str, Any]: + metrics = perf.get("Metrics") or {} + endpoint_info = perf.get("EndpointInfo") or {} + return { + "Metrics": { + "MaxInvocations": metrics.get("MaxInvocations"), + "ModelLatency": metrics.get("ModelLatency"), + }, + "EndpointInfo": {"EndpointName": endpoint_info.get("EndpointName")}, + } + + +def _build_outputs(describe_response: Dict[str, Any]) -> Dict[str, Any]: + """Project describe_inference_recommendations_job down to a stable, downstream-friendly dict.""" + recommendations: List[Dict[str, Any]] = [ + _project_recommendation(rec) for rec in (describe_response.get("InferenceRecommendations") or []) + ] + endpoint_performances: List[Dict[str, Any]] = [ + _project_endpoint_performance(perf) for perf in (describe_response.get("EndpointPerformances") or []) + ] + + return { + "JobArn": describe_response.get("JobArn"), + "JobName": describe_response.get("JobName"), + "JobType": describe_response.get("JobType"), + "InferenceRecommendations": recommendations, + "EndpointPerformances": endpoint_performances, + "CompletionTime": _isoformat(describe_response.get("CompletionTime")), + } + + +class SageMakerInferenceRecommenderJobConnector(Boto3ConnectorMixin, AsyncConnectorBase): + """Long-running connector for SageMaker Inference Recommender jobs.""" + + name = "SageMaker Inference Recommender Job Connector" + + def __init__(self): + super().__init__( + service="sagemaker", + task_type_name="sagemaker-inference-recommender-job", + metadata_type=SageMakerInferenceRecommenderJobMetadata, + ) + + async def create( + self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs + ) -> SageMakerInferenceRecommenderJobMetadata: + custom = task_template.custom + config = custom.get("config") + region = custom.get("region") + + try: + await self._call( + method="create_inference_recommendations_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"] + + # Idempotent re-runs: SageMaker rejects duplicate job names. Treat as already-running. + if e.idempotence_token and ( + error_code == "ResourceInUse" + or (error_code == "ValidationException" and "Cannot create already existing" in error_message) + ): + return SageMakerInferenceRecommenderJobMetadata(config=config, region=region, inputs=inputs) + raise e + + return SageMakerInferenceRecommenderJobMetadata(config=config, region=region, inputs=inputs) + + async def get(self, resource_meta: SageMakerInferenceRecommenderJobMetadata, **kwargs) -> Resource: + describe_response, _ = await self._call( + method="describe_inference_recommendations_job", + config={"JobName": resource_meta.config.get("JobName")}, + inputs=resource_meta.inputs, + region=resource_meta.region, + ) + + current_state = describe_response.get("Status") + flyte_phase = _STATE_MAP.get(current_state, TaskExecution.RUNNING) + + # Inference Recommender has no SecondaryStatus, but FailureReason is the + # most useful single line on terminal failure. + message: Optional[str] = None + if current_state in ("FAILED", "STOPPED", "DELETING", "DELETED"): + 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: SageMakerInferenceRecommenderJobMetadata, **kwargs): + try: + await self._call( + method="stop_inference_recommendations_job", + config={"JobName": resource_meta.config.get("JobName")}, + 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"] + + # Flyte may invoke delete() after the job has naturally completed (or already + # been stopped). SageMaker rejects stop on a non-running job - swallow that + # specific error since there's nothing to cancel. + if error_code == "ResourceNotFound" or ( + error_code == "ValidationException" and "non-running" in error_message + ): + return + raise e + + +ConnectorRegistry.register(SageMakerInferenceRecommenderJobConnector()) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/task.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/task.py new file mode 100644 index 0000000000..566a988438 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference_recommender/task.py @@ -0,0 +1,104 @@ +"""User-facing tasks for SageMaker Inference Recommender 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 SageMakerInferenceRecommenderJobTask(AsyncConnectorExecutorMixin, PythonTask): + """Run a SageMaker Inference Recommender job and emit its ranked recommendations. + + Outputs a single ``result: dict`` literal containing ``JobArn``, ``JobName``, + ``JobType`` (``Default`` or ``Advanced``), ``InferenceRecommendations`` (ranked + list with ``EndpointConfiguration.InstanceType``, ``InitialInstanceCount`` and + cost / latency / throughput metrics - feed the top entry into + ``SageMakerEndpointConfigTask`` to deploy on the recommended instance type), + ``EndpointPerformances`` (populated for Default jobs that benchmark existing + endpoints supplied through ``InputConfig.Endpoints``), and ``CompletionTime``. + + Use ``JobType: "Default"`` for a quick instance-type sweep (~45 min) keyed off + a ``ModelPackageVersionArn``; use ``JobType: "Advanced"`` to run a custom + traffic pattern + ``StoppingConditions`` over user-supplied + ``EndpointConfigurations``. + + ``name`` identifies the Flyte task. ``config`` is the boto3 + ``create_inference_recommendations_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-inference-recommender-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 SageMakerStopInferenceRecommenderJobTask(BotoTask): + """Sync helper task that stops a running SageMaker Inference Recommender 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_inference_recommendations_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) + + +class SageMakerDescribeInferenceRecommenderJobTask(BotoTask): + """Sync helper task that returns the full ``describe_inference_recommendations_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_inference_recommendations_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/__init__.py new file mode 100644 index 0000000000..74236c8fb7 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/__init__.py @@ -0,0 +1,27 @@ +""" +.. currentmodule:: flytekitplugins.awssagemaker_processing + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + SageMakerProcessingJobConnector + SageMakerProcessingJobTask + SageMakerStopProcessingJobTask + SageMakerDescribeProcessingJobTask +""" + +from .connector import SageMakerProcessingJobConnector, SageMakerProcessingJobMetadata +from .task import ( + SageMakerDescribeProcessingJobTask, + SageMakerProcessingJobTask, + SageMakerStopProcessingJobTask, +) + +__all__ = [ + "SageMakerProcessingJobConnector", + "SageMakerProcessingJobMetadata", + "SageMakerProcessingJobTask", + "SageMakerStopProcessingJobTask", + "SageMakerDescribeProcessingJobTask", +] diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/connector.py new file mode 100644 index 0000000000..16338a7e98 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/connector.py @@ -0,0 +1,186 @@ +"""SageMaker processing-job connector. + +Mirrors ``awssagemaker_training.connector`` (long-running async lifecycle: +create → describe-poll → stop) but targets ``CreateProcessingJob`` instead of +``CreateTrainingJob``. Surfaces the processed-output S3 URIs in outputs so +downstream Flyte tasks (a ``SageMakerTrainingJobTask`` consuming engineered +features, a ``SageMakerModelTask``, or a custom gate task) can consume them +without any extra plumbing. + +Note: ``ProcessingJobStatus`` values (``InProgress`` / ``Completed`` / +``Failed`` / ``Stopping`` / ``Stopped``) match the PascalCase training-job +convention, but processing jobs have no ``SecondaryStatus`` — ``ExitMessage`` / +``FailureReason`` are the useful single lines on terminal states. +""" + +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, 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 SageMakerProcessingJobMetadata(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) -> "SageMakerProcessingJobMetadata": + return cloudpickle.loads(data) + + +# ProcessingJobStatus → Flyte phase (verified against current boto3 reference). +# - Stopping is "still in flight" so we report Running while SageMaker tears the job down. +# - Stopped covers both user-stop and MaxRuntimeExceeded — a user-visible failure from +# the workflow's perspective. +# Processing jobs have no "Deleting" state (unlike training jobs). +_STATE_MAP = { + "InProgress": TaskExecution.RUNNING, + "Stopping": TaskExecution.RUNNING, + "Completed": TaskExecution.SUCCEEDED, + "Failed": TaskExecution.FAILED, + "Stopped": TaskExecution.FAILED, +} + + +def _isoformat(value: Any) -> Any: + """Best-effort ISO8601 string for datetime values, leave everything else alone.""" + if isinstance(value, datetime): + return value.isoformat() + return value + + +def _build_outputs(describe_response: Dict[str, Any]) -> Dict[str, Any]: + """Project the describe_processing_job response down to a stable, downstream-friendly dict.""" + output_config = describe_response.get("ProcessingOutputConfig") or {} + outputs: List[Dict[str, Any]] = [] + for output in output_config.get("Outputs") or []: + s3_output = output.get("S3Output") or {} + feature_store_output = output.get("FeatureStoreOutput") or {} + projected_output = {"OutputName": output.get("OutputName")} + if s3_output: + projected_output["S3Uri"] = s3_output.get("S3Uri") + if feature_store_output: + projected_output["FeatureGroupName"] = feature_store_output.get("FeatureGroupName") + outputs.append(projected_output) + + return { + "ProcessingJobArn": describe_response.get("ProcessingJobArn"), + "ProcessingJobName": describe_response.get("ProcessingJobName"), + "Outputs": outputs, + "ExitMessage": describe_response.get("ExitMessage"), + "ProcessingStartTime": _isoformat(describe_response.get("ProcessingStartTime")), + "ProcessingEndTime": _isoformat(describe_response.get("ProcessingEndTime")), + } + + +class SageMakerProcessingJobConnector(Boto3ConnectorMixin, AsyncConnectorBase): + """Long-running connector for SageMaker processing jobs.""" + + name = "SageMaker Processing Job Connector" + + def __init__(self): + super().__init__( + service="sagemaker", + task_type_name="sagemaker-processing-job", + metadata_type=SageMakerProcessingJobMetadata, + ) + + async def create( + self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs + ) -> SageMakerProcessingJobMetadata: + custom = task_template.custom + config = custom.get("config") + region = custom.get("region") + images = custom.get("images") + + try: + await self._call( + method="create_processing_job", + config=config, + images=images, + 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"] + + # Idempotent re-runs: SageMaker rejects duplicate job names. Treat as already-running. + if e.idempotence_token and ( + error_code == "ResourceInUse" + or (error_code == "ValidationException" and "Cannot create already existing" in error_message) + ): + return SageMakerProcessingJobMetadata(config=config, region=region, inputs=inputs) + raise e + + return SageMakerProcessingJobMetadata(config=config, region=region, inputs=inputs) + + async def get(self, resource_meta: SageMakerProcessingJobMetadata, **kwargs) -> Resource: + describe_response, _ = await self._call( + method="describe_processing_job", + config={"ProcessingJobName": resource_meta.config.get("ProcessingJobName")}, + inputs=resource_meta.inputs, + region=resource_meta.region, + ) + + current_state = describe_response.get("ProcessingJobStatus") + flyte_phase = _STATE_MAP.get(current_state, TaskExecution.RUNNING) + + # Processing jobs expose no SecondaryStatus, so there's no live sub-status to + # surface while running. On Failed/Stopped, FailureReason (falling back to + # ExitMessage) is the most useful single line. + message: Optional[str] = None + if current_state in ("Failed", "Stopped"): + message = describe_response.get("FailureReason") or describe_response.get("ExitMessage") + + 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: SageMakerProcessingJobMetadata, **kwargs): + try: + await self._call( + method="stop_processing_job", + config={"ProcessingJobName": resource_meta.config.get("ProcessingJobName")}, + 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"] + + # Flyte may invoke delete() after the job has naturally completed (or already + # been stopped). SageMaker rejects stop on a non-running job — swallow that + # specific error since there's nothing to cancel. + if error_code == "ResourceNotFound" or ( + error_code == "ValidationException" and "non-running" in error_message + ): + return + raise e + + +ConnectorRegistry.register(SageMakerProcessingJobConnector()) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/task.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/task.py new file mode 100644 index 0000000000..b97badefaa --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/task.py @@ -0,0 +1,112 @@ +"""User-facing tasks for SageMaker processing jobs.""" + +from typing import Any, Dict, Optional, Type, Union + +from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask + +from flytekit import ImageSpec, 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 +from flytekit.image_spec.image_spec import ImageBuildEngine + + +class SageMakerProcessingJobTask(AsyncConnectorExecutorMixin, PythonTask): + """Run a SageMaker processing job and emit its output S3 URIs. + + Processing jobs cover data pre/post-processing, feature engineering, model + evaluation, and SageMaker Clarify (bias / explainability) — the steps that + bookend training. The container image lives at ``AppSpecification.ImageUri``; + inputs are commonly S3-resident, while outputs can target S3 or SageMaker + Feature Store through ``ProcessingOutputConfig``. + + Outputs a single ``result: dict`` literal containing ``ProcessingJobArn``, + ``ProcessingJobName``, ``Outputs`` (a list containing ``OutputName`` plus + either ``S3Uri`` for S3 destinations or ``FeatureGroupName`` for Feature + Store destinations), and ``ExitMessage``. + + ``name`` identifies the Flyte task. ``config`` is the boto3 + ``create_processing_job`` request and may contain ``{inputs.X}``, + ``{images.X}``, and ``{idempotence_token}`` placeholders. ``region`` selects + the AWS region. ``images`` maps image placeholders to image URIs or + ``ImageSpec`` objects, and ``inputs`` maps input placeholders to Flyte types. + """ + + _TASK_TYPE = "sagemaker-processing-job" + + def __init__( + self, + name: str, + config: Dict[str, Any], + region: Optional[str] = None, + images: Optional[Dict[str, Union[str, ImageSpec]]] = 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 + self._images = images + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + images = self._images + if images is not None: + for key, image in images.items(): + if isinstance(image, ImageSpec): + ImageBuildEngine.build(image) + images[key] = image.image_name() + return {"config": self._config, "region": self._region, "images": images} + + +class SageMakerStopProcessingJobTask(BotoTask): + """Sync helper task that stops a running SageMaker processing 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_processing_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) + + +class SageMakerDescribeProcessingJobTask(BotoTask): + """Sync helper task that returns the full ``describe_processing_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_processing_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/__init__.py new file mode 100644 index 0000000000..5bef5c37b6 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/__init__.py @@ -0,0 +1,27 @@ +""" +.. currentmodule:: flytekitplugins.awssagemaker_training + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + SageMakerTrainingJobConnector + SageMakerTrainingJobTask + SageMakerStopTrainingJobTask + SageMakerDescribeTrainingJobTask +""" + +from .connector import SageMakerTrainingJobConnector, SageMakerTrainingJobMetadata +from .task import ( + SageMakerDescribeTrainingJobTask, + SageMakerStopTrainingJobTask, + SageMakerTrainingJobTask, +) + +__all__ = [ + "SageMakerTrainingJobConnector", + "SageMakerTrainingJobMetadata", + "SageMakerTrainingJobTask", + "SageMakerStopTrainingJobTask", + "SageMakerDescribeTrainingJobTask", +] diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/connector.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/connector.py new file mode 100644 index 0000000000..16e7e4795b --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/connector.py @@ -0,0 +1,186 @@ +"""SageMaker training-job connector. + +Mirrors the pattern in ``awssagemaker_inference.connector`` (long-running async +lifecycle: create → describe-poll → stop) but targets ``CreateTrainingJob`` +instead of ``CreateEndpoint``. Surfaces the trained ``S3ModelArtifacts`` URI and +final metrics in outputs so downstream Flyte tasks (a ``SageMakerModelTask`` for +deployment, or a custom Flyte gate task for accuracy thresholds) can consume them +without any extra plumbing. +""" + +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 SageMakerTrainingJobMetadata(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) -> "SageMakerTrainingJobMetadata": + return cloudpickle.loads(data) + + +# TrainingJobStatus → Flyte phase (verified against current boto3 reference). +# - Stopping is "still in flight" so we report Running while SageMaker tears the job down. +# - Stopped covers both user-stop and MaxRuntimeExceeded / MaxWaitTimeExceeded — all of +# these are user-visible failures from the workflow's perspective. +# - Deleting is a terminal admin state; treat as failure. +_STATE_MAP = { + "InProgress": TaskExecution.RUNNING, + "Stopping": TaskExecution.RUNNING, + "Completed": TaskExecution.SUCCEEDED, + "Failed": TaskExecution.FAILED, + "Stopped": TaskExecution.FAILED, + "Deleting": TaskExecution.FAILED, +} + + +def _isoformat(value: Any) -> Any: + """Best-effort ISO8601 string for datetime values, leave everything else alone.""" + if isinstance(value, datetime): + return value.isoformat() + return value + + +def _build_outputs(describe_response: Dict[str, Any]) -> Dict[str, Any]: + """Project the describe_training_job response down to a stable, downstream-friendly dict.""" + metrics = [] + for metric in describe_response.get("FinalMetricDataList") or []: + metrics.append( + { + "MetricName": metric.get("MetricName"), + "Value": metric.get("Value"), + "Timestamp": _isoformat(metric.get("Timestamp")), + } + ) + + model_artifacts = describe_response.get("ModelArtifacts") or {} + output_data_config = describe_response.get("OutputDataConfig") or {} + + return { + "TrainingJobArn": describe_response.get("TrainingJobArn"), + "TrainingJobName": describe_response.get("TrainingJobName"), + "ModelArtifacts": {"S3ModelArtifacts": model_artifacts.get("S3ModelArtifacts")}, + "OutputDataConfig": {"S3OutputPath": output_data_config.get("S3OutputPath")}, + "FinalMetricDataList": metrics, + "BillableTimeInSeconds": describe_response.get("BillableTimeInSeconds"), + "TrainingTimeInSeconds": describe_response.get("TrainingTimeInSeconds"), + } + + +class SageMakerTrainingJobConnector(Boto3ConnectorMixin, AsyncConnectorBase): + """Long-running connector for SageMaker training jobs.""" + + name = "SageMaker Training Job Connector" + + def __init__(self): + super().__init__( + service="sagemaker", + task_type_name="sagemaker-training-job", + metadata_type=SageMakerTrainingJobMetadata, + ) + + async def create( + self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs + ) -> SageMakerTrainingJobMetadata: + custom = task_template.custom + config = custom.get("config") + region = custom.get("region") + images = custom.get("images") + + try: + await self._call( + method="create_training_job", + config=config, + images=images, + 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"] + + # Idempotent re-runs: SageMaker rejects duplicate job names. Treat as already-running. + if e.idempotence_token and ( + error_code == "ResourceInUse" + or (error_code == "ValidationException" and "Cannot create already existing" in error_message) + ): + return SageMakerTrainingJobMetadata(config=config, region=region, inputs=inputs) + raise e + + return SageMakerTrainingJobMetadata(config=config, region=region, inputs=inputs) + + async def get(self, resource_meta: SageMakerTrainingJobMetadata, **kwargs) -> Resource: + describe_response, _ = await self._call( + method="describe_training_job", + config={"TrainingJobName": resource_meta.config.get("TrainingJobName")}, + inputs=resource_meta.inputs, + region=resource_meta.region, + ) + + current_state = describe_response.get("TrainingJobStatus") + flyte_phase = _STATE_MAP.get(current_state, TaskExecution.RUNNING) + + # Surface SecondaryStatus while running so the Flyte UI shows live progress + # (Starting → Downloading → Training → Uploading → Completed). On Failed/Stopped, + # FailureReason is the most useful single line. + message: Optional[str] = None + if current_state == "InProgress": + message = describe_response.get("SecondaryStatus") + elif current_state in ("Failed", "Stopped"): + message = describe_response.get("FailureReason") or describe_response.get("SecondaryStatus") + + 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: SageMakerTrainingJobMetadata, **kwargs): + try: + await self._call( + method="stop_training_job", + config={"TrainingJobName": resource_meta.config.get("TrainingJobName")}, + 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"] + + # Flyte may invoke delete() after the job has naturally completed (or already + # been stopped). SageMaker rejects stop on a non-running job — swallow that + # specific error since there's nothing to cancel. + if error_code == "ResourceNotFound" or ( + error_code == "ValidationException" and "non-running" in error_message + ): + return + raise e + + +ConnectorRegistry.register(SageMakerTrainingJobConnector()) diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/task.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/task.py new file mode 100644 index 0000000000..79e1f7de06 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/task.py @@ -0,0 +1,108 @@ +"""User-facing tasks for SageMaker training jobs.""" + +from typing import Any, Dict, Optional, Type, Union + +from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask + +from flytekit import ImageSpec, 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 +from flytekit.image_spec.image_spec import ImageBuildEngine + + +class SageMakerTrainingJobTask(AsyncConnectorExecutorMixin, PythonTask): + """Run a SageMaker training job and emit its model artefact URI plus final metrics. + + Outputs a single ``result: dict`` literal containing ``TrainingJobArn``, + ``TrainingJobName``, ``ModelArtifacts.S3ModelArtifacts`` (the S3 URI of the + trained ``model.tar.gz`` — feed this into ``SageMakerModelTask`` to deploy), + ``OutputDataConfig.S3OutputPath``, ``FinalMetricDataList`` (last value of every + metric defined in ``AlgorithmSpecification.MetricDefinitions``), + ``BillableTimeInSeconds`` and ``TrainingTimeInSeconds``. + + ``name`` identifies the Flyte task. ``config`` is the boto3 + ``create_training_job`` request and may contain ``{inputs.X}``, + ``{images.X}``, and ``{idempotence_token}`` placeholders. ``region`` selects + the AWS region. ``images`` maps image placeholders to image URIs or + ``ImageSpec`` objects, and ``inputs`` maps input placeholders to Flyte types. + """ + + _TASK_TYPE = "sagemaker-training-job" + + def __init__( + self, + name: str, + config: Dict[str, Any], + region: Optional[str] = None, + images: Optional[Dict[str, Union[str, ImageSpec]]] = 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 + self._images = images + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + images = self._images + if images is not None: + for key, image in images.items(): + if isinstance(image, ImageSpec): + ImageBuildEngine.build(image) + images[key] = image.image_name() + return {"config": self._config, "region": self._region, "images": images} + + +class SageMakerStopTrainingJobTask(BotoTask): + """Sync helper task that stops a running SageMaker training 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_training_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) + + +class SageMakerDescribeTrainingJobTask(BotoTask): + """Sync helper task that returns the full ``describe_training_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_training_job", + config=config, + region=region, + ), + inputs=inputs, + **kwargs, + ) diff --git a/plugins/flytekit-aws-sagemaker/setup.py b/plugins/flytekit-aws-sagemaker/setup.py index 0078086d43..3d41cfe7e6 100644 --- a/plugins/flytekit-aws-sagemaker/setup.py +++ b/plugins/flytekit-aws-sagemaker/setup.py @@ -2,6 +2,11 @@ PLUGIN_NAME = "awssagemaker" INFERENCE_PACKAGE = "awssagemaker_inference" +TRAINING_PACKAGE = "awssagemaker_training" +BATCH_TRANSFORM_PACKAGE = "awssagemaker_batch_transform" +INFERENCE_RECOMMENDER_PACKAGE = "awssagemaker_inference_recommender" +HYPERPARAMETER_TUNING_PACKAGE = "awssagemaker_hyperparameter_tuning" +PROCESSING_PACKAGE = "awssagemaker_processing" microlib_name = f"flytekitplugins-{PLUGIN_NAME}" @@ -18,7 +23,14 @@ author_email="admin@flyte.org", description="Flytekit AWS SageMaker Plugin", namespace_packages=["flytekitplugins"], - packages=[f"flytekitplugins.{INFERENCE_PACKAGE}"], + packages=[ + f"flytekitplugins.{INFERENCE_PACKAGE}", + f"flytekitplugins.{TRAINING_PACKAGE}", + f"flytekitplugins.{BATCH_TRANSFORM_PACKAGE}", + f"flytekitplugins.{INFERENCE_RECOMMENDER_PACKAGE}", + f"flytekitplugins.{HYPERPARAMETER_TUNING_PACKAGE}", + f"flytekitplugins.{PROCESSING_PACKAGE}", + ], install_requires=plugin_requires, license="apache2", python_requires=">=3.10", @@ -35,5 +47,14 @@ "Topic :: Software Development :: Libraries", "Topic :: Software Development :: Libraries :: Python Modules", ], - entry_points={"flytekit.plugins": [f"{INFERENCE_PACKAGE}=flytekitplugins.{INFERENCE_PACKAGE}"]}, + entry_points={ + "flytekit.plugins": [ + f"{INFERENCE_PACKAGE}=flytekitplugins.{INFERENCE_PACKAGE}", + f"{TRAINING_PACKAGE}=flytekitplugins.{TRAINING_PACKAGE}", + f"{BATCH_TRANSFORM_PACKAGE}=flytekitplugins.{BATCH_TRANSFORM_PACKAGE}", + f"{INFERENCE_RECOMMENDER_PACKAGE}=flytekitplugins.{INFERENCE_RECOMMENDER_PACKAGE}", + f"{HYPERPARAMETER_TUNING_PACKAGE}=flytekitplugins.{HYPERPARAMETER_TUNING_PACKAGE}", + f"{PROCESSING_PACKAGE}=flytekitplugins.{PROCESSING_PACKAGE}", + ] + }, ) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_connector.py new file mode 100644 index 0000000000..4520c5cf50 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_connector.py @@ -0,0 +1,262 @@ +from datetime import datetime, timedelta +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_batch_transform.connector import ( + SageMakerTransformJobMetadata, +) +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException + +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.interfaces.cli_identifiers import Identifier +from flytekit.models import literals +from flytekit.models.core.identifier import ResourceType +from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate + +idempotence_token = "74443947857331f7" + +REGION = "us-east-2" +TRANSFORM_JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:transform-job/score-74443947857331f7" +) +S3_OUTPUT_PATH = "s3://my-bucket/predictions/score-74443947857331f7/" + + +def _task_config(): + return { + "config": { + "TransformJobName": "score-{idempotence_token}", + "ModelName": "{inputs.model_name}", + "TransformInput": { + "DataSource": { + "S3DataSource": { + "S3DataType": "S3Prefix", + "S3Uri": "{inputs.input_data}", + } + }, + "ContentType": "text/csv", + "SplitType": "Line", + }, + "TransformOutput": { + "S3OutputPath": "{inputs.output_prefix}", + "AssembleWith": "Line", + }, + "TransformResources": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + }, + "DataProcessing": {"JoinSource": "Input"}, + }, + "region": REGION, + } + + +def _task_template(): + task_id = Identifier( + resource_type=ResourceType.TASK, + project="project", + domain="domain", + name="name", + version="version", + ) + task_metadata = TaskMetadata( + discoverable=True, + runtime=RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), + timeout=timedelta(days=1), + retries=literals.RetryStrategy(3), + interruptible=True, + discovery_version="0.1.1b0", + deprecated_error_message="This is deprecated!", + cache_serializable=True, + pod_template_name="A", + cache_ignore_input_vars=(), + ) + return TaskTemplate( + id=task_id, + custom=_task_config(), + metadata=task_metadata, + interface=None, + type="sagemaker-transform-job", + ) + + +def _completed_describe_response(): + return { + "TransformJobName": "score-74443947857331f7", + "TransformJobArn": TRANSFORM_JOB_ARN, + "TransformJobStatus": "Completed", + "ModelName": "ranker-prod", + "TransformOutput": {"S3OutputPath": S3_OUTPUT_PATH, "AssembleWith": "Line"}, + "TransformStartTime": datetime(2026, 4, 30, 10, 0, 0), + "TransformEndTime": datetime(2026, 4, 30, 10, 12, 0), + } + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_create_get_delete_happy_path(mock_call): + mock_call.return_value = (_completed_describe_response(), idempotence_token) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + metadata = SageMakerTransformJobMetadata( + config=_task_config()["config"], region=REGION + ) + + response = await connector.create(_task_template()) + assert response == metadata + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["TransformJobArn"] == TRANSFORM_JOB_ARN + assert result["TransformJobName"] == "score-74443947857331f7" + assert result["ModelName"] == "ranker-prod" + assert result["TransformOutput"] == {"S3OutputPath": S3_OUTPUT_PATH} + assert result["TransformStartTime"] == "2026-04-30T10:00:00" + assert result["TransformEndTime"] == "2026-04-30T10:12:00" + + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_get_inprogress_no_message(mock_call): + """Transform jobs have no SecondaryStatus, so message stays None during InProgress.""" + mock_call.return_value = ( + { + "TransformJobName": "score-x", + "TransformJobArn": TRANSFORM_JOB_ARN, + "TransformJobStatus": "InProgress", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + metadata = SageMakerTransformJobMetadata( + config={"TransformJobName": "score-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.message is None + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_get_failed_surfaces_failure_reason(mock_call): + mock_call.return_value = ( + { + "TransformJobName": "score-x", + "TransformJobArn": TRANSFORM_JOB_ARN, + "TransformJobStatus": "Failed", + "FailureReason": "ClientError: container exited with code 1", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + metadata = SageMakerTransformJobMetadata( + config={"TransformJobName": "score-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "ClientError: container exited with code 1" + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_create_already_exists_returns_metadata(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Transform job score-74443947857331f7 already exists", + } + }, + operation_name="CreateTransformJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + response = await connector.create(_task_template()) + assert response.config == _task_config()["config"] + assert response.region == REGION + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_create_resource_limit_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceLimitExceeded", + "Message": "Transform job quota exceeded", + } + }, + operation_name="CreateTransformJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_terminal_job_error(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ValidationException", + "Message": "The request was rejected because the transform job is not in a non-running state", + } + }, + operation_name="StopTransformJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + metadata = SageMakerTransformJobMetadata( + config={"TransformJobName": "score-x"}, region=REGION + ) + + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_batch_transform.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceNotFound", + "Message": "Transform job does not exist", + } + }, + operation_name="StopTransformJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-transform-job") + metadata = SageMakerTransformJobMetadata( + config={"TransformJobName": "score-x"}, region=REGION + ) + assert await connector.delete(metadata) is None diff --git a/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_task.py b/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_task.py new file mode 100644 index 0000000000..13cb3c73d7 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_batch_transform_task.py @@ -0,0 +1,70 @@ +import pytest +from flytekitplugins.awssagemaker_batch_transform import ( + SageMakerDescribeTransformJobTask, + SageMakerStopTransformJobTask, + SageMakerTransformJobTask, +) + +from flytekit import kwtypes +from flytekit.configuration import Image, ImageConfig, SerializationSettings + + +def _ser_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return SerializationSettings( + project="project", + domain="domain", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + +def test_transform_job_task_interface_and_custom(): + task = SageMakerTransformJobTask( + name="batch_score", + config={ + "TransformJobName": "score-{idempotence_token}", + "ModelName": "{inputs.model_name}", + "TransformInput": { + "DataSource": { + "S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": "{inputs.input_data}"} + }, + "ContentType": "text/csv", + "SplitType": "Line", + }, + "TransformOutput": {"S3OutputPath": "{inputs.output_prefix}"}, + "TransformResources": {"InstanceType": "ml.m5.xlarge", "InstanceCount": 1}, + }, + region="us-east-2", + inputs=kwtypes(model_name=str, input_data=str, output_prefix=str), + ) + + assert len(task.interface.inputs) == 3 + assert len(task.interface.outputs) == 1 + assert "result" in task.interface.outputs + + custom = task.get_custom(_ser_settings()) + assert custom["region"] == "us-east-2" + assert custom["config"]["TransformJobName"] == "score-{idempotence_token}" + + +@pytest.mark.parametrize( + "task_cls,method", + [ + (SageMakerStopTransformJobTask, "stop_transform_job"), + (SageMakerDescribeTransformJobTask, "describe_transform_job"), + ], +) +def test_helper_boto_tasks_use_correct_method(task_cls, method): + task = task_cls( + name="helper", + config={"TransformJobName": "{inputs.transform_job_name}"}, + region="us-east-2", + inputs=kwtypes(transform_job_name=str), + ) + + custom = task.get_custom(_ser_settings()) + assert custom["service"] == "sagemaker" + assert custom["method"] == method + assert custom["region"] == "us-east-2" diff --git a/plugins/flytekit-aws-sagemaker/tests/test_boto3_mixin.py b/plugins/flytekit-aws-sagemaker/tests/test_boto3_mixin.py index 39f91f32a2..537a5bead9 100644 --- a/plugins/flytekit-aws-sagemaker/tests/test_boto3_mixin.py +++ b/plugins/flytekit-aws-sagemaker/tests/test_boto3_mixin.py @@ -5,7 +5,9 @@ from flytekitplugins.awssagemaker_inference import triton_image_uri from flytekitplugins.awssagemaker_inference.boto3_mixin import ( Boto3ConnectorMixin, + convert_floats_with_no_fraction_to_ints, format_dict, + sorted_dict_str, ) from flytekit import FlyteContext, StructuredDataset @@ -245,3 +247,97 @@ async def test_call_with_truncated_idempotence_token_as_input(mock_session): assert result == mock_method.return_value assert idempotence_token == "ce735d6a183643f1" + + +def test_convert_floats_with_no_fraction_to_ints_recursive(): + """Whole-number floats become ints; non-whole floats and other types are left alone.""" + data = { + "InstanceCount": 1.0, + "MaxRuntimeInSeconds": 3600.0, + "ResourceConfig": {"VolumeSizeInGB": 30.0, "Ratio": 0.75}, + "InstanceGroups": [ + {"InstanceCount": 2.0, "InstanceType": "ml.m5.xlarge"}, + ], + } + + result = convert_floats_with_no_fraction_to_ints(data) + + assert result["InstanceCount"] == 1 + assert isinstance(result["InstanceCount"], int) + assert result["MaxRuntimeInSeconds"] == 3600 + assert isinstance(result["MaxRuntimeInSeconds"], int) + assert result["ResourceConfig"]["VolumeSizeInGB"] == 30 + assert isinstance(result["ResourceConfig"]["VolumeSizeInGB"], int) + # Non-whole float is preserved. + assert result["ResourceConfig"]["Ratio"] == 0.75 + assert isinstance(result["ResourceConfig"]["Ratio"], float) + # Lists recurse. + assert result["InstanceGroups"][0]["InstanceCount"] == 2 + assert isinstance(result["InstanceGroups"][0]["InstanceCount"], int) + assert result["InstanceGroups"][0]["InstanceType"] == "ml.m5.xlarge" + + +def test_sorted_dict_str_preserves_semantically_significant_list_order(): + first = {"ContainerArguments": ["--mode", "train"]} + second = {"ContainerArguments": ["train", "--mode"]} + + assert sorted_dict_str(first) != sorted_dict_str(second) + + +@pytest.mark.asyncio +@patch("flytekitplugins.awssagemaker_inference.boto3_mixin.aioboto3.Session") +async def test_call_normalises_whole_number_floats_to_ints(mock_session): + """Regression: the async path now applies the float->int conversion that + was previously only run by the sync BotoConnector. Without this, configs like + ``InstanceCount: 1.0`` (e.g. from JSON-decoded inputs) get rejected by boto3.""" + mixin = Boto3ConnectorMixin(service="sagemaker", region="us-east-1") + + mock_client = AsyncMock() + mock_session.return_value.client.return_value.__aenter__.return_value = mock_client + mock_method = mock_client.create_training_job + + config = { + "TrainingJobName": "t", + "ResourceConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1.0, + "VolumeSizeInGB": 30.0, + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600.0}, + } + + await mixin._call(method="create_training_job", config=config) + + mock_method.assert_called_with( + TrainingJobName="t", + ResourceConfig={ + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + }, + StoppingCondition={"MaxRuntimeInSeconds": 3600}, + ) + + +@pytest.mark.asyncio +@patch("flytekitplugins.awssagemaker_inference.boto3_mixin.aioboto3.Session") +async def test_call_hashes_normalised_integer_values_consistently(mock_session): + mixin = Boto3ConnectorMixin(service="sagemaker", region="us-east-1") + + mock_client = AsyncMock() + mock_session.return_value.client.return_value.__aenter__.return_value = mock_client + mock_client.create_training_job.return_value = {} + + float_config = { + "TrainingJobName": "train-{idempotence_token}", + "ResourceConfig": {"InstanceCount": 1.0}, + } + int_config = { + "TrainingJobName": "train-{idempotence_token}", + "ResourceConfig": {"InstanceCount": 1}, + } + + _, float_token = await mixin._call(method="create_training_job", config=float_config) + _, int_token = await mixin._call(method="create_training_job", config=int_config) + + assert float_token == int_token diff --git a/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_connector.py new file mode 100644 index 0000000000..eed0e25c7c --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_connector.py @@ -0,0 +1,560 @@ +from datetime import datetime, timedelta +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_hyperparameter_tuning.connector import ( + SageMakerHyperParameterTuningJobMetadata, +) +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException + +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.interfaces.cli_identifiers import Identifier +from flytekit.models import literals +from flytekit.models.core.identifier import ResourceType +from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate + +idempotence_token = "74443947857331f7" + +REGION = "us-east-2" +TUNING_JOB_NAME = "xgb-tune-{idempotence_token}" +TUNING_JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:hyper-parameter-tuning-job/" + "xgb-tune-74443947857331f7" +) +BEST_TRAINING_JOB_NAME = "xgb-tune-74443947857331f7-007-3f4a5b6c" +BEST_TRAINING_JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:training-job/" + "xgb-tune-74443947857331f7-007-3f4a5b6c" +) +S3_MODEL_ARTIFACTS = ( + "s3://my-bucket/output/xgb-tune-74443947857331f7-007-3f4a5b6c/output/model.tar.gz" +) + + +def _task_config(): + return { + "config": { + "HyperParameterTuningJobName": TUNING_JOB_NAME, + "HyperParameterTuningJobConfig": { + "Strategy": "Bayesian", + "HyperParameterTuningJobObjective": { + "Type": "Maximize", + "MetricName": "validation:auc", + }, + "ResourceLimits": { + "MaxNumberOfTrainingJobs": 4, + "MaxParallelTrainingJobs": 2, + }, + "ParameterRanges": { + "ContinuousParameterRanges": [ + {"Name": "eta", "MinValue": "0.01", "MaxValue": "0.5"}, + ], + "IntegerParameterRanges": [ + {"Name": "max_depth", "MinValue": "3", "MaxValue": "9"}, + ], + }, + }, + "TrainingJobDefinition": { + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + "MetricDefinitions": [ + {"Name": "validation:auc", "Regex": "auc=([0-9\\.]+)"}, + ], + }, + "RoleArn": "{inputs.execution_role_arn}", + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "ResourceConfig": { + "InstanceType": "ml.m5.large", + "InstanceCount": 1, + "VolumeSizeInGB": 10, + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 1800}, + }, + }, + "region": REGION, + "images": { + "training_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/xgboost:latest" + }, + } + + +def _task_template(): + task_id = Identifier( + resource_type=ResourceType.TASK, + project="project", + domain="domain", + name="name", + version="version", + ) + task_metadata = TaskMetadata( + discoverable=True, + runtime=RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), + timeout=timedelta(days=1), + retries=literals.RetryStrategy(3), + interruptible=True, + discovery_version="0.1.1b0", + deprecated_error_message="This is deprecated!", + cache_serializable=True, + pod_template_name="A", + cache_ignore_input_vars=(), + ) + return TaskTemplate( + id=task_id, + custom=_task_config(), + metadata=task_metadata, + interface=None, + type="sagemaker-hyperparameter-tuning-job", + ) + + +def _completed_describe_tuning_response(): + return { + "HyperParameterTuningJobName": "xgb-tune-74443947857331f7", + "HyperParameterTuningJobArn": TUNING_JOB_ARN, + "HyperParameterTuningJobStatus": "Completed", + "TrainingJobStatusCounters": { + "Completed": 4, + "InProgress": 0, + "RetryableError": 0, + "NonRetryableError": 0, + "Stopped": 0, + }, + "ObjectiveStatusCounters": {"Succeeded": 4, "Pending": 0, "Failed": 0}, + "BestTrainingJob": { + "TrainingJobName": BEST_TRAINING_JOB_NAME, + "TrainingJobArn": BEST_TRAINING_JOB_ARN, + "TrainingJobStatus": "Completed", + "ObjectiveStatus": "Succeeded", + "FinalHyperParameterTuningJobObjectiveMetric": { + "MetricName": "validation:auc", + "Value": 0.93, + }, + "TunedHyperParameters": {"eta": "0.21", "max_depth": "7"}, + "TrainingStartTime": datetime(2026, 4, 30, 12, 0, 0), + "TrainingEndTime": datetime(2026, 4, 30, 12, 8, 0), + }, + } + + +def _describe_training_response(): + return { + "TrainingJobName": BEST_TRAINING_JOB_NAME, + "TrainingJobArn": BEST_TRAINING_JOB_ARN, + "TrainingJobStatus": "Completed", + "ModelArtifacts": {"S3ModelArtifacts": S3_MODEL_ARTIFACTS}, + } + + +def _routing_side_effect(method_to_response): + """Build a side_effect that dispatches on the ``method`` kwarg of _call. + + The HPO connector's get() makes up to two boto3 calls per poll: + describe_hyper_parameter_tuning_job, and on Completed, describe_training_job. + Tests need to return the right payload for each. + """ + + async def _side_effect(*args, **kwargs): + method = kwargs.get("method") + if method not in method_to_response: + raise AssertionError(f"unexpected _call(method={method!r})") + return method_to_response[method] + + return _side_effect + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_create_get_delete_happy_path(mock_call): + mock_call.side_effect = _routing_side_effect( + { + "create_hyper_parameter_tuning_job": (None, idempotence_token), + "describe_hyper_parameter_tuning_job": ( + _completed_describe_tuning_response(), + idempotence_token, + ), + "describe_training_job": ( + _describe_training_response(), + idempotence_token, + ), + "stop_hyper_parameter_tuning_job": (None, idempotence_token), + } + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config=_task_config()["config"], region=REGION + ) + + response = await connector.create(_task_template()) + assert response == metadata + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["HyperParameterTuningJobArn"] == TUNING_JOB_ARN + assert result["HyperParameterTuningJobName"] == "xgb-tune-74443947857331f7" + + # BestTrainingJob carries the metric, tuned params, and — crucially — the + # S3ModelArtifacts resolved via the follow-up describe_training_job call. + best = result["BestTrainingJob"] + assert best["TrainingJobName"] == BEST_TRAINING_JOB_NAME + assert best["TrainingJobArn"] == BEST_TRAINING_JOB_ARN + assert best["ObjectiveStatus"] == "Succeeded" + assert best["FinalHyperParameterTuningJobObjectiveMetric"] == { + "MetricName": "validation:auc", + "Value": 0.93, + } + assert best["TunedHyperParameters"] == {"eta": "0.21", "max_depth": "7"} + assert best["ModelArtifacts"] == {"S3ModelArtifacts": S3_MODEL_ARTIFACTS} + assert best["TrainingStartTime"] == "2026-04-30T12:00:00" + assert best["TrainingEndTime"] == "2026-04-30T12:08:00" + + # Top-level convenience copy of the best artifacts so callers can consume + # `result["ModelArtifacts"]["S3ModelArtifacts"]` symmetric with training. + assert result["ModelArtifacts"] == {"S3ModelArtifacts": S3_MODEL_ARTIFACTS} + + assert result["TrainingJobStatusCounters"]["Completed"] == 4 + assert result["ObjectiveStatusCounters"] == { + "Succeeded": 4, + "Pending": 0, + "Failed": 0, + } + + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_get_inprogress_surfaces_counter_summary(mock_call): + """While running we surface a compact 'N Completed / M InProgress / K Failed' line.""" + mock_call.side_effect = _routing_side_effect( + { + "describe_hyper_parameter_tuning_job": ( + { + "HyperParameterTuningJobName": "xgb-tune-x", + "HyperParameterTuningJobArn": TUNING_JOB_ARN, + "HyperParameterTuningJobStatus": "InProgress", + "TrainingJobStatusCounters": { + "Completed": 3, + "InProgress": 1, + "RetryableError": 0, + "NonRetryableError": 1, + "Stopped": 0, + }, + }, + idempotence_token, + ), + } + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.message == "3 Completed / 1 InProgress / 1 Failed trials" + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_get_failed_surfaces_failure_reason(mock_call): + mock_call.side_effect = _routing_side_effect( + { + "describe_hyper_parameter_tuning_job": ( + { + "HyperParameterTuningJobName": "xgb-tune-x", + "HyperParameterTuningJobArn": TUNING_JOB_ARN, + "HyperParameterTuningJobStatus": "Failed", + "FailureReason": "All trials failed with ClientError", + }, + idempotence_token, + ), + } + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "All trials failed with ClientError" + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_get_completed_propagates_describe_training_failure(mock_call): + """The promised best-model artifact must not silently become ``None``.""" + + async def _side_effect(*args, **kwargs): + method = kwargs.get("method") + if method == "describe_hyper_parameter_tuning_job": + return (_completed_describe_tuning_response(), idempotence_token) + if method == "describe_training_job": + raise CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "AccessDeniedException", + "Message": "secondary describe blocked", + } + }, + operation_name="DescribeTrainingJob", + ), + ) + raise AssertionError(f"unexpected _call(method={method!r})") + + mock_call.side_effect = _side_effect + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + with pytest.raises(CustomException): + await connector.get(metadata) + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_get_completed_without_best_training_job(mock_call): + """Edge case: completed tuning with no BestTrainingJob (every trial failed + its objective). We still return a structured result with None artifacts.""" + mock_call.side_effect = _routing_side_effect( + { + "describe_hyper_parameter_tuning_job": ( + { + "HyperParameterTuningJobName": "xgb-tune-x", + "HyperParameterTuningJobArn": TUNING_JOB_ARN, + "HyperParameterTuningJobStatus": "Completed", + "TrainingJobStatusCounters": { + "Completed": 4, + "InProgress": 0, + "RetryableError": 0, + "NonRetryableError": 0, + "Stopped": 0, + }, + "ObjectiveStatusCounters": { + "Succeeded": 0, + "Pending": 0, + "Failed": 4, + }, + }, + idempotence_token, + ), + } + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["BestTrainingJob"]["TrainingJobName"] is None + assert result["BestTrainingJob"]["ModelArtifacts"] == {"S3ModelArtifacts": None} + assert result["ModelArtifacts"] == {"S3ModelArtifacts": None} + assert result["ObjectiveStatusCounters"]["Failed"] == 4 + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_get_stopped_maps_to_failed(mock_call): + mock_call.side_effect = _routing_side_effect( + { + "describe_hyper_parameter_tuning_job": ( + { + "HyperParameterTuningJobName": "xgb-tune-x", + "HyperParameterTuningJobArn": TUNING_JOB_ARN, + "HyperParameterTuningJobStatus": "Stopped", + "FailureReason": "User requested stop", + }, + idempotence_token, + ), + } + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "User requested stop" + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_create_already_exists_returns_metadata(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Hyperparameter tuning job xgb-tune-74443947857331f7 already exists", + } + }, + operation_name="CreateHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + response = await connector.create(_task_template()) + assert response.config == _task_config()["config"] + assert response.region == REGION + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_create_resource_limit_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceLimitExceeded", + "Message": ( + "The account-level service limit ... has been reached. " + "Please use AWS Service Quotas to request an increase for this quota." + ), + } + }, + operation_name="CreateHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_create_unknown_error_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="CreateHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_delete_swallows_terminal_job_error(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ValidationException", + "Message": ( + "The request was rejected because the hyperparameter " + "tuning job is not in a non-running state" + ), + } + }, + operation_name="StopHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceNotFound", + "Message": "Hyperparameter tuning job does not exist", + } + }, + operation_name="StopHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_hyperparameter_tuning.connector.Boto3ConnectorMixin._call" +) +async def test_delete_propagates_other_errors(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="StopHyperParameterTuningJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-hyperparameter-tuning-job") + metadata = SageMakerHyperParameterTuningJobMetadata( + config={"HyperParameterTuningJobName": "xgb-tune-x"}, region=REGION + ) + with pytest.raises(CustomException): + await connector.delete(metadata) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_task.py b/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_task.py new file mode 100644 index 0000000000..d1c4bf807c --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_hyperparameter_tuning_task.py @@ -0,0 +1,95 @@ +import pytest +from flytekitplugins.awssagemaker_hyperparameter_tuning import ( + SageMakerDescribeHyperParameterTuningJobTask, + SageMakerHyperParameterTuningJobTask, + SageMakerStopHyperParameterTuningJobTask, +) + +from flytekit import kwtypes +from flytekit.configuration import Image, ImageConfig, SerializationSettings + + +def _ser_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return SerializationSettings( + project="project", + domain="domain", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + +def test_hyperparameter_tuning_job_task_interface_and_custom(): + task = SageMakerHyperParameterTuningJobTask( + name="tune_xgb", + config={ + "HyperParameterTuningJobName": "xgb-tune-{idempotence_token}", + "HyperParameterTuningJobConfig": { + "Strategy": "Bayesian", + "HyperParameterTuningJobObjective": { + "Type": "Maximize", + "MetricName": "validation:auc", + }, + "ResourceLimits": { + "MaxNumberOfTrainingJobs": 4, + "MaxParallelTrainingJobs": 2, + }, + "ParameterRanges": { + "ContinuousParameterRanges": [ + {"Name": "eta", "MinValue": "0.01", "MaxValue": "0.5"}, + ], + }, + }, + "TrainingJobDefinition": { + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + }, + "RoleArn": "{inputs.execution_role_arn}", + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "ResourceConfig": { + "InstanceType": "ml.m5.large", + "InstanceCount": 1, + "VolumeSizeInGB": 10, + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 1800}, + }, + }, + region="us-east-2", + images={"training_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/xgb:latest"}, + inputs=kwtypes(execution_role_arn=str, output_prefix=str), + ) + + assert len(task.interface.inputs) == 2 + assert len(task.interface.outputs) == 1 + assert "result" in task.interface.outputs + + custom = task.get_custom(_ser_settings()) + assert custom["region"] == "us-east-2" + assert custom["config"]["HyperParameterTuningJobName"] == "xgb-tune-{idempotence_token}" + assert custom["images"]["training_image"].endswith("/xgb:latest") + + +@pytest.mark.parametrize( + "task_cls,method", + [ + (SageMakerStopHyperParameterTuningJobTask, "stop_hyper_parameter_tuning_job"), + ( + SageMakerDescribeHyperParameterTuningJobTask, + "describe_hyper_parameter_tuning_job", + ), + ], +) +def test_helper_boto_tasks_use_correct_method(task_cls, method): + task = task_cls( + name="helper", + config={"HyperParameterTuningJobName": "{inputs.tuning_job_name}"}, + region="us-east-2", + inputs=kwtypes(tuning_job_name=str), + ) + + custom = task.get_custom(_ser_settings()) + assert custom["service"] == "sagemaker" + assert custom["method"] == method + assert custom["region"] == "us-east-2" diff --git a/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_connector.py new file mode 100644 index 0000000000..999de15269 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_connector.py @@ -0,0 +1,346 @@ +from datetime import datetime, timedelta +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException +from flytekitplugins.awssagemaker_inference_recommender.connector import ( + SageMakerInferenceRecommenderJobMetadata, +) + +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.interfaces.cli_identifiers import Identifier +from flytekit.models import literals +from flytekit.models.core.identifier import ResourceType +from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate + +idempotence_token = "74443947857331f7" + +REGION = "us-east-2" +JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:inference-recommendations-job/" + "rec-74443947857331f7" +) + + +def _task_config(): + return { + "config": { + "JobName": "rec-{idempotence_token}", + "JobType": "Default", + "JobDescription": "Smoke recommendations for ranker-prod", + "RoleArn": "{inputs.role_arn}", + # Default-job InputConfig allows only ModelPackageVersionArn (or + # ModelName + ContainerConfig). JobDurationInSeconds / + # TrafficPattern / ResourceLimit / EndpointConfigurations and the + # top-level StoppingConditions are all Advanced-only — AWS rejects + # them with a ValidationException if set here. + "InputConfig": { + "ModelPackageVersionArn": "{inputs.model_package_version_arn}", + }, + }, + "region": REGION, + } + + +def _task_template(): + task_id = Identifier( + resource_type=ResourceType.TASK, + project="project", + domain="domain", + name="name", + version="version", + ) + task_metadata = TaskMetadata( + discoverable=True, + runtime=RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), + timeout=timedelta(days=1), + retries=literals.RetryStrategy(3), + interruptible=True, + discovery_version="0.1.1b0", + deprecated_error_message="This is deprecated!", + cache_serializable=True, + pod_template_name="A", + cache_ignore_input_vars=(), + ) + return TaskTemplate( + id=task_id, + custom=_task_config(), + metadata=task_metadata, + interface=None, + type="sagemaker-inference-recommender-job", + ) + + +def _completed_describe_response(): + return { + "JobName": "rec-74443947857331f7", + "JobArn": JOB_ARN, + "JobType": "Default", + "Status": "COMPLETED", + "CompletionTime": datetime(2026, 4, 30, 10, 45, 0), + "InferenceRecommendations": [ + { + "RecommendationId": "rec-74443947857331f7/1", + "Metrics": { + "CostPerHour": 0.42, + "CostPerInference": 0.0000012, + "MaxInvocations": 1200, + "ModelLatency": 38, + "CpuUtilization": 71.4, + "MemoryUtilization": 55.2, + "ModelSetupTime": 17, + }, + "EndpointConfiguration": { + "EndpointName": "sm-epc-1", + "VariantName": "AllTraffic", + "InstanceType": "ml.m5.xlarge", + "InitialInstanceCount": 1, + }, + "ModelConfiguration": { + "InferenceSpecificationName": "default", + "CompilationJobName": None, + }, + "InvocationStartTime": datetime(2026, 4, 30, 10, 5, 0), + "InvocationEndTime": datetime(2026, 4, 30, 10, 15, 0), + } + ], + "EndpointPerformances": [], + } + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_create_get_delete_happy_path(mock_call): + mock_call.return_value = (_completed_describe_response(), idempotence_token) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config=_task_config()["config"], region=REGION + ) + + response = await connector.create(_task_template()) + assert response == metadata + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["JobArn"] == JOB_ARN + assert result["JobName"] == "rec-74443947857331f7" + assert result["JobType"] == "Default" + assert result["CompletionTime"] == "2026-04-30T10:45:00" + + assert len(result["InferenceRecommendations"]) == 1 + top = result["InferenceRecommendations"][0] + assert top["RecommendationId"] == "rec-74443947857331f7/1" + assert top["EndpointConfiguration"]["InstanceType"] == "ml.m5.xlarge" + assert top["EndpointConfiguration"]["InitialInstanceCount"] == 1 + assert top["Metrics"]["CostPerHour"] == 0.42 + assert top["Metrics"]["ModelLatency"] == 38 + assert top["InvocationStartTime"] == "2026-04-30T10:05:00" + assert top["InvocationEndTime"] == "2026-04-30T10:15:00" + assert result["EndpointPerformances"] == [] + + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_get_pending_and_inprogress_map_to_running(mock_call): + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config={"JobName": "rec-x"}, region=REGION + ) + + mock_call.return_value = ( + {"JobName": "rec-x", "JobArn": JOB_ARN, "Status": "PENDING"}, + idempotence_token, + ) + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.message is None + assert resource.outputs is None + + mock_call.return_value = ( + {"JobName": "rec-x", "JobArn": JOB_ARN, "Status": "IN_PROGRESS"}, + idempotence_token, + ) + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.message is None + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_get_failed_surfaces_failure_reason(mock_call): + mock_call.return_value = ( + { + "JobName": "rec-x", + "JobArn": JOB_ARN, + "Status": "FAILED", + "FailureReason": "Model failed to load on ml.m5.large", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config={"JobName": "rec-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "Model failed to load on ml.m5.large" + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_create_already_exists_returns_metadata(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Inference recommendations job rec-74443947857331f7 already exists", + } + }, + operation_name="CreateInferenceRecommendationsJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + response = await connector.create(_task_template()) + assert response.config == _task_config()["config"] + assert response.region == REGION + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_create_resource_limit_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceLimitExceeded", + "Message": "Inference Recommender job quota exceeded", + } + }, + operation_name="CreateInferenceRecommendationsJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_delete_swallows_terminal_job_error(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ValidationException", + "Message": ( + "The request was rejected because the inference " + "recommendations job is not in a non-running state" + ), + } + }, + operation_name="StopInferenceRecommendationsJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config={"JobName": "rec-x"}, region=REGION + ) + + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceNotFound", + "Message": "Inference recommendations job does not exist", + } + }, + operation_name="StopInferenceRecommendationsJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config={"JobName": "rec-x"}, region=REGION + ) + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch( + "flytekitplugins.awssagemaker_inference_recommender.connector.Boto3ConnectorMixin._call" +) +async def test_get_existing_endpoint_job_emits_endpoint_performances(mock_call): + """Default jobs can benchmark existing endpoints and report their performance.""" + mock_call.return_value = ( + { + "JobName": "rec-existing-endpoint", + "JobArn": JOB_ARN, + "JobType": "Default", + "Status": "COMPLETED", + "InferenceRecommendations": [], + "EndpointPerformances": [ + { + "Metrics": {"MaxInvocations": 800, "ModelLatency": 52}, + "EndpointInfo": {"EndpointName": "ranker-prod-canary"}, + } + ], + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-inference-recommender-job") + metadata = SageMakerInferenceRecommenderJobMetadata( + config={"JobName": "rec-existing-endpoint"}, region=REGION + ) + resource = await connector.get(metadata) + result = resource.outputs["result"] + assert result["JobType"] == "Default" + assert result["InferenceRecommendations"] == [] + assert result["EndpointPerformances"] == [ + { + "Metrics": {"MaxInvocations": 800, "ModelLatency": 52}, + "EndpointInfo": {"EndpointName": "ranker-prod-canary"}, + } + ] diff --git a/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_task.py b/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_task.py new file mode 100644 index 0000000000..77d316f95c --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_inference_recommender_task.py @@ -0,0 +1,72 @@ +import pytest +from flytekitplugins.awssagemaker_inference_recommender import ( + SageMakerDescribeInferenceRecommenderJobTask, + SageMakerInferenceRecommenderJobTask, + SageMakerStopInferenceRecommenderJobTask, +) + +from flytekit import kwtypes +from flytekit.configuration import Image, ImageConfig, SerializationSettings + + +def _ser_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return SerializationSettings( + project="project", + domain="domain", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + +def test_inference_recommender_job_task_interface_and_custom(): + task = SageMakerInferenceRecommenderJobTask( + name="recommend", + config={ + "JobName": "rec-{idempotence_token}", + "JobType": "Default", + "RoleArn": "{inputs.role_arn}", + # Minimal valid Default-job InputConfig — AWS rejects + # JobDurationInSeconds / TrafficPattern / ResourceLimit / + # EndpointConfigurations / top-level StoppingConditions for Default. + "InputConfig": { + "ModelPackageVersionArn": "{inputs.model_package_version_arn}", + }, + }, + region="us-east-2", + inputs=kwtypes(role_arn=str, model_package_version_arn=str), + ) + + assert len(task.interface.inputs) == 2 + assert len(task.interface.outputs) == 1 + assert "result" in task.interface.outputs + + custom = task.get_custom(_ser_settings()) + assert custom["region"] == "us-east-2" + assert custom["config"]["JobName"] == "rec-{idempotence_token}" + assert custom["config"]["JobType"] == "Default" + + +@pytest.mark.parametrize( + "task_cls,method", + [ + (SageMakerStopInferenceRecommenderJobTask, "stop_inference_recommendations_job"), + ( + SageMakerDescribeInferenceRecommenderJobTask, + "describe_inference_recommendations_job", + ), + ], +) +def test_helper_boto_tasks_use_correct_method(task_cls, method): + task = task_cls( + name="helper", + config={"JobName": "{inputs.job_name}"}, + region="us-east-2", + inputs=kwtypes(job_name=str), + ) + + custom = task.get_custom(_ser_settings()) + assert custom["service"] == "sagemaker" + assert custom["method"] == method + assert custom["region"] == "us-east-2" diff --git a/plugins/flytekit-aws-sagemaker/tests/test_processing_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_processing_connector.py new file mode 100644 index 0000000000..6a2ca0f9bc --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_processing_connector.py @@ -0,0 +1,398 @@ +from datetime import datetime, timedelta +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException +from flytekitplugins.awssagemaker_processing.connector import ( + SageMakerProcessingJobMetadata, + _build_outputs, +) + +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.interfaces.cli_identifiers import Identifier +from flytekit.models import literals +from flytekit.models.core.identifier import ResourceType +from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate + +idempotence_token = "74443947857331f7" + +REGION = "us-east-2" +PROCESSING_JOB_NAME = "prep-{idempotence_token}" +PROCESSING_JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:processing-job/prep-74443947857331f7" +) +S3_OUTPUT = "s3://my-bucket/processing/prep-74443947857331f7/output/train" + + +def _task_config(): + return { + "config": { + "ProcessingJobName": PROCESSING_JOB_NAME, + "AppSpecification": { + "ImageUri": "{images.processing_image}", + "ContainerEntrypoint": ["python3", "/opt/ml/processing/preprocess.py"], + }, + "RoleArn": "{inputs.execution_role_arn}", + "ProcessingResources": { + "ClusterConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + } + }, + "ProcessingOutputConfig": { + "Outputs": [ + { + "OutputName": "train", + "S3Output": { + "S3Uri": "{inputs.output_prefix}", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + }, + } + ] + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + "region": REGION, + "images": {"processing_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/sklearn:latest"}, + } + + +def _task_template(): + task_id = Identifier( + resource_type=ResourceType.TASK, + project="project", + domain="domain", + name="name", + version="version", + ) + task_metadata = TaskMetadata( + discoverable=True, + runtime=RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), + timeout=timedelta(days=1), + retries=literals.RetryStrategy(3), + interruptible=True, + discovery_version="0.1.1b0", + deprecated_error_message="This is deprecated!", + cache_serializable=True, + pod_template_name="A", + cache_ignore_input_vars=(), + ) + return TaskTemplate( + id=task_id, + custom=_task_config(), + metadata=task_metadata, + interface=None, + type="sagemaker-processing-job", + ) + + +def _completed_describe_response(): + return { + "ProcessingJobName": "prep-74443947857331f7", + "ProcessingJobArn": PROCESSING_JOB_ARN, + "ProcessingJobStatus": "Completed", + "ProcessingOutputConfig": { + "Outputs": [ + { + "OutputName": "train", + "S3Output": {"S3Uri": S3_OUTPUT, "LocalPath": "/opt/ml/processing/output"}, + } + ] + }, + "ExitMessage": "Completed: Job completed successfully", + "ProcessingStartTime": datetime(2026, 4, 30, 12, 0, 0), + "ProcessingEndTime": datetime(2026, 4, 30, 12, 5, 0), + } + + +def test_build_outputs_preserves_feature_store_destination(): + result = _build_outputs( + { + "ProcessingOutputConfig": { + "Outputs": [ + { + "OutputName": "features", + "FeatureStoreOutput": {"FeatureGroupName": "customer-features"}, + } + ] + } + } + ) + + assert result["Outputs"] == [ + { + "OutputName": "features", + "FeatureGroupName": "customer-features", + } + ] + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_create_get_delete_happy_path(mock_call): + mock_call.return_value = (_completed_describe_response(), idempotence_token) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config=_task_config()["config"], region=REGION + ) + + # CREATE — returns metadata; mock return value is ignored by create(). + response = await connector.create(_task_template()) + assert response == metadata + + # GET — parses describe response, returns Completed with structured outputs. + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["ProcessingJobArn"] == PROCESSING_JOB_ARN + assert result["ProcessingJobName"] == "prep-74443947857331f7" + assert result["Outputs"] == [{"OutputName": "train", "S3Uri": S3_OUTPUT}] + assert result["ExitMessage"] == "Completed: Job completed successfully" + + # Timestamps must be ISO strings (datetime is not JSON-friendly). + assert result["ProcessingStartTime"] == "2026-04-30T12:00:00" + assert result["ProcessingEndTime"] == "2026-04-30T12:05:00" + + # DELETE — happy path returns None. + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_get_inprogress_has_no_outputs(mock_call): + mock_call.return_value = ( + { + "ProcessingJobName": "prep-x", + "ProcessingJobArn": PROCESSING_JOB_ARN, + "ProcessingJobStatus": "InProgress", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_get_failed_surfaces_failure_reason(mock_call): + mock_call.return_value = ( + { + "ProcessingJobName": "prep-x", + "ProcessingJobArn": PROCESSING_JOB_ARN, + "ProcessingJobStatus": "Failed", + "FailureReason": "AlgorithmError: script returned non-zero exit code", + "ExitMessage": "Traceback ...", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "AlgorithmError: script returned non-zero exit code" + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_get_failed_falls_back_to_exit_message(mock_call): + mock_call.return_value = ( + { + "ProcessingJobName": "prep-x", + "ProcessingJobArn": PROCESSING_JOB_ARN, + "ProcessingJobStatus": "Failed", + "ExitMessage": "Container exited with code 1", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "Container exited with code 1" + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_get_stopped_maps_to_failed(mock_call): + mock_call.return_value = ( + { + "ProcessingJobName": "prep-x", + "ProcessingJobArn": PROCESSING_JOB_ARN, + "ProcessingJobStatus": "Stopped", + "FailureReason": "MaxRuntimeExceeded", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "MaxRuntimeExceeded" + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_create_already_exists_returns_metadata(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Processing job prep-74443947857331f7 already exists", + } + }, + operation_name="CreateProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + response = await connector.create(_task_template()) + assert response.config == _task_config()["config"] + assert response.region == REGION + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_create_resource_limit_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceLimitExceeded", + "Message": ( + "The account-level service limit ... has been reached. " + "Please use AWS Service Quotas to request an increase for this quota." + ), + } + }, + operation_name="CreateProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_create_unknown_error_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="CreateProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_terminal_job_error(mock_call): + """If Flyte calls delete() after the job naturally finished, stop_processing_job + raises ValidationException — the connector must swallow that specific case.""" + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ValidationException", + "Message": "The request was rejected because the processing job is not in a non-running state", + } + }, + operation_name="StopProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + # Should NOT raise. + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceNotFound", + "Message": "Processing job does not exist", + } + }, + operation_name="StopProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_processing.connector.Boto3ConnectorMixin._call") +async def test_delete_propagates_other_errors(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="StopProcessingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-processing-job") + metadata = SageMakerProcessingJobMetadata( + config={"ProcessingJobName": "prep-x"}, region=REGION + ) + + with pytest.raises(CustomException): + await connector.delete(metadata) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_processing_task.py b/plugins/flytekit-aws-sagemaker/tests/test_processing_task.py new file mode 100644 index 0000000000..21aae20010 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_processing_task.py @@ -0,0 +1,87 @@ +import pytest +from flytekitplugins.awssagemaker_processing import ( + SageMakerDescribeProcessingJobTask, + SageMakerProcessingJobTask, + SageMakerStopProcessingJobTask, +) + +from flytekit import kwtypes +from flytekit.configuration import Image, ImageConfig, SerializationSettings + + +def _ser_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return SerializationSettings( + project="project", + domain="domain", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + +def test_processing_job_task_interface_and_custom(): + task = SageMakerProcessingJobTask( + name="preprocess", + config={ + "ProcessingJobName": "prep-{idempotence_token}", + "AppSpecification": { + "ImageUri": "{images.processing_image}", + "ContainerEntrypoint": ["python3", "/opt/ml/processing/preprocess.py"], + }, + "RoleArn": "{inputs.execution_role_arn}", + "ProcessingResources": { + "ClusterConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + } + }, + "ProcessingOutputConfig": { + "Outputs": [ + { + "OutputName": "train", + "S3Output": { + "S3Uri": "{inputs.output_prefix}", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + }, + } + ] + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + region="us-east-2", + images={"processing_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/sklearn:latest"}, + inputs=kwtypes(execution_role_arn=str, output_prefix=str), + ) + + assert len(task.interface.inputs) == 2 + assert len(task.interface.outputs) == 1 + assert "result" in task.interface.outputs + + custom = task.get_custom(_ser_settings()) + assert custom["region"] == "us-east-2" + assert custom["config"]["ProcessingJobName"] == "prep-{idempotence_token}" + assert custom["images"]["processing_image"].endswith("/sklearn:latest") + + +@pytest.mark.parametrize( + "task_cls,method", + [ + (SageMakerStopProcessingJobTask, "stop_processing_job"), + (SageMakerDescribeProcessingJobTask, "describe_processing_job"), + ], +) +def test_helper_boto_tasks_use_correct_method(task_cls, method): + task = task_cls( + name="helper", + config={"ProcessingJobName": "{inputs.processing_job_name}"}, + region="us-east-2", + inputs=kwtypes(processing_job_name=str), + ) + + custom = task.get_custom(_ser_settings()) + assert custom["service"] == "sagemaker" + assert custom["method"] == method + assert custom["region"] == "us-east-2" diff --git a/plugins/flytekit-aws-sagemaker/tests/test_training_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_training_connector.py new file mode 100644 index 0000000000..888aea8f17 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_training_connector.py @@ -0,0 +1,370 @@ +from datetime import datetime, timedelta +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException +from flytekitplugins.awssagemaker_training.connector import ( + SageMakerTrainingJobMetadata, +) + +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.interfaces.cli_identifiers import Identifier +from flytekit.models import literals +from flytekit.models.core.identifier import ResourceType +from flytekit.models.task import RuntimeMetadata, TaskMetadata, TaskTemplate + +idempotence_token = "74443947857331f7" + +REGION = "us-east-2" +TRAINING_JOB_NAME = "xgb-{idempotence_token}" +TRAINING_JOB_ARN = ( + "arn:aws:sagemaker:us-east-2:1234567890:training-job/xgb-74443947857331f7" +) +S3_MODEL_ARTIFACTS = "s3://my-bucket/output/xgb-74443947857331f7/output/model.tar.gz" + + +def _task_config(): + return { + "config": { + "TrainingJobName": TRAINING_JOB_NAME, + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + }, + "RoleArn": "{inputs.execution_role_arn}", + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "ResourceConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + }, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + "region": REGION, + "images": {"training_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/xgboost:latest"}, + } + + +def _task_template(): + task_id = Identifier( + resource_type=ResourceType.TASK, + project="project", + domain="domain", + name="name", + version="version", + ) + task_metadata = TaskMetadata( + discoverable=True, + runtime=RuntimeMetadata(RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), + timeout=timedelta(days=1), + retries=literals.RetryStrategy(3), + interruptible=True, + discovery_version="0.1.1b0", + deprecated_error_message="This is deprecated!", + cache_serializable=True, + pod_template_name="A", + cache_ignore_input_vars=(), + ) + return TaskTemplate( + id=task_id, + custom=_task_config(), + metadata=task_metadata, + interface=None, + type="sagemaker-training-job", + ) + + +def _completed_describe_response(): + return { + "TrainingJobName": "xgb-74443947857331f7", + "TrainingJobArn": TRAINING_JOB_ARN, + "TrainingJobStatus": "Completed", + "SecondaryStatus": "Completed", + "ModelArtifacts": {"S3ModelArtifacts": S3_MODEL_ARTIFACTS}, + "OutputDataConfig": {"S3OutputPath": "s3://my-bucket/output/"}, + "FinalMetricDataList": [ + { + "MetricName": "validation:auc", + "Value": 0.87, + "Timestamp": datetime(2026, 4, 30, 12, 0, 0), + } + ], + "BillableTimeInSeconds": 120, + "TrainingTimeInSeconds": 100, + } + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_create_get_delete_happy_path(mock_call): + mock_call.return_value = (_completed_describe_response(), idempotence_token) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config=_task_config()["config"], region=REGION + ) + + # CREATE — returns metadata; mock return value is ignored by create(). + response = await connector.create(_task_template()) + assert response == metadata + + # GET — parses describe response, returns Completed with structured outputs. + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.SUCCEEDED + + result = resource.outputs["result"] + assert result["TrainingJobArn"] == TRAINING_JOB_ARN + assert result["TrainingJobName"] == "xgb-74443947857331f7" + assert result["ModelArtifacts"] == {"S3ModelArtifacts": S3_MODEL_ARTIFACTS} + assert result["OutputDataConfig"] == {"S3OutputPath": "s3://my-bucket/output/"} + assert result["BillableTimeInSeconds"] == 120 + assert result["TrainingTimeInSeconds"] == 100 + + # FinalMetricDataList timestamps must be ISO strings (datetime is not JSON-friendly). + assert result["FinalMetricDataList"] == [ + { + "MetricName": "validation:auc", + "Value": 0.87, + "Timestamp": "2026-04-30T12:00:00", + } + ] + + # DELETE — happy path returns None. + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_get_inprogress_surfaces_secondary_status(mock_call): + mock_call.return_value = ( + { + "TrainingJobName": "xgb-x", + "TrainingJobArn": TRAINING_JOB_ARN, + "TrainingJobStatus": "InProgress", + "SecondaryStatus": "Downloading", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.RUNNING + assert resource.message == "Downloading" + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_get_failed_surfaces_failure_reason(mock_call): + mock_call.return_value = ( + { + "TrainingJobName": "xgb-x", + "TrainingJobArn": TRAINING_JOB_ARN, + "TrainingJobStatus": "Failed", + "FailureReason": "AlgorithmError: out of memory", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "AlgorithmError: out of memory" + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_get_stopped_maps_to_failed(mock_call): + mock_call.return_value = ( + { + "TrainingJobName": "xgb-x", + "TrainingJobArn": TRAINING_JOB_ARN, + "TrainingJobStatus": "Stopped", + "FailureReason": "MaxRuntimeExceeded", + }, + idempotence_token, + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + + resource = await connector.get(metadata) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "MaxRuntimeExceeded" + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_create_already_exists_returns_metadata(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Training job xgb-74443947857331f7 already exists", + } + }, + operation_name="CreateTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + response = await connector.create(_task_template()) + assert response.config == _task_config()["config"] + assert response.region == REGION + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_create_static_name_resource_in_use_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token="", + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceInUse", + "Message": "Training job static-name already exists", + } + }, + operation_name="CreateTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_create_resource_limit_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceLimitExceeded", + "Message": ( + "The account-level service limit ... has been reached. " + "Please use AWS Service Quotas to request an increase for this quota." + ), + } + }, + operation_name="CreateTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_create_unknown_error_propagates(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="CreateTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + with pytest.raises(CustomException): + await connector.create(_task_template()) + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_terminal_job_error(mock_call): + """If Flyte calls delete() after the job naturally finished, stop_training_job + raises ValidationException — the connector must swallow that specific case.""" + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ValidationException", + "Message": "The request was rejected because the training job is not in a non-running state", + } + }, + operation_name="StopTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + + # Should NOT raise. + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": { + "Code": "ResourceNotFound", + "Message": "Training job does not exist", + } + }, + operation_name="StopTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + assert await connector.delete(metadata) is None + + +@pytest.mark.asyncio +@mock.patch("flytekitplugins.awssagemaker_training.connector.Boto3ConnectorMixin._call") +async def test_delete_propagates_other_errors(mock_call): + mock_call.side_effect = CustomException( + message="An error occurred", + idempotence_token=idempotence_token, + original_exception=ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "nope"} + }, + operation_name="StopTrainingJob", + ), + ) + + connector = ConnectorRegistry.get_connector("sagemaker-training-job") + metadata = SageMakerTrainingJobMetadata( + config={"TrainingJobName": "xgb-x"}, region=REGION + ) + + with pytest.raises(CustomException): + await connector.delete(metadata) diff --git a/plugins/flytekit-aws-sagemaker/tests/test_training_task.py b/plugins/flytekit-aws-sagemaker/tests/test_training_task.py new file mode 100644 index 0000000000..4d24692094 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_training_task.py @@ -0,0 +1,74 @@ +import pytest +from flytekitplugins.awssagemaker_training import ( + SageMakerDescribeTrainingJobTask, + SageMakerStopTrainingJobTask, + SageMakerTrainingJobTask, +) + +from flytekit import kwtypes +from flytekit.configuration import Image, ImageConfig, SerializationSettings + + +def _ser_settings(): + default_img = Image(name="default", fqn="test", tag="tag") + return SerializationSettings( + project="project", + domain="domain", + version="123", + image_config=ImageConfig(default_image=default_img, images=[default_img]), + env={}, + ) + + +def test_training_job_task_interface_and_custom(): + task = SageMakerTrainingJobTask( + name="train_xgb", + config={ + "TrainingJobName": "xgb-{idempotence_token}", + "AlgorithmSpecification": { + "TrainingImage": "{images.training_image}", + "TrainingInputMode": "File", + }, + "RoleArn": "{inputs.execution_role_arn}", + "ResourceConfig": { + "InstanceType": "ml.m5.xlarge", + "InstanceCount": 1, + "VolumeSizeInGB": 30, + }, + "OutputDataConfig": {"S3OutputPath": "{inputs.output_prefix}"}, + "StoppingCondition": {"MaxRuntimeInSeconds": 3600}, + }, + region="us-east-2", + images={"training_image": "1234567890.dkr.ecr.us-east-2.amazonaws.com/xgb:latest"}, + inputs=kwtypes(execution_role_arn=str, output_prefix=str), + ) + + assert len(task.interface.inputs) == 2 + assert len(task.interface.outputs) == 1 + assert "result" in task.interface.outputs + + custom = task.get_custom(_ser_settings()) + assert custom["region"] == "us-east-2" + assert custom["config"]["TrainingJobName"] == "xgb-{idempotence_token}" + assert custom["images"]["training_image"].endswith("/xgb:latest") + + +@pytest.mark.parametrize( + "task_cls,method", + [ + (SageMakerStopTrainingJobTask, "stop_training_job"), + (SageMakerDescribeTrainingJobTask, "describe_training_job"), + ], +) +def test_helper_boto_tasks_use_correct_method(task_cls, method): + task = task_cls( + name="helper", + config={"TrainingJobName": "{inputs.training_job_name}"}, + region="us-east-2", + inputs=kwtypes(training_job_name=str), + ) + + custom = task.get_custom(_ser_settings()) + assert custom["service"] == "sagemaker" + assert custom["method"] == method + assert custom["region"] == "us-east-2"