1 files changed, 450 insertions, 0 deletions
diff --git a/src/mlia/backend/install.py b/src/mlia/backend/install.py
new file mode 100644
index 0000000..eea3403
--- /dev/null
+++ b/src/mlia/backend/install.py
@@ -0,0 +1,450 @@
+# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Module for installation process."""
+from __future__ import annotations
+
+import logging
+import platform
+import tarfile
+from abc import ABC
+from abc import abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+from typing import Iterable
+from typing import Optional
+from typing import Union
+
+from mlia.backend.executor.runner import BackendRunner
+from mlia.backend.executor.system import remove_system
+from mlia.utils.download import DownloadArtifact
+from mlia.utils.filesystem import all_files_exist
+from mlia.utils.filesystem import all_paths_valid
+from mlia.utils.filesystem import copy_all
+from mlia.utils.filesystem import get_mlia_resources
+from mlia.utils.filesystem import temp_directory
+from mlia.utils.filesystem import working_directory
+from mlia.utils.py_manager import get_package_manager
+
+
+logger = logging.getLogger(__name__)
+
+
+# Mapping backend -> device_type -> system_name
+_SUPPORTED_SYSTEMS = {
+    "Corstone-300": {
+        "ethos-u55": "Corstone-300: Cortex-M55+Ethos-U55",
+        "ethos-u65": "Corstone-300: Cortex-M55+Ethos-U65",
+    },
+    "Corstone-310": {
+        "ethos-u55": "Corstone-310: Cortex-M85+Ethos-U55",
+        "ethos-u65": "Corstone-310: Cortex-M85+Ethos-U65",
+    },
+}
+
+# Mapping system_name -> application
+_SYSTEM_TO_APP_MAP = {
+    "Corstone-300: Cortex-M55+Ethos-U55": "Generic Inference Runner: Ethos-U55",
+    "Corstone-300: Cortex-M55+Ethos-U65": "Generic Inference Runner: Ethos-U65",
+    "Corstone-310: Cortex-M85+Ethos-U55": "Generic Inference Runner: Ethos-U55",
+    "Corstone-310: Cortex-M85+Ethos-U65": "Generic Inference Runner: Ethos-U65",
+}
+
+
+def get_system_name(backend: str, device_type: str) -> str:
+    """Get the system name for the given backend and device type."""
+    return _SUPPORTED_SYSTEMS[backend][device_type]
+
+
+def get_application_name(system_name: str) -> str:
+    """Get application name for the provided system name."""
+    return _SYSTEM_TO_APP_MAP[system_name]
+
+
+def is_supported(backend: str, device_type: str | None = None) -> bool:
+    """Check if the backend (and optionally device type) is supported."""
+    if device_type is None:
+        return backend in _SUPPORTED_SYSTEMS
+
+    try:
+        get_system_name(backend, device_type)
+        return True
+    except KeyError:
+        return False
+
+
+def supported_backends() -> list[str]:
+    """Get a list of all backends supported by the backend manager."""
+    return list(_SUPPORTED_SYSTEMS.keys())
+
+
+def get_all_system_names(backend: str) -> list[str]:
+    """Get all systems supported by the backend."""
+    return list(_SUPPORTED_SYSTEMS.get(backend, {}).values())
+
+
+def get_all_application_names(backend: str) -> list[str]:
+    """Get all applications supported by the backend."""
+    app_set = {_SYSTEM_TO_APP_MAP[sys] for sys in get_all_system_names(backend)}
+    return list(app_set)
+
+
+@dataclass
+class InstallFromPath:
+    """Installation from the local path."""
+
+    backend_path: Path
+
+
+@dataclass
+class DownloadAndInstall:
+    """Download and install."""
+
+    eula_agreement: bool = True
+
+
+InstallationType = Union[InstallFromPath, DownloadAndInstall]
+
+
+class Installation(ABC):
+    """Base class for the installation process of the backends."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return name of the backend."""
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Return description of the backend."""
+
+    @property
+    @abstractmethod
+    def could_be_installed(self) -> bool:
+        """Return true if backend could be installed in current environment."""
+
+    @property
+    @abstractmethod
+    def already_installed(self) -> bool:
+        """Return true if backend is already installed."""
+
+    @abstractmethod
+    def supports(self, install_type: InstallationType) -> bool:
+        """Return true if installation supports requested installation type."""
+
+    @abstractmethod
+    def install(self, install_type: InstallationType) -> None:
+        """Install the backend."""
+
+    @abstractmethod
+    def uninstall(self) -> None:
+        """Uninstall the backend."""
+
+
+@dataclass
+class BackendInfo:
+    """Backend information."""
+
+    backend_path: Path
+    copy_source: bool = True
+    system_config: str | None = None
+
+
+PathChecker = Callable[[Path], Optional[BackendInfo]]
+BackendInstaller = Callable[[bool, Path], Path]
+
+
+class BackendMetadata:
+    """Backend installation metadata."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        system_config: str,
+        apps_resources: list[str],
+        fvp_dir_name: str,
+        download_artifact: DownloadArtifact | None,
+        supported_platforms: list[str] | None = None,
+    ) -> None:
+        """
+        Initialize BackendMetadata.
+
+        Members expected_systems and expected_apps are filled automatically.
+        """
+        self.name = name
+        self.description = description
+        self.system_config = system_config
+        self.apps_resources = apps_resources
+        self.fvp_dir_name = fvp_dir_name
+        self.download_artifact = download_artifact
+        self.supported_platforms = supported_platforms
+
+        self.expected_systems = get_all_system_names(name)
+        self.expected_apps = get_all_application_names(name)
+
+    @property
+    def expected_resources(self) -> Iterable[Path]:
+        """Return list of expected resources."""
+        resources = [self.system_config, *self.apps_resources]
+
+        return (get_mlia_resources() / resource for resource in resources)
+
+    @property
+    def supported_platform(self) -> bool:
+        """Return true if current platform supported."""
+        if not self.supported_platforms:
+            return True
+
+        return platform.system() in self.supported_platforms
+
+
+class BackendInstallation(Installation):
+    """Backend installation."""
+
+    def __init__(
+        self,
+        backend_runner: BackendRunner,
+        metadata: BackendMetadata,
+        path_checker: PathChecker,
+        backend_installer: BackendInstaller | None,
+    ) -> None:
+        """Init the backend installation."""
+        self.backend_runner = backend_runner
+        self.metadata = metadata
+        self.path_checker = path_checker
+        self.backend_installer = backend_installer
+
+    @property
+    def name(self) -> str:
+        """Return name of the backend."""
+        return self.metadata.name
+
+    @property
+    def description(self) -> str:
+        """Return description of the backend."""
+        return self.metadata.description
+
+    @property
+    def already_installed(self) -> bool:
+        """Return true if backend already installed."""
+        return self.backend_runner.all_installed(
+            self.metadata.expected_systems, self.metadata.expected_apps
+        )
+
+    @property
+    def could_be_installed(self) -> bool:
+        """Return true if backend could be installed."""
+        if not self.metadata.supported_platform:
+            return False
+
+        return all_paths_valid(self.metadata.expected_resources)
+
+    def supports(self, install_type: InstallationType) -> bool:
+        """Return true if backends supported type of the installation."""
+        if isinstance(install_type, DownloadAndInstall):
+            return self.metadata.download_artifact is not None
+
+        if isinstance(install_type, InstallFromPath):
+            return self.path_checker(install_type.backend_path) is not None
+
+        return False  # type: ignore
+
+    def install(self, install_type: InstallationType) -> None:
+        """Install the backend."""
+        if isinstance(install_type, DownloadAndInstall):
+            download_artifact = self.metadata.download_artifact
+            assert download_artifact is not None, "No artifact provided"
+
+            self.download_and_install(download_artifact, install_type.eula_agreement)
+        elif isinstance(install_type, InstallFromPath):
+            backend_path = self.path_checker(install_type.backend_path)
+            assert backend_path is not None, "Unable to resolve backend path"
+
+            self.install_from(backend_path)
+        else:
+            raise Exception(f"Unable to install {install_type}")
+
+    def install_from(self, backend_info: BackendInfo) -> None:
+        """Install backend from the directory."""
+        mlia_resources = get_mlia_resources()
+
+        with temp_directory() as tmpdir:
+            fvp_dist_dir = tmpdir / self.metadata.fvp_dir_name
+
+            system_config = self.metadata.system_config
+            if backend_info.system_config:
+                system_config = backend_info.system_config
+
+            resources_to_copy = [mlia_resources / system_config]
+            if backend_info.copy_source:
+                resources_to_copy.append(backend_info.backend_path)
+
+            copy_all(*resources_to_copy, dest=fvp_dist_dir)
+
+            self.backend_runner.install_system(fvp_dist_dir)
+
+        for app in self.metadata.apps_resources:
+            self.backend_runner.install_application(mlia_resources / app)
+
+    def download_and_install(
+        self, download_artifact: DownloadArtifact, eula_agrement: bool
+    ) -> None:
+        """Download and install the backend."""
+        with temp_directory() as tmpdir:
+            try:
+                downloaded_to = download_artifact.download_to(tmpdir)
+            except Exception as err:
+                raise Exception("Unable to download backend artifact") from err
+
+            with working_directory(tmpdir / "dist", create_dir=True) as dist_dir:
+                with tarfile.open(downloaded_to) as archive:
+                    archive.extractall(dist_dir)
+
+                assert self.backend_installer, (
+                    f"Backend '{self.metadata.name}' does not support "
+                    "download and installation."
+                )
+                backend_path = self.backend_installer(eula_agrement, dist_dir)
+                if self.path_checker(backend_path) is None:
+                    raise Exception("Downloaded artifact has invalid structure")
+
+                self.install(InstallFromPath(backend_path))
+
+    def uninstall(self) -> None:
+        """Uninstall the backend."""
+        remove_system(self.metadata.fvp_dir_name)
+
+
+class PackagePathChecker:
+    """Package path checker."""
+
+    def __init__(
+        self, expected_files: list[str], backend_subfolder: str | None = None
+    ) -> None:
+        """Init the path checker."""
+        self.expected_files = expected_files
+        self.backend_subfolder = backend_subfolder
+
+    def __call__(self, backend_path: Path) -> BackendInfo | None:
+        """Check if directory contains all expected files."""
+        resolved_paths = (backend_path / file for file in self.expected_files)
+        if not all_files_exist(resolved_paths):
+            return None
+
+        if self.backend_subfolder:
+            subfolder = backend_path / self.backend_subfolder
+
+            if not subfolder.is_dir():
+                return None
+
+            return BackendInfo(subfolder)
+
+        return BackendInfo(backend_path)
+
+
+class StaticPathChecker:
+    """Static path checker."""
+
+    def __init__(
+        self,
+        static_backend_path: Path,
+        expected_files: list[str],
+        copy_source: bool = False,
+        system_config: str | None = None,
+    ) -> None:
+        """Init static path checker."""
+        self.static_backend_path = static_backend_path
+        self.expected_files = expected_files
+        self.copy_source = copy_source
+        self.system_config = system_config
+
+    def __call__(self, backend_path: Path) -> BackendInfo | None:
+        """Check if directory equals static backend path with all expected files."""
+        if backend_path != self.static_backend_path:
+            return None
+
+        resolved_paths = (backend_path / file for file in self.expected_files)
+        if not all_files_exist(resolved_paths):
+            return None
+
+        return BackendInfo(
+            backend_path,
+            copy_source=self.copy_source,
+            system_config=self.system_config,
+        )
+
+
+class CompoundPathChecker:
+    """Compound path checker."""
+
+    def __init__(self, *path_checkers: PathChecker) -> None:
+        """Init compound path checker."""
+        self.path_checkers = path_checkers
+
+    def __call__(self, backend_path: Path) -> BackendInfo | None:
+        """Iterate over checkers and return first non empty backend info."""
+        first_resolved_backend_info = (
+            backend_info
+            for path_checker in self.path_checkers
+            if (backend_info := path_checker(backend_path)) is not None
+        )
+
+        return next(first_resolved_backend_info, None)
+
+
+class PyPackageBackendInstallation(Installation):
+    """Backend based on the python package."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        packages_to_install: list[str],
+        packages_to_uninstall: list[str],
+        expected_packages: list[str],
+    ) -> None:
+        """Init the backend installation."""
+        self._name = name
+        self._description = description
+        self._packages_to_install = packages_to_install
+        self._packages_to_uninstall = packages_to_uninstall
+        self._expected_packages = expected_packages
+
+        self.package_manager = get_package_manager()
+
+    @property
+    def name(self) -> str:
+        """Return name of the backend."""
+        return self._name
+
+    @property
+    def description(self) -> str:
+        """Return description of the backend."""
+        return self._description
+
+    @property
+    def could_be_installed(self) -> bool:
+        """Check if backend could be installed."""
+        return True
+
+    @property
+    def already_installed(self) -> bool:
+        """Check if backend already installed."""
+        return self.package_manager.packages_installed(self._expected_packages)
+
+    def supports(self, install_type: InstallationType) -> bool:
+        """Return true if installation supports requested installation type."""
+        return isinstance(install_type, DownloadAndInstall)
+
+    def install(self, install_type: InstallationType) -> None:
+        """Install the backend."""
+        if not self.supports(install_type):
+            raise Exception(f"Unsupported installation type {install_type}")
+
+        self.package_manager.install(self._packages_to_install)
+
+    def uninstall(self) -> None:
+        """Uninstall the backend."""
+        self.package_manager.uninstall(self._packages_to_uninstall)