Avoid reference cycles in graphs

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.

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 == []