|
| 1 | +// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | +// |
| 4 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +// you may not use this file except in compliance with the License. |
| 6 | +// You may obtain a copy of the License at |
| 7 | +// |
| 8 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +// |
| 10 | +// Unless required by applicable law or agreed to in writing, software |
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +// See the License for the specific language governing permissions and |
| 14 | +// limitations under the License. |
| 15 | + |
| 16 | +use system_metrics::{MyStats, DEFAULT_NAMESPACE}; |
| 17 | + |
| 18 | +use dynamo_runtime::{ |
| 19 | + logging, |
| 20 | + metrics::MetricsRegistry, |
| 21 | + pipeline::{ |
| 22 | + async_trait, network::Ingress, AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, |
| 23 | + ResponseStream, SingleIn, |
| 24 | + }, |
| 25 | + protocols::annotated::Annotated, |
| 26 | + stream, DistributedRuntime, Result, Runtime, Worker, |
| 27 | +}; |
| 28 | + |
| 29 | +use prometheus::{Counter, Histogram}; |
| 30 | +use std::sync::Arc; |
| 31 | + |
| 32 | +/// Service metrics struct using the metric classes from metrics.rs |
| 33 | +pub struct MySystemStatsMetrics { |
| 34 | + pub request_counter: Arc<Counter>, |
| 35 | + pub request_duration: Arc<Histogram>, |
| 36 | +} |
| 37 | + |
| 38 | +impl MySystemStatsMetrics { |
| 39 | + /// Create a new ServiceMetrics instance using the metric backend |
| 40 | + pub fn new<R: MetricsRegistry>( |
| 41 | + metrics_registry: Arc<R>, |
| 42 | + ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> { |
| 43 | + let request_counter = metrics_registry.create_counter( |
| 44 | + "service_requests_total", |
| 45 | + "Total number of requests processed", |
| 46 | + &[("service", "backend")], |
| 47 | + )?; |
| 48 | + let request_duration = metrics_registry.create_histogram( |
| 49 | + "service_request_duration_seconds", |
| 50 | + "Time spent processing requests", |
| 51 | + &[("service", "backend")], |
| 52 | + None, |
| 53 | + )?; |
| 54 | + Ok(Self { |
| 55 | + request_counter, |
| 56 | + request_duration, |
| 57 | + }) |
| 58 | + } |
| 59 | +} |
| 60 | + |
| 61 | +fn main() -> Result<()> { |
| 62 | + logging::init(); |
| 63 | + let worker = Worker::from_settings()?; |
| 64 | + worker.execute(app) |
| 65 | +} |
| 66 | + |
| 67 | +async fn app(runtime: Runtime) -> Result<()> { |
| 68 | + let distributed = DistributedRuntime::from_settings(runtime.clone()).await?; |
| 69 | + backend(distributed).await |
| 70 | +} |
| 71 | + |
| 72 | +struct RequestHandler { |
| 73 | + metrics: Arc<MySystemStatsMetrics>, |
| 74 | +} |
| 75 | + |
| 76 | +impl RequestHandler { |
| 77 | + fn new(metrics: Arc<MySystemStatsMetrics>) -> Arc<Self> { |
| 78 | + Arc::new(Self { metrics }) |
| 79 | + } |
| 80 | +} |
| 81 | + |
| 82 | +#[async_trait] |
| 83 | +impl AsyncEngine<SingleIn<String>, ManyOut<Annotated<String>>, Error> for RequestHandler { |
| 84 | + async fn generate(&self, input: SingleIn<String>) -> Result<ManyOut<Annotated<String>>> { |
| 85 | + let start_time = std::time::Instant::now(); |
| 86 | + |
| 87 | + // Record request start |
| 88 | + self.metrics.request_counter.inc(); |
| 89 | + |
| 90 | + let (data, ctx) = input.into_parts(); |
| 91 | + |
| 92 | + let chars = data |
| 93 | + .chars() |
| 94 | + .map(|c| Annotated::from_data(c.to_string())) |
| 95 | + .collect::<Vec<_>>(); |
| 96 | + |
| 97 | + let stream = stream::iter(chars); |
| 98 | + |
| 99 | + // Record request duration |
| 100 | + let duration = start_time.elapsed(); |
| 101 | + self.metrics |
| 102 | + .request_duration |
| 103 | + .observe(duration.as_secs_f64()); |
| 104 | + |
| 105 | + Ok(ResponseStream::new(Box::pin(stream), ctx.context())) |
| 106 | + } |
| 107 | +} |
| 108 | + |
| 109 | +async fn backend(drt: DistributedRuntime) -> Result<()> { |
| 110 | + let endpoint = drt |
| 111 | + .namespace(DEFAULT_NAMESPACE)? |
| 112 | + .component("component")? |
| 113 | + .service_builder() |
| 114 | + .create() |
| 115 | + .await? |
| 116 | + .endpoint("endpoint"); |
| 117 | + |
| 118 | + // make the ingress discoverable via a component service |
| 119 | + // we must first create a service, then we can attach one more more endpoints |
| 120 | + // attach an ingress to an engine, with the RequestHandler using the metrics struct |
| 121 | + let endpoint_metrics = Arc::new( |
| 122 | + MySystemStatsMetrics::new(Arc::new(endpoint.clone())) |
| 123 | + .map_err(|e| Error::msg(e.to_string()))?, |
| 124 | + ); |
| 125 | + let ingress = Ingress::for_engine(RequestHandler::new(endpoint_metrics.clone()))?; |
| 126 | + |
| 127 | + endpoint |
| 128 | + .endpoint_builder() |
| 129 | + .stats_handler(|_stats| { |
| 130 | + println!("Stats handler called with stats: {:?}", _stats); |
| 131 | + let stats = MyStats { val: 10 }; |
| 132 | + serde_json::to_value(stats).unwrap() |
| 133 | + }) |
| 134 | + .handler(ingress) |
| 135 | + .start() |
| 136 | + .await?; |
| 137 | + |
| 138 | + Ok(()) |
| 139 | +} |
0 commit comments