From b71adc349b0faa6f6afab6ab46ab40ee7ef814cf Mon Sep 17 00:00:00 2001
From: arjkesh <33526713+arjkesh@users.noreply.github.com>
Date: Wed, 29 Mar 2023 16:56:47 -0700
Subject: [PATCH 1/5] Add end of support messaging for DLCs

---
 src/sagemaker/image_uri_config/pytorch.json   |  57 +++++---
 src/sagemaker/image_uris.py                   |  43 ++++++
 .../image_uris/test_dlc_frameworks.py         | 122 ++++++++++++++++--
 3 files changed, 191 insertions(+), 31 deletions(-)

diff --git a/src/sagemaker/image_uri_config/pytorch.json b/src/sagemaker/image_uri_config/pytorch.json
index b833c1d7a6..e125445b53 100644
--- a/src/sagemaker/image_uri_config/pytorch.json
+++ b/src/sagemaker/image_uri_config/pytorch.json
@@ -34,7 +34,8 @@
                     "us-east-2": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-eia"
+                "repository": "pytorch-inference-eia",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.5.1": {
                 "py_versions": [
@@ -146,7 +147,8 @@
                     "us-west-1": "520713654638",
                     "us-west-2": "520713654638"
                 },
-                "repository": "sagemaker-pytorch"
+                "repository": "sagemaker-pytorch",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.1.0": {
                 "py_versions": [
@@ -180,7 +182,8 @@
                     "us-west-1": "520713654638",
                     "us-west-2": "520713654638"
                 },
-                "repository": "sagemaker-pytorch"
+                "repository": "sagemaker-pytorch",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.2.0": {
                 "py_versions": [
@@ -222,7 +225,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.3.1": {
                 "py_versions": [
@@ -264,7 +268,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.4.0": {
                 "py_versions": [
@@ -305,7 +310,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.5.0": {
                 "py_versions": [
@@ -346,7 +352,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.6.0": {
                 "py_versions": [
@@ -388,7 +395,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.7.1": {
                 "py_versions": [
@@ -430,7 +438,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.8.0": {
                 "py_versions": [
@@ -472,7 +481,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.8.1": {
                 "py_versions": [
@@ -514,7 +524,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.9.0": {
                 "py_versions": [
@@ -555,7 +566,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.9.1": {
                 "py_versions": [
@@ -596,7 +608,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-06-21T00:00:00.0Z"
             },
             "1.10.0": {
                 "py_versions": [
@@ -637,7 +650,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-10-26T00:00:00.0Z"
             },
             "1.10.2": {
                 "py_versions": [
@@ -678,7 +692,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-10-26T00:00:00.0Z"
             },
             "1.11.0": {
                 "py_versions": [
@@ -719,7 +734,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2023-03-10T00:00:00.0Z"
             },
             "1.12.0": {
                 "py_versions": [
@@ -760,7 +776,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2022-09-16T00:00:00.0Z"
             },
             "1.12.1": {
                 "py_versions": [
@@ -800,7 +817,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2023-07-01T00:00:00.0Z"
             },
             "1.13.1": {
                 "py_versions": [
@@ -837,7 +855,8 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference"
+                "repository": "pytorch-inference",
+                "end_of_support": "2023-10-28T00:00:00.0Z"
             }
         }
     },
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index b9101acd96..e1ea74c444 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -17,6 +17,7 @@
 import logging
 import os
 import re
+import datetime
 from typing import Optional
 from packaging.version import Version
 
@@ -446,6 +447,42 @@ def _validate_accelerator_type(accelerator_type):
         )
 
 
+def _get_end_of_support_warn_message(end_of_support, framework, version):
+    """
+    Get end of support warning message if needed.
+
+    Args:
+        end_of_support (str): json datetime string
+        framework (str): ML framework
+        version (str): ML framework version
+
+    Returns:
+        str: Warning message if version is nearing or out of support, else empty string
+    """
+    if not end_of_support:
+        return ""
+    dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
+    # Convert json object to UTC timezone string
+    end_of_support_dt = datetime.datetime.strptime(
+        end_of_support, '%Y-%m-%dT%H:%M:%S.%fZ'
+        ).replace(tzinfo=None).astimezone(tz=datetime.timezone.utc)
+    # Ensure that the version is still supported
+    current_dt = datetime.datetime.now(datetime.timezone.utc)
+    time_delt_days = (end_of_support_dt - current_dt).days
+    if current_dt >= end_of_support_dt:
+        return (
+            f"Unsupported DLC {framework} version: {version}." 
+            f"Please choose a supported version from our support policy - {dlc_support_policy}"
+            )
+    if time_delt_days <= 60:
+        return (
+            f"The {framework} {version} DLC is approaching end of support, "
+            f"and patching will stop on {end_of_support}. "
+            f"Please choose a supported version from our support policy - {dlc_support_policy}"
+            )
+    return ""
+
+
 def _validate_version_and_set_if_needed(version, config, framework):
     """Checks if the framework/algorithm version is one of the supported versions."""
     available_versions = list(config["versions"].keys())
@@ -463,6 +500,12 @@ def _validate_version_and_set_if_needed(version, config, framework):
         return available_versions[0]
 
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
+
+    # For DLCs, warn if image is out of support
+    end_of_support = config.get("end_of_support")
+    end_of_support_warning = _get_end_of_support_warn_message(end_of_support, framework, version)
+    if end_of_support_warning:
+        logger.warning(end_of_support_warning)
     return version
 
 
diff --git a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
index 0f00f7eb22..468ba6e66f 100644
--- a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
+++ b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
@@ -12,6 +12,10 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import io
+import contextlib
+import datetime
+
 from packaging.version import Version
 
 from sagemaker import image_uris
@@ -89,6 +93,42 @@ def _test_image_uris(
         assert expected == uri
 
 
+def _test_end_of_support_messaging(
+    framework,
+    fw_version,
+    py_version,
+    scope,
+    accelerator_type=None
+):
+    base_args = {
+        "framework": framework,
+        "version": fw_version,
+        "py_version": py_version,
+        "image_scope": scope,
+    }
+
+    config = image_uris._config_for_framework_and_scope(framework=framework, image_scope=scope, accelerator_type=accelerator_type)
+    end_of_support = config.get("end_of_support")
+    expected_warning_message = image_uris._get_end_of_support_warn_message(end_of_support, framework, fw_version)
+
+    TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
+    if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
+        """Handle P2 deprecation"""
+        TYPES_AND_PROCESSORS = RENEWED_PYTORCH_INSTANCE_TYPES_AND_PROCESSORS
+
+    for instance_type, _processor in TYPES_AND_PROCESSORS:
+        retrieve_io = io.StringIO()
+        with contextlib.redirect_stderr(retrieve_io):
+            _uri = image_uris.retrieve(region=REGION, instance_type=instance_type, **base_args)
+        assert expected_warning_message in retrieve_io.getvalue()
+
+    for region in SAGEMAKER_ALTERNATE_REGION_ACCOUNTS.keys():
+        retrieve_io = io.StringIO()
+        with contextlib.redirect_stderr(retrieve_io):
+            _uri = image_uris.retrieve(region=region, instance_type="ml.c4.xlarge", **base_args)
+        assert expected_warning_message in retrieve_io.getvalue()
+
+
 def test_chainer(chainer_version, chainer_py_version):
     expected_fn_args = {
         "chainer_version": chainer_version,
@@ -122,16 +162,26 @@ def test_tensorflow_training(tensorflow_training_version, tensorflow_training_py
         "tf_training_version": tensorflow_training_version,
         "py_version": tensorflow_training_py_version,
     }
+    
+    framework = "tensorflow"
+    scope = "training"
 
     _test_image_uris(
-        "tensorflow",
+        framework,
         tensorflow_training_version,
         tensorflow_training_py_version,
-        "training",
+        scope,
         _expected_tf_training_uri,
         expected_fn_args,
     )
 
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=tensorflow_training_version,
+        py_version=tensorflow_training_py_version,
+        scope=scope,
+    )
+
 
 def _expected_tf_training_uri(tf_training_version, py_version, processor="cpu", region=REGION):
     version = Version(tf_training_version)
@@ -155,15 +205,25 @@ def _expected_tf_training_uri(tf_training_version, py_version, processor="cpu",
 
 
 def test_tensorflow_inference(tensorflow_inference_version, tensorflow_inference_py_version):
+    framework = "tensorflow"
+    scope = "inference"
+
     _test_image_uris(
-        "tensorflow",
+        framework,
         tensorflow_inference_version,
         tensorflow_inference_py_version,
-        "inference",
+        scope,
         _expected_tf_inference_uri,
         {"tf_inference_version": tensorflow_inference_version},
     )
 
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=tensorflow_inference_version,
+        py_version=tensorflow_inference_py_version,
+        scope=scope,
+    )
+
 
 def test_tensorflow_eia(tensorflow_eia_version):
     base_args = {
@@ -222,16 +282,25 @@ def test_mxnet_training(mxnet_training_version, mxnet_training_py_version):
         "mxnet_version": mxnet_training_version,
         "py_version": mxnet_training_py_version,
     }
+    framework = "mxnet"
+    scope = "training"
 
     _test_image_uris(
-        "mxnet",
+        framework,
         mxnet_training_version,
         mxnet_training_py_version,
-        "training",
+        scope,
         _expected_mxnet_training_uri,
         expected_fn_args,
     )
 
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=mxnet_training_version,
+        py_version=mxnet_training_py_version,
+        scope=scope,
+    )
+
 
 def _expected_mxnet_training_uri(mxnet_version, py_version, processor="cpu", region=REGION):
     version = Version(mxnet_version)
@@ -258,15 +327,25 @@ def test_mxnet_inference(mxnet_inference_version, mxnet_inference_py_version):
         "py_version": mxnet_inference_py_version,
     }
 
+    framework = "mxnet"
+    scope = "inference"
+
     _test_image_uris(
-        "mxnet",
+        framework,
         mxnet_inference_version,
         mxnet_inference_py_version,
-        "inference",
+        scope,
         _expected_mxnet_inference_uri,
         expected_fn_args,
     )
 
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=mxnet_inference_version,
+        py_version=mxnet_inference_py_version,
+        scope=scope,
+    )
+
 
 def test_mxnet_eia(mxnet_eia_version, mxnet_eia_py_version):
     base_args = {
@@ -319,17 +398,26 @@ def _expected_mxnet_inference_uri(
 
 
 def test_pytorch_training(pytorch_training_version, pytorch_training_py_version):
+    framework = "pytorch"
+    scope = "training"
+
     _test_image_uris(
-        "pytorch",
+        framework,
         pytorch_training_version,
         pytorch_training_py_version,
-        "training",
+        scope,
         _expected_pytorch_training_uri,
         {
             "pytorch_version": pytorch_training_version,
             "py_version": pytorch_training_py_version,
         },
     )
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=pytorch_training_version,
+        py_version=pytorch_training_py_version,
+        scope=scope,
+    )
 
 
 def _expected_pytorch_training_uri(pytorch_version, py_version, processor="cpu", region=REGION):
@@ -350,11 +438,14 @@ def _expected_pytorch_training_uri(pytorch_version, py_version, processor="cpu",
 
 
 def test_pytorch_inference(pytorch_inference_version, pytorch_inference_py_version):
+    framework = "pytorch"
+    scope = "inference"
+
     _test_image_uris(
-        "pytorch",
+        framework,
         pytorch_inference_version,
         pytorch_inference_py_version,
-        "inference",
+        scope,
         _expected_pytorch_inference_uri,
         {
             "pytorch_version": pytorch_inference_version,
@@ -362,6 +453,13 @@ def test_pytorch_inference(pytorch_inference_version, pytorch_inference_py_versi
         },
     )
 
+    _test_end_of_support_messaging(
+        framework=framework,
+        fw_version=pytorch_inference_version,
+        py_version=pytorch_inference_py_version,
+        scope=scope,
+    )
+
 
 def _expected_pytorch_inference_uri(pytorch_version, py_version, processor="cpu", region=REGION):
     version = Version(pytorch_version)

From aee40bca0423ea6d7c6d5ef4ee8f204eed0736ec Mon Sep 17 00:00:00 2001
From: arjkesh <33526713+arjkesh@users.noreply.github.com>
Date: Wed, 29 Mar 2023 17:06:44 -0700
Subject: [PATCH 2/5] add additional end of support test

---
 .../unit/sagemaker/image_uris/test_dlc_frameworks.py  | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
index 468ba6e66f..4e1d99a2a4 100644
--- a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
+++ b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
@@ -520,3 +520,14 @@ def _sagemaker_or_dlc_account(repo, region):
         )
     else:
         return DLC_ACCOUNT if region == REGION else DLC_ALTERNATE_REGION_ACCOUNTS[region]
+
+
+def test_end_of_support():
+    end_of_support = "2021-06-21T00:00:00.0Z"
+    dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
+    warn_message = image_uris._get_end_of_support_warn_message(end_of_support=end_of_support, framework="pytorch", version="1.6.0")
+
+    assert warn_message == (
+        f"Unsupported DLC pytorch version: 1.6.0" 
+        f"Please choose a supported version from our support policy - {dlc_support_policy}"
+        )

From 61fc611049b15fa15eea77ae140bf8c26e1fa44a Mon Sep 17 00:00:00 2001
From: arjkesh <33526713+arjkesh@users.noreply.github.com>
Date: Wed, 29 Mar 2023 17:44:46 -0700
Subject: [PATCH 3/5] update image uris and tests

---
 src/sagemaker/image_uris.py                      | 16 +++++++++++-----
 .../sagemaker/image_uris/test_dlc_frameworks.py  |  4 ++--
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index e1ea74c444..734a5f699c 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -471,14 +471,14 @@ def _get_end_of_support_warn_message(end_of_support, framework, version):
     time_delt_days = (end_of_support_dt - current_dt).days
     if current_dt >= end_of_support_dt:
         return (
-            f"Unsupported DLC {framework} version: {version}." 
-            f"Please choose a supported version from our support policy - {dlc_support_policy}"
+            f"The {framework} {version} DLC has reached end of support. " 
+            f"Please choose a supported version from our support policy - {dlc_support_policy}."
             )
     if time_delt_days <= 60:
         return (
             f"The {framework} {version} DLC is approaching end of support, "
-            f"and patching will stop on {end_of_support}. "
-            f"Please choose a supported version from our support policy - {dlc_support_policy}"
+            f"and patching will stop on {end_of_support_dt.strftime('%Y-%m-%d %Z')}. "
+            f"Please choose a supported version from our support policy - {dlc_support_policy}."
             )
     return ""
 
@@ -488,6 +488,8 @@ def _validate_version_and_set_if_needed(version, config, framework):
     available_versions = list(config["versions"].keys())
     aliased_versions = list(config.get("version_aliases", {}).keys())
 
+
+
     if len(available_versions) == 1 and version not in aliased_versions:
         log_message = "Defaulting to the only supported framework/algorithm version: {}.".format(
             available_versions[0]
@@ -502,10 +504,14 @@ def _validate_version_and_set_if_needed(version, config, framework):
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
 
     # For DLCs, warn if image is out of support
-    end_of_support = config.get("end_of_support")
+    long_version = version
+    if version in aliased_versions:
+        long_version = config["version_aliases"][version]
+    end_of_support = config["versions"][long_version].get("end_of_support")
     end_of_support_warning = _get_end_of_support_warn_message(end_of_support, framework, version)
     if end_of_support_warning:
         logger.warning(end_of_support_warning)
+
     return version
 
 
diff --git a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
index 4e1d99a2a4..1534981258 100644
--- a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
+++ b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
@@ -528,6 +528,6 @@ def test_end_of_support():
     warn_message = image_uris._get_end_of_support_warn_message(end_of_support=end_of_support, framework="pytorch", version="1.6.0")
 
     assert warn_message == (
-        f"Unsupported DLC pytorch version: 1.6.0" 
-        f"Please choose a supported version from our support policy - {dlc_support_policy}"
+        f"The pytorch 1.6.0 DLC has reached end of support. " 
+        f"Please choose a supported version from our support policy - {dlc_support_policy}."
         )

From 07fa0a9a76f15b23516f8a9b8c75547c047988e5 Mon Sep 17 00:00:00 2001
From: arjkesh <33526713+arjkesh@users.noreply.github.com>
Date: Wed, 29 Mar 2023 21:42:17 -0700
Subject: [PATCH 4/5] refactor for ease of adding to support policy

---
 src/sagemaker/image_uri_config/pytorch.json   | 73 +++++++++----------
 src/sagemaker/image_uris.py                   | 60 ++++++++++-----
 .../image_uris/test_dlc_frameworks.py         | 26 +++++--
 3 files changed, 96 insertions(+), 63 deletions(-)

diff --git a/src/sagemaker/image_uri_config/pytorch.json b/src/sagemaker/image_uri_config/pytorch.json
index e125445b53..3193524056 100644
--- a/src/sagemaker/image_uri_config/pytorch.json
+++ b/src/sagemaker/image_uri_config/pytorch.json
@@ -34,8 +34,7 @@
                     "us-east-2": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference-eia",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference-eia"
             },
             "1.5.1": {
                 "py_versions": [
@@ -80,6 +79,14 @@
             "1.12": "1.12.1",
             "1.13": "1.13.1"
         },
+        "version_support": {
+            "1.12.1": {
+                "end_of_support": "2023-07-01T00:00:00.0Z"
+            },
+            "1.13.1": {
+                "end_of_support": "2023-10-28T00:00:00.0Z"
+            }
+        },
         "versions": {
             "0.4.0": {
                 "py_versions": [
@@ -147,8 +154,7 @@
                     "us-west-1": "520713654638",
                     "us-west-2": "520713654638"
                 },
-                "repository": "sagemaker-pytorch",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "sagemaker-pytorch"
             },
             "1.1.0": {
                 "py_versions": [
@@ -182,8 +188,7 @@
                     "us-west-1": "520713654638",
                     "us-west-2": "520713654638"
                 },
-                "repository": "sagemaker-pytorch",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "sagemaker-pytorch"
             },
             "1.2.0": {
                 "py_versions": [
@@ -225,8 +230,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.3.1": {
                 "py_versions": [
@@ -268,8 +272,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.4.0": {
                 "py_versions": [
@@ -310,8 +313,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.5.0": {
                 "py_versions": [
@@ -352,8 +354,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference" 
             },
             "1.6.0": {
                 "py_versions": [
@@ -395,8 +396,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"        
             },
             "1.7.1": {
                 "py_versions": [
@@ -438,8 +438,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.8.0": {
                 "py_versions": [
@@ -481,8 +480,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.8.1": {
                 "py_versions": [
@@ -524,8 +522,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.9.0": {
                 "py_versions": [
@@ -566,8 +563,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.9.1": {
                 "py_versions": [
@@ -608,8 +604,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-06-21T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.10.0": {
                 "py_versions": [
@@ -650,8 +645,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-10-26T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.10.2": {
                 "py_versions": [
@@ -692,8 +686,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-10-26T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.11.0": {
                 "py_versions": [
@@ -734,8 +727,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2023-03-10T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.12.0": {
                 "py_versions": [
@@ -776,8 +768,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2022-09-16T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.12.1": {
                 "py_versions": [
@@ -817,8 +808,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2023-07-01T00:00:00.0Z"
+                "repository": "pytorch-inference"
             },
             "1.13.1": {
                 "py_versions": [
@@ -855,8 +845,7 @@
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
                 },
-                "repository": "pytorch-inference",
-                "end_of_support": "2023-10-28T00:00:00.0Z"
+                "repository": "pytorch-inference"
             }
         }
     },
@@ -933,6 +922,14 @@
             "1.12": "1.12.1",
             "1.13": "1.13.1"
         },
+        "version_support": {
+            "1.12.1": {
+                "end_of_support": "2023-07-01T00:00:00.0Z"
+            },
+            "1.13.1": {
+                "end_of_support": "2023-10-28T00:00:00.0Z"
+            }
+        },
         "versions": {
             "0.4.0": {
                 "py_versions": [
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index 734a5f699c..7fec19976c 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -447,39 +447,65 @@ def _validate_accelerator_type(accelerator_type):
         )
 
 
-def _get_end_of_support_warn_message(end_of_support, framework, version):
+def _get_end_of_support_warn_message(version, config, framework):
     """
     Get end of support warning message if needed.
 
     Args:
-        end_of_support (str): json datetime string
-        framework (str): ML framework
         version (str): ML framework version
+        config (dict): config json loaded into python dictionary
+        framework (str): ML framework
 
     Returns:
         str: Warning message if version is nearing or out of support, else empty string
     """
-    if not end_of_support:
+    version_support = config.get("version_support")
+    aliases = config.get("version_aliases")
+
+    # If version_support and version_aliases are not implemented, do not return warning message
+    if not (version_support and aliases):
         return ""
+
     dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
+
+    end_of_support_msg = (
+        f"The {framework} {version} DLC has reached end of support. "
+        f"Please choose a supported version from our support policy - {dlc_support_policy}."
+    )
+
+    long_version = version
+    if version in aliases:
+        long_version = aliases[version]
+
+    # If we have version support, but the long version is not in it, assume this is out of support
+    if not version_support.get(long_version):
+        return end_of_support_msg
+
+    end_of_support = version_support[long_version].get("end_of_support")
+
+    # If no end of support specified, assume there is indefinite support
+    if not end_of_support:
+        return ""
+
     # Convert json object to UTC timezone string
-    end_of_support_dt = datetime.datetime.strptime(
-        end_of_support, '%Y-%m-%dT%H:%M:%S.%fZ'
-        ).replace(tzinfo=None).astimezone(tz=datetime.timezone.utc)
+    end_of_support_dt = (
+        datetime.datetime.strptime(end_of_support, "%Y-%m-%dT%H:%M:%S.%fZ")
+        .replace(tzinfo=None)
+        .astimezone(tz=datetime.timezone.utc)
+    )
     # Ensure that the version is still supported
     current_dt = datetime.datetime.now(datetime.timezone.utc)
     time_delt_days = (end_of_support_dt - current_dt).days
     if current_dt >= end_of_support_dt:
-        return (
-            f"The {framework} {version} DLC has reached end of support. " 
-            f"Please choose a supported version from our support policy - {dlc_support_policy}."
-            )
+        return end_of_support_msg
     if time_delt_days <= 60:
         return (
             f"The {framework} {version} DLC is approaching end of support, "
             f"and patching will stop on {end_of_support_dt.strftime('%Y-%m-%d %Z')}. "
             f"Please choose a supported version from our support policy - {dlc_support_policy}."
-            )
+        )
+
+    # Version is still well in support window, do not warn
     return ""
 
 
@@ -488,8 +514,6 @@ def _validate_version_and_set_if_needed(version, config, framework):
     available_versions = list(config["versions"].keys())
     aliased_versions = list(config.get("version_aliases", {}).keys())
 
-
-
     if len(available_versions) == 1 and version not in aliased_versions:
         log_message = "Defaulting to the only supported framework/algorithm version: {}.".format(
             available_versions[0]
@@ -503,12 +527,8 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
 
-    # For DLCs, warn if image is out of support
-    long_version = version
-    if version in aliased_versions:
-        long_version = config["version_aliases"][version]
-    end_of_support = config["versions"][long_version].get("end_of_support")
-    end_of_support_warning = _get_end_of_support_warn_message(end_of_support, framework, version)
+    # For DLCs, warn if image is out of support. Use aliased version to grab information about the latest
+    end_of_support_warning = _get_end_of_support_warn_message(version, config, framework)
     if end_of_support_warning:
         logger.warning(end_of_support_warning)
 
diff --git a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
index 1534981258..0b6b46bea8 100644
--- a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
+++ b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
@@ -108,8 +108,15 @@ def _test_end_of_support_messaging(
     }
 
     config = image_uris._config_for_framework_and_scope(framework=framework, image_scope=scope, accelerator_type=accelerator_type)
-    end_of_support = config.get("end_of_support")
-    expected_warning_message = image_uris._get_end_of_support_warn_message(end_of_support, framework, fw_version)
+    version_support = config.get("version_support")
+    aliases = config.get("version_aliases")
+
+    if version_support:
+        for v, _ in version_support.items():
+            # Version in version_support MUST be in an aliased value
+            assert v in aliases.values()
+
+    expected_warning_message = image_uris._get_end_of_support_warn_message(fw_version, config, framework)
 
     TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
     if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
@@ -122,7 +129,7 @@ def _test_end_of_support_messaging(
             _uri = image_uris.retrieve(region=REGION, instance_type=instance_type, **base_args)
         assert expected_warning_message in retrieve_io.getvalue()
 
-    for region in SAGEMAKER_ALTERNATE_REGION_ACCOUNTS.keys():
+    for region in DLC_ALTERNATE_REGION_ACCOUNTS:
         retrieve_io = io.StringIO()
         with contextlib.redirect_stderr(retrieve_io):
             _uri = image_uris.retrieve(region=region, instance_type="ml.c4.xlarge", **base_args)
@@ -523,9 +530,18 @@ def _sagemaker_or_dlc_account(repo, region):
 
 
 def test_end_of_support():
-    end_of_support = "2021-06-21T00:00:00.0Z"
+    config = {
+        "version_aliases": {
+            "1.6": "1.6.0"
+        },
+        "version_support": {
+            "1.6.0": {
+                "end_of_support": "2021-06-21T00:00:00.0Z"
+            }
+        }
+    }
     dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
-    warn_message = image_uris._get_end_of_support_warn_message(end_of_support=end_of_support, framework="pytorch", version="1.6.0")
+    warn_message = image_uris._get_end_of_support_warn_message(config=config, framework="pytorch", version="1.6.0")
 
     assert warn_message == (
         f"The pytorch 1.6.0 DLC has reached end of support. " 

From abff3778e5a55a1937a52722528977aa4b139d55 Mon Sep 17 00:00:00 2001
From: arjkesh <33526713+arjkesh@users.noreply.github.com>
Date: Wed, 29 Mar 2023 22:25:26 -0700
Subject: [PATCH 5/5] run formatting

---
 src/sagemaker/image_uris.py                   | 32 ++++++++++++----
 .../image_uris/test_dlc_frameworks.py         | 37 ++++++++-----------
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
index 7fec19976c..b3c93f11bf 100644
--- a/src/sagemaker/image_uris.py
+++ b/src/sagemaker/image_uris.py
@@ -164,7 +164,9 @@ def retrieve(
         config = _config_for_framework_and_scope(_framework, final_image_scope, accelerator_type)
 
     original_version = version
-    version = _validate_version_and_set_if_needed(version, config, framework)
+    version = _validate_version_and_set_if_needed(
+        version, config, framework, image_scope, base_framework_version
+    )
     version_config = config["versions"][_version_for_config(version, config)]
 
     if framework == HUGGING_FACE_FRAMEWORK:
@@ -199,9 +201,8 @@ def retrieve(
         container_version = sdk_version + "-" + container_version
 
     if framework == HUGGING_FACE_FRAMEWORK:
-        pt_or_tf_version = (
-            re.compile("^(pytorch|tensorflow)(.*)$").match(base_framework_version).group(2)
-        )
+        hf_base_fw_regex = _get_hf_base_framework_version_regex()
+        pt_or_tf_version = hf_base_fw_regex.match(base_framework_version).group(2)
         _version = original_version
 
         if repo in [
@@ -255,6 +256,10 @@ def retrieve(
     return ECR_URI_TEMPLATE.format(registry=registry, hostname=hostname, repository=repo)
 
 
+def _get_hf_base_framework_version_regex():
+    return re.compile("^(pytorch|tensorflow)(.*)$")
+
+
 def _get_instance_type_family(instance_type):
     """Return the family of the instance type.
 
@@ -509,7 +514,7 @@ def _get_end_of_support_warn_message(version, config, framework):
     return ""
 
 
-def _validate_version_and_set_if_needed(version, config, framework):
+def _validate_version_and_set_if_needed(version, config, framework, scope, base_framework_version):
     """Checks if the framework/algorithm version is one of the supported versions."""
     available_versions = list(config["versions"].keys())
     aliased_versions = list(config.get("version_aliases", {}).keys())
@@ -527,8 +532,21 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
 
-    # For DLCs, warn if image is out of support. Use aliased version to grab information about the latest
-    end_of_support_warning = _get_end_of_support_warn_message(version, config, framework)
+    # For DLCs, warn if image is out of support.
+    eos_version, eos_config, eos_framework = version, config, framework
+
+    # Default to underlying base framework for HF, except for neuron which has its own policies.
+    inf_or_trn = "inf" in config.get("processors", []) or "trn" in config.get("processors", [])
+    if framework == HUGGING_FACE_FRAMEWORK and base_framework_version and not inf_or_trn:
+        hf_base_fw_regex = _get_hf_base_framework_version_regex()
+        hf_match = hf_base_fw_regex.search(base_framework_version)
+        if hf_match:
+            eos_framework = hf_match.group(1)
+            eos_version = hf_match.group(2)
+            eos_config = _config_for_framework_and_scope(eos_framework, scope)
+    end_of_support_warning = _get_end_of_support_warn_message(
+        eos_version, eos_config, eos_framework
+    )
     if end_of_support_warning:
         logger.warning(end_of_support_warning)
 
diff --git a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
index 0b6b46bea8..a61180d3de 100644
--- a/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
+++ b/tests/unit/sagemaker/image_uris/test_dlc_frameworks.py
@@ -14,7 +14,6 @@
 
 import io
 import contextlib
-import datetime
 
 from packaging.version import Version
 
@@ -93,13 +92,7 @@ def _test_image_uris(
         assert expected == uri
 
 
-def _test_end_of_support_messaging(
-    framework,
-    fw_version,
-    py_version,
-    scope,
-    accelerator_type=None
-):
+def _test_end_of_support_messaging(framework, fw_version, py_version, scope, accelerator_type=None):
     base_args = {
         "framework": framework,
         "version": fw_version,
@@ -107,7 +100,9 @@ def _test_end_of_support_messaging(
         "image_scope": scope,
     }
 
-    config = image_uris._config_for_framework_and_scope(framework=framework, image_scope=scope, accelerator_type=accelerator_type)
+    config = image_uris._config_for_framework_and_scope(
+        framework=framework, image_scope=scope, accelerator_type=accelerator_type
+    )
     version_support = config.get("version_support")
     aliases = config.get("version_aliases")
 
@@ -116,7 +111,9 @@ def _test_end_of_support_messaging(
             # Version in version_support MUST be in an aliased value
             assert v in aliases.values()
 
-    expected_warning_message = image_uris._get_end_of_support_warn_message(fw_version, config, framework)
+    expected_warning_message = image_uris._get_end_of_support_warn_message(
+        fw_version, config, framework
+    )
 
     TYPES_AND_PROCESSORS = INSTANCE_TYPES_AND_PROCESSORS
     if framework == "pytorch" and Version(fw_version) >= Version("1.13"):
@@ -169,7 +166,7 @@ def test_tensorflow_training(tensorflow_training_version, tensorflow_training_py
         "tf_training_version": tensorflow_training_version,
         "py_version": tensorflow_training_py_version,
     }
-    
+
     framework = "tensorflow"
     scope = "training"
 
@@ -531,19 +528,15 @@ def _sagemaker_or_dlc_account(repo, region):
 
 def test_end_of_support():
     config = {
-        "version_aliases": {
-            "1.6": "1.6.0"
-        },
-        "version_support": {
-            "1.6.0": {
-                "end_of_support": "2021-06-21T00:00:00.0Z"
-            }
-        }
+        "version_aliases": {"1.6": "1.6.0"},
+        "version_support": {"1.6.0": {"end_of_support": "2021-06-21T00:00:00.0Z"}},
     }
     dlc_support_policy = "https://aws.amazon.com/releasenotes/dlc-support-policy/"
-    warn_message = image_uris._get_end_of_support_warn_message(config=config, framework="pytorch", version="1.6.0")
+    warn_message = image_uris._get_end_of_support_warn_message(
+        config=config, framework="pytorch", version="1.6.0"
+    )
 
     assert warn_message == (
-        f"The pytorch 1.6.0 DLC has reached end of support. " 
+        f"The pytorch 1.6.0 DLC has reached end of support. "
         f"Please choose a supported version from our support policy - {dlc_support_policy}."
-        )
+    )