296 lines
11 KiB
Rust
296 lines
11 KiB
Rust
//! What a classification on the ACP runtime cost, exported as the counter
|
|
//! `hyperhive.review_classifier.cost.usage` with the run's `repo`, `pr`,
|
|
//! `model` and `currency`, and the tokens it used, as the counter
|
|
//! `hyperhive.review_classifier.token.usage` with the same labels bar
|
|
//! `currency` and a `type` of `input`, `output`, `cacheRead` or
|
|
//! `cacheCreation`. Those are `claude_code.token.usage`'s types, which
|
|
//! `claude` exports for itself on the claude runtime, so the cache hit rate is
|
|
//! `cacheRead` over the three input types on either runtime.
|
|
//!
|
|
//! The amount is computed here from the prices the run carries and the
|
|
//! tokens the agent reports, so it does not depend on how, or whether, the
|
|
//! agent prices a custom provider.
|
|
//!
|
|
//! The exporter is built on the first record, from the standard `OTEL_*`
|
|
//! environment, as `hive-agent`'s `otel_turn_metrics` does; with no
|
|
//! `OTEL_EXPORTER_OTLP_ENDPOINT` a record does nothing.
|
|
|
|
use std::sync::OnceLock;
|
|
use std::time::Duration;
|
|
|
|
use hive_runtime::TokenUsage;
|
|
use opentelemetry::KeyValue;
|
|
use opentelemetry::metrics::{Counter, MeterProvider as _};
|
|
use opentelemetry_otlp::{MetricExporter, Protocol, WithExportConfig};
|
|
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
|
|
use swarm_wire::review_classifier::Prices;
|
|
|
|
/// The currency every price is in.
|
|
pub const CURRENCY: &str = "USD";
|
|
|
|
/// How often the counter is exported.
|
|
const INTERVAL: Duration = Duration::from_mins(1);
|
|
|
|
/// One classification's cost and tokens, and what they are labelled with.
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct Cost {
|
|
pub repo: String,
|
|
pub pr: u64,
|
|
pub model: String,
|
|
/// In [`CURRENCY`].
|
|
pub amount: f64,
|
|
pub usage: TokenUsage,
|
|
}
|
|
|
|
/// Where a classification's cost goes: the metrics store in production, a
|
|
/// list in the tests.
|
|
pub trait CostSink: Send + Sync {
|
|
fn record(&self, cost: &Cost);
|
|
}
|
|
|
|
/// `usage` at `prices`. Cache reads and writes without a price of their own
|
|
/// are priced as input.
|
|
#[must_use]
|
|
pub fn amount(prices: &Prices, usage: &TokenUsage) -> f64 {
|
|
let per_token = |price: f64| price / 1_000_000.0;
|
|
// One model call's count is far below `u32::MAX`.
|
|
let tokens = |n: u64| f64::from(u32::try_from(n).unwrap_or(u32::MAX));
|
|
tokens(usage.input_tokens) * per_token(prices.input)
|
|
+ tokens(usage.output_tokens) * per_token(prices.output)
|
|
+ tokens(usage.cache_read_input_tokens)
|
|
* per_token(prices.cache_read.unwrap_or(prices.input))
|
|
+ tokens(usage.cache_creation_input_tokens)
|
|
* per_token(prices.cache_write.unwrap_or(prices.input))
|
|
}
|
|
|
|
/// The counters in the metrics store.
|
|
pub struct Otel;
|
|
|
|
impl CostSink for Otel {
|
|
fn record(&self, cost: &Cost) {
|
|
static COUNTERS: OnceLock<Option<(SdkMeterProvider, Counters)>> = OnceLock::new();
|
|
let Some((_, counters)) = COUNTERS.get_or_init(build) else {
|
|
return;
|
|
};
|
|
add(counters, cost);
|
|
}
|
|
}
|
|
|
|
/// The cost counter, and the token counter beside it.
|
|
struct Counters {
|
|
cost: Counter<f64>,
|
|
tokens: Counter<u64>,
|
|
}
|
|
|
|
/// The counters on `provider`.
|
|
fn counters(provider: &SdkMeterProvider) -> Counters {
|
|
let meter = provider.meter("hyperhive.review_classifier");
|
|
Counters {
|
|
cost: meter
|
|
.f64_counter("hyperhive.review_classifier.cost.usage")
|
|
.build(),
|
|
tokens: meter
|
|
.u64_counter("hyperhive.review_classifier.token.usage")
|
|
.build(),
|
|
}
|
|
}
|
|
|
|
/// Adds `cost` to `counters` with the run's labels.
|
|
fn add(counters: &Counters, cost: &Cost) {
|
|
let labels = [
|
|
KeyValue::new("repo", cost.repo.clone()),
|
|
KeyValue::new("pr", i64::try_from(cost.pr).unwrap_or(i64::MAX)),
|
|
KeyValue::new("model", cost.model.clone()),
|
|
];
|
|
let with = |extra: KeyValue| {
|
|
let mut all = labels.to_vec();
|
|
all.push(extra);
|
|
all
|
|
};
|
|
counters
|
|
.cost
|
|
.add(cost.amount, &with(KeyValue::new("currency", CURRENCY)));
|
|
let usage = &cost.usage;
|
|
for (kind, n) in [
|
|
("input", usage.input_tokens),
|
|
("output", usage.output_tokens),
|
|
("cacheRead", usage.cache_read_input_tokens),
|
|
("cacheCreation", usage.cache_creation_input_tokens),
|
|
] {
|
|
counters.tokens.add(n, &with(KeyValue::new("type", kind)));
|
|
}
|
|
}
|
|
|
|
/// The provider, kept alive with the counters since it exports only while it
|
|
/// lives, or `None` when no collector is configured or the exporter failed.
|
|
fn build() -> Option<(SdkMeterProvider, Counters)> {
|
|
if !std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_ok_and(|s| !s.trim().is_empty()) {
|
|
return None;
|
|
}
|
|
// The blocking client: `PeriodicReader` exports from a thread of its own
|
|
// with no Tokio reactor. The resource comes from `OTEL_RESOURCE_ATTRIBUTES`.
|
|
let exporter = MetricExporter::builder()
|
|
.with_http()
|
|
.with_protocol(Protocol::HttpJson)
|
|
.build();
|
|
let exporter = match exporter {
|
|
Ok(exporter) => exporter,
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "building the cost metric exporter failed; no cost is exported");
|
|
return None;
|
|
}
|
|
};
|
|
let reader = PeriodicReader::builder(exporter)
|
|
.with_interval(INTERVAL)
|
|
.build();
|
|
let provider = SdkMeterProvider::builder().with_reader(reader).build();
|
|
let counters = counters(&provider);
|
|
Some((provider, counters))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use opentelemetry::Value;
|
|
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
|
|
use opentelemetry_sdk::metrics::data::{
|
|
AggregatedMetrics, MetricData, ResourceMetrics, ScopeMetrics,
|
|
};
|
|
|
|
fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64) -> TokenUsage {
|
|
let mut usage = TokenUsage::default();
|
|
usage.input_tokens = input;
|
|
usage.output_tokens = output;
|
|
usage.cache_read_input_tokens = cache_read;
|
|
usage.cache_creation_input_tokens = cache_write;
|
|
usage
|
|
}
|
|
|
|
#[test]
|
|
fn each_token_class_is_charged_at_its_price_per_million() {
|
|
let prices = Prices {
|
|
input: 3.0,
|
|
output: 15.0,
|
|
cache_read: Some(0.3),
|
|
cache_write: Some(3.75),
|
|
};
|
|
let cost = amount(&prices, &usage(1_000, 200, 10_000, 2_000));
|
|
assert!(
|
|
(cost - (0.003 + 0.003 + 0.003 + 0.0075)).abs() < 1e-12,
|
|
"{cost}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn cache_without_a_price_of_its_own_is_charged_as_input() {
|
|
let prices = Prices {
|
|
input: 2.0,
|
|
output: 0.0,
|
|
cache_read: None,
|
|
cache_write: None,
|
|
};
|
|
let cost = amount(&prices, &usage(0, 0, 500_000, 500_000));
|
|
assert!((cost - 2.0).abs() < 1e-12, "{cost}");
|
|
}
|
|
|
|
#[test]
|
|
fn a_record_adds_the_amount_under_the_counter_name_and_the_run_labels() {
|
|
let exporter = InMemoryMetricExporter::default();
|
|
let reader = PeriodicReader::builder(exporter.clone()).build();
|
|
let provider = SdkMeterProvider::builder().with_reader(reader).build();
|
|
let counters = counters(&provider);
|
|
let prices = Prices {
|
|
input: 3.0,
|
|
output: 15.0,
|
|
cache_read: None,
|
|
cache_write: None,
|
|
};
|
|
let cost = Cost {
|
|
repo: "hyperhive/hyperhive".to_owned(),
|
|
pr: 5214,
|
|
model: "some-model".to_owned(),
|
|
amount: amount(&prices, &usage(1_000_000, 100_000, 0, 0)),
|
|
usage: usage(1_000_000, 100_000, 0, 0),
|
|
};
|
|
add(&counters, &cost);
|
|
provider.force_flush().expect("flush");
|
|
|
|
let metrics = exporter.get_finished_metrics().expect("metrics");
|
|
let metric = metrics
|
|
.iter()
|
|
.flat_map(ResourceMetrics::scope_metrics)
|
|
.flat_map(ScopeMetrics::metrics)
|
|
.find(|m| m.name() == "hyperhive.review_classifier.cost.usage")
|
|
.expect("the cost counter was exported");
|
|
let AggregatedMetrics::F64(MetricData::Sum(sum)) = metric.data() else {
|
|
panic!("not an f64 sum: {:?}", metric.data());
|
|
};
|
|
let points: Vec<_> = sum.data_points().collect();
|
|
assert_eq!(points.len(), 1);
|
|
let point = points[0];
|
|
assert!((point.value() - 4.5).abs() < 1e-12, "{}", point.value());
|
|
let labels: Vec<(String, Value)> = point
|
|
.attributes()
|
|
.map(|kv| (kv.key.to_string(), kv.value.clone()))
|
|
.collect();
|
|
let label = |key: &str| {
|
|
labels
|
|
.iter()
|
|
.find(|(k, _)| k == key)
|
|
.map(|(_, v)| v.clone())
|
|
};
|
|
assert_eq!(label("repo"), Some(Value::from("hyperhive/hyperhive")));
|
|
assert_eq!(label("pr"), Some(Value::from(5214_i64)));
|
|
assert_eq!(label("model"), Some(Value::from("some-model")));
|
|
assert_eq!(label("currency"), Some(Value::from("USD")));
|
|
assert_eq!(labels.len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn a_record_adds_each_token_class_under_its_type() {
|
|
let exporter = InMemoryMetricExporter::default();
|
|
let reader = PeriodicReader::builder(exporter.clone()).build();
|
|
let provider = SdkMeterProvider::builder().with_reader(reader).build();
|
|
let cost = Cost {
|
|
repo: "hyperhive/hyperhive".to_owned(),
|
|
pr: 5245,
|
|
model: "some-model".to_owned(),
|
|
amount: 0.0,
|
|
usage: usage(100, 20, 9_000, 400),
|
|
};
|
|
add(&counters(&provider), &cost);
|
|
provider.force_flush().expect("flush");
|
|
|
|
let metrics = exporter.get_finished_metrics().expect("metrics");
|
|
let metric = metrics
|
|
.iter()
|
|
.flat_map(ResourceMetrics::scope_metrics)
|
|
.flat_map(ScopeMetrics::metrics)
|
|
.find(|m| m.name() == "hyperhive.review_classifier.token.usage")
|
|
.expect("the token counter was exported");
|
|
let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
|
|
panic!("not a u64 sum: {:?}", metric.data());
|
|
};
|
|
let mut by_type: Vec<(String, u64, usize)> = sum
|
|
.data_points()
|
|
.map(|point| {
|
|
let kind = point
|
|
.attributes()
|
|
.find(|kv| kv.key.as_str() == "type")
|
|
.map(|kv| kv.value.to_string())
|
|
.unwrap_or_default();
|
|
(kind, point.value(), point.attributes().count())
|
|
})
|
|
.collect();
|
|
by_type.sort();
|
|
assert_eq!(
|
|
by_type,
|
|
[
|
|
("cacheCreation".to_owned(), 400, 4),
|
|
("cacheRead".to_owned(), 9_000, 4),
|
|
("input".to_owned(), 100, 4),
|
|
("output".to_owned(), 20, 4),
|
|
]
|
|
);
|
|
}
|
|
}
|