pingtai_agent/plans/gcp_batch_inference_agent_plan.md

GCP 批量推理 Agent 设计方案

概述

基于 Vertex AI Flex PayGo 模式的批量推理 Agent，用户传入 Cloud Storage 地址，Agent 自动提交批量作业并监控状态。

---

架构设计

认证方案

┌─────────────────────────────────────────────────────────────────┐
│                        用户操作流程                              │
├─────────────────────────────────────────────────────────────────┤
│  1. 用户在 GCP 控制台给 Agent 服务账号授权存储桶访问权限           │
│  2. 用户准备 JSONL 输入文件并上传到自己的 Cloud Storage           │
│  3. 用户调用 Agent API，传入 GCS 输入/输出路径                    │
│  4. Agent 使用固定服务账号提交批量作业                            │
│  5. Agent 监控作业状态并返回结果                                  │
└─────────────────────────────────────────────────────────────────┘

系统架构

flowchart TB
    subgraph User[用户侧]
        U1[准备 JSONL 文件]
        U2[上传到 GCS]
        U3[授权 Agent 服务账号]
    end
    
    subgraph Agent[GCP Batch Agent]
        A1[API Server]
        A2[MCP Server]
        A3[Batch Job Manager]
        A4[Job Status Monitor]
        A5[Service Account Credentials]
    end
    
    subgraph GCP[Google Cloud Platform]
        G1[Cloud Storage - 输入]
        G2[Vertex AI Batch API]
        G3[Gemini Model]
        G4[Cloud Storage - 输出]
    end
    
    U1 --> U2
    U2 --> G1
    U3 --> A5
    
    A1 --> A2
    A2 --> A3
    A3 --> G2
    G2 --> G3
    G3 --> G4
    A4 --> G2
    A5 --> A3

---

核心功能

1. 提交批量作业 - submit_batch_job

输入参数：

参数	类型	必填	说明
input_uri	string	是	输入文件 GCS 路径，如 gs://bucket/input.jsonl
output_uri	string	是	输出目录 GCS 路径，如 gs://bucket/output/
model	string	否	模型名称，默认 gemini-2.0-flash
display_name	string	否	作业显示名称

输出：

{
  "success": true,
  "job_id": "projects/xxx/locations/us-central1/batchPredictionJobs/123456",
  "job_name": "batch-job-20240305-123456",
  "state": "JOB_STATE_PENDING",
  "create_time": "2024-03-05T07:30:00Z"
}

2. 查询作业状态 - get_job_status

输入参数：

参数	类型	必填	说明
job_id	string	是	批量作业 ID

输出：

{
  "success": true,
  "job_id": "projects/xxx/locations/us-central1/batchPredictionJobs/123456",
  "state": "JOB_STATE_SUCCEEDED",
  "progress": {
    "total_count": 1000,
    "succeeded_count": 998,
    "failed_count": 2
  },
  "output_uri": "gs://bucket/output/",
  "create_time": "2024-03-05T07:30:00Z",
  "end_time": "2024-03-05T08:15:00Z"
}

3. 列出作业 - list_jobs

输入参数：

参数	类型	必填	说明
filter	string	否	过滤条件
page_size	int	否	每页数量，默认 10

4. 取消作业 - cancel_job

输入参数：

参数	类型	必填	说明
job_id	string	是	批量作业 ID

---

输入文件格式

用户需要准备符合 Vertex AI 批量推理格式的 JSONL 文件：

文本请求示例

{"request":{"contents":[{"role":"user","parts":[{"text":"What is the capital of France?"}]}]}}
{"request":{"contents":[{"role":"user","parts":[{"text":"Explain quantum computing in simple terms."}]}]}}

多模态请求示例（图片）

{"request":{"contents":[{"role":"user","parts":[{"text":"Describe this image"},{"fileData":{"mimeType":"image/jpeg","fileUri":"gs://bucket/image1.jpg"}}]}]}}

多模态请求示例（视频）

{"request":{"contents":[{"role":"user","parts":[{"text":"Summarize this video"},{"fileData":{"mimeType":"video/mp4","fileUri":"gs://bucket/video1.mp4"}}]}]}}

---

项目结构

gcp_batch_agent/
├── Dockerfile
├── README.md
├── requirements.txt
├── run_api_server.py
├── USAGE.md
└── src/
    ├── __init__.py
    └── server/
        ├── __init__.py
        ├── api_server.py      # FastAPI + MCP HTTP 端点
        ├── mcp_server.py      # MCP 工具定义
        └── core/
            ├── __init__.py
            ├── config.py      # 配置管理
            ├── credentials.py # GCP 凭证管理
            ├── batch_client.py # Vertex AI 批量推理客户端
            └── schemas.py     # 数据模型

---

环境变量配置

变量	必需	说明
GCP_PROJECT_ID	是	GCP 项目 ID
GCP_LOCATION	否	GCP 区域，默认 us-central1
GCP_SERVICE_ACCOUNT_JSON	是	服务账号 JSON 凭证（Base64 编码或文件路径）
DEFAULT_MODEL	否	默认模型，默认 gemini-2.0-flash
API_PORT	否	API 端口，默认 8000

---

API 端点

端点	方法	说明
/	GET	服务状态
/health	GET	健康检查
/mcp	POST	MCP JSON-RPC
/mcp/sse	GET/POST	MCP SSE 流式
/api/v1/batch/submit	POST	提交批量作业
/api/v1/batch/status/{job_id}	GET	查询作业状态
/api/v1/batch/list	GET	列出作业
/api/v1/batch/cancel/{job_id}	POST	取消作业

---

用户使用前提

1. 授权 Agent 服务账号访问存储桶

用户需要在 GCP 控制台执行以下操作：

# 获取 Agent 服务账号邮箱（由 Agent 管理员提供）
AGENT_SA="batch-agent@your-project.iam.gserviceaccount.com"

# 授权存储桶访问权限
gsutil iam ch serviceAccount:${AGENT_SA}:objectViewer gs://your-bucket
gsutil iam ch serviceAccount:${AGENT_SA}:objectCreator gs://your-bucket

或在 GCP 控制台：

1. 进入 Cloud Storage > 存储桶 > 权限
2. 添加主账号：Agent 服务账号邮箱
3. 角色：Storage Object Viewer + Storage Object Creator

2. 准备输入文件

按照上述 JSONL 格式准备输入文件并上传到 GCS。

---

依赖包

# Pydantic AI
pydantic-ai>=0.0.14

# MCP
mcp>=0.9.0
fastmcp>=0.1.0

# FastAPI
fastapi>=0.109.0
uvicorn[standard]>=0.27.0

# HTTP Client
aiohttp>=3.9.0

# GCP
google-cloud-aiplatform>=1.38.0
google-auth>=2.23.0

---

实现计划

Phase 1: 基础框架 ✅

• 创建项目目录结构
• 配置 Dockerfile 和 requirements.txt
• 实现 GCP 凭证管理模块

Phase 2: 核心功能 ✅

• 实现 Vertex AI 批量推理客户端
• 实现 submit_batch_job 工具
• 实现 get_job_status 工具
• 实现 list_jobs 工具
• 实现 cancel_job 工具

Phase 3: API 集成 ✅

• 实现 MCP Server 工具注册
• 实现 REST API 端点
• 添加错误处理和日志

Phase 4: 文档和测试 ✅

• 编写 README.md
• 编写 USAGE.md 使用文档
• 本地测试验证（需要 GCP 凭证）

---

调用示例

REST API 调用

# 提交批量作业
curl -X POST http://localhost:8000/api/v1/batch/submit \
  -H "Content-Type: application/json" \
  -H "api-key: your-api-key" \
  -d '{
    "input_uri": "gs://my-bucket/input.jsonl",
    "output_uri": "gs://my-bucket/output/",
    "model": "gemini-2.0-flash",
    "display_name": "my-batch-job"
  }'

# 查询作业状态
curl http://localhost:8000/api/v1/batch/status/123456 \
  -H "api-key: your-api-key"

MCP 调用

{
  "jsonrpc": "2.0",
  "id": 1,
  "method": "tools/call",
  "params": {
    "name": "submit_batch_job",
    "arguments": {
      "input_uri": "gs://my-bucket/input.jsonl",
      "output_uri": "gs://my-bucket/output/",
      "model": "gemini-2.0-flash"
    }
  }
}

---

安全考虑

1. 服务账号权限最小化：Agent 服务账号只需要 Vertex AI User 和 Storage Object Admin 权限
2. 用户存储桶隔离：每个用户的数据存储在自己的存储桶中
3. API Key 验证：所有 API 调用需要有效的 API Key
4. 日志审计：记录所有批量作业的提交和状态变更

---

注意事项

1. Flex PayGo 限制：

   • 作业可能需要排队等待资源
   • 大型作业可能需要较长时间完成

2. 存储桶权限：

   • 用户必须提前授权 Agent 服务账号
   • 输入文件必须存在且可读
   • 输出目录必须可写

3. 文件格式：

   • 输入必须是有效的 JSONL 格式
   • 每行必须符合 Vertex AI 批量推理的请求格式