Merge pull request #302 from storytold/feb-16-mithoo

slightly tweakier instrumentation

crates/service/job/inference_job/src/main.rs

@@ -58,6 +58,7 @@ use crate::job_dependencies::{BucketDependencies, ClientDependencies, DatabaseDe
 use crate::job_specific_dependencies::JobSpecificDependencies;
 use crate::util::instrumentation::JobInstruments;
 use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory;
+use crate::util::instrumentation::{init_otel_metrics_pipeline, JobInstrumentLabels };
 use crate::util::scoped_execution::ScopedExecution;
 use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
 
@@ -81,40 +82,23 @@ const ENV_TTS_INFERENCE_SIDECAR_HOSTNAME: &str = "TTS_INFERENCE_SIDECAR_HOSTNAME
 
 const OTEL_METER_NAME: &str = "inference-job";
 
-// TODO@madhukar93: labels for model etc
-fn init_otel_metrics_pipeline() -> Result<(), opentelemetry::metrics::MetricsError>  {
- let provider = opentelemetry_otlp::new_pipeline()
-     .metrics(opentelemetry_sdk::runtime::Tokio)
-     // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset?
-     .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint("http://adot-collector.adot-collector-kubeprometheus:4317"))
-     .with_resource(Resource::new(vec![KeyValue::new("service.name", "inference-job")]))
-     .with_period(Duration::from_secs(3))
-     .with_timeout(Duration::from_secs(10))
-     .with_aggregation_selector(DefaultAggregationSelector::new())
-     .with_temporality_selector(DefaultTemporalitySelector::new())
-     .build()?;
-  opentelemetry::global::set_meter_provider(provider);
-  Ok(())
-}
+
 
 //#[tokio::main]
 #[actix_web::main]
 async fn main() -> AnyhowResult<()> {
 
-  if let Err(e) = init_otel_metrics_pipeline() {
-    warn!("Failed to initialize OpenTelemetry metrics pipeline, continuing execution: {}", e);
-  }
-
-  let meter = opentelemetry::global::meter("inference-job");
+  let app_name = "inference-job";
 
   let container_environment = bootstrap(BootstrapArgs {
-    app_name: "inference-job",
+    app_name: app_name,
     default_logging_override: Some(DEFAULT_RUST_LOG),
     config_search_directories: &[".", "./config", "crates/service/job/inference_job/config"],
   })?;
 
   info!("Hostname: {}", &container_environment.hostname);
 
+
   // NB: These are non-standard env vars we're injecting ourselves.
   let _k8s_node_name = easyenv::get_env_string_optional("K8S_NODE_NAME");
   let _k8s_pod_name = easyenv::get_env_string_optional("K8S_POD_NAME");
@@ -233,6 +217,25 @@ async fn main() -> AnyhowResult<()> {
 
   let scoped_execution = ScopedExecution::new_from_env()?;
 
+  let build_sha = std::fs::read_to_string("/GIT_SHA")
+      .unwrap_or(String::from("unknown"))
+      .trim()
+      .to_string();
+
+  if let Err(e) = init_otel_metrics_pipeline(
+    JobInstrumentLabels{
+      service_name: app_name.to_string(),
+      service_namespace: container_environment.cluster_name.clone(),
+      service_version: build_sha,
+      service_instance_id: container_environment.hostname.clone(),
+      service_job_scope: easyenv::get_env_string_optional("SCOPED_EXECUTION_MODEL_TYPES").unwrap_or_else(|| "unknown".to_string()),
+    }
+  ) {
+    warn!("Failed to initialize OpenTelemetry metrics pipeline, continuing execution: {}", e);
+  }
+
+  let meter = opentelemetry::global::meter("inference-job");
+
   let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution)?;
 
   let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory(

crates/service/job/inference_job/src/util/instrumentation.rs

@@ -1,5 +1,42 @@
-use opentelemetry::metrics::{Counter, Histogram, Meter, Unit, UpDownCounter};
-use opentelemetry_sdk::metrics::data::Gauge;
+use std::time::Duration;
+use opentelemetry::KeyValue;
+use opentelemetry::metrics::{Counter, Histogram, Meter, Unit};
+use opentelemetry_otlp::WithExportConfig;
+use opentelemetry_sdk::metrics::reader::{DefaultAggregationSelector, DefaultTemporalitySelector};
+use opentelemetry_sdk::Resource;
+
+
+pub struct JobInstrumentLabels {
+    pub service_name: String,
+    pub service_namespace: String,
+    pub service_version: String,
+    pub service_instance_id: String,
+    pub service_job_scope: String,
+}
+
+// TODO(kasisnu, 12/02/24): move this into a config struct/better type once more binaries are instrumented
+pub fn init_otel_metrics_pipeline(
+    job_instrument_labels: JobInstrumentLabels,
+) -> Result<(), opentelemetry::metrics::MetricsError>  {
+    let provider = opentelemetry_otlp::new_pipeline()
+        .metrics(opentelemetry_sdk::runtime::Tokio)
+        // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset?
+        .with_exporter(opentelemetry_otlp::new_exporter().tonic().with_endpoint("http://adot-collector.adot-collector-kubeprometheus:4317"))
+        .with_resource(Resource::new(vec![
+            KeyValue::new("service.name", job_instrument_labels.service_name),
+            KeyValue::new("service.namespace", job_instrument_labels.service_namespace),
+            KeyValue::new("service.version", job_instrument_labels.service_version),
+            KeyValue::new("service.instance.id", job_instrument_labels.service_instance_id),
+            KeyValue::new("service.job.scope", job_instrument_labels.service_job_scope),
+        ]))
+        .with_period(Duration::from_secs(3))
+        .with_timeout(Duration::from_secs(10))
+        .with_aggregation_selector(DefaultAggregationSelector::new())
+        .with_temporality_selector(DefaultTemporalitySelector::new())
+        .build()?;
+    opentelemetry::global::set_meter_provider(provider);
+    Ok(())
+}
 
 pub struct JobInstruments {
     pub job_duration: Histogram<u64>,