[vulnops][data] fix: Replace unsafe pickle.loads with restricted unpickler in Qwen VL pipeline - #3139
Conversation
…Qwen VL data pipeline The videohandler.__call__ method used pickle.loads() directly on data from WebDataset shards, enabling arbitrary code execution via crafted pickle payloads. Replace with a RestrictedUnpickler that only allows safe built-in types (list, dict, tuple, str, int, float, etc.). The restricted unpickler is placed in a shared utility module (megatron.bridge.utils.safe_pickle) for reuse across other fixes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
yaoyu-33
changed the title
yaoyu-33
added
area:data
Resolve conflict in auto_bridge.py: take main's logger.warning() approach and improved trust_remote_code kwargs handling. Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test 666aa36 |
📝 WalkthroughWalkthroughThis PR introduces security hardening for pickle deserialization by adding a restricted unpickler module that limits deserialization to whitelisted safe types, integrating it into the video data handler, and adding safety warnings and configurable parameters for remote code execution settings in model loading. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/megatron/bridge/data/energon/task_encoder_utils.py`:
- Around line 229-233: The code extracts extension into the local variable
extension, validates using extension.lower() but then looks up
self.extensions_mapping with the original extension which can cause a KeyError
for mixed-case extensions; update the logic in the task_encoder_utils function
(where extension is set and used) to normalize the extension once (e.g.,
extension_lower = extension.lower() or assign back to extension) and use that
normalized value for both the membership check against self.extensions and the
mapping lookup self.extensions_mapping[...], ensuring consistent lowercase keys
are used for validation and retrieval.
In `@src/megatron/bridge/utils/safe_pickle.py`:
- Around line 22-42: The _SAFE_MODULES allowlist is currently a mutable dict of
mutable sets; make it immutable by converting each inner set to frozenset and
wrapping the outer mapping with types.MappingProxyType (import MappingProxyType
from types). Update the _SAFE_MODULES definition so values are frozenset(...)
for each module entry (e.g., "builtins", "collections") and replace the literal
dict with MappingProxyType({...}) to prevent runtime mutation; ensure any
references to _SAFE_MODULES remain the same.
- Around line 23-40: Remove "type" from the builtins whitelist in safe_pickle
(the set literal under "builtins") because it's unnecessary for current payloads
and increases attack surface; edit the whitelist in
src/megatron/bridge/utils/safe_pickle.py (the set containing
"list","dict",...,"NoneType") to delete the "type" entry so unpickling still
permits NoneType but not the builtin type object.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 81827af7-b5dd-45d2-a047-52b594681254
📒 Files selected for processing (3)
src/megatron/bridge/data/energon/task_encoder_utils.pysrc/megatron/bridge/models/conversion/auto_bridge.pysrc/megatron/bridge/utils/safe_pickle.py
| extension = re.sub(r".*[.]", "", key) | ||
| if extension.lower() not in self.extensions: | ||
| return None | ||
| data = pickle.loads(data) | ||
| data = safe_pickle_loads(data) | ||
| key = self.extensions_mapping[extension] |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n src/megatron/bridge/data/energon/task_encoder_utils.py | sed -n '225,240p'Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 787
🏁 Script executed:
grep -n "self.extensions" src/megatron/bridge/data/energon/task_encoder_utils.py | head -20Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 326
Normalize extension before lookup to avoid mixed-case KeyError.
Line 230 validates with extension.lower(), but Line 233 indexes self.extensions_mapping using the original extension. Since both self.extensions and self.extensions_mapping are initialized with lowercase keys only ("jpgs", "mp4s", "videos"), any file with mixed-case extension (e.g., file.JPGS) will pass validation but fail at line 233 with a KeyError.
✅ Proposed fix
extension = re.sub(r".*[.]", "", key)
- if extension.lower() not in self.extensions:
+ extension = extension.lower()
+ if extension not in self.extensions:
return None
data = safe_pickle_loads(data)
key = self.extensions_mapping[extension]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| extension = re.sub(r".*[.]", "", key) | |
| if extension.lower() not in self.extensions: | |
| return None | |
| data = pickle.loads(data) | |
| data = safe_pickle_loads(data) | |
| key = self.extensions_mapping[extension] | |
| extension = re.sub(r".*[.]", "", key) | |
| extension = extension.lower() | |
| if extension not in self.extensions: | |
| return None | |
| data = safe_pickle_loads(data) | |
| key = self.extensions_mapping[extension] |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/data/energon/task_encoder_utils.py` around lines 229 -
233, The code extracts extension into the local variable extension, validates
using extension.lower() but then looks up self.extensions_mapping with the
original extension which can cause a KeyError for mixed-case extensions; update
the logic in the task_encoder_utils function (where extension is set and used)
to normalize the extension once (e.g., extension_lower = extension.lower() or
assign back to extension) and use that normalized value for both the membership
check against self.extensions and the mapping lookup
self.extensions_mapping[...], ensuring consistent lowercase keys are used for
validation and retrieval.
| _SAFE_MODULES = { | ||
| "builtins": { | ||
| "list", | ||
| "dict", | ||
| "tuple", | ||
| "set", | ||
| "frozenset", | ||
| "bytes", | ||
| "bytearray", | ||
| "str", | ||
| "int", | ||
| "float", | ||
| "bool", | ||
| "complex", | ||
| "slice", | ||
| "range", | ||
| "type", | ||
| "NoneType", | ||
| }, | ||
| "collections": {"OrderedDict"}, | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify Ruff warning is resolved for this file.
uv run ruff check src/megatron/bridge/utils/safe_pickle.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 113
🏁 Script executed:
cat -n src/megatron/bridge/utils/safe_pickle.py | head -50Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1987
🌐 Web query:
Ruff RUF012 mutable class attribute
💡 Result:
Ruff RUF012 (mutable-class-default) flags mutable default values assigned to class attributes (e.g., [], {}, set() in the class body), because that mutable object is shared by all instances, which can cause surprising cross-instance bugs. [1]
Example that triggers RUF012
class A:
items: list[int] = [] # shared across all A instancesCommon fixes
- Make it an instance attribute (usually preferred):
class A:
def __init__(self) -> None:
self.items: list[int] = []- If it’s truly intended to be shared, annotate as
ClassVar:
from typing import ClassVar
class A:
items: ClassVar[list[int]] = []- Use an immutable default instead (tuple/frozenset/proxy mapping, etc.). [1]
RUF012 is listed in Ruff’s rules as: “Mutable default value for class attribute.” [2]
Sources: [1] [2]
Make the allowlist immutable (RUF012).
Line 22 defines a mutable class attribute (dict of sets). Ruff's RUF012 rule flags this because the shared mutable object can be accidentally modified at runtime, affecting all instances. Use MappingProxyType and frozenset to make it immutable.
♻️ Proposed fix
+from types import MappingProxyType
+
class _RestrictedUnpickler(pickle.Unpickler):
@@
- _SAFE_MODULES = {
- "builtins": {
+ _SAFE_MODULES = MappingProxyType({
+ "builtins": frozenset({
"list",
"dict",
"tuple",
"set",
"frozenset",
"bytes",
"bytearray",
"str",
"int",
"float",
"bool",
"complex",
"slice",
"range",
"type",
"NoneType",
- },
- "collections": {"OrderedDict"},
- }
+ }),
+ "collections": frozenset({"OrderedDict"}),
+ })🧰 Tools
🪛 Ruff (0.15.9)
[warning] 22-42: Mutable default value for class attribute
(RUF012)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/utils/safe_pickle.py` around lines 22 - 42, The
_SAFE_MODULES allowlist is currently a mutable dict of mutable sets; make it
immutable by converting each inner set to frozenset and wrapping the outer
mapping with types.MappingProxyType (import MappingProxyType from types). Update
the _SAFE_MODULES definition so values are frozenset(...) for each module entry
(e.g., "builtins", "collections") and replace the literal dict with
MappingProxyType({...}) to prevent runtime mutation; ensure any references to
_SAFE_MODULES remain the same.
| "builtins": { | ||
| "list", | ||
| "dict", | ||
| "tuple", | ||
| "set", | ||
| "frozenset", | ||
| "bytes", | ||
| "bytearray", | ||
| "str", | ||
| "int", | ||
| "float", | ||
| "bool", | ||
| "complex", | ||
| "slice", | ||
| "range", | ||
| "type", | ||
| "NoneType", | ||
| }, |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -type f -name "safe_pickle.py" | head -5Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 113
🏁 Script executed:
find . -type f -path "*/bridge/data/datasets/utils.py" | head -5Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 115
🏁 Script executed:
cat -n src/megatron/bridge/utils/safe_pickle.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 2419
🏁 Script executed:
sed -n '1230,1280p' src/megatron/bridge/data/datasets/utils.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1994
🏁 Script executed:
sed -n '1239,1270p' src/megatron/bridge/data/datasets/utils.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1299
🏁 Script executed:
rg "safe_pickle_load" src/Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 448
🏁 Script executed:
rg -B 5 -A 5 "safe_pickle_loads" src/megatron/bridge/data/energon/task_encoder_utils.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1086
🏁 Script executed:
rg "type" src/megatron/bridge/utils/safe_pickle.py -A 3 -B 3Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1144
🏁 Script executed:
rg "NoneType\|type\(" src/megatron/bridge/data/datasets/utils.py src/megatron/bridge/data/energon/task_encoder_utils.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
🏁 Script executed:
rg "type\(" src/megatron/bridge/ -t py | grep -i pickleRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
🏁 Script executed:
python3 << 'EOF'
import pickle
import io
# Test what's needed to unpickle a simple dict with int values
data = dict(newline_int=100, version="1.0")
pickled = pickle.dumps(data)
# Check what classes are referenced in the pickle
import pickletools
print("Pickle opcodes for dict(newline_int=100, version='1.0'):")
pickletools.dis(io.BytesIO(pickled))
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 602
🏁 Script executed:
python3 << 'EOF'
import pickle
import io
import pickletools
# Test what happens when we pickle a list with None values
data = [1, 2, None, "test", 3.14]
pickled = pickle.dumps(data)
print("Pickle opcodes for list with None:")
pickletools.dis(io.BytesIO(pickled))
# Check if type is ever referenced
print("\n\nSearching for 'type' in pickle representation:")
if b'type' in pickled:
print("Found 'type' in pickle data")
else:
print("No 'type' found in pickle data")
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 583
🏁 Script executed:
python3 << 'EOF'
import pickle
import io
import pickletools
# Test if type is ever needed during unpickling
# The whitelist includes both 'type' and 'NoneType'
# Check if we can unpickle None without 'type'
data = None
pickled = pickle.dumps(data)
print("Pickle opcodes for None:")
pickletools.dis(io.BytesIO(pickled))
print("\n\n---")
# Now test a more complex scenario
data = {"version": __import__('sys').version_info.major}
pickled = pickle.dumps(data)
print("Pickle opcodes for dict with sys module reference:")
pickletools.dis(io.BytesIO(pickled))
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 516
🏁 Script executed:
python3 << 'EOF'
# Check if NoneType is actually referenced during unpickling
import pickle
import io
import pickletools
# NoneType is the type of None, but let's see if it's actually used
data = type(None)
pickled = pickle.dumps(data)
print("Pickle opcodes for NoneType (type(None)):")
pickletools.dis(io.BytesIO(pickled))
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 521
🏁 Script executed:
python3 << 'EOF'
# Check how to handle NoneType in a restricted unpickler
import pickle
import io
# NoneType is type(None), but how is it represented?
print("NoneType:", type(None))
print("NoneType name:", type(None).__name__)
print("NoneType module:", type(None).__module__)
# Can we access builtins.NoneType?
import builtins
try:
print("builtins.NoneType:", builtins.NoneType)
except AttributeError as e:
print(f"Error accessing builtins.NoneType: {e}")
# What about the actual way to get NoneType?
NoneType = type(None)
print(f"\ntype(None).__name__ = {NoneType.__name__}")
print(f"type(None).__module__ = {NoneType.__module__}")
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 295
🏁 Script executed:
python3 << 'EOF'
import pickle
import io
# Test if 'type' is actually needed for the metadata dict case
class RestrictedUnpickler(pickle.Unpickler):
"""Test without 'type' in whitelist"""
_SAFE_MODULES = {
"builtins": {
"list",
"dict",
"tuple",
"set",
"frozenset",
"bytes",
"bytearray",
"str",
"int",
"float",
"bool",
"complex",
"slice",
"range",
# "type", # REMOVED
"NoneType",
},
"collections": {"OrderedDict"},
}
def find_class(self, module: str, name: str) -> type:
if module in self._SAFE_MODULES and name in self._SAFE_MODULES[module]:
return super().find_class(module, name)
raise pickle.UnpicklingError(
f"Restricted unpickler refused to load '{module}.{name}'."
)
# Test 1: Simple metadata dict (the known payload)
print("Test 1: Unpickling metadata dict (known payload)")
data = dict(newline_int=100, version="1.0")
pickled = pickle.dumps(data)
try:
result = RestrictedUnpickler(io.BytesIO(pickled)).load()
print(f"✓ SUCCESS: {result}")
except pickle.UnpicklingError as e:
print(f"✗ FAILED: {e}")
# Test 2: Unpickling NoneType
print("\nTest 2: Unpickling type(None)")
data = type(None)
pickled = pickle.dumps(data)
try:
result = RestrictedUnpickler(io.BytesIO(pickled)).load()
print(f"✓ SUCCESS: {result}")
except pickle.UnpicklingError as e:
print(f"✗ FAILED: {e}")
# Test 3: Unpickling dict with None values
print("\nTest 3: Unpickling dict with None values")
data = dict(value=None, count=10)
pickled = pickle.dumps(data)
try:
result = RestrictedUnpickler(io.BytesIO(pickled)).load()
print(f"✓ SUCCESS: {result}")
except pickle.UnpicklingError as e:
print(f"✗ FAILED: {e}")
EOFRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 346
🏁 Script executed:
rg "type\(None\)" src/ -t pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 53
Remove builtins.type from the whitelist: it is not required for the known payloads and weakens the security boundary.
The metadata dict pickled at lines 1239-1270 (dict(newline_int=newline_int, version=__idx_version__)) unpickles successfully without type. While NoneType is listed in the whitelist, unpickling None values does not require the type builtin; type is only needed if someone explicitly pickles type(None) itself, which does not occur in this codebase. Including type expands the callable construction surface during unpickling unnecessarily.
🔒 Suggested hardening
- "type",🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/megatron/bridge/utils/safe_pickle.py` around lines 23 - 40, Remove "type"
from the builtins whitelist in safe_pickle (the set literal under "builtins")
because it's unnecessary for current payloads and increases attack surface; edit
the whitelist in src/megatron/bridge/utils/safe_pickle.py (the set containing
"list","dict",...,"NoneType") to delete the "type" entry so unpickling still
permits NoneType but not the builtin type object.
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| "complex", | ||
| "slice", | ||
| "range", | ||
| "type", |
There was a problem hiding this comment.
Bug: builtins.type should not be in the allowlist. Pickle can resolve builtins.type via find_class and then invoke it via REDUCE as type(name, bases, namespace_dict), which constructs an arbitrary new class. This undermines the purpose of the restricted unpickler.
Unless there's a known data format that pickles type objects in the Qwen VL shards, this should be removed:
| "type", | |
| "type", |
→ delete this line.
If you do need it, please add a comment explaining why and what types are expected.
|
|
||
| import dataclasses | ||
| import logging | ||
| import warnings |
There was a problem hiding this comment.
Nit: import warnings is added here but never used — only logger.warning() (from logging) is called in this file. Looks like a leftover from the merge conflict resolution.
| import warnings |
- Make _SAFE_MODULES immutable with MappingProxyType/frozenset (RUF012) - Remove builtins.type from allowlist to reduce attack surface - Normalize extension to lowercase before mapping lookup to prevent KeyError - Add unit tests for safe_pickle (round-trip, rejection, immutability) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
/ok to test 858ed24 |
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
/ok to test ef47a9d |
…rtability The hardcoded byte-string payloads used protocol 5 framing that caused "pickle data was truncated" on Python 3.12 instead of hitting find_class. Use pickle.dumps() to generate valid payloads at runtime. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
/ok to test f795f1a |
Manual bump since PR #3139 passed on mcore-dev. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
cuichenx
left a comment
cuichenx
left a comment
There was a problem hiding this comment.
tested on toy video dataset and ran without issues
chtruong814
added
the
r0.4.0
QA RCCA Analysis1. Fix Reference
2. Root CauseSECURITY VULNERABILITY: 3. Trigger Configuration
4. Nature of the BugClassification: CODE BUG (SECURITY) - Arbitrary code execution via malicious pickle in WebDataset 5. Existing Test CoverageIn Fix PR: YES - 1 test file:
In NMFW Tests: N/A for security fixes 6. Coverage Assessment
7. New Regression TestNOT NEEDED - Fix PR includes unit tests for 8. ConclusionVerdict: ADEQUATE COVERAGE - Security fix with unit tests for restricted unpickler. |
…ckler in Qwen VL pipeline (NVIDIA-NeMo#3139) Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
4 participants
Summary
safe_pickle.pyutility with_RestrictedUnpicklerthat only allows safe built-in types (list, dict, tuple, str, int, float, bool, bytes, etc.)pickle.loads()in Qwen VLvideohandlerwithsafe_pickle_loads()to prevent arbitrary code execution from malicious WebDataset shardspickle.UnpicklingErrorif any non-whitelisted type is encounteredTest plan
🤖 Generated with Claude Code
Summary by CodeRabbit
trust_remote_codeparameter for model loading with secure default setting