fix: restrict SafeUnpickler to explicit globals (#40259)
Co-authored-by: yhzhuang <yhzhuang@fb.com> Co-authored-by: Jihui Yang <16509088+jihuiyang@users.noreply.github.com>
This commit is contained in:
Yonghao Zhuang
co-authored by
yhzhuang
Jihui Yang
yhzhuang
Jihui Yang
parent
5e9342d16f
commit
5b42d10edf
@@ -2966,6 +2966,15 @@ def normalize_serialized_named_tensor_payloads(
|
||||
return [normalize_serialized_named_tensor_payload(data) for data in payloads]
|
||||
|
||||
|
||||
def _safe_load_torch_storage(data: bytes):
|
||||
storage = torch.load(io.BytesIO(data), weights_only=True)
|
||||
if not isinstance(storage, (torch.storage.TypedStorage, torch.UntypedStorage)):
|
||||
raise pickle.UnpicklingError(
|
||||
f"Expected a Torch storage, got {type(storage).__name__}"
|
||||
)
|
||||
return storage
|
||||
|
||||
|
||||
class SafeUnpickler(pickle.Unpickler):
|
||||
# Standard-library modules expose powerful callables alongside harmless data
|
||||
# types. Keep these globals exact so a newly added callable is denied by
|
||||
@@ -2989,6 +2998,9 @@ class SafeUnpickler(pickle.Unpickler):
|
||||
("collections", "OrderedDict"),
|
||||
("collections", "defaultdict"),
|
||||
("collections", "deque"),
|
||||
("collections", "Counter"),
|
||||
("copyreg", "__newobj__"),
|
||||
("copyreg", "__newobj_ex__"),
|
||||
("functools", "partial"),
|
||||
("itertools", "chain"),
|
||||
("itertools", "repeat"),
|
||||
@@ -2996,41 +3008,193 @@ class SafeUnpickler(pickle.Unpickler):
|
||||
("multiprocessing.reduction", "_rebuild_socket"),
|
||||
("multiprocessing.resource_sharer", "DupFd"),
|
||||
("types", "SimpleNamespace"),
|
||||
}
|
||||
|
||||
ALLOWED_MODULE_PREFIXES = {
|
||||
# --- PyTorch types ---
|
||||
"torch.",
|
||||
"torch._tensor.",
|
||||
"torch.storage.",
|
||||
"torch.nn.parameter.",
|
||||
"torch.autograd.function.",
|
||||
# --- torch distributed ---
|
||||
"torch.distributed.",
|
||||
"torch.distributed._shard.",
|
||||
"torch.distributed._composable.",
|
||||
"torch._C._distributed_c10d.",
|
||||
"torch._C._distributed_fsdp.",
|
||||
"torch.distributed.optim.",
|
||||
# --- PEFT / LoRA ---
|
||||
"peft.",
|
||||
"transformers.",
|
||||
"huggingface_hub.",
|
||||
# --- SGLang & Unitest ---
|
||||
"sglang.srt.weight_sync.tensor_bucket.",
|
||||
"sglang.srt.model_executor.model_runner.",
|
||||
"sglang.srt.model_executor.model_runner_components.weight_updater.",
|
||||
"sglang.srt.layers.",
|
||||
"sglang.srt.utils.",
|
||||
"sglang.srt.disaggregation.",
|
||||
"sglang.srt.managers.",
|
||||
"torch_npu.",
|
||||
("_codecs", "encode"),
|
||||
# --- PyTorch data containers & rebuild functions ---
|
||||
# Code-module prefixes (torch.*, sglang.srt.*) are NOT allowed: they
|
||||
# contain gadgets like sglang.srt.utils.common.dynamic_import
|
||||
("torch", "Tensor"),
|
||||
("torch", "BFloat16Tensor"),
|
||||
("torch", "BoolTensor"),
|
||||
("torch", "ByteTensor"),
|
||||
("torch", "CharTensor"),
|
||||
("torch", "DoubleTensor"),
|
||||
("torch", "FloatTensor"),
|
||||
("torch", "HalfTensor"),
|
||||
("torch", "IntTensor"),
|
||||
("torch", "LongTensor"),
|
||||
("torch", "ShortTensor"),
|
||||
("torch.cuda", "BFloat16Tensor"),
|
||||
("torch.cuda", "BoolTensor"),
|
||||
("torch.cuda", "ByteTensor"),
|
||||
("torch.cuda", "CharTensor"),
|
||||
("torch.cuda", "DoubleTensor"),
|
||||
("torch.cuda", "FloatTensor"),
|
||||
("torch.cuda", "HalfTensor"),
|
||||
("torch.cuda", "IntTensor"),
|
||||
("torch.cuda", "LongTensor"),
|
||||
("torch.cuda", "ShortTensor"),
|
||||
("torch.cuda.sparse", "BFloat16Tensor"),
|
||||
("torch.cuda.sparse", "ByteTensor"),
|
||||
("torch.cuda.sparse", "CharTensor"),
|
||||
("torch.cuda.sparse", "DoubleTensor"),
|
||||
("torch.cuda.sparse", "FloatTensor"),
|
||||
("torch.cuda.sparse", "HalfTensor"),
|
||||
("torch.cuda.sparse", "IntTensor"),
|
||||
("torch.cuda.sparse", "LongTensor"),
|
||||
("torch.cuda.sparse", "ShortTensor"),
|
||||
("torch.sparse", "BFloat16Tensor"),
|
||||
("torch.sparse", "ByteTensor"),
|
||||
("torch.sparse", "CharTensor"),
|
||||
("torch.sparse", "DoubleTensor"),
|
||||
("torch.sparse", "FloatTensor"),
|
||||
("torch.sparse", "HalfTensor"),
|
||||
("torch.sparse", "IntTensor"),
|
||||
("torch.sparse", "LongTensor"),
|
||||
("torch.sparse", "ShortTensor"),
|
||||
("torch", "Size"),
|
||||
("torch", "device"),
|
||||
("torch", "dtype"),
|
||||
("torch", "bfloat16"),
|
||||
("torch", "bit"),
|
||||
("torch", "bits16"),
|
||||
("torch", "bits1x8"),
|
||||
("torch", "bits2x4"),
|
||||
("torch", "bits4x2"),
|
||||
("torch", "bits8"),
|
||||
("torch", "bool"),
|
||||
("torch", "cdouble"),
|
||||
("torch", "cfloat"),
|
||||
("torch", "chalf"),
|
||||
("torch", "complex128"),
|
||||
("torch", "complex32"),
|
||||
("torch", "complex64"),
|
||||
("torch", "double"),
|
||||
("torch", "float"),
|
||||
("torch", "float16"),
|
||||
("torch", "float32"),
|
||||
("torch", "float4_e2m1fn_x2"),
|
||||
("torch", "float64"),
|
||||
("torch", "float8_e4m3fn"),
|
||||
("torch", "float8_e4m3fnuz"),
|
||||
("torch", "float8_e5m2"),
|
||||
("torch", "float8_e5m2fnuz"),
|
||||
("torch", "float8_e8m0fnu"),
|
||||
("torch", "half"),
|
||||
("torch", "int"),
|
||||
("torch", "int1"),
|
||||
("torch", "int16"),
|
||||
("torch", "int2"),
|
||||
("torch", "int3"),
|
||||
("torch", "int32"),
|
||||
("torch", "int4"),
|
||||
("torch", "int5"),
|
||||
("torch", "int6"),
|
||||
("torch", "int64"),
|
||||
("torch", "int7"),
|
||||
("torch", "int8"),
|
||||
("torch", "long"),
|
||||
("torch", "qint32"),
|
||||
("torch", "qint8"),
|
||||
("torch", "quint2x4"),
|
||||
("torch", "quint4x2"),
|
||||
("torch", "quint8"),
|
||||
("torch", "short"),
|
||||
("torch", "uint1"),
|
||||
("torch", "uint16"),
|
||||
("torch", "uint2"),
|
||||
("torch", "uint3"),
|
||||
("torch", "uint32"),
|
||||
("torch", "uint4"),
|
||||
("torch", "uint5"),
|
||||
("torch", "uint6"),
|
||||
("torch", "uint64"),
|
||||
("torch", "uint7"),
|
||||
("torch", "uint8"),
|
||||
("torch.nn.parameter", "Parameter"),
|
||||
("torch.serialization", "_get_layout"),
|
||||
("torch._utils", "_rebuild_tensor"),
|
||||
("torch._utils", "_rebuild_tensor_v2"),
|
||||
("torch._utils", "_rebuild_tensor_v3"),
|
||||
("torch._utils", "_rebuild_parameter"),
|
||||
("torch._utils", "_rebuild_parameter_with_state"),
|
||||
("torch._utils", "_rebuild_qtensor"),
|
||||
("torch._utils", "_rebuild_sparse_tensor"),
|
||||
("torch._utils", "_rebuild_meta_tensor_no_storage"),
|
||||
("torch._utils", "_rebuild_wrapper_subclass"),
|
||||
("torch._utils", "_rebuild_device_tensor_from_numpy"),
|
||||
("torch._utils", "_rebuild_device_tensor_from_cpu_tensor"),
|
||||
("torch._tensor", "_rebuild_from_type_v2"),
|
||||
("torch.storage", "UntypedStorage"),
|
||||
("torch.storage", "_UntypedStorage"),
|
||||
("torch.storage", "TypedStorage"),
|
||||
("torch", "UntypedStorage"),
|
||||
("torch", "BFloat16Storage"),
|
||||
("torch", "BoolStorage"),
|
||||
("torch", "ByteStorage"),
|
||||
("torch", "CharStorage"),
|
||||
("torch", "ComplexDoubleStorage"),
|
||||
("torch", "ComplexFloatStorage"),
|
||||
("torch", "DoubleStorage"),
|
||||
("torch", "FloatStorage"),
|
||||
("torch", "HalfStorage"),
|
||||
("torch", "IntStorage"),
|
||||
("torch", "LongStorage"),
|
||||
("torch", "QInt32Storage"),
|
||||
("torch", "QInt8Storage"),
|
||||
("torch", "QUInt2x4Storage"),
|
||||
("torch", "QUInt4x2Storage"),
|
||||
("torch", "QUInt8Storage"),
|
||||
("torch", "ShortStorage"),
|
||||
("torch.cuda", "BFloat16Storage"),
|
||||
("torch.cuda", "BoolStorage"),
|
||||
("torch.cuda", "ByteStorage"),
|
||||
("torch.cuda", "CharStorage"),
|
||||
("torch.cuda", "ComplexDoubleStorage"),
|
||||
("torch.cuda", "ComplexFloatStorage"),
|
||||
("torch.cuda", "DoubleStorage"),
|
||||
("torch.cuda", "FloatStorage"),
|
||||
("torch.cuda", "HalfStorage"),
|
||||
("torch.cuda", "IntStorage"),
|
||||
("torch.cuda", "LongStorage"),
|
||||
("torch.cuda", "ShortStorage"),
|
||||
("torch.multiprocessing.reductions", "rebuild_tensor"),
|
||||
("torch.multiprocessing.reductions", "rebuild_meta_tensor"),
|
||||
("torch.multiprocessing.reductions", "rebuild_cuda_tensor"),
|
||||
("sglang.srt.utils.patch_torch", "_rebuild_cuda_tensor_modified"),
|
||||
("torch_npu.multiprocessing.reductions", "rebuild_npu_tensor"),
|
||||
("sglang.srt.utils.patch_torch", "_rebuild_npu_tensor_modified"),
|
||||
("torch.multiprocessing.reductions", "rebuild_nested_tensor"),
|
||||
("torch.multiprocessing.reductions", "rebuild_sparse_coo_tensor"),
|
||||
("torch.multiprocessing.reductions", "rebuild_sparse_compressed_tensor"),
|
||||
("torch.multiprocessing.reductions", "rebuild_storage_fd"),
|
||||
("torch.multiprocessing.reductions", "rebuild_storage_filename"),
|
||||
("torch.multiprocessing.reductions", "rebuild_storage_empty"),
|
||||
("torch.multiprocessing.reductions", "rebuild_typed_storage"),
|
||||
("torch.multiprocessing.reductions", "rebuild_typed_storage_child"),
|
||||
("torch", "per_tensor_affine"),
|
||||
("torch", "per_tensor_symmetric"),
|
||||
("torch", "per_channel_affine"),
|
||||
("torch", "per_channel_symmetric"),
|
||||
("torch", "per_channel_affine_float_qparams"),
|
||||
# --- SGLang data containers only (no code modules) ---
|
||||
("sglang.srt.managers.io_struct", "GenerateReqInput"),
|
||||
("sglang.srt.managers.io_struct", "EmbeddingReqInput"),
|
||||
("sglang.srt.disaggregation.encoder.receiver", "EmbeddingData"),
|
||||
("sglang.srt.managers.schedule_batch", "Modality"),
|
||||
("sglang.srt.weight_sync.tensor_bucket", "FlattenedTensorMetadata"),
|
||||
("sglang.srt.weight_sync.tensor_bucket", "FlattenedTensorBucket"),
|
||||
(
|
||||
"sglang.srt.model_executor.model_runner_components.weight_updater",
|
||||
"LocalSerializedTensor",
|
||||
),
|
||||
("sglang.srt.model_executor.model_runner", "LocalSerializedTensor"),
|
||||
}
|
||||
|
||||
def find_class(self, module, name):
|
||||
if (module, name) in self.ALLOWED_GLOBALS or any(
|
||||
(module + ".").startswith(prefix) for prefix in self.ALLOWED_MODULE_PREFIXES
|
||||
):
|
||||
if (module, name) == ("torch.storage", "_load_from_bytes"):
|
||||
# Torch's helper calls an unrestricted nested torch.load.
|
||||
return _safe_load_torch_storage
|
||||
if (module, name) in self.ALLOWED_GLOBALS:
|
||||
return super().find_class(module, name)
|
||||
|
||||
raise RuntimeError(
|
||||
|
||||
Reference in New Issue
Block a user