Merge branch 'master' into pt21_refactor

pytorch/training/buildspec.yml

@@ -44,6 +44,7 @@ images:
     latest_release_tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-sagemaker" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: sagemaker
+    # build_tag_override: "beta:2.2.0-cpu-py310-ubuntu20.04-sagemaker"
     context:
       <<: *TRAINING_CONTEXT
   BuildSageMakerGPUPTTrainPy3DockerImage:
@@ -60,5 +61,6 @@ images:
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]
     target: sagemaker
+    # build_tag_override: "beta:2.2.0-gpu-py310-cu121-ubuntu20.04-sagemaker"
     context:
       <<: *TRAINING_CONTEXT

release_images_inference.yml

@@ -209,3 +209,16 @@ release_images:
       example: False
       disable_sm_tag: False
       force_release: False
+  17:
+    framework: "pytorch"
+    version: "2.2.0"
+    arch_type: "x86"
+    customer_type: "ec2"
+    inference:
+      device_types: [ "cpu", "gpu" ]
+      python_versions: [ "py310" ]
+      os_version: "ubuntu20.04"
+      cuda_version: "cu118"
+      example: False
+      disable_sm_tag: False
+      force_release: False

src/config.py

@@ -56,16 +56,13 @@ def is_autopatch_build_enabled_in_image_buildspec(buildspec_path=None):
     image_buildspec_object = Buildspec()
     image_buildspec_object.load(buildspec_path)
     autopatch_build_flag = image_buildspec_object.get("autopatch_build", "False").lower() == "true"
-    print(f"Here: {buildspec_path}")
     return autopatch_build_flag
 
 
 def is_autopatch_build_enabled(buildspec_path=None):
-    return (
-        parse_dlc_developer_configs("build", "autopatch_build")
-        or os.getenv("AUTOPATCH")
-        or is_autopatch_build_enabled_in_image_buildspec(buildspec_path=buildspec_path)
-    )
+    return parse_dlc_developer_configs(
+        "build", "autopatch_build"
+    ) or is_autopatch_build_enabled_in_image_buildspec(buildspec_path=buildspec_path)
 
 
 def is_ec2_test_enabled():

src/image.py

@@ -192,6 +192,7 @@ def docker_build(self, fileobj=None, custom_context=False):
         :return: int, Build Status
         """
         response = [f"Starting the Build Process for {self.repository}:{self.tag}"]
+
         for line in self.client.build(
             fileobj=fileobj,
             path=self.dockerfile,

src/image_builder.py

@@ -18,6 +18,7 @@
 import os
 import re
 import json
+import tempfile
 
 from copy import deepcopy
 
@@ -140,10 +141,7 @@ def image_builder(buildspec, image_types=[], device_types=[]):
         if is_nightly_build_context():
             additional_image_tags.append(tag_image_with_date(image_tag))
 
-        if build_context != "PR":
-            image_tag = tag_image_with_datetime(image_tag)
-        # If build is not enabled, we don't care about the datetime tag
-        elif is_build_enabled():
+        if is_build_enabled() or build_context != "PR":
             # Order appears to matter in datetime tagging, so tag with no datetime first, then
             # set image_tag to have datetime
             no_datetime = image_tag
@@ -211,10 +209,34 @@ def image_builder(buildspec, image_types=[], device_types=[]):
         if inference_toolkit_version:
             extra_build_args["SM_TOOLKIT_VERSION"] = inference_toolkit_version
 
+        tag_override = image_config.get("build_tag_override")
+        dockerfile = image_config["docker_file"]
+        target = image_config.get("target")
+        tag_override_regex = r"^(beta|pr):\S+$"
+        if tag_override and build_context == "PR":
+            if is_autopatch_build_enabled(buildspec_path=buildspec):
+                FORMATTER.print("AUTOPATCH ENABLED IN BUILDSPEC, CANNOT OVERRIDE WITH TAG, SORRY!")
+            elif not re.match(tag_override_regex, tag_override):
+                FORMATTER.print(
+                    f"TAG OVERRIDE MUST BE OF FORMAT {tag_override_regex}, but got {tag_override}. Proceeding with regular build."
+                )
+            else:
+                repo_override, t_override = tag_override.split(":")
+                with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file_handle:
+                    source_uri = (
+                        f"{image_repo_uri.replace('pr-', f'{repo_override}-')}:{t_override}"
+                    )
+                    temp_file_handle.write(
+                        f"FROM {source_uri}\nLABEL dlc.dev.source_uri={source_uri}"
+                    )
+                    dockerfile = temp_file_handle.name
+                    target = None
+                FORMATTER.print(f"USING TAG OVERRIDE {source_uri}")
+
         ARTIFACTS.update(
             {
                 "dockerfile": {
-                    "source": image_config["docker_file"],
+                    "source": dockerfile,
                     "target": "Dockerfile",
                 }
             }
@@ -308,14 +330,14 @@ def image_builder(buildspec, image_types=[], device_types=[]):
         # Create pre_push stage docker object
         pre_push_stage_image_object = DockerImage(
             info=info,
-            dockerfile=image_config["docker_file"],
+            dockerfile=dockerfile,
             repository=image_repo_uri,
             tag=append_tag(image_tag, "pre-push"),
             to_build=image_config["build"],
             stage=constants.PRE_PUSH_STAGE,
             context=context,
             additional_tags=additional_image_tags,
-            target=image_config.get("target"),
+            target=target,
         )
 
         ##### Create Common stage docker object #####

src/patch_helper.py

@@ -84,7 +84,7 @@ def get_impacted_os_packages(image_uri, python_version=None):
     ) = helper_function_for_leftover_vulnerabilities_from_enhanced_scanning(
         image_uri,
         python_version=python_version,
-        minimum_sev_threshold="LOW",
+        minimum_sev_threshold="UNDEFINED",
         allowlist_removal_enabled=False,
     )
     impacted_packages = set()

test/dlc_tests/sanity/quick_checks/test_buildspecs.py

@@ -37,25 +37,35 @@ def test_train_inference_buildspec():
 
                 # Don't look for framework buildspecs in the top level directory - these are not framework buildspecs
                 if os.path.split(buildspec_path)[0] != dlc_base_dir:
-                    _assert_single_image_type_buildspec(
+                    _assert_single_image_type_no_tag_override_buildspec(
                         buildspec_path, inference_pattern, training_pattern
                     )
 
 
-def _assert_single_image_type_buildspec(buildspec_path, inference_pattern, training_pattern):
+def _assert_single_image_type_no_tag_override_buildspec(
+    buildspec_path, inference_pattern, training_pattern
+):
     """
     Isolate condition for checking whether an buildspec is consistent with its image type (training or inference).
+    Also check to make sure we don't have any build tag overrides in there.
     Require that images are nested under training or inference, if not, raise error.
     """
+    build_tag_override_regex = re.compile(r"^\s*build_tag_override:\s\S*")
     if "training" in buildspec_path:
         with open(buildspec_path) as trn_buildspec_handle:
             for line in trn_buildspec_handle:
+                assert not build_tag_override_regex.search(
+                    line
+                ), f"Found valid build_tag_overrides in {buildspec_path} - please remove before merge."
                 assert not inference_pattern.search(
                     line
                 ), f"Found inference reference in training buildspec {buildspec_path}. Please check the file and remove them."
     elif "inference" in buildspec_path:
         with open(buildspec_path) as inf_buildspec_handle:
             for line in inf_buildspec_handle:
+                assert not build_tag_override_regex.search(
+                    line
+                ), f"Found valid build_tag_overrides in {buildspec_path} - please remove before merge."
                 assert not training_pattern.search(
                     line
                 ), f"Found training reference in inference buildspec {buildspec_path}. Please check the file and remove them."

test/dlc_tests/sanity/test_ecr_scan.py

@@ -221,10 +221,18 @@ def helper_function_for_leftover_vulnerabilities_from_enhanced_scanning(
     remaining_vulnerabilities = ecr_image_vulnerability_list - image_scan_allowlist
     LOGGER.info(f"ECR Enhanced Scanning test completed for image: {image}")
 
-    if remove_non_patchable_vulns and remaining_vulnerabilities:
-        non_patchable_vulnerabilities = extract_non_patchable_vulnerabilities(
-            remaining_vulnerabilities, ecr_enhanced_repo_uri
+    if remove_non_patchable_vulns:
+        non_patchable_vulnerabilities = ECREnhancedScanVulnerabilityList(
+            minimum_severity=CVESeverity[minimum_sev_threshold]
         )
+
+        ## non_patchable_vulnerabilities is a subset of remaining_vulnerabilities that cannot be patched.
+        ## Thus, if remaining_vulnerabilities exists, we need to find the non_patchable_vulnerabilities from it.
+        if remaining_vulnerabilities:
+            non_patchable_vulnerabilities = extract_non_patchable_vulnerabilities(
+                remaining_vulnerabilities, ecr_enhanced_repo_uri
+            )
+
         future_allowlist = generate_future_allowlist(
             ecr_image_vulnerability_list=ecr_image_vulnerability_list,
             image_scan_allowlist=image_scan_allowlist,
@@ -254,7 +262,10 @@ def helper_function_for_leftover_vulnerabilities_from_enhanced_scanning(
             s3_filepath=future_allowlist_upload_path,
             tag_set=upload_tag_set,
         )
-        remaining_vulnerabilities = remaining_vulnerabilities - non_patchable_vulnerabilities
+
+        if remaining_vulnerabilities:
+            remaining_vulnerabilities = remaining_vulnerabilities - non_patchable_vulnerabilities
+
         LOGGER.info(
             f"[FutureAllowlist][image_uri:{ecr_enhanced_repo_uri}] {json.dumps(future_allowlist.vulnerability_list, cls= test_utils.EnhancedJSONEncoder)}"
         )

test/test_utils/security.py

@@ -382,7 +382,7 @@ def __sub__(self, other):
         if not self.vulnerability_list:
             return None
         if not other or not other.vulnerability_list:
-            return self
+            return copy.deepcopy(self)
 
         missing_vulnerabilities = [
             vulnerability
@@ -407,7 +407,9 @@ def __add__(self, other):
         :return: Union of vulnerabilites exisiting in self and other
         """
         flattened_vulnerability_list_self = self.get_flattened_vulnerability_list()
-        flattened_vulnerability_list_other = other.get_flattened_vulnerability_list()
+        flattened_vulnerability_list_other = (
+            other.get_flattened_vulnerability_list() if other else []
+        )
         all_vulnerabilities = flattened_vulnerability_list_self + flattened_vulnerability_list_other
         if not all_vulnerabilities:
             return None