Skip to content

[ENH] Add save/load serialization for all skpro objects (#1072) - #1073

Open
patelchaitany wants to merge 6 commits into
sktime:mainfrom
patelchaitany:enh/save-load-serialization
Open

[ENH] Add save/load serialization for all skpro objects (#1072)#1073
patelchaitany wants to merge 6 commits into
sktime:mainfrom
patelchaitany:enh/save-load-serialization

Conversation

@patelchaitany

@patelchaitany patelchaitany commented Jun 18, 2026

Copy link
Copy Markdown
Member

Reference Issues/PRs

Fixes #1072

What does this implement/fix? Explain your changes.

  • Add save and load methods to BaseObject (pickle and joblib backends)
  • Standalone load() function in skpro.base
  • capability:serializable tag for opt-out
  • Round-trip tests for in-memory and file-based persistence

Does your contribution introduce a new dependency? If yes, which one?

What should a reviewer concentrate their feedback on?

Did you add any tests for the change?

Any other comments?

PR checklist

For all contributions
  • I've added myself to the list of contributors with any new badges I've earned :-)
    How to: add yourself to the all-contributors file in the skpro root directory (not the CONTRIBUTORS.md). Common badges: code - fixing a bug, or adding code logic. doc - writing or improving documentation or docstrings. bug - reporting or diagnosing a bug (get this plus code if you also fixed the bug in the PR).maintenance - CI, test framework, release.
    See here for full badge reference
  • The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.
For new estimators
  • I've added the estimator to the API reference - in docs/source/api_reference/taskname.rst, follow the pattern.
  • I've added one or more illustrative usage examples to the docstring, in a pydocstyle compliant Examples section.
  • If the estimator relies on a soft dependency, I've set the python_dependencies tag and ensured
    dependency isolation, see the estimator dependencies guide.

@patelchaitany
patelchaitany marked this pull request as draft June 18, 2026 11:08
@patelchaitany
patelchaitany force-pushed the enh/save-load-serialization branch from bac6c44 to 9440194 Compare June 18, 2026 16:38
@patelchaitany
patelchaitany force-pushed the enh/save-load-serialization branch from 9440194 to 007da46 Compare June 18, 2026 17:20
@patelchaitany
patelchaitany marked this pull request as ready for review June 19, 2026 05:21
@patelchaitany

Copy link
Copy Markdown
Member Author

The save() method recursively traverses the object tree. Each sub-component (e.g., each estimator in a VotingProbaRegressor) gets saved to its own subfolder inside the zip. A manifest.json at the root acts as a blueprint - lists every component, its class, its location in the zip, and the topological load order (leaves first, root last). On load, read the manifest, load in topological order, reassemble the tree.

Recursion & stop condition:

Each class that has save/load capability (skpro BaseObject descendants) is responsible for its own serialization - save() is called recursively down the tree. The recursion stops when:

  1. A component implements its own save/load but has no sub-components (leaf node) -saved as a structured entry in its own subfolder.
  2. A component does not implement save/load (e.g., sklearn estimators, third-party objects) - falls back to pickle as a .pkl blob and is marked as "serialized_as": "pickle_fallback" in the manifest so it's transparent to the user.
  3. If an object is not even picklable - raise an error, since there's no way to serialize it.

This avoids requiring every class to implement custom save/load logic (too much to ask), while keeping the manifest honest about what used structured save vs raw pickle fallback.

File structure (example: fitted VotingProbaRegressor with 2 sub-estimators):

voter_model.zip
├── manifest.json
├── root/
│   ├── _metadata
│   └── _obj
└── components/
    ├── estimators__0__1/    # r1 (unfitted)
    ├── estimators__1__1/    # r2 (unfitted)
    ├── estimators___0__1/   # r1 (fitted)
    └── estimators___1__1/   # r2 (fitted)

Inspired by the artifacts.yaml manifest pattern from the pytorch-forecasting EP - same core idea of manifest as single source of truth for what was saved, where, and how to reconstruct it.

@patelchaitany

Copy link
Copy Markdown
Member Author

@fkiraly What do you think?

Signed-off-by: Chaitany Patel <patelchaitany93@gmail.com>
@patelchaitany
patelchaitany force-pushed the enh/save-load-serialization branch from 46341b0 to 20ebe70 Compare June 25, 2026 15:36

@fkiraly fkiraly left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice! I think this is great!

  • may I kindly ask to document the file format somewhere? I think the save docstring should either link, or contain, the description of what is being saved exactly where.
    • in sktime, we will document the format in the docs (next to the ts format) and reference it there
  • I would also kindly like a review by @geetu040 and @phoeenniixx to ensure we are working towards a consistent format in all three packages
    • in particular, if we have composites that use objects from two or three of then, the recursive format should match

@fkiraly fkiraly added enhancement module:base-framework BaseObject, registry, base framework labels Jun 29, 2026
Comment thread skpro/base/_base.py Outdated
from io import BytesIO

if serialization_format == "pickle":
serialized = pickle.dumps(self)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I have a silly question: Why dont we save the copy of the class instance itself in a var and return that? Why do we need use pickle here? What I am able to understand is we are storing the class inside serialized (inside RAM) and if the session is complete, it would be lost anyway?

Or is it to provide data sharing over a network or something?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the serialized model is only saved when the path is not provided so the user can use that for cashing as well as the network transfer as i know it is not full saving but it is build for this purpose in mind only.

Comment thread skpro/base/_recursive_serialize.py Outdated


def is_v1_zip(namelist):
"""Check if a zip file uses the v1 flat format.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is a v1 flat format?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is like for the simpler format if any one used pickle to save them, then current version of load can load them successfully.

@phoeenniixx phoeenniixx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!
I think the superficial design is similar to ptf (using a manifest, saving the artifacts in a nested structure). I have a few doubts (see above) so that I can better understand the requirements you have thought of - would be really helpful for me in ptf!

@geetu040 geetu040 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewing in context of sktime/sktime#10582

Is there a design discussion or document that I am missing on, would be nice to share if there is any. For now I am relying on serialization.rst in docs/ to understand the pattern.

Nice implementation. I think the recursive handling of composite is good enough. We might need to consider the overall design and format for saving/loading.

Current sktime layout

I have documented the current design in the PR description for sktime/sktime#10453, but I'll add some for context here.

For an ordinary estimator, sktime keeps _metadata and _obj at the root:

estimator.zip
├── _metadata
└── _obj

When the estimator owns native artifacts, those artifacts are stored alongside the regular object state:

estimator.zip
├── _metadata
├── _obj
└── _artifacts/
    ├── index.json
    └── model_/
        ├── config.json
        └── model.safetensors

Another native backend might produce:

estimator.zip
├── _metadata
├── _obj
└── _artifacts/
    ├── index.json
    └── network_/
        └── state_dict.pt

The important property is that a serialized estimator is a self-contained node:

node
├── _metadata
├── _obj
└── _artifacts/       # optional

_obj contains the regular estimator state, while _artifacts contains attributes that must be saved using framework-native serialization.

Current skpro layout

The recursive layout in this PR is different:

estimator.zip
├── manifest.json
├── _format
├── _version
├── root/
│   ├── _metadata
│   └── _obj
└── components/
    ├── estimators__0__1/
    │   ├── _metadata
    │   └── _obj
    └── estimators__1__1/
        ├── _metadata
        └── _obj

I do not think the additional root/ directory is needed. The archive root already represents the root estimator, so _metadata and _obj can remain there. This also preserves the existing sktime layout and makes a non-composite archive a natural subset of a composite archive.

Proposed recursive layout

I suggest treating every estimator as the same self-contained serialization node and applying that node structure recursively:

estimator.zip
├── _metadata
├── _obj
├── _artifacts/                   # optional native artifacts of root
│   ├── index.json
│   └── model_/
│       └── model.safetensors
└── _components/                  # optional child estimators
    ├── index.json
    ├── forecaster_/
    │   ├── _metadata
    │   ├── _obj
    │   ├── _artifacts/           # optional native artifacts of child
    │   │   ├── index.json
    │   │   └── network_/
    │   │       └── state_dict.pt
    │   └── _components/          # optional grandchildren
    │       ├── index.json
    │       └── transformer_/
    │           ├── _metadata
    │           └── _obj
    └── transformer_/
        ├── _metadata
        └── _obj

_components/index.json could describe how child nodes are connected to the parent:

{
  "forecaster_": {
    "path": "forecaster_",
    "attribute": "forecaster_",
    "class": "sktime.forecasting.some_module.SomeForecaster"
  },
  "transformer_": {
    "path": "transformer_",
    "attribute": "transformer_",
    "class": "sktime.transformations.some_module.SomeTransformer"
  }
}

The parent _obj would contain placeholders or references for extracted child estimators. Loading would then:

  1. Unpickle the parent _obj.
  2. Restore the parent's native artifacts.
  3. Read _components/index.json.
  4. Recursively load each child using the same node contract.
  5. Reattach each child at the recorded location.

This makes the format recursive without introducing a different serialization format for the root.

Comment on lines +23 to +41
In-memory format
----------------

When ``obj.save()`` (or ``obj.save(path=None)``) is called, the return
value is a tuple:

.. code-block:: text

(cls, serialized_bytes, serialization_format)

where:

* ``cls`` is ``type(obj)``
* ``serialized_bytes`` is the full object blob encoded with
``pickle`` or ``joblib`` (see ``serialization_format``)
* ``serialization_format`` is ``"pickle"`` or ``"joblib"``

Restore with ``load(serial)`` or ``cls.load_from_serial(...)``.

@geetu040 geetu040 Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think skpro should preserve the established two-element API:

That is:

(metadata, serialized_bytes)

And all estimator-related info can go in metadata:

metadata = {
    "class": type(self),
    "serialization_format": "pickle",
    "format_version": 2,
}

Comment thread skpro/base/_recursive_serialize.py Outdated
Comment on lines +73 to +112
def recursive_load_from_zip(path):
"""Load an object from a recursive zip file.

Parameters
----------
path : str or Path
Path to the zip archive.

Returns
-------
obj : BaseObject
The deserialized object.
"""
path = Path(path)

with ZipFile(path, "r") as zf:
manifest = json.loads(zf.read("manifest.json"))
fmt = manifest["serialization_format"]

loaded = {}

for comp_id in manifest["load_order"]:
comp_info = manifest["components"][comp_id]
prefix = comp_info["path"]

obj_bytes = zf.read(prefix + "_obj")

if comp_info["is_leaf"]:
loaded[comp_id] = _deserialize_blob(obj_bytes, fmt)
else:
state_dict = _deserialize_blob(obj_bytes, fmt)
state_dict = _replace_placeholders(state_dict, loaded)

cls_bytes = zf.read(prefix + "_metadata")
cls = pickle.loads(cls_bytes) # noqa: S301
obj = cls.__new__(cls)
obj.__dict__.update(state_dict)
loaded[comp_id] = obj

return loaded["root"]

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is also an important difference in reconstruction.

The current skpro recursive loader reconstructs composite nodes approximately as:

obj = cls.__new__(cls)
obj.__dict__.update(state_dict)

By contrast, sktime first unpickles _obj normally and then restores native artifacts.

I am not against this design, I just though this should be mentioned and if preffered, we could have the same thing for sktime.

Signed-off-by: Chaitany Patel <patelchaitany93@gmail.com>
Replaces the flat-archive-plus-global-manifest design with the
serialization-node format specified in STEP 27, "Recursive composite and
native serialization" (sktime/enhancement-proposals#52), which explicitly
rejected the previous approach.

Archive format

- the archive root is the root object's own node; the `root/` wrapper and
  the flat `components/` directory are gone
- `manifest.json` is removed entirely, along with the `_format` and
  `_version` sidecars; each node now carries its own `_metadata`,
  `_artifacts/index.json` and `_components/index.json`, so no global index
  can drift out of sync with the directory tree
- `_components/` is recursive: each child is a full node that may hold its
  own artifacts and children
- `_metadata` is a versioned mapping of format version, class, and
  serialization format; it is always readable with plain pickle, so a
  reader can determine a node's format before it knows the serializer

Component references

- children are referenced by pickle persistent IDs rather than encoded
  attribute paths, via a Pickler subclass recognising children through the
  shared skbase base-object protocol
- this fixes three defects in the previous walker: children held in dict
  attributes were silently absorbed into the parent, custom and immutable
  containers were rebuilt by type and mangled, and a child referenced
  twice was restored as two distinct objects
- component IDs are opaque and node-local, and a node resolves them only
  through its own index

Native artifacts

- adds `_artifacts/` with the pretrained, keras, lightning_checkpoint and
  torch_state_dict backends, matching sktime#10453 so archives are
  readable across packages
- adds the `serialization:skip` and `serialization:native_artifacts` tags,
  with classification ordered skip, artifacts, components, then `_obj`;
  an attribute carrying both tags raises

Safety and compatibility

- saving never mutates the source object, restoring removed attributes
  even on partial failure
- ownership cycles and cross-branch aliases raise clear, archive-relative
  errors instead of recursing forever; self-references round-trip via the
  pickle memo
- component paths escaping their own node are rejected
- an unsupported format version is rejected up front
- readers still accept legacy bare-class `_metadata`, legacy in-memory
  tuples, and minimal nodes
- in-memory save now matches disk: a leaf stays lightweight pickle bytes,
  a composite becomes an in-memory zip of the same layout

Other

- moves the serialization API into a private `_SerializationMixin`, keeping
  the loose `load` thin
- replaces the `joblib` format with `cloudpickle`, per STEP 27 and sktime;
  cloudpickle support is what motivates storing the class object in
  `_metadata` rather than a qualified name
- drops the orphaned `capability:pred_int` tag, which was unrelated to
  serialization and referenced nowhere
- adds skpro/base/tests/test_serialize.py covering the archive layout,
  container shapes, identity, unsupported graphs, tags, and compatibility
@patelchaitany
patelchaitany force-pushed the enh/save-load-serialization branch from 63dd5cf to a0e77d1 Compare August 19, 2026 16:39
The previous commit replaced the `joblib` format with `cloudpickle` per
STEP 27, but nothing exercised it: `cloudpickle` was not declared in any
dependency group, and no test passed `serialization_format="cloudpickle"`.
The format was therefore untested on every CI job.

- declares `cloudpickle` in `all_extras`, so the all-extras matrix jobs
  install it
- adds round-trip coverage for the format: nested composites, per-node
  recording and child inheritance of the format name, and both in-memory
  paths
- checks that `_metadata` written by cloudpickle is still readable with
  plain pickle, which is what lets a reader determine a node's format
  before it knows the serializer
- checks that a class defined outside module scope, which plain pickle
  cannot serialize, round-trips under cloudpickle; this is the case that
  motivates storing the class object in `_metadata` rather than a
  qualified name
- checks that requesting the format without the dependency installed
  raises a clear error

The cloudpickle tests skip when the dependency is absent, and the
missing-dependency test skips when it is present, so both environments are
covered.
The previous commit declared `cloudpickle` in `all_extras`, on the
assumption that the `run-tests-all-extras` CI jobs install that group.
They do not: despite the name, those jobs run `uv pip install .[dev]`.
The only job installing `all_extras` is `run-notebook-examples`, which
does not run the test suite, so the cloudpickle tests would have skipped
on every job.

Adds `cloudpickle` to `dev`, which every test job installs, so the format
is actually exercised. It stays in `all_extras` as well, since it is a
genuine optional runtime dependency rather than test-only tooling.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement module:base-framework BaseObject, registry, base framework

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ENH] Add save/load serialization for all skpro objects

4 participants