lib.rs 40.1 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
5
use futures::StreamExt;
use once_cell::sync::OnceCell;
6
use pyo3::IntoPyObjectExt;
7
use pyo3::exceptions::PyStopAsyncIteration;
8
9
use pyo3::types::PyBytes;
use pyo3::types::{PyDict, PyList, PyString};
10
use pyo3::{exceptions::PyException, prelude::*};
11
use rand::seq::IteratorRandom as _;
12
use rs::pipeline::network::Ingress;
13
use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
14
use std::path::PathBuf;
15
use std::time::Duration;
16
17
18
use std::{fmt::Display, sync::Arc};
use tokio::sync::Mutex;

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

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

32
33
use crate::llm::local_model::ModelRuntimeConfig;

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#[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,
        }
    }
}
51

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

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);

/// 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<()> {
Ryan Olson's avatar
Ryan Olson committed
71
    logging::init();
Yan Ru Pei's avatar
Yan Ru Pei committed
72
    m.add_function(wrap_pyfunction!(llm::kv::compute_block_hash_for_seq_py, m)?)?;
Ryan Olson's avatar
Ryan Olson committed
73
    m.add_function(wrap_pyfunction!(log_message, m)?)?;
74
    m.add_function(wrap_pyfunction!(register_llm, m)?)?;
75
76
    m.add_function(wrap_pyfunction!(llm::entrypoint::make_engine, m)?)?;
    m.add_function(wrap_pyfunction!(llm::entrypoint::run_input, m)?)?;
77

78
79
80
81
82
83
    m.add_class::<DistributedRuntime>()?;
    m.add_class::<CancellationToken>()?;
    m.add_class::<Namespace>()?;
    m.add_class::<Component>()?;
    m.add_class::<Endpoint>()?;
    m.add_class::<Client>()?;
84
    m.add_class::<EtcdClient>()?;
85
    m.add_class::<AsyncResponseStream>()?;
86
    m.add_class::<llm::disagg_router::DisaggregatedRouter>()?;
87
88
89
    m.add_class::<llm::entrypoint::EntrypointArgs>()?;
    m.add_class::<llm::entrypoint::EngineConfig>()?;
    m.add_class::<llm::entrypoint::EngineType>()?;
90
91
    m.add_class::<llm::entrypoint::RouterConfig>()?;
    m.add_class::<llm::entrypoint::KvRouterConfig>()?;
92
    m.add_class::<llm::kv::WorkerMetricsPublisher>()?;
93
    m.add_class::<llm::model_card::ModelDeploymentCard>()?;
94
    m.add_class::<llm::local_model::ModelRuntimeConfig>()?;
95
96
    m.add_class::<llm::preprocessor::OAIChatPreprocessor>()?;
    m.add_class::<llm::backend::Backend>()?;
97
98
    m.add_class::<llm::kv::OverlapScores>()?;
    m.add_class::<llm::kv::KvIndexer>()?;
99
    m.add_class::<llm::kv::ApproxKvIndexer>()?;
100
101
102
    m.add_class::<llm::kv::EndpointKvMetrics>()?;
    m.add_class::<llm::kv::AggregatedMetrics>()?;
    m.add_class::<llm::kv::KvMetricsAggregator>()?;
103
    m.add_class::<llm::kv::KvEventPublisher>()?;
Yan Ru Pei's avatar
Yan Ru Pei committed
104
105
    m.add_class::<llm::kv::RadixTree>()?;
    m.add_class::<llm::kv::ZmqKvEventListener>()?;
106
107
    m.add_class::<llm::kv::ZmqKvEventPublisher>()?;
    m.add_class::<llm::kv::ZmqKvEventPublisherConfig>()?;
108
    m.add_class::<llm::kv::KvRecorder>()?;
109
110
111
    m.add_class::<http::HttpService>()?;
    m.add_class::<http::HttpError>()?;
    m.add_class::<http::HttpAsyncEngine>()?;
112
    m.add_class::<context::Context>()?;
113
    m.add_class::<EtcdKvCache>()?;
114
    m.add_class::<ModelType>()?;
115
    m.add_class::<ModelInput>()?;
116
117
118
119
    m.add_class::<llm::kv::ForwardPassMetrics>()?;
    m.add_class::<llm::kv::WorkerStats>()?;
    m.add_class::<llm::kv::KvStats>()?;
    m.add_class::<llm::kv::SpecDecodeStats>()?;
Yan Ru Pei's avatar
Yan Ru Pei committed
120
    m.add_class::<llm::kv::KvRouter>()?;
121
122
    m.add_class::<llm::kv::KvPushRouter>()?;
    m.add_class::<llm::kv::KvPushRouterStream>()?;
123
    m.add_class::<RouterMode>()?;
124
125
126
    m.add_class::<planner::VirtualConnectorCoordinator>()?;
    m.add_class::<planner::VirtualConnectorClient>()?;
    m.add_class::<planner::PlannerDecision>()?;
127
128

    engine::add_to_module(m)?;
129
    parsers::add_to_module(m)?;
130

131
132
133
    #[cfg(feature = "block-manager")]
    llm::block_manager::add_to_module(m)?;

134
135
136
137
138
139
140
141
142
143
    Ok(())
}

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

Ryan Olson's avatar
Ryan Olson committed
144
145
146
147
148
149
150
/// 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);
}

151
#[pyfunction]
152
#[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))]
153
#[allow(clippy::too_many_arguments)]
154
155
fn register_llm<'p>(
    py: Python<'p>,
156
    model_input: ModelInput,
157
    model_type: ModelType,
158
159
160
    endpoint: Endpoint,
    model_path: &str,
    model_name: Option<&str>,
161
162
    context_length: Option<u32>,
    kv_cache_block_size: Option<u32>,
163
    router_mode: Option<RouterMode>,
164
    migration_limit: u32,
165
    runtime_config: Option<ModelRuntimeConfig>,
166
    user_data: Option<&Bound<'p, PyDict>>,
167
    custom_template_path: Option<&str>,
168
) -> PyResult<Bound<'p, PyAny>> {
169
170
171
    let model_input = match model_input {
        ModelInput::Text => llm_rs::model_type::ModelInput::Text,
        ModelInput::Tokens => llm_rs::model_type::ModelInput::Tokens,
172
        ModelInput::Tensor => llm_rs::model_type::ModelInput::Tensor,
173
174
    };

175
176
    let model_type_obj = model_type.inner;

177
178
    let inner_path = model_path.to_string();
    let model_name = model_name.map(|n| n.to_string());
179
180
181
    let router_mode = router_mode.unwrap_or(RouterMode::RoundRobin);
    let router_config = RouterConfig::new(router_mode.into(), KvRouterConfig::default());

182
183
184
185
186
187
188
189
190
191
192
193
194
    // 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()?;

195
196
197
198
199
200
201
    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))
        })?;

202
    pyo3_async_runtimes::tokio::future_into_py(py, async move {
203
204
205
206
207
        let mut builder = dynamo_llm::local_model::LocalModelBuilder::default();
        builder
            .model_path(Some(PathBuf::from(inner_path)))
            .model_name(model_name)
            .context_length(context_length)
208
            .kv_cache_block_size(kv_cache_block_size)
209
            .router_config(Some(router_config))
210
            .migration_limit(Some(migration_limit))
211
            .runtime_config(runtime_config.unwrap_or_default().inner)
212
213
            .user_data(user_data_json)
            .custom_template_path(custom_template_path_owned);
214
        // Download from HF, load the ModelDeploymentCard
215
        let mut local_model = builder.build().await.map_err(to_pyerr)?;
216
217
        // Advertise ourself on etcd so ingress can find us
        local_model
218
            .attach(&endpoint.inner, model_type_obj, model_input)
219
220
221
222
223
224
225
            .await
            .map_err(to_pyerr)?;

        Ok(())
    })
}

226
227
228
229
230
231
#[pyclass]
#[derive(Clone)]
struct EtcdKvCache {
    inner: Arc<rs::transports::etcd::KvCache>,
}

232
233
#[pyclass]
#[derive(Clone)]
Ryan Olson's avatar
Ryan Olson committed
234
pub struct DistributedRuntime {
235
236
237
238
    inner: rs::DistributedRuntime,
    event_loop: PyObject,
}

Ryan Olson's avatar
Ryan Olson committed
239
240
impl DistributedRuntime {
    #[allow(dead_code)]
241
    pub(crate) fn inner(&self) -> &rs::DistributedRuntime {
Ryan Olson's avatar
Ryan Olson committed
242
243
244
245
        &self.inner
    }
}

246
247
248
249
250
251
#[pyclass]
#[derive(Clone)]
struct EtcdClient {
    inner: rs::transports::etcd::Client,
}

252
#[pyclass]
253
#[derive(Clone)]
254
255
256
257
258
struct CancellationToken {
    inner: rs::CancellationToken,
}

#[pyclass]
259
#[derive(Clone)]
260
261
262
263
264
265
struct Namespace {
    inner: rs::component::Namespace,
    event_loop: PyObject,
}

#[pyclass]
266
#[derive(Clone)]
267
268
269
270
271
272
struct Component {
    inner: rs::component::Component,
    event_loop: PyObject,
}

#[pyclass]
273
#[derive(Clone)]
274
275
276
277
278
279
struct Endpoint {
    inner: rs::component::Endpoint,
    event_loop: PyObject,
}

#[pyclass]
280
#[derive(Clone)]
281
struct Client {
282
    router: rs::pipeline::PushRouter<serde_json::Value, RsAnnotated<serde_json::Value>>,
283
284
}

285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
#[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,
    };
306
307
308
309
    #[classattr]
    const TensorBased: Self = ModelType {
        inner: llm_rs::model_type::ModelType::TensorBased,
    };
310
311
312
313
314
315
316
317
318
319
320
321

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

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

322
323
#[pyclass(eq, eq_int)]
#[derive(Clone, PartialEq)]
324
325
326
enum ModelInput {
    Text = 1,
    Tokens = 2,
327
    Tensor = 3,
328
329
}

330
331
332
#[pymethods]
impl DistributedRuntime {
    #[new]
333
    fn new(event_loop: PyObject, is_static: bool) -> PyResult<Self> {
334
        let worker = rs::Worker::from_settings().map_err(to_pyerr)?;
335
        INIT.get_or_try_init(|| {
336
            let primary = worker.tokio_runtime()?;
337
338
339
340
341
342
            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)?;

343
        let runtime = worker.runtime().clone();
344

345
346
347
348
349
350
351
352
353
354
355
        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)?;
356
357
358
359

        Ok(DistributedRuntime { inner, event_loop })
    }

Ryan Olson's avatar
Ryan Olson committed
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
    #[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(),
        })
    }

375
376
377
378
379
380
381
382
383
    /// Remove everything in an etcd namespace.
    /// Will be removed once we can clear the MDC automatically.
    fn temp_clear_namespace<'p>(&self, py: Python<'p>, name: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            inner.temp_clear_namespace(&name).await.map_err(to_pyerr)
        })
    }

384
385
386
387
388
389
390
    fn namespace(&self, name: String) -> PyResult<Namespace> {
        Ok(Namespace {
            inner: self.inner.namespace(name).map_err(to_pyerr)?,
            event_loop: self.event_loop.clone(),
        })
    }

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
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    /// 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!(
460
461
462
                                "Failed to bind to port block starting at {start_port} (attempt {}): {e}",
                                attempt_idx + 1,
                            );
463
464
465
466
467
468
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
                            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!(
518
519
                "Failed to allocate and reserve a port block of size {block_size} from range {min}-{max} after {candidate_count} attempts"
            )))
520
521
522
        })
    }

523
    fn do_not_use_etcd_client(&self) -> PyResult<Option<EtcdClient>> {
524
525
526
527
        match self.inner.etcd_client().clone() {
            Some(etcd_client) => Ok(Some(EtcdClient { inner: etcd_client })),
            None => Ok(None),
        }
528
529
    }

530
531
532
533
534
535
536
537
538
    fn shutdown(&self) {
        self.inner.runtime().shutdown();
    }

    fn event_loop(&self) -> PyObject {
        self.event_loop.clone()
    }
}

539
540
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
// 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> {
    Ok(format!("dyn://{namespace}/ports/{node_ip}/{port}"))
}

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),
    })
}

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
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
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
667
668
669
670
671
672
673
674
675
#[pymethods]
impl EtcdKvCache {
    #[new]
    fn py_new(
        _etcd_client: &EtcdClient,
        _prefix: String,
        _initial_values: &Bound<'_, PyDict>,
    ) -> PyResult<Self> {
        // We can't create the KvCache here because it's async, so we'll return an error
        Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
            "EtcdKvCache must be created using the 'new' class method",
        ))
    }

    #[staticmethod]
    #[allow(clippy::new_ret_no_self)]
    fn create<'p>(
        py: Python<'p>,
        etcd_client: &EtcdClient,
        prefix: String,
        initial_values: &Bound<'p, PyDict>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let client = etcd_client.inner.clone();

        // Convert Python dict to Rust HashMap
        let mut rust_initial_values = std::collections::HashMap::new();
        for (key, value) in initial_values.iter() {
            let key_str = key.extract::<String>()?;

            // Handle both string and bytes values
            let value_bytes = if let Ok(bytes) = value.extract::<Vec<u8>>() {
                bytes
            } else if let Ok(string) = value.extract::<String>() {
                string.into_bytes()
            } else {
                return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                    "Values must be either strings or bytes",
                ));
            };

            rust_initial_values.insert(key_str, value_bytes);
        }

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let kv_cache = rs::transports::etcd::KvCache::new(client, prefix, rust_initial_values)
                .await
                .map_err(to_pyerr)?;

            Ok(EtcdKvCache {
                inner: Arc::new(kv_cache),
            })
        })
    }

    fn get<'p>(&self, py: Python<'p>, key: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            if let Some(value) = inner.get(&key).await {
                match Python::with_gil(|py| {
                    let py_obj = PyBytes::new(py, &value).into_pyobject(py)?;
                    Ok(py_obj.unbind().into_any())
                }) {
                    Ok(result) => Ok(result),
                    Err(e) => Err(e),
                }
            } else {
                Ok(Python::with_gil(|py| py.None()))
            }
        })
    }

    fn get_all<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let all_values = inner.get_all().await;

            Python::with_gil(|py| {
                let dict = PyDict::new(py);
                for (key, value) in all_values {
                    // Strip the prefix from the key
                    let stripped_key = if let Some(stripped) = key.strip_prefix(&inner.prefix) {
                        stripped.to_string()
                    } else {
                        key
                    };
                    dict.set_item(stripped_key, PyBytes::new(py, &value))?;
                }
                let py_obj = dict.into_pyobject(py)?;
                Ok(py_obj.unbind().into_any())
            })
        })
    }

    #[pyo3(signature = (key, value, lease_id=None))]
    fn put<'p>(
        &self,
        py: Python<'p>,
        key: String,
        value: Vec<u8>,
        lease_id: Option<i64>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            inner.put(&key, value, lease_id).await.map_err(to_pyerr)?;
            Ok(())
        })
    }
676
677
678
679
680
681
682
683
684

    fn delete<'p>(&self, py: Python<'p>, key: String) -> PyResult<Bound<'p, PyAny>> {
        let inner = self.inner.clone();

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            inner.delete(&key).await.map_err(to_pyerr)?;
            Ok(())
        })
    }
685
686
}

687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
#[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(),
        })
    }

    fn create_service<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
        let builder = self.inner.service_builder();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let _ = builder.create().await.map_err(to_pyerr)?;
            Ok(())
        })
    }
}

#[pymethods]
impl Endpoint {
723
    #[pyo3(signature = (generator, graceful_shutdown = true, metrics_labels = None, health_check_payload = None))]
724
725
726
727
    fn serve_endpoint<'p>(
        &self,
        py: Python<'p>,
        generator: PyObject,
728
        graceful_shutdown: Option<bool>,
729
        metrics_labels: Option<Vec<(String, String)>>,
730
        health_check_payload: Option<&Bound<'p, PyDict>>,
731
732
733
734
735
736
    ) -> 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)?;
737
738
739
740
741
742
743
744
745
746
747
748
749

        // 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
750
751
752
753
754
755
        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)",
            ));
756
757
758
        }

        let mut builder = self
759
760
761
762
            .inner
            .endpoint_builder()
            .metrics_labels(metrics_labels)
            .handler(ingress);
763
764
765
766
767

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

768
        let graceful_shutdown = graceful_shutdown.unwrap_or(true);
769
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
770
771
772
773
774
            builder
                .graceful_shutdown(graceful_shutdown)
                .start()
                .await
                .map_err(to_pyerr)?;
775
776
777
778
            Ok(())
        })
    }

779
780
781
    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 {
782
            let client = inner.client().await.map_err(to_pyerr)?;
783
784
785
786
787
788
            let push_router = rs::pipeline::PushRouter::<
                serde_json::Value,
                RsAnnotated<serde_json::Value>,
            >::from_client(client, Default::default())
            .await
            .map_err(to_pyerr)?;
789
790
791
            Ok(Client {
                router: push_router,
            })
792
793
        })
    }
794
795

    fn lease_id(&self) -> i64 {
796
797
798
799
800
        self.inner
            .drt()
            .primary_lease()
            .map(|l| l.id())
            .unwrap_or(0)
801
    }
802
803
804
805
806
807
808
809
810
811
812
813
814
}

#[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(),
        })
    }
}

815
816
#[pymethods]
impl EtcdClient {
817
818
819
820
821
822
823
824
825
826
827
    #[pyo3(signature = (key, value, lease_id=None))]
    fn kv_create<'p>(
        &self,
        py: Python<'p>,
        key: String,
        value: Vec<u8>,
        lease_id: Option<i64>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let client = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            client
828
                .kv_create(&key, value, lease_id)
829
830
831
832
833
834
                .await
                .map_err(to_pyerr)?;
            Ok(())
        })
    }

835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
    #[pyo3(signature = (key, value, lease_id=None))]
    fn kv_create_or_validate<'p>(
        &self,
        py: Python<'p>,
        key: String,
        value: Vec<u8>,
        lease_id: Option<i64>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let client = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            client
                .kv_create_or_validate(key, value, lease_id)
                .await
                .map_err(to_pyerr)?;
            Ok(())
        })
    }

853
854
855
856
    fn primary_lease_id(&self) -> i64 {
        self.inner.lease_id()
    }

857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
    #[pyo3(signature = (key, value, lease_id=None))]
    fn kv_put<'p>(
        &self,
        py: Python<'p>,
        key: String,
        value: Vec<u8>,
        lease_id: Option<i64>,
    ) -> PyResult<Bound<'p, PyAny>> {
        let client = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            client
                .kv_put(key, value, lease_id)
                .await
                .map_err(to_pyerr)?;
            Ok(())
        })
    }

    fn kv_get_prefix<'p>(&self, py: Python<'p>, prefix: String) -> PyResult<Bound<'p, PyAny>> {
        let client = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let result = client
                .kv_get_prefix(prefix)
                .await
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;

            // Convert Vec<KeyValue> to a list of dictionaries
            let py_list = Python::with_gil(|py| {
                let list = PyList::empty(py);
                for kv in result {
                    let dict = PyDict::new(py);
                    dict.set_item("key", String::from_utf8_lossy(kv.key()).to_string())?;
                    dict.set_item("value", PyBytes::new(py, kv.value()))?;
                    dict.set_item("create_revision", kv.create_revision())?;
                    dict.set_item("mod_revision", kv.mod_revision())?;
                    dict.set_item("version", kv.version())?;
                    dict.set_item("lease", kv.lease())?;
                    list.append(dict)?;
                }
                Ok::<Py<PyList>, PyErr>(list.into())
            })?;

            Ok(py_list)
        })
    }
902
903
904
905
906
907
908
909

    fn revoke_lease<'p>(&self, py: Python<'p>, lease_id: i64) -> PyResult<Bound<'p, PyAny>> {
        let client = self.inner.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            client.revoke_lease(lease_id).await.map_err(to_pyerr)?;
            Ok(())
        })
    }
910
911
}

912
913
#[pymethods]
impl Client {
914
915
916
917
    /// Get list of current instances.
    /// Replaces endpoint_ids.
    fn instance_ids(&self) -> Vec<i64> {
        self.router.client.instance_ids()
918
919
    }

920
921
922
    /// 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>> {
923
        let inner = self.router.client.clone();
924
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
925
            inner
926
                .wait_for_instances()
927
928
929
                .await
                .map(|v| v.into_iter().map(|cei| cei.id()).collect::<Vec<i64>>())
                .map_err(to_pyerr)
930
931
932
933
        })
    }

    /// Issue a request to the endpoint using the default routing strategy.
934
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
935
936
937
938
939
    fn generate<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
940
        context: Option<context::Context>,
941
    ) -> PyResult<Bound<'p, PyAny>> {
942
        if self.router.client.is_static() {
943
            self.r#static(py, request, annotated, context)
944
        } else {
945
            self.random(py, request, annotated, context)
946
        }
947
948
949
    }

    /// Send a request to the next endpoint in a round-robin fashion.
950
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
951
952
953
954
955
    fn round_robin<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
956
        context: Option<context::Context>,
957
958
959
960
961
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
962
        let client = self.router.clone();
963
964

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
965
966
967
968
969
970
971
972
973
            let stream = match context {
                Some(context) => {
                    let request = RsContext::with_id(request, context.inner().id().to_string());
                    let stream = client.round_robin(request).await.map_err(to_pyerr)?;
                    context.inner().link_child(stream.context());
                    stream
                }
                _ => client.round_robin(request.into()).await.map_err(to_pyerr)?,
            };
974
975
976
977
978
979
980
981
982
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Send a request to a random endpoint.
983
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
984
985
986
987
988
    fn random<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
989
        context: Option<context::Context>,
990
991
992
993
994
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
995
        let client = self.router.clone();
996
997

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
998
999
1000
1001
1002
1003
1004
1005
1006
            let stream = match context {
                Some(context) => {
                    let request = RsContext::with_id(request, context.inner().id().to_string());
                    let stream = client.random(request).await.map_err(to_pyerr)?;
                    context.inner().link_child(stream.context());
                    stream
                }
                _ => client.random(request.into()).await.map_err(to_pyerr)?,
            };
1007
1008
1009
1010
1011
1012
1013
1014
1015
            tokio::spawn(process_stream(stream, tx));
            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }

    /// Directly send a request to a specific endpoint.
1016
    #[pyo3(signature = (request, instance_id, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
1017
1018
1019
1020
    fn direct<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
1021
        instance_id: i64,
1022
        annotated: Option<bool>,
1023
        context: Option<context::Context>,
1024
1025
1026
1027
1028
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
1029
        let client = self.router.clone();
1030
1031

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
            let stream = match context {
                Some(context) => {
                    let request = RsContext::with_id(request, context.inner().id().to_string());
                    let stream = client
                        .direct(request, instance_id)
                        .await
                        .map_err(to_pyerr)?;
                    context.inner().link_child(stream.context());
                    stream
                }
                _ => client
                    .direct(request.into(), instance_id)
                    .await
                    .map_err(to_pyerr)?,
            };
1047
1048
1049
1050
1051
1052
1053
1054
1055

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

            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }
1056
1057

    /// Directly send a request to a pre-defined static worker
1058
    #[pyo3(signature = (request, annotated=DEFAULT_ANNOTATED_SETTING, context=None))]
1059
1060
1061
1062
1063
    fn r#static<'p>(
        &self,
        py: Python<'p>,
        request: PyObject,
        annotated: Option<bool>,
1064
        context: Option<context::Context>,
1065
1066
1067
1068
1069
    ) -> PyResult<Bound<'p, PyAny>> {
        let request: serde_json::Value = pythonize::depythonize(&request.into_bound(py))?;
        let annotated = annotated.unwrap_or(false);

        let (tx, rx) = tokio::sync::mpsc::channel(32);
1070
        let client = self.router.clone();
1071
1072

        pyo3_async_runtimes::tokio::future_into_py(py, async move {
1073
1074
1075
1076
1077
1078
1079
1080
1081
            let stream = match context {
                Some(context) => {
                    let request = RsContext::with_id(request, context.inner().id().to_string());
                    let stream = client.r#static(request).await.map_err(to_pyerr)?;
                    context.inner().link_child(stream.context());
                    stream
                }
                _ => client.r#static(request.into()).await.map_err(to_pyerr)?,
            };
1082
1083
1084
1085
1086
1087
1088
1089
1090

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

            Ok(AsyncResponseStream {
                rx: Arc::new(Mutex::new(rx)),
                annotated,
            })
        })
    }
1091
1092
1093
}

async fn process_stream(
1094
    stream: EngineStream<RsAnnotated<serde_json::Value>>,
1095
1096
1097
1098
1099
    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
1100
        let annotated: RsAnnotated<serde_json::Value> = response;
1101
        let annotated: RsAnnotated<PyObject> = annotated.map_data(|data| {
1102
            Python::with_gil(|py| match pythonize::pythonize(py, &data) {
1103
1104
                Ok(pyobj) => Ok(pyobj.into()),
                Err(e) => Err(e.to_string()),
1105
            })
1106
1107
1108
1109
1110
1111
        });

        let is_error = annotated.is_error();

        // Send the PyObject through the channel or log an error
        if let Err(e) = tx.send(annotated).await {
1112
            tracing::error!("Failed to send response: {:?}", e);
1113
            break;
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
        }

        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")
        )
    }
}