orizu 0.3.5 → 0.3.7

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 = 8
+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,6 +142,34 @@ class ReflectionResult:
     candidate_text: str
 
 
+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
+
+
 @dataclass
 class Budget:
     budget_kind: str
@@ -122,20 +179,35 @@ class Budget:
     used_candidate_proposals: int = 0
 
     @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)
         if config.max_full_evals is not None:
-            return cls("max_full_evals", config.max_full_evals)
+            return cls("max_metric_calls", config.max_full_evals * (trainset_size + valset_size))
         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)
+        preset = "medium" if config.budget == "auto" else config.budget
+        preset_candidates = {
+            "light": 6,
+            "medium": 12,
+            "heavy": 18,
+            "high": 18,
+        }.get(preset, 6)
+        return cls(
+            "max_metric_calls",
+            _dspy_auto_metric_budget(
+                num_components=num_components,
+                num_candidates=preset_candidates,
+                valset_size=valset_size,
+            ),
+        )
 
     @property
     def used(self) -> int:
@@ -152,9 +224,6 @@ class Budget:
     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 to_payload(self) -> dict[str, Any]:
         return {
             "budget_kind": self.budget_kind,
@@ -172,6 +241,7 @@ class TextGepaConfig:
     budget: str = "light"
     max_iterations: int = 3
     minibatch_size: int = 3
+    num_threads: int | str = "auto"
     candidate_selection_strategy: str = "pareto"
     epsilon: float = 0.1
     max_metric_calls: int | None = None
@@ -189,6 +259,7 @@ class TextGepaConfig:
     fail_on_log_error: bool = True
     log_row_snapshots: bool = False
     cache_evaluations: bool = True
+    skip_perfect_parent_reflection: bool = True
 
 
 @dataclass(frozen=True)
@@ -315,10 +386,241 @@ 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()
+    cpu_bound = max(1, detected_cpu_count or 1)
+    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
 
 
+def _is_perfect_minibatch(results: list[RowEvaluation], higher_is_better: bool) -> bool:
+    if not results:
+        return False
+    if higher_is_better:
+        return all(result.score >= 1.0 for result in results)
+    return all(result.score <= 0.0 for result in results)
+
+
 class EvaluationCache:
     def __init__(self):
         self._values: dict[str, RowEvaluation] = {}
@@ -555,12 +857,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(
@@ -572,9 +875,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)
@@ -601,11 +907,57 @@ 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 _completed_all_rows(results: list[RowEvaluation], rows: list[DatasetRow]) -> bool:
@@ -627,7 +979,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]] = {}
@@ -655,6 +1017,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(),
         })
 
@@ -675,6 +1038,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(
@@ -685,7 +1049,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}
@@ -716,6 +1080,20 @@ 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):
             if not budget.allows_iteration():
@@ -778,6 +1156,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(
@@ -819,6 +1198,58 @@ def optimize_loaded_text_candidate(
                 candidate_id=parent_candidate_id,
             )
 
+            parent_is_perfect = _is_perfect_minibatch(parent_results, higher_is_better)
+            if config.skip_perfect_parent_reflection and parent_is_perfect:
+                event_sink.log_event(
+                    "reflection_skipped",
+                    {
+                        "reason": "parent_minibatch_perfect",
+                        "parent_candidate_id": parent_candidate_id,
+                        "parent_score_total": parent_total,
+                        "parent_score_mean": _mean(parent_results),
+                        "row_ids": [row.id for row in minibatch],
+                        "higher_is_better": higher_is_better,
+                        "decision_rule": "skip reflection when every parent minibatch row is already perfect",
+                    },
+                    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,
+                        "skipped_reflection": True,
+                        "skip_reason": "parent_minibatch_perfect",
+                        "budget": budget.to_payload(),
+                    },
+                    iteration=iteration,
+                    candidate_id=best_candidate_id,
+                )
+                if iteration < config.max_iterations 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
+                continue
+
             event_sink.log_event(
                 "reflection_started",
                 {
@@ -900,6 +1331,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(
@@ -982,6 +1414,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(
@@ -1106,6 +1539,21 @@ def optimize_loaded_text_candidate(
                 iteration=iteration,
                 candidate_id=best_candidate_id,
             )
+            if iteration < config.max_iterations 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
 
         if budget_exhausted:
             event_sink.log_event(