refactor(BA-6632): move image scanning out of the models layer into the image repository

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.

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:

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",
 )

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],