diff --git a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py index 587ae4930dede..d8f619e1b8d38 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py +++ b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py @@ -163,7 +163,11 @@ def stop(self): self._rpc_handler.stop() if self._rpc_service: self._rpc_service.stop() - self._rpc_service.join(1) - if self._rpc_service.is_alive(): - logger.warning("RPC service thread failed to stop in time.") + for retry in range(30): + self._rpc_service.join(2) + if not self._rpc_service.is_alive(): + logger.warning( + "RPC service thread failed to stop in time, retrying..." + ) + break logger.info("IoTDB-AINode has successfully stopped.") diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py index 00bb3b5568b0f..54580402ec293 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py @@ -504,7 +504,7 @@ def get_load(self, model_id, device_id, pool_id) -> int: return pool_group.get_load(pool_id) return -1 - def shutdown(self): + def stop(self): self._stop_event.set() # shutdown pool controller diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py index 841159d9b4c92..a67d576b0ec8c 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py @@ -366,9 +366,9 @@ def inference(self, req: TInferenceReq): single_output=False, ) - def shutdown(self): + def stop(self): self._stop_event.set() - self._pool_controller.shutdown() + self._pool_controller.stop() while not self._result_queue.empty(): self._result_queue.get_nowait() self._result_queue.close() diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py index 04ac139e7d481..f01e1594f0698 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py @@ -69,7 +69,7 @@ def __init__(self, ainode): def stop(self) -> None: logger.info("Stopping the RPC service handler of IoTDB-AINode...") - self._inference_manager.shutdown() + self._inference_manager.stop() def stopAINode(self) -> TSStatus: self._ainode.stop() diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py index c556ffd4ba0d5..fc93a0cce7d5b 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py @@ -125,3 +125,4 @@ def stop(self) -> None: logger.info("Stopping the RPC service of IoTDB-AINode...") self._stop_event.set() self.__pool_server.stop() + self._handler.stop() diff --git a/scripts/sbin/start-ainode.sh b/scripts/sbin/start-ainode.sh index 671de54ba1a86..454a0389c3866 100644 --- a/scripts/sbin/start-ainode.sh +++ b/scripts/sbin/start-ainode.sh @@ -27,10 +27,14 @@ export IOTDB_AINODE_HOME echo "IOTDB_AINODE_HOME: $IOTDB_AINODE_HOME" # fetch parameters with names -while getopts "i:rn" opt; do +daemon_mode=false +while getopts "i:rnd" opt; do case $opt in n) ;; + d) + daemon_mode=true + ;; \?) echo "Invalid option -$OPTARG" >&2 exit 1 ;; @@ -41,6 +45,11 @@ ain_ainode_executable="$IOTDB_AINODE_HOME/lib/ainode" echo Script got ainode executable: "$ain_ainode_executable" -echo Starting AINode... - -$ain_ainode_executable start +if [ "$daemon_mode" = true ]; then + echo Starting AINode in daemon mode... + nohup $ain_ainode_executable start > /dev/null 2>&1 & + echo AINode started in background +else + echo Starting AINode... + $ain_ainode_executable start +fi diff --git a/scripts/sbin/windows/start-ainode.bat b/scripts/sbin/windows/start-ainode.bat index 1d9a4306bafb7..09ce914f22672 100644 --- a/scripts/sbin/windows/start-ainode.bat +++ b/scripts/sbin/windows/start-ainode.bat @@ -30,8 +30,20 @@ set ain_ainode_executable=%IOTDB_AINODE_HOME%\lib\ainode echo Script got ainode executable: %ain_ainode_executable% -echo Starting AINode... - -%$ain_ainode_executable% start - -pause \ No newline at end of file +set daemon_mode=false +:parse_args +if "%~1"=="" goto end_parse +if /i "%~1"=="-d" set daemon_mode=true +shift +goto parse_args +:end_parse + +if "%daemon_mode%"=="true" ( + echo Starting AINode in daemon mode... + start /B "" %ain_ainode_executable% start + echo AINode started in background +) else ( + echo Starting AINode... + %ain_ainode_executable% start + pause +) \ No newline at end of file