FastAPI endpoint — accept .pkl file upload

The FastAPI route that receives the .pkl file + form fields. Uses UploadFile = File() + Form().

main.py v3.1

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
import pickle
import io
import time
import pandas as pd
import numpy as np

app = FastAPI(title="Benchmark Engine")

@app.post("/benchmark")
async def benchmark_model(
    model_file: UploadFile = File(...),
    dataset: str = Form(...),
    framework: str = Form(default="sklearn"),
):
    """Receive a .pkl model + dataset name, run benchmark, return metrics."""
    # 1. Validate the file is a .pkl
    if not model_file.filename.endswith(".pkl"):
        raise HTTPException(400, "File must be a .pkl file")

    # 2. Read the file bytes (async — streams the file)
    file_bytes = await model_file.read()
    print(f"Received {model_file.filename}: {len(file_bytes)} bytes")

    # 3. Load the model from bytes (pickle.loads, not pickle.load)
    try:
        model = pickle.loads(file_bytes)
    except Exception as e:
        raise HTTPException(400, f"Failed to load model: {e}")

    # 4. Check it's actually a model (duck typing)
    if not hasattr(model, "predict"):
        raise HTTPException(400, "File is not a valid ML model (no predict() method)")

    # 5. Load the dataset (in real life, pick based on `dataset` param)
    # For demo: generate fake data
    from sklearn.datasets import make_regression
    X, y = make_regression(n_samples=200, n_features=5, noise=10, random_state=42)

    # 6. Run predictions + time it
    start = time.perf_counter()
    y_pred = model.predict(X)
    inference_time = time.perf_counter() - start

    # 7. Compute metrics
    from sklearn.metrics import mean_absolute_error, r2_score
    mae = float(mean_absolute_error(y, y_pred))
    r2 = float(r2_score(y, y_pred))

    # 8. Return benchmark result as JSON
    return JSONResponse({
        "status": "success",
        "model_type": type(model).__name__,
        "dataset": dataset,
        "framework": framework,
        "n_samples": len(X),
        "MAE": mae,
        "R2": r2,
        "inference_time_sec": inference_time,
        "throughput_samples_per_sec": len(X) / inference_time,
    })

# Test with curl:
# curl -X POST http://localhost:8000/benchmark \
#   -F "model_file=@model.pkl" \
#   -F "dataset=house" \
#   -F "framework=sklearn"

API simulator click Simulate

Click "Simulate" to see what this FastAPI app does.

🧠 How it works

The 8 steps of a benchmark endpoint: (1) Validate file extension — defend against non-.pkl uploads. (2) await model_file.read() — streams the file as bytes. Async, doesn't load all at once. (3) pickle.loads(bytes) — deserialize. NOTE: loads (from bytes), not load (from file). (4) Duck-type check — hasattr(model, 'predict'). Rejects non-model files. (5) Load dataset — could be from CSV, DB, or generated. Here we generate fake data. (6) Time the prediction — perf_counter before/after. (7) Compute metrics — MAE, R² (add more as needed). (8) Return JSONResponse — the benchmark result. Form(...) reads from form-encoded body. File(...) reads from multipart upload. Both are required by the HTML form's enctype='multipart/form-data'.

✏️ Try these changes

Change the validation to also accept .joblib files
→ model_file.filename.endswith(('.pkl', '.joblib'))
Add a try/except around model.predict()
→ catches models that crash on the dataset
Add cross-validation (5-fold) before returning
→ more honest metrics
Change the response to include a 'model_size_bytes' field
→ len(file_bytes)

← Previous lab

Handle ANY model type (sklearn, xgboost, custom)

Next lab →

HTML form — upload .pkl + select dataset