mlonmcu 0.3.0.dev0__3-none-any.whl

mlonmcu/platform/espidf/espidf.py

@@ -0,0 +1,489 @@
+#
+# 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.
+#
+"""ESP-IDF Platform"""
+
+import os
+import time
+import signal
+import shutil
+import tempfile
+import subprocess
+from pathlib import Path
+import pkg_resources
+
+
+from mlonmcu.setup import utils
+from mlonmcu.artifact import Artifact, ArtifactFormat
+from mlonmcu.logging import get_logger
+from mlonmcu.target.target import Target
+from mlonmcu.config import str2bool
+
+from ..platform import CompilePlatform, TargetPlatform
+from .espidf_target import create_espidf_platform_target, get_espidf_platform_targets
+
+logger = get_logger()
+
+
+def get_project_template(name="project"):
+    espidf_templates = pkg_resources.resource_listdir("mlonmcu", os.path.join("..", "resources", "platforms", "espidf"))
+    if name not in espidf_templates:
+        return None
+    fname = pkg_resources.resource_filename("mlonmcu", os.path.join("..", "resources", "platforms", "espidf", name))
+    return fname
+
+
+class EspIdfPlatform(CompilePlatform, TargetPlatform):
+    """ESP-IDF Platform class."""
+
+    FEATURES = CompilePlatform.FEATURES + TargetPlatform.FEATURES + ["benchmark"]
+
+    DEFAULTS = {
+        **CompilePlatform.DEFAULTS,
+        **TargetPlatform.DEFAULTS,
+        "project_template": None,
+        "project_dir": None,
+        "port": None,
+        "baud": 115200,
+        "use_idf_monitor": True,
+        "wait_for_user": True,
+        "flash_only": False,
+    }
+
+    REQUIRED = ["espidf.install_dir", "espidf.src_dir"]
+
+    def __init__(self, features=None, config=None):
+        super().__init__(
+            "espidf",
+            features=features,
+            config=config,
+        )
+        self.tempdir = None
+        self.project_name = "app"
+        self.project_dir = None
+
+    @property
+    def espidf_install_dir(self):
+        return Path(self.config["espidf.install_dir"])
+
+    @property
+    def espidf_src_dir(self):
+        return Path(self.config["espidf.src_dir"])
+
+    @property
+    def idf_exe(self):
+        return self.espidf_src_dir / "tools" / "idf.py"
+
+    @property
+    def use_idf_monitor(self):
+        value = self.config["use_idf_monitor"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    @property
+    def wait_for_user(self):
+        value = self.config["wait_for_user"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    @property
+    def flash_only(self):
+        # TODO: get rid of this
+        value = self.config["flash_only"]
+        return str2bool(value) if not isinstance(value, (bool, int)) else value
+
+    def invoke_idf_exe(self, *args, **kwargs):
+        env = os.environ.copy()
+        env["IDF_PATH"] = str(self.espidf_src_dir)
+        env["IDF_TOOLS_PATH"] = str(self.espidf_install_dir)
+        cmd = (
+            ". "
+            + str(self.espidf_src_dir / "export.sh")
+            + f" && {self.idf_exe} "
+            # + f" > /dev/null && {self.idf_exe} "
+            + " ".join([str(arg) for arg in args])
+        )
+        out = utils.exec_getout(
+            cmd, shell=True, env=env, **kwargs, executable="/bin/bash"
+        )  # TODO: using shell=True is insecure but right now we can not avoid it?
+        return out
+
+    def init_directory(self, path=None, context=None):
+        if self.project_dir is not None:
+            self.project_dir.mkdir(exist_ok=True)
+            logger.debug("Project directory already initialized")
+            return
+        dir_name = self.name
+        if path is not None:
+            self.project_dir = Path(path)
+        elif self.config["project_dir"] is not None:
+            self.project_dir = Path(self.config["project_dir"])
+        else:
+            if context:
+                assert "temp" in context.environment.paths
+                self.project_dir = (
+                    context.environment.paths["temp"].path / dir_name
+                )  # TODO: Need to lock this for parallel builds
+            else:
+                logger.debug(
+                    "Creating temporary directory because no context was available "
+                    "and 'espidf.project_dir' was not supplied"
+                )
+                self.tempdir = tempfile.TemporaryDirectory()
+                self.project_dir = Path(self.tempdir.name) / dir_name
+                logger.debug("Temporary project directory: %s", self.project_dir)
+        self.project_dir.mkdir(exist_ok=True)
+
+    def get_supported_targets(self):
+        text = self.invoke_idf_exe("--list-targets", live=self.print_outputs, print_output=False)
+        # Warning: This will fail if a python executable is NOT available in the system. Aliasing
+        # python3 to python will not work. Not sure how this would handle a system which only has python2 installed?
+        target_names = text.split("\n")
+
+        def filter_names(x):
+            ret = []
+            for x in reversed(x):
+                if len(x) > 0:
+                    ret.append(x)
+                else:
+                    if len(ret) > 0:
+                        break
+                    else:
+                        continue
+            return reversed(ret)
+
+        target_names = filter_names(target_names)
+
+        return [name for name in target_names if len(name) > 0 and " " not in name]
+
+    def create_target(self, name):
+        assert name in self.get_supported_targets(), f"{name} is not a valid ESP-IDF target"
+        targets = get_espidf_platform_targets()
+        if name in targets:
+            base = targets[name]
+        else:
+            base = Target
+        return create_espidf_platform_target(name, self, base=base)
+
+    @property
+    def project_template(self):
+        return self.config["project_template"]
+
+    @property
+    def port(self):
+        return self.config["port"]
+
+    @property
+    def baud(self):
+        return self.config["baud"]
+
+    def close(self):
+        if self.tempdir:
+            self.tempdir.cleanup()
+
+    def check(self):
+        if not shutil.which(self.idf_exe):
+            raise RuntimeError(f"It seems like '{self.idf_exe}' is not available. Make sure to setup your environment!")
+
+    def prepare(self, target, src):
+        self.init_directory()
+        self.check()
+        template_dir = self.project_template
+        if template_dir is None:
+            template_dir = get_project_template()
+        else:
+            template_dir = Path(template_dir)
+            if not template_dir.is_dir():
+                template_dir = get_project_template(name=str(template_dir))
+        assert template_dir is not None, f"Provided project template does not exists: {template_dir}"
+        shutil.copytree(template_dir, self.project_dir, dirs_exist_ok=True)
+
+        def write_defaults(filename):
+            defs = self.definitions
+            with open(filename, "w", encoding="utf-8") as f:
+                f.write("CONFIG_PARTITION_TABLE_SINGLE_APP_LARGE=y\n")
+                if self.debug:
+                    f.write("CONFIG_OPTIMIZATION_LEVEL_DEBUG=y\n")
+                    f.write("CONFIG_COMPILER_OPTIMIZATION_LEVEL_DEBUG=y\n")
+                else:
+                    # Trying to reduce the binary size as much as possible
+                    # (https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/performance/size.html)
+                    f.write("CONFIG_BOOTLOADER_LOG_LEVEL_NONE=y\n")
+                    f.write("CONFIG_BOOTLOADER_LOG_LEVEL=0\n")
+                    f.write("CONFIG_BOOT_ROM_LOG_ALWAYS_OFF=y\n")
+                    f.write("CONFIG_COMPILER_OPTIMIZATION_ASSERTIONS_DISABLE=y\n")
+                    f.write("CONFIG_COMPILER_OPTIMIZATION_ASSERTION_LEVEL=0\n")
+                    f.write("CONFIG_COMPILER_OPTIMIZATION_CHECKS_SILENT=y\n")
+                    f.write("CONFIG_HAL_DEFAULT_ASSERTION_LEVEL=0\n")
+                    f.write("CONFIG_LOG_DEFAULT_LEVEL_NONE=y\n")
+                    f.write("CONFIG_LOG_DEFAULT_LEVEL=0\n")
+                    f.write("CONFIG_LOG_MAXIMUM_LEVEL=0\n")
+                    f.write("CONFIG_NEWLIB_NANO_FORMAT=y\n")
+                    f.write("CONFIG_COMPILER_OPTIMIZATION_LEVEL_RELEASE=y\n")
+                    f.write("CONFIG_OPTIMIZATION_LEVEL_RELEASE=y\n")
+                    optimize_for_size = True
+                    if optimize_for_size:
+                        f.write("CONFIG_COMPILER_OPTIMIZATION_SIZE=y\n")
+                    else:
+                        f.write("CONFIG_COMPILER_OPTIMIZATION_PERF=y\n")
+                watchdog_sec = 60
+                f.write(f"CONFIG_ESP_TASK_WDT_TIMEOUT_S={watchdog_sec}\n")
+                f.write(f'CONFIG_MLONMCU_CODEGEN_DIR="{src}"\n')
+                for key, value in defs.items():
+                    if isinstance(value, bool):
+                        value = "y" if value else "n"
+                    else:
+                        value = f'"{value}"'
+                    f.write(f"CONFIG_{key}={value}\n")
+
+        write_defaults(self.project_dir / "sdkconfig.defaults")
+        idfArgs = [
+            "-C",
+            self.project_dir,
+            "set-target",
+            target.name,
+        ]
+        out = self.invoke_idf_exe(*idfArgs, live=self.print_outputs)
+        return out
+
+    def get_idf_cmake_args(self):
+        cmake_defs = {"CMAKE_BUILD_TYPE": "Debug" if self.debug else "Release"}
+        return [f"-D{key}={value}" for key, value in cmake_defs.items()]
+
+    def compile(self, target, src=None):
+        out = ""
+        # TODO: build with cmake options
+        out += self.prepare(target, src)
+        # TODO: support self.num_threads (e.g. patch esp-idf)
+        idfArgs = [
+            "-C",
+            self.project_dir,
+            *self.get_idf_cmake_args(),
+            "build",
+        ]
+        out += self.invoke_idf_exe(*idfArgs, live=self.print_outputs)
+        return out
+
+    def generate_elf(self, src, target, model=None, data_file=None):
+        artifacts = []
+        out = self.compile(target, src=src)
+        elf_name = self.project_name + ".elf"
+        elf_file = self.project_dir / "build" / elf_name
+        # TODO: just use path instead of raw data?
+        if self.tempdir:
+            # Warning: The porject context will get destroyed afterwards wehen using  a temporory directory
+            with open(elf_file, "rb") as handle:
+                data = handle.read()
+                artifact = Artifact("generic_mlif", raw=data, fmt=ArtifactFormat.RAW)
+                artifacts.append(artifact)
+        else:
+            artifact = Artifact(elf_name, path=elf_file, fmt=ArtifactFormat.PATH)
+            artifacts.append(artifact)
+        metrics = self.get_metrics(elf_file)
+        content = metrics.to_csv(include_optional=True)  # TODO: store df instead?
+        metrics_artifact = Artifact("metrics.csv", content=content, fmt=ArtifactFormat.TEXT)
+        artifacts.append(metrics_artifact)
+        stdout_artifact = Artifact(
+            "espidf_out.log", content=out, fmt=ArtifactFormat.TEXT  # TODO: split into one file per command
+        )  # TODO: rename to tvmaot_out.log?
+        artifacts.append(stdout_artifact)
+        self.artifacts = artifacts
+
+    def get_idf_serial_args(self, monitor=False):
+        args = []
+        if self.port:
+            args.extend(["-p", self.port])
+        if self.baud:
+            args.extend(["-B" if monitor else "-b", self.baud])
+        return args
+
+    def flash(self, elf, target, timeout=120):
+        # Ignore elf, as we use self.project_dir instead
+        # TODO: add alternative approach which allows passing elf instead
+        if elf is not None:
+            logger.debug("Ignoring ELF file for espidf platform")
+        # TODO: implement timeout
+        # TODO: make sure that already compiled? -> error or just call compile routine?
+        if self.wait_for_user:  # INTERACTIVE
+            answer = input(
+                f"Make sure that the device '{target.name}' is connected before you press [Enter]"
+                + " (Type 'Abort' to cancel)"
+            )
+            if answer.lower() == "abort":
+                return ""
+        logger.debug("Flashing target software")
+
+        idfArgs = [
+            "-C",
+            self.project_dir,
+            *self.get_idf_cmake_args(),
+            "flash",
+            *self.get_idf_serial_args(),
+        ]
+        self.invoke_idf_exe(*idfArgs, live=self.print_outputs)
+
+    def monitor(self, target, timeout=60):
+        if self.flash_only:
+            return ""
+
+        if self.use_idf_monitor:
+            import psutil
+
+            def _kill_monitor():
+                import psutil
+                for proc in psutil.process_iter():
+                    # check whether the process name matches
+                    cmdline = " ".join(proc.cmdline())
+                    if "idf_monitor.py" in cmdline:  # TODO: do something less "dangerous"?
+                        proc.kill()
+
+            def _monitor_helper(*args, verbose=False, start_match=None, end_match=None, timeout=60):
+                # start_match and end_match are inclusive
+                if timeout:
+                    pass  # TODO: implement timeout
+                found_start = start_match is None
+                logger.debug("- Executing: %s", str(args))
+                outStr = ""
+                env = os.environ.copy()
+                env["IDF_PATH"] = str(self.espidf_src_dir)
+                env["IDF_TOOLS_PATH"] = str(self.espidf_install_dir)
+                cmd = (
+                    ". "
+                    + str(self.espidf_src_dir / "export.sh")
+                    + f"> /dev/null && {self.idf_exe} "
+                    + " ".join([str(arg) for arg in args])
+                )
+                process = subprocess.Popen(
+                    cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, shell=True, executable="/bin/bash", env=env
+                )
+                try:
+                    exit_code = None
+                    for line in process.stdout:
+                        new_line = line.decode(errors="replace")
+                        if verbose:
+                            print(new_line.replace("\n", ""))
+                        if start_match and start_match in new_line:
+                            outStr = new_line
+                            found_start = True
+                        else:
+                            outStr = outStr + new_line
+                        if found_start:
+                            if end_match and end_match in new_line:
+                                _kill_monitor()
+                                process.terminate()
+                                exit_code = 0
+                    while exit_code is None:
+                        exit_code = process.poll()
+                    if not verbose and exit_code != 0:
+                        logger.error(outStr)
+                    assert exit_code == 0, "The process returned an non-zero exit code {}! (CMD: `{}`)".format(
+                        exit_code, cmd
+                    )
+                except KeyboardInterrupt:
+                    logger.debug("Interrupted subprocess. Sending SIGINT signal...")
+                    _kill_monitor()
+                    pid = process.pid
+                    os.kill(pid, signal.SIGINT)
+                os.system("reset")
+                return outStr
+
+            logger.debug("Monitoring target software")
+            # TODO: implement timeout
+            idfArgs = [
+                "-C",
+                self.project_dir,
+                *self.get_idf_cmake_args(),
+                "monitor",
+                *self.get_idf_serial_args(monitor=True),
+            ]
+            return _monitor_helper(
+                *idfArgs,
+                verbose=self.print_outputs,
+                start_match="MLonMCU: START",
+                end_match="MLonMCU: STOP",
+                timeout=timeout,
+            )
+        else:
+            port = target.port
+            if port is None:
+                # Falling back to 'espidf.baud'
+                assert self.port is not None, f"If using custom serial monitor, please provide '{target.name}.port'"
+                port = self.port
+            baud = target.baud
+            if baud is None:
+                # Falling back to 'espidf.baud'
+                assert self.baud is not None, f"If using custom serial monitor, please provide '{target.name}.baud'"
+                baud = self.baud
+
+            def _monitor_helper2(port, baud, verbose=False, start_match=None, end_match=None, timeout=60):
+                # Local import to make this only required for real HW targets
+                import serial
+
+                # start_match and end_match are inclusive
+                found_start = start_match is None
+                outStr = ""
+                if timeout:
+                    pass  # TODO: implement timeout
+                # The following is a custom initialization sequence inspired by
+                # (https://github.com/espressif/esp-idf/blob/master/tools/idf_monitor_base/serial_reader.py)
+                high = False
+                low = True
+                ser = serial.Serial(port, baud)
+                ser.close()
+                ser.dtr = False
+                ser.rts = False
+                time.sleep(1)
+                ser.dtr = low
+                ser.rts = high
+                ser.dtr = ser.dtr
+                ser.open()
+                ser.dtr = high
+                ser.rts = low
+                ser.dtr = ser.dtr
+                time.sleep(0.002)
+                ser.rts = high
+                ser.dtr = ser.dtr
+                try:
+                    while True:
+                        try:
+                            ser_bytes = ser.readline()
+                            new_line = ser_bytes.decode("utf-8", errors="replace")
+                            if verbose:
+                                print(new_line.replace("\n", ""))
+                            if start_match and start_match in new_line:
+                                outStr = new_line
+                                found_start = True
+                            else:
+                                outStr = outStr + new_line
+                            if found_start:
+                                if end_match and end_match in new_line:
+                                    break
+                        except KeyboardInterrupt:
+                            logger.warning("Stopped processing serial port (KeyboardInterrupt)")
+                finally:
+                    ser.close()
+                return outStr
+
+            logger.debug("Monitoring target software")
+            # TODO: implement timeout
+            return _monitor_helper2(
+                port,
+                baud,
+                verbose=self.print_outputs,
+                start_match="MLonMCU: START",
+                end_match="MLonMCU: STOP",
+                timeout=timeout,
+            )