change: enable neo framework version support on ml_inf2 and ml_trn1 (#3909)

Author: HappyAmazonian

src/sagemaker/model.py

@@ -844,15 +844,20 @@ def multi_version_compilation_supported(
         ):
             if target_instance_type and framework and framework_version:
                 framework = framework.lower()
+
                 multi_version_frameworks_support_mapping = {
-                    "inferentia": ["pytorch", "tensorflow", "mxnet"],
+                    "ml_inf1": ["pytorch", "tensorflow", "mxnet"],
+                    "ml_inf2": ["pytorch", "tensorflow"],
+                    "ml_trn1": ["pytorch", "tensorflow"],
                     "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 in ["ml_inf1", "ml_inf2", "ml_trn1"]:
+                    return (
+                        framework in multi_version_frameworks_support_mapping[target_instance_type]
+                    )
                 if target_instance_type not in NEO_MULTIVERSION_UNSUPPORTED:
                     return framework in multi_version_frameworks_support_mapping["neo_edge_targets"]
             return False

tests/unit/sagemaker/model/test_neo.py

@@ -362,59 +362,35 @@ def test_compile_with_tensorflow_neo_in_ml_inf(session):
     )
 
 
-def test_compile_validates_framework_version(sagemaker_session):
-    sagemaker_session.wait_for_compilation_job = Mock(
-        return_value={
-            "CompilationJobStatus": "Completed",
-            "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"},
-            "InferenceImage": None,
-        }
-    )
+@pytest.mark.parametrize(
+    "target,framework,fx_version,expected_fx_version",
+    [
+        ("ml_c4", "pytorch", "1.6", "1.6"),
+        ("rasp3b", "pytorch", "1.6.1", "1.6"),
+        ("amba_cv2", "pytorch", "1.6.1", None),
+        ("ml_c4", "tensorflow", "1.15.1", "1.15"),
+        ("ml_c4", "tensorflow", "2.15.1", "2.15"),
+        ("ml_inf1", "tensorflow", "2.15.1", "2.15"),
+        ("ml_inf2", "pytorch", "2.0", "2.0"),
+        ("ml_inf2", "pytorch", "2.0.1", "2.0"),
+        ("ml_trn1", "pytorch", "2.0.1", "2.0"),
+        ("ml_trn1", "tensorflow", "2.0.1", "2.0"),
+    ],
+)
+def test_compile_validates_framework_version(
+    sagemaker_session, target, framework, fx_version, expected_fx_version
+):
     model = _create_model(sagemaker_session)
-    model.compile(
-        target_instance_family="ml_c4",
-        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
-
-    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",
+        target,
         {"data": [1, 3, 1024, 1024]},
         "s3://output",
         "role",
         900,
         "compile-model",
-        "pytorch",
+        framework,
         None,
-        framework_version="1.6.1",
+        framework_version=fx_version,
     )
 
-    assert config["input_model_config"].get("FrameworkVersion", None) is None
+    assert config["input_model_config"].get("FrameworkVersion", None) == expected_fx_version