components/_impl/workers/subprocess_worker.py

import contextlib
import datetime
import io
import os
import marshal
import pathlib
import shutil
import signal
import subprocess
import sys
import tempfile
import textwrap
import time
import typing

import components
from components._impl.workers import base
from components._impl.workers import subprocess_rpc


class SubprocessWorker(base.WorkerBase):
    """Open a subprocess using `python -i`, and use it to execute code.

        This class wraps a subprocess which runs a clean instance of Python.
    This enables hermetic execution of stateful code, GIL free concurrent
    benchmarking, and easy use of command line tools from Python.

        When using SubprocessWorker, it is important to remember that while the
    environment is (or at least tries to be) identical to the parent, it does
    not share state or initialization with the parent process. Imports must be
    re-run in the worker, and shared resources (such as file descriptors) will
    not be available. For most applications this mirrors the semantics of
    `timeit.Timer`.

        The principle extension point for SubprocessWorker is the `args`
    property. By overriding it, subclasses can change the nature of the
    underlying subprocess while reusing all of the generic communication and
    fault handling facilities of the base class. For example, suppose we want
    to use TaskSet to pin the worker to a single core. The code is simply:

    ```
        class TasksetZeroWorker(SubprocessWorker):

            @property
            def args(self) -> typing.List[str]:
                return ["taskset", "--cpu-list", "0"] + super().args
    ```
    """

    _working_dir: str
    _alive: bool = False
    _bootstrap_timeout: int = 10  # seconds

    def __init__(self, timeout: typing.Optional[float] = None) -> None:
        super().__init__()

        # Log inputs and outputs for debugging.
        self._command_log = os.path.join(self.working_dir, "commands.log")
        pathlib.Path(self._command_log).touch()

        self._stdout_f: io.FileIO = io.FileIO(
            os.path.join(self.working_dir, "stdout.txt"), mode="w",
        )
        self._stderr_f: io.FileIO = io.FileIO(
            os.path.join(self.working_dir, "stderr.txt"), mode="w",
        )

        # `self._run` has strong assumptions about how `_input_pipe` and
        # `_output_pipe` are used. They should not be accessed in any other
        # context. (The same is true for `self.load` and `_load_pipe`.)
        self._input_pipe = subprocess_rpc.Pipe()
        self._output_pipe = subprocess_rpc.Pipe(
            timeout=timeout,
            timeout_callback=self._kill_proc,
        )
        self._load_pipe = subprocess_rpc.Pipe(
            timeout=timeout,
            timeout_callback=self._kill_proc,
        )

        # Windows and Unix differ in how pipes are shared with children.
        # In Unix they are inherited, while in Windows the child consults the
        # OS to get access. Most of this complexity is handled by
        # `subprocess_rpc.Pipe`, however we also have to make sure Popen
        # exposes the pipes in a platform appropriate way.
        child_fds = [
            self._input_pipe.read_fd,
            self._output_pipe.write_fd,
            self._load_pipe.write_fd,
        ]
        if subprocess_rpc.IS_WINDOWS:
            for fd in child_fds:
                os.set_inheritable(fd, True)

            startupinfo = subprocess.STARTUPINFO()
            startupinfo.lpAttributeList["handle_list"].extend(
                [subprocess_rpc.to_handle(fd) for fd in child_fds])

            popen_kwargs = {
                "startupinfo": startupinfo,
            }

        else:
            popen_kwargs = {
                "close_fds": True,
                "pass_fds": child_fds,
            }

        self._proc = subprocess.Popen(
            args=self.args,
            stdin=subprocess.PIPE,
            stdout=self._stdout_f,
            stderr=self._stderr_f,
            encoding=subprocess_rpc.ENCODING,
            bufsize=1,
            cwd=os.getcwd(),
            **popen_kwargs,
        )

        # setup the pid of child process in the output pipe
        self._output_pipe.set_writer_pid(self._proc.pid)

        self._worker_bootstrap_finished: bool = False
        self._bootstrap_worker()
        self._alive = True

    @property
    def working_dir(self) -> str:
        # A subclass might need to access `self.working_dir` before calling
        # `super().__init__` in order to properly construct `args`, so we need
        # to lazily initialize it.
        if getattr(self, "_working_dir", None) is None:
            self._working_dir = tempfile.mkdtemp()
        return self._working_dir

    @property
    def args(self) -> typing.List[str]:
        return [sys.executable, "-i", "-u"]

    def run(self, snippet: str) -> None:
        self._run(snippet)

    def store(self, name: str, value: typing.Any, in_memory: bool = False) -> None:
        if in_memory:
            raise NotImplementedError("SubprocessWorker does not support `in_memory`")

        # NB: we convert the bytes to a hex string to avoid encoding issues.
        self._run(f"""
            {name} = {subprocess_rpc.WORKER_IMPL_NAMESPACE}["marshal"].loads(
                bytes.fromhex({repr(marshal.dumps(value).hex())})
            )
        """)

    def load(self, name: str) -> typing.Any:
        self._run(f"""
            {subprocess_rpc.WORKER_IMPL_NAMESPACE}["load_pipe"].write(
                {subprocess_rpc.WORKER_IMPL_NAMESPACE}["marshal"].dumps({name})
            )
        """)

        return marshal.loads(self._load_pipe.read())

    @property
    def in_process(self) -> bool:
        return False

    @property
    def alive(self) -> bool:
        return self._alive and self._proc.poll() is None

    def _bootstrap_worker(self) -> None:
        """Import subprocess_rpc in the worker, and start the work loop.

        Commands are executed by writing to `self._input_pipe`, and waiting for
        a response on `self._output_pipe`. This presumes, however, that there
        is a worker doing the opposite: listening to the input pipe and writing
        to the output pipe. At startup `self._proc` is a simple interactive
        Python process, so we have to bootstrap it to start the work loop or
        else `self._run` will hang waiting for jobs to be processed.
        """

        # NB: This gets sent directly to `self._proc`'s stdin, so it MUST be
        #     a single expression and may NOT contain any empty lines. (Due to
        #     how Python processes commands.)
        bootstrap_command = textwrap.dedent(f"""
            try:
                import marshal
                import sys
                sys_path_old = list(sys.path)
                sys.path = marshal.loads(
                    bytes.fromhex({repr(marshal.dumps(sys.path).hex())})
                )
                # The parent gets priority, but a subclass could set PYTHONPATH
                # so we have to respect extra paths.
                sys.path.extend([i for i in sys_path_old if i and i not in sys.path])
                from components._impl.workers import subprocess_rpc
                output_pipe = subprocess_rpc.Pipe(
                    write_handle={self._output_pipe.write_handle})
                output_pipe.write(subprocess_rpc.BOOTSTRAP_IMPORT_SUCCESS)
                subprocess_rpc.run_loop(
                    input_handle={self._input_pipe.read_handle},
                    output_pipe=output_pipe,
                    load_handle={self._load_pipe.write_handle},
                )
            except:
                sys.exit(1)
        """).strip()

        if self._proc.poll() is not None:
            raise ValueError("Process has already exited.")

        proc_stdin = self._proc.stdin
        assert proc_stdin is not None

        self._log_cmd(bootstrap_command)

        # We need two newlines for Python to stop waiting for more input.
        proc_stdin.write(f"{bootstrap_command}\n\n")
        proc_stdin.flush()

        with self.watch_stdout_stderr() as get_output:
            try:
                # Bootstrapping is very fast. (Unlike user code where we have
                # no a priori expected upper bound.) If we don't get a response
                # prior to the timeout, it is overwhelmingly likely that the
                # worker died or the bootstrap failed. (E.g. failed to resolve
                # import path.) This simply allows us to raise a good error.
                bootstrap_pipe = subprocess_rpc.Pipe(
                    read_handle=self._output_pipe.read_handle,
                    write_handle=self._output_pipe.write_handle,
                    timeout=self._bootstrap_timeout,
                )
                result = bootstrap_pipe.read()
                assert result == subprocess_rpc.BOOTSTRAP_IMPORT_SUCCESS, result

                result = bootstrap_pipe.read()
                assert result == subprocess_rpc.BOOTSTRAP_INPUT_LOOP_SUCCESS, result

                self._worker_bootstrap_finished = True
                assert self._proc.poll() is None
            except (Exception, KeyboardInterrupt) as e:
                stdout, stderr = get_output()
                cause = "import failed" if self._proc.poll() else "timeout"
                raise e from RuntimeError(
                    f"Failed to bootstrap worker ({cause}):\n"
                    f"    working_dir: {self.working_dir}\n"
                    f"    stdout:\n{textwrap.indent(stdout, ' ' * 8)}\n\n"
                    f"    stderr:\n{textwrap.indent(stderr, ' ' * 8)}"
                )

    def _log_cmd(self, snippet: str) -> None:
        with open(self._command_log, "at", encoding="utf-8") as f:
            now = datetime.datetime.now().strftime("[%Y-%m-%d] %H:%M:%S.%f")
            f.write(f"# {now}\n{snippet}\n\n")

    @contextlib.contextmanager
    def watch_stdout_stderr(self):
        # Get initial state for stdout and stderr, since we only want to
        # capture output since the contextmanager started.
        stdout_stat = os.stat(self._stdout_f.name)
        stderr_stat = os.stat(self._stderr_f.name)

        def get() -> typing.Tuple[str, str]:
            with open(self._stdout_f.name, "rb") as f:
                _ = f.seek(stdout_stat.st_size)
                stdout = f.read().decode("utf-8").strip()

            with open(self._stderr_f.name, "rb") as f:
                _ = f.seek(stderr_stat.st_size)
                stderr = f.read().decode("utf-8").strip()

            return stdout, stderr

        yield get

    def _run(self, snippet: str) -> None:
        """Helper method for running code in a subprocess."""
        assert self._worker_bootstrap_finished
        assert self.alive, "Process has exited"

        snippet = textwrap.dedent(snippet)
        with self.watch_stdout_stderr() as get_output:
            self._input_pipe.write(snippet.encode(subprocess_rpc.ENCODING))
            self._log_cmd(snippet)

            result = marshal.loads(self._output_pipe.read())
            if isinstance(result, str):
                assert result == subprocess_rpc.SUCCESS
                return

            assert isinstance(result, dict)
            if not result:
                stdout, stderr = get_output()
                raise subprocess.SubprocessError(
                    "Uncaught Exception in worker:"
                    f"    working_dir: {self.working_dir}\n"
                    f"    stdout:\n{textwrap.indent(stdout, ' ' * 8)}\n\n"
                    f"    stderr:\n{textwrap.indent(stderr, ' ' * 8)}")

            serialized_e = subprocess_rpc.SerializedException(**result)
            stdout, stderr = get_output()
            subprocess_rpc.SerializedException.raise_from(
                serialized_e=serialized_e,
                extra_context=(
                    f"    working_dir: {self.working_dir}\n"
                    f"    stdout:\n{textwrap.indent(stdout, ' ' * 8)}\n\n"
                    f"    stderr:\n{textwrap.indent(stderr, ' ' * 8)}"
                )
            )

    def _kill_proc(self) -> None:
        """Best effort to kill subprocess."""
        if getattr(self, "_proc", None) is None:
            # We failed in the constructor, so there's nothing to clean up.
            return

        self._input_pipe.write(subprocess_rpc.HARD_EXIT)
        try:
            self._proc.wait(timeout=1)

        except subprocess.TimeoutExpired:
            if not subprocess_rpc.IS_WINDOWS:
                self._proc.send_signal(signal.SIGINT)

            try:
                self._proc.terminate()

            except PermissionError:
                # NoisePoliceWorker runs under sudo, and thus will not allow
                # SIGTERM to be sent.
                print(f"Failed to clean up process {self._proc.pid}")

        # Unfortunately Popen does not clean up stdin when using PIPE. However
        # we also can't unconditionally close the fd as it could interfere with
        # the orderly teardown of the process. We try our best to kill
        # `self._proc` in the previous block; if `self._proc` is terminated we
        # make sure its stdin TextIOWrapper is closed as well.
        try:
            self._proc.wait(timeout=1)
            proc_stdin = self._proc.stdin
            if proc_stdin is not None:
                proc_stdin.close()
        except subprocess.TimeoutExpired:
            pass

        self._alive = False

    def __del__(self) -> None:
        self._kill_proc()

        # We own these fd's, and it seems that we can unconditionally close
        # them without impacting the shutdown of `self._proc`.
        self._stdout_f.close()
        self._stderr_f.close()

        # Finally, make sure we don't leak any files.
        shutil.rmtree(self._working_dir, ignore_errors=True)