fix(weights): change how models are serialized and deserialized by ashwinvaidya17 · Pull Request #3801 · open-edge-platform/anomalib

ashwinvaidya17 · 2026-09-17T10:58:13Z

📝 Description

Addresses PTK0009203
Currently anomalib engine loads checkpoints using weights_only=False. This will end up executing malicious coded if the checkpoint comes from untrusted sources. This PR proposes changes to AnomalibModule to mitigate this.
An alternative to this design is to use TRUST_REMOTE_CODE=True but this might make the users set this is their default env variable which is risky.

Design

This design introduces a scoped allowlist used at both checkpoint load seams:

Shared first-party enums (ANOMALIB_SAFE_GLOBALS)
Shared numpy leaf types (NUMPY_SAFE_GLOBALS)
Optional per-model extras via AnomalibModule.checkpoint_safe_globals()

In Anomalib Module

Override

@classmethod
 def load_from_checkpoint(
      ...
  ) -> dict[str, Any]:
    with anomalib_safe_globals(extra=cls.checkpoint_safe_globals()):
            return super().load_from_checkpoint(

Introduce

@classmethod
def checkpoint_safe_globals(cls) -> Sequence[Any]:
    """Model-specific types persisted in ``hyper_parameters``.
    """
    return ()

Model specific globals live here.

Add new plugin for model checkpoint AnomalibCheckpointIO

Single context manager, layered allowlist

anomalib.utils.serialization.anomalib_safe_globals
merges three layers for the duration of a load:

ANOMALIB_SAFE_GLOBALS  +  NUMPY_SAFE_GLOBALS  +  extra (optional)
        │                         │                      │
        ▼                         ▼                      ▼
  PrecisionType            numpy scalar/ndarray     model enums
                           (scheduler / opt state)

with anomalib_safe_globals(extra=cls.checkpoint_safe_globals()):
    torch.load(..., weights_only=True)  # via Lightning / CheckpointIO

Rules for what may be added:

Shared anomalib list: first-party enums reused across models.
Numpy list: inert reconstruction helpers / dtypes required for scheduler and
optimizer state under NumPy 2.x (numpy._core). Both numpy.core and
numpy._core pickle path aliases are registered so older pickles still resolve.
The project’s dependency floor is NumPy ≥ 2 (via opencv-python-headless), so
importing numpy._core.multiarray is intentional.
Model extras: small first-party types persisted only by that model’s
constructor hyperparameters. Not shared third-party types, not modules, not
pathlib.Path.

Dual load seams

Checkpoint restore happens in two places; both must use the same allowlist policy.

AnomalibModule.load_from_checkpoint
        │
        └─ anomalib_safe_globals(extra=cls.checkpoint_safe_globals())

Engine → AnomalibCheckpointIO.load_checkpoint
        │
        └─ anomalib_safe_globals(extra=self.extra_safe_globals)
              ▲
              │ Engine._ensure_checkpoint_io_plugin(model)
              │ copies model.checkpoint_safe_globals()

Direct API / export: AnomalibModule.load_from_checkpoint wraps Lightning’s
loader with extra=cls.checkpoint_safe_globals().
Trainer ckpt_path: AnomalibCheckpointIO wraps TorchCheckpointIO. The
Engine installs it when the user has not supplied a CheckpointIO, and refreshes
extra_safe_globals from the active model before fit / validate / test / predict.

Caller code stays unchanged:

model = VlmAd.load_from_checkpoint("model.ckpt")
engine.test(VlmAd(), ckpt_path="model.ckpt", datamodule=datamodule)

Per-model extension point

class AnomalibModule(...):
    @classmethod
    def checkpoint_safe_globals(cls) -> Sequence[Any]:
        return ()

class VlmAd(AnomalibModule):
    @classmethod
    def checkpoint_safe_globals(cls) -> Sequence[Any]:
        return (ModelName,)

✨ Changes

Select what type of change your PR is:

✅ Checklist

Before you submit your pull request, please make sure you have completed the following steps:

📚 I have made the necessary updates to the documentation (if applicable).
🧪 I have written tests that support my changes and prove that my fix is effective or my feature works (if applicable).
🏷️ My PR title follows conventional commit format.

For more information about code review checklists, see the Code Review Checklist.

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

🟡 Changes recommended

The security goal is not reliably enforced yet (weights_only defaults still defer to upstream defaults), and there are import/compatibility and repo-header convention issues that should be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR hardens anomalib checkpoint loading by introducing a scoped allowlist mechanism intended to support torch.load(..., weights_only=True) without executing arbitrary code from untrusted checkpoints, applying the policy both to direct AnomalibModule.load_from_checkpoint usage and Trainer ckpt_path restores via a custom CheckpointIO plugin.

Changes:

Added anomalib_safe_globals context manager plus shared allowlists for first-party enums and NumPy leaf types used during optimizer/scheduler state restore.
Added an AnomalibCheckpointIO plugin and Engine wiring to ensure Trainer restores use the same allowlist policy.
Added per-model checkpoint_safe_globals() overrides for models that persist model-local enums in hyper_parameters.

File summaries

File	Description
src/anomalib/utils/serialization.py	New safe-globals allowlist/context manager used to constrain checkpoint unpickling under `weights_only`.
src/anomalib/models/components/base/anomalib_module.py	Adds `checkpoint_safe_globals()` extension point and wraps `load_from_checkpoint` with the allowlist context.
src/anomalib/engine/plugins/checkpoint_io.py	New `TorchCheckpointIO` subclass that wraps checkpoint loads in the allowlist context.
src/anomalib/engine/plugins/init.py	Exposes the new `AnomalibCheckpointIO` plugin.
src/anomalib/engine/engine.py	Installs/refreshes the CheckpointIO plugin automatically and removes explicit `weights_only=False` on Trainer calls.
src/anomalib/models/image/vlm_ad/lightning_model.py	Adds model-specific `checkpoint_safe_globals()` to allowlist `ModelName`.
src/anomalib/models/image/reverse_distillation/lightning_model.py	Adds model-specific `checkpoint_safe_globals()` to allowlist `AnomalyMapGenerationMode`.
src/anomalib/models/image/efficient_ad/lightning_model.py	Adds model-specific `checkpoint_safe_globals()` to allowlist `EfficientAdModelSize`.
src/anomalib/models/image/dfkde/lightning_model.py	Adds model-specific `checkpoint_safe_globals()` to allowlist `FeatureScalingMethod`.

Review details

Files reviewed: 9/9 changed files
Comments generated: 8
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

rajeshgangireddy

Minor comments.

Can you also make sure tiling checkpoints are also loaded okay with this this new changes. get_ensemble_model() Tiling loading passes a preprocessor and i feel this might fail
Some files might need year update in copyright headers.
Do we need a similar weight_only plugin for TorchInferencer as well ? It still uses torch.load(path, map_location=self.device, weights_only=False) # nosec B614

rajeshgangireddy · 2026-09-18T14:51:17Z

+        self,
+        path: _PATH,
+        map_location: Callable | None = lambda storage, _loc: storage,
+        weights_only: bool | None = None,


similar to the copilot's comment.
I think it might be better to explicitly enforce weights_only=True by default
here and in
anomalib/src/anomalib/models/components/base/anomalib_module.py

Lightning might change None to False for HTTP/HTTPS checkpoint URLS?

https://docs.pytorch.org/docs/main/notes/serialization.html#id3

- fix pickle for some models - address tiling based models Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

One or more issues must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 5

Open (5)

Preserve or document custom component configurations across reloads · New
Support Path hyperparameters with weights-only checkpoint loading · New
Provide a legacy loading path for existing checkpoints · New
Restrict transform deserialization to an explicit safe registry · New
Serialize Path hyperparameters as safe plain data · New

Resolved since last review (8)

np.dtypes.Float64DType is a NumPy 2.x-specific symbol; anomalib doesn't declare a numpy>=2…
The PR intent is to mitigate untrusted checkpoint RCE by using weights_only=True, but this…
AnomalibCheckpointIO is intended to enable safe checkpoint loading, but it currently forwards…
Avoid importing from NumPy's private numpy._core package at module import time; this path is…
This file was modified in 2026 (new checkpoint_safe_globals method), but the copyright header is…
This file was modified in 2026 (new checkpoint_safe_globals method), but the copyright header…
This file was modified in 2026 (new checkpoint_safe_globals method), but the copyright header…
load_from_checkpoint is a public API; leaving **kwargs untyped weakens typing and makes the…

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved checkpoint compatibility, serialization, and checkpoint-IO integration issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 6

Open (6)

Normalize Path hyperparameters before checkpoint serialization · New
Register emitted transform class paths · New
Serialize Path hyperparameters as safe plain data
Provide a legacy loading path for existing checkpoints
Support Path hyperparameters with weights-only checkpoint loading
Preserve or document custom component configurations across reloads

Resolved since last review (1)

Restrict transform deserialization to an explicit safe registry

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

Critical checkpoint-save failures and unresolved restore and compatibility issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 5 · 1

Open (6)

Handle custom preprocessors without transform safely · New
Support SquarePad in preprocessor serialization · New
Serialize Path hyperparameters as safe plain data
Provide a legacy loading path for existing checkpoints
Preserve or document custom component configurations across reloads
Document safe-load default and migration requirements · New

Resolved since last review (3)

Register emitted transform class paths
Normalize Path hyperparameters before checkpoint serialization
Support Path hyperparameters with weights-only checkpoint loading

Previously missed (1)

In code that hasn't changed since last review

Add Engine-level checkpoint restore test

src/anomalib/engine/engine.py:460

The new tests exercise AnomalibCheckpointIO.load_checkpoint and plugin installation directly, but none drives Lightning's fit/validate/test/predict with ckpt_path through these changed Engine calls. Add one deterministic Engine-level restore test to catch plugin wiring or positional-argument regressions.

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved critical and moderate findings affect checkpoint restoration, compatibility, and safe plugin integration.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2 · 1

Open (3)

Handle empty transform lists during serialization · New
Decode enum values nested in union annotations · New
Remove network dependencies from model registration tests · New

Resolved since last review (6)

Support SquarePad in preprocessor serialization
Handle custom preprocessors without transform safely
Serialize Path hyperparameters as safe plain data
Provide a legacy loading path for existing checkpoints
Preserve or document custom component configurations across reloads
Document safe-load default and migration requirements

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved checkpoint compatibility, restore integration, component round-trip, and dependency issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 3 · 2

Open (5)

Add Engine restore regression test · New
Persist SuperADD post-processor settings · New
Validate or support all torchvision transforms · New
Update copyright range to 2026 · New
Update copyright header to 2026 · New

Resolved since last review (3)

Decode enum values nested in union annotations
Handle empty transform lists during serialization
Remove network dependencies from model registration tests

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot

Copilot review overview

🟡 Changes recommended

Component-state regressions, skipped subclass restoration hooks, legacy compatibility, and unsupported NumPy environments must be addressed.

Review effort: Balanced
Findings: 1 · 2

Open (3)

Preserve disabled component flags during checkpoint round-trips · New
Preserve Lightning's full map_location type contract · New
Document both TypeError and ValueError validation failures · New

Resolved since last review (5)

Validate or support all torchvision transforms
Persist SuperADD post-processor settings
Add Engine restore regression test

-        self.save_hyperparameters()
+        # Components are restored from safe checkpoint data below; evaluator and
+        # visualizer configuration is intentionally rebuilt from model defaults.
+        self.save_hyperparameters(ignore=["pre_processor", "post_processor", "evaluator", "visualizer"])


+    def load_checkpoint(
+        self,
+        path: _PATH,
+        map_location: Callable | None = lambda storage, _loc: storage,


+    Raises:
+        ValueError: If the specification is malformed, refers to an
+            unregistered transform, or has invalid constructor arguments.


Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Add new design

b6f33ac

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

ashwinvaidya17 requested review from rajeshgangireddy and samet-akcay as code owners September 17, 2026 10:58

Copilot AI lite review requested due to automatic review settings September 17, 2026 10:58

Copilot started reviewing on behalf of ashwinvaidya17 September 17, 2026 10:58 View session

Copilot AI reviewed Sep 17, 2026

View reviewed changes

samet-akcay previously approved these changes Sep 18, 2026

View reviewed changes

rajeshgangireddy reviewed Sep 18, 2026

View reviewed changes

- add tests

5246bba

- fix pickle for some models - address tiling based models Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot AI review requested due to automatic review settings September 23, 2026 08:22

ashwinvaidya17 dismissed samet-akcay’s stale review via 5246bba September 23, 2026 08:22

remove file

c1f9453

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot started reviewing on behalf of ashwinvaidya17 September 23, 2026 08:25 View session

github-advanced-security Bot found potential problems Sep 23, 2026

View reviewed changes

address security bot

4a9dd0a

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot AI reviewed Sep 23, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 23, 2026 08:36

Copilot started reviewing on behalf of ashwinvaidya17 September 23, 2026 08:37 View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

Comment thread src/anomalib/models/components/base/anomalib_module.py

Comment thread src/anomalib/pre_processing/utils/spec.py Outdated

address copilot comments

b23abf3

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot AI review requested due to automatic review settings September 23, 2026 09:13

Copilot started reviewing on behalf of ashwinvaidya17 September 23, 2026 09:14 View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

Comment thread src/anomalib/models/components/base/anomalib_module.py Outdated

Comment thread src/anomalib/pre_processing/utils/spec.py Outdated

Comment thread src/anomalib/models/components/base/anomalib_module.py

ashwinvaidya17 added 2 commits September 23, 2026 11:41

refactor

c2c3e22

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

more refactor

79d863f

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot AI review requested due to automatic review settings September 23, 2026 10:01

Copilot started reviewing on behalf of ashwinvaidya17 September 23, 2026 10:02 View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

Comment thread src/anomalib/models/components/base/anomalib_module.py

Comment thread src/anomalib/pre_processing/utils/spec.py Outdated

Comment thread tests/unit/pre_processing/test_spec.py Outdated

address copilot comments

5a46d9f

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Copilot AI review requested due to automatic review settings September 23, 2026 11:19

Merge branch 'main' into fix/engine/weights-loading

c885b37

Copilot started reviewing on behalf of ashwinvaidya17 September 23, 2026 11:20 View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

ashwinvaidya17 and others added 4 commits September 24, 2026 09:27

address copilot comments + refactor tests

2c53d04

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

fix guide

7da2de7

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Merge branch 'main' into fix/engine/weights-loading

5e62dc1

refactor tests

ada0a1c

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

rajeshgangireddy requested a balanced review from Copilot September 29, 2026 09:38

Copilot started reviewing on behalf of rajeshgangireddy September 29, 2026 09:39 View session

Copilot AI reviewed Sep 29, 2026

View reviewed changes

samet-akcay previously approved these changes Sep 29, 2026

View reviewed changes

address PR comments

3dafbed

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

ashwinvaidya17 dismissed samet-akcay’s stale review via 3dafbed September 29, 2026 10:47

ashwinvaidya17 merged commit 95fc959 into open-edge-platform:main Sep 29, 2026
47 of 70 checks passed

ashwinvaidya17 deleted the fix/engine/weights-loading branch September 29, 2026 11:21

Conversation

ashwinvaidya17 commented Sep 17, 2026

📝 Description

Design

In Anomalib Module

Single context manager, layered allowlist

Dual load seams

Per-model extension point

✨ Changes

✅ Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rajeshgangireddy left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rajeshgangireddy Sep 18, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rajeshgangireddy left a comment •

edited

Loading