onnxruntime/tools/python/util/android/android.py

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import collections
import contextlib
import logging
import os
import shutil
import signal
import subprocess
import time
import typing

from ..run import run
from ..platform_helpers import is_windows


_log = logging.getLogger("util.android")


SdkToolPaths = collections.namedtuple(
    "SdkToolPaths", ["emulator", "adb", "sdkmanager", "avdmanager"])


def get_sdk_tool_paths(sdk_root: str):
    def filename(name, windows_extension):
        if is_windows():
            return "{}.{}".format(name, windows_extension)
        else:
            return name

    def resolve_path(dirnames, basename):
        dirnames.insert(0, "")
        for dirname in dirnames:
            path = shutil.which(os.path.join(dirname, basename))
            if path is not None:
                path = os.path.realpath(path)
                _log.debug("Found {} at {}".format(basename, path))
                return path
        _log.warning("Failed to resolve path for {}".format(basename))
        return None

    return SdkToolPaths(
        emulator=resolve_path(
            [os.path.join(sdk_root, "emulator")],
            filename("emulator", "exe")),
        adb=resolve_path(
            [os.path.join(sdk_root, "platform-tools")],
            filename("adb", "exe")),
        sdkmanager=resolve_path(
            [os.path.join(sdk_root, "tools", "bin"),
             os.path.join(sdk_root, "cmdline-tools", "tools", "bin")],
            filename("sdkmanager", "bat")),
        avdmanager=resolve_path(
            [os.path.join(sdk_root, "tools", "bin"),
             os.path.join(sdk_root, "cmdline-tools", "tools", "bin")],
            filename("avdmanager", "bat")))


def create_virtual_device(
        sdk_tool_paths: SdkToolPaths,
        system_image_package_name: str,
        avd_name: str):
    run(sdk_tool_paths.sdkmanager, "--install", system_image_package_name,
        input=b"y")

    run(sdk_tool_paths.avdmanager, "create", "avd",
        "--name", avd_name,
        "--package", system_image_package_name,
        "--force",
        input=b"no")


_process_creationflags = subprocess.CREATE_NEW_PROCESS_GROUP if is_windows() else 0


def _start_process(*args) -> subprocess.Popen:
    _log.debug("Starting process - args: {}".format([*args]))
    return subprocess.Popen([*args], creationflags=_process_creationflags)


_stop_signal = signal.CTRL_BREAK_EVENT if is_windows() else signal.SIGTERM


def _stop_process(proc: subprocess.Popen):
    _log.debug("Stopping process - args: {}".format(proc.args))
    proc.send_signal(_stop_signal)

    try:
        proc.wait(30)
    except subprocess.TimeoutExpired:
        _log.warning("Timeout expired, forcibly stopping process...")
        proc.kill()


def _stop_process_with_pid(pid: int):
    # not attempting anything fancier than just sending _stop_signal for now
    _log.debug("Stopping process - pid: {}".format(pid))
    os.kill(pid, _stop_signal)


def start_emulator(
        sdk_tool_paths: SdkToolPaths,
        avd_name: str,
        extra_args: typing.Optional[typing.Sequence[str]] = None) -> subprocess.Popen:
    with contextlib.ExitStack() as emulator_stack, \
         contextlib.ExitStack() as waiter_stack:
        emulator_args = [
            sdk_tool_paths.emulator, "-avd", avd_name,
            "-memory", "4096",
            "-timezone", "America/Los_Angeles",
            "-no-snapshot",
            "-no-audio",
            "-no-boot-anim",
            "-no-window"]
        if extra_args is not None:
            emulator_args += extra_args

        emulator_process = emulator_stack.enter_context(
            _start_process(*emulator_args))
        emulator_stack.callback(_stop_process, emulator_process)

        waiter_process = waiter_stack.enter_context(
            _start_process(
                sdk_tool_paths.adb, "wait-for-device", "shell",
                "while [[ -z $(getprop sys.boot_completed) ]]; do sleep 1; done; input keyevent 82"))
        waiter_stack.callback(_stop_process, waiter_process)

        # poll subprocesses
        sleep_interval_seconds = 1
        while True:
            waiter_ret, emulator_ret = waiter_process.poll(), emulator_process.poll()

            if emulator_ret is not None:
                # emulator exited early
                raise RuntimeError("Emulator exited early with return code: {}".format(emulator_ret))

            if waiter_ret is not None:
                if waiter_ret == 0:
                    break
                raise RuntimeError("Waiter process exited with return code: {}".format(waiter_ret))

            time.sleep(sleep_interval_seconds)

        # emulator is ready now
        emulator_stack.pop_all()
        return emulator_process


def stop_emulator(emulator_proc_or_pid: typing.Union[subprocess.Popen, int]):
    if isinstance(emulator_proc_or_pid, subprocess.Popen):
        _stop_process(emulator_proc_or_pid)
    elif isinstance(emulator_proc_or_pid, int):
        _stop_process_with_pid(emulator_proc_or_pid)
    else:
        raise ValueError("Expected either a PID or subprocess.Popen instance.")