[ENH] Add save/load serialization for all skpro objects (#1072) - #1073
[ENH] Add save/load serialization for all skpro objects (#1072)#1073patelchaitany wants to merge 6 commits into
Conversation
bac6c44 to
9440194
Compare
9440194 to
007da46
Compare
|
The Recursion & stop condition: Each class that has
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 Inspired by the |
|
@fkiraly What do you think? |
Signed-off-by: Chaitany Patel <patelchaitany93@gmail.com>
46341b0 to
20ebe70
Compare
There was a problem hiding this comment.
Very nice! I think this is great!
- may I kindly ask to document the file format somewhere? I think the
savedocstring 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 thetsformat) and reference it there
- in
- 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
| from io import BytesIO | ||
|
|
||
| if serialization_format == "pickle": | ||
| serialized = pickle.dumps(self) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
|
|
||
|
|
||
| def is_v1_zip(namelist): | ||
| """Check if a zip file uses the v1 flat format. |
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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:
- Unpickle the parent
_obj. - Restore the parent's native artifacts.
- Read
_components/index.json. - Recursively load each child using the same node contract.
- Reattach each child at the recorded location.
This makes the format recursive without introducing a different serialization format for the root.
| 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(...)``. | ||
|
|
There was a problem hiding this comment.
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,
}
| 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"] |
There was a problem hiding this comment.
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>
e4b3d5a to
374793e
Compare
25dba03 to
63dd5cf
Compare
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
63dd5cf to
a0e77d1
Compare
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.
Reference Issues/PRs
Fixes #1072
What does this implement/fix? Explain your changes.
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
How to: add yourself to the all-contributors file in the
skproroot directory (not theCONTRIBUTORS.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 pluscodeif you also fixed the bug in the PR).maintenance- CI, test framework, release.See here for full badge reference
For new estimators
docs/source/api_reference/taskname.rst, follow the pattern.Examplessection.python_dependenciestag and ensureddependency isolation, see the estimator dependencies guide.