metaflow
metaflow copied to clipboard
Support token refresh in the Python client
We have a reverse proxy in front of the Metaflow service implementing JWT based authentication. To pass the token, we're currently doing this:
export METAFLOW_SERVICE_HEADERS="{\"Authorization\": \"Bearer $METAFLOW_TOKEN\"}"
The problem is that these tokens have a short lifetime and can easily expire during a training session. Thus some refresh mechanism would need to be supported.
I suggest following the route that MLflow has taken: https://www.mlflow.org/docs/latest/plugins.html#run-code-using-the-plugin. This would allow us to register a plugin, which would then implement the token refresh functionality.
I agree that this would be a very important feature since especially production ready mlflow servers usually need authentification and especially deep learning experiemnts can usually run for days. In such cases a token will most likely expire. I guess a simple callback that can be passed to create a new token on such event would be a possible way to implement that.
In our case we currently renew the token with the following wrapper function to call mlflow functions:
def __call_and_check_token(fun, *args):
try:
return fun(*args)
except mlflow.exceptions.RestException as e:
print("Token probably expired. Create new token...")
os.environ["MLFLOW_TRACKING_TOKEN"] = create_token()
return fun(*args)
__call_and_check_token(mlflow.log_param, 'a', 5)