@@ -50,7 +50,6 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
5050
5151 CONDITION = "Condition"
5252 CREATE_MODEL = "Model"
53- FAIL = "Fail"
5453 PROCESSING = "Processing"
5554 REGISTER_MODEL = "RegisterModel"
5655 TRAINING = "Training"
@@ -93,6 +92,38 @@ def ref(self) -> Dict[str, str]:
9392 return {"Name" : self .name }
9493
9594
95+ @attr .s
96+ class CacheConfig :
97+ """Configuration class to enable caching in pipeline workflow.
98+
99+ If caching is enabled, the pipeline attempts to find a previous execution of a step
100+ that was called with the same arguments. Step caching only considers successful execution.
101+ If a successful previous execution is found, the pipeline propagates the values
102+ from previous execution rather than recomputing the step. When multiple successful executions
103+ exist within the timeout period, it uses the result for the most recent successful execution.
104+
105+
106+ Attributes:
107+ enable_caching (bool): To enable step caching. Defaults to `False`.
108+ expire_after (str): If step caching is enabled, a timeout also needs to defined.
109+ It defines how old a previous execution can be to be considered for reuse.
110+ Value should be an ISO 8601 duration string. Defaults to `None`.
111+ """
112+
113+ enable_caching : bool = attr .ib (default = False )
114+ expire_after = attr .ib (
115+ default = None , validator = attr .validators .optional (attr .validators .instance_of (str ))
116+ )
117+
118+ @property
119+ def config (self ):
120+ """Configures caching in pipeline steps."""
121+ config = {"Enabled" : self .enable_caching }
122+ if self .expire_after is not None :
123+ config ["ExpireAfter" ] = self .expire_after
124+ return {"CacheConfig" : config }
125+
126+
96127class TrainingStep (Step ):
97128 """Training step for workflow."""
98129
@@ -101,6 +132,7 @@ def __init__(
101132 name : str ,
102133 estimator : EstimatorBase ,
103134 inputs : TrainingInput = None ,
135+ cache_config : CacheConfig = None ,
104136 ):
105137 """Construct a TrainingStep, given an `EstimatorBase` instance.
106138
@@ -111,14 +143,15 @@ def __init__(
111143 name (str): The name of the training step.
112144 estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
113145 inputs (TrainingInput): A `sagemaker.inputs.TrainingInput` instance. Defaults to `None`.
146+ cache_config (CacheConfig): A `sagemaker.workflow.steps.CacheConfig` instance.
114147 """
115148 super (TrainingStep , self ).__init__ (name , StepTypeEnum .TRAINING )
116149 self .estimator = estimator
117150 self .inputs = inputs
118-
119151 self ._properties = Properties (
120152 path = f"Steps.{ name } , shape_name = "DescribeTrainingJobResponse"
121153 )
154+ self .cache_config = cache_config
122155
123156 @property
124157 def arguments (self ) -> RequestType :
@@ -145,6 +178,14 @@ def properties(self):
145178 """A Properties object representing the DescribeTrainingJobResponse data model."""
146179 return self ._properties
147180
181+ def to_request (self ) -> RequestType :
182+ """Updates the dictionary with cache configuration."""
183+ request_dict = super ().to_request ()
184+ if self .cache_config :
185+ request_dict .update (self .cache_config .config )
186+
187+ return request_dict
188+
148189
149190class CreateModelStep (Step ):
150191 """CreateModel step for workflow."""
@@ -208,6 +249,7 @@ def __init__(
208249 name : str ,
209250 transformer : Transformer ,
210251 inputs : TransformInput ,
252+ cache_config : CacheConfig = None ,
211253 ):
212254 """Constructs a TransformStep, given an `Transformer` instance.
213255
@@ -218,11 +260,12 @@ def __init__(
218260 name (str): The name of the transform step.
219261 transformer (Transformer): A `sagemaker.transformer.Transformer` instance.
220262 inputs (TransformInput): A `sagemaker.inputs.TransformInput` instance.
263+ cache_config (CacheConfig): A `sagemaker.workflow.steps.CacheConfig` instance.
221264 """
222265 super (TransformStep , self ).__init__ (name , StepTypeEnum .TRANSFORM )
223266 self .transformer = transformer
224267 self .inputs = inputs
225-
268+ self . cache_config = cache_config
226269 self ._properties = Properties (
227270 path = f"Steps.{ name } , shape_name = "DescribeTransformJobResponse"
228271 )
@@ -258,6 +301,14 @@ def properties(self):
258301 """A Properties object representing the DescribeTransformJobResponse data model."""
259302 return self ._properties
260303
304+ def to_request (self ) -> RequestType :
305+ """Updates the dictionary with cache configuration."""
306+ request_dict = super ().to_request ()
307+ if self .cache_config :
308+ request_dict .update (self .cache_config .config )
309+
310+ return request_dict
311+
261312
262313class ProcessingStep (Step ):
263314 """Processing step for workflow."""
@@ -271,6 +322,7 @@ def __init__(
271322 job_arguments : List [str ] = None ,
272323 code : str = None ,
273324 property_files : List [PropertyFile ] = None ,
325+ cache_config : CacheConfig = None ,
274326 ):
275327 """Construct a ProcessingStep, given a `Processor` instance.
276328
@@ -290,6 +342,7 @@ def __init__(
290342 script to run. Defaults to `None`.
291343 property_files (List[PropertyFile]): A list of property files that workflow looks
292344 for and resolves from the configured processing output list.
345+ cache_config (CacheConfig): A `sagemaker.workflow.steps.CacheConfig` instance.
293346 """
294347 super (ProcessingStep , self ).__init__ (name , StepTypeEnum .PROCESSING )
295348 self .processor = processor
@@ -306,6 +359,7 @@ def __init__(
306359 self ._properties = Properties (
307360 path = f"Steps.{ name } , shape_name = "DescribeProcessingJobResponse"
308361 )
362+ self .cache_config = cache_config
309363
310364 @property
311365 def arguments (self ) -> RequestType :
@@ -336,49 +390,10 @@ def properties(self):
336390 def to_request (self ) -> RequestType :
337391 """Get the request structure for workflow service calls."""
338392 request_dict = super (ProcessingStep , self ).to_request ()
393+ if self .cache_config :
394+ request_dict .update (self .cache_config .config )
339395 if self .property_files :
340396 request_dict ["PropertyFiles" ] = [
341397 property_file .expr for property_file in self .property_files
342398 ]
343399 return request_dict
344-
345-
346- class FailStep (Step ):
347- """Pipeline step to indicate failure."""
348-
349- def __init__ (self , name : str = "Fail" ):
350- """Construct a FailStep.
351-
352- Causes the pipeline execution to terminate in a failed state.
353-
354- Args:
355- name (str): The name of the step.
356- """
357- super (FailStep , self ).__init__ (name , StepTypeEnum .FAIL )
358- root_path = f"Steps.{ name }
359- root_prop = Properties (path = root_path )
360- root_prop .__dict__ ["Fail" ] = Properties (f"{ root_path } )
361- self ._properties = root_prop
362-
363- @property
364- def arguments (self ) -> RequestType :
365- """The arguments to the particular step service call."""
366- return {}
367-
368- @property
369- def properties (self ):
370- """The properties of the particular step."""
371- return self ._properties
372-
373- def to_request (self ) -> RequestType :
374- """Get the request structure for workflow service calls."""
375- return {
376- "Name" : self .name ,
377- "Type" : self .step_type .value ,
378- "Arguments" : self .arguments ,
379- }
380-
381- @property
382- def ref (self ) -> Dict [str , str ]:
383- """Get a reference dict for steps"""
384- return {"Name" : self .name }
0 commit comments