ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/runner/metaflow_runner.py

@@ -1,11 +1,14 @@
 import importlib
+import inspect
 import os
 import sys
 import json
 
 from typing import Dict, Iterator, Optional, Tuple
 
-from metaflow import Run
+from metaflow import Run, Task
+
+from metaflow.metaflow_config import CLICK_API_PROCESS_CONFIG
 
 from metaflow.plugins import get_runner_cli
 
@@ -13,34 +16,41 @@ from .utils import (
     temporary_fifo,
     handle_timeout,
     async_handle_timeout,
+    with_dir,
 )
 from .subprocess_manager import CommandManager, SubprocessManager
 
 
-class ExecutingRun(object):
+class ExecutingProcess(object):
     """
-    This class contains a reference to a `metaflow.Run` object representing
-    the currently executing or finished run, as well as metadata related
-    to the process.
+    This is a base class for `ExecutingRun` and `ExecutingTask` classes.
+    The `ExecutingRun` and `ExecutingTask` classes are returned by methods
+    in `Runner` and `NBRunner`, and they are subclasses of this class.
 
-    `ExecutingRun` is returned by methods in `Runner` and `NBRunner`. It is not
-    meant to be instantiated directly.
+    The `ExecutingRun` class for instance contains a reference to a `metaflow.Run`
+    object representing the currently executing or finished run, as well as the metadata
+    related to the process.
+
+    Similarly, the `ExecutingTask` class contains a reference to a `metaflow.Task`
+    object representing the currently executing or finished task, as well as the metadata
+    related to the process.
+
+    This class or its subclasses are not meant to be instantiated directly. The class
+    works as a context manager, allowing you to use a pattern like:
 
-    This class works as a context manager, allowing you to use a pattern like
     ```python
     with Runner(...).run() as running:
         ...
     ```
-    Note that you should use either this object as the context manager or
-    `Runner`, not both in a nested manner.
+
+    Note that you should use either this object as the context manager or `Runner`, not both
+    in a nested manner.
     """
 
-    def __init__(
-        self, runner: "Runner", command_obj: CommandManager, run_obj: Run
-    ) -> None:
+    def __init__(self, runner: "Runner", command_obj: CommandManager) -> None:
         """
         Create a new ExecutingRun -- this should not be done by the user directly but
-        instead user Runner.run()
+        instead use Runner.run()
 
         Parameters
         ----------
@@ -53,9 +63,8 @@ class ExecutingRun(object):
         """
         self.runner = runner
         self.command_obj = command_obj
-        self.run = run_obj
 
-    def __enter__(self) -> "ExecutingRun":
+    def __enter__(self) -> "ExecutingProcess":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -63,7 +72,7 @@ class ExecutingRun(object):
 
     async def wait(
         self, timeout: Optional[float] = None, stream: Optional[str] = None
-    ) -> "ExecutingRun":
+    ) -> "ExecutingProcess":
         """
         Wait for this run to finish, optionally with a timeout
         and optionally streaming its output.
@@ -82,7 +91,7 @@ class ExecutingRun(object):
 
         Returns
         -------
-        ExecutingRun
+        ExecutingProcess
             This object, allowing you to chain calls.
         """
         await self.command_obj.wait(timeout, stream)
@@ -189,6 +198,76 @@ class ExecutingRun(object):
             yield position, line
 
 
+class ExecutingTask(ExecutingProcess):
+    """
+    This class contains a reference to a `metaflow.Task` object representing
+    the currently executing or finished task, as well as metadata related
+    to the process.
+    `ExecutingTask` is returned by methods in `Runner` and `NBRunner`. It is not
+    meant to be instantiated directly.
+    This class works as a context manager, allowing you to use a pattern like
+    ```python
+    with Runner(...).spin() as running:
+        ...
+    ```
+    Note that you should use either this object as the context manager or
+    `Runner`, not both in a nested manner.
+    """
+
+    def __init__(
+        self, runner: "Runner", command_obj: CommandManager, task_obj: Task
+    ) -> None:
+        """
+        Create a new ExecutingTask -- this should not be done by the user directly but
+        instead use Runner.spin()
+        Parameters
+        ----------
+        runner : Runner
+            Parent runner for this task.
+        command_obj : CommandManager
+            CommandManager containing the subprocess executing this task.
+        task_obj : Task
+            Task object corresponding to this task.
+        """
+        super().__init__(runner, command_obj)
+        self.task = task_obj
+
+
+class ExecutingRun(ExecutingProcess):
+    """
+    This class contains a reference to a `metaflow.Run` object representing
+    the currently executing or finished run, as well as metadata related
+    to the process.
+    `ExecutingRun` is returned by methods in `Runner` and `NBRunner`. It is not
+    meant to be instantiated directly.
+    This class works as a context manager, allowing you to use a pattern like
+    ```python
+    with Runner(...).run() as running:
+        ...
+    ```
+    Note that you should use either this object as the context manager or
+    `Runner`, not both in a nested manner.
+    """
+
+    def __init__(
+        self, runner: "Runner", command_obj: CommandManager, run_obj: Run
+    ) -> None:
+        """
+        Create a new ExecutingRun -- this should not be done by the user directly but
+        instead use Runner.run()
+        Parameters
+        ----------
+        runner : Runner
+            Parent runner for this run.
+        command_obj : CommandManager
+            CommandManager containing the subprocess executing this run.
+        run_obj : Run
+            Run object corresponding to this run.
+        """
+        super().__init__(runner, command_obj)
+        self.run = run_obj
+
+
 class RunnerMeta(type):
     def __new__(mcs, name, bases, dct):
         cls = super().__new__(mcs, name, bases, dct)
@@ -197,8 +276,22 @@ class RunnerMeta(type):
             def f(self, *args, **kwargs):
                 return runner_subcommand(self, *args, **kwargs)
 
-            f.__doc__ = runner_subcommand.__doc__ or ""
+            f.__doc__ = runner_subcommand.__init__.__doc__ or ""
             f.__name__ = subcommand_name
+            sig = inspect.signature(runner_subcommand)
+            # We take all the same parameters except replace the first with
+            # simple "self"
+            new_parameters = {}
+            for name, param in sig.parameters.items():
+                if new_parameters:
+                    new_parameters[name] = param
+                else:
+                    new_parameters["self"] = inspect.Parameter(
+                        "self", inspect.Parameter.POSITIONAL_OR_KEYWORD
+                    )
+            f.__signature__ = inspect.Signature(
+                list(new_parameters.values()), return_annotation=runner_subcommand
+            )
 
             return f
 
@@ -257,7 +350,7 @@ class Runner(metaclass=RunnerMeta):
         env: Optional[Dict[str, str]] = None,
         cwd: Optional[str] = None,
         file_read_timeout: int = 3600,
-        **kwargs
+        **kwargs,
     ):
         # these imports are required here and not at the top
         # since they interfere with the user defined Parameters
@@ -299,7 +392,7 @@ class Runner(metaclass=RunnerMeta):
         if profile:
             self.env_vars["METAFLOW_PROFILE"] = profile
 
-        self.cwd = cwd
+        self.cwd = cwd or os.getcwd()
         self.file_read_timeout = file_read_timeout
         self.spm = SubprocessManager()
         self.top_level_kwargs = kwargs
@@ -359,9 +452,15 @@ class Runner(metaclass=RunnerMeta):
             ExecutingRun containing the results of the run.
         """
         with temporary_fifo() as (attribute_file_path, attribute_file_fd):
-            command = self.api(**self.top_level_kwargs).run(
-                runner_attribute_file=attribute_file_path, **kwargs
-            )
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).run(
+                        runner_attribute_file=attribute_file_path, **kwargs
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).run(
+                    runner_attribute_file=attribute_file_path, **kwargs
+                )
 
             pid = self.spm.run_command(
                 [sys.executable, *command],
@@ -373,6 +472,78 @@ class Runner(metaclass=RunnerMeta):
 
             return self.__get_executing_run(attribute_file_fd, command_obj)
 
+    def __get_executing_task(self, attribute_file_fd, command_obj):
+        content = handle_timeout(attribute_file_fd, command_obj, self.file_read_timeout)
+
+        command_obj.sync_wait()
+
+        content = json.loads(content)
+        pathspec = f"{content.get('flow_name')}/{content.get('run_id')}/{content.get('step_name')}/{content.get('task_id')}"
+
+        # Set the correct metadata from the runner_attribute file corresponding to this run.
+        metadata_for_flow = content.get("metadata")
+
+        task_object = Task(
+            pathspec, _namespace_check=False, _current_metadata=metadata_for_flow
+        )
+        return ExecutingTask(self, command_obj, task_object)
+
+    async def __async_get_executing_task(self, attribute_file_fd, command_obj):
+        content = await async_handle_timeout(
+            attribute_file_fd, command_obj, self.file_read_timeout
+        )
+        content = json.loads(content)
+        pathspec = f"{content.get('flow_name')}/{content.get('run_id')}/{content.get('step_name')}/{content.get('task_id')}"
+
+        # Set the correct metadata from the runner_attribute file corresponding to this run.
+        metadata_for_flow = content.get("metadata")
+
+        task_object = Task(
+            pathspec, _namespace_check=False, _current_metadata=metadata_for_flow
+        )
+        return ExecutingTask(self, command_obj, task_object)
+
+    def spin(self, pathspec, **kwargs) -> ExecutingTask:
+        """
+        Blocking spin execution of the run.
+        This method will wait until the spun run has completed execution.
+        Parameters
+        ----------
+        pathspec : str
+            The pathspec of the step/task to spin.
+        **kwargs : Any
+            Additional arguments that you would pass to `python ./myflow.py` after
+            the `spin` command.
+        Returns
+        -------
+        ExecutingTask
+            ExecutingTask containing the results of the spun task.
+        """
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).spin(
+                        pathspec=pathspec,
+                        runner_attribute_file=attribute_file_path,
+                        **kwargs,
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).spin(
+                    pathspec=pathspec,
+                    runner_attribute_file=attribute_file_path,
+                    **kwargs,
+                )
+
+            pid = self.spm.run_command(
+                [sys.executable, *command],
+                env=self.env_vars,
+                cwd=self.cwd,
+                show_output=self.show_output,
+            )
+            command_obj = self.spm.get(pid)
+
+            return self.__get_executing_task(attribute_file_fd, command_obj)
+
     def resume(self, **kwargs) -> ExecutingRun:
         """
         Blocking resume execution of the run.
@@ -390,9 +561,15 @@ class Runner(metaclass=RunnerMeta):
             ExecutingRun containing the results of the resumed run.
         """
         with temporary_fifo() as (attribute_file_path, attribute_file_fd):
-            command = self.api(**self.top_level_kwargs).resume(
-                runner_attribute_file=attribute_file_path, **kwargs
-            )
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).resume(
+                        runner_attribute_file=attribute_file_path, **kwargs
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).resume(
+                    runner_attribute_file=attribute_file_path, **kwargs
+                )
 
             pid = self.spm.run_command(
                 [sys.executable, *command],
@@ -423,9 +600,15 @@ class Runner(metaclass=RunnerMeta):
             ExecutingRun representing the run that was started.
         """
         with temporary_fifo() as (attribute_file_path, attribute_file_fd):
-            command = self.api(**self.top_level_kwargs).run(
-                runner_attribute_file=attribute_file_path, **kwargs
-            )
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).run(
+                        runner_attribute_file=attribute_file_path, **kwargs
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).run(
+                    runner_attribute_file=attribute_file_path, **kwargs
+                )
 
             pid = await self.spm.async_run_command(
                 [sys.executable, *command],
@@ -455,9 +638,15 @@ class Runner(metaclass=RunnerMeta):
             ExecutingRun representing the resumed run that was started.
         """
         with temporary_fifo() as (attribute_file_path, attribute_file_fd):
-            command = self.api(**self.top_level_kwargs).resume(
-                runner_attribute_file=attribute_file_path, **kwargs
-            )
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).resume(
+                        runner_attribute_file=attribute_file_path, **kwargs
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).resume(
+                    runner_attribute_file=attribute_file_path, **kwargs
+                )
 
             pid = await self.spm.async_run_command(
                 [sys.executable, *command],
@@ -468,6 +657,50 @@ class Runner(metaclass=RunnerMeta):
 
             return await self.__async_get_executing_run(attribute_file_fd, command_obj)
 
+    async def async_spin(self, pathspec, **kwargs) -> ExecutingTask:
+        """
+        Non-blocking spin execution of the run.
+        This method will return as soon as the spun task has launched.
+
+        Note that this method is asynchronous and needs to be `await`ed.
+
+        Parameters
+        ----------
+        pathspec : str
+            The pathspec of the step/task to spin.
+        **kwargs : Any
+            Additional arguments that you would pass to `python ./myflow.py` after
+            the `spin` command.
+
+        Returns
+        -------
+        ExecutingTask
+            ExecutingTask representing the spun task that was started.
+        """
+        with temporary_fifo() as (attribute_file_path, attribute_file_fd):
+            if CLICK_API_PROCESS_CONFIG:
+                with with_dir(self.cwd):
+                    command = self.api(**self.top_level_kwargs).spin(
+                        pathspec=pathspec,
+                        runner_attribute_file=attribute_file_path,
+                        **kwargs,
+                    )
+            else:
+                command = self.api(**self.top_level_kwargs).spin(
+                    pathspec=pathspec,
+                    runner_attribute_file=attribute_file_path,
+                    **kwargs,
+                )
+
+            pid = await self.spm.async_run_command(
+                [sys.executable, *command],
+                env=self.env_vars,
+                cwd=self.cwd,
+            )
+            command_obj = self.spm.get(pid)
+
+            return await self.__async_get_executing_task(attribute_file_fd, command_obj)
+
     def __exit__(self, exc_type, exc_value, traceback):
         self.spm.cleanup()
 

metaflow/runner/subprocess_manager.py

@@ -9,6 +9,8 @@ import tempfile
 import threading
 from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
+from metaflow.packaging_sys import MetaflowCodeContent
+from metaflow.util import get_metaflow_root
 from .utils import check_process_exited
 
 
@@ -150,6 +152,19 @@ class SubprocessManager(object):
         int
             The process ID of the subprocess.
         """
+        env = env or {}
+        installed_root = os.environ.get("METAFLOW_EXTRACTED_ROOT", get_metaflow_root())
+
+        for k, v in MetaflowCodeContent.get_env_vars_for_packaged_metaflow(
+            installed_root
+        ).items():
+            if k.endswith(":"):
+                # Override
+                env[k[:-1]] = v
+            elif k in env:
+                env[k] = "%s:%s" % (v, env[k])
+            else:
+                env[k] = v
 
         command_obj = CommandManager(command, env, cwd)
         pid = command_obj.run(show_output=show_output)
@@ -181,6 +196,11 @@ class SubprocessManager(object):
         int
             The process ID of the subprocess.
         """
+        env = env or {}
+        if "PYTHONPATH" in env:
+            env["PYTHONPATH"] = "%s:%s" % (get_metaflow_root(), env["PYTHONPATH"])
+        else:
+            env["PYTHONPATH"] = get_metaflow_root()
 
         command_obj = CommandManager(command, env, cwd)
         pid = await command_obj.async_run()
@@ -237,7 +257,7 @@ class CommandManager(object):
         self.command = command
 
         self.env = env if env is not None else os.environ.copy()
-        self.cwd = cwd if cwd is not None else os.getcwd()
+        self.cwd = cwd or os.getcwd()
 
         self.process = None
         self.stdout_thread = None

metaflow/runner/utils.py

@@ -4,6 +4,7 @@ import time
 import asyncio
 import tempfile
 import select
+import fcntl
 from contextlib import contextmanager
 from subprocess import CalledProcessError
 from typing import Any, Dict, TYPE_CHECKING, ContextManager, Tuple
@@ -109,7 +110,6 @@ def read_from_fifo_when_ready(
     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(
@@ -130,6 +130,21 @@ def read_from_fifo_when_ready(
             data = os.read(fifo_fd, 8192)
             if data:
                 content += data
+                # We got data! Now switch to blocking mode for guaranteed complete reads.
+                # In blocking mode, read() won't return 0 until writer closes AND all
+                # kernel buffers are drained - this is POSIX guaranteed.
+                flags = fcntl.fcntl(fifo_fd, fcntl.F_GETFL)
+                fcntl.fcntl(fifo_fd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
+
+                # Now do blocking reads until true EOF
+                while True:
+                    chunk = os.read(fifo_fd, 8192)
+                    if not chunk:
+                        # True EOF - all data drained
+                        break
+                    content += chunk
+                # All data read, exit main loop
+                break
             else:
                 if len(events):
                     # We read an EOF -- consider the file done
@@ -137,22 +152,10 @@ def read_from_fifo_when_ready(
                 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
+                    pass
         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.
+            # File not ready yet, continue waiting
+            pass
 
     if not content and check_process_exited(command_obj):
         raise CalledProcessError(command_obj.process.returncode, command_obj.command)
@@ -322,3 +325,11 @@ def get_lower_level_group(
         raise ValueError(f"Sub-command '{sub_command}' not found in API '{api.name}'")
 
     return sub_command_obj(**sub_command_kwargs)
+
+
+@contextmanager
+def with_dir(new_dir):
+    current_dir = os.getcwd()
+    os.chdir(new_dir)
+    yield new_dir
+    os.chdir(current_dir)