トレースとプロファイリング#

OMMXは主要なエントリポイントから OpenTelemetry のスパンを出力します。これを人間が読める形にまとめるための薄いラッパーが ommx.tracing に2つ用意されています。

  • %%ommx_trace — 単一セルの実行中に発生したスパンをネストしたテキストツリーとして表示するJupyterセルマジック。加えてChrome Trace Event Format形式のJSONダウンロードリンクも表示されます。

  • capture_trace / traced() — 同じ仕組みを通常のPythonスクリプト、テスト、CIから利用するためのコンテキストマネージャとデコレータ。

両者は1つのインプロセスコレクタを共有しています。import時に何かを設定する必要はなく、OTelエクスポータをインストールする必要もありません。コレクタは最初の利用時に遅延インストールされます。OTelバックエンドへトレースを送信したい場合のみ、後述の 外部OTelバックエンドへ送信する セクションを参照してください。

クイックツアー#

セルマジック (%%ommx_trace)#

ノートブックごとに一度、拡張を読み込みます(通常は最初のセル)。

%load_ext ommx.tracing

以降、任意のセルの先頭に %%ommx_trace を付けるだけで計測されます。

%%ommx_trace
from ommx.v1 import Instance, DecisionVariable

x = DecisionVariable.binary(0, name="x")
y = DecisionVariable.binary(1, name="y")
instance = Instance.from_components(
    decision_variables=[x, y],
    objective=x + y,
    constraints={},
    sense=Instance.MAXIMIZE,
)
solution = instance.evaluate({0: 1.0, 1: 1.0})
└── ommx_trace_cell (14.32 ms) {scope=ommx.tracing}
    └── evaluate (1.53 ms) {scope=ommx}

Download Chrome Trace JSON (0.6 KB) — open in Perfetto, speedscope, or chrome://tracing.

セルの出力として以下の2つが表示されます。

  1. セル内で発生した全スパンをネストしたテキストツリーとして描画。各ノードには持続時間と代表的な属性が付与されます。

  2. Chrome Trace Event Format形式のトレース全体のダウンロードリンク。生成されたJSONファイルを Perfetto / speedscope / chrome://tracing に読み込ませることでフレームグラフとして閲覧できます。

注釈

このセル出力(テキストツリー+ダウンロードリンク)は最小限の開始点として提供しているものであり、今後改善される予定です(例: セル内へのインタラクティブなフレームグラフの埋め込みなど)。レイアウトやマークアップの細部は不安定として扱ってください。

セル内で例外が発生した場合も、トレースのHTMLは通常どおり描画されます(失敗したスパンには [ERROR] マーカーが付きます)。そしてその後例外は再送出されるので、nbconvert --execute、papermill、pytest-nbval などのノートブック自動化ツールから見てもセルは失敗扱いとなります。

コンテキストマネージャ (capture_trace)#

通常のPythonスクリプトからも同じ仕組みを利用できます。

from ommx.tracing import capture_trace, save_chrome_trace
from ommx.v1 import Instance, DecisionVariable

x = DecisionVariable.binary(0, name="x")
y = DecisionVariable.binary(1, name="y")

instance = Instance.from_components(
    decision_variables=[x, y],
    objective=x + y,
    constraints={},
    sense=Instance.MAXIMIZE,
)

with capture_trace() as trace:
    solution = instance.evaluate({0: 1.0, 1: 1.0})

trace
└── ommx_trace_block (242.0 µs) {scope=ommx.tracing}
    └── evaluate (9.0 µs) {scope=ommx}

Download Chrome Trace JSON (0.5 KB) — open in Perfetto, speedscope, or chrome://tracing.

trace はブロック終了時に値が埋められる TraceResult です。

  • request — 結果としてexportされたOTLP ExportTraceServiceRequest

  • spansrequest 内のexport済みOTLP protobuf spanを平坦化したリスト

  • otlp_protobuf() — Experimentのtraceに保存されるOTLP export requestのprotobuf bytes

Notebookでは trace を評価すると、__repr__ によって同じテキストツリーが表示されます。明示的な表示や保存には次の関数も使えます。

ブロック内で例外が発生した場合でも trace.request は埋められており(rendererでは失敗したスパンに [ERROR] マーカーが付く)、外側の exceptfinally から内容を調査・保存できます。元の例外はそのまま伝播します。OMMXが例外を握り潰すことはありません。

import tempfile
from pathlib import Path

output_path = Path(tempfile.gettempdir()) / "ommx_trace.json"
save_chrome_trace(trace, output_path)

デコレータ (@traced)#

traced()capture_trace の糖衣構文です。

import tempfile
from pathlib import Path

from ommx.tracing import traced

evaluate_output = Path(tempfile.gettempdir()) / "evaluate_trace.json"

@traced(output=str(evaluate_output))
def evaluate_once(inst):
    return inst.evaluate({0: 1.0, 1: 1.0})

solution = evaluate_once(instance)

3つの呼び出し形式すべてがサポートされています。

@traced
def f(): ...

@traced()
def f(): ...

@traced(name="build_qubo", output="qubo.json")
def f(): ...

主な挙動:

  • name を省略するとルートスパン名は fn.__qualname__ になります。複数のデコレート関数のトレースを区別しやすくするためです。

  • output を指定した場合、正常終了時はChrome Trace JSONが書き出されます。例外発生時も書き出しを試みますが、保存時のI/Oエラーなどは元の例外を上書きしないよう抑制されるため、例外パスでの保存は best-effort です。

  • async def もサポートされています。inspect.iscoroutinefunction でコルーチン関数を検知し、トレースブロック内で await します。この検知がないと、コルーチン生成直後にキャプチャウィンドウが閉じてしまい、スパンが全て失われます。

ExperimentにRunトレースを保存する#

ここまでのAPIは、トレースをその場で確認したり、単体のChrome Trace JSONとして保存したりするためのものです。実験記録と一緒にトレースも共有したい場合は、Experiment でtrace storageを有効にします。

from ommx.experiment import Experiment
from ommx_highs_adapter import OMMXHighsAdapter

with Experiment.with_temp_local_registry(store_trace=True) as experiment:
    with experiment.run() as run:
        run.log_parameter("solver", "highs")
        solution = run.log_solve(OMMXHighsAdapter, instance, verbose=False)
        run.log_parameter("objective", solution.objective)

trace = experiment.runs[0].trace
assert trace is not None
trace
└── Run (18.87 ms) {scope=ommx.experiment} [ommx.run.id=0]
    └── solve (6.26 ms) {scope=ommx.adapter.highs} [adapter='ommx_highs_adapter.adapter.OMMXHighsAdapter']
        ├── convert (5.90 ms) {scope=ommx.adapter.highs}
        │   └── reduce_capabilities (6.0 µs) {scope=ommx}
        ├── call (88.0 µs) {scope=ommx.adapter.highs}
        └── decode (196.0 µs) {scope=ommx.adapter.highs}
            └── evaluate (17.0 µs) {scope=ommx}

Download Chrome Trace JSON (1.6 KB) — open in Perfetto, speedscope, or chrome://tracing.

store_trace=True を指定すると、そのExperimentから作られる各 Run コンテキストマネージャの周囲に Run という名前の capture_trace ブロックが開かれます。このスパンが表すのは experiment.run() メソッド呼び出し自体の所要時間ではなく、with experiment.run() as run: が作る Run のライフタイムです。Runが正常終了すると、完了した TraceResult がOTLP protobuf bytesとしてExperiment Artifactに保存されます。あとからExperimentを読み込むと、同じデータを trace から取得できます。

loaded = Experiment.from_artifact(experiment.artifact)
trace = loaded.runs[0].trace
assert trace is not None
trace
└── Run (18.87 ms) {scope=ommx.experiment} [ommx.run.id=0]
    └── solve (6.26 ms) {scope=ommx.adapter.highs} [adapter='ommx_highs_adapter.adapter.OMMXHighsAdapter']
        ├── convert (5.90 ms) {scope=ommx.adapter.highs}
        │   └── reduce_capabilities (6.0 µs) {scope=ommx}
        ├── call (88.0 µs) {scope=ommx.adapter.highs}
        └── decode (196.0 µs) {scope=ommx.adapter.highs}
            └── evaluate (17.0 µs) {scope=ommx}

Download Chrome Trace JSON (1.6 KB) — open in Perfetto, speedscope, or chrome://tracing.

保存されるトレースの範囲は意図的にRun単位に限定されています。

  • デフォルトでは保存されません。store_trace=True を指定しない場合でも、アクティブなOpenTelemetry providerがあればOMMXのスパンはそこへ流れますが、Artifactにはtrace payloadを保存しません。

  • 保存境界は with experiment.run() as run: です。Experiment全体の寿命、手動の commit()、ノートブックのセル間の待ち時間は保存対象に含めません。

  • store_trace=True では Run をコンテキストマネージャとして使う必要があります。Run境界が閉じた時点でトレースを確定して保存するため、このモードではRunを作ってから手動で finish() する使い方は拒否されます。

  • commit済みExperimentをforkすると、既存Runのトレースはfork先にも引き継がれます。fork先で新しく追加するRunにもトレースを保存したい場合は、fork() にも store_trace=True を渡してください。

テキストツリーはトレースの要約です。各行にはスパン名、所要時間、instrumentation scope、主要な属性を表示します。たとえば Runwith experiment.run() as run: のライフタイム、solve はadapter API呼び出し、convert / call / decode はadapter内部の段階です。{scope=...} はそのスパンを発生させたpackageまたはmoduleを表し、adapterommx.run.id などの属性はRunやバックエンドを識別するための情報です。

外部OTelバックエンドへ送信する#

Run trace storageは、外部OTelバックエンドへのexport設定とは独立しています。OTLP、Jaeger、Honeycombなどへスパンを送信したい場合は、OMMXへの最初の呼び出し前にSDK TracerProviderを設定し、そこへexporterを追加してください。

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter

provider = TracerProvider()
provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))
trace.set_tracer_provider(provider)

from ommx.v1 import Instance

%%ommx_tracecapture_trace、Experimentのtrace storageは、アクティブなSDK providerにcollectorを追加します。このため同じスパンをNotebook上で表示しつつ、外部バックエンドにも送信できます。providerの設定は、スクリプトやNotebook kernelの先頭で、OMMXをimportまたは呼び出す前に行ってください。