Add support for public release only (#5408)

Yadan-Wei · Yadan Wei · web-flow · commit cbf299bfd42c · 2025-11-03T11:47:25.000-08:00
* add support for public release only

* format and add 2.4 for test

* update ecr public describe images

* change saved release information

* revert temp change and format

* remove blank line

---------

Co-authored-by: Yadan Wei &lt;yadanwei@amazon.com&gt;
diff --git a/generate_dlc_image_release_information.py b/generate_dlc_image_release_information.py
@@ -69,18 +69,15 @@ def parse_args():
         github_publishing_metadata = json.loads(f.read())
 
     dlc_account_id = github_publishing_metadata.get("target_account_id_classic")
+    dlc_public_account_id = github_publishing_metadata.get("target_account_id_public")
     dlc_tag = github_publishing_metadata.get("tag_with_dlc_version")
+    is_private_release = github_publishing_metadata.get("private_registry_release", True)
     dlc_soci_tag = github_publishing_metadata.get("tag_with_dlc_version_soci", None)
     dlc_repository = github_publishing_metadata.get("target_ecr_repository")
     dlc_release_successful = github_publishing_metadata.get("release_successful")
     dlc_region = os.getenv("REGION")
 
     dlc_public_registry = github_publishing_metadata.get("target_ecr_public_registry")
-    public_registry_image_uri_with_dlc_version = None
-    if dlc_public_registry is not None:
-        public_registry_image_uri_with_dlc_version = (
-            f"{dlc_public_registry}/{dlc_repository}:{dlc_tag}"
-        )
 
     if dlc_release_successful != "1":
         LOGGER.error(
@@ -89,7 +86,14 @@ def parse_args():
         sys.exit(0)
 
     dlc_release_information = DLCReleaseInformation(
-        dlc_account_id, dlc_region, dlc_repository, dlc_tag, dlc_soci_tag=dlc_soci_tag
+        dlc_account_id,
+        dlc_public_account_id,
+        dlc_region,
+        dlc_repository,
+        dlc_tag,
+        dlc_soci_tag=dlc_soci_tag,
+        is_private_release=is_private_release,
+        public_registry=dlc_public_registry,
     )
 
     # bom objects below are used to create .tar.gz file to be uploaded as an asset, 'imp' objects are used as release information
@@ -102,8 +106,10 @@ def parse_args():
         "image_digest": dlc_release_information.image_digest,
         "soci_image_digest": dlc_release_information.soci_image_digest,
         "image_uri_with_dlc_version": dlc_release_information.image,
+        "is_private_registry_release": is_private_release,
         "soci_image_uri_with_dlc_version": dlc_release_information.soci_image,
-        "public_registry_image_uri_with_dlc_version": public_registry_image_uri_with_dlc_version,
+        "public_image_uri_with_dlc_version": dlc_release_information.public_image,
+        "public_soci_image_uri_with_dlc_version": dlc_release_information.public_soci_image,
         "image_tags": dlc_release_information.image_tags,
         "soci_image_tags": dlc_release_information.soci_image_tags,
         "dlc_release_successful": dlc_release_successful,
@@ -161,5 +167,10 @@ def parse_args():
     os.remove(tarfile_name)
     shutil.rmtree(directory)
 
-    LOGGER.info(f"Release Information collected for image: {dlc_release_information.image}")
+    if is_private_release:
+        LOGGER.info(f"Release Information collected for image: {dlc_release_information.image}")
+    else:
+        LOGGER.info(
+            f"Release Information collected for public image: {dlc_release_information.public_image}"
+        )
     LOGGER.info(f"Release information and BOM uploaded to: {s3BucketURI}")
diff --git a/release/dlc_release_information.py b/release/dlc_release_information.py
@@ -16,18 +16,37 @@
 
 
 class DLCReleaseInformation:
-    def __init__(self, dlc_account_id, dlc_region, dlc_repository, dlc_tag, dlc_soci_tag=None):
-        if not all([dlc_account_id, dlc_tag, dlc_repository, dlc_region]):
+    def __init__(
+        self,
+        dlc_account_id,
+        dlc_public_account_id,
+        dlc_region,
+        dlc_repository,
+        dlc_tag,
+        dlc_soci_tag=None,
+        is_private_release=True,
+        public_registry=None,
+    ):
+        if not all([dlc_tag, dlc_repository, dlc_region]):
             raise ValueError(
-                "One or multiple environment variables TARGET_ACCOUNT_ID_CLASSIC, TAG_WITH_DLC_VERSION, "
-                "TARGET_ECR_REPOSITORY, REGION  not set. This environment variable is expected to be set by the promoter stage."
+                "One or multiple environment variables TAG_WITH_DLC_VERSION, "
+                "TARGET_ECR_REPOSITORY, REGION not set. This environment variable is expected to be set by the promoter stage."
             )
 
+        if is_private_release and not dlc_account_id:
+            raise ValueError("dlc_account_id is required for private releases")
+
+        if not is_private_release and not dlc_public_account_id:
+            raise ValueError("dlc_public_account_id is required for public releases only")
+
         self.dlc_account_id = dlc_account_id
+        self.dlc_public_account_id = dlc_public_account_id
         self.dlc_region = dlc_region
         self.dlc_repository = dlc_repository
         self.dlc_tag = dlc_tag
         self.dlc_soci_tag = dlc_soci_tag
+        self.is_private_release = is_private_release
+        self.public_registry = public_registry
 
         self.container_name = self.run_container()
 
@@ -43,6 +62,9 @@ def __init__(self, dlc_account_id, dlc_region, dlc_repository, dlc_tag, dlc_soci
     def get_boto3_ecr_client(self):
         return boto3.Session(region_name=self.dlc_region).client("ecr")
 
+    def get_boto3_ecr_public_client(self):
+        return boto3.Session(region_name="us-east-1").client("ecr-public")
+
     def run_container(self):
         """
         Quickly run a container and assign it a name, so that different commands may be run on it
@@ -53,8 +75,12 @@ def run_container(self):
 
         run(f"docker rm -f {container_name}", warn=True, hide=True)
 
+        if self.is_private_release:
+            image_uri = self.image
+        else:
+            image_uri = self.public_image
         run(
-            f"docker run -id --privileged --name {container_name} --entrypoint='/bin/bash' {self.image}",
+            f"docker run -id --privileged --name {container_name} --entrypoint='/bin/bash' {image_uri}",
             hide=True,
         )
 
@@ -75,23 +101,42 @@ def get_container_command_output(self, command):
     def get_image_details_from_ecr(self, tag):
         if tag is None:
             return None
-        _ecr = self.get_boto3_ecr_client()
 
-        try:
-            response = _ecr.describe_images(
-                registryId=self.dlc_account_id,
-                repositoryName=self.dlc_repository,
-                imageIds=[{"imageTag": tag}],
-            )
-        except ClientError as e:
-            LOGGER.error("ClientError when performing ECR operation. Exception: {}".format(e))
-
-        return response["imageDetails"][0]
+        if self.is_private_release:
+            _ecr = self.get_boto3_ecr_client()
+            try:
+                response = _ecr.describe_images(
+                    registryId=self.dlc_account_id,
+                    repositoryName=self.dlc_repository,
+                    imageIds=[{"imageTag": tag}],
+                )
+            except ClientError as e:
+                LOGGER.error("ClientError when performing ECR operation. Exception: {}".format(e))
+            return response["imageDetails"][0]
+        else:
+            _ecr_public = self.get_boto3_ecr_public_client()
+            try:
+                response = _ecr_public.describe_images(
+                    registryId=self.dlc_public_account_id,
+                    repositoryName=self.dlc_repository,
+                    imageIds=[{"imageTag": tag}],
+                )
+            except ClientError as e:
+                LOGGER.error(
+                    "ClientError when performing ECR Public operation. Exception: {}".format(e)
+                )
+            return response["imageDetails"][0]
 
     @property
     def image(self):
         return f"{self.dlc_account_id}.dkr.ecr.{self.dlc_region}.amazonaws.com/{self.dlc_repository}:{self.dlc_tag}"
 
+    @property
+    def public_image(self):
+        if self.public_registry is None:
+            return None
+        return f"{self.public_registry}/{self.dlc_repository}:{self.dlc_tag}"
+
     @property
     def image_tags(self):
         return self._image_details["imageTags"]
@@ -106,6 +151,12 @@ def soci_image(self):
             return None
         return f"{self.dlc_account_id}.dkr.ecr.{self.dlc_region}.amazonaws.com/{self.dlc_repository}:{self.dlc_soci_tag}"
 
+    @property
+    def public_soci_image(self):
+        if self.dlc_soci_tag is None or self.public_registry is None:
+            return None
+        return f"{self.public_registry}/{self.dlc_repository}:{self.dlc_soci_tag}"
+
     @property
     def soci_image_tags(self):
         if self.dlc_soci_tag is None:
diff --git a/release_images_training.yml b/release_images_training.yml
@@ -103,4 +103,4 @@ release_images:
       neuron_sdk_version: "sdk2.24.1"
       example: False
       disable_sm_tag: True
-      force_release: False
+      force_release: False
