fix unit test

Author: knikure

src/sagemaker/fw_utils.py

@@ -913,9 +913,10 @@ def _instance_type_supports_profiler(instance_type):
     Returns:
         bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
     """
-    match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
-    if match and match[1].startswith("trn"):
-        return False
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match and match[1].startswith("trn"):
+            return False
     return True
 
 

src/sagemaker/image_uris.py

@@ -152,6 +152,7 @@ def retrieve(
             inference_tool = _get_inference_tool(inference_tool, instance_type)
             if inference_tool == "neuron":
                 _framework = f"{framework}-{inference_tool}"
+        _validate_for_suppported_frameworks_and_instance_type(framework, instance_type)
         config = _config_for_framework_and_scope(_framework, image_scope, accelerator_type)
 
     original_version = version
@@ -283,13 +284,20 @@ def _config_for_framework_and_scope(framework, image_scope, accelerator_type=Non
         )
         image_scope = available_scopes[0]
 
-    if image_scope is None and framework not in TRAINIUM_ALLOWED_FRAMEWORKS:
-        _validate_framework(framework, TRAINIUM_ALLOWED_FRAMEWORKS, "framework")
-
     _validate_arg(image_scope, available_scopes, "image scope")
     return config if "scope" in config else config[image_scope]
 
 
+def _validate_for_suppported_frameworks_and_instance_type(framework, instace_type):
+    """Validate if framework is supported for the instance_type"""
+    if (
+        instace_type is not None
+        and "trn" in instace_type
+        and framework not in TRAINIUM_ALLOWED_FRAMEWORKS
+    ):
+        _validate_framework(framework, TRAINIUM_ALLOWED_FRAMEWORKS, "framework")
+
+
 def config_for_framework(framework):
     """Loads the JSON config for the given framework."""
     fname = os.path.join(os.path.dirname(__file__), "image_uri_config", "{}.json".format(framework))