tao-run-on-local-docker

本地 Docker

单节点执行平台，可在本地 Docker 守护进程上将 TAO 作业作为具名 Docker 容器运行。适用于开发、调试、小规模运行，以及 Agent 主机已配备所需 GPU、NVIDIA 驱动、Docker 和 NVIDIA Container Toolkit 的机器。

当数据位于 Docker 主机本地，或可通过挂载卷／云凭据访问时，请使用本地 Docker。不要将其用于远程集群调度、多节点训练或需要通过 SLURM 排队的作业。

预检

工作流必须先验证主机 GPU 运行时，再启动 Docker 作业。如果检查失败，请提示用户批准安装，运行显示的安装命令，然后重新执行预检。

# Host GPU runtime: NVIDIA driver 580, CUDA 13.0, NVIDIA Container Toolkit 1.19.0.
TAO_SKILL_BANK_ROOT="${TAO_SKILL_BANK_ROOT:-$PWD}"
SETUP_SCRIPT="${TAO_SKILL_BANK_ROOT}/skills/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh"
[ -x "$SETUP_SCRIPT" ] || SETUP_SCRIPT="${TAO_SKILL_BANK_ROOT}/platform/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh"

bash "$SETUP_SCRIPT" --backend docker --check-only || {
  echo "MISSING: TAO GPU host runtime is not ready."
  echo "After user approval, run:"
  echo "  bash \"$SETUP_SCRIPT\" --backend docker --install --yes"
  exit 1
}

# Mode 1 — direct docker (no Python). All you need is docker + the GPU runtime.
docker info >/dev/null 2>&1 || { echo "MISSING: docker daemon not reachable. Start Docker."; exit 1; }
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi >/dev/null 2>&1 || {
  echo "MISSING: NVIDIA Container Toolkit not installed/configured. See:"
  echo "  bash \"$SETUP_SCRIPT\" --backend docker --install --yes"
  exit 1
}

# Mode 2 — TAO SDK wrapper. Adds Job handles, S3 I/O wrapping, ActionWorkflow.
# Skip this block if Mode 1 is sufficient for the user's request.
# When Mode 2 is in scope, read `tao-skill-bank:tao-run-platform` for the DockerSDK
# kwarg contract, build_entrypoint, and monitoring patterns.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_docker).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_docker)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}
python -c "import docker" 2>/dev/null || {
  echo "MISSING: docker Python client not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

# DockerSDK attaches every job container to ${DOCKER_NETWORK:-tao_default}. If
# the network does not exist, container start fails instantly with
# `network <name> not found` for every create_job.
DOCKER_NETWORK_NAME="${DOCKER_NETWORK:-tao_default}"
docker network ls --format '{{.Name}}' | grep -qx "$DOCKER_NETWORK_NAME" || {
  echo "MISSING: docker network '$DOCKER_NETWORK_NAME' not found. After user approval, run:"
  echo "  docker network create $DOCKER_NETWORK_NAME"
  exit 1
}

如果检查失败，Agent 会提示用户通过 Bash 授权安装／修复，然后再继续。

凭据

除了访问 Docker 守护进程外，无需平台凭据。

可选环境变量：

• DOCKER_HOST：可选的 Docker 守护进程 URL。如果未设置，SDK 会使用 Docker Python 客户端的常规环境变量／默认套接字解析方式。
• DOCKER_NETWORK：作业容器使用的 Docker 网络。默认为 tao_default。
• DOCKER_USERNAME：镜像仓库用户名。NGC 默认使用 $oauthtoken。
• NGC_KEY：从 nvcr.io 拉取私有镜像时使用。
• HOST_SSH_PATH：当 AutoML brain 容器需要 SSH 密钥来监控远程 SLURM 子作业时，将此路径挂载到容器中。
• ACCESS_KEY、SECRET_KEY、S3_ENDPOINT_URL、S3_BUCKET_NAME：可选的兼容 S3 的存储设置，供仍需从本地容器读写云存储的作业使用。

启动前预检

生成脚本或启动容器之前：

1. 验证 Docker 守护进程可访问，且 NVIDIA 运行时能够识别 GPU。
2. 验证 Docker 主机上存在所有本地／文件数据集标注和媒体路径。
3. 对于 s3:// 数据集／结果，验证已设置 ACCESS_KEY 和 SECRET_KEY，并使用 aws s3 ls 确认能够读取相应路径。
4. 启动前验证特定模型所需的凭据，例如 HF_TOKEN。

多 GPU 和多节点

本地 Docker 不支持多节点。 作业仅在本地 Docker 守护进程所在主机上运行，不进行跨主机协调。

支持本地主机上的多 GPU，可通过 NVIDIA Container Toolkit 的 --gpus 标志实现（--gpus all 或 --gpus '"device=0,1,2,3"'）。DockerSDK.create_job(gpu_count=N) 会将参数传递给 --gpus。单主机分布式初始化使用 localhost；torchrun --nproc-per-node=N 或 PyTorch DDP 均可照常运行。

后端详情

使用 SDK 后端值 local-docker。本地后端架构没有额外的后端详情，因此大多数路由由环境和作业参数控制：

{
  "backend_type": "local-docker",
  "num_gpu": 1
}

遵循 Lepton/Brev SDK 的设计，平台／控制平面值保留在 SDK 状态和 Docker 标签中。SDK 不会将 BACKEND、HOST_PLATFORM、MONGOSECRET、DOCKER_HOST 或 DOCKER_NETWORK 注入训练容器。

容器执行

TAO SDK 本地 Docker 处理程序通过 Docker Python 客户端启动容器：

• 后端作业名称采用 SDK 处理程序使用的 tao-job-<job_id> 格式。
• 命令通常为 ["/bin/bash", "-c", "<job command>"]。
• 容器以分离模式运行。默认情况下，SDK 会保留容器，以便检查状态和日志，除非设置了 DOCKER_AUTO_REMOVE=true。
• /dev/shm 挂载为 tmpfs。
• 配置的 Docker 网络由 Docker 守护进程应用于作业容器；它不会作为进程环境变量传入。
• 启动替代容器前，会先停止并移除具有相同作业 ID 的现有容器。

对于 GPU 访问，处理程序会自动检测主机类型：

• Tegra 或 Jetson 主机使用 runtime="nvidia"，并设置 NVIDIA_VISIBLE_DEVICES 和 NVIDIA_DRIVER_CAPABILITIES=all。
• 标准 x86 主机使用 Docker device_requests 和 GPU 功能。

如果 num_gpus 为 0，则不会分配 GPU。如果 num_gpus 为 -1，则会请求所有可见 GPU。在共享开发机器上，建议明确指定 GPU 数量。

存储

本地 Docker 支持本地路径和 file:// 路径，因为容器运行在同一台 Docker 主机上。请确保规范中的每个路径都满足以下条件之一：

• 由处理程序或外围服务挂载到容器中，
• 已可从容器内部访问，或
• 是具有匹配凭据的云 URI。

对于远程／共享文件系统，建议使用负责该文件系统的平台。 例如，在集群上处理 Lustre 路径时，使用 SLURM 和 lustre:///...。

监控

• SDK 处理程序直接映射 Docker 容器状态：created -> Pending，running/restarting -> Running，paused -> Paused，退出码 0 -> Complete，非零退出码 -> Error。
• 日志通过 Docker Python 客户端直接从命名容器获取（docker logs tao-job-<job_id>）。

如果容器已退出、已终止、正在移除或无法找到，状态协调会将后端进程视为已终止。

取消

取消操作会停止命名容器。GPU 所有权由 Docker／NVIDIA runtime 管理，而非由 TAO Core 的本地 GPU 管理器管理。

可选：通过 TAO SDK 使用

如果你需要 Job 句柄、通过 SDK 的 script_runner 封装 S3 I/O，或需要跨会话持久化：

from tao_sdk.platforms.docker import DockerSDK

sdk = DockerSDK()  # reads DOCKER_HOST, NGC_KEY, S3 creds from env
job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:6.26.3-pyt',
    command='dino train -e /tmp/spec.yaml',
    gpu_count=1,
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

status = sdk.get_job_status(job.id)
logs = sdk.get_job_logs(job.id, tail=200)

这会将相同的 docker run 调用封装在 Job 句柄下，并通过 script_runner 调用入口点，以便自动从 S3 下载 inputs／将 outputs 上传到 S3。如果不需要这些功能，直接使用 docker run 即可——无需安装 SDK。

故障模式

Docker 客户端未初始化：确认已安装 Docker Python 包；如果未使用默认本地套接字，请设置 DOCKER_HOST；并确认进程可以与守护进程通信。

GPU 分配失败：请求的 GPU 不可用、未配置 NVIDIA Container Toolkit，或 Docker 守护进程无法创建设备请求。请减少 GPU 数量、等待其他作业完成，或确认主机上的 docker run --gpus ... 可以正常运行。

镜像拉取身份验证失败：为私有 nvcr.io 镜像设置有效的 NGC_KEY，或在 Docker 主机上运行 docker login nvcr.io -u '$oauthtoken'。

容器意外退出：检查 docker logs tao-job-<job_id>、已配置的 DOCKER_NETWORK，以及 SDK 操作运行器生成的命令。

容器内缺少路径：主机上的本地路径不一定会挂载到作业容器中。请使用操作运行器支持的路径约定，或通过外围服务配置显式卷。