orizu 0.3.7 → 0.4.0

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

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import sys
 from pathlib import Path
 from typing import Any
 
@@ -11,6 +12,8 @@ from .optimizer import TextGepaConfig, optimize_loaded_text_candidate
 from .reflection import reflect_with_provider
 from .runner import make_candidate_runner, make_scorer_runner
 
+BUDGET_OPTION_FLAGS = ("--budget", "--max-metric-calls", "--max-full-evals", "--max-iterations")
+
 
 def positive_int(value: str) -> int:
     parsed = int(value)
@@ -46,6 +49,38 @@ def read_json_object_arg(value: str | None, flag_name: str) -> dict[str, Any] |
     return parsed
 
 
+def _flag_occurrences(raw_args: list[str], flag: str) -> int:
+    return sum(1 for arg in raw_args if arg == flag or arg.startswith(f"{flag}="))
+
+
+def _format_budget_option(flag: str, count: int) -> str:
+    return flag if count == 1 else f"{flag} ({count} times)"
+
+
+def apply_budget_defaults(args: argparse.Namespace, raw_args: list[str]) -> None:
+    occurrences = {
+        flag: _flag_occurrences(raw_args, flag)
+        for flag in BUDGET_OPTION_FLAGS
+    }
+    provided = [
+        _format_budget_option(flag, count)
+        for flag, count in occurrences.items()
+        if count > 0
+    ]
+    total = sum(occurrences.values())
+
+    if total > 1:
+        allowed = ", ".join(BUDGET_OPTION_FLAGS)
+        received = ", ".join(provided)
+        raise ValueError(
+            f"Budget options are mutually exclusive; choose at most one of {allowed}. "
+            f"Received: {received}."
+        )
+
+    if total == 0:
+        args.budget = TextGepaConfig.budget
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(prog="orizu-gepa")
     parser.add_argument("--project", required=True)
@@ -60,8 +95,8 @@ def main() -> None:
     parser.add_argument("--split-set-id", required=True)
     parser.add_argument("--train-split", default="train")
     parser.add_argument("--val-split", default="validation")
-    parser.add_argument("--budget", default="light", choices=["auto", "light", "medium", "heavy"])
-    parser.add_argument("--max-iterations", type=positive_int, default=3)
+    parser.add_argument("--budget", choices=["auto", "light", "medium", "heavy"])
+    parser.add_argument("--max-iterations", type=positive_int)
     parser.add_argument("--minibatch-size", type=positive_int, default=3)
     parser.add_argument("--num-threads", type=num_threads_arg, default=TextGepaConfig.num_threads)
     parser.add_argument(
@@ -76,6 +111,18 @@ def main() -> None:
     parser.add_argument("--reflection-model", default=TextGepaConfig.reflection_model)
     parser.add_argument("--reflection-temperature", type=float)
     parser.add_argument("--reflection-max-tokens", type=positive_int)
+    parser.add_argument(
+        "--reflection-retry-attempts",
+        type=positive_int,
+        default=TextGepaConfig.reflection_retry_attempts,
+        help="Total attempts for retryable reflection HTTP failures.",
+    )
+    parser.add_argument(
+        "--reflection-http-timeout-seconds",
+        type=positive_int,
+        default=TextGepaConfig.reflection_http_timeout_seconds,
+        help="Per-attempt timeout for reflection HTTP calls.",
+    )
     parser.add_argument("--reflection-prompt-template")
     parser.add_argument("--reflection-provider-settings")
     parser.add_argument(
@@ -99,7 +146,12 @@ def main() -> None:
     parser.add_argument("--log-dir", default="logs")
     parser.add_argument("--no-local-log", action="store_true")
     parser.add_argument("--metadata", default="{}")
-    args = parser.parse_args()
+    raw_args = sys.argv[1:]
+    args = parser.parse_args(raw_args)
+    try:
+        apply_budget_defaults(args, raw_args)
+    except ValueError as error:
+        parser.error(str(error))
 
     metadata = json.loads(args.metadata)
     if not isinstance(metadata, dict):
@@ -117,6 +169,8 @@ def main() -> None:
         reflection_model=args.reflection_model,
         reflection_temperature=args.reflection_temperature,
         reflection_max_tokens=args.reflection_max_tokens,
+        reflection_retry_attempts=args.reflection_retry_attempts,
+        reflection_http_timeout_seconds=args.reflection_http_timeout_seconds,
         reflection_prompt_template=read_text_arg(args.reflection_prompt_template),
         reflection_provider_settings=read_json_object_arg(args.reflection_provider_settings, "--reflection-provider-settings") or {},
         skip_perfect_parent_reflection=args.skip_perfect_parent_reflection,

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

@@ -12,6 +12,39 @@ from typing import Any
 from .optimizer import DatasetRow, PromptContext
 
 
+def _format_number(value: Any) -> str | None:
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return str(value)
+    if isinstance(value, float):
+        return f"{value:g}"
+    return None
+
+
+def _progress_log_suffix(event_type: str, payload: dict[str, Any] | None) -> str:
+    if event_type != "optimization_progress" or not isinstance(payload, dict):
+        return ""
+
+    percent_value = payload.get("percent")
+    if percent_value is None:
+        percent_value = payload.get("progress_percent")
+    percent = _format_number(percent_value)
+    metric_calls_remaining = _format_number(payload.get("metric_calls_remaining"))
+    metric_call_budget_value = payload.get("metric_call_budget")
+    if metric_call_budget_value is None:
+        metric_call_budget_value = payload.get("approx_metric_call_budget")
+    metric_call_budget = _format_number(metric_call_budget_value)
+    parts: list[str] = []
+    if percent is not None:
+        parts.append(f"{percent}%")
+    if metric_calls_remaining is not None:
+        budget_suffix = f" / {metric_call_budget}" if metric_call_budget is not None else ""
+        parts.append(f"{metric_calls_remaining}{budget_suffix} metric calls left")
+
+    return f" {'; '.join(parts)}" if parts else ""
+
+
 @dataclass(frozen=True)
 class OrizuClient:
     api_url: str
@@ -306,7 +339,10 @@ class OrizuEventSink:
                     parent_candidate_id=parent_candidate_id,
                     child_candidate_id=child_candidate_id,
                 )
-                print(f"[orizu-gepa] {self.sequence:03d} {event_type}", flush=True)
+                print(
+                    f"[orizu-gepa] {self.sequence:03d} {event_type}{_progress_log_suffix(event_type, payload)}",
+                    flush=True,
+                )
                 return
             except Exception as exc:
                 last_error = exc

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

@@ -44,7 +44,7 @@ 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_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
@@ -142,6 +142,10 @@ 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,
@@ -170,13 +174,37 @@ def _dspy_auto_metric_budget(
     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(
@@ -188,11 +216,33 @@ class Budget:
         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_metric_calls", config.max_full_evals * (trainset_size + valset_size))
+            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)
+            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,
@@ -200,46 +250,106 @@ class Budget:
             "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",
-            _dspy_auto_metric_budget(
-                num_components=num_components,
-                num_candidates=preset_candidates,
-                valset_size=valset_size,
-            ),
+            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:
         if self.budget_kind == "max_full_evals":
             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 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 {
             "budget_kind": self.budget_kind,
             "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"
@@ -250,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."
@@ -524,7 +636,8 @@ def resolve_num_threads(
     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)
+    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,
@@ -960,10 +1073,30 @@ def evaluate_candidate(
     ]
 
 
+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,
@@ -998,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", {
@@ -1095,7 +1229,8 @@ def optimize_loaded_text_candidate(
             )
             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
 
@@ -1215,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",
@@ -1233,7 +1369,13 @@ def optimize_loaded_text_candidate(
                     iteration=iteration,
                     candidate_id=best_candidate_id,
                 )
-                if iteration < config.max_iterations and not budget.allows_iteration():
+                _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",
                         {
@@ -1248,6 +1390,7 @@ def optimize_loaded_text_candidate(
                     )
                     budget_exhausted = True
                     break
+                iteration += 1
                 continue
 
             event_sink.log_event(
@@ -1261,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
@@ -1523,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",
@@ -1539,7 +1744,13 @@ def optimize_loaded_text_candidate(
                 iteration=iteration,
                 candidate_id=best_candidate_id,
             )
-            if iteration < config.max_iterations and not budget.allows_iteration():
+            _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",
                     {
@@ -1554,6 +1765,7 @@ def optimize_loaded_text_candidate(
                 )
                 budget_exhausted = True
                 break
+            iteration += 1
 
         if budget_exhausted:
             event_sink.log_event(
@@ -1565,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(
@@ -1575,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(
@@ -1613,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(
@@ -1625,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(