Compartilhar via

Rastreamento do MLflow para agentes

  • Artigo

Importante

Esse recurso está em uma versão prévia.

Este artigo descreve o Rastreamento de MLflow no Databricks e como usá-lo para adicionar observabilidade aos seus aplicativos de IA generativos.

O que é o Rastreamento do MLflow?

O MLflow Tracing captura informações detalhadas sobre a execução de aplicativos de IA generativa. O rastreamento registra entradas, saídas e metadados associados a cada etapa intermediária de uma solicitação para que você possa identificar a origem dos bugs e o comportamento inesperado. Por exemplo, se seu modelo alucinar, você pode inspecionar rapidamente cada etapa que levou à alucinação.

O Tracing do MLflow é integrado às ferramentas e à infraestrutura do Databricks, permitindo que você armazene e visualize rastreamentos nos notebooks do Databricks ou na interface do usuário do experimento do MLflow.

O rastreamento em tempo real captura informações detalhadas para cada etapa em um aplicativo de IA generativa

Por que usar o Rastreamento do MLflow?

O Rastreamento do MLflow oferece vários benefícios:

  • Examine uma visualização de rastreamento interativa e use a ferramenta de investigação para diagnosticar problemas.
  • Verifique se os modelos de prompt e os guardrails produzem resultados razoáveis.
  • Analise a latência de estruturas, modelos e tamanhos de partes diferentes.
  • Estimar os custos do aplicativo medindo o uso de token em diferentes modelos.
  • Estabeleça conjuntos de dados "padrão ouro" de referência para avaliar o desempenho de diferentes versões.
  • Armazene rastreamentos de pontos de extremidade do modelo de produção para depurar problemas e executar análise e avaliação offline.

Adicionar rastreamentos ao seu agente

O MLflow Tracing dá suporte a três métodos para adicionar rastreamentos aos aplicativos de IA generativos. Para detalhes da referência de API, veja a documentação do MLflow.

API Caso de uso recomendado Descrição
Registro em log automático do MLflow Desenvolvimento com bibliotecas integradas do GenAI O log automático registra automaticamente rastreamentos para estruturas de software livre com suporte, como LangChain, LlamaIndex e OpenAI.
APIs fluentes Agente personalizado com Pyfunc APIs de baixa codificação para adicionar rastreamentos sem se preocupar em gerenciar a estrutura de árvore do rastreamento. O MLflow determina as relações de intervalo pai-filho apropriadas automaticamente usando a pilha do Python.
APIs de cliente do MLflow Casos de uso avançados, como multithreading MLflowClient fornece APIs granulares e thread-safe para casos de uso avançados. Você deve gerenciar manualmente a relação pai-filho dos spans. Isso proporciona um controle melhor sobre o ciclo de vida do rastreamento, especialmente para cenários de uso com múltiplos threads.

Instalar o Rastreamento do MLflow

O Rastreamento do MLflow está disponível nas versões 2.13.0 do MLflow e superiores, que estão pré-instaladas no <DBR< 15.4 LTS ML e em versões superiores. Se necessário, instale o MLflow com o seguinte código:

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Como alternativa, você pode instalar a versão mais recente do databricks-agents, que inclui uma versão compatível do MLflow:

%pip install databricks-agents

Usar o registro automático para adicionar rastreamentos aos seus agentes

Se sua biblioteca GenAI der suporte ao rastreamento, como LangChain ou OpenAI, habilite o registro automático adicionando mlflow.<library>.autolog() ao seu código. Por exemplo:

mlflow.langchain.autolog()

Observação

A partir do Databricks Runtime 15.4 LTS ML, o rastreamento do MLflow é habilitado por padrão em notebooks. Para desabilitar o rastreamento, por exemplo, com LangChain, você pode executar mlflow.langchain.autolog(log_traces=False) no notebook.

O MLflow oferece suporte a bibliotecas adicionais para o registro automático de rastreamentos. Para obter uma lista completa de bibliotecas integradas, consulte a documentação MLflow Tracing.

Usar APIs do Fluent Design para adicionar rastreamentos manualmente ao seu agente

As APIs fluentes no MLflow criam automaticamente hierarquias de rastreamento com base no fluxo de execução do código.

Decorar sua função

Use o decorador @mlflow.trace para criar um intervalo para o escopo da função decorada.

O objeto MLflow Span organiza as etapas de rastreamento. O Spans captura informações sobre operações ou etapas individuais, como chamadas à API ou consultas de repositório de vetores, dentro de um fluxo de trabalho.

O intervalo começa quando a função é invocada e termina quando retorna. O MLflow registra a entrada e a saída da função e todas as exceções geradas da função.

Por exemplo, o código a seguir cria um intervalo chamado my_function que captura argumentos de entrada x e y e a saída.

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

Usar o gerenciador de contexto de rastreamento

Se você quiser criar um intervalo para um bloco arbitrário de código, não apenas uma função, poderá usar mlflow.start_span() como um gerenciador de contexto que encapsula o bloco de código. O intervalo começa quando o contexto é inserido e termina quando o contexto é encerrado. A entrada e as saídas do span devem ser fornecidas manualmente usando métodos setter do objeto span gerado pelo gerenciador de contexto.

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

Encapsular uma função externa

Para rastrear funções de biblioteca externas, encapsule a função com 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)

Exemplo de API fluente

O exemplo a seguir mostra como usar as APIs fluentes mlflow.trace e mlflow.start_span para rastrear o 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

Depois de adicionar o rastreamento, execute a função. O exemplo a seguir continua o exemplo da função predict() na seção anterior. Os rastros são mostrados automaticamente quando você executa o método de invocação, 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,
  }
)

APIs de cliente do MLflow

MlflowClient expõe APIs granulares e thread-safe para iniciar e encerrar rastreamentos, gerenciar intervalos e definir campos de intervalo. Ela dá controle total sobre o ciclo de vida e a estrutura de rastreamento. Essas APIs são úteis quando as APIs do Fluent são insuficientes para seus requisitos, como aplicativos com vários threads e retornos de chamada.

Veja a seguir as etapas para criar um rastreamento completo usando o Cliente MLflow.

  1. Crie uma instância de MLflowClient por client = MlflowClient().

  2. Inicie um rastreamento usando o método client.start_trace(). Isso inicia o contexto de rastreamento, inicia um intervalo raiz absoluto e retorna um objeto de extensão raiz. Esse método deve ser executado antes da API de start_span().

    1. Defina seus atributos, entradas e saídas para o rastreamento em client.start_trace().

    Observação

    Não há um equivalente ao método start_trace() nas APIs fluentes. Isso ocorre porque as APIs fluentes inicializam automaticamente o contexto de rastreamento e determinam se ele é o intervalo raiz com base no estado gerenciado.

  3. A API start_trace() retorna um intervalo. Obtenha a ID da solicitação, um identificador exclusivo do rastreamento também conhecido como trace_ide a ID do intervalo retornado usando span.request_id e span.span_id.

  4. Inicie um intervalo filho usando client.start_span(request_id, parent_id=span_id) para definir seus atributos, entradas e saídas para o intervalo.

    1. Esse método requer request_id e parent_id para associar o intervalo à posição correta na hierarquia de rastreamento. Ele retorna outro objeto de intervalo.

  5. Termine o intervalo filho chamando client.end_span(request_id, span_id).

  6. Repita as Etapas 3 a 5 para os subintervalos a serem criados.

  7. Depois que todos os intervalos filho terminarem, chame client.end_trace(request_id) para fechar o rastreamento e registrá-lo.

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

Revisão de rastreamentos

Para examinar os rastreamentos após a execução do agente, use uma das seguintes opções:

  • A visualização de rastreamento é renderizada de forma embutida na saída da célula.
  • Os rastreamentos são registrados em seu experimento do MLflow. Você pode examinar e pesquisar a lista completa de rastreamentos históricos na guia Rastreamentos na página Experimento. Quando o agente é executado em uma execução ativa do MLflow, os rastreamentos aparecem na página Executar.
  • Recupere rastreamentos programaticamente usando a API search_traces().

Usar o rastreamento do MLflow na produção

O MLflow Tracing também é integrado ao Mosaic AI Model Serving, permitindo depurar problemas com eficiência, monitorar o desempenho e criar um conjunto de dados principal para avaliação offline. Quando o Rastreamento do MLflow está habilitado para o ponto de extremidade de serviço, os rastreamentos são registrados em uma tabela de inferência na coluna response.

Você pode visualizar os rastreamentos registrados em tabelas de inferência consultando a tabela e exibindo os resultados em um notebook. Use display(<the request logs table>) no notebook e selecione as linhas individuais dos rastreamentos que você quer visualizar.

Para habilitar o rastreamento do MLflow para o ponto de extremidade de serviço, você deve definir a variável de ambiente ENABLE_MLFLOW_TRACING na configuração do ponto de extremidade como True. Para saber como implantar um endpoint com variáveis de ambiente personalizadas, consulte Adicionar variáveis de ambiente de texto simples. Se você implantou seu agente usando a API deploy(), os rastreamentos serão registrados automaticamente em uma tabela de inferência. Confira Implantar um agente para um aplicativo de IA generativa.

Observação

A gravação de rastreamentos em uma tabela de inferência é feita de forma assíncrona, portanto, não adiciona a mesma sobrecarga que no ambiente de notebook durante o desenvolvimento. No entanto, ainda pode introduzir alguma sobrecarga na velocidade de resposta dos pontos de extremidade, especialmente quando o tamanho do rastreamento para cada solicitação de inferência for grande. A Databricks não garante qualquer contrato de nível de serviço (SLA) para o impacto real da latência no ponto de extremidade do seu modelo, uma vez que depende fortemente do ambiente e da implementação do modelo. A Databricks recomenda testar o desempenho do seu ponto de extremidade e obter insights sobre a sobrecarga de rastreio antes de implementar numa aplicação de produção.

A tabela a seguir fornece uma indicação aproximada do impacto na latência de inferência para diferentes tamanhos de rastreamento.

Tamanho do rastreamento por solicitação Impacto na latência (ms)
~10 KB ~ 1 ms
~ 1 MB 50 ~ 100 ms
10 MB 150 ms ~

Limitações

  • O rastreamento do MLflow está disponível em notebooks do Databricks, trabalhos de notebook e Model Serving.

O registro automático do LangChain pode não suportar todas as APIs de previsão do LangChain. Para obter a lista completa de APIs com suporte, confira documentação do MLflow.