synth-ai 0.2.8.dev4__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

examples/warming_up_to_rl/run_local_rollout_traced.py

@@ -0,0 +1,477 @@
+#!/usr/bin/env python3
+"""Run a local Crafter rollout, capture tracing metadata, and optionally persist the trace."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import httpx
+from synth_ai.task import (
+    RolloutEnvSpec,
+    RolloutPolicySpec,
+    RolloutRecordConfig,
+    RolloutRequest,
+    RolloutSafetyConfig,
+    TaskAppClient,
+)
+
+
+def build_rollout_request(
+    *,
+    seed: int,
+    run_id: str,
+    model: str,
+    inference_url: str,
+    inference_api_key: str,
+    ops: list[str],
+    return_trace: bool,
+    trace_format: str,
+    max_policy_tokens: int | None,
+) -> RolloutRequest:
+    policy_config = {
+        "model": model,
+        "inference_url": inference_url,
+        "api_key": inference_api_key,
+    }
+    if max_policy_tokens is not None:
+        policy_config.update(
+            {
+                "max_completion_tokens": max_policy_tokens,
+                "max_tokens": max_policy_tokens,
+            }
+        )
+
+    record = RolloutRecordConfig(
+        trajectories=True,
+        return_trace=return_trace,
+        trace_format=trace_format,
+    )
+
+    from synth_ai.task.contracts import RolloutMode
+    return RolloutRequest(
+        run_id=run_id,
+        env=RolloutEnvSpec(env_name="crafter", seed=seed, config={}),
+        policy=RolloutPolicySpec(policy_name="crafter-react", config=policy_config),
+        ops=ops,
+        record=record,
+        mode=RolloutMode.EVAL,
+        on_done="reset",
+        safety=RolloutSafetyConfig(),
+    )
+
+
+def summarise_rollout(response: Any) -> dict[str, Any]:
+    metrics = (
+        response.metrics.model_dump()
+        if hasattr(response, "metrics")
+        else response.get("metrics", {})
+    )
+    return {
+        "run_id": getattr(response, "run_id", None) or response.get("run_id"),
+        "num_episodes": metrics.get("num_episodes"),
+        "num_steps": metrics.get("num_steps"),
+        "episode_returns": metrics.get("episode_returns"),
+        "outcome_score": metrics.get("outcome_score"),
+        "events_score": metrics.get("events_score"),
+    }
+
+
+def summarise_trace(trace: Any) -> dict[str, Any]:
+    if trace is None:
+        return {"trace": None}
+    if not isinstance(trace, dict):
+        return {"trace_type": type(trace).__name__}
+
+    format_hint = "compact" if "events_count" in trace or "lm_calls" in trace else "full"
+    events_count = trace.get("events_count")
+    if (
+        events_count is None
+        and "event_history" in trace
+        and isinstance(trace["event_history"], list)
+    ):
+        events_count = len(trace["event_history"])
+    messages_count = trace.get("messages_count")
+    if (
+        messages_count is None
+        and "markov_blanket_message_history" in trace
+        and isinstance(trace["markov_blanket_message_history"], list)
+    ):
+        messages_count = len(trace["markov_blanket_message_history"])
+
+    metadata = trace.get("metadata") if isinstance(trace.get("metadata"), dict) else {}
+    lm_calls = trace.get("lm_calls") if isinstance(trace.get("lm_calls"), list) else []
+    decision_rewards = (
+        trace.get("decision_rewards") if isinstance(trace.get("decision_rewards"), list) else []
+    )
+
+    return {
+        "session_id": trace.get("session_id"),
+        "format": format_hint,
+        "events_count": events_count,
+        "messages_count": messages_count,
+        "metadata_keys": sorted(metadata.keys()),
+        "lm_calls_count": len(lm_calls),
+        "decision_turns": len(decision_rewards),
+    }
+
+
+def ensure_ops(ops_arg: str | None, max_llm_calls: int) -> list[str]:
+    if ops_arg:
+        ops = [op.strip() for op in ops_arg.split(",") if op.strip()]
+        if not ops:
+            raise ValueError("--ops must contain at least one entry when provided")
+        return ops
+    max_llm_calls = max(max_llm_calls, 1)
+    ops: list[str] = []
+    for _ in range(max_llm_calls):
+        ops.extend(["agent", "env"])
+    return ops
+
+
+def dump_trace(trace: dict[str, Any], *, path: Path, pretty: bool) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as fh:
+        json.dump(trace, fh, indent=2 if pretty else None)
+        fh.write("\n")
+
+
+def extract_environment_rewards(trace_payload: dict[str, Any] | None) -> list[float]:
+    if not trace_payload:
+        return []
+
+    rewards: list[float] = []
+
+    def _collect(events: list[dict[str, Any]]) -> None:
+        for event in events:
+            reward = event.get("reward")
+            if reward is not None:
+                try:
+                    rewards.append(float(reward))
+                except Exception:
+                    continue
+
+    if isinstance(trace_payload.get("event_history"), list):
+        _collect(trace_payload["event_history"])
+    if isinstance(trace_payload.get("session_time_steps"), list):
+        for step in trace_payload["session_time_steps"]:
+            _collect(step.get("events", []))
+
+    return rewards
+
+
+def extract_decision_rewards(trace_payload: dict[str, Any] | None) -> list[dict[str, Any]]:
+    if not trace_payload:
+        return []
+    rewards = trace_payload.get("decision_rewards")
+    return rewards if isinstance(rewards, list) else []
+
+
+def extract_trajectory_rewards(response: Any) -> list[float]:
+    """Extract per-step rewards directly from the rollout trajectories."""
+
+    rewards: list[float] = []
+
+    if response is None:
+        return rewards
+
+    trajectories = getattr(response, "trajectories", None)
+    if trajectories is None and isinstance(response, dict):
+        trajectories = response.get("trajectories")
+
+    if not trajectories:
+        return rewards
+
+    for traj in trajectories:
+        steps = getattr(traj, "steps", None)
+        if steps is None and isinstance(traj, dict):
+            steps = traj.get("steps")
+        if not steps:
+            continue
+        for step in steps:
+            reward_val = getattr(step, "reward", None)
+            if reward_val is None and isinstance(step, dict):
+                reward_val = step.get("reward")
+            if reward_val is None:
+                continue
+            try:
+                rewards.append(float(reward_val))
+            except Exception:
+                continue
+
+    return rewards
+
+
+def print_reward_summary(
+    trace_payload: dict[str, Any] | None,
+    rollout_summary: dict[str, Any],
+    trajectory_rewards: list[float],
+) -> None:
+    print("Reward summary:")
+
+    env_rewards = extract_environment_rewards(trace_payload)
+    reward_source = "trace"
+    if not env_rewards and trajectory_rewards:
+        env_rewards = trajectory_rewards
+        reward_source = "trajectory"
+
+    if env_rewards:
+        print(f"  Environment rewards per step ({reward_source}): {env_rewards}")
+        print(f"  Environment reward total: {sum(env_rewards):.3f}")
+    else:
+        print("  Environment rewards per step: none recorded")
+
+    decision_rewards = extract_decision_rewards(trace_payload)
+    if decision_rewards:
+        print("  Decision rewards:")
+        for entry in decision_rewards:
+            turn = entry.get("turn")
+            ach_delta = entry.get("ach_delta")
+            unique_delta = entry.get("unique_delta")
+            achievements = entry.get("achievements") or []
+            print(
+                f"    turn={turn}, ach_delta={ach_delta}, unique_delta={unique_delta}, achievements={achievements}"
+            )
+    else:
+        print("  Decision rewards: none recorded")
+
+    episode_returns = rollout_summary.get("episode_returns")
+    if episode_returns:
+        print(f"  Outcome rewards (episode returns): {episode_returns}")
+        if env_rewards:
+            try:
+                total_env_reward = float(sum(env_rewards))
+                target = float(episode_returns[0]) if episode_returns else 0.0
+                if abs(total_env_reward - target) > 1e-6:
+                    print(
+                        "  ⚠️  Reward mismatch: sum(environment rewards)"
+                        f"={total_env_reward:.3f} vs episode return={target:.3f}"
+                    )
+            except Exception:
+                pass
+    else:
+        print("  Outcome rewards: none recorded")
+
+
+async def main() -> None:
+    # Load .env file from current directory if it exists
+    env_file = Path.cwd() / ".env"
+    if env_file.exists():
+        from dotenv import load_dotenv
+
+        load_dotenv(env_file)
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base-url", default="http://localhost:8001", help="Task app base URL")
+    parser.add_argument("--api-key", help="RL Environment API key (will prompt if not provided)")
+    parser.add_argument(
+        "--inference-api-key", help="Inference provider API key (will prompt if not provided)"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Environment seed")
+    parser.add_argument("--run-id", default="local-trace", help="Run identifier")
+    parser.add_argument("--model", default="gpt-4o-mini", help="OpenAI-compatible model id")
+    parser.add_argument(
+        "--inference-url", default="https://api.openai.com", help="Inference base URL (OpenAI/Groq)"
+    )
+    parser.add_argument(
+        "--ops", help="Comma-separated rollout ops (fallback: alternating agent/env)"
+    )
+    parser.add_argument(
+        "--max-llm-calls",
+        type=int,
+        default=1,
+        help="Number of agent/env pairs when --ops not supplied",
+    )
+    parser.add_argument(
+        "--max-policy-tokens",
+        type=int,
+        default=None,
+        help="Optional max token budget forwarded to policy",
+    )
+    parser.add_argument(
+        "--trace-format",
+        choices=["compact", "full"],
+        default="compact",
+        help="Trace payload format requested from the server",
+    )
+    parser.add_argument(
+        "--trace-path",
+        type=Path,
+        help="Path to write the trace JSON (defaults to ./<run_id>_trace.json unless --no-trace-file is set)",
+    )
+    parser.add_argument(
+        "--no-trace-file",
+        action="store_true",
+        help="Do not write the trace JSON to disk",
+    )
+    parser.add_argument(
+        "--no-print-trace",
+        action="store_true",
+        help="Do not print the full trace payload to stdout",
+    )
+    parser.add_argument(
+        "--no-trace",
+        action="store_true",
+        help="Disable return_trace (useful for comparing behaviour without tracing)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=60.0,
+        help="HTTP timeout in seconds for the client (default: 60)",
+    )
+    args = parser.parse_args()
+
+    # Prompt for required parameters if not provided
+    base_url = args.base_url
+    if args.base_url == "http://localhost:8001":
+        print("\nTask app configuration:")
+        base_url_input = input("Task app base URL [http://localhost:8001]: ").strip()
+        base_url = base_url_input if base_url_input else "http://localhost:8001"
+
+    api_key = args.api_key or os.getenv("ENVIRONMENT_API_KEY")
+    if not api_key:
+        api_key = input("RL Environment API key (from ENVIRONMENT_API_KEY): ").strip()
+        if not api_key:
+            parser.error("RL Environment API key is required")
+
+    # Use Groq by default
+    model = "llama-3.3-70b-versatile"
+    inference_url = "https://api.groq.com/openai"
+
+    print("\nInference configuration (Groq):")
+    inference_api_key = args.inference_api_key or os.getenv("GROQ_API_KEY")
+    if not inference_api_key:
+        inference_api_key = input("Groq API key: ").strip()
+        if not inference_api_key:
+            parser.error("Groq API key is required")
+
+        # Save to .env for future use
+        env_path = Path.cwd() / ".env"
+        try:
+            # Read existing .env
+            existing_lines = []
+            if env_path.exists():
+                existing_lines = env_path.read_text().splitlines()
+
+            # Check if GROQ_API_KEY already exists
+            key_exists = any(line.strip().startswith("GROQ_API_KEY=") for line in existing_lines)
+
+            if not key_exists:
+                # Append to .env
+                with open(env_path, "a") as f:
+                    if existing_lines and not existing_lines[-1].strip():
+                        # File exists and last line is not empty
+                        pass
+                    elif existing_lines:
+                        # Add newline before appending
+                        f.write("\n")
+                    f.write(f"GROQ_API_KEY={inference_api_key}\n")
+                print(f"[INFO] Saved GROQ_API_KEY to {env_path}")
+        except Exception as e:
+            print(f"[WARN] Could not save GROQ_API_KEY to .env: {e}")
+
+    print("\nRollout configuration:")
+    max_llm_calls = args.max_llm_calls
+    if args.max_llm_calls == 1:
+        max_llm_calls_input = input("Max LLM calls [10]: ").strip()
+        max_llm_calls = int(max_llm_calls_input) if max_llm_calls_input else 10
+
+    # Override args with prompted values
+    args.base_url = base_url
+    args.max_llm_calls = max_llm_calls
+
+    ops = ensure_ops(args.ops, args.max_llm_calls)
+    return_trace = not args.no_trace
+
+    async with TaskAppClient(args.base_url, api_key=api_key, timeout=args.timeout) as client:
+        try:
+            print(f"Fetching task_info for seed {args.seed}…")
+            task_info = await client.task_info(seeds=[args.seed])
+            info_payload = task_info[0] if isinstance(task_info, list) else task_info
+            try:
+                print(json.dumps(info_payload.model_dump(), indent=2)[:600])
+            except Exception:
+                print(info_payload)
+
+            request = build_rollout_request(
+                seed=args.seed,
+                run_id=args.run_id,
+                model=model,
+                inference_url=inference_url,
+                inference_api_key=inference_api_key,
+                ops=ops,
+                return_trace=return_trace,
+                trace_format=args.trace_format,
+                max_policy_tokens=args.max_policy_tokens,
+            )
+
+            print("Requesting rollout…")
+            response = await client.rollout(request)
+            summary = summarise_rollout(response)
+            print(json.dumps(summary, indent=2))
+
+            trace_payload: dict[str, Any] | None = getattr(response, "trace", None)
+            if return_trace:
+                if trace_payload is None:
+                    print(
+                        "⚠️  Server did not include a trace. Ensure TASKAPP_TRACING_ENABLED=1 when starting the task app.",
+                        file=sys.stderr,
+                    )
+                else:
+                    trace_summary = summarise_trace(trace_payload)
+                    print("Trace summary:")
+                    print(json.dumps(trace_summary, indent=2))
+
+                    trace_path = args.trace_path
+                    if not args.no_trace_file:
+                        if trace_path is None:
+                            trace_path = Path(f"{args.run_id}_trace.json")
+                        dump_trace(trace_payload, path=trace_path, pretty=True)
+                        print(f"Trace written to {trace_path}")
+
+                    if not args.no_print_trace:
+                        print("Full trace payload:")
+                        print(json.dumps(trace_payload, indent=2))
+
+            trajectory_rewards = extract_trajectory_rewards(response)
+            print_reward_summary(
+                trace_payload if return_trace else None,
+                summary,
+                trajectory_rewards,
+            )
+
+            print(f"Ops executed: {ops}")
+            print(
+                "Tip: export TASKAPP_TRACING_ENABLED=1 and optionally TASKAPP_SFT_OUTPUT_DIR before running `uvx synth-ai deploy --runtime uvicorn …` to persist traces/SFT."
+            )
+        except httpx.HTTPStatusError as exc:
+            detail = (
+                exc.response.json()
+                if exc.response.headers.get("content-type", "").startswith("application/json")
+                else exc.response.text
+            )
+            print(f"HTTP error {exc.response.status_code}: {detail}", file=sys.stderr)
+            if exc.response.status_code in (401, 503):
+                print(
+                    "Hint: ensure the task app process is using the same ENVIRONMENT_API_KEY passed via --api-key.",
+                    file=sys.stderr,
+                )
+            if exc.response.status_code == 500:
+                print(
+                    "Hint: verify tracing is enabled server-side (TASKAPP_TRACING_ENABLED=1) and the inference credentials are configured.",
+                    file=sys.stderr,
+                )
+            raise
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("Interrupted", file=sys.stderr)

examples/warming_up_to_rl/run_rl_and_save.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import tomllib
+from pathlib import Path
+from typing import Any
+
+import requests
+from synth_ai._utils.base_url import PROD_BASE_URL_DEFAULT
+
+
+def _load_toml(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        print(f"config not found: {path}", file=sys.stderr)
+        sys.exit(2)
+    with path.open("rb") as fh:
+        return tomllib.load(fh)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(
+        description="Create clustered RL training job via backend RL endpoint"
+    )
+    p.add_argument(
+        "--backend", default=os.getenv("BACKEND_BASE_URL", f"{PROD_BASE_URL_DEFAULT}/api")
+    )
+    p.add_argument("--config", required=True, help="Path to RL TOML config")
+    p.add_argument(
+        "--task-url",
+        default=os.getenv("TASK_APP_URL", ""),
+        help="Override task service URL (or set TASK_APP_URL)",
+    )
+    p.add_argument(
+        "--idempotency",
+        default=os.getenv("RL_IDEMPOTENCY_KEY", ""),
+        help="Optional Idempotency-Key header value",
+    )
+    args = p.parse_args()
+
+    cfg_path = Path(args.config).expanduser()
+    cfg = _load_toml(cfg_path)
+
+    services = cfg.get("services", {}) if isinstance(cfg.get("services"), dict) else {}
+
+    # Resolve task app base URL for the job
+    cli_task_url = (args.task_url or "").strip()
+    env_task_url = (os.getenv("TASK_APP_URL") or "").strip()
+    task_url = (
+        cli_task_url
+        or env_task_url
+        or ((services.get("task_url") or "").strip() if isinstance(services, dict) else "")
+    )
+    if not task_url:
+        print(
+            "Missing task service URL. Provide --task-url or set TASK_APP_URL or services.task_url in TOML",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+    # TOML-only model selection validation
+    model_cfg = cfg.get("model", {}) if isinstance(cfg.get("model"), dict) else {}
+    has_source = bool((model_cfg.get("source") or "").strip())
+    has_base = bool((model_cfg.get("base") or "").strip())
+    if has_source == has_base:
+        print(
+            "Model selection must specify exactly one of [model].source or [model].base in TOML",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+    # Build create-job payload. Send full TOML under data.config, plus endpoint_base_url.
+    payload: dict[str, Any] = {
+        "job_type": "rl",
+        # Optional: compute pass-through
+        "compute": cfg.get("compute", {}) if isinstance(cfg.get("compute"), dict) else {},
+        "data": {
+            "endpoint_base_url": task_url,
+            "config": cfg,
+        },
+        "tags": {"source": "warming_up_to_rl"},
+    }
+
+    backend = str(args.backend).rstrip("/")
+    url = f"{backend}/rl/jobs"
+    api_key = (os.getenv("SYNTH_API_KEY") or os.getenv("SYNTH_KEY") or "").strip()
+    if not api_key:
+        print("Missing SYNTH_API_KEY in env", file=sys.stderr)
+        sys.exit(2)
+
+    headers = {
+        "content-type": "application/json",
+        "authorization": f"Bearer {api_key}",
+    }
+    idem = (args.idempotency or "").strip()
+    if idem:
+        headers["Idempotency-Key"] = idem
+
+    print(f"[INFO] POST {url}")
+    try:
+        preview = dict(payload)
+        preview_data = dict(preview.get("data", {}))
+        cfg_keys = list(cfg.keys())
+        preview_data["config"] = {"keys": cfg_keys}
+        preview["data"] = preview_data
+        print(f"[INFO] Payload: {json.dumps(preview)[:500]}")
+    except Exception:
+        print("[INFO] Payload: <unavailable>")
+
+    r = requests.post(url, headers=headers, json=payload, timeout=120)
+    ok = r.status_code in (200, 201)
+    try:
+        snippet = r.json()
+    except Exception:
+        snippet = r.text[:300]
+    print(f"[INFO] Response: {r.status_code} {snippet}")
+    sys.exit(0 if ok else 1)
+
+
+if __name__ == "__main__":
+    main()