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.
 

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,7 +518,7 @@ def register(
 
         if image_uri is not None:
             self.image_uri = image_uri
-        if self.model_type is not JumpStartModelType.PROPRIETARY:
+        if 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
@@ -547,7 +548,7 @@ def register(
             if self.model_data is not None:
                 container_def["ModelDataUrl"] = self.model_data
 
-        if self.model_type is JumpStartModelType.PROPRIETARY:
+        if model_type is JumpStartModelType.PROPRIETARY:
             source_uri = self.model_package_arn
             model_package_group_name = self.model_id
 

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
-