metaflow icon indicating copy to clipboard operation
metaflow copied to clipboard

Support token refresh in the Python client

Open dennari opened this issue 4 years ago • 1 comments

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.

dennari avatar Dec 09 '20 10:12 dennari

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)

mbcel avatar Dec 29 '21 13:12 mbcel