diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index 4647c23fcaf..7c4f051b274 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -511,6 +511,28 @@ 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, @@ -518,6 +540,7 @@ def _dispatch_load_leaf(payload, _gpu_id): "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]: diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index f9331288698..b097cda4138 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/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 diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index d182b496032..4ffc418c578 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/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: diff --git a/src/main/python/tests/scuro/test_unimodal_optimizer.py b/src/main/python/tests/scuro/test_unimodal_optimizer.py index 215b4712466..27c19c5e1cf 100644 --- a/src/main/python/tests/scuro/test_unimodal_optimizer.py +++ b/src/main/python/tests/scuro/test_unimodal_optimizer.py @@ -25,11 +25,13 @@ import numpy as np from systemds.scuro.representations.color_histogram import ColorHistogram -from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators from systemds.scuro.drsearch.node_executor import ResultEntry from systemds.scuro.drsearch.unimodal_optimizer import ( UnimodalOptimizer, UnimodalResults, + _lowest_fitted_key, + evaluate_fitted_dags, ) from systemds.scuro.representations.covarep_audio_features import ZeroCrossing @@ -63,6 +65,7 @@ from tests.scuro.data_generator import ( ModalityRandomDataGenerator, TestDataLoader, + TestSVM, TestTask, ) import copy @@ -256,6 +259,124 @@ def decode(name): for lazy_output, once_output in zip(lazy_outputs, once_outputs): np.testing.assert_array_equal(lazy_output, once_output) + def _text_and_two_tasks(self): + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 10 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + second = copy.copy(self.tasks[0]) + second.model = TestSVM("Test2") + return text, [self.tasks[0], second] + + def test_fitted_dags_fit_once_per_fold_for_all_dags_and_tasks(self): + text, tasks = self._text_and_two_tasks() + builder = RepresentationDAGBuilder() + tfidf = builder.create_operation_node( + TfIdf, [builder.create_leaf_node(text.modality_id)], {} + ) + plain = builder.build(tfidf, dag_id=0) + reduced = builder.build( + builder.create_operation_node( + MLPAveraging, [tfidf], MLPAveraging().get_current_parameters() + ), + dag_id=1, + ) + self.assertEqual(_lowest_fitted_key(plain), _lowest_fitted_key(reduced)) + + fit_sets, inputs = [], {} + tfidf_fit, svm_fit = TfIdf.fit, TestSVM.fit + + def recording_fit(op, data, labels=None): + fit_sets.append(set(data.instance_ids)) + return tfidf_fit(op, data, labels=labels) + + def recording_svm(model, X, y, X_test, y_test): + inputs.setdefault(model.name, []).append(np.asarray(X)) + return svm_fit(model, X, y, X_test, y_test) + + with patch.object(TfIdf, "fit", recording_fit), patch.object( + TestSVM, "fit", recording_svm + ): + results = evaluate_fitted_dags([plain, reduced], text, tasks) + shared_inputs, inputs = inputs, {} + for dag in (plain, reduced): + SplitAwareDAGEvaluator(dag, text, tasks[0]).evaluate() + + folds = [set(split.fit_indices) for split in tasks[0].evaluation_splits] + self.assertEqual(fit_sets[: len(folds)], folds) + self.assertEqual(len(fit_sets), 3 * len(folds)) + self.assertEqual( + [(dag_id, task_idx) for dag_id, task_idx, _, _ in results], + [(0, 0), (0, 1), (1, 0), (1, 1)], + ) + self.assertTrue(all(r[2].representation_time > 0 for r in results)) + reference = inputs["Test1"] + for name in ("Test1", "Test2"): + self.assertEqual(len(shared_inputs[name]), len(reference)) + for shared, fresh in zip(shared_inputs[name], reference): + np.testing.assert_array_equal(shared, fresh) + + def test_fitted_dags_are_evaluated_on_executor_workers(self): + text, tasks = self._text_and_two_tasks() + load_builtin_operators((ModalityType.TEXT, ModalityType.EMBEDDING)) + with patch.object( + Registry, "_representations", registry_for(ModalityType.TEXT, [TfIdf]) + ), patch( + "systemds.scuro.drsearch.unimodal_optimizer.evaluate_fitted_dags", + side_effect=AssertionError("fitted DAGs evaluated in the driver"), + ): + optimizer = UnimodalOptimizer( + [text], tasks, False, max_num_workers=2, enable_checkpointing=False + ) + optimizer.optimize() + + fitted_ids = [ + { + entry.dag.dag_id + for entry in optimizer.operator_performance.results[text.modality_id][ + task.model.name + ] + if dag_requires_fitting(entry.dag) + } + for task in tasks + ] + self.assertTrue(fitted_ids[0]) + self.assertEqual(fitted_ids[0], fitted_ids[1]) + + def test_k_best_tasks_share_an_unsupervised_leaf_fit(self): + text, tasks = self._text_and_two_tasks() + results = UnimodalResults([text], tasks, k=1) + builder = RepresentationDAGBuilder() + leaf = builder.create_leaf_node(text.modality_id) + dag = builder.build(builder.create_operation_node(TfIdf, [leaf], {})) + for task in tasks: + results.results[text.modality_id][task.model.name].append( + SimpleNamespace(dag=dag, val_score={"accuracy": 0.9}) + ) + fits = [] + fit_transform = SplitAwareDAGEvaluator.fit_transform + + def recording(evaluator, *args): + fits.append(evaluator.task.model.name) + return fit_transform(evaluator, *args) + + with patch.object(SplitAwareDAGEvaluator, "fit_transform", recording): + leaves = [ + results.get_k_best_results(text, task, "accuracy")[1][0] + for task in tasks + ] + + self.assertEqual(fits, ["Test1"]) + self.assertIsNot(leaves[0], leaves[1]) + self.assertEqual(leaves[1].fold_source, (dag, text)) + np.testing.assert_array_equal( + np.asarray(leaves[0].data), np.asarray(leaves[1].data) + ) + def test_robust_results_ignore_non_finite_scores(self): modality = SimpleNamespace(modality_id="modality") task = SimpleNamespace(model=SimpleNamespace(name="task"))