Skip to content

feat: expose StageRouter baseline route#35

Merged
bbednarski9 merged 3 commits into
NVIDIA-NeMo:topic/nemo-relay-integrationfrom
bbednarski9:feature/relay-cost-baseline
Jul 9, 2026
Merged

feat: expose StageRouter baseline route#35
bbednarski9 merged 3 commits into
NVIDIA-NeMo:topic/nemo-relay-integrationfrom
bbednarski9:feature/relay-cost-baseline

Conversation

@bbednarski9

@bbednarski9 bbednarski9 commented Jul 8, 2026 •
edited
Loading

Copy link
Copy Markdown

What

Add an optional baseline_route to switchyard.routing_decision.v1, have StageRouter report its configured capable target as the explicit counterfactual route, and make Relay-fed StageRouter snapshots expire deterministically.

The freshness policy uses monotonic server ingestion time and tracks the newest material routing signal separately from turn lifecycle activity. A default five-minute maximum age is configurable through --atof-max-snapshot-age-millis, SWITCHYARD_ATOF_MAX_SNAPSHOT_AGE_MILLIS, and the Python binding. Stale snapshots bypass the dimensions scorer and classifier and return the configured picker default.

Every StageRouter decision exposes additive router_metadata with feature_state (cold, fresh, or stale), snapshot age and limit, event count, turn depth, and decision source. Turn start/end events cannot refresh old routing evidence.

Other routing profiles omit baseline_route until they define a defensible counterfactual. Existing v1 clients remain compatible because all fields are additive and optional.

Why

Relay needs an explicit baseline to calculate model-routing cost effects without guessing from the client model or selected route. It also needs to know whether a Decision API result came from current ATOF evidence or a safe default after the evidence expired. Keeping both policies in Switchyard preserves routing-policy ownership.

Relates to NVIDIA/NeMo-Relay #369 and NVIDIA/NeMo-Relay #389.

Review guidelines

  1. Start with crates/switchyard-components-v2/src/decision.rs for the additive baseline contract.
  2. Review crates/switchyard-components-v2/src/profiles/stage_router.rs for material-signal timestamps, freshness state, and stale fallback.
  3. Review crates/switchyard-server/src/atof.rs and CLI/config wiring for monotonic ingestion and the maximum-age setting.
  4. Verify the policy boundary: StageRouter owns the capable baseline; Relay must not infer one. Profiles without an explicit policy omit it.
  5. Verify lifecycle-only events do not extend snapshot freshness.

How tested

  • cargo fmt --all -- --check.
  • cargo test --ignore-rust-version -p switchyard-components-v2 — 74 passed.
  • Full profile suites, including cold/fresh/stale and capable/efficient selections.
  • cargo test --ignore-rust-version -p switchyard-server --test server — 34 passed.
  • cargo test --ignore-rust-version -p switchyard-server --test profile_migration — 10 passed.
  • cargo check --ignore-rust-version -p switchyard-py.
  • cargo clippy --ignore-rust-version -p switchyard-components-v2 -p switchyard-server -p switchyard-py --all-targets -- -D warnings.
  • detect-secrets pre-commit hook after refreshing baseline line metadata.
  • Real Relay/Claude Code E2E: one cold and four fresh dimensions decisions, snapshot age metadata preserved end to end, and Sonnet/Opus committed routes observed.

Previous feedback addressed

  • Kept the Decision API schema version stable and the new fields additive.
  • Omitted baselines for profiles that cannot yet state a defensible counterfactual.
  • Replaced remaining user-facing Cascade terminology with StageRouter.
  • Added an explicit freshness limit rather than allowing accumulated Relay state to influence routing indefinitely.
  • Kept freshness based on material evidence so lifecycle traffic cannot accidentally revive an old snapshot.
  • Preserved cross-protocol selected/baseline targets and Python serialization compatibility.

Checklist

  • Unit, server integration, and real E2E coverage added.
  • Legacy payloads without baseline_route or freshness metadata remain valid.
  • No secret values are stored; the secret baseline change only refreshes line numbers.
  • Commits signed off per DCO.

@bbednarski9
feat: expose StageRouter baseline route
2aa72d5 
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
@bbednarski9 bbednarski9 mentioned this pull request Jul 8, 2026
Merged
2 tasks
@bbednarski9 bbednarski9 mentioned this pull request Jul 8, 2026
Merged
2 tasks
@bbednarski9
bbednarski9 marked this pull request as ready for review July 9, 2026 00:55
@bbednarski9
bbednarski9 requested a review from a team as a code owner July 9, 2026 00:55
rapids-bot Bot pushed a commit to NVIDIA/NeMo-Relay that referenced this pull request Jul 9, 2026
@bbednarski9
feat: add extensible LLM optimization accounting (#389)
aecaee1 
#### Overview

This adds a plugin-neutral Relay contract for recording arbitrary LLM optimizations, retaining their raw token effects, and calculating a combined baseline-versus-actual cost estimate when the managed LLM call closes.

- [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license.
- [x] I searched existing issues and open pull requests, and this does not duplicate existing work.

#### Summary

- Add an open, string-backed `LlmOptimizationKind`, schema-tagged custom payloads, shared model transitions, and explicit baseline/effective/saved token counts.
- Extend request-intercept outcomes with ordered optimization contributions and normalized responses with a typed optimization summary.
- Add a per-call recorder for first-party execution intercepts without changing native ABI v1 host structs.
- Emit canonical `nemo_relay.llm.optimization` marks parented to the managed LLM handle while preserving plugin-specific marks.
- Calculate one combined baseline-versus-actual estimate at buffered completion, streaming finalization, or failure close.
- Export complete evidence to ATOF/ATIF and focused attributes to OTEL/OpenInference.
- Carry the additive JSON contract through worker/native JSON, FFI, Python, Node, and plugin-authoring surfaces.

#### Runtime behavior

The close path performs no network, file, refresh, subprocess, exporter-flush, or async work. It uses an `Arc` snapshot of the active in-memory pricing resolver and a short-held non-async mutex. Work is bounded to 64 attempted/accepted contributions, 16 KiB per complete serialized contribution envelope, and 256 KiB total retained contribution data per call. A breached bound seals accounting for that call, retains the accepted prefix, produces a typed partial summary, and never changes a provider result.

Token counts remain first-class evidence. Downstream systems can reprice stored summaries from baseline/effective models, observed usage, explicit token savings, and pricing provenance. Checked arithmetic exposes overflow, missing observed fields, and inconsistent totals as limitations rather than clamping or fabricating evidence. Missing usage, model identity, pricing, calculable cost totals, or currency agreement produces a partial summary with null unavailable cost fields.

Exactly one applied `model_routing` contribution may define the counterfactual baseline. Relay prefers that contribution's effective model because it is the exact target committed by the execution interceptor; provider response aliases and the requested model are fallbacks only when no applied routing target exists. Multiple applied routing contributions produce `multiple_routing_contributions` and are excluded from route arithmetic while all original evidence remains visible. Non-applied decisions remain visible but have no accounting authority. When there is no routing contribution, the terminal effective model is also the baseline so compression-only savings remain calculable.

Canonical optimization marks run through the mark-sanitizer chain captured with the managed LLM scope, including streaming close/drop paths. The summary embedded in the LLM end event independently runs through Relay's lifecycle sanitizer policy. Sanitizer or subscriber-dispatch failures remain best effort, retain unacknowledged mark state through the next lifecycle boundary, and never alter the LLM result.

#### Where should the reviewer start?

1. `crates/types/src/codec/optimization.rs` — open contribution and summary wire contracts.
2. `crates/core/src/api/optimization.rs` — bounded recorder and close-time arithmetic.
3. `crates/core/src/api/llm.rs` and `crates/core/src/stream.rs` — buffered/streaming lifecycle integration.
4. `crates/core/src/observability/{atif,otel,openinference}.rs` — exporter projections.
5. `crates/core/tests/integration/middleware_tests.rs` — ordering, mark parentage, and end-summary integration.

#### Review guidelines

- Verify custom kinds and payload schemas round-trip without a Relay code change.
- Treat raw token counts as authoritative evidence and pricing as a versioned estimate.
- Check that applied routing evidence uses only the actually dispatched effective model.
- Confirm close-time work is bounded, in-memory, synchronous CPU work with asynchronous exporter delivery.
- Confirm missing data yields a partial summary rather than affecting provider behavior.
- This PR contains no Switchyard, Visor, or other plugin-specific policy or fixtures.

#### Testing summary

- Full `cargo test --workspace` after rebasing onto current `main`: passed, including 825 core library, 75 FFI unit, 65 FFI integration, 63 PII redaction, 57 middleware, 26 worker integration, 15 stream integration tests, and all doctests.
- Python: 499/499; Python plugin: 113/113 with 96.94% coverage; Python worker host E2E: 1/1.
- Node: 246/246; all Go packages passed.
- Full-envelope and attempt budgets, concurrent isolation, atomic close, sanitizer/queue failure, captured streaming sanitizer scope, interrupted streams, exactly-one routing authority, checked token arithmetic, incomplete costs, routed-model repricing, and provider-cost preservation.
- One canonical all-fields fixture round-trips through shared Rust types, native/worker SDKs, C FFI v2, Python native/worker, Node, and Go while preserving a custom kind and unknown top-level field.
- ATIF non-USD/no-optimization branches and matching OTEL/OpenInference per-cost currency and independent provenance, including complete non-USD summaries.
- A real cumulative Claude Code/InferenceHub StageRouter run produced five committed routing contributions, four Visor compression contributions, 1,868 explicit saved tokens, five complete summaries, zero provider errors/fallbacks, and exact ATOF/ATIF/OTEL parity.
- `cargo clippy --workspace --all-targets -- -D warnings`, `cargo deny check`, formatting, Ruff, `ty`, Go vet, Node formatting/docstrings, lockfile, protobuf, FFI-header, and non-attribution pre-commit checks passed.

#### Compatibility and scope

- Existing interceptors that omit `optimization_contributions` remain wire-compatible.
- Unknown optimization kinds and future top-level fields round-trip unchanged.
- Native ABI v1 structs are unchanged; the existing C helper is unchanged and an additive `_v2` helper accepts contribution JSON.
- Python native/worker, TypeScript/Node, Go, Rust worker/native, and C surfaces share the same canonical conformance fixture. The language models remain idiomatic and hand-maintained; schema generation is intentionally out of scope.
- Request contributions materialize after LLM start; execution contributions materialize at close with deterministic LLM parentage.
- Core arithmetic computes only the combined pipeline effect and never assigns monetary savings to individual plugins.

#### Review feedback addressed

- Moved new unit and wire tests out of `src/` following repository convention.
- Kept `LlmOptimizationKind` open and string-backed so third-party kinds do not require a Relay release; standard constants retain exact wire values.
- Bound the full contribution envelope and intercept aggregation, not only custom payload JSON.
- Routed canonical marks through captured-scope sanitizers and covered sanitizer/dispatcher failure behavior.
- Made model-pricing estimates follow the committed route while preserving provider-reported values.
- Added independently inspectable baseline/actual currency and provenance to OTEL/OpenInference.
- Completed Python response access, Python worker, Rust worker/native, Node, Go, and additive C FFI contribution surfaces.
- Added requested concurrency, ATIF, incomplete-cost, stream-lifecycle, token-overflow, and regression coverage.
- Made an omitted `applied` field fail-safe and consistent across Rust, Go, and Python; only explicit applied evidence can affect accounting.
- Preserved unknown future Python evidence-quality strings and prevented stale reserved fields in flattened `extra` data from reappearing on the wire.
- Timestamped execution contributions at acceptance so canonical routing marks sort strictly after their originating decision and no later than LLM end.

#### Related PRs

- Enables the accounting layer used by the Switchyard Decision API plugin in [#369](#369).
- [NVIDIA-NeMo/Switchyard #35](NVIDIA-NeMo/Switchyard#35) supplies an explicit StageRouter counterfactual baseline.



## Summary by CodeRabbit

* **New Features**
  * Added managed LLM optimization evidence capture and close-time `optimization_summary` for buffered and streaming calls.
  * Emit namespaced optimization marks/events and include finalized optimization summaries in end-of-call responses.
  * Extended request-intercept contracts to carry ordered `optimizationContributions` across Node, Python, Go, and native/FFI; added `optimization_summary`/`optimization_contributions` accessors and OpenInference/OTEL/ATIF optimization attribute exports.

* **Bug Fixes**
  * Improved stream termination/interrupt accounting and event delivery queueing to avoid incorrect/missing optimization evidence.

* **Tests**
  * Expanded unit/integration coverage with new optimization fixtures and serialization/ordering/sanitization assertions.

Authors:
  - Bryan Bednarski (https://github.com/bbednarski9)

Approvers:
  - Will Killian (https://github.com/willkill07)

URL: #389
@bbednarski9
bbednarski9 merged commit 6f5b530 into NVIDIA-NeMo:topic/nemo-relay-integration Jul 9, 2026
18 checks passed
@bbednarski9 bbednarski9 mentioned this pull request Jul 9, 2026
Merged
13 tasks
rapids-bot Bot pushed a commit to NVIDIA/NeMo-Relay that referenced this pull request Jul 13, 2026
@bbednarski9
feat(plugin)!: Add Switchyard Decision API routing plugin (#369)
658f287 
> [!WARNING]
> **BREAKING CHANGE:** **[Rust source compatibility]** This PR adds `FlowError::Upstream` and new fields to the public `AtofEndpointConfig`, `AtofEndpointSectionConfig`, and `EmitMarkEventParams` types. Downstream exhaustive `FlowError` matches must handle the new variant, and direct struct literals must initialize the new fields or migrate to the provided constructor/builder APIs. Builder-based mark emission and default Relay configurations remain compatible.

#### Overview

This adds Relay's first-party integration with the [Switchyard](https://github.com/NVIDIA-NeMo/Switchyard) Decision API. The plugin routes buffered and streaming LLM requests through an external Switchyard service, executes the selected Relay-owned backend, and records causal model-routing optimization evidence.

Switchyard owns provider-protocol translation through its `switchyard-translation` Rust library. Relay consumes that library only when the native Switchyard plugin is enabled; default Relay builds do not include either the plugin or the translation dependency.

- [x] I confirm this contribution is my own work, or I have the right to submit it under this project's license.
- [x] I searched existing issues and open pull requests, and this does not duplicate existing work.

#### Details

##### Plugin and execution behavior

- Add the standalone `nemo-relay-switchyard` first-party plugin.
- Normalize stable Relay identity and build `switchyard.routing_request.v1` requests.
- Validate selected routes and optional counterfactual baselines against exact Relay-owned backend bindings.
- Support enforce and observe-only modes, six request-materialization modes, bounded retry-aware dispatch, protocol-specific trusted fallbacks, and no retry after the first upstream stream item.
- Record one model-routing optimization contribution only after the selected route commits. Failed decisions, failed provider attempts, and fallback dispatches do not claim routing savings.
- Preserve Switchyard routing metadata, including snapshot freshness, decision source, event count, and turn depth.
- Continue using Relay's existing ATOF exporter and lifecycle; the plugin does not install local routers, feature accumulators, or a duplicate event subscriber.

##### Translation ownership and optional dependency

Request, buffered-response, and SSE translation across OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages use Switchyard's public `switchyard-translation` library instead of a Relay-local translation engine.

The dependency is pinned to the immutable topic-branch commit [`8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8`](NVIDIA-NeMo/Switchyard@8f9db9a), which includes the Decision API series, StageRouter baseline/freshness behavior, and [Switchyard #46](NVIDIA-NeMo/Switchyard#46).

The Relay CLI exposes this through `--features switchyard`:

- Without that feature, `nemo-relay-cli` contains neither `nemo-relay-switchyard` nor `switchyard-translation`.
- With that feature, the plugin and its pinned Switchyard translation library are compiled and registered.
- Relay continues to own target bindings, credentials, dispatch, retries, fallbacks, lifecycle, and observability.
- Relay retains narrow fail-open policy checks for provider extensions that cannot be translated without data loss.

##### Runtime requirements

The current experimental deployment requires a separately running Switchyard service from the pinned topic commit. The plugin does not start or supervise that service.

Switchyard availability is not a Relay startup requirement. It is a request-time dependency: connection failures, timeouts, malformed decisions, selected-route drift, unsupported extensions, and missing required identity fail open to the configured same-protocol default. ATOF-backed profiles additionally require Relay's ATOF HTTP exporter to target Switchyard's `/v1/atof/events` endpoint with preserved field names and environment-referenced authentication.

The public examples document this service boundary, verify the exact Switchyard revision before startup, and intentionally remain manual compatibility workflows rather than production orchestration.

##### Optimization accounting and observability

This PR builds on the plugin-neutral accounting contract merged in [NeMo Relay #389](#389). Applied routing contributions carry the selected and baseline model transition, decision identity, backend and tier, attempt, rollout mode, reason, and router metadata. Observe-only decisions remain visible as non-applied evidence.

Relay core combines routing and other plugin contributions at LLM close while retaining explicit token evidence for downstream repricing. Existing requested, decision, retry, fallback, and error routing marks remain available alongside the canonical optimization contribution.

The generic ATIF change that attributes `step.model_name` to the effective response model is intentionally separated into [NeMo Relay #399](#399). This PR still ensures translated provider responses carry the actual selected model; #399 independently controls how ATIF presents it.

##### Testing summary

- `cargo fmt --all -- --check`
- `cargo clippy -p nemo-relay-switchyard --all-targets -- -D warnings`
- `cargo clippy -p nemo-relay-cli --features switchyard --all-targets -- -D warnings`
- `cargo check -p nemo-relay-cli` — proves the default CLI builds without Switchyard.
- `cargo check -p nemo-relay-cli --features switchyard` — proves the opt-in dependency path.
- `cargo test -p nemo-relay-switchyard` — 31 passed.
- `cargo test -p nemo-relay-cli gateway::tests` — 33 passed without the feature.
- `cargo test -p nemo-relay-cli --features switchyard gateway::tests` — 33 passed with the feature.
- `cargo test -p nemo-relay-cli --features switchyard --test switchyard_process_e2e` — buffered, streaming, retry, and fail-open process boundary passed after the split.
- `examples/switchyard/run-real-e2e.sh` against the exact pinned Switchyard service — passed with the expected StageRouter route sequence `provider/weak`, `provider/strong`, `provider/strong`.
- The pinned Switchyard translation crate separately passes 61 request, response, streaming, extension, and round-trip tests under Rust 1.93.

##### Previous feedback addressed

- Split the implementation at repository ownership boundaries.
- Replaced Relay's custom translation engine with the Switchyard-owned reusable library.
- Kept the dependency compile-time optional for users who do not enable the Relay-native Switchyard plugin.
- Removed in-process Random, LLM, and Cascade routers, Relay-side feature accumulation, duplicate Switchyard subscriptions, and the global fallback target.
- Replaced literal credentials with environment-referenced headers and activation-time validation.
- Hardened internal dispatch headers and structured upstream-failure metadata.
- Standardized routing terminology on StageRouter.
- Standardized model-routing evidence on Relay's shared optimization contract instead of adding a second post-call savings mark.
- Moved generic ATIF effective-model attribution into focused PR #399.
- Moved Fern documentation into focused experimental docs PR #400 so this functional PR remains code-focused.
- Kept public examples free of internal-only providers, credentials, Visor source references, and private infrastructure assumptions.

#### Where should the reviewer start?

1. `crates/switchyard/src/component.rs` — Decision API lifecycle, exact target and baseline validation, retries, fallback, and terminal route commitment.
2. `crates/switchyard/src/translation.rs` and `stream_translation.rs` — thin policy adapters over the Switchyard-owned translation library.
3. `crates/cli/Cargo.toml` and the `cfg(feature = "switchyard")` registration paths — compile-time optionality.
4. `crates/cli/src/gateway.rs` — reserved dispatch headers and structured retry-aware upstream failures.
5. `crates/cli/tests/switchyard_process_e2e.rs` — process-boundary buffered, streaming, and fail-open coverage.
6. `examples/switchyard/README.md` and `run-real-e2e.sh` — pinned experimental service workflow.

#### Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

- Relates to [NVIDIA/NeMo-Relay #389](#389).
- Relates to [NVIDIA/NeMo-Relay #399](#399).
- Relates to [NVIDIA/NeMo-Relay #400](#400).
- Relates to [NVIDIA-NeMo/Switchyard #46](NVIDIA-NeMo/Switchyard#46).
- Paired with the Switchyard Decision API topic series:
  - [NVIDIA-NeMo/Switchyard #8](NVIDIA-NeMo/Switchyard#8)
  - [NVIDIA-NeMo/Switchyard #9](NVIDIA-NeMo/Switchyard#9)
  - [NVIDIA-NeMo/Switchyard #11](NVIDIA-NeMo/Switchyard#11)
  - [NVIDIA-NeMo/Switchyard #35](NVIDIA-NeMo/Switchyard#35)

Authors:
  - Bryan Bednarski (https://github.com/bbednarski9)

Approvers:
  - Will Killian (https://github.com/willkill07)
  - Maryam Najafian (https://github.com/mnajafian-nv)
  - https://github.com/Salonijain27

URL: #369
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@bbednarski9