MLflow - Tapedrive

MLflow tracks machine learning experiments and stores each run’s artifacts in an artifact store. Pointing its S3 artifact store at the gateway takes one environment variable.

Before you start

You need a bucket and a credential. A bucket is a tape that has delegated writes to the gateway, and a credential is an access key pair the gateway operator issues you. Set up a bucket walks through both. A training run logs many files. Run it against your own gateway (gateway setup) rather than a public one.

Configure

export MLFLOW_S3_ENDPOINT_URL=https://s3.example.com
export AWS_ACCESS_KEY_ID=TAPE4Q7ZK2M9XH3PD8AW
export AWS_SECRET_ACCESS_KEY=mZ9vQ3nL8kR2tW6yB4cX7dF1hJ5pS0aG3uE8iO2q

Keep the tracking store local and point the experiment’s artifact location at the bucket:

import mlflow

mlflow.set_tracking_uri("sqlite:///mlflow.db")
experiment_id = mlflow.create_experiment("demo", artifact_location="s3://my-bucket/mlflow")

Everyday calls

with mlflow.start_run(experiment_id=experiment_id) as run:
    # one file, a whole directory, a dict, a string
    mlflow.log_artifact("model.pkl")
    mlflow.log_artifacts("./plots", artifact_path="plots")
    mlflow.log_dict({"lr": 0.001}, "params.json")
    mlflow.log_text("notes", "notes.txt")

# list and fetch
mlflow.artifacts.list_artifacts(run_id=run.info.run_id)
path = mlflow.artifacts.download_artifacts(run_id=run.info.run_id, artifact_path="model.pkl")

A 24 MiB artifact uploads in about 2 s and downloads in well under a second. Downloaded files match the originals byte for byte.

Good to know

The bucket must exist before MLflow starts. MLflow cannot create one. Buckets are created with tape create.
Every artifact is a transaction on the chain, so the cost of a run depends on how many files it logs, not how large they are.

Next

AWS CLI and boto3 for the client MLflow uses underneath.
Uploads for multipart uploads and overwrites.