diff --git a/scripts/serve.py b/scripts/serve.py new file mode 100644 index 0000000..a977dc0 --- /dev/null +++ b/scripts/serve.py @@ -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() \ No newline at end of file diff --git a/src/disk_checker/web/models.py b/src/disk_checker/web/models.py new file mode 100644 index 0000000..26b27bc --- /dev/null +++ b/src/disk_checker/web/models.py @@ -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 \ No newline at end of file diff --git a/src/disk_checker/web/routes.py b/src/disk_checker/web/routes.py new file mode 100644 index 0000000..c226f7b --- /dev/null +++ b/src/disk_checker/web/routes.py @@ -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()) \ No newline at end of file diff --git a/src/disk_checker/web/server.py b/src/disk_checker/web/server.py new file mode 100644 index 0000000..b5c669a --- /dev/null +++ b/src/disk_checker/web/server.py @@ -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") \ No newline at end of file diff --git a/tests/test_web.py b/tests/test_web.py new file mode 100644 index 0000000..ec8d8f8 --- /dev/null +++ b/tests/test_web.py @@ -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 \ No newline at end of file