Comience: Trazado de MLflow para GenAI (Cuaderno de Databricks)

Abrir la versión 'notebook' de esta página

Este inicio rápido le ayuda a integrar la aplicación genAI con el seguimiento de MLflow si usa un cuaderno de Databricks como entorno de desarrollo. Si usa un IDE local, use el inicio rápido del IDE en su lugar.

Al final de este tutorial, tendrá:

• Un cuaderno de Databricks con un experimento de MLflow vinculado para la aplicación GenAI
• Una sencilla aplicación GenAI instrumentada con el seguimiento de MLflow
• Seguimiento de esa aplicación en el experimento de MLflow

Configuración del entorno

1. Cree un cuaderno en el área de trabajo de Databricks. El cuaderno contará con un experimento predeterminado de MLflow que actuará como el contenedor de tu aplicación GenAI. Obtenga más información sobre los experimentos de MLflow en la sección Conceptos de MLflow.

2. Instale los paquetes necesarios:

   • mlflow[databricks]: use la versión más reciente de MLflow para obtener más características y mejoras.
   • openai: este tutorial usará el cliente de api de OpenAI para llamar a modelos hospedados en Databricks.

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

Paso 1: Instrumentar la aplicación con rastreo

Los fragmentos de código siguientes definen una aplicación genAI sencilla que completa las plantillas de oración mediante un LLM.

En primer lugar, cree un cliente de OpenAI para conectarse a los modelos de base hospedados en 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"

Como alternativa, puede usar el SDK de OpenAI para conectarse a modelos hospedados en OpenAI:

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"

En segundo lugar, defina y ejecute la aplicación. Instrumentar la aplicación con seguimiento se realiza utilizando:

• mlflow.openai.autolog(): instrumentación automática para capturar los detalles de la invocación al SDK de OpenAI
• @mlflow.trace: Decorador que facilita el seguimiento de cualquier función de Python.

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)

Para más información sobre cómo agregar seguimiento a las aplicaciones, consulte la guía de instrumentación de seguimiento y las más de 20 integraciones de bibliotecas.

Paso 2: Visualización del seguimiento en MLflow

La traza aparecerá por debajo de la celda del notebook.

Opcionalmente, puede ir a la interfaz de usuario del experimento de MLflow para ver el seguimiento:

1. Haga clic en el icono en la barra lateral derecha.
2. Haga clic en el icono de apertura junto a las ejecuciones de los experimentos.
3. El seguimiento generado aparece en la pestaña Seguimientos .
4. Haga clic en la traza para ver sus detalles.

Descripción del seguimiento

El seguimiento que acaba de crear muestra:

• Intervalo raíz: representa las entradas de la my_app(...) función.
  • Span secundario: representa la solicitud de finalización de OpenAI.
• Atributos: contiene metadatos como el nombre del modelo, los recuentos de tokens y la información de tiempo.
• Entradas: los mensajes enviados al modelo
• Salidas: la respuesta recibida del modelo

Este seguimiento sencillo ya proporciona información valiosa sobre el comportamiento de la aplicación, como:

• ¿Qué se preguntó?
• ¿Qué respuesta se generó?
• Cuánto tiempo tardó la solicitud
• Cuántos tokens se usaron (lo que afecta al costo)

Para aplicaciones más complejas como sistemas RAG o agentes de varios pasos, el seguimiento de MLflow proporciona aún más valor al revelar los trabajos internos de cada componente y paso.

Pasos siguientes

• Guía de seguimiento de MLflow : comience aquí para obtener más información detallada sobre el seguimiento de MLflow
• Integraciones de seguimiento de MLflow : más de 20 bibliotecas con integraciones de seguimiento automático
• Conceptos de seguimiento : descripción de los aspectos básicos del seguimiento de MLflow

Cuaderno de ejemplo

Comience: Trazado de MLflow para GenAI (Cuaderno de Databricks)

Obtener el cuaderno