[AINode] Fix cp errors

Author: CRZbulabula

.github/workflows/cluster-it-1c1d1a.yml

@@ -41,9 +41,6 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Build AINode
-        shell: bash
-        run: mvn clean package -DskipTests -P with-ainode
       - name: IT Test
         shell: bash
         run: |

iotdb-core/ainode/.flake8

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+[flake8]
+ignore =
+    E203,
+    W503
+max-line-length=200
+exclude =
+    .git,
+    tests/*,
+    iotdb/thrift/**/*
+extend-exclude =
+    dist,
+    build,
+    venv
+show-source = True
+statistics = True
+format = %(path)s:%(row)d,%(col)d:%(code)s:%(text)s:https://lintlyci.github.io/Flake8Rules/rules/%(code)s.html

iotdb-core/ainode/build_binary.py

@@ -31,7 +31,7 @@
 def get_venv_base_dir():
     """
     Get the base directory for virtual environments outside the project.
-    
+
     Returns:
         Path: Base directory path
         - Linux/macOS: ~/.cache/iotdb-ainode-build/
@@ -51,13 +51,13 @@ def get_venv_base_dir():
 def setup_venv():
     """
     Create virtual environment outside the project directory.
-    
+
     The virtual environment is created in a platform-specific location:
     - Linux/macOS: ~/.cache/iotdb-ainode-build/<project-name>/
     - Windows: %LOCALAPPDATA%\\iotdb-ainode-build\\<project-name>\\
-    
+
     The same venv is reused across multiple builds of the same project.
-    
+
     Returns:
         Path: Path to the virtual environment directory
     """

iotdb-core/ainode/iotdb/Session.py

@@ -18,46 +18,48 @@
 
 import logging
 import random
-import sys
 import struct
+import sys
 import warnings
+
 from thrift.protocol import TBinaryProtocol, TCompactProtocol
 from thrift.transport import TSocket, TTransport
 from tzlocal import get_localzone_name
 
 from iotdb.utils.SessionDataSet import SessionDataSet
+
 from .template.Template import Template
 from .template.TemplateQueryType import TemplateQueryType
 from .thrift.common.ttypes import TEndPoint
 from .thrift.rpc.IClientRPCService import (
     Client,
-    TSCreateTimeseriesReq,
+    TSAppendSchemaTemplateReq,
+    TSCloseSessionReq,
     TSCreateAlignedTimeseriesReq,
+    TSCreateMultiTimeseriesReq,
+    TSCreateSchemaTemplateReq,
+    TSCreateTimeseriesReq,
+    TSDropSchemaTemplateReq,
+    TSExecuteStatementReq,
     TSInsertRecordReq,
+    TSInsertRecordsOfOneDeviceReq,
+    TSInsertRecordsReq,
     TSInsertStringRecordReq,
     TSInsertTabletReq,
-    TSExecuteStatementReq,
-    TSOpenSessionReq,
-    TSCreateMultiTimeseriesReq,
-    TSCloseSessionReq,
     TSInsertTabletsReq,
-    TSInsertRecordsReq,
-    TSInsertRecordsOfOneDeviceReq,
-    TSCreateSchemaTemplateReq,
-    TSDropSchemaTemplateReq,
-    TSAppendSchemaTemplateReq,
+    TSOpenSessionReq,
     TSPruneSchemaTemplateReq,
+    TSQueryTemplateReq,
     TSSetSchemaTemplateReq,
     TSUnsetSchemaTemplateReq,
-    TSQueryTemplateReq,
 )
 from .thrift.rpc.ttypes import (
     TSDeleteDataReq,
+    TSInsertStringRecordsOfOneDeviceReq,
+    TSLastDataQueryReq,
     TSProtocolVersion,
-    TSSetTimeZoneReq,
     TSRawDataQueryReq,
-    TSLastDataQueryReq,
-    TSInsertStringRecordsOfOneDeviceReq,
+    TSSetTimeZoneReq,
 )
 from .tsfile.utils.date_utils import parse_date_to_int
 from .utils import rpc_utils
@@ -245,6 +247,7 @@ def init_connection(self, endpoint):
     def __get_transport(self, endpoint):
         if self.__use_ssl:
             import ssl
+
             from thrift.transport import TSSLSocket
 
             if sys.version_info >= (3, 10):

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

@@ -0,0 +1,17 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#

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

@@ -23,18 +23,19 @@
 import psutil
 
 from iotdb.ainode.core.config import AINodeDescriptor
-from ainode.core.constant import AINODE_SYSTEM_FILE_NAME
-from ainode.core.log import Logger
-from ainode.core.rpc.client import ClientManager
-from ainode.core.rpc.handler import AINodeRPCServiceHandler
-from ainode.core.rpc.service import AINodeRPCService
-from ainode.thrift.common.ttypes import (
+from iotdb.ainode.core.constant import AINODE_SYSTEM_FILE_NAME
+from iotdb.ainode.core.exception import MissingConfigError
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.rpc.client import ClientManager
+from iotdb.ainode.core.rpc.handler import AINodeRPCServiceHandler
+from iotdb.ainode.core.rpc.service import AINodeRPCService
+from iotdb.thrift.common.ttypes import (
     TAINodeConfiguration,
     TAINodeLocation,
     TEndPoint,
     TNodeResource,
 )
-from ainode.thrift.confignode.ttypes import TNodeVersionInfo
+from iotdb.thrift.confignode.ttypes import TNodeVersionInfo
 
 logger = Logger()
 

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

@@ -17,8 +17,9 @@
 #
 import os
 import re
+import threading
 
-from ainode.core.constant import (
+from iotdb.ainode.core.constant import (
     AINODE_BUILD_INFO,
     AINODE_BUILTIN_MODELS_DIR,
     AINODE_CLUSTER_INGRESS_ADDRESS,
@@ -47,9 +48,9 @@
     AINODE_VERSION_INFO,
 )
 from iotdb.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
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.util.decorator import singleton
+from iotdb.thrift.common.ttypes import TEndPoint
 
 logger = Logger()
 
@@ -101,15 +102,37 @@ def __init__(self):
         # Whether to enable compression for thrift
         self._ain_thrift_compression_enabled = AINODE_THRIFT_COMPRESSION_ENABLED
 
+        # use for ssl
+        self._ain_cluster_ingress_ssl_enabled = False
+        self._ain_internal_ssl_enabled = False
+        self._ain_thrift_ssl_cert_file = None
+        self._ain_thrift_ssl_key_file = None
+
         # Cache number of model storage to avoid repeated loading
         self._ain_model_storage_cache_size = 30
 
+        # activation
+        self._ain_activated = threading.Event()
+
     def get_cluster_name(self) -> str:
         return self._cluster_name
 
     def set_cluster_name(self, cluster_name: str) -> None:
         self._cluster_name = cluster_name
 
+    def is_activated(self) -> bool:
+        return self._ain_activated.is_set()
+
+    def set_activated(self, is_activated: bool) -> None:
+        if is_activated:
+            if not self._ain_activated.is_set():
+                logger.info("TimechoDB-AINode is activated.")
+            self._ain_activated.set()
+        else:
+            if self._ain_activated.is_set():
+                logger.info("TimechoDB-AINode is deactivated.")
+            self._ain_activated.clear()
+
     def get_version_info(self) -> str:
         return self._version_info
 
@@ -218,6 +241,32 @@ def set_ain_thrift_compression_enabled(
     ) -> None:
         self._ain_thrift_compression_enabled = ain_thrift_compression_enabled
 
+    def get_ain_cluster_ingress_ssl_enabled(self) -> bool:
+        return self._ain_cluster_ingress_ssl_enabled
+
+    def set_ain_cluster_ingress_ssl_enabled(
+        self, ain_cluster_ingress_ssl_enabled: int
+    ) -> None:
+        self._ain_cluster_ingress_ssl_enabled = ain_cluster_ingress_ssl_enabled
+
+    def get_ain_internal_ssl_enabled(self) -> bool:
+        return self._ain_internal_ssl_enabled
+
+    def set_ain_internal_ssl_enabled(self, ain_internal_ssl_enabled: int) -> None:
+        self._ain_internal_ssl_enabled = ain_internal_ssl_enabled
+
+    def get_ain_thrift_ssl_cert_file(self) -> str:
+        return self._ain_thrift_ssl_cert_file
+
+    def set_ain_thrift_ssl_cert_file(self, ain_thrift_ssl_cert_file: str) -> None:
+        self._ain_thrift_ssl_cert_file = ain_thrift_ssl_cert_file
+
+    def get_ain_thrift_ssl_key_file(self) -> str:
+        return self._ain_thrift_ssl_key_file
+
+    def set_ain_thrift_ssl_key_file(self, ain_thrift_ssl_key_file: str) -> None:
+        self._ain_thrift_ssl_key_file = ain_thrift_ssl_key_file
+
     def get_ain_model_storage_cache_size(self) -> int:
         return self._ain_model_storage_cache_size
 
@@ -358,6 +407,26 @@ def _load_config_from_file(self) -> None:
                     int(file_configs["ain_thrift_compression_enabled"])
                 )
 
+            if "ain_cluster_ingress_ssl_enabled" in config_keys:
+                self._config.set_ain_cluster_ingress_ssl_enabled(
+                    int(file_configs["ain_cluster_ingress_ssl_enabled"])
+                )
+
+            if "ain_internal_ssl_enabled" in config_keys:
+                self._config.set_ain_internal_ssl_enabled(
+                    int(file_configs["ain_internal_ssl_enabled"])
+                )
+
+            if "ain_thrift_ssl_cert_file" in config_keys:
+                self._config.set_ain_thrift_ssl_cert_file(
+                    file_configs["ain_thrift_ssl_cert_file"]
+                )
+
+            if "ain_thrift_ssl_key_file" in config_keys:
+                self._config.set_ain_thrift_ssl_key_file(
+                    file_configs["ain_thrift_ssl_key_file"]
+                )
+
             if "ain_logs_dir" in config_keys:
                 log_dir = file_configs["ain_logs_dir"]
                 self._config.set_ain_logs_dir(log_dir)

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

@@ -15,13 +15,14 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+
 import logging
 import os
 from enum import Enum
 from typing import List
 
 from iotdb.ainode.core.model.model_enums import BuiltInModelType
-from ainode.thrift.common.ttypes import TEndPoint
+from iotdb.thrift.common.ttypes import TEndPoint
 
 IOTDB_AINODE_HOME = os.getenv("IOTDB_AINODE_HOME", "")
 AINODE_VERSION_INFO = "UNKNOWN"
@@ -40,7 +41,7 @@
 AINODE_CLUSTER_INGRESS_ADDRESS = "127.0.0.1"
 AINODE_CLUSTER_INGRESS_PORT = 6667
 AINODE_CLUSTER_INGRESS_USERNAME = "root"
-AINODE_CLUSTER_INGRESS_PASSWORD = "root"
+AINODE_CLUSTER_INGRESS_PASSWORD = "TimechoDB@2021"
 AINODE_CLUSTER_INGRESS_TIME_ZONE = "UTC+8"
 
 # RPC config
@@ -76,6 +77,9 @@
 INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE = (
     "log_inference_rank_{}_"  # example: log_inference_rank_0_all.log
 )
+TRAINING_LOG_FILE_NAME_PREFIX_TEMPLATE = (
+    "log_training_rank_{}_"  # example: log_training_rank_0_all.log
+)
 
 # AINode model management
 MODEL_WEIGHTS_FILE_IN_SAFETENSORS = "model.safetensors"
@@ -96,6 +100,10 @@ class TSStatusCode(Enum):
     INVALID_INFERENCE_CONFIG = 1512
     INFERENCE_INTERNAL_ERROR = 1520
 
+    # Training status codes
+    INVALID_TRAINING_CONFIG = 1550
+    TRAINING_INTERNAL_ERROR = 1551
+
     def get_status_code(self) -> int:
         return self.value
 

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

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