diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index dd4ded4748..b81f97ce3a 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/src/sagemaker/jumpstart/constants.py @@ -15,6 +15,7 @@ import logging import os from typing import Dict, Set, Type +import json import boto3 from sagemaker.base_deserializers import BaseDeserializer, JSONDeserializer from sagemaker.jumpstart.enums import ( @@ -35,214 +36,58 @@ from sagemaker.session import Session +JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart") + +# disable logging if env var is set +JUMPSTART_LOGGER.addHandler( + type( + "", + (logging.StreamHandler,), + { + "emit": lambda self, *args, **kwargs: ( + logging.StreamHandler.emit(self, *args, **kwargs) + if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING) + else None + ) + }, + )() +) + + +_CURRENT_FILE_DIRECTORY_PATH = os.path.dirname(os.path.realpath(__file__)) +REGION_CONFIG_JSON_FILENAME = "region_config.json" +REGION_CONFIG_JSON_FILEPATH = os.path.join( + _CURRENT_FILE_DIRECTORY_PATH, REGION_CONFIG_JSON_FILENAME +) + + +def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]: + """Load the JumpStart region config from a JSON file.""" + debug_msg = f"Loading JumpStart region config from '{filepath}'." + JUMPSTART_LOGGER.debug(debug_msg) + try: + with open(filepath) as f: + config = json.load(f) + + return { + JumpStartLaunchedRegionInfo( + region_name=region, + content_bucket=data["content_bucket"], + gated_content_bucket=data.get("gated_content_bucket"), + neo_content_bucket=data.get("neo_content_bucket"), + ) + for region, data in config.items() + } + except Exception: # pylint: disable=W0703 + JUMPSTART_LOGGER.error("Unable to load JumpStart region config.", exc_info=True) + return set() + + ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING = "DISABLE_JUMPSTART_LOGGING" ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY = "DISABLE_JUMPSTART_TELEMETRY" -JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set( - [ - JumpStartLaunchedRegionInfo( - region_name="us-west-2", - content_bucket="jumpstart-cache-prod-us-west-2", - gated_content_bucket="jumpstart-private-cache-prod-us-west-2", - neo_content_bucket="sagemaker-sd-models-prod-us-west-2", - ), - JumpStartLaunchedRegionInfo( - region_name="us-east-1", - content_bucket="jumpstart-cache-prod-us-east-1", - gated_content_bucket="jumpstart-private-cache-prod-us-east-1", - neo_content_bucket="sagemaker-sd-models-prod-us-east-1", - ), - JumpStartLaunchedRegionInfo( - region_name="us-east-2", - content_bucket="jumpstart-cache-prod-us-east-2", - gated_content_bucket="jumpstart-private-cache-prod-us-east-2", - neo_content_bucket="sagemaker-sd-models-prod-us-east-2", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-west-1", - content_bucket="jumpstart-cache-prod-eu-west-1", - gated_content_bucket="jumpstart-private-cache-prod-eu-west-1", - neo_content_bucket="sagemaker-sd-models-prod-eu-west-1", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-central-1", - content_bucket="jumpstart-cache-prod-eu-central-1", - gated_content_bucket="jumpstart-private-cache-prod-eu-central-1", - neo_content_bucket="sagemaker-sd-models-prod-eu-central-1", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-central-2", - content_bucket="jumpstart-cache-prod-eu-central-2", - gated_content_bucket="jumpstart-private-cache-prod-eu-central-2", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-north-1", - content_bucket="jumpstart-cache-prod-eu-north-1", - gated_content_bucket="jumpstart-private-cache-prod-eu-north-1", - neo_content_bucket="sagemaker-sd-models-prod-eu-north-1", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-south-2", - content_bucket="jumpstart-cache-prod-eu-south-2", - gated_content_bucket="jumpstart-private-cache-prod-eu-south-2", - neo_content_bucket="sagemaker-sd-models-prod-eu-south-2", - ), - JumpStartLaunchedRegionInfo( - region_name="me-south-1", - content_bucket="jumpstart-cache-prod-me-south-1", - gated_content_bucket="jumpstart-private-cache-prod-me-south-1", - ), - JumpStartLaunchedRegionInfo( - region_name="me-central-1", - content_bucket="jumpstart-cache-prod-me-central-1", - gated_content_bucket="jumpstart-private-cache-prod-me-central-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-south-1", - content_bucket="jumpstart-cache-prod-ap-south-1", - gated_content_bucket="jumpstart-private-cache-prod-ap-south-1", - neo_content_bucket="sagemaker-sd-models-prod-ap-south-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-south-2", - content_bucket="jumpstart-cache-prod-ap-south-2", - gated_content_bucket="jumpstart-private-cache-prod-ap-south-2", - neo_content_bucket="sagemaker-sd-models-prod-ap-south-2", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-west-3", - content_bucket="jumpstart-cache-prod-eu-west-3", - gated_content_bucket="jumpstart-private-cache-prod-eu-west-3", - neo_content_bucket="sagemaker-sd-models-prod-eu-west-3", - ), - JumpStartLaunchedRegionInfo( - region_name="af-south-1", - content_bucket="jumpstart-cache-prod-af-south-1", - gated_content_bucket="jumpstart-private-cache-prod-af-south-1", - ), - JumpStartLaunchedRegionInfo( - region_name="sa-east-1", - content_bucket="jumpstart-cache-prod-sa-east-1", - gated_content_bucket="jumpstart-private-cache-prod-sa-east-1", - neo_content_bucket="sagemaker-sd-models-prod-sa-east-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-east-1", - content_bucket="jumpstart-cache-prod-ap-east-1", - gated_content_bucket="jumpstart-private-cache-prod-ap-east-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-northeast-2", - content_bucket="jumpstart-cache-prod-ap-northeast-2", - gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-2", - neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-2", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-northeast-3", - content_bucket="jumpstart-cache-prod-ap-northeast-3", - gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-3", - neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-3", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-3", - content_bucket="jumpstart-cache-prod-ap-southeast-3", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-3", - neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-3", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-4", - content_bucket="jumpstart-cache-prod-ap-southeast-4", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-4", - neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-4", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-5", - content_bucket="jumpstart-cache-prod-ap-southeast-5", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-5", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-7", - content_bucket="jumpstart-cache-prod-ap-southeast-7", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-7", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-west-2", - content_bucket="jumpstart-cache-prod-eu-west-2", - gated_content_bucket="jumpstart-private-cache-prod-eu-west-2", - neo_content_bucket="sagemaker-sd-models-prod-eu-west-2", - ), - JumpStartLaunchedRegionInfo( - region_name="eu-south-1", - content_bucket="jumpstart-cache-prod-eu-south-1", - gated_content_bucket="jumpstart-private-cache-prod-eu-south-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-northeast-1", - content_bucket="jumpstart-cache-prod-ap-northeast-1", - gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-1", - neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-1", - ), - JumpStartLaunchedRegionInfo( - region_name="us-west-1", - content_bucket="jumpstart-cache-prod-us-west-1", - gated_content_bucket="jumpstart-private-cache-prod-us-west-1", - neo_content_bucket="sagemaker-sd-models-prod-us-west-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-1", - content_bucket="jumpstart-cache-prod-ap-southeast-1", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-1", - neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ap-southeast-2", - content_bucket="jumpstart-cache-prod-ap-southeast-2", - gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-2", - neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-2", - ), - JumpStartLaunchedRegionInfo( - region_name="ca-central-1", - content_bucket="jumpstart-cache-prod-ca-central-1", - gated_content_bucket="jumpstart-private-cache-prod-ca-central-1", - neo_content_bucket="sagemaker-sd-models-prod-ca-central-1", - ), - JumpStartLaunchedRegionInfo( - region_name="ca-west-1", - content_bucket="jumpstart-cache-prod-ca-west-1", - gated_content_bucket="jumpstart-private-cache-prod-ca-west-1", - neo_content_bucket="sagemaker-sd-models-prod-ca-west-1", - ), - JumpStartLaunchedRegionInfo( - region_name="cn-north-1", - content_bucket="jumpstart-cache-prod-cn-north-1", - gated_content_bucket="jumpstart-private-cache-prod-cn-north-1", - ), - JumpStartLaunchedRegionInfo( - region_name="cn-northwest-1", - content_bucket="jumpstart-cache-prod-cn-northwest-1", - gated_content_bucket="jumpstart-private-cache-prod-cn-northwest-1", - ), - JumpStartLaunchedRegionInfo( - region_name="il-central-1", - content_bucket="jumpstart-cache-prod-il-central-1", - gated_content_bucket="jumpstart-private-cache-prod-il-central-1", - ), - JumpStartLaunchedRegionInfo( - region_name="mx-central-1", - content_bucket="jumpstart-cache-prod-mx-central-1", - gated_content_bucket="jumpstart-private-cache-prod-mx-central-1", - ), - JumpStartLaunchedRegionInfo( - region_name="us-gov-east-1", - content_bucket="jumpstart-cache-prod-us-gov-east-1", - gated_content_bucket="jumpstart-private-cache-prod-us-gov-east-1", - ), - JumpStartLaunchedRegionInfo( - region_name="us-gov-west-1", - content_bucket="jumpstart-cache-prod-us-gov-west-1", - gated_content_bucket="jumpstart-private-cache-prod-us-gov-west-1", - ), - ] +JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = _load_region_config( + REGION_CONFIG_JSON_FILEPATH ) JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = { @@ -331,23 +176,6 @@ MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html" -JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart") - -# disable logging if env var is set -JUMPSTART_LOGGER.addHandler( - type( - "", - (logging.StreamHandler,), - { - "emit": lambda self, *args, **kwargs: ( - logging.StreamHandler.emit(self, *args, **kwargs) - if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING) - else None - ) - }, - )() -) - try: DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session( boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME) diff --git a/src/sagemaker/jumpstart/region_config.json b/src/sagemaker/jumpstart/region_config.json new file mode 100644 index 0000000000..30bea6ee70 --- /dev/null +++ b/src/sagemaker/jumpstart/region_config.json @@ -0,0 +1,163 @@ +{ + "af-south-1": { + "content_bucket": "jumpstart-cache-prod-af-south-1", + "gated_content_bucket": "jumpstart-private-cache-prod-af-south-1" + }, + "ap-east-1": { + "content_bucket": "jumpstart-cache-prod-ap-east-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-east-1" + }, + "ap-northeast-1": { + "content_bucket": "jumpstart-cache-prod-ap-northeast-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-1", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-1" + }, + "ap-northeast-2": { + "content_bucket": "jumpstart-cache-prod-ap-northeast-2", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-2", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-2" + }, + "ap-northeast-3": { + "content_bucket": "jumpstart-cache-prod-ap-northeast-3", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-3", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-3" + }, + "ap-south-1": { + "content_bucket": "jumpstart-cache-prod-ap-south-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-south-1", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-south-1" + }, + "ap-south-2": { + "content_bucket": "jumpstart-cache-prod-ap-south-2", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-south-2" + }, + "ap-southeast-1": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-1", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-southeast-1" + }, + "ap-southeast-2": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-2", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-2", + "neo_content_bucket": "sagemaker-sd-models-prod-ap-southeast-2" + }, + "ap-southeast-3": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-3", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-3" + }, + "ap-southeast-4": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-4", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-4" + }, + "ap-southeast-5": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-5", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-5" + }, + "ap-southeast-7": { + "content_bucket": "jumpstart-cache-prod-ap-southeast-7", + "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-7" + }, + "ca-central-1": { + "content_bucket": "jumpstart-cache-prod-ca-central-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ca-central-1", + "neo_content_bucket": "sagemaker-sd-models-prod-ca-central-1" + }, + "ca-west-1": { + "content_bucket": "jumpstart-cache-prod-ca-west-1", + "gated_content_bucket": "jumpstart-private-cache-prod-ca-west-1" + }, + "cn-north-1": { + "content_bucket": "jumpstart-cache-prod-cn-north-1", + "gated_content_bucket": "jumpstart-private-cache-prod-cn-north-1" + }, + "cn-northwest-1": { + "content_bucket": "jumpstart-cache-prod-cn-northwest-1", + "gated_content_bucket": "jumpstart-private-cache-prod-cn-northwest-1" + }, + "eu-central-1": { + "content_bucket": "jumpstart-cache-prod-eu-central-1", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-central-1", + "neo_content_bucket": "sagemaker-sd-models-prod-eu-central-1" + }, + "eu-central-2": { + "content_bucket": "jumpstart-cache-prod-eu-central-2", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-central-2" + }, + "eu-north-1": { + "content_bucket": "jumpstart-cache-prod-eu-north-1", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-north-1", + "neo_content_bucket": "sagemaker-sd-models-prod-eu-north-1" + }, + "eu-south-1": { + "content_bucket": "jumpstart-cache-prod-eu-south-1", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-south-1" + }, + "eu-south-2": { + "content_bucket": "jumpstart-cache-prod-eu-south-2", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-south-2" + }, + "eu-west-1": { + "content_bucket": "jumpstart-cache-prod-eu-west-1", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-1", + "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-1" + }, + "eu-west-2": { + "content_bucket": "jumpstart-cache-prod-eu-west-2", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-2", + "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-2" + }, + "eu-west-3": { + "content_bucket": "jumpstart-cache-prod-eu-west-3", + "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-3", + "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-3" + }, + "il-central-1": { + "content_bucket": "jumpstart-cache-prod-il-central-1", + "gated_content_bucket": "jumpstart-private-cache-prod-il-central-1" + }, + "me-central-1": { + "content_bucket": "jumpstart-cache-prod-me-central-1", + "gated_content_bucket": "jumpstart-private-cache-prod-me-central-1" + }, + "me-south-1": { + "content_bucket": "jumpstart-cache-prod-me-south-1", + "gated_content_bucket": "jumpstart-private-cache-prod-me-south-1" + }, + "mx-central-1": { + "content_bucket": "jumpstart-cache-prod-mx-central-1", + "gated_content_bucket": "jumpstart-private-cache-prod-mx-central-1" + }, + "sa-east-1": { + "content_bucket": "jumpstart-cache-prod-sa-east-1", + "gated_content_bucket": "jumpstart-private-cache-prod-sa-east-1", + "neo_content_bucket": "sagemaker-sd-models-prod-sa-east-1" + }, + "us-east-1": { + "content_bucket": "jumpstart-cache-prod-us-east-1", + "gated_content_bucket": "jumpstart-private-cache-prod-us-east-1", + "neo_content_bucket": "sagemaker-sd-models-prod-us-east-1" + }, + "us-east-2": { + "content_bucket": "jumpstart-cache-prod-us-east-2", + "gated_content_bucket": "jumpstart-private-cache-prod-us-east-2", + "neo_content_bucket": "sagemaker-sd-models-prod-us-east-2" + }, + "us-gov-east-1": { + "content_bucket": "jumpstart-cache-prod-us-gov-east-1", + "gated_content_bucket": "jumpstart-private-cache-prod-us-gov-east-1" + }, + "us-gov-west-1": { + "content_bucket": "jumpstart-cache-prod-us-gov-west-1", + "gated_content_bucket": "jumpstart-private-cache-prod-us-gov-west-1" + }, + "us-west-1": { + "content_bucket": "jumpstart-cache-prod-us-west-1", + "gated_content_bucket": "jumpstart-private-cache-prod-us-west-1", + "neo_content_bucket": "sagemaker-sd-models-prod-us-west-1" + }, + "us-west-2": { + "content_bucket": "jumpstart-cache-prod-us-west-2", + "gated_content_bucket": "jumpstart-private-cache-prod-us-west-2", + "neo_content_bucket": "sagemaker-sd-models-prod-us-west-2" + } +} \ No newline at end of file diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py index ea4d64f289..e3e3110da8 100644 --- a/tests/unit/sagemaker/jumpstart/test_utils.py +++ b/tests/unit/sagemaker/jumpstart/test_utils.py @@ -13,10 +13,9 @@ from __future__ import absolute_import import os from unittest import TestCase -from unittest.mock import call - +from unittest.mock import call, mock_open, Mock, patch +import json from botocore.exceptions import ClientError -from mock.mock import Mock, patch import pytest import boto3 import random @@ -24,6 +23,7 @@ from sagemaker import session from sagemaker.jumpstart import utils from sagemaker.jumpstart.constants import ( + _load_region_config, DEFAULT_JUMPSTART_SAGEMAKER_SESSION, ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING, ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE, @@ -38,6 +38,7 @@ JUMPSTART_RESOURCE_BASE_NAME, NEO_DEFAULT_REGION_NAME, JumpStartScriptScope, + JUMPSTART_LAUNCHED_REGIONS, ) from functools import partial from sagemaker.jumpstart.enums import JumpStartTag, MIMEType, JumpStartModelType @@ -49,6 +50,7 @@ JumpStartBenchmarkStat, JumpStartModelHeader, JumpStartVersionedModelId, + JumpStartLaunchedRegionInfo, ) from tests.unit.sagemaker.jumpstart.utils import ( get_base_spec_with_prototype_configs, @@ -1569,6 +1571,109 @@ def test_multiple_config_names_found_aliases_inconsistent(self): mock_list_tags.assert_called_once_with("some-arn") +class TestJumpStartLaunchedRegions(TestCase): + def test_regions_not_empty(self): + self.assertTrue(len(JUMPSTART_LAUNCHED_REGIONS) > 0) + + +class TestLoadRegionConfig(TestCase): + def setUp(self): + # Sample valid config that matches the expected structure + self.valid_config = { + "us-east-1": { + "content_bucket": "jumpstart-cache-prod-us-east-1", + "gated_content_bucket": "jumpstart-private-cache-prod-us-east-1", + "neo_content_bucket": "jumpstart-neo-cache-prod-us-east-1", + }, + "us-west-2": { + "content_bucket": "jumpstart-cache-prod-us-west-2", + }, + } + self.config_json = json.dumps(self.valid_config) + + @patch("builtins.open", new_callable=mock_open) + def test_successful_config_load(self, mock_file): + # Setup mock to return valid config + mock_file.return_value.__enter__().read.return_value = self.config_json + + result = _load_region_config("dummy/path") + + # Verify the returned dictionary contains JumpStartLaunchedRegionInfo objects + self.assertTrue(all(isinstance(region, JumpStartLaunchedRegionInfo) for region in result)) + + for region in result: + if region.region_name == "us-east-1": + self.assertEqual(region.region_name, "us-east-1") + self.assertEqual(region.content_bucket, "jumpstart-cache-prod-us-east-1") + self.assertEqual( + region.gated_content_bucket, "jumpstart-private-cache-prod-us-east-1" + ) + self.assertEqual(region.neo_content_bucket, "jumpstart-neo-cache-prod-us-east-1") + + elif region.region_name == "us-west-2": + self.assertEqual(region.region_name, "us-west-2") + self.assertEqual(region.content_bucket, "jumpstart-cache-prod-us-west-2") + self.assertIsNone(region.gated_content_bucket) + self.assertIsNone(region.neo_content_bucket) + else: + raise AssertionError(f"Unexpected region name found: {region.region_name}") + + @patch("builtins.open", new_callable=mock_open) + def test_missing_required_field(self, mock_file): + # Config missing required content_bucket field + invalid_config = { + "us-east-1": { + "gated_content_bucket": "XXXXXXXXXXX", + "neo_content_bucket": "some-other-bucket", + } + } + mock_file.return_value.__enter__().read.return_value = json.dumps(invalid_config) + + # Should return empty dict due to exception handling + result = _load_region_config("dummy/path") + self.assertEqual(result, set()) + + @patch("builtins.open") + def test_file_not_found(self, mock_file): + # Simulate file not found + mock_file.side_effect = FileNotFoundError() + + # Should return empty dict due to exception handling + result = _load_region_config("dummy/path") + self.assertEqual(result, set()) + + @patch("builtins.open", new_callable=mock_open) + def test_invalid_json(self, mock_file): + # Setup mock to return invalid JSON + mock_file.return_value.__enter__().read.return_value = "invalid json content" + + # Should return empty dict due to exception handling + result = _load_region_config("dummy/path") + self.assertEqual(result, set()) + + @patch("builtins.open", new_callable=mock_open) + def test_empty_config(self, mock_file): + # Setup mock to return empty JSON object + mock_file.return_value.__enter__().read.return_value = "{}" + + result = _load_region_config("dummy/path") + self.assertEqual(result, set()) + + @patch("sagemaker.jumpstart.constants.JUMPSTART_LOGGER") + @patch("builtins.open") + def test_logging_on_error(self, mock_file, mock_logger): + + # Simulate an error + mock_file.side_effect = Exception("Test error") + + result = _load_region_config("dummy/path") + + self.assertEqual(result, set()) + + # Verify error was logged + mock_logger.error.assert_called_once() + + class TestJumpStartLogger(TestCase): @patch.dict("os.environ", {}) @patch("logging.StreamHandler.emit")