endpoint.rs 17 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
// SPDX-License-Identifier: Apache-2.0
Ryan Olson's avatar
Ryan Olson committed
3

4
5
6
7
use std::sync::Arc;

use anyhow::Result;
use derive_builder::Builder;
Ryan Olson's avatar
Ryan Olson committed
8
use derive_getters::Dissolve;
9
use educe::Educe;
10
use tokio_util::sync::CancellationToken;
Ryan Olson's avatar
Ryan Olson committed
11

12
use crate::{
13
    component::{DeviceType, Endpoint, Instance, TransportType},
14
    distributed::RequestPlaneMode,
15
    pipeline::network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint},
16
    protocols::EndpointId,
17
    traits::DistributedRuntimeProvider,
18
    transports::nats,
19
};
20

21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
fn endpoint_device_type() -> Option<DeviceType> {
    // Common CUDA masks that explicitly disable GPU visibility.
    if std::env::var("CUDA_VISIBLE_DEVICES")
        .ok()
        .map(|v| {
            let l = v.trim().to_ascii_lowercase();
            l.is_empty() || l == "-1" || l == "none" || l == "void"
        })
        .unwrap_or(false)
    {
        return Some(DeviceType::Cpu);
    }

    // Container runtimes often use NVIDIA_VISIBLE_DEVICES to gate GPU visibility.
    if std::env::var("NVIDIA_VISIBLE_DEVICES")
        .ok()
        .map(|v| {
            let l = v.trim().to_ascii_lowercase();
            l == "none" || l == "void"
        })
        .unwrap_or(false)
    {
        return Some(DeviceType::Cpu);
    }

    // Default: no explicit CPU override means this endpoint is CUDA-capable.
    Some(DeviceType::Cuda)
}

Ryan Olson's avatar
Ryan Olson committed
50
51
52
53
54
55
56
57
58
59
#[derive(Educe, Builder, Dissolve)]
#[educe(Debug)]
#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
pub struct EndpointConfig {
    #[builder(private)]
    endpoint: Endpoint,

    /// Endpoint handler
    #[educe(Debug(ignore))]
    handler: Arc<dyn PushWorkHandler>,
60

61
62
63
64
    /// Additional labels for metrics
    #[builder(default, setter(into))]
    metrics_labels: Option<Vec<(String, String)>>,

65
66
67
    /// Whether to wait for inflight requests to complete during shutdown
    #[builder(default = "true")]
    graceful_shutdown: bool,
68
69
70
71
72
73
74

    /// Health check payload for this endpoint
    /// This payload will be sent to the endpoint during health checks
    /// to verify it's responding properly
    #[educe(Debug(ignore))]
    #[builder(default, setter(into, strip_option))]
    health_check_payload: Option<serde_json::Value>,
Ryan Olson's avatar
Ryan Olson committed
75
76
77
78
79
80
81
}

impl EndpointConfigBuilder {
    pub(crate) fn from_endpoint(endpoint: Endpoint) -> Self {
        Self::default().endpoint(endpoint)
    }

82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    /// Register an async engine in the local endpoint registry for direct in-process calls
    pub fn register_local_engine(
        self,
        engine: crate::local_endpoint_registry::LocalAsyncEngine,
    ) -> Result<Self> {
        if let Some(endpoint) = &self.endpoint {
            let registry = endpoint.drt().local_endpoint_registry();
            registry.register(endpoint.name.clone(), engine);
            tracing::debug!(
                "Registered engine for endpoint '{}' in local registry",
                endpoint.name
            );
        }
        Ok(self)
    }

Ryan Olson's avatar
Ryan Olson committed
98
    pub async fn start(self) -> Result<()> {
99
100
        let (endpoint, handler, metrics_labels, graceful_shutdown, health_check_payload) =
            self.build_internal()?.dissolve();
101
        let connection_id = endpoint.drt().connection_id();
102
        let endpoint_id = endpoint.id();
Ryan Olson's avatar
Ryan Olson committed
103

104
        tracing::debug!("Starting endpoint: {endpoint_id}");
Ryan Olson's avatar
Ryan Olson committed
105

106
107
108
        let metrics_labels: Option<Vec<(&str, &str)>> = metrics_labels
            .as_ref()
            .map(|v| v.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect());
109
        // Add metrics to the handler. The endpoint provides additional information to the handler.
110
        handler.add_metrics(&endpoint, metrics_labels.as_deref())?;
111

112
113
114
        // This creates a child token of the runtime's endpoint_shutdown_token. That token is
        // cancelled first as part of graceful shutdown. See Runtime::shutdown.
        let endpoint_shutdown_token = endpoint.drt().child_token();
115

116
        let system_health = endpoint.drt().system_health();
117
118
119
120
121
122
123
124
125
126
127
128

        // Register with graceful shutdown tracker if needed
        if graceful_shutdown {
            tracing::debug!(
                "Registering endpoint '{}' with graceful shutdown tracker",
                endpoint.name
            );
            let tracker = endpoint.drt().graceful_shutdown_tracker();
            tracker.register_endpoint();
        } else {
            tracing::debug!("Endpoint '{}' has graceful_shutdown=false", endpoint.name);
        }
Ryan Olson's avatar
Ryan Olson committed
129

130
        // Launch endpoint based on request plane mode
131
132
133
134
135
136
137
        let tracker_clone = if graceful_shutdown {
            Some(endpoint.drt().graceful_shutdown_tracker())
        } else {
            None
        };

        // Create clones for the async closure
138
139
140
        let namespace_name_for_task = endpoint_id.namespace.clone();
        let component_name_for_task = endpoint_id.component.clone();
        let endpoint_name_for_task = endpoint_id.name.clone();
141

142
        // Get the unified request plane server
143
144
        let server = endpoint.drt().request_plane_server().await?;

145
146
        // Register health check target in SystemHealth if provided
        if let Some(health_check_payload) = &health_check_payload {
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
            if system_health.lock().health_check_enabled()
                && endpoint
                    .drt()
                    .local_endpoint_registry()
                    .get(&endpoint.name)
                    .is_none()
            {
                anyhow::bail!(
                    "Endpoint '{}' has a health_check_payload and canary is enabled, \
                     but no local engine is registered. Call .register_local_engine() \
                     before .start() so the canary health check can function.",
                    endpoint.name
                );
            }

162
163
164
165
166
167
168
169
170
            // Build transport based on request plane mode
            let transport = build_transport_type(&endpoint, &endpoint_id, connection_id).await?;

            let instance = Instance {
                component: endpoint_id.component.clone(),
                endpoint: endpoint_id.name.clone(),
                namespace: endpoint_id.namespace.clone(),
                instance_id: connection_id,
                transport,
171
                device_type: endpoint_device_type(),
172
173
174
175
176
177
178
179
180
181
182
183
184
            };
            tracing::debug!(endpoint_name = %endpoint.name, "Registering endpoint health check target");
            let guard = system_health.lock();
            guard.register_health_check_target(
                &endpoint.name,
                instance,
                health_check_payload.clone(),
            );
            if let Some(notifier) = guard.get_endpoint_health_check_notifier(&endpoint.name) {
                handler.set_endpoint_health_check_notifier(notifier)?;
            }
        }

185
        tracing::debug!(
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
            endpoint = %endpoint_name_for_task,
            transport = server.transport_name(),
            "Registering endpoint with request plane server"
        );

        // Register endpoint with the server (unified interface)
        server
            .register_endpoint(
                endpoint_name_for_task.clone(),
                handler,
                connection_id,
                namespace_name_for_task.clone(),
                component_name_for_task.clone(),
                system_health.clone(),
            )
            .await?;

        // Create cleanup task that unregisters on cancellation
        let endpoint_name_for_cleanup = endpoint_name_for_task.clone();
        let server_for_cleanup = server.clone();
        let cancel_token_for_cleanup = endpoint_shutdown_token.clone();

        let task: tokio::task::JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
            cancel_token_for_cleanup.cancelled().await;

            tracing::debug!(
                endpoint = %endpoint_name_for_cleanup,
                "Unregistering endpoint from request plane server"
            );

            // Unregister from server
            if let Err(e) = server_for_cleanup
                .unregister_endpoint(&endpoint_name_for_cleanup)
                .await
            {
                tracing::warn!(
                    endpoint = %endpoint_name_for_cleanup,
                    error = %e,
                    "Failed to unregister endpoint"
                );
            }
227
228
229

            // Unregister from graceful shutdown tracker
            if let Some(tracker) = tracker_clone {
230
                tracing::debug!("Unregister endpoint from graceful shutdown tracker");
231
232
233
                tracker.unregister_endpoint();
            }

234
            anyhow::Ok(())
235
        });
Ryan Olson's avatar
Ryan Olson committed
236

237
238
239
240
241
        // Register this endpoint instance in the discovery plane
        // The discovery interface abstracts storage backend (etcd, k8s, etc) and provides
        // consistent registration/discovery across the system.
        let discovery = endpoint.drt().discovery();

242
        // Build transport for discovery service based on request plane mode
243
        let transport = build_transport_type(&endpoint, &endpoint_id, connection_id).await?;
244

245
        let discovery_spec = crate::discovery::DiscoverySpec::Endpoint {
246
247
248
            namespace: endpoint_id.namespace.clone(),
            component: endpoint_id.component.clone(),
            endpoint: endpoint_id.name.clone(),
249
            transport,
250
            device_type: endpoint_device_type(),
Ryan Olson's avatar
Ryan Olson committed
251
252
        };

253
        if let Err(e) = discovery.register(discovery_spec).await {
254
            tracing::error!(
255
                %endpoint_id,
256
                error = %e,
257
258
                "Unable to register service for discovery"
            );
259
            endpoint_shutdown_token.cancel();
260
            anyhow::bail!(
261
                "Unable to register service for discovery. Check discovery service status"
262
            );
Ryan Olson's avatar
Ryan Olson committed
263
        }
264

Ryan Olson's avatar
Ryan Olson committed
265
266
267
268
269
        task.await??;

        Ok(())
    }
}
270

271
/// Build transport type based on request plane mode
272
///
273
274
275
/// This function handles both health check and discovery transport building.
/// All transport modes use consistent addressing:
/// - HTTP: Uses full URL path including endpoint name (e.g., http://host:port/v1/rpc/endpoint_name)
276
/// - TCP: Includes instance_id and endpoint name for routing (e.g., host:port/instance_id_hex/endpoint_name)
277
/// - NATS: Uses subject-based addressing (unique per endpoint)
278
279
280
281
///
/// # Errors
/// Returns an error if TCP mode is used but the TCP server hasn't been started yet.
fn build_transport_type_inner(
282
    mode: RequestPlaneMode,
283
284
    endpoint_id: &EndpointId,
    connection_id: u64,
285
) -> Result<TransportType> {
286
287
288
    match mode {
        RequestPlaneMode::Http => {
            let http_host = crate::utils::get_http_rpc_host_from_env();
289
290
            // If a fixed port is explicitly configured, use it directly.
            // Otherwise, use the actual bound port (set by HTTP server after binding when port 0 is used).
291
292
293
            let http_port = std::env::var("DYN_HTTP_RPC_PORT")
                .ok()
                .and_then(|p| p.parse::<u16>().ok())
294
295
                .filter(|&p| p != 0)
                .unwrap_or(crate::pipeline::network::manager::get_actual_http_rpc_port()?);
296
297
298
299
            let rpc_root =
                std::env::var("DYN_HTTP_RPC_ROOT_PATH").unwrap_or_else(|_| "/v1/rpc".to_string());

            let http_endpoint = format!(
300
301
                "http://{http_host}:{http_port}{rpc_root}/{}",
                endpoint_id.name
302
303
            );

304
            Ok(TransportType::Http(http_endpoint))
305
306
307
        }
        RequestPlaneMode::Tcp => {
            let tcp_host = crate::utils::get_tcp_rpc_host_from_env();
308
309
            // If a fixed port is explicitly configured, use it directly (no init ordering dependency).
            // Otherwise, use the actual bound port (set by TCP server after binding when port 0 is used).
310
311
312
            let tcp_port = std::env::var("DYN_TCP_RPC_PORT")
                .ok()
                .and_then(|p| p.parse::<u16>().ok())
313
                .filter(|&p| p != 0)
314
                .unwrap_or(crate::pipeline::network::manager::get_actual_tcp_rpc_port()?);
315

316
317
318
319
320
321
322
323
            // Include instance_id and endpoint name for proper TCP routing.
            // Format: host:port/instance_id_hex/endpoint_name
            // This ensures each worker has a unique routing key when multiple workers
            // share the same TCP server (e.g., --num-workers > 1).
            let tcp_endpoint = format!(
                "{}:{}/{:x}/{}",
                tcp_host, tcp_port, connection_id, endpoint_id.name
            );
324

325
            Ok(TransportType::Tcp(tcp_endpoint))
326
        }
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
        RequestPlaneMode::Nats => Ok(TransportType::Nats(nats::instance_subject(
            endpoint_id,
            connection_id,
        ))),
    }
}

/// Build transport type, ensuring TCP server is initialized when needed.
///
/// In TCP mode with an OS-assigned port (`DYN_TCP_RPC_PORT` unset or invalid), the server must bind
/// before we can construct a correct transport address. This helper ensures that initialization
/// occurs, then delegates to the internal builder.
pub async fn build_transport_type(
    endpoint: &Endpoint,
    endpoint_id: &EndpointId,
    connection_id: u64,
) -> Result<TransportType> {
    let mode = endpoint.drt().request_plane();

346
347
348
349
    // For TCP and HTTP with OS-assigned ports, we must ensure the server is initialized
    // (bound to a port) before we can construct a correct transport address.
    let has_fixed_port = match mode {
        RequestPlaneMode::Tcp => std::env::var("DYN_TCP_RPC_PORT")
350
351
            .ok()
            .and_then(|p| p.parse::<u16>().ok())
352
353
354
355
356
357
358
359
360
361
362
363
364
            .filter(|&p| p != 0)
            .is_some(),
        RequestPlaneMode::Http => std::env::var("DYN_HTTP_RPC_PORT")
            .ok()
            .and_then(|p| p.parse::<u16>().ok())
            .filter(|&p| p != 0)
            .is_some(),
        RequestPlaneMode::Nats => true, // NATS doesn't need port init
    };

    if !has_fixed_port {
        // Ensure request plane server is initialized before building transport.
        let _ = endpoint.drt().request_plane_server().await?;
365
    }
366
367

    build_transport_type_inner(mode, endpoint_id, connection_id)
368
}
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389

impl Endpoint {
    /// Unregister this endpoint instance from discovery.
    ///
    /// This removes the endpoint from the instances bucket, preventing the router
    /// from sending requests to this worker. Use this when a worker is sleeping
    /// and should not receive any requests.
    pub async fn unregister_endpoint_instance(&self) -> anyhow::Result<()> {
        let drt = self.drt();
        let instance_id = drt.connection_id();
        let endpoint_id = self.id();

        // Get the transport type for the endpoint
        let transport = build_transport_type(self, &endpoint_id, instance_id).await?;

        let instance = crate::discovery::DiscoveryInstance::Endpoint(Instance {
            namespace: endpoint_id.namespace,
            component: endpoint_id.component,
            endpoint: endpoint_id.name,
            instance_id,
            transport,
390
            device_type: endpoint_device_type(),
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
        });

        let discovery = drt.discovery();
        if let Err(e) = discovery.unregister(instance).await {
            let endpoint_id = self.id();
            tracing::error!(
                %endpoint_id,
                error = %e,
                "Unable to unregister endpoint instance from discovery"
            );
            anyhow::bail!(
                "Unable to unregister endpoint instance from discovery. Check discovery service status"
            );
        }

        tracing::info!(
            instance_id = instance_id,
            "Successfully unregistered endpoint instance from discovery - worker removed from routing pool"
        );

        Ok(())
    }

    /// Re-register this endpoint instance to discovery.
    ///
    /// This adds the endpoint back to the instances bucket, allowing the router
    /// to send requests to this worker again. Use this when a worker wakes up
    /// and should start receiving requests.
    pub async fn register_endpoint_instance(&self) -> anyhow::Result<()> {
        let drt = self.drt();
        let instance_id = drt.connection_id();
        let endpoint_id = self.id();

        // Get the transport type for the endpoint
        let transport = build_transport_type(self, &endpoint_id, instance_id).await?;

        let spec = crate::discovery::DiscoverySpec::Endpoint {
            namespace: endpoint_id.namespace,
            component: endpoint_id.component,
            endpoint: endpoint_id.name,
            transport,
432
            device_type: endpoint_device_type(),
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
        };

        let discovery = drt.discovery();
        if let Err(e) = discovery.register(spec).await {
            let endpoint_id = self.id();
            tracing::error!(
                %endpoint_id,
                error = %e,
                "Unable to re-register endpoint instance to discovery"
            );
            anyhow::bail!(
                "Unable to re-register endpoint instance to discovery. Check discovery service status"
            );
        }

        tracing::info!(
            instance_id = instance_id,
            "Successfully re-registered endpoint instance to discovery - worker added back to routing pool"
        );

        Ok(())
    }
}