MLflow Tracing para agentes

Importante

Esse recurso está em Prévia Pública.

Este artigo descreve o MLflow Tracing em Databricks e como usá-lo para adicionar observabilidade aos seus aplicativos generativos AI.

O que é MLflow Tracing?

MLflow O rastreamento captura informações detalhadas sobre a execução de aplicativos do site AI. Rastreamento logs entradas, saídas e metadados associados a cada etapa intermediária de uma solicitação para que o senhor possa identificar a origem de bugs e comportamentos inesperados. Por exemplo, se seu modelo tiver alucinações, você poderá inspecionar rapidamente cada etapa que levou à alucinação.

MLflow O rastreamento é integrado às ferramentas e à infraestrutura do Databricks, permitindo que o senhor armazene e exiba rastreamentos no Databricks Notebook ou na UI do experimento MLflow.

O rastreamento em linha captura informações detalhadas de cada etapa em um aplicativo gen AI

Por que usar o MLflow Tracing?

O MLflow Tracing oferece vários benefícios:

  • Analise uma visualização de rastreamento interativa e use a ferramenta de investigação para diagnosticar problemas.

  • Verifique se os padrões e grades de proteção imediatos produzem resultados razoáveis.

  • Analise a latência de diferentes estruturas, modelos e tamanhos de blocos.

  • Estimar os custos de aplicação medindo o uso de tokens em diferentes modelos.

  • Estabelecer um conjunto de dados "dourado" de referência para avaliar o desempenho de diferentes versões.

  • Armazene traços do endpoint do modelo de produção para depurar problemas e realizar revisão e avaliação off-line.

Adicione rastreamentos ao seu agente

MLflow O rastreamento é compatível com três métodos para adicionar rastreamentos aos aplicativos generativos do AI. Para obter detalhes de referência da API, consulte a documentação do MLflow.

API

Caso de uso recomendado

Descrição

Autologação do MLflow

Desenvolvimento com biblioteca GenAI integrada

Autologging automaticamente logs traces for supported código aberto frameworks like LangChain, LlamaIndex, and OpenAI.

APIs fluentes

Agente personalizado com Pyfunc

APIs de baixo código para adicionar traços sem se preocupar com o gerenciamento da estrutura de árvore do traço. O MLflow determina automaticamente as relações apropriadas de extensão pai-filho usando a pilha 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. O senhor deve gerenciar manualmente a relação pai-filho dos períodos. Isso proporciona um melhor controle sobre o ciclo de vida do rastreamento, especialmente para casos de uso com vários segmentos.

Instalar o MLflow Tracing

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

%pip install mlflow>=2.13.0 -qqqU
%restart_python

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

%pip install databricks-agents

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

Se a sua biblioteca GenAI for compatível com rastreamento, como LangChain ou OpenAI, ative o autologging adicionando mlflow.<library>.autolog() ao seu código. Por exemplo:

mlflow.langchain.autolog()

Observação

A partir de Databricks Runtime 15.4 LTS ML, o rastreamento MLflow é ativado por default no Notebook. Para desativar o rastreamento, por exemplo, com LangChain, o senhor pode executar mlflow.langchain.autolog(log_traces=False) no Notebook.

MLflow suporta biblioteca adicional para registro automático de rastreamento. Para obter uma lista completa de bibliotecas integradas, consulte a documentação de rastreamento do siteMLflow .

Use as APIs do Fluent para adicionar manualmente traços ao seu agente

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

Decore sua função

Use o decorador @mlflow.trace para criar uma extensão 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 para o site API ou consultas ao vector store, dentro de um fluxo de trabalho.

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

Por exemplo, o código a seguir cria um intervalo chamado my_function que captura os 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

Use o gerenciador de contexto de rastreamento

Se você quiser criar uma extensão para um bloco arbitrário de código, não apenas uma função, você pode usar mlflow.start_span() como um gerenciador de contexto que envolve 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 de span devem ser fornecidas manualmente usando os 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")

Envolva uma função externa

Para rastrear funções de biblioteca externas, envolva 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)

### 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

Depois de adicionar o rastreamento, execute a função. O exemplo a seguir continua com a função predict() na seção anterior. Os traços são mostrados automaticamente quando o senhor 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 o APIs granular e thread-safe para começar e terminar rastreamentos, gerenciar intervalos e definir campos de intervalo. Ele fornece controle completo do ciclo de vida e da estrutura do rastreamento. Essas APIs são úteis quando as APIs do Fluent não são suficientes para seus requisitos, como aplicativos multithread e retornos de chamada.

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

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

  2. Começar um rastreamento usando o método client.start_trace(). Isso inicia o contexto de rastreamento, inicia um intervalo de raiz absoluto e retorna um objeto de intervalo de raiz. Esse método deve ser executado antes do start_span() API.

    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 do Fluent. Isso ocorre porque o Fluent APIs inicializa automaticamente o contexto de rastreamento e determina se ele é a extensão raiz com base no estado de gerenciar.

  3. O começar() API retorna um intervalo. Obtenha o ID da solicitação, um identificador exclusivo do rastreamento também conhecido como trace_id e o ID do intervalo retornado usando span.request_id e span.span_id.

  4. Comece 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 a extensão à posição correta na hierarquia de rastreamento. Ele retorna outro objeto de extensão.

  5. Encerre a duração da criança ligando para client.end_span(request_id, span_id).

  6. Repita as etapas 3 a 5 para qualquer período infantil que você queira criar.

  7. Depois que todas as extensões da criança terminarem, chame client.end_trace(request_id) para fechar o rastreamento e gravá-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})

Revisando traços

Para revisar os rastreamentos depois de executar o agente, use uma das seguintes opções:

  • A visualização do traço é renderizada em linha na saída da célula.

  • Os traços são registros do seu experimento MLflow. O senhor pode revisar e pesquisar a lista completa de traços históricos em Traces tab na página Experiment. Quando o agente executa sob uma execução ativa do MLflow, os rastros aparecem na página de execução.

  • Recupere programaticamente os rastros usando a API search_traces().

Usar o MLflow Tracing na produção

MLflow O rastreamento também está integrado ao Mosaic AI Model Servingpermitindo que o senhor depure problemas com eficiência, monitore o desempenho e crie um golden dataset para avaliação off-line. Quando o rastreamento do MLflow está ativado para o seu endpoint de serviço, os rastreamentos são registrados em uma tabela de inferência na coluna response.

Para ativar o MLflow Tracing para seu endpoint de serviço, o senhor deve definir a variável de ambiente ENABLE_MLFLOW_TRACING na configuração do endpoint como True. Para saber como implantar um endpoint com variável de ambiente personalizada, consulte Adicionar variável de ambiente de texto simples. Se o senhor implantar seu agente usando o deploy() API os rastros são automaticamente registrados em uma tabela de inferência. Veja implantado um agente para aplicação generativa do AI .

Observação

A gravação de traços em uma tabela de inferência é feita de forma assíncrona, portanto, não adiciona a mesma sobrecarga que no ambiente do Notebook durante o desenvolvimento. No entanto, ele ainda pode introduzir alguma sobrecarga na velocidade de resposta do endpoint, principalmente quando o tamanho do rastreamento para cada solicitação de inferência é grande. A Databricks não garante nenhum acordo de nível de serviço (SLA) para o impacto real da latência em seu endpoint de modelo, pois isso depende muito do ambiente e da implementação do modelo. Databricks recomenda testar seu endpoint desempenho e obter percepções sobre a sobrecarga de rastreamento antes de implantá-lo em um aplicativo de produção.

A tabela a seguir fornece uma indicação aproximada do impacto na latência da 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

  • MLflow O rastreamento está disponível em Databricks Notebook, Notebook Job e servindo modelo.

O autologging da LangChain pode não ser compatível com todas as APIs de previsão da LangChain. Para obter a lista completa de APIs compatíveis, consulte a documentação do MLflow.