Skip to content

Commit 8ead59a

Browse files
evakravibenieric
evakravi
and
benieric
authored
chore: move jumpstart region definitions to json file (aws#5095)
* chore: move jumpstart region definitions to json file * chore: address formatting issues * fix: neo regions not ga in 5 regions * chore: make variable private --------- Co-authored-by: Erick Benitez-Ramos <[email protected]>
1 parent e6b498c commit 8ead59a

File tree

3 files changed

+321
-225
lines changed

3 files changed

+321
-225
lines changed

‎src/sagemaker/jumpstart/constants.py

Lines changed: 50 additions & 222 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import logging
1616
import os
1717
from typing import Dict, Set, Type
18+
import json
1819
import boto3
1920
from sagemaker.base_deserializers import BaseDeserializer, JSONDeserializer
2021
from sagemaker.jumpstart.enums import (
@@ -35,214 +36,58 @@
3536
from sagemaker.session import Session
3637

3738

39+
JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
40+
41+
# disable logging if env var is set
42+
JUMPSTART_LOGGER.addHandler(
43+
type(
44+
"",
45+
(logging.StreamHandler,),
46+
{
47+
"emit": lambda self, *args, **kwargs: (
48+
logging.StreamHandler.emit(self, *args, **kwargs)
49+
if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING)
50+
else None
51+
)
52+
},
53+
)()
54+
)
55+
56+
57+
_CURRENT_FILE_DIRECTORY_PATH = os.path.dirname(os.path.realpath(__file__))
58+
REGION_CONFIG_JSON_FILENAME = "region_config.json"
59+
REGION_CONFIG_JSON_FILEPATH = os.path.join(
60+
_CURRENT_FILE_DIRECTORY_PATH, REGION_CONFIG_JSON_FILENAME
61+
)
62+
63+
64+
def _load_region_config(filepath: str) -> Set[JumpStartLaunchedRegionInfo]:
65+
"""Load the JumpStart region config from a JSON file."""
66+
debug_msg = f"Loading JumpStart region config from '{filepath}'."
67+
JUMPSTART_LOGGER.debug(debug_msg)
68+
try:
69+
with open(filepath) as f:
70+
config = json.load(f)
71+
72+
return {
73+
JumpStartLaunchedRegionInfo(
74+
region_name=region,
75+
content_bucket=data["content_bucket"],
76+
gated_content_bucket=data.get("gated_content_bucket"),
77+
neo_content_bucket=data.get("neo_content_bucket"),
78+
)
79+
for region, data in config.items()
80+
}
81+
except Exception: # pylint: disable=W0703
82+
JUMPSTART_LOGGER.error("Unable to load JumpStart region config.", exc_info=True)
83+
return set()
84+
85+
3886
ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING = "DISABLE_JUMPSTART_LOGGING"
3987
ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY = "DISABLE_JUMPSTART_TELEMETRY"
4088

41-
JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set(
42-
[
43-
JumpStartLaunchedRegionInfo(
44-
region_name="us-west-2",
45-
content_bucket="jumpstart-cache-prod-us-west-2",
46-
gated_content_bucket="jumpstart-private-cache-prod-us-west-2",
47-
neo_content_bucket="sagemaker-sd-models-prod-us-west-2",
48-
),
49-
JumpStartLaunchedRegionInfo(
50-
region_name="us-east-1",
51-
content_bucket="jumpstart-cache-prod-us-east-1",
52-
gated_content_bucket="jumpstart-private-cache-prod-us-east-1",
53-
neo_content_bucket="sagemaker-sd-models-prod-us-east-1",
54-
),
55-
JumpStartLaunchedRegionInfo(
56-
region_name="us-east-2",
57-
content_bucket="jumpstart-cache-prod-us-east-2",
58-
gated_content_bucket="jumpstart-private-cache-prod-us-east-2",
59-
neo_content_bucket="sagemaker-sd-models-prod-us-east-2",
60-
),
61-
JumpStartLaunchedRegionInfo(
62-
region_name="eu-west-1",
63-
content_bucket="jumpstart-cache-prod-eu-west-1",
64-
gated_content_bucket="jumpstart-private-cache-prod-eu-west-1",
65-
neo_content_bucket="sagemaker-sd-models-prod-eu-west-1",
66-
),
67-
JumpStartLaunchedRegionInfo(
68-
region_name="eu-central-1",
69-
content_bucket="jumpstart-cache-prod-eu-central-1",
70-
gated_content_bucket="jumpstart-private-cache-prod-eu-central-1",
71-
neo_content_bucket="sagemaker-sd-models-prod-eu-central-1",
72-
),
73-
JumpStartLaunchedRegionInfo(
74-
region_name="eu-central-2",
75-
content_bucket="jumpstart-cache-prod-eu-central-2",
76-
gated_content_bucket="jumpstart-private-cache-prod-eu-central-2",
77-
),
78-
JumpStartLaunchedRegionInfo(
79-
region_name="eu-north-1",
80-
content_bucket="jumpstart-cache-prod-eu-north-1",
81-
gated_content_bucket="jumpstart-private-cache-prod-eu-north-1",
82-
neo_content_bucket="sagemaker-sd-models-prod-eu-north-1",
83-
),
84-
JumpStartLaunchedRegionInfo(
85-
region_name="eu-south-2",
86-
content_bucket="jumpstart-cache-prod-eu-south-2",
87-
gated_content_bucket="jumpstart-private-cache-prod-eu-south-2",
88-
neo_content_bucket="sagemaker-sd-models-prod-eu-south-2",
89-
),
90-
JumpStartLaunchedRegionInfo(
91-
region_name="me-south-1",
92-
content_bucket="jumpstart-cache-prod-me-south-1",
93-
gated_content_bucket="jumpstart-private-cache-prod-me-south-1",
94-
),
95-
JumpStartLaunchedRegionInfo(
96-
region_name="me-central-1",
97-
content_bucket="jumpstart-cache-prod-me-central-1",
98-
gated_content_bucket="jumpstart-private-cache-prod-me-central-1",
99-
),
100-
JumpStartLaunchedRegionInfo(
101-
region_name="ap-south-1",
102-
content_bucket="jumpstart-cache-prod-ap-south-1",
103-
gated_content_bucket="jumpstart-private-cache-prod-ap-south-1",
104-
neo_content_bucket="sagemaker-sd-models-prod-ap-south-1",
105-
),
106-
JumpStartLaunchedRegionInfo(
107-
region_name="ap-south-2",
108-
content_bucket="jumpstart-cache-prod-ap-south-2",
109-
gated_content_bucket="jumpstart-private-cache-prod-ap-south-2",
110-
neo_content_bucket="sagemaker-sd-models-prod-ap-south-2",
111-
),
112-
JumpStartLaunchedRegionInfo(
113-
region_name="eu-west-3",
114-
content_bucket="jumpstart-cache-prod-eu-west-3",
115-
gated_content_bucket="jumpstart-private-cache-prod-eu-west-3",
116-
neo_content_bucket="sagemaker-sd-models-prod-eu-west-3",
117-
),
118-
JumpStartLaunchedRegionInfo(
119-
region_name="af-south-1",
120-
content_bucket="jumpstart-cache-prod-af-south-1",
121-
gated_content_bucket="jumpstart-private-cache-prod-af-south-1",
122-
),
123-
JumpStartLaunchedRegionInfo(
124-
region_name="sa-east-1",
125-
content_bucket="jumpstart-cache-prod-sa-east-1",
126-
gated_content_bucket="jumpstart-private-cache-prod-sa-east-1",
127-
neo_content_bucket="sagemaker-sd-models-prod-sa-east-1",
128-
),
129-
JumpStartLaunchedRegionInfo(
130-
region_name="ap-east-1",
131-
content_bucket="jumpstart-cache-prod-ap-east-1",
132-
gated_content_bucket="jumpstart-private-cache-prod-ap-east-1",
133-
),
134-
JumpStartLaunchedRegionInfo(
135-
region_name="ap-northeast-2",
136-
content_bucket="jumpstart-cache-prod-ap-northeast-2",
137-
gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-2",
138-
neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-2",
139-
),
140-
JumpStartLaunchedRegionInfo(
141-
region_name="ap-northeast-3",
142-
content_bucket="jumpstart-cache-prod-ap-northeast-3",
143-
gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-3",
144-
neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-3",
145-
),
146-
JumpStartLaunchedRegionInfo(
147-
region_name="ap-southeast-3",
148-
content_bucket="jumpstart-cache-prod-ap-southeast-3",
149-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-3",
150-
neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-3",
151-
),
152-
JumpStartLaunchedRegionInfo(
153-
region_name="ap-southeast-4",
154-
content_bucket="jumpstart-cache-prod-ap-southeast-4",
155-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-4",
156-
neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-4",
157-
),
158-
JumpStartLaunchedRegionInfo(
159-
region_name="ap-southeast-5",
160-
content_bucket="jumpstart-cache-prod-ap-southeast-5",
161-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-5",
162-
),
163-
JumpStartLaunchedRegionInfo(
164-
region_name="ap-southeast-7",
165-
content_bucket="jumpstart-cache-prod-ap-southeast-7",
166-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-7",
167-
),
168-
JumpStartLaunchedRegionInfo(
169-
region_name="eu-west-2",
170-
content_bucket="jumpstart-cache-prod-eu-west-2",
171-
gated_content_bucket="jumpstart-private-cache-prod-eu-west-2",
172-
neo_content_bucket="sagemaker-sd-models-prod-eu-west-2",
173-
),
174-
JumpStartLaunchedRegionInfo(
175-
region_name="eu-south-1",
176-
content_bucket="jumpstart-cache-prod-eu-south-1",
177-
gated_content_bucket="jumpstart-private-cache-prod-eu-south-1",
178-
),
179-
JumpStartLaunchedRegionInfo(
180-
region_name="ap-northeast-1",
181-
content_bucket="jumpstart-cache-prod-ap-northeast-1",
182-
gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-1",
183-
neo_content_bucket="sagemaker-sd-models-prod-ap-northeast-1",
184-
),
185-
JumpStartLaunchedRegionInfo(
186-
region_name="us-west-1",
187-
content_bucket="jumpstart-cache-prod-us-west-1",
188-
gated_content_bucket="jumpstart-private-cache-prod-us-west-1",
189-
neo_content_bucket="sagemaker-sd-models-prod-us-west-1",
190-
),
191-
JumpStartLaunchedRegionInfo(
192-
region_name="ap-southeast-1",
193-
content_bucket="jumpstart-cache-prod-ap-southeast-1",
194-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-1",
195-
neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-1",
196-
),
197-
JumpStartLaunchedRegionInfo(
198-
region_name="ap-southeast-2",
199-
content_bucket="jumpstart-cache-prod-ap-southeast-2",
200-
gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-2",
201-
neo_content_bucket="sagemaker-sd-models-prod-ap-southeast-2",
202-
),
203-
JumpStartLaunchedRegionInfo(
204-
region_name="ca-central-1",
205-
content_bucket="jumpstart-cache-prod-ca-central-1",
206-
gated_content_bucket="jumpstart-private-cache-prod-ca-central-1",
207-
neo_content_bucket="sagemaker-sd-models-prod-ca-central-1",
208-
),
209-
JumpStartLaunchedRegionInfo(
210-
region_name="ca-west-1",
211-
content_bucket="jumpstart-cache-prod-ca-west-1",
212-
gated_content_bucket="jumpstart-private-cache-prod-ca-west-1",
213-
neo_content_bucket="sagemaker-sd-models-prod-ca-west-1",
214-
),
215-
JumpStartLaunchedRegionInfo(
216-
region_name="cn-north-1",
217-
content_bucket="jumpstart-cache-prod-cn-north-1",
218-
gated_content_bucket="jumpstart-private-cache-prod-cn-north-1",
219-
),
220-
JumpStartLaunchedRegionInfo(
221-
region_name="cn-northwest-1",
222-
content_bucket="jumpstart-cache-prod-cn-northwest-1",
223-
gated_content_bucket="jumpstart-private-cache-prod-cn-northwest-1",
224-
),
225-
JumpStartLaunchedRegionInfo(
226-
region_name="il-central-1",
227-
content_bucket="jumpstart-cache-prod-il-central-1",
228-
gated_content_bucket="jumpstart-private-cache-prod-il-central-1",
229-
),
230-
JumpStartLaunchedRegionInfo(
231-
region_name="mx-central-1",
232-
content_bucket="jumpstart-cache-prod-mx-central-1",
233-
gated_content_bucket="jumpstart-private-cache-prod-mx-central-1",
234-
),
235-
JumpStartLaunchedRegionInfo(
236-
region_name="us-gov-east-1",
237-
content_bucket="jumpstart-cache-prod-us-gov-east-1",
238-
gated_content_bucket="jumpstart-private-cache-prod-us-gov-east-1",
239-
),
240-
JumpStartLaunchedRegionInfo(
241-
region_name="us-gov-west-1",
242-
content_bucket="jumpstart-cache-prod-us-gov-west-1",
243-
gated_content_bucket="jumpstart-private-cache-prod-us-gov-west-1",
244-
),
245-
]
89+
JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = _load_region_config(
90+
REGION_CONFIG_JSON_FILEPATH
24691
)
24792

24893
JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = {
@@ -331,23 +176,6 @@
331176

332177
MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
333178

334-
JUMPSTART_LOGGER = logging.getLogger("sagemaker.jumpstart")
335-
336-
# disable logging if env var is set
337-
JUMPSTART_LOGGER.addHandler(
338-
type(
339-
"",
340-
(logging.StreamHandler,),
341-
{
342-
"emit": lambda self, *args, **kwargs: (
343-
logging.StreamHandler.emit(self, *args, **kwargs)
344-
if not os.environ.get(ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING)
345-
else None
346-
)
347-
},
348-
)()
349-
)
350-
351179
try:
352180
DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
353181
boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)

0 commit comments

Comments
 (0)