Skip to content

fix(models): leave the scored splits out of a saved model - #950

Open
Felipedino wants to merge 1 commit into
developfrom
fix/model-artifact-without-data
Open

Felipedino wants to merge 1 commit into
developfrom
fix/model-artifact-without-data

Conversation

@Felipedino

Copy link
Copy Markdown
Collaborator

Problem

When a model is fitted, its train/validation/test splits are attached to it as x_data / y_data so that calculate_metrics can read them. Those splits are in-memory Arrow tables. Every model that saves itself with joblib.dump(self) therefore wrote the whole dataset into its .joblib. On large datasets, SaveModelUnit took about a minute, and every run kept its own copy of the data.

Fix

  • BaseModel.__getstate__ drops x_data / y_data from the pickled state. This covers every caller of model.save, including the pipeline Train node.
  • BaseModel declares both attributes as None, so a loaded model behaves like one that was never given data.
  • FaissBaseModel.__getstate__ now calls super(), so FAISS models drop the data too.

The model in memory is unchanged. Artifacts saved before this fix still load; they just keep the extra data.

Tests

New tests/back/models/test_saved_model_leaves_out_data.py (sklearn and FAISS KMeans): the restored model has no data, the artifact size does not grow with the data, the live model keeps its data, predictions match after reload, and a model with no attributes can still be pickled. 5 of its 9 tests fail without the fix.

Whoever fits a model points x_data and y_data at the splits it trains on,
so calculate_metrics can read them. The splits are in-memory Arrow tables,
and every model that saves itself with joblib.dump(self) wrote them into
its artifact: a copy of train, validation and test on every run, and a
save that grew with the dataset instead of with the model.

BaseModel.__getstate__ now drops both attributes, which covers every
caller of model.save (SaveModelUnit and the pipeline Train node) at once.
BaseModel declares both as None, so a loaded model, like a model never
pointed at any data, scores nothing on its own instead of raising
AttributeError. FaissBaseModel built its state from __dict__ directly and
now goes through super(), so FAISS models drop the data too.

Artifacts saved before this still load; they only carry the extra data.
Copilot AI balanced review requested due to automatic review settings October 28, 2026 12:30

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

This branch has not been deployed

No deployments
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.

2 participants