Observability with MLflow

Synalinks provides built-in observability through MLflow, enabling you to trace and monitor your LM programs in production.

Overview

The observability system automatically creates spans for each module call, capturing:

• Inputs and outputs of each module
• Duration of each call
• Cost information (when available from the language model)
• Success/failure status
• Parent-child relationships between nested module calls

Quick Start

Enable Observability

Important: You must call enable_observability() BEFORE creating any modules. Hooks are registered when modules are instantiated, so enabling observability after module creation will not trace those modules.

import synalinks

# Enable FIRST, before creating any modules
synalinks.enable_observability(
    tracking_uri="http://localhost:5000", experiment_name="my_experiment"
)

# Now create your modules - they will be automatically traced
inputs = synalinks.Input(data_model=Question)
outputs = await synalinks.Generator(...)(inputs)

Once enabled, all module calls in your program will be automatically traced.

Example Usage

import asyncio
import synalinks

# Enable observability before creating your program
synalinks.enable_observability(
    tracking_uri="http://localhost:5000", experiment_name="question_answering"
)


class Question(synalinks.DataModel):
    question: str = synalinks.Field(description="The user's question")


class Answer(synalinks.DataModel):
    answer: str = synalinks.Field(description="The answer to the question")


async def main():
    language_model = synalinks.LanguageModel(model="gemini/gemini-3.1-flash-lite-preview")

    # Create a simple question-answering program
    inputs = synalinks.Input(data_model=Question)
    outputs = await synalinks.Generator(
        data_model=Answer,
        language_model=language_model,
    )(inputs)

    program = synalinks.Program(
        inputs=inputs,
        outputs=outputs,
        name="qa_program",
        description="A simple QA program",
    )

    # Run the program - traces will be sent to MLflow
    result = await program(Question(question="What is the capital of France?"))
    if result:
        print(result.prettify_json())


if __name__ == "__main__":
    asyncio.run(main())

Running MLflow with Docker

Using Docker

Run MLflow tracking server locally with artifact proxying enabled:

docker run -d \
    --name mlflow \
    -p 5000:5000 \
    -v mlflow-data:/mlflow \
    ghcr.io/mlflow/mlflow:latest \
    mlflow server \
    --host 0.0.0.0 \
    --port 5000 \
    --backend-store-uri sqlite:///mlflow/mlflow.db \
    --default-artifact-root mlflow-artifacts:/ \
    --serve-artifacts \
    --artifacts-destination /mlflow/artifacts

Important flags:

• --serve-artifacts: Enables the MLflow server to proxy artifact uploads from clients
• --default-artifact-root mlflow-artifacts:/: Tells clients to use the server as an artifact proxy
• --artifacts-destination /mlflow/artifacts: Where the server stores artifacts on disk

Then configure Synalinks to use it:

import synalinks

synalinks.enable_observability(
    tracking_uri="http://localhost:5000", experiment_name="synalinks_traces"
)

Using Docker Compose

For a more complete setup with persistent storage, create a docker-compose.yml:

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    container_name: mlflow
    ports:
      - "5000:5000"
    volumes:
      - ./mlflow-data:/mlflow
    command: >
      mlflow server
      --host 0.0.0.0
      --port 5000
      --backend-store-uri sqlite:///mlflow/mlflow.db
      --default-artifact-root mlflow-artifacts:/
      --serve-artifacts
      --artifacts-destination /mlflow/artifacts
    restart: unless-stopped

Start the services:

docker compose up -d

Access the MLflow UI at http://localhost:5000.

Production Setup with PostgreSQL

For production deployments, use PostgreSQL as the backend store:

services:
  postgres:
    image: postgres:16-alpine
    container_name: mlflow-postgres
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: mlflow
      POSTGRES_DB: mlflow
    volumes:
      - postgres-data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mlflow"]
      interval: 5s
      timeout: 5s
      retries: 5
    restart: unless-stopped

  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    container_name: mlflow
    depends_on:
      postgres:
        condition: service_healthy
    ports:
      - "5000:5000"
    volumes:
      - mlflow-artifacts:/mlflow/artifacts
    environment:
      MLFLOW_BACKEND_STORE_URI: postgresql://mlflow:mlflow@postgres:5432/mlflow
    command: >
      mlflow server
      --host 0.0.0.0
      --port 5000
      --backend-store-uri postgresql://mlflow:mlflow@postgres:5432/mlflow
      --default-artifact-root mlflow-artifacts:/
      --serve-artifacts
      --artifacts-destination /mlflow/artifacts
    restart: unless-stopped

volumes:
  postgres-data:
  mlflow-artifacts:

Understanding Traces

When you run a Synalinks program with observability enabled, MLflow captures detailed traces.

Span Types

Synalinks automatically categorizes spans based on module type for better visualization in MLflow:

Span Attributes

Each span includes these attributes:

Exception Events

When a module call fails, the span automatically records an exception event with: - exception.type: The exception class name - exception.message: The exception message

Users and Sessions

MLflow can associate traces with a user and a chat session, which is useful for analyzing multi-turn conversations: you can then filter the traces of one session in the UI and inspect what happened at each turn. MLflow reads two reserved metadata keys for that, mlflow.trace.user and mlflow.trace.session.

Synalinks creates its spans outside of MLflow's fluent context (nested module calls are linked explicitly), so mlflow.update_current_trace() does not see them. Use synalinks.trace_context() instead: every trace started inside the block carries the given user and session.

import synalinks

synalinks.enable_observability()

# ... build your program ...

with synalinks.trace_context(user_id="user-123", session_id="session-123"):
    result = await program(inputs)

The context is stored in a contextvars.ContextVar, so it is safe to set per request in an async server: concurrent requests each see their own value, and the tasks spawned inside the block (parallel branches, agent tool calls) inherit it. Nested blocks merge with the enclosing one, the innermost value winning.

@app.post("/chat")
async def chat(request: ChatRequest):
    with synalinks.trace_context(
        user_id=request.user_id, session_id=request.session_id
    ):
        return await program(request.messages)

You can also attach free-form trace metadata (immutable once the trace is logged) and tags (editable afterwards in the MLflow UI):

with synalinks.trace_context(
    user_id="user-123",
    session_id="session-123",
    metadata={"turn": "3"},
    tags={"env": "production"},
):
    result = await program(inputs)

Viewing Traces

1. Open MLflow UI at http://localhost:5000
2. Navigate to your experiment
3. Click on a run to see detailed traces
4. Use the trace view to explore the call hierarchy

Configuration Options

Environment Variables

You can also configure MLflow using environment variables:

export MLFLOW_TRACKING_URI=http://localhost:5000

Then in your code:

import synalinks

# Will use MLFLOW_TRACKING_URI from environment
synalinks.enable_observability(experiment_name="my_experiment")

Direct Monitor Hook

For fine-grained control, you can create a Monitor hook directly:

import synalinks

monitor = synalinks.hooks.Monitor(
    tracking_uri="http://localhost:5000", experiment_name="custom_experiment"
)

# Add to a specific module
generator = synalinks.Generator(data_model=Answer, hooks=[monitor])

Training Metrics and Artifacts

The Monitor callback logs training metrics and program artifacts to MLflow during fit().

Basic Usage

import synalinks

# Create the monitor callback
monitor = synalinks.callbacks.Monitor(
    tracking_uri="http://localhost:5000",
    experiment_name="training_experiment",
    run_name="my_training_run",
    log_program_plot=True,  # Save program visualization as artifact
)

# Use during training
program.fit(x=train_inputs, y=train_labels, epochs=10, callbacks=[monitor])

Program Plot Artifact

When log_program_plot=True (the default), the Monitor callback automatically saves a visualization of your program architecture as an MLflow artifact at the start of training.

The plot is saved under program_plots/ in the artifacts folder and includes:

• Module names and types
• Input/output schemas
• Trainable status of each module

You can view the program plot in the MLflow UI under the "Artifacts" tab of your run.

Program Model Artifact

When log_program_model=True (the default), the Monitor callback saves the program's trainable state at the end of training. This includes:

• model/state_tree.json: Contains all trainable variables (few-shot examples, optimized prompts, etc.)
• model/model_info.json: Metadata about the program (name, description, number of trainable variables)

This is useful for:

• Checkpointing learned parameters during optimization
• Comparing different training runs
• Restoring program state for inference

Callback Parameters

Example with Full Configuration

import synalinks

monitor = synalinks.callbacks.Monitor(
    tracking_uri="http://localhost:5000",
    experiment_name="gsm8k_optimization",
    run_name="chain_of_thought_v1",
    log_batch_metrics=True,
    log_epoch_metrics=True,
    log_program_plot=True,
    tags={"model": "gpt-4o-mini", "optimizer": "RandomFewShot", "dataset": "gsm8k"},
)

program.fit(x=train_questions, y=train_answers, epochs=5, callbacks=[monitor])

Evaluation Results Over Time

program.evaluate() with the Monitor callback logs the reward and metrics of that evaluation to an MLflow run. Runs without a logged model are listed in the experiment's Evaluation runs tab (/#/experiments/<id>/evaluation-runs), where the charts view compares the metrics of every evaluation. The run is tagged mlflow.runType = genai_evaluate, like the runs created by mlflow.genai.evaluate(), and the traces of the evaluated calls are linked to it. Each evaluated sample's reward is also logged as a reward feedback assessment on the sample's trace, so the run's traces can be sorted and filtered by score (disable with log_assessments=False).

When synalinks.enable_observability() was called, a Monitor callback created without experiment_name uses that same experiment, so the evaluation runs appear next to the traces:

import synalinks

synalinks.enable_observability(
    tracking_uri="http://localhost:5000", experiment_name="my_app"
)

# ... build and compile your program ...

monitor = synalinks.callbacks.Monitor(run_name="eval")
await program.evaluate(x=test_x, y=test_y, callbacks=[monitor])

Each evaluate() call creates one run, so the Evaluation runs charts show one point per evaluation. To draw the metrics as a curve inside a single run instead, like the epoch curves of fit(), resume the same run on every evaluation with resume=True: the callback looks up the run named run_name in the experiment (creating it the first time), and logs each evaluation at the next step. This works across processes, so a scheduled evaluation appends one point per execution:

monitor = synalinks.callbacks.Monitor(run_name="nightly_eval", resume=True)
await program.evaluate(x=test_x, y=test_y, callbacks=[monitor])

In the MLflow UI, open the run and plot the metric against its step, or against wall time to see the results over time. Pass run_id instead of resume to resume a specific run; monitor.run_id holds the id of the run in use.

Combining Tracing with Training

When using both enable_observability() and the Monitor callback for training:

1. During program building (symbolic calls): Traces go to the experiment specified in enable_observability()

2. During training (fit()): Traces are associated with the training run and go to the experiment of the Monitor callback. When the callback has no experiment_name, that is the experiment of enable_observability(), so traces and runs stay together; give the callback its own experiment_name to keep them apart.

Full Example

import synalinks

# Enable tracing for all module calls
synalinks.enable_observability(
    tracking_uri="http://localhost:5000",
    experiment_name="synalinks_traces",  # Traces during setup go here
)

# Create your program (symbolic traces created here)
inputs = synalinks.Input(data_model=Question)
outputs = await synalinks.Generator(
    data_model=Answer,
    language_model=language_model,
)(inputs)

program = synalinks.Program(inputs=inputs, outputs=outputs, name="my_program")

# Create Monitor callback for training
monitor = synalinks.callbacks.Monitor(
    tracking_uri="http://localhost:5000",
    experiment_name="training_runs",  # Training metrics + traces go here
    run_name="experiment_v1",
)

# Train - traces during fit() are associated with the training run
program.compile(reward=reward, optimizer=optimizer)
await program.fit(x=train_x, y=train_y, epochs=5, callbacks=[monitor])

After training, you'll have: - synalinks_traces experiment: Setup traces (symbolic module calls) - training_runs experiment: Training run with metrics, artifacts, and execution traces

Best Practices

1. Enable observability early in your script, before creating any modules
2. Use meaningful experiment names to organize your traces by project or feature
3. Use persistent storage (PostgreSQL) for production deployments
4. Set up retention policies to manage storage for long-running applications

Troubleshooting

No traces being created

If you don't see any traces in MLflow:

1. Check call order: Ensure enable_observability() is called before creating any modules

   # Wrong - modules created before enabling observability
   inputs = synalinks.Input(data_model=Question)
   synalinks.enable_observability()  # Too late!
   
   # Correct - enable first
   synalinks.enable_observability()
   inputs = synalinks.Input(data_model=Question)  # Now traces will be created
   

2. Verify observability is enabled: Check with synalinks.is_observability_enabled()

3. Check the correct experiment: During training, traces go to the training experiment, not the observability experiment

MLflow not receiving traces

1. Verify the MLflow server is running: curl http://localhost:5000/health
2. Check the tracking URI is correct
3. Ensure mlflow package is installed: pip install mlflow

Artifacts not showing in MLflow UI

If artifacts are uploaded but don't appear in the MLflow UI:

1. Check server configuration: Ensure the MLflow server is started with --serve-artifacts flag
2. Verify artifact root: The server must use --default-artifact-root mlflow-artifacts:/ for remote clients
3. Check permissions: The server needs write access to --artifacts-destination path

Correct server configuration:

mlflow server \
    --serve-artifacts \
    --default-artifact-root mlflow-artifacts:/ \
    --artifacts-destination /mlflow/artifacts

Common mistake - Missing --serve-artifacts causes clients to try writing directly to the server's local filesystem, resulting in permission errors like:

PermissionError: [Errno 13] Permission denied: '/mlflow'

Missing cost information

Cost tracking requires the language model to return usage information. Ensure your LLM provider supports this feature.