Obserwowalność agenta za pomocą śledzenia MLflow

W tym artykule opisano, jak dodać obserwowalność do swoich aplikacji generacyjnych AI za pomocą śledzenia MLflow w Databricks.

Co to jest śledzenie MLflow?

MLflow Tracing zapewnia kompleksową obserwowalność aplikacji generatywnej sztucznej inteligencji od etapu programowania do wdrożenia. Śledzenie jest w pełni zintegrowane z zestawem narzędzi AI generowanej przez Databricks, przechwytując szczegółowe informacje w całym cyklu rozwoju i produkcji.

Poniżej przedstawiono kluczowe przypadki użycia śledzenia w aplikacjach generatywnej sztucznej inteligencji:

• Usprawnione debugowanie: Śledzenie zapewnia wgląd w każdy krok aplikacji sztucznej inteligencji generatywnej, co ułatwia diagnozowanie i rozwiązywanie problemów.

• Ocena w trybie offline: Śledzenie generuje cenne dane do oceny agentów, umożliwiając pomiar i poprawę ich jakości w czasie.

• monitorowanie produkcji: śledzenie zapewnia wgląd w zachowanie agenta i szczegółowe kroki wykonywania, umożliwiając monitorowanie i optymalizowanie wydajności agenta w środowisku produkcyjnym.

• dzienniki inspekcji: Śledzenie MLflow generuje kompleksowe dzienniki inspekcji akcji i decyzji agenta. Jest to niezbędne do zapewnienia zgodności i obsługi debugowania w przypadku wystąpienia nieoczekiwanych problemów.

Wymagania

Śledzenie MLflow jest dostępne w systemie MLflow w wersji 2.13.0 lub nowszej. Usługa Databricks zaleca zainstalowanie najnowszej wersji biblioteki MLflow w celu uzyskania dostępu do najnowszych funkcji i ulepszeń.

%pip install mlflow>=2.13.0 -qqqU
%restart_python

automatyczne śledzenie

automatycznego rejestrowania MLflow umożliwia szybkie instrumentowanie agenta przez dodanie pojedynczego wiersza do kodu mlflow.<library>.autolog().

Biblioteka MLflow obsługuje automatyczne rejestrowanie dla najpopularniejszych bibliotek tworzenia agentów. Aby uzyskać więcej informacji na temat każdej biblioteki tworzenia, zobacz dokumentację automatycznego rejestrowania MLflow:

Wyłącz automatyczne rejestrowanie

Automatyczne śledzenie logów jest domyślnie włączone w wersji Databricks Runtime 15.4 ML i wersjach nowszych dla następujących bibliotek:

• LangChain
• Langgraph
• OpenAI
• LlamaIndex

Aby wyłączyć śledzenie automatyczne dla tych bibliotek, uruchom następujące polecenie w notesie:

`mlflow.<library>.autolog(log_traces=False)`

ręczne dodawanie śladów

Automatyczne rejestrowanie zapewnia wygodny sposób instrumentowania agentów, ale możesz bardziej szczegółowo instrumentować agenta lub dodać dodatkowe ślady, które nie są przechwytywane przez funkcję automatycznego rejestrowania. W takich przypadkach użyj interfejsów API śledzenia MLflow, aby ręcznie dodać ślady.

API śledzenia MLflow to niskokodowe interfejsy do dodawania śladów bez konieczności martwienia się o zarządzanie strukturą drzewa śledzenia. MLflow automatycznie określa odpowiednie relacje nadrzędno-podrzędne, korzystając z stacku języka Python.

Łączenie automatycznego rejestrowania i ręcznego śledzenia

Interfejsy API śledzenia ręcznego mogą służyć do automatycznego rejestrowania. MLflow łączy zakresy utworzone przez automatyczne rejestrowanie i ręczne śledzenie w celu utworzenia pełnego śledzenia wykonania agenta. Aby zapoznać się z przykładem łączenia automatycznego rejestrowania i ręcznego śledzenia, zobacz Instrumentacja agenta wywołującego narzędzie za pomocą śledzenia MLflow.

Śledź funkcje używając dekoratora @mlflow.trace

Najprostszym sposobem ręcznego instrumentowania kodu jest dekorowanie funkcji za pomocą dekoratora @mlflow.trace. Dekorator śledzenia MLflow tworzy „span” obejmujący zakres funkcji dekorowanej, który reprezentuje jednostkę wykonawczą w śledzeniu i jest wyświetlany jako pojedynczy wiersz w wizualizacji śledzenia. Zakres przechwytuje dane wejściowe i wyjściowe funkcji, opóźnienia i wszelkie wyjątki zgłaszane z funkcji.

Na przykład poniższy kod tworzy zakres o nazwie my_function, który przechwytuje argumenty wejściowe x i y oraz dane wyjściowe.

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

Możesz również dostosować nazwę zakresu, typ zakresu i dodać atrybuty niestandardowe do zakresu:

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

Śledzenie dowolnych bloków kodu przy użyciu menedżera kontekstu

Aby utworzyć zakres dla dowolnego bloku kodu, a nie tylko funkcji, użyj mlflow.start_span() jako menedżera kontekstu, który opakowuje blok kodu. Zakres rozpoczyna się po wejściu w kontekst i kończy się po wyjściu z kontekstu. Dane wejściowe i wyjściowe dla obiektu span powinny być podawane ręcznie, przy użyciu metod ustawiających dane tego obiektu span uzyskanego za pomocą menedżera kontekstu. Aby uzyskać więcej informacji, zobacz dokumentację MLflow — program obsługi kontekstu.

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

Biblioteki śledzenia niższego poziomu

Platforma MLflow udostępnia również interfejsy API niskiego poziomu do jawnego kontrolowania struktury drzewa śledzenia. Zobacz dokumentację MLflow — Instrumentacja ręczna.

Przykład śledzenia : Połączenie automatycznego rejestrowania i ręcznego śledzenia

W poniższym przykładzie przedstawiono połączenie automatycznego rejestrowania openAI i ręcznego śledzenia w celu pełnego instrumentowania agenta wywołującego narzędzia.

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

Dodawanie adnotacji do śladów przy użyciu tagów

tagów śledzenia MLflow to pary klucz-wartość, które umożliwiają dodawanie niestandardowych metadanych do śladów, takich jak identyfikator konwersacji, identyfikator użytkownika, skrót zatwierdzenia usługi Git itp. Tagi są wyświetlane w interfejsie użytkownika platformy MLflow w celu filtrowania i wyszukiwania śladów.

Tagi można ustawić na bieżący lub ukończony ślad przy użyciu API MLflow lub interfejsu użytkownika MLflow. W poniższym przykładzie pokazano dodawanie tagu do trwającego śladu za pomocą interfejsu API mlflow.update_current_trace().

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

Aby dowiedzieć się więcej na temat tagowania śladów i sposobu ich używania do filtrowania i wyszukiwania śladów, zobacz dokumentację MLflow — Ustawianie tagów śledzenia.

Przegląd śladów

Aby przejrzeć ślady po uruchomieniu agenta, użyj jednej z następujących opcji:

• wizualizacji wbudowanej: w notesach usługi Databricks ślady są renderowane śródwierszowo w danych wyjściowych komórki.
• eksperymentu MLflow: w usłudze Databricks przejdź do Eksperymenty> Wybierz eksperyment >Ślady, aby wyświetlić i przeszukać wszystkie ślady eksperymentu.
• MLflow Run: Gdy agent działa w ramach aktywnego MLflow Run, ślady są wyświetlane na stronie Run w MLflow UI.
• Interfejs użytkownika oceny agenta: W narzędziu Mosaic AI Agent Evaluation można przejrzeć ślady każdego wykonania agenta, klikając Zobacz szczegółowy widok śledzenia w wyniku oceny.
• Trace Search API: Aby programowo pobrać ślady, użyj Trace Search API.

Ocenianie agentów przy użyciu śladów

Dane śledzenia służą jako cenny zasób do oceny agentów. Przechwytując szczegółowe informacje na temat wykonywania modeli, MLflow Tracing pełni kluczową rolę w ocenie offline. Możesz użyć danych śledzenia, aby ocenić wydajność agenta względem złotego zestawu danych, zidentyfikować problemy i poprawić wydajność agenta.

%pip install -U mlflow databricks-agents
%restart_python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

Aby dowiedzieć się więcej informacji o ocenie agenta, zobacz Jak przeprowadzić ocenę i wyświetlić wyniki.

Monitor wdrożonych agentów z tabelami wnioskowania

Po wdrożeniu agenta w Mosaic AI Model Serving można użyć tabel wnioskowania do monitorowania agenta. Tabele wniosków zawierają szczegółowe dzienniki żądań, odpowiedzi, śledzenia agenta i informacje zwrotne od agenta z aplikacji przeglądowej. Te informacje umożliwiają debugowanie problemów, monitorowanie wydajności i tworzenie złotego zestawu danych na potrzeby oceny w trybie offline.

Aby włączyć tabele wnioskowania dla wdrożeń agentów, zobacz Włączanie tabel wnioskowania dla agentów sztucznej inteligencji.

Wykonywanie zapytań dotyczących śladów online

Użyj notesu, aby wysłać zapytanie do tabeli wnioskowania i przeanalizować wyniki.

Aby wizualizować ślady, uruchom display(<the request logs table>) i wybierz wiersze, aby sprawdzić:

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

Monitorowanie agentów

Zobacz Jak monitorować aplikację generatywnej sztucznej inteligencji.

Opóźnienie związane z narzutem trasowania

Ślady są zapisywane asynchronicznie, aby zminimalizować wpływ na wydajność. Jednak śledzenie nadal dodaje opóźnienie do szybkości odpowiedzi punktu końcowego, szczególnie gdy rozmiar śledzenia dla każdego żądania wnioskowania jest duży. Databricks zaleca przetestowanie punktu końcowego, aby zrozumieć wpływ opóźnień w śledzeniu, zanim wdrożysz go do środowiska produkcyjnego.

W poniższej tabeli przedstawiono przybliżone oszacowania wpływu opóźnienia według rozmiaru śledzenia:

Rozwiązywanie problemów

Aby uzyskać informacje dotyczące rozwiązywania problemów i często zadawanych pytań, zobacz dokumentację platformy MLflow: Śledzenie instrukcji i dokumentacji platformy MLflow : często zadawane pytania