Safe-by-default model loading for torch models (weights_only=True) by hackchang · Pull Request #3183 · unit8co/darts

hackchang · 2026-08-19T14:14:45Z

What

Makes torch-model checkpoint loading safe-by-default against CWE-502 (a malicious .ckpt executing arbitrary code during unpickling), while keeping legitimate models loading unchanged.

Registers Darts' own checkpoint-serialized types (likelihoods + LikelihoodType, optimizer/scheduler/loss/activation classes, torchmetrics metric classes + their tensor-reduction helpers) via torch.serialization.add_safe_globals, so a weights_only=True load of a legitimate checkpoint succeeds. Only class/type and argument-less library helpers are allow-listed — never callables that run on unpickle — so this does not re-open the RCE surface. No-op on torch < 2.6.
Lightning-checkpoint paths (_DartsCheckpointIO, load_from_checkpoint, load_weights, load_weights_from_checkpoint) now default to weights_only=True, with a user-facing weights_only parameter to opt back into full unpickling for trusted files.
The full-object .pt path in load() (an arbitrary object graph that cannot use weights_only=True) keeps weights_only=False by default but exposes the parameter and a security note.

Behavior changes (please note)

load_from_checkpoint(), load_weights(), load_weights_from_checkpoint() now default to weights_only=True. Checkpoints using custom loss/optimizer/likelihood classes defined outside darts/torch may need weights_only=False.
load_weights_from_checkpoint() previously raised on weights_only=True; it now works and is the default.

Testing

New security regression test (test_load_checkpoint_weights_only_blocks_malicious_payload): a normal model still loads under the safe default; a crafted __reduce__ payload is refused under weights_only=True (no code execution) and only loads under the explicit weights_only=False opt-out.
Full test_torch_forecasting_model.py suite passes locally (350 passed) on torch 2.13 / lightning 2.6.5.

Transparency

This contribution was prepared with AI assistance; I reviewed, tested, and take responsibility for every change, and remain in the loop throughout, per the project's guidelines.

dennisbader · 2026-08-20T09:10:28Z

Hi @hackchang not sure whether you've seen it but the checks failed (linting and unit tests). Could you fix the issues? After that I can review

hackchang · 2026-08-20T11:08:57Z

Thanks @dennisbader — fixed both: ran ruff format for the lint failure, and the test failures were Darts' NeuralForecast-wrapped models (neuralforecast.models.*) not being in the weights_only=True allow-list — now registered when neuralforecast is installed. Verified the previously-failing save/load + nf_model tests pass locally. Pushed; CI should be green now.

codecov · 2026-08-20T11:36:57Z

Codecov Report

❌ Patch coverage is 96.49123% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.05%. Comparing base (3d4c9e7) to head (631ea58).

Files with missing lines	Patch %	Lines
...arts/models/forecasting/torch_forecasting_model.py	96.49%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3183      +/-   ##
==========================================
- Coverage   97.11%   97.05%   -0.07%     
==========================================
  Files         167      167              
  Lines       18429    18480      +51     
==========================================
+ Hits        17898    17936      +38     
- Misses        531      544      +13

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

dennisbader

Thanks a lot @hackchang for this PR 🚀

The direction makes sense and looks like a good start. I do have some comments / suggestions:

The .ckpt default to weights_only=True is good
The import-time allow-list of half of torch/torchmetrics is not ideal: too much maintenance, runs on every import, and we'll miss edge cases (custom losses, new metric subpackages).
Most loading paths still fully unpickle the .pt files so the security fix is narrower than the description suggests.

I'd prefer a slightly reduced PR: with flipping weights_only defaults, adding docs/warnings, minimal Darts-only allow-list, and maybe checkpoint-driven registration on failure.

dennisbader · 2026-09-07T09:28:28Z

+        # ``torchmetrics`` metrics store references to their own tensor-reduction /
+        # distributed helper *functions* in their pickled state (``_reductions``,
+        # ``dist_sync_fn``, ...). A ``weights_only=True`` load of a metric-bearing
+        # checkpoint needs these specific, known-safe library callables allow-listed.


We're allow-listing functions here, not just classes. PyTorch's serialization docs say registered functions can be called during unpickling. I don't think that's equivalent to "does not re-open the RCE surface" — worth being honest about in the PR description.

dennisbader · 2026-09-07T09:28:34Z


-class _DartsCheckpointIO(TorchCheckpointIO):
-    """Custom CheckpointIO that defaults ``weights_only`` to ``False``.
+def _register_safe_globals() -> None:


This runs at import time for every torch model. That means we unconditionally import torchmetrics (and optionally neuralforecast) even if the user never touches checkpoint loading. Could we defer this to first load, or use the safe_globals context manager scoped around the actual torch.load call?

Also, scanning vars(module) for every loss/activation/optimizer/metric subclass feels brittle — torchmetrics adds subpackages over time and we'll keep chasing them. Have you tried get_unsafe_globals_in_checkpoint() on the specific file being loaded?

dennisbader · 2026-09-07T09:29:57Z

        if _PL_2_6_OR_ABOVE:
            weights_only_kwargs["weights_only"] = (
-                False if weights_only is None else weights_only
+                True if weights_only is None else weights_only


Flipping the default from False to True is a behavior change for all Lightning-internal checkpoint loads, not just our public APIs. fit(ckpt_path=...) still passes weights_only=False explicitly, but what about other trainer paths that go through the plugin without overriding? Worth a test for resume-from-checkpoint with optimizer state.

dennisbader · 2026-09-07T09:32:24Z

        if _PL_2_6_OR_ABOVE:
            weights_only_kwargs["weights_only"] = (
-                False if weights_only is None else weights_only
+                True if weights_only is None else weights_only


Comment on lines 772-774 still says we force weights_only=False for internal loading, but _DartsCheckpointIO now defaults to True. Stale, or did I miss something?

dennisbader · 2026-09-07T09:34:10Z

            base_model_path, weights_only=False, map_location=kwargs.get("map_location")
        )

        # load PyTorch LightningModule from checkpoint


Just flagging for the PR description: even with this change, load_from_checkpoint still fully unpickles the base model file first. So the "safe-by-default against CWE-502" claim only really applies to the .ckpt half. The docstring warning helps but the top-level description oversells it.

dennisbader · 2026-09-07T09:35:12Z

        """
        Loads a model from a given file path.

+        .. warning::


Fine for backwards compat, but means RNNModel.load("foo.pt") is still arbitrary code execution on a malicious file. If we're closing #3177 with this PR, we should be clear in the issue / CHANGELOG that .pt loads are out of scope.

dennisbader · 2026-09-07T09:36:43Z

+        reloaded.load_weights(ckpt_path, map_location="cpu")
+        reloaded.predict(n=2, series=self.series[:20])
+
+        # 2) craft a malicious checkpoint whose `__reduce__` writes a marker file --------


Part 1 (legitimate model loads with default) is good. Part 2 only tests raw torch.load though — can you add a case where the evil payload is loaded via reloaded.load_weights(evil_path) or similar? That's the path users actually hit.

dennisbader · 2026-09-07T09:41:43Z

+        )
+        # cleanup the benign marker
+        os.remove(marker_path)
+


On test_load_weights_from_checkpoint (~line 1344–1353)

This expects ValueError when weights_only=True, but you removed that guard and made True the default -> the the load_weights_from_checkpoint above should then fail now. Did you mean to update this block (now it's probably handled via the safe globals)?

@dennisbader

) Address @dennisbader's review: - Drop the import-time `add_safe_globals` scan of torch/torchmetrics/ neuralforecast. Allow-listing is now scoped to each load via `torch.serialization.safe_globals(...)` (`_load_ckpt_safely`) -- nothing is imported or registered unless a checkpoint is actually loaded. - Checkpoint-driven, per file: use `get_unsafe_globals_in_checkpoint()` and only allow-list the globals that file references, filtered to classes under torch/torchmetrics/darts/neuralforecast that subclass a known-safe base (nn.Module / Optimizer / LR scheduler / torchmetrics Metric|MetricCollection / Darts likelihood). Anything else stays blocked and the load fails loudly. - Stop scanning modules for functions. torchmetrics genuinely pickles references to its own reduction helpers, so allow-list only an explicit, audited set of 6 such functions by exact qualified name (added only when referenced), and document that registering a callable is not zero-risk. - Honesty: docstrings + CHANGELOG state this protects the `.ckpt` only; the Darts base model `.pt` is still fully unpickled (a malicious `.pt` is still RCE and is out of scope). Marked the change with the breaking-change flag and unit8co#3177 as "partially addresses". - Fix the now-stale internal-loading comment (internal loads are safe-by-default `weights_only=True`; only the training-resume path stays explicit `False`). - Tests: fix the now-stale `weights_only=True` block (it was actually catching an architecture mismatch); add a malicious-payload case through the real user path `load_weights_from_checkpoint(skip_checks=True)`; add a resume-from-checkpoint test that exercises optimizer/scheduler state under `weights_only=True`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015RFWdogyAg3fWKPv9pxXAd

hackchang · 2026-09-08T23:29:39Z

Thanks for the detailed review — this was the right call. Pushed a reduced version in 2888876:

No import-time registration. The old add_safe_globals scan (torch/torchmetrics/
neuralforecast) is gone. Allow-listing is now scoped to each load via
torch.serialization.safe_globals(...) — nothing is imported or registered unless a
checkpoint is actually loaded.
Checkpoint-driven, per file. Using get_unsafe_globals_in_checkpoint(path) I only
allow-list the globals that specific file references, filtered to classes under
torch/torchmetrics/darts/neuralforecast that subclass a known-safe base (nn.Module /
Optimizer / LR scheduler / torchmetrics Metric|MetricCollection / Darts likelihood).
Anything else stays blocked and the load fails loudly.
Functions: you're right. No more module scans. torchmetrics does pickle references to
its own reduction helpers (e.g. dim_zero_sum, jit_distributed_available), so I allow-list
only an explicit, audited set of 6 such functions by exact qualified name, added only when
the checkpoint references them — and I documented that registering a callable is not zero-risk.
Honesty on scope. Docstrings + CHANGELOG now state this protects the .ckpt only; the
Darts base .pt is still fully unpickled, so a malicious .pt is still RCE and is out of
scope. Marked the change 🔴 and set Safe-by-default model loading #3177 to "partially addresses".
Tests. Fixed the now-stale weights_only=True block (its ValueError was actually an
architecture mismatch), added a malicious-payload case through the real user path
(load_weights_from_checkpoint), and added a resume-from-checkpoint test that exercises
optimizer/scheduler state under weights_only=True. The torch checkpoint/save-load suite
passes locally.

Happy to trim further to a Darts-classes-only allow-list (requiring weights_only=False for
metric/optimizer-bearing checkpoints) if you'd prefer even less surface — let me know.

…/hackchang/safe-model-loading

hackchang added 2 commits August 19, 2026 23:12

Add files via upload

f8c5372

Add files via upload

42f9d6b

hackchang requested a review from dennisbader as a code owner August 19, 2026 14:14

hackchang added 2 commits August 19, 2026 23:18

Add files via upload

db1d679

Add files via upload

8c9cd26

hackchang added 2 commits August 20, 2026 20:07

Add files via upload

76e1576

Add files via upload

bc2e1e3

hackchang and others added 4 commits August 20, 2026 23:47

Merge branch 'master' into safe-model-loading

b5790f8

Merge branch 'master' into safe-model-loading

a59df20

Merge branch 'master' into fork/hackchang/safe-model-loading

c616335

Merge branch 'master' into safe-model-loading

631ea58

dennisbader requested changes Sep 7, 2026

View reviewed changes

hackchang requested a review from dennisbader September 8, 2026 23:57

dennisbader added 2 commits September 21, 2026 14:29

Merge remote-tracking branch 'hackchang/safe-model-loading' into fork…

538d26b

…/hackchang/safe-model-loading

Merge branch 'master' into fork/hackchang/safe-model-loading

8974d63

Conversation

hackchang commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Behavior changes (please note)

Testing

Transparency

Uh oh!

dennisbader commented Aug 20, 2026

Uh oh!

hackchang commented Aug 20, 2026

Uh oh!

codecov Bot commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

dennisbader left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hackchang commented Sep 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hackchang commented Aug 19, 2026 •

edited

Loading

codecov Bot commented Aug 20, 2026 •

edited

Loading