feat: create PyTorch Geometric dataloader from RocksDB (#13) - #38
feat: create PyTorch Geometric dataloader from RocksDB (#13)#38mangopep wants to merge 1 commit into
Conversation
WalkthroughIntroduces a new ai package with dataset loaders (core and PyTorch variants), temporal graph construction, RocksDB-backed event storage and streaming utilities, package metadata, and AI-specific dependencies. Adds conditional exports for PyTorch availability, a dataloader factory, and a .gitignore for AI artifacts. No modifications to existing code outside the new ai directory. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant DS as ai.datasets
participant Core as NERRFDatasetCore
participant Torch as NERRFDataset (PyG)
participant FS as JSONL File
participant RDB as RocksDB (optional)
User->>DS: import NERRFDataset, create_dataloader
Note over DS: Conditional import\nHAS_PYTORCH?
alt PyTorch available
User->>Torch: NERRFDataset(data_path, window_size, stride, format)
Torch->>Torch: _load_events()
alt format == "jsonl"
Torch->>FS: read JSONL lines
else format == "rocksdb"
Torch->>RDB: query time range (placeholder)
end
Torch->>Torch: _create_windows()
loop per window
Torch->>Torch: build_graph(events in window)
Torch->>Torch: extract_features()
end
User->>Torch: create_dataloader(...)\n(shuffle, batch_size)
else PyTorch not available
User->>Core: NERRFDatasetCore(jsonl_path, window_size, stride)
Core->>FS: read JSONL lines
Core->>Core: _create_windows() + build_graph + extract_features
end
sequenceDiagram
autonumber
actor Ingest
participant ESP as EventStreamProcessor
participant Store as RocksDBEventStore
participant DB as RocksDB
Ingest->>ESP: process_protobuf_event(pb_event)
ESP->>Store: store_event(event_dict, timestamp)
Store->>DB: Put(time-key, JSON)
actor Trainer
Trainer->>ESP: create_training_windows(start, end)
loop sliding windows
ESP->>Store: query_events(win_start, win_end)
Store->>DB: RangeIter(time-key prefix)
Store-->>ESP: (timestamp, event_dict)*
ESP-->>Trainer: [events in window]
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
ai/.gitignore(1 hunks)ai/__init__.py(1 hunks)ai/datasets/__init__.py(1 hunks)ai/datasets/nerrf_dataset.py(1 hunks)ai/datasets/nerrf_dataset_core.py(1 hunks)ai/requirements.txt(1 hunks)ai/utils/__init__.py(1 hunks)ai/utils/rocksdb_store.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
ai/utils/__init__.py (1)
ai/utils/rocksdb_store.py (3)
RocksDBEventStore(42-255)EventStreamProcessor(258-332)migrate_jsonl_to_rocksdb(335-374)
ai/datasets/nerrf_dataset_core.py (2)
ai/datasets/nerrf_dataset.py (2)
NERRFEvent(40-83)_load_events(254-261)ai/utils/rocksdb_store.py (1)
open(34-35)
ai/datasets/nerrf_dataset.py (2)
ai/datasets/nerrf_dataset_core.py (2)
NERRFEvent(29-72)_load_events(238-266)ai/utils/rocksdb_store.py (1)
open(34-35)
ai/datasets/__init__.py (2)
ai/datasets/nerrf_dataset.py (4)
NERRFEvent(40-83)TemporalGraphBuilder(86-217)NERRFDataset(220-361)create_dataloader(364-387)ai/datasets/nerrf_dataset_core.py (3)
NERRFEvent(29-72)TemporalGraphBuilder(75-203)NERRFDatasetCore(206-374)
🪛 OSV Scanner (2.2.3)
ai/requirements.txt
[CRITICAL] 1-1: torch 2.0.0: undefined
(PYSEC-2024-250)
[CRITICAL] 1-1: torch 2.0.0: undefined
(PYSEC-2024-251)
[CRITICAL] 1-1: torch 2.0.0: undefined
(PYSEC-2024-252)
[CRITICAL] 1-1: torch 2.0.0: undefined
(PYSEC-2024-259)
[CRITICAL] 1-1: torch 2.0.0: undefined
(PYSEC-2025-41)
[CRITICAL] 1-1: torch 2.0.0: PyTorch susceptible to local Denial of Service
[CRITICAL] 1-1: torch 2.0.0: PyTorch: torch.load with weights_only=True leads to remote code execution
[CRITICAL] 1-1: torch 2.0.0: PyTorch heap buffer overflow vulnerability
[CRITICAL] 1-1: torch 2.0.0: PyTorch Improper Resource Shutdown or Release vulnerability
[CRITICAL] 1-1: torch 2.0.0: Pytorch use-after-free vulnerability
🪛 Ruff (0.13.3)
ai/utils/__init__.py
5-5: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
ai/datasets/nerrf_dataset_core.py
176-176: Consider [*node_type, path_hash, *time_features] instead of concatenation
Replace with [*node_type, path_hash, *time_features]
(RUF005)
181-181: Loop control variable u not used within loop body
Rename unused u to _u
(B007)
181-181: Loop control variable v not used within loop body
Rename unused v to _v
(B007)
196-200: Consider iterable unpacking instead of concatenation
(RUF005)
262-262: Do not catch blind exception: Exception
(BLE001)
302-302: Avoid specifying long messages outside the exception class
(TRY003)
ai/datasets/nerrf_dataset.py
189-189: Consider [*node_type, path_hash, *time_features] instead of concatenation
Replace with [*node_type, path_hash, *time_features]
(RUF005)
194-194: Loop control variable u not used within loop body
Rename unused u to _u
(B007)
194-194: Loop control variable v not used within loop body
Rename unused v to _v
(B007)
209-213: Consider iterable unpacking instead of concatenation
(RUF005)
261-261: Avoid specifying long messages outside the exception class
(TRY003)
ai/datasets/__init__.py
16-16: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
ai/utils/rocksdb_store.py
34-34: Unused static method argument: args
(ARG004)
34-34: Unused static method argument: kwargs
(ARG004)
35-35: Avoid specifying long messages outside the exception class
(TRY003)
86-86: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
132-132: Consider moving this statement to an else block
(TRY300)
134-134: Do not catch blind exception: Exception
(BLE001)
135-135: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
161-161: Consider moving this statement to an else block
(TRY300)
163-163: Do not catch blind exception: Exception
(BLE001)
164-164: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
193-193: Do not catch blind exception: Exception
(BLE001)
194-194: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
222-222: Do not catch blind exception: Exception
(BLE001)
223-223: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
234-234: Do not catch blind exception: Exception
(BLE001)
235-235: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
246-246: Consider moving this statement to an else block
(TRY300)
247-247: Do not catch blind exception: Exception
(BLE001)
248-248: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
300-300: Do not catch blind exception: Exception
(BLE001)
301-301: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
370-370: Consider moving this statement to an else block
(TRY300)
372-372: Do not catch blind exception: Exception
(BLE001)
373-373: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
| if node_data["node_type"] == "process": | ||
| node_type = [1.0, 0.0] # [process, file] | ||
| path_hash = hash(f"pid_{node_data['pid']}") % 1000 / 1000.0 | ||
| else: # file | ||
| node_type = [0.0, 1.0] | ||
| path_hash = hash(node_data["path"]) % 1000 / 1000.0 | ||
|
|
||
| # Time features (hour of day, day of week) | ||
| timestamp = node_data["first_seen"] | ||
| time_features = [ | ||
| timestamp.hour / 24.0, | ||
| timestamp.weekday() / 7.0 | ||
| ] | ||
|
|
||
| features = node_type + [path_hash] + time_features | ||
| node_features.append(features) |
There was a problem hiding this comment.
Stabilize node feature hashing.
hash(...) is salted per Python process, so the same path or PID produces different feature values run-to-run, breaking reproducibility between training, validation, and inference. Replace it with a deterministic digest (e.g., SHA256/MD5 modulo a fixed scale) so identical inputs always yield the same embedding.
🧰 Tools
🪛 Ruff (0.13.3)
189-189: Consider [*node_type, path_hash, *time_features] instead of concatenation
Replace with [*node_type, path_hash, *time_features]
(RUF005)
🤖 Prompt for AI Agents
In ai/datasets/nerrf_dataset.py around lines 175-190, the code uses Python's
built-in hash() for path and PID which is randomized per process; replace those
calls with a deterministic hash (e.g., hashlib.sha256 or md5) by encoding the
string (for PID use f"pid_{pid}" as before), computing a hex/int digest,
reducing it with modulo 1000, and dividing by 1000.0 to preserve the existing
scale; add the required import (import hashlib) at the top and apply the same
deterministic function for both process and file branches so features are stable
across runs.
| start_time = min(event.timestamp for event in self.events) | ||
| end_time = max(event.timestamp for event in self.events) | ||
|
|
||
| windows = [] | ||
| current_start = start_time | ||
|
|
||
| while current_start + self.window_size <= end_time: | ||
| current_end = current_start + self.window_size | ||
| windows.append((current_start, current_end)) | ||
| current_start += self.stride | ||
|
|
||
| return windows |
There was a problem hiding this comment.
Always emit at least one time window.
while current_start + self.window_size <= end_time drops every dataset whose span is smaller than window_size—even a single event yields zero windows, making the dataset unusable. Adjust the loop to include the final (possibly partial) window, for example by iterating while current_start <= end_time and breaking once you pass the tail.
🤖 Prompt for AI Agents
In ai/datasets/nerrf_dataset.py around lines 296 to 307, the window generation
loop uses `while current_start + self.window_size <= end_time` which yields zero
windows when the total span is smaller than window_size; change the loop
condition to `while current_start <= end_time` so the final (possibly partial)
window is emitted, compute `current_end = min(current_start + self.window_size,
end_time)` for each iteration, append that window, then advance `current_start`
by `self.stride` and optionally break when `current_start > end_time` to avoid
infinite loops.
| G = self.graph_builder.build_graph(window_events) | ||
|
|
||
| # Convert to PyTorch Geometric | ||
| pyg_data = from_networkx(G) | ||
|
|
||
| # Extract features | ||
| node_features, edge_features = self.graph_builder.extract_features(G) | ||
|
|
There was a problem hiding this comment.
Prevent PyG conversion crash on string attributes.
from_networkx tries to convert all node/edge attributes to tensors. Because the graph carries strings (paths, node types) and pandas timestamps, this call raises TypeError: new(): invalid data type 'str' on the very first batch. Please strip the attributes before conversion, e.g. by passing empty attr lists:
- pyg_data = from_networkx(G)
+ pyg_data = from_networkx(G, node_attrs=[], edge_attrs=[])📝 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.
| G = self.graph_builder.build_graph(window_events) | |
| # Convert to PyTorch Geometric | |
| pyg_data = from_networkx(G) | |
| # Extract features | |
| node_features, edge_features = self.graph_builder.extract_features(G) | |
| G = self.graph_builder.build_graph(window_events) | |
| # Convert to PyTorch Geometric | |
| pyg_data = from_networkx(G, node_attrs=[], edge_attrs=[]) | |
| # Extract features | |
| node_features, edge_features = self.graph_builder.extract_features(G) |
🤖 Prompt for AI Agents
In ai/datasets/nerrf_dataset.py around lines 339 to 346, from_networkx is
currently converting all node/edge attributes (including strings and pandas
timestamps) which causes a TypeError; fix this by removing or excluding
non-tensorable attributes before conversion — call from_networkx with explicit
empty attribute lists (e.g. from_networkx(G, node_attrs=[], edge_attrs=[])) or
strip those attributes from G prior to conversion so only numeric attributes are
turned into tensors.
| torch>=2.0.0 | ||
| torch-geometric>=2.3.0 | ||
| torch-scatter>=2.1.0 | ||
| torch-sparse>=0.6.0 |
There was a problem hiding this comment.
Tighten the torch lower bound to avoid known CVEs.
Allowing torch 2.0.0 pulls in multiple critical vulnerabilities (RCE, DoS, use-after-free; e.g., GHSA-53q9-r3pm-6pq6, GHSA-3749-ghw9-m3mg). Please raise the minimum to a patched release (e.g., torch>=2.2.1) or explicitly exclude the affected 2.0.x range.
🤖 Prompt for AI Agents
In ai/requirements.txt around lines 3 to 6, the torch lower bound is too
permissive (torch>=2.0.0) and may pull in versions with known CVEs; update the
requirement to a patched minimum such as torch>=2.2.1 (or use a version spec
that excludes the vulnerable 2.0.x range, e.g. >=2.0.0,!=2.0.*,<2.2.1 is not
acceptable — prefer >=2.2.1) and keep the other packages as-is so deployments no
longer install the vulnerable torch releases.
| class rocksdb: | ||
| @staticmethod | ||
| def open(*args, **kwargs): | ||
| raise ImportError("RocksDB not installed. Install with: pip install rocksdb") | ||
| from google.protobuf.timestamp_pb2 import Timestamp |
There was a problem hiding this comment.
Fix the RocksDB fallback path.
When the native bindings are absent we define a dummy rocksdb with only open(), but the rest of the module still calls rocksdb.Options, WriteBatch, etc. The first instantiation hits AttributeError instead of the intended helpful ImportError, so the optional dependency story is broken. Please guard RocksDBEventStore.__init__ (and other entry points) to raise a clear ImportError when ROCKSDB_AVAILABLE is false, or provide stub attributes that mirror the real API surface just enough to fail gracefully. For example:
@@
-except ImportError:
- ROCKSDB_AVAILABLE = False
- # Create dummy class for environments without RocksDB
- class rocksdb:
- @staticmethod
- def open(*args, **kwargs):
- raise ImportError("RocksDB not installed. Install with: pip install rocksdb")
+except ImportError:
+ ROCKSDB_AVAILABLE = False
+ rocksdb = None
@@
def __init__(self, db_path: str, create_if_missing: bool = True):
"""
Initialize RocksDB connection.
"""
+ if not ROCKSDB_AVAILABLE:
+ raise ImportError("RocksDB bindings not installed. Install with: pip install rocksdb")This keeps consumers on dependency-free paths from tripping over missing attributes.
Also applies to: 53-87
🧰 Tools
🪛 Ruff (0.13.3)
34-34: Unused static method argument: args
(ARG004)
34-34: Unused static method argument: kwargs
(ARG004)
35-35: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In ai/utils/rocksdb_store.py around lines 32-36 and also covering lines 53-87,
the current fallback only defines rocksdb.open() so later references to
rocksdb.Options, WriteBatch, DB, etc. raise AttributeError instead of a clear
ImportError; update the module so that when the native bindings are missing
(ROCKSDB_AVAILABLE is false) you either 1) raise ImportError at
RocksDBEventStore.__init__ and any other public entry points (guard constructors
and factory functions) with a clear message instructing to install the rocksdb
package, or 2) expand the fallback rocksdb stub to include minimal placeholder
attributes (Options, WriteBatch, DB with open method signatures) that
immediately raise ImportError with the same clear message when any of them are
accessed/instantiated; choose one approach and apply it consistently for lines
32-36 and 53-87 so consumers get a single, helpful ImportError instead of
AttributeError.
1 participant
Implements the PyTorch Geometric dataset loader for M2 AI Spike milestone. This delivers the core data pipeline that converts filesystem events from RocksDB into temporal graphs ready for neural network training.
What's Changed
Testing Commands
1. Quick Functionality Test
2. Dataset Pipeline Test
3. Check Clean Structure
4. Verify Dependencies
Expected Output
Usage Example
M2 AI Spike Ready
Status
Fixes #13 - Ready for M2 AI Spike neural network development.
Summary by CodeRabbit
New Features
Chores