Experiment: Add breakpoint mvp

docs/source/logging_and_errors.md

@@ -117,6 +117,40 @@ The data contains the following information:
 - `env` the current user environment when launching the task.
 - `packages` a list of installed python packages. Only used for the [`UvExecutor`](#tierkreis.controller.executor.uv_executor.UvExecutor)
 
+### Breakpoints
+
+All nodes can be declared ar breakpoints by adding `NodeMetaData` as follows:
+```python
+from tierkreis.controller.data.graph import NodeMetaData
+g.task(..., NodeMetaData(has_breakpoint=True))
+```
+By default the controller will ignore this information unless you set `enable_breakpoints=True` in `run_graph`. 
+When running with breakpoints the graph execution will stop as soon as it hits a breakpoint node e.g. after running:
+```python
+run_graph(
+    storage, executor, graph, None, enable_breakpoints=enable_breakpoints
+)
+```
+you can examine the storage and current values of all previous nodes.
+Afterward you can resume the execution with
+```python
+resume_graph(storage, executor)
+```
+
+### Debug Mode
+
+If you want to debug a graph with a python debugger you can use [](#tierkreis.controller.storgare.debug_graph.debug_graph).
+It acts similar to `run_workflow` with some defaults enabled:
+- Enables all set breakpoints
+- Sets up logging
+- Adds a specific storage and executor to enable python debugging
+```{important}
+This will only work with python based workers.
+All workers need to be installed locally.
+```
+
+
+
 ## Visualizer
 
 If you're using the visualize to debug workflow, error information will be immediately visible to you.

tierkreis/tests/controller/test_breakpoint.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+from uuid import UUID
+
+import pytest
+from pytket._tket.circuit import Circuit
+from tierkreis.builder import GraphBuilder
+from tierkreis.builtins import iadd
+from tierkreis.controller import resume_graph, run_graph
+from tierkreis.controller.data.graph import NodeMetaData
+from tierkreis.controller.data.location import Loc
+from tierkreis.controller.data.models import TKR
+from tierkreis.controller.executor.in_memory_executor import InMemoryExecutor
+from tierkreis.controller.storage.filestorage import ControllerFileStorage
+from tierkreis.controller.storage.in_memory import ControllerInMemoryStorage
+from tierkreis.executor import ShellExecutor
+from pytket_worker import n_qubits
+from tierkreis.storage import read_outputs
+
+
+def breakpoint_graph() -> GraphBuilder[TKR[Circuit], TKR[int]]:
+    g = GraphBuilder(TKR[Circuit], TKR[int])
+    test = g.const(5)
+    nq = g.task(n_qubits(g.inputs), NodeMetaData(has_breakpoint=True))  # type: ignore
+    out = g.task(iadd(test, nq))
+    g.outputs(out)
+    return g
+
+
+storage_classes = [ControllerFileStorage, ControllerInMemoryStorage]
+storage_ids = ["FileStorage", "In-memory"]
+
+
+@pytest.mark.parametrize("storage_class", storage_classes, ids=storage_ids)
+@pytest.mark.parametrize("enable_breakpoints", [True, False], ids=["True", "False"])
+def test_breakpoint(
+    storage_class: type[ControllerFileStorage | ControllerInMemoryStorage],
+    enable_breakpoints: bool,
+) -> None:
+    graph = breakpoint_graph()
+    storage = storage_class(UUID(int=400), name="breakpoints")
+    executor = ShellExecutor(registry_path=None, workflow_dir=storage.workflow_dir)
+    if isinstance(storage, ControllerInMemoryStorage):
+        executor = InMemoryExecutor(Path("./tierkreis/tierkreis"), storage=storage)
+    storage.clean_graph_files()
+    run_graph(
+        storage, executor, graph, Circuit(2), enable_breakpoints=enable_breakpoints
+    )
+    if enable_breakpoints:
+        assert not storage.is_node_finished(Loc())
+        assert storage.exists(storage._breakpoint(Loc("-.N2")))
+        resume_graph(storage, executor)
+    assert storage.is_node_finished(Loc())
+    out = read_outputs(graph, storage)
+    assert out == 7

tierkreis/tierkreis/builder.py

@@ -5,6 +5,7 @@
 from collections.abc import Callable
 from copy import copy
 from dataclasses import dataclass
+from functools import partial
 from inspect import isclass
 from typing import (
     Any,
@@ -15,8 +16,8 @@
     runtime_checkable,
 )
 
-from tierkreis.controller.data.core import EmptyModel
-from tierkreis.controller.data.graph import GraphData, ValueRef, reindex_inputs
+from tierkreis.controller.data.core import EmptyModel, ValueRef
+from tierkreis.controller.data.graph import GraphData, NodeMetaData, reindex_inputs
 from tierkreis.controller.data.models import (
     TKR,
     TModel,
@@ -126,16 +127,23 @@ class GraphBuilder[Inputs: TModel, Outputs: TModel]:
 
     outputs_type: type
     inputs: Inputs
+    _breakpoints_on_outputs: bool
 
     def __init__(
         self,
         inputs_type: type[Inputs] = EmptyModel,
         outputs_type: type[Outputs] = EmptyModel,
+        breakpoints_on_inputs: bool = False,
+        breakpoints_on_outputs: bool = False,
     ) -> None:
         self.data = GraphData()
         self.inputs_type = inputs_type
         self.outputs_type = outputs_type
-        self.inputs = init_tmodel(self.inputs_type, self.data.input)
+        input_fn = partial(
+            self.data.input, metadata=NodeMetaData(breakpoints_on_inputs)
+        )
+        self.inputs = init_tmodel(self.inputs_type, input_fn)
+        self._breakpoints_on_outputs = breakpoints_on_outputs
 
     def get_data(self) -> GraphData:
         """Return the underlying graph from the builder.
@@ -159,7 +167,10 @@ def outputs(self, outputs: Outputs) -> None:
         :param outputs: The output nodes.
         :type outputs: Outputs
         """
-        self.data.output(inputs=dict_from_tmodel(outputs))
+        self.data.output(
+            inputs=dict_from_tmodel(outputs),
+            metadata=NodeMetaData(self._breakpoints_on_outputs),
+        )
 
     def embed[A: TModel, B: TModel](self, other: "GraphBuilder[A, B]", inputs: A) -> B:
         if other.data.graph_output_idx is None:
@@ -196,14 +207,15 @@ def const[T: PType](self, value: T) -> TKR[T]:
         :return: The constant value.
         :rtype: TKR[T]
         """
-        idx, port = self.data.const(value)
+        idx, port = self.data.const(value, NodeMetaData())
         return TKR[T](idx, port)
 
     def ifelse[A: PType, B: PType](
         self,
         pred: TKR[bool],
         if_true: TKR[A],
         if_false: TKR[B],
+        metadata: NodeMetaData | None = None,
     ) -> TKR[A] | TKR[B]:
         """Add an if-else node to the graph.
 
@@ -216,13 +228,16 @@ def ifelse[A: PType, B: PType](
         :type if_true: TKR[A]
         :param if_false: The value if the predicate is false.
         :type if_false: TKR[B]
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the if-else expression.
         :rtype: TKR[A] | TKR[B]
         """
         idx, port = self.data.if_else(
             pred.value_ref(),
             if_true.value_ref(),
             if_false.value_ref(),
+            metadata,
         )("value")
         return TKR(idx, port)
 
@@ -231,6 +246,7 @@ def eifelse[A: PType, B: PType](
         pred: TKR[bool],
         if_true: TKR[A],
         if_false: TKR[B],
+        metadata: NodeMetaData | None = None,
     ) -> TKR[A] | TKR[B]:
         """Add an eager if-else node to the graph.
 
@@ -243,13 +259,16 @@ def eifelse[A: PType, B: PType](
         :type if_true: TKR[A]
         :param if_false: The value if the predicate is false.
         :type if_false: TKR[B]
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the if-else expression.
         :rtype: TKR[A] | TKR[B]
         """
         idx, port = self.data.eager_if_else(
             pred.value_ref(),
             if_true.value_ref(),
             if_false.value_ref(),
+            metadata,
         )("value")
         return TKR(idx, port)
 
@@ -258,31 +277,38 @@ def _graph_const[A: TModel, B: TModel](
         graph: GraphBuilder[A, B],
     ) -> TypedGraphRef[A, B]:
         # TODO @philipp-seitz: Turn this into a public method?
-        idx, port = self.data.const(graph.data.model_dump())
+        idx, port = self.data.const(graph.data.model_dump(), NodeMetaData())
         return TypedGraphRef[A, B](
             graph_ref=(idx, port),
             outputs_type=graph.outputs_type,
             inputs_type=graph.inputs_type,
         )
 
-    def task[Out: TModel](self, func: Function[Out]) -> Out:
+    def task[Out: TModel](
+        self,
+        func: Function[Out],
+        metadata: NodeMetaData | None = None,
+    ) -> Out:
         """Add a worker task node to the graph.
 
         :param func: The worker function.
         :type func: Function[Out]
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the task.
         :rtype: Out
         """
         name = f"{func.namespace}.{func.__class__.__name__}"
         inputs = dict_from_tmodel(func)
-        idx, _ = self.data.func(name, inputs)("dummy")
+        idx, _ = self.data.func(name, inputs, metadata)("dummy")
         OutModel = func.out()  # noqa: N806
         return init_tmodel(OutModel, lambda p: (idx, p))
 
     def eval[A: TModel, B: TModel](
         self,
         body: GraphBuilder[A, B] | TypedGraphRef[A, B],
         eval_inputs: A,
+        metadata: NodeMetaData | None = None,
     ) -> B:
         """Add a evaluation node to the graph.
 
@@ -293,20 +319,25 @@ def eval[A: TModel, B: TModel](
             where A are the input type and B the output type of the graph.
         :param eval_inputs: The inputs to the graph.
         :type eval_inputs: A
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the evaluation.
         :rtype: B
         """
         if isinstance(body, GraphBuilder):
             body = self._graph_const(body)
 
-        idx, _ = self.data.eval(body.graph_ref, dict_from_tmodel(eval_inputs))("dummy")
+        idx, _ = self.data.eval(
+            body.graph_ref, dict_from_tmodel(eval_inputs), metadata
+        )("dummy")
         return init_tmodel(body.outputs_type, lambda p: (idx, p))
 
     def loop[A: TModel, B: LoopOutput](
         self,
         body: TypedGraphRef[A, B] | GraphBuilder[A, B],
         loop_inputs: A,
         name: str | None = None,
+        metadata: NodeMetaData | None = None,
     ) -> B:
         """Add a loop node to the graph.
 
@@ -321,6 +352,8 @@ def loop[A: TModel, B: LoopOutput](
         :type loop_inputs: A
         :param name: An optional name for the loop.
         :type name: str | None
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the loop.
         :rtype: B
         """
@@ -333,6 +366,7 @@ def loop[A: TModel, B: LoopOutput](
             dict_from_tmodel(loop_inputs),
             "should_continue",
             name,
+            metadata=metadata,
         )(
             "dummy",
         )
@@ -369,9 +403,10 @@ def _map_graph_full[A: TModel, B: TModel](
         self,
         map_inputs: TList[A],
         body: TypedGraphRef[A, B],
+        metadata: NodeMetaData | None = None,
     ) -> TList[B]:
         ins = dict_from_tmodel(map_inputs._value)  # noqa: SLF001
-        idx, _ = self.data.map(body.graph_ref, ins)("x")
+        idx, _ = self.data.map(body.graph_ref, ins, metadata)("x")
 
         return TList(init_tmodel(body.outputs_type, lambda s: (idx, s + "-*")))
 
@@ -382,6 +417,7 @@ def map[A: PType, B: TNamedModel](
             Callable[[TKR[A]], B] | TypedGraphRef[TKR[A], B] | GraphBuilder[TKR[A], B]
         ),
         map_inputs: TKR[list[A]],
+        metadata: NodeMetaData | None = None,
     ) -> TList[B]: ...
 
     @overload
@@ -391,13 +427,15 @@ def map[A: TNamedModel, B: PType](
             Callable[[A], TKR[B]] | TypedGraphRef[A, TKR[B]] | GraphBuilder[A, TKR[B]]
         ),
         map_inputs: TList[A],
+        metadata: NodeMetaData | None = None,
     ) -> TKR[list[B]]: ...
 
     @overload
     def map[A: TNamedModel, B: TNamedModel](
         self,
         body: TypedGraphRef[A, B] | GraphBuilder[A, B],
         map_inputs: TList[A],
+        metadata: NodeMetaData | None = None,
     ) -> TList[B]: ...
 
     @overload
@@ -409,19 +447,23 @@ def map[A: PType, B: PType](
             | GraphBuilder[TKR[A], TKR[B]]
         ),
         map_inputs: TKR[list[A]],
+        metadata: NodeMetaData | None = None,
     ) -> TKR[list[B]]: ...
 
     def map(
         self,
         body: TypedGraphRef | Callable | GraphBuilder,
         map_inputs: TKR | TList,
+        metadata: NodeMetaData | None = None,
     ) -> Any:
         """Add a map node to the graph.
 
         :param body: The graph to map over.
         :type body: TypedGraphRef | Callable | GraphBuilder
         :param map_inputs: The values to map over.
         :type map_inputs: TKR | TList
+        :param metadata: Optional metadata for the node, defaults to None
+        :type metadata: NodeMetaData | None, optional
         :return: The outputs of the map.
         :rtype: Any
         """
@@ -437,7 +479,7 @@ def map(
         if isinstance(map_inputs, TKR):
             map_inputs = self._unfold_list(map_inputs)
 
-        out = self._map_graph_full(map_inputs, body)
+        out = self._map_graph_full(map_inputs, body, metadata)
 
         if not isclass(body.outputs_type) or not issubclass(
             body.outputs_type,