suggested changes

Author: madhur-ob

metaflow/cli.py

@@ -678,11 +678,11 @@ def common_run_options(func):
         help="Write the ID of this run to the file specified.",
     )
     @click.option(
-        "--flow-name-file",
+        "--pathspec-file",
         default=None,
         show_default=True,
         type=str,
-        help="Write the flow name of this run to the file specified.",
+        help="Write the pathspec of this run to the file specified.",
     )
     @wraps(func)
     def wrapper(*args, **kwargs):
@@ -742,7 +742,7 @@ def resume(
     decospecs=None,
     run_id_file=None,
     resume_identifier=None,
-    flow_name_file=None,
+    pathspec_file=None,
 ):
     before_run(obj, tags, decospecs + obj.environment.decospecs())
 
@@ -799,7 +799,7 @@ def resume(
         resume_identifier=resume_identifier,
     )
     write_file(run_id_file, runtime.run_id)
-    write_file(flow_name_file, obj.flow.name)
+    write_file(pathspec_file, "/".join((obj.flow.name, runtime.run_id)))
     runtime.print_workflow_info()
 
     runtime.persist_constants()
@@ -833,7 +833,7 @@ def run(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
-    flow_name_file=None,
+    pathspec_file=None,
     user_namespace=None,
     **kwargs
 ):
@@ -858,7 +858,7 @@ def run(
     )
     write_latest_run_id(obj, runtime.run_id)
     write_file(run_id_file, runtime.run_id)
-    write_file(flow_name_file, obj.flow.name)
+    write_file(pathspec_file, "/".join((obj.flow.name, runtime.run_id)))
 
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.print_workflow_info()

metaflow/click_api.py

@@ -75,7 +75,7 @@ def _method_sanity_check(
         if supplied_k not in possible_params:
             raise ValueError(
                 "Unknown argument: '%s', possible args are: %s"
-                % (supplied_k, list(possible_params.keys()))
+                % (supplied_k, ", ".join(possible_params.keys()))
             )
 
         try:

metaflow/metaflow_runner.py

@@ -8,16 +8,33 @@
 from metaflow.subprocess_manager import SubprocessManager, CommandManager
 
 
-def read_from_file_when_ready(file_path):
+def read_from_file_when_ready(file_path: str, timeout: float = 5):
+    start_time = time.time()
     with open(file_path, "r") as file_pointer:
         content = file_pointer.read()
         while not content:
+            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
 
 
 class ExecutingRun(object):
+    """
+    An object that encapsulates both the:
+        - CommandManager Object (that has ability to stream logs, kill the underlying process, etc.)
+        - Metaflow's Run object
+
+    This is a user facing object exposing methods and properties for:
+        - waiting (with an optional timeout and optional streaming of logs)
+        - current snapshot of stdout
+        - current snapshot of stderr
+        - ability to stream logs
+    """
+
     def __init__(
         self, runner: "Runner", command_obj: CommandManager, run_obj: Run
     ) -> None:
@@ -32,21 +49,25 @@ def __exit__(self, exc_type, exc_value, traceback):
         self.runner.__exit__(exc_type, exc_value, traceback)
 
     async def wait(self, timeout: Optional[float] = None, stream: Optional[str] = None):
+        """Wait for the run to finish, optionally with a timeout and optionally streaming its output."""
         await self.command_obj.wait(timeout, stream)
         return self
 
     @property
     def stdout(self):
+        """Get the current snapshot of stdout from the log file."""
         with open(self.command_obj.log_files.get("stdout"), "r") as fp:
             return fp.read()
 
     @property
     def stderr(self):
+        """Get the current snapshot of stderr from the log file."""
         with open(self.command_obj.log_files.get("stderr"), "r") as fp:
             return fp.read()
 
-    async def stream_logs(self, stream: str, position: Optional[int] = None):
-        async for position, line in self.command_obj.stream_logs(stream, position):
+    async def stream_log(self, stream: str, position: Optional[int] = None):
+        """Stream logs from the run using the log files. Used with an async for loop"""
+        async for position, line in self.command_obj.stream_log(stream, position):
             yield position, line
 
 
@@ -58,6 +79,22 @@ def __init__(
         env: Dict = {},
         **kwargs,
     ):
+        """
+        Metaflow's Runner API that presents a programmatic interface
+        to run flows either synchronously or asynchronously. The class expects
+        a path to the flow file along with optional top level parameters such as:
+        metadata, environment, datastore, etc.
+
+        The run() method expects run-level parameters. The object returned by the
+        'run()' method has access to Metaflow's Run object using `.run` along with
+        other abilities such as streaming logs, etc.
+
+        Example:
+            with metaflow_runner.Runner('../try.py', metadata="local") as runner:
+                result = runner.run(alpha=5, tags=["abc", "def"], max_workers=5)
+                print(result.run.finished)
+        """
+
         # these imports are required here and not at the top
         # since they interfere with the user defined Parameters
         # in the flow file, this is related to the ability of
@@ -95,40 +132,29 @@ def run(self, **kwargs):
 
     async def async_run(self, **kwargs):
         with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_flow = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
-            tfp_run_id = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
+            tfp_pathspec = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
 
-            command = self.runner(
-                run_id_file=tfp_run_id.name, flow_name_file=tfp_flow.name, **kwargs
-            )
+            command = self.runner(pathspec_file=tfp_pathspec.name, **kwargs)
 
             pid = await self.spm.run_command(
                 [sys.executable, *command], env=self.env_vars
             )
             command_obj = self.spm.get(pid)
 
-            # detect failures even before writing to the run_id and flow_name files
-            # the error (if any) must happen within the first 0.5 seconds
             try:
-                await asyncio.wait_for(command_obj.process.wait(), timeout=0.5)
-            except asyncio.TimeoutError:
-                pass
-
-            # if the returncode is None, the process has encountered no error within the
-            # initial 0.5 seconds and we proceed to run it in the background
-            # during which it would have written to the run_id and flow_name files
-            if command_obj.process.returncode is not None:
+                pathspec = read_from_file_when_ready(tfp_pathspec.name, timeout=5)
+                run_object = Run(pathspec, _namespace_check=False)
+                return ExecutingRun(self, command_obj, run_object)
+            except TimeoutError as e:
+                stdout_log = open(command_obj.log_files["stdout"]).read()
                 stderr_log = open(command_obj.log_files["stderr"]).read()
                 command = " ".join(command_obj.command)
-                raise RuntimeError(f"Error executing: '{command}':\n\n{stderr_log}")
-            else:
-                flow_name = read_from_file_when_ready(tfp_flow.name)
-                run_id = read_from_file_when_ready(tfp_run_id.name)
-
-                pathspec_components = (flow_name, run_id)
-                run_object = Run("/".join(pathspec_components), _namespace_check=False)
-
-                return ExecutingRun(self, command_obj, run_object)
+                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
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.spm.cleanup()

metaflow/subprocess_manager.py

@@ -53,9 +53,9 @@ async def run_command(
         """Run a command asynchronously and return its process ID."""
 
         command_obj = CommandManager(command, env, cwd)
-        process = await command_obj.run()
-        self.commands[process.pid] = command_obj
-        return process.pid
+        pid = await command_obj.run()
+        self.commands[pid] = command_obj
+        return pid
 
     def get(self, pid: int) -> "CommandManager":
         """Get the CommandManager object for a given process ID."""
@@ -98,7 +98,6 @@ async def __aexit__(self, exc_type, exc_value, traceback):
     def handle_sigint(self, signum, frame):
         """Handle the SIGINT signal."""
 
-        print("SIGINT received.")
         asyncio.create_task(self.kill())
 
     async def wait(
@@ -121,38 +120,45 @@ async def wait(
                 command_string = " ".join(self.command)
                 await self.kill()
                 print(
-                    "Timeout: The process: %s with command: '%s' didn't complete within %s seconds."
+                    "Timeout: The process (PID %d; command: '%s') did not complete within %s seconds."
                     % (self.process.pid, command_string, timeout)
                 )
 
     async def run(self):
         """Run the subprocess, streaming the logs to temporary files"""
 
-        self.temp_dir = tempfile.mkdtemp()
-        stdout_logfile = os.path.join(self.temp_dir, "stdout.log")
-        stderr_logfile = os.path.join(self.temp_dir, "stderr.log")
-
-        try:
-            # returns when process has been started,
-            # not when it is finished...
-            self.process = await asyncio.create_subprocess_exec(
-                *self.command,
-                cwd=self.cwd,
-                env=self.env,
-                stdout=open(stdout_logfile, "w"),
-                stderr=open(stderr_logfile, "w"),
-            )
+        if not self.run_called:
+            self.temp_dir = tempfile.mkdtemp()
+            stdout_logfile = os.path.join(self.temp_dir, "stdout.log")
+            stderr_logfile = os.path.join(self.temp_dir, "stderr.log")
+
+            try:
+                # returns when process has been started,
+                # not when it is finished...
+                self.process = await asyncio.create_subprocess_exec(
+                    *self.command,
+                    cwd=self.cwd,
+                    env=self.env,
+                    stdout=open(stdout_logfile, "w"),
+                    stderr=open(stderr_logfile, "w"),
+                )
 
-            self.log_files["stdout"] = stdout_logfile
-            self.log_files["stderr"] = stderr_logfile
+                self.log_files["stdout"] = stdout_logfile
+                self.log_files["stderr"] = stderr_logfile
 
-            self.run_called = True
-            return self.process
-        except Exception as e:
-            print("Error starting subprocess: %s" % e)
-            self.cleanup()
+                self.run_called = True
+                return self.process.pid
+            except Exception as e:
+                print("Error starting subprocess: %s" % e)
+                self.cleanup()
+        else:
+            command_string = " ".join(self.command)
+            print(
+                "Command '%s' has already been called. Please create another CommandManager object."
+                % command_string
+            )
 
-    async def stream_logs(
+    async def stream_log(
         self,
         stream: str,
         position: Optional[int] = None,
@@ -161,13 +167,13 @@ async def stream_logs(
     ):
         """Stream logs from the subprocess using the log files"""
 
-        if self.run_called is False:
-            raise ValueError("No command run yet to get the logs for...")
+        if not self.run_called:
+            raise RuntimeError("No command run yet to get the logs for...")
 
         if stream not in self.log_files:
             raise ValueError(
                 "No log file found for '%s', valid values are: %s"
-                % (stream, list(self.log_files.keys()))
+                % (stream, ", ".join(self.log_files.keys()))
             )
 
         log_file = self.log_files[stream]
@@ -206,18 +212,18 @@ async def stream_logs(
                         continue
 
                 position = f.tell()
-                yield position, line.strip()
+                yield position, line.rstrip()
 
     async def emit_logs(self, stream: str = "stdout", custom_logger: Callable = print):
-        """Helper function to iterate over stream_logs"""
+        """Helper function to iterate over stream_log"""
 
-        async for _, line in self.stream_logs(stream):
+        async for _, line in self.stream_log(stream):
             custom_logger(line)
 
     def cleanup(self):
         """Clean up log files for a running subprocesses."""
 
-        if hasattr(self, "temp_dir"):
+        if self.run_called:
             shutil.rmtree(self.temp_dir, ignore_errors=True)
 
     async def kill(self, termination_timeout: float = 5):
@@ -272,7 +278,7 @@ async def main():
 
         # stream logs line by line and check for existence of a string, noting down the position
         interesting_position = 0
-        async for position, line in command_obj.stream_logs(stream="stdout"):
+        async for position, line in command_obj.stream_log(stream="stdout"):
             print(line)
             if "alpha is" in line:
                 interesting_position = position
@@ -290,21 +296,21 @@ async def main():
             "resuming streaming from: %s while process is still running..."
             % interesting_position
         )
-        async for position, line in command_obj.stream_logs(
+        async for position, line in command_obj.stream_log(
             stream="stdout", position=interesting_position
         ):
             print(line)
 
         # this will be instantaneous since the process has finished and we just read from the log file
         print("process has ended by now... streaming again from scratch..")
-        async for position, line in command_obj.stream_logs(stream="stdout"):
+        async for position, line in command_obj.stream_log(stream="stdout"):
             print(line)
 
         # this will be instantaneous since the process has finished and we just read from the log file
         print(
             "process has ended by now... streaming again but from position of choice.."
         )
-        async for position, line in command_obj.stream_logs(
+        async for position, line in command_obj.stream_log(
             stream="stdout", position=interesting_position
         ):
             print(line)