From 1475b1eefa44d3c14076d29038398c13efb5d67a Mon Sep 17 00:00:00 2001
From: ochafik <ochafik@google.com>
Date: Wed, 10 Apr 2024 08:05:03 +0100
Subject: [PATCH] agent: fix killing of subprocesses

subprocesses again
---
 examples/agent/agent.py         |  7 ++-----
 examples/openai/server.py       |  8 +++++---
 examples/openai/subprocesses.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 examples/openai/subprocesses.py

diff --git a/examples/agent/agent.py b/examples/agent/agent.py
index 8e95a7c40..1cbc254fe 100644
--- a/examples/agent/agent.py
+++ b/examples/agent/agent.py
@@ -1,7 +1,4 @@
-import atexit
-import os
 from pathlib import Path
-import subprocess
 import sys
 from time import sleep
 import typer
@@ -14,6 +11,7 @@ from examples.agent.tools.std_tools import StandardTools
 from examples.openai.api import ChatCompletionRequest, ChatCompletionResponse, Message, ResponseFormat, Tool, ToolFunction
 from examples.agent.utils import collect_functions, load_module
 from examples.openai.prompting import ToolsPromptStyle
+from examples.openai.subprocesses import spawn_subprocess
 
 def _get_params_schema(fn: Callable, verbose):
     if isinstance(fn, OpenAPIMethod):
@@ -185,8 +183,7 @@ def main(
             *([f'--context-length={context_length}'] if context_length else []),
             *([f'--style={style.value}'] if style else []),
         ]
-        server_process = subprocess.Popen(cmd, stdout=sys.stderr)
-        atexit.register(server_process.kill)
+        spawn_subprocess(cmd)
         sleep(5)
 
     tool_functions = []
diff --git a/examples/openai/server.py b/examples/openai/server.py
index e12a1d737..a23f3bb4b 100644
--- a/examples/openai/server.py
+++ b/examples/openai/server.py
@@ -1,7 +1,7 @@
 # https://gist.github.com/ochafik/a3d4a5b9e52390544b205f37fb5a0df3
 # pip install "fastapi[all]" "uvicorn[all]" sse-starlette jsonargparse jinja2 pydantic
 
-import json, sys, subprocess, atexit
+import json, sys
 from pathlib import Path
 import time
 
@@ -22,6 +22,8 @@ from starlette.responses import StreamingResponse
 from typing import Annotated, Optional
 import typer
 
+from examples.openai.subprocesses import spawn_subprocess
+
 def generate_id(prefix):
     return f"{prefix}{random.randint(0, 1 << 32)}"
 
@@ -71,8 +73,8 @@ def main(
             "-c", f"{context_length}",
             *([] if verbose else ["--log-disable"]),
         ]
-        server_process = subprocess.Popen(cmd, stdout=sys.stderr)
-        atexit.register(server_process.kill)
+
+        spawn_subprocess(cmd)
         endpoint = f"http://{server_host}:{server_port}"
 
     app = FastAPI()
diff --git a/examples/openai/subprocesses.py b/examples/openai/subprocesses.py
new file mode 100644
index 000000000..33ee8a507
--- /dev/null
+++ b/examples/openai/subprocesses.py
@@ -0,0 +1,30 @@
+
+import atexit
+import os
+import signal
+import subprocess
+import sys
+
+
+def _cleanup_process(p):
+    pid = p.pid
+
+    if sys.platform == 'win32':
+        os.system(f'taskkill /PID {pid} /T /F')
+    else:
+        pgid = os.getpgid(pid)
+        os.killpg(pgid, signal.SIGTERM)
+
+        p.wait()
+        if p.poll() is None:
+            os.killpg(pgid, signal.SIGKILL)
+
+def spawn_subprocess(cmd, **kwargs):
+    server_process = subprocess.Popen(
+        cmd,
+        stdout=sys.stderr,
+        start_new_session=True,
+        **kwargs
+    )
+    atexit.register(_cleanup_process, server_process)
+    return server_process