fix: Invalid use of subproceess.Popen API with wrapped file objects (…

…#3101)

Co-authored-by: Joongi Kim <[email protected]>

changes/3101.fix.md

@@ -0,0 +1 @@
+agent not loading `krunner-extractor` image when Docker instance does not support loading XZ compressed images

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

@@ -478,8 +478,12 @@ async def prepare_krunner_env_impl(distro: str, entrypoint_name: str) -> Tuple[s
                 "ai.backend.runner", f"krunner-extractor.img.{arch}.tar.xz"
             )
             with lzma.open(extractor_archive, "rb") as reader:
-                proc = await asyncio.create_subprocess_exec(*["docker", "load"], stdin=reader)
-                if await proc.wait() != 0:
+                image_tar = reader.read()
+                proc = await asyncio.create_subprocess_exec(
+                    *["docker", "load"], stdin=asyncio.subprocess.PIPE
+                )
+                await proc.communicate(input=image_tar)
+                if proc.returncode != 0:
                     raise RuntimeError("loading krunner extractor image has failed!")
 
         log.info("checking krunner-env for {}...", distro)

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

@@ -3,7 +3,7 @@
 import logging
 import subprocess
 from pathlib import Path
-from typing import Any, BinaryIO, Mapping, Optional, Tuple, cast
+from typing import Any, Final, Mapping, Optional, Tuple
 
 import pkg_resources
 from aiodocker.docker import Docker
@@ -17,6 +17,9 @@
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
 
+IMAGE_CHUNK_SIZE: Final[int] = 1 * 1024 * 1024 * 1024  # 1MiB
+
+
 class PersistentServiceContainer:
     def __init__(
         self,
@@ -104,14 +107,19 @@ async def install_latest(self) -> None:
         with gzip.open(self.img_path, "rb") as reader:
             proc = await asyncio.create_subprocess_exec(
                 *["docker", "load"],
-                stdin=cast(BinaryIO, reader),
+                stdin=subprocess.PIPE,
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.PIPE,
             )
-            if await proc.wait() != 0:
-                stderr = b"(unavailable)"
-                if proc.stderr is not None:
-                    stderr = await proc.stderr.read()
+            assert proc.stdin is not None
+            while True:
+                chunk = reader.read(IMAGE_CHUNK_SIZE)
+                if not chunk:
+                    break
+                proc.stdin.write(chunk)
+                await proc.stdin.drain()
+            _, stderr = await proc.communicate()
+            if proc.returncode != 0:
                 raise RuntimeError(
                     "loading the image has failed!",
                     self.image_ref,