AWS AI Capacity - GPU Capacity Management Agent

A Pydantic AI-powered agent for managing and querying AWS GPU capacity, with a Chainlit chat interface and CLI for automated reporting.

Features

• SageMaker Training Plans: Search offerings, list plans, check utilization
• EC2 Capacity: Query reservations, check instance availability
• Spot Capacity: Spot placement scores (1-10) and price history across regions
• On-Demand Capacity: Launch-and-terminate availability check (opt-in, costs money)
• GPU Utilization: View running GPU instances with spot vs on-demand breakdown
• GPU Specs: Compare instance types and their capabilities
• Chat Interface: Interactive Chainlit UI for ad-hoc queries
• CLI Reports: Automated report generation for cron jobs

Quick Start (uvx)

Run directly without installing:

# With inline environment variables
AWS_PROFILE=myprofile uvx aws-ai-capacity chat "What GPU capacity is available?"

# Or export first
export AWS_PROFILE=myprofile
uvx aws-ai-capacity chat "What p5 training plans are available?"
uvx aws-ai-capacity report daily

Installation (Development)

# Clone and install
git clone https://github.com/drewdresser/aws-ai-capacity.git
cd aws-ai-capacity
uv sync

# Copy environment template
cp .env.example .env

Configuration

The agent uses AWS Bedrock with Claude Opus 4.5 (anthropic.claude-opus-4-5-20251101-v1:0) by default.

Required AWS permissions:

• bedrock:InvokeModel - For Claude model access
• sagemaker:SearchTrainingPlanOfferings - Query available plans
• sagemaker:ListTrainingPlans - List existing plans
• sagemaker:DescribeTrainingPlan - Get plan details
• ec2:DescribeCapacityReservations - Query reservations
• ec2:DescribeInstanceTypeOfferings - Check availability
• ec2:DescribeInstances - List running GPU instances
• ec2:DescribeSpotPriceHistory - Spot price trends
• ec2:GetSpotPlacementScores - Spot capacity scores
• ec2:RunInstances / ec2:TerminateInstances - On-demand capacity check (opt-in)
• ssm:GetParameter - AMI lookup for on-demand check

Environment Variables

For development or to override defaults, create a .env file:

# AWS credentials
AWS_PROFILE=myprofile
AWS_REGION=us-east-1

# Override the default model
BEDROCK_MODEL_ID=us.anthropic.claude-opus-4-5-20251101-v1:0
BEDROCK_REGION=us-east-1

# Other options
AGENT_MAX_RETRIES=3
LOG_LEVEL=DEBUG

See .env.example for all available options.

Usage

Chat Interface

Start the Chainlit UI:

uv run aws-ai-capacity serve

Then open http://localhost:8000

CLI Commands

# Single query
uv run aws-ai-capacity chat "What p5 training plans are available?"

# Debug mode - shows all tool calls made by the agent
uv run aws-ai-capacity chat "What p5 training plans are available?" --debug

# Generate reports
uv run aws-ai-capacity report daily
uv run aws-ai-capacity report availability
uv run aws-ai-capacity report training-plans
uv run aws-ai-capacity report spot

# Save report to file
uv run aws-ai-capacity report daily -o capacity-report.md

# Generate all reports (for cron)
uv run aws-ai-capacity cron-report -d ./reports

# List GPU instance types
uv run aws-ai-capacity list-instance-types

Cron Job Example

# Daily capacity report at 8am
0 8 * * * cd /path/to/ai-capacity && uv run aws-ai-capacity cron-report -d ./reports

Example Questions

• "What p5.48xlarge training plans are available for the next week?"
• "Show me all active capacity reservations"
• "Which regions have p4d.24xlarge available?"
• "Compare the specs of p5 vs p4d instances"
• "What are the spot placement scores for p5.48xlarge across all regions?"
• "Show me spot price history for p4d.24xlarge in us-east-1"
• "Show me all running GPU instances in my account"
• "Check on-demand capacity for g5.xlarge in us-east-1"

Project Structure

ai-capacity/
├── src/ai_capacity/
│   ├── agent/          # Pydantic AI agent definition
│   ├── tools/          # AWS API tools (SageMaker, EC2, Spot/On-Demand)
│   ├── cli/            # Typer CLI commands
│   ├── ui/             # Chainlit chat interface
│   └── config.py       # Settings management
├── chainlit.md         # Chat welcome message
└── .env.example        # Environment template

Development

# Install with dev dependencies
uv sync --all-extras

# Run tests
uv run pytest

# Lint
uv run ruff check .