Concepts 🚀 Project 🧮 Labs

Error handling — when the .pkl is broken or wrong

Real users upload broken files. Your benchmark engine must handle: corrupt .pkl, wrong model type, feature mismatch, predict() crashes.

script.py v3.1

import pickle
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.datasets import make_regression

# ── A robust benchmark function with full error handling ──
def safe_benchmark(model_bytes, X_test, y_test=None):
    """Benchmark a .pkl model with comprehensive error handling.

    Returns dict with 'status' = 'success' or 'error'.
    Never raises — always returns a dict (safe for FastAPI).
    """
    result = {"status": "error"}

    # ── 1. Empty file ──
    if not model_bytes or len(model_bytes) == 0:
        result["message"] = "Empty file — no model data received"
        return result

    # ── 2. Too large (> 100 MB) ──
    if len(model_bytes) > 100 * 1024 * 1024:
        result["message"] = f"File too large: {len(model_bytes) / 1e6:.1f} MB (max 100 MB)"
        return result

    # ── 3. Not a valid pickle file ──
    try:
        model = pickle.loads(model_bytes)
    except pickle.UnpicklingError as e:
        result["message"] = f"Not a valid .pkl file: {e}"
        return result
    except Exception as e:
        result["message"] = f"Failed to load model: {type(e).__name__}: {e}"
        return result

    result["model_type"] = type(model).__name__

    # ── 4. Not an ML model (no predict method) ──
    if not hasattr(model, "predict"):
        result["message"] = f"Loaded object is {type(model).__name__}, not an ML model (no predict method)"
        return result

    # ── 5. Feature mismatch (model expects different number of features) ──
    try:
        y_pred = model.predict(X_test)
    except ValueError as e:
        if "feature" in str(e).lower() or "shape" in str(e).lower():
            n_expected = getattr(model, "n_features_in_", "?")
            result["message"] = (
                f"Feature mismatch: model expects {n_expected} features, "
                f"got {X_test.shape[1]}. Retrain with the correct dataset."
            )
        else:
            result["message"] = f"Prediction failed: {e}"
        return result
    except Exception as e:
        result["message"] = f"Prediction failed: {type(e).__name__}: {e}"
        return result

    # ── 6. Success — compute metrics if y_test provided ──
    result["status"] = "success"
    result["n_predictions"] = len(y_pred)
    result["predictions_sample"] = y_pred[:5].tolist() if hasattr(y_pred, "tolist") else list(y_pred[:5])

    if y_test is not None and len(y_test) == len(y_pred):
        try:
            from sklearn.metrics import mean_absolute_error, r2_score
            result["MAE"] = float(mean_absolute_error(y_test, y_pred))
            result["R2"] = float(r2_score(y_test, y_pred))
        except Exception as e:
            result["metrics_error"] = str(e)

    return result


# ── TEST: try breaking it in 5 different ways ──
X, y = make_regression(n_samples=100, n_features=5, noise=10, random_state=42)
good_model = LinearRegression().fit(X, y)

print('Test 1: Valid model')
r = safe_benchmark(pickle.dumps(good_model), X[:10], y[:10])
print(f'  status={r["status"]}, model={r.get("model_type")}, MAE={r.get("MAE")}')

print('\nTest 2: Empty file')
r = safe_benchmark(b"", X[:10])
print(f'  status={r["status"]}, message={r["message"]}')

print('\nTest 3: Not a pickle (random bytes)')
r = safe_benchmark(b"this is not a pickle", X[:10])
print(f'  status={r["status"]}, message={r["message"][:60]}...')

print('\nTest 4: Valid pickle but not a model (a list)')
r = safe_benchmark(pickle.dumps([1, 2, 3]), X[:10])
print(f'  status={r["status"]}, message={r["message"]}')

print('\nTest 5: Feature mismatch (model trained on 5 features, test on 3)')
wrong_X = X[:10, :3]  # only 3 features instead of 5
r = safe_benchmark(pickle.dumps(good_model), wrong_X)
print(f'  status={r["status"]}, message={r["message"]}')

print('\nAll tests passed — engine never crashed!')

Python output click Run to execute

Click "Run" to execute the code on the left.
The output will appear here.

🧠 How it works

Production benchmark engines MUST handle broken inputs gracefully. 6 failure modes to handle: (1) Empty file — user uploaded nothing. (2) Too large — DoS protection. Cap at 100 MB. (3) Not a pickle — random bytes, wrong format. (4) Not a model — valid pickle but it's a list/dict/string, not an ML model. (5) Feature mismatch — model trained on 5 features, test data has 3. predict() crashes. (6) Prediction crash — model is broken, predict() raises. KEY PRINCIPLE: never raise an exception. Always return a dict with status='error' + message. FastAPI will serialize it to JSON — the user sees a clear error instead of a 500. The 5 tests at the bottom prove every failure mode is handled.

✏️ Try these changes

Change the size limit from 100 MB to 10 MB
→ tighter limit — rejects large models
Add a check for pickle's security (only allow certain classes)
→ prevents arbitrary code execution
Add a timeout: if predict() takes > 10 seconds, kill it
→ prevents infinite loops
Test with a model that has predict_proba (classifier)
→ the function handles it — it just calls predict()

← Previous lab

Save benchmark results + compare models