chore(coprocessor): add more metrics for sns-worker (#1369)

Also, move code around to put metrics in one place.

Author: dartdart26

charts/coprocessor/Chart.yaml

@@ -1,6 +1,6 @@
 name: coprocessor
 description: A helm chart to distribute and deploy Zama fhevm Co-Processor services
-version: 0.7.2
+version: 0.7.3
 apiVersion: v2
 keywords:
   - fhevm

charts/coprocessor/values.yaml

@@ -601,10 +601,12 @@ snsWorker:
     - --cleanup-interval=120s
     - --liveness-threshold=70s
     - --lifo=false
-    - --enable-compssion=true
+    - --enable-compression=true
     - --schedule-policy=rayon_parallel
     - --service-name=sns-worker
     - --log-level=INFO
+    # Only enable `gauge-update-interval-secs` for some of the workers to reduce DB load and not have duplicate metrics data for no reason.
+    # --gauge-update-interval-secs=10
     # - --pg-auto-explain-with-min-duration="30s"
     # - --keys-file-path
 

coprocessor/fhevm-engine/sns-worker/src/aws_upload.rs

@@ -1,3 +1,4 @@
+use crate::metrics::{AWS_UPLOAD_FAILURE_COUNTER, AWS_UPLOAD_SUCCESS_COUNTER};
 use crate::{
     BigCiphertext, Ciphertext128Format, Config, ExecutionError, HandleItem, S3Config, UploadJob,
 };
@@ -166,16 +167,19 @@ async fn run_uploader_loop(
                 let h = tokio::spawn(async move {
                     let s = item.otel.child_span("upload_s3");
                     match upload_ciphertexts(trx, item, &client, &conf).instrument(error_span!("upload_s3")).await {
-                        Ok(()) => telemetry::end_span(s),
+                        Ok(()) => {
+                            telemetry::end_span(s);
+                            AWS_UPLOAD_SUCCESS_COUNTER.inc();
+                        }
                         Err(err) => {
                             if let ExecutionError::S3TransientError(_) = err {
                                 ready_flag.store(false, Ordering::Release);
                                 info!(error = %err, "S3 setup is not ready, due to transient error");
                             } else {
                                 error!(error = %err, "Failed to upload ciphertexts");
                             }
-
                             telemetry::end_span_with_err(s, err.to_string());
+                            AWS_UPLOAD_FAILURE_COUNTER.inc();
                         }
                     }
                     drop(permit);

coprocessor/fhevm-engine/sns-worker/src/bin/sns_worker.rs

@@ -1,4 +1,4 @@
-use sns_worker::{Config, DBConfig, HealthCheckConfig, S3Config, S3RetryPolicy};
+use sns_worker::{Config, DBConfig, HealthCheckConfig, S3Config, S3RetryPolicy, SNSMetricsConfig};
 
 use tokio::signal::unix;
 use tokio_util::sync::CancellationToken;
@@ -21,7 +21,10 @@ fn construct_config() -> Config {
     Config {
         tenant_api_key: args.tenant_api_key,
         service_name: args.service_name,
-        metrics_addr: args.metrics_addr,
+        metrics: SNSMetricsConfig {
+            addr: args.metrics_addr,
+            gauge_update_interval_secs: args.gauge_update_interval_secs,
+        },
         db: DBConfig {
             url: db_url,
             listen_channels: args.pg_listen_channels,

coprocessor/fhevm-engine/sns-worker/src/bin/utils/daemon_cli.rs

@@ -4,7 +4,8 @@ use clap::{command, Parser};
 use fhevm_engine_common::telemetry::MetricsConfig;
 use fhevm_engine_common::utils::DatabaseURL;
 use humantime::parse_duration;
-use sns_worker::{SchedulePolicy, SNS_LATENCY_OP_HISTOGRAM_CONF};
+use sns_worker::metrics::SNS_LATENCY_OP_HISTOGRAM_CONF;
+use sns_worker::SchedulePolicy;
 use tracing::Level;
 
 #[derive(Parser, Debug, Clone)]
@@ -128,6 +129,9 @@ pub struct Args {
     /// Prometheus metrics: coprocessor_sns_op_latency_seconds
     #[arg(long, default_value = "0.1:10.0:0.1", value_parser = clap::value_parser!(MetricsConfig))]
     pub metric_sns_op_latency: MetricsConfig,
+
+    #[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
+    pub gauge_update_interval_secs: Option<u32>,
 }
 
 pub fn parse_args() -> Args {

coprocessor/fhevm-engine/sns-worker/src/executor.rs

@@ -1,5 +1,8 @@
 use crate::aws_upload::check_is_ready;
 use crate::keyset::fetch_keyset;
+use crate::metrics::SNS_LATENCY_OP_HISTOGRAM;
+use crate::metrics::TASK_EXECUTE_FAILURE_COUNTER;
+use crate::metrics::TASK_EXECUTE_SUCCESS_COUNTER;
 use crate::squash_noise::SquashNoiseCiphertext;
 use crate::BigCiphertext;
 use crate::Ciphertext128Format;
@@ -8,7 +11,6 @@ use crate::InternalEvents;
 use crate::KeySet;
 use crate::SchedulePolicy;
 use crate::UploadJob;
-use crate::SNS_LATENCY_OP_HISTOGRAM;
 use crate::{Config, ExecutionError};
 use aws_sdk_s3::Client;
 use fhevm_engine_common::healthz_server::{HealthCheckService, HealthStatus, Version};
@@ -238,7 +240,15 @@ pub(crate) async fn run_loop(
             continue;
         };
 
-        let maybe_remaining = fetch_and_execute_sns_tasks(&pool, &tx, keys, &conf, &token).await?;
+        let (maybe_remaining, _tasks_processed) =
+            fetch_and_execute_sns_tasks(&pool, &tx, keys, &conf, &token)
+                .await
+                .inspect(|(_, tasks_processed)| {
+                    TASK_EXECUTE_SUCCESS_COUNTER.inc_by(*tasks_processed as u64);
+                })
+                .inspect_err(|_| {
+                    TASK_EXECUTE_FAILURE_COUNTER.inc();
+                })?;
         if maybe_remaining {
             if token.is_cancelled() {
                 return Ok(());
@@ -317,13 +327,14 @@ pub async fn garbage_collect(pool: &PgPool, limit: u32) -> Result<(), ExecutionE
 }
 
 /// Fetch and process SnS tasks from the database.
+/// Returns (maybe_remaining, number_of_tasks_processed) on success.
 async fn fetch_and_execute_sns_tasks(
     pool: &PgPool,
     tx: &Sender<UploadJob>,
     keys: &KeySet,
     conf: &Config,
     token: &CancellationToken,
-) -> Result<bool, ExecutionError> {
+) -> Result<(bool, usize), ExecutionError> {
     let mut db_txn = match pool.begin().await {
         Ok(txn) => txn,
         Err(err) => {
@@ -341,8 +352,10 @@ async fn fetch_and_execute_sns_tasks(
     let trx = &mut db_txn;
 
     let mut maybe_remaining = false;
+    let tasks_processed;
     if let Some(mut tasks) = query_sns_tasks(trx, conf.db.batch_limit, order).await? {
         maybe_remaining = conf.db.batch_limit as usize == tasks.len();
+        tasks_processed = tasks.len();
 
         let t = telemetry::tracer("batch_execution", &None);
         t.set_attribute("count", tasks.len().to_string());
@@ -375,10 +388,11 @@ async fn fetch_and_execute_sns_tasks(
             }
         }
     } else {
+        tasks_processed = 0;
         db_txn.rollback().await?;
     }
 
-    Ok(maybe_remaining)
+    Ok((maybe_remaining, tasks_processed))
 }
 
 /// Queries the database for a fixed number of tasks.

coprocessor/fhevm-engine/sns-worker/src/lib.rs

@@ -3,6 +3,8 @@ mod executor;
 mod keyset;
 mod squash_noise;
 
+pub mod metrics;
+
 #[cfg(test)]
 mod tests;
 
@@ -21,7 +23,6 @@ use fhevm_engine_common::{
     metrics_server,
     pg_pool::{PostgresPoolManager, ServiceError},
     telemetry::{self, OtelTracer},
-    telemetry::{register_histogram, MetricsConfig},
     types::FhevmError,
     utils::{to_hex, DatabaseURL},
 };
@@ -43,11 +44,9 @@ use tracing::{error, info, Level};
 use crate::{
     aws_upload::{check_is_ready, spawn_resubmit_task, spawn_uploader},
     executor::SwitchNSquashService,
+    metrics::spawn_gauge_update_routine,
 };
 
-use prometheus::Histogram;
-use std::sync::{LazyLock, OnceLock};
-
 pub const UPLOAD_QUEUE_SIZE: usize = 20;
 pub const SAFE_SER_LIMIT: u64 = 1024 * 1024 * 66;
 pub type InternalEvents = Option<tokio::sync::mpsc::Sender<&'static str>>;
@@ -75,6 +74,12 @@ pub struct DBConfig {
     pub lifo: bool,
 }
 
+#[derive(Clone, Default, Debug)]
+pub struct SNSMetricsConfig {
+    pub addr: Option<String>,
+    pub gauge_update_interval_secs: Option<u32>,
+}
+
 #[derive(Clone, Default, Debug)]
 pub struct S3Config {
     pub bucket_ct128: String,
@@ -106,7 +111,7 @@ pub struct Config {
     pub s3: S3Config,
     pub log_level: Level,
     pub health_checks: HealthCheckConfig,
-    pub metrics_addr: Option<String>,
+    pub metrics: SNSMetricsConfig,
     pub enable_compression: bool,
     pub schedule_policy: SchedulePolicy,
     pub pg_auto_explain_with_min_duration: Option<Duration>,
@@ -390,9 +395,6 @@ pub async fn run_computation_loop(
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let port = conf.health_checks.port;
 
-    // Start metrics server
-    metrics_server::spawn(conf.metrics_addr.clone(), token.child_token());
-
     let service = Arc::new(
         SwitchNSquashService::create(
             pool_mngr,
@@ -532,6 +534,21 @@ pub async fn run_all(
 
     let pg_mngr = pool_mngr.clone();
 
+    // Start metrics server
+    metrics_server::spawn(conf.metrics.addr.clone(), token.child_token());
+
+    // Start gauge update routine.
+    if let Some(interval_secs) = conf.metrics.gauge_update_interval_secs {
+        info!(
+            interval_secs = interval_secs,
+            "Starting gauge update routine"
+        );
+        spawn_gauge_update_routine(
+            Duration::from_secs(interval_secs.into()),
+            pg_mngr.pool().clone(),
+        );
+    }
+
     // Spawns a task to handle S3 uploads
     spawn(async move {
         if let Err(err) = run_uploader_loop(&pg_mngr, &conf, jobs_rx, tx, s3, is_ready).await {
@@ -552,12 +569,3 @@ pub async fn run_all(
 
     Ok(())
 }
-
-pub static SNS_LATENCY_OP_HISTOGRAM_CONF: OnceLock<MetricsConfig> = OnceLock::new();
-pub static SNS_LATENCY_OP_HISTOGRAM: LazyLock<Histogram> = LazyLock::new(|| {
-    register_histogram(
-        SNS_LATENCY_OP_HISTOGRAM_CONF.get(),
-        "coprocessor_sns_op_latency_seconds",
-        "Squash_noise computation latencies in seconds",
-    )
-});

coprocessor/fhevm-engine/sns-worker/src/metrics.rs

@@ -0,0 +1,104 @@
+use std::sync::{LazyLock, OnceLock};
+
+use fhevm_engine_common::telemetry::{register_histogram, MetricsConfig};
+use prometheus::{register_int_counter, IntCounter};
+use prometheus::{register_int_gauge, Histogram, IntGauge};
+use sqlx::PgPool;
+use tokio::task::JoinHandle;
+use tokio::time::sleep;
+use tracing::{error, info};
+
+pub static SNS_LATENCY_OP_HISTOGRAM_CONF: OnceLock<MetricsConfig> = OnceLock::new();
+pub(crate) static SNS_LATENCY_OP_HISTOGRAM: LazyLock<Histogram> = LazyLock::new(|| {
+    register_histogram(
+        SNS_LATENCY_OP_HISTOGRAM_CONF.get(),
+        "coprocessor_sns_op_latency_seconds",
+        "Squash_noise computation latencies in seconds",
+    )
+});
+
+pub(crate) static TASK_EXECUTE_SUCCESS_COUNTER: LazyLock<IntCounter> = LazyLock::new(|| {
+    register_int_counter!(
+        "coprocessor_sns_worker_task_execute_success_counter",
+        "Number of successful task execute operations in sns-worker (including persistence to DB)"
+    )
+    .unwrap()
+});
+
+pub(crate) static TASK_EXECUTE_FAILURE_COUNTER: LazyLock<IntCounter> = LazyLock::new(|| {
+    register_int_counter!(
+        "coprocessor_sns_worker_task_execute_failure_counter",
+        "Number of failed task execute operations in sns-worker (including persistence to DB)"
+    )
+    .unwrap()
+});
+
+pub(crate) static AWS_UPLOAD_SUCCESS_COUNTER: LazyLock<IntCounter> = LazyLock::new(|| {
+    register_int_counter!(
+        "coprocessor_sns_worker_aws_upload_success_counter",
+        "Number of successful AWS uploads in sns-worker"
+    )
+    .unwrap()
+});
+
+pub(crate) static AWS_UPLOAD_FAILURE_COUNTER: LazyLock<IntCounter> = LazyLock::new(|| {
+    register_int_counter!(
+        "coprocessor_sns_worker_aws_upload_failure_counter",
+        "Number of failed AWS uploads in sns-worker"
+    )
+    .unwrap()
+});
+
+pub(crate) static UNCOMPLETE_TASKS: LazyLock<IntGauge> = LazyLock::new(|| {
+    register_int_gauge!(
+        "coprocessor_sns_worker_uncomplete_tasks_gauge",
+        "Number of uncomplete tasks in sns-worker"
+    )
+    .unwrap()
+});
+
+pub(crate) static UNCOMPLETE_AWS_UPLOADS: LazyLock<IntGauge> = LazyLock::new(|| {
+    register_int_gauge!(
+        "coprocessor_sns_worker_uncomplete_aws_uploads_gauge",
+        "Number of uncomplete AWS uploads in sns-worker"
+    )
+    .unwrap()
+});
+
+pub fn spawn_gauge_update_routine(period: std::time::Duration, db_pool: PgPool) -> JoinHandle<()> {
+    tokio::spawn(async move {
+        loop {
+            match sqlx::query_scalar(
+                "SELECT COUNT(*) FROM pbs_computations WHERE is_completed = FALSE",
+            )
+            .fetch_one(&db_pool)
+            .await
+            {
+                Ok(count) => {
+                    info!(uncomplete_tasks = %count, "Fetched uncomplete tasks count");
+                    UNCOMPLETE_TASKS.set(count);
+                }
+                Err(e) => {
+                    error!(error = %e, "Failed to fetch uncomplete tasks count");
+                }
+            }
+
+            match sqlx::query_scalar(
+                "SELECT COUNT(*) FROM ciphertext_digest WHERE (ciphertext128 IS NULL OR ciphertext IS NULL)",
+            )
+            .fetch_one(&db_pool)
+            .await
+            {
+                Ok(count) => {
+                    info!(uncomplete_aws_uploads = %count, "Fetched uncomplete AWS uploads count");
+                    UNCOMPLETE_AWS_UPLOADS.set(count);
+                }
+                Err(e) => {
+                    error!(error = %e, "Failed to fetch uncomplete AWS uploads count");
+                }
+            }
+
+            sleep(period).await;
+        }
+    })
+}

coprocessor/fhevm-engine/sns-worker/src/tests/mod.rs

@@ -755,6 +755,6 @@ fn build_test_config(url: DatabaseURL, enable_compression: bool) -> Config {
         enable_compression,
         schedule_policy,
         pg_auto_explain_with_min_duration: Some(Duration::from_secs(1)),
-        metrics_addr: None,
+        metrics: Default::default(),
     }
 }