feat(mlops): Add CleanupStep pipeline step

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
@@ -108,6 +109,8 @@
     "CallbackStep",
     "CallbackOutput",
     "ClarifyCheckStep",
+    "CleanupStep",
+    "S3CleanupSource",
     "ConditionStep",
     "EMRStep",
     "EMRStepConfig",

sagemaker-mlops/src/sagemaker/mlops/workflow/cleanup_step.py

@@ -0,0 +1,142 @@
+# 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.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, 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):
+        # 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.
+
+    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, StepOutput, StepDependency]]
+        ] = 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 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.")
+        # 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 = {_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)
+                dependency_names.add(source.step_name)
+        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.")

sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py

@@ -69,6 +69,7 @@ class StepTypeEnum(Enum):
     INFERENCE_COMPONENT = "InferenceComponent"
     LINEAGE = "Lineage"
     JOB = "Job"
+    CLEANUP = "Cleanup"
 
 
 class Step(Entity):

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"])