Microsoft has open-sourced TauGrid, a cloud-native platform designed to manage, schedule, and monitor AI workloads on GPU-enabled Kubernetes clusters.
Running AI workloads on Kubernetes often requires platform teams to assemble and maintain multiple open-source projects, custom scripts, and operational tools. This includes the glue between those components: submission scripts, queue wrappers, health checks, and result retrieval.
According to Microsoft, TauGrid is designed for both engineering and research teams and provided a unified stack. Platform teams can use advanced capabilities such as workspaces, queues, compute profiles, storage, identity, and observability, while researchers can submit workloads without needing to learn Kubernetes.
The company says TauGrid provides end-to-end management for AI workloads, covering everything from initial data preparation to distributed training, fine-tuning, and inference. Built on Kubernetes, it uses specialized queuing and topology-aware scheduling to efficiently manage intensive GPU workloads.
Besides its tau CLI, TauGrid includes Kueue for workload queuing and resource management, KubeRay for orchestration, GPU-node health monitoring, and observability capabilities.
Instead of building and maintaining these components and the integration between them separately, TauGrid provides a single Helm install with clear ownership boundaries.
TauGrid uses a yaml configuration file to define workloads, which can be submitted with tau run. The command validates the configuration and creates a Kubernetes Job or KubeRay RayJob, which is then queued using Kueue based on remaining quota and priority. When executed, TauGrid tracks workload status, logs, and checkpoints. It also collects and stores experiment evidence so that experiments can be reproduced and failures diagnosed later. The following is a tau.yaml configuration example for a PyTorch training job running on a single A100 GPU:
schema_version: 1
name: aks-gpu-quickstart
run:
entrypoint: train.py
workload_kind: rayjob
compute:
gpus: 1
workers: 1
cpus: 16
memory: 64Gi
runtime:
image: mcr.microsoft.com/aks/ai-runtime/ray:py3.12-ray2.56.0-cuda13.0
pip:
- torch>=2.4.0
When a job fails, TauGrid can resume it from a checkpoint.
TauGrid is still a work in progress, with its roadmap outlining a range of planned features. These include multi-tenant workspaces, RBAC and quotas, support for PyTorch DDP/FSDP, DeepSpeed and LoRA/QLoRA workflows, dataset lifecycle management, multi-cluster/multi-cloud execution, and others.
The TauGrid codebase is written primarily in Go, with development and contributions managed openly within the Azure ecosystem. Running TauGrid requires a Kubernetes cluster (1.30+) with GPU nodes, kubectl, and Helm 3.0 or later.
TauGrid is not the only Kubernetes-based AI-workload platform currently available. Alternatives include Kubeflow, which is advancing toward CNCF graduation as a "mature, production-ready ML system", Nvidia Run:AI, and others.