add memory measuring functionality at runtime

christinadionysio · christinadionysio · commit e282a657d28f · 2026-03-13T13:45:26.000+01:00
diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py
@@ -45,6 +45,35 @@
 from systemds.scuro.representations.unimodal import UnimodalRepresentation
 from systemds.scuro.utils.checkpointing import CheckpointManager
 import sys
+import threading
+import time
+import psutil
+import os
+
+
+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()
+    try:
+        out = fn(*args, **kwargs)
+    finally:
+        stop.set()
+        t.join()
+
+    return out, (peak - baseline), peak
 
 
 class RefCountResultCache:
@@ -86,60 +115,52 @@ def get_memory_total_memory_usage(self):
         return sum(self.memory_usage_per_node.values())
 
 
-def _execute_node_worker(
-    node: RepresentationNode,
-    input_mods: List[Any],
-    task: Any,
-    rep_cache: Optional[Dict[str, Any]],
-    gpu_id: Optional[int],
-):
-    proc = psutil.Process(os.getpid())
-    before = proc.memory_info().rss  # bytes
-
+def _execute_node_worker(node, input_mods, task, rep_cache, gpu_id):
     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)
 
-    result = None
     node_operation = node.operation(params=node.parameters)
     operation_name = node_operation.name
-    # print(
-    #     f"Executing node {node.node_id} inputs: {input_mods[0].modality_id}, gpu: {gpu_id}, operation: {operation_name}"
-    # )
+
     if gpu_id is not None and hasattr(node_operation, "gpu_id"):
         node_operation.gpu_id = gpu_id
 
-    if len(input_mods) == 1:
-        if isinstance(node_operation, Context):
-            result = input_mods[0].context(node_operation)
-        elif isinstance(node_operation, DimensionalityReduction):
-            result = input_mods[0].dimensionality_reduction(node_operation)
-        elif isinstance(node_operation, AggregatedRepresentation):
-            result = node_operation.transform(input_mods[0])
-        elif isinstance(node_operation, UnimodalRepresentation):
-            if rep_cache is not None and node_operation.name in rep_cache:
-                result = rep_cache[node_operation.name]
-            else:
-                result = input_mods[0].apply_representation(node_operation)
-        else:
-            result = input_mods[0].apply_representation(node_operation)
-    else:
-        fusion_op = node_operation
-        if hasattr(fusion_op, "needs_training") and fusion_op.needs_training:
-            result = input_mods[0].combine_with_training(
-                input_mods[1:], fusion_op, task
-            )
+    def _run_node_op():
+        if len(input_mods) == 1:
+            if isinstance(node_operation, Context):
+                return input_mods[0].context(node_operation)
+            elif isinstance(node_operation, DimensionalityReduction):
+                return input_mods[0].dimensionality_reduction(node_operation)
+            elif isinstance(node_operation, AggregatedRepresentation):
+                return node_operation.transform(input_mods[0])
+            elif isinstance(node_operation, UnimodalRepresentation):
+                if rep_cache is not None and node_operation.name in rep_cache:
+                    return rep_cache[node_operation.name]
+                return input_mods[0].apply_representation(node_operation)
+            return input_mods[0].apply_representation(node_operation)
         else:
-            result = input_mods[0].combine(input_mods[1:], fusion_op)
-    delta_bytes = proc.memory_info().rss - before
+            fusion_op = node_operation
+            if hasattr(fusion_op, "needs_training") and fusion_op.needs_training:
+                return input_mods[0].combine_with_training(
+                    input_mods[1:], fusion_op, task
+                )
+            return input_mods[0].combine(input_mods[1:], fusion_op)
+
+    result, peak_delta_bytes, peak_abs_rss = measure_peak_rss_during(
+        _run_node_op,
+        sample_s=0.01,
+    )
+
     gpu_peak_bytes = (
         torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0
     )
-    # print(f"Node {node.node_id}: {operation_name} has a CPU peak memory usage of {delta_bytes/1024**3:.2f} GB, and a GPU peak memory usage of {gpu_peak_bytes/1024**3:.2f} GB")
+
     return {
         "result": result,
-        "peak_bytes": delta_bytes,
+        "peak_bytes": peak_delta_bytes,  # per-call CPU peak over baseline
+        "peak_abs_rss_bytes": peak_abs_rss,  # optional
         "gpu_peak_bytes": gpu_peak_bytes,
         "operation_name": operation_name,
     }
@@ -148,8 +169,6 @@ def _execute_node_worker(
 def _execute_task_worker(
     task_node_id: str, task: Any, data: Any, gpu_id: Optional[int]
 ) -> Dict[str, Any]:
-    proc = psutil.Process(os.getpid())
-    before = proc.memory_info().rss  # bytes
 
     # print(f"Executing task {task_node_id} on GPU {gpu_id}")
     if gpu_id is not None:
@@ -159,18 +178,25 @@ def _execute_task_worker(
 
     if gpu_id is not None and hasattr(task, "model") and hasattr(task.model, "device"):
         task.model.device = torch.device(f"cuda:{gpu_id}")
-    start = time.perf_counter()
-    scores = task.run(data)
-    end = time.perf_counter()
-    delta_bytes = proc.memory_info().rss - before
+
+    def _run_task():
+        start = time.perf_counter()
+        scores = task.run(data)
+        end = time.perf_counter()
+        return scores, end - start
+
     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 {delta_bytes/1024**3:.2f} GB, and a GPU peak memory usage of {gpu_peak_bytes/1024**3:.2f} GB")
     return {
-        "scores": scores,
-        "task_time": end - start,
-        "peak_bytes": delta_bytes,
+        "scores": result[0],
+        "task_time": result[1],
+        "peak_bytes": peak_delta_bytes,
         "gpu_peak_bytes": gpu_peak_bytes,
     }
 
diff --git a/src/main/python/systemds/scuro/representations/aggregated_representation.py b/src/main/python/systemds/scuro/representations/aggregated_representation.py
@@ -49,35 +49,31 @@ def __init__(self, aggregation="mean", target_dimensions=None, params=None):
         self.data_type = np.float32
 
     def get_output_stats(self, input_stats: RepresentationStats) -> RepresentationStats:
-        if len(input_stats.output_shape) == 1 or len(input_stats.output_shape) == 2:
-            return RepresentationStats(
-                input_stats.num_instances, (input_stats.output_shape[0],)
-            )
-        elif len(input_stats.output_shape) == 3:
-            return RepresentationStats(
-                input_stats.num_instances,
-                (
-                    input_stats.output_shape[0],
-                    input_stats.output_shape[1],
-                ),
-            )
+        if len(input_stats.output_shape) == 0:
+            out_shape = (1,)
+        elif len(input_stats.output_shape) == 1:
+            out_shape = (1,)
         else:
-            raise ValueError(f"Invalid output shape: {input_stats.output_shape}")
+            out_shape = input_stats.output_shape[:-1]
+        return RepresentationStats(input_stats.num_instances, out_shape)
 
     def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
-        output_memory_bytes = 1
-        output_shape = self.get_output_stats(input_stats).output_shape
-        for dim in output_shape:
-            output_memory_bytes *= dim
-        return (
-            input_stats.num_instances
-            * output_memory_bytes
-            * np.dtype(self.data_type).itemsize
-        )
+        out_shape = self.get_output_stats(input_stats).output_shape
+        out_numel = int(np.prod(out_shape)) if len(out_shape) > 0 else 1
+        dtype_size = 8
+        return int(input_stats.num_instances * out_numel * dtype_size)
 
     def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
+        dtype_size = np.dtype(self.data_type).itemsize
+        in_shape = tuple(input_stats.output_shape)
+        in_numel = int(np.prod(in_shape)) if len(in_shape) > 0 else 1
+        input_bytes = int(input_stats.num_instances * in_numel * dtype_size)
+        output_bytes = self.estimate_output_memory_bytes(input_stats)
+        safety = 1.2
+        cpu_peak = input_bytes * 2 + output_bytes * 2
+
         return {
-            "cpu_peak_bytes": self.estimate_output_memory_bytes(input_stats),
+            "cpu_peak_bytes": int(cpu_peak * safety),
             "gpu_peak_bytes": 0,
         }
 
diff --git a/src/main/python/systemds/scuro/representations/bert.py b/src/main/python/systemds/scuro/representations/bert.py
@@ -98,12 +98,8 @@ def get_output_stats(self, input_stats) -> RepresentationStats:
             )
 
     def estimate_output_memory_bytes(self, input_stats):
-        return (
-            input_stats.num_instances
-            * self.max_seq_length
-            * 768
-            * self.data_type.itemsize
-        )
+        output_stats = self.get_output_stats(input_stats).output_shape
+        return int(input_stats.num_instances * np.prod(output_stats) * 8)
 
     def estimate_peak_memory_bytes(self, input_stats):
         model = AutoModel.from_pretrained(self.model_name)
@@ -112,7 +108,7 @@ def estimate_peak_memory_bytes(self, input_stats):
 
         output_bytes = self.estimate_output_memory_bytes(input_stats)
 
-        per_instance_input_bytes = self.max_seq_length * 3 * 8
+        per_instance_input_bytes = int(np.prod(input_stats.output_shape)) * 8
         input_bytes_all_instances = input_stats.num_instances * per_instance_input_bytes
 
         safety_margin_bytes = 64 * 1024 * 1024  # 64 MB
diff --git a/src/main/python/systemds/scuro/representations/bow.py b/src/main/python/systemds/scuro/representations/bow.py
@@ -52,10 +52,12 @@ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
         return RepresentationStats(input_stats.num_instances, (vocab_estimate,))
 
     def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+        output_bytes = 1
+        output_shape = self.get_output_stats(input_stats).output_shape
+        for dim in output_shape:
+            output_bytes *= dim
         return (
-            input_stats.num_instances
-            * self.get_output_stats(input_stats).output_shape[0]
-            * np.dtype(self.data_type).itemsize
+            input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
         )
 
     def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py
@@ -223,7 +223,12 @@ def gpu_id(self, gpu_id):
         self.device = get_device(gpu_id)
 
     def estimate_output_memory_bytes(self, input_stats) -> int:
-        return input_stats.num_instances * 512 * self.data_type.itemsize
+        output_stats = self.get_output_stats(input_stats)
+        output_bytes = 1
+        for dim in output_stats.output_shape:
+            output_bytes *= dim
+
+        return input_stats.num_instances * output_bytes * self.data_type.itemsize
 
     def get_output_stats(self, input_stats) -> RepresentationStats:
         if not isinstance(input_stats, RepresentationStats):
diff --git a/src/main/python/systemds/scuro/representations/glove.py b/src/main/python/systemds/scuro/representations/glove.py
@@ -62,10 +62,12 @@ def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
         return RepresentationStats(input_stats.num_instances, (self.embedding_dim,))
 
     def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
+        output_bytes = 1
+        output_shape = self.get_output_stats(input_stats).output_shape
+        for dim in output_shape:
+            output_bytes *= dim
         return (
-            input_stats.num_instances
-            * self.embedding_dim
-            * np.dtype(self.data_type).itemsize
+            input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
         )
 
     def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
diff --git a/src/main/python/systemds/scuro/representations/mlp_averaging.py b/src/main/python/systemds/scuro/representations/mlp_averaging.py
@@ -58,22 +58,73 @@ def __init__(self, output_dim=512, batch_size=32, params=None):
         self.batch_size = batch_size
         self.device = None
         self.data_type = np.float32
+        self.gpu_id = None
+
+    @property
+    def gpu_id(self):
+        return self._gpu_id
+
+    @gpu_id.setter
+    def gpu_id(self, gpu_id):
+        self._gpu_id = gpu_id
+        self.device = get_device(gpu_id)
 
     def get_output_stats(self, input_stats: RepresentationStats) -> RepresentationStats:
         return RepresentationStats(input_stats.num_instances, (self.output_dim,))
 
     def estimate_output_memory_bytes(self, input_stats: RepresentationStats) -> int:
+        output_bytes = 1
+        for dim in input_stats.output_shape:
+            output_bytes *= dim
         return (
-            input_stats.num_instances
-            * self.output_dim
-            * np.dtype(self.data_type).itemsize
+            input_stats.num_instances * output_bytes * np.dtype(self.data_type).itemsize
         )
 
     def estimate_peak_memory_bytes(self, input_stats: RepresentationStats) -> dict:
-        return {
-            "cpu_peak_bytes": self.estimate_output_memory_bytes(input_stats),
-            "gpu_peak_bytes": 0,
-        }
+        n = int(input_stats.num_instances)
+        input_dim = (
+            int(np.prod(input_stats.output_shape)) if input_stats.output_shape else 0
+        )
+        elem_size = np.dtype(self.data_type).itemsize
+
+        if input_dim < self.output_dim or n == 0 or input_dim == 0:
+            input_bytes = n * input_dim * elem_size
+            cpu_peak = int(input_bytes * 1.05 + 8 * 1024**2)  # small safety margin
+            return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}
+
+        out_dim = int(self.output_dim)
+        batch = int(max(1, min(self.batch_size, n)))
+
+        input_bytes = n * input_dim * elem_size
+        output_bytes = n * out_dim * elem_size
+        weight_bytes = out_dim * input_dim * elem_size
+
+        batch_input_bytes = batch * input_dim * elem_size
+        batch_output_bytes = batch * out_dim * elem_size
+
+        num_batches = (n + batch - 1) // batch
+        python_overhead = num_batches * 1024 
+
+        cpu_working = (
+            input_bytes
+            + 2 * output_bytes
+            + weight_bytes
+            + batch_input_bytes
+            + batch_output_bytes
+            + python_overhead
+        )
+        cpu_peak = int(
+            cpu_working * 1.20 + 64 * 1024**2
+        ) 
+
+        device_type = getattr(self.device, "type", "cpu")
+        if device_type == "cuda":
+            gpu_working = weight_bytes + batch_input_bytes + batch_output_bytes
+            gpu_peak = int(gpu_working * 1.35 + 64 * 1024**2)
+        else:
+            gpu_peak = 0
+
+        return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": gpu_peak}
 
     def execute(self, data):
         set_random_seeds(42)
@@ -86,20 +137,9 @@ def execute(self, data):
             return data
 
         dim_reduction_model = AggregationMLP(input_dim, self.output_dim)
-        self.device = get_device_for_model(dim_reduction_model, memory_factor=1.5)
         dim_reduction_model = dim_reduction_model.to(self.device)
         dim_reduction_model.eval()
 
-        # sample = data[0] if data else ""
-        # self.batch_size = compute_batch_size(
-        #     model=dim_reduction_model,
-        #     device=self.device,
-        #     sample_data=sample,
-        #     tokenizer=None,
-        #     max_seq_length=None,
-        #     max_batch_size=self.batch_size,
-        # )
-
         tensor_data = torch.from_numpy(data).float()
 
         dataset = TensorDataset(tensor_data)
diff --git a/src/main/python/systemds/scuro/representations/text_context_with_indices.py b/src/main/python/systemds/scuro/representations/text_context_with_indices.py
diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py
diff --git a/src/main/python/systemds/scuro/representations/word2vec.py b/src/main/python/systemds/scuro/representations/word2vec.py
