Make nodes and operations pickable.

python/mlcroissant/mlcroissant/_src/operation_graph/base_operation.py

@@ -23,33 +23,6 @@ def __init__(self):
         super().__init__(self)
         self.last_operations: dict[Node, Operation] = {}
 
-    def add_node(self, operation: Operation) -> None:
-        """Overloads nx.add_node to keep track of the last operations."""
-        if not self.has_node(operation):
-            super().add_node(operation)
-
-    def add_edge(self, operation1: Operation, operation2: Operation) -> None:
-        """Overloads nx.add_node to keep track of the last operations."""
-        if not self.has_edge(operation1, operation2):
-            super().add_edge(operation1, operation2)
-
-    @property
-    def nodes(self) -> Iterable[Operation]:
-        """Overloads nx.nodes to return an interator of operations."""
-        return super().nodes()
-
-    def is_leaf(self, operation: Operation | None) -> bool:
-        """Tests whether an operation is a leaf in the graph."""
-        return self.out_degree(operation) == 0
-
-    def entry_operations(self) -> list[Operation]:
-        """Lists all operations without a parent in the graph of operations."""
-        return [
-            operation
-            for operation, indegree in self.in_degree(self.nodes)
-            if indegree == 0 and isinstance(operation, Operation)
-        ]
-
 
 @dataclasses.dataclass(frozen=True, repr=False)
 class Operation(abc.ABC):
@@ -116,3 +89,31 @@ def __rrshift__(
         for left_operation in left_operations:
             self.operations.add_edge(left_operation, right_operation)
         return right_operation
+
+    def __reduce__(self):
+        """Allows pickling the operation.
+
+        `self.operations` is stored separately in the state to break the cyclic
+        reference between Operation and Operations. We could refactor the codebase to
+        not have this dependency, but it would be a bigger change and it's convenient
+        to be able to reference all operations from a single operation.
+        """
+        state = self.__getstate__()
+        args = tuple(state.values())
+        return (
+            self.__class__,
+            args,
+            {"operations": self.operations},
+        )
+
+    def __getstate__(self):
+        state = {}
+        for field in dataclasses.fields(self):
+            if field.name == "operations":
+                state[field.name] = Operations()
+            else:
+                state[field.name] = getattr(self, field.name)
+        return state
+
+    def __setstate__(self, state):
+        object.__setattr__(self, "operations", state["operations"])

python/mlcroissant/mlcroissant/_src/operation_graph/graph.py

@@ -208,7 +208,11 @@ def from_nodes(cls, ctx: Context, metadata: Node) -> "OperationGraph":
                 _add_operations_for_field(operations, node)
 
         # Attach all entry nodes to a single `start` node
-        entry_operations = operations.entry_operations()
+        entry_operations = [
+            operation
+            for operation, indegree in operations.in_degree(operations.nodes)
+            if indegree == 0 and isinstance(operation, Operation)
+        ]
         init_operation = InitOperation(operations=operations, node=metadata)
         for entry_operation in entry_operations:
             operations.add_edge(init_operation, entry_operation)

python/mlcroissant/mlcroissant/_src/operation_graph/operations/read_test.py

@@ -1,6 +1,7 @@
 """read_test module."""
 
 import pathlib
+import pickle
 import tempfile
 from unittest import mock
 
@@ -84,3 +85,14 @@ def test_reading_method():
     assert _reading_method(empty_file_object, (filename,)) == ReadingMethod.NONE
     with pytest.raises(ValueError):
         _reading_method(empty_file_object, (content_field, lines_field))
+
+
+def test_pickable():
+    operation = Read(
+        operations=operations(),
+        node=empty_file_object,
+        folder=epath.Path("/foo/bar"),
+        fields=(),
+    )
+    operation = pickle.loads(pickle.dumps(operation))
+    assert operation.folder == epath.Path("/foo/bar")

python/mlcroissant/mlcroissant/_src/structure_graph/base_node.py

@@ -330,6 +330,27 @@ def __deepcopy__(self, memo):
         memo[id(self)] = copy
         return copy
 
+    def __reduce__(self):
+        """Allows pickling the node.
+
+        `self.ctx` is stored separately in the state because it's not pickable directly.
+        """
+        state = self.__getstate__()
+        args = tuple(state.values())
+        return (self.__class__, args, {"ctx": self.ctx})
+
+    def __getstate__(self):
+        state = {}
+        for field in dataclasses.fields(self):
+            if field.name == "ctx":
+                state[field.name] = Context()
+            else:
+                state[field.name] = getattr(self, field.name)
+        return state
+
+    def __setstate__(self, state):
+        self.ctx = state["ctx"]
+
     def to_json(self) -> Json:
         """Converts the Python class to JSON."""
         cls = self.__class__

python/mlcroissant/mlcroissant/_src/structure_graph/nodes/file_object_test.py

@@ -1,5 +1,6 @@
 """Tests for FileObjects."""
 
+import pickle
 from unittest import mock
 
 from etils import epath
@@ -91,3 +92,15 @@ def test_from_jsonld(encoding):
         == "48a7c257f3c90b2a3e529ddd2cca8f4f1bd8e49ed244ef53927649504ac55354"
     )
     assert not ctx.issues.errors
+
+
+@pytest.mark.parametrize(
+    ["conforms_to"],
+    [[CroissantVersion.V_0_8], [CroissantVersion.V_1_0]],
+)
+def test_pickable(conforms_to):
+    ctx = Context(conforms_to=conforms_to)
+    file_object = create_test_node(FileObject, ctx=ctx)
+    file_object = pickle.loads(pickle.dumps(file_object))
+    # Test that the context was successfully restored:
+    assert file_object.ctx.conforms_to == conforms_to