wehub-resource-sync
36 lines
932 B
Python
36 lines
932 B
Python
import logging
|
|
|
|
from mlflow.projects.backend.local import LocalBackend
|
|
from mlflow.utils.plugins import get_entry_points
|
|
|
|
ENTRYPOINT_GROUP_NAME = "mlflow.project_backend"
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Statically register backend defined in mlflow
|
|
MLFLOW_BACKENDS = {
|
|
"local": LocalBackend,
|
|
}
|
|
|
|
|
|
def load_backend(backend_name):
|
|
# Static backends
|
|
if backend_name in MLFLOW_BACKENDS:
|
|
return MLFLOW_BACKENDS[backend_name]()
|
|
|
|
# backends from plugin
|
|
entrypoints = get_entry_points(ENTRYPOINT_GROUP_NAME)
|
|
if entrypoint := next((e for e in entrypoints if e.name == backend_name), None):
|
|
builder = entrypoint.load()
|
|
return builder()
|
|
|
|
# TODO Should be a error when all backends are migrated here
|
|
_logger.warning(
|
|
"Backend '%s' is not available. Available plugins are %s",
|
|
backend_name,
|
|
[*entrypoints, *MLFLOW_BACKENDS.keys()],
|
|
)
|
|
|
|
return None
|