Updated profiler_config and test

Marius Moisescu · Marius Moisescu · commit 8bbe6534c2a8 · 2022-12-13T20:13:33.000Z
diff --git a/src/sagemaker/debugger/profiler_config.py b/src/sagemaker/debugger/profiler_config.py
@@ -32,8 +32,8 @@ def __init__(
         s3_output_path: Optional[Union[str, PipelineVariable]] = None,
         system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None,
         framework_profile_params: Optional[FrameworkProfile] = None,
-        disable_profiler: Optional[FrameworkProfile] = False,
-    ):
+        disable_profiler: Optional[Union[str, PipelineVariable]] = False,
+    )
         """Initialize a ``ProfilerConfig`` instance.
 
         Pass the output of this class
diff --git a/tests/unit/sagemaker/workflow/test_steps.py b/tests/unit/sagemaker/workflow/test_steps.py
@@ -919,7 +919,7 @@ def test_create_model_step_with_model_pipeline(tfo, time, sagemaker_session):
                 },
                 {
                     "Environment": {"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},
-                    "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.4",
+                    "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:3.3",
                     "ModelDataUrl": "s3://bucket/model_2.tar.gz",
                 },
             ],

Original file line number	Diff line number	Diff line change
`@@ -919,7 +919,7 @@ def test_create_model_step_with_model_pipeline(tfo, time, sagemaker_session):`
`919`	`919`	`},`
`920`	`920`	`{`
`921`	`921`	`"Environment": {"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},`
`922`		`- "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.4",`
	`922`	`+ "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:3.3",`
`923`	`923`	`"ModelDataUrl": "s3://bucket/model_2.tar.gz",`
`924`	`924`	`},`
`925`	`925`	`],`