Datacenter Thermal Management AI

A Hybrid Artificial Intelligence system designed to dynamically optimize GPU datacenter cooling. This project utilizes a "Two-Brain" architecture, combining Sequence Forecasting (LSTM) and Deep Reinforcement Learning (PPO) to proactively manage fan speeds, water pump flow, and GPU power limits.

The AI learns to balance three critical objectives:

1. Preventing thermal crashes (keeping GPU core temps strictly under 95°C).
2. Minimizing facility power consumption (optimizing PUE by reducing unnecessary fan/pump usage).
3. Maximizing GPU compute performance (avoiding artificial power throttling unless absolutely necessary).

---

The "Two-Brain" Architecture

Instead of reacting to heat after the hardware is already burning, this system uses predictive modeling to act preemptively.

• Brain 1: The Oracle (LSTM Predictor)

  • A PyTorch-based Long Short-Term Memory network.
  • Consumes a rolling 60-second window of datacenter telemetry (ambient temp, core temp, workload, pump flow, fan speed, power).
  • Objective: Forecasts the GPU core temperature 3 minutes into the future.

• Brain 2: The Controller (PPO Agent)

  • A Reinforcement Learning agent built with stable-baselines3.
  • Trained in a custom Gymnasium thermal physics sandbox.
  • Consumes real-time telemetry plus the Oracle's future prediction (a 7-variable observation space).
  • Objective: Outputs mechanical actions (deltas) to adjust Fans (RPM), Pumps (L/min), and Power Limits (kW) to alter the predicted future.

---

📂 Project Structure

thermal_ai/
│
├── data/
│   └── raw_telemetry.csv            # Synthetic/Historical telemetry dataset
│
├── src/
│   ├── simulator/
│   │   ├── gym_env.py               # Custom Gymnasium environment with thermal physics & reward logic
│   │   └── thermal_physics.py       # Math engine for calculating heat transfer and cooling power
│   │
│   ├── models/
│   │   └── lstm_predictor.py        # PyTorch neural network architecture for The Oracle
│   │
│   ├── train/
│   │   ├── train_lstm.py            # Supervised learning script for the Oracle
│   │   └── train_rl.py              # PPO training script for the Controller
│   │
│   └── deploy/
│       └── inference_loop.py        # Live shadow-mode loop linking both brains and reading CSV data
│
├── ppo_telemetry_predictor.pt       # Saved Oracle weights (Generated)
├── ppo_datacenter_thermal_agent.zip # Saved PPO Agent brain (Generated)
├── .gitignore                       
└── README.md

Installation

pip install torch pandas numpy stable-baselines3[extra] gymnasium

Usage Guide

1. Train the Oracle (LSTM)

python src/train/train_lstm.py

2. Train Controller (RL Agent)

python src/train/train_rl.py

3. Deploy Inference Loop

python src/deploy/inference_loop.py

Reward Function Highlights

• +5.0: Base survival reward per step.
• -0.05 * Power: Light penalty for wasting electricity on fans/pumps.
• Exponential Heat Penalty: If temp > 75°C, penalty squares rapidly to force immediate pump/fan usage.
• -2000.0: Massive terminal penalty for allowing the server to crash (> 95°C).