Rastreamento de extensão com gerenciadores de contexto

O gerenciador de contexto mlflow.start_span() permite criar intervalos para blocos de código arbitrários. Enquanto os decoradores de função rastreiam na granularidade das funções, start_span() pode capturar interações mais refinadas e complexas dentro do seu código.

O rastreamento de trechos de código com gerenciadores de contexto oferece controle preciso sobre qual código será rastreado:

Blocos de código arbitrários: Rastreie qualquer bloco de código, não apenas funções inteiras.
Gerenciamento automático de contexto: o MLflow lida com relacionamentos pai-filho e limpeza.
Funciona com decoradores de função: combine com @mlflow.trace para abordagens híbridas.
Tratamento de exceções: Captura automática de erros, semelhante aos decoradores de função.

Pré-requisitos

Este tutorial requer o seguinte pacote:

mlflow[databricks] 3.1 e versões superiores: Funcionalidade principal MLflow com recurso GenAI e conectividade com Databricks .
openai 1.0.0 e acima: O aplicativo de exemplo abaixo usa o cliente OpenAI . No seu próprio código, substitua isso por outros SDKs conforme necessário.

Instale os requisitos básicos:

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

dbutils.library.restartPython()

Pré-requisitos para o MLflow 2

A Databricks recomenda fortemente a instalação do MLflow 3.1 ou mais recente se estiver usando mlflow[databricks].

Para MLflow 2, o rastreamento de spans com gerenciadores de contexto requer o seguinte pacote:

mlflow[databricks] 2.15.0 e versões superiores: Funcionalidade principal MLflow com conectividade Databricks .
openai 1.0.0 e versões superiores: (Opcional) Instale se o seu código personalizado usar o cliente 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 rastreamento de extensão

De forma semelhante ao decorador de função, o gerenciador de contexto captura automaticamente a relação pai-filho, exceções e tempo de execução. Também é compatível com rastreamento automático .

Diferentemente do que ocorre com o decorador de função, o nome, as entradas e as saídas do span devem ser fornecidos manualmente. Você pode configurá-los usando o objeto LiveSpan que é retornado pelo gerenciador de contexto.

O trecho de código abaixo ilustra o rastreamento básico de intervalos (span tracing):

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)

Eventos de abrangência

Os objetosSpanEvent registram ocorrências específicas durante a vida útil de um intervalo. O trecho de código abaixo mostra:

Criando um evento com o carimbo de data/hora atual.
Criar um evento com um carimbo de data/hora específico (nanossegundos)
Criar um evento a partir de um 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={
            "records_validated": 1000,
            "errors_found": 3,
            "validation_type": "schema"
        }
    )
    span.add_event(event)

    # Create an event with specific timestamp (nanoseconds)
    specific_time_event = SpanEvent(
        name="data_checkpoint",
        timestamp=int(time.time() * 1e9),
        attributes={"checkpoint_id": "ckpt_123"}
    )
    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)

Estado do span

SpanStatus define o estado de um intervalo. Note que o gerenciador de contexto mlflow.start_span() sobrescreve o status ao sair. O trecho de código abaixo mostra várias maneiras de definir o status em 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"

Status da consulta a partir de intervalos concluídos:

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` abrange

Use spansRETRIEVER ao recuperar documentos de um armazenamento de dados. RETRIEVER spans deve produzir uma lista de Documents conforme mostrado no exemplo a seguir:

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={
                "doc_uri": "path/to/document.md",
                "chunk_id": "chunk_001",
                "relevance_score": 0.95,
                "source": "knowledge_base"
            },
            id="doc_123"  # Optional document ID
        ),
        Document(
            page_content="Another relevant section...",
            metadata={
                "doc_uri": "path/to/other.md",
                "chunk_id": "chunk_042",
                "relevance_score": 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?")

Saídas do recuperador de acesso:

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})")

Exemplo avançado

Abaixo, um exemplo mais complexo que combina:

mlflow.start_span() gerenciador de contexto
@mlflow.trace decorador de função
Rastreamento automático para OpenAI

Python
from databricks.sdk import WorkspaceClient
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.
workspace = WorkspaceClient()
client = workspace.serving_endpoints.get_open_ai_client()

@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={
                "stage": "initialization",
                "memory_usage_mb": 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?",
)

Para ver um exemplo de rastreamento aninhado para conversas mais longas, remova o comentário do exemplo abaixo:

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()

Próximos passos

Decoradores de função - Uma abordagem mais simples para rastrear funções inteiras.
APIs de cliente de baixo nível - Aprenda cenários avançados que exigem controle total.
Depure e analise seu aplicativo - Consulte e analise os registros de logs.

Exemplo de caderno

Rastreamento de extensão com gerenciadores de contexto

Open notebook in new tab

Pré-requisitos​

Pré-requisitos para o MLflow 2​

API de rastreamento de extensão​

Eventos de abrangência​

Estado do span​

RETRIEVER abrange​

Exemplo avançado​

Próximos passos​

Exemplo de caderno​

Rastreamento de extensão com gerenciadores de contexto

Pré-requisitos

Pré-requisitos para o MLflow 2

API de rastreamento de extensão

Eventos de abrangência

Estado do span

`RETRIEVER` abrange

Exemplo avançado

Próximos passos

Exemplo de caderno