Pemantauan agen dengan Pelacakan MLflow

Artikel ini menjelaskan cara menambahkan pengamatan ke aplikasi AI generatif Anda dengan MLflow Tracing on Databricks.

Apa itu Pelacakan MLflow?

MLflow Tracing memberikan pengamatan end-to-end untuk aplikasi AI generatif dari pengembangan hingga penyebaran. Pelacakan sepenuhnya terintegrasi dengan toolset AI gen Databricks, menangkap wawasan terperinci di seluruh siklus hidup pengembangan dan produksi.

Berikut ini adalah kasus penggunaan utama untuk pelacakan dalam aplikasi AI generatif:

• Penelusuran kesalahan yang disederhanakan: Pelacakan memberikan pandangan ke setiap langkah aplikasi kecerdasan buatan generatif Anda, sehingga mempermudah diagnosis dan penyelesaian masalah.

• Evaluasi Offline: Pelacakan menghasilkan data berharga untuk evaluasi agen, memungkinkan Anda mengukur dan meningkatkan kualitas agen seiring waktu.

• Pemantauan produksi: Pelacakan memberikan visibilitas ke dalam perilaku agen dan langkah-langkah eksekusi terperinci, memungkinkan Anda memantau dan mengoptimalkan performa agen dalam produksi.

• Log audit: Pelacakan MLflow menghasilkan log audit yang komprehensif mengenai tindakan dan keputusan agen. Ini sangat penting untuk memastikan kepatuhan dan mendukung penelusuran kesalahan ketika masalah tak terduga muncul.

Persyaratan

Pelacakan MLflow tersedia pada MLflow versi 2.13.0 ke atas. Databricks merekomendasikan penginstalan MLflow versi terbaru untuk mengakses fitur dan peningkatan terbaru.

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Pelacakan otomatis

autologging MLflow memungkinkan Anda dengan cepat melengkapi agen Anda dengan menambahkan satu baris ke kode Anda, mlflow.<library>.autolog().

MLflow mendukung autologging untuk pustaka pengembangan agen paling populer. Untuk informasi selengkapnya tentang setiap pustaka penulisan, lihat dokumentasi MLflow autologging:

Nonaktifkan pembuatan log otomatis

Pelacakan pencatatan otomatis diaktifkan secara default di Databricks Runtime 15.4 ML dan versi di atasnya untuk perpustakaan berikut:

• LangChain
• Langgraph
• OpenAI
• LlamaIndex

Untuk menonaktifkan pelacakan pengelogan otomatis untuk pustaka ini, jalankan perintah berikut ini di buku catatan:

`mlflow.<library>.autolog(log_traces=False)`

Menambahkan jejak secara manual

Meskipun autologging menyediakan cara yang nyaman untuk menginstrumentasi agen, Anda mungkin ingin menginstrumentasi agen Anda secara lebih terperinci atau menambahkan trace tambahan yang tidak ditangkap oleh autologging. Dalam kasus ini, gunakan API Pelacakan MLflow untuk menambahkan jejak secara manual.

API Pelacakan MLflow adalah API kode rendah untuk menambahkan jejak tanpa khawatir mengelola struktur pohon pelacakan. MLflow menentukan hubungan rentang induk-anak yang sesuai secara otomatis dengan memanfaatkan stack Python.

Menggabungkan autologging dan pelacakan manual

API pelacakan manual dapat digunakan dengan autologging. MLflow menggabungkan rentang yang dibuat oleh autologging dan pelacakan manual untuk membuat jejak lengkap eksekusi agen Anda. Untuk contoh menggabungkan autologging dan pelacakan manual, lihat Instrumenting agen panggilan alat dengan MLflow Tracing.

Melacak fungsi menggunakan dekorator @mlflow.trace

Cara paling sederhana untuk melengkapi kode Anda secara manual adalah dengan mendekorasi fungsi dengan dekorator @mlflow.trace. Dekorator jejak MLflow membuat "span" dengan cakupan fungsi yang didekorasi, yang mewakili suatu satuan eksekusi dalam jejak dan ditampilkan sebagai satu baris dalam visualisasi jejak. Rentang ini menangkap input dan output fungsi, latensi, dan pengecualian apa pun yang dimunculkan dari fungsi.

Misalnya, kode berikut membuat elemen bernama my_function yang menangkap argumen input x dan y, serta outputnya.

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

Anda juga dapat menyesuaikan nama rentang, jenis rentang, dan menambahkan atribut kustom ke rentang:

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

Melacak blok kode arbitrer menggunakan manajer konteks

Untuk membuat rentang untuk blok kode arbitrer, bukan hanya fungsi, gunakan mlflow.start_span() sebagai manajer konteks yang membungkus blok kode. Rentang dimulai saat konteks dimasuki dan berakhir saat konteks ditinggalkan. Rentang input dan output harus ditentukan secara manual dengan menggunakan metode setter dari objek rentang yang diperoleh melalui pengelola konteks. Untuk informasi selengkapnya, lihat dokumentasi MLflow - penangan konteks.

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

Pustaka pelacakan tingkat bawah

MLflow juga menyediakan API tingkat rendah untuk mengontrol struktur pohon jejak secara eksplisit. Lihat dokumentasi MLflow - Instrumentasi Manual.

Contoh pelacakan: Menggabungkan pencatatan otomatis dan jejak manual

Contoh berikut menggabungkan autologging OpenAI dan pelacakan manual untuk sepenuhnya menginstrumen agen pemanggil alat.

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

Membuat anotasi jejak dengan tag

tag pelacakan MLflow adalah pasangan kunci-nilai yang memungkinkan Anda menambahkan metadata kustom ke pelacakan, seperti ID percakapan, ID pengguna, hash komit Git, dll. Tag ditampilkan di UI MLflow untuk memfilter dan mencari pelacakan.

Tag dapat dikonfigurasi untuk pelacakan yang sedang berjalan atau yang telah selesai menggunakan API MLflow atau UI MLflow. Contoh berikut menunjukkan penambahan tag ke pelacakan yang sedang berlangsung menggunakan API mlflow.update_current_trace().

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

Untuk mempelajari selengkapnya tentang menandai jejak dan cara menggunakannya untuk memfilter dan mencari jejak, lihat dokumentasi MLflow - Mengatur Tag Pelacakan.

Tinjau jejak

Untuk meninjau jejak setelah menjalankan agen, Anda dapat menggunakan salah satu opsi berikut:

• Visualisasi sebaris: Di buku catatan Databricks, jejak dirender sebaris dalam output sel.
• eksperimen MLflow: Di Databricks, buka Eksperimen > Pilih eksperimen >Jejak untuk melihat dan mencari semua jejak untuk eksperimen.
• MLflow Run: Saat agen berjalan di bawah MLflow Run yang aktif, jejak muncul di halaman Run dari UI MLflow.
• UI Evaluasi Agen : Dalam Evaluasi Agen AI Mosaik, Anda dapat meninjau jejak untuk setiap eksekusi agen dengan mengklik Lihat tampilan jejak terperinci dalam hasil evaluasi.
• Trace Search API: Untuk mengambil jejak secara terprogram, gunakan Trace Search API.

Mengevaluasi agen dengan menggunakan jejak

Data pelacakan berfungsi sebagai sumber daya yang berharga untuk mengevaluasi agen Anda. Dengan menangkap informasi terperinci tentang eksekusi model Anda, MLflow Tracing berperan penting dalam evaluasi offline. Anda dapat menggunakan data pelacakan untuk mengevaluasi performa agen Anda terhadap himpunan data emas, mengidentifikasi masalah, dan meningkatkan performa agen Anda.

%pip install -U mlflow databricks-agents
%restart_python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

Untuk mempelajari selengkapnya tentang evaluasi agen, lihat Menjalankan evaluasi dan melihat hasilnya.

Memantau agen yang telah diaktifkan menggunakan tabel inferensi

Setelah agen disebarkan ke Mosaic AI Model Serving, Anda dapat menggunakan tabel inferensi untuk memantau agen. Tabel inferensi berisi log terperinci permintaan, respons, jejak agen, dan umpan balik agen dari aplikasi ulasan. Informasi ini memungkinkan Anda men-debug masalah, memantau performa, dan membuat himpunan data emas untuk evaluasi offline.

Untuk mengaktifkan tabel inferensi untuk penyebaran agen, lihat Mengaktifkan tabel inferensi untuk agen AI.

Mengkueri jejak online

Gunakan buku catatan untuk mengkueri tabel inferensi dan menganalisis hasilnya.

Untuk memvisualisasikan jejak, jalankan display(<the request logs table>) dan pilih baris untuk diperiksa:

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

Memantau agen

Lihat Cara memantau aplikasi AI generatif Anda.

Overhead latensi pelacakan

Jejak ditulis secara asinkron untuk meminimalkan dampak kinerja. Namun, pelacakan masih menambahkan latensi ke kecepatan respons titik akhir, terutama ketika ukuran jejak untuk setiap permintaan inferensi besar. Databricks merekomendasikan pengujian titik akhir Anda untuk memahami dampak latensi penelusuran sebelum menerapkan ke lingkungan produksi.

Tabel berikut ini memberikan perkiraan kasar untuk dampak latensi berdasarkan ukuran jejak:

Pemecahan masalah

Untuk pemecahan masalah dan pertanyaan umum, lihat dokumentasi MLflow: Panduan Cara Melacak dan dokumentasi MLflow : Tanya Jawab Umum