ob-metaflow 2.12.30.2__py2.py3-none-any.whl → 2.13.6.1__py2.py3-none-any.whl

metaflow/runner/nbrun.py

@@ -44,7 +44,7 @@ class NBRunner(object):
         The directory to run the subprocess in; if not specified, the current
         working directory is used.
     file_read_timeout : int, default 3600
-        The timeout until which we try to read the runner attribute file.
+        The timeout until which we try to read the runner attribute file (in seconds).
     **kwargs : Any
         Additional arguments that you would pass to `python myflow.py` before
         the `run` command.

metaflow/runner/subprocess_manager.py

@@ -9,26 +9,61 @@ import tempfile
 import threading
 from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
+from .utils import check_process_exited
 
-def kill_process_and_descendants(pid, termination_timeout):
+
+def kill_processes_and_descendants(pids: List[str], termination_timeout: float):
     # TODO: there's a race condition that new descendants might
     # spawn b/w the invocations of 'pkill' and 'kill'.
     # Needs to be fixed in future.
     try:
-        subprocess.check_call(["pkill", "-TERM", "-P", str(pid)])
-        subprocess.check_call(["kill", "-TERM", str(pid)])
+        subprocess.check_call(["pkill", "-TERM", "-P", *pids])
+        subprocess.check_call(["kill", "-TERM", *pids])
     except subprocess.CalledProcessError:
         pass
 
     time.sleep(termination_timeout)
 
     try:
-        subprocess.check_call(["pkill", "-KILL", "-P", str(pid)])
-        subprocess.check_call(["kill", "-KILL", str(pid)])
+        subprocess.check_call(["pkill", "-KILL", "-P", *pids])
+        subprocess.check_call(["kill", "-KILL", *pids])
     except subprocess.CalledProcessError:
         pass
 
 
+async def async_kill_processes_and_descendants(
+    pids: List[str], termination_timeout: float
+):
+    # TODO: there's a race condition that new descendants might
+    # spawn b/w the invocations of 'pkill' and 'kill'.
+    # Needs to be fixed in future.
+    try:
+        sub_term = await asyncio.create_subprocess_exec("pkill", "-TERM", "-P", *pids)
+        await sub_term.wait()
+    except Exception:
+        pass
+
+    try:
+        main_term = await asyncio.create_subprocess_exec("kill", "-TERM", *pids)
+        await main_term.wait()
+    except Exception:
+        pass
+
+    await asyncio.sleep(termination_timeout)
+
+    try:
+        sub_kill = await asyncio.create_subprocess_exec("pkill", "-KILL", "-P", *pids)
+        await sub_kill.wait()
+    except Exception:
+        pass
+
+    try:
+        main_kill = await asyncio.create_subprocess_exec("kill", "-KILL", *pids)
+        await main_kill.wait()
+    except Exception:
+        pass
+
+
 class LogReadTimeoutError(Exception):
     """Exception raised when reading logs times out."""
 
@@ -46,14 +81,28 @@ class SubprocessManager(object):
             loop = asyncio.get_running_loop()
             loop.add_signal_handler(
                 signal.SIGINT,
-                lambda: self._handle_sigint(signum=signal.SIGINT, frame=None),
+                lambda: asyncio.create_task(self._async_handle_sigint()),
             )
         except RuntimeError:
             signal.signal(signal.SIGINT, self._handle_sigint)
 
+    async def _async_handle_sigint(self):
+        pids = [
+            str(command.process.pid)
+            for command in self.commands.values()
+            if command.process and not check_process_exited(command)
+        ]
+        if pids:
+            await async_kill_processes_and_descendants(pids, termination_timeout=2)
+
     def _handle_sigint(self, signum, frame):
-        for each_command in self.commands.values():
-            each_command.kill(termination_timeout=2)
+        pids = [
+            str(command.process.pid)
+            for command in self.commands.values()
+            if command.process and not check_process_exited(command)
+        ]
+        if pids:
+            kill_processes_and_descendants(pids, termination_timeout=2)
 
     async def __aenter__(self) -> "SubprocessManager":
         return self
@@ -71,6 +120,9 @@ class SubprocessManager(object):
         """
         Run a command synchronously and return its process ID.
 
+        Note: in no case does this wait for the process to *finish*. Use sync_wait()
+        to wait for the command to finish.
+
         Parameters
         ----------
         command : List[str]
@@ -96,7 +148,6 @@ class SubprocessManager(object):
         command_obj = CommandManager(command, env, cwd)
         pid = command_obj.run(show_output=show_output)
         self.commands[pid] = command_obj
-        command_obj.sync_wait()
         return pid
 
     async def async_run_command(
@@ -472,7 +523,7 @@ class CommandManager(object):
         """
 
         if self.process is not None:
-            kill_process_and_descendants(self.process.pid, termination_timeout)
+            kill_processes_and_descendants([str(self.process.pid)], termination_timeout)
         else:
             print("No process to kill.")
 

metaflow/runner/utils.py

@@ -2,9 +2,11 @@ import os
 import ast
 import time
 import asyncio
-
+import tempfile
+import select
+from contextlib import contextmanager
 from subprocess import CalledProcessError
-from typing import Any, Dict, TYPE_CHECKING
+from typing import Any, Dict, TYPE_CHECKING, ContextManager, Tuple
 
 if TYPE_CHECKING:
     import tempfile
@@ -39,45 +41,216 @@ def format_flowfile(cell):
     return "\n".join(lines)
 
 
-def check_process_status(
+def check_process_exited(
     command_obj: "metaflow.runner.subprocess_manager.CommandManager",
-):
+) -> bool:
     if isinstance(command_obj.process, asyncio.subprocess.Process):
         return command_obj.process.returncode is not None
     else:
         return command_obj.process.poll() is not None
 
 
-def read_from_file_when_ready(
-    file_path: str,
+@contextmanager
+def temporary_fifo() -> ContextManager[Tuple[str, int]]:
+    """
+    Create and open the read side of a temporary FIFO in a non-blocking mode.
+
+    Returns
+    -------
+    str
+        Path to the temporary FIFO.
+    int
+        File descriptor of the temporary FIFO.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        path = os.path.join(temp_dir, "fifo")
+        os.mkfifo(path)
+        # Blocks until the write side is opened unless in non-blocking mode
+        fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK)
+        try:
+            yield path, fd
+        finally:
+            os.close(fd)
+
+
+def read_from_fifo_when_ready(
+    fifo_fd: int,
     command_obj: "metaflow.runner.subprocess_manager.CommandManager",
-    timeout: float = 5,
-):
-    start_time = time.time()
-    with open(file_path, "r", encoding="utf-8") as file_pointer:
-        content = file_pointer.read()
-        while not content:
-            if check_process_status(command_obj):
-                # Check to make sure the file hasn't been read yet to avoid a race
-                # where the file is written between the end of this while loop and the
-                # poll call above.
-                content = file_pointer.read()
-                if content:
+    encoding: str = "utf-8",
+    timeout: int = 3600,
+) -> str:
+    """
+    Read the content from the FIFO file descriptor when it is ready.
+
+    Parameters
+    ----------
+    fifo_fd : int
+        File descriptor of the FIFO.
+    command_obj : CommandManager
+        Command manager object that handles the write side of the FIFO.
+    encoding : str, optional
+        Encoding to use while reading the file, by default "utf-8".
+    timeout : int, optional
+        Timeout for reading the file in seconds, by default 3600.
+
+    Returns
+    -------
+    str
+        Content read from the FIFO.
+
+    Raises
+    ------
+    TimeoutError
+        If no event occurs on the FIFO within the timeout.
+    CalledProcessError
+        If the process managed by `command_obj` has exited without writing any
+        content to the FIFO.
+    """
+    content = bytearray()
+    poll = select.poll()
+    poll.register(fifo_fd, select.POLLIN)
+    max_timeout = 3  # Wait for 10 * 3 = 30 ms after last write
+    while True:
+        if check_process_exited(command_obj) and command_obj.process.returncode != 0:
+            raise CalledProcessError(
+                command_obj.process.returncode, command_obj.command
+            )
+
+        if timeout < 0:
+            raise TimeoutError("Timeout while waiting for the file content")
+
+        poll_begin = time.time()
+        # We poll for a very short time to be also able to check if the file was closed
+        # If the file is closed, we assume that we only have one writer so if we have
+        # data, we break out. This is to work around issues in macos
+        events = poll.poll(min(10, timeout * 1000))
+        timeout -= time.time() - poll_begin
+
+        try:
+            data = os.read(fifo_fd, 8192)
+            if data:
+                content += data
+            else:
+                if len(events):
+                    # We read an EOF -- consider the file done
                     break
-                raise CalledProcessError(
-                    command_obj.process.returncode, command_obj.command
-                )
-            if time.time() - start_time > timeout:
-                raise TimeoutError(
-                    "Timeout while waiting for file content from '%s'" % file_path
-                )
-            time.sleep(0.1)
-            content = file_pointer.read()
-        return content
+                else:
+                    # We had no events (just a timeout) and the read didn't return
+                    # an exception so the file is still open; we continue waiting for data
+                    # Unfortunately, on MacOS, it seems that even *after* the file is
+                    # closed on the other end, we still don't get a BlockingIOError so
+                    # we hack our way and timeout if there is no write in 30ms which is
+                    # a relative eternity for file writes.
+                    if content:
+                        if max_timeout <= 0:
+                            break
+                        max_timeout -= 1
+                        continue
+        except BlockingIOError:
+            has_blocking_error = True
+            if content:
+                # The file was closed
+                break
+            # else, if we have no content, we continue waiting for the file to be open
+            # and written to.
+
+    if not content and check_process_exited(command_obj):
+        raise CalledProcessError(command_obj.process.returncode, command_obj.command)
+
+    return content.decode(encoding)
+
+
+async def async_read_from_fifo_when_ready(
+    fifo_fd: int,
+    command_obj: "metaflow.runner.subprocess_manager.CommandManager",
+    encoding: str = "utf-8",
+    timeout: int = 3600,
+) -> str:
+    """
+    Read the content from the FIFO file descriptor when it is ready.
+
+    Parameters
+    ----------
+    fifo_fd : int
+        File descriptor of the FIFO.
+    command_obj : CommandManager
+        Command manager object that handles the write side of the FIFO.
+    encoding : str, optional
+        Encoding to use while reading the file, by default "utf-8".
+    timeout : int, optional
+        Timeout for reading the file in seconds, by default 3600.
+
+    Returns
+    -------
+    str
+        Content read from the FIFO.
+
+    Raises
+    ------
+    TimeoutError
+        If no event occurs on the FIFO within the timeout.
+    CalledProcessError
+        If the process managed by `command_obj` has exited without writing any
+        content to the FIFO.
+    """
+    return await asyncio.to_thread(
+        read_from_fifo_when_ready, fifo_fd, command_obj, encoding, timeout
+    )
+
+
+def make_process_error_message(
+    command_obj: "metaflow.runner.subprocess_manager.CommandManager",
+):
+    stdout_log = open(command_obj.log_files["stdout"], encoding="utf-8").read()
+    stderr_log = open(command_obj.log_files["stderr"], encoding="utf-8").read()
+    command = " ".join(command_obj.command)
+    error_message = "Error executing: '%s':\n" % command
+    if stdout_log.strip():
+        error_message += "\nStdout:\n%s\n" % stdout_log
+    if stderr_log.strip():
+        error_message += "\nStderr:\n%s\n" % stderr_log
+    return error_message
 
 
 def handle_timeout(
-    tfp_runner_attribute: "tempfile._TemporaryFileWrapper[str]",
+    attribute_file_fd: int,
+    command_obj: "metaflow.runner.subprocess_manager.CommandManager",
+    file_read_timeout: int,
+):
+    """
+    Handle the timeout for a running subprocess command that reads a file
+    and raises an error with appropriate logs if a TimeoutError occurs.
+
+    Parameters
+    ----------
+    attribute_file_fd : int
+        File descriptor belonging to the FIFO containing the attribute data.
+    command_obj : CommandManager
+        Command manager object that encapsulates the running command details.
+    file_read_timeout : int
+        Timeout for reading the file, in seconds
+
+    Returns
+    -------
+    str
+        Content read from the temporary file.
+
+    Raises
+    ------
+    RuntimeError
+        If a TimeoutError occurs, it raises a RuntimeError with the command's
+        stdout and stderr logs.
+    """
+    try:
+        return read_from_fifo_when_ready(
+            attribute_file_fd, command_obj=command_obj, timeout=file_read_timeout
+        )
+    except (CalledProcessError, TimeoutError) as e:
+        raise RuntimeError(make_process_error_message(command_obj)) from e
+
+
+async def async_handle_timeout(
+    attribute_file_fd: "int",
     command_obj: "metaflow.runner.subprocess_manager.CommandManager",
     file_read_timeout: int,
 ):
@@ -87,12 +260,12 @@ def handle_timeout(
 
     Parameters
     ----------
-    tfp_runner_attribute : NamedTemporaryFile
-        Temporary file that stores runner attribute data.
+    attribute_file_fd : int
+        File descriptor belonging to the FIFO containing the attribute data.
     command_obj : CommandManager
         Command manager object that encapsulates the running command details.
     file_read_timeout : int
-        Timeout for reading the file.
+        Timeout for reading the file, in seconds
 
     Returns
     -------
@@ -106,20 +279,11 @@ def handle_timeout(
         stdout and stderr logs.
     """
     try:
-        content = read_from_file_when_ready(
-            tfp_runner_attribute.name, command_obj, timeout=file_read_timeout
+        return await async_read_from_fifo_when_ready(
+            attribute_file_fd, command_obj=command_obj, timeout=file_read_timeout
         )
-        return content
     except (CalledProcessError, TimeoutError) as e:
-        stdout_log = open(command_obj.log_files["stdout"], encoding="utf-8").read()
-        stderr_log = open(command_obj.log_files["stderr"], encoding="utf-8").read()
-        command = " ".join(command_obj.command)
-        error_message = "Error executing: '%s':\n" % command
-        if stdout_log.strip():
-            error_message += "\nStdout:\n%s\n" % stdout_log
-        if stderr_log.strip():
-            error_message += "\nStderr:\n%s\n" % stderr_log
-        raise RuntimeError(error_message) from e
+        raise RuntimeError(make_process_error_message(command_obj)) from e
 
 
 def get_lower_level_group(