feat: expose Python binding for KVEventPublisher. Use event pub/sub trait for KV events (ai-dynamo#169)

Author: GuanLuo

lib/bindings/c/src/lib.rs

@@ -148,7 +148,7 @@ fn dynamo_create_kv_publisher(
     {
         Ok(drt) => {
             let backend = drt.namespace(namespace)?.component(component)?;
-            KvEventPublisher::new(drt.clone(), backend, worker_id, kv_block_size)
+            KvEventPublisher::new(backend, worker_id, kv_block_size)
         }
         Err(e) => Err(e),
     }

lib/bindings/python/rust/lib.rs

@@ -77,6 +77,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<llm::kv::EndpointKvMetrics>()?;
     m.add_class::<llm::kv::AggregatedMetrics>()?;
     m.add_class::<llm::kv::KvMetricsAggregator>()?;
+    m.add_class::<llm::kv::KvEventPublisher>()?;
     m.add_class::<http::HttpService>()?;
     m.add_class::<http::HttpError>()?;
     m.add_class::<http::HttpAsyncEngine>()?;
@@ -227,10 +228,6 @@ impl Component {
             Ok(())
         })
     }
-
-    fn event_subject(&self, name: String) -> String {
-        self.inner.event_subject(name)
-    }
 }
 
 #[pymethods]

lib/bindings/python/rust/llm/kv.rs

@@ -17,8 +17,11 @@ use std::collections::HashMap;
 
 use super::*;
 use llm_rs::kv_router::indexer::KvIndexerInterface;
+use rs::traits::events::EventSubscriber;
 use tracing;
 
+use llm_rs::kv_router::{indexer::compute_block_hash_for_seq, protocols::*};
+
 #[pyclass]
 pub(crate) struct KvRouter {
     inner: Arc<llm_rs::kv_router::KvRouter>,
@@ -119,6 +122,114 @@ impl KvMetricsPublisher {
     }
 }
 
+#[pyclass]
+pub(crate) struct KvEventPublisher {
+    inner: Arc<llm_rs::kv_router::publisher::KvEventPublisher>,
+    warning_count: u32,
+}
+
+#[pymethods]
+impl KvEventPublisher {
+    #[new]
+    fn new(component: Component, worker_id: i64, kv_block_size: usize) -> PyResult<Self> {
+        let inner = llm_rs::kv_router::publisher::KvEventPublisher::new(
+            component.inner.clone(),
+            worker_id,
+            kv_block_size,
+        )
+        .map_err(to_pyerr)?;
+        Ok(Self {
+            inner: inner.into(),
+            warning_count: 0,
+        })
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    #[pyo3(signature = (event_id, token_ids, num_block_tokens, block_hashes, lora_id, parent_hash=None))]
+    fn publish_stored(
+        &mut self,
+        _py: Python,
+        event_id: u64,
+        token_ids: Vec<u32>,
+        num_block_tokens: Vec<u64>,
+        block_hashes: Vec<u64>,
+        lora_id: u64,
+        parent_hash: Option<u64>,
+    ) -> PyResult<()> {
+        let event = KvCacheEvent {
+            event_id,
+            data: KvCacheEventData::Stored(KvCacheStoreData {
+                parent_hash: parent_hash.map(ExternalSequenceBlockHash),
+                blocks: self.create_stored_blocks(
+                    &token_ids,
+                    &num_block_tokens,
+                    &block_hashes,
+                    lora_id,
+                ),
+            }),
+        };
+
+        self.inner.publish(event).map_err(to_pyerr)
+    }
+
+    fn publish_removed(&self, _py: Python, event_id: u64, block_hashes: Vec<u64>) -> PyResult<()> {
+        let block_hashes: Vec<ExternalSequenceBlockHash> = block_hashes
+            .iter()
+            .map(|&v| ExternalSequenceBlockHash(v))
+            .collect();
+        let event = KvCacheEvent {
+            event_id,
+            data: KvCacheEventData::Removed(KvCacheRemoveData { block_hashes }),
+        };
+
+        self.inner.publish(event).map_err(to_pyerr)
+    }
+}
+
+impl KvEventPublisher {
+    fn create_stored_block_from_parts(
+        &self,
+        block_hash: u64,
+        token_ids: &[u32],
+        _lora_id: u64,
+    ) -> KvCacheStoredBlockData {
+        let tokens_hash = compute_block_hash_for_seq(token_ids, self.inner.kv_block_size())[0];
+        KvCacheStoredBlockData {
+            block_hash: ExternalSequenceBlockHash(block_hash),
+            tokens_hash,
+        }
+    }
+
+    fn create_stored_blocks(
+        &mut self,
+        token_ids: &[u32],
+        num_block_tokens: &[u64],
+        block_hashes: &[u64],
+        lora_id: u64,
+    ) -> Vec<KvCacheStoredBlockData> {
+        let mut blocks: Vec<KvCacheStoredBlockData> = Vec::new();
+
+        let mut token_offset: usize = 0;
+        for (num_tokens_it, block_hash_it) in num_block_tokens.iter().zip(block_hashes.iter()) {
+            if (self.warning_count < 3) && (*num_tokens_it != self.inner.kv_block_size() as u64) {
+                tracing::warn!(
+                    "Block not published. Block size must be {} tokens to be published. Block size is: {}",
+                    self.inner.kv_block_size(),
+                    *num_tokens_it
+                );
+                self.warning_count += 1;
+                break;
+            }
+
+            let tokens = &token_ids[token_offset..(token_offset + *num_tokens_it as usize)];
+            blocks.push(self.create_stored_block_from_parts(*block_hash_it, tokens, lora_id));
+            token_offset += *num_tokens_it as usize;
+        }
+
+        blocks
+    }
+}
+
 #[pyclass]
 #[derive(Clone)]
 pub(crate) struct OverlapScores {
@@ -149,21 +260,17 @@ impl KvIndexer {
     fn new(component: Component, kv_block_size: usize) -> PyResult<Self> {
         let runtime = pyo3_async_runtimes::tokio::get_runtime();
         runtime.block_on(async {
-            let kv_subject = component
-                .inner
-                .event_subject(llm_rs::kv_router::KV_EVENT_SUBJECT);
             let inner: Arc<llm_rs::kv_router::indexer::KvIndexer> =
                 llm_rs::kv_router::indexer::KvIndexer::new(
                     component.inner.drt().runtime().child_token(),
                     kv_block_size,
                 )
                 .into();
+            // [gluo TODO] try subscribe_with_type::<RouterEvent>,
+            // error checking below will be different.
             let mut kv_events_rx = component
                 .inner
-                .drt()
-                .nats_client()
-                .client()
-                .subscribe(kv_subject)
+                .subscribe(llm_rs::kv_router::KV_EVENT_SUBJECT)
                 .await
                 .map_err(to_pyerr)?;
             let kv_events_tx = inner.event_sender();

lib/bindings/python/src/dynamo/_core.pyi

@@ -102,11 +102,6 @@ class Component:
         """
         ...
 
-    def event_subject(self, name: str) -> str:
-        """
-        Create an event subject
-        """
-        ...
 
 class Endpoint:
     """
@@ -354,6 +349,30 @@ class KvMetricsAggregator:
         """
         ...
 
+class KvEventPublisher:
+    """
+    A KV event publisher will publish KV events corresponding to the component.
+    """
+
+    ...
+
+    def __init__(self, component: Component, worker_id: int, kv_block_size: int) -> None:
+        """
+        Create a `KvEventPublisher` object
+        """
+
+    def publish_stored(self, event_id, int, token_ids: List[int], num_block_tokens: List[int], block_hashes: List[int], lora_id: int, parent_hash: Optional[int] = None) -> None:
+        """
+        Publish a KV stored event.
+        """
+        ...
+
+    def publish_removed(self, event_id, int, block_hashes: List[int]) -> None:
+        """
+        Publish a KV removed event.
+        """
+        ...
+
 class HttpService:
     """
     A HTTP service for dynamo applications.

lib/bindings/python/src/dynamo/llm/__init__.py

@@ -18,6 +18,7 @@
 from dynamo._core import HttpAsyncEngine as HttpAsyncEngine
 from dynamo._core import HttpError as HttpError
 from dynamo._core import HttpService as HttpService
+from dynamo._core import KvEventPublisher as KvEventPublisher
 from dynamo._core import KvIndexer as KvIndexer
 from dynamo._core import KvMetricsAggregator as KvMetricsAggregator
 from dynamo._core import KvMetricsPublisher as KvMetricsPublisher

lib/bindings/python/src/dynamo/runtime/__init__.py

@@ -24,8 +24,8 @@
 # import * causes "unable to detect undefined names"
 from dynamo._core import Backend as Backend
 from dynamo._core import Client as Client
+from dynamo._core import Component as Component
 from dynamo._core import DistributedRuntime as DistributedRuntime
-from dynamo._core import KvRouter as KvRouter
 from dynamo._core import ModelDeploymentCard as ModelDeploymentCard
 from dynamo._core import OAIChatPreprocessor as OAIChatPreprocessor
 

lib/bindings/python/tests/test_kv_bindings.py

@@ -24,8 +24,13 @@
 
 import pytest
 
-from dynamo.llm import KvIndexer, KvMetricsAggregator, KvMetricsPublisher
-from dynamo.runtime import DistributedRuntime
+from dynamo.llm import (
+    KvEventPublisher,
+    KvIndexer,
+    KvMetricsAggregator,
+    KvMetricsPublisher,
+)
+from dynamo.runtime import Component, DistributedRuntime
 
 pytestmark = pytest.mark.pre_merge
 
@@ -64,14 +69,14 @@ async def test_event_handler(distributed_runtime):
     kv_block_size = 32
     namespace = "kv_test"
     component = "event"
+    kv_listener = distributed_runtime.namespace(namespace).component(component)
+    await kv_listener.create_service()
 
     # publisher
     worker_id = 233
-    event_publisher = EventPublisher(namespace, component, worker_id, kv_block_size)
+    event_publisher = EventPublisher(kv_listener, worker_id, kv_block_size)
 
     # indexer
-    kv_listener = distributed_runtime.namespace(namespace).component(component)
-    await kv_listener.create_service()
     indexer = KvIndexer(kv_listener, kv_block_size)
 
     test_token = [3] * kv_block_size
@@ -93,16 +98,48 @@ async def test_event_handler(distributed_runtime):
     scores = await indexer.find_matches_for_request(test_token, lora_id)
     assert not scores.scores
 
-    event_publisher.shutdown()
+
+class EventPublisher:
+    def __init__(self, component: Component, worker_id: int, kv_block_size: int):
+        self.publisher = KvEventPublisher(component, worker_id, kv_block_size)
+        self.event_id_counter = 0
+        self.block_hashes: List[int] = []
+
+    def store_event(self, tokens, lora_id):
+        parent_hash = self.event_id_counter if self.event_id_counter > 0 else None
+        self.publisher.publish_stored(
+            self.event_id_counter,  # event_id
+            tokens,  # token_ids
+            [
+                len(tokens),
+            ],  # num_block_tokens
+            [
+                self.event_id_counter,
+            ],  # block_hashes
+            lora_id,  # lora_id
+            parent_hash,  # parent_hash
+        )
+        self.block_hashes.append(self.event_id_counter)
+        self.event_id_counter += 1
+
+    def remove_event(self):
+        self.publisher.publish_removed(
+            self.event_id_counter,  # event_id
+            [
+                self.block_hashes[-1],
+            ],  # block_hashes
+        )
+        self.event_id_counter += 1
 
 
+# [TODO] to be deprecated
 # KV events
 class DynamoResult:
     OK = 0
     ERR = 1
 
 
-class EventPublisher:
+class CtypesEventPublisher:
     def __init__(
         self, namespace: str, component: str, worker_id: int, kv_block_size: int
     ):

lib/llm/src/kv_router.rs

@@ -29,14 +29,18 @@ pub mod scoring;
 
 use crate::kv_router::{
     indexer::{KvIndexer, KvIndexerInterface, RouterEvent},
-    metrics_aggregator::collect_endpoints,
+    metrics_aggregator::collect_endpoints_task,
     scheduler::KvScheduler,
     scoring::ProcessedEndpoints,
 };
 
-// this should be discovered from the backend
+use dynamo_runtime::traits::events::{EventPublisher, EventSubscriber};
+
+// [gluo TODO] shouldn't need to be public
+// this should be discovered from the component
 pub const KV_EVENT_SUBJECT: &str = "kv_events";
 pub const KV_HIT_RATE_SUBJECT: &str = "kv-hit-rate";
+pub const KV_METRICS_ENDPOINT: &str = "load_metrics";
 
 pub struct KvRouter {
     // properties of request plane
@@ -54,49 +58,37 @@ pub struct KvRouter {
 impl KvRouter {
     pub async fn from_runtime(
         runtime: DistributedRuntime,
-        backend: Component,
+        component: Component,
         kv_block_size: usize,
     ) -> Result<Arc<Self>> {
-        let nats_client = runtime.nats_client();
-        let service_name = backend.service_name();
-        let kv_subject = backend.event_subject(KV_EVENT_SUBJECT);
-        let namespace = runtime.namespace(backend.namespace())?;
-
-        tracing::info!("Component Namespace {}", backend.namespace());
-        tracing::info!("Component Service Name {}", service_name);
-        tracing::info!("KV Subject {}", kv_subject);
-        Self::new(
-            nats_client,
-            service_name,
-            kv_subject,
-            namespace,
-            kv_block_size,
-        )
-        .await
+        let namespace = runtime.namespace(component.namespace().name())?;
+
+        tracing::info!("Component Namespace {}", component.namespace());
+        tracing::info!("Component Service Name {}", component.service_name());
+        tracing::info!("KV Subject {}.{}", component.subject(), KV_EVENT_SUBJECT);
+        Self::new(component, namespace, kv_block_size).await
     }
 
     pub async fn new(
-        nats_client: dynamo_runtime::transports::nats::Client,
-        service_name: String,
-        kv_subject: String,
+        component: Component,
         namespace: Namespace,
         kv_block_size: usize,
     ) -> Result<Arc<Self>> {
         let cancellation_token = CancellationToken::new();
         let (ep_tx, ep_rx) = tokio::sync::mpsc::channel(128);
 
-        tokio::spawn(collect_endpoints(
-            nats_client.clone(),
-            service_name.clone(),
+        tokio::spawn(collect_endpoints_task(
+            component.clone(),
             ep_tx,
             cancellation_token.clone(),
         ));
 
         let indexer = KvIndexer::new(cancellation_token.clone(), kv_block_size);
         let scheduler = KvScheduler::start(ep_rx, namespace, kv_block_size).await?;
 
-        tracing::debug!("subscribing to kv events: {}", kv_subject);
-        let mut kv_events_rx = nats_client.client().subscribe(kv_subject).await?;
+        // [gluo TODO] try subscribe_with_type::<RouterEvent>,
+        // error checking below will be different.
+        let mut kv_events_rx = component.subscribe(KV_EVENT_SUBJECT).await?;
         let kv_events_tx = indexer.event_sender();
 
         tokio::spawn(async move {
@@ -120,7 +112,7 @@ impl KvRouter {
         });
 
         Ok(Arc::new(Self {
-            service_name,
+            service_name: component.service_name(),
             cancellation_token,
             scheduler,
             indexer,

lib/llm/src/kv_router/metrics_aggregator.rs

@@ -16,11 +16,12 @@
 use std::sync::{Arc, Mutex};
 
 pub use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::kv_router::KV_METRICS_ENDPOINT;
 
-use crate::kv_router::scheduler::{Endpoint, Service};
+use crate::kv_router::scheduler::Endpoint;
 use crate::kv_router::ProcessedEndpoints;
 use dynamo_runtime::component::Component;
-use std::time::Duration;
+use dynamo_runtime::{service::EndpointInfo, utils::Duration, Result};
 use tokio_util::sync::CancellationToken;
 
 pub struct KvMetricsAggregator {
@@ -32,9 +33,8 @@ impl KvMetricsAggregator {
     pub async fn new(component: Component, cancellation_token: CancellationToken) -> Self {
         let (ep_tx, mut ep_rx) = tokio::sync::mpsc::channel(128);
 
-        tokio::spawn(collect_endpoints(
-            component.drt().nats_client().clone(),
-            component.service_name(),
+        tokio::spawn(collect_endpoints_task(
+            component.clone(),
             ep_tx,
             cancellation_token.clone(),
         ));
@@ -80,13 +80,41 @@ impl KvMetricsAggregator {
     }
 }
 
+/// [gluo TODO] 'collect_endpoints' is from component/metrics,
+/// should consolidate these functions into generic metrics aggregator
+/// functions and shared by KvMetricsAggregator and component/metrics.
+/// Collect endpoints from a component
 pub async fn collect_endpoints(
-    nats_client: dynamo_runtime::transports::nats::Client,
-    service_name: String,
+    component: &Component,
+    subject: &str,
+    timeout: Duration,
+) -> Result<Vec<EndpointInfo>> {
+    // Collect stats from each backend
+    let stream = component.scrape_stats(timeout).await?;
+
+    // Filter the stats by the service subject
+    let endpoints = stream
+        .into_endpoints()
+        .filter(|e| e.subject.starts_with(subject))
+        .collect::<Vec<_>>();
+    tracing::debug!("Endpoints: {endpoints:?}");
+
+    if endpoints.is_empty() {
+        tracing::warn!("No endpoints found matching subject {subject}");
+    }
+
+    Ok(endpoints)
+}
+
+pub async fn collect_endpoints_task(
+    component: Component,
     ep_tx: tokio::sync::mpsc::Sender<ProcessedEndpoints>,
     cancel: CancellationToken,
 ) {
     let backoff_delay = Duration::from_millis(100);
+    let scrape_timeout = Duration::from_millis(300);
+    let endpoint = component.endpoint(KV_METRICS_ENDPOINT);
+    let service_subject = endpoint.subject();
 
     loop {
         tokio::select! {
@@ -95,48 +123,41 @@ pub async fn collect_endpoints(
                 break;
             }
             _ = tokio::time::sleep(backoff_delay) => {
-                tracing::trace!("collecting endpoints for service: {}", service_name);
-                let values = match nats_client
-                    .get_endpoints(&service_name, Duration::from_millis(300))
-                    .await
-                {
-                    Ok(v) => v,
-                    Err(e) => {
-                        tracing::warn!("Failed to retrieve endpoints for {}: {:?}", service_name, e);
-                        continue;
-                    }
-                };
-
-                tracing::debug!("values: {:?}", values);
-                let services: Vec<Service> = values
-                    .into_iter()
-                    .filter(|v| !v.is_empty())
-                    .filter_map(|v| match serde_json::from_slice::<Service>(&v) {
-                        Ok(service) => Some(service),
+                tracing::trace!("collecting endpoints for service: {}", service_subject);
+                let unfiltered_endpoints =
+                    match collect_endpoints(&component, &service_subject, scrape_timeout).await
+                    {
+                        Ok(v) => v,
                         Err(e) => {
-                            tracing::warn!("For value: {:?} \nFailed to parse service: {:?}", v, e);
-                            None
+                            tracing::warn!("Failed to retrieve endpoints for {}: {:?}", service_subject, e);
+                            continue;
                         }
-                    })
-                    .collect();
-                tracing::debug!("services: {:?}", services);
+                    };
+                tracing::debug!("unfiltered endpoints: {:?}", unfiltered_endpoints);
 
-                let endpoints: Vec<Endpoint> = services
+                let endpoints: Vec<Endpoint> = unfiltered_endpoints
                     .into_iter()
-                    .flat_map(|s| s.endpoints)
                     .filter(|s| s.data.is_some())
-                    .map(|s| Endpoint {
-                        name: s.name,
-                        subject: s.subject,
-                        data: s.data.unwrap(),
-                    })
+                    .filter_map(|s|
+                        match s.data.unwrap().decode::<ForwardPassMetrics>() {
+                            Ok(data) => Some(Endpoint {
+                                name: s.name,
+                                subject: s.subject,
+                                data,
+                            }),
+                            Err(e) => {
+                                tracing::debug!("skip endpoint data that can't be parsed as ForwardPassMetrics: {:?}", e);
+                                None
+                            }
+                        }
+                    )
                     .collect();
                 tracing::debug!("endpoints: {:?}", endpoints);
 
                 tracing::trace!(
                     "found {} endpoints for service: {}",
                     endpoints.len(),
-                    service_name
+                    service_subject
                 );
 
                 let processed = ProcessedEndpoints::new(endpoints);

lib/llm/src/kv_router/publisher.rs

@@ -13,16 +13,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::kv_router::{indexer::RouterEvent, protocols::*, KV_EVENT_SUBJECT};
+use crate::kv_router::{indexer::RouterEvent, protocols::*, KV_EVENT_SUBJECT, KV_METRICS_ENDPOINT};
 use async_trait::async_trait;
+use dynamo_runtime::traits::{events::EventPublisher, DistributedRuntimeProvider};
 use dynamo_runtime::{
     component::Component,
     pipeline::{
         network::Ingress, AsyncEngine, AsyncEngineContextProvider, ManyOut, ResponseStream,
         SingleIn,
     },
     protocols::annotated::Annotated,
-    DistributedRuntime, Error, Result,
+    Error, Result,
 };
 use futures::stream;
 use std::sync::Arc;
@@ -35,16 +36,11 @@ pub struct KvEventPublisher {
 }
 
 impl KvEventPublisher {
-    pub fn new(
-        drt: DistributedRuntime,
-        backend: Component,
-        worker_id: i64,
-        kv_block_size: usize,
-    ) -> Result<Self> {
+    pub fn new(component: Component, worker_id: i64, kv_block_size: usize) -> Result<Self> {
         let (tx, rx) = mpsc::unbounded_channel::<KvCacheEvent>();
         let p = KvEventPublisher { tx, kv_block_size };
 
-        start_publish_task(drt, backend, worker_id, rx);
+        start_publish_task(component, worker_id, rx);
         Ok(p)
     }
 
@@ -59,21 +55,18 @@ impl KvEventPublisher {
 }
 
 fn start_publish_task(
-    drt: DistributedRuntime,
-    backend: Component,
+    component: Component,
     worker_id: i64,
     mut rx: mpsc::UnboundedReceiver<KvCacheEvent>,
 ) {
-    let client = drt.nats_client().client().clone();
-    let kv_subject = backend.event_subject(KV_EVENT_SUBJECT);
-    log::info!("Publishing KV Events to subject: {}", kv_subject);
+    let component_clone = component.clone();
+    log::info!("Publishing KV Events to subject: {}", KV_EVENT_SUBJECT);
 
-    _ = drt.runtime().secondary().spawn(async move {
+    _ = component.drt().runtime().secondary().spawn(async move {
         while let Some(event) = rx.recv().await {
             let router_event = RouterEvent::new(worker_id, event);
-            let data = serde_json::to_string(&router_event).unwrap();
-            client
-                .publish(kv_subject.to_string(), data.into())
+            component_clone
+                .publish(KV_EVENT_SUBJECT, &router_event)
                 .await
                 .unwrap();
         }
@@ -105,7 +98,7 @@ impl KvMetricsPublisher {
         let handler = Ingress::for_engine(handler)?;
 
         component
-            .endpoint("load_metrics")
+            .endpoint(KV_METRICS_ENDPOINT)
             .endpoint_builder()
             .stats_handler(move |_| {
                 let metrics = metrics_rx.borrow_and_update().clone();

lib/llm/src/kv_router/scheduler.rs

@@ -43,13 +43,8 @@ pub enum KvSchedulerError {
     SubscriberShutdown,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct FlexibleEndpoint {
-    pub name: String,
-    pub subject: String,
-    pub data: Option<ForwardPassMetrics>,
-}
-
+/// [gluo FIXME] exactly the same as EndpointInfo except that 'data'
+/// is cleaned (not optional)
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Endpoint {
     pub name: String,
@@ -72,24 +67,6 @@ impl Endpoint {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct FlexibleService {
-    pub name: String,
-    pub id: String,
-    pub version: String,
-    pub started: String,
-    pub endpoints: Vec<FlexibleEndpoint>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct Service {
-    pub name: String,
-    pub id: String,
-    pub version: String,
-    pub started: String,
-    pub endpoints: Vec<FlexibleEndpoint>,
-}
-
 pub struct SchedulingRequest {
     isl_tokens: usize,
     overlap: OverlapScores,

lib/runtime/src/component.rs

@@ -61,6 +61,8 @@ use std::{collections::HashMap, sync::Arc};
 use validator::{Validate, ValidationError};
 
 mod client;
+#[allow(clippy::module_inception)]
+mod component;
 mod endpoint;
 mod namespace;
 mod registry;
@@ -115,12 +117,12 @@ pub struct Component {
     // todo - restrict the namespace to a-z0-9-_A-Z
     /// Namespace
     #[builder(setter(into))]
-    namespace: String,
+    namespace: Namespace,
 }
 
 impl std::fmt::Display for Component {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}.{}", self.namespace, self.name)
+        write!(f, "{}.{}", self.namespace.name(), self.name)
     }
 }
 
@@ -138,30 +140,21 @@ impl RuntimeProvider for Component {
 
 impl Component {
     pub fn etcd_path(&self) -> String {
-        format!("{}/components/{}", self.namespace, self.name)
+        format!("{}/components/{}", self.namespace.name(), self.name)
     }
 
     pub fn service_name(&self) -> String {
-        Slug::from_string(format!("{}|{}", self.namespace, self.name)).to_string()
-    }
-
-    // todo - move to EventPlane
-    pub fn event_subject(&self, name: impl AsRef<str>) -> String {
-        format!("{}.events.{}", self.service_name(), name.as_ref())
+        Slug::from_string(format!("{}|{}", self.namespace.name(), self.name)).to_string()
     }
 
     pub fn path(&self) -> String {
-        format!("{}/{}", self.namespace, self.name)
+        format!("{}/{}", self.namespace.name(), self.name)
     }
 
-    pub fn namespace(&self) -> &str {
+    pub fn namespace(&self) -> &Namespace {
         &self.namespace
     }
 
-    pub fn drt(&self) -> &DistributedRuntime {
-        &self.drt
-    }
-
     pub fn endpoint(&self, endpoint: impl Into<String>) -> Endpoint {
         Endpoint {
             component: self.clone(),
@@ -300,6 +293,12 @@ impl RuntimeProvider for Namespace {
     }
 }
 
+impl std::fmt::Display for Namespace {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}
+
 impl Namespace {
     pub(crate) fn new(runtime: DistributedRuntime, name: String) -> Result<Self> {
         Ok(NamespaceBuilder::default()
@@ -312,7 +311,7 @@ impl Namespace {
     pub fn component(&self, name: impl Into<String>) -> Result<Component> {
         Ok(ComponentBuilder::from_runtime(self.runtime.clone())
             .name(name)
-            .namespace(self.name.clone())
+            .namespace(self.clone())
             .build()?)
     }
 

lib/runtime/src/component/component.rs

@@ -0,0 +1,121 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use anyhow::Context;
+use async_trait::async_trait;
+use futures::stream::StreamExt;
+use futures::{Stream, TryStreamExt};
+
+use super::*;
+
+use crate::traits::events::{EventPublisher, EventSubscriber};
+
+#[async_trait]
+impl EventPublisher for Component {
+    fn subject(&self) -> String {
+        format!("namespace.{}.component.{}", self.namespace.name, self.name)
+    }
+
+    async fn publish(
+        &self,
+        event_name: impl AsRef<str> + Send + Sync,
+        event: &(impl Serialize + Send + Sync),
+    ) -> Result<()> {
+        let bytes = serde_json::to_vec(event)?;
+        self.publish_bytes(event_name, bytes).await
+    }
+
+    async fn publish_bytes(
+        &self,
+        event_name: impl AsRef<str> + Send + Sync,
+        bytes: Vec<u8>,
+    ) -> Result<()> {
+        let subject = format!("{}.{}", self.subject(), event_name.as_ref());
+        Ok(self
+            .drt()
+            .nats_client()
+            .client()
+            .publish(subject, bytes.into())
+            .await?)
+    }
+}
+
+#[async_trait]
+impl EventSubscriber for Component {
+    async fn subscribe(
+        &self,
+        event_name: impl AsRef<str> + Send + Sync,
+    ) -> Result<async_nats::Subscriber> {
+        let subject = format!("{}.{}", self.subject(), event_name.as_ref());
+        Ok(self.drt().nats_client().client().subscribe(subject).await?)
+    }
+
+    async fn subscribe_with_type<T: for<'de> Deserialize<'de> + Send + 'static>(
+        &self,
+        event_name: impl AsRef<str> + Send + Sync,
+    ) -> Result<impl Stream<Item = Result<T>> + Send> {
+        let subscriber = self.subscribe(event_name).await?;
+
+        // Transform the subscriber into a stream of deserialized events
+        let stream = subscriber.map(move |msg| {
+            serde_json::from_slice::<T>(&msg.payload)
+                .with_context(|| format!("Failed to deserialize event payload: {:?}", msg.payload))
+        });
+
+        Ok(stream)
+    }
+}
+
+#[cfg(feature = "integration")]
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // todo - make a distributed runtime fixture
+    // todo - two options - fully mocked or integration test
+    #[tokio::test]
+    async fn test_publish() {
+        let rt = Runtime::from_current().unwrap();
+        let dtr = DistributedRuntime::from_settings(rt.clone()).await.unwrap();
+        let ns = dtr.namespace("test".to_string()).unwrap();
+        let cp = ns.component("component".to_string()).unwrap();
+        cp.publish("test", &"test".to_string()).await.unwrap();
+        rt.shutdown();
+    }
+
+    #[tokio::test]
+    async fn test_subscribe() {
+        let rt = Runtime::from_current().unwrap();
+        let dtr = DistributedRuntime::from_settings(rt.clone()).await.unwrap();
+        let ns = dtr.namespace("test".to_string()).unwrap();
+        let cp = ns.component("component".to_string()).unwrap();
+
+        // Create a subscriber
+        let mut subscriber = ns.subscribe("test").await.unwrap();
+
+        // Publish a message
+        cp.publish("test", &"test_message".to_string())
+            .await
+            .unwrap();
+
+        // Receive the message
+        if let Some(msg) = subscriber.next().await {
+            let received = String::from_utf8(msg.payload.to_vec()).unwrap();
+            assert_eq!(received, "\"test_message\"");
+        }
+
+        rt.shutdown();
+    }
+}

lib/runtime/src/component/endpoint.rs

@@ -113,7 +113,7 @@ impl EndpointConfigBuilder {
         let info = ComponentEndpointInfo {
             component: endpoint.component.name.clone(),
             endpoint: endpoint.name.clone(),
-            namespace: endpoint.component.namespace.clone(),
+            namespace: endpoint.component.namespace.name.clone(),
             lease_id: lease.id(),
             transport: TransportType::NatsTcp(endpoint.subject_to(lease.id())),
         };

lib/runtime/src/component/namespace.rs

@@ -101,7 +101,7 @@ mod tests {
         let ns = dtr.namespace("test".to_string()).unwrap();
 
         // Create a subscriber
-        let subscriber = ns.subscribe("test").await.unwrap();
+        let mut subscriber = ns.subscribe("test").await.unwrap();
 
         // Publish a message
         ns.publish("test", &"test_message".to_string())

lib/runtime/src/service.rs

@@ -73,12 +73,28 @@ impl EndpointInfo {
         i64::from_str_radix(id, 16).map_err(|e| error!("Invalid id format: {}", e))
     }
 }
+
+// TODO: This is _really_ close to the async_nats::service::Stats object,
+// but it's missing a few fields like "name", so use a temporary struct
+// for easy deserialization. Ideally, this type already exists or can
+// be exposed in the library somewhere.
+/// Stats structure returned from NATS service API
 #[derive(Debug, Clone, Serialize, Deserialize, Dissolve)]
-pub struct Metrics(pub serde_json::Value);
+pub struct Metrics {
+    // Standard NATS Service API fields
+    pub average_processing_time: f64,
+    pub last_error: String,
+    pub num_errors: u64,
+    pub num_requests: u64,
+    pub processing_time: u64,
+    pub queue_group: String,
+    // Field containing custom stats handler data
+    pub data: serde_json::Value,
+}
 
 impl Metrics {
     pub fn decode<T: for<'de> Deserialize<'de>>(self) -> Result<T> {
-        serde_json::from_value(self.0).map_err(Into::into)
+        serde_json::from_value(self.data).map_err(Into::into)
     }
 }
 
@@ -153,12 +169,28 @@ mod tests {
                     EndpointInfo {
                         name: "endpoint1".to_string(),
                         subject: "subject1".to_string(),
-                        data: Some(Metrics(serde_json::json!({"key": "value1"}))),
+                        data: Some(Metrics {
+                            average_processing_time: 0.1,
+                            last_error: "none".to_string(),
+                            num_errors: 0,
+                            num_requests: 10,
+                            processing_time: 100,
+                            queue_group: "group1".to_string(),
+                            data: serde_json::json!({"key": "value1"}),
+                        }),
                     },
                     EndpointInfo {
                         name: "endpoint2-foo".to_string(),
                         subject: "subject2".to_string(),
-                        data: Some(Metrics(serde_json::json!({"key": "value1"}))),
+                        data: Some(Metrics {
+                            average_processing_time: 0.1,
+                            last_error: "none".to_string(),
+                            num_errors: 0,
+                            num_requests: 10,
+                            processing_time: 100,
+                            queue_group: "group1".to_string(),
+                            data: serde_json::json!({"key": "value1"}),
+                        }),
                     },
                 ],
             },
@@ -171,12 +203,28 @@ mod tests {
                     EndpointInfo {
                         name: "endpoint1".to_string(),
                         subject: "subject1".to_string(),
-                        data: Some(Metrics(serde_json::json!({"key": "value1"}))),
+                        data: Some(Metrics {
+                            average_processing_time: 0.1,
+                            last_error: "none".to_string(),
+                            num_errors: 0,
+                            num_requests: 10,
+                            processing_time: 100,
+                            queue_group: "group1".to_string(),
+                            data: serde_json::json!({"key": "value1"}),
+                        }),
                     },
                     EndpointInfo {
                         name: "endpoint2-bar".to_string(),
                         subject: "subject2".to_string(),
-                        data: Some(Metrics(serde_json::json!({"key": "value2"}))),
+                        data: Some(Metrics {
+                            average_processing_time: 0.1,
+                            last_error: "none".to_string(),
+                            num_errors: 0,
+                            num_requests: 10,
+                            processing_time: 100,
+                            queue_group: "group1".to_string(),
+                            data: serde_json::json!({"key": "value2"}),
+                        }),
                     },
                 ],
             },

lib/runtime/src/transports/nats.rs

@@ -108,35 +108,6 @@ impl Client {
         Ok(subscription)
     }
 
-    // todo - deprecate - move to service subscriber
-    pub async fn get_endpoints(
-        &self,
-        service_name: &str,
-        timeout: time::Duration,
-    ) -> Result<Vec<Bytes>, anyhow::Error> {
-        let subject = format!("$SRV.STATS.{}", service_name);
-        let reply_subject = format!("_INBOX.{}", nuid::next());
-        let mut subscription = self.client.subscribe(reply_subject.clone()).await?;
-
-        let deadline = tokio::time::Instant::now() + timeout;
-
-        // Publish the request with the reply-to subject
-        self.client
-            .publish_with_reply(subject, reply_subject, "".into())
-            .await?;
-
-        // Set a timeout to gather responses
-        let mut responses = Vec::new();
-        // let mut response_stream = subscription.take_while(|_| futures::future::ready(true));
-
-        while let Ok(Some(message)) = time::timeout_at(deadline, subscription.next()).await {
-            // log::debug!("get endpoint received message before timeout: {:?}", message);
-            responses.push(message.payload);
-        }
-
-        Ok(responses)
-    }
-
     // /// create a new stream
     // async fn get_or_create_work_queue_stream(
     //     &self,
@@ -272,35 +243,6 @@ impl Client {
     //     Ok(())
     // }
 
-    // pub async fn get_endpoints(
-    //     &self,
-    //     service_name: &str,
-    //     timeout: Duration,
-    // ) -> Result<Vec<Bytes>, anyhow::Error> {
-    //     let subject = format!("$SRV.STATS.{}", service_name);
-    //     let reply_subject = format!("_INBOX.{}", nuid::next());
-    //     let mut subscription = self.client.subscribe(reply_subject.clone()).await?;
-
-    //     // Publish the request with the reply-to subject
-    //     self.client
-    //         .publish_with_reply(subject, reply_subject, "".into())
-    //         .await?;
-
-    //     // Set a timeout to gather responses
-    //     let mut responses = Vec::new();
-    //     // let mut response_stream = subscription.take_while(|_| futures::future::ready(true));
-
-    //     let start = time::Instant::now();
-    //     while let Ok(Some(message)) = time::timeout(timeout, subscription.next()).await {
-    //         responses.push(message.payload);
-    //         if start.elapsed() > timeout {
-    //             break;
-    //         }
-    //     }
-
-    //     Ok(responses)
-    // }
-
     // pub fn frontend_client(&self, request_id: String) -> SpecializedClient {
     //     SpecializedClient::new(self.client.clone(), ClientKind::Frontend, request_id)
     // }
@@ -691,35 +633,6 @@ mod tests {
 //     assert_eq!(initial_work_queue_count, work_queue_count);
 // }
 
-// pub async fn get_endpoints(
-//     &self,
-//     service_name: &str,
-//     timeout: Duration,
-// ) -> Result<Vec<Bytes>, anyhow::Error> {
-//     let subject = format!("$SRV.STATS.{}", service_name);
-//     let reply_subject = format!("_INBOX.{}", nuid::next());
-//     let mut subscription = self.client.subscribe(reply_subject.clone()).await?;
-
-//     // Publish the request with the reply-to subject
-//     self.client
-//         .publish_with_reply(subject, reply_subject, "".into())
-//         .await?;
-
-//     // Set a timeout to gather responses
-//     let mut responses = Vec::new();
-//     // let mut response_stream = subscription.take_while(|_| futures::future::ready(true));
-
-//     let start = time::Instant::now();
-//     while let Ok(Some(message)) = time::timeout(timeout, subscription.next()).await {
-//         responses.push(message.payload);
-//         if start.elapsed() > timeout {
-//             break;
-//         }
-//     }
-
-//     Ok(responses)
-// }
-
 // async fn connect(config: Arc<Config>) -> Result<NatsClient> {
 //     let client = ClientOptions::builder()
 //         .server(config.nats_address.clone())
@@ -852,35 +765,6 @@ mod tests {
 //     pub fn service_builder(&self) -> NatsServiceBuilder {
 //         self.client.service_builder()
 //     }
-
-//     pub async fn get_endpoints(
-//         &self,
-//         service_name: &str,
-//         timeout: Duration,
-//     ) -> Result<Vec<Bytes>, anyhow::Error> {
-//         let subject = format!("$SRV.STATS.{}", service_name);
-//         let reply_subject = format!("_INBOX.{}", nuid::next());
-//         let mut subscription = self.client.subscribe(reply_subject.clone()).await?;
-
-//         // Publish the request with the reply-to subject
-//         self.client
-//             .publish_with_reply(subject, reply_subject, "".into())
-//             .await?;
-
-//         // Set a timeout to gather responses
-//         let mut responses = Vec::new();
-//         // let mut response_stream = subscription.take_while(|_| futures::future::ready(true));
-
-//         let start = tokio::time::Instant::now();
-//         while let Ok(Some(message)) = tokio::time::timeout(timeout, subscription.next()).await {
-//             responses.push(message.payload);
-//             if start.elapsed() > timeout {
-//                 break;
-//             }
-//         }
-
-//         Ok(responses)
-//     }
 // }
 
 // #[derive(Debug, Clone, Serialize, Deserialize)]