diff --git a/CHANGELOG.md b/CHANGELOG.md index 47e1a6b41..e949c4f98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- **NeMo Relay native plugin** — a dynamically loaded integration that loads + Switchyard's standard TOML deployment and executes its named routes in + process. Managed calls require NeMo Relay 0.8.0 or newer; unknown models use + Relay's continuation unchanged. + +- **NeMo Relay routing marks** — routing-model usage, measured routing + overhead, and selected-model decisions are emitted as ATOF marks. The final + serving call remains represented only by Relay's outer LLM lifecycle event to + prevent double-counting. + - **Advisor-gate routing** — new `advisor` route type pairing the serving executor with a stronger judge-only advisor that reviews terminal turns: APPROVE releases the buffered turn, REDO discards it and feeds the advisor's diff --git a/Cargo.lock b/Cargo.lock index c3029243f..f5528189e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -274,6 +274,9 @@ name = "bitflags" version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +dependencies = [ + "serde_core", +] [[package]] name = "borrow-or-share" @@ -334,6 +337,16 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "num-traits", + "serde", +] + [[package]] name = "clap" version = "4.6.2" @@ -1226,6 +1239,32 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nemo-relay-plugin" +version = "0.8.0" +source = "git+https://github.com/NVIDIA/NeMo-Relay.git?tag=0.8.0-rc.3#4b3e66e8f42b113be2b241914674082fa104cdc4" +dependencies = [ + "futures", + "nemo-relay-types", + "serde", + "serde_json", + "tokio", + "tokio-util", +] + +[[package]] +name = "nemo-relay-types" +version = "0.8.0" +source = "git+https://github.com/NVIDIA/NeMo-Relay.git?tag=0.8.0-rc.3#4b3e66e8f42b113be2b241914674082fa104cdc4" +dependencies = [ + "bitflags", + "chrono", + "serde", + "serde_json", + "typed-builder", + "uuid", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2325,6 +2364,24 @@ dependencies = [ "wiremock", ] +[[package]] +name = "switchyard-nemo-relay-plugin" +version = "0.2.0" +dependencies = [ + "futures-util", + "http", + "nemo-relay-plugin", + "serde", + "serde_json", + "switchyard-llm-client", + "switchyard-protocol", + "switchyard-runner", + "switchyard-translation", + "tempfile", + "tokio", + "toml", +] + [[package]] name = "switchyard-protocol" version = "0.2.0" @@ -2642,6 +2699,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -2815,6 +2873,26 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typed-builder" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31aa81521b70f94402501d848ccc0ecaa8f93c8eb6999eb9747e72287757ffda" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "076a02dc54dd46795c2e9c8282ed40bcfb1e22747e955de9389a1de28190fb26" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "unicode-general-category" version = "1.1.0" @@ -2857,6 +2935,18 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" +dependencies = [ + "getrandom 0.3.4", + "js-sys", + "serde", + "wasm-bindgen", +] + [[package]] name = "uuid-simd" version = "0.8.0" diff --git a/Cargo.toml b/Cargo.toml index 4a46fe6ff..1e68dcce6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "crates/prefill-router", "crates/switchyard-py", "crates/protocol", + "crates/switchyard-nemo-relay-plugin", "crates/switchyard-runner", "crates/switchyard-server", "crates/switchyard-skill-distillation", @@ -33,6 +34,7 @@ http = "1" httpdate = "1" jsonschema = { version = "0.49.4", default-features = false } jsonptr = { version = "0.8.1", default-features = false, features = ["std", "json", "resolve"] } +nemo-relay-plugin = { git = "https://github.com/NVIDIA/NeMo-Relay.git", tag = "0.8.0-rc.3" } parking_lot = "0.12" rand = "0.10" regex = "1" diff --git a/README.md b/README.md index 142f80ee7..bc83ad51e 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ algorithm you write yourself. - **Protocol Translation**: convert between OpenAI Chat, Anthropic Messages, and OpenAI Responses formats - **Multi-Backend Routing**: random routing, LLM-as-classifier routing, signal-driven stage-router, or your own algorithm - **Operational Metrics**: Prometheus metrics cover requests, errors, latency, tokens, and routing overhead +- **NeMo Relay Plugin**: run random, classifier, escalation, or stage routing in Relay while Switchyard owns provider HTTP dispatch ## Maturity @@ -123,6 +124,7 @@ configured LLM client selects one upstream format. - **[`switchyard-libsy`](crates/libsy/README.md)**: embed routing algorithms in a Rust application - **[`switchyard-protocol`](crates/protocol/README.md)**: provider-neutral request, response, and streaming types - **[`switchyard-translation`](crates/switchyard-translation/README.md)**: request, response, and stream translation +- **[`switchyard-nemo-relay-plugin`](crates/switchyard-nemo-relay-plugin/README.md)**: install Switchyard as a native NeMo Relay plugin ## Community diff --git a/crates/switchyard-nemo-relay-plugin/Cargo.toml b/crates/switchyard-nemo-relay-plugin/Cargo.toml new file mode 100644 index 000000000..283477c86 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/Cargo.toml @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "switchyard-nemo-relay-plugin" +version.workspace = true +description = "Switchyard-owned HTTP routing plugin for NeMo Relay" +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true +publish = false + +[lib] +crate-type = ["cdylib"] + +[dependencies] +futures-util.workspace = true +http.workspace = true +nemo-relay-plugin.workspace = true +serde.workspace = true +serde_json.workspace = true +switchyard-llm-client.workspace = true +switchyard-protocol.workspace = true +switchyard-runner.workspace = true +switchyard-translation.workspace = true +tokio.workspace = true +toml = "1.1" + +[dev-dependencies] +tempfile = "3" diff --git a/crates/switchyard-nemo-relay-plugin/README.md b/crates/switchyard-nemo-relay-plugin/README.md new file mode 100644 index 000000000..aa9d465ce --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/README.md @@ -0,0 +1,118 @@ +# Switchyard NeMo Relay Plugin + +`switchyard-nemo-relay-plugin` is a native NeMo Relay dynamic plugin. It loads +a standard Switchyard TOML deployment from a file or Relay's nested plugin +configuration and executes its configured routes through `switchyard-runner`. + +The plugin does not define a second routing or target configuration language. +`switchyard-server` and Relay therefore use the same targets, client pooling, +algorithm construction, retry policy, and route validation. + +## Install + +Build the platform bundle with the package script, then configure Relay to load +the generated `relay-plugin.toml` manifest. The plugin requires NeMo Relay +`>=0.8.0,<1.0`. + +## Configure Relay + +Use exactly one Switchyard deployment source. To share an existing deployment +file with `switchyard-server`, configure its path: + +```toml +[[plugins.dynamic]] +manifest = "./plugins/switchyard/relay-plugin.toml" + +[plugins.dynamic.config] +priority = 0 +switchyard_config_path = "/etc/switchyard/routes.toml" +``` + +`switchyard_config_path` is a Switchyard version-1 TOML deployment, accepted by both +`switchyard-server` and `switchyard-runner`. See the +[server configuration guide](../switchyard-server/CONFIGURATION.md) for the +deployment schema and routing algorithms. + +To keep the deployment in the Relay configuration, nest the same version-1 +Switchyard configuration under `switchyard_config`: + +```toml +[[plugins.dynamic]] +manifest = "./plugins/switchyard/relay-plugin.toml" + +[plugins.dynamic.config] +priority = 0 + +[plugins.dynamic.config.switchyard_config] +schema_version = 1 + +[plugins.dynamic.config.switchyard_config.llm_clients.primary] +format = "openai_chat" +base_url = "https://example.test/v1" + +[plugins.dynamic.config.switchyard_config.targets.default] +id = "example/model" +llm_client = "primary" + +[plugins.dynamic.config.switchyard_config.routes.default] +id = "switchyard/default" +type = "passthrough" +target = "default" +``` + +## Request handling + +For OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages calls, +the plugin decodes the Relay request and checks the requested model against the +deployment's route IDs. + +- A configured route is executed by `switchyard-runner`. +- An unknown model calls Relay's continuation unchanged. +- The returned provider response is encoded back into the caller's wire format. +- Streaming responses are returned as unpolled translated streams; Relay owns + cancellation and the outer serving-call lifecycle. + +Each route's target client must use the caller's wire format: `openai_chat`, +`openai_responses`, or `anthropic_messages`. The runner selects the upstream +backend from that format rather than translating a route to a different +provider API. When one upstream model must serve multiple caller formats, +declare a target and route for each corresponding client format. + +The plugin emits a routing request mark, routing-model call marks, measured +routing-overhead marks, and a selected-model decision mark. Token usage is +emitted as Switchyard metrics for both routing-model and answer-model calls; +Relay retains ownership of the outer LLM lifecycle. + +## Observability + +When Relay is configured with OTLP logs and metrics exporters, the plugin emits +typed telemetry through Relay's native plugin runtime: + +- Routing request, decision, and overhead marks are Info logs. +- Per-routing-model call marks are Debug logs, including their outcome and + latency, but not token usage. +- Terminal routing and response-finalization failures are Error logs. Their + payload contains only the safe Switchyard failure summary; it excludes + provider response bodies and free-form provider messages. +- Metrics use bounded attributes only: algorithm for + `switchyard.routing.requests`; outcome for `switchyard.routing.llm_calls` + and `switchyard.routing.llm_call.duration`; + and safe failure kind, category, phase, and optional upstream HTTP status for + `switchyard.routing.failures`. `switchyard.routing.overhead` records total + routing latency, including routing-model calls; durations use milliseconds. +- `switchyard.routing.llm_tokens` records normalized token usage with + `call_role` (`routing` or `answer`), configured `target_model`, and + `token_type` attributes. A provider may omit usage for streaming responses; + the plugin does not synthesize zero-value measurements. + +The plugin does not attach sessions, requests, or provider messages as metric +attributes. `target_model` comes from the configured Switchyard target set, +rather than arbitrary caller input, keeping the metric cardinality bounded by +the deployment. + +## Failure policy + +`switchyard-llm-client` owns provider retry and route-candidate fallback +behavior. The plugin does not maintain a separate trusted-default target or +rerun routing after an execution failure. Failures outside the shared runner, +including response translation failures, are returned to Relay. diff --git a/crates/switchyard-nemo-relay-plugin/config.schema.json b/crates/switchyard-nemo-relay-plugin/config.schema.json new file mode 100644 index 000000000..61f178a12 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/config.schema.json @@ -0,0 +1,40 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "Switchyard NeMo Relay Plugin", + "description": "Runs a Switchyard deployment through its shared in-process runner.", + "type": "object", + "additionalProperties": false, + "oneOf": [ + { "required": ["switchyard_config_path"] }, + { "required": ["switchyard_config"] } + ], + "properties": { + "priority": { + "type": "integer", + "default": 0, + "description": "NeMo Relay execution-intercept priority." + }, + "executor": { + "type": "object", + "additionalProperties": false, + "description": "Optional NeMo Relay SDK executor override.", + "properties": { + "worker_threads": { + "type": "integer", + "minimum": 1, + "description": "Number of worker threads for the plugin's SDK-owned executor." + } + } + }, + "switchyard_config_path": { + "type": "string", + "minLength": 1, + "description": "Path to a Switchyard version-1 TOML deployment shared with switchyard-server." + }, + "switchyard_config": { + "type": "object", + "minProperties": 1, + "description": "Inline Switchyard version-1 deployment using the same schema as switchyard-server." + } + } +} diff --git a/crates/switchyard-nemo-relay-plugin/relay-plugin.toml b/crates/switchyard-nemo-relay-plugin/relay-plugin.toml new file mode 100644 index 000000000..ba5992bd6 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/relay-plugin.toml @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +manifest_version = 1 + +[plugin] +id = "nvidia.switchyard" +kind = "rust_dynamic" + +[compat] +relay = ">=0.8.0,<1.0" +native_api = "1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_native", "config_schema"] + +[config_schema] +path = "config.schema.json" + +[source] +artifact = "" + +[integrity] +sha256 = "sha256:" + +[load] +library = "" +symbol = "nemo_relay_register_plugin" diff --git a/crates/switchyard-nemo-relay-plugin/scripts/package_bundle.py b/crates/switchyard-nemo-relay-plugin/scripts/package_bundle.py new file mode 100644 index 000000000..3b79af769 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/scripts/package_bundle.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Materialize the minimal Relay plugin bundle from a compiled cdylib.""" + +from __future__ import annotations + +import argparse +import hashlib +import shutil +import tarfile +import zipfile +from pathlib import Path + +CRATE_ROOT = Path(__file__).resolve().parents[1] +REPOSITORY_ROOT = CRATE_ROOT.parents[1] +PACKAGE_NAME = "switchyard-nemo-relay-plugin" + + +def digest(path: Path) -> str: + """Return the lowercase SHA-256 digest for a file.""" + value = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + value.update(chunk) + return value.hexdigest() + + +def archive_bundle(bundle: Path, archive: Path) -> None: + """Archive a materialized bundle under the stable package directory name.""" + if archive.parent.is_relative_to(bundle): + raise ValueError(f"bundle archive must be outside output directory: {archive}") + archive.parent.mkdir(parents=True, exist_ok=True) + if archive.exists(): + raise ValueError(f"bundle archive already exists: {archive}") + + if archive.name.endswith(".tar.gz"): + with tarfile.open(archive, "w:gz") as stream: + stream.add(bundle, arcname=PACKAGE_NAME) + return + + if archive.suffix == ".zip": + with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as stream: + for path in sorted(bundle.rglob("*")): + if path.is_file(): + stream.write(path, Path(PACKAGE_NAME) / path.relative_to(bundle)) + return + + raise ValueError("bundle archive must end in .tar.gz or .zip") + + +def main() -> None: + """Materialize a Relay-loadable plugin bundle in an empty directory.""" + parser = argparse.ArgumentParser() + parser.add_argument("--library", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--archive", type=Path) + args = parser.parse_args() + + library = args.library.resolve() + if not library.is_file(): + parser.error(f"compiled plugin library does not exist: {library}") + + manifest = (CRATE_ROOT / "relay-plugin.toml").read_text(encoding="utf-8") + placeholders = ("", "") + missing = [placeholder for placeholder in placeholders if placeholder not in manifest] + if missing: + parser.error(f"plugin manifest is missing placeholders: {', '.join(missing)}") + + output = args.output.resolve() + if output.exists() and not output.is_dir(): + parser.error(f"bundle output exists and is not a directory: {output}") + if output.is_dir() and any(output.iterdir()): + parser.error(f"bundle output directory must be empty: {output}") + + archive = args.archive.resolve() if args.archive is not None else None + if archive is not None and archive.parent.is_relative_to(output): + parser.error(f"bundle archive must be outside output directory: {archive}") + output.mkdir(parents=True, exist_ok=True) + + artifact = output / library.name + shutil.copy2(library, artifact) + shutil.copy2(CRATE_ROOT / "config.schema.json", output / "config.schema.json") + for filename in ("LICENSE", "NOTICE"): + shutil.copy2(REPOSITORY_ROOT / filename, output / filename) + + artifact_digest = digest(artifact) + manifest = manifest.replace("", artifact.name) + manifest = manifest.replace("", artifact_digest) + (output / "relay-plugin.toml").write_text(manifest, encoding="utf-8") + + if archive is None: + print(output) + return + + try: + archive_bundle(output, archive) + except ValueError as error: + parser.error(str(error)) + print(archive) + + +if __name__ == "__main__": + main() diff --git a/crates/switchyard-nemo-relay-plugin/src/config.rs b/crates/switchyard-nemo-relay-plugin/src/config.rs new file mode 100644 index 000000000..2df6386b3 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/config.rs @@ -0,0 +1,152 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::path::PathBuf; + +use serde::Deserialize; +use serde_json::{Map, Value}; +use switchyard_protocol::WireFormat; +use switchyard_runner::Runner; + +pub(crate) fn protocol_from_call(name: &str) -> Option { + match name { + "openai.chat_completions" => Some(WireFormat::OpenAiChat), + "openai.responses" => Some(WireFormat::OpenAiResponses), + "anthropic.messages" => Some(WireFormat::AnthropicMessages), + _ => None, + } +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct SwitchyardConfig { + #[serde(default)] + pub(crate) priority: i32, + #[serde(default)] + pub(crate) switchyard_config_path: Option, + #[serde(default)] + pub(crate) switchyard_config: Option>, +} + +impl SwitchyardConfig { + pub(crate) fn load_runner(&self) -> Result { + match (&self.switchyard_config_path, &self.switchyard_config) { + (Some(path), None) => Runner::load(path).map_err(|error| error.to_string()), + (None, Some(config)) => toml::to_string(config) + .map_err(|error| format!("failed to serialize Switchyard configuration: {error}")) + .and_then(|source| Runner::from_toml(&source).map_err(|error| error.to_string())), + (Some(_), Some(_)) => Err( + "configure exactly one of switchyard_config_path or switchyard_config".to_string(), + ), + (None, None) => { + Err("configure one of switchyard_config_path or switchyard_config".to_string()) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use serde_json::json; + use tempfile::NamedTempFile; + + use super::*; + + fn inline_deployment() -> Map { + json!({ + "schema_version": 1, + "llm_clients": { + "primary": { + "format": "openai_chat", + "base_url": "https://example.test/v1" + } + }, + "targets": { + "default": { + "id": "example/model", + "llm_client": "primary" + } + }, + "routes": { + "default": { + "id": "switchyard/default", + "type": "passthrough", + "target": "default" + } + } + }) + .as_object() + .unwrap() + .clone() + } + + #[test] + fn maps_supported_relay_call_names_to_wire_formats() { + let cases = [ + ("openai.chat_completions", Some(WireFormat::OpenAiChat)), + ("openai.responses", Some(WireFormat::OpenAiResponses)), + ("anthropic.messages", Some(WireFormat::AnthropicMessages)), + ("unsupported.call", None), + ]; + + for (name, expected) in cases { + assert_eq!(protocol_from_call(name), expected, "{name}"); + } + } + + #[test] + fn builds_runner_from_inline_deployment() { + let config = SwitchyardConfig { + priority: 0, + switchyard_config_path: None, + switchyard_config: Some(inline_deployment()), + }; + + let runner = config.load_runner().unwrap(); + assert!(runner.route("switchyard/default").is_some()); + } + + #[test] + fn builds_runner_from_deployment_path() { + let mut deployment = NamedTempFile::new().unwrap(); + write!( + deployment, + "{}", + toml::to_string(&inline_deployment()).unwrap() + ) + .unwrap(); + let config = SwitchyardConfig { + priority: 0, + switchyard_config_path: Some(deployment.path().to_path_buf()), + switchyard_config: None, + }; + + let runner = config.load_runner().unwrap(); + assert!(runner.route("switchyard/default").is_some()); + } + + #[test] + fn requires_exactly_one_configuration_source() { + let neither = SwitchyardConfig { + priority: 0, + switchyard_config_path: None, + switchyard_config: None, + }; + assert!(matches!( + neither.load_runner(), + Err(error) if error.contains("configure one of switchyard_config_path or switchyard_config") + )); + + let both = SwitchyardConfig { + priority: 0, + switchyard_config_path: Some(PathBuf::from("routes.toml")), + switchyard_config: Some(inline_deployment()), + }; + assert!(matches!( + both.load_runner(), + Err(error) if error.contains("configure exactly one of switchyard_config_path or switchyard_config") + )); + } +} diff --git a/crates/switchyard-nemo-relay-plugin/src/lib.rs b/crates/switchyard-nemo-relay-plugin/src/lib.rs new file mode 100644 index 000000000..79ade407a --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/lib.rs @@ -0,0 +1,156 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +mod config; +mod runtime; +mod translation; + +use std::sync::Arc; + +use nemo_relay_plugin::{ + ConfigDiagnostic, DiagnosticLevel, Json, LlmJsonAsyncStream, NativePlugin, PluginContext, + PluginRuntime, +}; +use serde_json::Map; + +use crate::config::{SwitchyardConfig, protocol_from_call}; +use crate::runtime::{SwitchyardRuntime, emit_event, emit_events}; + +#[derive(Default)] +struct SwitchyardPlugin; + +impl NativePlugin for SwitchyardPlugin { + fn plugin_kind(&self) -> &str { + "nvidia.switchyard" + } + + fn allows_multiple_components(&self) -> bool { + false + } + + fn validate(&self, plugin_config: &Map) -> Vec { + match parse_config(plugin_config).and_then(SwitchyardRuntime::new) { + Ok(_) => Vec::new(), + Err(message) => vec![ConfigDiagnostic { + level: DiagnosticLevel::Error, + code: "switchyard.invalid_config".into(), + component: Some("nvidia.switchyard".into()), + field: Some("config".into()), + message, + }], + } + } + + fn register( + &mut self, + plugin_config: &Map, + ctx: &mut PluginContext<'_>, + ) -> nemo_relay_plugin::Result<()> { + let config = parse_config(plugin_config)?; + let priority = config.priority; + let runtime = Arc::new(SwitchyardRuntime::new(config)?); + let plugin_runtime = ctx.runtime(); + register_buffered(ctx, priority, Arc::clone(&runtime), plugin_runtime.clone())?; + register_stream(ctx, priority, runtime, plugin_runtime)?; + Ok(()) + } +} + +fn register_buffered( + ctx: &mut PluginContext<'_>, + priority: i32, + runtime: Arc, + plugin_runtime: PluginRuntime, +) -> Result<(), String> { + ctx.register_llm_execution_intercept( + "switchyard.runner.buffered", + priority, + move |name, request, next| { + let runtime = Arc::clone(&runtime); + let plugin_runtime = plugin_runtime.clone(); + async move { + let Some(inbound) = protocol_from_call(&name) else { + return next.call(request).await; + }; + let decoded = runtime.decode_request(inbound, &request, false)?; + if !runtime.manages(&decoded) { + return next.call(request).await; + } + let execution = runtime.execute_buffered(inbound, decoded).await; + emit_events(&plugin_runtime, execution.events); + execution.result + } + }, + ) +} + +fn register_stream( + ctx: &mut PluginContext<'_>, + priority: i32, + runtime: Arc, + plugin_runtime: PluginRuntime, +) -> Result<(), String> { + ctx.register_llm_stream_execution_intercept( + "switchyard.runner.streaming", + priority, + move |name, request, next| { + let runtime = Arc::clone(&runtime); + let plugin_runtime = plugin_runtime.clone(); + async move { + let Some(inbound) = protocol_from_call(&name) else { + return next.call(request).await; + }; + let decoded = runtime.decode_request(inbound, &request, true)?; + if !runtime.manages(&decoded) { + return next.call(request).await; + } + let stream_plugin_runtime = plugin_runtime.clone(); + let execution = runtime + .execute_stream( + inbound, + decoded, + Arc::new(move |event| emit_event(&stream_plugin_runtime, event)), + ) + .await; + emit_events(&plugin_runtime, execution.events); + execution + .result + .map(|stream| Box::pin(stream) as LlmJsonAsyncStream) + } + }, + ) +} + +fn parse_config(plugin_config: &Map) -> Result { + let mut config = plugin_config.clone(); + config.remove("executor"); + serde_json::from_value(Json::Object(config)) + .map_err(|error| format!("invalid Switchyard configuration: {error}")) +} + +nemo_relay_plugin::nemo_relay_plugin!(nemo_relay_register_plugin, SwitchyardPlugin::default); + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn plugin_configuration_requires_a_switchyard_source() { + let config = json!({"priority": 0}); + assert!(matches!( + parse_config(config.as_object().unwrap()).and_then(SwitchyardRuntime::new), + Err(error) if error.contains("switchyard_config_path or switchyard_config") + )); + } + + #[test] + fn plugin_configuration_ignores_the_sdk_executor_override() { + let config = json!({ + "executor": {"worker_threads": 4}, + "switchyard_config_path": "routes.toml" + }); + assert!(parse_config(config.as_object().unwrap()).is_ok()); + } +} diff --git a/crates/switchyard-nemo-relay-plugin/src/runtime.rs b/crates/switchyard-nemo-relay-plugin/src/runtime.rs new file mode 100644 index 000000000..811ef7811 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/runtime.rs @@ -0,0 +1,927 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::pin::Pin; +use std::sync::{Arc, Mutex}; + +use futures_util::{Stream, StreamExt}; +use nemo_relay_plugin::{ + Json, LlmRequest as RelayRequest, LogSeverity, MetricKind, MetricMeasurement, MetricValueType, + PluginRuntime, +}; +use serde_json::{Map, json}; +use switchyard_llm_client::{LlmCallObservation, RunObservation, RunObserver}; +use switchyard_protocol::{LlmResponse, Metadata, Request, Response, WireFormat}; +use switchyard_runner::{Route, RouteErrorSummary, Runner, stream_error_summary}; +use switchyard_translation::{TranslationEngine, encode_stream}; + +use crate::config::SwitchyardConfig; +use crate::translation; + +#[derive(Debug)] +pub(crate) struct RoutingMark { + pub(crate) name: String, + pub(crate) data: Json, + pub(crate) metadata: Json, + pub(crate) severity: Option, +} + +#[derive(Debug)] +pub(crate) struct RoutingMetric { + pub(crate) name: String, + pub(crate) measurements: Vec, + pub(crate) metadata: Json, +} + +#[derive(Debug)] +pub(crate) enum RoutingEvent { + Mark(RoutingMark), + Metric(RoutingMetric), +} + +struct MetricDescriptor<'a> { + name: &'a str, + kind: MetricKind, + value_type: MetricValueType, + unit: Option<&'a str>, + description: &'a str, +} + +pub(crate) type ReturnedEventStream = Pin> + Send>>; +pub(crate) type RoutingEventEmitter = Arc; + +pub(crate) struct Execution { + pub(crate) result: Result, + pub(crate) events: Vec, +} + +pub(crate) struct SwitchyardRuntime { + runner: Runner, + translation: TranslationEngine, +} + +impl SwitchyardRuntime { + pub(crate) fn new(config: SwitchyardConfig) -> Result { + Ok(Self { + runner: config.load_runner()?, + translation: TranslationEngine::default(), + }) + } + + pub(crate) fn manages(&self, request: &Request) -> bool { + request + .llm_request + .model + .as_deref() + .is_some_and(|model| self.runner.route(model).is_some()) + } + + pub(crate) fn decode_request( + &self, + inbound: WireFormat, + request: &RelayRequest, + streaming: bool, + ) -> Result { + let mut llm_request = translation::decode_request(&self.translation, inbound, request)?; + llm_request.stream = streaming; + let headers = string_headers(&request.headers); + let mut metadata = Metadata::from_headers(&headers); + let relay_gateway_placeholder = !headers.contains_key("x-switchyard-session-id") + && headers + .get("x-nemo-relay-source") + .and_then(|value| value.to_str().ok()) + == Some("gateway") + && metadata.session_id.as_deref() == Some("gateway-gateway"); + if relay_gateway_placeholder { + metadata.session_id = None; + } + metadata.http_headers = Some(headers); + metadata.wire_format = Some(inbound); + Ok(Request { + llm_request, + raw_request: Some(request.content.clone()), + metadata: Some(metadata), + }) + } + + pub(crate) async fn execute_buffered( + &self, + inbound: WireFormat, + request: Request, + ) -> Execution { + let Execution { result, mut events } = self.execute(request).await; + let (result, finalization_failed) = match result { + Ok(response) => { + let result = finalize_buffered_response(&self.translation, inbound, response); + let failed = result.is_err(); + (result, failed) + } + Err(error) => (Err(error), false), + }; + if finalization_failed { + self.error_mark(&mut events, "response_finalization", None); + } + Execution { result, events } + } + + pub(crate) async fn execute_stream( + &self, + inbound: WireFormat, + request: Request, + emit_event: RoutingEventEmitter, + ) -> Execution { + let Execution { result, mut events } = self.execute(request).await; + let (result, finalization_failed) = match result { + Ok(response) => { + let metadata = events + .iter() + .find_map(|event| match event { + RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), + RoutingEvent::Metric(_) => None, + }) + .unwrap_or_else(|| Json::Object(Map::new())); + let result = returned_events(response, inbound, metadata, emit_event); + let failed = result.is_err(); + (result, failed) + } + Err(error) => (Err(error), false), + }; + if finalization_failed { + self.error_mark(&mut events, "response_finalization", None); + } + Execution { result, events } + } + + async fn execute(&self, request: Request) -> Execution { + let Some(route) = self.route(&request) else { + return Execution { + result: Err("Switchyard has no route for this request model".into()), + events: Vec::new(), + }; + }; + let metadata = identity_metadata(request.metadata.as_ref()); + let mut events = vec![RoutingEvent::Mark(RoutingMark { + name: "switchyard.routing.requested".into(), + data: json!({"algorithm": route.algorithm_name()}), + metadata: metadata.clone(), + severity: Some(LogSeverity::Info), + })]; + events.push(request_metric(route.algorithm_name(), metadata.clone())); + if let Err(error) = route.check_caller_format(metadata_wire_format(&request)) { + self.error_mark(&mut events, "caller_format", None); + return Execution { + result: Err(format!("Switchyard caller format is incompatible: {error}")), + events, + }; + } + let observations = Arc::new(Mutex::new(Vec::new())); + let observed = Arc::clone(&observations); + let observer: RunObserver = Arc::new(move |observation| { + observed + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(observation); + }); + match route.execute(request, Some(observer)).await { + Ok(output) => { + self.emit_observations(&mut events, take_observations(&observations), &metadata); + events.push(RoutingEvent::Mark(RoutingMark { + name: "switchyard.routing.decision".into(), + data: json!({ + "algorithm": route.algorithm_name(), + "selected_model": output.selected_model, + }), + metadata, + severity: Some(LogSeverity::Info), + })); + Execution { + result: Ok(output.response), + events, + } + } + Err(error) => { + self.emit_observations(&mut events, take_observations(&observations), &metadata); + self.route_execution_error_mark( + &mut events, + &error.execution_error_summary(), + None, + ); + Execution { + result: Err("Switchyard route execution failed".into()), + events, + } + } + } + } + + fn route(&self, request: &Request) -> Option<&Route> { + request + .llm_request + .model + .as_deref() + .and_then(|model| self.runner.route(model)) + } + + fn emit_observations( + &self, + events: &mut Vec, + observations: Vec, + metadata: &Json, + ) { + let mut call_index = 0; + for observation in observations { + match observation { + RunObservation::LlmCall(call) => { + call_index += 1; + self.routing_call_events(events, call, call_index, metadata); + } + RunObservation::RoutingOverhead(duration) => { + let latency_ms = duration.as_secs_f64() * 1_000.0; + events.push(RoutingEvent::Mark(RoutingMark { + name: "switchyard.routing.overhead".into(), + data: json!({"latency_ms": latency_ms}), + metadata: metadata.clone(), + severity: Some(LogSeverity::Info), + })); + events.push(routing_overhead_metric(latency_ms, metadata.clone())); + } + RunObservation::AnswerCall(call) => { + events.extend(token_usage_metrics("answer", &call, metadata)); + } + } + } + } + + fn routing_call_events( + &self, + events: &mut Vec, + call: LlmCallObservation, + call_index: usize, + metadata: &Json, + ) { + let outcome = if call.is_success { "ok" } else { "error" }; + let latency_ms = call.duration.as_secs_f64() * 1_000.0; + let token_metrics = token_usage_metrics("routing", &call, metadata); + events.push(RoutingEvent::Mark(RoutingMark { + name: "switchyard.routing.llm_call".into(), + data: json!({ + "call_index": call_index, + "selected_model": call.selected_model.as_str(), + "call_role": "routing", + "outcome": outcome, + "latency_ms": latency_ms, + }), + metadata: metadata.clone(), + severity: Some(LogSeverity::Debug), + })); + events.extend(routing_call_metrics(outcome, latency_ms, metadata.clone())); + events.extend(token_metrics); + } + + fn error_mark( + &self, + events: &mut Vec, + failure_kind: &str, + metadata: Option<&Json>, + ) { + let metadata = metadata + .cloned() + .unwrap_or_else(|| event_metadata(events).unwrap_or_else(|| Json::Object(Map::new()))); + events.push(RoutingEvent::Mark(RoutingMark { + name: "switchyard.routing.error".into(), + data: json!({"failure_kind": failure_kind}), + metadata: metadata.clone(), + severity: Some(LogSeverity::Error), + })); + events.push(failure_metric(failure_kind, None, None, None, metadata)); + } + + fn route_execution_error_mark( + &self, + events: &mut Vec, + summary: &RouteErrorSummary, + metadata: Option<&Json>, + ) { + let metadata = metadata + .cloned() + .unwrap_or_else(|| event_metadata(events).unwrap_or_else(|| Json::Object(Map::new()))); + events.extend(route_execution_error_events(summary, metadata)); + } +} + +pub(crate) fn emit_events(runtime: &PluginRuntime, events: Vec) { + for event in events { + emit_event(runtime, event); + } +} + +pub(crate) fn emit_event(runtime: &PluginRuntime, event: RoutingEvent) { + let result = match event { + RoutingEvent::Mark(mark) => runtime + .emit_mark_with_options( + &mark.name, + Some(&mark.data), + Some(&mark.metadata), + None, + mark.severity, + ) + .map_err(|error| ("routing mark", mark.name, error)), + RoutingEvent::Metric(metric) => runtime + .emit_metric(&metric.name, metric.measurements, Some(&metric.metadata)) + .map_err(|error| ("routing metric", metric.name, error)), + }; + if let Err((kind, name, error)) = result { + eprintln!("Switchyard could not emit {kind} {name:?}: {error}"); + } +} + +fn take_observations(observations: &Mutex>) -> Vec { + std::mem::take( + &mut *observations + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + ) +} + +fn metadata_wire_format(request: &Request) -> WireFormat { + request + .metadata + .as_ref() + .and_then(|metadata| metadata.wire_format) + .expect("decoded Relay requests always carry a wire format") +} + +fn finalize_buffered_response( + translation_engine: &TranslationEngine, + inbound: WireFormat, + response: Response, +) -> Result { + let LlmResponse::Agg(response) = response.llm_response else { + return Err("Switchyard returned a stream for a buffered request".into()); + }; + translation::encode_response(translation_engine, inbound, &response) +} + +fn returned_events( + response: Response, + inbound: WireFormat, + metadata: Json, + emit_event: RoutingEventEmitter, +) -> Result { + let served_model = response.served_model().cloned(); + let chunks = match response.llm_response { + LlmResponse::Agg(response) => response.into_stream(), + LlmResponse::Stream(chunks) => chunks, + }; + let chunks = Box::pin(chunks.map(move |item| { + if let Err(error) = &item { + for event in route_execution_error_events( + &stream_error_summary(error, served_model.as_ref()), + metadata.clone(), + ) { + emit_event(event); + } + } + item + })); + let events = encode_stream(chunks, inbound, None) + .map_err(|error| format!("Switchyard response stream setup failed: {error}"))?; + Ok(Box::pin(events.map(|item| { + item.map_err(|error| format!("Switchyard response stream failed: {error}")) + }))) +} + +fn route_execution_error_mark(summary: &RouteErrorSummary, metadata: Json) -> RoutingMark { + RoutingMark { + name: "switchyard.routing.error".into(), + data: json!({ + "failure_kind": "route_execution", + "category": summary.kind.as_str(), + "phase": summary.phase.as_str(), + "upstream_status": summary.upstream_status, + "target": summary.target.as_ref().map(|target| target.as_str()), + }), + metadata, + severity: Some(LogSeverity::Error), + } +} + +fn route_execution_error_events(summary: &RouteErrorSummary, metadata: Json) -> Vec { + vec![ + RoutingEvent::Mark(route_execution_error_mark(summary, metadata.clone())), + failure_metric( + "route_execution", + Some(summary.kind.as_str()), + Some(summary.phase.as_str()), + summary.upstream_status, + metadata, + ), + ] +} + +fn event_metadata(events: &[RoutingEvent]) -> Option { + events.iter().find_map(|event| match event { + RoutingEvent::Mark(mark) => Some(mark.metadata.clone()), + RoutingEvent::Metric(_) => None, + }) +} + +fn request_metric(algorithm: &str, metadata: Json) -> RoutingEvent { + counter_metric( + "switchyard.routing.requests", + "Requests managed by Switchyard routing.", + json!({"algorithm": algorithm}), + metadata, + ) +} + +fn routing_call_metrics(outcome: &str, latency_ms: f64, metadata: Json) -> [RoutingEvent; 2] { + let attributes = json!({"outcome": outcome}); + [ + counter_metric( + "switchyard.routing.llm_calls", + "Switchyard model calls made while routing.", + attributes.clone(), + metadata.clone(), + ), + histogram_metric( + "switchyard.routing.llm_call.duration", + "Duration of Switchyard model calls made while routing.", + latency_ms, + attributes, + metadata, + ), + ] +} + +fn routing_overhead_metric(latency_ms: f64, metadata: Json) -> RoutingEvent { + histogram_metric( + "switchyard.routing.overhead", + "Time needed to produce the Switchyard routing outcome, including routing model calls.", + latency_ms, + json!({}), + metadata, + ) +} + +fn token_usage_metrics( + call_role: &str, + call: &LlmCallObservation, + metadata: &Json, +) -> Vec { + let Some(usage) = call.usage.as_ref() else { + return Vec::new(); + }; + [ + ("input", usage.input_tokens), + ("cached_input", usage.cached_input_tokens()), + ("cache_creation_input", usage.cache_creation_input_tokens()), + ("output", usage.output_tokens), + ("reasoning", usage.reasoning_tokens), + ("total", usage.total_tokens), + ] + .into_iter() + .filter_map(|(token_type, value)| { + value.map(|value| { + metric( + MetricDescriptor { + name: "switchyard.routing.llm_tokens", + kind: MetricKind::Counter, + value_type: MetricValueType::U64, + unit: Some("{token}"), + description: "Normalized tokens used by Switchyard model calls.", + }, + json!(value), + json!({ + "call_role": call_role, + "target_model": call.selected_model.as_str(), + "token_type": token_type, + }), + metadata.clone(), + ) + }) + }) + .collect() +} + +fn failure_metric( + failure_kind: &str, + category: Option<&str>, + phase: Option<&str>, + upstream_status: Option, + metadata: Json, +) -> RoutingEvent { + let mut attributes = Map::new(); + attributes.insert("failure_kind".into(), Json::String(failure_kind.into())); + if let Some(category) = category { + attributes.insert("category".into(), Json::String(category.into())); + } + if let Some(phase) = phase { + attributes.insert("phase".into(), Json::String(phase.into())); + } + if let Some(upstream_status) = upstream_status { + attributes.insert("upstream_status".into(), Json::from(upstream_status)); + } + counter_metric( + "switchyard.routing.failures", + "Terminal Switchyard routing failures.", + Json::Object(attributes), + metadata, + ) +} + +fn counter_metric(name: &str, description: &str, attributes: Json, metadata: Json) -> RoutingEvent { + metric( + MetricDescriptor { + name, + kind: MetricKind::Counter, + value_type: MetricValueType::U64, + unit: Some("{event}"), + description, + }, + json!(1), + attributes, + metadata, + ) +} + +fn histogram_metric( + name: &str, + description: &str, + value: f64, + attributes: Json, + metadata: Json, +) -> RoutingEvent { + metric( + MetricDescriptor { + name, + kind: MetricKind::Histogram, + value_type: MetricValueType::F64, + unit: Some("ms"), + description, + }, + json!(value), + attributes, + metadata, + ) +} + +fn metric( + descriptor: MetricDescriptor<'_>, + value: Json, + attributes: Json, + metadata: Json, +) -> RoutingEvent { + RoutingEvent::Metric(RoutingMetric { + name: descriptor.name.into(), + measurements: vec![MetricMeasurement { + name: descriptor.name.into(), + kind: descriptor.kind, + value_type: descriptor.value_type, + value, + unit: descriptor.unit.map(Into::into), + description: Some(descriptor.description.into()), + attributes: Some(attributes), + boundaries: None, + }], + metadata, + }) +} + +fn string_headers(headers: &Map) -> http::HeaderMap { + let mut parsed = http::HeaderMap::with_capacity(headers.len()); + for (name, value) in headers { + let Some(value) = value.as_str() else { + continue; + }; + let (Ok(name), Ok(value)) = ( + http::HeaderName::from_bytes(name.as_bytes()), + http::HeaderValue::from_str(value), + ) else { + continue; + }; + parsed.insert(name, value); + } + parsed +} + +fn identity_metadata(metadata: Option<&Metadata>) -> Json { + json!({ + "session_id": metadata.and_then(|value| value.session_id.as_deref()), + "agent_id": metadata.and_then(|value| value.agent_id.as_deref()), + "parent_agent_id": metadata.and_then(|value| value.parent_agent_id.as_deref()), + "task_id": metadata.and_then(|value| value.task_id.as_deref()), + "turn_id": metadata.and_then(|value| value.turn_id.as_deref()), + "correlation_id": metadata.and_then(|value| value.correlation_id.as_deref()), + }) +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, HashMap}; + + use switchyard_llm_client::ClientRouter; + use switchyard_protocol::{ + LlmClientError, LlmResponseStreamEvent, ModelId, Usage, text_request, + }; + use switchyard_runner::{AlgorithmSpec, ModelCapabilities, RunnerError}; + + use super::*; + + fn runtime_for(model: &str) -> SwitchyardRuntime { + let algorithm = AlgorithmSpec::Noop {} + .build("relay", &BTreeMap::new()) + .expect("noop route should build"); + let route = Route::new( + algorithm, + ClientRouter::new(HashMap::new()), + None, + ModelCapabilities::default(), + None, + Vec::new(), + ); + SwitchyardRuntime { + runner: Runner::new(vec![(ModelId::from(model), route)]), + translation: TranslationEngine::default(), + } + } + + #[test] + fn only_configured_route_models_are_managed() { + let runtime = runtime_for("switchyard"); + let configured = Request { + llm_request: text_request(Some("switchyard".into()), "hello"), + ..Request::default() + }; + let other = Request { + llm_request: text_request(Some("other".into()), "hello"), + ..Request::default() + }; + + assert!(runtime.manages(&configured)); + assert!(!runtime.manages(&other)); + } + + #[test] + fn execution_failure_mark_uses_the_safe_runner_summary() { + let secret = "provider response body"; + let error = RunnerError::Client(LlmClientError::ContextWindowExceeded { + model: ModelId::from("weak"), + message: secret.into(), + }); + + let mark = route_execution_error_mark( + &error.execution_error_summary(), + json!({"session_id": "session"}), + ); + + assert_eq!(mark.name, "switchyard.routing.error"); + assert_eq!(mark.data["failure_kind"], "route_execution"); + assert_eq!(mark.data["category"], "context_window_exceeded"); + assert_eq!(mark.data["phase"], "before_response"); + assert_eq!(mark.data["target"], "weak"); + assert_eq!(mark.data["upstream_status"], Json::Null); + assert_eq!(mark.severity, Some(LogSeverity::Error)); + assert!(!mark.data.to_string().contains(secret)); + } + + #[test] + fn routing_observations_emit_debug_marks_and_metrics() { + let runtime = runtime_for("switchyard"); + let mut events = Vec::new(); + runtime.emit_observations( + &mut events, + vec![ + RunObservation::LlmCall(LlmCallObservation { + selected_model: ModelId::from("routing-model"), + is_success: false, + duration: std::time::Duration::from_millis(12), + usage: Some(Usage { + input_tokens: Some(4), + ..Usage::default() + }), + }), + RunObservation::RoutingOverhead(std::time::Duration::from_millis(3)), + ], + &json!({"session_id": "session"}), + ); + + assert_eq!(events.len(), 6); + let RoutingEvent::Mark(call_mark) = &events[0] else { + panic!("first event should be the routing call mark"); + }; + assert_eq!(call_mark.name, "switchyard.routing.llm_call"); + assert_eq!(call_mark.severity, Some(LogSeverity::Debug)); + assert_eq!(call_mark.data["outcome"], "error"); + assert!(call_mark.data.get("usage").is_none()); + + let RoutingEvent::Metric(call_count) = &events[1] else { + panic!("second event should be the routing call counter"); + }; + assert_eq!(call_count.name, "switchyard.routing.llm_calls"); + assert_eq!(call_count.measurements[0].kind, MetricKind::Counter); + assert_eq!( + call_count.measurements[0].attributes, + Some(json!({"outcome": "error"})) + ); + + let RoutingEvent::Metric(call_duration) = &events[2] else { + panic!("third event should be the routing call histogram"); + }; + assert_eq!(call_duration.name, "switchyard.routing.llm_call.duration"); + assert_eq!(call_duration.measurements[0].kind, MetricKind::Histogram); + assert_eq!(call_duration.measurements[0].value, json!(12.0)); + + let RoutingEvent::Metric(tokens) = &events[3] else { + panic!("fourth event should be the routing token counter"); + }; + assert_eq!(tokens.name, "switchyard.routing.llm_tokens"); + + let RoutingEvent::Mark(overhead_mark) = &events[4] else { + panic!("fifth event should be the routing overhead mark"); + }; + assert_eq!(overhead_mark.severity, Some(LogSeverity::Info)); + + let RoutingEvent::Metric(overhead) = &events[5] else { + panic!("sixth event should be the routing overhead histogram"); + }; + assert_eq!(overhead.name, "switchyard.routing.overhead"); + assert_eq!(overhead.measurements[0].attributes, Some(json!({}))); + } + + #[test] + fn request_and_failure_metrics_use_bounded_attributes() { + let RoutingEvent::Metric(request) = request_metric("stage_router", json!({})) else { + panic!("request should be a metric"); + }; + assert_eq!(request.name, "switchyard.routing.requests"); + assert_eq!( + request.measurements[0].attributes, + Some(json!({"algorithm": "stage_router"})) + ); + + let RoutingEvent::Metric(failure) = failure_metric( + "route_execution", + Some("upstream_http"), + Some("before_response"), + Some(503), + json!({}), + ) else { + panic!("failure should be a metric"); + }; + assert_eq!(failure.name, "switchyard.routing.failures"); + assert_eq!( + failure.measurements[0].attributes, + Some(json!({ + "failure_kind": "route_execution", + "category": "upstream_http", + "phase": "before_response", + "upstream_status": 503, + })) + ); + } + + #[test] + fn token_usage_metrics_distinguish_routing_and_answer_targets() { + let call = LlmCallObservation { + selected_model: ModelId::from("judge-model"), + is_success: true, + duration: std::time::Duration::from_millis(1), + usage: Some(Usage { + input_tokens: Some(11), + cache: Usage::cache_details(Some(3), Some(2)), + output_tokens: Some(7), + total_tokens: Some(23), + reasoning_tokens: Some(5), + }), + }; + + let routing = token_usage_metrics("routing", &call, &json!({"session_id": "session"})); + assert_eq!(routing.len(), 6); + for event in &routing { + let RoutingEvent::Metric(metric) = event else { + panic!("token usage should be emitted as a metric"); + }; + assert_eq!(metric.name, "switchyard.routing.llm_tokens"); + assert_eq!(metric.measurements[0].kind, MetricKind::Counter); + assert_eq!(metric.measurements[0].unit.as_deref(), Some("{token}")); + assert_eq!( + metric.measurements[0].attributes.as_ref().unwrap()["call_role"], + "routing" + ); + assert_eq!( + metric.measurements[0].attributes.as_ref().unwrap()["target_model"], + "judge-model" + ); + } + let token_values = routing + .iter() + .map(|event| { + let RoutingEvent::Metric(metric) = event else { + panic!("token usage should be emitted as a metric"); + }; + metric.measurements[0].value.clone() + }) + .collect::>(); + assert_eq!( + token_values, + vec![json!(11), json!(3), json!(2), json!(7), json!(5), json!(23)] + ); + + let answer = token_usage_metrics("answer", &call, &json!({})); + let RoutingEvent::Metric(metric) = &answer[0] else { + panic!("answer usage should be emitted as a metric"); + }; + assert_eq!( + metric.measurements[0].attributes.as_ref().unwrap()["call_role"], + "answer" + ); + } + + #[test] + fn answer_observations_emit_token_metrics_without_answer_logs() { + let runtime = runtime_for("switchyard"); + let mut events = Vec::new(); + runtime.emit_observations( + &mut events, + vec![RunObservation::AnswerCall(LlmCallObservation { + selected_model: ModelId::from("selected-target"), + is_success: true, + duration: std::time::Duration::from_millis(2), + usage: Some(Usage { + output_tokens: Some(9), + ..Usage::default() + }), + })], + &json!({}), + ); + + assert_eq!(events.len(), 1); + let RoutingEvent::Metric(metric) = &events[0] else { + panic!("answer observation should only emit a token metric"); + }; + assert_eq!(metric.name, "switchyard.routing.llm_tokens"); + assert_eq!( + metric.measurements[0].attributes, + Some(json!({ + "call_role": "answer", + "target_model": "selected-target", + "token_type": "output", + })) + ); + } + + #[tokio::test] + async fn stream_failure_emits_a_safe_failure_mark() { + let secret = "provider response body"; + let response = Response { + llm_response: LlmResponse::Stream(Box::pin(futures_util::stream::iter([Err::< + LlmResponseStreamEvent, + LlmClientError, + >( + LlmClientError::ContextWindowExceeded { + model: ModelId::from("weak"), + message: secret.into(), + }, + )]))), + metadata: Some(Metadata { + served_model: Some(ModelId::from("strong")), + ..Default::default() + }), + }; + let captured = Arc::new(Mutex::new(Vec::new())); + let emitted = Arc::clone(&captured); + let stream = returned_events( + response, + WireFormat::OpenAiChat, + json!({"session_id": "session"}), + Arc::new(move |mark| emitted.lock().unwrap().push(mark)), + ) + .expect("stream setup should succeed"); + + let events = stream.collect::>().await; + assert!(events[0].is_err()); + let events = captured.lock().unwrap(); + assert_eq!(events.len(), 2); + let RoutingEvent::Mark(mark) = &events[0] else { + panic!("first event should be the safe failure mark"); + }; + assert_eq!(mark.data["category"], "context_window_exceeded"); + assert_eq!(mark.data["phase"], "during_stream"); + assert_eq!(mark.data["target"], "strong"); + assert_eq!(mark.severity, Some(LogSeverity::Error)); + assert!(!mark.data.to_string().contains(secret)); + let RoutingEvent::Metric(metric) = &events[1] else { + panic!("second event should be the failure counter"); + }; + assert_eq!(metric.name, "switchyard.routing.failures"); + assert_eq!( + metric.measurements[0].attributes, + Some(json!({ + "failure_kind": "route_execution", + "category": "context_window_exceeded", + "phase": "during_stream", + })) + ); + } +} diff --git a/crates/switchyard-nemo-relay-plugin/src/translation.rs b/crates/switchyard-nemo-relay-plugin/src/translation.rs new file mode 100644 index 000000000..cad0c988e --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/src/translation.rs @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use nemo_relay_plugin::LlmRequest as RelayRequest; +use serde_json::Value as Json; +use switchyard_protocol::{AggLlmResponse, LlmRequest, WireFormat}; +use switchyard_translation::{ + DeterministicIdPolicy, DiagnosticSeverity, LossyConversionPolicy, PreservationPolicy, + TargetCapabilities, TranslationDiagnostic, TranslationEngine, TranslationPolicy, + UnknownFieldPolicy, +}; + +pub(crate) fn decode_request( + engine: &TranslationEngine, + protocol: WireFormat, + request: &RelayRequest, +) -> Result { + let output = engine + .decode_request(protocol, &request.content, &policy()) + .map_err(error)?; + safe(&output.diagnostics)?; + Ok(output.request) +} + +pub(crate) fn encode_response( + engine: &TranslationEngine, + protocol: WireFormat, + response: &AggLlmResponse, +) -> Result { + let output = engine + .encode_response(protocol, response, &policy()) + .map_err(error)?; + safe(&output.diagnostics)?; + Ok(output.body) +} + +fn policy() -> TranslationPolicy { + TranslationPolicy { + unknown_field_policy: UnknownFieldPolicy::Preserve, + lossy_conversion_policy: LossyConversionPolicy::Reject, + deterministic_ids: DeterministicIdPolicy::GenerateStable { + prefix: "relay".into(), + }, + preservation: PreservationPolicy::InMemory, + target_capabilities: TargetCapabilities::default(), + } +} + +fn safe(diagnostics: &[TranslationDiagnostic]) -> Result<(), String> { + let unsafe_diagnostics = diagnostics + .iter() + .filter(|diagnostic| diagnostic.severity != DiagnosticSeverity::Info) + .collect::>(); + if unsafe_diagnostics.is_empty() { + Ok(()) + } else { + Err(format!( + "Switchyard translation was not lossless: {unsafe_diagnostics:?}" + )) + } +} + +fn error(error: switchyard_translation::TranslationError) -> String { + format!("Switchyard translation failed: {error}") +} diff --git a/crates/switchyard-nemo-relay-plugin/tests/test_package_bundle.py b/crates/switchyard-nemo-relay-plugin/tests/test_package_bundle.py new file mode 100644 index 000000000..f1dc6d479 --- /dev/null +++ b/crates/switchyard-nemo-relay-plugin/tests/test_package_bundle.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the native Relay plugin bundle packager.""" + +from __future__ import annotations + +import hashlib +import subprocess +import sys +import tarfile +import tempfile +import unittest +import zipfile +from pathlib import Path + +CRATE_ROOT = Path(__file__).resolve().parents[1] +PACKAGER = CRATE_ROOT / "scripts" / "package_bundle.py" +PACKAGE_NAME = "switchyard-nemo-relay-plugin" + + +class PackageBundleTest(unittest.TestCase): + """Verify materialized and archived plugin bundle contents.""" + + def test_materializes_and_archives_supported_formats(self) -> None: + for archive_suffix in (".tar.gz", ".zip"): + with self.subTest(archive_suffix=archive_suffix), tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + library = root / "libswitchyard_nemo_relay_plugin.so" + library.write_bytes(b"compiled plugin") + output = root / "bundle" + archive = root / f"{PACKAGE_NAME}-0.2.0-linux-x86_64{archive_suffix}" + + subprocess.run( + [ + sys.executable, + str(PACKAGER), + "--library", + str(library), + "--output", + str(output), + "--archive", + str(archive), + ], + check=True, + capture_output=True, + text=True, + ) + + expected = { + "LICENSE", + "NOTICE", + "config.schema.json", + library.name, + "relay-plugin.toml", + } + self.assertEqual({path.name for path in output.iterdir()}, expected) + manifest = (output / "relay-plugin.toml").read_text(encoding="utf-8") + self.assertIn(f'artifact = "{library.name}"', manifest) + self.assertIn(hashlib.sha256(library.read_bytes()).hexdigest(), manifest) + self.assertNotIn("", manifest) + self.assertNotIn("", manifest) + self.assertEqual(self.archive_members(archive), {f"{PACKAGE_NAME}/{name}" for name in expected}) + + def test_rejects_archive_inside_output_before_materializing_bundle(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + library = root / "libswitchyard_nemo_relay_plugin.so" + library.write_bytes(b"compiled plugin") + output = root / "bundle" + archive = output / "plugin.zip" + + result = subprocess.run( + [ + sys.executable, + str(PACKAGER), + "--library", + str(library), + "--output", + str(output), + "--archive", + str(archive), + ], + check=False, + capture_output=True, + text=True, + ) + + self.assertNotEqual(result.returncode, 0) + self.assertIn("bundle archive must be outside output directory", result.stderr) + self.assertFalse(output.exists()) + self.assertFalse(archive.exists()) + + @staticmethod + def archive_members(archive: Path) -> set[str]: + """Return regular-file paths from a supported bundle archive.""" + if archive.name.endswith(".tar.gz"): + with tarfile.open(archive) as stream: + return {member.name for member in stream.getmembers() if member.isfile()} + with zipfile.ZipFile(archive) as stream: + return {member.filename for member in stream.infolist() if not member.is_dir()} + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/index.md b/docs/index.md index e68ed477b..9b8fb17f9 100644 --- a/docs/index.md +++ b/docs/index.md @@ -9,6 +9,7 @@ It supports OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages. |---|---|---| | Run Switchyard as a standalone proxy for API clients | Server Path | [Build and run the Rust server](getting_started.md#server-path) | | Add Switchyard routing to a Rust application | Library Path | [`switchyard-libsy`](../crates/libsy/README.md) | +| Add Switchyard routing to NeMo Relay | Native Plugin Path | [`switchyard-nemo-relay-plugin`](../crates/switchyard-nemo-relay-plugin/README.md) | The Server Path builds and runs the standalone `switchyard-server` binary. @@ -27,3 +28,4 @@ The Server Path builds and runs the standalone `switchyard-server` binary. - [`switchyard-libsy`](reference/rust_api.md#switchyard-libsy): embeddable routing algorithms - [`switchyard-protocol`](reference/rust_api.md#switchyard-protocol): provider-neutral API types - [`switchyard-translation`](../crates/switchyard-translation/README.md): protocol translation +- [`switchyard-nemo-relay-plugin`](../crates/switchyard-nemo-relay-plugin/README.md): native NeMo Relay integration