feat: add opentelemetry to report tensor ready order

Cargo.toml

@@ -1 +1,5 @@
-workspace = {members = ["bagua-core-internal", "bagua-core-py", "bagua-core-c"],exclude = []}
+workspace = { members = [
+    "bagua-core-internal",
+    "bagua-core-py",
+    "bagua-opentelemetry",
+], exclude = [] }

bagua-core-internal/Cargo.toml

@@ -29,6 +29,8 @@ ureq = "2.1"
 num-traits = "0.2"
 num-derive = "0.3"
 display_utils = "0.4.0"
+opentelemetry = { version = "0.15", features = ["serialize", "metrics"] }
+bagua-opentelemetry = { path = "../bagua-opentelemetry" }
 
 [dependencies.pyo3]
 version = "0.14.1"

bagua-core-internal/src/datatypes/mod.rs

@@ -8,7 +8,6 @@ use crate::comm_ops::python_ffi_op::PythonFFIOp;
 use crate::comm_ops::CommOpTrait;
 use crate::communicators::{BaguaCommunicator, BaguaSingleCommunicator};
 use crate::resource_pool::{CudaMemory, CUDA_DEVICE_MEMORY_POOL};
-use crate::telemetry::TELEMETRY;
 use crate::torch_ffi::root::c10::{DeviceType, StorageImpl, TensorImpl};
 use crate::{kernels, BaguaCoreError};
 use itertools::Itertools;
@@ -769,12 +768,6 @@ impl BaguaTensor {
         if cuda_event_ptr == 0 {
             tracing::info!("mark comm ready with an event 0, ignoring event");
         }
-        match TELEMETRY.as_ref() {
-            None => {}
-            Some(ref x) => {
-                x.lock().new_tensor_ready(self.inner.read().name.as_str());
-            }
-        }
         let mut guard = self.inner.write();
         guard.ready_for_comm = true;
         guard.ready_cuda_event_ptr = cuda_event_ptr;

bagua-core-internal/src/lib.rs

@@ -10,16 +10,20 @@ pub mod datatypes;
 pub mod events;
 pub mod kernels;
 pub mod resource_pool;
-pub mod telemetry;
 mod torch_ffi;
 
 use crate::comm_ops::CommOpTrait;
-use crate::telemetry::{SCHEDULED_THREAD_POOL, TELEMETRY};
+use bagua_opentelemetry;
 use cpp::cpp;
 use datatypes::{BaguaBucket, BaguaTensor};
 use events::BaguaEventChannel;
 use flume::RecvTimeoutError;
 use hashbrown::{HashMap, HashSet};
+use opentelemetry::{
+    global,
+    trace::{Span, Tracer},
+    KeyValue,
+};
 use std::collections::VecDeque;
 use std::fmt::Debug;
 use std::sync::Arc;
@@ -167,6 +171,8 @@ impl BaguaCommBackend {
     }
 }
 
+static TELEMETRY_INIT_ONCE: std::sync::Once = std::sync::Once::new();
+
 impl BaguaCommBackend {
     pub fn new(schedule_channel_cap: usize, device_id: usize) -> BaguaCommBackend {
         unsafe {
@@ -181,6 +187,20 @@ impl BaguaCommBackend {
         let (monitor_op_finish_channel_sender, monitor_op_finish_channel_receiver) =
             flume::unbounded();
 
+        TELEMETRY_INIT_ONCE.call_once(|| {
+            match std::env::var("AUTO_TUNE_SERVER_ADDR") {
+                Ok(server_addr) => {
+                    tracing::info!("detected auto tuning server, connecting");
+                    bagua_opentelemetry::init_tracer(&server_addr);
+                }
+                Err(_) => {
+                    tracing::warn!(
+                        "auto tuning server not detected, may experience degraded performance"
+                    );
+                }
+            };
+        });
+
         BaguaCommBackend {
             ordered_buckets: Default::default(),
             bucket_mapping: Default::default(),
@@ -282,6 +302,10 @@ impl BaguaCommBackend {
         tensor: &BaguaTensor,
         ready_cuda_event_ptr: u64,
     ) -> Result<(), BaguaCoreError> {
+        let tracer = global::tracer("bagua-core");
+        let mut span = tracer.start("tensor_ready");
+        span.set_attribute(KeyValue::new("tensor_name", tensor.name()));
+
         tensor.mark_comm_ready(ready_cuda_event_ptr);
         while self.should_schedule()? {
             let bucket = self.ordered_buckets.pop_front().unwrap();
@@ -311,27 +335,4 @@ impl BaguaCommBackend {
             }
         }
     }
-
-    pub fn start_upload_telemetry(&self, skip: bool) -> Result<(), BaguaCoreError> {
-        SCHEDULED_THREAD_POOL.execute(move || match TELEMETRY.as_ref() {
-            None => {}
-            Some(x) => {
-                let mut guard = x.lock();
-                match skip {
-                    true => {
-                        guard.clear();
-                    }
-                    false => {
-                        match guard.push_payload_and_clear() {
-                            Ok(_) => {}
-                            Err(x) => {
-                                tracing::error!("{:?}", x)
-                            }
-                        };
-                    }
-                }
-            }
-        });
-        Ok(())
-    }
 }

bagua-core-py/src/lib.rs

@@ -339,11 +339,6 @@ impl BaguaCommBackendPy {
         py.allow_threads(|| self.inner.wait_pending_comm_ops())
             .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
     }
-
-    pub fn start_upload_telemetry(&self, skip: bool, py: Python) -> PyResult<()> {
-        py.allow_threads(|| self.inner.start_upload_telemetry(skip))
-            .map_err(|e| PyRuntimeError::new_err(format!("{:?}", e)))
-    }
 }
 
 #[pyclass(dict)]

bagua-opentelemetry/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "bagua-opentelemetry"
+version = "0.1.0"
+edition = "2018"
+publish = ["private"]
+
+[dependencies]
+tracing = "0.1"
+async-std = { version = "1.6", features = ["attributes", "tokio1"] }
+async-trait = { version = "0.1" }
+hyper = { version = "0.14", features = ["full"] }
+opentelemetry = { version = "0.15", default-features = false, features = [
+    "trace",
+    "rt-async-std",
+] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+reqwest = { version = "0.11", features = ["json"] }
+tokio = { version = "1", features = ["full"] }
+futures = { version = "0.3" }
+
+tokio-stream = { version = "0.1" }

bagua-opentelemetry/src/exporter/agent.rs

@@ -0,0 +1,44 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Clone, Debug, Hash)]
+pub struct BaguaSpan {
+    pub trace_id: u128,
+    pub action: String,
+    pub tensor_name: String,
+    pub start_time: u128,
+    pub end_time: u128,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, Hash)]
+pub struct BaguaBatch {
+    pub spans: Vec<BaguaSpan>,
+}
+
+#[derive(Debug)]
+pub struct AgentAsyncClientHTTP {
+    server_addr: String,
+    client: reqwest::Client,
+}
+
+impl AgentAsyncClientHTTP {
+    pub fn new(server_addr: String) -> AgentAsyncClientHTTP {
+        Self {
+            server_addr: server_addr,
+            client: reqwest::Client::new(),
+        }
+    }
+
+    pub async fn emit_batch(
+        &mut self,
+        batch: BaguaBatch,
+    ) -> Result<reqwest::Response, reqwest::Error> {
+        let uri = format!(
+            "http://{}/api/v1/report_tensor_execution_order",
+            self.server_addr
+        );
+
+        let resp = self.client.post(uri).json(&batch).send().await?;
+
+        Ok(resp)
+    }
+}

bagua-opentelemetry/src/exporter/mod.rs

@@ -0,0 +1,62 @@
+pub mod agent;
+
+use crate::exporter::agent::{AgentAsyncClientHTTP, BaguaBatch, BaguaSpan};
+use async_trait::async_trait;
+use opentelemetry::{sdk::export::trace, Key};
+use reqwest::StatusCode;
+use std::time::UNIX_EPOCH;
+
+#[derive(Debug)]
+pub struct Exporter {
+    pub uploader: AgentAsyncClientHTTP,
+}
+
+#[async_trait]
+impl trace::SpanExporter for Exporter {
+    async fn export(&mut self, batch: Vec<trace::SpanData>) -> trace::ExportResult {
+        let mut bagua_spans = Vec::new();
+        for span in batch {
+            let bagua_span = BaguaSpan {
+                trace_id: span.span_context.trace_id().to_u128(),
+                action: span.name.into_owned(),
+                tensor_name: span
+                    .attributes
+                    .get(&Key::new("tensor_name"))
+                    .unwrap()
+                    .as_str()
+                    .to_string(),
+                start_time: span
+                    .start_time
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_millis(),
+                end_time: span
+                    .end_time
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_millis(),
+            };
+
+            bagua_spans.push(bagua_span);
+        }
+
+        let resp = self
+            .uploader
+            .emit_batch(BaguaBatch { spans: bagua_spans })
+            .await;
+        match resp {
+            Ok(resp) => {
+                if resp.status() != StatusCode::OK {
+                    tracing::warn!("upload bagua span failed, resp={:?}", resp);
+                }
+            }
+            Err(err) => {
+                tracing::warn!("upload bagua span failed, err={:?}", err);
+            }
+        }
+
+        Ok(())
+    }
+
+    fn shutdown(&mut self) {}
+}

bagua-opentelemetry/src/lib.rs

@@ -0,0 +1,21 @@
+pub mod exporter;
+
+use crate::exporter::agent::AgentAsyncClientHTTP;
+use crate::exporter::Exporter;
+use opentelemetry;
+use opentelemetry::{global, sdk, trace::Tracer, trace::TracerProvider};
+
+pub fn init_tracer(autotune_server_addr: &str) -> impl Tracer {
+    let exporter = Exporter {
+        uploader: AgentAsyncClientHTTP::new(autotune_server_addr.to_string()),
+    };
+
+    let builder = sdk::trace::TracerProvider::builder()
+        .with_batch_exporter(exporter, opentelemetry::runtime::AsyncStd);
+
+    let tracer_provider = builder.build();
+    let tracer = tracer_provider.get_tracer("bagua-opentelemetry", Some(env!("CARGO_PKG_VERSION")));
+    let _ = global::set_tracer_provider(tracer_provider);
+
+    tracer
+}

setup.py

@@ -194,6 +194,9 @@ def install_dependency_library():
         ],
         author="Kuaishou AI Platform & DS3 Lab",
         author_email="admin@mail.xrlian.com",
-        install_requires=[],
+        install_requires=[
+            "setuptools_rust",
+            "colorama",
+        ],
         zip_safe=False,
     )