AI Infrastructure on Azure

Table of Contents

1. Overview
2. Infrastructure References Catalog
3. Storage References Catalog
4. AI Training Example Catalog
5. Infrastructure Validation Catalog
6. Scheduling and Workload Management
7. Utilities Catalog
8. AI Infrastructure MCP Server
9. Contributing
10. Trademarks
11. Contributors

1. Overview

This repository collects architectural guidance and AI training examples meant to run on Azure AI Infrastructure.

This includes infrastructure examples and real use case scenarios on Azure AI Infrastructure involving different orchestration solutions:

• Azure CycleCloud Workspace for Slurm
• Azure Kubernetes Service
• Azure Machine Learning

For each scenario and architecture, the repository will include storage recommendations among Azure Storage services (Azure Blob Storage, Azure Managed Lustre, Azure NetApp Files), monitoring and observability.

2. Infrastructure References Catalog

1. Azure CycleCloud Slurm Workspace AI Cluster - Prototypes for the creation of Azure CycleCloud Slurm Workspace AI Clusters using CLI deployment
2. Azure Kubernetes Service Cluster - Deployment script for AKS cluster

3. Storage References Catalog

1. AKS Shared Storage - Helm charts for deploying shared storage on AKS using Azure Blob Storage (BlobFuse) and Azure Managed Lustre File System (AMLFS)
2. Slurm Squashed Images - Tuning guidance for container squashed image files on Slurm clusters, including Azure Managed Lustre striping optimization and local NVME staging

4. AI Training Example Catalog

1. MegatronLM GPT3-175B with Slimpajama 627B dataset - Example of an end-to-end training workflow based on MegatronLM, including data pre-processing from Slimpajama 627B dataset
   • Slurm version
   • AKS version
2. LLM Foundry MPT Training - Example of an end-to-end training workflow of Mosaic Pretrained Transformer (MPT) model on C4 dataset, based on LLM Foundry
   • Slurm version
   • AKS version

5. Infrastructure Validation Catalog

1. NCCL All-reduce - Testing distributed communication performance for multi-GPU training
   • Slurm version
   • AKS version
2. Node Health Checks - Automated system validation and monitoring for compute nodes
   • Slurm version
   • AKS version
3. Thermal Test - GPU thermal stress testing and monitoring
   • Slurm version
4. FIO Storage Performance Testing - I/O performance testing with Azure Container Storage, blobfuse, and other storage types
   • AKS version

6. Scheduling and Workload Management

1. Kueue for AKS - Kubernetes-native job queueing and quota management for batch workloads on AKS. Provides a simple Helm chart example for setting up a GPU queue. All Helm charts in this repository support optional Kueue integration via the kueue.queueName parameter.

7. Utilities Catalog

1. Node Labeler - Automatically labels nodes with host information and InfiniBand HCA GUIDs for network topology awareness
   • AKS version
2. Torset Labeler - Discovers and labels nodes with torset (InfiniBand switching domain) information using SHARP topology discovery
   • AKS version

8. AI Infrastructure MCP Server

The AI Infrastructure MCP Server is a Model Context Protocol (MCP) server that provides tools for managing and monitoring Slurm-based HPC clusters. It enables AI assistants like GitHub Copilot to interact with cluster infrastructure through a standardized protocol, offering capabilities such as:

• Slurm job management - Query job status, accounting data, and cluster information
• System monitoring - Check systemd services and logs across cluster nodes
• File operations - Read and search files on the cluster
• Azure VM metadata - Retrieve physical hostnames and VMSS information

Currently targeting Slurm clusters with SSH-based connectivity. See the full documentation for setup and usage details.

9. Contributing

This project welcomes contributions and suggestions. 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.

When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact [email protected] with any additional questions or comments.

10. Trademarks

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 is subject to those third-party's policies.

11. Contributors

Please join us in contributing to the project