fix: remove all JUMPSTART_DEFAULT_REGION_NAME default arguments

evakravi · evakravi · commit 7279028b3b8b · 2024-03-12T19:04:09.000Z
diff --git a/src/sagemaker/jumpstart/cache.py b/src/sagemaker/jumpstart/cache.py
@@ -15,7 +15,7 @@
 import datetime
 from difflib import get_close_matches
 import os
-from typing import List, Optional, Set, Tuple, Union
+from typing import List, Optional, Tuple, Union
 import json
 import boto3
 import botocore
@@ -25,9 +25,7 @@
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
     JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
-    JUMPSTART_DEFAULT_REGION_NAME,
     JUMPSTART_LOGGER,
-    JUMPSTART_REGION_NAME_SET,
     MODEL_ID_LIST_WEB_URL,
 )
 from sagemaker.jumpstart.exceptions import get_wildcard_model_version_msg
@@ -94,7 +92,7 @@ def __init__(
             s3_client (Optional[boto3.client]): s3 client to use. Default: None.
         """
 
-        self._region = region or self._get_region_fallback(
+        self._region = region or utils.get_region_fallback(
             s3_bucket_name=s3_bucket_name, s3_client=s3_client
         )
 
@@ -122,35 +120,6 @@ def __init__(
             else boto3.client("s3", region_name=self._region)
         )
 
-    def _get_region_fallback(
-        self, s3_bucket_name: Optional[str], s3_client: Optional[boto3.client]
-    ) -> str:
-        """Returns region to use throughout cache in the absence of one specified in constructor."""
-        regions_in_s3_bucket_name: Set[str] = {
-            region
-            for region in JUMPSTART_REGION_NAME_SET
-            if s3_bucket_name is not None
-            if region in s3_bucket_name
-        }
-        regions_in_s3_client_endpoint_url: Set[str] = {
-            region
-            for region in JUMPSTART_REGION_NAME_SET
-            if s3_client is not None
-            if region in s3_client._endpoint.host
-        }
-
-        combined_regions = regions_in_s3_client_endpoint_url.union(regions_in_s3_bucket_name)
-
-        if len(combined_regions) > 1:
-            raise ValueError(
-                "Unable to resolve a region name from  the s3 bucket and client provided."
-            )
-
-        if len(combined_regions) == 0:
-            return JUMPSTART_DEFAULT_REGION_NAME
-
-        return list(combined_regions)[0]
-
     def set_region(self, region: str) -> None:
         """Set region for cache. Clears cache after new region is set."""
         if region != self._region:
diff --git a/src/sagemaker/jumpstart/notebook_utils.py b/src/sagemaker/jumpstart/notebook_utils.py
@@ -23,7 +23,6 @@
 from sagemaker.jumpstart import accessors
 from sagemaker.jumpstart.constants import (
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-    JUMPSTART_DEFAULT_REGION_NAME,
 )
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.filters import (
@@ -36,6 +35,7 @@
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs
 from sagemaker.jumpstart.utils import (
     get_jumpstart_content_bucket,
+    get_region_fallback,
     get_sagemaker_version,
     verify_model_region_and_return_specs,
 )
@@ -143,7 +143,7 @@ def extract_framework_task_model(model_id: str) -> Tuple[str, str, str]:
 
 def list_jumpstart_tasks(  # pylint: disable=redefined-builtin
     filter: Union[Operator, str] = Constant(BooleanValues.TRUE),
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> List[str]:
     """List tasks for JumpStart, and optionally apply filters to result.
@@ -155,11 +155,14 @@ def list_jumpstart_tasks(  # pylint: disable=redefined-builtin
             (e.g. ``"task == ic"``). If this argument is not supplied, all tasks will be listed.
             (Default: Constant(BooleanValues.TRUE)).
         region (str): Optional. The AWS region from which to retrieve JumpStart metadata regarding
-            models. (Default: JUMPSTART_DEFAULT_REGION_NAME).
+            models. (Default: None).
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session to
             use to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     tasks: Set[str] = set()
     for model_id, _ in _generate_jumpstart_model_versions(
         filter=filter, region=region, sagemaker_session=sagemaker_session
@@ -171,7 +174,7 @@ def list_jumpstart_tasks(  # pylint: disable=redefined-builtin
 
 def list_jumpstart_frameworks(  # pylint: disable=redefined-builtin
     filter: Union[Operator, str] = Constant(BooleanValues.TRUE),
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> List[str]:
     """List frameworks for JumpStart, and optionally apply filters to result.
@@ -183,11 +186,14 @@ def list_jumpstart_frameworks(  # pylint: disable=redefined-builtin
             (eg. ``"task == ic"``). If this argument is not supplied, all frameworks will be listed.
             (Default: Constant(BooleanValues.TRUE)).
         region (str): Optional. The AWS region from which to retrieve JumpStart metadata regarding
-            models. (Default: JUMPSTART_DEFAULT_REGION_NAME).
+            models. (Default: None).
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session
             to use to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     frameworks: Set[str] = set()
     for model_id, _ in _generate_jumpstart_model_versions(
         filter=filter, region=region, sagemaker_session=sagemaker_session
@@ -199,7 +205,7 @@ def list_jumpstart_frameworks(  # pylint: disable=redefined-builtin
 
 def list_jumpstart_scripts(  # pylint: disable=redefined-builtin
     filter: Union[Operator, str] = Constant(BooleanValues.TRUE),
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> List[str]:
     """List scripts for JumpStart, and optionally apply filters to result.
@@ -211,10 +217,13 @@ def list_jumpstart_scripts(  # pylint: disable=redefined-builtin
             (e.g. ``"task == ic"``). If this argument is not supplied, all scripts will be listed.
             (Default: Constant(BooleanValues.TRUE)).
         region (str): Optional. The AWS region from which to retrieve JumpStart metadata regarding
-            models. (Default: JUMPSTART_DEFAULT_REGION_NAME).
+            models. (Default: None).
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session to
             use to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     if (isinstance(filter, Constant) and filter.resolved_value == BooleanValues.TRUE) or (
         isinstance(filter, str) and filter.lower() == BooleanValues.TRUE.lower()
     ):
@@ -242,7 +251,7 @@ def list_jumpstart_scripts(  # pylint: disable=redefined-builtin
 
 def list_jumpstart_models(  # pylint: disable=redefined-builtin
     filter: Union[Operator, str] = Constant(BooleanValues.TRUE),
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     list_incomplete_models: bool = False,
     list_old_models: bool = False,
     list_versions: bool = False,
@@ -257,7 +266,7 @@ def list_jumpstart_models(  # pylint: disable=redefined-builtin
             (e.g. ``"task == ic"``). If this argument is not supplied, all models will be listed.
             (Default: Constant(BooleanValues.TRUE)).
         region (str): Optional. The AWS region from which to retrieve JumpStart metadata regarding
-            models. (Default: JUMPSTART_DEFAULT_REGION_NAME).
+            models. (Default: None).
         list_incomplete_models (bool): Optional. If a model does not contain metadata fields
             requested by the filter, and the filter cannot be resolved to a include/not include,
             whether the model should be included. By default, these models are omitted from results.
@@ -270,6 +279,9 @@ def list_jumpstart_models(  # pylint: disable=redefined-builtin
             to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     model_id_version_dict: Dict[str, List[str]] = dict()
     for model_id, version in _generate_jumpstart_model_versions(
         filter=filter,
@@ -299,7 +311,7 @@ def list_jumpstart_models(  # pylint: disable=redefined-builtin
 
 def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
     filter: Union[Operator, str] = Constant(BooleanValues.TRUE),
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     list_incomplete_models: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> Generator:
@@ -312,7 +324,7 @@ def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
             (e.g. ``"task == ic"``). If this argument is not supplied, all models will be generated.
             (Default: Constant(BooleanValues.TRUE)).
         region (str): Optional. The AWS region from which to retrieve JumpStart metadata regarding
-            models. (Default: JUMPSTART_DEFAULT_REGION_NAME).
+            models. (Default: None).
         list_incomplete_models (bool): Optional. If a model does not contain metadata fields
             requested by the filter, and the filter cannot be resolved to a include/not include,
             whether the model should be included. By default, these models are omitted from
@@ -321,6 +333,10 @@ def _generate_jumpstart_model_versions(  # pylint: disable=redefined-builtin
             to use to perform the model search. (Default: DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
+
     models_manifest_list = accessors.JumpStartModelsAccessor._get_manifest(
         region=region, s3_client=sagemaker_session.s3_client
     )
@@ -453,7 +469,7 @@ def evaluate_model(model_manifest: JumpStartModelHeader) -> Optional[Tuple[str,
 def get_model_url(
     model_id: str,
     model_version: str,
-    region: str = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> str:
     """Retrieve web url describing pretrained model.
@@ -462,11 +478,14 @@ def get_model_url(
         model_id (str): The model ID for which to retrieve the url.
         model_version (str): The model version for which to retrieve the url.
         region (str): Optional. The region from which to retrieve metadata.
-            (Default: JUMPSTART_DEFAULT_REGION_NAME)
+            (Default: None)
         sagemaker_session (sagemaker.session.Session): Optional. The SageMaker Session to use
             to retrieve the model url.
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     model_specs = verify_model_region_and_return_specs(
         region=region,
         model_id=model_id,
diff --git a/src/sagemaker/jumpstart/payload_utils.py b/src/sagemaker/jumpstart/payload_utils.py
@@ -22,12 +22,12 @@
 from sagemaker.jumpstart.artifacts.payloads import _retrieve_example_payloads
 from sagemaker.jumpstart.constants import (
     DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-    JUMPSTART_DEFAULT_REGION_NAME,
 )
 from sagemaker.jumpstart.enums import MIMEType
 from sagemaker.jumpstart.types import JumpStartSerializablePayload
 from sagemaker.jumpstart.utils import (
     get_jumpstart_content_bucket,
+    get_region_fallback,
 )
 from sagemaker.session import Session
 
@@ -125,12 +125,14 @@ class PayloadSerializer:
     def __init__(
         self,
         bucket: Optional[str] = None,
-        region: str = JUMPSTART_DEFAULT_REGION_NAME,
+        region: Optional[str] = None,
         s3_client: Optional[boto3.client] = None,
     ) -> None:
         """Initializes PayloadSerializer object."""
         self.bucket = bucket or get_jumpstart_content_bucket()
-        self.region = region
+        self.region = region or get_region_fallback(
+            s3_client=s3_client,
+        )
         self.s3_client = s3_client
 
     def get_bytes_payload_with_s3_references(
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 import logging
 import os
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 from urllib.parse import urlparse
 import boto3
 from packaging.version import Version
@@ -813,3 +813,42 @@ def get_jumpstart_model_id_version_from_resource_arn(
             model_version = model_version_from_tag
 
     return model_id, model_version
+
+
+def get_region_fallback(
+    s3_bucket_name: Optional[str] = None,
+    s3_client: Optional[boto3.client] = None,
+    sagemaker_session: Optional[Session] = None,
+) -> str:
+    """Returns region to use for JumpStart functionality implicitly via session objects."""
+    regions_in_s3_bucket_name: Set[str] = {
+        region
+        for region in constants.JUMPSTART_REGION_NAME_SET
+        if s3_bucket_name is not None
+        if region in s3_bucket_name
+    }
+    regions_in_s3_client_endpoint_url: Set[str] = {
+        region
+        for region in constants.JUMPSTART_REGION_NAME_SET
+        if s3_client is not None
+        if region in s3_client._endpoint.host
+    }
+
+    regions_in_sagemaker_session: Set[str] = {
+        region
+        for region in constants.JUMPSTART_REGION_NAME_SET
+        if sagemaker_session
+        if region == sagemaker_session.boto_region_name
+    }
+
+    combined_regions = regions_in_s3_client_endpoint_url.union(
+        regions_in_s3_bucket_name, regions_in_sagemaker_session
+    )
+
+    if len(combined_regions) > 1:
+        raise ValueError("Unable to resolve a region name from the s3 bucket and client provided.")
+
+    if len(combined_regions) == 0:
+        return constants.JUMPSTART_DEFAULT_REGION_NAME
+
+    return list(combined_regions)[0]
diff --git a/src/sagemaker/jumpstart/validators.py b/src/sagemaker/jumpstart/validators.py
@@ -14,7 +14,6 @@
 from __future__ import absolute_import
 from typing import Any, Dict, List, Optional
 from sagemaker import session
-from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
 from sagemaker.jumpstart.enums import (
     HyperparameterValidationMode,
@@ -24,7 +23,7 @@
 )
 from sagemaker.jumpstart.exceptions import JumpStartHyperparametersError
 from sagemaker.jumpstart.types import JumpStartHyperparameter
-from sagemaker.jumpstart.utils import verify_model_region_and_return_specs
+from sagemaker.jumpstart.utils import get_region_fallback, verify_model_region_and_return_specs
 
 
 def _validate_hyperparameter(
@@ -168,7 +167,7 @@ def validate_hyperparameters(
     model_version: str,
     hyperparameters: Dict[str, Any],
     validation_mode: HyperparameterValidationMode = HyperparameterValidationMode.VALIDATE_PROVIDED,
-    region: Optional[str] = JUMPSTART_DEFAULT_REGION_NAME,
+    region: Optional[str] = None,
     sagemaker_session: Optional[session.Session] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -184,8 +183,7 @@ def validate_hyperparameters(
           to this function will be validated, the missing hyperparameters will be ignored.
           If set to``VALIDATE_ALGORITHM``, all algorithm hyperparameters will be validated.
           If set to ``VALIDATE_ALL``, all hyperparameters for the model will be validated.
-        region (str): Region for which to validate hyperparameters. (Default: JumpStart
-          default region).
+        region (str): Region for which to validate hyperparameters. (Default: None).
         sagemaker_session (Optional[Session]): Custom SageMaker Session to use.
           (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -202,6 +200,9 @@ def validate_hyperparameters(
 
     """
 
+    region = region or get_region_fallback(
+        sagemaker_session=sagemaker_session,
+    )
     if validation_mode is None:
         validation_mode = HyperparameterValidationMode.VALIDATE_PROVIDED
 
diff --git a/tests/unit/sagemaker/jumpstart/test_cache.py b/tests/unit/sagemaker/jumpstart/test_cache.py
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
