fix: jumpstart works without region in aws config

Author: evakravi

src/sagemaker/jumpstart/constants.py

@@ -112,7 +112,7 @@
 }
 JUMPSTART_REGION_NAME_SET = {region.region_name for region in JUMPSTART_LAUNCHED_REGIONS}
 
-JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name
+JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 

tests/integ/sagemaker/jumpstart/retrieve_uri/conftest.py

@@ -27,6 +27,8 @@
     JUMPSTART_TAG,
 )
 
+from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
+
 
 def _setup():
     print("Setting up...")
@@ -41,7 +43,9 @@ def _teardown():
     test_suite_id = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
 
     sagemaker_client = boto3.client(
-        "sagemaker", config=Config(retries={"max_attempts": 10, "mode": "standard"})
+        "sagemaker",
+        config=Config(retries={"max_attempts": 10, "mode": "standard"}),
+        region_name=JUMPSTART_DEFAULT_REGION_NAME,
     )
 
     search_endpoints_result = sagemaker_client.search(

tests/integ/sagemaker/jumpstart/retrieve_uri/inference.py

@@ -23,8 +23,8 @@
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
     get_test_artifact_bucket,
+    get_sm_session,
 )
-from sagemaker.session import Session
 
 from sagemaker.utils import repack_model
 from sagemaker.model import (
@@ -59,7 +59,7 @@ def __init__(
         self.region = region
         self.config = boto_config
         self.base_name = base_name
-        self.execution_role = execution_role or Session().get_caller_identity_arn()
+        self.execution_role = execution_role or get_sm_session().get_caller_identity_arn()
         self.account_id = boto3.client("sts").get_caller_identity()["Account"]
         self.image_uri = image_uri
         self.script_uri = script_uri
@@ -102,7 +102,7 @@ def package_artifacts(self):
             dependencies=None,
             model_uri=self.model_uri,
             repacked_model_uri=repacked_model_uri,
-            sagemaker_session=Session(),
+            sagemaker_session=get_sm_session(),
             kms_key=None,
         )
 

tests/integ/sagemaker/jumpstart/retrieve_uri/training.py

@@ -21,15 +21,14 @@
 from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
     get_full_hyperparameters,
     get_test_artifact_bucket,
+    get_sm_session,
 )
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 
 from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
 )
 
-from sagemaker.session import Session
-
 
 class TrainingJobLauncher:
     def __init__(
@@ -53,7 +52,7 @@ def __init__(
         self.region = region
         self.config = boto_config
         self.base_name = base_name
-        self.execution_role = execution_role or Session().get_caller_identity_arn()
+        self.execution_role = execution_role or get_sm_session().get_caller_identity_arn()
         self.image_uri = image_uri
         self.script_uri = script_uri
         self.model_uri = model_uri

tests/integ/sagemaker/jumpstart/retrieve_uri/utils.py

@@ -67,8 +67,12 @@ def get_training_dataset_for_model_and_version(model_id: str, version: str) -> d
     return TRAINING_DATASET_MODEL_DICT[(model_id, version)]
 
 
+def get_sm_session() -> Session:
+    return Session(boto_session=boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME))
+
+
 def get_test_artifact_bucket() -> str:
-    bucket_name = Session().default_bucket()
+    bucket_name = get_sm_session().default_bucket()
     return bucket_name