[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)
]
