suggested improvements

metaflow/subprocess_manager.py

@@ -1,94 +1,112 @@
 import os
 import sys
-import time
 import signal
 import shutil
-import hashlib
 import asyncio
 import tempfile
-from typing import List
-
-
-def hash_command_invocation(command: List[str]):
-    concatenated_string = "".join(command)
-    current_time = str(time.time())
-    concatenated_string += current_time
-    hash_object = hashlib.sha256(concatenated_string.encode())
-    return hash_object.hexdigest()
+from typing import List, Dict, Optional, Callable
 
 
 class LogReadTimeoutError(Exception):
+    """Exception raised when reading logs times out."""
+
     pass
 
 
 class SubprocessManager(object):
+    """A manager for subprocesses."""
+
     def __init__(self):
-        self.commands = {}
+        self.commands: Dict[int, CommandManager] = {}
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "SubprocessManager":
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.cleanup()
 
-    async def run_command(self, command: List[str], env=None, cwd=None):
-        command_id = hash_command_invocation(command)
-        self.commands[command_id] = CommandManager(command, env, cwd)
-        await self.commands[command_id].run()
-        return command_id
+    async def run_command(
+        self,
+        command: List[str],
+        env: Optional[Dict[str, str]] = None,
+        cwd: Optional[str] = None,
+    ) -> int:
+        """Run a command asynchronously and return its process ID."""
 
-    def get(self, command_id: str) -> "CommandManager":
-        return self.commands.get(command_id, None)
+        command_obj = CommandManager(command, env, cwd)
+        process = await command_obj.run()
+        self.commands[process.pid] = command_obj
+        return process.pid
 
-    async def cleanup(self):
-        for _, v in self.commands.items():
+    def get(self, pid: int) -> "CommandManager":
+        """Get the CommandManager object for a given process ID."""
+
+        return self.commands.get(pid, None)
+
+    async def cleanup(self) -> None:
+        """Clean up log files for all running subprocesses."""
+
+        for v in self.commands.values():
             await v.cleanup()
 
 
 class CommandManager(object):
-    def __init__(self, command: List[str], env=None, cwd=None):
-        self.command = command
+    """A manager for an individual subprocess."""
 
-        if env is None:
-            env = os.environ.copy()
-        self.env = env
+    def __init__(
+        self,
+        command: List[str],
+        env: Optional[Dict[str, str]] = None,
+        cwd: Optional[str] = None,
+    ):
+        self.command = command
 
-        if cwd is None:
-            cwd = os.getcwd()
-        self.cwd = cwd
+        self.env = env if env is not None else os.environ.copy()
+        self.cwd = cwd if cwd is not None else os.getcwd()
 
         self.process = None
-        self.run_called = False
-        self.log_files = {}
+        self.run_called: bool = False
+        self.log_files: Dict[str, str] = {}
 
         signal.signal(signal.SIGINT, self.handle_sigint)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "CommandManager":
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.cleanup()
 
     def handle_sigint(self, signum, frame):
+        """Handle the SIGINT signal."""
+
         print("SIGINT received.")
         asyncio.create_task(self.kill())
 
-    async def wait(self, timeout=None, stream=None):
+    async def wait(
+        self, timeout: Optional[float] = None, stream: Optional[str] = None
+    ) -> None:
+        """Wait for the subprocess to finish, optionally with a timeout and optionally streaming its output."""
+
         if timeout is None:
             if stream is None:
                 await self.process.wait()
             else:
                 await self.emit_logs(stream)
         else:
-            tasks = [asyncio.create_task(asyncio.sleep(timeout))]
-            if stream is None:
-                tasks.append(asyncio.create_task(self.process.wait()))
-            else:
-                tasks.append(asyncio.create_task(self.emit_logs(stream)))
-
-            await asyncio.wait(tasks, return_when="FIRST_COMPLETED")
+            try:
+                if stream is None:
+                    await asyncio.wait_for(self.process.wait(), timeout)
+                else:
+                    await asyncio.wait_for(self.emit_logs(stream), timeout)
+            except asyncio.TimeoutError:
+                command_string = " ".join(self.command)
+                print(
+                    f"Timeout: The process: '{command_string}' didn't complete within {timeout} seconds."
+                )
 
     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")
@@ -114,14 +132,20 @@ async def run(self):
             await self.cleanup()
 
     async def stream_logs(
-        self, stream, position=None, timeout_per_line=None, log_write_delay=0.01
+        self,
+        stream: str,
+        position: Optional[int] = None,
+        timeout_per_line: Optional[float] = None,
+        log_write_delay: float = 0.01,
     ):
+        """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 stream not in self.log_files:
             raise ValueError(
-                f"No log file found for {stream}, valid values are: {list(self.log_files.keys())}"
+                f"No log file found for '{stream}', valid values are: {list(self.log_files.keys())}"
             )
 
         log_file = self.log_files[stream]
@@ -161,15 +185,21 @@ async def stream_logs(
                 position = f.tell()
                 yield position, line.strip()
 
-    async def emit_logs(self, stream="stdout", custom_logger=print):
+    async def emit_logs(self, stream: str = "stdout", custom_logger: Callable = print):
+        """Helper function to iterate over stream_logs"""
+
         async for _, line in self.stream_logs(stream):
             custom_logger(line)
 
     async def cleanup(self):
+        """Clean up log files for a running subprocesses."""
+
         if hasattr(self, "temp_dir"):
             shutil.rmtree(self.temp_dir, ignore_errors=True)
 
-    async def kill(self, termination_timeout=5):
+    async def kill(self, termination_timeout: float = 5):
+        """Kill the subprocess."""
+
         if self.process is not None:
             if self.process.returncode is None:
                 self.process.terminate()