diff --git a/plugins/nemo-optimization/src/nemo_optimization/artifact_utils.py b/plugins/nemo-optimization/src/nemo_optimization/artifact_utils.py new file mode 100644 index 0000000000..5e0ae99720 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/artifact_utils.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared artifact helpers for optimizer backends.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +_SECRET_VALUE_KEYS = frozenset( + { + "api_key", + "apikey", + "password", + "passwd", + "secret", + "token", + "authorization", + "access_token", + "refresh_token", + "client_secret", + "nvidia_api_key", + } +) + + +def sanitize_config_for_artifact(config: Mapping[str, Any]) -> dict[str, Any]: + """Return a deep copy with secret-bearing fields redacted for persistent YAML/JSON.""" + + def _redact(value: Any, *, key: str | None = None) -> Any: + if isinstance(value, Mapping): + return {str(k): _redact(v, key=str(k)) for k, v in value.items()} + if isinstance(value, list): + return [_redact(v, key=key) for v in value] + if key is not None and isinstance(value, str) and value and not value.startswith("${"): + if _is_secret_value_key(key): + return "${REDACTED}" + return value + + return _redact(config) + + +def _is_secret_value_key(key: str) -> bool: + lowered = key.lower().replace("-", "_") + # Reference fields hold env/secret *names*, not credentials. + if lowered.endswith("_env") or lowered in {"api_key_secret", "api_key_env"}: + return False + if lowered in _SECRET_VALUE_KEYS: + return True + return any(lowered.endswith(f"_{suffix}") for suffix in ("api_key", "password", "token", "secret")) + + +__all__ = ["sanitize_config_for_artifact"] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/backend.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/backend.py index 2e4e70215a..fb6744ee9e 100644 --- a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/backend.py +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/backend.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Prompt GA backend stub.""" +"""Prompt GA optimize backend.""" from __future__ import annotations @@ -12,6 +12,22 @@ from nemo_platform import NeMoPlatform from nemo_platform_plugin.job_context import JobContext +from nemo_optimization.atif_metadata import resolve_experiment_id +from nemo_optimization.backends.ga.config import ( + GaConfigError, + GaPromptOptimizerConfig, + parse_ga_prompt_optimizer_config, +) +from nemo_optimization.backends.ga.driver import ( + GaPromptOptimizationResult, + GaPromptOptimizerError, + run_ga_prompt_optimization, +) +from nemo_optimization.backends.ga.transform import ( + ModelPromptTransformer, + PromptTransformer, + PromptTransformError, +) from nemo_optimization.backends.protocol import ( OptimizationBackendCapabilities, OptimizationPhase, @@ -19,13 +35,15 @@ OptimizationPhaseResult, OptimizationPhaseStatus, ) -from nemo_optimization.search_space import parse_prompt_optimizer_config +from nemo_optimization.candidate import CandidateEvaluationError, CandidateEvaluator +from nemo_optimization.config import generate_optimize_id +from nemo_optimization.fabric_evaluator import FabricCandidateEvaluator RESULT_NAME = "optimizer_results" class GaBackendError(RuntimeError): - """Raised when prompt GA is requested before the backend ships.""" + """Raised when prompt GA backend usage is invalid.""" class GaBackend: @@ -58,28 +76,160 @@ def run_phase( del sdk if request.phase is not OptimizationPhase.PROMPT: raise GaBackendError(f"GA backend does not support the {request.phase.value!r} phase.") - parse_prompt_optimizer_config(request.payload) - message = ( - "optimizer.prompt.enabled is not supported yet. " - "Prompt GA is tracked separately and will be implemented in the GA algorithm stack." - ) + payload = request.payload output_dir = ctx.storage.persistent / "results" / RESULT_NAME - output_dir.mkdir(parents=True, exist_ok=True) - failure = { - "status": OptimizationPhaseStatus.FAILED.value, - "backend": self.name, - "phase": OptimizationPhase.PROMPT.value, - "error": message, - } - (output_dir / "prompt_phase_failure.json").write_text(json.dumps(failure, indent=2) + "\n", encoding="utf-8") + + try: + config = parse_ga_prompt_optimizer_config(payload) + except GaConfigError: + raise + except KeyError as exc: + raise GaConfigError(f"payload optimizer section is missing required key: {exc}") from exc + + experiment_id = request.experiment_id or resolve_experiment_id(payload, generate_id=generate_optimize_id) + try: + evaluator = _build_prompt_evaluator( + payload, + metric_names=config.metric_names, + output_dir=output_dir, + experiment_id=experiment_id, + ) + transformer = _build_prompt_transformer(payload, model_name=config.model) + result = run_ga_prompt_optimization( + payload, + output_dir, + evaluator, + transformer, + config=config, + trial_number_offset=request.trial_number_offset, + ) + except (CandidateEvaluationError, GaPromptOptimizerError, PromptTransformError) as exc: + return _failed_prompt_phase_result( + payload, + backend=self.name, + output_dir=output_dir, + ctx=ctx, + experiment_id=experiment_id, + error=str(exc), + optimized_payload=getattr(exc, "optimized_payload", None), + trial_count=getattr(exc, "trial_count", 0), + trial_number_offset=request.trial_number_offset, + ) + + summary = _phase_summary(result, experiment_id=experiment_id, payload=payload, config=config) + summary_path = output_dir / "prompt_phase_summary.json" + summary_path.write_text( + json.dumps( + { + "status": OptimizationPhaseStatus.COMPLETED.value, + "backend": self.name, + "phase": OptimizationPhase.PROMPT.value, + **summary, + }, + indent=2, + default=str, + ) + + "\n", + encoding="utf-8", + ) ref = ctx.results.save(RESULT_NAME, output_dir) return OptimizationPhaseResult( phase=OptimizationPhase.PROMPT, backend=self.name, - status=OptimizationPhaseStatus.FAILED, - optimized_payload=copy.deepcopy(request.payload), - summary={"error": message}, + status=OptimizationPhaseStatus.COMPLETED, + optimized_payload=result.optimized_payload, + summary=summary, artifacts={"result": ref.model_dump(mode="json")}, - trial_count=0, + trial_count=result.executed_trials, trial_number_offset=request.trial_number_offset, ) + + +def _phase_summary( + result: GaPromptOptimizationResult, + *, + experiment_id: str, + payload: dict[str, Any], + config: GaPromptOptimizerConfig, +) -> dict[str, Any]: + best = result.best_individual + return { + "experiment_id": experiment_id, + "population_size": config.population_size, + "generations": config.generations, + "generations_completed": result.generations_completed, + "executed_trials": result.executed_trials, + "best_individual": best.individual_id, + "best_phase_trial_number": best.phase_trial_number, + "best_global_trial_number": best.global_trial_number, + "best_prompts": dict(best.prompts), + "best_metrics": dict(best.aggregate_metrics), + "best_fitness": best.fitness, + "metric_names": list(result.metric_names), + "agent": payload.get("metadata", {}).get("name"), + } + + +def _failed_prompt_phase_result( + payload: dict[str, Any], + *, + backend: str, + output_dir, + ctx: JobContext, + experiment_id: str, + error: str, + optimized_payload: dict[str, Any] | None, + trial_count: int, + trial_number_offset: int, +) -> OptimizationPhaseResult: + output_dir.mkdir(parents=True, exist_ok=True) + failure = { + "status": OptimizationPhaseStatus.FAILED.value, + "backend": backend, + "phase": OptimizationPhase.PROMPT.value, + "experiment_id": experiment_id, + "error": error, + "executed_trials": trial_count, + } + failure_path = output_dir / "prompt_phase_failure.json" + if failure_path.is_file(): + try: + existing = json.loads(failure_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + existing = {} + if isinstance(existing, dict): + failure = {**existing, **failure} + failure_path.write_text(json.dumps(failure, indent=2) + "\n", encoding="utf-8") + ref = ctx.results.save(RESULT_NAME, output_dir) + return OptimizationPhaseResult( + phase=OptimizationPhase.PROMPT, + backend=backend, + status=OptimizationPhaseStatus.FAILED, + optimized_payload=copy.deepcopy(optimized_payload if optimized_payload is not None else payload), + summary={"experiment_id": experiment_id, "error": error, "executed_trials": trial_count}, + artifacts={"result": ref.model_dump(mode="json")}, + trial_count=trial_count, + trial_number_offset=trial_number_offset, + ) + + +def _build_prompt_evaluator( + payload: dict[str, Any], + *, + metric_names: tuple[str, ...], + output_dir, + experiment_id: str, +) -> CandidateEvaluator: + if isinstance(payload.get("eval"), dict): + return FabricCandidateEvaluator( + payload=payload, + metric_names=metric_names, + output_dir=output_dir, + experiment_id=experiment_id, + ) + + raise GaConfigError("Prompt GA optimization requires payload.eval; prompt-only means no numeric phase.") + + +def _build_prompt_transformer(payload: dict[str, Any], *, model_name: str) -> PromptTransformer: + return ModelPromptTransformer(payload=payload, model_name=model_name) diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/config.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/config.py new file mode 100644 index 0000000000..d612a4b966 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/config.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration parsing for prompt GA optimization.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import Any + +from nemo_optimization.search_space import PromptSearchSpaceSpec, SearchSpaceError, parse_prompt_optimizer_config + +DEFAULT_POPULATION_SIZE = 10 +DEFAULT_GENERATIONS = 5 +DEFAULT_CROSSOVER_RATE = 0.7 +DEFAULT_MUTATION_RATE = 0.1 +DEFAULT_SELECTION_METHOD = "tournament" +DEFAULT_TOURNAMENT_SIZE = 3 +DEFAULT_PARALLEL_EVALUATIONS = 8 +DEFAULT_MULTI_OBJECTIVE_MODE = "harmonic" +DEFAULT_ORACLE_FEEDBACK_MODE = "never" +DEFAULT_ORACLE_FEEDBACK_WORST_N = 5 +DEFAULT_ORACLE_FEEDBACK_MAX_CHARS = 4000 +DEFAULT_ORACLE_FEEDBACK_FITNESS_THRESHOLD = 0.3 +DEFAULT_ORACLE_FEEDBACK_STAGNATION_GENERATIONS = 3 +DEFAULT_ORACLE_FEEDBACK_VARIANCE_THRESHOLD = 0.01 +DEFAULT_ORACLE_FEEDBACK_DIVERSITY_THRESHOLD = 0.5 + +SELECTION_METHODS = frozenset({"tournament", "roulette"}) +MULTI_OBJECTIVE_MODES = frozenset({"harmonic", "weighted_sum", "chebyshev"}) +ORACLE_FEEDBACK_MODES = frozenset({"never", "always", "failing_only", "adaptive"}) + + +class GaConfigError(ValueError): + """Raised when prompt GA configuration is invalid.""" + + +class MetricDirection(str, Enum): + """Optimization direction for one metric.""" + + MAXIMIZE = "maximize" + MINIMIZE = "minimize" + + +@dataclass(frozen=True) +class MetricSpec: + """One metric objective consumed by the GA fitness reducer.""" + + name: str + direction: MetricDirection + weight: float = 1.0 + + +@dataclass(frozen=True) +class GaPromptOptimizerConfig: + """Validated prompt GA optimizer settings.""" + + backend: str + model: str + search_space: dict[str, PromptSearchSpaceSpec] + metrics: tuple[MetricSpec, ...] + population_size: int + generations: int + crossover_rate: float + mutation_rate: float + elitism: int + selection_method: str + tournament_size: int + diversity_lambda: float + parallel_evaluations: int + seed: int | None + target: float | None + multi_objective_mode: str + reps_per_param_set: int + oracle_feedback_mode: str + oracle_feedback_worst_n: int + oracle_feedback_max_chars: int + oracle_feedback_fitness_threshold: float + oracle_feedback_stagnation_generations: int + oracle_feedback_fitness_variance_threshold: float + oracle_feedback_diversity_threshold: float + + @property + def metric_names(self) -> tuple[str, ...]: + return tuple(metric.name for metric in self.metrics) + + +def parse_ga_prompt_optimizer_config(payload: Mapping[str, Any]) -> GaPromptOptimizerConfig: + """Parse prompt GA settings from a Fabric optimize payload.""" + + try: + base = parse_prompt_optimizer_config(payload) + except SearchSpaceError as exc: + raise GaConfigError(str(exc)) from exc + + optimizer = payload.get("optimizer") + if not isinstance(optimizer, Mapping): + raise GaConfigError("payload.optimizer must be a mapping.") + prompt = optimizer.get("prompt") + if not isinstance(prompt, Mapping): + raise GaConfigError("optimizer.prompt must be a mapping.") + + population_size = _positive_int( + prompt, + "population_size", + "ga_population_size", + default=DEFAULT_POPULATION_SIZE, + ) + generations = _positive_int(prompt, "generations", "ga_generations", default=DEFAULT_GENERATIONS) + elitism = _non_negative_int( + prompt, + "elitism", + "ga_elitism", + default=min(2, max(0, population_size - 1)), + ) + if elitism >= population_size and population_size > 1: + raise GaConfigError("optimizer.prompt.elitism must be less than population_size.") + if population_size == 1 and elitism: + raise GaConfigError("optimizer.prompt.elitism must be 0 when population_size is 1.") + + selection_method = _choice( + prompt, + "selection_method", + "ga_selection_method", + choices=SELECTION_METHODS, + default=DEFAULT_SELECTION_METHOD, + ) + multi_objective_mode = _choice( + optimizer, + "multi_objective_combination_mode", + choices=MULTI_OBJECTIVE_MODES, + default=DEFAULT_MULTI_OBJECTIVE_MODE, + normalize=_normalize_mode, + ) + oracle_feedback_mode = _choice( + prompt, + "oracle_feedback_mode", + choices=ORACLE_FEEDBACK_MODES, + default=DEFAULT_ORACLE_FEEDBACK_MODE, + normalize=_normalize_mode, + ) + + return GaPromptOptimizerConfig( + backend=base.backend, + model=base.model, + search_space=base.search_space, + metrics=_parse_metrics(optimizer), + population_size=population_size, + generations=generations, + crossover_rate=_rate(prompt, "crossover_rate", "ga_crossover_rate", default=DEFAULT_CROSSOVER_RATE), + mutation_rate=_rate(prompt, "mutation_rate", "ga_mutation_rate", default=DEFAULT_MUTATION_RATE), + elitism=elitism, + selection_method=selection_method, + tournament_size=_positive_int( + prompt, + "tournament_size", + "ga_tournament_size", + default=DEFAULT_TOURNAMENT_SIZE, + ), + diversity_lambda=_non_negative_float(prompt, "diversity_lambda", default=0.0), + parallel_evaluations=_positive_int( + prompt, + "parallel_evaluations", + "ga_parallel_evaluations", + default=DEFAULT_PARALLEL_EVALUATIONS, + ), + seed=_optional_int(prompt, "seed"), + target=_optional_float(optimizer, "target"), + multi_objective_mode=multi_objective_mode, + reps_per_param_set=_positive_int(optimizer, "reps_per_param_set", default=1), + oracle_feedback_mode=oracle_feedback_mode, + oracle_feedback_worst_n=_positive_int( + prompt, + "oracle_feedback_worst_n", + default=DEFAULT_ORACLE_FEEDBACK_WORST_N, + ), + oracle_feedback_max_chars=_positive_int( + prompt, + "oracle_feedback_max_chars", + default=DEFAULT_ORACLE_FEEDBACK_MAX_CHARS, + ), + oracle_feedback_fitness_threshold=_non_negative_float( + prompt, + "oracle_feedback_fitness_threshold", + default=DEFAULT_ORACLE_FEEDBACK_FITNESS_THRESHOLD, + ), + oracle_feedback_stagnation_generations=_positive_int( + prompt, + "oracle_feedback_stagnation_generations", + default=DEFAULT_ORACLE_FEEDBACK_STAGNATION_GENERATIONS, + ), + oracle_feedback_fitness_variance_threshold=_non_negative_float( + prompt, + "oracle_feedback_fitness_variance_threshold", + default=DEFAULT_ORACLE_FEEDBACK_VARIANCE_THRESHOLD, + ), + oracle_feedback_diversity_threshold=_non_negative_float( + prompt, + "oracle_feedback_diversity_threshold", + default=DEFAULT_ORACLE_FEEDBACK_DIVERSITY_THRESHOLD, + ), + ) + + +def _parse_metrics(optimizer: Mapping[str, Any]) -> tuple[MetricSpec, ...]: + eval_metrics = optimizer.get("eval_metrics") + if not isinstance(eval_metrics, Mapping) or not eval_metrics: + raise GaConfigError("optimizer.eval_metrics must declare at least one metric.") + + metrics: list[MetricSpec] = [] + for name, raw in eval_metrics.items(): + if not isinstance(name, str): + raise GaConfigError("optimizer.eval_metrics keys must be strings.") + if not isinstance(raw, Mapping): + raise GaConfigError(f"optimizer.eval_metrics[{name!r}] must be a mapping.") + direction_raw = raw.get("direction", MetricDirection.MAXIMIZE.value) + if not isinstance(direction_raw, str): + raise GaConfigError(f"Metric {name!r} direction must be 'maximize' or 'minimize'.") + direction = _normalize_mode(direction_raw) + if direction not in {MetricDirection.MAXIMIZE.value, MetricDirection.MINIMIZE.value}: + raise GaConfigError(f"Metric {name!r} direction must be 'maximize' or 'minimize'.") + metric_name = raw.get("evaluator_name") or name + if not isinstance(metric_name, str) or not metric_name.strip(): + raise GaConfigError(f"Metric {name!r} evaluator_name must be a non-empty string.") + weight = _float_value(raw.get("weight", 1.0), path=f"optimizer.eval_metrics[{name!r}].weight") + if weight <= 0: + raise GaConfigError(f"Metric {name!r} weight must be greater than 0.") + metrics.append( + MetricSpec( + name=metric_name.strip(), + direction=MetricDirection(direction), + weight=weight, + ) + ) + return tuple(metrics) + + +def _choice( + mapping: Mapping[str, Any], + *keys: str, + choices: frozenset[str], + default: str, + normalize: Callable[[str], str] | None = None, +) -> str: + value = _first_present(mapping, *keys, default=default) + if not isinstance(value, str): + raise GaConfigError(f"{_path(keys)} must be a string.") + normalized = (normalize or _normalize_choice)(value) + if normalized not in choices: + supported = ", ".join(sorted(choices)) + raise GaConfigError(f"{_path(keys)} has unsupported value {value!r}; supported values: {supported}.") + return normalized + + +def _rate(mapping: Mapping[str, Any], *keys: str, default: float) -> float: + value = _float_value(_first_present(mapping, *keys, default=default), path=_path(keys)) + if value < 0 or value > 1: + raise GaConfigError(f"{_path(keys)} must be between 0 and 1.") + return value + + +def _positive_int(mapping: Mapping[str, Any], *keys: str, default: int) -> int: + value = _int_value(_first_present(mapping, *keys, default=default), path=_path(keys)) + if value < 1: + raise GaConfigError(f"{_path(keys)} must be greater than 0.") + return value + + +def _non_negative_int(mapping: Mapping[str, Any], *keys: str, default: int) -> int: + value = _int_value(_first_present(mapping, *keys, default=default), path=_path(keys)) + if value < 0: + raise GaConfigError(f"{_path(keys)} must be greater than or equal to 0.") + return value + + +def _non_negative_float(mapping: Mapping[str, Any], *keys: str, default: float) -> float: + value = _float_value(_first_present(mapping, *keys, default=default), path=_path(keys)) + if value < 0: + raise GaConfigError(f"{_path(keys)} must be greater than or equal to 0.") + return value + + +def _optional_int(mapping: Mapping[str, Any], *keys: str) -> int | None: + value = _first_present(mapping, *keys, default=None) + if value is None: + return None + return _int_value(value, path=_path(keys)) + + +def _optional_float(mapping: Mapping[str, Any], *keys: str) -> float | None: + value = _first_present(mapping, *keys, default=None) + if value is None: + return None + return _float_value(value, path=_path(keys)) + + +def _first_present(mapping: Mapping[str, Any], *keys: str, default: Any) -> Any: + for key in keys: + if key in mapping: + return mapping[key] + return default + + +def _int_value(value: Any, *, path: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise GaConfigError(f"{path} must be an integer.") + return value + + +def _float_value(value: Any, *, path: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise GaConfigError(f"{path} must be a number.") + return float(value) + + +def _normalize_mode(value: str) -> str: + return value.strip().lower().replace("-", "_") + + +def _normalize_choice(value: str) -> str: + return value.strip().lower() + + +def _path(keys: tuple[str, ...]) -> str: + if len(keys) == 1: + return keys[0] + return "/".join(keys) + + +__all__ = [ + "GaConfigError", + "GaPromptOptimizerConfig", + "MetricDirection", + "MetricSpec", + "parse_ga_prompt_optimizer_config", +] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/driver.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/driver.py new file mode 100644 index 0000000000..2a45997976 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/driver.py @@ -0,0 +1,777 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prompt genetic-algorithm optimizer.""" + +from __future__ import annotations + +import copy +import csv +import json +import logging +import random +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + +from nemo_optimization.artifact_utils import sanitize_config_for_artifact +from nemo_optimization.backends.ga.config import GaPromptOptimizerConfig, parse_ga_prompt_optimizer_config +from nemo_optimization.backends.ga.fitness import ( + FitnessSnapshot, + GaFitnessError, + assign_generation_fitness, + best_individual, + rank_valid_individuals, + target_met, +) +from nemo_optimization.backends.ga.individual import GaIndividual +from nemo_optimization.backends.ga.oracle_feedback import ( + OracleFeedbackState, + build_oracle_feedback, + should_use_oracle_feedback, +) +from nemo_optimization.backends.ga.transform import PromptTransformer, PromptTransformError +from nemo_optimization.candidate import CandidateEvaluationError, CandidateEvaluationResult, CandidateEvaluator +from nemo_optimization.config_overlay import apply_suggestions +from nemo_optimization.search_space import PromptSearchSpaceSpec, suggestions_by_path + +logger = logging.getLogger(__name__) + + +class GaPromptOptimizerError(RuntimeError): + """Raised when prompt GA optimization cannot produce a valid candidate.""" + + def __init__( + self, + message: str, + *, + optimized_payload: dict[str, Any] | None = None, + trial_count: int = 0, + ) -> None: + super().__init__(message) + self.optimized_payload = optimized_payload + self.trial_count = trial_count + + +@dataclass(frozen=True) +class GaPromptOptimizationResult: + """Completed prompt GA run details.""" + + optimized_payload: dict[str, Any] + best_individual: GaIndividual + metric_names: tuple[str, ...] + executed_trials: int + generations_completed: int + history: tuple[GaIndividual, ...] + output_dir: Path + + +def run_ga_prompt_optimization( + payload: Mapping[str, Any], + output_dir: Path, + evaluator: CandidateEvaluator, + transformer: PromptTransformer, + *, + config: GaPromptOptimizerConfig | None = None, + trial_number_offset: int = 0, +) -> GaPromptOptimizationResult: + """Execute prompt GA against the prompt dimensions in ``optimizer.search_space``.""" + + base_payload = copy.deepcopy(dict(payload)) + ga_config = config or parse_ga_prompt_optimizer_config(base_payload) + output_dir.mkdir(parents=True, exist_ok=True) + + rng = random.Random(ga_config.seed) + next_phase_trial_number = 0 + population = _initial_population(base_payload, config=ga_config, transformer=transformer, rng=rng) + history: list[GaIndividual] = [] + best_so_far: GaIndividual | None = None + stagnation_generations = 0 + generations_completed = 0 + + for generation in range(ga_config.generations): + logger.info("Evaluating prompt GA generation %d/%d", generation + 1, ga_config.generations) + next_phase_trial_number = _evaluate_population( + base_payload, + population, + config=ga_config, + evaluator=evaluator, + output_dir=output_dir, + next_phase_trial_number=next_phase_trial_number, + trial_number_offset=trial_number_offset, + ) + + try: + snapshot = assign_generation_fitness( + population, + metrics=ga_config.metrics, + mode=ga_config.multi_objective_mode, + diversity_lambda=ga_config.diversity_lambda, + ) + except GaFitnessError as exc: + optimized_payload = _optimized_payload_for_best(base_payload, best_so_far, ga_config) + _write_history_artifacts(output_dir, history=(*history, *population)) + _write_failure_artifact( + output_dir, + config=ga_config, + message=str(exc), + trial_count=next_phase_trial_number, + best=best_so_far, + ) + raise GaPromptOptimizerError( + str(exc), + optimized_payload=optimized_payload, + trial_count=next_phase_trial_number, + ) from exc + + history.extend(copy.deepcopy(population)) + previous_best_id = best_so_far.individual_id if best_so_far is not None else None + best_so_far = best_individual( + history, + metrics=ga_config.metrics, + mode=ga_config.multi_objective_mode, + ) + if best_so_far is None: + message = "Prompt GA completed a generation but no valid best individual was available." + _write_failure_artifact( + output_dir, + config=ga_config, + message=message, + trial_count=next_phase_trial_number, + best=None, + ) + raise GaPromptOptimizerError( + message, + optimized_payload=copy.deepcopy(base_payload), + trial_count=next_phase_trial_number, + ) + + if previous_best_id is not None and best_so_far.individual_id == previous_best_id: + stagnation_generations += 1 + else: + stagnation_generations = 0 + generations_completed = generation + 1 + _write_generation_artifacts( + output_dir, + generation=generation, + population=population, + snapshot=snapshot, + best=best_so_far, + config=ga_config, + ) + + if target_met(best_so_far, metrics=ga_config.metrics, target=ga_config.target): + logger.info("Prompt GA stopped early after meeting target") + break + if generation == ga_config.generations - 1: + break + + population = _next_generation( + population, + generation=generation + 1, + config=ga_config, + transformer=transformer, + rng=rng, + oracle_state=OracleFeedbackState( + stagnation_generations=stagnation_generations, + fitness_variance=snapshot.fitness_variance, + duplicate_ratio=snapshot.duplicate_ratio, + ), + ) + + if best_so_far is None: + message = "Prompt GA did not produce a valid individual." + _write_failure_artifact(output_dir, config=ga_config, message=message, trial_count=next_phase_trial_number) + raise GaPromptOptimizerError( + message, + optimized_payload=copy.deepcopy(base_payload), + trial_count=next_phase_trial_number, + ) + + optimized_payload = _optimized_payload_for_best(base_payload, best_so_far, ga_config) + _write_final_artifacts( + output_dir, + base_payload=base_payload, + optimized_payload=optimized_payload, + best=best_so_far, + config=ga_config, + history=history, + executed_trials=next_phase_trial_number, + generations_completed=generations_completed, + ) + return GaPromptOptimizationResult( + optimized_payload=optimized_payload, + best_individual=copy.deepcopy(best_so_far), + metric_names=ga_config.metric_names, + executed_trials=next_phase_trial_number, + generations_completed=generations_completed, + history=tuple(history), + output_dir=output_dir, + ) + + +def _initial_population( + payload: Mapping[str, Any], + *, + config: GaPromptOptimizerConfig, + transformer: PromptTransformer, + rng: random.Random, +) -> list[GaIndividual]: + prompts = {name: _initial_prompt(payload, prompt_spec) for name, prompt_spec in sorted(config.search_space.items())} + population = [GaIndividual(prompts=copy.deepcopy(prompts), generation=0, individual_index=0)] + for individual_index in range(1, config.population_size): + mutated_prompts: dict[str, str] = {} + failures: list[str] = [] + for name, spec in sorted(config.search_space.items()): + prompt = prompts[name] + mutated_prompts[name] = _mutate_prompt( + transformer, + prompt_name=name, + prompt=prompt, + spec=spec, + feedback=None, + failures=failures, + rng=rng, + ) + population.append( + GaIndividual( + prompts=mutated_prompts, + generation=0, + individual_index=individual_index, + parent_ids=(population[0].individual_id,), + transform_failures=failures, + ) + ) + return population + + +def _next_generation( + population: Sequence[GaIndividual], + *, + generation: int, + config: GaPromptOptimizerConfig, + transformer: PromptTransformer, + rng: random.Random, + oracle_state: OracleFeedbackState, +) -> list[GaIndividual]: + parents = rank_valid_individuals(population) + if not parents: + raise GaPromptOptimizerError("Cannot create next generation without a valid parent.") + + next_population: list[GaIndividual] = [ + elite.clone_as_elite(generation=generation, individual_index=index) + for index, elite in enumerate(parents[: config.elitism]) + ] + while len(next_population) < config.population_size: + parent_a = _select_parent(parents, config=config, rng=rng) + parent_b = _select_parent(parents, config=config, rng=rng) + if parent_b.individual_id == parent_a.individual_id and len(parents) > 1: + alternatives = [parent for parent in parents if parent.individual_id != parent_a.individual_id] + parent_b = rng.choice(alternatives) + child_prompts: dict[str, str] = {} + failures: list[str] = [] + for name, spec in sorted(config.search_space.items()): + prompt = _child_prompt( + transformer, + name=name, + spec=spec, + parent_a=parent_a, + parent_b=parent_b, + config=config, + rng=rng, + oracle_state=oracle_state, + failures=failures, + ) + child_prompts[name] = prompt + next_population.append( + GaIndividual( + prompts=child_prompts, + generation=generation, + individual_index=len(next_population), + parent_ids=(parent_a.individual_id, parent_b.individual_id), + transform_failures=failures, + ) + ) + return next_population + + +def _child_prompt( + transformer: PromptTransformer, + *, + name: str, + spec: PromptSearchSpaceSpec, + parent_a: GaIndividual, + parent_b: GaIndividual, + config: GaPromptOptimizerConfig, + rng: random.Random, + oracle_state: OracleFeedbackState, + failures: list[str], +) -> str: + parent_prompt = parent_a.prompts[name] + prompt_sources = (parent_a,) + if rng.random() < config.crossover_rate and parent_a.individual_id != parent_b.individual_id: + feedback = _feedback_for_sources((parent_a, parent_b), config=config, state=oracle_state) + try: + parent_prompt = transformer.recombine( + prompt_name=name, + parent_a=parent_a.prompts[name], + parent_b=parent_b.prompts[name], + purpose=spec.purpose, + prompt_format=spec.format, + feedback=feedback, + ) + prompt_sources = (parent_a, parent_b) + except PromptTransformError as exc: + fallback_parent = rng.choice([parent_a, parent_b]) + parent_prompt = fallback_parent.prompts[name] + prompt_sources = (fallback_parent,) + failures.append( + f"recombine:{name}:{parent_a.individual_id},{parent_b.individual_id}:" + f"{exc}; used {fallback_parent.individual_id}" + ) + elif rng.random() < 0.5: + parent_prompt = parent_b.prompts[name] + prompt_sources = (parent_b,) + + if rng.random() < config.mutation_rate: + feedback = _feedback_for_sources(prompt_sources, config=config, state=oracle_state) + parent_prompt = _mutate_prompt( + transformer, + prompt_name=name, + prompt=parent_prompt, + spec=spec, + feedback=feedback, + failures=failures, + rng=rng, + ) + return parent_prompt + + +def _mutate_prompt( + transformer: PromptTransformer, + *, + prompt_name: str, + prompt: str, + spec: PromptSearchSpaceSpec, + feedback: str | None, + failures: list[str], + rng: random.Random, +) -> str: + del rng + try: + return transformer.mutate( + prompt_name=prompt_name, + prompt=prompt, + purpose=spec.purpose, + prompt_format=spec.format, + feedback=feedback, + ) + except PromptTransformError as exc: + failures.append(f"mutate:{prompt_name}:{exc}; retained parent prompt") + return prompt + + +def _select_parent( + parents: Sequence[GaIndividual], + *, + config: GaPromptOptimizerConfig, + rng: random.Random, +) -> GaIndividual: + if config.selection_method == "tournament": + contenders = rng.sample(list(parents), k=min(config.tournament_size, len(parents))) + return rank_valid_individuals(contenders)[0] + + positive_weights = [max(0.0, parent.fitness or 0.0) for parent in parents] + total = sum(positive_weights) + if total <= 0: + return rng.choice(list(parents)) + needle = rng.random() * total + cumulative = 0.0 + for parent, weight in zip(parents, positive_weights, strict=True): + cumulative += weight + if cumulative >= needle: + return parent + return parents[-1] + + +def _evaluate_population( + payload: Mapping[str, Any], + population: Sequence[GaIndividual], + *, + config: GaPromptOptimizerConfig, + evaluator: CandidateEvaluator, + output_dir: Path, + next_phase_trial_number: int, + trial_number_offset: int, +) -> int: + pending = [ + individual + for individual in population + if not (individual.status == "completed" and individual.aggregate_metrics) + ] + for individual in pending: + individual.phase_trial_number = next_phase_trial_number + individual.global_trial_number = trial_number_offset + next_phase_trial_number + next_phase_trial_number += 1 + + if len(pending) <= 1 or config.parallel_evaluations <= 1: + for individual in pending: + _evaluate_assigned_individual( + payload, + individual, + config=config, + evaluator=evaluator, + output_dir=output_dir, + ) + return next_phase_trial_number + + max_workers = min(config.parallel_evaluations, len(pending)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + _evaluate_assigned_individual, + payload, + individual, + config=config, + evaluator=evaluator, + output_dir=output_dir, + ) + for individual in pending + ] + for future in futures: + future.result() + return next_phase_trial_number + + +def _evaluate_assigned_individual( + payload: Mapping[str, Any], + individual: GaIndividual, + *, + config: GaPromptOptimizerConfig, + evaluator: CandidateEvaluator, + output_dir: Path, +) -> None: + if individual.global_trial_number is None: + raise GaPromptOptimizerError("Prompt GA individual was evaluated before trial numbering.") + path_suggestions = suggestions_by_path(config.search_space, individual.prompts) + write_candidate_config( + output_dir, + trial_number=individual.global_trial_number, + trial_config=apply_suggestions(payload, path_suggestions), + ) + try: + rep_results = [ + evaluator.evaluate( + trial_number=individual.global_trial_number, + suggestions=dict(path_suggestions), + trial_overlay=_trial_overlay(individual), + rep=rep, + ) + for rep in range(config.reps_per_param_set) + ] + individual.aggregate_metrics = _average_rep_metrics(rep_results, config.metric_names) + individual.raw_scores = tuple(score for result in rep_results for score in result.scores) + individual.status = "completed" + individual.failure_reason = None + except CandidateEvaluationError as exc: + individual.status = "failed" + individual.failure_reason = str(exc) + individual.aggregate_metrics = {} + individual.raw_scores = () + individual.fitness = None + individual.normalized_metrics = {} + + +def _average_rep_metrics( + rep_results: Sequence[CandidateEvaluationResult], + metric_names: Sequence[str], +) -> dict[str, float]: + if not rep_results: + raise CandidateEvaluationError("Cannot average prompt GA metrics from zero repetitions.") + averaged: dict[str, float] = {} + for metric_name in metric_names: + values: list[float] = [] + for rep_index, result in enumerate(rep_results): + if metric_name not in result.aggregate_metrics: + raise CandidateEvaluationError(f"Prompt GA rep {rep_index} missing metric {metric_name!r}.") + values.append(float(result.aggregate_metrics[metric_name])) + averaged[metric_name] = sum(values) / len(values) + return averaged + + +def _trial_overlay(individual: GaIndividual) -> dict[str, Any]: + return { + "metadata": { + "nemo.optimizer.phase": "prompt", + "nemo.optimizer.phase_trial_number": individual.phase_trial_number, + "nemo.optimizer.global_trial_number": individual.global_trial_number, + "nemo.optimizer.generation": individual.generation, + "nemo.optimizer.individual_index": individual.individual_index, + } + } + + +def _feedback_for_sources( + sources: Sequence[GaIndividual], + *, + config: GaPromptOptimizerConfig, + state: OracleFeedbackState, +) -> str | None: + sections: list[str] = [] + seen: set[str] = set() + for source in sources: + if source.individual_id in seen: + continue + seen.add(source.individual_id) + if not should_use_oracle_feedback(config=config, parent=source, state=state): + continue + feedback = build_oracle_feedback(individual=source, config=config) + if feedback: + sections.append(f"Feedback from {source.individual_id}:\n{feedback}") + if not sections: + return None + return "\n\n".join(sections) + + +def _initial_prompt(payload: Mapping[str, Any], spec: PromptSearchSpaceSpec) -> str: + cursor: Any = payload + for segment in spec.path.split("."): + if not isinstance(cursor, Mapping) or segment not in cursor: + raise GaPromptOptimizerError(f"Prompt path {spec.path!r} did not resolve during GA initialization.") + cursor = cursor[segment] + if not isinstance(cursor, str): + raise GaPromptOptimizerError(f"Prompt path {spec.path!r} must resolve to a string.") + return cursor + + +def _optimized_payload_for_best( + payload: Mapping[str, Any], + best: GaIndividual | None, + config: GaPromptOptimizerConfig, +) -> dict[str, Any]: + if best is None: + return copy.deepcopy(dict(payload)) + return apply_suggestions( + payload, + suggestions_by_path(config.search_space, best.prompts), + strip_optimizer=False, + ) + + +def write_candidate_config( + output_dir: Path, + *, + trial_number: int, + trial_config: Mapping[str, Any], +) -> Path: + path = output_dir / f"config_prompt_trial_{trial_number:03d}.yml" + path.write_text( + yaml.safe_dump(sanitize_config_for_artifact(trial_config), sort_keys=False), + encoding="utf-8", + ) + return path + + +def _write_generation_artifacts( + output_dir: Path, + *, + generation: int, + population: Sequence[GaIndividual], + snapshot: FitnessSnapshot, + best: GaIndividual, + config: GaPromptOptimizerConfig, +) -> None: + payload = { + "generation": generation, + "best_individual": _individual_payload(best), + "fitness": { + "valid_count": snapshot.valid_count, + "failed_count": snapshot.failed_count, + "best_fitness": snapshot.best_fitness, + "fitness_variance": snapshot.fitness_variance, + "duplicate_ratio": snapshot.duplicate_ratio, + }, + "population": [_individual_payload(individual) for individual in population], + "prompt_paths": {name: spec.path for name, spec in sorted(config.search_space.items())}, + } + (output_dir / f"optimized_prompts_gen{generation}.json").write_text( + json.dumps(payload, indent=2, default=str) + "\n", + encoding="utf-8", + ) + checkpoint_dir = output_dir / "checkpoints" + checkpoint_dir.mkdir(parents=True, exist_ok=True) + (checkpoint_dir / f"generation_{generation:03d}.json").write_text( + json.dumps(payload, indent=2, default=str) + "\n", + encoding="utf-8", + ) + + +def _write_final_artifacts( + output_dir: Path, + *, + base_payload: Mapping[str, Any], + optimized_payload: Mapping[str, Any], + best: GaIndividual, + config: GaPromptOptimizerConfig, + history: Sequence[GaIndividual], + executed_trials: int, + generations_completed: int, +) -> None: + optimized_config = apply_suggestions( + base_payload, + suggestions_by_path(config.search_space, best.prompts), + ) + (output_dir / "optimized_config.yml").write_text( + yaml.safe_dump(sanitize_config_for_artifact(optimized_config), sort_keys=False), + encoding="utf-8", + ) + (output_dir / "optimized_prompts.json").write_text( + json.dumps( + { + "best_individual": _individual_payload(best), + "prompts": dict(best.prompts), + "prompt_paths": {name: spec.path for name, spec in sorted(config.search_space.items())}, + "metrics": dict(best.aggregate_metrics), + "fitness": best.fitness, + }, + indent=2, + default=str, + ) + + "\n", + encoding="utf-8", + ) + _write_history_artifacts(output_dir, history=history) + (output_dir / "ga_summary.json").write_text( + json.dumps( + { + "status": "completed", + "backend": "ga", + "phase": "prompt", + "population_size": config.population_size, + "generations": config.generations, + "generations_completed": generations_completed, + "executed_trials": executed_trials, + "metric_names": list(config.metric_names), + "best_individual": best.individual_id, + "best_phase_trial_number": best.phase_trial_number, + "best_global_trial_number": best.global_trial_number, + "best_metrics": dict(best.aggregate_metrics), + "best_fitness": best.fitness, + "optimized_payload": sanitize_config_for_artifact(optimized_payload), + }, + indent=2, + default=str, + ) + + "\n", + encoding="utf-8", + ) + + +def _write_history_artifacts(output_dir: Path, *, history: Sequence[GaIndividual]) -> None: + rows = [individual.to_history_row() for individual in history] + (output_dir / "ga_history_prompts.json").write_text( + json.dumps(rows, indent=2, default=str) + "\n", + encoding="utf-8", + ) + (output_dir / "ga_score_records.json").write_text( + json.dumps(_score_record_rows(history), indent=2, default=str) + "\n", + encoding="utf-8", + ) + csv_path = output_dir / "ga_history_prompts.csv" + fieldnames = _history_fieldnames(rows) + with csv_path.open("w", encoding="utf-8", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def _write_failure_artifact( + output_dir: Path, + *, + config: GaPromptOptimizerConfig, + message: str, + trial_count: int, + best: GaIndividual | None = None, +) -> None: + payload = { + "status": "failed", + "backend": "ga", + "phase": "prompt", + "error": message, + "executed_trials": trial_count, + "metric_names": list(config.metric_names), + "best_individual": _individual_payload(best) if best is not None else None, + } + (output_dir / "prompt_phase_failure.json").write_text( + json.dumps(payload, indent=2, default=str) + "\n", + encoding="utf-8", + ) + + +def _history_fieldnames(rows: Sequence[Mapping[str, Any]]) -> list[str]: + prefix = [ + "individual_id", + "phase", + "generation", + "individual_index", + "phase_trial_number", + "global_trial_number", + "status", + "fitness", + "failure_reason", + "parent_ids", + "carried_from", + "transform_failures", + ] + fields = set().union(*(row.keys() for row in rows)) if rows else set() + dynamic = sorted(field for field in fields if field not in prefix) + return [*prefix, *dynamic] + + +def _score_record_rows(history: Sequence[GaIndividual]) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for individual in history: + if individual.carried_from is not None: + continue + metadata = { + "individual_id": individual.individual_id, + "phase": "prompt", + "generation": individual.generation, + "individual_index": individual.individual_index, + "phase_trial_number": individual.phase_trial_number, + "global_trial_number": individual.global_trial_number, + } + for score in individual.raw_scores: + rows.append({**metadata, "score": score.model_dump(mode="json")}) + return rows + + +def _individual_payload(individual: GaIndividual) -> dict[str, Any]: + return { + "individual_id": individual.individual_id, + "generation": individual.generation, + "individual_index": individual.individual_index, + "phase_trial_number": individual.phase_trial_number, + "global_trial_number": individual.global_trial_number, + "status": individual.status, + "fitness": individual.fitness, + "aggregate_metrics": dict(individual.aggregate_metrics), + "normalized_metrics": dict(individual.normalized_metrics), + "parent_ids": list(individual.parent_ids), + "carried_from": individual.carried_from, + "failure_reason": individual.failure_reason, + "transform_failures": list(individual.transform_failures), + "prompts": dict(individual.prompts), + } + + +__all__ = [ + "GaPromptOptimizationResult", + "GaPromptOptimizerError", + "run_ga_prompt_optimization", + "write_candidate_config", +] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/fitness.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/fitness.py new file mode 100644 index 0000000000..0954d64093 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/fitness.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fitness normalization, scalarization, and ranking for prompt GA.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass +from statistics import pvariance + +from nemo_optimization.backends.ga.config import MetricDirection, MetricSpec +from nemo_optimization.backends.ga.individual import GaIndividual + +EPSILON = 1e-12 + + +class GaFitnessError(RuntimeError): + """Raised when GA fitness cannot be computed.""" + + +@dataclass(frozen=True) +class FitnessSnapshot: + """Aggregate fitness/diversity signals for one generation.""" + + valid_count: int + failed_count: int + best_fitness: float + fitness_variance: float + duplicate_ratio: float + + +def assign_generation_fitness( + population: Sequence[GaIndividual], + *, + metrics: Sequence[MetricSpec], + mode: str, + diversity_lambda: float, +) -> FitnessSnapshot: + """Normalize completed individuals within one generation and assign higher-is-better fitness.""" + + valid = _valid_for_metrics(population, metrics) + if not valid: + failed_count = len([individual for individual in population if individual.status == "failed"]) + raise GaFitnessError( + f"Generation {population[0].generation if population else 0} has no valid individuals " + f"({failed_count} failed)." + ) + + normalized = _normalized_metric_values(valid, metrics) + scalar_scores = _scalar_scores(normalized, metrics=metrics, mode=mode) + penalties = _diversity_penalties(valid, diversity_lambda=diversity_lambda) + for index, individual in enumerate(valid): + individual.normalized_metrics = normalized[index] + individual.fitness = max(0.0, scalar_scores[index] - penalties[index]) + + fitness_values = [individual.fitness for individual in valid if individual.fitness is not None] + return FitnessSnapshot( + valid_count=len(valid), + failed_count=len(population) - len(valid), + best_fitness=max(fitness_values), + fitness_variance=pvariance(fitness_values) if len(fitness_values) > 1 else 0.0, + duplicate_ratio=duplicate_ratio(valid), + ) + + +def rank_valid_individuals(population: Sequence[GaIndividual]) -> list[GaIndividual]: + """Return selectable individuals best-first.""" + + valid = [individual for individual in population if individual.is_valid] + return sorted(valid, key=_rank_key, reverse=True) + + +def best_individual( + individuals: Sequence[GaIndividual], + *, + metrics: Sequence[MetricSpec], + mode: str, +) -> GaIndividual | None: + """Pick the best completed individual across generations using global normalization.""" + + valid = _valid_for_metrics(individuals, metrics) + if not valid: + return None + normalized = _normalized_metric_values(valid, metrics) + scalar_scores = _scalar_scores(normalized, metrics=metrics, mode=mode) + ranked = sorted( + zip(valid, scalar_scores, strict=True), + key=lambda item: ( + _safe_score(item[1]), + -_trial_number(item[0]), + -item[0].generation, + -item[0].individual_index, + ), + reverse=True, + ) + return ranked[0][0] + + +def target_met(individual: GaIndividual, *, metrics: Sequence[MetricSpec], target: float | None) -> bool: + """Return true when a single-objective target has been met.""" + + if target is None or len(metrics) != 1: + return False + metric = metrics[0] + if metric.name not in individual.aggregate_metrics: + return False + value = individual.aggregate_metrics[metric.name] + if metric.direction is MetricDirection.MINIMIZE: + return value <= target + return value >= target + + +def duplicate_ratio(individuals: Sequence[GaIndividual]) -> float: + """Return the share of valid individuals that duplicate another prompt signature.""" + + if not individuals: + return 0.0 + unique_count = len({individual.prompt_signature for individual in individuals}) + return 1.0 - (unique_count / len(individuals)) + + +def _valid_for_metrics(individuals: Sequence[GaIndividual], metrics: Sequence[MetricSpec]) -> list[GaIndividual]: + metric_names = {metric.name for metric in metrics} + return [ + individual + for individual in individuals + if individual.status == "completed" and metric_names.issubset(individual.aggregate_metrics) + ] + + +def _normalized_metric_values( + individuals: Sequence[GaIndividual], + metrics: Sequence[MetricSpec], +) -> list[dict[str, float]]: + by_metric: dict[str, list[float]] = { + metric.name: [float(individual.aggregate_metrics[metric.name]) for individual in individuals] + for metric in metrics + } + normalized = [{metric.name: 0.0 for metric in metrics} for _ in individuals] + for metric in metrics: + values = by_metric[metric.name] + low = min(values) + high = max(values) + if abs(high - low) < EPSILON: + for row in normalized: + row[metric.name] = 1.0 + continue + for index, value in enumerate(values): + if metric.direction is MetricDirection.MINIMIZE: + normalized[index][metric.name] = (high - value) / (high - low) + else: + normalized[index][metric.name] = (value - low) / (high - low) + return normalized + + +def _scalar_scores( + normalized: Sequence[dict[str, float]], + *, + metrics: Sequence[MetricSpec], + mode: str, +) -> list[float]: + total_weight = sum(metric.weight for metric in metrics) + weights = [metric.weight / total_weight for metric in metrics] + scores: list[float] = [] + for row in normalized: + values = [max(0.0, min(1.0, float(row[metric.name]))) for metric in metrics] + if mode == "weighted_sum": + score = sum(weight * value for weight, value in zip(weights, values, strict=True)) + elif mode == "chebyshev": + score = 1.0 - max(weight * (1.0 - value) for weight, value in zip(weights, values, strict=True)) + elif mode == "harmonic": + score = 1.0 / sum(weight / max(value, EPSILON) for weight, value in zip(weights, values, strict=True)) + else: + raise GaFitnessError(f"Unsupported multi-objective combination mode: {mode!r}") + scores.append(score) + return scores + + +def _diversity_penalties(individuals: Sequence[GaIndividual], *, diversity_lambda: float) -> list[float]: + if diversity_lambda <= 0 or len(individuals) <= 1: + return [0.0 for _ in individuals] + counts = Counter(individual.prompt_signature for individual in individuals) + denominator = max(1, len(individuals) - 1) + return [diversity_lambda * ((counts[individual.prompt_signature] - 1) / denominator) for individual in individuals] + + +def _rank_key(individual: GaIndividual) -> tuple[float, int, int, int]: + return ( + _safe_score(individual.fitness), + -_trial_number(individual), + -individual.generation, + -individual.individual_index, + ) + + +def _safe_score(score: float | None) -> float: + return score if score is not None else float("-inf") + + +def _trial_number(individual: GaIndividual) -> int: + if individual.global_trial_number is not None: + return individual.global_trial_number + if individual.phase_trial_number is not None: + return individual.phase_trial_number + return 1_000_000_000 + + +__all__ = [ + "FitnessSnapshot", + "GaFitnessError", + "assign_generation_fitness", + "best_individual", + "duplicate_ratio", + "rank_valid_individuals", + "target_met", +] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/individual.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/individual.py new file mode 100644 index 0000000000..abcc0e5354 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/individual.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prompt GA individual state.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass, field +from typing import Any + +from nemo_evaluator_sdk.agent_eval.scores import AgentEvalTaskScore + + +@dataclass +class GaIndividual: + """One prompt candidate in a GA population.""" + + prompts: dict[str, str] + generation: int + individual_index: int + parent_ids: tuple[str, ...] = () + carried_from: str | None = None + phase_trial_number: int | None = None + global_trial_number: int | None = None + aggregate_metrics: dict[str, float] = field(default_factory=dict) + normalized_metrics: dict[str, float] = field(default_factory=dict) + raw_scores: tuple[AgentEvalTaskScore, ...] = () + fitness: float | None = None + status: str = "pending" + failure_reason: str | None = None + transform_failures: list[str] = field(default_factory=list) + + @property + def individual_id(self) -> str: + """Stable generation-local individual identifier.""" + + return f"g{self.generation:03d}-i{self.individual_index:03d}" + + @property + def is_valid(self) -> bool: + """Whether this individual can participate in selection.""" + + return self.status == "completed" and self.fitness is not None + + @property + def prompt_signature(self) -> tuple[tuple[str, str], ...]: + """Deterministic prompt tuple for duplicate/diversity checks.""" + + return tuple(sorted(self.prompts.items())) + + def clone_as_elite(self, *, generation: int, individual_index: int) -> GaIndividual: + """Copy an already evaluated elite into a later generation without reevaluation.""" + + return GaIndividual( + prompts=copy.deepcopy(self.prompts), + generation=generation, + individual_index=individual_index, + parent_ids=(self.individual_id,), + carried_from=self.individual_id, + phase_trial_number=self.phase_trial_number, + global_trial_number=self.global_trial_number, + aggregate_metrics=dict(self.aggregate_metrics), + normalized_metrics=dict(self.normalized_metrics), + raw_scores=self.raw_scores, + fitness=self.fitness, + status=self.status, + failure_reason=self.failure_reason, + transform_failures=list(self.transform_failures), + ) + + def to_history_row(self) -> dict[str, Any]: + """Return a JSON/CSV-safe summary row for this individual.""" + + return { + "individual_id": self.individual_id, + "phase": "prompt", + "generation": self.generation, + "individual_index": self.individual_index, + "phase_trial_number": self.phase_trial_number, + "global_trial_number": self.global_trial_number, + "status": self.status, + "fitness": self.fitness, + "failure_reason": self.failure_reason, + "parent_ids": ",".join(self.parent_ids), + "carried_from": self.carried_from, + "transform_failures": " | ".join(self.transform_failures), + **{f"metric.{name}": value for name, value in sorted(self.aggregate_metrics.items())}, + **{f"normalized.{name}": value for name, value in sorted(self.normalized_metrics.items())}, + **{f"prompt.{name}": value for name, value in sorted(self.prompts.items())}, + } + + +__all__ = ["GaIndividual"] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/oracle_feedback.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/oracle_feedback.py new file mode 100644 index 0000000000..067d1b3868 --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/oracle_feedback.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Oracle feedback extraction for prompt GA operators.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from nemo_optimization.backends.ga.config import GaPromptOptimizerConfig, MetricDirection, MetricSpec +from nemo_optimization.backends.ga.individual import GaIndividual +from nemo_optimization.candidate import CandidateEvaluationResult, RowReasoning + + +@dataclass(frozen=True) +class OracleFeedbackState: + """Generation-level signals for adaptive oracle feedback.""" + + stagnation_generations: int + fitness_variance: float + duplicate_ratio: float + + +def should_use_oracle_feedback( + *, + config: GaPromptOptimizerConfig, + parent: GaIndividual, + state: OracleFeedbackState, +) -> bool: + """Decide whether a transform should include evaluator reasoning feedback.""" + + mode = config.oracle_feedback_mode + if mode == "never": + return False + if mode == "always": + return True + if mode == "failing_only": + return parent.fitness is not None and parent.fitness < config.oracle_feedback_fitness_threshold + if mode == "adaptive": + return ( + state.stagnation_generations >= config.oracle_feedback_stagnation_generations + or state.fitness_variance <= config.oracle_feedback_fitness_variance_threshold + or state.duplicate_ratio >= config.oracle_feedback_diversity_threshold + ) + return False + + +def build_oracle_feedback( + *, + individual: GaIndividual, + config: GaPromptOptimizerConfig, +) -> str | None: + """Build compact row-level feedback from evaluator reasoning outputs.""" + + if not individual.raw_scores: + return None + + evaluation = CandidateEvaluationResult( + aggregate_metrics=dict(individual.aggregate_metrics), + scores=individual.raw_scores, + ) + sections: list[str] = [] + weighted_metrics = sorted(config.metrics, key=lambda metric: (metric.weight, metric.name), reverse=True) + total_weight = sum(metric.weight for metric in weighted_metrics) + remaining_chars = config.oracle_feedback_max_chars + for metric in weighted_metrics: + section = _metric_feedback_section(evaluation, individual=individual, metric=metric, config=config) + if not section: + continue + separator_chars = 2 if sections else 0 + available_chars = remaining_chars - separator_chars + if available_chars <= 0: + break + budget = int(config.oracle_feedback_max_chars * (metric.weight / total_weight)) + budget = max(1, min(budget, available_chars)) + sections.append(section[:budget]) + remaining_chars -= separator_chars + len(sections[-1]) + + feedback = "\n\n".join(sections).strip() + if not feedback: + return None + return feedback + + +def _metric_feedback_section( + evaluation: CandidateEvaluationResult, + *, + individual: GaIndividual, + metric: MetricSpec, + config: GaPromptOptimizerConfig, +) -> str | None: + if metric.name not in individual.aggregate_metrics: + return None + reasoning_rows = list(evaluation.reasoning_for_metric(metric.name)) + if not reasoning_rows: + return None + rows = _worst_rows(reasoning_rows, direction=metric.direction, limit=config.oracle_feedback_worst_n) + lines = [ + f"Metric {metric.name} ({metric.direction.value}, weight={metric.weight:.6g}); " + f"aggregate={individual.aggregate_metrics[metric.name]:.6g}:" + ] + for row in rows: + lines.append(f"- row={row.task_id} score={row.objective_value:.6g} reasoning={row.reasoning.strip()}") + return "\n".join(lines) + + +def _worst_rows( + rows: list[RowReasoning], + *, + direction: MetricDirection, + limit: int, +) -> list[RowReasoning]: + reverse = direction is MetricDirection.MINIMIZE + return sorted(rows, key=lambda row: row.objective_value, reverse=reverse)[:limit] + + +__all__ = [ + "OracleFeedbackState", + "build_oracle_feedback", + "should_use_oracle_feedback", +] diff --git a/plugins/nemo-optimization/src/nemo_optimization/backends/ga/transform.py b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/transform.py new file mode 100644 index 0000000000..b76d80cf6f --- /dev/null +++ b/plugins/nemo-optimization/src/nemo_optimization/backends/ga/transform.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prompt mutation and recombination interfaces for GA optimization.""" + +from __future__ import annotations + +import json +import os +import re +import urllib.error +import urllib.request +from collections.abc import Mapping +from typing import Any, Protocol + + +class PromptTransformError(RuntimeError): + """Raised when a prompt transform cannot be produced.""" + + +class PromptTransformer(Protocol): + """Generate prompt variants for GA operators.""" + + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + """Return one mutated prompt.""" + ... + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + """Return one child prompt from two parent prompts.""" + ... + + +class ModelPromptTransformer: + """Direct OpenAI-compatible model caller used by the prompt GA backend. + + This adapter intentionally supports only model declarations with a direct + ``url``/``base_url`` and optional ``api_key_env``. It does not resolve + platform secret references; unsupported Fabric model shapes fail before the + first mutation instead of falling back to guessed credentials. + """ + + def __init__( + self, + *, + payload: Mapping[str, Any], + model_name: str, + temperature: float | None = None, + max_tokens: int | None = None, + timeout_s: float | None = None, + ) -> None: + model_config = _model_config(payload, model_name) + settings = _model_settings(model_config, model_name) + self._model_name = str(model_config.get("model") or model_config.get("model_name") or model_name) + self._base_url = _chat_completions_url(model_config, model_name) + self._api_key = _api_key(model_config) + self._temperature = _float_setting( + settings, + "temperature", + default=0.7 if temperature is None else temperature, + ) + self._max_tokens = _int_setting( + settings, + "max_tokens", + default=2048 if max_tokens is None else max_tokens, + ) + self._timeout_s = _float_setting( + settings, + "timeout_s", + default=60.0 if timeout_s is None else timeout_s, + ) + + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + user_prompt = _mutation_user_prompt( + prompt_name=prompt_name, + prompt=prompt, + purpose=purpose, + prompt_format=prompt_format, + feedback=feedback, + ) + return self._complete(user_prompt) + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + user_prompt = _recombination_user_prompt( + prompt_name=prompt_name, + parent_a=parent_a, + parent_b=parent_b, + purpose=purpose, + prompt_format=prompt_format, + feedback=feedback, + ) + return self._complete(user_prompt) + + def _complete(self, user_prompt: str) -> str: + body = { + "model": self._model_name, + "messages": [ + { + "role": "system", + "content": ( + "You are optimizing an agent prompt. Return only the revised prompt text, " + "with no markdown fences, commentary, or scoring rationale." + ), + }, + {"role": "user", "content": user_prompt}, + ], + "temperature": self._temperature, + "max_tokens": self._max_tokens, + } + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + request = urllib.request.Request( + self._base_url, + data=json.dumps(body).encode("utf-8"), + headers=headers, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=self._timeout_s) as response: # noqa: S310 + response_payload = json.loads(response.read().decode("utf-8")) + except (OSError, urllib.error.HTTPError, urllib.error.URLError, json.JSONDecodeError) as exc: + raise PromptTransformError(f"Prompt optimizer model call failed: {exc}") from exc + + content = _extract_chat_content(response_payload) + if not content: + raise PromptTransformError("Prompt optimizer model returned an empty prompt.") + return content + + +def _model_config(payload: Mapping[str, Any], model_name: str) -> Mapping[str, Any]: + models = payload.get("models") + if not isinstance(models, Mapping): + raise PromptTransformError("Fabric payload must declare models for prompt GA optimization.") + model_config = models.get(model_name) + if not isinstance(model_config, Mapping): + raise PromptTransformError(f"Prompt optimizer model {model_name!r} not found under payload.models.") + provider = str(model_config.get("provider") or "openai").strip().lower().replace("-", "_") + supported = {"openai", "openai_compatible", "nvidia", "nim", "nvidia_nim"} + if provider not in supported: + raise PromptTransformError( + f"Prompt optimizer model {model_name!r} has unsupported provider {provider!r}; " + "the GA direct adapter currently supports OpenAI-compatible providers with url/base_url." + ) + return model_config + + +def _model_settings(model_config: Mapping[str, Any], model_name: str) -> Mapping[str, Any]: + settings = model_config.get("settings") + if settings is None: + return {} + if not isinstance(settings, Mapping): + raise PromptTransformError(f"Prompt optimizer model {model_name!r} settings must be a mapping.") + return settings + + +def _chat_completions_url(model_config: Mapping[str, Any], model_name: str) -> str: + raw_url = model_config.get("url") or model_config.get("base_url") + if not isinstance(raw_url, str) or not raw_url.strip(): + raise PromptTransformError(f"Prompt optimizer model {model_name!r} must declare 'url' or 'base_url'.") + url = raw_url.strip().rstrip("/") + if url.endswith("/chat/completions"): + return url + if url.endswith("/v1"): + return f"{url}/chat/completions" + return f"{url}/v1/chat/completions" + + +def _api_key(model_config: Mapping[str, Any]) -> str | None: + if model_config.get("api_key_secret") is not None and model_config.get("api_key_env") is None: + raise PromptTransformError( + "Prompt GA direct model adapter supports api_key_env, not api_key_secret; " + "use an environment-backed model declaration for prompt optimization." + ) + env_name = model_config.get("api_key_env") + if not isinstance(env_name, str) or not env_name.strip(): + return None + value = os.getenv(env_name.strip()) + if value is None: + raise PromptTransformError(f"Prompt optimizer model api_key_env {env_name!r} is not set.") + return value + + +def _float_setting(settings: Mapping[str, Any], key: str, *, default: float) -> float: + value = settings.get(key, default) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise PromptTransformError(f"Prompt optimizer model setting {key!r} must be a number.") + return float(value) + + +def _int_setting(settings: Mapping[str, Any], key: str, *, default: int) -> int: + value = settings.get(key, default) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise PromptTransformError(f"Prompt optimizer model setting {key!r} must be a positive integer.") + return value + + +def _extract_chat_content(payload: Mapping[str, Any]) -> str: + choices = payload.get("choices") + if not isinstance(choices, list) or not choices: + return "" + first = choices[0] + if not isinstance(first, Mapping): + return "" + message = first.get("message") + content: Any + if isinstance(message, Mapping): + content = message.get("content") + else: + content = first.get("text") + if not isinstance(content, str): + return "" + return _strip_prompt_wrapper(content) + + +def _strip_prompt_wrapper(content: str) -> str: + stripped = content.strip() + fenced = re.fullmatch(r"```(?:[A-Za-z0-9_-]+)?\s*(.*?)\s*```", stripped, flags=re.DOTALL) + if fenced: + stripped = fenced.group(1).strip() + return stripped + + +def _mutation_user_prompt( + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, +) -> str: + sections = [ + f"Prompt dimension: {prompt_name}", + f"Purpose: {purpose}", + f"Required format: {prompt_format or 'Preserve the current prompt format.'}", + "Current prompt:", + prompt, + ] + if feedback: + sections.extend(["Evaluation feedback to address:", feedback]) + sections.append("Rewrite the prompt to improve the objective while preserving the agent's role and constraints.") + return "\n\n".join(sections) + + +def _recombination_user_prompt( + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, +) -> str: + sections = [ + f"Prompt dimension: {prompt_name}", + f"Purpose: {purpose}", + f"Required format: {prompt_format or 'Preserve the current prompt format.'}", + "Parent prompt A:", + parent_a, + "Parent prompt B:", + parent_b, + ] + if feedback: + sections.extend(["Evaluation feedback to address:", feedback]) + sections.append("Combine the strongest parts of both parents into one improved prompt.") + return "\n\n".join(sections) + + +__all__ = [ + "ModelPromptTransformer", + "PromptTransformError", + "PromptTransformer", +] diff --git a/plugins/nemo-optimization/src/nemo_optimization/fabric_evaluator.py b/plugins/nemo-optimization/src/nemo_optimization/fabric_evaluator.py index 1c0ff9cecf..70e2f115c8 100644 --- a/plugins/nemo-optimization/src/nemo_optimization/fabric_evaluator.py +++ b/plugins/nemo-optimization/src/nemo_optimization/fabric_evaluator.py @@ -8,6 +8,7 @@ import copy import json import logging +import threading from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any @@ -67,6 +68,7 @@ def __init__( general = general if isinstance(general, Mapping) else {} self._parallelism = int(general.get("max_concurrency", default_parallelism)) self._trace_map: list[dict[str, Any]] = [] + self._trace_lock = threading.Lock() # Validate dataset/metrics once at construction so config errors fail before the study loop. build_agent_eval_tasks(self._payload) @@ -78,7 +80,7 @@ def evaluate( trial_overlay: dict[str, Any], rep: int, ) -> CandidateEvaluationResult: - del trial_overlay # reserved for profile overlays; runtime uses path-resolved payload + overlay_metadata = _trial_overlay_metadata(trial_overlay) # ``suggestions`` are Fabric dotted paths. trial_payload = apply_suggestions(self._payload, suggestions) # Rebuild tasks from the path-resolved payload so search-space paths under @@ -90,11 +92,14 @@ def evaluate( work_root=self._trial_work_root(trial_number, rep), timeout_s=self._timeout_s, capture_trajectory=self._capture_trajectory, - trajectory_extra=build_atif_trial_tags( - experiment_id=self._experiment_id, - trial_number=trial_number, - rep=rep, - ), + trajectory_extra={ + **build_atif_trial_tags( + experiment_id=self._experiment_id, + trial_number=trial_number, + rep=rep, + ), + **overlay_metadata, + }, task_hook=self._task_hook, ) result = AgentEvaluator().run_sync( @@ -111,8 +116,9 @@ def evaluate( # ``run`` no longer stores the bundle; keep writing one per trial as before. No dashboard — # each trial bundle is intermediate evidence for the study, not something anyone opens. result.persist(write_dashboard=False) - self._record_traces(result, trial_number=trial_number, rep=rep) - self._write_trace_map() + with self._trace_lock: + self._record_traces(result, trial_number=trial_number, rep=rep, metadata=overlay_metadata) + self._write_trace_map() return CandidateEvaluationResult( aggregate_metrics=reduce_agent_eval_scores(result.scores, self._metric_names), scores=tuple(result.scores), @@ -124,7 +130,14 @@ def _trial_work_root(self, trial_number: int, rep: int) -> Path: def _trial_output_dir(self, trial_number: int, rep: int) -> Path: return self._output_dir / "agent_eval" / f"trial-{trial_number:03d}" / f"rep-{rep:03d}" - def _record_traces(self, result: AgentEvalResult, *, trial_number: int, rep: int) -> None: + def _record_traces( + self, + result: AgentEvalResult, + *, + trial_number: int, + rep: int, + metadata: Mapping[str, Any], + ) -> None: for trial in result.trials: trace = trial.evidence.descriptors.get(EVIDENCE_TRACE) if trial.evidence is not None else None if trace is None: @@ -139,6 +152,7 @@ def _record_traces(self, result: AgentEvalResult, *, trial_number: int, rep: int "trial_id": trial.id, "trace_ref": trace.ref, "trace_format": trace.format, + **metadata, } ) @@ -326,6 +340,15 @@ def _eval_config(payload: Mapping[str, Any]) -> Mapping[str, Any]: return eval_config +def _trial_overlay_metadata(trial_overlay: Mapping[str, Any] | None) -> dict[str, Any]: + if not isinstance(trial_overlay, Mapping): + return {} + metadata = trial_overlay.get("metadata") + if not isinstance(metadata, Mapping): + return {} + return {str(key): value for key, value in metadata.items()} + + def _runtime_agent_config(config: Mapping[str, Any]) -> dict[str, Any]: runtime_config = copy.deepcopy(dict(config)) runtime_config.pop("eval", None) diff --git a/plugins/nemo-optimization/tests/test_fabric_trial.py b/plugins/nemo-optimization/tests/test_fabric_trial.py index 26c0810e3c..b17c8c094f 100644 --- a/plugins/nemo-optimization/tests/test_fabric_trial.py +++ b/plugins/nemo-optimization/tests/test_fabric_trial.py @@ -324,6 +324,7 @@ def run_sync(self, *, tasks, target, config): # noqa: ANN001 "nemo.optimizer.experiment_id": "exp-test", "nemo.optimizer.trial_number": 7, "nemo.optimizer.rep": 0, + "name": "trial-007", } assert "profiles" not in captured["runtime"] assert captured["runtime"]["task_hook"] is None diff --git a/plugins/nemo-optimization/tests/test_ga_driver.py b/plugins/nemo-optimization/tests/test_ga_driver.py new file mode 100644 index 0000000000..80b867ad5d --- /dev/null +++ b/plugins/nemo-optimization/tests/test_ga_driver.py @@ -0,0 +1,691 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path +from typing import Any + +import pytest +from nemo_evaluator_sdk.agent_eval.scores import AgentEvalScoreStatus, AgentEvalTaskScore +from nemo_evaluator_sdk.metrics.protocol import MetricOutput +from nemo_optimization.backends.ga.config import parse_ga_prompt_optimizer_config +from nemo_optimization.backends.ga.driver import GaPromptOptimizerError, _child_prompt, run_ga_prompt_optimization +from nemo_optimization.backends.ga.fitness import assign_generation_fitness, rank_valid_individuals +from nemo_optimization.backends.ga.individual import GaIndividual +from nemo_optimization.backends.ga.oracle_feedback import OracleFeedbackState, build_oracle_feedback +from nemo_optimization.backends.ga.transform import PromptTransformError +from nemo_optimization.candidate import CandidateEvaluationError, CandidateEvaluationResult + + +def test_ga_prompt_optimizer_runs_prompt_only_and_writes_artifacts(tmp_path: Path) -> None: + payload = _payload( + prompt={ + "population_size": 3, + "generations": 2, + "mutation_rate": 1.0, + "crossover_rate": 1.0, + "elitism": 1, + "seed": 11, + "parallel_evaluations": 1, + } + ) + evaluator = RecordingEvaluator() + + result = run_ga_prompt_optimization( + payload, + tmp_path, + evaluator, + DeterministicTransformer(), + trial_number_offset=10, + ) + + assert result.executed_trials == 5 + assert result.generations_completed == 2 + assert result.best_individual.global_trial_number in {10, 11, 12, 13, 14} + assert "[system_prompt:" in result.optimized_payload["instructions"]["system"]["content"] + assert "optimizer" in result.optimized_payload + assert (tmp_path / "optimized_config.yml").is_file() + assert (tmp_path / "optimized_prompts.json").is_file() + assert (tmp_path / "optimized_prompts_gen0.json").is_file() + assert (tmp_path / "optimized_prompts_gen1.json").is_file() + assert (tmp_path / "ga_history_prompts.csv").is_file() + assert (tmp_path / "ga_score_records.json").is_file() + assert (tmp_path / "checkpoints" / "generation_001.json").is_file() + + assert [call["metadata"]["nemo.optimizer.global_trial_number"] for call in evaluator.calls] == [ + 10, + 11, + 12, + 13, + 14, + ] + assert {call["metadata"]["nemo.optimizer.phase"] for call in evaluator.calls} == {"prompt"} + assert any(individual.carried_from for individual in result.history) + + +def test_ga_transform_failures_fallback_without_aborting(tmp_path: Path) -> None: + payload = _payload( + prompt={ + "population_size": 3, + "generations": 2, + "mutation_rate": 0.0, + "crossover_rate": 1.0, + "elitism": 0, + "seed": 7, + "parallel_evaluations": 1, + } + ) + + result = run_ga_prompt_optimization( + payload, + tmp_path, + RecordingEvaluator(), + FailingTransformer(), + ) + + assert result.executed_trials == 6 + assert all(individual.status == "completed" for individual in result.history) + assert all(individual.prompts["system_prompt"] == "Base prompt." for individual in result.history) + assert any("mutate:system_prompt" in " ".join(individual.transform_failures) for individual in result.history) + assert any("recombine:system_prompt" in " ".join(individual.transform_failures) for individual in result.history) + + +def test_ga_uses_bounded_parallel_candidate_evaluation(tmp_path: Path) -> None: + payload = _payload( + prompt={ + "population_size": 6, + "generations": 1, + "parallel_evaluations": 3, + "seed": 13, + } + ) + evaluator = ConcurrencyEvaluator() + + result = run_ga_prompt_optimization( + payload, + tmp_path, + evaluator, + DeterministicTransformer(), + ) + + assert result.executed_trials == 6 + assert evaluator.max_active > 1 + assert evaluator.max_active <= 3 + assert sorted(evaluator.trial_numbers) == list(range(6)) + + +def test_ga_excludes_failed_individuals_from_parent_selection(tmp_path: Path) -> None: + payload = _payload( + prompt={ + "population_size": 3, + "generations": 2, + "mutation_rate": 0.0, + "crossover_rate": 0.0, + "elitism": 0, + "seed": 3, + "parallel_evaluations": 1, + } + ) + + result = run_ga_prompt_optimization( + payload, + tmp_path, + FailsBadPromptEvaluator(), + OneBadMutationTransformer(), + ) + + failed = [individual for individual in result.history if individual.status == "failed"] + assert [individual.individual_id for individual in failed] == ["g000-i001"] + generation_one = [individual for individual in result.history if individual.generation == 1] + assert generation_one + assert all("g000-i001" not in individual.parent_ids for individual in generation_one) + + +def test_ga_fails_when_every_individual_in_generation_fails(tmp_path: Path) -> None: + payload = _payload(prompt={"population_size": 3, "generations": 1, "seed": 5}) + + with pytest.raises(GaPromptOptimizerError) as exc_info: + run_ga_prompt_optimization( + payload, + tmp_path, + AlwaysFailEvaluator(), + DeterministicTransformer(), + ) + + assert exc_info.value.trial_count == 3 + failure = json.loads((tmp_path / "prompt_phase_failure.json").read_text(encoding="utf-8")) + assert failure["status"] == "failed" + history = json.loads((tmp_path / "ga_history_prompts.json").read_text(encoding="utf-8")) + assert {row["status"] for row in history} == {"failed"} + + +def test_ga_fitness_supports_minimize_and_diversity_penalty() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + optimizer={ + "multi_objective_combination_mode": "weighted_sum", + "eval_metrics": {"latency": {"direction": "minimize", "weight": 1.0}}, + }, + prompt={"population_size": 3, "generations": 1, "diversity_lambda": 0.6}, + ) + ) + population = [ + _completed_individual(prompt="same", score=3.0, index=0, metric_name="latency"), + _completed_individual(prompt="same", score=1.0, index=1, metric_name="latency"), + _completed_individual(prompt="other", score=2.0, index=2, metric_name="latency"), + ] + + snapshot = assign_generation_fitness( + population, + metrics=config.metrics, + mode=config.multi_objective_mode, + diversity_lambda=config.diversity_lambda, + ) + + assert snapshot.duplicate_ratio == pytest.approx(1 / 3) + ranked = rank_valid_individuals(population) + assert ranked[0].individual_index == 1 + assert ranked[0].normalized_metrics["latency"] == 1.0 + + +def test_ga_config_accepts_ga_prefixed_aliases() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + optimizer={"multi_objective_combination_mode": "weighted-sum"}, + prompt={ + "ga_population_size": 4, + "ga_generations": 2, + "ga_crossover_rate": 0.4, + "ga_mutation_rate": 0.5, + "ga_elitism": 1, + "ga_selection_method": "roulette", + "ga_tournament_size": 2, + "ga_parallel_evaluations": 3, + }, + ) + ) + + assert config.population_size == 4 + assert config.generations == 2 + assert config.crossover_rate == 0.4 + assert config.mutation_rate == 0.5 + assert config.elitism == 1 + assert config.selection_method == "roulette" + assert config.tournament_size == 2 + assert config.parallel_evaluations == 3 + assert config.multi_objective_mode == "weighted_sum" + + +def test_oracle_feedback_uses_worst_reasoning_rows_for_metric() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + optimizer={"eval_metrics": {"average_score": {"direction": "maximize", "weight": 1.0}}}, + prompt={"oracle_feedback_worst_n": 1, "oracle_feedback_mode": "always"}, + ) + ) + individual = GaIndividual( + prompts={"system_prompt": "Base prompt."}, + generation=0, + individual_index=0, + status="completed", + aggregate_metrics={"average_score": 0.6}, + raw_scores=( + _score(task_id="bad-row", metric_name="average_score", value=0.2, reasoning="Missed the key fact."), + _score(task_id="good-row", metric_name="average_score", value=0.9, reasoning="Correct and complete."), + ), + ) + + feedback = build_oracle_feedback(individual=individual, config=config) + + assert feedback is not None + assert "bad-row" in feedback + assert "Missed the key fact." in feedback + assert "good-row" not in feedback + + +def test_oracle_feedback_prioritizes_high_weight_metrics_when_truncated() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + optimizer={ + "eval_metrics": { + "low_value": {"direction": "maximize", "weight": 1.0}, + "high_value": {"direction": "maximize", "weight": 9.0}, + } + }, + prompt={"oracle_feedback_max_chars": 160, "oracle_feedback_worst_n": 1}, + ) + ) + individual = GaIndividual( + prompts={"system_prompt": "Base prompt."}, + generation=0, + individual_index=0, + status="completed", + aggregate_metrics={"low_value": 0.1, "high_value": 0.2}, + raw_scores=( + _score( + task_id="low-row", + metric_name="low_value", + value=0.1, + reasoning="low-weight feedback should not lead the prompt context.", + ), + _score( + task_id="high-row", + metric_name="high_value", + value=0.2, + reasoning="high-weight feedback should lead the prompt context.", + ), + ), + ) + + feedback = build_oracle_feedback(individual=individual, config=config) + + assert feedback is not None + assert feedback.startswith("Metric high_value") + assert len(feedback) <= 160 + + +def test_mutation_feedback_follows_selected_prompt_source() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + prompt={ + "population_size": 2, + "generations": 2, + "crossover_rate": 0.0, + "mutation_rate": 1.0, + "oracle_feedback_mode": "always", + } + ) + ) + transformer = FeedbackRecordingTransformer() + parent_a = _completed_parent(prompt="Parent A", index=0, reasoning="reasoning from A") + parent_b = _completed_parent(prompt="Parent B", index=1, reasoning="reasoning from B") + + prompt = _child_prompt( + transformer, + name="system_prompt", + spec=config.search_space["system_prompt"], + parent_a=parent_a, + parent_b=parent_b, + config=config, + rng=FixedRandom([0.99, 0.0, 0.0]), + oracle_state=OracleFeedbackState(stagnation_generations=0, fitness_variance=0.0, duplicate_ratio=0.0), + failures=[], + ) + + assert prompt == "Parent B" + assert transformer.mutation_feedback is not None + assert "Feedback from g000-i001" in transformer.mutation_feedback + assert "reasoning from B" in transformer.mutation_feedback + assert "reasoning from A" not in transformer.mutation_feedback + + +def test_recombination_feedback_combines_both_parents() -> None: + config = parse_ga_prompt_optimizer_config( + _payload( + prompt={ + "population_size": 2, + "generations": 2, + "crossover_rate": 1.0, + "mutation_rate": 0.0, + "oracle_feedback_mode": "always", + } + ) + ) + transformer = FeedbackRecordingTransformer() + parent_a = _completed_parent(prompt="Parent A", index=0, reasoning="reasoning from A") + parent_b = _completed_parent(prompt="Parent B", index=1, reasoning="reasoning from B") + + prompt = _child_prompt( + transformer, + name="system_prompt", + spec=config.search_space["system_prompt"], + parent_a=parent_a, + parent_b=parent_b, + config=config, + rng=FixedRandom([0.0, 0.99]), + oracle_state=OracleFeedbackState(stagnation_generations=0, fitness_variance=0.0, duplicate_ratio=0.0), + failures=[], + ) + + assert prompt == "Parent A + Parent B" + assert transformer.recombination_feedback is not None + assert "reasoning from A" in transformer.recombination_feedback + assert "reasoning from B" in transformer.recombination_feedback + + +def test_raw_score_records_skip_carried_elites(tmp_path: Path) -> None: + payload = _payload( + prompt={ + "population_size": 2, + "generations": 2, + "mutation_rate": 1.0, + "crossover_rate": 1.0, + "elitism": 1, + "parallel_evaluations": 1, + } + ) + + result = run_ga_prompt_optimization( + payload, + tmp_path, + RawScoreEvaluator(), + DeterministicTransformer(), + ) + + records = json.loads((tmp_path / "ga_score_records.json").read_text(encoding="utf-8")) + assert len(records) == result.executed_trials + carried_ids = {individual.individual_id for individual in result.history if individual.carried_from} + assert carried_ids + assert carried_ids.isdisjoint({record["individual_id"] for record in records}) + + +class RecordingEvaluator: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def evaluate( + self, + *, + trial_number: int, + suggestions: dict[str, Any], + trial_overlay: dict[str, Any], + rep: int, + ) -> CandidateEvaluationResult: + self.calls.append( + { + "trial_number": trial_number, + "suggestions": dict(suggestions), + "metadata": dict(trial_overlay["metadata"]), + "rep": rep, + } + ) + prompt = str(suggestions["instructions.system.content"]) + return CandidateEvaluationResult(aggregate_metrics={"average_score": float(len(prompt))}) + + +class ConcurrencyEvaluator: + def __init__(self) -> None: + self._lock = threading.Lock() + self._active = 0 + self.max_active = 0 + self.trial_numbers: list[int] = [] + + def evaluate( + self, + *, + trial_number: int, + suggestions: dict[str, Any], + trial_overlay: dict[str, Any], + rep: int, + ) -> CandidateEvaluationResult: + del trial_overlay, rep + with self._lock: + self._active += 1 + self.max_active = max(self.max_active, self._active) + self.trial_numbers.append(trial_number) + time.sleep(0.02) + with self._lock: + self._active -= 1 + prompt = str(suggestions["instructions.system.content"]) + return CandidateEvaluationResult(aggregate_metrics={"average_score": float(len(prompt))}) + + +class DeterministicTransformer: + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del purpose, prompt_format, feedback + return f"{prompt}\n\n[{prompt_name}: mutation]" + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del purpose, prompt_format, feedback + return f"{parent_a}\n\n[{prompt_name}: recombined]\n{parent_b}" + + +class FeedbackRecordingTransformer: + def __init__(self) -> None: + self.mutation_feedback: str | None = None + self.recombination_feedback: str | None = None + + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, purpose, prompt_format + self.mutation_feedback = feedback + return prompt + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, purpose, prompt_format + self.recombination_feedback = feedback + return f"{parent_a} + {parent_b}" + + +class FixedRandom: + def __init__(self, values: list[float]) -> None: + self._values = iter(values) + + def random(self) -> float: + return next(self._values) + + +class FailsBadPromptEvaluator: + def evaluate( + self, + *, + trial_number: int, + suggestions: dict[str, Any], + trial_overlay: dict[str, Any], + rep: int, + ) -> CandidateEvaluationResult: + del trial_number, trial_overlay, rep + prompt = str(suggestions["instructions.system.content"]) + if "bad candidate" in prompt: + raise CandidateEvaluationError("candidate crashed") + return CandidateEvaluationResult(aggregate_metrics={"average_score": float(len(prompt))}) + + +class AlwaysFailEvaluator: + def evaluate( + self, + *, + trial_number: int, + suggestions: dict[str, Any], + trial_overlay: dict[str, Any], + rep: int, + ) -> CandidateEvaluationResult: + del trial_number, suggestions, trial_overlay, rep + raise CandidateEvaluationError("all candidates failed") + + +class RawScoreEvaluator: + def evaluate( + self, + *, + trial_number: int, + suggestions: dict[str, Any], + trial_overlay: dict[str, Any], + rep: int, + ) -> CandidateEvaluationResult: + del trial_overlay, rep + prompt = str(suggestions["instructions.system.content"]) + score_value = float(len(prompt)) + return CandidateEvaluationResult( + aggregate_metrics={"average_score": score_value}, + scores=( + _score( + task_id=f"row-{trial_number}", + metric_name="average_score", + value=score_value, + reasoning=f"reasoning for trial {trial_number}", + ), + ), + ) + + +class FailingTransformer: + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, prompt, purpose, prompt_format, feedback + raise PromptTransformError("mutation model failed") + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, parent_a, parent_b, purpose, prompt_format, feedback + raise PromptTransformError("recombination model failed") + + +class OneBadMutationTransformer: + def __init__(self) -> None: + self._mutation_count = 0 + + def mutate( + self, + *, + prompt_name: str, + prompt: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, purpose, prompt_format, feedback + self._mutation_count += 1 + if self._mutation_count == 1: + return f"{prompt}\n\nbad candidate" + return f"{prompt}\n\nbetter candidate" + + def recombine( + self, + *, + prompt_name: str, + parent_a: str, + parent_b: str, + purpose: str, + prompt_format: str | None, + feedback: str | None, + ) -> str: + del prompt_name, parent_b, purpose, prompt_format, feedback + return parent_a + + +def _payload( + *, + optimizer: dict[str, Any] | None = None, + prompt: dict[str, Any] | None = None, +) -> dict[str, Any]: + optimizer_config: dict[str, Any] = { + "eval_metrics": {"average_score": {"direction": "maximize", "weight": 1.0}}, + } + if optimizer: + optimizer_config.update(optimizer) + prompt_config = {"enabled": True, "backend": "ga", "model": "prompt_optimizer"} + if prompt: + prompt_config.update(prompt) + optimizer_config["prompt"] = prompt_config + optimizer_config["search_space"] = { + "system_prompt": { + "type": "fabric", + "path": "instructions.system.content", + "is_prompt": True, + "purpose": "Answer accurately.", + } + } + return { + "schema_version": "fabric.agent/v1alpha1", + "metadata": {"name": "demo"}, + "models": {"prompt_optimizer": {"provider": "openai", "model": "gpt-5-mini"}}, + "instructions": {"system": {"content": "Base prompt."}}, + "optimizer": optimizer_config, + } + + +def _completed_individual(*, prompt: str, score: float, index: int, metric_name: str) -> GaIndividual: + return GaIndividual( + prompts={"system_prompt": prompt}, + generation=0, + individual_index=index, + status="completed", + aggregate_metrics={metric_name: score}, + ) + + +def _completed_parent(*, prompt: str, index: int, reasoning: str) -> GaIndividual: + return GaIndividual( + prompts={"system_prompt": prompt}, + generation=0, + individual_index=index, + status="completed", + aggregate_metrics={"average_score": 0.5}, + fitness=1.0, + raw_scores=( + _score( + task_id=f"row-{index}", + metric_name="average_score", + value=0.5, + reasoning=reasoning, + ), + ), + ) + + +def _score(*, task_id: str, metric_name: str, value: float, reasoning: str) -> AgentEvalTaskScore: + return AgentEvalTaskScore( + id=f"score-{task_id}", + run_id="run", + task_id=task_id, + trial_id=f"trial-{task_id}", + metric_type="tunable-rag-evaluator", + status=AgentEvalScoreStatus.COMPLETED, + outputs=[ + MetricOutput(name=metric_name, value=value), + MetricOutput(name="reasoning", value=reasoning), + ], + ) diff --git a/plugins/nemo-optimization/tests/test_ga_transform.py b/plugins/nemo-optimization/tests/test_ga_transform.py new file mode 100644 index 0000000000..78a47ff174 --- /dev/null +++ b/plugins/nemo-optimization/tests/test_ga_transform.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from typing import Any + +import pytest +from nemo_optimization.backends.ga.transform import ModelPromptTransformer, PromptTransformError + + +def test_model_prompt_transformer_uses_openai_compatible_settings(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + class FakeResponse: + def __enter__(self) -> FakeResponse: + return self + + def __exit__(self, exc_type, exc, traceback) -> None: # noqa: ANN001 + del exc_type, exc, traceback + + def read(self) -> bytes: + return b'{"choices":[{"message":{"content":"```text\\nImproved prompt\\n```"}}]}' + + def fake_urlopen(request, timeout): # noqa: ANN001 + captured["url"] = request.full_url + captured["timeout"] = timeout + captured["authorization"] = request.get_header("Authorization") + captured["body"] = json.loads(request.data.decode("utf-8")) + return FakeResponse() + + monkeypatch.setenv("PROMPT_OPTIMIZER_API_KEY", "test-key") + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + + transformer = ModelPromptTransformer( + payload={ + "models": { + "prompt_optimizer": { + "provider": "openai", + "base_url": "https://example.test/v1", + "model": "gpt-test", + "api_key_env": "PROMPT_OPTIMIZER_API_KEY", + "settings": {"temperature": 0.2, "max_tokens": 77, "timeout_s": 12.0}, + } + } + }, + model_name="prompt_optimizer", + ) + + result = transformer.mutate( + prompt_name="system_prompt", + prompt="Base prompt.", + purpose="Answer accurately.", + prompt_format=None, + feedback=None, + ) + + assert result == "Improved prompt" + assert captured["url"] == "https://example.test/v1/chat/completions" + assert captured["timeout"] == 12.0 + assert captured["authorization"] == "Bearer test-key" + assert captured["body"]["model"] == "gpt-test" + assert captured["body"]["temperature"] == 0.2 + assert captured["body"]["max_tokens"] == 77 + + +def test_model_prompt_transformer_rejects_secret_refs_for_direct_adapter() -> None: + with pytest.raises(PromptTransformError, match="api_key_env, not api_key_secret"): + ModelPromptTransformer( + payload={ + "models": { + "prompt_optimizer": { + "provider": "openai", + "base_url": "https://example.test/v1", + "model": "gpt-test", + "api_key_secret": "PROMPT_OPTIMIZER_SECRET", + } + } + }, + model_name="prompt_optimizer", + ) diff --git a/plugins/nemo-optimization/tests/test_router.py b/plugins/nemo-optimization/tests/test_router.py index 5610d49faa..d8066003d9 100644 --- a/plugins/nemo-optimization/tests/test_router.py +++ b/plugins/nemo-optimization/tests/test_router.py @@ -6,6 +6,7 @@ import json import pytest +from nemo_optimization.backends.ga.config import GaConfigError from nemo_optimization.backends.protocol import ( OptimizationBackend, OptimizationBackendCapabilities, @@ -80,13 +81,22 @@ def test_dispatch_uses_executed_trial_count_for_phase_range(ctx: JobContext) -> assert result["trial_number_range"] == {"start": 0, "end_exclusive": 1, "count": 1} -def test_dispatch_prompt_enabled_returns_failed_phase_result(ctx: JobContext) -> None: +def test_dispatch_prompt_enabled_requires_eval(ctx: JobContext) -> None: payload = { "schema_version": "fabric.agent/v1alpha1", "models": {"prompt_optimizer": {"provider": "openai", "model": "gpt-5-mini"}}, "instructions": {"system": {"content": "Base prompt."}}, "optimizer": { - "prompt": {"enabled": True, "backend": "ga", "model": "prompt_optimizer"}, + "prompt": { + "enabled": True, + "backend": "ga", + "model": "prompt_optimizer", + "population_size": 3, + "generations": 1, + }, + "eval_metrics": { + "average_score": {"direction": "maximize", "weight": 1.0}, + }, "search_space": { "system_prompt": { "type": "fabric", @@ -97,13 +107,9 @@ def test_dispatch_prompt_enabled_returns_failed_phase_result(ctx: JobContext) -> }, }, } - result = OptimizeRouter.dispatch_payload(payload, ctx=ctx) - assert result["status"] == "failed" - assert result["backend"] == "ga" - assert result["phase"] == "prompt" - assert "not supported yet" in result["error"] - assert (ctx.storage.persistent / "results" / "optimizer_results" / "prompt_phase_failure.json").is_file() + with pytest.raises(GaConfigError, match="requires payload.eval"): + OptimizeRouter.dispatch_payload(payload, ctx=ctx) def test_dispatch_prompt_backend_must_support_prompt_phase(ctx: JobContext) -> None: