Tool_Agent/test_docker_deployment.py

"""测试 Docker 工具部署 — 在远程 Docker 上部署 rembg 并注册"""
import httpx
import time

BASE_URL = "http://localhost:8001"

def main():
    task_spec = """在远程 Docker 上部署 rembg 图像背景移除工具（使用 git clone 方式）。

背景信息：
- GitHub: https://github.com/danielgatis/rembg
- rembg 是一个 AI 图像背景移除工具，22k+ stars
- 项目自带 HTTP 服务器，运行 'rembg s' 即可启动
- 需要部署到远程 Docker 容器（不是本地 uv）
- rembg 首次启动会下载 AI 模型到 /root/.u2net/ 目录

任务要求：
1. 使用 create_docker_env 创建 Docker 容器：
   - 镜像: python:3.12-slim
   - 端口: [7000]
   - volumes: {
       "/opt/tool_agent/docker/rembg/app": "/app",  # 代码目录
       "/opt/tool_agent/docker/rembg/models": "/root/.u2net"  # 模型目录
     }
2. 在容器内 git clone 项目：
   - cd /app
   - git clone https://github.com/danielgatis/rembg.git
   - cd rembg
3. 安装依赖：
   - 先查看 README.md 或 setup.py 了解安装方式
   - 安装项目：pip install -e .[cpu,cli]
4. 后台启动服务: rembg s --host 0.0.0.0 --port 7000
5. 等待服务启动完成（首次会下载 AI 模型，需要几分钟）
6. 使用 test_deployment 工具验证部署（容器状态、路由、连通性）
7. 验证通过后注册到工具库
8. 容器需要持久化运行，不要销毁

注册信息：
- tool_id: rembg_bg_remover
- name: Rembg Background Remover
- category: cv
- runtime_type: docker
- endpoint_path: /api/remove
- http_method: POST
- input_schema: {"type": "object", "properties": {"file": {"type": "string", "format": "binary", "description": "image file"}}, "required": ["file"]}
"""

    print("Submitting rembg Docker deployment task...")
    resp = httpx.post(
        f"{BASE_URL}/create_tool",
        json={"task_spec": task_spec, "tool_id": "rembg_bg_remover"},
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    task_id = data["task_id"]
    print(f"Task submitted: {task_id}")

    print("Waiting for task completion (polling every 30s)...")
    while True:
        time.sleep(30)
        try:
            resp = httpx.get(f"{BASE_URL}/task/{task_id}", timeout=30)
            resp.raise_for_status()
            task_data = resp.json()
            status = task_data["status"]
            print(f"[{time.strftime('%H:%M:%S')}] status: {status}")

            if status == "completed":
                print("\nDone!")
                print(task_data.get("result", ""))
                break
            elif status == "failed":
                print("\nFailed!")
                print(task_data.get("error", ""))
                break
        except Exception as e:
            print(f"[{time.strftime('%H:%M:%S')}] poll error: {e}")

if __name__ == "__main__":
    main()