From bc637ec2eb38623c3e44d8c8264321e97d6c899e Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Tue, 13 Oct 2026 14:52:17 -0700 Subject: [PATCH 1/4] feat(mlops): Add CleanupStep pipeline step Add a CleanupStep that deletes the Amazon S3 outputs of upstream steps. Each S3CleanupSource names a source step and, for Processing and Transform steps, the exact object URIs to delete. Other rules (supported source types, URI limits, VpcConfig) are left to the service, which validates them at CreatePipeline. depends_on defaults to the source steps since only outputs of completed steps can be deleted. Integ test passed against a production region. --- X-AI-Prompt: Add a PySDK CleanupStep with unit and integration tests and raise a PR X-AI-Tool: kiro-cli --- .../src/sagemaker/mlops/workflow/__init__.py | 3 + .../sagemaker/mlops/workflow/cleanup_step.py | 117 ++++++++++++++ .../src/sagemaker/mlops/workflow/steps.py | 1 + .../tests/integ/workflow/test_cleanup_step.py | 134 ++++++++++++++++ .../tests/unit/workflow/test_cleanup_step.py | 149 ++++++++++++++++++ .../tests/unit/workflow/test_job_step.py | 1 + 6 files changed, 405 insertions(+) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py create mode 100644 sagemaker-mlops/tests/integ/workflow/test_cleanup_step.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index d438f4f1db..d9a1d6dfd0 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -46,6 +46,7 @@ from sagemaker.mlops.workflow.automl_step import AutoMLStep from sagemaker.mlops.workflow.callback_step import CallbackStep, CallbackOutput from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep +from sagemaker.mlops.workflow.cleanup_step import CleanupStep, S3CleanupSource from sagemaker.mlops.workflow.condition_step import ConditionStep from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep @@ -105,6 +106,8 @@ "CallbackStep", "CallbackOutput", "ClarifyCheckStep", + "CleanupStep", + "S3CleanupSource", "ConditionStep", "EMRStep", "EMRStepConfig", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py new file mode 100644 index 0000000000..b61cf5ebf1 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py @@ -0,0 +1,117 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""The step definitions for the Cleanup step in SageMaker Pipelines.""" + +from __future__ import absolute_import + +from typing import List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import PipelineVariable, RequestType + +from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class S3CleanupSource: + """An upstream step whose Amazon S3 outputs a `CleanupStep` deletes. + + Omit `s3_uris` to let the service find the outputs from the source step + (Training and Athena steps). Set `s3_uris` to delete exact objects + written by a Processing or Transform step; each URI must be under that + step's output location. + """ + + def __init__( + self, + step: Union[str, Step], + s3_uris: Optional[List[Union[str, PipelineVariable]]] = None, + ): + """Constructs an `S3CleanupSource`. + + Args: + step (str or Step): The source step, or its name. + s3_uris (List[str or PipelineVariable]): Exact `s3://bucket/key` + object URIs to delete. Omit for Training and Athena steps + (default: None). + """ + if s3_uris is not None and len(s3_uris) == 0: + raise ValueError("s3_uris must be omitted or contain at least one URI.") + self.step = step + self.s3_uris = s3_uris + + @property + def step_name(self) -> str: + """The name of the source step.""" + return self.step.name if isinstance(self.step, Step) else self.step + + def to_request(self) -> RequestType: + """Gets the request structure for this cleanup source.""" + request = {"StepName": self.step_name} + if self.s3_uris is not None: + request["S3Uris"] = list(self.s3_uris) + return request + + +class CleanupStep(Step): + """`CleanupStep` for SageMaker Pipelines Workflows. + + Deletes the Amazon S3 outputs of upstream steps. The pipeline execution + role needs `s3:DeleteObject` on the target objects. A source step that + sets a `VpcConfig` can't be cleaned up. + """ + + def __init__( + self, + name: str, + s3_sources: List[S3CleanupSource], + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + ): + """Constructs a `CleanupStep`. + + Args: + name (str): The name of the `CleanupStep`. + s3_sources (List[S3CleanupSource]): The steps whose S3 outputs + to delete. + display_name (str): The display name of the `CleanupStep` + (default: None). + description (str): The description of the `CleanupStep` + (default: None). + depends_on (List[str or Step or StepCollection]): The steps this + `CleanupStep` depends on. Defaults to the source steps + (default: None). + """ + if not s3_sources: + raise ValueError("s3_sources must contain at least one S3CleanupSource.") + # The service can only delete outputs of steps that already ran, so the + # sources are the natural dependencies. + if depends_on is None: + depends_on = [source.step for source in s3_sources] + super(CleanupStep, self).__init__( + name, display_name, description, StepTypeEnum.CLEANUP, depends_on + ) + self.s3_sources = s3_sources + + @property + def arguments(self) -> RequestType: + """The arguments dictionary that is used to define the `CleanupStep`.""" + return {"Resources": {"S3": [source.to_request() for source in self.s3_sources]}} + + @property + def properties(self): + """A `Properties` object is not available for the `CleanupStep`. + + Cleanup results are reported in the execution step metadata. + """ + raise RuntimeError("The Properties object is not available for the CleanupStep.") diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 333f32616c..614101d890 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -68,6 +68,7 @@ class StepTypeEnum(Enum): INFERENCE_COMPONENT = "InferenceComponent" LINEAGE = "Lineage" JOB = "Job" + CLEANUP = "Cleanup" class Step(Entity): diff --git a/sagemaker-mlops/tests/integ/workflow/test_cleanup_step.py b/sagemaker-mlops/tests/integ/workflow/test_cleanup_step.py new file mode 100644 index 0000000000..74bee72eb0 --- /dev/null +++ b/sagemaker-mlops/tests/integ/workflow/test_cleanup_step.py @@ -0,0 +1,134 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Integration test for CleanupStep. + +Runs a Processing step that writes two objects, then a CleanupStep that +deletes one of them by URI. Requires the execution role to have +``s3:DeleteObject`` on the session's default bucket. +""" + +from __future__ import absolute_import + +import time +import uuid + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.core import image_uris +from sagemaker.core.processing import Processor +from sagemaker.core.shapes import ProcessingOutput, ProcessingS3Output +from sagemaker.mlops.workflow import CleanupStep, S3CleanupSource +from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.steps import ProcessingStep + + +def _object_exists(s3_client, bucket, key): + try: + s3_client.head_object(Bucket=bucket, Key=key) + return True + except ClientError as e: + if e.response["Error"]["Code"] in ("404", "NoSuchKey", "NotFound"): + return False + raise + + +def test_cleanup_step_deletes_processing_output( + sagemaker_session, pipeline_session, role, sklearn_latest_version +): + stamp = uuid.uuid4().hex[:8] + bucket = sagemaker_session.default_bucket() + prefix = f"integ-test-cleanup-step/{stamp}" + output_uri = f"s3://{bucket}/{prefix}/output" + region = pipeline_session.boto_region_name + + processor = Processor( + role=role, + image_uri=image_uris.retrieve( + framework="sklearn", + region=region, + version=sklearn_latest_version, + py_version="py3", + instance_type="ml.m5.large", + ), + instance_count=1, + instance_type="ml.m5.large", + entrypoint=[ + "/bin/sh", + "-c", + "echo keep > /opt/ml/processing/output/keep.txt && " + "echo delete > /opt/ml/processing/output/delete.txt", + ], + sagemaker_session=pipeline_session, + ) + step_process = ProcessingStep( + name="WriteOutput", + step_args=processor.run( + outputs=[ + ProcessingOutput( + output_name="output", + s3_output=ProcessingS3Output( + s3_uri=output_uri, + local_path="/opt/ml/processing/output", + s3_upload_mode="EndOfJob", + ), + ) + ], + ), + ) + step_cleanup = CleanupStep( + name="CleanupOutput", + s3_sources=[S3CleanupSource(step=step_process, s3_uris=[f"{output_uri}/delete.txt"])], + ) + pipeline = Pipeline( + name=f"integ-cleanup-step-{stamp}", + steps=[step_process, step_cleanup], + sagemaker_session=pipeline_session, + ) + + client = sagemaker_session.sagemaker_client + s3_client = sagemaker_session.boto_session.client("s3") + try: + pipeline.upsert(role_arn=role) + execution = pipeline.start() + + final_status = None + deadline = time.time() + 1200 + while time.time() < deadline: + final_status = execution.describe()["PipelineExecutionStatus"] + if final_status in ("Succeeded", "Failed", "Stopped", "CompletedWithFailures"): + break + time.sleep(30) + + steps = client.list_pipeline_execution_steps(PipelineExecutionArn=execution.arn)[ + "PipelineExecutionSteps" + ] + if final_status != "Succeeded": + failures = "\n".join( + f"{s['StepName']}: {s.get('FailureReason', 'no reason')}" + for s in steps + if s.get("StepStatus") == "Failed" + ) + pytest.fail(f"Pipeline execution status={final_status}. Details:\n{failures}") + + cleanup = next(s for s in steps if s["StepName"] == "CleanupOutput") + assert cleanup["Metadata"]["Cleanup"]["CleanedSteps"] == ["WriteOutput"] + assert not _object_exists(s3_client, bucket, f"{prefix}/output/delete.txt") + assert _object_exists(s3_client, bucket, f"{prefix}/output/keep.txt") + finally: + try: + pipeline.delete() + except Exception: + pass + for obj in s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix).get("Contents", []): + s3_client.delete_object(Bucket=bucket, Key=obj["Key"]) diff --git a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py new file mode 100644 index 0000000000..44ec42da6e --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py @@ -0,0 +1,149 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for workflow cleanup_step.""" + +from __future__ import absolute_import + +import json +from unittest.mock import Mock + +import pytest + +from sagemaker.core.workflow.functions import Join +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.mlops.workflow import CleanupStep, S3CleanupSource +from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class CustomStep(Step): + def __init__(self, name, step_type=StepTypeEnum.TRAINING): + super(CustomStep, self).__init__(name=name, step_type=step_type, depends_on=[]) + + @property + def arguments(self): + return {} + + @property + def properties(self): + return None + + +@pytest.fixture +def mock_session(): + session = Mock() + session.boto_region_name = "us-west-2" + session.local_mode = False + session.sagemaker_config = {} + return session + + +def test_cleanup_step_resolved_source(): + training = CustomStep("TrainingJob1") + step = CleanupStep(name="CleanupStep1", s3_sources=[S3CleanupSource(step=training)]) + + assert step.step_type == StepTypeEnum.CLEANUP + assert step.arguments == {"Resources": {"S3": [{"StepName": "TrainingJob1"}]}} + assert step.depends_on == [training] + + +def test_cleanup_step_direct_uri_source_with_step_name(): + step = CleanupStep( + name="CleanupStep1", + s3_sources=[S3CleanupSource(step="ProcessingJob1", s3_uris=["s3://bucket/out/data.json"])], + ) + + assert step.arguments == { + "Resources": { + "S3": [{"StepName": "ProcessingJob1", "S3Uris": ["s3://bucket/out/data.json"]}] + } + } + assert step.depends_on == ["ProcessingJob1"] + + +def test_cleanup_step_explicit_depends_on_overrides_sources(): + step = CleanupStep( + name="CleanupStep1", + s3_sources=[S3CleanupSource(step="TrainingJob1")], + depends_on=["OtherStep"], + ) + + assert step.depends_on == ["OtherStep"] + + +def test_cleanup_step_requires_sources(): + with pytest.raises(ValueError, match="s3_sources"): + CleanupStep(name="CleanupStep1", s3_sources=[]) + + +def test_s3_cleanup_source_rejects_empty_uris(): + with pytest.raises(ValueError, match="s3_uris"): + S3CleanupSource(step="ProcessingJob1", s3_uris=[]) + + +def test_cleanup_step_properties_raises_error(): + step = CleanupStep(name="CleanupStep1", s3_sources=[S3CleanupSource(step="TrainingJob1")]) + with pytest.raises(RuntimeError, match="not available"): + _ = step.properties + + +def test_cleanup_step_pipeline_definition(mock_session): + output_path = ParameterString(name="ProcessingOutputPath", default_value="s3://bucket/out") + training = CustomStep("TrainingJob1") + processing = CustomStep("ProcessingJob1", step_type=StepTypeEnum.PROCESSING) + cleanup = CleanupStep( + name="CleanupStep1", + s3_sources=[ + S3CleanupSource(step=training), + S3CleanupSource( + step=processing, + s3_uris=[Join(on="", values=[output_path, "/processing-job-data.json"])], + ), + ], + ) + pipeline = Pipeline( + name="MyPipeline", + parameters=[output_path], + steps=[training, processing, cleanup], + sagemaker_session=mock_session, + ) + + steps = json.loads(pipeline.definition())["Steps"] + cleanup_definition = next(s for s in steps if s["Name"] == "CleanupStep1") + + assert cleanup_definition == { + "Name": "CleanupStep1", + "Type": "Cleanup", + "DependsOn": ["TrainingJob1", "ProcessingJob1"], + "Arguments": { + "Resources": { + "S3": [ + {"StepName": "TrainingJob1"}, + { + "StepName": "ProcessingJob1", + "S3Uris": [ + { + "Std:Join": { + "On": "", + "Values": [ + {"Get": "Parameters.ProcessingOutputPath"}, + "/processing-job-data.json", + ], + } + } + ], + }, + ] + } + }, + } diff --git a/sagemaker-mlops/tests/unit/workflow/test_job_step.py b/sagemaker-mlops/tests/unit/workflow/test_job_step.py index 4e6e572409..2a699197ba 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_job_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_job_step.py @@ -50,6 +50,7 @@ "INFERENCE_COMPONENT": "InferenceComponent", "LINEAGE": "Lineage", "JOB": "Job", + "CLEANUP": "Cleanup", } From d6f9df417edf40011fb8f1dd888513b116f8c83f Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Thu, 15 Oct 2026 12:00:49 -0700 Subject: [PATCH 2/4] fix(mlops): Always add CleanupStep sources to depends_on The service does not require a Cleanup step's DependsOn to include its sources, so an explicit depends_on could let Cleanup run before its source steps finished. --- X-AI-Prompt: Address jkasiraj's PR comment to merge sources into depends_on X-AI-Tool: kiro-cli --- .../src/sagemaker/mlops/workflow/cleanup_step.py | 16 ++++++++++------ .../tests/unit/workflow/test_cleanup_step.py | 8 ++++---- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py index b61cf5ebf1..9f230a96c8 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py @@ -88,16 +88,20 @@ def __init__( (default: None). description (str): The description of the `CleanupStep` (default: None). - depends_on (List[str or Step or StepCollection]): The steps this - `CleanupStep` depends on. Defaults to the source steps - (default: None). + depends_on (List[str or Step or StepCollection]): Additional steps + this `CleanupStep` depends on. The source steps are always + added (default: None). """ if not s3_sources: raise ValueError("s3_sources must contain at least one S3CleanupSource.") # The service can only delete outputs of steps that already ran, so the - # sources are the natural dependencies. - if depends_on is None: - depends_on = [source.step for source in s3_sources] + # sources are always dependencies, in addition to any the caller passes. + depends_on = list(depends_on or []) + dependency_names = {dep if isinstance(dep, str) else dep.name for dep in depends_on} + for source in s3_sources: + if source.step_name not in dependency_names: + depends_on.append(source.step) + dependency_names.add(source.step_name) super(CleanupStep, self).__init__( name, display_name, description, StepTypeEnum.CLEANUP, depends_on ) diff --git a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py index 44ec42da6e..b8961fdc52 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py @@ -71,14 +71,14 @@ def test_cleanup_step_direct_uri_source_with_step_name(): assert step.depends_on == ["ProcessingJob1"] -def test_cleanup_step_explicit_depends_on_overrides_sources(): +def test_cleanup_step_explicit_depends_on_adds_sources(): step = CleanupStep( name="CleanupStep1", - s3_sources=[S3CleanupSource(step="TrainingJob1")], - depends_on=["OtherStep"], + s3_sources=[S3CleanupSource(step="TrainingJob1"), S3CleanupSource(step="AthenaStep1")], + depends_on=["OtherStep", "AthenaStep1"], ) - assert step.depends_on == ["OtherStep"] + assert step.depends_on == ["OtherStep", "AthenaStep1", "TrainingJob1"] def test_cleanup_step_requires_sources(): From 043d0c78f13815c630a2e98c82c93b46394b529d Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Thu, 15 Oct 2026 14:45:04 -0700 Subject: [PATCH 3/4] fix(mlops): Resolve every depends_on form in CleanupStep A StepOutput entry raised AttributeError and a StepCollection was matched by its own name instead of its last step, so the source could be appended twice. --- X-AI-Prompt: Address jkasiraj's PR comment to handle all depends_on types X-AI-Tool: kiro-cli --- .../sagemaker/mlops/workflow/cleanup_step.py | 23 ++++++++--- .../tests/unit/workflow/test_cleanup_step.py | 40 +++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py index 9f230a96c8..3c7166cee7 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py @@ -17,11 +17,24 @@ from typing import List, Optional, Union from sagemaker.core.helper.pipeline_variable import PipelineVariable, RequestType +from sagemaker.core.workflow.step_outputs import StepOutput, get_step from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum +def _dependency_name(dep: Union[str, Step, StepCollection, StepOutput]) -> str: + """The name of the step a `depends_on` entry refers to.""" + if isinstance(dep, str): + return dep + if isinstance(dep, StepCollection): + # Matches how the pipeline resolves a StepCollection dependency. + return dep.steps[-1].name + if isinstance(dep, StepOutput): + return get_step(dep).name + return dep.name + + class S3CleanupSource: """An upstream step whose Amazon S3 outputs a `CleanupStep` deletes. @@ -76,7 +89,7 @@ def __init__( s3_sources: List[S3CleanupSource], display_name: Optional[str] = None, description: Optional[str] = None, - depends_on: Optional[List[Union[str, Step, StepCollection]]] = None, + depends_on: Optional[List[Union[str, Step, StepCollection, StepOutput]]] = None, ): """Constructs a `CleanupStep`. @@ -88,16 +101,16 @@ def __init__( (default: None). description (str): The description of the `CleanupStep` (default: None). - depends_on (List[str or Step or StepCollection]): Additional steps - this `CleanupStep` depends on. The source steps are always - added (default: None). + depends_on (List[str or Step or StepCollection or StepOutput]): + Additional steps this `CleanupStep` depends on. The source + steps are always added (default: None). """ if not s3_sources: raise ValueError("s3_sources must contain at least one S3CleanupSource.") # The service can only delete outputs of steps that already ran, so the # sources are always dependencies, in addition to any the caller passes. depends_on = list(depends_on or []) - dependency_names = {dep if isinstance(dep, str) else dep.name for dep in depends_on} + dependency_names = {_dependency_name(dep) for dep in depends_on} for source in s3_sources: if source.step_name not in dependency_names: depends_on.append(source.step) diff --git a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py index b8961fdc52..a417fdc638 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py @@ -21,8 +21,10 @@ from sagemaker.core.workflow.functions import Join from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.core.workflow.step_outputs import StepOutput from sagemaker.mlops.workflow import CleanupStep, S3CleanupSource from sagemaker.mlops.workflow.pipeline import Pipeline +from sagemaker.mlops.workflow.step_collections import StepCollection from sagemaker.mlops.workflow.steps import Step, StepTypeEnum @@ -81,6 +83,44 @@ def test_cleanup_step_explicit_depends_on_adds_sources(): assert step.depends_on == ["OtherStep", "AthenaStep1", "TrainingJob1"] +class _StepOutput(StepOutput): + @property + def expr(self): + return {} + + +def _dependency_forms(target): + """Every form a depends_on entry can take, each referring to `target`.""" + return { + "str": target.name, + "step": target, + "step_collection": StepCollection(name="Collection", steps=[CustomStep("First"), target]), + "step_output": _StepOutput(target), + } + + +@pytest.mark.parametrize("form", ["str", "step", "step_collection", "step_output"]) +def test_cleanup_step_depends_on_already_has_source(form): + training = CustomStep("TrainingJob1") + dep = _dependency_forms(training)[form] + step = CleanupStep( + name="CleanupStep1", s3_sources=[S3CleanupSource(step=training)], depends_on=[dep] + ) + + assert step.depends_on == [dep] + + +@pytest.mark.parametrize("form", ["str", "step", "step_collection", "step_output"]) +def test_cleanup_step_depends_on_other_step_adds_source(form): + training = CustomStep("TrainingJob1") + dep = _dependency_forms(CustomStep("OtherStep"))[form] + step = CleanupStep( + name="CleanupStep1", s3_sources=[S3CleanupSource(step=training)], depends_on=[dep] + ) + + assert step.depends_on == [dep, training] + + def test_cleanup_step_requires_sources(): with pytest.raises(ValueError, match="s3_sources"): CleanupStep(name="CleanupStep1", s3_sources=[]) From b9132bf5baa1ab43a606ca3dfd7e41a7d3d50f12 Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Fri, 16 Oct 2026 09:50:48 -0700 Subject: [PATCH 4/4] fix(mlops): Honor StepDependency entries in CleanupStep depends_on A source listed with a run condition was not recognized, so a second plain Completed edge was appended and compilation failed on conflicting triggers. --- X-AI-Prompt: Fix CleanupStep for the merged run conditions PR X-AI-Tool: kiro-cli --- .../sagemaker/mlops/workflow/cleanup_step.py | 18 +++++++++---- .../tests/unit/workflow/test_cleanup_step.py | 26 +++++++++++++++++-- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py index 3c7166cee7..94c3e989d4 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py @@ -20,11 +20,14 @@ from sagemaker.core.workflow.step_outputs import StepOutput, get_step from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.step_dependency import StepDependency from sagemaker.mlops.workflow.steps import Step, StepTypeEnum -def _dependency_name(dep: Union[str, Step, StepCollection, StepOutput]) -> str: +def _dependency_name(dep: Union[str, Step, StepCollection, StepOutput, StepDependency]) -> str: """The name of the step a `depends_on` entry refers to.""" + if isinstance(dep, StepDependency): + return _dependency_name(dep.step) if isinstance(dep, str): return dep if isinstance(dep, StepCollection): @@ -89,7 +92,9 @@ def __init__( s3_sources: List[S3CleanupSource], display_name: Optional[str] = None, description: Optional[str] = None, - depends_on: Optional[List[Union[str, Step, StepCollection, StepOutput]]] = None, + depends_on: Optional[ + List[Union[str, Step, StepCollection, StepOutput, StepDependency]] + ] = None, ): """Constructs a `CleanupStep`. @@ -101,9 +106,12 @@ def __init__( (default: None). description (str): The description of the `CleanupStep` (default: None). - depends_on (List[str or Step or StepCollection or StepOutput]): - Additional steps this `CleanupStep` depends on. The source - steps are always added (default: None). + depends_on (List[str or Step or StepCollection or StepOutput or + StepDependency]): Additional steps this `CleanupStep` depends + on. Source steps not already listed are added with the + default `Completed` trigger; list a source in a + `StepDependency` to clean up on another outcome, for example + `TriggerCondition.ANY` (default: None). """ if not s3_sources: raise ValueError("s3_sources must contain at least one S3CleanupSource.") diff --git a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py index a417fdc638..4b9b314b92 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py @@ -25,6 +25,7 @@ from sagemaker.mlops.workflow import CleanupStep, S3CleanupSource from sagemaker.mlops.workflow.pipeline import Pipeline from sagemaker.mlops.workflow.step_collections import StepCollection +from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition from sagemaker.mlops.workflow.steps import Step, StepTypeEnum @@ -96,10 +97,14 @@ def _dependency_forms(target): "step": target, "step_collection": StepCollection(name="Collection", steps=[CustomStep("First"), target]), "step_output": _StepOutput(target), + "step_dependency": StepDependency(target, TriggerCondition.ANY), } -@pytest.mark.parametrize("form", ["str", "step", "step_collection", "step_output"]) +_FORMS = ["str", "step", "step_collection", "step_output", "step_dependency"] + + +@pytest.mark.parametrize("form", _FORMS) def test_cleanup_step_depends_on_already_has_source(form): training = CustomStep("TrainingJob1") dep = _dependency_forms(training)[form] @@ -110,7 +115,7 @@ def test_cleanup_step_depends_on_already_has_source(form): assert step.depends_on == [dep] -@pytest.mark.parametrize("form", ["str", "step", "step_collection", "step_output"]) +@pytest.mark.parametrize("form", _FORMS) def test_cleanup_step_depends_on_other_step_adds_source(form): training = CustomStep("TrainingJob1") dep = _dependency_forms(CustomStep("OtherStep"))[form] @@ -121,6 +126,23 @@ def test_cleanup_step_depends_on_other_step_adds_source(form): assert step.depends_on == [dep, training] +def test_cleanup_step_source_with_any_trigger_pipeline_definition(mock_session): + training = CustomStep("TrainingJob1") + cleanup = CleanupStep( + name="CleanupStep1", + s3_sources=[S3CleanupSource(step=training)], + depends_on=[StepDependency(training, TriggerCondition.ANY)], + ) + pipeline = Pipeline( + name="MyPipeline", steps=[training, cleanup], sagemaker_session=mock_session + ) + + steps = json.loads(pipeline.definition())["Steps"] + cleanup_definition = next(s for s in steps if s["Name"] == "CleanupStep1") + + assert cleanup_definition["DependsOn"] == [{"Step": "TrainingJob1", "Trigger": "Any"}] + + def test_cleanup_step_requires_sources(): with pytest.raises(ValueError, match="s3_sources"): CleanupStep(name="CleanupStep1", s3_sources=[])