Skip to content

feat: add Relay-aware Cascade routing#11

Merged
bbednarski9 merged 4 commits into
NVIDIA-NeMo:topic/nemo-relay-integrationfrom
bbednarski9:feature/relay-cascade-integration
Jul 8, 2026
Merged

feat: add Relay-aware Cascade routing#11
bbednarski9 merged 4 commits into
NVIDIA-NeMo:topic/nemo-relay-integrationfrom
bbednarski9:feature/relay-cascade-integration

Conversation

@bbednarski9

@bbednarski9 bbednarski9 commented Jul 1, 2026

Copy link
Copy Markdown

What

Summary

Adds the third NeMo Relay integration layer on top of #8 and #9:

  • an owned DecisionContext that carries the typed request and one immutable, exact Relay snapshot through the existing profile registry;
  • Relay-aware Cascade decisions with typed FeatureFreshness, explicit cold-default behavior, and no selected-target dispatch;
  • exact, canonicalized (session_id, owner_id) lookup without parent, root, or session-only fallback;
  • unified typed session metadata for Decision, normal Rust endpoints, and Python profile calls;
  • strict reconciliation of proxy_x_session_id and x-nemo-relay-session-id, including repeated-value, conflict, empty-value, invalid-UTF-8, and body/header mismatch handling.

This is the Cascade and request-metadata slice of a previous end-to-end implementation, split into a focused public-repository change.

Why

The Decision API and router-neutral ATOF snapshots become end-to-end useful when an existing router can consume them safely. This PR lets Cascade route from Relay history while preserving its normal inference path, the single configured profile runtime, strict tenant/session isolation, and Relay's current body-only Decision identity behavior.

How tested

Testing summary

  • cargo fmt --all -- --check
  • cargo clippy --workspace --all-targets -- -D warnings
  • cargo check --workspace --all-targets
  • cargo test -p switchyard-components-v2 (70 unit tests plus 51 profile integration tests)
  • cargo test -p switchyard-server (1 CLI, 10 migration, and 32 server tests)
  • uv run ruff check .
  • uv run mypy switchyard (160 source files)
  • uv run pytest tests/ -q -o addopts= (1,880 passed, 50 skipped)
  • One-record Relay http_post ATOF ingestion followed by a body-only Decision request warms Cascade and returns a strong route while both selected target URLs are unreachable.
  • Cold state, turn-only state, exact owner isolation, padded/blank owner canonicalization, no owner/session fallback, header alias reconciliation, invalid UTF-8, turn-depth saturation, and unchanged normal Cascade dispatch are covered.

cargo test -p switchyard-py reaches the known macOS PyO3 cdylib link limitation (undefined _Py* symbols); the crate passes workspace check/clippy, and the binding behavior passes in the full 1,880-test Python suite.

Checklist

  • One class per file; filename = snake_case of the primary class. (N/A: Rust/Python functional API change.)
  • New public symbols exported from switchyard/__init__.py.__all__ if intended for downstream use. (N/A: Rust symbols are exported from the crate root; the existing Python module exposes the extended metadata type.)
  • Unit tests added for new components / bug fixes.
  • README / --help updated if customer-facing surface changed. (DESIGN.md documents decision-scoped resources and session metadata.)
  • Commits signed off (Signed-off-by: Your Name <email>) per the DCO.

Notes for reviewers

Dependency and delta

This PR depends on #8 and #9. Because the author has read-only access to the upstream repository, this draft must target upstream main and temporarily contains both prerequisite commits.

Review the #11 delta commits 2bb2abc0 (Cascade integration), 8425e3ab (Decision bearer auth and E2E coverage), c5701405 (envelope-validation coverage), and ae5a982a (shared Relay auth documentation), plus the bounded-memory/reporting fixes inherited from #9. Use the exact fork delta. After #9 merges, this branch will be rebased onto updated main so the PR contains only the Cascade/auth delta.

Review guidelines

Suggested review order:

  1. decision.rs, profile.rs, and registry.rs for the owned DecisionContext, typed freshness/session metadata, and same-runtime dispatch.
  2. profiles/cascade.rs for snapshot projection, cold/fresh behavior, source accounting, and the decision-only no-dispatch boundary.
  3. switchyard-server for exact canonical identity lookup and session-header/body reconciliation.
  4. The PyO3 binding and stub for the append-only session_id argument and matching alias semantics.
  5. Component and server tests for one-record ATOF-to-Decision behavior, isolation, and unchanged normal Cascade inference.

Fresh snapshots use Cascade's existing override/scorer/classifier policy. Missing or turn-only snapshots select the configured picker default immediately, report cascade_feature_cold_default, retain the existing fall_open source, omit confidence, and do not call the classifier or selected target.

Previous feedback addressed

  • Completes the requested three-PR split: Decision API contracts/profiles, ATOF subscriptions/accumulation, then Cascade and request-metadata integration.
  • Keeps one server-owned router-neutral accumulator and the same configured Arc<dyn Profile> for inference and decisions; no parallel registry, duplicate runtime, or Cascade-owned event state is introduced.
  • Replaces string-based freshness control flow with Option<FeatureFreshness> and an exhaustive Fresh match.
  • Uses exact canonical (session_id, owner_id) lookup only. Parent/root scopes and session-only state cannot satisfy an owner-specific miss; padded owners normalize and blank owners match ownerless ATOF state.
  • Preserves the previous implementation's full ATOF identity compatibility order implemented in feat: add bounded Relay ATOF snapshots #9 while making request-side identity canonical.
  • Accepts Decision requests without a session header because Relay currently sends body identity; when either supported header is present, every value must agree with the body and with the other alias.
  • Keeps legacy Relay's additive decision_profile.router field compatible while profile_id remains authoritative.
  • Makes cold Cascade behavior explicit and deterministic, records its routing source, and preserves the existing normal run() signal extraction and selected-target dispatch path.
  • Adds typed Rust/Python session metadata with an append-only Python constructor argument and matching conflict/empty-value validation.
  • Avoids new fallible .expect(), .unwrap(), or panic paths in the integration code.

Follow-up fixes

  • Applies the configured Relay bearer token consistently to both ATOF ingestion and the Decision endpoint.
  • Adds an end-to-end server test covering missing, invalid, and accepted Decision authorization.

@bbednarski9
bbednarski9 marked this pull request as ready for review July 1, 2026 04:04
@bbednarski9
bbednarski9 force-pushed the feature/relay-cascade-integration branch 3 times, most recently from df3a7e8 to fd105d2 Compare July 6, 2026 17:20
@bbednarski9
bbednarski9 changed the base branch from main to topic/nemo-relay-integration July 8, 2026 00:59
@bbednarski9
bbednarski9 requested a review from a team as a code owner July 8, 2026 00:59
@bbednarski9
bbednarski9 force-pushed the feature/relay-cascade-integration branch from fd105d2 to 32e38c5 Compare July 8, 2026 13:27
@bbednarski9
bbednarski9 force-pushed the feature/relay-cascade-integration branch from bb7f93e to c40f49c Compare July 8, 2026 13:59
@bbednarski9
bbednarski9 merged commit 5e61cb7 into NVIDIA-NeMo:topic/nemo-relay-integration Jul 8, 2026
16 checks passed
rapids-bot Bot pushed a commit to NVIDIA/NeMo-Relay that referenced this pull request Jul 13, 2026
> [!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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant