Model and artifact supply-chain security
Keep hostile weights, packages, and build outputs from becoming code execution in your training and serving environments.
What you'll learn
- Why pickle deserialization is code execution, and why an unsafe torch.load can compromise a training cluster
- What safetensors prevents, and the attacks it cannot prevent
- How typosquatting, dependency confusion, and public fine-tunes reach ML systems
- How SBOMs, attestations, signatures, and admission checks fit together
- The smallest practical supply-chain security baseline for an ML team
Before you start
At 3:07 a.m., a training job starts on four GPUs.
The job downloads model.pt from a public model hub, calls torch.load, and never reaches the first training step. The worker instead opens an outbound connection using the pod’s cloud credentials. It can read the training bucket, the experiment database, and the Kubernetes API token mounted into the container.
Nothing was wrong with the model’s accuracy. The model was the attack.
A machine-learning pipeline consumes far more than a neural network. It pulls Python packages, container images, model weights, tokenizer files, configuration, preprocessing code, datasets, and sometimes custom code. Several of those things are executable, even when they look like passive data.
Supply-chain security is the practice of proving what entered that pipeline, who built it, what changed, and whether the exact item is allowed to run. The important word is exact. A tag such as latest is a moving label. A cryptographic digest identifies one exact byte sequence.
This is different from the model-behaviour attacks covered in ML security. A prompt can trigger bad model behaviour. A malicious checkpoint can run a command before the model produces a prediction.
The dangerous file that looks like weights
Python’s pickle format does not merely store data. It stores instructions for reconstructing Python objects.
During unpickling, a class can provide a __reduce__ method that tells the unpickler which callable to invoke and with which arguments. That callable might construct a tensor—or run a shell command, open a socket, or copy credentials.
That is why “never unpickle untrusted data” is a code-execution rule.
PyTorch historically used pickle for the Python object structure around tensors. torch.save can therefore create files whose loading path invokes Python reconstruction logic. An explicit torch.load(..., weights_only=False) accepts that behaviour.
This example uses print rather than a destructive command, but demonstrates the mechanism:
import torch
from pathlib import Path
class Marker:
def __reduce__(self):
return (print, ("code ran while loading",))
path = Path("untrusted-demo.pt")
torch.save(Marker(), path)
torch.load(path, weights_only=False)
path.unlink()
On a compatible PyTorch installation, the load prints:
code ran while loading
The print happens during deserialization, before the application has a model object to inspect. Replacing print with a process-launching callable produces arbitrary code execution.
A training worker is a valuable victim: it often has GPU and network access, package-install tools, and credentials for object storage and experiment tracking. A serving pod may have less access, but can still affect predictions or reach internal services. The file does not need to escape the container; the container’s own permissions are usually enough.
What weights_only changes
Starting with PyTorch 2.6, torch.load defaults to weights_only=True when a custom pickle module is not supplied. Its restricted unpickler is designed for tensor state dictionaries, primitive values, dictionaries, and explicitly allowlisted types. It does not freely import and execute arbitrary globals like ordinary pickle loading.
Only allowlist classes and functions from reviewed, trusted code. Allowlisting a type from an untrusted checkpoint can execute its code during reconstruction; never treat an allowlist error as a reason to trust the checkpoint.
This is a meaningful improvement, not a universal safety certificate. Older versions used a different default, existing code may pass weights_only=False, and some checkpoints pressure engineers to disable the restriction. Restricted loading can still consume excessive memory or CPU, contain parser vulnerabilities, or be accompanied by custom Python code executed through another path.
The safest answer is to avoid pickle for tensor exchange. That is where safetensors helps.
Safetensors: safer loading, not trustworthy weights
A safetensors file contains tensor names, shapes, data types, offsets, and raw tensor bytes. Its header is structured metadata rather than a Python object graph, so a loader can map tensor bytes without invoking arbitrary Python constructors.
For a normal state dictionary, that removes the pickle deserialization route and makes the file easier to inspect or load incrementally. A 7.8 GB file need not become an opaque Python object before anyone knows which tensors it contains.
Safetensors does not answer three different questions:
- Who approved these exact bytes? A verified signature can answer that. Establishing who produced or built them requires trusted provenance or an acquisition attestation.
- Are the numbers benign? The weights may contain a backdoor, a deliberately biased classifier, or values that make inference consume extreme resources.
- What else does the repository execute? A model hub entry may include configuration, tokenizer code, preprocessing code, or a custom model implementation.
Some high-level loaders offer options such as trust_remote_code=True. That means code from the remote repository may be imported. Do not enable it for an unreviewed repository merely because its weights end in .safetensors.
Safetensors makes loading the numbers less dangerous. It does not establish that the numbers, surrounding code, or publisher deserve trust.
| Choice | Code execution during ordinary weight loading | Main limitation |
|---|---|---|
| Pickle or legacy PyTorch checkpoint | High if loaded unsafely | Deserialization can execute code |
| PyTorch restricted weights load | Lower for supported tensor state dictionaries | Compatibility pressure and resource attacks remain |
| Safetensors | No Python object reconstruction in the tensor format | Does not stop poisoned numbers, custom code, or replacement |
| Rebuilt artifact from reviewed source | Depends on the build and dependencies | Reproducibility and review cost are real |
The key distinction is between format safety and artifact trust. They solve different problems.
The package path is just as dangerous
A model pipeline may install:
- a deep-learning framework;
- tokenizer and dataset libraries;
- GPU extensions;
- evaluation tools; and
- internal packages.
A typosquatted package resembles a popular package. An engineer intends to install transformers, types a slightly different name, and receives the attacker’s package.
Dependency confusion targets internal names. If your company has an internal feature-prep package and the resolver also searches a public index, an attacker can publish a public package with that name and a version designed to win resolution. The package can execute code during its build or import.
PEP 517 build isolation limits dependency contamination; it does not make a malicious build backend safe.
Inspection after installation may be too late: a build step can execute before the wheel exists, and module-level code can run on import.
Use these controls:
- Lock exact versions and hashes. Pinning
torchto2.6.0does not pin the bytes served for every platform; hash checking does. - Use a private package index or explicit index policy for internal names. Do not let a public package silently satisfy an internal dependency.
- Require hashes in automated installs. For pip,
--require-hasheshelps only when every requirement has an approved hash. - Build in a disposable environment without cloud credentials and with restricted network access.
- Record the resolved package set in the image’s SBOM and scan before promotion.
A clean vulnerability scan is not proof of trust. A scanner may know about a published CVE, but generally cannot identify a newly uploaded package with no CVE as a dependency-confusion attack.
Provenance, signatures, and the admission gate
A hash tells you whether bytes changed. A signature tells you whether the holder of a particular private key approved those bytes.
Neither tells you whether the build was sensible. That requires provenance: evidence about the source revision, builder identity, workflow, inputs, and build time.
Four records cover the common production boundary:
- An SBOM, or software bill of materials, lists the packages and libraries inside the container image.
- A build attestation records which source, workflow, and builder produced the image.
- A model signature approves the exact model artifact or model-package digest.
- An admission check verifies those records immediately before the image or job runs.
A container signature does not cover a model downloaded later from object storage. Sign the model itself, or package it into an immutable, signed OCI artifact and verify that artifact. Tools such as Sigstore’s Cosign can sign and verify OCI images and artifacts; key management remains an architectural decision.
Here is the practical sequence:
Suppose trainer:latest pointed to image digest A at 09:00, then to digest B at 09:20. A download-time check may approve A while a later deployment resolves the tag to B. A promotion record containing the image and model digests prevents that ambiguity.
Admission is the last trustworthy moment before execution. Reject a workload unless:
- the image uses a digest, not a mutable tag;
- the image signature chains to an approved identity;
- its SBOM and build attestation meet policy;
- the model digest is signed by an approved identity;
- the model format is allowed, preferably safetensors;
- the manifest names the expected tokenizer, configuration, and preprocessing versions;
- remote custom code is disabled unless explicitly reviewed.
Verify during download too. It catches bad data early, but registries copy artifacts, tags move, and deployment configuration drifts. Verify the exact bytes at promotion and admission.
This belongs beside, not inside, your model registry. The registry records model versions, evaluation results, owners, and promotion state. Supply-chain controls prove that the version evaluated is the one you intend to run.
Public weights are third-party software
A pretrained model from a public hub is not safe merely because thousands of people downloaded it. Popularity is not provenance.
A poisoned model can behave normally on evaluation data and fail on a trigger:
- a phrase;
- an image feature;
- a rare token sequence; or
- an input format.
A fine-tune can introduce that behaviour while preserving benchmark scores. A LoRA adapter is smaller than a full checkpoint, but can still change behaviour substantially.
Treat every fine-tune as a new release. Record:
- the exact base-model digest;
- the adapter or merged-weight digest;
- the training dataset and code revision;
- intended and prohibited use, license, and evaluation scope;
- tests for suspicious triggers and unexpected output changes.
A model card describes intended use, limitations, evaluation, training approach, and known risks. Dataset documentation records source, collection method, licensing or consent basis, preprocessing, exclusions, and quality problems.
These documents do not stop code execution, but make omissions visible. Make them part of the promotion record, not an unchecked wiki page.
A small team’s workable baseline
For every production model, keep a manifest containing:
- the model format;
- byte size;
- SHA-256 digest;
- base-model digest;
- tokenizer digest;
- package lock;
- container digest;
- owner; and
- evaluation record.
A 6.2 GB model with one changed byte is a different artifact; the digest makes that difference cheap to detect.
Use this path:
- Download into quarantine without cloud credentials, Kubernetes tokens, or production network access.
- Prefer safetensors. Convert legacy checkpoints only in a disposable, isolated environment; conversion still crosses the unsafe loading boundary.
- Build training and serving images from locked dependencies, scan them, and store an SBOM.
- Pin image and model references by digest and packages by version plus hash.
- Sign the image and model artifact. Attach:
- build provenance;
- model-card; and
- dataset documentation.
- Gate promotion and admission on signature verification and expected digests. A failed check blocks the workload.
- Run with least privilege:
- read-only model storage where possible;
- minimal egress;
- no runtime package installation; and
- only the cloud permissions the job needs.
The honest limitation
No serialization format can tell you that a set of numbers is good. No signature can tell you that the signer reviewed the training data carefully. A signed poisoned model is still poisoned.
These controls cost time:
- key rotation;
- identity policies;
- evidence storage;
- exception review; and
- rebuilds when dependencies become vulnerable.
Safetensors may not preserve every custom Python object, and admission gates can block an emergency if nobody owns approval.
That is still a good trade. The aim is not to make every artifact harmless, but to make the trust decision explicit, reproducible, and difficult to skip at 3 a.m.
What to remember
- Loading pickle is deserialization plus execution. An unsafe
torch.loadcan run code before training begins. - Safetensors removes Python object reconstruction from tensor loading, but does not prove origin, prevent poisoned weights, or secure custom code.
- Packages, model-hub files, containers, tokenizers, and adapters all belong in supply-chain review.
- Hashes identify exact bytes. Signatures identify an approving identity. Provenance explains how bytes were built.
- Verify image and model digests at promotion and admission. A download check alone cannot control a mutable tag or later registry copy.
Quick check
Practice this in an interview
All questionsMLSecOps extends security practices across the whole ML lifecycle rather than just the deployed app, covering data, training, the model artifact, and serving. Key threats include data and model poisoning, adversarial evasion inputs, model theft or extraction, membership-inference and privacy leakage, and supply-chain risks like malicious model files and dependencies. Defenses span provenance and validation, robustness testing, access control and signing of artifacts, input monitoring, and scanning, integrated into the MLOps pipeline.
A model registry is a central catalog for deployable model artifacts and their metadata, lineage, approvals, and immutable versions. Production systems promote a tested version through deployment controls, usually using a mutable alias such as champion while retaining the exact version for rollback and audit.
Register every candidate as an immutable, versioned artifact, then move it through environments (dev to staging to prod) gated by automated checks rather than promoting straight to prod. In modern MLflow you use aliases like champion and challenger instead of the deprecated stage labels, and promotion is a governed, auditable action with sign-off and an easy rollback by repointing the alias. Always validate in staging and roll out progressively (canary or shadow) before full traffic.
Tool poisoning is malicious instruction content in a tool description or result; cross-server shadowing uses one server’s names or content to influence or misroute another server’s capability; a rug pull changes a previously reviewed capability later. Defenses combine origin-aware namespaces, isolated trust contexts, capability snapshots and change review, exact-argument authorization, sandboxing, and runtime monitoring.