diff --git a/changes/12449.fix.md b/changes/12449.fix.md new file mode 100644 index 00000000000..c1ed2fdc205 --- /dev/null +++ b/changes/12449.fix.md @@ -0,0 +1 @@ +Retry committing container registry image rescan results after transient database serialization conflicts. diff --git a/src/ai/backend/manager/container_registry/base.py b/src/ai/backend/manager/container_registry/base.py index 560bc83e08a..e156666d4bf 100644 --- a/src/ai/backend/manager/container_registry/base.py +++ b/src/ai/backend/manager/container_registry/base.py @@ -7,8 +7,8 @@ from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import asynccontextmanager as actxmgr from contextvars import ContextVar +from dataclasses import dataclass from typing import ( - TYPE_CHECKING, Any, Final, cast, @@ -30,14 +30,18 @@ ) from ai.backend.common.docker import login as registry_login from ai.backend.common.exception import ( + BackendAIError, InvalidImageName, InvalidImageTag, ProjectMismatchWithCanonical, ) from ai.backend.common.json import read_json +from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy +from ai.backend.common.resilience.resilience import Resilience from ai.backend.common.types import SlotName, SSLContextType from ai.backend.common.utils import join_non_empty from ai.backend.logging import BraceStyleAdapter +from ai.backend.logging.types import LogLevel from ai.backend.manager.data.image.types import ( ImageData, ImageStatus, @@ -47,6 +51,7 @@ from ai.backend.manager.data.permission.types import RBACElementRef from ai.backend.manager.defs import INTRINSIC_SLOTS_MIN from ai.backend.manager.exceptions import ScanImageError, ScanTagError +from ai.backend.manager.models.container_registry import ContainerRegistryRow from ai.backend.manager.models.image import ImageIdentifier, ImageRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.repositories.base.rbac.entity_creator import ( @@ -61,9 +66,30 @@ "progress_reporter", default=None ) all_updates: ContextVar[dict[ImageIdentifier, dict[str, Any]]] = ContextVar("all_updates") +commit_rescan_result_resilience = Resilience( + policies=[ + RetryPolicy( + RetryArgs( + max_retries=10, + retry_delay=0.1, + backoff_strategy=BackoffStrategy.FIXED, + non_retryable_exceptions=(BackendAIError,), + ) + ), + ] +) + + +@dataclass(frozen=True) +class _RescanProgressEvent: + """ + A rescan progress message buffered during the commit transaction and emitted + only after a successful commit, so retries do not duplicate logs or over-count + the progress reporter. + """ -if TYPE_CHECKING: - from ai.backend.manager.models.container_registry import ContainerRegistryRow + level: LogLevel + message: str class BaseContainerRegistry(metaclass=ABCMeta): @@ -176,97 +202,129 @@ def _determine_additional_image_scopes( async def commit_rescan_result(self) -> list[ImageData]: scanned_images: list[ImageData] = [] - _all_updates = all_updates.get() - if not _all_updates: + original_updates = all_updates.get() + if not original_updates: log.info("No images found in registry {0}", self.registry_url) else: - image_identifiers = [(k.canonical, k.architecture) for k in _all_updates.keys()] - async with self.db.begin_session() as session: - existing_images = await session.scalars( - sa.select(ImageRow).where( - sa.func.ROW(ImageRow.name, ImageRow.architecture).in_(image_identifiers), - ) + image_identifiers = [(k.canonical, k.architecture) for k in original_updates.keys()] + + scanned_images, progress_events = await self._commit_rescan_result_once( + original_updates, + image_identifiers, + ) + for event in progress_events: + match event.level: + case LogLevel.WARNING: + log.warning(event.message) + case _: + log.info(event.message) + if (reporter := progress_reporter.get()) is not None: + await reporter.update(1, message=event.message) + return scanned_images + + @commit_rescan_result_resilience.apply() + async def _commit_rescan_result_once( + self, + original_updates: dict[ImageIdentifier, dict[str, Any]], + image_identifiers: list[tuple[str, str]], + ) -> tuple[list[ImageData], list[_RescanProgressEvent]]: + scanned_images: list[ImageData] = [] + progress_events: list[_RescanProgressEvent] = [] + pending_updates = dict(original_updates) + + async with self.db.begin_session() as session: + # Serialize commits per registry. Some images may not exist yet, so locking + # only ImageRow cannot prevent concurrent insert races for the same registry. + await session.execute( + sa.select(ContainerRegistryRow.id) + .where(ContainerRegistryRow.id == self.registry_info.id) + .with_for_update() + ) + existing_images = await session.scalars( + sa.select(ImageRow).where( + sa.func.ROW(ImageRow.name, ImageRow.architecture).in_(image_identifiers), ) - is_local = self.registry_name == "local" - - for image_row in existing_images: - image_ref = image_row.image_ref - update_key = ImageIdentifier(image_ref.canonical, image_ref.architecture) - if update := _all_updates.pop(update_key, None): - image_row.config_digest = update["config_digest"] - image_row.size_bytes = update["size_bytes"] - image_row.accelerators = update.get("accels") - image_row.labels = update["labels"] - image_row.is_local = is_local - scanned_images.append(image_row.to_dataclass()) - - if image_row.status == ImageStatus.DELETED: - image_row.status = ImageStatus.ALIVE - - progress_msg = f"Restored deleted image - {image_ref.canonical}/{image_ref.architecture} ({update['config_digest']})" - log.info(progress_msg) - - if (reporter := progress_reporter.get()) is not None: - await reporter.update(1, message=progress_msg) - - rbac_creators: list[RBACEntityCreator[ImageRow]] = [] - for image_identifier, update in _all_updates.items(): - try: - parsed_img = ImageRef.from_image_str( - image_identifier.canonical, - self.registry_info.project, - self.registry_info.registry_name, - is_local=is_local, + ) + is_local = self.registry_name == "local" + + for image_row in existing_images: + image_ref = image_row.image_ref + update_key = ImageIdentifier(image_ref.canonical, image_ref.architecture) + if update := pending_updates.pop(update_key, None): + image_row.config_digest = update["config_digest"] + image_row.size_bytes = update["size_bytes"] + image_row.accelerators = update.get("accels") + image_row.labels = update["labels"] + image_row.is_local = is_local + scanned_images.append(image_row.to_dataclass()) + + if image_row.status == ImageStatus.DELETED: + image_row.status = ImageStatus.ALIVE + + progress_msg = ( + f"Restored deleted image - {image_ref.canonical}/" + f"{image_ref.architecture} ({update['config_digest']})" ) - except (ProjectMismatchWithCanonical, ValueError) as e: - skip_reason = str(e) - progress_msg = f"Skipped image - {image_identifier.canonical}/{image_identifier.architecture} ({skip_reason})" - log.warning(progress_msg) - if (reporter := progress_reporter.get()) is not None: - await reporter.update(1, message=progress_msg) - continue - - rbac_creators.append( - RBACEntityCreator( - spec=ImageRowCreatorSpec( - name=parsed_img.canonical, - project=self.registry_info.project, - architecture=image_identifier.architecture, - registry_id=self.registry_info.id, - is_local=is_local, - registry=parsed_img.registry, - image=join_non_empty(parsed_img.project, parsed_img.name, sep="/"), - tag=parsed_img.tag, - config_digest=update["config_digest"], - size_bytes=update["size_bytes"], - type=ImageType.COMPUTE, - accelerators=update.get("accels"), - labels=update["labels"], - status=ImageStatus.ALIVE, - ), - scope_ref=RBACElementRef( - RBACElementType.CONTAINER_REGISTRY, - str(self.registry_info.id), - ), - additional_scope_refs=self._determine_additional_image_scopes( - update["labels"] - ), - element_type=RBACElementType.IMAGE, - ), - ) + progress_events.append(_RescanProgressEvent(LogLevel.INFO, progress_msg)) - bulk_result = await execute_rbac_entity_creators(session, rbac_creators) - for row in bulk_result.rows: - scanned_images.append(row.to_dataclass()) + rbac_creators: list[RBACEntityCreator[ImageRow]] = [] + for image_identifier, update in pending_updates.items(): + try: + parsed_img = ImageRef.from_image_str( + image_identifier.canonical, + self.registry_info.project, + self.registry_info.registry_name, + is_local=is_local, + ) + except (ProjectMismatchWithCanonical, ValueError) as e: + skip_reason = str(e) progress_msg = ( - f"Updated image - {row.name}/{row.architecture} ({row.config_digest})" + f"Skipped image - {image_identifier.canonical}/" + f"{image_identifier.architecture} ({skip_reason})" ) - log.info(progress_msg) - if (reporter := progress_reporter.get()) is not None: - await reporter.update(1, message=progress_msg) + progress_events.append(_RescanProgressEvent(LogLevel.WARNING, progress_msg)) + continue - await session.flush() - return scanned_images + rbac_creators.append( + RBACEntityCreator( + spec=ImageRowCreatorSpec( + name=parsed_img.canonical, + project=self.registry_info.project, + architecture=image_identifier.architecture, + registry_id=self.registry_info.id, + is_local=is_local, + registry=parsed_img.registry, + image=join_non_empty(parsed_img.project, parsed_img.name, sep="/"), + tag=parsed_img.tag, + config_digest=update["config_digest"], + size_bytes=update["size_bytes"], + type=ImageType.COMPUTE, + accelerators=update.get("accels"), + labels=update["labels"], + status=ImageStatus.ALIVE, + ), + scope_ref=RBACElementRef( + RBACElementType.CONTAINER_REGISTRY, + str(self.registry_info.id), + ), + additional_scope_refs=self._determine_additional_image_scopes( + update["labels"] + ), + element_type=RBACElementType.IMAGE, + ), + ) + + bulk_result = await execute_rbac_entity_creators(session, rbac_creators) + for row in bulk_result.rows: + scanned_images.append(row.to_dataclass()) + progress_msg = ( + f"Updated image - {row.name}/{row.architecture} ({row.config_digest})" + ) + progress_events.append(_RescanProgressEvent(LogLevel.INFO, progress_msg)) + + await session.flush() + + return scanned_images, progress_events async def scan_single_ref(self, image: str) -> RescanImagesResult: all_updates_token = all_updates.set({}) diff --git a/tests/unit/manager/test_container_registry_base.py b/tests/unit/manager/test_container_registry_base.py new file mode 100644 index 00000000000..1e61d289089 --- /dev/null +++ b/tests/unit/manager/test_container_registry_base.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Any, cast +from uuid import uuid4 + +import aiohttp +import pytest +from sqlalchemy.exc import DBAPIError + +from ai.backend.manager.container_registry import base as registry_base +from ai.backend.manager.container_registry.base import BaseContainerRegistry, all_updates +from ai.backend.manager.data.image.types import ImageData, ImageIdentifier, ImageStatus +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine + + +class _SqlStateError(Exception): + sqlstate: str + pgcode: str + + def __init__(self, sqlstate: str) -> None: + self.sqlstate = sqlstate + self.pgcode = sqlstate + + +class _FakeImageRow: + image_ref: SimpleNamespace + status: ImageStatus + config_digest: str | None + size_bytes: int | None + accelerators: str | None + labels: dict[str, str] + is_local: bool + _image_data: ImageData + + def __init__(self, canonical: str, architecture: str, image_data: ImageData) -> None: + self.image_ref = SimpleNamespace(canonical=canonical, architecture=architecture) + self.status = ImageStatus.ALIVE + self.config_digest = None + self.size_bytes = None + self.accelerators = None + self.labels = {} + self.is_local = False + self._image_data = image_data + + def to_dataclass(self) -> ImageData: + return self._image_data + + +class _FakeSession: + _db: _FakeDB + + def __init__(self, db: _FakeDB) -> None: + self._db = db + + async def execute(self, _stmt: Any) -> None: + self._db.lock_count += 1 + self._db.operations.append("lock") + + async def scalars(self, _stmt: Any) -> list[_FakeImageRow]: + self._db.operations.append("scan") + return [self._db.image_row] + + async def flush(self) -> None: + self._db.flush_count += 1 + self._db.operations.append("flush") + if self._db.flush_count == 1: + raise DBAPIError(None, None, _SqlStateError("40001")) + + +class _FakeDB: + image_row: _FakeImageRow + begin_count: int + lock_count: int + flush_count: int + operations: list[str] + + def __init__(self, image_row: _FakeImageRow) -> None: + self.image_row = image_row + self.begin_count = 0 + self.lock_count = 0 + self.flush_count = 0 + self.operations = [] + + @asynccontextmanager + async def begin_session(self) -> AsyncIterator[_FakeSession]: + self.begin_count += 1 + yield _FakeSession(self) + + +class _TestRegistry(BaseContainerRegistry): + async def fetch_repositories( + self, + _sess: aiohttp.ClientSession, + ) -> AsyncIterator[str]: + repositories: tuple[str, ...] = () + for repository in repositories: + yield repository + + +class TestCommitRescanResult: + @pytest.fixture + def canonical(self) -> str: + return "test-registry/test-project/python:3.12" + + @pytest.fixture + def architecture(self) -> str: + return "x86_64" + + @pytest.fixture + def image_data(self) -> ImageData: + return cast(ImageData, object()) + + @pytest.fixture + def image_row( + self, + canonical: str, + architecture: str, + image_data: ImageData, + ) -> _FakeImageRow: + return _FakeImageRow(canonical, architecture, image_data) + + @pytest.fixture + def db(self, image_row: _FakeImageRow) -> _FakeDB: + return _FakeDB(image_row) + + @pytest.fixture + def registry_info(self) -> SimpleNamespace: + return SimpleNamespace( + url="https://registry.example.com", + ssl_verify=True, + username=None, + password=None, + project="test-project", + registry_name="test-registry", + id=uuid4(), + ) + + @pytest.fixture + def registry( + self, + db: _FakeDB, + registry_info: SimpleNamespace, + ) -> _TestRegistry: + return _TestRegistry( + cast(ExtendedAsyncSAEngine, db), + "test-registry", + cast(Any, registry_info), + ) + + @pytest.fixture + def update_key(self, canonical: str, architecture: str) -> ImageIdentifier: + return ImageIdentifier(canonical, architecture) + + @pytest.fixture + def updates(self, update_key: ImageIdentifier) -> dict[ImageIdentifier, dict[str, Any]]: + return { + update_key: { + "config_digest": "sha256:updated", + "size_bytes": 1024, + "labels": {}, + }, + } + + @pytest.fixture + def active_updates( + self, + updates: dict[ImageIdentifier, dict[str, Any]], + ) -> Iterator[dict[ImageIdentifier, dict[str, Any]]]: + token = all_updates.set(updates) + try: + yield updates + finally: + all_updates.reset(token) + + @pytest.fixture + def no_rbac_creators(self, monkeypatch: pytest.MonkeyPatch) -> None: + async def fake_execute_rbac_entity_creators( + _session: _FakeSession, + creators: list[Any], + ) -> SimpleNamespace: + assert creators == [] + return SimpleNamespace(rows=[]) + + monkeypatch.setattr( + registry_base, + "execute_rbac_entity_creators", + fake_execute_rbac_entity_creators, + ) + + async def test_retries_serialization_failure_without_losing_updates( + self, + registry: _TestRegistry, + db: _FakeDB, + image_row: _FakeImageRow, + image_data: ImageData, + update_key: ImageIdentifier, + active_updates: dict[ImageIdentifier, dict[str, Any]], + no_rbac_creators: None, + ) -> None: + result = await registry.commit_rescan_result() + assert update_key in all_updates.get() + + assert active_updates[update_key]["config_digest"] == "sha256:updated" + assert result == [image_data] + assert db.begin_count == 2 + assert db.lock_count == 2 + assert db.operations[:2] == ["lock", "scan"] + assert db.flush_count == 2 + assert image_row.config_digest == "sha256:updated" + assert image_row.size_bytes == 1024