Fix Langfuse ckpt/heldout child trace_context

kaggle/train_grpo.py

@@ -344,12 +344,20 @@ def run_heldout_agent_eval(
                     gen_ids = out_ids[0][inputs_local["input_ids"].shape[-1] :]
                     return tokenizer.decode(gen_ids, skip_special_tokens=True)
 
-                if lf is not None and getattr(env, "_root", None) is not None:
+                child_ctx = None
+                if hasattr(env, "_child_trace_context"):
+                    child_ctx = env._child_trace_context()
+                if lf is not None and child_ctx is not None:
                     with lf.start_as_current_observation(
                         name=f"llm-turn-{turn}",
                         as_type="generation",
+                        trace_context=child_ctx,
                         input={"messages_tail": messages[-2:], "turn": turn},
-                        metadata={"phase": "ckpt_eval", "turn": turn},
+                        metadata={
+                            "phase": "ckpt_eval",
+                            "turn": turn,
+                            "trace_id": env._trace_id,
+                        },
                     ) as gen_span:
                         completion = _generate_once()
                         calls = _parse_tool_calls(completion)
@@ -486,12 +494,10 @@ def langfuse_log_ckpt_eval(
     n_perfect = sum(1 for r in results if float(r.get("pass_frac") or 0) >= 1.0)
     n_resolved = sum(1 for r in results if r.get("resolved"))
 
-    # Own root trace (do not nest under / swallow per-rollout traces).
-    agg_trace_id = lf.create_trace_id()
+    # Clean root (no bare trace_context — that creates phantom orphan parents).
     with lf.start_as_current_observation(
         name=f"kaggle-ckpt-eval:step-{step}",
         as_type="evaluator",
-        trace_context={"trace_id": agg_trace_id},
         metadata={
             "step": step,
             "hf_revision": hf_revision,
@@ -503,10 +509,18 @@ def langfuse_log_ckpt_eval(
         input={"step": step, "hf_revision": hf_revision, "n": n},
         end_on_exit=True,
     ) as root:
+        agg_trace_id = str(getattr(root, "trace_id", "") or "")
+        root_span_id = str(getattr(root, "id", "") or "")
         for r in results:
+            child_ctx = (
+                {"trace_id": agg_trace_id, "parent_span_id": root_span_id}
+                if agg_trace_id and root_span_id
+                else None
+            )
             with lf.start_as_current_observation(
                 name=f"ckpt-summary:{r['task_id']}",
                 as_type="span",
+                trace_context=child_ctx,
                 input={"prompt": str(r.get("prompt", ""))[:2000]},
             ) as span:
                 span.update(