[Offload] Key tensors by id

[kylesayrs]

No description provided.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/compressed_tensors/offload/cache/disk.py (1)

94-96: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Membership test uses tensor instead of id(tensor) — will always fail.

The self.index is now keyed by int (via id()), but this assertion checks membership using the tensor object itself. A tensor will never be in an int-keyed dict, so this assertion will incorrectly fail for valid meta tensors that are in the index.

🐛 Proposed fix

         if tensor.device.type == "meta":
-            assert tensor in self.index
+            assert id(tensor) in self.index
             return tensor

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/compressed_tensors/offload/cache/disk.py` around lines 94 - 96, The
assertion that checks if a tensor is in self.index is using the tensor object
directly, but self.index is keyed by id() integers. Change the assertion from
`assert tensor in self.index` to `assert id(tensor) in self.index` so that it
correctly checks if the integer ID of the tensor exists in the index dictionary
rather than checking for the tensor object itself.

🧹 Nitpick comments (1)

src/compressed_tensors/offload/cache/disk.py (1)

35-36: 🧹 Nitpick | 🔵 Trivial | 💤 Low value

Add ClassVar annotation to match base.py pattern.

This class attribute is intentionally shared across instances (used by create_checkpoint_symlink classmethod), but the type annotation should use ClassVar for consistency with OffloadCache.keep_onloaded_values and to address the Ruff RUF012 warning.

♻️ Suggested change

+from typing import ClassVar
+
 class DiskCache(OffloadCache):
     ...
     # id(offloaded tensor) -> weight info
-    index: dict[int, dict[str, str]] = dict()
+    index: ClassVar[dict[int, dict[str, str]]] = dict()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/compressed_tensors/offload/cache/disk.py` around lines 35 - 36, The index
class attribute at the module level needs to be annotated with ClassVar to
indicate it is a shared class-level attribute rather than an instance attribute.
Import ClassVar from typing if not already imported, then update the type
annotation for the index attribute from dict[int, dict[str, str]] to
ClassVar[dict[int, dict[str, str]]] to match the pattern used in base.py and
resolve the Ruff RUF012 warning.

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/compressed_tensors/offload/cache/disk.py`:
- Around line 136-137: The assertion on the line with assert offloaded in
self.index is checking membership using the tensor object directly, but
self.index is keyed by integer ids (as shown by the next line that accesses
self.index[id(offloaded)]). Change the membership test in the assertion from
checking offloaded in self.index to checking id(offloaded) in self.index to
ensure it correctly verifies that the offloaded tensor exists in the index
before attempting to update it.
- Around line 35-36: The index dictionary is now keyed by id() values instead of
tensor objects, but three membership tests still check using tensor objects
directly without id(). Update the three assertion and conditional checks: the
assertion at disk.py line 95 using tensor, the assertion at disk.py line 136
using offloaded, and the conditional at from_accelerate.py line 189 using
offloaded. For each, replace the direct membership check (e.g., `tensor in
self.index`) with the id-based check (e.g., `id(tensor) in self.index`) to be
consistent with how the index is keyed.

---

Outside diff comments:
In `@src/compressed_tensors/offload/cache/disk.py`:
- Around line 94-96: The assertion that checks if a tensor is in self.index is
using the tensor object directly, but self.index is keyed by id() integers.
Change the assertion from `assert tensor in self.index` to `assert id(tensor) in
self.index` so that it correctly checks if the integer ID of the tensor exists
in the index dictionary rather than checking for the tensor object itself.

---

Nitpick comments:
In `@src/compressed_tensors/offload/cache/disk.py`:
- Around line 35-36: The index class attribute at the module level needs to be
annotated with ClassVar to indicate it is a shared class-level attribute rather
than an instance attribute. Import ClassVar from typing if not already imported,
then update the type annotation for the index attribute from dict[int, dict[str,
str]] to ClassVar[dict[int, dict[str, str]]] to match the pattern used in
base.py and resolve the Ruff RUF012 warning.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a8704142-9152-4fd5-927a-aa06c97d8ddc

📥 Commits

Reviewing files that changed from the base of the PR and between 72b0765 and 077b0c0.

📒 Files selected for processing (3)

• src/compressed_tensors/offload/cache/base.py
• src/compressed_tensors/offload/cache/disk.py
• src/compressed_tensors/offload/convert/to_accelerate.py

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

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: a4d2d416-93e2-430e-9a9f-c5a482152a29

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 kylesayrs/key-by-id

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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