[offload] DiskCache.clean_offload_dir by brian-dellabetta · Pull Request #678 · vllm-project/compressed-tensors

brian-dellabetta · 2026-04-10T18:38:10Z

Corequisite:

[oneshot] clean offload_dir during post-processing llm-compressor#2605

The DiskCache creates 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 an entrypoint.

PR adds helpers to clean any intermediate files that were added during normal operation.

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

coderabbitai · 2026-04-10T18:38:19Z

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: 4dc0a2c0-e1ad-4215-8577-1a3ba798d4ae

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.}

kylesayrs · 2026-04-13T14:33:23Z

        }

+    @classmethod
+    def clean_offload_dir(cls, offload_dir: str | os.PathLike | None = None) -> int:


Do we ever expect there to be files in the offload dir that aren't created by CT? Or files created by CT which are not in the offload dir? Why not just delete the whole directly?

I went with a design goal of scoping to only files created in this session, as the user might be running multiple sessions or jobs and clearing out a user's offload_dir entirely seems like a bad side-effect / UX

Makes sense, #666 will help handle the case where two sessions share the same folder

kylesayrs · 2026-04-13T14:34:03Z

+                except FileNotFoundError:
+                    pass  # Already deleted


Why would this case occur?

if a file is deleted by some other means? seems like a place where it's worth a few lines to fool-proof it

Consider using os.path.exists instead of try; catch

kylesayrs · 2026-04-13T14:34:49Z

        }

+    @classmethod
+    def clean_offload_dir(cls, offload_dir: str | os.PathLike | None = None) -> int:


How do you expect this to be called in distributed cases?

we need to discuss distributed case, how it would get called, but implementation is below in dist_disk.py

kylesayrs · 2026-04-13T14:36:07Z

+            file_path = cls.index[offloaded]["safetensors_file"]
+
+            # Only delete files in the specified dir (if provided)
+            if offload_dir is not None and not file_path.startswith(offload_dir):


Technically Path.resolve will resolve symlinks right? So maybe the startswith would fail here?

i can remove this if we're fine to just use cls.offload_dir

This should also be an assert

kylesayrs · 2026-04-13T14:38:06Z

+                           If None, clean all files in the shared index.
+        :return: Number of files cleaned up
+        """
+        if offload_dir is not None:


Why provide offload_dir and test for nullability? It's guaranteed at this point that cls.offload_dir is populated, right?

I can remove this, was following the convention used in create_checkpoint_symlink which provides optional offload_dir. same question arises there

I think this is fine, but we should raise major warnings/errors if the caller provided offload_dir ever conflicts with cls.offload_dir

kylesayrs

Looks good, my only points are

There are 3 cases where "if" statements should be assertions. If any of these assumptions are broken, then something has seriously gone wrong
os.unlink == os.remove?
I think we should remove cls.index, since it no longer contains valid data. You might consider renaming to DiskCache.clear_index()
I think you do need to think more about when this is going to be called. As I see it, there are two cases for when this could be called, and both conflict with this implementation

Case A

If we delete files before the model is cleaned up, then we've created a model which will error if any parameters are accessed, which is a bad state to be in. Assuming no user shenanigans this should be fine if we make it clear to users that the model is invalid after this and will error if used.

Case B

If we delete files during model cleanup, then we'll want to delete incrementally, so the implementation will look very different than this. One solution would be to add a finalizer to all instances of DiskCache which attempts file deletion on garbage collection. This is the safest but not what's implemented.

Case C

If we delete files after model cleanup, this is safe, but requires the user to del their model, then call this, which like case A needs to made clear to the user.

kylesayrs · 2026-04-17T14:33:11Z

+                continue
+
+            # Only delete files we created (with our prefix)
+            if os.path.basename(file_path).startswith(cls._new_file_prefix):


This should be an assert, not an if

kylesayrs · 2026-04-17T14:52:48Z

+            file_path = cls.index[offloaded]["safetensors_file"]
+
+            # Only delete files in the specified dir (if provided)
+            if offload_dir is not None and not file_path.startswith(offload_dir):


This should also be an assert

kylesayrs · 2026-04-17T15:01:02Z

+            # Only delete files we created (with our prefix)
+            if os.path.basename(file_path).startswith(cls._new_file_prefix):
+                try:
+                    # os.unlink removes symlink itself, not target


I don't think I understand this point? What does it mean to remove the symlink itself? os.remove and os.unlink are exactly equivalent from my testing and the documentation.

kylesayrs · 2026-04-17T15:04:41Z

+                           If None, clean all files in the shared index.
+        :return: Number of files cleaned up
+        """
+        if offload_dir is not None:


I think this is fine, but we should raise major warnings/errors if the caller provided offload_dir ever conflicts with cls.offload_dir

kylesayrs · 2026-04-17T15:05:39Z

        }

+    @classmethod
+    def clean_offload_dir(cls, offload_dir: str | os.PathLike | None = None) -> int:


Makes sense, #666 will help handle the case where two sessions share the same folder

mergify · 2026-06-22T21:38:32Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @brian-dellabetta.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

clean offload dir

584945b

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

brian-dellabetta requested review from HDCharles and kylesayrs April 10, 2026 18:38

brian-dellabetta mentioned this pull request Apr 10, 2026

[oneshot] clean offload_dir during post-processing vllm-project/llm-compressor#2605

Draft

3 tasks

kylesayrs reviewed Apr 13, 2026

View reviewed changes

kylesayrs reviewed Apr 17, 2026

View reviewed changes

kylesayrs mentioned this pull request Jun 22, 2026

[Offload] use weakref.finalize to handle shared tensor deletion #698

Open

mergify Bot added the needs-rebase label Jun 22, 2026

Conversation

brian-dellabetta commented Apr 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Apr 10, 2026

Review skipped

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kylesayrs left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Case A

Case B

Case C

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

brian-dellabetta commented Apr 10, 2026 •

edited

Loading

brian-dellabetta Apr 13, 2026 •

edited

Loading

kylesayrs left a comment •

edited

Loading