run formatting

Author: arjkesh

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)
 

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,21 +92,17 @@ 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,
         "py_version": py_version,
         "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}."
-        )
+    )