Fix unsafe deserialization in ETRecord and export_serialize#18133
Fix unsafe deserialization in ETRecord and export_serialize#18133
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18133
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 Awaiting Approval, 1 New Failure, 3 Pending, 9 Unrelated FailuresAs of commit 1bf0180 with merge base eaf0d65 ( FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
This PR needs a
|
Add weights_only=True to torch.load() in deserialize_torch_artifact() to prevent arbitrary code execution via malicious serialized artifacts. Add security warnings to parse_etrecord() docstring and inline comments for the pickle.loads() calls on reference_outputs and representative_inputs, which remain unsafe but cannot be replaced without breaking the serialization format. Addresses TOB-EXECUTORCH-11. This PR was authored with the assistance of Claude.
Switch reference_outputs and representative_inputs serialization from pickle.dumps/loads to torch.save/torch.load(weights_only=True). The data types (Dict[str, List[List[Tensor]]] and List[Union[Tensor, int, float, bool]]) are all in the weights_only allowlist. ETRecord files are transient dev artifacts with a single writer and reader, so no backward-compatibility fallback is needed. Old ETRecord files can be regenerated from source. Addresses TOB-EXECUTORCH-11. This PR was authored with the assistance of Claude.
lucylq
requested review from
Gasoonjia,
JacobSzwejbka and
larryliu0820
as code owners
There was a problem hiding this comment.
Pull request overview
Hardens serialization/deserialization paths by switching away from raw pickle usage and enabling safer torch.load(..., weights_only=True) to mitigate arbitrary code execution risks from untrusted artifacts.
Changes:
- Use
torch.load(..., weights_only=True)indeserialize_torch_artifact. - Replace
pickle.dumps/loadsfor ETRecord metadata fields withtorch.save/loadusing in-memory buffers. - Remove
pickleimport and addiousage for byte buffering.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| exir/serde/export_serialize.py | Enables weights_only=True for safer torch artifact deserialization. |
| devtools/etrecord/_etrecord.py | Replaces pickle-based storage/loading of ETRecord metadata with torch serialization and safer loading. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| reference_outputs = torch.load( | ||
| io.BytesIO( | ||
| etrecord_zip.read(ETRecordReservedFileNames.REFERENCE_OUTPUTS) | ||
| ), | ||
| weights_only=True, | ||
| ) | ||
| elif entry == ETRecordReservedFileNames.REPRESENTATIVE_INPUTS: | ||
| # @lint-ignore PYTHONPICKLEISBAD | ||
| representative_inputs = pickle.loads( | ||
| etrecord_zip.read(ETRecordReservedFileNames.REPRESENTATIVE_INPUTS) | ||
| representative_inputs = torch.load( | ||
| io.BytesIO( | ||
| etrecord_zip.read(ETRecordReservedFileNames.REPRESENTATIVE_INPUTS) | ||
| ), | ||
| weights_only=True, | ||
| ) |
There was a problem hiding this comment.
This change makes parse_etrecord unable to read existing ETRecords that were written with pickle.dumps for these entries (the on-disk format changed). Consider adding a backward-compatible read path: attempt torch.load(..., weights_only=True) first, and if it fails with a known/expected exception, fall back to legacy pickle.loads only under an explicit opt-in (e.g., a trust_legacy_pickle: bool param or env var) plus a clear warning/error. Without this, previously-generated ETRecords will fail to load.
| buffer = io.BytesIO(serialized) | ||
| buffer.seek(0) | ||
| artifact = torch.load(buffer) | ||
| artifact = torch.load(buffer, weights_only=True) |
There was a problem hiding this comment.
weights_only is not supported in some older PyTorch versions. If this repository supports a wider PyTorch range, this will raise TypeError: load() got an unexpected keyword argument 'weights_only'. Consider a small compatibility shim (e.g., feature-detect via inspect.signature(torch.load) or try/except TypeError) so the code fails gracefully with an actionable message (or uses a safe alternative) when weights_only is unavailable.
| artifact = torch.load(buffer, weights_only=True) | |
| # `weights_only` is not supported in some older PyTorch versions. | |
| # Feature-detect support and fall back gracefully if unavailable. | |
| try: | |
| load_sig = inspect.signature(torch.load) | |
| if "weights_only" in load_sig.parameters: | |
| artifact = torch.load(buffer, weights_only=True) | |
| else: | |
| logging.warning( | |
| "torch.load does not support the 'weights_only' argument in this " | |
| "PyTorch version; falling back to torch.load(buffer) without it." | |
| ) | |
| artifact = torch.load(buffer) | |
| except (TypeError, ValueError): | |
| # As an extra safeguard, fall back to calling without `weights_only` | |
| logging.warning( | |
| "Unable to inspect torch.load signature; falling back to " | |
| "torch.load(buffer) without 'weights_only'." | |
| ) | |
| artifact = torch.load(buffer) |
| buf = io.BytesIO() | ||
| torch.save(self._reference_outputs, buf) | ||
| etrecord_zip.writestr( | ||
| ETRecordReservedFileNames.REFERENCE_OUTPUTS, | ||
| pickle.dumps(self._reference_outputs), | ||
| buf.getvalue(), | ||
| ) | ||
|
|
||
| if self._representative_inputs is not None: | ||
| buf = io.BytesIO() | ||
| torch.save(self._representative_inputs, buf) | ||
| etrecord_zip.writestr( | ||
| ETRecordReservedFileNames.REPRESENTATIVE_INPUTS, | ||
| pickle.dumps(self._representative_inputs), | ||
| buf.getvalue(), | ||
| ) |
There was a problem hiding this comment.
torch.save uses a zip-based format by default; embedding those bytes inside another ZipFile can noticeably increase size and CPU overhead (zip-inside-zip). If these fields are large, consider using torch.save(..., _use_new_zipfile_serialization=False) to avoid nested zip containers, or an alternative serialization for purely-tensor content to reduce overhead.
| reference_outputs = torch.load( | ||
| io.BytesIO( | ||
| etrecord_zip.read(ETRecordReservedFileNames.REFERENCE_OUTPUTS) | ||
| ), | ||
| weights_only=True, | ||
| ) |
There was a problem hiding this comment.
Consider forcing map_location='cpu' when loading from ETRecords. If tensors were saved on CUDA, torch.load can require GPU availability or attempt GPU allocations during parsing, which is undesirable for a devtool reader and can break in CPU-only environments.
|
can you trigger import to pay attention to the internal ci? |
Add weights_only=True to torch.load() in deserialize_torch_artifact()
to prevent arbitrary code execution via malicious serialized artifacts.
Replace usages of pickle.dumps and pickle.loads with torch.save and torch.load
As generating and reading etrecord should be done as debugging steps and usually on the same version of et, BC is not a big concern.
This PR was authored with the assistance of Claude.