infer framework and its version for mxnet, pytorch, sklearn and tensorflow models

Author: rahven14

src/sagemaker/mxnet/model.py

@@ -238,8 +238,8 @@ def register(
             domain=domain,
             sample_payload_url=sample_payload_url,
             task=task,
-            framework=framework,
-            framework_version=framework_version,
+            framework=framework or self._framework_name,
+            framework_version=framework_version or self.framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

src/sagemaker/pytorch/model.py

@@ -239,8 +239,8 @@ def register(
             domain=domain,
             sample_payload_url=sample_payload_url,
             task=task,
-            framework=framework,
-            framework_version=framework_version,
+            framework=framework or self._framework_name,
+            framework_version=framework_version or self.framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

src/sagemaker/sklearn/model.py

@@ -233,8 +233,8 @@ def register(
             domain=domain,
             sample_payload_url=sample_payload_url,
             task=task,
-            framework=framework,
-            framework_version=framework_version,
+            framework=framework or self._framework_name,
+            framework_version=framework_version or self.framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

src/sagemaker/tensorflow/model.py

@@ -285,8 +285,8 @@ def register(
             domain=domain,
             sample_payload_url=sample_payload_url,
             task=task,
-            framework=framework,
-            framework_version=framework_version,
+            framework=framework or self._framework_name,
+            framework_version=framework_version or self.framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

tests/unit/test_sklearn.py

@@ -185,6 +185,23 @@ def test_create_model(sagemaker_session, sklearn_version):
     assert model_values["Image"] == image_uri
 
 
+def test_register_model(sagemaker_session, sklearn_version):
+    source_dir = "s3://mybucket/source"
+
+    sklearn_model = SKLearnModel(
+        model_data=source_dir,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        entry_point=SCRIPT_PATH,
+        framework_version=sklearn_version,
+    )
+    
+    model = sklearn_model.register(
+        content_types=["application/json"],
+        response_types=["application/json"],
+    )
+
+
 @patch("sagemaker.model.FrameworkModel._upload_code")
 def test_create_model_with_network_isolation(upload, sagemaker_session, sklearn_version):
     source_dir = "s3://mybucket/source"