Repository: Significant-Gravitas/AutoGPT
Pull request: feat(server): Clean up resources when spinning down services/processes
- Resolves #7937
- Blocked by #7892

### Changes 🏗️

- Add SIGTERM handler and `cleanup()` hook to `AppProcess`
- Implement `cleanup()` on `AppService` to close DB and Redis connections
- Implement `cleanup()` on `ExecutionManager` to shut down worker pool
- Add `atexit` and SIGTERM handlers to node executor to close DB connection and shut down node workers
- Improve logging in `.executor.manager`
- Fix shutdown order of `.util.test:SpinTestServer`

### PR Quality Scorecard ✨

<!--
Check out our contribution guide:
https://github.com/Significant-Gravitas/AutoGPT/wiki/Contributing

1. Avoid duplicate work, issues, PRs etc.
2. Also consider contributing something other than code; see the [contribution guide]
   for options.
3. Clearly explain your changes.
4. Avoid making unnecessary changes, especially if they're purely based on personal
   preferences. Doing so is the maintainers' job. ;-)
-->

- [x] Have you used the PR description template? &ensp; `+2 pts`
- [x] Is your pull request atomic, focusing on a single change? &ensp; `+5 pts`
- [x] Have you linked the GitHub issue(s) that this PR addresses? &ensp; `+5 pts`
- [x] Have you documented your changes clearly and comprehensively? &ensp; `+5 pts`
- [ ] Have you changed or added a feature? &ensp; `-4 pts`
  - [ ] Have you added/updated corresponding documentation? &ensp; `+4 pts`
  - [ ] Have you added/updated corresponding integration tests? &ensp; `+5 pts`
- [ ] Have you changed the behavior of AutoGPT? &ensp; `-5 pts`
  - [ ] Have you also run `agbenchmark` to verify that these changes do not regress performance? &ensp; `+10 pts`


Problem statement:
[Server] Clean up resources when spinning down services/processes
* `AppService` instances don't close their DB connection when they are stopped, leaving dozens of zombie Prisma processes running in the background after repeatedly starting/stopping the application.
* `ExecutionManager` doesn't clean up its graph execution workers when it is stopped, leaving dozens of RAM-hungry zombie worker processes (haha the irony) behind.

Hints from issue discussion:

Commit message: Merge branch 'master' into reinier/open-1772-server-clean-up-resources-when-spinning-down

Patch to review:
diff --git a/rnd/autogpt_server/autogpt_server/executor/manager.py b/rnd/autogpt_server/autogpt_server/executor/manager.py
index 3b6d817a2894..3c44825615cc 100644
--- a/rnd/autogpt_server/autogpt_server/executor/manager.py
+++ b/rnd/autogpt_server/autogpt_server/executor/manager.py
@@ -1,5 +1,7 @@
 import asyncio
 import logging
+import signal
+import sys
 from concurrent.futures import Future, ProcessPoolExecutor, TimeoutError
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Coroutine, Generator, TypeVar
@@ -336,6 +338,14 @@ def on_node_executor_start(cls):
         cls.loop.run_until_complete(db.connect())
         cls.agent_server_client = get_agent_server_client()
 
+        # Set up shutdown handler
+        signal.signal(signal.SIGTERM, lambda _, __: cls.node_executor_stop())
+
+    @classmethod
+    def node_executor_stop(cls):
+        cls.loop.run_until_complete(db.disconnect())
+        sys.exit(0)
+
     @classmethod
     def on_node_execution(cls, q: ExecutionQueue[NodeExecution], data: NodeExecution):
         prefix = get_log_prefix(data.graph_exec_id, data.node_exec_id)
@@ -415,12 +425,18 @@ def run_service(self):
         with ProcessPoolExecutor(
             max_workers=self.pool_size,
             initializer=Executor.on_graph_executor_start,
-        ) as executor:
-            logger.warning(
+        ) as self.executor:
+            logger.info(
                 f"Execution manager started with max-{self.pool_size} graph workers."
             )
             while True:
-                executor.submit(Executor.on_graph_execution, self.queue.get())
+                self.executor.submit(Executor.on_graph_execution, self.queue.get())
+
+    def cleanup(self):
+        logger.info("⏳ Shutting down graph executor pool...")
+        self.executor.shutdown()
+
+        super().cleanup()
 
     @property
     def agent_server_client(self) -> "AgentServer":
diff --git a/rnd/autogpt_server/autogpt_server/util/process.py b/rnd/autogpt_server/autogpt_server/util/process.py
index 995ce8a48dd9..b734b4c601d7 100644
--- a/rnd/autogpt_server/autogpt_server/util/process.py
+++ b/rnd/autogpt_server/autogpt_server/util/process.py
@@ -1,4 +1,5 @@
 import os
+import signal
 import sys
 from abc import ABC, abstractmethod
 from multiprocessing import Process, set_start_method
@@ -13,6 +14,8 @@ class AppProcess(ABC):
     process: Optional[Process] = None
     set_start_method("spawn", force=True)
 
+    # Methods that are executed INSIDE the process #
+
     @abstractmethod
     def run(self):
         """
@@ -20,7 +23,16 @@ def run(self):
         """
         pass
 
+    def cleanup(self):
+        """
+        Implement this method on a subclass to do post-execution cleanup,
+        e.g. disconnecting from a database or terminating child processes.
+        """
+        pass
+
     def execute_run_command(self, silent):
+        signal.signal(signal.SIGTERM, self._self_terminate)
+
         try:
             if silent:
                 sys.stdout = open(os.devnull, "w")
@@ -29,6 +41,12 @@ def execute_run_command(self, silent):
         except KeyboardInterrupt or SystemExit as e:
             print(f"Process terminated: {e}")
 
+    def _self_terminate(self, signum: int, frame):
+        self.cleanup()
+        sys.exit(0)
+
+    # Methods that are executed OUTSIDE the process #
+
     def __enter__(self):
         self.start(background=True)
         return self
diff --git a/rnd/autogpt_server/autogpt_server/util/service.py b/rnd/autogpt_server/autogpt_server/util/service.py
index e048c27f2167..e78553eeea34 100644
--- a/rnd/autogpt_server/autogpt_server/util/service.py
+++ b/rnd/autogpt_server/autogpt_server/util/service.py
@@ -87,6 +87,11 @@ def run(self):
         # Run the main service (if it's not implemented, just sleep).
         self.run_service()
 
+    def cleanup(self):
+        if self.use_db:
+            logger.info("⏳ Disconnecting DB")
+            self.run_and_wait(db.disconnect())
+
     @conn_retry
     def __start_pyro(self):
         daemon = pyro.Daemon(host=pyro_host)
@@ -107,7 +112,6 @@ def get_service_client(service_type: Type[AS]) -> AS:
     service_name = service_type.service_name
 
     class DynamicClient:
-
         @conn_retry
         def __init__(self):
             ns = pyro.locate_ns()


Code review comments:
diff_hunk: @@ -336,6 +338,14 @@ def on_node_executor_start(cls):
         cls.loop.run_until_complete(db.connect())
         cls.agent_server_client = get_agent_server_client()
 
+        # Set up shutdown handler
+        signal.signal(signal.SIGTERM, lambda _, __: cls.node_executor_stop())
+
+    @classmethod
+    def node_executor_stop(cls):
path: rnd/autogpt_server/autogpt_server/executor/manager.py
text: @user1:
nit: can we keep the same naming convention (event), e.g: `on_node_executor_stop` ?

@user1:
Is the exit that triggers the sigterm?

@author:
* Changed to cleanup triggered by `atexit`
* Renamed cleanup function to `on_(node|graph)_executor_stop`

@author:
I had to add back `stop_node_executor` to handle SIGTERM. A SIGTERM handler must call `sys.exit(0)`, while an `atexit` handler shouldn't.