はじめに
MLflowでは、サードパーティのOpenTelemetry(OTel)クライアントからマネージドOTLPコレクター経由でトレースを取り込む際に、span属性gen_ai.usage.input_tokens / gen_ai.usage.output_tokensを設定することでトークン使用量をトレースに載せられます。
このドキュメントのトークン使用量テーブルでは、合計(Total)の欄が(not set — calculated automatically)となっており、「合計は自動計算される」と読めます。ところが実際には、child spanにトークン量を設定してもトレースレベルの合計に出てこない、という状況に出会うことがあります。
ポイントは、「自動計算される」が指すのは"合計値(= 入力 + 出力)の導出"だけで、その材料をどのspanから読むかは別の話、という点です。この挙動を実際にワークスペースで検証したので、対処法も含めて書いていきます。
合計に含まれるのはroot spanの入力・出力だけ
先に結論です。トレースレベルで自動計算される合計トークン数に含まれるのは、root spanに設定されたgen_ai.usage.input_tokensとgen_ai.usage.output_tokensだけです。
- child spanに設定したトークン量は、この合計に足し込まれません
- child spanにだけ値があってroot spanに無い場合、合計を計算する材料がroot spanに存在しないため、トレースレベルのトークン量は空になります
これは不具合ではなく集計仕様どおりの動作です。ドキュメントも、トレースサマリの集計元がroot spanであることを明記しています。
MLflow reads inputs, outputs, token usage, and session ID from the root span to populate the trace summary in the UI.
(訳: MLflowは、UIのトレースサマリを生成するために、入力・出力・トークン使用量・セッションIDをroot spanから読み取ります。)
実際に送って確かめる
「root spanだけが集計対象」を、生のOpenTelemetry SDKからトレースを送って確認します。第三者OTelクライアントを模して、gen_ai.usage.*を付けたspanをマネージドOTLPコレクター(/api/2.0/otel/v1/traces)に送信します。
ケースは2つ用意しました。
- ケースA: rootに設定(入力150/出力42)、childには別の値(999/999)
- ケースB: rootには設定なし、childのみ(入力90/出力30)
# pip install "mlflow[databricks]>=3.14.0" opentelemetry-sdk opentelemetry-exporter-otlp-proto-http
from opentelemetry import trace as otrace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
TABLE = "<catalog>.<schema>.<prefix>_otel_spans" # Experimentのトレース保存場所に紐付いたテーブル
exporter = OTLPSpanExporter(
endpoint="https://<workspace-host>/api/2.0/otel/v1/traces",
headers={"authorization": f"Bearer {token}", "X-Databricks-UC-Table-Name": TABLE},
)
provider = TracerProvider(resource=Resource.create({"service.name": "token-usage-verify"}))
provider.add_span_processor(SimpleSpanProcessor(exporter))
otrace.set_tracer_provider(provider)
tracer = otrace.get_tracer("verify")
# ケースA: rootに設定(150/42)、childには別値(999/999)
with tracer.start_as_current_span("caseA_root") as root:
root.set_attribute("gen_ai.usage.input_tokens", 150)
root.set_attribute("gen_ai.usage.output_tokens", 42)
with tracer.start_as_current_span("caseA_child") as child:
child.set_attribute("gen_ai.usage.input_tokens", 999)
child.set_attribute("gen_ai.usage.output_tokens", 999)
# ケースB: rootには設定なし、childのみ(90/30)
with tracer.start_as_current_span("caseB_root") as root:
with tracer.start_as_current_span("caseB_child") as child:
child.set_attribute("gen_ai.usage.input_tokens", 90)
child.set_attribute("gen_ai.usage.output_tokens", 30)
provider.force_flush()
これでトレースが送信されます。トレースレベルのトークン量は、mlflow.get_trace(trace_id).info.token_usage(内部的にはmlflow.trace.tokenUsage)で読み戻せます。
import mlflow
tr = mlflow.get_trace(trace_id)
print(tr.info.token_usage)
検証結果
結果は次のとおりでした。
| ケース | 送信内容(root / child) | トレースレベル合計mlflow.trace.tokenUsage
|
|---|---|---|
| A: rootに設定 | root=入力150/出力42、child=999/999 | {"input_tokens": 150, "output_tokens": 42, "total_tokens": 192} |
| B: childのみ設定 | root=なし、child=入力90/出力30 | 空(token_usage = null) |
-
ケースA: 合計が
150 + 999ではなく150 + 42 = 192。root span自身の値だけが集計され、childの999は含まれていません。total(合計)は自動算出されています - ケースB: childに値があってもroot spanに無いため、トレースレベルの合計は空。childからは拾い上げられません
いずれも「root spanだけが集計対象」という仕様どおりの結果でした。
混同しやすい「別経路」との違い
ここがハマりどころです。「MLflowはトークン量を全spanから合算してくれるはず」という感覚は、実は間違いではありません。ただし、それが効くのは別の経路です。「合計に何が含まれるか」を軸に並べると、2つの経路の違いが明確になります。
| 経路 | 読む属性キー | 合計に含まれる範囲 |
|---|---|---|
| サードパーティOTel取り込み | gen_ai.usage.* |
root span自身の入力・出力のみ(子孫は含まない) |
| MLflowネイティブ / autolog | mlflow.chat.tokenUsage |
全spanを走査して合算(親が値を持つ場合は二重計上を避ける) |
全spanを合算するのはmlflow.chat.tokenUsageという別の属性キーで、これはMLflow自身のインストルメンテーションが各spanに付与するものです。サードパーティOTelでgen_ai.usage.*だけを付けたspanにはmlflow.chat.tokenUsageが無いため、全span合算のロジックは働きません。結果として、サードパーティOTel経路では合計に含まれるのはroot span自身の値だけになります。
両経路とも最終的に同じトレースレベルの項目(mlflow.trace.tokenUsage、UIに表示される値)を埋めますが、取り込み元によって合算範囲が異なるわけです。
対処: root spanに「合算値」を設定する
対処は簡単です。トレースレベルの合計にトークン量を載せたい場合は、root spanに合算値を設定します。
-
gen_ai.usage.input_tokens… トレース内の全LLM呼び出しの入力トークンの合計 -
gen_ai.usage.output_tokens… トレース内の全LLM呼び出しの出力トークンの合計
合計(total)はroot spanの入力＋出力から自動算出されるので、設定するのは入力・出力の2つだけで十分です。
# child spanが2回のLLM呼び出し(入力90+60、出力30+12)だった場合、
# その合算値 入力150・出力42をroot spanに設定する。
root.set_attribute("gen_ai.usage.input_tokens", 150) # = 90 + 60
root.set_attribute("gen_ai.usage.output_tokens", 42) # = 30 + 12
これでトレースレベルの合計にトークン量が表示されます。MLflowネイティブ経路が内部で行う「全span合算」に相当する処理を、サードパーティOTel経路では利用側で行ってroot spanに集約する、というイメージです。
まとめ
サードパーティOTel取り込みでのトークン量集計を検証して分かったことをまとめます。
- サードパーティOTelの
gen_ai.usage.*では、トレースレベルの合計はroot spanの入力・出力だけから算出される(childは含まない) - childにだけ設定すると合計は空になる。エラーは出ず、仕様どおりの動作
- 合計に載せたいトークン量は、root spanに合算して設定する。total(合計)は自動算出
- 全span合算が効くのはMLflowネイティブ / autologの
mlflow.chat.tokenUsage経路であり、属性キーも処理層も別物
一番の収穫は、「自動計算」という言葉の射程を検証で切り分けられたことでした。ドキュメントの一文だけを読むと「全部足してくれる」と期待しがちですが、自動なのはtotalの導出だけで、材料を集める責任は利用側にあります。サードパーティOTelクライアントからMLflowにトレースを送る構成では、root spanへの合算値の設定を計装の設計に最初から織り込んでおくのが良さそうです。