Add tokens_per_joule to telemetry store and aggregator (Task 3)

- Add tokens_per_joule column to TelemetryStore schema, INSERT, and migration
- Add avg_tokens_per_joule field to ModelStats and EngineStats
- Add AVG(tokens_per_joule) to per_model_stats() and per_engine_stats() SQL
  queries using _safe_col() pattern for backward compatibility
- Add 3 tests for store/retrieve, multi-record aggregation, and engine stats

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Jon Saad-Falcon
2026-02-28 06:53:23 +00:00
co-authored by Claude Opus 4.6
parent 4e5709d478
commit 45892539bc
3 changed files with 86 additions and 1 deletions
+16
View File
@@ -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
+5 -1
View File
@@ -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,
@@ -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()