diff --git a/mlcube/mlcube/platform.py b/mlcube/mlcube/platform.py index c6aa0e4..63e03a6 100644 --- a/mlcube/mlcube/platform.py +++ b/mlcube/mlcube/platform.py @@ -6,28 +6,31 @@ import logging import pkgutil import typing as t - -from mlcube.runner import Runner +from types import ModuleType from omegaconf import DictConfig +from mlcube.runner import Runner logger = logging.getLogger(__name__) -__all__ = ['Platform'] +__all__ = ["Platform"] class Platform(object): """Class to manage reference Python-based MLCube runners.""" - reference_runners = { - 'docker': {'pkg': 'mlcube_docker'}, - 'singularity': {'pkg': 'mlcube_singularity'}, - 'ssh': {'pkg': 'mlcube_ssh'}, - 'gcp': {'pkg': 'mlcube_gcp'}, - 'k8s': {'pkg': 'mlcube_k8s'}, - 'kubeflow': {'pkg': 'mlcube_kubeflow'}, - } + @staticmethod + def get_package_info(package: ModuleType) -> t.Dict: + """Return information about the Python package""" + info = {"name": package.__name__, "location": package.__file__} + try: + from importlib.metadata import version + + info["version"] = version(info["name"]) + except: + pass + return info @staticmethod def get_installed_runners() -> t.Dict[str, t.Dict]: @@ -44,32 +47,37 @@ def get_installed_runners() -> t.Dict[str, t.Dict]: Installed runners are found by inspecting Python packages. MLCube system settings file is not used. """ - candidate_runners = { - name: importlib.import_module(name) for _, name, _ in pkgutil.iter_modules() if name.startswith('mlcube_') - } installed_runners = {} - for pkg_name, module in candidate_runners.items(): + for _, pkg_name, _ in pkgutil.iter_modules(): + if not pkg_name.startswith("mlcube_"): + continue + module_info: t.Optional[t.Dict] = None try: - get_runner_class: t.Optional[t.Callable] = getattr(module, 'get_runner_class', None) - if get_runner_class is None: - logger.warning( - "Platform.get_installed_runners candidate MLCube runner package (%s) does not provide runner " - "class function (get_runner_class).", pkg_name - ) - continue - runner_cls: t.Type[Runner] = get_runner_class() + module = importlib.import_module(pkg_name) + module_info = Platform.get_package_info(module) + runner_cls: t.Type[Runner] = module.get_runner_class() if not issubclass(runner_cls, Runner): - raise TypeError(f"Invalid type of a runner ({runner_cls}). Expecting subclass of {Runner}.") + raise TypeError( + f"Invalid runner type (expected: {Runner}, actual: {runner_cls})." + ) runner_name = runner_cls.CONFIG.DEFAULT.runner - installed_runners[runner_name] = {'config': {'pkg': pkg_name}, 'runner_cls': runner_cls} + installed_runners[runner_name] = { + "config": {"pkg": pkg_name}, + "runner_cls": runner_cls, + } logger.info( - "Platform.get_installed_runners found installed MLCube runner: platform=%s, pkg=%s", - runner_name, pkg_name + "Platform.get_installed_runners found installed MLCube runner (platform=%s, pkg=%s, info=%s)", + runner_name, + pkg_name, + module_info, ) - except (AttributeError, TypeError) as e: + except (ImportError, AttributeError, TypeError) as e: logger.warning( - "Platform.get_installed_runners package (%s) is not a valid MLCube runner. Error=\"%s\"", - pkg_name, str(e) + "Platform.get_installed_runners package (pkg_name=%s, info=%s) is not a valid MLCube runner. " + 'Error="%s".', + pkg_name, + module_info, + str(e), ) return installed_runners @@ -84,10 +92,14 @@ def get_runner(runner_config: t.Optional[DictConfig]) -> t.Type[Runner]: """ if not runner_config: raise RuntimeError("Can't create runner. Runner config is null or empty.") - if 'pkg' not in runner_config: - raise RuntimeError(f"Do not know how to instantiate a runner. Runner config={str(runner_config)}") + if "pkg" not in runner_config: + raise RuntimeError( + f"Do not know how to instantiate a runner. Runner config={str(runner_config)}" + ) module = importlib.import_module(runner_config.pkg) - get_runner_class: t.Optional[t.Callable] = getattr(module, 'get_runner_class', None) + get_runner_class: t.Optional[t.Callable] = getattr( + module, "get_runner_class", None + ) if get_runner_class is None: raise RuntimeError( f"Imported module ({runner_config.pkg}) does not provide runner class function (get_runner_class)."