bridge.utils.safe_pickle

Module Contents

Classes

_SafeEnumToken	Hold a validated Enum member until the pickle VM has finished.
_SafeEnumResolver	Resolve a validated Enum value to an inert token.
_RestrictedUnpickler	Unpickler that only allows safe built-in types to prevent arbitrary code execution.
_TorchTensorRestrictedUnpickler	Restricted unpickler for plain tensors and containers of tensors.
_NumpyRestrictedUnpickler	Unpickler that allows safe builtins and the narrow set of numpy types needed for object array reconstruction.
_EnergonUnpickler	Unpickler for Energon dataloader state files (.pt).

Functions

_restore_legacy_bytes	Reconstruct bytes from the fixed representation emitted by pickle protocols 0-2.
_find_safe_loaded_enum	Resolve a plain Enum through already-loaded namespaces without importing modules.
_restore_safe_enum_tokens	Replace inert Enum tokens after all pickle opcodes have completed.
_safe_torch_load_from_bytes	Reconstruct tensor storage bytes without enabling arbitrary pickle globals.
energon_torch_load	Load an Energon dataloader state .pt file through a restricted unpickler.
safe_pickle_load	Deserialize from a file using a restricted unpickler that only allows safe types.
safe_pickle_loads	Deserialize pickle data using a restricted unpickler that only allows safe types.
safe_torch_tensor_pickle_loads	Deserialize raw pickle data containing only safe containers and plain torch tensors.
safe_load_npy	Load a .npy file from raw bytes without enabling unrestricted pickle.

Data

_BUILTIN_SAFE_TYPES
_ENERGON_SAFE_STATE_GLOBALS
_TRAVERSAL_IN_PROGRESS

API

bridge.utils.safe_pickle._BUILTIN_SAFE_TYPES

‘frozenset(…)’

bridge.utils.safe_pickle._ENERGON_SAFE_STATE_GLOBALS

‘MappingProxyType(…)’

bridge.utils.safe_pickle._TRAVERSAL_IN_PROGRESS

‘object(…)’

bridge.utils.safe_pickle._restore_legacy_bytes(value: object, encoding: object) → bytes

Reconstruct bytes from the fixed representation emitted by pickle protocols 0-2.

class bridge.utils.safe_pickle._SafeEnumToken(member: enum.Enum)

Hold a validated Enum member until the pickle VM has finished.

Initialization

__slots__

(‘_member’,)

__call__(*_args: object) → None

__setstate__(_state: object) → None

class bridge.utils.safe_pickle._SafeEnumResolver(members: tuple[tuple[object, enum.Enum], ...])

Resolve a validated Enum value to an inert token.

Initialization

__slots__

(‘_members’,)

__call__(*args: object) → bridge.utils.safe_pickle._SafeEnumToken

__setstate__(_state: object) → None

bridge.utils.safe_pickle._find_safe_loaded_enum(

module: str,

name: str,

) → bridge.utils.safe_pickle._SafeEnumResolver | None

Resolve a plain Enum through already-loaded namespaces without importing modules.

bridge.utils.safe_pickle._restore_safe_enum_tokens(

value: object,

memo: dict[int, object] | None = None,

) → object

Replace inert Enum tokens after all pickle opcodes have completed.

class bridge.utils.safe_pickle._RestrictedUnpickler

Bases: pickle.Unpickler

Unpickler that only allows safe built-in types to prevent arbitrary code execution.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_SAFE_MODULES

‘MappingProxyType(…)’

find_class(module: str, name: str) → object

bridge.utils.safe_pickle._safe_torch_load_from_bytes(data: bytes) → object

Reconstruct tensor storage bytes without enabling arbitrary pickle globals.

class bridge.utils.safe_pickle._TorchTensorRestrictedUnpickler

Bases: bridge.utils.safe_pickle._RestrictedUnpickler

Restricted unpickler for plain tensors and containers of tensors.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_SAFE_MODULES

‘MappingProxyType(…)’

find_class(module: str, name: str) → object

class bridge.utils.safe_pickle._NumpyRestrictedUnpickler

Bases: pickle.Unpickler

Unpickler that allows safe builtins and the narrow set of numpy types needed for object array reconstruction.

NumPy object arrays (dtype=’O’) are serialized via pickle inside .npy files. The pickle stream references numpy.core.multiarray._reconstruct, numpy.ndarray, and numpy.dtype to rebuild the array container, while the elements (dicts, lists, ints, …) use only standard builtins.

This unpickler permits exactly those types and nothing else — in particular, os, subprocess, builtins.eval, etc. are blocked, preventing arbitrary-code-execution attacks via crafted .npy files.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_SAFE_MODULES

‘MappingProxyType(…)’

find_class(module: str, name: str) → object

class bridge.utils.safe_pickle._EnergonUnpickler

Bases: bridge.utils.safe_pickle._NumpyRestrictedUnpickler

Unpickler for Energon dataloader state files (.pt).

Extends the NumPy-safe unpickler with the exact Energon dataclass types that Energon serialises into dataloader checkpoint files and inert tokens for narrowly validated, already-loaded Enum members used as grouping keys. All other globals — including os, subprocess, and any __reduce__ payload callable outside these rules — are blocked, preventing arbitrary code execution from attacker-controlled checkpoint files.

Use via :func:energon_torch_load rather than instantiating directly.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_SAFE_MODULES: types.MappingProxyType

‘MappingProxyType(…)’

find_class(module: str, name: str) → object

bridge.utils.safe_pickle.energon_torch_load(

path: str | BinaryIO,

*,

map_location: str = 'cpu',

) → object

Load an Energon dataloader state .pt file through a restricted unpickler.

Parses the torch zip format directly without calling torch.load. Security is enforced by :class:_EnergonUnpickler: a GLOBAL opcode must resolve to an explicitly allowlisted type or a narrowly validated, already-loaded Enum. Enum members remain inert private tokens until every pickle opcode has completed, blocking application hooks during deserialization.

torch.load(weights_only=True) is not used because PyTorch ≥ 2.13 restricts SETITEM/SETITEMS to exact dict, OrderedDict, and Counter types, rejecting dict subclasses such as Energon’s FlexState — which is always present in real Energon checkpoints (SavableDatasetState.dataset_state is typed FlexState, not Optional).

torch.save writes a zip archive whose directory prefix is the file stem. This function opens the zip, runs :class:_EnergonUnpickler on the pickle stream, and reconstructs tensor storages from the raw blobs via persistent_load. Storages are cached by key so that tensors sharing a storage (views, slices) remain aliased after restore.

Parameters:

• path –

  Path to or binary stream for the .pt file written by

  func:

  ~megatron.bridge.training.checkpointing.maybe_save_dataloader_state.

• map_location – Device to map tensor storages to; defaults to "cpu" to avoid GPU allocation during restore.

Returns:

The deserialized object (a dict containing "dataloader_state_dict").

bridge.utils.safe_pickle.safe_pickle_load(fp) → object

Deserialize from a file using a restricted unpickler that only allows safe types.

bridge.utils.safe_pickle.safe_pickle_loads(data: bytes) → object

Deserialize pickle data using a restricted unpickler that only allows safe types.

bridge.utils.safe_pickle.safe_torch_tensor_pickle_loads(data: bytes) → object

Deserialize raw pickle data containing only safe containers and plain torch tensors.

bridge.utils.safe_pickle.safe_load_npy(data: bytes)

Load a .npy file from raw bytes without enabling unrestricted pickle.

For numeric arrays the fast allow_pickle=False path is used. For object arrays (packed datasets storing dicts of variable-length lists) the pickle payload is deserialized through :class:_NumpyRestrictedUnpickler, which blocks dangerous modules like os and subprocess.

Parameters:

data – Raw bytes of a .npy file.

Returns:

numpy.ndarray loaded from the file.