Skip to content

Commit 07fa0a9

Browse files
arjkesharjkesh
committed
refactor for ease of adding to support policy
1 parent 61fc611 commit 07fa0a9

File tree

3 files changed

+96
-63
lines changed

3 files changed

+96
-63
lines changed

‎src/sagemaker/image_uri_config/pytorch.json‎

Lines changed: 35 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,7 @@
3434
"us-east-2": "763104351884",
3535
"us-west-2": "763104351884"
3636
},
37-
"repository": "pytorch-inference-eia",
38-
"end_of_support": "2022-06-21T00:00:00.0Z"
37+
"repository": "pytorch-inference-eia"
3938
},
4039
"1.5.1": {
4140
"py_versions": [
@@ -80,6 +79,14 @@
8079
"1.12": "1.12.1",
8180
"1.13": "1.13.1"
8281
},
82+
"version_support": {
83+
"1.12.1": {
84+
"end_of_support": "2023-07-01T00:00:00.0Z"
85+
},
86+
"1.13.1": {
87+
"end_of_support": "2023-10-28T00:00:00.0Z"
88+
}
89+
},
8390
"versions": {
8491
"0.4.0": {
8592
"py_versions": [
@@ -147,8 +154,7 @@
147154
"us-west-1": "520713654638",
148155
"us-west-2": "520713654638"
149156
},
150-
"repository": "sagemaker-pytorch",
151-
"end_of_support": "2022-06-21T00:00:00.0Z"
157+
"repository": "sagemaker-pytorch"
152158
},
153159
"1.1.0": {
154160
"py_versions": [
@@ -182,8 +188,7 @@
182188
"us-west-1": "520713654638",
183189
"us-west-2": "520713654638"
184190
},
185-
"repository": "sagemaker-pytorch",
186-
"end_of_support": "2022-06-21T00:00:00.0Z"
191+
"repository": "sagemaker-pytorch"
187192
},
188193
"1.2.0": {
189194
"py_versions": [
@@ -225,8 +230,7 @@
225230
"us-west-1": "763104351884",
226231
"us-west-2": "763104351884"
227232
},
228-
"repository": "pytorch-inference",
229-
"end_of_support": "2022-06-21T00:00:00.0Z"
233+
"repository": "pytorch-inference"
230234
},
231235
"1.3.1": {
232236
"py_versions": [
@@ -268,8 +272,7 @@
268272
"us-west-1": "763104351884",
269273
"us-west-2": "763104351884"
270274
},
271-
"repository": "pytorch-inference",
272-
"end_of_support": "2022-06-21T00:00:00.0Z"
275+
"repository": "pytorch-inference"
273276
},
274277
"1.4.0": {
275278
"py_versions": [
@@ -310,8 +313,7 @@
310313
"us-west-1": "763104351884",
311314
"us-west-2": "763104351884"
312315
},
313-
"repository": "pytorch-inference",
314-
"end_of_support": "2022-06-21T00:00:00.0Z"
316+
"repository": "pytorch-inference"
315317
},
316318
"1.5.0": {
317319
"py_versions": [
@@ -352,8 +354,7 @@
352354
"us-west-1": "763104351884",
353355
"us-west-2": "763104351884"
354356
},
355-
"repository": "pytorch-inference",
356-
"end_of_support": "2022-06-21T00:00:00.0Z"
357+
"repository": "pytorch-inference"
357358
},
358359
"1.6.0": {
359360
"py_versions": [
@@ -395,8 +396,7 @@
395396
"us-west-1": "763104351884",
396397
"us-west-2": "763104351884"
397398
},
398-
"repository": "pytorch-inference",
399-
"end_of_support": "2022-06-21T00:00:00.0Z"
399+
"repository": "pytorch-inference"
400400
},
401401
"1.7.1": {
402402
"py_versions": [
@@ -438,8 +438,7 @@
438438
"us-west-1": "763104351884",
439439
"us-west-2": "763104351884"
440440
},
441-
"repository": "pytorch-inference",
442-
"end_of_support": "2022-06-21T00:00:00.0Z"
441+
"repository": "pytorch-inference"
443442
},
444443
"1.8.0": {
445444
"py_versions": [
@@ -481,8 +480,7 @@
481480
"us-west-1": "763104351884",
482481
"us-west-2": "763104351884"
483482
},
484-
"repository": "pytorch-inference",
485-
"end_of_support": "2022-06-21T00:00:00.0Z"
483+
"repository": "pytorch-inference"
486484
},
487485
"1.8.1": {
488486
"py_versions": [
@@ -524,8 +522,7 @@
524522
"us-west-1": "763104351884",
525523
"us-west-2": "763104351884"
526524
},
527-
"repository": "pytorch-inference",
528-
"end_of_support": "2022-06-21T00:00:00.0Z"
525+
"repository": "pytorch-inference"
529526
},
530527
"1.9.0": {
531528
"py_versions": [
@@ -566,8 +563,7 @@
566563
"us-west-1": "763104351884",
567564
"us-west-2": "763104351884"
568565
},
569-
"repository": "pytorch-inference",
570-
"end_of_support": "2022-06-21T00:00:00.0Z"
566+
"repository": "pytorch-inference"
571567
},
572568
"1.9.1": {
573569
"py_versions": [
@@ -608,8 +604,7 @@
608604
"us-west-1": "763104351884",
609605
"us-west-2": "763104351884"
610606
},
611-
"repository": "pytorch-inference",
612-
"end_of_support": "2022-06-21T00:00:00.0Z"
607+
"repository": "pytorch-inference"
613608
},
614609
"1.10.0": {
615610
"py_versions": [
@@ -650,8 +645,7 @@
650645
"us-west-1": "763104351884",
651646
"us-west-2": "763104351884"
652647
},
653-
"repository": "pytorch-inference",
654-
"end_of_support": "2022-10-26T00:00:00.0Z"
648+
"repository": "pytorch-inference"
655649
},
656650
"1.10.2": {
657651
"py_versions": [
@@ -692,8 +686,7 @@
692686
"us-west-1": "763104351884",
693687
"us-west-2": "763104351884"
694688
},
695-
"repository": "pytorch-inference",
696-
"end_of_support": "2022-10-26T00:00:00.0Z"
689+
"repository": "pytorch-inference"
697690
},
698691
"1.11.0": {
699692
"py_versions": [
@@ -734,8 +727,7 @@
734727
"us-west-1": "763104351884",
735728
"us-west-2": "763104351884"
736729
},
737-
"repository": "pytorch-inference",
738-
"end_of_support": "2023-03-10T00:00:00.0Z"
730+
"repository": "pytorch-inference"
739731
},
740732
"1.12.0": {
741733
"py_versions": [
@@ -776,8 +768,7 @@
776768
"us-west-1": "763104351884",
777769
"us-west-2": "763104351884"
778770
},
779-
"repository": "pytorch-inference",
780-
"end_of_support": "2022-09-16T00:00:00.0Z"
771+
"repository": "pytorch-inference"
781772
},
782773
"1.12.1": {
783774
"py_versions": [
@@ -817,8 +808,7 @@
817808
"us-west-1": "763104351884",
818809
"us-west-2": "763104351884"
819810
},
820-
"repository": "pytorch-inference",
821-
"end_of_support": "2023-07-01T00:00:00.0Z"
811+
"repository": "pytorch-inference"
822812
},
823813
"1.13.1": {
824814
"py_versions": [
@@ -855,8 +845,7 @@
855845
"us-west-1": "763104351884",
856846
"us-west-2": "763104351884"
857847
},
858-
"repository": "pytorch-inference",
859-
"end_of_support": "2023-10-28T00:00:00.0Z"
848+
"repository": "pytorch-inference"
860849
}
861850
}
862851
},
@@ -933,6 +922,14 @@
933922
"1.12": "1.12.1",
934923
"1.13": "1.13.1"
935924
},
925+
"version_support": {
926+
"1.12.1": {
927+
"end_of_support": "2023-07-01T00:00:00.0Z"
928+
},
929+
"1.13.1": {
930+
"end_of_support": "2023-10-28T00:00:00.0Z"
931+
}
932+
},
936933
"versions": {
937934
"0.4.0": {
938935
"py_versions": [

‎src/sagemaker/image_uris.py‎

Lines changed: 40 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -447,39 +447,65 @@ def _validate_accelerator_type(accelerator_type):
447447
)
448448

449449

450-
def _get_end_of_support_warn_message(end_of_support, framework, version):
450+
def _get_end_of_support_warn_message(version, config, framework):
451451
"""
452452
Get end of support warning message if needed.
453453
454454
Args:
455-
end_of_support (str): json datetime string
456-
framework (str): ML framework
457455
version (str): ML framework version
456+
config (dict): config json loaded into python dictionary
457+
framework (str): ML framework
458458
459459
Returns:
460460
str: Warning message if version is nearing or out of support, else empty string
461461
"""
462-
if not end_of_support:
462+
version_support = config.get("version_support")
463+
aliases = config.get("version_aliases")
464+
465+
# If version_support and version_aliases are not implemented, do not return warning message
466+
if not (version_support and aliases):
463467
return ""
468+
464469
dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
470+
471+
end_of_support_msg = (
472+
f"The {framework} {version} DLC has reached end of support. "
473+
f"Please choose a supported version from our support policy - {dlc_support_policy}."
474+
)
475+
476+
long_version = version
477+
if version in aliases:
478+
long_version = aliases[version]
479+
480+
# If we have version support, but the long version is not in it, assume this is out of support
481+
if not version_support.get(long_version):
482+
return end_of_support_msg
483+
484+
end_of_support = version_support[long_version].get("end_of_support")
485+
486+
# If no end of support specified, assume there is indefinite support
487+
if not end_of_support:
488+
return ""
489+
465490
# Convert json object to UTC timezone string
466-
end_of_support_dt = datetime.datetime.strptime(
467-
end_of_support, '%Y-%m-%dT%H:%M:%S.%fZ'
468-
).replace(tzinfo=None).astimezone(tz=datetime.timezone.utc)
491+
end_of_support_dt = (
492+
datetime.datetime.strptime(end_of_support, "%Y-%m-%dT%H:%M:%S.%fZ")
493+
.replace(tzinfo=None)
494+
.astimezone(tz=datetime.timezone.utc)
495+
)
469496
# Ensure that the version is still supported
470497
current_dt = datetime.datetime.now(datetime.timezone.utc)
471498
time_delt_days = (end_of_support_dt - current_dt).days
472499
if current_dt >= end_of_support_dt:
473-
return (
474-
f"The {framework} {version} DLC has reached end of support. "
475-
f"Please choose a supported version from our support policy - {dlc_support_policy}."
476-
)
500+
return end_of_support_msg
477501
if time_delt_days <= 60:
478502
return (
479503
f"The {framework} {version} DLC is approaching end of support, "
480504
f"and patching will stop on {end_of_support_dt.strftime('%Y-%m-%d %Z')}. "
481505
f"Please choose a supported version from our support policy - {dlc_support_policy}."
482-
)
506+
)
507+
508+
# Version is still well in support window, do not warn
483509
return ""
484510

485511

@@ -488,8 +514,6 @@ def _validate_version_and_set_if_needed(version, config, framework):
488514
available_versions = list(config["versions"].keys())
489515
aliased_versions = list(config.get("version_aliases", {}).keys())
490516

491-
492-
493517
if len(available_versions) == 1 and version not in aliased_versions:
494518
log_message = "Defaulting to the only supported framework/algorithm version: {}.".format(
495519
available_versions[0]
@@ -503,12 +527,8 @@ def _validate_version_and_set_if_needed(version, config, framework):
503527

504528
_validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
505529

506-
# For DLCs, warn if image is out of support
507-
long_version = version
508-
if version in aliased_versions:
509-
long_version = config["version_aliases"][version]
510-
end_of_support = config["versions"][long_version].get("end_of_support")
511-
end_of_support_warning = _get_end_of_support_warn_message(end_of_support, framework, version)
530+
# For DLCs, warn if image is out of support. Use aliased version to grab information about the latest
531+
end_of_support_warning = _get_end_of_support_warn_message(version, config, framework)
512532
if end_of_support_warning:
513533
logger.warning(end_of_support_warning)
514534

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

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,15 @@ def _test_end_of_support_messaging(
108108
}
109109

110110
config = image_uris._config_for_framework_and_scope(framework=framework, image_scope=scope, accelerator_type=accelerator_type)
111-
end_of_support = config.get("end_of_support")
112-
expected_warning_message = image_uris._get_end_of_support_warn_message(end_of_support, framework, fw_version)
111+
version_support = config.get("version_support")
112+
aliases = config.get("version_aliases")
113+
114+
if version_support:
115+
for v, _ in version_support.items():
116+
# Version in version_support MUST be in an aliased value
117+
assert v in aliases.values()
118+
119+
expected_warning_message = image_uris._get_end_of_support_warn_message(fw_version, config, framework)
113120

114121
TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
115122
if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
@@ -122,7 +129,7 @@ def _test_end_of_support_messaging(
122129
_uri = image_uris.retrieve(region=REGION, instance_type=instance_type, **base_args)
123130
assert expected_warning_message in retrieve_io.getvalue()
124131

125-
for region in SAGEMAKER_ALTERNATE_REGION_ACCOUNTS.keys():
132+
for region in DLC_ALTERNATE_REGION_ACCOUNTS:
126133
retrieve_io = io.StringIO()
127134
with contextlib.redirect_stderr(retrieve_io):
128135
_uri = image_uris.retrieve(region=region, instance_type="ml.c4.xlarge", **base_args)
@@ -523,9 +530,18 @@ def _sagemaker_or_dlc_account(repo, region):
523530

524531

525532
def test_end_of_support():
526-
end_of_support = "2021-06-21T00:00:00.0Z"
533+
config = {
534+
"version_aliases": {
535+
"1.6": "1.6.0"
536+
},
537+
"version_support": {
538+
"1.6.0": {
539+
"end_of_support": "2021-06-21T00:00:00.0Z"
540+
}
541+
}
542+
}
527543
dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
528-
warn_message = image_uris._get_end_of_support_warn_message(end_of_support=end_of_support, framework="pytorch", version="1.6.0")
544+
warn_message = image_uris._get_end_of_support_warn_message(config=config, framework="pytorch", version="1.6.0")
529545

530546
assert warn_message == (
531547
f"The pytorch 1.6.0 DLC has reached end of support. "

0 commit comments

Comments
 (0)