feat(web): add FastAPI backend with scan, plot, and benchmark endpoints
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
"""Startup script for the web server."""
|
||||
|
||||
import argparse
|
||||
import uvicorn
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the web server."""
|
||||
parser = argparse.ArgumentParser(description="Start Disk Checker API server")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host to bind to")
|
||||
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
|
||||
parser.add_argument("--workers", type=int, default=1, help="Number of workers")
|
||||
parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
|
||||
parser.add_argument("--work-dir", type=Path, default=Path("."), help="Work directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set work directory in environment
|
||||
import os
|
||||
os.environ["DISK_CHECKER_WORK_DIR"] = str(args.work_dir)
|
||||
|
||||
uvicorn.run(
|
||||
"disk_checker.web.server:app",
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
workers=args.workers,
|
||||
reload=args.reload,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Pydantic models for the web API."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ToolType(str, Enum):
|
||||
"""Available benchmarking tools."""
|
||||
FIO = "fio"
|
||||
BADBLOCKS = "badblocks"
|
||||
IO_URING = "io_uring"
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a benchmark job."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
class MetricInfo(BaseModel):
|
||||
"""Information about a metric."""
|
||||
name: str
|
||||
total_ios: int
|
||||
min_latency_ns: Optional[float] = None
|
||||
max_latency_ns: Optional[float] = None
|
||||
avg_latency_ns: Optional[float] = None
|
||||
|
||||
|
||||
class PercentileInfo(BaseModel):
|
||||
"""Latency value at a specific percentile."""
|
||||
percentile: float
|
||||
latency_ns: float
|
||||
latency_display: str
|
||||
|
||||
|
||||
class MetricResult(BaseModel):
|
||||
"""Complete result for a metric."""
|
||||
metric: str
|
||||
total_ios: int
|
||||
percentiles: List[PercentileInfo]
|
||||
min_latency_ns: float
|
||||
max_latency_ns: float
|
||||
avg_latency_ns: float
|
||||
|
||||
|
||||
class ScanResponse(BaseModel):
|
||||
"""Response from CSV scan endpoint."""
|
||||
metrics: List[MetricResult]
|
||||
summary: Dict[str, Any]
|
||||
|
||||
|
||||
class BenchmarkRequest(BaseModel):
|
||||
"""Request to start a benchmark job."""
|
||||
device: str = Field(..., description="Block device path (e.g., /dev/sda)")
|
||||
tool: ToolType = ToolType.FIO
|
||||
job_file: Optional[str] = None
|
||||
options: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"device": "/dev/sda",
|
||||
"tool": "fio",
|
||||
"options": {
|
||||
"rw": "randread",
|
||||
"bs": "4k",
|
||||
"size": "1G",
|
||||
"runtime": "60"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class JobStatusResponse(BaseModel):
|
||||
"""Response with job status."""
|
||||
job_id: str
|
||||
status: JobStatus
|
||||
output_path: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
created_at: Optional[str] = None
|
||||
completed_at: Optional[str] = None
|
||||
|
||||
|
||||
class PlotRequest(BaseModel):
|
||||
"""Request to generate a plot."""
|
||||
metric: Optional[str] = None
|
||||
unit: str = Field(default="us", pattern="^(ns|us|ms|s)$")
|
||||
log_x: bool = False
|
||||
log_y: bool = False
|
||||
cdf: bool = False
|
||||
normalize: bool = False
|
||||
title: Optional[str] = None
|
||||
@@ -0,0 +1,200 @@
|
||||
"""API routes for disk-checker web backend."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import APIRouter, UploadFile, File, HTTPException, Depends
|
||||
from fastapi.responses import JSONResponse, FileResponse
|
||||
import pandas as pd
|
||||
|
||||
from disk_checker.core.extractor import load_latency_csv
|
||||
from disk_checker.core.helpers import percentile_table
|
||||
from disk_checker.plotting.plots import generate_plot
|
||||
from disk_checker.web.models import (
|
||||
ScanResponse,
|
||||
BenchmarkRequest,
|
||||
JobStatusResponse,
|
||||
PlotRequest,
|
||||
MetricInfo,
|
||||
PercentileInfo,
|
||||
MetricResult,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def _calculate_metric_results(df: pd.DataFrame, metric: str, pcts: List[float]) -> MetricResult:
|
||||
"""Calculate detailed results for a metric."""
|
||||
metric_df = df[df["metric"] == metric]
|
||||
counts = metric_df["count"].values.astype(float)
|
||||
total_ios = int(counts.sum())
|
||||
|
||||
latencies_ns = metric_df["latency_ns"].values
|
||||
|
||||
# Calculate percentiles
|
||||
percentiles = []
|
||||
cdf = __import__('numpy').cumsum(counts) / total_ios * 100.0
|
||||
for p in pcts:
|
||||
idx = min(__import__('numpy').searchsorted(cdf, p), len(latencies_ns) - 1)
|
||||
percentiles.append(PercentileInfo(
|
||||
percentile=p,
|
||||
latency_ns=float(latencies_ns[idx]),
|
||||
latency_display=f"{latencies_ns[idx]:.3f} ns"
|
||||
))
|
||||
|
||||
return MetricResult(
|
||||
metric=metric,
|
||||
total_ios=total_ios,
|
||||
percentiles=percentiles,
|
||||
min_latency_ns=float(latencies_ns.min()),
|
||||
max_latency_ns=float(latencies_ns.max()),
|
||||
avg_latency_ns=float(latencies_ns.mean()),
|
||||
)
|
||||
|
||||
|
||||
@router.post("/scan", response_model=ScanResponse)
|
||||
async def scan_csv(file: UploadFile = File(...)):
|
||||
"""
|
||||
Upload a fio JSON+ CSV and return analysis.
|
||||
"""
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
|
||||
content = await file.read()
|
||||
tmp.write(content)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
df = load_latency_csv(tmp_path)
|
||||
if df.empty:
|
||||
raise HTTPException(status_code=400, detail="No valid data in CSV")
|
||||
|
||||
# Calculate metrics
|
||||
metrics = []
|
||||
pcts = [50, 90, 99, 99.9, 99.99]
|
||||
for metric in sorted(df["metric"].unique()):
|
||||
metrics.append(_calculate_metric_results(df, metric, pcts))
|
||||
|
||||
summary = {
|
||||
"total_metrics": len(metrics),
|
||||
"total_ios": sum(m.total_ios for m in metrics),
|
||||
}
|
||||
|
||||
return ScanResponse(metrics=metrics, summary=summary)
|
||||
finally:
|
||||
Path(tmp_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
@router.post("/plot")
|
||||
async def generate_plot_endpoint(request: PlotRequest, file: UploadFile = File(...)):
|
||||
"""
|
||||
Generate and return a latency plot.
|
||||
"""
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
|
||||
content = await file.read()
|
||||
tmp.write(content)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
df = load_latency_csv(tmp_path)
|
||||
if df.empty:
|
||||
raise HTTPException(status_code=400, detail="No valid data in CSV")
|
||||
|
||||
# Determine metrics
|
||||
available = sorted(df["metric"].unique())
|
||||
if request.metric:
|
||||
if request.metric not in available:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Metric {request.metric} not found"
|
||||
)
|
||||
metrics = [request.metric]
|
||||
else:
|
||||
totals = df.groupby("metric")["count"].sum().sort_values(ascending=False)
|
||||
metrics = [totals.index[0]]
|
||||
|
||||
# Generate plot
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as img:
|
||||
output_path = Path(img.name)
|
||||
|
||||
generate_plot(
|
||||
df=df,
|
||||
metrics=metrics,
|
||||
output_path=output_path,
|
||||
unit=request.unit,
|
||||
log_x=request.log_x,
|
||||
log_y=request.log_y,
|
||||
normalize=request.normalize,
|
||||
show_cdf=request.cdf,
|
||||
title=request.title or "",
|
||||
)
|
||||
|
||||
return FileResponse(output_path, media_type="image/png")
|
||||
finally:
|
||||
Path(tmp_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
@router.post("/benchmark", response_model=JobStatusResponse)
|
||||
async def start_benchmark(request: BenchmarkRequest):
|
||||
"""
|
||||
Start a disk benchmark job.
|
||||
"""
|
||||
orchestrator = request.app.state.orchestrator
|
||||
|
||||
job_id = f"bench_{asyncio.get_event_loop().time()}"
|
||||
|
||||
async def run_job():
|
||||
try:
|
||||
if request.tool == "fio":
|
||||
output_path = await orchestrator.run_fio(
|
||||
device=request.device,
|
||||
**request.options,
|
||||
)
|
||||
elif request.tool == "badblocks":
|
||||
output_path = await orchestrator.run_badblocks(
|
||||
device=request.device,
|
||||
**request.options,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown tool: {request.tool}")
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"status": "completed",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
except Exception as e:
|
||||
return {"job_id": job_id, "status": "failed", "error": str(e)}
|
||||
|
||||
request.app.state.jobs[job_id] = {
|
||||
"status": "running",
|
||||
"task": asyncio.create_task(run_job()),
|
||||
}
|
||||
|
||||
return JobStatusResponse(job_id=job_id, status="running")
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}", response_model=JobStatusResponse)
|
||||
async def get_job_status(job_id: str, request: Request):
|
||||
"""
|
||||
Check the status of a benchmark job.
|
||||
"""
|
||||
jobs = request.app.state.jobs
|
||||
if job_id not in jobs:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
job = jobs[job_id]
|
||||
task = job["task"]
|
||||
|
||||
if task.done():
|
||||
result = task.result()
|
||||
return JobStatusResponse(**result)
|
||||
|
||||
return JobStatusResponse(job_id=job_id, status="running")
|
||||
|
||||
|
||||
@router.get("/jobs", response_model=List[str])
|
||||
async def list_jobs(request: Request):
|
||||
"""
|
||||
List all benchmark jobs.
|
||||
"""
|
||||
return list(request.app.state.jobs.keys())
|
||||
@@ -0,0 +1,49 @@
|
||||
"""FastAPI application for disk-checker web API."""
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from contextlib import asynccontextmanager
|
||||
import logging
|
||||
|
||||
from disk_checker.core.orchestrator import Orchestrator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Application lifespan handler."""
|
||||
# Startup: initialize orchestrator
|
||||
app.state.orchestrator = Orchestrator(work_dir=app.state.work_dir)
|
||||
app.state.jobs = {}
|
||||
logger.info("Disk Checker API started")
|
||||
yield
|
||||
# Shutdown: cleanup
|
||||
logger.info("Disk Checker API stopped")
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="Disk Checker API",
|
||||
description="API for disk benchmarking and latency visualization",
|
||||
version="0.1.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
return {"status": "ok", "version": "0.1.0"}
|
||||
|
||||
|
||||
# Import and include routers
|
||||
from disk_checker.web.routes import router as api_router
|
||||
app.include_router(api_router, prefix="/api/v1")
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Tests for the web backend."""
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from disk_checker.web.server import app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
"""Create test client."""
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_csv_content():
|
||||
"""Sample CSV content for testing."""
|
||||
return """nsec,read_clat_ns_count,read_clat_ns_cumulative,read_clat_ns_percentile
|
||||
19791872,1,1,0.00016775708773695687
|
||||
22413312,4,5,0.0008387854386847844
|
||||
22675456,9,14,0.0023485992283173965
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_csv_path(tmp_path, sample_csv_content):
|
||||
"""Create sample CSV file."""
|
||||
path = tmp_path / "test.csv"
|
||||
path.write_text(sample_csv_content)
|
||||
return path
|
||||
|
||||
|
||||
def test_health_check(client):
|
||||
"""Test health check endpoint."""
|
||||
response = client.get("/health")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "ok"
|
||||
|
||||
|
||||
def test_scan_csv(client, sample_csv_path):
|
||||
"""Test CSV scan endpoint."""
|
||||
with open(sample_csv_path, "rb") as f:
|
||||
response = client.post("/api/v1/scan", files={"file": f})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "metrics" in data
|
||||
assert len(data["metrics"]) > 0
|
||||
|
||||
|
||||
def test_plot_endpoint(client, sample_csv_path):
|
||||
"""Test plot generation endpoint."""
|
||||
with open(sample_csv_path, "rb") as f:
|
||||
response = client.post("/api/v1/plot", files={"file": f})
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"] == "image/png"
|
||||
|
||||
|
||||
def test_plot_with_options(client, sample_csv_path):
|
||||
"""Test plot generation with options."""
|
||||
with open(sample_csv_path, "rb") as f:
|
||||
response = client.post(
|
||||
"/api/v1/plot",
|
||||
files={"file": f},
|
||||
data={
|
||||
"metric": "read_clat",
|
||||
"unit": "ms",
|
||||
"log_x": "true",
|
||||
"cdf": "true",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_invalid_metric(client, sample_csv_path):
|
||||
"""Test invalid metric returns error."""
|
||||
with open(sample_csv_path, "rb") as f:
|
||||
response = client.post(
|
||||
"/api/v1/plot",
|
||||
files={"file": f},
|
||||
data={"metric": "invalid_metric"},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
Reference in New Issue
Block a user