dynamo_runtime/local_endpoint_registry.rs
1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Local Endpoint Registry
5//!
6//! Provides a registry for locally registered endpoints that can be called in-process
7//! without going through the network stack.
8
9use crate::engine::AsyncEngine;
10use dashmap::DashMap;
11use std::sync::Arc;
12
13/// Type alias for a boxed async engine that can handle generic requests and responses
14pub type LocalAsyncEngine = Arc<
15 dyn AsyncEngine<
16 crate::pipeline::SingleIn<serde_json::Value>,
17 crate::pipeline::ManyOut<crate::protocols::annotated::Annotated<serde_json::Value>>,
18 anyhow::Error,
19 > + Send
20 + Sync,
21>;
22
23/// Registry for locally registered endpoints
24///
25/// This registry stores endpoints that are registered locally (in the same process)
26/// and allows them to be called directly without going through the network transport layer.
27#[derive(Clone, Default)]
28pub struct LocalEndpointRegistry {
29 /// Map of endpoint name to async engine
30 engines: Arc<DashMap<String, LocalAsyncEngine>>,
31}
32
33impl LocalEndpointRegistry {
34 /// Create a new local endpoint registry
35 pub fn new() -> Self {
36 Self {
37 engines: Arc::new(DashMap::new()),
38 }
39 }
40
41 /// Register a local endpoint
42 ///
43 /// # Arguments
44 ///
45 /// * `endpoint_name` - Name of the endpoint (e.g., "load_lora", "generate")
46 /// * `engine` - The async engine that handles requests for this endpoint
47 pub fn register(&self, endpoint_name: String, engine: LocalAsyncEngine) {
48 tracing::debug!("Registering local endpoint: {}", endpoint_name);
49 self.engines.insert(endpoint_name, engine);
50 }
51
52 /// Get a registered local endpoint
53 ///
54 /// The async engine if found, None otherwise
55 pub fn get(&self, endpoint_name: &str) -> Option<LocalAsyncEngine> {
56 self.engines.get(endpoint_name).map(|e| e.clone())
57 }
58}