fix(BA-450): do not set timeout on docker push API execution (#3391)

changes/3391.fix.md

@@ -0,0 +1 @@
+Fix certain customized images not being pushed to registry properly

src/ai/backend/agent/agent.py

@@ -1616,7 +1616,13 @@ async def _scan_images_wrapper(self, interval: float) -> None:
         self.images = await self.scan_images()
 
     @abstractmethod
-    async def push_image(self, image_ref: ImageRef, registry_conf: ImageRegistry) -> None:
+    async def push_image(
+        self,
+        image_ref: ImageRef,
+        registry_conf: ImageRegistry,
+        *,
+        timeout: float | None | Sentinel = Sentinel.TOKEN,
+    ) -> None:
         """
         Push the given image to the given registry.
         """

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

@@ -1538,7 +1538,13 @@ async def handle_agent_socket(self):
                 else:
                     zmq_ctx.destroy()
 
-    async def push_image(self, image_ref: ImageRef, registry_conf: ImageRegistry) -> None:
+    async def push_image(
+        self,
+        image_ref: ImageRef,
+        registry_conf: ImageRegistry,
+        *,
+        timeout: float | None | Sentinel = Sentinel.TOKEN,
+    ) -> None:
         if image_ref.is_local:
             return
         auth_config = None
@@ -1554,7 +1560,10 @@ async def push_image(self, image_ref: ImageRef, registry_conf: ImageRegistry) ->
             }
 
         async with closing_async(Docker()) as docker:
-            result = await docker.images.push(image_ref.canonical, auth=auth_config)
+            kwargs: dict[str, Any] = {"auth": auth_config}
+            if timeout != Sentinel.TOKEN:
+                kwargs["timeout"] = timeout
+            result = await docker.images.push(image_ref.canonical, **kwargs)
 
             if not result:
                 raise RuntimeError("Failed to push image: unexpected return value from aiodocker")

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

@@ -31,6 +31,7 @@
     MountPermission,
     MountTypes,
     ResourceSlot,
+    Sentinel,
     ServicePort,
     SessionId,
     SlotName,
@@ -293,7 +294,13 @@ async def pull_image(
         delay = self.dummy_agent_cfg["delay"]["pull-image"]
         await asyncio.sleep(delay)
 
-    async def push_image(self, image_ref: ImageRef, registry_conf: ImageRegistry) -> None:
+    async def push_image(
+        self,
+        image_ref: ImageRef,
+        registry_conf: ImageRegistry,
+        *,
+        timeout: float | None | Sentinel = Sentinel.TOKEN,
+    ) -> None:
         delay = self.dummy_agent_cfg["delay"]["push-image"]
         await asyncio.sleep(delay)
 

src/ai/backend/agent/server.py

@@ -819,6 +819,7 @@ async def _push_image(reporter: ProgressReporter) -> None:
             await self.agent.push_image(
                 image_ref,
                 registry_conf,
+                timeout=None,
             )
 
         task_id = await bgtask_mgr.start(_push_image)