feat: enable save/load after recovery without additional steps by psyberck · Pull Request #619 · PrunaAI/pruna

psyberck · 2026-04-07T05:35:37Z

Description

Extend save and load to also allow saving and loading after recovery without additional steps.

Related Issue

Fixes #603

Type of Change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Refactor (no functional change)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Documentation update

Testing

I added or updated tests covering my changes
Existing tests pass locally (uv run pytest -m "cpu and not slow")

For full setup and testing instructions, see the Contributing Guide.

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my code, especially for agent-assisted changes
I updated the documentation where necessary

Thanks for contributing to Pruna! We're excited to review your work.

New to contributing? Check out our Contributing Guide for everything you need to get started.

Note:

Draft PRs or PRs without a clear and detailed overview may be delayed.

Please mark your PR as Ready for Review and ensure the sections above are filled out.

Contributions that are entirely AI-generated without meaningful human review are discouraged.

First Prune (1-year OSS anniversary)

First Prune marks one year of Pruna’s open-source work. During the initiative window, qualifying merged contributions count toward First Prune. You can earn credits for our performance models via our API.

If you’d like your contribution to count toward First Prune, here’s how it works:

Initiative window: First Prune starts on March 31.
Issue assignment: For your PR to count toward First Prune, the related issue must be assigned to the contributor opening the PR. Issues are labeled with first-prune.
Open for review: Please open your PR and mark it ready for review by April 30 (end of April).
Review priority: We’ll make our best effort to review quickly any PR that is open and has a review request before April 30.
Credits: Each qualifying merged PR earns 30 credits. We’ll be in touch after all qualifying PRs for First Prune have been merged.
To get started: Have a look at all models. You’ll need to sign up on the dashboard before you can redeem your credits.

codacy-production · 2026-04-07T05:36:35Z

Not up to standards ⛔

🔴 Issues 3 high

Alerts:
⚠ 3 issues (≤ 0 issues of at least minor severity)

Results:
3 new issues

Category Results

Security 3 high

View in Codacy

🟢 Metrics 12 complexity · 0 duplication

Metric Results

Complexity 12

Duplication 0

View in Codacy

_{TIP This summary will be updated as you push new changes. Give us feedback}

sdiazlor · 2026-04-08T09:13:49Z

Hi @psyberck! Thank you for your contribution! I've run the tests. Please take a look before we jump into any review :)

psyberck · 2026-04-13T17:25:40Z

Hey @sdiazlor , can you help kindly re-run the tests? There seems some network failures during the tests:

ERROR tests/algorithms/test_algorithms.py::test_full_integration[TestIFW_whisper_tiny_random-cpu] - huggingface_hub.errors.LocalEntryNotFoundError: An error happened while trying to locate the files on the Hub and we cannot find the appropriate snapshot folder for the specified revision on the local disk. Please check your internet connection and try again.

sdiazlor · 2026-04-13T20:57:41Z

Sure thing! I ran them again :)

github-actions · 2026-05-02T00:22:58Z

This PR has been inactive for 10 days and is now marked as stale.

simlang

Thanks for looking at this! I left a view requests :)

sdiazlor · 2026-05-11T16:13:17Z

Hi @psyberck! When you have time, could you take a look at the comments? :)

psyberck · 2026-05-11T18:09:03Z

Thanks for the comments @simlang. I will update accordingly tonight @sdiazlor :)

psyberck · 2026-05-13T18:22:30Z

Code updated @simlang cc @sdiazlor

github-actions · 2026-05-25T00:28:41Z

This PR has been inactive for 10 days and is now marked as stale.

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Comment @cursor review or bugbot run to trigger another review on this PR

^{Reviewed by Cursor Bugbot for commit 81ca7d0. Configure here.}

github-actions · 2026-06-12T00:35:41Z

This PR has been inactive for 10 days and is now marked as stale.

github-actions · 2026-06-19T00:40:01Z

Closing this PR due to inactivity. Please reopen it or create a new PR if this is still relevant.

sdiazlor · 2026-06-24T12:47:33Z

Hi @psyberck, sorry for the delay, would it be possible to review the conflicts :)

…: get_save_before_smash_dir

…'s weights for "save_before_apply" algos

github-actions · 2026-07-14T00:20:36Z

This PR has been inactive for 10 days and is now marked as stale. It will be closed in 7 days if there is no further activity.

sdiazlor · 2026-07-14T08:34:36Z

Could the errors be fixed? :)

simlang · 2026-07-14T15:04:37Z



-class PERPRecoverer(PrunaAlgorithmBase):
+class PERPRecoverer(PrunaAlgorithmBase, metaclass=ABCMeta):


we don't use this anywhere else in pruna, to keep it consistent we can remove this

simlang · 2026-07-14T15:06:55Z


        super().__init__()  # self.adapters need to be set before calling get_hyperparameters

+    @property


please also revert this, to keep it consistent with the rest of the code

simlang · 2026-07-14T15:07:59Z


+    def post_apply_hook(self, model: Any, smash_config: SmashConfig):
+        """
+        Override to run side effects after the algorithm has been applied.


can you please change the docstring to what the override does specifically here?

simlang · 2026-07-14T15:13:33Z

+    ori_save_fns = smash_config.save_fns[:]
+    smash_config.save_fns = [fn for fn in smash_config.save_fns if fn != SAVE_FUNCTIONS.save_before_apply.name]
+    # Re-save with recovered weights
+    shutil.rmtree(model_path, ignore_errors=True)


i'm not super comfortable with ignoring errors. Is there a particular reason why you went for this - otherwise i would prefer to at least log something if something went wrong

simlang · 2026-07-14T15:15:11Z

+    smash_config.save_fns = [fn for fn in smash_config.save_fns if fn != SAVE_FUNCTIONS.save_before_apply.name]
+    # Re-save with recovered weights
+    shutil.rmtree(model_path, ignore_errors=True)
+    save_pruna_model(model, model_path, smash_config)


do you think it's possible to first save the new pruna model before removing the old one? if something goes wrong in the save here at least we didn't lose the original model
so i guess the order would be:

save new model
remove old model
move new model to the location of the old one

simlang · 2026-07-14T15:23:03Z

+    from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
+    from pruna.algorithms.global_utils.recovery.perp_recoverer import PERPRecoverer
+
+    class FakeRecoverer(PERPRecoverer):


i like the idea. but i would feel safer, if we actually use the smash-workflow instead of re-implementing the workflow in the test.

you could use monkeypatching to make these overrides and then add an algorithm which uses saving before apply to test things and finally call smash with this algorithm + monkeypatched PERPRecoverer

psyberck changed the title ~~Feat/save load recovery~~ feat: enable save/load after recovery without additional steps Apr 8, 2026

sdiazlor requested review from johannaSommer and simlang April 21, 2026 13:59

github-actions Bot added the stale label May 2, 2026

sdiazlor removed the stale label May 4, 2026

simlang requested changes May 4, 2026

View reviewed changes

sdiazlor requested a review from simlang May 14, 2026 08:23

github-actions Bot added the stale label May 25, 2026

sdiazlor removed the stale label May 25, 2026

psyberck force-pushed the feat/save-load-recovery branch from 8975c49 to 81ca7d0 Compare June 1, 2026 15:22

cursor Bot reviewed Jun 1, 2026

View reviewed changes

Comment thread src/pruna/algorithms/base/pruna_base.py

github-actions Bot added the stale label Jun 12, 2026

github-actions Bot closed this Jun 19, 2026

sdiazlor reopened this Jun 24, 2026

github-actions Bot removed the stale label Jun 25, 2026

psyberck added 3 commits June 26, 2026 09:08

refactor: add prepare_saving as config property and get util function…

3bfd375

…: get_save_before_smash_dir

refactor: save_fns won't append any new save_fn; save recovered model…

ad13bbd

…'s weights for "save_before_apply" algos

test: add unit tests to verify main ideas

89a1cbd

psyberck added 5 commits June 26, 2026 09:09

fix: remove unused imports from test_save.py

c49d50a

docs: add docstring for get_save_before_smash_dir

d4058ed

refactor: move cache refresh logic into post_apply_hook func

e5d7d6a

test: update recover unit tests

43f4839

style: docstring styles

aec268e

psyberck force-pushed the feat/save-load-recovery branch from 81ca7d0 to aec268e Compare June 26, 2026 16:10

fix: post_apply_hook parameter

b8173c1

PrunaAI deleted a comment from cursor Bot Jul 3, 2026

github-actions Bot added the stale label Jul 14, 2026

simlang requested changes Jul 14, 2026

View reviewed changes

github-actions Bot removed the stale label Jul 15, 2026



		class PERPRecoverer(PrunaAlgorithmBase):
		class PERPRecoverer(PrunaAlgorithmBase, metaclass=ABCMeta):


		super().__init__() # self.adapters need to be set before calling get_hyperparameters

		@property

Uh oh!

Conversation

psyberck commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

Testing

Checklist

First Prune (1-year OSS anniversary)

Uh oh!

codacy-production Bot commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Not up to standards ⛔

Uh oh!

sdiazlor commented Apr 8, 2026

Uh oh!

psyberck commented Apr 13, 2026

Uh oh!

sdiazlor commented Apr 13, 2026

Uh oh!

github-actions Bot commented May 2, 2026

Uh oh!

simlang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sdiazlor commented May 11, 2026

Uh oh!

psyberck commented May 11, 2026

Uh oh!

psyberck commented May 13, 2026

Uh oh!

github-actions Bot commented May 25, 2026

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Jun 12, 2026

Uh oh!

github-actions Bot commented Jun 19, 2026

Uh oh!

sdiazlor commented Jun 24, 2026

Uh oh!

github-actions Bot commented Jul 14, 2026

Uh oh!

sdiazlor commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

simlang Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

psyberck commented Apr 7, 2026 •

edited

Loading

codacy-production Bot commented Apr 7, 2026 •

edited

Loading

sdiazlor commented Jul 14, 2026 •

edited

Loading