[SYSTEMDS-3977] Add Parallel Execution of CV-Splits

src/main/python/systemds/scuro/drsearch/node_executor.py

@@ -511,13 +511,36 @@ def _dispatch_load_leaf(payload, _gpu_id):
     return _load_leaf_worker(modality, materialize)
 
 
+def _dispatch_split_dags(payload, gpu_id):
+    from systemds.scuro.drsearch.unimodal_optimizer import evaluate_fitted_dags
+
+    start_time = time.perf_counter()
+    allowed = os.sched_getaffinity(0)
+    if gpu_id is not None:
+        cores = sorted(allowed)
+        os.sched_setaffinity(0, {cores[os.getpid() % len(cores)]})
+    try:
+        results = evaluate_fitted_dags(*payload)
+    finally:
+        os.sched_setaffinity(0, allowed)
+        if gpu_id is not None:
+            cleanup_gpu(gpu_id)
+    return {
+        "results": results,
+        "start_time": start_time,
+        "end_time": time.perf_counter(),
+        "pid": os.getpid(),
+    }
+
+
 _WORKER_DISPATCH = {
     "node": _dispatch_node,
     "task": _dispatch_task,
     "task_batch": _dispatch_task_batch,
     "leaf_batch": _dispatch_leaf_batch,
     "load_leaf": _dispatch_load_leaf,
     "release_models": _dispatch_release_models,
+    "split_dags": _dispatch_split_dags,
 }
 
 
@@ -542,6 +565,12 @@ class _TaskBatchUnit(_BatchUnit):
     gpu_bytes: int = 0
 
 
+@dataclass
+class _SplitUnit:
+    index: int
+    gpu_id: Optional[int] = None
+
+
 @dataclass
 class ResultEntry:
     val_score: PerformanceMeasure = None
@@ -576,6 +605,7 @@ def __init__(
         enable_checkpointing: bool = False,
         worker_pool: Optional[PersistentWorkerPool] = None,
         search_start: Optional[float] = None,
+        split_jobs: Optional[List[tuple]] = None,
     ):
         self.enable_checkpointing = enable_checkpointing
         available_total_cpu = cpu_memory_budget_bytes(
@@ -633,6 +663,8 @@ def __init__(
         self._task_results: Dict[str, ResultEntry] = {}
         self._fused_shapes: Dict[tuple, tuple] = {}
         self._fused_ready_since: Dict[str, float] = {}
+        self._split_pending = list(enumerate(split_jobs or []))
+        self._split_results: List[Any] = [None] * len(self._split_pending)
 
         self._owns_pool = worker_pool is None
         if worker_pool is None:
@@ -973,9 +1005,30 @@ def _submit_fused_task_batches(self, node_ids: List[str]) -> None:
                 self._job_units[job_id] = _TaskBatchUnit(batch, gpu_id, gpu_bytes)
                 self._job_retained_shm[job_id] = retained
 
+    def _submit_split_jobs(self) -> None:
+        while self._split_pending and self._pool.has_idle_worker:
+            index, (payload, uses_gpu) = self._split_pending[0]
+            gpu_id = None
+            if uses_gpu and self.scheduler.gpu_devices:
+                gpu_id = self.scheduler.reserve_gpu(
+                    _FUSED_CONTEXT_BYTES,
+                    lambda g: self._pool.has_idle_worker_for(g),
+                    force=not self._job_units,
+                )
+                if gpu_id is None and self._job_units:
+                    return
+            if gpu_id is None and not self._pool.has_idle_worker_for(None, True):
+                return
+            self._split_pending.pop(0)
+            job_id = self._pool.submit(
+                "split_dags", payload, gpu_id=gpu_id, allow_gpu_worker_for_cpu=True
+            )
+            self._job_units[job_id] = _SplitUnit(index, gpu_id)
+
     def _fill_pipeline(self) -> None:
         if not self._pool.has_idle_worker:
             return
+        self._submit_split_jobs()
         ready = self.scheduler.get_runnable().copy()
         fused = self._fused_task_nodes(ready)
         if fused:
@@ -1052,6 +1105,21 @@ def _process_result(self, jr) -> None:
                 return
             if isinstance(unit, _TaskBatchUnit):
                 self.scheduler.release_gpu(unit.gpu_id, unit.gpu_bytes)
+            if isinstance(unit, _SplitUnit):
+                if unit.gpu_id is not None:
+                    self.scheduler.release_gpu(unit.gpu_id, _FUSED_CONTEXT_BYTES)
+                if jr.ok:
+                    value = jr.value
+                    self._record_stats(
+                        f"split_{unit.index}",
+                        value["pid"],
+                        value["start_time"],
+                        value["end_time"],
+                    )
+                    self._split_results[unit.index] = value["results"]
+                else:
+                    self._split_results[unit.index] = jr.error or "worker failure"
+                return
             if not jr.ok:
                 self._handle_job_failure(unit, jr)
                 return
@@ -1411,7 +1479,11 @@ def run(self) -> Dict[str, Any]:
         self._task_results = {}
         try:
             self._load_leaf_modalities()
-            while self._job_units or not self.scheduler.is_finished():
+            while (
+                self._job_units
+                or self._split_pending
+                or not self.scheduler.is_finished()
+            ):
                 self._fill_pipeline()
                 if not self._job_units:
                     continue
@@ -1436,6 +1508,7 @@ def run(self) -> Dict[str, Any]:
         return {
             "task_results": list(self._task_results.values()),
             "statistics": self.statistics,
+            "split_results": self._split_results,
         }
 
     def _reuse_report(self) -> Dict[str, Any]:

src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py

@@ -21,12 +21,14 @@
 import copy
 from collections import Counter
 import math
+import os
 import pickle
 import time
 from concurrent.futures import ProcessPoolExecutor, as_completed
 import multiprocessing as mp
 from typing import List, Any, Optional, Dict
 from functools import lru_cache
+from weakref import WeakValueDictionary
 
 from systemds.scuro.modality.type import ModalityType
 from systemds.scuro.drsearch.node_executor import (
@@ -40,7 +42,11 @@
     SplitAwareDAGEvaluator,
     dag_requires_fitting,
 )
-from systemds.scuro.representations.representation import RepresentationStats
+from systemds.scuro.drsearch.trial_result_cache import TrialNodeCache, TrialResultCache
+from systemds.scuro.representations.representation import (
+    FitRequirement,
+    RepresentationStats,
+)
 from systemds.scuro.drsearch.ranking import rank_by_tradeoff, rank_by_robustness
 from systemds.scuro.drsearch.task import PerformanceMeasure
 from systemds.scuro.representations.concatenation import Concatenation
@@ -51,7 +57,7 @@
     AggregatedRepresentation,
 )
 from systemds.scuro.modality.modality import Modality
-from systemds.scuro.drsearch.algebra import AlgebraConfig
+from systemds.scuro.drsearch.algebra import AlgebraConfig, canonical_signature
 from systemds.scuro.drsearch.operator_registry import Registry
 from systemds.scuro.representations.window_aggregation import WindowAggregation
 from systemds.scuro.representations.lstm import LSTM
@@ -65,6 +71,69 @@
 from bisect import bisect_left
 
 
+def _fit_requirement(node):
+    return getattr(node.operation, "fit_requirement", FitRequirement.NONE)
+
+
+def _lowest_fitted_key(dag: RepresentationDag) -> str:
+    keys = set()
+
+    def visit(node_id):
+        node = dag.get_node_by_id(node_id)
+        below = [visit(input_id) for input_id in node.inputs]
+        fitted = _fit_requirement(node) is not FitRequirement.NONE
+        if fitted and not any(below):
+            keys.add(repr(canonical_signature(dag, frozenset(), False, node_id)))
+        return fitted or any(below)
+
+    visit(dag.root_node_id)
+    return repr(sorted(keys))
+
+
+def _task_uses_gpu(task) -> bool:
+    uses_gpu = getattr(task.model, "uses_gpu", None)
+    if uses_gpu is not None:
+        return bool(uses_gpu)
+    device = getattr(task.model, "device", None)
+    return str(getattr(device, "type", device) or "").startswith("cuda")
+
+
+def evaluate_fitted_dags(dags, modality, tasks) -> list:
+    budget_gb = float(os.environ.get("SCURO_FITTED_DAG_CACHE_GB", "4"))
+    store = TrialResultCache(int(budget_gb * 1024**3), min_compute_seconds=0)
+    results = []
+    for dag in dags:
+        eval_dag, eval_modality, encode_time = UnimodalOptimizer._encode_below_lstm(
+            dag, modality
+        )
+        supervised = any(
+            _fit_requirement(node) is FitRequirement.SUPERVISED
+            for node in eval_dag.nodes
+        )
+        for task_idx, task in enumerate(tasks):
+            owner = task_idx if supervised else None
+            signatures = {
+                node.node_id: (
+                    owner,
+                    canonical_signature(eval_dag, frozenset(), False, node.node_id),
+                )
+                for node in eval_dag.nodes
+                if node.inputs
+            }
+            saved = store.saved_seconds
+            evaluation = SplitAwareDAGEvaluator(
+                eval_dag,
+                eval_modality,
+                task,
+                node_cache=TrialNodeCache(store, signatures),
+                share_fitted=True,
+            ).evaluate()
+            evaluation.representation_time += encode_time + store.saved_seconds - saved
+            timing = dict(getattr(task, "last_run_timing", {}))
+            results.append((dag.dag_id, task_idx, evaluation, timing))
+    return results
+
+
 class UnimodalOptimizer:
     def __init__(
         self,
@@ -386,10 +455,19 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None):
 
         task_results = []
         statistics = {"worker_stats": {}, "node_stats": {}, "reuse": {}}
+        split_jobs = []
         if fitted_dags:
             SplitAwareDAGEvaluator.materialize(modality)
+            groups = {}
+            for dag in fitted_dags:
+                groups.setdefault(_lowest_fitted_key(dag), []).append(dag)
+            uses_gpu = any(_task_uses_gpu(task) for task in self.tasks)
+            split_jobs = [
+                ((group, modality, self.tasks), uses_gpu) for group in groups.values()
+            ]
 
-        if stateless_task_dags:
+        split_results = []
+        if stateless_task_dags or split_jobs:
             node_executor = NodeExecutor(
                 stateless_task_dags,
                 [modality],
@@ -399,12 +477,14 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None):
                 enable_checkpointing=self.enable_checkpointing,
                 worker_pool=self._persistent_worker_pool,
                 search_start=self._search_start,
+                split_jobs=split_jobs,
             )
             if self._reuse_worker_pool and self._persistent_worker_pool is None:
                 self._persistent_worker_pool = node_executor.detach_worker_pool()
             exec_out = node_executor.run()
             task_results.extend(exec_out["task_results"])
             statistics = exec_out["statistics"]
+            split_results = exec_out["split_results"]
             reuse = statistics.get("reuse", {})
             planned = reuse.get("planned_task_nodes", 0)
             completed = reuse.get("completed_task_nodes", planned)
@@ -419,15 +499,23 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None):
 
         if fitted_dags:
             fitted_task_dags = self._expand_dags_with_task_roots(fitted_dags)
+            evaluations = {}
+            for (payload, _), result in zip(split_jobs, split_results):
+                if isinstance(result, list):
+                    for dag_id, task_idx, evaluation, timing in result:
+                        evaluations[dag_id, task_idx] = evaluation, timing
+                else:
+                    print(
+                        f"[unimodal] WARNING: modality {modality.modality_id} "
+                        f"fitted DAGs {[dag.dag_id for dag in payload[0]]} "
+                        f"failed: {result}",
+                        flush=True,
+                    )
             for dag in fitted_dags:
-                eval_dag, eval_modality, encode_time = self._encode_below_lstm(
-                    dag, modality
-                )
                 for task_idx, task in enumerate(self.tasks):
-                    evaluation = SplitAwareDAGEvaluator(
-                        eval_dag, eval_modality, task
-                    ).evaluate()
-                    evaluation.representation_time += encode_time
+                    if (dag.dag_id, task_idx) not in evaluations:
+                        continue
+                    evaluation, timing = evaluations[dag.dag_id, task_idx]
                     result_dag = next(
                         task_dag
                         for task_dag in fitted_task_dags
@@ -445,7 +533,7 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None):
                             train_score_stds=evaluation.scores[0].score_stds(),
                             val_score_stds=evaluation.scores[1].score_stds(),
                             test_score_stds=evaluation.scores[2].score_stds(),
-                            task_timing=getattr(task, "last_run_timing", {}),
+                            task_timing=timing,
                             representation_time=evaluation.representation_time,
                             task_time=evaluation.task_time,
                             dag=result_dag,
@@ -1212,15 +1300,9 @@ def get_k_best_results(
                 try:
                     for result in results:
                         if dag_requires_fitting(result.dag):
-                            SplitAwareDAGEvaluator.materialize(modality)
-                            transformed = SplitAwareDAGEvaluator(
-                                result.dag, modality, task
-                            ).fit_transform(
-                                task.train_indices,
-                                range(len(modality.data)),
+                            cache.append(
+                                self._fit_k_best_leaf(result.dag, modality, task)
                             )
-                            transformed.modality.fold_source = (result.dag, modality)
-                            cache.append(transformed.modality)
                         else:
                             cache.append(result.dag.execute([modality]))
                 finally:
@@ -1252,6 +1334,31 @@ def get_k_best_results(
 
         return results, cache
 
+    def _fit_k_best_leaf(self, dag, modality, task):
+        supervised = any(
+            _fit_requirement(node) is FitRequirement.SUPERVISED for node in dag.nodes
+        )
+        key = (
+            repr(canonical_signature(dag, frozenset(), False)),
+            str(modality.modality_id),
+            tuple(task.train_indices),
+            task.model.name if supervised else None,
+        )
+        leaves = self.__dict__.setdefault("_fitted_leaves", WeakValueDictionary())
+        leaf = leaves.get(key)
+        if leaf is None:
+            SplitAwareDAGEvaluator.materialize(modality)
+            leaf = (
+                SplitAwareDAGEvaluator(dag, modality, task)
+                .fit_transform(task.train_indices, range(len(modality.data)))
+                .modality
+            )
+            leaves[key] = leaf
+        else:
+            leaf = copy.copy(leaf)
+        leaf.fold_source = (dag, modality)
+        return leaf
+
     @staticmethod
     def _decode_lazy_leaf_once(modality, results):
         data = modality.data if modality.has_data() else None

src/main/python/systemds/scuro/representations/lstm.py

@@ -23,7 +23,11 @@
 from torch import nn
 from torch.utils.data import DataLoader, TensorDataset
 from typing import List, Dict, Any, Optional
-from systemds.scuro.utils.static_variables import get_device
+from systemds.scuro.utils.static_variables import (
+    get_device,
+    get_global_seed,
+    get_train_seed,
+)
 from systemds.scuro.utils.utils import set_random_seeds
 import numpy as np
 
@@ -93,7 +97,7 @@ def __init__(
         self.model_state = None
         self.is_multilabel = False
 
-        set_random_seeds()
+        set_random_seeds(get_train_seed(get_global_seed()))
 
     @staticmethod
     def _shape_description(modality: Modality) -> str: