fix: updating neo multiversion support to include edge devices

Author: mnuyens

src/sagemaker/model.py

@@ -69,6 +69,17 @@
 
 NEO_IOC_TARGET_DEVICES = ["ml_c4", "ml_c5", "ml_m4", "ml_m5", "ml_p2", "ml_p3", "ml_g4dn"]
 
+NEO_MULTIVERSION_UNSUPPORTED = [
+    "imx8mplus",
+    "jacinto_tda4vm",
+    "coreml",
+    "sitara_am57x",
+    "amba_cv2",
+    "amba_cv22",
+    "amba_cv25",
+    "lambda",
+]
+
 
 class ModelBase(abc.ABC):
     """An object that encapsulates a trained model.
@@ -836,11 +847,14 @@ def multi_version_compilation_supported(
                 multi_version_frameworks_support_mapping = {
                     "inferentia": ["pytorch", "tensorflow", "mxnet"],
                     "neo_ioc_targets": ["pytorch", "tensorflow"],
+                    "neo_edge_targets": ["pytorch", "tensorflow"],
                 }
                 if target_instance_type in NEO_IOC_TARGET_DEVICES:
                     return framework in multi_version_frameworks_support_mapping["neo_ioc_targets"]
                 if target_instance_type == "ml_inf1":
                     return framework in multi_version_frameworks_support_mapping["inferentia"]
+                if target_instance_type not in NEO_UNSUPPORTED_MULTI_VERSION:
+                    return framework in multi_version_frameworks_support_mapping["neo_edge_targets"]
             return False
 
         if multi_version_compilation_supported(target_instance_type, framework, framework_version):

tests/unit/sagemaker/model/test_neo.py

@@ -382,3 +382,42 @@ def test_compile_validates_framework_version(sagemaker_session):
     )
 
     assert model.image_uri is None
+
+    model = _create_model(sagemaker_session)
+    model.compile(
+        target_instance_family="rasp3b",
+        input_shape={"data": [1, 3, 1024, 1024]},
+        output_path="s3://output",
+        role="role",
+        framework="pytorch",
+        framework_version="1.6.1",
+        job_name="compile-model",
+    )
+
+    assert model.image_uri is None
+
+    config = model._compilation_job_config(
+        target_instance_family="rasp3b",
+        input_shape={"data": [1, 3, 1024, 1024]},
+        output_path="s3://output",
+        role="role",
+        framework="pytorch",
+        framework_version="1.6.1",
+        job_name="compile-model",
+        tags=None,
+    )
+
+    assert config["input_model_config"]["FrameworkVersion"] == "1.6"
+
+    config = model._compilation_job_config(
+        target_instance_family="amba_cv2",
+        input_shape={"data": [1, 3, 1024, 1024]},
+        output_path="s3://output",
+        role="role",
+        framework="pytorch",
+        framework_version="1.6.1",
+        job_name="compile-model",
+        tags=None,
+    )
+
+    assert config["input_model_config"].get("FrameworkVersion", None) is None