diff --git a/changes/12439.enhance.md b/changes/12439.enhance.md new file mode 100644 index 00000000000..9a0a6caf296 --- /dev/null +++ b/changes/12439.enhance.md @@ -0,0 +1 @@ +Move container-registry image scanning out of `models/image/row.py` into the image repository layer (scan/rescan on the DB source, the registry untag on the repository) so the models layer no longer imports the `container_registry` package. diff --git a/src/ai/backend/manager/cli/image_impl.py b/src/ai/backend/manager/cli/image_impl.py index 7235d39473d..5654724322e 100644 --- a/src/ai/backend/manager/cli/image_impl.py +++ b/src/ai/backend/manager/cli/image_impl.py @@ -10,14 +10,17 @@ from tabulate import tabulate from ai.backend.common.arch import CURRENT_ARCH +from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.docker import validate_image_labels from ai.backend.common.exception import UnknownImageReference from ai.backend.common.types import ImageAlias, ImageID from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.container_registry.harbor import HarborRegistry_v2 from ai.backend.manager.data.image.types import ImageStatus +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.image import ImageAliasRow, ImageIdentifier, ImageRow -from ai.backend.manager.models.image import rescan_images as rescan_images_func from ai.backend.manager.models.utils import connect_database +from ai.backend.manager.repositories.image.db_source.db_source import ImageDBSource from .context import CLIContext, redis_ctx @@ -138,7 +141,15 @@ async def purge_image( await session.delete(image_row) if remove_from_registry: - await image_row.untag_image_from_registry(db=db, session=session) + registry_info = await session.get(ContainerRegistryRow, image_row.registry_id) + if registry_info is None: + raise click.ClickException(f"Registry not found for image {image_row.name}") + if registry_info.type != ContainerRegistryType.HARBOR2: + raise click.ClickException( + "Untagging from the registry is only supported for Harbor v2 registries" + ) + scanner = HarborRegistry_v2(db, image_row.image_ref.registry, registry_info) + await scanner.untag(image_row.image_ref) except UnknownImageReference: log.exception("Image not found.") @@ -191,7 +202,7 @@ async def rescan_images( connect_database(bootstrap_config.db) as db, ): try: - result = await rescan_images_func(db, registry_or_image, project) + result = await ImageDBSource(db).rescan_images(registry_or_image, project) for error in result.errors: log.error(f"Failed to scan registries: {error}") except Exception as e: diff --git a/src/ai/backend/manager/models/image/__init__.py b/src/ai/backend/manager/models/image/__init__.py index bbaaf446d28..482803ae7ae 100644 --- a/src/ai/backend/manager/models/image/__init__.py +++ b/src/ai/backend/manager/models/image/__init__.py @@ -8,8 +8,6 @@ PublicImageLoadFilter, Resources, bulk_get_image_configs, - rescan_images, - scan_single_image, ) from .row import ( get_permission_ctx as get_permission_ctx, @@ -26,6 +24,4 @@ "PublicImageLoadFilter", "Resources", "bulk_get_image_configs", - "rescan_images", - "scan_single_image", ) diff --git a/src/ai/backend/manager/models/image/row.py b/src/ai/backend/manager/models/image/row.py index dfada730712..807f6bcdb52 100644 --- a/src/ai/backend/manager/models/image/row.py +++ b/src/ai/backend/manager/models/image/row.py @@ -31,7 +31,6 @@ ) from sqlalchemy.sql.expression import true -from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.docker import ImageRef from ai.backend.common.exception import UnknownImageReference from ai.backend.common.types import ( @@ -48,7 +47,6 @@ from ai.backend.common.utils import join_non_empty from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.config.loader.legacy_etcd_loader import LegacyEtcdLoader -from ai.backend.manager.container_registry import get_container_registry_cls from ai.backend.manager.data.image.types import ( ImageAliasData, ImageData, @@ -60,7 +58,6 @@ ImageTagEntry, ImageType, KVPair, - RescanImagesResult, ResourceLimit, ) from ai.backend.manager.data.permission.types import EntityType @@ -91,10 +88,8 @@ AssociationScopesEntitiesRow, ) from ai.backend.manager.models.user import UserRole, UserRow -from ai.backend.manager.models.utils import ExtendedAsyncSAEngine if TYPE_CHECKING: - from ai.backend.common.bgtask.reporter import ProgressReporter from ai.backend.manager.models.container_registry import ContainerRegistryRow log = BraceStyleAdapter(logging.getLogger(__spec__.name)) @@ -106,7 +101,6 @@ "ImageLoadFilter", "ImageRow", "PublicImageLoadFilter", - "rescan_images", ) @@ -151,157 +145,6 @@ def _apply_loading_option( return query_stmt -async def load_configured_registries( - db: ExtendedAsyncSAEngine, - project: str | None, -) -> dict[str, ContainerRegistryRow]: - join = functools.partial(join_non_empty, sep="/") - - async with db.begin_readonly_session() as session: - result = await session.execute(sa.select(ContainerRegistryRow)) - if project: - registries = cast( - dict[str, ContainerRegistryRow], - { - join(row.registry_name, row.project): row - for row in result.scalars().all() - if row.project == project - }, - ) - else: - registries = { - join(row.registry_name, row.project): row for row in result.scalars().all() - } - - return registries - - -async def scan_registries( - db: ExtendedAsyncSAEngine, - registries: dict[str, ContainerRegistryRow], - reporter: ProgressReporter | None = None, -) -> RescanImagesResult: - """ - Performs an image rescan for all images in the registries. - """ - from ai.backend.manager.data.image.types import RescanImagesResult - - images, errors = [], [] - - for registry_key, registry_row in registries.items(): - registry_name = ImageRef.parse_image_str(registry_key, "*").registry - log.info('Scanning kernel images from the registry "{0}"', registry_name) - - scanner_cls = get_container_registry_cls(registry_row) - scanner = scanner_cls(db, registry_name, registry_row) - - try: - scan_result = await scanner.rescan_single_registry(reporter) - images.extend(scan_result.images or []) - errors.extend(scan_result.errors or []) - except Exception as e: - errors.append(str(e)) - - return RescanImagesResult(images=images, errors=errors) - - -async def scan_single_image( - db: ExtendedAsyncSAEngine, - registry_key: str, - registry_row: ContainerRegistryRow, - image_canonical: str, -) -> RescanImagesResult: - """ - Performs a scan for a single image. - """ - registry_name = ImageRef.parse_image_str(registry_key, "*").registry - image_name = image_canonical.removeprefix(registry_name + "/") - - log.debug("running a per-image metadata scan: {}, {}", registry_name, image_name) - - scanner_cls = get_container_registry_cls(registry_row) - scanner = scanner_cls(db, registry_name, registry_row) - return await scanner.scan_single_ref(image_name) - - -def filter_registry_dict( - registries: dict[str, ContainerRegistryRow], - condition: Callable[[str, ContainerRegistryRow], bool], -) -> dict[str, ContainerRegistryRow]: - return { - registry_key: registry_row - for registry_key, registry_row in registries.items() - if condition(registry_key, registry_row) - } - - -def filter_registries_by_img_canonical( - registries: dict[str, ContainerRegistryRow], registry_or_image: str -) -> dict[str, ContainerRegistryRow]: - """ - Filters the matching registry assuming `registry_or_image` is an image canonical name. - """ - return filter_registry_dict( - registries, - lambda registry_key, _row: registry_or_image.startswith(registry_key + "/"), - ) - - -def filter_registries_by_registry_name( - registries: dict[str, ContainerRegistryRow], registry_or_image: str -) -> dict[str, ContainerRegistryRow]: - """ - Filters the matching registry assuming `registry_or_image` is a registry name. - """ - return filter_registry_dict( - registries, - lambda registry_key, _row: registry_key.startswith(registry_or_image), - ) - - -async def rescan_images( - db: ExtendedAsyncSAEngine, - registry_or_image: str | None = None, - project: str | None = None, - *, - reporter: ProgressReporter | None = None, -) -> RescanImagesResult: - """ - Rescan container registries and the update images table. - Refer to the comments below for details on the function's behavior. - - If registry name is provided for `registry_or_image`, scans all images in the specified registry. - If image canonical name is provided for `registry_or_image`, only scan the image. - If the `registry_or_image` is not provided, scan all configured registries. - - If `project` is provided, only scan the registries associated with the project. - """ - registries = await load_configured_registries(db, project) - - if registry_or_image is None: - return await scan_registries(db, registries, reporter=reporter) - - matching_registries = filter_registries_by_img_canonical(registries, registry_or_image) - - if matching_registries: - if len(matching_registries) > 1: - raise RuntimeError( - "ContainerRegistryRows exist with the same registry_name and project!", - ) - - registry_key, registry_row = next(iter(matching_registries.items())) - return await scan_single_image(db, registry_key, registry_row, registry_or_image) - - matching_registries = filter_registries_by_registry_name(registries, registry_or_image) - - if not matching_registries: - raise RuntimeError("It is an unknown registry.", registry_or_image) - - log.debug("running a per-registry metadata scan") - return await scan_registries(db, matching_registries, reporter=reporter) - # TODO: delete images removed from registry? - - type Resources = dict[SlotName, dict[str, Any]] @@ -953,25 +796,6 @@ def to_detailed_dataclass(self) -> ImageDataWithDetails: hash=self.trimmed_digest or None, ) - async def untag_image_from_registry( - self, db: ExtendedAsyncSAEngine, session: AsyncSession - ) -> None: - """ - Works only for HarborV2 registries. - """ - from ai.backend.manager.container_registry.harbor import HarborRegistry_v2 - - query = sa.select(ContainerRegistryRow).where(ContainerRegistryRow.id == self.registry_id) - - registry_info = (await session.execute(query)).scalar() - if registry_info is None: - raise RuntimeError(f"Registry not found for image {self.name}") - if registry_info.type != ContainerRegistryType.HARBOR2: - raise NotImplementedError("This feature is only supported for Harbor 2 registries") - - scanner = HarborRegistry_v2(db, self.image_ref.registry, registry_info) - await scanner.untag(self.image_ref) - async def bulk_get_image_configs( image_refs: Iterable[ImageRef], diff --git a/src/ai/backend/manager/repositories/image/db_source/db_source.py b/src/ai/backend/manager/repositories/image/db_source/db_source.py index b21515a5d3a..76b7705d419 100644 --- a/src/ai/backend/manager/repositories/image/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/image/db_source/db_source.py @@ -1,7 +1,8 @@ from __future__ import annotations +import functools import logging -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import cast from uuid import UUID @@ -15,7 +16,9 @@ from ai.backend.common.docker import ImageRef from ai.backend.common.exception import UnknownImageReference from ai.backend.common.types import ImageAlias, ImageID +from ai.backend.common.utils import join_non_empty from ai.backend.logging.utils import BraceStyleAdapter +from ai.backend.manager.container_registry import get_container_registry_cls from ai.backend.manager.data.image.types import ( ImageAliasData, ImageAliasListResult, @@ -41,8 +44,6 @@ ImageAliasRow, ImageIdentifier, ImageRow, - rescan_images, - scan_single_image, ) from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.repositories.base import BatchQuerier, execute_batch_querier @@ -306,16 +307,22 @@ async def scan_and_upsert_image( if not registry_row: raise RegistryNotFoundForImage(f"Registry not found for image {image_canonical}") - # Call the original scan function - return await scan_single_image(self._db, registry_key, registry_row, image_canonical) + return await self.scan_single_image(registry_key, registry_row, image_canonical) - async def remove_tag_from_registry(self, image_id: UUID) -> ImageData: + async def fetch_image_and_registry( + self, image_id: UUID + ) -> tuple[ImageData, ImageRef, ContainerRegistryRow]: + """Read the image (as data + ref) and its container registry row. + + Used by the repository to orchestrate a registry untag without holding a + DB session open across the external registry call. + """ async with self._db.begin_readonly_session() as session: image_row = await self._get_image_by_id(session, image_id, load_aliases=True) - if not image_row: - raise ImageNotFound() - await image_row.untag_image_from_registry(self._db, session) - return image_row.to_dataclass() + registry_row = await session.get(ContainerRegistryRow, image_row.registry_id) + if registry_row is None: + raise RegistryNotFoundForImage(f"Registry not found for image {image_id}") + return image_row.to_dataclass(), image_row.image_ref, registry_row async def modify_image_properties(self, updater: Updater[ImageRow]) -> ImageData: try: @@ -475,12 +482,128 @@ async def rescan_images( reporter: ProgressReporter | None = None, ) -> RescanImagesResult: """ - Rescan container registries and update images table. - Wraps the rescan_images function from models/image/row.py. + Rescan container registries and update the images table. + + If registry name is provided for ``registry_or_image``, scans all images in that registry. + If image canonical name is provided, only scan that image. + If ``registry_or_image`` is not provided, scan all configured registries. + If ``project`` is provided, only scan the registries associated with the project. """ - return await rescan_images( - self._db, - registry_or_image, - project, - reporter=reporter, + registries = await self._load_configured_registries(project) + + if registry_or_image is None: + return await self._scan_registries(registries, reporter=reporter) + + matching_registries = self._filter_by_img_canonical(registries, registry_or_image) + + if matching_registries: + if len(matching_registries) > 1: + raise RuntimeError( + "ContainerRegistryRows exist with the same registry_name and project!", + ) + + registry_key, registry_row = next(iter(matching_registries.items())) + return await self.scan_single_image(registry_key, registry_row, registry_or_image) + + matching_registries = self._filter_by_registry_name(registries, registry_or_image) + + if not matching_registries: + raise RuntimeError("It is an unknown registry.", registry_or_image) + + log.debug("running a per-registry metadata scan") + return await self._scan_registries(matching_registries, reporter=reporter) + # TODO: delete images removed from registry? + + async def scan_single_image( + self, + registry_key: str, + registry_row: ContainerRegistryRow, + image_canonical: str, + ) -> RescanImagesResult: + """Performs a scan for a single image.""" + registry_name = ImageRef.parse_image_str(registry_key, "*").registry + image_name = image_canonical.removeprefix(registry_name + "/") + + log.debug("running a per-image metadata scan: {}, {}", registry_name, image_name) + + scanner_cls = get_container_registry_cls(registry_row) + scanner = scanner_cls(self._db, registry_name, registry_row) + return await scanner.scan_single_ref(image_name) + + async def _load_configured_registries( + self, project: str | None + ) -> dict[str, ContainerRegistryRow]: + join = functools.partial(join_non_empty, sep="/") + + async with self._db.begin_readonly_session() as session: + result = await session.execute(sa.select(ContainerRegistryRow)) + if project: + registries = cast( + dict[str, ContainerRegistryRow], + { + join(row.registry_name, row.project): row + for row in result.scalars().all() + if row.project == project + }, + ) + else: + registries = { + join(row.registry_name, row.project): row for row in result.scalars().all() + } + + return registries + + async def _scan_registries( + self, + registries: dict[str, ContainerRegistryRow], + reporter: ProgressReporter | None = None, + ) -> RescanImagesResult: + """Performs an image rescan for all images in the registries.""" + images, errors = [], [] + + for registry_key, registry_row in registries.items(): + registry_name = ImageRef.parse_image_str(registry_key, "*").registry + log.info('Scanning kernel images from the registry "{0}"', registry_name) + + scanner_cls = get_container_registry_cls(registry_row) + scanner = scanner_cls(self._db, registry_name, registry_row) + + try: + scan_result = await scanner.rescan_single_registry(reporter) + images.extend(scan_result.images or []) + errors.extend(scan_result.errors or []) + except Exception as e: + errors.append(str(e)) + + return RescanImagesResult(images=images, errors=errors) + + @staticmethod + def _filter_registry_dict( + registries: dict[str, ContainerRegistryRow], + condition: Callable[[str, ContainerRegistryRow], bool], + ) -> dict[str, ContainerRegistryRow]: + return { + registry_key: registry_row + for registry_key, registry_row in registries.items() + if condition(registry_key, registry_row) + } + + @classmethod + def _filter_by_img_canonical( + cls, registries: dict[str, ContainerRegistryRow], registry_or_image: str + ) -> dict[str, ContainerRegistryRow]: + """Filter registries assuming ``registry_or_image`` is an image canonical name.""" + return cls._filter_registry_dict( + registries, + lambda registry_key, _row: registry_or_image.startswith(registry_key + "/"), + ) + + @classmethod + def _filter_by_registry_name( + cls, registries: dict[str, ContainerRegistryRow], registry_or_image: str + ) -> dict[str, ContainerRegistryRow]: + """Filter registries assuming ``registry_or_image`` is a registry name.""" + return cls._filter_registry_dict( + registries, + lambda registry_key, _row: registry_key.startswith(registry_or_image), ) diff --git a/src/ai/backend/manager/repositories/image/repository.py b/src/ai/backend/manager/repositories/image/repository.py index c9cbc896c86..1b89df337d0 100644 --- a/src/ai/backend/manager/repositories/image/repository.py +++ b/src/ai/backend/manager/repositories/image/repository.py @@ -5,6 +5,7 @@ from ai.backend.common.bgtask.reporter import ProgressReporter from ai.backend.common.clients.valkey_client.valkey_image.client import ValkeyImageClient +from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.docker import ImageRef from ai.backend.common.exception import BackendAIError from ai.backend.common.metrics.metric import DomainType, LayerType @@ -13,6 +14,7 @@ from ai.backend.common.resilience.resilience import Resilience from ai.backend.common.types import AgentId, ImageAlias, ImageID from ai.backend.manager.config.provider import ManagerConfigProvider +from ai.backend.manager.container_registry.harbor import HarborRegistry_v2 from ai.backend.manager.data.image.types import ( ImageAgentInstallStatus, ImageAliasData, @@ -55,6 +57,7 @@ class ImageRepository: + _db: ExtendedAsyncSAEngine _db_source: ImageDBSource _stateful_source: ImageStatefulSource _config_provider: ManagerConfigProvider @@ -65,6 +68,7 @@ def __init__( valkey_image: ValkeyImageClient, config_provider: ManagerConfigProvider, ) -> None: + self._db = db self._db_source = ImageDBSource(db) self._stateful_source = ImageStatefulSource(valkey_image) self._config_provider = config_provider @@ -265,7 +269,14 @@ async def scan_image_by_identifier( @image_repository_resilience.apply() async def untag_image_from_registry(self, image_id: UUID) -> ImageData: - return await self._db_source.remove_tag_from_registry(image_id) + image_data, image_ref, registry_row = await self._db_source.fetch_image_and_registry( + image_id + ) + if registry_row.type != ContainerRegistryType.HARBOR2: + raise NotImplementedError("This feature is only supported for Harbor 2 registries") + scanner = HarborRegistry_v2(self._db, image_ref.registry, registry_row) + await scanner.untag(image_ref) + return image_data @image_repository_resilience.apply() async def update_image_properties(self, updater: Updater[ImageRow]) -> ImageData: