fix: Allow to use local image when registry_id is None

lablup · Oct 6, 2024 · 872fb3b · 872fb3b
1 parent 5297138
commit 872fb3b
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 15 deletions.
diff --git a/src/ai/backend/agent/agent.py b/src/ai/backend/agent/agent.py
@@ -100,7 +100,7 @@
     VolumeUnmounted,
 )
 from ai.backend.common.events_experimental import EventDispatcher as ExperimentalEventDispatcher
-from ai.backend.common.exception import VolumeMountFailed
+from ai.backend.common.exception import UnknownImageRegistry, VolumeMountFailed
 from ai.backend.common.lock import FileLock
 from ai.backend.common.plugin.monitor import ErrorPluginContext, StatsPluginContext
 from ai.backend.common.service_ports import parse_service_ports
@@ -1856,6 +1856,10 @@ async def create_kernel(
                 await self.produce_event(
                     KernelPullingEvent(kernel_id, session_id, ctx.image_ref.canonical),
                 )
+
+                if not kernel_config["image"]["registry"]:
+                    raise UnknownImageRegistry(None)
+
                 try:
                     await self.pull_image(
                         ctx.image_ref,

diff --git a/src/ai/backend/common/types.py b/src/ai/backend/common/types.py
@@ -1011,7 +1011,7 @@ class ImageConfig(TypedDict):
     architecture: str
     digest: str
     repo_digest: Optional[str]
-    registry: ImageRegistry
+    registry: Optional[ImageRegistry]
     labels: Mapping[str, str]
     is_local: bool
     auto_pull: str  # AutoPullBehavior value

diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py
@@ -31,7 +31,7 @@
 from ai.backend.common import redis_helper
 from ai.backend.common.docker import ImageRef
 from ai.backend.common.etcd import AsyncEtcd
-from ai.backend.common.exception import UnknownImageReference, UnknownImageRegistry
+from ai.backend.common.exception import UnknownImageReference
 from ai.backend.common.types import (
     AutoPullBehavior,
     BinarySize,
@@ -557,7 +557,7 @@ async def bulk_get_image_configs(
         for ref in image_refs:
             resolved_image_info = await ImageRow.resolve(db_session, [ref])
 
-            registry_info: ImageRegistry
+            registry_info: ImageRegistry | None
             if resolved_image_info.image_ref.is_local:
                 registry_info = {
                     "name": ref.registry,
@@ -567,17 +567,17 @@ async def bulk_get_image_configs(
                 }
             else:
                 if resolved_image_info.registry_id is None:
-                    raise UnknownImageRegistry(None)
-
-                url, credential = await ContainerRegistryRow.get_container_registry_info(
-                    db_session, resolved_image_info.registry_id
-                )
-                registry_info = {
-                    "name": ref.registry,
-                    "url": str(url),
-                    "username": credential["username"],
-                    "password": credential["password"],
-                }
+                    registry_info = None
+                else:
+                    url, credential = await ContainerRegistryRow.get_container_registry_info(
+                        db_session, resolved_image_info.registry_id
+                    )
+                    registry_info = {
+                        "name": ref.registry,
+                        "url": str(url),
+                        "username": credential["username"],
+                        "password": credential["password"],
+                    }
 
             image_conf: ImageConfig = {
                 "architecture": ref.architecture,