Repository navigation
feat(telemetry): give each rollout its own trace, linked to the job span - #4751
Conversation
The eval driver's `gym.job` span wraps a whole collection run, and every `/run` it dispatched carried that span as `traceparent`, so all rollouts of a run shared one trace. That trace grows with the run and stays open until the job ends; backends cap and buffer traces as units, so long runs lost spans and showed no root until the end. Each dispatch now opens `gym.rollout.dispatch` as the root of a new trace with a span link to `gym.job`, and the agent's `gym.rollout` span and everything below it join that trace. A run produces one small job trace and one bounded trace per rollout, all carrying the same `run_id`. The job ↔ rollout relation is kept as a link instead of a parent edge. The new root is only opened when the current span is itself a trace root, that is, when the driver owns the trace. Under an embedder's own span the dispatch nests as before, so a caller that runs the driver inside its own trace keeps that shape. Span groups no longer decide trace shape: the `per_rollout` preset is left as a plain span selection and the docs stop recommending it for long runs. Signed-off-by: Ingo Guehring <iguehring@nvidia.com>
|
🌿 Preview your docs: https://nvidia-preview-pr-4751-d78f5645f30f.docs.buildwithfern.com/nemo/gym |
iguehring-nvidia
added
feature
|
/review |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
HOLD — 3 RISK findings affect embedded trace ownership, materialized-rollout correlation, and testing without nemo-lens; 1 NOTE concerns the disabled instrumentation path.
Reviewed all six changed files and traced dispatch, HTTP propagation, span-group gating, and optional-dependency behavior. Static review only; tests were not executed.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex
|
|
||
| current = otel_trace.get_current_span() | ||
| current_ctx = current.get_span_context() | ||
| detach = current_ctx.is_valid and getattr(current, "parent", None) is None |
There was a problem hiding this comment.
RISK — A root span does not establish driver ownership.
Failure / blast radius: Configure span_groups: per_rollout and invoke run_from_config() inside an embedder's root embedder.step span. The JOB-disabled branch creates no gym.job, so the current span here is the embedder's root. This predicate nevertheless detaches it, and the outgoing /run joins a new trace instead of retaining the embedder's parent/child chain. Calling the public run_examples() under an external root has the same problem. A span link does not preserve that nesting.
Correction: Give rollout_root_span() explicit ownership information rather than inferring it from .parent. In run_from_config(), capture whether a valid external span exists before opening gym.job; set/reset a context-local owned-job SpanContext only for a standalone job. Pass that context through _dispatch_span() as an explicit helper keyword. With that plumbing, replace this predicate with:
detach = current_ctx.is_valid and current_ctx == owned_job_contextHere owned_job_context must default to None, not the current context. The invariant is that only Gym's standalone job is split; external roots and externally propagated/non-recording contexts remain parents. The predicate alone is not a drop-in fix without the dependent caller changes.
Acceptance / preservation — proposed, not executed: Add a JOB-disabled embedded-driver case and a direct run_examples() case: dispatch and HTTP spans must retain the external trace ID, with the dispatch parent equal to the external span and no links. Also exercise an attached remote NonRecordingSpan, which has no SDK .parent property and must not establish ownership. Retain the existing span tests: a standalone owned job produces a different trace ID, no parent, and one job link; an externally parented job nests without links; no current context produces an unlinked root; exceptions still end the span and restore context. Update the owned-job test to supply the explicit ownership signal.
There was a problem hiding this comment.
Agreed: a root is not ownership. Fixed in 2fc479d along the lines suggested. run_from_config records whether a valid span is already current before opening gym.job; for a standalone job it keeps the job's SpanContext in a ContextVar for the run, otherwise None. _dispatch_span passes it as owner= and rollout_root_span detaches only when current_ctx == owner. Tests added: the job group off under an embedder's root (nests, no link), an attached remote NonRecordingSpan (nests, no link), and owner set but a child of the job current (nests). The owned-job, embedded-job, no-context and exception tests are kept, the owned one now supplying the owner explicitly.
|
|
||
| attributes: Dict[str, Any] = {"nemo.gym.rollout.dispatch.name": dispatch_name} | ||
| try: | ||
| rollout_id = maybe_rollout_id_from_run_body(request_body) |
There was a problem hiding this comment.
RISK — Materialized episode dispatches lose their rollout correlation attribute.
Failure / blast radius: For a materialized-task row, the dispatcher passes _episode_request_body(row) to this helper. That conversion retains identity only as episode_id. However, maybe_rollout_id_from_run_body() reads the flat explicit ID or task/rollout indices, not episode_id. It therefore returns None, and every such gym.rollout.dispatch omits nemo.gym.rollout.id. The new trace root cannot be found by the rollout ID even though the original row contains it.
Correction: Derive telemetry identity from the original collector row while leaving the HTTP payload unchanged. At the dispatch call site, the minimal replacement is:
with _dispatch_span(row, server_name):
res = await server_client.post(server_name=server_name, url_path="/run", json=request_body)Rename the helper parameter/docstring to clarify that it receives the collector row. This preserves flat requests, where request_body already is row, and uses the same identity resolver as _episode_request_body() without adding collector-private fields to the typed wire contract.
Acceptance / preservation — proposed, not executed: Dispatch a materialized row with task index 3, rollout index 0, and attempt index 2; assert the root attribute is 3-0-a2 while the posted body still contains only episode_id and task. Repeat with an explicit rollout ID custom and attempt 2, expecting custom-a2. Retain flat-row derivation (3-0), explicit-ID precedence, dispatch-name attribution, and omission of the ID attribute for an otherwise valid flat request without identity.
There was a problem hiding this comment.
Fixed in 2fc479d: the dispatch span takes its identity from the collector row; the posted body is unchanged. Parametrized tests cover 3-0, 3-0-a2, explicit custom with attempt 2 → custom-a2, a materialized row (task_id.taskset set) → 5-1, and a row without identity, which gets no id attribute.
| dispatch = next(s for s in exporter.get_finished_spans() if s.name == "gym.rollout.dispatch") | ||
| (link,) = dispatch.links | ||
| assert link.context.span_id == job.get_span_context().span_id | ||
| assert dispatch.attributes["nemo.gym.rollout.id"] == "3-0" |
There was a problem hiding this comment.
RISK — This assertion fails in the supported nemo-lens-absent configuration.
Failure / blast radius: Run this test with the OTel SDK installed but without nemo-lens. The fallback resolver deliberately makes safe_set_span_attributes() a no-op, so dispatch.attributes has no rollout ID and this assertion raises KeyError. The topology behavior still works. This contradicts this module's SDK-only contract and the existing attribute test, which explicitly permits absent attributes without lens.
Correction: Keep the topology assertions running without lens, but distinguish the attribute contract. Import LENS_INSTALLED from the telemetry conftest and replace this assertion with:
if LENS_INSTALLED:
assert dispatch.attributes["nemo.gym.rollout.id"] == "3-0"
else:
assert not dispatch.attributesDo not change the production no-op fallback or skip the entire topology test merely to satisfy the attribute assertion.
Acceptance / preservation — proposed, not executed: Run this test in SDK-only and SDK-plus-lens environments. Both must assert the new trace ID, root parent, active-span identity, and job link. SDK-only must finish successfully with no attributes; lens-present must still assert the exact 3-0 attribute.
There was a problem hiding this comment.
Fixed in 2fc479d: the attribute assertion is now guarded by LENS_INSTALLED from the telemetry conftest, with assert not dispatch.attributes in the SDK-only configuration; the topology assertions run in both.
| try: | ||
| request_body = _episode_request_body(row) if _materialized_taskset(row) else row | ||
| res = await server_client.post(server_name=server_name, url_path="/run", json=request_body) | ||
| with _dispatch_span(request_body, server_name): |
There was a problem hiding this comment.
NOTE — Gate before constructing the dispatch context manager.
Failure / blast radius: Every rollout now constructs and enters the generator-backed _dispatch_span() even when telemetry is disabled. Its internal gate avoids imports and attributes, but not context-manager allocation or enter/exit work. This deviates from the established disabled-site contract in spans.py and the overhead regression tests. The impact is unnecessary per-dispatch overhead, not a demonstrated rollout hang or scoring failure.
Correction: Put the group check outside the with, preserving the same single HTTP call in either branch:
if is_span_group_enabled(GymSpanGroup.ROLLOUT):
with _dispatch_span(request_body, server_name):
res = await server_client.post(server_name=server_name, url_path="/run", json=request_body)
else:
res = await server_client.post(server_name=server_name, url_path="/run", json=request_body)The identity argument can independently be changed to the original row; the HTTP JSON must remain request_body. Avoid replacing this with nullcontext(), which still enters a context manager on the disabled path.
Acceptance / preservation — proposed, not executed: With ROLLOUT disabled, make _dispatch_span raise if invoked and dispatch a valid row: the request must still occur exactly once and return its unchanged result. With ROLLOUT enabled, assert exactly one dispatch span and one request. Source inspection should confirm the disabled branch performs no dispatch context-manager construction, identity extraction, or telemetry import.
There was a problem hiding this comment.
Fixed in 2fc479d: the rollout group check sits in front of the with, with the same single request in both branches, matching the other call sites; _dispatch_span no longer gates internally.
… owns Ownership is now explicit instead of inferred from the current span having no parent: run_from_config records, before opening gym.job, whether a caller's span is already active, keeps the owned job's SpanContext in a ContextVar for the run, and _dispatch_span hands it to rollout_root_span, which opens a new root only when the current span is that owner. An embedder's root with the job group off, and a remote non-recording context, now nest as they should. The dispatch span takes its identity from the collector row, not the posted body, so a materialized episode keeps its rollout id. The span group gate moved in front of the context manager, as at the other call sites, and the attribute assertion in the span test respects the nemo-lens-absent configuration. Signed-off-by: Ingo Guehring <iguehring@nvidia.com>
|
/review |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
SHIP — all three prior RISK findings and the NOTE are resolved; no new escalation-eligible findings.
Verified against the previous reviewed head:
- Driver ownership is explicit and reset after the run; the span helper preserves nesting for external and remote contexts.
- Dispatch reads correlation identity from the collector row, leaves the HTTP payload unchanged, and gates before constructing the context manager.
- The SDK-only assertion now permits absent attributes while retaining topology checks.
Inspected the changed sections of all seven files and directly affected propagation and fallback contracts. This was a follow-up review of the fixes and subsequent changes, not a fresh whole-PR review. Static inspection only; tests were not executed.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex
yaoyu-33
added
the
complexity:medium
ananthsub
left a comment
ananthsub
left a comment
There was a problem hiding this comment.
Thanks for this change. A root trace per rollout, linked back to gym.job, fixes the unbounded run-long trace that the default preset produced, and it keeps the job span instead of dropping it the way per_rollout did.
The cross-process wiring checks out. A reproduction with a real agent server and an in-memory exporter showed each gym.rollout.dispatch starting a new trace with one link to gym.job. The HTTP POST client span is its child and carries the traceparent. The agent's gym.rollout joins that trace. The context variable is cleared after the run returns, raises, or is cancelled, and nothing per rollout outlives its rollout. The four fixes from the automated review are correct.
flowchart TB
subgraph J["job trace"]
job["gym.job"]
end
subgraph R["one trace per rollout"]
d["gym.rollout.dispatch (root)"] --> h["HTTP POST"]
h --> s["POST /run on the agent server"]
s --> r["gym.rollout"]
end
d -. "span link" .-> job
PR description. Three statements no longer match the code:
- "In-process callers such as NeMo-RL see no difference" and "the set of spans and their attributes ... does not change" are not accurate. NeMo RL calls
run_examples, which reaches the same dispatch site with no owned job. So every rollout in NeMo RL's trace gains one nestedgym.rollout.dispatchspan, and the span-metrics connector gains a series for it. The nested span is useful because it carries the rollout id, so the description should say it is added rather than claim nothing changes. - The first "How" bullet still says the helper splits when the current span is itself a trace root. Since 2fc479d it splits only when the current span is the job span the driver opened.
Related open work.
- #4612 moves Gym's metrics off span-group checks. It touches the same area of
rollout_collection.pynear_record_rollout_outcome. A trial merge conflicts there, but it conflicts the same way against currentmainwithout this PR, so this PR adds no new conflict. - NVIDIA-NeMo/Lens#87 makes nemo-lens enable no span group while traces are off. Once Gym's pin includes it, the
rolloutgroup check alone becomes a cheap gate. Until then, the trace-exporter check requested inline is what keeps metrics-only runs from paying for this span, and nothing enforces that ordering. - The suggested order is this PR (with the exporter check), then #4612, then the Lens#87 pin bump. This PR does not depend on either of the other two.
Context leak noted in the description. After this change, a model call that carries another rollout's context appears as a foreign span inside the wrong rollout's trace. Before, it was a misplaced branch in the single job trace. Could you link the tracking issue from the description so readers of these traces can find it?
| try: | ||
| request_body = _episode_request_body(row) if _materialized_taskset(row) else row | ||
| res = await server_client.post(server_name=server_name, url_path="/run", json=request_body) | ||
| if is_span_group_enabled(GymSpanGroup.ROLLOUT): |
There was a problem hiding this comment.
This gate checks only the rollout span group, so the dispatch span still runs in a process that exports metrics but not traces.
At the pinned nemo-lens, setup_telemetry keeps span groups enabled while traces are off, and it installs no SDK tracer provider.
Each rollout then imports rollout_root_span, derives the rollout id, starts a non-recording span, and attaches and detaches it, all for a span that nothing receives.
No traceparent is sent either, because _traced_request skips its span when no trace exporter is active.
A reproduction with metrics on and traces off measured about 6 µs per rollout through this gate, and about 0.05 µs with the exporter check added.
The amount is small next to a /run call, but every other span-only site follows #4481 and does no span work while traces are off.
Could you add the same check that _traced_request uses, importing is_trace_exporter_active from nemo_gym.telemetry.setup?
| if is_span_group_enabled(GymSpanGroup.ROLLOUT): | |
| if is_span_group_enabled(GymSpanGroup.ROLLOUT) and is_trace_exporter_active(): |
Once NVIDIA-NeMo/Lens#87 lands and Gym bumps its pin, the group check alone becomes cheap.
Nothing enforces that ordering, so the check is needed here until then.
There was a problem hiding this comment.
Done in 5bd5391: the gate is now is_span_group_enabled(GymSpanGroup.ROLLOUT) and is_trace_exporter_active(), and a test asserts _dispatch_span is not entered while traces are off.
| request_body = _episode_request_body(row) if _materialized_taskset(row) else row | ||
| res = await server_client.post(server_name=server_name, url_path="/run", json=request_body) | ||
| if is_span_group_enabled(GymSpanGroup.ROLLOUT): | ||
| with _dispatch_span(row, server_name): |
There was a problem hiding this comment.
The span's nemo.gym.rollout.dispatch.name is set to server_name, which is not always the value the rollout metrics record under the same key.
For a row that names an agent but no environment server, server_name is the environment server that agent routes to.
gym.rollout.completed_total and gym.rollout.rewarded_total take the key from _dispatch_name(row), which is the agent name for that row, as gym_metrics.py documents.
On the verifif config, for example, the span records verifif_environment_server while the metrics record verifif_simple_agent.
A dashboard that joins span metrics to the rollout counters on this attribute then misses every row routed by agent, which is the common case.
Could you pass the value the metrics use?
| with _dispatch_span(row, server_name): | |
| with _dispatch_span(row, self._dispatch_name(row)): |
If the routed environment server is useful on the span as well, it can go under a separate key.
There was a problem hiding this comment.
Done in 5bd5391: the span records self._dispatch_name(row) under nemo.gym.rollout.dispatch.name, same as the counters; the posting server is kept under nemo.gym.rollout.dispatch.server.
| """ | ||
| from nemo_gym.telemetry.spans import rollout_root_span | ||
|
|
||
| attributes: Dict[str, Any] = {"nemo.gym.rollout.dispatch.name": dispatch_name} |
There was a problem hiding this comment.
Both attribute keys in this function are string literals, while Gym already defines constants for them.
DISPATCH_NAME_ATTRIBUTE is in nemo_gym/telemetry/gym_metrics.py, which this module already imports from, and ROLLOUT_ID_ATTRIBUTE is in nemo_gym/telemetry/endpoints.py.
Spans and metrics are meant to be joined on these keys, so renaming either constant later would silently split them.
Could you use the two constants here and at the nemo.gym.rollout.id assignment below?
There was a problem hiding this comment.
Done in 5bd5391: both keys come from DISPATCH_NAME_ATTRIBUTE and ROLLOUT_ID_ATTRIBUTE.
| attributes: Dict[str, Any] = {"nemo.gym.rollout.dispatch.name": dispatch_name} | ||
| try: | ||
| rollout_id = maybe_rollout_id_from_run_body(row) | ||
| except ValueError: |
There was a problem hiding this comment.
maybe_rollout_id_from_run_body can raise TypeError as well as ValueError, because it calls int() on the attempt index.
A list or dict there raises TypeError, which escapes _post_subroutine and aborts the run.
That only happens while the rollout group is on, since this code does not run otherwise.
The driver stamps integer attempt indexes, so this needs a malformed input row.
Telemetry should still never decide whether a rollout runs, so could you catch (TypeError, ValueError) here?
There was a problem hiding this comment.
Done in 5bd5391: catches (TypeError, ValueError); a test passes a list as attempt index and the dispatch proceeds without an id.
|
|
||
| #: ``SpanContext`` of the ``gym.job`` span the driver opened as a standalone trace root, for the | ||
| #: duration of that run; None under an embedder's trace or when the ``job`` group is off. | ||
| _OWNED_JOB_SPAN: ContextVar[Any] = ContextVar("nemo_gym_owned_job_span", default=None) |
There was a problem hiding this comment.
This variable holds a SpanContext, not a span, but both the Python name and the ContextVar name say "span".
A later reader can easily take it for the span itself and call span methods on it.
The comment also misses the traces-off case.
managed_span still yields a non-recording span when no trace exporter is active, so the stored value becomes an invalid span context rather than None.
That is harmless today because rollout_root_span checks is_valid.
Could you rename it to _OWNED_JOB_SPAN_CONTEXT, with the ContextVar name "nemo_gym_owned_job_span_context", and mention the invalid context in the comment?
ContextVar[Optional[SpanContext]], with the import under TYPE_CHECKING, would also state the contract better than Any.
There was a problem hiding this comment.
Done in 5bd5391: _OWNED_JOB_SPAN_CONTEXT, typed ContextVar[Optional[SpanContext]] with the import under TYPE_CHECKING, and the comment covers the invalid context when traces are off.
| ({"task_id": {"taskset": "t"}, "_ng_task_index": 5, "_ng_rollout_index": 1}, "5-1"), | ||
| ], | ||
| ) | ||
| def test_dispatch_span_takes_its_identity_from_the_collector_row(recorded_spans, row, expected): |
There was a problem hiding this comment.
These tests call _dispatch_span directly with no current span and no owner, so nothing tests the driver wiring.
Changing _OWNED_JOB_SPAN.set(owned) to .set(None) in run_from_config turns the split off for every run, and the telemetry and rollout collection suites still pass.
The same is true for passing owner=None at the call, dropping the reset, treating an embedded job as owned, and passing request_body instead of row.
The last one would bring back the materialized-row bug fixed in 2fc479d.
This module is also marked requires_lens, and the default unit test job does not install the telemetry extra.
Only the full test suite runs these tests.
Could you add a test that runs run_from_config with the job group on against a fake /run?
It should check three things.
Each dispatch is a new trace linked to gym.job.
The traceparent sent to /run belongs to that trace.
The context variable is None afterwards.
A version that passes on this branch, reusing helpers from test_rollout_collection.py:
async def test_standalone_run_gives_each_rollout_its_own_trace_linked_to_the_job(
tmp_path, monkeypatch, recorded_spans, empty_global_config, traces_on
):
sent = _install_recording_post(monkeypatch)
await RolloutCollectionHelper().run_from_config(_write_input(tmp_path, 3))
spans = recorded_spans()
(job,) = [s for s in spans if s.name == "gym.job"]
dispatches = [s for s in spans if s.name == "gym.rollout.dispatch"]
assert len(dispatches) == 3
assert job.parent is None
trace_ids = {d.context.trace_id for d in dispatches}
assert len(trace_ids) == 3 and job.context.trace_id not in trace_ids
for dispatch in dispatches:
assert dispatch.parent is None
(link,) = dispatch.links
assert (link.context.trace_id, link.context.span_id) == (job.context.trace_id, job.context.span_id)
# Each `/run` is sent under its rollout's dispatch span, so the agent joins that trace.
by_span_id = {d.context.span_id: d for d in dispatches}
for request in sent:
dispatch = by_span_id[request["span"].span_id]
assert request["traceparent"].split("-")[1] == f"{dispatch.context.trace_id:032x}"
assert rollout_collection._OWNED_JOB_SPAN.get() is NoneHere _install_recording_post patches the server client so that /run records the current span and the traceparent it would inject.
traces_on sets the trace exporter as active.
A second test that runs the same driver under an outer span and checks that each dispatch nests with no link would cover the embedder case.
There was a problem hiding this comment.
Done in 5bd5391: test_standalone_run_gives_each_rollout_its_own_trace_linked_to_the_job runs run_from_config against a fake /run with the three assertions, plus the embedded case (nests under the caller's span, no links) and a traces-off case (the dispatch span is never entered). Thanks for the sketch; the helpers it needed are in the same module.
| `per_rollout` omits `job` on purpose. `job` wraps a whole rollout-collection run, so | ||
| including it nests every rollout of that run under one span. That is useful for seeing a | ||
| run's overall shape and unusable for a run with thousands of rollouts. | ||
| `per_rollout` omits `job`, the span that wraps a whole rollout-collection run. Trace shape |
There was a problem hiding this comment.
This sentence says each rollout is its own trace under every preset, but two setups still produce one trace per run.
A caller that runs the driver inside its own span keeps the old nesting, which is intended.
With nemo-lens's ENABLE_SINGLE_TRACEID set, the deterministic trace id generator gives every new root the job's trace id.
The run is then one trace again, and it grows with every rollout.
Readers choosing a preset from this page will expect bounded traces in both cases.
Could you qualify the sentence to say it holds for a standalone gym eval run without ENABLE_SINGLE_TRACEID?
The per_rollout docstring in span_groups.py makes the same claim.
There was a problem hiding this comment.
Done in 5bd5391: the sentence and the preset docstring now say it holds for a standalone gym eval run with ENABLE_SINGLE_TRACEID unset, and name the two cases that still give one trace per run.
| [Span Groups](/main/observability/span-groups). | ||
| ### One Trace Per Rollout, Linked to the Job | ||
|
|
||
| The eval driver's `gym.job` span wraps a whole rollout-collection run, but rollouts are not |
There was a problem hiding this comment.
The span lists elsewhere still describe the old shape.
The Groups table in span-groups.mdx lists only gym.rollout for the rollout group.
The parent chain diagram earlier on this page still marks gym.rollout as the root.
The GymSpanGroup.ROLLOUT docstring says "driver side", which now fits gym.rollout.dispatch but not gym.rollout.
Someone searching a backend by the names in that table will not find the new root span.
Could you add gym.rollout.dispatch to the table, on the rollout collection driver once per dispatch, and show it above gym.rollout in the diagram?
There was a problem hiding this comment.
Done in 5bd5391: gym.rollout.dispatch is in the groups table and at the top of the parent-chain diagram; the ROLLOUT docstring names both spans.
| its children. Each dispatch opens `gym.rollout.dispatch` as the root of a new trace with a | ||
| span link to `gym.job`, and the agent's `gym.rollout` span and everything below it join that | ||
| trace. A run therefore produces one small job trace and one bounded trace per rollout, all | ||
| carrying the same `run_id`, instead of one trace that grows with the run and stays open until |
There was a problem hiding this comment.
Each rollout trace now gets its own sampling decision, because a new root does not inherit the job's.
Under OTEL_TRACES_SAMPLER=parentbased_traceidratio at 0.25, a reproduction dropped the gym.job span and kept about a quarter of the rollout traces.
Each of those rollout traces links to a job span that does not exist in the backend.
Sampling per rollout is arguably better, but a reader who follows the link will not expect a missing target.
Could you add a sentence saying that sampling is decided per rollout trace and that the linked gym.job span may not have been sampled?
There was a problem hiding this comment.
Done in 5bd5391: added a note that sampling is decided per rollout trace and the linked gym.job span may not have been sampled.
| _PRESETS: ClassVar[dict] = { | ||
| "default": frozenset([JOB, STARTUP]) | CROSS_PROCESS_SPINE, | ||
| # NOTE: ``per_rollout`` deliberately omits ``job`` so each rollout is its own root trace with a bounded span count. | ||
| # NOTE: ``per_rollout`` omits the run-level ``job`` and ``startup`` spans; rollouts are |
There was a problem hiding this comment.
This note now says that per_rollout omits startup, and the line three below says it again.
The new sentence also wraps across two comment lines, unlike the neighbouring lines.
Could you merge the two statements into one sentence on one line?
iguehring-nvidia
left a comment
iguehring-nvidia
left a comment
There was a problem hiding this comment.
Thanks for the thorough pass and the reproduction. All eleven inline items and the three description points are in 5bd5391; the branch is merged with current main. The driver-level tests you sketched are in test_call_sites.py for the standalone and embedded cases, plus one asserting the dispatch span is skipped while traces are off. On ordering: agreed, this PR first, then #4612, then the Lens #87 pin bump. An end-to-end run on this head with sandboxed terminal-bench 4 is in the description: 8 rollout traces linked to the job, dispatch attributes as the counters record them, all 361 sandbox and 277 engine spans inside their own rollout's tree, and with #4755 in main no span in a foreign rollout.
|
|
||
|
|
||
| #: The server a dispatch is posted to; differs from the dispatch name for a row routed by agent. | ||
| DISPATCH_SERVER_ATTRIBUTE = "nemo.gym.rollout.dispatch.server" |
There was a problem hiding this comment.
The new key nemo.gym.rollout.dispatch.server names the server that the /run call goes to, but Gym already has a key for that concept.
HTTP_DESTINATION_SERVER_NAME_ATTRIBUTE in nemo_gym/telemetry/gym_metrics.py is nemo.gym.http.destination.server.name, and the connection pool metrics record the destination server under it.
The other server keys also end in .name: nemo.gym.server.name, nemo.gym.rollout.dispatch.name, and the destination key itself.
As written, a backend user who filters by destination server needs one key for the metrics and another for this span.
This constant is also the only attribute key defined outside nemo_gym/telemetry/, and it appears in none of the observability pages.
Could you reuse HTTP_DESTINATION_SERVER_NAME_ATTRIBUTE for this attribute and drop the new constant?
There was a problem hiding this comment.
Done in d78f564: the destination is recorded under HTTP_DESTINATION_SERVER_NAME_ATTRIBUTE, the new constant is gone, and the tests assert the shared key.
| return rollout_collection | ||
|
|
||
|
|
||
| async def test_standalone_run_gives_each_rollout_its_own_trace_linked_to_the_job( |
There was a problem hiding this comment.
These tests now cover the driver wiring.
Replaying last round's mutations against this head, every one now fails a test except one.
The mutation that still passes changes the call site to give _dispatch_span the request_body instead of the row.
For a materialized row the posted body carries the identity only as episode_id, so that change drops nemo.gym.rollout.id from every materialized dispatch.
That is the bug fixed in 2fc479d.
The parametrized test above passes a materialized row to _dispatch_span directly, so it cannot tell which value the call site passes.
Could you add a call site test for a materialized row?
This one passes on this head and fails with that mutation:
async def test_materialized_dispatch_takes_its_id_from_the_row_not_the_episode_body(
monkeypatch, recorded_spans, driver_telemetry
):
"""The posted episode body carries the identity only as `episode_id`; the dispatch span still gets the id."""
import json as _json
from unittest.mock import AsyncMock
from nemo_gym.rollout_collection import RolloutCollectionConfig, RolloutCollectionHelper
from tests.unit_tests import test_rollout_collection as collection_tests
post = AsyncMock(return_value=collection_tests.FakeResponse(200, {"reward": 1.0}))
client = collection_tests.install_fake_server_client(monkeypatch, post)
client.global_config_dict = collection_tests.TestEnvironmentServerRouting._mixed_batch_config()
row = collection_tests.TestEnvironmentServerRouting._materialized_row()
config = RolloutCollectionConfig(
input_jsonl_fpath="input.jsonl",
output_jsonl_fpath="output.jsonl",
environment_server_routes={"swe_pro": "environment"},
num_repeats=1,
)
rows = RolloutCollectionHelper._preprocess_raw_rows([(0, _json.dumps(row), row)], config)
for future in RolloutCollectionHelper().run_examples(rows):
await future
(dispatch,) = [s for s in recorded_spans() if s.name == "gym.rollout.dispatch"]
assert dispatch.attributes["nemo.gym.rollout.id"] == "0-0"
assert "episode_id" in post.call_args.kwargs["json"]There was a problem hiding this comment.
Done in d78f564: added as sketched, without the extra requires_lens since the module already carries it. Confirmed that passing request_body instead of row at the call site fails it.
…er-name key and pin the materialized-row call site Review follow-ups on #4751: the span's destination server now uses HTTP_DESTINATION_SERVER_NAME_ATTRIBUTE, the key the connection-pool metrics already record it under, instead of a new constant; and a driver-level test runs a materialized row through run_examples and asserts the dispatch span takes its rollout id from the collector row, not from the posted episode body (passing the body instead fails it). Signed-off-by: Ingo Guehring <iguehring@nvidia.com>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39632969348 |
What does this PR do?
Gives each rollout its own trace, linked to the run's
gym.jobspan, instead of nesting every rollout of a run under it.Why:
gym.jobwraps a whole rollout-collection run, and every/runthe driver dispatched carried it astraceparent, so a run was one trace. That trace grows with the number of rollouts (a day-long run with engine spans reaches hundreds of thousands of spans) and stays open until the job ends. Trace backends cap and buffer a trace as one unit, so long runs lost spans past the cap and showed "root span not yet received" until the end. Theper_rolloutpreset avoided this only by not emitting thejobspan at all, which also drops the run-levelstartupspans and ties trace shape to a span-selection knob.How:
nemo_gym.telemetry.spans.linked_root_span: starts a span as the root of a new trace (context=Context()) with a span link to the current span, and attaches it. It splits only when the current span is thegym.jobspan the driver opened as a standalone root; under an embedder's span, with no owned job, under a remote non-recording context, or with a child of the job current, the span nests as before./rundispatch ingym.rollout.dispatchusing that helper, gated on therolloutspan group and an active trace exporter, with the rollout id, the dispatch name the rollout counters use (nemo.gym.rollout.dispatch.name) and the server the call went to (nemo.gym.rollout.dispatch.server) as attributes. The agent'sgym.rolloutspan and everything below it adopt the new trace through the existingtraceparentinjection.run_id, link not parent edge) and no longer recommendsper_rolloutfor long runs; the span-groups page and the preset docstrings say trace shape does not depend on the preset.What does not change: metrics, the span-metrics connector's existing series, search by
run_id(a resource attribute on every span), and Gym's servers (a server still parents its rollout under whatever context the caller sends). What changes for in-process callers of the driver such as NeMo RL: no job is owned there, so the dispatch does not split, but each rollout gains one nestedgym.rollout.dispatchspan carrying the rollout id and the dispatch name, and the span-metrics connector gains a series for it.No issue; follows the instrumentation in #2647.
Trace shape, before and after
Dotted boxes are traces. Startup spans already form their own trace today and are unchanged by this PR.
1. Today (
defaultpreset): every rollout is a child ofgym.job. One trace per run, growing with the run and open until the job ends.flowchart TB subgraph S["startup trace"] st["gym.startup"] --> ss["gym.server.startup ×N"] end subgraph J["job trace · N rollouts × ~50 spans · open until job end"] job["gym.job (root)"] --> r0["gym.rollout 0-0"] --> m0["gym.model.* · vLLM · gym.verify"] job --> r1["gym.rollout 1-0"] --> m1["gym.model.* · vLLM · gym.verify"] job --> rn["… gym.rollout N"] end style S stroke-dasharray: 4 3 style J stroke-dasharray: 4 32. Existing escape hatch (
per_rolloutpreset): bounded traces by not emittingjob. Each rollout is a root because nothing is above it; the run-leveljobandstartupspans are not emitted at all.flowchart TB subgraph T0["trace 1"] r0["gym.rollout 0-0 (root)"] --> m0["gym.model.* · vLLM · gym.verify"] end subgraph T1["trace 2"] r1["gym.rollout 1-0 (root)"] --> m1["gym.model.* · vLLM · gym.verify"] end subgraph TN["… trace N"] rn["gym.rollout N (root)"] end nojob["gym.job, gym.startup: not emitted"] style T0 stroke-dasharray: 4 3 style T1 stroke-dasharray: 4 3 style TN stroke-dasharray: 4 3 style nojob stroke-dasharray: 4 3,fill:none3. This PR: a root per rollout, linked to a kept
gym.job. The dispatch opensgym.rollout.dispatchas a new trace root with a span link (dashed) togym.job; the agent'sgym.rolloutand everything below join that trace. Same spans as today, different wiring. Applies only when the driver owns the trace; under an embedder's span the dispatch nests as before.flowchart TB subgraph S["startup trace · unchanged"] st["gym.startup"] --> ss["gym.server.startup ×N"] end subgraph J["job trace · one span · closes at job end"] job["gym.job (root)"] end subgraph T0["trace 1"] d0["gym.rollout.dispatch 0-0 (root)"] --> r0["gym.rollout 0-0"] --> m0["gym.model.* · vLLM · gym.verify"] end subgraph T1["trace 2"] d1["gym.rollout.dispatch 1-0 (root)"] --> r1["gym.rollout 1-0"] --> m1["gym.model.* · vLLM · gym.verify"] end d0 -. span link .-> job d1 -. span link .-> job style S stroke-dasharray: 4 3 style J stroke-dasharray: 4 3 style T0 stroke-dasharray: 4 3 style T1 stroke-dasharray: 4 3Validation
pytest tests/unit_tests/telemetry/test_spans.py tests/unit_tests/telemetry/test_span_groups.py tests/unit_tests/telemetry/test_propagation.py tests/unit_tests/telemetry/test_call_sites.py: 52 passed; the one failure is the pre-existing macOS-onlytest_supervisor_is_the_only_writer_of_the_active_servers_gauge, identical onmain.pytest tests/unit_tests/test_rollout_collection.py: 498 passed.pre-commit run --files <changed files>: passed.Rollout evidence
N/A for rollout behaviour: the change affects only trace ids and span links on the driver's dispatch; requests, retries and results are unchanged. The trace shape will be confirmed on the next collector validation run together with #4749.
Combined end-to-end run (added after opening)
Run on a Slurm cluster from a temporary branch merging #4749, #4750 and #4751 on top of
main(7719379), withotel.vllm_traces: trueand no hand-set flags: Nano 3.5 on vLLM 0.26.0, 2 GPUs (TP2), 16 scicode rollouts, 8 in parallel. Completed in 8 min 31 s, 16 of 16 scored, collector log clean, no token in the job directory.otlp_traces_endpoint; 72llm_requestspans arrived, each a child of the policy model's client span.Received request chatcmpl-<rollout id>-<8 hex>for all 72 requests, andgen_ai.request.idon every engine span carries the same id.gym.rollout.dispatchroots in 16 distinct traces, every one with a span link to thegym.jobspan; each trace holds itsgym.rolloutand its engine spans. Thegym.jobtrace holds only the job span and the aggregate-metrics calls; startup spans are their own trace as before. Largest trace: 271 spans, versus the whole run in one trace before.Known pre-existing issue this change makes visible
In both validation runs a small fraction of model calls carry another rollout's trace context: the agent's outgoing client span has rollout A's context and rollout-id ContextVar while the request body and URL belong to rollout B (10 of 72 engine calls on a plain-
mainrun, 5 of 72 on the combined run). The context leaks across concurrently handled rollouts inside the agent process; it predates this PR and is independent of it. Today such a call is a wrong branch inside the one job trace. After this change it appears as a foreign model call inside another rollout's trace. The rollout-id attributes on Gym's spans, and the request-id prefix from #4750, are correct on every affected span, so attribute-based correlation is unaffected; only parent-chain attribution is off by that rate. Tracked and fixed in #4755 (now merged); not a reason to keep the unbounded job trace.Review follow-up
The automated review raised four points, all taken (2fc479d): ownership of the trace is now explicit rather than inferred from the current span having no parent, so an embedder's root with the
jobgroup off and a remote non-recording context both keep nesting; the dispatch span reads its rollout id from the collector row, so materialized episodes keep it; the span-group gate sits in front of the context manager; and the attribute assertion respects the nemo-lens-absent configuration. 523 collection and propagation tests and the telemetry span, call-site and overhead tests pass.Sandboxed end-to-end run (added after opening)
Terminal-bench 4 miniswe on OpenSandbox sandboxes, 8 tasks, capture on, with the final review fixes (2fc479d): 8
gym.rollout.dispatchroots in 8 traces, every one linked togym.joband carrying its rollout id (identity now taken from the collector row), each trace holding itsgym.rollout,gym.verifyand sandbox spans. Largest trace 1108 spans. No cross-rollout mismatch among 3278 attributed spans with this agent.Second review round
All items from the human review are in: the dispatch span is gated on an active trace exporter as well as the span group; its dispatch-name attribute is the value the rollout counters record, with the posting server as a separate attribute; attribute keys come from the shared constants; a
TypeErrorin the row no longer aborts a rollout;linked_root_span(split_from=...)replacesrollout_root_span(owner=...)with typed boundaries and the ContextVar is named and typed for theSpanContextit holds; docs describe both helpers, listgym.rollout.dispatchin the groups table and diagram, qualify the per-rollout claim for embedders andENABLE_SINGLE_TRACEID, and note per-trace sampling. New tests runrun_from_configend to end against a fake/runfor the standalone and embedded cases, check thetraceparentactually sent, the cleared ContextVar, and that no dispatch span runs while traces are off. Branch merged with currentmain(which now includes #4749, #4750, #4752 and #4755).End-to-end on this head (5bd5391) with sandboxed terminal-bench 4 miniswe, 8 tasks, capture on,
vllm_traces: true: 8gym.rollout.dispatchroots in 8 traces, each linked togym.job, carrying the rollout id, the dispatch name the counters use (terminal_bench_4_miniswe) and the posting server (terminal_bench_4_miniswe_environment_server); all 361 sandbox spans and all 277 engine spans inside their own rollout's tree; no span under a rollout other than its own among 3385 attributed spans, so with #4755 inmainthe leak the description mentions no longer shows; collector clean, no token in the job directory.