gpu-container-setup-flagos

GPU 容器设置技能

此技能可自动为 PyTorch 工作负载设置多厂商 GPU 容器。

支持的 GPU 厂商

厂商	PyTorch 后端	检测方式
NVIDIA	CUDA	nvidia-smi
AMD	ROCm (HIP)	rocm-smi、/opt/rocm
Ascend	torch_npu	npu-smi、/usr/local/Ascend
Metax	torch_musa	mx-smi、/opt/metax
Iluvatar	torch_corex	ixsmi、/opt/iluvatar

执行流程

调用时，请按以下步骤操作：

步骤 1：解析参数

检查用户是否提供了：

• --vendor <name> - 指定厂商（跳过检测）
• --image <image> - 指定容器镜像
• --data <path> - 指定数据挂载路径
• --name <name> - 容器名称（默认值：pytorch-gpu）

步骤 2：检测 GPU 厂商

运行检测脚本：

python3 .claude/skills/gpu-container-setup/scripts/detect_gpu.py

预期输出：

{"vendor": "ascend", "devices": ["Ascend 910B"], "count": 8}

如果检测失败且未提供 --vendor 标志，请询问用户要使用哪个厂商。

步骤 3：查找数据盘

运行数据盘检测：

python3 .claude/skills/gpu-container-setup/scripts/find_data_disk.py

预期输出：

{"data_disk": "/mnt/data", "found": true, "size": "2.0T", "available": "1.5T"}

如果未找到合适的数据盘，请询问用户数据挂载路径。

步骤 4：查找容器镜像

严格按优先级顺序执行（仅在当前步骤失败时才进行下一步）：

1. Primary Vendor Hub (hardcoded) → 2. BAAI Harbor → 3. Web Search → 4. Local Images → 5. Ask User

步骤 4.1：主要厂商仓库（硬编码 URL）

厂商	镜像仓库	API/查询
NVIDIA	nvcr.io	https://api.ngc.nvidia.com/v2/repos/nvidia/pytorch/tags
Ascend	ascendhub.huawei.com	门户：https://ascendhub.huawei.com
Metax	registry.metax-tech.com	https://registry.metax-tech.com/v2/pytorch/metax-pytorch/tags/list
Iluvatar	hub.iluvatar.com	https://hub.iluvatar.com/v2/pytorch/iluvatar-pytorch/tags/list
AMD	docker.io (rocm/pytorch)	https://hub.docker.com/v2/repositories/rocm/pytorch/tags

# Example: Query NGC for latest NVIDIA PyTorch
TAG=$(curl -s "https://api.ngc.nvidia.com/v2/repos/nvidia/pytorch/tags" | jq -r '.tags[].name' | grep -E '^[0-9]{2}\.[0-9]{2}-py3$' | sort -rV | head -1)
IMAGE="nvcr.io/nvidia/pytorch:${TAG}"

步骤 4.2：BAAI Harbor（备用）

仅当步骤 4.1 失败时使用（无法访问、没有镜像或拉取失败）。

# Query BAAI Harbor
curl -s "https://harbor.baai.ac.cn/api/v2.0/projects/flagrelease-public/repositories?page_size=100" | jq -r '.[].name' | grep "flagrelease-<vendor>"

步骤 4.3：网络搜索（备用）

仅当步骤 4.1 和 4.2 均失败时，搜索 "<vendor> pytorch docker official"。

步骤 4.4：本地镜像（备用）

仅当步骤 4.1–4.3 均失败时，检查 docker images | grep pytorch。

使用前测试

docker pull "${IMAGE}" && docker run --rm "${IMAGE}" python -c "import torch; print(torch.__version__)"

如果测试失败，请尝试下一个来源。如果全部失败，请向用户索要镜像。

步骤 4.5：更新技能（自我改进）

重要提示：如果通过网络搜索（步骤 4.3）找到的镜像通过了所有测试，请更新 references/image-sources.md，将新发现的厂商仓库添加为主要来源。这能加快今后的查找速度。

# After successful web search discovery:
# 1. Verify image works (pull + pytorch test + GPU test)
# 2. Extract registry URL pattern
# 3. Update references/image-sources.md Step 1 section with new vendor hub

步骤 5：构建 Docker 命令

请参阅 references/mount-requirements.md，了解各厂商的具体要求。

NVIDIA：

docker run -d --gpus all \
  --name pytorch-gpu \
  --shm-size=16g \
  -v <data_disk>:/data \
  <image> sleep infinity

AMD/ROCm：

docker run -d \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --group-add render \
  --name pytorch-gpu \
  --shm-size=16g \
  -v <data_disk>:/data \
  <image> sleep infinity

Ascend：

docker run -d \
  --device=/dev/davinci0 --device=/dev/davinci1 ... \
  --device=/dev/davinci_manager \
  --device=/dev/devmm_svm \
  --device=/dev/hisi_hdc \
  -v /usr/local/Ascend:/usr/local/Ascend:ro \
  -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi:ro \
  --name pytorch-gpu \
  --shm-size=16g \
  -v <data_disk>:/data \
  <image> sleep infinity

Metax：

docker run -d \
  --device=/dev/mx0 --device=/dev/mx1 ... \
  -v /opt/metax:/opt/metax:ro \
  --name pytorch-gpu \
  --shm-size=16g \
  -v <data_disk>:/data \
  <image> sleep infinity

Iluvatar：

docker run -d \
  --device=/dev/bi0 --device=/dev/bi1 ... \
  -v /opt/iluvatar:/opt/iluvatar:ro \
  --name pytorch-gpu \
  --shm-size=16g \
  -v <data_disk>:/data \
  <image> sleep infinity

步骤 6：启动容器

执行 docker run 命令。如果存在同名容器：

1. 检查容器是否正在运行——询问用户是使用现有容器还是替换
2. 如果容器已停止——询问用户是重新启动还是替换

步骤 7：验证 PyTorch GPU

将验证脚本复制到容器内并运行：

docker cp .claude/skills/gpu-container-setup/scripts/validate_pytorch.py pytorch-gpu:/tmp/
docker exec pytorch-gpu python3 /tmp/validate_pytorch.py

预期输出：

{
  "status": "PASS",
  "backend": "npu",
  "device_count": 8,
  "device_names": ["Ascend 910B", ...],
  "tests": {
    "device_detection": true,
    "tensor_creation": true,
    "matrix_multiply": true,
    "gpu_to_cpu_transfer": true
  }
}

步骤 8：报告结果

向用户总结：

• 检测到的 GPU 厂商和设备
• 使用的容器名称和镜像
• 数据挂载路径
• 验证状态
• 访问方式：docker exec -it pytorch-gpu bash

错误处理

错误	操作
未检测到 GPU	询问用户 GPU 厂商或检查驱动
镜像拉取失败	尝试备用镜像仓库或网络搜索
容器启动失败	检查设备权限并显示错误
验证失败	显示详细错误并提供修复建议

参考文件

• references/gpu-detection.md - 各厂商的检测方法
• references/image-sources.md - 镜像发现指南（镜像仓库 API、优先级顺序、选择标准）
• references/mount-requirements.md - 厂商挂载规范

示例用法

User: /gpu-container-setup
User: setup a pytorch container
User: start container with ascend GPU
User: /gpu-container-setup --image nvcr.io/nvidia/pytorch:24.01-py3
User: /gpu-container-setup --image harbor.baai.ac.cn/flagrelease-public/ngctorch:2601