Traçage de span avec des gestionnaires de contexte

Ouvrir dans Databricks

Le gestionnaire de contexte mlflow.start_span() vous permet de créer des spans pour des blocs de code arbitraires. Tandis que les décorateurs de fonction tracent à la granularité des fonctions, start_span() peut capturer des interactions plus fines et plus complexes au sein de votre code.

Le traçage de span avec les gestionnaires de contexte vous offre un contrôle précis sur le code tracé :

Blocs de code arbitraires : tracez n'importe quel bloc de code, pas seulement des fonctions entières
Gestion automatique du contexte : MLflow gère les relations parent-enfant et le nettoyage.
Fonctionne avec les décorateurs de fonction : mélangez et associez-les avec @mlflow.trace pour les approches hybrides
Gestion des exceptions : capture automatique des erreurs, de manière similaire aux décorateurs de fonctions

Prérequis

Ce didacticiel nécessite les packages suivants :

mlflow[databricks] 3.1 et supérieur : Fonctionnalités MLflow de base avec des fonctionnalités GenAI et la connectivité Databricks.
openai 1.0.0 et versions ultérieures : L'application d'exemple ci-dessous utilise le client OpenAI. Dans votre propre code, remplacez-le par d’autres SDK, au besoin.

Installer les exigences de base :

Python
%pip install --upgrade "mlflow[databricks]>=3.1" "openai>=1.0.0"

dbutils.library.restartPython()

Prérequis pour MLflow 2

Databricks recommande vivement d'installer MLflow 3.1 ou une version plus récente si vous utilisez mlflow[databricks].

Pour MLflow 2, le traçage des spans avec les gestionnaires de contexte nécessite les packages suivants :

mlflow[databricks] 2.15.0 et versions ultérieures : fonctionnalité MLflow essentielle avec connectivité Databricks.
openai 1.0.0 et versions ultérieures : (Facultatif) Installez si votre code personnalisé utilise le client OpenAI.

Python
%pip install --upgrade "mlflow[databricks]>=2.15.0,<3.0.0"
pip install --upgrade openai>=1.0.0 # Install if needed
dbutils.library.restartPython()

API de traçage de spans

De la même manière que le décorateur de fonction, le gestionnaire de contexte capture automatiquement la relation parent-enfant, les exceptions et le temps d'exécution. Il est également compatible avec le traçage automatique.

Contrairement au décorateur de fonctions, le nom, les entrées et les sorties du span doivent être fournis manuellement. Vous pouvez les définir en utilisant l'objet LiveSpan qui est renvoyé par le context manager.

L'extrait de code ci-dessous illustre le traçage de portée de base :

Python
import mlflow

with mlflow.start_span(name="my_span") as span:
    x = 1
    y = 2
    span.set_inputs({"x":x, "y": y})
    z = x + y
    span.set_outputs(z)

Événements d'étendue

SpanEvent objets enregistrent des occurrences spécifiques pendant la durée de vie d'une étendue. L'extrait de code ci-dessous montre :

Création d'un événement avec le timestamp actuel
Création d'un événement avec un Timestamp spécifique (nanosecondes)
Création d'un événement à partir d'un Exception

Python
import mlflow
from mlflow.entities import SpanEvent, SpanType
import time

with mlflow.start_span(name="manual_span", span_type=SpanType.CHAIN) as span:
    # Create an event with current timestamp
    event = SpanEvent(
        name="validation_completed",
        attributes={
            &quot;records_validated&quot;: 1000,
            &quot;errors_found&quot;: 3,
            &quot;validation_type&quot;: &quot;schema&quot;
        }
    )
    span.add_event(event)

    # Create an event with specific timestamp (nanoseconds)
    specific_time_event = SpanEvent(
        name="data_checkpoint",
        timestamp=int(time.time() * 1e9),
        attributes={&quot;checkpoint_id&quot;: &quot;ckpt_123&quot;}
    )
    span.add_event(specific_time_event)

    # Create an event from an exception
    try:
        raise ValueError("Invalid input format")
    except Exception as e:
        error_event = SpanEvent.from_exception(e)
        # This creates an event with name="exception" and attributes containing:
        # - exception.message
        # - exception.type
        # - exception.stacktrace

        # Add to current span
        span = mlflow.get_current_active_span()
        span.add_event(error_event)

Statut de portée

SpanStatus définit l'état d'une étendue. Notez que le gestionnaire de contexte mlflow.start_span() écrase le statut à la sortie. L'extrait de code ci-dessous montre différentes façons de définir le statut des spans :

Python
import mlflow
from mlflow.entities import SpanStatus, SpanStatusCode, SpanType

with mlflow.start_span(name="manual_span", span_type=SpanType.CHAIN) as span:

    # Create status objects
    success_status = SpanStatus(SpanStatusCode.OK)
    error_status = SpanStatus(
        SpanStatusCode.ERROR,
        description="Failed to connect to database"
    )

    # Set status on a live span
    span.set_status(success_status)

    # Or use string shortcuts
    span.set_status("OK")
    span.set_status("ERROR")

    # When the context manager exits successfully, the status is overwritten with status "OK"

État de la query des étendues terminées :

Python
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(last_trace_id)

for span in trace.data.spans:
    print(span.status.status_code)

`RETRIEVER` spans

Utilisez RETRIEVER spans lors de la récupération de documents à partir d'un magasin de données. RETRIEVER spans doivent produire une liste de Documents, comme le montre l'exemple suivant :

Python
import mlflow
from mlflow.entities import Document, SpanType

@mlflow.trace(span_type=SpanType.RETRIEVER)
def retrieve_documents(query: str):
    span = mlflow.get_current_active_span()

    # Create Document objects (required for RETRIEVER spans)
    documents = [
        Document(
            page_content="The content of the document...",
            metadata={
                &quot;doc_uri&quot;: &quot;path/to/document.md&quot;,
                &quot;chunk_id&quot;: &quot;chunk_001&quot;,
                &quot;relevance_score&quot;: 0.95,
                &quot;source&quot;: &quot;knowledge_base&quot;
            },
            id="doc_123"  # Optional document ID
        ),
        Document(
            page_content="Another relevant section...",
            metadata={
                &quot;doc_uri&quot;: &quot;path/to/other.md&quot;,
                &quot;chunk_id&quot;: &quot;chunk_042&quot;,
                &quot;relevance_score&quot;: 0.87
            }
        )
    ]

    # Set outputs as Document objects for proper UI rendering
    span.set_outputs(documents)

    # Return in your preferred format
    return [doc.to_dict() for doc in documents]

retrieve_documents(query="What is ML?")

Accéder aux sorties du récupérateur :

Python
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(last_trace_id)

retriever_span = trace.search_spans(span_type=SpanType.RETRIEVER)[0]

if retriever_span.outputs:
    for doc in retriever_span.outputs:
        if isinstance(doc, dict):
            content = doc.get('page_content', '')
            uri = doc.get('metadata', {}).get('doc_uri', '')
            score = doc.get('metadata', {}).get('relevance_score', 0)
            print(f"Document from {uri} (score: {score})")

Exemple avancé

Vous trouverez ci-dessous un exemple plus complexe qui combine :

mlflow.start_span() gestionnaire de contexte
@mlflow.trace décorateur de fonction
Traçage automatique pour OpenAI

Python
from databricks_openai import DatabricksOpenAI
import mlflow
from mlflow.entities import SpanEvent, SpanType
import openai
import time

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

# Create an OpenAI client that is connected to Databricks-hosted LLMs.
client = DatabricksOpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def chat_iteration(messages, user_input):
    with mlflow.start_span(name="User", span_type=SpanType.CHAIN) as span:
        span.set_inputs({
            "messages": messages,
            "timestamp": time.time(),
        })

        # Set individual attribute
        span.set_attribute("messages_length", len(messages))

        # Set multiple attributes at once
        span.set_attributes({
            "environment": "production",
            "custom_metadata": {"key": "value"}
        })

        # Add events during execution
        span.add_event(SpanEvent(
            name="processing_started",
            attributes={
                &quot;stage&quot;: &quot;initialization&quot;,
                &quot;memory_usage_mb&quot;: 256,
            }
        ))

        span.set_outputs(user_input)

    messages.append({"role": "user", "content": user_input})

    response = client.chat.completions.create(
        model="databricks-claude-sonnet-4-5",
        max_tokens=100,
        messages=messages,
    )
    answer = response.choices[0].message.content
    print(f"Assistant: {answer}")

    messages.append({"role": "assistant", "content": answer})


chat_iteration(
    messages = [{"role": "system", "content": "You are a friendly chat bot"}],
    user_input="What is your favorite color?",
)

Pour voir un exemple de traçage imbriqué pour des conversations plus longues, décommentez l'exemple ci-dessous :

Python
# @mlflow.trace(span_type=SpanType.CHAIN)
# def start_session():
#     messages = [{"role": "system", "content": "You are a friendly chat bot"}]
#     while True:
#         user_input = input(">> ")
#         chat_iteration(messages, user_input)

#         if user_input == "BYE":
#             break


# start_session()

Étapes suivantes

Décorateurs de fonction — Approche plus simple pour le traçage de fonctions entières
APIs client de bas niveau — Apprenez des scénarios avancés nécessitant un contrôle total
Déboguer et analyser votre application - Interroger et analyser les traces journalisées

Exemple de Notebook

Traçage des spans avec les gestionnaires de contexte

Ouvrir le notebook dans un nouvel onglet Ouvrir dans Databricks

Prérequis​

Prérequis pour MLflow 2​

API de traçage de spans​

Événements d'étendue​

Statut de portée​

RETRIEVER spans​

Exemple avancé​

Étapes suivantes​

Exemple de Notebook​