Skip to content

Commit 3ecabcf

Browse files
nayan3107nayan3107
andauthored
feat(mlops): Add CleanupStep pipeline step (#6468)
* 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 * 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 * 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 * 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 --------- Co-authored-by: nayan3107 <nayancho@amazon.com>
1 parent 697f2a7 commit 3ecabcf

6 files changed

Lines changed: 492 additions & 0 deletions

File tree

‎sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from sagemaker.mlops.workflow.automl_step import AutoMLStep
4747
from sagemaker.mlops.workflow.callback_step import CallbackStep, CallbackOutput
4848
from sagemaker.mlops.workflow.clarify_check_step import ClarifyCheckStep
49+
from sagemaker.mlops.workflow.cleanup_step import CleanupStep, S3CleanupSource
4950
from sagemaker.mlops.workflow.condition_step import ConditionStep
5051
from sagemaker.mlops.workflow.emr_step import EMRStep, EMRStepConfig
5152
from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep
@@ -108,6 +109,8 @@
108109
"CallbackStep",
109110
"CallbackOutput",
110111
"ClarifyCheckStep",
112+
"CleanupStep",
113+
"S3CleanupSource",
111114
"ConditionStep",
112115
"EMRStep",
113116
"EMRStepConfig",
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""The step definitions for the Cleanup step in SageMaker Pipelines."""
14+
15+
from __future__ import absolute_import
16+
17+
from typing import List, Optional, Union
18+
19+
from sagemaker.core.helper.pipeline_variable import PipelineVariable, RequestType
20+
from sagemaker.core.workflow.step_outputs import StepOutput, get_step
21+
22+
from sagemaker.mlops.workflow.step_collections import StepCollection
23+
from sagemaker.mlops.workflow.step_dependency import StepDependency
24+
from sagemaker.mlops.workflow.steps import Step, StepTypeEnum
25+
26+
27+
def _dependency_name(dep: Union[str, Step, StepCollection, StepOutput, StepDependency]) -> str:
28+
"""The name of the step a `depends_on` entry refers to."""
29+
if isinstance(dep, StepDependency):
30+
return _dependency_name(dep.step)
31+
if isinstance(dep, str):
32+
return dep
33+
if isinstance(dep, StepCollection):
34+
# Matches how the pipeline resolves a StepCollection dependency.
35+
return dep.steps[-1].name
36+
if isinstance(dep, StepOutput):
37+
return get_step(dep).name
38+
return dep.name
39+
40+
41+
class S3CleanupSource:
42+
"""An upstream step whose Amazon S3 outputs a `CleanupStep` deletes.
43+
44+
Omit `s3_uris` to let the service find the outputs from the source step
45+
(Training and Athena steps). Set `s3_uris` to delete exact objects
46+
written by a Processing or Transform step; each URI must be under that
47+
step's output location.
48+
"""
49+
50+
def __init__(
51+
self,
52+
step: Union[str, Step],
53+
s3_uris: Optional[List[Union[str, PipelineVariable]]] = None,
54+
):
55+
"""Constructs an `S3CleanupSource`.
56+
57+
Args:
58+
step (str or Step): The source step, or its name.
59+
s3_uris (List[str or PipelineVariable]): Exact `s3://bucket/key`
60+
object URIs to delete. Omit for Training and Athena steps
61+
(default: None).
62+
"""
63+
if s3_uris is not None and len(s3_uris) == 0:
64+
raise ValueError("s3_uris must be omitted or contain at least one URI.")
65+
self.step = step
66+
self.s3_uris = s3_uris
67+
68+
@property
69+
def step_name(self) -> str:
70+
"""The name of the source step."""
71+
return self.step.name if isinstance(self.step, Step) else self.step
72+
73+
def to_request(self) -> RequestType:
74+
"""Gets the request structure for this cleanup source."""
75+
request = {"StepName": self.step_name}
76+
if self.s3_uris is not None:
77+
request["S3Uris"] = list(self.s3_uris)
78+
return request
79+
80+
81+
class CleanupStep(Step):
82+
"""`CleanupStep` for SageMaker Pipelines Workflows.
83+
84+
Deletes the Amazon S3 outputs of upstream steps. The pipeline execution
85+
role needs `s3:DeleteObject` on the target objects. A source step that
86+
sets a `VpcConfig` can't be cleaned up.
87+
"""
88+
89+
def __init__(
90+
self,
91+
name: str,
92+
s3_sources: List[S3CleanupSource],
93+
display_name: Optional[str] = None,
94+
description: Optional[str] = None,
95+
depends_on: Optional[
96+
List[Union[str, Step, StepCollection, StepOutput, StepDependency]]
97+
] = None,
98+
):
99+
"""Constructs a `CleanupStep`.
100+
101+
Args:
102+
name (str): The name of the `CleanupStep`.
103+
s3_sources (List[S3CleanupSource]): The steps whose S3 outputs
104+
to delete.
105+
display_name (str): The display name of the `CleanupStep`
106+
(default: None).
107+
description (str): The description of the `CleanupStep`
108+
(default: None).
109+
depends_on (List[str or Step or StepCollection or StepOutput or
110+
StepDependency]): Additional steps this `CleanupStep` depends
111+
on. Source steps not already listed are added with the
112+
default `Completed` trigger; list a source in a
113+
`StepDependency` to clean up on another outcome, for example
114+
`TriggerCondition.ANY` (default: None).
115+
"""
116+
if not s3_sources:
117+
raise ValueError("s3_sources must contain at least one S3CleanupSource.")
118+
# The service can only delete outputs of steps that already ran, so the
119+
# sources are always dependencies, in addition to any the caller passes.
120+
depends_on = list(depends_on or [])
121+
dependency_names = {_dependency_name(dep) for dep in depends_on}
122+
for source in s3_sources:
123+
if source.step_name not in dependency_names:
124+
depends_on.append(source.step)
125+
dependency_names.add(source.step_name)
126+
super(CleanupStep, self).__init__(
127+
name, display_name, description, StepTypeEnum.CLEANUP, depends_on
128+
)
129+
self.s3_sources = s3_sources
130+
131+
@property
132+
def arguments(self) -> RequestType:
133+
"""The arguments dictionary that is used to define the `CleanupStep`."""
134+
return {"Resources": {"S3": [source.to_request() for source in self.s3_sources]}}
135+
136+
@property
137+
def properties(self):
138+
"""A `Properties` object is not available for the `CleanupStep`.
139+
140+
Cleanup results are reported in the execution step metadata.
141+
"""
142+
raise RuntimeError("The Properties object is not available for the CleanupStep.")

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ class StepTypeEnum(Enum):
6969
INFERENCE_COMPONENT = "InferenceComponent"
7070
LINEAGE = "Lineage"
7171
JOB = "Job"
72+
CLEANUP = "Cleanup"
7273

7374

7475
class Step(Entity):
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integration test for CleanupStep.
14+
15+
Runs a Processing step that writes two objects, then a CleanupStep that
16+
deletes one of them by URI. Requires the execution role to have
17+
``s3:DeleteObject`` on the session's default bucket.
18+
"""
19+
20+
from __future__ import absolute_import
21+
22+
import time
23+
import uuid
24+
25+
import pytest
26+
from botocore.exceptions import ClientError
27+
28+
from sagemaker.core import image_uris
29+
from sagemaker.core.processing import Processor
30+
from sagemaker.core.shapes import ProcessingOutput, ProcessingS3Output
31+
from sagemaker.mlops.workflow import CleanupStep, S3CleanupSource
32+
from sagemaker.mlops.workflow.pipeline import Pipeline
33+
from sagemaker.mlops.workflow.steps import ProcessingStep
34+
35+
36+
def _object_exists(s3_client, bucket, key):
37+
try:
38+
s3_client.head_object(Bucket=bucket, Key=key)
39+
return True
40+
except ClientError as e:
41+
if e.response["Error"]["Code"] in ("404", "NoSuchKey", "NotFound"):
42+
return False
43+
raise
44+
45+
46+
def test_cleanup_step_deletes_processing_output(
47+
sagemaker_session, pipeline_session, role, sklearn_latest_version
48+
):
49+
stamp = uuid.uuid4().hex[:8]
50+
bucket = sagemaker_session.default_bucket()
51+
prefix = f"integ-test-cleanup-step/{stamp}"
52+
output_uri = f"s3://{bucket}/{prefix}/output"
53+
region = pipeline_session.boto_region_name
54+
55+
processor = Processor(
56+
role=role,
57+
image_uri=image_uris.retrieve(
58+
framework="sklearn",
59+
region=region,
60+
version=sklearn_latest_version,
61+
py_version="py3",
62+
instance_type="ml.m5.large",
63+
),
64+
instance_count=1,
65+
instance_type="ml.m5.large",
66+
entrypoint=[
67+
"/bin/sh",
68+
"-c",
69+
"echo keep > /opt/ml/processing/output/keep.txt && "
70+
"echo delete > /opt/ml/processing/output/delete.txt",
71+
],
72+
sagemaker_session=pipeline_session,
73+
)
74+
step_process = ProcessingStep(
75+
name="WriteOutput",
76+
step_args=processor.run(
77+
outputs=[
78+
ProcessingOutput(
79+
output_name="output",
80+
s3_output=ProcessingS3Output(
81+
s3_uri=output_uri,
82+
local_path="/opt/ml/processing/output",
83+
s3_upload_mode="EndOfJob",
84+
),
85+
)
86+
],
87+
),
88+
)
89+
step_cleanup = CleanupStep(
90+
name="CleanupOutput",
91+
s3_sources=[S3CleanupSource(step=step_process, s3_uris=[f"{output_uri}/delete.txt"])],
92+
)
93+
pipeline = Pipeline(
94+
name=f"integ-cleanup-step-{stamp}",
95+
steps=[step_process, step_cleanup],
96+
sagemaker_session=pipeline_session,
97+
)
98+
99+
client = sagemaker_session.sagemaker_client
100+
s3_client = sagemaker_session.boto_session.client("s3")
101+
try:
102+
pipeline.upsert(role_arn=role)
103+
execution = pipeline.start()
104+
105+
final_status = None
106+
deadline = time.time() + 1200
107+
while time.time() < deadline:
108+
final_status = execution.describe()["PipelineExecutionStatus"]
109+
if final_status in ("Succeeded", "Failed", "Stopped", "CompletedWithFailures"):
110+
break
111+
time.sleep(30)
112+
113+
steps = client.list_pipeline_execution_steps(PipelineExecutionArn=execution.arn)[
114+
"PipelineExecutionSteps"
115+
]
116+
if final_status != "Succeeded":
117+
failures = "\n".join(
118+
f"{s['StepName']}: {s.get('FailureReason', 'no reason')}"
119+
for s in steps
120+
if s.get("StepStatus") == "Failed"
121+
)
122+
pytest.fail(f"Pipeline execution status={final_status}. Details:\n{failures}")
123+
124+
cleanup = next(s for s in steps if s["StepName"] == "CleanupOutput")
125+
assert cleanup["Metadata"]["Cleanup"]["CleanedSteps"] == ["WriteOutput"]
126+
assert not _object_exists(s3_client, bucket, f"{prefix}/output/delete.txt")
127+
assert _object_exists(s3_client, bucket, f"{prefix}/output/keep.txt")
128+
finally:
129+
try:
130+
pipeline.delete()
131+
except Exception:
132+
pass
133+
for obj in s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix).get("Contents", []):
134+
s3_client.delete_object(Bucket=bucket, Key=obj["Key"])

0 commit comments

Comments
 (0)