MLflow token refresh failure during long-running training jobs

Problem – Intermittent 401 Unauthorized Errors During Long‑Running Training

In a multi‑node GPU cluster orchestrated by Kubernetes, a distributed PyTorch/TensorFlow training job pushes metrics, parameters, and artifacts to a remote MLflow tracking server that is protected by an OIDC/OAuth2 provider. After several hours of execution the training process receives repeated 401 Unauthorized responses:

HTTPError 401 – Unauthorized while sending POST to /api/2.0/mlflow/runs/log-metric
mlflow.exceptions.RestException: Request failed with status code 401: Invalid or expired token
WARN mlflow.tracking.client: Token refresh failed: response status 401 (invalid_grant)
ERROR mlflow.tracking: Failed to log artifact: HTTPError 401 – Unauthorized (token expired)

Symptoms observed in the pod logs:

• Metric logging stops after ~2 hours (token TTL = 1 hour in the OIDC provider).
• Training pod crashes with an unhandled RestException inside the logging callback.
• Subsequent epochs produce no artifact uploads.

Impact:

• Loss of telemetry for the remainder of the experiment.
• Job termination or manual restart required.
• Potential waste of expensive GPU time.

Root Cause – Token Refresh Path Broken in the MLflow Client

The MLflow tracking client authenticates by sending the Authorization: Bearer <access_token> header on every REST call (MLflow REST API authentication guide). When the access token expires, the client attempts to obtain a new token using the configured MLFLOW_AUTHENTICATION_PROVIDER (e.g., mlflow.auth.oidc.OIDCAuthenticationProvider) and the MLFLOW_TRACKING_TOKEN environment variable as a refresh token.

Two conditions cause the failure:

1. Short‑lived access tokens with no automatic refresh. By default the client does not schedule a refresh; it only retries after receiving a 401. In long‑running jobs the refresh request is sent after the token has already been invalidated, causing the provider to return invalid_grant (see the warning log above).
2. Kubernetes sidecar or network partition prevents the refresh request. As described in GitHub issue mlflow#7321, a sidecar proxy that injects tokens may lose connectivity, so the refresh HTTP call fails with 401 before the client can obtain a new token.

Consequently, the client never receives a valid token again, and all subsequent logging attempts fail.

Debug – Systematic Investigation Steps

1. Verify token lifetime and expiration

# Inside a training pod
echo $MLFLOW_TRACKING_TOKEN | cut -d'.' -f2 | base64 -d | jq .
{
  "exp": 1698765432,
  "iat": 1698761832,
  "iss": "https://auth.mycompany.com/"
}

Compare the exp value with the job start time to confirm the token expires before the job finishes.

2. Capture the failing HTTP request

kubectl exec -it $POD_NAME -- tcpdump -i any -w /tmp/mlflow.pcap host tracking.mycompany.com and port 443
# later, analyze with Wireshark:
# Look for POST /api/2.0/mlflow/runs/log-metric with 401 response

3. Check MLflow client configuration

# Environment in the container
printenv | grep MLFLOW
MLFLOW_TRACKING_URI=https://tracking.mycompany.com
MLFLOW_AUTHENTICATION_PROVIDER=mlflow.auth.oidc.OIDCAuthenticationProvider
MLFLOW_TRACKING_TOKEN=/var/run/secrets/oidc/token
MLFLOW_TRACKING_REFRESH_INTERVAL=1800   # missing? default is 0

If MLFLOW_TRACKING_REFRESH_INTERVAL is unset or zero, the client will not proactively refresh.

4. Review sidecar logs (if used)

kubectl logs $POD_NAME -c oidc-sidecar
2024-07-28T12:45:02Z WARN token refresh failed: response status 401 (invalid_grant)

5. Reproduce locally

Run a short script that sleeps longer than the token TTL while logging a metric every minute. Observe the same 401 after token expiration, confirming the issue is client‑side.

Solution – Enable Reliable Token Refresh

Option 1: Use the built‑in periodic refresh interval

Set MLFLOW_TRACKING_REFRESH_INTERVAL (seconds) to a value smaller than the token TTL. The client will request a new access token before the current one expires.

# Deployment manifest snippet
env:
  - name: MLFLOW_TRACKING_URI
    value: "https://tracking.mycompany.com"
  - name: MLFLOW_AUTHENTICATION_PROVIDER
    value: "mlflow.auth.oidc.OIDCAuthenticationProvider"
  - name: MLFLOW_TRACKING_TOKEN
    valueFrom:
      secretKeyRef:
        name: oidc-token
        key: token
  - name: MLFLOW_TRACKING_REFRESH_INTERVAL
    value: "1500"   # 25 minutes, for a 1‑hour token TTL

Option 2: Provide a custom token provider script

When sidecar injection is unreliable, replace it with a script that fetches a fresh token before each logging call.

# token_provider.py
import os, subprocess, time, json, base64

def get_token():
    # Use client credentials grant to fetch a new access token
    result = subprocess.check_output([
        "curl", "-s", "-X", "POST",
        "-d", "grant_type=client_credentials",
        "-d", "client_id=$OIDC_CLIENT_ID",
        "-d", "client_secret=$OIDC_CLIENT_SECRET",
        "https://auth.mycompany.com/oauth/token"
    ])
    return json.loads(result)['access_token']

def refresh_token():
    token = get_token()
    os.environ["MLFLOW_TRACKING_TOKEN"] = token

# Hook into MLflow's tracking client
import mlflow
mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
mlflow.tracking._tracking_client._TrackingClient._refresh_token = refresh_token

Import this module at the start of the training script:

import token_provider   # ensures token refresh hook is installed
mlflow.start_run()

Option 3: Increase token TTL on the OIDC provider (if policy permits)

Configure the provider to issue tokens with a TTL longer than the expected maximum job runtime (e.g., 12 hours). This eliminates the need for frequent refreshes but may conflict with security policies.

Before / After Comparison

MLFLOW_TRACKING_REFRESH_INTERVAL=0
MLFLOW_TRACKING_TOKEN=/var/run/secrets/oidc/token

MLFLOW_TRACKING_REFRESH_INTERVAL=1500
MLFLOW_TRACKING_TOKEN=/var/run/secrets/oidc/token

Verify – Confirming the Fix Works

1. Observe token refresh logs

2024-07-28T13:20:12Z INFO mlflow.tracking.client: Refreshing access token
2024-07-28T13:20:13Z INFO mlflow.tracking.client: Token refreshed successfully, expires at 2024-07-28T14:20:13Z

2. Check that metric logging continues after the original expiry time

# Tail the training pod logs for the full duration
kubectl logs -f $POD_NAME | grep "log-metric"
2024-07-28T12:00:00Z INFO mlflow.tracking.client: Logged metric step=0
...
2024-07-28T14:00:00Z INFO mlflow.tracking.client: Logged metric step=120

3. Verify no 401 entries in the pod logs

kubectl logs $POD_NAME | grep "401"
# No output

Prevent – Operational Guardrails and Best Practices

• Set a refresh interval. Always define MLFLOW_TRACKING_REFRESH_INTERVAL to a value < ½ × token TTL.
• Centralize token acquisition. Use a sidecar only if it can guarantee network connectivity; otherwise prefer a script‑based provider.
• Monitor token‑related metrics. Export a custom Prometheus metric (e.g., mlflow_token_refresh_success_total) and alert on consecutive failures.
• Graceful error handling. Wrap logging calls in try/except blocks that retry on RestException with 401, allowing the client to re‑authenticate.
• Document token TTLs. Include the provider’s access_token_lifetime in the run configuration so engineers can size MLFLOW_TRACKING_REFRESH_INTERVAL appropriately.

Related Topic Hub: Model Serving Troubleshooting Hub

FAQ

1. Why does the MLflow client still receive 401 after I set MLFLOW_TRACKING_REFRESH_INTERVAL?
   The refresh interval is ignored if the environment variable is misspelled or if the client is instantiated before the variable is set. Ensure the variable is defined at pod start and that the MLflow client is imported after the environment is ready.
2. Can I rely on the sidecar token injection without a refresh interval?
   Only if the sidecar itself performs proactive refreshes and updates MLFLOW_TRACKING_TOKEN in place. Most sidecars expose a static token, so the client must handle refreshes.
3. How do I debug a failed token refresh when the OIDC provider returns invalid_grant?
   Inspect the refresh request payload (client_id, client_secret, grant_type). A common cause is a revoked refresh token or clock skew between the pod and the provider. Synchronize the node time (e.g., chrony) and verify the refresh token is still valid.
4. Is increasing the token TTL a safe long‑term solution?
   Longer TTL reduces refresh traffic but expands the attack window if a token is compromised. Prefer periodic refresh with short TTL for production environments.
5. What metric should I alert on to catch token‑related failures early?
   Create an alert on the rate of mlflow_rest_exception_total{status="401"} exceeding a threshold (e.g., > 5 per minute) or on the absence of successful mlflow_token_refresh_success_total within the last refresh interval.