orizu 0.3.6 → 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)
@@ -19,6 +22,12 @@ def positive_int(value: str) -> int:
     return parsed
 
 
+def num_threads_arg(value: str) -> int | str:
+    if value.strip().lower() == "auto":
+        return "auto"
+    return positive_int(value)
+
+
 def read_text_arg(value: str | None) -> str | None:
     if value is None:
         return None
@@ -40,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)
@@ -54,9 +95,10 @@ 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", "high"])
-    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(
         "--candidate-selection-strategy",
         default=TextGepaConfig.candidate_selection_strategy,
@@ -69,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(
@@ -92,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):
@@ -102,6 +161,7 @@ def main() -> None:
         budget=args.budget,
         max_iterations=args.max_iterations,
         minibatch_size=args.minibatch_size,
+        num_threads=args.num_threads,
         candidate_selection_strategy=args.candidate_selection_strategy,
         epsilon=args.epsilon,
         max_metric_calls=args.max_metric_calls,
@@ -109,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,
@@ -153,6 +215,7 @@ def main() -> None:
         "dataset_size": len(trainset) + len(valset),
         "train_count": len(trainset),
         "validation_count": len(valset),
+        "num_threads": args.num_threads,
     }
 
     run_id = client.start_run(

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