Skill: wbavon/gpu-container-setup-flagos (version 1.0.0)

--- name: gpu-container-setup-flagos description: | Automatically detect GPU vendor, find appropriate PyTorch container image, launch with correct mounts, and validate GPU functionality. Supports NVIDIA, Ascend, Metax, Iluvatar, and AMD/ROCm. Use when user says "setup container", "start pytorch container", or invokes /gpu-container-setup. user-invokable: true allowed-tools: "Bash(*) Read Edit Write Glob Grep WebSearch WebFetch AskUserQuestion" --- # GPU Container Setup Skill This skill automates multi-vendor GPU container setup for PyTorch workloads. ## Supported GPU Vendors | Vendor | PyTorch Backend | Detection | |--------|-----------------|-----------| | 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` | ## Execution Flow When invoked, follow these steps: ### Step 1: Parse Arguments Check if user provided: - `--vendor <name>` - Force specific vendor (skip detection) - `--image <image>` - Force specific container image - `--data <path>` - Force specific data mount path - `--name <name>` - Container name (default: `pytorch-gpu`) ### Step 2: Detect GPU Vendor Run the detection script: ```bash python3 .claude/skills/gpu-container-setup/scripts/detect_gpu.py ``` Expected output: ```json {"vendor": "ascend", "devices": ["Ascend 910B"], "count": 8} ``` If detection fails and no `--vendor` flag provided, ask user which vendor to use. ### Step 3: Find Data Disk Run the data disk detection: ```bash python3 .claude/skills/gpu-container-setup/scripts/find_data_disk.py ``` Expected output: ```json {"data_disk": "/mnt/data", "found": true, "size": "2.0T", "available": "1.5T"} ``` If no suitable disk found, ask user for data mount path. ### Step 4: Find Container Image Follow strict priority order (only proceed to next if current fails): ``` 1. Primary Vendor Hub (hardcoded) → 2. BAAI Harbor → 3. Web Search → 4. Local Images → 5. Ask User ``` #### Step 4.1: Primary Vendor Hub (hardcoded URLs) | Vendor | Registry | API/Query | |--------|----------|-----------| | NVIDIA | `nvcr.io` | `https://api.ngc.nvidia.com/v2/repos/nvidia/pytorch/tags` | | Ascend | `ascendhub.huawei.com` | Portal: 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` | ```bash # 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}" ``` #### Step 4.2: BAAI Harbor (fallback) Only if Step 4.1 fails (unreachable, no image, pull fails). ```bash # 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>" ``` #### Step 4.3: Web Search (fallback) Only if Steps 4.1 and 4.2 fail. Search for `"<vendor> pytorch docker official"`. #### Step 4.4: Local Images (fallback) Only if Steps 4.1-4.3 fail. Check `docker images | grep pytorch`. #### Test Before Use ```bash docker pull "${IMAGE}" && docker run --rm "${IMAGE}" python -c "import torch; print(torch.__version__)" ``` If test fails, try next source. If all fail, ask user for image. #### Step 4.5: Update Skill (self-improvement) **IMPORTANT**: If image found via Web Search (Step 4.3) passes all tests, update `references/image-sources.md` to add the newly discovered vendor hub as a primary source. This makes future lookups faster. ```bash # 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 ``` ### Step 5: Build Docker Command Refer to `references/mount-requirements.md` for vendor-specific requirements. **NVIDIA:** ```bash docker run -d --gpus all \ --name pytorch-gpu \ --shm-size=16g \ -v <data_disk>:/data \ <image> sleep infinity ``` **AMD/ROCm:** ```bash 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:** ```bash 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:** ```bash 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:** ```bash 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 ``` ### Step 6: Start Container Execute the docker run command. If container with same name exists: 1. Check if it's running - offer to use existing or replace 2. If stopped - offer to restart or replace ### Step 7: Validate PyTorch GPU Copy and run validation script inside container: ```bash docker cp .claude/skills/gpu-container-setup/scripts/validate_pytorch.py pytorch-gpu:/tmp/ docker exec pytorch-gpu python3 /tmp/validate_pytorch.py ``` Expected output: ```json { "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 } } ``` ### Step 8: Report Results Summarize to user: - GPU vendor and devices detected - Container name and image used - Data mount path - Validation status - How to access: `docker exec -it pytorch-gpu bash` ## Error Handling | Error | Action | |-------|--------| | No GPU detected | Ask user for vendor or check drivers | | Image pull fails | Try alternative registry or web search | | Container start fails | Check device permissions, show error | | Validation fails | Show detailed error, suggest fixes | ## Reference Files - `references/gpu-detection.md` - Detection methods by vendor - `references/image-sources.md` - Image discovery guide (registry APIs, priority order, selection criteria) - `references/mount-requirements.md` - Vendor mount specifications ## Example Usage ``` 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 ```


Security scan verdicts — ClawScan: suspicious | static analysis: clean | VirusTotal:  | SkillSpector: suspicious (CRITICAL)