refactor(backend): Un-share resource initializations from AppService + Remove Pyro (#9750)

This is a prerequisite infra change for
#9714.

We will need a service where we can maintain our own client (db, redis,
rabbitmq, be it async/sync) and configure our own cadence of
initialization and cleanup.

While refactoring the service.py, an option to use Pyro as an RPC
protocol is also removed.

### Changes 🏗️

* Decouple resource initialization and cleanup from the parent
AppService logic.
* Removed Pyro.

### Checklist 📋

#### For code changes:
- [x] I have clearly listed my changes in the PR description
- [x] I have made a test plan
- [x] I have tested my changes according to the test plan:
  <!-- Put your test plan here: -->
  - [x] CI

Author: majdyz

autogpt_platform/backend/backend/executor/database.py

@@ -1,3 +1,6 @@
+import logging
+
+from backend.data import db, redis
 from backend.data.credit import UsageTransactionMetadata, get_user_credit_model
 from backend.data.execution import (
     GraphExecution,
@@ -44,6 +47,7 @@
 
 config = Config()
 _user_credit_model = get_user_credit_model()
+logger = logging.getLogger(__name__)
 
 
 async def _spend_credits(
@@ -55,10 +59,22 @@ async def _spend_credits(
 class DatabaseManager(AppService):
     def __init__(self):
         super().__init__()
-        self.use_db = True
-        self.use_redis = True
         self.execution_event_bus = RedisExecutionEventBus()
 
+    def run_service(self) -> None:
+        logger.info(f"[{self.service_name}] ⏳ Connecting to Database...")
+        self.run_and_wait(db.connect())
+        logger.info(f"[{self.service_name}] ⏳ Connecting to Redis...")
+        redis.connect()
+        super().run_service()
+
+    def cleanup(self):
+        super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Disconnecting Redis...")
+        redis.disconnect()
+        logger.info(f"[{self.service_name}] ⏳ Disconnecting Database...")
+        self.run_and_wait(db.disconnect())
+
     @classmethod
     def get_port(cls) -> int:
         return config.database_api_port

autogpt_platform/backend/backend/executor/manager.py

@@ -930,8 +930,6 @@ def _handle_low_balance_notif(
 class ExecutionManager(AppService):
     def __init__(self):
         super().__init__()
-        self.use_redis = True
-        self.use_supabase = True
         self.pool_size = settings.config.num_graph_workers
         self.queue = ExecutionQueue[GraphExecutionEntry]()
         self.active_graph_runs: dict[str, tuple[Future, threading.Event]] = {}
@@ -944,14 +942,17 @@ def run_service(self):
         from backend.integrations.credentials_store import IntegrationCredentialsStore
 
         self.credentials_store = IntegrationCredentialsStore()
+
+        logger.info(f"[{self.service_name}] ⏳ Spawn max-{self.pool_size} workers...")
         self.executor = ProcessPoolExecutor(
             max_workers=self.pool_size,
             initializer=Executor.on_graph_executor_start,
         )
+
+        logger.info(f"[{self.service_name}] ⏳ Connecting to Redis...")
+        redis.connect()
+
         sync_manager = multiprocessing.Manager()
-        logger.info(
-            f"[{self.service_name}] Started with max-{self.pool_size} graph workers"
-        )
         while True:
             graph_exec_data = self.queue.get()
             graph_exec_id = graph_exec_data.graph_exec_id
@@ -968,10 +969,13 @@ def run_service(self):
             )
 
     def cleanup(self):
-        logger.info(f"[{__class__.__name__}] ⏳ Shutting down graph executor pool...")
+        super().cleanup()
+
+        logger.info(f"[{self.service_name}] ⏳ Shutting down graph executor pool...")
         self.executor.shutdown(cancel_futures=True)
 
-        super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Disconnecting Redis...")
+        redis.disconnect()
 
     @property
     def db_client(self) -> "DatabaseManager":

autogpt_platform/backend/backend/executor/scheduler.py

@@ -206,6 +206,12 @@ def run_service(self):
         self.scheduler.add_listener(job_listener, EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
         self.scheduler.start()
 
+    def cleanup(self):
+        super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Shutting down scheduler...")
+        if self.scheduler:
+            self.scheduler.shutdown(wait=False)
+
     @expose
     def add_execution_schedule(
         self,

autogpt_platform/backend/backend/notifications/notifications.py

@@ -9,6 +9,7 @@
 from prisma.enums import NotificationType
 from pydantic import BaseModel
 
+from backend.data import rabbitmq
 from backend.data.notifications import (
     BaseSummaryData,
     BaseSummaryParams,
@@ -128,6 +129,20 @@ def __init__(self):
         self.running = True
         self.email_sender = EmailSender()
 
+    @property
+    def rabbit(self) -> rabbitmq.AsyncRabbitMQ:
+        """Access the RabbitMQ service. Will raise if not configured."""
+        if not self.rabbitmq_service:
+            raise RuntimeError("RabbitMQ not configured for this service")
+        return self.rabbitmq_service
+
+    @property
+    def rabbit_config(self) -> rabbitmq.RabbitMQConfig:
+        """Access the RabbitMQ config. Will raise if not configured."""
+        if not self.rabbitmq_config:
+            raise RuntimeError("RabbitMQ not configured for this service")
+        return self.rabbitmq_config
+
     @classmethod
     def get_port(cls) -> int:
         return settings.config.notification_service_port
@@ -688,10 +703,14 @@ def _run_queue(
                 )
 
     def run_service(self):
+        logger.info(f"[{self.service_name}] ⏳ Configuring RabbitMQ...")
+        self.rabbitmq_service = rabbitmq.AsyncRabbitMQ(self.rabbitmq_config)
+        self.run_and_wait(self.rabbitmq_service.connect())
+
         logger.info(f"[{self.service_name}] Started notification service")
 
         # Set up scheduler for batch processing of all notification types
-        # this can be changed later to spawn differnt cleanups on different schedules
+        # this can be changed later to spawn different cleanups on different schedules
         try:
             get_scheduler().add_batched_notification_schedule(
                 notification_types=list(NotificationType),
@@ -753,3 +772,5 @@ def cleanup(self):
         """Cleanup service resources"""
         self.running = False
         super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Disconnecting RabbitMQ...")
+        self.run_and_wait(self.rabbitmq_service.disconnect())

autogpt_platform/backend/backend/server/rest_api.py

@@ -145,6 +145,10 @@ def run(self):
             log_config=generate_uvicorn_config(),
         )
 
+    def cleanup(self):
+        super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Shutting down Agent Server...")
+
     @staticmethod
     async def test_execute_graph(
         graph_id: str,

autogpt_platform/backend/backend/server/ws_api.py

@@ -10,7 +10,6 @@
 from fastapi import Depends, FastAPI, WebSocket, WebSocketDisconnect
 from starlette.middleware.cors import CORSMiddleware
 
-from backend.data import redis
 from backend.data.execution import AsyncRedisExecutionEventBus
 from backend.data.user import DEFAULT_USER_ID
 from backend.server.conn_manager import ConnectionManager
@@ -56,15 +55,12 @@ def get_db_client():
 
 async def event_broadcaster(manager: ConnectionManager):
     try:
-        redis.connect()
         event_queue = AsyncRedisExecutionEventBus()
         async for event in event_queue.listen("*"):
             await manager.send_execution_update(event)
     except Exception as e:
         logger.exception(f"Event broadcaster error: {e}")
         raise
-    finally:
-        redis.disconnect()
 
 
 async def authenticate_websocket(websocket: WebSocket) -> str:
@@ -294,3 +290,7 @@ def run(self):
             port=Config().websocket_server_port,
             log_config=generate_uvicorn_config(),
         )
+
+    def cleanup(self):
+        super().cleanup()
+        logger.info(f"[{self.service_name}] ⏳ Shutting down WebSocket Server...")

autogpt_platform/backend/backend/util/process.py

@@ -48,6 +48,7 @@ def run(self):
     def service_name(cls) -> str:
         return cls.__name__
 
+    @abstractmethod
     def cleanup(self):
         """
         Implement this method on a subclass to do post-execution cleanup,