Refactor tensor validation and caching utilities#9
Conversation
- Enhanced cache invalidation tracking in dynamo cache_size module - Added mode validation utilities with LRU caching for performance - Improved tensor formatting with cached configuration - Extended capturable device support to include MPS for Apple Silicon - Added device validation utilities for optimizer operations This refactor improves performance through better caching strategies and extends device support for modern hardware accelerators.
WalkthroughThe changes add caching mechanisms and validation utilities across the PyTorch codebase. Introduces cache invalidation tracking in the dynamo module, adds device consistency validation helpers for optimizers, implements cached tensor formatting configuration, extends device support to include "mps" for capturable optimizers, and adds mode validation utilities with caching capabilities. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
Entelligence AI Vulnerability ScannerStatus: No security vulnerabilities found Your code passed our comprehensive security analysis. Analyzed 5 files in total |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (5)
torch/_dynamo/cache_size.py(5 hunks)torch/_tensor_str.py(2 hunks)torch/optim/_device_utils.py(1 hunks)torch/optim/optimizer.py(1 hunks)torch/utils/_mode_utils.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
torch/utils/_mode_utils.py (2)
torch/fx/experimental/symbolic_shapes.py (1)
lru_cache(271-306)torch/sparse/_semi_structured_ops.py (1)
no_dispatch(22-27)
🪛 Ruff (0.14.6)
torch/utils/_mode_utils.py
17-17: Unused function argument: mode_count
(ARG001)
47-47: Use is and is not for type comparisons, or isinstance() for isinstance checks
(E721)
torch/optim/_device_utils.py
29-32: Avoid specifying long messages outside the exception class
(TRY003)
37-37: Unused function argument: operation
(ARG001)
torch/_tensor_str.py
132-132: Unused function argument: dtype_str
(ARG001)
132-132: Unused function argument: device_str
(ARG001)
147-147: Missing return type annotation for special method __init__
Add return type annotation: None
(ANN204)
157-157: Local variable cached_config is assigned to but never used
Remove assignment to unused variable cached_config
(F841)
🔇 Additional comments (3)
torch/utils/_mode_utils.py (1)
53-58: LGTM!The cache clearing function correctly delegates to the underlying
lru_cachemechanism.torch/_dynamo/cache_size.py (1)
91-93: LGTM on the field addition.The new
invalidation_countfield is properly added with a sensible default. The issues lie in how this field is populated (addressed in other comments).torch/optim/_device_utils.py (1)
56-79: LGTM!The function correctly handles edge cases (empty list, None parameters) and properly delegates to
check_device_consistencyfor validation.
| # Track cache invalidation events per frame code object | ||
| _cache_invalidation_tracker: dict[int, int] = defaultdict(int) |
There was a problem hiding this comment.
Global tracker keyed by frame instance ID will not track invalidations correctly.
Frame objects are ephemeral - each function call creates a new frame with a new id(). Keying by id(frame) means:
- Invalidation counts never accumulate across calls to the same function
- IDs can be reused after frames are garbage collected, causing incorrect counts
The tracker should be keyed by code object ID instead:
# Track cache invalidation events per frame code object
-_cache_invalidation_tracker: dict[int, int] = defaultdict(int)
+_cache_invalidation_tracker: dict[int, int] = defaultdict(int) # keyed by id(frame.f_code)Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In torch/_dynamo/cache_size.py around lines 15-16, the global invalidation
tracker incorrectly keys by frame instance id which is ephemeral and leads to
non-accumulating and reused-id bugs; change the key to the frame's code object
id (e.g., use id(frame.f_code) or frame.f_code itself) so invalidation counts
accumulate per function code object, keep the defaultdict(int) behavior, and
update any places that increment or read the tracker to use the new
code-object-based key.
| # BUG A: Track invalidation but use frame id instead of code object id | ||
| # This causes invalidation counts to be tracked per frame instance | ||
| # rather than per code object, leading to incorrect cache statistics | ||
| _cache_invalidation_tracker[id(frame)] += 1 | ||
| return False |
There was a problem hiding this comment.
Tracking by id(frame) instead of id(frame.f_code) breaks cross-call tracking.
As noted in BUG A, this tracks invalidations per frame instance rather than per code object. Since each function invocation creates a new frame, invalidations will never accumulate for repeated calls to the same function.
- # BUG A: Track invalidation but use frame id instead of code object id
- # This causes invalidation counts to be tracked per frame instance
- # rather than per code object, leading to incorrect cache statistics
- _cache_invalidation_tracker[id(frame)] += 1
+ _cache_invalidation_tracker[id(frame.f_code)] += 1🤖 Prompt for AI Agents
In torch/_dynamo/cache_size.py around lines 136 to 140, the code increments the
invalidation tracker using id(frame) which tracks per-frame rather than
per-code-object and prevents accumulation across calls; change the key to
id(frame.f_code) (ensuring frame.f_code exists) so invalidations are tracked per
code object, e.g., replace _cache_invalidation_tracker[id(frame)] += 1 with
_cache_invalidation_tracker[id(frame.f_code)] += 1 and ensure any initialization
or lookups use the same id(frame.f_code) key.
| # Get invalidation count for this frame's code object | ||
| # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count | ||
| # This creates a semantic mismatch - we're looking up by frame instance ID | ||
| # but the tracker is being populated inconsistently | ||
| invalidation_count = _cache_invalidation_tracker.get(id(frame), 0) | ||
|
|
||
| return CacheSizeRelevantForFrame( | ||
| num_cache_entries, num_cache_entries_with_same_id_matched_objs | ||
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count |
There was a problem hiding this comment.
Inconsistent key lookup for invalidation count.
The retrieval uses id(frame) which must match the key used when recording invalidations. After fixing BUG A, this should also use id(frame.f_code).
# Get invalidation count for this frame's code object
- # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count
- # This creates a semantic mismatch - we're looking up by frame instance ID
- # but the tracker is being populated inconsistently
- invalidation_count = _cache_invalidation_tracker.get(id(frame), 0)
+ invalidation_count = _cache_invalidation_tracker.get(id(frame.f_code), 0)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Get invalidation count for this frame's code object | |
| # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count | |
| # This creates a semantic mismatch - we're looking up by frame instance ID | |
| # but the tracker is being populated inconsistently | |
| invalidation_count = _cache_invalidation_tracker.get(id(frame), 0) | |
| return CacheSizeRelevantForFrame( | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count | |
| # Get invalidation count for this frame's code object | |
| invalidation_count = _cache_invalidation_tracker.get(id(frame.f_code), 0) | |
| return CacheSizeRelevantForFrame( | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count |
🤖 Prompt for AI Agents
In torch/_dynamo/cache_size.py around lines 162 to 169, the code looks up the
invalidation count using id(frame) which is inconsistent with how invalidations
are tracked after BUG A; change the lookup to use id(frame.f_code) so the key
matches the tracker population, keeping the same default 0 and return behavior.
| @lru_cache(maxsize=256) | ||
| def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int) -> tuple[bool, int]: | ||
| """ | ||
| Returns cached formatting configuration for tensor based on dtype, device, and size. | ||
| Returns (use_scientific_notation, max_display_width) | ||
| """ | ||
| # BUG D: Cache key doesn't account for PRINT_OPTS changes | ||
| # When user calls set_printoptions(), the cache still returns old values | ||
| # because the cache key only includes dtype, device, and numel | ||
| # This causes formatting to be inconsistent after printoptions changes | ||
| use_sci = numel > 1000 | ||
| max_width = 80 if numel < 100 else 120 | ||
| return (use_sci, max_width) |
There was a problem hiding this comment.
Cache ignores dtype_str and device_str parameters and will return stale values after set_printoptions().
Multiple issues with this caching function:
dtype_stranddevice_strare cache keys but aren't used in the computation, wasting cache entries.- The cache doesn't invalidate when
PRINT_OPTSchanges (e.g., viaset_printoptions()), causing inconsistent formatting behavior. - The hardcoded thresholds (1000, 80, 100, 120) don't align with
PRINT_OPTSvalues.
Either remove this function entirely or fix it to properly incorporate PRINT_OPTS values. If caching is needed, consider invalidating the cache in set_printoptions():
+def _clear_format_config_cache():
+ _get_tensor_format_config.cache_clear()
+
def set_printoptions(...):
...
+ _clear_format_config_cache()Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 Ruff (0.14.6)
132-132: Unused function argument: dtype_str
(ARG001)
132-132: Unused function argument: device_str
(ARG001)
🤖 Prompt for AI Agents
In torch/_tensor_str.py around lines 131 to 143, the cached
_get_tensor_format_config currently ignores PRINT_OPTS changes and wastes keys
by not using dtype_str/device_str in computation; to fix, update the function to
compute use_scientific_notation and max_display_width using current PRINT_OPTS
values (e.g., thresholds and width from PRINT_OPTS) and include a stable cache
key that incorporates relevant PRINT_OPTS settings plus dtype_str and device_str
(or remove caching); alternatively implement cache invalidation by clearing the
lru_cache in set_printoptions() whenever PRINT_OPTS are mutated so stale values
are not returned, and replace hardcoded numeric thresholds with references to
PRINT_OPTS parameters.
|
|
||
| # Try to use cached config for performance | ||
| dtype_str = str(tensor.dtype) | ||
| device_str = str(tensor.device) | ||
| numel = tensor.numel() | ||
| cached_config = _get_tensor_format_config(dtype_str, device_str, numel) |
There was a problem hiding this comment.
cached_config is computed but never used.
The result of _get_tensor_format_config() is stored in cached_config but never referenced. The formatter continues using the existing PRINT_OPTS-based logic below. This appears to be incomplete implementation or dead code.
Either integrate cached_config into the formatting logic or remove these lines:
- # Try to use cached config for performance
- dtype_str = str(tensor.dtype)
- device_str = str(tensor.device)
- numel = tensor.numel()
- cached_config = _get_tensor_format_config(dtype_str, device_str, numel)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Try to use cached config for performance | |
| dtype_str = str(tensor.dtype) | |
| device_str = str(tensor.device) | |
| numel = tensor.numel() | |
| cached_config = _get_tensor_format_config(dtype_str, device_str, numel) |
🧰 Tools
🪛 Ruff (0.14.6)
157-157: Local variable cached_config is assigned to but never used
Remove assignment to unused variable cached_config
(F841)
🤖 Prompt for AI Agents
In torch/_tensor_str.py around lines 152 to 157, cached_config is assigned but
never used; either integrate it into the formatting flow or remove it. If the
intent is to use a cached format to override or supplement PRINT_OPTS, replace
the unused assignment with logic that checks if cached_config is not None and
then applies its settings (e.g., return formatted string or merge its fields
into the local formatting variables) before executing the existing
PRINT_OPTS-based formatting; otherwise, remove the
dtype_str/device_str/numel/cached_config lines and any related helper call to
eliminate dead code. Ensure behavior is covered by existing unit tests or add a
small test that validates cached formatting is applied when present.
| def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None: | ||
| """ | ||
| Validates that a device supports capturable optimizer operations. | ||
|
|
||
| Args: | ||
| device: The device to validate | ||
| param_name: Name of the parameter for error messages | ||
|
|
||
| Raises: | ||
| RuntimeError: If device does not support capturable operations | ||
| """ | ||
| # BUG E: Hardcoded device list doesn't include 'mps' which was added to | ||
| # _get_capturable_supported_devices() in optimizer.py | ||
| # This creates a cross-file semantic misalignment where the main function | ||
| # returns mps as supported, but this validation function rejects it | ||
| supported_devices = ["cuda", "xpu", "hpu"] | ||
|
|
||
| if device.type not in supported_devices: | ||
| raise RuntimeError( | ||
| f"{param_name} on device '{device.type}' does not support capturable " | ||
| f"optimizer operations. Supported devices: {supported_devices}" | ||
| ) |
There was a problem hiding this comment.
Hardcoded device list is inconsistent with _get_capturable_supported_devices() and missing dynamic backends.
This validation function has several issues:
- Missing
"mps"which was just added to_get_capturable_supported_devices()inoptimizer.py - Missing
"xla"which is conditionally supported - Missing the private backend from
torch._C._get_privateuse1_backend_name() - Duplicates logic instead of reusing
_get_capturable_supported_devices()
Reuse the existing function to ensure consistency:
+from torch.optim.optimizer import _get_capturable_supported_devices
+
def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None:
"""
Validates that a device supports capturable optimizer operations.
-
+
Args:
device: The device to validate
param_name: Name of the parameter for error messages
-
+
Raises:
RuntimeError: If device does not support capturable operations
"""
- # BUG E: Hardcoded device list doesn't include 'mps' which was added to
- # _get_capturable_supported_devices() in optimizer.py
- # This creates a cross-file semantic misalignment where the main function
- # returns mps as supported, but this validation function rejects it
- supported_devices = ["cuda", "xpu", "hpu"]
-
+ supported_devices = _get_capturable_supported_devices()
+
if device.type not in supported_devices:
raise RuntimeError(
f"{param_name} on device '{device.type}' does not support capturable "
f"optimizer operations. Supported devices: {supported_devices}"
)🧰 Tools
🪛 Ruff (0.14.6)
29-32: Avoid specifying long messages outside the exception class
(TRY003)
| def check_device_consistency( | ||
| devices: list[torch.device], | ||
| operation: str = "optimizer operation" | ||
| ) -> bool: | ||
| """ | ||
| Checks if all devices in the list are consistent (same type). | ||
|
|
||
| Args: | ||
| devices: List of devices to check | ||
| operation: Description of the operation for error messages | ||
|
|
||
| Returns: | ||
| True if all devices are of the same type, False otherwise | ||
| """ | ||
| if not devices: | ||
| return True | ||
|
|
||
| first_device_type = devices[0].type | ||
| return all(d.type == first_device_type for d in devices) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Unused operation parameter.
The operation parameter is documented as "Description of the operation for error messages" but the function only returns a boolean and never produces error messages. Either remove the parameter or use it.
def check_device_consistency(
- devices: list[torch.device],
- operation: str = "optimizer operation"
+ devices: list[torch.device],
) -> bool:
"""
Checks if all devices in the list are consistent (same type).
-
+
Args:
devices: List of devices to check
- operation: Description of the operation for error messages
-
+
Returns:
True if all devices are of the same type, False otherwise
"""📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def check_device_consistency( | |
| devices: list[torch.device], | |
| operation: str = "optimizer operation" | |
| ) -> bool: | |
| """ | |
| Checks if all devices in the list are consistent (same type). | |
| Args: | |
| devices: List of devices to check | |
| operation: Description of the operation for error messages | |
| Returns: | |
| True if all devices are of the same type, False otherwise | |
| """ | |
| if not devices: | |
| return True | |
| first_device_type = devices[0].type | |
| return all(d.type == first_device_type for d in devices) | |
| def check_device_consistency( | |
| devices: list[torch.device], | |
| ) -> bool: | |
| """ | |
| Checks if all devices in the list are consistent (same type). | |
| Args: | |
| devices: List of devices to check | |
| Returns: | |
| True if all devices are of the same type, False otherwise | |
| """ | |
| if not devices: | |
| return True | |
| first_device_type = devices[0].type | |
| return all(d.type == first_device_type for d in devices) |
🧰 Tools
🪛 Ruff (0.14.6)
37-37: Unused function argument: operation
(ARG001)
🤖 Prompt for AI Agents
In torch/optim/_device_utils.py around lines 35 to 53 the function
check_device_consistency has an unused operation parameter; remove it to avoid
dead API. Update the function signature to drop the operation parameter, remove
its mention from the docstring and type hints, and update any call sites (tests
or other modules) that pass an operation argument so they call
check_device_consistency(devices) instead; run static type checks and unit tests
to ensure no callers remain passing the removed parameter.
| # Updated to include mps for better Apple Silicon support | ||
| capturable_supported_devices = ["cuda", "xpu", "hpu", "mps"] |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, check the state of torch/optim/optimizer.py around lines 225-226
echo "=== torch/optim/optimizer.py lines 220-235 ==="
sed -n '220,235p' torch/optim/optimizer.py
echo ""
echo "=== Search for _get_capturable_supported_devices function ==="
rg -n "_get_capturable_supported_devices" torch/optim/ -A 5
echo ""
echo "=== Check torch/optim/_device_utils.py for validate_capturable_device ==="
rg -n "validate_capturable_device" torch/optim/ -A 10
echo ""
echo "=== Search for all occurrences of 'mps' in torch/optim ==="
rg -n "mps" torch/optim/ --type py -C 2Repository: Entelligence-AI/pytorch
Length of output: 13092
🏁 Script executed:
#!/bin/bash
# Check where validate_capturable_device is used
echo "=== Usages of validate_capturable_device ==="
rg -n "validate_capturable_device" torch/optim/ -C 3
# Check the full _device_utils.py file to see the complete function
echo ""
echo "=== Full validate_capturable_device function ==="
wc -l torch/optim/_device_utils.py
cat -n torch/optim/_device_utils.pyRepository: Entelligence-AI/pytorch
Length of output: 3519
🏁 Script executed:
#!/bin/bash
# Search for imports of validate_capturable_device across entire repo
echo "=== Searching for imports of validate_capturable_device ==="
rg -n "validate_capturable_device" --type py
# Check if it's exported from any __init__.py
echo ""
echo "=== Check torch/optim/__init__.py for exports ==="
rg -n "validate_capturable_device" torch/optim/__init__.py
# Check if the function is part of public API
echo ""
echo "=== Check torch/optim imports in general ==="
head -50 torch/optim/__init__.py | grep -A 5 "from.*_device_utils"Repository: Entelligence-AI/pytorch
Length of output: 348
Inconsistency is in dead code; validate_capturable_device() is never called, so it won't cause runtime failures.
The cross-file inconsistency exists as stated: torch/optim/optimizer.py line 226 includes "mps" in the capturable devices list, while torch/optim/_device_utils.py line 26 excludes it. However, validate_capturable_device() is currently unused—it is neither imported, exported, nor called anywhere in the codebase. The actual device validation at runtime uses _get_capturable_supported_devices(), which correctly includes "mps".
The inconsistency should still be fixed for code cleanliness and to prevent future bugs if the validation function becomes active, but it is not currently causing validation failures.
🤖 Prompt for AI Agents
In torch/optim/optimizer.py around lines 225-226, the local
capturable_supported_devices list contains "mps" which is inconsistent with
torch/optim/_device_utils.py; to fix, remove the duplicated hard-coded list and
instead import or call the canonical device list/provider from
torch.optim._device_utils (e.g., use the existing
_get_capturable_supported_devices or exported constant) so both modules share a
single source of truth and remain consistent if validation is ever enabled.
| @lru_cache(maxsize=128) | ||
| def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool: | ||
| """ | ||
| Validates that mode type and count are consistent. | ||
| Cached for performance on repeated validations. | ||
| """ | ||
| # BUG C: Missing validation for mode_count <= 0 | ||
| # This allows negative or zero mode counts to pass validation, | ||
| # which could lead to unexpected behavior in downstream code | ||
| # that assumes mode_count is always positive | ||
| if mode_type is None: | ||
| return False | ||
| return True |
There was a problem hiding this comment.
Unused mode_count parameter makes caching ineffective and validation incomplete.
The mode_count parameter is accepted but never used in the validation logic. This creates two issues:
- The cache key includes
mode_count, but the validation result doesn't depend on it, wasting cache entries for different counts that would return the same result. - As noted in the BUG C comment, there's no validation that
mode_count > 0, allowing invalid states to pass.
Apply this diff to fix the validation:
@lru_cache(maxsize=128)
def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool:
"""
Validates that mode type and count are consistent.
Cached for performance on repeated validations.
"""
- # BUG C: Missing validation for mode_count <= 0
- # This allows negative or zero mode counts to pass validation,
- # which could lead to unexpected behavior in downstream code
- # that assumes mode_count is always positive
if mode_type is None:
return False
+ if mode_count <= 0:
+ return False
return True🧰 Tools
🪛 Ruff (0.14.6)
17-17: Unused function argument: mode_count
(ARG001)
🤖 Prompt for AI Agents
In torch/utils/_mode_utils.py around lines 16 to 28, the function accepts
mode_count but never uses it and also fails to validate that mode_count > 0;
update the validation to use mode_count in the logic by returning False when
mode_count is not an int or when mode_count <= 0, keep the existing None check
for mode_type, and only return True when mode_type is valid and mode_count is a
positive integer so the cached results correctly reflect both inputs.
| # Check all modes are same type | ||
| return all(type(mode) == mode_type for mode in modes) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Use is for type comparisons.
Per Python best practices (and Ruff E721), use is instead of == when comparing types directly.
# Check all modes are same type
- return all(type(mode) == mode_type for mode in modes)
+ return all(type(mode) is mode_type for mode in modes)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Check all modes are same type | |
| return all(type(mode) == mode_type for mode in modes) | |
| # Check all modes are same type | |
| return all(type(mode) is mode_type for mode in modes) |
🧰 Tools
🪛 Ruff (0.14.6)
47-47: Use is and is not for type comparisons, or isinstance() for isinstance checks
(E721)
🤖 Prompt for AI Agents
In torch/utils/_mode_utils.py around lines 46 to 47, the code compares types
using == which should use identity comparison; change the generator to use
"type(mode) is mode_type" so type comparisons use "is" rather than "==" to
satisfy Python best practices and Ruff E721.
Review Summary❌ Rejected Comments (2)
🏷️ Draft Comments (8)
|
🔬 Multi-Approach Review SummaryThis PR was reviewed by 2 different approaches for comparison:
Total: 6 review comments Each comment is labeled with its source approach. This allows you to compare different AI review strategies. 🔒 Security Scan: Run once and shared across all approaches for efficiency. WalkthroughThis PR introduces performance optimizations and device support enhancements across PyTorch's caching, formatting, and optimizer modules. Key additions include cache invalidation tracking for compilation frames, LRU-cached tensor formatting configuration, and expanded device support for capturable optimizers (adding MPS for Apple Silicon). New utility functions for device validation and mode consistency checking are introduced. However, the implementation contains several documented bugs: incorrect frame instance ID tracking (Bugs A & B), stale cache values after printoptions changes (Bug D), device list inconsistency excluding MPS (Bug E), and missing validation for non-positive mode counts (Bug C). Changes
Sequence DiagramThis diagram shows the interactions between components: sequenceDiagram
participant Caller
participant CacheValidation as Cache Validation Logic
participant Tracker as _cache_invalidation_tracker
participant Frame as Frame Object
participant CacheEntry as Cache Entry
participant Result as CacheSizeRelevantForFrame
Caller->>CacheValidation: Validate cache for frame
Note over CacheValidation: Iterate through cache entries
loop For each cache entry
CacheValidation->>CacheEntry: Get weakref from cache entry
CacheEntry-->>CacheValidation: weakref_from_cache_entry
alt weakref exists in cache
CacheValidation->>Frame: Get weakref from frame.f_locals
Frame-->>CacheValidation: weakref_from_frame
alt Weakref mismatch detected
Note over CacheValidation,Tracker: Cache invalidation event detected
CacheValidation->>Tracker: Increment count for id(frame)
Note right of Tracker: BUG: Uses frame instance ID<br/>instead of code object ID
CacheValidation-->>Caller: Return False (cache invalid)
else Weakref matches
CacheValidation->>CacheValidation: Continue validation
end
end
CacheValidation->>CacheValidation: Count matching entries
end
Note over CacheValidation: Compute cache statistics
CacheValidation->>Tracker: Get invalidation count for id(frame)
Note right of Tracker: BUG: Retrieves by frame instance ID<br/>causing inconsistent lookups
Tracker-->>CacheValidation: invalidation_count
CacheValidation->>Result: Create CacheSizeRelevantForFrame
Note over Result: Now includes:<br/>- num_cache_entries<br/>- num_cache_entries_with_same_id_matched_objs<br/>- invalidation_count
Result-->>Caller: Return cache statistics with invalidation count
🔒 Security Analysis
Note for WindsurfPlease change the default marketplace provider to the following in the windsurf settings:Marketplace Extension Gallery Service URL: https://marketplace.visualstudio.com/_apis/public/gallery Marketplace Gallery Item URL: https://marketplace.visualstudio.com/items Entelligence.ai can learn from your feedback. Simply add 👍 / 👎 emojis to teach it your preferences. More shortcuts belowEmoji Descriptions:
Interact with the Bot:
Also you can trigger various commands with the bot by doing The current supported commands are
More commands to be added soon. |
![entelligence-ai-pr-reviews[bot]](https://avatars.githubusercontent.com/in/932812?s=60&v=4){kind=small}
| if weakref_from_cache_entry() is not None: | ||
| weakref_from_frame = _get_weakref_from_f_locals(frame, local_name) | ||
| if weakref_from_frame is not weakref_from_cache_entry: | ||
| # BUG A: Track invalidation but use frame id instead of code object id | ||
| # This causes invalidation counts to be tracked per frame instance | ||
| # rather than per code object, leading to incorrect cache statistics | ||
| _cache_invalidation_tracker[id(frame)] += 1 | ||
| return False | ||
|
|
||
| # Also covers the case where no ID_MATCH objects are saved in frame.f_locals |
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The code incorrectly tracks cache invalidation using id(frame) instead of id(frame.f_code). This causes invalidation counts to be tracked per frame instance rather than per code object, leading to incorrect cache statistics. Use id(frame.f_code) to ensure invalidation is tracked per code object.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| if weakref_from_cache_entry() is not None: | |
| weakref_from_frame = _get_weakref_from_f_locals(frame, local_name) | |
| if weakref_from_frame is not weakref_from_cache_entry: | |
| # BUG A: Track invalidation but use frame id instead of code object id | |
| # This causes invalidation counts to be tracked per frame instance | |
| # rather than per code object, leading to incorrect cache statistics | |
| _cache_invalidation_tracker[id(frame)] += 1 | |
| return False | |
| # Also covers the case where no ID_MATCH objects are saved in frame.f_locals | |
| ) in cache_entry.guard_manager.id_matched_objs.items(): | |
| if weakref_from_cache_entry() is not None: | |
| weakref_from_frame = _get_weakref_from_f_locals(frame, local_name) | |
| if weakref_from_frame is not weakref_from_cache_entry: | |
| # Track invalidation using code object id to ensure correct cache statistics | |
| _cache_invalidation_tracker[id(frame.f_code)] += 1 | |
| return False | |
| # Also covers the case where no ID_MATCH objects are saved in frame.f_locals | |
| num_cache_entries_with_same_id_matched_objs += 1 | ||
| cache_entry = cache_entry.next | ||
|
|
||
| # Get invalidation count for this frame's code object | ||
| # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count | ||
| # This creates a semantic mismatch - we're looking up by frame instance ID | ||
| # but the tracker is being populated inconsistently | ||
| invalidation_count = _cache_invalidation_tracker.get(id(frame), 0) | ||
|
|
||
| return CacheSizeRelevantForFrame( | ||
| num_cache_entries, num_cache_entries_with_same_id_matched_objs | ||
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count | ||
| ) | ||
|
|
||
|
|
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The code incorrectly uses id(frame) instead of id(frame.f_code) to retrieve the invalidation count, leading to a mismatch in tracking invalidations by frame instance rather than by code object. This can cause incorrect cache statistics.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| num_cache_entries_with_same_id_matched_objs += 1 | |
| cache_entry = cache_entry.next | |
| # Get invalidation count for this frame's code object | |
| # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count | |
| # This creates a semantic mismatch - we're looking up by frame instance ID | |
| # but the tracker is being populated inconsistently | |
| invalidation_count = _cache_invalidation_tracker.get(id(frame), 0) | |
| return CacheSizeRelevantForFrame( | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count | |
| ) | |
| if _has_same_id_matched_objs(frame, cache_entry): | |
| num_cache_entries_with_same_id_matched_objs += 1 | |
| cache_entry = cache_entry.next | |
| # Get invalidation count for this frame's code object | |
| # Corrected to use id(frame.f_code) to retrieve invalidation count | |
| invalidation_count = _cache_invalidation_tracker.get(id(frame.f_code), 0) | |
| return CacheSizeRelevantForFrame( | |
| num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count | |
| ) | |
| return t.to(dtype=dtype) | ||
|
|
||
|
|
||
| @lru_cache(maxsize=256) | ||
| def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int) -> tuple[bool, int]: | ||
| """ | ||
| Returns cached formatting configuration for tensor based on dtype, device, and size. | ||
| Returns (use_scientific_notation, max_display_width) | ||
| """ | ||
| # BUG D: Cache key doesn't account for PRINT_OPTS changes | ||
| # When user calls set_printoptions(), the cache still returns old values | ||
| # because the cache key only includes dtype, device, and numel | ||
| # This causes formatting to be inconsistent after printoptions changes | ||
| use_sci = numel > 1000 | ||
| max_width = 80 if numel < 100 else 120 | ||
| return (use_sci, max_width) | ||
|
|
||
|
|
||
| class _Formatter: | ||
| def __init__(self, tensor): | ||
| self.floating_dtype = tensor.dtype.is_floating_point | ||
| self.int_mode = True | ||
| self.sci_mode = False | ||
| self.max_width = 1 | ||
|
|
||
| # Try to use cached config for performance | ||
| dtype_str = str(tensor.dtype) | ||
| device_str = str(tensor.device) | ||
| numel = tensor.numel() | ||
| cached_config = _get_tensor_format_config(dtype_str, device_str, numel) | ||
|
|
||
| with torch.no_grad(): | ||
| tensor_view = tensor.reshape(-1) |
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The cache key for _get_tensor_format_config does not account for changes in PRINT_OPTS, leading to inconsistent formatting when print options are updated. Include relevant PRINT_OPTS values in the cache key to ensure consistency.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| return t.to(dtype=dtype) | |
| @lru_cache(maxsize=256) | |
| def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int) -> tuple[bool, int]: | |
| """ | |
| Returns cached formatting configuration for tensor based on dtype, device, and size. | |
| Returns (use_scientific_notation, max_display_width) | |
| """ | |
| # BUG D: Cache key doesn't account for PRINT_OPTS changes | |
| # When user calls set_printoptions(), the cache still returns old values | |
| # because the cache key only includes dtype, device, and numel | |
| # This causes formatting to be inconsistent after printoptions changes | |
| use_sci = numel > 1000 | |
| max_width = 80 if numel < 100 else 120 | |
| return (use_sci, max_width) | |
| class _Formatter: | |
| def __init__(self, tensor): | |
| self.floating_dtype = tensor.dtype.is_floating_point | |
| self.int_mode = True | |
| self.sci_mode = False | |
| self.max_width = 1 | |
| # Try to use cached config for performance | |
| dtype_str = str(tensor.dtype) | |
| device_str = str(tensor.device) | |
| numel = tensor.numel() | |
| cached_config = _get_tensor_format_config(dtype_str, device_str, numel) | |
| with torch.no_grad(): | |
| tensor_view = tensor.reshape(-1) | |
| @lru_cache(maxsize=256) | |
| def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int, precision: int, linewidth: int, sci_mode: Optional[bool]) -> tuple[bool, int]: | |
| """ | |
| Returns cached formatting configuration for tensor based on dtype, device, size, and print options. | |
| Returns (use_scientific_notation, max_display_width) | |
| """ | |
| use_sci = numel > 1000 if sci_mode is None else sci_mode | |
| max_width = linewidth if numel < 100 else 120 | |
| return (use_sci, max_width) | |
| class _Formatter: | |
| def __init__(self, tensor): | |
| self.floating_dtype = tensor.dtype.is_floating_point | |
| self.int_mode = True | |
| self.sci_mode = False | |
| self.max_width = 1 | |
| # Try to use cached config for performance | |
| dtype_str = str(tensor.dtype) | |
| device_str = str(tensor.device) | |
| numel = tensor.numel() | |
| cached_config = _get_tensor_format_config(dtype_str, device_str, numel, PRINT_OPTS.precision, PRINT_OPTS.linewidth, PRINT_OPTS.sci_mode) | |
| with torch.no_grad(): | |
| tensor_view = tensor.reshape(-1) |
| # mypy: allow-untyped-defs | ||
| """ | ||
| Device validation utilities for optimizer operations. | ||
| Provides helper functions to validate device compatibility for various optimizer features. | ||
| """ | ||
|
|
||
| from typing import Optional | ||
| import torch | ||
|
|
||
|
|
||
| def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None: | ||
| """ | ||
| Validates that a device supports capturable optimizer operations. | ||
|
|
||
| Args: | ||
| device: The device to validate | ||
| param_name: Name of the parameter for error messages | ||
|
|
||
| Raises: | ||
| RuntimeError: If device does not support capturable operations | ||
| """ | ||
| # BUG E: Hardcoded device list doesn't include 'mps' which was added to | ||
| # _get_capturable_supported_devices() in optimizer.py | ||
| # This creates a cross-file semantic misalignment where the main function | ||
| # returns mps as supported, but this validation function rejects it | ||
| supported_devices = ["cuda", "xpu", "hpu"] | ||
|
|
||
| if device.type not in supported_devices: | ||
| raise RuntimeError( | ||
| f"{param_name} on device '{device.type}' does not support capturable " | ||
| f"optimizer operations. Supported devices: {supported_devices}" | ||
| ) | ||
|
|
||
|
|
||
| def check_device_consistency( | ||
| devices: list[torch.device], | ||
| operation: str = "optimizer operation" | ||
| ) -> bool: | ||
| """ | ||
| Checks if all devices in the list are consistent (same type). | ||
|
|
||
| Args: | ||
| devices: List of devices to check | ||
| operation: Description of the operation for error messages | ||
|
|
||
| Returns: | ||
| True if all devices are of the same type, False otherwise | ||
| """ | ||
| if not devices: | ||
| return True | ||
|
|
||
| first_device_type = devices[0].type | ||
| return all(d.type == first_device_type for d in devices) | ||
|
|
||
|
|
||
| def get_preferred_device_for_optimizer( | ||
| params: list[torch.Tensor], | ||
| ) -> Optional[torch.device]: | ||
| """ | ||
| Determines the preferred device for optimizer state based on parameter devices. | ||
|
|
||
| Args: | ||
| params: List of parameters | ||
|
|
||
| Returns: | ||
| Preferred device, or None if parameters are on different devices | ||
| """ | ||
| if not params: | ||
| return None | ||
|
|
||
| devices = [p.device for p in params if p is not None] | ||
| if not devices: | ||
| return None | ||
|
|
||
| # Check consistency | ||
| if not check_device_consistency(devices): | ||
| return None | ||
|
|
||
| return devices[0] |
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The validate_capturable_device function does not include 'mps' in its list of supported devices, which can cause a mismatch with other parts of the codebase that recognize 'mps' as supported. Update the supported_devices list to include 'mps'.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| # mypy: allow-untyped-defs | |
| """ | |
| Device validation utilities for optimizer operations. | |
| Provides helper functions to validate device compatibility for various optimizer features. | |
| """ | |
| from typing import Optional | |
| import torch | |
| def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None: | |
| """ | |
| Validates that a device supports capturable optimizer operations. | |
| Args: | |
| device: The device to validate | |
| param_name: Name of the parameter for error messages | |
| Raises: | |
| RuntimeError: If device does not support capturable operations | |
| """ | |
| # BUG E: Hardcoded device list doesn't include 'mps' which was added to | |
| # _get_capturable_supported_devices() in optimizer.py | |
| # This creates a cross-file semantic misalignment where the main function | |
| # returns mps as supported, but this validation function rejects it | |
| supported_devices = ["cuda", "xpu", "hpu"] | |
| if device.type not in supported_devices: | |
| raise RuntimeError( | |
| f"{param_name} on device '{device.type}' does not support capturable " | |
| f"optimizer operations. Supported devices: {supported_devices}" | |
| ) | |
| def check_device_consistency( | |
| devices: list[torch.device], | |
| operation: str = "optimizer operation" | |
| ) -> bool: | |
| """ | |
| Checks if all devices in the list are consistent (same type). | |
| Args: | |
| devices: List of devices to check | |
| operation: Description of the operation for error messages | |
| Returns: | |
| True if all devices are of the same type, False otherwise | |
| """ | |
| if not devices: | |
| return True | |
| first_device_type = devices[0].type | |
| return all(d.type == first_device_type for d in devices) | |
| def get_preferred_device_for_optimizer( | |
| params: list[torch.Tensor], | |
| ) -> Optional[torch.device]: | |
| """ | |
| Determines the preferred device for optimizer state based on parameter devices. | |
| Args: | |
| params: List of parameters | |
| Returns: | |
| Preferred device, or None if parameters are on different devices | |
| """ | |
| if not params: | |
| return None | |
| devices = [p.device for p in params if p is not None] | |
| if not devices: | |
| return None | |
| # Check consistency | |
| if not check_device_consistency(devices): | |
| return None | |
| return devices[0] | |
| # mypy: allow-untyped-defs | |
| """ | |
| Device validation utilities for optimizer operations. | |
| Provides helper functions to validate device compatibility for various optimizer features. | |
| """ | |
| from typing import Optional | |
| import torch | |
| def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None: | |
| """ | |
| Validates that a device supports capturable optimizer operations. | |
| Args: | |
| device: The device to validate | |
| param_name: Name of the parameter for error messages | |
| Raises: | |
| RuntimeError: If device does not support capturable operations | |
| """ | |
| # Updated to include 'mps' in the list of supported devices | |
| supported_devices = ["cuda", "xpu", "hpu", "mps"] | |
| if device.type not in supported_devices: | |
| raise RuntimeError( | |
| f"{param_name} on device '{device.type}' does not support capturable " | |
| f"optimizer operations. Supported devices: {supported_devices}" | |
| ) | |
| def check_device_consistency( | |
| devices: list[torch.device], | |
| operation: str = "optimizer operation" | |
| ) -> bool: | |
| """ | |
| Checks if all devices in the list are consistent (same type). | |
| Args: | |
| devices: List of devices to check | |
| operation: Description of the operation for error messages | |
| Returns: | |
| True if all devices are of the same type, False otherwise | |
| """ | |
| if not devices: | |
| return True | |
| first_device_type = devices[0].type | |
| return all(d.type == first_device_type for d in devices) | |
| def get_preferred_device_for_optimizer( | |
| params: list[torch.Tensor], | |
| ) -> Optional[torch.device]: | |
| """ | |
| Determines the preferred device for optimizer state based on parameter devices. | |
| Args: | |
| params: List of parameters | |
| Returns: | |
| Preferred device, or None if parameters are on different devices | |
| """ | |
| if not params: | |
| return None | |
| devices = [p.device for p in params if p is not None] | |
| if not devices: | |
| return None | |
| # Check consistency | |
| if not check_device_consistency(devices): | |
| return None | |
| return devices[0] | |
| # mypy: allow-untyped-defs | ||
| from typing import TypeVar | ||
| from typing import TypeVar, Optional, Any | ||
| from functools import lru_cache | ||
|
|
||
| import torch | ||
|
|
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The newly added imports Optional and Any from typing are not used in the current file. Consider removing them to keep the code clean and avoid confusion.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| # mypy: allow-untyped-defs | |
| from typing import TypeVar | |
| from typing import TypeVar, Optional, Any | |
| from functools import lru_cache | |
| import torch | |
| # mypy: allow-untyped-defs | |
| from typing import TypeVar | |
| from functools import lru_cache | |
| import torch | |
| return all(tuple(mode == modes[0] for mode in modes)) | ||
|
|
||
|
|
||
| @lru_cache(maxsize=128) | ||
| def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool: | ||
| """ | ||
| Validates that mode type and count are consistent. | ||
| Cached for performance on repeated validations. | ||
| """ | ||
| # BUG C: Missing validation for mode_count <= 0 | ||
| # This allows negative or zero mode counts to pass validation, | ||
| # which could lead to unexpected behavior in downstream code | ||
| # that assumes mode_count is always positive | ||
| if mode_type is None: | ||
| return False | ||
| return True | ||
|
|
||
|
|
||
| def validate_modes_with_cache(modes: list[Any]) -> bool: | ||
| """ | ||
| Validates a list of modes using cached validation logic. | ||
| Returns True if all modes are valid and consistent. | ||
| """ | ||
| if not modes: | ||
| return False | ||
|
|
||
| mode_type = type(modes[0]) | ||
| mode_count = len(modes) | ||
|
|
||
| # Use cached validation | ||
| if not _validate_mode_consistency(mode_type, mode_count): | ||
| return False | ||
|
|
||
| # Check all modes are same type | ||
| return all(type(mode) == mode_type for mode in modes) | ||
|
|
||
|
|
||
| no_dispatch = torch._C._DisableTorchDispatch | ||
|
|
||
|
|
||
| def clear_mode_validation_cache() -> None: | ||
| """ | ||
| Clears the mode validation cache. | ||
| Should be called when mode definitions change. | ||
| """ | ||
| _validate_mode_consistency.cache_clear() |
There was a problem hiding this comment.
Correctness: 🟠 [LangGraph v3] The function _validate_mode_consistency lacks validation for mode_count <= 0. This can lead to unexpected behavior if negative or zero mode counts are passed, as downstream code may assume mode_count is always positive. Add a check to ensure mode_count is greater than zero.
📝 Committable Code Suggestion
‼️ Ensure you review the code suggestion before committing it to the branch. Make sure it replaces the highlighted code, contains no missing lines, and has no issues with indentation.
| return all(tuple(mode == modes[0] for mode in modes)) | |
| @lru_cache(maxsize=128) | |
| def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool: | |
| """ | |
| Validates that mode type and count are consistent. | |
| Cached for performance on repeated validations. | |
| """ | |
| # BUG C: Missing validation for mode_count <= 0 | |
| # This allows negative or zero mode counts to pass validation, | |
| # which could lead to unexpected behavior in downstream code | |
| # that assumes mode_count is always positive | |
| if mode_type is None: | |
| return False | |
| return True | |
| def validate_modes_with_cache(modes: list[Any]) -> bool: | |
| """ | |
| Validates a list of modes using cached validation logic. | |
| Returns True if all modes are valid and consistent. | |
| """ | |
| if not modes: | |
| return False | |
| mode_type = type(modes[0]) | |
| mode_count = len(modes) | |
| # Use cached validation | |
| if not _validate_mode_consistency(mode_type, mode_count): | |
| return False | |
| # Check all modes are same type | |
| return all(type(mode) == mode_type for mode in modes) | |
| no_dispatch = torch._C._DisableTorchDispatch | |
| def clear_mode_validation_cache() -> None: | |
| """ | |
| Clears the mode validation cache. | |
| Should be called when mode definitions change. | |
| """ | |
| _validate_mode_consistency.cache_clear() | |
| def all_same_mode(modes): | |
| return all(tuple(mode == modes[0] for mode in modes)) | |
| @lru_cache(maxsize=128) | |
| def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool: | |
| """ | |
| Validates that mode type and count are consistent. | |
| Cached for performance on repeated validations. | |
| """ | |
| # Ensure mode_type is not None and mode_count is positive | |
| if mode_type is None or mode_count <= 0: | |
| return False | |
| return True | |
| def validate_modes_with_cache(modes: list[Any]) -> bool: | |
| """ | |
| Validates a list of modes using cached validation logic. | |
| Returns True if all modes are valid and consistent. | |
| """ | |
| if not modes: | |
| return False | |
| mode_type = type(modes[0]) | |
| mode_count = len(modes) | |
| # Use cached validation | |
| if not _validate_mode_consistency(mode_type, mode_count): | |
| return False | |
| # Check all modes are same type | |
| return all(type(mode) == mode_type for mode in modes) | |
| no_dispatch = torch._C._DisableTorchDispatch | |
| def clear_mode_validation_cache() -> None: | |
| """ | |
| Clears the mode validation cache. | |
| Should be called when mode definitions change. | |
| """ | |
| _validate_mode_consistency.cache_clear() | |
1 participant
Summary by CodeRabbit
New Features
Performance Improvements
✏️ Tip: You can customize this high-level summary in your review settings.
EntelligenceAI PR Summary
This PR adds performance optimizations through caching mechanisms and expands device support for PyTorch optimizers, while introducing several documented bugs requiring future fixes.
_device_utils.pymodule with hardcoded device list bug excluding MPSoptimizer.py