This is an automated email from the ASF dual-hosted git repository.
yongzao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 593520ab1f0 [AINode] Add callback and release resources (#16070)
593520ab1f0 is described below
commit 593520ab1f034f0a73e446ba68a5a6d62fd244b0
Author: Leo <[email protected]>
AuthorDate: Thu Jul 31 19:49:41 2025 +0800
[AINode] Add callback and release resources (#16070)
---
iotdb-core/ainode/ainode/core/ainode.py | 1 +
.../ainode/core/inference/inference_request_pool.py | 16 ++++------------
.../ainode/ainode/core/manager/inference_manager.py | 2 ++
iotdb-core/ainode/ainode/core/rpc/handler.py | 4 ++++
4 files changed, 11 insertions(+), 12 deletions(-)
diff --git a/iotdb-core/ainode/ainode/core/ainode.py
b/iotdb-core/ainode/ainode/core/ainode.py
index 82c6f6988e2..3ec1f099131 100644
--- a/iotdb-core/ainode/ainode/core/ainode.py
+++ b/iotdb-core/ainode/ainode/core/ainode.py
@@ -159,6 +159,7 @@ class AINode:
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)
diff --git a/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py
b/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py
index b6c0edb7f60..e7e84cca7ff 100644
--- a/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py
+++ b/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py
@@ -175,18 +175,10 @@ class InferenceRequestPool(mp.Process):
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")
diff --git a/iotdb-core/ainode/ainode/core/manager/inference_manager.py
b/iotdb-core/ainode/ainode/core/manager/inference_manager.py
index 5a853ac4e72..3cac5315b6f 100644
--- a/iotdb-core/ainode/ainode/core/manager/inference_manager.py
+++ b/iotdb-core/ainode/ainode/core/manager/inference_manager.py
@@ -310,6 +310,8 @@ class InferenceManager:
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()
diff --git a/iotdb-core/ainode/ainode/core/rpc/handler.py
b/iotdb-core/ainode/ainode/core/rpc/handler.py
index d3948020ab3..ce7ce26abfc 100644
--- a/iotdb-core/ainode/ainode/core/rpc/handler.py
+++ b/iotdb-core/ainode/ainode/core/rpc/handler.py
@@ -46,6 +46,10 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
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.")
