Agent: support CPU-only mode in parallel_extract.py

[luotao1]

PR Category

Feature Enhancement

Description

问题

在没有这些修改的情况下，parallel_extract.py 可以跑 CPU，但受限制：

1. 自动触发：只有当 torch.cuda.is_available() 返回 False 时（即机器没有 GPU 或 CUDA 未安装），才会自动进入 CPU 模式
2. 无法强制切换：如果机器有 GPU，当前代码会强制使用 GPU，不能通过 --gpus "" 参数切换到纯 CPU 模式。因为 get_gpu_ids() 对空字符串的处理是 fallback 到 _get_default_gpus()，返回 [0]
3. 变通方法：可以通过设置环境变量 CUDA_VISIBLE_DEVICES="" 来隐藏 GPU，让 torch 检测不到 CUDA，从而进入 CPU 模式

因此，如果需要在有 GPU 的机器上主动选择 CPU-only 模式，就必须做这些修改。

修改内容

• --cpu-workers 触发 CPU-only 模式：当 --cpu-workers 被设置时（且值 > 0），强制进入 CPU-only 模式，--gpus 参数不再生效
• 无 CUDA 时自动降级：torch.cuda.is_available() 返回 False 时自动进入 CPU-only 模式
• 默认 worker 数取 CPU 核心数的一半：避免打满 CPU 导致系统过载（每个 worker 是重负载进程，涉及模型加载 + 图抽取），用户可通过 --cpu-workers 自定义
• 重命名参数 --num-workers → --cpu-workers：参数语义更明确，专用于 CPU-only 模式下的 worker 数量
• 修正 worker 计数和启动日志：正确显示 CPU worker 数量，日志中区分 GPU/CPU worker 类型

使用方式

# CPU-only 模式（有 GPU 的机器上主动选择）
python graph_net/agent/parallel_extract.py \
    --model-list /tmp/test_one_model.txt \
    --workspace /tmp/graphnet_test_cleanup \
    --cpu-workers 1

# GPU 模式（默认行为，自动检测可用 GPU）
python graph_net/agent/parallel_extract.py \
    --model-list /tmp/test_one_model.txt \
    --workspace /tmp/graphnet_test_cleanup \
    --gpus 0,1

# 输出
[INFO] CPU-only mode: 1 workers
[INFO] Total models: 1, CPU workers: 1
[START] 2026-05-19 20:31:19 — launching 1 CPU workers
[Worker-0 CPU] Started (extract_timeout=1200s, verify_timeout=600s)
[Worker-0 CPU] OK sshleifer/tiny-gpt2 (37.4s)

[paddle-bot]

Thanks for your contribution!