はじめに
MLflowのトレーシングでspan link (OpenTelemetryのSpan Link) が使えるようになっていました。最新の3.16.0で確認しています。
入ったのは2段階で、mlflow.entities.Link と span.add_link() が3.13.0、start_span(links=[...]) と @mlflow.trace(links=[...]) の引数が3.15.0です。3.13.0のリリースハイライトにはspan linkの記載がなく、気づかないうちに増えていた機能でした。
span linkが何なのかというところから、OSSのMLflowで実際に手元に環境を作って動かすところまで、ハマったポイントも含めて書いていきます。
span linkとは何か
トレースの中でスパン同士を繋ぐ方法は、通常は親子関係です。関数Aの中で関数Bを呼べば、AのスパンがBのスパンの親になる。1本のトレースの中で木構造ができます。
span linkは、この親子関係とは別に、直接の呼び出し関係にないスパン同士を繋ぐための仕組みです。OpenTelemetryの仕様に含まれているもので、MLflowのSpanオブジェクトもこの仕様に沿っています。
| 親子関係 (parent_span_id) | span link (links) | |
|---|---|---|
| 繋げる範囲 | 同一トレース内 | トレースをまたげる |
| 個数 | 親は1つ | 何本でも張れる |
| 意味 | 呼び出しの包含関係 | 任意の関連づけ |
| 付けるタイミング | スパン生成時に自動 | 明示的に指定 |
効いてくるのは、処理が1本のトレースに収まらない構成です。マルチエージェントで、オーケストレーターとサブエージェントがそれぞれ別トレースとして記録される場合。非同期のキューを挟んで、投入側と処理側でトレースが分かれる場合。バッチ処理が個別のリクエストをまとめて処理する場合。リトライで作り直したトレースを元のトレースに紐づけたい場合。
いずれも「呼び出しの木」としては繋がらないけれど、デバッグするときには繋がっていてほしい関係です。ここをlinkで埋める、という位置づけになります。
どのバージョンから使えるのか
リリースノートを追いかけても分からなかったので、各バージョンのwheelを展開して確認しました。
- 3.12.0以前:
mlflow/entities/link.pyが存在しない - 3.13.0:
Linkエンティティとspan.add_link()が追加。mlflow.entitiesからimportできる - 3.15.0:
start_span(links=[...])と@mlflow.trace(links=[...])の引数が追加
3.13系でも add_link() でリンクは張れますが、デコレータやコンテキストマネージャの引数で宣言的に渡したいなら3.15.0以上が必要です。この記事は3.16.0で検証しています。
Link の中身は3つのフィールドだけのシンプルなdataclassです。
from mlflow.entities import Link
Link(
trace_id="tr-xxxxxxxx", # リンク先のトレースID
span_id="0123456789abcdef", # そのトレース内のスパンID
attributes={"relation": "triggered_by"}, # 任意の属性
)
trace_id はMLflow形式の tr-xxx、v4形式の trace:/<location>/<hex>、生のhex文字列のいずれも受け付けます。attributes はJSONシリアライズ可能な値であれば自由に付けられるので、「どういう関係でリンクしているのか」をここに書いておけます。
検証環境を用意する
ローカルのMacで試します。仮想環境を切って、フル版のMLflowを入れます。
python -m venv .venv
source .venv/bin/activate
pip install "mlflow>=3.15"
ここで軽量版の mlflow-tracing を選ぶと後で詰まります。最初こちらで試したのですが、トラッキングストアを使おうとした時点で落ちました。
ImportError: cannot import name 'EvaluationDataset' from 'mlflow.entities'
mlflow-tracing は計装のAPIだけを持つパッケージで、トラッキングストアと mlflow ui のサーバは含まれていません。UIまで確認したいならフル版が必要です。
もう1つ、バックエンドはSQLiteにしてください。従来よく使っていた ./mlruns のファイルストアは、3.16ではメンテナンスモード扱いで例外になります。
MlflowException: The filesystem tracking backend (e.g., './mlruns') is in maintenance mode
and will not receive further updates. Please migrate to a database backend
(e.g., 'sqlite:///mlflow.db') to access the latest MLflow features.
MLFLOW_ALLOW_FILE_STORE=true で回避もできますが、素直にSQLiteにしたほうが早いです。
2つのトレースをspan linkで繋ぐ
上流の検索処理と下流の生成処理が別トレースになっている、という想定でコードを書きます。実際のマルチエージェント構成の最小形です。
まずトラッキング先を決めます。
import mlflow
from mlflow.entities import Link
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("span-link-demo")
sqlite:///mlflow.db は相対パスなので、DBは実行時のカレントディレクトリに作られます。あとでUIを起動するときも同じディレクトリにいる必要があります。
上流のトレースを1本作ります。
@mlflow.trace(name="retriever")
def retrieve(query: str):
return [f"doc for {query}"]
retrieve("span link とは")
producer_trace_id = mlflow.get_last_active_trace_id()
ここがハマりどころです。このあとすぐ mlflow.get_trace() で読み直すと、こう言われて None が返ってきます。
Failed to get trace from the tracking store: Trace with ID tr-xxx cannot be loaded
because its span data is corrupted.
「corrupted」と出るのでデータが壊れたのかと思いますが、そうではなく、トレースの書き込みが非同期なだけです。読み直す前にフラッシュすれば解決します。メッセージからは原因にたどり着きにくいので、先に知っておくと時間を節約できます。
mlflow.flush_trace_async_logging()
producer_trace = mlflow.get_trace(producer_trace_id)
producer_span = producer_trace.data.spans[0]
取得したスパンのIDでリンクを組み立てます。
upstream = Link(
trace_id=producer_span.trace_id,
span_id=producer_span.span_id,
attributes={"relation": "triggered_by", "note": "別トレースの検索処理"},
)
これを下流のトレースに渡します。デコレータの引数で渡す方法と、実行中のスパンに後付けする方法の両方を入れてみます。
@mlflow.trace(name="generator", links=[upstream])
def generate(docs: list[str]):
with mlflow.start_span("post_process") as child:
child.add_link(Link(trace_id=upstream.trace_id, span_id=upstream.span_id))
return "answer"
generate(["doc"])
consumer_trace_id = mlflow.get_last_active_trace_id()
mlflow.flush_trace_async_logging()
読み直すと、どちらの方法でもリンクが記録されていることが確認できます。
consumer_trace = mlflow.get_trace(consumer_trace_id)
for span in consumer_trace.data.spans:
print(span.name, [link.to_dict() for link in span.links])
generator [{'trace_id': 'tr-586ae293...', 'span_id': '1f755ef0c5f52234',
'attributes': {'relation': 'triggered_by', 'note': '別トレースの検索処理'}}]
post_process [{'trace_id': 'tr-586ae293...', 'span_id': '1f755ef0c5f52234',
'attributes': None}]
属性は日本語もそのまま入ります。add_link() で後付けしたほうは属性を渡していないので None です。
UIでリンクを確認する
MLflow UIを起動します。DBを作ったディレクトリで実行してください。
mlflow ui --backend-store-uri sqlite:///mlflow.db
span-link-demo エクスペリメントのTracesタブを開くと、トレースが2本並んでいます。下流の generator のトレースを開きます。
スパンツリーを見ると、generator と post_process の両方に鎖のアイコンと 1 のバッジが付いています。リンクを持っているスパンがツリーの時点で分かるので、深いトレースでもリンク付きのスパンを探せます。地味ですが、実際に使うときにはこれが効きます。
スパンを選ぶと、右ペインのタブに Links が増えています。ここも件数バッジ付きです。
Linksタブの中身は、リンク1件につき1枚のカードです。Trace IDとSpan IDが短縮表示でコピーボタン付き、その下にAttributesとして relation: triggered_by と note: 別トレースの検索処理 が展開されています。右上の「Jump to linked span」で、リンク先のスパンに移動できます。外部リンクのアイコンが付いているのは、別トレース宛なので新しいタブで開くためです。同一トレース内のリンクであれば、そのままツリー内でジャンプします。
post_process のほうも見てみます。
こちらは add_link() で属性なしで張ったリンクなので、Attributesが「No attributes」と表示されます。セクションごと消えるのではなく空状態が明示されるので、「リンクはあるが属性が無い」のか「そもそもリンクが無い」のかが区別できます。
このスパンだけタブが Attributes と Links の2つしかないのは、with mlflow.start_span("post_process") の中で set_inputs() / set_outputs() を呼んでいないためです。span linkとは関係ありませんが、タブ構成が動的に変わることは知っておくと混乱しません。
UIに出ないときの切り分け
もしUIでLinksタブが出ない場合、データ側の問題かフロント側の問題かを切り分けたくなります。UIが実際に叩いているのはこのエンドポイントです。
curl -s "http://127.0.0.1:5000/ajax-api/2.0/mlflow/get-trace-artifact?request_id=<trace_id>"
レスポンスの各スパンに links が入っていれば、データは記録できています。
{"spans": [{"name": "generator", "links": [{"trace_id": "tr-586ae293...",
"span_id": "1f755ef0c5f52234", "attributes": {"relation": "triggered_by"}}]}]}
スパンのシリアライズはOTLPのproto形式と相互変換できる作りになっているので、外部のOpenTelemetryコレクタにエクスポートした場合もリンクは引き継がれます。
Unity Catalogのトレースでは現時点でリンクは記録されない
ここは実務上いちばん重要な制約です。トレースの記録先がUnity Catalogスキーマの場合、span linkは記録されません。
リンクを張ろうとすると、警告が出てリンクだけがスキップされます。スパン自体は通常どおり記録されます。
Span links are not currently supported for Unity Catalog traces.
The link to trace 'trace:/...' will be skipped.
判定はトレースIDのプレフィックスだけで行われています。trace:/ で始まるv4形式のトレースIDが対象です。ソースコードのコメントによると、v4のトレースIDを生のOTel hexへ黙って正規化してしまうと不整合が起きるため、変換せずにスキップして警告を出す、という判断のようです。曖昧なまま記録されるより、落ちたことが分かるほうが後で追いやすいので、妥当な作りだと思います。
このv4形式になるのは、記録先をUnity Catalogスキーマにしたときです。具体的には次の3つのいずれかを設定している場合です。
mlflow.set_experiment(..., trace_location=mlflow.entities.UnityCatalog(...))- 環境変数
MLFLOW_TRACING_DESTINATION=<catalog>.<schema> -
mlflow.tracing.set_destination()にUnity Catalogのロケーションを指定
これらを設定していなければ通常のMLflowエクスペリメントが記録先になり、tr- で始まるトレースIDが振られて、リンクはそのまま保持されます。
もう1点、トレースを別のトレースにコピーする操作でも、リンクは引き継がれません。こちらも警告が出ます。
Span links on span '<span名>' will not be retained when copying to a new trace
span linkを前提に組むなら、記録先がどちらなのかを最初に確認しておくのが安全です。
まとめ
MLflowのspan linkを試して分かったことをまとめます。
- span linkは親子関係にないスパン同士を繋ぐ仕組み。トレースをまたげて、1つのスパンに何本でも張れる
-
Linkエンティティとspan.add_link()は3.13.0から、start_span(links=[...])と@mlflow.trace(links=[...])は3.15.0から使える - OSSのMLflow UIにもLinksタブとして反映される。スパンツリーにもリンク件数のバッジが出る
- 「Jump to linked span」で別トレースのスパンへ移動できる。リンクの
attributesもUIに表示される - 記録先がUnity Catalogスキーマの場合、リンクは警告付きでスキップされる。判定は
trace:/プレフィックスのみ - 軽量版の
mlflow-tracingにはトラッキングストアとUIが含まれないので、UIまで見るならフル版が必要 get_trace()の前にflush_trace_async_logging()を呼ばないと「span data is corrupted」で読めない
一番の収穫は、リリースノートに載っていない機能でもUIまで含めてきちんと作り込まれていた、という点でした。トレーシングまわりは変化が速いので、気になった機能はリリースノートを待たずに手元のwheelを覗いてみると発見があります。マルチエージェントやキュー越しの非同期処理をMLflowで追いかけている方は、一度試してみる価値があると思います。