add cache config and unit test

Author: ahsan-z-khan

src/sagemaker/workflow/steps.py

@@ -64,6 +64,7 @@ class Step(Entity):
     Attributes:
         name (str): The name of the step.
         step_type (StepTypeEnum): The type of the step.
+
     """
 
     name: str = attr.ib(factory=str)
@@ -93,6 +94,26 @@ def ref(self) -> Dict[str, str]:
         return {"Name": self.name}
 
 
+@attr.s
+class CacheConfig:
+    """Step to cache pipeline workflow.
+
+    Attributes:
+        enable_caching (bool): To enable step caching. Off by default.
+        expire_after (str): If step caching is enabled, a timeout also needs to defined.
+            It defines how old a previous execution can be to be considered for reuse.
+            Needs to be ISO 8601 duration string.
+    """
+
+    enable_caching: bool = attr.ib(default=False)
+    expire_after: str = attr.ib(factory=str)
+
+    @property
+    def config(self):
+        """Enables caching in pipeline steps."""
+        return {"CacheConfig": {"Enabled": self.enable_caching, "ExpireAfter": self.expire_after}}
+
+
 class TrainingStep(Step):
     """Training step for workflow."""
 
@@ -101,6 +122,7 @@ def __init__(
         name: str,
         estimator: EstimatorBase,
         inputs: TrainingInput = None,
+        cache_config: CacheConfig = None,
     ):
         """Construct a TrainingStep, given an `EstimatorBase` instance.
 
@@ -111,14 +133,15 @@ def __init__(
             name (str): The name of the training step.
             estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
             inputs (TrainingInput): A `sagemaker.inputs.TrainingInput` instance. Defaults to `None`.
+            cache_config (CacheConfig):  An instance to enable caching.
         """
         super(TrainingStep, self).__init__(name, StepTypeEnum.TRAINING)
         self.estimator = estimator
         self.inputs = inputs
-
         self._properties = Properties(
             path=f"Steps.{name}", shape_name="DescribeTrainingJobResponse"
         )
+        self.cache_config = cache_config
 
     @property
     def arguments(self) -> RequestType:
@@ -145,6 +168,13 @@ def properties(self):
         """A Properties object representing the DescribeTrainingJobResponse data model."""
         return self._properties
 
+    def to_request(self) -> RequestType:
+        """Updates the dictionary with cache configuration."""
+        request_dict = super().to_request()
+        request_dict.update(self.cache_config.config)
+
+        return request_dict
+
 
 class CreateModelStep(Step):
     """CreateModel step for workflow."""
@@ -208,6 +238,7 @@ def __init__(
         name: str,
         transformer: Transformer,
         inputs: TransformInput,
+        cache_config: CacheConfig = None,
     ):
         """Constructs a TransformStep, given an `Transformer` instance.
 
@@ -218,11 +249,12 @@ def __init__(
             name (str): The name of the transform step.
             transformer (Transformer): A `sagemaker.transformer.Transformer` instance.
             inputs (TransformInput): A `sagemaker.inputs.TransformInput` instance.
+            cache_config (CacheConfig):  An instance to enable caching.
         """
         super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM)
         self.transformer = transformer
         self.inputs = inputs
-
+        self.cache_config = cache_config
         self._properties = Properties(
             path=f"Steps.{name}", shape_name="DescribeTransformJobResponse"
         )
@@ -258,6 +290,13 @@ def properties(self):
         """A Properties object representing the DescribeTransformJobResponse data model."""
         return self._properties
 
+    def to_request(self) -> RequestType:
+        """Updates the dictionary with cache configuration."""
+        request_dict = super().to_request()
+        request_dict.update(self.cache_config.config)
+
+        return request_dict
+
 
 class ProcessingStep(Step):
     """Processing step for workflow."""
@@ -271,6 +310,7 @@ def __init__(
         job_arguments: List[str] = None,
         code: str = None,
         property_files: List[PropertyFile] = None,
+        cache_config: CacheConfig = None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -290,6 +330,7 @@ def __init__(
                 script to run. Defaults to `None`.
             property_files (List[PropertyFile]): A list of property files that workflow looks
                 for and resolves from the configured processing output list.
+            cache_config (CacheConfig):  An instance to enable caching.
         """
         super(ProcessingStep, self).__init__(name, StepTypeEnum.PROCESSING)
         self.processor = processor
@@ -306,6 +347,7 @@ def __init__(
         self._properties = Properties(
             path=f"Steps.{name}", shape_name="DescribeProcessingJobResponse"
         )
+        self.cache_config = cache_config
 
     @property
     def arguments(self) -> RequestType:
@@ -336,6 +378,7 @@ def properties(self):
     def to_request(self) -> RequestType:
         """Get the request structure for workflow service calls."""
         request_dict = super(ProcessingStep, self).to_request()
+        request_dict.update(self.cache_config.config)
         if self.property_files:
             request_dict["PropertyFiles"] = [
                 property_file.expr for property_file in self.property_files

tests/integ/test_workflow.py

@@ -48,6 +48,7 @@
     CreateModelStep,
     ProcessingStep,
     TrainingStep,
+    CacheConfig
 )
 from sagemaker.workflow.step_collections import RegisterModel
 from sagemaker.workflow.pipeline import Pipeline
@@ -551,3 +552,80 @@ def test_training_job_with_debugger(
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_cache_hit_expired_entry(
+        sagemaker_session,
+        workflow_session,
+        region_name,
+        role,
+        script_dir,
+        pipeline_name,
+):
+
+    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_count = ParameterInteger(name="InstanceCount", default_value=1)
+
+
+    estimator =
+
+    step_train = TrainingStep(
+        name="my-train",
+        estimator=sklearn_train,
+        inputs=TrainingInput(
+            s3_data=step_process.properties.ProcessingOutputConfig.Outputs[
+                "train_data"
+            ].S3Output.S3Uri
+        ),
+        cache_config=
+    )
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_type, instance_count],
+        steps=[step_train],
+        sagemaker_session=workflow_session,
+    )
+
+    try:
+        # NOTE: We should exercise the case when role used in the pipeline execution is
+        # different than that required of the steps in the pipeline itself. The role in
+        # the pipeline definition needs to create training and processing jobs and other
+        # sagemaker entities. However, the jobs created in the steps themselves execute
+        # under a potentially different role, often requiring access to S3 and other
+        # artifacts not required to during creation of the jobs in the pipeline steps.
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role)
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+
+        execution = pipeline.start(parameters={})
+        assert re.match(
+            fr"arn:aws:sagemaker:{region}:\d{{12}}:pipeline/{pipeline_name}/execution/",
+            execution.arn,
+        )
+
+        response = execution.describe()
+        assert response["PipelineArn"] == create_arn
+
+        try:
+            execution.wait(delay=30, max_attempts=3)
+        except WaiterError:
+            pass
+        execution_steps = execution.list_steps()
+        assert len(execution_steps) == 1
+        assert execution_steps[0]["StepName"] == "sklearn-process"
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass

tests/unit/sagemaker/workflow/test_steps.py

@@ -37,6 +37,7 @@
     TrainingStep,
     TransformStep,
     CreateModelStep,
+    CacheConfig,
 )
 
 REGION = "us-west-2"
@@ -114,10 +115,9 @@ def test_training_step(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
     inputs = TrainingInput(f"s3://{BUCKET}/train_manifest")
+    cache_config = CacheConfig(enable_caching=False, expire_after="PT1H")
     step = TrainingStep(
-        name="MyTrainingStep",
-        estimator=estimator,
-        inputs=inputs,
+        name="MyTrainingStep", estimator=estimator, inputs=inputs, cache_config=cache_config
     )
     assert step.to_request() == {
         "Name": "MyTrainingStep",
@@ -145,6 +145,7 @@ def test_training_step(sagemaker_session):
             "RoleArn": ROLE,
             "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
         },
+        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
     }
     assert step.properties.TrainingJobName.expr == {"Get": "Steps.MyTrainingStep.TrainingJobName"}
 
@@ -163,11 +164,13 @@ def test_processing_step(sagemaker_session):
             destination="processing_manifest",
         )
     ]
+    cache_config = CacheConfig(enable_caching=False, expire_after="PT1H")
     step = ProcessingStep(
         name="MyProcessingStep",
         processor=processor,
         inputs=inputs,
         outputs=[],
+        cache_config=cache_config,
     )
     assert step.to_request() == {
         "Name": "MyProcessingStep",
@@ -197,6 +200,7 @@ def test_processing_step(sagemaker_session):
             },
             "RoleArn": "DummyRole",
         },
+        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
     }
     assert step.properties.ProcessingJobName.expr == {
         "Get": "Steps.MyProcessingStep.ProcessingJobName"
@@ -238,10 +242,9 @@ def test_transform_step(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
     inputs = TransformInput(data=f"s3://{BUCKET}/transform_manifest")
+    cache_config = CacheConfig(enable_caching=False, expire_after="PT1H")
     step = TransformStep(
-        name="MyTransformStep",
-        transformer=transformer,
-        inputs=inputs,
+        name="MyTransformStep", transformer=transformer, inputs=inputs, cache_config=cache_config
     )
     assert step.to_request() == {
         "Name": "MyTransformStep",
@@ -262,6 +265,7 @@ def test_transform_step(sagemaker_session):
                 "InstanceType": "c4.4xlarge",
             },
         },
+        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
     }
     assert step.properties.TransformJobName.expr == {
         "Get": "Steps.MyTransformStep.TransformJobName"