mlonmcu 0.3.0.dev0__3-none-any.whl

mlonmcu/setup/setup.py

@@ -0,0 +1,285 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import os
+import shutil
+import multiprocessing
+from tqdm import tqdm
+
+from mlonmcu.logging import get_logger
+from mlonmcu.feature.type import FeatureType
+from mlonmcu.feature.features import get_matching_features
+from mlonmcu.config import filter_config, str2bool
+from .tasks import get_task_factory
+from .task import TaskGraph
+from mlonmcu.utils import ask_user
+
+logger = get_logger()
+
+
+class Setup:
+    """MLonMCU dependency management interface."""
+
+    FEATURES = []
+
+    DEFAULTS = {
+        "print_outputs": False,
+        "num_threads": None,
+    }
+
+    REQUIRED = []
+    OPTIONAL = []
+
+    def __init__(self, features=None, config=None, context=None, tasks_factory=None):
+        if not tasks_factory:
+            tasks_factory = get_task_factory()
+        self.config = config if config else {}
+        self.features = self.process_features(features)
+        self.config = filter_config(self.config, "setup", self.DEFAULTS, self.OPTIONAL, self.REQUIRED)
+        self.context = context
+        self.tasks_factory = tasks_factory
+        self.num_threads = int(
+            self.config["num_threads"] if self.config["num_threads"] else multiprocessing.cpu_count()
+        )
+
+    @property
+    def verbose(self):
+        value = self.config["print_outputs"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    def clean_cache(self, interactive=True):
+        assert self.context is not None
+        deps_dir = self.context.environment.lookup_path("deps").path
+        cache_file = deps_dir / "cache.ini"
+        if cache_file.is_file():
+            print(f"The dependency cache file ({cache_file}) will be removed.")
+            if ask_user("Are your sure?", default=not interactive, interactive=interactive):
+                print(f"Removing {cache_file} ...")
+                os.remove(cache_file)
+
+    def clean_dependencies(self, interactive=True):
+        assert self.context is not None
+        self.clean_cache(interactive=interactive)
+        deps_dir = self.context.environment.lookup_path("deps").path
+        subdirs = ["src", "build", "install"]
+        print(f"All dependencies will be removed from {deps_dir}.")
+        if ask_user("Are your sure?", default=not interactive, interactive=interactive):
+            for subdir in subdirs:
+                full_path = deps_dir / subdir
+                print(f"Removing contents of {full_path} ...")
+                shutil.rmtree(full_path, ignore_errors=True)
+                full_path.mkdir(exist_ok=True)
+
+    def process_features(self, features):
+        if features is None:
+            return []
+        features = get_matching_features(features, FeatureType.SETUP)
+        for feature in features:
+            # Not need to list features explicitly
+            # assert (
+            #     feature.name in self.FEATURES
+            # ), f"Incompatible feature: {feature.name}"
+            feature.add_setup_config(self.config)
+        return features
+
+    def _get_task_graph(self):
+        self.tasks_factory.reset_changes()
+        task_graph = TaskGraph(
+            self.tasks_factory.registry.keys(),
+            self.tasks_factory.dependencies,
+            self.tasks_factory.providers,
+        )
+        return task_graph
+
+    def get_dependency_order(self):
+        task_graph = self._get_task_graph()
+        V, E = task_graph.get_graph()
+        order = task_graph.get_order()
+        order_str = " -> ".join(order)
+        logger.debug("Determined dependency order: %s" % order_str)
+        return order
+
+    def setup_progress_bar(self, enabled):
+        if enabled:
+            pbar = tqdm(
+                total=len(self.tasks_factory.registry),
+                desc="Installing dependencies",
+                ncols=100,
+                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}s]",
+            )
+            return pbar
+        else:
+            logger.info("Installing dependencies...")
+            return None
+
+    def write_cache_file(self):
+        logger.debug("Updating dependency cache")
+        cache_file = self.context.environment.paths["deps"].path / "cache.ini"
+        self.context.cache.write_to_file(cache_file)
+
+    def visualize(self, path, ordered=False):
+        task_graph = self._get_task_graph()
+        task_graph.export_dot(path)
+        logger.debug("Written task graph to file: %s" % path)
+
+    def invoke_single_task(self, name, progress=False, write_cache=True, rebuild=False):
+        assert name in self.tasks_factory.registry, f"Invalid task name: {name}"
+        func = self.tasks_factory.registry[name]
+        func(self.context, progress=progress, rebuild=rebuild, verbose=self.verbose, threads=self.num_threads)
+        if write_cache:
+            self.write_cache_file()
+
+    def install_dependencies(
+        self,
+        progress=False,
+        write_cache=True,
+        rebuild=False,
+    ):
+        assert self.context is not None
+        order = self.get_dependency_order()
+        pbar = self.setup_progress_bar(progress)
+        for task in order:
+            func = self.tasks_factory.registry[task]
+            func(self.context, progress=progress, rebuild=rebuild, verbose=self.verbose, threads=self.num_threads)
+            if pbar:
+                pbar.update(1)
+        if pbar:
+            pbar.close()
+        if write_cache:
+            self.write_cache_file()
+        logger.info("Finished installing dependencies")
+        return True
+
+    def generate_requirements(
+        self,
+    ):
+        from .gen_requirements import join_requirements
+        import os
+
+        requirements = join_requirements()
+        output_dir = os.path.join(self.context.environment.home, "requirements_addition.txt")
+
+        def feature_enabled_and_supported(obj, feature):
+            if getattr(obj, "enabled", True):
+                if getattr(obj, "name", None) == feature and getattr(obj, "supported", False):
+                    return True
+                else:
+                    enable_and_supported = False
+                    for k, v in obj.__dict__.items():
+                        if isinstance(v, list):
+                            for i in v:
+                                enable_and_supported = enable_and_supported or feature_enabled_and_supported(i, feature)
+                    return enable_and_supported
+            else:
+                return False
+
+        with open(output_dir, "w") as f:
+            f.write(f"# AUTOGENERATED by 'mlonmcu setup -g' {os.linesep}")
+            config_pools = (
+                self.context.environment.frameworks
+                + self.context.environment.frontends
+                + self.context.environment.platforms
+                + self.context.environment.targets
+            )
+            for config in config_pools:
+                if "espidf" in config.name and config.enabled:
+                    for d in requirements["espidf"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for espidf")
+                    break
+            for config in config_pools:
+                if "zephyr" in config.name and config.enabled:
+                    for d in requirements["zephyr"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for zephyr")
+                    break
+            for config in config_pools:
+                if "tflm" in config.name and config.enabled:
+                    for d in requirements["tflm"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for tflm")
+                    break
+            for config in config_pools:
+                if "tflite" in config.name and config.enabled:
+                    for d in requirements["tflite"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for tflite")
+                    break
+            for config in config_pools:
+                if "onnx" in config.name and config.enabled:
+                    for d in requirements["onnx"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for onnx")
+                    break
+            for config in config_pools:
+                if "etiss" in config.name and config.enabled:
+                    for d in requirements["etiss"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for etiss")
+                    break
+            for config in config_pools:
+                if "gvsoc_pulp" in config.name and config.enabled:
+                    for d in requirements["gvsoc_pulp"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for gvsoc_pulp")
+                    break
+            for config in config_pools:
+                if "visualize" in config.name and config.enabled:
+                    for d in requirements["visualize"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for visualize")
+                    break
+            for config in config_pools:
+                if "microtvm" in config.name and config.enabled:
+                    for d in requirements["microtvm"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for microtvm")
+                    break
+            for config in config_pools:
+                if feature_enabled_and_supported(config, "moiopt"):
+                    for d in requirements["moiopt"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for moiopt")
+                    break
+            # for config in config_pools:
+            # if feature_enabled_and_supported(config, "relayviz"):
+            #       # equirements for relay visualization
+            #     for d in requirements['relayviz'][1]:
+            #         f.write(f"{d}{os.linesep}")
+            # break
+            for config in config_pools:
+                if "tvm" in config.name and config.enabled:
+                    for d in requirements["tvm"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for tvm")
+                    break
+            for config in config_pools:
+                if "tvm" in config.name and feature_enabled_and_supported(config, "autotuned"):
+                    for d in requirements["tvm-autotuning"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for tvm-autotuning")
+                    break
+            for config in config_pools:
+                if "tf" in config.name and feature_enabled_and_supported(config, "visualize"):
+                    # requirements for tvm tflite visualization
+                    for d in requirements["visualize_tflite"][1]:
+                        f.write(f"{d}{os.linesep}")
+                    logger.info("add dependencies for tflite visualization")
+                    break
+        logger.info("Finished generating requirements_addition.txt")
+        return True

mlonmcu/setup/task.py

@@ -0,0 +1,380 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Definitions of a task registry used to automatically install dependencies."""
+
+from functools import wraps
+import itertools
+from enum import Enum
+import time
+from typing import List, Tuple
+import networkx as nx
+from networkx.drawing.nx_agraph import write_dot
+from tqdm import tqdm
+
+from mlonmcu.logging import get_logger
+
+logger = get_logger()
+
+
+def get_combs(data) -> List[dict]:
+    """Utility which returns combinations of the input data.
+
+    Parameters
+    ----------
+    data : dict
+        Input dictionary
+
+    Returns
+    -------
+    combs : list
+        All combinations of the input data.
+
+    Examples
+    --------
+    >>> get_combs({"foo": [False, True], "bar": [5, 10]})
+    [{"foo": False, "bar": 5}, {"foo": False, "bar": 10}, {"foo": True, "bar": 5}, {"foo": True, "bar": 10}]
+    """
+    keys = list(data.keys())
+    values = list(data.values())
+    prod = list(itertools.product(*values))
+    if len(prod) == 1:
+        if len(prod[0]) == 0:
+            prod = []
+    combs = [dict(zip(keys, p)) for p in prod]
+    return combs
+
+
+class TaskType(Enum):
+    """Enumeration for the task type."""
+
+    MISC = 0
+    FRAMEWORK = 1
+    BACKEND = 2
+    TOOLCHAIN = 3
+    TARGET = 4
+    FRONTEND = 5
+    OPT = 6
+    FEATURE = 7
+    PLATFORM = 8
+
+
+class TaskGraph:
+    """Task graph object.
+
+    Attributes
+    ----------
+    names : list
+        list of task names in the graph
+    dependencies : dict
+        Dependencies between task artifacts
+    providers : dict
+        Providers for all the artifacts
+
+    Examples
+    -------
+    TODO
+    """
+
+    def __init__(self, names: List[str], dependencies: dict, providers: dict):
+        self.names = names
+        self.dependencies = dependencies
+        self.providers = providers
+
+    def get_graph(self) -> Tuple[list, list]:
+        """Get nodes and edges of the task graph.
+
+        Returns
+        -------
+        nodes : list
+            List of edges
+        edges : list
+            List of edge tuples.
+        """
+        nodes = list(self.names)
+        edges = []
+        for dest, deps in self.dependencies.items():
+            for dep in deps:
+                if dep not in self.providers.keys():
+                    raise RuntimeError(f"Unable to resolve dependency '{dep}'")
+                src = self.providers[dep]
+                edge = (src, dest)
+                edges.append(edge)
+        # Remove duplicates
+        edges = list(dict.fromkeys(edges))
+        return nodes, edges
+
+    def get_order(self) -> list:
+        """Get execution order of tasks via topological sorting."""
+        nodes, edges = self.get_graph()
+        graph = nx.DiGraph(edges)
+        graph.add_nodes_from(nodes)
+        order = list(nx.topological_sort(graph))
+        return order
+
+    def export_dot(self, path):
+        """Visualize the task dependency graph."""
+        nodes, edges = self.get_graph()
+        graph = nx.DiGraph(edges)
+        graph.add_nodes_from(nodes)
+        # order = list(nx.topological_sort(graph))
+        # TODO: annotate with order
+        # TODO: also export order as extra graph
+        write_dot(graph, path)
+
+
+class TaskFactory:
+    """Class which is used to register all available tasks and their annotations.
+
+    Attributes
+    ----------
+    registry : dict
+        Mapping of task names and their actual function
+    dependencies : dict
+        Mapping of task dependencies
+    providers : dict
+        Mapping of which task provides which artifacts
+    types : dict
+        Mapping of task types
+    validates : dict
+        Mapping of validation functions for the tasks
+    changed : list
+        List of tasks?artifacts which have changed recently
+    """
+
+    def __init__(self):
+        self.registry = {}
+        self.dependencies = {}
+        self.providers = {}
+        self.types = {}
+        self.params = {}
+        self.validates = {}
+        self.changed = []  # Main problem: per
+
+    def reset_changes(self):
+        """Reset all pending changes."""
+        self.changed = []
+
+    def needs(self, keys, force=True):
+        """Decorator which registers the artifacts a task needs to be processed."""
+
+        def real_decorator(function):
+            name = function.__name__
+            if name in self.dependencies:
+                self.dependencies[name].extend(keys)
+            else:
+                self.dependencies[name] = keys
+
+            @wraps(function)
+            def wrapper(*args, **kwargs):
+                # logger.debug("Checking inputs...")
+                if force:
+                    context = args[0]
+                    variables = context.cache._vars
+                    for key in keys:
+                        if key not in variables.keys() or variables[key] is None:
+                            raise RuntimeError(f"Task '{name}' needs the value of '{key}' which is not set")
+                retval = function(*args, **kwargs)
+                return retval
+
+            return wrapper
+
+        return real_decorator
+
+    def optional(self, keys):
+        """Decorator for optional task requirements."""
+        return self.needs(keys, force=False)
+
+    # def optional(self, keys):
+    #     def real_decorator(function):
+    #         name = function.__name__
+    #         if name in self.dependencies:
+    #             self.dependencies[name].extend(keys)
+    #         else:
+    #             self.dependencies[name] = keys
+    #         @wraps(function)
+    #         def wrapper(*args, **kwargs):
+    #             retval = function(*args, **kwargs)
+    #             return retval
+    #         return wrapper
+    #     return real_decorator
+
+    def provides(self, keys):
+        """Decorator which registers what a task provides."""
+
+        def real_decorator(function):
+            name = function.__name__
+            for key in keys:
+                self.providers[key] = name
+
+            @wraps(function)
+            def wrapper(*args, **kwargs):
+                context = args[0]
+                for key in keys:
+                    if key in context.cache._vars:
+                        del context.cache._vars[key]  # Unset the value before calling function
+                retval = function(*args, **kwargs)
+                if retval is not False:
+                    # logger.debug("Checking outputs...")
+                    variables = context.cache._vars
+                    for key in keys:
+                        if key not in variables.keys() or variables[key] is None:
+                            raise RuntimeError(f"Task '{name}' did not set the value of '{key}'")
+                return retval
+
+            self.registry[name] = wrapper
+            return wrapper
+
+        return real_decorator
+
+    def param(self, flag, options):
+        """Decorator which registers available task parameters."""
+        if not isinstance(options, list):
+            options = [options]
+
+        def real_decorator(function):
+            name = function.__name__
+            if name in self.params:
+                self.params[name][flag] = options
+            else:
+                self.params[name] = {flag: options}
+
+            @wraps(function)
+            def wrapper(*args, **kwargs):
+                retval = function(*args, **kwargs)
+                return retval
+
+            return wrapper
+
+        return real_decorator
+
+    def validate(self, func):
+        """Decorator which registers validattion functions for a task."""
+
+        def real_decorator(function):
+            name = function.__name__
+            self.validates[name] = func
+            return function
+
+        return real_decorator
+
+    def register(self, category=TaskType.MISC):
+        """Decorator which actually registers a task in the registry."""
+
+        def real_decorator(function):
+            name = function.__name__
+
+            @wraps(function)
+            def wrapper(*args, rebuild=False, progress=False, **kwargs):
+                combs = get_combs(self.params[name])
+
+                def get_valid_combs(combs):
+                    ret = []
+                    for comb in combs:
+                        if name in self.validates:
+                            check = self.validates[name](args[0], params=comb)
+                            if not check:
+                                continue
+                        ret.append(comb)
+                    return ret
+
+                combs = get_valid_combs(combs)
+
+                def process(name_, params=None, rebuild=False):
+                    if not params:
+                        params = []
+                    rebuild = rebuild
+                    if name in self.dependencies:
+                        for dep in self.dependencies[name]:
+                            if dep in self.changed:
+                                rebuild = True
+                                break
+                    retval = function(*args, params=params, rebuild=rebuild, **kwargs)
+                    if retval:
+                        keys = [key for key, provider in self.providers.items() if provider == name]
+                        for key in keys:
+                            if key not in self.changed:
+                                self.changed.append(key)
+                    # logger.debug("Processed task:", function.__name__)
+                    return retval
+
+                if progress:
+                    pbar = tqdm(
+                        total=max(len(combs), 1),
+                        desc="Processing",
+                        ncols=100,
+                        bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}s]",
+                        leave=None,
+                    )
+                else:
+                    pbar = None
+                if len(combs) == 0:
+                    if pbar:
+                        pbar.set_description(f"Processing: {name}")
+                    else:
+                        logger.info("Processing task: %s", name)
+                    time.sleep(0.1)
+                    check = True
+                    if name in self.validates:
+                        check = self.validates[name](args[0], params={})
+                    if check:
+                        start = time.time()
+                        retval = process(name, rebuild=rebuild)
+                        end = time.time()
+                        diff = end - start
+                        minutes = int(diff // 60)
+                        seconds = int(diff % 60)
+                        duration_str = f"{seconds}s" if minutes == 0 else f"{minutes}m{seconds}s"
+                        if not pbar:
+                            logger.debug("-> Done (%s)", duration_str)
+                        # TODO: move this to helper func
+                    else:
+                        logger.debug("-> Skipped")
+                        retval = False
+                    if pbar:
+                        pbar.update(1)
+                else:
+                    for comb in combs:  # TODO process in parallel?
+                        extended_name = name + str(comb)
+                        if pbar:
+                            pbar.set_description(f"Processing - {extended_name}")
+                        else:
+                            logger.info("Processing task: %s", extended_name)
+                        time.sleep(0.1)
+                        start = time.time()
+                        retval = process(extended_name, params=comb, rebuild=rebuild)
+                        end = time.time()
+                        diff = end - start
+                        minutes = int(diff // 60)
+                        seconds = int(diff % 60)
+                        duration_str = f"{seconds}s" if minutes == 0 else f"{minutes}m{seconds}s"
+                        if not pbar:
+                            logger.debug("-> Done (%s)", duration_str)
+                            # TODO: move this to helper func
+                        else:
+                            pbar.update(1)
+                if pbar:
+                    pbar.close()
+                return retval
+
+            self.registry[name] = wrapper
+            self.types[name] = category
+            self.params[name] = {}
+            return wrapper
+
+        return real_decorator