feat: jsch jumpstart estimator support (aws#4439)

Author: bencrabtree

src/sagemaker/environment_variables.py

@@ -30,6 +30,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     include_aws_sdk_env_vars: bool = True,
@@ -46,6 +47,8 @@ def retrieve_default(
             retrieve the default environment variables. (Default: None).
         model_version (str): Optional. The version of the model for which to retrieve the
             default environment variables. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -80,6 +83,7 @@ def retrieve_default(
     return artifacts._retrieve_default_environment_variables(
         model_id,
         model_version,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

src/sagemaker/hyperparameters.py

@@ -31,6 +31,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     instance_type: Optional[str] = None,
     include_container_hyperparameters: bool = False,
     tolerate_vulnerable_model: bool = False,
@@ -46,6 +47,8 @@ def retrieve_default(
             retrieve the default hyperparameters. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default hyperparameters. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         instance_type (str): An instance type to optionally supply in order to get hyperparameters
             specific for the instance type.
         include_container_hyperparameters (bool): ``True`` if the container hyperparameters
@@ -80,6 +83,7 @@ def retrieve_default(
     return artifacts._retrieve_default_hyperparameters(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         instance_type=instance_type,
         region=region,
         include_container_hyperparameters=include_container_hyperparameters,
@@ -93,6 +97,7 @@ def validate(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     hyperparameters: Optional[dict] = None,
     validation_mode: HyperparameterValidationMode = HyperparameterValidationMode.VALIDATE_PROVIDED,
     tolerate_vulnerable_model: bool = False,
@@ -107,6 +112,8 @@ def validate(
             (Default: None).
         model_version (str): The version of the model for which to validate hyperparameters.
             (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         hyperparameters (dict): Hyperparameters to validate.
             (Default: None).
         validation_mode (HyperparameterValidationMode): Method of validation to use with
@@ -148,6 +155,7 @@ def validate(
     return validate_hyperparameters(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         hyperparameters=hyperparameters,
         validation_mode=validation_mode,
         region=region,