Coverage for core / src / sensorkit / data / graph.py: 89%
258 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import asyncio
5import collections
6import contextlib
7import io
8import weakref
9from abc import ABC, abstractmethod
10from collections.abc import Buffer
11from typing import (
12 Any,
13 AsyncGenerator,
14 Callable,
15 ClassVar,
16 Literal,
17 Self,
18 Unpack,
19 cast,
20 final,
21 overload,
22)
24from loguru import logger
25from pydantic import BaseModel, Field, model_validator
27from sensorkit.common.model import ModelRegistry, RegistryBaseModel
28from sensorkit.data.context import Context
29from sensorkit.data.streams import (
30 BufferReader,
31 BufferWriter,
32 StreamReader,
33 StreamWriter,
34 create_connected_streams,
35)
38class DataGraphError(Exception): ...
39class DataGraphCycleError(DataGraphError): ...
40class DataGraphSourceError(DataGraphError): ...
43type ReceiveKind = Literal["stream", "buffer"] | None
46class DataFlow:
47 """An edge in a DataGraph."""
49 def __init__(self):
50 self.send_called = False
51 self.send_ready = asyncio.Event()
52 self.receive_ready = asyncio.Event()
53 self.receive_kind: ReceiveKind = None
54 self.receive_result: tuple[Context, Buffer | StreamReader] | None = None
56 @overload
57 async def send(
58 self,
59 context: Context,
60 arg: StreamReader,
61 ): ...
63 @overload
64 async def send(
65 self,
66 context: Context,
67 arg: Buffer,
68 ): ...
70 @overload
71 async def send(
72 self,
73 context: Context,
74 ) -> StreamWriter: ...
76 async def send(
77 self,
78 context: Context,
79 arg: Buffer | StreamReader | None = None,
80 ) -> StreamWriter | None:
81 """Send data through this edge to the consumer.
83 Pass a Buffer or StreamReader directly, or omit *arg* to receive a StreamWriter that
84 the caller can write into. May only be called once per edge instance.
85 """
86 if self.send_called:
87 raise RuntimeError("send may only be called once")
89 self.send_called = True
91 # Wait for the consumer to call the `receive` method.
92 await self.receive_ready.wait()
94 # Implement each of the data passing combinations based on the type of the input argument
95 # and the requested receive kind. Note that a `None` receive kind means that the argument
96 # type dictates the effective receive kind, enabling "passthrough" semantics.
97 match (arg, self.receive_kind):
98 # When a buffer is given, we either pass it straight through to the consumer by
99 # reference or we return a StreamReader backed by it. In the former case, ownership
100 # of the buffer is transferred to the consumer. In the latter case, the buffer must
101 # be immutable.
102 case (Buffer(), "buffer" | None):
103 # Buffer to buffer. We do this by reference to avoid a copy.
104 self._send_ready(context, arg)
105 case (Buffer(), "stream"):
106 # Buffer to stream. Return a StreamReader backed by the input buffer.
107 self._send_ready(context, BufferReader(arg))
109 # When a StreamReader is given, we either drain it immediately (buffer case) or pass
110 # it straight through to the consumer.
111 case (StreamReader(), "buffer"):
112 bio = io.BytesIO()
114 async for chunk in arg:
115 # FIXME: event loop starvation possible here
116 bio.write(chunk)
118 self._send_ready(context, bio.getvalue())
119 case (StreamReader(), "stream" | None):
120 self._send_ready(context, arg)
122 # When no argument is supplied, we return a StreamWriter that writes data to the
123 # requested destination.
124 case (None, "buffer"):
125 # Stream to buffer. Return a StreamWriter that will fill a buffer.
126 writer = BufferWriter()
127 fut = writer.get_future()
128 fut.add_done_callback(lambda _: self._send_ready(context, fut.result()))
129 return writer
130 case (None, "stream" | None):
131 # Stream to stream.
132 reader, writer = create_connected_streams()
133 self._send_ready(context, reader)
134 return writer
135 case _:
136 raise TypeError()
138 return None
140 def _send_ready(self, *result: Unpack[tuple[Context, Buffer | StreamReader]]):
141 # Signal the consumer that the receive result is ready.
142 self.receive_result = result
143 self.send_ready.set()
145 @overload
146 async def receive(
147 self,
148 kind: Literal["stream"],
149 ) -> tuple[Context, StreamReader]: ...
151 @overload
152 async def receive(
153 self,
154 kind: Literal["buffer"],
155 ) -> tuple[Context, Buffer]: ...
157 @overload
158 async def receive(
159 self,
160 ) -> tuple[Context, Buffer | StreamReader]: ...
162 async def receive(
163 self,
164 kind: ReceiveKind | None = None,
165 ) -> tuple[Context, StreamReader | Buffer]:
166 """Wait for the producer to send data and return `(context, data)`.
168 *kind* controls whether data is delivered as a buffer or stream; `None` defers to
169 the producer's choice. May only be called once per edge instance.
170 """
171 if self.receive_ready.is_set():
172 raise RuntimeError("receive may only be called once")
174 # Signal that we are ready to receive so the producer can proceed.
175 self.receive_kind = kind
176 self.receive_ready.set()
178 # Wait for the producer to provide a result for us and return it.
179 await self.send_ready.wait()
180 return self.receive_result
183class DataOp(RegistryBaseModel, ABC):
184 """A node in a DataGraph."""
185 op: Literal[None] | str = None
186 output: list[str] = Field(default_factory=list)
187 registry: ClassVar[ModelRegistry[DataOp]] = ModelRegistry(discriminator="op")
189 @classmethod
190 def model_registry(cls):
191 return cls.registry
193 @abstractmethod
194 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
195 """Execute this node, consuming *incoming* edges and producing *outgoing* edges."""
196 ...
199class SourceOp(DataOp, ABC):
200 """A DataGraph node that triggers graph runs."""
202 @abstractmethod
203 def graph_source(self) -> AsyncGenerator[None, DataFlow]:
204 """Async generator that triggers each graph run by sending a DataFlow edge."""
205 ...
207 @final
208 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
209 assert len(outgoing) == 1
210 context, data = await incoming[0].receive()
211 await outgoing[0].send(context, data)
214class SinkOp(DataOp, ABC):
215 """A DataGraph node that consumes graph runs."""
217 @abstractmethod
218 async def graph_sink(self, incoming: DataFlow):
219 """Consume the final DataFlow edge to complete a graph run."""
220 ...
222 @final
223 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
224 assert len(incoming) == 1 and not outgoing
225 await self.graph_sink(incoming[0])
228def _parse_simple(simple: list[dict[str, Any]]):
229 """Parse a simple representation of a linear DataGraph."""
230 counts = collections.Counter()
231 nodes = {}
232 prev_id = None
234 for node in simple:
235 if "output" in node:
236 raise ValueError("Simple DataGraph representations cannot contain 'output' fields")
238 op = node["op"]
239 counts[op] += 1
240 node_id = f"{op}_{counts[op]}"
242 if prev_id:
243 nodes[prev_id]["output"] = [node_id]
245 nodes[node_id] = dict(node)
246 prev_id = node_id
248 return nodes
251def _describe_simple_form(schema: dict[str, Any]):
252 """Declare the shorthand accepted in place of a node mapping.
254 Args:
255 schema: The generated JSON Schema for the model.
256 """
257 schema["properties"]["simple"] = {
258 "type": "array",
259 "items": {"type": "object"},
260 "description": "Operations to run in order, in place of 'nodes'.",
261 }
264class DataGraph(BaseModel, json_schema_extra=_describe_simple_form):
265 """A graph of operations on data."""
267 nodes: dict[str, DataOp] = Field(default_factory=dict)
269 @model_validator(mode="before")
270 def _validate(cls, values):
271 if simple := values.get("simple"):
272 if "nodes" in values:
273 raise ValueError("Only one of 'nodes' or 'simple' can be defined")
275 del values["simple"]
276 values["nodes"] = dict(_parse_simple(simple))
278 return values
280 def __init__(self, /, **data: Any):
281 super().__init__(**data)
282 self._inputs: dict[str, list[DataOp]] = collections.defaultdict(list)
283 self._sources: set[str] = set()
284 self._sinks: set[str] = set()
285 self._runners: set[DataGraphRunner] = set()
287 for name, node in self.nodes.items():
288 self._index_node(name, node)
290 def add(self, name: str, node: DataOp):
291 """Add a named node to the graph. Raises DataGraphError if already started or duplicate."""
292 if self._runners:
293 raise DataGraphError("cannot modify graph while executing")
295 if name in self.nodes:
296 raise DataGraphError(f"node '{name}' already defined")
298 self.nodes[name] = node
299 self._index_node(name, node)
301 def _index_node(self, name: str, node: DataOp):
302 if isinstance(node, SourceOp):
303 self._sources.add(name)
305 if isinstance(node, SinkOp):
306 self._sinks.add(name)
308 for output in node.output:
309 self._inputs[output].append(node)
311 def find_source[M: SourceOp](self, kind: type[M] | None = None) -> M:
312 """Return the source node of the given type, or the sole source if *kind* is omitted."""
313 if kind:
314 for source in self._sources:
315 source_node = self.nodes[source]
317 if isinstance(source_node, kind):
318 return source_node
320 raise DataGraphSourceError(f"no source of type '{kind.__name__}' defined")
322 if len(self._sources) != 1:
323 raise DataGraphError(f"expected exactly one source, got {len(self._sources)}")
325 return cast(SourceOp, self.nodes[next(iter(self._sources))])
327 def app_source(self):
328 """Return the AppSource node in this graph."""
329 from sensorkit.data.local import AppSource
331 return self.find_source(AppSource)
333 def find_sink[M: SinkOp](self, kind: type[M] | None = None) -> M:
334 """Return the sink node of the given type, or the sole sink if *kind* is omitted."""
335 if kind:
336 for sink in self._sinks:
337 sink_node = self.nodes[sink]
339 if isinstance(sink_node, kind):
340 return sink_node
342 raise DataGraphSourceError(f"no sink of type '{kind.__name__}' defined")
344 if len(self._sinks) != 1:
345 raise DataGraphError(f"expected exactly one sink, got {len(self._sinks)}")
347 return cast(SinkOp, self.nodes[next(iter(self._sinks))])
349 def app_sink(self):
350 """Return the AppSink node in this graph."""
351 from sensorkit.data.local import AppSink
353 return self.find_sink(AppSink)
355 def _validate_component(self, source: str):
356 if source not in self._sources:
357 raise DataGraphSourceError(f"source '{source}' is undefined")
359 queue = collections.deque([source])
360 visited = set()
362 while queue:
363 node_name = queue.popleft()
365 if node_name in visited:
366 raise DataGraphCycleError(f"cycle detected at node '{node_name}'")
368 visited.add(node_name)
370 if node_name not in self.nodes:
371 raise DataGraphError(f"node '{node_name}' is undefined")
373 node = self.nodes[node_name]
375 if node_name != source and isinstance(node, SourceOp):
376 # FIXME: This may actually not be an error condition. There may be legitimate use
377 # cases for having multiple sources in a traversal (only one actually acting
378 # as a source at a time), and the way SourceOp is implemented would
379 # naturally support no-op passthrough of intermediate sources.
380 raise DataGraphSourceError(f"multiple sources in traversal at '{source}'")
382 for next_name in node.output:
383 queue.append(next_name)
385 return visited
387 def start(self, *, task_group: asyncio.TaskGroup = asyncio):
388 """Validate and start all source components of the graph in the given task group."""
389 if not self._sources:
390 raise DataGraphSourceError("no sources defined")
392 for source in self._sources:
393 self._validate_component(source)
394 runner = DataGraphRunner(self, source, task_group)
395 self._runners.add(runner)
396 runner.start(done_callback=self._runners.discard)
398 async def stop(self):
399 """Cancel all running DataGraphRunner instances and wait for them to finish."""
400 await asyncio.gather(
401 *(runner.stop() for runner in self._runners.copy()),
402 return_exceptions=True
403 )
406class DataGraphRunner:
407 """Invokes runs of one component of a DataGraph."""
409 def __init__(self, graph: DataGraph, source: str, task_group: asyncio.TaskGroup):
410 self._graph = weakref.ref(graph)
411 self._source = source
412 self._task_group = task_group
413 self._exec_task: asyncio.Task | None = None
415 def start(self, done_callback: Callable[[Self], Any] | None = None):
416 """Start executing the graph component, optionally invoking *done_callback* on exit."""
417 if self._exec_task:
418 raise RuntimeError("DataGraph runner already started")
420 # Get a hard reference to the DataGraph.
421 graph = self._graph()
423 if not graph:
424 raise RuntimeError("DataGraph was disposed")
426 # Ensure that exceptions are not propagated to the main task group, as runs should
427 # keep the service alive until they complete but errors should not kill the service.
428 self._exec_task = self._task_group.create_task(self._run_graph(graph))
429 logger.info("DataGraph runner started")
431 def _finalize_graph_run(t: asyncio.Task):
432 if t.cancelled():
433 logger.info("DataGraph runner shut down")
434 elif e := t.exception():
435 logger.error("DataGraph runner exited due to error")
436 logger.opt(exception=e).debug("DataGraph runner failed")
437 else:
438 logger.warning("DataGraph runner exited")
440 if done_callback:
441 done_callback(self)
443 self._exec_task.add_done_callback(_finalize_graph_run)
445 async def stop(self):
446 """Cancel the graph runner task and wait for it to finish."""
447 self._exec_task.cancel()
449 with contextlib.suppress(asyncio.CancelledError):
450 await self._exec_task
452 async def _run_graph(self, graph: DataGraph):
453 # Invoke the source and prime the generator it returns.
454 source = cast(SourceOp, graph.nodes[self._source])
455 gen = source.graph_source()
456 running = True
458 await gen.asend(None)
460 # Execute the graph for each iteration of the source op.
461 while running:
462 shutdown = False
464 # Create a new edge to be injected as the incoming edge to the source node.
465 edge = DataFlow()
467 try:
468 # Run the graph in a TaskGroup so that if any node fails, all nodes for this
469 # run are cancelled.
470 async with asyncio.TaskGroup() as tg:
471 self._run_tasks(graph, source, edge, task_group=tg)
473 try:
474 # Wait for the SourceOp to trigger and feed the incoming edge.
475 await gen.asend(edge)
476 except asyncio.CancelledError:
477 # FIXME: This check is likely not sufficient to guarantee no deadlock.
478 if not edge.send_called:
479 raise
481 # Wait for tasks to complete before shutting down. We do this by
482 # suppressing the cancellation and allowing the task group to complete.
483 # Then we fall through to re-raise via the `shutdown` flag.
484 asyncio.current_task().uncancel()
485 shutdown = True
486 except* StopAsyncIteration:
487 logger.debug("DataGraph source stopped producing")
488 running = False
489 except* Exception:
490 # We log this as an exception so that the traceback is preserved.
491 logger.exception("DataGraph run encountered an error")
493 if shutdown:
494 raise asyncio.CancelledError()
496 def _run_tasks(
497 self,
498 graph: DataGraph,
499 source: DataOp,
500 source_incoming: DataFlow,
501 *,
502 task_group: asyncio.TaskGroup,
503 ):
504 tasks: list[asyncio.Task] = []
505 incoming: dict[str, list[DataFlow]] = collections.defaultdict(list)
506 outgoing: dict[str, list[DataFlow]] = collections.defaultdict(list)
508 for name, node in graph.nodes.items():
509 for output in node.output:
510 edge = DataFlow()
511 incoming[output].append(edge)
512 outgoing[name].append(edge)
514 if node is source:
515 incoming[name].append(source_incoming)
517 for name, node in graph.nodes.items():
518 tasks.append(
519 task_group.create_task(node.process(incoming[name], outgoing[name]))
520 )
522 return tasks