Coverage for src/mlopus/mlflow/utils.py: 85%

1import os 

2from typing import Dict, Any, Type, TypeVar, List 

3 

4from mlopus.utils import import_utils, dicts 

5from .api.base import BaseMlflowApi 

6 

7API = TypeVar("API", bound=BaseMlflowApi) 

8"""Type of :class:`BaseMlflowApi`""" 

9 

10PLUGIN_GROUP = "mlopus.mlflow_api_providers" 

11 

12 

13def list_api_plugins() -> List[import_utils.EntryPoint]: 

14 """List all API plugins available in this environment.""" 

15 return import_utils.list_plugins(PLUGIN_GROUP) 

16 

17 

18def get_api( 

19 plugin: str | None = None, 

20 cls: Type[API] | str | None = None, 

21 conf: Dict[str, Any] | None = None, 

22) -> BaseMlflowApi | API: 

23 """Load MLflow API class or plugin with specified configuration. 

24 

25 The default API class is :class:`~mlopus.mlflow.providers.mlflow.MlflowApi` 

26 (registered under the plugin name `mlflow`). 

27 

28 :param plugin: | Plugin name from group `mlopus.mlflow_api_providers`. 

29 | Incompatible with :paramref:`cls`. 

30 

31 :param cls: | A type that implements :class:`BaseMlflowApi` or a fully qualified class name of such a type 

32 (e.g.: `package.module:Class`). 

33 | Incompatible with :paramref:`plugin`. 

34 

35 :param conf: | A `dict` of keyword arguments for the resolved API class. 

36 | See :func:`api_conf_schema` 

37 

38 :return: API instance. 

39 """ 

40 return _get_api_cls(plugin, cls).parse_obj(dicts.deep_merge(_get_env_conf(), conf or {})) 

41 

42 

43def api_conf_schema( 

44 plugin: str | None = None, 

45 cls: Type[API] | str | None = None, 

46) -> dicts.AnyDict: 

47 """Get configuration schema for MLflow API class or plugin. 

48 

49 :param plugin: | See :paramref:`get_api.plugin`. 

50 

51 :param cls: | See :paramref:`get_api.cls`. 

52 """ 

53 return _get_api_cls(plugin, cls).schema() 

54 

55 

56def _get_api_cls( 

57 plugin: str | None = None, 

58 cls: Type[API] | str | None = None, 

59) -> Type[API]: 

60 assert None in (plugin, cls), "`plugin` and `cls` are mutually excluding." 

61 

62 if isinstance(cls, str): 

63 cls = import_utils.find_type(cls, BaseMlflowApi) 

64 elif cls is None: 

65 cls = import_utils.load_plugin(PLUGIN_GROUP, plugin or "mlflow", BaseMlflowApi) 

66 

67 return cls 

68 

69 

70def _get_env_conf() -> Dict[str, Any]: 

71 conf = {} 

72 for k, v in os.environ.items(): 

73 if v and k.startswith(prefix := "MLOPUS_MLFLOW_CONF_"): 

74 dicts.set_nested(conf, k.removeprefix(prefix).lower().split("__"), v) 

75 return conf