Store models in formats that cannot run code

How to do it

Pickle files, and formats built on them such as joblib and older PyTorch files, run code while they load. A model downloaded from a hub or a shared drive can take over the training or serving machine. Convert outside models to a weights-only format in an isolated job before they enter the registry. Where a library still needs pickle, such as some scikit-learn pipelines, load only files your own pipeline wrote. Check their hash against the registry first. In PyTorch, keep weights_only=True on every torch.load. Run ONNX models in a sandbox with memory limits, because a crafted graph can still exhaust them. Outside models also follow use models and data from trusted sources.

Examples

Python example

from safetensors.torch import load_file

def load_registered_model(version: ModelVersion) -> dict[str, torch.Tensor]:
    path = registry.download(version)
    if sha256_file(path) != version.weights_sha256:
        raise ModelHashMismatch(version.id)
    return load_file(path)

TypeScript example

import * as ort from 'onnxruntime-node';

export async function loadRegisteredModel(version: ModelVersion): Promise<ort.InferenceSession> {
  const bytes = await registry.download(version);
  if (sha256(bytes) !== version.weightsSha256) throw new ModelHashMismatch(version.id);
  return ort.InferenceSession.create(bytes);
}

How to check

A code search finds no pickle.load, joblib.load or torch.load with weights_only=False outside the one loader that checks hashes. The registry refuses an upload with a .pkl or .pt file from outside the training pipeline. A test loads a file whose hash differs from the registry, and the loader refuses it.

Sources

• PyTorch torch.load security warning
• scikit-learn model persistence and security

Related rules

• MLO-21 Keep every model in a registry with versions and aliases in Machine learning
• APS-77 Use models and data from trusted sources in Application security