modify unit test to enable cache to True

Author: ahsan-z-khan

tests/integ/test_workflow.py

@@ -289,6 +289,8 @@ def test_one_step_sklearn_processing_pipeline(
         ProcessingInput(dataset_definition=athena_dataset_definition),
     ]
 
+    cache_config = CacheConfig(enable_caching=True, expire_after="T30m")
+
     sklearn_processor = SKLearnProcessor(
         framework_version=sklearn_latest_version,
         role=role,
@@ -304,6 +306,7 @@ def test_one_step_sklearn_processing_pipeline(
         processor=sklearn_processor,
         inputs=inputs,
         code=script_path,
+        cache_config=cache_config,
     )
     pipeline = Pipeline(
         name=pipeline_name,
@@ -343,6 +346,11 @@ def test_one_step_sklearn_processing_pipeline(
         response = execution.describe()
         assert response["PipelineArn"] == create_arn
 
+        # Check CacheConfig
+        response = json.loads(pipeline.describe()["PipelineDefinition"])["Steps"][0]["CacheConfig"]
+        assert response["Enabled"] == cache_config.enable_caching
+        assert response["ExpireAfter"] == cache_config.expire_after
+
         try:
             execution.wait(delay=30, max_attempts=3)
         except WaiterError:
@@ -547,213 +555,3 @@ def test_training_job_with_debugger(
             pipeline.delete()
         except Exception:
             pass
-
-
-def test_cache_hit(
-    sagemaker_session,
-    workflow_session,
-    region_name,
-    role,
-    script_dir,
-    pipeline_name,
-    athena_dataset_definition,
-):
-
-    cache_config = CacheConfig(enable_caching=True, expire_after="T30m")
-
-    framework_version = "0.20.0"
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
-    instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-
-    input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv"
-
-    sklearn_processor = SKLearnProcessor(
-        framework_version=framework_version,
-        instance_type=instance_type,
-        instance_count=instance_count,
-        base_job_name="test-sklearn",
-        sagemaker_session=sagemaker_session,
-        role=role,
-    )
-
-    step_process = ProcessingStep(
-        name="my-cache-test",
-        processor=sklearn_processor,
-        inputs=[
-            ProcessingInput(source=input_data, destination="/opt/ml/processing/input"),
-            ProcessingInput(dataset_definition=athena_dataset_definition),
-        ],
-        outputs=[
-            ProcessingOutput(output_name="train_data", source="/opt/ml/processing/train"),
-            ProcessingOutput(output_name="test_data", source="/opt/ml/processing/test"),
-        ],
-        code=os.path.join(script_dir, "preprocessing.py"),
-        cache_config=cache_config,
-    )
-
-    pipeline = Pipeline(
-        name=pipeline_name,
-        parameters=[instance_count, instance_type],
-        steps=[step_process],
-        sagemaker_session=workflow_session,
-    )
-
-    try:
-        response = pipeline.create(role)
-        create_arn = response["PipelineArn"]
-        pytest.set_trace()
-
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
-            create_arn,
-        )
-
-        # Run pipeline for the first time to get an entry in the cache
-        execution1 = pipeline.start(parameters={})
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
-            execution1.arn,
-        )
-
-        response = execution1.describe()
-        assert response["PipelineArn"] == create_arn
-
-        try:
-            execution1.wait(delay=30, max_attempts=10)
-        except WaiterError:
-            pass
-        execution1_steps = execution1.list_steps()
-        assert len(execution1_steps) == 1
-        assert execution1_steps[0]["StepName"] == "my-cache-test"
-
-        # Run pipeline for the second time and expect cache hit
-        execution2 = pipeline.start(parameters={})
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
-            execution2.arn,
-        )
-
-        response = execution2.describe()
-        assert response["PipelineArn"] == create_arn
-
-        try:
-            execution2.wait(delay=30, max_attempts=10)
-        except WaiterError:
-            pass
-        execution2_steps = execution2.list_steps()
-        assert len(execution2_steps) == 1
-        assert execution2_steps[0]["StepName"] == "my-cache-test"
-
-        assert execution1_steps[0] == execution2_steps[0]
-
-    finally:
-        try:
-            pipeline.delete()
-        except Exception:
-            pass
-
-
-def test_cache_expiry(
-    sagemaker_session,
-    workflow_session,
-    region_name,
-    role,
-    script_dir,
-    pipeline_name,
-    athena_dataset_definition,
-):
-
-    cache_config = CacheConfig(enable_caching=True, expire_after="T1m")
-
-    framework_version = "0.20.0"
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
-    instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-
-    input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv"
-
-    sklearn_processor = SKLearnProcessor(
-        framework_version=framework_version,
-        instance_type=instance_type,
-        instance_count=instance_count,
-        base_job_name="test-sklearn",
-        sagemaker_session=sagemaker_session,
-        role=role,
-    )
-
-    step_process = ProcessingStep(
-        name="my-cache-test-expiry",
-        processor=sklearn_processor,
-        inputs=[
-            ProcessingInput(source=input_data, destination="/opt/ml/processing/input"),
-            ProcessingInput(dataset_definition=athena_dataset_definition),
-        ],
-        outputs=[
-            ProcessingOutput(output_name="train_data", source="/opt/ml/processing/train"),
-            ProcessingOutput(output_name="test_data", source="/opt/ml/processing/test"),
-        ],
-        code=os.path.join(script_dir, "preprocessing.py"),
-        cache_config=cache_config,
-    )
-
-    pipeline = Pipeline(
-        name=pipeline_name,
-        parameters=[instance_count, instance_type],
-        steps=[step_process],
-        sagemaker_session=workflow_session,
-    )
-
-    try:
-        response = pipeline.create(role)
-        create_arn = response["PipelineArn"]
-
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
-            create_arn,
-        )
-
-        # Run pipeline for the first time to get an entry in the cache
-        execution1 = pipeline.start(parameters={})
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
-            execution1.arn,
-        )
-
-        response = execution1.describe()
-        assert response["PipelineArn"] == create_arn
-
-        try:
-            execution1.wait(delay=30, max_attempts=3)
-        except WaiterError:
-            pass
-        execution1_steps = execution1.list_steps()
-        assert len(execution1_steps) == 1
-        assert execution1_steps[0]["StepName"] == "my-cache-test-expiry"
-
-        # wait 1 minute for cache to expire
-        time.sleep(60)
-
-        # Run pipeline for the second time and expect cache miss
-        execution2 = pipeline.start(parameters={})
-        assert re.match(
-            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
-            execution2.arn,
-        )
-
-        response = execution2.describe()
-        assert response["PipelineArn"] == create_arn
-
-        try:
-            execution2.wait(delay=30, max_attempts=3)
-        except WaiterError:
-            pass
-        execution2_steps = execution2.list_steps()
-        assert len(execution2_steps) == 1
-        assert execution2_steps[0]["StepName"] == "my-cache-test-expiry"
-
-        assert execution1_steps[0] != execution2_steps[0]
-
-    finally:
-        try:
-            pipeline.delete()
-        except Exception:
-            pass

tests/unit/sagemaker/workflow/test_steps.py

@@ -115,7 +115,7 @@ 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")
+    cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
     step = TrainingStep(
         name="MyTrainingStep", estimator=estimator, inputs=inputs, cache_config=cache_config
     )
@@ -145,7 +145,7 @@ def test_training_step(sagemaker_session):
             "RoleArn": ROLE,
             "StoppingCondition": {"MaxRuntimeInSeconds": 86400},
         },
-        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
+        "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }
     assert step.properties.TrainingJobName.expr == {"Get": "Steps.MyTrainingStep.TrainingJobName"}
 
@@ -164,7 +164,7 @@ def test_processing_step(sagemaker_session):
             destination="processing_manifest",
         )
     ]
-    cache_config = CacheConfig(enable_caching=False, expire_after="PT1H")
+    cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
     step = ProcessingStep(
         name="MyProcessingStep",
         processor=processor,
@@ -200,7 +200,7 @@ def test_processing_step(sagemaker_session):
             },
             "RoleArn": "DummyRole",
         },
-        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
+        "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }
     assert step.properties.ProcessingJobName.expr == {
         "Get": "Steps.MyProcessingStep.ProcessingJobName"
@@ -242,7 +242,7 @@ 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")
+    cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
     step = TransformStep(
         name="MyTransformStep", transformer=transformer, inputs=inputs, cache_config=cache_config
     )
@@ -265,7 +265,7 @@ def test_transform_step(sagemaker_session):
                 "InstanceType": "c4.4xlarge",
             },
         },
-        "CacheConfig": {"Enabled": False, "ExpireAfter": "PT1H"},
+        "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }
     assert step.properties.TransformJobName.expr == {
         "Get": "Steps.MyTransformStep.TransformJobName"