Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions lib/backend-common/examples/mocker/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,7 @@ impl MockerBackend {
custom_jinja_template: args.common.custom_jinja_template,
disaggregation_mode,
route_to_encoder: args.common.route_to_encoder,
enable_rl: args.common.enable_rl,
model_name: args.model_path,
served_model_name: Some(args.model_name),
tool_call_parser,
Expand Down
4 changes: 4 additions & 0 deletions lib/backend-common/src/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,4 +90,8 @@ pub struct CommonArgs {
/// shim does not read this env var.
#[arg(long, default_value_t = false, env = "DYN_ROUTE_TO_ENCODER")]
pub route_to_encoder: bool,

/// Publish this worker's engine control/update routes on the RL request-plane endpoint.
#[arg(long, default_value_t = false, env = "DYN_ENABLE_RL")]
pub enable_rl: bool,
}
1 change: 1 addition & 0 deletions lib/backend-common/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ pub mod engine;
pub mod error;
pub mod metrics;
mod publisher;
mod rl;
pub mod run;
pub mod snapshot_publisher;
pub mod telemetry;
Expand Down
174 changes: 174 additions & 0 deletions lib/backend-common/src/rl.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,174 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::sync::Arc;

use async_trait::async_trait;
use dynamo_runtime::component::{Endpoint, StartedEndpoint};
use dynamo_runtime::engine_routes::EngineRouteRegistry;
use dynamo_runtime::pipeline::network::Ingress;
use dynamo_runtime::pipeline::{
AsyncEngine, AsyncEngineContextProvider, ManyOut, ResponseStream, SingleIn,
};
use dynamo_runtime::protocols::annotated::Annotated;
use dynamo_runtime::traits::DistributedRuntimeProvider;
use futures::stream;
use serde_json::{Value, json};

const DEFAULT_RL_ENDPOINT: &str = "rl";

pub(crate) struct RlServeEndpoint {
started: StartedEndpoint,
}

pub(crate) struct RlEndpointConfig {
endpoint_name: String,
system_url: String,
}

impl RlServeEndpoint {
pub(crate) async fn shutdown(self) -> anyhow::Result<()> {
self.started.shutdown().await
}
}

pub(crate) fn prepare_endpoint(primary: &Endpoint) -> anyhow::Result<RlEndpointConfig> {
let endpoint_name = resolve_endpoint_name(&primary.id().name)?;
let system_url = self_host_base_url(primary.drt()).ok_or_else(|| {
anyhow::anyhow!(
"RL discovery requires the Dynamo system server; set DYN_SYSTEM_PORT to 0 or a positive port"
)
})?;
Ok(RlEndpointConfig {
endpoint_name,
system_url,
})
}

pub(crate) async fn serve_endpoint(
primary: &Endpoint,
config: RlEndpointConfig,
) -> anyhow::Result<RlServeEndpoint> {
let endpoint = primary.component().endpoint(config.endpoint_name);
let handler = Arc::new(RlRouteHandler {
routes: primary.drt().engine_routes().clone(),
system_url: config.system_url,
});
let ingress = Ingress::for_engine(handler)?;
let started = endpoint
.endpoint_builder()
.handler(ingress)
.graceful_shutdown(true)
.start_with_registration()
.await?;
Ok(RlServeEndpoint { started })
}

fn self_host_base_url(drt: &dynamo_runtime::DistributedRuntime) -> Option<String> {
let info = drt.system_status_server_info()?;
let socket_addr = info.socket_addr;
if socket_addr.ip().is_unspecified() {
let host = dynamo_runtime::utils::local_ip_for_advertise();
Some(format!("http://{host}:{}", socket_addr.port()))
} else {
Some(format!("http://{socket_addr}"))
}
}

fn resolve_endpoint_name(primary_name: &str) -> anyhow::Result<String> {
let endpoint_name =
std::env::var("DYN_RL_ENDPOINT").unwrap_or_else(|_| DEFAULT_RL_ENDPOINT.into());
validate_endpoint_name(endpoint_name.trim(), primary_name)
}

fn validate_endpoint_name(endpoint_name: &str, primary_name: &str) -> anyhow::Result<String> {
if endpoint_name.is_empty() {
anyhow::bail!("DYN_RL_ENDPOINT must not be empty");
}
if endpoint_name == primary_name {
anyhow::bail!("DYN_RL_ENDPOINT `{endpoint_name}` collides with the serving endpoint");
}
if !endpoint_name
.chars()
.all(|character| character.is_ascii_alphanumeric() || matches!(character, '-' | '_'))
{
anyhow::bail!("DYN_RL_ENDPOINT must contain only letters, digits, '-' or '_'");
}
Ok(endpoint_name.to_string())
}

struct RlRouteHandler {
routes: EngineRouteRegistry,
system_url: String,
}

impl RlRouteHandler {
fn dispatch(&self, request: &Value) -> Value {
let Some(method) = request
.as_object()
.and_then(|request| request.get("method"))
.and_then(Value::as_str)
else {
return json!({"status": "error", "message": "rl_dispatch: missing 'method' (str)"});
};
if method != "routes" {
return json!({
"status": "error",
"method": method,
"message": "rl request-plane endpoint only supports method='routes'",
});
}

let mut routes = self.routes.routes().into_iter().collect::<Vec<_>>();
routes.sort();
routes.dedup();
json!({
"status": "ok",
"routes": routes,
"system_url": self.system_url,
})
Comment thread
connorcarpenter15 marked this conversation as resolved.
}
}

#[async_trait]
impl AsyncEngine<SingleIn<Value>, ManyOut<Annotated<Value>>, anyhow::Error> for RlRouteHandler {
async fn generate(&self, input: SingleIn<Value>) -> anyhow::Result<ManyOut<Annotated<Value>>> {
let (request, context) = input.into_parts();
let response = self.dispatch(&request);
Ok(ResponseStream::new(
Box::pin(stream::once(async move { Annotated::from_data(response) })),
context.context(),
))
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn rl_dispatch_only_describes_the_worker_engine_surface() {
let routes = EngineRouteRegistry::new();
routes.register(
"control/pause_generation",
Arc::new(|_| Box::pin(async { Ok(json!({"status": "ok"})) })),
);
let handler = RlRouteHandler {
routes,
system_url: "http://worker:8080".to_string(),
};

assert_eq!(
handler.dispatch(&json!({"method": "routes"})),
json!({
"status": "ok",
"routes": ["control/pause_generation"],
"system_url": "http://worker:8080",
})
);
assert_eq!(
handler.dispatch(&json!({"method": "control/pause_generation"}))["status"],
"error"
);
}
}
38 changes: 38 additions & 0 deletions lib/backend-common/src/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,8 @@ pub struct WorkerConfig {
/// roles -- setting it on `Decode` or `Encode` is rejected at
/// `Worker::run` validation time with `BackendError::InvalidArgument`.
pub route_to_encoder: bool,
/// Publish the worker's engine routes through an auxiliary RL discovery endpoint.
pub enable_rl: bool,
/// Optional frontend media decoding and fetch policy advertised on the
/// model deployment card.
pub media_decoder: Option<MediaDecoder>,
Expand Down Expand Up @@ -235,6 +237,7 @@ impl Default for WorkerConfig {
structural_tag_schema: StructuralTagSchemaMode::Auto,
runtime: RuntimeConfig::default(),
route_to_encoder: false,
enable_rl: false,
media_decoder: None,
media_fetcher: None,
default_thinking_mode: None,
Expand Down Expand Up @@ -922,6 +925,16 @@ impl Worker {
) -> Result<(), DynamoError> {
let model_type = resolve_model_type(&self.config)?;
let (worker_type, needs) = resolve_worker_type_and_needs(&self.config);
let rl_config = if self.config.enable_rl {
Some(crate::rl::prepare_endpoint(&endpoint).map_err(|error| {
err(
ErrorType::Backend(BackendError::InvalidArgument),
format!("RL endpoint configuration: {error}"),
)
})?)
} else {
None
};
let mut local_model =
build_local_model(&self.config, engine_config, self.engine.is_raw()).await?;
tracing::debug!("local model built");
Expand Down Expand Up @@ -1096,6 +1109,25 @@ impl Worker {
// the exact primary discovery instance is callable.
self.activate_engine_routes().await;

let rl_endpoint = if let Some(rl_config) = rl_config {
match crate::rl::serve_endpoint(&endpoint, rl_config).await {
Ok(endpoint) => Some(endpoint),
Err(error) => {
self.begin_engine_route_shutdown().await;
if let Err(shutdown_error) = primary_endpoint.shutdown().await {
tracing::warn!(%shutdown_error, "primary endpoint shutdown failed");
}
self.orchestrator_steps(&endpoint).await;
return Err(err(
ErrorType::Backend(BackendError::Unknown),
format!("RL endpoint setup: {error}"),
));
}
}
} else {
None
};

let serve_fut = primary_endpoint.wait();
tokio::pin!(serve_fut);

Expand Down Expand Up @@ -1132,6 +1164,12 @@ impl Worker {
// routes. No resume callback can re-register after the final unregister.
self.begin_engine_route_shutdown().await;

if let Some(rl_endpoint) = rl_endpoint
&& let Err(error) = rl_endpoint.shutdown().await
{
tracing::warn!(%error, "RL discovery endpoint shutdown failed");
}

self.orchestrator_steps(&endpoint).await;
serve_result
}
Expand Down
3 changes: 3 additions & 0 deletions lib/bindings/python/rust/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -466,6 +466,9 @@ impl WorkerConfig {
structural_tag_schema: st_schema,
runtime: runtime.map(|r| r.inner).unwrap_or_default(),
route_to_encoder,
// Python vLLM owns and serves its existing `.rl` endpoint.
// The shared Rust endpoint is opt-in for Rust sidecars only.
enable_rl: false,
media_decoder: media_decoder.map(|decoder| decoder.inner),
media_fetcher: media_fetcher.map(|fetcher| fetcher.inner),
},
Expand Down
1 change: 1 addition & 0 deletions lib/sidecar/sglang/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ impl SglangSidecarEngine {
.or_else(|| discovery_string(&discovery.server_info, "tool_call_parser")),
exclude_tools_when_tool_choice_none: common.exclude_tools_when_tool_choice_none,
route_to_encoder: false,
enable_rl: common.enable_rl,
..Default::default()
};

Expand Down
1 change: 1 addition & 0 deletions lib/sidecar/trtllm/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ impl TrtllmSidecarEngine {
enable_kv_routing: false,
disaggregation_mode: DisaggregationMode::Aggregated,
route_to_encoder: false,
enable_rl: args.sidecar.common.enable_rl,
..Default::default()
};
Ok((engine, config))
Expand Down
1 change: 1 addition & 0 deletions lib/sidecar/vllm/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ impl VllmSidecarEngine {
enable_kv_routing: true,
disaggregation_mode: mode,
route_to_encoder: false,
enable_rl: args.sidecar.common.enable_rl,
..Default::default()
};
Ok((engine, config))
Expand Down
Loading