Streamline in-memory span eviction thresholds - #161
Yuge Zhang (ultmaster) merged 7 commits into
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR streamlines memory management for span eviction in the in-memory store by replacing explicit capacity overrides with automatic memory detection and consolidating span size accounting. The key changes improve memory handling efficiency and add better test coverage for span-related functionality.
- Remove explicit memory capacity overrides and rely on detected system memory totals
- Merge span size accounting into unified helpers with strict start-time-based eviction
- Add comprehensive test coverage for span eviction behavior and memory threshold validation
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| agentlightning/store/memory.py | Implements new memory detection, threshold resolution, and span eviction logic with consolidated size accounting |
| tests/store/test_memory.py | Adds test coverage for span eviction, memory threshold validation, and model size estimation functionality |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| def _detect_total_memory_bytes() -> int: | ||
| """Best-effort detection of the total available system memory in bytes.""" | ||
|
|
||
| psutil_spec = importlib.util.find_spec("psutil") | ||
| if psutil_spec is not None: | ||
| psutil = importlib.import_module("psutil") | ||
| return int(psutil.virtual_memory().total) | ||
|
|
||
| # Fallback to 8GB if memory cannot be detected. | ||
| return 8 * 1024**3 |
There was a problem hiding this comment.
Missing import statement for importlib.util in the imports section. The function uses importlib.util.find_spec but only importlib is imported.
| self._rollouts.get(rollout_id).start_time | ||
| if self._rollouts.get(rollout_id) | ||
| else spans[0].start_time or 0.0 |
There was a problem hiding this comment.
Potential AttributeError if self._rollouts.get(rollout_id) returns None. The code attempts to access .start_time on a potentially None value. Consider using a single variable to store the rollout result.
| self._rollouts.get(rollout_id).start_time | |
| if self._rollouts.get(rollout_id) | |
| else spans[0].start_time or 0.0 | |
| (self._rollouts.get(rollout_id).start_time if (rollout := self._rollouts.get(rollout_id)) is not None else spans[0].start_time or 0.0) |
| + sys.getsizeof(context.span_id) | ||
| + sys.getsizeof(context.is_remote) | ||
| + sys.getsizeof(context.trace_state) | ||
| + sum(sys.getsizeof(v) for v in context.trace_state.values()) |
There was a problem hiding this comment.
The test assumes context.trace_state has dictionary values, but line 664 shows it's initialized as an empty dict {}. The sum() over .values() will work but the size calculation logic may be incorrect for empty dictionaries.
| if link.attributes is None: | ||
| link_attributes_expected = sys.getsizeof(None) | ||
| else: | ||
| link_attributes_expected = sys.getsizeof(link.attributes) + sum( | ||
| sys.getsizeof(v) for v in link.attributes.values() | ||
| ) |
There was a problem hiding this comment.
Potential AttributeError when link.attributes is None (as set in line 666). The code attempts to call .values() on None, which will raise an AttributeError.
| if link.attributes is None: | |
| link_attributes_expected = sys.getsizeof(None) | |
| else: | |
| link_attributes_expected = sys.getsizeof(link.attributes) + sum( | |
| sys.getsizeof(v) for v in link.attributes.values() | |
| ) | |
| link_attributes_expected = sys.getsizeof(link.attributes) + sum( | |
| sys.getsizeof(v) for v in (link.attributes.values() if link.attributes is not None else []) | |
| ) |
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
|
|
||
| Args: | ||
| eviction_memory_threshold: The threshold for evicting spans in bytes. | ||
| By default, it's 70% of the total VRAM available. |
There was a problem hiding this comment.
Corrected 'VRAM' to 'RAM' as this refers to system memory, not video memory.
| By default, it's 70% of the total VRAM available. | |
| By default, it's 70% of the total RAM available. |
| ) | ||
| assert estimate_model_size(inner) == inner_expected | ||
|
|
||
| mapping_expected = sys.getsizeof(outer.mapping) + sum(sys.getsizeof(v) for v in outer.mapping.values()) |
There was a problem hiding this comment.
The size calculation is incomplete. It should include the size of dictionary keys as well as values.
| mapping_expected = sys.getsizeof(outer.mapping) + sum(sys.getsizeof(v) for v in outer.mapping.values()) | |
| mapping_expected = ( | |
| sys.getsizeof(outer.mapping) | |
| + sum(sys.getsizeof(k) for k in outer.mapping.keys()) | |
| + sum(sys.getsizeof(v) for v in outer.mapping.values()) | |
| ) |
| return sum(estimate_model_size(value) for value in values) + sys.getsizeof(cast(object, obj)) | ||
| if isinstance(obj, MappingABC): | ||
| mapping = cast(Mapping[Any, Any], obj) | ||
| return sum(estimate_model_size(value) for value in mapping.values()) + sys.getsizeof(cast(object, obj)) |
There was a problem hiding this comment.
The mapping size estimation is incomplete. It should include the size of dictionary keys as well as values.
| return sum(estimate_model_size(value) for value in mapping.values()) + sys.getsizeof(cast(object, obj)) | |
| return sum(estimate_model_size(key) + estimate_model_size(value) for key, value in mapping.items()) + sys.getsizeof(cast(object, obj)) |
…x/add-span-eviction-for-inmemorylightningstore
Summary
Testing
https://chatgpt.com/codex/tasks/task_e_68efc4dadcb0832eb776177b8647c3a0