tao-run-on-local-docker

本地 Docker

单独安装？ 如果此会话不是由 TAO 技能库插件初始化的，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

单节点执行平台，可在 Docker 守护进程上将 TAO 作业作为命名 Docker 容器运行。守护进程可以位于 Agent 主机本地，也可以通过 DOCKER_HOST=ssh://user@host 或 Docker 上下文连接远程守护进程。它适用于开发、调试、小规模运行，以及由本地编码 Agent 向远程 GPU 机器提交作业的工作流。

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

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

预检

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

# Host 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
}

# 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; the pin below is stamped from the 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 attaches every job container to ${DOCKER_NETWORK:-tao_default}.
# Create the network if it is missing; the operation is local and idempotent.
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 会提示用户通过 Bash 授权安装／修复，然后再继续。可通过 pip 安装的 Python 依赖项以及上文所述的 Docker 网络创建操作除外：请自动安装／创建它们，然后重新运行预检。

凭据

除了访问 Docker 守护进程外，平台不需要其他凭据。

可选环境变量：

• DOCKER_HOST：可选的 Docker 守护进程 URL。如果未设置，SDK 会使用 Docker Python 客户端的常规环境／默认套接字解析方式。使用 remote-docker 平台选项时必须设置。
• 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 Container Toolkit 已注册为 Docker 运行时、GPU 和驱动程序版本均有报告，并且启动前的冒烟测试容器能够看到 GPU。对于远程 Docker，应通过针对远程守护进程运行 docker run ... nvidia-smi 来查询 GPU；不要在 Agent 机器上使用本地 nvidia-smi。
2. 确认 Docker 主机上存在每个本地/文件数据集的标注和媒体路径。
3. 将每个绑定挂载分类为只读或可写。默认情况下，可写挂载必须使用 Docker 主机用户的数字 UID:GID，并且必须为该身份设置可写的容器身份（USER/LOGNAME）和 HOME/框架缓存路径。直接使用 Docker 时必须显式传递这些值（参见“非 root 容器身份”）；SDK 会在可写的 /results 绑定挂载下准备这些路径（仅当强制使用非 root 身份的作业没有此挂载时，才会使用隔离的 /tmp 作为备用路径）。对于远程 Docker，应在远程主机上解析身份，而不是照搬 Agent 笔记本电脑上的数字 ID。
4. 对于 s3:// 数据集/结果，确认已设置 ACCESS_KEY 和 SECRET_KEY，并且可以通过 aws s3 ls 读取准确路径。如果缺少 aws，请报告缺少的依赖并先征求安装许可；安装后重新执行启动前检查。
5. 启动前确认模型所需的凭据，例如 HF_TOKEN。
6. 使用 nvidia-smi 检查当前 GPU 占用情况；如果用户要求避开正在被其他运行作业使用的 GPU，就不要选择这些 GPU。在启动审核中显示所选 GPU ID。
7. 对于已知存在架构限制的模型/容器组合，启动前比较主机 GPU 的计算能力与容器堆栈。如果所选镜像无法为主机架构执行 JIT 或运行内核，则应尽早阻止启动，并请求用户提供兼容的镜像或平台。

尽可能使用随附的辅助工具完成这些检查：

${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 守护进程，请使用 remote-docker 平台并传入或导出 DOCKER_HOST。辅助工具会验证远程 GPU/运行时是否就绪，并通过只读绑定挂载检查 远程主机上的数据集路径：

${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
# Use the same SSH account represented by DOCKER_HOST=ssh://user@gpu-host,
# or obtain these two values from the remote administrator.
REMOTE_UID="$(ssh user@gpu-host id -u)"
REMOTE_GID="$(ssh user@gpu-host id -g)"
case "$REMOTE_UID" in
  ''|*[!0-9]*|0)
    echo "A verified non-root remote submitting UID is required."
    exit 1
    ;;
esac
case "$REMOTE_GID" in
  ''|*[!0-9]*)
    echo "A verified numeric remote submitting GID is required."
    exit 1
    ;;
esac
TAO_DOCKER_CONTAINER_USER="$REMOTE_UID:$REMOTE_GID"
export TAO_DOCKER_CONTAINER_USER

# Prove that this exact identity can create and remove a child in the 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 submitting identity cannot write/delete under $REMOTE_RESULTS."
  exit 1
}

非 root 容器身份

仅设置 --user <uid>:<gid> 还不够。TAO 镜像仅为 UID 1000 配置了非 root 账户（ubuntu 和 taotoolkituser 在此处发生冲突），因此其他所有数字 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        # must live on a writable mount
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 条目，而且仅用于缓存路径命名。不要为了绕过启动时的 getpwuid 错误而删除 --user——这样只会把启动错误换成由 root 拥有的输出；后者的修复代价更高。

多 GPU 和多节点

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

支持在本地主机上使用 NVIDIA Container Toolkit 的 --gpus 标志启用多 GPU（--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
}

遵循 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。
• 在 run_as_user=None（默认值）时，如果本地作业有一个可写的绝对路径 /results 绑定，SDK 会将其映射到调用者的 UID:GID，保留本地补充组，并在 /results/.tao-runtime/home 下准备 HOME／框架缓存。设置 run_as_user=True 可为其他本地挂载布局启用 用户映射。如果 SDK 进程本身以 root 身份运行，自动映射会安全失败，而不会映射到 0:0；请通过 container_user 提供已验证的提交者非 root UID:GID。container_user 也可显式覆盖远程主机上的非 root Docker 用户。对于已确认需要 root 权限的镜像，可通过 run_as_user=False 有意选择不进行映射。
• /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 数量。如果有明确的设备 ID，在共享机器上应优先使用设备 ID，而不是仅按数量选择，以免启动时占用其他任务正在使用的 GPU。

存储

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

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

对于绑定挂载的输出，以主机用户身份运行是启动时的必要条件，而不是用于绕过权限错误的方法。以 root 身份运行的容器通常会以 root:root 所有权和 0755 权限创建检查点子目录；即使顶层输出目录是预先创建的，主机用户也无法删除其中的文件。容器自动移除也不会影响绑定挂载的输出。

只有在能够明确证明所选镜像需要 root 权限时，才可关闭主机用户映射（run_as_user=False）。请在启动审核中记录这一例外，隔离其可写挂载，并在所有终端退出和取消操作后，将每个输出/缓存挂载的所有权恢复为 Docker 主机的 UID:GID。对于远程 Docker，请通过 container_user 传入经过验证的远程主机非 root 身份；切勿根据客户端机器或输出目录所有者推断该身份。所有权恢复成功之前，不得开始其他实验。如果 Agent 无权执行或验证此修复，则不能在本地 Docker 上启动需要 root 权限的镜像。

AutoML 的默认检查点保留策略更为严格：在启动试验之前，其预检会拒绝 run_as_user=False、命名卷、远程绑定挂载或不兼容的显式 container_user，因为 SDK 无法保证在主机端删除文件。仅当明确禁用保留策略且由外部操作员负责清理产物时，AutoML 才可使用这些方式。

对于远程/共享文件系统，优先使用管理该文件系统的平台。例如，对于集群上的 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，或需要跨会话保持状态：

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)

这会在 Job 句柄下封装相同的 docker run 调用。对于 S3 I/O，请先调用 build_entrypoint(...)，并传入其命令，以便 script_runner 执行声明的下载/上传操作。如果你不需要作业跟踪或该封装，直接使用 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 操作运行器生成的命令。

KeyError: getpwuid(): uid not found：启动时传入了 --user，但该 UID 在镜像的 /etc/passwd 中没有对应条目，且未传入 USER/LOGNAME。请添加“非 root 容器身份”中所述的身份信息和缓存环境变量；不要改为以 root 身份运行。

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

绑定挂载的结果文件归 root 所有：停止启动新实验，通过 docker inspect 找出所有可写挂载点，并让主机管理员一次性修复现有文件的所有权。今后的启动必须使用主机 UID:GID 映射，并将 HOME/缓存重定向到可写位置。docker rm 和 DOCKER_AUTO_REMOVE 不会修复或删除绑定挂载的文件。