From cfa9558a332b93db3f55f5cb08e8bbdb8bf3c568 Mon Sep 17 00:00:00 2001
From: Madhukar Mishra <madhukar93@gmail.com>
Date: Mon, 8 Jan 2024 19:13:37 +0530
Subject: [PATCH] add a metric to try otel

---
 crates/service/job/inference_job/Cargo.toml    |  2 +-
 .../src/job/job_loop/main_loop.rs              | 18 ++++++++++++++++++
 crates/service/job/inference_job/src/main.rs   |  3 +++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml
index 313e392991..7e714de4ec 100644
--- a/crates/service/job/inference_job/Cargo.toml
+++ b/crates/service/job/inference_job/Cargo.toml
@@ -85,4 +85,4 @@ rand = { version = "0.8.5", features = [] }
 jsonpath_lib = "0.3.0"
 walkdir = "2.4.0"
 reqwest = "0.11.22"
-opentelemetry = "0.21.0"
+opentelemetry = {version = "0.21.0", features = ["metrics"]}
diff --git a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
index a1e11adb6b..c59cd19beb 100644
--- a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
+++ b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs
@@ -9,13 +9,17 @@ use jobs_common::noop_logger::NoOpLogger;
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::{AvailableInferenceJob, list_available_generic_inference_jobs, ListAvailableGenericInferenceJobArgs};
 use mysql_queries::queries::generic_inference::job::mark_generic_inference_job_completely_failed::mark_generic_inference_job_completely_failed;
 use mysql_queries::queries::generic_inference::job::mark_generic_inference_job_failure::mark_generic_inference_job_failure;
+use opentelemetry::metrics::AsyncInstrument;
 
+use crate::OTEL_METER_NAME;
 use crate::job::job_loop::clear_full_filesystem::clear_full_filesystem;
 use crate::job::job_loop::process_single_job::process_single_job;
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 use crate::job::job_loop::process_single_job_success_case::ProcessSingleJobSuccessCase;
 use crate::job_dependencies::JobDependencies;
 
+use opentelemetry::{global as otel, KeyValue as OtelAttribute};
+
 // Job runner timeouts (guards MySQL)
 const START_TIMEOUT_MILLIS : u64 = 500;
 const INCREASE_TIMEOUT_MILLIS : u64 = 1000;
@@ -118,7 +122,12 @@ pub async fn main_loop(job_dependencies: JobDependencies) {
 
 async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec<AvailableInferenceJob>) -> AnyhowResult<()> {
   let job_count = jobs.len();
+  let meter = otel::meter(OTEL_METER_NAME);
+  // TODO: total_duration as histogram for min/max/avg/percentiles
+  let job_duration = meter.i64_observable_gauge("job_duration").with_unit(metrics::Unit::new("ms")).init();
+
   for (i, job) in jobs.into_iter().enumerate() {
+    let start_time = Instant::now();
     let result = process_single_job(job_dependencies, &job).await;
 
     match result {
@@ -145,6 +154,15 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec<Availab
         let _r = handle_error(&job_dependencies, &job, err).await?;
       }
     }
+    let job_duration = Instant::now().duration_since(start_time);;
+    job_duration.observe(job_duration.as_millis(),
+      &[
+        // TODO: need job_ prefix?
+        OtelAttribute::new("job_id", job.id),
+        OtelAttribute::new("job_inference_category", job.inference_category),
+        OtelAttribute::new("job_inference_token", job.inference_job_token),
+      ]
+    )
   }
 
   Ok(())
diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs
index d783427f3c..7132d28e90 100644
--- a/crates/service/job/inference_job/src/main.rs
+++ b/crates/service/job/inference_job/src/main.rs
@@ -22,6 +22,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use log::{info, warn};
+use opentelemetry::metrics;
 use r2d2_redis::r2d2;
 use r2d2_redis::RedisConnectionManager;
 use sqlx::mysql::MySqlPoolOptions;
@@ -72,6 +73,8 @@ const ENV_PUBLIC_BUCKET_NAME : &str = "PUBLIC_BUCKET_NAME";
 // HTTP sidecar
 const ENV_TTS_INFERENCE_SIDECAR_HOSTNAME: &str = "TTS_INFERENCE_SIDECAR_HOSTNAME";
 
+const OTEL_METER_NAME: &str = "inference-job";
+
 //#[tokio::main]
 #[actix_web::main]
 async fn main() -> AnyhowResult<()> {