diff --git a/flytekit/extend/backend/utils.py b/flytekit/extend/backend/utils.py index 9bcc654927..d0161fc327 100644 --- a/flytekit/extend/backend/utils.py +++ b/flytekit/extend/backend/utils.py @@ -37,7 +37,12 @@ def is_terminal_phase(phase: TaskExecution.Phase) -> bool: """ Return true if the phase is terminal. """ - return phase in [TaskExecution.SUCCEEDED, TaskExecution.ABORTED, TaskExecution.FAILED] + return phase in [ + TaskExecution.SUCCEEDED, + TaskExecution.ABORTED, + TaskExecution.FAILED, + TaskExecution.RETRYABLE_FAILED, + ] def get_connector_secret(secret_key: str) -> str: 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..bdbab16d76 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,607 @@ 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. + +## Pythonic Training and Processing + +Training and Processing also support a Flyte-native mode for code that is more +naturally expressed as a typed Python function than as a complete boto3 job +request. Use `SageMakerProcessing` or `SageMakerTraining` as the `task_config` +on a normal `@task`: + +```python +from flytekit import ImageSpec, task +from flytekitplugins.awssagemaker_processing import SageMakerProcessing +from flytekitplugins.awssagemaker_training import SageMakerTraining + +ROLE = "arn:aws:iam:::role/" +REGION = "us-east-1" + +# The registry must resolve to Amazon ECR. When base_image is omitted, Flytekit +# supplies its version-compatible default image before building and pushing. +image = ImageSpec( + name="sagemaker-pythonic", + registry=".dkr.ecr.us-east-1.amazonaws.com", + packages=["numpy"], +) + + +@task( + task_config=SageMakerProcessing( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.large", + ), + container_image=image, +) +def preprocess(values: list[float]) -> list[float]: + mean = sum(values) / len(values) + return [value - mean for value in values] + + +@task( + task_config=SageMakerTraining( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.xlarge", + ), + container_image=image, +) +def train(values: list[float]) -> float: + return sum(value * value for value in values) +``` + +The connector puts Flyte's rendered container arguments into SageMaker's +`ContainerEntrypoint`. Inside the SageMaker container, `pyflyte-execute` runs +the function and writes its typed result to Flyte's `outputs.pb`. User failures +are written to `error.pb`, preserve recoverable/non-recoverable semantics, and +fail the SageMaker job. + +Requirements and current constraints: + +- `container_image` is required. An `ImageSpec` is the simplest option; a plain + image URI must already contain a compatible Flytekit runtime and must be + available through a SageMaker-supported ECR registry. +- Pythonic jobs currently require `instance_count=1`. Running the same Flyte + function on every SageMaker host would duplicate side effects and race on + Flyte output files. +- The connector identity needs the relevant SageMaker lifecycle permissions and + `iam:PassRole`. The SageMaker execution role needs ECR pull, CloudWatch Logs, + and read/write access to Flyte's S3 input, fast-registration, and output + prefixes. +- Kubernetes-mounted secrets and Flyte pod environment injection are not + available inside SageMaker. Use the SageMaker execution role and an AWS secret + service for runtime credentials. Do not place secrets in `environment`; task + configuration and container environment values are serialized in the Flyte + task template. +- `SageMakerProcessing.network_config` accepts the boto3 `NetworkConfig` shape. + `EnableNetworkIsolation=True` is not supported because the Flyte entrypoint + must access S3. `SageMakerTraining.vpc_config` accepts the training-job + `VpcConfig` shape. +- Pythonic Training returns the function's typed Flyte output; it does not treat + SageMaker's generated `model.tar.gz` as the task result. Set `output_s3_path` + only when that SageMaker-side archive is also needed. + +## 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/pythonic_base.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/pythonic_base.py new file mode 100644 index 0000000000..4cd40a047b --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/pythonic_base.py @@ -0,0 +1,436 @@ +"""Shared foundation for SageMaker "Pythonic mode" job tasks and connectors. + +Pythonic mode lets users run a Flyte ``@task`` function body directly inside a +SageMaker job container (Processing or Training) instead of supplying a boto3 +config + prebuilt algorithm image. It mirrors ``flytekit-aws-emr-serverless`` +but is simpler: SageMaker runs the container directly via the job's +``ContainerEntrypoint``, so there is no S3 entrypoint shim to upload — the +connector points ``ContainerEntrypoint`` at the rendered Flyte command +(``task_template.container.args``) on the user's flytekit-containing image. + +The mechanics are identical for Processing and Training; the only per-service +differences are captured as overridable class attributes / methods on +``PythonicSageMakerJobConnector``: + +* ``create_method`` / ``describe_method`` / ``stop_method`` — boto3 method names +* ``job_name_key`` / ``status_key`` / ``secondary_status_key`` — response fields +* ``_state_map`` — service status -> Flyte phase +* ``_build_request`` — assemble the ``create_*_job`` request (``AppSpecification`` + vs ``AlgorithmSpecification``, ``ProcessingResources`` vs ``ResourceConfig``) + +Outputs flow back the Flyte-native way: the inner ``pyflyte-execute`` writes +``outputs.pb`` to the rendered ``--output-prefix`` and the connector returns +``outputs=None`` so flytekit materializes the typed return. Because SageMaker +only surfaces a job status (and ``pyflyte-execute`` exits 0 even on user error), +the connector checks for an ``error.pb`` at the output prefix on ``Completed`` to +avoid reporting a false success — the SageMaker-native equivalent of EMR's +exit-code translation shim. +""" + +import dataclasses +import hashlib +import re +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +import cloudpickle +from flyteidl.core.execution_pb2 import TaskExecution + +from flytekit import FlyteContext, ImageSpec, PythonFunctionTask +from flytekit.configuration import DefaultImages, SerializationSettings +from flytekit.core.constants import FLYTE_FAIL_ON_ERROR +from flytekit.core.context_manager import FlyteContextManager +from flytekit.extend.backend.base_connector import ( + AsyncConnectorBase, + AsyncConnectorExecutorMixin, + Resource, + ResourceMeta, +) +from flytekit.models.literals import LiteralMap +from flytekit.models.task import TaskTemplate + +from .boto3_mixin import Boto3ConnectorMixin, CustomException + +# Default base image for Pythonic-mode tasks. Applied (EMR-style) when the user +# passes an ``ImageSpec`` without an explicit ``base_image`` so the built image +# is guaranteed to contain flytekit, which the in-container ``pyflyte-execute`` +# needs. Users can always override by setting ``base_image`` or passing a plain +# ECR image URI string as ``container_image``. +SAGEMAKER_PYTHONIC_BASE_IMAGE = DefaultImages.default_image() + +# SageMaker job names: <= 63 chars, must match ^[a-zA-Z0-9](-*[a-zA-Z0-9])*. +_MAX_JOB_NAME_LEN = 63 +_INVALID_JOB_NAME_CHARS = re.compile(r"[^a-zA-Z0-9-]") + + +@dataclass +class PythonicJobConfig: + """Shared task configuration for SageMaker Pythonic-mode jobs. + + :param execution_role_arn: IAM role SageMaker assumes to run the job (S3/ECR/logs). + :param region: AWS region for the SageMaker client. + :param instance_type: SageMaker ML instance type for the job. + :param instance_count: Number of instances. + :param volume_size_in_gb: Attached EBS volume size. + :param max_runtime_in_seconds: Hard stop for the job. + :param environment: Extra environment variables for the container. + :param kms_key_id: Optional KMS key for the attached storage volume. + :param tags: Resource tags (dict; converted to the boto3 ``[{Key, Value}]`` shape). + :param job_name_prefix: Prefix for the generated SageMaker job name. + """ + + execution_role_arn: str + region: str + instance_type: str = "ml.m5.large" + instance_count: int = 1 + volume_size_in_gb: int = 30 + max_runtime_in_seconds: int = 3600 + environment: Optional[Dict[str, str]] = None + kms_key_id: Optional[str] = None + tags: Optional[Dict[str, str]] = None + job_name_prefix: str = "flyte-" + + def __post_init__(self) -> None: + if not self.execution_role_arn: + raise ValueError("execution_role_arn is required") + if not self.region: + raise ValueError("region is required") + if self.instance_count != 1: + raise ValueError("Pythonic SageMaker jobs currently require instance_count=1") + if self.volume_size_in_gb < 1: + raise ValueError("volume_size_in_gb must be at least 1") + if self.max_runtime_in_seconds < 1: + raise ValueError("max_runtime_in_seconds must be at least 1") + + def to_dict(self) -> Dict[str, Any]: + """Serialize to a plain dict for ``task_template.custom`` (drops Nones).""" + return {k: v for k, v in dataclasses.asdict(self).items() if v is not None} + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "PythonicJobConfig": + """Deserialize from a dict (inverse of ``to_dict``).""" + field_names = {f.name for f in dataclasses.fields(cls)} + return cls(**{k: v for k, v in data.items() if k in field_names}) + + +class PythonicSageMakerJobTask(AsyncConnectorExecutorMixin, PythonFunctionTask): + """Base class for SageMaker Pythonic-mode tasks (Processing / Training). + + Subclasses only set ``_TASK_TYPE``. This base applies the default base image + to bare ``ImageSpec``s, serializes the task config into ``custom``, and + dispatches ``execute()`` between local-mimic (connector) and in-container + (run the user function) — the same shape as ``EMRServerlessTask``. + """ + + _TASK_TYPE = "pythonic-sagemaker-job" + + def __init__(self, task_config, task_function, container_image=None, **kwargs): + if container_image is None: + raise ValueError( + "Pythonic SageMaker jobs require an explicit container_image that can be pushed to or resolved from ECR." + ) + if isinstance(container_image, ImageSpec) and container_image.base_image is None: + container_image = dataclasses.replace(container_image, base_image=SAGEMAKER_PYTHONIC_BASE_IMAGE) + + super().__init__( + task_config=task_config, + task_function=task_function, + task_type=self._TASK_TYPE, + container_image=container_image, + **kwargs, + ) + + def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: + return self.task_config.to_dict() + + def execute(self, **kwargs: Any) -> Any: + """Local ``pyflyte run`` mimics the backend via the connector; on the + SageMaker worker (dispatched by the rendered entrypoint) run the user + function directly.""" + ctx = FlyteContextManager.current_context() + if ctx.execution_state and ctx.execution_state.is_local_execution(): + return AsyncConnectorExecutorMixin.execute(self, **kwargs) + return PythonFunctionTask.execute(self, **kwargs) + + +@dataclass +class PythonicJobMetadata(ResourceMeta): + """Metadata persisted by FlytePropeller between connector calls. + + ``output_prefix`` is retained so ``get()`` can check for ``error.pb`` written + by the in-container ``pyflyte-execute``. + """ + + job_name: str + output_prefix: str + region: Optional[str] = None + has_outputs: bool = False + + def encode(self) -> bytes: + return cloudpickle.dumps(self) + + @classmethod + def decode(cls, data: bytes) -> "PythonicJobMetadata": + return cloudpickle.loads(data) + + +def _make_job_name( + prefix: str, + task_execution_metadata: Optional[Any], + task_template: TaskTemplate, + output_prefix: str = "", +) -> str: + """Build a deterministic, unique, SageMaker-valid (<=63 char) job name. + + Derived from the Flyte execution + node + retry so retries get fresh names + while a connector re-invocation for the same attempt reuses the same name + (the create() idempotency handler then treats it as already-running). + """ + seed_parts = [output_prefix] + teid = getattr(task_execution_metadata, "task_execution_id", None) if task_execution_metadata else None + if teid is not None: + node = getattr(teid, "node_execution_id", None) + if node is not None: + ex = getattr(node, "execution_id", None) + if ex is not None: + seed_parts.extend( + [ + getattr(ex, "project", "") or "", + getattr(ex, "domain", "") or "", + getattr(ex, "name", "") or "", + ] + ) + seed_parts.append(getattr(node, "node_id", "") or "") + seed_parts.append(str(getattr(teid, "retry_attempt", "") or "")) + + task_id = task_template.id + if task_id is not None: + seed_parts.extend( + [ + getattr(task_id, "project", "") or "", + getattr(task_id, "domain", "") or "", + getattr(task_id, "name", "") or "", + getattr(task_id, "version", "") or "", + ] + ) + + seed = "-".join(p for p in seed_parts if p) + if not seed: + seed = "job" + + digest = hashlib.sha256(seed.encode("utf-8")).hexdigest()[:20] + normalized_prefix = _INVALID_JOB_NAME_CHARS.sub("-", prefix).strip("-") or "flyte" + suffix = f"-{digest}" + normalized_prefix = normalized_prefix[: _MAX_JOB_NAME_LEN - len(suffix)].rstrip("-") or "j" + return f"{normalized_prefix}{suffix}" + + +def _tags_to_list(tags: Optional[Dict[str, str]]) -> Optional[list]: + if not tags: + return None + return [{"Key": str(k), "Value": str(v)} for k, v in tags.items()] + + +def _build_environment(container: Any, config: Dict[str, Any]) -> Dict[str, str]: + """Merge the serialized task environment with explicit SageMaker overrides.""" + environment = dict(getattr(container, "env", None) or {}) + environment.update(config.get("environment") or {}) + environment[FLYTE_FAIL_ON_ERROR] = "true" + return environment + + +def _container_entrypoint(container: Any) -> list[str]: + """Validate and return the rendered entrypoint accepted by SageMaker.""" + entrypoint = list(getattr(container, "args", None) or []) + if not entrypoint: + raise ValueError("Pythonic SageMaker jobs require a rendered Flyte container entrypoint.") + if len(entrypoint) > 100: + raise ValueError("SageMaker ContainerEntrypoint supports at most 100 arguments.") + if any(len(argument) > 256 for argument in entrypoint): + raise ValueError("Each SageMaker ContainerEntrypoint argument must be at most 256 characters.") + return entrypoint + + +@dataclass(frozen=True) +class _PythonicJobError: + message: str + recoverable: bool + + +class PythonicSageMakerJobConnector(Boto3ConnectorMixin, AsyncConnectorBase): + """Base async connector for SageMaker Pythonic-mode jobs. + + Subclasses supply the per-service deltas (method names, response keys, state + map) and implement ``_build_request``. ``create``/``get``/``delete`` and the + ``outputs.pb`` / ``error.pb`` handling are shared. + """ + + # --- per-service deltas (override in subclasses) --- + task_type_name: str = "pythonic-sagemaker-job" + create_method: str = "" + describe_method: str = "" + stop_method: str = "" + job_name_key: str = "" + status_key: str = "" + secondary_status_key: Optional[str] = None + _state_map: Dict[str, Any] = field(default_factory=dict) + + def __init__(self): + super().__init__( + service="sagemaker", + task_type_name=self.task_type_name, + metadata_type=PythonicJobMetadata, + ) + + def _build_request( + self, + *, + container: Any, + config: Dict[str, Any], + job_name: str, + output_prefix: str, + ) -> Dict[str, Any]: + """Assemble the ``create_*_job`` boto3 request. Implemented per service.""" + raise NotImplementedError + + async def create( + self, + task_template: TaskTemplate, + output_prefix: str, + inputs: Optional[LiteralMap] = None, + task_execution_metadata: Optional[Any] = None, + **kwargs: Any, + ) -> PythonicJobMetadata: + container = task_template.container + if container is None or not getattr(container, "image", None): + raise ValueError( + "Pythonic mode requires a container image. Pass container_image= " + "to the @task so SageMaker can run the function inside a flytekit-containing image." + ) + + config = dict(task_template.custom or {}) + region = config.get("region") + job_name = _make_job_name( + config.get("job_name_prefix") or "flyte-", + task_execution_metadata, + task_template, + output_prefix, + ) + request = self._build_request( + container=container, config=config, job_name=job_name, output_prefix=output_prefix + ) + + try: + await self._call(method=self.create_method, config=request, 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. + already_exists = error_code == "ResourceInUse" or ( + error_code == "ValidationException" and "Cannot create already existing" in error_message + ) + if not already_exists: + raise e + + interface = getattr(task_template, "interface", None) + has_outputs = bool(getattr(interface, "outputs", None)) + return PythonicJobMetadata( + job_name=job_name, + output_prefix=output_prefix, + region=region, + has_outputs=has_outputs, + ) + + async def get(self, resource_meta: PythonicJobMetadata, **kwargs: Any) -> Resource: + describe_response, _ = await self._call( + method=self.describe_method, + config={self.job_name_key: resource_meta.job_name}, + region=resource_meta.region, + ) + + current_state = describe_response.get(self.status_key) + flyte_phase = self._state_map.get(current_state, TaskExecution.RUNNING) + message = self._status_message(describe_response, current_state) + + if current_state in ("Completed", "Failed", "Stopped"): + task_error = self._read_error(resource_meta.output_prefix) + if task_error is not None: + phase = TaskExecution.RETRYABLE_FAILED if task_error.recoverable else TaskExecution.FAILED + return Resource(phase=phase, message=task_error.message) + + if current_state == "Completed": + if resource_meta.has_outputs and not self._artifact_exists(resource_meta.output_prefix, "outputs.pb"): + return Resource( + phase=TaskExecution.FAILED, + message="SageMaker job completed without producing Flyte outputs.pb.", + ) + # outputs=None -> flytekit reads the typed return from outputs.pb. + return Resource(phase=TaskExecution.SUCCEEDED, outputs=None, message=message) + + return Resource(phase=flyte_phase, message=message) + + async def delete(self, resource_meta: PythonicJobMetadata, **kwargs: Any) -> None: + try: + await self._call( + method=self.stop_method, + config={self.job_name_key: resource_meta.job_name}, + region=resource_meta.region, + ) + except CustomException as e: + original_exception = e.original_exception + error_code = original_exception.response["Error"]["Code"] + error_message = original_exception.response["Error"]["Message"] + + # delete() may run after the job already finished/stopped — SageMaker + # rejects stop on a non-running job; nothing to cancel. + if error_code == "ResourceNotFound" or ( + error_code == "ValidationException" and "non-running" in error_message + ): + return + raise e + + def _status_message(self, describe_response: Dict[str, Any], state: Optional[str]) -> Optional[str]: + if self.secondary_status_key and state == "InProgress": + return describe_response.get(self.secondary_status_key) + if state in ("Failed", "Stopped"): + secondary = describe_response.get(self.secondary_status_key) if self.secondary_status_key else None + return describe_response.get("FailureReason") or secondary or describe_response.get("ExitMessage") + return None + + @staticmethod + def _artifact_exists(output_prefix: str, name: str) -> bool: + ctx = FlyteContext.current_context() + return ctx.file_access.exists(f"{output_prefix.rstrip('/')}/{name}") + + @classmethod + def _read_error(cls, output_prefix: str) -> Optional[_PythonicJobError]: + """Return structured user-error details when ``error.pb`` exists.""" + error_path = f"{output_prefix.rstrip('/')}/error.pb" + if not cls._artifact_exists(output_prefix, "error.pb"): + return None + + try: + from flyteidl.core import errors_pb2 + + ctx = FlyteContext.current_context() + local_path = ctx.file_access.get_random_local_path() + ctx.file_access.get_data(error_path, local_path) + with open(local_path, "rb") as f: + doc = errors_pb2.ErrorDocument() + doc.ParseFromString(f.read()) + if doc.error and doc.error.message: + return _PythonicJobError( + message=doc.error.message, + recoverable=doc.error.kind == errors_pb2.ContainerError.RECOVERABLE, + ) + except Exception: + return _PythonicJobError( + message="User task raised an error, but error.pb could not be decoded.", + recoverable=False, + ) + return _PythonicJobError( + message="User task raised an error (see error.pb in the task output prefix).", + recoverable=False, + ) 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..a6a384e1a5 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/__init__.py @@ -0,0 +1,39 @@ +""" +.. currentmodule:: flytekitplugins.awssagemaker_processing + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + SageMakerProcessingJobConnector + SageMakerProcessingJobTask + SageMakerStopProcessingJobTask + SageMakerDescribeProcessingJobTask + SageMakerProcessing + SageMakerProcessingTask + SageMakerProcessingTaskConnector +""" + +from .connector import ( + SageMakerProcessingJobConnector, + SageMakerProcessingJobMetadata, + SageMakerProcessingTaskConnector, +) +from .task import ( + SageMakerDescribeProcessingJobTask, + SageMakerProcessing, + SageMakerProcessingJobTask, + SageMakerProcessingTask, + SageMakerStopProcessingJobTask, +) + +__all__ = [ + "SageMakerProcessingJobConnector", + "SageMakerProcessingJobMetadata", + "SageMakerProcessingJobTask", + "SageMakerStopProcessingJobTask", + "SageMakerDescribeProcessingJobTask", + "SageMakerProcessing", + "SageMakerProcessingTask", + "SageMakerProcessingTaskConnector", +] 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..e8322ae6e9 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/connector.py @@ -0,0 +1,245 @@ +"""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 flytekitplugins.awssagemaker_inference.pythonic_base import ( + PythonicSageMakerJobConnector, + _build_environment, + _container_entrypoint, + _tags_to_list, +) + +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()) + + +class SageMakerProcessingTaskConnector(PythonicSageMakerJobConnector): + """Pythonic-mode connector: runs a Flyte ``@task`` function inside a processing job.""" + + name = "SageMaker Processing Task Connector" + + task_type_name = "sagemaker-processing-task" + create_method = "create_processing_job" + describe_method = "describe_processing_job" + stop_method = "stop_processing_job" + job_name_key = "ProcessingJobName" + status_key = "ProcessingJobStatus" + secondary_status_key = None # processing jobs have no SecondaryStatus + _state_map = _STATE_MAP + + def _build_request(self, *, container, config, job_name, output_prefix): + cluster_config: Dict[str, Any] = { + "InstanceType": config.get("instance_type", "ml.m5.large"), + "InstanceCount": config.get("instance_count", 1), + "VolumeSizeInGB": config.get("volume_size_in_gb", 30), + } + kms_key_id = config.get("kms_key_id") + if kms_key_id: + cluster_config["VolumeKmsKeyId"] = kms_key_id + + request: Dict[str, Any] = { + "ProcessingJobName": job_name, + "RoleArn": config["execution_role_arn"], + "AppSpecification": { + "ImageUri": container.image, + "ContainerEntrypoint": _container_entrypoint(container), + }, + "ProcessingResources": {"ClusterConfig": cluster_config}, + "StoppingCondition": {"MaxRuntimeInSeconds": config.get("max_runtime_in_seconds", 3600)}, + } + + request["Environment"] = _build_environment(container, config) + network_config = config.get("network_config") + if network_config: + if network_config.get("EnableNetworkIsolation"): + raise ValueError( + "EnableNetworkIsolation=True is incompatible with Pythonic mode because pyflyte-execute " + "must access Flyte's S3 input and output prefixes." + ) + request["NetworkConfig"] = network_config + tags = _tags_to_list(config.get("tags")) + if tags: + request["Tags"] = tags + return request + + +ConnectorRegistry.register(SageMakerProcessingTaskConnector()) 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..28b542dd35 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_processing/task.py @@ -0,0 +1,143 @@ +"""User-facing tasks for SageMaker processing jobs.""" + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Type, Union + +from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask +from flytekitplugins.awssagemaker_inference.pythonic_base import ( + PythonicJobConfig, + PythonicSageMakerJobTask, +) + +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 import TaskPlugins +from flytekit.extend.backend.base_connector import AsyncConnectorExecutorMixin +from flytekit.image_spec.image_spec import ImageBuildEngine + + +@dataclass +class SageMakerProcessing(PythonicJobConfig): + """Pythonic-mode config for a SageMaker processing job. + + Use as ``@task(task_config=SageMakerProcessing(...), container_image=...)`` + to run the decorated Python function inside a SageMaker processing job. See + :class:`~flytekitplugins.awssagemaker_inference.pythonic_base.PythonicJobConfig` + for the available fields. + + :param network_config: Optional boto3 ``NetworkConfig`` request shape, + including an optional nested ``VpcConfig``. + """ + + network_config: Optional[Dict[str, Any]] = None + + +class SageMakerProcessingTask(PythonicSageMakerJobTask): + """Pythonic-mode SageMaker processing task (runs a ``@task`` function in a processing job).""" + + _TASK_TYPE = "sagemaker-processing-task" + + +TaskPlugins.register_pythontask_plugin(SageMakerProcessing, SageMakerProcessingTask) + + +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..cf1c4b75f0 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/__init__.py @@ -0,0 +1,39 @@ +""" +.. currentmodule:: flytekitplugins.awssagemaker_training + +.. autosummary:: + :template: custom.rst + :toctree: generated/ + + SageMakerTrainingJobConnector + SageMakerTrainingJobTask + SageMakerStopTrainingJobTask + SageMakerDescribeTrainingJobTask + SageMakerTraining + SageMakerTrainingTask + SageMakerTrainingTaskConnector +""" + +from .connector import ( + SageMakerTrainingJobConnector, + SageMakerTrainingJobMetadata, + SageMakerTrainingTaskConnector, +) +from .task import ( + SageMakerDescribeTrainingJobTask, + SageMakerStopTrainingJobTask, + SageMakerTraining, + SageMakerTrainingJobTask, + SageMakerTrainingTask, +) + +__all__ = [ + "SageMakerTrainingJobConnector", + "SageMakerTrainingJobMetadata", + "SageMakerTrainingJobTask", + "SageMakerStopTrainingJobTask", + "SageMakerDescribeTrainingJobTask", + "SageMakerTraining", + "SageMakerTrainingTask", + "SageMakerTrainingTaskConnector", +] 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..6345aab888 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/connector.py @@ -0,0 +1,247 @@ +"""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 flytekitplugins.awssagemaker_inference.pythonic_base import ( + PythonicSageMakerJobConnector, + _build_environment, + _container_entrypoint, + _tags_to_list, +) + +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()) + + +class SageMakerTrainingTaskConnector(PythonicSageMakerJobConnector): + """Pythonic-mode connector: runs a Flyte ``@task`` function inside a training job.""" + + name = "SageMaker Training Task Connector" + + task_type_name = "sagemaker-training-task" + create_method = "create_training_job" + describe_method = "describe_training_job" + stop_method = "stop_training_job" + job_name_key = "TrainingJobName" + status_key = "TrainingJobStatus" + secondary_status_key = "SecondaryStatus" + _state_map = _STATE_MAP + + def _build_request(self, *, container, config, job_name, output_prefix): + resource_config: Dict[str, Any] = { + "InstanceType": config.get("instance_type", "ml.m5.large"), + "InstanceCount": config.get("instance_count", 1), + "VolumeSizeInGB": config.get("volume_size_in_gb", 30), + } + kms_key_id = config.get("kms_key_id") + if kms_key_id: + resource_config["VolumeKmsKeyId"] = kms_key_id + + # Pythonic mode returns results via Flyte outputs.pb; SageMaker still + # requires OutputDataConfig, so default it to the Flyte output prefix + # (the resulting model.tar.gz is harmless and unused). + output_s3_path = config.get("output_s3_path") or f"{output_prefix}/_sagemaker_model" + + request: Dict[str, Any] = { + "TrainingJobName": job_name, + "RoleArn": config["execution_role_arn"], + "AlgorithmSpecification": { + "TrainingImage": container.image, + "ContainerEntrypoint": _container_entrypoint(container), + "TrainingInputMode": "File", + }, + "ResourceConfig": resource_config, + "OutputDataConfig": {"S3OutputPath": output_s3_path}, + "StoppingCondition": {"MaxRuntimeInSeconds": config.get("max_runtime_in_seconds", 3600)}, + } + + request["Environment"] = _build_environment(container, config) + vpc_config = config.get("vpc_config") + if vpc_config: + request["VpcConfig"] = vpc_config + tags = _tags_to_list(config.get("tags")) + if tags: + request["Tags"] = tags + return request + + +ConnectorRegistry.register(SageMakerTrainingTaskConnector()) 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..ac4c83588f --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_training/task.py @@ -0,0 +1,144 @@ +"""User-facing tasks for SageMaker training jobs.""" + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Type, Union + +from flytekitplugins.awssagemaker_inference.boto3_task import BotoConfig, BotoTask +from flytekitplugins.awssagemaker_inference.pythonic_base import ( + PythonicJobConfig, + PythonicSageMakerJobTask, +) + +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 import TaskPlugins +from flytekit.extend.backend.base_connector import AsyncConnectorExecutorMixin +from flytekit.image_spec.image_spec import ImageBuildEngine + + +@dataclass +class SageMakerTraining(PythonicJobConfig): + """Pythonic-mode config for a SageMaker training job. + + Use as ``@task(task_config=SageMakerTraining(...), container_image=...)`` to + run the decorated Python function inside a SageMaker training job. Inherits + all fields from + :class:`~flytekitplugins.awssagemaker_inference.pythonic_base.PythonicJobConfig`. + + :param output_s3_path: Optional ``OutputDataConfig.S3OutputPath`` for the + SageMaker model tar. Pythonic mode returns results via Flyte outputs, so + this is rarely needed; when unset the connector points it at the Flyte + output prefix (the resulting ``model.tar.gz`` is harmless and unused). + :param vpc_config: Optional boto3 ``VpcConfig`` request shape containing + ``Subnets`` and ``SecurityGroupIds``. + """ + + output_s3_path: Optional[str] = None + vpc_config: Optional[Dict[str, Any]] = None + + +class SageMakerTrainingTask(PythonicSageMakerJobTask): + """Pythonic-mode SageMaker training task (runs a ``@task`` function in a training job).""" + + _TASK_TYPE = "sagemaker-training-task" + + +TaskPlugins.register_pythontask_plugin(SageMakerTraining, SageMakerTrainingTask) + + +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_pythonic_task.py b/plugins/flytekit-aws-sagemaker/tests/test_processing_pythonic_task.py new file mode 100644 index 0000000000..e3a7f95417 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_processing_pythonic_task.py @@ -0,0 +1,113 @@ +"""Unit tests for the Pythonic-mode SageMaker processing task.""" + +import types + +import pytest +from flytekitplugins.awssagemaker_inference.pythonic_base import ( + SAGEMAKER_PYTHONIC_BASE_IMAGE, + PythonicSageMakerJobTask, +) +from flytekitplugins.awssagemaker_processing import SageMakerProcessing, SageMakerProcessingTask + +from flytekit import ImageSpec, PythonFunctionTask, task +from flytekit.extend.backend.base_connector import AsyncConnectorExecutorMixin + +ROLE = "arn:aws:iam::123456789012:role/sm-exec" +REGION = "us-east-1" + + +def _fn() -> int: + return 1 + + +def _build_task(container_image): + return SageMakerProcessingTask( + task_config=SageMakerProcessing(execution_role_arn=ROLE, region=REGION), + task_function=_fn, + container_image=container_image, + ) + + +def test_dataclass_validation(): + with pytest.raises(ValueError): + SageMakerProcessing(execution_role_arn="", region=REGION) + with pytest.raises(ValueError): + SageMakerProcessing(execution_role_arn=ROLE, region="") + with pytest.raises(ValueError): + SageMakerProcessing(execution_role_arn=ROLE, region=REGION, instance_count=2) + with pytest.raises(ValueError): + SageMakerProcessing(execution_role_arn=ROLE, region=REGION, max_runtime_in_seconds=0) + + +def test_to_from_dict_round_trip(): + cfg = SageMakerProcessing( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.2xlarge", + tags={"team": "ml"}, + ) + as_dict = cfg.to_dict() + # None-valued fields are dropped. + assert "network_config" not in as_dict + assert as_dict["instance_count"] == 1 + assert SageMakerProcessing.from_dict(as_dict) == cfg + + +def test_default_base_image_applied_to_bare_imagespec(): + img = ImageSpec(name="smoke", registry="r") + assert img.base_image is None + t = _build_task(img) + assert isinstance(t.container_image, ImageSpec) + assert t.container_image.base_image == SAGEMAKER_PYTHONIC_BASE_IMAGE + + +def test_user_base_image_preserved(): + img = ImageSpec(name="smoke", registry="r", base_image="my/base:1") + t = _build_task(img) + assert t.container_image.base_image == "my/base:1" + + +def test_string_image_preserved(): + t = _build_task("123.dkr.ecr.us-east-1.amazonaws.com/img:tag") + assert t.container_image == "123.dkr.ecr.us-east-1.amazonaws.com/img:tag" + + +def test_explicit_container_image_is_required(): + with pytest.raises(ValueError, match="explicit container_image"): + _build_task(None) + + +def test_get_custom_returns_config_dict(): + t = _build_task("img:tag") + custom = t.get_custom(None) + assert custom["execution_role_arn"] == ROLE + assert SageMakerProcessing.from_dict(custom) == t.task_config + + +def test_register_pythontask_plugin_wiring(): + @task( + task_config=SageMakerProcessing(execution_role_arn=ROLE, region=REGION), + container_image="img:tag", + ) + def my_processing() -> int: + return 1 + + assert isinstance(my_processing, SageMakerProcessingTask) + assert isinstance(my_processing, PythonicSageMakerJobTask) + assert my_processing.task_type == "sagemaker-processing-task" + + +@pytest.mark.parametrize("is_local,expected", [(True, "local"), (False, "worker")]) +def test_execute_dispatches_local_vs_worker(monkeypatch, is_local, expected): + import flytekitplugins.awssagemaker_inference.pythonic_base as base + + t = _build_task("img:tag") + monkeypatch.setattr(AsyncConnectorExecutorMixin, "execute", lambda self, **kw: "local") + monkeypatch.setattr(PythonFunctionTask, "execute", lambda self, **kw: "worker") + + ctx = types.SimpleNamespace( + execution_state=types.SimpleNamespace(is_local_execution=lambda: is_local) + ) + monkeypatch.setattr(base.FlyteContextManager, "current_context", lambda: ctx) + + assert t.execute() == expected 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_pythonic_connector.py b/plugins/flytekit-aws-sagemaker/tests/test_pythonic_connector.py new file mode 100644 index 0000000000..f104293b81 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_pythonic_connector.py @@ -0,0 +1,584 @@ +"""Unit tests for the Pythonic-mode SageMaker connectors (Processing + Training) +and their shared base (request building, outputs.pb/error.pb resolution, job +name generation, idempotency, stop-on-delete).""" + +import types +from collections import OrderedDict +from pathlib import Path +from unittest import mock + +import pytest +from botocore.exceptions import ClientError +from flyteidl.core import errors_pb2 +from flyteidl.core.execution_pb2 import TaskExecution +from flytekitplugins.awssagemaker_inference.boto3_mixin import CustomException +from flytekitplugins.awssagemaker_inference.pythonic_base import ( + PythonicJobMetadata, + _PythonicJobError, + _make_job_name, + _tags_to_list, +) +from flytekitplugins.awssagemaker_processing import SageMakerProcessing +from flytekitplugins.awssagemaker_training import SageMakerTraining + +from flytekit import task +from flytekit.configuration import Image, ImageConfig, SerializationSettings +from flytekit.core.constants import FLYTE_FAIL_ON_ERROR +from flytekit.extend.backend.base_connector import ConnectorRegistry +from flytekit.extend.backend.utils import render_task_template +from flytekit.tools.translator import get_serializable + +ROLE = "arn:aws:iam::123456789012:role/sm-exec" +REGION = "us-east-1" +OUTPUT_PREFIX = "s3://bucket/flyte/exec/n0" +IMAGE = "123.dkr.ecr.us-east-1.amazonaws.com/img:tag" +ARGS = ["pyflyte-fast-execute", "--", "pyflyte-execute", "--output-prefix", OUTPUT_PREFIX] + +_CALL = "flytekitplugins.awssagemaker_inference.boto3_mixin.Boto3ConnectorMixin._call" +_READ_ERROR = ( + "flytekitplugins.awssagemaker_inference.pythonic_base." + "PythonicSageMakerJobConnector._read_error" +) +_ARTIFACT_EXISTS = ( + "flytekitplugins.awssagemaker_inference.pythonic_base." + "PythonicSageMakerJobConnector._artifact_exists" +) + + +def _template(custom, args=ARGS, image=IMAGE, env=None, outputs=None): + container = types.SimpleNamespace(image=image, args=args, env=env or {}) + tid = types.SimpleNamespace( + project="project", + domain="domain", + name="project.domain.my_task", + version="v1", + ) + interface = types.SimpleNamespace(outputs=outputs or {}) + return types.SimpleNamespace( + container=container, + custom=custom, + id=tid, + interface=interface, + ) + + +def _client_error(code, message, op): + return CustomException( + message="boom", + idempotence_token="tok", + original_exception=ClientError( + error_response={"Error": {"Code": code, "Message": message}}, + operation_name=op, + ), + ) + + +# --------------------------- shared base helpers --------------------------- + + +def test_tags_to_list(): + assert _tags_to_list(None) is None + assert _tags_to_list({}) is None + assert _tags_to_list({"a": "b"}) == [{"Key": "a", "Value": "b"}] + + +def test_metadata_encode_decode_round_trip(): + meta = PythonicJobMetadata( + job_name="j", + output_prefix=OUTPUT_PREFIX, + region=REGION, + has_outputs=True, + ) + assert PythonicJobMetadata.decode(meta.encode()) == meta + + +def test_make_job_name_valid_and_unique_per_retry(): + tt = _template({}) + md_retry0 = types.SimpleNamespace( + task_execution_id=types.SimpleNamespace( + node_execution_id=types.SimpleNamespace( + execution_id=types.SimpleNamespace(name="exec1"), node_id="n0" + ), + retry_attempt=0, + ) + ) + md_retry1 = types.SimpleNamespace( + task_execution_id=types.SimpleNamespace( + node_execution_id=types.SimpleNamespace( + execution_id=types.SimpleNamespace(name="exec1"), node_id="n0" + ), + retry_attempt=1, + ) + ) + name0 = _make_job_name("flyte-", md_retry0, tt) + name1 = _make_job_name("flyte-", md_retry1, tt) + + assert name0 != name1 # retries get distinct names + for name in (name0, name1): + assert name.startswith("flyte-") + assert len(name) <= 63 + assert name[0].isalnum() + assert all(c.isalnum() or c == "-" for c in name) + + +def test_make_job_name_falls_back_to_task_id_without_metadata(): + name = _make_job_name("flyte-", None, _template({})) + assert name.startswith("flyte-") + assert len(name) <= 63 + + +def test_make_job_name_reserves_digest_for_long_prefixes(): + first = _make_job_name("x" * 100, None, _template({}), "s3://bucket/exec-1") + second = _make_job_name("x" * 100, None, _template({}), "s3://bucket/exec-2") + + assert len(first) <= 63 + assert len(second) <= 63 + assert first != second + + +def test_read_error_preserves_recoverable_kind(monkeypatch, tmp_path): + import flytekitplugins.awssagemaker_inference.pythonic_base as base + + payload = errors_pb2.ErrorDocument( + error=errors_pb2.ContainerError( + code="USER:Recoverable", + message="try again", + kind=errors_pb2.ContainerError.RECOVERABLE, + ) + ).SerializeToString() + local_path = tmp_path / "error.pb" + file_access = types.SimpleNamespace( + exists=lambda _path: True, + get_random_local_path=lambda: str(local_path), + get_data=lambda _remote, target: Path(target).write_bytes(payload), + ) + monkeypatch.setattr( + base.FlyteContext, + "current_context", + lambda: types.SimpleNamespace(file_access=file_access), + ) + + error = base.PythonicSageMakerJobConnector._read_error(OUTPUT_PREFIX) + + assert error == _PythonicJobError(message="try again", recoverable=True) + + +def test_read_error_propagates_storage_failures(monkeypatch): + import flytekitplugins.awssagemaker_inference.pythonic_base as base + + file_access = types.SimpleNamespace( + exists=mock.Mock(side_effect=PermissionError("S3 access denied")) + ) + monkeypatch.setattr( + base.FlyteContext, + "current_context", + lambda: types.SimpleNamespace(file_access=file_access), + ) + + with pytest.raises(PermissionError, match="S3 access denied"): + base.PythonicSageMakerJobConnector._read_error(OUTPUT_PREFIX) + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_real_serialized_task_renders_entrypoint_and_environment(mock_call): + @task( + task_config=SageMakerProcessing(execution_role_arn=ROLE, region=REGION), + container_image=IMAGE, + environment={"FROM_TASK": "yes"}, + ) + def serialized_processing(value: int) -> int: + return value + 1 + + default_image = Image(name="default", fqn=IMAGE.rsplit(":", 1)[0], tag="tag") + settings = SerializationSettings( + project="project", + domain="domain", + version="version", + image_config=ImageConfig( + default_image=default_image, + images=[default_image], + ), + ) + task_spec = get_serializable(OrderedDict(), settings, serialized_processing) + template = render_task_template(task_spec.template, OUTPUT_PREFIX) + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + mock_call.return_value = ({}, "tok") + + metadata = await connector.create(template, output_prefix=OUTPUT_PREFIX) + + request = mock_call.call_args.kwargs["config"] + entrypoint = request["AppSpecification"]["ContainerEntrypoint"] + assert any(OUTPUT_PREFIX in argument for argument in entrypoint) + assert not any("{{." in argument for argument in entrypoint) + assert request["Environment"]["FROM_TASK"] == "yes" + assert request["Environment"][FLYTE_FAIL_ON_ERROR] == "true" + assert metadata.has_outputs is True + + +# --------------------------- processing connector --------------------------- + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_create_builds_request(mock_call): + mock_call.return_value = ({}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + custom = SageMakerProcessing( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.xlarge", + tags={"team": "ml"}, + environment={"K": "V"}, + ).to_dict() + + meta = await connector.create( + _template( + custom, + env={"FROM_TASK": "yes", "K": "old"}, + outputs={"result": object()}, + ), + output_prefix=OUTPUT_PREFIX, + ) + + assert meta.output_prefix == OUTPUT_PREFIX + assert meta.region == REGION + assert meta.has_outputs is True + assert meta.job_name.startswith("flyte-") + + req = mock_call.call_args.kwargs["config"] + assert mock_call.call_args.kwargs["method"] == "create_processing_job" + assert req["ProcessingJobName"] == meta.job_name + assert req["RoleArn"] == ROLE + assert req["AppSpecification"] == {"ImageUri": IMAGE, "ContainerEntrypoint": ARGS} + cc = req["ProcessingResources"]["ClusterConfig"] + assert cc["InstanceType"] == "ml.m5.xlarge" + assert cc["InstanceCount"] == 1 + assert req["Environment"] == { + "FROM_TASK": "yes", + "K": "V", + FLYTE_FAIL_ON_ERROR: "true", + } + assert req["Tags"] == [{"Key": "team", "Value": "ml"}] + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=None) +@mock.patch(_CALL) +async def test_processing_get_completed_no_error_succeeds(mock_call, _mock_err): + mock_call.return_value = ({"ProcessingJobStatus": "Completed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + + resource = await connector.get(meta) + assert resource.phase == TaskExecution.SUCCEEDED + # outputs=None -> flytekit materializes the typed return from outputs.pb. + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=_PythonicJobError("ValueError: boom", False)) +@mock.patch(_CALL) +async def test_processing_get_completed_with_error_pb_fails(mock_call, _mock_err): + mock_call.return_value = ({"ProcessingJobStatus": "Completed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + + resource = await connector.get(meta) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "ValueError: boom" + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=_PythonicJobError("temporary failure", True)) +@mock.patch(_CALL) +async def test_processing_recoverable_error_is_retryable(mock_call, _mock_err): + mock_call.return_value = ({"ProcessingJobStatus": "Failed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + + resource = await connector.get(meta) + + assert resource.phase == TaskExecution.RETRYABLE_FAILED + assert resource.message == "temporary failure" + + +@pytest.mark.asyncio +@mock.patch(_ARTIFACT_EXISTS, return_value=False) +@mock.patch(_READ_ERROR, return_value=None) +@mock.patch(_CALL) +async def test_processing_completed_without_declared_outputs_fails( + mock_call, + _mock_err, + _mock_exists, +): + mock_call.return_value = ({"ProcessingJobStatus": "Completed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata( + job_name="j", + output_prefix=OUTPUT_PREFIX, + region=REGION, + has_outputs=True, + ) + + resource = await connector.get(meta) + + assert resource.phase == TaskExecution.FAILED + assert "without producing Flyte outputs.pb" in resource.message + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_get_in_progress_running(mock_call): + mock_call.return_value = ({"ProcessingJobStatus": "InProgress"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + + resource = await connector.get(meta) + assert resource.phase == TaskExecution.RUNNING + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_create_already_exists_is_idempotent(mock_call): + mock_call.side_effect = _client_error( + "ResourceInUse", "Processing job flyte-x already exists", "CreateProcessingJob" + ) + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = await connector.create( + _template(SageMakerProcessing(execution_role_arn=ROLE, region=REGION).to_dict()), + output_prefix=OUTPUT_PREFIX, + ) + assert meta.output_prefix == OUTPUT_PREFIX + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_create_resource_limit_propagates(mock_call): + mock_call.side_effect = _client_error( + "ResourceLimitExceeded", + "Processing job quota exceeded", + "CreateProcessingJob", + ) + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + with pytest.raises(CustomException): + await connector.create( + _template(SageMakerProcessing(execution_role_arn=ROLE, region=REGION).to_dict()), + output_prefix=OUTPUT_PREFIX, + ) + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_create_unknown_error_propagates(mock_call): + mock_call.side_effect = _client_error("AccessDeniedException", "nope", "CreateProcessingJob") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + with pytest.raises(CustomException): + await connector.create( + _template(SageMakerProcessing(execution_role_arn=ROLE, region=REGION).to_dict()), + output_prefix=OUTPUT_PREFIX, + ) + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_delete_stops_job(mock_call): + mock_call.return_value = ({}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + assert await connector.delete(meta) is None + assert mock_call.call_args.kwargs["method"] == "stop_processing_job" + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_delete_swallows_non_running(mock_call): + mock_call.side_effect = _client_error( + "ValidationException", "the processing job is not in a non-running state", "StopProcessingJob" + ) + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + assert await connector.delete(meta) is None + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_processing_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = _client_error( + "ResourceNotFound", + "Processing job does not exist", + "StopProcessingJob", + ) + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + assert await connector.delete(meta) is None + + +# --------------------------- training connector --------------------------- + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_training_create_builds_request(mock_call): + mock_call.return_value = ({}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + custom = SageMakerTraining( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.large", + vpc_config={"Subnets": ["subnet-1"], "SecurityGroupIds": ["sg-1"]}, + ).to_dict() + + meta = await connector.create(_template(custom), output_prefix=OUTPUT_PREFIX) + + req = mock_call.call_args.kwargs["config"] + assert mock_call.call_args.kwargs["method"] == "create_training_job" + assert req["TrainingJobName"] == meta.job_name + assert req["RoleArn"] == ROLE + assert req["AlgorithmSpecification"] == { + "TrainingImage": IMAGE, + "ContainerEntrypoint": ARGS, + "TrainingInputMode": "File", + } + assert req["ResourceConfig"]["InstanceType"] == "ml.m5.large" + # OutputDataConfig defaults to the Flyte output prefix when unset. + assert req["OutputDataConfig"]["S3OutputPath"] == f"{OUTPUT_PREFIX}/_sagemaker_model" + assert req["VpcConfig"] == {"Subnets": ["subnet-1"], "SecurityGroupIds": ["sg-1"]} + assert req["Environment"][FLYTE_FAIL_ON_ERROR] == "true" + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_training_create_honours_explicit_output_path(mock_call): + mock_call.return_value = ({}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + custom = SageMakerTraining( + execution_role_arn=ROLE, + region=REGION, + output_s3_path="s3://b/model/", + ).to_dict() + await connector.create(_template(custom), output_prefix=OUTPUT_PREFIX) + req = mock_call.call_args.kwargs["config"] + assert req["OutputDataConfig"]["S3OutputPath"] == "s3://b/model/" + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_training_get_in_progress_surfaces_secondary_status(mock_call): + mock_call.return_value = ( + {"TrainingJobStatus": "InProgress", "SecondaryStatus": "Downloading"}, + "tok", + ) + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + resource = await connector.get(meta) + assert resource.phase == TaskExecution.RUNNING + assert resource.message == "Downloading" + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=None) +@mock.patch(_CALL) +async def test_training_get_completed_no_error_succeeds(mock_call, _mock_err): + mock_call.return_value = ({"TrainingJobStatus": "Completed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + resource = await connector.get(meta) + assert resource.phase == TaskExecution.SUCCEEDED + assert resource.outputs is None + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=_PythonicJobError("RuntimeError: nan loss", False)) +@mock.patch(_CALL) +async def test_training_get_completed_with_error_pb_fails(mock_call, _mock_err): + mock_call.return_value = ({"TrainingJobStatus": "Completed"}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + resource = await connector.get(meta) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "RuntimeError: nan loss" + + +@pytest.mark.asyncio +@mock.patch(_READ_ERROR, return_value=None) +@mock.patch(_CALL) +async def test_training_get_failed_surfaces_failure_reason(mock_call, _mock_err): + mock_call.return_value = ( + {"TrainingJobStatus": "Failed", "FailureReason": "ClientError: boom"}, + "tok", + ) + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + resource = await connector.get(meta) + assert resource.phase == TaskExecution.FAILED + assert resource.message == "ClientError: boom" + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_training_delete_stops_job(mock_call): + mock_call.return_value = ({}, "tok") + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX) + assert await connector.delete(meta) is None + assert mock_call.call_args.kwargs["method"] == "stop_training_job" + + +@pytest.mark.asyncio +@mock.patch(_CALL) +async def test_training_delete_swallows_resource_not_found(mock_call): + mock_call.side_effect = _client_error( + "ResourceNotFound", + "Training job does not exist", + "StopTrainingJob", + ) + connector = ConnectorRegistry.get_connector("sagemaker-training-task") + meta = PythonicJobMetadata(job_name="j", output_prefix=OUTPUT_PREFIX, region=REGION) + assert await connector.delete(meta) is None + + +# --------------------------- container and network guards --------------------------- + + +@pytest.mark.asyncio +async def test_create_without_container_image_raises(): + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + tt = _template( + SageMakerProcessing(execution_role_arn=ROLE, region=REGION).to_dict(), + image=None, + ) + with pytest.raises(ValueError): + await connector.create(tt, output_prefix=OUTPUT_PREFIX) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "args,error", + [ + ([], "rendered Flyte container entrypoint"), + (["x"] * 101, "at most 100 arguments"), + (["x" * 257], "at most 256 characters"), + ], +) +async def test_create_rejects_invalid_container_entrypoint(args, error): + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + tt = _template( + SageMakerProcessing(execution_role_arn=ROLE, region=REGION).to_dict(), + args=args, + ) + with pytest.raises(ValueError, match=error): + await connector.create(tt, output_prefix=OUTPUT_PREFIX) + + +@pytest.mark.asyncio +async def test_processing_rejects_network_isolation(): + connector = ConnectorRegistry.get_connector("sagemaker-processing-task") + config = SageMakerProcessing( + execution_role_arn=ROLE, + region=REGION, + network_config={"EnableNetworkIsolation": True}, + ).to_dict() + with pytest.raises(ValueError, match="incompatible with Pythonic mode"): + await connector.create(_template(config), output_prefix=OUTPUT_PREFIX) 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_pythonic_task.py b/plugins/flytekit-aws-sagemaker/tests/test_training_pythonic_task.py new file mode 100644 index 0000000000..b121bfc2b8 --- /dev/null +++ b/plugins/flytekit-aws-sagemaker/tests/test_training_pythonic_task.py @@ -0,0 +1,96 @@ +"""Unit tests for the Pythonic-mode SageMaker training task.""" + +import types + +import pytest +from flytekitplugins.awssagemaker_inference.pythonic_base import ( + SAGEMAKER_PYTHONIC_BASE_IMAGE, + PythonicSageMakerJobTask, +) +from flytekitplugins.awssagemaker_training import SageMakerTraining, SageMakerTrainingTask + +from flytekit import ImageSpec, PythonFunctionTask, task +from flytekit.extend.backend.base_connector import AsyncConnectorExecutorMixin + +ROLE = "arn:aws:iam::123456789012:role/sm-exec" +REGION = "us-east-1" + + +def _fn() -> int: + return 1 + + +def _build_task(container_image): + return SageMakerTrainingTask( + task_config=SageMakerTraining(execution_role_arn=ROLE, region=REGION), + task_function=_fn, + container_image=container_image, + ) + + +def test_dataclass_validation(): + with pytest.raises(ValueError): + SageMakerTraining(execution_role_arn="", region=REGION) + with pytest.raises(ValueError): + SageMakerTraining(execution_role_arn=ROLE, region="") + with pytest.raises(ValueError): + SageMakerTraining(execution_role_arn=ROLE, region=REGION, volume_size_in_gb=0) + + +def test_to_from_dict_round_trip_with_output_path(): + cfg = SageMakerTraining( + execution_role_arn=ROLE, + region=REGION, + instance_type="ml.m5.large", + output_s3_path="s3://bucket/model/", + ) + as_dict = cfg.to_dict() + assert as_dict["output_s3_path"] == "s3://bucket/model/" + assert SageMakerTraining.from_dict(as_dict) == cfg + + +def test_output_s3_path_optional_and_dropped_when_none(): + cfg = SageMakerTraining(execution_role_arn=ROLE, region=REGION) + assert "output_s3_path" not in cfg.to_dict() + + +def test_default_base_image_applied_to_bare_imagespec(): + img = ImageSpec(name="smoke", registry="r") + t = _build_task(img) + assert t.container_image.base_image == SAGEMAKER_PYTHONIC_BASE_IMAGE + + +def test_get_custom_returns_config_dict(): + t = _build_task("img:tag") + custom = t.get_custom(None) + assert custom["execution_role_arn"] == ROLE + assert SageMakerTraining.from_dict(custom) == t.task_config + + +def test_register_pythontask_plugin_wiring(): + @task( + task_config=SageMakerTraining(execution_role_arn=ROLE, region=REGION), + container_image="img:tag", + ) + def my_training() -> int: + return 1 + + assert isinstance(my_training, SageMakerTrainingTask) + assert isinstance(my_training, PythonicSageMakerJobTask) + assert my_training.task_type == "sagemaker-training-task" + + +@pytest.mark.parametrize("is_local,expected", [(True, "local"), (False, "worker")]) +def test_execute_dispatches_local_vs_worker(monkeypatch, is_local, expected): + import flytekitplugins.awssagemaker_inference.pythonic_base as base + + t = _build_task("img:tag") + monkeypatch.setattr(AsyncConnectorExecutorMixin, "execute", lambda self, **kw: "local") + monkeypatch.setattr(PythonFunctionTask, "execute", lambda self, **kw: "worker") + + ctx = types.SimpleNamespace( + execution_state=types.SimpleNamespace(is_local_execution=lambda: is_local) + ) + monkeypatch.setattr(base.FlyteContextManager, "current_context", lambda: ctx) + + assert t.execute() == expected 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" diff --git a/tests/flytekit/unit/extend/test_connector.py b/tests/flytekit/unit/extend/test_connector.py index 217210171b..5f7d768df8 100644 --- a/tests/flytekit/unit/extend/test_connector.py +++ b/tests/flytekit/unit/extend/test_connector.py @@ -374,6 +374,7 @@ def test_is_terminal_phase(): assert is_terminal_phase(TaskExecution.SUCCEEDED) assert is_terminal_phase(TaskExecution.ABORTED) assert is_terminal_phase(TaskExecution.FAILED) + assert is_terminal_phase(TaskExecution.RETRYABLE_FAILED) assert not is_terminal_phase(TaskExecution.RUNNING)