Skip to main content

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}