-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdynamic_graph.py
More file actions
290 lines (230 loc) · 9.41 KB
/
Copy pathdynamic_graph.py
File metadata and controls
290 lines (230 loc) · 9.41 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
"""Lock-free dynamic graph for live GNN training.
This module provides a Pythonic wrapper around the Rust DynamicGraph,
which uses C-tree neighbor lists with arena allocation for zero-alloc
reads and lock-free concurrent access.
"""
from __future__ import annotations
import os
from typing import TYPE_CHECKING, Any
import numpy as np
import numpy.typing as npt
from aethergraph._core import DynamicGraph as _DynamicGraph
from aethergraph._ids import _to_uint32_ids
if TYPE_CHECKING:
from aethergraph._core import GraphSnapshot
from aethergraph.graph import Graph
class DynamicGraph:
"""Lock-free dynamic graph with C-tree neighbor lists.
Supports concurrent edge inserts and neighbor reads for live
GNN training on evolving graphs (e.g., Reddit's social graph).
Each vertex's neighbor list is a balanced tree of sorted,
cache-line-sized chunks. 90% of nodes (degree < 15) fit in a
single chunk -- one cache line read, identical cost to static CSR.
Attributes:
_inner: The underlying Rust DynamicGraph instance.
Example:
>>> g = DynamicGraph(num_vertices=1_000_000, arena_mb=512)
>>> g.insert_edge(0, 42)
True
>>> g.degree(0)
1
>>> g.neighbors(0)
array([42])
"""
_inner: _DynamicGraph
def __init__(self, num_vertices: int, arena_mb: int = 256) -> None:
"""Create an empty dynamic graph.
Args:
num_vertices: Number of vertices (fixed at construction).
arena_mb: Arena capacity in megabytes (max 32768). Superseded
nodes are recycled once no reader or snapshot can observe
them, so steady-state usage tracks live edges plus a
bounded recycling lag.
"""
self._inner = _DynamicGraph(num_vertices=num_vertices, arena_mb=arena_mb)
@classmethod
def open_with_wal(
cls,
path: str | os.PathLike[str],
num_vertices: int,
arena_mb: int = 256,
) -> DynamicGraph:
"""Open a DynamicGraph backed by an append-only write-ahead log.
Existing records at ``path`` are replayed before this returns; if the
file ends in a torn record (mid-write crash), the trailing bytes are
truncated. Every subsequent ``insert_edge*`` call appends to the log
and fsyncs at writer-guard close.
Args:
path: WAL file path. Created if it does not exist.
num_vertices: Number of vertices (fixed at construction).
arena_mb: Arena capacity in megabytes.
Returns:
DynamicGraph with ``current_epoch`` equal to the number of
replayed records, ready to accept new writers.
Raises:
OSError: I/O failure opening, reading, or writing the WAL.
ValueError: File exists but is not a valid AetherGraph WAL, or
a replayed record exceeds ``num_vertices``.
RuntimeError: Arena filled up during replay.
"""
obj = cls.__new__(cls)
obj._inner = _DynamicGraph.open_with_wal(path, num_vertices, arena_mb)
return obj
@property
def current_epoch(self) -> int:
"""Monotonic version counter — advances on every successful writer
commit. Pin before a multi-source read to coordinate consistency
with other subsystems sharing the same ``EpochClock``."""
return self._inner.current_epoch
@classmethod
def from_edges(
cls,
num_vertices: int,
src: npt.NDArray[Any],
dst: npt.NDArray[Any],
arena_mb: int = 256,
) -> DynamicGraph:
"""Build a DynamicGraph from edge arrays.
Args:
num_vertices: Number of vertices.
src: Source vertex array. Range-checked and converted to uint32.
dst: Destination vertex array. Range-checked and converted to uint32.
arena_mb: Arena capacity in megabytes.
Returns:
DynamicGraph with all edges inserted.
Raises:
ValueError: If src and dst have different lengths, or an ID is
negative or exceeds the uint32 range.
"""
src_arr = _to_uint32_ids(src, "src")
dst_arr = _to_uint32_ids(dst, "dst")
obj = cls.__new__(cls)
obj._inner = _DynamicGraph.from_edges(num_vertices, src_arr, dst_arr, arena_mb)
return obj
def insert_edge(self, src: int, dst: int) -> bool:
"""Insert a directed edge from src to dst.
Args:
src: Source vertex ID.
dst: Destination vertex ID.
Returns:
True if the edge was new, False if it already existed.
Raises:
RuntimeError: If the arena is full.
ValueError: If src or dst is >= num_vertices.
"""
return self._inner.insert_edge(src, dst)
def insert_edges(self, src: npt.NDArray[Any], dst: npt.NDArray[Any]) -> int:
"""Batch-insert edges from arrays.
Args:
src: Source vertex array. Range-checked and converted to uint32.
dst: Destination vertex array. Range-checked and converted to uint32.
Returns:
Number of new edges inserted (duplicates are skipped).
Raises:
ValueError: If src and dst have different lengths, an ID is
negative or exceeds the uint32 range, or an edge references
a vertex >= num_vertices.
RuntimeError: If the arena is full.
"""
src_arr = _to_uint32_ids(src, "src")
dst_arr = _to_uint32_ids(dst, "dst")
return self._inner.insert_edges(src_arr, dst_arr)
@property
def num_vertices(self) -> int:
"""Number of vertices (fixed at construction)."""
return self._inner.num_vertices
@property
def num_edges(self) -> int:
"""Total number of edges."""
return self._inner.num_edges
@property
def arena_used(self) -> int:
"""Arena bytes currently used."""
return self._inner.arena_used
@property
def arena_capacity(self) -> int:
"""Arena total capacity in bytes."""
return self._inner.arena_capacity
def degree(self, vertex: int) -> int:
"""Get the degree (number of outgoing edges) for a vertex.
Args:
vertex: Vertex ID to query.
Returns:
Number of outgoing edges from this vertex.
"""
return self._inner.degree(vertex)
def has_edge(self, src: int, dst: int) -> bool:
"""Check whether edge (src -> dst) exists.
Args:
src: Source vertex ID.
dst: Destination vertex ID.
Returns:
True if the edge exists.
"""
return self._inner.has_edge(src, dst)
def neighbors(self, vertex: int) -> npt.NDArray[np.int64]:
"""Get sorted neighbor array for a vertex.
Returns int64 dtype for PyTorch compatibility.
Args:
vertex: Vertex ID to query.
Returns:
Sorted numpy array of neighbor IDs (dtype=int64).
"""
return self._inner.neighbors(vertex)
def snapshot(self) -> Graph:
"""Create a frozen CSR snapshot for use with NeighborLoader.
Collects all edges from the C-tree neighbor lists into a static
CSR graph. O(V + E) time -- call once per epoch, not per batch.
The returned Graph is completely independent of this DynamicGraph.
Edges inserted after the snapshot is taken will not appear in it.
Returns:
A static :class:`~aethergraph.graph.Graph` usable with
:class:`~aethergraph.pytorch.loader.NeighborLoader` and all
existing sampling infrastructure.
Example:
>>> dg = DynamicGraph(num_vertices=1000)
>>> dg.insert_edge(0, 1)
True
>>> graph = dg.snapshot()
>>> graph.num_nodes
1000
>>> graph.degree(0)
1
"""
# `Graph` is `aethergraph._core.CsrGraph` directly; `snapshot()`
# already returns one, so this is just a passthrough.
return self._inner.snapshot()
def acquire(self) -> GraphSnapshot:
"""Pin the latest committed snapshot.
The snapshot is immutable and strictly serializable: it reflects
every committed insert up to its epoch and never changes while
inserts continue concurrently. Reads on it are lock-free. Holding
it defers arena recycling of its state -- drop it when done.
Returns:
A :class:`GraphSnapshot` pinned at the current epoch.
Example:
>>> dg = DynamicGraph(num_vertices=100)
>>> dg.insert_edge(0, 1)
True
>>> snap = dg.acquire()
>>> dg.insert_edge(0, 2)
True
>>> snap.degree(0)
1
>>> dg.degree(0)
2
"""
return self._inner.acquire()
def __repr__(self) -> str:
"""Return detailed string representation."""
return (
f"DynamicGraph(num_vertices={self.num_vertices}, "
f"num_edges={self.num_edges}, "
f"arena={self.arena_used // (1024 * 1024)}/{self.arena_capacity // (1024 * 1024)}MB)"
)
def __str__(self) -> str:
"""Return short string description."""
return f"DynamicGraph with {self.num_vertices:,} vertices and {self.num_edges:,} edges"
def __len__(self) -> int:
"""Return number of vertices."""
return self.num_vertices