Integracje śledzenia MLflow

Śledzenie MLflow jest zintegrowane z szeroką gamą popularnych bibliotek i frameworków generatywnej sztucznej inteligencji, oferując jednowierszowe środowisko automatycznego śledzenia dla wszystkich z nich. Dzięki temu można uzyskać bezpośredni wgląd w aplikacje GenAI przy minimalnej konfiguracji.

Ta szeroka obsługa oznacza, że można uzyskać wgląd bez znaczących zmian w kodzie, wykorzystując już używane narzędzia. W przypadku składników niestandardowych lub nieobsługiwanych bibliotek MLflow udostępnia również zaawansowane interfejsy API śledzenia ręcznego.

Automatyczne śledzenie rejestruje logikę i kroki pośrednie aplikacji, takie jak wywołania LLM, użycie narzędzi i interakcje agentów, opierając się na implementacji konkretnej biblioteki lub SDK.

Najważniejsze integracje na pierwszy rzut oka

Poniżej przedstawiono przykłady szybkiego startu dla niektórych najczęściej używanych integracji. Kliknij kartę, aby wyświetlić podstawowy przykład użycia. Aby zapoznać się ze szczegółowymi wymaganiami wstępnymi i bardziej zaawansowanymi scenariuszami dla każdego z nich, odwiedź dedykowane strony integracji (dostępne w kartach lub na liście poniżej).

OpenAI

import mlflow
import openai

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)
# View trace in MLflow UI

Pełny przewodnik integracji interfejsu OpenAI

LangChain

import mlflow
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.langchain.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langchain-tracing-demo")

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7, max_tokens=1000)
prompt = PromptTemplate.from_template("Tell me a joke about {topic}.")
chain = prompt | llm | StrOutputParser()

chain.invoke({"topic": "artificial intelligence"})
# View trace in MLflow UI

Pełny przewodnik integracji langchain

LangGraph

import mlflow
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.langchain.autolog() # LangGraph uses LangChain's autolog

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/langgraph-tracing-demo")

@tool
def get_weather(city: str):
    """Use this to get weather information."""
    return f"It might be cloudy in {city}"

llm = ChatOpenAI(model="gpt-4o-mini")
graph = create_react_agent(llm, [get_weather])
result = graph.invoke({"messages": [("user", "what is the weather in sf?")]})
# View trace in MLflow UI

Pełny przewodnik integracji langgraph

Antropiczny

import mlflow
import anthropic
import os

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.anthropic.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/anthropic-tracing-demo")

client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
message = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    messages=[{"role": "user", "content": "Hello, Claude"}],
)
# View trace in MLflow UI

Pełny przewodnik po integracji antropotycznej

DSPy

import mlflow
import dspy

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.dspy.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/dspy-tracing-demo")

lm = dspy.LM("openai/gpt-4o-mini") # Assumes OPENAI_API_KEY is set
dspy.configure(lm=lm)

class SimpleSignature(dspy.Signature):
    input_text: str = dspy.InputField()
    output_text: str = dspy.OutputField()

program = dspy.Predict(SimpleSignature)
result = program(input_text="Summarize MLflow Tracing.")
# View trace in MLflow UI

Pełny przewodnik integracji dsPy

Databricks

import mlflow
import os
from openai import OpenAI # Databricks FMAPI uses OpenAI client

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.openai.autolog() # Traces Databricks FMAPI using OpenAI client

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-fmapi-tracing")

client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[{"role": "user", "content": "Key features of MLflow?"}],
)
# View trace in MLflow UI

Pełny przewodnik integracji usługi Databricks

Bedrock

import mlflow
import boto3

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# import os
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.bedrock.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-tracing-demo")

bedrock = boto3.client(
    service_name="bedrock-runtime",
    region_name="us-east-1" # Replace with your region
)
response = bedrock.converse(
    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[{"role": "user", "content": "Hello World in one line."}]
)
# View trace in MLflow UI

Pełny przewodnik integracji Bedrock

Autogen

import mlflow
from autogen import ConversableAgent
import os

# If running this code outside of a Databricks notebook (e.g., locally),
# uncomment and set the following environment variables to point to your Databricks workspace:
# os.environ["DATABRICKS_HOST"] = "https://your-workspace.cloud.databricks.com"
# os.environ["DATABRICKS_TOKEN"] = "your-personal-access-token"

mlflow.autogen.autolog()

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/autogen-tracing-demo")

config_list = [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]
assistant = ConversableAgent("assistant", llm_config={"config_list": config_list})
user_proxy = ConversableAgent("user_proxy", human_input_mode="NEVER", code_execution_config=False)

user_proxy.initiate_chat(assistant, message="What is 2+2?")
# View trace in MLflow UI

Pełny przewodnik po integracji z rozwiązaniem AutoGen

Bezpieczne zarządzanie kluczami interfejsu API

W środowiskach produkcyjnych Databricks zaleca używanie AI Gateway lub tajnych danych Databricks do zarządzania kluczami API. Preferowana metoda to AI Gateway, oferujący dodatkowe funkcje zarządzania.

Brama sztucznej inteligencji (zalecane)

Databricks rekomenduje usługę Mosaic AI Gateway do zarządzania i monitorowania dostępu do modeli generatywnej sztucznej inteligencji.

Utwórz punkt końcowy modelu bazowego skonfigurowany przy użyciu bramy sztucznej inteligencji.

1. W obszarze roboczym usługi Databricks wybierz pozycję Obsługa>nowego punktu końcowego.
2. Wybierz typ punktu końcowego i dostawcę.
3. Skonfiguruj punkt końcowy za pomocą klucza API.
4. Podczas konfigurowania punktu końcowego włącz AI Gateway i skonfiguruj ograniczanie szybkości, mechanizmy awaryjne i zabezpieczenia według potrzeb.
5. Możesz uzyskać automatycznie wygenerowany kod, aby szybko rozpocząć wykonywanie zapytań względem punktu końcowego. Przejdź do udostępniania> swojego punktu końcowego >użyj>zapytania. Pamiętaj, aby dodać kod śledzenia:

import mlflow
from openai import OpenAI
import os

# How to get your Databricks token: https://docs.databricks.com/en/dev-tools/auth/pat.html
# DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')
# Alternatively in a Databricks notebook you can use this:
DATABRICKS_TOKEN = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking (if running outside Databricks)
# If running in a Databricks notebook, these are not needed.
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/my-genai-app")

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url="<YOUR_HOST_URL>/serving-endpoints"
)

chat_completion = client.chat.completions.create(
  messages=[
  {
    "role": "system",
    "content": "You are an AI assistant"
  },
  {
    "role": "user",
    "content": "What is MLflow?"
  }
  ],
  model="<YOUR_ENDPOINT_NAME>",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

Tajne dane Databricks

Użyj sekretów Databricks do zarządzania kluczami API:

1. Utwórz tajny zakres i zapisz klucz API:

   from databricks.sdk import WorkspaceClient
   
   # Set your secret scope and key names
   secret_scope_name = "llm-secrets"  # Choose an appropriate scope name
   secret_key_name = "api-key"        # Choose an appropriate key name
   
   # Create the secret scope and store your API key
   w = WorkspaceClient()
   w.secrets.create_scope(scope=secret_scope_name)
   w.secrets.put_secret(
       scope=secret_scope_name,
       key=secret_key_name,
       string_value="your-api-key-here"  # Replace with your actual API key
   )
   

2. Pobierz i użyj tajemnicy w kodzie.

   import mlflow
   import openai
   import os
   
   # Configure your secret scope and key names
   secret_scope_name = "llm-secrets"
   secret_key_name = "api-key"
   
   # Retrieve the API key from Databricks secrets
   os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(
       scope=secret_scope_name,
       key=secret_key_name
   )
   
   # Enable automatic tracing
   mlflow.openai.autolog()
   
   # Use OpenAI client with securely managed API key
   client = openai.OpenAI()
   response = client.chat.completions.create(
       model="gpt-4o-mini",
       messages=[{"role": "user", "content": "Explain MLflow Tracing"}],
       max_tokens=100
   )
   

Włączanie wielu integracji śledzenia automatycznego

Ponieważ aplikacje GenAI często łączą wiele bibliotek, śledzenie MLflow umożliwia jednoczesne włączanie automatycznego śledzenia dla kilku integracji, zapewniając ujednolicone środowisko śledzenia.

Aby na przykład włączyć zarówno śledzenie LangChain, jak i bezpośrednie śledzenie OpenAI:

import mlflow

# Enable MLflow Tracing for both LangChain and OpenAI
mlflow.langchain.autolog()
mlflow.openai.autolog()

# Your code using both LangChain and OpenAI directly...
# ... an example can be found on the Automatic Tracing page ...

Narzędzie MLflow wygeneruje pojedynczy, spójny ślad, który łączy kroki zarówno z wywołań LangChain, jak i bezpośrednich wywołań OpenAI LLM, co umożliwia przeanalizowanie pełnego przepływu. Więcej przykładów łączenia integracji można znaleźć na stronie Automatyczne śledzenie .

Wyłączanie automatycznego śledzenia

Automatyczne śledzenie dla dowolnej konkretnej biblioteki można wyłączyć przez wywołanie metody mlflow.<library>.autolog(disable=True). Aby wyłączyć wszystkie integracje automatycznego rejestrowania jednocześnie, użyj polecenia mlflow.autolog(disable=True).

import mlflow

# Disable for a specific library
mlflow.openai.autolog(disable=True)

# Disable all autologging
mlflow.autolog(disable=True)