Forward parent process interrupts to child process

Author: ncoghlan

.github/workflows/test.yml

@@ -41,8 +41,10 @@ jobs:
       max-parallel: 8
       matrix:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
-        # There's no platform specific SDK code, but explicitly check Windows
-        # to ensure there aren't any inadvertent POSIX-only assumptions
+        # While the main SDK is platform independent, the subprocess execution
+        # in the plugin runner and tests requires some Windows-specific code
+        # Note: a green tick in CI is currently misleading due to
+        # https://github.com/lmstudio-ai/lmstudio-python/issues/140
         os: [ubuntu-22.04, windows-2022]
 
     # Check https://github.com/actions/action-versions/tree/main/config/actions

src/lmstudio/plugin/_dev_runner.py

@@ -3,13 +3,15 @@
 import asyncio
 import io
 import os
+import signal
 import subprocess
 import sys
 
 from contextlib import asynccontextmanager
-from pathlib import Path
 from functools import partial
-from typing import AsyncGenerator, Iterable, TypeAlias
+from pathlib import Path
+from threading import Event as SyncEvent
+from typing import Any, AsyncGenerator, Iterable, TypeAlias
 
 from typing_extensions import (
     # Native in 3.11+
@@ -116,6 +118,7 @@ async def register_dev_plugin(self) -> AsyncGenerator[tuple[str, str], None]:
     async def _run_plugin_task(
         self, result_queue: asyncio.Queue[int], debug: bool = False
     ) -> None:
+        notify_subprocess_thread = SyncEvent()
         async with self.register_dev_plugin() as (client_id, client_key):
             wait_for_subprocess = asyncio.ensure_future(
                 asyncio.to_thread(
@@ -124,18 +127,20 @@ async def _run_plugin_task(
                         self._plugin_path,
                         client_id,
                         client_key,
-                        debug,
+                        notify_subprocess_thread,
+                        debug=debug,
                     )
                 )
             )
             try:
                 result = await wait_for_subprocess
             except asyncio.CancelledError:
                 # Likely a Ctrl-C press, which is the expected termination process
+                notify_subprocess_thread.set()
                 result_queue.put_nowait(0)
                 raise
             # Subprocess terminated, pass along its return code in the parent process
-            await result_queue.put(result.returncode)
+            await result_queue.put(result)
 
     async def run_plugin(
         self, *, allow_local_imports: bool = True, debug: bool = False
@@ -150,10 +155,49 @@ async def run_plugin(
         return await result_queue.get()
 
 
+def _get_creation_flags() -> int:
+    if sys.platform == "win32":
+        return subprocess.CREATE_NEW_PROCESS_GROUP
+    return 0
+
+
+def _start_child_process(
+    command: list[str], *, text: bool | None = True, **kwds: Any
+) -> subprocess.Popen[str]:
+    creationflags = kwds.pop("creationflags", 0)
+    creationflags |= _get_creation_flags()
+    return subprocess.Popen(command, text=text, creationflags=creationflags, **kwds)
+
+
+def _get_interrupt_signal() -> signal.Signals:
+    if sys.platform == "win32":
+        return signal.CTRL_C_EVENT
+    return signal.SIGINT
+
+
+_PLUGIN_INTERRUPT_SIGNAL = _get_interrupt_signal()
+_PLUGIN_STATUS_POLL_INTERVAL = 1
+_PLUGIN_STOP_TIMEOUT = 2
+
+
+def _interrupt_child_process(process: subprocess.Popen[Any], timeout: float) -> int:
+    process.send_signal(_PLUGIN_INTERRUPT_SIGNAL)
+    try:
+        return process.wait(timeout)
+    except TimeoutError:
+        process.kill()
+        raise
+
+
 # TODO: support the same source code change monitoring features as `lms dev`
 def _run_plugin_in_child_process(
-    plugin_path: Path, client_id: str, client_key: str, debug: bool = False
-) -> subprocess.CompletedProcess[str]:
+    plugin_path: Path,
+    client_id: str,
+    client_key: str,
+    abort_event: SyncEvent,
+    *,
+    debug: bool = False,
+) -> int:
     env = os.environ.copy()
     env[ENV_CLIENT_ID] = client_id
     env[ENV_CLIENT_KEY] = client_key
@@ -176,7 +220,17 @@ def _run_plugin_in_child_process(
         *debug_option,
         os.fspath(plugin_path),
     ]
-    return subprocess.run(command, text=True, env=env)
+    process = _start_child_process(command, env=env)
+    while True:
+        result = process.poll()
+        if result is not None:
+            print("Child process terminated unexpectedly")
+            break
+        if abort_event.wait(_PLUGIN_STATUS_POLL_INTERVAL):
+            print("Gracefully terminating child process...")
+            result = _interrupt_child_process(process, _PLUGIN_STOP_TIMEOUT)
+            break
+    return result
 
 
 async def run_plugin_async(

src/lmstudio/plugin/cli.py

@@ -50,7 +50,7 @@ def main(argv: Sequence[str] | None = None) -> int:
         try:
             runner.run_plugin(plugin_path, allow_local_imports=True)
         except KeyboardInterrupt:
-            print("Plugin execution terminated by user", flush=True)
+            print("Plugin execution terminated by console interrupt", flush=True)
     else:
         # Retrieve args from API host, spawn plugin in subprocess
         try:

tests/test_plugin_examples.py

@@ -1,6 +1,5 @@
 """Test plugin examples can run as dev plugins."""
 
-import signal
 import subprocess
 import sys
 import time
@@ -13,6 +12,12 @@
 
 import pytest
 
+from lmstudio.plugin._dev_runner import (
+    _interrupt_child_process,
+    _start_child_process,
+    _PLUGIN_STOP_TIMEOUT,
+)
+
 _THIS_DIR = Path(__file__).parent.resolve()
 _PLUGIN_EXAMPLES_DIR = (_THIS_DIR / "../examples/plugins").resolve()
 
@@ -47,30 +52,18 @@ def _exec_plugin(plugin_path: Path) -> subprocess.Popen[str]:
         "--dev",
         str(plugin_path),
     ]
-    return subprocess.Popen(
-        cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1
-    )
+    return _start_child_process(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
 
 _PLUGIN_START_TIMEOUT = 5
-_PLUGIN_STOP_TIMEOUT = 5
-
-
-def _get_interrupt_signal() -> signal.Signals:
-    if sys.platform == "win32":
-        return signal.CTRL_C_EVENT
-    return signal.SIGINT
 
 
-_INTERRUPT_SIGNAL = _get_interrupt_signal()
-
 def _exec_and_interrupt(plugin_path: Path) -> tuple[list[str], list[str], list[str]]:
+    # Start the plugin in a child process
     process = _exec_plugin(plugin_path)
     # Ensure pipes don't fill up and block subprocess execution
     stdout_q: Queue[str] = Queue()
-    stdout_thread = Thread(
-        target=_monitor_stream, args=[process.stdout, stdout_q], kwargs={"debug": True}
-    )
+    stdout_thread = Thread(target=_monitor_stream, args=[process.stdout, stdout_q])
     stdout_thread.start()
     stderr_q: Queue[str] = Queue()
     stderr_thread = Thread(target=_monitor_stream, args=[process.stderr, stderr_q])
@@ -94,15 +87,14 @@ def _exec_and_interrupt(plugin_path: Path) -> tuple[list[str], list[str], list[s
     finally:
         # Instruct the process to terminate
         print("Sending termination request to plugin subprocess")
-        process.send_signal(_INTERRUPT_SIGNAL)
+        stop_deadline = time.monotonic() + _PLUGIN_STOP_TIMEOUT
+        _interrupt_child_process(process, (stop_deadline - time.monotonic()))
         # Give threads a chance to halt their file reads
         # (process terminating will close the pipes)
-        stop_deadline = time.monotonic() + _PLUGIN_STOP_TIMEOUT
         stdout_thread.join(timeout=(stop_deadline - time.monotonic()))
         stderr_thread.join(timeout=(stop_deadline - time.monotonic()))
-        process.wait(timeout=(stop_deadline - time.monotonic()))
         with process:
-            # Close pipes
+            # Closes open pipes
             pass
     # Collect remainder of subprocess output
     shutdown_lines = [*_drain_queue(stdout_q)]
@@ -121,9 +113,4 @@ def test_plugin_execution(plugin_path: Path) -> None:
     for log_line in stderr_lines:
         assert log_line.startswith("INFO:")
     assert startup_lines[-1].endswith("Ctrl-C to terminate...\n")
-    # Outside an actual terminal, pipe may be closed before the termination is reported
-    # TODO: Consider migrating to using pexpect, so this check can be more robust
-    assert (
-        not shutdown_lines
-        or shutdown_lines[-1] == "Plugin execution terminated by user\n"
-    )
+    assert shutdown_lines[-1] == "Plugin execution terminated by console interrupt\n"