fix: deserialization, base criterion in graph

Author: glichtner

execution_engine/execution_graph/graph.py

@@ -90,8 +90,12 @@ def traverse(
                 )
                 subgraph.set_sink_nodes_store(bind_params=dict(pi_pair_id=expr.id))
 
+            elif expr == base_node:
+                # don't need to do anything - only non-base criteria are connected to the base criterion,
+                # otherwise we get a cyclic graph
+                pass
             elif expr.is_Atom:
-                assert expr in graph.nodes
+                assert expr in graph.nodes, "Node not found in graph"
                 graph.nodes[expr]["store_result"] = True
                 graph.add_edge(base_node, expr)
             else:

execution_engine/omop/cohort/graph_builder.py

@@ -36,6 +36,8 @@ def filter_symbols(cls, node: logic.Expr, filter_: logic.Expr) -> logic.Expr:
         :rtype: logic.Expr
         """
 
+        node = copy.copy(node)
+
         if isinstance(node, logic.Symbol):
             return logic.LeftDependentToggle(left=filter_, right=node)
         elif isinstance(node, logic.Expr):
@@ -140,7 +142,10 @@ def build(cls, expr: logic.Expr, base_criterion: Criterion) -> ExecutionGraph:
         )
         graph.add_edges_from((src, p_combination_node) for src in p_sink_nodes)
 
+        if graph.in_degree(base_criterion) != 0:
+            raise AssertionError("Base criterion must not have incoming edges")
+
         if not nx.is_directed_acyclic_graph(graph):
-            raise ValueError("Graph is not acyclic")
+            raise AssertionError("Graph is not acyclic")
 
         return graph

execution_engine/task/creator.py

@@ -1,9 +1,96 @@
+import pickle  # nosec
+
 import networkx as nx
+from typing_extensions import Any
 
 import execution_engine.util.logic as logic
 from execution_engine.task.task import Task
 
 
+def assert_pickle_roundtrip(obj: logic.BaseExpr) -> None:
+    """
+    Serializes 'obj' via pickle (the same method multiprocessing would use),
+    then deserializes it, and finally compares the original object to the result.
+
+    :param obj: The object to serialize/deserialize.
+    :raises AssertionError: If the object does not match its clone.
+    :return: The deserialized clone (for further inspection if needed).
+    """
+    pickled = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)  # nosec
+    clone = pickle.loads(pickled)  # nosec
+
+    if isinstance(obj, logic.CountOperator):
+        if obj.dict()["data"]["threshold"] is None:
+            raise AssertionError("Threshold must be set")
+
+    if obj == clone:
+        return  # They are considered equal, so nothing to do.
+
+    # If they're unequal, compare their dict() representations to find differences.
+    d1 = obj.dict()
+    d2 = clone.dict()
+
+    if d1 == d2:
+        # If they're unequal but dicts are the same, there's some internal difference
+        # not visible via .dict(). Just alert that we can't show details.
+        raise AssertionError(
+            f"Objects differ in __eq__, but their dict() representations are identical.\n"
+            f"obj:   {obj}\n"
+            f"clone: {clone}"
+        )
+
+    # Otherwise, gather all leaf-level differences in d1 vs d2.
+    diffs = _compare_dicts_leaf_level(d1, d2)
+    diff_msg = "\n".join(diffs)
+
+    raise AssertionError(
+        f"Object does not match its clone after round-trip!\n\n"
+        f"Differences at leaf level in .dict() representations:\n{diff_msg}"
+    )
+
+
+def _compare_dicts_leaf_level(d1: Any, d2: Any, path: str = "") -> list[str]:
+    """
+    Recursively compare two dict/list/tuple/scalar structures and return
+    a list of strings describing differences at the leaf level.
+
+    :param d1, d2: Potentially nested structures (dict, list, tuple, scalar).
+    :param path: Path string to locate the current point in the structure.
+    :return: List of difference descriptions.
+    """
+    differences = []
+
+    # If both are dicts, recurse into matching keys
+    if isinstance(d1, dict) and isinstance(d2, dict):
+        all_keys = set(d1.keys()) | set(d2.keys())
+        for key in sorted(all_keys):
+            sub_path = f"{path}.{key}" if path else str(key)
+            if key not in d1:
+                differences.append(f"[MISSING IN ORIGINAL] {sub_path} => {d2[key]!r}")
+            elif key not in d2:
+                differences.append(f"[MISSING IN CLONE] {sub_path} => {d1[key]!r}")
+            else:
+                differences.extend(
+                    _compare_dicts_leaf_level(d1[key], d2[key], sub_path)
+                )
+
+    # If both are lists/tuples, compare element by element
+    elif isinstance(d1, (list, tuple)) and isinstance(d2, (list, tuple)):
+        if len(d1) != len(d2):
+            differences.append(f"[LEN MISMATCH] {path} => {len(d1)} vs {len(d2)}")
+        else:
+            for i, (item1, item2) in enumerate(zip(d1, d2)):
+                sub_path = f"{path}[{i}]"
+                differences.extend(_compare_dicts_leaf_level(item1, item2, sub_path))
+
+    # Otherwise, treat them as leaf values and compare directly
+    else:
+        if d1 != d2:
+            differences.append(f"[VALUE MISMATCH] {path} => {d1!r} vs {d2!r}")
+
+    return differences
+
+
 class TaskCreator:
     """
     A TaskCreator object creates a Task tree for an expression and its dependencies.
@@ -43,12 +130,11 @@ def node_to_task(expr: logic.Expr, attr: dict) -> Task:
 
         flattened_tasks = list(tasks.values())
 
-        # we will make sure all tasks are depickled correctly
-        for i, node in enumerate(tasks):
-            if logic.Expr.from_dict(node.dict(include_id=True)) != node:
-                raise RuntimeError(
-                    "Expected depickled node to be the same as initial node."
-                )
+        # we will make sure all tasks are depickled correctly [commented out for performance reasons]
+        # from tqdm import tqdm
+        #
+        # for node in tqdm(tasks):
+        #     assert_pickle_roundtrip(node)
 
         assert (
             len(set(flattened_tasks))

execution_engine/util/logic.py

@@ -58,25 +58,31 @@ def _recreate(cls, args: Any, kwargs: dict) -> "Expr":
         self.set_id(_id)
         return self
 
-    def _reduce_helper(self, ivars_map: dict | None = None) -> tuple[Callable, tuple]:
+    def _reduce_helper(
+        self, ivars_map: dict | None = None, args: tuple | None = None
+    ) -> tuple[Callable, tuple]:
         """
         Return a picklable tuple that calls self._recreate and passes in
         (self.args, combined_ivars).
 
         :param ivars_map: A dictionary for renaming keys in the instance variables.
                           For each old_key -> new_key in ivars_map, if old_key exists
                           in data, it will be removed and stored under new_key instead.
+        :param args: Optionally, the args parameter can be specified. If it is None, self.args is used.
         """
-        data = dict(self.get_instance_variables())
-        data["_id"] = self._id
+        data = self.get_instance_variables()
+        data["_id"] = self._id  # type: ignore[index]
 
         # Apply key renaming if ivars_map is provided
         if ivars_map:
             for old_key, new_key in ivars_map.items():
                 if old_key in data:
-                    data[new_key] = data.pop(old_key)
+                    data[new_key] = data.pop(old_key)  # type: ignore[index,union-attr]
 
-        return (self._recreate, (self.args, data))
+        if args is None:
+            args = self.args
+
+        return (self._recreate, (args, data))
 
     def __reduce__(self) -> tuple[Callable, tuple]:
         """
@@ -404,15 +410,30 @@ def __new__(
 
         return self
 
+    def _replace_map(self) -> dict:
+        replace = {}
+
+        if self.count_min and self.count_max:
+            assert self.count_min == self.count_max
+            replace["count_min"] = "threshold"
+        elif self.count_min:
+            replace["count_min"] = "threshold"
+        elif self.count_max:
+            replace["count_max"] = "threshold"
+        else:
+            raise AttributeError("At least one of count_min or count_max must be set")
+
+        return replace
+
     def __reduce__(self) -> tuple[Callable, tuple]:
         """
         Reduce the expression to its arguments and category.
 
         Required for pickling (e.g. when using multiprocessing).
 
-        :return: Tuple of the class, arguments, and category.
+        :return: Tuple of the class, argument.
         """
-        return self._reduce_helper({"count_min": "threshold", "count_max": "threshold"})
+        return self._reduce_helper(self._replace_map())
 
 
 class Count(CountOperator):
@@ -683,6 +704,31 @@ def _validate_time_inputs(
                 f"Invalid criterion - expected Criterion or CriterionCombination, got {type(interval_criterion)}"
             )
 
+    def __reduce__(self) -> tuple[Callable, tuple]:
+        """
+        Reduce the expression to its arguments and category.
+
+        Required for pickling (e.g. when using multiprocessing).
+
+        :return: Tuple of the class, argument.
+        """
+        if self.interval_criterion:
+
+            if len(self.args) <= 1:
+                raise ValueError(
+                    "More than one argument required if interval_criterion is set"
+                )
+
+            args, pop = self.args[:-1], self.args[-1]
+
+            if pop != self.interval_criterion:
+                raise ValueError(
+                    f"Expected last argument to be the interval_criterion, got {str(pop)}"
+                )
+            return self._reduce_helper(self._replace_map(), args=args)
+
+        return super().__reduce__()
+
     def dict(self, include_id: bool = False) -> dict:
         """
         Get a dictionary representation of the object.
@@ -746,11 +792,11 @@ class TemporalMinCount(TemporalCount):
     def __new__(
         cls,
         *args: Any,
-        threshold: int | None,
-        start_time: time | None,
-        end_time: time | None,
-        interval_type: TimeIntervalType | None,
-        interval_criterion: BaseExpr | None,
+        threshold: int,
+        start_time: time | None = None,
+        end_time: time | None = None,
+        interval_type: TimeIntervalType | None = None,
+        interval_criterion: BaseExpr | None = None,
         **kwargs: Any,
     ) -> "TemporalMinCount":
         """
@@ -776,7 +822,7 @@ def dict(self, include_id: bool = False) -> dict:
         Get a dictionary representation of the object.
         """
         data = super().dict(include_id=include_id)
-        data.update({"threshold": self.count_min})
+        data["data"].update({"threshold": self.count_min})
         return data
 
 
@@ -788,11 +834,11 @@ class TemporalMaxCount(TemporalCount):
     def __new__(
         cls,
         *args: Any,
-        threshold: int | None,
-        start_time: time | None,
-        end_time: time | None,
-        interval_type: TimeIntervalType | None,
-        interval_criterion: BaseExpr | None,
+        threshold: int,
+        start_time: time | None = None,
+        end_time: time | None = None,
+        interval_type: TimeIntervalType | None = None,
+        interval_criterion: BaseExpr | None = None,
         **kwargs: Any,
     ) -> "TemporalMaxCount":
         """
@@ -818,7 +864,7 @@ def dict(self, include_id: bool = False) -> dict:
         Get a dictionary representation of the object.
         """
         data = super().dict(include_id=include_id)
-        data.update({"threshold": self.count_max})
+        data["data"].update({"threshold": self.count_max})
         return data
 
 
@@ -830,11 +876,11 @@ class TemporalExactCount(TemporalCount):
     def __new__(
         cls,
         *args: Any,
-        threshold: int | None,
-        start_time: time | None,
-        end_time: time | None,
-        interval_type: TimeIntervalType | None,
-        interval_criterion: BaseExpr | None,
+        threshold: int,
+        start_time: time | None = None,
+        end_time: time | None = None,
+        interval_type: TimeIntervalType | None = None,
+        interval_criterion: BaseExpr | None = None,
         **kwargs: Any,
     ) -> "TemporalExactCount":
         """
@@ -860,7 +906,7 @@ def dict(self, include_id: bool = False) -> dict:
         Get a dictionary representation of the object.
         """
         data = super().dict(include_id=include_id)
-        data.update({"threshold": self.count_min})
+        data["data"].update({"threshold": self.count_min})
         return data
 
 

execution_engine/util/serializable.py

@@ -105,6 +105,24 @@ def __new__(mcs, name: str, bases: tuple[type, ...], attrs: dict[str, Any]) -> S
         return new_class
 
 
+def immutable_setattr(self: Self, key: str, value: Any) -> None:
+    """
+    Prevent setting attributes on an immutable object.
+
+    This function is assigned to an instance's __setattr__ method in order to enforce
+    immutability after object creation. Any attempt to set an attribute on the instance
+    after initialization will raise an AttributeError.
+
+    :param self: The instance on which the attribute assignment was attempted.
+    :param key: The name of the attribute being set.
+    :param value: The value being assigned.
+    :raises AttributeError: Always, to enforce immutability.
+    """
+    raise AttributeError(
+        f"Cannot set attribute {key} on immutable object {self.__class__.__name__}"
+    )
+
+
 class Serializable(metaclass=RegisteredPostInitMeta):
     """
     Base class for making objects serializable.
@@ -137,11 +155,6 @@ def __post_init__(self) -> None:
         """
         self.rehash()
 
-        def immutable_setattr(self: Self, key: str, value: Any) -> None:
-            raise AttributeError(
-                f"Cannot set attribute {key} on immutable object {self.__class__.__name__}"
-            )
-
         self.__setattr__ = immutable_setattr  # type: ignore[assignment]
 
     def set_id(self, value: int, overwrite: bool = False) -> None: