Skip to content

Test all keys missed in storage in caching and eval mode#311

Closed
jiashuy wants to merge 1 commit intoNVIDIA:mainfrom
jiashuy:test_empty_batch
Closed

Test all keys missed in storage in caching and eval mode#311
jiashuy wants to merge 1 commit intoNVIDIA:mainfrom
jiashuy:test_empty_batch

Conversation

@jiashuy
Copy link
Copy Markdown
Collaborator

@jiashuy jiashuy commented Feb 26, 2026

Description

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.

@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps bot commented Feb 26, 2026

Greptile Summary

This PR adds a defensive check to prevent calling cache operations with empty tensors. In eval mode, when all keys are missing from storage, the code would previously attempt to update the cache with empty key tensors, which could cause issues. The fix adds an early return in update_cache() when there are no keys to process.

  • Added early return check in update_cache() when missing_keys.numel() == 0
  • Added test coverage for the corner case where all requested keys are missing in eval mode
  • The fix prevents potential issues with cache.insert_and_evict() being called with empty tensors

Confidence Score: 5/5

  • This PR is safe to merge with minimal risk
  • The change is a simple defensive check that prevents edge case errors, includes appropriate test coverage, and follows good engineering practices. The logic is straightforward and the test validates the fix works correctly.
  • No files require special attention

Important Files Changed

Filename Overview
corelib/dynamicemb/dynamicemb/key_value_table.py Added early return in update_cache() to handle empty missing_keys tensor, preventing potential issues with cache operations
corelib/dynamicemb/test/test_batched_dynamic_embedding_tables_v2.py Added test case for corner case where all keys are missing from storage in eval mode

Last reviewed commit: 5f44bc6

Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps bot left a comment

Choose a reason for hiding this comment

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

2 files reviewed, no comments

Edit Code Review Agent Settings | Greptile

missing_values: torch.Tensor,
missing_scores: Optional[torch.Tensor] = None,
):
batch = missing_keys.numel()
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

There is no problem even if I remove the check. Adding it to get more robust.

@shijieliu
Copy link
Copy Markdown
Collaborator

merge this PR into #313 and close this one.

@jiashuy jiashuy added the dynamicemb Related with dynamicemb label Apr 7, 2026
@jiashuy
Copy link
Copy Markdown
Collaborator Author

jiashuy commented Apr 10, 2026

closed as already fixed in #349

@jiashuy jiashuy closed this Apr 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dynamicemb Related with dynamicemb

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants