Rastreamento de modelos da Databricks Foundation
O MLflow Tracing fornece o recurso de rastreamento automático para os Databricks Foundation Models. Como os Databricks Foundation Models usam um API compatível com OpenAI, o senhor pode ativar o rastreamento automático chamando a função mlflow.openai.autolog
, e o MLflow capturará os rastreamentos para a invocação do LLM e os log para o experimento MLflow ativo.
import mlflow
mlflow.openai.autolog()
MLflow O rastreamento captura automaticamente as seguintes informações sobre as chamadas do Databricks Foundation Model:
- Solicitações e respostas de conclusão
- Latências
- Nome do modelo e endpoint
- Metadados adicionais, como
temperature
,max_tokens
, se especificados - Chamada de função se retornada na resposta
- Qualquer exceção, se levantada
Pré-requisitos
Para usar o MLflow Tracing com o Databricks Foundation Models, o senhor precisa instalar o MLflow e o OpenAI SDK (já que o Databricks Foundation Models usa uma API compatível com o OpenAI).
- Development
- Production
Para ambientes de desenvolvimento, instale o pacote completo do MLflow com os extras do Databricks e o OpenAI SDK:
pip install --upgrade "mlflow[databricks]>=3.1" openai
O pacote completo do mlflow[databricks]
inclui todos os recursos para desenvolvimento local e experimentação no Databricks.
Para implementações de produção, instale o site mlflow-tracing
e o OpenAI SDK:
pip install --upgrade mlflow-tracing openai
O pacote mlflow-tracing
é otimizado para uso na produção.
O MLflow 3 é altamente recomendado para obter a melhor experiência de rastreamento com o Databricks Foundation Models.
Antes de executar os exemplos, você precisará configurar seu ambiente:
Para usuários fora do Databricks Notebook : Defina seu Databricks variável de ambiente:
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"
Para usuários do Databricks Notebook : Essas credenciais são definidas automaticamente para o senhor.
APIs suportadas
O MLflow oferece suporte ao rastreamento automático para as seguintes APIs do Databricks Foundation Model:
Conclusão do bate-papo | Chamada de função | transmissão | Assíncrono |
---|---|---|---|
✅ | ✅ | ✅ | ✅ |
Para solicitar suporte para APIs adicionais, abra uma solicitação de recurso no GitHub.
Exemplo básico
import mlflow
import os
from openai import OpenAI
# Databricks Foundation Model APIs use Databricks authentication.
# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()
# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/databricks-foundation-models-demo")
# Create OpenAI client configured for Databricks
client = OpenAI(
api_key=os.environ.get("DATABRICKS_TOKEN"),
base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
messages = [
{
"role": "user",
"content": "What is the capital of France?",
}
]
response = client.chat.completions.create(
model="databricks-llama-4-maverick",
messages=messages,
temperature=0.1,
max_tokens=100,
)
transmissão
MLflow Tracing suporta a transmissão API de Databricks Foundation Models. Com a mesma configuração de rastreamento automático, o site MLflow rastreia automaticamente a resposta de transmissão e renderiza a saída concatenada na UI do span.
import mlflow
import os
from openai import OpenAI
# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()
# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/databricks-streaming-demo")
# Create OpenAI client configured for Databricks
client = OpenAI(
api_key=os.environ.get("DATABRICKS_TOKEN"),
base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
stream = client.chat.completions.create(
model="databricks-llama-4-maverick",
messages=[
{"role": "user", "content": "Explain the benefits of using Databricks Foundation Models"}
],
stream=True, # Enable streaming response
)
for chunk in stream:
print(chunk.choices[0].delta.content or "", end="")
Chamada de função
O MLflow Tracing captura automaticamente a resposta de chamada de função do Databricks Foundation Models. A instrução da função na resposta será destacada na interface do usuário de rastreamento. Além disso, o senhor pode anotar a função da ferramenta com o decorador @mlflow.trace
para criar um intervalo para a execução da ferramenta.
O exemplo a seguir implementa um agente de chamada de função simples usando Databricks Foundation Models e MLflow Tracing.
import json
import os
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType
# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()
# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/databricks-function-agent-demo")
# Create OpenAI client configured for Databricks
client = OpenAI(
api_key=os.environ.get("DATABRICKS_TOKEN"),
base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)
# 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="databricks-llama-4-maverick",
messages=messages,
tools=tools,
)
ai_msg = response.choices[0].message
# If the model requests tool call(s), invoke the function 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="databricks-llama-4-maverick", 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)
Modelos disponíveis
O Databricks Foundation Models fornece acesso a uma variedade de modelos de última geração, incluindo Llama, Anthropic e outros modelos de fundação líderes.
Para obter a lista completa e mais atualizada dos modelos disponíveis e suas IDs de modelo, consulte a documentação Databricks Foundation Models.
Desativar o rastreamento automático
O rastreamento automático para Databricks Foundation Models pode ser desativado globalmente ligando para mlflow.openai.autolog(disable=True)
ou mlflow.autolog(disable=True)
.