bridge.utils.safe_pickle

Module Contents

Classes

_RestrictedUnpickler	Unpickler that only allows safe built-in types to prevent arbitrary code execution.
_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

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_load_npy	Load a .npy file from raw bytes without enabling unrestricted pickle.

Data

_BUILTIN_SAFE_TYPES

API

bridge.utils.safe_pickle._BUILTIN_SAFE_TYPES

‘frozenset(…)’

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) → type

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) → type

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. All other globals — including os, subprocess, and any __reduce__ payload callable outside this allowlist — 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) → type

bridge.utils.safe_pickle.energon_torch_load(path: str, *, 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: any GLOBAL opcode whose (module, name) is not in the explicit allowlist raises pickle.UnpicklingError, blocking __reduce__-based code execution from attacker-controlled checkpoint files.

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 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_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.