Skip to content

feat: add bounded Relay ATOF snapshots#9

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

feat: add bounded Relay ATOF snapshots#9
bbednarski9 merged 4 commits into
NVIDIA-NeMo:topic/nemo-relay-integrationfrom
bbednarski9:feature/relay-atof-snapshots

Conversation

@bbednarski9

@bbednarski9 bbednarski9 commented Jul 1, 2026

Copy link
Copy Markdown

What

Summary

Adds the second independently functional NeMo Relay layer on top of #8:

  • a router-neutral RelaySnapshotAccumulator in switchyard-components-v2::features;
  • atomic finite-NDJSON ingestion at POST /v1/atof/events;
  • exact (session_id, owner_id) snapshots with Relay-compatible identity fallbacks;
  • deterministic identity, history, dedupe, request-size, and 64 MiB total retained-state bounds;
  • optional bearer authentication through the Rust CLI, PyO3 binding, Python wrapper, and Python CLI;
  • a checked-in Relay http_post / field_name_policy = "preserve" fixture.

This is the ATOF subscription and accumulation slice of a previous end-to-end implementation, split into a focused public-repository change.

Why

History-aware routers need one bounded, reusable snapshot source rather than per-router event state. This PR deliberately stops at ingestion and router-neutral reconstruction: Cascade projection and decision behavior remain in PR 3.

How tested

Testing summary

  • cargo fmt --all --check
  • cargo clippy --workspace --all-targets -- -D warnings
  • cargo test --workspace --exclude switchyard-py
  • cargo check -p switchyard-py
  • uv run ruff check on all changed Python/test files
  • uv run mypy switchyard/cli/switchyard_cli.py switchyard_rust/core.py switchyard_rust/server.py
  • uv run maturin develop
  • uv run pytest tests/test_serve_profile_config.py tests/test_route_bundle.py tests/test_profile_migration.py -q -o addopts= (97 passed)
  • One-record Relay-style POST, multiline atomic rejection, auth, media type, event size, batch size, count pruning, total-byte pruning/accounting, global dedupe, canonical-event validation, and every identity fallback are covered.

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; Python constants are exported from switchyard_rust.server.)
  • Unit tests added for new components / bug fixes.
  • README / --help updated if customer-facing surface changed.
  • Commits signed off (Signed-off-by: Your Name <email>) per the DCO.

Notes for reviewers

Dependency and delta

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

Review the #9 delta commits edd342d0 (snapshot memory/reporting fixes), 3857a695 (envelope-validation coverage), and e81cd9ec (shared Relay auth documentation), or use the exact fork delta. The earlier commits are prerequisites from #8 and the original ATOF implementation; after #8 merges, this branch will be rebased onto updated main.

Review guidelines

Suggested review order:

  1. features.rs for atomic preparation, canonical event validation, identity resolution, global dedupe, and retained-byte accounting.
  2. switchyard-server for optional auth, bounded body parsing, and the single server-owned accumulator.
  3. Rust/Python CLI and binding propagation for every limit and the bearer token.
  4. examples/relay_atof_switchyard.toml and CLI docs for the supported Relay transport contract.

The endpoint intentionally supports finite Relay transport = "http_post" requests (one event per POST or a finite multiline batch). Relay's long-lived transport = "ndjson" upload is out of scope because atomic batch application waits for request completion.

Previous feedback addressed

  • Replaced the Cascade-specific accumulator with a reusable router-neutral snapshot contract while honoring the selected source location in components-v2::features.
  • Bounded identity count, per-identity history, global dedupe entries, event bytes, batch bytes, and total retained-memory accounting bytes (64 MiB default).
  • Parses and semantically validates the entire batch before taking the mutation lock; malformed recognized scopes cannot partially warm state.
  • Restored global UUID/phase idempotency so changed identity metadata cannot bypass dedupe.
  • Preserved the previous implementation's full compatibility fallback order for four session and seven owner sources, with precedence tests.
  • Threads optional bearer auth through both Rust and Python startup paths; explicitly blank tokens fail configuration instead of silently disabling auth.
  • Removes fallible test .expect() usage and documents the auth/body-validation contract.
  • The Relay fixture pins http_post, field_name_policy = "preserve", and a replaceable bearer header without assuming unsupported header environment interpolation.

Relay-aware Cascade consumption, typed freshness, and unified request session metadata are intentionally deferred to #11.

Follow-up fixes

  • Uses a read/write lock so concurrent snapshot readers do not serialize behind one global mutex.
  • Makes the retained-state cap a conservative heap-accounting budget (serialized payload plus value/container overhead), rather than a wire-byte-only limit.
  • Defines accepted_events as state-mutating ingests and exposes parsed received_events separately.

@bbednarski9
bbednarski9 marked this pull request as ready for review July 1, 2026 04:04
@bbednarski9
bbednarski9 force-pushed the feature/relay-atof-snapshots branch 3 times, most recently from 4647814 to 24567e6 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-atof-snapshots branch 2 times, most recently from 32e38c5 to 5a021d0 Compare July 8, 2026 13:28
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
Signed-off-by: Bryan Bednarski <bbednarski@nvidia.com>
@bbednarski9
bbednarski9 force-pushed the feature/relay-atof-snapshots branch from 4072270 to a2f03d5 Compare July 8, 2026 13:53
@bbednarski9
bbednarski9 merged commit 7669268 into NVIDIA-NeMo:topic/nemo-relay-integration Jul 8, 2026
19 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