pass hub_arn into all estimator utils/artifacts

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,

src/sagemaker/image_uris.py

@@ -61,6 +61,7 @@ def retrieve(
     training_compiler_config=None,
     model_id=None,
     model_version=None,
+    hub_arn=None,
     tolerate_vulnerable_model=False,
     tolerate_deprecated_model=False,
     sdk_version=None,
@@ -101,6 +102,8 @@ def retrieve(
             (default: None).
         model_version (str): The version of the JumpStart model for which to retrieve the
             image URI (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 without an exception raised. If ``False``, raises an exception if
             the script used by this version of the model has dependencies with known security
@@ -146,6 +149,7 @@ def retrieve(
             model_id,
             model_version,
             image_scope,
+            hub_arn,
             framework,
             region,
             version,

src/sagemaker/instance_types.py

@@ -45,6 +45,8 @@ def retrieve_default(
             retrieve the default instance type. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default instance type. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         scope (str): The model type, i.e. what it is used for.
             Valid values: "training" and "inference".
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -110,6 +112,8 @@ def retrieve(
             retrieve the supported instance types. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             supported instance types. (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

src/sagemaker/jumpstart/artifacts/environment_variables.py

@@ -31,6 +31,7 @@
 def _retrieve_default_environment_variables(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -46,6 +47,8 @@ def _retrieve_default_environment_variables(
             retrieve the default environment variables.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default environment variables.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (Optional[str]): Region for which to retrieve default environment variables.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -77,6 +80,7 @@ def _retrieve_default_environment_variables(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=script,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -113,6 +117,7 @@ def _retrieve_default_environment_variables(
             gated_model_env_var: Optional[str] = _retrieve_gated_model_uri_env_var_value(
                 model_id=model_id,
                 model_version=model_version,
+                hub_arn=hub_arn,
                 region=region,
                 tolerate_vulnerable_model=tolerate_vulnerable_model,
                 tolerate_deprecated_model=tolerate_deprecated_model,
@@ -131,6 +136,7 @@ def _retrieve_default_environment_variables(
 def _retrieve_gated_model_uri_env_var_value(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -144,6 +150,8 @@ def _retrieve_gated_model_uri_env_var_value(
             retrieve the gated model env var URI.
         model_version (str): Version of the JumpStart model for which to retrieve the
             gated model env var URI.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (Optional[str]): Region for which to retrieve the gated model env var URI.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -174,6 +182,7 @@ def _retrieve_gated_model_uri_env_var_value(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/hyperparameters.py

@@ -30,6 +30,7 @@
 def _retrieve_default_hyperparameters(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     include_container_hyperparameters: bool = False,
     tolerate_vulnerable_model: bool = False,
@@ -44,6 +45,8 @@ def _retrieve_default_hyperparameters(
             retrieve the default hyperparameters.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default hyperparameters.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (str): Region for which to retrieve default hyperparameters.
             (Default: None).
         include_container_hyperparameters (bool): True if container hyperparameters
@@ -76,6 +79,7 @@ def _retrieve_default_hyperparameters(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -33,6 +33,7 @@ def _retrieve_image_uri(
     model_id: str,
     model_version: str,
     image_scope: str,
+    hub_arn: Optional[str] = None,
     framework: Optional[str] = None,
     region: Optional[str] = None,
     version: Optional[str] = None,
@@ -57,6 +58,8 @@ def _retrieve_image_uri(
         model_id (str): JumpStart model ID for which to retrieve image URI.
         model_version (str): Version of the JumpStart model for which to retrieve
             the image URI.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         image_scope (str): The image type, i.e. what it is used for.
             Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
             ``image_scope`` is ignored.
@@ -110,6 +113,7 @@ def _retrieve_image_uri(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=image_scope,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/incremental_training.py

@@ -30,6 +30,7 @@ def _model_supports_incremental_training(
     model_id: str,
     model_version: str,
     region: Optional[str],
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -43,6 +44,8 @@ def _model_supports_incremental_training(
             support status for incremental training.
         region (Optional[str]): Region for which to retrieve the
             support status for incremental training.
+        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
@@ -64,6 +67,7 @@ def _model_supports_incremental_training(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/instance_types.py

@@ -49,6 +49,8 @@ def _retrieve_default_instance_type(
             default instance type.
         scope (str): The script type, i.e. what it is used for.
             Valid values: "training" and "inference".
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (Optional[str]): Region for which to retrieve default instance type.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -137,6 +139,8 @@ def _retrieve_instance_types(
             supported instance types.
         scope (str): The script type, i.e. what it is used for.
             Valid values: "training" and "inference".
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (Optional[str]): Region for which to retrieve supported instance types.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -140,6 +140,7 @@ def _retrieve_estimator_init_kwargs(
     model_id: str,
     model_version: str,
     instance_type: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -154,6 +155,8 @@ def _retrieve_estimator_init_kwargs(
             kwargs.
         instance_type (str): Instance type of the training job, to determine if volume size is
             supported.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         region (Optional[str]): Region for which to retrieve kwargs.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -177,6 +180,7 @@ def _retrieve_estimator_init_kwargs(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/metric_definitions.py

@@ -31,6 +31,7 @@ def _retrieve_default_training_metric_definitions(
     model_id: str,
     model_version: str,
     region: Optional[str],
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -45,6 +46,8 @@ def _retrieve_default_training_metric_definitions(
             default training metric definitions.
         region (Optional[str]): Region for which to retrieve default training metric
             definitions.
+        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
@@ -68,6 +71,7 @@ def _retrieve_default_training_metric_definitions(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -110,6 +110,7 @@ def _retrieve_model_package_model_artifact_s3_uri(
     model_id: str,
     model_version: str,
     region: Optional[str],
+    hub_arn: Optional[str] = None,
     scope: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -124,6 +125,8 @@ def _retrieve_model_package_model_artifact_s3_uri(
             model package artifact.
         region (Optional[str]): Region for which to retrieve the model package artifact.
             (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from (default: None).
         scope (Optional[str]): Scope for which to retrieve the model package artifact.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -152,6 +155,7 @@ def _retrieve_model_package_model_artifact_s3_uri(
         model_specs = verify_model_region_and_return_specs(
             model_id=model_id,
             version=model_version,
+            hub_arn=hub_arn,
             scope=scope,
             region=region,
             tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -178,6 +178,7 @@ def _model_supports_training_model_uri(
     model_id: str,
     model_version: str,
     region: Optional[str],
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -191,6 +192,8 @@ def _model_supports_training_model_uri(
             support status for model uri with training.
         region (Optional[str]): Region for which to retrieve the
             support status for model uri with training.
+        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
@@ -212,6 +215,7 @@ def _model_supports_training_model_uri(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/payloads.py

@@ -32,6 +32,7 @@ def _retrieve_example_payloads(
     model_id: str,
     model_version: str,
     region: Optional[str],
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -45,6 +46,8 @@ def _retrieve_example_payloads(
             example payloads.
         region (Optional[str]): Region for which to retrieve the
             example payloads.
+        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
@@ -67,6 +70,7 @@ def _retrieve_example_payloads(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,