Cloud-native AI infrastructure for GPU workloads on Kubernetes

TauGrid runs GPU workloads on Kubernetes, including data preparation, distributed training, fine-tuning, and inference.

It combines the tau CLI, workload queueing and admission with Kueue, Ray cluster orchestration with KubeRay, node-level GPU health monitoring, and cluster and workload observability. Platform teams install this stack instead of assembling each component separately. Researchers use the CLI to submit and manage workloads without configuring Kubernetes directly.

TauGrid is built on open Kubernetes-native components:

┌─────────────────────────────────────────────────────────────────┐
│  Researcher                                                     │
│  tau run · status · logs · get · cancel                         │
└──────────────────────────┬──────────────────────────────────────┘
                           │
┌──────────────────────────▼──────────────────────────────────────┐
│  TauGrid Control Plane                                          │
│  ┌────────────┐  ┌──────────────┐  ┌────────────────────────┐  │
│  │   Kueue    │  │   KubeRay    │  │   GPU Health Monitor   │  │
│  │  Admission │  │ Orchestration│  │   Node Diagnostics     │  │
│  └────────────┘  └──────────────┘  └────────────────────────┘  │
└──────────────────────────┬──────────────────────────────────────┘
                           │
┌──────────────────────────▼──────────────────────────────────────┐
│  Compute, Data, and Evidence                                    │
│  CPU + GPU  ·  Datasets + Checkpoints  ·  Observability         │
└─────────────────────────────────────────────────────────────────┘

• A Kubernetes cluster (1.30+) with GPU nodes
• kubectl configured for your cluster
• Helm 3.0 or later

TauGrid is built on Kubernetes and is tested end-to-end on AKS. Some integrations, such as observability through Azure Data Explorer (Kusto), remain Azure-specific. The project intends to support cloud and on-premises Kubernetes environments without an Azure dependency; contributions toward that goal are welcome.

helm install taugrid \
  oci://mcr.microsoft.com/aks/ai-runtime/helm/taugrid \
  --version 0.4.0 \
  --namespace tau-system \
  --create-namespace

Install the latest stable GitHub Release with curl:

TAU_RELEASE_URL="$(
  curl -fsSL -o /dev/null -w '%{url_effective}' \
    https://github.com/Azure/taugrid/releases/latest
)"
TAU_VERSION="${TAU_RELEASE_URL##*/}"

curl -fsSL \
  "https://github.com/Azure/taugrid/releases/download/$TAU_VERSION/install.sh" |
  TAU_VERSION="$TAU_VERSION" bash

export PATH="$HOME/.local/bin:$PATH"
command -v tau
tau version --short

See the installation guide for PATH persistence, pinned versions, the Python SDK wheel, upgrades, and the advanced source installation.

The historical v0.3.0 Release predates the release installer. If it is still the latest Release, use the guide's source path until a newer Release is available.

Create a minimal tau.yaml:

name: hello-gpu
image: mcr.microsoft.com/aks/ai-runtime/tau:0.4.0
command: ["python", "train.py"]
resources:
  gpu: 4

Then submit, check status, and stream logs:

tau run --config tau.yaml
tau run status hello-gpu
tau run logs hello-gpu

Browse the runnable examples for checked-in configurations and step-by-step guides covering local, CPU, GPU, and Ray workloads.

TauGrid first-party images are published under the following MCR repositories:

Use the versioned tag or immutable digest documented by each release rather than a mutable latest tag.

TauGrid-owned charts are published as public OCI artifacts under oci://mcr.microsoft.com/aks/ai-runtime/helm. Chart versions come from each chart's Chart.yaml; published versions are immutable.

Full documentation is available at https://azure.github.io/taugrid.

ROADMAP.md describes where the project is going and what is deliberately out of scope.

TauGrid does not send telemetry to Microsoft by default. Platform operators can explicitly enable optional observability components and configure destinations such as Azure Data Explorer. To prevent remote telemetry export, keep those integrations disabled and do not configure a remote telemetry endpoint. Local Kubernetes logs, events, and metrics remain under the cluster operator's control.

TauGrid integrates third-party open-source software including Kubernetes, Kueue, KubeRay, Ray, Node Problem Detector, and Prometheus components. These projects remain under their respective licenses. Exact dependency versions are recorded in the repository's module manifests, Helm chart metadata and lock files, and container build definitions.

This project welcomes contributions and suggestions. Please see CONTRIBUTING.md for details.

Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.

For bug reports and feature requests, please open a GitHub issue.

Please see SECURITY.md for reporting security vulnerabilities.

This project is licensed under the MIT License - see the LICENSE file for details.

This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow Microsoft's Trademark & Brand Guidelines. Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party's policies.