|
| 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.") |
0 commit comments