diff --git a/src/openjarvis/telemetry/aggregator.py b/src/openjarvis/telemetry/aggregator.py index f68b49be..986cdf9f 100644 --- a/src/openjarvis/telemetry/aggregator.py +++ b/src/openjarvis/telemetry/aggregator.py @@ -24,6 +24,7 @@ class ModelStats: total_energy_joules: float = 0.0 avg_gpu_utilization_pct: float = 0.0 avg_throughput_tok_per_sec: float = 0.0 + avg_tokens_per_joule: float = 0.0 avg_energy_per_output_token_joules: float = 0.0 avg_throughput_per_watt: float = 0.0 total_prefill_energy_joules: float = 0.0 @@ -47,6 +48,7 @@ class EngineStats: total_energy_joules: float = 0.0 avg_gpu_utilization_pct: float = 0.0 avg_throughput_tok_per_sec: float = 0.0 + avg_tokens_per_joule: float = 0.0 avg_energy_per_output_token_joules: float = 0.0 avg_throughput_per_watt: float = 0.0 total_prefill_energy_joules: float = 0.0 @@ -122,10 +124,15 @@ class TelemetryAggregator: # Build optional columns for new fields (graceful on old DBs) extra_cols = "" + has_tpj = self._safe_col("tokens_per_joule") has_derived = self._safe_col("energy_per_output_token_joules") has_phase = self._safe_col("prefill_energy_joules") has_itl = self._safe_col("mean_itl_ms") + if has_tpj: + extra_cols += ( + ", AVG(tokens_per_joule) AS avg_tokens_per_joule" + ) if has_derived: extra_cols += ( ", AVG(energy_per_output_token_joules)" @@ -178,6 +185,8 @@ class TelemetryAggregator: avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0, avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0, ) + if has_tpj: + ms.avg_tokens_per_joule = r["avg_tokens_per_joule"] or 0.0 if has_derived: ms.avg_energy_per_output_token_joules = ( r["avg_energy_per_output_token_joules"] or 0.0 @@ -206,10 +215,15 @@ class TelemetryAggregator: where, params = self._time_filter(since, until) extra_cols = "" + has_tpj = self._safe_col("tokens_per_joule") has_derived = self._safe_col("energy_per_output_token_joules") has_phase = self._safe_col("prefill_energy_joules") has_itl = self._safe_col("mean_itl_ms") + if has_tpj: + extra_cols += ( + ", AVG(tokens_per_joule) AS avg_tokens_per_joule" + ) if has_derived: extra_cols += ( ", AVG(energy_per_output_token_joules)" @@ -258,6 +272,8 @@ class TelemetryAggregator: avg_gpu_utilization_pct=r["avg_gpu_utilization_pct"] or 0.0, avg_throughput_tok_per_sec=r["avg_throughput_tok_per_sec"] or 0.0, ) + if has_tpj: + es.avg_tokens_per_joule = r["avg_tokens_per_joule"] or 0.0 if has_derived: es.avg_energy_per_output_token_joules = ( r["avg_energy_per_output_token_joules"] or 0.0 diff --git a/src/openjarvis/telemetry/store.py b/src/openjarvis/telemetry/store.py index 01cccd90..c7601134 100644 --- a/src/openjarvis/telemetry/store.py +++ b/src/openjarvis/telemetry/store.py @@ -37,6 +37,7 @@ CREATE TABLE IF NOT EXISTS telemetry ( cpu_energy_joules REAL NOT NULL DEFAULT 0.0, gpu_energy_joules REAL NOT NULL DEFAULT 0.0, dram_energy_joules REAL NOT NULL DEFAULT 0.0, + tokens_per_joule REAL NOT NULL DEFAULT 0.0, energy_per_output_token_joules REAL NOT NULL DEFAULT 0.0, throughput_per_watt REAL NOT NULL DEFAULT 0.0, prefill_energy_joules REAL NOT NULL DEFAULT 0.0, @@ -61,6 +62,7 @@ INSERT INTO telemetry ( throughput_tok_per_sec, prefill_latency_seconds, decode_latency_seconds, energy_method, energy_vendor, batch_id, is_warmup, cpu_energy_joules, gpu_energy_joules, dram_energy_joules, + tokens_per_joule, energy_per_output_token_joules, throughput_per_watt, prefill_energy_joules, decode_energy_joules, mean_itl_ms, median_itl_ms, p90_itl_ms, p95_itl_ms, p99_itl_ms, std_itl_ms, @@ -70,7 +72,7 @@ INSERT INTO telemetry ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ?, ?, ? ) """ @@ -88,6 +90,7 @@ _MIGRATE_COLUMNS = [ ("cpu_energy_joules", "REAL NOT NULL DEFAULT 0.0"), ("gpu_energy_joules", "REAL NOT NULL DEFAULT 0.0"), ("dram_energy_joules", "REAL NOT NULL DEFAULT 0.0"), + ("tokens_per_joule", "REAL NOT NULL DEFAULT 0.0"), ("energy_per_output_token_joules", "REAL NOT NULL DEFAULT 0.0"), ("throughput_per_watt", "REAL NOT NULL DEFAULT 0.0"), ("prefill_energy_joules", "REAL NOT NULL DEFAULT 0.0"), @@ -153,6 +156,7 @@ class TelemetryStore: rec.cpu_energy_joules, rec.gpu_energy_joules, rec.dram_energy_joules, + rec.tokens_per_joule, rec.energy_per_output_token_joules, rec.throughput_per_watt, rec.prefill_energy_joules, diff --git a/tests/telemetry/test_store_tokens_per_joule.py b/tests/telemetry/test_store_tokens_per_joule.py new file mode 100644 index 00000000..3f3d3971 --- /dev/null +++ b/tests/telemetry/test_store_tokens_per_joule.py @@ -0,0 +1,65 @@ +"""Tests for tokens_per_joule storage and aggregation.""" + +from __future__ import annotations + +import time + +import pytest + +from openjarvis.core.types import TelemetryRecord +from openjarvis.telemetry.store import TelemetryStore +from openjarvis.telemetry.aggregator import TelemetryAggregator + + +class TestTokensPerJouleStorage: + def test_store_and_retrieve(self, tmp_path): + db = tmp_path / "tel.db" + store = TelemetryStore(db_path=db) + rec = TelemetryRecord( + timestamp=time.time(), + model_id="test-model", + completion_tokens=50, + energy_joules=2.5, + tokens_per_joule=20.0, + ) + store.record(rec) + store.close() + agg = TelemetryAggregator(db) + stats = agg.per_model_stats() + assert len(stats) == 1 + assert stats[0].avg_tokens_per_joule == pytest.approx(20.0, rel=0.1) + agg.close() + + def test_aggregate_multiple(self, tmp_path): + db = tmp_path / "tel.db" + store = TelemetryStore(db_path=db) + for tpj in [10.0, 20.0, 30.0]: + rec = TelemetryRecord( + timestamp=time.time(), + model_id="m1", + tokens_per_joule=tpj, + ) + store.record(rec) + store.close() + agg = TelemetryAggregator(db) + stats = agg.per_model_stats() + assert stats[0].avg_tokens_per_joule == pytest.approx(20.0, rel=0.1) + agg.close() + + def test_engine_stats_aggregate(self, tmp_path): + db = tmp_path / "tel.db" + store = TelemetryStore(db_path=db) + for tpj in [15.0, 25.0]: + rec = TelemetryRecord( + timestamp=time.time(), + model_id="m1", + engine="ollama", + tokens_per_joule=tpj, + ) + store.record(rec) + store.close() + agg = TelemetryAggregator(db) + stats = agg.per_engine_stats() + assert len(stats) == 1 + assert stats[0].avg_tokens_per_joule == pytest.approx(20.0, rel=0.1) + agg.close()