dataflow_agent/toolkits/dockertool/mini_docker.py

from __future__ import annotations
import asyncio
import argparse
import io
import json
import os
import tarfile
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

import docker  # pip install docker
from docker.errors import DockerException, NotFound, APIError, ContainerError

# 使用 build_docker.py 生成的自定义镜像
DOCKER_IMAGE = "myorg/dataflow-py:3.11"
RUN_TIMEOUT = 120  # s


def _run_in_container_sync(
    file_path: Path,
    image: str = DOCKER_IMAGE,
    timeout: int = RUN_TIMEOUT,
) -> Dict[str, Any]:
    """
    旧版：只读挂载、禁网运行（保留兼容）。实际请使用 run_file_in_minidocker。
    """
    client = docker.from_env()

    try:
        client.images.get(image)
    except NotFound:
        client.close()
        return {
            "success": False,
            "return_code": -1,
            "stdout": "",
            "stderr": (
                f"image not found: {image}. 请先在宿主机预置镜像（例如 docker load），"
                f"或用 build_docker.py 构建并保存/加载。"
            ),
            "file_path": str(file_path),
        }

    container = None
    try:
        container = client.containers.run(
            image=image,
            command=["python", "/app/script.py"],
            volumes={str(file_path): {"bind": "/app/script.py", "mode": "ro"}},
            detach=True,
            stdout=True,
            stderr=True,
            network_disabled=True,
            read_only=True,
            mem_limit="512m",
            pids_limit=128,
            cpu_quota=100000,
        )

        exit_status = container.wait(timeout=timeout)
        code = exit_status.get("StatusCode", 137)
        result = {
            "success": code == 0,
            "return_code": code,
            "stdout": container.logs(stdout=True, stderr=False).decode(),
            "stderr": container.logs(stdout=False, stderr=True).decode(),
            "file_path": str(file_path),
        }
    except DockerException as e:
        result = {
            "success": False,
            "return_code": -1,
            "stdout": "",
            "stderr": f"Docker error: {e}",
            "file_path": str(file_path),
        }
    finally:
        if container:
            try:
                container.remove(force=True)
            except Exception:
                pass
        client.close()

    return result


def _make_tar_bytes(files: Dict[str, bytes]) -> bytes:
    """
    将内存文件打包为 tar 字节流，用于 put_archive。
    """
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w") as tf:
        for rel_path, data in files.items():
            ti = tarfile.TarInfo(name=rel_path)
            ti.size = len(data)
            ti.mtime = int(datetime.now().timestamp())
            tf.addfile(ti, io.BytesIO(data))
    buf.seek(0)
    return buf.read()


def run_file_in_minidocker(
    file_path: str | Path,
    image: str = DOCKER_IMAGE,
    workdir: Optional[str] = None,
    requirements: Optional[str] = None,
    save_tar_path: Optional[str] = None,
    enable_network: bool = True,
    timeout: int = RUN_TIMEOUT,
) -> Dict[str, Any]:
    """
    将 Python 文件复制进容器，联网安装依赖（可选），运行并在成功后提交镜像并保存为 tar。
    仅保存容器镜像 tar，不额外保存脚本输出文件。
    """
    src = Path(file_path).expanduser().resolve()
    if not src.exists():
        return {"success": False, "return_code": -1, "stdout": "", "stderr": f"file not found: {src}"}

    req_path: Optional[Path] = None
    if requirements:
        rp = Path(requirements).expanduser().resolve()
        if not rp.exists():
            return {"success": False, "return_code": -1, "stdout": "", "stderr": f"requirements not found: {rp}"}
        req_path = rp

    client = docker.from_env()

    try:
        client.images.get(image)
    except NotFound:
        client.close()
        return {
            "success": False,
            "return_code": -1,
            "stdout": "",
            "stderr": (
                f"image not found: {image}. 请先在宿主机预置镜像（例如 docker load），"
                f"或用 build_docker.py 构建并保存/加载。"
            ),
        }

    # 可选：挂载宿主机工作目录到 /workspace（便于用户脚本访问外部数据）。脚本本身与 requirements 会复制到容器内部。
    volumes = {}
    if workdir:
        volumes[os.path.abspath(workdir)] = {"bind": "/workspace", "mode": "rw"}

    # 通过覆盖 entrypoint 为 sh -lc，一次性执行：pip 安装（可选）+ 运行脚本。
    # 简化为直接依赖容器 stdout/stderr（docker 日志）来获取输出，并在宿主机保存一份副本。
    run_ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    out_filename = f"{Path(file_path).stem}_{run_ts}.log"
    cmd_parts = []
    if req_path:
        cmd_parts.append("python -m pip install --user -r /app/requirements.txt")
    # 直接运行脚本，输出由 docker 日志采集
    cmd_parts.append("python -u /app/script.py")
    cmd_str = " && ".join(cmd_parts)

    # 挂载仅保留用户指定的工作目录，其它文件通过 put_archive 复制进容器，确保 commit 后可用
    mount_volumes = {}
    if volumes:
        mount_volumes.update(volumes)

    artifacts_dir = Path(__file__).resolve().parent / "artifacts"
    run_logs_dir = artifacts_dir / "run_logs"
    run_logs_dir.mkdir(parents=True, exist_ok=True)

    # 改为 run(detach) + exec_run，避免 shell 行为影响退出码和输出
    # 尝试多种保活策略，确保容器处于 running
    container = None
    start_error: Optional[Exception] = None
    strategies = [
        {"entrypoint": ["tail", "-f", "/dev/null"], "command": None},
        {"entrypoint": ["/bin/sh", "-c"], "command": "sleep infinity"},
        {"entrypoint": ["python"], "command": ["-c", "import time; time.sleep(10**9)"]},
    ]
    for strat in strategies:
        try:
            container = client.containers.run(
                image=image,
                entrypoint=strat["entrypoint"],
                command=strat["command"],
                detach=True,
                volumes=mount_volumes,
                network_disabled=not enable_network,
                mem_limit="2g",
                pids_limit=512,
                cpu_quota=200000,
                working_dir="/app",
                environment={"PYTHONUNBUFFERED": "1"},
            )
            start_error = None
            break
        except (APIError, ContainerError, DockerException) as e:
            start_error = e
            continue
    if container is None:
        client.close()
        return {"success": False, "return_code": -1, "stdout": "", "stderr": f"container start error: {start_error}"}

    # 等待容器进入 running 状态，避免后续 exec_run 409 错误
    try:
        waited = 0.0
        max_wait = float(timeout)
        while waited < max_wait:
            container.reload()
            status = getattr(container, "status", None)
            if status == "running":
                break
            if status in {"exited", "dead"}:
                logs = ""
                try:
                    logs = container.logs().decode()
                except Exception:
                    pass
                try:
                    container.remove(force=True)
                except Exception:
                    pass
                client.close()
                return {"success": False, "return_code": -1, "stdout": "", "stderr": f"container not running, status={status}. logs: {logs}"}
            await_time = 0.2
            import time as _t
            _t.sleep(await_time)
            waited += await_time
        # 最后再检查一次
        container.reload()
        if getattr(container, "status", None) != "running":
            client.close()
            return {"success": False, "return_code": -1, "stdout": "", "stderr": f"container not running after wait, status={getattr(container, 'status', None)}"}
    except Exception as e:
        logs = ""
        try:
            logs = container.logs().decode()
        except Exception:
            pass
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": -1, "stdout": "", "stderr": f"container wait error: {e}. logs: {logs}"}

    combined_logs: list[str] = []
    # 创建工作目录并复制脚本/依赖到容器内部，确保 commit 后仍可用
    job_dir = "/app_job"
    try:
        _code, _out = container.exec_run(["/bin/sh", "-c", f"mkdir -p {job_dir}"])
        # 打包文件为 tar 并上传
        files: Dict[str, bytes] = {}
        files["script.py"] = src.read_bytes()
        if req_path:
            files["requirements.txt"] = req_path.read_bytes()
        tar_bytes = _make_tar_bytes(files)
        container.put_archive(job_dir, tar_bytes)
    except Exception as e:
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": -1, "stdout": "", "stderr": f"put files error: {e}"}
    # 先安装依赖（如有）
    if req_path:
        try:
            pip_exit, pip_out = container.exec_run(
                ["python", "-m", "pip", "install", "--user", "-r", f"{job_dir}/requirements.txt"],
                stream=False,
            )
            try:
                combined_logs.append(pip_out.decode())
            except Exception:
                combined_logs.append(pip_out.decode(errors="ignore"))
            if pip_exit not in (0, None):
                # 依赖安装失败
                try:
                    container.remove(force=True)
                except Exception:
                    pass
                client.close()
                return {"success": False, "return_code": pip_exit or -1, "stdout": "".join(combined_logs), "stderr": "pip install failed"}
        except Exception as e:
            try:
                container.remove(force=True)
            except Exception:
                pass
            client.close()
            return {"success": False, "return_code": -1, "stdout": "".join(combined_logs), "stderr": f"pip exec error: {e}"}

    # 运行脚本
    try:
        code, run_out = container.exec_run(["python", "-u", f"{job_dir}/script.py"], stream=False)
        try:
            combined_logs.append(run_out.decode())
        except Exception:
            combined_logs.append(run_out.decode(errors="ignore"))
        if code is None:
            code = 137
    except Exception as e:
        out = "".join(combined_logs)
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": -1, "stdout": out, "stderr": f"run exec error: {e}"}

    stdout_s = "".join(combined_logs)
    stderr_s = ""  # exec_run 已合并输出
    if code != 0 or ("Traceback (most recent call last)" in stdout_s):
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": code, "stdout": stdout_s, "stderr": stderr_s}

    # 成功提交镜像并保存 tar
    repo = f"minidocker/{src.stem}"
    tag = datetime.now().strftime("%Y%m%d-%H%M%S")
    try:
        image_obj = container.commit(repository=repo, tag=tag)
    except DockerException as e:
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": -1, "stdout": stdout_s, "stderr": f"commit error: {e}"}

    artifacts_dir = Path(__file__).resolve().parent / "artifacts"
    artifacts_dir.mkdir(exist_ok=True)
    # 将运行日志保存到本地，便于排查问题
    run_logs_dir = artifacts_dir / "run_logs"
    run_logs_dir.mkdir(parents=True, exist_ok=True)
    log_path = run_logs_dir / out_filename
    try:
        with open(log_path, "w", encoding="utf-8") as f:
            f.write(stdout_s)
            if stderr_s:
                f.write("\n[stderr]\n")
                f.write(stderr_s)
    except Exception:
        # 如果保存日志失败，不影响后续保存 tar
        pass

    tar_path = Path(save_tar_path).expanduser().resolve() if save_tar_path else artifacts_dir / f"{repo.replace('/', '_')}_{tag}.tar"
    try:
        with open(tar_path, "wb") as f:
            for chunk in image_obj.save(named=True):
                f.write(chunk)
    except Exception as e:
        try:
            container.remove(force=True)
        except Exception:
            pass
        client.close()
        return {"success": False, "return_code": -1, "stdout": stdout_s, "stderr": f"save tar error: {e}"}

    try:
        container.remove(force=True)
    except Exception:
        pass
    client.close()
    return {
        "success": True,
        "return_code": 0,
        "stdout": stdout_s,
        "stderr": stderr_s,
        "image": repo,
        "tag": tag,
        "tar_path": str(tar_path),
        "log_path": str(log_path),
    }


async def _run_py_in_docker(
    file_path: Path,
    image: str = DOCKER_IMAGE,
    timeout: int = RUN_TIMEOUT,
) -> Dict[str, Any]:
    """
    异步封装（旧版），建议换用 run_file_in_minidocker。
    """
    return await asyncio.to_thread(_run_in_container_sync, file_path, image, timeout)


def _cli_main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="在 minidocker 中运行 Python 文件并保存镜像 tar")
    parser.add_argument("--file", required=True, help="要运行的 Python 脚本路径")
    parser.add_argument("--image", default=DOCKER_IMAGE, help="基础镜像名:TAG")
    parser.add_argument("--requirements", default=None, help="requirements.txt 路径（可选）")
    parser.add_argument("--save", dest="save_tar", default=None, help="保存 tar 的目标路径（可选）")
    parser.add_argument("--workdir", default=None, help="宿主机工作目录，映射到 /workspace（可选）")
    parser.add_argument("--timeout", type=int, default=RUN_TIMEOUT, help="运行超时（秒）")
    parser.add_argument("--enable-network", action="store_true", default=True, help="启用容器网络（默认开启）")
    args = parser.parse_args(argv)

    res = run_file_in_minidocker(
        file_path=args.file,
        image=args.image,
        workdir=args.workdir,
        requirements=args.requirements,
        save_tar_path=args.save_tar,
        enable_network=args.enable_network,
        timeout=args.timeout,
    )
    print(json.dumps(res, ensure_ascii=False, indent=2))
    return 0 if res.get("success") else 1


if __name__ == "__main__":
    sys.exit(_cli_main())