Tao Run On Docker

Docker for NVIDIA GPU Workloads

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

The Docker execution platform: a consumer that runs a model/data skill's spec-bundle by implementing four verbs (submit/status/logs/cancel) over the docker CLI, on a local daemon or a remote GPU box via DOCKER_HOST=ssh://. The verbs (§ Execution) sit on top of the docker conventions in the rest of this file — GPU flags, mounts, NGC auth, inspection, error modes — which are the how the model/data skill defers to. Single-node only; for multi-node use SLURM or Kubernetes.

Sources: official Docker CLI reference (https://docs.docker.com/reference/cli/docker/) and NVIDIA Container Toolkit docs.

Prerequisites

1. Host GPU runtime — by default, NVIDIA driver >=580, CUDA Toolkit >=13.0, and NVIDIA Container Toolkit >=1.19.0. If the selected model's references/skill_info.yaml declares runtime_requirements.gpu_host, pass those values to tao-setup-nvidia-gpu-host instead. Model requirements override the defaults for that workflow.
2. Docker — docker --version must return ≥ 20.10. Install: https://docs.docker.com/engine/install/.
3. NGC API key for nvcr.io/* pulls. Get from https://ngc.nvidia.com/.

set -a; source /path/to/.env; set +a   # omit if already exported
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 (append --yes for non-interactive agent runs):"
  echo "  bash \"$SETUP_SCRIPT\" --backend docker --install"
  exit 1
}

docker --version
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi
[ -n "$NGC_KEY" ] || echo "NGC_KEY unset — cannot pull nvcr.io images"

If the selected model declares runtime_requirements.gpu_host, append the corresponding --min-driver-version, --min-cuda-version, and --min-container-toolkit-version values to both the check and any approved install command. Do not apply one model's override to unrelated workflows.

NGC authentication

set -a; source /path/to/.env; set +a   # omit if already exported
echo "$NGC_KEY" | docker login nvcr.io -u '$oauthtoken' --password-stdin

Persists in ~/.docker/config.json across reboots. Re-run on unauthorized errors.

Execution — the four verbs

Run a spec-bundle by implementing exactly these four verbs, mutating only the job-record. Status values are the fixed vocabulary from tao-artifacts (PENDING RUNNING COMPLETE ERROR CANCELED UNKNOWN); native docker states map below, with the raw state carried in the transition message. $BANK = ${TAO_SKILL_BANK_PATH}.

submit

1. Stage inputs via tao-data-io: it picks the storage tier and returns the mount args + compute-frame paths. Docker uses tier A (bind-mount a host dir, -v /host/data:/data) as the norm, or tier C (pass S3 creds, the container fetches). Author the spec file at <stage>/spec.yaml with those compute-frame paths.
2. Lint the assembled command — redact_secrets.py lint must pass (no inline secrets; pass creds as -e VAR with no value).
3. Open the record — this mints the id and binds results_dir BEFORE launch:

   JOB_ID=$("$BANK/scripts/tao_job_record.py" open \
     --platform docker --image "$IMAGE" \
     --network-arch "$ARCH" --action "$ACTION" \
     --storage-tier "$TIER" --results-root "$RESULTS_ROOT")

4. Launch detached, naming the container after the id so the other verbs find it (keep --rm OFF so an exited container stays inspectable):

   set -a; source /path/to/.env; set +a   # omit if already exported
   CID=$(docker run -d --name "$JOB_ID" --label "tao-job=$JOB_ID" \
     --gpus "$GPUS" --shm-size=8g \
     -v "$STAGE:/workspace" \
     -e AWS_ACCESS_KEY_ID -e AWS_SECRET_ACCESS_KEY -e HF_TOKEN -e NGC_KEY \
     "$IMAGE" <bundle command, reading /workspace/spec.yaml>)

5. Record RUNNING: "$BANK/scripts/tao_job_record.py" mark "$JOB_ID" --state RUNNING --backend-ref "$CID".

A submit that skipped step 3 has no id, so it cannot launch — that is the record-then-launch invariant.

status

read -r st code < <(docker inspect --format '{{.State.Status}} {{.State.ExitCode}}' "$JOB_ID" 2>/dev/null) || st=missing

On a terminal state, mark it — and for tier C, tao-data-io uploads results before you docker rm (the container is the only copy).

logs

docker logs --tail "${N:-200}" "$JOB_ID"    # add -f to follow in-turn

cancel

docker rm -f "$JOB_ID"
"$BANK/scripts/tao_job_record.py" mark "$JOB_ID" --state CANCELED --source agent

Local vs remote (DOCKER_HOST)

There is no separate "remote docker" — point the daemon at an SSH-reachable box: export DOCKER_HOST=ssh://user@gpu-host. Every verb above is byte-identical; the docker CLI marshals the request over SSH (reuses your key, avoids nested-quoting the command). One consequence: -v bind-mount sources then refer to the remote host's filesystem, not the launcher — stage data there (tier A) or fetch in-container (tier C). Fallback for sudo docker-only hosts: ssh host 'sudo docker …' (same skill, different prefix).

docker run — canonical flags

set -a; source /path/to/.env; set +a   # omit if already exported
HOST_RESULTS=/host/results
HOST_UID="$(id -u)"
HOST_GID="$(id -g)"
HOST_USER_NAME="$(id -un)"
[ "$HOST_UID" -ne 0 ] || { echo "Refusing writable Docker launch as UID 0" >&2; exit 1; }
HOST_IDENTITY_ARGS=(--user "$HOST_UID:$HOST_GID")
for group_id in $(id -G); do
  [ "$group_id" = "$HOST_GID" ] || HOST_IDENTITY_ARGS+=(--group-add "$group_id")
done
mkdir -p "$HOST_RESULTS/.tao-runtime/home/.cache"/{huggingface,torch,triton,torchinductor,matplotlib}

docker run \
  --gpus all \
  --rm \
  --shm-size=8g \
  "${HOST_IDENTITY_ARGS[@]}" \
  -v /host/data:/data \
  -v "$HOST_RESULTS:/results" \
  -e HOME=/results/.tao-runtime/home \
  -e USER="$HOST_USER_NAME" -e LOGNAME="$HOST_USER_NAME" \
  -e XDG_CACHE_HOME=/results/.tao-runtime/home/.cache \
  -e HF_HOME=/results/.tao-runtime/home/.cache/huggingface \
  -e TORCH_HOME=/results/.tao-runtime/home/.cache/torch \
  -e TRITON_CACHE_DIR=/results/.tao-runtime/home/.cache/triton \
  -e TORCHINDUCTOR_CACHE_DIR=/results/.tao-runtime/home/.cache/torchinductor \
  -e MPLCONFIGDIR=/results/.tao-runtime/home/.cache/matplotlib \
  -e HF_TOKEN -e NGC_KEY \
  <image> \
  <command>

Notes:

• --gpus '"device=0,1"' — select GPUs by id, not by count, on any shared host (double-quote-escaped). A count-based request resolves to the first N devices, so --gpus 1 can only ever land on GPU 0: if GPU 0 is busy, every job OOMs there while the other GPUs sit idle, and there is no way to steer it — -e NVIDIA_VISIBLE_DEVICES is overwritten by --gpus. Read current occupancy (nvidia-smi --query-gpu=index,memory.used --format=csv) and pass the free ids. Ids may also be GPU UUIDs. Without nvidia-container-toolkit: could not select device driver "" with capabilities: [[gpu]].
• --rm — clean up the container at exit; omit when you want docker logs after exit.
• --shm-size=8g — torchrun + PyTorch DataLoaders exhaust the default 64 MB /dev/shm otherwise; size it for multi-GPU training and raise (e.g. 16g) if you still hit Bus error.
• --user "$(id -u):$(id -g)" — required by default whenever a bind mount is writable. It prevents root-owned checkpoint trees that the submitting host user cannot clean up.
• Refuse UID 0 for the canonical writable-bind path. If the launcher itself is root, obtain the verified non-root submitting UID:GID explicitly; never infer it from the output-directory owner.
• --group-add <gid> — preserve supplementary host-group access to shared datasets and workspaces. The canonical array adds every host group except the primary GID.
• HOME, USER, LOGNAME, and cache redirects — keep frameworks from writing to image-owned locations such as /root after the user override. Prepare these directories on the writable mount before launch. USER/LOGNAME are load-bearing, not cosmetic: an arbitrary --user UID has no /etc/passwd entry in the image, and torch 2.x calls getpass.getuser() at import (torch/_dynamo → inductor cache-dir setup) — with neither env var set the container crashes with KeyError: 'getpwuid(): uid not found: <uid>' before any workload code runs. Any non-empty name satisfies it; the name does not need to exist in the image.
• -v host:container — bind mount; the command references container paths only.
• -e VAR — passthrough from parent shell (no value needed if already set). Use this form for secrets.

Container name collision

docker run --name X fails if a container named X already exists. Defensive pattern before reusing a name:

docker stop my-worker 2>/dev/null; docker rm my-worker 2>/dev/null
docker run --name my-worker ...

Detached + exec pattern

For multi-step workflows on the same container (download → run → post-process), avoid restart cost:

HOST_RESULTS=/host/results
HOST_UID="$(id -u)"
HOST_GID="$(id -g)"
[ "$HOST_UID" -ne 0 ] || { echo "Refusing writable Docker launch as UID 0" >&2; exit 1; }
HOST_IDENTITY_ARGS=(--user "$HOST_UID:$HOST_GID")
for group_id in $(id -G); do
  [ "$group_id" = "$HOST_GID" ] || HOST_IDENTITY_ARGS+=(--group-add "$group_id")
done
mkdir -p "$HOST_RESULTS/.tao-runtime/home/.cache"/{huggingface,torch,triton,torchinductor,matplotlib}

docker run -d --name <worker> \
  --gpus all --shm-size=8g \
  "${HOST_IDENTITY_ARGS[@]}" \
  -v <host-data>:/data \
  -v "$HOST_RESULTS:/results" \
  -e HOME=/results/.tao-runtime/home \
  -e USER="$(id -un)" -e LOGNAME="$(id -un)" \
  -e XDG_CACHE_HOME=/results/.tao-runtime/home/.cache \
  -e HF_HOME=/results/.tao-runtime/home/.cache/huggingface \
  -e TORCH_HOME=/results/.tao-runtime/home/.cache/torch \
  -e TRITON_CACHE_DIR=/results/.tao-runtime/home/.cache/triton \
  -e TORCHINDUCTOR_CACHE_DIR=/results/.tao-runtime/home/.cache/torchinductor \
  -e MPLCONFIGDIR=/results/.tao-runtime/home/.cache/matplotlib \
  --entrypoint sh \
  <image> -c "tail -f /dev/null"

docker exec <worker> <step_1>
docker exec <worker> <step_2>

docker stop <worker> && docker rm <worker>

Pull-if-missing idiom

docker image inspect <image> >/dev/null 2>&1 || docker pull <image>

Labels for discovery

Tag containers for filtered listing later:

docker run --label tao-toolkit ...
docker ps --filter 'label=tao-toolkit'

Mount patterns

The container expects its data at conventional paths defined by the image (often /data, /results, /workspace/checkpoints). The host side is arbitrary. The command inside docker run references container paths only.