Collection-based Lightning Store #315
Conversation
There was a problem hiding this comment.
Pull Request Overview
This pull request introduces a collection-based architecture for the Lightning Store, refactoring the in-memory store implementation to use a more modular and extensible design pattern.
Key Changes:
- Introduces abstract collection interfaces (
Collection,Queue,KeyValue) that define storage operations - Implements
CollectionBasedLightningStoreas a base class that uses collections for data storage - Refactors
InMemoryLightningStoreto inherit fromCollectionBasedLightningStoreand useInMemoryLightningCollections - Adds comprehensive test coverage for collection structures with 600+ lines of new tests
Reviewed Changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| agentlightning/store/collection/base.py | Defines abstract interfaces for Collection, Queue, KeyValue, and LightningCollections |
| agentlightning/store/collection/memory.py | Implements in-memory collection classes using Python data structures (lists, deques, dicts) |
| agentlightning/store/collection/init.py | Exports collection classes and interfaces |
| agentlightning/store/collection_based.py | Implements CollectionBasedLightningStore that uses collections as storage backend |
| agentlightning/store/memory.py | Refactors InMemoryLightningStore to extend CollectionBasedLightningStore |
| agentlightning/store/init.py | Exports CollectionBasedLightningStore |
| tests/store/test_collection_structures.py | Comprehensive tests for collection implementations covering CRUD operations, filtering, sorting, and pagination |
| tests/store/conftest.py | Updates mock_readable_span fixture to generate unique span IDs per call |
| tests/store/test_implementation.py | Adds test for duplicate span ID validation |
| docs/reference/store.md | Documents new collection-based store classes |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # If not completed, wait for completion | ||
| while deadline is None or time.time() < deadline: | ||
| # Poll every 10 seconds by default | ||
| rest_time = max(0.01, deadline - time.time()) if deadline is not None else 10.0 |
There was a problem hiding this comment.
The sleep time calculation could cause the loop to sleep for the entire remaining time until the deadline. When deadline - time.time() is large (e.g., 60 seconds), the loop will sleep for that entire duration without polling. Consider capping the sleep time to a maximum value (e.g., 10 seconds) to ensure regular polling:
rest_time = min(10.0, max(0.01, deadline - time.time())) if deadline is not None else 10.0| rest_time = max(0.01, deadline - time.time()) if deadline is not None else 10.0 | |
| rest_time = min(10.0, max(0.01, deadline - time.time())) if deadline is not None else 10.0 |
| current_attempt.last_heartbeat_time = time.time() | ||
| if current_attempt.status in ["preparing", "unresponsive"]: | ||
| current_attempt.status = "running" |
There was a problem hiding this comment.
The current_attempt is modified (status and last_heartbeat_time) but never persisted to the collection. After modifying the attempt, you should call await self.collections.attempts.update([current_attempt]) before the return statement to persist the changes.
| def _get_sort_value(item: T, sort_by: str) -> Any: | ||
| """Get a sort key for the given item/field. | ||
|
|
||
| - If the field name ends with '_time', values are treated as comparable timestamps.s |
There was a problem hiding this comment.
There's a typo in the comment: "timestamps.s" should be "timestamps."
| - If the field name ends with '_time', values are treated as comparable timestamps.s | |
| - If the field name ends with '_time', values are treated as comparable timestamps |
| # Filter out the exceptions | ||
| return [rollout for rollout in rollouts if isinstance(rollout, Rollout)] | ||
|
|
||
| async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]: |
There was a problem hiding this comment.
Mixing implicit and explicit returns may indicate an error, as implicit returns always return None.
| assert await dict_key_value.pop("beta") == 2 | ||
| assert dict_key_value.size() == 1 | ||
| assert await dict_key_value.pop("missing", 42) == 42 |
There was a problem hiding this comment.
This 'assert' statement contains an expression which may have side effects.
| assert await dict_key_value.pop("beta") == 2 | |
| assert dict_key_value.size() == 1 | |
| assert await dict_key_value.pop("missing", 42) == 42 | |
| result = await dict_key_value.pop("beta") | |
| assert result == 2 | |
| assert dict_key_value.size() == 1 | |
| result_missing = await dict_key_value.pop("missing", 42) | |
| assert result_missing == 42 |
| async def test_dict_key_value_pop_returns_default(dict_key_value: DictBasedKeyValue[str, int]) -> None: | ||
| assert await dict_key_value.pop("beta") == 2 | ||
| assert dict_key_value.size() == 1 | ||
| assert await dict_key_value.pop("missing", 42) == 42 |
There was a problem hiding this comment.
This 'assert' statement contains an expression which may have side effects.
| assert await dict_key_value.pop("missing", 42) == 42 | |
| result = await dict_key_value.pop("missing", 42) | |
| assert result == 42 |
|
/ci |
|
🚀 CI Watcher for correlation id-3539935279-mi2nwph1 triggered by comment 3539935279
✅ All runs completed. |
|
/ci |
|
🚀 CI Watcher for correlation id-3540180574-mi2rkq1z triggered by comment 3540180574
✅ All runs completed. |
|
/ci |
|
🚀 CI Watcher for correlation id-3540794432-mi2y66mj triggered by comment 3540794432
✅ All runs completed. |
No description provided.