From b021c1ab23e1e01704ef9b3748a9196806354718 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 12:40:20 +0200 Subject: [PATCH 1/8] initial ga optimizer --- .../drsearch/multimodal_ga_deap_optimizer.py | 1277 +++++++++++++++++ .../scuro/drsearch/multimodal_ga_optimizer.py | 1235 ++++++++++++++++ 2 files changed, 2512 insertions(+) create mode 100644 src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py create mode 100644 src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py new file mode 100644 index 00000000000..1b8c2223ee7 --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py @@ -0,0 +1,1277 @@ +# ------------------------------------------------------------- +# +# 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 __future__ import annotations +import copy +import multiprocessing as mp +import os +import pickle +import random +import tempfile +import time +import traceback +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait +from dataclasses import dataclass, field +from itertools import chain +from typing import Any, Dict, List, Optional, Tuple +from deap import base, creator, tools +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + _collect_internal_paths, + _get_subtree, + _replace_subtree, + _rebuild_fusion_ops, + _remove_leaf_from_tree, + _reindex_tree, + _collect_leaf_indices, +) +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.drsearch.representation_dag import ( + RepresentationDAGBuilder, + RepresentationDag, +) +from systemds.scuro.drsearch.task import Task +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.utils.schema_helpers import get_shape + +Tree = int | Tuple["Tree", "Tree"] + + +@dataclass +class FusionSearchResult: + dag: RepresentationDag + train_score: dict + val_score: dict + test_score: dict + runtime: float = 0.0 + task_time: float = 0.0 + representation_time: float = 0.0 + task_name: str = "" + + # Per-fold score vectors, {metric: [fold0, ...]}. Needed for error bars and + # for any selection rule that has to tell a consistent winner apart from a + # lucky-fold one; the averaged dicts above cannot. + val_fold_scores: dict = field(default_factory=dict) + train_fold_scores: dict = field(default_factory=dict) + test_fold_scores: dict = field(default_factory=dict) + + # Fit/inference wall-clock from the fold loop, per fold and averaged. + task_timing: dict = field(default_factory=dict) + + # Search trajectory: generation, position in the evaluation sequence, and + # seconds since this GA run began. Results are appended in evaluation + # order, but nothing recorded *when* -- which is what a best-so-far curve + # against a wall-clock budget needs. + generation: int = -1 + eval_index: int = -1 + t_since_search_start_s: float = 0.0 + t_eval_end_unix: float = 0.0 + + +@dataclass +class DagGenome: + leaves: List[Tuple[str, int]] + tree: Tree + fusion_ops: Dict[str, type] + + +# ---------------------------------------------------------------------- +# Module-level evaluation helpers. +# +# These are intentionally free functions (not bound methods) so that +# `_evaluate_dag_worker` can be shipped to a separate process and pickled +# safely. `_evaluate_genome_body` holds the actual logic and is shared by +# both the in-process (serial) and cross-process (parallel) evaluation +# paths, so the two can never silently diverge in behavior. +# ---------------------------------------------------------------------- + + +# Objectives that come from evaluation timing rather than from the task's +# score dict (scores[i].average_scores). Anything not listed here is looked +# up as a key in the validation score dict (e.g. "accuracy", "f1"). +_TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"} + +ObjectiveSpec = Tuple[str, str] # (name, "max" | "min") + + +def _objective_value( + name: str, val_score: Dict[str, float], timing: Dict[str, float] +) -> float: + if name in _TIMING_OBJECTIVES: + return timing[name] + return val_score[name] + + +def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: + """ + The worst possible value for each objective, direction-aware: -inf for a + "max" objective, +inf for a "min" objective (e.g. runtime), so that a + failed evaluation always ranks last regardless of DEAP's per-objective + fitness weight sign (wvalue = raw_value * weight). + """ + return tuple( + float("-inf") if direction == "max" else float("inf") + for _, direction in objective_specs + ) + + +def _evaluate_genome_body( + dag: RepresentationDag, + task: Task, + modalities: List[Any], + objective_specs: List[ObjectiveSpec], +) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: + start = time.time() + fused = dag.execute(modalities, task, enable_cache=False) + if fused is None: + return None, None + + if isinstance(fused, dict): + fused = fused[list(fused.keys())[-1]] + + if task.expected_dim == 1 and get_shape(fused.metadata) > 1: + fused = AggregatedRepresentation().transform(fused) + + t0 = time.time() + scores = task.run(fused.data) + task_time = time.time() - t0 + total = time.time() - start + + val_score = scores[1].average_scores + timing = { + "runtime": total, + "task_time": task_time, + "representation_time": total - task_time, + } + fitness = tuple( + _objective_value(name, val_score, timing) for name, _ in objective_specs + ) + payload = { + "train_score": scores[0].average_scores, + "val_score": val_score, + "test_score": scores[2].average_scores, + "train_fold_scores": scores[0].fold_scores(), + "val_fold_scores": scores[1].fold_scores(), + "test_fold_scores": scores[2].fold_scores(), + "task_timing": getattr(task, "last_run_timing", {}), + **timing, + } + return fitness, payload + + +def _evaluate_dag_worker( + dag_bytes: bytes, + task_bytes: bytes, + modalities_bytes: bytes, + objective_specs: List[ObjectiveSpec], +) -> Tuple[Tuple[float, ...], Optional[Dict[str, Any]], Optional[str]]: + """ + Runs in a worker process. Never raises: any failure (a shape-incompatible + fusion, an OOM, a bug in a representation) is caught and reported back as + plain, picklable data, so a single bad genome can never take down the + rest of the population's evaluation. + """ + try: + dag = pickle.loads(dag_bytes) + task = pickle.loads(task_bytes) + modalities = pickle.loads(modalities_bytes) + fitness, payload = _evaluate_genome_body(dag, task, modalities, objective_specs) + if fitness is None: + fitness = _failure_fitness(objective_specs) + return fitness, payload, None + except Exception: + return _failure_fitness(objective_specs), None, traceback.format_exc() + + +class MultimodalDeapOptimizer: + """ + Multimodal optimizer that uses DEAP's genome/fitness containers and + selection operators to run a genetic search over fusion DAGs. + + @param modalities: List of modalities to optimize. + @param unimodal_optimization_results: Unimodal optimization results. + @param tasks: List of tasks to optimize. + @param debug: Whether to print debug information. + @param min_modalities: Minimum number of modalities to use. + @param max_modalities: Maximum number of modalities to use. + @param metric: Metric to optimize (used when objectives is not given). + @param objectives: Optional list of (name, direction) pairs to run a + multi-objective (Pareto/NSGA-II) search instead of optimizing a + single scalar metric, e.g. [("accuracy", "max"), ("runtime", + "min")]. "runtime", "task_time", "representation_time" are read + from evaluation timing; any other name is looked up in the task's + validation-score dict (e.g. "accuracy", "f1"). When given, this + overrides metric/maximize_metric. Defaults to None (single + objective, unchanged behavior). + @param population_size: Population size. + @param generations: Maximum number of generations. + @param crossover_probability: Crossover probability. + @param mutation_probability: Mutation probability. + @param random_seed: Random seed. + @param elite_size: Number of top individuals carried over unchanged + (with their fitness preserved) into the next generation. + @param max_workers: Number of worker processes used to evaluate a + generation's population. 1 (default) evaluates serially in-process. + @param batch_size: Maximum number of genome evaluations kept in flight + at once when max_workers > 1. Bounds peak memory/GPU usage + independently of max_workers; defaults to max_workers. + @param early_stopping_patience: Stop once this many consecutive + generations pass without a validation-metric improvement larger + than early_stopping_min_delta. None disables early stopping. + @param early_stopping_min_delta: Minimum improvement to reset the + early-stopping counter. + @param allow_repeated_modalities: Allow one modality to contribute more + than one leaf to a genome, each with a different representation + (intra-modal fusion). Off by default, which keeps the historical + behaviour of at most one leaf per modality and caps the leaf count at + len(modalities). Turning it on raises that cap to the total number of + distinct (modality, representation) leaves available. + @param hall_of_fame_size: How many incumbents to keep per task in + single-objective mode (ignored in multi-objective mode, where the + whole non-dominated front is kept). See get_hall_of_fame. + @param novelty_breeding: Reject offspring whose genome was already + evaluated in an *earlier* generation, not just earlier in the + current one. Without this a converged population refills itself + with re-treads that hit the fitness cache, cost no wall clock and + explore nothing -- on the regret-analysis replays 55-84% of all + fitness requests were such duplicates, so a population of 16 + behaved like a population of ~5. Elites are exempt (they are + carried over before the archive is consulted), and once the + archive covers most of the reachable space the retry budget runs + out and breeding falls back to duplicates, so the search + degrades gracefully rather than stalling. Defaults to True; set + False to reproduce pre-fix runs. + """ + + def __init__( + self, + modalities: List[Any], + unimodal_optimization_results: Any, + tasks: List[Task], + debug: bool = True, + min_modalities: int = 2, + max_modalities: int = None, + metric: str = "accuracy", + objectives: Optional[List[ObjectiveSpec]] = None, + checkpoint_every: int = None, + resume: bool = True, + population_size: int = 32, + generations: int = 20, + crossover_probability: float = 0.7, + mutation_probability: float = 0.4, + random_seed: int = 42, + maximize_metric: bool = True, + elite_size: int = 2, + max_workers: int = 1, + batch_size: Optional[int] = None, + early_stopping_patience: Optional[int] = 5, + early_stopping_min_delta: float = 1e-6, + novelty_breeding: bool = True, + hall_of_fame_size: int = 5, + allow_repeated_modalities: bool = False, + ): + self.modalities = modalities + self.tasks = tasks + self.debug = debug + # A genome may hold several leaves of the same modality (different + # representations of it) only when this is on. With it off, one + # modality contributes at most one leaf, so the leaf count can never + # exceed len(modalities) and max_modalities is clamped to that. + self.allow_repeated_modalities = allow_repeated_modalities + + # min_modalities=1 admits single-leaf (unimodal) genomes. It is not the + # default -- a "multimodal" search that returns a unimodal pipeline is + # almost always a bug, which is why the floor used to be a hard 2 -- + # but it is a legitimate configuration for a sweep that asks what the + # fusion actually buys, so it is now the caller's choice. + self.min_modalities = max(1, min_modalities) + requested_max = max_modalities or len(modalities) + self.max_modalities = ( + requested_max + if allow_repeated_modalities + else min(requested_max, len(modalities)) + ) + if self.max_modalities < self.min_modalities: + raise ValueError( + f"max_modalities ({self.max_modalities}) is below min_modalities " + f"({self.min_modalities})" + ) + self.metric_name = metric + self.maximize_metric = maximize_metric + + if objectives is not None: + if len(objectives) < 1: + raise ValueError( + "objectives must contain at least one (name, direction) pair" + ) + for name, direction in objectives: + if direction not in ("max", "min"): + raise ValueError( + f"objective direction must be 'max' or 'min', got " + f"{direction!r} for objective {name!r}" + ) + self.objective_specs: List[ObjectiveSpec] = list(objectives) + # Keep metric_name/maximize_metric meaningful for callers that + # still read them (e.g. _extract_k_best's ranking metric) - must + # be resolved before _extract_k_best() runs below. + self.metric_name = self.objective_specs[0][0] + self.maximize_metric = self.objective_specs[0][1] == "max" + else: + self.objective_specs = [ + (self.metric_name, "max" if self.maximize_metric else "min") + ] + self.is_multi_objective = len(self.objective_specs) > 1 + + if len(self.modalities) < self.min_modalities: + raise ValueError( + f"MultimodalDeapOptimizer requires at least {self.min_modalities} " + f"modalities, got {len(self.modalities)}." + ) + + self.operator_registry = Registry() + self.fusion_operators = self.operator_registry.get_fusion_operators() + if not self.fusion_operators: + raise ValueError( + "MultimodalDeapOptimizer requires at least one registered " + "fusion operator." + ) + self.k_best_representations = self._extract_k_best( + unimodal_optimization_results + ) + + self.optimization_results: Dict[str, List[FusionSearchResult]] = {} + self.evaluation_errors: Dict[str, int] = {} + + # Incumbent(s) per task, maintained as evaluations come in, so the + # winner does not have to be recovered by re-ranking every result the + # search ever produced. Single-objective: the hall_of_fame_size best + # by the configured objective. Multi-objective: the (unbounded) + # non-dominated front over the objective tuples. + self.hall_of_fame_size = max(1, hall_of_fame_size) + self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} + self._hof_fitness: Dict[str, List[Tuple[float, ...]]] = {} + self.rng = random.Random(random_seed) + # Search-trajectory bookkeeping: see FusionSearchResult.eval_index. + self._eval_counter = 0 + self._current_generation = -1 + self._search_start = time.perf_counter() + self.population_size = max(1, population_size) + self.generations = max(1, generations) + self.crossover_probability = crossover_probability + self.mutation_probability = mutation_probability + self.random_seed = random_seed + self._fitness_cache: Dict[str, Dict[Tuple, Tuple[float, ...]]] = {} + + # Leave room for at least one non-elite offspring per generation, + # otherwise the next generation would just be a sorted clone of the + # current one and crossover/mutation would never run. Unused in + # multi-objective mode (NSGA-II selection is implicitly elitist). + self.elite_size = max(0, min(elite_size, self.population_size - 1)) + self.max_workers = max(1, max_workers) + self.batch_size = max(1, batch_size or self.max_workers) + self.early_stopping_patience = early_stopping_patience + self.early_stopping_min_delta = early_stopping_min_delta + self.novelty_breeding = novelty_breeding + self._current_task_name = None + + desired_weights = tuple( + 1.0 if direction == "max" else -1.0 for _, direction in self.objective_specs + ) + self._objective_weights = desired_weights + existing_weights = getattr( + getattr(creator, "FitnessMax", None), "weights", None + ) + if existing_weights != desired_weights: + if hasattr(creator, "Individual"): + del creator.Individual + if hasattr(creator, "FitnessMax"): + del creator.FitnessMax + creator.create("FitnessMax", base.Fitness, weights=desired_weights) + creator.create("Individual", list, fitness=creator.FitnessMax) + elif not hasattr(creator, "Individual"): + creator.create("Individual", list, fitness=creator.FitnessMax) + + # ------------------------------------------------------------------ + # Public entrypoint + # ------------------------------------------------------------------ + + def optimize( + self, + ) -> Dict[str, List[FusionSearchResult]]: + """ + Optimize the multimodal representations for the tasks. + @return: Dictionary of optimization results for each task. + """ + for task in self.tasks: + task_name = task.model.name + self._current_task_name = task_name + self.optimization_results.setdefault(task_name, []) + self.evaluation_errors.setdefault(task_name, 0) + + # Each task restarts the search, so its trajectory restarts too. + self._eval_counter = 0 + self._search_start = time.perf_counter() + + population = self._build_initial_population(task_name) + best_ever = None + no_improve = 0 + + for gen in range(self.generations): + self._current_generation = gen + self._evaluate_population(population, task) + + if self.is_multi_objective: + # No single "best" individual exists with multiple + # objectives - track the non-dominated (Pareto) front + # instead; progress means that front actually changed. + front = tools.sortNondominated( + population, len(population), first_front_only=True + )[0] + front_signature = frozenset( + self._genome_signature(ind[0]) for ind in front + ) + if best_ever is None or front_signature != best_ever: + best_ever = front_signature + no_improve = 0 + else: + no_improve += 1 + debug_msg = f"front_size={len(front)}" + else: + gen_best = max(population, key=lambda ind: ind.fitness.values[0]) + if ( + best_ever is None + or gen_best.fitness.values[0] + > best_ever.fitness.values[0] + self.early_stopping_min_delta + ): + best_ever = self._clone_individual(gen_best) + no_improve = 0 + else: + no_improve += 1 + debug_msg = f"best={gen_best.fitness.values[0]:.4f}" + + if self.debug: + print( + f"[GA] task={task_name} gen={gen} {debug_msg} " + f"no_improve={no_improve} " + f"errors={self.evaluation_errors.get(task_name, 0)}" + ) + + stagnated = ( + self.early_stopping_patience is not None + and no_improve >= self.early_stopping_patience + ) + if stagnated or gen == self.generations - 1: + if self.debug and stagnated: + print( + f"[GA] task={task_name} early stopping after " + f"{no_improve} generations without improvement" + ) + break + + population = self._next_generation(population, task_name, task) + + return self.optimization_results + + # ------------------------------------------------------------------ + # Population lifecycle + # ------------------------------------------------------------------ + + def _make_individual(self, genome: DagGenome): + return creator.Individual([genome]) + + def _clone_individual(self, ind): + clone = creator.Individual([copy.deepcopy(ind[0])]) + if ind.fitness.valid: + clone.fitness.values = ind.fitness.values + return clone + + def _append_if_unique( + self, + population: List[Any], + genome: DagGenome, + seen_signatures: set, + ) -> bool: + if len(population) >= self.population_size: + return False + sig = self._genome_signature(genome) + if sig in seen_signatures: + return False + seen_signatures.add(sig) + population.append(self._make_individual(genome)) + return True + + def _build_initial_population(self, task_name: str) -> List[Any]: + population: List[Any] = [] + seen: set = set() + retry_budget = max(20, self.population_size * 10) + retries = 0 + while len(population) < self.population_size and retries < retry_budget: + genome = self._random_genome(task_name) + if self._append_if_unique(population, genome, seen): + retries = 0 + else: + retries += 1 + # Search space smaller than population_size: fall back to + # (possibly duplicated) random individuals rather than looping + # forever. + while len(population) < self.population_size: + population.append(self._make_individual(self._random_genome(task_name))) + return population + + def _next_generation( + self, population: List[Any], task_name: str, task: Task + ) -> List[Any]: + if self.is_multi_objective: + # Classic (mu+lambda) NSGA-II: breed a full offspring pool, + # evaluate it, then select the next generation from parents + + # offspring combined via non-dominated sorting + crowding + # distance. This is inherently elitist (non-dominated parents + # survive on their own merit), so there's no separate elite_size + # concept here. + offspring = self._breed_offspring( + population, task_name, seen=self._novelty_archive(task_name) + ) + self._evaluate_population(offspring, task) + combined = list(population) + list(offspring) + return list(tools.selNSGA2(combined, self.population_size)) + + ranked = sorted(population, key=lambda ind: ind.fitness.values[0], reverse=True) + elite = [self._clone_individual(ind) for ind in ranked[: self.elite_size]] + seen = {self._genome_signature(ind[0]) for ind in elite} + seen |= self._novelty_archive(task_name) + return self._breed_offspring(population, task_name, initial=elite, seen=seen) + + def _novelty_archive(self, task_name: str) -> set: + """Signatures already evaluated for this task, or empty if disabled. + + `_breed_offspring` dedups against `seen`, which without this holds + only the current generation's elite -- so an offspring identical to + something evaluated three generations ago is accepted, served from + `_fitness_cache`, and occupies a population slot that explores + nothing. Feeding the whole cache in makes the dedup global. + """ + if not self.novelty_breeding: + return set() + return set(self._fitness_cache.get(task_name, {})) + + def _breed_offspring( + self, + population: List[Any], + task_name: str, + initial: Optional[List[Any]] = None, + seen: Optional[set] = None, + ) -> List[Any]: + """ + Fills a new population of size population_size via tournament + selection + crossover + mutation, starting from `initial` (e.g. an + elite carryover) if given. Falls back to fresh random immigrants + (allowing duplicates) once the retry budget is exhausted, so a + search space smaller than population_size can never hang. + """ + next_population = list(initial) if initial else [] + seen = set(seen) if seen else set() + + retry_budget = max(20, self.population_size * 10) + retries = 0 + tournsize = max(1, min(3, len(population))) + while len(next_population) < self.population_size and retries < retry_budget: + p1, p2 = tools.selTournament(population, 2, tournsize=tournsize) + + if self.rng.random() < self.crossover_probability: + g1, g2 = self._crossover_genomes(p1[0], p2[0]) + else: + g1, g2 = copy.deepcopy(p1[0]), copy.deepcopy(p2[0]) + + if self.rng.random() < self.mutation_probability: + g1 = self._mutate_genome(g1, task_name) + if self.rng.random() < self.mutation_probability: + g2 = self._mutate_genome(g2, task_name) + + added1 = self._append_if_unique(next_population, g1, seen) + added2 = self._append_if_unique(next_population, g2, seen) + retries = 0 if (added1 or added2) else retries + 1 + + # Ran out of retries (tiny search space): top up with fresh random + # immigrants, allowing duplicates rather than looping forever. + while len(next_population) < self.population_size: + genome = self._random_genome(task_name) + if not self._append_if_unique(next_population, genome, seen): + next_population.append(self._make_individual(genome)) + + return next_population + + # ------------------------------------------------------------------ + # Evaluation (serial + bounded parallel) + # ------------------------------------------------------------------ + + def _evaluate_population(self, population: List[Any], task: Task) -> None: + to_evaluate = [ind for ind in population if not ind.fitness.valid] + if not to_evaluate: + return + if self.max_workers > 1 and len(to_evaluate) > 1: + self._evaluate_individuals_parallel(to_evaluate, task) + else: + for ind in to_evaluate: + fitness = self._evaluate_genome(ind[0], task) + ind.fitness.values = fitness + + def _evaluate_individuals_parallel( + self, individuals: List[Any], task: Task + ) -> None: + task_name = task.model.name + cache = self._fitness_cache.setdefault(task_name, {}) + ctx = mp.get_context("spawn") + task_bytes = pickle.dumps(task) + futures: Dict[Any, Tuple[Any, RepresentationDag, Tuple]] = {} + pending_followers: Dict[Tuple, List[Any]] = {} + + def _drain(done_futures): + for fut in done_futures: + ind, dag, sig = futures.pop(fut) + fitness, payload, error = fut.result() + ind.fitness.values = fitness + self._record_evaluation(task_name, dag, payload, error) + cache[sig] = fitness + for follower in pending_followers.pop(sig, []): + follower.fitness.values = fitness + + with ProcessPoolExecutor( + max_workers=self.max_workers, mp_context=ctx + ) as executor: + for ind in individuals: + genome = ind[0] + sig = self._genome_signature(genome) + cached = cache.get(sig) + if cached is not None: + ind.fitness.values = cached + continue + if sig in pending_followers: + pending_followers[sig].append(ind) + continue + + dag = self._genome_to_dag(genome) + modalities = list( + chain.from_iterable(self.k_best_representations[task_name].values()) + ) + fut = executor.submit( + _evaluate_dag_worker, + pickle.dumps(dag), + task_bytes, + pickle.dumps(modalities), + self.objective_specs, + ) + futures[fut] = (ind, dag, sig) + pending_followers[sig] = [] + + if len(futures) >= self.batch_size: + done, _ = wait(set(futures.keys()), return_when=FIRST_COMPLETED) + _drain(done) + + if futures: + done, _ = wait(set(futures.keys())) + _drain(done) + + def _record_evaluation( + self, + task_name: str, + dag: RepresentationDag, + payload: Optional[Dict[str, Any]], + error: Optional[str], + ) -> None: + self.optimization_results.setdefault(task_name, []) + if error is not None or payload is None: + self.evaluation_errors[task_name] = ( + self.evaluation_errors.get(task_name, 0) + 1 + ) + if self.debug and error is not None: + last_line = error.strip().splitlines()[-1] if error.strip() else error + print( + f"[GA] genome evaluation failed for task={task_name}: {last_line}" + ) + return + + result = FusionSearchResult( + dag=dag, + train_score=payload["train_score"], + val_score=payload["val_score"], + test_score=payload["test_score"], + train_fold_scores=payload.get("train_fold_scores", {}), + val_fold_scores=payload.get("val_fold_scores", {}), + test_fold_scores=payload.get("test_fold_scores", {}), + task_timing=payload.get("task_timing", {}), + runtime=payload["runtime"], + task_time=payload["task_time"], + representation_time=payload["representation_time"], + task_name=task_name, + generation=self._current_generation, + eval_index=self._eval_counter, + t_since_search_start_s=time.perf_counter() - self._search_start, + t_eval_end_unix=time.time(), + ) + self.optimization_results.setdefault(task_name, []).append(result) + self._update_hall_of_fame(task_name, result) + self._eval_counter += 1 + + # ------------------------------------------------------------------ + # Hall of fame + # ------------------------------------------------------------------ + + def _dominates(self, a: Tuple[float, ...], b: Tuple[float, ...]) -> bool: + """True if objective tuple `a` Pareto-dominates `b`, direction-aware.""" + wa = [w * v for w, v in zip(self._objective_weights, a)] + wb = [w * v for w, v in zip(self._objective_weights, b)] + return all(x >= y for x, y in zip(wa, wb)) and any( + x > y for x, y in zip(wa, wb) + ) + + def _update_hall_of_fame(self, task_name: str, result: FusionSearchResult) -> None: + timing = { + "runtime": result.runtime, + "task_time": result.task_time, + "representation_time": result.representation_time, + } + try: + fitness = tuple( + _objective_value(name, result.val_score, timing) + for name, _ in self.objective_specs + ) + except KeyError: + # An objective the task did not report for this candidate: it + # cannot be ranked, so it simply does not enter the hall of fame. + return + + hof = self.hall_of_fame.setdefault(task_name, []) + fits = self._hof_fitness.setdefault(task_name, []) + + if self.is_multi_objective: + if any(self._dominates(f, fitness) or f == fitness for f in fits): + return + keep = [i for i, f in enumerate(fits) if not self._dominates(fitness, f)] + self.hall_of_fame[task_name] = [hof[i] for i in keep] + [result] + self._hof_fitness[task_name] = [fits[i] for i in keep] + [fitness] + return + + weight = self._objective_weights[0] + hof.append(result) + fits.append(fitness) + order = sorted( + range(len(fits)), key=lambda i: weight * fits[i][0], reverse=True + ) + order = order[: self.hall_of_fame_size] + self.hall_of_fame[task_name] = [hof[i] for i in order] + self._hof_fitness[task_name] = [fits[i] for i in order] + + def get_hall_of_fame(self, task_name: str) -> List[FusionSearchResult]: + """The incumbent(s) for a task: the best `hall_of_fame_size` results by + the configured objective, or the non-dominated front in multi-objective + mode. Empty if nothing evaluated successfully.""" + return list(self.hall_of_fame.get(task_name, [])) + + def _extract_k_best(self, unimodal_results) -> Dict[str, Dict[str, List[Any]]]: + """ + Extract the k best representations for each modality and task. + @param unimodal_results: Unimodal optimization results. + @return: Dictionary of k best representations for each modality and task. + """ + k_best = {} + for task in self.tasks: + name = task.model.name + k_best[name] = {} + for modality in self.modalities: + _, cached_data = unimodal_results.get_k_best_results( + modality, task, self.metric_name + ) + k_best[name][modality.modality_id] = cached_data + return k_best + + def _available_modality_ids(self, task_name: str) -> List[Any]: + """Modalities with at least one representation for this task. + + Datasets like StressID can be missing a modality for some + tasks/instances, and a leaf pointing at an empty representation list + is unresolvable. + """ + reps = self.k_best_representations[task_name] + return [ + m.modality_id + for m in self.modalities + if len(reps.get(m.modality_id, [])) > 0 + ] + + def _leaf_capacity(self, task_name: str) -> int: + """How many distinct leaves a genome may hold at most. + + One per modality normally; with allow_repeated_modalities, one per + (modality, representation) pair, since that is what makes a leaf + distinct. + """ + reps = self.k_best_representations[task_name] + ids = self._available_modality_ids(task_name) + if not self.allow_repeated_modalities: + return len(ids) + return sum(len(reps[mid]) for mid in ids) + + def _random_genome(self, task_name: str) -> DagGenome: + """ + Generate a random genome for a given task. + @param task_name: Name of the task. + @return: Random genome. + """ + reps = self.k_best_representations[task_name] + available_modality_ids = self._available_modality_ids(task_name) + capacity = self._leaf_capacity(task_name) + if capacity < self.min_modalities: + raise ValueError( + f"Need at least {self.min_modalities} distinct leaves for task " + f"'{task_name}', but only {capacity} are available across " + f"{len(available_modality_ids)} modalities." + ) + + upper = min(self.max_modalities, capacity) + lower = min(self.min_modalities, upper) + r = self.rng.randint(lower, upper) + + if self.allow_repeated_modalities: + # Sample distinct (modality, representation) leaves; a modality may + # appear more than once as long as it contributes a different + # representation each time. + pool = [ + (mid, idx) + for mid in available_modality_ids + for idx in range(len(reps[mid])) + ] + leaves = self.rng.sample(pool, r) + else: + chosen = self.rng.sample(available_modality_ids, r) + leaves = [(mid, self.rng.randrange(len(reps[mid]))) for mid in chosen] + + tree = self._random_binary_tree(len(leaves)) + fusion_ops = {} + self._assign_fusion_ops(tree, fusion_ops, "") + return DagGenome(leaves=leaves, tree=tree, fusion_ops=fusion_ops) + + def _internal_paths(self, tree): + return _collect_internal_paths(tree) + + def _random_binary_tree(self, n: int) -> Tree: + """ + Generate a random binary tree with n leaves. + @param n: Number of leaves. + @return: Random binary tree. + """ + nodes: List[Tree] = list(range(n)) + while len(nodes) > 1: + i, j = self.rng.sample(range(len(nodes)), 2) + a, b = nodes.pop(max(i, j)), nodes.pop(min(i, j)) + nodes.append((a, b)) + return nodes[0] + + def _assign_fusion_ops(self, subtree: Tree, ops: Dict[str, Any], path: str) -> None: + """ + Assign fusion operators to the internal nodes of a binary tree. + @param subtree: Binary tree. + @param ops: Dictionary of fusion operators. + @param path: Path to the current node. + """ + if isinstance(subtree, int): + return + ops[path] = self.rng.choice(self.fusion_operators) + left, right = subtree + self._assign_fusion_ops(left, ops, path + "L") + self._assign_fusion_ops(right, ops, path + "R") + + def _genome_to_dag(self, genome: DagGenome) -> RepresentationDag: + """ + Convert a genome to a DAG. + @param genome: Genome. + @return: RepresentationDAG. + """ + builder = RepresentationDAGBuilder() + leaf_ids = [ + builder.create_leaf_node(mod_id, repr_idx) + for mod_id, repr_idx in genome.leaves + ] + + def build(subtree: Tree, path: str) -> str: + if isinstance(subtree, int): + return leaf_ids[subtree] + left, right = subtree + left_id = build(left, path + "L") + right_id = build(right, path + "R") + op_cls = genome.fusion_ops[path] + op = op_cls() + return builder.create_operation_node( + op.__class__, [left_id, right_id], op.get_current_parameters() + ) + + return builder.build(build(genome.tree, "")) + + def _genome_signature(self, g: DagGenome) -> Tuple: + """ + Generate a signature for a genome to be used as a cache key. + @param g: Genome. + @return: Signature. + """ + + def norm(t: Tree): + return t if isinstance(t, int) else (norm(t[0]), norm(t[1])) + + return ( + tuple(g.leaves), + norm(g.tree), + tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), + ) + + def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: + """ + Evaluate a genome for a given task (in-process). Stores the result + in the optimization results, or records a failure without raising + if the DAG/task fails to execute (e.g. an incompatible fusion). + @param genome: Genome. + @param task: Task. + @return: Fitness tuple, one value per configured objective. + """ + task_name = task.model.name + sig = self._genome_signature(genome) + cache = self._fitness_cache.setdefault(task_name, {}) + if sig in cache: + return cache[sig] + + dag = self._genome_to_dag(genome) + modalities = list( + chain.from_iterable(self.k_best_representations[task_name].values()) + ) + + try: + fitness, payload = _evaluate_genome_body( + dag, task, modalities, self.objective_specs + ) + error = None + if fitness is None: + fitness = _failure_fitness(self.objective_specs) + except Exception: + fitness = _failure_fitness(self.objective_specs) + payload, error = None, traceback.format_exc() + + self._record_evaluation(task_name, dag, payload, error) + cache[sig] = fitness + return fitness + + # ------------------------------------------------------------------ + # Crossover + # ------------------------------------------------------------------ + + def _crossover_genomes( + self, g1: DagGenome, g2: DagGenome + ) -> Tuple[DagGenome, DagGenome]: + """ + Recombine two genomes. + + When both parents agree on the exact same modality/representation + leaves, a classic reciprocal subtree swap is performed. Otherwise + (the common case, since leaves are sampled independently per + genome) subtree crossover over leaf indices would be meaningless, + so fusion-operator choices are mixed at whichever internal-path + keys the two trees happen to share instead - every non-trivial + tree has at least the root path "" in common, so this always + performs real recombination rather than silently returning the + parents unchanged. + @param g1: First parent genome. + @param g2: Second parent genome. + @return: Two child genomes. + """ + c1, c2 = copy.deepcopy(g1), copy.deepcopy(g2) + + if c1.leaves == c2.leaves: + paths1 = self._internal_paths(c1.tree) + paths2 = self._internal_paths(c2.tree) + if paths1 and paths2: + path1 = self.rng.choice(paths1) + path2 = self.rng.choice(paths2) + subtree1 = _get_subtree(c1.tree, path1) + subtree2 = _get_subtree(c2.tree, path2) + c1.tree = _replace_subtree(c1.tree, path1, subtree2) + c2.tree = _replace_subtree(c2.tree, path2, subtree1) + c1.fusion_ops = _rebuild_fusion_ops( + c1.tree, + {**c1.fusion_ops, **c2.fusion_ops}, + self.rng, + self.fusion_operators, + randomized_prefixes=[path1], + ) + c2.fusion_ops = _rebuild_fusion_ops( + c2.tree, + {**c2.fusion_ops, **c1.fusion_ops}, + self.rng, + self.fusion_operators, + randomized_prefixes=[path2], + ) + return c1, c2 + + shared_paths = set(c1.fusion_ops) & set(c2.fusion_ops) + for path in shared_paths: + if self.rng.random() < 0.5: + c1.fusion_ops[path], c2.fusion_ops[path] = ( + c2.fusion_ops[path], + c1.fusion_ops[path], + ) + return c1, c2 + + # ------------------------------------------------------------------ + # Mutation + # ------------------------------------------------------------------ + + def _mutate_genome(self, g: DagGenome, task_name: str) -> DagGenome: + op = self.rng.choice( + [ + self._mutate_change_fusion, + lambda gg: self._mutate_swap_leaf_repr(gg, task_name), + lambda gg: self.mutate_add_leaf(gg, task_name), + self.mutate_remove_leaf, + self.mutate_replace_subtree, + ] + ) + return op(g) + + def _mutate_change_fusion(self, g: DagGenome) -> DagGenome: + """ + Change a fusion operator at a random internal node. + @param g: Genome. + @return: Mutated genome. + """ + g = copy.deepcopy(g) + paths = [p for p in g.fusion_ops] + if not paths: + return g + path = self.rng.choice(paths) + choices = [op for op in self.fusion_operators if op != g.fusion_ops[path]] + if choices: + g.fusion_ops[path] = self.rng.choice(choices) + return g + + def _mutate_swap_leaf_repr(self, g: DagGenome, task_name: str) -> DagGenome: + """ + Change which k-best unimodal repr a leaf uses (same modality). + @param g: Genome. + @param task_name: Name of the task. + @return: Mutated genome. + """ + g = copy.deepcopy(g) + i = self.rng.randrange(len(g.leaves)) + mod_id, current = g.leaves[i] + k = len(self.k_best_representations[task_name][mod_id]) + if k <= 1: + return g + # With repeated modalities allowed the genome may already hold another + # leaf of this modality; swapping onto that representation would make + # the two leaves identical, which wastes a slot on a duplicate. + taken = { + idx for j, (mid, idx) in enumerate(g.leaves) if mid == mod_id and j != i + } + choices = [idx for idx in range(k) if idx != current and idx not in taken] + if not choices: + return g + g.leaves[i] = (mod_id, self.rng.choice(choices)) + return g + + def mutate_add_leaf(self, g: DagGenome, task_name: str) -> DagGenome: + """ + Add a new leaf to the genome (adding a new modality). Only + modalities that have at least one representation for this task are + considered. + @param g: Genome. + @param task_name: Name of the task. + @return: Mutated genome. + """ + if len(g.leaves) >= self.max_modalities: + return g + g = copy.deepcopy(g) + reps = self.k_best_representations[task_name] + if self.allow_repeated_modalities: + # Any (modality, representation) leaf the genome does not already + # hold, so a modality can be added a second time under a different + # representation. + existing_leaves = set(g.leaves) + candidates = [ + (mid, idx) + for mid in self._available_modality_ids(task_name) + for idx in range(len(reps[mid])) + if (mid, idx) not in existing_leaves + ] + if not candidates: + return g + new_leaf = self.rng.choice(candidates) + else: + existing = {l[0] for l in g.leaves} + available = [ + m.modality_id + for m in self.modalities + if m.modality_id not in existing + and len(reps.get(m.modality_id, [])) > 0 + ] + if not available: + return g + mod_id = self.rng.choice(available) + new_leaf = (mod_id, self.rng.randrange(len(reps[mod_id]))) + new_idx = len(g.leaves) + g.leaves.append(new_leaf) + if isinstance(g.tree, int): + g.tree = (g.tree, new_idx) + g.fusion_ops = {"": self.rng.choice(self.fusion_operators)} + else: + paths = self._internal_paths(g.tree) + path = self.rng.choice(paths) + sub = _get_subtree(g.tree, path) + g.tree = _replace_subtree(g.tree, path, (sub, new_idx)) + # `sub` (and everything under it) just shifted one level deeper + # in the tree (from `path` to `path + "L"`), so its existing + # fusion_ops entries are keyed under stale paths and must be + # rebuilt rather than patched in place - only the brand-new + # node at `path` needs a freshly chosen operator. + g.fusion_ops = _rebuild_fusion_ops( + g.tree, + g.fusion_ops, + self.rng, + self.fusion_operators, + randomized_prefixes=[path], + ) + return g + + def mutate_remove_leaf(self, g: DagGenome) -> DagGenome: + """ + Remove a leaf from the genome (removing a modality). + @param g: Genome. + @return: Mutated genome. + """ + if len(g.leaves) <= self.min_modalities: + return g + g = copy.deepcopy(g) + drop = self.rng.randrange(len(g.leaves)) + new_tree = _remove_leaf_from_tree(g.tree, drop) + if new_tree is None: + return g + # Reindex leaves + fusion_ops paths + keep = [i for i in range(len(g.leaves)) if i != drop] + index_map = {old: new for new, old in enumerate(keep)} + g.leaves = [g.leaves[i] for i in keep] + g.tree = _reindex_tree(new_tree, index_map) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, g.fusion_ops, self.rng, self.fusion_operators + ) + return g + + def mutate_replace_subtree(self, g: DagGenome) -> DagGenome: + """ + Replace a subtree with a random subtree. + + The collapse branch (replacing an internal node by one of its + children) drops every leaf on the other side of that node, so it is + only taken when at least min_modalities leaves survive -- otherwise a + search configured min_modalities=2 silently evaluates unimodal + pipelines. The dropped leaves are also removed from `genome.leaves` + and the tree reindexed; leaving them in place used to desynchronise + the genome from its own tree. + @param g: Genome. + @return: Mutated genome. + """ + g = copy.deepcopy(g) + paths = self._internal_paths(g.tree) + if not paths: + return g + path = self.rng.choice(paths) + sub = _get_subtree(g.tree, path) + if isinstance(sub, int): + return g + + if self.rng.random() < 0.5: + child = sub[0] if self.rng.random() < 0.5 else sub[1] + candidate = _replace_subtree(g.tree, path, child) + kept = sorted(set(_collect_leaf_indices(candidate))) + if len(kept) >= self.min_modalities: + index_map = {old: new for new, old in enumerate(kept)} + g.leaves = [g.leaves[i] for i in kept] + g.tree = _reindex_tree(candidate, index_map) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, {}, self.rng, self.fusion_operators + ) + return g + # Collapsing here would leave fewer leaves than min_modalities, + # so fall through to the reshuffle branch, which is leaf-preserving. + + leaf_idxs = _collect_leaf_indices(sub) + new_sub = self._random_binary_tree(len(leaf_idxs)) + local_map = {i: leaf_idxs[i] for i in range(len(leaf_idxs))} + new_sub = _reindex_tree(new_sub, local_map) + g.tree = _replace_subtree(g.tree, path, new_sub) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, + g.fusion_ops, + self.rng, + self.fusion_operators, + randomized_prefixes=[path], + ) + return g + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + def store_results(self, file_name: str = None, overwrite: bool = False) -> str: + """ + Persist optimization_results to disk. + + Refuses to clobber an existing file unless overwrite=True is + passed explicitly. The write itself is atomic (temp file + + os.replace), so a crash mid-write can never leave a corrupted or + truncated results file behind. + @param file_name: Destination path. A timestamped name is + generated if omitted. + @param overwrite: Set True to explicitly replace an existing file + at file_name. + @return: The path the results were written to. + """ + if file_name is None: + timestr = time.strftime("%Y%m%d-%H%M%S") + file_name = f"multimodal_optimizer_{timestr}.pkl" + + directory = os.path.dirname(file_name) or "." + os.makedirs(directory, exist_ok=True) + + if os.path.exists(file_name) and not overwrite: + raise FileExistsError( + f"Refusing to overwrite existing results file '{file_name}'. " + "Pass overwrite=True if this is intentional, or choose a " + "different file_name." + ) + + fd, tmp_path = tempfile.mkstemp( + dir=directory, prefix=".tmp_multimodal_results_", suffix=".pkl" + ) + try: + with os.fdopen(fd, "wb") as f: + pickle.dump(self.optimization_results, f) + os.replace(tmp_path, file_name) + except Exception: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + return file_name diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py new file mode 100644 index 00000000000..378c2618043 --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -0,0 +1,1235 @@ +from __future__ import annotations + +import copy +import multiprocessing as mp +import pickle +import random +import time +from dataclasses import dataclass, field +from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED +from typing import Any, Dict, Generator, List, Optional, Tuple + +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.drsearch.representation_dag import ( + RepresentationDag, + RepresentationDAGBuilder, +) +from systemds.scuro.drsearch.task import Task +from systemds.scuro.modality.modality import Modality +from systemds.scuro.utils.checkpointing import CheckpointManager +from systemds.scuro.utils.static_variables import DEBUG + +# ---------------------------- +# Genome / Individual encoding +# ---------------------------- + +Tree = Any # int leaf index OR tuple(left_subtree, right_subtree) + + +def genome_to_dag(genome) -> RepresentationDag: + builder = RepresentationDAGBuilder() + leaf_ids = [ + builder.create_leaf_node(mod_id, repr_idx) for mod_id, repr_idx in genome.leaves + ] + + def build(subtree: Tree, path: str) -> str: + if isinstance(subtree, int): + return leaf_ids[subtree] + left, right = subtree + left_id = build(left, path + "L") + right_id = build(right, path + "R") + op_cls = genome.fusion_ops[path] + op = op_cls() + return builder.create_operation_node( + op.__class__, [left_id, right_id], op.get_current_parameters() + ) + + return builder.build(build(genome.tree, "")) + + +def _evaluate_individual_worker( + dag_pickle: bytes, + task_pickle: bytes, + modalities_pickle: bytes, + metric_name: str, +) -> Tuple[float, Dict[str, Any]]: + dag = pickle.loads(dag_pickle) + task = pickle.loads(task_pickle) + modalities = pickle.loads(modalities_pickle) + + start_time = time.time() + fused_representation = dag.execute(modalities, task) + scores = task.run(fused_representation.data) + runtime = time.time() - start_time + fitness = scores[1].average_scores[metric_name] + + objective = { + "train_score": scores[0].average_scores, + "val_score": scores[1].average_scores, + "test_score": scores[2].average_scores, + "val": fitness, + "runtime": runtime, + "representation_time": runtime, + "task_time": 0.0, + } + return fitness, objective + + +@dataclass +class Individual: + # Leaves are concrete unimodal choices: (modality_id, representation_index) + leaves: List[Tuple[str, int]] + # Binary tree over leaf indices in `leaves` + tree: Tree + # Fusion op class per internal tree-path key, e.g. "L", "RLL", ... + fusion_ops: Dict[str, Any] + fitness: float = float("-inf") + objective: Dict[str, float] = field(default_factory=dict) + dag: Optional[RepresentationDag] = None + + +# ---------------------------- +# Optional helper abstraction +# ---------------------------- + + +class MutationOperator: + name: str = "base" + + def __call__( + self, + individual: Individual, + rng: random.Random, + context: Dict[str, Any], + ) -> Individual: + raise NotImplementedError() + + +def _collect_internal_paths(subtree: Tree, path: str = "") -> List[str]: + if isinstance(subtree, int): + return [] + left, right = subtree + return ( + [path] + + _collect_internal_paths(left, path + "L") + + _collect_internal_paths(right, path + "R") + ) + + +def _get_subtree(subtree: Tree, target_path: str) -> Tree: + if target_path == "": + return copy.deepcopy(subtree) + if isinstance(subtree, int): + raise ValueError(f"Path '{target_path}' does not exist in leaf subtree") + left, right = subtree + if target_path[0] == "L": + return _get_subtree(left, target_path[1:]) + return _get_subtree(right, target_path[1:]) + + +def _replace_subtree(subtree: Tree, target_path: str, replacement: Tree) -> Tree: + if target_path == "": + return copy.deepcopy(replacement) + if isinstance(subtree, int): + raise ValueError(f"Path '{target_path}' does not exist in leaf subtree") + left, right = subtree + if target_path[0] == "L": + return ( + _replace_subtree(left, target_path[1:], replacement), + copy.deepcopy(right), + ) + return (copy.deepcopy(left), _replace_subtree(right, target_path[1:], replacement)) + + +def _collect_leaf_indices(subtree: Tree) -> List[int]: + if isinstance(subtree, int): + return [subtree] + left, right = subtree + return _collect_leaf_indices(left) + _collect_leaf_indices(right) + + +def _sample_random_binary_tree_from_leaves( + leaf_indices: List[int], rng: random.Random +) -> Tree: + nodes: List[Tree] = list(leaf_indices) + while len(nodes) > 1: + i, j = rng.sample(range(len(nodes)), 2) + a = nodes.pop(max(i, j)) + b = nodes.pop(min(i, j)) + nodes.append((a, b)) + return nodes[0] + + +def _reindex_tree(subtree: Tree, index_map: Dict[int, int]) -> Tree: + if isinstance(subtree, int): + return index_map[subtree] + left, right = subtree + return (_reindex_tree(left, index_map), _reindex_tree(right, index_map)) + + +def _remove_leaf_from_tree(subtree: Tree, leaf_idx: int) -> Optional[Tree]: + if isinstance(subtree, int): + if subtree == leaf_idx: + return None + return subtree + + left, right = subtree + new_left = _remove_leaf_from_tree(left, leaf_idx) + new_right = _remove_leaf_from_tree(right, leaf_idx) + + if new_left is None: + return new_right + if new_right is None: + return new_left + return (new_left, new_right) + + +def _path_is_under_prefix(path: str, prefix: str) -> bool: + return prefix == "" or path == prefix or path.startswith(prefix) + + +def _rebuild_fusion_ops( + subtree: Tree, + existing_fusion_ops: Dict[str, Any], + rng: random.Random, + fusion_operators: List[Any], + randomized_prefixes: Optional[List[str]] = None, +) -> Dict[str, Any]: + fusion_ops: Dict[str, Any] = {} + randomized_prefixes = randomized_prefixes or [] + for path in _collect_internal_paths(subtree): + if any(_path_is_under_prefix(path, prefix) for prefix in randomized_prefixes): + fusion_ops[path] = rng.choice(fusion_operators) + elif path in existing_fusion_ops: + fusion_ops[path] = existing_fusion_ops[path] + else: + fusion_ops[path] = rng.choice(fusion_operators) + return fusion_ops + + +# ---------------------------- +# GA Optimizer skeleton +# ---------------------------- + + +class MultimodalGAOptimizer: + """ + GA-based multimodal optimizer skeleton. + Keeps constructor inputs compatible with original MultimodalOptimizer. + """ + + def __init__( + self, + modalities: List[Any], + unimodal_optimization_results: Any, + tasks: List[Any], + k: int = 2, + debug: bool = False, + min_modalities: int = 2, + max_modalities: int = None, + metric: str = "accuracy", + checkpoint_every: int = None, + resume: bool = True, + # --- GA controls (new) --- + population_size: int = 32, + generations: int = 20, + elite_size: int = 4, + tournament_size: int = 3, + crossover_rate: float = 0.9, + mutation_rate: float = 0.3, + random_seed: int = 42, + early_stopping_patience: Optional[int] = 5, + early_stopping_min_delta: float = 1e-6, + ): + self.modalities = modalities + self.tasks = tasks + self.k = k + self.debug = debug + if DEBUG: + self.debug = True + self.metric_name = metric + self.min_modalities = max(2, min_modalities) + self.max_modalities = max_modalities or len(modalities) + + self.population_size = population_size + self.generations = generations + self.elite_size = elite_size + self.tournament_size = tournament_size + self.crossover_rate = crossover_rate + self.mutation_rate = mutation_rate + self.random_seed = random_seed + self.rng = random.Random(random_seed) + self.early_stopping_patience = early_stopping_patience + self.early_stopping_min_delta = early_stopping_min_delta + + self.operator_registry = Registry() + self.fusion_operators = self.operator_registry.get_fusion_operators() + + self.k_best_representations = self._extract_k_best_representations( + unimodal_optimization_results + ) + + self.optimization_results: Dict[str, List[OptimizationResult]] = {} + self._seen_results_by_task: Dict[str, Dict[str, OptimizationResult]] = {} + self._fitness_cache_by_task: Dict[ + str, Dict[Tuple[Any, ...], Tuple[float, Dict[str, Any]]] + ] = {} + self._stats_by_task: Dict[str, Dict[str, int]] = {} + self._checkpoint_manager = CheckpointManager( + ".", + "multimodal_ga_checkpoint_", + checkpoint_every=checkpoint_every, + resume=resume, + ) + + # Register mutation ops you want enabled + self.mutation_ops: List[MutationOperator] = [ + MutateFusionOperator(), + MutateRepresentationIndex(), + MutateTreeRotation(), + MutateSubtreeResample(), + # Optional global search mutations: + MutateAddModality(), + MutateDropModality(), + ] + + # ---------------------------- + # Public entrypoint + # ---------------------------- + + def optimize( + self, + max_evaluations_per_task: Optional[int] = None, + # If you later wire this to NodeExecutor, add params like: + # use_node_executor: bool = True, + # max_workers: int = ... + ) -> Dict[str, List[OptimizationResult]]: + self.rng = random.Random(self.random_seed) + self._resume_if_available() + self._ensure_results_initialized() + + for task in self.tasks: + task_name = task.model.name + if self.debug: + print(f"[GA] Task={task_name} initialization") + + # 1) Initialize population stochastically + population = self._initialize_population(task_name) + + eval_budget = 0 + best_seen: Optional[Individual] = None + no_improvement_gens = 0 + + # 2) Evolution loop + for gen in range(self.generations): + if self.debug: + print(f"[GA] Task={task_name}, generation={gen}") + + # Evaluate all individuals (or only unevaluated) + self._evaluate_population(population, task) + + # Track best + gen_best = max(population, key=lambda ind: ind.fitness) + if best_seen is None or ( + gen_best.fitness > best_seen.fitness + self.early_stopping_min_delta + ): + best_seen = copy.deepcopy(gen_best) + no_improvement_gens = 0 + else: + no_improvement_gens += 1 + + # Optional budget stop + eval_budget += self._stats_by_task[task_name][ + "last_generation_executed" + ] + if ( + max_evaluations_per_task is not None + and eval_budget >= max_evaluations_per_task + ): + break + if ( + self.early_stopping_patience is not None + and no_improvement_gens >= self.early_stopping_patience + ): + if self.debug: + print( + f"[GA] Task={task_name} early stopping after " + f"{no_improvement_gens} no-improvement generations" + ) + break + + # Selection + variation + replacement + next_population = self._elitism(population) + seen_genotypes = { + self._individual_signature(ind) for ind in next_population + } + duplicate_retry_budget = self.population_size * 10 + duplicate_retries = 0 + while len(next_population) < self.population_size: + p1 = self._tournament_select(population) + p2 = self._tournament_select(population) + + if self.rng.random() < self.crossover_rate: + c1, c2 = self._crossover(p1, p2) + else: + c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) + + c1 = self._mutate(c1, task_name) + c2 = self._mutate(c2, task_name) + + added_c1 = self._append_if_unique( + next_population, c1, seen_genotypes, self.population_size + ) + added_c2 = self._append_if_unique( + next_population, c2, seen_genotypes, self.population_size + ) + if not added_c1 and not added_c2: + duplicate_retries += 1 + else: + duplicate_retries = 0 + + if duplicate_retries >= duplicate_retry_budget: + break + + while len(next_population) < self.population_size: + immigrant = self._sample_random_individual(task_name) + self._append_if_unique( + next_population, immigrant, seen_genotypes, self.population_size + ) + + population = next_population + + self._checkpoint_manager.checkpoint_if_due( + self.optimization_results, "eval_count_by_task" + ) + + if self.debug and best_seen is not None: + print(f"[GA] Task={task_name}, best_fitness={best_seen.fitness:.6f}") + + return self.optimization_results + + def optimize_parallel( + self, + max_combinations: Optional[int] = None, + max_workers: int = 2, + batch_size: int = 8, + ) -> Dict[str, List[OptimizationResult]]: + self.rng = random.Random(self.random_seed) + self._resume_if_available() + self._ensure_results_initialized() + + for task in self.tasks: + task_name = task.model.name + if self.debug: + print(f"[GA-P] Task={task_name} initialization") + + population = self._initialize_population(task_name) + eval_budget = 0 + best_seen: Optional[Individual] = None + no_improvement_gens = 0 + + for gen in range(self.generations): + if self.debug: + print(f"[GA-P] Task={task_name}, generation={gen}") + + self._evaluate_population_parallel( + population, + task, + max_workers=max_workers, + batch_size=max(1, batch_size), + ) + + gen_best = max(population, key=lambda ind: ind.fitness) + if best_seen is None or ( + gen_best.fitness > best_seen.fitness + self.early_stopping_min_delta + ): + best_seen = copy.deepcopy(gen_best) + no_improvement_gens = 0 + else: + no_improvement_gens += 1 + + eval_budget += self._stats_by_task[task_name][ + "last_generation_executed" + ] + if max_combinations is not None and eval_budget >= max_combinations: + break + if ( + self.early_stopping_patience is not None + and no_improvement_gens >= self.early_stopping_patience + ): + if self.debug: + print( + f"[GA-P] Task={task_name} early stopping after " + f"{no_improvement_gens} no-improvement generations" + ) + break + + next_population = self._elitism(population) + seen_genotypes = { + self._individual_signature(ind) for ind in next_population + } + duplicate_retry_budget = self.population_size * 10 + duplicate_retries = 0 + while len(next_population) < self.population_size: + p1 = self._tournament_select(population) + p2 = self._tournament_select(population) + + if self.rng.random() < self.crossover_rate: + c1, c2 = self._crossover(p1, p2) + else: + c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) + + c1 = self._mutate(c1, task_name) + c2 = self._mutate(c2, task_name) + + added_c1 = self._append_if_unique( + next_population, c1, seen_genotypes, self.population_size + ) + added_c2 = self._append_if_unique( + next_population, c2, seen_genotypes, self.population_size + ) + if not added_c1 and not added_c2: + duplicate_retries += 1 + else: + duplicate_retries = 0 + + if duplicate_retries >= duplicate_retry_budget: + break + + while len(next_population) < self.population_size: + immigrant = self._sample_random_individual(task_name) + self._append_if_unique( + next_population, immigrant, seen_genotypes, self.population_size + ) + + population = next_population + + self._checkpoint_manager.checkpoint_if_due( + self.optimization_results, "eval_count_by_task" + ) + + if self.debug and best_seen is not None: + print(f"[GA-P] Task={task_name}, best_fitness={best_seen.fitness:.6f}") + + return self.optimization_results + + # ---------------------------- + # Initialization + # ---------------------------- + + def _initialize_population(self, task_name: str) -> List[Individual]: + population = [] + seen_signatures = set() + retry_budget = self.population_size * 10 + retries = 0 + for _ in range(self.population_size): + while retries < retry_budget: + candidate = self._sample_random_individual(task_name) + signature = self._individual_signature(candidate) + if signature not in seen_signatures: + seen_signatures.add(signature) + population.append(candidate) + break + retries += 1 + if retries >= retry_budget: + # Fall back to potentially duplicated random individuals rather than + # stalling initialization on very small search spaces. + population.append(self._sample_random_individual(task_name)) + return population + + def _sample_random_individual(self, task_name: str) -> Individual: + leaves = self._sample_leaf_set(task_name) + tree = self._sample_random_binary_tree(len(leaves)) + fusion_ops = {} + self._assign_random_fusion_ops(tree, fusion_ops, path="") + return Individual(leaves=leaves, tree=tree, fusion_ops=fusion_ops) + + def _sample_leaf_set(self, task_name: str) -> List[Tuple[str, int]]: + # Only sample from modalities that have at least one representation + # for the current task. + task_reps = self.k_best_representations.get(task_name, {}) + available_modalities = [ + m.modality_id + for m in self.modalities + if len(task_reps.get(m.modality_id, [])) > 0 + ] + + if len(available_modalities) < 2: + raise ValueError( + f"Need at least 2 modalities with non-empty representations for task " + f"'{task_name}', found {len(available_modalities)}." + ) + + # Clamp sampling range to available modalities. + lower = max(2, self.min_modalities) + upper = min(self.max_modalities, len(available_modalities)) + if lower > upper: + lower = upper + + r = self.rng.randint(lower, upper) + chosen_modalities = self.rng.sample(available_modalities, r) + + # Sample one representation index from top-k for each chosen modality. + return [ + (mod_id, self.rng.randrange(len(task_reps[mod_id]))) + for mod_id in chosen_modalities + ] + + def _sample_random_binary_tree(self, n_leaves: int) -> Tree: + nodes: List[Tree] = list(range(n_leaves)) + while len(nodes) > 1: + i, j = self.rng.sample(range(len(nodes)), 2) + a = nodes.pop(max(i, j)) + b = nodes.pop(min(i, j)) + nodes.append((a, b)) + return nodes[0] + + def _assign_random_fusion_ops( + self, subtree: Tree, fusion_ops: Dict[str, Any], path: str + ): + if isinstance(subtree, int): + return + fusion_ops[path] = self.rng.choice(self.fusion_operators) + left, right = subtree + self._assign_random_fusion_ops(left, fusion_ops, path + "L") + self._assign_random_fusion_ops(right, fusion_ops, path + "R") + + # ---------------------------- + # Evaluation + # ---------------------------- + + def _evaluate_population(self, population: List[Individual], task: Task) -> None: + # Placeholder: currently serial + per-individual eval. + # You can batch DAGs and evaluate with NodeExecutor later. + task_name = task.model.name + cache = self._fitness_cache_by_task.setdefault(task_name, {}) + evals_before = self._stats_by_task[task_name]["executed_evaluations"] + for ind in population: + # Already evaluated and unchanged. + if ind.fitness != float("-inf"): + continue + + signature = self._individual_signature(ind) + cached = cache.get(signature) + if cached is not None: + self._stats_by_task[task_name]["cache_hits"] += 1 + ind.fitness, ind.objective = cached[0], dict(cached[1]) + continue + + dag = self._individual_to_dag(ind) + ind.dag = dag + modalities = [ + self.k_best_representations[task_name][mod_id][repr_idx] + for mod_id, repr_idx in ind.leaves + ] + self._stats_by_task[task_name]["executed_evaluations"] += 1 + fitness, objective = self._evaluate_individual(dag, task, modalities) + ind.fitness = fitness + ind.objective = objective + cache[signature] = (fitness, dict(objective)) + self._record_result(task_name, dag, objective) + self._stats_by_task[task_name]["last_generation_executed"] = ( + self._stats_by_task[task_name]["executed_evaluations"] - evals_before + ) + + def _evaluate_population_parallel( + self, + population: List[Individual], + task: Task, + max_workers: int, + batch_size: int, + ) -> None: + ctx = mp.get_context("spawn") + task_name = task.model.name + cache = self._fitness_cache_by_task.setdefault(task_name, {}) + evals_before = self._stats_by_task[task_name]["executed_evaluations"] + task_pickle = pickle.dumps(copy.deepcopy(task)) + futures = {} + pending_followers: Dict[Tuple[Any, ...], List[Individual]] = {} + + def _collect_ready(done_futures): + for done in done_futures: + ind, dag, signature = futures.pop(done) + fitness, objective = done.result() + ind.fitness = fitness + ind.objective = objective + cache[signature] = (fitness, dict(objective)) + self._record_result(task_name, dag, objective) + for follower in pending_followers.pop(signature, []): + follower.fitness = fitness + follower.objective = dict(objective) + + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: + for ind in population: + # Already evaluated and unchanged. + if ind.fitness != float("-inf"): + continue + + signature = self._individual_signature(ind) + cached = cache.get(signature) + if cached is not None: + self._stats_by_task[task_name]["cache_hits"] += 1 + ind.fitness, ind.objective = cached[0], dict(cached[1]) + continue + + if signature in pending_followers: + pending_followers[signature].append(ind) + self._stats_by_task[task_name]["pending_signature_hits"] += 1 + continue + + dag = self._individual_to_dag(ind) + ind.dag = dag + modalities = [ + self.k_best_representations[task_name][mod_id][repr_idx] + for mod_id, repr_idx in ind.leaves + ] + + fut = executor.submit( + _evaluate_individual_worker, + pickle.dumps(dag), + task_pickle, + pickle.dumps(modalities), + self.metric_name, + ) + futures[fut] = (ind, dag, signature) + pending_followers[signature] = [] + self._stats_by_task[task_name]["executed_evaluations"] += 1 + + if len(futures) >= batch_size: + done, _ = wait(set(futures.keys()), return_when=FIRST_COMPLETED) + _collect_ready(done) + + if futures: + done, _ = wait(set(futures.keys())) + _collect_ready(done) + self._stats_by_task[task_name]["last_generation_executed"] = ( + self._stats_by_task[task_name]["executed_evaluations"] - evals_before + ) + + def _evaluate_individual( + self, dag: RepresentationDag, task: Task, modalities: List[Modality] + ) -> Tuple[float, Dict[str, float]]: + start_time = time.time() + fused_representation = dag.execute(modalities, task) + task_start_time = time.time() + scores = task.run(fused_representation.data) + task_end_time = time.time() + runtime = time.time() - start_time + fitness = scores[1].average_scores[self.metric_name] + objective = { + "train_score": scores[0].average_scores, + "val_score": scores[1].average_scores, + "test_score": scores[2].average_scores, + "val": fitness, + "runtime": runtime, + "representation_time": fused_representation.transform_time, + "task_time": task_end_time - task_start_time, + } + return fitness, objective + + def _record_result( + self, task_name: str, dag: RepresentationDag, objective: Dict[str, Any] + ) -> None: + key = str(dag.compute_full_node_signature(dag.root_node_id)) + if key in self._seen_results_by_task[task_name]: + return + result = OptimizationResult( + dag=dag, + val_score=objective.get("val_score", {}), + train_score=objective.get("train_score", {}), + test_score=objective.get("test_score", {}), + runtime=objective.get("runtime", 0.0), + representation_time=objective.get("representation_time", 0.0), + task_time=objective.get("task_time", 0.0), + task_name=task_name, + ) + self._seen_results_by_task[task_name][key] = result + self.optimization_results[task_name].append(result) + + def get_task_stats(self, task_name: str) -> Dict[str, int]: + return dict(self._stats_by_task.get(task_name, {})) + + def _individual_signature(self, ind: Individual) -> Tuple[Any, ...]: + def _normalize_tree(subtree: Tree) -> Any: + if isinstance(subtree, int): + return subtree + left, right = subtree + return (_normalize_tree(left), _normalize_tree(right)) + + leaves_sig = tuple(ind.leaves) + tree_sig = _normalize_tree(ind.tree) + ops_sig = tuple( + sorted((path, op.__name__) for path, op in ind.fusion_ops.items()) + ) + return leaves_sig, tree_sig, ops_sig + + def _append_if_unique( + self, + population: List[Individual], + candidate: Individual, + seen_genotypes: set, + target_size: int, + ) -> bool: + if len(population) >= target_size: + return False + sig = self._individual_signature(candidate) + if sig in seen_genotypes: + return False + seen_genotypes.add(sig) + population.append(candidate) + return True + + # ---------------------------- + # Selection / crossover / mutation + # ---------------------------- + + def _elitism(self, population: List[Individual]) -> List[Individual]: + ranked = sorted(population, key=lambda x: x.fitness, reverse=True) + return [copy.deepcopy(ind) for ind in ranked[: self.elite_size]] + + def _tournament_select(self, population: List[Individual]) -> Individual: + contestants = self.rng.sample( + population, k=min(self.tournament_size, len(population)) + ) + return copy.deepcopy(max(contestants, key=lambda x: x.fitness)) + + def _crossover( + self, p1: Individual, p2: Individual + ) -> Tuple[Individual, Individual]: + """ + Skeleton crossover: + - Safe version assumes compatible leaves (same modality set/order). + - If incompatible, fallback to op-only crossover. + """ + c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) + + if self._compatible_for_subtree_crossover(c1, c2): + c1_paths = _collect_internal_paths(c1.tree) + c2_paths = _collect_internal_paths(c2.tree) + if c1_paths and c2_paths: + path1 = self.rng.choice(c1_paths) + path2 = self.rng.choice(c2_paths) + subtree1 = _get_subtree(c1.tree, path1) + subtree2 = _get_subtree(c2.tree, path2) + c1.tree = _replace_subtree(c1.tree, path1, subtree2) + c2.tree = _replace_subtree(c2.tree, path2, subtree1) + c1.fusion_ops = _rebuild_fusion_ops( + c1.tree, + c1.fusion_ops, + self.rng, + self.fusion_operators, + randomized_prefixes=[path1], + ) + c2.fusion_ops = _rebuild_fusion_ops( + c2.tree, + c2.fusion_ops, + self.rng, + self.fusion_operators, + randomized_prefixes=[path2], + ) + else: + # op-only crossover: mix operator assignments + all_keys = set(c1.fusion_ops.keys()) | set(c2.fusion_ops.keys()) + for k in all_keys: + if self.rng.random() < 0.5: + if k in c2.fusion_ops: + c1.fusion_ops[k] = c2.fusion_ops[k] + else: + if k in c1.fusion_ops: + c2.fusion_ops[k] = c1.fusion_ops[k] + + # invalidate stale eval + c1.fitness, c1.objective, c1.dag = float("-inf"), {}, None + c2.fitness, c2.objective, c2.dag = float("-inf"), {}, None + return c1, c2 + + def _compatible_for_subtree_crossover(self, a: Individual, b: Individual) -> bool: + return [m for m, _ in a.leaves] == [m for m, _ in b.leaves] and len( + a.leaves + ) == len(b.leaves) + + def _mutate(self, ind: Individual, task_name: str) -> Individual: + out = copy.deepcopy(ind) + if self.rng.random() >= self.mutation_rate: + return out + + op = self.rng.choice(self.mutation_ops) + context = { + "task_name": task_name, + "fusion_operators": self.fusion_operators, + "k_best_representations": self.k_best_representations, + "min_modalities": self.min_modalities, + "max_modalities": self.max_modalities, + "all_modalities": [m.modality_id for m in self.modalities], + } + out = op(out, self.rng, context) + out.fitness, out.objective, out.dag = float("-inf"), {}, None + return out + + # ---------------------------- + # Genome -> DAG + # ---------------------------- + + def _individual_to_dag(self, ind: Individual) -> RepresentationDag: + builder = RepresentationDAGBuilder() + leaf_ids = [] + for modality_id, repr_idx in ind.leaves: + # Leaves reference the cached transformed modality directly. + # We collapse the unimodal path from raw->representation here, so + # downstream execution only needs to run multimodal fusion nodes. + leaf_ids.append(builder.create_leaf_node(modality_id, repr_idx)) + + def build(subtree: Tree, path: str) -> str: + if isinstance(subtree, int): + return leaf_ids[subtree] + left, right = subtree + left_id = build(left, path + "L") + right_id = build(right, path + "R") + op_cls = ind.fusion_ops[path] + op = op_cls() + return builder.create_operation_node( + op.__class__, [left_id, right_id], op.get_current_parameters() + ) + + root_id = build(ind.tree, "") + dag = builder.build(root_id) + return self._collapse_cached_unimodal_nodes(dag) + + def _collapse_cached_unimodal_nodes( + self, dag: RepresentationDag + ) -> RepresentationDag: + """ + Remove unary unimodal nodes that originate at leaves. + + In GA multimodal search, leaf inputs are already transformed unimodal + representations taken from the unimodal optimizer cache. Any unary chain + attached to those leaves is redundant and can be bypassed to ensure only + multimodal (fusion) operations are executed. + """ + node_by_id = {node.node_id: copy.deepcopy(node) for node in dag.nodes} + if dag.root_node_id not in node_by_id: + return dag + + changed = True + while changed: + changed = False + node_ids = list(node_by_id.keys()) + for node_id in node_ids: + if node_id not in node_by_id: + continue + node = node_by_id[node_id] + if len(node.inputs) != 1: + continue + + parent_id = node.inputs[0] + parent = node_by_id.get(parent_id) + if parent is None: + continue + if parent.inputs: + continue + + # Bypass unary node by rewiring all consumers to the leaf input. + for consumer in node_by_id.values(): + consumer.inputs = [ + parent_id if input_id == node_id else input_id + for input_id in consumer.inputs + ] + if dag.root_node_id == node_id: + dag.root_node_id = parent_id + del node_by_id[node_id] + changed = True + + return RepresentationDag( + list(node_by_id.values()), dag.root_node_id, dag.dag_id + ) + + # ---------------------------- + # Existing helper compatibility + # ---------------------------- + + def _extract_k_best_representations( + self, unimodal_optimization_results: Any + ) -> Dict[str, Dict[str, List[Any]]]: + k_best = {} + for task in self.tasks: + task_name = task.model.name + k_best[task_name] = {} + for modality in self.modalities: + _, cached_data = unimodal_optimization_results.get_k_best_results( + modality, task, self.metric_name + ) + k_best[task_name][modality.modality_id] = cached_data + return k_best + + def _resume_if_available(self) -> None: + loaded = self._checkpoint_manager.resume_from_checkpoint( + "eval_count_by_task", + lambda results: { + t.model.name: len(results.get(t.model.name, [])) for t in self.tasks + }, + ) + if loaded: + results, _, _ = loaded + self.optimization_results = results + + def _ensure_results_initialized(self): + if not isinstance(self.optimization_results, dict): + self.optimization_results = {} + for task in self.tasks: + task_name = task.model.name + self.optimization_results.setdefault(task_name, []) + seen = self._seen_results_by_task.setdefault(task_name, {}) + self._fitness_cache_by_task.setdefault(task_name, {}) + self._stats_by_task.setdefault( + task_name, + { + "executed_evaluations": 0, + "last_generation_executed": 0, + "cache_hits": 0, + "pending_signature_hits": 0, + }, + ) + if self.optimization_results[task_name]: + for result in self.optimization_results[task_name]: + if result.dag is None: + continue + key = str( + result.dag.compute_full_node_signature(result.dag.root_node_id) + ) + seen.setdefault(key, result) + + def _to_optimization_results( + self, individuals: List[Individual], task_name: str + ) -> List[OptimizationResult]: + out = [] + for ind in individuals: + # Keep consistent with old output shape: OptimizationResult list per task + out.append( + OptimizationResult( + dag=ind.dag, + val_score={ + self.metric_name: ind.objective.get("val", float("-inf")) + }, + train_score={}, + test_score={}, + runtime=ind.objective.get("runtime", 0.0), + representation_time=ind.objective.get("representation_time", 0.0), + task_time=ind.objective.get("task_time", 0.0), + task_name=task_name, + ) + ) + return out + + +# ---------------------------- +# Mutation operators (outline) +# ---------------------------- + + +class MutateFusionOperator(MutationOperator): + """ + Change one internal fusion op class. + """ + + name = "mutate_fusion_operator" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + if not out.fusion_ops: + return out + key = rng.choice(list(out.fusion_ops.keys())) + current = out.fusion_ops[key] + candidates = [op for op in context["fusion_operators"] if op != current] + if candidates: + out.fusion_ops[key] = rng.choice(candidates) + return out + + +class MutateRepresentationIndex(MutationOperator): + """ + Keep modality set fixed, change repr_idx for one modality leaf. + """ + + name = "mutate_representation_index" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + if not out.leaves: + return out + i = rng.randrange(len(out.leaves)) + modality_id, cur_idx = out.leaves[i] + task_name = context["task_name"] + reps = context["k_best_representations"][task_name][modality_id] + if len(reps) <= 1: + return out + new_idx = rng.randrange(len(reps)) + while new_idx == cur_idx: + new_idx = rng.randrange(len(reps)) + out.leaves[i] = (modality_id, new_idx) + return out + + +class MutateTreeRotation(MutationOperator): + """ + Local reassociation (e.g., ((a,b),c) <-> (a,(b,c))) on a random eligible subtree. + """ + + name = "mutate_tree_rotation" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + candidates: List[Tuple[str, str]] = [] + for path in _collect_internal_paths(out.tree): + subtree = _get_subtree(out.tree, path) + if isinstance(subtree, int): + continue + left, right = subtree + if not isinstance(left, int): + candidates.append((path, "left")) + if not isinstance(right, int): + candidates.append((path, "right")) + + if not candidates: + return out + + path, direction = rng.choice(candidates) + subtree = _get_subtree(out.tree, path) + left, right = subtree + + if direction == "left": + left_left, left_right = left + rotated = (left_left, (left_right, right)) + else: + right_left, right_right = right + rotated = ((left, right_left), right_right) + + out.tree = _replace_subtree(out.tree, path, rotated) + out.fusion_ops = _rebuild_fusion_ops( + out.tree, + out.fusion_ops, + rng, + context["fusion_operators"], + randomized_prefixes=[path], + ) + return out + + +class MutateSubtreeResample(MutationOperator): + """ + Pick a random internal subtree, keep the same leaves, but resample: + - subtree topology + - fusion operators inside subtree + """ + + name = "mutate_subtree_resample" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + internal_paths = _collect_internal_paths(out.tree) + if not internal_paths: + return out + + path = rng.choice(internal_paths) + subtree = _get_subtree(out.tree, path) + leaf_indices = _collect_leaf_indices(subtree) + if len(leaf_indices) < 2: + return out + + new_subtree = _sample_random_binary_tree_from_leaves(leaf_indices, rng) + out.tree = _replace_subtree(out.tree, path, new_subtree) + out.fusion_ops = _rebuild_fusion_ops( + out.tree, + out.fusion_ops, + rng, + context["fusion_operators"], + randomized_prefixes=[path], + ) + return out + + +class MutateAddModality(MutationOperator): + """ + Optional global mutation: + Add one new modality (if below max_modalities), choose repr_idx, fuse with current root. + """ + + name = "mutate_add_modality" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + if len(out.leaves) >= context["max_modalities"]: + return out + + existing = {m for m, _ in out.leaves} + candidates = [m for m in context["all_modalities"] if m not in existing] + if not candidates: + return out + + m = rng.choice(candidates) + task_name = context["task_name"] + reps = context["k_best_representations"][task_name][m] + if len(reps) == 0: + return out + + repr_idx = rng.randrange(len(reps)) + new_leaf_idx = len(out.leaves) + out.leaves.append((m, repr_idx)) + + # Wrap old tree with new root fusion + old_tree = out.tree + old_fusion_ops = copy.deepcopy(out.fusion_ops) + out.tree = (old_tree, new_leaf_idx) + out.fusion_ops = {"": rng.choice(context["fusion_operators"])} + for path, op in old_fusion_ops.items(): + out.fusion_ops["L" + path] = op + return out + + +class MutateDropModality(MutationOperator): + """ + Optional global mutation: + Remove one modality leaf (if above min_modalities) and collapse tree. + """ + + name = "mutate_drop_modality" + + def __call__( + self, individual: Individual, rng: random.Random, context: Dict[str, Any] + ) -> Individual: + out = copy.deepcopy(individual) + if len(out.leaves) <= context["min_modalities"]: + return out + + leaf_idx = rng.randrange(len(out.leaves)) + new_tree = _remove_leaf_from_tree(out.tree, leaf_idx) + if new_tree is None: + return out + + remaining_leaves = [leaf for i, leaf in enumerate(out.leaves) if i != leaf_idx] + index_map = {} + next_idx = 0 + for old_idx in range(len(out.leaves)): + if old_idx == leaf_idx: + continue + index_map[old_idx] = next_idx + next_idx += 1 + + out.leaves = remaining_leaves + out.tree = _reindex_tree(new_tree, index_map) + out.fusion_ops = _rebuild_fusion_ops( + out.tree, + out.fusion_ops, + rng, + context["fusion_operators"], + ) + return out From ca7ba60c0960a36bb62d693f2b6196b93e7afff5 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 13:32:36 +0200 Subject: [PATCH 2/8] consolidate ga files --- .../drsearch/multimodal_ga_deap_optimizer.py | 1277 ----------- .../scuro/drsearch/multimodal_ga_optimizer.py | 1926 ++++++++--------- 2 files changed, 857 insertions(+), 2346 deletions(-) delete mode 100644 src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py deleted file mode 100644 index 1b8c2223ee7..00000000000 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_deap_optimizer.py +++ /dev/null @@ -1,1277 +0,0 @@ -# ------------------------------------------------------------- -# -# 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 __future__ import annotations -import copy -import multiprocessing as mp -import os -import pickle -import random -import tempfile -import time -import traceback -from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait -from dataclasses import dataclass, field -from itertools import chain -from typing import Any, Dict, List, Optional, Tuple -from deap import base, creator, tools -from systemds.scuro.drsearch.multimodal_ga_optimizer import ( - _collect_internal_paths, - _get_subtree, - _replace_subtree, - _rebuild_fusion_ops, - _remove_leaf_from_tree, - _reindex_tree, - _collect_leaf_indices, -) -from systemds.scuro.drsearch.operator_registry import Registry -from systemds.scuro.drsearch.representation_dag import ( - RepresentationDAGBuilder, - RepresentationDag, -) -from systemds.scuro.drsearch.task import Task -from systemds.scuro.representations.aggregated_representation import ( - AggregatedRepresentation, -) -from systemds.scuro.utils.schema_helpers import get_shape - -Tree = int | Tuple["Tree", "Tree"] - - -@dataclass -class FusionSearchResult: - dag: RepresentationDag - train_score: dict - val_score: dict - test_score: dict - runtime: float = 0.0 - task_time: float = 0.0 - representation_time: float = 0.0 - task_name: str = "" - - # Per-fold score vectors, {metric: [fold0, ...]}. Needed for error bars and - # for any selection rule that has to tell a consistent winner apart from a - # lucky-fold one; the averaged dicts above cannot. - val_fold_scores: dict = field(default_factory=dict) - train_fold_scores: dict = field(default_factory=dict) - test_fold_scores: dict = field(default_factory=dict) - - # Fit/inference wall-clock from the fold loop, per fold and averaged. - task_timing: dict = field(default_factory=dict) - - # Search trajectory: generation, position in the evaluation sequence, and - # seconds since this GA run began. Results are appended in evaluation - # order, but nothing recorded *when* -- which is what a best-so-far curve - # against a wall-clock budget needs. - generation: int = -1 - eval_index: int = -1 - t_since_search_start_s: float = 0.0 - t_eval_end_unix: float = 0.0 - - -@dataclass -class DagGenome: - leaves: List[Tuple[str, int]] - tree: Tree - fusion_ops: Dict[str, type] - - -# ---------------------------------------------------------------------- -# Module-level evaluation helpers. -# -# These are intentionally free functions (not bound methods) so that -# `_evaluate_dag_worker` can be shipped to a separate process and pickled -# safely. `_evaluate_genome_body` holds the actual logic and is shared by -# both the in-process (serial) and cross-process (parallel) evaluation -# paths, so the two can never silently diverge in behavior. -# ---------------------------------------------------------------------- - - -# Objectives that come from evaluation timing rather than from the task's -# score dict (scores[i].average_scores). Anything not listed here is looked -# up as a key in the validation score dict (e.g. "accuracy", "f1"). -_TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"} - -ObjectiveSpec = Tuple[str, str] # (name, "max" | "min") - - -def _objective_value( - name: str, val_score: Dict[str, float], timing: Dict[str, float] -) -> float: - if name in _TIMING_OBJECTIVES: - return timing[name] - return val_score[name] - - -def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: - """ - The worst possible value for each objective, direction-aware: -inf for a - "max" objective, +inf for a "min" objective (e.g. runtime), so that a - failed evaluation always ranks last regardless of DEAP's per-objective - fitness weight sign (wvalue = raw_value * weight). - """ - return tuple( - float("-inf") if direction == "max" else float("inf") - for _, direction in objective_specs - ) - - -def _evaluate_genome_body( - dag: RepresentationDag, - task: Task, - modalities: List[Any], - objective_specs: List[ObjectiveSpec], -) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: - start = time.time() - fused = dag.execute(modalities, task, enable_cache=False) - if fused is None: - return None, None - - if isinstance(fused, dict): - fused = fused[list(fused.keys())[-1]] - - if task.expected_dim == 1 and get_shape(fused.metadata) > 1: - fused = AggregatedRepresentation().transform(fused) - - t0 = time.time() - scores = task.run(fused.data) - task_time = time.time() - t0 - total = time.time() - start - - val_score = scores[1].average_scores - timing = { - "runtime": total, - "task_time": task_time, - "representation_time": total - task_time, - } - fitness = tuple( - _objective_value(name, val_score, timing) for name, _ in objective_specs - ) - payload = { - "train_score": scores[0].average_scores, - "val_score": val_score, - "test_score": scores[2].average_scores, - "train_fold_scores": scores[0].fold_scores(), - "val_fold_scores": scores[1].fold_scores(), - "test_fold_scores": scores[2].fold_scores(), - "task_timing": getattr(task, "last_run_timing", {}), - **timing, - } - return fitness, payload - - -def _evaluate_dag_worker( - dag_bytes: bytes, - task_bytes: bytes, - modalities_bytes: bytes, - objective_specs: List[ObjectiveSpec], -) -> Tuple[Tuple[float, ...], Optional[Dict[str, Any]], Optional[str]]: - """ - Runs in a worker process. Never raises: any failure (a shape-incompatible - fusion, an OOM, a bug in a representation) is caught and reported back as - plain, picklable data, so a single bad genome can never take down the - rest of the population's evaluation. - """ - try: - dag = pickle.loads(dag_bytes) - task = pickle.loads(task_bytes) - modalities = pickle.loads(modalities_bytes) - fitness, payload = _evaluate_genome_body(dag, task, modalities, objective_specs) - if fitness is None: - fitness = _failure_fitness(objective_specs) - return fitness, payload, None - except Exception: - return _failure_fitness(objective_specs), None, traceback.format_exc() - - -class MultimodalDeapOptimizer: - """ - Multimodal optimizer that uses DEAP's genome/fitness containers and - selection operators to run a genetic search over fusion DAGs. - - @param modalities: List of modalities to optimize. - @param unimodal_optimization_results: Unimodal optimization results. - @param tasks: List of tasks to optimize. - @param debug: Whether to print debug information. - @param min_modalities: Minimum number of modalities to use. - @param max_modalities: Maximum number of modalities to use. - @param metric: Metric to optimize (used when objectives is not given). - @param objectives: Optional list of (name, direction) pairs to run a - multi-objective (Pareto/NSGA-II) search instead of optimizing a - single scalar metric, e.g. [("accuracy", "max"), ("runtime", - "min")]. "runtime", "task_time", "representation_time" are read - from evaluation timing; any other name is looked up in the task's - validation-score dict (e.g. "accuracy", "f1"). When given, this - overrides metric/maximize_metric. Defaults to None (single - objective, unchanged behavior). - @param population_size: Population size. - @param generations: Maximum number of generations. - @param crossover_probability: Crossover probability. - @param mutation_probability: Mutation probability. - @param random_seed: Random seed. - @param elite_size: Number of top individuals carried over unchanged - (with their fitness preserved) into the next generation. - @param max_workers: Number of worker processes used to evaluate a - generation's population. 1 (default) evaluates serially in-process. - @param batch_size: Maximum number of genome evaluations kept in flight - at once when max_workers > 1. Bounds peak memory/GPU usage - independently of max_workers; defaults to max_workers. - @param early_stopping_patience: Stop once this many consecutive - generations pass without a validation-metric improvement larger - than early_stopping_min_delta. None disables early stopping. - @param early_stopping_min_delta: Minimum improvement to reset the - early-stopping counter. - @param allow_repeated_modalities: Allow one modality to contribute more - than one leaf to a genome, each with a different representation - (intra-modal fusion). Off by default, which keeps the historical - behaviour of at most one leaf per modality and caps the leaf count at - len(modalities). Turning it on raises that cap to the total number of - distinct (modality, representation) leaves available. - @param hall_of_fame_size: How many incumbents to keep per task in - single-objective mode (ignored in multi-objective mode, where the - whole non-dominated front is kept). See get_hall_of_fame. - @param novelty_breeding: Reject offspring whose genome was already - evaluated in an *earlier* generation, not just earlier in the - current one. Without this a converged population refills itself - with re-treads that hit the fitness cache, cost no wall clock and - explore nothing -- on the regret-analysis replays 55-84% of all - fitness requests were such duplicates, so a population of 16 - behaved like a population of ~5. Elites are exempt (they are - carried over before the archive is consulted), and once the - archive covers most of the reachable space the retry budget runs - out and breeding falls back to duplicates, so the search - degrades gracefully rather than stalling. Defaults to True; set - False to reproduce pre-fix runs. - """ - - def __init__( - self, - modalities: List[Any], - unimodal_optimization_results: Any, - tasks: List[Task], - debug: bool = True, - min_modalities: int = 2, - max_modalities: int = None, - metric: str = "accuracy", - objectives: Optional[List[ObjectiveSpec]] = None, - checkpoint_every: int = None, - resume: bool = True, - population_size: int = 32, - generations: int = 20, - crossover_probability: float = 0.7, - mutation_probability: float = 0.4, - random_seed: int = 42, - maximize_metric: bool = True, - elite_size: int = 2, - max_workers: int = 1, - batch_size: Optional[int] = None, - early_stopping_patience: Optional[int] = 5, - early_stopping_min_delta: float = 1e-6, - novelty_breeding: bool = True, - hall_of_fame_size: int = 5, - allow_repeated_modalities: bool = False, - ): - self.modalities = modalities - self.tasks = tasks - self.debug = debug - # A genome may hold several leaves of the same modality (different - # representations of it) only when this is on. With it off, one - # modality contributes at most one leaf, so the leaf count can never - # exceed len(modalities) and max_modalities is clamped to that. - self.allow_repeated_modalities = allow_repeated_modalities - - # min_modalities=1 admits single-leaf (unimodal) genomes. It is not the - # default -- a "multimodal" search that returns a unimodal pipeline is - # almost always a bug, which is why the floor used to be a hard 2 -- - # but it is a legitimate configuration for a sweep that asks what the - # fusion actually buys, so it is now the caller's choice. - self.min_modalities = max(1, min_modalities) - requested_max = max_modalities or len(modalities) - self.max_modalities = ( - requested_max - if allow_repeated_modalities - else min(requested_max, len(modalities)) - ) - if self.max_modalities < self.min_modalities: - raise ValueError( - f"max_modalities ({self.max_modalities}) is below min_modalities " - f"({self.min_modalities})" - ) - self.metric_name = metric - self.maximize_metric = maximize_metric - - if objectives is not None: - if len(objectives) < 1: - raise ValueError( - "objectives must contain at least one (name, direction) pair" - ) - for name, direction in objectives: - if direction not in ("max", "min"): - raise ValueError( - f"objective direction must be 'max' or 'min', got " - f"{direction!r} for objective {name!r}" - ) - self.objective_specs: List[ObjectiveSpec] = list(objectives) - # Keep metric_name/maximize_metric meaningful for callers that - # still read them (e.g. _extract_k_best's ranking metric) - must - # be resolved before _extract_k_best() runs below. - self.metric_name = self.objective_specs[0][0] - self.maximize_metric = self.objective_specs[0][1] == "max" - else: - self.objective_specs = [ - (self.metric_name, "max" if self.maximize_metric else "min") - ] - self.is_multi_objective = len(self.objective_specs) > 1 - - if len(self.modalities) < self.min_modalities: - raise ValueError( - f"MultimodalDeapOptimizer requires at least {self.min_modalities} " - f"modalities, got {len(self.modalities)}." - ) - - self.operator_registry = Registry() - self.fusion_operators = self.operator_registry.get_fusion_operators() - if not self.fusion_operators: - raise ValueError( - "MultimodalDeapOptimizer requires at least one registered " - "fusion operator." - ) - self.k_best_representations = self._extract_k_best( - unimodal_optimization_results - ) - - self.optimization_results: Dict[str, List[FusionSearchResult]] = {} - self.evaluation_errors: Dict[str, int] = {} - - # Incumbent(s) per task, maintained as evaluations come in, so the - # winner does not have to be recovered by re-ranking every result the - # search ever produced. Single-objective: the hall_of_fame_size best - # by the configured objective. Multi-objective: the (unbounded) - # non-dominated front over the objective tuples. - self.hall_of_fame_size = max(1, hall_of_fame_size) - self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} - self._hof_fitness: Dict[str, List[Tuple[float, ...]]] = {} - self.rng = random.Random(random_seed) - # Search-trajectory bookkeeping: see FusionSearchResult.eval_index. - self._eval_counter = 0 - self._current_generation = -1 - self._search_start = time.perf_counter() - self.population_size = max(1, population_size) - self.generations = max(1, generations) - self.crossover_probability = crossover_probability - self.mutation_probability = mutation_probability - self.random_seed = random_seed - self._fitness_cache: Dict[str, Dict[Tuple, Tuple[float, ...]]] = {} - - # Leave room for at least one non-elite offspring per generation, - # otherwise the next generation would just be a sorted clone of the - # current one and crossover/mutation would never run. Unused in - # multi-objective mode (NSGA-II selection is implicitly elitist). - self.elite_size = max(0, min(elite_size, self.population_size - 1)) - self.max_workers = max(1, max_workers) - self.batch_size = max(1, batch_size or self.max_workers) - self.early_stopping_patience = early_stopping_patience - self.early_stopping_min_delta = early_stopping_min_delta - self.novelty_breeding = novelty_breeding - self._current_task_name = None - - desired_weights = tuple( - 1.0 if direction == "max" else -1.0 for _, direction in self.objective_specs - ) - self._objective_weights = desired_weights - existing_weights = getattr( - getattr(creator, "FitnessMax", None), "weights", None - ) - if existing_weights != desired_weights: - if hasattr(creator, "Individual"): - del creator.Individual - if hasattr(creator, "FitnessMax"): - del creator.FitnessMax - creator.create("FitnessMax", base.Fitness, weights=desired_weights) - creator.create("Individual", list, fitness=creator.FitnessMax) - elif not hasattr(creator, "Individual"): - creator.create("Individual", list, fitness=creator.FitnessMax) - - # ------------------------------------------------------------------ - # Public entrypoint - # ------------------------------------------------------------------ - - def optimize( - self, - ) -> Dict[str, List[FusionSearchResult]]: - """ - Optimize the multimodal representations for the tasks. - @return: Dictionary of optimization results for each task. - """ - for task in self.tasks: - task_name = task.model.name - self._current_task_name = task_name - self.optimization_results.setdefault(task_name, []) - self.evaluation_errors.setdefault(task_name, 0) - - # Each task restarts the search, so its trajectory restarts too. - self._eval_counter = 0 - self._search_start = time.perf_counter() - - population = self._build_initial_population(task_name) - best_ever = None - no_improve = 0 - - for gen in range(self.generations): - self._current_generation = gen - self._evaluate_population(population, task) - - if self.is_multi_objective: - # No single "best" individual exists with multiple - # objectives - track the non-dominated (Pareto) front - # instead; progress means that front actually changed. - front = tools.sortNondominated( - population, len(population), first_front_only=True - )[0] - front_signature = frozenset( - self._genome_signature(ind[0]) for ind in front - ) - if best_ever is None or front_signature != best_ever: - best_ever = front_signature - no_improve = 0 - else: - no_improve += 1 - debug_msg = f"front_size={len(front)}" - else: - gen_best = max(population, key=lambda ind: ind.fitness.values[0]) - if ( - best_ever is None - or gen_best.fitness.values[0] - > best_ever.fitness.values[0] + self.early_stopping_min_delta - ): - best_ever = self._clone_individual(gen_best) - no_improve = 0 - else: - no_improve += 1 - debug_msg = f"best={gen_best.fitness.values[0]:.4f}" - - if self.debug: - print( - f"[GA] task={task_name} gen={gen} {debug_msg} " - f"no_improve={no_improve} " - f"errors={self.evaluation_errors.get(task_name, 0)}" - ) - - stagnated = ( - self.early_stopping_patience is not None - and no_improve >= self.early_stopping_patience - ) - if stagnated or gen == self.generations - 1: - if self.debug and stagnated: - print( - f"[GA] task={task_name} early stopping after " - f"{no_improve} generations without improvement" - ) - break - - population = self._next_generation(population, task_name, task) - - return self.optimization_results - - # ------------------------------------------------------------------ - # Population lifecycle - # ------------------------------------------------------------------ - - def _make_individual(self, genome: DagGenome): - return creator.Individual([genome]) - - def _clone_individual(self, ind): - clone = creator.Individual([copy.deepcopy(ind[0])]) - if ind.fitness.valid: - clone.fitness.values = ind.fitness.values - return clone - - def _append_if_unique( - self, - population: List[Any], - genome: DagGenome, - seen_signatures: set, - ) -> bool: - if len(population) >= self.population_size: - return False - sig = self._genome_signature(genome) - if sig in seen_signatures: - return False - seen_signatures.add(sig) - population.append(self._make_individual(genome)) - return True - - def _build_initial_population(self, task_name: str) -> List[Any]: - population: List[Any] = [] - seen: set = set() - retry_budget = max(20, self.population_size * 10) - retries = 0 - while len(population) < self.population_size and retries < retry_budget: - genome = self._random_genome(task_name) - if self._append_if_unique(population, genome, seen): - retries = 0 - else: - retries += 1 - # Search space smaller than population_size: fall back to - # (possibly duplicated) random individuals rather than looping - # forever. - while len(population) < self.population_size: - population.append(self._make_individual(self._random_genome(task_name))) - return population - - def _next_generation( - self, population: List[Any], task_name: str, task: Task - ) -> List[Any]: - if self.is_multi_objective: - # Classic (mu+lambda) NSGA-II: breed a full offspring pool, - # evaluate it, then select the next generation from parents + - # offspring combined via non-dominated sorting + crowding - # distance. This is inherently elitist (non-dominated parents - # survive on their own merit), so there's no separate elite_size - # concept here. - offspring = self._breed_offspring( - population, task_name, seen=self._novelty_archive(task_name) - ) - self._evaluate_population(offspring, task) - combined = list(population) + list(offspring) - return list(tools.selNSGA2(combined, self.population_size)) - - ranked = sorted(population, key=lambda ind: ind.fitness.values[0], reverse=True) - elite = [self._clone_individual(ind) for ind in ranked[: self.elite_size]] - seen = {self._genome_signature(ind[0]) for ind in elite} - seen |= self._novelty_archive(task_name) - return self._breed_offspring(population, task_name, initial=elite, seen=seen) - - def _novelty_archive(self, task_name: str) -> set: - """Signatures already evaluated for this task, or empty if disabled. - - `_breed_offspring` dedups against `seen`, which without this holds - only the current generation's elite -- so an offspring identical to - something evaluated three generations ago is accepted, served from - `_fitness_cache`, and occupies a population slot that explores - nothing. Feeding the whole cache in makes the dedup global. - """ - if not self.novelty_breeding: - return set() - return set(self._fitness_cache.get(task_name, {})) - - def _breed_offspring( - self, - population: List[Any], - task_name: str, - initial: Optional[List[Any]] = None, - seen: Optional[set] = None, - ) -> List[Any]: - """ - Fills a new population of size population_size via tournament - selection + crossover + mutation, starting from `initial` (e.g. an - elite carryover) if given. Falls back to fresh random immigrants - (allowing duplicates) once the retry budget is exhausted, so a - search space smaller than population_size can never hang. - """ - next_population = list(initial) if initial else [] - seen = set(seen) if seen else set() - - retry_budget = max(20, self.population_size * 10) - retries = 0 - tournsize = max(1, min(3, len(population))) - while len(next_population) < self.population_size and retries < retry_budget: - p1, p2 = tools.selTournament(population, 2, tournsize=tournsize) - - if self.rng.random() < self.crossover_probability: - g1, g2 = self._crossover_genomes(p1[0], p2[0]) - else: - g1, g2 = copy.deepcopy(p1[0]), copy.deepcopy(p2[0]) - - if self.rng.random() < self.mutation_probability: - g1 = self._mutate_genome(g1, task_name) - if self.rng.random() < self.mutation_probability: - g2 = self._mutate_genome(g2, task_name) - - added1 = self._append_if_unique(next_population, g1, seen) - added2 = self._append_if_unique(next_population, g2, seen) - retries = 0 if (added1 or added2) else retries + 1 - - # Ran out of retries (tiny search space): top up with fresh random - # immigrants, allowing duplicates rather than looping forever. - while len(next_population) < self.population_size: - genome = self._random_genome(task_name) - if not self._append_if_unique(next_population, genome, seen): - next_population.append(self._make_individual(genome)) - - return next_population - - # ------------------------------------------------------------------ - # Evaluation (serial + bounded parallel) - # ------------------------------------------------------------------ - - def _evaluate_population(self, population: List[Any], task: Task) -> None: - to_evaluate = [ind for ind in population if not ind.fitness.valid] - if not to_evaluate: - return - if self.max_workers > 1 and len(to_evaluate) > 1: - self._evaluate_individuals_parallel(to_evaluate, task) - else: - for ind in to_evaluate: - fitness = self._evaluate_genome(ind[0], task) - ind.fitness.values = fitness - - def _evaluate_individuals_parallel( - self, individuals: List[Any], task: Task - ) -> None: - task_name = task.model.name - cache = self._fitness_cache.setdefault(task_name, {}) - ctx = mp.get_context("spawn") - task_bytes = pickle.dumps(task) - futures: Dict[Any, Tuple[Any, RepresentationDag, Tuple]] = {} - pending_followers: Dict[Tuple, List[Any]] = {} - - def _drain(done_futures): - for fut in done_futures: - ind, dag, sig = futures.pop(fut) - fitness, payload, error = fut.result() - ind.fitness.values = fitness - self._record_evaluation(task_name, dag, payload, error) - cache[sig] = fitness - for follower in pending_followers.pop(sig, []): - follower.fitness.values = fitness - - with ProcessPoolExecutor( - max_workers=self.max_workers, mp_context=ctx - ) as executor: - for ind in individuals: - genome = ind[0] - sig = self._genome_signature(genome) - cached = cache.get(sig) - if cached is not None: - ind.fitness.values = cached - continue - if sig in pending_followers: - pending_followers[sig].append(ind) - continue - - dag = self._genome_to_dag(genome) - modalities = list( - chain.from_iterable(self.k_best_representations[task_name].values()) - ) - fut = executor.submit( - _evaluate_dag_worker, - pickle.dumps(dag), - task_bytes, - pickle.dumps(modalities), - self.objective_specs, - ) - futures[fut] = (ind, dag, sig) - pending_followers[sig] = [] - - if len(futures) >= self.batch_size: - done, _ = wait(set(futures.keys()), return_when=FIRST_COMPLETED) - _drain(done) - - if futures: - done, _ = wait(set(futures.keys())) - _drain(done) - - def _record_evaluation( - self, - task_name: str, - dag: RepresentationDag, - payload: Optional[Dict[str, Any]], - error: Optional[str], - ) -> None: - self.optimization_results.setdefault(task_name, []) - if error is not None or payload is None: - self.evaluation_errors[task_name] = ( - self.evaluation_errors.get(task_name, 0) + 1 - ) - if self.debug and error is not None: - last_line = error.strip().splitlines()[-1] if error.strip() else error - print( - f"[GA] genome evaluation failed for task={task_name}: {last_line}" - ) - return - - result = FusionSearchResult( - dag=dag, - train_score=payload["train_score"], - val_score=payload["val_score"], - test_score=payload["test_score"], - train_fold_scores=payload.get("train_fold_scores", {}), - val_fold_scores=payload.get("val_fold_scores", {}), - test_fold_scores=payload.get("test_fold_scores", {}), - task_timing=payload.get("task_timing", {}), - runtime=payload["runtime"], - task_time=payload["task_time"], - representation_time=payload["representation_time"], - task_name=task_name, - generation=self._current_generation, - eval_index=self._eval_counter, - t_since_search_start_s=time.perf_counter() - self._search_start, - t_eval_end_unix=time.time(), - ) - self.optimization_results.setdefault(task_name, []).append(result) - self._update_hall_of_fame(task_name, result) - self._eval_counter += 1 - - # ------------------------------------------------------------------ - # Hall of fame - # ------------------------------------------------------------------ - - def _dominates(self, a: Tuple[float, ...], b: Tuple[float, ...]) -> bool: - """True if objective tuple `a` Pareto-dominates `b`, direction-aware.""" - wa = [w * v for w, v in zip(self._objective_weights, a)] - wb = [w * v for w, v in zip(self._objective_weights, b)] - return all(x >= y for x, y in zip(wa, wb)) and any( - x > y for x, y in zip(wa, wb) - ) - - def _update_hall_of_fame(self, task_name: str, result: FusionSearchResult) -> None: - timing = { - "runtime": result.runtime, - "task_time": result.task_time, - "representation_time": result.representation_time, - } - try: - fitness = tuple( - _objective_value(name, result.val_score, timing) - for name, _ in self.objective_specs - ) - except KeyError: - # An objective the task did not report for this candidate: it - # cannot be ranked, so it simply does not enter the hall of fame. - return - - hof = self.hall_of_fame.setdefault(task_name, []) - fits = self._hof_fitness.setdefault(task_name, []) - - if self.is_multi_objective: - if any(self._dominates(f, fitness) or f == fitness for f in fits): - return - keep = [i for i, f in enumerate(fits) if not self._dominates(fitness, f)] - self.hall_of_fame[task_name] = [hof[i] for i in keep] + [result] - self._hof_fitness[task_name] = [fits[i] for i in keep] + [fitness] - return - - weight = self._objective_weights[0] - hof.append(result) - fits.append(fitness) - order = sorted( - range(len(fits)), key=lambda i: weight * fits[i][0], reverse=True - ) - order = order[: self.hall_of_fame_size] - self.hall_of_fame[task_name] = [hof[i] for i in order] - self._hof_fitness[task_name] = [fits[i] for i in order] - - def get_hall_of_fame(self, task_name: str) -> List[FusionSearchResult]: - """The incumbent(s) for a task: the best `hall_of_fame_size` results by - the configured objective, or the non-dominated front in multi-objective - mode. Empty if nothing evaluated successfully.""" - return list(self.hall_of_fame.get(task_name, [])) - - def _extract_k_best(self, unimodal_results) -> Dict[str, Dict[str, List[Any]]]: - """ - Extract the k best representations for each modality and task. - @param unimodal_results: Unimodal optimization results. - @return: Dictionary of k best representations for each modality and task. - """ - k_best = {} - for task in self.tasks: - name = task.model.name - k_best[name] = {} - for modality in self.modalities: - _, cached_data = unimodal_results.get_k_best_results( - modality, task, self.metric_name - ) - k_best[name][modality.modality_id] = cached_data - return k_best - - def _available_modality_ids(self, task_name: str) -> List[Any]: - """Modalities with at least one representation for this task. - - Datasets like StressID can be missing a modality for some - tasks/instances, and a leaf pointing at an empty representation list - is unresolvable. - """ - reps = self.k_best_representations[task_name] - return [ - m.modality_id - for m in self.modalities - if len(reps.get(m.modality_id, [])) > 0 - ] - - def _leaf_capacity(self, task_name: str) -> int: - """How many distinct leaves a genome may hold at most. - - One per modality normally; with allow_repeated_modalities, one per - (modality, representation) pair, since that is what makes a leaf - distinct. - """ - reps = self.k_best_representations[task_name] - ids = self._available_modality_ids(task_name) - if not self.allow_repeated_modalities: - return len(ids) - return sum(len(reps[mid]) for mid in ids) - - def _random_genome(self, task_name: str) -> DagGenome: - """ - Generate a random genome for a given task. - @param task_name: Name of the task. - @return: Random genome. - """ - reps = self.k_best_representations[task_name] - available_modality_ids = self._available_modality_ids(task_name) - capacity = self._leaf_capacity(task_name) - if capacity < self.min_modalities: - raise ValueError( - f"Need at least {self.min_modalities} distinct leaves for task " - f"'{task_name}', but only {capacity} are available across " - f"{len(available_modality_ids)} modalities." - ) - - upper = min(self.max_modalities, capacity) - lower = min(self.min_modalities, upper) - r = self.rng.randint(lower, upper) - - if self.allow_repeated_modalities: - # Sample distinct (modality, representation) leaves; a modality may - # appear more than once as long as it contributes a different - # representation each time. - pool = [ - (mid, idx) - for mid in available_modality_ids - for idx in range(len(reps[mid])) - ] - leaves = self.rng.sample(pool, r) - else: - chosen = self.rng.sample(available_modality_ids, r) - leaves = [(mid, self.rng.randrange(len(reps[mid]))) for mid in chosen] - - tree = self._random_binary_tree(len(leaves)) - fusion_ops = {} - self._assign_fusion_ops(tree, fusion_ops, "") - return DagGenome(leaves=leaves, tree=tree, fusion_ops=fusion_ops) - - def _internal_paths(self, tree): - return _collect_internal_paths(tree) - - def _random_binary_tree(self, n: int) -> Tree: - """ - Generate a random binary tree with n leaves. - @param n: Number of leaves. - @return: Random binary tree. - """ - nodes: List[Tree] = list(range(n)) - while len(nodes) > 1: - i, j = self.rng.sample(range(len(nodes)), 2) - a, b = nodes.pop(max(i, j)), nodes.pop(min(i, j)) - nodes.append((a, b)) - return nodes[0] - - def _assign_fusion_ops(self, subtree: Tree, ops: Dict[str, Any], path: str) -> None: - """ - Assign fusion operators to the internal nodes of a binary tree. - @param subtree: Binary tree. - @param ops: Dictionary of fusion operators. - @param path: Path to the current node. - """ - if isinstance(subtree, int): - return - ops[path] = self.rng.choice(self.fusion_operators) - left, right = subtree - self._assign_fusion_ops(left, ops, path + "L") - self._assign_fusion_ops(right, ops, path + "R") - - def _genome_to_dag(self, genome: DagGenome) -> RepresentationDag: - """ - Convert a genome to a DAG. - @param genome: Genome. - @return: RepresentationDAG. - """ - builder = RepresentationDAGBuilder() - leaf_ids = [ - builder.create_leaf_node(mod_id, repr_idx) - for mod_id, repr_idx in genome.leaves - ] - - def build(subtree: Tree, path: str) -> str: - if isinstance(subtree, int): - return leaf_ids[subtree] - left, right = subtree - left_id = build(left, path + "L") - right_id = build(right, path + "R") - op_cls = genome.fusion_ops[path] - op = op_cls() - return builder.create_operation_node( - op.__class__, [left_id, right_id], op.get_current_parameters() - ) - - return builder.build(build(genome.tree, "")) - - def _genome_signature(self, g: DagGenome) -> Tuple: - """ - Generate a signature for a genome to be used as a cache key. - @param g: Genome. - @return: Signature. - """ - - def norm(t: Tree): - return t if isinstance(t, int) else (norm(t[0]), norm(t[1])) - - return ( - tuple(g.leaves), - norm(g.tree), - tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), - ) - - def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: - """ - Evaluate a genome for a given task (in-process). Stores the result - in the optimization results, or records a failure without raising - if the DAG/task fails to execute (e.g. an incompatible fusion). - @param genome: Genome. - @param task: Task. - @return: Fitness tuple, one value per configured objective. - """ - task_name = task.model.name - sig = self._genome_signature(genome) - cache = self._fitness_cache.setdefault(task_name, {}) - if sig in cache: - return cache[sig] - - dag = self._genome_to_dag(genome) - modalities = list( - chain.from_iterable(self.k_best_representations[task_name].values()) - ) - - try: - fitness, payload = _evaluate_genome_body( - dag, task, modalities, self.objective_specs - ) - error = None - if fitness is None: - fitness = _failure_fitness(self.objective_specs) - except Exception: - fitness = _failure_fitness(self.objective_specs) - payload, error = None, traceback.format_exc() - - self._record_evaluation(task_name, dag, payload, error) - cache[sig] = fitness - return fitness - - # ------------------------------------------------------------------ - # Crossover - # ------------------------------------------------------------------ - - def _crossover_genomes( - self, g1: DagGenome, g2: DagGenome - ) -> Tuple[DagGenome, DagGenome]: - """ - Recombine two genomes. - - When both parents agree on the exact same modality/representation - leaves, a classic reciprocal subtree swap is performed. Otherwise - (the common case, since leaves are sampled independently per - genome) subtree crossover over leaf indices would be meaningless, - so fusion-operator choices are mixed at whichever internal-path - keys the two trees happen to share instead - every non-trivial - tree has at least the root path "" in common, so this always - performs real recombination rather than silently returning the - parents unchanged. - @param g1: First parent genome. - @param g2: Second parent genome. - @return: Two child genomes. - """ - c1, c2 = copy.deepcopy(g1), copy.deepcopy(g2) - - if c1.leaves == c2.leaves: - paths1 = self._internal_paths(c1.tree) - paths2 = self._internal_paths(c2.tree) - if paths1 and paths2: - path1 = self.rng.choice(paths1) - path2 = self.rng.choice(paths2) - subtree1 = _get_subtree(c1.tree, path1) - subtree2 = _get_subtree(c2.tree, path2) - c1.tree = _replace_subtree(c1.tree, path1, subtree2) - c2.tree = _replace_subtree(c2.tree, path2, subtree1) - c1.fusion_ops = _rebuild_fusion_ops( - c1.tree, - {**c1.fusion_ops, **c2.fusion_ops}, - self.rng, - self.fusion_operators, - randomized_prefixes=[path1], - ) - c2.fusion_ops = _rebuild_fusion_ops( - c2.tree, - {**c2.fusion_ops, **c1.fusion_ops}, - self.rng, - self.fusion_operators, - randomized_prefixes=[path2], - ) - return c1, c2 - - shared_paths = set(c1.fusion_ops) & set(c2.fusion_ops) - for path in shared_paths: - if self.rng.random() < 0.5: - c1.fusion_ops[path], c2.fusion_ops[path] = ( - c2.fusion_ops[path], - c1.fusion_ops[path], - ) - return c1, c2 - - # ------------------------------------------------------------------ - # Mutation - # ------------------------------------------------------------------ - - def _mutate_genome(self, g: DagGenome, task_name: str) -> DagGenome: - op = self.rng.choice( - [ - self._mutate_change_fusion, - lambda gg: self._mutate_swap_leaf_repr(gg, task_name), - lambda gg: self.mutate_add_leaf(gg, task_name), - self.mutate_remove_leaf, - self.mutate_replace_subtree, - ] - ) - return op(g) - - def _mutate_change_fusion(self, g: DagGenome) -> DagGenome: - """ - Change a fusion operator at a random internal node. - @param g: Genome. - @return: Mutated genome. - """ - g = copy.deepcopy(g) - paths = [p for p in g.fusion_ops] - if not paths: - return g - path = self.rng.choice(paths) - choices = [op for op in self.fusion_operators if op != g.fusion_ops[path]] - if choices: - g.fusion_ops[path] = self.rng.choice(choices) - return g - - def _mutate_swap_leaf_repr(self, g: DagGenome, task_name: str) -> DagGenome: - """ - Change which k-best unimodal repr a leaf uses (same modality). - @param g: Genome. - @param task_name: Name of the task. - @return: Mutated genome. - """ - g = copy.deepcopy(g) - i = self.rng.randrange(len(g.leaves)) - mod_id, current = g.leaves[i] - k = len(self.k_best_representations[task_name][mod_id]) - if k <= 1: - return g - # With repeated modalities allowed the genome may already hold another - # leaf of this modality; swapping onto that representation would make - # the two leaves identical, which wastes a slot on a duplicate. - taken = { - idx for j, (mid, idx) in enumerate(g.leaves) if mid == mod_id and j != i - } - choices = [idx for idx in range(k) if idx != current and idx not in taken] - if not choices: - return g - g.leaves[i] = (mod_id, self.rng.choice(choices)) - return g - - def mutate_add_leaf(self, g: DagGenome, task_name: str) -> DagGenome: - """ - Add a new leaf to the genome (adding a new modality). Only - modalities that have at least one representation for this task are - considered. - @param g: Genome. - @param task_name: Name of the task. - @return: Mutated genome. - """ - if len(g.leaves) >= self.max_modalities: - return g - g = copy.deepcopy(g) - reps = self.k_best_representations[task_name] - if self.allow_repeated_modalities: - # Any (modality, representation) leaf the genome does not already - # hold, so a modality can be added a second time under a different - # representation. - existing_leaves = set(g.leaves) - candidates = [ - (mid, idx) - for mid in self._available_modality_ids(task_name) - for idx in range(len(reps[mid])) - if (mid, idx) not in existing_leaves - ] - if not candidates: - return g - new_leaf = self.rng.choice(candidates) - else: - existing = {l[0] for l in g.leaves} - available = [ - m.modality_id - for m in self.modalities - if m.modality_id not in existing - and len(reps.get(m.modality_id, [])) > 0 - ] - if not available: - return g - mod_id = self.rng.choice(available) - new_leaf = (mod_id, self.rng.randrange(len(reps[mod_id]))) - new_idx = len(g.leaves) - g.leaves.append(new_leaf) - if isinstance(g.tree, int): - g.tree = (g.tree, new_idx) - g.fusion_ops = {"": self.rng.choice(self.fusion_operators)} - else: - paths = self._internal_paths(g.tree) - path = self.rng.choice(paths) - sub = _get_subtree(g.tree, path) - g.tree = _replace_subtree(g.tree, path, (sub, new_idx)) - # `sub` (and everything under it) just shifted one level deeper - # in the tree (from `path` to `path + "L"`), so its existing - # fusion_ops entries are keyed under stale paths and must be - # rebuilt rather than patched in place - only the brand-new - # node at `path` needs a freshly chosen operator. - g.fusion_ops = _rebuild_fusion_ops( - g.tree, - g.fusion_ops, - self.rng, - self.fusion_operators, - randomized_prefixes=[path], - ) - return g - - def mutate_remove_leaf(self, g: DagGenome) -> DagGenome: - """ - Remove a leaf from the genome (removing a modality). - @param g: Genome. - @return: Mutated genome. - """ - if len(g.leaves) <= self.min_modalities: - return g - g = copy.deepcopy(g) - drop = self.rng.randrange(len(g.leaves)) - new_tree = _remove_leaf_from_tree(g.tree, drop) - if new_tree is None: - return g - # Reindex leaves + fusion_ops paths - keep = [i for i in range(len(g.leaves)) if i != drop] - index_map = {old: new for new, old in enumerate(keep)} - g.leaves = [g.leaves[i] for i in keep] - g.tree = _reindex_tree(new_tree, index_map) - g.fusion_ops = _rebuild_fusion_ops( - g.tree, g.fusion_ops, self.rng, self.fusion_operators - ) - return g - - def mutate_replace_subtree(self, g: DagGenome) -> DagGenome: - """ - Replace a subtree with a random subtree. - - The collapse branch (replacing an internal node by one of its - children) drops every leaf on the other side of that node, so it is - only taken when at least min_modalities leaves survive -- otherwise a - search configured min_modalities=2 silently evaluates unimodal - pipelines. The dropped leaves are also removed from `genome.leaves` - and the tree reindexed; leaving them in place used to desynchronise - the genome from its own tree. - @param g: Genome. - @return: Mutated genome. - """ - g = copy.deepcopy(g) - paths = self._internal_paths(g.tree) - if not paths: - return g - path = self.rng.choice(paths) - sub = _get_subtree(g.tree, path) - if isinstance(sub, int): - return g - - if self.rng.random() < 0.5: - child = sub[0] if self.rng.random() < 0.5 else sub[1] - candidate = _replace_subtree(g.tree, path, child) - kept = sorted(set(_collect_leaf_indices(candidate))) - if len(kept) >= self.min_modalities: - index_map = {old: new for new, old in enumerate(kept)} - g.leaves = [g.leaves[i] for i in kept] - g.tree = _reindex_tree(candidate, index_map) - g.fusion_ops = _rebuild_fusion_ops( - g.tree, {}, self.rng, self.fusion_operators - ) - return g - # Collapsing here would leave fewer leaves than min_modalities, - # so fall through to the reshuffle branch, which is leaf-preserving. - - leaf_idxs = _collect_leaf_indices(sub) - new_sub = self._random_binary_tree(len(leaf_idxs)) - local_map = {i: leaf_idxs[i] for i in range(len(leaf_idxs))} - new_sub = _reindex_tree(new_sub, local_map) - g.tree = _replace_subtree(g.tree, path, new_sub) - g.fusion_ops = _rebuild_fusion_ops( - g.tree, - g.fusion_ops, - self.rng, - self.fusion_operators, - randomized_prefixes=[path], - ) - return g - - # ------------------------------------------------------------------ - # Persistence - # ------------------------------------------------------------------ - - def store_results(self, file_name: str = None, overwrite: bool = False) -> str: - """ - Persist optimization_results to disk. - - Refuses to clobber an existing file unless overwrite=True is - passed explicitly. The write itself is atomic (temp file + - os.replace), so a crash mid-write can never leave a corrupted or - truncated results file behind. - @param file_name: Destination path. A timestamped name is - generated if omitted. - @param overwrite: Set True to explicitly replace an existing file - at file_name. - @return: The path the results were written to. - """ - if file_name is None: - timestr = time.strftime("%Y%m%d-%H%M%S") - file_name = f"multimodal_optimizer_{timestr}.pkl" - - directory = os.path.dirname(file_name) or "." - os.makedirs(directory, exist_ok=True) - - if os.path.exists(file_name) and not overwrite: - raise FileExistsError( - f"Refusing to overwrite existing results file '{file_name}'. " - "Pass overwrite=True if this is intentional, or choose a " - "different file_name." - ) - - fd, tmp_path = tempfile.mkstemp( - dir=directory, prefix=".tmp_multimodal_results_", suffix=".pkl" - ) - try: - with os.fdopen(fd, "wb") as f: - pickle.dump(self.optimization_results, f) - os.replace(tmp_path, file_name) - except Exception: - if os.path.exists(tmp_path): - os.remove(tmp_path) - raise - return file_name diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 378c2618043..718c342d573 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -1,114 +1,59 @@ +# ------------------------------------------------------------- +# +# 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 __future__ import annotations import copy import multiprocessing as mp +import os import pickle import random +import tempfile import time +import traceback +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait from dataclasses import dataclass, field -from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED -from typing import Any, Dict, Generator, List, Optional, Tuple +from itertools import chain +from typing import Any, Dict, List, Optional, Tuple + +from deap import base, creator, tools from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.drsearch.representation_dag import ( - RepresentationDag, RepresentationDAGBuilder, + RepresentationDag, ) from systemds.scuro.drsearch.task import Task -from systemds.scuro.modality.modality import Modality -from systemds.scuro.utils.checkpointing import CheckpointManager -from systemds.scuro.utils.static_variables import DEBUG - -# ---------------------------- -# Genome / Individual encoding -# ---------------------------- - -Tree = Any # int leaf index OR tuple(left_subtree, right_subtree) - - -def genome_to_dag(genome) -> RepresentationDag: - builder = RepresentationDAGBuilder() - leaf_ids = [ - builder.create_leaf_node(mod_id, repr_idx) for mod_id, repr_idx in genome.leaves - ] - - def build(subtree: Tree, path: str) -> str: - if isinstance(subtree, int): - return leaf_ids[subtree] - left, right = subtree - left_id = build(left, path + "L") - right_id = build(right, path + "R") - op_cls = genome.fusion_ops[path] - op = op_cls() - return builder.create_operation_node( - op.__class__, [left_id, right_id], op.get_current_parameters() - ) - - return builder.build(build(genome.tree, "")) - - -def _evaluate_individual_worker( - dag_pickle: bytes, - task_pickle: bytes, - modalities_pickle: bytes, - metric_name: str, -) -> Tuple[float, Dict[str, Any]]: - dag = pickle.loads(dag_pickle) - task = pickle.loads(task_pickle) - modalities = pickle.loads(modalities_pickle) - - start_time = time.time() - fused_representation = dag.execute(modalities, task) - scores = task.run(fused_representation.data) - runtime = time.time() - start_time - fitness = scores[1].average_scores[metric_name] - - objective = { - "train_score": scores[0].average_scores, - "val_score": scores[1].average_scores, - "test_score": scores[2].average_scores, - "val": fitness, - "runtime": runtime, - "representation_time": runtime, - "task_time": 0.0, - } - return fitness, objective - - -@dataclass -class Individual: - # Leaves are concrete unimodal choices: (modality_id, representation_index) - leaves: List[Tuple[str, int]] - # Binary tree over leaf indices in `leaves` - tree: Tree - # Fusion op class per internal tree-path key, e.g. "L", "RLL", ... - fusion_ops: Dict[str, Any] - fitness: float = float("-inf") - objective: Dict[str, float] = field(default_factory=dict) - dag: Optional[RepresentationDag] = None - - -# ---------------------------- -# Optional helper abstraction -# ---------------------------- - +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.utils.schema_helpers import get_shape -class MutationOperator: - name: str = "base" +Tree = int | Tuple["Tree", "Tree"] - def __call__( - self, - individual: Individual, - rng: random.Random, - context: Dict[str, Any], - ) -> Individual: - raise NotImplementedError() - -def _collect_internal_paths(subtree: Tree, path: str = "") -> List[str]: - if isinstance(subtree, int): +def _collect_internal_paths(tree: Tree, path: str = "") -> List[str]: + if isinstance(tree, int): return [] - left, right = subtree + left, right = tree return ( [path] + _collect_internal_paths(left, path + "L") @@ -116,768 +61,707 @@ def _collect_internal_paths(subtree: Tree, path: str = "") -> List[str]: ) -def _get_subtree(subtree: Tree, target_path: str) -> Tree: - if target_path == "": - return copy.deepcopy(subtree) - if isinstance(subtree, int): - raise ValueError(f"Path '{target_path}' does not exist in leaf subtree") - left, right = subtree - if target_path[0] == "L": - return _get_subtree(left, target_path[1:]) - return _get_subtree(right, target_path[1:]) +def _get_subtree(tree: Tree, path: str) -> Tree: + if not path: + return copy.deepcopy(tree) + left, right = tree + branch = left if path[0] == "L" else right + return _get_subtree(branch, path[1:]) -def _replace_subtree(subtree: Tree, target_path: str, replacement: Tree) -> Tree: - if target_path == "": +def _replace_subtree(tree: Tree, path: str, replacement: Tree) -> Tree: + if not path: return copy.deepcopy(replacement) - if isinstance(subtree, int): - raise ValueError(f"Path '{target_path}' does not exist in leaf subtree") - left, right = subtree - if target_path[0] == "L": - return ( - _replace_subtree(left, target_path[1:], replacement), - copy.deepcopy(right), - ) - return (copy.deepcopy(left), _replace_subtree(right, target_path[1:], replacement)) + left, right = tree + if path[0] == "L": + return _replace_subtree(left, path[1:], replacement), copy.deepcopy(right) + return copy.deepcopy(left), _replace_subtree(right, path[1:], replacement) -def _collect_leaf_indices(subtree: Tree) -> List[int]: - if isinstance(subtree, int): - return [subtree] - left, right = subtree - return _collect_leaf_indices(left) + _collect_leaf_indices(right) +def _collect_leaf_indices(tree: Tree) -> List[int]: + if isinstance(tree, int): + return [tree] + return _collect_leaf_indices(tree[0]) + _collect_leaf_indices(tree[1]) -def _sample_random_binary_tree_from_leaves( - leaf_indices: List[int], rng: random.Random -) -> Tree: - nodes: List[Tree] = list(leaf_indices) - while len(nodes) > 1: - i, j = rng.sample(range(len(nodes)), 2) - a = nodes.pop(max(i, j)) - b = nodes.pop(min(i, j)) - nodes.append((a, b)) - return nodes[0] +def _remove_leaf_from_tree(tree: Tree, leaf: int) -> Optional[Tree]: + if isinstance(tree, int): + return None if tree == leaf else tree + left = _remove_leaf_from_tree(tree[0], leaf) + right = _remove_leaf_from_tree(tree[1], leaf) + if left is None: + return right + if right is None: + return left + return left, right -def _reindex_tree(subtree: Tree, index_map: Dict[int, int]) -> Tree: - if isinstance(subtree, int): - return index_map[subtree] - left, right = subtree - return (_reindex_tree(left, index_map), _reindex_tree(right, index_map)) +def _reindex_tree(tree: Tree, index_map: Dict[int, int]) -> Tree: + if isinstance(tree, int): + return index_map[tree] + return _reindex_tree(tree[0], index_map), _reindex_tree(tree[1], index_map) -def _remove_leaf_from_tree(subtree: Tree, leaf_idx: int) -> Optional[Tree]: - if isinstance(subtree, int): - if subtree == leaf_idx: - return None - return subtree +def _rebuild_fusion_ops( + tree: Tree, + existing: Dict[str, type], + rng: random.Random, + operators: List[type], + randomized_prefixes: Optional[List[str]] = None, +) -> Dict[str, type]: + prefixes = randomized_prefixes or [] + rebuilt = {} + for path in _collect_internal_paths(tree): + randomized = any(not prefix or path.startswith(prefix) for prefix in prefixes) + rebuilt[path] = ( + rng.choice(operators) + if randomized or path not in existing + else existing[path] + ) + return rebuilt - left, right = subtree - new_left = _remove_leaf_from_tree(left, leaf_idx) - new_right = _remove_leaf_from_tree(right, leaf_idx) - if new_left is None: - return new_right - if new_right is None: - return new_left - return (new_left, new_right) +@dataclass +class FusionSearchResult: + dag: RepresentationDag + train_score: dict + val_score: dict + test_score: dict + runtime: float = 0.0 + task_time: float = 0.0 + representation_time: float = 0.0 + task_name: str = "" + val_fold_scores: dict = field(default_factory=dict) + train_fold_scores: dict = field(default_factory=dict) + test_fold_scores: dict = field(default_factory=dict) -def _path_is_under_prefix(path: str, prefix: str) -> bool: - return prefix == "" or path == prefix or path.startswith(prefix) + task_timing: dict = field(default_factory=dict) + generation: int = -1 + eval_index: int = -1 + t_since_search_start_s: float = 0.0 + t_eval_end_unix: float = 0.0 + + +@dataclass +class DagGenome: + leaves: List[Tuple[str, int]] + tree: Tree + fusion_ops: Dict[str, type] -def _rebuild_fusion_ops( - subtree: Tree, - existing_fusion_ops: Dict[str, Any], - rng: random.Random, - fusion_operators: List[Any], - randomized_prefixes: Optional[List[str]] = None, -) -> Dict[str, Any]: - fusion_ops: Dict[str, Any] = {} - randomized_prefixes = randomized_prefixes or [] - for path in _collect_internal_paths(subtree): - if any(_path_is_under_prefix(path, prefix) for prefix in randomized_prefixes): - fusion_ops[path] = rng.choice(fusion_operators) - elif path in existing_fusion_ops: - fusion_ops[path] = existing_fusion_ops[path] - else: - fusion_ops[path] = rng.choice(fusion_operators) - return fusion_ops +_TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"} -# ---------------------------- -# GA Optimizer skeleton -# ---------------------------- +ObjectiveSpec = Tuple[str, str] # (name, "max" | "min") -class MultimodalGAOptimizer: - """ - GA-based multimodal optimizer skeleton. - Keeps constructor inputs compatible with original MultimodalOptimizer. - """ +def _objective_value( + name: str, val_score: Dict[str, float], timing: Dict[str, float] +) -> float: + if name in _TIMING_OBJECTIVES: + return timing[name] + return val_score[name] + +def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: + return tuple( + float("-inf") if direction == "max" else float("inf") + for _, direction in objective_specs + ) + + +def _evaluate_genome_body( + dag: RepresentationDag, + task: Task, + modalities: List[Any], + objective_specs: List[ObjectiveSpec], +) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: + start = time.time() + fused = dag.execute(modalities, task, enable_cache=False) + if fused is None: + return None, None + + if isinstance(fused, dict): + fused = fused[list(fused.keys())[-1]] + + if task.expected_dim == 1 and get_shape(fused.metadata) > 1: + fused = AggregatedRepresentation().transform(fused) + + t0 = time.time() + scores = task.run(fused.data) + task_time = time.time() - t0 + total = time.time() - start + + val_score = scores[1].average_scores + timing = { + "runtime": total, + "task_time": task_time, + "representation_time": total - task_time, + } + fitness = tuple( + _objective_value(name, val_score, timing) for name, _ in objective_specs + ) + payload = { + "train_score": scores[0].average_scores, + "val_score": val_score, + "test_score": scores[2].average_scores, + "train_fold_scores": scores[0].fold_scores(), + "val_fold_scores": scores[1].fold_scores(), + "test_fold_scores": scores[2].fold_scores(), + "task_timing": getattr(task, "last_run_timing", {}), + **timing, + } + return fitness, payload + + +def _evaluate_dag_worker( + dag_bytes: bytes, + task_bytes: bytes, + modalities_bytes: bytes, + objective_specs: List[ObjectiveSpec], +) -> Tuple[Tuple[float, ...], Optional[Dict[str, Any]], Optional[str]]: + try: + dag = pickle.loads(dag_bytes) + task = pickle.loads(task_bytes) + modalities = pickle.loads(modalities_bytes) + fitness, payload = _evaluate_genome_body(dag, task, modalities, objective_specs) + if fitness is None: + fitness = _failure_fitness(objective_specs) + return fitness, payload, None + except Exception: + return _failure_fitness(objective_specs), None, traceback.format_exc() + + +class MultimodalDeapOptimizer: def __init__( self, modalities: List[Any], unimodal_optimization_results: Any, - tasks: List[Any], - k: int = 2, - debug: bool = False, + tasks: List[Task], + debug: bool = True, min_modalities: int = 2, max_modalities: int = None, metric: str = "accuracy", - checkpoint_every: int = None, - resume: bool = True, - # --- GA controls (new) --- + objectives: Optional[List[ObjectiveSpec]] = None, population_size: int = 32, generations: int = 20, - elite_size: int = 4, - tournament_size: int = 3, - crossover_rate: float = 0.9, - mutation_rate: float = 0.3, + crossover_probability: float = 0.7, + mutation_probability: float = 0.4, random_seed: int = 42, + maximize_metric: bool = True, + elite_size: int = 2, + max_workers: int = 1, + batch_size: Optional[int] = None, early_stopping_patience: Optional[int] = 5, early_stopping_min_delta: float = 1e-6, + novelty_breeding: bool = True, + hall_of_fame_size: int = 5, + allow_repeated_modalities: bool = False, ): self.modalities = modalities self.tasks = tasks - self.k = k self.debug = debug - if DEBUG: - self.debug = True + self.allow_repeated_modalities = allow_repeated_modalities + + self.min_modalities = max(1, min_modalities) + requested_max = max_modalities or len(modalities) + self.max_modalities = ( + requested_max + if allow_repeated_modalities + else min(requested_max, len(modalities)) + ) + if self.max_modalities < self.min_modalities: + raise ValueError( + f"max_modalities ({self.max_modalities}) is below min_modalities " + f"({self.min_modalities})" + ) self.metric_name = metric - self.min_modalities = max(2, min_modalities) - self.max_modalities = max_modalities or len(modalities) - - self.population_size = population_size - self.generations = generations - self.elite_size = elite_size - self.tournament_size = tournament_size - self.crossover_rate = crossover_rate - self.mutation_rate = mutation_rate - self.random_seed = random_seed - self.rng = random.Random(random_seed) - self.early_stopping_patience = early_stopping_patience - self.early_stopping_min_delta = early_stopping_min_delta + self.maximize_metric = maximize_metric + + if objectives is not None: + if len(objectives) < 1: + raise ValueError( + "objectives must contain at least one (name, direction) pair" + ) + for name, direction in objectives: + if direction not in ("max", "min"): + raise ValueError( + f"objective direction must be 'max' or 'min', got " + f"{direction!r} for objective {name!r}" + ) + self.objective_specs: List[ObjectiveSpec] = list(objectives) + self.metric_name = self.objective_specs[0][0] + self.maximize_metric = self.objective_specs[0][1] == "max" + else: + self.objective_specs = [ + (self.metric_name, "max" if self.maximize_metric else "min") + ] + self.is_multi_objective = len(self.objective_specs) > 1 + + if len(self.modalities) < self.min_modalities: + raise ValueError( + f"MultimodalDeapOptimizer requires at least {self.min_modalities} " + f"modalities, got {len(self.modalities)}." + ) self.operator_registry = Registry() self.fusion_operators = self.operator_registry.get_fusion_operators() - - self.k_best_representations = self._extract_k_best_representations( + if not self.fusion_operators: + raise ValueError( + "MultimodalDeapOptimizer requires at least one registered " + "fusion operator." + ) + self.k_best_representations = self._extract_k_best( unimodal_optimization_results ) - self.optimization_results: Dict[str, List[OptimizationResult]] = {} - self._seen_results_by_task: Dict[str, Dict[str, OptimizationResult]] = {} - self._fitness_cache_by_task: Dict[ - str, Dict[Tuple[Any, ...], Tuple[float, Dict[str, Any]]] - ] = {} - self._stats_by_task: Dict[str, Dict[str, int]] = {} - self._checkpoint_manager = CheckpointManager( - ".", - "multimodal_ga_checkpoint_", - checkpoint_every=checkpoint_every, - resume=resume, - ) + self.optimization_results: Dict[str, List[FusionSearchResult]] = {} + self.evaluation_errors: Dict[str, int] = {} - # Register mutation ops you want enabled - self.mutation_ops: List[MutationOperator] = [ - MutateFusionOperator(), - MutateRepresentationIndex(), - MutateTreeRotation(), - MutateSubtreeResample(), - # Optional global search mutations: - MutateAddModality(), - MutateDropModality(), - ] + self.hall_of_fame_size = max(1, hall_of_fame_size) + self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} + self._hof_fitness: Dict[str, List[Tuple[float, ...]]] = {} + self.rng = random.Random(random_seed) + self._eval_counter = 0 + self._current_generation = -1 + self._search_start = time.perf_counter() + self.population_size = max(1, population_size) + self.generations = max(1, generations) + self.crossover_probability = crossover_probability + self.mutation_probability = mutation_probability + self.random_seed = random_seed + self._fitness_cache: Dict[str, Dict[Tuple, Tuple[float, ...]]] = {} - # ---------------------------- - # Public entrypoint - # ---------------------------- + self.elite_size = max(0, min(elite_size, self.population_size - 1)) + self.max_workers = max(1, max_workers) + self.batch_size = max(1, batch_size or self.max_workers) + self.early_stopping_patience = early_stopping_patience + self.early_stopping_min_delta = early_stopping_min_delta + self.novelty_breeding = novelty_breeding + self._current_task_name = None + + desired_weights = tuple( + 1.0 if direction == "max" else -1.0 for _, direction in self.objective_specs + ) + self._objective_weights = desired_weights + existing_weights = getattr( + getattr(creator, "FitnessMax", None), "weights", None + ) + if existing_weights != desired_weights: + if hasattr(creator, "Individual"): + del creator.Individual + if hasattr(creator, "FitnessMax"): + del creator.FitnessMax + creator.create("FitnessMax", base.Fitness, weights=desired_weights) + creator.create("Individual", list, fitness=creator.FitnessMax) + elif not hasattr(creator, "Individual"): + creator.create("Individual", list, fitness=creator.FitnessMax) def optimize( self, - max_evaluations_per_task: Optional[int] = None, - # If you later wire this to NodeExecutor, add params like: - # use_node_executor: bool = True, - # max_workers: int = ... - ) -> Dict[str, List[OptimizationResult]]: - self.rng = random.Random(self.random_seed) - self._resume_if_available() - self._ensure_results_initialized() - + ) -> Dict[str, List[FusionSearchResult]]: for task in self.tasks: task_name = task.model.name - if self.debug: - print(f"[GA] Task={task_name} initialization") + self._current_task_name = task_name + self.optimization_results.setdefault(task_name, []) + self.evaluation_errors.setdefault(task_name, 0) - # 1) Initialize population stochastically - population = self._initialize_population(task_name) + self._eval_counter = 0 + self._search_start = time.perf_counter() - eval_budget = 0 - best_seen: Optional[Individual] = None - no_improvement_gens = 0 + population = self._build_initial_population(task_name) + best_ever = None + no_improve = 0 - # 2) Evolution loop for gen in range(self.generations): - if self.debug: - print(f"[GA] Task={task_name}, generation={gen}") - - # Evaluate all individuals (or only unevaluated) + self._current_generation = gen self._evaluate_population(population, task) - # Track best - gen_best = max(population, key=lambda ind: ind.fitness) - if best_seen is None or ( - gen_best.fitness > best_seen.fitness + self.early_stopping_min_delta - ): - best_seen = copy.deepcopy(gen_best) - no_improvement_gens = 0 - else: - no_improvement_gens += 1 - - # Optional budget stop - eval_budget += self._stats_by_task[task_name][ - "last_generation_executed" - ] - if ( - max_evaluations_per_task is not None - and eval_budget >= max_evaluations_per_task - ): - break - if ( - self.early_stopping_patience is not None - and no_improvement_gens >= self.early_stopping_patience - ): - if self.debug: - print( - f"[GA] Task={task_name} early stopping after " - f"{no_improvement_gens} no-improvement generations" - ) - break - - # Selection + variation + replacement - next_population = self._elitism(population) - seen_genotypes = { - self._individual_signature(ind) for ind in next_population - } - duplicate_retry_budget = self.population_size * 10 - duplicate_retries = 0 - while len(next_population) < self.population_size: - p1 = self._tournament_select(population) - p2 = self._tournament_select(population) - - if self.rng.random() < self.crossover_rate: - c1, c2 = self._crossover(p1, p2) - else: - c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) - - c1 = self._mutate(c1, task_name) - c2 = self._mutate(c2, task_name) - - added_c1 = self._append_if_unique( - next_population, c1, seen_genotypes, self.population_size - ) - added_c2 = self._append_if_unique( - next_population, c2, seen_genotypes, self.population_size + if self.is_multi_objective: + front = tools.sortNondominated( + population, len(population), first_front_only=True + )[0] + front_signature = frozenset( + self._genome_signature(ind[0]) for ind in front ) - if not added_c1 and not added_c2: - duplicate_retries += 1 + if best_ever is None or front_signature != best_ever: + best_ever = front_signature + no_improve = 0 else: - duplicate_retries = 0 - - if duplicate_retries >= duplicate_retry_budget: - break - - while len(next_population) < self.population_size: - immigrant = self._sample_random_individual(task_name) - self._append_if_unique( - next_population, immigrant, seen_genotypes, self.population_size - ) - - population = next_population - - self._checkpoint_manager.checkpoint_if_due( - self.optimization_results, "eval_count_by_task" - ) - - if self.debug and best_seen is not None: - print(f"[GA] Task={task_name}, best_fitness={best_seen.fitness:.6f}") - - return self.optimization_results - - def optimize_parallel( - self, - max_combinations: Optional[int] = None, - max_workers: int = 2, - batch_size: int = 8, - ) -> Dict[str, List[OptimizationResult]]: - self.rng = random.Random(self.random_seed) - self._resume_if_available() - self._ensure_results_initialized() - - for task in self.tasks: - task_name = task.model.name - if self.debug: - print(f"[GA-P] Task={task_name} initialization") - - population = self._initialize_population(task_name) - eval_budget = 0 - best_seen: Optional[Individual] = None - no_improvement_gens = 0 + no_improve += 1 + debug_msg = f"front_size={len(front)}" + else: + gen_best = max(population, key=lambda ind: ind.fitness.values[0]) + if ( + best_ever is None + or gen_best.fitness.values[0] + > best_ever.fitness.values[0] + self.early_stopping_min_delta + ): + best_ever = self._clone_individual(gen_best) + no_improve = 0 + else: + no_improve += 1 + debug_msg = f"best={gen_best.fitness.values[0]:.4f}" - for gen in range(self.generations): if self.debug: - print(f"[GA-P] Task={task_name}, generation={gen}") - - self._evaluate_population_parallel( - population, - task, - max_workers=max_workers, - batch_size=max(1, batch_size), - ) - - gen_best = max(population, key=lambda ind: ind.fitness) - if best_seen is None or ( - gen_best.fitness > best_seen.fitness + self.early_stopping_min_delta - ): - best_seen = copy.deepcopy(gen_best) - no_improvement_gens = 0 - else: - no_improvement_gens += 1 + print( + f"[GA] task={task_name} gen={gen} {debug_msg} " + f"no_improve={no_improve} " + f"errors={self.evaluation_errors.get(task_name, 0)}" + ) - eval_budget += self._stats_by_task[task_name][ - "last_generation_executed" - ] - if max_combinations is not None and eval_budget >= max_combinations: - break - if ( + stagnated = ( self.early_stopping_patience is not None - and no_improvement_gens >= self.early_stopping_patience - ): - if self.debug: + and no_improve >= self.early_stopping_patience + ) + if stagnated or gen == self.generations - 1: + if self.debug and stagnated: print( - f"[GA-P] Task={task_name} early stopping after " - f"{no_improvement_gens} no-improvement generations" + f"[GA] task={task_name} early stopping after " + f"{no_improve} generations without improvement" ) break - next_population = self._elitism(population) - seen_genotypes = { - self._individual_signature(ind) for ind in next_population - } - duplicate_retry_budget = self.population_size * 10 - duplicate_retries = 0 - while len(next_population) < self.population_size: - p1 = self._tournament_select(population) - p2 = self._tournament_select(population) - - if self.rng.random() < self.crossover_rate: - c1, c2 = self._crossover(p1, p2) - else: - c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) - - c1 = self._mutate(c1, task_name) - c2 = self._mutate(c2, task_name) + population = self._next_generation(population, task_name, task) - added_c1 = self._append_if_unique( - next_population, c1, seen_genotypes, self.population_size - ) - added_c2 = self._append_if_unique( - next_population, c2, seen_genotypes, self.population_size - ) - if not added_c1 and not added_c2: - duplicate_retries += 1 - else: - duplicate_retries = 0 - - if duplicate_retries >= duplicate_retry_budget: - break - - while len(next_population) < self.population_size: - immigrant = self._sample_random_individual(task_name) - self._append_if_unique( - next_population, immigrant, seen_genotypes, self.population_size - ) - - population = next_population - - self._checkpoint_manager.checkpoint_if_due( - self.optimization_results, "eval_count_by_task" - ) + return self.optimization_results - if self.debug and best_seen is not None: - print(f"[GA-P] Task={task_name}, best_fitness={best_seen.fitness:.6f}") + def _make_individual(self, genome: DagGenome): + return creator.Individual([genome]) - return self.optimization_results + def _clone_individual(self, ind): + clone = creator.Individual([copy.deepcopy(ind[0])]) + if ind.fitness.valid: + clone.fitness.values = ind.fitness.values + return clone - # ---------------------------- - # Initialization - # ---------------------------- + def _append_if_unique( + self, + population: List[Any], + genome: DagGenome, + seen_signatures: set, + ) -> bool: + if len(population) >= self.population_size: + return False + sig = self._genome_signature(genome) + if sig in seen_signatures: + return False + seen_signatures.add(sig) + population.append(self._make_individual(genome)) + return True - def _initialize_population(self, task_name: str) -> List[Individual]: - population = [] - seen_signatures = set() - retry_budget = self.population_size * 10 + def _build_initial_population(self, task_name: str) -> List[Any]: + population: List[Any] = [] + seen: set = set() + retry_budget = max(20, self.population_size * 10) retries = 0 - for _ in range(self.population_size): - while retries < retry_budget: - candidate = self._sample_random_individual(task_name) - signature = self._individual_signature(candidate) - if signature not in seen_signatures: - seen_signatures.add(signature) - population.append(candidate) - break + while len(population) < self.population_size and retries < retry_budget: + genome = self._random_genome(task_name) + if self._append_if_unique(population, genome, seen): + retries = 0 + else: retries += 1 - if retries >= retry_budget: - # Fall back to potentially duplicated random individuals rather than - # stalling initialization on very small search spaces. - population.append(self._sample_random_individual(task_name)) + while len(population) < self.population_size: + population.append(self._make_individual(self._random_genome(task_name))) return population - def _sample_random_individual(self, task_name: str) -> Individual: - leaves = self._sample_leaf_set(task_name) - tree = self._sample_random_binary_tree(len(leaves)) - fusion_ops = {} - self._assign_random_fusion_ops(tree, fusion_ops, path="") - return Individual(leaves=leaves, tree=tree, fusion_ops=fusion_ops) - - def _sample_leaf_set(self, task_name: str) -> List[Tuple[str, int]]: - # Only sample from modalities that have at least one representation - # for the current task. - task_reps = self.k_best_representations.get(task_name, {}) - available_modalities = [ - m.modality_id - for m in self.modalities - if len(task_reps.get(m.modality_id, [])) > 0 - ] - - if len(available_modalities) < 2: - raise ValueError( - f"Need at least 2 modalities with non-empty representations for task " - f"'{task_name}', found {len(available_modalities)}." + def _next_generation( + self, population: List[Any], task_name: str, task: Task + ) -> List[Any]: + if self.is_multi_objective: + offspring = self._breed_offspring( + population, task_name, seen=self._novelty_archive(task_name) ) - - # Clamp sampling range to available modalities. - lower = max(2, self.min_modalities) - upper = min(self.max_modalities, len(available_modalities)) - if lower > upper: - lower = upper - - r = self.rng.randint(lower, upper) - chosen_modalities = self.rng.sample(available_modalities, r) - - # Sample one representation index from top-k for each chosen modality. - return [ - (mod_id, self.rng.randrange(len(task_reps[mod_id]))) - for mod_id in chosen_modalities - ] - - def _sample_random_binary_tree(self, n_leaves: int) -> Tree: - nodes: List[Tree] = list(range(n_leaves)) - while len(nodes) > 1: - i, j = self.rng.sample(range(len(nodes)), 2) - a = nodes.pop(max(i, j)) - b = nodes.pop(min(i, j)) - nodes.append((a, b)) - return nodes[0] - - def _assign_random_fusion_ops( - self, subtree: Tree, fusion_ops: Dict[str, Any], path: str - ): - if isinstance(subtree, int): + self._evaluate_population(offspring, task) + combined = list(population) + list(offspring) + return list(tools.selNSGA2(combined, self.population_size)) + + ranked = sorted(population, key=lambda ind: ind.fitness.values[0], reverse=True) + elite = [self._clone_individual(ind) for ind in ranked[: self.elite_size]] + seen = {self._genome_signature(ind[0]) for ind in elite} + seen |= self._novelty_archive(task_name) + return self._breed_offspring(population, task_name, initial=elite, seen=seen) + + def _novelty_archive(self, task_name: str) -> set: + if not self.novelty_breeding: + return set() + return set(self._fitness_cache.get(task_name, {})) + + def _breed_offspring( + self, + population: List[Any], + task_name: str, + initial: Optional[List[Any]] = None, + seen: Optional[set] = None, + ) -> List[Any]: + next_population = list(initial) if initial else [] + seen = set(seen) if seen else set() + + retry_budget = max(20, self.population_size * 10) + retries = 0 + tournsize = max(1, min(3, len(population))) + while len(next_population) < self.population_size and retries < retry_budget: + p1, p2 = tools.selTournament(population, 2, tournsize=tournsize) + + if self.rng.random() < self.crossover_probability: + g1, g2 = self._crossover_genomes(p1[0], p2[0]) + else: + g1, g2 = copy.deepcopy(p1[0]), copy.deepcopy(p2[0]) + + if self.rng.random() < self.mutation_probability: + g1 = self._mutate_genome(g1, task_name) + if self.rng.random() < self.mutation_probability: + g2 = self._mutate_genome(g2, task_name) + + added1 = self._append_if_unique(next_population, g1, seen) + added2 = self._append_if_unique(next_population, g2, seen) + retries = 0 if (added1 or added2) else retries + 1 + + while len(next_population) < self.population_size: + genome = self._random_genome(task_name) + if not self._append_if_unique(next_population, genome, seen): + next_population.append(self._make_individual(genome)) + + return next_population + + def _evaluate_population(self, population: List[Any], task: Task) -> None: + to_evaluate = [ind for ind in population if not ind.fitness.valid] + if not to_evaluate: return - fusion_ops[path] = self.rng.choice(self.fusion_operators) - left, right = subtree - self._assign_random_fusion_ops(left, fusion_ops, path + "L") - self._assign_random_fusion_ops(right, fusion_ops, path + "R") - - # ---------------------------- - # Evaluation - # ---------------------------- - - def _evaluate_population(self, population: List[Individual], task: Task) -> None: - # Placeholder: currently serial + per-individual eval. - # You can batch DAGs and evaluate with NodeExecutor later. - task_name = task.model.name - cache = self._fitness_cache_by_task.setdefault(task_name, {}) - evals_before = self._stats_by_task[task_name]["executed_evaluations"] - for ind in population: - # Already evaluated and unchanged. - if ind.fitness != float("-inf"): - continue - - signature = self._individual_signature(ind) - cached = cache.get(signature) - if cached is not None: - self._stats_by_task[task_name]["cache_hits"] += 1 - ind.fitness, ind.objective = cached[0], dict(cached[1]) - continue - - dag = self._individual_to_dag(ind) - ind.dag = dag - modalities = [ - self.k_best_representations[task_name][mod_id][repr_idx] - for mod_id, repr_idx in ind.leaves - ] - self._stats_by_task[task_name]["executed_evaluations"] += 1 - fitness, objective = self._evaluate_individual(dag, task, modalities) - ind.fitness = fitness - ind.objective = objective - cache[signature] = (fitness, dict(objective)) - self._record_result(task_name, dag, objective) - self._stats_by_task[task_name]["last_generation_executed"] = ( - self._stats_by_task[task_name]["executed_evaluations"] - evals_before - ) + if self.max_workers > 1 and len(to_evaluate) > 1: + self._evaluate_individuals_parallel(to_evaluate, task) + else: + for ind in to_evaluate: + fitness = self._evaluate_genome(ind[0], task) + ind.fitness.values = fitness - def _evaluate_population_parallel( - self, - population: List[Individual], - task: Task, - max_workers: int, - batch_size: int, + def _evaluate_individuals_parallel( + self, individuals: List[Any], task: Task ) -> None: - ctx = mp.get_context("spawn") task_name = task.model.name - cache = self._fitness_cache_by_task.setdefault(task_name, {}) - evals_before = self._stats_by_task[task_name]["executed_evaluations"] - task_pickle = pickle.dumps(copy.deepcopy(task)) - futures = {} - pending_followers: Dict[Tuple[Any, ...], List[Individual]] = {} - - def _collect_ready(done_futures): - for done in done_futures: - ind, dag, signature = futures.pop(done) - fitness, objective = done.result() - ind.fitness = fitness - ind.objective = objective - cache[signature] = (fitness, dict(objective)) - self._record_result(task_name, dag, objective) - for follower in pending_followers.pop(signature, []): - follower.fitness = fitness - follower.objective = dict(objective) - - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor: - for ind in population: - # Already evaluated and unchanged. - if ind.fitness != float("-inf"): - continue - - signature = self._individual_signature(ind) - cached = cache.get(signature) + cache = self._fitness_cache.setdefault(task_name, {}) + ctx = mp.get_context("spawn") + task_bytes = pickle.dumps(task) + futures: Dict[Any, Tuple[Any, RepresentationDag, Tuple]] = {} + pending_followers: Dict[Tuple, List[Any]] = {} + + def _drain(done_futures): + for fut in done_futures: + ind, dag, sig = futures.pop(fut) + fitness, payload, error = fut.result() + ind.fitness.values = fitness + self._record_evaluation(task_name, dag, payload, error) + cache[sig] = fitness + for follower in pending_followers.pop(sig, []): + follower.fitness.values = fitness + + with ProcessPoolExecutor( + max_workers=self.max_workers, mp_context=ctx + ) as executor: + for ind in individuals: + genome = ind[0] + sig = self._genome_signature(genome) + cached = cache.get(sig) if cached is not None: - self._stats_by_task[task_name]["cache_hits"] += 1 - ind.fitness, ind.objective = cached[0], dict(cached[1]) + ind.fitness.values = cached continue - - if signature in pending_followers: - pending_followers[signature].append(ind) - self._stats_by_task[task_name]["pending_signature_hits"] += 1 + if sig in pending_followers: + pending_followers[sig].append(ind) continue - dag = self._individual_to_dag(ind) - ind.dag = dag - modalities = [ - self.k_best_representations[task_name][mod_id][repr_idx] - for mod_id, repr_idx in ind.leaves - ] - + dag = self._genome_to_dag(genome) + modalities = list( + chain.from_iterable(self.k_best_representations[task_name].values()) + ) fut = executor.submit( - _evaluate_individual_worker, + _evaluate_dag_worker, pickle.dumps(dag), - task_pickle, + task_bytes, pickle.dumps(modalities), - self.metric_name, + self.objective_specs, ) - futures[fut] = (ind, dag, signature) - pending_followers[signature] = [] - self._stats_by_task[task_name]["executed_evaluations"] += 1 + futures[fut] = (ind, dag, sig) + pending_followers[sig] = [] - if len(futures) >= batch_size: + if len(futures) >= self.batch_size: done, _ = wait(set(futures.keys()), return_when=FIRST_COMPLETED) - _collect_ready(done) + _drain(done) if futures: done, _ = wait(set(futures.keys())) - _collect_ready(done) - self._stats_by_task[task_name]["last_generation_executed"] = ( - self._stats_by_task[task_name]["executed_evaluations"] - evals_before - ) + _drain(done) - def _evaluate_individual( - self, dag: RepresentationDag, task: Task, modalities: List[Modality] - ) -> Tuple[float, Dict[str, float]]: - start_time = time.time() - fused_representation = dag.execute(modalities, task) - task_start_time = time.time() - scores = task.run(fused_representation.data) - task_end_time = time.time() - runtime = time.time() - start_time - fitness = scores[1].average_scores[self.metric_name] - objective = { - "train_score": scores[0].average_scores, - "val_score": scores[1].average_scores, - "test_score": scores[2].average_scores, - "val": fitness, - "runtime": runtime, - "representation_time": fused_representation.transform_time, - "task_time": task_end_time - task_start_time, - } - return fitness, objective - - def _record_result( - self, task_name: str, dag: RepresentationDag, objective: Dict[str, Any] + def _record_evaluation( + self, + task_name: str, + dag: RepresentationDag, + payload: Optional[Dict[str, Any]], + error: Optional[str], ) -> None: - key = str(dag.compute_full_node_signature(dag.root_node_id)) - if key in self._seen_results_by_task[task_name]: + self.optimization_results.setdefault(task_name, []) + if error is not None or payload is None: + self.evaluation_errors[task_name] = ( + self.evaluation_errors.get(task_name, 0) + 1 + ) + if self.debug and error is not None: + last_line = error.strip().splitlines()[-1] if error.strip() else error + print( + f"[GA] genome evaluation failed for task={task_name}: {last_line}" + ) return - result = OptimizationResult( + + result = FusionSearchResult( dag=dag, - val_score=objective.get("val_score", {}), - train_score=objective.get("train_score", {}), - test_score=objective.get("test_score", {}), - runtime=objective.get("runtime", 0.0), - representation_time=objective.get("representation_time", 0.0), - task_time=objective.get("task_time", 0.0), + train_score=payload["train_score"], + val_score=payload["val_score"], + test_score=payload["test_score"], + train_fold_scores=payload.get("train_fold_scores", {}), + val_fold_scores=payload.get("val_fold_scores", {}), + test_fold_scores=payload.get("test_fold_scores", {}), + task_timing=payload.get("task_timing", {}), + runtime=payload["runtime"], + task_time=payload["task_time"], + representation_time=payload["representation_time"], task_name=task_name, + generation=self._current_generation, + eval_index=self._eval_counter, + t_since_search_start_s=time.perf_counter() - self._search_start, + t_eval_end_unix=time.time(), ) - self._seen_results_by_task[task_name][key] = result - self.optimization_results[task_name].append(result) + self.optimization_results.setdefault(task_name, []).append(result) + self._update_hall_of_fame(task_name, result) + self._eval_counter += 1 + + def _dominates(self, a: Tuple[float, ...], b: Tuple[float, ...]) -> bool: + """True if objective tuple `a` Pareto-dominates `b`, direction-aware.""" + wa = [w * v for w, v in zip(self._objective_weights, a)] + wb = [w * v for w, v in zip(self._objective_weights, b)] + return all(x >= y for x, y in zip(wa, wb)) and any( + x > y for x, y in zip(wa, wb) + ) + + def _update_hall_of_fame(self, task_name: str, result: FusionSearchResult) -> None: + timing = { + "runtime": result.runtime, + "task_time": result.task_time, + "representation_time": result.representation_time, + } + try: + fitness = tuple( + _objective_value(name, result.val_score, timing) + for name, _ in self.objective_specs + ) + except KeyError: + return - def get_task_stats(self, task_name: str) -> Dict[str, int]: - return dict(self._stats_by_task.get(task_name, {})) + hof = self.hall_of_fame.setdefault(task_name, []) + fits = self._hof_fitness.setdefault(task_name, []) - def _individual_signature(self, ind: Individual) -> Tuple[Any, ...]: - def _normalize_tree(subtree: Tree) -> Any: - if isinstance(subtree, int): - return subtree - left, right = subtree - return (_normalize_tree(left), _normalize_tree(right)) + if self.is_multi_objective: + if any(self._dominates(f, fitness) or f == fitness for f in fits): + return + keep = [i for i, f in enumerate(fits) if not self._dominates(fitness, f)] + self.hall_of_fame[task_name] = [hof[i] for i in keep] + [result] + self._hof_fitness[task_name] = [fits[i] for i in keep] + [fitness] + return - leaves_sig = tuple(ind.leaves) - tree_sig = _normalize_tree(ind.tree) - ops_sig = tuple( - sorted((path, op.__name__) for path, op in ind.fusion_ops.items()) + weight = self._objective_weights[0] + hof.append(result) + fits.append(fitness) + order = sorted( + range(len(fits)), key=lambda i: weight * fits[i][0], reverse=True ) - return leaves_sig, tree_sig, ops_sig + order = order[: self.hall_of_fame_size] + self.hall_of_fame[task_name] = [hof[i] for i in order] + self._hof_fitness[task_name] = [fits[i] for i in order] - def _append_if_unique( - self, - population: List[Individual], - candidate: Individual, - seen_genotypes: set, - target_size: int, - ) -> bool: - if len(population) >= target_size: - return False - sig = self._individual_signature(candidate) - if sig in seen_genotypes: - return False - seen_genotypes.add(sig) - population.append(candidate) - return True + def get_hall_of_fame(self, task_name: str) -> List[FusionSearchResult]: + return list(self.hall_of_fame.get(task_name, [])) - # ---------------------------- - # Selection / crossover / mutation - # ---------------------------- + def _extract_k_best(self, unimodal_results) -> Dict[str, Dict[str, List[Any]]]: + k_best = {} + for task in self.tasks: + name = task.model.name + k_best[name] = {} + for modality in self.modalities: + _, cached_data = unimodal_results.get_k_best_results( + modality, task, self.metric_name + ) + k_best[name][modality.modality_id] = cached_data + return k_best - def _elitism(self, population: List[Individual]) -> List[Individual]: - ranked = sorted(population, key=lambda x: x.fitness, reverse=True) - return [copy.deepcopy(ind) for ind in ranked[: self.elite_size]] + def _available_modality_ids(self, task_name: str) -> List[Any]: + reps = self.k_best_representations[task_name] + return [ + m.modality_id + for m in self.modalities + if len(reps.get(m.modality_id, [])) > 0 + ] - def _tournament_select(self, population: List[Individual]) -> Individual: - contestants = self.rng.sample( - population, k=min(self.tournament_size, len(population)) - ) - return copy.deepcopy(max(contestants, key=lambda x: x.fitness)) - - def _crossover( - self, p1: Individual, p2: Individual - ) -> Tuple[Individual, Individual]: - """ - Skeleton crossover: - - Safe version assumes compatible leaves (same modality set/order). - - If incompatible, fallback to op-only crossover. - """ - c1, c2 = copy.deepcopy(p1), copy.deepcopy(p2) - - if self._compatible_for_subtree_crossover(c1, c2): - c1_paths = _collect_internal_paths(c1.tree) - c2_paths = _collect_internal_paths(c2.tree) - if c1_paths and c2_paths: - path1 = self.rng.choice(c1_paths) - path2 = self.rng.choice(c2_paths) - subtree1 = _get_subtree(c1.tree, path1) - subtree2 = _get_subtree(c2.tree, path2) - c1.tree = _replace_subtree(c1.tree, path1, subtree2) - c2.tree = _replace_subtree(c2.tree, path2, subtree1) - c1.fusion_ops = _rebuild_fusion_ops( - c1.tree, - c1.fusion_ops, - self.rng, - self.fusion_operators, - randomized_prefixes=[path1], - ) - c2.fusion_ops = _rebuild_fusion_ops( - c2.tree, - c2.fusion_ops, - self.rng, - self.fusion_operators, - randomized_prefixes=[path2], - ) + def _leaf_capacity(self, task_name: str) -> int: + reps = self.k_best_representations[task_name] + ids = self._available_modality_ids(task_name) + if not self.allow_repeated_modalities: + return len(ids) + return sum(len(reps[mid]) for mid in ids) + + def _random_genome(self, task_name: str) -> DagGenome: + reps = self.k_best_representations[task_name] + available_modality_ids = self._available_modality_ids(task_name) + capacity = self._leaf_capacity(task_name) + if capacity < self.min_modalities: + raise ValueError( + f"Need at least {self.min_modalities} distinct leaves for task " + f"'{task_name}', but only {capacity} are available across " + f"{len(available_modality_ids)} modalities." + ) + + upper = min(self.max_modalities, capacity) + lower = min(self.min_modalities, upper) + r = self.rng.randint(lower, upper) + + if self.allow_repeated_modalities: + pool = [ + (mid, idx) + for mid in available_modality_ids + for idx in range(len(reps[mid])) + ] + leaves = self.rng.sample(pool, r) else: - # op-only crossover: mix operator assignments - all_keys = set(c1.fusion_ops.keys()) | set(c2.fusion_ops.keys()) - for k in all_keys: - if self.rng.random() < 0.5: - if k in c2.fusion_ops: - c1.fusion_ops[k] = c2.fusion_ops[k] - else: - if k in c1.fusion_ops: - c2.fusion_ops[k] = c1.fusion_ops[k] + chosen = self.rng.sample(available_modality_ids, r) + leaves = [(mid, self.rng.randrange(len(reps[mid]))) for mid in chosen] - # invalidate stale eval - c1.fitness, c1.objective, c1.dag = float("-inf"), {}, None - c2.fitness, c2.objective, c2.dag = float("-inf"), {}, None - return c1, c2 + tree = self._random_binary_tree(len(leaves)) + fusion_ops = {} + self._assign_fusion_ops(tree, fusion_ops, "") + return DagGenome(leaves=leaves, tree=tree, fusion_ops=fusion_ops) - def _compatible_for_subtree_crossover(self, a: Individual, b: Individual) -> bool: - return [m for m, _ in a.leaves] == [m for m, _ in b.leaves] and len( - a.leaves - ) == len(b.leaves) - - def _mutate(self, ind: Individual, task_name: str) -> Individual: - out = copy.deepcopy(ind) - if self.rng.random() >= self.mutation_rate: - return out - - op = self.rng.choice(self.mutation_ops) - context = { - "task_name": task_name, - "fusion_operators": self.fusion_operators, - "k_best_representations": self.k_best_representations, - "min_modalities": self.min_modalities, - "max_modalities": self.max_modalities, - "all_modalities": [m.modality_id for m in self.modalities], - } - out = op(out, self.rng, context) - out.fitness, out.objective, out.dag = float("-inf"), {}, None - return out + def _internal_paths(self, tree): + return _collect_internal_paths(tree) - # ---------------------------- - # Genome -> DAG - # ---------------------------- + def _random_binary_tree(self, n: int) -> Tree: + nodes: List[Tree] = list(range(n)) + while len(nodes) > 1: + i, j = self.rng.sample(range(len(nodes)), 2) + a, b = nodes.pop(max(i, j)), nodes.pop(min(i, j)) + nodes.append((a, b)) + return nodes[0] - def _individual_to_dag(self, ind: Individual) -> RepresentationDag: + def _assign_fusion_ops(self, subtree: Tree, ops: Dict[str, Any], path: str) -> None: + if isinstance(subtree, int): + return + ops[path] = self.rng.choice(self.fusion_operators) + left, right = subtree + self._assign_fusion_ops(left, ops, path + "L") + self._assign_fusion_ops(right, ops, path + "R") + + def _genome_to_dag(self, genome: DagGenome) -> RepresentationDag: builder = RepresentationDAGBuilder() - leaf_ids = [] - for modality_id, repr_idx in ind.leaves: - # Leaves reference the cached transformed modality directly. - # We collapse the unimodal path from raw->representation here, so - # downstream execution only needs to run multimodal fusion nodes. - leaf_ids.append(builder.create_leaf_node(modality_id, repr_idx)) + leaf_ids = [ + builder.create_leaf_node(mod_id, repr_idx) + for mod_id, repr_idx in genome.leaves + ] def build(subtree: Tree, path: str) -> str: if isinstance(subtree, int): @@ -885,351 +769,255 @@ def build(subtree: Tree, path: str) -> str: left, right = subtree left_id = build(left, path + "L") right_id = build(right, path + "R") - op_cls = ind.fusion_ops[path] + op_cls = genome.fusion_ops[path] op = op_cls() return builder.create_operation_node( op.__class__, [left_id, right_id], op.get_current_parameters() ) - root_id = build(ind.tree, "") - dag = builder.build(root_id) - return self._collapse_cached_unimodal_nodes(dag) - - def _collapse_cached_unimodal_nodes( - self, dag: RepresentationDag - ) -> RepresentationDag: - """ - Remove unary unimodal nodes that originate at leaves. - - In GA multimodal search, leaf inputs are already transformed unimodal - representations taken from the unimodal optimizer cache. Any unary chain - attached to those leaves is redundant and can be bypassed to ensure only - multimodal (fusion) operations are executed. - """ - node_by_id = {node.node_id: copy.deepcopy(node) for node in dag.nodes} - if dag.root_node_id not in node_by_id: - return dag - - changed = True - while changed: - changed = False - node_ids = list(node_by_id.keys()) - for node_id in node_ids: - if node_id not in node_by_id: - continue - node = node_by_id[node_id] - if len(node.inputs) != 1: - continue + return builder.build(build(genome.tree, "")) - parent_id = node.inputs[0] - parent = node_by_id.get(parent_id) - if parent is None: - continue - if parent.inputs: - continue + def _genome_signature(self, g: DagGenome) -> Tuple: + def norm(t: Tree): + return t if isinstance(t, int) else (norm(t[0]), norm(t[1])) - # Bypass unary node by rewiring all consumers to the leaf input. - for consumer in node_by_id.values(): - consumer.inputs = [ - parent_id if input_id == node_id else input_id - for input_id in consumer.inputs - ] - if dag.root_node_id == node_id: - dag.root_node_id = parent_id - del node_by_id[node_id] - changed = True - - return RepresentationDag( - list(node_by_id.values()), dag.root_node_id, dag.dag_id + return ( + tuple(g.leaves), + norm(g.tree), + tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), ) - # ---------------------------- - # Existing helper compatibility - # ---------------------------- - - def _extract_k_best_representations( - self, unimodal_optimization_results: Any - ) -> Dict[str, Dict[str, List[Any]]]: - k_best = {} - for task in self.tasks: - task_name = task.model.name - k_best[task_name] = {} - for modality in self.modalities: - _, cached_data = unimodal_optimization_results.get_k_best_results( - modality, task, self.metric_name - ) - k_best[task_name][modality.modality_id] = cached_data - return k_best - - def _resume_if_available(self) -> None: - loaded = self._checkpoint_manager.resume_from_checkpoint( - "eval_count_by_task", - lambda results: { - t.model.name: len(results.get(t.model.name, [])) for t in self.tasks - }, + def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: + task_name = task.model.name + sig = self._genome_signature(genome) + cache = self._fitness_cache.setdefault(task_name, {}) + if sig in cache: + return cache[sig] + + dag = self._genome_to_dag(genome) + modalities = list( + chain.from_iterable(self.k_best_representations[task_name].values()) ) - if loaded: - results, _, _ = loaded - self.optimization_results = results - def _ensure_results_initialized(self): - if not isinstance(self.optimization_results, dict): - self.optimization_results = {} - for task in self.tasks: - task_name = task.model.name - self.optimization_results.setdefault(task_name, []) - seen = self._seen_results_by_task.setdefault(task_name, {}) - self._fitness_cache_by_task.setdefault(task_name, {}) - self._stats_by_task.setdefault( - task_name, - { - "executed_evaluations": 0, - "last_generation_executed": 0, - "cache_hits": 0, - "pending_signature_hits": 0, - }, + try: + fitness, payload = _evaluate_genome_body( + dag, task, modalities, self.objective_specs ) - if self.optimization_results[task_name]: - for result in self.optimization_results[task_name]: - if result.dag is None: - continue - key = str( - result.dag.compute_full_node_signature(result.dag.root_node_id) - ) - seen.setdefault(key, result) - - def _to_optimization_results( - self, individuals: List[Individual], task_name: str - ) -> List[OptimizationResult]: - out = [] - for ind in individuals: - # Keep consistent with old output shape: OptimizationResult list per task - out.append( - OptimizationResult( - dag=ind.dag, - val_score={ - self.metric_name: ind.objective.get("val", float("-inf")) - }, - train_score={}, - test_score={}, - runtime=ind.objective.get("runtime", 0.0), - representation_time=ind.objective.get("representation_time", 0.0), - task_time=ind.objective.get("task_time", 0.0), - task_name=task_name, + error = None + if fitness is None: + fitness = _failure_fitness(self.objective_specs) + except Exception: + fitness = _failure_fitness(self.objective_specs) + payload, error = None, traceback.format_exc() + + self._record_evaluation(task_name, dag, payload, error) + cache[sig] = fitness + return fitness + + def _crossover_genomes( + self, g1: DagGenome, g2: DagGenome + ) -> Tuple[DagGenome, DagGenome]: + c1, c2 = copy.deepcopy(g1), copy.deepcopy(g2) + + if c1.leaves == c2.leaves: + paths1 = self._internal_paths(c1.tree) + paths2 = self._internal_paths(c2.tree) + if paths1 and paths2: + path1 = self.rng.choice(paths1) + path2 = self.rng.choice(paths2) + subtree1 = _get_subtree(c1.tree, path1) + subtree2 = _get_subtree(c2.tree, path2) + c1.tree = _replace_subtree(c1.tree, path1, subtree2) + c2.tree = _replace_subtree(c2.tree, path2, subtree1) + c1.fusion_ops = _rebuild_fusion_ops( + c1.tree, + {**c1.fusion_ops, **c2.fusion_ops}, + self.rng, + self.fusion_operators, + randomized_prefixes=[path1], ) - ) - return out - - -# ---------------------------- -# Mutation operators (outline) -# ---------------------------- - - -class MutateFusionOperator(MutationOperator): - """ - Change one internal fusion op class. - """ - - name = "mutate_fusion_operator" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - if not out.fusion_ops: - return out - key = rng.choice(list(out.fusion_ops.keys())) - current = out.fusion_ops[key] - candidates = [op for op in context["fusion_operators"] if op != current] - if candidates: - out.fusion_ops[key] = rng.choice(candidates) - return out - - -class MutateRepresentationIndex(MutationOperator): - """ - Keep modality set fixed, change repr_idx for one modality leaf. - """ - - name = "mutate_representation_index" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - if not out.leaves: - return out - i = rng.randrange(len(out.leaves)) - modality_id, cur_idx = out.leaves[i] - task_name = context["task_name"] - reps = context["k_best_representations"][task_name][modality_id] - if len(reps) <= 1: - return out - new_idx = rng.randrange(len(reps)) - while new_idx == cur_idx: - new_idx = rng.randrange(len(reps)) - out.leaves[i] = (modality_id, new_idx) - return out - - -class MutateTreeRotation(MutationOperator): - """ - Local reassociation (e.g., ((a,b),c) <-> (a,(b,c))) on a random eligible subtree. - """ - - name = "mutate_tree_rotation" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - candidates: List[Tuple[str, str]] = [] - for path in _collect_internal_paths(out.tree): - subtree = _get_subtree(out.tree, path) - if isinstance(subtree, int): - continue - left, right = subtree - if not isinstance(left, int): - candidates.append((path, "left")) - if not isinstance(right, int): - candidates.append((path, "right")) - - if not candidates: - return out - - path, direction = rng.choice(candidates) - subtree = _get_subtree(out.tree, path) - left, right = subtree + c2.fusion_ops = _rebuild_fusion_ops( + c2.tree, + {**c2.fusion_ops, **c1.fusion_ops}, + self.rng, + self.fusion_operators, + randomized_prefixes=[path2], + ) + return c1, c2 + + shared_paths = set(c1.fusion_ops) & set(c2.fusion_ops) + for path in shared_paths: + if self.rng.random() < 0.5: + c1.fusion_ops[path], c2.fusion_ops[path] = ( + c2.fusion_ops[path], + c1.fusion_ops[path], + ) + return c1, c2 - if direction == "left": - left_left, left_right = left - rotated = (left_left, (left_right, right)) + def _mutate_genome(self, g: DagGenome, task_name: str) -> DagGenome: + op = self.rng.choice( + [ + self._mutate_change_fusion, + lambda gg: self._mutate_swap_leaf_repr(gg, task_name), + lambda gg: self.mutate_add_leaf(gg, task_name), + self.mutate_remove_leaf, + self.mutate_replace_subtree, + ] + ) + return op(g) + + def _mutate_change_fusion(self, g: DagGenome) -> DagGenome: + g = copy.deepcopy(g) + paths = [p for p in g.fusion_ops] + if not paths: + return g + path = self.rng.choice(paths) + choices = [op for op in self.fusion_operators if op != g.fusion_ops[path]] + if choices: + g.fusion_ops[path] = self.rng.choice(choices) + return g + + def _mutate_swap_leaf_repr(self, g: DagGenome, task_name: str) -> DagGenome: + g = copy.deepcopy(g) + i = self.rng.randrange(len(g.leaves)) + mod_id, current = g.leaves[i] + k = len(self.k_best_representations[task_name][mod_id]) + if k <= 1: + return g + taken = { + idx for j, (mid, idx) in enumerate(g.leaves) if mid == mod_id and j != i + } + choices = [idx for idx in range(k) if idx != current and idx not in taken] + if not choices: + return g + g.leaves[i] = (mod_id, self.rng.choice(choices)) + return g + + def mutate_add_leaf(self, g: DagGenome, task_name: str) -> DagGenome: + if len(g.leaves) >= self.max_modalities: + return g + g = copy.deepcopy(g) + reps = self.k_best_representations[task_name] + if self.allow_repeated_modalities: + existing_leaves = set(g.leaves) + candidates = [ + (mid, idx) + for mid in self._available_modality_ids(task_name) + for idx in range(len(reps[mid])) + if (mid, idx) not in existing_leaves + ] + if not candidates: + return g + new_leaf = self.rng.choice(candidates) else: - right_left, right_right = right - rotated = ((left, right_left), right_right) - - out.tree = _replace_subtree(out.tree, path, rotated) - out.fusion_ops = _rebuild_fusion_ops( - out.tree, - out.fusion_ops, - rng, - context["fusion_operators"], - randomized_prefixes=[path], + existing = {l[0] for l in g.leaves} + available = [ + m.modality_id + for m in self.modalities + if m.modality_id not in existing + and len(reps.get(m.modality_id, [])) > 0 + ] + if not available: + return g + mod_id = self.rng.choice(available) + new_leaf = (mod_id, self.rng.randrange(len(reps[mod_id]))) + new_idx = len(g.leaves) + g.leaves.append(new_leaf) + if isinstance(g.tree, int): + g.tree = (g.tree, new_idx) + g.fusion_ops = {"": self.rng.choice(self.fusion_operators)} + else: + paths = self._internal_paths(g.tree) + path = self.rng.choice(paths) + sub = _get_subtree(g.tree, path) + g.tree = _replace_subtree(g.tree, path, (sub, new_idx)) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, + g.fusion_ops, + self.rng, + self.fusion_operators, + randomized_prefixes=[path], + ) + return g + + def mutate_remove_leaf(self, g: DagGenome) -> DagGenome: + if len(g.leaves) <= self.min_modalities: + return g + g = copy.deepcopy(g) + drop = self.rng.randrange(len(g.leaves)) + new_tree = _remove_leaf_from_tree(g.tree, drop) + if new_tree is None: + return g + keep = [i for i in range(len(g.leaves)) if i != drop] + index_map = {old: new for new, old in enumerate(keep)} + g.leaves = [g.leaves[i] for i in keep] + g.tree = _reindex_tree(new_tree, index_map) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, g.fusion_ops, self.rng, self.fusion_operators ) - return out - - -class MutateSubtreeResample(MutationOperator): - """ - Pick a random internal subtree, keep the same leaves, but resample: - - subtree topology - - fusion operators inside subtree - """ - - name = "mutate_subtree_resample" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - internal_paths = _collect_internal_paths(out.tree) - if not internal_paths: - return out - - path = rng.choice(internal_paths) - subtree = _get_subtree(out.tree, path) - leaf_indices = _collect_leaf_indices(subtree) - if len(leaf_indices) < 2: - return out - - new_subtree = _sample_random_binary_tree_from_leaves(leaf_indices, rng) - out.tree = _replace_subtree(out.tree, path, new_subtree) - out.fusion_ops = _rebuild_fusion_ops( - out.tree, - out.fusion_ops, - rng, - context["fusion_operators"], + return g + + def mutate_replace_subtree(self, g: DagGenome) -> DagGenome: + g = copy.deepcopy(g) + paths = self._internal_paths(g.tree) + if not paths: + return g + path = self.rng.choice(paths) + sub = _get_subtree(g.tree, path) + if isinstance(sub, int): + return g + + if self.rng.random() < 0.5: + child = sub[0] if self.rng.random() < 0.5 else sub[1] + candidate = _replace_subtree(g.tree, path, child) + kept = sorted(set(_collect_leaf_indices(candidate))) + if len(kept) >= self.min_modalities: + index_map = {old: new for new, old in enumerate(kept)} + g.leaves = [g.leaves[i] for i in kept] + g.tree = _reindex_tree(candidate, index_map) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, {}, self.rng, self.fusion_operators + ) + return g + + leaf_idxs = _collect_leaf_indices(sub) + new_sub = self._random_binary_tree(len(leaf_idxs)) + local_map = {i: leaf_idxs[i] for i in range(len(leaf_idxs))} + new_sub = _reindex_tree(new_sub, local_map) + g.tree = _replace_subtree(g.tree, path, new_sub) + g.fusion_ops = _rebuild_fusion_ops( + g.tree, + g.fusion_ops, + self.rng, + self.fusion_operators, randomized_prefixes=[path], ) - return out - - -class MutateAddModality(MutationOperator): - """ - Optional global mutation: - Add one new modality (if below max_modalities), choose repr_idx, fuse with current root. - """ - - name = "mutate_add_modality" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - if len(out.leaves) >= context["max_modalities"]: - return out - - existing = {m for m, _ in out.leaves} - candidates = [m for m in context["all_modalities"] if m not in existing] - if not candidates: - return out - - m = rng.choice(candidates) - task_name = context["task_name"] - reps = context["k_best_representations"][task_name][m] - if len(reps) == 0: - return out - - repr_idx = rng.randrange(len(reps)) - new_leaf_idx = len(out.leaves) - out.leaves.append((m, repr_idx)) - - # Wrap old tree with new root fusion - old_tree = out.tree - old_fusion_ops = copy.deepcopy(out.fusion_ops) - out.tree = (old_tree, new_leaf_idx) - out.fusion_ops = {"": rng.choice(context["fusion_operators"])} - for path, op in old_fusion_ops.items(): - out.fusion_ops["L" + path] = op - return out - - -class MutateDropModality(MutationOperator): - """ - Optional global mutation: - Remove one modality leaf (if above min_modalities) and collapse tree. - """ - - name = "mutate_drop_modality" - - def __call__( - self, individual: Individual, rng: random.Random, context: Dict[str, Any] - ) -> Individual: - out = copy.deepcopy(individual) - if len(out.leaves) <= context["min_modalities"]: - return out - - leaf_idx = rng.randrange(len(out.leaves)) - new_tree = _remove_leaf_from_tree(out.tree, leaf_idx) - if new_tree is None: - return out - - remaining_leaves = [leaf for i, leaf in enumerate(out.leaves) if i != leaf_idx] - index_map = {} - next_idx = 0 - for old_idx in range(len(out.leaves)): - if old_idx == leaf_idx: - continue - index_map[old_idx] = next_idx - next_idx += 1 - - out.leaves = remaining_leaves - out.tree = _reindex_tree(new_tree, index_map) - out.fusion_ops = _rebuild_fusion_ops( - out.tree, - out.fusion_ops, - rng, - context["fusion_operators"], + return g + + def store_results(self, file_name: str = None, overwrite: bool = False) -> str: + if file_name is None: + timestr = time.strftime("%Y%m%d-%H%M%S") + file_name = f"multimodal_optimizer_{timestr}.pkl" + + directory = os.path.dirname(file_name) or "." + os.makedirs(directory, exist_ok=True) + + if os.path.exists(file_name) and not overwrite: + raise FileExistsError( + f"Refusing to overwrite existing results file '{file_name}'. " + "Pass overwrite=True if this is intentional, or choose a " + "different file_name." + ) + + fd, tmp_path = tempfile.mkstemp( + dir=directory, prefix=".tmp_multimodal_results_", suffix=".pkl" ) - return out + try: + with os.fdopen(fd, "wb") as f: + pickle.dump(self.optimization_results, f) + os.replace(tmp_path, file_name) + except Exception: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + return file_name From e83a20986ae5ce07738fb4f819d0aae934ef29ff Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 13:48:04 +0200 Subject: [PATCH 3/8] add test --- .../scuro/test_multimodal_ga_optimizer.py | 1129 +++++++++++++++++ 1 file changed, 1129 insertions(+) create mode 100644 src/main/python/tests/scuro/test_multimodal_ga_optimizer.py diff --git a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py new file mode 100644 index 00000000000..d325361723e --- /dev/null +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -0,0 +1,1129 @@ +# ------------------------------------------------------------- +# +# 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 copy +import os +import pickle +import random +import tempfile +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from deap import tools + +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + DagGenome, + MultimodalDeapOptimizer, + _collect_leaf_indices, + _failure_fitness, + _objective_value, +) +from systemds.scuro.drsearch.operator_registry import Registry, register_fusion_operator +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.fusion import Fusion +from tests.scuro.data_generator import ModalityRandomDataGenerator, TestTask + +MODULE = "systemds.scuro.drsearch.multimodal_ga_optimizer" + + +@register_fusion_operator() +class _AlwaysFailingFusion(Fusion): + """A fusion operator that always raises - used to prove that one bad + genome can no longer crash the rest of a population's evaluation.""" + + def __init__(self, params=None): + super().__init__("AlwaysFailingFusion") + + def execute(self, modalities): + raise RuntimeError("intentional failure for testing") + + +class _FakeModality: + def __init__(self, modality_id): + self.modality_id = modality_id + + +class _FakeUnimodalResults: + """Stand-in for UnimodalOptimizer.operator_performance: hands back a + fixed list of representations per modality without running any real + unimodal search.""" + + def __init__(self, reps_per_modality): + self.reps_per_modality = reps_per_modality + + def get_k_best_results(self, modality, task, performance_metric_name): + reps = self.reps_per_modality.get(modality.modality_id, []) + return list(range(len(reps))), reps + + +def _make_task(name="task0"): + return SimpleNamespace(model=SimpleNamespace(name=name)) + + +def _make_optimizer(n_modalities=3, reps_per_modality=2, task=None, **kwargs): + modalities = [_FakeModality(f"m{i}") for i in range(n_modalities)] + reps = { + f"m{i}": [object() for _ in range(reps_per_modality)] + for i in range(n_modalities) + } + task = task or _make_task() + kwargs.setdefault("debug", False) + optimizer = MultimodalDeapOptimizer( + modalities, _FakeUnimodalResults(reps), [task], **kwargs + ) + return optimizer, modalities, task + + +def _fake_success_body(_dag, _task, _modalities, _objective_specs, value=0.5): + return (value,), { + "train_score": {}, + "val_score": {"accuracy": value}, + "test_score": {}, + "runtime": 0.0, + "task_time": 0.0, + "representation_time": 0.0, + } + + +def _make_real_representation(modality_id, num_instances, num_features): + gen = ModalityRandomDataGenerator() + rep = gen.create1DModality(num_instances, num_features, ModalityType.TIMESERIES) + rep.modality_id = modality_id + return rep + + +def _build_real_optimizer(num_instances=10, fusion_ops=None, **kwargs): + """Builds an optimizer wired to real (but tiny/synthetic) modalities, + a real cheap task/model, and real fusion operators, so it can execute + genuine RepresentationDag.execute() calls - including across process + boundaries, which rules out unittest.mock patches (a spawned worker + re-imports the module fresh and never sees main-process patches).""" + modality_ids = [0, 1] + modalities = [_FakeModality(mid) for mid in modality_ids] + reps = { + mid: [_make_real_representation(mid, num_instances, 4)] for mid in modality_ids + } + task = TestTask("mm_ga_test_task", "mm_ga_test_model", num_instances) + kwargs.setdefault("min_modalities", 2) + kwargs.setdefault("max_modalities", 2) + kwargs.setdefault("population_size", 4) + kwargs.setdefault("generations", 3) + kwargs.setdefault("elite_size", 1) + kwargs.setdefault("debug", False) + optimizer = MultimodalDeapOptimizer( + modalities, _FakeUnimodalResults(reps), [task], **kwargs + ) + optimizer.fusion_operators = fusion_ops or [Concatenation] + return optimizer, task + + +class TestConstructorValidation(unittest.TestCase): + def test_rejects_too_few_modalities(self): + modalities = [_FakeModality("m0")] + task = _make_task() + with self.assertRaises(ValueError): + MultimodalDeapOptimizer( + modalities, + _FakeUnimodalResults({"m0": [object()]}), + [task], + min_modalities=2, + debug=False, + ) + + def test_rejects_no_registered_fusion_operators(self): + modalities = [_FakeModality("m0"), _FakeModality("m1")] + reps = {"m0": [object()], "m1": [object()]} + task = _make_task() + with patch.object(Registry, "_fusion_operators", []): + with self.assertRaises(ValueError): + MultimodalDeapOptimizer( + modalities, _FakeUnimodalResults(reps), [task], debug=False + ) + + def test_elite_size_clamped_below_population_size(self): + optimizer, _, _ = _make_optimizer(population_size=3, elite_size=10) + self.assertLessEqual(optimizer.elite_size, 2) + + def test_min_modalities_clamped_to_available_when_too_high(self): + # 2 modalities available but caller asks for min_modalities=5: + # construction itself should not explode, and genome sampling + # must clamp instead of calling randint(5, 2). + optimizer, _, task = _make_optimizer( + n_modalities=2, reps_per_modality=2, min_modalities=2, max_modalities=2 + ) + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + self.assertEqual(len(genome.leaves), 2) + + def test_batch_size_defaults_to_max_workers(self): + optimizer, _, _ = _make_optimizer(max_workers=4) + self.assertEqual(optimizer.batch_size, 4) + + def test_batch_size_respects_explicit_value(self): + optimizer, _, _ = _make_optimizer(max_workers=8, batch_size=2) + self.assertEqual(optimizer.batch_size, 2) + self.assertEqual(optimizer.max_workers, 8) + + +class TestGenomeGeneration(unittest.TestCase): + def test_random_genome_respects_min_max_modalities(self): + optimizer, _, task = _make_optimizer( + n_modalities=4, min_modalities=2, max_modalities=3 + ) + optimizer.fusion_operators = [Concatenation, Average] + for _ in range(50): + genome = optimizer._random_genome(task.model.name) + self.assertGreaterEqual(len(genome.leaves), 2) + self.assertLessEqual(len(genome.leaves), 3) + self.assertEqual(len(genome.fusion_ops), len(genome.leaves) - 1) + self.assertEqual( + set(genome.fusion_ops.keys()), + set(optimizer._internal_paths(genome.tree)), + ) + + def test_random_genome_skips_modalities_without_representations(self): + modalities = [_FakeModality("m0"), _FakeModality("m1"), _FakeModality("m2")] + reps = {"m0": [object(), object()], "m1": [], "m2": [object()]} + task = _make_task() + optimizer = MultimodalDeapOptimizer( + modalities, + _FakeUnimodalResults(reps), + [task], + debug=False, + min_modalities=2, + max_modalities=3, + ) + optimizer.fusion_operators = [Concatenation] + for _ in range(50): + genome = optimizer._random_genome(task.model.name) + used = {mod_id for mod_id, _ in genome.leaves} + self.assertNotIn("m1", used) + + def test_random_genome_raises_when_not_enough_modalities_have_reps(self): + modalities = [_FakeModality("m0"), _FakeModality("m1")] + reps = {"m0": [object()], "m1": []} + task = _make_task() + optimizer = MultimodalDeapOptimizer( + modalities, + _FakeUnimodalResults(reps), + [task], + debug=False, + min_modalities=2, + max_modalities=2, + ) + optimizer.fusion_operators = [Concatenation] + with self.assertRaises(ValueError): + optimizer._random_genome(task.model.name) + + +class TestMutations(unittest.TestCase): + def test_add_leaf_noop_beyond_max_modalities(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, min_modalities=2, max_modalities=2 + ) + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + mutated = optimizer.mutate_add_leaf(genome, task.model.name) + self.assertEqual(mutated.leaves, genome.leaves) + + def test_add_leaf_skips_modalities_without_representations(self): + modalities = [_FakeModality("m0"), _FakeModality("m1"), _FakeModality("m2")] + reps = {"m0": [object()], "m1": [], "m2": [object()]} + task = _make_task() + optimizer = MultimodalDeapOptimizer( + modalities, + _FakeUnimodalResults(reps), + [task], + debug=False, + min_modalities=2, + max_modalities=3, + ) + optimizer.fusion_operators = [Concatenation] + genome = DagGenome(leaves=[("m0", 0)], tree=0, fusion_ops={}) + for _ in range(20): + mutated = optimizer.mutate_add_leaf(genome, task.model.name) + used = {mod_id for mod_id, _ in mutated.leaves} + self.assertNotIn("m1", used) + + def test_remove_leaf_noop_at_min_modalities(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, min_modalities=2, max_modalities=2 + ) + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + mutated = optimizer.mutate_remove_leaf(genome) + self.assertEqual(mutated.leaves, genome.leaves) + + def test_remove_leaf_reduces_and_stays_consistent(self): + optimizer, _, task = _make_optimizer( + n_modalities=4, reps_per_modality=2, min_modalities=2, max_modalities=4 + ) + optimizer.fusion_operators = [Concatenation, Average] + genome = None + for _ in range(50): + candidate = optimizer._random_genome(task.model.name) + if len(candidate.leaves) > 2: + genome = candidate + break + self.assertIsNotNone(genome) + mutated = optimizer.mutate_remove_leaf(genome) + self.assertEqual(len(mutated.leaves), len(genome.leaves) - 1) + self.assertEqual( + set(mutated.fusion_ops.keys()), + set(optimizer._internal_paths(mutated.tree)), + ) + + def test_replace_subtree_never_drops_below_min_modalities(self): + """Regression test: the collapse branch used to replace a two-leaf + root by one of its children, leaving a bare leaf index as the tree + while genome.leaves kept both entries -- so a min_modalities=2 search + evaluated unimodal pipelines.""" + optimizer, _, task = _make_optimizer( + n_modalities=3, reps_per_modality=2, min_modalities=2, max_modalities=2 + ) + optimizer.fusion_operators = [Concatenation, Average] + for _ in range(200): + genome = optimizer._random_genome(task.model.name) + mutated = optimizer.mutate_replace_subtree(genome) + self.assertGreaterEqual(len(mutated.leaves), optimizer.min_modalities) + self.assertNotIsInstance(mutated.tree, int) + self.assertEqual( + sorted(set(_collect_leaf_indices(mutated.tree))), + list(range(len(mutated.leaves))), + ) + + def test_replace_subtree_collapse_prunes_and_reindexes_leaves(self): + """When the collapse is allowed (enough leaves survive), the dropped + leaves must leave genome.leaves too, and the tree must be reindexed + onto the surviving ones.""" + optimizer, _, task = _make_optimizer( + n_modalities=4, reps_per_modality=2, min_modalities=2, max_modalities=4 + ) + optimizer.fusion_operators = [Concatenation, Average] + saw_collapse = False + for _ in range(400): + genome = optimizer._random_genome(task.model.name) + if len(genome.leaves) < 3: + continue + mutated = optimizer.mutate_replace_subtree(genome) + leaf_idxs = _collect_leaf_indices(mutated.tree) + self.assertEqual(sorted(set(leaf_idxs)), list(range(len(mutated.leaves)))) + self.assertEqual( + set(mutated.fusion_ops.keys()), + set(optimizer._internal_paths(mutated.tree)), + ) + if len(mutated.leaves) < len(genome.leaves): + saw_collapse = True + # every surviving leaf still names a leaf of the parent genome + for leaf in mutated.leaves: + self.assertIn(leaf, genome.leaves) + self.assertTrue(saw_collapse, "collapse branch never taken") + + +class TestRepeatedModalities(unittest.TestCase): + """allow_repeated_modalities lets one modality contribute several leaves, + each a different representation of it (intra-modal fusion).""" + + def test_off_by_default_one_leaf_per_modality(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, reps_per_modality=4, min_modalities=2, max_modalities=3 + ) + for _ in range(100): + g = optimizer._random_genome(task.model.name) + mods = [mid for mid, _ in g.leaves] + self.assertEqual(len(mods), len(set(mods))) + + def test_max_modalities_clamped_to_modality_count_when_off(self): + optimizer, _, _ = _make_optimizer( + n_modalities=3, reps_per_modality=5, max_modalities=12 + ) + self.assertEqual(optimizer.max_modalities, 3) + + def test_max_modalities_can_exceed_modality_count_when_on(self): + optimizer, _, _ = _make_optimizer( + n_modalities=3, + reps_per_modality=5, + max_modalities=12, + allow_repeated_modalities=True, + ) + self.assertEqual(optimizer.max_modalities, 12) + + def test_random_genome_may_repeat_a_modality_and_leaves_stay_distinct(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=4, + min_modalities=2, + max_modalities=6, + allow_repeated_modalities=True, + ) + saw_repeat = False + for _ in range(200): + g = optimizer._random_genome(task.model.name) + self.assertEqual(len(set(g.leaves)), len(g.leaves)) + self.assertLessEqual(len(g.leaves), 8) # 2 modalities x 4 reps + if len({mid for mid, _ in g.leaves}) < len(g.leaves): + saw_repeat = True + self.assertTrue(saw_repeat, "no genome ever repeated a modality") + + def test_leaf_capacity_bounds_genome_size(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=3, + min_modalities=2, + max_modalities=99, + allow_repeated_modalities=True, + ) + self.assertEqual(optimizer._leaf_capacity(task.model.name), 6) + for _ in range(100): + g = optimizer._random_genome(task.model.name) + self.assertLessEqual(len(g.leaves), 6) + + def test_add_leaf_can_repeat_a_modality_without_duplicating_a_leaf(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=3, + min_modalities=2, + max_modalities=6, + allow_repeated_modalities=True, + ) + name = task.model.name + g = optimizer._random_genome(name) + for _ in range(20): + g = optimizer.mutate_add_leaf(g, name) + self.assertEqual(len(set(g.leaves)), len(g.leaves)) + self.assertEqual(len(g.leaves), 6) # saturates at capacity, no duplicates + + def test_swap_leaf_repr_never_creates_a_duplicate_leaf(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=2, + min_modalities=2, + max_modalities=4, + allow_repeated_modalities=True, + ) + name = task.model.name + for _ in range(200): + g = optimizer._random_genome(name) + mutated = optimizer._mutate_swap_leaf_repr(g, name) + self.assertEqual(len(set(mutated.leaves)), len(mutated.leaves)) + + def test_min_modalities_one_admits_unimodal_genomes(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, reps_per_modality=2, min_modalities=1, max_modalities=3 + ) + self.assertEqual(optimizer.min_modalities, 1) + sizes = { + len(optimizer._random_genome(task.model.name).leaves) for _ in range(200) + } + self.assertIn(1, sizes) + + def test_rejects_max_below_min(self): + with self.assertRaises(ValueError): + _make_optimizer(n_modalities=4, min_modalities=3, max_modalities=2) + + +class TestHallOfFame(unittest.TestCase): + def test_keeps_best_single_objective_results_in_order(self): + optimizer, _, task = _make_optimizer(hall_of_fame_size=2) + optimizer.fusion_operators = [Concatenation] + name = task.model.name + + for value in (0.3, 0.9, 0.5, 0.7): + genome = optimizer._random_genome(name) + with patch( + f"{MODULE}._evaluate_genome_body", + side_effect=lambda *a, v=value, **kw: _fake_success_body(*a, value=v), + ): + optimizer._evaluate_genome(genome, task) + + hof = optimizer.get_hall_of_fame(name) + self.assertEqual([r.val_score["accuracy"] for r in hof], [0.9, 0.7]) + # the full result list is untouched by the hall of fame + self.assertEqual(len(optimizer.optimization_results[name]), 4) + + def test_direction_aware_for_a_minimised_objective(self): + optimizer, _, task = _make_optimizer( + objectives=[("accuracy", "min")], hall_of_fame_size=1 + ) + optimizer.fusion_operators = [Concatenation] + name = task.model.name + for value in (0.8, 0.2, 0.6): + genome = optimizer._random_genome(name) + with patch( + f"{MODULE}._evaluate_genome_body", + side_effect=lambda *a, v=value, **kw: _fake_success_body(*a, value=v), + ): + optimizer._evaluate_genome(genome, task) + hof = optimizer.get_hall_of_fame(name) + self.assertEqual([r.val_score["accuracy"] for r in hof], [0.2]) + + def test_multi_objective_keeps_non_dominated_front_only(self): + optimizer, _, task = _make_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + optimizer.fusion_operators = [Concatenation] + name = task.model.name + + # (accuracy, runtime): B dominates C; A and B are mutually non-dominated. + points = [(0.9, 10.0), (0.5, 1.0), (0.4, 2.0)] + for accuracy, runtime in points: + genome = optimizer._random_genome(name) + + def body(*_a, acc=accuracy, rt=runtime, **_kw): + return (acc, rt), { + "train_score": {}, + "val_score": {"accuracy": acc}, + "test_score": {}, + "runtime": rt, + "task_time": 0.0, + "representation_time": rt, + } + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=body): + optimizer._evaluate_genome(genome, task) + + front = { + (r.val_score["accuracy"], r.runtime) + for r in optimizer.get_hall_of_fame(name) + } + self.assertEqual(front, {(0.9, 10.0), (0.5, 1.0)}) + + def test_failed_evaluations_never_enter_the_hall_of_fame(self): + optimizer, _, task = _make_optimizer() + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + with patch(f"{MODULE}._evaluate_genome_body", side_effect=RuntimeError("boom")): + optimizer._evaluate_genome(genome, task) + self.assertEqual(optimizer.get_hall_of_fame(task.model.name), []) + + +class TestCrossover(unittest.TestCase): + def test_same_leaves_produces_structurally_valid_children(self): + optimizer, _, task = _make_optimizer(n_modalities=3, reps_per_modality=2) + optimizer.fusion_operators = [Concatenation, Average] + g1 = optimizer._random_genome(task.model.name) + g2 = copy.deepcopy(g1) + g2.tree = optimizer._random_binary_tree(len(g2.leaves)) + g2.fusion_ops = {} + optimizer._assign_fusion_ops(g2.tree, g2.fusion_ops, "") + + c1, c2 = optimizer._crossover_genomes(g1, g2) + for child in (c1, c2): + self.assertEqual(sorted(child.leaves), sorted(g1.leaves)) + self.assertEqual( + set(child.fusion_ops.keys()), set(optimizer._internal_paths(child.tree)) + ) + + def test_mismatched_leaves_still_recombines(self): + """Regression test: the original crossover returned both parents + completely unchanged whenever they didn't select the exact same + modality subset/order/index - which, since leaves are sampled + independently per genome, made crossover a near total no-op even + though it fired with 70% probability every generation.""" + optimizer, _, task = _make_optimizer( + n_modalities=4, reps_per_modality=2, min_modalities=2, max_modalities=4 + ) + optimizer.fusion_operators = [Concatenation, Average, Sum] + optimizer.rng = random.Random(0) + + g1 = optimizer._random_genome(task.model.name) + g2 = optimizer._random_genome(task.model.name) + for _ in range(50): + if g1.leaves != g2.leaves: + break + g2 = optimizer._random_genome(task.model.name) + self.assertNotEqual(g1.leaves, g2.leaves, "test setup needs mismatched parents") + + changed = False + for _ in range(100): + c1, c2 = optimizer._crossover_genomes(g1, g2) + if c1.fusion_ops != g1.fusion_ops or c2.fusion_ops != g2.fusion_ops: + changed = True + break + self.assertTrue( + changed, "crossover with mismatched leaf sets never recombined anything" + ) + # leaves/tree topology are untouched by the op-only fallback + self.assertEqual(c1.leaves, g1.leaves) + self.assertEqual(c2.leaves, g2.leaves) + + +class TestGenomeSignature(unittest.TestCase): + def test_signature_ignores_fusion_ops_dict_insertion_order(self): + optimizer, _, task = _make_optimizer(n_modalities=3, reps_per_modality=2) + optimizer.fusion_operators = [Concatenation, Average] + genome = optimizer._random_genome(task.model.name) + reordered = DagGenome( + leaves=list(genome.leaves), + tree=genome.tree, + fusion_ops=dict(reversed(list(genome.fusion_ops.items()))), + ) + self.assertEqual( + optimizer._genome_signature(genome), optimizer._genome_signature(reordered) + ) + + def test_signature_differs_for_different_fusion_op(self): + optimizer, _, task = _make_optimizer(n_modalities=2, reps_per_modality=1) + optimizer.fusion_operators = [Concatenation, Average] + genome = DagGenome( + leaves=[("m0", 0), ("m1", 0)], tree=(0, 1), fusion_ops={"": Concatenation} + ) + other = DagGenome( + leaves=[("m0", 0), ("m1", 0)], tree=(0, 1), fusion_ops={"": Average} + ) + self.assertNotEqual( + optimizer._genome_signature(genome), optimizer._genome_signature(other) + ) + + +class TestEvaluateGenome(unittest.TestCase): + def test_records_result_on_success(self): + optimizer, _, task = _make_optimizer() + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=_fake_success_body): + fitness = optimizer._evaluate_genome(genome, task) + + self.assertEqual(fitness, (0.5,)) + self.assertEqual(len(optimizer.optimization_results[task.model.name]), 1) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + + def test_survives_exception_without_crashing(self): + """Regression test for the crash bug: a failure used to raise + NameError (missing `traceback` import in the except-block) and then + TypeError (unpacking the commented-out None return), instead of + just scoring the genome -inf and moving on.""" + optimizer, _, task = _make_optimizer() + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=RuntimeError("boom")): + fitness = optimizer._evaluate_genome(genome, task) + + self.assertEqual(fitness, (float("-inf"),)) + self.assertEqual(len(optimizer.optimization_results[task.model.name]), 0) + self.assertEqual(optimizer.evaluation_errors[task.model.name], 1) + + def test_uses_cache_for_repeated_signature(self): + optimizer, _, task = _make_optimizer() + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + + call_count = {"n": 0} + + def counting_body(*args, **kwargs): + call_count["n"] += 1 + return _fake_success_body(*args, **kwargs) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=counting_body): + optimizer._evaluate_genome(genome, task) + optimizer._evaluate_genome(copy.deepcopy(genome), task) + + self.assertEqual(call_count["n"], 1) + + +class TestNextGenerationAndPopulation(unittest.TestCase): + def test_next_generation_carries_over_elite_individuals(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, reps_per_modality=3, population_size=5, elite_size=2 + ) + optimizer.fusion_operators = [Concatenation, Average] + population = optimizer._build_initial_population(task.model.name) + for i, ind in enumerate(population): + ind.fitness.values = (float(i),) + + ranked = sorted(population, key=lambda ind: ind.fitness.values[0], reverse=True) + top_signatures = { + optimizer._genome_signature(ind[0]) + for ind in ranked[: optimizer.elite_size] + } + + next_population = optimizer._next_generation(population, task.model.name, task) + next_signatures = { + optimizer._genome_signature(ind[0]) for ind in next_population + } + self.assertTrue(top_signatures.issubset(next_signatures)) + + elites_in_next = [ + ind + for ind in next_population + if optimizer._genome_signature(ind[0]) in top_signatures + ] + for ind in elites_in_next: + self.assertTrue(ind.fitness.valid) + + def test_novelty_breeding_rejects_genomes_from_earlier_generations(self): + # The dedup set must include the whole fitness cache, not just the + # current elite. Without that, offspring identical to something + # scored in an earlier generation are accepted, served from the + # cache, and occupy a population slot that explores nothing. + optimizer, _, task = _make_optimizer( + n_modalities=3, reps_per_modality=4, population_size=6, elite_size=1 + ) + optimizer.fusion_operators = [Concatenation, Average] + name = task.model.name + population = optimizer._build_initial_population(name) + for i, ind in enumerate(population): + ind.fitness.values = (float(i),) + + # Pretend a previous generation already scored these genomes. + cache = optimizer._fitness_cache.setdefault(name, {}) + for ind in population: + cache[optimizer._genome_signature(ind[0])] = ind.fitness.values + stale = set(cache) + + nxt = optimizer._next_generation(population, name, task) + elite = { + optimizer._genome_signature(ind[0]) + for ind in sorted( + population, key=lambda i: i.fitness.values[0], reverse=True + )[: optimizer.elite_size] + } + # Elites are exempt; every other slot must be a genome never scored. + non_elite = [ + optimizer._genome_signature(ind[0]) + for ind in nxt + if optimizer._genome_signature(ind[0]) not in elite + ] + self.assertTrue(non_elite) + self.assertEqual([g for g in non_elite if g in stale], []) + + def test_novelty_breeding_can_be_disabled(self): + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=1, + population_size=4, + min_modalities=2, + max_modalities=2, + novelty_breeding=False, + ) + optimizer.fusion_operators = [Concatenation] + name = task.model.name + self.assertEqual(optimizer._novelty_archive(name), set()) + + def test_novelty_breeding_still_terminates_when_space_is_exhausted(self): + # Archive covers the only genome that exists: breeding must fall back + # to duplicates rather than spinning on the retry budget forever. + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=1, + population_size=5, + min_modalities=2, + max_modalities=2, + ) + optimizer.fusion_operators = [Concatenation] + name = task.model.name + population = optimizer._build_initial_population(name) + for ind in population: + ind.fitness.values = (0.5,) + cache = optimizer._fitness_cache.setdefault(name, {}) + for ind in population: + cache[optimizer._genome_signature(ind[0])] = ind.fitness.values + + nxt = optimizer._next_generation(population, name, task) + self.assertEqual(len(nxt), 5) + + def test_next_generation_terminates_with_tiny_search_space(self): + # 2 modalities x 1 representation x 1 fusion op => exactly one + # distinct genome is possible; requesting a bigger population must + # not hang. + optimizer, _, task = _make_optimizer( + n_modalities=2, + reps_per_modality=1, + population_size=8, + min_modalities=2, + max_modalities=2, + ) + optimizer.fusion_operators = [Concatenation] + population = optimizer._build_initial_population(task.model.name) + self.assertEqual(len(population), 8) + for ind in population: + ind.fitness.values = (0.5,) + + next_population = optimizer._next_generation(population, task.model.name, task) + self.assertEqual(len(next_population), 8) + + +class TestOptimizeLoop(unittest.TestCase): + def test_end_to_end_with_stubbed_fitness(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, + reps_per_modality=3, + population_size=6, + generations=8, + elite_size=1, + early_stopping_patience=3, + ) + optimizer.fusion_operators = [Concatenation, Average] + + def fake_body(dag, _task, _modalities, _metric): + h = hash(str(dag.nodes)) % 1000 / 1000.0 + return _fake_success_body(dag, _task, _modalities, _metric, value=h) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=fake_body): + results = optimizer.optimize() + + self.assertIn(task.model.name, results) + self.assertGreater(len(results[task.model.name]), 0) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + + def test_survives_partial_evaluation_failures(self): + optimizer, _, task = _make_optimizer( + population_size=6, generations=4, elite_size=1, early_stopping_patience=None + ) + optimizer.fusion_operators = [Concatenation, Average] + + call_counter = {"n": 0} + + def flaky_body(dag, _task, _modalities, _metric): + call_counter["n"] += 1 + if call_counter["n"] % 3 == 0: + raise RuntimeError("simulated fusion failure") + return _fake_success_body(dag, _task, _modalities, _metric, value=0.6) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=flaky_body): + results = optimizer.optimize() + + self.assertGreater(optimizer.evaluation_errors.get(task.model.name, 0), 0) + self.assertGreater(len(results[task.model.name]), 0) + + def test_early_stopping_triggers_before_generation_budget(self): + optimizer, _, task = _make_optimizer( + population_size=6, + generations=50, + elite_size=1, + early_stopping_patience=2, + early_stopping_min_delta=1e-6, + ) + optimizer.fusion_operators = [Concatenation] + + original_next_gen = optimizer._next_generation + gens_run = {"n": 0} + + def counting_next_gen(pop, name, task_): + gens_run["n"] += 1 + return original_next_gen(pop, name, task_) + + optimizer._next_generation = counting_next_gen + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=_fake_success_body): + optimizer.optimize() + + self.assertLess(gens_run["n"], 10) + + +class TestStoreResults(unittest.TestCase): + def test_refuses_overwrite_by_default(self): + optimizer, _, task = _make_optimizer() + optimizer.optimization_results[task.model.name] = ["dummy"] + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "results.pkl") + optimizer.store_results(path) + with open(path, "rb") as f: + original = f.read() + + optimizer.optimization_results[task.model.name] = ["different"] + with self.assertRaises(FileExistsError): + optimizer.store_results(path) + + with open(path, "rb") as f: + self.assertEqual(f.read(), original) + + def test_overwrite_true_replaces_file(self): + optimizer, _, task = _make_optimizer() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "results.pkl") + optimizer.optimization_results[task.model.name] = ["v1"] + optimizer.store_results(path) + + optimizer.optimization_results[task.model.name] = ["v2"] + optimizer.store_results(path, overwrite=True) + + with open(path, "rb") as f: + loaded = pickle.load(f) + self.assertEqual(loaded[task.model.name], ["v2"]) + + def test_write_is_atomic_no_partial_file_on_failure(self): + optimizer, _, _ = _make_optimizer() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "results.pkl") + with patch("pickle.dump", side_effect=RuntimeError("disk full")): + with self.assertRaises(RuntimeError): + optimizer.store_results(path) + self.assertFalse(os.path.exists(path)) + self.assertEqual(os.listdir(d), []) + + +class TestRealFusionIntegration(unittest.TestCase): + """Exercises the actual RepresentationDag.execute() path end-to-end, + including across process boundaries for the parallel case (a spawned + worker re-imports this module fresh, so unittest.mock patches from the + parent process cannot reach it - these tests need real, picklable + fusion ops/tasks/modalities).""" + + def test_optimize_end_to_end_real_serial(self): + optimizer, task = _build_real_optimizer(max_workers=1) + results = optimizer.optimize() + self.assertGreater(len(results[task.model.name]), 0) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + + def test_optimize_end_to_end_real_parallel(self): + optimizer, task = _build_real_optimizer(max_workers=2, batch_size=2) + results = optimizer.optimize() + self.assertGreater(len(results[task.model.name]), 0) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + + def test_parallel_evaluation_dedupes_identical_genomes_in_same_batch(self): + optimizer, task = _build_real_optimizer(max_workers=2, batch_size=2) + genome = optimizer._random_genome(task.model.name) + ind1 = optimizer._make_individual(copy.deepcopy(genome)) + ind2 = optimizer._make_individual(copy.deepcopy(genome)) + + optimizer._evaluate_individuals_parallel([ind1, ind2], task) + + self.assertTrue(ind1.fitness.valid) + self.assertTrue(ind2.fitness.valid) + self.assertEqual(ind1.fitness.values, ind2.fitness.values) + self.assertEqual(len(optimizer.optimization_results[task.model.name]), 1) + + def test_parallel_evaluation_survives_a_failing_genome(self): + """A genome whose fusion operator always raises must not take down + the rest of the (parallel) batch.""" + optimizer, task = _build_real_optimizer( + max_workers=2, + batch_size=2, + fusion_ops=[Concatenation, _AlwaysFailingFusion], + ) + good_genome = DagGenome( + leaves=[(0, 0), (1, 0)], tree=(0, 1), fusion_ops={"": Concatenation} + ) + bad_genome = DagGenome( + leaves=[(0, 0), (1, 0)], + tree=(0, 1), + fusion_ops={"": _AlwaysFailingFusion}, + ) + ind_good = optimizer._make_individual(good_genome) + ind_bad = optimizer._make_individual(bad_genome) + + optimizer._evaluate_individuals_parallel([ind_good, ind_bad], task) + + self.assertTrue(ind_good.fitness.valid) + self.assertTrue(ind_bad.fitness.valid) + self.assertEqual(ind_bad.fitness.values[0], float("-inf")) + self.assertGreater(ind_good.fitness.values[0], float("-inf")) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 1) + self.assertEqual(len(optimizer.optimization_results[task.model.name]), 1) + + +class TestMultiObjective(unittest.TestCase): + def test_objective_value_reads_timing_vs_val_score(self): + val_score = {"accuracy": 0.8, "f1": 0.7} + timing = {"runtime": 1.5, "task_time": 1.0, "representation_time": 0.5} + self.assertEqual(_objective_value("accuracy", val_score, timing), 0.8) + self.assertEqual(_objective_value("f1", val_score, timing), 0.7) + self.assertEqual(_objective_value("runtime", val_score, timing), 1.5) + self.assertEqual(_objective_value("task_time", val_score, timing), 1.0) + + def test_failure_fitness_is_direction_aware(self): + """A failed evaluation must always be the worst possible candidate, + regardless of whether an objective is maximized or minimized - + using -inf for a 'min' objective (e.g. runtime) would make a failure + look infinitely fast and win every tournament/dominance check.""" + specs = [("accuracy", "max"), ("runtime", "min")] + worst = _failure_fitness(specs) + self.assertEqual(worst, (float("-inf"), float("inf"))) + + def test_constructor_rejects_invalid_direction(self): + with self.assertRaises(ValueError): + _make_optimizer(objectives=[("accuracy", "sideways")]) + + def test_constructor_rejects_empty_objectives(self): + with self.assertRaises(ValueError): + _make_optimizer(objectives=[]) + + def test_is_multi_objective_flag_and_weights(self): + optimizer, _, _ = _make_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + self.assertTrue(optimizer.is_multi_objective) + self.assertEqual( + optimizer.objective_specs, [("accuracy", "max"), ("runtime", "min")] + ) + from deap import creator + + self.assertEqual(creator.FitnessMax.weights, (1.0, -1.0)) + + def test_single_objective_by_default(self): + optimizer, _, _ = _make_optimizer() + self.assertFalse(optimizer.is_multi_objective) + self.assertEqual(optimizer.objective_specs, [("accuracy", "max")]) + + def test_evaluate_genome_returns_tuple_per_objective(self): + optimizer, _, task = _make_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + + def fake_body(_dag, _task, _modalities, objective_specs): + self.assertEqual(objective_specs, optimizer.objective_specs) + return (0.9, 1.2), { + "train_score": {}, + "val_score": {"accuracy": 0.9}, + "test_score": {}, + "runtime": 1.2, + "task_time": 1.0, + "representation_time": 0.2, + } + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=fake_body): + fitness = optimizer._evaluate_genome(genome, task) + + self.assertEqual(fitness, (0.9, 1.2)) + self.assertEqual(len(optimizer.optimization_results[task.model.name]), 1) + + def test_evaluate_genome_failure_uses_direction_aware_sentinel(self): + optimizer, _, task = _make_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + optimizer.fusion_operators = [Concatenation] + genome = optimizer._random_genome(task.model.name) + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=RuntimeError("boom")): + fitness = optimizer._evaluate_genome(genome, task) + + self.assertEqual(fitness, (float("-inf"), float("inf"))) + + def test_next_generation_multi_objective_returns_population_sized_front(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, + reps_per_modality=3, + population_size=6, + objectives=[("accuracy", "max"), ("runtime", "min")], + ) + optimizer.fusion_operators = [Concatenation, Average] + population = optimizer._build_initial_population(task.model.name) + for i, ind in enumerate(population): + ind.fitness.values = ( + float(i) / len(population), + float(len(population) - i), + ) + + def fake_body(dag, _task, _modalities, _objective_specs): + h = hash(str(dag.nodes)) % 1000 / 1000.0 + return (h, 1.0 - h), { + "train_score": {}, + "val_score": {"accuracy": h}, + "test_score": {}, + "runtime": 1.0 - h, + "task_time": 0.0, + "representation_time": 0.0, + } + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=fake_body): + next_population = optimizer._next_generation( + population, task.model.name, task + ) + + self.assertEqual(len(next_population), optimizer.population_size) + for ind in next_population: + self.assertTrue(ind.fitness.valid) + self.assertEqual(len(ind.fitness.values), 2) + + def test_optimize_end_to_end_multi_objective_stubbed(self): + optimizer, _, task = _make_optimizer( + n_modalities=3, + reps_per_modality=3, + population_size=6, + generations=6, + objectives=[("accuracy", "max"), ("runtime", "min")], + early_stopping_patience=3, + ) + optimizer.fusion_operators = [Concatenation, Average] + + def fake_body(dag, _task, _modalities, _objective_specs): + h = hash(str(dag.nodes)) % 1000 / 1000.0 + return (h, 1.0 - h), { + "train_score": {}, + "val_score": {"accuracy": h}, + "test_score": {}, + "runtime": 1.0 - h, + "task_time": 0.0, + "representation_time": 0.0, + } + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=fake_body): + results = optimizer.optimize() + + self.assertIn(task.model.name, results) + self.assertGreater(len(results[task.model.name]), 0) + self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + + def test_optimize_multi_objective_survives_partial_failures(self): + optimizer, _, task = _make_optimizer( + population_size=6, + generations=4, + objectives=[("accuracy", "max"), ("runtime", "min")], + early_stopping_patience=None, + ) + optimizer.fusion_operators = [Concatenation, Average] + + call_counter = {"n": 0} + + def flaky_body(dag, _task, _modalities, _objective_specs): + call_counter["n"] += 1 + if call_counter["n"] % 3 == 0: + raise RuntimeError("simulated fusion failure") + return (0.7, 0.3), { + "train_score": {}, + "val_score": {"accuracy": 0.7}, + "test_score": {}, + "runtime": 0.3, + "task_time": 0.0, + "representation_time": 0.0, + } + + with patch(f"{MODULE}._evaluate_genome_body", side_effect=flaky_body): + results = optimizer.optimize() + + self.assertGreater(optimizer.evaluation_errors.get(task.model.name, 0), 0) + self.assertGreater(len(results[task.model.name]), 0) + + def test_real_fusion_multi_objective_end_to_end(self): + """Exercises the real dag.execute() path (not stubbed) with two + objectives to make sure runtime is actually threaded through from + real evaluation timing, not just from stubbed payloads.""" + optimizer, task = _build_real_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + results = optimizer.optimize() + self.assertGreater(len(results[task.model.name]), 0) + for result in results[task.model.name]: + self.assertGreaterEqual(result.runtime, 0.0) + + +if __name__ == "__main__": + unittest.main() From 1a1d53d8b60b2c9bd976f814915004fddb37087a Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 15:37:39 +0200 Subject: [PATCH 4/8] add multiprocessing to multimodal optimizer --- .github/workflows/python.yml | 3 +- .../scuro/drsearch/multimodal_ga_optimizer.py | 223 ++++++++++++------ .../scuro/test_multimodal_ga_optimizer.py | 103 +++++++- 3 files changed, 247 insertions(+), 82 deletions(-) diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index d55f9adc6c0..4a2f4b8a672 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -179,7 +179,8 @@ jobs: fvcore \ scikit-optimize \ flair \ - optuna + optuna \ + deap kill $KA cd src/main/python python -m unittest discover -s tests/scuro -p 'test_*.py' -v diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 718c342d573..47a26e83276 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -22,19 +22,23 @@ from __future__ import annotations import copy -import multiprocessing as mp import os import pickle import random import tempfile +import threading import time import traceback -from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait from dataclasses import dataclass, field from itertools import chain from typing import Any, Dict, List, Optional, Tuple -from deap import base, creator, tools +from deap import base, tools + +from systemds.scuro.drsearch.modality_shared_memory import ( + add_shared_memory_candidate, + unlink_shm, +) from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.drsearch.representation_dag import ( @@ -42,6 +46,7 @@ RepresentationDag, ) from systemds.scuro.drsearch.task import Task +from systemds.scuro.drsearch.worker_pool import PersistentWorkerPool, create_mp_context from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) @@ -215,22 +220,21 @@ def _evaluate_genome_body( return fitness, payload -def _evaluate_dag_worker( - dag_bytes: bytes, - task_bytes: bytes, - modalities_bytes: bytes, - objective_specs: List[ObjectiveSpec], -) -> Tuple[Tuple[float, ...], Optional[Dict[str, Any]], Optional[str]]: - try: - dag = pickle.loads(dag_bytes) - task = pickle.loads(task_bytes) - modalities = pickle.loads(modalities_bytes) - fitness, payload = _evaluate_genome_body(dag, task, modalities, objective_specs) - if fitness is None: - fitness = _failure_fitness(objective_specs) - return fitness, payload, None - except Exception: - return _failure_fitness(objective_specs), None, traceback.format_exc() +def _dispatch_genome_evaluation(payload, _gpu_id): + dag, task, modalities, objective_specs = payload + fitness, result = _evaluate_genome_body(dag, task, modalities, objective_specs) + if fitness is None: + fitness = _failure_fitness(objective_specs) + return fitness, result + + +_WORKER_DISPATCH = {"genome": _dispatch_genome_evaluation} + + +class _FusionIndividual(list): + def __init__(self, values, fitness_type): + super().__init__(values) + self.fitness = fitness_type() class MultimodalDeapOptimizer: @@ -258,6 +262,7 @@ def __init__( novelty_breeding: bool = True, hall_of_fame_size: int = 5, allow_repeated_modalities: bool = False, + threads_per_worker: Optional[int] = None, ): self.modalities = modalities self.tasks = tasks @@ -336,30 +341,44 @@ def __init__( self.elite_size = max(0, min(elite_size, self.population_size - 1)) self.max_workers = max(1, max_workers) self.batch_size = max(1, batch_size or self.max_workers) + cpu_count = os.cpu_count() or 1 + self.threads_per_worker = max( + 1, + ( + threads_per_worker + if threads_per_worker is not None + else cpu_count // self.max_workers + ), + ) self.early_stopping_patience = early_stopping_patience self.early_stopping_min_delta = early_stopping_min_delta self.novelty_breeding = novelty_breeding self._current_task_name = None + self._optimize_lock = threading.Lock() + self._worker_pool: Optional[PersistentWorkerPool] = None + self._parallel_task_name: Optional[str] = None + self._parallel_modalities: Optional[List[Any]] = None + self._parallel_shm_names: List[str] = [] desired_weights = tuple( 1.0 if direction == "max" else -1.0 for _, direction in self.objective_specs ) self._objective_weights = desired_weights - existing_weights = getattr( - getattr(creator, "FitnessMax", None), "weights", None + self._fitness_type = type( + f"FusionFitness_{id(self)}", (base.Fitness,), {"weights": desired_weights} ) - if existing_weights != desired_weights: - if hasattr(creator, "Individual"): - del creator.Individual - if hasattr(creator, "FitnessMax"): - del creator.FitnessMax - creator.create("FitnessMax", base.Fitness, weights=desired_weights) - creator.create("Individual", list, fitness=creator.FitnessMax) - elif not hasattr(creator, "Individual"): - creator.create("Individual", list, fitness=creator.FitnessMax) def optimize( self, + ) -> Dict[str, List[FusionSearchResult]]: + with self._optimize_lock: + try: + return self._optimize() + finally: + self._shutdown_parallel_runtime() + + def _optimize( + self, ) -> Dict[str, List[FusionSearchResult]]: for task in self.tasks: task_name = task.model.name @@ -370,6 +389,9 @@ def optimize( self._eval_counter = 0 self._search_start = time.perf_counter() + if self.max_workers > 1: + self._start_parallel_runtime(task_name) + population = self._build_initial_population(task_name) best_ever = None no_improve = 0 @@ -428,10 +450,10 @@ def optimize( return self.optimization_results def _make_individual(self, genome: DagGenome): - return creator.Individual([genome]) + return _FusionIndividual([genome], self._fitness_type) def _clone_individual(self, ind): - clone = creator.Individual([copy.deepcopy(ind[0])]) + clone = self._make_individual(copy.deepcopy(ind[0])) if ind.fitness.valid: clone.fitness.values = ind.fitness.values return clone @@ -536,61 +558,112 @@ def _evaluate_population(self, population: List[Any], task: Task) -> None: fitness = self._evaluate_genome(ind[0], task) ind.fitness.values = fitness + def _start_parallel_runtime(self, task_name: str) -> None: + if self._worker_pool is not None and self._parallel_task_name == task_name: + return + self._shutdown_parallel_runtime() + modalities = list( + chain.from_iterable(self.k_best_representations[task_name].values()) + ) + shared_modalities = [] + shm_names = [] + try: + for modality in modalities: + shared_modality = copy.copy(modality) + resident_bytes = 0 + try: + resident_bytes = modality.calculate_memory_usage() + except Exception: + pass + wrapped, shm_name, _, _ = add_shared_memory_candidate( + modality.data, resident_bytes + ) + if wrapped is not None: + shared_modality._data = wrapped + shm_names.append(shm_name) + shared_modalities.append(shared_modality) + worker_pool = PersistentWorkerPool( + self.max_workers, + _WORKER_DISPATCH, + ctx=create_mp_context(), + threads_per_worker=self.threads_per_worker, + ) + except Exception: + for shm_name in shm_names: + unlink_shm(shm_name) + raise + self._worker_pool = worker_pool + self._parallel_task_name = task_name + self._parallel_modalities = shared_modalities + self._parallel_shm_names = shm_names + + def _shutdown_parallel_runtime(self) -> None: + if self._worker_pool is not None: + self._worker_pool.shutdown() + self._worker_pool = None + self._parallel_task_name = None + self._parallel_modalities = None + for shm_name in self._parallel_shm_names: + unlink_shm(shm_name) + self._parallel_shm_names = [] + def _evaluate_individuals_parallel( self, individuals: List[Any], task: Task ) -> None: task_name = task.model.name + manage_runtime = self._worker_pool is None + if manage_runtime: + self._start_parallel_runtime(task_name) cache = self._fitness_cache.setdefault(task_name, {}) - ctx = mp.get_context("spawn") - task_bytes = pickle.dumps(task) - futures: Dict[Any, Tuple[Any, RepresentationDag, Tuple]] = {} pending_followers: Dict[Tuple, List[Any]] = {} + pending_work = [] + jobs: Dict[int, Tuple[Any, RepresentationDag, Tuple]] = {} + + for ind in individuals: + genome = ind[0] + sig = self._genome_signature(genome) + cached = cache.get(sig) + if cached is not None: + ind.fitness.values = cached + continue + if sig in pending_followers: + pending_followers[sig].append(ind) + continue + dag = self._genome_to_dag(genome) + pending_followers[sig] = [] + pending_work.append((ind, dag, sig)) - def _drain(done_futures): - for fut in done_futures: - ind, dag, sig = futures.pop(fut) - fitness, payload, error = fut.result() + try: + while pending_work or jobs: + while ( + pending_work + and self._worker_pool.has_idle_worker + and len(jobs) < self.batch_size + ): + ind, dag, sig = pending_work.pop(0) + job_id = self._worker_pool.submit( + "genome", + (dag, task, self._parallel_modalities, self.objective_specs), + ) + jobs[job_id] = (ind, dag, sig) + + jr = self._worker_pool.wait() + ind, dag, sig = jobs.pop(jr.job_id) + if jr.ok: + fitness, payload = jr.value + error = None + else: + fitness = _failure_fitness(self.objective_specs) + payload = None + error = jr.error ind.fitness.values = fitness self._record_evaluation(task_name, dag, payload, error) cache[sig] = fitness for follower in pending_followers.pop(sig, []): follower.fitness.values = fitness - - with ProcessPoolExecutor( - max_workers=self.max_workers, mp_context=ctx - ) as executor: - for ind in individuals: - genome = ind[0] - sig = self._genome_signature(genome) - cached = cache.get(sig) - if cached is not None: - ind.fitness.values = cached - continue - if sig in pending_followers: - pending_followers[sig].append(ind) - continue - - dag = self._genome_to_dag(genome) - modalities = list( - chain.from_iterable(self.k_best_representations[task_name].values()) - ) - fut = executor.submit( - _evaluate_dag_worker, - pickle.dumps(dag), - task_bytes, - pickle.dumps(modalities), - self.objective_specs, - ) - futures[fut] = (ind, dag, sig) - pending_followers[sig] = [] - - if len(futures) >= self.batch_size: - done, _ = wait(set(futures.keys()), return_when=FIRST_COMPLETED) - _drain(done) - - if futures: - done, _ = wait(set(futures.keys())) - _drain(done) + finally: + if manage_runtime: + self._shutdown_parallel_runtime() def _record_evaluation( self, diff --git a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py index d325361723e..f79157f0247 100644 --- a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -107,6 +107,43 @@ def _fake_success_body(_dag, _task, _modalities, _objective_specs, value=0.5): } +class _SynchronousPool: + instances = [] + + def __init__(self, n_workers, dispatch, ctx=None, threads_per_worker=1): + self.n_workers = n_workers + self.dispatch = dispatch + self.threads_per_worker = threads_per_worker + self.pending = None + self.next_job_id = 0 + self.shutdown_called = False + self.instances.append(self) + + @property + def has_idle_worker(self): + return self.pending is None + + def submit(self, kind, payload, gpu_id=None): + job_id = self.next_job_id + self.next_job_id += 1 + try: + value = self.dispatch[kind](payload, gpu_id) + self.pending = SimpleNamespace(job_id=job_id, ok=True, value=value) + except Exception as exc: + self.pending = SimpleNamespace( + job_id=job_id, ok=False, value=None, error=str(exc) + ) + return job_id + + def wait(self): + result = self.pending + self.pending = None + return result + + def shutdown(self): + self.shutdown_called = True + + def _make_real_representation(modality_id, num_instances, num_features): gen = ModalityRandomDataGenerator() rep = gen.create1DModality(num_instances, num_features, ModalityType.TIMESERIES) @@ -896,6 +933,58 @@ def test_optimize_end_to_end_real_parallel(self): self.assertGreater(len(results[task.model.name]), 0) self.assertEqual(optimizer.evaluation_errors.get(task.model.name, 0), 0) + def test_optimize_reuses_one_pool_across_generations(self): + optimizer, task = _build_real_optimizer( + max_workers=2, batch_size=2, threads_per_worker=3 + ) + _SynchronousPool.instances = [] + with patch(f"{MODULE}.PersistentWorkerPool", _SynchronousPool), patch( + f"{MODULE}.create_mp_context", return_value=None + ): + results = optimizer.optimize() + + self.assertGreater(len(results[task.model.name]), 0) + self.assertEqual(len(_SynchronousPool.instances), 1) + pool = _SynchronousPool.instances[0] + self.assertEqual(pool.threads_per_worker, 3) + self.assertTrue(pool.shutdown_called) + self.assertIsNone(optimizer._worker_pool) + + def test_parallel_runtime_uses_shared_copies_and_unlinks_them(self): + optimizer, task = _build_real_optimizer(max_workers=2, batch_size=2) + task_name = task.model.name + source_modalities = list( + modality + for reps in optimizer.k_best_representations[task_name].values() + for modality in reps + ) + original_data = [modality.data for modality in source_modalities] + wrappers = [object(), object()] + shared_results = [ + (wrappers[0], "shm-0", 1024, 0), + (wrappers[1], "shm-1", 1024, 0), + ] + _SynchronousPool.instances = [] + + with patch( + f"{MODULE}.add_shared_memory_candidate", side_effect=shared_results + ), patch(f"{MODULE}.unlink_shm") as unlink, patch( + f"{MODULE}.PersistentWorkerPool", _SynchronousPool + ), patch( + f"{MODULE}.create_mp_context", return_value=None + ): + optimizer._start_parallel_runtime(task_name) + self.assertEqual( + [modality.data for modality in optimizer._parallel_modalities], + wrappers, + ) + self.assertEqual([m.data for m in source_modalities], original_data) + optimizer._shutdown_parallel_runtime() + + self.assertEqual( + [call.args[0] for call in unlink.call_args_list], ["shm-0", "shm-1"] + ) + def test_parallel_evaluation_dedupes_identical_genomes_in_same_batch(self): optimizer, task = _build_real_optimizer(max_workers=2, batch_size=2) genome = optimizer._random_genome(task.model.name) @@ -965,16 +1054,18 @@ def test_constructor_rejects_empty_objectives(self): _make_optimizer(objectives=[]) def test_is_multi_objective_flag_and_weights(self): - optimizer, _, _ = _make_optimizer( + multi, _, _ = _make_optimizer( objectives=[("accuracy", "max"), ("runtime", "min")] ) - self.assertTrue(optimizer.is_multi_objective) + single, _, _ = _make_optimizer() + self.assertTrue(multi.is_multi_objective) self.assertEqual( - optimizer.objective_specs, [("accuracy", "max"), ("runtime", "min")] + multi.objective_specs, [("accuracy", "max"), ("runtime", "min")] ) - from deap import creator - - self.assertEqual(creator.FitnessMax.weights, (1.0, -1.0)) + multi_ind = multi._make_individual(DagGenome([("m0", 0)], 0, {})) + single_ind = single._make_individual(DagGenome([("m0", 0)], 0, {})) + self.assertEqual(multi_ind.fitness.weights, (1.0, -1.0)) + self.assertEqual(single_ind.fitness.weights, (1.0,)) def test_single_objective_by_default(self): optimizer, _, _ = _make_optimizer() From f702bfaa62e0a96f382f3c952b63d81b7fdcdc95 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 15:56:46 +0200 Subject: [PATCH 5/8] fix dependencies --- .github/workflows/python.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 182bbf9c6d8..d6e66813708 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -180,7 +180,7 @@ jobs: scikit-optimize \ flair \ optuna \ - deap + deap \ openface-test \ imagebind \ "pytorchvideo @ git+https://github.com/facebookresearch/pytorchvideo.git@eb04d1b" From d84309e4c386030d7947d4ad5c8a30c325ea39f2 Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 16:59:25 +0200 Subject: [PATCH 6/8] fix python3.8 imports --- .../python/systemds/scuro/drsearch/multimodal_ga_optimizer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 47a26e83276..68ba8b586e9 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -31,7 +31,7 @@ import traceback from dataclasses import dataclass, field from itertools import chain -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union from deap import base, tools @@ -52,7 +52,7 @@ ) from systemds.scuro.utils.schema_helpers import get_shape -Tree = int | Tuple["Tree", "Tree"] +Tree = Union[int, Tuple["Tree", "Tree"]] def _collect_internal_paths(tree: Tree, path: str = "") -> List[str]: From 29f642735058aeac949c3e730c7447d02f67c12f Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 17:39:48 +0200 Subject: [PATCH 7/8] add missing function --- .../scuro/drsearch/multimodal_ga_optimizer.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py index 68ba8b586e9..44858d96a0e 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -176,6 +176,13 @@ def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: ) +def _fold_scores(performance_measure) -> Dict[str, List[float]]: + return { + metric: [float(value) for value in values] + for metric, values in performance_measure.scores.items() + } + + def _evaluate_genome_body( dag: RepresentationDag, task: Task, @@ -211,9 +218,9 @@ def _evaluate_genome_body( "train_score": scores[0].average_scores, "val_score": val_score, "test_score": scores[2].average_scores, - "train_fold_scores": scores[0].fold_scores(), - "val_fold_scores": scores[1].fold_scores(), - "test_fold_scores": scores[2].fold_scores(), + "train_fold_scores": _fold_scores(scores[0]), + "val_fold_scores": _fold_scores(scores[1]), + "test_fold_scores": _fold_scores(scores[2]), "task_timing": getattr(task, "last_run_timing", {}), **timing, } From 9a7b89edf3a42d6b59f70530e4f44b1ffc6c73bf Mon Sep 17 00:00:00 2001 From: Christina Dionysio Date: Mon, 31 Aug 2026 18:00:00 +0200 Subject: [PATCH 8/8] fix multimodal operator registry --- .../systemds/scuro/drsearch/operator_registry.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/main/python/systemds/scuro/drsearch/operator_registry.py b/src/main/python/systemds/scuro/drsearch/operator_registry.py index 7a80aafa913..d9ae798f867 100644 --- a/src/main/python/systemds/scuro/drsearch/operator_registry.py +++ b/src/main/python/systemds/scuro/drsearch/operator_registry.py @@ -47,9 +47,9 @@ def __new__(cls): def set_fusion_operators(self, fusion_operators): if isinstance(fusion_operators, list): - self._fusion_operators = fusion_operators + type(self)._fusion_operators = fusion_operators else: - self._fusion_operators = [fusion_operators] + type(self)._fusion_operators = [fusion_operators] def set_representations(self, modality_type, representations): if isinstance(representations, list): @@ -89,7 +89,7 @@ def add_context_operator(self, context_operator, modality_type): self._context_operators[m_type].append(context_operator) def add_fusion_operator(self, fusion_operator): - self._fusion_operators.append(fusion_operator) + type(self)._fusion_operators.append(fusion_operator) def add_dimensionality_reduction_operator( self, dimensionality_reduction_operator, modality_type @@ -134,10 +134,10 @@ def get_dimensionality_reduction_operators(self, modality_type): return self._dimensionality_reduction_operators.get(modality_type, []) def get_fusion_operators(self): - return self._fusion_operators + return type(self)._fusion_operators def get_fusion_operator_by_name(self, fusion_name): - for fusion in self._fusion_operators: + for fusion in type(self)._fusion_operators: if fusion.__name__ == fusion_name: return fusion return None