ensure code_interpreter subprocesses are properly terimnated

ye7love7 · Apr 1, 2024 · 8b049e8 · 8b049e8
1 parent 5dc3dee
commit 8b049e8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 25 deletions.
diff --git a/qwen_agent/tools/code_interpreter.py b/qwen_agent/tools/code_interpreter.py
@@ -65,6 +65,7 @@ def _fix_secure_write_for_code_interpreter():
     'AlibabaPuHuiTi-3-45-Light.ttf')
 
 _KERNEL_CLIENTS: Dict[int, BlockingKernelClient] = {}
+_MISC_SUBPROCESSES: Dict[str, subprocess.Popen] = {}
 
 
 def _start_kernel(pid) -> BlockingKernelClient:
@@ -80,15 +81,18 @@ def _start_kernel(pid) -> BlockingKernelClient:
     with open(launch_kernel_script, 'w') as fout:
         fout.write(LAUNCH_KERNEL_PY)
 
-    kernel_process = subprocess.Popen([
-        sys.executable,
-        launch_kernel_script,
-        '--IPKernelApp.connection_file',
-        connection_file,
-        '--matplotlib=inline',
-        '--quiet',
-    ],
-                                      cwd=WORK_DIR)
+    kernel_process = subprocess.Popen(
+        [
+            sys.executable,
+            launch_kernel_script,
+            '--IPKernelApp.connection_file',
+            connection_file,
+            '--matplotlib=inline',
+            '--quiet',
+        ],
+        cwd=WORK_DIR,
+    )
+    _MISC_SUBPROCESSES[f'kc_{kernel_process.pid}'] = kernel_process
     logger.info(f"INFO: kernel process's PID = {kernel_process.pid}")
 
     # Wait for kernel connection file to be written
@@ -113,16 +117,24 @@ def _start_kernel(pid) -> BlockingKernelClient:
     return kc
 
 
-def _kill_kernels():
+def _kill_kernels_and_subprocesses(sig_num=None, _frame=None):
     for v in _KERNEL_CLIENTS.values():
         v.shutdown()
     for k in list(_KERNEL_CLIENTS.keys()):
         del _KERNEL_CLIENTS[k]
 
+    for v in _MISC_SUBPROCESSES.values():
+        v.terminate()
+    for k in list(_MISC_SUBPROCESSES.keys()):
+        del _MISC_SUBPROCESSES[k]
 
-atexit.register(_kill_kernels)
-signal.signal(signal.SIGTERM, _kill_kernels)
-signal.signal(signal.SIGINT, _kill_kernels)
+    if sig_num == signal.SIGINT:
+        raise KeyboardInterrupt()
+
+
+atexit.register(_kill_kernels_and_subprocesses)
+signal.signal(signal.SIGTERM, _kill_kernels_and_subprocesses)
+signal.signal(signal.SIGINT, _kill_kernels_and_subprocesses)
 
 
 def _serve_image(image_base64: str) -> str:
@@ -140,17 +152,16 @@ def _serve_image(image_base64: str) -> str:
     # Hotfix: Temporarily generate image URL proxies for code interpreter to display in gradio
     # Todo: Generate real url
     if static_url == 'http://127.0.0.1:7865/static':
-        try:
-            # run a fastapi server for image show in gradio demo by http://127.0.0.1:7865/figure_name
-            subprocess.Popen([
-                'python',
-                Path(__file__).absolute().parent / 'resource' /
-                'image_service.py'
-            ])
-        except OSError as ex:
-            logger.warning(ex)
-        except Exception:
-            print_traceback()
+        if 'image_service' not in _MISC_SUBPROCESSES:
+            try:
+                # run a fastapi server for image show in gradio demo by http://127.0.0.1:7865/figure_name
+                _MISC_SUBPROCESSES['image_service'] = subprocess.Popen([
+                    'python',
+                    Path(__file__).absolute().parent / 'resource' /
+                    'image_service.py'
+                ])
+            except Exception:
+                print_traceback()
 
     image_url = f'{static_url}/{image_file}'
 

diff --git a/run_server.py b/run_server.py
@@ -109,11 +109,13 @@ def main():
         ]),
     }
 
-    def signal_handler(_sig, _frame):
+    def signal_handler(sig_num, _frame):
         for v in servers.values():
             v.terminate()
         for k in list(servers.keys()):
             del servers[k]
+        if sig_num == signal.SIGINT:
+            raise KeyboardInterrupt()
 
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)