Skip to content

Rotbaum: fall back to legacy predictor.pkl in TreePredictor.deserialize - #3294

Open
shaun0927 wants to merge 2 commits into
awslabs:devfrom
shaun0927:fix/rotbaum-legacy-pickle
Open

Rotbaum: fall back to legacy predictor.pkl in TreePredictor.deserialize#3294
shaun0927 wants to merge 2 commits into
awslabs:devfrom
shaun0927:fix/rotbaum-legacy-pickle

Conversation

@shaun0927

@shaun0927 shaun0927 commented Apr 17, 2026

Copy link
Copy Markdown

Issue #, if available: follow-up to #3176

Description of changes:

#3176 migrated TreePredictor persistence from pickle (predictor.pkl) to JSON (model_list.json). TreePredictor.deserialize was rewritten to hardcode the new path:

with (path / "model_list.json").open("r") as fp:
    predictor.model_list = load_json(fp.read())

Every Rotbaum predictor that was serialized with a v0.15.x release (which wrote predictor.pkl and nothing else) now fails with FileNotFoundError: .../model_list.json when loaded on v0.16+.

Reproducer

>>> import pathlib, pickle, tempfile, sys, types
>>> sys.modules["xgboost"] = types.ModuleType("xgboost")   # avoid libomp in the demo
>>> from gluonts.ext.rotbaum._predictor import TreePredictor
>>> p = pathlib.Path(tempfile.mkdtemp())
>>> with (p / "predictor.pkl").open("wb") as f:
...     pickle.dump([["sample-model"]], f)
>>> TreePredictor.deserialize(p)
FileNotFoundError: [...]/model_list.json

Design choice: opt-in legacy load

Reintroducing pickle as an automatic fallback would resurrect the very reason #3176 migrated away from it — pickle.load executes arbitrary code on load. The first revision of this PR took that simpler path; I've changed it to require explicit opt-in so we don't regress on that front.

TreePredictor.deserialize now:

  1. loads model_list.json if present (the post-Rotbaum: turn to json-based serialization #3176 format);
  2. otherwise, if predictor.pkl is present AND the caller passed allow_legacy_pickle=True, loads it and emits a DeprecationWarning telling the user to re-save so they migrate on their own schedule;
  3. otherwise, if predictor.pkl is present but the flag is not set, raises FileNotFoundError with a one-line message that points users at the opt-in flag;
  4. raises FileNotFoundError when neither file is present.

Pickle is only imported lazily inside the fallback branch, so the JSON happy path keeps the module's current zero-pickle import surface. The stale docstring that still said "loads the trained model list by reading the pickle file" is updated to describe the new behaviour.

Trade-offs considered

  • Silent pickle auto-fallback (first revision): simpler, but anyone who can drop a predictor.pkl into a path the user calls deserialize on (shared model hub, S3 bucket with broad write perms) gets RCE via pickle. The opt-in flag keeps that surface closed by default.
  • Remove the legacy path entirely and ship a standalone migration CLI: rejected. Users who rediscover a v0.15.x artifact months later and are willing to trust it still deserve a one-call way in. The opt-in is one extra keyword argument.
  • Gate on a class attribute instead of a kwarg (e.g. TreePredictor.allow_legacy_pickle = True): rejected. A per-call kwarg scopes the trust decision to the single load.

Verification

Dynamic matrix with super().deserialize mocked out (that's the parent path, unchanged by this PR):

A json-only:              model_list=[['dummy']]                        (no warning)
B pkl-only default:       FileNotFoundError; opt-in hint present -> True
C pkl-only opt-in:        model_list=[['dummy']]                        depr=1
D both default:           model_list=[['dummy']]                        JSON wins, no warning
E neither:                FileNotFoundError                             OK

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Please tag this pr with at least one of these labels to make our release process faster: BREAKING, new feature, bug fix, other change, dev setup

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant