amazon-science/ammo

amazon-science/ammo

Description: Multi-agent system that autonomously optimizes vLLM GPU kernels for a specific deployment — model, hardware, dtype, parallelism — over multi-hour unattended campaigns

Language: Python

License: NOASSERTION

Stars: 0

Forks: 0

Open issues: 1

Created: 2026-08-19T20:59:32Z

Pushed: 2026-08-27T18:15:56Z

Default branch: main

Fork: no

Archived: no

README:

AMMO Sessions Server

> This code is being released solely for academic and scientific > reproducibility purposes, in support of the methods and findings described > in the associated publication. Pull requests are not being accepted in > order to maintain the code exactly as it was used in the paper.

AMMO (Agentic Model-on-Machine Optimizer) is a multi-agent system that autonomously optimizes vLLM GPU kernels for a specific deployment — model, hardware, dtype, parallelism — over multi-hour unattended campaigns. This server creates and manages the isolated sessions those campaigns run in.

Each session is a git worktree of vLLM with its own venv, a slice of the host's GPUs, and an AI CLI (Claude Code or Codex) running in a hardened terminal you drive from your browser.

What You Need

• Hardware: Linux host with an NVIDIA GPU (CUDA capability 8.0+ — A100,

L40S, H100, H200, B200), CUDA 12.0+, and the NVIDIA container toolkit.

• An AI CLI subscription or API key: sessions run Claude Code

(ANTHROPIC_API_KEY) or Codex (OPENAI_API_KEY / an existing ~/.codex/auth.json login). The AMMO campaign template pins a high-effort frontier model and spawns agent teams; a full optimization campaign runs for hours and consumes API credits accordingly. The shipped model-per-agent assignments are the recommended setup, and every pin is user-changeable — see [Which Model Runs Which Agent](docs/AMMO_DEEP_DIVE.md#which-model-runs-which-agent) for the roster, the reasoning, and concrete ways to reduce cost.

• Disk: the Docker image is ~42 GB (CUDA toolchain + a Python-ready

vLLM checkout).

Host prerequisites for GPU profiling

Campaigns drive NVIDIA Nsight profilers (ncu, nsys) from unprivileged session users. On a stock host the NVIDIA driver restricts GPU performance counters to admin users, so ncu fails with ERR_NVGPUCTRPERM and the profiling stages of a campaign cannot run. Enable counters once per host:

echo 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' | \
sudo tee /etc/modprobe.d/nvidia-profiling.conf
sudo sysctl -w kernel.perf_event_paranoid=2
# reboot (or reload the nvidia kernel modules) for the modprobe option to apply

Quick Start

git clone ammo-server
cd ammo-server

# Make your AI CLI credentials available to sessions
export ANTHROPIC_API_KEY=sk-ant-... # for Claude Code sessions
# and/or: log in to Codex once so ~/.codex/auth.json exists

./docker-build.sh && ./docker-run.sh --gpu all

# Verify
curl http://localhost:8000/health
# {"status": "healthy", "gpu_available": true, ...}

Then open http://localhost:8000/ui, click New Session, pick a HuggingFace model (the server auto-suggests TP/DP/dtype), and you get a browser terminal with the AI CLI running inside the session's vLLM worktree. To start an optimization campaign, tell it:

Use $ammo for model_id=Qwen/Qwen3-8B TP=1 dtype=bf16

The campaign runs unattended from there — profiling, debating candidates, writing kernels, validating, and shipping only what survives end-to-end measurement. The final deliverable is a REPORT.md you can view from the UI or fetch via GET /sessions/{id}/report.

Without Docker

pip install -r requirements.txt -r requirements.server.txt
python main.py # serves on :8000

The local path expects git, tmux, ttyd, and the AI CLIs on the host; the Docker image provisions all of this, so prefer it.

Using the Server

Web UI

http://localhost:8000/ui is the primary interface: create sessions from a modal with live HuggingFace model search, watch up to two terminals side-by-side, pause/resume/terminate, download session archives, and read optimization reports. Set AMMO_API_KEY to put the UI and API behind a login.

API

| Method | Endpoint | Description | |--------|----------|-------------| | POST | /sessions | Create a session (503 if insufficient GPUs, 429 at the per-client limit) | | GET | /sessions | List sessions | | GET | /sessions/{id} | Get session info | | POST | /sessions/{id}/pause | Pause (frees GPUs; state preserved) | | POST | /sessions/{id}/resume | Resume (conversation auto-continues) | | DELETE | /sessions/{id} | Terminate and clean up | | GET | /sessions/{id}/report | The campaign's REPORT.md | | POST | /sessions/{id}/prepare-download → GET .../download | Sanitized ZIP of the worktree | | GET | /health | GPU + vLLM build info (no auth) | | GET | /api/hf-model-config/{model_id} | Suggested TP/DP/dtype for an HF model |

import requests

r = requests.post("http://localhost:8000/sessions", json={
"model_name": "Qwen/Qwen3-8B",
"dtype": "bf16",
"gpu_count": 1,
"cli_tool": "claude",
"initial_prompt": "Use $ammo for model_id=Qwen/Qwen3-8B TP=1 dtype=bf16",
}, headers={"X-Client-ID": "my-client"})
print(r.json()["session_id"])

Session Lifecycle

CREATE → ACTIVE ↔ PAUSED → TERMINATED

Sessions survive browser disconnects (the terminal runs in tmux). Pause frees the GPUs and preserves the worktree; resume restarts the CLI with the conversation intact. Idle sessions auto-pause after SESSION_INACTIVITY_TIMEOUT_MINS (default 24 h).

Custom vLLM Forks

Point a session at your own vLLM fork by passing vllm_fork_url (and vllm_fork_token for private forks) at create time — the server clones it and runs a full source build before the session opens. See [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md#custom-vllm-forks).

Configuration

The variables you are most likely to set:

| Variable | Default | Description | |----------|---------|-------------| | ANTHROPIC_API_KEY | — | Claude Code auth, inherited by sessions | | OPENAI_API_KEY / CODEX_AUTH_JSON_PATH | — / ~/.codex/auth.json | Codex auth | | AMMO_API_KEY | — | Gate the API and UI behind a key (unset = open) | | MAX_SESSIONS_PER_CLIENT | 8 | Concurrent active sessions per client | | SESSION_S3_BUCKET | — | Optional: persist paused sessions to S3 (enables cross-host resume) | | SESSION_INACTIVITY_TIMEOUT_MINS | 1440 | Auto-pause idle sessions |...