From bc9c5a8405cdaa2b0a05422c35b57627180e8a3c Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Thu, 13 Aug 2026 15:51:09 +0200 Subject: [PATCH 1/6] make node executor more efficient --- .../scuro/drsearch/modality_result_cache.py | 136 ++ .../scuro/drsearch/modality_shared_memory.py | 28 +- .../systemds/scuro/drsearch/node_executor.py | 1240 +++++++++-------- .../systemds/scuro/drsearch/node_scheduler.py | 224 ++- .../scuro/drsearch/unimodal_optimizer.py | 5 +- .../systemds/scuro/drsearch/worker_pool.py | 277 ++++ 6 files changed, 1271 insertions(+), 639 deletions(-) create mode 100644 src/main/python/systemds/scuro/drsearch/modality_result_cache.py create mode 100644 src/main/python/systemds/scuro/drsearch/worker_pool.py 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 00000000000..0848702cbe1 --- /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 d4092b90cfc..e68d7195a2d 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 __init__( 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 __init__( 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 __init__(self, shm_name: str, dtype_str: str, shape: tuple): 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 __init__( 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 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: 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 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: 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 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: 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 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: 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 a6a7ffe2ca4..8f177344c01 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): - resident_bytes = result.calculate_memory_usage() - shared_backing_bytes = 0 +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) - 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 _offload_to_shared_memory(result: Any): + if result is None or not hasattr(result, "data"): + return None, None, 0, None - 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 + actual_stats = _infer_actual_output_stats(result) + shm_name = None + resident_bytes = None + shm_bytes = 0 + try: + resident_bytes = result.calculate_memory_usage() + 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}") - 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) - - -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, - } + return shm_name, resident_bytes, int(shm_bytes or 0), actual_stats -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 @@ def _run_node_op(): 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 @@ def _execute_task_worker( 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 @@ def _run_task(): return scores, end - start gpu_peak_bytes = -1 - peak_delta_bytes = -1 + measurement = None if DEBUG: - 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, + result, measurement = measure_memory_during( + lambda: _run_gpu_op(_run_task, gpu_id), + input_resident_bytes=estimate_modality_bytes(modality), 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" + gpu_peak_bytes = ( + torch.cuda.max_memory_allocated(device) if gpu_id is not None else 0 ) 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 @@ def __init__( 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": {}} + + 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 _shm_names_for_submit( - self, parent_node_ids: List[str], payload_data: Any - ) -> List[str]: + 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)) - - 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) + 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 _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) - - def run(self) -> None: - task_results = {} - memory_usage_data = {} - - self._materialize_leaf_modalities_in_shared_memory() - - 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() + 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) + + self._job_units[job_id] = _NodeUnit(node_id) + self._job_retained_shm[job_id] = retained + + 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 + + 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) + + 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"] + ) - done, _ = wait( - set(future_to_node_id.keys()), return_when=FIRST_COMPLETED + 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,214 @@ def _handle_modality_result( 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) + 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, + 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}" + ) + 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(data) + 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,) - ) - return RepresentationStats( - num_instances, output_shape, output_shape_is_known=True - ) - - 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), + 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 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 209f4503860..1ca681e88ad 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 @@ def __init__( 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 @@ def get_runnable(self) -> List[RepresentationNode]: 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 @@ def _score(node_id: str): 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 @@ def update_node_stats_and_reestimate_descendants( 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 @@ def update_node_stats_and_reestimate_descendants( 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 @@ def update_node_stats_and_reestimate_descendants( 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 @@ def _is_deadlock(self) -> bool: 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 @@ def _gpu_with_most_free_memory(self, memory_needed): 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 - ] - - 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 + return list(self._candidates) - 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 @@ def _estimate_node_resources(self): 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 @@ def _stats_for_overhead(input_stats: Any) -> Any: 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 @@ def _stats_to_bytes(stats: Optional[Any], dtype_size: int = 4) -> int: 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 a67cbe12029..6527635ea36 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -338,9 +338,8 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): 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() 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 00000000000..7e78862a104 --- /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() From ffe1b01f5292b4145d929adcfb2f8fa3d56525ad Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Fri, 14 Aug 2026 19:02:51 +0200 Subject: [PATCH 2/6] fix tests --- .../scuro/drsearch/hyperparameter_tuner.py | 20 +-- .../systemds/scuro/drsearch/node_executor.py | 54 ++++---- .../scuro/drsearch/unimodal_optimizer.py | 92 +++++++------- .../systemds/scuro/utils/memory_utility.py | 115 +++++++++++++++++- src/main/python/tests/scuro/data_generator.py | 4 +- 5 files changed, 206 insertions(+), 79 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 9407bdf1250..772c5813b1c 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 @@ def __init__( 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 @@ def __init__( self.wandb_group = wandb_group self.wandb_tags = wandb_tags or [] self._wandb_run = None + self def get_modalities_by_id(self, modality_ids: List[int]) -> Modality: modalities = [] @@ -415,14 +417,18 @@ def tune_unimodal_representations(self, max_eval_per_rep: Optional[int] = None): ) ) 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/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 8f177344c01..ec5d9f40c36 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -737,7 +737,8 @@ def _checkpoint_memory_usage( 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 @@ -755,7 +756,7 @@ def _checkpoint_memory_usage( 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}" + 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: @@ -766,28 +767,35 @@ def _checkpoint_memory_usage( 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}" + f"{est_gpu/1024**3:.2f} GB, >200% of estimated for node {node_id}: {operation_name}" ) - 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) + 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, diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index 6527635ea36..a632c97e973 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 @@ def optimize_parallel(self, n_workers=None): 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 - } - - 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 + 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 + + traceback.print_exc() + self._checkpoint_manager.save_checkpoint( + self.operator_performance.results, + {}, + ) + continue def optimize(self): if self.resume: @@ -265,10 +261,11 @@ def optimize(self): ) 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 @@ def optimize(self): 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 @@ -616,7 +614,7 @@ def temporal_context_operators(self, modality, builder, leaf_id): 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/utils/memory_utility.py b/src/main/python/systemds/scuro/utils/memory_utility.py index 0d3cd9d786c..88698fa53cc 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 @@ def estimate_modality_bytes(modality) -> int: 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 937fd622d85..b30946fb7df 100644 --- a/src/main/python/tests/scuro/data_generator.py +++ b/src/main/python/tests/scuro/data_generator.py @@ -222,7 +222,9 @@ def create_audio_data(self, num_instances, max_audio_length): 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) ] From 7b74c5bb080649c093762a54fd7a5b1e53bef315 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Fri, 14 Aug 2026 19:10:29 +0200 Subject: [PATCH 3/6] add missing member to class --- src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index 772c5813b1c..a62b990fa99 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -334,7 +334,7 @@ def __init__( self.wandb_group = wandb_group self.wandb_tags = wandb_tags or [] self._wandb_run = None - self + self.enable_checkpointing = enable_checkpointing def get_modalities_by_id(self, modality_ids: List[int]) -> Modality: modalities = [] From 881210d3b4b00dd3f1c0cd989b25151f3e19b25d Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Fri, 14 Aug 2026 19:31:16 +0200 Subject: [PATCH 4/6] add missing file --- .../scuro/representations/representation.py | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/src/main/python/systemds/scuro/representations/representation.py b/src/main/python/systemds/scuro/representations/representation.py index d83553ec6e1..c7b6d69d730 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 @@ class RepresentationStats: 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: From 3783ce27243651912d0df856482387ae5f334bb5 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Fri, 14 Aug 2026 19:32:39 +0200 Subject: [PATCH 5/6] fix test assertion --- src/main/python/tests/scuro/test_data_loaders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py index 0d5ace01a47..3c31901bd7c 100644 --- a/src/main/python/tests/scuro/test_data_loaders.py +++ b/src/main/python/tests/scuro/test_data_loaders.py @@ -99,7 +99,7 @@ def test_video_loader_stats(self): self.assertIsInstance(stats, VideoStats) self.assertEqual(stats.num_instances, 2) - self.assertEqual(stats.max_length, 60) + self.assertEqual(stats.max_length, 45) self.assertEqual(stats.max_width, 160) self.assertEqual(stats.max_height, 120) self.assertEqual(stats.max_channels, 3) From 76ad5fb5aae384a00b8282f2b9d6907d16c2fffb Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Fri, 14 Aug 2026 19:44:56 +0200 Subject: [PATCH 6/6] revert test assertion (local error) --- src/main/python/tests/scuro/test_data_loaders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py index 3c31901bd7c..0d5ace01a47 100644 --- a/src/main/python/tests/scuro/test_data_loaders.py +++ b/src/main/python/tests/scuro/test_data_loaders.py @@ -99,7 +99,7 @@ def test_video_loader_stats(self): self.assertIsInstance(stats, VideoStats) self.assertEqual(stats.num_instances, 2) - self.assertEqual(stats.max_length, 45) + self.assertEqual(stats.max_length, 60) self.assertEqual(stats.max_width, 160) self.assertEqual(stats.max_height, 120) self.assertEqual(stats.max_channels, 3)