orizu 0.3.6 → 0.4.0

package/vendor/orizu-gepa-python/src/orizu_gepa/optimizer.py

@@ -4,9 +4,13 @@ import dataclasses
 import hashlib
 import json
 import math
+import os
 import random
+import subprocess
+import sys
 import time
 import uuid
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 from dataclasses import dataclass, field
 from typing import Any, Callable, Protocol
 
@@ -40,6 +44,31 @@ Return only the complete updated parameter text, exactly as it should be used by
 Do not include analysis, explanation, labels, XML tags, or a surrounding markdown code fence.
 Your response will be used verbatim as the next candidate."""
 
+AUTO_NUM_THREADS_HARD_CAP = 64
+AUTO_NUM_THREADS_WORKER_MEMORY_BYTES = 512 * 1024 * 1024
+AUTO_NUM_THREADS_MEMORY_RESERVE_BYTES = 1024 * 1024 * 1024
+AUTO_NUM_THREADS_FD_HEADROOM = 64
+AUTO_NUM_THREADS_FDS_PER_WORKER = 16
+
+
+@dataclass(frozen=True)
+class NumThreadsPlan:
+    requested: int | str
+    resolved: int
+    row_bound: int
+    cpu_bound: int
+    memory_bound: int | None
+    fd_bound: int | None
+    hard_cap: int
+    worker_memory_bytes: int
+    available_memory_bytes: int | None
+    total_memory_bytes: int | None
+    memory_reserve_bytes: int | None
+    limiting_factor: str
+
+    def to_payload(self) -> dict[str, Any]:
+        return dataclasses.asdict(self)
+
 
 @dataclass(frozen=True)
 class DatasetRow:
@@ -113,29 +142,128 @@ class ReflectionResult:
     candidate_text: str
 
 
+class RetryableReflectionError(RuntimeError):
+    """Raised when a transient reflection-provider failure remains after retries."""
+
+
+def _dspy_auto_metric_budget(
+    *,
+    num_components: int,
+    num_candidates: int,
+    valset_size: int,
+    minibatch_size: int = 35,
+    full_eval_steps: int = 5,
+) -> int:
+    if num_candidates <= 0:
+        raise ValueError("num_candidates must be > 0.")
+    if num_components < 0 or valset_size < 0 or minibatch_size < 0:
+        raise ValueError("num_components, valset_size, and minibatch_size must be >= 0.")
+    if full_eval_steps < 1:
+        raise ValueError("full_eval_steps must be >= 1.")
+
+    num_trials = int(max(2 * (num_components * 2) * math.log2(num_candidates), 1.5 * num_candidates))
+    total = valset_size
+    total += num_candidates * 5
+    total += num_trials * minibatch_size
+    if num_trials == 0:
+        return total
+
+    periodic_fulls = (num_trials + 1) // full_eval_steps + 1
+    extra_final = 1 if num_trials < full_eval_steps else 0
+    total += (periodic_fulls + extra_final) * valset_size
+    return total
+
+
+def _approx_metric_calls_for_candidate_budget(
+    *,
+    candidate_proposals: int,
+    trainset_size: int,
+    valset_size: int,
+    minibatch_size: int,
+) -> int:
+    """Estimate row-level metric calls for proposal-count budgets."""
+    safe_proposals = max(0, candidate_proposals)
+    safe_valset_size = max(0, valset_size)
+    safe_trainset_size = max(0, trainset_size)
+    safe_minibatch_size = max(0, minibatch_size)
+    sampled_train_rows = min(safe_trainset_size, safe_minibatch_size)
+    per_iteration = (2 * sampled_train_rows) + safe_valset_size
+    return safe_valset_size + (safe_proposals * per_iteration)
+
+
 @dataclass
 class Budget:
     budget_kind: str
     limit: int
+    approx_metric_call_limit: int | None = None
     used_metric_calls: int = 0
+    used_reflection_failure_metric_charges: int = 0
     used_full_evals: int = 0
     used_candidate_proposals: int = 0
+    used_iterations: int = 0
+
+    def __post_init__(self) -> None:
+        if self.approx_metric_call_limit is None and self.budget_kind == "max_metric_calls":
+            self.approx_metric_call_limit = self.limit
 
     @classmethod
-    def from_config(cls, config: "TextGepaConfig") -> "Budget":
+    def from_config(
+        cls,
+        config: "TextGepaConfig",
+        *,
+        trainset_size: int = 0,
+        valset_size: int = 0,
+        num_components: int = 1,
+    ) -> "Budget":
         if config.max_metric_calls is not None:
-            return cls("max_metric_calls", config.max_metric_calls)
+            return cls(
+                "max_metric_calls",
+                config.max_metric_calls,
+                approx_metric_call_limit=config.max_metric_calls,
+            )
         if config.max_full_evals is not None:
-            return cls("max_full_evals", config.max_full_evals)
+            metric_limit = config.max_full_evals * (trainset_size + valset_size)
+            return cls(
+                "max_metric_calls",
+                metric_limit,
+                approx_metric_call_limit=metric_limit,
+            )
         if config.max_candidate_proposals is not None:
-            return cls("max_candidate_proposals", config.max_candidate_proposals)
-        preset_metric_calls = {
-            "auto": 64,
-            "light": 48,
-            "medium": 96,
-            "high": 192,
-        }.get(config.budget, 48)
-        return cls("max_metric_calls", preset_metric_calls)
+            return cls(
+                "max_candidate_proposals",
+                config.max_candidate_proposals,
+                approx_metric_call_limit=_approx_metric_calls_for_candidate_budget(
+                    candidate_proposals=config.max_candidate_proposals,
+                    trainset_size=trainset_size,
+                    valset_size=valset_size,
+                    minibatch_size=config.minibatch_size,
+                ),
+            )
+        if config.max_iterations is not None:
+            return cls("max_iterations", config.max_iterations)
+        if not config.budget:
+            raise ValueError("No optimization budget configured.")
+        preset = "medium" if config.budget == "auto" else config.budget
+        preset_candidates = {
+            "light": 6,
+            "medium": 12,
+            "heavy": 18,
+            "high": 18,
+        }.get(preset, 6)
+        metric_limit = _dspy_auto_metric_budget(
+            num_components=num_components,
+            num_candidates=preset_candidates,
+            valset_size=valset_size,
+        )
+        return cls(
+            "max_metric_calls",
+            metric_limit,
+            approx_metric_call_limit=metric_limit,
+        )
+
+    @property
+    def metric_budget_used(self) -> int:
+        return self.used_metric_calls + self.used_reflection_failure_metric_charges
 
     @property
     def used(self) -> int:
@@ -143,17 +271,41 @@ class Budget:
             return self.used_full_evals
         if self.budget_kind == "max_candidate_proposals":
             return self.used_candidate_proposals
-        return self.used_metric_calls
+        if self.budget_kind == "max_iterations":
+            return self.used_iterations
+        return self.metric_budget_used
 
     @property
     def remaining(self) -> int:
         return max(0, self.limit - self.used)
 
+    @property
+    def metric_calls_remaining(self) -> int:
+        if self.approx_metric_call_limit is None:
+            return 0
+        if self.remaining == 0:
+            return 0
+        return max(0, self.approx_metric_call_limit - self.metric_budget_used)
+
+    @property
+    def progress_percent(self) -> float:
+        if self.remaining == 0:
+            return 100.0
+        if self.approx_metric_call_limit is None:
+            if self.limit <= 0:
+                return 100.0 if self.used > 0 else 0.0
+            return min(100.0, max(0.0, (self.used / self.limit) * 100.0))
+        if self.approx_metric_call_limit <= 0:
+            # A zero-sized budget with any recorded metric work should still read as complete.
+            return 100.0 if self.metric_budget_used > 0 else 0.0
+        return min(100.0, max(0.0, (self.metric_budget_used / self.approx_metric_call_limit) * 100.0))
+
     def allows_iteration(self) -> bool:
         return self.remaining > 0
 
-    def allows_metric_call(self) -> bool:
-        return self.budget_kind != "max_metric_calls" or self.remaining > 0
+    def charge_retryable_reflection_failure(self) -> None:
+        if self.budget_kind == "max_metric_calls":
+            self.used_reflection_failure_metric_charges += 1
 
     def to_payload(self) -> dict[str, Any]:
         return {
@@ -161,17 +313,45 @@ class Budget:
             "limit": self.limit,
             "used": self.used,
             "remaining": self.remaining,
+            "approx_metric_call_limit": self.approx_metric_call_limit,
+            "metric_call_budget": self.approx_metric_call_limit,
+            "metric_calls_remaining": self.metric_calls_remaining,
+            "metric_budget_used": self.metric_budget_used,
             "used_metric_calls": self.used_metric_calls,
+            "used_reflection_failure_metric_charges": self.used_reflection_failure_metric_charges,
             "used_full_evals": self.used_full_evals,
             "used_candidate_proposals": self.used_candidate_proposals,
+            "used_iterations": self.used_iterations,
+            "iteration_budget": self.limit if self.budget_kind == "max_iterations" else None,
+            "iterations_remaining": self.remaining if self.budget_kind == "max_iterations" else None,
+        }
+
+    def progress_payload(self, *, stage: str, iteration: int | None) -> dict[str, Any]:
+        percent = round(self.progress_percent, 1)
+        return {
+            "stage": stage,
+            "iteration": iteration,
+            "percent": percent,
+            "progress_percent": percent,
+            "metric_calls_used": self.used_metric_calls,
+            "metric_budget_used": self.metric_budget_used,
+            "metric_call_budget": self.approx_metric_call_limit,
+            "approx_metric_call_budget": self.approx_metric_call_limit,
+            "metric_calls_remaining": self.metric_calls_remaining,
+            "is_over_budget": (
+                self.approx_metric_call_limit is not None and
+                self.metric_budget_used > self.approx_metric_call_limit
+            ),
+            "budget": self.to_payload(),
         }
 
 
 @dataclass(frozen=True)
 class TextGepaConfig:
-    budget: str = "light"
-    max_iterations: int = 3
+    budget: str | None = "auto"
+    max_iterations: int | None = None
     minibatch_size: int = 3
+    num_threads: int | str = "auto"
     candidate_selection_strategy: str = "pareto"
     epsilon: float = 0.1
     max_metric_calls: int | None = None
@@ -180,6 +360,8 @@ class TextGepaConfig:
     reflection_model: str = "anthropic/claude-opus-4-7"
     reflection_temperature: float | None = None
     reflection_max_tokens: int | None = None
+    reflection_retry_attempts: int = 3
+    reflection_http_timeout_seconds: int = 180
     reflection_prompt_template: str | None = None
     reflection_provider_settings: dict[str, Any] = field(default_factory=dict)
     objective: str = "Improve this text candidate to maximize evaluator score while preserving intended behavior."
@@ -316,6 +498,230 @@ def _metric_key(scorer_context: PromptContext) -> str:
     return scorer_context.metric_key or "score"
 
 
+def _positive_int_env(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        value = int(raw)
+    except ValueError:
+        return default
+    return value if value > 0 else default
+
+
+def _proc_meminfo_bytes(key: str) -> int | None:
+    try:
+        with open("/proc/meminfo", encoding="utf-8") as handle:
+            for line in handle:
+                if not line.startswith(f"{key}:"):
+                    continue
+                parts = line.split()
+                if len(parts) >= 2:
+                    return int(parts[1]) * 1024
+    except (OSError, ValueError):
+        return None
+    return None
+
+
+def _sysconf_memory_bytes(pages_name: str) -> int | None:
+    try:
+        pages = os.sysconf(pages_name)
+        page_size = os.sysconf("SC_PAGE_SIZE")
+    except (AttributeError, OSError, ValueError):
+        return None
+    if not isinstance(pages, int) or not isinstance(page_size, int):
+        return None
+    if pages <= 0 or page_size <= 0:
+        return None
+    return pages * page_size
+
+
+def _darwin_available_memory_bytes() -> int | None:
+    if sys.platform != "darwin":
+        return None
+    try:
+        output = subprocess.check_output(["vm_stat"], text=True)
+    except (OSError, subprocess.SubprocessError):
+        return None
+
+    lines = output.splitlines()
+    if not lines:
+        return None
+    page_size_marker = "page size of "
+    if page_size_marker not in lines[0]:
+        return None
+    try:
+        page_size = int(lines[0].split(page_size_marker, 1)[1].split()[0])
+    except (IndexError, ValueError):
+        return None
+
+    available_pages = 0
+    available_keys = {
+        "Pages free",
+        "Pages inactive",
+        "Pages speculative",
+        "Pages purgeable",
+    }
+    for line in lines[1:]:
+        key, separator, value = line.partition(":")
+        if separator != ":" or key not in available_keys:
+            continue
+        normalized_value = value.strip().rstrip(".").replace(",", "")
+        try:
+            available_pages += int(normalized_value)
+        except ValueError:
+            continue
+    return available_pages * page_size if available_pages > 0 else None
+
+
+def _available_memory_bytes() -> int | None:
+    return (
+        _proc_meminfo_bytes("MemAvailable")
+        or _darwin_available_memory_bytes()
+        or _sysconf_memory_bytes("SC_AVPHYS_PAGES")
+        or _total_memory_bytes()
+    )
+
+
+def _total_memory_bytes() -> int | None:
+    return _proc_meminfo_bytes("MemTotal") or _sysconf_memory_bytes("SC_PHYS_PAGES")
+
+
+def _open_file_soft_limit() -> int | None:
+    try:
+        import resource
+    except ImportError:
+        return None
+    try:
+        soft_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+    except (OSError, ValueError):
+        return None
+    if soft_limit <= 0 or soft_limit == resource.RLIM_INFINITY:
+        return None
+    return int(soft_limit)
+
+
+def _parse_requested_num_threads(requested: int | str) -> int | str:
+    if isinstance(requested, int):
+        if requested <= 0:
+            raise ValueError("num_threads must be positive or 'auto'")
+        return requested
+    if isinstance(requested, str):
+        normalized = requested.strip().lower()
+        if normalized == "auto":
+            return "auto"
+        try:
+            parsed = int(normalized)
+        except ValueError as exc:
+            raise ValueError("num_threads must be positive or 'auto'") from exc
+        if parsed <= 0:
+            raise ValueError("num_threads must be positive or 'auto'")
+        return parsed
+    raise ValueError("num_threads must be positive or 'auto'")
+
+
+def resolve_num_threads(
+    requested: int | str,
+    *,
+    minibatch_size: int,
+    validation_count: int,
+    cpu_count: int | None = None,
+    available_memory_bytes: int | None = None,
+    total_memory_bytes: int | None = None,
+    fd_limit: int | None = None,
+    hard_cap: int | None = None,
+    worker_memory_bytes: int | None = None,
+) -> NumThreadsPlan:
+    """Resolve an explicit or automatic row-evaluation thread count for one GEPA run."""
+    parsed = _parse_requested_num_threads(requested)
+    row_bound = max(1, max(minibatch_size, validation_count))
+    detected_cpu_count = cpu_count if cpu_count is not None else os.cpu_count()
+    normalized_cpu_count = max(1, detected_cpu_count or 1)
+    cpu_bound = normalized_cpu_count * 2
+    resolved_hard_cap = hard_cap or _positive_int_env(
+        "ORIZU_GEPA_AUTO_THREADS_MAX",
+        AUTO_NUM_THREADS_HARD_CAP,
+    )
+    resolved_worker_memory_bytes = worker_memory_bytes or (
+        _positive_int_env(
+            "ORIZU_GEPA_WORKER_MEMORY_MB",
+            AUTO_NUM_THREADS_WORKER_MEMORY_BYTES // (1024 * 1024),
+        ) * 1024 * 1024
+    )
+
+    detected_total_memory_bytes = (
+        total_memory_bytes if total_memory_bytes is not None else _total_memory_bytes()
+    )
+    detected_available_memory_bytes = (
+        available_memory_bytes
+        if available_memory_bytes is not None
+        else _available_memory_bytes()
+    )
+    memory_reserve_bytes = None
+    memory_bound = None
+    if detected_available_memory_bytes is not None:
+        total_reserve = (
+            detected_total_memory_bytes // 4
+            if detected_total_memory_bytes is not None
+            else 0
+        )
+        memory_reserve_bytes = max(AUTO_NUM_THREADS_MEMORY_RESERVE_BYTES, total_reserve)
+        usable_memory_bytes = max(0, detected_available_memory_bytes - memory_reserve_bytes)
+        memory_bound = max(1, usable_memory_bytes // max(1, resolved_worker_memory_bytes))
+
+    detected_fd_limit = fd_limit if fd_limit is not None else _open_file_soft_limit()
+    fd_bound = None
+    if detected_fd_limit is not None:
+        usable_fds = max(0, detected_fd_limit - AUTO_NUM_THREADS_FD_HEADROOM)
+        fd_bound = max(1, usable_fds // AUTO_NUM_THREADS_FDS_PER_WORKER)
+
+    if parsed != "auto":
+        return NumThreadsPlan(
+            requested=requested,
+            resolved=parsed,
+            row_bound=row_bound,
+            cpu_bound=cpu_bound,
+            memory_bound=memory_bound,
+            fd_bound=fd_bound,
+            hard_cap=resolved_hard_cap,
+            worker_memory_bytes=resolved_worker_memory_bytes,
+            available_memory_bytes=detected_available_memory_bytes,
+            total_memory_bytes=detected_total_memory_bytes,
+            memory_reserve_bytes=memory_reserve_bytes,
+            limiting_factor="explicit",
+        )
+
+    limits = {
+        "rows": row_bound,
+        "cpu": cpu_bound,
+        "hard_cap": resolved_hard_cap,
+    }
+    if memory_bound is not None:
+        limits["memory"] = memory_bound
+    if fd_bound is not None:
+        limits["file_descriptors"] = fd_bound
+    resolved = max(1, min(limits.values()))
+    limiting_factor = next(
+        name
+        for name in ("rows", "cpu", "memory", "file_descriptors", "hard_cap")
+        if limits.get(name) == resolved
+    )
+    return NumThreadsPlan(
+        requested=requested,
+        resolved=resolved,
+        row_bound=row_bound,
+        cpu_bound=cpu_bound,
+        memory_bound=memory_bound,
+        fd_bound=fd_bound,
+        hard_cap=resolved_hard_cap,
+        worker_memory_bytes=resolved_worker_memory_bytes,
+        available_memory_bytes=detected_available_memory_bytes,
+        total_memory_bytes=detected_total_memory_bytes,
+        memory_reserve_bytes=memory_reserve_bytes,
+        limiting_factor=limiting_factor,
+    )
+
+
 def _is_better(challenger: float, incumbent: float, higher_is_better: bool) -> bool:
     return challenger > incumbent if higher_is_better else challenger < incumbent
 
@@ -564,12 +970,13 @@ def evaluate_candidate(
     scorer_runner: ScorerRunner,
     budget: Budget,
     evaluation_cache: EvaluationCache | None = None,
+    num_threads: int = 1,
 ) -> list[RowEvaluation]:
-    evaluations: list[RowEvaluation] = []
-    for row in rows:
-        if not budget.allows_metric_call():
-            break
+    """Evaluate one candidate over rows, preserving row order while parallelizing uncached calls."""
+    evaluations_by_index: dict[int, RowEvaluation] = {}
+    pending: list[tuple[int, DatasetRow, str | None]] = []
 
+    for index, row in enumerate(rows):
         cache_key = None
         if evaluation_cache is not None:
             cache_key = evaluation_cache.key(
@@ -581,9 +988,12 @@ def evaluate_candidate(
             )
             cached = evaluation_cache.get(cache_key)
             if cached is not None:
-                evaluations.append(cached)
+                evaluations_by_index[index] = cached
                 continue
 
+        pending.append((index, row, cache_key))
+
+    def run_uncached_row(index: int, row: DatasetRow, cache_key: str | None) -> tuple[int, str | None, RowEvaluation]:
         started = time.time()
         candidate_result = candidate_runner(candidate_text, row, prompt_context, candidate_id)
         scorer_result = scorer_runner(row, candidate_result, scorer_context, candidate_id)
@@ -610,17 +1020,83 @@ def evaluate_candidate(
             if candidate_result.cost_usd is not None or scorer_result.cost_usd is not None else None,
             error=candidate_result.error or scorer_result.error,
         )
-        evaluations.append(evaluation)
+        return index, cache_key, evaluation
+
+    worker_count = max(1, min(max(1, num_threads), len(pending) or 1))
+    if worker_count == 1:
+        completed = [run_uncached_row(index, row, cache_key) for index, row, cache_key in pending]
+    else:
+        completed = []
+        next_pending_index = 0
+        executor = ThreadPoolExecutor(max_workers=worker_count)
+        futures = set()
+
+        def submit_next_pending() -> bool:
+            nonlocal next_pending_index
+            if next_pending_index >= len(pending):
+                return False
+            index, row, cache_key = pending[next_pending_index]
+            next_pending_index += 1
+            futures.add(executor.submit(run_uncached_row, index, row, cache_key))
+            return True
+
+        try:
+            for _ in range(worker_count):
+                if not submit_next_pending():
+                    break
+            while futures:
+                done, _ = wait(futures, return_when=FIRST_COMPLETED)
+                batch = []
+                for future in done:
+                    futures.remove(future)
+                    batch.append(future.result())
+                completed.extend(batch)
+                for _ in batch:
+                    submit_next_pending()
+        except BaseException:
+            for future in futures:
+                future.cancel()
+            executor.shutdown(wait=True)
+            raise
+        else:
+            executor.shutdown(wait=True)
+
+    for index, cache_key, evaluation in completed:
+        evaluations_by_index[index] = evaluation
         if evaluation_cache is not None and cache_key is not None:
             evaluation_cache.set(cache_key, evaluation)
         budget.used_metric_calls += 1
-    return evaluations
+
+    return [
+        evaluations_by_index[index]
+        for index in sorted(evaluations_by_index)
+    ]
+
+
+def _log_iteration_progress(
+    event_sink: EventSink,
+    *,
+    budget: Budget,
+    iteration: int,
+    candidate_id: str | None,
+) -> None:
+    event_sink.log_event(
+        "optimization_progress",
+        budget.progress_payload(stage="iteration_completed", iteration=iteration),
+        event_layer="extension",
+        iteration=iteration,
+        candidate_id=candidate_id,
+    )
 
 
 def _completed_all_rows(results: list[RowEvaluation], rows: list[DatasetRow]) -> bool:
     return len(results) == len(rows)
 
 
+def _has_future_iteration(config: TextGepaConfig, iteration: int) -> bool:
+    return config.max_iterations is None or iteration < config.max_iterations
+
+
 def optimize_loaded_text_candidate(
     *,
     run_id: str,
@@ -636,7 +1112,17 @@ def optimize_loaded_text_candidate(
     local_logger: Any | None = None,
 ) -> TextGepaResult:
     seed_text = prompt_context.body or ""
-    budget = Budget.from_config(config)
+    budget = Budget.from_config(
+        config,
+        trainset_size=len(trainset),
+        valset_size=len(valset),
+        num_components=1,
+    )
+    num_threads_plan = resolve_num_threads(
+        config.num_threads,
+        minibatch_size=config.minibatch_size,
+        validation_count=len(valset),
+    )
     evaluation_cache = EvaluationCache() if config.cache_evaluations else None
     candidate_text_by_id = {"seed": seed_text}
     val_scores_by_candidate: dict[str, dict[str, float]] = {}
@@ -645,6 +1131,7 @@ def optimize_loaded_text_candidate(
     metric_key = _metric_key(scorer_context)
     higher_is_better = scorer_context.higher_is_better
     budget_exhausted = False
+    failed_reflection_count = 0
 
     try:
         event_sink.log_event("run_started", {
@@ -664,6 +1151,7 @@ def optimize_loaded_text_candidate(
             "train_count": len(trainset),
             "validation_count": len(valset),
             "config": dataclasses.asdict(config),
+            "num_threads": num_threads_plan.to_payload(),
             "budget": budget.to_payload(),
         })
 
@@ -684,6 +1172,7 @@ def optimize_loaded_text_candidate(
             scorer_runner=scorer_runner,
             budget=budget,
             evaluation_cache=evaluation_cache,
+            num_threads=num_threads_plan.resolved,
         )
         if local_logger is not None:
             local_logger.append_evaluations(
@@ -694,7 +1183,7 @@ def optimize_loaded_text_candidate(
                 results=seed_results,
             )
         if not _completed_all_rows(seed_results, valset):
-            raise RuntimeError("Budget exhausted before seed validation completed")
+            raise RuntimeError("Seed validation did not complete all requested rows")
         budget.used_full_evals += 1
         seed_score = _mean(seed_results)
         val_scores_by_candidate["seed"] = {result.row_id: result.score for result in seed_results}
@@ -725,8 +1214,23 @@ def optimize_loaded_text_candidate(
             candidate_id=best_candidate_id,
         )
         event_sink.log_event("budget_updated", budget.to_payload(), event_layer="extension")
+        if not budget.allows_iteration():
+            event_sink.log_event(
+                "budget_exhausted",
+                {
+                    "stage": "seed_val_set_completed",
+                    "requested_rows": len(valset),
+                    "completed_rows": len(seed_results),
+                    "budget": budget.to_payload(),
+                    "decision_rule": "budget is checked only after completing seed validation or a full iteration",
+                },
+                event_layer="extension",
+                candidate_id="seed",
+            )
+            budget_exhausted = True
 
-        for iteration in range(1, config.max_iterations + 1):
+        iteration = 1
+        while config.max_iterations is None or iteration <= config.max_iterations:
             if not budget.allows_iteration():
                 break
 
@@ -787,6 +1291,7 @@ def optimize_loaded_text_candidate(
                 scorer_runner=scorer_runner,
                 budget=budget,
                 evaluation_cache=evaluation_cache,
+                num_threads=num_threads_plan.resolved,
             )
             if local_logger is not None:
                 local_logger.append_evaluations(
@@ -845,6 +1350,7 @@ def optimize_loaded_text_candidate(
                     iteration=iteration,
                     candidate_id=parent_candidate_id,
                 )
+                budget.used_iterations += 1
                 event_sink.log_event("budget_updated", budget.to_payload(), event_layer="extension", iteration=iteration)
                 event_sink.log_event(
                     "iteration_completed",
@@ -863,6 +1369,28 @@ def optimize_loaded_text_candidate(
                     iteration=iteration,
                     candidate_id=best_candidate_id,
                 )
+                _log_iteration_progress(
+                    event_sink,
+                    budget=budget,
+                    iteration=iteration,
+                    candidate_id=best_candidate_id,
+                )
+                if _has_future_iteration(config, iteration) and not budget.allows_iteration():
+                    event_sink.log_event(
+                        "budget_exhausted",
+                        {
+                            "stage": "iteration_completed",
+                            "completed_iteration": iteration,
+                            "budget": budget.to_payload(),
+                            "decision_rule": "budget is checked only between iterations",
+                        },
+                        event_layer="extension",
+                        iteration=iteration,
+                        candidate_id=best_candidate_id,
+                    )
+                    budget_exhausted = True
+                    break
+                iteration += 1
                 continue
 
             event_sink.log_event(
@@ -876,7 +1404,68 @@ def optimize_loaded_text_candidate(
                 iteration=iteration,
                 candidate_id=parent_candidate_id,
             )
-            reflection = reflector(parent_text, parent_results, config)
+            try:
+                reflection = reflector(parent_text, parent_results, config)
+            except RetryableReflectionError as error:
+                failed_reflection_count += 1
+                budget.used_candidate_proposals += 1
+                budget.used_iterations += 1
+                budget.charge_retryable_reflection_failure()
+                event_sink.log_event(
+                    "reflection_failed",
+                    {
+                        "parent_candidate_id": parent_candidate_id,
+                        "row_ids": [row.id for row in minibatch],
+                        "error": str(error),
+                        "error_type": type(error).__name__,
+                        "budget": budget.to_payload(),
+                    },
+                    event_layer="extension",
+                    iteration=iteration,
+                    candidate_id=parent_candidate_id,
+                )
+                event_sink.log_event("budget_updated", budget.to_payload(), event_layer="extension", iteration=iteration)
+                event_sink.log_event(
+                    "iteration_completed",
+                    {
+                        "parent_candidate_id": parent_candidate_id,
+                        "child_candidate_id": None,
+                        "parent_train_score_total": parent_total,
+                        "child_train_score_total": None,
+                        "child_validation_score": None,
+                        "best_candidate_id": best_candidate_id,
+                        "best_validation_score": best_score,
+                        "reflection_failed": True,
+                        "reflection_error": str(error),
+                        "reflection_error_type": type(error).__name__,
+                        "budget": budget.to_payload(),
+                    },
+                    iteration=iteration,
+                    candidate_id=best_candidate_id,
+                )
+                _log_iteration_progress(
+                    event_sink,
+                    budget=budget,
+                    iteration=iteration,
+                    candidate_id=best_candidate_id,
+                )
+                if _has_future_iteration(config, iteration) and not budget.allows_iteration():
+                    event_sink.log_event(
+                        "budget_exhausted",
+                        {
+                            "stage": "iteration_completed",
+                            "completed_iteration": iteration,
+                            "budget": budget.to_payload(),
+                            "decision_rule": "budget is checked only between iterations",
+                        },
+                        event_layer="extension",
+                        iteration=iteration,
+                        candidate_id=best_candidate_id,
+                    )
+                    budget_exhausted = True
+                    break
+                iteration += 1
+                continue
             child_id = _candidate_id(iteration, reflection.candidate_text)
             child_text = reflection.candidate_text
             candidate_text_by_id[child_id] = child_text
@@ -946,6 +1535,7 @@ def optimize_loaded_text_candidate(
                 scorer_runner=scorer_runner,
                 budget=budget,
                 evaluation_cache=evaluation_cache,
+                num_threads=num_threads_plan.resolved,
             )
             if local_logger is not None:
                 local_logger.append_evaluations(
@@ -1028,6 +1618,7 @@ def optimize_loaded_text_candidate(
                     scorer_runner=scorer_runner,
                     budget=budget,
                     evaluation_cache=evaluation_cache,
+                    num_threads=num_threads_plan.resolved,
                 )
                 if local_logger is not None:
                     local_logger.append_evaluations(
@@ -1136,6 +1727,7 @@ def optimize_loaded_text_candidate(
                     parent_candidate_id=parent_candidate_id,
                 )
 
+            budget.used_iterations += 1
             event_sink.log_event("budget_updated", budget.to_payload(), event_layer="extension", iteration=iteration)
             event_sink.log_event(
                 "iteration_completed",
@@ -1152,6 +1744,28 @@ def optimize_loaded_text_candidate(
                 iteration=iteration,
                 candidate_id=best_candidate_id,
             )
+            _log_iteration_progress(
+                event_sink,
+                budget=budget,
+                iteration=iteration,
+                candidate_id=best_candidate_id,
+            )
+            if _has_future_iteration(config, iteration) and not budget.allows_iteration():
+                event_sink.log_event(
+                    "budget_exhausted",
+                    {
+                        "stage": "iteration_completed",
+                        "completed_iteration": iteration,
+                        "budget": budget.to_payload(),
+                        "decision_rule": "budget is checked only between iterations",
+                    },
+                    event_layer="extension",
+                    iteration=iteration,
+                    candidate_id=best_candidate_id,
+                )
+                budget_exhausted = True
+                break
+            iteration += 1
 
         if budget_exhausted:
             event_sink.log_event(
@@ -1163,6 +1777,7 @@ def optimize_loaded_text_candidate(
                     "seed_validation_score": seed_score,
                     "higher_is_better": higher_is_better,
                     "budget": budget.to_payload(),
+                    "failed_reflection_count": failed_reflection_count,
                 },
             )
             event_sink.finish_run(
@@ -1173,6 +1788,7 @@ def optimize_loaded_text_candidate(
                     "seed_score": seed_score,
                     "budget": budget.to_payload(),
                     "pause_reason": "budget_exhausted",
+                    "failed_reflection_count": failed_reflection_count,
                 },
             )
             return TextGepaResult(
@@ -1211,6 +1827,7 @@ def optimize_loaded_text_candidate(
                 "metric_key": metric_key,
                 "higher_is_better": higher_is_better,
                 "budget": budget.to_payload(),
+                "failed_reflection_count": failed_reflection_count,
             },
         )
         event_sink.finish_run(
@@ -1223,6 +1840,7 @@ def optimize_loaded_text_candidate(
                 "metric_key": metric_key,
                 "higher_is_better": higher_is_better,
                 "budget": budget.to_payload(),
+                "failed_reflection_count": failed_reflection_count,
             },
         )
         return TextGepaResult(