MLflow

RIME integrates with MLflow Tracking. This integration unlocks the true power of AI Stress Testing when experimenting with many models.

With the RIME SDK you are able to query for RIME metrics and include them in your MLflow experiments. You can see in the below code snippet. To get a better understanding of the integration, take a look at the MLflow Demo Video.

import pandas as pd
import mlflow
from rime_sdk import Client

# Set these before beginning!
BACKEND_URL = "rime-backend.<YOUR_ORG_NAME>.rime.dev"
API_KEY = "<YOUR_API_KEY>"

# Connect to your cluster.
rime_client = Client(BACKEND_URL, api_key=API_KEY)

# Run the test!
stress_test = rime_client.start_stress_test(test_run_config=config)

# Query the results.
test_run_result = stress_test.get_test_run_result()
test_cases_result = stress_test.get_test_cases_result()

# Preparing the results to log.
test_run_metrics = test_run_result.columns
test_cases_result.to_csv("test_cases_results.csv")

# Log experiment results to MLflow.
with mlflow.start_run():
    for metric in test_run_metrics:
        mlflow.log_metric(column, test_run_result[column][0])
    mlflow.log_artifact("test_cases_results.csv")
mlflow.end_run()