feat(BA-674): Implement Image Soft/Hard Delete APIs

changes/3628.feature.md

@@ -0,0 +1 @@
+Update `ForgetImage`, `ForgetImageById`, `ClearImages` to perform soft delete and add `PurgeImageById` API for hard delete.

docs/manager/graphql-reference/schema.graphql

@@ -1942,7 +1942,12 @@
 
   """Added in 24.03.0"""
   forget_image_by_id(image_id: String!): ForgetImageById
-  forget_image(architecture: String = "x86_64", reference: String!): ForgetImage
+
+  """Deprecated since 25.4.0. Use `forget_image_by_id` instead."""
+  forget_image(architecture: String = "x86_64", reference: String!): ForgetImage @deprecated(reason: "Deprecated since 25.4.0. Use `forget_image_by_id` instead.")
+
+  """Added in 25.4.0"""
+  purge_image_by_id(image_id: String!): PurgeImageById
 
   """Added in 24.03.1"""
   untag_image_from_registry(image_id: String!): UntagImageFromRegistry
@@ -2496,7 +2501,8 @@
   image: ImageNode
 }
 
+"""Deprecated since 25.4.0. Use `forget_image_by_id` instead."""
 type ForgetImage {
   ok: Boolean
   msg: String
 
@@ -2504,6 +2510,11 @@
   image: ImageNode
 }
 
+"""Added in 25.4.0."""
+type PurgeImageById {
+  image: ImageNode
+}
+
 """Added in 24.03.1"""
 type UntagImageFromRegistry {
   ok: Boolean

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

@@ -68,7 +68,7 @@ def list(ctx: CLIContext, customized: bool) -> None:
 @click.argument("reference_or_id", type=str)
 @click.option("--arch", type=str, default=None, help="Set an explicit architecture.")
 def forget(reference_or_id, arch):
-    """Forget image from server. This command will only work for image customized by user
+    """Mark image as deleted from server. This command will only work for image customized by user
     unless callee has superadmin privileges.
 
     REFERENCE_OR_ID: Canonical string of image (<registry>/<project>/<name>:<tag>)"""

src/ai/backend/common/docker.py

@@ -329,6 +329,12 @@ def __init__(self, tags: Iterable[str], value: Optional[str] = None) -> None:
                 value = ""
             self._data[key] = value
 
+    def __repr__(self):
+        return self.__str__()
+
+    def __str__(self):
+        return f"PlatformTagSet({str(self._data)})"
+
     def has(self, key: str, version: Optional[str] = None):
         if version is None:
             return key in self._data

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

@@ -15,6 +15,7 @@
 from .image_impl import forget_image as forget_image_impl
 from .image_impl import inspect_image as inspect_image_impl
 from .image_impl import list_images as list_images_impl
+from .image_impl import purge_image as purge_image_impl
 from .image_impl import rescan_images as rescan_images_impl
 from .image_impl import set_image_resource_limit as set_image_resource_limit_impl
 from .image_impl import validate_image_alias as validate_image_alias_impl
@@ -51,10 +52,19 @@ def inspect(cli_ctx, canonical_or_alias, architecture) -> None:
 @click.argument("architecture")
 @click.pass_obj
 def forget(cli_ctx, canonical_or_alias, architecture) -> None:
-    """Forget (delete) a specific image."""
+    """Forget (soft-delete) a specific image."""
     asyncio.run(forget_image_impl(cli_ctx, canonical_or_alias, architecture))
 
 
+@cli.command()
+@click.argument("canonical_or_alias")
+@click.argument("architecture")
+@click.pass_obj
+def purge(cli_ctx, canonical_or_alias, architecture) -> None:
+    """Purge (hard-delete) a specific image."""
+    asyncio.run(purge_image_impl(cli_ctx, canonical_or_alias, architecture))
+
+
 @cli.command()
 @click.argument("canonical_or_alias")
 @click.argument("slot_type")

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

@@ -81,8 +81,8 @@ async def _build_smembers_pipeline(redis: Redis) -> Pipeline:
                         pprint(item)
             if short:
                 print(tabulate(displayed_items, tablefmt="plain"))
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def inspect_image(cli_ctx, canonical_or_alias, architecture):
@@ -101,8 +101,8 @@ async def inspect_image(cli_ctx, canonical_or_alias, architecture):
             pprint(await image_row.inspect())
         except UnknownImageReference:
             log.exception("Image not found.")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def forget_image(cli_ctx, canonical_or_alias, architecture):
@@ -118,11 +118,32 @@ async def forget_image(cli_ctx, canonical_or_alias, architecture):
                     ImageAlias(canonical_or_alias),
                 ],
             )
+            await image_row.mark_as_deleted(session)
+        except UnknownImageReference:
+            log.exception("Image not found.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
+
+
+async def purge_image(cli_ctx, canonical_or_alias, architecture):
+    async with (
+        connect_database(cli_ctx.local_config) as db,
+        db.begin_session() as session,
+    ):
+        try:
+            image_row = await ImageRow.resolve(
+                session,
+                [
+                    ImageIdentifier(canonical_or_alias, architecture),
+                    ImageAlias(canonical_or_alias),
+                ],
+                load_only_active=False,
+            )
             await session.delete(image_row)
         except UnknownImageReference:
             log.exception("Image not found.")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def set_image_resource_limit(
@@ -147,8 +168,8 @@ async def set_image_resource_limit(
             await image_row.set_resource_limit(slot_type, range_value)
         except UnknownImageReference:
             log.exception("Image not found.")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def rescan_images(
@@ -161,8 +182,8 @@ async def rescan_images(
     ):
         try:
             await rescan_images_func(db, registry_or_image, project)
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def alias(cli_ctx, alias, target, architecture):
@@ -180,8 +201,8 @@ async def alias(cli_ctx, alias, target, architecture):
             await ImageAliasRow.create(session, alias, image_row)
         except UnknownImageReference:
             log.exception("Image not found.")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def dealias(cli_ctx, alias):
@@ -213,8 +234,8 @@ async def validate_image_alias(cli_ctx, alias: str) -> None:
 
         except UnknownImageReference:
             log.error(f"No images were found with alias: {alias}")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")
 
 
 async def validate_image_canonical(
@@ -257,5 +278,5 @@ async def validate_image_canonical(
 
         except UnknownImageReference as e:
             log.error(f"{e}")
-        except Exception:
-            log.exception("An error occurred.")
+        except Exception as e:
+            log.exception(f"An error occurred. Error: {e}")

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

@@ -119,6 +119,7 @@
     ImagePermissionValueField,
     ModifyImage,
     PreloadImage,
+    PurgeImageById,
     RescanImages,
     UnloadImage,
     UntagImageFromRegistry,
@@ -296,7 +297,10 @@ class Mutations(graphene.ObjectType):
     unload_image = UnloadImage.Field()
     modify_image = ModifyImage.Field()
     forget_image_by_id = ForgetImageById.Field(description="Added in 24.03.0")
-    forget_image = ForgetImage.Field()
+    forget_image = ForgetImage.Field(
+        deprecation_reason="Deprecated since 25.4.0. Use `forget_image_by_id` instead."
+    )
+    purge_image_by_id = PurgeImageById.Field(description="Added in 25.4.0")
     untag_image_from_registry = UntagImageFromRegistry.Field(description="Added in 24.03.1")
     alias_image = AliasImage.Field()
     dealias_image = DealiasImage.Field()

src/ai/backend/manager/models/gql_models/base.py

@@ -2,6 +2,7 @@
 
 import uuid
 from typing import Any, Optional
+from uuid import UUID
 
 import graphene
 import graphql
@@ -11,6 +12,9 @@
 from graphql.language.ast import FloatValueNode, IntValueNode, ObjectValueNode, ValueNode
 from graphql.language.printer import print_ast
 
+from ...api.exceptions import ObjectNotFound
+from ..gql_relay import AsyncNode
+
 SAFE_MIN_INT = -9007199254740991
 SAFE_MAX_INT = 9007199254740991
 
@@ -151,3 +155,19 @@ def parse_value(value: Any) -> dict[str, float]:
                 raise GraphQLError(f"UUIDFloatMap cannot represent value {v} as a float")
             validated[key_str] = v
         return validated
+
+
+def extract_object_uuid(info: graphene.ResolveInfo, global_id: str, object_name: str) -> UUID:
+    """
+    Converts a GraphQL global ID to its corresponding UUID.
+    If the global ID is not valid, raises an error using the provided object name.
+    """
+
+    _, raw_id = AsyncNode.resolve_global_id(info, global_id)
+    if not raw_id:
+        raw_id = global_id
+
+    try:
+        return UUID(raw_id)
+    except ValueError:
+        raise ObjectNotFound(object_name)

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

@@ -43,7 +43,7 @@
 from ai.backend.manager.models.rbac.context import ClientContext
 from ai.backend.manager.models.rbac.permission_defs import ImagePermission
 
-from ...api.exceptions import ImageNotFound, ObjectNotFound
+from ...api.exceptions import GenericForbidden, ImageNotFound, ObjectNotFound
 from ...defs import DEFAULT_IMAGE_ARCH
 from ..base import (
     FilterExprArg,
@@ -58,6 +58,7 @@
     ImageIdentifier,
     ImageLoadFilter,
     ImageRow,
+    ImageStatus,
     ImageType,
     get_permission_ctx,
     rescan_images,
@@ -70,6 +71,7 @@
     KVPairInput,
     ResourceLimit,
     ResourceLimitInput,
+    extract_object_uuid,
 )
 
 if TYPE_CHECKING:
@@ -86,6 +88,7 @@
     "RescanImages",
     "ForgetImage",
     "ForgetImageById",
+    "PurgeImageById",
     "UntagImageFromRegistry",
     "ModifyImage",
     "AliasImage",
@@ -675,37 +678,28 @@ async def mutate(
         info: graphene.ResolveInfo,
         image_id: str,
     ) -> ForgetImageById:
-        _, raw_image_id = AsyncNode.resolve_global_id(info, image_id)
-        if not raw_image_id:
-            raw_image_id = image_id
-
-        try:
-            _image_id = UUID(raw_image_id)
-        except ValueError:
-            raise ObjectNotFound("image")
-
         log.info("forget image {0} by API request", image_id)
+        image_uuid = extract_object_uuid(info, image_id, "image")
+
         ctx: GraphQueryContext = info.context
         client_role = ctx.user["role"]
 
         async with ctx.db.begin_session() as session:
-            image_row = await ImageRow.get(session, _image_id, load_aliases=True)
+            image_row = await ImageRow.get(session, image_uuid, load_aliases=True)
             if not image_row:
                 raise ObjectNotFound("image")
             if client_role != UserRole.SUPERADMIN:
-                customized_image_owner = (image_row.labels or {}).get(
-                    "ai.backend.customized-image.owner"
-                )
-                if (
-                    not customized_image_owner
-                    or customized_image_owner != f"user:{ctx.user['uuid']}"
-                ):
+                if not image_row.is_customized_by(ctx.user["uuid"]):
                     return ForgetImageById(ok=False, msg="Forbidden")
-            await session.delete(image_row)
+            await image_row.mark_as_deleted(session)
             return ForgetImageById(ok=True, msg="", image=ImageNode.from_row(ctx, image_row))
 
 
 class ForgetImage(graphene.Mutation):
+    """
+    Deprecated since 25.4.0. Use `forget_image_by_id` instead.
+    """
+
     allowed_roles = (
         UserRole.SUPERADMIN,
         UserRole.ADMIN,
@@ -740,18 +734,49 @@ async def mutate(
                 ],
             )
             if client_role != UserRole.SUPERADMIN:
-                customized_image_owner = (image_row.labels or {}).get(
-                    "ai.backend.customized-image.owner"
-                )
-                if (
-                    not customized_image_owner
-                    or customized_image_owner != f"user:{ctx.user['uuid']}"
-                ):
+                if not image_row.is_customized_by(ctx.user["uuid"]):
                     return ForgetImage(ok=False, msg="Forbidden")
-            await session.delete(image_row)
+            await image_row.mark_as_deleted(session)
             return ForgetImage(ok=True, msg="", image=ImageNode.from_row(ctx, image_row))
 
 
+class PurgeImageById(graphene.Mutation):
+    """Added in 25.4.0."""
+
+    allowed_roles = (
+        UserRole.SUPERADMIN,
+        UserRole.ADMIN,
+        UserRole.USER,
+    )
+
+    class Arguments:
+        image_id = graphene.String(required=True)
+
+    image = graphene.Field(ImageNode)
+
+    @staticmethod
+    async def mutate(
+        root: Any,
+        info: graphene.ResolveInfo,
+        image_id: str,
+    ) -> PurgeImageById:
+        log.info("purge image {0} by API request", image_id)
+        image_uuid = extract_object_uuid(info, image_id, "image")
+
+        ctx: GraphQueryContext = info.context
+        client_role = ctx.user["role"]
+
+        async with ctx.db.begin_session() as session:
+            image_row = await ImageRow.get(session, image_uuid, load_aliases=True)
+            if not image_row:
+                raise ObjectNotFound("image")
+            if client_role != UserRole.SUPERADMIN:
+                if not image_row.is_customized_by(ctx.user["uuid"]):
+                    raise GenericForbidden("Image is not owned by your account.")
+            await session.delete(image_row)
+            return PurgeImageById(image=ImageNode.from_row(ctx, image_row))
+
+
 class UntagImageFromRegistry(graphene.Mutation):
     """Added in 24.03.1"""
 
@@ -776,31 +801,18 @@ async def mutate(
     ) -> UntagImageFromRegistry:
         from ai.backend.manager.container_registry.harbor import HarborRegistry_v2
 
-        _, raw_image_id = AsyncNode.resolve_global_id(info, image_id)
-        if not raw_image_id:
-            raw_image_id = image_id
+        image_uuid = extract_object_uuid(info, image_id, "image")
 
-        try:
-            _image_id = UUID(raw_image_id)
-        except ValueError:
-            raise ObjectNotFound("image")
-
-        log.info("remove image from registry {0} by API request", str(_image_id))
+        log.info("remove image from registry {0} by API request", str(image_uuid))
         ctx: GraphQueryContext = info.context
         client_role = ctx.user["role"]
 
         async with ctx.db.begin_readonly_session() as session:
-            image_row = await ImageRow.get(session, _image_id, load_aliases=True)
+            image_row = await ImageRow.get(session, image_uuid, load_aliases=True)
             if not image_row:
                 raise ImageNotFound
             if client_role != UserRole.SUPERADMIN:
-                customized_image_owner = (image_row.labels or {}).get(
-                    "ai.backend.customized-image.owner"
-                )
-                if (
-                    not customized_image_owner
-                    or customized_image_owner != f"user:{ctx.user['uuid']}"
-                ):
+                if not image_row.is_customized_by(ctx.user["uuid"]):
                     return UntagImageFromRegistry(ok=False, msg="Forbidden")
 
             query = sa.select(ContainerRegistryRow).where(
@@ -975,15 +987,12 @@ async def mutate(
         ctx: GraphQueryContext = info.context
         try:
             async with ctx.db.begin_session() as session:
-                result = await session.execute(
-                    sa.select(ImageRow).where(ImageRow.registry == registry)
-                )
-                image_ids = [x.id for x in result.scalars().all()]
-
                 await session.execute(
-                    sa.delete(ImageAliasRow).where(ImageAliasRow.image_id.in_(image_ids))
+                    sa.update(ImageRow)
+                    .where(ImageRow.registry == registry)
+                    .where(ImageRow.status != ImageStatus.DELETED)
+                    .values(status=ImageStatus.DELETED)
                 )
-                await session.execute(sa.delete(ImageRow).where(ImageRow.registry == registry))
         except ValueError as e:
             return ClearImages(ok=False, msg=str(e))
         return ClearImages(ok=True, msg="")

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

@@ -696,6 +696,10 @@ async def inspect(self) -> Mapping[str, Any]:
         parsed_image_info["reverse_aliases"] = [x.alias for x in self.aliases]
         return parsed_image_info
 
+    async def mark_as_deleted(self, db_session: AsyncSession) -> None:
+        self.status = ImageStatus.DELETED
+        await db_session.flush()
+
     def set_resource_limit(
         self,
         slot_type: str,
@@ -711,6 +715,9 @@ def set_resource_limit(
 
         self.resources = resources
 
+    def is_customized_by(self, user_id: str) -> bool:
+        return (self.labels or {}).get("ai.backend.customized-image.owner") == f"user:{user_id}"
+
 
 async def bulk_get_image_configs(
     image_refs: Iterable[ImageRef],