Seguimiento de Amazon Bedrock con MLflow

MLflow admite el seguimiento automático de Amazon Bedrock, un servicio totalmente administrado en AWS que ofrece fundamentos de alto rendimiento de proveedores líderes de inteligencia artificial como Anthropic, Cohere, Meta, Mistral AI, etc. Al habilitar el seguimiento automático para Amazon Bedrock llamando a la función mlflow.bedrock.autolog, MLflow capturará rastros de la llamada a LLM y los registrará en el experimento de MLflow activo.

El seguimiento de MLflow captura automáticamente la siguiente información sobre las llamadas de Amazon Bedrock:

• Preguntas y respuestas
• Latencias
• Nombre del modelo
• Metadatos adicionales, como temperatura, max_tokens, si se especifica.
• Llamada a función si se devuelve en la respuesta
• Cualquier excepción si se genera

Prerrequisitos

Para usar el seguimiento de MLflow con Amazon Bedrock, debe instalar MLflow y el SDK de AWS para Python (Boto3).

Desarrollo

Para entornos de desarrollo, instale el paquete MLflow completo con complementos de Databricks y boto3:

pip install --upgrade "mlflow[databricks]>=3.1" boto3

El paquete completo mlflow[databricks] incluye todas las características para el desarrollo local y la experimentación en Databricks.

Producción

Para las implementaciones de producción, instale mlflow-tracing y boto3:

pip install --upgrade mlflow-tracing boto3

El mlflow-tracing paquete está optimizado para su uso en producción.

Nota:

MLflow 3 es muy recomendable para la mejor experiencia de seguimiento con Amazon Bedrock.

Antes de ejecutar los ejemplos siguientes, deberá configurar el entorno:

Para los usuarios que no utilizan los cuadernos de Databricks: establece las variables de entorno de Databricks:

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Para usuarios dentro de cuadernos Databricks: Estas credenciales se establecen automáticamente para usted.

Credenciales de AWS: asegúrese de que las credenciales de AWS para el acceso a Bedrock están configuradas. Para su uso en producción, considere la posibilidad de usar roles de IAM, AWS Secrets Manager o secretos de Databricks en lugar de variables de entorno (por ejemplo, a través de la CLI de AWS, roles de IAM o variables de entorno).

API admitidas

MLflow admite el seguimiento automático de las siguientes API de Amazon Bedrock:

• converse
• converse_stream
• invoke_model
• invoke_model_with_response_stream

Ejemplo 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,
    },
)

El seguimiento registrado, asociado al experimento, se puede ver en la interfaz de usuario de MLflow.

Entradas y salidas sin formato

De forma predeterminada, MLflow representa la interfaz de usuario tipo chat enriquecido para los mensajes de entrada y salida en la pestaña Chat. Para ver la carga útil de entrada y salida sin procesar, incluidos los parámetros de configuración, haga clic en la pestaña Inputs / Outputs de la interfaz de usuario.

Nota:

El Chat panel solo se admite para las converse y converse_stream APIs. Para las otras API, MLflow solo muestra la Inputs / Outputs pestaña .

Transmisión en línea

MLflow admite el seguimiento de llamadas de streaming a las API de Amazon Bedrock. El seguimiento generado muestra el mensaje de salida agregado en la Chat pestaña, mientras que los fragmentos individuales se muestran en la Events pestaña.

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)

Advertencia

MLflow no crea un intervalo inmediatamente cuando se devuelve la respuesta de streaming. En su lugar, crea un intervalo cuando se consumen los fragmentos de streaming, por ejemplo, el bucle for en el fragmento de código anterior.

Agente de llamada a funciones

El seguimiento de MLflow captura automáticamente los metadatos de la llamada de función al utilizar las API de Amazon Bedrock. La definición de función y la instrucción de la respuesta se resaltarán en la pestaña Chat de la IU de seguimiento.

Al combinarlo con la característica de seguimiento manual, puede definir un agente de llamada a funciones (ReAct) y realizar un seguimiento de su ejecución. La implementación completa del agente podría parecer complicada, pero la parte de trazado es bastante sencilla: (1) agrega el @mlflow.trace decorador a las funciones a trazar y (2) habilita el trazado automático para Amazon Bedrock con mlflow.bedrock.autolog(). MLflow se ocupará de la complejidad, como resolver cadenas de llamadas y registrar metadatos de ejecución.

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)

Al ejecutar el código anterior, se creará un único seguimiento que implica todas las invocaciones de LLM y las llamadas a la herramienta.

Advertencia

Para entornos de producción, use siempre secretos de AI Gateway o Databricks en lugar de valores codificados de forma rígida para la administración segura de claves de API.

Deshabilitar el seguimiento automático

El seguimiento automático de Amazon Bedrock se puede deshabilitar globalmente llamando a mlflow.bedrock.autolog(disable=True) o mlflow.autolog(disable=True).