Get Started with the MLflow Integration

Install the mlflow-mongodb package.

Run the following command to install the mlflow-mongodb package. The package also installs MLflow.

python -m pip install mlflow-mongodb

Prepare your connection string.

The plugin requires the name of the database that stores your MLflow data in the path of the connection string. Replace the placeholders in the following connection strings with your values:

mongodb+srv://<db_username>:<db_password>@<cluster-host>/<database>?retryWrites=true&w=majority

mongodb://<db_username>:<db_password>@<host>:<port>/<database>?retryWrites=true&w=majority

Start the MLflow server.

Run the following commands to start the server. Replace <connection-string> with the connection string from the previous step:

export TRACKING_STORE_URI="<connection-string>"export REGISTRY_STORE_URI="<connection-string>"export MLFLOW_BIND_HOST="127.0.0.1"export MLFLOW_BIND_PORT="5000"
mlflow server \  --backend-store-uri "${TRACKING_STORE_URI}" \  --registry-store-uri "${REGISTRY_STORE_URI}" \  --host "${MLFLOW_BIND_HOST}" \  --port "${MLFLOW_BIND_PORT}"

MLflow stores artifacts in the mlartifacts directory on the server. To view the MLflow UI, open http://127.0.0.1:5000 in your browser.

Create the example file.

The following example creates an experiment named mongodb-backend-demo, then starts a run by calling the mlflow.start_run() method. The run logs the following items:

• model_type parameter

• accuracy metric

• run-summary.txt artifact

The examples in this tutorial connect to the server by calling the mlflow.set_tracking_uri() method. To connect MLflow clients without changing the client code, set the MLFLOW_TRACKING_URI environment variable to http://127.0.0.1:5000.

Save the following code in a file named log_run.py:

import mlflow
mlflow.set_tracking_uri("http://127.0.0.1:5000")mlflow.set_experiment("mongodb-backend-demo")
with mlflow.start_run() as run:    mlflow.log_param("model_type", "example")    mlflow.log_metric("accuracy", 0.92)    mlflow.log_text("Hello from the MongoDB backend", "run-summary.txt")    print("run_id:", run.info.run_id)

Run the file.

python log_run.py

The output includes the run ID. MLflow stores the experiment, run, parameters, and metrics in MongoDB, and the artifact in the artifact store that you configured.

View the run in the MLflow UI.

In your browser, go to http://127.0.0.1:5000, and open the mongodb-backend-demo experiment. On the Runs tab, select the run to view its parameters, metrics, and artifacts. The run in the MLflow UI resembles the following image:

Create the example file.

The following example calls the log_model() method and registers the model as mongodb-backend-demo-model. The code then sets the alias and tag. Save the file as register_model.py:

import mlflowfrom mlflow import MlflowClientfrom mlflow.pyfunc import PythonModel
mlflow.set_tracking_uri("http://127.0.0.1:5000")mlflow.set_experiment("mongodb-backend-demo")
class EchoModel(PythonModel):    def predict(self, context, model_input):        return model_input
with mlflow.start_run():    mlflow.pyfunc.log_model(        name="model",        python_model=EchoModel(),        registered_model_name="mongodb-backend-demo-model",    )
client = MlflowClient()client.set_registered_model_alias("mongodb-backend-demo-model", "mongo-verified", 1)client.set_model_version_tag("mongodb-backend-demo-model", 1, "validated_by", "doc-test")

Run the file.

python register_model.py

View the model in the MLflow UI.

In your browser, go to http://127.0.0.1:5000, click Registered Models, and select mongodb-backend-demo-model to view its versions, aliases, and tags. The new model version in the MLflow UI resembles the following image:

Enable tracing.

To enable tracing, use one of the following approaches:

• Call the mlflow.<flavor>.autolog() method to trace calls to a supported framework or provider, such as OpenAI or LangChain. To learn more, see Automatic Tracing in the MLflow documentation.

• Add the @mlflow.trace decorator to the functions that you want to trace.

The following example uses the @mlflow.trace decorator to trace a function that calls an agent. Save the file as trace_agent.py:

import mlflow
mlflow.set_tracking_uri("http://127.0.0.1:5000")mlflow.set_experiment("agent-traces")
@mlflow.tracedef call_agent(prompt: str) -> str:    # Replace this line with your actual LLM or agent call    return f"Echo: {prompt}"
call_agent("Summarize this document")

Run the file.

python trace_agent.py

View the trace in the MLflow UI.

Open the agent-traces experiment in the MLflow UI and click Traces. The trace in the MLflow UI resembles the following image: