Aller au contenu principal
Dernière mise à jour le

Suivi d'Amazon Bedrock avec MLflow

MLflow prend en charge le traçage automatique pour Amazon Bedrock, un service entièrement managé sur AWS qui fournit des fondations haute performance de la part de principaux fournisseurs d'IA tels qu'Anthropic, Cohere, Meta, Mistral et d'autres. En activant le traçage automatique pour Amazon Bedrock en appelant la fonction mlflow.bedrock.autolog, MLflow capturera les traces pour l'invocation de LLM et les consignera dans l'Experimentation MLflow active.

Traçage d'agents Bedrock DIY

La trace MLflow capture automatiquement les informations suivantes concernant les appels Amazon Bedrock :

  • Prompts et réponses de complétion
  • Latences
  • Nom du modèle
  • Métadonnées supplémentaires telles que la température, max_tokens, si spécifié.
  • Appel de fonction si renvoyé dans la réponse
  • Toute exception si levée
remarque

Sur les clusters de compute serverless, l'autologging n'est pas activé automatiquement. Vous devez appeler explicitement mlflow.bedrock.autolog() pour activer le traçage automatique de cette intégration.

Prérequis​

Pour utiliser MLflow Tracing avec Amazon Bedrock, vous devez installer MLflow et le SDK AWS pour Python (Boto3).

Pour les environnements de développement, installez le package MLflow complet avec les extras Databricks et boto3:

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

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est fortement recommandé pour une expérience de traçage optimale avec Amazon Bedrock.

Avant d'exécuter les exemples ci-dessous, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

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

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Identifiants AWS : assurez-vous que vos identifiants AWS pour l'accès à Bedrock sont configurés. Pour une utilisation en production, il est recommandé d'utiliser les rôles IAM, AWS Secrets Manager ou les secrets Databricks plutôt que des variables d'environnement (par exemple, en utilisant AWS CLI, les rôles IAM ou des variables d'environnement).

APIs prises en charge​

MLflow prend en charge le traçage automatique pour les APIs Amazon Bedrock suivantes :

Exemple de base​

Python
import boto3
import mlflow
import os

# Ensure your AWS credentials are configured in your environment
# (e.g., using 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={
        &quot;maxTokens&quot;: 512,
        &quot;temperature&quot;: 0.1,
        &quot;topP&quot;: 0.9,
    },
)

La trace enregistrée, associée à l'expérimentation, est visible dans l'interface utilisateur de MLflow.

Entrées et sorties brutes​

Par default, MLflow affiche l'interface utilisateur riche de type conversation pour les messages d'entrée et de sortie dans l'tab Chat. Pour afficher la charge utile d'entrée et de sortie brute, y compris les parameters de configuration, cliquez sur l'tab Inputs / Outputs de l'interface utilisateur.

remarque

Le panneau Chat n'est pris en charge que pour les APIs converse et converse_stream. Pour les autres APIs, MLflow affiche uniquement le tab Inputs / Outputs.

streaming​

MLflow prend en charge le suivi des appels de streaming vers les APIs Amazon Bedrock. La trace générée affiche le message de sortie agrégé dans le tab Chat, tandis que les morceaux individuels sont affichés dans le tab Events.

Python
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={
        &quot;maxTokens&quot;: 300,
        &quot;temperature&quot;: 0.1,
        &quot;topP&quot;: 0.9,
    },
)

for chunk in response["stream"]:
    print(chunk)

Traçage de Stream Bedrock

attention

MLflow ne crée pas de span immédiatement lorsque la réponse en streaming est renvoyée. Au lieu de cela, il crée une étendue lorsque les blocs de streaming sont consommés , par exemple, la boucle for dans l'extrait de code ci-dessus.

Agent d'appel de fonction​

MLflow Tracing capture automatiquement les métadonnées d'appel de fonction lors de l'appel des APIs Amazon Bedrock. La définition et l'instruction de la fonction dans la réponse seront mises en évidence dans la tab Chat de l'interface utilisateur de trace.

En combinant cela avec la fonctionnalité de traçage manuel, vous pouvez définir un agent d'appel de fonctions (ReAct) et en tracer l'exécution. L'implémentation complète de l'agent peut sembler compliquée, mais la partie traçage est assez simple : (1) ajoutez le décorateur @mlflow.trace aux fonctions à tracer et (2) activez le traçage automatique pour Amazon Bedrock avec mlflow.bedrock.autolog(). MLflow prendra en charge la complexité, telle que la résolution des chaînes d'appels et l'enregistrement des métadonnées d'exécution.

Python
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={&quot;tools&quot;: 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={&quot;tools&quot;: 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)

L'exécution du code ci-dessus créera une trace unique qui implique toutes les invocations LLM et les appels d'outil.

Traçage d&#39;agents Bedrock DIY

attention

Pour les environnements de production, utilisez toujours AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour une gestion sécurisée des clés API.

Désactiver le suivi automatique​

Le traçage automatique pour Amazon Bedrock peut être désactivé globalement en appelant mlflow.bedrock.autolog(disable=True) ou mlflow.autolog(disable=True).

Sur cette page