← back to Exo
exo-bench (Benchmark model pp & tg speed) (#1099)
077b1bc73214716bce35442eb8eb0ba661088ebd · 2026-01-06 17:39:09 +0000 · rltakashige
## Motivation
This PR implements benchmarking in the style of llama-bench. The main
difficulty here is the fact that exo is not a library - it exposes an
endpoint. This means that benchmarking numbers will be inaccurate if the
API is measured.
The solution assumes nodes are set up with uv run exo (or via the app),
and then hits the new endpoint /bench/chat/completions to retrieve
generation statistics directly from mlx_lm.
<!-- Why is this change needed? What problem does it solve? -->
This will allow us to release benchmarks for models and perform
regression tests.
TODO: Performance benchmarking.
<!-- If it fixes an open issue, please link to the issue here -->
## Changes
<!-- Describe what you changed in detail -->
- Adds /bench/chat/completions endpoint
- Adds BenchChatCompletion/Response
- Adds a logits processor to prevent response from ending early
- Adds a "Prompt Sizer" which downloads the tokenizer and dynamically
adjusts the prompt of "a" to fit the desired prompt size.
- Reduce prefill step size to 2048 for now (in future, dynamically
adjust this value)
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
Benchmarked Llama, Qwen, DeepSeek and Kimi models. Will require several
fixes to run consistently on all configurations (to be done in the
future).
Manually tested the normal API to verify chat requests complete as
expected.
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
Not really possible. Type checker passes.
Files touched
D .github/benchmark-dashboard/README.mdD .github/benchmark-dashboard/index.htmlD .github/configs/README.mdD .github/configs/bench_config.yamlD .github/configs/bench_simple.yamlD .github/scripts/bench.pyD .github/scripts/build_matrix.pyD .github/workflows/BENCH_USAGE.mdD .github/workflows/bench.ymlA bench/exo_bench.pyM pyproject.tomlM src/exo/master/api.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/worker/runner_response.pyM src/exo/worker/engines/mlx/generator/generate.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/runner.pyD tmp/disable_bridge_enable_dhcp.sh
Diff
commit 077b1bc73214716bce35442eb8eb0ba661088ebd
Author: rltakashige <rl.takashige@gmail.com>
Date: Tue Jan 6 17:39:09 2026 +0000
exo-bench (Benchmark model pp & tg speed) (#1099)
## Motivation
This PR implements benchmarking in the style of llama-bench. The main
difficulty here is the fact that exo is not a library - it exposes an
endpoint. This means that benchmarking numbers will be inaccurate if the
API is measured.
The solution assumes nodes are set up with uv run exo (or via the app),
and then hits the new endpoint /bench/chat/completions to retrieve
generation statistics directly from mlx_lm.
<!-- Why is this change needed? What problem does it solve? -->
This will allow us to release benchmarks for models and perform
regression tests.
TODO: Performance benchmarking.
<!-- If it fixes an open issue, please link to the issue here -->
## Changes
<!-- Describe what you changed in detail -->
- Adds /bench/chat/completions endpoint
- Adds BenchChatCompletion/Response
- Adds a logits processor to prevent response from ending early
- Adds a "Prompt Sizer" which downloads the tokenizer and dynamically
adjusts the prompt of "a" to fit the desired prompt size.
- Reduce prefill step size to 2048 for now (in future, dynamically
adjust this value)
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
Benchmarked Llama, Qwen, DeepSeek and Kimi models. Will require several
fixes to run consistently on all configurations (to be done in the
future).
Manually tested the normal API to verify chat requests complete as
expected.
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
Not really possible. Type checker passes.
---
.github/benchmark-dashboard/README.md | 159 ---
.github/benchmark-dashboard/index.html | 1641 ----------------------
.github/configs/README.md | 186 ---
.github/configs/bench_config.yaml | 49 -
.github/configs/bench_simple.yaml | 125 --
.github/scripts/bench.py | 1399 ------------------
.github/scripts/build_matrix.py | 70 -
.github/workflows/BENCH_USAGE.md | 156 --
.github/workflows/bench.yml | 305 ----
bench/exo_bench.py | 526 +++++++
pyproject.toml | 2 +-
src/exo/master/api.py | 70 +
src/exo/shared/types/api.py | 21 +
src/exo/shared/types/chunks.py | 2 +
src/exo/shared/types/worker/runner_response.py | 3 +-
src/exo/worker/engines/mlx/generator/generate.py | 71 +-
src/exo/worker/engines/mlx/utils_mlx.py | 10 +
src/exo/worker/runner/runner.py | 3 +
tmp/disable_bridge_enable_dhcp.sh | 24 -
19 files changed, 696 insertions(+), 4126 deletions(-)
diff --git a/.github/benchmark-dashboard/README.md b/.github/benchmark-dashboard/README.md
deleted file mode 100644
index 1db78344..00000000
--- a/.github/benchmark-dashboard/README.md
+++ /dev/null
@@ -1,159 +0,0 @@
-# EXO Benchmark Dashboard
-
-A fully self-contained, browser-based dashboard for tracking EXO benchmark performance over time.
-
-## Features
-
-- 📊 **Success Rate Tracking**: Monitor cluster reliability across commits
-- ⚡ **Response Time Analysis**: Track average request completion times
-- 🎯 **Throughput Metrics**: Tokens per second visualization
-- 📈 **Request Distribution**: Success/failure breakdown over time
-- 🔄 **Auto-Refresh**: Updates every 60 seconds
-- 📺 **TV-Ready**: Large, clear visualizations perfect for display
-- 🔐 **Secure**: Credentials stored in browser localStorage only
-- 🌐 **No Backend**: Directly accesses S3 from the browser
-
-## Quick Start
-
-### Option 1: Direct File Access (Simplest)
-
-Just open the HTML file directly in your browser:
-
-```bash
-open .github/benchmark-dashboard/index.html
-```
-
-Then click "Configure AWS Credentials" and enter your keys.
-
-### Option 2: URL Parameters (For Quick Setup)
-
-```bash
-# Serve with credentials in URL (they'll be moved to localStorage)
-open ".github/benchmark-dashboard/index.html?accessKey=YOUR_KEY&secretKey=YOUR_SECRET®ion=us-east-1"
-```
-
-The credentials will be saved to localStorage and removed from the URL immediately.
-
-### Option 3: Simple HTTP Server
-
-```bash
-# From repo root
-python3 -m http.server 8080
-
-# Then open: http://localhost:8080/.github/benchmark-dashboard/
-```
-
-## AWS Credentials
-
-The dashboard needs read-only access to the `exo-benchmark-results` S3 bucket.
-
-### Required IAM Permissions
-
-```json
-{
- "Version": "2012-10-17",
- "Statement": [
- {
- "Effect": "Allow",
- "Action": [
- "s3:GetObject",
- "s3:ListBucket"
- ],
- "Resource": [
- "arn:aws:s3:::exo-benchmark-results",
- "arn:aws:s3:::exo-benchmark-results/*"
- ]
- }
- ]
-}
-```
-
-### Security Notes
-
-- ✅ Credentials stored in browser `localStorage` only
-- ✅ Never sent to any server (except AWS)
-- ✅ All S3 access happens client-side
-- ✅ Use read-only IAM credentials
-- ⚠️ Don't commit credentials to git
-- ⚠️ Use a dedicated read-only IAM user
-
-## TV/Kiosk Mode
-
-For permanent display on a TV:
-
-### macOS
-```bash
-open -a "Google Chrome" --args --kiosk ".github/benchmark-dashboard/index.html"
-```
-
-### Linux
-```bash
-chromium-browser --kiosk --app="file://$(pwd)/.github/benchmark-dashboard/index.html"
-```
-
-### Auto-start on Boot
-
-Create a simple startup script:
-
-```bash
-#!/bin/bash
-# /usr/local/bin/start-benchmark-dashboard.sh
-
-cd /path/to/exo
-python3 -m http.server 8080 &
-sleep 2
-chromium-browser --kiosk http://localhost:8080/.github/benchmark-dashboard/
-```
-
-## Data Displayed
-
-### Summary Cards
-- **Latest Success Rate**: Most recent benchmark success percentage with trend
-- **Avg Response Time**: Latest average response time in ms with trend
-- **Total Benchmarks**: Count of all benchmarks run
-- **Active Configurations**: Number of unique benchmark configs
-
-### Charts
-1. **Success Rate Over Time**: Line chart showing reliability trends
-2. **Average Response Time**: Performance over time (lower is better)
-3. **Throughput**: Tokens/second metric (higher is better)
-4. **Request Distribution**: Stacked bar chart of successes/failures
-
-## How It Works
-
-1. **Loads AWS SDK**: Uses AWS SDK for JavaScript (browser version)
-2. **Lists S3 Objects**: Fetches all files from `s3://exo-benchmark-results/bench/`
-3. **Downloads Results**: Fetches each JSON result file
-4. **Parses & Visualizes**: Uses Chart.js to create interactive charts
-5. **Auto-Refreshes**: Polls S3 every 60 seconds for new results
-
-## Customization
-
-To modify the dashboard:
-
-1. Edit `index.html`
-2. Adjust `REFRESH_INTERVAL` for different polling frequency
-3. Modify chart colors/styles in the Chart.js configuration
-4. Add new metrics by extending the results parsing
-
-## Troubleshooting
-
-**"AWS credentials not configured"**
-- Click "Configure AWS Credentials" and enter your keys
-
-**"Error loading benchmark data"**
-- Check AWS credentials are correct
-- Verify S3 bucket name is `exo-benchmark-results`
-- Ensure IAM user has read permissions
-- Check browser console for detailed errors
-
-**"No benchmark results found"**
-- Wait for benchmark workflows to run
-- Verify results are being uploaded to S3
-- Check S3 bucket has files in `bench/` prefix
-
-**Charts not updating**
-- Check browser console for errors
-- Verify network connectivity to S3
-- Try refreshing the page manually
-
diff --git a/.github/benchmark-dashboard/index.html b/.github/benchmark-dashboard/index.html
deleted file mode 100644
index 5f72a831..00000000
--- a/.github/benchmark-dashboard/index.html
+++ /dev/null
@@ -1,1641 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-<head>
- <meta charset="UTF-8">
- <meta name="viewport" content="width=device-width, initial-scale=1.0">
- <title>EXO Benchmark Dashboard</title>
- <script src="https://cdn.jsdelivr.net/npm/chart.js@4.4.0/dist/chart.umd.min.js"></script>
- <script src="https://cdn.jsdelivr.net/npm/luxon@3.4.4/build/global/luxon.min.js"></script>
- <script src="https://cdn.jsdelivr.net/npm/chartjs-adapter-luxon@1.3.1/dist/chartjs-adapter-luxon.umd.min.js"></script>
- <script src="https://cdn.plot.ly/plotly-2.27.0.min.js"></script>
- <script src="https://sdk.amazonaws.com/js/aws-sdk-2.1691.0.min.js"></script>
- <style>
- * {
- margin: 0;
- padding: 0;
- box-sizing: border-box;
- }
-
- body {
- font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
- background: linear-gradient(135deg, #0a0a0a 0%, #1a1a2e 100%);
- color: #ffffff;
- padding: 2rem;
- min-height: 100vh;
- }
-
- .header {
- text-align: center;
- margin-bottom: 3rem;
- padding: 2rem;
- background: rgba(255, 255, 255, 0.05);
- border-radius: 20px;
- backdrop-filter: blur(10px);
- }
-
- .header h1 {
- font-size: 3.5rem;
- font-weight: 700;
- background: linear-gradient(90deg, #00d4ff, #7b2ff7);
- -webkit-background-clip: text;
- -webkit-text-fill-color: transparent;
- margin-bottom: 0.5rem;
- }
-
- .header .subtitle {
- font-size: 1.2rem;
- color: #aaaaaa;
- margin-top: 0.5rem;
- }
-
- .credentials-banner {
- background: rgba(255, 165, 0, 0.1);
- border: 2px solid rgba(255, 165, 0, 0.3);
- border-radius: 10px;
- padding: 1rem;
- margin-bottom: 2rem;
- text-align: center;
- }
-
- .credentials-banner button {
- background: #00d4ff;
- color: #0a0a0a;
- border: none;
- padding: 0.5rem 1.5rem;
- border-radius: 8px;
- font-size: 1rem;
- font-weight: 600;
- cursor: pointer;
- margin: 0.5rem;
- }
-
- .credentials-banner button:hover {
- background: #00b8e6;
- }
-
- .modal {
- display: none;
- position: fixed;
- top: 0;
- left: 0;
- right: 0;
- bottom: 0;
- background: rgba(0, 0, 0, 0.8);
- z-index: 1000;
- align-items: center;
- justify-content: center;
- }
-
- .modal.active {
- display: flex;
- }
-
- .modal-content {
- background: #1a1a2e;
- padding: 2rem;
- border-radius: 15px;
- max-width: 500px;
- width: 90%;
- border: 1px solid rgba(255, 255, 255, 0.1);
- }
-
- .modal-content h2 {
- margin-bottom: 1rem;
- color: #00d4ff;
- }
-
- .modal-content input {
- width: 100%;
- padding: 0.75rem;
- margin: 0.5rem 0;
- background: rgba(255, 255, 255, 0.05);
- border: 1px solid rgba(255, 255, 255, 0.1);
- border-radius: 8px;
- color: #ffffff;
- font-size: 1rem;
- }
-
- .modal-content .button-group {
- display: flex;
- gap: 1rem;
- margin-top: 1.5rem;
- }
-
- .modal-content button {
- flex: 1;
- padding: 0.75rem;
- border: none;
- border-radius: 8px;
- font-size: 1rem;
- font-weight: 600;
- cursor: pointer;
- }
-
- .modal-content .btn-primary {
- background: #00d4ff;
- color: #0a0a0a;
- }
-
- .modal-content .btn-secondary {
- background: rgba(255, 255, 255, 0.1);
- color: #ffffff;
- }
-
- .stats-grid {
- display: grid;
- grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
- gap: 1.5rem;
- margin-bottom: 3rem;
- }
-
- .stat-card {
- background: rgba(255, 255, 255, 0.05);
- border-radius: 15px;
- padding: 1.5rem;
- backdrop-filter: blur(10px);
- border: 1px solid rgba(255, 255, 255, 0.1);
- transition: transform 0.3s ease, box-shadow 0.3s ease;
- }
-
- .stat-card:hover {
- transform: translateY(-5px);
- box-shadow: 0 10px 30px rgba(0, 212, 255, 0.3);
- }
-
- .stat-card .label {
- font-size: 0.9rem;
- color: #aaaaaa;
- text-transform: uppercase;
- letter-spacing: 1px;
- margin-bottom: 0.5rem;
- }
-
- .stat-card .value {
- font-size: 2.5rem;
- font-weight: 700;
- color: #00d4ff;
- }
-
- .stat-card .trend {
- font-size: 0.9rem;
- margin-top: 0.5rem;
- }
-
- .trend.up {
- color: #00ff88;
- }
-
- .trend.down {
- color: #ff4444;
- }
-
- .charts-container {
- display: grid;
- grid-template-columns: 1fr;
- gap: 2rem;
- }
-
- .chart-card {
- background: rgba(255, 255, 255, 0.05);
- border-radius: 20px;
- padding: 2rem;
- backdrop-filter: blur(10px);
- border: 1px solid rgba(255, 255, 255, 0.1);
- }
-
- .chart-card h2 {
- font-size: 1.5rem;
- margin-bottom: 1.5rem;
- color: #00d4ff;
- }
-
- .chart-wrapper {
- position: relative;
- height: 400px;
- }
-
- .loading {
- text-align: center;
- padding: 4rem;
- font-size: 1.5rem;
- color: #aaaaaa;
- }
-
- .error {
- text-align: center;
- padding: 4rem;
- font-size: 1.2rem;
- color: #ff4444;
- background: rgba(255, 68, 68, 0.1);
- border-radius: 15px;
- }
-
- .last-update {
- text-align: center;
- margin-top: 2rem;
- color: #666;
- font-size: 0.9rem;
- }
-
- .summary-table {
- background: rgba(255, 255, 255, 0.05);
- border-radius: 20px;
- padding: 2rem;
- backdrop-filter: blur(10px);
- border: 1px solid rgba(255, 255, 255, 0.1);
- margin-bottom: 3rem;
- overflow-x: auto;
- }
-
- .summary-table h2 {
- font-size: 1.5rem;
- margin-bottom: 1.5rem;
- color: #00d4ff;
- }
-
- .summary-table table {
- width: 100%;
- border-collapse: collapse;
- }
-
- .summary-table th {
- background: rgba(0, 212, 255, 0.1);
- padding: 1rem;
- text-align: left;
- font-weight: 600;
- color: #00d4ff;
- border-bottom: 2px solid rgba(0, 212, 255, 0.3);
- white-space: nowrap;
- }
-
- .summary-table td {
- padding: 0.75rem 1rem;
- border-bottom: 1px solid rgba(255, 255, 255, 0.05);
- }
-
- .summary-table tr:hover {
- background: rgba(255, 255, 255, 0.03);
- }
-
- .summary-table .config-name {
- font-family: 'Courier New', monospace;
- color: #ffd93d;
- font-size: 0.9rem;
- }
-
- .summary-table .metric-value {
- font-weight: 600;
- color: #00ff88;
- text-align: right;
- }
-
- .summary-table .metric-value.warning {
- color: #ff9800;
- }
-
- .summary-table .metric-value.error {
- color: #ff4444;
- }
- </style>
-</head>
-<body>
- <div class="credentials-banner" id="credentialsBanner" style="display: none;">
- <p>⚠️ AWS credentials not configured. Click below to set them.</p>
- <button onclick="showCredentialsModal()">Configure AWS Credentials</button>
- </div>
-
- <div class="modal" id="credentialsModal">
- <div class="modal-content">
- <h2>Configure AWS Credentials</h2>
- <p style="color: #aaa; margin-bottom: 1rem;">Enter your AWS credentials to access benchmark results from S3.</p>
- <input type="text" id="accessKeyInput" placeholder="AWS Access Key ID" />
- <input type="password" id="secretKeyInput" placeholder="AWS Secret Access Key" />
- <input type="text" id="regionInput" placeholder="AWS Region (default: us-east-1)" value="us-east-1" />
- <div class="button-group">
- <button class="btn-primary" onclick="saveCredentials()">Save & Load Data</button>
- <button class="btn-secondary" onclick="closeCredentialsModal()">Cancel</button>
- </div>
- <p style="color: #666; font-size: 0.85rem; margin-top: 1rem;">
- Credentials are stored in browser localStorage only. They never leave your device.
- </p>
- </div>
- </div>
-
- <div class="header">
- <h1>🚀 EXO Benchmark Dashboard</h1>
- <p class="subtitle">Real-time performance tracking across commits</p>
- <div class="last-update" id="lastUpdate">Loading...</div>
- </div>
-
- <div class="stats-grid" id="statsGrid">
- <div class="stat-card">
- <div class="label">Latest Success Rate</div>
- <div class="value" id="latestSuccessRate">--%</div>
- <div class="trend" id="successRateTrend"></div>
- </div>
- <div class="stat-card">
- <div class="label">Avg Response Time</div>
- <div class="value" id="latestAvgTime">-- ms</div>
- <div class="trend" id="avgTimeTrend"></div>
- </div>
- <div class="stat-card">
- <div class="label">Time to First Token</div>
- <div class="value" id="latestTTFT">-- ms</div>
- <div class="trend" id="ttftTrend"></div>
- </div>
- <div class="stat-card">
- <div class="label">Decode Speed</div>
- <div class="value" id="latestDecodeTPS">-- t/s</div>
- <div class="trend" id="decodeTpsTrend"></div>
- </div>
- <div class="stat-card">
- <div class="label">Total Benchmarks</div>
- <div class="value" id="totalBenchmarks">--</div>
- </div>
- <div class="stat-card">
- <div class="label">Active Configurations</div>
- <div class="value" id="activeConfigs">--</div>
- </div>
- </div>
-
- <div class="summary-table">
- <h2>📋 All Tests Summary</h2>
- <table>
- <thead>
- <tr>
- <th>Name</th>
- <th>Strategy</th>
- <th>Success Rate</th>
- <th>Prefill Time</th>
- <th>ms per token</th>
- </tr>
- </thead>
- <tbody id="summaryTableBody">
- <tr>
- <td colspan="5" style="text-align: center; color: #aaa;">Loading...</td>
- </tr>
- </tbody>
- </table>
- </div>
-
- <div style="text-align: center; margin: 2rem 0;">
- <button id="showDetailsButton" onclick="loadDetailedCharts()" style="
- background: linear-gradient(135deg, #00d4ff, #7b2ff7);
- color: #ffffff;
- border: none;
- padding: 1rem 2rem;
- border-radius: 12px;
- font-size: 1.1rem;
- font-weight: 600;
- cursor: pointer;
- box-shadow: 0 4px 15px rgba(0, 212, 255, 0.4);
- transition: transform 0.2s ease, box-shadow 0.2s ease;
- display: none;
- " onmouseover="this.style.transform='translateY(-2px)'; this.style.boxShadow='0 6px 20px rgba(0, 212, 255, 0.6)'" onmouseout="this.style.transform=''; this.style.boxShadow='0 4px 15px rgba(0, 212, 255, 0.4)'">
- 📊 Load Detailed Charts
- </button>
- </div>
-
- <div class="charts-container" id="chartsContainer">
- <div class="loading">Loading benchmark data...</div>
- </div>
-
- <script>
- const BUCKET_NAME = 'exo-benchmark-results';
- const REFRESH_INTERVAL = 60000; // 1 minute
-
- let charts = {};
- let allResults = [];
- let s3Client = null;
- let detailedChartsLoaded = false;
-
- // Check for credentials in URL parameters (for easy setup)
- const urlParams = new URLSearchParams(window.location.search);
- if (urlParams.get('accessKey')) {
- localStorage.setItem('aws_access_key_id', urlParams.get('accessKey'));
- localStorage.setItem('aws_secret_access_key', urlParams.get('secretKey'));
- localStorage.setItem('aws_region', urlParams.get('region') || 'us-east-1');
- // Remove credentials from URL for security
- window.history.replaceState({}, document.title, window.location.pathname);
- }
-
- function getStoredCredentials() {
- return {
- accessKeyId: localStorage.getItem('aws_access_key_id'),
- secretAccessKey: localStorage.getItem('aws_secret_access_key'),
- region: localStorage.getItem('aws_region') || 'us-east-1'
- };
- }
-
- function hasCredentials() {
- const creds = getStoredCredentials();
- return creds.accessKeyId && creds.secretAccessKey;
- }
-
- function initializeS3Client() {
- const creds = getStoredCredentials();
- if (!creds.accessKeyId || !creds.secretAccessKey) {
- return null;
- }
-
- AWS.config.update({
- accessKeyId: creds.accessKeyId,
- secretAccessKey: creds.secretAccessKey,
- region: creds.region
- });
-
- return new AWS.S3();
- }
-
- function showCredentialsModal() {
- document.getElementById('credentialsModal').classList.add('active');
- const creds = getStoredCredentials();
- document.getElementById('accessKeyInput').value = creds.accessKeyId || '';
- document.getElementById('secretKeyInput').value = creds.secretAccessKey || '';
- document.getElementById('regionInput').value = creds.region || 'us-east-1';
- }
-
- function closeCredentialsModal() {
- document.getElementById('credentialsModal').classList.remove('active');
- }
-
- function saveCredentials() {
- const accessKey = document.getElementById('accessKeyInput').value.trim();
- const secretKey = document.getElementById('secretKeyInput').value.trim();
- const region = document.getElementById('regionInput').value.trim() || 'us-east-1';
-
- if (!accessKey || !secretKey) {
- alert('Please enter both Access Key ID and Secret Access Key');
- return;
- }
-
- localStorage.setItem('aws_access_key_id', accessKey);
- localStorage.setItem('aws_secret_access_key', secretKey);
- localStorage.setItem('aws_region', region);
-
- closeCredentialsModal();
- document.getElementById('credentialsBanner').style.display = 'none';
-
- // Reinitialize and load data
- s3Client = initializeS3Client();
- loadAndDisplayData();
- }
-
- async function fetchBenchmarkResults() {
- if (!s3Client) {
- throw new Error('AWS credentials not configured');
- }
-
- try {
- // List all objects in the bench/ prefix
- const listParams = {
- Bucket: BUCKET_NAME,
- Prefix: 'bench/',
- };
-
- const listResponse = await s3Client.listObjectsV2(listParams).promise();
-
- if (!listResponse.Contents || listResponse.Contents.length === 0) {
- return [];
- }
-
- // Fetch each result file
- const results = [];
- for (const obj of listResponse.Contents) {
- if (!obj.Key.endsWith('.json')) continue;
-
- try {
- const getParams = {
- Bucket: BUCKET_NAME,
- Key: obj.Key
- };
- const data = await s3Client.getObject(getParams).promise();
- const result = JSON.parse(data.Body.toString('utf-8'));
- results.push(result);
- } catch (error) {
- console.error(`Error fetching ${obj.Key}:`, error);
- }
- }
-
- // Sort by timestamp
- results.sort((a, b) =>
- a.metadata.benchmark_completed_at - b.metadata.benchmark_completed_at
- );
-
- return results;
-
- } catch (error) {
- console.error('Error fetching from S3:', error);
- throw error;
- }
- }
-
- // Chart.js default configuration
- Chart.defaults.color = '#ffffff';
- Chart.defaults.borderColor = 'rgba(255, 255, 255, 0.1)';
-
- const chartOptions = {
- responsive: true,
- maintainAspectRatio: false,
- plugins: {
- legend: {
- display: true,
- labels: {
- color: '#ffffff',
- font: { size: 14 }
- }
- }
- },
- scales: {
- x: {
- type: 'time',
- time: {
- unit: 'day',
- displayFormats: {
- day: 'MMM dd'
- }
- },
- ticks: { color: '#aaaaaa' },
- grid: { color: 'rgba(255, 255, 255, 0.05)' }
- },
- y: {
- ticks: { color: '#aaaaaa' },
- grid: { color: 'rgba(255, 255, 255, 0.05)' }
- }
- }
- };
-
- function updateSummaryTable(results) {
- const tableBody = document.getElementById('summaryTableBody');
-
- if (results.length === 0) {
- tableBody.innerHTML = '<tr><td colspan="5" style="text-align: center; color: #aaa;">No test data available</td></tr>';
- return;
- }
-
- let rows = '';
-
- results.forEach((result, idx) => {
- const cluster = result.cluster || {};
- const config = result.configuration || {};
- const stages = config.stages || [];
- const metadata = result.metadata || {};
- const resultStages = result.results?.stages || [];
-
- // Get model IDs
- const modelIds = cluster.model_ids || ['unknown'];
- const modelName = modelIds.length === 1 ? modelIds[0] : `${modelIds.length} models`;
-
- // Get strategy (backwards compatible with old format)
- // New format: sharding + instance_meta, e.g. "Pipeline (MLX Ring)"
- // Old format: strategy field
- let strategy = 'N/A';
- if (cluster.strategy) {
- // Backwards compatibility: use old strategy field
- strategy = cluster.strategy;
- } else if (cluster.sharding || cluster.instance_meta) {
- // New format: combine sharding and instance_meta
- const sharding = cluster.sharding || '';
- const instanceMeta = cluster.instance_meta || '';
-
- // Format instance_meta: convert camelCase/PascalCase to readable format
- const formatInstanceMeta = (meta) => {
- if (!meta) return '';
- // Insert spaces before capital letters and handle common acronyms
- return meta
- .replace(/([A-Z])/g, ' $1')
- .trim()
- .replace(/\bMlx\b/g, 'MLX')
- .replace(/\bIbv\b/g, 'IBV');
- };
-
- if (sharding && instanceMeta) {
- strategy = `${sharding} (${formatInstanceMeta(instanceMeta)})`;
- } else if (sharding) {
- strategy = sharding;
- } else if (instanceMeta) {
- strategy = formatInstanceMeta(instanceMeta);
- }
- }
-
- // For each stage in the configuration, create a row
- stages.forEach((stageConfig, stageIdx) => {
- const resultStage = resultStages[stageIdx];
-
- if (!resultStage) return;
-
- // Format: stage_name: model [prompt_len/generation_len] iterations every time_between_requests secs
- const stageName = stageConfig.name || `Stage ${stageIdx + 1}`;
- const name = `${stageName}: ${modelName} [${stageConfig.prompt_length}/${stageConfig.generation_length}] ${stageConfig.iterations}× @ ${stageConfig.time_between_requests}s`;
-
- // Success Rate
- let successRate = 'N/A';
- let successRateClass = '';
- if (resultStage.success_rate !== null && resultStage.success_rate !== undefined) {
- const rate = resultStage.success_rate * 100;
- successRate = `${rate.toFixed(1)}%`;
-
- // Color code based on success rate
- if (rate >= 95) {
- successRateClass = 'metric-value';
- } else if (rate >= 80) {
- successRateClass = 'metric-value warning';
- } else {
- successRateClass = 'metric-value error';
- }
- }
-
- // Prefill Time (TTFT)
- let prefillTime = 'N/A';
- if (resultStage.avg_time_to_first_token !== null && resultStage.avg_time_to_first_token !== undefined) {
- const ttftMs = resultStage.avg_time_to_first_token * 1000;
- if (resultStage.std_time_to_first_token !== null && resultStage.std_time_to_first_token !== undefined) {
- const stdMs = resultStage.std_time_to_first_token * 1000;
- prefillTime = `${ttftMs.toFixed(0)} ± ${stdMs.toFixed(0)} ms`;
- } else {
- prefillTime = `${ttftMs.toFixed(0)} ms`;
- }
- }
-
- // ms per token (1000 / decode_tps)
- let msPerToken = 'N/A';
- let msPerTokenClass = '';
- if (resultStage.avg_ms_per_token !== null && resultStage.avg_ms_per_token !== undefined) {
- const ms = resultStage.avg_ms_per_token;
- if (resultStage.std_ms_per_token !== null && resultStage.std_ms_per_token !== undefined) {
- const stdMs = resultStage.std_ms_per_token;
- msPerToken = `${ms.toFixed(1)} ± ${stdMs.toFixed(1)} ms`;
- } else {
- msPerToken = `${ms.toFixed(1)} ms`;
- }
-
- // Color code based on performance
- if (ms < 50) {
- msPerTokenClass = 'metric-value';
- } else if (ms < 100) {
- msPerTokenClass = 'metric-value warning';
- } else {
- msPerTokenClass = 'metric-value error';
- }
- }
-
- rows += `
- <tr>
- <td class="config-name">${name}</td>
- <td class="metric-value">${strategy}</td>
- <td class="${successRateClass || 'metric-value'}">${successRate}</td>
- <td class="metric-value">${prefillTime}</td>
- <td class="${msPerTokenClass || 'metric-value'}">${msPerToken}</td>
- </tr>
- `;
- });
- });
-
- tableBody.innerHTML = rows || '<tr><td colspan="5" style="text-align: center; color: #aaa;">No valid test data</td></tr>';
- }
-
- function updateStats(results) {
- if (results.length === 0) return;
-
- const latest = results[results.length - 1];
- const latestStage = latest.results.stages[0];
-
- document.getElementById('latestSuccessRate').textContent =
- `${(latestStage.success_rate * 100).toFixed(1)}%`;
-
- document.getElementById('latestAvgTime').textContent =
- `${(latestStage.avg_time_per_request * 1000).toFixed(0)} ms`;
-
- // Display TTFT and Decode TPS
- if (latestStage.avg_time_to_first_token !== null && latestStage.avg_time_to_first_token !== undefined) {
- document.getElementById('latestTTFT').textContent =
- `${(latestStage.avg_time_to_first_token * 1000).toFixed(0)} ms`;
- } else {
- document.getElementById('latestTTFT').textContent = 'N/A';
- }
-
- if (latestStage.avg_decode_tps !== null && latestStage.avg_decode_tps !== undefined) {
- document.getElementById('latestDecodeTPS').textContent =
- `${latestStage.avg_decode_tps.toFixed(1)} t/s`;
- } else {
- document.getElementById('latestDecodeTPS').textContent = 'N/A';
- }
-
- document.getElementById('totalBenchmarks').textContent = results.length;
-
- const uniqueConfigs = new Set(results.map(r => r.metadata.config_file || 'unknown'));
- document.getElementById('activeConfigs').textContent = uniqueConfigs.size;
-
- // Calculate trends
- if (results.length >= 2) {
- const previous = results[results.length - 2].results.stages[0];
-
- const successDiff = latestStage.success_rate - previous.success_rate;
- const successTrend = document.getElementById('successRateTrend');
- if (successDiff > 0) {
- successTrend.textContent = `↑ ${(successDiff * 100).toFixed(1)}%`;
- successTrend.className = 'trend up';
- } else if (successDiff < 0) {
- successTrend.textContent = `↓ ${(Math.abs(successDiff) * 100).toFixed(1)}%`;
- successTrend.className = 'trend down';
- } else {
- successTrend.textContent = '→ No change';
- successTrend.className = 'trend';
- }
-
- const timeDiff = latestStage.avg_time_per_request - previous.avg_time_per_request;
- const timeTrend = document.getElementById('avgTimeTrend');
- if (timeDiff < 0) {
- timeTrend.textContent = `↑ ${(Math.abs(timeDiff) * 1000).toFixed(0)}ms faster`;
- timeTrend.className = 'trend up';
- } else if (timeDiff > 0) {
- timeTrend.textContent = `↓ ${(timeDiff * 1000).toFixed(0)}ms slower`;
- timeTrend.className = 'trend down';
- } else {
- timeTrend.textContent = '→ No change';
- timeTrend.className = 'trend';
- }
-
- // TTFT trend
- if (latestStage.avg_time_to_first_token !== null && previous.avg_time_to_first_token !== null) {
- const ttftDiff = latestStage.avg_time_to_first_token - previous.avg_time_to_first_token;
- const ttftTrend = document.getElementById('ttftTrend');
- if (ttftDiff < 0) {
- ttftTrend.textContent = `↑ ${(Math.abs(ttftDiff) * 1000).toFixed(0)}ms faster`;
- ttftTrend.className = 'trend up';
- } else if (ttftDiff > 0) {
- ttftTrend.textContent = `↓ ${(ttftDiff * 1000).toFixed(0)}ms slower`;
- ttftTrend.className = 'trend down';
- } else {
- ttftTrend.textContent = '→ No change';
- ttftTrend.className = 'trend';
- }
- }
-
- // Decode TPS trend
- if (latestStage.avg_decode_tps !== null && previous.avg_decode_tps !== null) {
- const tpsDiff = latestStage.avg_decode_tps - previous.avg_decode_tps;
- const tpsTrend = document.getElementById('decodeTpsTrend');
- if (tpsDiff > 0) {
- tpsTrend.textContent = `↑ ${tpsDiff.toFixed(1)} t/s faster`;
- tpsTrend.className = 'trend up';
- } else if (tpsDiff < 0) {
- tpsTrend.textContent = `↓ ${Math.abs(tpsDiff).toFixed(1)} t/s slower`;
- tpsTrend.className = 'trend down';
- } else {
- tpsTrend.textContent = '→ No change';
- tpsTrend.className = 'trend';
- }
- }
- }
-
- const updateTime = new Date(latest.metadata.benchmark_completed_at * 1000);
- document.getElementById('lastUpdate').textContent =
- `Last benchmark: ${updateTime.toLocaleString()}`;
- }
-
- function createAggregateChartsHTML() {
- const container = document.getElementById('chartsContainer');
-
- // Create aggregate charts section only
- let html = `
- <div style="margin-bottom: 3rem;">
- <h1 style="font-size: 2rem; margin-bottom: 2rem; color: #00d4ff;">📊 Aggregate Performance Metrics</h1>
-
- <div class="chart-card">
- <h2>📊 Success Rate Over Time</h2>
- <div class="chart-wrapper">
- <canvas id="successRateChart"></canvas>
- </div>
- </div>
-
- <div class="chart-card">
- <h2>⚡ Average Response Time</h2>
- <div class="chart-wrapper">
- <canvas id="responseTimeChart"></canvas>
- </div>
- </div>
-
- <div class="chart-card">
- <h2>🎯 Throughput (Tokens/Second)</h2>
- <div class="chart-wrapper">
- <canvas id="throughputChart"></canvas>
- </div>
- </div>
-
- <div class="chart-card">
- <h2>⏱️ Time to First Token (TTFT)</h2>
- <div class="chart-wrapper">
- <canvas id="ttftChart"></canvas>
- </div>
- </div>
-
- <div class="chart-card">
- <h2>🚀 Decode Speed (Tokens/Second)</h2>
- <div class="chart-wrapper">
- <canvas id="decodeTpsChart"></canvas>
- </div>
- </div>
-
- <div class="chart-card">
- <h2>📈 Request Success/Failure Distribution</h2>
- <div class="chart-wrapper">
- <canvas id="requestDistChart"></canvas>
- </div>
- </div>
- </div>
- `;
-
- container.innerHTML = html;
- }
-
- function createIndividualChartsHTML(results) {
- const container = document.getElementById('chartsContainer');
-
- // Append individual test sections to existing content
- let html = container.innerHTML;
- html += `<div id="individualChartsSection" style="margin-top: 4rem;">
- <h1 style="font-size: 2rem; margin-bottom: 2rem; color: #00d4ff;">🔬 Individual Test Results</h1>
- `;
-
- results.forEach((result, idx) => {
- const timestamp = new Date(result.metadata.benchmark_completed_at * 1000);
- const config = result.configuration || {};
- const stages = config.stages || [];
- const metadata = result.metadata || {};
- const cluster = result.cluster || {};
-
- html += `
- <div class="chart-card" style="margin-bottom: 2rem;">
- <h2 style="margin-bottom: 1rem;">Test #${idx + 1} - ${timestamp.toLocaleString()}</h2>
-
- <div style="background: rgba(0, 0, 0, 0.3); padding: 1rem; border-radius: 10px; margin-bottom: 1.5rem;">
- <h3 style="color: #00d4ff; margin-bottom: 0.75rem;">Configuration</h3>
- <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 0.75rem; font-size: 0.9rem;">
- <div><strong>Git Commit:</strong> ${metadata.git_commit || 'N/A'}</div>
- <div><strong>Config File:</strong> ${metadata.config_file || 'N/A'}</div>
- <div><strong>Expected Nodes:</strong> ${metadata.expected_nodes || 'N/A'}</div>
- <div><strong>Model IDs:</strong> ${cluster.model_ids ? cluster.model_ids.join(', ') : 'N/A'}</div>
- <div><strong>Instances:</strong> ${cluster.instance_count || 0}</div>
- <div><strong>Runners:</strong> ${cluster.runner_count || 0}</div>
- <div><strong>Duration:</strong> ${metadata.total_duration_s ? metadata.total_duration_s.toFixed(1) + 's' : 'N/A'}</div>
- </div>
-
- ${stages.length > 0 ? `
- <h4 style="color: #00d4ff; margin-top: 1rem; margin-bottom: 0.5rem;">Stages:</h4>
- <div style="display: grid; gap: 0.5rem;">
- ${stages.map((stage, sidx) => `
- <div style="background: rgba(255, 255, 255, 0.05); padding: 0.5rem; border-radius: 5px; font-size: 0.85rem;">
- <strong>${sidx + 1}. ${stage.name}</strong> -
- Prompt: ${stage.prompt_length} tokens,
- Gen: ${stage.generation_length} tokens,
- Iterations: ${stage.iterations},
- Interval: ${stage.time_between_requests}s
- </div>
- `).join('')}
- </div>
- ` : ''}
- </div>
-
- <div style="margin-bottom: 1.5rem;">
- <h3 style="color: #ffd93d; margin-bottom: 0.5rem;">📊 Request Success/Failure Breakdown</h3>
- <div id="requestBreakdownPlot_${idx}" style="width: 100%; height: 400px;"></div>
- </div>
-
- ${result.metrics && result.metrics.snapshots && result.metrics.snapshots.length > 0 ? `
- <div style="margin-bottom: 1.5rem;">
- <h3 style="color: #00ff88; margin-bottom: 0.5rem;">🖥️ Running Tasks Per Node</h3>
- <p style="color: #aaa; font-size: 0.85rem; margin-bottom: 0.5rem;">Interactive chart - click and drag to zoom, double-click to reset</p>
- <div id="nodeTasksPlot_${idx}" style="width: 100%; height: 500px;"></div>
- </div>
-
- <div style="margin-bottom: 1.5rem;">
- <h3 style="color: #ff8c00; margin-bottom: 0.5rem;">📦 Running Tasks Per Instance</h3>
- <p style="color: #aaa; font-size: 0.85rem; margin-bottom: 0.5rem;">Interactive chart - click and drag to zoom, double-click to reset</p>
- <div id="instanceTasksPlot_${idx}" style="width: 100%; height: 500px;"></div>
- </div>
-
- <div style="margin-bottom: 1.5rem;">
- <h3 style="color: #00d4ff; margin-bottom: 0.5rem;">📋 Cluster-wide Task Summary</h3>
- <p style="color: #aaa; font-size: 0.85rem; margin-bottom: 0.5rem;">Interactive chart - click and drag to zoom, double-click to reset</p>
- <div id="clusterTasksPlot_${idx}" style="width: 100%; height: 400px;"></div>
- </div>
-
- <div style="margin-bottom: 1.5rem;">
- <h3 style="color: #7b2ff7; margin-bottom: 0.5rem;">💾 Memory Usage Per Node</h3>
- <p style="color: #aaa; font-size: 0.85rem; margin-bottom: 0.5rem;">Interactive chart - click and drag to zoom, double-click to reset</p>
- <div id="memoryPlot_${idx}" style="width: 100%; height: 400px;"></div>
- </div>
- ` : '<p style="color: #aaa;">No metrics data available for this test.</p>'}
- </div>
- `;
- });
-
- html += `</div>`;
- container.innerHTML = html;
- }
-
- function createAggregateCharts(results) {
- createAggregateChartsHTML();
-
- const timestamps = results.map(r => new Date(r.metadata.benchmark_completed_at * 1000));
- const successRates = results.map(r => r.results.stages[0].success_rate * 100);
- const avgTimes = results.map(r => r.results.stages[0].avg_time_per_request * 1000);
- const throughput = results.map(r => {
- const stage = r.results.stages[0];
- return stage.avg_tokens_per_request / stage.avg_time_per_request;
- });
- const ttftValues = results.map(r => {
- const ttft = r.results.stages[0].avg_time_to_first_token;
- return ttft !== null && ttft !== undefined ? ttft * 1000 : null;
- });
- const decodeTpsValues = results.map(r => {
- const tps = r.results.stages[0].avg_decode_tps;
- return tps !== null && tps !== undefined ? tps : null;
- });
- const successful = results.map(r => r.results.stages[0].successful_requests);
- const failed = results.map(r => r.results.stages[0].failed_requests);
-
- charts.successRate = new Chart(
- document.getElementById('successRateChart'),
- {
- type: 'line',
- data: {
- labels: timestamps,
- datasets: [{
- label: 'Success Rate (%)',
- data: successRates,
- borderColor: '#00ff88',
- backgroundColor: 'rgba(0, 255, 136, 0.1)',
- fill: true,
- tension: 0.4,
- borderWidth: 3,
- pointRadius: 6,
- pointHoverRadius: 8,
- }]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- min: 0,
- max: 100,
- ticks: {
- ...chartOptions.scales.y.ticks,
- callback: (value) => value + '%'
- }
- }
- }
- }
- }
- );
-
- charts.responseTime = new Chart(
- document.getElementById('responseTimeChart'),
- {
- type: 'line',
- data: {
- labels: timestamps,
- datasets: [{
- label: 'Avg Response Time (ms)',
- data: avgTimes,
- borderColor: '#7b2ff7',
- backgroundColor: 'rgba(123, 47, 247, 0.1)',
- fill: true,
- tension: 0.4,
- borderWidth: 3,
- pointRadius: 6,
- pointHoverRadius: 8,
- }]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- ticks: {
- ...chartOptions.scales.y.ticks,
- callback: (value) => value.toFixed(0) + ' ms'
- }
- }
- }
- }
- }
- );
-
- charts.throughput = new Chart(
- document.getElementById('throughputChart'),
- {
- type: 'line',
- data: {
- labels: timestamps,
- datasets: [{
- label: 'Tokens/Second',
- data: throughput,
- borderColor: '#00d4ff',
- backgroundColor: 'rgba(0, 212, 255, 0.1)',
- fill: true,
- tension: 0.4,
- borderWidth: 3,
- pointRadius: 6,
- pointHoverRadius: 8,
- }]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- ticks: {
- ...chartOptions.scales.y.ticks,
- callback: (value) => value.toFixed(1) + ' t/s'
- }
- }
- }
- }
- }
- );
-
- charts.ttft = new Chart(
- document.getElementById('ttftChart'),
- {
- type: 'line',
- data: {
- labels: timestamps,
- datasets: [{
- label: 'Time to First Token (ms)',
- data: ttftValues,
- borderColor: '#ffd93d',
- backgroundColor: 'rgba(255, 217, 61, 0.1)',
- fill: true,
- tension: 0.4,
- borderWidth: 3,
- pointRadius: 6,
- pointHoverRadius: 8,
- }]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- ticks: {
- ...chartOptions.scales.y.ticks,
- callback: (value) => value.toFixed(0) + ' ms'
- }
- }
- }
- }
- }
- );
-
- charts.decodeTps = new Chart(
- document.getElementById('decodeTpsChart'),
- {
- type: 'line',
- data: {
- labels: timestamps,
- datasets: [{
- label: 'Decode Speed (tokens/s)',
- data: decodeTpsValues,
- borderColor: '#ff6b9d',
- backgroundColor: 'rgba(255, 107, 157, 0.1)',
- fill: true,
- tension: 0.4,
- borderWidth: 3,
- pointRadius: 6,
- pointHoverRadius: 8,
- }]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- ticks: {
- ...chartOptions.scales.y.ticks,
- callback: (value) => value.toFixed(1) + ' t/s'
- }
- }
- }
- }
- }
- );
-
- charts.requestDist = new Chart(
- document.getElementById('requestDistChart'),
- {
- type: 'bar',
- data: {
- labels: timestamps,
- datasets: [
- {
- label: 'Successful',
- data: successful,
- backgroundColor: 'rgba(0, 255, 136, 0.8)',
- borderColor: '#00ff88',
- borderWidth: 1,
- },
- {
- label: 'Failed',
- data: failed,
- backgroundColor: 'rgba(255, 68, 68, 0.8)',
- borderColor: '#ff4444',
- borderWidth: 1,
- }
- ]
- },
- options: {
- ...chartOptions,
- scales: {
- ...chartOptions.scales,
- y: {
- ...chartOptions.scales.y,
- stacked: true,
- },
- x: {
- ...chartOptions.scales.x,
- stacked: true,
- }
- }
- }
- }
- );
- }
-
- function createDetailedCharts(results) {
- createIndividualChartsHTML(results);
-
- // Create individual test charts for each result
- results.forEach((result, idx) => {
- createRequestBreakdownChart(result, idx);
- if (result.metrics && result.metrics.snapshots && result.metrics.snapshots.length > 0) {
- createIndividualTestCharts(result, idx);
- }
- });
- }
-
- function createRequestBreakdownChart(result, testIndex) {
- // Aggregate all requests from all stages
- const allRequests = [];
- if (result.results && result.results.stages) {
- result.results.stages.forEach(stage => {
- if (stage.requests) {
- allRequests.push(...stage.requests);
- }
- });
- }
-
- if (allRequests.length === 0) {
- return;
- }
-
- // Count successful requests
- const successfulCount = allRequests.filter(r => r.success).length;
-
- // Group failed requests by error reason
- const failureReasons = {};
- allRequests.filter(r => !r.success).forEach(r => {
- const reason = r.error || 'Unknown error';
- // Truncate long error messages for readability
- const shortReason = reason.length > 60 ? reason.substring(0, 57) + '...' : reason;
- failureReasons[shortReason] = (failureReasons[shortReason] || 0) + 1;
- });
-
- // Create bar chart data
- const labels = ['✅ Successful'];
- const values = [successfulCount];
- const colors = ['#00ff88'];
-
- // Add failure reasons
- Object.entries(failureReasons)
- .sort((a, b) => b[1] - a[1]) // Sort by count descending
- .forEach(([reason, count]) => {
- labels.push(`❌ ${reason}`);
- values.push(count);
- colors.push('#ff6b6b');
- });
-
- const trace = {
- x: values,
- y: labels,
- type: 'bar',
- orientation: 'h',
- marker: {
- color: colors,
- line: {
- color: colors.map(c => c === '#00ff88' ? '#00d47a' : '#ff5555'),
- width: 2
- }
- },
- text: values.map(v => v.toString()),
- textposition: 'outside',
- hovertemplate: '<b>%{y}</b><br>Count: %{x}<extra></extra>'
- };
-
- Plotly.newPlot(`requestBreakdownPlot_${testIndex}`, [trace], {
- paper_bgcolor: 'rgba(0,0,0,0)',
- plot_bgcolor: 'rgba(0,0,0,0.2)',
- font: { color: '#ffffff', size: 12 },
- xaxis: {
- title: 'Number of Requests',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- yaxis: {
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false,
- automargin: true
- },
- margin: { l: 250, r: 50, t: 20, b: 50 },
- showlegend: false
- }, {
- responsive: true,
- displayModeBar: true,
- modeBarButtonsToRemove: ['lasso2d', 'select2d', 'zoom2d', 'pan2d', 'zoomIn2d', 'zoomOut2d', 'autoScale2d'],
- displaylogo: false
- });
- }
-
- function createIndividualTestCharts(result, testIndex) {
- if (!result.metrics || !result.metrics.snapshots) {
- return;
- }
-
- const snapshots = result.metrics.snapshots;
- const startTime = result.metadata.benchmark_started_at;
-
- // Extract timestamps relative to benchmark start
- const relativeTimestamps = snapshots.map(s => (s.timestamp - startTime));
-
- // Distinct color palette (10 colors - very different from each other)
- const distinctColors = [
- '#FF6B6B', // Red
- '#4ECDC4', // Teal
- '#FFD93D', // Yellow
- '#6BCF7F', // Green
- '#A78BFA', // Purple
- '#FB8500', // Orange
- '#219EBC', // Blue
- '#F72585', // Pink
- '#95D5B2', // Mint
- '#FFB627', // Gold
- ];
-
- // 1. RUNNING TASKS PER NODE (Most Important!)
- if (snapshots[0].node_tasks && snapshots[0].node_tasks.length > 0) {
- const nodeIds = snapshots[0].node_tasks.map(n => n.node_id);
- const nodeTraces = nodeIds.map((nodeId, idx) => {
- const color = distinctColors[idx % distinctColors.length];
- return {
- x: relativeTimestamps,
- y: snapshots.map(s => {
- const nodeData = s.node_tasks.find(n => n.node_id === nodeId);
- return nodeData ? nodeData.running_tasks : 0;
- }),
- name: `...${nodeId.slice(-4)}`,
- type: 'scatter',
- mode: 'lines+markers',
- line: { color: color, width: 3 },
- marker: { size: 6 }
- };
- });
-
- Plotly.newPlot(`nodeTasksPlot_${testIndex}`, nodeTraces, {
- paper_bgcolor: 'rgba(0,0,0,0)',
- plot_bgcolor: 'rgba(0,0,0,0.2)',
- font: { color: '#ffffff' },
- xaxis: {
- title: 'Time (seconds)',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- yaxis: {
- title: 'Running Tasks',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- hovermode: 'closest',
- showlegend: true,
- legend: {
- bgcolor: 'rgba(0,0,0,0.5)',
- bordercolor: 'rgba(255,255,255,0.2)',
- borderwidth: 1
- }
- }, {
- responsive: true,
- displayModeBar: true,
- modeBarButtonsToRemove: ['lasso2d', 'select2d'],
- displaylogo: false
- });
- }
-
- // 2. RUNNING TASKS PER INSTANCE
- if (snapshots[0].instance_tasks && snapshots[0].instance_tasks.length > 0) {
- const instanceIds = snapshots[0].instance_tasks.map(i => i.instance_id);
- const instanceTraces = instanceIds.map((instanceId, idx) => {
- const color = distinctColors[idx % distinctColors.length];
- return {
- x: relativeTimestamps,
- y: snapshots.map(s => {
- const instData = s.instance_tasks.find(i => i.instance_id === instanceId);
- return instData ? instData.running_tasks : 0;
- }),
- name: `...${instanceId.slice(-4)}`,
- type: 'scatter',
- mode: 'lines+markers',
- line: { color: color, width: 3 },
- marker: { size: 6 }
- };
- });
-
- Plotly.newPlot(`instanceTasksPlot_${testIndex}`, instanceTraces, {
- paper_bgcolor: 'rgba(0,0,0,0)',
- plot_bgcolor: 'rgba(0,0,0,0.2)',
- font: { color: '#ffffff' },
- xaxis: {
- title: 'Time (seconds)',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- yaxis: {
- title: 'Running Tasks',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- hovermode: 'closest',
- showlegend: true,
- legend: {
- bgcolor: 'rgba(0,0,0,0.5)',
- bordercolor: 'rgba(255,255,255,0.2)',
- borderwidth: 1
- }
- }, {
- responsive: true,
- displayModeBar: true,
- modeBarButtonsToRemove: ['lasso2d', 'select2d'],
- displaylogo: false
- });
- }
-
- // 3. CLUSTER-WIDE TASK SUMMARY
- const totalRunning = snapshots.map(s =>
- s.instance_tasks.reduce((sum, inst) => sum + inst.running_tasks, 0)
- );
- const totalPending = snapshots.map(s =>
- s.instance_tasks.reduce((sum, inst) => sum + inst.pending_tasks, 0)
- );
- const totalActive = snapshots.map(s =>
- s.instance_tasks.reduce((sum, inst) => sum + inst.total_active_tasks, 0)
- );
-
- const clusterTraces = [
- {
- x: relativeTimestamps,
- y: totalActive,
- name: 'Total Active',
- type: 'scatter',
- mode: 'lines',
- line: { color: '#00d4ff', width: 3 },
- fill: 'tozeroy',
- fillcolor: 'rgba(0, 212, 255, 0.1)'
- },
- {
- x: relativeTimestamps,
- y: totalRunning,
- name: 'Running',
- type: 'scatter',
- mode: 'lines',
- line: { color: '#00ff88', width: 2, dash: 'dash' }
- },
- {
- x: relativeTimestamps,
- y: totalPending,
- name: 'Pending',
- type: 'scatter',
- mode: 'lines',
- line: { color: '#ffd700', width: 2, dash: 'dot' }
- }
- ];
-
- Plotly.newPlot(`clusterTasksPlot_${testIndex}`, clusterTraces, {
- paper_bgcolor: 'rgba(0,0,0,0)',
- plot_bgcolor: 'rgba(0,0,0,0.2)',
- font: { color: '#ffffff' },
- xaxis: {
- title: 'Time (seconds)',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- yaxis: {
- title: 'Task Count',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- hovermode: 'x unified',
- showlegend: true,
- legend: {
- bgcolor: 'rgba(0,0,0,0.5)',
- bordercolor: 'rgba(255,255,255,0.2)',
- borderwidth: 1
- }
- }, {
- responsive: true,
- displayModeBar: true,
- modeBarButtonsToRemove: ['lasso2d', 'select2d'],
- displaylogo: false
- });
-
- // 4. MEMORY USAGE PER NODE
- const nodeIds = Object.keys(snapshots[0].node_memory || {});
- if (nodeIds.length > 0) {
- const memoryTraces = nodeIds.map((nodeId, idx) => {
- const color = distinctColors[idx % distinctColors.length];
- return {
- x: relativeTimestamps,
- y: snapshots.map(s => {
- const mem = s.node_memory[nodeId];
- return mem ? mem.ram_used_bytes / (1024 ** 3) : 0; // Convert to GB
- }),
- name: `...${nodeId.slice(-4)}`,
- type: 'scatter',
- mode: 'lines',
- line: { color: color, width: 2 }
- };
- });
-
- Plotly.newPlot(`memoryPlot_${testIndex}`, memoryTraces, {
- paper_bgcolor: 'rgba(0,0,0,0)',
- plot_bgcolor: 'rgba(0,0,0,0.2)',
- font: { color: '#ffffff' },
- xaxis: {
- title: 'Time (seconds)',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- yaxis: {
- title: 'RAM Used (GB)',
- gridcolor: 'rgba(255,255,255,0.1)',
- zeroline: false
- },
- hovermode: 'closest',
- showlegend: true,
- legend: {
- bgcolor: 'rgba(0,0,0,0.5)',
- bordercolor: 'rgba(255,255,255,0.2)',
- borderwidth: 1
- }
- }, {
- responsive: true,
- displayModeBar: true,
- modeBarButtonsToRemove: ['lasso2d', 'select2d'],
- displaylogo: false
- });
- }
-
- }
-
- function loadDetailedCharts() {
- if (detailedChartsLoaded) {
- // Already loaded, just scroll to it
- document.getElementById('individualChartsSection')?.scrollIntoView({ behavior: 'smooth' });
- return;
- }
-
- if (allResults.length === 0) {
- console.error('No results available to create detailed charts');
- return;
- }
-
- // Hide the button and show loading
- const button = document.getElementById('showDetailsButton');
- button.textContent = '⏳ Loading detailed charts...';
- button.disabled = true;
-
- // Use setTimeout to allow UI to update before heavy rendering
- setTimeout(() => {
- createDetailedCharts(allResults);
- detailedChartsLoaded = true;
- button.style.display = 'none';
- }, 100);
- }
-
- function destroyAggregateCharts() {
- // Only destroy Chart.js charts (aggregate charts)
- Object.values(charts).forEach(chart => {
- if (chart && typeof chart.destroy === 'function') {
- chart.destroy();
- }
- });
- charts = {};
- }
-
- function destroyDetailedCharts() {
- // Destroy all Plotly charts (detailed individual charts)
- if (allResults.length > 0) {
- allResults.forEach((_, idx) => {
- const plotIds = [
- `requestBreakdownPlot_${idx}`,
- `nodeTasksPlot_${idx}`,
- `instanceTasksPlot_${idx}`,
- `clusterTasksPlot_${idx}`,
- `memoryPlot_${idx}`
- ];
- plotIds.forEach(id => {
- const element = document.getElementById(id);
- if (element) {
- Plotly.purge(element);
- }
- });
- });
- }
- }
-
- function updateCharts(results) {
- // Destroy existing charts
- destroyAggregateCharts();
-
- if (detailedChartsLoaded) {
- destroyDetailedCharts();
- detailedChartsLoaded = false;
- }
-
- // Recreate aggregate charts
- createAggregateCharts(results);
-
- // Show the button again if we have results
- if (results.length > 0) {
- document.getElementById('showDetailsButton').style.display = 'inline-block';
- }
- }
-
- async function loadAndDisplayData() {
- if (!hasCredentials()) {
- document.getElementById('credentialsBanner').style.display = 'block';
- document.getElementById('chartsContainer').innerHTML =
- '<div class="error">AWS credentials required. Please configure them above.</div>';
- document.getElementById('summaryTableBody').innerHTML =
- '<tr><td colspan="5" style="text-align: center; color: #aaa;">AWS credentials required</td></tr>';
- return;
- }
-
- if (!s3Client) {
- s3Client = initializeS3Client();
- }
-
- try {
- document.getElementById('lastUpdate').textContent = 'Fetching latest results...';
-
- const results = await fetchBenchmarkResults();
- allResults = results;
-
- if (allResults.length === 0) {
- document.getElementById('chartsContainer').innerHTML =
- '<div class="error">No benchmark results found in S3 bucket</div>';
- document.getElementById('lastUpdate').textContent = 'No data available';
- return;
- }
-
- updateStats(allResults);
- updateSummaryTable(allResults);
-
- if (Object.keys(charts).length === 0) {
- createAggregateCharts(allResults);
- // Show the button to load detailed charts
- document.getElementById('showDetailsButton').style.display = 'inline-block';
- } else {
- updateCharts(allResults);
- }
- } catch (error) {
- console.error('Error loading data:', error);
- document.getElementById('chartsContainer').innerHTML =
- `<div class="error">Error loading benchmark data: ${error.message}<br><br>
- <button onclick="showCredentialsModal()" style="padding: 0.75rem 1.5rem; background: #00d4ff; color: #0a0a0a; border: none; border-radius: 8px; font-size: 1rem; cursor: pointer;">
- Reconfigure Credentials
- </button></div>`;
- document.getElementById('lastUpdate').textContent = 'Error loading data';
- }
- }
-
- // Initialize
- if (hasCredentials()) {
- s3Client = initializeS3Client();
- loadAndDisplayData();
- // Auto-refresh
- setInterval(loadAndDisplayData, REFRESH_INTERVAL);
- } else {
- document.getElementById('credentialsBanner').style.display = 'block';
- document.getElementById('chartsContainer').innerHTML =
- '<div class="error">AWS credentials required. Please configure them above.</div>';
- document.getElementById('summaryTableBody').innerHTML =
- '<tr><td colspan="5" style="text-align: center; color: #aaa;">AWS credentials required</td></tr>';
- }
- </script>
-</body>
-</html>
diff --git a/.github/configs/README.md b/.github/configs/README.md
deleted file mode 100644
index 4a399c88..00000000
--- a/.github/configs/README.md
+++ /dev/null
@@ -1,186 +0,0 @@
-# EXO Benchmark Configurations
-
-This directory contains configuration files for the EXO staged benchmark system.
-
-## Overview
-
-The staged benchmark system allows you to run complex, multi-stage load tests against EXO clusters. Each stage can have different characteristics:
-
-- **Prompt Length**: Number of tokens in the input prompt
-- **Generation Length**: Maximum tokens to generate in the response
-- **Time Between Requests**: Delay (in seconds) between firing consecutive requests
-- **Iterations**: Number of requests to send in this stage
-
-Requests are **fire-and-forget** - they don't wait for the previous request to complete. This allows you to test overlapping request handling and measure success rates under load.
-
-## Configuration Files
-
-### `bench_simple.yaml`
-A minimal configuration that replicates the behavior of the original `bench.py` script:
-- Single stage with 1 iteration
-- Short prompt (~20 tokens)
-- Generates up to 100 tokens
-
-This is useful for quick smoke tests.
-
-### `bench_config.yaml`
-A comprehensive multi-stage benchmark with:
-1. **Warmup** (10 requests): Light load with short prompts
-2. **Medium Load** (20 requests): Moderate load with medium prompts
-3. **Stress Test** (30 requests): Heavy overlapping requests with long prompts
-4. **Cooldown** (5 requests): Light load to wind down
-
-This tests the cluster's behavior under varying load patterns.
-
-## Configuration Schema
-
-```yaml
-# Hardware configuration - maps runner labels to instance counts
-hardware_plan:
- M3ULTRA_GPU80_512GB: 4
-
-# Environment variables to set on each node (optional)
-environment:
- OVERRIDE_MEMORY_MB: 512
-
-# Timeout for instance and runner readiness (seconds)
-timeout_seconds: 600
-
-# Model instances to run concurrently
-model_ids:
- - "mlx-community/Llama-3.2-1B-Instruct-4bit"
-
-# Benchmark stages
-stages:
- - name: "stage_name" # Human-readable name for this stage
- prompt_length: 100 # Target prompt length in tokens
- generation_length: 200 # Max tokens to generate
- time_between_requests: 2.0 # Seconds between firing requests
- iterations: 10 # Number of requests in this stage
-```
-
-## Running Benchmarks
-
-### Via GitHub Actions
-
-**Automatic (every commit):**
-- The **`bench`** workflow runs automatically on every push
-- Uses `bench_simple.yaml` as the default configuration
-- All settings (hardware plan, timeout, environment variables, models, stages) are defined in the config file
-
-**Manual (on-demand):**
-1. Go to **Actions** → **bench** workflow
-2. Click **Run workflow**
-3. Configure:
- - **Config File**: Path to your YAML config (default: `.github/configs/bench_simple.yaml`)
- - `.github/configs/bench_simple.yaml` for quick tests
- - `.github/configs/bench_config.yaml` for complex multi-stage tests
-
-All other settings (hardware plan, timeout, environment variables, models, stages) are read from the specified config file.
-
-### Via Command Line
-
-```bash
-# Start EXO on localhost:8000
-uv run exo --api-port 8000
-
-# Run simple benchmark (1 stage, 1 iteration)
-python3 .github/scripts/bench.py \
- --api-port 8000 \
- --config .github/configs/bench_simple.yaml \
- --expected-nodes 1 \
- --is-primary true \
- --timeout-seconds 600
-
-# Run complex staged benchmark (4 stages, multiple iterations)
-python3 .github/scripts/bench.py \
- --api-port 8000 \
- --config .github/configs/bench_config.yaml \
- --expected-nodes 1 \
- --is-primary true \
- --timeout-seconds 600
-```
-
-## Output Metrics
-
-For each stage, the benchmark reports:
-
-- **Total Requests**: Number of requests fired
-- **Successful Requests**: Requests that completed successfully
-- **Failed Requests**: Requests that encountered errors
-- **Success Rate**: Percentage of successful requests
-- **Total Tokens**: Sum of all tokens generated across successful requests
-- **Avg Tokens/Request**: Average tokens per successful request
-- **Avg Time/Request**: Average completion time per successful request
-
-A JSON summary is also printed for easy parsing and storage.
-
-## Creating Custom Benchmarks
-
-To create a custom benchmark:
-
-1. Copy an existing config file (e.g., `bench_config.yaml`)
-2. Modify the stages to match your test scenario
-3. Save it in this directory with a descriptive name
-4. Run it using the workflow or command line
-
-### Example: Sustained Load Test
-
-```yaml
-hardware_plan:
- M3ULTRA_GPU80_512GB: 2
-
-environment:
- OVERRIDE_MEMORY_MB: 1024
-
-timeout_seconds: 600
-
-model_ids:
- - "mlx-community/Llama-3.2-1B-Instruct-4bit"
-
-stages:
- - name: "sustained_load"
- prompt_length: 200
- generation_length: 150
- time_between_requests: 0.5 # Very fast - 2 requests/second
- iterations: 100 # Run for ~50 seconds
-```
-
-### Example: Varying Prompt Sizes
-
-```yaml
-hardware_plan:
- M4PRO_GPU16_24GB: 3
-
-timeout_seconds: 900
-
-model_ids:
- - "mlx-community/Llama-3.2-1B-Instruct-4bit"
-
-stages:
- - name: "tiny_prompts"
- prompt_length: 10
- generation_length: 100
- time_between_requests: 1.0
- iterations: 10
-
- - name: "medium_prompts"
- prompt_length: 200
- generation_length: 100
- time_between_requests: 1.0
- iterations: 10
-
- - name: "large_prompts"
- prompt_length: 1000
- generation_length: 100
- time_between_requests: 1.0
- iterations: 10
-```
-
-## Tips
-
-- **Overlapping Requests**: Set `time_between_requests` < expected completion time to test concurrent request handling
-- **Sequential Requests**: Set `time_between_requests` > expected completion time to ensure requests don't overlap
-- **Realistic Load**: Model real usage patterns by varying prompt/generation lengths across stages
-- **Success Rate**: A 100% success rate indicates the cluster handled the load well; lower rates suggest capacity limits
-
diff --git a/.github/configs/bench_config.yaml b/.github/configs/bench_config.yaml
deleted file mode 100644
index 2477a4ff..00000000
--- a/.github/configs/bench_config.yaml
+++ /dev/null
@@ -1,49 +0,0 @@
-# EXO Staged Benchmark Configuration
-# This configuration defines a multi-stage load test for EXO clusters
-
-# Hardware configuration - maps runner labels to instance counts
-hardware_plan:
- M3ULTRA_GPU80_512GB: 4
-
-# Environment variables to set on each node (optional)
-environment:
- OVERRIDE_MEMORY_MB: 512
-
-# Timeout for instance and runner readiness (seconds)
-timeout_seconds: 600
-
-# Multiple instances run concurrently on the cluster
-model_ids:
- - "mlx-community/Qwen3-0.6B-4bit"
- - "mlx-community/Qwen3-0.6B-4bit"
-
-# Stages run sequentially, each with its own characteristics
-stages:
- # Stage 1: Light load with short prompts
- - name: "warmup"
- prompt_length: 50 # Number of tokens in prompt
- generation_length: 100 # Max tokens to generate
- time_between_requests: 5.0 # Seconds between firing requests
- iterations: 10 # Number of requests to send in this stage
-
- # Stage 2: Medium load with medium prompts
- - name: "medium_load"
- prompt_length: 200
- generation_length: 150
- time_between_requests: 3.0
- iterations: 20
-
- # Stage 3: Heavy load with long prompts - requests will overlap
- - name: "stress_test"
- prompt_length: 500
- generation_length: 200
- time_between_requests: 1.0 # Fast firing - will definitely overlap
- iterations: 30
-
- # Stage 4: Cool down with simple prompts
- - name: "cooldown"
- prompt_length: 50
- generation_length: 50
- time_between_requests: 10.0
- iterations: 5
-
diff --git a/.github/configs/bench_simple.yaml b/.github/configs/bench_simple.yaml
deleted file mode 100644
index 9a76b6db..00000000
--- a/.github/configs/bench_simple.yaml
+++ /dev/null
@@ -1,125 +0,0 @@
-# Simple single-shot benchmark
-# Tests 2 instances concurrently on 2 nodes
-
-# Hardware configuration - maps runner labels to instance counts
-hardware_plan:
- puffin4: 1
- puffin8: 1
-
-# Environment variables to set on each node
-environment:
- PLACEHOLDER: "placeholder"
- # OVERRIDE_MEMORY_MB: 50000
- MLX_METAL_FAST_SYNCH: 1
-
-# Timeout for instance and runner readiness (seconds)
-timeout_seconds: 1800
-
-# Model instances to run concurrently
-model_ids:
- # - "mlx-community/DeepSeek-V3.1-8bit"
- # - "mlx-community/Kimi-K2-Instruct-4bit"
- - "mlx-community/Kimi-K2-Thinking"
- # - "mlx-community/Qwen3-235B-A22B-4bit"
- # - "mlx-community/Llama-3.3-70B-Instruct-4bit"
- # - "mlx-community/Llama-3.3-70B-Instruct-8bit"
- # - "mlx-community/Llama-3.2-1B-Instruct-4bit"
-
-# Sharding strategy: "Pipeline" or "Tensor"
-sharding: "Tensor"
-
-# Instance type: "MlxRing" or "MlxIbv"
-instance_meta: "MlxIbv"
-
-# If true, run requests sequentially (no overlap); if false, fire-and-forget (default: false)
-no_overlap: true
-
-# Benchmark stages
-# pp: 64, 256, 1024, 2048, 4096, 8192, 16384
-# g: 64, 512
-stages:
- # - name: "simple"
- # prompt_length: 512
- # generation_length: 10
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp64_g64"
- # prompt_length: 64
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp64_g64"
- # prompt_length: 64
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp64_g512"
- # prompt_length: 64
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp256_g64"
- # prompt_length: 256
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- - name: "pp256_g64"
- prompt_length: 256
- generation_length: 64
- time_between_requests: 2.0
- iterations: 5
- # - name: "pp256_g512"
- # prompt_length: 256
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp1024_g64"
- # prompt_length: 1024
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp1024_g512"
- # prompt_length: 1024
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp2048_g64"
- # prompt_length: 2048
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp2048_g512"
- # prompt_length: 2048
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp4096_g64"
- # prompt_length: 4096
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 4
- # - name: "pp4096_g512"
- # prompt_length: 4096
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp8192_g64"
- # prompt_length: 8192
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp8192_g512"
- # prompt_length: 8192
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 5
- # - name: "pp16384_g64"
- # prompt_length: 16384
- # generation_length: 64
- # time_between_requests: 2.0
- # iterations: 10
- # - name: "pp16384_g512"
- # prompt_length: 16384
- # generation_length: 512
- # time_between_requests: 2.0
- # iterations: 10
diff --git a/.github/scripts/bench.py b/.github/scripts/bench.py
deleted file mode 100644
index 0ba73d58..00000000
--- a/.github/scripts/bench.py
+++ /dev/null
@@ -1,1399 +0,0 @@
-#!/usr/bin/env python3
-
-# type: ignore
-"""
-Unified benchmark script for EXO.
-Runs single or multi-stage benchmarks with configurable load patterns.
-Requests are fire-and-forget, allowing overlapping execution.
-
-Simple benchmark (1 iteration): --config .github/configs/bench_simple.yaml
-Complex benchmark (multiple stages): --config .github/configs/bench_config.yaml
-"""
-
-# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownVariableType=false
-from __future__ import annotations
-
-import argparse
-import asyncio
-import contextlib
-import json
-import sys
-import time
-import urllib.error
-import urllib.request
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Mapping
-
-import yaml
-
-
-def _format_http_error(error: Exception) -> str:
- """Format HTTP error with full response details for debugging."""
- if isinstance(error, urllib.error.HTTPError):
- try:
- body = error.read().decode("utf-8", errors="replace")
- except Exception:
- body = "<unable to read body>"
-
- headers_str = (
- "\n".join(f" {k}: {v}" for k, v in error.headers.items())
- if error.headers
- else "<no headers>"
- )
-
- return (
- f"HTTP {error.code} {error.reason}\n"
- f"URL: {error.url}\n"
- f"Headers:\n{headers_str}\n"
- f"Body: {body}"
- )
- elif isinstance(error, urllib.error.URLError):
- return f"URLError: {error.reason}"
- else:
- return str(error)
-
-
-def _http_request(
- url: str, *, method: str = "GET", data: Mapping[str, Any] | None = None
-) -> dict[str, Any]:
- headers = {"Content-Type": "application/json"}
- payload: bytes | None = None
- if data is not None:
- payload = json.dumps(data).encode("utf-8")
- req = urllib.request.Request(url, data=payload, headers=headers, method=method)
- try:
- with urllib.request.urlopen(req, timeout=300) as resp: # nosec - runner-local API
- body = resp.read().decode("utf-8")
- try:
- return json.loads(body)
- except json.JSONDecodeError:
- return {"raw": body}
- except Exception as e:
- error_details = _format_http_error(e)
- print(f"HTTP request failed:\n{error_details}")
- raise
-
-
-async def _http_request_async(
- url: str, *, method: str = "GET", data: Mapping[str, Any] | None = None
-) -> dict[str, Any]:
- """Async version that runs in executor to not block event loop."""
- loop = asyncio.get_event_loop()
- return await loop.run_in_executor(
- None, lambda: _http_request(url, method=method, data=data)
- )
-
-
-async def _http_stream_async(
- url: str, *, method: str = "POST", data: Mapping[str, Any], timeout: int = 300
-) -> list[tuple[str, float]]:
- """Async streaming request. Returns list of (line, timestamp) tuples."""
-
- def _stream() -> list[tuple[str, float]]:
- headers = {"Content-Type": "application/json"}
- payload = json.dumps(data).encode("utf-8")
- req = urllib.request.Request(url, data=payload, headers=headers, method=method)
- lines: list[tuple[str, float]] = []
- try:
- with urllib.request.urlopen(req, timeout=timeout) as resp: # nosec - runner-local API
- for raw_line in resp:
- timestamp = time.monotonic()
- line = raw_line.decode("utf-8", errors="replace").rstrip("\n\r")
- if line:
- lines.append((line, timestamp))
- return lines
- except Exception as e:
- error_details = _format_http_error(e)
- print(f"HTTP request failed:\n{error_details}")
- raise
-
- loop = asyncio.get_event_loop()
- return await loop.run_in_executor(None, _stream)
-
-
-def fetch_state(api_base: str) -> dict[str, Any]:
- return _http_request(f"{api_base}/state")
-
-
-def unwrap_tagged_union(obj: Any) -> tuple[str | None, Any]:
- """Extract tag and payload from tagged union format {Tag: {fields...}}.
-
- Returns (tag_name, payload) if the object is a tagged union, otherwise (None, obj).
- """
- if not isinstance(obj, dict):
- return None, obj
-
- keys = list(obj.keys())
- if len(keys) == 1 and isinstance(keys[0], str):
- tag = keys[0]
- payload = obj[tag]
- return tag, payload
-
- return None, obj
-
-
-def collect_metrics_snapshot(state: Mapping[str, Any]) -> MetricsSnapshot:
- """Collect current metrics snapshot from state."""
- timestamp = time.time()
-
- # Collect memory for each node
- node_memory: dict[str, MemorySnapshot] = {}
- node_profiles: Mapping[str, Any] = state.get("nodeProfiles", {})
-
- for node_id, profile in node_profiles.items():
- if not isinstance(profile, dict):
- continue
-
- memory = profile.get("memory", {})
- if not isinstance(memory, dict):
- continue
-
- # Parse memory values - they're objects with 'inBytes' field
- def get_bytes(mem_obj: Any) -> int:
- if isinstance(mem_obj, dict):
- return int(mem_obj.get("inBytes", 0))
- return 0
-
- ram_total = get_bytes(memory.get("ramTotal"))
- ram_available = get_bytes(memory.get("ramAvailable"))
- swap_total = get_bytes(memory.get("swapTotal"))
- swap_available = get_bytes(memory.get("swapAvailable"))
-
- node_memory[node_id] = MemorySnapshot(
- ram_total_bytes=ram_total,
- ram_available_bytes=ram_available,
- ram_used_bytes=max(ram_total - ram_available, 0),
- swap_total_bytes=swap_total,
- swap_available_bytes=swap_available,
- swap_used_bytes=max(swap_total - swap_available, 0),
- )
-
- # Collect task counts per instance and per node
- instance_tasks: list[InstanceTaskSnapshot] = []
- instances: Mapping[str, Any] = state.get("instances", {})
- tasks: Mapping[str, Any] = state.get("tasks", {})
- print(f"[DEBUG] Num tasks: {len(tasks)}. Num instances: {len(instances)}.")
-
- # Map instance_id -> node_ids (instances can span multiple nodes)
- instance_to_nodes: dict[str, set[str]] = {}
- for instance_id, instance_wrapped in instances.items():
- # Unwrap tagged Instance union (MlxRingInstance or MlxIbvInstance)
- _instance_tag, instance_data = unwrap_tagged_union(instance_wrapped)
- if not isinstance(instance_data, dict):
- continue
-
- shard_assignments = instance_data.get("shardAssignments", {})
- if not isinstance(shard_assignments, dict):
- continue
-
- # Get all nodes that this instance uses
- node_to_runner = shard_assignments.get("nodeToRunner", {})
- if isinstance(node_to_runner, dict):
- instance_to_nodes[instance_id] = set(node_to_runner.keys())
-
- # Count tasks per instance (only Pending and Running exist in state; completed tasks are deleted)
- instance_task_counts: dict[str, dict[str, int]] = {}
- for instance_id in instances:
- instance_task_counts[instance_id] = {
- "Pending": 0,
- "Running": 0,
- }
-
- # Iterate through tasks and count by instance and status
- tasks_matched = 0
- tasks_skipped = 0
-
- for _task_id, task_wrapper in tasks.items():
- if not isinstance(task_wrapper, dict):
- print(f"[DEBUG] Task wrapper is not a dict: {task_wrapper}")
- tasks_skipped += 1
- continue
-
- # Extract actual task from wrapper (e.g., {"ChatCompletion": {...}})
- if len(task_wrapper) != 1:
- print(
- f"[DEBUG] Task wrapper has unexpected number of keys: {len(task_wrapper)}"
- )
- tasks_skipped += 1
- continue
-
- _task_type, task_data = next(iter(task_wrapper.items()))
-
- if not isinstance(task_data, dict):
- print(f"[DEBUG] Task data is not a dict: {task_data}")
- tasks_skipped += 1
- continue
-
- instance_id = task_data.get("instanceId")
- task_status = task_data.get("taskStatus")
-
- if not instance_id or instance_id not in instance_task_counts:
- tasks_skipped += 1
- continue
-
- if task_status not in ["Pending", "Running"]:
- tasks_skipped += 1
- continue
-
- # Count this task
- instance_task_counts[instance_id][task_status] += 1
- tasks_matched += 1
-
- if tasks_skipped > 0:
- print(
- f"[DEBUG] Task matching: {tasks_matched} matched, {tasks_skipped} skipped (from {len(tasks)} total)"
- )
-
- # Build snapshots for each instance (assign to primary node - first in sorted order)
- for instance_id, counts in instance_task_counts.items():
- pending = counts["Pending"]
- running = counts["Running"]
- total_active = pending + running
-
- node_ids = instance_to_nodes.get(instance_id, set())
- primary_node = sorted(node_ids)[0] if node_ids else "unknown"
-
- instance_tasks.append(
- InstanceTaskSnapshot(
- instance_id=instance_id,
- node_id=primary_node,
- pending_tasks=pending,
- running_tasks=running,
- total_active_tasks=total_active,
- )
- )
-
- # Aggregate tasks per node
- node_task_counts: dict[str, dict[str, int]] = {}
- node_instance_counts: dict[str, int] = {}
-
- for instance_snapshot in instance_tasks:
- node_id = instance_snapshot.node_id
-
- if node_id not in node_task_counts:
- node_task_counts[node_id] = {
- "Pending": 0,
- "Running": 0,
- }
- node_instance_counts[node_id] = 0
-
- node_task_counts[node_id]["Pending"] += instance_snapshot.pending_tasks
- node_task_counts[node_id]["Running"] += instance_snapshot.running_tasks
- node_instance_counts[node_id] += 1
-
- # Build node snapshots
- node_tasks: list[NodeTaskSnapshot] = []
- for node_id, counts in node_task_counts.items():
- pending = counts["Pending"]
- running = counts["Running"]
- total_active = pending + running
-
- node_tasks.append(
- NodeTaskSnapshot(
- node_id=node_id,
- pending_tasks=pending,
- running_tasks=running,
- total_active_tasks=total_active,
- instance_count=node_instance_counts.get(node_id, 0),
- )
- )
-
- return MetricsSnapshot(
- timestamp=timestamp,
- node_memory=node_memory,
- instance_tasks=instance_tasks,
- node_tasks=node_tasks,
- )
-
-
-def get_topology_node_count(state: Mapping[str, Any]) -> int:
- """Get the number of nodes in the topology."""
- topology = state.get("topology", {})
- nodes = topology.get("nodes", [])
- return len(nodes)
-
-
-def count_instances_by_model(state: Mapping[str, Any], model_id: str) -> int:
- """Count how many instances exist for a given model_id."""
- instances: Mapping[str, Any] = state.get("instances", {})
- count = 0
- for instance_wrapped in instances.values():
- # Unwrap tagged Instance union
- _instance_tag, instance_data = unwrap_tagged_union(instance_wrapped)
- if not isinstance(instance_data, dict):
- continue
-
- shard = instance_data.get("shardAssignments", {})
- if isinstance(shard, dict) and shard.get("modelId") == model_id:
- count += 1
- return count
-
-
-def get_all_instance_ids_for_model(
- state: Mapping[str, Any], model_id: str
-) -> list[str]:
- """Get all instance IDs for a given model_id."""
- instances: Mapping[str, Any] = state.get("instances", {})
- instance_ids = []
- for instance_id, instance_wrapped in instances.items():
- # Unwrap tagged Instance union
- _instance_tag, instance_data = unwrap_tagged_union(instance_wrapped)
- if not isinstance(instance_data, dict):
- continue
-
- shard = instance_data.get("shardAssignments", {})
- if isinstance(shard, dict) and shard.get("modelId") == model_id:
- instance_ids.append(instance_id)
- return instance_ids
-
-
-def count_ready_instances_by_model(state: Mapping[str, Any], model_id: str) -> int:
- """Count how many instances for a model have all their runners ready."""
- instances: Mapping[str, Any] = state.get("instances", {})
- ready_count = 0
-
- for instance_id, instance_wrapped in instances.items():
- # Unwrap tagged Instance union
- _instance_tag, instance_data = unwrap_tagged_union(instance_wrapped)
- if not isinstance(instance_data, dict):
- continue
-
- shard = instance_data.get("shardAssignments", {})
- if not isinstance(shard, dict) or shard.get("modelId") != model_id:
- continue
-
- # Check if all runners for this instance are ready
- runner_ids = get_runner_ids_for_instance(state, instance_id)
- if len(runner_ids) == 0:
- continue
-
- # Fixed runner status names: RunnerReady and RunnerRunning (not LoadedRunnerStatus/RunningRunnerStatus)
- all_ready = all(
- get_runner_status_kind(state, rid) in {"RunnerReady", "RunnerRunning"}
- for rid in runner_ids
- )
-
- if all_ready:
- ready_count += 1
-
- return ready_count
-
-
-def get_runner_ids_for_instance(
- state: Mapping[str, Any], instance_id: str
-) -> list[str]:
- instances: Mapping[str, Any] = state.get("instances", {})
- instance_wrapped = instances.get(instance_id, {})
-
- # Unwrap tagged Instance union
- _instance_tag, instance_data = unwrap_tagged_union(instance_wrapped)
- if not isinstance(instance_data, dict):
- return []
-
- shard_assignments = instance_data.get("shardAssignments", {})
- if not isinstance(shard_assignments, dict):
- return []
-
- r2s = shard_assignments.get("runnerToShard", {})
- if isinstance(r2s, dict):
- return list(r2s.keys())
- return []
-
-
-def get_runner_status_kind(state: Mapping[str, Any], runner_id: str) -> str | None:
- runners: Mapping[str, Any] = state.get("runners", {})
- status_obj = runners.get(runner_id)
- if not isinstance(status_obj, dict):
- return None
- if len(status_obj) == 1:
- return next(iter(status_obj.keys()))
- return None
-
-
-async def wait_for_topology_ready(
- api_base: str, expected_nodes: int, timeout_s: int
-) -> None:
- """Wait for all expected nodes to appear in the topology."""
- print(
- f"Waiting for {expected_nodes} node(s) to appear in topology (timeout: {timeout_s}s)..."
- )
- start = time.monotonic()
- while True:
- state = fetch_state(api_base)
- node_count = get_topology_node_count(state)
- elapsed = time.monotonic() - start
- print(
- f" Topology has {node_count}/{expected_nodes} nodes (elapsed: {elapsed:.1f}s)"
- )
-
- if node_count >= expected_nodes:
- print(f"All {expected_nodes} node(s) are in topology!")
- return
-
- if elapsed > timeout_s:
- raise TimeoutError(
- f"Timed out waiting for topology. Expected {expected_nodes} nodes, got {node_count}"
- )
- await asyncio.sleep(2)
-
-
-async def wait_for_instances_ready(
- api_base: str, model_id: str, expected_count: int, timeout_s: int
-) -> list[str]:
- """Wait for a specific count of instances for a model to be fully ready."""
- print(
- f"Waiting for {expected_count} instance(s) of {model_id} to be ready (timeout: {timeout_s}s)..."
- )
- start = time.monotonic()
- while True:
- state = fetch_state(api_base)
-
- total_count = count_instances_by_model(state, model_id)
- ready_count = count_ready_instances_by_model(state, model_id)
- elapsed = time.monotonic() - start
-
- print(
- f" Model {model_id}: {ready_count}/{expected_count} ready ({total_count} total) (elapsed: {elapsed:.1f}s)"
- )
-
- if ready_count >= expected_count:
- instance_ids = get_all_instance_ids_for_model(state, model_id)
- print(
- f"All {expected_count} instance(s) ready! Instance IDs: {instance_ids}"
- )
- return instance_ids
-
- if elapsed > timeout_s:
- raise TimeoutError(
- f"Timed out waiting for instances. Expected {expected_count} ready instances of {model_id}, "
- f"got {ready_count} ready out of {total_count} total"
- )
- await asyncio.sleep(2)
-
-
-async def wait_for_all_instances_deleted(api_base: str, model_id: str) -> None:
- """Wait for all instances of a model to be deleted."""
- print(f"Waiting for all instances of {model_id} to be deleted...")
- start = time.monotonic()
- while True:
- state = fetch_state(api_base)
- count = count_instances_by_model(state, model_id)
- if count == 0:
- elapsed = time.monotonic() - start
- print(f"All instances of {model_id} deleted after {elapsed:.1f}s")
- return
- await asyncio.sleep(2)
-
-
-async def wait_for_tasks_drained(api_base: str, timeout_s: int = 600) -> None:
- """Wait for all tasks in the cluster to be drained (completed or failed).
-
- Tasks are deleted from state when complete, so we wait until there are no
- pending or running tasks remaining.
- """
- print(f"\n{'=' * 80}")
- print("⏳ WAITING FOR ALL TASKS TO DRAIN")
- print(f"{'=' * 80}")
- start = time.monotonic()
-
- while True:
- state = fetch_state(api_base)
- snapshot = collect_metrics_snapshot(state)
-
- # Count total active tasks across all nodes
- total_pending = sum(node.pending_tasks for node in snapshot.node_tasks)
- total_running = sum(node.running_tasks for node in snapshot.node_tasks)
- total_active = total_pending + total_running
-
- elapsed = time.monotonic() - start
-
- if total_active == 0:
- print(f"✅ All tasks drained after {elapsed:.1f}s")
- return
-
- print(
- f" [{elapsed:.1f}s] Still draining: {total_active} active tasks ({total_pending} pending, {total_running} running)"
- )
-
- # Print per-node breakdown if there are active tasks
- if snapshot.node_tasks:
- for node_snapshot in snapshot.node_tasks:
- if node_snapshot.total_active_tasks > 0:
- node_short = node_snapshot.node_id[-4:]
- print(
- f" Node ...{node_short}: {node_snapshot.running_tasks} running, {node_snapshot.pending_tasks} pending"
- )
-
- if elapsed > timeout_s:
- print(
- f"⚠️ WARNING: Timed out waiting for tasks to drain after {timeout_s}s"
- )
- print(
- f" Remaining: {total_active} tasks ({total_pending} pending, {total_running} running)"
- )
- return
-
- await asyncio.sleep(2)
-
-
-def generate_prompt(length: int) -> str:
- """Generate a prompt of approximately the specified token length."""
- # Rough approximation: 1 token ≈ 4 characters
- # Use a repeating pattern that's easy to generate
- base_text = "The quick brown fox jumps over the lazy dog. "
- target_chars = length * 4
- repetitions = (target_chars // len(base_text)) + 1
- return (base_text * repetitions)[:target_chars]
-
-
-@dataclass(frozen=True)
-class StageConfig:
- name: str
- prompt_length: int
- generation_length: int
- time_between_requests: float
- iterations: int
-
-
-@dataclass
-class RequestResult:
- request_id: int
- success: bool
- tokens: int
- elapsed_s: float
- started_at: float
- completed_at: float
- time_to_first_token_s: float | None = None
- decode_tps: float | None = None
- error: str | None = None
-
-
-@dataclass
-class StageResult:
- name: str
- total_requests: int
- successful_requests: int
- failed_requests: int
- success_rate: float
- total_tokens: int
- total_time: float
- avg_tokens_per_request: float
- avg_time_per_request: float
- avg_time_to_first_token: float | None
- std_time_to_first_token: float | None
- avg_decode_tps: float | None
- avg_ms_per_token: float | None
- std_ms_per_token: float | None
- request_results: list[RequestResult]
- stage_started_at: float
- stage_completed_at: float
-
-
-@dataclass(frozen=True)
-class MemorySnapshot:
- """Memory snapshot for a node at a point in time."""
-
- ram_total_bytes: int
- ram_available_bytes: int
- ram_used_bytes: int
- swap_total_bytes: int
- swap_available_bytes: int
- swap_used_bytes: int
-
-
-@dataclass(frozen=True)
-class InstanceTaskSnapshot:
- """Task counts for an instance at a point in time.
-
- Note: Tasks are deleted from state when complete, so we only track active tasks.
- total_active_tasks = pending + running.
- """
-
- instance_id: str
- node_id: str
- pending_tasks: int
- running_tasks: int
- total_active_tasks: int
-
-
-@dataclass(frozen=True)
-class NodeTaskSnapshot:
- """Task counts for a node at a point in time.
-
- Note: Tasks are deleted from state when complete, so we only track active tasks.
- total_active_tasks = pending + running across all instances on this node.
- """
-
- node_id: str
- pending_tasks: int
- running_tasks: int
- total_active_tasks: int
- instance_count: int
-
-
-@dataclass(frozen=True)
-class MetricsSnapshot:
- """System metrics snapshot at a point in time."""
-
- timestamp: float
- node_memory: dict[str, MemorySnapshot]
- instance_tasks: list[InstanceTaskSnapshot]
- node_tasks: list[NodeTaskSnapshot]
-
-
-async def run_single_request(
- api_base: str,
- model_id: str,
- prompt: str,
- max_tokens: int,
- request_id: int,
- timeout: int = 300,
-) -> RequestResult:
- """Run a single chat completion request and return its result."""
- started_at = time.time()
- start = time.monotonic()
- try:
- lines = await _http_stream_async(
- f"{api_base}/v1/chat/completions",
- method="POST",
- data={
- "model": model_id,
- "messages": [{"role": "user", "content": prompt}],
- "max_tokens": max_tokens,
- "temperature": 0.7,
- },
- timeout=timeout,
- )
-
- tokens = 0
- got_done = False
- first_token_time: float | None = None
- last_token_time: float | None = None
-
- for line, timestamp in lines:
- if not line.startswith("data:"):
- continue
- payload = line[len("data:") :].strip()
- if payload == "[DONE]":
- got_done = True
- break
- try:
- obj = json.loads(payload)
- content = obj.get("choices", [{}])[0].get("delta", {}).get("content")
- if content:
- if first_token_time is None:
- first_token_time = timestamp
- last_token_time = timestamp
- tokens += 1
- except json.JSONDecodeError:
- continue
-
- elapsed = time.monotonic() - start
- completed_at = time.time()
-
- # Calculate TTFT and decode TPS
- time_to_first_token: float | None = None
- decode_tps: float | None = None
-
- if first_token_time is not None:
- time_to_first_token = first_token_time - start
-
- # Decode TPS: tokens per second after first token
- if last_token_time is not None and tokens > 1:
- decode_time = last_token_time - first_token_time
- if decode_time > 0:
- decode_tps = (tokens - 1) / decode_time
-
- # Request is only successful if we got at least one token AND a [DONE] marker
- if tokens == 0:
- print(
- f" Request #{request_id}: FAILED - no tokens generated in {elapsed:.2f}s"
- )
- return RequestResult(
- request_id=request_id,
- success=False,
- tokens=0,
- elapsed_s=elapsed,
- started_at=started_at,
- completed_at=completed_at,
- time_to_first_token_s=time_to_first_token,
- decode_tps=decode_tps,
- error="No tokens generated",
- )
-
- if not got_done:
- print(
- f" Request #{request_id}: FAILED - incomplete response (no [DONE]) after {elapsed:.2f}s"
- )
- return RequestResult(
- request_id=request_id,
- success=False,
- tokens=tokens,
- elapsed_s=elapsed,
- started_at=started_at,
- completed_at=completed_at,
- time_to_first_token_s=time_to_first_token,
- decode_tps=decode_tps,
- error="Incomplete response (no [DONE] marker)",
- )
-
- ttft_str = (
- f"{time_to_first_token:.3f}s" if time_to_first_token is not None else "N/A"
- )
- tps_str = f"{decode_tps:.1f} t/s" if decode_tps is not None else "N/A"
- print(
- f" Request #{request_id}: SUCCESS - {tokens} tokens in {elapsed:.2f}s (TTFT: {ttft_str}, Decode: {tps_str})"
- )
- return RequestResult(
- request_id=request_id,
- success=True,
- tokens=tokens,
- elapsed_s=elapsed,
- started_at=started_at,
- completed_at=completed_at,
- time_to_first_token_s=time_to_first_token,
- decode_tps=decode_tps,
- )
-
- except Exception as e:
- elapsed = time.monotonic() - start
- completed_at = time.time()
- error_details = _format_http_error(e)
- print(f" Request #{request_id}: FAILED - {error_details}")
- return RequestResult(
- request_id=request_id,
- success=False,
- tokens=0,
- elapsed_s=elapsed,
- started_at=started_at,
- completed_at=completed_at,
- time_to_first_token_s=None,
- decode_tps=None,
- error=error_details,
- )
-
-
-async def monitor_metrics(
- api_base: str,
- metrics_snapshots: list[MetricsSnapshot],
- stop_event: asyncio.Event,
- interval_seconds: float = 5.0,
-) -> None:
- """Background task that collects metrics snapshots every interval_seconds."""
- print(f"\n{'=' * 80}")
- print(f"🔍 METRICS MONITORING STARTED (polling every {interval_seconds}s)")
- print(f"{'=' * 80}\n")
-
- snapshot_count = 0
- while not stop_event.is_set():
- try:
- snapshot_count += 1
- state = fetch_state(api_base)
- snapshot = collect_metrics_snapshot(state)
- metrics_snapshots.append(snapshot)
-
- # Print detailed summary
- node_count = len(snapshot.node_memory)
- instance_count = len(snapshot.instance_tasks)
-
- # Aggregate task counts from node level (only active tasks in state)
- total_pending = sum(node.pending_tasks for node in snapshot.node_tasks)
- total_running = sum(node.running_tasks for node in snapshot.node_tasks)
- total_active = sum(node.total_active_tasks for node in snapshot.node_tasks)
-
- # Print detailed breakdown
- print(
- f"\n[METRICS #{snapshot_count}] {node_count} nodes, {instance_count} instances | Active Tasks: {total_active} ({total_pending} pending, {total_running} running)"
- )
-
- # Print per-node breakdown (only if there are nodes)
- if snapshot.node_tasks:
- for node_snapshot in snapshot.node_tasks:
- node_short = node_snapshot.node_id[-4:]
- print(
- f" Node ...{node_short}: {node_snapshot.total_active_tasks} active ({node_snapshot.pending_tasks} pending, {node_snapshot.running_tasks} running) across {node_snapshot.instance_count} instances"
- )
-
- except Exception as e:
- print(f"[METRICS] Error collecting snapshot: {e}")
- import traceback
-
- traceback.print_exc()
-
- # Wait for interval or until stopped
- with contextlib.suppress(asyncio.TimeoutError):
- await asyncio.wait_for(stop_event.wait(), timeout=interval_seconds)
-
-
-async def run_stage(
- api_base: str,
- model_id: str,
- stage: StageConfig,
- no_overlap: bool = False,
-) -> StageResult:
- """Run a single benchmark stage with fire-and-forget requests (or sequential if no_overlap=True)."""
- print("=" * 80)
- print(f"STAGE: {stage.name}")
- print("=" * 80)
- print(f" Prompt Length: {stage.prompt_length} tokens")
- print(f" Generation Length: {stage.generation_length} tokens")
- print(f" Time Between Reqs: {stage.time_between_requests}s")
- print(f" Iterations: {stage.iterations}")
- print(f" No Overlap: {no_overlap}")
- print("=" * 80)
-
- stage_started_at = time.time()
- prompt = generate_prompt(stage.prompt_length)
- results: list[RequestResult] = []
-
- if no_overlap:
- # Sequential execution: wait for each request to complete before starting next
- print("\nRunning requests sequentially (no overlap)...")
- for i in range(stage.iterations):
- result = await run_single_request(
- api_base, model_id, prompt, stage.generation_length, i + 1
- )
- results.append(result)
-
- # Wait before starting next request (except after last one)
- if i < stage.iterations - 1:
- await asyncio.sleep(stage.time_between_requests)
- else:
- # Concurrent execution: fire-and-forget with delays between starts
- print("\nRunning requests concurrently (with overlap)...")
- tasks: list[asyncio.Task[RequestResult]] = []
-
- # Fire off requests with delays between them
- for i in range(stage.iterations):
- task = asyncio.create_task(
- run_single_request(
- api_base, model_id, prompt, stage.generation_length, i + 1
- )
- )
- tasks.append(task)
-
- # Wait before firing next request (except after last one)
- if i < stage.iterations - 1:
- await asyncio.sleep(stage.time_between_requests)
-
- # Wait for all requests to complete
- print(f"\nWaiting for all {len(tasks)} HTTP requests to complete...")
- results = list(await asyncio.gather(*tasks))
-
- # Wait for all tasks in the cluster to be drained
- print("\nHTTP requests completed. Now waiting for cluster tasks to drain...")
- await wait_for_tasks_drained(api_base, timeout_s=600)
-
- stage_completed_at = time.time()
-
- # Compute statistics
- successful = sum(1 for r in results if r.success)
- failed = len(results) - successful
- success_rate = successful / len(results) if results else 0.0
- total_tokens = sum(r.tokens for r in results)
- total_time = sum(r.elapsed_s for r in results)
- avg_tokens = total_tokens / successful if successful > 0 else 0.0
- avg_time = total_time / successful if successful > 0 else 0.0
-
- # Calculate average TTFT and decode TPS for successful requests only
- successful_results = [r for r in results if r.success]
-
- # Skip first iteration if there are more than 1 iterations (warmup)
- results_for_stats = (
- successful_results[1:] if len(successful_results) > 1 else successful_results
- )
-
- # TTFT statistics
- ttft_values = [
- r.time_to_first_token_s
- for r in results_for_stats
- if r.time_to_first_token_s is not None
- ]
- avg_ttft = sum(ttft_values) / len(ttft_values) if ttft_values else None
-
- if avg_ttft is not None and len(ttft_values) > 1:
- variance_ttft = sum((x - avg_ttft) ** 2 for x in ttft_values) / len(ttft_values)
- std_ttft = variance_ttft**0.5
- else:
- std_ttft = None
-
- # Decode TPS and ms per token statistics
- decode_tps_values = [
- r.decode_tps for r in results_for_stats if r.decode_tps is not None
- ]
- avg_decode_tps = (
- sum(decode_tps_values) / len(decode_tps_values) if decode_tps_values else None
- )
-
- # Convert to ms per token
- ms_per_token_values = (
- [1000.0 / tps for tps in decode_tps_values] if decode_tps_values else []
- )
- avg_ms_per_token = (
- sum(ms_per_token_values) / len(ms_per_token_values)
- if ms_per_token_values
- else None
- )
-
- if avg_ms_per_token is not None and len(ms_per_token_values) > 1:
- variance_ms_per_token = sum(
- (x - avg_ms_per_token) ** 2 for x in ms_per_token_values
- ) / len(ms_per_token_values)
- std_ms_per_token = variance_ms_per_token**0.5
- else:
- std_ms_per_token = None
-
- return StageResult(
- name=stage.name,
- total_requests=len(results),
- successful_requests=successful,
- failed_requests=failed,
- success_rate=success_rate,
- total_tokens=total_tokens,
- total_time=total_time,
- avg_tokens_per_request=avg_tokens,
- avg_time_per_request=avg_time,
- avg_time_to_first_token=avg_ttft,
- std_time_to_first_token=std_ttft,
- avg_decode_tps=avg_decode_tps,
- avg_ms_per_token=avg_ms_per_token,
- std_ms_per_token=std_ms_per_token,
- request_results=list(results),
- stage_started_at=stage_started_at,
- stage_completed_at=stage_completed_at,
- )
-
-
-async def run_benchmark(
- api_base: str,
- config_path: Path,
- expected_nodes: int,
- is_primary: bool,
- timeout_seconds: int,
- results_output_path: Path | None = None,
- git_commit: str | None = None,
- hardware_labels: list[str] | None = None,
-) -> int:
- """Run the full staged benchmark."""
- benchmark_started_at = time.time()
-
- # Load configuration
- with open(config_path) as f:
- config = yaml.safe_load(f)
-
- # Support both model_id (legacy) and model_ids (new)
- if "model_ids" in config:
- model_ids = config["model_ids"]
- elif "model_id" in config:
- model_ids = [config["model_id"]]
- else:
- raise ValueError("Config must contain either 'model_id' or 'model_ids'")
-
- # Get sharding and instance_meta (optional, defaults to None if not specified)
- sharding: str | None = config.get("sharding")
- instance_meta: str | None = config.get("instance_meta")
-
- # Get no_overlap flag (optional, defaults to False)
- no_overlap: bool = config.get("no_overlap", False)
-
- stages = [StageConfig(**s) for s in config["stages"]]
-
- print("=" * 80)
- print("EXO BENCHMARK")
- print("=" * 80)
- print(f"Configuration File: {config_path}")
- print(f"Model IDs: {model_ids}")
- print(f"Instance Count: {len(model_ids)}")
- print(
- f"Sharding: {sharding if sharding else 'not specified (defaults to Pipeline)'}"
- )
- print(
- f"Instance Type: {instance_meta if instance_meta else 'not specified (defaults to MlxRing)'}"
- )
- print(f"No Overlap: {no_overlap}")
- print(f"Stages: {len(stages)}")
- print(f"Expected Nodes: {expected_nodes}")
- print(f"Is Primary: {is_primary}")
- print("=" * 80)
-
- try:
- # Wait for all nodes to join the topology first
- await wait_for_topology_ready(
- api_base, expected_nodes, timeout_s=timeout_seconds
- )
-
- # Add 30 second delay to allow topology to stabilize before creating instances
- print(
- "\nWaiting 30 seconds for topology to stabilize before creating instances..."
- )
- await asyncio.sleep(30)
- print("Proceeding with instance creation\n")
-
- # Count how many instances we need for each unique model_id
- from collections import Counter
-
- model_counts = Counter(model_ids)
-
- print("\nTarget instance counts by model:")
- for model_id, count in model_counts.items():
- print(f" {model_id}: {count} instance(s)")
- print()
-
- # Track all instance IDs (collected at the end)
- all_instance_ids: list[str] = []
-
- if is_primary:
- # Primary: create instances one at a time, waiting for count to increase
- for idx, model_id in enumerate(model_ids):
- # Determine current and target counts for this model
- current_state = fetch_state(api_base)
- current_ready = count_ready_instances_by_model(current_state, model_id)
- target_count = current_ready + 1
-
- print("=" * 80)
- print(
- f"[PRIMARY] Creating instance {idx + 1}/{len(model_ids)} for model: {model_id}"
- )
- print(
- f"[PRIMARY] Current ready count for {model_id}: {current_ready}, target: {target_count}"
- )
-
- # Build instance creation request data
- instance_data: dict[str, Any] = {"model_id": model_id}
- if sharding is not None:
- instance_data["sharding"] = sharding
- if instance_meta is not None:
- instance_data["instance_meta"] = instance_meta
-
- response = await _http_request_async(
- f"{api_base}/instance", method="POST", data=instance_data
- )
- print(f"[PRIMARY] Instance creation response: {response}")
-
- # Wait for one more instance of this model to be ready
- await wait_for_instances_ready(
- api_base, model_id, target_count, timeout_s=timeout_seconds
- )
- print(f"[PRIMARY] Instance {idx + 1}/{len(model_ids)} is ready")
- print("=" * 80)
- else:
- # Secondary: wait for expected counts of each model to be ready
- print("[SECONDARY] Waiting for all instances to be created and ready...")
- for model_id, expected_count in model_counts.items():
- await wait_for_instances_ready(
- api_base, model_id, expected_count, timeout_s=timeout_seconds
- )
-
- # Collect all instance IDs for all models
- state = fetch_state(api_base)
- for model_id in model_counts:
- ids = get_all_instance_ids_for_model(state, model_id)
- all_instance_ids.extend(ids)
-
- # Count total runners
- total_runners = 0
- for instance_id in all_instance_ids:
- runner_ids = get_runner_ids_for_instance(state, instance_id)
- total_runners += len(runner_ids)
-
- print(
- f"\nAll {len(all_instance_ids)} instance(s) with {total_runners} total runner(s) are ready!"
- )
- print(f"Instance IDs: {all_instance_ids}")
-
- if is_primary:
- # Run all stages once (requests will use available instances)
- # We use the first model_id for the benchmark requests
- benchmark_model_id = model_ids[0]
- print(f"\n{'=' * 80}")
- print(f"RUNNING BENCHMARK (using model: {benchmark_model_id})")
- print(f"Instances available: {len(all_instance_ids)}")
- print(f"{'=' * 80}")
-
- # Start metrics monitoring with 500ms interval to catch fast-completing tasks
- metrics_snapshots: list[MetricsSnapshot] = []
- stop_monitoring = asyncio.Event()
- monitoring_task = asyncio.create_task(
- monitor_metrics(
- api_base, metrics_snapshots, stop_monitoring, interval_seconds=0.5
- )
- )
-
- stage_results: list[StageResult] = []
- for stage in stages:
- result = await run_stage(
- api_base, benchmark_model_id, stage, no_overlap=no_overlap
- )
- stage_results.append(result)
-
- # Stop metrics monitoring
- print("\nStopping metrics monitoring...")
- stop_monitoring.set()
- await monitoring_task
- print(f"Collected {len(metrics_snapshots)} metrics snapshots")
-
- # Print final results
- print("\n" + "=" * 80)
- print("BENCHMARK COMPLETE - RESULTS SUMMARY")
- print("=" * 80)
- print(f"Instances tested: {len(all_instance_ids)}")
- print(f"Model IDs: {model_ids}")
- print(f"Instance IDs: {all_instance_ids}")
-
- for result in stage_results:
- print(f"\nStage: {result.name}")
- print(f" Total Requests: {result.total_requests}")
- print(f" Successful: {result.successful_requests}")
- print(f" Failed: {result.failed_requests}")
- print(f" Success Rate: {result.success_rate * 100:.1f}%")
- print(f" Total Tokens: {result.total_tokens}")
- print(f" Avg Tokens/Request: {result.avg_tokens_per_request:.1f}")
- print(f" Avg Time/Request: {result.avg_time_per_request:.2f}s")
- if result.avg_time_to_first_token is not None:
- if result.std_time_to_first_token is not None:
- print(
- f" Avg TTFT: {result.avg_time_to_first_token:.3f}s ± {result.std_time_to_first_token:.3f}s"
- )
- else:
- print(
- f" Avg TTFT: {result.avg_time_to_first_token:.3f}s"
- )
- if result.avg_ms_per_token is not None:
- if result.std_ms_per_token is not None:
- print(
- f" Avg ms/token: {result.avg_ms_per_token:.2f}ms ± {result.std_ms_per_token:.2f}ms"
- )
- else:
- print(f" Avg ms/token: {result.avg_ms_per_token:.2f}ms")
- if result.avg_decode_tps is not None:
- print(f" Avg Decode TPS: {result.avg_decode_tps:.2f} tokens/s")
-
- benchmark_completed_at = time.time()
-
- # Build comprehensive results document
- results_doc = {
- "metadata": {
- "benchmark_started_at": benchmark_started_at,
- "benchmark_completed_at": benchmark_completed_at,
- "total_duration_s": benchmark_completed_at - benchmark_started_at,
- "git_commit": git_commit,
- "config_file": str(config_path),
- "hardware_labels": hardware_labels or [],
- "expected_nodes": expected_nodes,
- "timeout_seconds": timeout_seconds,
- },
- "cluster": {
- "model_ids": model_ids,
- "instance_ids": all_instance_ids,
- "instance_count": len(all_instance_ids),
- "runner_count": total_runners,
- "sharding": sharding,
- "instance_meta": instance_meta,
- },
- "configuration": {
- "stages": [
- {
- "name": stage.name,
- "prompt_length": stage.prompt_length,
- "generation_length": stage.generation_length,
- "time_between_requests": stage.time_between_requests,
- "iterations": stage.iterations,
- }
- for stage in stages
- ]
- },
- "results": {
- "stages": [
- {
- "name": r.name,
- "total_requests": r.total_requests,
- "successful_requests": r.successful_requests,
- "failed_requests": r.failed_requests,
- "success_rate": round(r.success_rate, 4),
- "total_tokens": r.total_tokens,
- "avg_tokens_per_request": round(
- r.avg_tokens_per_request, 2
- ),
- "avg_time_per_request": round(r.avg_time_per_request, 3),
- "avg_time_to_first_token": round(
- r.avg_time_to_first_token, 3
- )
- if r.avg_time_to_first_token is not None
- else None,
- "std_time_to_first_token": round(
- r.std_time_to_first_token, 3
- )
- if r.std_time_to_first_token is not None
- else None,
- "avg_decode_tps": round(r.avg_decode_tps, 2)
- if r.avg_decode_tps is not None
- else None,
- "avg_ms_per_token": round(r.avg_ms_per_token, 2)
- if r.avg_ms_per_token is not None
- else None,
- "std_ms_per_token": round(r.std_ms_per_token, 2)
- if r.std_ms_per_token is not None
- else None,
- "stage_started_at": r.stage_started_at,
- "stage_completed_at": r.stage_completed_at,
- "stage_duration_s": r.stage_completed_at
- - r.stage_started_at,
- "requests": [
- {
- "request_id": req.request_id,
- "success": req.success,
- "tokens": req.tokens,
- "elapsed_s": round(req.elapsed_s, 3),
- "started_at": req.started_at,
- "completed_at": req.completed_at,
- "time_to_first_token_s": round(
- req.time_to_first_token_s, 3
- )
- if req.time_to_first_token_s is not None
- else None,
- "decode_tps": round(req.decode_tps, 2)
- if req.decode_tps is not None
- else None,
- "error": req.error,
- }
- for req in r.request_results
- ],
- }
- for r in stage_results
- ]
- },
- "metrics": {
- "snapshots": [
- {
- "timestamp": snapshot.timestamp,
- "node_memory": {
- node_id: {
- "ram_total_bytes": mem.ram_total_bytes,
- "ram_available_bytes": mem.ram_available_bytes,
- "ram_used_bytes": mem.ram_used_bytes,
- "swap_total_bytes": mem.swap_total_bytes,
- "swap_available_bytes": mem.swap_available_bytes,
- "swap_used_bytes": mem.swap_used_bytes,
- }
- for node_id, mem in snapshot.node_memory.items()
- },
- "instance_tasks": [
- {
- "instance_id": inst.instance_id,
- "node_id": inst.node_id,
- "pending_tasks": inst.pending_tasks,
- "running_tasks": inst.running_tasks,
- "total_active_tasks": inst.total_active_tasks,
- }
- for inst in snapshot.instance_tasks
- ],
- "node_tasks": [
- {
- "node_id": node.node_id,
- "pending_tasks": node.pending_tasks,
- "running_tasks": node.running_tasks,
- "total_active_tasks": node.total_active_tasks,
- "instance_count": node.instance_count,
- }
- for node in snapshot.node_tasks
- ],
- }
- for snapshot in metrics_snapshots
- ]
- },
- }
-
- # Output JSON summary
- print("\n" + "=" * 80)
- print("JSON RESULTS")
- print("=" * 80)
- print(json.dumps(results_doc, indent=2))
- print("=" * 80)
-
- # Save to file if path provided
- if results_output_path:
- print(f"Saving results to: {results_output_path}")
- with open(results_output_path, "w") as f:
- json.dump(results_doc, f, indent=2)
- print("Results saved successfully")
-
- # Cleanup all instances
- for instance_id in all_instance_ids:
- print(f"[PRIMARY] Cleaning up instance: {instance_id}")
- await _http_request_async(
- f"{api_base}/instance/{instance_id}", method="DELETE"
- )
- print(f"[PRIMARY] Instance {instance_id} deleted successfully")
- else:
- print(
- "[SECONDARY] Waiting with cluster (primary handles benchmark execution)"
- )
- # Secondary nodes wait until all instances of all models are deleted
- for model_id in model_counts:
- await wait_for_all_instances_deleted(api_base, model_id)
-
- return 0
-
- except TimeoutError as e:
- print("=" * 80)
- print(f"TIMEOUT ERROR: {e}")
- print("=" * 80)
- return 1
- except Exception as e:
- print("=" * 80)
- print(f"ERROR: {e}")
- import traceback
-
- traceback.print_exc()
- print("=" * 80)
- return 1
-
-
-def main() -> int:
- parser = argparse.ArgumentParser(
- description="Run unified benchmark for EXO (single or multi-stage)"
- )
- parser.add_argument("--api-port", type=int, required=True)
- parser.add_argument(
- "--config", type=Path, required=True, help="Path to YAML config file"
- )
- parser.add_argument(
- "--expected-nodes",
- type=int,
- required=True,
- help="Total number of nodes expected in the cluster",
- )
- parser.add_argument(
- "--is-primary", type=str, choices=["true", "false"], required=True
- )
- parser.add_argument("--timeout-seconds", type=int, default=1800)
- parser.add_argument(
- "--output", type=Path, help="Path to save detailed results JSON"
- )
- parser.add_argument("--git-commit", type=str, help="Git commit hash for metadata")
- parser.add_argument(
- "--hardware-labels", type=str, help="Comma-separated hardware labels"
- )
- args = parser.parse_args()
-
- api_base = f"http://localhost:{args.api_port}"
- is_primary = args.is_primary.lower() == "true"
- hardware_labels = args.hardware_labels.split(",") if args.hardware_labels else None
-
- return asyncio.run(
- run_benchmark(
- api_base,
- args.config,
- args.expected_nodes,
- is_primary,
- args.timeout_seconds,
- results_output_path=args.output,
- git_commit=args.git_commit,
- hardware_labels=hardware_labels,
- )
- )
-
-
-if __name__ == "__main__":
- sys.exit(main())
diff --git a/.github/scripts/build_matrix.py b/.github/scripts/build_matrix.py
deleted file mode 100644
index 2f139350..00000000
--- a/.github/scripts/build_matrix.py
+++ /dev/null
@@ -1,70 +0,0 @@
-#!/usr/bin/env python3
-import json
-import os
-from typing import NotRequired, TypedDict, cast
-
-import yaml
-
-
-class MatrixEntry(TypedDict):
- label: str
- index: int
-
-
-class MatrixInclude(TypedDict):
- label: str
- index: int
- is_primary: bool
- expected_nodes: int
-
-
-class Config(TypedDict):
- hardware_plan: dict[str, int]
- timeout_seconds: NotRequired[int]
- environment: NotRequired[dict[str, str]]
-
-
-# Read the config file
-config_file: str = os.environ["CONFIG_FILE"]
-with open(config_file, "r") as f:
- config: Config = cast(Config, yaml.safe_load(f))
-
-# Extract hardware plan from config
-plan: dict[str, int] = config["hardware_plan"]
-if not plan:
- raise ValueError(f"No hardware_plan found in {config_file}")
-
-# Build matrix entries
-entries: list[MatrixEntry] = []
-for label, count in plan.items():
- for idx in range(count):
- entries.append({"label": label, "index": idx})
-
-total_nodes: int = len(entries)
-matrix: dict[str, list[MatrixInclude]] = {
- "include": [
- {
- "label": e["label"],
- "index": e["index"],
- "is_primary": (i == 0),
- "expected_nodes": total_nodes,
- }
- for i, e in enumerate(entries)
- ]
-}
-
-# Extract other config values
-timeout_seconds: int = config.get("timeout_seconds", 600)
-environment: dict[str, str] = config.get("environment", {})
-
-# Output to GitHub Actions
-with open(os.environ["GITHUB_OUTPUT"], "a") as f:
- f.write(f"matrix={json.dumps(matrix)}\n")
- f.write(f"config_file={config_file}\n")
- f.write(f"timeout_seconds={timeout_seconds}\n")
- f.write(f"environment={json.dumps(environment)}\n")
-
-print(f"Matrix: {json.dumps(matrix)}")
-print(f"Config file: {config_file}")
-print(f"Timeout: {timeout_seconds}")
-print(f"Environment: {json.dumps(environment)}")
diff --git a/.github/workflows/BENCH_USAGE.md b/.github/workflows/BENCH_USAGE.md
deleted file mode 100644
index b61d31da..00000000
--- a/.github/workflows/BENCH_USAGE.md
+++ /dev/null
@@ -1,156 +0,0 @@
-# Benchmark Workflow Usage
-
-## Overview
-
-The `bench_matrix.yml` workflow enables distributed benchmarking of models across multiple self-hosted macOS runners with different hardware configurations.
-
-## Workflow Inputs
-
-| Input | Description | Default | Required |
-|-------|-------------|---------|----------|
-| `model_id` | Model ID to benchmark | `mlx-community/Llama-3.2-1B-Instruct-4bit` | Yes |
-| `hardware_plan` | JSON mapping of runner labels to counts | `{"M4PRO_GPU16_24GB": 1}` | Yes |
-| `prompt` | Benchmark prompt text | `What is the capital of France?` | No |
-| `timeout_seconds` | Timeout for instance/runner readiness | `600` | No |
-
-## Hardware Plan Format
-
-The `hardware_plan` input is a JSON object mapping runner labels to the number of machines:
-
-```json
-{
- "M4PRO_GPU16_24GB": 2,
- "M3ULTRA_GPU80_512GB": 1
-}
-```
-
-This example would:
-- Start 2 runners with the `M4PRO_GPU16_24GB` label
-- Start 1 runner with the `M3ULTRA_GPU80_512GB` label
-- Total of 3 runners coordinating on a single distributed inference instance
-
-## How It Works
-
-1. **Planning Job** (`plan`)
- - Runs on `ubuntu-latest`
- - Parses the `hardware_plan` JSON
- - Generates a dynamic matrix with one entry per runner
- - Only the first runner (index 0) is marked as `is_primary`
-
-2. **Benchmark Worker Jobs** (`bench_worker`)
- - Each job runs on a self-hosted macOS runner with the specified label
- - All runners start EXO in parallel
- - The primary runner creates the model instance
- - All runners wait for their assigned runner to be ready (Loaded/Running status)
- - The primary runner executes the benchmark and prints results
- - The primary runner deletes the instance
-
-## Example Usage
-
-### Single Machine Benchmark
-
-```yaml
-model_id: mlx-community/Llama-3.2-1B-Instruct-4bit
-hardware_plan: '{"M4PRO_GPU16_24GB": 1}'
-prompt: What is the capital of France?
-timeout_seconds: 600
-```
-
-### Multi-Machine Distributed Benchmark
-
-```yaml
-model_id: mlx-community/Llama-3.2-3B-Instruct-4bit
-hardware_plan: '{"M4PRO_GPU16_24GB": 2, "M3ULTRA_GPU80_512GB": 1}'
-prompt: Explain quantum computing in simple terms.
-timeout_seconds: 900
-```
-
-## Benchmark Output
-
-The primary runner outputs a JSON object with benchmark results:
-
-```json
-{
- "model_id": "mlx-community/Llama-3.2-1B-Instruct-4bit",
- "instance_id": "abc-123-def",
- "tokens": 42,
- "elapsed_s": 2.451,
- "tps": 17.136
-}
-```
-
-Where:
-- `tokens`: Number of chunks/tokens generated
-- `elapsed_s`: Total elapsed time in seconds
-- `tps`: Tokens per second (tokens / elapsed_s)
-
-## Runner Requirements
-
-Each self-hosted runner must:
-- Be labeled with appropriate hardware tags (e.g., `M4PRO_GPU16_24GB`)
-- Have the `self-hosted` and `macOS` labels
-- Have Nix installed with flakes enabled
-- Have network connectivity to other runners in the same job
-
-## Architecture
-
-```
-┌─────────────────────────────────────────────────────────────┐
-│ GitHub Actions Workflow (bench_matrix.yml) │
-├─────────────────────────────────────────────────────────────┤
-│ │
-│ ┌────────────────┐ │
-│ │ Plan Job │ │
-│ │ (ubuntu) │──┬─► Matrix: [{label, index, primary}] │
-│ └────────────────┘ │ │
-│ │ │
-│ ┌───────────────────▼──────────────────────────────────┐ │
-│ │ Bench Worker Jobs (Matrix) │ │
-│ ├──────────────────────────────────────────────────────┤ │
-│ │ │ │
-│ │ Runner 0 (Primary) Runner 1 Runner 2 │ │
-│ │ ┌─────────────┐ ┌─────────────┐ ┌──────────┐ │ │
-│ │ │ Start EXO │ │ Start EXO │ │ Start EXO│ │ │
-│ │ │ Create Inst │ │ Wait... │ │ Wait... │ │ │
-│ │ │ Wait Ready │ │ Wait Ready │ │ Wait... │ │ │
-│ │ │ Run Bench │ │ (idle) │ │ (idle) │ │ │
-│ │ │ Print TPS │ │ │ │ │ │ │
-│ │ │ Delete Inst │ │ │ │ │ │ │
-│ │ └─────────────┘ └─────────────┘ └──────────┘ │ │
-│ └───────────────────────────────────────────────────────┘ │
-└─────────────────────────────────────────────────────────────┘
-```
-
-## Implementation Details
-
-### `scripts/bench.py`
-
-A standalone Python script that:
-- Creates instance (primary only)
-- Polls `/state` endpoint until instance and all runners are ready
-- Executes chat completion with timing (primary only)
-- Parses SSE stream and counts tokens
-- Computes TPS metrics
-- Cleans up instance (primary only)
-
-### Key Functions
-
-- `wait_for_instance()`: Polls until instance with model_id appears
-- `wait_for_runners_ready()`: Polls until expected number of runners reach Loaded/Running status
-- `run_benchmark()`: Executes chat completion, measures time, counts tokens
-
-## Troubleshooting
-
-### Instance never becomes ready
-- Check EXO logs in the workflow output
-- Verify model_id is valid and accessible
-- Increase `timeout_seconds`
-
-### Runner mismatch
-- Ensure hardware_plan counts match available labeled runners
-- Check runner labels match exactly (case-sensitive)
-
-### Network issues
-- Verify runners can communicate on the network
-- Check firewall rules between runner hosts
-
diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml
deleted file mode 100644
index dda16435..00000000
--- a/.github/workflows/bench.yml
+++ /dev/null
@@ -1,305 +0,0 @@
-name: bench
-
-on: [push]
-
-jobs:
- plan:
- if: contains(github.event.head_commit.message, '/bench')
- runs-on: ubuntu-latest
- outputs:
- matrix: ${{ steps.build.outputs.matrix }}
- config_file: ${{ steps.build.outputs.config_file }}
- timeout_seconds: ${{ steps.build.outputs.timeout_seconds }}
- environment: ${{ steps.build.outputs.environment }}
- steps:
- - name: Checkout repository
- uses: actions/checkout@v4
-
- - name: Build matrix from config file
- id: build
- shell: bash
- run: |
- set -euo pipefail
- CONFIG_FILE='.github/configs/bench_simple.yaml'
- export CONFIG_FILE
- echo "Config file: $CONFIG_FILE"
- python3 .github/scripts/build_matrix.py
-
- bench_worker:
- needs: plan
- strategy:
- fail-fast: false
- matrix: ${{ fromJSON(needs.plan.outputs.matrix) }}
- name: "bench on ${{ matrix.label }} [${{ matrix.index }}]"
- runs-on: [self-hosted, macOS, "${{ matrix.label }}"]
- steps:
- - name: Checkout repository
- uses: actions/checkout@v4
- with:
- lfs: false
-
- - name: Configure git user
- run: |
- git config --local user.email "github-actions@users.noreply.github.com"
- git config --local user.name "github-actions bot"
- shell: bash
-
- # TODO: this is mega hacky and I'd like a simpler solution.
- - name: Setup Nix Environment
- run: |
- echo "Checking for nix installation..."
-
- # Check if nix is already available
- if command -v nix >/dev/null 2>&1; then
- echo "Nix already in PATH"
- # Try sourcing profile scripts to set up environment properly
- elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
- echo "Sourcing multi-user nix-daemon profile script"
- source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
- elif [ -f "$HOME/.nix-profile/etc/profile.d/nix.sh" ]; then
- echo "Sourcing single-user nix profile script"
- source "$HOME/.nix-profile/etc/profile.d/nix.sh"
- elif [ -f /nix/var/nix/profiles/per-user/$USER/profile/etc/profile.d/nix.sh ]; then
- echo "Sourcing per-user nix profile script"
- source /nix/var/nix/profiles/per-user/$USER/profile/etc/profile.d/nix.sh
- elif [ -f /etc/profile.d/nix.sh ]; then
- echo "Sourcing system-wide nix profile script"
- source /etc/profile.d/nix.sh
- # Fallback: manually add nix to PATH if binary exists
- elif [ -f /nix/var/nix/profiles/default/bin/nix ]; then
- echo "Found nix binary, manually adding to PATH"
- export PATH="/nix/var/nix/profiles/default/bin:$PATH"
- elif [ -f "$HOME/.nix-profile/bin/nix" ]; then
- echo "Found nix binary in user profile, manually adding to PATH"
- export PATH="$HOME/.nix-profile/bin:$PATH"
- else
- echo "Nix not found. Debugging info:"
- echo "USER: $USER"
- echo "HOME: $HOME"
- echo "Current PATH: $PATH"
- echo ""
- echo "Checking common Nix locations:"
- echo " /nix/var/nix/profiles/default/bin/nix:"
- ls -la /nix/var/nix/profiles/default/bin/nix 2>/dev/null || echo " Not found"
- echo " /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh:"
- ls -la /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh 2>/dev/null || echo " Not found"
- echo " ~/.nix-profile/etc/profile.d/nix.sh:"
- ls -la "$HOME/.nix-profile/etc/profile.d/nix.sh" 2>/dev/null || echo " Not found"
- echo " /nix/var/nix/profiles/per-user/$USER/profile/etc/profile.d/nix.sh:"
- ls -la "/nix/var/nix/profiles/per-user/$USER/profile/etc/profile.d/nix.sh" 2>/dev/null || echo " Not found"
- echo ""
- echo "/nix directory structure:"
- ls -la /nix 2>/dev/null || echo " /nix directory not found"
- echo ""
- echo "/nix/var:"
- ls -la /nix/var 2>/dev/null || echo " /nix/var not found"
- echo ""
- echo "/nix/store:"
- ls -la /nix/store 2>/dev/null | head -20 || echo " /nix/store not found"
- echo ""
- echo "GitHub Actions runner is running as user '$USER'."
- echo "If Nix is installed for a different user, either:"
- echo " 1. Install Nix for user '$USER' (multi-user install recommended)"
- echo " 2. Configure the runner service to run as the user with Nix installed"
- echo " 3. Ensure Nix is installed system-wide with proper daemon setup"
- exit 1
- fi
-
- # Verify nix is available and persist to GITHUB_ENV
- if command -v nix >/dev/null 2>&1; then
- echo "✓ Nix is available"
- nix --version
- echo "PATH=$PATH" >> $GITHUB_ENV
- if [ -n "$NIX_PATH" ]; then
- echo "NIX_PATH=$NIX_PATH" >> $GITHUB_ENV
- fi
- else
- echo "ERROR: Failed to set up Nix"
- echo "PATH after setup attempt: $PATH"
- exit 1
- fi
- shell: bash
-
- - name: Setup EXO_HOME and API_PORT
- run: |
- EXO_HOME=$(mktemp -d -t exo-e2e-XXXXXXXX)
- API_PORT=$((49152 + RANDOM % (65535 - 49152 + 1)))
- EXO_MODELS_DIR="$HOME/.exo/models"
- EXO_LIBP2P_NAMESPACE="bench-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}"
- echo "EXO_HOME=$EXO_HOME" >> "$GITHUB_ENV"
- echo "API_PORT=$API_PORT" >> "$GITHUB_ENV"
- echo "EXO_MODELS_DIR=$EXO_MODELS_DIR" >> "$GITHUB_ENV"
- echo "EXO_LIBP2P_NAMESPACE=$EXO_LIBP2P_NAMESPACE" >> "$GITHUB_ENV"
- echo "Created EXO_HOME: $EXO_HOME"
- echo "Generated API_PORT: $API_PORT"
- echo "Using models from: $EXO_MODELS_DIR"
- echo "Using libp2p namespace: $EXO_LIBP2P_NAMESPACE"
- shell: bash
-
- - name: Configure local MLX if available
- run: |
- echo "=== DEBUG: Checking for local MLX configuration ==="
- MODIFIED=false
-
- echo "Checking for /Users/Shared/mlx directory..."
- if [ -d "/Users/Shared/mlx" ]; then
- echo "✓ Found /Users/Shared/mlx"
- ls -la /Users/Shared/mlx | head -5
- echo "Enabling local mlx path in pyproject.toml"
- sed -i.bak 's|^# mlx = { path = "/Users/Shared/mlx", editable=true }$|mlx = { path = "/Users/Shared/mlx", editable=true }|' pyproject.toml
- MODIFIED=true
- else
- echo "✗ /Users/Shared/mlx not found, will use PyPI version"
- fi
-
- echo "Checking for /Users/Shared/mlx-lm directory..."
- if [ -d "/Users/Shared/mlx-lm" ]; then
- echo "✓ Found /Users/Shared/mlx-lm"
- ls -la /Users/Shared/mlx-lm | head -5
- echo "Enabling local mlx-lm path in pyproject.toml"
- sed -i.bak 's|^# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }$|mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }|' pyproject.toml
- MODIFIED=true
- else
- echo "✗ /Users/Shared/mlx-lm not found, will use PyPI version"
- fi
-
- if [ "$MODIFIED" = true ]; then
- echo "=== Modified pyproject.toml [tool.uv.sources] section: ==="
- sed -n '/\[tool\.uv\.sources\]/,/^\[/{/^\[tool\.uv\.sources\]/p; /^\[/!p;}' pyproject.toml
- echo "=== Regenerating uv.lock with local MLX paths... ==="
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command uv lock --upgrade-package mlx --upgrade-package mlx-lm
- echo "✓ Lock file regenerated"
- else
- echo "⚠ No local MLX directories found, using PyPI packages"
- fi
- echo "=== DEBUG: Local MLX configuration complete ==="
- shell: bash
-
- - name: Sync dependencies
- run: |
- if [ -d "/Users/Shared/test" ]; then
- pushd /Users/Shared/test
- uv sync --reinstall
- popd
- fi
- echo "Running just sync to ensure clean dependencies..."
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command just sync
- shell: bash
-
- - name: Start EXO and run bench script
- shell: bash
- env:
- IS_PRIMARY: ${{ matrix.is_primary }}
- EXPECTED_NODES: ${{ matrix.expected_nodes }}
- HARDWARE_LABEL: ${{ matrix.label }}
- CONFIG_FILE: ${{ needs.plan.outputs.config_file }}
- TIMEOUT_SECONDS: ${{ needs.plan.outputs.timeout_seconds }}
- ENVIRONMENT_JSON: ${{ needs.plan.outputs.environment }}
- run: |
- set -euo pipefail
-
- # Parse environment variables from config
- ENV_VARS=""
- if [ -n "$ENVIRONMENT_JSON" ] && [ "$ENVIRONMENT_JSON" != "{}" ]; then
- ENV_VARS=$(echo "$ENVIRONMENT_JSON" | python3 -c "import sys, json; env = json.load(sys.stdin); print(' '.join([f'{k}={v}' for k, v in env.items()]))")
- fi
-
- echo "Starting EXO with API_PORT=${API_PORT} EXO_HOME=${EXO_HOME} EXO_LIBP2P_NAMESPACE=${EXO_LIBP2P_NAMESPACE}"
- echo "Environment variables from config: $ENV_VARS"
- LOG_FILE=/tmp/exo.log
- : > "$LOG_FILE"
-
- MASTER_FLAG=""
- if [ "$IS_PRIMARY" = "true" ]; then
- MASTER_FLAG="-m"
- fi
-
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command bash -c \
- "EXO_HOME=$EXO_HOME EXO_MODELS_DIR=$EXO_MODELS_DIR EXO_LIBP2P_NAMESPACE=$EXO_LIBP2P_NAMESPACE $ENV_VARS PYTHONUNBUFFERED=1 PYTHONDEBUG=1 PYTHONPATH=. uv run exo $MASTER_FLAG --api-port $API_PORT" \
- >> "$LOG_FILE" 2>&1 &
-
- EXO_PID=$!
- echo "Started EXO in background with PID: $EXO_PID"
- echo "Log file: $LOG_FILE"
-
- cleanup() {
- echo '=== EXO log (tail) ==='
- tail -n 300 "$LOG_FILE" || true
- if ps -p "$EXO_PID" >/dev/null 2>&1; then
- echo "Killing EXO (PID $EXO_PID)"
- kill "$EXO_PID" || true
- fi
- }
- trap cleanup EXIT
-
- for i in $(seq 1 60); do
- if curl -s "http://localhost:${API_PORT}/state" >/dev/null 2>&1; then
- echo "EXO API ready"
- break
- fi
- if ! ps -p "$EXO_PID" >/dev/null 2>&1; then
- echo "EXO terminated early"; sed -n '1,200p' "$LOG_FILE" || true; exit 1
- fi
- sleep 1
- done
-
- RESULTS_FILE="/tmp/bench_results_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}_$(date +%s).json"
- echo "Results will be saved to: $RESULTS_FILE"
- echo "RESULTS_FILE=$RESULTS_FILE" >> "$GITHUB_ENV"
-
- echo "Running bench script with config: $CONFIG_FILE, timeout: $TIMEOUT_SECONDS"
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command bash -c \
- "PYTHONUNBUFFERED=1 uv run --no-project --with pyyaml --with pydantic python .github/scripts/bench.py \
- --api-port $API_PORT \
- --config $CONFIG_FILE \
- --expected-nodes ${EXPECTED_NODES} \
- --is-primary ${IS_PRIMARY} \
- --timeout-seconds ${TIMEOUT_SECONDS} \
- --output $RESULTS_FILE \
- --git-commit ${GITHUB_SHA} \
- --hardware-labels ${HARDWARE_LABEL}"
-
- - name: Install AWS CLI
- if: always() && env.RESULTS_FILE && matrix.is_primary
- run: |
- if ! command -v aws &> /dev/null; then
- echo "AWS CLI not found, installing..."
- brew install awscli
- else
- echo "AWS CLI already installed"
- fi
- shell: bash
-
- - name: Upload results to S3
- if: always() && env.RESULTS_FILE && matrix.is_primary
- env:
- AWS_ACCESS_KEY_ID: ${{ secrets.S3_BENCHMARKS_AWS_ACCESS_KEY_ID }}
- AWS_SECRET_ACCESS_KEY: ${{ secrets.S3_BENCHMARKS_AWS_SECRET_ACCESS_KEY }}
- AWS_DEFAULT_REGION: us-east-1
- run: |
- echo "Checking for results file: $RESULTS_FILE"
- echo "Is primary: ${{ matrix.is_primary }}"
-
- if [ -f "$RESULTS_FILE" ]; then
- TIMESTAMP=$(date -u +%Y/%m/%d/%H%M%S)
- S3_KEY="bench/${TIMESTAMP}_${GITHUB_SHA:0:8}_${GITHUB_RUN_ID}.json"
- echo "Uploading results to s3://exo-benchmark-results/$S3_KEY"
-
- aws s3 cp "$RESULTS_FILE" "s3://exo-benchmark-results/$S3_KEY" \
- --content-type application/json \
- --metadata "commit=${GITHUB_SHA},run_id=${GITHUB_RUN_ID},branch=${GITHUB_REF_NAME}"
-
- echo "Results uploaded successfully"
- echo "View at: https://exo-benchmark-results.s3.amazonaws.com/$S3_KEY"
- else
- echo "Results file not found at: $RESULTS_FILE"
- echo "Skipping upload"
- fi
- shell: bash
-
- - name: Cleanup EXO_HOME
- run: |
- echo "Cleaning up EXO_HOME: $EXO_HOME"
- rm -rf "$EXO_HOME"
- shell: bash
- if: always()
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
new file mode 100644
index 00000000..2443b03d
--- /dev/null
+++ b/bench/exo_bench.py
@@ -0,0 +1,526 @@
+#!/usr/bin/env python3
+# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
+from __future__ import annotations
+
+import argparse
+import http.client
+import json
+import os
+import time
+from collections.abc import Callable
+from statistics import mean
+from typing import Any
+from urllib.parse import urlencode
+
+from loguru import logger
+from transformers import AutoTokenizer
+
+from exo.shared.models.model_cards import MODEL_CARDS
+from exo.shared.types.memory import Memory
+
+
+class ExoHttpError(RuntimeError):
+ def __init__(self, status: int, reason: str, body_preview: str):
+ super().__init__(f"HTTP {status} {reason}: {body_preview}")
+ self.status = status
+
+
+class ExoClient:
+ def __init__(self, host: str, port: int, timeout_s: float = 2400.0):
+ self.host = host
+ self.port = port
+ self.timeout_s = timeout_s
+
+ def request_json(
+ self,
+ method: str,
+ path: str,
+ params: dict[str, Any] | None = None,
+ body: dict[str, Any] | None = None,
+ headers: dict[str, str] | None = None,
+ ) -> Any:
+ if not path.startswith("/"):
+ path = "/" + path
+ if params:
+ path = path + "?" + urlencode(params)
+
+ conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
+ try:
+ payload: bytes | None = None
+ hdrs: dict[str, str] = {"Accept": "application/json"}
+
+ if body is not None:
+ payload = json.dumps(body).encode("utf-8")
+ hdrs["Content-Type"] = "application/json"
+ if headers:
+ hdrs.update(headers)
+
+ conn.request(method.upper(), path, body=payload, headers=hdrs)
+ resp = conn.getresponse()
+ raw = resp.read()
+ text = raw.decode("utf-8", errors="replace") if raw else ""
+
+ if resp.status >= 400:
+ raise ExoHttpError(resp.status, resp.reason, text[:300])
+
+ if not text:
+ return None
+ return json.loads(text)
+ finally:
+ conn.close()
+
+ def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
+ return self.request_json("POST", "/bench/chat/completions", body=payload)
+
+
+def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
+ if len(instance) != 1:
+ raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
+
+ tag = next(iter(instance))
+ inner = instance[tag]
+ if not isinstance(inner, dict):
+ raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
+ return inner
+
+
+def instance_id_from_instance(instance: dict[str, Any]) -> str:
+ inner = unwrap_instance(instance)
+ return str(inner["instanceId"])
+
+
+def nodes_used_in_instance(instance: dict[str, Any]) -> int:
+ inner = unwrap_instance(instance)
+ return len(inner["shardAssignments"]["nodeToRunner"])
+
+
+def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
+ inner = unwrap_instance(instance)
+ runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+ return list(runner_to_shard.keys())
+
+
+def runner_ready(runner: dict[str, Any]) -> bool:
+ return "RunnerReady" in runner
+
+
+def wait_for_instance_ready(
+ client: ExoClient, instance_id: str, timeout: float = 24000.0
+) -> None:
+ start_time = time.time()
+ while time.time() - start_time < timeout:
+ state = client.request_json("GET", "/state")
+ instances = state.get("instances", {})
+
+ if instance_id not in instances:
+ time.sleep(0.1)
+ continue
+
+ instance = instances[instance_id]
+ runner_ids = runner_ids_from_instance(instance)
+ runners = state.get("runners", {})
+
+ if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+ return
+
+ time.sleep(0.1)
+
+ raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
+
+
+def wait_for_instance_gone(
+ client: ExoClient, instance_id: str, timeout: float = 3.0
+) -> None:
+ start_time = time.time()
+ while time.time() - start_time < timeout:
+ try:
+ client.request_json("GET", f"/instance/{instance_id}")
+ time.sleep(0.4)
+ except ExoHttpError as e:
+ if e.status == 404:
+ return
+
+ raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
+
+
+def format_peak_memory(b: float) -> str:
+ for unit in ["B", "KB", "MB", "GB", "TB"]:
+ if b < 1024.0:
+ return f"{b:.2f}{unit}"
+ b /= 1024.0
+ raise ValueError("You're using petabytes of memory. Something went wrong...")
+
+
+def parse_int_list(values: list[str]) -> list[int]:
+ items: list[int] = []
+ for v in values:
+ for part in v.split(","):
+ part = part.strip()
+ if part:
+ items.append(int(part))
+
+ seen: set[int] = set()
+ out: list[int] = []
+ for x in items:
+ if x not in seen:
+ out.append(x)
+ seen.add(x)
+ return out
+
+
+def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
+ models = client.request_json("GET", "/models") or {}
+ data = models.get("data") or []
+
+ for m in data:
+ if m.get("id") == model_arg:
+ short_id = str(m["id"])
+ full_id = str(m.get("hugging_face_id") or m["id"])
+ return short_id, full_id
+
+ for m in data:
+ if m.get("hugging_face_id") == model_arg:
+ short_id = str(m["id"])
+ full_id = str(m["hugging_face_id"])
+ return short_id, full_id
+
+ raise ValueError(f"Model not found in /models: {model_arg}")
+
+
+def placement_filter(instance_meta: str, wanted: str) -> bool:
+ s = (instance_meta or "").lower()
+ if wanted == "both":
+ return ("ring" in s) or ("jaccl" in s)
+ return wanted in s
+
+
+def sharding_filter(sharding: str, wanted: str) -> bool:
+ s = (sharding or "").lower()
+ if wanted == "both":
+ return ("pipeline" in s) or ("tensor" in s)
+ return wanted in s
+
+
+def run_one_completion(
+ client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
+) -> tuple[dict[str, Any], int]:
+ content, pp_tokens = prompt_sizer.build(pp_hint)
+ payload: dict[str, Any] = {
+ "model": model_id,
+ "messages": [{"role": "user", "content": content}],
+ "stream": False,
+ "max_tokens": tg,
+ }
+
+ t0 = time.perf_counter()
+ out = client.post_bench_chat_completions(payload)
+ elapsed = time.perf_counter() - t0
+
+ stats = out.get("generation_stats")
+
+ preview = (out.get("choices") or [{}])[0]["message"]["content"][:200]
+
+ return {
+ "elapsed_s": elapsed,
+ "output_text_preview": preview,
+ "stats": stats,
+ }, pp_tokens
+
+
+class PromptSizer:
+ def __init__(self, tokenizer: Any, atom: str = "a "):
+ self.tokenizer = tokenizer
+ self.atom = atom
+ self.count_fn = PromptSizer._make_counter(tokenizer)
+ self.base_tokens = self.count_fn("")
+
+ @staticmethod
+ def _make_counter(tokenizer: Any) -> Callable[[str], int]:
+ def count_fn(user_content: str) -> int:
+ messages = [{"role": "user", "content": user_content}]
+ ids = tokenizer.apply_chat_template(
+ messages, tokenize=True, add_generation_prompt=True
+ )
+ return int(len(ids))
+
+ return count_fn
+
+ def build(self, target_prompt_tokens: int) -> tuple[str, int]:
+ target = int(target_prompt_tokens)
+ if target < self.base_tokens:
+ raise RuntimeError(
+ f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
+ )
+
+ content = ""
+ tok = self.count_fn(content)
+
+ while tok < target:
+ content += self.atom
+ tok = self.count_fn(content)
+
+ if tok != target:
+ raise RuntimeError(
+ f"Overshot: got {tok} tokens (target {target}). "
+ f"Pick a different atom (try ' a' or '\\n' or '0 ')."
+ )
+
+ return content, tok
+
+
+def main() -> int:
+ ap = argparse.ArgumentParser(
+ prog="exo-bench",
+ description="Benchmark exo model throughput across placement previews.",
+ )
+ ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
+ ap.add_argument(
+ "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
+ )
+ ap.add_argument("--model", required=True, help="Model short id or huggingface id")
+ ap.add_argument(
+ "--pp",
+ nargs="+",
+ required=True,
+ help="Prompt-size hints (ints). Accepts commas.",
+ )
+ ap.add_argument(
+ "--tg",
+ nargs="+",
+ required=True,
+ help="Generation lengths (ints). Accepts commas.",
+ )
+ ap.add_argument(
+ "--max-nodes",
+ type=int,
+ default=4,
+ help="Only consider placements using <= this many nodes.",
+ )
+ ap.add_argument(
+ "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
+ )
+ ap.add_argument(
+ "--sharding", choices=["pipeline", "tensor", "both"], default="both"
+ )
+ ap.add_argument(
+ "--skip-pipeline-jaccl",
+ action="store_true",
+ help="Pipeline jaccl is often pointless, skip by default",
+ )
+ ap.add_argument(
+ "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
+ )
+ ap.add_argument(
+ "--warmup",
+ type=int,
+ default=0,
+ help="Warmup runs per placement (uses first pp/tg).",
+ )
+ ap.add_argument(
+ "--timeout", type=float, default=2400.0, help="HTTP timeout (seconds)."
+ )
+ ap.add_argument(
+ "--json-out",
+ default="bench/results.json",
+ help="Write raw per-run results JSON to this path.",
+ )
+ ap.add_argument(
+ "--dry-run", action="store_true", help="List selected placements and exit."
+ )
+ args = ap.parse_args()
+
+ pp_list = parse_int_list(args.pp)
+ tg_list = parse_int_list(args.tg)
+ if not pp_list or not tg_list:
+ logger.error("pp and tg lists must be non-empty")
+ return 2
+ if args.repeat <= 0:
+ logger.error("--repeat must be >= 1")
+ return 2
+
+ client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+ short_id, full_model_id = resolve_model_short_id(client, args.model)
+
+ previews_resp = client.request_json(
+ "GET", "/instance/previews", params={"model_id": short_id}
+ )
+ previews = previews_resp.get("previews") or []
+
+ tokenizer = AutoTokenizer.from_pretrained(
+ full_model_id,
+ trust_remote_code=True,
+ )
+ if tokenizer is None:
+ raise RuntimeError("[exo-bench] tokenizer load failed")
+
+ try:
+ prompt_sizer = PromptSizer(tokenizer)
+ logger.debug(f"[exo-bench] loaded tokenizer: {full_model_id} for prompt sizer")
+ except Exception:
+ logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
+ raise
+
+ selected: list[dict[str, Any]] = []
+ for p in previews:
+ if p.get("error") is not None:
+ continue
+ if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
+ continue
+ if not sharding_filter(str(p.get("sharding", "")), args.sharding):
+ continue
+
+ instance = p.get("instance")
+ if not isinstance(instance, dict):
+ continue
+
+ n = nodes_used_in_instance(instance)
+ # Skip tensor ring single node as it is pointless when pipeline ring
+ if n == 1 and (
+ (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
+ or (
+ args.instance_meta == "both"
+ and "jaccl" in p.get("instance_meta", "").lower()
+ )
+ ):
+ continue
+
+ if (
+ args.skip_pipeline_jaccl
+ and (
+ args.instance_meta == "both"
+ and "jaccl" in p.get("instance_meta", "").lower()
+ )
+ and (
+ args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
+ )
+ ):
+ continue
+
+ if 0 < n <= args.max_nodes:
+ selected.append(p)
+
+ if not selected:
+ logger.error("No valid placements matched your filters.")
+ return 1
+
+ selected.sort(
+ key=lambda p: (
+ str(p.get("instance_meta", "")),
+ str(p.get("sharding", "")),
+ -nodes_used_in_instance(p["instance"]),
+ ),
+ reverse=True,
+ )
+
+ logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
+ logger.info(f"placements: {len(selected)}")
+ for p in selected:
+ logger.info(
+ f" - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
+ )
+
+ if args.dry_run:
+ return 0
+
+ all_rows: list[dict[str, Any]] = []
+
+ for preview in selected:
+ instance = preview["instance"]
+ instance_id = instance_id_from_instance(instance)
+
+ sharding = str(preview["sharding"])
+ instance_meta = str(preview["instance_meta"])
+ n_nodes = nodes_used_in_instance(instance)
+
+ logger.info("=" * 80)
+ logger.info(
+ f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
+ )
+
+ client.request_json("POST", "/instance", body={"instance": instance})
+ wait_for_instance_ready(client, instance_id)
+
+ time.sleep(1)
+
+ try:
+ for i in range(args.warmup):
+ run_one_completion(
+ client, full_model_id, pp_list[0], tg_list[0], prompt_sizer
+ )
+ logger.debug(f" warmup {i + 1}/{args.warmup} done")
+
+ for pp in pp_list:
+ if (
+ pp * n_nodes > 2048
+ and "ring" in instance_meta.lower()
+ and "tensor" in sharding.lower()
+ ):
+ model_card = MODEL_CARDS[short_id]
+ if model_card.metadata.storage_size > Memory.from_gb(10):
+ logger.info(
+ f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
+ )
+ continue
+ for tg in tg_list:
+ runs: list[dict[str, Any]] = []
+ for r in range(args.repeat):
+ time.sleep(3)
+ try:
+ row, actual_pp_tokens = run_one_completion(
+ client, full_model_id, pp, tg, prompt_sizer
+ )
+ except Exception as e:
+ logger.error(e)
+ continue
+ row.update(
+ {
+ "model_short_id": short_id,
+ "model_id": full_model_id,
+ "placement_sharding": sharding,
+ "placement_instance_meta": instance_meta,
+ "placement_nodes": n_nodes,
+ "instance_id": instance_id,
+ "pp_tokens": actual_pp_tokens,
+ "tg": tg,
+ "repeat_index": r,
+ }
+ )
+ runs.append(row)
+ all_rows.append(row)
+
+ if runs:
+ prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+ gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+ ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+ gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+ peak = mean(
+ x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+ )
+
+ logger.info(
+ f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
+ f"prompt_tokens={ptok} gen_tokens={gtok} "
+ f"peak_memory={format_peak_memory(peak)}\n"
+ )
+ time.sleep(2)
+ finally:
+ try:
+ client.request_json("DELETE", f"/instance/{instance_id}")
+ except ExoHttpError as e:
+ if e.status != 404:
+ raise
+ wait_for_instance_gone(client, instance_id)
+ logger.debug(f"Deleted instance {instance_id}")
+
+ time.sleep(5)
+
+ if args.json_out:
+ with open(args.json_out, "w", encoding="utf-8") as f:
+ json.dump(all_rows, f, indent=2, ensure_ascii=False)
+ logger.debug(f"\nWrote results JSON: {args.json_out}")
+
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/pyproject.toml b/pyproject.toml
index 749bbadc..e51c9ec2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -82,7 +82,7 @@ build-backend = "uv_build"
###
[tool.basedpyright]
-include = [".venv/lib/mlx", ".venv/lib/mlx_lm", "src"]
+include = [".venv/lib/mlx", ".venv/lib/mlx_lm", "src", "bench"]
typeCheckingMode = "strict"
failOnWarnings = true
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index fcb58a12..30f87e2a 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -27,6 +27,8 @@ from exo.shared.logging import InterceptLogger
from exo.shared.models.model_cards import MODEL_CARDS
from exo.shared.models.model_meta import get_model_meta
from exo.shared.types.api import (
+ BenchChatCompletionResponse,
+ BenchChatCompletionTaskParams,
ChatCompletionChoice,
ChatCompletionMessage,
ChatCompletionResponse,
@@ -34,6 +36,7 @@ from exo.shared.types.api import (
CreateInstanceResponse,
DeleteInstanceResponse,
FinishReason,
+ GenerationStats,
ModelList,
ModelListModel,
PlaceInstanceParams,
@@ -172,6 +175,7 @@ class API:
self.app.post("/v1/chat/completions", response_model=None)(
self.chat_completions
)
+ self.app.post("/bench/chat/completions")(self.bench_chat_completions)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(lambda: self._event_log)
@@ -490,6 +494,45 @@ class API:
],
)
+ async def _collect_chat_completion_with_stats(
+ self, command_id: CommandId, parse_gpt_oss: bool
+ ) -> BenchChatCompletionResponse:
+ text_parts: list[str] = []
+ model: str | None = None
+ finish_reason: FinishReason | None = None
+
+ stats: GenerationStats | None = None
+
+ async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
+ if model is None:
+ model = chunk.model
+
+ text_parts.append(chunk.text)
+ stats = chunk.stats or stats
+
+ if chunk.finish_reason is not None:
+ finish_reason = chunk.finish_reason
+
+ combined_text = "".join(text_parts)
+ assert model is not None
+
+ resp = BenchChatCompletionResponse(
+ id=command_id,
+ created=int(time.time()),
+ model=model,
+ choices=[
+ ChatCompletionChoice(
+ index=0,
+ message=ChatCompletionMessage(
+ role="assistant", content=combined_text
+ ),
+ finish_reason=finish_reason,
+ )
+ ],
+ generation_stats=stats,
+ )
+ return resp
+
async def _trigger_notify_user_to_download_model(self, model_id: str) -> None:
logger.warning(
"TODO: we should send a notification to the user to download the model"
@@ -525,6 +568,33 @@ class API:
return await self._collect_chat_completion(command.command_id, parse_gpt_oss)
+ async def bench_chat_completions(
+ self, payload: BenchChatCompletionTaskParams
+ ) -> BenchChatCompletionResponse:
+ model_meta = await resolve_model_meta(payload.model)
+ parse_gpt_oss = "gpt-oss" in model_meta.model_id.lower()
+ payload.model = model_meta.model_id
+
+ if not any(
+ instance.shard_assignments.model_id == payload.model
+ for instance in self.state.instances.values()
+ ):
+ await self._trigger_notify_user_to_download_model(payload.model)
+ raise HTTPException(
+ status_code=404, detail=f"No instance found for model {payload.model}"
+ )
+
+ payload.stream = False
+
+ command = ChatCompletion(request_params=payload)
+ await self._send(command)
+
+ response = await self._collect_chat_completion_with_stats(
+ command.command_id,
+ parse_gpt_oss,
+ )
+ return response
+
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
total_available = Memory()
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 6e3acdab..5073f12b 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
from exo.shared.types.common import CommandId
+from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
@@ -51,6 +52,10 @@ class ChatCompletionMessage(BaseModel):
function_call: dict[str, Any] | None = None
+class BenchChatCompletionMessage(ChatCompletionMessage):
+ pass
+
+
class TopLogprobItem(BaseModel):
token: str
logprob: float
@@ -113,6 +118,18 @@ class ChatCompletionResponse(BaseModel):
service_tier: str | None = None
+class GenerationStats(BaseModel):
+ prompt_tps: float
+ generation_tps: float
+ prompt_tokens: int
+ generation_tokens: int
+ peak_memory_usage: Memory
+
+
+class BenchChatCompletionResponse(ChatCompletionResponse):
+ generation_stats: GenerationStats | None = None
+
+
class ChatCompletionTaskParams(BaseModel):
model: str
frequency_penalty: float | None = None
@@ -135,6 +152,10 @@ class ChatCompletionTaskParams(BaseModel):
user: str | None = None
+class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
+ pass
+
+
class PlaceInstanceParams(BaseModel):
model_id: str
sharding: Sharding = Sharding.Pipeline
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index ac90d20c..420e18eb 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -1,5 +1,6 @@
from enum import Enum
+from exo.shared.types.api import GenerationStats
from exo.utils.pydantic_ext import TaggedModel
from .api import FinishReason
@@ -20,6 +21,7 @@ class TokenChunk(BaseChunk):
text: str
token_id: int
finish_reason: FinishReason | None = None
+ stats: GenerationStats | None = None
class ImageChunk(BaseChunk):
diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py
index 8c2d3754..1be5abca 100644
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/src/exo/shared/types/worker/runner_response.py
@@ -1,4 +1,4 @@
-from exo.shared.types.api import FinishReason
+from exo.shared.types.api import FinishReason, GenerationStats
from exo.utils.pydantic_ext import TaggedModel
@@ -15,6 +15,7 @@ class GenerationResponse(BaseRunnerResponse):
token: int
# logprobs: list[float] | None = None # too big. we can change to be top-k
finish_reason: FinishReason | None = None
+ stats: GenerationStats | None = None
class FinishedResponse(BaseRunnerResponse):
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 9d90da06..cbcc288d 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -6,7 +6,13 @@ from mlx_lm.models.cache import KVCache
from mlx_lm.tokenizer_utils import TokenizerWrapper
# from exo.engines.mlx.cache import KVPrefixCache
-from exo.shared.types.api import ChatCompletionMessage, FinishReason
+from exo.shared.types.api import (
+ BenchChatCompletionTaskParams,
+ ChatCompletionMessage,
+ FinishReason,
+ GenerationStats,
+)
+from exo.shared.types.memory import Memory
from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.runner_response import (
GenerationResponse,
@@ -72,7 +78,7 @@ def warmup_inference(
max_tokens=50,
sampler=sampler,
prompt_cache=cache,
- prefill_step_size=65536,
+ prefill_step_size=2048,
kv_group_size=KV_GROUP_SIZE,
kv_bits=KV_BITS,
):
@@ -80,17 +86,42 @@ def warmup_inference(
tokens_generated += 1
logger.info("Generated ALL warmup tokens")
+
+ # TODO: Do we want an mx_barrier?
+ # At least this version is actively incorrect, as it should use mx_barrier(group)
mx_barrier()
return tokens_generated
+def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]:
+ token_ids = [int(t) for t in token_ids]
+
+ def proc(_history: mx.array, logits: mx.array) -> mx.array:
+ for tid in token_ids:
+ logits[..., tid] = -1e9
+ return logits
+
+ return proc
+
+
+def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
+ eos: list[int] | None = getattr(tokenizer, "eos_token_ids", None)
+ if eos is None:
+ return []
+ return eos
+
+
def mlx_generate(
model: Model,
tokenizer: TokenizerWrapper,
sampler: Callable[[mx.array], mx.array],
task: ChatCompletionTaskParams,
) -> Generator[GenerationResponse]:
+ # Ensure that generation stats only contains peak memory for this generation
+ mx.reset_peak_memory()
+ is_bench: bool = isinstance(task, BenchChatCompletionTaskParams)
+
# Currently we support chat-completion tasks only.
logger.info(f"task_params: {task}")
@@ -101,6 +132,12 @@ def mlx_generate(
caches = make_kv_cache(model=model)
+ logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
+ if is_bench:
+ # Only sample length eos tokens
+ eos_ids = eos_ids_from_tokenizer(tokenizer)
+ logits_processors = [ban_token_ids(eos_ids)]
+
max_tokens = task.max_tokens or MAX_TOKENS
for out in stream_generate(
model=model,
@@ -108,26 +145,40 @@ def mlx_generate(
prompt=prompt,
max_tokens=max_tokens,
sampler=sampler,
+ logits_processors=logits_processors,
prompt_cache=caches,
- prefill_step_size=65536,
+ # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
+ prefill_step_size=2048,
kv_group_size=KV_GROUP_SIZE,
kv_bits=KV_BITS,
):
logger.info(out.text)
- if out.finish_reason is not None and out.finish_reason not in get_args(
- FinishReason
- ):
- # We don't throw here as this failure case is really not all that bad
- # Just log the error and move on
- logger.warning(
- f"Model generated unexpected finish_reason: {out.finish_reason}"
+
+ stats: GenerationStats | None = None
+ if out.finish_reason is not None:
+ stats = GenerationStats(
+ prompt_tps=float(out.prompt_tps),
+ generation_tps=float(out.generation_tps),
+ prompt_tokens=int(out.prompt_tokens),
+ generation_tokens=int(out.generation_tokens),
+ peak_memory_usage=Memory.from_gb(out.peak_memory),
)
+ if out.finish_reason not in get_args(FinishReason):
+ # We don't throw here as this failure case is really not all that bad
+ # Just log the error and move on
+ logger.warning(
+ f"Model generated unexpected finish_reason: {out.finish_reason}"
+ )
+
yield GenerationResponse(
text=out.text,
token=out.token,
finish_reason=cast(FinishReason | None, out.finish_reason),
+ stats=stats,
)
if out.finish_reason is not None:
break
+
+ # TODO: Do we want an mx_barrier?
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 8d5d3a16..96c8893f 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -397,3 +397,13 @@ def set_wired_limit_for_model(model_size: Memory):
)
mx.set_wired_limit(max_rec_size)
logger.info(f"Wired limit set to {max_rec_size}.")
+
+
+def mlx_cleanup(
+ model: Model | None, tokenizer: TokenizerWrapper | None, group: Group | None
+) -> None:
+ del model, tokenizer, group
+ mx.clear_cache()
+ import gc
+
+ gc.collect()
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 14510c1f..a3732065 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -41,6 +41,7 @@ from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_infer
from exo.worker.engines.mlx.utils_mlx import (
initialize_mlx,
load_mlx_items,
+ mlx_cleanup,
mlx_force_oom,
)
from exo.worker.runner.bootstrap import logger
@@ -179,6 +180,7 @@ def main(
text=response.text,
token_id=response.token,
finish_reason=response.finish_reason,
+ stats=response.stats,
),
)
)
@@ -190,6 +192,7 @@ def main(
case Shutdown():
current_status = RunnerShuttingDown()
logger.info("runner shutting down")
+ mlx_cleanup(model, tokenizer, group)
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
diff --git a/tmp/disable_bridge_enable_dhcp.sh b/tmp/disable_bridge_enable_dhcp.sh
deleted file mode 100755
index 8bce9333..00000000
--- a/tmp/disable_bridge_enable_dhcp.sh
+++ /dev/null
@@ -1,24 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-networksetup -listallnetworkservices | grep -q '^Thunderbolt Bridge$' \
- && echo "Disabling bridge in networksetup" \
- && networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
-
-networksetup -listallnetworkservices | grep -q '^\*Thunderbolt Bridge$' \
- && echo "Bridge disabled in networksetup"
-
-ifconfig bridge0 &>/dev/null && {
- ifconfig bridge0 | grep -q 'member' && echo "Removing bridge members in ifconfig" && {
- ifconfig bridge0 | \
- awk '/member/ {print $2}' | \
- xargs -n1 sudo ifconfig bridge0 deletem
- }
- ifconfig bridge0 | grep -q 'status: active' && sudo ifconfig bridge0 down
- ifconfig bridge0 | grep -q 'status: inactive' && echo "Bridge disabled in ifconfig"
-}
-
-for iface in $(seq 2 7); do
- sudo ipconfig set "en$iface" dhcp && echo "enabled dhcp on en$iface" || echo "failed to enable dhcp on en$iface"
-done
-
← 4963c331 Fix Discord link in README.md. Fixes #1096 (#1097)
·
back to Exo
·
local network check (#1103) ea841aca →