[oneshot] clean offload_dir during post-processing

[brian-dellabetta]

SUMMARY:
Corequisite:

• [offload] DiskCache.clean_offload_dir compressed-tensors#678

The DiskCache creates potentially several hundreds of intermediate files in user-provided offload_dir during normal operation with large models. These are either symlinks or intermediate tensors that are not necessary/valuable to keep around after successfully running oneshot.

PR adds an input to oneshot to allow user to clean the offload dir during post-processing.

Some open questions:

• Do we want to run this in try/catch instead?

• Do this with a context instead? with disk_cache as DiskCache(): ...

• How to handle DistributedDiskCache?

TEST PLAN:
"please outline how the changes were tested"

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: baa608f2-eca7-4fc3-9dc1-f1c1caa17c49

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch bdellabe/offload-dir-prune

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces a clean_offload_dir parameter to the oneshot entrypoint and the post_process utility, enabling the automatic cleanup of intermediate disk offload files. The changes include updating the Oneshot class, the oneshot function, and the post_process function to handle this new flag. Feedback was provided to correct a typo in a warning message and to refine the logic in post_process to prevent misleading warnings when an output directory is not provided.

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.