add client side kubernetes auth plugin

This mlflow plugin is intended to be installed and used by the client, when headers are missing workspace/tokens it will search for these values in the pod's filesystem (when running in a pod), or use the local kubeconfig.

The plugin is intentionally added as part of the mlflow package to avoid having to package it separately. This is something we may do in the future.

(cherry picked from commit 40260db)

Signed-off-by: Humair Khan <HumairAK@users.noreply.github.com>

Author: HumairAK

mlflow/tracking/request_auth/kubernetes_request_auth_provider.md

@@ -0,0 +1,30 @@
+# Request Auth Provider Plugin
+
+This plugin system allows you to add custom authentication to MLflow tracking requests.
+
+## Usage
+
+Set the `MLFLOW_TRACKING_AUTH` environment variable to the name of your auth provider:
+
+```bash
+export MLFLOW_TRACKING_AUTH=kubernetes
+```
+
+MLflow will then use that provider to add authentication headers to all outgoing tracking requests.
+
+## Built-in Providers
+
+### Kubernetes
+
+Adds authentication headers for Kubernetes environments.
+
+```bash
+export MLFLOW_TRACKING_AUTH=kubernetes
+```
+
+This provider automatically adds:
+
+- `X-MLFLOW-WORKSPACE`: The Kubernetes namespace (from service account or kubeconfig)
+- `Authorization`: Bearer token (from service account or kubeconfig)
+
+It first checks for in-cluster service account credentials, then falls back to your local kubeconfig.

mlflow/tracking/request_auth/kubernetes_request_auth_provider.py

@@ -0,0 +1,190 @@
+"""Request auth provider for Kubernetes environments.
+
+This provider automatically adds workspace and authorization headers when running
+in Kubernetes environments or when kubeconfig is available.
+
+To use this provider, set the environment variable:
+    MLFLOW_TRACKING_AUTH=kubernetes
+"""
+
+import logging
+from pathlib import Path
+
+from cachetools import TTLCache
+from kubernetes import client, config
+
+from mlflow.exceptions import MlflowException
+from mlflow.tracking.request_auth.abstract_request_auth_provider import RequestAuthProvider
+
+# Cache for file reads (1 minute TTL)
+_FILE_CACHE_TTL = 60
+_file_cache: TTLCache = TTLCache(maxsize=10, ttl=_FILE_CACHE_TTL)
+
+_logger = logging.getLogger(__name__)
+
+# Kubernetes service account paths
+_SERVICE_ACCOUNT_NAMESPACE_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
+_SERVICE_ACCOUNT_TOKEN_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/token")
+
+# Header names
+WORKSPACE_HEADER_NAME = "X-MLFLOW-WORKSPACE"
+AUTHORIZATION_HEADER_NAME = "Authorization"
+
+
+def _read_file_if_exists(path: Path) -> str | None:
+    """Read a file and return its contents stripped, or None if it doesn't exist."""
+    cache_key = str(path)
+    if cache_key in _file_cache:
+        return _file_cache[cache_key]
+
+    result = None
+    try:
+        if path.exists():
+            result = path.read_text().strip() or None
+    except (OSError, PermissionError) as e:
+        _logger.debug("Could not read file %s: %s", path, e)
+
+    _file_cache[cache_key] = result
+    return result
+
+
+def _get_credentials_from_service_account() -> tuple[str, str] | None:
+    """Get namespace and token from mounted service account files.
+
+    Returns:
+        Tuple of (namespace, token) if both are available, None otherwise.
+    """
+    namespace = _read_file_if_exists(_SERVICE_ACCOUNT_NAMESPACE_PATH)
+    token = _read_file_if_exists(_SERVICE_ACCOUNT_TOKEN_PATH)
+
+    if namespace and token:
+        return namespace, f"Bearer {token}"
+    return None
+
+
+def _get_credentials_from_kubeconfig() -> tuple[str, str] | None:
+    """Get namespace and token from kubeconfig.
+
+    Uses ApiClient so kubeconfig exec auth is resolved when possible
+    (EKS, GKE, AKS, OpenShift, OIDC).
+
+    Returns:
+        Tuple of (namespace, token) if both are available, None otherwise.
+    """
+    config.load_kube_config()
+
+    # Get namespace from context
+    _, active_context = config.list_kube_config_contexts()
+    if not active_context:
+        return None
+
+    namespace = active_context.get("context", {}).get("namespace", "").strip() or None
+    if not namespace:
+        return None
+
+    # Get token from ApiClient
+    api_client = client.ApiClient()
+    token = None
+
+    # 1) Try default_headers first (where most auth flows put the resolved token)
+    auth = api_client.default_headers.get("Authorization") or api_client.default_headers.get(
+        "authorization"
+    )
+    if isinstance(auth, str) and auth.lower().startswith("bearer "):
+        token = auth[7:].strip()
+
+    # 2) Fallback: configuration.api_key (some versions/auth flows store it here)
+    if not token:
+        api_key = api_client.configuration.api_key.get("authorization")
+        if isinstance(api_key, str):
+            token = api_key.strip()
+            if token.lower().startswith("bearer "):
+                token = token[7:].strip()
+
+    if not token:
+        return None
+
+    return namespace, f"Bearer {token}"
+
+
+def _get_credentials() -> tuple[str, str] | None:
+    """Get workspace and authorization credentials.
+
+    Tries service account files first, then falls back to kubeconfig.
+    Both values must come from the same source for consistency.
+
+    Returns:
+        Tuple of (namespace, authorization) if available, None otherwise.
+    """
+    # Try service account files first (running in a pod)
+    if creds := _get_credentials_from_service_account():
+        return creds
+
+    # Fallback to kubeconfig
+    return _get_credentials_from_kubeconfig()
+
+
+class KubernetesAuth:
+    """Custom authentication class for Kubernetes environments.
+
+    This class is callable and will be invoked by the requests library
+    to add authentication headers to each request.
+    """
+
+    def __call__(self, request):
+        """Add Kubernetes authentication headers to the request.
+
+        Args:
+            request: The prepared request object from the requests library.
+
+        Returns:
+            The modified request object with authentication headers.
+
+        Raises:
+            MlflowException: If workspace or authorization cannot be determined.
+        """
+        # Skip if both headers are already set
+        if (
+            WORKSPACE_HEADER_NAME in request.headers
+            and AUTHORIZATION_HEADER_NAME in request.headers
+        ):
+            return request
+
+        credentials = _get_credentials()
+        if not credentials:
+            raise MlflowException(
+                "Could not determine Kubernetes credentials. "
+                "Ensure you are running in a Kubernetes pod with a service account "
+                "or have a valid kubeconfig with a namespace and credentials set."
+            )
+
+        namespace, authorization = credentials
+
+        if WORKSPACE_HEADER_NAME not in request.headers:
+            request.headers[WORKSPACE_HEADER_NAME] = namespace
+
+        if AUTHORIZATION_HEADER_NAME not in request.headers:
+            request.headers[AUTHORIZATION_HEADER_NAME] = authorization
+
+        return request
+
+
+class KubernetesRequestAuthProvider(RequestAuthProvider):
+    """Provides authentication for Kubernetes environments.
+
+    This provider adds headers based on Kubernetes environment:
+    - X-MLFLOW-WORKSPACE: Set from service account namespace file or kubeconfig context
+    - Authorization: Set from service account token file or kubeconfig credentials
+
+    Credentials are sourced consistently - either both from mounted service account
+    files (when running in a pod) or both from kubeconfig (when running locally).
+
+    To enable this provider, set:
+        MLFLOW_TRACKING_AUTH=kubernetes
+    """
+
+    def get_name(self) -> str:
+        return "kubernetes"
+
+    def get_auth(self):
+        return KubernetesAuth()

mlflow/tracking/request_auth/registry.py

@@ -1,5 +1,8 @@
 import warnings
 
+from mlflow.tracking.request_auth.kubernetes_request_auth_provider import (
+    KubernetesRequestAuthProvider,
+)
 from mlflow.utils.plugins import get_entry_points
 
 REQUEST_AUTH_PROVIDER_ENTRYPOINT = "mlflow.request_auth_provider"
@@ -29,6 +32,7 @@ def __iter__(self):
 
 
 _request_auth_provider_registry = RequestAuthProviderRegistry()
+_request_auth_provider_registry.register(KubernetesRequestAuthProvider)
 _request_auth_provider_registry.register_entrypoints()
 
 

pyproject.release.toml

@@ -37,6 +37,7 @@ dependencies = [
   "graphene<4",
   "gunicorn<24; platform_system != 'Windows'",
   "huey<3,>=2.5.0",
+  "kubernetes<35",
   "matplotlib<4",
   "numpy<3",
   "pandas<3",

pyproject.toml

@@ -42,6 +42,7 @@ dependencies = [
   "gunicorn<24; platform_system != 'Windows'",
   "huey<3,>=2.5.0",
   "importlib_metadata<9,>=3.7.0,!=4.7.0",
+  "kubernetes<35",
   "matplotlib<4",
   "numpy<3",
   "opentelemetry-api<3,>=1.9.0",

requirements/core-requirements.yaml

@@ -80,3 +80,7 @@ huey:
   pip_release: huey
   minimum: "2.5.0"
   max_major_version: 2
+
+kubernetes:
+  pip_release: kubernetes
+  max_major_version: 34