Pelacakan MLflow untuk agen
Penting
Fitur ini ada di Pratinjau Publik.
Artikel ini menjelaskan Pelacakan MLflow pada Databricks dan bagaimana menggunakannya untuk menambahkan kemampuan observasi ke dalam aplikasi AI generatif Anda.
Apa itu Pelacakan MLflow?
MLflow Tracing menangkap informasi terperinci tentang pelaksanaan aplikasi AI generatif. Melacak input, output, dan metadata dalam log yang terkait dengan setiap langkah perantara dari sebuah permintaan sehingga Anda dapat mengidentifikasi sumber bug dan perilaku yang tak terduga. Misalnya, jika model Berhalusinasi, Anda dapat dengan cepat memeriksa setiap langkah yang menyebabkan halusinasi.
Pelacakan MLflow terintegrasi dengan alat dan infrastruktur Databricks, memungkinkan Anda menyimpan dan menampilkan jejak di notebook Databricks atau UI eksperimen MLflow.
Mengapa menggunakan Pelacakan MLflow?
Pelacakan MLflow memberikan beberapa manfaat:
- Tinjau visualisasi pelacakan interaktif dan gunakan alat investigasi untuk mendiagnosis masalah.
- Verifikasi bahwa templat dan pagar pembatas yang diminta menghasilkan hasil yang wajar.
- Analisis latensi berbagai kerangka kerja, model, dan ukuran gugus.
- Perkirakan biaya aplikasi dengan mengukur penggunaan token di berbagai model.
- Tetapkan himpunan data tolok ukur "emas" untuk mengevaluasi performa versi yang berbeda.
- Simpan jejak dari titik akhir model produksi untuk memecahkan masalah dan lakukan tinjauan serta evaluasi secara offline.
Menambahkan jejak ke agen Anda
MLflow Tracing mendukung tiga metode untuk menambahkan jejak ke aplikasi AI generatif Anda. Untuk detail referensi API, lihat dokumentasi MLflow.
| API | Kasus penggunaan yang direkomendasikan | Deskripsi |
|---|---|---|
| MLflow autologging | Pengembangan dengan pustaka GenAI terintegrasi | Pengelogan otomatis secara otomatis mencatat jejak untuk kerangka kerja sumber terbuka yang didukung seperti LangChain, LlamaIndex, dan OpenAI. |
| API Fasih |
Agen kustom dengan Pyfunc | 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. |
| API Klien MLflow | Kasus penggunaan tingkat lanjut seperti multi-utas |
MLflowClient menyediakan API granular dan aman terhadap utas untuk skenario penggunaan tingkat lanjut. Anda harus mengelola hubungan rentang induk-turunan secara manual. Ini memberi Anda kontrol yang lebih baik atas siklus hidup pelacakan, terutama untuk kasus penggunaan multi-utas. |
Menginstal Pelacakan MLflow
MLflow Tracing tersedia dalam MLflow versi 2.13.0 ke atas, yang sudah diinstal sebelumnya di <DBR< 15.4 LTS ML dan ke atas. Jika perlu, instal MLflow dengan kode berikut:
%pip install mlflow>=2.13.0 -qqqU
%restart_python
Atau, Anda dapat menginstal versi terbaru databricks-agents, yang mencakup versi MLflow yang kompatibel:
%pip install databricks-agents
Gunakan autologging untuk menambahkan pelacakan ke agen Anda
Jika pustaka GenAI Anda mendukung pelacakan, seperti LangChain atau OpenAI, aktifkan autologging dengan menambahkan mlflow.<library>.autolog() ke kode Anda. Misalnya:
mlflow.langchain.autolog()
Catatan
Pada Databricks Runtime 15.4 LTS ML, pelacakan MLflow diaktifkan secara default dalam notebook. Untuk menonaktifkan pelacakan, misalnya, dengan LangChain, Anda dapat menjalankan mlflow.langchain.autolog(log_traces=False) di buku catatan Anda.
MLflow mendukung pustaka tambahan untuk melacak pengelogan otomatis. Untuk daftar lengkap pustaka terintegrasi, lihat dokumentasi MLflow Tracking.
Gunakan Fluent API untuk menambahkan jejak secara manual ke agen Anda
API fasih di MLflow secara otomatis membuat hierarki pelacakan berdasarkan alur eksekusi kode Anda.
Menghias acara Anda
Gunakan dekorator @mlflow.trace untuk membuat rentang untuk lingkup fungsi yang didekorasi.
Objek Span MLflow mengorganisir langkah-langkah pelacakan. "Span menangkap informasi terkait operasi atau langkah-langkah individu, seperti panggilan API atau kueri penyimpanan vektor, dalam alur kerja."
Rentang dimulai ketika fungsi dipanggil dan berakhir ketika fungsi tersebut mengembalikan hasil. MLflow merekam input dan output fungsi serta setiap pengecualian yang dibangkitkan dari fungsi.
Misalnya, kode berikut membuat elemen bernama my_function yang menangkap argumen input x dan y, serta outputnya.
@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
return x + y
Menggunakan pengelola konteks pelacakan
Jika Anda ingin membuat rentang untuk blok kode arbitrer, bukan hanya fungsi, Anda dapat menggunakan mlflow.start_span() sebagai manajer konteks yang membungkus blok kode. Rentang dimulai ketika konteks dimasukkan dan berakhir saat konteks keluar. Rentang input dan output harus ditentukan secara manual dengan menggunakan metode setter dari objek rentang yang diperoleh melalui pengelola konteks.
with mlflow.start_span("my_span") as span:
span.set_inputs({"x": x, "y": y})
result = x + y
span.set_outputs(result)
span.set_attribute("key", "value")
Membungkus fungsi eksternal
Untuk menelusuri fungsi pustaka eksternal, bungkus fungsi tersebut dengan mlflow.trace.
from sklearn.metrics import accuracy_score
y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]
traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)
Contoh API fasih
Contoh berikut menunjukkan cara menggunakan api Fasih mlflow.trace dan mlflow.start_span untuk melacak quickstart-agent:
import mlflow
from mlflow.deployments import get_deploy_client
class QAChain(mlflow.pyfunc.PythonModel):
def __init__(self):
self.client = get_deploy_client("databricks")
@mlflow.trace(name="quickstart-agent")
def predict(self, model_input, system_prompt, params):
messages = [
{
"role": "system",
"content": system_prompt,
},
{
"role": "user",
"content": model_input[0]["query"]
}
]
traced_predict = mlflow.trace(self.client.predict)
output = traced_predict(
endpoint=params["model_name"],
inputs={
"temperature": params["temperature"],
"max_tokens": params["max_tokens"],
"messages": messages,
},
)
with mlflow.start_span(name="_final_answer") as span:
# Initiate another span generation
span.set_inputs({"query": model_input[0]["query"]})
answer = output["choices"][0]["message"]["content"]
span.set_outputs({"generated_text": answer})
# Attributes computed at runtime can be set using the set_attributes() method.
span.set_attributes({
"model_name": params["model_name"],
"prompt_tokens": output["usage"]["prompt_tokens"],
"completion_tokens": output["usage"]["completion_tokens"],
"total_tokens": output["usage"]["total_tokens"]
})
return answer
Setelah menambahkan jejak, jalankan fungsi . Berikut ini melanjutkan contoh dengan fungsi predict() di bagian sebelumnya. Jejak secara otomatis ditampilkan saat Anda menjalankan metode pemanggilan, predict().
SYSTEM_PROMPT = """
You are an assistant for Databricks users. You answer Python, coding, SQL, data engineering, spark, data science, DW and platform, API, or infrastructure administration questions related to Databricks. If the question is unrelated to one of these topics, kindly decline to answer. If you don't know the answer, say that you don't know; don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""
model = QAChain()
prediction = model.predict(
[
{"query": "What is in MLflow 5.0"},
],
SYSTEM_PROMPT,
{
# Using Databricks Foundation Model for easier testing, feel free to replace it.
"model_name": "databricks-dbrx-instruct",
"temperature": 0.1,
"max_tokens": 1000,
}
)
API Klien MLflow
MlflowClient memaparkan API granular, aman utas untuk memulai dan mengakhiri jejak, mengelola rentang, dan mengatur bidang rentang. Ini memberikan kontrol penuh atas siklus hidup dan struktur jejak. API ini berguna jika API Fluent tidak mencukupi untuk kebutuhan Anda, seperti aplikasi multithread dan callback.
Berikut ini adalah langkah-langkah untuk membuat jejak lengkap menggunakan Klien MLflow.
Buat instans MLflowClient dengan
client = MlflowClient().Mulai pelacakan menggunakan metode
client.start_trace(). Ini memulai konteks pelacakan, memulai rentang akar absolut, dan mengembalikan objek rentang akar. Metode ini harus dijalankan sebelum APIstart_span().- Tetapkan atribut, input, dan output Anda untuk jejak pemrosesan di
client.start_trace().
Catatan
Tidak ada yang setara dengan metode
start_trace()dalam API Fasih. Ini karena Fluent API secara otomatis menginisialisasi konteks pelacakan dan menentukan apakah itu adalah span akar berdasarkan status yang dikelola.- Tetapkan atribut, input, dan output Anda untuk jejak pemrosesan di
API start_trace() mengembalikan rentang. Dapatkan ID permintaan, pengidentifikasi unik penelusuran yang juga disebut sebagai
trace_id, dan ID rentang yang diperoleh menggunakanspan.request_iddanspan.span_id.Mulai span anak menggunakan
client.start_span(request_id, parent_id=span_id)untuk mengatur atribut, input, dan output Anda untuk span tersebut.- Metode ini memerlukan
request_iddanparent_iduntuk mengaitkan rentang dengan posisi yang benar dalam hierarki jejak. Ini mengembalikan objek rentang lain.
- Metode ini memerlukan
Akhiri span anak dengan memanggil
client.end_span(request_id, span_id).Ulangi Langkah 3 - 5 untuk rentang anak apa pun yang ingin Anda buat.
Setelah semua rentang anak berakhir, panggil
client.end_trace(request_id)untuk menutup jejak dan merekamnya.
from mlflow.client import MlflowClient
mlflow_client = MlflowClient()
root_span = mlflow_client.start_trace(
name="simple-rag-agent",
inputs={
"query": "Demo",
"model_name": "DBRX",
"temperature": 0,
"max_tokens": 200
}
)
request_id = root_span.request_id
# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)
span_ss = mlflow_client.start_span(
"search",
# Specify request_id and parent_id to create the span at the right position in the trace
request_id=request_id,
parent_id=root_span.span_id,
inputs=similarity_search_input
)
retrieved = ["Test Result"]
# You must explicitly end the span
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)
root_span.end_trace(request_id, outputs={"output": retrieved})
Meninjau jejak
Untuk meninjau jejak setelah menjalankan agen, Anda dapat menggunakan salah satu opsi berikut:
- Visualisasi jejak dirender sebaris dalam output sel.
- Jejak dicatat ke eksperimen MLflow Anda. anda dapat meninjau dan mencari daftar lengkap jejak historis di tab Jejak di halaman Eksperimen. Saat agen berjalan di bawah MLflow Run aktif, jejak muncul di halaman Run.
- Ambil jejak secara terprogram menggunakan API search_traces().
Menggunakan Pelacakan MLflow dalam produksi
MLflow Tracing juga terintegrasi dengan Mosaic AI Model Serving, memungkinkan Anda untuk men-debug masalah secara efisien, memantau performa, dan membuat himpunan data emas untuk evaluasi offline. Saat Pelacakan MLflow diaktifkan untuk titik akhir penyajian Anda, jejak direkam dalam tabel inferensi di bawah kolom response.
Anda bisa memvisualisasikan jejak yang dicatat ke tabel inferensi dengan mengkueri tabel dan menampilkan hasilnya dalam buku catatan. Gunakan display(<the request logs table>) di buku catatan Anda dan pilih baris individual dari jejak yang ingin Anda visualisasikan.
Untuk mengaktifkan Pelacakan MLflow untuk titik akhir penyajian, Anda harus mengatur variabel lingkungan ENABLE_MLFLOW_TRACING dalam konfigurasi titik akhir ke True. Untuk mempelajari cara menyebarkan titik akhir dengan variabel lingkungan kustom, lihat Menambahkan variabel lingkungan teks biasa. Jika Anda menyebarkan agen menggunakan API deploy(), jejak secara otomatis dicatat ke tabel inferensi. Lihat Menyebarkan agen untuk aplikasi AI generatif.
Catatan
Menulis jejak ke tabel inferensi dilakukan secara asinkron, sehingga tidak menambahkan overhead yang sama seperti di lingkungan notebook selama pengembangan. Namun, mungkin masih memperkenalkan beberapa overhead ke kecepatan respons titik akhir, terutama ketika ukuran jejak untuk setiap permintaan inferensi besar. Databricks tidak menjamin perjanjian tingkat layanan (SLA) untuk dampak latensi aktual pada titik akhir model Anda, karena sangat tergantung pada lingkungan dan implementasi model. Databricks merekomendasikan pengujian performa titik akhir Anda dan mendapatkan wawasan tentang overhead pelacakan sebelum menyebarkan ke aplikasi produksi.
Tabel berikut memberikan indikasi kasar tentang dampak pada latensi inferensi untuk ukuran jejak yang berbeda.
| Ukuran pelacakan per permintaan | Dampak ke latensi (ms) |
|---|---|
| ~10 KB | ~ 1 mdtk |
| ~ 1 MB | 50 ~ 100 mdtk |
| 10 MB | 150 ms ~ |
Batasan
- Pelacakan MLflow tersedia di notebook Databricks, pekerjaan notebook, dan Model Serving.
Autologging LangChain mungkin tidak mendukung semua API prediksi LangChain. Untuk daftar lengkap API yang didukung, lihat dokumentasi MLflow.