[AINode] Fix cp errors

Author: CRZbulabula

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/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):

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

@@ -17,17 +17,17 @@
 #
 import numpy as np
 import torch
+from torch.utils.data import Dataset
+
+from iotdb.ainode.core.config import AINodeDescriptor
+from iotdb.ainode.core.ingress.dataset import BasicDatabaseForecastDataset
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.util.cache import MemoryLRUCache
+from iotdb.ainode.core.util.decorator import singleton
 from iotdb.Session import Session
 from iotdb.table_session import TableSession, TableSessionConfig
 from iotdb.utils.Field import Field
 from iotdb.utils.IoTDBConstants import TSDataType
-from torch.utils.data import Dataset
-
-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.ingress import BasicDatabaseForecastDataset
-from ainode.core.log import Logger
-from ainode.core.util.cache import MemoryLRUCache
-from ainode.core.util.decorator import singleton
 
 logger = Logger()
 
@@ -91,6 +91,10 @@ def __init__(
             user=username,
             password=password,
             zone_id=time_zone,
+            use_ssl=AINodeDescriptor()
+            .get_config()
+            .get_ain_cluster_ingress_ssl_enabled(),
+            ca_certs=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(),
         )
         self.session.open(False)
         self.use_rate = use_rate
@@ -269,6 +273,10 @@ def __init__(
             username=username,
             password=password,
             time_zone=time_zone,
+            use_ssl=AINodeDescriptor()
+            .get_config()
+            .get_ain_cluster_ingress_ssl_enabled(),
+            ca_certs=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(),
         )
         self.session = TableSession(table_session_config)
         self.use_rate = use_rate

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

@@ -26,14 +26,14 @@
 import threading
 from logging.handlers import TimedRotatingFileHandler
 
-from ainode.core.constant import (
+from iotdb.ainode.core.constant import (
     AINODE_LOG_DIR,
     AINODE_LOG_FILE_LEVELS,
     AINODE_LOG_FILE_NAME_PREFIX,
     DEFAULT_LOG_LEVEL,
     LOG_FILE_TYPE,
 )
-from ainode.core.util.decorator import singleton
+from iotdb.ainode.core.util.decorator import singleton
 
 
 class BaseLogger:

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

@@ -17,8 +17,8 @@
 #
 import psutil
 
-from ainode.thrift.ainode.ttypes import TAIHeartbeatReq, TAIHeartbeatResp
-from ainode.thrift.common.ttypes import TLoadSample
+from iotdb.thrift.ainode.ttypes import TAIHeartbeatReq, TAIHeartbeatResp
+from iotdb.thrift.common.ttypes import TLoadSample
 
 
 class ClusterManager:
@@ -42,5 +42,6 @@ def get_heart_beat(req: TAIHeartbeatReq) -> TAIHeartbeatResp:
             )
         else:
             return TAIHeartbeatResp(
-                heartbeatTimestamp=req.heartbeatTimestamp, status="Running"
+                heartbeatTimestamp=req.heartbeatTimestamp,
+                status="Running",
             )