feat: Replace image rescan command's --local flag with the local type ContainerRegistry record

changes/2665.feature.md

@@ -0,0 +1 @@
+Replace rescan command's `--local` flag with local container registry record.

fixtures/manager/example-container-registries-local.json

@@ -0,0 +1,11 @@
+{
+    "container_registries": [
+        {
+            "id": "fe878f09-06cc-4b91-9242-4c71015cce05",
+            "registry_name": "local",
+            "url": "http://localhost",
+            "type": "local",
+            "project": "stable"
+        }
+    ]
+}

src/ai/backend/agent/docker/agent.py

@@ -823,8 +823,14 @@ async def start_container(
         container_log_size = self.local_config["agent"]["container-logs"]["max-length"]
         container_log_file_count = 5
         container_log_file_size = BinarySize(container_log_size // container_log_file_count)
+
+        if self.image_ref.is_local:
+            image = self.image_ref.short
+        else:
+            image = self.image_ref.canonical
+
         container_config: MutableMapping[str, Any] = {
-            "Image": self.image_ref.canonical,
+            "Image": image,
             "Tty": True,
             "OpenStdin": True,
             "Privileged": False,

src/ai/backend/manager/api/session.py

@@ -1307,7 +1307,6 @@ async def _commit_and_upload(reporter: ProgressReporter) -> None:
             await rescan_images(
                 root_ctx.db,
                 new_image_ref.canonical,
-                local=new_image_ref.is_local,
             )
             await reporter.update(increment=1, message="Completed")
         except BackendError:

src/ai/backend/manager/cli/etcd.py

@@ -261,7 +261,7 @@ def rescan_images(cli_ctx: CLIContext, registry: str) -> None:
     Pass the name (usually hostname or "lablup") of the Docker registry configured as REGISTRY.
     """
     log.warning("etcd rescan-images command is deprecated, use image rescan instead")
-    asyncio.run(rescan_images_impl(cli_ctx, registry, False))
+    asyncio.run(rescan_images_impl(cli_ctx, registry))
 
 
 @cli.command()

src/ai/backend/manager/cli/image.py

@@ -82,20 +82,14 @@ def set_resource_limit(
 
 @cli.command()
 @click.argument("registry_or_image", required=False, default="")
-@click.option(
-    "--local",
-    is_flag=True,
-    default=False,
-    help="Scan the local Docker daemon instead of a registry",
-)
 @click.pass_obj
-def rescan(cli_ctx, registry_or_image: str, local: bool) -> None:
+def rescan(cli_ctx, registry_or_image: str) -> None:
     """
     Update the kernel image metadata from all configured docker registries.
 
     Pass the name (usually hostname or "lablup") of the Docker registry configured as REGISTRY.
     """
-    asyncio.run(rescan_images_impl(cli_ctx, registry_or_image, local))
+    asyncio.run(rescan_images_impl(cli_ctx, registry_or_image))
 
 
 @cli.command()

src/ai/backend/manager/cli/image_impl.py

@@ -151,14 +151,14 @@ async def set_image_resource_limit(
             log.exception("An error occurred.")
 
 
-async def rescan_images(cli_ctx: CLIContext, registry_or_image: str, local: bool) -> None:
-    if not registry_or_image and not local:
+async def rescan_images(cli_ctx: CLIContext, registry_or_image: str) -> None:
+    if not registry_or_image:
         raise click.BadArgumentUsage("Please specify a valid registry or full image name.")
     async with (
         connect_database(cli_ctx.local_config) as db,
     ):
         try:
-            await rescan_images_func(db, registry_or_image, local=local)
+            await rescan_images_func(db, registry_or_image)
         except Exception:
             log.exception("An error occurred.")
 

src/ai/backend/manager/container_registry/base.py

@@ -516,14 +516,12 @@ async def _read_manifest(
                 except ValueError as e:
                     skip_reason = str(e)
                     continue
-                if self.registry_name == "local":
-                    if image.partition("/")[1] == "":
-                        image = "library/" + image
-                    update_key = ImageIdentifier(f"{image}:{tag}", architecture)
-                else:
-                    update_key = ImageIdentifier(
-                        f"{self.registry_name}/{image}:{tag}", architecture
-                    )
+
+                update_key = ImageIdentifier(
+                    f"{self.registry_name}/{image}:{tag}",
+                    architecture,
+                )
+
                 updates = {
                     "config_digest": manifest["digest"],
                     "size_bytes": manifest["size"],

src/ai/backend/manager/models/image.py

@@ -131,51 +131,41 @@ async def rescan_images(
     db: ExtendedAsyncSAEngine,
     registry_or_image: str | None = None,
     *,
-    local: bool | None = False,
     reporter: ProgressReporter | None = None,
 ) -> None:
-    if local:
-        registries = {
-            "local": ContainerRegistryRow(
-                registry_name="local",
-                url="http://localhost",
-                type=ContainerRegistryType.LOCAL,
-            )
-        }
-    else:
-        async with db.begin_readonly_session() as session:
-            result = await session.execute(sa.select(ContainerRegistryRow))
-            latest_registry_config = cast(
-                dict[str, ContainerRegistryRow],
-                {row.registry_name: row for row in result.scalars().all()},
-            )
+    async with db.begin_readonly_session() as session:
+        result = await session.execute(sa.select(ContainerRegistryRow))
+        latest_registry_config = cast(
+            dict[str, ContainerRegistryRow],
+            {row.registry_name: row for row in result.scalars().all()},
+        )
 
-        # TODO: delete images from registries removed from the previous config?
-        if registry_or_image is None:
-            # scan all configured registries
-            registries = latest_registry_config
+    # TODO: delete images from registries removed from the previous config?
+    if registry_or_image is None:
+        # scan all configured registries
+        registries = latest_registry_config
+    else:
+        # find if it's a full image ref of one of configured registries
+        for registry_name, registry_info in latest_registry_config.items():
+            if registry_or_image.startswith(registry_name + "/"):
+                repo_with_tag = registry_or_image.removeprefix(registry_name + "/")
+                log.debug(
+                    "running a per-image metadata scan: {}, {}",
+                    registry_name,
+                    repo_with_tag,
+                )
+                scanner_cls = get_container_registry_cls(registry_info)
+                scanner = scanner_cls(db, registry_name, registry_info)
+                await scanner.scan_single_ref(repo_with_tag)
+                return
         else:
-            # find if it's a full image ref of one of configured registries
-            for registry_name, registry_info in latest_registry_config.items():
-                if registry_or_image.startswith(registry_name + "/"):
-                    repo_with_tag = registry_or_image.removeprefix(registry_name + "/")
-                    log.debug(
-                        "running a per-image metadata scan: {}, {}",
-                        registry_name,
-                        repo_with_tag,
-                    )
-                    scanner_cls = get_container_registry_cls(registry_info)
-                    scanner = scanner_cls(db, registry_name, registry_info)
-                    await scanner.scan_single_ref(repo_with_tag)
-                    return
-            else:
-                # treat it as a normal registry name
-                registry = registry_or_image
-                try:
-                    registries = {registry: latest_registry_config[registry]}
-                    log.debug("running a per-registry metadata scan")
-                except KeyError:
-                    raise RuntimeError("It is an unknown registry.", registry)
+            # treat it as a normal registry name
+            registry = registry_or_image
+            try:
+                registries = {registry: latest_registry_config[registry]}
+                log.debug("running a per-registry metadata scan")
+            except KeyError:
+                raise RuntimeError("It is an unknown registry.", registry)
     async with aiotools.TaskGroup() as tg:
         for registry_name, registry_info in registries.items():
             log.info('Scanning kernel images from the registry "{0}"', registry_name)