Skip to main content

MLflow

MLflow provides management of Machine Learning models and experiment metrics and logs. Including calls to Metaflow in your code is similar to Weights and Biases.

In order for the MLflow client library to log metrics and register models with MLflow, provide one or more environment variables to Anyscale.

Open Source MLflow

If you have created your own MLflow server in your cloud account, then you can configure you Anyscale applications to track to it. Here's a ray.init() call that initializes an environment for tracking to MLflow.

ray.init(runtime_env={"pip":["mlflow"],
"env_vars":{"MLFLOW_TRACKING_URI":'YOUR_MLFLOW_TRACKING_URI'},
"excludes":["tests", "yello*"],
"working_dir"="."})

MLflow Hosted by Databricks

If you have a Databricks account, then include a hostname, token, and experiment name from Databricks and MLflow will log to your Databricks instance. For example:

ray.init(runtime_env={"pip":["mlflow"],
"env_vars":{"MLFLOW_TRACKING_URI":'YOUR_MLFLOW_TRACKING_URI',
"DATABRICKS_HOST":"http://databricks....",
"DATABRICKS_TOKEN":"YOURDATABRICKSTOKEN",
"MLFLOW_EXPERIMENT_NAME":"/Users/xxx@yyy.com/first-experiment"},
"excludes":["tests", "yello*"],
"working_dir"="."})

Here's an example of a task that logs some parameters and metrics to Databricks's MLflow:

@ray.remote
def logging_task():
with mlflow.start_run():
alpha = "ALPHA"
l1_ratio = "L1"
rmse = 0.211
r2 = 0.122
mae = 30
mlflow.log_param("alpha", alpha)
mlflow.log_param("l1_ratio", l1_ratio)
mlflow.log_metric("rmse", rmse)
mlflow.log_metric("r2", r2)
mlflow.log_metric("mae", mae)
return "Done"

print(ray.get(logging_task.remote()))