fix: Update neo multiversion support to include edge devices (#3875)

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_MULTIVERSION_UNSUPPORTED:
+                    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,39 @@ def test_compile_validates_framework_version(sagemaker_session):
     )
 
     assert model.image_uri is None
+
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value={
+            "CompilationJobStatus": "Completed",
+            "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"},
+            "InferenceImage": None,
+        }
+    )
+
+    config = model._compilation_job_config(
+        "rasp3b",
+        {"data": [1, 3, 1024, 1024]},
+        "s3://output",
+        "role",
+        900,
+        "compile-model",
+        "pytorch",
+        None,
+        framework_version="1.6.1",
+    )
+
+    assert config["input_model_config"]["FrameworkVersion"] == "1.6"
+
+    config = model._compilation_job_config(
+        "amba_cv2",
+        {"data": [1, 3, 1024, 1024]},
+        "s3://output",
+        "role",
+        900,
+        "compile-model",
+        "pytorch",
+        None,
+        framework_version="1.6.1",
+    )
+
+    assert config["input_model_config"].get("FrameworkVersion", None) is None