Skip to content

Commit 4477343

Browse files
malte-awsmalte-aws
committed
fix: linting
1 parent 67852da commit 4477343

File tree

1 file changed

+163
-25
lines changed

1 file changed

+163
-25
lines changed

‎tests/unit/sagemaker/image_uris/test_djl.py‎

Lines changed: 163 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -42,17 +42,45 @@ def _test_djl_uris(account, region, version, tag, djl_framework):
4242
)
4343
assert expected == uri
4444

45+
4546
# Expected regions for DJL LMI based on documentation
4647
# https://github.com/aws/deep-learning-containers/blob/master/available_images.md
4748
EXPECTED_DJL_LMI_REGIONS = {
48-
"us-east-1", "us-east-2", "us-west-1", "us-west-2",
49-
"af-south-1", "ap-east-1", "ap-east-2", "ap-south-1", "ap-south-2",
50-
"ap-southeast-1", "ap-southeast-2", "ap-southeast-3", "ap-southeast-4",
51-
"ap-southeast-5", "ap-southeast-7", "ap-northeast-1", "ap-northeast-2",
52-
"ap-northeast-3", "ca-central-1", "ca-west-1", "eu-central-1",
53-
"eu-central-2", "eu-west-1", "eu-west-2", "eu-west-3", "eu-north-1",
54-
"eu-south-1", "eu-south-2", "il-central-1", "mx-central-1",
55-
"me-south-1", "me-central-1", "sa-east-1", "cn-north-1", "cn-northwest-1"
49+
"us-east-1",
50+
"us-east-2",
51+
"us-west-1",
52+
"us-west-2",
53+
"af-south-1",
54+
"ap-east-1",
55+
"ap-east-2",
56+
"ap-south-1",
57+
"ap-south-2",
58+
"ap-southeast-1",
59+
"ap-southeast-2",
60+
"ap-southeast-3",
61+
"ap-southeast-4",
62+
"ap-southeast-5",
63+
"ap-southeast-7",
64+
"ap-northeast-1",
65+
"ap-northeast-2",
66+
"ap-northeast-3",
67+
"ca-central-1",
68+
"ca-west-1",
69+
"eu-central-1",
70+
"eu-central-2",
71+
"eu-west-1",
72+
"eu-west-2",
73+
"eu-west-3",
74+
"eu-north-1",
75+
"eu-south-1",
76+
"eu-south-2",
77+
"il-central-1",
78+
"mx-central-1",
79+
"me-south-1",
80+
"me-central-1",
81+
"sa-east-1",
82+
"cn-north-1",
83+
"cn-northwest-1",
5684
}
5785

5886
# Known missing framework:version:region combinations that don't exist in ECR
@@ -63,14 +91,118 @@ def _test_djl_uris(account, region, version, tag, djl_framework):
6391
"0.28.0-lmi10.0.0-cu124": {"ap-east-2"},
6492
},
6593
"djl-neuronx": {
66-
"0.29.0-neuronx-sdk2.19.1": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
67-
"0.28.0-neuronx-sdk2.18.2": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
68-
"0.27.0-neuronx-sdk2.18.1": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
69-
"0.26.0-neuronx-sdk2.16.0": {"ap-east-1", "me-central-1", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
70-
"0.25.0-neuronx-sdk2.15.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
71-
"0.24.0-neuronx-sdk2.14.1": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
72-
"0.23.0-neuronx-sdk2.12.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
73-
"0.22.1-neuronx-sdk2.10.0": {"eu-north-1", "ap-east-1", "me-central-1", "eu-west-2", "ap-east-2", "ap-southeast-3", "eu-south-1", "ca-central-1", "us-west-1", "ap-northeast-3", "ap-northeast-2", "af-south-1", "me-south-1"},
94+
"0.29.0-neuronx-sdk2.19.1": {
95+
"ap-east-1",
96+
"me-central-1",
97+
"ap-east-2",
98+
"ap-southeast-3",
99+
"eu-south-1",
100+
"ca-central-1",
101+
"us-west-1",
102+
"ap-northeast-3",
103+
"ap-northeast-2",
104+
"af-south-1",
105+
"me-south-1",
106+
},
107+
"0.28.0-neuronx-sdk2.18.2": {
108+
"ap-east-1",
109+
"me-central-1",
110+
"ap-east-2",
111+
"ap-southeast-3",
112+
"eu-south-1",
113+
"ca-central-1",
114+
"us-west-1",
115+
"ap-northeast-3",
116+
"ap-northeast-2",
117+
"af-south-1",
118+
"me-south-1",
119+
},
120+
"0.27.0-neuronx-sdk2.18.1": {
121+
"ap-east-1",
122+
"me-central-1",
123+
"ap-east-2",
124+
"ap-southeast-3",
125+
"eu-south-1",
126+
"ca-central-1",
127+
"us-west-1",
128+
"ap-northeast-3",
129+
"ap-northeast-2",
130+
"af-south-1",
131+
"me-south-1",
132+
},
133+
"0.26.0-neuronx-sdk2.16.0": {
134+
"ap-east-1",
135+
"me-central-1",
136+
"ap-east-2",
137+
"ap-southeast-3",
138+
"eu-south-1",
139+
"ca-central-1",
140+
"us-west-1",
141+
"ap-northeast-3",
142+
"ap-northeast-2",
143+
"af-south-1",
144+
"me-south-1",
145+
},
146+
"0.25.0-neuronx-sdk2.15.0": {
147+
"eu-north-1",
148+
"ap-east-1",
149+
"me-central-1",
150+
"eu-west-2",
151+
"ap-east-2",
152+
"ap-southeast-3",
153+
"eu-south-1",
154+
"ca-central-1",
155+
"us-west-1",
156+
"ap-northeast-3",
157+
"ap-northeast-2",
158+
"af-south-1",
159+
"me-south-1",
160+
},
161+
"0.24.0-neuronx-sdk2.14.1": {
162+
"eu-north-1",
163+
"ap-east-1",
164+
"me-central-1",
165+
"eu-west-2",
166+
"ap-east-2",
167+
"ap-southeast-3",
168+
"eu-south-1",
169+
"ca-central-1",
170+
"us-west-1",
171+
"ap-northeast-3",
172+
"ap-northeast-2",
173+
"af-south-1",
174+
"me-south-1",
175+
},
176+
"0.23.0-neuronx-sdk2.12.0": {
177+
"eu-north-1",
178+
"ap-east-1",
179+
"me-central-1",
180+
"eu-west-2",
181+
"ap-east-2",
182+
"ap-southeast-3",
183+
"eu-south-1",
184+
"ca-central-1",
185+
"us-west-1",
186+
"ap-northeast-3",
187+
"ap-northeast-2",
188+
"af-south-1",
189+
"me-south-1",
190+
},
191+
"0.22.1-neuronx-sdk2.10.0": {
192+
"eu-north-1",
193+
"ap-east-1",
194+
"me-central-1",
195+
"eu-west-2",
196+
"ap-east-2",
197+
"ap-southeast-3",
198+
"eu-south-1",
199+
"ca-central-1",
200+
"us-west-1",
201+
"ap-northeast-3",
202+
"ap-northeast-2",
203+
"af-south-1",
204+
"me-south-1",
205+
},
74206
},
75207
"djl-tensorrtllm": {
76208
"0.30.0-tensorrtllm0.12.0-cu125": {"ap-east-2"},
@@ -96,28 +228,34 @@ def _test_djl_uris(account, region, version, tag, djl_framework):
96228
"0.21.0-deepspeed0.8.3-cu117": {"ap-east-2"},
97229
"0.20.0-deepspeed0.7.5-cu116": {"ap-east-2"},
98230
"0.19.0-deepspeed0.7.3-cu113": {"ap-east-2"},
99-
}
231+
},
100232
}
101233

102234

103-
@pytest.mark.parametrize("framework", ["djl-deepspeed", "djl-fastertransformer", "djl-lmi", "djl-neuronx", "djl-tensorrtllm"])
235+
@pytest.mark.parametrize(
236+
"framework",
237+
["djl-deepspeed", "djl-fastertransformer", "djl-lmi", "djl-neuronx", "djl-tensorrtllm"],
238+
)
104239
def test_djl_lmi_config_for_framework_has_all_regions(framework):
105240
"""Test that config_for_framework returns all expected regions for each version."""
106241
config = image_uris.config_for_framework(framework)
107-
242+
108243
# Check that each version has all expected regions, excluding known missing combinations
109244
for version, version_config in config["versions"].items():
110245
actual_regions = set(version_config["registries"].keys())
111246
expected_regions_for_version = EXPECTED_DJL_LMI_REGIONS.copy()
112-
247+
113248
# Use tag_prefix for lookup if available, otherwise fall back to version
114249
lookup_key = version_config.get("tag_prefix", version)
115-
250+
116251
# Remove regions that are known to be missing for this framework:version combination
117-
missing_regions_for_version = KNOWN_MISSING_COMBINATIONS.get(framework, {}).get(lookup_key, set())
252+
missing_regions_for_version = KNOWN_MISSING_COMBINATIONS.get(framework, {}).get(
253+
lookup_key, set()
254+
)
118255
expected_regions_for_version -= missing_regions_for_version
119-
256+
120257
missing_regions = expected_regions_for_version - actual_regions
121-
122-
assert not missing_regions, f"Framework {framework} version {version} missing regions: {missing_regions}"
123258

259+
assert (
260+
not missing_regions
261+
), f"Framework {framework} version {version} missing regions: {missing_regions}"

0 commit comments

Comments
 (0)