Support Training Module of AINode

Author: ycycse

iotdb-client/client-py/resources/pyproject.toml

@@ -42,7 +42,7 @@ dependencies = [
     "thrift>=0.14.1",
     "pandas>=1.0.0",
     "numpy>=1.0.0",
-    "sqlalchemy<1.5,>=1.4",
+    "sqlalchemy>=1.4",
     "sqlalchemy-utils>=0.37.8"
 ]
 

iotdb-core/ainode/.gitignore

@@ -1,8 +1,8 @@
 # generated by Thrift
-/iotdb/thrift/
+/ainode/thrift/
 
 # generated by maven
-/iotdb/conf/
+/ainode/conf/
 
 # .whl of ainode, generated by Poetry
 /dist/

iotdb-core/ainode/iotdb/__init__.py iotdb-core/ainode/ainode/__init__.pyiotdb-core/ainode/iotdb/__init__.py renamed to iotdb-core/ainode/ainode/__init__.py

@@ -21,15 +21,16 @@
 from thrift.protocol import TCompactProtocol, TBinaryProtocol
 from thrift.transport import TSocket, TTransport
 
-from iotdb.ainode.config import AINodeDescriptor
-from iotdb.ainode.constant import TSStatusCode
-from iotdb.ainode.log import Logger
-from iotdb.ainode.util.decorator import singleton
-from iotdb.ainode.util.status import verify_success
-from iotdb.thrift.common.ttypes import TEndPoint, TSStatus, TAINodeLocation, TAINodeConfiguration
-from iotdb.thrift.confignode import IConfigNodeRPCService
-from iotdb.thrift.confignode.ttypes import (TAINodeRemoveReq, TNodeVersionInfo,
-                                            TAINodeRegisterReq, TAINodeRestartReq)
+from ainode.core.config import AINodeDescriptor
+from ainode.core.constant import TSStatusCode
+from ainode.core.log import Logger
+from ainode.core.util.decorator import singleton
+from ainode.core.util.status import verify_success
+from ainode.thrift.common.ttypes import TEndPoint, TSStatus, TAINodeLocation, TAINodeConfiguration
+from ainode.thrift.confignode import IConfigNodeRPCService
+from ainode.thrift.confignode.ttypes import (TAINodeRemoveReq, TNodeVersionInfo,
+                                             TAINodeRegisterReq, TAINodeRestartReq)
+from ainode.thrift.confignode.ttypes import TUpdateModelInfoReq
 
 logger = Logger()
 
@@ -202,3 +203,26 @@ def get_ainode_configuration(self, node_id: int) -> map:
                 self._config_leader = None
             self._wait_and_reconnect()
         raise TException(self._MSG_RECONNECTION_FAIL)
+
+    def update_model_info(self, model_id:str, model_status:int, attribute:str = "", ainode_id=None, input_length=0, output_length=0) -> None:
+        if ainode_id is None:
+            ainode_id = []
+        for _ in range(0, self._RETRY_NUM):
+            try:
+                req = TUpdateModelInfoReq(
+                    model_id, model_status, attribute
+                )
+                if ainode_id is not None:
+                    req.aiNodeIds = ainode_id
+                req.inputLength = input_length
+                req.outputLength = output_length
+                status = self._client.updateModelInfo(req)
+                if not self._update_config_node_leader(status):
+                    verify_success(status, "An error occurs when calling update model info")
+                    return status
+            except TTransport.TException:
+                logger.warning("Failed to connect to ConfigNode {} from AINode when executing update model info",
+                               self._config_leader)
+                self._config_leader = None
+            self._wait_and_reconnect()
+        raise TException(self._MSG_RECONNECTION_FAIL)

…tdb-core/ainode/iotdb/ainode/__init__.py iotdb-core/ainode/ainode/core/__init__.pyiotdb-core/ainode/iotdb/ainode/__init__.py renamed to iotdb-core/ainode/ainode/core/__init__.py iotdb-core/ainode/iotdb/ainode/__init__.py renamed to iotdb-core/ainode/ainode/core/__init__.py

@@ -17,17 +17,17 @@
 #
 import os
 
-from iotdb.ainode.constant import (AINODE_CONF_DIRECTORY_NAME,
-                                   AINODE_CONF_FILE_NAME,
-                                   AINODE_MODELS_DIR, AINODE_LOG_DIR, AINODE_SYSTEM_DIR, AINODE_INFERENCE_RPC_ADDRESS,
-                                   AINODE_INFERENCE_RPC_PORT, AINODE_THRIFT_COMPRESSION_ENABLED,
-                                   AINODE_SYSTEM_FILE_NAME, AINODE_CLUSTER_NAME, AINODE_VERSION_INFO, AINODE_BUILD_INFO,
-                                   AINODE_CONF_GIT_FILE_NAME, AINODE_CONF_POM_FILE_NAME, AINODE_ROOT_DIR,
-                                   AINODE_ROOT_CONF_DIRECTORY_NAME)
-from iotdb.ainode.exception import BadNodeUrlError
-from iotdb.ainode.log import Logger
-from iotdb.ainode.util.decorator import singleton
-from iotdb.thrift.common.ttypes import TEndPoint
+from ainode.core.constant import (AINODE_CONF_DIRECTORY_NAME,
+                                  AINODE_CONF_FILE_NAME,
+                                  AINODE_MODELS_DIR, AINODE_LOG_DIR, AINODE_SYSTEM_DIR, AINODE_INFERENCE_RPC_ADDRESS,
+                                  AINODE_INFERENCE_RPC_PORT, AINODE_THRIFT_COMPRESSION_ENABLED,
+                                  AINODE_SYSTEM_FILE_NAME, AINODE_CLUSTER_NAME, AINODE_VERSION_INFO, AINODE_BUILD_INFO,
+                                  AINODE_CONF_GIT_FILE_NAME, AINODE_CONF_POM_FILE_NAME, AINODE_ROOT_DIR,
+                                  AINODE_ROOT_CONF_DIRECTORY_NAME)
+from ainode.core.exception import BadNodeUrlError
+from ainode.core.log import Logger
+from ainode.core.util.decorator import singleton
+from ainode.thrift.common.ttypes import TEndPoint
 
 logger = Logger()
 

iotdb-core/ainode/iotdb/ainode/client.py iotdb-core/ainode/ainode/core/client.pyiotdb-core/ainode/iotdb/ainode/client.py renamed to iotdb-core/ainode/ainode/core/client.py

@@ -17,7 +17,7 @@
 #
 import re
 
-from iotdb.ainode.constant import DEFAULT_MODEL_FILE_NAME, DEFAULT_CONFIG_FILE_NAME
+from ainode.core.constant import DEFAULT_MODEL_FILE_NAME, DEFAULT_CONFIG_FILE_NAME
 
 
 class _BaseError(Exception):

iotdb-core/ainode/iotdb/ainode/config.py iotdb-core/ainode/ainode/core/config.pyiotdb-core/ainode/iotdb/ainode/config.py renamed to iotdb-core/ainode/ainode/core/config.py

@@ -16,14 +16,14 @@
 # under the License.
 #
 
-from iotdb.ainode.manager.cluster_manager import ClusterManager
-from iotdb.ainode.manager.inference_manager import InferenceManager
-from iotdb.ainode.manager.model_manager import ModelManager
-from iotdb.thrift.ainode import IAINodeRPCService
-from iotdb.thrift.ainode.ttypes import (TDeleteModelReq, TRegisterModelReq,
-                                        TAIHeartbeatReq, TInferenceReq, TRegisterModelResp, TInferenceResp,
-                                        TAIHeartbeatResp)
-from iotdb.thrift.common.ttypes import TSStatus
+from ainode.core.manager.cluster_manager import ClusterManager
+from ainode.core.manager.inference_manager import InferenceManager
+from ainode.core.manager.model_manager import ModelManager
+from ainode.thrift.ainode import IAINodeRPCService
+from ainode.thrift.ainode.ttypes import (TDeleteModelReq, TRegisterModelReq,
+                                         TAIHeartbeatReq, TInferenceReq, TRegisterModelResp, TInferenceResp,
+                                         TAIHeartbeatResp, TTrainingReq)
+from ainode.thrift.common.ttypes import TSStatus
 
 
 class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
@@ -41,3 +41,6 @@ def inference(self, req: TInferenceReq) -> TInferenceResp:
 
     def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp:
         return ClusterManager.get_heart_beat(req)
+
+    def createTrainingTask(self, req: TTrainingReq) -> TSStatus:
+        pass

…tdb-core/ainode/iotdb/ainode/constant.py iotdb-core/ainode/ainode/core/constant.pyiotdb-core/ainode/iotdb/ainode/constant.py renamed to iotdb-core/ainode/ainode/core/constant.py iotdb-core/ainode/iotdb/ainode/constant.py renamed to iotdb-core/ainode/ainode/core/constant.py

@@ -23,8 +23,8 @@
 import sys
 import threading
 
-from iotdb.ainode.constant import STD_LEVEL, AINODE_LOG_FILE_NAMES, AINODE_LOG_FILE_LEVELS
-from iotdb.ainode.util.decorator import singleton
+from ainode.core.constant import STD_LEVEL, AINODE_LOG_FILE_NAMES, AINODE_LOG_FILE_LEVELS
+from ainode.core.util.decorator import singleton
 
 
 class LoggerFilter(logging.Filter):