Skip to content

Mlflow load_model support for runs

Context

Mlflow load_model supports a myriad of ways to load artifacts for a model:

We have support for models:/<model_name>/<model_version>

This issue is to add support for runs:/<mlflow_run_id>/run-relative/path/to/model

The location, in URI format, of the MLflow model. For example:

  • /Users/me/path/to/local/model
  • relative/path/to/local/model
  • s3://my_bucket/path/to/model
  • runs:/<mlflow_run_id>/run-relative/path/to/model
  • models:/<model_name>/<model_version>
  • models:/<model_name>/
  • mlflow-artifacts:/path/to/model

We have an error currently like this.

  File "/usr/local/lib/python3.10/dist-packages/mlflow/store/artifact/http_artifact_repo.py", line 84, in list_artifacts
    url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
ValueError: not enough values to unpack (expected 2, got 1)`

Discussion from #506768 (comment 2257465885)

mlflow.pytorch.load_model('runs:/3edb1ff9-ee52-4903-bfac-21fef5629b3f/model')
mlflow.pytorch.load_model('https://gitlab.com/api/v4/projects/<:project_id>/packages/generic/ml_experiment_14/104/unet')

Proposal

  • Add support for load_model with runs:/<mlflow_run_id>/run-relative/path/to/model
  • Ensure subfolders can be loaded
  • Add documentation
Edited by Alper Akgun