"""Shared helpers for running benchmark tasks locally or on Modal.""" from __future__ import annotations import time import traceback from dataclasses import dataclass from typing import Any from collections.abc import Callable import warnings import importlib from typing import cast import os import logging from datetime import datetime import importlib.metadata import hashlib from benchmarks.packages import registry as benchmark_config from benchmarks.packages.registry import BenchmarkDatasetConfig import opik_optimizer.datasets from benchmarks.core.types import ( DatasetMetadata, TaskEvaluationResult, EvaluationSet, EvaluationStage, TaskResult, TASK_STATUS_FAILED, TASK_STATUS_SUCCESS, ) from benchmarks.core.types import TaskSpec from opik_optimizer import BaseOptimizer, ChatPrompt from opik_optimizer.utils import reporting as reporting_utils from benchmarks.utils.display import display_preflight_report from benchmarks.utils.logging import console from benchmarks.core.types import ( PreflightContext, PreflightEntry, PreflightReport, PreflightSummary, ) from benchmarks.packages.registry import resolve_package _SPLIT_SUFFIXES = { "train": "_train", "validation": "_validation", "test": "_test", } logger = logging.getLogger(__name__) @dataclass(frozen=True) class DatasetBundle: """Resolved dataset bundle for a benchmark task.""" train_name: str train: Any validation_name: str | None validation: Any | None test_name: str | None test: Any | None evaluation_name: str evaluation_role: str evaluation: Any requested_split: str | None def _parse_base_name(dataset_name: str) -> tuple[str, str | None]: for split, suffix in _SPLIT_SUFFIXES.items(): if dataset_name.endswith(suffix): return dataset_name[: -len(suffix)], split return dataset_name, None def _resolve_initial_prompt(dataset_name: str) -> list[dict[str, Any]]: """ Return an initial prompt for a dataset, tolerating sample-suffixed names. Tries, in order: 1) Exact dataset_name 2) Base name (stripping _train/_validation/_test) 3) Base name + '_train' """ # Exact match if dataset_name in benchmark_config.INITIAL_PROMPTS: return benchmark_config.INITIAL_PROMPTS[dataset_name] base, _ = _parse_base_name(dataset_name) # Base name if base in benchmark_config.INITIAL_PROMPTS: return benchmark_config.INITIAL_PROMPTS[base] # Base train fallback candidate = f"{base}_train" if candidate in benchmark_config.INITIAL_PROMPTS: return benchmark_config.INITIAL_PROMPTS[candidate] raise KeyError(f"No initial prompt configured for dataset '{dataset_name}'") def _load_dataset( dataset_name: str, split: str | None, test_mode: bool, *, dry_run: bool = False ) -> Any: """Load a dataset by name, falling back to base loader when split-specific helpers are absent.""" if dry_run: return None loader = getattr(opik_optimizer.datasets, dataset_name, None) if callable(loader): return loader(test_mode=test_mode) base_name, _ = _parse_base_name(dataset_name) base_loader = getattr(opik_optimizer.datasets, base_name, None) if callable(base_loader): kwargs: dict[str, Any] = {"test_mode": test_mode} if split in ("train", "validation", "test"): kwargs["split"] = split kwargs["dataset_name"] = dataset_name if "prefer_presets" in base_loader.__code__.co_varnames: kwargs["prefer_presets"] = True return base_loader(**kwargs) raise ValueError(f"Unknown dataset loader for '{dataset_name}'.") def resolve_dataset_bundle( dataset_name: str, test_mode: bool, datasets: dict[str, Any] | None = None, *, dry_run: bool = False, ) -> DatasetBundle: """Return train/validation/test dataset objects for a given benchmark dataset key. When ``datasets`` is provided, the loader kwargs (train/validation/test) are used instead of the registered preset slices. If only one override is given, it is reused for all splits (with a warning). """ if datasets: if "train" not in datasets and any( k in datasets for k in ("validation", "test") ): raise ValueError( "datasets config must include a train split when validation/test are provided." ) explicit_roles = any( role in datasets for role in ("train", "validation", "test") ) if explicit_roles: role_specs = datasets else: # If the user provided a single override object, apply it to train only. # Callers should explicitly specify validation/test if they need them. role_specs = {"train": datasets} warnings.warn( "Dataset overrides provided without explicit splits; applying overrides to train only " "and skipping validation/test.", stacklevel=2, ) def _load_override(role: str) -> tuple[str, Any] | tuple[None, None]: spec = role_specs.get(role) if spec is None: return None, None loader_name = spec.get("loader") if isinstance(spec, dict) else None kwargs = dict(spec) if isinstance(spec, dict) else {} loader_name = loader_name or dataset_name kwargs.pop("loader", None) kwargs.setdefault("dataset_name", f"{loader_name}_{role}") if role in ("train", "validation", "test"): kwargs.setdefault("split", role) kwargs["test_mode"] = test_mode loader = getattr(opik_optimizer.datasets, loader_name, None) if callable(loader): return ( kwargs["dataset_name"], None if dry_run else loader(**kwargs), ) raise ValueError( f"Unknown dataset loader '{loader_name}' for role '{role}'." ) train_name, train_ds = _load_override("train") validation_name, validation_ds = _load_override("validation") test_name, test_ds = _load_override("test") evaluation_ds = validation_ds or train_ds evaluation_name = validation_name or train_name or dataset_name evaluation_role = "validation" if validation_ds is not None else "train" return DatasetBundle( train_name=train_name or dataset_name, train=train_ds, validation_name=validation_name, validation=validation_ds, test_name=test_name, test=test_ds, evaluation_name=evaluation_name, evaluation_role=evaluation_role, evaluation=evaluation_ds, requested_split=None, ) base_name, requested_split = _parse_base_name(dataset_name) def _candidate(split: str) -> str | None: candidate_name = f"{base_name}_{split}" return ( candidate_name if candidate_name in benchmark_config.DATASET_CONFIG else None ) train_name = _candidate("train") or dataset_name validation_name = _candidate("validation") test_name = _candidate("test") if train_name not in benchmark_config.DATASET_CONFIG: raise ValueError( f"Dataset '{dataset_name}' is not registered in benchmark_config.DATASET_CONFIG." ) train_dataset = _load_dataset( train_name, "train", test_mode=test_mode, dry_run=dry_run ) validation_dataset = ( _load_dataset( validation_name, "validation", test_mode=test_mode, dry_run=dry_run ) if validation_name else None ) test_dataset = ( _load_dataset(test_name, "test", test_mode=test_mode, dry_run=dry_run) if test_name else None ) evaluation_dataset = validation_dataset or train_dataset evaluation_name = validation_name or train_name evaluation_role = "validation" if validation_dataset is not None else "train" return DatasetBundle( train_name=train_name, train=train_dataset, validation_name=validation_name, validation=validation_dataset, test_name=test_name, test=test_dataset, evaluation_name=evaluation_name, evaluation_role=evaluation_role, evaluation=evaluation_dataset, requested_split=requested_split, ) def _safe_version(pkg: str) -> str | None: try: return importlib.metadata.version(pkg) except Exception: return None def preflight_tasks( task_specs: list[TaskSpec], info: dict[str, Any] | None = None ) -> PreflightReport: """Validate datasets/metrics/optimizers before scheduling to fail fast.""" errors: list[str] = [] had_error = False datasets_seen: set[str] = set() optimizers_seen: set[str] = set() models_seen: set[str] = set() entries: list[PreflightEntry] = [] logger.info("🔎 Preflight: validating %d tasks", len(task_specs)) console.print( f"[bold blue]Preflight:[/bold blue] validating {len(task_specs)} tasks" ) now_iso = datetime.now().isoformat(timespec="seconds") manifest_path = None checkpoint_dir = None run_id = None if info: manifest_path = info.get("manifest_path") checkpoint_dir = info.get("checkpoint_dir") run_id = info.get("run_id") def _role_display( role: str, ds_name: str | None, spec: dict[str, Any] | None, present: bool, ) -> str: # Prefer explicit dataset_name in the manifest, then fall back to the loader, # then whatever name came from the resolved bundle. base = None if spec and isinstance(spec, dict): base = spec.get("dataset_name") or spec.get("loader") base = base or ds_name if spec is not None and isinstance(spec, dict): count = spec.get("count") if count is not None: return f"{role}={base or 'None'}({count})" if not present: return f"{role}=None" return f"{role}={base or 'None'}" def _format_splits(bundle: DatasetBundle, task: TaskSpec) -> str: """Human-friendly split summary with dataset names and counts.""" tokens: list[str] = [] train_spec = task.datasets.get("train") if task.datasets else None val_spec = task.datasets.get("validation") if task.datasets else None test_spec = task.datasets.get("test") if task.datasets else None tokens.append( _role_display( "train", bundle.train_name, train_spec, bundle.train is not None ) ) tokens.append( _role_display( "val", bundle.validation_name, val_spec, bundle.validation is not None ) ) tokens.append( _role_display("test", bundle.test_name, test_spec, bundle.test is not None) ) return ", ".join(tokens) for task in task_specs: if task.optimizer_name not in benchmark_config.OPTIMIZER_CONFIGS: msg = f"Unknown optimizer '{task.optimizer_name}'" logger.error(msg) errors.append(msg) had_error = True entries.append( PreflightEntry( task_id=task.task_id, short_id=hashlib.sha1( f"{run_id or 'run'}:{task.task_id}".encode() ).hexdigest()[:5], dataset_name=task.dataset_name, evaluation_name=None, optimizer_name=task.optimizer_name, model_name=task.model_name, status="error", splits=None, error=msg, ) ) continue try: bundle = resolve_dataset_bundle( dataset_name=task.dataset_name, test_mode=task.test_mode, datasets=task.datasets, dry_run=True, ) split_summary = _format_splits(bundle, task) dataset_config = benchmark_config.DATASET_CONFIG.get( bundle.evaluation_name, benchmark_config.DATASET_CONFIG.get(task.dataset_name), ) if dataset_config is None: raise ValueError( f"Dataset '{task.dataset_name}' is not registered in benchmark_config.DATASET_CONFIG." ) _resolve_metrics( dataset_config, cast(list[str | dict[str, Any]] | None, task.metrics), ) datasets_seen.add(bundle.evaluation_name or task.dataset_name) optimizers_seen.add(task.optimizer_name) models_seen.add(task.model_name) short_id = hashlib.sha1( f"{run_id or 'run'}:{task.task_id}".encode() ).hexdigest()[:5] entries.append( PreflightEntry( task_id=task.task_id, short_id=short_id, dataset_name=task.dataset_name, evaluation_name=bundle.evaluation_name, optimizer_name=task.optimizer_name, model_name=task.model_name, status="ok", splits=split_summary, error=None, ) ) logger.info( "✅ Preflight ok: dataset=%s (eval=%s) optimizer=%s model=%s", task.dataset_name, bundle.evaluation_name, task.optimizer_name, task.model_name, ) except Exception as exc: err = f"Preflight failed for dataset '{task.dataset_name}': {exc}" logger.error(err) errors.append(err) had_error = True short_id = hashlib.sha1( f"{run_id or 'run'}:{task.task_id}".encode() ).hexdigest()[:5] entries.append( PreflightEntry( task_id=task.task_id, short_id=short_id, dataset_name=task.dataset_name, evaluation_name=None, optimizer_name=task.optimizer_name, model_name=task.model_name, status="error", splits=None, error=str(exc), ) ) summary = PreflightSummary( total_tasks=len(task_specs), datasets=sorted(datasets_seen), optimizers=sorted(optimizers_seen), models=sorted(models_seen), errors=errors, ) report = PreflightReport( context=PreflightContext( system_time=now_iso, cwd=os.getcwd(), manifest_path=manifest_path, checkpoint_dir=checkpoint_dir, run_id=run_id, opik_version=_safe_version("opik"), opik_optimizer_version=_safe_version("opik-optimizer"), ), summary=summary, entries=entries, ) display_preflight_report(report, had_error=had_error, console=console) if had_error: raise ValueError("Benchmark preflight checks failed:\n- " + "\n- ".join(errors)) return report def _dataset_metadata(dataset: Any, dataset_name: str, role: str) -> DatasetMetadata: return DatasetMetadata( name=getattr(dataset, "name", dataset_name), id=getattr(dataset, "id", None), split=role, ) def _resolve_metrics( dataset_config: BenchmarkDatasetConfig, custom_metrics: list[str | dict[str, Any]] | None, ) -> list[Callable]: if not custom_metrics: return dataset_config.metrics resolved: list[Callable] = [] for entry in custom_metrics: path: str args: list[Any] = [] kwargs: dict[str, Any] = {} if isinstance(entry, str): path = entry elif isinstance(entry, dict): path = str(entry.get("path")) args = entry.get("args", []) or [] kwargs = entry.get("kwargs", {}) or {} else: raise ValueError( "Metrics entries must be strings or objects with a 'path' key." ) module_path, _, attr = path.rpartition(".") if not module_path or not attr: raise ValueError( f"Invalid metric path '{path}'. Expected module.attr format." ) module = importlib.import_module(module_path) metric_obj = getattr(module, attr, None) if metric_obj is None: raise ValueError(f"Metric '{path}' not found.") metric_fn = ( metric_obj(*args, **kwargs) if (args or kwargs) and callable(metric_obj) else metric_obj ) if not callable(metric_fn): raise ValueError( f"Metric '{path}' is not callable after applying args/kwargs." ) resolved.append(metric_fn) return resolved def collect_dataset_metadata(bundle: DatasetBundle) -> dict[str, DatasetMetadata]: """Build a metadata map keyed by split name.""" metadata = {"train": _dataset_metadata(bundle.train, bundle.train_name, "train")} if bundle.validation is not None and bundle.validation_name: metadata["validation"] = _dataset_metadata( bundle.validation, bundle.validation_name, "validation" ) if bundle.test is not None and bundle.test_name: metadata["test"] = _dataset_metadata(bundle.test, bundle.test_name, "test") return metadata def evaluate_prompt_on_dataset( *, optimizer: BaseOptimizer, prompt: ChatPrompt | dict[str, ChatPrompt], dataset: Any, dataset_name: str, dataset_role: str, metrics: list[Callable], n_threads: int, agent: Any | None = None, ) -> TaskEvaluationResult: """Run all metrics for a prompt on a dataset and return a structured result. Args: optimizer: The optimizer instance to use for evaluation. prompt: Either a single ChatPrompt or a dict of ChatPrompts for multi-prompt agents. dataset: The dataset to evaluate on. dataset_name: Name of the dataset. dataset_role: Role of the dataset (train/validation/test). metrics: List of metric functions to evaluate. n_threads: Number of threads for parallel evaluation. agent: Optional agent instance for multi-prompt/compound AI system evaluation. """ start_time = time.time() metric_entries = [] for metric_fn in metrics: metric_name = getattr(metric_fn, "__name__", metric_fn.__class__.__name__) score = optimizer.evaluate_prompt( prompt=prompt, dataset=dataset, metric=metric_fn, n_threads=n_threads, agent=agent, ) metric_entries.append( { "metric_name": metric_name, "score": score, "timestamp": time.time(), } ) return TaskEvaluationResult( dataset=_dataset_metadata(dataset, dataset_name, dataset_role), metrics=metric_entries, # type: ignore[arg-type] duration_seconds=time.time() - start_time, ) def _serialize_optimization_result(result: Any) -> Any: """Serialize an OptimizationResult to a dict for storage. Handles both Pydantic v2 (model_dump) and v1 (dict) models, as well as plain dicts or other serializable objects. """ if hasattr(result, "model_dump"): # Pydantic v2 try: return result.model_dump(mode="json") except Exception: # Fallback to default mode if json mode fails return result.model_dump() if hasattr(result, "dict"): # Pydantic v1 return result.dict() if isinstance(result, dict): return result return result def execute_task( *, task_id: str, dataset_name: str, optimizer_name: str, model_name: str, model_parameters: dict[str, Any] | None, test_mode: bool, optimizer_params_override: dict[str, Any] | None, optimizer_prompt_params_override: dict[str, Any] | None, datasets: dict[str, Any] | None = None, metrics: list[str | dict[str, Any]] | None = None, prompt_messages: list[dict[str, Any]] | None = None, ) -> TaskResult: """Shared execution path used by local and Modal runners.""" timestamp_start = time.time() initial_prompt: ChatPrompt | dict[str, ChatPrompt] | None = None optimized_prompt: ChatPrompt | dict[str, ChatPrompt] | None = None optimize_kwargs: dict[str, Any] | None = None constructor_kwargs: dict[str, Any] | None = None test_initial_evaluation: TaskEvaluationResult | None = None steps: list[dict[str, Any]] = [] history_state = opik_optimizer.core.results.OptimizationHistoryState() def _record_step(step: dict[str, Any]) -> None: round_handle = history_state.start_round( round_index=step.get("index"), extras={ "step_id": step.get("step_id"), "kind": step.get("kind"), }, ) metrics_payload = step.get("metrics") or {} score_value: float | None = None if isinstance(metrics_payload, dict): for split_metrics in metrics_payload.values(): candidate_metrics = None if isinstance(split_metrics, dict) and split_metrics: candidate_metrics = split_metrics elif isinstance(split_metrics, (list, tuple)) and split_metrics: first_entry = split_metrics[0] if isinstance(first_entry, dict) and first_entry: candidate_metrics = first_entry if candidate_metrics: first_value = next(iter(candidate_metrics.values())) try: score_value = float(first_value) except Exception: score_value = None break history_state.record_trial( round_handle=round_handle, score=score_value, candidate=step.get("prompt_snapshot"), dataset_split=step.get("split"), extras={ "metrics": metrics_payload, "llm_calls": step.get("llm_calls"), "meta": step.get("meta"), }, ) history_state.end_round( round_handle=round_handle, best_score=score_value, best_prompt=step.get("prompt_snapshot"), ) with reporting_utils.suppress_opik_logs(): try: if test_mode and os.getenv("OPIK_DATASET_SKIP_EXISTING") is None: # Avoid brittle failures in smoke runs when datasets already exist with different sizes. os.environ["OPIK_DATASET_SKIP_EXISTING"] = "true" bundle = resolve_dataset_bundle( dataset_name=dataset_name, test_mode=test_mode, datasets=datasets, ) # Resolve dataset config defensively, tolerating base names and override names. _dataset_config: BenchmarkDatasetConfig | None = None for candidate in ( bundle.evaluation_name, dataset_name, bundle.train_name, bundle.validation_name, bundle.test_name, ): if candidate and candidate in benchmark_config.DATASET_CONFIG: _dataset_config = benchmark_config.DATASET_CONFIG[candidate] break if _dataset_config is None: raise KeyError( f"Dataset config not found for '{dataset_name}' " f"(evaluation={bundle.evaluation_name}, train={bundle.train_name}, " f"validation={bundle.validation_name}, test={bundle.test_name}). " "Ensure the dataset is registered in benchmark_config.DATASET_CONFIG." ) dataset_config = _dataset_config metrics_resolved = _resolve_metrics( dataset_config, cast(list[str | dict[str, Any]] | None, metrics) ) if not metrics_resolved: raise ValueError( f"No metrics configured for dataset '{dataset_config.name}'. " "Provide at least one metric via dataset config or manifest overrides." ) optimizer_config = benchmark_config.OPTIMIZER_CONFIGS[optimizer_name] constructor_kwargs = dict(optimizer_config.params) if optimizer_params_override: constructor_kwargs.update(optimizer_params_override) # Ensure we only supply model/model_parameters once. If provided in the config/overrides, # respect those; otherwise inject the runner-specified values. constructor_kwargs.setdefault("model", model_name) constructor_kwargs.setdefault("model_parameters", model_parameters) optimizer: BaseOptimizer = getattr( opik_optimizer, optimizer_config.class_name )( **constructor_kwargs, ) # Check if this dataset is handled by a package-specific agent/prompt wiring. agent = None package_resolution = resolve_package(dataset_name) if package_resolution is not None: package = package_resolution.package agent = package.build_agent( model_name=model_name, model_parameters=model_parameters, ) package_prompt = package.build_initial_prompt() if package_prompt is not None: initial_prompt = package_prompt logger.info( "Resolved package %s for dataset %s", package_resolution.key, dataset_name, ) # Standard single-prompt benchmark fallback if initial_prompt is None: messages = prompt_messages or _resolve_initial_prompt(bundle.train_name) # Bind the optimizer's model/model_parameters to the prompt so evaluations # use the requested model instead of ChatPrompt defaults. initial_prompt = ChatPrompt( messages=messages, # type: ignore[arg-type] model=getattr(optimizer, "model", model_name), model_parameters=getattr( optimizer, "model_parameters", model_parameters ), ) if initial_prompt is None: raise RuntimeError( f"Failed to initialize benchmark prompt for dataset '{dataset_name}'." ) initial_evaluation = evaluate_prompt_on_dataset( optimizer=optimizer, prompt=initial_prompt, dataset=bundle.evaluation, dataset_name=bundle.evaluation_name, dataset_role=bundle.evaluation_role, metrics=metrics_resolved, n_threads=4, agent=agent, ) steps.append( { "step_id": "initial-eval", "kind": "baseline", "index": 0, "split": bundle.evaluation_role, "prompt_snapshot": initial_prompt, "metrics": {bundle.evaluation_role: initial_evaluation.metrics}, "llm_calls": 0, "meta": {}, } ) _record_step(steps[-1]) if bundle.test is not None and bundle.test_name is not None: test_initial_evaluation = evaluate_prompt_on_dataset( optimizer=optimizer, prompt=initial_prompt, dataset=bundle.test, dataset_name=bundle.test_name, dataset_role="test", metrics=metrics_resolved, n_threads=4, agent=agent, ) steps.append( { "step_id": "initial-test", "kind": "baseline", "index": 0, "split": "test", "prompt_snapshot": initial_prompt, "metrics": {"test": test_initial_evaluation.metrics}, "llm_calls": 0, "meta": {}, } ) _record_step(steps[-1]) optimize_kwargs = dict(optimizer_config.optimizer_prompt_params) if optimizer_prompt_params_override: optimize_kwargs.update(optimizer_prompt_params_override) optimization_results = optimizer.optimize_prompt( prompt=initial_prompt, dataset=bundle.train, validation_dataset=bundle.validation, metric=metrics_resolved[0], agent=agent, **optimize_kwargs, ) # Handle the optimized prompt - may be dict for multi-prompt agents result_prompt = optimization_results.prompt if isinstance(result_prompt, dict): optimized_prompt = result_prompt elif isinstance(result_prompt, ChatPrompt): optimized_prompt = result_prompt else: optimized_prompt = ChatPrompt( messages=result_prompt, # type: ignore[arg-type] model=getattr(optimizer, "model", model_name), model_parameters=getattr( optimizer, "model_parameters", model_parameters ), ) optimized_evaluation = evaluate_prompt_on_dataset( optimizer=optimizer, prompt=optimized_prompt, dataset=bundle.evaluation, dataset_name=bundle.evaluation_name, dataset_role=bundle.evaluation_role, metrics=metrics_resolved, n_threads=4, agent=agent, ) steps.append( { "step_id": "final-eval", "kind": "post_opt", "index": 1, "split": bundle.evaluation_role, "prompt_snapshot": optimized_prompt, "metrics": {bundle.evaluation_role: optimized_evaluation.metrics}, "llm_calls": optimization_results.llm_calls or 0, "meta": {}, } ) _record_step(steps[-1]) test_evaluation = None if bundle.test is not None and bundle.test_name is not None: test_evaluation = evaluate_prompt_on_dataset( optimizer=optimizer, prompt=optimized_prompt, dataset=bundle.test, dataset_name=bundle.test_name, dataset_role="test", metrics=metrics_resolved, n_threads=4, agent=agent, ) steps.append( { "step_id": "final-test", "kind": "post_opt", "index": 1, "split": "test", "prompt_snapshot": optimized_prompt, "metrics": {"test": test_evaluation.metrics}, "llm_calls": 0, "meta": {}, } ) _record_step(steps[-1]) evaluations = { "initial": EvaluationSet( **{ bundle.evaluation_role: EvaluationSet.EvaluationEntry( step_id="initial-eval", result=initial_evaluation ), "test": EvaluationSet.EvaluationEntry( step_id="initial-test", result=test_initial_evaluation ) if test_initial_evaluation else None, } ), "final": EvaluationSet( **{ bundle.evaluation_role: EvaluationSet.EvaluationEntry( step_id="final-eval", result=optimized_evaluation ), "test": EvaluationSet.EvaluationEntry( step_id="final-test", result=test_evaluation ) if test_evaluation else None, } ), } stages: list[EvaluationStage] = [] stages.append( EvaluationStage( stage="initial", split=bundle.evaluation_role, evaluation=initial_evaluation, prompt_snapshot=initial_prompt, step_ref="initial-eval", ) ) if test_initial_evaluation: stages.append( EvaluationStage( stage="initial", split="test", evaluation=test_initial_evaluation, prompt_snapshot=initial_prompt, step_ref="initial-test", ) ) stages.append( EvaluationStage( stage="final", split=bundle.evaluation_role, evaluation=optimized_evaluation, prompt_snapshot=optimized_prompt, step_ref="final-eval", ) ) if test_evaluation: stages.append( EvaluationStage( stage="final", split="test", evaluation=test_evaluation, prompt_snapshot=optimized_prompt, step_ref="final-test", ) ) return TaskResult( id=task_id, dataset_name=dataset_name, optimizer_name=optimizer_name, model_name=model_name, status=TASK_STATUS_SUCCESS, timestamp_start=timestamp_start, initial_prompt=initial_prompt, optimized_prompt=optimized_prompt, evaluations=evaluations, stages=stages, optimization_history={"rounds": optimization_results.history}, error_message=None, llm_calls_total_optimization=optimization_results.llm_calls or 0, optimization_raw_result=optimization_results, optimization_summary=_serialize_optimization_result( optimization_results ), timestamp_end=time.time(), dataset_metadata=collect_dataset_metadata(bundle), evaluation_split=bundle.evaluation_role, requested_split=bundle.requested_split, optimizer_prompt_params_used=optimize_kwargs, optimizer_params_used=constructor_kwargs, ) except Exception: if test_mode: raise return TaskResult( id=task_id, dataset_name=dataset_name, optimizer_name=optimizer_name, model_name=model_name, status=TASK_STATUS_FAILED, timestamp_start=timestamp_start, initial_prompt=initial_prompt, optimized_prompt=optimized_prompt, error_message=traceback.format_exc(), timestamp_end=time.time(), dataset_metadata={}, evaluation_split=None, requested_split=None, evaluations={}, stages=[], optimization_history={"rounds": history_state.get_entries()}, optimizer_prompt_params_used=optimize_kwargs, optimizer_params_used=constructor_kwargs, )