Feature: register proprietary models from jumpstart

Author: selvask-aws

src/sagemaker/jumpstart/factory/model.py

@@ -623,6 +623,7 @@ def get_deploy_kwargs(
 def get_register_kwargs(
     model_id: str,
     model_version: Optional[str] = None,
+    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_WEIGHTS,
     region: Optional[str] = None,
     tolerate_deprecated_model: Optional[bool] = None,
     tolerate_vulnerable_model: Optional[bool] = None,
@@ -656,6 +657,7 @@ def get_register_kwargs(
     register_kwargs = JumpStartModelRegisterKwargs(
         model_id=model_id,
         model_version=model_version,
+        model_type=model_type,
         region=region,
         tolerate_deprecated_model=tolerate_deprecated_model,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -688,6 +690,7 @@ def get_register_kwargs(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        model_type=model_type,
         region=region,
         scope=JumpStartScriptScope.INFERENCE,
         sagemaker_session=sagemaker_session,

src/sagemaker/jumpstart/model.py

@@ -796,6 +796,7 @@ def register(
         register_kwargs = get_register_kwargs(
             model_id=self.model_id,
             model_version=self.model_version,
+            model_type=self.model_type,
             region=self.region,
             tolerate_deprecated_model=self.tolerate_deprecated_model,
             tolerate_vulnerable_model=self.tolerate_vulnerable_model,

src/sagemaker/jumpstart/types.py

@@ -2094,6 +2094,7 @@ class JumpStartModelRegisterKwargs(JumpStartKwargs):
         "tolerate_deprecated_model",
         "region",
         "model_id",
+        "model_type",
         "model_version",
         "sagemaker_session",
         "content_types",
@@ -2128,13 +2129,15 @@ class JumpStartModelRegisterKwargs(JumpStartKwargs):
         "model_id",
         "model_version",
         "sagemaker_session",
+        "model_type",
     }
 
     def __init__(
         self,
         model_id: str,
         model_version: Optional[str] = None,
         region: Optional[str] = None,
+        model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_WEIGHTS,
         tolerate_deprecated_model: Optional[bool] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
         sagemaker_session: Optional[Any] = None,
@@ -2166,6 +2169,7 @@ def __init__(
 
         self.model_id = model_id
         self.model_version = model_version
+        self.model_type = model_type
         self.region = region
         self.image_uri = image_uri
         self.sagemaker_session = sagemaker_session

src/sagemaker/model.py

@@ -44,6 +44,7 @@
     ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH,
     load_sagemaker_config,
 )
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.model_card import (
     ModelCard,
     ModelPackageModelCard,
@@ -512,7 +513,7 @@ def register(
         if image_uri is not None:
             self.image_uri = image_uri
 
-        if model_package_group_name is None and model_package_name is None:
+        if model_package_group_name is None and model_package_name is None and self.model_type is not JumpStartModelType.PROPRIETARY:
             # 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(
@@ -542,6 +543,10 @@ 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

@@ -291,3 +291,32 @@ def test_proprietary_jumpstart_model(setup):
     response = predictor.predict(payload)
 
     assert response is not None
+
+@pytest.mark.skipif(
+    True,
+    reason="Only enable if test account is subscribed to the proprietary model",
+)
+def test_register_proprietary_jumpstart_model(setup):
+
+    model_id = "ai21-jurassic-2-light"
+
+    model = JumpStartModel(
+        model_id=model_id,
+        model_version="2.0.004",
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+    )
+
+    pp = model.register()
+
+    predictor = pp.deploy(
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}]
+    )
+    payload = {"prompt": "To be, or", "maxTokens": 4, "temperature": 0, "numResults": 1}
+
+    response = predictor.predict(payload)
+
+    predictor.delete_predictor()
+
+    assert response is not None
+

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

@@ -473,9 +473,11 @@ def test_eula_gated_conditional_s3_prefix_metadata_model(
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
     @mock.patch("sagemaker.jumpstart.model.Model.__init__")
     @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    @mock.patch("sagemaker.jumpstart.model.Model.register")
     @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
     def test_proprietary_model_endpoint(
         self,
+        mock_model_register: mock.Mock,
         mock_model_deploy: mock.Mock,
         mock_model_init: mock.Mock,
         mock_get_model_specs: mock.Mock,
@@ -507,8 +509,14 @@ def test_proprietary_model_endpoint(
             enable_network_isolation=False,
         )
 
+        model.register()
         model.deploy()
 
+        mock_model_register.assert_called_once_with(
+            content_types=["application/json"],
+            response_types=["application/json"],
+        )
+
         mock_model_deploy.assert_called_once_with(
             initial_instance_count=1,
             instance_type="ml.p4de.24xlarge",