From 0f62e81ba144f0c1ddc4d0531a36539ca4b9649a Mon Sep 17 00:00:00 2001
From: Sergey Serebryakov <serebryakov.sergey@gmail.com>
Date: Wed, 2 Aug 2023 15:43:25 -0700
Subject: [PATCH] Refactoring Platform.get_installed_runner method. (#329)

- Removing unused dictionary of reference runners.
- Moving the import of runner packages withing try/except block not to fail when candidate mLCube runner fails to load.
- Simplifying the reporting logic.
- Adding new method (Platform.get_package_info) that returns additional info for the runner package for logging purposes.
---
 mlcube/mlcube/platform.py | 78 ++++++++++++++++++++++-----------------
 1 file changed, 45 insertions(+), 33 deletions(-)

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)."