cmpnd 0.76.0__tar.gz → 0.78.0__tar.gz

{cmpnd-0.76.0 → cmpnd-0.78.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cmpnd
-Version: 0.76.0
+Version: 0.78.0
 Summary: DSPy observability and deployment SDK for cmpnd
 Project-URL: Homepage, https://cmpnd.ai
 Author-email: cmpnd <hello@cmpnd.ai>

{cmpnd-0.76.0 → cmpnd-0.78.0}/cmpnd/_gepa_patch.py

@@ -9,6 +9,10 @@ module-style (``from . import identity``); basedpyright can't resolve
 """
 
 import contextvars
+import functools
+from collections.abc import Callable
+
+from . import context
 
 _gepa_patched = False
 
@@ -39,6 +43,43 @@ def _gepa_supports_callbacks() -> bool:
         return False
 
 
+# Marks a metric wrapper _judge_metric_calls installed, so a compile nested in another
+# (TrackedGEPA under auto_instrument) leaves the outer wrapper in place. An attribute
+# on our own wrapper, read with getattr, because the user's metric may be any
+# callable: a dataclass instance is unhashable, so no set or dict can hold it.
+_JUDGED_METRIC_MARK = "_cmpnd_judges_metric_calls"
+
+
+def _judge_metric_calls(gepa) -> Callable[[], None]:
+    """Run each of ``gepa``'s metric calls as the metric phase; return the undo.
+
+    dspy.GEPA reads its metric from ``metric_fn`` once its compile starts, so
+    swapping in a wrapper for the compile's span reaches every metric call, those
+    DSPy Evaluate makes and those GEPA's feedback makes. A trace a grading module
+    roots inside a call is then tagged ``optimization_phase: metric``, as a hosted
+    run's judge traces are, while the rollout the metric grades ran before it and
+    carries no phase. The wrapper keeps the metric's signature (``functools.wraps``)
+    for DSPy's checks of it, and the undo restores the optimizer's own metric.
+    """
+    metric = getattr(gepa, "metric_fn", None)
+    if not callable(metric) or getattr(metric, _JUDGED_METRIC_MARK, False) is True:
+        return lambda: None
+
+    @functools.wraps(metric)
+    def judged_metric(*args, **kwargs):
+        with context.optimization_phase("metric"):
+            return metric(*args, **kwargs)
+
+    setattr(judged_metric, _JUDGED_METRIC_MARK, True)
+    gepa.metric_fn = judged_metric
+
+    def restore() -> None:
+        if getattr(gepa, "metric_fn", None) is judged_metric:
+            gepa.metric_fn = metric
+
+    return restore
+
+
 def _run_tracked_compile(
     gepa,
     student,
@@ -137,8 +178,10 @@ def _run_tracked_compile(
         trk.generated_program_name = program_naming.generated_program_name(student)
 
     restore_attribution = callback.start_trace_attribution()
+    restore_metric = None
 
     try:
+        restore_metric = _judge_metric_calls(gepa)
         result = run_original()
 
         if not has_callbacks:
@@ -164,6 +207,8 @@ def _run_tracked_compile(
         raise
 
     finally:
+        if restore_metric is not None:
+            restore_metric()
         restore_attribution()
 
 

{cmpnd-0.76.0 → cmpnd-0.78.0}/cmpnd/callback.py

@@ -517,6 +517,11 @@ def _finish_reason(response: Any) -> str | None:
     return "stop"
 
 
+# The `optimization.role` an LM call in each GEPA phase is priced under. A metric scores a
+# rollout, so its calls are rollout spend; reflection proposes the next candidate.
+_OPTIMIZATION_PHASE_ROLES = {"metric": "rollout", "reflection": "reflection"}
+
+
 class _CallbackBase:
     """Shared state and span-tracking primitives for the CmpndCallback mixins."""
 
@@ -532,6 +537,7 @@ class _CallbackBase:
         max_output_field_size: int | None = None,
         stream_trace_starts: bool | None = None,
         own_config: configuration.Config | None = None,
+        optimization_phase: str | None = None,
     ) -> None:
         # A client passes its own config here; the auto_instrument callback leaves
         # it None and reads the process-global config live, so a reconfigure()
@@ -541,6 +547,19 @@ class _CallbackBase:
         # later cmpnd.configure(). That fixedness is what keeps a client
         # orthogonal to the global config.
         self._own_config = own_config
+        # The optimization phase this callback traces for a run it does not own: an
+        # executor that scores a hosted rollout with the author's metric ("metric") or runs
+        # the author's GEPA proposer ("reflection"). Such a process knows the program the
+        # run optimizes and not the run's id, which the ingest key it pushes with decides.
+        # With a phase set, an LM call outside any trace becomes a standalone GEPA trace as
+        # a local optimization's judge and reflection calls do, and a trace a module opens
+        # is tagged with the phase. None is every other caller.
+        if optimization_phase not in (None, *_OPTIMIZATION_PHASE_ROLES):
+            raise ValueError(
+                f"optimization_phase must be one of {sorted(_OPTIMIZATION_PHASE_ROLES)} or None, "
+                f"got {optimization_phase!r}"
+            )
+        self._optimization_phase = optimization_phase
         # Set once this callback has warned that a declared name replaced the program
         # id its config read from the environment, so the warning appears once per client.
         self._warned_name_over_id = False
@@ -1348,6 +1367,12 @@ class _ModuleMixin(_CallbackBase):
         # Apply default tags
         if config and config.default_tags:
             trace.tags.update(config.default_tags)
+        # A trace opened while this callback traces a phase, or while a local run
+        # calls its metric, is the metric's or the proposer's, never a rollout, even
+        # when its module is the program's own class.
+        phase = self._optimization_phase or context.get_current_optimization_phase()
+        if phase is not None:
+            trace.tags["optimization_phase"] = phase
 
         return trace
 
@@ -1401,7 +1426,9 @@ class _LMMixin(_CallbackBase):
         span.model_name = getattr(instance, "model", None)
         span.model_provider = helpers.infer_provider(instance)
 
-        if self._active_optimization_run_id is not None or context.get_current_optimization_run_id() is not None:
+        if self._optimization_phase is not None:
+            span.attributes["optimization.role"] = _OPTIMIZATION_PHASE_ROLES[self._optimization_phase]
+        elif self._active_optimization_run_id is not None or context.get_current_optimization_run_id() is not None:
             span.attributes["optimization.role"] = "rollout"
 
         # The usage fields read one named value each out of the raw kwargs. The
@@ -1603,7 +1630,9 @@ class _LMMixin(_CallbackBase):
         # must NOT also be added to _orphaned_lm_usage / untraced_lm_usage — that
         # would double-count them in the optimization total.
         current_trace = context.get_current_trace()
-        if current_trace is None and self._active_optimization_run_id is not None:
+        if current_trace is None and (
+            self._active_optimization_run_id is not None or self._optimization_phase is not None
+        ):
             self._emit_optimization_lm_trace(span)
             return  # handled as a standalone GEPA-process trace
 
@@ -1644,14 +1673,18 @@ class _LMMixin(_CallbackBase):
         carries the distinction. Each becomes a standalone one-span
         trace so the prompt, response, tokens and cost are visible alongside the
         student traces instead of being reduced to an aggregate token count.
+
+        A callback built with an ``optimization_phase`` takes its phase from that
+        instead, and may have no run id: it traces the metric or the proposer of a
+        run whose id the ingest key the trace is pushed with decides.
         """
         run_id = self._active_optimization_run_id
-        if not run_id:
+        if not run_id and self._optimization_phase is None:
             return
 
-        phase = "metric" if self._active_eval_run is not None else "reflection"
+        phase = self._optimization_phase or ("metric" if self._active_eval_run is not None else "reflection")
         span.attributes["gepa.phase"] = phase
-        span.attributes["optimization.role"] = "reflection" if phase == "reflection" else "rollout"
+        span.attributes["optimization.role"] = _OPTIMIZATION_PHASE_ROLES[phase]
 
         trace = models.TraceBuilder(
             trace_id=span.trace_id or uuid4(),

{cmpnd-0.76.0 → cmpnd-0.78.0}/cmpnd/context.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
+from contextlib import contextmanager
 from contextvars import ContextVar
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any
@@ -27,6 +29,12 @@ _current_eval_run: ContextVar[models.EvalRunBuilder | None] = ContextVar("cmpnd_
 # evals). Prefixed text id (`op_…`), minted via cmpnd.identity.
 _current_optimization_run_id: ContextVar[str | None] = ContextVar("cmpnd_current_optimization_run_id", default=None)
 
+# The optimization phase the code running now serves, when that code is an
+# optimizer's metric rather than the program it optimizes. A local GEPA run sets it
+# around each metric call, so a trace the metric's grading module roots is tagged as
+# the metric's, as a hosted run's kernel tags it through its callback's own phase.
+_current_optimization_phase: ContextVar[str | None] = ContextVar("cmpnd_current_optimization_phase", default=None)
+
 
 def get_current_trace() -> models.TraceBuilder | None:
     """Get the current trace."""
@@ -79,6 +87,21 @@ def set_current_optimization_run_id(optimization_run_id: str | None) -> None:
     _current_optimization_run_id.set(optimization_run_id)
 
 
+def get_current_optimization_phase() -> str | None:
+    """The optimization phase the running code serves, or None outside a metric call."""
+    return _current_optimization_phase.get()
+
+
+@contextmanager
+def optimization_phase(phase: str) -> Generator[None]:
+    """Run the block as the optimization phase ``phase``, restoring the outer phase after."""
+    token = _current_optimization_phase.set(phase)
+    try:
+        yield
+    finally:
+        _current_optimization_phase.reset(token)
+
+
 # The program each optimization run this process started optimizes, by run id. A trace
 # linked to a run that names no program of its own takes its run's, so everything a run
 # emits is reachable through that program: GEPA's reflection and judge calls, a judge
@@ -148,4 +171,5 @@ def clear_context() -> None:
     _span_stack.set(())
     _current_eval_run.set(None)
     _current_optimization_run_id.set(None)
+    _current_optimization_phase.set(None)
     _optimization_program_ids.clear()

{cmpnd-0.76.0 → cmpnd-0.78.0}/cmpnd/identity.py

@@ -1,9 +1,10 @@
 """Identity and hash computations.
 
-The 64-bit program-identity hash is produced by the structural IR in
+A program's identity is its server-issued program id (``is_program_id``).
+The 64-bit program signature hash is produced by the structural IR in
 :mod:`cmpnd.ir`. ``compute_program_hash`` reflects a live ``dspy.Module``
-through that pipeline; trace, eval, and optimization paths share the
-one encoder.
+through that pipeline, and ``program_naming`` uses the hash only to
+generate a name for a program its user never named.
 
 Sibling hashes (``compute_signature_hash`` for model-sweep grouping,
 ``compute_valset_hash``, ``compute_dataset_hash``, ``compute_metric_hash``,
@@ -84,10 +85,9 @@ def _unsigned_to_signed_int64(value: int) -> int:  # pyright: ignore[reportUnuse
     Written as arithmetic rather than as ``ctypes.c_int64(value).value``, which
     is the identical operation spelled with a C type: mask to 64 bits, then
     subtract 2^64 from anything at or above the sign bit. ``ctypes`` is not
-    merely heavy for a two-line cast, it is *absent* where this has to run — the
-    wasm guest (`worker-wasm/guest/`) imports this module to reproduce the SDK's
-    ``program_signature_hash``, and its CPython has no ``_ctypes`` because there
-    is no libffi for wasi, so the module was never compiled. Equivalence with
+    merely heavy for a two-line cast, it is *absent* from the worker's wasm
+    kernel (`worker-wasm/guest/`), which imports this module, and whose CPython
+    has no ``_ctypes`` because there is no libffi for wasi. Equivalence with
     the ctypes form is asserted over every boundary in ``tests/test_identity.py``.
 
     Imported by ``cmpnd.ir.hash`` and exercised by ``tests/test_identity.py``;
@@ -156,8 +156,9 @@ def validate_program_id(value: object, *, source: str = "program_id") -> str:
 def compute_program_hash(module: Any) -> int | None:
     """Reflect a live ``dspy.Module`` to IR and hash it.
 
-    Returns ``None`` when ``module`` is ``None`` (callers commonly fold
-    this with ``or 0`` for trace-side DEFAULT-0 columns).
+    Returns ``None`` when ``module`` is ``None``. The worker's kernel does not
+    bundle :mod:`cmpnd.ir`, so this raises ``ModuleNotFoundError`` there; the
+    kernel generates no program names, so it never calls this.
     """
     if module is None:
         return None

{cmpnd-0.76.0 → cmpnd-0.78.0}/cmpnd/optimizers/gepa.py

@@ -205,9 +205,11 @@ class TrackedGEPA:
         existing_callbacks = self._gepa.gepa_kwargs.get("callbacks") or []
         self._gepa.gepa_kwargs["callbacks"] = existing_callbacks + [gcb]
         restore_attribution = None
+        restore_metric = None
 
         try:
             restore_attribution = gcb.start_trace_attribution()
+            restore_metric = _gepa_patch._judge_metric_calls(self._gepa)
             # Run optimization (trainset/valset are keyword-only in GEPA)
             compile_token = _gepa_patch._tracked_compile_optimizer.set(self._gepa)
             try:
@@ -246,6 +248,8 @@ class TrackedGEPA:
             raise
 
         finally:
+            if restore_metric is not None:
+                restore_metric()
             if restore_attribution is not None:
                 restore_attribution()
             # Remove our callback from gepa_kwargs

{cmpnd-0.76.0 → cmpnd-0.78.0}/docs/sdk-instrumentation.md

@@ -146,7 +146,7 @@ optimization run name their program by `program_id`.
 [`CLUSTERING.md`](../../docs/CLUSTERING.md) explains how the hash is computed.
 
 - **On a span** — `_capture_module_attributes` in `callback.py` walks `instance.signature.input_fields` / `output_fields` and records names, types, descriptions and schemas. It reflects nothing, because a span never names a program.
-- **On a trace** — `_create_trace` in `callback.py` records the root signature, or for a composite module the `dspy.Predict` children's signatures as `component_signatures`. A root module call that names no program reflects the module once, for the name the SDK generates (`program_naming.generated_program_name`). The hash goes into that name and nowhere else.
+- **On a trace** — `_create_trace` in `callback.py` records the root signature, or for a composite module the `dspy.Predict` children's signatures as `component_signatures`. A root module call that names no program reflects the module once, for the name the SDK generates (`program_naming.generated_program_name`). The hash goes into that name and nowhere else. The worker's kernel replaces that function with one that generates no name, so a kernel never reflects a module.
 
 `python-sdk/cmpnd/identity.py` is the home of every hash the SDK computes.
 
@@ -180,6 +180,7 @@ Everything lives in `python-sdk/cmpnd/context.py`:
 - `_span_stack` — an immutable tuple of `SpanBuilder`s. `push_span` and `pop_span` set a new tuple rather than mutating the old one, because a context copy shares the objects its parent context holds: a mutable stack would be one list shared by every sibling copy, so concurrent runs (GEPA's `DspyAdapter.batch_evaluate`, which runs each candidate in a `copy_context()`, or `asyncio.gather` over `acall`) would read each other's spans as their parent. `get_current_span` returns the innermost. `tests/test_trace_context_concurrency.py` pins this.
 - `_current_eval_run` (`context.py:24`) — set while an `EvalHandler` is active.
 - `_current_optimization_run_id` (`context.py:29`) — set while a patched `GEPA.compile` is running, so evals triggered inside it can be linked back.
+- `_current_optimization_phase` — set to `"metric"` for the span of each metric call a tracked GEPA compile makes, so a trace the metric's grading module roots is tagged as the metric's. `context.optimization_phase(phase)` sets it for a block.
 
 `create_child_span` at `context.py:83` is the standard constructor: it reads the current trace + span, assigns `trace_id` and `parent_span_id`, and returns an unpushed `SpanBuilder`. The callback uses it for every child span; so does the `start_span` context manager below. Pass `root=True` to have it ignore the stack and leave `parent_span_id` unset, for a span that starts a trace beneath one that does not parent it.
 
@@ -202,7 +203,41 @@ callback never fires. This prevents an ordinary evaluation after a failure from
 being queued under the failed run or being claimed as the next optimization's
 seed evidence.
 
-For tests that need a clean slate, `clear_context()` at `context.py:122` resets all four vars — see the `cmpnd-testing` skill and [`../AGENTS.md`](../AGENTS.md) § "Running Tests & Checks".
+### A callback that traces one phase of a run it does not own
+
+`CmpndCallback(optimization_phase=...)` takes `"metric"` or `"reflection"`. A worker
+kernel builds one around a hosted run's metric and around its custom GEPA proposer.
+That process knows the program the run optimizes, from its config, and not the run's
+id, which serve stamps from the ingest key the traces are pushed with. Such a callback
+makes those calls look like a local optimization's:
+
+- An LM call outside any trace becomes the standalone one-span `GEPA` trace that
+  `_emit_optimization_lm_trace` writes for a local run. Its span carries `gepa.phase`
+  and `optimization.role`, and its trace carries the `optimization_phase` tag. The
+  callback takes the phase from its constructor, not from an active eval run, and
+  leaves `optimization_run_id` unset.
+- A trace that a dspy module opens is tagged `optimization_phase`, so a program's trace
+  list tells a judge module from a rollout of the same class.
+- Every LM span gets `optimization.role`: `rollout` for a metric, `reflection` for a
+  proposer.
+
+A local run owns its run, so its callback declares no phase. It tags its judge traces
+through context instead. While a tracked GEPA compile runs, through `auto_instrument`,
+`TrackedGEPA` or a client's `optimize_program`, `_gepa_patch._judge_metric_calls`
+swaps the optimizer's `metric_fn` for a wrapper that runs each call under
+`context.optimization_phase("metric")`, and restores the metric when the compile ends.
+`_create_trace` reads that phase when its callback has none. So a judge module the
+metric grades through roots a trace tagged `optimization_phase: metric`, as a hosted
+judge's module trace is. The rollout the metric grades ran before the call and carries
+no tag. The context phase changes only that tag: a direct judge or reflection call is
+still the trace `_emit_optimization_lm_trace` builds, and module spans get no
+`optimization.role`.
+
+`tests/test_optimization_phase_callback.py` drives both paths, and the context phase,
+through real DSPy dispatch. `tests/test_run_attribution_paths.py` runs a real GEPA
+search whose metric grades through a module.
+
+For tests that need a clean slate, `clear_context()` in `context.py` resets these vars — see the `cmpnd-testing` skill and [`../AGENTS.md`](../AGENTS.md) § "Running Tests & Checks".
 
 ### Across a real fork
 
@@ -1035,7 +1070,7 @@ A new hand-rolled serializer on any of those paths reopens the hole. `tests/test
 
 ## See also
 
-- [`./clustering.md`](../../docs/CLUSTERING.md) — the canonical explainer for `signature_hash`. Start here if you're wondering why two identical-looking modules have different hashes. A hash never decides a trace's program.
+- [`CLUSTERING.md`](../../docs/CLUSTERING.md) explains how traces and runs group under a program, and how the program signature hash is computed. A hash never decides a trace's program.
 - [`./optimization-tracking.md`](../../docs/optimization-tracking.md) — how `auto_instrument`'s GEPA patch and `cmpnd/optimizers/tracker.py` turn a `GEPA.compile` call into an `OptimizationRun` with linked evals.
 - The `cmpnd-testing` skill + [`../AGENTS.md`](../AGENTS.md) § "Running Tests & Checks" — patterns for writing callback tests against real databases, including `clear_context()` usage.
 - [`serve/DB.md`](../../serve/DB.md) — the Postgres schema these payloads land in.

{cmpnd-0.76.0 → cmpnd-0.78.0}/pyproject.toml

@@ -5,7 +5,7 @@ name = "cmpnd"
 # it's new) — no git tags. Bump with `uv version --bump {minor,patch}` or by
 # hand. Runtime `cmpnd.__version__` reads installed metadata, so it tracks
 # this automatically. See docs/plans/2026-07-02-sdk-pypi-tag-free-publish.md.
-version = "0.76.0"
+version = "0.78.0"
 description = "DSPy observability and deployment SDK for cmpnd"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -41,7 +41,7 @@ dependencies = [
     # Still required, and still the fast path — but `cmpnd/ir/xxh64.py` no
     # longer *needs* it: it falls back to a pure-Python xxh64 that the tests
     # check against this module, so a platform with no wheel and no compiler can
-    # still compute a program's identity. Don't drop the dependency on that
+    # still compute a program's signature hash. Don't drop the dependency on that
     # basis; the C module is an order of magnitude faster and every ordinary
     # install should have it.
     "xxhash>=3.0.0",

cmpnd-0.78.0/tests/test_optimization_phase_callback.py

@@ -0,0 +1,196 @@
+"""A callback that traces one optimization phase on behalf of a run it does not own.
+
+A worker kernel scores a hosted rollout with the author's metric and runs the author's
+GEPA proposer, and both may call a model. The kernel knows the program the run
+optimizes but not the run's id: serve stamps that from the run key the traces are
+pushed with. So the kernel builds its callback with ``optimization_phase`` and these
+calls must come out as a local optimization's judge and reflection calls do:
+
+* an LM call outside any trace is a standalone one-span ``GEPA`` trace, with
+  ``gepa.phase`` and ``optimization.role`` on its span and ``optimization_phase`` in
+  its tags;
+* a call made through a dspy module roots a trace at that module, tagged with the phase
+  so a program's trace list tells it from a rollout of the same module class;
+* every such trace carries the configured program id, and no run id.
+
+A local GEPA run owns its run, so its callback declares no phase. It runs each metric
+call under ``context.optimization_phase("metric")`` instead, and a trace the metric's
+grading module roots there carries the same tag a hosted judge trace does.
+
+Driven through real DSPy dispatch (``dspy.Predict`` and ``DummyLM``), because the
+paths under test are which callback hooks DSPy fires and with what trace in scope.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+import dspy
+import pytest
+from dspy.utils.dummies import DummyLM
+
+from cmpnd import configuration, context
+from cmpnd.configuration import configure
+from cmpnd.context import clear_context
+from tests._traced_run import make_callback, run
+
+PROGRAM_ID = "p_0123456789Ab"
+
+
+@pytest.fixture(autouse=True)
+def configured_program():
+    """The kernel's config: the SDK configured, holding the deployment's program id."""
+    config = configure(api_key="test_key", endpoint="http://localhost:8000")
+    with configuration._config_lock:
+        configuration._config = dataclasses.replace(config, program_id=PROGRAM_ID)
+    yield
+    clear_context()
+    configuration._config = None
+    from cmpnd.exporter import shutdown_exporter
+
+    shutdown_exporter()
+
+
+def _lm(answers: list[dict[str, str]]) -> DummyLM:
+    return DummyLM(answers, adapter=dspy.ChatAdapter())
+
+
+def _direct_judge(lm: Any):
+    """A metric that asks its model directly, with no dspy module in between."""
+
+    def judge() -> Any:
+        return lm("Is 'four' right for 2+2? Say yes or no.")
+
+    return judge
+
+
+class TestADirectCallIsAStandaloneTrace:
+    @pytest.mark.parametrize(("phase", "role"), [("metric", "rollout"), ("reflection", "reflection")])
+    def test_it_is_one_gepa_span_carrying_the_phase_and_the_program(self, phase: str, role: str):
+        lm = _lm([{"answer": "yes"}])
+        _, exp = run(_direct_judge(lm), lm, dspy.ChatAdapter(), callback=make_callback(optimization_phase=phase))
+
+        assert len(exp.traces) == 1, exp.traces
+        trace = exp.traces[0]
+        assert trace.module_class == "GEPA"
+        assert trace.tags.get("optimization_phase") == phase
+        assert trace.program_id == PROGRAM_ID
+        # Serve stamps the run from the key the trace is pushed with; the kernel claims none.
+        assert trace.optimization_run_id is None
+        assert len(trace._batch_spans) == 1
+        span = trace._batch_spans[0]
+        assert span.span_type.value == "lm_call"
+        assert span.parent_span_id is None
+        assert span.trace_id == trace.trace_id
+        wire = span.to_api_dict()
+        assert wire["attributes"]["gepa.phase"] == phase
+        assert wire["attributes"]["optimization.role"] == role
+
+    def test_each_call_is_its_own_trace(self):
+        lm = _lm([{"answer": "yes"}, {"answer": "no"}])
+        judge = _direct_judge(lm)
+
+        def judge_twice() -> None:
+            judge()
+            judge()
+
+        _, exp = run(judge_twice, lm, dspy.ChatAdapter(), callback=make_callback(optimization_phase="metric"))
+
+        assert len(exp.traces) == 2
+        assert len({trace.trace_id for trace in exp.traces}) == 2
+
+
+class TestACallThroughAModuleRootsItsOwnTrace:
+    @pytest.mark.parametrize(("phase", "role"), [("metric", "rollout"), ("reflection", "reflection")])
+    def test_the_module_trace_is_tagged_with_the_phase(self, phase: str, role: str):
+        lm = _lm([{"verdict": "right"}])
+        grader = dspy.Predict("question, answer -> verdict: str")
+        _, exp = run(
+            grader,
+            lm,
+            dspy.ChatAdapter(),
+            callback=make_callback(optimization_phase=phase),
+            question="2+2?",
+            answer="four",
+        )
+
+        assert len(exp.traces) == 1, exp.traces
+        trace = exp.traces[0]
+        assert trace.module_class == "Predict"
+        assert trace.tags.get("optimization_phase") == phase
+        assert trace.program_id == PROGRAM_ID
+        assert trace.optimization_run_id is None
+        lm_spans = [span.to_api_dict() for span in exp.spans if span.span_type.value == "lm_call"]
+        assert lm_spans, exp.wire()
+        assert all(span["attributes"]["optimization.role"] == role for span in lm_spans)
+        assert all("gepa.phase" not in span["attributes"] for span in lm_spans)
+
+
+class TestWithoutAPhase:
+    def test_a_program_trace_carries_no_phase(self):
+        """A rollout's own callback declares no phase, so its traces stay rollouts."""
+        lm = _lm([{"answer": "4"}])
+        _, exp = run(dspy.Predict("question -> answer"), lm, dspy.ChatAdapter(), question="2+2?")
+
+        assert len(exp.traces) == 1
+        assert "optimization_phase" not in exp.traces[0].tags
+        assert all("optimization.role" not in span.to_api_dict()["attributes"] for span in exp.spans)
+
+    def test_a_direct_call_outside_an_optimization_is_no_trace(self):
+        lm = _lm([{"answer": "yes"}])
+        _, exp = run(_direct_judge(lm), lm, dspy.ChatAdapter())
+
+        assert exp.traces == []
+
+
+class TestALocalMetricCallsPhase:
+    """The phase a local run sets around its metric, read by a callback that declares none."""
+
+    def test_a_module_the_metric_grades_through_is_tagged_as_the_metric(self):
+        lm = _lm([{"verdict": "right"}])
+        grader = dspy.Predict("question, answer -> verdict: str")
+
+        def grade() -> Any:
+            with context.optimization_phase("metric"):
+                return grader(question="2+2?", answer="four")
+
+        _, exp = run(grade, lm, dspy.ChatAdapter())
+
+        assert len(exp.traces) == 1, exp.traces
+        assert exp.traces[0].module_class == "Predict"
+        assert exp.traces[0].tags.get("optimization_phase") == "metric"
+        # Only the trace's tag changes: its spans are built as any module trace's are.
+        assert all("optimization.role" not in span.to_api_dict()["attributes"] for span in exp.spans)
+
+    def test_a_program_traced_after_the_metric_returns_carries_no_phase(self):
+        lm = _lm([{"verdict": "right"}, {"answer": "4"}])
+        grader = dspy.Predict("question, answer -> verdict: str")
+        program = dspy.Predict("question -> answer")
+
+        def grade_then_run() -> None:
+            with context.optimization_phase("metric"):
+                grader(question="2+2?", answer="four")
+            program(question="2+2?")
+
+        _, exp = run(grade_then_run, lm, dspy.ChatAdapter())
+
+        assert [trace.tags.get("optimization_phase") for trace in exp.traces] == ["metric", None]
+
+    def test_a_direct_call_outside_an_optimization_stays_untraced(self):
+        """The phase tags module traces only: a direct call still needs a run to trace."""
+        lm = _lm([{"answer": "yes"}])
+        judge = _direct_judge(lm)
+
+        def judge_in_phase() -> None:
+            with context.optimization_phase("metric"):
+                judge()
+
+        _, exp = run(judge_in_phase, lm, dspy.ChatAdapter())
+
+        assert exp.traces == []
+
+
+def test_an_unknown_phase_is_refused():
+    with pytest.raises(ValueError, match="optimization_phase"):
+        make_callback(optimization_phase="rollout")