live stream of command output in Shell driver

streams stdout and stderr for long-running methods

Author: michalskrivanek

packages/jumpstarter-driver-shell/jumpstarter_driver_shell/client.py

@@ -1,3 +1,4 @@
+import sys
 from dataclasses import dataclass
 
 import click
@@ -13,8 +14,8 @@ class ShellClient(DriverClient):
     Client interface for Shell driver.
 
     This client dynamically checks that the method is configured
-    on the driver, and if it is, it will call it and get the results
-    in the form of (stdout, stderr, returncode).
+    on the driver, and if it is, it will call it with live streaming output.
+    Output chunks are displayed as they arrive.
     """
 
     def _check_method_exists(self, method):
@@ -26,7 +27,17 @@ def _check_method_exists(self, method):
     ## capture any method calls dynamically
     def __getattr__(self, name):
         self._check_method_exists(name)
-        return lambda *args, **kwargs: tuple(self.call("call_method", name, kwargs, *args))
+        def execute(*args, **kwargs):
+            returncode = 0
+            for stdout, stderr, code in self.streamingcall("call_method", name, kwargs, *args):
+                if stdout:
+                    print(stdout, end='', flush=True)
+                if stderr:
+                    print(stderr, end='', file=sys.stderr, flush=True)
+                if code is not None:
+                    returncode = code
+            return returncode
+        return execute
 
     def cli(self):
         """Create CLI interface for dynamically configured shell methods"""
@@ -61,14 +72,7 @@ def method_command(args, env):
                 else:
                     raise click.BadParameter(f"Invalid --env value '{env_var}'. Use KEY=VALUE.")
 
-            # Call the method
-            stdout, stderr, returncode = self.call("call_method", method_name, env_dict, *args)
-
-            # Display results
-            if stdout:
-                click.echo(stdout, nl=not stdout.endswith("\n"))
-            if stderr:
-                click.echo(stderr, err=True, nl=not stderr.endswith("\n"))
+            returncode = getattr(self, method_name)(*args, **env_dict)
 
             # Exit with the same return code as the shell command
             if returncode != 0:

packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver.py

@@ -1,6 +1,8 @@
+import asyncio
 import os
 import subprocess
 from dataclasses import dataclass, field
+from typing import AsyncGenerator
 
 from jumpstarter.driver import Driver, export
 
@@ -27,41 +29,38 @@ def get_methods(self) -> list[str]:
         return methods
 
     @export
-    def call_method(self, method: str, env, *args):
+    async def call_method(self, method: str, env, *args) -> AsyncGenerator[tuple[str, str, int | None], None]:
+        """
+        Execute a shell method with live streaming output.
+        Yields (stdout_chunk, stderr_chunk, returncode) tuples.
+        returncode is None until the process completes, then it's the final return code.
+        """
         self.logger.info(f"calling {method} with args: {args} and kwargs as env: {env}")
         if method not in self.methods:
             raise ValueError(f"Method '{method}' not found in available methods: {list(self.methods.keys())}")
         script = self.methods[method]
         self.logger.debug(f"running script: {script}")
+
         try:
-            result = self._run_inline_shell_script(method, script, *args, env_vars=env)
-            if result.returncode != 0:
-                self.logger.info(f"{method} return code: {result.returncode}")
-            if result.stderr != "":
-                stderr = result.stderr.rstrip("\n")
-                self.logger.debug(f"{method} stderr:\n{stderr}")
-            if result.stdout != "":
-                stdout = result.stdout.rstrip("\n")
-                self.logger.debug(f"{method} stdout:\n{stdout}")
-            return result.stdout, result.stderr, result.returncode
+            async for stdout_chunk, stderr_chunk, returncode in self._run_inline_shell_script(
+                method, script, *args, env_vars=env
+            ):
+                if stdout_chunk:
+                    self.logger.debug(f"{method} stdout:\n{stdout_chunk.rstrip()}")
+                if stderr_chunk:
+                    self.logger.debug(f"{method} stderr:\n{stderr_chunk.rstrip()}")
+
+                if returncode is not None and returncode != 0:
+                    self.logger.info(f"{method} return code: {returncode}")
+
+                yield stdout_chunk, stderr_chunk, returncode
         except subprocess.TimeoutExpired as e:
             self.logger.error(f"Timeout expired while running {method}: {e}")
-            return "", f"Timeout expired while running {method}: {e}", 199
-
-    def _run_inline_shell_script(self, method, script, *args, env_vars=None):
-        """
-        Run the given shell script (as a string) with optional arguments and
-        environment variables. Returns a CompletedProcess with stdout, stderr, and returncode.
-
-        :param script:      The shell script contents as a string.
-        :param args:        Arguments to pass to the script (mapped to $1, $2, etc. in the script).
-        :param env_vars:    A dict of environment variables to make available to the script.
-
-        :return:            A subprocess.CompletedProcess object (Python 3.5+).
-        """
+            yield "", f"\nTimeout expired while running {method}: {e}\n", 199
 
+    def _validate_script_params(self, script, args, env_vars):
+        """Validate script parameters and return combined environment."""
         # Merge parent environment with the user-supplied env_vars
-        # so that we don't lose existing environment variables.
         combined_env = os.environ.copy()
         if env_vars:
             # Validate environment variable names
@@ -82,16 +81,94 @@ def _run_inline_shell_script(self, method, script, *args, env_vars=None):
         if self.cwd and not os.path.isdir(self.cwd):
             raise ValueError(f"Working directory does not exist: {self.cwd}")
 
+        return combined_env
+
+    async def _read_process_output(self, process, read_all=False):
+        """Read data from stdout and stderr streams.
+
+        :param process: The subprocess to read from
+        :param read_all: If True, read all remaining data. If False, read with timeout.
+        :return: Tuple of (stdout_data, stderr_data)
+        """
+        stdout_data = ""
+        stderr_data = ""
+
+        # Read from stdout
+        if process.stdout:
+            try:
+                if read_all:
+                    chunk = await process.stdout.read()
+                else:
+                    chunk = await asyncio.wait_for(process.stdout.read(1024), timeout=0.01)
+                if chunk:
+                    stdout_data = chunk.decode('utf-8', errors='replace')
+            except (asyncio.TimeoutError, Exception):
+                pass
+
+        # Read from stderr
+        if process.stderr:
+            try:
+                if read_all:
+                    chunk = await process.stderr.read()
+                else:
+                    chunk = await asyncio.wait_for(process.stderr.read(1024), timeout=0.01)
+                if chunk:
+                    stderr_data = chunk.decode('utf-8', errors='replace')
+            except (asyncio.TimeoutError, Exception):
+                pass
+
+        return stdout_data, stderr_data
+
+    async def _run_inline_shell_script(
+        self, method, script, *args, env_vars=None
+    ) -> AsyncGenerator[tuple[str, str, int | None], None]:
+        """
+        Run the given shell script with live streaming output.
+
+        :param method:      The method name (for logging).
+        :param script:      The shell script contents as a string.
+        :param args:        Arguments to pass to the script (mapped to $1, $2, etc. in the script).
+        :param env_vars:    A dict of environment variables to make available to the script.
+
+        :yields:            Tuples of (stdout_chunk, stderr_chunk, returncode).
+                           returncode is None until the process completes.
+        """
+        combined_env = self._validate_script_params(script, args, env_vars)
         cmd = self.shell + [script, method] + list(args)
 
-        # Run the command
-        result = subprocess.run(
-            cmd,
-            capture_output=True,  # Captures stdout and stderr
-            text=True,  # Returns stdout/stderr as strings (not bytes)
-            env=combined_env,  # Pass our merged environment
-            cwd=self.cwd,  # Run in the working directory (if set)
-            timeout=self.timeout,
+        # Start the process with pipes for streaming
+        process = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=combined_env,
+            cwd=self.cwd,
         )
 
-        return result
+        # Create a task to monitor the process timeout
+        start_time = asyncio.get_event_loop().time()
+
+        # Read output in real-time
+        while process.returncode is None:
+            if asyncio.get_event_loop().time() - start_time > self.timeout:
+                process.kill()
+                await process.wait()
+                raise subprocess.TimeoutExpired(cmd, self.timeout) from None
+
+            try:
+                stdout_data, stderr_data = await self._read_process_output(process, read_all=False)
+
+                # Yield any data we got
+                if stdout_data or stderr_data:
+                    yield stdout_data, stderr_data, None
+
+                # Small delay to prevent busy waiting
+                await asyncio.sleep(0.1)
+
+            except Exception:
+                break
+
+        # Process completed, get return code and final output
+        returncode = process.returncode
+        remaining_stdout, remaining_stderr = await self._read_process_output(process, read_all=True)
+        yield remaining_stdout, remaining_stderr, returncode

packages/jumpstarter-driver-shell/jumpstarter_driver_shell/driver_test.py

@@ -1,9 +1,28 @@
+
 import pytest
 
 from .driver import Shell
 from jumpstarter.common.utils import serve
 
 
+def _collect_streaming_output(client, method_name, env_vars=None, *args):
+    """Helper function to collect streaming output for testing"""
+    stdout_parts = []
+    stderr_parts = []
+    final_returncode = None
+
+    env_vars = env_vars or {}
+    for stdout_chunk, stderr_chunk, returncode in client.streamingcall("call_method", method_name, env_vars, *args):
+        if stdout_chunk:
+            stdout_parts.append(stdout_chunk)
+        if stderr_chunk:
+            stderr_parts.append(stderr_chunk)
+        if returncode is not None:
+            final_returncode = returncode
+
+    return "".join(stdout_parts), "".join(stderr_parts), final_returncode
+
+
 @pytest.fixture
 def client():
     instance = Shell(
@@ -21,23 +40,38 @@ def client():
 
 
 def test_normal_args(client):
-    assert client.echo("hello") == ("hello\n", "", 0)
+    stdout, stderr, returncode = _collect_streaming_output(client, "echo", {}, "hello")
+    assert stdout == "hello\n"
+    assert stderr == ""
+    assert returncode == 0
 
 
 def test_env_vars(client):
-    assert client.env(ENV1="world") == ("world\n", "", 0)
+    stdout, stderr, returncode = _collect_streaming_output(client, "env", {"ENV1": "world"})
+    assert stdout == "world\n"
+    assert stderr == ""
+    assert returncode == 0
 
 
 def test_multi_line_scripts(client):
-    assert client.multi_line("a", "b", "c") == ("a\nb\nc\n", "", 0)
+    stdout, stderr, returncode = _collect_streaming_output(client, "multi_line", {}, "a", "b", "c")
+    assert stdout == "a\nb\nc\n"
+    assert stderr == ""
+    assert returncode == 0
 
 
 def test_return_codes(client):
-    assert client.exit1() == ("", "", 1)
+    stdout, stderr, returncode = _collect_streaming_output(client, "exit1")
+    assert stdout == ""
+    assert stderr == ""
+    assert returncode == 1
 
 
 def test_stderr(client):
-    assert client.stderr("error") == ("", "error\n", 0)
+    stdout, stderr, returncode = _collect_streaming_output(client, "stderr", {}, "error")
+    assert stdout == ""
+    assert stderr == "error\n"
+    assert returncode == 0
 
 
 def test_unknown_method(client):