Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add support for openTelemetry #3576

Merged
merged 6 commits into from
Dec 17, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
chore: update span names
Signed-off-by: Wei Zhang <kweizh@tabbyml.com>
zwpaper committed Dec 17, 2024
commit fa700490528b2ea2beadc75b865e2a0483250da8
2 changes: 1 addition & 1 deletion crates/http-api-bindings/src/embedding/openai.rs
Original file line number Diff line number Diff line change
@@ -41,12 +41,12 @@
user: None,
dimensions: None,
};
let resp = self
.client
.embeddings()
.create(request)
.instrument(info_span!("embed_openai"))
.instrument(info_span!("embedding", kind = "openai"))
.await?;

Check warning on line 49 in crates/http-api-bindings/src/embedding/openai.rs

Codecov / codecov/patch

crates/http-api-bindings/src/embedding/openai.rs#L44-L49

Added lines #L44 - L49 were not covered by tests
let data = resp
.data
.into_iter()
2 changes: 1 addition & 1 deletion crates/http-api-bindings/src/rate_limit.rs
Original file line number Diff line number Diff line change
@@ -36,10 +36,10 @@
impl Embedding for RateLimitedEmbedding {
async fn embed(&self, prompt: &str) -> anyhow::Result<Vec<f32>> {
self.rate_limiter.acquire(1).await;
self.embedding
.embed(prompt)
.instrument(info_span!("rate_limit_embed"))
.instrument(info_span!("rate_limited_compute_embedding"))
.await

Check warning on line 42 in crates/http-api-bindings/src/rate_limit.rs

Codecov / codecov/patch

crates/http-api-bindings/src/rate_limit.rs#L39-L42

Added lines #L39 - L42 were not covered by tests
}
}

2 changes: 1 addition & 1 deletion crates/tabby-index/src/code/mod.rs
Original file line number Diff line number Diff line change
@@ -128,7 +128,7 @@ async fn build_binarize_embedding_tokens(
) -> Result<Vec<String>> {
let embedding = match embedding
.embed(body)
.instrument(info_span!("embed", corpus = corpus::CODE))
.instrument(info_span!("index_compute_embedding", corpus = corpus::CODE))
.await
{
Ok(x) => x,

Unchanged files with check annotations Beta

color_eyre::install().expect("Must be able to install color_eyre");
let cli = Cli::parse();
let _guard = otel::init_tracing_subscriber(cli.otlp_endpoint);

Check warning on line 58 in crates/tabby/src/main.rs

Codecov / codecov/patch

crates/tabby/src/main.rs#L58

Added line #L58 was not covered by tests
let config = Config::load().expect("Must be able to load config");
let root = tabby_common::path::tabby_root();
};
// Create a Resource that captures information about the entity for which telemetry is recorded.
fn resource() -> Resource {
Resource::from_schema_url(
[
KeyValue::new(SERVICE_NAME, env!("CARGO_PKG_NAME")),
KeyValue::new(SERVICE_VERSION, env!("CARGO_PKG_VERSION")),
KeyValue::new(DEPLOYMENT_ENVIRONMENT_NAME, "develop"),
],
SCHEMA_URL,
)
}

Check warning on line 28 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L19-L28

Added lines #L19 - L28 were not covered by tests
// Construct TracerProvider for OpenTelemetryLayer
fn init_tracer_provider(otlp_endpoint: String) -> TracerProvider {
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(otlp_endpoint)
.with_timeout(std::time::Duration::from_secs(3))
.build()
.unwrap();
TracerProvider::builder()
// Customize sampling strategy
.with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(
1.0,
))))
// If export trace to AWS X-Ray, you can use XrayIdGenerator
.with_id_generator(RandomIdGenerator::default())
.with_resource(resource())
.with_batch_exporter(exporter, runtime::Tokio)
.build()
}

Check warning on line 49 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L31-L49

Added lines #L31 - L49 were not covered by tests
// Initialize tracing-subscriber and return OtelGuard for opentelemetry-related termination processing
pub fn init_tracing_subscriber(otlp_endpoint: Option<String>) -> OtelGuard {
let mut layers: Vec<Box<dyn Layer<Registry> + Send + Sync>> = Vec::new();

Check warning on line 53 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L52-L53

Added lines #L52 - L53 were not covered by tests
let tracer_provider = if let Some(endpoint) = otlp_endpoint {
let tracer_provider = init_tracer_provider(endpoint);
let tracer = tracer_provider.tracer("tracing-otel-subscriber");
layers.push(Box::new(OpenTelemetryLayer::new(tracer)));
Some(tracer_provider)

Check warning on line 59 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L55-L59

Added lines #L55 - L59 were not covered by tests
} else {
None

Check warning on line 61 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L61

Added line #L61 was not covered by tests
};
let fmt_layer = tracing_subscriber::fmt::layer()
.with_file(true)
.with_line_number(true)
.boxed();
layers.push(fmt_layer);

Check warning on line 68 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L64-L68

Added lines #L64 - L68 were not covered by tests
let mut dirs = if cfg!(feature = "prod") {
"tabby=info,otel=debug,http_api_bindings=info,llama_cpp_server=info".into()

Check warning on line 71 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L70-L71

Added lines #L70 - L71 were not covered by tests
} else {
"tabby=debug,otel=debug,http_api_bindings=debug,llama_cpp_server=debug".into()

Check warning on line 73 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L73

Added line #L73 was not covered by tests
};
if let Ok(env) = std::env::var(EnvFilter::DEFAULT_ENV) {
dirs = format!("{dirs},{env}")
};

Check warning on line 78 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L76-L78

Added lines #L76 - L78 were not covered by tests
let env_filter = EnvFilter::builder()
.with_default_directive(LevelFilter::WARN.into())
.parse_lossy(dirs);
tracing_subscriber::registry()
.with(layers)
.with(env_filter)
.init();
OtelGuard { tracer_provider }
}

Check warning on line 90 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L80-L90

Added lines #L80 - L90 were not covered by tests
pub struct OtelGuard {
tracer_provider: Option<TracerProvider>,
}
impl Drop for OtelGuard {
fn drop(&mut self) {
if let Some(tracer_provider) = self.tracer_provider.take() {
if let Err(err) = tracer_provider.shutdown() {
eprintln!("{err:?}");
}
}
}

Check warning on line 103 in crates/tabby/src/otel.rs

Codecov / codecov/patch

crates/tabby/src/otel.rs#L97-L103

Added lines #L97 - L103 were not covered by tests
}