feat: Improve local type container registry handling

lablup · Oct 25, 2024 · 4dc04ff · 4dc04ff
1 parent 3924686
commit 4dc04ff
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/fixtures/manager/example-container-registries-local.json b/fixtures/manager/example-container-registries-local.json
@@ -4,8 +4,7 @@
             "id": "fe878f09-06cc-4b91-9242-4c71015cce05",
             "registry_name": "local",
             "url": "http://localhost",
-            "type": "local",
-            "project": "stable"
+            "type": "local"
         }
     ]
 }
diff --git a/src/ai/backend/agent/docker/agent.py b/src/ai/backend/agent/docker/agent.py
@@ -824,13 +824,8 @@ async def start_container(
         container_log_file_count = 5
         container_log_file_size = BinarySize(container_log_size // container_log_file_count)
 
-        if self.image_ref.is_local:
-            image = self.image_ref.short
-        else:
-            image = self.image_ref.canonical
-
         container_config: MutableMapping[str, Any] = {
-            "Image": image,
+            "Image": self.image_ref.canonical,
             "Tty": True,
             "OpenStdin": True,
             "Privileged": False,

diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py
@@ -247,6 +247,18 @@ def trimmed_digest(self) -> str:
 
     @property
     def image_ref(self) -> ImageRef:
+        if self.is_local:
+            image_name, tag = ImageRef.parse_image_tag(self.name)
+
+            return ImageRef(
+                image_name,
+                self.project,
+                tag,
+                self.registry,
+                self.architecture,
+                self.is_local,
+            )
+
         # Empty image name
         if self.project == self.image:
             image_name = ""