Add Deep Canary For Public Registry (#5045)

* add public registry deep canary

* reformat

* fix typo

* handle lookup function issue

* remove unuseful assert

---------

Co-authored-by: Yadan Wei <yadanwei@amazon.com>

Author: Yadan-Wei

test/dlc_tests/conftest.py

@@ -1799,6 +1799,13 @@ def lookup_condition(lookup, image):
 def pytest_generate_tests(metafunc):
     images = metafunc.config.getoption("--images")
 
+    # Check for public registry canary first
+    if os.getenv("IS_PUBLIC_REGISTRY_CANARY", "false").lower() == "true":
+        # Only handle framework agnostic tests for public registry
+        if "image" in metafunc.fixturenames:
+            metafunc.parametrize("image", images)
+        return
+
     # Parametrize framework specific tests
     for fixture in FRAMEWORK_FIXTURES:
         if fixture in metafunc.fixturenames:

test/dlc_tests/sanity/test_canary_integration.py

@@ -25,13 +25,25 @@
 @pytest.mark.integration("deep_canary")
 def test_deep_canary_integration(image, region):
     ctx = Context()
-    image_account_id = get_account_id_from_image_uri(image)
-    image_region = get_region_from_image_uri(image)
-    assert image_region == region, f"Problem: Test region {region} != canary region {image_region}"
+    if os.getenv("IS_PUBLIC_REGISTRY_CANARY", "false").lower() == "true":
+        # If the image is a public registry canary, we don't need to login to ECR
+        LOGGER.info(f"Deep Canary pull test for public registry canary {image}")
+        try:
+            ctx.run(f"docker pull {image}", hide="out")
+            LOGGER.info(f"Deep Canary pull test succeeded for {image}")
+        finally:
+            ctx.run(f"docker rmi {image}", warn=True, hide="out")
+        return
+    else:
+        image_account_id = get_account_id_from_image_uri(image)
+        image_region = get_region_from_image_uri(image)
+        assert (
+            image_region == region
+        ), f"Problem: Test region {region} != canary region {image_region}"
 
-    try:
-        login_to_ecr_registry(ctx, image_account_id, region)
-        ctx.run(f"docker pull {image}", hide="out")
-        LOGGER.info(f"Deep Canary pull test succeeded for {image}")
-    finally:
-        ctx.run(f"docker rmi {image}", warn=True, hide="out")
+        try:
+            login_to_ecr_registry(ctx, image_account_id, region)
+            ctx.run(f"docker pull {image}", hide="out")
+            LOGGER.info(f"Deep Canary pull test succeeded for {image}")
+        finally:
+            ctx.run(f"docker rmi {image}", warn=True, hide="out")

test/test_utils/__init__.py

@@ -239,6 +239,7 @@ def get_ami_id_ssm(region_name, parameter_path):
 SAGEMAKER_REMOTE_TEST_TYPE = "sagemaker"
 
 PUBLIC_DLC_REGISTRY = "763104351884"
+DLC_PUBLIC_REGISTRY_ALIAS = "public.ecr.aws/deep-learning-containers"
 
 SAGEMAKER_EXECUTION_REGIONS = ["us-west-2", "us-east-1", "eu-west-1"]
 # Before SM GA with Trn1, they support launch of ml.trn1 instance only in us-east-1. After SM GA this can be removed
@@ -1318,6 +1319,7 @@ def get_dlc_images():
             canary_arch_type=get_test_job_arch_type(),
             canary_region=os.getenv("AWS_REGION"),
             canary_region_prod_account=os.getenv("REGIONAL_PROD_ACCOUNT", PUBLIC_DLC_REGISTRY),
+            is_public_registry=os.getenv("IS_PUBLIC_REGISTRY_CANARY", "false").lower() == "true",
         )
         return " ".join(deep_canary_images)
     elif is_pr_context() or is_empty_build_context():
@@ -1345,6 +1347,7 @@ def get_deep_canary_images(
     canary_arch_type,
     canary_region,
     canary_region_prod_account,
+    is_public_registry=False,
 ):
     """
     For an input combination of canary job specs, find a matching list of image uris to be tested
@@ -1382,9 +1385,17 @@ def get_deep_canary_images(
             and canary_image_type == image_type
             and canary_arch_type == image_arch_type
         ):
-            regionalized_image_uri = image_uri.replace(image_region, canary_region).replace(
-                image_account_id, canary_region_prod_account
-            )
+            if is_public_registry:
+                # For public registry, we use the public account ID
+                image_repository = image_uri.split("/")[-1].split(":")[0]
+                image_tag = image_uri.split(":")[-1]
+                regionalized_image_uri = (
+                    f"{DLC_PUBLIC_REGISTRY_ALIAS}/{image_repository}:{image_tag}"
+                )
+            else:
+                regionalized_image_uri = image_uri.replace(image_region, canary_region).replace(
+                    image_account_id, canary_region_prod_account
+                )
             matching_images.append(regionalized_image_uri)
     return matching_images