Blog post: Versioning Machine Learning Models with Metaflow Tags

filter_runs.py

from metaflow import Flow


def get_generate_data_run():
    return Flow('GenerateData').latest_successful_run


def get_train_run(tags: list):
    print(f'Retrieving run data with tags {tags}')
    run_list = [
        run
        for run in list(Flow('Train').runs(*tags))
        if run.successful
    ]
    if len(run_list) == 0:
        raise RuntimeError(f'No runs with tags {tags} found')
    return run_list[0]

train_predict_quickstart.sh

get_tagged_train_run_id() {
  flow_name=$1
  tag_name=$2
  result=`python -c "from metaflow import Flow;print([run for run in list(Flow('$flow_name').runs('$tag_name')) if run.successful][0].data.train_run_id)"`
  echo $result
}

get_tagged_run_id() {
  flow_name=$1
  tag_name=$2
  result=`python -c "from metaflow import Flow;print([run for run in list(Flow('$flow_name').runs('$tag_name')) if run.successful][0].id)"`
  echo $result
}

validate_run_ids() {
  if [ "$1" == "$2" ]; then echo "Run IDs match"; return 0; else echo "Run IDs do not match"; return 1; fi
}

# Make sure your venv is activated if you're using one
python -m pip install -r example-requirements.txt

# Generate training and prediction data
python examples/model-versioning/generate_data.py run --configuration test_config 

# Do a v1 production model training run
python examples/model-versioning/train.py run --configuration test_config --tag prod:v1

# Now do a v1 experimental model training run
python examples/model-versioning/train.py run --configuration test_config --tag test:v1

# Predict on the v1 production model
python examples/model-versioning/predict.py run --configuration test_config --train_tag prod:v1 --tag predict_prod_v1

# Predict on the v1 exerimental model
python examples/model-versioning/predict.py run --configuration test_config --train_tag test:v1 --tag predict_test_v1

# Validate the tagged train run IDs and the one used by the corresponding tagged Predict runs
# These should succeed
validate_run_ids `get_tagged_train_run_id Predict predict_prod_v1` `get_tagged_run_id Train prod:v1`
validate_run_ids `get_tagged_train_run_id Predict predict_test_v1` `get_tagged_run_id Train test:v1`
# This should fail
validate_run_ids `get_tagged_train_run_id Predict predict_prod_v1` `get_tagged_run_id Train test:v1`