diff --git a/src/main/python/systemds/scuro/dataloader/json_loader.py b/src/main/python/systemds/scuro/dataloader/json_loader.py index d2802894ecb..c8361643046 100644 --- a/src/main/python/systemds/scuro/dataloader/json_loader.py +++ b/src/main/python/systemds/scuro/dataloader/json_loader.py @@ -81,7 +81,7 @@ def extract(self, file: str, index: Optional[Union[str, List[str]]] = None): self.data.append(text) self.metadata.append( - self.modality_type.create_metadata(len(text), text) | record + {**self.modality_type.create_metadata(len(text), text), **record} ) def get_stats(self, source_path: str): diff --git a/src/main/python/systemds/scuro/drsearch/algebra.py b/src/main/python/systemds/scuro/drsearch/algebra.py new file mode 100644 index 00000000000..880e7afe7ea --- /dev/null +++ b/src/main/python/systemds/scuro/drsearch/algebra.py @@ -0,0 +1,577 @@ +# ------------------------------------------------------------- +# +# 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 hashlib +import itertools +import json +from collections import Counter +from dataclasses import dataclass +from fractions import Fraction +from functools import lru_cache +from math import comb, factorial, prod +from typing import NamedTuple, Optional + +import numpy as np + +from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.context import Context +from systemds.scuro.representations.dimensionality_reduction import ( + DimensionalityReduction, +) +from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.representation import ( + FitRequirement, + RepresentationStats, +) +from systemds.scuro.representations.zero_pad import ZeroPad + +ALL_LAWS = frozenset( + { + "comm", + "assoc", + "avg_flatten", + "concat_perm", + "agg_pushdown", + "implicit_agg", + "implicit_pad", + } +) + +EXACT_LAWS = frozenset({"agg_pushdown", "implicit_agg", "implicit_pad"}) + + +@dataclass(frozen=True) +class AlgebraConfig: + """Switches for the Beta-algebra rules; the defaults are the legacy search.""" + + laws: frozenset = frozenset() + feature_order_invariant: Optional[bool] = None + cse_law_aware: bool = False + ga_canonical_dedup: bool = False + exhaustive_dedup_permutations: bool = False + exhaustive_canonical: bool = False + typecheck: str = "off" # off | warn | strict + explicit_fusion_prep: bool = False + explicit_padding: str = "off" # off | pad | reject + typed_chaining: bool = False + share_fused_subtrees: bool = False + + @classmethod + def full(cls) -> "AlgebraConfig": + return cls( + laws=ALL_LAWS, + feature_order_invariant=True, + cse_law_aware=True, + ga_canonical_dedup=True, + exhaustive_dedup_permutations=True, + exhaustive_canonical=True, + typecheck="warn", + explicit_fusion_prep=True, + explicit_padding="pad", + share_fused_subtrees=True, + ) + + +class OperatorSignature(NamedTuple): + kind: str # encode | context | aggregate | reduce | fuse | task + input_types: frozenset + output_type: Optional[ModalityType] # None: the input's type + arity: object # 1, or "n" for fusion + input_rank: str # any | seq | windows | vector + output_rank: str + + +ANY_TYPE = frozenset(m for m in ModalityType if isinstance(m.value, int)) + + +def operator_signature(op_cls, parameters=None) -> OperatorSignature: + """`op_cls` None is a task node.""" + if op_cls is None: + return OperatorSignature("task", ANY_TYPE, None, 1, "any", "any") + if issubclass(op_cls, AggregatedRepresentation): + return OperatorSignature("aggregate", ANY_TYPE, None, 1, "seq", "vector") + if issubclass(op_cls, ZeroPad): + return OperatorSignature("reduce", ANY_TYPE, None, 1, "vector", "vector") + if issubclass(op_cls, Fusion): + keep_axis = (parameters or {}).get("preserve_leading_axis") + seq_in = op_cls.preserves_leading_axis if keep_axis is None else keep_axis + ranks = ["seq" if seq else "vector" for seq in (seq_in, keep_axis)] + return OperatorSignature("fuse", ANY_TYPE, _output_type(op_cls), "n", *ranks) + if issubclass(op_cls, Context): + kind, lists, ranks = "context", Registry._context_operators, ("seq", "windows") + elif issubclass(op_cls, DimensionalityReduction): + kind, lists = "reduce", Registry._dimensionality_reduction_operators + ranks = ("vector", "vector") + else: + kind, lists, ranks = "encode", Registry._representations, ("any", "any") + load_builtin_operators(ANY_TYPE) + input_types = frozenset(m for m, ops in lists.items() if op_cls in ops) + return OperatorSignature(kind, input_types, _output_type(op_cls), 1, *ranks) + + +@lru_cache(maxsize=None) +def _needs_alignment(op_cls) -> bool: + return op_cls().needs_alignment + + +@lru_cache(maxsize=None) +def _output_type(op_cls): + return getattr(op_cls, "OUTPUT_MODALITY_TYPE", None) or getattr( + op_cls(), "output_modality_type", None + ) + + +class Violation(NamedTuple): + severity: str # error | implicit + kind: str # input_type | arity | hidden_aggregation | zero_padding + node_id: str + operator: str + detail: str + + +def typecheck(dag, leaf_types, leaf_stats=None) -> list: + nodes = {node.node_id: node for node in dag.nodes} + types, violations = {}, [] + stats = output_stats(nodes.__getitem__, leaf_stats) + + def key(node): + return node.modality_id, node.representation_index + + def flag(severity, kind, node, detail): + name = node.operation.__name__ + violations.append(Violation(severity, kind, node.node_id, name, detail)) + + def check_fusion(node): + operands = [stats(i) for i in node.inputs] + if None in operands: + return + op = node.operation(params=node.parameters) + if not op.preserves_leading_axis: + ranks = [len(s.output_shape) for s in operands] + if max(ranks) > 1: + flag("implicit", "hidden_aggregation", node, f"operand ranks {ranks}") + if op.needs_alignment: + widths = [s.output_shape[-1] for s in op._fusion_input_stats(operands)] + if len(set(widths)) > 1: + flag("implicit", "zero_padding", node, f"operand widths {widths}") + + def out_type(nid): + if nid in types: + return types[nid] + node = nodes[nid] + if not node.inputs: + types[nid] = leaf_types.get(key(node)) + return types[nid] + in_types = [out_type(i) for i in node.inputs] + if node.operation is None: + types[nid] = in_types[0] + return types[nid] + sig = operator_signature(node.operation, node.parameters) + if (sig.arity == "n") != (len(node.inputs) > 1): + flag("error", "arity", node, f"{len(node.inputs)} input(s)") + for t in in_types: + if t is not None and t not in sig.input_types: + flag("error", "input_type", node, f"{t.name} input") + if sig.kind == "fuse" and leaf_stats is not None: + check_fusion(node) + types[nid] = sig.output_type or in_types[0] + return types[nid] + + out_type(dag.root_node_id) + return violations + + +def output_stats(get_node, leaf_stats): + from systemds.scuro.drsearch.representation_dag import ( + pushdown_aggregation_for_node, + ) + + memo = {} + + def stats(nid): + if nid not in memo: + node = get_node(nid) + memo[nid] = None + if not node.inputs: + key = node.modality_id, node.representation_index + memo[nid] = (leaf_stats or {}).get(key) + elif node.operation is not None: + inputs = [stats(i) for i in node.inputs] + if None not in inputs: + try: + out = node.operation(params=node.parameters).get_output_stats( + inputs if len(inputs) > 1 else inputs[0] + ) + agg = pushdown_aggregation_for_node(node.parameters) + memo[nid] = agg.get_output_stats(out) if agg else out + except Exception: + pass + return memo[nid] + + return stats + + +def sample_stats(modality) -> Optional[RepresentationStats]: + data = getattr(modality, "data", None) + if data is None or len(data) == 0: + return None + try: + sample = np.asarray(data[0]) + except (TypeError, ValueError): + return None + if sample.dtype == object: + return None + shape = tuple(sample.shape) if sample.ndim else (1,) + return RepresentationStats(len(data), shape, output_shape_is_known=True) + + +def explicit_prep_on(algebra) -> bool: + return algebra.explicit_fusion_prep or algebra.explicit_padding == "pad" + + +def explicit_fusion_operands( + builder, operation, parameters, inputs, leaf_stats, algebra +) -> list: + if not issubclass(operation, Fusion): + return inputs + op = operation(params=parameters) + stats = output_stats(builder.get_node, leaf_stats) + if algebra.explicit_fusion_prep and not op.preserves_leading_axis: + agg = AggregatedRepresentation(target_dimensions=1, aggregate_leading=True) + inputs = [ + ( + builder.create_operation_node( + AggregatedRepresentation, [i], agg.get_current_parameters() + ) + if stats(i) is not None and len(stats(i).output_shape) > 1 + else i + ) + for i in inputs + ] + shapes = [stats(i) for i in inputs] + if ( + algebra.explicit_padding == "pad" + and op.needs_alignment + and None not in shapes + and all(len(s.output_shape) == 1 for s in shapes) + ): + width = max(s.output_shape[0] for s in shapes) + inputs = [ + ( + builder.create_operation_node( + ZeroPad, + [i], + ZeroPad(width, s.output_shape[0]).get_current_parameters(), + ) + if s.output_shape[0] < width + else i + ) + for i, s in zip(inputs, shapes) + ] + return inputs + + +def count_explicit_prep(counter, dag) -> None: + """Counts the explicit operand preparation in a plan about to be evaluated.""" + nodes = {n.node_id: n for n in dag.filter_connected_nodes(dag.nodes)} + padded = False + for node in nodes.values(): + if not (node.operation and issubclass(node.operation, Fusion)): + continue + for operand in (nodes[i] for i in node.inputs): + if operand.operation is ZeroPad: + padded = True + ratio = operand.parameters["width"] / operand.parameters["input_width"] + counter["explicit_paddings"] += 1 + counter["explicit_padding_ratio_sum"] += ratio + counter["explicit_padding_ratio_max"] = max( + counter["explicit_padding_ratio_max"], ratio + ) + operand = nodes[operand.inputs[0]] + if operand.operation is AggregatedRepresentation: + counter["explicit_aggregations"] += 1 + counter["explicit_prep_plans"] += 1 + counter["explicit_padding_plans"] += padded + + +def count_violations(counter, violations) -> list: + """Adds per-kind and per-operator counts; returns the errors.""" + counter["typecheck_plans"] += 1 + for v in violations: + counter[f"typecheck_{v.kind}"] += 1 + counter[f"typecheck_{v.kind}:{v.operator}"] += 1 + return [v for v in violations if v.severity == "error"] + + +def order_invariant(config: AlgebraConfig, tasks) -> bool: + if config.feature_order_invariant is not None: + return config.feature_order_invariant + return bool(tasks) and all( + getattr(task.model, "feature_order_invariant", None) is True for task in tasks + ) + + +def canonical_signature(dag, laws, invariant, node_id=None) -> tuple: + # representation_dag imports this module. + from systemds.scuro.drsearch.representation_dag import _spec_safe + + def view(nid): + node = dag.get_node_by_id(nid) + op = node.operation + params = {k: _spec_safe(v) for k, v in (node.parameters or {}).items()} + name = op.__name__ if op is not None else None + return ( + name, + op, + params, + node.inputs, + node.modality_id, + node.representation_index, + ) + + return _signature(view, node_id or dag.root_node_id, laws, invariant) + + +def canonical_signature_from_spec(spec: dict, laws, invariant) -> tuple: + nodes = {node["id"]: node for node in spec["nodes"]} + classes = _fusion_classes() + + def view(nid): + n = nodes[nid] + return ( + n["op"], + classes.get(n["op"]), + n["params"], + n["inputs"], + n["modality_id"], + n["representation_index"], + ) + + return _signature(view, spec["root"], laws, invariant) + + +def canonical_id(dag, laws, invariant, node_id=None) -> str: + """`dag` is a RepresentationDag or a to_spec() dict.""" + if isinstance(dag, dict): + sig = canonical_signature_from_spec(dag, laws, invariant) + else: + sig = canonical_signature(dag, laws, invariant, node_id) + blob = json.dumps(sig, sort_keys=True, separators=(",", ":")) + return hashlib.sha1(blob.encode()).hexdigest()[:16] + + +def _fusion_classes() -> dict: + classes = {cls.__name__: cls for cls in Registry().get_fusion_operators()} + stack = [Fusion] + while stack: + cls = stack.pop() + classes.setdefault(cls.__name__, cls) + stack.extend(cls.__subclasses__()) + return classes + + +def _params_key(params) -> str: + return json.dumps(params, sort_keys=True, separators=(",", ":")) + + +def _signature(view, root, laws, invariant) -> tuple: + # The aggregation Fusion._prepare_modalities applies to every operand of a + # fusion that does not preserve the leading axis. + implicit_agg = ( + "op", + AggregatedRepresentation.__name__, + _params_key( + AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True + ).get_current_parameters() + ), + ) + memo = {} + + def sig(nid, laws): + if (nid, laws) not in memo: + memo[nid, laws] = compute(nid, laws) + return memo[nid, laws] + + def compute(nid, laws): + name, op, params, inputs, modality_id, representation_index = view(nid) + if not inputs: + return ("leaf", modality_id, representation_index) + + children = [sig(i, laws & EXACT_LAWS if _fitted(op) else laws) for i in inputs] + params = dict(params) + pushed = ( + params.pop("_pushdown_aggregation", None) + if "agg_pushdown" in laws + else None + ) + head = ("op", name, _params_key(params)) + fusion = isinstance(op, type) and issubclass(op, Fusion) + op_laws = op.algebraic_laws(params) & laws if fusion else frozenset() + + if ( + fusion + and "implicit_agg" in laws + and not params.get("preserve_leading_axis", op.preserves_leading_axis) + ): + children = [c[3][0] if c[:3] == implicit_agg else c for c in children] + if fusion and "implicit_pad" in laws and _needs_alignment(op): + children = [ + c[3][0] if c[:2] == ("op", ZeroPad.__name__) else c for c in children + ] + if "assoc" in op_laws: + children = [g for c in children for g in (c[3] if c[:3] == head else (c,))] + if "avg_flatten" in op_laws: + weights = Counter() + for c in children: + for w, g in c[3] if c[:3] == head else (("1", c),): + weights[g] += Fraction(w) / len(children) + children = sorted(((str(w), g) for g, w in weights.items()), key=repr) + elif "comm" in op_laws or ("concat_perm" in op_laws and invariant): + children = sorted(children, key=repr) + + out = head + (tuple(children),) + if pushed is not None: + out = ("op", AggregatedRepresentation.__name__, _params_key(pushed), (out,)) + return out + + return sig(root, frozenset(laws)) + + +def count_fusion_plans(n, fusion_ops, laws, invariant) -> dict: + heads = _fusion_heads(fusion_ops, laws) + f = len(heads) + reorders = sum(_reorders(op_laws, invariant) for *_, op_laws in heads) + ordered = factorial(2 * n - 2) // factorial(n - 1) * f ** (n - 1) + if n == 1: + canonical = 1 + elif reorders == f and all("assoc" in op_laws for *_, op_laws in heads): + canonical = f * _flattened_trees(n, f) + elif n <= 4: + canonical = _brute_force_canonical(n, fusion_ops, laws, invariant) + else: + canonical = None + return { + "ordered": ordered, + "legacy_exhaustive": factorial(n) * ordered, + "commutative": _commutative_plans(n, heads, invariant), + "canonical": canonical, + } + + +def count_search_space( + ks, fusion_ops, laws, invariant, min_modalities=2, max_modalities=None +) -> dict: + """`ks[i]` is the number of representations of modality i.""" + totals = dict.fromkeys( + ("ordered", "legacy_exhaustive", "commutative", "canonical"), 0 + ) + for r in range(min_modalities, (max_modalities or len(ks)) + 1): + leaves = sum(prod(c) for c in itertools.combinations(ks, r)) + for key, plans in count_fusion_plans(r, fusion_ops, laws, invariant).items(): + if plans is None or totals[key] is None: + totals[key] = None + else: + totals[key] += leaves * plans + return totals + + +def _fitted(op) -> bool: + return getattr(op, "fit_requirement", None) not in (None, FitRequirement.NONE) + + +def _commutative_plans(n, heads, invariant) -> int: + plans = {(1, True): 1, (1, False): 1} + for m in range(2, n + 1): + for free in (True, False): + plans[m, free] = 0 + for *_, op, _, op_laws in heads: + below = free and not _fitted(op) + splits = sum( + comb(m, k) * plans[k, below] * plans[m - k, below] + for k in range(1, m) + ) + plans[m, free] += splits // ( + 2 if free and _reorders(op_laws, invariant) else 1 + ) + return plans[n, True] + + +def _reorders(op_laws, invariant) -> bool: + return "comm" in op_laws or ("concat_perm" in op_laws and invariant) + + +def _fusion_heads(fusion_ops, laws) -> list: + from systemds.scuro.drsearch.representation_dag import _spec_safe + + heads = [] + for op in fusion_ops: + params = {k: _spec_safe(v) for k, v in op().get_current_parameters().items()} + heads.append((op.__name__, op, params, op.algebraic_laws(params) & laws)) + return heads + + +def _flattened_trees(n, f) -> int: + tree, block, partitions = [0, 1], [0, 1], [1, 1] + for m in range(2, n + 1): + tree.append( + sum(comb(m - 1, j - 1) * block[j] * partitions[m - j] for j in range(1, m)) + ) + block.append((f - 1) * tree[m]) + partitions.append(tree[m] + block[m]) + return tree[n] + + +def _ordered_trees(leaves): + if len(leaves) == 1: + yield leaves[0] + return + for split in range(1, len(leaves)): + for left in itertools.combinations(leaves, split): + right = tuple(i for i in leaves if i not in left) + for l_tree in _ordered_trees(left): + for r_tree in _ordered_trees(right): + yield (l_tree, r_tree) + + +def _brute_force_canonical(n, fusion_ops, laws, invariant) -> int: + heads = _fusion_heads(fusion_ops, laws) + signatures = set() + for tree in _ordered_trees(tuple(range(n))): + for choice in itertools.product(heads, repeat=n - 1): + nodes, ops = {}, iter(choice) + + def add(t): + if isinstance(t, int): + nodes[t] = (None, None, {}, (), t, 0) + return t + inputs = (add(t[0]), add(t[1])) + name, op, params, _ = next(ops) + nid = f"n{len(nodes)}" + nodes[nid] = (name, op, params, inputs, None, None) + return nid + + root = add(tree) + signatures.add(_signature(nodes.__getitem__, root, laws, invariant)) + return len(signatures) diff --git a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py index bfb62bcc890..7805ac66b3f 100644 --- a/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py +++ b/src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py @@ -32,7 +32,10 @@ import random from systemds.scuro.drsearch.representation_dag import ( RepresentationNode, + get_modality_by_id_and_instance_id, ) +from systemds.scuro.representations.representation import infer_stats_from_data +from systemds.scuro.utils.memory_utility import cpu_memory_budget_bytes from systemds.scuro.modality.modality import Modality from systemds.scuro.drsearch.split_aware_evaluator import ( SplitAwareDAGEvaluator, @@ -151,6 +154,65 @@ def _is_window_operation(op: Any) -> bool: return False +def _is_leading_axis_concatenation(node: RepresentationNode) -> bool: + from systemds.scuro.representations.concatenation import Concatenation + + return ( + inspect.isclass(node.operation) + and issubclass(node.operation, Concatenation) + and bool((node.parameters or {}).get("preserve_leading_axis")) + ) + + +def _tied_window_params(dag) -> Dict[str, str]: + parent = {} + + def find(node_id): + while parent[node_id] != node_id: + node_id = parent[node_id] + return node_id + + def windows_under(node_id): + node = dag.get_node_by_id(node_id) + if _is_window_operation(node.operation): + return [node_id] + if _is_leading_axis_concatenation(node): + return [w for input_id in node.inputs for w in windows_under(input_id)] + return [] + + for node in dag.nodes: + if not _is_leading_axis_concatenation(node): + continue + windows = windows_under(node.node_id) + for window_id in windows: + parent.setdefault(window_id, window_id) + for window_id in windows[1:]: + parent[find(window_id)] = find(windows[0]) + + tied = {} + canonical = {} + for window_id in parent: + name = dag.get_node_by_id(window_id).operation.granularity_parameter + if name is None: + continue + key = (find(window_id), name) + canonical.setdefault(key, f"{window_id}-{name}") + if canonical[key] != f"{window_id}-{name}": + tied[f"{window_id}-{name}"] = canonical[key] + return tied + + +def _expand_tied_params( + config: Dict[str, Any], tied_params: Dict[str, str] +) -> Dict[str, Any]: + return { + **config, + **{ + alias: config[name] for alias, name in tied_params.items() if name in config + }, + } + + def _materialize_node_params( node: RepresentationNode, flat_params: Dict[str, Any] ) -> Dict[str, Any]: @@ -581,6 +643,8 @@ def visit_node(node_id): for node_id in hyperparams } param_specs = self._build_param_specs(hyperparams, node_parameters) + tied_params = _tied_window_params(dag) + param_specs = [s for s in param_specs if s["name"] not in tied_params] discrete_size = self._estimate_discrete_search_size(param_specs) n_calls = min(discrete_size, max_evals) if max_evals else discrete_size all_results = self._search_best_configs( @@ -593,6 +657,7 @@ def visit_node(node_id): budget=n_calls, initial_config=None, rep_name=rep_name, + tied_params=tied_params, ) all_results.append(baseline) @@ -924,6 +989,7 @@ def _search_best_configs( budget: int, initial_config: Optional[Dict[str, Any]], rep_name: str = "", + tied_params: Optional[Dict[str, str]] = None, ) -> List[Tuple[Dict[str, Any], Any]]: import optuna from optuna.trial import TrialState @@ -1002,12 +1068,26 @@ def _search_best_configs( trial_results: List[Tuple[Dict[str, Any], Any]] = [] def objective(trial: optuna.Trial) -> float: - config = self._suggest_config_from_specs(trial, param_specs) + config = _expand_tied_params( + self._suggest_config_from_specs(trial, param_specs), tied_params or {} + ) key = self._config_key(config) if key in seen: trial.set_user_attr("duplicate", True) raise optuna.TrialPruned() + estimate = self._estimate_trial_peak_bytes( + dag, config, modality_ids, task, modalities_override + ) + budget_bytes = cpu_memory_budget_bytes() + if estimate is not None and estimate > budget_bytes: + self.logger.warning( + f"{rep_name}: skipping trial estimated at {estimate / 1024**3:.1f} GB " + f"(budget {budget_bytes / 1024**3:.1f} GB): {config}" + ) + trial.set_user_attr("over_memory_budget", True) + raise optuna.TrialPruned() + params, scores = self.evaluate_dag_config( dag, config, @@ -1072,6 +1152,47 @@ def _get_cached_modalities_for_task(self, task, modality_ids): ) return cached_modalities + def _estimate_trial_peak_bytes( + self, dag, params, modality_ids, task, modalities_override=None + ) -> Optional[int]: + modalities = ( + modalities_override + if modalities_override is not None + else self.get_modalities_by_id(modality_ids) + ) + stats = {} + + def visit(node_id): + if node_id in stats: + return stats[node_id] + node = dag.get_node_by_id(node_id) + if not node.inputs: + modality = get_modality_by_id_and_instance_id( + modalities, int(node.modality_id), node.representation_index + ) + stats[node_id] = ( + modality.get_stats() + if hasattr(modality, "get_stats") + else infer_stats_from_data(modality.data) + ) + return stats[node_id] + inputs = [visit(input_id) for input_id in node.inputs] + if any(s is None for s in inputs): + stats[node_id] = None + return None + inputs = inputs[0] if len(inputs) == 1 else inputs + operation = node.operation(params=_apply_trial_params_to_node(node, params)) + peaks.append(operation.estimate_peak_memory_bytes(inputs)["cpu_peak_bytes"]) + stats[node_id] = operation.get_output_stats(inputs) + return stats[node_id] + + peaks = [] + output_stats = visit(dag.root_node_id) + if output_stats is None: + return None + peaks.append(task.estimate_peak_memory_bytes(output_stats)["cpu_peak_bytes"]) + return int(max(peaks)) + def _trial_node_cache(self, dag, params): store = getattr(self, "_trial_cache", None) if store is None: diff --git a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py index b7472a35a6d..1f78ff67421 100644 --- a/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py +++ b/src/main/python/systemds/scuro/drsearch/modality_shared_memory.py @@ -18,11 +18,28 @@ # under the License. # # ------------------------------------------------------------- +import os +import shutil from typing import Any, List, Tuple import numpy as np from multiprocessing import shared_memory SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024 +_SHM_DIR = "/dev/shm" + + +def _shm_reserve_bytes() -> int: + return int(float(os.environ.get("SCURO_SHM_RESERVE_GB", "8")) * 1024**3) + + +def _should_offload(data_nbytes: int) -> bool: + if data_nbytes < SHARED_MEMORY_MIN_BYTES: + return False + try: + free = shutil.disk_usage(_SHM_DIR).free + except OSError: + return True + return data_nbytes + _shm_reserve_bytes() <= free def unlink_shm(name: str) -> None: @@ -324,7 +341,7 @@ def _is_nested_shared_memory_candidate(data: Any) -> bool: def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: if _is_shared_memory_candidate(data): data_nbytes = _estimate_ndarray_list_nbytes(data) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): dtype = data[0].dtype offsets = [] total_elems = 0 @@ -351,7 +368,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: elif _is_shared_ndarray_candidate(data): arr = data data_nbytes = int(arr.nbytes) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): shm = shared_memory.SharedMemory(create=True, size=data_nbytes) shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf) np.copyto(shm_arr, arr, casting="no") @@ -371,7 +388,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: leaves.append(arr) group_bounds.append((lo, len(leaves))) data_nbytes = _estimate_ndarray_list_nbytes(leaves) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): dtype = leaves[0].dtype offsets = [] total_elems = 0 @@ -400,7 +417,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool: elif _is_string_list_shared_memory_candidate(data): encoded = [s.encode("utf-8") for s in data] data_nbytes = int(sum(len(b) for b in encoded)) - if data_nbytes >= SHARED_MEMORY_MIN_BYTES: + if _should_offload(data_nbytes): shm = shared_memory.SharedMemory(create=True, size=data_nbytes) mv = shm.buf cursor = 0 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 79a04ca262c..0c182338101 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_ga_optimizer.py @@ -22,6 +22,7 @@ from __future__ import annotations import copy +from collections import Counter import gc import os import pickle @@ -42,6 +43,18 @@ unlink_shm, ) +from systemds.scuro.drsearch.algebra import ( + AlgebraConfig, + canonical_id, + canonical_signature, + count_explicit_prep, + count_violations, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + sample_stats, + typecheck, +) from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.drsearch.representation_dag import ( RepresentationDAGBuilder, @@ -53,6 +66,7 @@ SplitAwareDAGEvaluator, dag_requires_fitting, ) +from systemds.scuro.drsearch.trial_result_cache import TrialNodeCache, TrialResultCache from systemds.scuro.drsearch.worker_pool import ( PersistentWorkerPool, create_mp_context, @@ -62,7 +76,7 @@ AggregatedRepresentation, ) from systemds.scuro.representations.fusion import FusionInputTooLargeError -from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.utils.memory_utility import gpu_memory_info from systemds.scuro.utils.schema_helpers import get_shape Tree = Union[int, Tuple["Tree", "Tree"]] @@ -175,6 +189,7 @@ class DagGenome: _TIMING_OBJECTIVES = {"runtime", "task_time", "representation_time"} +_MAX_ERROR_SAMPLES = 5 ObjectiveSpec = Tuple[str, str] # (name, "max" | "min") @@ -187,6 +202,14 @@ def _objective_value( return val_score[name] +def _task_uses_gpu(task: Task) -> bool: + uses_gpu = getattr(task.model, "uses_gpu", None) + if uses_gpu is not None: + return bool(uses_gpu) + device = getattr(task.model, "device", None) + return str(getattr(device, "type", device) or "").startswith("cuda") + + def _failure_fitness(objective_specs: List[ObjectiveSpec]) -> Tuple[float, ...]: return tuple( float("-inf") if direction == "max" else float("inf") @@ -205,15 +228,45 @@ def _release_cuda_memory() -> None: pass +_SUBTREE_CACHE: Optional[TrialResultCache] = None + + +def _subtree_cache() -> TrialResultCache: + global _SUBTREE_CACHE + if _SUBTREE_CACHE is None: + gb = float(os.environ.get("SCURO_GA_SUBTREE_CACHE_GB", "2")) + _SUBTREE_CACHE = TrialResultCache(int(gb * 1024**3)) + return _SUBTREE_CACHE + + +def _subtree_node_cache(dag: RepresentationDag, task: Task) -> TrialNodeCache: + signatures = { + node.node_id: ( + task.model.name, + canonical_signature(dag, frozenset(), False, node.node_id), + ) + for node in dag.nodes + if node.inputs + } + return TrialNodeCache(_subtree_cache(), signatures) + + def _evaluate_genome_body( dag: RepresentationDag, task: Task, modalities: List[Any], objective_specs: List[ObjectiveSpec], + share_subtrees: bool = False, ) -> Tuple[Optional[Tuple[float, ...]], Optional[Dict[str, Any]]]: start = time.time() if dag_requires_fitting(dag): - evaluation = SplitAwareDAGEvaluator(dag, modalities, task).evaluate() + node_cache = _subtree_node_cache(dag, task) if share_subtrees else None + if node_cache is not None: + store = node_cache.store + before = store.hits, store.misses, store.saved_seconds + evaluation = SplitAwareDAGEvaluator( + dag, modalities, task, node_cache=node_cache, share_fitted=share_subtrees + ).evaluate() scores = evaluation.scores total = evaluation.representation_time + evaluation.task_time val_score = scores[1].average_scores @@ -238,6 +291,13 @@ def _evaluate_genome_body( "task_timing": getattr(task, "last_run_timing", {}), **timing, } + if node_cache is not None: + payload["subtree_cache"] = { + "hits": store.hits - before[0], + "misses": store.misses - before[1], + "saved_seconds": store.saved_seconds - before[2], + "bytes": store.bytes, + } return fitness, payload fused = dag.execute(modalities, task, enable_cache=False) @@ -281,11 +341,11 @@ def _evaluate_genome_body( def _dispatch_genome_evaluation(payload, _gpu_id): - dag, task, modalities, objective_specs = payload + dag, task, modalities, objective_specs, *share = payload try: try: fitness, result = _evaluate_genome_body( - dag, task, modalities, objective_specs + dag, task, modalities, objective_specs, *share ) if fitness is None: fitness = _failure_fitness(objective_specs) @@ -310,6 +370,9 @@ def __init__(self, values, fitness_type): class MultimodalDeapOptimizer: + # For instances made without __init__, as regret_analysis/exhaustive_enumerator.py does. + algebra = AlgebraConfig() + def __init__( self, modalities: List[Any], @@ -336,6 +399,7 @@ def __init__( allow_repeated_modalities: bool = False, threads_per_worker: Optional[int] = None, max_feature_dim: Optional[int] = DEFAULT_MAX_FUSED_FEATURE_DIM, + algebra: Optional[AlgebraConfig] = None, ): self.modalities = modalities self.tasks = tasks @@ -343,6 +407,13 @@ def __init__( self.allow_repeated_modalities = allow_repeated_modalities self.max_feature_dim = None if max_feature_dim is None else int(max_feature_dim) self.oversized_candidates: Dict[str, int] = {} + self.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) + self._canonical_signatures: Dict[Tuple, Tuple] = {} + self._seen_legacy_signatures: set = set() + # (task, canonical signature) -> [first evaluation runtime, evaluations avoided] + self._canonical_classes: Dict[Tuple, List] = {} self.min_modalities = max(1, min_modalities) requested_max = max_modalities or len(modalities) @@ -399,6 +470,7 @@ def __init__( self.optimization_results: Dict[str, List[FusionSearchResult]] = {} self.evaluation_errors: Dict[str, int] = {} + self.evaluation_error_samples: Dict[str, List[str]] = {} self.hall_of_fame_size = max(1, hall_of_fame_size) self.hall_of_fame: Dict[str, List[FusionSearchResult]] = {} @@ -659,11 +731,18 @@ def _start_parallel_runtime(self, task_name: str) -> None: shared_modality._data = wrapped shm_names.append(shm_name) shared_modalities.append(shared_modality) + # Workers are spawned with CUDA_VISIBLE_DEVICES="" unless they get + # a GPU lane, so without these a CUDA task model cannot run at all. worker_pool = PersistentWorkerPool( self.max_workers, _WORKER_DISPATCH, ctx=create_mp_context(), threads_per_worker=self.threads_per_worker, + gpu_devices=[info["index"] for info in gpu_memory_info()], + gpu_slots_per_device=int( + os.environ.get("SCURO_GPU_SLOTS_PER_DEVICE", "1") + ), + gpu_demand_fraction=1.0, ) except Exception: for shm_name in shm_names: @@ -684,6 +763,17 @@ def _shutdown_parallel_runtime(self) -> None: unlink_shm(shm_name) self._parallel_shm_names = [] + def _idle_lane(self, uses_gpu: bool): + pool = self._worker_pool + if uses_gpu: + for gpu_id in pool.gpu_worker_devices: + if pool.has_idle_worker_for(gpu_id): + return gpu_id + return None if pool.has_idle_cpu_worker else False + if pool.has_idle_worker_for(None, allow_gpu_worker_for_cpu=True): + return None + return False + def _evaluate_individuals_parallel( self, individuals: List[Any], task: Task ) -> None: @@ -711,20 +801,25 @@ def _evaluate_individuals_parallel( ind.fitness.values = self._reject_oversized(task_name, genome, reason) continue dag = self._genome_to_dag(genome) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) pending_followers[sig] = [] pending_work.append((ind, dag, sig)) + uses_gpu = _task_uses_gpu(task) try: while pending_work or jobs: - while ( - pending_work - and self._worker_pool.has_idle_worker - and len(jobs) < self.batch_size - ): + while pending_work and len(jobs) < self.batch_size: + lane = self._idle_lane(uses_gpu) + if lane is False: + break ind, dag, sig = pending_work.pop(0) job_id = self._worker_pool.submit( "genome", - (dag, task, self._parallel_modalities, self.objective_specs), + (dag, task, self._parallel_modalities, self.objective_specs) + + self._share_args(), + gpu_id=lane, + allow_gpu_worker_for_cpu=not uses_gpu, ) jobs[job_id] = (ind, dag, sig) @@ -771,10 +866,22 @@ def _record_evaluation( self.evaluation_errors[task_name] = ( self.evaluation_errors.get(task_name, 0) + 1 ) - if self.debug and error is not None: + samples = self.evaluation_error_samples.setdefault(task_name, []) + if error is not None and len(samples) < _MAX_ERROR_SAMPLES: + spec = dag.to_spec() if dag is not None else None + samples.append(f"{spec}\n{error}") + if error is not None and (self.debug or len(samples) == 1): print(f"[GA] genome evaluation failed for task={task_name}:\n{error}") return + cache_stats = payload.get("subtree_cache") + if cache_stats: + for key in ("hits", "misses", "saved_seconds"): + self.algebra_stats[f"subtree_cache_{key}"] += cache_stats[key] + self.algebra_stats["subtree_cache_max_bytes"] = max( + self.algebra_stats["subtree_cache_max_bytes"], cache_stats["bytes"] + ) + result = FusionSearchResult( dag=dag, train_score=payload["train_score"], @@ -798,6 +905,13 @@ def _record_evaluation( ) self.optimization_results.setdefault(task_name, []).append(result) self._update_hall_of_fame(task_name, result) + if self.algebra.ga_canonical_dedup: + cls = self._canonical_classes.setdefault( + (task_name, self._canonical_signature(dag)), [None, 0] + ) + if cls[0] is None: + cls[0] = result.runtime + self.algebra_stats["ga_seconds_avoided"] += cls[0] * cls[1] self._eval_counter += 1 def _dominates(self, a: Tuple[float, ...], b: Tuple[float, ...]) -> bool: @@ -962,18 +1076,35 @@ def _genome_to_dag(self, genome: DagGenome) -> RepresentationDag: builder.create_leaf_node(mod_id, repr_idx) for mod_id, repr_idx in genome.leaves ] + leaf_stats = ( + { + leaf: self._leaf_stats(self._current_task_name, leaf) + for leaf in genome.leaves + } + if self is not None + and explicit_prep_on(self.algebra) + and self._current_task_name is not None + else None + ) 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") + inputs = [build(left, path + "L"), 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() - ) + params = op.get_current_parameters() + if leaf_stats is not None: + inputs = explicit_fusion_operands( + builder, + op.__class__, + params, + inputs, + leaf_stats, + self.algebra, + ) + return builder.create_operation_node(op.__class__, inputs, params) return builder.build(build(genome.tree, "")) @@ -981,11 +1112,41 @@ 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 ( + legacy = ( tuple(g.leaves), norm(g.tree), tuple(sorted((p, c.__name__) for p, c in g.fusion_ops.items())), ) + if not self.algebra.ga_canonical_dedup: + return legacy + canon = self._canonical_signatures.get(legacy) + if canon is None: + canon = self._canonical_signature(self._genome_to_dag(g)) + self._canonical_signatures[legacy] = canon + self._count_signature_request(legacy, canon) + return canon + + def _canonical_signature(self, dag: RepresentationDag) -> Tuple: + return ( + "canon", + canonical_id(dag, self.algebra.laws, self._order_invariant), + ) + + def _count_signature_request(self, legacy: Tuple, canon: Tuple) -> None: + task_name = self._current_task_name + self.algebra_stats["ga_signature_requests"] += 1 + if (task_name, legacy) in self._seen_legacy_signatures: + return + self._seen_legacy_signatures.add((task_name, legacy)) + self.algebra_stats["ga_unique_legacy"] += 1 + cls = self._canonical_classes.get((task_name, canon)) + if cls is None: + self._canonical_classes[(task_name, canon)] = [None, 0] + self.algebra_stats["ga_unique_canonical"] += 1 + return + cls[1] += 1 + self.algebra_stats["ga_evaluations_avoided"] += 1 + self.algebra_stats["ga_seconds_avoided"] += cls[0] or 0.0 def _leaf_stats(self, task_name: str, leaf: Tuple[Any, int]): """Shape of one candidate leaf representation, or None if unknown.""" @@ -996,18 +1157,7 @@ def _leaf_stats(self, task_name: str, leaf: Tuple[Any, int]): ] except (KeyError, IndexError, TypeError): return None - - data = getattr(modality, "data", None) - if data is None or len(data) == 0: - return None - try: - sample = np.asarray(data[0]) - except (TypeError, ValueError): - return None - if sample.dtype == object: - return None - shape = tuple(sample.shape) if sample.ndim else (1,) - return RepresentationStats(len(data), shape, output_shape_is_known=True) + return sample_stats(modality) def _predicted_feature_dim( self, task_name: str, genome: DagGenome @@ -1071,7 +1221,24 @@ def walk(subtree: Tree, path: str): except Exception: return None + def _type_error(self, task_name: str, genome: DagGenome) -> Optional[str]: + reps = self.k_best_representations[task_name] + types = { + (mid, idx): getattr(reps[mid][idx], "modality_type", None) + for mid, idx in genome.leaves + } + stats = {leaf: self._leaf_stats(task_name, leaf) for leaf in genome.leaves} + violations = typecheck(self._genome_to_dag(genome), types, stats) + errors = count_violations(self.algebra_stats, violations) + if errors and self.algebra.typecheck == "strict": + return f"type_error: {errors}" + return None + def _oversized_reason(self, task_name: str, genome: DagGenome) -> Optional[str]: + if self.algebra.typecheck != "off" and task_name is not None: + reason = self._type_error(task_name, genome) + if reason is not None: + return reason reason = self._operator_input_violation(task_name, genome) if reason is not None: return f"{reason}; candidate rejected before execution" @@ -1112,13 +1279,15 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: return self._reject_oversized(task_name, genome, reason) dag = self._genome_to_dag(genome) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) modalities = list( chain.from_iterable(self.k_best_representations[task_name].values()) ) try: fitness, payload = _evaluate_genome_body( - dag, task, modalities, self.objective_specs + dag, task, modalities, self.objective_specs, *self._share_args() ) error = None if fitness is None: @@ -1135,6 +1304,10 @@ def _evaluate_genome(self, genome: DagGenome, task: Task) -> Tuple[float, ...]: cache[sig] = fitness return fitness + def _share_args(self) -> Tuple: + # Legacy passes no extra argument, so 4-argument overrides keep working. + return (True,) if self.algebra.share_fused_subtrees else () + def _crossover_genomes( self, g1: DagGenome, g2: DagGenome ) -> Tuple[DagGenome, DagGenome]: diff --git a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py index 54cddfd142b..22476bce60c 100644 --- a/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/multimodal_optimizer.py @@ -37,12 +37,23 @@ from systemds.scuro.representations.aggregated_representation import ( AggregatedRepresentation, ) +from systemds.scuro.drsearch.algebra import ( + AlgebraConfig, + canonical_id, + count_explicit_prep, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + sample_stats, +) from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.utils.schema_helpers import get_shape from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED +import math import pickle import copy +from collections import Counter import time import traceback from itertools import chain @@ -128,6 +139,7 @@ def __init__( checkpoint_every: int = None, checkpoint_dir: Optional[str] = None, resume: bool = False, + algebra: Optional[AlgebraConfig] = None, ): self.modalities = modalities self.tasks = tasks @@ -135,6 +147,9 @@ def __init__( self.debug = debug self.min_modalities = max(2, min_modalities) self.max_modalities = max_modalities or len(modalities) + self.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) self.operator_registry = Registry() self.fusion_operators = self.operator_registry.get_fusion_operators() @@ -217,7 +232,7 @@ def optimize_parallel( break for dag in self._generate_fusion_dags( - modality_subset, repr_combo + modality_subset, repr_combo, task.model.name ): if max_combinations and evaluated_count >= max_combinations: stop_generation = True @@ -353,9 +368,20 @@ def _generate_representation_combinations( } def _generate_fusion_dags( - self, modality_subset: List[str], representation_combo: Dict[str, int] + self, + modality_subset: List[str], + representation_combo: Dict[str, int], + task_name: str = None, ) -> Generator[RepresentationDag, None, None]: leaf_infos = [(m, representation_combo[m]) for m in modality_subset] + leaf_stats = ( + { + (m, i): sample_stats(self.k_best_representations[task_name][m][i]) + for m, i in leaf_infos + } + if explicit_prep_on(self.algebra) and task_name is not None + else None + ) def gen_trees(indices: List[int]): if len(indices) == 1: @@ -393,18 +419,35 @@ def build_variants( for fusion_op_class in self.fusion_operators: new_builder = copy.deepcopy(right_builder) fusion_op = fusion_op_class() + params = fusion_op.get_current_parameters() + inputs = [left_root, right_root] + if leaf_stats is not None: + inputs = explicit_fusion_operands( + new_builder, + fusion_op.__class__, + params, + inputs, + leaf_stats, + self.algebra, + ) fusion_id = new_builder.create_operation_node( - fusion_op.__class__, - [left_root, right_root], - fusion_op.get_current_parameters(), + fusion_op.__class__, inputs, params ) variants.append((new_builder, fusion_id)) return variants n = len(leaf_infos) + dedup_permutations = self.algebra.exhaustive_dedup_permutations + # gen_trees already yields every leaf order, so permuting the leaves repeats each tree n! times. + leaf_orders = ( + [leaf_infos] + if dedup_permutations + else itertools.permutations(leaf_infos, n) + ) + seen_canonical = set() - for permuted_leaf_infos in itertools.permutations(leaf_infos, n): + for permuted_leaf_infos in leaf_orders: base_builder = RepresentationDAGBuilder() leaf_id_map = {} for idx, (modality_id, repr_idx) in enumerate(permuted_leaf_infos): @@ -416,12 +459,28 @@ def build_variants( for tree in gen_trees(indices): variants = build_variants(tree, base_builder, leaf_id_map) for builder_variant, root_id in variants: + self.algebra_stats["exhaustive_generated"] += ( + math.factorial(n) if dedup_permutations else 1 + ) + if dedup_permutations: + self.algebra_stats["exhaustive_after_permutation_dedup"] += 1 try: - yield builder_variant.build(root_id) + dag = builder_variant.build(root_id) except ValueError: if self.debug: print(f"[DEBUG] Skipping invalid DAG for root {root_id}") continue + if self.algebra.exhaustive_canonical: + cid = canonical_id( + dag, self.algebra.laws, self._order_invariant + ) + if cid in seen_canonical: + continue + seen_canonical.add(cid) + self.algebra_stats["exhaustive_after_canonical"] += 1 + if leaf_stats is not None: + count_explicit_prep(self.algebra_stats, dag) + yield dag def _evaluate_dag(self, dag: RepresentationDag, task: Task) -> "OptimizationResult": start_time = time.time() @@ -525,7 +584,9 @@ def optimize( modality_subset, task.model.name ): - for dag in self._generate_fusion_dags(modality_subset, repr_combo): + for dag in self._generate_fusion_dags( + modality_subset, repr_combo, task.model.name + ): if max_combinations and evaluated_count >= max_combinations: break if skip_remaining > 0: diff --git a/src/main/python/systemds/scuro/drsearch/node_executor.py b/src/main/python/systemds/scuro/drsearch/node_executor.py index a498c2752e3..4be1b669af0 100644 --- a/src/main/python/systemds/scuro/drsearch/node_executor.py +++ b/src/main/python/systemds/scuro/drsearch/node_executor.py @@ -163,6 +163,10 @@ def _offload_to_shared_memory(result: Any): return shm_name, resident_bytes, int(shm_bytes or 0), actual_stats +def _measure_memory() -> bool: + return DEBUG or os.environ.get("SCURO_MEASURE_MEMORY", "0") == "1" + + def _execute_node_worker(node, input_mods: List[Any], gpu_id: Optional[int]): start_time = time.perf_counter() if gpu_id is not None: @@ -205,7 +209,7 @@ def _run_node_op(): gpu_peak_bytes = -1 measurement = None - if DEBUG: + if _measure_memory(): input_resident = sum(estimate_modality_bytes(m) for m in input_mods) result, measurement = measure_memory_during( lambda: _run_gpu_op(_run_node_op, gpu_id, cleanup_after=False), @@ -271,7 +275,7 @@ def _run_task(): gpu_peak_bytes = -1 measurement = None - if DEBUG: + if _measure_memory(): result, measurement = measure_memory_during( lambda: _run_gpu_op(_run_task, gpu_id), input_resident_bytes=estimate_modality_bytes(modality), @@ -480,6 +484,7 @@ def __init__( self._cache_misses = 0 self._nodes_executed = 0 self._peak_cache_bytes = 0 + self._leaf_cached_bytes = self.scheduler.memory_stats["cpu_cached"] self._eval_counter = 0 self._nodes_executed = 0 self._search_start = ( @@ -741,6 +746,12 @@ def _process_result(self, jr) -> None: finally: if retained: self._result_cache.release_shm_names(retained) + self._sync_cache_accounting() + + def _sync_cache_accounting(self) -> None: + cache_bytes = self._result_cache.get_memory_total_memory_usage() + self._peak_cache_bytes = max(self._peak_cache_bytes, cache_bytes) + self.scheduler.set_cached_cpu_bytes(self._leaf_cached_bytes + cache_bytes) def _handle_job_failure(self, unit: Union[_NodeUnit, _BatchUnit], jr): reason = jr.error or "unknown worker failure" @@ -860,8 +871,8 @@ def _handle_modality_result( operation_name, actual_stats, measurement=measurement, + estimated_stats=estimated_stats, ) - before_bytes = self._result_cache.get_memory_total_memory_usage() self._manage_result_cache( node_id, transformed_modality, @@ -869,9 +880,7 @@ def _handle_modality_result( resident_bytes=resident_bytes, shm_bytes=shm_bytes, ) - after_bytes = self._result_cache.get_memory_total_memory_usage() - self._peak_cache_bytes = max(self._peak_cache_bytes, after_bytes) - self.scheduler.update_cpu_memory_in_use(after_bytes - before_bytes) + self._sync_cache_accounting() self.scheduler.complete_node(node_id) def _manage_result_cache( @@ -906,9 +915,12 @@ def _checkpoint_memory_usage( operation_name: str, actual_stats: Optional[RepresentationStats], measurement: Optional[MemoryMeasurement] = None, + estimated_stats: Optional[RepresentationStats] = None, ): if self.enable_checkpointing: self._memory_usage_checkpoint.increment(node_id) + if estimated_stats is None: + estimated_stats = self.scheduler.node_stats.get(node_id) if measurement is not None: peak_bytes = measurement.footprint_bytes @@ -962,6 +974,12 @@ def _checkpoint_memory_usage( "rss_delta_bytes": measurement.rss_delta_bytes if measurement else -1, "num_instances": getattr(actual_stats, "num_instances", None), "output_shape": getattr(actual_stats, "output_shape", None), + "estimated_output_shape": getattr( + estimated_stats, "output_shape", None + ), + "estimated_num_instances": getattr( + estimated_stats, "num_instances", None + ), "dtype": str(getattr(actual_stats, "dtype", None)), "container": str(getattr(actual_stats, "container", None)), } @@ -1100,6 +1118,13 @@ def _reuse_report(self) -> Dict[str, Any]: "cpu_fallback_nodes": len( getattr(self.scheduler, "cpu_fallback_nodes", []) or [] ), + "planned_task_nodes": len(getattr(self.scheduler, "roots", ()) or ()), + "completed_task_nodes": self._eval_counter, + "exclusive_memory_nodes": len( + getattr(self.scheduler, "exclusive_memory_nodes", {}) or {} + ), + "deadlock": bool(getattr(self.scheduler, "deadlock", False)), + "deadlock_reason": getattr(self.scheduler, "deadlock_reason", None), } ) return report diff --git a/src/main/python/systemds/scuro/drsearch/node_scheduler.py b/src/main/python/systemds/scuro/drsearch/node_scheduler.py index 78ef5b953fa..b69c50f43cc 100644 --- a/src/main/python/systemds/scuro/drsearch/node_scheduler.py +++ b/src/main/python/systemds/scuro/drsearch/node_scheduler.py @@ -44,6 +44,10 @@ def _exclusive_over_budget_enabled() -> bool: return os.environ.get("SCURO_EXCLUSIVE_OVER_BUDGET", "0") == "1" +def _cache_drain_fraction() -> float: + return float(os.environ.get("SCURO_CACHE_DRAIN_FRACTION", "0.5")) + + class MemoryAwareNodeScheduler: def __init__( @@ -141,16 +145,45 @@ def can_start_now(self, node_id: str) -> bool: def update_cpu_memory_in_use(self, delta_bytes: int): self.memory_stats["cpu_cached"] += delta_bytes + def set_cached_cpu_bytes(self, cached_bytes: int) -> None: + self.memory_stats["cpu_cached"] = cached_bytes + + def _cache_draining(self) -> bool: + budget = self.memory_budget["cpu"] + if _exclusive_over_budget_enabled() and self._total_cpu_in_use() >= budget: + return True + return self.memory_stats["cpu_cached"] >= _cache_drain_fraction() * budget + + def _release_bytes(self, node_id: str) -> int: + release_bytes = 0 + for parent_id in self.parents.get(node_id, set()): + if ( + parent_id not in self.leaves + and self.remaining_children.get(parent_id, 0) == 1 + ): + release_bytes += self.node_resources[parent_id][0] + return release_bytes + + def _grows_cache(self, node_id: str) -> bool: + adds_bytes = ( + self.node_resources[node_id][0] if self.children.get(node_id) else 0 + ) + return adds_bytes > self._release_bytes(node_id) + def get_runnable(self) -> List[RepresentationNode]: runnable_nodes = self._get_runnable_nodes() admitted_bytes = self._pending_admitted_cpu_bytes() pending_gpu = {gpu_id: 0 for gpu_id in self.gpu_devices} pending_slots = {gpu_id: 0 for gpu_id in self.gpu_devices} + draining = self._cache_draining() + hold_growing_nodes = False for node in runnable_nodes: if node in self._ready_set: continue + if hold_growing_nodes and self._grows_cache(node): + continue ok, gpu_id = self._check_memory_constraints( node, admitted_bytes, pending_gpu, pending_slots ) @@ -163,6 +196,13 @@ def get_runnable(self) -> List[RepresentationNode]: self._candidates.discard(node) self.ready_nodes.append(node) self._ready_set.add(node) + elif ( + draining + and not self._grows_cache(node) + and self.node_resources[node][0] + > self.memory_budget["cpu"] - self._total_cpu_in_use() - admitted_bytes + ): + hold_growing_nodes = True chunked_leaf_by_device = defaultdict(list) for node in self.ready_nodes: if isinstance(node, list): @@ -189,28 +229,17 @@ def get_runnable(self) -> List[RepresentationNode]: def _get_runnable_nodes(self) -> List[str]: runnable_nodes = list(self._candidates) - # With exclusive admission a full budget no longer stops the run, so - # the cache of intermediates must be drained deliberately: while the - # CPU budget is exceeded, nodes whose completion does not grow the - # cache (task roots, or last children releasing at least their own - # output) run before nodes that would cache another output. - draining_first = ( - _exclusive_over_budget_enabled() - and self._total_cpu_in_use() >= self.memory_budget["cpu"] - ) + # Upward rank always puts a producer ahead of its consumers, so without + # a counterweight every large intermediate is materialized before the + # tasks that would release it. Once the cache holds a large share of + # the budget (or, with exclusive admission, the budget is exceeded), + # nodes whose completion does not grow the cache (task roots, or last + # children releasing at least their own output) run first. + draining_first = self._cache_draining() def _score(node_id: str): - release_bytes = 0 - for parent_id in self.parents.get(node_id, set()): - if ( - parent_id not in self.leaves - and self.remaining_children.get(parent_id, 0) == 1 - ): - release_bytes += self.node_resources[parent_id][0] - adds_bytes = ( - self.node_resources[node_id][0] if self.children.get(node_id) else 0 - ) - grows_cache = draining_first and adds_bytes > release_bytes + release_bytes = self._release_bytes(node_id) + grows_cache = draining_first and self._grows_cache(node_id) return ( grows_cache, -self.node_priorities.get(node_id, 0.0), diff --git a/src/main/python/systemds/scuro/drsearch/representation_dag.py b/src/main/python/systemds/scuro/drsearch/representation_dag.py index d84d092bbaf..efc4692dcdf 100644 --- a/src/main/python/systemds/scuro/drsearch/representation_dag.py +++ b/src/main/python/systemds/scuro/drsearch/representation_dag.py @@ -36,8 +36,20 @@ DimensionalityReduction, ) from systemds.scuro.utils.identifier import get_op_id, get_node_id +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + AlgebraConfig, + canonical_id, + count_explicit_prep, + count_violations, + explicit_fusion_operands, + explicit_prep_on, + order_invariant, + typecheck, +) +from systemds.scuro.representations.fusion import Fusion import time -from collections import OrderedDict, defaultdict, deque +from collections import Counter, OrderedDict, defaultdict, deque def pushdown_aggregation_for_node( @@ -156,6 +168,9 @@ def pipeline_id(self) -> str: blob = json.dumps(self.to_spec(), sort_keys=True, separators=(",", ":")) return hashlib.sha1(blob.encode()).hexdigest()[:16] + def canonical_id(self, laws=ALL_LAWS, invariant=False) -> str: + return canonical_id(self, laws, invariant) + def _topological_order(self) -> List[str]: node_map = {node.node_id: node for node in self.nodes} visited = [] @@ -670,7 +685,13 @@ def pushdown_aggregation(dag_group: List[RepresentationDag]) -> List[Representat class CSEAwareDAGBuilder: - def __init__(self): + def __init__(self, algebra: Optional[AlgebraConfig] = None, tasks=None, leaves=()): + self.algebra = algebra or AlgebraConfig() + self.leaves = leaves + self.algebra_stats = Counter() + self._order_invariant = order_invariant(self.algebra, tasks) + self._seen_ordered = set() + self._seen_legacy = set() self.global_nodes: List[RepresentationNode] = [] self.signature_to_node: Dict[Hashable, str] = {} self.node_to_signature: Dict[str, Hashable] = {} @@ -685,8 +706,28 @@ def _compute_node_signature( input_sigs = tuple(sorted(ip, key=repr)) if inputs else () op_cls = operation().name params_items = tuple(sorted((parameters or {}).items())) + if self.algebra.cse_law_aware and not self._may_reorder(operation, parameters): + input_sigs = tuple(ip) return ("op", op_cls, params_items, input_sigs) + def _may_reorder(self, operation: Any, parameters: Dict[str, Any]) -> bool: + if not issubclass(operation, Fusion): + return False + laws = operation.algebraic_laws(parameters or {}) + return "comm" in laws or ("concat_perm" in laws and self._order_invariant) + + def _count_law_merges(self, signature: Hashable, inputs: List[str]): + ip = tuple(self.node_to_signature[inp] for inp in inputs) + ordered = signature[:3] + (ip,) + legacy = signature[:3] + (tuple(sorted(ip, key=repr)),) + if signature in self.signature_to_node: + if ordered not in self._seen_ordered: + self.algebra_stats["cse_merges_by_law"] += 1 + elif legacy in self._seen_legacy: + self.algebra_stats["cse_unsafe_merges_prevented"] += 1 + self._seen_ordered.add(ordered) + self._seen_legacy.add(legacy) + def _compute_leaf_signature( self, modality_id: str, representation_index: int = -1 ) -> Hashable: @@ -705,6 +746,8 @@ def _get_or_create_node( signature = self._compute_leaf_signature(modality_id, representation_index) else: signature = self._compute_node_signature(operation, inputs, parameters) + if self.algebra.cse_law_aware: + self._count_law_merges(signature, inputs) try: if signature in self.signature_to_node: @@ -751,6 +794,15 @@ def create_leaf_node( def create_operation_node( self, operation: Any, inputs: List[str], parameters: Dict[str, Any] = None ): + if explicit_prep_on(self.algebra): + inputs = explicit_fusion_operands( + self, + operation, + parameters, + inputs, + self._leaf_stats(), + self.algebra, + ) return self._get_or_create_node( operation=operation, inputs=inputs, parameters=parameters, is_leaf=False ) @@ -769,11 +821,31 @@ def build(self, root_node_id: str, dag_id: int = None) -> RepresentationDag: self.dag_counter += 1 if not dag.validate(): raise ValueError("Invalid DAG construction") + if self.algebra.typecheck != "off": + self._typecheck(dag) + if explicit_prep_on(self.algebra): + count_explicit_prep(self.algebra_stats, dag) if dag_id is None: self._dag_by_root[root_node_id] = dag return dag + def _leaf_stats(self) -> Dict[Hashable, Any]: + return {(m.modality_id, -1): getattr(m, "stats", None) for m in self.leaves} + + def _typecheck(self, dag: RepresentationDag) -> None: + violations = typecheck( + dag, + { + (m.modality_id, -1): getattr(m, "modality_type", None) + for m in self.leaves + }, + self._leaf_stats(), + ) + errors = count_violations(self.algebra_stats, violations) + if errors and self.algebra.typecheck == "strict": + raise ValueError(f"Ill-typed plan: {errors}") + def get_node(self, node_id: str) -> Optional[RepresentationNode]: for node in self.global_nodes: if node.node_id == node_id: diff --git a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py index 97de3da3a58..05b70d03af6 100644 --- a/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py +++ b/src/main/python/systemds/scuro/drsearch/split_aware_evaluator.py @@ -63,9 +63,17 @@ def dag_requires_fitting(dag: RepresentationDag) -> bool: class SplitAwareDAGEvaluator: - def __init__(self, dag: RepresentationDag, modalities, task, node_cache=None): + def __init__( + self, + dag: RepresentationDag, + modalities, + task, + node_cache=None, + share_fitted=False, + ): self.dag = dag self.node_cache = node_cache + self.share_fitted = share_fitted self.modalities = ( list(modalities) if isinstance(modalities, (list, tuple)) else [modalities] ) @@ -231,10 +239,12 @@ def execute_node(node_id): return result input_results = [execute_node(input_id) for input_id in node.inputs] - reusable = ( - self.node_cache is not None - and node_id != self.dag.root_node_id - and not any(input_id in fitted for input_id in node.inputs) + reusable = self.node_cache is not None and ( + self.share_fitted + or ( + node_id != self.dag.root_node_id + and not any(input_id in fitted for input_id in node.inputs) + ) ) if reusable: cached = self.node_cache.get(node_id, partition_key) @@ -254,7 +264,7 @@ def execute_node(node_id): if requirement is not FitRequirement.NONE: fitted.add(node_id) - reusable = False + reusable = reusable and self.share_fitted labels = None fit_input = input_partitions[0] fit_ids = ( diff --git a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py index ced501ba095..8b717e5ced3 100644 --- a/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py +++ b/src/main/python/systemds/scuro/drsearch/unimodal_optimizer.py @@ -19,6 +19,7 @@ # # ------------------------------------------------------------- import copy +from collections import Counter import math import pickle import time @@ -44,6 +45,7 @@ AggregatedRepresentation, ) from systemds.scuro.modality.modality import Modality +from systemds.scuro.drsearch.algebra import AlgebraConfig from systemds.scuro.drsearch.operator_registry import Registry from systemds.scuro.representations.window_aggregation import WindowAggregation from systemds.scuro.utils.checkpointing import CheckpointManager @@ -71,10 +73,13 @@ def __init__( max_num_workers: int = -1, enable_checkpointing: bool = False, window_combination_chains: int = 1, + algebra: Optional[AlgebraConfig] = None, ): self._node_stats: Dict[str, Any] = {} self.pruned = [] self.window_combination_chains = window_combination_chains + self.algebra = algebra or AlgebraConfig() + self.algebra_stats = Counter() self.enable_checkpointing = enable_checkpointing self.modalities = modalities self.tasks = tasks @@ -96,7 +101,8 @@ def __init__( ) # TODO: check if we should make this a local variable (might keep unnecessary memory usage) self.builders = { - modality.modality_id: CSEAwareDAGBuilder() for modality in modalities + modality.modality_id: CSEAwareDAGBuilder(self.algebra, tasks, [modality]) + for modality in modalities } self.debug = debug @@ -392,6 +398,17 @@ def _process_modality(self, modality, skip_remaining: int = 0, scheduler=None): exec_out = node_executor.run() task_results.extend(exec_out["task_results"]) statistics = exec_out["statistics"] + reuse = statistics.get("reuse", {}) + planned = reuse.get("planned_task_nodes", 0) + completed = reuse.get("completed_task_nodes", planned) + if reuse.get("deadlock") or completed < planned: + print( + f"[unimodal] WARNING: modality {modality.modality_id} " + f"evaluated {completed} of {planned} planned task nodes; " + f"deadlock={reuse.get('deadlock')} " + f"reason={reuse.get('deadlock_reason')}", + flush=True, + ) if fitted_dags: fitted_task_dags = self._expand_dags_with_task_roots(fitted_dags) diff --git a/src/main/python/systemds/scuro/modality/transformed.py b/src/main/python/systemds/scuro/modality/transformed.py index 62cc03eec04..87cc3e91434 100644 --- a/src/main/python/systemds/scuro/modality/transformed.py +++ b/src/main/python/systemds/scuro/modality/transformed.py @@ -185,10 +185,13 @@ def combine(self, other: Union[Modality, List[Modality]], fusion_method): fused_modality = TransformedModality( self, fusion_method, ModalityType.EMBEDDING ) + modalities = self.create_modality_list(other) start_time = time.time() - fused_modality.data = fusion_method.transform(self.create_modality_list(other)) + fused_modality.data = fusion_method.transform(modalities) end_time = time.time() - fused_modality.transform_time = end_time - start_time + fused_modality.transform_time += (end_time - start_time) + sum( + getattr(m, "transform_time", 0) or 0 for m in modalities[1:] + ) return fused_modality def create_modality_list(self, other: Union[Modality, List[Modality]]): diff --git a/src/main/python/systemds/scuro/models/model.py b/src/main/python/systemds/scuro/models/model.py index fcd4150cab2..965058eb62c 100644 --- a/src/main/python/systemds/scuro/models/model.py +++ b/src/main/python/systemds/scuro/models/model.py @@ -22,6 +22,7 @@ class Model: uses_gpu = None + feature_order_invariant = None def __init__(self, name: str): """ diff --git a/src/main/python/systemds/scuro/representations/aggregate.py b/src/main/python/systemds/scuro/representations/aggregate.py index 366bb9f3ce9..2c54913445c 100644 --- a/src/main/python/systemds/scuro/representations/aggregate.py +++ b/src/main/python/systemds/scuro/representations/aggregate.py @@ -51,18 +51,22 @@ def _mode_agg(data, aggregate_dim=0): if arr.ndim == 1: values, counts = np.unique(arr, return_counts=True) return values[np.argmax(counts)] - moved = np.moveaxis(arr, axis, -1) - flat = moved.reshape(-1, moved.shape[-1]) + axes = tuple( + a % arr.ndim for a in (axis if isinstance(axis, tuple) else (axis,)) + ) + moved = np.moveaxis(arr, axes, tuple(range(arr.ndim - len(axes), arr.ndim))) + kept = moved.shape[: arr.ndim - len(axes)] + flat = moved.reshape(int(np.prod(kept)), -1) modes = np.empty(flat.shape[0], dtype=arr.dtype) for i in range(flat.shape[0]): values, counts = np.unique(flat[i], return_counts=True) modes[i] = values[np.argmax(counts)] - return modes.reshape(moved.shape[:-1]) + return modes.reshape(kept) @staticmethod def _normalize_axis(aggregate_dim): if isinstance(aggregate_dim, tuple): - return aggregate_dim[0] if aggregate_dim else 0 + return aggregate_dim if aggregate_dim else 0 return aggregate_dim _aggregation_function = { diff --git a/src/main/python/systemds/scuro/representations/average.py b/src/main/python/systemds/scuro/representations/average.py index 06c435d6580..ade8b03eae0 100644 --- a/src/main/python/systemds/scuro/representations/average.py +++ b/src/main/python/systemds/scuro/representations/average.py @@ -32,14 +32,14 @@ @register_fusion_operator() class Average(Fusion): + _laws = frozenset({"comm", "avg_flatten"}) + def __init__(self, params=None): """ Combines modalities using averaging """ super().__init__("Average") self.needs_alignment = True - self.associative = True - self.commutative = True def execute(self, modalities: List[Modality]): data = np.array(modalities[0].data, dtype=np.float64) diff --git a/src/main/python/systemds/scuro/representations/clip.py b/src/main/python/systemds/scuro/representations/clip.py index a64354b1ace..55f841e42e8 100644 --- a/src/main/python/systemds/scuro/representations/clip.py +++ b/src/main/python/systemds/scuro/representations/clip.py @@ -63,12 +63,13 @@ @register_representation([ModalityType.VIDEO, ModalityType.IMAGE]) @register_expensive_representation([ModalityType.VIDEO, ModalityType.IMAGE]) class CLIPVisual(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True def __init__(self, output_file=None, batch_size=32, layer_name="", params=None): parameters = self._get_parameters() - super().__init__("CLIPVisual", ModalityType.EMBEDDING, parameters) + super().__init__("CLIPVisual", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params self._activation_hook = None self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") diff --git a/src/main/python/systemds/scuro/representations/concatenation.py b/src/main/python/systemds/scuro/representations/concatenation.py index cee30382963..f599aba0142 100644 --- a/src/main/python/systemds/scuro/representations/concatenation.py +++ b/src/main/python/systemds/scuro/representations/concatenation.py @@ -34,6 +34,14 @@ @register_fusion_operator() class Concatenation(Fusion): + _laws = frozenset({"assoc", "concat_perm"}) + + @classmethod + def algebraic_laws(cls, parameters: dict) -> frozenset: + if parameters.get("preserve_leading_axis"): + return frozenset({"assoc"}) + return cls._laws + def __init__(self, params=None, preserve_leading_axis=False): super().__init__("Concatenation") if params is not None: diff --git a/src/main/python/systemds/scuro/representations/fusion.py b/src/main/python/systemds/scuro/representations/fusion.py index 5cbe128cae0..95bdec2c1a9 100644 --- a/src/main/python/systemds/scuro/representations/fusion.py +++ b/src/main/python/systemds/scuro/representations/fusion.py @@ -55,22 +55,28 @@ class FusionInputTooLargeError(ValueError): class Fusion(Representation): + _laws = frozenset() + preserves_leading_axis = False + def __init__(self, name, parameters=None): """ Parent class for different multimodal fusion types :param name: Name of the fusion type """ super().__init__(name, parameters) - self.associative = False - self.commutative = False + self.associative = "assoc" in self._laws + self.commutative = "comm" in self._laws self.needs_alignment = False self.needs_training = False self.needs_instance_alignment = False - self.preserves_leading_axis = False self.output_modality_type = ModalityType.EMBEDDING self.data_type = np.float32 self.max_alignment_pad_ratio = DEFAULT_MAX_ALIGNMENT_PAD_RATIO + @classmethod + def algebraic_laws(cls, parameters: dict) -> frozenset: + return cls._laws + @staticmethod def _as_stats_list(input_stats) -> List[RepresentationStats]: if isinstance(input_stats, RepresentationStats): diff --git a/src/main/python/systemds/scuro/representations/glove.py b/src/main/python/systemds/scuro/representations/glove.py index 9513760fb9f..8e9cf0f8a72 100644 --- a/src/main/python/systemds/scuro/representations/glove.py +++ b/src/main/python/systemds/scuro/representations/glove.py @@ -49,8 +49,10 @@ def load_glove_embeddings(file_path): @register_representation(ModalityType.TEXT) @register_expensive_representation(ModalityType.TEXT) class GloVe(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING + def __init__(self, output_file=None, params=None): - super().__init__("GloVe", ModalityType.TEXT) + super().__init__("GloVe", self.OUTPUT_MODALITY_TYPE) file_path = hf_hub_download( repo_id="stanfordnlp/glove", filename="glove.6B.zip" ) diff --git a/src/main/python/systemds/scuro/representations/hadamard.py b/src/main/python/systemds/scuro/representations/hadamard.py index 0697a9f67ea..a88d7d68634 100644 --- a/src/main/python/systemds/scuro/representations/hadamard.py +++ b/src/main/python/systemds/scuro/representations/hadamard.py @@ -33,14 +33,14 @@ @register_fusion_operator() class Hadamard(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities using elementwise multiply (Hadamard product) """ super().__init__("Hadamard") self.needs_alignment = True # zero padding falsifies the result - self.commutative = True - self.associative = True def execute(self, modalities: List[Modality], train_indices=None): fused_data = np.prod([m.data for m in modalities], axis=0) diff --git a/src/main/python/systemds/scuro/representations/lstm.py b/src/main/python/systemds/scuro/representations/lstm.py index 97b318e8bad..46e0842c0c1 100644 --- a/src/main/python/systemds/scuro/representations/lstm.py +++ b/src/main/python/systemds/scuro/representations/lstm.py @@ -42,6 +42,7 @@ @register_fusion_operator() class LSTM(Fusion): fit_requirement = FitRequirement.SUPERVISED + preserves_leading_axis = True def __init__( self, @@ -89,7 +90,6 @@ def __init__( self.device = get_device() self.needs_training = True self.needs_alignment = False - self.preserves_leading_axis = True self.model = None self.input_dim = None self.num_classes = None diff --git a/src/main/python/systemds/scuro/representations/max.py b/src/main/python/systemds/scuro/representations/max.py index 2e0df4c464e..f221a0c882c 100644 --- a/src/main/python/systemds/scuro/representations/max.py +++ b/src/main/python/systemds/scuro/representations/max.py @@ -31,6 +31,8 @@ @register_fusion_operator() class RowMax(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities by computing the outer product of a modality combination and @@ -38,8 +40,6 @@ def __init__(self, params=None): """ super().__init__("RowMax") self.needs_alignment = True - self.associative = True - self.commutative = True def execute( self, @@ -64,3 +64,14 @@ def get_output_stats(self, input_stats_list) -> RepresentationStats: output_shape = max(stats_list, key=self._stats_num_elements).output_shape output_shape_is_known = all(s.output_shape_is_known for s in stats_list) return RepresentationStats(num_instances, output_shape, output_shape_is_known) + + def estimate_peak_memory_bytes(self, input_stats) -> dict: + stats_list = self._as_stats_list(input_stats) + input_bytes = sum(self._stats_bytes(s) for s in stats_list) + output_bytes = self._stats_bytes(self.get_output_stats(input_stats)) + + raw_bytes = self._raw_input_bytes(input_stats) + cpu_peak = ( + int((raw_bytes + 2 * input_bytes + output_bytes) * 1.15) + 8 * 1024 * 1024 + ) + return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0} diff --git a/src/main/python/systemds/scuro/representations/openface.py b/src/main/python/systemds/scuro/representations/openface.py index 60599e806c3..b2d82cc6e91 100644 --- a/src/main/python/systemds/scuro/representations/openface.py +++ b/src/main/python/systemds/scuro/representations/openface.py @@ -37,6 +37,7 @@ from systemds.scuro.representations.representation import ( CONTAINER_LIST, RepresentationStats, + stats_itemsize, ) from systemds.scuro.representations.unimodal import UnimodalRepresentation from systemds.scuro.representations.utils import get_sequence_lengths @@ -170,6 +171,9 @@ class OpenFace(UnimodalRepresentation): FEATURE_DIM = len(FEATURE_COLUMNS) FRAME_BATCH_SIZE = 32 + MAX_SIDE = 1024 + SIZE_WINDOW_CHUNKS = 8 + DETECTOR_MEAN_BGR = (104, 117, 123) FACE_MODEL_MEMORY_BYTES = 8 * 1024 * 1024 MULTITASK_MODEL_MEMORY_BYTES = 128 * 1024 * 1024 @@ -177,6 +181,7 @@ class OpenFace(UnimodalRepresentation): CPU_RUNTIME_OVERHEAD_BYTES = 128 * 1024 * 1024 GPU_RUNTIME_OVERHEAD_BYTES = 64 * 1024 * 1024 LANDMARK_RUNTIME_BYTES = 2400 * 1024 * 1024 + DETECTOR_BYTES_PER_PIXEL = 144 LANDMARK_BYTES_PER_FRAME = 52 * 1024 * 1024 MULTITASK_BYTES_PER_FRAME = 26 * 1024 * 1024 @@ -186,6 +191,7 @@ def __init__( confidence_threshold=0.02, nms_threshold=0.4, vis_threshold=0.5, + max_side=MAX_SIDE, params=None, ): if params is not None: @@ -195,6 +201,7 @@ def __init__( ) nms_threshold = params.get("nms_threshold", nms_threshold) vis_threshold = params.get("vis_threshold", vis_threshold) + max_side = params.get("max_side", max_side) parameters = { "feature_set": list(self.FEATURE_SETS), @@ -207,6 +214,7 @@ def __init__( self.confidence_threshold = float(confidence_threshold) self.nms_threshold = float(nms_threshold) self.vis_threshold = float(vis_threshold) + self.max_side = int(max_side) self.params = params self.data_type = np.float32 self._gpu_id = None @@ -300,9 +308,28 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: max_height = int(getattr(input_stats, "max_height", 224)) max_width = int(getattr(input_stats, "max_width", 224)) max_channels = int(getattr(input_stats, "max_channels", 3)) - frame_bytes = ( - max_height * max_width * max_channels * np.dtype(np.float32).itemsize + float_size = np.dtype(np.float32).itemsize + + if isinstance(input_stats, VideoStats): + video_bytes = ( + int(input_stats.max_length) + * max_height + * max_width + * max_channels + * stats_itemsize(input_stats) + ) + sample_bytes = 2 * video_bytes + 2 * max_height * max_width * 3 * float_size + else: + sample_bytes = 2 * max_height * max_width * max_channels + + canvas_pixels = min(max_height, self.max_side) * min(max_width, self.max_side) + window_bytes = ( + self.FRAME_BATCH_SIZE * self.SIZE_WINDOW_CHUNKS * canvas_pixels * 3 ) + chunk_bytes = window_bytes + ( + self.FRAME_BATCH_SIZE * canvas_pixels * 3 * 2 * float_size + ) + model_bytes = self.FACE_MODEL_MEMORY_BYTES + self.MULTITASK_MODEL_MEMORY_BYTES batch_bytes = self.MULTITASK_BYTES_PER_FRAME runtime_bytes = self.GPU_RUNTIME_OVERHEAD_BYTES @@ -315,13 +342,15 @@ def estimate_peak_memory_bytes(self, input_stats) -> dict: return { "cpu_peak_bytes": int( output_bytes - + frame_bytes + + sample_bytes + + chunk_bytes + model_bytes + self.CPU_RUNTIME_OVERHEAD_BYTES ), "gpu_peak_bytes": int( model_bytes - + self.FRAME_BATCH_SIZE * (frame_bytes + batch_bytes) + + self.FRAME_BATCH_SIZE + * (canvas_pixels * self.DETECTOR_BYTES_PER_PIXEL + batch_bytes) + runtime_bytes ), } @@ -366,15 +395,54 @@ def transform(self, modality, aggregation=None): return transformed_modality def _extract_image_features(self, images): - return list(self._extract_batch_features([self._to_bgr(i) for i in images])) + return list(self._extract_streamed_features(images)) def _extract_video_features(self, frames, owner_id): if len(frames) == 0: raise ValueError(f"Video instance {owner_id} contains no frames") - return self._extract_batch_features([self._to_bgr(frame) for frame in frames]) + return self._extract_streamed_features(frames) def _extract_features(self, image): - return self._extract_batch_features([self._to_bgr(image)])[0] + return self._extract_batch_features([self._prepare(image)])[0] + + def _extract_streamed_features(self, samples): + features = np.zeros((len(samples), self.feature_dim), dtype=np.float32) + window = self.FRAME_BATCH_SIZE * self.SIZE_WINDOW_CHUNKS + for window_start in range(0, len(samples), window): + indices = range(window_start, min(window_start + window, len(samples))) + prepared = {index: self._prepare(samples[index]) for index in indices} + order = sorted(indices, key=lambda index: prepared[index].shape[:2]) + for start, stop in self._chunks(len(order)): + chunk = order[start:stop] + features[chunk] = self._extract_batch_features( + [prepared[index] for index in chunk] + ) + return features + + def _prepare(self, image): + image = self._to_bgr(image) + height, width = image.shape[:2] + longer = max(height, width) + if longer <= self.max_side: + return image + scale = self.max_side / longer + size = (max(1, round(width * scale)), max(1, round(height * scale))) + return cv2.resize(image, size, interpolation=cv2.INTER_AREA) + + def _pad_to(self, image, height, width): + pad_bottom = height - image.shape[0] + pad_right = width - image.shape[1] + if not pad_bottom and not pad_right: + return image + return cv2.copyMakeBorder( + image, + 0, + pad_bottom, + 0, + pad_right, + cv2.BORDER_CONSTANT, + value=self.DETECTOR_MEAN_BGR, + ) def _extract_batch_features(self, images_bgr): self._load_models() @@ -460,53 +528,52 @@ def _detect_faces(self, images_bgr): detections = [None] * len(images_bgr) priors_by_shape = {} - by_shape = {} - for index, image in enumerate(images_bgr): - by_shape.setdefault(image.shape, []).append(index) - - for shape, indices in by_shape.items(): - height, width = shape[:2] - if shape not in priors_by_shape: - priors_by_shape[shape] = ( + for start, stop in self._chunks(len(images_bgr)): + chunk = range(start, stop) + height = max(images_bgr[i].shape[0] for i in chunk) + width = max(images_bgr[i].shape[1] for i in chunk) + if (height, width) not in priors_by_shape: + priors_by_shape[(height, width)] = ( PriorBox(detector.cfg, image_size=(height, width)) .forward() .to(detector.device) .data ) - priors = priors_by_shape[shape] + priors = priors_by_shape[(height, width)] box_scale = torch.Tensor([width, height, width, height]).to(detector.device) landmark_scale = torch.Tensor([width, height] * 5).to(detector.device) - for start, stop in self._chunks(len(indices)): - chunk = indices[start:stop] - batch = torch.cat( - [detector.preprocess_image(images_bgr[i])[0] for i in chunk] + batch = torch.cat( + [ + detector.preprocess_image( + self._pad_to(images_bgr[i], height, width) + )[0] + for i in chunk + ] + ) + with torch.no_grad(): + locations, confidences, landmark_offsets = detector.model(batch) + + for row, index in enumerate(chunk): + boxes = decode(locations.data[row], priors, variance) + boxes = (boxes * box_scale).cpu().numpy() + scores = confidences[row].data.cpu().numpy()[:, 1] + points = decode_landm(landmark_offsets.data[row], priors, variance) + points = (points * landmark_scale).cpu().numpy() + + keep = np.where(scores > detector.confidence_threshold)[0] + boxes, points, scores = boxes[keep], points[keep], scores[keep] + dets = np.hstack((boxes, scores[:, np.newaxis])).astype( + np.float32, copy=False ) - with torch.no_grad(): - locations, confidences, landmark_offsets = detector.model(batch) - - for row, index in enumerate(chunk): - boxes = decode(locations.data[row], priors, variance) - boxes = (boxes * box_scale).cpu().numpy() - scores = confidences[row].data.cpu().numpy()[:, 1] - points = decode_landm(landmark_offsets.data[row], priors, variance) - points = (points * landmark_scale).cpu().numpy() - - keep = np.where(scores > detector.confidence_threshold)[0] - boxes, points, scores = boxes[keep], points[keep], scores[keep] - dets = np.hstack((boxes, scores[:, np.newaxis])).astype( - np.float32, copy=False - ) - keep = py_cpu_nms(dets, detector.nms_threshold) - dets = np.concatenate((dets[keep], points[keep]), axis=1) - - if len(dets) == 0 or dets[0][4] < detector.vis_threshold: - continue - bbox = dets[0][:4].astype(int) - faces[index] = images_bgr[index][ - bbox[1] : bbox[3], bbox[0] : bbox[2] - ] - detections[index] = dets + keep = py_cpu_nms(dets, detector.nms_threshold) + dets = np.concatenate((dets[keep], points[keep]), axis=1) + + if len(dets) == 0 or dets[0][4] < detector.vis_threshold: + continue + bbox = dets[0][:4].astype(int) + faces[index] = images_bgr[index][bbox[1] : bbox[3], bbox[0] : bbox[2]] + detections[index] = dets return faces, detections def _predict_multitask(self, faces, needs_backbone): diff --git a/src/main/python/systemds/scuro/representations/sum.py b/src/main/python/systemds/scuro/representations/sum.py index 58cac2e7210..0b9ce2e399e 100644 --- a/src/main/python/systemds/scuro/representations/sum.py +++ b/src/main/python/systemds/scuro/representations/sum.py @@ -33,6 +33,8 @@ @register_fusion_operator() class Sum(Fusion): + _laws = frozenset({"comm", "assoc"}) + def __init__(self, params=None): """ Combines modalities using colum-wise sum @@ -41,16 +43,13 @@ def __init__(self, params=None): self.needs_alignment = True def execute(self, modalities: List[Modality]): - data = np.array( - modalities[0].data, - dtype=modalities[0].metadata[0]["data_layout"]["type"], - ) - - for m in range(1, len(modalities)): - data += np.asarray( - modalities[m].data, - dtype=modalities[m].metadata[0]["data_layout"]["type"], - ) + arrays = [ + np.asarray(m.data, dtype=m.metadata[0]["data_layout"]["type"]) + for m in modalities + ] + data = np.array(arrays[0], dtype=np.result_type(*arrays)) + for array in arrays[1:]: + data += array return data def get_output_stats(self, input_stats_list) -> RepresentationStats: diff --git a/src/main/python/systemds/scuro/representations/swin_video_transformer.py b/src/main/python/systemds/scuro/representations/swin_video_transformer.py index 5bf4a6d0add..2172e6af4ab 100644 --- a/src/main/python/systemds/scuro/representations/swin_video_transformer.py +++ b/src/main/python/systemds/scuro/representations/swin_video_transformer.py @@ -55,6 +55,7 @@ @register_representation([ModalityType.VIDEO]) @register_expensive_representation([ModalityType.VIDEO]) class SwinVideoTransformer(FrameBudgetedVideoModel, UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING _EMBED_DIM = 768 cache_in_worker = True _MEMORY_PROFILE = VideoMemoryProfile( @@ -80,7 +81,7 @@ def __init__( # "batch_size": [1, 2, 4, 8, 16, 32], } self.data_type = torch.float32 - super().__init__("SwinVideoTransformer", ModalityType.EMBEDDING, parameters) + super().__init__("SwinVideoTransformer", self.OUTPUT_MODALITY_TYPE, parameters) if params is not None: layer_name = params.get("layer_name", layer_name) batch_size = int(params.get("batch_size", batch_size)) diff --git a/src/main/python/systemds/scuro/representations/vgg.py b/src/main/python/systemds/scuro/representations/vgg.py index 6f82937952a..214b56a59b0 100644 --- a/src/main/python/systemds/scuro/representations/vgg.py +++ b/src/main/python/systemds/scuro/representations/vgg.py @@ -56,6 +56,7 @@ def forward(self, input_: torch.Tensor) -> torch.Tensor: @register_representation([ModalityType.IMAGE, ModalityType.VIDEO]) @register_expensive_representation([ModalityType.IMAGE, ModalityType.VIDEO]) class VGG19(UnimodalRepresentation): + OUTPUT_MODALITY_TYPE = ModalityType.EMBEDDING supports_aggregation_pushdown = True cache_in_worker = True @@ -71,7 +72,7 @@ def __init__( self.model = models.vgg19(weights=models.VGG19_Weights.DEFAULT) self.model = self.model.to(self.device) parameters = self._get_parameters() - super().__init__("VGG19", ModalityType.EMBEDDING, parameters) + super().__init__("VGG19", self.OUTPUT_MODALITY_TYPE, parameters) self.params = params if params is not None: batch_size = int((params or {}).get("batch_size", batch_size)) diff --git a/src/main/python/systemds/scuro/representations/window_aggregation.py b/src/main/python/systemds/scuro/representations/window_aggregation.py index e8189e94312..60b5b467ff3 100644 --- a/src/main/python/systemds/scuro/representations/window_aggregation.py +++ b/src/main/python/systemds/scuro/representations/window_aggregation.py @@ -693,10 +693,11 @@ def _feature_shape(self, in_shape): if isinstance(self.aggregation_function, Aggregation): return in_shape[1:] length = in_shape[0] if in_shape else 0 - approx_window_size = ( - max(1, int(length / self._effective_num_windows(length))) if in_shape else 1 - ) - return self._per_window_feature_shape(approx_window_size) + # Window sizes grow geometrically and _pad_stack pads every window's + # features to the largest one, so size them from the largest window + # (about 4x the mean), not the mean. + max_window_size = int(self._window_sizes(length).max()) if in_shape else 1 + return self._per_window_feature_shape(max_window_size) def get_output_stats(self, input_stats: RepresentationStats) -> tuple: in_shape = tuple(int(s) for s in input_stats.output_shape) diff --git a/src/main/python/systemds/scuro/representations/zero_pad.py b/src/main/python/systemds/scuro/representations/zero_pad.py new file mode 100644 index 00000000000..505545e35a3 --- /dev/null +++ b/src/main/python/systemds/scuro/representations/zero_pad.py @@ -0,0 +1,48 @@ +# ------------------------------------------------------------- +# +# 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 systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.representations.fusion import Fusion +from systemds.scuro.representations.representation import derive_stats +from systemds.scuro.representations.unimodal import UnimodalRepresentation + + +class ZeroPad(UnimodalRepresentation): + def __init__(self, width=None, input_width=None, params=None): + super().__init__("ZeroPad", None) + params = params or {} + self.width = params.get("width", width) + self.input_width = params.get("input_width", input_width) + + def transform(self, modality, aggregation=None): + padded = TransformedModality( + modality, self, self_contained=modality.self_contained + ) + padded.data = modality.data + Fusion._normalize_for_fusion(padded) + padded.pad(max_len=self.width) + return padded + + def get_output_stats(self, input_stats): + shape = (*input_stats.output_shape[:-1], self.width) + return derive_stats(input_stats, output_shape=shape) + + def get_current_parameters(self): + return {"width": self.width, "input_width": self.input_width} diff --git a/src/main/python/tests/scuro/test_algebra_canonical.py b/src/main/python/tests/scuro/test_algebra_canonical.py new file mode 100644 index 00000000000..9801d4817a3 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_canonical.py @@ -0,0 +1,271 @@ +# ------------------------------------------------------------- +# +# 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 json +import unittest + +from systemds.scuro.drsearch.algebra import ALL_LAWS, canonical_id +from systemds.scuro.drsearch.representation_dag import RepresentationDAGBuilder +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.bow import BoW +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.max import RowMax +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.zero_pad import ZeroPad + +IMPLICIT = AggregatedRepresentation( + target_dimensions=1, aggregate_leading=True +).get_current_parameters() +FEATURE_AGG = AggregatedRepresentation(target_dimensions=1).get_current_parameters() +NONE = frozenset() + + +def S(*c, op=Sum): + return (op, {}, *c) + + +def A(*c): + return (Average, {}, *c) + + +def C(*c): + return (Concatenation, {"preserve_leading_axis": False}, *c) + + +def W(*c): + return (Concatenation, {"preserve_leading_axis": True}, *c) + + +def L(*c): + return (LSTM, {}, *c) + + +def G(c, params=IMPLICIT): + return (AggregatedRepresentation, params, c) + + +def P(c): + return (ZeroPad, ZeroPad(8, 3).get_current_parameters(), c) + + +def pushed(c, params): + return (BoW, {"_pushdown_aggregation": params}, c) + + +# (name, left, right, laws, invariant, equal) +CASES = [ + ("leaf", "a", "a", ALL_LAWS, False, True), + ("leaf_index", "a", "b", ALL_LAWS, False, False), + ("params", (BoW, {"x": 1}, "a"), (BoW, {"x": 2}, "a"), ALL_LAWS, False, False), + *[ + ( + f"comm_{op.__name__}", + S("a", "b", op=op), + S("b", "a", op=op), + ALL_LAWS, + False, + True, + ) + for op in (Sum, RowMax, Hadamard) + ], + ("comm_off", S("a", "b"), S("b", "a"), NONE, False, False), + ("assoc", S(S("a", "b"), "c"), S("a", S("b", "c")), ALL_LAWS, False, True), + ("assoc_nary", S(S("a", "b"), "c"), S("a", "b", "c"), ALL_LAWS, False, True), + ("assoc_off", S(S("a", "b"), "c"), S("a", "b", "c"), {"comm"}, False, False), + ( + "assoc_other_op", + S(S("a", "b", op=RowMax), "c"), + S("a", "b", "c"), + ALL_LAWS, + False, + False, + ), + ("concat_order_kept", C("a", "b"), C("b", "a"), ALL_LAWS, False, False), + ("concat_perm", C("a", "b"), C("b", "a"), ALL_LAWS, True, True), + ("concat_assoc", C(C("a", "b"), "c"), C("a", C("b", "c")), ALL_LAWS, False, True), + ("window_concat_never_sorted", W("a", "b"), W("b", "a"), ALL_LAWS, True, False), + ( + "window_concat_assoc", + W(W("a", "b"), "c"), + W("a", "b", "c"), + ALL_LAWS, + True, + True, + ), + ("window_concat_not_concat", W("a", "b"), C("a", "b"), ALL_LAWS, True, False), + ("lstm_never_sorted", L("a", "b"), L("b", "a"), ALL_LAWS, True, False), + ( + "no_comm_below_lstm", + L(S("a", "b"), "c"), + L(S("b", "a"), "c"), + ALL_LAWS, + True, + False, + ), + ( + "no_perm_below_lstm", + L(C("a", "b"), "c"), + L(C("b", "a"), "c"), + ALL_LAWS, + True, + False, + ), + ( + "pad_below_lstm", + L(S(P("a"), "b"), "c"), + L(S("a", "b"), "c"), + ALL_LAWS, + False, + True, + ), + ( + "comm_above_lstm", + S(L("a", "b"), "c"), + S("c", L("a", "b")), + ALL_LAWS, + False, + True, + ), + ("avg_flatten", A(A("a", "b"), "c"), A("c", A("b", "a")), ALL_LAWS, False, True), + ("avg_not_assoc", A(A("a", "b"), "c"), A("a", A("b", "c")), ALL_LAWS, False, False), + ( + "avg_weights_add", + A(A("a", "b"), A("a", "c")), + A("a", A("b", "c")), + ALL_LAWS, + False, + True, + ), + ( + "avg_flatten_off", + A(A("a", "b"), A("a", "c")), + A("a", A("b", "c")), + {"comm"}, + False, + False, + ), + ( + "agg_pushdown", + pushed("a", FEATURE_AGG), + G((BoW, {}, "a"), FEATURE_AGG), + ALL_LAWS, + False, + True, + ), + ( + "agg_pushdown_off", + pushed("a", FEATURE_AGG), + G((BoW, {}, "a"), FEATURE_AGG), + NONE, + False, + False, + ), + ("implicit_agg", S(G("a"), "b"), S("a", "b"), ALL_LAWS, False, True), + ( + "implicit_agg_pushed", + S(pushed("a", IMPLICIT), "b"), + S((BoW, {}, "a"), "b"), + ALL_LAWS, + False, + True, + ), + ( + "implicit_agg_off", + S(G("a"), "b"), + S("a", "b"), + ALL_LAWS - {"implicit_agg"}, + False, + False, + ), + ( + "implicit_agg_other_params", + S(G("a", FEATURE_AGG), "b"), + S("a", "b"), + ALL_LAWS, + False, + False, + ), + ("implicit_agg_window_concat", W(G("a"), "b"), W("a", "b"), ALL_LAWS, False, False), + ("implicit_agg_lstm", L(G("a"), "b"), L("a", "b"), ALL_LAWS, False, False), + ("implicit_agg_unary", (BoW, {}, G("a")), (BoW, {}, "a"), ALL_LAWS, False, False), + ("implicit_pad", S(P("a"), "b"), S("a", "b"), ALL_LAWS, False, True), + ("implicit_pad_avg", A(P("a"), "b"), A("a", "b"), ALL_LAWS, False, True), + ( + "implicit_pad_off", + S(P("a"), "b"), + S("a", "b"), + ALL_LAWS - {"implicit_pad"}, + False, + False, + ), + ("implicit_pad_concat", C(P("a"), "b"), C("a", "b"), ALL_LAWS, True, False), +] + + +def build(expr): + builder = RepresentationDAGBuilder() + leaves = {} + + def make(e): + if isinstance(e, str): + if e not in leaves: + i = "abc".index(e) + leaves[e] = builder.create_leaf_node(f"m{i}", i) + return leaves[e] + op, params, *children = e + return builder.create_operation_node(op, [make(c) for c in children], params) + + return builder.build(make(expr)) + + +class TestCanonicalId(unittest.TestCase): + def test_canonical_pairs(self): + for name, left, right, laws, invariant, equal in CASES: + with self.subTest(name): + laws = frozenset(laws) + left, right = build(left), build(right) + ids = [canonical_id(dag, laws, invariant) for dag in (left, right)] + self.assertEqual(ids[0] == ids[1], equal) + for dag, dag_id in zip((left, right), ids): + spec = json.loads(json.dumps(dag.to_spec())) + self.assertEqual(canonical_id(spec, laws, invariant), dag_id) + + def test_dag_method_defaults(self): + dag = build(C("a", S("c", "b"))) + self.assertEqual(dag.canonical_id(), canonical_id(dag, ALL_LAWS, False)) + self.assertNotEqual(dag.canonical_id(laws=NONE), dag.canonical_id()) + + def test_sub_plan_signature(self): + dag = build(C("a", S("b", "c"))) + sub = dag.get_node_by_id(dag.get_node_by_id(dag.root_node_id).inputs[1]) + self.assertEqual( + canonical_id(dag, ALL_LAWS, False, sub.node_id), + canonical_id(build(S("c", "b")), ALL_LAWS, False), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_cse.py b/src/main/python/tests/scuro/test_algebra_cse.py new file mode 100644 index 00000000000..df86a55bd21 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_cse.py @@ -0,0 +1,82 @@ +# ------------------------------------------------------------- +# +# 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 unittest +from types import SimpleNamespace + +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.unimodal_optimizer import UnimodalOptimizer +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.sum import Sum + +LEGACY = AlgebraConfig() +LAW_AWARE = AlgebraConfig(cse_law_aware=True) +FULL = AlgebraConfig.full() + + +def _builder(algebra): + optimizer = UnimodalOptimizer( + [SimpleNamespace(modality_id=0)], + [SimpleNamespace(model=SimpleNamespace(name="task"), expected_dim=1)], + debug=False, + window_combination_chains=2, + algebra=algebra, + ) + builder = optimizer.builders[0] + return optimizer, builder, [builder.create_leaf_node(0, i) for i in range(2)] + + +def _nodes(builder, op): + return [n for n in builder.global_nodes if n.operation is op] + + +class TestLawAwareCSE(unittest.TestCase): + def test_window_concatenation_rotations(self): + for algebra, expected in ((LEGACY, 1), (LAW_AWARE, 2), (FULL, 2)): + with self.subTest(algebra=algebra): + optimizer, builder, windows = _builder(algebra) + optimizer.combine_windowed_representations(builder, windows) + self.assertEqual(len(_nodes(builder, Concatenation)), expected) + prevented = builder.algebra_stats["cse_unsafe_merges_prevented"] + self.assertEqual(prevented, int(algebra.cse_law_aware)) + + def test_commutative_fusion_merges_in_both_modes(self): + cases = [ + (Sum, {}, LEGACY, 1), + (Sum, {}, LAW_AWARE, 1), + (Concatenation, {"preserve_leading_axis": False}, LEGACY, 1), + (Concatenation, {"preserve_leading_axis": False}, LAW_AWARE, 2), + (Concatenation, {"preserve_leading_axis": False}, FULL, 1), + ] + for op, params, algebra, expected in cases: + with self.subTest(op=op.__name__, algebra=algebra): + _, builder, (a, b) = _builder(algebra) + builder.create_operation_node(op, [a, b], params) + builder.create_operation_node(op, [b, a], params) + self.assertEqual(len(_nodes(builder, op)), expected) + merged_by_law = builder.algebra_stats["cse_merges_by_law"] + self.assertEqual( + merged_by_law, int(algebra.cse_law_aware and expected == 1) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_exhaustive.py b/src/main/python/tests/scuro/test_algebra_exhaustive.py new file mode 100644 index 00000000000..249720ed8c2 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_exhaustive.py @@ -0,0 +1,76 @@ +# ------------------------------------------------------------- +# +# 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 tempfile +import unittest + +from systemds.scuro.drsearch.algebra import ALL_LAWS, AlgebraConfig, canonical_id +from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import ( + _FakeModality, + _FakeUnimodalResults, + _make_task, +) + +PERMUTATIONS = AlgebraConfig(exhaustive_dedup_permutations=True) +CANONICAL = AlgebraConfig( + laws=ALL_LAWS, + feature_order_invariant=True, + exhaustive_dedup_permutations=True, + exhaustive_canonical=True, +) + + +def _fusion_dags(algebra, n=3): + ids = [f"m{i}" for i in range(n)] + with tempfile.TemporaryDirectory() as tmp: + optimizer = MultimodalOptimizer( + [_FakeModality(m) for m in ids], + _FakeUnimodalResults({m: [object()] for m in ids}), + [_make_task()], + debug=False, + checkpoint_dir=tmp, + algebra=algebra, + ) + optimizer.fusion_operators = [Concatenation, Sum, Average] + dags = list(optimizer._generate_fusion_dags(ids, {m: 0 for m in ids})) + return dags, optimizer.algebra_stats + + +class TestExhaustiveDedup(unittest.TestCase): + def test_counts(self): + legacy, _ = _fusion_dags(AlgebraConfig()) + ordered, stats = _fusion_dags(PERMUTATIONS) + self.assertEqual((len(legacy), len(ordered)), (648, 108)) + self.assertEqual(stats["exhaustive_generated"], 648) + self.assertEqual(stats["exhaustive_after_permutation_dedup"], 108) + + def test_canonical_matches_brute_force(self): + ordered, _ = _fusion_dags(PERMUTATIONS) + ids = {canonical_id(d, ALL_LAWS, True) for d in ordered} + canonical, stats = _fusion_dags(CANONICAL) + self.assertEqual( + sorted(canonical_id(d, ALL_LAWS, True) for d in canonical), sorted(ids) + ) + self.assertEqual(stats["exhaustive_after_canonical"], len(ids)) diff --git a/src/main/python/tests/scuro/test_algebra_explicit_prep.py b/src/main/python/tests/scuro/test_algebra_explicit_prep.py new file mode 100644 index 00000000000..0f6ff056567 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_explicit_prep.py @@ -0,0 +1,242 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import itertools +import tempfile +import unittest +from collections import Counter +from types import SimpleNamespace + +import numpy as np + +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + AlgebraConfig, + canonical_id, + count_explicit_prep, + sample_stats, + typecheck, +) +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + DagGenome, + MultimodalDeapOptimizer, +) +from systemds.scuro.drsearch.multimodal_optimizer import MultimodalOptimizer +from systemds.scuro.drsearch.representation_dag import CSEAwareDAGBuilder +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.zero_pad import ZeroPad +from tests.scuro.data_generator import ModalityRandomDataGenerator +from tests.scuro.test_multimodal_ga_optimizer import ( + _FakeModality, + _FakeUnimodalResults, + _make_task, +) + +AGG = AlgebraConfig(explicit_fusion_prep=True) +PAD = AlgebraConfig(explicit_padding="pad") +BOTH = AlgebraConfig(explicit_fusion_prep=True, explicit_padding="pad") +FUSIONS = [Concatenation, Sum, Average, Hadamard] +ALIGNED = {Sum, Average, Hadamard} + + +def _reps(): + """Modality 0 is a (4, 8) sequence; 1 and 2 are vectors of width 3 and 8.""" + gen = ModalityRandomDataGenerator() + data, _ = gen.create_2d_modality(6, (4, 8)) + seq = TransformedModality(gen, "test") + seq.data = data + reps = {0: seq} + for mid, width in ((1, 3), (2, 8)): + reps[mid] = gen.create1DModality(6, width, ModalityType.TIMESERIES) + for mid, rep in reps.items(): + rep.modality_id = mid + return {mid: [rep] for mid, rep in reps.items()} + + +def _expected(algebra, op, leaves): + out = [] + if algebra.explicit_fusion_prep and 0 in leaves: + out.append(("AggregatedRepresentation", 0)) + # Padding needs every operand to be a vector first. + if algebra.explicit_padding == "pad" and op in ALIGNED: + if algebra.explicit_fusion_prep or 0 not in leaves: + out.append(("ZeroPad", 1)) + return out + + +def _prep_nodes(dag): + return sorted( + (n.operation.__name__, dag.get_node_by_id(n.inputs[0]).modality_id) + for n in dag.filter_connected_nodes(dag.nodes) + if n.operation in (AggregatedRepresentation, ZeroPad) + ) + + +def _fused(dag, reps): + out = dag.execute([m for ms in reps.values() for m in ms], enable_cache=False) + return np.asarray(out[dag.root_node_id] if isinstance(out, dict) else out.data) + + +def _assert_equivalent(legacy, explicit, reps, expected): + assert _prep_nodes(legacy) == [] + assert _prep_nodes(explicit) == expected + assert canonical_id(legacy, ALL_LAWS, False) == canonical_id( + explicit, ALL_LAWS, False + ) + np.testing.assert_array_equal(_fused(legacy, reps), _fused(explicit, reps)) + + +def _ga_dag(algebra, reps, genome): + optimizer = MultimodalDeapOptimizer( + [_FakeModality(m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + algebra=algebra, + ) + optimizer._current_task_name = "task0" + return optimizer._genome_to_dag(genome) + + +class TestExplicitPrep(unittest.TestCase): + def test_ga_genome_to_dag(self): + cases = itertools.product([AGG, PAD, BOTH], FUSIONS, [(0, 1), (1, 2)]) + for algebra, op, leaves in cases: + with self.subTest(algebra=algebra, op=op.__name__, leaves=leaves): + reps = _reps() + genome = DagGenome([(m, 0) for m in leaves], (0, 1), {"": op}) + _assert_equivalent( + _ga_dag(AlgebraConfig(), reps, genome), + _ga_dag(algebra, reps, genome), + reps, + _expected(algebra, op, leaves), + ) + + def test_exhaustive_enumerator(self): + for algebra in (AGG, BOTH): + with self.subTest(algebra=algebra): + reps = _reps() + runs = [] + for explicit in (AlgebraConfig(), algebra): + with tempfile.TemporaryDirectory() as tmp: + optimizer = MultimodalOptimizer( + [_FakeModality(m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + checkpoint_dir=tmp, + algebra=AlgebraConfig( + exhaustive_dedup_permutations=True, + explicit_fusion_prep=explicit.explicit_fusion_prep, + explicit_padding=explicit.explicit_padding, + ), + ) + optimizer.fusion_operators = FUSIONS + runs.append( + list( + optimizer._generate_fusion_dags( + [0, 1], {0: 0, 1: 0}, "task0" + ) + ) + ) + self.assertEqual(len(runs[0]), 2 * len(FUSIONS)) + self.assertEqual(len(runs[1]), 2 * len(FUSIONS)) + for legacy, explicit, op in zip(*runs, FUSIONS * 2): + _assert_equivalent( + legacy, explicit, reps, _expected(algebra, op, (0, 1)) + ) + self.assertEqual( + optimizer.algebra_stats["explicit_prep_plans"], 2 * len(FUSIONS) + ) + + def test_counts_and_typechecks_padded_plan(self): + reps = _reps() + dag = _ga_dag(BOTH, reps, DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum})) + stats = Counter() + count_explicit_prep(stats, dag) + self.assertEqual( + stats, + { + "explicit_prep_plans": 1, + "explicit_padding_plans": 1, + "explicit_aggregations": 1, + "explicit_paddings": 1, + "explicit_padding_ratio_sum": 8 / 3, + "explicit_padding_ratio_max": 8 / 3, + }, + ) + leaves = {(m, 0): r[0] for m, r in reps.items()} + types = {k: r.modality_type for k, r in leaves.items()} + leaf_stats = {k: sample_stats(r) for k, r in leaves.items()} + self.assertEqual(typecheck(dag, types, leaf_stats), []) + + def test_zero_pad_is_not_registered(self): + registry = Registry() + lists = [ + registry._representations, + registry._context_operators, + registry._dimensionality_reduction_operators, + ] + for ops in (ops for lst in lists for ops in lst.values()): + self.assertNotIn(ZeroPad, ops) + + def test_unbound_genome_to_dag_and_bare_checker(self): + reps = _reps() + genome = DagGenome([(0, 0), (1, 0)], (0, 1), {"": Sum}) + dag = MultimodalDeapOptimizer._genome_to_dag(None, genome) + self.assertEqual(_prep_nodes(dag), []) + checker = object.__new__(MultimodalDeapOptimizer) + checker.max_feature_dim = None + checker.k_best_representations = {"task0": reps} + self.assertIsNone(checker._oversized_reason("task0", genome)) + + def test_cse_builder(self): + for params, wrapped in (({}, True), ({"preserve_leading_axis": True}, False)): + with self.subTest(params=params): + leaf = SimpleNamespace( + modality_id="t", stats=RepresentationStats(5, (4, 8)) + ) + builder = CSEAwareDAGBuilder(BOTH, leaves=[leaf]) + leaf_id = builder.create_leaf_node("t") + root = builder.create_operation_node( + Concatenation, [leaf_id, leaf_id], params + ) + builder.build(root) + self.assertEqual( + builder.get_node(root).inputs != [leaf_id, leaf_id], wrapped + ) + self.assertEqual( + builder.algebra_stats["explicit_aggregations"], 2 * wrapped + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_ga_dedup.py b/src/main/python/tests/scuro/test_algebra_ga_dedup.py new file mode 100644 index 00000000000..8bc457dd530 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_ga_dedup.py @@ -0,0 +1,65 @@ +# ------------------------------------------------------------- +# +# 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 unittest +from unittest.mock import patch + +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.multimodal_ga_optimizer import DagGenome +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import MODULE, _make_optimizer + + +def _body(_dag, _task, _modalities, _objective_specs): + return (0.5,), { + "train_score": {}, + "val_score": {"accuracy": 0.5}, + "test_score": {}, + "runtime": 2.0, + "task_time": 1.0, + "representation_time": 1.0, + } + + +class TestCanonicalGADedup(unittest.TestCase): + def test_commuted_sum_evaluated_once_when_on(self): + leaves = [("m0", 0), ("m1", 0)] + genomes = [DagGenome(leaves, tree, {"": Sum}) for tree in ((0, 1), (1, 0))] + for algebra, evaluations in ( + (AlgebraConfig(), 2), + (AlgebraConfig(laws=frozenset({"comm"}), ga_canonical_dedup=True), 1), + ): + with self.subTest(algebra=algebra): + optimizer, _, task = _make_optimizer( + n_modalities=2, max_workers=1, algebra=algebra + ) + optimizer._current_task_name = task.model.name + population = [optimizer._make_individual(g) for g in genomes] + with patch(f"{MODULE}._evaluate_genome_body", side_effect=_body): + optimizer._evaluate_population(population, task) + results = optimizer.optimization_results["task0"] + self.assertEqual(len(results), evaluations) + avoided = 2 - evaluations + stats = optimizer.algebra_stats + self.assertEqual(stats["ga_evaluations_avoided"], avoided) + self.assertEqual(stats["ga_seconds_avoided"], 2.0 * avoided) + self.assertEqual(stats["ga_unique_legacy"], 2 * avoided) + self.assertEqual(stats["ga_unique_canonical"], avoided) diff --git a/src/main/python/tests/scuro/test_algebra_laws.py b/src/main/python/tests/scuro/test_algebra_laws.py new file mode 100644 index 00000000000..c0d55ddc6b6 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_laws.py @@ -0,0 +1,144 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import itertools +import unittest + +import numpy as np + +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.modality.modality import Modality +from systemds.scuro.modality.transformed import TransformedModality +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.concatenation import Concatenation + +OPERATORS = [(op, {}) for op in Registry().get_fusion_operators()] + [ + (Concatenation, {"preserve_leading_axis": True}) +] +DECLARED = [ + (op, params, law) + for op, params in OPERATORS + for law in sorted(op.algebraic_laws(params)) +] +WIDTHS = {"equal": (6, 6, 6), "unequal": (6, 3, 9)} +DTYPES = { + "float32": (np.float32, np.float32, np.float32), + "float32/64": (np.float32, np.float64, np.float32), + "int/float": (np.int64, np.float32, np.float64), +} + + +def _operands(widths, dtypes, rank): + rng = np.random.default_rng(0) + shape = (5,) + (4,) * (rank - 1) + return [ + (rng.standard_normal(shape + (w,)) * 5).astype(dtype) + for w, dtype in zip(widths, dtypes) + ] + + +def _fuse(op, params, *operands): + modalities = [] + for data in operands: + metadata = [ModalityType.EMBEDDING.create_metadata(d) for d in data] + modality = TransformedModality( + Modality(ModalityType.EMBEDDING, 0, metadata), "operand" + ) + modality.data = data + modalities.append(modality) + return np.asarray(op(params).transform(modalities)) + + +def _assert_same(x, y, exact): + assert x.shape == y.shape + if exact: + np.testing.assert_array_equal(x, y) + else: + np.testing.assert_allclose(x, y, rtol=1e-5, atol=1e-5) + + +def _prepared(x, width): + # What _prepare_modalities hands to execute: hidden mean aggregation, then zero-padding. + x = x.mean(axis=1) if x.ndim == 3 else x + return np.pad(x, ((0, 0), (0, width - x.shape[-1]))) + + +def _comm(f, a, b, c, exact): + _assert_same(f(a, b), f(b, a), exact) + + +def _assoc(f, a, b, c, exact): + left = f(f(a, b), c) + _assert_same(left, f(a, f(b, c)), exact) + _assert_same(left, f(a, b, c), exact) + + +def _avg_flatten(f, a, b, c, exact): + width = max(x.shape[-1] for x in (a, b, c)) + expected = sum(w * _prepared(x, width) for w, x in ((0.25, a), (0.25, b), (0.5, c))) + _assert_same(f(f(a, b), c), expected, exact) + + +def _concat_perm(f, a, b, c, exact): + _assert_same(f(a, b), np.roll(f(b, a), -b.shape[-1], axis=-1), exact=True) + + +LAW_CHECKS = { + "comm": _comm, + "assoc": _assoc, + "avg_flatten": _avg_flatten, + "concat_perm": _concat_perm, +} + + +class TestAlgebraLaws(unittest.TestCase): + def test_every_declared_law_is_checked(self): + self.assertLessEqual({law for _, _, law in DECLARED}, set(LAW_CHECKS)) + + def test_declared_law_holds(self): + cases = itertools.product(DECLARED, WIDTHS, DTYPES, (1, 2)) + for (op, params, law), widths, dtypes, rank in cases: + with self.subTest( + op=f"{op.__name__}{params or ''}", + law=law, + widths=widths, + dtypes=dtypes, + rank=rank, + ): + LAW_CHECKS[law]( + lambda *xs: _fuse(op, params, *xs), + *_operands(WIDTHS[widths], DTYPES[dtypes], rank), + exact=op is Concatenation, + ) + + def test_window_axis_concatenation_is_order_sensitive(self): + params = {"preserve_leading_axis": True} + a, b, _ = _operands(WIDTHS["equal"], DTYPES["float32"], rank=2) + self.assertFalse( + np.array_equal( + _fuse(Concatenation, params, a, b), _fuse(Concatenation, params, b, a) + ) + ) + self.assertFalse({"comm", "concat_perm"} & Concatenation.algebraic_laws(params)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_order_invariant.py b/src/main/python/tests/scuro/test_algebra_order_invariant.py new file mode 100644 index 00000000000..1f6c1264300 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_order_invariant.py @@ -0,0 +1,56 @@ +# ------------------------------------------------------------- +# +# 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 unittest +from types import SimpleNamespace + +from systemds.scuro.drsearch.algebra import AlgebraConfig, order_invariant +from systemds.scuro.models.model import Model + + +def _tasks(*flags): + return [ + SimpleNamespace(model=SimpleNamespace(feature_order_invariant=f)) for f in flags + ] + + +class TestOrderInvariant(unittest.TestCase): + def test_model_default_is_unknown(self): + self.assertIsNone(Model("m").feature_order_invariant) + + def test_cases(self): + cases = [ + (None, _tasks(True, True), True), + (None, _tasks(True, None), False), + (None, _tasks(True, False), False), + (None, [], False), + (None, [SimpleNamespace(model=object())], False), + (True, _tasks(False), True), + (False, _tasks(True), False), + ] + for override, tasks, expected in cases: + with self.subTest(override=override, tasks=tasks): + config = AlgebraConfig(feature_order_invariant=override) + self.assertIs(order_invariant(config, tasks), expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_plan_counts.py b/src/main/python/tests/scuro/test_algebra_plan_counts.py new file mode 100644 index 00000000000..8861d89b766 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_plan_counts.py @@ -0,0 +1,71 @@ +# ------------------------------------------------------------- +# +# 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 unittest + +from systemds.scuro.drsearch.algebra import ( + ALL_LAWS, + _brute_force_canonical, + count_fusion_plans, + count_search_space, +) +from systemds.scuro.representations.average import Average +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.hadamard import Hadamard +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.max import RowMax +from systemds.scuro.representations.sum import Sum + +COMM = frozenset({"comm", "concat_perm"}) +COMM_ASSOC = [Sum, Hadamard, RowMax] +MIXED = [Sum, Average, Concatenation, LSTM] + + +class TestPlanCounts(unittest.TestCase): + def test_closed_forms_match_brute_force(self): + for ops in (COMM_ASSOC, MIXED): + for n in (2, 3, 4): + with self.subTest(ops=[op.__name__ for op in ops], n=n): + counts = count_fusion_plans(n, ops, ALL_LAWS, True) + for key, laws in ( + ("ordered", frozenset()), + ("commutative", COMM), + ("canonical", ALL_LAWS), + ): + brute = _brute_force_canonical(n, ops, laws, True) + self.assertEqual(counts[key], brute, key) + + def test_known_values(self): + counts = count_fusion_plans(3, [Concatenation, Sum, Average], ALL_LAWS, True) + self.assertEqual((counts["legacy_exhaustive"], counts["ordered"]), (648, 108)) + # comm alone divides the ordered trees by 2^(n-1) + self.assertEqual( + count_fusion_plans(4, COMM_ASSOC, COMM, True)["commutative"], 3240 // 8 + ) + self.assertIsNone(count_fusion_plans(5, MIXED, ALL_LAWS, True)["canonical"]) + + def test_search_space_sums_over_subsets(self): + per_size = {n: count_fusion_plans(n, MIXED, ALL_LAWS, False) for n in (2, 3)} + total = count_search_space([2, 2, 3], MIXED, ALL_LAWS, False) + for key, value in total.items(): + self.assertEqual( + value, (4 + 6 + 6) * per_size[2][key] + 12 * per_size[3][key], key + ) diff --git a/src/main/python/tests/scuro/test_algebra_signatures.py b/src/main/python/tests/scuro/test_algebra_signatures.py new file mode 100644 index 00000000000..d18097861a1 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_signatures.py @@ -0,0 +1,90 @@ +# ------------------------------------------------------------- +# +# 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 unittest + +from systemds.scuro.drsearch.algebra import ANY_TYPE, operator_signature +from systemds.scuro.drsearch.operator_registry import Registry, load_builtin_operators +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.aggregated_representation import ( + AggregatedRepresentation, +) +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.lstm import LSTM +from systemds.scuro.representations.mlp_averaging import MLPAveraging +from systemds.scuro.representations.sum import Sum +from systemds.scuro.representations.window_aggregation import WindowAggregation + +load_builtin_operators(ANY_TYPE, include_fusion=True) +REGISTERED = sorted( + { + (op, kind, modality) + for kind, lists in [ + ("encode", Registry._representations), + ("context", Registry._context_operators), + ("reduce", Registry._dimensionality_reduction_operators), + ] + for modality, ops in lists.items() + for op in ops + }, + key=lambda t: (t[0].__name__, t[1], t[2].name), +) + [(op, "fuse", None) for op in Registry().get_fusion_operators()] + +RANK_CASES = [ + (Sum, {}, ("fuse", "n", "vector", "vector")), + ( + Concatenation, + {"preserve_leading_axis": False}, + ("fuse", "n", "vector", "vector"), + ), + (Concatenation, {"preserve_leading_axis": True}, ("fuse", "n", "seq", "seq")), + (LSTM, {}, ("fuse", "n", "seq", "vector")), + (AggregatedRepresentation, {}, ("aggregate", 1, "seq", "vector")), + (WindowAggregation, {}, ("context", 1, "seq", "windows")), + (MLPAveraging, {}, ("reduce", 1, "vector", "vector")), + (None, {}, ("task", 1, "any", "any")), +] + + +class TestOperatorSignatures(unittest.TestCase): + def test_every_registered_operator_has_a_signature(self): + for op, kind, modality in REGISTERED: + with self.subTest( + op=op.__name__, kind=kind, modality=getattr(modality, "name", None) + ): + sig = operator_signature(op) + self.assertEqual(sig.kind, kind) + if modality is not None: + self.assertIn(modality, sig.input_types) + if kind in ("encode", "fuse"): + self.assertIsInstance(sig.output_type, ModalityType) + + def test_rank_rules(self): + for op, params, expected in RANK_CASES: + with self.subTest(op=getattr(op, "__name__", None), params=params): + sig = operator_signature(op, params) + self.assertEqual( + (sig.kind, sig.arity, sig.input_rank, sig.output_rank), expected + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_stats.py b/src/main/python/tests/scuro/test_algebra_stats.py new file mode 100644 index 00000000000..619c686e37b --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_stats.py @@ -0,0 +1,65 @@ +# ------------------------------------------------------------- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +# ------------------------------------------------------------- + +import itertools +import unittest + +import numpy as np + +from systemds.scuro.drsearch.operator_registry import Registry +from systemds.scuro.representations.representation import RepresentationStats +from tests.scuro.test_algebra_laws import _fuse + +# Other test modules register fusion operators of their own. +OPERATORS = [ + op + for op in Registry().get_fusion_operators() + if op.__module__.startswith("systemds.") and not op().needs_training +] +OPERANDS = { + "equal": [(5, 6), (5, 6)], + "unequal": [(5, 6), (5, 3)], + "sequence": [(5, 4, 6), (5, 3)], +} + + +class TestPredictedStats(unittest.TestCase): + def test_predicted_stats_match_execution(self): + for op, case in itertools.product(OPERATORS, OPERANDS): + with self.subTest(op=op.__name__, case=case): + rng = np.random.default_rng(0) + operands = [ + rng.standard_normal(shape).astype(np.float32) + for shape in OPERANDS[case] + ] + stats = [RepresentationStats(len(x), x.shape[1:]) for x in operands] + fused = _fuse(op, {}, *operands) + self.assertEqual( + tuple(op().get_output_stats(stats).output_shape), fused.shape[1:] + ) + self.assertGreater( + op().estimate_peak_memory_bytes(stats)["cpu_peak_bytes"], + fused.nbytes, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_subtree_sharing.py b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py new file mode 100644 index 00000000000..fbaa7f7721c --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_subtree_sharing.py @@ -0,0 +1,127 @@ +# ------------------------------------------------------------- +# +# 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 unittest +from unittest.mock import patch + +import numpy as np + +from systemds.scuro.drsearch import multimodal_ga_optimizer as ga +from systemds.scuro.drsearch.algebra import AlgebraConfig +from systemds.scuro.drsearch.representation_dag import RepresentationDAGBuilder +from systemds.scuro.drsearch.trial_result_cache import TrialResultCache +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.lstm import LSTM +from tests.scuro.data_generator import TestSVM, TestTask +from tests.scuro.test_multimodal_ga_optimizer import ( + _build_real_optimizer, + _make_real_representation, +) + +LSTM_PARAMS = {"width": 8, "epochs": 2} +SPECS = [("accuracy", "max")] + + +class _RecordingSVM(TestSVM): + inputs = [] + + def fit(self, X, y, X_test, y_test): + self.inputs.append(np.array(X)) + return super().fit(X, y, X_test, y_test) + + +def _dag(leaves, root_op=None): + b = RepresentationDAGBuilder() + ids = [b.create_leaf_node(i, 0) for i in leaves] + node = b.create_operation_node(LSTM, ids[:2], dict(LSTM_PARAMS)) + if root_op is not None: + node = b.create_operation_node(root_op, [node, ids[2]], {}) + return b.build(node) + + +def _evaluate(task, modalities, dag, share): + _RecordingSVM.inputs = [] + _, payload = ga._evaluate_genome_body(dag, task, modalities, SPECS, share) + return payload, _RecordingSVM.inputs + + +class TestSubtreeSharing(unittest.TestCase): + def setUp(self): + cache = patch.object(ga, "_SUBTREE_CACHE", TrialResultCache(2**30, 0.0)) + cache.start() + self.addCleanup(cache.stop) + self.task = TestTask("share_task", "share_model", 40) + self.task.model = _RecordingSVM("share_model") + self.modalities = [_make_real_representation(i, 40, 4) for i in range(3)] + _RecordingSVM.inputs = [] + + def test_fitted_subtree_is_reused_bit_identically(self): + task, modalities = self.task, self.modalities + splits = len(task.evaluation_splits) + _evaluate(task, modalities, _dag([0, 1]), True) + shared, shared_inputs = _evaluate( + task, modalities, _dag([0, 1, 2], Concatenation), True + ) + _, fresh_inputs = _evaluate( + task, modalities, _dag([0, 1, 2], Concatenation), False + ) + + self.assertEqual(shared["subtree_cache"]["hits"], splits) + self.assertEqual(len(shared_inputs), splits) + self.assertEqual(len(fresh_inputs), splits) + for a, b in zip(shared_inputs, fresh_inputs): + np.testing.assert_array_equal(a, b) + + def test_operand_order_is_part_of_the_key(self): + _evaluate(self.task, self.modalities, _dag([0, 1]), True) + swapped, _ = _evaluate(self.task, self.modalities, _dag([1, 0]), True) + self.assertEqual(swapped["subtree_cache"]["hits"], 0) + + def test_switch_controls_call_and_stats(self): + payload = { + "train_score": {}, + "val_score": {"accuracy": 0.5}, + "test_score": {}, + "runtime": 1.0, + "task_time": 0.5, + "representation_time": 0.5, + "subtree_cache": { + "hits": 3, + "misses": 2, + "saved_seconds": 4.0, + "bytes": 10, + }, + } + for share in (False, True): + with self.subTest(share=share): + optimizer, task = _build_real_optimizer( + fusion_ops=[LSTM], algebra=AlgebraConfig(share_fused_subtrees=share) + ) + genome = optimizer._random_genome(task.model.name) + with patch.object( + ga, "_evaluate_genome_body", return_value=((0.5,), payload) + ) as body: + optimizer._evaluate_genome(genome, task) + self.assertEqual(len(body.call_args.args), 4 + share) + self.assertEqual(optimizer.algebra_stats["subtree_cache_hits"], 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_algebra_typecheck.py b/src/main/python/tests/scuro/test_algebra_typecheck.py new file mode 100644 index 00000000000..4e34a3be095 --- /dev/null +++ b/src/main/python/tests/scuro/test_algebra_typecheck.py @@ -0,0 +1,133 @@ +# ------------------------------------------------------------- +# +# 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 unittest +from types import SimpleNamespace + +import numpy as np + +from systemds.scuro.drsearch.algebra import AlgebraConfig, typecheck +from systemds.scuro.drsearch.multimodal_ga_optimizer import ( + DagGenome, + MultimodalDeapOptimizer, +) +from systemds.scuro.drsearch.representation_dag import ( + CSEAwareDAGBuilder, + RepresentationDAGBuilder, +) +from systemds.scuro.modality.type import ModalityType +from systemds.scuro.representations.bow import BoW +from systemds.scuro.representations.concatenation import Concatenation +from systemds.scuro.representations.mfcc import MFCC +from systemds.scuro.representations.representation import RepresentationStats +from systemds.scuro.representations.sum import Sum +from tests.scuro.test_multimodal_ga_optimizer import _FakeUnimodalResults, _make_task + +TEXT, EMB = ModalityType.TEXT, ModalityType.EMBEDDING +SEQ, VEC = RepresentationStats(5, (4, 8)), RepresentationStats(5, (3,)) + + +def _check(op, params, leaves, with_stats=True): + builder = RepresentationDAGBuilder() + ids = [builder.create_leaf_node(f"m{i}", 0) for i in range(len(leaves))] + dag = builder.build(builder.create_operation_node(op, ids, params)) + types = {(f"m{i}", 0): t for i, (t, _) in enumerate(leaves)} + stats = {(f"m{i}", 0): s for i, (_, s) in enumerate(leaves)} + violations = typecheck(dag, types, stats if with_stats else None) + return sorted((v.severity, v.kind) for v in violations) + + +VIOLATION_CASES = [ + (BoW, {}, [(TEXT, None)], []), + (MFCC, {}, [(TEXT, None)], [("error", "input_type")]), + (Sum, {}, [(EMB, VEC)], [("error", "arity")]), + (BoW, {}, [(TEXT, None), (TEXT, None)], [("error", "arity")]), + (Sum, {}, [(EMB, VEC), (EMB, VEC)], []), + ( + Sum, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation"), ("implicit", "zero_padding")], + ), + ( + Concatenation, + {}, + [(EMB, SEQ), (EMB, VEC)], + [("implicit", "hidden_aggregation")], + ), + (Concatenation, {"preserve_leading_axis": True}, [(EMB, SEQ), (EMB, SEQ)], []), +] + + +class TestTypecheck(unittest.TestCase): + def test_violations(self): + for i, (op, params, leaves, expected) in enumerate(VIOLATION_CASES): + with self.subTest(case=i, op=op.__name__): + self.assertEqual(_check(op, params, leaves), expected) + + def test_shapes_unknown_without_leaf_stats(self): + violations = _check(Sum, {}, [(EMB, SEQ), (EMB, VEC)], with_stats=False) + self.assertEqual(violations, []) + + def test_builder_modes(self): + for mode, raises, counted in [ + ("off", False, 0), + ("warn", False, 1), + ("strict", True, 1), + ]: + with self.subTest(mode=mode): + modality = SimpleNamespace( + modality_id="t", modality_type=TEXT, stats=None + ) + builder = CSEAwareDAGBuilder( + AlgebraConfig(typecheck=mode), leaves=[modality] + ) + root = builder.create_operation_node( + MFCC, [builder.create_leaf_node("t")], {} + ) + if raises: + with self.assertRaisesRegex(ValueError, "input_type"): + builder.build(root) + else: + builder.build(root) + self.assertEqual(builder.algebra_stats["typecheck_input_type"], counted) + + def test_ga_counts_hidden_operators(self): + reps = { + "m0": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 4, 8)))], + "m1": [SimpleNamespace(modality_type=EMB, data=np.zeros((5, 3)))], + } + optimizer = MultimodalDeapOptimizer( + [SimpleNamespace(modality_id=m) for m in reps], + _FakeUnimodalResults(reps), + [_make_task()], + debug=False, + algebra=AlgebraConfig(typecheck="strict"), + ) + genome = DagGenome([("m0", 0), ("m1", 0)], (0, 1), {"": Sum}) + self.assertIsNone(optimizer._oversized_reason("task0", genome)) + stats = optimizer.algebra_stats + self.assertEqual(stats["typecheck_hidden_aggregation"], 1) + self.assertEqual(stats["typecheck_zero_padding:Sum"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/main/python/tests/scuro/test_data_loaders.py b/src/main/python/tests/scuro/test_data_loaders.py index 0d5ace01a47..4affe2d3b7f 100644 --- a/src/main/python/tests/scuro/test_data_loaders.py +++ b/src/main/python/tests/scuro/test_data_loaders.py @@ -19,6 +19,8 @@ # # ------------------------------------------------------------- +import json +import os import shutil import unittest import numpy as np @@ -26,6 +28,7 @@ from tests.scuro.data_generator import setup_data from systemds.scuro.dataloader.audio_loader import AudioLoader, AudioStats from systemds.scuro.dataloader.image_loader import ImageLoader, ImageStats +from systemds.scuro.dataloader.json_loader import JSONLoader from systemds.scuro.dataloader.text_loader import TextLoader, TextStats from systemds.scuro.dataloader.video_loader import VideoLoader, VideoStats from systemds.scuro.modality.type import ModalityType @@ -183,5 +186,36 @@ def test_image_loader_stats(self): # self.assertEqual(stats.num_signals, 2) +class TestJSONLoaderFields(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.test_file_path = "test_json_fields_data" + os.makedirs(cls.test_file_path, exist_ok=True) + cls.indices = ["0", "1"] + cls.json_path = os.path.join(cls.test_file_path, "dialogue.json") + with open(cls.json_path, "w") as f: + json.dump( + { + index: {"utterance": f"utt {index}", "context": ["a", "b"]} + for index in cls.indices + }, + f, + ) + + @classmethod + def tearDownClass(cls): + shutil.rmtree(cls.test_file_path) + + def test_single_field(self): + data, _ = JSONLoader(self.json_path, self.indices, "utterance").load() + self.assertEqual(data, ["utt 0", "utt 1"]) + + def test_list_of_fields_is_joined_in_order(self): + loader = JSONLoader(self.json_path, self.indices, ["context", "utterance"]) + data, _ = loader.load() + self.assertEqual(data, ["a b utt 0", "a b utt 1"]) + self.assertEqual(loader.stats.max_length, len("a b utt 0")) + + if __name__ == "__main__": unittest.main() 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 f79157f0247..e5eac1c5cfb 100644 --- a/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py +++ b/src/main/python/tests/scuro/test_multimodal_ga_optimizer.py @@ -36,6 +36,7 @@ _collect_leaf_indices, _failure_fitness, _objective_value, + _task_uses_gpu, ) from systemds.scuro.drsearch.operator_registry import Registry, register_fusion_operator from systemds.scuro.modality.type import ModalityType @@ -110,10 +111,20 @@ 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): + def __init__( + self, + n_workers, + dispatch, + ctx=None, + threads_per_worker=1, + gpu_devices=None, + gpu_slots_per_device=1, + gpu_demand_fraction=1.0, + ): self.n_workers = n_workers self.dispatch = dispatch self.threads_per_worker = threads_per_worker + self.gpu_devices = list(gpu_devices or []) self.pending = None self.next_job_id = 0 self.shutdown_called = False @@ -123,7 +134,18 @@ def __init__(self, n_workers, dispatch, ctx=None, threads_per_worker=1): def has_idle_worker(self): return self.pending is None - def submit(self, kind, payload, gpu_id=None): + @property + def has_idle_cpu_worker(self): + return self.pending is None + + @property + def gpu_worker_devices(self): + return [] + + def has_idle_worker_for(self, gpu_id, allow_gpu_worker_for_cpu=False): + return self.pending is None + + def submit(self, kind, payload, gpu_id=None, allow_gpu_worker_for_cpu=False): job_id = self.next_job_id self.next_job_id += 1 try: @@ -1216,5 +1238,53 @@ def test_real_fusion_multi_objective_end_to_end(self): self.assertGreaterEqual(result.runtime, 0.0) +class _LanePool: + """Idle-worker bookkeeping of PersistentWorkerPool, without processes.""" + + def __init__(self, idle_gpu, idle_cpu): + self.idle_gpu = set(idle_gpu) + self.idle_cpu = idle_cpu + self.gpu_worker_devices = [0, 1] + + @property + def has_idle_cpu_worker(self): + return self.idle_cpu + + def has_idle_worker_for(self, gpu_id, allow_gpu_worker_for_cpu=False): + if gpu_id is not None: + return gpu_id in self.idle_gpu + return self.idle_cpu or (allow_gpu_worker_for_cpu and bool(self.idle_gpu)) + + +class TestGpuLanes(unittest.TestCase): + def _lane(self, pool, uses_gpu): + return MultimodalDeapOptimizer._idle_lane( + SimpleNamespace(_worker_pool=pool), uses_gpu + ) + + def test_gpu_task_prefers_idle_gpu_worker(self): + self.assertEqual(self._lane(_LanePool({1}, idle_cpu=True), True), 1) + + def test_gpu_task_overflows_to_cpu_worker(self): + self.assertIsNone(self._lane(_LanePool(set(), idle_cpu=True), True)) + + def test_cpu_task_may_take_gpu_worker(self): + self.assertIsNone(self._lane(_LanePool({0}, idle_cpu=False), False)) + + def test_no_idle_worker(self): + self.assertIs(self._lane(_LanePool(set(), idle_cpu=False), True), False) + self.assertIs(self._lane(_LanePool(set(), idle_cpu=False), False), False) + + def test_task_uses_gpu_from_flag_or_device(self): + def task(**model_attrs): + return SimpleNamespace(model=SimpleNamespace(**model_attrs)) + + self.assertTrue(_task_uses_gpu(task(uses_gpu=True))) + self.assertFalse(_task_uses_gpu(task(uses_gpu=False, device="cuda"))) + self.assertTrue(_task_uses_gpu(task(device=SimpleNamespace(type="cuda")))) + self.assertFalse(_task_uses_gpu(task(device=SimpleNamespace(type="cpu")))) + self.assertFalse(_task_uses_gpu(task())) + + if __name__ == "__main__": unittest.main() diff --git a/src/main/python/tests/scuro/test_unimodal_optimizer.py b/src/main/python/tests/scuro/test_unimodal_optimizer.py index f27c721aa25..6e002c1688e 100644 --- a/src/main/python/tests/scuro/test_unimodal_optimizer.py +++ b/src/main/python/tests/scuro/test_unimodal_optimizer.py @@ -47,6 +47,7 @@ from systemds.scuro.representations.bow import BoW from systemds.scuro.representations.bert import Bert from systemds.scuro.representations.word2vec import W2V +from systemds.scuro.representations.window_aggregation import Window from tests.scuro.test_unimodal_representations import ( PHYSIOLOGICAL_REPRESENTATIONS, TIMESERIES_REPRESENTATIONS, @@ -162,6 +163,43 @@ def test_robust_results_ignore_non_finite_scores(self): self.assertEqual(indices, [0, 5]) self.assertTrue(all(np.isfinite(entry.robustness_score) for entry in robust)) + def test_robust_results_penalise_fold_variation(self): + modality = SimpleNamespace(modality_id="modality") + task = SimpleNamespace(model=SimpleNamespace(name="task")) + results = UnimodalResults([modality], [task], k=2) + entries = [ + ResultEntry( + val_score={"accuracy": 0.90}, + val_fold_scores={"accuracy": [0.70, 0.80, 0.90, 1.00, 1.10]}, + dag=SimpleNamespace( + nodes=[SimpleNamespace(operation=UnimodalOptimizer)] + ), + ), + ResultEntry( + val_score={"accuracy": 0.86}, + val_fold_scores={"accuracy": [0.85, 0.86, 0.87, 0.85, 0.87]}, + dag=SimpleNamespace( + nodes=[SimpleNamespace(operation=UnimodalOptimizer)] + ), + ), + ] + results.results[modality.modality_id][task.model.name] = entries + + robust, indices = results.get_k_most_robust_results( + modality, task, "accuracy", one_se_parsimony=False + ) + + self.assertEqual(robust, [entries[1], entries[0]]) + self.assertEqual(indices, [1, 0]) + self.assertAlmostEqual( + entries[0].robustness_score, + 0.90 - np.std(entries[0].val_fold_scores["accuracy"], ddof=1), + ) + self.assertAlmostEqual( + entries[1].robustness_score, + 0.86 - np.std(entries[1].val_fold_scores["accuracy"], ddof=1), + ) + def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): text_data, text_md = ModalityRandomDataGenerator().create_text_data( self.num_instances, 10 @@ -194,6 +232,84 @@ def test_bow_and_tfidf_require_dimensionality_reduction_before_task(self): task_input = dag.get_node_by_id(task_node.inputs[0]) self.assertIs(task_input.operation, MLPAveraging) + def _window_inputs(self, modality, representations): + """(input operation, input rank) of every window node in the search.""" + with patch.object( + Registry, + "_representations", + registry_for(modality.modality_type, representations), + ): + optimizer = UnimodalOptimizer( + [modality], self.tasks, False, enable_checkpointing=False + ) + dags, _, _ = optimizer._build_execution_dags_for_modality(modality) + stats = {} + for dag in dags: + stats.update(optimizer._node_output_stats(dag)) + inputs = [] + for dag in dags: + for node in dag.nodes: + if node.operation is None or not issubclass(node.operation, Window): + continue + source = dag.get_node_by_id(node.inputs[0]) + inputs.append( + ( + source.operation, + len(stats[source.node_id].output_shape), + ) + ) + return optimizer, inputs + + def test_windows_are_not_applied_to_rank_one_representation_outputs(self): + text_data, text_md = ModalityRandomDataGenerator().create_text_data( + self.num_instances, 10 + ) + text = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.TEXT, text_data, str, text_md + ) + ) + audio_data, audio_md = ModalityRandomDataGenerator().create_audio_data( + self.num_instances, 4000 + ) + audio = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.AUDIO, audio_data, np.float32, audio_md + ) + ) + # W2V pools to one vector per instance; Concatenation flattens the + # (frames, features) outputs of Spectral and ZeroCrossing to rank 1. + for modality, representations in ( + (text, [W2V]), + (audio, [Spectral, ZeroCrossing]), + ): + with self.subTest(modality=modality.modality_type.name): + optimizer, inputs = self._window_inputs(modality, representations) + representation_inputs = [ + (operation, rank) + for operation, rank in inputs + if operation is not None + ] + self.assertEqual(representation_inputs, []) + self.assertTrue( + any("no time axis" in p["reason"] for p in optimizer.pruned) + ) + + def test_windows_are_kept_for_representations_with_a_time_axis(self): + video_data, video_md = ModalityRandomDataGenerator().create_visual_modality( + self.num_instances, 10, 10, 10 + ) + video = UnimodalModality( + TestDataLoader( + self.indices, None, ModalityType.VIDEO, video_data, np.float32, video_md + ) + ) + _, inputs = self._window_inputs(video, [ResNet]) + self.assertGreater(len(inputs), 0) + for operation, rank in inputs: + self.assertIs(operation, ResNet) + self.assertGreaterEqual(rank, 2) + def test_unimodal_optimizer_for_image_modality(self): image_data, image_md = ModalityRandomDataGenerator().create_visual_modality( self.num_instances, 1, 10, 10 @@ -453,3 +569,71 @@ def optimize_unimodal_representation_for_modality(self, modalities): modalities[0], self.tasks[0], "accuracy" ) assert len(result) == 1 + + +class TestSelectionStrategies(unittest.TestCase): + """The three selection strategies must genuinely disagree. + + They are only worth carrying separately through HPO and fusion if they can + pick different winners out of the same pool, so the fixture is built so + each strategy has a different answer: A wins on raw score, B on robustness + (low fold variation, small DAG), and B again once runtime counts, with A + demoted for being slow. + """ + + @staticmethod + def _entry(accuracy, std, runtime, n_nodes): + dag = SimpleNamespace( + nodes=[SimpleNamespace(operation=object()) for _ in range(n_nodes)] + ) + return SimpleNamespace( + dag=dag, + val_score={"accuracy": accuracy}, + val_fold_scores={"accuracy": [accuracy - std, accuracy + std]}, + val_score_stds={"accuracy": std}, + representation_time=runtime, + task_time=0.0, + ) + + def setUp(self): + self.modality = SimpleNamespace(modality_id="m1") + self.task = SimpleNamespace(model=SimpleNamespace(name="t1")) + self.entries = { + "noisy_slow": self._entry(0.90, 0.10, 100.0, 5), + "stable_fast": self._entry(0.88, 0.01, 10.0, 2), + "weak_fastest": self._entry(0.85, 0.05, 1.0, 1), + } + self.results = UnimodalResults([self.modality], [self.task], k=3) + self.results.results["m1"]["t1"] = list(self.entries.values()) + + def _winner(self, strategy, **kwargs): + view = self.results.with_selection_strategy(strategy, **kwargs) + ranked, _ = view.get_k_best_results( + self.modality, self.task, "accuracy", cache_needed=False + ) + return ranked[0] + + def test_performance_takes_the_validation_argmax(self): + self.assertIs(self._winner("performance"), self.entries["noisy_slow"]) + + def test_robustness_penalises_fold_variation(self): + self.assertIs(self._winner("robustness"), self.entries["stable_fast"]) + + def test_tradeoff_weighs_runtime_against_accuracy(self): + self.assertIs(self._winner("tradeoff"), self.entries["stable_fast"]) + + def test_tradeoff_with_all_weight_on_accuracy_matches_performance(self): + self.assertIs( + self._winner("tradeoff", tradeoff_weights=(1.0, 0.0)), + self.entries["noisy_slow"], + ) + + def test_views_share_the_candidate_pool_and_leave_the_source_alone(self): + view = self.results.with_selection_strategy("robustness") + self.assertIs(view.results, self.results.results) + self.assertIs(view.cache, self.results.cache) + self.assertEqual(self.results.selection_strategy, "performance") + + def test_unknown_strategy_is_rejected(self): + with self.assertRaises(ValueError): + self.results.with_selection_strategy("cheapest")