tao-run-on-local-docker

Par nvidia · skills

Exécution Docker locale pour les conteneurs de tâches du TAO SDK, utilisant le daemon Docker hôte et le runtime GPU NVIDIA. Utilisation

npx skills add https://github.com/nvidia/skills --skill tao-run-on-local-docker

Docker local

Plateforme d'exécution mononode qui exécute les jobs TAO en tant que conteneurs Docker nommés sur le démon Docker local. Elle est utile pour le développement, le débogage, les petites exécutions et les machines où l'hôte de l'agent dispose déjà des GPU requis, du pilote NVIDIA, de Docker et du NVIDIA Container Toolkit.

Utilisez Docker local quand les données sont locales sur l'hôte Docker ou accessibles via des volumes montés/identifiants cloud. Ne l'utilisez pas pour l'ordonnancement de clusters distants, l'entraînement multi-nœud ou les jobs nécessitant la file d'attente SLURM.

Préparation

Le workflow doit vérifier le runtime GPU de l'hôte avant de démarrer les jobs Docker. Si la vérification échoue, invitez l'utilisateur à approuver l'installation, exécutez la commande d'installation affichée et relancez la préparation.

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

Si une vérification échoue, l'agent invite l'utilisateur à autoriser l'installation/correction via Bash avant de continuer.

Identifiants

Aucun identifiant de plateforme n'est requis au-delà de l'accès au démon Docker.

Environnement optionnel :

  • DOCKER_HOST : URL optionnelle du démon Docker. Si non défini, le SDK utilise la résolution normale du client Docker Python pour l'environnement/socket par défaut.
  • DOCKER_NETWORK : Réseau Docker pour les conteneurs de jobs. Par défaut : tao_default.
  • DOCKER_USERNAME : Nom d'utilisateur du registre. Par défaut : $oauthtoken pour NGC.
  • NGC_KEY : Utilisé lors du téléchargement d'images privées depuis nvcr.io.
  • HOST_SSH_PATH : Monté dans les conteneurs du cerveau AutoML quand ils ont besoin de clés SSH pour surveiller les jobs enfants SLURM distants.
  • ACCESS_KEY, SECRET_KEY, S3_ENDPOINT_URL, S3_BUCKET_NAME : Paramètres optionnels de stockage compatible S3 pour les jobs qui lisent/écrivent toujours dans le stockage cloud depuis un conteneur local.

Préparation avant lancement

Avant de générer des scripts ou de démarrer des conteneurs :

  1. Vérifiez que le démon Docker est accessible et que le runtime NVIDIA peut voir les GPU.
  2. Vérifiez que chaque annotation de dataset local/fichier et chemin média existent sur l'hôte Docker.
  3. Pour les datasets/résultats s3://, vérifiez que ACCESS_KEY et SECRET_KEY sont définis et que les chemins exacts sont lisibles avec aws s3 ls.
  4. Vérifiez les identifiants spécifiques au modèle comme HF_TOKEN avant le lancement.

Multi-GPU et multi-nœud

Le multi-nœud n'est pas supporté sur Docker local. Un job s'exécute sur l'hôte du démon Docker local sans coordination interhost.

Le multi-GPU sur l'hôte local est supporté via le flag --gpus du NVIDIA Container Toolkit (--gpus all ou --gpus '"device=0,1,2,3"'). DockerSDK.create_job(gpu_count=N) se connecte à --gpus. L'initialisation distribuée monohost utilise localhost ; torchrun --nproc-per-node=N ou PyTorch DDP fonctionnent normalement.

Détails du backend

Utilisez la valeur de backend local-docker. Le schéma du backend local n'a pas de détails de backend supplémentaires, donc la majorité du routage est contrôlée par l'environnement et les paramètres du job :

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

Suivant la conception du SDK Lepton/Brev, les valeurs plateforme/plan de contrôle restent dans l'état du SDK et les labels Docker. Le SDK n'injecte pas BACKEND, HOST_PLATFORM, MONGOSECRET, DOCKER_HOST ou DOCKER_NETWORK dans le conteneur d'entraînement.

Exécution de conteneur

Le gestionnaire Docker local du SDK TAO démarre les conteneurs via le client Docker Python :

  • Le nom du job backend utilise la forme tao-job-<job_id> utilisée par les gestionnaires SDK.
  • La commande est généralement ["/bin/bash", "-c", "<job command>"].
  • Les conteneurs s'exécutent en détaché. Le SDK conserve les conteneurs par défaut afin que le statut et les logs restent inspectables, sauf si DOCKER_AUTO_REMOVE=true.
  • /dev/shm est monté en tmpfs.
  • Le réseau Docker configuré est appliqué par le démon Docker pour le conteneur du job ; il n'est pas passé en tant que variable d'environnement de processus.
  • Les conteneurs existants avec le même ID de job sont arrêtés et supprimés avant qu'un remplacement ne démarre.

Pour l'accès GPU, le gestionnaire détecte automatiquement le type d'hôte :

  • Les hôtes Tegra ou Jetson utilisent runtime="nvidia" plus NVIDIA_VISIBLE_DEVICES et NVIDIA_DRIVER_CAPABILITIES=all.
  • Les hôtes x86 standard utilisent Docker device_requests avec les capacités GPU.

Si num_gpus est 0, aucun GPU n'est assigné. Si num_gpus est -1, tous les GPU visibles sont demandés. Préférez les décomptes GPU explicites pour les machines de développement partagées.

Stockage

Docker local accepte les chemins locaux et file:// car le conteneur s'exécute sur le même hôte Docker. Assurez-vous que chaque chemin de la spécification est soit :

  • monté dans le conteneur par le gestionnaire ou le service environnant,
  • déjà accessible de l'intérieur du conteneur, ou
  • un URI cloud avec les identifiants correspondants.

Pour les systèmes de fichiers distants/partagés, préférez la plateforme qui possède ce système de fichiers. Par exemple, utilisez SLURM plus lustre:///... pour les chemins Lustre sur un cluster.

Surveillance

  • Le gestionnaire SDK mappe directement l'état du conteneur Docker : created → Pending, running/restarting → Running, paused → Paused, code de sortie 0 → Complete, code de sortie non nul → Error.
  • Les logs proviennent directement du conteneur nommé via le client Docker Python (docker logs tao-job-<job_id>).

Si le conteneur a quitté, est mort, est supprimé ou ne peut pas être trouvé, la réconciliation de statut traite le processus backend comme terminé.

Annulation

L'annulation arrête le conteneur nommé. La propriété GPU est gérée par Docker / le runtime NVIDIA, non par le gestionnaire GPU local de TAO Core.

Optionnel : via le SDK TAO

Si vous voulez des handles Job, l'enveloppe S3 I/O via le script_runner du SDK, ou la durabilité entre les sessions :

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)

Ceci enveloppe la même invocation docker run sous un handle Job et achemine le point d'entrée via script_runner afin que inputs/outputs soient téléchargés depuis / chargés vers S3 automatiquement. Si vous n'en avez pas besoin, utilisez simplement docker run directement — aucune installation SDK requise.

Modes de défaillance

Client Docker non initialisé : Vérifiez que le package Docker Python est installé, définissez DOCKER_HOST si vous n'utilisez pas le socket local par défaut, et confirmez que le processus peut communiquer avec le démon.

L'assignation GPU a échoué : Les GPU demandés ne sont pas disponibles, le NVIDIA Container Toolkit n'est pas configuré, ou le démon Docker ne peut pas créer les demandes de périphérique GPU. Utilisez moins de GPU, attendez qu'un autre job se termine, ou vérifiez que docker run --gpus ... fonctionne sur l'hôte.

L'authentification de tirage d'image a échoué : Définissez un NGC_KEY valide pour les images privées nvcr.io ou exécutez docker login nvcr.io -u '$oauthtoken' sur l'hôte Docker.

Le conteneur a quitté de manière inattendue : Vérifiez docker logs tao-job-<job_id>, le DOCKER_NETWORK configuré et la commande produite par le gestionnaire d'actions SDK.

Chemin manquant dans le conteneur : Un chemin local sur l'hôte n'est pas nécessairement monté dans le conteneur du job. Utilisez une convention de chemin supportée par le gestionnaire d'actions ou configurez un volume explicite via le service environnant.

Skills similaires