wehub-resource-sync
67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
import warnings
|
|
|
|
from mlflow.tracking.request_auth.kubernetes_request_auth_provider import (
|
|
KubernetesNamespacedRequestAuthProvider,
|
|
KubernetesRequestAuthProvider,
|
|
)
|
|
from mlflow.utils.plugins import get_entry_points
|
|
|
|
REQUEST_AUTH_PROVIDER_ENTRYPOINT = "mlflow.request_auth_provider"
|
|
|
|
|
|
class RequestAuthProviderRegistry:
|
|
def __init__(self):
|
|
self._registry = []
|
|
|
|
def register(self, request_auth_provider):
|
|
self._registry.append(request_auth_provider())
|
|
|
|
def register_entrypoints(self):
|
|
for entrypoint in get_entry_points(REQUEST_AUTH_PROVIDER_ENTRYPOINT):
|
|
try:
|
|
self.register(entrypoint.load())
|
|
except (AttributeError, ImportError) as exc:
|
|
warnings.warn(
|
|
'Failure attempting to register request auth provider "{}": {}'.format(
|
|
entrypoint.name, str(exc)
|
|
),
|
|
stacklevel=2,
|
|
)
|
|
|
|
def __iter__(self):
|
|
return iter(self._registry)
|
|
|
|
|
|
_request_auth_provider_registry = RequestAuthProviderRegistry()
|
|
_request_auth_provider_registry.register(KubernetesRequestAuthProvider)
|
|
_request_auth_provider_registry.register(KubernetesNamespacedRequestAuthProvider)
|
|
_request_auth_provider_registry.register_entrypoints()
|
|
|
|
|
|
def fetch_auth(request_auth):
|
|
"""
|
|
Find the request auth from registered providers based on the auth provider's name.
|
|
The auth provider's name can be provided through environment variable `MLFLOW_TRACKING_AUTH`.
|
|
|
|
This function iterates through all request auth providers in the registry. Additional context
|
|
providers can be registered as described in
|
|
:py:class:`mlflow.tracking.request_auth.RequestAuthProvider`.
|
|
|
|
Args:
|
|
request_auth: The name of request auth provider.
|
|
|
|
Returns:
|
|
The auth object.
|
|
"""
|
|
|
|
for auth_provider in _request_auth_provider_registry:
|
|
if auth_provider.get_name() == request_auth:
|
|
return auth_provider.get_auth()
|
|
|
|
warnings.warn(
|
|
f"Could not find any registered plugin for {request_auth}. "
|
|
"No authentication header will be added. Please check your "
|
|
"provider documentation for installing the right plugin or "
|
|
"correct provider name."
|
|
)
|