lib.rs 37.2 KB
Newer Older
1
2
3
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

4
use dynamo_llm::local_model::LocalModel;
5
6
use futures::StreamExt;
use once_cell::sync::OnceCell;
7
use pyo3::IntoPyObjectExt;
8
use pyo3::exceptions::PyStopAsyncIteration;
9
use pyo3::types::{PyDict, PyString};
10
use pyo3::{exceptions::PyException, prelude::*};
11
use rand::seq::IteratorRandom as _;
12
use rs::pipeline::network::Ingress;
13
use std::fs;
14
use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
15
use std::path::PathBuf;
16
use std::time::Duration;
17
18
use std::{fmt::Display, sync::Arc};
use tokio::sync::Mutex;
19
use tracing::Instrument;
20

Neelay Shah's avatar
Neelay Shah committed
21
use dynamo_runtime::{
Ryan Olson's avatar
Ryan Olson committed
22
    self as rs, logging,
23
    pipeline::{
24
        AsyncEngineContextProvider, EngineStream, ManyOut, SingleIn, context::Context as RsContext,
25
        network::egress::push_router::RouterMode as RsRouterMode,
26
    },
27
    protocols::annotated::Annotated as RsAnnotated,
28
    traits::DistributedRuntimeProvider,
29
30
};

Neelay Shah's avatar
Neelay Shah committed
31
use dynamo_llm::{self as llm_rs};
32
33
use dynamo_llm::{entrypoint::RouterConfig, kv_router::KvRouterConfig};

34
35
use crate::llm::local_model::ModelRuntimeConfig;

36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
#[pyclass(eq, eq_int)]
#[derive(Clone, Debug, PartialEq)]
pub enum RouterMode {
    RoundRobin,
    Random,
    KV,
}

impl From<RouterMode> for RsRouterMode {
    fn from(mode: RouterMode) -> Self {
        match mode {
            RouterMode::RoundRobin => Self::RoundRobin,
            RouterMode::Random => Self::Random,
            RouterMode::KV => Self::KV,
        }
    }
}
53

54
mod context;
55
mod engine;
56
mod http;
57
mod llm;
58
mod parsers;
59
mod planner;
60
mod prometheus_metrics;
61
62
63
64
65
66
67
68

type JsonServerStreamingIngress =
    Ingress<SingleIn<serde_json::Value>, ManyOut<RsAnnotated<serde_json::Value>>>;

static INIT: OnceCell<()> = OnceCell::new();

const DEFAULT_ANNOTATED_SETTING: Option<bool> = Some(true);

69
70
// Helper to get appropriate span for instrumentation - always emit spans
fn get_span_for_context(context: &context::Context, operation: &str) -> tracing::Span {
71
72
73
74
75
76
    logging::make_client_request_span(
        operation,
        context.inner().id(),
        context.trace_context(),
        None,
    )
77
78
79
80
81
82
83
84
}

// Helper to create span for direct method with instance_id
fn get_span_for_direct_context(
    context: &context::Context,
    operation: &str,
    instance_id: &str,
) -> tracing::Span {
85
86
87
88
89
90
    logging::make_client_request_span(
        operation,
        context.inner().id(),
        context.trace_context(),
        Some(instance_id),
    )
91
92
}

93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
// Helper to create request context with proper linking and cancellation handling
fn create_request_context(
    request: serde_json::Value,
    parent_ctx: &Option<context::Context>,
) -> RsContext<serde_json::Value> {
    match parent_ctx {
        // If there is a parent context, link the request as a child context of it
        Some(parent_ctx) => {
            let child_ctx = RsContext::with_id(request, parent_ctx.inner().id().to_string());
            parent_ctx.inner().link_child(child_ctx.context());
            if parent_ctx.inner().is_stopped() || parent_ctx.inner().is_killed() {
                // Let the server handle the cancellation for now since not all backends are
                // properly handling request exceptions
                // TODO: (DIS-830) Return an error if context is cancelled
                child_ctx.context().stop_generating();
            }
            child_ctx
        }
        // Otherwise if there is no parent context, use the request as-is
        _ => request.into(),
    }
}

116
117
118
119
120
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
Yan Ru Pei's avatar
Yan Ru Pei committed
121
    m.add_function(wrap_pyfunction!(llm::kv::compute_block_hash_for_seq_py, m)?)?;
Ryan Olson's avatar
Ryan Olson committed
122
    m.add_function(wrap_pyfunction!(log_message, m)?)?;
123
    m.add_function(wrap_pyfunction!(register_llm, m)?)?;
124
    m.add_function(wrap_pyfunction!(fetch_llm, m)?)?;
125
126
    m.add_function(wrap_pyfunction!(llm::entrypoint::make_engine, m)?)?;
    m.add_function(wrap_pyfunction!(llm::entrypoint::run_input, m)?)?;
127

128
129
130
131
132
133
134
    m.add_class::<DistributedRuntime>()?;
    m.add_class::<CancellationToken>()?;
    m.add_class::<Namespace>()?;
    m.add_class::<Component>()?;
    m.add_class::<Endpoint>()?;
    m.add_class::<Client>()?;
    m.add_class::<AsyncResponseStream>()?;
135
    m.add_class::<llm::disagg_router::DisaggregatedRouter>()?;
136
137
138
    m.add_class::<llm::entrypoint::EntrypointArgs>()?;
    m.add_class::<llm::entrypoint::EngineConfig>()?;
    m.add_class::<llm::entrypoint::EngineType>()?;
139
140
    m.add_class::<llm::entrypoint::RouterConfig>()?;
    m.add_class::<llm::entrypoint::KvRouterConfig>()?;
141
    m.add_class::<llm::kv::WorkerMetricsPublisher>()?;
142
    m.add_class::<llm::model_card::ModelDeploymentCard>()?;
143
    m.add_class::<llm::local_model::ModelRuntimeConfig>()?;
144
145
    m.add_class::<llm::preprocessor::OAIChatPreprocessor>()?;
    m.add_class::<llm::backend::Backend>()?;
146
147
    m.add_class::<llm::kv::OverlapScores>()?;
    m.add_class::<llm::kv::KvIndexer>()?;
148
    m.add_class::<llm::kv::ApproxKvIndexer>()?;
149
    m.add_class::<llm::kv::KvEventPublisher>()?;
Yan Ru Pei's avatar
Yan Ru Pei committed
150
151
    m.add_class::<llm::kv::RadixTree>()?;
    m.add_class::<llm::kv::ZmqKvEventListener>()?;
152
153
    m.add_class::<llm::kv::ZmqKvEventPublisher>()?;
    m.add_class::<llm::kv::ZmqKvEventPublisherConfig>()?;
154
    m.add_class::<llm::kv::KvRecorder>()?;
155
156
    m.add_class::<http::HttpService>()?;
    m.add_class::<http::HttpAsyncEngine>()?;
157
    m.add_class::<context::Context>()?;
158
    m.add_class::<ModelType>()?;
159
    m.add_class::<ModelInput>()?;
160
161
162
163
    m.add_class::<llm::kv::ForwardPassMetrics>()?;
    m.add_class::<llm::kv::WorkerStats>()?;
    m.add_class::<llm::kv::KvStats>()?;
    m.add_class::<llm::kv::SpecDecodeStats>()?;
164
165
    m.add_class::<llm::kv::KvPushRouter>()?;
    m.add_class::<llm::kv::KvPushRouterStream>()?;
166
    m.add_class::<RouterMode>()?;
167
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
168
169
170
    m.add_class::<planner::VirtualConnectorCoordinator>()?;
    m.add_class::<planner::VirtualConnectorClient>()?;
    m.add_class::<planner::PlannerDecision>()?;
171
172

    engine::add_to_module(m)?;
173
    parsers::add_to_module(m)?;
174

175
    m.add_class::<prometheus_metrics::RuntimeMetrics>()?;
176
177
178
179
    let prometheus_metrics = PyModule::new(m.py(), "prometheus_metrics")?;
    prometheus_metrics::add_to_module(&prometheus_metrics)?;
    m.add_submodule(&prometheus_metrics)?;

180
181
182
    #[cfg(feature = "block-manager")]
    llm::block_manager::add_to_module(m)?;

183
184
185
186
187
188
189
190
191
192
    Ok(())
}

pub fn to_pyerr<E>(err: E) -> PyErr
where
    E: Display,
{
    PyException::new_err(format!("{}", err))
}

Ryan Olson's avatar
Ryan Olson committed
193
194
195
196
197
198
199
/// Log a message from Python with file and line info
#[pyfunction]
#[pyo3(text_signature = "(level, message, module, file, line)")]
fn log_message(level: &str, message: &str, module: &str, file: &str, line: u32) {
    logging::log_message(level, message, module, file, line);
}

200
201
/// Create an engine and attach it to an endpoint to make it visible to the frontend.
/// This is the main way you create a Dynamo worker / backend.
202
#[pyfunction]
203
#[pyo3(signature = (model_input, model_type, endpoint, model_path, model_name=None, context_length=None, kv_cache_block_size=None, router_mode=None, migration_limit=0, runtime_config=None, user_data=None, custom_template_path=None))]
204
#[allow(clippy::too_many_arguments)]
205
206
fn register_llm<'p>(
    py: Python<'p>,
207
    model_input: ModelInput,
208
    model_type: ModelType,
209
210
211
    endpoint: Endpoint,
    model_path: &str,
    model_name: Option<&str>,
212
213
    context_length: Option<u32>,
    kv_cache_block_size: Option<u32>,
214
    router_mode: Option<RouterMode>,
215
    migration_limit: u32,
216
    runtime_config: Option<ModelRuntimeConfig>,
217
    user_data: Option<&Bound<'p, PyDict>>,
218
    custom_template_path: Option<&str>,
219
) -> PyResult<Bound<'p, PyAny>> {
220
221
222
223
224
225
226
227
228
229
230
231
232
233
    // Validate Prefill model type requirements
    if model_type.inner == llm_rs::model_type::ModelType::Prefill {
        if !matches!(model_input, ModelInput::Tokens) {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "ModelType::Prefill requires model_input to be ModelInput::Tokens",
            ));
        }
        if migration_limit != 0 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "ModelType::Prefill requires migration_limit to be 0",
            ));
        }
    }

234
235
236
    let model_input = match model_input {
        ModelInput::Text => llm_rs::model_type::ModelInput::Text,
        ModelInput::Tokens => llm_rs::model_type::ModelInput::Tokens,
237
        ModelInput::Tensor => llm_rs::model_type::ModelInput::Tensor,
238
239
    };

240
241
    let model_type_obj = model_type.inner;

242
    let inner_path = model_path.to_string();
243
    let mut model_name = model_name.map(|n| n.to_string());
244
245
246
    let router_mode = router_mode.unwrap_or(RouterMode::RoundRobin);
    let router_config = RouterConfig::new(router_mode.into(), KvRouterConfig::default());

247
248
249
250
251
252
253
254
255
256
257
258
259
    // Early validation of custom template path
    let custom_template_path_owned = custom_template_path
        .map(|s| {
            let path = PathBuf::from(s);
            if !path.exists() {
                return Err(PyErr::new::<pyo3::exceptions::PyFileNotFoundError, _>(
                    format!("Custom template file does not exist: {}", path.display()),
                ));
            }
            Ok(path)
        })
        .transpose()?;

260
261
262
263
264
265
266
    let user_data_json = user_data
        .map(|dict| pythonize::depythonize(dict))
        .transpose()
        .map_err(|err| {
            PyErr::new::<PyException, _>(format!("Failed to convert user_data: {}", err))
        })?;

267
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
268
269
270
271
272
273
274
275
276
277
278
279
280
        let model_path = if fs::exists(&inner_path)? {
            PathBuf::from(inner_path)
        } else {
            // Preserve the model name
            if model_name.is_none() {
                model_name = Some(inner_path.clone());
            }
            // Likely it's a Hugging Face repo, download it
            LocalModel::fetch(&inner_path, false)
                .await
                .map_err(to_pyerr)?
        };

281
282
        let mut builder = dynamo_llm::local_model::LocalModelBuilder::default();
        builder
283
            .model_path(model_path)
284
285
            .model_name(model_name)
            .context_length(context_length)
286
            .kv_cache_block_size(kv_cache_block_size)
287
            .router_config(Some(router_config))
288
            .migration_limit(Some(migration_limit))
289
            .runtime_config(runtime_config.unwrap_or_default().inner)
290
291
            .user_data(user_data_json)
            .custom_template_path(custom_template_path_owned);
292
        // Load the ModelDeploymentCard
293
        let mut local_model = builder.build().await.map_err(to_pyerr)?;
294
295
        // Advertise ourself on etcd so ingress can find us
        local_model
296
            .attach(&endpoint.inner, model_type_obj, model_input)
297
298
299
300
301
302
303
            .await
            .map_err(to_pyerr)?;

        Ok(())
    })
}

304
305
306
307
308
309
310
311
312
313
314
/// Download a model from Hugging Face, returning it's local path
/// Example: `model_path = await fetch_llm("Qwen/Qwen3-0.6B")`
#[pyfunction]
#[pyo3(signature = (remote_name))]
fn fetch_llm<'p>(py: Python<'p>, remote_name: &str) -> PyResult<Bound<'p, PyAny>> {
    let repo = remote_name.to_string();
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
        LocalModel::fetch(&repo, false).await.map_err(to_pyerr)
    })
}

315
316
#[pyclass]
#[derive(Clone)]
Ryan Olson's avatar
Ryan Olson committed
317
pub struct DistributedRuntime {
318
319
320
321
    inner: rs::DistributedRuntime,
    event_loop: PyObject,
}

Ryan Olson's avatar
Ryan Olson committed
322
323
impl DistributedRuntime {
    #[allow(dead_code)]
324
    pub(crate) fn inner(&self) -> &rs::DistributedRuntime {
Ryan Olson's avatar
Ryan Olson committed
325
326
327
328
        &self.inner
    }
}

329
#[pyclass]
330
#[derive(Clone)]
331
332
333
334
335
struct CancellationToken {
    inner: rs::CancellationToken,
}

#[pyclass]
336
#[derive(Clone)]
337
338
339
340
341
342
struct Namespace {
    inner: rs::component::Namespace,
    event_loop: PyObject,
}

#[pyclass]
343
#[derive(Clone)]
344
345
346
347
348
349
struct Component {
    inner: rs::component::Component,
    event_loop: PyObject,
}

#[pyclass]
350
#[derive(Clone)]
351
352
353
354
355
356
struct Endpoint {
    inner: rs::component::Endpoint,
    event_loop: PyObject,
}

#[pyclass]
357
#[derive(Clone)]
358
struct Client {
359
    router: rs::pipeline::PushRouter<serde_json::Value, RsAnnotated<serde_json::Value>>,
360
361
}

362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
#[pyclass]
#[derive(Clone, PartialEq)]
struct ModelType {
    inner: llm_rs::model_type::ModelType,
}

#[pymethods]
#[allow(non_upper_case_globals)]
impl ModelType {
    #[classattr]
    const Chat: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Chat,
    };
    #[classattr]
    const Completions: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Completions,
    };
    #[classattr]
    const Embedding: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Embedding,
    };
383
384
385
386
    #[classattr]
    const TensorBased: Self = ModelType {
        inner: llm_rs::model_type::ModelType::TensorBased,
    };
387
388
389
390
    #[classattr]
    const Prefill: Self = ModelType {
        inner: llm_rs::model_type::ModelType::Prefill,
    };
391
392
393
394
395
396
397
398
399
400
401
402

    fn __or__(&self, other: &Self) -> Self {
        ModelType {
            inner: self.inner | other.inner,
        }
    }

    fn __str__(&self) -> String {
        self.inner.to_string()
    }
}

403
404
#[pyclass(eq, eq_int)]
#[derive(Clone, PartialEq)]
405
406
407
enum ModelInput {
    Text = 1,
    Tokens = 2,
408
    Tensor = 3,
409
410
}

411
412
413
#[pymethods]
impl DistributedRuntime {
    #[new]
414
    fn new(event_loop: PyObject, is_static: bool) -> PyResult<Self> {
415
        let worker = rs::Worker::from_settings().map_err(to_pyerr)?;
416
        INIT.get_or_try_init(|| {
417
            let primary = worker.tokio_runtime()?;
418
419
420
421
422
423
            pyo3_async_runtimes::tokio::init_with_runtime(primary)
                .map_err(|e| rs::error!("failed to initialize pyo3 static runtime: {:?}", e))?;
            rs::OK(())
        })
        .map_err(to_pyerr)?;

424
        let runtime = worker.runtime().clone();
425

426
427
428
429
430
431
        // Initialize logging in context where tokio runtime is available
        // otel exporter requires it
        runtime.secondary().block_on(async {
            rs::logging::init();
        });

432
433
434
435
436
437
438
439
440
441
442
        let inner =
            if is_static {
                runtime.secondary().block_on(
                    rs::DistributedRuntime::from_settings_without_discovery(runtime),
                )
            } else {
                runtime
                    .secondary()
                    .block_on(rs::DistributedRuntime::from_settings(runtime))
            };
        let inner = inner.map_err(to_pyerr)?;
443
444
445
446

        Ok(DistributedRuntime { inner, event_loop })
    }

Ryan Olson's avatar
Ryan Olson committed
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
    #[staticmethod]
    fn detached(py: Python) -> PyResult<Self> {
        let rt = rs::Worker::runtime_from_existing().map_err(to_pyerr)?;
        let handle = rt.primary();

        let inner = handle
            .block_on(rs::DistributedRuntime::from_settings(rt))
            .map_err(to_pyerr)?;

        Ok(DistributedRuntime {
            inner,
            event_loop: py.None(),
        })
    }

462
463
464
465
466
467
468
    fn namespace(&self, name: String) -> PyResult<Namespace> {
        Ok(Namespace {
            inner: self.inner.namespace(name).map_err(to_pyerr)?,
            event_loop: self.event_loop.clone(),
        })
    }

469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
    /// Allocate a contiguous block of ports from the specified range and atomically reserve them.
    /// Returns a list of all allocated ports in order.
    #[pyo3(signature = (namespace, port_min, port_max, block_size, context=None))]
    fn allocate_port_block<'p>(
        &self,
        py: Python<'p>,
        namespace: &str,
        port_min: u16,
        port_max: u16,
        block_size: u16,
        context: Option<String>, // Optional info to store alongside the reservation
    ) -> PyResult<Bound<'p, PyAny>> {
        const MAX_ALLOCATE_ATTEMPTS: usize = 100;
        if block_size == 0 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "Block size must be at least 1",
            ));
        }

        let Some(etcd_client) = self.inner.etcd_client() else {
            return Err(PyErr::new::<PyException, _>(
                "Static workers should not need to reserve ports",
            ));
        };

        let min = port_min;
        let max = port_max;

        // Compute maximum valid starting port (inclusive)
        let max_start_port = max.saturating_sub(block_size.saturating_sub(1));
        if max_start_port < min {
            return Err(PyErr::new::<PyException, _>(format!(
                "Port range {min}-{max} is too small for block size {block_size}",
            )));
        }

        // Randomize candidate starting ports to reduce contention/races
        let candidate_count =
            (max_start_port - port_min + 1).min(MAX_ALLOCATE_ATTEMPTS as u16) as usize;
        let mut rng = rand::rng();
        let candidate_ports: Vec<u16> =
            (port_min..=max_start_port).choose_multiple(&mut rng, candidate_count);

        let local_ip = match local_ip() {
            Ok(ip) => ip,
            Err(err) => {
                return Err(PyErr::new::<PyException, _>(format!(
                    "Failed fetching local IP address: {err}"
                )));
            }
        };

        let context_bytes = context.map(|s| s.as_bytes().to_vec()).unwrap_or_default();
        let namespace = namespace.to_owned();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            for (attempt_idx, start_port) in candidate_ports.into_iter().enumerate() {
                let end_port_exclusive = start_port + block_size;
                let ports_to_reserve: Vec<u16> = (start_port..end_port_exclusive).collect();

                // Hold/bind all ports in the block
                let mut sockets = Vec::with_capacity(ports_to_reserve.len());
                let mut bind_failed = false;

                for &port in &ports_to_reserve {
                    match bind_tcp_port(port) {
                        Ok(sock) => sockets.push(sock),
                        Err(e) => {
                            tracing::error!(
538
539
540
                                "Failed to bind to port block starting at {start_port} (attempt {}): {e}",
                                attempt_idx + 1,
                            );
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
                            bind_failed = true;
                            break;
                        }
                    }
                }

                if bind_failed {
                    // Let previously bound sockets drop here
                    if attempt_idx < candidate_count - 1 {
                        tokio::time::sleep(Duration::from_millis(10)).await;
                    }
                    continue;
                }

                // With sockets held, reserve in ETCD
                let mut reserved_keys = Vec::with_capacity(ports_to_reserve.len());
                let mut reservation_failed = false;
                for port in &ports_to_reserve {
                    let key = make_port_key(&namespace, local_ip, *port).map_err(to_pyerr)?;

                    if let Err(e) = etcd_client
                        .kv_create(&key, context_bytes.clone(), None)
                        .await
                    {
                        tracing::error!(
                            "Failed to reserve port block starting at {start_port} (attempt {}): {e}",
                            attempt_idx + 1,
                        );
                        reservation_failed = true;
                        break;
                    }
                    reserved_keys.push(key);
                }

                if reservation_failed {
                    // Cleanup partial reservations
                    for key in reserved_keys {
                        if let Err(e) = etcd_client.kv_delete(key.as_str(), None).await {
                            tracing::warn!("Failed to cleanup reserved port {key}: {e}");
                        }
                    }

                    // Sockets automatically released via RAII
                    if attempt_idx < candidate_count - 1 {
                        tokio::time::sleep(Duration::from_millis(10)).await;
                    }
                    continue;
                }

                // Success - sockets will be released automatically
                tracing::debug!("Reserved port block {ports_to_reserve:?}");
                return Ok(ports_to_reserve);
            }

            Err(PyErr::new::<PyException, _>(format!(
596
597
                "Failed to allocate and reserve a port block of size {block_size} from range {min}-{max} after {candidate_count} attempts"
            )))
598
599
600
        })
    }

601
602
603
604
605
606
607
    fn shutdown(&self) {
        self.inner.runtime().shutdown();
    }

    fn event_loop(&self) -> PyObject {
        self.event_loop.clone()
    }
608
609
610
611
612

    fn child_token(&self) -> CancellationToken {
        let inner = self.inner.runtime().child_token();
        CancellationToken { inner }
    }
613
614
}

615
616
617
618
619
620
621
622
623
624
625
626
627
628
// Bind a TCP port and return a socket held until dropped.
fn bind_tcp_port(port: u16) -> std::io::Result<socket2::Socket> {
    let sock = socket2::Socket::new(
        socket2::Domain::IPV4,
        socket2::Type::STREAM,
        Some(socket2::Protocol::TCP),
    )?;
    sock.set_reuse_address(true)?;
    let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port));
    sock.bind(&addr.into())?;
    Ok(sock)
}

fn make_port_key(namespace: &str, node_ip: IpAddr, port: u16) -> anyhow::Result<String> {
629
    Ok(format!("v1/{namespace}/ports/{node_ip}/{port}"))
630
631
632
633
634
635
636
637
638
639
640
641
}

fn local_ip() -> Result<IpAddr, local_ip_address::Error> {
    local_ip_address::local_ip().or_else(|err| match err {
        local_ip_address::Error::LocalIpAddressNotFound => {
            // Fall back to IPv6 if no IPv4 addresses are found
            local_ip_address::local_ipv6()
        }
        _ => Err(err),
    })
}

642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
#[pymethods]
impl CancellationToken {
    fn cancel(&self) {
        self.inner.cancel();
    }

    fn cancelled<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let token = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            token.cancelled().await;
            Ok(())
        })
    }
}

#[pymethods]
impl Component {
    fn endpoint(&self, name: String) -> PyResult<Endpoint> {
        let inner = self.inner.endpoint(name);
        Ok(Endpoint {
            inner,
            event_loop: self.event_loop.clone(),
        })
    }

667
    /// NATS specific stats/metrics call
668
    fn create_service<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
669
        let mut inner = self.inner.clone();
670
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
671
            inner.add_stats_service().await.map_err(to_pyerr)?;
672
673
674
            Ok(())
        })
    }
675
676
677

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
678
679
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_component(self.inner.clone())
680
    }
681
682
683
684
}

#[pymethods]
impl Endpoint {
685
    #[pyo3(signature = (generator, graceful_shutdown = true, metrics_labels = None, health_check_payload = None))]
686
687
688
689
    fn serve_endpoint<'p>(
        &self,
        py: Python<'p>,
        generator: PyObject,
690
        graceful_shutdown: Option<bool>,
691
        metrics_labels: Option<Vec<(String, String)>>,
692
        health_check_payload: Option<&Bound<'p, PyDict>>,
693
694
695
696
697
698
    ) -> PyResult<Bound<'p, PyAny>> {
        let engine = Arc::new(engine::PythonAsyncEngine::new(
            generator,
            self.event_loop.clone(),
        )?);
        let ingress = JsonServerStreamingIngress::for_engine(engine).map_err(to_pyerr)?;
699
700
701
702
703
704
705
706
707
708
709
710
711

        // Convert Python dict to serde_json::Value if provided and validate it's an object
        let health_payload_json = health_check_payload
            .map(|dict| pythonize::depythonize::<serde_json::Value>(dict))
            .transpose()
            .map_err(|err| {
                pyo3::exceptions::PyTypeError::new_err(format!(
                    "Failed to convert health_check_payload: {}",
                    err
                ))
            })?;

        // Require an object/dict
712
713
714
715
716
717
        if let Some(ref payload) = health_payload_json
            && !payload.is_object()
        {
            return Err(pyo3::exceptions::PyTypeError::new_err(
                "health_check_payload must be a JSON object (dict)",
            ));
718
719
720
        }

        let mut builder = self
721
722
723
724
            .inner
            .endpoint_builder()
            .metrics_labels(metrics_labels)
            .handler(ingress);
725
726
727
728
729

        if let Some(payload) = health_payload_json {
            builder = builder.health_check_payload(payload);
        }

730
        let graceful_shutdown = graceful_shutdown.unwrap_or(true);
731
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
732
733
734
735
736
            builder
                .graceful_shutdown(graceful_shutdown)
                .start()
                .await
                .map_err(to_pyerr)?;
737
738
739
740
            Ok(())
        })
    }

741
742
743
    fn client<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
744
            let client = inner.client().await.map_err(to_pyerr)?;
745
746
747
748
749
750
            let push_router = rs::pipeline::PushRouter::<
                serde_json::Value,
                RsAnnotated<serde_json::Value>,
            >::from_client(client, Default::default())
            .await
            .map_err(to_pyerr)?;
751
752
753
            Ok(Client {
                router: push_router,
            })
754
755
        })
    }
756
757

    fn lease_id(&self) -> i64 {
758
759
760
761
762
        self.inner
            .drt()
            .primary_lease()
            .map(|l| l.id())
            .unwrap_or(0)
763
    }
764
765
766

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
767
768
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_endpoint(self.inner.clone())
769
    }
770
771
772
773
774
775
776
777
778
779
780
}

#[pymethods]
impl Namespace {
    fn component(&self, name: String) -> PyResult<Component> {
        let inner = self.inner.component(name).map_err(to_pyerr)?;
        Ok(Component {
            inner,
            event_loop: self.event_loop.clone(),
        })
    }
781
782
783

    /// Get a RuntimeMetrics helper for creating Prometheus metrics
    #[getter]
784
785
    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
        prometheus_metrics::RuntimeMetrics::from_namespace(self.inner.clone())
786
    }
787
788
789
790
}

#[pymethods]
impl Client {
791
792
793
794
    /// Get list of current instances.
    /// Replaces endpoint_ids.
    fn instance_ids(&self) -> Vec<i64> {
        self.router.client.instance_ids()
795
796
    }

797
798
799
    /// Wait for an instance to be available for work.
    /// Replaces wait_for_endpoints.
    fn wait_for_instances<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
800
        let inner = self.router.client.clone();
801
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
802
            inner
803
                .wait_for_instances()
804
805
806
                .await
                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<i64>>())
                .map_err(to_pyerr)
807
808
809
810
        })
    }

    /// Issue a request to the endpoint using the default routing strategy.
811
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
812
813
814
815
816
    fn generate<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
817
        context: Option<context::Context>,
818
    ) -> PyResult<Bound<'p, PyAny>> {
819
        if self.router.client.is_static() {
820
            self.r#static(py, request, annotated, context)
821
        } else {
822
            self.random(py, request, annotated, context)
823
        }
824
825
826
    }

    /// Send a request to the next endpoint in a round-robin fashion.
827
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
828
829
830
831
832
    fn round_robin<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
833
        context: Option<context::Context>,
834
835
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
836
        let request_ctx = create_request_context(request, &context);
837
838
839
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
840
        let client = self.router.clone();
841
842

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
843
844
            let stream = match context {
                Some(context) => {
845
846
                    // Always instrument with appropriate span (none if no trace context)
                    let span = get_span_for_context(&context, "round_robin");
847
848
849
850
851
                    client
                        .round_robin(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
852
                }
853
                _ => client.round_robin(request_ctx).await.map_err(to_pyerr)?,
854
            };
855
856
857
858
859
860
861
862
863
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Send a request to a random endpoint.
864
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
865
866
867
868
869
    fn random<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
870
        context: Option<context::Context>,
871
872
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
873
        let request_ctx = create_request_context(request, &context);
874
875
876
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
877
        let client = self.router.clone();
878
879

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
880
881
            let stream = match context {
                Some(context) => {
882
                    // Always instrument with appropriate span (none if no trace context)
883
                    let span = get_span_for_context(&context, "random");
884
885
886
887
888
                    client
                        .random(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
889
                }
890
                _ => client.random(request_ctx).await.map_err(to_pyerr)?,
891
            };
892
893
894
895
896
897
898
899
900
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Directly send a request to a specific endpoint.
901
    #[pyo3(signature = (request, instance_id, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
902
903
904
905
    fn direct<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
906
        instance_id: i64,
907
        annotated: Option<bool>,
908
        context: Option<context::Context>,
909
910
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
911
        let request_ctx = create_request_context(request, &context);
912
913
914
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
915
        let client = self.router.clone();
916
917

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
918
919
            let stream = match context {
                Some(context) => {
920
                    // Always instrument with appropriate span (none if no trace context)
921
922
                    let span =
                        get_span_for_direct_context(&context, "direct", &instance_id.to_string());
923
924
925
926
927
                    client
                        .direct(request_ctx, instance_id)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
928
929
                }
                _ => client
930
                    .direct(request_ctx, instance_id)
931
932
933
                    .await
                    .map_err(to_pyerr)?,
            };
934
935
936
937
938
939
940
941
942

            tokio::spawn(process_stream(stream, tx));

            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }
943
944

    /// Directly send a request to a pre-defined static worker
945
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
946
947
948
949
950
    fn r#static<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
951
        context: Option<context::Context>,
952
953
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
954
        let request_ctx = create_request_context(request, &context);
955
956
957
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
958
        let client = self.router.clone();
959
960

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
961
962
            let stream = match context {
                Some(context) => {
963
                    // Always instrument with appropriate span (none if no trace context)
964
                    let span = get_span_for_context(&context, "static");
965
966
967
968
969
                    client
                        .r#static(request_ctx)
                        .instrument(span)
                        .await
                        .map_err(to_pyerr)?
970
                }
971
                _ => client.r#static(request_ctx).await.map_err(to_pyerr)?,
972
            };
973
974
975
976
977
978
979
980
981

            tokio::spawn(process_stream(stream, tx));

            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }
982
983
984
}

async fn process_stream(
985
    stream: EngineStream<RsAnnotated<serde_json::Value>>,
986
987
988
989
990
    tx: tokio::sync::mpsc::Sender<RsAnnotated<PyObject>>,
) {
    let mut stream = stream;
    while let Some(response) = stream.next().await {
        // Convert the response to a PyObject using Python's GIL
991
        let annotated: RsAnnotated<serde_json::Value> = response;
992
        let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
993
            Python::with_gil(|py| match pythonize::pythonize(py, &data) {
994
995
                Ok(pyobj) => Ok(pyobj.into()),
                Err(e) => Err(e.to_string()),
996
            })
997
998
999
1000
1001
1002
        });

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
1003
            tracing::error!("Failed to send response: {:?}", e);
1004
            break;
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
        }

        if is_error {
            break;
        }
    }
}

#[pyclass]
struct AsyncResponseStream {
    rx: Arc<Mutex<tokio::sync::mpsc::Receiver<RsAnnotated<PyObject>>>>,
    annotated: bool,
}

#[pymethods]
impl AsyncResponseStream {
    /// This method is required to implement the `AsyncIterator` protocol.
    #[pyo3(name = "__aiter__")]
    fn aiter(slf: PyRef<Self>, py: Python) -> PyResult<Py<PyAny>> {
        slf.into_py_any(py)
    }
    /// This method is required to implement the `AsyncIterator` protocol.
    #[pyo3(name = "__anext__")]
    fn next<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let rx = self.rx.clone();
        let annotated = self.annotated;

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            loop {
                let value = rx.lock().await.recv().await;
                match value {
                    Some(pyobj) => {
                        let pyobj = match pyobj.ok() {
                            Ok(pyobj) => pyobj,
                            Err(e) => {
                                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(e));
                            }
                        };

                        if annotated {
                            let object = Annotated { inner: pyobj };
                            #[allow(deprecated)]
                            let object = Python::with_gil(|py| object.into_py(py));
                            return Ok(object);
                        } else {
                            match pyobj.data {
                                Some(data) => return Ok(data),
                                None => continue,
                            }
                        }
                    }
                    None => return Err(PyStopAsyncIteration::new_err("Stream exhausted")),
                }
            }
        })
    }
}

#[pyclass]
struct Annotated {
    inner: RsAnnotated<PyObject>,
}

#[pymethods]
impl Annotated {
    #[new]
    fn new(data: PyObject) -> Self {
        Annotated {
            inner: RsAnnotated::from_data(data),
        }
    }

    fn is_error(&self) -> bool {
        self.inner.is_error()
    }

    fn data(&self) -> Option<PyObject> {
        self.inner.data.clone()
    }

    fn event(&self) -> Option<String> {
        self.inner.event.clone()
    }

    fn comments(&self) -> Option<Vec<String>> {
        self.inner.comment.clone()
    }

    fn id(&self) -> Option<String> {
        self.inner.id.clone()
    }

    #[pyo3(name = "__repr__")]
    fn _repr(&self, py: Python) -> String {
        let data = self.inner.data.clone().map(|obj| {
            obj.call_method0(py, "__repr__")
                .and_then(|repr_obj| repr_obj.extract::<Py<PyString>>(py))
                .map(|py_str| py_str.to_string_lossy(py).into_owned())
                .unwrap_or_else(|_| "<failed_repr>".to_string())
        });

        format!(
            "Annotated(data={}, event={}, comment={:?}, id={})",
            data.unwrap_or_else(|| "<no_data>".to_string()),
            self.inner.event.as_deref().unwrap_or("None"),
            self.inner.comment.as_deref().unwrap_or(&[]),
            self.inner.id.as_deref().unwrap_or("None")
        )
    }
}