[fix] Check py_version existence in RegisterModel step (#2320)

* Check py_version existence in RegisterModel step

* black-format

Co-authored-by: Payton Staub <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>
Co-authored-by: icywang86rui <[email protected]>

Author: 4 people

src/sagemaker/workflow/_utils.py

@@ -307,7 +307,7 @@ def arguments(self) -> RequestType:
                     model._framework_name,
                     region_name,
                     version=model.framework_version,
-                    py_version=model.py_version,
+                    py_version=model.py_version if hasattr(model, "py_version") else None,
                     instance_type=self.kwargs.get("instance_type", self.estimator.instance_type),
                     accelerator_type=self.kwargs.get("accelerator_type"),
                     image_scope="inference",

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -28,6 +28,7 @@
 )
 
 from sagemaker.estimator import Estimator
+from sagemaker.tensorflow import TensorFlow
 from sagemaker.inputs import CreateModelInput, TransformInput
 from sagemaker.model_metrics import (
     MetricsSource,
@@ -120,6 +121,19 @@ def estimator(sagemaker_session):
     )
 
 
+@pytest.fixture
+def estimator_tf(sagemaker_session):
+    return TensorFlow(
+        entry_point="/some/script.py",
+        framework_version="1.15.2",
+        py_version="py3",
+        role=ROLE,
+        instance_type="ml.c4.2xlarge",
+        instance_count=1,
+        sagemaker_session=sagemaker_session,
+    )
+
+
 @pytest.fixture
 def model_metrics():
     return ModelMetrics(
@@ -202,6 +216,56 @@ def test_register_model(estimator, model_metrics):
     )
 
 
+def test_register_model_tf(estimator_tf, model_metrics):
+    model_data = f"s3://{BUCKET}/model.tar.gz"
+    register_model = RegisterModel(
+        name="RegisterModelStep",
+        estimator=estimator_tf,
+        model_data=model_data,
+        content_types=["content_type"],
+        response_types=["response_type"],
+        inference_instances=["inference_instance"],
+        transform_instances=["transform_instance"],
+        model_package_group_name="mpg",
+        model_metrics=model_metrics,
+        approval_status="Approved",
+        description="description",
+    )
+    assert ordered(register_model.request_dicts()) == ordered(
+        [
+            {
+                "Name": "RegisterModelStep",
+                "Type": "RegisterModel",
+                "Arguments": {
+                    "InferenceSpecification": {
+                        "Containers": [
+                            {
+                                "Image": "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-inference:1.15.2-cpu",
+                                "ModelDataUrl": f"s3://{BUCKET}/model.tar.gz",
+                            }
+                        ],
+                        "SupportedContentTypes": ["content_type"],
+                        "SupportedRealtimeInferenceInstanceTypes": ["inference_instance"],
+                        "SupportedResponseMIMETypes": ["response_type"],
+                        "SupportedTransformInstanceTypes": ["transform_instance"],
+                    },
+                    "ModelApprovalStatus": "Approved",
+                    "ModelMetrics": {
+                        "ModelQuality": {
+                            "Statistics": {
+                                "ContentType": "text/csv",
+                                "S3Uri": f"s3://{BUCKET}/metrics.csv",
+                            },
+                        },
+                    },
+                    "ModelPackageDescription": "description",
+                    "ModelPackageGroupName": "mpg",
+                },
+            },
+        ]
+    )
+
+
 def test_register_model_with_model_repack(estimator, model_metrics):
     model_data = f"s3://{BUCKET}/model.tar.gz"
     register_model = RegisterModel(