Dry run performance optimization for huge process graphs (#427)

* add test that shows large graph is slow

* replace double-for loop with list comprehension: is faster when processing 4000000 traces

* logging impacts performance: reduce it

* #426 changelog

* skip long test

* Fix pytest skip #427

* Further optimization of DryRunDataTracer.get_trace_leaves #426/#427

---------

Co-authored-by: Stefaan Lippens <[email protected]>

Author: jdries

CHANGELOG.md

@@ -22,7 +22,7 @@ and start a new "In Progress" section above it.
 ## In progress: 0.136.0
 
 - Start supporting custom `UdfRuntimes` implementation in `OpenEoBackendImplementation` ([#415](https://github.com/Open-EO/openeo-python-driver/issues/415))
-
+- Process graph parsing (dry-run) for very large graphs got faster. ([#426](https://github.com/Open-EO/openeo-python-driver/issues/426))
 
 ## 0.135.0
 

openeo_driver/ProcessGraphDeserializer.py

@@ -1914,14 +1914,13 @@ def check_subgraph_for_data_mask_optimization(args: dict) -> bool:
 
 
 def apply_process(process_id: str, args: dict, namespace: Union[str, None], env: EvalEnv) -> DriverDataCube:
-    _log.debug(f"apply_process {process_id} with {args}")
+    _log.debug(f"apply_process {process_id} ")
     parameters = env.collect_parameters()
 
     if process_id == "mask" and args.get("replacement", None) is None \
             and smart_bool(env.get("data_mask_optimization", True)):
         mask_node = args.get("mask", None)
         # evaluate the mask
-        _log.debug(f"data_mask: convert_node(mask_node): {mask_node}")
         the_mask = convert_node(mask_node, env=env)
         dry_run_tracer: DryRunDataTracer = env.get(ENV_DRY_RUN_TRACER)
         if not dry_run_tracer and check_subgraph_for_data_mask_optimization(args):
@@ -1958,7 +1957,6 @@ def apply_process(process_id: str, args: dict, namespace: Union[str, None], env:
 
     try:
         process_function = process_registry.get_function(process_id, namespace=(namespace or "backend"))
-        _log.debug(f"Applying process {process_id} to arguments {args}")
         #TODO: for API compliance, we would actually first need to check if a UDP with same name exists.
         # we would however prefer to avoid overriding predefined functions with UDP's.
         # if we want to do this, we require caching in UDP registry to avoid expensive UDP lookups. We only need to cache the list of UDP names for a given user.

openeo_driver/dry_run.py

@@ -37,7 +37,7 @@
 
 import logging
 from enum import Enum
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, Iterator
 
 import numpy
 import shapely.geometry.base
@@ -100,8 +100,10 @@
 class DataTraceBase:
     """Base class for data traces."""
 
+    __slots__ = ["children"]
+
     def __init__(self):
-        self.children = []
+        self.children: List[DataTraceBase] = []
 
     def __hash__(self):
         # Identity hash (e.g. memory address)
@@ -369,19 +371,19 @@ def get_trace_leaves(self) -> List[DataTraceBase]:
         Get all nodes in the tree of traces that are not parent of another trace.
         In openEO this could be for instance a save_result process that ends the workflow.
         """
-        leaves = []
-
-        def get_leaves(tree: DataTraceBase) -> List[DataTraceBase]:
-            return (
-                [tree] if len(tree.children) == 0 else [leaf for child in tree.children for leaf in get_leaves(child)]
-            )
-
-        for trace in self._traces:
-            for leaf in get_leaves(trace):
-                if leaf not in leaves:
-                    leaves.append(leaf)
+        visited = set()
+
+        def get_leaves(trace: DataTraceBase) -> Iterator[DataTraceBase]:
+            nonlocal visited
+            if trace not in visited:
+                visited.add(trace)
+                if trace.children:
+                    for child in trace.children:
+                        yield from get_leaves(child)
+                else:
+                    yield trace
 
-        return leaves
+        return [leave for trace in self._traces for leave in get_leaves(trace)]
 
     def get_metadata_links(self):
         result = {}