fix: CreateModelPackage API error for Scikit-learn and XGBoost frameworkss (aws#3400)

Author: knikure

src/sagemaker/sklearn/model.py

@@ -248,7 +248,7 @@ def register(
             sample_payload_url=sample_payload_url,
             task=task,
             framework=framework,
-            framework_version=framework_version or self.framework_version,
+            framework_version=framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

src/sagemaker/xgboost/model.py

@@ -236,7 +236,7 @@ def register(
             sample_payload_url=sample_payload_url,
             task=task,
             framework=framework,
-            framework_version=framework_version or self.framework_version,
+            framework_version=framework_version,
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
         )

tests/unit/test_sklearn.py

@@ -17,7 +17,7 @@
 import os
 import pytest
 
-from mock import Mock, ANY
+from mock import Mock
 from mock import patch
 
 from sagemaker.fw_utils import UploadedCode
@@ -587,49 +587,3 @@ def test_model_py2_raises(sagemaker_session, sklearn_version):
             framework_version=sklearn_version,
             py_version="py2",
         )
-
-
-def test_register_sklearn_model_auto_infer_framework(sagemaker_session, sklearn_version):
-    source_dir = "s3://mybucket/source"
-
-    model_package_group_name = "test-sklearn-register-model"
-    content_types = ["application/json"]
-    response_types = ["application/json"]
-    image_uri = "fakeimage"
-
-    sklearn_model = SKLearnModel(
-        model_data=source_dir,
-        role=ROLE,
-        sagemaker_session=sagemaker_session,
-        entry_point=SCRIPT_PATH,
-        framework_version=sklearn_version,
-    )
-
-    sklearn_model.register(
-        content_types,
-        response_types,
-        model_package_group_name=model_package_group_name,
-        marketplace_cert=True,
-        image_uri=image_uri,
-    )
-
-    expected_create_model_package_request = {
-        "containers": [
-            {
-                "Image": image_uri,
-                "Environment": ANY,
-                "ModelDataUrl": source_dir,
-                "Framework": "SKLEARN",
-                "FrameworkVersion": sklearn_version,
-            },
-        ],
-        "content_types": content_types,
-        "response_types": response_types,
-        "inference_instances": None,
-        "transform_instances": None,
-        "model_package_group_name": model_package_group_name,
-        "marketplace_cert": True,
-    }
-    sagemaker_session.create_model_package_from_containers.assert_called_with(
-        **expected_create_model_package_request
-    )

tests/unit/test_xgboost.py

@@ -19,7 +19,6 @@
 
 from mock import Mock
 from mock import patch
-from mock import ANY
 from packaging.version import Version
 
 
@@ -673,54 +672,3 @@ def test_unsupported_xgboost_version_error(sagemaker_session):
     error_message = "XGBoost 1.1 is not supported"
     assert error_message in str(error1)
     assert error_message in str(error2)
-
-
-def test_register_xgboost_model_auto_infer_framework(sagemaker_session, xgboost_framework_version):
-    source_dir = "s3://mybucket/source"
-
-    model_package_group_name = "test-pytorch-register-model"
-    content_types = ["application/json"]
-    response_types = ["application/json"]
-    inference_instances = ["ml.m4.xlarge"]
-    transform_instances = ["ml.m4.xlarge"]
-    image_uri = "fakeimage"
-
-    xgboost_model = XGBoostModel(
-        model_data=source_dir,
-        role=ROLE,
-        sagemaker_session=sagemaker_session,
-        entry_point=SCRIPT_PATH,
-        framework_version=xgboost_framework_version,
-    )
-
-    xgboost_model.register(
-        content_types,
-        response_types,
-        inference_instances,
-        transform_instances,
-        model_package_group_name=model_package_group_name,
-        marketplace_cert=True,
-        image_uri=image_uri,
-    )
-
-    expected_create_model_package_request = {
-        "containers": [
-            {
-                "Image": image_uri,
-                "Environment": ANY,
-                "ModelDataUrl": ANY,
-                "Framework": "XGBOOST",
-                "FrameworkVersion": xgboost_framework_version,
-            },
-        ],
-        "content_types": content_types,
-        "response_types": response_types,
-        "inference_instances": inference_instances,
-        "transform_instances": transform_instances,
-        "model_package_group_name": model_package_group_name,
-        "marketplace_cert": True,
-    }
-
-    sagemaker_session.create_model_package_from_containers.assert_called_with(
-        **expected_create_model_package_request
-    )