diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..0d60700 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,81 @@ +name: CI + +# Runs on every PR (and on pushes to the default branches, so a merge that +# bypasses PR checks — e.g. an admin merge — still gets caught). See +# ruff.toml for why the lint job's rule set is deliberately narrow, and +# tests/conftest.py for why the test job needs a real Postgres service +# rather than a mock. + +on: + pull_request: + push: + branches: [master, main] + +permissions: + contents: read + +jobs: + lint: + name: Ruff + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install dev dependencies + run: pip install -r requirements-dev.txt + + - name: Ruff check + run: ruff check . + + test: + name: Pytest + runs-on: ubuntu-latest + + # tests/conftest.py requires a real PostgreSQL instance — it creates the + # aistudio_test database on first run and truncates tables between + # tests. A service container is the GitHub Actions equivalent of the + # `docker compose up -d postgres` step in local development. + services: + postgres: + image: postgres:15-alpine + env: + POSTGRES_USER: aistudio + POSTGRES_PASSWORD: aistudio + POSTGRES_DB: aistudio + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U aistudio" + --health-interval 5s + --health-timeout 3s + --health-retries 5 + + env: + # Matches the "pytest directly on the host machine" case documented at + # the top of tests/conftest.py — the service container is reachable on + # localhost:5432 from the job's steps. + POSTGRES_HOST: localhost + POSTGRES_PORT: 5432 + POSTGRES_USERNAME: aistudio + POSTGRES_PASSWORD: aistudio + POSTGRES_DATABASE: aistudio_test + RABBITMQ_URL: localhost + RABBITMQ_USERNAME: aistudio + RABBITMQ_PASSWORD: aistudio + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install dev dependencies + run: pip install -r requirements-dev.txt + + - name: Run tests + run: pytest tests/ -v diff --git a/Dockerfile b/Dockerfile index 398b791..7db5cf5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -4,8 +4,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends openssh-client WORKDIR /AIStudio -COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt +COPY requirements.txt requirements-dev.txt ./ +RUN pip install --no-cache-dir -r requirements-dev.txt COPY . . diff --git a/README.md b/README.md index 5cac6e7..b39cce5 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # aistudio-server +[![CI](https://github.com/corespan/aistudio-server/actions/workflows/ci.yml/badge.svg)](https://github.com/corespan/aistudio-server/actions/workflows/ci.yml) + Open-source LLM benchmarking and workload orchestration backend. SSHes into GPU nodes, runs vLLM benchmarks or launches Jupyter Lab, streams live logs via SSE, and stores results in PostgreSQL with a full leaderboard API. **Licence:** Apache-2.0 for this repository's source. The workload container diff --git a/app/main.py b/app/main.py index 38c6491..4d75185 100644 --- a/app/main.py +++ b/app/main.py @@ -3,6 +3,8 @@ from fastapi.responses import JSONResponse from fastapi.exceptions import RequestValidationError +from app.routers import system, ingest, benchmarks, results, jupyter, gpu_specs + # The app instance app = FastAPI( title="AIStudio API", @@ -60,8 +62,6 @@ async def generic_exception_handler(request: Request, exc: Exception): }, ) -from app.routers import system, ingest, benchmarks, results, jupyter, gpu_specs - # ── Routers ─────────────────────────────────────────────────────────────────── app.include_router(system.router) app.include_router(ingest.router) diff --git a/app/models/workload.py b/app/models/workload.py index 56962db..ce2b7d4 100644 --- a/app/models/workload.py +++ b/app/models/workload.py @@ -3,7 +3,7 @@ from datetime import datetime from typing import Optional -from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String, Text, func +from sqlalchemy import DateTime, Enum, Integer, String, Text, func from sqlalchemy.dialects.postgresql import JSONB, UUID from sqlalchemy.orm import Mapped, mapped_column, relationship diff --git a/app/routers/benchmarks.py b/app/routers/benchmarks.py index 146b2ef..c06b2a4 100644 --- a/app/routers/benchmarks.py +++ b/app/routers/benchmarks.py @@ -79,6 +79,10 @@ async def get_benchmark_status( return { "task_id": workload.workload_id, + # Kept alongside task_id (same value) for callers that key off the + # DB column name directly — e.g. Composer lifecycle-event tracking, + # which correlates status polls against workload_id elsewhere. + "workload_id": workload.workload_id, "state": workload.state, "error_message": workload.error_message, "updated_at": workload.updated_at, diff --git a/app/services/node_inspector.py b/app/services/node_inspector.py index b6a7fde..8f0814d 100644 --- a/app/services/node_inspector.py +++ b/app/services/node_inspector.py @@ -1,5 +1,4 @@ from app.config import settings -import json import logging from app.services.ssh_executor import SSHExecutor diff --git a/app/services/ssh_executor.py b/app/services/ssh_executor.py index 958ae2b..e5ed393 100644 --- a/app/services/ssh_executor.py +++ b/app/services/ssh_executor.py @@ -1,7 +1,7 @@ import re import uuid import paramiko -from typing import Generator, Optional +from typing import Optional from app.database import SyncSessionLocal from app.models.task_log import TaskLog diff --git a/app/utils/sse.py b/app/utils/sse.py index 0dd7eda..0efe467 100644 --- a/app/utils/sse.py +++ b/app/utils/sse.py @@ -8,14 +8,6 @@ from datetime import datetime from typing import AsyncGenerator, Optional -# Masks the last two octets of any IPv4 address in log lines so node IPs are -# never exposed to the client. e.g. 10.6.12.26 → 10.6.x.x -_IP_RE = re.compile(r'(? str: - return _IP_RE.sub(r'\1.x.x', text) - from fastapi import Request from sqlalchemy import select @@ -24,6 +16,14 @@ def _mask_ip(text: str) -> str: from app.models.task_log import TaskLog from app.models.workload import Workload +# Masks the last two octets of any IPv4 address in log lines so node IPs are +# never exposed to the client. e.g. 10.6.12.26 → 10.6.x.x +_IP_RE = re.compile(r'(? str: + return _IP_RE.sub(r'\1.x.x', text) + async def task_log_stream( workload_db_id: uuid.UUID, diff --git a/docs/api.md b/docs/api.md new file mode 100644 index 0000000..9ae90e3 --- /dev/null +++ b/docs/api.md @@ -0,0 +1,325 @@ +# API Reference + +Base URL: `http://localhost:8002` (local) or your deployment URL. + +Interactive docs (Swagger UI): `GET /docs` + +All request and response bodies are JSON. All timestamps are UTC ISO-8601. + +--- + +## Health + +### `GET /health` + +Liveness and readiness probe. Returns database connectivity status. + +**Response** +```json +{ + "status": "healthy", + "database": "ok" +} +``` +`status` is `"healthy"` when the database is reachable, `"degraded"` otherwise. + +--- + +## Benchmarks + +### `POST /api/v1/benchmarks/start` + +Start a new LLM benchmark run. The server SSHes into each node and runs the workload container asynchronously. + +**Request body** +```json +{ + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "node_ips": ["10.0.0.12"], + "config": { + "precision": "fp16", + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 128, + "gpu_count": 1, + "max_model_len": 2048, + "dataset_path": "/home/ubuntu/datasets/sharegpt.json" + } +} +``` + +| Field | Type | Required | Description | +|-------|------|----------|-------------| +| `model_name` | string | ✓ | HuggingFace repo ID, e.g. `meta-llama/Meta-Llama-3-8B-Instruct` | +| `node_ips` | string[] | ✓ | GPU node IP addresses to run on | +| `config` | object | ✓ | Benchmark parameters — see table below | + +**Config fields** + +| Field | Default | Description | +|-------|---------|-------------| +| `precision` | `"fp16"` | Model precision: `fp16`, `bf16`, `fp8` | +| `concurrency` | `4` | Number of concurrent requests | +| `input_tokens` | `512` | Prompt length in tokens | +| `output_tokens` | `256` | Generated tokens per request | +| `gpu_count` | `1` | Number of GPUs (tensor parallelism) | +| `max_model_len` | _(model default)_ | Max sequence length override | +| `dataset_path` | _(required)_ | Absolute path to a ShareGPT-format JSON file on the GPU node | + +**Response `200`** +```json +{ + "status": "queued", + "task_id": "wl-20260810-a1b2c3", + "message": "Benchmark started" +} +``` + +--- + +### `GET /api/v1/benchmarks/{task_id}/status` + +Poll the status of a running or completed benchmark. + +**Response `200`** +```json +{ + "task_id": "wl-20260810-a1b2c3", + "state": "running", + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "node_ips": ["10.0.x.x"], + "created_at": "2026-08-10T10:00:00Z" +} +``` + +`state` values: `queued` → `running` → `completed` | `failed` + +**Response `404`** — task_id not found. + +--- + +### `GET /api/v1/benchmarks/{task_id}/logs/stream` + +Server-Sent Events (SSE) stream of all log lines for a run. Supports `Last-Event-ID` for browser reconnects — only logs after the last received event are sent. + +**Headers** + +``` +Accept: text/event-stream +``` + +**Event format** +``` +id: 142 +data: [10:01:23] Starting vLLM server on port 9123... + +id: 143 +data: BENCH_RESULT:{"concurrency":4,"throughput_tok_s":1234.5,...} +``` + +The stream ends with a `data: [DONE]` sentinel when the run completes or fails. + +--- + +### `GET /api/v1/benchmarks` + +Leaderboard — paginated list of completed benchmark results. + +**Query parameters** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `model` | string | Filter by model name | +| `gpu_type` | string | Filter by GPU type | +| `server_name` | string | Filter by server name | +| `precision` | string | Filter by precision | +| `concurrency` | int | Filter by concurrency level | +| `date` | string | Filter by date (YYYY-MM-DD) | + +**Response `200`** — array of benchmark result objects. Key fields: + +```json +[ + { + "run_id": "wl-20260810-a1b2c3", + "sub_run_index": 0, + "workload_type": "LLMInference", + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "gpu_type": "T4", + "gpu_count": 1, + "precision": "fp16", + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 128, + "total_token_throughput": 1234.5, + "per_gpu_throughput_tok_s": 1234.5, + "mean_ttft_ms": 45.2, + "mean_tpot_ms": 12.1, + "mean_e2el_ms": 520.3, + "status": "completed", + "started_at": "2026-08-10T10:01:00Z", + "completed_at": "2026-08-10T10:12:00Z", + "duration_seconds": 660.0 + } +] +``` + +> Node IPs in responses have the last two octets masked (`10.6.x.x`) for privacy. + +--- + +### `GET /api/v1/benchmarks/{run_id}` + +Full detail for a single run, including all sub-runs (one per concurrency level) and the raw metrics blob. + +--- + +### `GET /api/v1/benchmarks/compare` + +Compare two runs side by side. + +**Query parameters:** `run_id_a`, `run_id_b` + +--- + +### `DELETE /api/v1/benchmarks/{run_id}` + +Delete a single benchmark run and all its sub-runs. + +### `DELETE /api/v1/benchmarks/bulk` + +Delete multiple runs by ID. **Body:** `{"run_ids": ["wl-...", "wl-..."]}` + +### `DELETE /api/v1/benchmarks/all` + +Delete all benchmark data. Irreversible. + +--- + +## System + +### `GET /api/v1/models/config` + +Returns the default vLLM configuration for a given model. Used by the UI wizard after the user selects a model. Falls back to generic defaults for unknown models. + +**Query parameter:** `model` (string, required) — HuggingFace repo ID + +**Response `200`** +```json +{ + "precision": "fp16", + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 128, + "max_model_len": 2048, + "gpu_count": 1, + "dataset_path": "", + "gated": false, + "license": "Apache-2.0", + "license_url": "https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "hf_repo": "TinyLlama/TinyLlama-1.1B-Chat-v1.0" +} +``` + +--- + +### `GET /api/v1/workload-types` + +Returns the list of supported workload types, seeded from `catalog.json`. + +**Response `200`** +```json +[ + { + "id": "uuid", + "name": "LLMInference", + "display_name": "LLM Inference (vLLM)", + "description": "Benchmark LLM inference throughput and latency using vLLM." + } +] +``` + +--- + +## Reference Dropdowns + +These endpoints return distinct values from completed benchmark results, used to populate UI filter dropdowns. + +| Endpoint | Returns | +|----------|---------| +| `GET /api/v1/models` | Distinct model names | +| `GET /api/v1/gpu-types` | Distinct GPU types | +| `GET /api/v1/servers` | Distinct server names | +| `GET /api/v1/nodes` | Distinct node IPs (masked) | +| `GET /api/v1/concurrencies` | Distinct concurrency levels | + +All accept an optional `date` query parameter to filter by run date. + +--- + +## Jupyter + +### `POST /api/v1/jupyter/launch` + +Launch a Jupyter Lab instance on a GPU node. + +**Request body** +```json +{ + "node_ip": "10.0.0.12", + "gpu_type": "T4" +} +``` + +**Response `200`** +```json +{ + "task_id": "jup-20260810-xyz", + "url": "http://10.0.0.12:7008/lab" +} +``` + +When `NGINX_ENABLED=true`, the `url` uses the public `PROXY_BASE_URL` with a path-based route. + +### `GET /api/v1/jupyter/instances` + +List all active Jupyter instances. + +### `GET /api/v1/jupyter/instances/{task_id}/status` + +Status of a specific Jupyter instance. + +### `GET /api/v1/jupyter/instances/{task_id}/logs/stream` + +SSE log stream for Jupyter launch (same format as benchmark logs). + +### `GET /api/v1/jupyter/instances/{task_id}/health` + +Checks if the Jupyter server is responding on the node. + +### `DELETE /api/v1/jupyter/instances/{task_id}` + +Stop and remove a Jupyter instance. + +--- + +## Metrics Ingestion (Legacy) + +### `POST /api/v1/metrics` + +Legacy ingestion endpoint for external runners (Kubeflow, CLI scripts). Accepts the old metrics payload format and maps it into the benchmark results table. + +```json +{ + "run_id": "my-run-001", + "timestamp": "2026-08-10T10:00:00Z", + "workload": { "name": "llama3-8b", "type": "llm" }, + "metrics": { "throughput_tok_s": 1234.5, "ttft_ms": 45.2 }, + "status": "completed", + "gpu_type": "A100", + "node_ip": "10.0.0.12", + "config": { "concurrency": 8 } +} +``` + +**Response `202 Accepted`** diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000..955a530 --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,132 @@ +# Architecture + +AIStudio Server is a Python backend that orchestrates GPU workloads on remote nodes over SSH. It does not run workloads itself — it builds shell commands, executes them remotely, and collects results. + +--- + +## Service Overview + +``` +┌─────────────────────────────────────────────────────┐ +│ Client (UI / API) │ +└──────────────────────┬──────────────────────────────┘ + │ HTTP / SSE +┌──────────────────────▼──────────────────────────────┐ +│ FastAPI (api container) │ +│ POST /benchmarks/start → enqueues Celery task │ +│ GET /benchmarks/{id}/logs/stream → SSE log tail │ +│ GET /benchmarks → leaderboard results │ +└──────────┬───────────────────────┬──────────────────┘ + │ PostgreSQL │ RabbitMQ + │ (state + results) │ (task queue) +┌──────────▼──────────────────────▼──────────────────┐ +│ Celery Worker (worker container) │ +│ execute_benchmark() — the main orchestration task │ +│ 1. SSHExecutor.run_command(manifest_builder cmd) │ +│ 2. Streams stdout/stderr → TaskLog rows │ +│ 3. Parses BENCH_RESULT:{json} → BenchmarkResult │ +└──────────────────────┬──────────────────────────────┘ + │ SSH +┌──────────────────────▼──────────────────────────────┐ +│ GPU Node (remote) │ +│ docker run ... llminference:1.0.0-nvidia │ +│ → benchmark.py starts vLLM server │ +│ → sweeps concurrency levels │ +│ → prints BENCH_RESULT:{json} to stdout │ +│ → writes /results// on the node │ +└─────────────────────────────────────────────────────┘ +``` + +--- + +## Components + +### FastAPI (`api` container) + +Entry point: `app/main.py`. Handles all HTTP traffic: + +- **Benchmark routes** (`app/routers/benchmarks.py`) — start, status, log streaming, leaderboard CRUD +- **Jupyter routes** (`app/routers/jupyter.py`) — launch, status, health, delete +- **System routes** (`app/routers/system.py`) — health check, model config, workload types +- **Results routes** (`app/routers/results.py`) — leaderboard filters, compare, distinct values +- **Ingest routes** (`app/routers/ingest.py`) — legacy `POST /api/v1/metrics` for external runners +- **GPU specs routes** (`app/routers/gpu_specs.py`) — GPU hardware metadata + +The API is stateless — all persistent state lives in PostgreSQL. + +### Celery Worker (`worker` container) + +Entry point: `app/worker.py`. Picks up tasks from RabbitMQ and runs the benchmark orchestration: + +1. Validates the request and reads workload config from PostgreSQL +2. Calls `ManifestBuilder.build_llm_benchmark_command()` to assemble the `docker run` shell command +3. Runs it on the GPU node via `SSHExecutor` +4. Tails stdout/stderr, writing each line as a `TaskLog` row in PostgreSQL (the SSE endpoint reads these) +5. Parses `BENCH_RESULT:{json}` lines from stdout into `BenchmarkResult` rows + +### ManifestBuilder (`app/services/manifest_builder.py`) + +Builds the exact shell string executed on the GPU node. Responsibilities: + +- Sources `~/.aistudio/env` on the node (loads `HF_TOKEN` for gated models — see [gpu-nodes.md](./gpu-nodes.md)) +- Constructs `docker run` with the correct GPU flags, volume mounts, env vars, and `benchmark.py` arguments +- Bind-mounts the user-supplied dataset file at the same path inside the container +- Bind-mounts `~/.cache/huggingface` so models are cached across runs + +### SSHExecutor + +Opens an SSH connection to the GPU node using the key at `SSH_KEY_PATH`, runs the command, and yields stdout/stderr lines as they arrive. Non-interactive shell — `~/.bashrc` is not sourced (see [gpu-nodes.md](./gpu-nodes.md) for why this matters). + +### PostgreSQL + +Stores all persistent state. Key tables: + +| Table | Purpose | +|-------|---------| +| `workloads` | One row per benchmark run — state machine, config, node IPs | +| `nodes` | GPU nodes used per run | +| `benchmark_results` | Parsed metric rows (one per concurrency level per run) | +| `task_logs` | Streamed log lines, read by the SSE endpoint | +| `workload_types` | Catalog of supported workload types, seeded from `catalog.json` | + +### RabbitMQ + +Message broker for Celery. The API enqueues tasks; the worker picks them up. No custom exchanges — uses Celery's default queue (`celery`). + +### Nginx (`nginx` container) + +Optional reverse proxy for Jupyter instances. When `NGINX_ENABLED=true`, each Jupyter session gets a path-based public URL (`{PROXY_BASE_URL}/jupyter/{gpu_type}/{task_id}/lab`). The worker writes an nginx location config file; `inotifywait` inside the nginx container auto-reloads on changes — no manual reload needed. + +--- + +## Workload Lifecycle + +``` +POST /benchmarks/start + → Workload row created (state: queued) + → Celery task dispatched + +Worker picks up task + → state: running + → SSHExecutor runs docker pull (if needed) + → SSHExecutor runs benchmark.py in container + → Logs stream to TaskLog rows + → BENCH_RESULT lines parsed → BenchmarkResult rows + → state: completed (or failed) + +Client polls GET /benchmarks/{id}/status +Client streams GET /benchmarks/{id}/logs/stream (SSE) +Client reads GET /benchmarks → leaderboard +``` + +--- + +## Key Design Decisions + +**No agent on GPU nodes.** The server pushes work via SSH; there is no persistent daemon on the GPU node. Any node reachable over SSH with Docker installed is a valid target. + +**BENCH_RESULT protocol.** `benchmark.py` prints `BENCH_RESULT:{json}` lines to stdout. The worker reads these from the SSH stream and inserts them as `BenchmarkResult` rows. This keeps the workload image decoupled from the server's database. + +**SSE log streaming.** Logs are stored as `TaskLog` rows, not tailed from a live SSH stream. The SSE endpoint polls the database every 500ms, supporting browser reconnects via `Last-Event-ID`. + +**User-supplied datasets.** No dataset is bundled or downloaded by the server. The operator places a ShareGPT-format JSON file on the GPU node and provides its absolute path in `dataset_path`. The path is bind-mounted into the container at the same location. diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 0000000..5ad2580 --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,121 @@ +# Configuration + +All configuration is read from environment variables. For local development, copy `.env.example` to `.env` — docker-compose loads it automatically. + +```bash +cp .env.example .env +``` + +--- + +## PostgreSQL + +| Variable | Default | Description | +|----------|---------|-------------| +| `POSTGRES_HOST` | `postgres` | Hostname of the PostgreSQL server. Use `localhost` for local dev without docker-compose. | +| `POSTGRES_PORT` | `5432` | PostgreSQL port | +| `POSTGRES_USERNAME` | `aistudio` | Database user | +| `POSTGRES_PASSWORD` | `aistudio` | Database password — change for production | +| `POSTGRES_DATABASE` | `aistudio` | Database name | + +--- + +## RabbitMQ / Celery + +| Variable | Default | Description | +|----------|---------|-------------| +| `RABBITMQ_URL` | `rabbitmq` | Hostname of the RabbitMQ broker | +| `RABBITMQ_PORT` | `5672` | AMQP port | +| `RABBITMQ_USERNAME` | `aistudio` | RabbitMQ user | +| `RABBITMQ_PASSWORD` | `aistudio` | RabbitMQ password — change for production | + +--- + +## Workload Images + +The workload container images are pulled from Google Artifact Registry (`us-docker.pkg.dev/aimlworkbench/aistudio`). This registry has public read access — no authentication is needed on GPU nodes. + +| Variable | Default | Description | +|----------|---------|-------------| +| `GCP_REGISTRY_URL` | `us-docker.pkg.dev` | Registry hostname | +| `GCP_PROJECT_ID` | `aimlworkbench` | GCP project ID | +| `GCP_REPOSITORY` | `aistudio` | Artifact Registry repository name | +| `WORKLOAD_IMAGE_TAG` | `1.0.0-nvidia` | Tag for the `llminference` image | +| `JUPYTER_IMAGE_TAG` | `1.0.0-nvidia` | Tag for the `jupyternotebook` image | + +Full image paths: +``` +us-docker.pkg.dev/aimlworkbench/aistudio/llminference:1.0.0-nvidia +us-docker.pkg.dev/aimlworkbench/aistudio/jupyternotebook:1.0.0-nvidia +``` + +--- + +## Model Storage + +Controls how the workload container accesses model weights on the GPU node. + +| Variable | Default | Description | +|----------|---------|-------------| +| `MODEL_STORAGE_MODE` | `huggingface` | Storage backend: `huggingface`, `local`, or `gcs` | +| `MODEL_LOCAL_PATH` | `/home/ubuntu/models` | Local model directory (used when `MODEL_STORAGE_MODE=local`) | +| `MODEL_GCS_BUCKET` | _(empty)_ | GCS bucket URI, e.g. `gs://my-bucket` (used when `MODEL_STORAGE_MODE=gcs`) | + +**Mode details:** + +- `huggingface` (default) — mounts `~/.cache/huggingface` from the GPU node into the container. Models are downloaded on first run and cached for subsequent runs. +- `local` — mounts `MODEL_LOCAL_PATH` from the node as `/models` inside the container. Use when weights are pre-downloaded and stored locally. +- `gcs` — passes `GCS_BUCKET` as an env var to the container. The workload image must support GCS model loading. + +--- + +## GPU Node Paths + +These paths are on the **GPU node**, not on the server. + +| Variable | Default | Description | +|----------|---------|-------------| +| `NODE_RESULTS_PATH` | `/results` | Where benchmark output is written on the GPU node. Each run creates a subdirectory `/results//` containing `benchmark_result.json`, `summary.json`, and `logs/`. | +| `NODE_JUPYTER_DATA_PATH` | `/data` | Where Jupyter notebooks are stored on the GPU node. Each session creates `/data//`. Any writable directory works — does not need shared storage. | + +--- + +## SSH + +The server connects to GPU nodes over SSH using a private key. The key must be mounted into the `api` and `worker` containers (done automatically by docker-compose when using the default `SSH_KEY_PATH`). + +| Variable | Default | Description | +|----------|---------|-------------| +| `SSH_KEY_PATH` | `~/.ssh/id_rsa` | Path to the private SSH key on the host machine | +| `SSH_DEFAULT_USER` | `ubuntu` | Default SSH username on GPU nodes | + +> **Important:** SSH commands run in a non-interactive shell. Do not rely on `~/.bashrc` for environment variables on GPU nodes — they are not sourced. Use `~/.aistudio/env` instead. See [gpu-nodes.md](./gpu-nodes.md). + +--- + +## Nginx Reverse Proxy + +When enabled, the Nginx container proxies Jupyter instances via path-based routing, hiding GPU node IPs from clients and providing public URLs. + +| Variable | Default | Description | +|----------|---------|-------------| +| `NGINX_ENABLED` | `false` | Set to `true` to enable the reverse proxy | +| `PROXY_BASE_URL` | `https://your-domain.com:8443` | Public base URL. Jupyter URLs become `{PROXY_BASE_URL}/jupyter/{gpu_type}/{task_id}/lab` | +| `NGINX_CONF_DIR` | `/etc/nginx/jupyter-locations` | Directory where the worker writes per-instance location configs. Shared via Docker volume with the nginx container. | +| `NGINX_RELOAD_CMD` | `true` | Command to reload nginx after config changes. Defaults to `true` (no-op) because the nginx container uses `inotifywait` for auto-reload. | + +--- + +## Server + +| Variable | Default | Description | +|----------|---------|-------------| +| `PORT` | `8001` | Port the FastAPI server listens on inside the container. The docker-compose default maps this to `8002` on the host. | + +--- + +## HuggingFace Token + +The HF token is **not** configured in `.env`. It belongs on the GPU node, not on the server — it is forwarded into the workload container at runtime and never stored in the database or a run manifest. + +See [gpu-nodes.md](./gpu-nodes.md#huggingface-token-for-gated-models) for setup instructions. diff --git a/docs/development.md b/docs/development.md new file mode 100644 index 0000000..78a5013 --- /dev/null +++ b/docs/development.md @@ -0,0 +1,217 @@ +# Development + +This page covers local development setup, running tests, and the contribution workflow. + +--- + +## Prerequisites + +| Tool | Version | +|------|---------| +| Python | 3.11+ | +| Docker Desktop | 24+ | +| PostgreSQL | 15+ (or run via docker-compose) | +| RabbitMQ | 3.13+ (or run via docker-compose) | + +--- + +## Local Setup + +### 1. Clone and install dependencies + +```bash +git clone https://github.com/corespan/aistudio-server.git +cd aistudio-server + +python -m venv .venv +source .venv/bin/activate # Windows: .venv\Scripts\activate +pip install -r requirements.txt +``` + +### 2. Start infrastructure services + +The easiest way is to start only PostgreSQL and RabbitMQ from docker-compose: + +```bash +docker compose up -d postgres rabbitmq +``` + +Or run a full local PostgreSQL: +```bash +docker run -d --name aistudio-pg \ + -e POSTGRES_USER=aistudio \ + -e POSTGRES_PASSWORD=aistudio \ + -e POSTGRES_DB=aistudio \ + -p 5432:5432 \ + postgres:15-alpine +``` + +### 3. Configure environment + +```bash +cp .env.example .env +``` + +For local dev, update `POSTGRES_HOST` and `RABBITMQ_URL`: +```env +POSTGRES_HOST=localhost +RABBITMQ_URL=localhost +``` + +### 4. Run migrations and seed the catalog + +```bash +alembic upgrade head +python -m app.services.catalog_seeder +``` + +### 5. Start the API and worker + +```bash +# Terminal 1 — FastAPI +uvicorn app.main:app --host 0.0.0.0 --port 8001 --reload + +# Terminal 2 — Celery worker +celery -A app.worker:celery_app worker --loglevel=info --concurrency=2 -Q celery +``` + +The API is available at **http://localhost:8001** and **http://localhost:8001/docs**. + +--- + +## Running Tests + +### AIStudio tests (79 tests) + +Tests live in `AIStudio/tests/` and use a dedicated `aistudio_test` database that is created and dropped automatically each run. + +**Run all tests:** +```bash +cd AIStudio +pytest tests/ -v --log-cli-level=INFO +``` + +**Run a specific test class:** +```bash +pytest tests/test_services.py::TestStateMachine -v +``` + +**Run by keyword:** +```bash +pytest tests/test_services.py -k "ingest or state_machine" -v +``` + +**Inside the docker-compose container:** +```bash +docker compose exec api pytest tests/ -v --log-cli-level=INFO +``` + +### Test coverage + +| # | Class | What it validates | +|---|-------|------------------| +| 1 | `TestHealthCheck` | `GET /health` — 200, shape, db=ok | +| 2 | `TestWorkloadTypes` | `GET /api/v1/workload-types` | +| 3 | `TestMetricsIngest` | `POST /api/v1/metrics` — 202, persistence, idempotent upsert | +| 4–5 | `TestBenchmarkStartAndStatus` | Start endpoint + status polling | +| 6 | `TestLogStreaming` | SSE `/logs/stream` — 404 when no task, emits all log lines | +| 7–9 | `TestResults` | Leaderboard, filters, single result, compare | +| 10 | `TestSummary` | Aggregate counts and success-rate maths | +| 11 | `TestReferenceDropdowns` | `/models`, `/gpu-types`, `/concurrencies` | +| 12 | `TestStateMachine` | Valid/invalid transitions, terminal states, audit trail | +| 13 | `TestManifestBuilder` | Shell script content — image selection, volume mounts | +| 14 | `TestDependencyInstaller` | SSH-mocked install — pull, fallback, NFS/GCS | +| 15 | `TestConfig` | Computed URL properties, storage defaults | + +--- + +## Code Style + +This project uses **ruff** for linting and formatting. + +```bash +# Check +ruff check . + +# Fix +ruff check --fix . + +# Format +ruff format . +``` + +Configuration is in `pyproject.toml` (or `ruff.toml` if present). CI will fail on linting errors. + +--- + +## Database Migrations + +Migrations are managed with **Alembic**. + +```bash +# Create a new migration after changing a model +alembic revision --autogenerate -m "describe your change" + +# Apply migrations +alembic upgrade head + +# Roll back one step +alembic downgrade -1 +``` + +Migration files live in `alembic/versions/`. Always review auto-generated migrations before committing — autogenerate can miss or misidentify changes. + +--- + +## Makefile Targets + +| Target | Description | +|--------|-------------| +| `make setup` | Full first-time setup (vendor assets, build, migrate, seed) | +| `make check-licenses` | Run licence compliance checks | +| `make check-node-env NODE=` | Verify a GPU node is correctly configured | +| `make sbom` | Generate SBOM for all workload images | + +Run `make help` to see all available targets. + +--- + +## Project Structure + +``` +aistudio-server/ +├── app/ +│ ├── main.py # FastAPI application factory +│ ├── config.py # Settings (pydantic-settings) +│ ├── worker.py # Celery app + benchmark task +│ ├── catalog.py # catalog.json loader +│ ├── database.py # SQLAlchemy async session +│ ├── models/ # SQLAlchemy ORM models +│ ├── routers/ # FastAPI route handlers +│ ├── schemas/ # Pydantic request/response schemas +│ └── services/ +│ ├── manifest_builder.py # Builds docker run shell commands +│ ├── ssh_executor.py # SSH connection + command execution +│ └── catalog_seeder.py # Seeds workload types from catalog.json +├── AIStudio/tests/ # pytest test suite +├── alembic/ # Database migrations +├── scripts/ # Utility scripts +├── sbom/ # Generated SBOMs (committed at release) +├── catalog.json # Model and workload catalog +├── docker-compose.yml # Local development stack +├── .env.example # Environment variable template +└── docs/ # This documentation +``` + +--- + +## Contribution Workflow + +1. Fork the repo and create a feature branch from `master` +2. Make your changes with tests +3. Run the full test suite: `pytest tests/ -v` +4. Run lint: `ruff check .` +5. Run licence compliance: `make check-licenses` +6. Open a pull request — describe what changed and why + +For larger changes (new workload types, schema changes, new API endpoints), open an issue first to align on the approach. See [CONTRIBUTING.md](../CONTRIBUTING.md) for full details. diff --git a/docs/gpu-nodes.md b/docs/gpu-nodes.md new file mode 100644 index 0000000..e902b0c --- /dev/null +++ b/docs/gpu-nodes.md @@ -0,0 +1,183 @@ +# GPU Node Setup + +AIStudio Server connects to GPU nodes over SSH and runs Docker commands on them. This page covers everything needed to prepare a node. + +--- + +## Requirements + +| Requirement | Notes | +|-------------|-------| +| NVIDIA GPU | Any size — T4, A100, H100, RTX series | +| Ubuntu 20.04+ or Debian 11+ | Other distros work but are untested | +| Docker 24+ with NVIDIA Container Toolkit | See below | +| SSH access from the server | Key-based, no password | +| Internet access (optional) | Only needed to download model weights | + +--- + +## 1. Install Docker and NVIDIA Container Toolkit + +```bash +# Docker +curl -fsSL https://get.docker.com | sh +sudo usermod -aG docker $USER +newgrp docker + +# NVIDIA Container Toolkit +curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \ + sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg +curl -sL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \ + sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \ + sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list +sudo apt update && sudo apt install -y nvidia-container-toolkit +sudo nvidia-ctk runtime configure --runtime=docker +sudo systemctl restart docker + +# Verify +docker run --rm --gpus all nvidia/cuda:12.1-base-ubuntu22.04 nvidia-smi +``` + +--- + +## 2. Authorize the server's SSH key + +The server connects using the private key at `SSH_KEY_PATH` (default: `~/.ssh/id_rsa`). Copy the corresponding public key to the node: + +```bash +# From your local machine / server host +ssh-copy-id -i ~/.ssh/id_rsa.pub @ + +# Verify — should run without a password prompt and show docker output +ssh @ docker ps +``` + +If you don't have an existing key pair: +```bash +ssh-keygen -t ed25519 -C "aistudio-server" +# Then copy the public key as above +``` + +--- + +## 3. Create required directories + +```bash +# Results directory — must match NODE_RESULTS_PATH in .env (default: /results) +sudo mkdir -p /results && sudo chown $USER:$USER /results + +# Jupyter data directory — must match NODE_JUPYTER_DATA_PATH in .env (default: /data) +sudo mkdir -p /data && sudo chown $USER:$USER /data + +# HuggingFace model cache (created automatically on first run, but can pre-create) +mkdir -p ~/.cache/huggingface +``` + +--- + +## 4. Place a dataset file + +LLM benchmarks require a ShareGPT-format JSON file on the GPU node. The path you provide in `dataset_path` is bind-mounted into the benchmark container at the same path. + +```bash +mkdir -p ~/datasets + +# Option A: OpenOrca (MIT licence — recommended) +# Download from HuggingFace and convert to ShareGPT format, or use a pre-converted copy + +# Option B: Dolly (CC-BY-SA-3.0) +wget -O ~/datasets/dolly.json \ + https://huggingface.co/datasets/databricks/databricks-dolly-15k/resolve/main/databricks-dolly-15k.jsonl + +# Option C: ShareGPT (contested provenance — see MODEL-LICENSES.md) +# Place a local copy at e.g. ~/datasets/sharegpt.json +``` + +When starting a benchmark, set `dataset_path` to the absolute path: +```json +"config": { + "dataset_path": "/home/ubuntu/datasets/sharegpt.json" +} +``` + +--- + +## 5. HuggingFace token (for gated models) + +Several models in `catalog.json` (all Meta Llama variants) require you to request access on HuggingFace and accept their licence. Downloads fail with HTTP 401 without an approved token — unless the weights are already in `~/.cache/huggingface`. + +**The token goes on the GPU node, not in the server's `.env`.** The Celery worker sources it from the node and forwards it into the workload container at run time. It is never stored in the database or written into a run manifest. + +```bash +# On the GPU node +mkdir -p ~/.aistudio +echo "HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxx" > ~/.aistudio/env +chmod 600 ~/.aistudio/env +``` + +> **Do not use `~/.bashrc`.** Ubuntu and Debian ship a `~/.bashrc` that starts with an early `return` for non-interactive shells: +> ```bash +> case $- in +> *i*) ;; +> *) return;; +> esac +> ``` +> Every command the server runs over SSH is non-interactive, so anything appended to `~/.bashrc` sits below that `return` and never runs. The token appears set when you log in interactively and unset for every benchmark — `~/.aistudio/env` is sourced explicitly by the worker and avoids this entirely. + +Verify the token is being picked up correctly: +```bash +# From the server machine +make check-node-env NODE= +``` + +Check which models require a token: +```bash +python3 scripts/check_model_access.py +``` + +--- + +## 6. Verify the node + +Run the full pre-flight check from the server: +```bash +make check-node-env NODE= +``` + +This checks: +- SSH connectivity +- Docker is running and the GPU is accessible +- `~/.aistudio/env` is present (if HF_TOKEN is configured) +- Required directories exist + +--- + +## Troubleshooting + +**`docker: command not found` on the node** +The SSH session doesn't inherit the same `PATH` as an interactive login. Check if Docker is installed: +```bash +ssh @ which docker +``` +If Docker is at a non-standard path, symlink it: `sudo ln -s /usr/local/bin/docker /usr/bin/docker` + +**Benchmark fails with `HF_TOKEN` not set for a gated model** +The token is unset in the SSH session. Verify `~/.aistudio/env` exists and is readable: +```bash +ssh @ 'cat ~/.aistudio/env' +``` +Do not use `~/.bashrc` — see above. + +**`permission denied` writing to `/results`** +The directory exists but the SSH user doesn't own it: +```bash +ssh @ 'sudo chown $USER:$USER /results' +``` + +**Container image pull fails** +The GCR registry is public. If the node is in an air-gapped environment, pull the image manually on a connected machine and transfer it: +```bash +docker save us-docker.pkg.dev/aimlworkbench/aistudio/llminference:1.0.0-nvidia | gzip > llminference.tar.gz +scp llminference.tar.gz @:~/ +ssh @ 'docker load < ~/llminference.tar.gz' +``` diff --git a/docs/models.md b/docs/models.md new file mode 100644 index 0000000..8ef1f7a --- /dev/null +++ b/docs/models.md @@ -0,0 +1,137 @@ +# Models + +AIStudio Server maintains a `catalog.json` at the repository root that defines which models are available to benchmark. This page explains how the catalog works and how to add new models. + +--- + +## How the Catalog Works + +`catalog.json` is the single source of truth for: +- Which models appear in the benchmark wizard UI +- The default vLLM configuration for each model (precision, concurrency, GPU count, etc.) +- Whether a model is gated (requires a HuggingFace token) +- License metadata surfaced in the UI + +On startup (`make setup` or `catalog_seeder`), the server reads `catalog.json` and seeds the `workload_types` table in PostgreSQL. Model configuration is served directly from the file at runtime via `GET /api/v1/models/config`. + +--- + +## Supported Models + +### Gated models — HuggingFace token required + +Access must be requested on HuggingFace and approved by the publisher. See [gpu-nodes.md](./gpu-nodes.md#huggingface-token-for-gated-models) for token setup. + +| Model | Licence | Min GPU memory | +|-------|---------|----------------| +| `meta-llama/Meta-Llama-3-8B-Instruct` | Llama 3 Community | 16 GB | +| `meta-llama/Meta-Llama-3-70B-Instruct` | Llama 3 Community | 80 GB (4× GPU) | +| `meta-llama/Meta-Llama-3.1-70B-Instruct` | Llama 3.1 Community | 80 GB (4× GPU) | +| `meta-llama/Llama-3.3-70B-Instruct` | Llama 3.3 Community | 80 GB (4× GPU) | + +### Ungated models — downloadable anonymously + +| Model | Licence | Min GPU memory | +|-------|---------|----------------| +| `TinyLlama/TinyLlama-1.1B-Chat-v1.0` | Apache-2.0 | 4 GB | +| `mistralai/Mistral-7B-Instruct-v0.3` | Apache-2.0 | 16 GB | +| `Qwen/Qwen2.5-7B-Instruct` | Apache-2.0 | 16 GB | +| `Qwen/Qwen2.5-32B-Instruct` | Apache-2.0 | 40 GB (2× GPU) | +| `deepseek-ai/DeepSeek-R1-Distill-Llama-70B` | MIT + Llama 3.3 (see below) | 80 GB (4× GPU) | + +**DeepSeek-R1-Distill-Llama-70B note:** Published under MIT but distilled from `Llama-3.3-70B-Instruct`. The Llama 3.3 Community Licence treats distilled outputs as derivative works, so Llama terms likely travel with it. Treat as subject to both MIT and Llama 3.3. Escalate to counsel before using in marketing material. + +See [MODEL-LICENSES.md](../MODEL-LICENSES.md) for full per-model licence details. + +--- + +## Adding a New Model + +Open `catalog.json` and add an entry to `supported_models`: + +```json +{ + "model_id": "my-model-id", + "display_name": "My Model 7B", + "hf_repo": "org/model-name-on-huggingface", + "license": "Apache-2.0", + "gated": false, + "license_url": "https://huggingface.co/org/model-name-on-huggingface", + "min_gpu_memory_gb": 16, + "default_config": { + "precision": "fp16", + "gpu_count": 1, + "max_model_len": 8192, + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 512 + } +} +``` + +| Field | Description | +|-------|-------------| +| `model_id` | Unique slug used internally. Lowercase, hyphens. | +| `display_name` | Human-readable name shown in the UI | +| `hf_repo` | Exact HuggingFace repo ID (used for the `docker run --model` arg) | +| `license` | SPDX licence identifier, e.g. `Apache-2.0`, `MIT` | +| `gated` | `true` if HuggingFace requires approval before downloading | +| `license_url` | Link to the model card or licence text | +| `min_gpu_memory_gb` | Minimum GPU VRAM in GB | +| `default_config` | Default benchmark parameters — shown pre-filled in the UI | + +After editing `catalog.json`, re-seed the database so the new model appears in the UI: + +```bash +docker compose exec api python -m app.services.catalog_seeder +``` + +Verify the model access is available (useful to run before attempting a benchmark): +```bash +python3 scripts/check_model_access.py +``` + +--- + +## Model Configuration + +When the UI wizard reaches the "Configure" step, it calls `GET /api/v1/models/config?model=`. The server returns the model's `default_config` from `catalog.json`, merged with licence metadata. The user can edit any field before starting the run. + +For models not in `catalog.json`, the server falls back to these defaults: + +```json +{ + "precision": "fp16", + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 256, + "max_model_len": 4096, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "batch_size": 32, + "dataset_path": "" +} +``` + +--- + +## Llama Community Licence — Key Obligations + +Not legal advice — read the full licence text linked from each model card. Provisions that most often surprise people: + +- **Attribution.** Derivative models must include "Llama" at the start of their name. Outputs must display "Built with Llama" prominently. +- **Acceptable Use Policy.** Incorporated by reference. Review before deploying in production. +- **Scale threshold.** Over 700M monthly active users requires a separate Meta licence. +- **Redistribution.** If you pass the weights on, the licence text and a specific attribution notice must travel with them. + +--- + +## Re-verifying Gate Status and Licences + +Model gate status and licence terms change without notice. Run before every release: + +```bash +python3 scripts/check_model_access.py +``` + +This is also wired into CI as a non-blocking weekly job — see `.github/workflows/compliance.yml`. diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 0000000..379cbdd --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,165 @@ +# Quickstart + +Get AIStudio Server running and complete your first LLM benchmark in under 10 minutes. + +## Prerequisites + +| Tool | Version | Install | +|------|---------|---------| +| Docker Desktop | 24+ | https://docs.docker.com/get-docker/ | +| `make` | any | Pre-installed on Linux/Mac. Windows: use WSL | +| A GPU node | NVIDIA, any size | Must be reachable over SSH from the server | + +> **Windows users:** Run all `make` and shell commands in a WSL terminal (`Win+R` → `wsl`), not PowerShell or CMD. + +--- + +## 1. Clone and configure + +```bash +git clone https://github.com/corespan/aistudio-server.git +cd aistudio-server +cp .env.example .env +``` + +Open `.env` and set at minimum: + +```env +SSH_KEY_PATH=~/.ssh/id_rsa # path to the private key that can SSH into your GPU node +SSH_DEFAULT_USER=ubuntu # SSH user on the GPU node +``` + +Everything else works with defaults for local development. See [configuration.md](./configuration.md) for the full reference. + +--- + +## 2. Start the server + +```bash +make setup +``` + +`make setup` does four things in one shot: +1. Vendors the demo-UI frontend assets (fonts + Chart.js, served locally — no CDN calls) +2. Builds and starts all Docker containers (API, Celery worker, PostgreSQL, RabbitMQ, Nginx) +3. Runs database migrations (`alembic upgrade head`) +4. Seeds the workload catalog (`python -m app.services.catalog_seeder`) + +**Without `make` (Linux/WSL):** +```bash +./scripts/vendor_frontend_assets.sh --if-missing +docker compose up --build -d +docker compose exec api alembic upgrade head +docker compose exec api python -m app.services.catalog_seeder +``` + +**Without `make` (Windows PowerShell):** +```powershell +copy .env.example .env +docker compose up --build -d +docker compose exec api alembic upgrade head +docker compose exec api python -m app.services.catalog_seeder +``` + +Verify the server is up: +```bash +curl http://localhost:8002/health +# → {"status":"healthy","database":"ok"} +``` + +Interactive API docs are available at **http://localhost:8002/docs**. + +--- + +## 3. Prepare your GPU node + +The server SSHes into the GPU node and runs Docker commands there. The node needs: + +**a) Docker installed and running:** +```bash +# On the GPU node +docker --version +``` + +**b) The server's SSH public key authorised:** +```bash +# On your local machine +ssh-copy-id -i ~/.ssh/id_rsa.pub @ + +# Verify +ssh @ docker ps +``` + +**c) A results directory:** +```bash +# On the GPU node — must match NODE_RESULTS_PATH in .env (default: /results) +sudo mkdir -p /results && sudo chown $USER:$USER /results +``` + +**d) A dataset file** (required for LLM benchmarks): +```bash +# On the GPU node — download a ShareGPT-format JSON dataset +# Example using OpenOrca (MIT licence): +wget -O /home/$USER/datasets/dataset.json \ + https://huggingface.co/datasets/Open-Orca/OpenOrca/resolve/main/1M-GPT4-Augmented.parquet +# Or place any ShareGPT-format JSON at a path of your choice +``` + +Verify the node is reachable: +```bash +make check-node-env NODE= +``` + +--- + +## 4. Run your first benchmark + +Using the API directly: + +```bash +curl -X POST http://localhost:8002/api/v1/benchmarks/start \ + -H "Content-Type: application/json" \ + -d '{ + "model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", + "node_ips": [""], + "config": { + "precision": "fp16", + "concurrency": 4, + "input_tokens": 512, + "output_tokens": 128, + "gpu_count": 1, + "dataset_path": "/home/ubuntu/datasets/dataset.json" + } + }' +# → {"status":"queued","task_id":"wl-20260810-a1b2c3","message":"..."} +``` + +Poll for status: +```bash +curl http://localhost:8002/api/v1/benchmarks/wl-20260810-a1b2c3/status +``` + +Stream live logs: +```bash +curl -N http://localhost:8002/api/v1/benchmarks/wl-20260810-a1b2c3/logs/stream +``` + +View results in the leaderboard: +```bash +curl http://localhost:8002/api/v1/benchmarks +``` + +--- + +## 5. Open the demo UI + +The demo UI is served at **http://localhost:3000** and connects to the local API automatically. It provides a benchmark wizard, live log streaming, and the results leaderboard. + +--- + +## Next steps + +- [gpu-nodes.md](./gpu-nodes.md) — SSH setup, HF token for gated models, troubleshooting +- [models.md](./models.md) — Adding models, gated vs ungated, catalog.json +- [configuration.md](./configuration.md) — Full environment variable reference +- [api.md](./api.md) — REST API reference diff --git a/docs/workloads.md b/docs/workloads.md new file mode 100644 index 0000000..94e3634 --- /dev/null +++ b/docs/workloads.md @@ -0,0 +1,138 @@ +# Workloads + +AIStudio Server supports two workload types. Each runs as a Docker container on the GPU node, managed by the Celery worker over SSH. + +--- + +## LLM Inference (`LLMInference`) + +**Image:** `us-docker.pkg.dev/aimlworkbench/aistudio/llminference:1.0.0-nvidia` +**Source:** [aistudio-workloads/llm-inference](https://github.com/corespan/aistudio-workloads/tree/main/llm-inference) + +Benchmarks LLM inference throughput and latency using vLLM. The container owns the entire workflow — it starts its own vLLM server, sweeps the requested concurrency levels, collects metrics, and writes results. + +### What it measures + +| Metric | Description | +|--------|-------------| +| `total_token_throughput` | Total output tokens per second across all concurrent requests | +| `per_gpu_throughput_tok_s` | `total_token_throughput ÷ gpu_count` — for cross-GPU comparison | +| `mean_ttft_ms` | Mean time to first token (ms) | +| `mean_tpot_ms` | Mean time per output token (ms) | +| `mean_e2el_ms` | Mean end-to-end latency per request (ms) | + +### How it works + +1. The Celery worker calls `ManifestBuilder.build_llm_benchmark_command()` to produce a `docker run` shell command. +2. The command is executed on the GPU node via SSH. +3. Inside the container, `benchmark.py`: + - Starts a vLLM server on port `9123` (internal to the container) + - Runs `vllm bench serve` for each concurrency level + - Prints `BENCH_RESULT:{json}` lines to stdout for each level + - Writes `benchmark_result.json` and `summary.json` to `/results//` +4. The worker parses the `BENCH_RESULT:` lines and inserts a `BenchmarkResult` row per concurrency level. + +### Volume mounts + +| Host path | Container path | Purpose | +|-----------|----------------|---------| +| `~/.cache/huggingface` | `/root/.cache/huggingface` | Model weight cache — shared across runs | +| `NODE_RESULTS_PATH/` | `/results/` | Benchmark output persistence | +| `dataset_path` | `dataset_path` (same path) | User-supplied dataset file | + +### Dataset requirement + +The benchmark requires a ShareGPT-format JSON file. No dataset is bundled or downloaded automatically. The operator provides an absolute path via `dataset_path` in the benchmark config. + +```json +"config": { + "dataset_path": "/home/ubuntu/datasets/sharegpt.json" +} +``` + +The path is bind-mounted into the container at the same location — `benchmark.py` reads it directly. + +### Configuration parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `precision` | `fp16` | Model precision | +| `concurrency` | `4` | Number of simultaneous requests | +| `input_tokens` | `512` | Prompt length | +| `output_tokens` | `256` | Generated tokens per request | +| `gpu_count` | `1` | Number of GPUs (tensor parallelism) | +| `max_model_len` | _(model default)_ | Max sequence length override | +| `dataset_path` | _(required)_ | Absolute path to dataset on GPU node | + +--- + +## Jupyter Notebook (`JupyterNotebook`) + +**Image:** `us-docker.pkg.dev/aimlworkbench/aistudio/jupyternotebook:1.0.0-nvidia` +**Source:** [aistudio-workloads/jupyter-notebook](https://github.com/corespan/aistudio-workloads/tree/main/jupyter-notebook) + +Launches a JupyterLab environment on the GPU node with pre-installed GPU profiling utilities and MLPerf microbenchmarks. + +### What's included + +- JupyterLab with GPU access (`--gpus all`) +- `jupyter-ai` — AI-assisted coding inside notebooks +- PyTorch + torchvision + CUDA +- vLLM for in-notebook inference experiments +- onnxruntime-gpu, accelerate, datasets, pycocotools (MLPerf dependencies) + +### How it works + +1. The Celery worker calls `ManifestBuilder.build_jupyter_command()`. +2. The container starts in detached mode (`-d`) — unlike benchmarks, it runs until explicitly stopped. +3. `script.sh` copies notebooks to `/data//` and starts JupyterLab from that directory. +4. The worker polls the Jupyter API endpoint until it responds (up to 5 minutes). +5. The UI shows the Jupyter URL once the health check passes. + +### URL and proxy + +When `NGINX_ENABLED=false` (default), the Jupyter URL is the GPU node's direct IP and port. This exposes the node's internal IP to the client. + +When `NGINX_ENABLED=true`, the server writes an nginx location config for the instance and the URL becomes a public path-based route: +``` +{PROXY_BASE_URL}/jupyter/{gpu_type}/{task_id}/lab +``` +The nginx container auto-reloads when new location configs are written (via `inotifywait`) — no manual restart needed. + +### Volume mounts + +| Host path | Container path | Purpose | +|-----------|----------------|---------| +| `NODE_JUPYTER_DATA_PATH` | `/data` | Notebook storage — persists across container restarts | + +### Stopping a Jupyter instance + +```bash +DELETE /api/v1/jupyter/instances/{task_id} +``` + +Or from the UI — click the delete button on the instance row. + +--- + +## Adding a New Workload Type + +New workload types require changes in both repositories: + +**aistudio-workloads** — create a new directory with: +- `Dockerfile` — base image + dependencies +- `script.sh` — entrypoint +- `requirements.txt` +- `version.py` + +**aistudio-server** — three changes: +1. Add the workload type to `catalog.json` under `workload_types` +2. Add a `build__command()` method to `ManifestBuilder` +3. Add a route and Celery task to handle the new type + +Re-seed the catalog after updating `catalog.json`: +```bash +docker compose exec api python -m app.services.catalog_seeder +``` + +See [CONTRIBUTING.md](../CONTRIBUTING.md) for the full contribution workflow. diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..32e1052 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,9 @@ +[pytest] +asyncio_mode = auto +asyncio_default_fixture_loop_scope = session +asyncio_default_test_loop_scope = session +testpaths = tests +; Safety net: without this, a hung async test (e.g. an SSE/streaming +; endpoint that never terminates) blocks the whole suite indefinitely +; instead of failing fast. +timeout = 60 diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..889d89a --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,10 @@ +# Development and testing dependencies — not needed in production containers. +# Install with: pip install -r requirements-dev.txt + +-r requirements.txt + +pytest>=8.0.0 +pytest-asyncio>=0.24.0 +pytest-timeout>=2.3.0 +anyio[trio]>=4.0.0 +ruff>=0.16.0 diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..4b5b8f5 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,41 @@ +# Ruff configuration — see .github/workflows/ci.yml for how this runs in CI. +# +# Scope is deliberately narrow: syntax errors and Pyflakes (undefined names, +# unused imports, etc.) rather than Ruff's full default rule set (which also +# pulls in pyupgrade, flake8-bugbear, flake8-datetimez and a dozen other +# plugins). Turning all of that on for a codebase that has never been linted +# produces hundreds of pre-existing findings unrelated to any given PR, which +# makes the check noise rather than a real gate. Start narrow and correct; +# widen deliberately later if the team wants stricter style enforcement. +target-version = "py311" + +[lint] +# E4/E7/E9 = pycodestyle error subset (real mistakes, not style nits). +# F = Pyflakes (undefined names, unused imports, etc). +select = ["E4", "E7", "E9", "F"] + +[lint.per-file-ignores] +# Models are imported here purely so Alembic's autogenerate can discover them +# via `Base.metadata` — the names are never referenced directly in this file. +"app/models/__init__.py" = ["F401"] + +# SQLAlchemy 2.0 `Mapped["ClassName"]` relationship annotations use forward +# references as plain strings resolved by SQLAlchemy at mapper-configuration +# time, not by Python's own name resolution — Pyflakes has no way to know +# that and flags them as undefined names. Known false positive; see +# https://github.com/astral-sh/ruff/issues/8811. +"app/models/*.py" = ["F821"] + +# alembic/env.py intentionally star-imports every model so Alembic's +# autogenerate can see them via Base.metadata; Pyflakes can't verify names +# from a star import. +"alembic/env.py" = ["F403"] + +# Historical, already-applied migrations are immutable — don't hand-edit a +# shipped revision just to satisfy a linter added long after it merged. +"alembic/versions/*.py" = ["F401"] + +# conftest.py must set POSTGRES_*/RABBITMQ_* environment variables BEFORE +# importing app.config, since pydantic-settings reads env vars at class-body +# time. That requires imports after statements — intentional, not an oversight. +"tests/conftest.py" = ["E402"] diff --git a/scripts/check_model_access.py b/scripts/check_model_access.py index e703cac..ed503f9 100755 --- a/scripts/check_model_access.py +++ b/scripts/check_model_access.py @@ -20,7 +20,6 @@ import json import pathlib import re -import sys import urllib.error import urllib.request diff --git a/scripts/generate_third_party_notices.py b/scripts/generate_third_party_notices.py index 3ed8e81..4f01d6b 100755 --- a/scripts/generate_third_party_notices.py +++ b/scripts/generate_third_party_notices.py @@ -260,9 +260,10 @@ def main() -> int: print(f"error: {OUTPUT.name} is missing. Run `make third-party`.") return 1 # Ignore the generated-on date line, which changes every run. - strip = lambda t: "\n".join( - l for l in t.splitlines() if not l.startswith("Generated:") - ) + def strip(text): + return "\n".join( + line for line in text.splitlines() if not line.startswith("Generated:") + ) if strip(OUTPUT.read_text(encoding="utf-8")) != strip(rendered): print( f"error: {OUTPUT.name} is out of date with requirements.txt. " diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..3f38626 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,165 @@ +""" +Shared pytest fixtures for aistudio-server tests. + +Requires a local PostgreSQL instance. The test database ('aistudio_test') is +created automatically on first run. + +Start PostgreSQL before running: + docker compose up -d postgres + +Run all tests: + pytest tests/ -v +""" + +import asyncio +import os + +# ── Point at the test DB BEFORE any app module is imported ───────────────── +# pydantic-settings reads env vars at class-body time; lru_cache must be +# cleared if settings was already cached by a prior import. +# ── PostgreSQL connection ───────────────────────────────────────────────────── +# Two execution contexts: +# +# make test (inside the api container — `docker compose exec api pytest tests/`) +# POSTGRES_HOST=postgres already set by docker-compose environment block +# POSTGRES_PORT NOT set by docker-compose → Settings default of 5432 is used +# +# pytest directly on the host machine +# docker-compose maps postgres to host port 5433 (5433:5432). +# Run as: POSTGRES_PORT=5433 pytest tests/ -v +# +# We do NOT setdefault POSTGRES_PORT here so that Settings.POSTGRES_PORT keeps +# its built-in default of 5432, which is correct inside the container. +# Host-side callers must pass POSTGRES_PORT=5433 explicitly. +os.environ.setdefault("POSTGRES_HOST", "localhost") +os.environ.setdefault("POSTGRES_USERNAME", "aistudio") +os.environ.setdefault("POSTGRES_PASSWORD", "aistudio") +# Force (not setdefault) — this MUST be an isolated database, never whatever +# POSTGRES_DATABASE the app's own .env configures. Inside the api container, +# docker-compose's `env_file: .env` already sets POSTGRES_DATABASE=aistudio +# (the real dev/seed database) before pytest ever runs, which would make +# setdefault() here a silent no-op. clean_tables below TRUNCATEs every table +# before every test — running that against the real database wipes seeded +# data instead of a disposable one. +os.environ["POSTGRES_DATABASE"] = "aistudio_test" +os.environ.setdefault("RABBITMQ_URL", "localhost") +os.environ.setdefault("RABBITMQ_USERNAME", "aistudio") +os.environ.setdefault("RABBITMQ_PASSWORD", "aistudio") + +# Clear the lru_cache so Settings re-reads our env vars +from app.config import get_settings +get_settings.cache_clear() + +import pytest +import pytest_asyncio +from sqlalchemy import create_engine, text +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_sessionmaker, + create_async_engine, +) +from httpx import AsyncClient, ASGITransport + +from app.main import app +from app.database import Base, get_db + +# ── Test DB connection strings ────────────────────────────────────────────── +_PG_USER = os.environ["POSTGRES_USERNAME"] +_PG_PASS = os.environ["POSTGRES_PASSWORD"] +_PG_HOST = os.environ["POSTGRES_HOST"] +_PG_PORT = os.environ["POSTGRES_PORT"] +_TEST_DB = os.environ["POSTGRES_DATABASE"] + +_ADMIN_SYNC_URL = f"postgresql://{_PG_USER}:{_PG_PASS}@{_PG_HOST}:{_PG_PORT}/postgres" +_TEST_ASYNC_URL = f"postgresql+asyncpg://{_PG_USER}:{_PG_PASS}@{_PG_HOST}:{_PG_PORT}/{_TEST_DB}" + + +# ── Create the test database once (outside pytest fixtures) ───────────────── +def _ensure_test_db_exists() -> None: + engine = create_engine(_ADMIN_SYNC_URL, isolation_level="AUTOCOMMIT") + with engine.connect() as conn: + exists = conn.execute( + text("SELECT 1 FROM pg_database WHERE datname = :name"), + {"name": _TEST_DB}, + ).scalar() + if not exists: + conn.execute(text(f'CREATE DATABASE "{_TEST_DB}"')) + engine.dispose() + + +_ensure_test_db_exists() + + +# ── Session-scoped event loop ──────────────────────────────────────────────── +@pytest.fixture(scope="session") +def event_loop(): + """Single event loop shared across the whole test session.""" + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +# ── Create / drop tables once per session ──────────────────────────────────── +@pytest_asyncio.fixture(scope="session", autouse=True) +async def create_tables(): + """DDL run once: create all ORM-mapped tables in the test DB.""" + engine = create_async_engine(_TEST_ASYNC_URL, echo=False) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + yield + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.drop_all) + await engine.dispose() + + +# ── Truncate all tables before every test ──────────────────────────────────── +@pytest_asyncio.fixture(autouse=True) +async def clean_tables(): + """Wipe all rows before each test so tests are fully isolated.""" + engine = create_async_engine(_TEST_ASYNC_URL, echo=False) + async with engine.begin() as conn: + # Reversed sorted_tables respects FK order; CASCADE handles the rest. + for table in reversed(Base.metadata.sorted_tables): + await conn.execute( + text(f'TRUNCATE TABLE "{table.name}" RESTART IDENTITY CASCADE') + ) + await engine.dispose() + yield + + +# ── Direct DB session for test setup ──────────────────────────────────────── +@pytest_asyncio.fixture +async def db_session() -> AsyncSession: + """ + Async SQLAlchemy session wired to the test DB. + + Use for inserting seed data in tests that need pre-existing rows. + Remember to call ``await db_session.commit()`` after inserts so the + HTTP client (which opens its own session) can read them. + """ + engine = create_async_engine(_TEST_ASYNC_URL, echo=False) + factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) + async with factory() as session: + yield session + await engine.dispose() + + +# ── HTTPX client wired to the test DB ──────────────────────────────────────── +@pytest_asyncio.fixture +async def http_client(db_session: AsyncSession): + """ + Async HTTPX client pointing at the FastAPI app. + + Overrides the ``get_db`` dependency so every request handler uses the + same test-DB session as the ``db_session`` fixture — inserts made in a + test are visible to the HTTP handler without a separate commit step. + """ + async def _override_get_db(): + yield db_session + + app.dependency_overrides[get_db] = _override_get_db + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as client: + yield client + app.dependency_overrides.clear() diff --git a/tests/test_services.py b/tests/test_services.py new file mode 100644 index 0000000..c7170d8 --- /dev/null +++ b/tests/test_services.py @@ -0,0 +1,1478 @@ +""" +Test suite for aistudio-server (open-source edition). + +Each class is self-contained: + - No shared mutable state between classes. + - Seed data is inserted in the test body or a local helper. + - Celery tasks are always mocked — tests never touch RabbitMQ. + +Run: + pytest tests/ -v + pytest tests/test_services.py::TestStateMachine -v + pytest tests/test_services.py -k "ingest or manifest" -v +""" + +import re +import uuid +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest + +# ── Helpers ─────────────────────────────────────────────────────────────────── + +def _make_ingest_payload( + *, + run_id: str | None = None, + model_name: str = "tinyllama-1.1b", + workload_type: str = "llm", + gpu_type: str = "t4", + node_ip: str = "10.6.12.99", + status: str = "success", + sub_run_index: int = 0, + concurrency: int = 4, + input_tokens: int = 512, + output_tokens: int = 128, + total_token_throughput: float = 5000.0, + **kwargs, +) -> dict: + """Return a valid BenchmarkIngestPayload dict.""" + return { + "run_id": run_id or ("run-%s" % uuid.uuid4().hex[:8]), + "timestamp": datetime.now(tz=timezone.utc).isoformat(), + "workload": {"name": model_name, "type": workload_type}, + "metrics": { + "total_token_throughput": total_token_throughput, + "mean_ttft_ms": 45.0, + "mean_tpot_ms": 12.0, + "mean_e2el_ms": 310.0, + **kwargs.pop("extra_metrics", {}), + }, + "status": status, + "gpu_type": gpu_type, + "node_ip": node_ip, + "config": { + "concurrency": concurrency, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "gpu_count": 1, + "precision": "fp16", + **kwargs.pop("extra_config", {}), + }, + "sub_run_index": sub_run_index, + **kwargs, + } + + +# ───────────────────────────────────────────────────────────────────────────── +# 1. Health Check +# ───────────────────────────────────────────────────────────────────────────── + +class TestHealthCheck: + """GET /health — DB connectivity probe.""" + + @pytest.mark.asyncio + async def test_health_returns_healthy(self, http_client): + r = await http_client.get("/health") + assert r.status_code == 200 + body = r.json() + assert body["status"] == "healthy" + assert body["database"] == "ok" + + @pytest.mark.asyncio + async def test_health_schema_has_expected_keys(self, http_client): + r = await http_client.get("/health") + body = r.json() + assert "status" in body + assert "database" in body + + +# ───────────────────────────────────────────────────────────────────────────── +# 2. Workload Types +# ───────────────────────────────────────────────────────────────────────────── + +class TestWorkloadTypes: + """GET /api/v1/workload-types — seeded from catalog.json.""" + + @pytest.mark.asyncio + async def test_empty_db_returns_empty_list(self, http_client): + r = await http_client.get("/api/v1/workload-types") + assert r.status_code == 200 + assert r.json() == [] + + @pytest.mark.asyncio + async def test_seeded_types_appear(self, http_client, db_session): + from app.models.workload_type import WorkloadType + db_session.add(WorkloadType( + name="LLMInference", + display_name="LLM Inference (vLLM)", + description="Benchmark LLM inference throughput.", + image_tag="1.0.0-nvidia", + )) + await db_session.commit() + + r = await http_client.get("/api/v1/workload-types") + assert r.status_code == 200 + types = r.json() + assert len(types) == 1 + assert types[0]["name"] == "LLMInference" + assert types[0]["display_name"] == "LLM Inference (vLLM)" + assert "id" in types[0] + + @pytest.mark.asyncio + async def test_multiple_types_all_returned(self, http_client, db_session): + from app.models.workload_type import WorkloadType + for name in ("LLMInference", "JupyterNotebook"): + db_session.add(WorkloadType( + name=name, + display_name=name, + description="", + image_tag="1.0.0-nvidia", + )) + await db_session.commit() + + r = await http_client.get("/api/v1/workload-types") + names = [t["name"] for t in r.json()] + assert "LLMInference" in names + assert "JupyterNotebook" in names + + +# ───────────────────────────────────────────────────────────────────────────── +# 3. Model Config +# ───────────────────────────────────────────────────────────────────────────── + +class TestModelConfig: + """GET /api/v1/models/config?model= — vLLM default config from catalog.json.""" + + @pytest.mark.asyncio + async def test_known_model_returns_config(self, http_client): + # TinyLlama is in catalog.json + r = await http_client.get( + "/api/v1/models/config", params={"model": "tinyllama-1.1b-chat"} + ) + assert r.status_code == 200 + body = r.json() + assert "precision" in body + assert "concurrency" in body + assert "tensor_parallel_size" in body + + @pytest.mark.asyncio + async def test_known_model_returns_gated_and_license_fields(self, http_client): + r = await http_client.get( + "/api/v1/models/config", params={"model": "tinyllama-1.1b-chat"} + ) + body = r.json() + assert "gated" in body + assert "license" in body + assert "hf_repo" in body + + @pytest.mark.asyncio + async def test_unknown_model_falls_back_to_defaults(self, http_client): + r = await http_client.get( + "/api/v1/models/config", params={"model": "totally-unknown-model"} + ) + assert r.status_code == 200 + body = r.json() + # Default config should still have mandatory keys + assert "precision" in body + assert body["gated"] is False + + @pytest.mark.asyncio + async def test_lookup_is_case_insensitive(self, http_client): + r_lower = await http_client.get( + "/api/v1/models/config", params={"model": "tinyllama-1.1b-chat"} + ) + r_upper = await http_client.get( + "/api/v1/models/config", params={"model": "TinyLlama-1.1B-Chat"} + ) + assert r_lower.status_code == 200 + assert r_upper.status_code == 200 + assert r_lower.json()["precision"] == r_upper.json()["precision"] + + +# ───────────────────────────────────────────────────────────────────────────── +# 4. Metrics Ingest +# ───────────────────────────────────────────────────────────────────────────── + +class TestMetricsIngest: + """POST /api/v1/metrics — legacy BenchmarkIngestPayload format.""" + + @pytest.mark.asyncio + async def test_happy_path_returns_202(self, http_client): + payload = _make_ingest_payload() + r = await http_client.post("/api/v1/metrics", json=payload) + assert r.status_code == 202 + assert r.json()["status"] == "success" + + @pytest.mark.asyncio + async def test_row_written_to_db(self, http_client, db_session): + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload(run_id="run-ingest-01") + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-ingest-01") + )).scalar_one_or_none() + assert row is not None + assert row.model_name == "tinyllama-1.1b" + assert row.gpu_type == "t4" + assert row.status == "success" + + @pytest.mark.asyncio + async def test_idempotent_upsert(self, http_client, db_session): + """Re-posting same (run_id, sub_run_index) must NOT create a second row.""" + from sqlalchemy import select, func + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload( + run_id="run-idem-01", + total_token_throughput=1000.0, + ) + await http_client.post("/api/v1/metrics", json=payload) + + # Update the throughput and re-post + payload["metrics"]["total_token_throughput"] = 9999.0 + r = await http_client.post("/api/v1/metrics", json=payload) + assert r.status_code == 202 + + count = (await db_session.execute( + select(func.count()).select_from(BenchmarkResult) + .where(BenchmarkResult.run_id == "run-idem-01") + )).scalar() + assert count == 1 + + # Updated value should win + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-idem-01") + )).scalar_one() + assert row.total_token_throughput == pytest.approx(9999.0) + + @pytest.mark.asyncio + async def test_parallelism_tensor_parallel(self, http_client, db_session): + """metrics.parallelism = {tp:4, pp:1} → hot column = 'tp4'.""" + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload( + run_id="run-tp4", + extra_metrics={"parallelism": {"tensor_parallel_size": 4, "pipeline_parallel_size": 1}}, + ) + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-tp4") + )).scalar_one() + assert row.parallelism == "tp4" + + @pytest.mark.asyncio + async def test_parallelism_pipeline_parallel(self, http_client, db_session): + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload( + run_id="run-pp4", + extra_metrics={"parallelism": {"tensor_parallel_size": 1, "pipeline_parallel_size": 4}}, + ) + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-pp4") + )).scalar_one() + assert row.parallelism == "pp4" + + @pytest.mark.asyncio + async def test_parallelism_both_tp_and_pp(self, http_client, db_session): + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload( + run_id="run-tp4pp2", + extra_metrics={"parallelism": {"tensor_parallel_size": 4, "pipeline_parallel_size": 2}}, + ) + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-tp4pp2") + )).scalar_one() + assert row.parallelism == "tp4pp2" + + @pytest.mark.asyncio + async def test_duration_computed_from_started_at(self, http_client, db_session): + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + started = datetime(2026, 1, 1, 10, 0, 0, tzinfo=timezone.utc) + completed = datetime(2026, 1, 1, 10, 1, 30, tzinfo=timezone.utc) # 90 s later + + payload = _make_ingest_payload( + run_id="run-dur-01", + extra_config={"started_at": started.isoformat()}, + ) + payload["timestamp"] = completed.isoformat() + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-dur-01") + )).scalar_one() + assert row.duration_seconds == pytest.approx(90.0) + + @pytest.mark.asyncio + async def test_gpu_type_stored_lowercase(self, http_client, db_session): + from sqlalchemy import select + from app.models.benchmark_result import BenchmarkResult + + payload = _make_ingest_payload(run_id="run-gpu-case", gpu_type="H100") + await http_client.post("/api/v1/metrics", json=payload) + + row = (await db_session.execute( + select(BenchmarkResult).where(BenchmarkResult.run_id == "run-gpu-case") + )).scalar_one() + assert row.gpu_type == "h100" + + +# ───────────────────────────────────────────────────────────────────────────── +# 5. Benchmark Start + Status +# ───────────────────────────────────────────────────────────────────────────── + +class TestBenchmarkStartAndStatus: + """POST /api/v1/benchmarks/start and GET /api/v1/benchmarks/{task_id}/status.""" + + @pytest.mark.asyncio + async def test_start_returns_task_id(self, http_client): + with patch("app.worker.start_benchmark_chain") as mock_chain: + mock_chain.delay.return_value = None + r = await http_client.post("/api/v1/benchmarks/start", json={ + "model_name": "tinyllama", + "node_ips": ["10.0.0.1"], + "config": {}, + }) + assert r.status_code == 200 + body = r.json() + assert "task_id" in body + assert body["task_id"].startswith("wl-") + assert body["status"] == "success" + + @pytest.mark.asyncio + async def test_start_creates_workload_in_created_state(self, http_client, db_session): + from sqlalchemy import select + from app.models.workload import Workload + + with patch("app.worker.start_benchmark_chain") as mock_chain: + mock_chain.delay.return_value = None + r = await http_client.post("/api/v1/benchmarks/start", json={ + "model_name": "llama3", + "node_ips": ["10.0.0.1"], + "config": {"gpu_count": 2}, + }) + task_id = r.json()["task_id"] + + workload = (await db_session.execute( + select(Workload).where(Workload.workload_id == task_id) + )).scalar_one_or_none() + assert workload is not None + assert workload.state == "CREATED" + assert workload.model_name == "llama3" + + @pytest.mark.asyncio + async def test_start_multi_node_creates_node_records(self, http_client, db_session): + from sqlalchemy import select + from app.models.node import Node + from app.models.workload import Workload + + ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3"] + with patch("app.worker.start_benchmark_chain") as mock_chain: + mock_chain.delay.return_value = None + r = await http_client.post("/api/v1/benchmarks/start", json={ + "model_name": "llama3", + "node_ips": ips, + "config": {}, + }) + task_id = r.json()["task_id"] + + workload = (await db_session.execute( + select(Workload).where(Workload.workload_id == task_id) + )).scalar_one() + + nodes = (await db_session.execute( + select(Node).where(Node.workload_id == workload.id) + )).scalars().all() + assert len(nodes) == 3 + node_ips = [n.machine_ip for n in nodes] + for ip in ips: + assert ip in node_ips + + @pytest.mark.asyncio + async def test_status_returns_created_for_new_workload(self, http_client): + with patch("app.worker.start_benchmark_chain") as mock_chain: + mock_chain.delay.return_value = None + r_start = await http_client.post("/api/v1/benchmarks/start", json={ + "model_name": "tinyllama", + "node_ips": ["10.0.0.1"], + "config": {}, + }) + task_id = r_start.json()["task_id"] + + r_status = await http_client.get(f"/api/v1/benchmarks/{task_id}/status") + assert r_status.status_code == 200 + body = r_status.json() + assert body["state"] == "CREATED" + assert body["workload_id"] == task_id + + @pytest.mark.asyncio + async def test_status_404_for_unknown_workload(self, http_client): + r = await http_client.get("/api/v1/benchmarks/wl-19990101-ffffff/status") + assert r.status_code == 404 + + @pytest.mark.asyncio + async def test_start_dispatches_celery_task(self, http_client): + with patch("app.worker.start_benchmark_chain") as mock_chain: + mock_chain.delay.return_value = None + r = await http_client.post("/api/v1/benchmarks/start", json={ + "model_name": "tinyllama", + "node_ips": ["10.0.0.1"], + "config": {}, + }) + task_id = r.json()["task_id"] + mock_chain.delay.assert_called_once_with(task_id) + + +# ───────────────────────────────────────────────────────────────────────────── +# 6. Log Streaming +# ───────────────────────────────────────────────────────────────────────────── + +class TestLogStreaming: + """GET /api/v1/benchmarks/{task_id}/logs/stream — SSE endpoint.""" + + @pytest.mark.asyncio + async def test_unknown_workload_returns_404(self, http_client): + r = await http_client.get( + "/api/v1/benchmarks/wl-doesnotexist/logs/stream" + ) + assert r.status_code == 404 + + @pytest.mark.asyncio + async def test_known_workload_returns_event_stream(self, http_client, db_session): + from app.models.workload import Workload + + # NOTE: httpx's ASGITransport has no real socket underneath it — it + # awaits the whole ASGI app call to completion and buffers every + # chunk before handle_async_request() returns a Response at all + # (see httpx._transports.asgi.ASGITransport). It cannot do + # incremental/partial reads the way a real HTTP transport can. + # That means the generator behind this endpoint MUST terminate on + # its own within the test, or client.stream() hangs forever waiting + # for the (infinite) body to finish. Use a workload that is already + # terminal so the first poll immediately emits the close event. + wl = Workload( + workload_id="wl-stream-01", + model_name="tinyllama", + workload_config={}, + state="READY", + ) + db_session.add(wl) + await db_session.commit() + + # Open the SSE stream briefly — just check Content-Type and 200 + async with http_client.stream( + "GET", "/api/v1/benchmarks/wl-stream-01/logs/stream" + ) as r: + assert r.status_code == 200 + ct = r.headers.get("content-type", "") + assert "text/event-stream" in ct + + @pytest.mark.asyncio + async def test_stream_delivers_inserted_log_lines(self, http_client, db_session): + from app.models.workload import Workload + from app.models.task import Task + from app.models.task_log import TaskLog + + wl = Workload( + workload_id="wl-stream-02", + model_name="tinyllama", + workload_config={}, + state="RUNNING", + ) + db_session.add(wl) + await db_session.flush() + + task = Task( + workload_id=wl.id, + run_name="run-stream-02", + task_config={}, + status="running", + ) + db_session.add(task) + await db_session.flush() + + log = TaskLog(task_id=task.id, line="GPU validation passed.") + db_session.add(log) + + # Mark the workload terminal so the generator's poll loop sees this + # one log line, then closes on the next pass (no new logs + terminal + # state). See the NOTE above: with ASGITransport the whole response + # body must finish within the test, or the stream never returns. + wl.state = "READY" + await db_session.commit() + + collected = [] + async with http_client.stream( + "GET", "/api/v1/benchmarks/wl-stream-02/logs/stream" + ) as r: + assert r.status_code == 200 + async for line in r.aiter_lines(): + if line.startswith("data:"): + text_val = line[len("data:"):].strip() + collected.append(text_val) + if text_val in ("READY", "FAILED"): + break + + assert any("GPU validation passed." in ln for ln in collected) + + +# ───────────────────────────────────────────────────────────────────────────── +# 7. Results Leaderboard +# ───────────────────────────────────────────────────────────────────────────── + +class TestResultsLeaderboard: + """GET /api/v1/benchmarks — leaderboard with filters and derived fields.""" + + async def _seed(self, http_client, **overrides): + """Helper: ingest one row and return its run_id.""" + payload = _make_ingest_payload(**overrides) + await http_client.post("/api/v1/metrics", json=payload) + return payload["run_id"] + + @pytest.mark.asyncio + async def test_empty_db_returns_empty_list(self, http_client): + r = await http_client.get("/api/v1/benchmarks") + assert r.status_code == 200 + assert r.json() == [] + + @pytest.mark.asyncio + async def test_ingested_rows_appear_in_list(self, http_client): + run_id = await self._seed(http_client) + r = await http_client.get("/api/v1/benchmarks") + run_ids = [row["run_id"] for row in r.json()] + assert run_id in run_ids + + @pytest.mark.asyncio + async def test_filter_by_model(self, http_client): + await self._seed(http_client, run_id="run-m-llama", model_name="llama3") + await self._seed(http_client, run_id="run-m-tiny", model_name="tinyllama") + + r = await http_client.get("/api/v1/benchmarks", params={"model": "llama3"}) + rows = r.json() + assert len(rows) == 1 + assert rows[0]["model_name"] == "llama3" + + @pytest.mark.asyncio + async def test_filter_by_gpu_type(self, http_client): + await self._seed(http_client, run_id="run-g-h100", gpu_type="h100") + await self._seed(http_client, run_id="run-g-t4", gpu_type="t4") + + r = await http_client.get("/api/v1/benchmarks", params={"gpu_type": "h100"}) + rows = r.json() + assert all(row["gpu_type"] == "h100" for row in rows) + + @pytest.mark.asyncio + async def test_filter_by_precision(self, http_client): + await self._seed(http_client, run_id="run-p-fp16", extra_config={"precision": "fp16"}) + await self._seed(http_client, run_id="run-p-bf16", extra_config={"precision": "bf16"}) + + r = await http_client.get("/api/v1/benchmarks", params={"precision": "fp16"}) + rows = r.json() + assert all(row["precision"] == "fp16" for row in rows) + + @pytest.mark.asyncio + async def test_filter_by_concurrency(self, http_client): + await self._seed(http_client, run_id="run-c-4", concurrency=4) + await self._seed(http_client, run_id="run-c-32", concurrency=32) + + r = await http_client.get("/api/v1/benchmarks", params={"concurrency": 32}) + rows = r.json() + assert all(row["concurrency"] == 32 for row in rows) + + @pytest.mark.asyncio + async def test_per_gpu_throughput_computed(self, http_client): + """per_gpu_throughput_tok_s = total_token_throughput / gpu_count.""" + payload = _make_ingest_payload( + run_id="run-pgpu", + total_token_throughput=8000.0, + extra_config={"gpu_count": 4}, + ) + await http_client.post("/api/v1/metrics", json=payload) + + r = await http_client.get("/api/v1/benchmarks", params={"run_id": "run-pgpu"}) + row = r.json()[0] + assert row["per_gpu_throughput_tok_s"] == pytest.approx(2000.0) + + @pytest.mark.asyncio + async def test_ip_masking_applied_to_node_ips(self, http_client): + """Last two octets of node IPs must be masked in responses.""" + payload = _make_ingest_payload(run_id="run-mask-ip", node_ip="10.6.12.26") + await http_client.post("/api/v1/metrics", json=payload) + + r = await http_client.get("/api/v1/benchmarks", params={"run_id": "run-mask-ip"}) + row = r.json()[0] + for ip in row["node_ips"]: + assert re.search(r'\d+\.\d+\.x\.x', ip), f"IP not masked: {ip}" + assert "12.26" not in ip + + @pytest.mark.asyncio + async def test_limit_respected(self, http_client): + for i in range(5): + await self._seed(http_client, run_id=f"run-lim-{i}") + r = await http_client.get("/api/v1/benchmarks", params={"limit": 3}) + assert len(r.json()) <= 3 + + +# ───────────────────────────────────────────────────────────────────────────── +# 8. Single Result Detail +# ───────────────────────────────────────────────────────────────────────────── + +class TestSingleResult: + """GET /api/v1/benchmarks/{run_id} — detail with sub_runs.""" + + @pytest.mark.asyncio + async def test_single_subrun_detail(self, http_client): + payload = _make_ingest_payload(run_id="run-detail-01") + await http_client.post("/api/v1/metrics", json=payload) + + r = await http_client.get("/api/v1/benchmarks/run-detail-01") + assert r.status_code == 200 + body = r.json() + assert body["run_id"] == "run-detail-01" + assert len(body["sub_runs"]) == 1 + + @pytest.mark.asyncio + async def test_multiple_subruns_grouped(self, http_client): + for i in range(3): + payload = _make_ingest_payload(run_id="run-multi-sub", sub_run_index=i) + await http_client.post("/api/v1/metrics", json=payload) + + r = await http_client.get("/api/v1/benchmarks/run-multi-sub") + assert r.status_code == 200 + body = r.json() + assert body["run_id"] == "run-multi-sub" + assert len(body["sub_runs"]) == 3 + + @pytest.mark.asyncio + async def test_detail_includes_metrics_blob(self, http_client): + payload = _make_ingest_payload(run_id="run-metrics-blob") + await http_client.post("/api/v1/metrics", json=payload) + + r = await http_client.get("/api/v1/benchmarks/run-metrics-blob") + sub = r.json()["sub_runs"][0] + assert "metrics" in sub + assert "total_token_throughput" in sub["metrics"] + + @pytest.mark.asyncio + async def test_unknown_run_returns_404(self, http_client): + r = await http_client.get("/api/v1/benchmarks/run-does-not-exist") + assert r.status_code == 404 + + +# ───────────────────────────────────────────────────────────────────────────── +# 9. Comparison +# ───────────────────────────────────────────────────────────────────────────── + +class TestComparison: + """GET /api/v1/benchmarks/compare?run_a=&run_b= — side-by-side.""" + + @pytest.mark.asyncio + async def test_compare_two_runs(self, http_client): + for run_id in ("run-cmp-a", "run-cmp-b"): + await http_client.post("/api/v1/metrics", json=_make_ingest_payload(run_id=run_id)) + + r = await http_client.get( + "/api/v1/benchmarks/compare", + params={"run_a": "run-cmp-a", "run_b": "run-cmp-b"}, + ) + assert r.status_code == 200 + body = r.json() + assert "run_a" in body + assert "run_b" in body + + @pytest.mark.asyncio + async def test_compare_missing_run_a_returns_404(self, http_client): + await http_client.post( + "/api/v1/metrics", json=_make_ingest_payload(run_id="run-cmp-only-b") + ) + r = await http_client.get( + "/api/v1/benchmarks/compare", + params={"run_a": "run-does-not-exist", "run_b": "run-cmp-only-b"}, + ) + assert r.status_code == 404 + + @pytest.mark.asyncio + async def test_compare_missing_run_b_returns_404(self, http_client): + await http_client.post( + "/api/v1/metrics", json=_make_ingest_payload(run_id="run-cmp-only-a") + ) + r = await http_client.get( + "/api/v1/benchmarks/compare", + params={"run_a": "run-cmp-only-a", "run_b": "run-does-not-exist"}, + ) + assert r.status_code == 404 + + +# ───────────────────────────────────────────────────────────────────────────── +# 10. Delete Endpoints +# ───────────────────────────────────────────────────────────────────────────── + +class TestDelete: + """DELETE single / bulk / by-filter / all.""" + + async def _seed(self, http_client, run_id, **kw): + await http_client.post("/api/v1/metrics", json=_make_ingest_payload(run_id=run_id, **kw)) + + @pytest.mark.asyncio + async def test_delete_single_removes_row(self, http_client): + await self._seed(http_client, run_id="run-del-single") + + r = await http_client.delete("/api/v1/benchmarks/run-del-single") + assert r.status_code == 200 + + r_check = await http_client.get("/api/v1/benchmarks/run-del-single") + assert r_check.status_code == 404 + + @pytest.mark.asyncio + async def test_delete_single_unknown_returns_404(self, http_client): + r = await http_client.delete("/api/v1/benchmarks/run-nonexistent") + assert r.status_code == 404 + + @pytest.mark.asyncio + async def test_delete_bulk(self, http_client): + ids = ["run-bulk-1", "run-bulk-2", "run-bulk-3"] + for rid in ids: + await self._seed(http_client, run_id=rid) + + r = await http_client.request( + "DELETE", "/api/v1/benchmarks/bulk", + json={"run_ids": ["run-bulk-1", "run-bulk-2"]}, + ) + assert r.status_code == 200 + + r_all = await http_client.get("/api/v1/benchmarks") + remaining = [row["run_id"] for row in r_all.json()] + assert "run-bulk-1" not in remaining + assert "run-bulk-2" not in remaining + assert "run-bulk-3" in remaining + + @pytest.mark.asyncio + async def test_delete_by_filter(self, http_client): + await self._seed(http_client, run_id="run-fil-h100", gpu_type="h100") + await self._seed(http_client, run_id="run-fil-t4", gpu_type="t4") + + r = await http_client.delete( + "/api/v1/benchmarks/filter", params={"gpu_type": "h100"} + ) + assert r.status_code == 200 + + r_all = await http_client.get("/api/v1/benchmarks") + ids = [row["run_id"] for row in r_all.json()] + assert "run-fil-h100" not in ids + assert "run-fil-t4" in ids + + @pytest.mark.asyncio + async def test_delete_all_requires_confirm(self, http_client): + await self._seed(http_client, run_id="run-all-1") + r = await http_client.delete("/api/v1/benchmarks/all") + assert r.status_code != 200 or r.json().get("deleted", 0) == 0 + # The row must still exist + r_all = await http_client.get("/api/v1/benchmarks") + assert len(r_all.json()) > 0 + + @pytest.mark.asyncio + async def test_delete_all_with_confirm(self, http_client): + for i in range(3): + await self._seed(http_client, run_id=f"run-del-all-{i}") + + r = await http_client.delete("/api/v1/benchmarks/all", params={"confirm": "true"}) + assert r.status_code == 200 + + r_all = await http_client.get("/api/v1/benchmarks") + assert r_all.json() == [] + + +# ───────────────────────────────────────────────────────────────────────────── +# 11. Summary +# ───────────────────────────────────────────────────────────────────────────── + +class TestSummary: + """GET /api/v1/summary — aggregate stats.""" + + @pytest.mark.asyncio + async def test_empty_db_returns_zeros(self, http_client): + r = await http_client.get("/api/v1/summary") + assert r.status_code == 200 + body = r.json() + assert body["total_runs"] == 0 + assert body["success_rate"] == 0.0 + + @pytest.mark.asyncio + async def test_all_success_gives_100_percent(self, http_client): + for i in range(3): + await http_client.post( + "/api/v1/metrics", + json=_make_ingest_payload(run_id=f"run-sum-ok-{i}", status="success"), + ) + + r = await http_client.get("/api/v1/summary") + body = r.json() + assert body["total_runs"] == 3 + assert body["successful_runs"] == 3 + assert body["success_rate"] == pytest.approx(100.0) + + @pytest.mark.asyncio + async def test_mixed_success_failure_rate(self, http_client): + for i in range(2): + await http_client.post( + "/api/v1/metrics", + json=_make_ingest_payload(run_id=f"run-mix-ok-{i}", status="success"), + ) + for i in range(2): + await http_client.post( + "/api/v1/metrics", + json=_make_ingest_payload(run_id=f"run-mix-fail-{i}", status="failed"), + ) + + r = await http_client.get("/api/v1/summary") + body = r.json() + assert body["total_runs"] == 4 + assert body["success_rate"] == pytest.approx(50.0) + + @pytest.mark.asyncio + async def test_avg_throughput_computed(self, http_client): + for i, tpt in enumerate([1000.0, 3000.0]): + await http_client.post( + "/api/v1/metrics", + json=_make_ingest_payload( + run_id=f"run-avg-{i}", total_token_throughput=tpt + ), + ) + + r = await http_client.get("/api/v1/summary") + body = r.json() + assert body["avg_throughput"] == pytest.approx(2000.0, rel=1e-2) + + +# ───────────────────────────────────────────────────────────────────────────── +# 12. Dropdown / Filter Reference Endpoints +# ───────────────────────────────────────────────────────────────────────────── + +class TestDropdowns: + """Distinct-value endpoints used to populate UI filter dropdowns.""" + + async def _seed_varied(self, http_client): + rows = [ + dict(run_id="run-dd-1", model_name="llama3", gpu_type="h100", concurrency=4, + input_tokens=512, output_tokens=128, node_ip="10.1.1.1", + extra_config={"precision": "fp16"}), + dict(run_id="run-dd-2", model_name="tinyllama", gpu_type="t4", concurrency=8, + input_tokens=256, output_tokens=64, node_ip="10.1.1.2", + extra_config={"precision": "bf16"}), + ] + for row in rows: + await http_client.post( + "/api/v1/metrics", json=_make_ingest_payload(**row) + ) + + @pytest.mark.asyncio + async def test_list_models(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/models") + assert r.status_code == 200 + names = r.json() + assert "llama3" in names + assert "tinyllama" in names + + @pytest.mark.asyncio + async def test_list_gpu_types(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/gpu-types") + assert r.status_code == 200 + types = r.json() + assert "h100" in types + assert "t4" in types + + @pytest.mark.asyncio + async def test_list_concurrencies(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/concurrencies") + assert r.status_code == 200 + vals = r.json() + assert "4" in vals or 4 in vals + assert "8" in vals or 8 in vals + + @pytest.mark.asyncio + async def test_list_precisions(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/precisions") + assert r.status_code == 200 + precisions = r.json() + assert "fp16" in precisions + assert "bf16" in precisions + + @pytest.mark.asyncio + async def test_list_input_tokens(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/input-tokens") + assert r.status_code == 200 + vals = r.json() + assert 512 in vals or "512" in vals + + @pytest.mark.asyncio + async def test_list_output_tokens(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/output-tokens") + assert r.status_code == 200 + vals = r.json() + assert 128 in vals or "128" in vals + + @pytest.mark.asyncio + async def test_list_nodes(self, http_client): + await self._seed_varied(http_client) + r = await http_client.get("/api/v1/nodes") + assert r.status_code == 200 + # node IPs are masked in responses — just check some are returned + assert len(r.json()) > 0 + + @pytest.mark.asyncio + async def test_empty_db_returns_empty_lists(self, http_client): + # NOTE: /api/v1/models is intentionally excluded here. Per its own + # docstring, it returns the union of the catalog (always present, + # regardless of run history) and historical model names — so it is + # never empty even on a freshly truncated DB. + for endpoint in ("/api/v1/gpu-types", "/api/v1/precisions"): + r = await http_client.get(endpoint) + assert r.status_code == 200 + assert r.json() == [] + + @pytest.mark.asyncio + async def test_empty_db_models_returns_only_catalog(self, http_client): + from app.catalog import _MODEL_CONFIGS + + r = await http_client.get("/api/v1/models") + assert r.status_code == 200 + assert r.json() == list(_MODEL_CONFIGS.keys()) + + +# ───────────────────────────────────────────────────────────────────────────── +# 13. GPU Specs Catalog +# ───────────────────────────────────────────────────────────────────────────── + +class TestGpuSpecs: + """GET / GET /{slug} / POST /api/v1/gpu-specs.""" + + _H100 = { + "gpu_type": "h100", + "display_name": "NVIDIA H100 NVL", + "vendor": "nvidia", + "arch": "hopper", + "vram_gb": 94, + "tdp_watts": 400, + "tier_rank": 1, + "fp16_tflops": 989.0, + "fp8_tflops": 1978.0, + } + _T4 = { + "gpu_type": "t4", + "display_name": "NVIDIA T4", + "vendor": "nvidia", + "arch": "turing", + "vram_gb": 16, + "tdp_watts": 70, + "tier_rank": 5, + "fp16_tflops": 65.0, + "fp8_tflops": None, + } + + @pytest.mark.asyncio + async def test_list_empty(self, http_client): + r = await http_client.get("/api/v1/gpu-specs") + assert r.status_code == 200 + assert r.json() == [] + + @pytest.mark.asyncio + async def test_list_ordered_by_tier_rank(self, http_client): + # Insert T4 (tier 5) first, H100 (tier 1) second + await http_client.post("/api/v1/gpu-specs", json=self._T4) + await http_client.post("/api/v1/gpu-specs", json=self._H100) + + r = await http_client.get("/api/v1/gpu-specs") + specs = r.json() + assert specs[0]["tier_rank"] < specs[1]["tier_rank"] + assert specs[0]["gpu_type"] == "h100" + + @pytest.mark.asyncio + async def test_get_by_slug_found(self, http_client): + await http_client.post("/api/v1/gpu-specs", json=self._H100) + r = await http_client.get("/api/v1/gpu-specs/h100") + assert r.status_code == 200 + body = r.json() + assert body["gpu_type"] == "h100" + assert body["vram_gb"] == 94 + + @pytest.mark.asyncio + async def test_get_by_slug_not_found(self, http_client): + r = await http_client.get("/api/v1/gpu-specs/rx-9090-xt") + assert r.status_code == 404 + + @pytest.mark.asyncio + async def test_upsert_creates_new_spec(self, http_client): + r = await http_client.post("/api/v1/gpu-specs", json=self._H100) + assert r.status_code == 201 + + r_get = await http_client.get("/api/v1/gpu-specs/h100") + assert r_get.status_code == 200 + assert r_get.json()["display_name"] == "NVIDIA H100 NVL" + + @pytest.mark.asyncio + async def test_upsert_updates_existing_spec(self, http_client): + await http_client.post("/api/v1/gpu-specs", json=self._H100) + + updated = {**self._H100, "display_name": "NVIDIA H100 NVL 94GB"} + await http_client.post("/api/v1/gpu-specs", json=updated) + + r = await http_client.get("/api/v1/gpu-specs/h100") + assert r.json()["display_name"] == "NVIDIA H100 NVL 94GB" + + @pytest.mark.asyncio + async def test_upsert_is_idempotent_no_duplicate(self, http_client): + for _ in range(3): + await http_client.post("/api/v1/gpu-specs", json=self._H100) + + r = await http_client.get("/api/v1/gpu-specs") + h100_rows = [s for s in r.json() if s["gpu_type"] == "h100"] + assert len(h100_rows) == 1 + + +# ───────────────────────────────────────────────────────────────────────────── +# 14. Catalog Seeder +# ───────────────────────────────────────────────────────────────────────────── + +class TestCatalogSeeder: + """seed_catalog() — populates workload_types from catalog.json.""" + + @pytest.mark.asyncio + async def test_seed_from_real_catalog_json(self, db_session): + from sqlalchemy import select + from app.models.workload_type import WorkloadType + from app.services.catalog_seeder import seed_catalog, CATALOG_PATH + + count = seed_catalog(CATALOG_PATH) + assert count >= 1 # at least one new type inserted + + rows = (await db_session.execute(select(WorkloadType))).scalars().all() + names = [r.name for r in rows] + assert "LLMInference" in names + + @pytest.mark.asyncio + async def test_seed_is_idempotent(self, db_session): + from sqlalchemy import select, func + from app.models.workload_type import WorkloadType + from app.services.catalog_seeder import seed_catalog, CATALOG_PATH + + seed_catalog(CATALOG_PATH) + first_count = (await db_session.execute( + select(func.count()).select_from(WorkloadType) + )).scalar() + + seed_catalog(CATALOG_PATH) # second call + second_count = (await db_session.execute( + select(func.count()).select_from(WorkloadType) + )).scalar() + + assert first_count == second_count + + @pytest.mark.asyncio + async def test_seed_updates_image_tag(self, db_session): + """Running seed again with a different tag updates the existing row.""" + import json + import tempfile + from sqlalchemy import select + from app.models.workload_type import WorkloadType + from app.services.catalog_seeder import seed_catalog + + catalog_v1 = { + "workload_types": [ + {"name": "TestType", "display_name": "Test", "image_tag": "1.0.0"} + ] + } + catalog_v2 = { + "workload_types": [ + {"name": "TestType", "display_name": "Test", "image_tag": "2.0.0"} + ] + } + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f1: + json.dump(catalog_v1, f1) + path_v1 = f1.name + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f2: + json.dump(catalog_v2, f2) + path_v2 = f2.name + + seed_catalog(path_v1) + seed_catalog(path_v2) + + # Refresh session to see committed changes from sync session + await db_session.rollback() + row = (await db_session.execute( + select(WorkloadType).where(WorkloadType.name == "TestType") + )).scalar_one() + assert row.image_tag == "2.0.0" + + @pytest.mark.asyncio + async def test_seed_missing_file_returns_zero(self): + from app.services.catalog_seeder import seed_catalog + result = seed_catalog("/tmp/this-file-does-not-exist-xyzzy.json") + assert result == 0 + + +# ───────────────────────────────────────────────────────────────────────────── +# 15. State Machine +# ───────────────────────────────────────────────────────────────────────────── + +class TestStateMachine: + """transition_workload_state() — valid/invalid transitions + audit log.""" + + def _make_sync_session(self): + """Return a synchronous SQLAlchemy session connected to the test DB.""" + from sqlalchemy import create_engine + from sqlalchemy.orm import sessionmaker + from app.config import get_sync_database_url + + # get_sync_database_url() already points at aistudio_test because + # conftest.py set POSTGRES_DATABASE=aistudio_test before any app import. + url = get_sync_database_url() + engine = create_engine(url) + Session = sessionmaker(bind=engine, expire_on_commit=False) + return Session() + + def _create_workload_sync(self, db): + """Insert a CREATED workload synchronously and return its workload_id.""" + from app.models.workload import Workload + wl_id = "wl-%s" % uuid.uuid4().hex[:8] + wl = Workload( + workload_id=wl_id, + model_name="tinyllama", + workload_config={}, + state="CREATED", + ) + db.add(wl) + db.commit() + return wl_id + + def test_valid_transition_created_to_validating(self): + from app.services.state_machine import transition_workload_state + from app.models.workload import Workload + + db = self._make_sync_session() + wl_id = self._create_workload_sync(db) + + transition_workload_state( + db=db, + workload_id=wl_id, + new_state="VALIDATING", + trigger="validate_node", + message="Starting GPU validation.", + ) + + wl = db.query(Workload).filter(Workload.workload_id == wl_id).one() + assert wl.state == "VALIDATING" + db.close() + + def test_invalid_transition_raises(self): + from app.services.state_machine import ( + transition_workload_state, + InvalidStateTransition, + ) + + db = self._make_sync_session() + wl_id = self._create_workload_sync(db) + + with pytest.raises(InvalidStateTransition): + # CREATED → RUNNING is not allowed + transition_workload_state( + db=db, + workload_id=wl_id, + new_state="RUNNING", + trigger="bad_skip", + message="skip all steps", + ) + db.close() + + def test_failed_is_terminal(self): + from app.services.state_machine import ( + transition_workload_state, + InvalidStateTransition, + ) + + db = self._make_sync_session() + wl_id = self._create_workload_sync(db) + + transition_workload_state(db, wl_id, "FAILED", "validate_node", "GPU OOM") + + with pytest.raises(InvalidStateTransition): + transition_workload_state(db, wl_id, "VALIDATING", "retry", "retrying") + db.close() + + def test_transition_writes_audit_event(self): + from app.services.state_machine import transition_workload_state + from app.models.workload import Workload + from app.models.workload_event import WorkloadEvent + + db = self._make_sync_session() + wl_id = self._create_workload_sync(db) + + transition_workload_state( + db, wl_id, "VALIDATING", "validate_node", "Checking GPU drivers." + ) + + wl = db.query(Workload).filter(Workload.workload_id == wl_id).one() + events = db.query(WorkloadEvent).filter( + WorkloadEvent.workload_id == wl.id + ).all() + assert len(events) == 1 + assert events[0].state == "VALIDATING" + assert events[0].trigger == "validate_node" + db.close() + + def test_workload_not_found_raises_value_error(self): + from app.services.state_machine import transition_workload_state + + db = self._make_sync_session() + with pytest.raises(ValueError, match="not found"): + transition_workload_state( + db, "wl-does-not-exist", "VALIDATING", "test", "msg" + ) + db.close() + + def test_failed_transition_sets_error_message(self): + from app.services.state_machine import transition_workload_state + from app.models.workload import Workload + + db = self._make_sync_session() + wl_id = self._create_workload_sync(db) + + transition_workload_state( + db, wl_id, "FAILED", "validate_node", "Driver version too old." + ) + + wl = db.query(Workload).filter(Workload.workload_id == wl_id).one() + assert wl.error_message == "Driver version too old." + db.close() + + +# ───────────────────────────────────────────────────────────────────────────── +# 16. Manifest Builder +# ───────────────────────────────────────────────────────────────────────────── + +class TestManifestBuilder: + """ManifestBuilder — shell command string generation (no SSH required).""" + + def test_llm_command_is_docker_run(self): + from app.services.manifest_builder import ManifestBuilder + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name="tinyllama/tinyllama-1.1b", + config={"gpu_count": 1, "concurrency": 4}, + run_id="run-manifest-01", + ) + assert "docker run" in cmd + + def test_llm_command_contains_model_name(self): + from app.services.manifest_builder import ManifestBuilder + model = "meta-llama/Meta-Llama-3-8B-Instruct" + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name=model, + config={"gpu_count": 1}, + run_id="run-manifest-02", + ) + assert model in cmd + + def test_llm_command_mounts_results_path(self): + from app.services.manifest_builder import ManifestBuilder + from app.config import settings + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name="tinyllama", + config={"gpu_count": 1}, + run_id="run-manifest-03", + ) + assert settings.NODE_RESULTS_PATH in cmd + assert "/results" in cmd + + def test_llm_command_respects_gpu_count_as_tp(self): + from app.services.manifest_builder import ManifestBuilder + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name="llama3", + config={"gpu_count": 4}, + run_id="run-manifest-04", + ) + # --tp 4 should appear + assert "--tp" in cmd + assert "4" in cmd + + def test_llm_command_starts_with_env_prelude(self): + from app.services.manifest_builder import ManifestBuilder + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name="tinyllama", + config={}, + run_id="run-manifest-env", + ) + # Must start with the node env sourcing command + assert cmd.startswith("if [ -f") + + def test_llm_command_uses_gcr_image(self): + from app.services.manifest_builder import ManifestBuilder + from app.config import get_workload_registry + cmd = ManifestBuilder.build_llm_benchmark_command( + model_name="tinyllama", + config={}, + run_id="run-manifest-img", + ) + registry = get_workload_registry() + assert registry in cmd + assert "llminference:" in cmd + + def test_jupyter_command_is_detached(self): + from app.services.manifest_builder import ManifestBuilder + cmd = ManifestBuilder.build_jupyter_command(run_id="jup-manifest-01") + assert "docker run" in cmd + assert " -d " in cmd or cmd.count("-d") >= 1 + + def test_jupyter_command_includes_workload_id_env(self): + from app.services.manifest_builder import ManifestBuilder + run_id = "jup-manifest-02" + cmd = ManifestBuilder.build_jupyter_command(run_id=run_id) + assert run_id in cmd + + def test_jupyter_command_custom_tag(self): + from app.services.manifest_builder import ManifestBuilder + from app.config import get_workload_registry + cmd = ManifestBuilder.build_jupyter_command(run_id="jup-manifest-03") + registry = get_workload_registry() + assert registry in cmd + assert "jupyternotebook:" in cmd + + +# ───────────────────────────────────────────────────────────────────────────── +# 17. Config +# ───────────────────────────────────────────────────────────────────────────── + +class TestConfig: + """Computed URL helpers and Settings defaults.""" + + def test_database_url_uses_asyncpg(self): + from app.config import get_database_url + url = get_database_url() + assert url.startswith("postgresql+asyncpg://") + + def test_sync_database_url_uses_plain_postgres(self): + from app.config import get_sync_database_url + url = get_sync_database_url() + assert url.startswith("postgresql://") + assert "+asyncpg" not in url + + def test_celery_broker_url_uses_amqp(self): + from app.config import get_celery_broker_url + url = get_celery_broker_url() + assert url.startswith("amqp://") + + def test_workload_registry_format(self): + from app.config import get_workload_registry, settings + registry = get_workload_registry() + assert settings.GCP_REGISTRY_URL in registry + assert settings.GCP_PROJECT_ID in registry + assert settings.GCP_REPOSITORY in registry + + def test_model_storage_mode_default(self): + from app.config import settings + assert settings.MODEL_STORAGE_MODE == "huggingface" + + def test_ssh_default_user_default(self): + # Verifies the shipped code default for open-source users who clone + # this repo without a customized .env — not the live runtime value, + # which Corespan's own deployments intentionally override via .env + # to match our node machines' actual SSH user. Reading the field + # default directly (instead of the live `settings` singleton) makes + # this test correct regardless of what .env is present on the + # machine running it. + from app.config import Settings + assert Settings.model_fields["SSH_DEFAULT_USER"].default == "ubuntu" + + def test_database_url_contains_test_db_name(self): + from app.config import get_database_url + url = get_database_url() + # Must point at the test DB (env was overridden in conftest) + assert "aistudio_test" in url + + +# ───────────────────────────────────────────────────────────────────────────── +# 18. IP Masking +# ───────────────────────────────────────────────────────────────────────────── + +class TestIPMasking: + """_mask_ip() and end-to-end masking in leaderboard responses.""" + + def test_masks_last_two_octets(self): + from app.schemas.benchmark import _mask_ip + assert _mask_ip("10.6.12.26") == "10.6.x.x" + assert _mask_ip("192.168.1.100") == "192.168.x.x" + assert _mask_ip("172.16.0.1") == "172.16.x.x" + + def test_does_not_change_already_masked(self): + from app.schemas.benchmark import _mask_ip + # Calling twice should not double-mask + masked = _mask_ip("10.6.12.26") + assert masked == "10.6.x.x" + + @pytest.mark.asyncio + async def test_masking_in_leaderboard_response(self, http_client): + """Real IP must not appear in GET /api/v1/benchmarks response.""" + payload = _make_ingest_payload(run_id="run-mask-e2e", node_ip="10.6.12.99") + await http_client.post("/api/v1/metrics", json=payload) + r = await http_client.get("/api/v1/benchmarks", params={"run_id": "run-mask-e2e"}) + row = r.json()[0] + for ip in row["node_ips"]: + assert "12.99" not in ip + assert re.match(r'\d+\.\d+\.x\.x', ip) + + @pytest.mark.asyncio + async def test_multiple_ips_all_masked(self, http_client, db_session): + """When node_ips has multiple addresses, every one must be masked.""" + from sqlalchemy.dialects.postgresql import insert as pg_insert + from app.models.benchmark_result import BenchmarkResult + + stmt = pg_insert(BenchmarkResult).values( + run_id="run-multi-ip", + sub_run_index=0, + model_name="tinyllama", + workload_type="llm", + node_ips=["10.1.2.3", "10.4.5.6"], + gpu_type="t4", + gpu_count=1, + precision="fp16", + input_tokens=512, + output_tokens=128, + concurrency=4, + status="success", + pipeline_version="unknown", + # started_at has no server_default on benchmark_results (unlike + # task_logs) — it must be supplied explicitly on every insert. + started_at=datetime.now(timezone.utc), + ).on_conflict_do_nothing() + await db_session.execute(stmt) + await db_session.commit() + + r = await http_client.get("/api/v1/benchmarks", params={"run_id": "run-multi-ip"}) + row = r.json()[0] + for ip in row["node_ips"]: + assert re.match(r'\d+\.\d+\.x\.x', ip), f"Unmasked IP: {ip}"