| """Langfuse logging for checkpoint eval batches (Kaggle + local).""" | |
| from __future__ import annotations | |
| import os | |
| from typing import Any | |
| def init_langfuse(): | |
| from langfuse import Langfuse, get_client | |
| Langfuse( | |
| public_key=os.environ["LANGFUSE_PUBLIC_KEY"], | |
| secret_key=os.environ["LANGFUSE_SECRET_KEY"], | |
| host=os.getenv("LANGFUSE_HOST") | |
| or os.getenv("LANGFUSE_BASE_URL") | |
| or "https://cloud.langfuse.com", | |
| ) | |
| lf = get_client() | |
| if not lf.auth_check(): | |
| raise RuntimeError("Langfuse auth_check failed") | |
| return lf | |
| def log_checkpoint_eval( | |
| *, | |
| step: int, | |
| hf_revision: str, | |
| train_model_id: str, | |
| results: list[dict[str, Any]], | |
| wandb_run_id: str | None = None, | |
| trace_name: str = "kaggle-ckpt-eval", | |
| ) -> dict[str, Any]: | |
| """ | |
| Write one Langfuse agent trace for a checkpoint eval suite. | |
| Each problem becomes a nested span; aggregate scores on the trace. | |
| MUST be called BEFORE pushing/saving the checkpoint artifact. | |
| """ | |
| lf = init_langfuse() | |
| n = len(results) | |
| mean_pass = sum(float(r.get("pass_frac", 0.0)) for r in results) / n if n else 0.0 | |
| mean_reward = sum(float(r.get("reward_total", 0.0)) for r in results) / n if n else 0.0 | |
| with lf.start_as_current_observation( | |
| name=trace_name, | |
| as_type="evaluator", | |
| metadata={ | |
| "step": step, | |
| "hf_revision": hf_revision, | |
| "train_model_id": train_model_id, | |
| "wandb_run_id": wandb_run_id, | |
| "phase": "ckpt_eval", | |
| "n_problems": n, | |
| }, | |
| input={"step": step, "hf_revision": hf_revision, "n_problems": n}, | |
| ) as root: | |
| for r in results: | |
| with lf.start_as_current_observation( | |
| name=f"task:{r.get('task_id')}", | |
| as_type="span", | |
| input={"prompt": r.get("prompt", "")[:2000]}, | |
| metadata={"task_id": r.get("task_id")}, | |
| ) as span: | |
| span.update( | |
| output={ | |
| "completion": (r.get("completion") or "")[:2000], | |
| "pass_frac": r.get("pass_frac"), | |
| "reward_total": r.get("reward_total"), | |
| "ok": r.get("pass_frac", 0) >= 1.0, | |
| } | |
| ) | |
| summary = { | |
| "mean_pass_frac": mean_pass, | |
| "mean_reward_total": mean_reward, | |
| "n_problems": n, | |
| "n_perfect": sum(1 for r in results if r.get("pass_frac", 0) >= 1.0), | |
| "step": step, | |
| "hf_revision": hf_revision, | |
| } | |
| root.update(output=summary) | |
| lf.score_current_trace(name="mean_pass_frac", value=mean_pass) | |
| lf.score_current_trace(name="mean_reward_total", value=mean_reward) | |
| lf.score_current_trace(name="n_perfect", value=float(summary["n_perfect"])) | |
| lf.score_current_trace(name="checkpoint_step", value=float(step)) | |
| trace_id = lf.get_current_trace_id() or "" | |
| lf.flush() | |
| return {"langfuse_trace_id": trace_id, **summary} | |