🚀 Step 5 of 10
← Previous: Step 4 — Forms + Password Hashing + JWT — Login
→
File Upload + Pickle — Upload a Model
Accept a .pkl file via multipart form, save it to disk with a UUID name, store metadata in the DB. This is how users submit trained models.
🎯 Goal — what you'll achieve in this step
Train a sklearn model, pickle it, upload via the form, and see it listed in /my-models.
📖 The basic concept
Learn this first. It's the pure Python/HTML idea, stripped of all app complexity.
# The basic concept: pickle saves a Python object to a file
import pickle
from sklearn.linear_model import LinearRegression
import numpy as np
# Train a model
model = LinearRegression()
X = np.array([[1], [2], [3]])
y = np.array([2, 4, 6])
model.fit(X, y)
# Save it to a file
with open('model.pkl', 'wb') as f:
pickle.dump(model, f)
# Later (even in a different process): load it back
with open('model.pkl', 'rb') as f:
loaded = pickle.load(f)
print(loaded.predict([[5]])) # [10.]
# The .pkl file contains the trained model's weights + structure.
# You can upload it, download it, and reuse it anywhere.
🧮 Apply it to OpenBenchML
main.py
python
# Step 5: Add /upload + /my-models on top of steps 1-4
import uuid, shutil, pickle
from pathlib import Path
from fastapi import UploadFile, File, Form, Depends
from sqlalchemy.orm import Session
from sqlalchemy import Column, Integer, String, ForeignKey
# Add to database.py:
# class Model(Base):
# __tablename__ = 'models'
# id = Column(Integer, primary_key=True)
# name = Column(String)
# filename = Column(String) # UUID.pkl on disk
# owner_id = Column(Integer, ForeignKey('users.id'))
UPLOAD_DIR = Path('uploads')
UPLOAD_DIR.mkdir(exist_ok=True)
@app.get('/upload', response_class=HTMLResponse)
async def upload_form(request: Request):
return templates.TemplateResponse('upload.html', {'request': request})
@app.post('/upload')
async def upload(
request: Request,
name: str = Form(...),
file: UploadFile = File(...),
db: Session = Depends(get_db),
):
# 1. Validate the file is a .pkl
if not file.filename.endswith('.pkl'):
return {'error': 'must be a .pkl file'}
# 2. Verify it's a valid pickle by loading it
try:
data = await file.read()
model = pickle.loads(data) # load from bytes
assert hasattr(model, 'predict')
except Exception:
return {'error': 'invalid pickle or not a model'}
# 3. Save with a UUID name (never trust user's filename)
safe_name = f'{uuid.uuid4().hex}.pkl'
with open(UPLOAD_DIR / safe_name, 'wb') as f:
f.write(data)
# 4. Store metadata in DB
from database import Model, User
token = request.cookies.get('access_token')
username = decode_token(token)
user = db.query(User).filter(User.username == username).first()
m = Model(name=name, filename=safe_name, owner_id=user.id)
db.add(m); db.commit()
return RedirectResponse(url='/my-models', status_code=303)
@app.get('/my-models', response_class=HTMLResponse)
async def my_models(request: Request, db: Session = Depends(get_db)):
from database import Model, User
token = request.cookies.get('access_token')
username = decode_token(token)
if not username:
return RedirectResponse(url='/login', status_code=303)
user = db.query(User).filter(User.username == username).first()
models = db.query(Model).filter(Model.owner_id == user.id).all()
return templates.TemplateResponse('my_models.html', {
'request': request, 'models': models,
})
train_model.py
python
# Train a model and save it as a .pkl — run this to create a file to upload
import pickle
from sklearn.linear_model import LinearRegression
import numpy as np
model = LinearRegression()
X = np.array([[i] for i in range(100)])
y = np.array([2 * i + 1 + np.random.normal(0, 0.5) for i in range(100)])
model.fit(X, y)
print(f'R^2 = {model.score(X, y):.3f}')
with open('my_model.pkl', 'wb') as f:
pickle.dump(model, f)
print('saved my_model.pkl — upload this via the form')
templates/upload.html
html
{% extends 'base.html' %}
{% block title %}Upload Model{% endblock %}
{% block content %}
<h1>Upload a Model</h1>
<form method='post' action='/upload' enctype='multipart/form-data'>
<input type='text' name='name' placeholder='Model name' required>
<input type='file' name='file' accept='.pkl' required>
<button type='submit'>Upload</button>
</form>
{% endblock %}
How it works
Pickle is Python's serialization format — it turns any Python object (including a trained sklearn model) into a byte stream you can save to disk or send over a network. The upload flow: (1) The form has enctype='multipart/form-data' — WITHOUT it, the file isn't sent. (2) FastAPI's UploadFile streams the file; we read() it into bytes. (3) pickle.loads(data) verifies it's a real model with a .predict() method — never trust user uploads blindly. (4) We save with a UUID name (uuid.uuid4().hex) instead of the user's filename — this prevents path traversal attacks (../../etc/passwd) and filename collisions. (5) Metadata (name, filename, owner_id) goes in the DB; the actual .pkl goes on disk. Storing big binaries in the DB is a performance disaster.
✏️ Try yourself
Add a file-size limit: reject files > 50 MB. Read the Content-Length header from `file.size` before reading the bytes. Add a /download/{model_id} route that streams the .pkl back using FileResponse.
⚠️ Common mistakes to avoid
- Missing `enctype='multipart/form-data'` on the form — the file silently isn't sent. The `file: UploadFile` parameter will be None.
- Trusting the user's filename — `file.filename` could be `../../etc/passwd`. ALWAYS generate a UUID name.
- Loading the whole file into memory — for large files, use `shutil.copyfileobj(file.file, dest)` to stream to disk. We read() here only because we need to verify the pickle.
- Not validating the pickle — `pickle.loads` can execute arbitrary code. In production, only accept pickles from trusted users, or use a safer format like ONNX or joblib with a hash check.
🔍 Where this is used in OpenBenchML
app/routes/models.py — /upload, /my-models, /download/{id}. app/database/models.py — Model table with owner_id, filename, name, benchmark_score.
Next — Step 6
Benchmark Engine — Run a Model on a Dataset