PR: T1 — Drop-in Instrumentation & End-to-End Optimization

.env.example

@@ -0,0 +1,8 @@
+# OpenRouter Configuration
+# Copy this file to .env and fill in your values
+# Get your API key from: https://openrouter.ai/keys
+
+OPENROUTER_API_KEY=sk-or-v1-your-key-here
+OPENROUTER_MODEL=meta-llama/llama-3.1-8b-instruct:free
+OPENROUTER_BASE_URL=https://openrouter.ai/api/v1
+USE_STUB_LLM=false

.gitignore

@@ -168,4 +168,6 @@ OAI_CONFIG_LIST
 *.gv.pdf
 
 # jupyter book API output
-docs/api/*
+docs/api/*
+
+uv.lock

T1_PR_M1_latest_fix.diff

@@ -0,0 +1,164 @@
+diff --git a/opto/trace/io/bindings.py b/opto/trace/io/bindings.py
+index 5f3e0bf..a7e7e4a 100644
+--- a/opto/trace/io/bindings.py
++++ b/opto/trace/io/bindings.py
+@@ -1,10 +1,12 @@
+ from __future__ import annotations
+
+ import logging
++import re
+ from dataclasses import dataclass
+ from typing import Any, Callable, Dict, Literal
+
+ logger = logging.getLogger(__name__)
+@@
+ def apply_updates(
+-    updates: Dict[str, Any],
++    updates: Dict[Any, Any],
+     bindings: Dict[str, Binding],
+     *,
+     strict: bool = True,
+-) -> None:
++) -> Dict[str, Any]:
+@@
+-    for key, value in updates.items():
+-        binding = bindings.get(key)
++    def _normalize_key(k: Any) -> str:
++        if isinstance(k, str):
++            s = k
++        else:
++            s = (
++                getattr(k, "name", None)
++                or getattr(k, "_name", None)
++                or getattr(k, "py_name", None)
++                or str(k)
++            )
++        s = str(s).strip()
++        if s.startswith("param."):
++            s = s[len("param."):]
++        s = s.split(":")[0].split("/")[-1]
++        if s not in bindings:
++            s2 = re.sub(r"\d+$", "", s)
++            if s2 in bindings:
++                s = s2
++        return s
++
++    applied: Dict[str, Any] = {}
++    for raw_key, value in updates.items():
++        key = _normalize_key(raw_key)
++        binding = bindings.get(key)
+         if binding is None:
+             if strict:
+                 raise KeyError(
+-                    f"apply_updates: no binding for key {key!r}. "
++                    f"apply_updates: no binding for key {key!r} (from {raw_key!r}). "
+                     f"Available bindings: {sorted(bindings.keys())}"
+                 )
+-            logger.debug("apply_updates: skipping unknown key %r (strict=False)", key)
++            logger.debug("apply_updates: skipping unknown key %r (from %r)", key, raw_key)
+             continue
+         try:
+             binding.set(value)
++            applied[key] = value
+             logger.debug("apply_updates: set %r (kind=%s)", key, binding.kind)
+         except Exception:
+             logger.exception("apply_updates: failed to set %r", key)
++            if strict:
++                raise
++    return applied
+
+diff --git a/opto/trace/io/tgj_ingest.py b/opto/trace/io/tgj_ingest.py
+index 95d3c70..c651d1b 100644
+--- a/opto/trace/io/tgj_ingest.py
++++ b/opto/trace/io/tgj_ingest.py
+@@ -1,6 +1,7 @@
+ from __future__ import annotations
+
+ from dataclasses import dataclass
++from typing import Optional
+ from typing import Any, Dict, List, Tuple
+@@
+-def ingest_tgj(doc: Dict[str,Any], port_index: Optional[Dict[str,Node]] = None) -> Dict[str,Node]:
++def ingest_tgj(
++    doc: Dict[str,Any],
++    port_index: Optional[Dict[str,Node]] = None,
++    *,
++    param_cache: Optional[Dict[str,"ParameterNode"]] = None,
++) -> Dict[str,Node]:
+@@
+         if k == "parameter":
+-            n = ParameterNode(
+-                rec.get("value"),
+-                name=nm,
+-                trainable=bool(rec.get("trainable", True)),
+-                description=rec.get("description", "[Parameter]")
+-            )
++            n = param_cache.get(nid) if param_cache is not None else None
++            if n is None:
++                n = ParameterNode(
++                    rec.get("value"),
++                    name=nm,
++                    trainable=bool(rec.get("trainable", True)),
++                    description=rec.get("description", "[Parameter]"),
++                )
++                if param_cache is not None:
++                    param_cache[nid] = n
++            else:
++                try:
++                    n._data = rec.get("value")
++                except Exception:
++                    pass
++                try:
++                    n.trainable = bool(rec.get("trainable", True))
++                except Exception:
++                    pass
+             nodes[nid] = n
+             nodes[nm] = n
+
+diff --git a/opto/trace/io/optimization.py b/opto/trace/io/optimization.py
+index 271e2b1..a9f0d44 100644
+--- a/opto/trace/io/optimization.py
++++ b/opto/trace/io/optimization.py
+@@ -210,6 +210,7 @@ def optimize_graph(
+     eval_fn = eval_fn or _default_eval_fn
+
++    graph.session.flush_otlp(clear=True)
+     score_history: List[float] = []
+@@
+     best_updates: Dict[str, Any] = {}
+     last_applied_updates: Dict[str, Any] = {}
+
++    param_cache: Dict[str, Any] = {}
+@@
+-                        nodes = _ingest_tgj(doc)
++                        nodes = _ingest_tgj(doc, param_cache=param_cache)
+@@
+                     if _optimizer is not None and all_output_nodes:
+                         try:
+                             _optimizer.zero_feedback()
+-                            _optimizer.backward(output_node, feedback_text)
++                            for output_node, run_for_output in all_output_nodes:
++                                feedback_text = run_for_output.feedback or (
++                                    f"Score: {run_for_output.score}"
++                                    if run_for_output.score is not None
++                                    else "No feedback"
++                                )
++                                _optimizer.backward(output_node, feedback_text)
+                             raw_updates = _optimizer.step()
+                             if isinstance(raw_updates, dict):
+                                 updates.update(raw_updates)
+                         except Exception as exc:
+                             logger.warning(
+                                 "Optimizer step failed: %s", exc, exc_info=True
+                             )
+@@
+             if updates and apply_updates_flag:
+                 try:
+-                    apply_updates(updates, effective_bindings, strict=False)
+-                    last_applied_updates = dict(updates)
+-                    logger.info("Applied updates: %s", sorted(updates.keys()))
++                    applied = apply_updates(updates, effective_bindings, strict=False)
++                    last_applied_updates = dict(applied)
++                    logger.info("Applied updates: %s", sorted(applied.keys()))
+                 except Exception as exc:
+                     logger.warning("apply_updates failed: %s", exc, exc_info=True)