Traccia di MLflow per gli agenti

Importante

Questa funzionalità è disponibile in anteprima pubblica.

Questo articolo descrive MLflow Tracing in Databricks e come usarlo per aggiungere osservabilità alle applicazioni di intelligenza artificiale generative.

Che cos'è MLflow Tracing?

MLflow Tracing acquisisce informazioni dettagliate sull'esecuzione di applicazioni di intelligenza artificiale di generazione. Tracciando gli input, gli output e i metadati associati a ogni passaggio intermedio di una richiesta, è possibile individuare l'origine dei bug e il comportamento imprevisto. Ad esempio, se il modello allucina, è possibile esaminare rapidamente ogni passaggio che ha portato allucinazione.

MLflow Tracing è integrato con gli strumenti e l'infrastruttura di Databricks, consentendo di archiviare e visualizzare tracce nei notebook di Databricks o nell'interfaccia utente dell'esperimento MLflow.

Perché usare MLflow Tracing?

MLflow Tracing offre diversi vantaggi:

• Esaminare una visualizzazione interattiva della traccia e usare lo strumento di analisi per diagnosticare i problemi.
• Verificare che i modelli di prompt e i guardrail producano risultati ragionevoli.
• Analizzare la latenza di framework, modelli e dimensioni dei blocchi diversi.
• Stimare i costi dell'applicazione misurando l'uso dei token tra modelli diversi.
• Stabilire set di dati di riferimento "golden" per valutare le prestazioni di versioni diverse.
• Archiviare tracce dagli endpoint del modello di produzione per eseguire il debug dei problemi ed eseguire la revisione e la valutazione offline.

Aggiungere tracce agli agenti

MLflow Tracing supporta tre metodi per aggiungere tracce alle applicazioni di intelligenza artificiale generative. Per dettagli di riferimento API, vedere la documentazione MLflow.

API	Caso d'uso consigliato	Descrizione
Autologging di MLflow	Sviluppo con librerie GenAI integrate	Autologging registra automaticamente le tracce per i framework open source supportati come LangChain, LlamaIndex e OpenAI.
API Fluenti	Agente personalizzato con Pyfunc	API a basso codice per l'aggiunta di tracce senza doversi preoccupare della gestione della struttura ad albero della traccia. MLflow determina automaticamente le relazioni di intervallo padre-figlio appropriate usando lo stack Python.
API Client di MLflow	Casi d'uso avanzati, ad esempio multithreading	MLflowClient fornisce API granulari e thread-safe per casi d'uso avanzati. È necessario gestire manualmente la relazione padre-figlio degli intervalli. È possibile avere un migliore controllo sul ciclo di vita della traccia, soprattutto nei casi d'uso multithread.

Installare MLflow Tracing

MLflow Tracing è disponibile nelle versioni 2.13.0 di MLflow e successive, le quali sono preinstallate in <DBR< 15.4 LTS ML e versioni successive. Se necessario, installare MLflow con il codice seguente:

%pip install mlflow>=2.13.0 -qqqU
%restart_python

In alternativa, è possibile installare la versione più recente di databricks-agents, che include una versione MLflow compatibile:

%pip install databricks-agents

Usare l'autologging per aggiungere tracce agli agenti

Se la libreria GenAI supporta la tracciabilità, ad esempio LangChain o OpenAI, abilita l'autologging aggiungendo mlflow.<library>.autolog() al codice. Per esempio:

mlflow.langchain.autolog()

Nota

A partire da Databricks Runtime 15.4 LTS ML, la traccia MLflow è abilitata per impostazione predefinita nei notebook. Per disabilitare il tracciamento, ad esempio, con LangChain, puoi eseguire mlflow.langchain.autolog(log_traces=False) nel notebook.

MLflow supporta librerie aggiuntive per la registrazione automatica delle tracce. Per un list completo delle librerie integrate, vedere la documentazione sul tracciamento di MLflow .

Usare le API Fluent per aggiungere manualmente tracce all'agente

Le API Fluent in MLflow creano automaticamente gerarchie di traccia in base al flusso di esecuzione del codice.

Decorare la funzione

Usare l'@mlflow.trace decorator per creare un intervallo per l'ambito della funzione decorata.

L'oggetto MLflow Span organizza le fasi di tracciamento. Span acquisisce informazioni su singole operazioni o passaggi, ad esempio chiamate API o query di archiviazione vettoriali, all'interno di un flusso di lavoro.

L'intervallo inizia quando viene richiamata la funzione e termina quando viene restituita. MLflow registra l'input e l'output della funzione ed eventuali eccezioni generate dalla funzione.

Ad esempio, il seguente codice crea uno span denominato my_function che raccoglie gli argomenti di input x e y e l'output.

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

Usare la gestione del contesto di tracciamento

Se si vuole creare un intervallo per un blocco arbitrario di codice, non solo una funzione, è possibile usare mlflow.start_span() come gestore del contesto che esegue il wrapping del blocco di codice. L'intervallo inizia quando si entra nel contesto e termina quando si esce dal contesto. L'input e l'output dello span devono essere forniti manualmente utilizzando i metodi setter dell'oggetto span restituito dal gestore del contesto.

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")

Incapsulare una funzione esterna

Per tracciare le funzioni della libreria esterna, racchiudere la funzione con 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)

### Fluent API example

The following example shows how to use the Fluent APIs `mlflow.trace` and `mlflow.start_span` to trace the `quickstart-agent`:

```python
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

Dopo aver aggiunto la traccia, eseguire la funzione . L'esempio seguente continua con la funzione predict() nella sezione precedente. Le tracce vengono visualizzate automaticamente quando si esegue il metodo di invocazione 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,
  }
)

Client API MLflow

MlflowClient espone API granulari e thread-safe per avviare e terminare le tracce, gestire gli intervalli e set campi di intervallo. Fornisce il controllo completo del ciclo di vita e della struttura della traccia. Queste API sono utili quando le API Fluent non sono sufficienti per le tue esigenze specifiche, ad esempio in applicazioni multithread e callback.

Di seguito sono riportati i passaggi per creare una traccia completa usando il client MLflow.

1. Crea un'istanza di MLflowClient client = MlflowClient().

2. Avviare una traccia usando il metodo client.start_trace(). Questo inizializza il contesto di traccia, avvia uno span radice assoluto e restituisce un oggetto span radice. Questo metodo deve essere eseguito prima dell'API start_span().

   1. Set attributi, input e output per la traccia in client.start_trace().

   Nota

   Non esiste un equivalente al metodo start_trace() nelle API Fluent. Ciò avviene perché le API Fluent inizializzano automaticamente il contesto di traccia e determinano se si tratta dell'intervallo radice in base allo stato gestito.

3. L'API start_trace() restituisce un intervallo. Get l'ID della richiesta, un identifier univoco della traccia, anche chiamato trace_id, e l'ID dell'intervallo restituito utilizzando span.request_id e span.span_id.

4. Avviare una traccia figlia utilizzando client.start_span(request_id, parent_id=span_id) a set per gli attributi, input e output della traccia.

   1. Questo metodo richiede request_id e parent_id a associare l'intervallo alla posizione corretta nella gerarchia di traccia. Restituisce un altro oggetto span.

5. Terminare il child span chiamando client.end_span(request_id, span_id).

6. Ripetere i passaggi da 3 a 5 per tutti gli intervalli figlio da creare.

7. Dopo che tutti gli intervalli figlio terminano, chiamare client.end_trace(request_id) per chiudere la traccia e registrarla.

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})

Revisione di tracce

Per esaminare le tracce dopo l'esecuzione dell'agente, usare una delle opzioni seguenti:

• Il rendering della visualizzazione di traccia viene eseguito inline nell'output della cella.
• Le tracce vengono registrate nell'esperimento MLflow. È possibile esaminare e cercare il completo delle tracce cronologiche nella scheda tracce nella pagina esperimento . Quando l'agente viene eseguito durante un'esecuzione MLflow attiva, le tracce appaiono sulla pagina del Run .
• Recuperare tracce usando l'API search_traces().

Usare MLflow Tracing nell'ambiente di produzione

MLflow Tracing è integrato anche con Mosaic AI Model Serving, consentendo di eseguire il debug dei problemi in modo efficiente, monitorare le prestazioni e creare un set di dati golden per la valutazione offline. Quando MLflow Tracing è abilitato per l'endpoint di servizio, i tracciamenti vengono registrati in un inferenza table nel responsecolumn.

Per abilitare il tracciamento di MLflow per l'endpoint di servizio, è necessario set la variabile di ambiente ENABLE_MLFLOW_TRACING nella configurazione dell'endpoint a True. Per informazioni su come distribuire un endpoint con variabili di ambiente personalizzate, vedere Aggiungere variabili di ambiente di testo normale. Se hai distribuito l'agente utilizzando l'API deploy(), le tracce vengono automaticamente registrate su un sistema di inferenza table. Vedere Distribuire un agente per un'applicazione di IA generativa.

Nota

La scrittura di tracce in un'inferenza table viene eseguita in modo asincrono, quindi non introduce lo stesso sovraccarico presente nell'ambiente notebook durante lo sviluppo. Tuttavia, potrebbe comunque comportare un sovraccarico alla velocità di risposta dell'endpoint, in particolare quando le dimensioni della traccia per ogni richiesta di inferenza sono elevate. Databricks non garantisce alcun contratto di servizio (SLA) per l'impatto effettivo sulla latenza sull'endpoint del modello, perché dipende in modo significativo dall'ambiente e dall'implementazione del modello. Databricks consiglia di testare le prestazioni dell'endpoint e ottenere informazioni dettagliate sul sovraccarico di traccia prima della distribuzione in un'applicazione di produzione.

L'table seguente fornisce un'indicazione approssimativa dell'impatto sulla latenza di inferenza per dimensioni di traccia diverse.

Dimensioni della traccia per richiesta	Impatto sulla latenza (ms)
~10 KB	~ 1 ms
~ 1 MB	50 ~ 100 ms
10 MB	150 ms ~

Limiti

• MLflow Tracing è disponibile nei notebook di Databricks, nei processi notebook e nella gestione dei modelli.

L'assegnazione automatica di LangChain potrebbe non supportare tutte le API di stima LangChain. Per la list completa delle API supportate, vedere documentazione di MLflow.