agent: fix killing of subprocesses

subprocesses again
This commit is contained in:
ochafik 2024-04-10 08:05:03 +01:00
parent 24e34f174b
commit 1475b1eefa
3 changed files with 37 additions and 8 deletions

View file

@ -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 = []

View file

@ -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()

View file

@ -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