diff --git a/src/ai/backend/manager/container_registry/base.py b/src/ai/backend/manager/container_registry/base.py index 7fa8b921eba..78b9aea9ab1 100644 --- a/src/ai/backend/manager/container_registry/base.py +++ b/src/ai/backend/manager/container_registry/base.py @@ -229,26 +229,21 @@ async def _scan_tag( resp.raise_for_status() resp_json = await resp.json() - async with aiotools.TaskGroup() as tg: - match content_type: - case self.MEDIA_TYPE_DOCKER_MANIFEST: - await self._process_docker_v2_image( - tg, sess, rqst_args, image, tag, resp_json - ) - case self.MEDIA_TYPE_DOCKER_MANIFEST_LIST: - await self._process_docker_v2_multiplatform_image( - tg, sess, rqst_args, image, tag, resp_json - ) - case self.MEDIA_TYPE_OCI_INDEX: - await self._process_oci_index( - tg, sess, rqst_args, image, tag, resp_json - ) - case _: - log.warn("Unknown content type: {}", content_type) - raise RuntimeError( - "The registry does not support the standard way of " - "listing multiarch images." - ) + match content_type: + case self.MEDIA_TYPE_DOCKER_MANIFEST: + await self._read_docker_v2_image(sess, rqst_args, image, tag, resp_json) + case self.MEDIA_TYPE_DOCKER_MANIFEST_LIST: + await self._read_docker_v2_multiplatform_image( + sess, rqst_args, image, tag, resp_json + ) + case self.MEDIA_TYPE_OCI_INDEX: + await self._read_oci_index(sess, rqst_args, image, tag, resp_json) + case _: + log.warn("Unknown content type: {}", content_type) + raise RuntimeError( + "The registry does not support the standard way of " + "listing multiarch images." + ) async def _read_manifest_list( self, @@ -321,9 +316,8 @@ async def _preprocess_manifest( "digest": config_digest, } - async def _process_oci_index( + async def _read_oci_index( self, - tg: aiotools.TaskGroup, sess: aiohttp.ClientSession, rqst_args: Mapping[str, Any], image: str, @@ -339,9 +333,8 @@ async def _process_oci_index( await self._read_manifest_list(sess, manifest_list, rqst_args, image, tag) - async def _process_docker_v2_multiplatform_image( + async def _read_docker_v2_multiplatform_image( self, - tg: aiotools.TaskGroup, sess: aiohttp.ClientSession, rqst_args: Mapping[str, Any], image: str, @@ -359,9 +352,8 @@ async def _process_docker_v2_multiplatform_image( tag, ) - async def _process_docker_v2_image( + async def _read_docker_v2_image( self, - tg: aiotools.TaskGroup, sess: aiohttp.ClientSession, rqst_args: Mapping[str, Any], image: str, diff --git a/src/ai/backend/manager/container_registry/harbor.py b/src/ai/backend/manager/container_registry/harbor.py index fcdb7afc0f2..2a4b09c7f51 100644 --- a/src/ai/backend/manager/container_registry/harbor.py +++ b/src/ai/backend/manager/container_registry/harbor.py @@ -243,15 +243,15 @@ async def _scan_image( match image_info["manifest_media_type"]: case self.MEDIA_TYPE_OCI_INDEX: await self._process_oci_index( - tg, sess, rqst_args, image, tag, image_info + tg, sess, rqst_args, image, image_info ) case self.MEDIA_TYPE_DOCKER_MANIFEST_LIST: await self._process_docker_v2_multiplatform_image( - tg, sess, rqst_args, image, tag, image_info + tg, sess, rqst_args, image, image_info ) case self.MEDIA_TYPE_DOCKER_MANIFEST: await self._process_docker_v2_image( - tg, sess, rqst_args, image, tag, image_info + tg, sess, rqst_args, image, image_info ) case _ as media_type: raise RuntimeError( @@ -292,19 +292,15 @@ async def _scan_tag( resp.raise_for_status() resp_json = await resp.json() async with aiotools.TaskGroup() as tg: - tag = resp_json["tags"][0]["name"] - match resp_json["manifest_media_type"]: case self.MEDIA_TYPE_OCI_INDEX: - await self._process_oci_index(tg, sess, rqst_args, image, tag, resp_json) + await self._process_oci_index(tg, sess, rqst_args, image, resp_json) case self.MEDIA_TYPE_DOCKER_MANIFEST_LIST: await self._process_docker_v2_multiplatform_image( - tg, sess, rqst_args, image, tag, resp_json + tg, sess, rqst_args, image, resp_json ) case self.MEDIA_TYPE_DOCKER_MANIFEST: - await self._process_docker_v2_image( - tg, sess, rqst_args, image, tag, resp_json - ) + await self._process_docker_v2_image(tg, sess, rqst_args, image, resp_json) case _ as media_type: raise RuntimeError(f"Unsupported artifact media-type: {media_type}") @@ -314,7 +310,6 @@ async def _process_oci_index( sess: aiohttp.ClientSession, _rqst_args: Mapping[str, Any], image: str, - tag: str, image_info: Mapping[str, Any], ) -> None: rqst_args = dict(_rqst_args) @@ -322,6 +317,7 @@ async def _process_oci_index( rqst_args["headers"] = {} rqst_args["headers"].update({"Accept": "application/vnd.oci.image.manifest.v1+json"}) digests: list[tuple[str, str]] = [] + tag_name = image_info["tags"][0]["name"] for reference in image_info["references"]: if ( reference["platform"]["os"] == "unknown" @@ -339,7 +335,7 @@ async def _process_oci_index( rqst_args, image, digest=digest, - tag=tag, + tag=tag_name, architecture=architecture, ) ) @@ -350,7 +346,6 @@ async def _process_docker_v2_multiplatform_image( sess: aiohttp.ClientSession, _rqst_args: Mapping[str, Any], image: str, - tag: str, image_info: Mapping[str, Any], ) -> None: rqst_args = dict(_rqst_args) @@ -360,6 +355,7 @@ async def _process_docker_v2_multiplatform_image( "Accept": "application/vnd.docker.distribution.manifest.v2+json" }) digests: list[tuple[str, str]] = [] + tag_name = image_info["tags"][0]["name"] for reference in image_info["references"]: if ( reference["platform"]["os"] == "unknown" @@ -377,7 +373,7 @@ async def _process_docker_v2_multiplatform_image( rqst_args, image, digest=digest, - tag=tag, + tag=tag_name, architecture=architecture, ) ) @@ -388,7 +384,6 @@ async def _process_docker_v2_image( sess: aiohttp.ClientSession, _rqst_args: Mapping[str, Any], image: str, - tag: str, image_info: Mapping[str, Any], ) -> None: rqst_args = dict(_rqst_args) @@ -399,13 +394,14 @@ async def _process_docker_v2_image( }) if (reporter := progress_reporter.get()) is not None: reporter.total_progress += 1 + tag_name = image_info["tags"][0]["name"] async with aiotools.TaskGroup() as tg: tg.create_task( self._harbor_scan_tag_single_arch( sess, rqst_args, image, - tag, + tag=tag_name, ) )