mlflow-export-import
mlflow-export-import copied to clipboard
Add token support
Proposed Changes:
I changed the function get_mlflow_host_token()
in the module mlflow_utils.py
, such that it loads a tracking token if it's found in the in environment variables.
The fix required only a very small change. The function _mk_headers()
in the module http_client.py
is already written in such a way that it automatically sets the tracking token if present, so no changes are further necessary.
Motivation:
We want to use mlflow-export-import
to migrate models between OOS tracking servers in an enterprise setting (at a bank). However, since our tracking servers are both behind oauth2 proxies, support for bearer tokens is essential for us to make it work. I decided to make a PR because it can also be useful to others.