Generated banner for Tao Run On Kubernetes
Generated logo for Tao Run On Kubernetes

Tao Run On Kubernetes

NVIDIAOfficialCreated by NVIDIA Corporation

  • Skill
  • Tool integration
  • Cloud and DevOps
  • Apache-2.0
  • Medium risk
  • Has scripts

From the author

Kubernetes execution platform — submits TAO container jobs as k8s Jobs with NVIDIA GPU scheduling; single-pod for one node, Indexed Jobs for multi-node distributed training. Use when running on EKS / GKE / AKS / on-prem clusters with the NVIDIA GPU Operator installed, or when integrating TAO into an existing k8s-native ML platform.
— the description by NVIDIA Corporation

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

The tao-run-on-kubernetes skill provides a platform-native way to run TAO container jobs as Kubernetes Jobs. It supports both single-pod execution and multi-node distributed training via Indexed Jobs. The skill handles GPU resource gating, storage tiering, and job lifecycle management using standard kubectl commands, making it suitable for integration into EKS, GKE, AKS, or on-premise clusters equipped with the NVIDIA GPU Operator.

What it does

  • Submits TAO jobs as standard Kubernetes Jobs with GPU resource requests.
  • Supports multi-node distributed training using Indexed Jobs and headless services.
  • Performs preflight checks for GPU runtime, cluster connectivity, and GPU operator availability.
  • Manages job lifecycle including submission, status tracking, log retrieval, and cancellation.
  • Integrates with S3 or PVC-based storage tiers for dataset and result handling.
  • Provides a mechanism for secure credential injection via Kubernetes secrets.

When to use it

  • When you need to run TAO training or inference workloads on a Kubernetes cluster.
  • When integrating TAO into an existing Kubernetes-native ML platform.
  • When performing multi-node distributed training that requires coordinated GPU scheduling.
  • When you prefer using standard Kubernetes primitives like Jobs and Secrets over proprietary SDKs.

Advantages

  • Native Kubernetes integration using standard kubectl workflows.
  • Flexible support for both single-node and distributed multi-node training.
  • Hard-fails on insufficient GPU capacity to prevent pending jobs.
  • Compatible with major cloud providers (EKS, GKE, AKS) and on-premise clusters.
  • No requirement for additional proprietary SDKs within the cluster.

Usage examples

“Submit a single-node TAO training job to my EKS cluster.”
The agent validates GPU capacity, mounts the required PVC, submits the job via kubectl, and returns the job ID for tracking.
“Run a 4-node distributed training job on my Kubernetes cluster.”
The agent renders an Indexed Job manifest, provisions a headless service for rendezvous, and monitors the job status across all nodes.

Install

  1. Download the skill as a ZIP:

    https://autogineer.com/agent-skills/nvidia/tao-run-on-kubernetes/download
  2. Unzip it into the cross-client skills folder of your project:

    mkdir -p .agents/skills && unzip tao-run-on-kubernetes.zip -d .agents/skills

Documentation by NVIDIA Corporation

Kubernetes

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).

Submits TAO container jobs as Kubernetes Jobs. Works on any cluster reachable via kubeconfig (EKS / GKE / AKS / on-prem) or in-cluster service account (when running inside a pod).

Single-pod by default; opt into multi-node distributed training via num_nodes > 1 (uses Indexed Job + headless Service, see Multi-node training below).

Preflight

Three checks: GPU host runtime ready, cluster reachable via kubectl, GPU Operator/device plugin present.

# 0. GPU node host runtime.
# Run this on each self-managed GPU worker node or in the node image build.
# Set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1 only when using managed GPU nodes whose
# driver/toolkit lifecycle is owned by the cloud provider or GPU Operator policy.
if [ "${TAO_K8S_SKIP_NODE_RUNTIME_CHECK:-0}" != "1" ]; then
  TAO_SKILL_BANK_ROOT="${TAO_SKILL_BANK_ROOT:-$PWD}"
  SETUP_SCRIPT="${TAO_SKILL_BANK_ROOT}/skills/platform/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh"

  bash "$SETUP_SCRIPT" --backend kubernetes --check-only || {
    echo "MISSING: TAO Kubernetes GPU node runtime is not ready."
    echo "For self-managed GPU nodes, run after user approval:"
    echo "  bash \"$SETUP_SCRIPT\" --backend kubernetes --install --yes"
    echo "For managed clusters, verify the node image/GPU Operator policy installs driver 580 and toolkit 1.19.0, then set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1."
    exit 1
  }
fi

# 1. Cluster reachable (kubeconfig OR in-cluster service account)
command -v kubectl >/dev/null 2>&1 || {
  echo "MISSING: kubectl not found on PATH. Install kubectl to submit Jobs."
  exit 1
}
kubectl cluster-info >/dev/null 2>&1 || {
  echo "MISSING: no reachable cluster (kubeconfig at ~/.kube/config, \$KUBECONFIG, or in-pod service account)."
  echo "Configure kubectl for your cluster, or set \$KUBECONFIG:"
  echo "  EKS: aws eks update-kubeconfig --name <cluster> --region <region>"
  echo "  GKE: gcloud container clusters get-credentials <cluster> --region <region>"
  echo "  AKS: az aks get-credentials --resource-group <rg> --name <cluster>"
  echo "  local: minikube start   (see 'Local cluster' below)"
  exit 1
}

# 2. NVIDIA GPU Operator present (soft check — warn, don't fail)
gpu=$(kubectl get nodes -o jsonpath='{range .items[*]}{.status.allocatable.nvidia\.com/gpu}{"\n"}{end}' 2>/dev/null | grep -v '^$' | head -1)
if [ -z "$gpu" ] || [ "$gpu" = "0" ]; then
  echo "WARN: no nvidia.com/gpu allocatable on this cluster."
  echo "Install the NVIDIA GPU Operator before submitting GPU jobs:"
  echo "  https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html"
fi

The GPU node runtime check is mandatory for self-managed nodes. For managed clusters where the client is not running on a GPU worker, verify the provider node image or GPU Operator policy and set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1 instead of running the installer on the client. The GPU-capacity warning here is a soft check; the submit verb re-checks allocatable nvidia.com/gpu and hard-fails before applying the manifest (there is no gang scheduling, so a too-big Job would sit Pending forever).

Credentials & configuration

  • Kubeconfig (one of):

    • ~/.kube/config — default discovery path

    • $KUBECONFIG — alternate path

    • In-cluster service account — used when running inside a pod (no kubeconfig needed)

  • TAO_K8S_NAMESPACE (optional): default namespace for Job submission. Defaults to default.

  • TAO_K8S_CONTEXT (optional): kubeconfig context name to switch clusters.

  • NGC_KEY (optional): for nvcr.io image pulls. If you've pre-created an image-pull secret in the target namespace, reference its name in the rendered manifest's imagePullSecrets.

  • AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL (optional): for S3 dataset I/O (storage tier C), injected into the pod via the per-job Secret (envFrom.secretRef), never inline. Legacy ACCESS_KEY/SECRET_KEY are mapped by tao-data-io.

Do not ask for Brev or SLURM credentials for Kubernetes runs. Ask for S3 credentials only when the selected workflow uses s3:// inputs or outputs, and ask for model-specific credentials such as HF_TOKEN only when the selected model requires them. Before launch, verify the selected namespace can create Jobs, dataset/result paths are visible from the pod, and PVC/mounted filesystem paths are proven to be mounted into the job container; an agent-host local path is not sufficient proof.

Execution — the four verbs

tao-run-on-kubernetes is a platform consumer: it runs a spec-bundle via kubectl, mutating only the job-record. No nvidia-tao-sdk, no tao_sdk import — jobs are submitted with plain kubectl apply. $BANK = ${TAO_SKILL_BANK_PATH}.

submit
  1. GPU-capacity gate — hard-fail first (no gang scheduling → a too-big Job sits Pending forever):

    ALLOC=$(kubectl get nodes -o jsonpath='{range .items[*]}{.status.allocatable.nvidia\.com/gpu}{"\n"}{end}' | awk '{s+=$1} END{print s+0}')
    [ "${ALLOC:-0}" -ge "$NUM_GPUS" ] || { echo "insufficient GPUs: need $NUM_GPUS, allocatable $ALLOC"; exit 1; }
  2. Storage tier (via tao-data-io): A = mount a bound PVC/NFS holding the data (author the mount paths, no fetch — the air-gap answer, and what the packaged template does); C = ephemeral: an initContainer fetches from S3 into a shared emptyDir and a final step uploads results to S3 before TTL.

    Tier C holds the GPU while it downloads. A pod reserves nvidia.com/gpu for its whole lifetime, initContainers included, so a large tier-C fetch — or a first-time multi-GB image pull — is billed and reaper-eligible idle GPU time, exactly like pulling inside a SLURM allocation. Prefer tier A when the data is already on a PVC; choose tier C knowingly, for small inputs.

    A producer action request may declare several mounts, including duplicate- source aliases for logical and embedded absolute paths. Read references/action-request.md and use its staging map plus packaged renderer. For mode=config, materialize the producer's nested spec first, stage that exact generated file, and pass it back to the renderer so the pod receives a verified read-only config mount. The legacy single-root template cannot represent that contract.

  3. Open the record — mints the id, binds results_dir, before launch:

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

    results_dir must be a mounted (surviving) volume path or an S3 prefix — ttlSecondsAfterFinished deletes the Job and its logs after it ends, so nothing is recoverable from the Job object later.

  4. Render, gate, apply, and record RUNNING. For a producer action request, follow references/action-request.md; it owns backend-name normalization, conditional Secret references, native argv rendering, server dry-run, and binding the applied object name to the job-record. For a simple one-root spec-bundle, render templates/k8s/single-pod-job.yaml.tmpl, run redact_secrets.py lint plus kubectl apply --dry-run=server, apply it, and mark the record with backend-ref=<namespace>/<actual-object-name>.

A submit that skipped the gate or the open has no id — so it cannot launch.

status

Keep K8S_JOB_NAME from submit. On reattach, read the job-record's backend_ref=<namespace>/<name> and recover both values from that field; do not assume the Kubernetes name equals the record id.

kubectl get job "$K8S_JOB_NAME" -n "$NAMESPACE" \
  -o jsonpath='{.status.conditions[0].type} {.status.active} {.status.succeeded} {.status.failed}'
kubectl signalvocab
no pods scheduledPENDING (kubectl get pods -n "$NAMESPACE" -l job-name="$K8S_JOB_NAME" → ImagePullBackOff / Insufficient nvidia.com/gpu in message)
active ≥ 1RUNNING
condition CompleteCOMPLETE
condition FailedERROR (classify from the pod's terminated reason — OOMKilled → ERR_INFRA)
Job/pod not foundUNKNOWN (may be TTL-deleted — the job-record is the source of truth)
logs
kubectl logs -n "$NAMESPACE" -l "job-name=$K8S_JOB_NAME" --tail "${N:-200}"
cancel
kubectl delete job "$K8S_JOB_NAME" -n "$NAMESPACE" --cascade=foreground
if [ -n "${CRED_SECRET:-}" ]; then
  kubectl delete secret "$CRED_SECRET" -n "$NAMESPACE" --ignore-not-found
fi
"$BANK/scripts/tao_job_record.py" mark "$JOB_ID" --state CANCELED --source agent
Multi-node (nodes > 1)

Same four verbs, plus:

  1. Version gate: require k8s ≥ 1.28 (kubectl version -o json) — the pod hostname <job>-<index> (PodIndexLabel) that MASTER_ADDR=<job>-0.<svc> resolves to needs it; on older clusters rank-0 hangs at rendezvous.

  2. Capacity gate ×nodes: hard-fail unless allocatable GPUs ≥ gpus_per_node × nodes (no gang scheduling → a partial start leaves rank-0 waiting forever).

  3. Render templates/k8s/indexed-job.yaml.tmpl — the headless Service + Indexed Job + rendezvous env (WORLD_SIZE = node count, NODE_RANK from JOB_COMPLETION_INDEX, MASTER_ADDR=<job>-0.<svc>, /dev/shm 16Gi so NCCL doesn't silently hang). kubectl apply -f creates the Service and Job together; cancel deletes the Job (Foreground) and the Service.

  4. NCCL probe first (as SLURM) — a 2-node all-reduce with a timeout; on hang, set the cluster NCCL env and re-probe; cache per cluster.

Local cluster (development, CI, and evals)

A throwaway minikube/kind cluster exercises admission, the four verbs, job-record wiring, and log plumbing without cluster quota — and is what an agent-driven eval should provision for itself. kubectl and minikube are single static binaries needing no root, so a non-root CI container can install them itself.

Two prerequisites keep a rendered Job Pending, and the first masks the second: the PVC the template mounts must exist (persistentvolumeclaim "<name>" not found fires before any GPU complaint), then a Job requesting nvidia.com/gpu on a GPU-less cluster reports Insufficient nvidia.com/gpu and waits forever. Render NUM_GPUS=0 for a lifecycle-only run and say GPU scheduling was not verified; on a Linux GPU host, minikube start --driver=docker --gpus all passes real GPUs through, so one GPU box suffices for a GPU-real smoke.

Install commands, driver choice, the container/host-networking caveat, and the fake-device-plugin middle option: references/local-cluster.md.

Container shell

The simple single-pod template invokes its command via /bin/sh -c (POSIX sh, present in busybox/distroless as well as TAO images). For producer action requests, an args-mode command and its arguments are native container argv; a producer that needs a shell declares the shell and its script explicitly. A simple config-mode command also becomes native argv after {config_path} is substituted. A producer-owned config command that is itself a multi-line shell script is preserved verbatim under /bin/sh -c; the renderer never constructs shell text from config values.

GPU Operator dependency

The submit verb refuses to launch GPU jobs on a cluster with no nvidia.com/gpu allocatable. For self-managed clusters, first run the tao-setup-nvidia-gpu-host install action on every GPU worker node or bake the same package set into the node image:

bash skills/platform/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh --backend kubernetes --install --yes

Then install the NVIDIA GPU Operator or device plugin:

helm repo add nvidia https://helm.ngc.nvidia.com/nvidia
helm repo update
helm install --wait gpu-operator -n gpu-operator --create-namespace nvidia/gpu-operator

Full guide: https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html

Multi-node training (distributed)

Set num_nodes > 1 (see the Multi-node (nodes > 1) verb steps above) to run distributed training across N pods. Rendering templates/k8s/indexed-job.yaml.tmpl provisions:

  1. A headless Service named after the Job (selector: job-name=<job-name>, clusterIP: None, publishNotReadyAddresses: true so pods can rendezvous before they're all Ready).

  2. An Indexed Job with parallelism = completions = num_nodes, completionMode: Indexed. Each pod gets JOB_COMPLETION_INDEX injected by k8s automatically (= the node rank).

  3. A command wrapper that exports the rendezvous env vars before invoking the user command. Two naming conventions are exported simultaneously:

    Env varValueRead by
    WORLD_SIZEnum_nodesTAO PyTorch container's nvidia_tao_pytorch/core/entrypoint.py (uses this to mean node count, even though PyTorch's own convention is total processes)
    NUM_GPU_PER_NODEgpu_countTAO PyTorch container's entrypoint
    NNODESnum_nodestorchrun and PyTorch-standard rendezvous
    NPROC_PER_NODEgpu_counttorchrun
    NODE_RANK$JOB_COMPLETION_INDEXboth
    MASTER_ADDR<job-name>-0.<job-name> (pod-0's DNS)both
    MASTER_PORT29500both (TAO's default)

    Both naming conventions are set so TAO entrypoints (dino train, etc.) and raw torchrun commands work without modification.

For a TAO entrypoint, the container reads spec.train.num_nodes and the wired env vars — e.g. dino train -e /tmp/spec.yaml with gpu_count=8, num_nodes=4 (4 × 8 = 32 GPUs total).

For raw torchrun-based commands (non-TAO containers), the wrapper invokes:

torchrun --nnodes=$NNODES --nproc-per-node=$NPROC_PER_NODE --node-rank=$NODE_RANK \
  --master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py

The capacity check sums across nodes: gpu_count × num_nodes ≤ cluster's allocatable nvidia.com/gpu.

Cluster requirements for multi-node
  • k8s 1.28+ is required for stable pod hostnames in Indexed Jobs (the PodIndexLabel feature). On older clusters the MASTER_ADDR=<job>-0.<svc> DNS lookup fails. Verify with kubectl version.

  • Pod-to-pod networking must be open on port 29500 (PyTorch default; configurable via MASTER_PORT env var). Most CNIs (Calico, Cilium, AWS VPC CNI) allow this by default; restrictive NetworkPolicies must be relaxed.

  • NCCL in the container talks GPU-to-GPU; if the cluster has multi-NIC nodes or RDMA, set NCCL_SOCKET_IFNAME / NCCL_IB_HCA in the container env of the rendered manifest.

Reference reading
Kubernetes operator alternatives

For more sophisticated topologies (gang scheduling, PyTorch elastic / fault-tolerant training, MPI / Horovod, RDMA setup), reach for an operator instead of plain Indexed Job:

This skill's Indexed Job path is intentionally simple and dependency-free; if you need elastic restart or gang scheduling, layer one of these on top and submit jobs through the operator's CRD instead.

Common error patterns

No nvidia.com/gpu resources allocatable on the cluster — the GPU Operator (or NVIDIA Device Plugin) isn't installed. Install per the link above; verify with kubectl get nodes -o jsonpath='{.items[*].status.allocatable}'.

ImagePullBackOff / ErrImagePull — the cluster can't pull the image. For nvcr.io: pre-create an image-pull secret in the namespace and reference it as the pod's imagePullSecrets in the rendered manifest: Feed the key over stdin — --docker-password=$NGC_KEY would put the secret in argv, where it is visible in the host's process table and shell history:

set -a; source /path/to/.env; set +a   # omit if already exported
kubectl create secret generic ngc-pull-secret -n tao-jobs \
  --type=kubernetes.io/dockerconfigjson \
  --from-file=.dockerconfigjson=/dev/stdin <<EOF
{"auths": {"nvcr.io": {"username": "\$oauthtoken", "password": "${NGC_KEY}"}}}
EOF
# Verify without reading the secret back:
kubectl get secret ngc-pull-secret -n tao-jobs >/dev/null && echo SECRET_OK

Pod stays Pending forever — kubectl describe pod -l job-name=$JOB_ID shows the scheduling reason in the Events. Common causes: insufficient GPU capacity (Insufficient nvidia.com/gpu), no node matches the pod's nodeSelector, missing image-pull secret, or PVC mount failure.

OOMKilled (exit 137) — container exceeded memory. Reduce batch size, lower max_length, or add a memory request/limit and target a larger node.

CredentialError: Could not authenticate to a Kubernetes cluster — neither kubeconfig nor in-cluster auth worked. Run kubectl get nodes to verify your config, or set $KUBECONFIG to the right path.

What this skill does NOT support (yet)

  • Elastic / fault-tolerant training. Indexed Job has backoff_limit=0 — failures fail the whole training run. For elastic restart (e.g., resume from checkpoint after a node death), use Kubeflow's PyTorchJob operator instead.

  • Gang scheduling. Indexed Job pods are scheduled independently — no all-or-nothing. Multi-node training will partially start if only some pods can be scheduled (rank-0 will hang waiting for peers). For all-or-nothing scheduling on shared clusters, use Volcano or Kueue.

  • MPI / Horovod. Use the MPI Operator. The Indexed Job path here is PyTorch-distributed-shaped (env-var rendezvous on MASTER_ADDR:MASTER_PORT).

  • Auto-creating image-pull secrets from $NGC_KEY. You pre-create the secret in the target namespace and pass the name. K8s namespace conventions vary widely, so we keep secret creation explicit.

Files

scripts/render_action_job.py

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Render one producer action request as a single-pod Kubernetes Job.

The data mover first mirrors each distinct request ``mounts[].source`` into a
job-owned directory on one PVC.  Its receipt records the relative PVC path for
each source.  This renderer preserves every producer-declared target and access
mode, including multiple target aliases for the same source, by mounting that
PVC subPath once per target.

Credentials are deliberately not accepted as values.  Each ``forward_env``
name is projected from the same key of an optional, pre-created Secret; keys
that the producer did not approve are never imported.  Registry authentication
is an independent optional image-pull Secret.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import pathlib
import posixpath
import re
import shlex
import stat
import sys
import tempfile
from typing import Any


REPO_ROOT = pathlib.Path(__file__).resolve().parents[4]
DEFAULT_TEMPLATE = REPO_ROOT / "templates/k8s/action-job.yaml.tmpl"
MARKER_RE = re.compile(r"@@[A-Z][A-Z0-9_]*@@")
ENV_NAME_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
DNS_LABEL_RE = re.compile(r"[a-z0-9](?:[-a-z0-9]*[a-z0-9])?")
DNS_SUBDOMAIN_RE = re.compile(
    r"[a-z0-9](?:[-a-z0-9.]*[a-z0-9])?"
)
QUANTITY_RE = re.compile(r"[1-9][0-9]*(?:Ki|Mi|Gi|Ti|Pi|Ei)")
MAX_JOB_NAME = 52
REQUEST_SCHEMA_VERSION = "1"
CONFIG_FORMATS = {"json": "json", "toml": "toml", "yaml": "yaml"}
SHELL_META = ("\n", ";", "&&", "||", "$(", "`", "|", ">", "<")


class RenderError(ValueError):
    """The action request or staging receipt cannot be rendered safely."""


def _load_object(path: pathlib.Path, label: str) -> dict[str, Any]:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise RenderError(f"cannot read {label} JSON {path}: {exc}") from exc
    if not isinstance(payload, dict):
        raise RenderError(f"{label} JSON root must be an object")
    return payload


def _nonempty_string(value: Any, label: str) -> str:
    if not isinstance(value, str) or not value or "\x00" in value:
        raise RenderError(f"{label} must be a non-empty string without NUL")
    return value


def _string_value(value: Any, label: str) -> str:
    if not isinstance(value, str) or "\x00" in value:
        raise RenderError(f"{label} must be a string without NUL")
    return value


def _dns_name(value: str, label: str, *, label_only: bool = False) -> str:
    pattern = DNS_LABEL_RE if label_only else DNS_SUBDOMAIN_RE
    maximum = 63 if label_only else 253
    if len(value) > maximum or pattern.fullmatch(value) is None:
        kind = "DNS label" if label_only else "DNS subdomain"
        raise RenderError(f"{label} must be a valid Kubernetes {kind}")
    return value


def kubernetes_job_name(job_id: str) -> str:
    """Map an immutable job-record id to a stable Kubernetes-safe name.

    Job-record components may contain underscores, uppercase letters, dots, or
    enough text to exceed Kubernetes' practical Job-name length. Preserve an
    already-safe short id; otherwise append a digest so normalization and
    truncation cannot merge two record ids into one backend object.
    """
    original = _nonempty_string(job_id, "job id")
    if len(original) <= MAX_JOB_NAME and DNS_LABEL_RE.fullmatch(original):
        return original
    slug = re.sub(r"[^a-z0-9]+", "-", original.lower()).strip("-")
    if not slug:
        slug = "tao-job"
    digest = hashlib.sha256(original.encode("utf-8")).hexdigest()[:10]
    prefix = slug[: MAX_JOB_NAME - len(digest) - 1].rstrip("-") or "tao-job"
    return f"{prefix}-{digest}"


def _absolute_clean_path(value: Any, label: str) -> str:
    path = _nonempty_string(value, label)
    if (
        not path.startswith("/")
        or path == "/"
        or "\\" in path
        or any(ord(ch) < 32 for ch in path)
    ):
        raise RenderError(f"{label} must be an absolute non-root POSIX path")
    normalized = posixpath.normpath(path)
    if normalized != path or "//" in path:
        raise RenderError(f"{label} must be normalized and contain no traversal")
    return path


def _safe_sub_path(value: Any, label: str) -> str:
    path = _nonempty_string(value, label)
    if path.startswith("/") or "\\" in path or any(ord(ch) < 32 for ch in path):
        raise RenderError(f"{label} must be a safe relative POSIX path")
    normalized = posixpath.normpath(path)
    if normalized != path or path in {".", ".."} or path.startswith("../"):
        raise RenderError(f"{label} must be normalized and contain no traversal")
    return path


def _toml_key(value: Any) -> str:
    key = _nonempty_string(value, "TOML key")
    if re.fullmatch(r"[A-Za-z0-9_-]+", key):
        return key
    return json.dumps(key)


def _toml_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, int):
        return str(value)
    if isinstance(value, float):
        if not math.isfinite(value):
            raise RenderError("TOML config values must be finite")
        return repr(value)
    if isinstance(value, str):
        return json.dumps(value)
    if isinstance(value, list):
        return "[" + ", ".join(_toml_value(item) for item in value) + "]"
    if isinstance(value, dict):
        return "{ " + ", ".join(
            f"{_toml_key(key)} = {_toml_value(item)}"
            for key, item in sorted(value.items())
        ) + " }"
    raise RenderError(f"no TOML representation for {type(value).__name__}")


def _toml_document(spec: dict[str, Any], prefix: tuple[str, ...] = ()) -> str:
    scalars: list[str] = []
    tables: list[tuple[str, dict[str, Any]]] = []
    for raw_key in sorted(spec):
        key = _toml_key(raw_key)
        value = spec[raw_key]
        if isinstance(value, dict):
            tables.append((key, value))
        else:
            scalars.append(f"{key} = {_toml_value(value)}")
    sections: list[str] = []
    if scalars:
        sections.append("\n".join(scalars))
    for key, value in tables:
        name = ".".join((*prefix, key))
        body = _toml_document(value, (*prefix, key)).rstrip("\n")
        sections.append(f"[{name}]" + (f"\n{body}" if body else ""))
    return "\n\n".join(sections) + "\n"


def _config_content(bundle: dict[str, Any]) -> tuple[str, str, str]:
    spec = bundle.get("spec")
    if not isinstance(spec, dict):
        raise RenderError("spec_bundle.spec must be a nested object for mode=config")
    raw_format = bundle.get("config_format")
    if raw_format not in CONFIG_FORMATS:
        raise RenderError(
            "spec_bundle.config_format must be one of json, toml, or yaml"
        )
    extension = CONFIG_FORMATS[raw_format]
    try:
        if raw_format == "toml":
            content = _toml_document(spec)
        else:
            # JSON is valid YAML. Keeping both formats on the stdlib serializer
            # avoids making Kubernetes submission depend on a host YAML package.
            content = json.dumps(
                spec, allow_nan=False, indent=2, sort_keys=True
            ) + "\n"
    except (TypeError, ValueError) as exc:
        raise RenderError(f"spec_bundle.spec is not {raw_format}-serializable: {exc}") from exc
    digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
    return content, extension, digest


def _read_regular_file(path: pathlib.Path, label: str) -> bytes:
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise RenderError(f"cannot open {label} {path}: {exc}") from exc
    try:
        if not stat.S_ISREG(os.fstat(descriptor).st_mode):
            raise RenderError(f"{label} must be a regular non-symlink file: {path}")
        with os.fdopen(descriptor, "rb") as handle:
            descriptor = -1
            return handle.read()
    except OSError as exc:
        raise RenderError(f"cannot read {label} {path}: {exc}") from exc
    finally:
        if descriptor >= 0:
            os.close(descriptor)


def materialize_config(
    request: dict[str, Any], output_dir: pathlib.Path
) -> pathlib.Path:
    """Write one content-addressed config file for later PVC staging."""
    if request.get("schema_version") != REQUEST_SCHEMA_VERSION:
        raise RenderError(
            f"request.schema_version must be {REQUEST_SCHEMA_VERSION!r}"
        )
    if request.get("platform") != "kubernetes":
        raise RenderError("request.platform must be 'kubernetes'")
    bundle = request.get("spec_bundle")
    if not isinstance(bundle, dict):
        raise RenderError("request.spec_bundle must be an object")
    if bundle.get("mode") != "config":
        raise RenderError("materialize-config requires spec_bundle.mode=config")
    if "args" in bundle:
        raise RenderError("spec_bundle.mode=config must not carry args")
    content, extension, digest = _config_content(bundle)

    raw_output_dir = output_dir.expanduser()
    if raw_output_dir.is_symlink():
        raise RenderError("config output directory must not be a symlink")
    try:
        raw_output_dir.mkdir(parents=True, exist_ok=True)
        root = raw_output_dir.resolve(strict=True)
    except OSError as exc:
        raise RenderError(f"cannot prepare config output directory: {exc}") from exc
    if root == pathlib.Path(root.anchor) or not root.is_dir():
        raise RenderError("config output directory must be a non-root directory")

    destination = root / f"tao-action-config-{digest}.{extension}"
    encoded = content.encode("utf-8")
    if destination.exists() or destination.is_symlink():
        if _read_regular_file(destination, "materialized config") != encoded:
            raise RenderError(f"materialized config content mismatch: {destination}")
        return destination

    temporary: pathlib.Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="wb",
            prefix=destination.name + ".",
            suffix=".tmp",
            dir=root,
            delete=False,
        ) as handle:
            temporary = pathlib.Path(handle.name)
            os.fchmod(handle.fileno(), 0o600)
            handle.write(encoded)
            handle.flush()
            os.fsync(handle.fileno())
        try:
            # link() publishes the complete file without ever replacing an
            # existing path. Concurrent materializers either create the same
            # content-addressed file or verify the winner below.
            os.link(temporary, destination)
        except FileExistsError:
            if _read_regular_file(destination, "materialized config") != encoded:
                raise RenderError(
                    f"materialized config content mismatch: {destination}"
                )
    except OSError as exc:
        raise RenderError(f"cannot materialize config in {root}: {exc}") from exc
    finally:
        if temporary is not None:
            try:
                temporary.unlink()
            except FileNotFoundError:
                pass
    return destination


def _is_shell_script(command: str) -> bool:
    return any(marker in command for marker in SHELL_META)


def _command_bundle(
    request: dict[str, Any], config_path: str | None
) -> tuple[list[str], list[str], str, int]:
    bundle = request.get("spec_bundle")
    if not isinstance(bundle, dict):
        raise RenderError("request.spec_bundle must be an object")
    mode = bundle.get("mode")
    command_text = _nonempty_string(bundle.get("command"), "spec_bundle.command")
    if mode == "args":
        if "spec" in bundle or "config_format" in bundle:
            raise RenderError(
                "spec_bundle.mode=args must not carry spec or config_format"
            )
        if config_path is not None:
            raise RenderError("--config-source must be omitted for spec_bundle.mode=args")
        raw_args = bundle.get("args")
        if not isinstance(raw_args, list):
            raise RenderError("spec_bundle.args must be an array")
        args = [
            _string_value(value, f"spec_bundle.args[{index}]")
            for index, value in enumerate(raw_args)
        ]
    elif mode == "config":
        if "args" in bundle:
            raise RenderError("spec_bundle.mode=config must not carry args")
        _config_content(bundle)
        if config_path is None:
            raise RenderError(
                "spec_bundle.mode=config requires --config-source; run "
                "materialize-config, stage the returned file, then render"
            )
        if "{config_path}" not in command_text:
            raise RenderError(
                "spec_bundle.mode=config requires {config_path} in command"
            )
        command_text = command_text.replace("{config_path}", config_path)
        args = []
    else:
        raise RenderError("spec_bundle.mode must be 'args' or 'config'")

    if mode == "config" and _is_shell_script(command_text):
        command = ["/bin/sh", "-c"]
        args = [command_text, "tao-action", *args]
    else:
        try:
            command = shlex.split(command_text)
        except ValueError as exc:
            raise RenderError(f"spec_bundle.command cannot be parsed: {exc}") from exc
        if not command:
            raise RenderError("spec_bundle.command must contain an executable")
    image = _nonempty_string(bundle.get("image"), "spec_bundle.image")
    workload_image = request.get("workload_image")
    if workload_image is not None and workload_image != image:
        raise RenderError("request.workload_image must equal spec_bundle.image")
    shape = bundle.get("compute_shape")
    if not isinstance(shape, dict):
        raise RenderError("spec_bundle.compute_shape must be an object")
    gpus = shape.get("gpus")
    nodes = shape.get("nodes")
    if not isinstance(gpus, int) or isinstance(gpus, bool) or gpus < 0:
        raise RenderError("spec_bundle.compute_shape.gpus must be a non-negative integer")
    if nodes != 1:
        raise RenderError("single-pod action rendering requires compute_shape.nodes=1")
    return command, args, image, gpus


def _staged_sources(payload: dict[str, Any]) -> dict[str, str]:
    if payload.get("schema_version") != "1":
        raise RenderError("staging map schema_version must be '1'")
    rows = payload.get("sources")
    if not isinstance(rows, list) or not rows:
        raise RenderError("staging map sources must be a non-empty array")
    result: dict[str, str] = {}
    sub_paths: dict[str, str] = {}
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            raise RenderError(f"staging map sources[{index}] must be an object")
        source = _absolute_clean_path(row.get("source"), f"staging sources[{index}].source")
        sub_path = _safe_sub_path(row.get("sub_path"), f"staging sources[{index}].sub_path")
        if source in result:
            raise RenderError(f"staging map repeats source: {source}")
        if sub_path in sub_paths:
            raise RenderError(
                "staging map assigns one PVC subPath to distinct sources: "
                f"{sub_paths[sub_path]} and {source}"
            )
        result[source] = sub_path
        sub_paths[sub_path] = source
    return result


def _mounts(
    request: dict[str, Any],
    staged: dict[str, str],
    config_mount: tuple[str, str] | None = None,
) -> tuple[list[dict[str, Any]], list[tuple[pathlib.PurePosixPath, bool]]]:
    rows = request.get("mounts")
    if not isinstance(rows, list) or not rows:
        raise RenderError("request.mounts must be a non-empty array")
    rendered: list[dict[str, Any]] = [
        {"name": "dshm", "mountPath": "/dev/shm", "readOnly": False}  # nosec B108 - k8s emptyDir shm mount, not a temp file
    ]
    used_sources: set[str] = set()
    targets: set[str] = {"/dev/shm"}  # nosec B108 - k8s emptyDir shm mount, not a temp file
    source_modes: list[tuple[pathlib.PurePosixPath, bool]] = []
    for index, row in enumerate(rows):
        if not isinstance(row, dict):
            raise RenderError(f"request.mounts[{index}] must be an object")
        source = _absolute_clean_path(row.get("source"), f"mounts[{index}].source")
        target = _absolute_clean_path(row.get("target"), f"mounts[{index}].target")
        read_only = row.get("read_only")
        if not isinstance(read_only, bool):
            raise RenderError(f"mounts[{index}].read_only must be boolean")
        if source not in staged:
            raise RenderError(f"mount source has no staged PVC subPath: {source}")
        if target in targets:
            raise RenderError(f"request repeats Kubernetes mount target: {target}")
        targets.add(target)
        used_sources.add(source)
        source_modes.append((pathlib.PurePosixPath(source), read_only))
        rendered.append(
            {
                "name": "workspace",
                "mountPath": target,
                "subPath": staged[source],
                "readOnly": read_only,
            }
        )
    if config_mount is not None:
        source, target = config_mount
        if source in used_sources:
            raise RenderError(
                "materialized config must have its own staged file source"
            )
        if source not in staged:
            raise RenderError(
                f"materialized config has no staged PVC subPath: {source}"
            )
        if target in targets:
            raise RenderError(
                f"materialized config repeats Kubernetes mount target: {target}"
            )
        used_sources.add(source)
        targets.add(target)
        rendered.append(
            {
                "name": "workspace",
                "mountPath": target,
                "subPath": staged[source],
                "readOnly": True,
            }
        )
    extras = sorted(set(staged) - used_sources)
    if extras:
        raise RenderError("staging map contains undeclared mount sources: " + ", ".join(extras))
    return rendered, source_modes


def _config_mount(
    request: dict[str, Any],
    staged: dict[str, str],
    config_source: pathlib.Path | None,
) -> tuple[str | None, tuple[str, str] | None]:
    bundle = request.get("spec_bundle")
    if not isinstance(bundle, dict):
        raise RenderError("request.spec_bundle must be an object")
    mode = bundle.get("mode")
    if mode != "config":
        if config_source is not None:
            raise RenderError("--config-source is valid only for spec_bundle.mode=config")
        return None, None
    if config_source is None:
        return None, None

    lexical = config_source.expanduser()
    source = _absolute_clean_path(str(lexical), "config source")
    path = pathlib.Path(source)
    content, extension, digest = _config_content(bundle)
    if _read_regular_file(path, "config source") != content.encode("utf-8"):
        raise RenderError(
            "config source content does not match spec_bundle.spec; rerun "
            "materialize-config and restage the returned file"
        )
    if source not in staged:
        raise RenderError(f"materialized config has no staged PVC subPath: {source}")
    target = f"/tao-action-config/spec-{digest}.{extension}"
    return target, (source, target)


def _require_writable_outputs(
    request: dict[str, Any], source_modes: list[tuple[pathlib.PurePosixPath, bool]]
) -> None:
    outputs = request.get("fresh_outputs", [])
    if not isinstance(outputs, list):
        raise RenderError("request.fresh_outputs must be an array")
    for index, raw in enumerate(outputs):
        output = pathlib.PurePosixPath(
            _absolute_clean_path(raw, f"fresh_outputs[{index}]")
        )
        writable = False
        for source, read_only in source_modes:
            try:
                output.relative_to(source)
            except ValueError:
                continue
            if not read_only:
                writable = True
                break
        if not writable:
            raise RenderError(
                f"fresh output is not covered by a writable declared mount: {output}"
            )


def _environment(
    request: dict[str, Any], credential_secret: str | None
) -> list[dict[str, Any]]:
    raw_environment = request.get("environment", {})
    if not isinstance(raw_environment, dict):
        raise RenderError("request.environment must be an object")
    environment: list[dict[str, Any]] = []
    for name in sorted(raw_environment):
        if ENV_NAME_RE.fullmatch(name) is None:
            raise RenderError(f"invalid environment variable name: {name!r}")
        value = _string_value(raw_environment[name], f"environment.{name}")
        environment.append({"name": name, "value": value})

    raw_forward = request.get("forward_env", [])
    if not isinstance(raw_forward, list):
        raise RenderError("request.forward_env must be an array")
    forward: list[str] = []
    for index, raw in enumerate(raw_forward):
        name = _nonempty_string(raw, f"forward_env[{index}]")
        if ENV_NAME_RE.fullmatch(name) is None:
            raise RenderError(f"invalid forwarded environment variable name: {name!r}")
        if name in forward:
            raise RenderError(f"request.forward_env repeats {name}")
        if name in raw_environment:
            raise RenderError(f"credential {name} must not be present in request.environment")
        forward.append(name)
    if forward and credential_secret is None:
        raise RenderError(
            "request.forward_env is non-empty but no --credential-secret was supplied"
        )
    if not forward and credential_secret is not None:
        raise RenderError(
            "--credential-secret must be omitted when request.forward_env is empty"
        )
    if credential_secret is not None:
        environment.extend(
            {
                "name": name,
                "valueFrom": {
                    "secretKeyRef": {"name": credential_secret, "key": name}
                },
            }
            for name in forward
        )
    return environment


def render_action_job(
    request: dict[str, Any],
    staging_map: dict[str, Any],
    *,
    job_id: str,
    namespace: str,
    pvc_claim: str,
    credential_secret: str | None = None,
    image_pull_secret: str | None = None,
    config_source: pathlib.Path | None = None,
    ttl_seconds: int = 3600,
    shm_size: str = "16Gi",
    template_path: pathlib.Path = DEFAULT_TEMPLATE,
) -> str:
    """Validate inputs and return a complete Kubernetes Job manifest."""
    if request.get("schema_version") != REQUEST_SCHEMA_VERSION:
        raise RenderError(
            f"request.schema_version must be {REQUEST_SCHEMA_VERSION!r}"
        )
    if request.get("platform") != "kubernetes":
        raise RenderError("request.platform must be 'kubernetes'")
    job_id = _nonempty_string(job_id, "job id")
    job_name = kubernetes_job_name(job_id)
    namespace = _dns_name(namespace, "namespace", label_only=True)
    pvc_claim = _dns_name(pvc_claim, "PVC claim")
    if credential_secret is not None:
        credential_secret = _dns_name(credential_secret, "credential secret")
    if image_pull_secret is not None:
        image_pull_secret = _dns_name(image_pull_secret, "image pull secret")
    if (
        not isinstance(ttl_seconds, int)
        or isinstance(ttl_seconds, bool)
        or not 0 <= ttl_seconds <= 604800
    ):
        raise RenderError("ttl_seconds must be an integer from 0 through 604800")
    if QUANTITY_RE.fullmatch(shm_size) is None:
        raise RenderError("shm_size must be a positive binary Kubernetes quantity such as 16Gi")

    staged = _staged_sources(staging_map)
    config_path, config_mount = _config_mount(request, staged, config_source)
    command, args, image, gpus = _command_bundle(request, config_path)
    volume_mounts, source_modes = _mounts(request, staged, config_mount)
    _require_writable_outputs(request, source_modes)
    environment = _environment(request, credential_secret)
    image_pull_secrets = (
        [{"name": image_pull_secret}] if image_pull_secret is not None else []
    )

    values: dict[str, Any] = {
        "JOB_NAME_JSON": job_name,
        "JOB_ID_JSON": job_id,
        "NAMESPACE_JSON": namespace,
        "TTL_SECONDS_JSON": ttl_seconds,
        "IMAGE_PULL_SECRETS_JSON": image_pull_secrets,
        "IMAGE_JSON": image,
        "COMMAND_JSON": command,
        "ARGS_JSON": args,
        "NUM_GPUS_JSON": str(gpus),
        "ENV_JSON": environment,
        "VOLUME_MOUNTS_JSON": volume_mounts,
        "SHM_SIZE_JSON": shm_size,
        "PVC_CLAIM_JSON": pvc_claim,
    }
    try:
        rendered = template_path.read_text(encoding="utf-8")
    except OSError as exc:
        raise RenderError(f"cannot read Kubernetes action template: {exc}") from exc
    for marker, value in values.items():
        rendered = rendered.replace(f"@@{marker}@@", json.dumps(value, separators=(",", ":")))
    remaining = MARKER_RE.findall(rendered)
    if remaining:
        raise RenderError(f"unresolved template markers: {sorted(set(remaining))}")
    return rendered


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="operation", required=True)
    name = commands.add_parser(
        "name", help="print the deterministic Kubernetes name for a job-record id"
    )
    name.add_argument("--job-id", required=True)

    materialize = commands.add_parser(
        "materialize-config",
        help="write a content-addressed mode=config file for PVC staging",
    )
    materialize.add_argument("--request", required=True, type=pathlib.Path)
    materialize.add_argument("--output-dir", required=True, type=pathlib.Path)

    render = commands.add_parser("render", help="render a complete Job manifest")
    render.add_argument("--request", required=True, type=pathlib.Path)
    render.add_argument("--staging-map", required=True, type=pathlib.Path)
    render.add_argument("--job-id", required=True)
    render.add_argument(
        "--namespace", default=os.environ.get("TAO_K8S_NAMESPACE", "default")
    )
    render.add_argument("--pvc-claim", required=True)
    render.add_argument("--credential-secret")
    render.add_argument("--image-pull-secret")
    render.add_argument("--config-source", type=pathlib.Path)
    render.add_argument("--ttl-seconds", type=int, default=3600)
    render.add_argument("--shm-size", default="16Gi")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.operation == "name":
            print(kubernetes_job_name(args.job_id))
            return 0
        request = _load_object(args.request, "action request")
        if args.operation == "materialize-config":
            print(materialize_config(request, args.output_dir))
            return 0
        staging_map = _load_object(args.staging_map, "staging map")
        manifest = render_action_job(
            request,
            staging_map,
            job_id=args.job_id,
            namespace=args.namespace,
            pvc_claim=args.pvc_claim,
            credential_secret=args.credential_secret,
            image_pull_secret=args.image_pull_secret,
            config_source=args.config_source,
            ttl_seconds=args.ttl_seconds,
            shm_size=args.shm_size,
        )
    except RenderError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2
    sys.stdout.write(manifest)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Declares allowed tools, Needs a command-line tool, Needs Python, Needs Node.js

The author says: Requires GPU worker nodes with NVIDIA driver branch 580, CUDA Toolkit 13.0, and NVIDIA Container Toolkit 1.19.0; a `kubectl` client authenticated to the cluster; and the NVIDIA GPU Operator or device plugin. No nvidia-tao-sdk required — jobs are submitted with plain `kubectl`.

Security

What our scanner noticed· Medium risk

  • Large base64-encoded blobs
  • Instructions that use the network
  • Scripts that your agent may run

This skill includes scripts that your agent may run. Read them before you use it.

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by
NVIDIA Corporation
Publisher
NVIDIA (Official)
Source
NVIDIA/skills · skills/tao-run-on-kubernetes · commit fce3cca
Licence
Apache License 2.0 (redistribution allowed)
  • LICENSE (added by AutoGineer from LICENSE-APACHE)
  • ATTRIBUTION.md (written by AutoGineer)
Synced

AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-run-on-kubernetes.zip