feat: Implement ID based client workflow to ContainerRegistry graphql API

changes/2615.feature.md

@@ -0,0 +1 @@
+Implement ID-based client workflow to ContainerRegistry API.

changes/2707.feature.md

@@ -0,0 +1 @@
+Add `project` column to the images table and refactoring `ImageRef` logic.

src/ai/backend/agent/agent.py

@@ -213,6 +213,7 @@ def __init__(
         session_id: SessionId,
         agent_id: AgentId,
         event_producer: EventProducer,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         distro: str,
         local_config: Mapping[str, Any],
@@ -229,7 +230,7 @@ def __init__(
         self.agent_id = agent_id
         self.event_producer = event_producer
         self.kernel_config = kernel_config
-        self.image_ref = ImageRef.from_image_config(kernel_config["image"])
+        self.image_ref = kernel_image
         self.distro = distro
         self.internal_data = kernel_config["internal_data"] or {}
         self.computers = computers
@@ -1699,6 +1700,7 @@ async def init_kernel_context(
         self,
         kernel_id: KernelId,
         session_id: SessionId,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         *,
         restarting: bool = False,
@@ -1790,6 +1792,7 @@ async def create_kernel(
         self,
         session_id: SessionId,
         kernel_id: KernelId,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         cluster_info: ClusterInfo,
         *,
@@ -1814,6 +1817,7 @@ async def create_kernel(
             ctx = await self.init_kernel_context(
                 kernel_id,
                 session_id,
+                kernel_image,
                 kernel_config,
                 restarting=restarting,
                 cluster_ssh_port_mapping=cluster_info.get("cluster_ssh_port_mapping"),
@@ -2352,7 +2356,7 @@ def get_public_service_ports(self, service_ports: list[ServicePort]) -> list[Ser
         return [port for port in service_ports if port["protocol"] != ServicePortProtocols.INTERNAL]
 
     @abstractmethod
-    async def extract_image_command(self, image_ref: str) -> str | None:
+    async def extract_image_command(self, image: str) -> str | None:
         raise NotImplementedError
 
     @abstractmethod
@@ -2439,6 +2443,7 @@ async def restart_kernel(
         self,
         session_id: SessionId,
         kernel_id: KernelId,
+        kernel_image: ImageRef,
         updating_kernel_config: KernelCreationConfig,
     ):
         tracker = self.restarting_kernels.get(kernel_id)
@@ -2486,6 +2491,7 @@ async def restart_kernel(
                     await self.create_kernel(
                         session_id,
                         kernel_id,
+                        kernel_image,
                         kernel_config,
                         existing_cluster_info,
                         restarting=True,

src/ai/backend/agent/docker/agent.py

@@ -50,7 +50,7 @@
 from ai.backend.common.cgroup import get_cgroup_mount_point
 from ai.backend.common.docker import MAX_KERNELSPEC, MIN_KERNELSPEC, ImageRef
 from ai.backend.common.events import EventProducer, KernelLifecycleEventReason
-from ai.backend.common.exception import ImageNotAvailable
+from ai.backend.common.exception import ImageNotAvailable, InvalidImageName, InvalidImageTag
 from ai.backend.common.plugin.monitor import ErrorPluginContext, StatsPluginContext
 from ai.backend.common.types import (
     AgentId,
@@ -200,6 +200,7 @@ def __init__(
         session_id: SessionId,
         agent_id: AgentId,
         event_producer: EventProducer,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         distro: str,
         local_config: Mapping[str, Any],
@@ -216,6 +217,7 @@ def __init__(
             session_id,
             agent_id,
             event_producer,
+            kernel_image,
             kernel_config,
             distro,
             local_config,
@@ -1234,10 +1236,10 @@ async def scan_available_resources(self) -> Mapping[SlotName, Decimal]:
             self.local_config, {name: cctx.instance for name, cctx in self.computers.items()}
         )
 
-    async def extract_image_command(self, image_ref: str) -> str | None:
+    async def extract_image_command(self, image: str) -> str | None:
         async with closing_async(Docker()) as docker:
-            image = await docker.images.get(image_ref)
-            return image["Config"].get("Cmd")
+            result = await docker.images.get(image)
+            return result["Config"].get("Cmd")
 
     async def enumerate_containers(
         self,
@@ -1347,11 +1349,12 @@ async def scan_images(self) -> Mapping[str, str]:
                     if repo_tag.endswith("<none>"):
                         continue
                     try:
-                        ImageRef(repo_tag, ["*"])
-                    except ValueError:
+                        ImageRef.parse_image_str(repo_tag, "*")
+                    except (InvalidImageName, InvalidImageTag) as e:
                         log.warning(
-                            "Image name {} does not conform to Backend.AI's image naming rule. This image will be ignored.",
+                            "Image name {} does not conform to Backend.AI's image naming rule. This image will be ignored. Details: {}",
                             repo_tag,
+                            e,
                         )
                         continue
 
@@ -1523,6 +1526,7 @@ async def init_kernel_context(
         self,
         kernel_id: KernelId,
         session_id: SessionId,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         *,
         restarting: bool = False,
@@ -1534,6 +1538,7 @@ async def init_kernel_context(
             session_id,
             self.id,
             self.event_producer,
+            kernel_image,
             kernel_config,
             distro,
             self.local_config,

src/ai/backend/agent/dummy/agent.py

@@ -55,6 +55,7 @@ def __init__(
         session_id: SessionId,
         agent_id: AgentId,
         event_producer: EventProducer,
+        kenrel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         distro: str,
         local_config: Mapping[str, Any],
@@ -68,6 +69,7 @@ def __init__(
             session_id,
             agent_id,
             event_producer,
+            kenrel_image,
             kernel_config,
             distro,
             local_config,
@@ -248,7 +250,7 @@ def get_public_service_ports(self, service_ports: list[ServicePort]) -> list[Ser
     async def sync_container_lifecycles(self, interval: float) -> None:
         return
 
-    async def extract_command(self, image_ref: str) -> str | None:
+    async def extract_command(self, image: str) -> str | None:
         return None
 
     async def enumerate_containers(
@@ -268,7 +270,7 @@ async def scan_available_resources(self) -> Mapping[SlotName, Decimal]:
             self.local_config, {name: cctx.instance for name, cctx in self.computers.items()}
         )
 
-    async def extract_image_command(self, image_ref: str) -> str | None:
+    async def extract_image_command(self, image: str) -> str | None:
         delay = self.dummy_agent_cfg["delay"]["scan-image"]
         await asyncio.sleep(delay)
         return "cr.backend.ai/stable/python:3.9-ubuntu20.04"
@@ -305,6 +307,7 @@ async def init_kernel_context(
         self,
         kernel_id: KernelId,
         session_id: SessionId,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         *,
         restarting: bool = False,
@@ -316,6 +319,7 @@ async def init_kernel_context(
             session_id,
             self.id,
             self.event_producer,
+            kernel_image,
             kernel_config,
             distro,
             self.local_config,

src/ai/backend/agent/kubernetes/agent.py

@@ -105,6 +105,7 @@ def __init__(
         session_id: SessionId,
         agent_id: AgentId,
         event_producer: EventProducer,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         distro: str,
         local_config: Mapping[str, Any],
@@ -119,6 +120,7 @@ def __init__(
             session_id,
             agent_id,
             event_producer,
+            kernel_image,
             kernel_config,
             distro,
             local_config,
@@ -1032,6 +1034,7 @@ async def init_kernel_context(
         self,
         kernel_id: KernelId,
         session_id: SessionId,
+        kernel_image: ImageRef,
         kernel_config: KernelCreationConfig,
         *,
         restarting: bool = False,
@@ -1043,6 +1046,7 @@ async def init_kernel_context(
             session_id,
             self.id,
             self.event_producer,
+            kernel_image,
             kernel_config,
             distro,
             self.local_config,

src/ai/backend/agent/kubernetes/utils.py

@@ -20,14 +20,14 @@ class PersistentServiceContainer:
     def __init__(
         self,
         docker: Docker,
-        image_ref: str,
+        image: str,
         container_config: Mapping[str, Any],
         *,
         name: Optional[str] = None,
     ) -> None:
         self.docker = docker
-        self.image_ref = image_ref
-        default_container_name = image_ref.split(":")[0].rsplit("/", maxsplit=1)[-1]
+        self.image = image
+        default_container_name = image.split(":")[0].rsplit("/", maxsplit=1)[-1]
         if name is None:
             self.container_name = default_container_name
         else:
@@ -69,7 +69,7 @@ async def get_container_version_and_status(self) -> Tuple[int, bool]:
 
     async def get_image_version(self) -> int:
         try:
-            img = await self.docker.images.inspect(self.image_ref)
+            img = await self.docker.images.inspect(self.image)
         except DockerError as e:
             if e.status == 404:
                 return 0
@@ -80,22 +80,22 @@ async def get_image_version(self) -> int:
     async def ensure_running_latest(self) -> None:
         image_version = await self.get_image_version()
         if image_version == 0:
-            log.info("PersistentServiceContainer({}): installing...", self.image_ref)
+            log.info("PersistentServiceContainer({}): installing...", self.image)
             await self.install_latest()
         elif image_version < self.img_version:
             log.info(
                 "PersistentServiceContainer({}): upgrading (v{} -> v{})",
-                self.image_ref,
+                self.image,
                 image_version,
                 self.img_version,
             )
             await self.install_latest()
         container_version, is_running = await self.get_container_version_and_status()
         if container_version == 0 or image_version != container_version or not is_running:
-            log.info("PersistentServiceContainer({}): recreating...", self.image_ref)
+            log.info("PersistentServiceContainer({}): recreating...", self.image)
             await self.recreate()
         if not is_running:
-            log.info("PersistentServiceContainer({}): starting...", self.image_ref)
+            log.info("PersistentServiceContainer({}): starting...", self.image)
             await self.start()
 
     async def install_latest(self) -> None:
@@ -112,7 +112,7 @@ async def install_latest(self) -> None:
                     stderr = await proc.stderr.read()
                 raise RuntimeError(
                     "loading the image has failed!",
-                    self.image_ref,
+                    self.image,
                     proc.returncode,
                     stderr,
                 )
@@ -130,7 +130,7 @@ async def recreate(self) -> None:
             else:
                 raise
         container_config: dict[str, Any] = {
-            "Image": self.image_ref,
+            "Image": self.image,
             "Tty": True,
             "Privileged": False,
             "AttachStdin": False,

src/ai/backend/agent/server.py

@@ -486,6 +486,7 @@ async def create_kernels(
         raw_kernel_ids: Sequence[str],
         raw_configs: Sequence[dict],
         raw_cluster_info: dict,
+        kernel_image_refs: dict[KernelId, ImageRef],
     ):
         cluster_info = cast(ClusterInfo, raw_cluster_info)
         session_id = SessionId(UUID(raw_session_id))
@@ -503,6 +504,7 @@ async def create_kernels(
                 self.agent.create_kernel(
                     session_id,
                     kernel_id,
+                    kernel_image_refs[kernel_id],
                     kernel_config,
                     cluster_info,
                     throttle_sema=throttle_sema,
@@ -583,12 +585,14 @@ async def restart_kernel(
         self,
         session_id: str,
         kernel_id: str,
+        kernel_image: ImageRef,
         updated_config: dict,
     ) -> dict[str, Any]:
         log.info("rpc::restart_kernel(s:{0}, k:{1})", session_id, kernel_id)
         return await self.agent.restart_kernel(
             SessionId(UUID(session_id)),
             KernelId(UUID(kernel_id)),
+            kernel_image,
             cast(KernelCreationConfig, updated_config),
         )
 
@@ -704,30 +708,22 @@ async def _commit(reporter: ProgressReporter) -> None:
     @collect_error
     async def push_image(
         self,
-        canonical: str,
-        architecture: str,
+        image_ref: ImageRef,
         registry_conf: ImageRegistry,
-        *,
-        is_local: bool = False,
     ) -> dict[str, Any]:
-        log.info("rpc::push_image(c:{})", canonical)
+        log.info("rpc::push_image(c:{})", image_ref.canonical)
         bgtask_mgr = self.agent.background_task_manager
 
         async def _push_image(reporter: ProgressReporter) -> None:
             await self.agent.push_image(
-                ImageRef(
-                    canonical,
-                    known_registries=["*"],
-                    is_local=is_local,
-                    architecture=architecture,
-                ),
+                image_ref,
                 registry_conf,
             )
 
         task_id = await bgtask_mgr.start(_push_image)
         return {
             "bgtask_id": str(task_id),
-            "canonical": canonical,
+            "canonical": image_ref.canonical,
         }
 
     @rpc_function