diff --git a/changes/2665.feature.md b/changes/2665.feature.md new file mode 100644 index 00000000000..73d78e3ee9d --- /dev/null +++ b/changes/2665.feature.md @@ -0,0 +1 @@ +Replace rescan command's `--local` flag with local container registry record. diff --git a/fixtures/manager/example-container-registries-local.json b/fixtures/manager/example-container-registries-local.json new file mode 100644 index 00000000000..04c0cd894fd --- /dev/null +++ b/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" + } + ] +} diff --git a/src/ai/backend/agent/docker/agent.py b/src/ai/backend/agent/docker/agent.py index 4bbf5567889..240feba7306 100644 --- a/src/ai/backend/agent/docker/agent.py +++ b/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, diff --git a/src/ai/backend/manager/api/session.py b/src/ai/backend/manager/api/session.py index adac34c5bb0..f37d98ee941 100644 --- a/src/ai/backend/manager/api/session.py +++ b/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: diff --git a/src/ai/backend/manager/cli/etcd.py b/src/ai/backend/manager/cli/etcd.py index 3b7bf304d78..c620bb435d9 100644 --- a/src/ai/backend/manager/cli/etcd.py +++ b/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() diff --git a/src/ai/backend/manager/cli/image.py b/src/ai/backend/manager/cli/image.py index fe87c3bdbf9..27a1b3e8282 100644 --- a/src/ai/backend/manager/cli/image.py +++ b/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() diff --git a/src/ai/backend/manager/cli/image_impl.py b/src/ai/backend/manager/cli/image_impl.py index 37a18407330..7268d53f71f 100644 --- a/src/ai/backend/manager/cli/image_impl.py +++ b/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.") diff --git a/src/ai/backend/manager/container_registry/base.py b/src/ai/backend/manager/container_registry/base.py index 4292cb9ac12..9698ac21d67 100644 --- a/src/ai/backend/manager/container_registry/base.py +++ b/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"], diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index c1463bee445..e2b2463debf 100644 --- a/src/ai/backend/manager/models/image.py +++ b/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)