From a7beaa5eafa9bd048f5c37b08b15d63516a53d9b Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Thu, 20 Mar 2025 18:00:01 +0000 Subject: [PATCH 1/4] chore: move jumpstart region definitions to json file --- src/sagemaker/jumpstart/constants.py | 269 ++++--------------- src/sagemaker/jumpstart/region_config.json | 169 ++++++++++++ tests/unit/sagemaker/jumpstart/test_utils.py | 111 +++++++- 3 files changed, 324 insertions(+), 225 deletions(-) create mode 100644 src/sagemaker/jumpstart/region_config.json diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index dd4ded4748..25a60be039 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/src/sagemaker/jumpstart/constants.py @@ -33,216 +33,58 @@ JSONSerializer, ) from sagemaker.session import Session +import json + + +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.""" + JUMPSTART_LOGGER.debug(f"Loading JumpStart region config from '{filepath}'.") + 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: + 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 +173,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..e951f0d59f --- /dev/null +++ b/src/sagemaker/jumpstart/region_config.json @@ -0,0 +1,169 @@ +{ + "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", + "neo_content_bucket": "sagemaker-sd-models-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", + "neo_content_bucket": "sagemaker-sd-models-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", + "neo_content_bucket": "sagemaker-sd-models-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", + "neo_content_bucket": "sagemaker-sd-models-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", + "neo_content_bucket": null + }, + "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", + "neo_content_bucket": "sagemaker-sd-models-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") From 6a48dccb00d987a74391f0a94d84d7e742be71f4 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Mon, 24 Mar 2025 20:38:30 +0000 Subject: [PATCH 2/4] chore: address formatting issues --- src/sagemaker/jumpstart/constants.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index 25a60be039..1f5d42166e 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 ( @@ -33,7 +34,6 @@ JSONSerializer, ) from sagemaker.session import Session -import json JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart") @@ -61,7 +61,8 @@ def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]: """Load the JumpStart region config from a JSON file.""" - JUMPSTART_LOGGER.debug(f"Loading JumpStart region config from '{filepath}'.") + debug_msg = f"Loading JumpStart region config from '{filepath}'." + JUMPSTART_LOGGER.debug(debug_msg) try: with open(filepath) as f: config = json.load(f) @@ -75,7 +76,7 @@ def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]: ) for region, data in config.items() } - except Exception: + except Exception: # pylint: disable=W0703 JUMPSTART_LOGGER.error("Unable to load JumpStart region config.", exc_info=True) return set() From af02b0511ad56d47ac5dcc7389ed18ff48573f21 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Tue, 25 Mar 2025 12:47:49 +0000 Subject: [PATCH 3/4] fix: neo regions not ga in 5 regions --- src/sagemaker/jumpstart/region_config.json | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/sagemaker/jumpstart/region_config.json b/src/sagemaker/jumpstart/region_config.json index e951f0d59f..30bea6ee70 100644 --- a/src/sagemaker/jumpstart/region_config.json +++ b/src/sagemaker/jumpstart/region_config.json @@ -29,8 +29,7 @@ }, "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" + "gated_content_bucket": "jumpstart-private-cache-prod-ap-south-2" }, "ap-southeast-1": { "content_bucket": "jumpstart-cache-prod-ap-southeast-1", @@ -44,13 +43,11 @@ }, "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" + "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", - "neo_content_bucket": "sagemaker-sd-models-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", @@ -67,13 +64,11 @@ }, "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" + "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", - "neo_content_bucket": null + "gated_content_bucket": "jumpstart-private-cache-prod-cn-north-1" }, "cn-northwest-1": { "content_bucket": "jumpstart-cache-prod-cn-northwest-1", @@ -99,8 +94,7 @@ }, "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" + "gated_content_bucket": "jumpstart-private-cache-prod-eu-south-2" }, "eu-west-1": { "content_bucket": "jumpstart-cache-prod-eu-west-1", From f0fdc1c907a58fa7238f3eecb133163a526af654 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Tue, 25 Mar 2025 12:51:46 +0000 Subject: [PATCH 4/4] chore: make variable private --- src/sagemaker/jumpstart/constants.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index 1f5d42166e..b81f97ce3a 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/src/sagemaker/jumpstart/constants.py @@ -54,9 +54,11 @@ ) -CURRENT_FILE_DIRECTORY_PATH = os.path.dirname(os.path.realpath(__file__)) +_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) +REGION_CONFIG_JSON_FILEPATH = os.path.join( + _CURRENT_FILE_DIRECTORY_PATH, REGION_CONFIG_JSON_FILENAME +) def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]: