Merge remote-tracking branch 'origin/main' into codex/pr5479-hybrid-model

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/energon/base_energon_datamodule.py

@@ -12,15 +12,156 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import logging
-from typing import Any, Literal, Optional
+from typing import Any, Callable, Literal, Optional
 
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.energon import WorkerConfig, get_savable_loader, get_train_dataset
+from megatron.energon import Sample, WorkerConfig, get_savable_loader, get_train_dataset
+from megatron.energon import dataset_config as _energon_dataset_config
+from megatron.energon.epathlib import EPath
+from megatron.energon.flavors.webdataset.default_generic_webdataset import DefaultGenericWebdatasetFactory
+from megatron.energon.typed_converter import JsonParser
 
 
 logger = logging.getLogger(__name__)
 
+_ORIGINAL_METHOD_ATTRIBUTE = "__megatron_bridge_original_method__"
+_TRUSTED_DATASET_FACTORY_MODULE_PREFIXES = (
+    "megatron.bridge.data.energon",
+    "megatron.energon",
+)
+_TRUSTED_DATASET_FACTORY_MODULE_ALIASES = frozenset(
+    {
+        "megatron.bridge.models.qwen_vl.data.energon",
+    }
+)
+_energon_factory_init = getattr(
+    DefaultGenericWebdatasetFactory.__init__,
+    _ORIGINAL_METHOD_ATTRIBUTE,
+    DefaultGenericWebdatasetFactory.__init__,
+)
+_energon_load_config = getattr(
+    _energon_dataset_config.load_config,
+    _ORIGINAL_METHOD_ATTRIBUTE,
+    _energon_dataset_config.load_config,
+)
+
+
+def _validate_energon_dataset_metadata(value: Any, *, root: bool = True) -> None:
+    """Reject executable object references in untrusted dataset factory metadata."""
+    if isinstance(value, list):
+        for item in value:
+            _validate_energon_dataset_metadata(item, root=False)
+        return
+    if not isinstance(value, dict):
+        return
+
+    module_name = value.get("__module__")
+    function_name = value.get("__function__")
+    class_name = value.get("__class__")
+    if function_name is not None:
+        raise ValueError(
+            "Energon dataset metadata cannot resolve serialized Python functions. "
+            "Use declarative configuration or pass a Python callable from trusted application code."
+        )
+    if class_name is not None:
+        trusted_class = (
+            isinstance(module_name, str)
+            and isinstance(class_name, str)
+            and (
+                module_name in _TRUSTED_DATASET_FACTORY_MODULE_ALIASES
+                or any(
+                    module_name == prefix or module_name.startswith(f"{prefix}.")
+                    for prefix in _TRUSTED_DATASET_FACTORY_MODULE_PREFIXES
+                )
+            )
+        )
+        if trusted_class:
+            module = importlib.import_module(module_name)
+            referenced_type = getattr(module, class_name, None)
+            required_base = DefaultGenericWebdatasetFactory if root else Sample
+            trusted_class = isinstance(referenced_type, type) and issubclass(referenced_type, required_base)
+        if not trusted_class:
+            raise ValueError(
+                "Energon dataset metadata cannot instantiate serialized Python classes. "
+                "Only packaged Energon dataset factory and sample classes are allowed."
+            )
+
+    for item in value.values():
+        _validate_energon_dataset_metadata(item, root=False)
+
+
+def _secure_energon_load_config(
+    path: EPath | dict[str, Any],
+    *,
+    default_type: type,
+    default_kwargs: dict[str, Any] | None = None,
+    parser: JsonParser = JsonParser(strict=True),
+) -> Any:
+    """Validate dataset metadata before Energon resolves serialized objects."""
+    is_dataset_factory = isinstance(default_type, type) and issubclass(default_type, DefaultGenericWebdatasetFactory)
+    if not is_dataset_factory:
+        return _energon_load_config(
+            path,
+            default_type=default_type,
+            default_kwargs=default_kwargs,
+            parser=parser,
+        )
+
+    if isinstance(path, dict):
+        data = path
+    else:
+        with path.open("rb") as config_file:
+            data = _energon_dataset_config.load_yaml(config_file)
+    _validate_energon_dataset_metadata(data)
+    return _energon_load_config(
+        data,
+        default_type=default_type,
+        default_kwargs=default_kwargs,
+        parser=parser,
+    )
+
+
+def _secure_energon_factory_init(
+    self: DefaultGenericWebdatasetFactory,
+    path: EPath,
+    *,
+    subflavors: dict[str, Any] | None = None,
+    field_map: dict[str, str] | None = None,
+    sample_loader: str | Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    part_filter: str | list[str] | Callable[[str], bool] | None = None,
+    **kwargs: Any,
+) -> None:
+    """Reject dataset-local Python hooks before Energon resolves their files."""
+    executable_fields = [
+        name
+        for name, value in (("sample_loader", sample_loader), ("part_filter", part_filter))
+        if isinstance(value, str)
+    ]
+    if executable_fields:
+        raise ValueError(
+            "Energon dataset metadata cannot load Python files through "
+            f"{', '.join(executable_fields)}. Use a declarative field_map instead."
+        )
+    _energon_factory_init(
+        self,
+        path,
+        subflavors=subflavors,
+        field_map=field_map,
+        sample_loader=sample_loader,
+        part_filter=part_filter,
+        **kwargs,
+    )
+
+
+# Energon constructs this factory internally after reading dataset.yaml, so
+# Bridge must install the guard before calling get_train_dataset().
+setattr(_secure_energon_load_config, _ORIGINAL_METHOD_ATTRIBUTE, _energon_load_config)
+_energon_dataset_config.load_config = _secure_energon_load_config
+setattr(_secure_energon_factory_init, _ORIGINAL_METHOD_ATTRIBUTE, _energon_factory_init)
+DefaultGenericWebdatasetFactory.__init__ = _secure_energon_factory_init
+
 
 class EnergonMultiModalDataModule:
     """

src/megatron/bridge/diffusion/data/common/diffusion_task_encoder_with_sp.py

@@ -24,6 +24,7 @@
 from megatron.bridge.data.energon.metadata import sample_metadata_kwargs
 from megatron.bridge.data.packing.algorithms import first_fit_decreasing
 from megatron.bridge.diffusion.data.common.diffusion_sample import DiffusionSample
+from megatron.bridge.diffusion.data.common.safe_decoder import SafeDiffusionSampleDecoder
 from megatron.bridge.diffusion.data.common.sequence_packing_utils import packing_length
 
 
@@ -51,6 +52,8 @@ def cook(sample: dict) -> dict:
 
 
 class DiffusionTaskEncoderWithSequencePacking(DefaultTaskEncoder, ABC):  # noqa: D101
+    decoder = SafeDiffusionSampleDecoder()
+
     cookers = [
         Cooker(cook),
     ]

src/megatron/bridge/diffusion/data/common/safe_decoder.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Safe decoding for Bridge diffusion WebDataset fields."""
+
+import io
+import re
+from typing import Any
+
+import torch
+from megatron.energon.flavors.webdataset.sample_decoder import SampleDecoder
+from webdataset.autodecode import basichandlers
+
+from megatron.bridge.utils.safe_pickle import safe_torch_tensor_pickle_loads
+
+
+_PICKLE_EXTENSIONS = frozenset({"pickle", "pkl", "pyd"})
+
+
+def _safe_basic_handler(key: str, data: bytes) -> Any:
+    """Decode tensor fields safely and delegate non-executable formats."""
+    extension = re.sub(r".*[.]", "", key).lower()
+    if extension in _PICKLE_EXTENSIONS:
+        return safe_torch_tensor_pickle_loads(data)
+    if extension == "pth":
+        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+    return basichandlers(key, data)
+
+
+class SafeDiffusionSampleDecoder(SampleDecoder):
+    """Energon decoder that never uses WebDataset's unrestricted pickle handlers."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        replaced = False
+        for index, handler in enumerate(self._decoder.handlers):
+            if handler is basichandlers:
+                self._decoder.handlers[index] = _safe_basic_handler
+                replaced = True
+        if not replaced:
+            raise RuntimeError("Energon SampleDecoder no longer exposes the expected WebDataset basic handler.")

src/megatron/bridge/utils/instantiate_utils.py

@@ -105,6 +105,16 @@
     },
 }
 
+# Resolve aliases before enforcing these entries. For example,
+# ``numpy.lib.npyio.load`` is the same callable as ``numpy.load`` and
+# ``torch.serialization.load`` is the same callable as ``torch.load``.
+_DISALLOWED_CANONICAL_TARGETS: set[str] = {
+    "numpy.load",
+    "torch.serialization.load",
+    "transformers.dynamic_module_utils.get_class_in_module",
+    "transformers.dynamic_module_utils.get_class_from_dynamic_module",
+}
+
 _DISALLOWED_CALLABLE_FIELD_NAMES: set[str] = {
     "collate_impl",
     "hf_filter_lambda",
@@ -182,7 +192,19 @@ def _resolve_target(
     """Resolve target string, type, or callable after Bridge validation."""
     if isinstance(target, str):
         _validate_target_prefix(target=target, full_key=full_key)
-    return _mcore_resolve_target(target, full_key, check_callable)
+    resolved_target = _mcore_resolve_target(target, full_key, check_callable)
+    if isinstance(target, str):
+        module = getattr(resolved_target, "__module__", None)
+        qualname = getattr(resolved_target, "__qualname__", None)
+        canonical_target = f"{module}.{qualname}" if isinstance(module, str) and isinstance(qualname, str) else None
+        if canonical_target in _DISALLOWED_CANONICAL_TARGETS:
+            raise InstantiationException(
+                f"Instantiation of '{target}' is not allowed because it resolves to the unsafe target "
+                f"'{canonical_target}'." + (f"\nfull_key: {full_key}" if full_key else "")
+            )
+        if canonical_target is not None:
+            _validate_target_prefix(target=canonical_target, full_key=full_key)
+    return resolved_target
 
 
 _mcore_instantiate_utils._resolve_target = _resolve_target

src/megatron/bridge/utils/safe_pickle.py

@@ -239,6 +239,32 @@ def find_class(self, module: str, name: str) -> object:
         )
 
 
+def _safe_torch_load_from_bytes(data: bytes) -> object:
+    """Reconstruct tensor storage bytes without enabling arbitrary pickle globals."""
+    import torch
+
+    return torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+
+
+class _TorchTensorRestrictedUnpickler(_RestrictedUnpickler):
+    """Restricted unpickler for plain tensors and containers of tensors."""
+
+    _SAFE_MODULES = MappingProxyType(
+        {
+            **_RestrictedUnpickler._SAFE_MODULES,
+            "torch._utils": frozenset({"_rebuild_tensor", "_rebuild_tensor_v2"}),
+        }
+    )
+
+    def find_class(self, module: str, name: str) -> object:
+        if (module, name) in {
+            ("megatron.core.safe_globals", "safe_load_from_bytes"),
+            ("torch.storage", "_load_from_bytes"),
+        }:
+            return _safe_torch_load_from_bytes
+        return super().find_class(module, name)
+
+
 class _NumpyRestrictedUnpickler(pickle.Unpickler):
     """Unpickler that allows safe builtins and the narrow set of numpy types needed for object array reconstruction.
 
@@ -428,6 +454,11 @@ def safe_pickle_loads(data: bytes) -> object:
     return _RestrictedUnpickler(io.BytesIO(data)).load()
 
 
+def safe_torch_tensor_pickle_loads(data: bytes) -> object:
+    """Deserialize raw pickle data containing only safe containers and plain torch tensors."""
+    return _TorchTensorRestrictedUnpickler(io.BytesIO(data)).load()
+
+
 def safe_load_npy(data: bytes):
     """Load a ``.npy`` file from raw bytes without enabling unrestricted pickle.