[SYSTEMDS-3940] Improve Scuro Node Executor 

In this patch the node executor for Representation DAGs is improved and made mode efficient. Additionally, the code was cleaned up.

Assisted-by: AI
diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
index 9407bdf..a62b990 100644
--- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
+++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py
@@ -297,6 +297,7 @@
         wandb_entity: Optional[str] = None,
         wandb_group: Optional[str] = None,
         wandb_tags: Optional[List[str]] = None,
+        enable_checkpointing: bool = False,
     ):
         self.tasks = tasks
         self.unimodal_optimization_results = optimization_results
@@ -333,6 +334,7 @@
         self.wandb_group = wandb_group
         self.wandb_tags = wandb_tags or []
         self._wandb_run = None
+        self.enable_checkpointing = enable_checkpointing
 
     def get_modalities_by_id(self, modality_ids: List[int]) -> Modality:
         modalities = []
@@ -415,14 +417,18 @@
                             )
                         )
                     self.optimization_results.add_result(results)
-                    self._checkpoint_manager.increment(task.model.name, len(results))
-                    self._checkpoint_manager.checkpoint_if_due(
-                        self.optimization_results.results,
-                    )
+                    if self.enable_checkpointing:
+                        self._checkpoint_manager.increment(
+                            task.model.name, len(results)
+                        )
+                        self._checkpoint_manager.checkpoint_if_due(
+                            self.optimization_results.results,
+                        )
                 except Exception:
-                    self._checkpoint_manager.save_checkpoint(
-                        self.optimization_results.results, {}
-                    )
+                    if self.enable_checkpointing:
+                        self._checkpoint_manager.save_checkpoint(
+                            self.optimization_results.results, {}
+                        )
                     raise
 
         if self.save_results:
diff --git a/src/main/python/systemds/scuro/drsearch/modality_result_cache.py b/src/main/python/systemds/scuro/drsearch/modality_result_cache.py
new file mode 100644
index 0000000..0848702
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/modality_result_cache.py
@@ -0,0 +1,136 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+from typing import Any, Dict, List, Optional
+
+from systemds.scuro.drsearch.modality_shared_memory import unlink_shm
+from systemds.scuro.utils.static_variables import DEBUG
+
+
+class RefCountResultCache:
+    def __init__(self):
+        self.cache: Dict[str, Any] = {}
+        self.ref_count: Dict[str, int] = {}
+        self.memory_usage_per_node: Dict[str, int] = {}
+        self.shared_memory_names: Dict[str, List[str]] = {}
+        self._shm_retain_count: Dict[str, int] = {}
+
+    def get(self, node_id: str) -> Any:
+        return self.cache[node_id]
+
+    def add_result(
+        self,
+        node_id: str,
+        result: Any,
+        shm_name: Optional[str] = None,
+        resident_bytes: Optional[int] = None,
+        shm_bytes: int = 0,
+    ):
+        if shm_name is not None:
+            self.shared_memory_names[node_id] = [shm_name]
+        self.cache[node_id] = result
+        self.memory_usage_per_node[node_id] = int(resident_bytes or 0) + int(
+            shm_bytes or 0
+        )
+        if DEBUG:
+            print(
+                f"Node {node_id} has a CPU memory usage of "
+                f"{self.memory_usage_per_node[node_id]/1024**3:.5f} GB"
+                + (
+                    f" ({int(shm_bytes or 0)/1024**3:.5f} GB of it shared memory)"
+                    if shm_name is not None
+                    else ""
+                )
+            )
+
+    def inc_ref(self, node_id: str):
+        self.ref_count[node_id] = self.ref_count.get(node_id, 0) + 1
+
+    def dec_ref(self, node_id: str):
+        if node_id not in self.ref_count:
+            return
+        self.ref_count[node_id] -= 1
+        if self.ref_count[node_id] <= 0:
+            self.ref_count[node_id] = 0
+            self._try_cleanup_node(node_id)
+
+    def clear(self, node_id: str):
+        self.ref_count[node_id] = 0
+        self._try_cleanup_node(node_id)
+
+    def retain_shm_names(self, shm_names: List[str]) -> List[str]:
+        retained: List[str] = []
+        for shm_name in shm_names:
+            if not shm_name:
+                continue
+            self._shm_retain_count[shm_name] = (
+                self._shm_retain_count.get(shm_name, 0) + 1
+            )
+            retained.append(shm_name)
+        return retained
+
+    def release_shm_names(self, shm_names: List[str]) -> None:
+        nodes_to_recheck: List[str] = []
+        for shm_name in shm_names:
+            if not shm_name:
+                continue
+            count = self._shm_retain_count.get(shm_name, 0) - 1
+            if count <= 0:
+                self._shm_retain_count.pop(shm_name, None)
+            else:
+                self._shm_retain_count[shm_name] = count
+            for node_id, node_names in self.shared_memory_names.items():
+                if shm_name in node_names and node_id not in nodes_to_recheck:
+                    nodes_to_recheck.append(node_id)
+        for node_id in nodes_to_recheck:
+            self._try_cleanup_node(node_id)
+
+    def __len__(self):
+        return len(self.cache)
+
+    def get_memory_total_memory_usage(self):
+        return sum(self.memory_usage_per_node.values())
+
+    def _shm_names_in_use(self, shm_names: List[str]) -> bool:
+        return any(self._shm_retain_count.get(name, 0) > 0 for name in shm_names)
+
+    def _try_cleanup_node(self, node_id: str) -> None:
+        if self.ref_count.get(node_id, 0) > 0:
+            return
+        shm_names = self.shared_memory_names.get(node_id, [])
+        if shm_names and self._shm_names_in_use(shm_names):
+            return
+        self.cache.pop(node_id, None)
+        self.ref_count.pop(node_id, None)
+        self.memory_usage_per_node.pop(node_id, None)
+        self._cleanup_shared_memory(node_id)
+
+    def _cleanup_shared_memory(self, node_id: str):
+        names = self.shared_memory_names.pop(node_id, [])
+        for shm_name in names:
+            unlink_shm(shm_name)
+
+    def cleanup_all(self):
+        self._shm_retain_count.clear()
+        for node_id in list(self.shared_memory_names.keys()):
+            self.ref_count.pop(node_id, None)
+            self.cache.pop(node_id, None)
+            self.memory_usage_per_node.pop(node_id, None)
+            self._cleanup_shared_memory(node_id)
diff --git a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py
index d4092b9..e68d719 100644
--- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py
+++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py
@@ -20,11 +20,29 @@
 # -------------------------------------------------------------
 from typing import Any, List, Tuple
 import numpy as np
-from multiprocessing import shared_memory
+from multiprocessing import shared_memory, resource_tracker
 
 SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024
 
 
+def _untrack(shm: shared_memory.SharedMemory) -> None:
+    try:
+        resource_tracker.unregister(shm._name, "shared_memory")
+    except Exception:
+        pass
+
+
+def unlink_shm(name: str) -> None:
+    try:
+        shm = shared_memory.SharedMemory(name=name)
+        shm.close()
+        shm.unlink()
+    except FileNotFoundError:
+        pass
+    except Exception:
+        pass
+
+
 class SharedStringList:
     def __init__(
         self, shm_name: str, offsets: List[Tuple[int, int]], payload_nbytes: int
@@ -37,6 +55,7 @@
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
 
     def __len__(self):
         return len(self.offsets)
@@ -91,6 +110,7 @@
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._buffer = np.ndarray(
                 (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
             )
@@ -146,6 +166,7 @@
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._arr = np.ndarray(self.shape, dtype=self._dtype, buffer=self._shm.buf)
             self._arr.setflags(write=False)
 
@@ -215,6 +236,7 @@
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._buffer = np.ndarray(
                 (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
             )
@@ -330,6 +352,7 @@
                 resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_shared_ndarray_candidate(data):
         arr = data
@@ -344,6 +367,7 @@
             resident_bytes = min(resident_bytes, 2 * 1024 * 1024)
 
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_nested_shared_memory_candidate(data):
         leaves: List[np.ndarray] = []
@@ -379,6 +403,7 @@
                 resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_string_list_shared_memory_candidate(data):
         encoded = [s.encode("utf-8") for s in data]
@@ -398,6 +423,7 @@
                 resident_bytes, max(2 * 1024 * 1024, len(str_offsets) * 32)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
 
     return None, None, 0, resident_bytes
diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py
index a6a7ffe..ec5d9f4 100644
--- a/src/main/python/systemds/scuro/drsearch/node_executor.py
+++ b/src/main/python/systemds/scuro/drsearch/node_executor.py
@@ -18,229 +18,131 @@
 # under the License.
 #
 # -------------------------------------------------------------
-from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
-from dataclasses import dataclass
+import multiprocessing as mp
 import os
-from multiprocessing import shared_memory
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+
 from systemds.scuro import Modality
+from systemds.scuro.drsearch.modality_result_cache import RefCountResultCache
 from systemds.scuro.drsearch.modality_shared_memory import (
     add_shared_memory_candidate,
     collect_shm_names_from_payload,
+    unlink_shm,
 )
 from systemds.scuro.drsearch.node_scheduler import MemoryAwareNodeScheduler
 from systemds.scuro.drsearch.representation_dag import (
     RepresentationDag,
     RepresentationNode,
 )
-
-import threading
-import numpy as np
-from typing import Any, Dict, List, Optional
-import multiprocessing as mp
-import psutil
-import time
-import torch
 from systemds.scuro.drsearch.task import PerformanceMeasure
+from systemds.scuro.drsearch.worker_pool import PersistentWorkerPool, create_mp_context
+from systemds.scuro.representations.aggregated_representation import (
+    AggregatedRepresentation,
+)
 from systemds.scuro.representations.context import Context
 from systemds.scuro.representations.dimensionality_reduction import (
     DimensionalityReduction,
 )
-from systemds.scuro.representations.aggregated_representation import (
-    AggregatedRepresentation,
+from systemds.scuro.representations.representation import (
+    RepresentationStats,
+    infer_stats_from_data,
 )
-from systemds.scuro.representations.representation import RepresentationStats
 from systemds.scuro.representations.unimodal import UnimodalRepresentation
 from systemds.scuro.utils.checkpointing import CheckpointManager
-import threading
-import time
-import psutil
-import os
+from systemds.scuro.utils.memory_utility import (
+    MemoryMeasurement,
+    cleanup_gpu,
+    cpu_memory_budget_bytes,
+    estimate_modality_bytes,
+    is_cuda_oom,
+    measure_memory_during,
+)
 from systemds.scuro.utils.static_variables import DEBUG
 
+_MAX_NODE_RETRIES = int(os.environ.get("SCURO_MAX_NODE_RETRIES", "3"))
 
-def measure_peak_rss_during(fn, *args, sample_s=0.01, **kwargs):
-    proc = psutil.Process(os.getpid())
-    baseline = proc.memory_info().rss
-    peak = baseline
-    stop = threading.Event()
 
-    def sampler():
-        nonlocal peak
-        while not stop.is_set():
-            rss = proc.memory_info().rss
-            if rss > peak:
-                peak = rss
-            time.sleep(sample_s)
-
-    t = threading.Thread(target=sampler, daemon=True)
-    t.start()
+def _run_gpu_op(fn, gpu_id: Optional[int]):
+    if gpu_id is None or not torch.cuda.is_available():
+        return fn()
     try:
-        out = fn(*args, **kwargs)
+        try:
+            return fn()
+        except Exception as e:
+            if is_cuda_oom(e):
+                cleanup_gpu(gpu_id)
+                _WORKER_OP_CACHE.clear()
+                return fn()
+            raise
     finally:
-        stop.set()
-        t.join()
-
-    return out, (peak - baseline), peak
+        cleanup_gpu(gpu_id)
 
 
-class RefCountResultCache:
-    def __init__(self):
-        self.cache = {}
-        self.ref_count = {}
-        self.memory_usage_per_node = {}
-        self.shared_memory_names = {}
-        self._shm_retain_count: Dict[str, int] = {}
+_WORKER_OP_CACHE: Dict[str, Any] = {}
 
-    def get(self, node_id: str) -> Any:
-        return self.cache[node_id]
 
-    def add_result(self, node_id: str, result: Any):
+def _instantiate_operation(node):
+    cache_key = None
+    if getattr(node.operation, "cache_in_worker", False):
+        try:
+            params_repr = repr(sorted(node.parameters.items(), key=lambda kv: kv[0]))
+            cache_key = (
+                f"{node.operation.__module__}.{node.operation.__qualname__}"
+                f"|{params_repr}"
+            )
+        except Exception:
+            cache_key = None
+    if cache_key is not None and cache_key in _WORKER_OP_CACHE:
+        return _WORKER_OP_CACHE[cache_key]
+    operation = node.operation(params=node.parameters)
+    if cache_key is not None:
+        _WORKER_OP_CACHE[cache_key] = operation
+    return operation
+
+
+def _infer_actual_output_stats(
+    transformed_modality: Any,
+) -> Optional[RepresentationStats]:
+    if transformed_modality is None or not hasattr(transformed_modality, "data"):
+        return None
+    return infer_stats_from_data(transformed_modality.data)
+
+
+def _offload_to_shared_memory(result: Any):
+    if result is None or not hasattr(result, "data"):
+        return None, None, 0, None
+
+    actual_stats = _infer_actual_output_stats(result)
+    shm_name = None
+    resident_bytes = None
+    shm_bytes = 0
+    try:
         resident_bytes = result.calculate_memory_usage()
-        shared_backing_bytes = 0
+        data, shm_name, shm_bytes, resident_bytes = add_shared_memory_candidate(
+            result.data, resident_bytes
+        )
+        if data is not None:
+            result._data = data
+    except Exception as e:
+        shm_name = None
+        shm_bytes = 0
+        print(f"Failed to move worker result to shared memory: {e}")
 
-        if hasattr(result, "data"):
-            try:
-                data, shm_name, data_nbytes, resident_bytes = (
-                    add_shared_memory_candidate(result.data, resident_bytes)
-                )
-                if data is not None:
-                    result.data = data
-                    self.shared_memory_names[node_id] = [shm_name]
-                    shared_backing_bytes = data_nbytes
-            except Exception as e:
-                print(
-                    f"Failed to move cache entry {node_id} to shared memory, falling back to RAM: {e}"
-                )
-
-        self.cache[node_id] = result
-        self.memory_usage_per_node[node_id] = int(resident_bytes)
-        if DEBUG:
-            print(
-                f"Node {node_id} has a CPU memory usage of {self.memory_usage_per_node[node_id]/1024**3:.5f} GB"
-                + (
-                    f" (shared-memory backing: {shared_backing_bytes/1024**3:.5f} GB)"
-                    if shared_backing_bytes > 0
-                    else ""
-                )
-            )
-
-    def inc_ref(self, node_id: str):
-        if node_id not in self.ref_count:
-            self.ref_count[node_id] = 0
-        self.ref_count[node_id] += 1
-
-    def dec_ref(self, node_id: str):
-        if node_id not in self.ref_count:
-            return
-        self.ref_count[node_id] -= 1
-        if self.ref_count[node_id] <= 0:
-            self.ref_count[node_id] = 0
-            self._try_cleanup_node(node_id)
-
-    def clear(self, node_id: str):
-        if node_id in self.ref_count:
-            self.ref_count[node_id] = 0
-        self._try_cleanup_node(node_id)
-
-    def retain_shm_names(self, shm_names: List[str]) -> List[str]:
-        retained: List[str] = []
-        for shm_name in shm_names:
-            if not shm_name:
-                continue
-            self._shm_retain_count[shm_name] = (
-                self._shm_retain_count.get(shm_name, 0) + 1
-            )
-            retained.append(shm_name)
-        return retained
-
-    def release_shm_names(self, shm_names: List[str]) -> None:
-        nodes_to_recheck: List[str] = []
-        for shm_name in shm_names:
-            if not shm_name:
-                continue
-            count = self._shm_retain_count.get(shm_name, 0) - 1
-            if count <= 0:
-                self._shm_retain_count.pop(shm_name, None)
-            else:
-                self._shm_retain_count[shm_name] = count
-            for node_id, node_names in self.shared_memory_names.items():
-                if shm_name in node_names and node_id not in nodes_to_recheck:
-                    nodes_to_recheck.append(node_id)
-        for node_id in nodes_to_recheck:
-            self._try_cleanup_node(node_id)
-
-    def __len__(self):
-        return len(self.cache)
-
-    def get_memory_total_memory_usage(self):
-        return sum(self.memory_usage_per_node.values())
-
-    def _shm_names_in_use(self, shm_names: List[str]) -> bool:
-        return any(self._shm_retain_count.get(name, 0) > 0 for name in shm_names)
-
-    def _try_cleanup_node(self, node_id: str) -> None:
-        if self.ref_count.get(node_id, 0) > 0:
-            return
-        shm_names = self.shared_memory_names.get(node_id, [])
-        if shm_names and self._shm_names_in_use(shm_names):
-            return
-        self.cache.pop(node_id, None)
-        self.ref_count.pop(node_id, None)
-        self.memory_usage_per_node.pop(node_id, None)
-        self._cleanup_shared_memory(node_id)
-
-    def _cleanup_shared_memory(self, node_id: str):
-        names = self.shared_memory_names.pop(node_id, [])
-        for shm_name in names:
-            try:
-                shm = shared_memory.SharedMemory(name=shm_name)
-                shm.close()
-                shm.unlink()
-            except FileNotFoundError:
-                pass
-            except Exception:
-                pass
-
-    def cleanup_all(self):
-        self._shm_retain_count.clear()
-        for node_id in list(self.shared_memory_names.keys()):
-            self.ref_count.pop(node_id, None)
-            self.cache.pop(node_id, None)
-            self.memory_usage_per_node.pop(node_id, None)
-            self._cleanup_shared_memory(node_id)
+    return shm_name, resident_bytes, int(shm_bytes or 0), actual_stats
 
 
-def _execute_multiple_reps_for_leaf_dependencies(
-    nodes: List[RepresentationNode],
-    modalities: List[Modality],
-    gpu_id: Optional[int],
-):
-    representations = []
-    node_id_by_representation = {}
-    for node in nodes:
-        operation = node.operation(params=node.parameters)
-        if hasattr(operation, "gpu_id"):
-            operation.gpu_id = gpu_id
-        representations.append(operation)
-        node_id_by_representation[operation.name] = node.node_id
-
-    modality_results = modalities[0].apply_representations(
-        representations, parallel=True
-    )
-    return {
-        "results": modality_results,
-        "node_id_by_representation": node_id_by_representation,
-    }
-
-
-def _execute_node_worker(node, input_mods, task, rep_cache, gpu_id):
+def _execute_node_worker(node, input_mods: List[Any], gpu_id: Optional[int]):
     start_time = time.perf_counter()
     if gpu_id is not None:
         device = torch.device(f"cuda:{gpu_id}")
         torch.cuda.set_device(device)
         torch.cuda.reset_peak_memory_stats(device)
 
-    node_operation = node.operation(params=node.parameters)
+    node_operation = _instantiate_operation(node)
     operation_name = node_operation.name
     if DEBUG:
         print(f"Executing node {node.node_id} {operation_name} on GPU {gpu_id}")
@@ -258,47 +160,55 @@
                 return node_operation.transform(input_mods[0])
             elif isinstance(node_operation, UnimodalRepresentation):
                 pushdown_config = node.parameters.get("_pushdown_aggregation", None)
-                agg = None
-                if pushdown_config is not None:
-                    agg = AggregatedRepresentation(params=pushdown_config)
-                if rep_cache is not None and node_operation.name in rep_cache:
-                    return rep_cache[node_operation.name]
+                agg = (
+                    AggregatedRepresentation(params=pushdown_config)
+                    if pushdown_config is not None
+                    else None
+                )
                 return input_mods[0].apply_representation(
                     node_operation, aggregation=agg
                 )
             return input_mods[0].apply_representation(node_operation)
         else:
             fusion_op = node_operation
-            if hasattr(fusion_op, "needs_training") and fusion_op.needs_training:
+            if getattr(fusion_op, "needs_training", False):
                 return input_mods[0].combine_with_training(
-                    input_mods[1:], fusion_op, task
+                    input_mods[1:], fusion_op, None
                 )
             return input_mods[0].combine(input_mods[1:], fusion_op)
 
     gpu_peak_bytes = -1
-    peak_delta_bytes = -1
-    peak_abs_rss = -1
+    measurement = None
     if DEBUG:
-        result, peak_delta_bytes, peak_abs_rss = measure_peak_rss_during(
-            _run_node_op,
+        input_resident = sum(estimate_modality_bytes(m) for m in input_mods)
+        result, measurement = measure_memory_during(
+            lambda: _run_gpu_op(_run_node_op, gpu_id),
+            input_resident_bytes=input_resident,
             sample_s=0.01,
         )
         gpu_peak_bytes = (
             torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0
         )
     else:
-        result = _run_node_op()
+        result = _run_gpu_op(_run_node_op, gpu_id)
+
+    shm_name, resident_bytes, shm_bytes, actual_stats = _offload_to_shared_memory(
+        result
+    )
     end_time = time.perf_counter()
-    pid = os.getpid()
     return {
         "result": result,
-        "peak_bytes": peak_delta_bytes,
-        "peak_abs_rss_bytes": peak_abs_rss,
+        "result_shm_name": shm_name,
+        "result_resident_bytes": resident_bytes,
+        "result_shm_bytes": shm_bytes,
+        "actual_stats": actual_stats,
+        "memory": measurement,
+        "peak_bytes": measurement.increment_bytes if measurement else -1,
         "gpu_peak_bytes": gpu_peak_bytes,
         "operation_name": operation_name,
         "start_time": start_time,
         "end_time": end_time,
-        "pid": pid,
+        "pid": os.getpid(),
     }
 
 
@@ -307,7 +217,7 @@
     task: Any,
     modality: Any,
     gpu_id: Optional[int],
-    aggregation: AggregatedRepresentation = None,
+    aggregation=None,
 ) -> Dict[str, Any]:
     start_time = time.perf_counter()
     if DEBUG:
@@ -335,55 +245,150 @@
         return scores, end - start
 
     gpu_peak_bytes = -1
-    peak_delta_bytes = -1
+    measurement = None
     if DEBUG:
+        result, measurement = measure_memory_during(
+            lambda: _run_gpu_op(_run_task, gpu_id),
+            input_resident_bytes=estimate_modality_bytes(modality),
+            sample_s=0.01,
+        )
         gpu_peak_bytes = (
             torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0
         )
-        result, peak_delta_bytes, peak_abs_rss = measure_peak_rss_during(
-            _run_task,
-            sample_s=0.01,
-        )
-
-        print(
-            f"Task {task_node_id} has a CPU peak memory usage of {peak_delta_bytes/1024**3:.2f} GB, and a GPU peak memory usage of {gpu_peak_bytes/1024**3:.2f} GB"
-        )
     else:
-        result = _run_task()
+        result = _run_gpu_op(_run_task, gpu_id)
     end_time = time.perf_counter()
-    pid = os.getpid()
     return {
         "scores": result[0],
         "task_time": result[1],
-        "peak_bytes": peak_delta_bytes,
+        "memory": measurement,
+        "peak_bytes": measurement.increment_bytes if measurement else -1,
         "gpu_peak_bytes": gpu_peak_bytes,
         "start_time": start_time,
         "end_time": end_time,
-        "pid": pid,
+        "pid": os.getpid(),
     }
 
 
+def _execute_leaf_batch_worker(nodes: List[Any], modality: Any, gpu_id: Optional[int]):
+    node_id_by_representation = {}
+
+    def _run():
+        representations = []
+        for node in nodes:
+            operation = node.operation(params=node.parameters)
+            if hasattr(operation, "gpu_id"):
+                operation.gpu_id = gpu_id
+            representations.append(operation)
+            node_id_by_representation[operation.name] = node.node_id
+        return modality.apply_representations(representations, parallel=True)
+
+    modality_results = _run_gpu_op(_run, gpu_id)
+    shm_info = {}
+    for representation_name, transformed_modality in modality_results.items():
+        shm_name, resident_bytes, shm_bytes, actual_stats = _offload_to_shared_memory(
+            transformed_modality
+        )
+        shm_info[representation_name] = {
+            "shm_name": shm_name,
+            "resident_bytes": resident_bytes,
+            "shm_bytes": shm_bytes,
+            "actual_stats": actual_stats,
+        }
+    return {
+        "results": modality_results,
+        "node_id_by_representation": node_id_by_representation,
+        "shm_info": shm_info,
+    }
+
+
+def _load_leaf_worker(modality: Any) -> Dict[str, Any]:
+    if hasattr(modality, "extract_raw_data") and not modality.has_data():
+        modality.extract_raw_data()
+
+    data = modality.data
+    resident_bytes = 0
+    try:
+        resident_bytes = modality.estimate_memory_bytes()
+    except Exception:
+        resident_bytes = 0
+
+    wrapped, shm_name, _, resident_bytes = add_shared_memory_candidate(
+        data, resident_bytes
+    )
+    if wrapped is not None:
+        data = wrapped
+
+    return {"data": data, "metadata": modality.metadata, "shm_name": shm_name}
+
+
+def _dispatch_node(payload, gpu_id):
+    node, input_mods = payload
+    return _execute_node_worker(node, input_mods, gpu_id)
+
+
+def _dispatch_task(payload, gpu_id):
+    task_node_id, task, modality, aggregation = payload
+    return _execute_task_worker(task_node_id, task, modality, gpu_id, aggregation)
+
+
+def _dispatch_leaf_batch(payload, gpu_id):
+    nodes, modality = payload
+    return _execute_leaf_batch_worker(nodes, modality, gpu_id)
+
+
+def _dispatch_load_leaf(payload, _gpu_id):
+    (modality,) = payload
+    return _load_leaf_worker(modality)
+
+
+_WORKER_DISPATCH = {
+    "node": _dispatch_node,
+    "task": _dispatch_task,
+    "leaf_batch": _dispatch_leaf_batch,
+    "load_leaf": _dispatch_load_leaf,
+}
+
+
+@dataclass
+class _NodeUnit:
+    node_id: str
+
+
+@dataclass
+class _BatchUnit:
+    node_ids: List[str]
+
+
+@dataclass
+class ResultEntry:
+    val_score: PerformanceMeasure = None
+    train_score: PerformanceMeasure = None
+    test_score: PerformanceMeasure = None
+    representation_time: float = 0.0
+    task_time: float = 0.0
+    dag: RepresentationDag = None
+    tradeoff_score: float = 0.0
+
+
 class NodeExecutor:
     def __init__(
         self,
         dags: List[RepresentationDag],
         modalities: List[Modality],
         tasks: List[Any],
-        checkpoint_manager: Optional[CheckpointManager] = None,
         max_num_workers: int = -1,
         result_path: Optional[str] = None,
-        enable_checkpointing: bool = True,
+        enable_checkpointing: bool = False,
+        worker_pool: Optional[PersistentWorkerPool] = None,
     ):
         self.enable_checkpointing = enable_checkpointing
-        available_total_cpu = (
-            float(psutil.virtual_memory().available)
-            - float(psutil.virtual_memory().available) * 0.30
-        )
+        available_total_cpu = cpu_memory_budget_bytes()
         self.dags = dags
         self.scheduler = MemoryAwareNodeScheduler(
             dags, modalities, tasks, available_total_cpu
         )
-        self.checkpoint_manager = CheckpointManager(
+        self._checkpoint_manager = CheckpointManager(
             checkpoint_dir=result_path if result_path is not None else os.getcwd(),
             prefix=f"node_executor_checkpoint_{modalities[0].modality_id}_",
             checkpoint_every=1,
@@ -394,257 +399,281 @@
             if max_num_workers != -1
             else mp.cpu_count()
         )
-        self.modalities = modalities
-        self.tasks = tasks
-        self.result_cache = RefCountResultCache()
-        self.memory_usage_checkpoint = CheckpointManager(
+        self._modalities = modalities
+        self._tasks = tasks
+        self._result_path = result_path
+        self._result_cache = RefCountResultCache()
+        self._memory_usage_checkpoint = CheckpointManager(
             checkpoint_dir=result_path if result_path is not None else os.getcwd(),
             prefix=f"memory_usage_checkpoint_{modalities[0].modality_id}_",
             checkpoint_every=1,
             resume=False,
         )
-        self.statistics = {}
-        self.statistics["worker_stats"] = {}
-        self.statistics["node_stats"] = {}
+        self._memory_usage_data: Dict[str, Any] = {}
+        self.statistics = {"worker_stats": {}, "node_stats": {}}
 
-    def _shm_names_for_submit(
-        self, parent_node_ids: List[str], payload_data: Any
-    ) -> List[str]:
+        self._node_attempts: Dict[str, int] = {}
+
+        self._job_units: Dict[int, Union[_NodeUnit, _BatchUnit]] = {}
+        self._job_retained_shm: Dict[int, List[str]] = {}
+        self._leaf_shm_names: List[str] = []
+        self._task_results: Dict[str, ResultEntry] = {}
+
+        self._owns_pool = worker_pool is None
+        if worker_pool is None:
+            cpu_count = os.cpu_count() or 1
+            threads_per_worker = max(1, cpu_count // max(1, self.max_num_workers))
+            worker_pool = PersistentWorkerPool(
+                self.max_num_workers,
+                _WORKER_DISPATCH,
+                ctx=create_mp_context(),
+                threads_per_worker=threads_per_worker,
+            )
+        self._pool = worker_pool
+
+    def _requeue_or_give_up(self, node_id: str, reason: str) -> bool:
+        attempts = self._node_attempts.get(node_id, 0) + 1
+        self._node_attempts[node_id] = attempts
+        if attempts > _MAX_NODE_RETRIES:
+            print(
+                f"[node_executor] giving up on node {node_id} after {attempts} "
+                f"failed attempts ({reason}); marking it permanently failed and "
+                f"continuing with the rest of the search.",
+                flush=True,
+            )
+            self.scheduler.add_failed_node(node_id, reason)
+            return False
+        print(
+            f"[node_executor] node {node_id} did not complete (attempt "
+            f"{attempts}/{_MAX_NODE_RETRIES + 1}): {reason}. Re-queuing it for "
+            f"another attempt.",
+            flush=True,
+        )
+        self.scheduler.requeue_node(node_id)
+        return True
+
+    def _release_parents(self, node_id: str) -> None:
+        for parent_id in self.scheduler.get_valid_parents(node_id):
+            self._result_cache.dec_ref(parent_id)
+
+    def _retain_for_submit(self, parent_ids: List[str], payload: Any) -> List[str]:
         names: List[str] = []
-        for parent_id in parent_node_ids or []:
-            names.extend(self.result_cache.shared_memory_names.get(parent_id, []))
-        if parent_node_ids:
-            names.extend(collect_shm_names_from_payload(payload_data))
-        else:
-            names.extend(getattr(self, "_leaf_shm_names", []))
-            names.extend(collect_shm_names_from_payload(payload_data))
-        # preserve order, drop duplicates
-        return list(dict.fromkeys(names))
+        for parent_id in parent_ids or []:
+            names.extend(self._result_cache.shared_memory_names.get(parent_id, []))
+        if not parent_ids:
+            names.extend(self._leaf_shm_names)
+        names.extend(collect_shm_names_from_payload(payload))
+        names = list(dict.fromkeys(names))
+        return self._result_cache.retain_shm_names(names)
 
-    def _retain_for_submit(
-        self, parent_node_ids: List[str], payload_data: Any
-    ) -> List[str]:
-        return self.result_cache.retain_shm_names(
-            self._shm_names_for_submit(parent_node_ids, payload_data)
+    def _load_leaf_modalities(self) -> None:
+        for modality in self._modalities:
+            if getattr(modality, "has_data", None) and modality.has_data():
+                continue
+            attempts = 0
+            while True:
+                self._pool.submit("load_leaf", (modality,), gpu_id=None)
+                jr = self._pool.wait()
+                if jr.ok:
+                    break
+                attempts += 1
+                if attempts > _MAX_NODE_RETRIES:
+                    raise RuntimeError(
+                        f"Failed to load leaf modality {modality.modality_id}: "
+                        f"{jr.error}"
+                    )
+            modality._data = jr.value["data"]
+            modality.metadata = jr.value["metadata"]
+            shm_name = jr.value.get("shm_name")
+            if shm_name is not None:
+                self._leaf_shm_names.append(shm_name)
+
+    def _cleanup_leaf_shared_memory(self) -> None:
+        for shm_name in self._leaf_shm_names:
+            unlink_shm(shm_name)
+        self._leaf_shm_names = []
+
+    def _submit_node(self, node_id: str) -> None:
+        node = self.scheduler.mapping[node_id]
+        gpu_id = node.gpu_id
+        parent_ids = self.scheduler.get_valid_parents(node_id)
+        parent_results = (
+            [self._result_cache.get(pid) for pid in parent_ids] if parent_ids else None
         )
 
-    def _release_for_future(self, retained_shm_names: List[str]) -> None:
-        if retained_shm_names:
-            self.result_cache.release_shm_names(retained_shm_names)
+        if self._is_task_node(node):
+            task_idx = int(node.parameters.get("_task_idx", 0))
+            payload = (
+                self._modalities[0] if parent_results is None else parent_results[0]
+            )
+            self._task_results[node_id] = ResultEntry(
+                dag=self._get_dag_from_node_ids(node_id),
+                representation_time=payload.transform_time,
+            )
+            retained = self._retain_for_submit(parent_ids, payload)
+            self.scheduler.begin_execution(node_id)
+            self.scheduler.move_to_running(node_id)
+            job_id = self._pool.submit(
+                "task",
+                (node_id, self._tasks[task_idx], payload, node.aggregation),
+                gpu_id=gpu_id,
+            )
+        else:
+            payload = self._modalities if parent_results is None else parent_results
+            retained = self._retain_for_submit(parent_ids, payload)
+            self.scheduler.begin_execution(node_id)
+            self.scheduler.move_to_running(node_id)
+            job_id = self._pool.submit("node", (node, payload), gpu_id=gpu_id)
 
-    def run(self) -> None:
-        task_results = {}
-        memory_usage_data = {}
+        self._job_units[job_id] = _NodeUnit(node_id)
+        self._job_retained_shm[job_id] = retained
 
-        self._materialize_leaf_modalities_in_shared_memory()
+    def _submit_leaf_batch(self, node_ids: List[str]) -> None:
+        nodes = [self.scheduler.mapping[nid] for nid in node_ids]
+        gpu_id = nodes[0].gpu_id
+        retained = self._retain_for_submit([], self._modalities[0].data)
+        for nid in node_ids:
+            self.scheduler.begin_execution(nid)
+        self.scheduler.move_to_running(node_ids)
+        job_id = self._pool.submit(
+            "leaf_batch", (nodes, self._modalities[0]), gpu_id=gpu_id
+        )
+        self._job_units[job_id] = _BatchUnit(node_ids)
+        self._job_retained_shm[job_id] = retained
 
-        ctx = mp.get_context("spawn")
-        with ProcessPoolExecutor(
-            max_workers=self.max_num_workers, mp_context=ctx
-        ) as executor:
-            future_to_node_id = {}
-            future_to_retained_shm: Dict[Any, List[str]] = {}
-
-            def submit_nodes_with_leaf_dependencies(node_ids: List[str]):
-                nodes = [self.scheduler.mapping[node_id] for node_id in node_ids]
-                gpu_id = nodes[0].gpu_id
-                self.scheduler.move_to_running(node_ids)
-
-                retained = self._retain_for_submit([], self.modalities[0].data)
-                future = executor.submit(
-                    _execute_multiple_reps_for_leaf_dependencies,
-                    nodes,
-                    self.modalities,
-                    gpu_id,
-                )
-                future_to_node_id[future] = node_ids
-                future_to_retained_shm[future] = retained
-
-            def submit_node(node_id: str):
-                node = self.scheduler.mapping[node_id]
-                gpu_id = node.gpu_id
-                parent_node_ids = self.scheduler.get_valid_parents(node_id)
-                parent_results = None
-                if parent_node_ids:
-                    parent_results = [
-                        self.result_cache.get(parent_node_id)
-                        for parent_node_id in parent_node_ids
-                    ]
-
-                if self._is_task_node(node):
-                    # potentially batch task nodes and then execute them together
-                    # by the the same task type (index, gpu vs cpu)
-                    # either enough nodes to batch or enough time to batch whatever happens first
-
-                    task_result = ResultEntry(
-                        dag=self._get_dag_from_node_ids(node_id),
-                        representation_time=parent_results[0].transform_time,
-                    )
-                    task_results[node_id] = task_result
-                    task_idx = int(node.parameters.get("_task_idx", 0))
-                    payload_data = (
-                        self.modalities[0]
-                        if parent_results is None
-                        else parent_results[0]
-                    )
-                    retained = self._retain_for_submit(parent_node_ids, payload_data)
-                    aggregation = node.aggregation
-
-                    future = executor.submit(
-                        _execute_task_worker,
-                        node_id,
-                        self.tasks[task_idx],
-                        payload_data,
-                        gpu_id,
-                        aggregation,
-                    )
-                else:
-                    payload_data = (
-                        self.modalities if parent_results is None else parent_results
-                    )
-                    retained = self._retain_for_submit(parent_node_ids, payload_data)
-                    future = executor.submit(
-                        _execute_node_worker,
-                        node,
-                        payload_data,
-                        None,
-                        None,
-                        gpu_id,
-                    )
-                self.scheduler.move_to_running(node_id)
-                future_to_node_id[future] = node_id
-                future_to_retained_shm[future] = retained
-
-            def submit_new_ready_nodes():
-                ready_nodes = self.scheduler.get_runnable().copy()
-                for node_id in ready_nodes:
-                    if isinstance(node_id, list):
-                        submit_nodes_with_leaf_dependencies(node_id)
-                        continue
-                    submit_node(node_id)
-
-            submit_new_ready_nodes()
-
-            while future_to_node_id or not self.scheduler.is_finished():
-                if not future_to_node_id:
-                    submit_new_ready_nodes()
+    def _fill_pipeline(self) -> None:
+        ready = self.scheduler.get_runnable().copy()
+        for entry in ready:
+            if not self._pool.has_idle_worker:
+                break
+            if isinstance(entry, list):
+                self._submit_leaf_batch(entry)
+            else:
+                if not self.scheduler.can_start_now(entry):
                     continue
+                self._submit_node(entry)
 
-                done, _ = wait(
-                    set(future_to_node_id.keys()), return_when=FIRST_COMPLETED
+    def _record_stats(self, node_id: str, pid: int, start_time: float, end_time: float):
+        node_stats = self.statistics["node_stats"]
+        worker_stats = self.statistics["worker_stats"]
+        node_stats[node_id] = {"start_time": start_time, "end_time": end_time}
+        entry = worker_stats.get(pid)
+        if entry is None:
+            worker_stats[pid] = {
+                "start_time": start_time,
+                "end_time": end_time,
+                "busy_time": end_time - start_time,
+                "num_jobs": 1,
+            }
+        else:
+            entry["start_time"] = min(entry["start_time"], start_time)
+            entry["end_time"] = max(entry["end_time"], end_time)
+            entry["busy_time"] += end_time - start_time
+            entry["num_jobs"] += 1
+
+    def _process_result(self, jr) -> None:
+        unit = self._job_units.pop(jr.job_id, None)
+        retained = self._job_retained_shm.pop(jr.job_id, [])
+        try:
+            if unit is None:
+                return
+            if not jr.ok:
+                self._handle_job_failure(unit, jr)
+                return
+            if isinstance(unit, _BatchUnit):
+                self._handle_batch_success(jr.value)
+            else:
+                self._handle_node_success(unit.node_id, jr.value)
+        finally:
+            if retained:
+                self._result_cache.release_shm_names(retained)
+
+    def _handle_job_failure(self, unit: Union[_NodeUnit, _BatchUnit], jr):
+        reason = jr.error or "unknown worker failure"
+        if jr.cuda_oom:
+            reason += " (CUDA out of memory)"
+        node_ids = unit.node_ids if isinstance(unit, _BatchUnit) else [unit.node_id]
+        for node_id in node_ids:
+            requeued = self._requeue_or_give_up(node_id, reason)
+            if not requeued:
+                self._release_parents(node_id)
+
+    def _handle_node_success(self, node_id: str, value: Dict[str, Any]) -> None:
+        node = self.scheduler.mapping[node_id]
+        if "pid" in value:
+            self._record_stats(
+                node_id, value["pid"], value["start_time"], value["end_time"]
+            )
+
+        if self._is_task_node(node):
+            entry = self._task_results[node_id]
+            entry.task_time = value["task_time"]
+            entry.train_score = value["scores"][0].average_scores
+            entry.val_score = value["scores"][1].average_scores
+            entry.test_score = value["scores"][2].average_scores
+            if self.enable_checkpointing:
+                self._checkpoint_manager.increment(node_id)
+                self._checkpoint_manager.checkpoint_if_due(
+                    self._task_results, self._discard_report()
                 )
+                self._checkpoint_memory_usage(
+                    node_id,
+                    value["peak_bytes"],
+                    value["gpu_peak_bytes"],
+                    "task",
+                    None,
+                    measurement=value.get("memory"),
+                )
+            self._release_parents(node_id)
+            self.scheduler.complete_node(node_id)
+        else:
+            self._handle_modality_result(
+                value["result"],
+                node_id,
+                value["peak_bytes"],
+                value["gpu_peak_bytes"],
+                value["operation_name"],
+                actual_stats=value.get("actual_stats"),
+                shm_name=value.get("result_shm_name"),
+                resident_bytes=value.get("result_resident_bytes"),
+                shm_bytes=value.get("result_shm_bytes", 0),
+                measurement=value.get("memory"),
+            )
 
-                for future in done:
-                    node_id = future_to_node_id.pop(future)
-                    retained_shm = future_to_retained_shm.pop(future, [])
-                    try:
-                        result = future.result()
-
-                        if isinstance(node_id, list):
-                            results = result["results"]
-                            node_id_by_representation = result[
-                                "node_id_by_representation"
-                            ]
-                            for (
-                                representation,
-                                transformed_modality,
-                            ) in results.items():
-                                batch_node_id = node_id_by_representation[
-                                    representation
-                                ]
-                                self._handle_modality_result(
-                                    transformed_modality,
-                                    batch_node_id,
-                                    None,
-                                    None,
-                                    memory_usage_data,
-                                    representation,
-                                )
-                            submit_new_ready_nodes()
-                            continue
-
-                        peak_bytes = result["peak_bytes"]
-                        gpu_peak_bytes = result["gpu_peak_bytes"]
-                        node = self.scheduler.mapping[node_id]
-                        self.statistics["worker_stats"][result["pid"]] = {
-                            "start_time": result["start_time"],
-                            "end_time": result["end_time"],
-                        }
-                        self.statistics["node_stats"][node_id] = {
-                            "start_time": result["start_time"],
-                            "end_time": result["end_time"],
-                        }
-                        if self._is_task_node(node):
-                            task_results[node_id].task_time = result["task_time"]
-                            task_results[node_id].train_score = result["scores"][
-                                0
-                            ].average_scores
-                            task_results[node_id].val_score = result["scores"][
-                                1
-                            ].average_scores
-                            task_results[node_id].test_score = result["scores"][
-                                2
-                            ].average_scores
-                            if self.enable_checkpointing:
-                                self.checkpoint_manager.increment(node_id)
-                                self.checkpoint_manager.checkpoint_if_due(task_results)
-                                self._checkpoint_memory_usage(
-                                    node_id,
-                                    peak_bytes,
-                                    gpu_peak_bytes,
-                                    "task",
-                                    memory_usage_data,
-                                    None,
-                                )
-
-                            parent_node_ids = self.scheduler.get_valid_parents(node_id)
-                            for parent_node_id in parent_node_ids:
-                                self.result_cache.dec_ref(parent_node_id)
-                            self.scheduler.complete_node(node_id)
-                        else:
-                            transformed_modality = result["result"]
-                            self._handle_modality_result(
-                                transformed_modality,
-                                node_id,
-                                peak_bytes,
-                                gpu_peak_bytes,
-                                memory_usage_data,
-                                result["operation_name"],
-                            )
-
-                        submit_new_ready_nodes()
-                    except Exception:
-                        parent_node_ids = []
-                        if not isinstance(node_id, list):
-                            parent_node_ids = self.scheduler.get_valid_parents(node_id)
-                        for parent_node_id in parent_node_ids:
-                            self.result_cache.dec_ref(parent_node_id)
-                        if not isinstance(node_id, list):
-                            self.scheduler.add_failed_node(node_id)
-                        raise
-                    finally:
-                        self._release_for_future(retained_shm)
-
-        assert not self.result_cache.ref_count
-        assert not self.result_cache._shm_retain_count
-
-        self.result_cache.cleanup_all()
-        self._cleanup_leaf_shared_memory()
-        return {
-            "task_results": list(task_results.values()),
-            "statistics": self.statistics,
-        }
+    def _handle_batch_success(self, value: Dict[str, Any]) -> None:
+        results = value["results"]
+        node_id_by_representation = value["node_id_by_representation"]
+        shm_info = value.get("shm_info", {})
+        for representation, transformed_modality in results.items():
+            node_id = node_id_by_representation[representation]
+            info = shm_info.get(representation, {})
+            self._handle_modality_result(
+                transformed_modality,
+                node_id,
+                None,
+                None,
+                representation,
+                actual_stats=info.get("actual_stats"),
+                shm_name=info.get("shm_name"),
+                resident_bytes=info.get("resident_bytes"),
+                shm_bytes=info.get("shm_bytes", 0),
+            )
 
     def _handle_modality_result(
         self,
         transformed_modality: Any,
         node_id: str,
-        peak_bytes: int,
-        gpu_peak_bytes: int,
-        memory_usage_data,
+        peak_bytes: Optional[int],
+        gpu_peak_bytes: Optional[int],
         operation_name: str,
+        actual_stats: Optional[RepresentationStats] = None,
+        shm_name: Optional[str] = None,
+        resident_bytes: Optional[int] = None,
+        shm_bytes: int = 0,
+        measurement: Optional[MemoryMeasurement] = None,
     ):
-        actual_stats = self._infer_actual_output_stats(transformed_modality)
+        if actual_stats is None:
+            actual_stats = _infer_actual_output_stats(transformed_modality)
         estimated_stats = self.scheduler.node_stats.get(node_id)
 
         if actual_stats is not None and (
@@ -660,169 +689,222 @@
                 peak_bytes,
                 gpu_peak_bytes,
                 operation_name,
-                memory_usage_data,
-                transformed_modality.data,
+                actual_stats,
+                measurement=measurement,
             )
-        before_bytes = self.result_cache.get_memory_total_memory_usage()
-        self._manage_result_cache(node_id, transformed_modality)
-        after_bytes = self.result_cache.get_memory_total_memory_usage()
+        before_bytes = self._result_cache.get_memory_total_memory_usage()
+        self._manage_result_cache(
+            node_id,
+            transformed_modality,
+            shm_name=shm_name,
+            resident_bytes=resident_bytes,
+            shm_bytes=shm_bytes,
+        )
+        after_bytes = self._result_cache.get_memory_total_memory_usage()
         self.scheduler.update_cpu_memory_in_use(after_bytes - before_bytes)
         self.scheduler.complete_node(node_id)
 
-    def _materialize_leaf_modalities_in_shared_memory(self):
-        self._leaf_shm_names = []
-        for modality in self.modalities:
-            if hasattr(modality, "extract_raw_data") and not modality.has_data():
-                modality.extract_raw_data()
-            data, shm_name, _, _ = add_shared_memory_candidate(modality.data)
-            if shm_name is not None:
-                modality.data = data
-                self._leaf_shm_names.append(shm_name)
+    def _manage_result_cache(
+        self,
+        node_id: str,
+        result: Any,
+        shm_name: Optional[str] = None,
+        resident_bytes: Optional[int] = None,
+        shm_bytes: int = 0,
+    ):
+        self._release_parents(node_id)
 
-    def _cleanup_leaf_shared_memory(self):
-        for shm_name in getattr(self, "_leaf_shm_names", []):
-            try:
-                shm = shared_memory.SharedMemory(name=shm_name)
-                shm.close()
-                shm.unlink()
-            except FileNotFoundError:
-                pass
-            except Exception:
-                pass
-        self._leaf_shm_names = []
+        children = self.scheduler.get_children(node_id)
+        if children:
+            for _ in children:
+                self._result_cache.inc_ref(node_id)
+            self._result_cache.add_result(
+                node_id,
+                result,
+                shm_name=shm_name,
+                resident_bytes=resident_bytes,
+                shm_bytes=shm_bytes,
+            )
+        elif shm_name is not None:
+            unlink_shm(shm_name)
 
     def _checkpoint_memory_usage(
         self,
         node_id: str,
-        peak_bytes: int,
-        gpu_peak_bytes: int,
+        peak_bytes: Optional[int],
+        gpu_peak_bytes: Optional[int],
         operation_name: str,
-        data,
-        result,
+        actual_stats: Optional[RepresentationStats],
+        measurement: Optional[MemoryMeasurement] = None,
     ):
-        self.memory_usage_checkpoint.increment(node_id)
+        if self.enable_checkpointing:
+            self._memory_usage_checkpoint.increment(node_id)
+
+        if measurement is not None:
+            peak_bytes = measurement.footprint_bytes
 
         shape = None
         if DEBUG:
-            shape = self._print_node_stats(node_id, result, operation_name)
-            if peak_bytes > self.scheduler.node_resources[node_id][0]:
-                print(
-                    f"UNDERESTIMATED PEAK MEMORY: Peak bytes: {peak_bytes/1024**3:.2f} GB, Estimated CPU bytes: {self.scheduler.node_resources[node_id][0]/1024**3:.2f} GB for node {node_id}: {operation_name}"
-                )
-            if gpu_peak_bytes > self.scheduler.node_resources[node_id][1]:
-                print(
-                    f"UNDERESTIMATED GPU PEAK MEMORY: GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, Estimated GPU bytes: {self.scheduler.node_resources[node_id][1]/1024**3:.2f} GB for node {node_id}: {operation_name}"
-                )
-            if self.scheduler.node_resources[node_id][0] >= peak_bytes * 2:
-                print(
-                    f"Peak bytes: {peak_bytes/1024**3:.2f} GB, Estimated CPU bytes: {self.scheduler.node_resources[node_id][0]/1024**3:.2f} GB, 200% of estimated for node {node_id}: {operation_name}"
-                )
-            if self.scheduler.node_resources[node_id][1] > gpu_peak_bytes * 2:
-                print(
-                    f"GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, Estimated GPU bytes: {self.scheduler.node_resources[node_id][1]/1024**3:.2f} GB, 200% of estimated for node {node_id}: {operation_name}"
-                )
-        data[node_id] = {
-            "cpu_peak_bytes": peak_bytes,
-            "gpu_peak_bytes": gpu_peak_bytes,
-            "operation_name": operation_name,
-            "estimated_cpu_bytes": self.scheduler.node_resources[node_id][0],
-            "estimated_gpu_bytes": self.scheduler.node_resources[node_id][1],
-            "shape": shape,
-        }
-        self.memory_usage_checkpoint.checkpoint_if_due(data)
+            shape = self._print_node_stats(node_id, actual_stats, operation_name)
+            est_cpu, est_gpu = self.scheduler.node_resources[node_id]
+            if peak_bytes is not None and peak_bytes >= 0:
+                if peak_bytes > est_cpu:
+                    print(
+                        f"UNDERESTIMATED PEAK MEMORY: Peak bytes: {peak_bytes/1024**3:.2f} GB, "
+                        f"Estimated CPU bytes: {est_cpu/1024**3:.2f} GB for node {node_id}: {operation_name}"
+                    )
+                if est_cpu >= peak_bytes * 2:
+                    print(
+                        f"Peak bytes: {peak_bytes/1024**3:.2f} GB, Estimated CPU bytes: "
+                        f"{est_cpu/1024**3:.2f} GB, >200% of estimated for node {node_id}: {operation_name}"
+                    )
+            if gpu_peak_bytes is not None and gpu_peak_bytes >= 0:
+                if gpu_peak_bytes > est_gpu:
+                    print(
+                        f"UNDERESTIMATED GPU PEAK MEMORY: GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, "
+                        f"Estimated GPU bytes: {est_gpu/1024**3:.2f} GB for node {node_id}: {operation_name}"
+                    )
+                if est_gpu > gpu_peak_bytes * 2:
+                    print(
+                        f"GPU peak bytes: {gpu_peak_bytes/1024**3:.2f} GB, Estimated GPU bytes: "
+                        f"{est_gpu/1024**3:.2f} GB, >200% of estimated for node {node_id}: {operation_name}"
+                    )
+        if self.enable_checkpointing:
+            self._memory_usage_data[node_id] = {
+                "cpu_peak_bytes": peak_bytes if peak_bytes is not None else -1,
+                "gpu_peak_bytes": gpu_peak_bytes if gpu_peak_bytes is not None else -1,
+                "operation_name": operation_name,
+                "estimated_cpu_bytes": self.scheduler.node_resources[node_id][0],
+                "estimated_gpu_bytes": self.scheduler.node_resources[node_id][1],
+                "shape": shape,
+                "cpu_increment_bytes": (
+                    measurement.increment_bytes if measurement else -1
+                ),
+                "cpu_footprint_bytes": (
+                    measurement.footprint_bytes if measurement else -1
+                ),
+                "input_resident_bytes": (
+                    measurement.input_resident_bytes if measurement else -1
+                ),
+                "traced_peak_bytes": (
+                    measurement.traced_peak_bytes if measurement else -1
+                ),
+                "rss_delta_bytes": measurement.rss_delta_bytes if measurement else -1,
+                "num_instances": getattr(actual_stats, "num_instances", None),
+                "output_shape": getattr(actual_stats, "output_shape", None),
+                "dtype": str(getattr(actual_stats, "dtype", None)),
+                "container": str(getattr(actual_stats, "container", None)),
+            }
+            self._memory_usage_checkpoint.checkpoint_if_due(self._memory_usage_data)
 
-    def _print_node_stats(self, node_id: str, result: Any, operation_name: str):
-        if (
-            result is not None
-            and operation_name != "BoW"
-            and not operation_name.endswith("Split")
-        ):
-            node_stats = self.scheduler.node_stats[node_id]
-            shape = None
-            if isinstance(result[0], list):
-                if isinstance(result[0][0], np.ndarray):
-                    shape = (len(result[0]), *result[0][0].shape)
-                elif isinstance(result[0][0], list):
-                    shape = (len(result[0]), *result[0][0][0].shape)
-                else:
-                    shape = (len(result[0]), *result[0][0].shape)
-            else:
-                shape = result[0].shape
-            print(
-                f"Node {node_id} {operation_name} should have shape of {node_stats.num_instances, node_stats.output_shape}, actual shape: {len(result), shape} output shape is known: {node_stats.output_shape_is_known}"
-            )
-            if node_stats.output_shape_is_known:
-                assert (
-                    len(result) == node_stats.num_instances
-                ), f"Node {node_id} {operation_name} should have {node_stats.num_instances} instances, actual: {len(result)}"
-                # assert (
-                #     shape == node_stats.output_shape
-                # ), f"Node {node_id} {operation_name} should have shape of {node_stats.output_shape}, actual shape: {shape}"
-            return shape
-
-    def _infer_actual_output_stats(
-        self, transformed_modality: Any
-    ) -> Optional[RepresentationStats]:
-        if transformed_modality is None or not hasattr(transformed_modality, "data"):
+    def _print_node_stats(
+        self,
+        node_id: str,
+        actual_stats: Optional[RepresentationStats],
+        operation_name: str,
+    ):
+        if actual_stats is None:
             return None
-
-        data = transformed_modality.data
-
-        if isinstance(data, np.ndarray):
-            if data.ndim == 0:
-                return RepresentationStats(1, (1,), output_shape_is_known=True)
-            num_instances = int(data.shape[0])
-            output_shape = (
-                tuple(int(d) for d in data.shape[1:]) if data.ndim > 1 else (1,)
+        node_stats = self.scheduler.node_stats.get(node_id)
+        shape = actual_stats.output_shape
+        if node_stats is not None:
+            print(
+                f"Node {node_id} {operation_name} should have shape of "
+                f"{node_stats.num_instances, node_stats.output_shape}, actual shape: "
+                f"{actual_stats.num_instances, shape} output shape is known: "
+                f"{node_stats.output_shape_is_known}"
             )
-            return RepresentationStats(
-                num_instances, output_shape, output_shape_is_known=True
-            )
+        return shape
 
-        if isinstance(data, list) and len(data) > 0 and isinstance(data[0], np.ndarray):
-            num_instances = len(data)
-            first_shape = tuple(int(d) for d in data[0].shape)
-            same_shape = all(
-                isinstance(x, np.ndarray) and x.shape == data[0].shape for x in data
-            )
-            return RepresentationStats(
-                num_instances,
-                first_shape,
-                output_shape_is_known=bool(same_shape),
-            )
-
-        return None
-
-    def _manage_result_cache(self, node_id: str, result: Any):
-        parent_node_ids = self.scheduler.get_valid_parents(node_id)
-        for parent_node_id in parent_node_ids:
-            self.result_cache.dec_ref(parent_node_id)
-
-        if self.scheduler.get_children(node_id):
-            for _ in self.scheduler.get_children(node_id):
-                self.result_cache.inc_ref(node_id)
-            self.result_cache.add_result(node_id, result)
-
-    def _get_nodes_by_ids(self, nodes_ids: List[str]) -> List[RepresentationNode]:
-        return [self.scheduler.mapping[node_id] for node_id in nodes_ids]
-
-    def _get_dag_from_node_ids(self, node_id: str) -> RepresentationDag:
+    def _get_dag_from_node_ids(self, node_id: str) -> Optional[RepresentationDag]:
         for dag in self.dags:
             if dag.root_node_id == node_id:
                 return dag
         return None
 
+    def _describe_node(self, node_id: str) -> str:
+        names = []
+        seen = set()
+        current = node_id
+        while current is not None and current not in seen:
+            seen.add(current)
+            node = self.scheduler.mapping.get(current)
+            if node is None:
+                break
+            if node.operation is not None:
+                try:
+                    names.append(node.operation().name)
+                except Exception:
+                    names.append(
+                        getattr(node.operation, "__name__", str(node.operation))
+                    )
+            parent_ids = [pid for pid in self.scheduler.get_valid_parents(current)]
+            current = parent_ids[0] if len(parent_ids) == 1 else None
+            if len(parent_ids) > 1:
+                names.append(
+                    "["
+                    + ", ".join(self._describe_node(pid) for pid in parent_ids)
+                    + "]"
+                )
+        names.reverse()
+        return " -> ".join(names) if names else node_id
+
+    def _discard_report(self) -> Dict[str, Any]:
+        return {
+            "failed_nodes": {
+                node_id: {
+                    "representation": self._describe_node(node_id),
+                    "reason": reason,
+                }
+                for node_id, reason in self.scheduler.failed_node_reasons.items()
+            },
+            "blocked_memory_nodes": {
+                node_id: {
+                    "representation": self._describe_node(node_id),
+                    "reason": reason,
+                }
+                for node_id, reason in self.scheduler.blocked_memory_reasons.items()
+            },
+            "cpu_fallback_nodes": {
+                node_id: {
+                    "representation": self._describe_node(node_id),
+                    "reason": reason,
+                }
+                for node_id, reason in self.scheduler.cpu_fallback_reasons.items()
+            },
+            "deadlock": self.scheduler.deadlock,
+            "deadlock_reason": self.scheduler.deadlock_reason,
+        }
+
     @staticmethod
     def _is_task_node(node: RepresentationNode) -> bool:
         return bool(getattr(node, "parameters", {}).get("_node_kind") == "task")
 
+    def run(self) -> Dict[str, Any]:
+        self._task_results = {}
+        self._load_leaf_modalities()
 
-@dataclass
-class ResultEntry:
-    val_score: PerformanceMeasure = None
-    train_score: PerformanceMeasure = None
-    test_score: PerformanceMeasure = None
-    representation_time: float = 0.0
-    task_time: float = 0.0
-    dag: RepresentationDag = None
-    tradeoff_score: float = 0.0
+        try:
+            self._fill_pipeline()
+            while self._job_units or not self.scheduler.is_finished():
+                self._fill_pipeline()
+                if not self._job_units:
+                    continue
+                jr = self._pool.wait()
+                self._process_result(jr)
+                self._fill_pipeline()
+        finally:
+            if self._owns_pool:
+                self._pool.shutdown()
+            self._result_cache.cleanup_all()
+            self._cleanup_leaf_shared_memory()
+
+        if self.enable_checkpointing:
+            self._checkpoint_manager.save_checkpoint(
+                self._task_results, self._discard_report()
+            )
+
+        return {
+            "task_results": list(self._task_results.values()),
+            "statistics": self.statistics,
+        }
diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py
index 209f450..1ca681e 100644
--- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py
+++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py
@@ -19,6 +19,7 @@
 #
 # -------------------------------------------------------------
 from __future__ import annotations
+import os
 import re
 from typing import List, Dict, Optional, Any
 from collections import defaultdict, deque
@@ -29,9 +30,15 @@
     RepresentationNode,
 )
 from systemds.scuro.modality.modality import Modality
+from systemds.scuro.representations.representation import (
+    stats_dtype,
+    stats_itemsize,
+)
 from systemds.scuro.utils.memory_utility import gpu_memory_info
 from systemds.scuro.utils.static_variables import DEBUG
 
+_MAX_GPU_SCHEDULE_ATTEMPTS = int(os.environ.get("SCURO_MAX_GPU_SCHEDULE_ATTEMPTS", "3"))
+
 
 class MemoryAwareNodeScheduler:
 
@@ -66,37 +73,82 @@
         self.success = False
         self.deadlock = False
         self.ready_nodes = []
-        self.running_nodes = []
+        self._ready_set = set()
+        self.running_nodes = set()
         self.completed_nodes = []
+        self._completed_set = set()
         self.failed_nodes = []
+        self.failed_node_reasons: Dict[str, str] = {}
         self.blocked_memory_nodes_perm = []
+        self.blocked_memory_reasons: Dict[str, str] = {}
         self.cancelled_nodes = []
+        self.deadlock_reason: Optional[str] = None
+        self.gpu_wait_attempts: Dict[str, int] = {}
+        self.cpu_fallback_nodes: List[str] = []
+        self.cpu_fallback_reasons: Dict[str, str] = {}
+        self._candidates = {
+            node_id
+            for node_id in self.topo_order
+            if node_id not in self.leaves and self.unresolved_parents[node_id] == 0
+        }
         self.n_gpu = (
             torch.cuda.device_count() if torch and torch.cuda.is_available() else 0
         )
+        leaf_cached = sum(self.node_resources[node][0] for node in self.leaves)
         self.memory_stats = {
-            "cpu_in_use": sum([self.node_resources[node][0] for node in self.leaves]),
+            "cpu_cached": leaf_cached,
+            "cpu_in_flight": 0,
             "gpu_in_use": {
                 info["index"]: int(info["total_b"] - info["free_b"])
                 for info in self.gpu_memory_info
             },
         }
+        self._cpu_reserved_nodes: Dict[str, int] = {}
         self._initialized = False
 
+    def _total_cpu_in_use(self) -> float:
+        return self.memory_stats["cpu_cached"] + self.memory_stats["cpu_in_flight"]
+
+    def _pending_admitted_cpu_bytes(self) -> int:
+        total = 0
+        for node_id in self._ready_set:
+            if node_id in self._cpu_reserved_nodes:
+                continue
+            resources = self.node_resources.get(node_id)
+            if resources:
+                total += resources[0]
+        return int(total)
+
+    def can_start_now(self, node_id: str) -> bool:
+        resources = self.node_resources.get(node_id)
+        if not resources:
+            return True
+        if not self._cpu_reserved_nodes:
+            return True
+        return resources[0] <= self.memory_budget["cpu"] - self._total_cpu_in_use()
+
     def update_cpu_memory_in_use(self, delta_bytes: int):
-        self.memory_stats["cpu_in_use"] += delta_bytes
+        self.memory_stats["cpu_cached"] += delta_bytes
 
     def get_runnable(self) -> List[RepresentationNode]:
         runnable_nodes = self._get_runnable_nodes()
 
+        admitted_bytes = self._pending_admitted_cpu_bytes()
+
         for node in runnable_nodes:
-            ok, gpu_id = self._check_memory_constraints(node)
+            if node in self._ready_set:
+                continue
+            ok, gpu_id = self._check_memory_constraints(node, admitted_bytes)
             if ok:
+                admitted_bytes += self.node_resources[node][0]
                 self.mapping[node].gpu_id = gpu_id
-                self._reserve_memory(node, gpu_id)
+                self._candidates.discard(node)
                 self.ready_nodes.append(node)
+                self._ready_set.add(node)
         contains_leaf = []
         for node in self.ready_nodes:
+            if isinstance(node, list):
+                continue
             if any(re.fullmatch(r"leaf_\d+", i) for i in self.mapping[node].inputs):
                 for mod in self.modalities:
                     if (
@@ -115,16 +167,7 @@
         return self.ready_nodes
 
     def _get_runnable_nodes(self) -> List[str]:
-        runnable_nodes = []
-        for node in self.topo_order:
-            if (
-                node not in self.leaves
-                and self.unresolved_parents[node] == 0
-                and node not in self.running_nodes
-                and node not in self.completed_nodes
-                and node not in self.ready_nodes
-            ):
-                runnable_nodes.append(node)
+        runnable_nodes = list(self._candidates)
 
         def _score(node_id: str):
             release_bytes = 0
@@ -134,32 +177,54 @@
                     and self.remaining_children.get(parent_id, 0) == 1
                 ):
                     release_bytes += self.node_resources[parent_id][0]
-            return (-release_bytes, node_id not in self.roots)
+            return (-release_bytes, node_id not in self.roots, node_id)
 
         runnable_nodes.sort(key=_score)
         return runnable_nodes
 
-    def add_failed_node(self, node_id: str):
+    def add_failed_node(self, node_id: str, reason: str = "unknown failure"):
         self.failed_nodes.append(node_id)
-        self.running_nodes.remove(node_id)
+        self.failed_node_reasons[node_id] = reason
+        self.running_nodes.discard(node_id)
+        self._release_execution_memory(node_id, self.mapping[node_id].gpu_id)
 
-        self._release_memory(node_id, self.mapping[node_id].gpu_id)
+    def requeue_node(self, node_id: str) -> None:
+        self.running_nodes.discard(node_id)
+        self._release_execution_memory(node_id, self.mapping[node_id].gpu_id)
+        self._candidates.add(node_id)
+
+    def begin_execution(self, node_id: str) -> None:
+        gpu_id = self.mapping[node_id].gpu_id
+        cpu_mem, gpu_mem = self.node_resources[node_id]
+        if gpu_id is not None and gpu_mem > 0:
+            self.memory_stats["gpu_in_use"][gpu_id] += gpu_mem
+        if cpu_mem > 0 and node_id not in self._cpu_reserved_nodes:
+            self._cpu_reserved_nodes[node_id] = int(cpu_mem)
+            self.memory_stats["cpu_in_flight"] += int(cpu_mem)
 
     def move_to_running(self, node_id: str | list):
         self.ready_nodes.remove(node_id)
         if isinstance(node_id, list):
-            self.running_nodes.extend(node_id)
+            self._ready_set.difference_update(node_id)
+            self.running_nodes.update(node_id)
         else:
-            self.running_nodes.append(node_id)
+            self._ready_set.discard(node_id)
+            self.running_nodes.add(node_id)
 
     def complete_node(self, node_id: str):
-        self.running_nodes.remove(node_id)
+        self.running_nodes.discard(node_id)
         self.completed_nodes.append(node_id)
-        self._release_memory(node_id, self.mapping[node_id].gpu_id)
-        self.topo_order.remove(node_id)
+        self._completed_set.add(node_id)
+        self._release_execution_memory(node_id, self.mapping[node_id].gpu_id)
         for child_id in self.children[node_id]:
             self.parent_refcounts[child_id] -= 1
             self.unresolved_parents[child_id] -= 1
+            if (
+                self.unresolved_parents[child_id] == 0
+                and child_id not in self.leaves
+                and child_id not in self._completed_set
+            ):
+                self._candidates.add(child_id)
 
         for parent_id in self.parents.get(node_id, set()):
             if self.remaining_children.get(parent_id, 0) > 0:
@@ -209,9 +274,9 @@
         for desc_id in descendants:
             if (
                 desc_id in self.leaves
-                or desc_id in self.ready_nodes
+                or desc_id in self._ready_set
                 or desc_id in self.running_nodes
-                or desc_id in self.completed_nodes
+                or desc_id in self._completed_set
             ):
                 continue
 
@@ -219,9 +284,10 @@
             if not parent_ids:
                 continue
 
-            input_stats = self.node_stats.get(parent_ids[0])
-            if input_stats is None:
+            parent_stats = [self.node_stats.get(pid) for pid in parent_ids]
+            if any(s is None for s in parent_stats):
                 continue
+            input_stats = parent_stats[0] if len(parent_stats) == 1 else parent_stats
 
             if desc_id not in self.roots:
                 operation = self.mapping[desc_id].operation(
@@ -232,7 +298,9 @@
                 peak_memory["cpu_peak_bytes"] += (
                     64 * 1024 + 512 * input_stats_for_overhead.num_instances
                 )
-                output_stats = operation.get_output_stats(input_stats)
+                output_stats = self._resolve_output_stats(
+                    operation.get_output_stats(input_stats), input_stats
+                )
 
                 self.node_resources[desc_id] = (
                     int(peak_memory["cpu_peak_bytes"]),
@@ -288,29 +356,65 @@
         return blocked
 
     def not_enough_memory(self) -> bool:
-        for node_id in self._get_pending_nodes():
+        if self.running_nodes or self.ready_nodes:
+            return False
+        for node_id in self._candidates:
             cpu_mem, gpu_mem = self.node_resources[node_id]
-            if cpu_mem > self.memory_budget["cpu"] - self.memory_stats["cpu_in_use"]:
+            if cpu_mem > self.memory_budget["cpu"] - self._total_cpu_in_use():
+                self.deadlock_reason = (
+                    f"node {node_id} needs {cpu_mem / 1024**3:.2f} GB CPU but only "
+                    f"{(self.memory_budget['cpu'] - self._total_cpu_in_use()) / 1024**3:.2f} "
+                    f"GB is free and nothing is running to release more"
+                )
                 return True
             if gpu_mem > 0.0 and self.n_gpu > 0:
                 gpu_id = self._gpu_with_most_free_memory(gpu_mem)
                 if gpu_id is None:
+                    if (
+                        self.gpu_wait_attempts.get(node_id, 0)
+                        <= _MAX_GPU_SCHEDULE_ATTEMPTS
+                    ):
+                        continue
+                    self.deadlock_reason = (
+                        f"node {node_id} needs {gpu_mem / 1024**3:.2f} GB GPU memory "
+                        f"but no GPU has enough free and nothing is running to release more"
+                    )
                     return True
-        return self.memory_stats["cpu_in_use"] > self.memory_budget["cpu"]
+        return False
 
-    def _check_memory_constraints(self, node_id: str) -> bool:
+    def _check_memory_constraints(self, node_id: str, pending_bytes: int = 0) -> bool:
         cpu_mem, gpu_mem = self.node_resources[node_id]
         gpu_id = None
-        if cpu_mem > self.memory_budget["cpu"] - self.memory_stats["cpu_in_use"]:
+        if (
+            cpu_mem
+            > self.memory_budget["cpu"] - self._total_cpu_in_use() - pending_bytes
+        ):
             if cpu_mem > self.memory_budget["cpu"]:
                 self.blocked_memory_nodes_perm.append(node_id)
-                self.topo_order.remove(node_id)
+                self.blocked_memory_reasons[node_id] = (
+                    f"estimated CPU peak {cpu_mem / 1024**3:.2f} GB exceeds the "
+                    f"total CPU memory budget of "
+                    f"{self.memory_budget['cpu'] / 1024**3:.2f} GB"
+                )
+                self._candidates.discard(node_id)
             return False, None
 
         if gpu_mem > 0.0 and self.n_gpu > 0:
             gpu_id = self._gpu_with_most_free_memory(gpu_mem)
 
             if gpu_id is None:
+                attempts = self.gpu_wait_attempts.get(node_id, 0) + 1
+                self.gpu_wait_attempts[node_id] = attempts
+                if attempts > _MAX_GPU_SCHEDULE_ATTEMPTS:
+                    reason = (
+                        f"no GPU had {gpu_mem / 1024**3:.2f} GB free after "
+                        f"{attempts} scheduling attempts; running on CPU instead"
+                    )
+                    if DEBUG:
+                        print(f"Node {node_id}: {reason}")
+                    self.cpu_fallback_nodes.append(node_id)
+                    self.cpu_fallback_reasons[node_id] = reason
+                    return True, None
                 if DEBUG:
                     print(f"Node {node_id} has no available GPU")
                 return False, None
@@ -330,23 +434,17 @@
         return free_memory.index(max(free_memory))
 
     def _get_pending_nodes(self) -> List[str]:
-        return [
-            node_id
-            for node_id in self.topo_order
-            if node_id not in self.leaves and self.unresolved_parents[node_id] == 0
-        ]
+        return list(self._candidates)
 
-    def _reserve_memory(self, node_id: str, gpu_id: int) -> bool:
-        cpu_mem, gpu_mem = self.node_resources[node_id]
-        self.memory_stats["cpu_in_use"] += cpu_mem
-        if gpu_id is not None:
-            self.memory_stats["gpu_in_use"][gpu_id] += gpu_mem
-
-    def _release_memory(self, node_id: str, gpu_id: int) -> bool:
-        cpu_mem, gpu_mem = self.node_resources[node_id]
-        self.memory_stats["cpu_in_use"] -= cpu_mem
-        if gpu_id is not None:
+    def _release_execution_memory(self, node_id: str, gpu_id: int) -> None:
+        _, gpu_mem = self.node_resources[node_id]
+        if gpu_id is not None and gpu_mem > 0:
             self.memory_stats["gpu_in_use"][gpu_id] -= gpu_mem
+        reserved = self._cpu_reserved_nodes.pop(node_id, 0)
+        if reserved:
+            self.memory_stats["cpu_in_flight"] = max(
+                0, self.memory_stats["cpu_in_flight"] - reserved
+            )
 
     def _get_nodes_from_dags(
         self, dags: List[RepresentationDag]
@@ -430,7 +528,9 @@
                         64 * 1024 + 512 * input_stats_for_overhead.num_instances
                     )  # Placeholder for transformed modality creation overhead
                     peak_memory["cpu_peak_bytes"] *= 1
-                    output_stats = operation.get_output_stats(input_stats)
+                    output_stats = self._resolve_output_stats(
+                        operation.get_output_stats(input_stats), input_stats
+                    )
                     node_resources[node] = (
                         int(peak_memory["cpu_peak_bytes"]),
                         int(peak_memory["gpu_peak_bytes"]),
@@ -461,9 +561,11 @@
         return input_stats
 
     @staticmethod
-    def _stats_to_bytes(stats: Optional[Any], dtype_size: int = 4) -> int:
+    def _stats_to_bytes(stats: Optional[Any], dtype_size: Optional[int] = None) -> int:
         if stats is None:
             return 0
+        if dtype_size is None:
+            dtype_size = stats_itemsize(stats)
         num_instances = int(getattr(stats, "num_instances", 0))
         output_shape = tuple(getattr(stats, "output_shape", ()))
         numel = 1
@@ -473,3 +575,21 @@
             except Exception:
                 numel *= 1
         return max(0, int(num_instances * numel * dtype_size))
+
+    @staticmethod
+    def _resolve_output_stats(output_stats: Any, input_stats: Any) -> Any:
+        if output_stats is None or getattr(output_stats, "dtype", None) is not None:
+            return output_stats
+
+        sources = input_stats if isinstance(input_stats, list) else [input_stats]
+        sources = [s for s in sources if s is not None]
+        sources = [s for s in sources if stats_itemsize(s) > 0]
+        if not sources:
+            return output_stats
+
+        widest = max(sources, key=lambda s: stats_itemsize(s))
+        try:
+            output_stats.dtype = stats_dtype(widest)
+        except AttributeError:
+            pass
+        return output_stats
diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
index a67cbe1..a632c97 100644
--- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
+++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py
@@ -202,45 +202,41 @@
         if n_workers is None:
             n_workers = min(len(self.modalities), mp.cpu_count())
 
-        with mp.Manager() as manager:
+        ctx = mp.get_context("spawn")
+        with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx) as executor:
+            future_to_modality = {
+                executor.submit(
+                    self._process_modality,
+                    modality,
+                    self._checkpoint_manager.skip_remaining_by_key.get(
+                        modality.modality_id, 0
+                    )
+                    / len(self.tasks),
+                    scheduler=None,
+                ): modality
+                for modality in self.modalities
+            }
 
-            ctx = mp.get_context("spawn")
-            with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx) as executor:
-                future_to_modality = {
-                    executor.submit(
-                        self._process_modality,
-                        modality,
-                        self._checkpoint_manager.skip_remaining_by_key.get(
-                            modality.modality_id, 0
-                        )
-                        / len(self.tasks),
-                        scheduler=None,
-                    ): modality
-                    for modality in self.modalities
-                }
+            for future in as_completed(future_to_modality):
+                modality = future_to_modality[future]
+                try:
+                    results = future.result()
+                    self._merge_results(results)
+                    new_count = self._count_results(results.results)
+                    self._checkpoint_manager.increment(modality.modality_id, new_count)
+                    self._checkpoint_manager.checkpoint_if_due(
+                        self.operator_performance.results,
+                    )
+                except Exception as e:
+                    print(f"Error processing modality {modality.modality_id}: {e}")
+                    import traceback
 
-                for future in as_completed(future_to_modality):
-                    modality = future_to_modality[future]
-                    try:
-                        results = future.result()
-                        self._merge_results(results)
-                        new_count = self._count_results(results.results)
-                        self._checkpoint_manager.increment(
-                            modality.modality_id, new_count
-                        )
-                        self._checkpoint_manager.checkpoint_if_due(
-                            self.operator_performance.results,
-                        )
-                    except Exception as e:
-                        print(f"Error processing modality {modality.modality_id}: {e}")
-                        import traceback
-
-                        traceback.print_exc()
-                        self._checkpoint_manager.save_checkpoint(
-                            self.operator_performance.results,
-                            {},
-                        )
-                        continue
+                    traceback.print_exc()
+                    self._checkpoint_manager.save_checkpoint(
+                        self.operator_performance.results,
+                        {},
+                    )
+                    continue
 
     def optimize(self):
         if self.resume:
@@ -265,10 +261,11 @@
                 )
                 self._merge_results(local_result)
                 new_count = self._count_results(local_result.results)
-                self._checkpoint_manager.increment(modality.modality_id, new_count)
-                self._checkpoint_manager.checkpoint_if_due(
-                    self.operator_performance.results
-                )
+                if self.enable_checkpointing:
+                    self._checkpoint_manager.increment(modality.modality_id, new_count)
+                    self._checkpoint_manager.checkpoint_if_due(
+                        self.operator_performance.results
+                    )
                 if self.save_all_results:
                     self.store_results(f"{modality.modality_id}_unimodal_results.pkl")
             except Exception as e:
@@ -276,9 +273,10 @@
                 import traceback
 
                 traceback.print_exc()
-                self._checkpoint_manager.save_checkpoint(
-                    self.operator_performance.results, {}
-                )
+                if self.enable_checkpointing:
+                    self._checkpoint_manager.save_checkpoint(
+                        self.operator_performance.results, {}
+                    )
                 raise
         return execution_time
 
@@ -338,9 +336,8 @@
             expanded_dags_with_task_roots,
             [modality],
             self.tasks,
-            self._checkpoint_manager,
-            self.max_num_workers,
-            self.result_path,
+            max_num_workers=self.max_num_workers,
+            result_path=self.result_path,
             enable_checkpointing=self.enable_checkpointing,
         )
         start_time = time.perf_counter()
@@ -617,7 +614,7 @@
             for context_operator in context_operators:
                 for window_size, num_window in zip(window_lengths, num_windows):
                     context_operator_instance = context_operator(agg())
-                    if hasattr(context_operator, "num_windows"):
+                    if hasattr(context_operator_instance, "num_windows"):
                         context_operator_instance.num_windows = num_window
                     elif hasattr(context_operator_instance, "window_size"):
                         context_operator_instance.window_size = window_size
diff --git a/src/main/python/systemds/scuro/drsearch/worker_pool.py b/src/main/python/systemds/scuro/drsearch/worker_pool.py
new file mode 100644
index 0000000..7e78862
--- /dev/null
+++ b/src/main/python/systemds/scuro/drsearch/worker_pool.py
@@ -0,0 +1,277 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+import itertools
+import multiprocessing as mp
+import multiprocessing.connection as mp_connection
+import os
+import signal
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+
+from systemds.scuro.utils.memory_utility import is_cuda_oom
+
+_THREAD_ENV_VARS = (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+    "BLIS_NUM_THREADS",
+)
+
+
+def _resolve_thread_count(num_threads: int) -> int:
+    explicit = os.environ.get("OMP_NUM_THREADS")
+    if explicit:
+        try:
+            num_threads = int(explicit)
+        except ValueError:
+            pass
+    return max(1, int(num_threads))
+
+
+def set_thread_env_before_spawn(num_threads: int) -> None:
+    num_threads = _resolve_thread_count(num_threads)
+    for var in _THREAD_ENV_VARS:
+        os.environ[var] = str(num_threads)
+
+
+def _worker_initializer(num_threads: int) -> None:
+    num_threads = _resolve_thread_count(num_threads)
+    for var in _THREAD_ENV_VARS:
+        os.environ[var] = str(num_threads)
+    try:
+        torch.set_num_threads(num_threads)
+    except Exception:
+        pass
+
+
+@dataclass
+class _Job:
+    job_id: int
+    kind: str
+    payload: tuple
+    gpu_id: Optional[int] = None
+
+
+@dataclass
+class _JobResult:
+    job_id: int
+    ok: bool
+    pid: Optional[int]
+    value: Any = None
+    error: Optional[str] = None
+    cuda_oom: bool = False
+    worker_died: bool = False
+
+
+def _worker_main(
+    job_q, result_q, dispatch: Dict[str, Callable], num_threads: int
+) -> None:
+    _worker_initializer(num_threads)
+    while True:
+        job = job_q.get()
+        if job is None:
+            return
+        try:
+            value = dispatch[job.kind](job.payload, job.gpu_id)
+            result_q.put(_JobResult(job.job_id, True, os.getpid(), value=value))
+        except Exception as e:
+            result_q.put(
+                _JobResult(
+                    job.job_id,
+                    False,
+                    os.getpid(),
+                    error=f"{type(e).__name__}: {e}",
+                    cuda_oom=is_cuda_oom(e),
+                )
+            )
+
+
+def _describe_worker_death(exitcode: Optional[int]) -> str:
+    if exitcode is None:
+        return "exit code unknown"
+    if exitcode < 0:
+        try:
+            sig = signal.Signals(-exitcode)
+        except ValueError:
+            return f"killed by signal {-exitcode}"
+        hint = {
+            signal.SIGKILL: " (often the OOM killer or an explicit kill -9)",
+            signal.SIGSEGV: " (segmentation fault, often a native library crash, e.g. CUDA/BLAS)",
+            signal.SIGABRT: " (abort, often a C-level assertion or CUDA error)",
+            signal.SIGBUS: " (bus error, often a full /dev/shm or a shared-memory issue)",
+        }.get(sig, "")
+        return f"killed by signal {sig.name} ({-exitcode}){hint}"
+    return f"exited with status {exitcode}"
+
+
+def create_mp_context():
+    ctx_name = os.environ.get("SCURO_MP_CONTEXT", "spawn")
+    try:
+        return mp.get_context(ctx_name)
+    except ValueError:
+        return mp.get_context("spawn")
+
+
+class PersistentWorkerPool:
+    def __init__(
+        self,
+        n_workers: int,
+        dispatch: Dict[str, Callable],
+        ctx=None,
+        threads_per_worker: int = 1,
+    ):
+        self._ctx = ctx or create_mp_context()
+        self._dispatch = dispatch
+        self._threads_per_worker = max(1, int(threads_per_worker))
+        self._result_q = self._ctx.Queue()
+        self._job_counter = itertools.count()
+        self._workers: Dict[int, Dict[str, Any]] = {}
+        self._idle_pids: List[int] = []
+        self._running: Dict[int, tuple] = {}
+        for _ in range(max(1, n_workers)):
+            self._spawn_worker()
+
+    def _spawn_worker(self) -> None:
+        set_thread_env_before_spawn(self._threads_per_worker)
+        job_q = self._ctx.Queue()
+        p = self._ctx.Process(
+            target=_worker_main,
+            args=(job_q, self._result_q, self._dispatch, self._threads_per_worker),
+            daemon=True,
+        )
+        p.start()
+        self._workers[p.pid] = {"process": p, "job_q": job_q}
+        self._idle_pids.append(p.pid)
+
+    @property
+    def has_idle_worker(self) -> bool:
+        return len(self._idle_pids) > 0
+
+    @property
+    def num_in_flight(self) -> int:
+        return len(self._running)
+
+    def submit(self, kind: str, payload: tuple, gpu_id: Optional[int] = None) -> int:
+        if not self._idle_pids:
+            raise RuntimeError("submit() called with no idle worker available")
+        job_id = next(self._job_counter)
+        job = _Job(job_id, kind, payload, gpu_id)
+        pid = self._idle_pids.pop()
+        self._running[job_id] = (pid, job)
+        self._workers[pid]["job_q"].put(job)
+        return job_id
+
+    def wait(self) -> _JobResult:
+        while True:
+            sentinel_to_pid = {
+                w["process"].sentinel: pid for pid, w in self._workers.items()
+            }
+            ready = mp_connection.wait(
+                [self._result_q._reader, *sentinel_to_pid.keys()]
+            )
+            if self._result_q._reader in ready:
+                jr = self._result_q.get()
+                entry = self._running.pop(jr.job_id, None)
+                if entry is not None:
+                    pid, _job = entry
+                    if pid in self._workers:
+                        self._idle_pids.append(pid)
+                return jr
+            for r in ready:
+                dead_pid = sentinel_to_pid.get(r)
+                if dead_pid is None:
+                    continue
+                result = self._replace_dead_worker(dead_pid)
+                if result is not None:
+                    return result
+
+    def _replace_dead_worker(self, pid: int) -> Optional[_JobResult]:
+        w = self._workers.pop(pid, None)
+        if w is None:
+            return None
+        try:
+            if pid in self._idle_pids:
+                self._idle_pids.remove(pid)
+        except ValueError:
+            pass
+        try:
+            w["process"].join(timeout=1)
+        except Exception:
+            pass
+        exitcode = w["process"].exitcode
+        try:
+            w["job_q"].close()
+            w["job_q"].join_thread()
+        except Exception:
+            pass
+
+        failed_job_id = None
+        for job_id, (running_pid, _job) in self._running.items():
+            if running_pid == pid:
+                failed_job_id = job_id
+                break
+        if failed_job_id is not None:
+            self._running.pop(failed_job_id, None)
+
+        self._spawn_worker()
+
+        if failed_job_id is None:
+            return None
+        return _JobResult(
+            failed_job_id,
+            False,
+            pid,
+            error=f"worker process died ({_describe_worker_death(exitcode)})",
+            worker_died=True,
+        )
+
+    def shutdown(self) -> None:
+        for w in self._workers.values():
+            try:
+                w["job_q"].put(None)
+            except Exception:
+                pass
+        for w in self._workers.values():
+            try:
+                w["process"].join(timeout=5)
+                if w["process"].is_alive():
+                    w["process"].kill()
+                    w["process"].join(timeout=2)
+            except Exception:
+                pass
+        for w in self._workers.values():
+            try:
+                w["job_q"].close()
+                w["job_q"].join_thread()
+            except Exception:
+                pass
+        try:
+            self._result_q.close()
+            self._result_q.join_thread()
+        except Exception:
+            pass
+        self._workers.clear()
+        self._idle_pids.clear()
+        self._running.clear()
diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py
index d83553e..c7b6d69 100644
--- a/src/main/python/systemds/scuro/representations/representation.py
+++ b/src/main/python/systemds/scuro/representations/representation.py
@@ -20,8 +20,18 @@
 # -------------------------------------------------------------
 import abc
 from dataclasses import dataclass
+from typing import Any, Optional
+
+import numpy as np
+
 from systemds.scuro.utils.identifier import Identifier
 
+CONTAINER_ARRAY = "ndarray"
+CONTAINER_LIST = "list_of_ndarray"
+CONTAINER_RAGGED = "ragged"
+NDARRAY_OBJECT_OVERHEAD_BYTES = 112
+DEFAULT_DTYPE = np.dtype(np.float32)
+
 
 @dataclass
 class RepresentationStats:
@@ -29,6 +39,117 @@
     output_shape: tuple
     output_shape_is_known: bool = True
     aggregate_dim: tuple = (0,)
+    dtype: Optional[Any] = None
+    container: str = CONTAINER_ARRAY
+    shape_variance: float = 0.0
+
+
+def stats_dtype(stats) -> np.dtype:
+    dtype = getattr(stats, "dtype", None)
+    if dtype is None:
+        return DEFAULT_DTYPE
+
+    if type(dtype).__module__ == "torch":
+        try:
+            resolved = np.dtype(str(dtype).rsplit(".", 1)[-1])
+            return resolved if resolved.itemsize else DEFAULT_DTYPE
+        except TypeError:
+            return DEFAULT_DTYPE
+    try:
+        resolved = np.dtype(dtype)
+    except TypeError:
+        return DEFAULT_DTYPE
+
+    if resolved.itemsize == 0:
+        return DEFAULT_DTYPE
+    return resolved
+
+
+def stats_itemsize(stats) -> int:
+    return int(stats_dtype(stats).itemsize)
+
+
+def stats_num_elements(stats) -> int:
+    n = 1
+    for dim in getattr(stats, "output_shape", ()) or ():
+        n *= int(dim)
+    return int(n)
+
+
+def stats_bytes(stats, quantile: float = 0.0) -> int:
+    num_instances = int(getattr(stats, "num_instances", 0) or 0)
+    per_instance = stats_num_elements(stats) * stats_itemsize(stats)
+    total = num_instances * per_instance
+
+    container = getattr(stats, "container", CONTAINER_ARRAY)
+    if container in (CONTAINER_LIST, CONTAINER_RAGGED):
+        total += num_instances * NDARRAY_OBJECT_OVERHEAD_BYTES
+
+    if quantile > 0.0:
+        variance = float(getattr(stats, "shape_variance", 0.0) or 0.0)
+        if variance > 0.0:
+            z = {0.9: 1.282, 0.95: 1.645, 0.99: 2.326}.get(quantile, 1.645)
+            total = int(total * (1.0 + z * variance))
+
+    return int(total)
+
+
+def infer_stats_from_data(data) -> Optional[RepresentationStats]:
+    if data is None:
+        return None
+
+    if isinstance(data, np.ndarray):
+        if data.ndim == 0:
+            return RepresentationStats(
+                1, (1,), output_shape_is_known=True, dtype=data.dtype
+            )
+        num_instances = int(data.shape[0])
+        output_shape = tuple(int(d) for d in data.shape[1:]) if data.ndim > 1 else (1,)
+        return RepresentationStats(
+            num_instances,
+            output_shape,
+            output_shape_is_known=True,
+            dtype=data.dtype,
+            container=CONTAINER_ARRAY,
+        )
+
+    if isinstance(data, list) and len(data) > 0 and isinstance(data[0], np.ndarray):
+        num_instances = len(data)
+        first_shape = tuple(int(d) for d in data[0].shape)
+        sizes = [x.size for x in data if isinstance(x, np.ndarray)]
+        same_shape = len(sizes) == num_instances and all(
+            x.shape == data[0].shape for x in data if isinstance(x, np.ndarray)
+        )
+        variance = 0.0
+        if not same_shape and sizes:
+            mean = sum(sizes) / len(sizes)
+            if mean > 0:
+                spread = (sum((s - mean) ** 2 for s in sizes) / len(sizes)) ** 0.5
+                variance = spread / mean
+        return RepresentationStats(
+            num_instances,
+            first_shape,
+            output_shape_is_known=bool(same_shape),
+            dtype=data[0].dtype,
+            container=CONTAINER_LIST if same_shape else CONTAINER_RAGGED,
+            shape_variance=variance,
+        )
+
+    return None
+
+
+def derive_stats(stats: RepresentationStats, **overrides) -> RepresentationStats:
+    fields = dict(
+        num_instances=stats.num_instances,
+        output_shape=stats.output_shape,
+        output_shape_is_known=getattr(stats, "output_shape_is_known", True),
+        aggregate_dim=getattr(stats, "aggregate_dim", (0,)),
+        dtype=getattr(stats, "dtype", None),
+        container=getattr(stats, "container", CONTAINER_ARRAY),
+        shape_variance=getattr(stats, "shape_variance", 0.0),
+    )
+    fields.update(overrides)
+    return RepresentationStats(**fields)
 
 
 class Representation:
diff --git a/src/main/python/systemds/scuro/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py
index 0d3cd9d..88698fa 100644
--- a/src/main/python/systemds/scuro/utils/memory_utility.py
+++ b/src/main/python/systemds/scuro/utils/memory_utility.py
@@ -18,13 +18,31 @@
 # under the License.
 #
 # -------------------------------------------------------------
+import os
 import resource
 import sys
+import threading
+import time
+import tracemalloc
+from dataclasses import dataclass
 import numpy as np
 from sympy import Dict
 import torch
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 import psutil
+import gc
+
+_CPU_MEMORY_BUDGET_FRACTION = float(
+    os.environ.get("SCURO_CPU_MEMORY_BUDGET_FRACTION", "0.5")
+)
+_CPU_MEMORY_BUDGET_GB = os.environ.get("SCURO_CPU_MEMORY_BUDGET_GB")
+
+
+def cpu_memory_budget_bytes() -> float:
+    budget = float(psutil.virtual_memory().available) * _CPU_MEMORY_BUDGET_FRACTION
+    if _CPU_MEMORY_BUDGET_GB:
+        budget = min(budget, float(_CPU_MEMORY_BUDGET_GB) * 1024**3)
+    return budget
 
 
 def get_model_size_mb(model: torch.nn.Module) -> float:
@@ -192,3 +210,98 @@
     metadata = getattr(modality, "metadata", None)
     metadata_bytes = estimate_numpy_like_bytes(metadata)
     return int(data_bytes + metadata_bytes)
+
+
+def is_cuda_oom(exc: BaseException) -> bool:
+    if isinstance(exc, torch.cuda.OutOfMemoryError):
+        return True
+    msg = str(exc).lower()
+    return (
+        "cuda out of memory" in msg
+        or "cudamalloc" in msg
+        or "cuda error: out of memory" in msg
+    )
+
+
+def cleanup_gpu(gpu_id: Optional[int]) -> None:
+    if gpu_id is None or not torch.cuda.is_available():
+        return
+    device = torch.device(f"cuda:{gpu_id}")
+    torch.cuda.set_device(device)
+    torch.cuda.synchronize(device)
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+@dataclass
+class MemoryMeasurement:
+    increment_bytes: int
+    footprint_bytes: int
+    input_resident_bytes: int
+    traced_peak_bytes: int
+    rss_delta_bytes: int
+    peak_abs_rss_bytes: int
+
+
+def merge_memory_measurements(
+    measurements: List[Optional[MemoryMeasurement]],
+) -> Optional[MemoryMeasurement]:
+    present = [m for m in measurements if m is not None]
+    if not present:
+        return None
+    return MemoryMeasurement(
+        increment_bytes=max(m.increment_bytes for m in present),
+        footprint_bytes=max(m.footprint_bytes for m in present),
+        input_resident_bytes=max(m.input_resident_bytes for m in present),
+        traced_peak_bytes=max(m.traced_peak_bytes for m in present),
+        rss_delta_bytes=max(m.rss_delta_bytes for m in present),
+        peak_abs_rss_bytes=max(m.peak_abs_rss_bytes for m in present),
+    )
+
+
+def measure_memory_during(
+    fn, *args, input_resident_bytes: int = 0, sample_s: float = 0.01, **kwargs
+):
+    proc = psutil.Process(os.getpid())
+    baseline_rss = proc.memory_info().rss
+    peak_rss = baseline_rss
+    stop = threading.Event()
+
+    def sampler():
+        nonlocal peak_rss
+        while not stop.is_set():
+            rss = proc.memory_info().rss
+            if rss > peak_rss:
+                peak_rss = rss
+            time.sleep(sample_s)
+
+    owns_tracing = not tracemalloc.is_tracing()
+    if owns_tracing:
+        tracemalloc.start()
+    else:
+        tracemalloc.reset_peak()
+    traced_baseline, _ = tracemalloc.get_traced_memory()
+
+    t = threading.Thread(target=sampler, daemon=True)
+    t.start()
+    try:
+        out = fn(*args, **kwargs)
+    finally:
+        stop.set()
+        t.join()
+        _, traced_peak = tracemalloc.get_traced_memory()
+        if owns_tracing:
+            tracemalloc.stop()
+
+    traced_increment = max(int(traced_peak) - int(traced_baseline), 0)
+    rss_delta = max(int(peak_rss) - int(baseline_rss), 0)
+    increment = max(traced_increment, rss_delta)
+
+    return out, MemoryMeasurement(
+        increment_bytes=increment,
+        footprint_bytes=increment + int(input_resident_bytes),
+        input_resident_bytes=int(input_resident_bytes),
+        traced_peak_bytes=traced_increment,
+        rss_delta_bytes=rss_delta,
+        peak_abs_rss_bytes=int(peak_rss),
+    )
diff --git a/src/main/python/tests/scuro/data_generator.py b/src/main/python/tests/scuro/data_generator.py
index 937fd62..b30946f 100644
--- a/src/main/python/tests/scuro/data_generator.py
+++ b/src/main/python/tests/scuro/data_generator.py
@@ -222,7 +222,9 @@
         data = [
             [
                 random.random()
-                for _ in range(random.randint(max_audio_length * 0.9, max_audio_length))
+                for _ in range(
+                    random.randint(int(max_audio_length * 0.9), max_audio_length)
+                )
             ]
             for _ in range(num_instances)
         ]