在本地Docker上运行TAOSkill tao-run-on-local-docker

本技能用于在本地或远程Docker环境上执行NVIDIA TAO SDK作业容器,支持GPU加速的模型训练、微调和推理。涵盖环境预检、GPU runtime配置、容器身份映射、多GPU运行、存储挂载、监控与失败处理等关键操作。关键词:TAO SDK、Docker、本地Docker、远程Docker、GPU训练、NVIDIA Container Toolkit、TAO作业容器、模型训练、DOCKER_HOST、容器执行、GPU加速、预检、bind挂载、DockerSDK、tao-setup、视觉模型训练、失败模式。

视觉模型训练 0 次安装 0 次浏览 更新于 9/6/2026

本地 Docker

独立安装? 如果此会话不是由 TAO skill bank 插件初始化的,请先运行 tao-setup 技能(主机预检、凭据、跨技能发现)。

单节点执行平台,通过 Docker 守护进程将 TAO 作业作为命名 Docker 容器运行。该守护进程可以位于 agent 主机本地,也可以通过 DOCKER_HOST=ssh://user@host / Docker context 远程连接。它适用于开发、调试、小型运行,以及本地 coding agent 将作业提交到远程 GPU 机器的工作流。

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

当 agent 运行在工作站或笔记本电脑上,但 Docker 守护进程和 GPU 位于另一台单 GPU 服务器上时,请使用远程 Docker。在远程 Docker 模式下,specs 中的所有本地文件系统路径都解释在远程 Docker 主机上,而不是在 agent 机器上。

预检

工作流必须在启动 Docker 作业前验证主机 GPU runtime。如果检查失败,提示用户批准安装、运行打印的安装命令并重新运行预检。

# 主机 GPU runtime:NVIDIA driver 580, CUDA 13.0, NVIDIA Container Toolkit 1.19.0。
SB="${TAO_SKILL_BANK_PATH:-${TAO_SKILL_BANK_ROOT:-$PWD}}"
SETUP_SCRIPT="${SB}/skills/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
}

# 模式 1 — 直接 docker(无需 Python)。你只需要 docker + 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
}

# 模式 2 — TAO SDK wrapper。增加 Job 句柄、S3 I/O 包装、ActionWorkflow。
# 如果用户请求仅靠模式 1 就满足,可跳过本块。
# 当模式 2 属于范围时,请阅读 `tao-skill-bank:tao-run-platform` 了解 DockerSDK
# kwarg 契约、build_entrypoint 和监控模式。
# nvidia-tao-sdk 在公共 PyPI 上;下面的 pin 来自 release manifest。
PIN="nvidia-tao-sdk[docker]==7.1.0rc42"  # versions-key: wheels.tao_sdk_docker
python -c "import tao_sdk" 2>/dev/null || python -m pip install "$PIN"
python -c "import docker" 2>/dev/null || python -m pip install "$PIN"
python -c "import tao_sdk, docker"

# DockerSDK 将每个作业容器附加到 ${DOCKER_NETWORK:-tao_default}。
# 如果该网络不存在,则创建它;操作是本地且幂等的。
DOCKER_NETWORK_NAME="${DOCKER_NETWORK:-tao_default}"
docker network inspect "$DOCKER_NETWORK_NAME" >/dev/null 2>&1 || \
  docker network create "$DOCKER_NETWORK_NAME" >/dev/null

如果检查失败,agent 在执行前提示用户授权安装/修复。可通过 pip 安装的 Python 要求和上面的 Docker 网络创建是例外:自动安装/创建它们,然后重新运行预检。

凭据

除了访问 Docker 守护进程外,没有平台级凭据要求。

可选环境变量:

  • DOCKER_HOST:可选的 Docker daemon URL。如未设置,SDK 使用 Docker Python client 的正常环境/默认 socket 解析。remote-docker 平台选项需要此项。
  • DOCKER_NETWORK:作业容器使用的 Docker 网络。默认 tao_default
  • DOCKER_USERNAME:注册表用户名。NGC 默认为 $oauthtoken
  • NGC_KEY:从 nvcr.io 拉取私有镜像时使用。
  • HOST_SSH_PATH:当 AutoML brain 容器需要 SSH 密钥以监控远程 SLURM 子作业时,挂载到其中。
  • ACCESS_KEYSECRET_KEYS3_ENDPOINT_URLS3_BUCKET_NAME:用于仍从本地容器读写云存储的作业的可选 S3 兼容存储设置。

启动预检

在生成脚本或启动容器之前:

  1. 验证 Docker daemon 可达、NVIDIA Container Toolkit 已注册为 Docker runtime、GPU 和驱动版本已报告、启动前一个 smoke 容器能看到 GPU。对于远程 Docker,通过 docker run ... nvidia-smi 查询远程 daemon 的 GPU;不要使用 agent 机器上的本地 nvidia-smi
  2. 验证每个本地/文件数据集的标注和媒体路径在 Docker 主机上存在。
  3. 将每个 bind mount 分类为只读或可写。默认情况下,可写挂载必须使用 Docker 主机用户的数字 UID:GID,同时必须设置容器身份(USER/LOGNAME)以及 HOME/framework-cache 路径,并对该身份可写。直接 Docker 必须显式传递这些(见“非根容器身份”);SDK 在可写的 /results bind(对于没有这种 bind 的强制非根作业,仅降级到隔离的 /tmp)下准备它们。对于远程 Docker,请在远程主机上解析身份,而不是复制 agent 笔记本电脑的数字 ID。
  4. 对于 s3:// 数据集/结果,验证 ACCESS_KEYSECRET_KEY 已设置,并且可以用 aws s3 ls 读取确切的路径。如果缺少 aws,报告缺失依赖并询问是否安装;安装后重新运行预检。
  5. 在启动前验证模型特定凭据,如 HF_TOKEN
  6. 使用 nvidia-smi 检查当前 GPU 占用,并当用户请求该约束时,避免已被其他运行作业使用的 GPU。在启动审查中显示所选 GPU ID。
  7. 对于已知架构限制的模型/容器组合,请在启动前将主机 GPU 计算能力与容器栈进行比较。如果所选镜像不能为主机架构 JIT 或运行 kernel,请提前阻止并要求提供兼容的镜像或平台。

尽可能使用打包的 helper 进行这些检查:

${TAO_SKILL_BANK_PATH:-~/tao-skills-external}/scripts/check_tao_launch_preflight.py \
  --platform local-docker \
  --container-image "<selected-image>" \
  --path train_annotation=/abs/path/to/annotations.json \
  --path train_media=/abs/path/to/media

对于远程 Docker daemon,使用 remote-docker 平台并传递或导出 DOCKER_HOST。helper 通过只读 bind mount 验证远程 GPU/runtime 就绪,并检查远程主机数据集路径:

${TAO_SKILL_BANK_PATH:-~/tao-skills-external}/scripts/check_tao_launch_preflight.py \
  --platform remote-docker \
  --docker-host ssh://user@gpu-host \
  --container-image "<selected-image>" \
  --gpu-smoke-image ubuntu:22.04 \
  --path train_annotation=/remote/data/train/annotations.json \
  --path train_media=/remote/data/train

上面的 --path 值必须存在于远程 Docker 主机上。不要传递仅在本地笔记本电脑或 Codex 主机上存在的路径。

在远程 Docker 主机上解析实际提交用户的 UID:GID,然后将该身份显式传给 SDK。不要复用客户端笔记本电脑的 UID:GID,也不要从共享输出目录的 stat 所有权推断容器用户:该目录可能属于 root:<shared-group> 或由另一个组成员拥有。

REMOTE_RESULTS=/remote/results
# 使用 DOCKER_HOST=ssh://user@gpu-host 表示的同一 SSH 账号,或从远程管理员处获得这两个值。
REMOTE_UID="$(ssh user@gpu-host id -u)"
REMOTE_GID="$(ssh user@gpu-host id -g)"
case "$REMOTE_UID" in
  ''|*[!0-9]*|0)
    echo "需要经过验证的非 root 远程提交 UID。"
    exit 1
    ;;
esac
case "$REMOTE_GID" in
  ''|*[!0-9]*)
    echo "需要经过验证的数字远程提交 GID。"
    exit 1
    ;;
esac
TAO_DOCKER_CONTAINER_USER="$REMOTE_UID:$REMOTE_GID"
export TAO_DOCKER_CONTAINER_USER

# 证明该确切身份可以在 bind 中创建和删除子项。
docker --host "$DOCKER_HOST" run --rm \
  --user "$TAO_DOCKER_CONTAINER_USER" \
  -v "$REMOTE_RESULTS:/ownership-probe" ubuntu:22.04 \
  sh -c 'p=/ownership-probe/.tao-write-delete-probe-$$; touch "$p" && rm "$p"' || {
  echo "远程提交身份无法在 $REMOTE_RESULTS 下写入/删除。"
  exit 1
}

非根容器身份

--user <uid>:<gid> 必要但不充分。TAO 镜像仅在 UID 1000 提供非根账号(ubuntutaotoolkituser,它们在那里冲突),因此其他每个数字 UID 在 /etc/passwd 中都没有条目。这导致在 UID-1000 工作站上失败不可见,而在其他环境可复现。getpass.getuser() 读取 LOGNAME/USER/LNAME/USERNAME,只有在这些都未设置时才回退到 pwd.getpwuid(),因此在任何 TAO 代码运行之前查询就会报错:

File "/usr/lib/python3.12/getpass.py", line 169, in getuser
    return pwd.getpwuid(os.getuid())[0]
KeyError: 'getpwuid(): uid not found: 1002'

Torch 在 import 期间初始化 inductor 缓存目录时到达该调用,因此容器在启动时退出 1。对于未知 UID,Docker 也会将 HOME=/ 留空,这会把框架缓存发送到镜像拥有的路径。

因此,每个传递 --user 的直接 Docker 启动也必须传递身份和缓存环境。它们与 SDK 在 docker_handler.py 中注入的内容一致;更改时保持两个列表同步。

HOST_UID="$(id -u)"; HOST_GID="$(id -g)"
TAO_HOME=/results/.tao-runtime/home        # 必须位于可写挂载上
mkdir -p "$RESULTS_DIR/.tao-runtime/home"

docker run --rm --gpus all --ipc=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --user "$HOST_UID:$HOST_GID" \
  -e USER="$HOST_UID" -e LOGNAME="$HOST_UID" \
  -e HOME="$TAO_HOME" \
  -e XDG_CACHE_HOME="$TAO_HOME/.cache" \
  -e HF_HOME="$TAO_HOME/.cache/huggingface" \
  -e TORCH_HOME="$TAO_HOME/.cache/torch" \
  -e TRITON_CACHE_DIR="$TAO_HOME/.cache/triton" \
  -e TORCHINDUCTOR_CACHE_DIR="$TAO_HOME/.cache/torchinductor" \
  -e MPLCONFIGDIR="$TAO_HOME/.cache/matplotlib" \
  -v "$DATA_DIR:/data:ro" -v "$RESULTS_DIR:/results" -v "$SPECS_DIR:/specs:ro" \
  "$IMAGE" <action> train -e /specs/<spec>.yaml

数字 USER/LOGNAME 值是故意的:它们描述了一个实际上没有 passwd 条目的身份,并且仅用于缓存路径命名。不要去掉 --user 来解决启动时的 getpwuid 失败——那会把启动错误变成 root 拥有输出,这是更昂贵的失败。

多 GPU 和多节点

本地 Docker 不支持多节点。 一个作业运行在本地 Docker daemon 的主机上,没有跨主机协调。

通过 NVIDIA Container Toolkit 的 --gpus 标志(--gpus all--gpus '"device=0,1,2,3"')支持本地主机上的多 GPU。DockerSDK.create_job(gpu_count=N) 传递到 --gpus。单主机分布式初始化使用 localhosttorchrun --nproc-per-node=N 或 PyTorch DDP 照常工作。

后端细节

使用 SDK 后端值 local-docker。本地后端 schema 没有额外的后端细节,因此大多数字段由环境和作业参数控制:

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

遵循 Brev SDK 设计,平台/控制平面值保留在 SDK 状态和 Docker labels 中。SDK 不向训练容器注入 BACKENDHOST_PLATFORMMONGOSECRETDOCKER_HOSTDOCKER_NETWORK

容器执行

TAO SDK 本地 Docker handler 通过 Docker Python client 启动容器:

  • 后端作业名使用 SDK handlers 使用的 tao-job-<job_id> 形式。
  • 命令通常是 ["/bin/bash", "-c", "<job command>"]
  • 容器以 detach 方式运行。默认情况下,SDK 保留容器,以便状态和日志保持可检查,除非 DOCKER_AUTO_REMOVE=true
  • 使用 run_as_user=None(默认)时,当本地作业具有绝对可写的 /results bind 时,SDK 将本地作业映射到调用 UID:GID,保留本地补充组,并在 /results/.tao-runtime/home 下准备 HOME/framework caches。run_as_user=True 选择其他本地挂载布局进行用户映射。如果 SDK 进程本身是 root,则自动映射 fail closed,而不是映射 0:0;通过 container_user 提供经过验证的提交非 root UID:GID。container_user 也是远程主机的显式非 root Docker 用户覆盖。run_as_user=False 是经过验证需要 root 的镜像的有意选择。
  • /dev/shm 挂载为 tmpfs。
  • 配置的 Docker 网络由 Docker daemon 应用于作业容器;它不是作为进程环境变量传递的。
  • 在替换启动前,会停止并移除具有相同 job id 的现有容器。

对于 GPU 访问,handler 自动检测主机类型:

  • Tegra 或 Jetson 主机使用 runtime="nvidia" 以及 NVIDIA_VISIBLE_DEVICESNVIDIA_DRIVER_CAPABILITIES=all
  • 标准 x86 主机使用带有 GPU capabilities 的 Docker device_requests

如果 num_gpus0,则不分配 GPU。如果 num_gpus-1,则请求所有可见 GPU。对于共享开发机器,建议使用显式的 GPU 计数。当显式设备 ID 可用时,优先于仅计数选择,这样启动不会占用其他任务占用的 GPU。

存储

本地 Docker 接受本地和 file:// 路径,因为容器运行在同一 Docker 主机上。确保 spec 中的每个路径:

  • 由 handler 或周围服务挂载到容器中,
  • 可从容器内部到达,或
  • 是带有匹配凭据的云 URI。

对于 bind 挂载的输出,主机用户所有权是启动不变量,而不是权限错误的变通方法。Root 容器通常创建 root:root 模式 0755 的 checkpoint 子目录;即使预先创建了顶层输出目录,主机用户也无法删除其中的文件。容器自动移除也会留下 bind 挂载的输出。

只有在所选镜像被证明需要 root 时,才选择退出主机用户映射(run_as_user=False)。在启动审查中记录该例外,隔离其可写挂载,并在所有最终退出和取消后将每个输出/缓存挂载规范化回 Docker 主机 UID:GID。对于远程 Docker,通过 container_user 传递远程主机经过验证的非 root 身份;永远不要从客户端机器或输出目录所有者推断。在所有权规范化成功之前,不要开始另一个实验。如果 agent 没有权限执行或验证该修复,则无法在本地 Docker 上启动需要 root 的镜像。

AutoML 的默认 checkpoint 保留更严格:其 preflight 在启动 trial 之前拒绝 run_as_user=False、命名卷、远程 bind mounts 或不兼容的显式 container_user,因为 SDK 无法保证 host 端删除。仅当显式禁用保留且外部操作员负责 artifact 清理时,才将那些路由用于 AutoML。

对于远程/共享文件系统,优先选择拥有该文件系统的平台。例如,在集群上使用 SLURM 加 lustre:///... 来处理 Lustre 路径。

监控

  • SDK handler 直接映射 Docker 容器状态:created -> Pending, running/restarting -> Running, paused -> Paused, exit code 0 -> Complete, nonzero exit -> Error。
  • 日志直接来自命名容器,通过 Docker Python client(docker logs tao-job-<job_id>)。

如果容器已退出、死亡、被删除或找不到,状态协调会将后端进程视为已终止。

取消

取消会停止指定容器。GPU 所有权由 Docker / NVIDIA runtime 管理,而不是 TAO Core 的本地 GPU manager。

可选:通过 TAO SDK

如果你想要 Job 句柄、通过 SDK 的 script_runner 进行 S3 I/O 包装,或跨会话持久性:

import os

from tao_sdk.platforms.docker import DockerSDK

docker_host = os.environ.get('DOCKER_HOST', '')
is_remote = bool(docker_host) and not docker_host.startswith(('unix://', 'npipe://', '/'))
container_user = os.environ.get('TAO_DOCKER_CONTAINER_USER')
if is_remote and not container_user:
    raise RuntimeError('Set TAO_DOCKER_CONTAINER_USER to the remote output owner UID:GID')

sdk = DockerSDK()  # reads DOCKER_HOST, NGC_KEY, S3 creds from env
job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:7.1.0-pyt',  # versions-key: images.tao_toolkit.pyt
    command='dino train -e /data/spec.yaml',
    gpu_count=1,
    mounts=[
        {'host_path': '/host/data', 'container_path': '/data', 'read_only': True},
        {'host_path': '/host/results', 'container_path': '/results'},
    ],
    container_user=container_user,
)

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

这包装了相同的 docker run 调用,并提供 Job 句柄。对于 S3 I/O,先调用 build_entrypoint(...) 并传递其命令,使 script_runner 能够执行声明的下载/上传。如果你不需要作业跟踪或该包装,请直接使用 docker run——无需安装 SDK。

失败模式

Docker client 未初始化:验证 Docker Python 包已安装,如果你不使用默认本地 socket,请设置 DOCKER_HOST,并确认进程可以与 daemon 通信。

GPU 分配失败:请求的 GPU 不可用、NVIDIA Container Toolkit 未配置,或 Docker daemon 无法创建设备请求。使用更少的 GPU、等待另一个作业完成,或验证 docker run --gpus ... 在主机上正常工作。

镜像拉取认证失败:为私有 nvcr.io 镜像设置有效的 NGC_KEY,或在 Docker 主机上运行 docker login nvcr.io -u '$oauthtoken'

容器意外退出:检查 docker logs tao-job-<job_id>、配置的 DOCKER_NETWORK 以及 SDK action runner 生成的命令。

KeyError: getpwuid(): uid not found:启动传递了 --user,但该 UID 在镜像的 /etc/passwd 中没有条目,且未传递 USER/LOGNAME。添加“非根容器身份”中的身份和缓存环境;不要回退到以 root 运行。

容器内缺少路径:主机上的本地路径不一定挂载到作业容器中。使用 action runner 支持的路径约定,或通过周围服务配置显式卷。

Root 拥有的 bind 挂载结果:停止启动新实验,通过 docker inspect 识别每个可写挂载,并让主机管理员一次性修复现有所有权。未来的启动必须使用主机 UID:GID 映射和可写的 HOME/cache 重定向。docker rmDOCKER_AUTO_REMOVE 不会修复或删除 bind 挂载的文件。