diff --git a/src/cyclebane/graph.py b/src/cyclebane/graph.py index 78d5da0..6814cc0 100644 --- a/src/cyclebane/graph.py +++ b/src/cyclebane/graph.py @@ -13,6 +13,38 @@ from .value_array import Grouping +def _uncached(name: str) -> property: + """Return an uncached version of one of the cached properties of `nx.DiGraph`.""" + return property(getattr(nx.DiGraph, name).func) + + +class DiGraph(nx.DiGraph): + """ + A NetworkX DiGraph that does not cache its report views. + + NetworkX caches views such as ``G.edges`` or ``G.out_degree`` in ``G.__dict__``. + Each of those views holds a reference back to ``G``, i.e., merely looking at a + graph places it in a reference cycle, so it can be freed only by the cyclic + garbage collector. That collector is triggered by the *number* of allocated + objects, so a handful of graphs holding large node data may survive for a long + time. Cyclebane builds and discards graphs in most of its operations, and the node + data can be arbitrarily large, which would lead to unbounded memory growth in, + e.g., streaming applications. The views are cheap to create, so not caching them + keeps every graph reclaimable by reference counting. + + As a consequence the views are not identical across accesses. They compare equal + where NetworkX defines equality for them, but the degree views do not, i.e., + ``G.degree == G.degree`` is False. + """ + + edges = _uncached('edges') + out_edges = _uncached('out_edges') + in_edges = _uncached('in_edges') + degree = _uncached('degree') + out_degree = _uncached('out_degree') + in_degree = _uncached('in_degree') + + def _get_unique_sink(graph: nx.DiGraph) -> Hashable: sink_nodes = [node for node in graph.nodes if graph.out_degree(node) == 0] if len(sink_nodes) != 1: @@ -214,13 +246,14 @@ def __init__(self, graph: nx.DiGraph, *, node_values: NodeValues | None = None): Parameters ---------- graph: - The directed graph representing the data flow. + The directed graph representing the data flow. Unless it is a + :py:class:`DiGraph`, a copy is made, as for `nx.DiGraph.copy`. node_values: A mapping from source node names to array-like objects. The implementation assumes that the graph has been setup correctly. Do not use this argument unless you know what you are doing. """ - self.graph = graph + self.graph = graph if isinstance(graph, DiGraph) else DiGraph(graph) self._node_values = node_values or NodeValues({}) def copy(self) -> Graph: @@ -364,7 +397,7 @@ def _from_orig_key( def by_position(self, index_name: IndexName) -> PositionalIndexer: return PositionalIndexer(self, index_name) - def to_networkx(self, value_attr: str = 'value') -> nx.DiGraph: + def to_networkx(self, value_attr: str = 'value') -> DiGraph: """ Convert to a NetworkX graph, spelling out the internal array structures as explicit nodes. diff --git a/tests/memory_test.py b/tests/memory_test.py new file mode 100644 index 0000000..1d658fc --- /dev/null +++ b/tests/memory_test.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +""" +Tests ensuring that graphs are reclaimable by reference counting. + +Graphs (and thus the data stored on their nodes) that are part of a reference cycle +are freed only by the cyclic garbage collector. That collector is triggered by the +number of allocated objects, so few but large objects may survive indefinitely. +""" + +import gc +import weakref +from collections.abc import Iterator +from contextlib import contextmanager + +import networkx as nx +import pytest + +import cyclebane as cb + + +@contextmanager +def refcount_only_gc() -> Iterator[None]: + """Disable the cyclic garbage collector, leaving only reference counting.""" + gc.collect() + gc.disable() + try: + yield + finally: + gc.enable() + + +def live_graphs() -> list[nx.DiGraph]: + return [obj for obj in gc.get_objects() if isinstance(obj, nx.DiGraph)] + + +class Payload: + """Stand-in for the (potentially large) data stored on a node.""" + + +def set_value(graph: cb.Graph, node: str, value: Payload) -> None: + """Set a node value the way, e.g., Sciline inserts a parameter into a pipeline.""" + branch = nx.DiGraph() + branch.add_node(node, value=value) + graph[node] = cb.Graph(branch) + + +@pytest.fixture +def graph() -> cb.Graph: + return cb.Graph(nx.DiGraph([('a', 'b'), ('b', 'c')])) + + +def test_graph_does_not_form_cycle_when_taking_views() -> None: + with refcount_only_gc(): + graph = cb.Graph(nx.DiGraph([('a', 'b'), ('b', 'c')])) + ref = weakref.ref(graph.graph) + list(graph.graph.edges) + list(graph.graph.in_edges('b')) + assert graph.graph.out_degree('c') == 0 + del graph + assert ref() is None + + +def test_setitem_leaves_no_garbage(graph: cb.Graph) -> None: + branch = cb.Graph(nx.DiGraph([('x', 'b')])) + with refcount_only_gc(): + before = {id(g) for g in live_graphs()} + graph['b'] = branch + garbage = [ + g for g in live_graphs() if id(g) not in before and g is not graph.graph + ] + assert garbage == [] + + +def test_setitem_frees_value_of_replaced_node(graph: cb.Graph) -> None: + with refcount_only_gc(): + set_value(graph, 'a', Payload()) + ref = weakref.ref(graph.graph.nodes['a']['value']) + set_value(graph, 'a', Payload()) + assert ref() is None + + +def test_delitem_leaves_graph_reclaimable(graph: cb.Graph) -> None: + with refcount_only_gc(): + del graph['c'] + ref = weakref.ref(graph.graph) + graph.graph = None + assert ref() is None + + +def test_to_networkx_result_is_freed_after_inspecting_it(graph: cb.Graph) -> None: + mapped = graph.map({'a': [1, 2, 3]}).reduce('c', name='sum') + with refcount_only_gc(): + result = mapped.to_networkx() + # Consumers of the returned graph inspect it, which materializes views. + assert len(list(result.edges)) > 0 + assert result.in_degree('sum') == 3 + ref = weakref.ref(result) + del result + assert ref() is None + + +def test_map_reduce_to_networkx_leaves_no_garbage(graph: cb.Graph) -> None: + graph.map({'a': [1, 2, 3]}).reduce('c', name='sum').to_networkx() + with refcount_only_gc(): + before = {id(g) for g in live_graphs()} + graph.map({'a': [1, 2, 3]}).reduce('c', name='sum').to_networkx() + garbage = [g for g in live_graphs() if id(g) not in before] + assert garbage == []