Migrating the MLflow Python RunData interface to expose metrics, params, tags as dicts

0_description.md

Background

With the next MLflow release, we'd like to update the Python RunData interface to better support querying metrics, params, and tags. In particular, the current RunData interface, which exposes flat lists of Metric, Param, and RunTag instances, falls short in that it:

1. requires a complicated O(n) scan to find a metric/param/tag by key, and
2. forces users to then make an additional field-access within a Metric/Param/RunTag instance to get at the actual metric/param/tag value.

Design Decisions & Proposed API

We can address point a) by migrating the metrics, params, tags fields of RunData to dictionaries whose keys are metric, param, and tag keys. This is viable as RunData contains at most one metric, param, and tag per key - for metrics, RunData is expected to contain the maximum metric value at the maximum timestamp for a given key.

Given this dict-based approach, we must also choose a value type for the dictionaries - two reasonable options are:

1. (Proposed) To use simple scalar values - strings for params and tags, and floats for metrics and
2. (Alternative) To use Python entity classes (Metric, Param, RunTag).

Option 1) makes it easier to access scalar metric/param/tag values, but harder to access non-value fields of metrics, params, and tags (e.g. the timestamp & eventually x-coordinate of metrics). Note additionally that option 1) doesn't preclude exposing full Metric / Param / RunTag objects via the RunData interface in the future (i.e. we could introduce new fields in RunData that expose such information). For example, we could handle the addition of new fields to params or tags by introducing lower-level metric_objs, param_objs, tag_objs fields to RunData that expose flat lists of metrics/params/tags (as the current API does) - however, we expect changes to params & tags to be unlikely. In brief, option 1) trades off flexibility of the RunData interface in exchange for user-friendliness, which is a compromise that seems worthwhile given the unlikeliness of API changes to Param or RunTag.

A small detail: with option 1) and the current set of fluent/client APIs, it won't be possible to access metric timestamps in Python. We can address this by exposing the AbstractStore's get_metric_history API via the Python client API, e.g. via an get_metric_history method in MlflowClient.

Option 2)'s strengths are option 1)'s weaknesses - 2) allows us to add new fields to metrics/params/tags (e.g. metric x coordinates) and access them via RunData without adding new fields to RunData, at the cost of requiring an extra field access to read metric/param/tag values.

Request for Feedback

Any feedback on the proposed APIs is much appreciated :) - see the gists below for an example of how user workflows might look with the new APIs. It'd be particularly helpful to hear if there are any use cases that would be helped/hurt by the APIs proposed above. Note also that we expect many query use cases to be addressed by the search_runs API, which allows for filtering & searching runs by metric/param/tag. We'd also eventually like to add support for getting a handle to all logged run data as a Pandas or Spark DataFrame, which we expect will also simplify query use cases as well as data export from a server.

1_new_api_proposed.py

"""
This file contains an example of querying a run's params, metrics, and tags using the newly proposed RunData interface
(Option 1). Any feedback is much appreciated, especially regarding query workflows that we should consider while
designing the new interface.
"""
run = mlflow.get_run("...")

# 1. Simple example of looking up metric/param/tag values
tag_key = "mlflow.docker"
metric_key = "rmse"
param_key = "alpha"
tag_val = run.data.tags[tag_key] # tag_val is a string, e.g. 'true'
metric_val = run.data.metrics[metric_key] # metric_obj is a float value, e.g. '0.1'
param_val = run.data.params[param_key] # param_value is a string, e.g. '1.0'
print("Tag key: %s, value: %s" % (tag_key, tag_val))
print("Metric key: %s, value: %s" % (metric_key, metric_val))
print("Param key: %s, value: %s" % (param_key, param_val))

# 2. We can access metric timestamps using the get_metric_history API, assuming it's exposed as described in the proposal
client = mlflow.tracking.MlflowClient()
metric_objs = client.get_metric_history(metric_key)
print([metric.timestamp for metric in metric_objs])

# 3. Some examples of comparing data across runs.
# Note that comparing `run.data.metrics` across runs only considers metric values, rather than timestamp/x coordinate.
# This may or may not be desirable - in general though, it seems unlikely for metrics to have the exact same values across
# runs.
new_run = mlflow.get_run("...")
if run.data.metrics == new_run.data.metrics:
  print("Two runs had identical metrics. We might use such a comparison to decide whether to update a model, etc "
        "(although in practice it's probably unlikely that metrics would have the exact same values across runs)")
  
if run.data.params == new_run.data.params:
  print("Two runs had identical params. I'm not entirely sure when such a comparison would be useful (more common "
        "might be to use the search API to query for an existing run with the same parameters before launching "
        "a new run")

2_new_api_alternative.py

"""
This file contains an example of querying a run's params, metrics, and tags using an alternative RunData interface
(Option 2). Any feedback is much appreciated, especially regarding query workflows that we should consider while
designing the new interface.
"""
run = mlflow.get_run("925dbee9028f415fb66d1005a9d90f2f")

# 1. Simple example of looking up metric/param/tag values
tag_key = "mlflow.docker"
metric_key = "rmse"
param_key = "alpha"
tag_obj = run.data.tags[tag_key]
metric_obj = run.data.metrics[metric_key]
param_obj = run.data.params[param_key]
print("Tag key: %s, value: %s" % (tag_key, tag_obj.value))
print("Metric key: %s, value: %s" % (metric_key, metric_obj.value))
print("Param key: %s, value: %s" % (param_key, param_obj.value))

# 2. We can access the most-recent timestamps for each metric key from the metric objects in RunData
client = mlflow.tracking.MlflowClient()
print([metric.timestamp for metric in run.data.metrics.values()])

# 3. Some examples of comparing data across runs. Directly comparing metrics/params/tags won't work with the current
# entity implementations, but we can make it work by adding an "equals" implementation to the entities e.g. as in
# https://github.com/mlflow/mlflow/pull/961/files.
new_run = mlflow.get_run("...")
if run.data.metrics == new_run.data.metrics:
  print("Two runs had identical metrics. We might use such a comparison to decide whether to update a model, etc "
        "(although in practice it's probably unlikely that metrics would have the exact same values across runs)")
  
if run.data.params == new_run.data.params:
  print("Two runs had identical params. I'm not entirely sure when such a comparison would be useful (more common "
        "might be to use the search API to query for an existing run with the same parameters before launching "
        "a new run")

3_old_api.py

"""
This file contains an example of querying a run's params, metrics, and tags using the existing RunData API.
"""
# 1. Simple example of looking up metric/param/tag values. Note that it's quite complex to look up
# a metric/param/tag by key.
run = mlflow.get_run("...")
tag_obj = [tag for tag in run.data.tags if tag.key == "mlflow.docker"][0]
metric_obj = [metric for metric in run.data.metrics if metric.key == "rmse"][0]
param_obj = [param for param in run.data.params if param.key == "alpha"][0]
print("Tag key: %s, value: %s" % (tag.key, tag.value))
print("Metric key: %s, value: %s, timestamp: %s" % (metric.key, metric.value, metric.timestamp))
print("Param key: %s, value: %s" % (param_obj.key, param_obj.value))