メインコンテンツまでスキップ

コンテキストマネージャーによるスパントレース

mlflow.start_spanコンテキスト マネージャーを使用すると、任意のコード ブロックのスパンを作成できます。これは、単一の関数の境界をキャプチャすることによって可能なものよりも詳細に、コード内の複雑な相互作用をキャプチャするのに役立ちます。

コンテキスト マネージャーを使用したスパン トレースにより、トレースされるコードをきめ細かく制御できます。

  • 任意のコードブロック :関数全体だけでなく、任意のコードブロックをトレースします。
  • 柔軟な境界 :スパンの開始点と終了点を正確に定義します
  • 自動コンテキスト管理 : MLflowは親子関係とクリーンアップを処理します
  • 関数デコレータと連携 : ハイブリッドアプローチのために@mlflow.traceと組み合わせて使用します
  • 例外処理 : デコレータのような自動エラーキャプチャ

前提条件

このページには次のパッケージが必要です:

  • mlflow[databricks] 3.1 以上: GenAI 機能と Databricks 接続を備えたコア MLflow 機能。
  • openai 1.0.0 以上: (オプション) カスタム コードが OpenAI と対話する場合のみ。必要に応じて他の SDK に置き換えます。

基本要件をインストールします。

Python
%pip install --upgrade "mlflow[databricks]>=3.1"
# %pip install --upgrade openai>=1.0.0 # Install if needed

コンテキストマネージャーAPI

デコレータと同様に、コンテキスト マネージャーは親子関係、例外、実行時間を自動的にキャプチャし、自動トレースで動作します。ただし、スパンの名前、入力、および出力は手動で提供する必要があります。コンテキスト マネージャーから返されるmlflow.entities.Spanオブジェクトを使用して設定できます。

Python
with mlflow.start_span(name="my_span") as span:
span.set_inputs({"x": 1, "y": 2})
z = x + y
span.set_outputs(z)

以下は、 mlflow.start_spanコンテキスト マネージャーを OpenAI のデコレータと自動トレースの両方と組み合わせて使用する、もう少し複雑な例です。

Python
import mlflow
import openai
from mlflow.entities import SpanType

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Create OpenAI client
client = openai.OpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def start_session():
messages = [{"role": "system", "content": "You are a friendly chat bot"}]
while True:
with mlflow.start_span(name="User") as span:
span.set_inputs(messages)
user_input = input(">> ")
span.set_outputs(user_input)

if user_input == "BYE":
break

messages.append({"role": "user", "content": user_input})

response = client.chat.completions.create(
model="gpt-4o-mini",
max_tokens=100,
messages=messages,
)
answer = response.choices[0].message.content
print(f"Assistant: {answer}")

messages.append({"role": "assistant", "content": answer})


start_session()

次のステップ