fix: register jumpstart models on model registry

Author: selvask-aws

src/sagemaker/jumpstart/model.py

@@ -761,7 +761,6 @@ def register(
         source_uri: Optional[Union[str, PipelineVariable]] = None,
         model_card: Optional[Union[ModelPackageModelCard, ModelCard]] = None,
         accept_eula: Optional[bool] = None,
-
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -820,8 +819,9 @@ def register(
             A `sagemaker.model.ModelPackage` instance.
         """
 
-        if model_package_group_name is None and self.model_type is JumpStartModelType.PROPRIETARY:
+        if model_package_group_name is None:
             model_package_group_name = self.model_id
+        if self.model_type is JumpStartModelType.PROPRIETARY:
             source_uri = self.model_package_arn
 
         register_kwargs = get_register_kwargs(

src/sagemaker/jumpstart/types.py

@@ -2410,7 +2410,6 @@ class JumpStartModelRegisterKwargs(JumpStartKwargs):
         "model_version",
         "hub_arn",
         "sagemaker_session",
-        "model_type",
     }
 
     def __init__(

src/sagemaker/model.py

@@ -450,6 +450,7 @@ def register(
         source_uri: Optional[Union[str, PipelineVariable]] = None,
         model_card: Optional[Union[ModelPackageModelCard, ModelCard]] = None,
         accept_eula: Optional[bool] = None,
+        model_type: Optional[JumpStartModelType] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -517,22 +518,22 @@ def register(
 
         if image_uri is not None:
             self.image_uri = image_uri
-        if self.model_type is not JumpStartModelType.PROPRIETARY:
-            if model_package_group_name is None and model_package_name is None:
-                # If model package group and model package name is not set
-                # then register to auto-generated model package group
-                model_package_group_name = utils.base_name_from_image(
-                    self.image_uri, default_base_name=ModelPackage.__name__
-                )
-            if model_package_group_name is not None:
-                container_def = self.prepare_container_def(accept_eula=accept_eula)
-                container_def = update_container_with_inference_params(
-                    framework=framework,
-                    framework_version=framework_version,
-                    nearest_model_name=nearest_model_name,
-                    data_input_configuration=data_input_configuration,
-                    container_def=container_def,
-                )
+
+        if model_package_group_name is None and model_package_name is None:
+            # If model package group and model package name is not set
+            # then register to auto-generated model package group
+            model_package_group_name = utils.base_name_from_image(
+                self.image_uri, default_base_name=ModelPackage.__name__
+            )
+        if model_package_group_name is not None:
+            container_def = self.prepare_container_def(accept_eula=accept_eula)
+            container_def = update_container_with_inference_params(
+                framework=framework,
+                framework_version=framework_version,
+                nearest_model_name=nearest_model_name,
+                data_input_configuration=data_input_configuration,
+                container_def=container_def,
+            )
         else:
             container_def = {
                 "Image": self.image_uri,
@@ -547,10 +548,6 @@ def register(
             if self.model_data is not None:
                 container_def["ModelDataUrl"] = self.model_data
 
-        if self.model_type is JumpStartModelType.PROPRIETARY:
-            source_uri = self.model_package_arn
-            model_package_group_name = self.model_id
-
         model_pkg_args = sagemaker.get_model_package_args(
             self.content_types,
             self.response_types,

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -294,6 +294,7 @@ def test_proprietary_jumpstart_model(setup):
 
     assert response is not None
 
+
 @pytest.mark.skipif(
     True,
     reason="Only enable if test account is subscribed to the proprietary model",
@@ -309,7 +310,6 @@ def test_register_proprietary_jumpstart_model(setup):
         sagemaker_session=get_sm_session(),
     )
     model_package = model.register()
-    
 
     predictor = model_package.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}]
@@ -329,7 +329,7 @@ def test_register_proprietary_jumpstart_model(setup):
 )
 def test_register_gated_jumpstart_model(setup):
 
-    model_id="meta-textgenerationneuron-llama-2-7b"
+    model_id = "meta-textgenerationneuron-llama-2-7b"
     model = JumpStartModel(
         model_id=model_id,
         model_version="1.1.0",
@@ -339,7 +339,8 @@ def test_register_gated_jumpstart_model(setup):
     model_package = model.register(accept_eula=True)
 
     predictor = model_package.deploy(
-        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], accept_eula=True
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        accept_eula=True,
     )
     payload = {"prompt": "To be, or", "maxTokens": 4, "temperature": 0, "numResults": 1}
 
@@ -348,4 +349,3 @@ def test_register_gated_jumpstart_model(setup):
     predictor.delete_predictor()
 
     assert response is not None
-

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -513,8 +513,11 @@ def test_proprietary_model_endpoint(
         model.deploy()
 
         mock_model_register.assert_called_once_with(
+            model_type=JumpStartModelType.PROPRIETARY,
             content_types=["application/json"],
             response_types=["application/json"],
+            model_package_group_name=model_id,
+            source_uri=model.model_package_arn,
         )
 
         mock_model_deploy.assert_called_once_with(
@@ -1416,8 +1419,10 @@ def test_model_registry_accept_and_response_types(
         model.register()
 
         mock_model_register.assert_called_once_with(
+            model_type=JumpStartModelType.OPEN_WEIGHTS,
             content_types=["application/x-text"],
             response_types=["application/json;verbose", "application/json"],
+            model_package_group_name=model.model_id,
         )
 
     @mock.patch("sagemaker.jumpstart.model.get_default_predictor")