Kom igång: MLflow-spårning för GenAI (Databricks Notebook)

Öppna notebook-versionen av den här sidan

Den här snabbstarten hjälper dig att integrera din GenAI-app med MLflow Tracing om du använder en Databricks-notebook-fil som utvecklingsmiljö. Om du använder en lokal IDE använder du snabbstarten för IDE i stället.

I slutet av den här självstudien har du:

• En Databricks-notebook-fil med ett länkat MLflow-experiment för din GenAI-app
• Ett enkelt GenAI-program som instrumenterats med MLflow Tracing
• En spårning från appen i ditt MLflow-experiment

Miljöinställningar

1. Skapa en ny notebook i din Databricks-arbetsyta. Notebook-filen kommer att ha ett MLflow-standardexperiment, vilket är standardbehållaren för din GenAI-applikation. Läs mer om MLflow-experiment i avsnittet MLflow-begrepp.

2. Installera de paket som krävs:

   • mlflow[databricks]: Använd den senaste versionen av MLflow för att få fler funktioner och förbättringar.
   • openai: I den här självstudien används OpenAI API-klienten för att anropa modeller som finns på Databricks.

%pip install -qq --upgrade "mlflow[databricks]>=3.1.0" openai
dbutils.library.restartPython()

Steg 1: Instrumentera ditt program med spårning

Kodfragmenten nedan definierar en enkel GenAI-app som slutför meningsmallar med hjälp av en LLM.

Först, skapa en klient för att ansluta till OpenAI-modeller som är värdbaserade på Databricks.

from databricks.sdk import WorkspaceClient

# Create an OpenAI client that is connected to Databricks-hosted LLMs
w = WorkspaceClient()
client = w.serving_endpoints.get_open_ai_client()

model_name = "databricks-claude-sonnet-4"

Du kan också använda OpenAI SDK för att ansluta till OpenAI-hostade modeller:

import openai

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured

client = openai.OpenAI()

model_name = "gpt-4o-mini"

För det andra definierar och kör du programmet. Instrumentering av appen med spårning använder helt enkelt:

• mlflow.openai.autolog(): Automatisk instrumentering för att samla in information om anropet till OpenAI SDK
• @mlflow.trace: Dekoratör som gör det enkelt att spåra alla Python-funktioner

import mlflow
import os

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

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

# Use the trace decorator to capture the application's entry point
@mlflow.trace
def my_app(input: str):
  # This call is automatically instrumented by `mlflow.openai.autolog()`
  response = client.chat.completions.create(
    model=model_name,
    temperature=0.1,
    max_tokens=200,
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant.",
      },
      {
        "role": "user",
        "content": input,
      },
    ]
  )
  return response.choices[0].message.content

result = my_app(input="What is MLflow?")
print(result)

Mer information om hur du lägger till spårning i appar finns i guiden för spårningsinstrumentation och 20+-biblioteksintegreringar.

Steg 2: Visa spårningen i MLflow

Spårningen visas under notebook-cellen.

Du kan också gå till användargränssnittet för MLflow-experimentet för att se spårningen:

1. Klicka på I det högra sidofältet.
2. Klicka på ikonen Bredvid experimentkörningar.
3. Den genererade spårningen visas på fliken Spårningar .
4. Klicka på spårningen för att visa dess information.

Förstå spårningen

Spårningen som du nyss skapade visar:

• Rotintervall: Representerar indata till my_app(...) funktionen
  • Child span: Representerar en begäran om färdigställande från OpenAI
• Attribut: Innehåller metadata som modellnamn, tokenantal och tidsinformation
• Indata: De meddelanden som skickas till modellen
• Utdata: Svaret som togs emot från modellen

Den här enkla spårningen ger redan värdefulla insikter om programmets beteende, till exempel:

• Det som efterfrågades
• Vilket svar genererades
• Hur lång tid begäran tog
• Hur många token som användes (vilket påverkar kostnaden)

För mer komplexa program som RAG-system eller multistegsagenter ger MLflow Tracing ännu mer värde genom att avslöja det inre arbetet för varje komponent och steg.

Nästa steg

• MLflow Tracing-guide – Börja här för mer djupgående utbildning om MLflow Tracing
• MLflow Tracing-integreringar – över 20 bibliotek med automatisk spårningsintegrering
• Spårningsbegrepp – Förstå grunderna i MLflow Tracing

Exempelanteckningsbok

Kom igång: MLflow-spårning för GenAI (Databricks Notebook)

Hämta anteckningsbok