diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 4f50c1be74..d09f8c33a0 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 @@ -108,6 +109,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..94c3e989d4 --- /dev/null +++ b/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.") diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index b5b45e77a4..306c991277 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/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): 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..4b9b314b92 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_cleanup_step.py @@ -0,0 +1,211 @@ +# 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.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.step_dependency import StepDependency, TriggerCondition +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_adds_sources(): + step = CleanupStep( + name="CleanupStep1", + s3_sources=[S3CleanupSource(step="TrainingJob1"), S3CleanupSource(step="AthenaStep1")], + depends_on=["OtherStep", "AthenaStep1"], + ) + + 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), + "step_dependency": StepDependency(target, TriggerCondition.ANY), + } + + +_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] + step = CleanupStep( + name="CleanupStep1", s3_sources=[S3CleanupSource(step=training)], depends_on=[dep] + ) + + assert step.depends_on == [dep] + + +@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] + step = CleanupStep( + name="CleanupStep1", s3_sources=[S3CleanupSource(step=training)], depends_on=[dep] + ) + + 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=[]) + + +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", }