Merge pull request #167 from casperdcl/docker-cli-stream

Dockerfile

@@ -1,13 +1,12 @@
 FROM python:3.10-slim-bullseye
-
 WORKDIR /usr/src/app/
 
-RUN apt-get update && apt-get install -y gcc
-
-COPY requirements.txt ./
-
-RUN pip install --no-cache-dir -r requirements.txt --upgrade pip
+RUN apt update -qq && apt install -yqq --no-install-recommends \
+    gcc curl \
+    && curl -fsSL https://get.docker.com | sed 's/if is_wsl/if false/' | sh \
+    && rm -rf /var/lib/apt/lists/*
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . .
-
 CMD python main.py

app/core/utils.py

@@ -1,4 +1,8 @@
+import errno
 import logging
+import os
+import pty
+import signal
 import subprocess
 import xml.etree.ElementTree as ET
 from http import HTTPStatus
@@ -129,6 +133,34 @@
 ]
 
 
+def subprocess_tty(cmd, encoding="utf-8", **kwargs):
+    """`subprocess.Popen` yielding stdout lines acting as a TTY"""
+    m, s = pty.openpty()
+    p = subprocess.Popen(cmd, stdout=s, stderr=s, **kwargs)
+    os.close(s)
+
+    try:
+        for line in open(m, encoding=encoding):
+            if not line:  # EOF
+                break
+            yield line
+    except OSError as e:
+        if errno.EIO != e.errno:  # EIO also means EOF
+            raise
+    finally:
+        if p.poll() is None:
+            p.send_signal(signal.SIGINT)
+            try:
+                p.wait(10)
+            except subprocess.TimeoutExpired:
+                p.terminate()
+                try:
+                    p.wait(10)
+                except subprocess.TimeoutExpired:
+                    p.kill()
+        p.wait()
+
+
 def get_docker_client():
     return docker.from_env()
 

app/routes.py

@@ -1,5 +1,7 @@
 import json
 import logging
+import re
+import shutil
 from enum import Enum
 from http import HTTPStatus
 
@@ -40,6 +42,14 @@ class Status(Enum):
 }
 
 
+def si2float(string, default=None):
+    """e.g. "42k" -> 42000.0"""
+    try:
+        return float(string.replace("k", "e3").replace("M", "e6").replace("G", "e9"))
+    except Exception:
+        return default
+
+
 @router.get("/", response_model=schemas.HealthResponse)
 async def health():
     return schemas.HealthResponse(status=True)
@@ -226,26 +236,57 @@ def generator():
 async def generator(service_object, request):
     layers = {}
 
-    client = utils.get_docker_client()
+    try:
+        cmd = shutil.which("docker")
+        if not cmd:
+            logger.info("docker not found (did you forget to install dind?)")
+            raise FileNotFoundError
+        logger.info("which docker: %r", cmd)
+    except FileNotFoundError:
+        logger.info("docker SDK")
+        client = utils.get_docker_client()
 
-    for line in client.api.pull(
-        service_object["dockerImage"], stream=True, decode=True
-    ):
-        status = line["status"]
-        if (
-            status == Status.ALREADY_EXISTS.value
-            or status == Status.DOWNLOAD_COMPLETE.value
-            or status.startswith("Pulling from")
-            or status.startswith("Pulling fs layer")
+        for line in client.api.pull(
+            service_object["dockerImage"], stream=True, decode=True
         ):
-            continue
-
-        if "id" in line and "status" in line and line["id"] != "latest":
-            layer_id = line["id"]
-            get_progress = progress_mapping.get(Status(status), lambda _: 100)
-            layers[layer_id] = get_progress(line)
-            line["percentage"] = round(sum(layers.values()) / len(layers), 2)
-            yield json.dumps(line) + "\n"
+            status = line["status"]
+            if (
+                status == Status.ALREADY_EXISTS.value
+                or status == Status.DOWNLOAD_COMPLETE.value
+                or status.startswith("Pulling from")
+                or status.startswith("Pulling fs layer")
+            ):
+                continue
+
+            if "id" in line and "status" in line and line["id"] != "latest":
+                layer_id = line["id"]
+                get_progress = progress_mapping.get(Status(status), lambda _: 100)
+                layers[layer_id] = get_progress(line)
+                line["percentage"] = round(sum(layers.values()) / len(layers), 2)
+                yield json.dumps(line) + "\n"
+    else:
+        RE_LAYERINFO = re.compile(
+            rf"(\w+): ({'|'.join(i.value for i in Status.__members__.values())})(?:\s+(.*)B/(.*)B)?"
+        )
+        logger.info("docker pull")
+        for line in utils.subprocess_tty([cmd, "pull", service_object["dockerImage"]]):
+            if match := RE_LAYERINFO.match(line):
+                layer_id, status, current, total = match.groups()
+                get_progress = progress_mapping.get(Status(status), lambda _: 100)
+                layers[layer_id] = get_progress(
+                    {
+                        "progressDetail": {
+                            "current": si2float(current, default=0),
+                            "total": si2float(total, default=1),
+                        }
+                    }
+                )
+                yield json.dumps(
+                    {
+                        "status": status,
+                        "percentage": round(sum(layers.values()) / len(layers), 2),
+                    }
+                ) + "\n"
 
     yield json.dumps(
         {"status": Status.DOWNLOAD_COMPLETE.value, "percentage": 100}