[AINode] Add callback and release resources (#16070)

Author: yunbow30944

iotdb-core/ainode/ainode/core/ainode.py

@@ -159,6 +159,7 @@ def _handle_signal(self, signum, frame):
     def stop(self):
         if not self._stop_event.is_set():
             self._stop_event.set()
+            self._rpc_handler.stop()
             if self._rpc_service:
                 self._rpc_service.stop()
                 self._rpc_service.join(1)

iotdb-core/ainode/ainode/core/inference/inference_request_pool.py

@@ -175,18 +175,10 @@ def run(self):
         execute_daemon.start()
         for thread in self._threads:
             thread.join()
+        logger.info(
+            f"[Inference][Device-{self.device}][Pool-{self.pool_id}] InferenceRequestPool exited cleanly."
+        )
 
     def stop(self):
         self._stop_event.set()
-        logger.info(
-            f"[Inference][Device-{self.device}][Pool-{self.pool_id}] Stopping and releasing resources."
-        )
-        try:
-            del self.model
-            if "cuda" in str(self.device):
-                torch.cuda.empty_cache()
-            gc.collect()
-        except Exception as e:
-            logger.warning(
-                f"[Inference][Device-{self.device}][Pool-{self.pool_id}] Failed to clean up: {e}"
-            )
+        logger.debug(f"[Inference][Pool-{self.pool_id}] stop() called")

iotdb-core/ainode/ainode/core/manager/inference_manager.py

@@ -310,6 +310,8 @@ def shutdown(self):
                 while not requestQueue.empty():
                     requestQueue.get_nowait()
                 requestQueue.close()
+            for requestPool, _ in pools:
+                requestPool.join(timeout=10)
         while not self._result_queue.empty():
             self._result_queue.get_nowait()
         self._result_queue.close()

iotdb-core/ainode/ainode/core/rpc/handler.py

@@ -46,6 +46,10 @@ def __init__(self, ainode):
         self._model_manager = ModelManager()
         self._inference_manager = InferenceManager()
 
+    def stop(self) -> None:
+        logger.info("Stopping the RPC service handler of IoTDB-AINode...")
+        self._inference_manager.shutdown()
+
     def stopAINode(self) -> TSStatus:
         self._ainode.stop()
         return get_status(TSStatusCode.SUCCESS_STATUS, "AINode stopped successfully.")