Updated profiler_config and test

Marius Moisescu · Marius Moisescu · commit c16a7c097192 · 2022-12-14T22:42:11.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
@@ -874,7 +874,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
`@@ -874,7 +874,7 @@ def test_create_model_step_with_model_pipeline(tfo, time, sagemaker_session):`
`874`	`874`	`},`
`875`	`875`	`{`
`876`	`876`	`"Environment": {"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},`
`877`		`- "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.4",`
	`877`	`+ "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:3.3",`
`878`	`878`	`"ModelDataUrl": "s3://bucket/model_2.tar.gz",`
`879`	`879`	`},`
`880`	`880`	`],`