Skip to content

Commit

Permalink
rename method and remove unused argument
Browse files Browse the repository at this point in the history
  • Loading branch information
kyujin-cho committed Aug 1, 2024
1 parent f1a4645 commit 52810f8
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 42 deletions.
44 changes: 18 additions & 26 deletions src/ai/backend/manager/container_registry/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
28 changes: 12 additions & 16 deletions src/ai/backend/manager/container_registry/harbor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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}")

Expand All @@ -314,14 +310,14 @@ 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)
if not rqst_args.get("headers"):
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"
Expand All @@ -339,7 +335,7 @@ async def _process_oci_index(
rqst_args,
image,
digest=digest,
tag=tag,
tag=tag_name,
architecture=architecture,
)
)
Expand All @@ -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)
Expand All @@ -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"
Expand All @@ -377,7 +373,7 @@ async def _process_docker_v2_multiplatform_image(
rqst_args,
image,
digest=digest,
tag=tag,
tag=tag_name,
architecture=architecture,
)
)
Expand All @@ -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)
Expand All @@ -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,
)
)

Expand Down

0 comments on commit 52810f8

Please sign in to comment.