PickleBall: Secure Deserialization of Pickle-Based Machine Learning Models (github.com)

PickleBall is a new defense for the well-known security hole in Python’s pickle-based model serialization: because pickle can execute arbitrary functions at load time, attackers can embed payloads that run when a model is deserialized. PickleBall automatically infers the minimal set of functions and imports required to load a given model class, then enforces that policy at load time so a model may only invoke the functions necessary for legitimate loading. The project (paper in ACM CCS ’25) targets real-world ML ecosystems (e.g., PyTorch-heavy Hugging Face models) and makes model-loading safer without changing model formats. Technically, PickleBall has two parts. The Policy Generator performs static analysis of an ML library’s source (using a Joern-based pipeline that can produce an AST or a richer CPG) to produce a JSON policy listing allowed imports and invocations for a specific model class. The Policy Enforcer is a drop-in replacement for Python’s pickle loader (enforce/enforce.py) that reads the policy and blocks any disallowed calls during unpickling. The project ships Docker images for generation and enforcement, supports cached precomputed policies (recommended: a Weights-Only Unpickler policy for PyTorch), and publishes artifacts and datasets on Zenodo for reproducibility. Caveats: Joern CPG mode can be more accurate but less stable; the authors provide evaluation and reproduction instructions in the repo.