diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index d679aa74a92..d6e66813708 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -180,6 +180,7 @@ jobs: scikit-optimize \ flair \ optuna \ + deap \ openface-test \ imagebind \ "pytorchvideo @ git+https://github.com/facebookresearch/pytorchvideo.git@eb04d1b" 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..44858d96a0e --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -0,0 +1,1103 @@ +# ------------------------------------------------------------- +# +# 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 os +import pickle +import random +import tempfile +import threading +import time +import traceback +from dataclasses import dataclass, field +from itertools import chain +from typing import Any, Dict, List, Optional, Tuple, Union + +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 ( + RepresentationDAGBuilder, + 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, +) +from systemds.scuro.utils.schema_helpers import get_shape + +Tree = Union[int, Tuple["Tree", "Tree"]] + + +def _collect_internal_paths(tree: Tree, path: str = "") -> List[str]: + if isinstance(tree, int): + return [] + left, right = tree + return ( + [path] + + _collect_internal_paths(left, path + "L") + + _collect_internal_paths(right, path + "R") + ) + + +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(tree: Tree, path: str, replacement: Tree) -> Tree: + if not path: + return copy.deepcopy(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(tree: Tree) -> List[int]: + if isinstance(tree, int): + return [tree] + return _collect_leaf_indices(tree[0]) + _collect_leaf_indices(tree[1]) + + +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(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 _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 + + +@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) + + 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] + + +_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, ...]: + return tuple( + float("-inf") if direction == "max" else float("inf") + for _, direction in objective_specs + ) + + +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, + 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": _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, + } + return fitness, payload + + +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: + 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, + 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, + threads_per_worker: Optional[int] = None, + ): + self.modalities = modalities + self.tasks = tasks + self.debug = debug + 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.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() + 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] = {} + + 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, ...]]] = {} + + 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 + self._fitness_type = type( + f"FusionFitness_{id(self)}", (base.Fitness,), {"weights": desired_weights} + ) + + 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 + self._current_task_name = task_name + self.optimization_results.setdefault(task_name, []) + self.evaluation_errors.setdefault(task_name, 0) + + 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 + + for gen in range(self.generations): + self._current_generation = gen + self._evaluate_population(population, task) + + 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 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 + + def _make_individual(self, genome: DagGenome): + return _FusionIndividual([genome], self._fitness_type) + + def _clone_individual(self, ind): + clone = self._make_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 + 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: + 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: + 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 + 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 _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, {}) + 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)) + + 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 + finally: + if manage_runtime: + self._shutdown_parallel_runtime() + + 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 + + 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 + + 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]: + return list(self.hall_of_fame.get(task_name, [])) + + 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 _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 _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: + 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: + 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: + 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 = [ + 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: + 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, ...]: + 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 + + 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], + ) + 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 + + 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: + 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 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 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" + ) + 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/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 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..f79157f0247 --- /dev/null +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -0,0 +1,1220 @@ +# ------------------------------------------------------------- +# +# 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, + } + + +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) + 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_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) + 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): + multi, _, _ = _make_optimizer( + objectives=[("accuracy", "max"), ("runtime", "min")] + ) + single, _, _ = _make_optimizer() + self.assertTrue(multi.is_multi_objective) + self.assertEqual( + multi.objective_specs, [("accuracy", "max"), ("runtime", "min")] + ) + 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() + 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()