Feature: register proprietary models from jumpstart

Author: selvask-aws

src/sagemaker/jumpstart/factory/model.py

@@ -694,6 +694,7 @@ def get_register_kwargs(
     model_version: Optional[str] = None,
     model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_WEIGHTS,
     hub_arn: 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,
@@ -729,6 +730,7 @@ def get_register_kwargs(
         model_version=model_version,
         model_type=model_type,
         hub_arn=hub_arn,
+        model_type=model_type,
         region=region,
         tolerate_deprecated_model=tolerate_deprecated_model,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -763,6 +765,7 @@ def get_register_kwargs(
         version=model_version,
         model_type=model_type,
         hub_arn=hub_arn,
+        model_type=model_type,
         region=region,
         scope=JumpStartScriptScope.INFERENCE,
         sagemaker_session=sagemaker_session,

src/sagemaker/jumpstart/model.py

@@ -819,6 +819,7 @@ def register(
             model_version=self.model_version,
             model_type=self.model_type,
             hub_arn=self.hub_arn,
+            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

@@ -2452,6 +2452,7 @@ def __init__(
         self.model_version = model_version
         self.model_type = model_type
         self.hub_arn = hub_arn
+        self.model_type = model_type
         self.region = region
         self.image_uri = image_uri
         self.sagemaker_session = sagemaker_session