Rastreando o Amazon Bedrock com o MLflow
MLflow oferece suporte ao rastreamento automático para Amazon Bedrock, um serviço totalmente gerenciado em AWS que fornece bases de alto desempenho em
dos principais provedores de AI, como Anthropic, Cohere, Meta, Mistral AI, entre outros. Ao ativar o rastreamento automático
para Amazon Bedrock chamando a função mlflow.bedrock.autolog
, MLflow capturará os rastreamentos para a invocação LLM e log para a experiência MLflow ativa.
MLflow O rastreamento captura automaticamente as seguintes informações sobre as chamadas do site Amazon Bedrock:
- Solicitações e respostas de conclusão
- Latências
- Nome do modelo
- Metadados adicionais, como temperatura, 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 Amazon Bedrock, o senhor precisa instalar o MLflow e o AWS SDK for Python (Boto3).
- Development
- Production
Para ambientes de desenvolvimento, instale o pacote completo do MLflow com os extras do Databricks e boto3
:
pip install --upgrade "mlflow[databricks]>=3.1" boto3
O pacote completo do mlflow[databricks]
inclui todos os recursos para desenvolvimento local e experimentação no Databricks.
Para implantações de produção, instale mlflow-tracing
e boto3
:
pip install --upgrade mlflow-tracing boto3
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 Amazon Bedrock.
Antes de executar os exemplos abaixo, o senhor 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.
AWS Credenciais : Certifique-se de que suas credenciais AWS para acesso ao Bedrock estejam configuradas (por exemplo, via AWS CLI, IAM role ou variável de ambiente).
APIs suportadas
O MLflow oferece suporte ao rastreamento automático para as seguintes APIs do Amazon Bedrock:
Exemplo básico
import boto3
import mlflow
import os
# Ensure your AWS credentials are configured in your environment
# (e.g., via AWS CLI `aws configure`, or by setting
# AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN, AWS_DEFAULT_REGION)
# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()
# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-tracing-demo")
# Create a boto3 client for invoking the Bedrock API
bedrock = boto3.client(
service_name="bedrock-runtime",
region_name="<REPLACE_WITH_YOUR_AWS_REGION>",
)
# MLflow will log a trace for Bedrock API call
response = bedrock.converse(
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
messages=[
{
"role": "user",
"content": "Describe the purpose of a 'hello world' program in one line.",
}
],
inferenceConfig={
"maxTokens": 512,
"temperature": 0.1,
"topP": 0.9,
},
)
O rastreamento dos registros, associado ao experimento, pode ser visto na UI MLflow.
Entradas e saídas brutas
Em default, MLflow renderiza a interface de usuário avançada do tipo chat para mensagens de entrada e saída no Chat
tab. Para acessar view a carga útil bruta de entrada e saída, inclusive os parâmetros de configuração, clique em Inputs / Outputs
tab na interface do usuário.
O painel Chat
é compatível apenas com as APIs converse
e converse_stream
. Para os outros APIs, MLflow exibe apenas o Inputs / Outputs
tab.
transmissão
MLflow suporta o rastreamento de chamadas de transmissão para Amazon Bedrock APIs. O rastreamento gerado mostra a mensagem de saída agregada no site Chat
tab, enquanto os blocos individuais são exibidos no site Events
tab.
response = bedrock.converse_stream(
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
messages=[
{
"role": "user",
"content": [
{"text": "Describe the purpose of a 'hello world' program in one line."}
],
}
],
inferenceConfig={
"maxTokens": 300,
"temperature": 0.1,
"topP": 0.9,
},
)
for chunk in response["stream"]:
print(chunk)
MLflow não cria um intervalo imediatamente quando a resposta de transmissão é retornada. Em vez disso, ele cria um intervalo quando os pedaços de transmissão são consumidos , por exemplo, o for-loop no trecho de código acima.
Agente de chamada de funções
O MLflow Tracing captura automaticamente os metadados de chamada de função ao chamar as APIs do Amazon Bedrock. A definição da função e a instrução na resposta serão destacadas no site Chat
tab na UI de rastreamento.
Combinando isso com o recurso de rastreamento manual, o senhor pode definir um agente de chamada de função (ReAct) e rastrear sua execução. Toda a implementação do agente pode parecer complicada, mas a parte do rastreamento é bem simples: (1) adicione o decorador @mlflow.trace
às funções a serem rastreadas e (2) ative o rastreamento automático para o Amazon Bedrock com mlflow.bedrock.autolog()
. O MLflow cuidará da complexidade, como a resolução de cadeias de chamadas e o registro de metadados de execução.
import boto3
import mlflow
from mlflow.entities import SpanType
import os
# Ensure your AWS credentials are configured in your environment
# Enable auto-tracing for Amazon Bedrock
mlflow.bedrock.autolog()
# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/bedrock-agent-demo")
# Create a boto3 client for invoking the Bedrock API
bedrock = boto3.client(
service_name="bedrock-runtime",
region_name="<REPLACE_WITH_YOUR_AWS_REGION>",
)
model_id = "anthropic.claude-3-5-sonnet-20241022-v2:0"
# 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:
""" "Get the current weather in a given location"""
return "sunny" if city == "San Francisco, CA" else "unknown"
# Define the tool configuration passed to Bedrock
tools = [
{
"toolSpec": {
"name": "get_weather",
"description": "Get the current weather in a given location",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The city and state, e.g., San Francisco, CA",
},
},
"required": ["city"],
}
},
}
}
]
tool_functions = {"get_weather": get_weather}
# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str) -> str:
messages = [{"role": "user", "content": [{"text": question}]}]
# Invoke the model with the given question and available tools
response = bedrock.converse(
modelId=model_id,
messages=messages,
toolConfig={"tools": tools},
)
assistant_message = response["output"]["message"]
messages.append(assistant_message)
# If the model requests tool call(s), invoke the function with the specified arguments
tool_use = next(
(c["toolUse"] for c in assistant_message["content"] if "toolUse" in c), None
)
if tool_use:
tool_func = tool_functions[tool_use["name"]]
tool_result = tool_func(**tool_use["input"])
messages.append(
{
"role": "user",
"content": [
{
"toolResult": {
"toolUseId": tool_use["toolUseId"],
"content": [{"text": tool_result}],
}
}
],
}
)
# Send the tool results to the model and get a new response
response = bedrock.converse(
modelId=model_id,
messages=messages,
toolConfig={"tools": tools},
)
return response["output"]["message"]["content"][0]["text"]
# Run the tool calling agent
question = "What's the weather like in San Francisco today?"
answer = run_tool_agent(question)
A execução do código acima criará um único rastreamento que envolve todas as invocações do LLM e as chamadas de ferramentas.
Desativar o rastreamento automático
O rastreamento automático para o Amazon Bedrock pode ser desativado globalmente ligando para mlflow.bedrock.autolog(disable=True)
ou mlflow.autolog(disable=True)
.