openspeechapi 0.1.0__py3-none-any.whl

openspeech/local_engines/backends/docker_backend.py

@@ -0,0 +1,490 @@
+"""Docker runtime backend for local engine lifecycle."""
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+import time
+import urllib.error
+import urllib.request
+from collections import deque
+from pathlib import Path
+
+from openspeech.local_engines.base import RuntimeBackend
+from openspeech.local_engines.models import EngineSpec, EngineStatus, RuntimeConfig
+
+_PULL_PROGRESS_RE = re.compile(
+    r"([0-9]+(?:\.[0-9]+)?)\s*([kKmMgGtTpP]?i?[bB])\s*/\s*([0-9]+(?:\.[0-9]+)?)\s*([kKmMgGtTpP]?i?[bB])"
+)
+_LAYER_PREFIX_RE = re.compile(r"^([a-f0-9]{6,64}):\s*(.*)$")
+_UNIT_MULTIPLIER = {
+    "B": 1,
+    "KB": 1000,
+    "MB": 1000**2,
+    "GB": 1000**3,
+    "TB": 1000**4,
+    "KIB": 1024,
+    "MIB": 1024**2,
+    "GIB": 1024**3,
+    "TIB": 1024**4,
+}
+
+
+class DockerBackend(RuntimeBackend):
+    runtime_name = "docker"
+
+    def _require_docker(self) -> None:
+        if shutil.which("docker") is None:
+            raise RuntimeError("Docker is not installed or not found in PATH")
+
+    def _run(self, cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
+        try:
+            return subprocess.run(cmd, capture_output=True, text=True, check=check)
+        except subprocess.CalledProcessError as exc:
+            out = (exc.stdout or "").strip()
+            err = (exc.stderr or "").strip()
+            detail = "\n".join(x for x in [out, err] if x)
+            if detail:
+                raise RuntimeError(
+                    f"Command failed: {' '.join(cmd)}\n{detail}"
+                ) from exc
+            raise RuntimeError(f"Command failed: {' '.join(cmd)} (exit={exc.returncode})") from exc
+
+    def _image_exists(self, image: str) -> bool:
+        result = self._run(["docker", "image", "inspect", image], check=False)
+        return result.returncode == 0
+
+    @staticmethod
+    def _to_bytes(value: float, unit: str) -> int:
+        unit_norm = unit.strip().upper().replace("IB", "IB").replace("B", "B")
+        if unit_norm not in _UNIT_MULTIPLIER:
+            unit_norm = unit_norm.replace("I", "")
+        mul = _UNIT_MULTIPLIER.get(unit_norm, 1)
+        return int(value * mul)
+
+    @staticmethod
+    def _extract_pull_progress(line: str) -> tuple[str, int, int] | None:
+        m_layer = _LAYER_PREFIX_RE.match(line.strip())
+        if not m_layer:
+            return None
+        layer, payload = m_layer.group(1), m_layer.group(2)
+        m = _PULL_PROGRESS_RE.search(payload)
+        if not m:
+            return None
+        cur = DockerBackend._to_bytes(float(m.group(1)), m.group(2))
+        total = DockerBackend._to_bytes(float(m.group(3)), m.group(4))
+        if total <= 0:
+            return None
+        return layer, min(cur, total), total
+
+    @staticmethod
+    def _extract_pull_status(line: str) -> tuple[str, str] | None:
+        m = _LAYER_PREFIX_RE.match(line.strip())
+        if not m:
+            return None
+        return m.group(1), m.group(2)
+
+    @staticmethod
+    def _parse_platform(platform: str) -> tuple[str, str]:
+        p = (platform or "").strip()
+        if "/" in p:
+            os_name, arch = p.split("/", 1)
+            return os_name.strip() or "linux", arch.strip() or "amd64"
+        return "linux", "amd64"
+
+    def _fetch_layer_sizes(self, image: str, platform: str) -> dict[str, int]:
+        """Fetch compressed layer sizes from manifest for progress estimation."""
+        os_name, arch = self._parse_platform(platform)
+        cmd = ["docker", "manifest", "inspect", "--verbose", image]
+        proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
+        if proc.returncode != 0 or not (proc.stdout or "").strip():
+            return {}
+        try:
+            manifests = json.loads(proc.stdout)
+        except json.JSONDecodeError:
+            return {}
+        if not isinstance(manifests, list):
+            return {}
+        for entry in manifests:
+            desc = entry.get("Descriptor", {})
+            plat = desc.get("platform", {})
+            if plat.get("os") == os_name and plat.get("architecture") == arch:
+                oci = entry.get("OCIManifest", {})
+                layers = oci.get("layers", [])
+                result: dict[str, int] = {}
+                for layer in layers:
+                    digest = str(layer.get("digest", ""))
+                    size = int(layer.get("size", 0) or 0)
+                    if not digest.startswith("sha256:") or size <= 0:
+                        continue
+                    key = digest.split(":", 1)[1][:12]
+                    result[key] = size
+                return result
+        return {}
+
+    def _pull_with_progress(self, image: str, platform: str, report) -> None:
+        cmd = ["docker", "pull"]
+        if platform:
+            cmd.extend(["--platform", platform])
+        cmd.append(image)
+
+        proc = subprocess.Popen(  # noqa: S603
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=False,
+            bufsize=0,
+        )
+        if proc.stdout is None:
+            raise RuntimeError("Failed to capture docker pull output")
+
+        layer_sizes = self._fetch_layer_sizes(image, platform)
+        layer_progress: dict[str, tuple[int, int]] = {}
+        layer_state: dict[str, str] = {}
+        layer_last_emit_at: dict[str, float] = {}
+        last_emit = 0.0
+        speed_window: deque[tuple[float, int]] = deque(maxlen=12)
+        transfer_started_at = time.monotonic()
+        tail = deque(maxlen=40)
+        try:
+            pending = ""
+
+            def _handle_line(raw: str) -> None:
+                nonlocal last_emit
+                line = raw.strip()
+                if not line:
+                    return
+                tail.append(line)
+                parsed = self._extract_pull_progress(line)
+                status_line = self._extract_pull_status(line)
+                now = time.monotonic()
+
+                def _current_transfer_stats() -> tuple[float, int | None]:
+                    done = sum(c for c, t in layer_progress.values())
+                    known_total = sum(t for c, t in layer_progress.values())
+                    manifest_total = sum(layer_sizes.values()) if layer_sizes else 0
+                    all_total = manifest_total if manifest_total > 0 else known_total
+
+                    speed_window.append((now, done))
+                    speed_bps = 0.0
+                    eta_seconds: int | None = None
+                    if len(speed_window) >= 2:
+                        t0, b0 = speed_window[0]
+                        t1, b1 = speed_window[-1]
+                        dt = max(0.001, t1 - t0)
+                        db = max(0, b1 - b0)
+                        speed_bps = db / dt
+                        remaining = max(0, all_total - done)
+                        if speed_bps > 1 and all_total > 0:
+                            eta_seconds = int(remaining / speed_bps)
+                    return speed_bps, eta_seconds
+
+                if status_line is not None:
+                    layer_id, status = status_line
+                    total_size = layer_sizes.get(layer_id)
+                    if total_size is not None and layer_id not in layer_progress:
+                        layer_progress[layer_id] = (0, total_size)
+                    if "Already exists" in status or "Download complete" in status:
+                        if total_size is not None:
+                            layer_progress[layer_id] = (total_size, total_size)
+                    if layer_state.get(layer_id) != status:
+                        layer_state[layer_id] = status
+                        speed_bps, eta_seconds = _current_transfer_stats()
+                        report(
+                            "pull_image",
+                            (
+                                f"layer {layer_id}: {status} "
+                                f"({self._format_speed(speed_bps)}, ETA {self._format_eta(eta_seconds)})"
+                            ),
+                            None,
+                            eta_seconds=eta_seconds,
+                        )
+
+                if parsed is not None:
+                    layer, cur, total = parsed
+                    layer_progress[layer] = (cur, total)
+                    prev = layer_last_emit_at.get(layer, 0.0)
+                    if (now - prev) >= 1.2:
+                        speed_bps, eta_seconds = _current_transfer_stats()
+                        ratio = (cur / total) if total > 0 else 0.0
+                        report(
+                            "pull_image",
+                            (
+                                f"layer {layer}: downloading "
+                                f"{self._format_bytes(cur)}/{self._format_bytes(total)} "
+                                f"({ratio * 100:.1f}%) at {self._format_speed(speed_bps)}, "
+                                f"ETA {self._format_eta(eta_seconds)}"
+                            ),
+                            None,
+                            eta_seconds=eta_seconds,
+                        )
+                        layer_last_emit_at[layer] = now
+
+                if (now - last_emit) >= 0.4 and layer_progress:
+                    done = sum(c for c, t in layer_progress.values())
+                    known_total = sum(t for c, t in layer_progress.values())
+                    manifest_total = sum(layer_sizes.values()) if layer_sizes else 0
+                    all_total = manifest_total if manifest_total > 0 else known_total
+                    ratio = (done / all_total) if all_total > 0 else 0.0
+                    progress = 35.0 + min(55.0, ratio * 55.0)
+                    speed_bps, eta_seconds = _current_transfer_stats()
+                    completed = sum(1 for c, t in layer_progress.values() if c >= t > 0)
+                    speed_txt = self._format_speed(speed_bps)
+                    eta_txt = self._format_eta(eta_seconds)
+                    report(
+                        "pull_image",
+                        (
+                            f"Pulling {image} ... {ratio * 100:.1f}% "
+                            f"({completed}/{len(layer_progress)} layers, {speed_txt}, ETA {eta_txt})"
+                        ),
+                        progress,
+                        eta_seconds=eta_seconds,
+                    )
+                    last_emit = now
+
+            while True:
+                chunk = proc.stdout.read(4096)
+                if not chunk:
+                    break
+                text = chunk.decode("utf-8", errors="ignore").replace("\r", "\n")
+                pending += text
+                lines = pending.split("\n")
+                pending = lines.pop() if lines else ""
+                for raw_line in lines:
+                    _handle_line(raw_line)
+
+            if pending.strip():
+                _handle_line(pending)
+
+            returncode = proc.wait()
+            if returncode != 0:
+                detail = "\n".join(tail)
+                raise RuntimeError(
+                    f"Command failed: {' '.join(cmd)}\n{detail}"
+                )
+            elapsed = max(0.001, time.monotonic() - transfer_started_at)
+            done = sum(c for c, t in layer_progress.values())
+            speed = done / elapsed if done > 0 else 0.0
+            report(
+                "pull_image",
+                f"Image {image} is ready. Avg speed {self._format_speed(speed)}.",
+                90.0,
+                eta_seconds=0,
+            )
+        except Exception:
+            try:
+                proc.kill()
+            except Exception:
+                pass
+            raise
+
+    @staticmethod
+    def _format_speed(speed_bps: float) -> str:
+        if speed_bps <= 0:
+            return "-- B/s"
+        units = ["B/s", "KB/s", "MB/s", "GB/s"]
+        v = float(speed_bps)
+        idx = 0
+        while v >= 1024 and idx < len(units) - 1:
+            v /= 1024
+            idx += 1
+        return f"{v:.1f} {units[idx]}"
+
+    @staticmethod
+    def _format_eta(eta_seconds: int | None) -> str:
+        if eta_seconds is None:
+            return "--:--"
+        if eta_seconds <= 0:
+            return "00:00"
+        mins, sec = divmod(eta_seconds, 60)
+        hours, mins = divmod(mins, 60)
+        if hours > 0:
+            return f"{hours:02d}:{mins:02d}:{sec:02d}"
+        return f"{mins:02d}:{sec:02d}"
+
+    @staticmethod
+    def _format_bytes(size: int) -> str:
+        if size <= 0:
+            return "0 B"
+        units = ["B", "KB", "MB", "GB", "TB"]
+        value = float(size)
+        idx = 0
+        while value >= 1024 and idx < len(units) - 1:
+            value /= 1024
+            idx += 1
+        return f"{value:.1f} {units[idx]}"
+
+    def _effective_options(self, spec: EngineSpec, cfg: RuntimeConfig) -> dict:
+        opts = dict(spec.options)
+        opts.update(cfg.options)
+        return opts
+
+    def _ensure_workdir(self, cfg: RuntimeConfig) -> Path:
+        p = Path(cfg.work_dir).expanduser().resolve()
+        p.mkdir(parents=True, exist_ok=True)
+        return p
+
+    def _health_url(self, opts: dict, cfg: RuntimeConfig) -> str:
+        host_port = int(opts.get("host_port", 8080))
+        health_path = str(opts.get("health_path", "/health"))
+        return cfg.options.get("health_url", f"http://127.0.0.1:{host_port}{health_path}")
+
+    def install(self, spec: EngineSpec, cfg: RuntimeConfig, report) -> None:
+        opts = self._effective_options(spec, cfg)
+        image = str(opts["docker_image"])
+        platform = str(opts.get("platform", "")).strip()
+        report("check_runtime", "Checking docker runtime...", 5)
+        self._require_docker()
+
+        report("pull_image", f"Pulling image {image} ...", 20)
+        self._pull_with_progress(image, platform, report)
+        report("done", "Image pull completed.", 100)
+
+    def update(self, spec: EngineSpec, cfg: RuntimeConfig, report) -> None:
+        opts = self._effective_options(spec, cfg)
+        image = str(opts["docker_image"])
+        platform = str(opts.get("platform", "")).strip()
+        report("check_runtime", "Checking docker runtime...", 5)
+        self._require_docker()
+        report("pull_image", f"Updating image {image} ...", 30)
+        self._pull_with_progress(image, platform, report)
+        report("done", "Image update completed.", 100)
+
+    def start(self, spec: EngineSpec, cfg: RuntimeConfig, report) -> None:
+        opts = self._effective_options(spec, cfg)
+        image = str(opts["docker_image"])
+        container = str(opts["container_name"])
+        platform = str(opts.get("platform", "")).strip()
+        host_port = int(opts.get("host_port", 8080))
+        container_port = int(opts.get("container_port", 8080))
+        start_command = str(opts.get("start_command", "")).strip()
+
+        report("check_runtime", "Checking docker runtime...", 5)
+        self._require_docker()
+
+        report("prepare_dirs", "Preparing local workspace for engine...", 15)
+        work_dir = self._ensure_workdir(cfg)
+        cache_dir = work_dir / spec.name / "cache"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        report("cleanup", "Removing stale container if present...", 25)
+        self._run(["docker", "rm", "-f", container], check=False)
+
+        pull_on_start = bool(opts.get("pull_on_start", True))
+        if pull_on_start:
+            report("pull_image", f"Ensuring image {image} is present...", 35)
+            self._pull_with_progress(image, platform, report)
+        else:
+            if self._image_exists(image):
+                report("pull_image", f"Skip pull (check updates off): using local image {image}.", 35)
+            else:
+                report(
+                    "pull_image",
+                    f"Local image not found, pulling once: {image} ...",
+                    35,
+                )
+                self._pull_with_progress(image, platform, report)
+
+        report("start_container", f"Starting container {container} ...", 40)
+        cmd = [
+            "docker",
+            "run",
+            "-d",
+            "--name",
+            container,
+            "-p",
+            f"{host_port}:{container_port}",
+            "-v",
+            f"{cache_dir}:/data",
+        ]
+        if platform:
+            cmd.extend(["--platform", platform])
+        cmd.append(image)
+        if start_command:
+            cmd.extend(["sh", "-lc", start_command])
+        self._run(cmd, check=True)
+
+        report("wait_health", "Waiting for engine health check...", 70)
+        health_url = self._health_url(opts, cfg)
+        deadline = time.monotonic() + max(cfg.timeout_s, 1.0)
+        next_heartbeat = 0.0
+        while time.monotonic() < deadline:
+            now = time.monotonic()
+            if now >= next_heartbeat:
+                total = max(cfg.timeout_s, 1.0)
+                elapsed = total - max(0.0, deadline - now)
+                ratio = min(1.0, max(0.0, elapsed / total))
+                progress = 70.0 + ratio * 25.0
+                report("wait_health", f"Waiting for health endpoint {health_url} ...", progress)
+                next_heartbeat = now + 1.0
+            try:
+                with urllib.request.urlopen(health_url, timeout=3) as resp:
+                    if resp.status < 500:
+                        report("ready", f"Engine is healthy: {health_url}", 100)
+                        return
+            except urllib.error.URLError:
+                pass
+            time.sleep(1.0)
+
+        raise RuntimeError(
+            f"Timed out waiting for health check at {health_url}. "
+            f"Check logs via `openspeech engine logs --name {spec.name}`."
+        )
+
+    def stop(self, spec: EngineSpec, cfg: RuntimeConfig, report) -> None:
+        opts = self._effective_options(spec, cfg)
+        container = str(opts["container_name"])
+        report("check_runtime", "Checking docker runtime...", 10)
+        self._require_docker()
+        report("stop_container", f"Stopping container {container} ...", 60)
+        self._run(["docker", "rm", "-f", container], check=False)
+        report("done", "Engine stopped.", 100)
+
+    def status(self, spec: EngineSpec, cfg: RuntimeConfig) -> EngineStatus:
+        opts = self._effective_options(spec, cfg)
+        container = str(opts["container_name"])
+        self._require_docker()
+
+        result = self._run(
+            [
+                "docker",
+                "ps",
+                "-a",
+                "--filter",
+                f"name=^{container}$",
+                "--format",
+                "{{.Status}}",
+            ],
+            check=False,
+        )
+        raw = (result.stdout or "").strip()
+        running = raw.lower().startswith("up")
+        healthy = False
+        detail = raw or "not found"
+        if running:
+            health_url = self._health_url(opts, cfg)
+            try:
+                with urllib.request.urlopen(health_url, timeout=3) as resp:
+                    healthy = resp.status < 500
+            except urllib.error.URLError:
+                healthy = False
+        return EngineStatus(
+            engine=spec.name,
+            runtime=self.runtime_name,
+            running=running,
+            healthy=healthy,
+            detail=detail,
+            metadata={"container_name": container},
+        )
+
+    def logs(self, spec: EngineSpec, cfg: RuntimeConfig, lines: int = 100) -> str:
+        opts = self._effective_options(spec, cfg)
+        container = str(opts["container_name"])
+        self._require_docker()
+        result = self._run(
+            ["docker", "logs", "--tail", str(lines), container],
+            check=False,
+        )
+        output = (result.stdout or "") + (result.stderr or "")
+        return output.strip()