|
15 | 15 | import logging
|
16 | 16 | import os
|
17 | 17 | from typing import Dict, Set, Type
|
| 18 | +import json |
18 | 19 | import boto3
|
19 | 20 | from sagemaker.base_deserializers import BaseDeserializer, JSONDeserializer
|
20 | 21 | from sagemaker.jumpstart.enums import (
|
|
35 | 36 | from sagemaker.session import Session
|
36 | 37 |
|
37 | 38 |
|
| 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 | + |
38 | 86 | ENV_VARIABLE_DISABLE_JUMPSTART_LOGGING = "DISABLE_JUMPSTART_LOGGING"
|
39 | 87 | ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY = "DISABLE_JUMPSTART_TELEMETRY"
|
40 | 88 |
|
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 |
246 | 91 | )
|
247 | 92 |
|
248 | 93 | JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = {
|
|
331 | 176 |
|
332 | 177 | MODEL_ID_LIST_WEB_URL = "https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html"
|
333 | 178 |
|
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 |
| - |
351 | 179 | try:
|
352 | 180 | DEFAULT_JUMPSTART_SAGEMAKER_SESSION = Session(
|
353 | 181 | boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
|
|
0 commit comments