This is an automated email from the ASF dual-hosted git repository. yongzao pushed a commit to branch cp-ain-to-206 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 97947e2ca76763f1a69e98db5387bf4b2bb0420f Author: Yongzao <[email protected]> AuthorDate: Tue Dec 16 13:58:24 2025 +0800 [AINode] Fix cp errors --- iotdb-core/ainode/build_binary.py | 8 +- iotdb-core/ainode/iotdb/Session.py | 33 +-- iotdb-core/ainode/iotdb/ainode/core/ai_node.py | 15 +- iotdb-core/ainode/iotdb/ainode/core/config.py | 77 ++++++- iotdb-core/ainode/iotdb/ainode/core/constant.py | 12 +- iotdb-core/ainode/iotdb/ainode/core/exception.py | 5 +- .../ainode/iotdb/ainode/core/ingress/iotdb.py | 22 +- iotdb-core/ainode/iotdb/ainode/core/log.py | 4 +- .../iotdb/ainode/core/manager/cluster_manager.py | 7 +- .../iotdb/ainode/core/manager/inference_manager.py | 244 +++++++++------------ .../iotdb/ainode/core/manager/model_manager.py | 21 +- .../ainode/iotdb/ainode/core/manager/utils.py | 33 ++- .../ainode/core/model/built_in_model_factory.py | 46 ++-- .../iotdb/ainode/core/model/model_factory.py | 8 +- .../ainode/iotdb/ainode/core/model/model_info.py | 2 +- .../iotdb/ainode/core/model/model_storage.py | 16 +- .../ainode/core/model/sundial/modeling_sundial.py | 11 +- .../ainode/core/model/timerxl/modeling_timer.py | 6 +- .../ainode/iotdb/ainode/core/model/uri_utils.py | 6 +- iotdb-core/ainode/iotdb/ainode/core/rpc/client.py | 63 +++++- iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py | 79 ++++++- iotdb-core/ainode/iotdb/ainode/core/rpc/service.py | 39 +++- iotdb-core/ainode/iotdb/ainode/core/rpc/status.py | 6 +- iotdb-core/ainode/iotdb/ainode/core/script.py | 87 ++------ iotdb-core/ainode/iotdb/ainode/core/util/cache.py | 2 +- .../ainode/iotdb/ainode/core/util/decorator.py | 15 ++ .../ainode/iotdb/sqlalchemy/IoTDBSQLCompiler.py | 3 +- iotdb-core/ainode/iotdb/table_session_pool.py | 2 +- .../ainode/iotdb/template/MeasurementNode.py | 5 +- iotdb-core/ainode/iotdb/template/Template.py | 2 +- .../iotdb/thrift/ainode/IAINodeRPCService-remote | 6 +- .../iotdb/thrift/ainode/IAINodeRPCService.py | 18 +- iotdb-core/ainode/iotdb/thrift/ainode/constants.py | 11 +- iotdb-core/ainode/iotdb/thrift/ainode/ttypes.py | 13 +- iotdb-core/ainode/iotdb/thrift/common/constants.py | 11 +- iotdb-core/ainode/iotdb/thrift/common/ttypes.py | 14 +- .../thrift/confignode/IConfigNodeRPCService-remote | 6 +- .../thrift/confignode/IConfigNodeRPCService.py | 18 +- .../ainode/iotdb/thrift/confignode/constants.py | 11 +- .../ainode/iotdb/thrift/confignode/ttypes.py | 13 +- .../thrift/datanode/IDataNodeRPCService-remote | 6 +- .../iotdb/thrift/datanode/IDataNodeRPCService.py | 18 +- .../thrift/datanode/MPPDataExchangeService-remote | 6 +- .../thrift/datanode/MPPDataExchangeService.py | 18 +- .../ainode/iotdb/thrift/datanode/constants.py | 11 +- iotdb-core/ainode/iotdb/thrift/datanode/ttypes.py | 13 +- .../iotdb/thrift/rpc/IClientRPCService-remote | 6 +- .../ainode/iotdb/thrift/rpc/IClientRPCService.py | 18 +- iotdb-core/ainode/iotdb/thrift/rpc/constants.py | 11 +- iotdb-core/ainode/iotdb/thrift/rpc/ttypes.py | 13 +- .../iotdb/tsfile/utils/read_write_io_utils.py | 2 +- .../ainode/iotdb/tsfile/utils/tsblock_serde.py | 1 - iotdb-core/ainode/iotdb/utils/Field.py | 8 +- iotdb-core/ainode/iotdb/utils/IoTDBConstants.py | 3 +- iotdb-core/ainode/iotdb/utils/NumpyTablet.py | 4 +- iotdb-core/ainode/iotdb/utils/SessionDataSet.py | 6 +- iotdb-core/ainode/iotdb/utils/Tablet.py | 2 +- iotdb-core/ainode/iotdb/utils/iotdb_rpc_dataset.py | 4 +- iotdb-core/ainode/pom.xml | 14 +- iotdb-core/ainode/pyproject.toml | 109 ++++----- .../InformationSchemaContentSupplierFactory.java | 1 - 61 files changed, 789 insertions(+), 485 deletions(-) diff --git a/iotdb-core/ainode/build_binary.py b/iotdb-core/ainode/build_binary.py index 01b0ef67e8f..6607521e152 100644 --- a/iotdb-core/ainode/build_binary.py +++ b/iotdb-core/ainode/build_binary.py @@ -31,7 +31,7 @@ from pathlib import Path 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 """ diff --git a/iotdb-core/ainode/iotdb/Session.py b/iotdb-core/ainode/iotdb/Session.py index 3f44bd41584..155c0ebe87b 100644 --- a/iotdb-core/ainode/iotdb/Session.py +++ b/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 @@ class Session(object): def __get_transport(self, endpoint): if self.__use_ssl: import ssl + from thrift.transport import TSSLSocket if sys.version_info >= (3, 10): diff --git a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py index 13ff1f485d4..438faf1e1c0 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py +++ b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py @@ -23,18 +23,19 @@ from datetime import datetime 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() diff --git a/iotdb-core/ainode/iotdb/ainode/core/config.py b/iotdb-core/ainode/iotdb/ainode/core/config.py index 43b6ce7ccda..f5021ca2c41 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/config.py +++ b/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 @@ from ainode.core.constant import ( 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 @@ class AINodeConfig(object): # 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 @@ class AINodeConfig(object): ) -> 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 @@ class AINodeDescriptor(object): 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) diff --git a/iotdb-core/ainode/iotdb/ainode/core/constant.py b/iotdb-core/ainode/iotdb/ainode/core/constant.py index fd8f10177a9..e6b213d6a5d 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/constant.py +++ b/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_RPC_PORT = 10810 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 @@ DEFAULT_LOG_LEVEL = logging.INFO 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 diff --git a/iotdb-core/ainode/iotdb/ainode/core/exception.py b/iotdb-core/ainode/iotdb/ainode/core/exception.py index f00900e8d01..bc89cdc3066 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/exception.py +++ b/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): diff --git a/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py index 792d64b2980..13c56ca9d2d 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py +++ b/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 @@ class IoTDBTreeModelDataset(BasicDatabaseForecastDataset): 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 @@ class IoTDBTableModelDataset(BasicDatabaseForecastDataset): 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 diff --git a/iotdb-core/ainode/iotdb/ainode/core/log.py b/iotdb-core/ainode/iotdb/ainode/core/log.py index 947fe820c30..fd121b26349 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/log.py +++ b/iotdb-core/ainode/iotdb/ainode/core/log.py @@ -26,14 +26,14 @@ import sys 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: diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py index b7d65f47dc1..8b8731cb92f 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py +++ b/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 @@ class ClusterManager: ) else: return TAIHeartbeatResp( - heartbeatTimestamp=req.heartbeatTimestamp, status="Running" + heartbeatTimestamp=req.heartbeatTimestamp, + status="Running", ) 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 65f4abacfca..a67d576b0ec 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py @@ -15,54 +15,58 @@ # specific language governing permissions and limitations # under the License. # + import threading import time from abc import ABC, abstractmethod -from typing import Dict, List +from typing import Dict import pandas as pd import torch import torch.multiprocessing as mp -from iotdb.tsfile.utils.tsblock_serde import deserialize from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.constant import TSStatusCode +from iotdb.ainode.core.constant import TSStatusCode from iotdb.ainode.core.exception import ( InferenceModelInternalError, InvalidWindowArgumentError, NumericalRangeException, runtime_error_extractor, ) -from ainode.core.inference.inference_request import ( +from iotdb.ainode.core.inference.inference_request import ( InferenceRequest, InferenceRequestProxy, ) -from iotdb.ainode.core.inference.inference_request_pool import InferenceRequestPool -from iotdb.ainode.core.inference.strategy import ( +from iotdb.ainode.core.inference.pool_controller import PoolController +from iotdb.ainode.core.inference.strategy.timer_sundial_inference_pipeline import ( TimerSundialInferencePipeline, ) -from iotdb.ainode.core.inference.strategy import ( +from iotdb.ainode.core.inference.strategy.timerxl_inference_pipeline import ( TimerXLInferencePipeline, ) -from ainode.core.inference.utils import generate_req_id -from ainode.core.log import Logger -from ainode.core.manager.model_manager import ModelManager -from ainode.core.manager.utils import ( - _estimate_pool_size, - _measure_model_memory, -) -from ainode.core.model.sundial.configuration_sundial import SundialConfig -from ainode.core.model.sundial.modeling_sundial import SundialForPrediction -from ainode.core.model.timerxl.configuration_timer import TimerConfig +from iotdb.ainode.core.inference.utils import generate_req_id +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.manager.model_manager import ModelManager +from iotdb.ainode.core.model.model_enums import BuiltInModelType +from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig +from iotdb.ainode.core.model.sundial.modeling_sundial import SundialForPrediction +from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig from iotdb.ainode.core.model.timerxl.modeling_timer import TimerForPrediction -from ainode.core.rpc.status import get_status +from iotdb.ainode.core.rpc.status import get_status +from iotdb.ainode.core.util.gpu_mapping import get_available_devices from iotdb.ainode.core.util.serde import convert_to_binary -from ainode.thrift.ainode.ttypes import ( +from iotdb.thrift.ainode.ttypes import ( TForecastReq, TForecastResp, TInferenceReq, TInferenceResp, + TLoadModelReq, + TShowLoadedModelsReq, + TShowLoadedModelsResp, + TUnloadModelReq, ) +from iotdb.thrift.common.ttypes import TSStatus +from iotdb.tsfile.utils.tsblock_serde import deserialize logger = Logger() @@ -145,95 +149,79 @@ class RegisteredStrategy(InferenceStrategy): class InferenceManager: - ACCELERATE_MODEL_ID = ["sundial", "timer_xl"] - DEFAULT_DEVICE = "cpu" - # DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") WAITING_INTERVAL_IN_MS = ( AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms() ) # How often to check for requests in the result queue def __init__(self): self._model_manager = ModelManager() + self._model_mem_usage_map: Dict[str, int] = ( + {} + ) # store model memory usage for each model self._result_queue = mp.Queue() self._result_wrapper_map = {} self._result_wrapper_lock = threading.RLock() - # structure: {model_id: [(InferenceRequestPool, request_queue), ...]} - self._request_pool_map: Dict[str, List[(InferenceRequestPool, mp.Queue)]] = {} + self._stop_event = mp.Event() self._result_handler_thread = threading.Thread( target=self._handle_results, daemon=True ) self._result_handler_thread.start() - self._model_mem_usage_map: Dict[str, int] = ( - {} - ) # store model memory usage for each model - # self._preload_model_benchmarks() - - def _preload_model_benchmarks(self): - if "cuda" in str(self.DEFAULT_DEVICE): - for model_id in self.ACCELERATE_MODEL_ID: - mem_usage = _measure_model_memory(self.DEFAULT_DEVICE, model_id) - self._model_mem_usage_map[model_id] = mem_usage - logger.info( - f"[Inference] Preloaded benchmark for {model_id}, mem_usage={mem_usage/1024**2:.2f} MB" - ) - else: - logger.warning( - f"[Inference] Skipped preloading benchmarks for {self.DEFAULT_DEVICE}, only supports CUDA currently" + self._pool_controller = PoolController(self._result_queue) + + def load_model(self, req: TLoadModelReq) -> TSStatus: + devices_to_be_processed = [] + devices_not_to_be_processed = [] + for device_id in req.deviceIdList: + if self._pool_controller.has_request_pools( + model_id=req.existingModelId, device_id=device_id + ): + devices_not_to_be_processed.append(device_id) + else: + devices_to_be_processed.append(device_id) + if len(devices_to_be_processed) > 0: + self._pool_controller.load_model( + model_id=req.existingModelId, device_id_list=devices_to_be_processed ) - - def _first_pool_init(self, model_id: str): - if model_id == "sundial": - config = SundialConfig() - elif model_id == "timer_xl": - config = TimerConfig() - first_queue = mp.Queue() - ready_event = mp.Event() - first_pool = InferenceRequestPool( - pool_id=0, - model_id=model_id, - config=config, - request_queue=first_queue, - result_queue=self._result_queue, - ready_event=ready_event, + logger.info( + f"[Inference] Start loading model [{req.existingModelId}] to devices [{devices_to_be_processed}], skipped devices [{devices_not_to_be_processed}] cause they have already loaded this model." + ) + return TSStatus( + code=TSStatusCode.SUCCESS_STATUS.value, + message='Successfully submitted load model task, please use "SHOW LOADED MODELS" to check progress.', ) - first_pool.start() - if not ready_event.wait(timeout=30): - logger.error( - f"[Inference][Device-{self.DEFAULT_DEVICE}][Pool-0] First pool failed to be ready in time" - ) - else: - self._request_pool_map[model_id] = [(first_pool, first_queue)] - logger.info( - f"[Inference][Device-{self.DEFAULT_DEVICE}][Pool-0] Initialized inference request pool for model {model_id}" - ) - def _expand_pools(self, model_id, start_idx, count): - for idx in range(count): - queue = mp.Queue() - if model_id == "sundial": - config = SundialConfig() - elif model_id == "timer_xl": - config = TimerConfig() - pool = InferenceRequestPool( - pool_id=start_idx + idx, - model_id=model_id, - config=config, - request_queue=queue, - result_queue=self._result_queue, - ready_event=mp.Event(), - ) - pool.start() - if not pool.ready_event.wait(timeout=30): - logger.error( - f"[Inference][Device-{self.DEFAULT_DEVICE}][Pool-{start_idx + idx}] Pool failed to be ready in time" - ) - continue + def unload_model(self, req: TUnloadModelReq) -> TSStatus: + devices_to_be_processed = [] + devices_not_to_be_processed = [] + for device_id in req.deviceIdList: + if self._pool_controller.has_request_pools( + model_id=req.modelId, device_id=device_id + ): + devices_to_be_processed.append(device_id) else: - self._request_pool_map[model_id].append((pool, queue)) - logger.info( - f"[Inference][Device-{self.DEFAULT_DEVICE}][Pool-{pool.pool_id}] New inference request pool started for model {model_id}" - ) + devices_not_to_be_processed.append(device_id) + if len(devices_to_be_processed) > 0: + self._pool_controller.unload_model( + model_id=req.modelId, device_id_list=req.deviceIdList + ) + logger.info( + f"[Inference] Start unloading model [{req.modelId}] from devices [{devices_to_be_processed}], skipped devices [{devices_not_to_be_processed}] cause they haven't loaded this model." + ) + return TSStatus( + code=TSStatusCode.SUCCESS_STATUS.value, + message='Successfully submitted unload model task, please use "SHOW LOADED MODELS" to check progress.', + ) + + def show_loaded_models(self, req: TShowLoadedModelsReq) -> TShowLoadedModelsResp: + return TShowLoadedModelsResp( + status=get_status(TSStatusCode.SUCCESS_STATUS), + deviceLoadedModelsMap=self._pool_controller.show_loaded_models( + req.deviceIdList + if len(req.deviceIdList) > 0 + else get_available_devices() + ), + ) def _handle_results(self): while not self._stop_event.is_set(): @@ -246,12 +234,29 @@ class InferenceManager: infer_req.get_final_output() ) + def _process_request(self, req): + req_id = req.req_id + infer_proxy = InferenceRequestProxy(req_id) + with self._result_wrapper_lock: + self._result_wrapper_map[req_id] = infer_proxy + try: + # dispatch request to the pool + self._pool_controller.add_request(req, infer_proxy) + outputs = infer_proxy.wait_for_result() + return outputs + except Exception as e: + logger.error(e) + raise InferenceModelInternalError(str(e)) + finally: + with self._result_wrapper_lock: + del self._result_wrapper_map[req_id] + def _get_strategy(self, model_id, model): if isinstance(model, TimerForPrediction): return TimerXLStrategy(model) if isinstance(model, SundialForPrediction): return SundialStrategy(model) - if self._model_manager.model_storage._is_built_in_or_fine_tuned(model_id): + if self._model_manager.model_storage.is_built_in_or_fine_tuned(model_id): return BuiltInStrategy(model) return RegisteredStrategy(model) @@ -284,26 +289,8 @@ class InferenceManager: predict_length, ) - if model_id in self.ACCELERATE_MODEL_ID and "cuda" in str( - self.DEFAULT_DEVICE - ): - # lazy initialization for first request - if model_id not in self._request_pool_map: - pool_num = _estimate_pool_size(self.DEFAULT_DEVICE, model_id) - if pool_num <= 0: - raise InferenceModelInternalError( - f"Not enough memory to run model {model_id}." - ) - # initialize the first pool - self._first_pool_init(model_id) - # start a background thread to expand pools - expand_thread = threading.Thread( - target=self._expand_pools, - args=(model_id, 1, pool_num - 1), - daemon=True, - ) - expand_thread.start() - # TODO: Logic in this branch shall handle all LTSM inferences + if self._pool_controller.has_request_pools(model_id): + # use request pool to accelerate inference when the model instance is already loaded. # TODO: TSBlock -> Tensor codes should be unified data = full_data[1][0] if data.dtype.byteorder not in ("=", "|"): @@ -311,30 +298,24 @@ class InferenceManager: data = np_data.view(np_data.dtype.newbyteorder()) # the inputs should be on CPU before passing to the inference request inputs = torch.tensor(data).unsqueeze(0).float().to("cpu") - if model_id == "sundial": + model_type = self._model_manager.get_model_info(model_id).model_type + if model_type == BuiltInModelType.SUNDIAL.value: inference_pipeline = TimerSundialInferencePipeline(SundialConfig()) - elif model_id == "timer_xl": + elif model_type == BuiltInModelType.TIMER_XL.value: inference_pipeline = TimerXLInferencePipeline(TimerConfig()) + else: + raise InferenceModelInternalError( + f"Unsupported model_id: {model_id}" + ) infer_req = InferenceRequest( req_id=generate_req_id(), + model_id=model_id, inputs=inputs, inference_pipeline=inference_pipeline, max_new_tokens=predict_length, ) - infer_proxy = InferenceRequestProxy(infer_req.req_id) - with self._result_wrapper_lock: - self._result_wrapper_map[infer_req.req_id] = infer_proxy - pool_idx = hash(infer_req.req_id) % len( - self._request_pool_map[model_id] - ) - self._request_pool_map[model_id][pool_idx][1].put(infer_req) - logger.debug( - f"[Inference][Device-{self.DEFAULT_DEVICE}][Pool-{pool_idx}][ID-{infer_req.req_id}] Request is queued for inference" - ) - outputs = infer_proxy.wait_for_completion() + outputs = self._process_request(infer_req) outputs = convert_to_binary(pd.DataFrame(outputs[0])) - with self._result_wrapper_lock: - del self._result_wrapper_map[infer_req.req_id] else: # load model accel = str(inference_attrs.get("acceleration", "")).lower() == "true" @@ -385,16 +366,9 @@ class InferenceManager: single_output=False, ) - def shutdown(self): + def stop(self): self._stop_event.set() - for model_id, pools in self._request_pool_map.items(): - for requestPool, requestQueue in pools: - requestPool.stop() - while not requestQueue.empty(): - requestQueue.get_nowait() - requestQueue.close() - for requestPool, _ in pools: - requestPool.join(timeout=10) + 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/manager/model_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py index 086ab69a465..d84bca77c84 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py @@ -20,25 +20,22 @@ from typing import Callable, Dict from torch import nn from yaml import YAMLError -from ainode.core.constant import TSStatusCode -from iotdb.ainode.core.exception import ( - BadConfigValueError, - InvalidUriError, -) -from ainode.core.log import Logger +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.exception import BadConfigValueError, InvalidUriError +from iotdb.ainode.core.log import Logger from iotdb.ainode.core.model.model_enums import BuiltInModelType, ModelStates -from ainode.core.model.model_info import ModelInfo -from ainode.core.model.model_storage import ModelStorage -from ainode.core.rpc.status import get_status -from ainode.core.util.decorator import singleton -from ainode.thrift.ainode.ttypes import ( +from iotdb.ainode.core.model.model_info import ModelInfo +from iotdb.ainode.core.model.model_storage import ModelStorage +from iotdb.ainode.core.rpc.status import get_status +from iotdb.ainode.core.util.decorator import singleton +from iotdb.thrift.ainode.ttypes import ( TDeleteModelReq, TRegisterModelReq, TRegisterModelResp, TShowModelsReq, TShowModelsResp, ) -from ainode.thrift.common.ttypes import TSStatus +from iotdb.thrift.common.ttypes import TSStatus logger = Logger() diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py index 65106e08e08..0264e27331a 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py @@ -22,9 +22,10 @@ import psutil import torch from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.log import Logger -from ainode.core.manager.model_manager import ModelManager -from ainode.core.model.model_info import BUILT_IN_LTSM_MAP +from iotdb.ainode.core.exception import ModelNotExistError +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.manager.model_manager import ModelManager +from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP logger = Logger() @@ -39,7 +40,7 @@ INFERENCE_EXTRA_MEMORY_RATIO = ( ) # the overhead ratio for inference, used to estimate the pool size -def _measure_model_memory(device: torch.device, model_id: str) -> int: +def measure_model_memory(device: torch.device, model_id: str) -> int: # TODO: support CPU in the future # TODO: we can estimate the memory usage by running a dummy inference torch.cuda.empty_cache() @@ -62,7 +63,7 @@ def _measure_model_memory(device: torch.device, model_id: str) -> int: return final -def _evaluate_system_resources(device: torch.device) -> dict: +def evaluate_system_resources(device: torch.device) -> dict: if torch.cuda.is_available(): free_mem, total_mem = torch.cuda.mem_get_info() logger.info( @@ -70,7 +71,6 @@ def _evaluate_system_resources(device: torch.device) -> dict: ) return {"device": "cuda", "free_mem": free_mem, "total_mem": total_mem} else: - # TODO: test cpu in the future free_mem = psutil.virtual_memory().available total_mem = psutil.virtual_memory().total logger.info( @@ -79,21 +79,20 @@ def _evaluate_system_resources(device: torch.device) -> dict: return {"device": "cpu", "free_mem": free_mem, "total_mem": total_mem} -def _estimate_pool_size(device: torch.device, model_id: str) -> int: +def estimate_pool_size(device: torch.device, model_id: str) -> int: model_info = BUILT_IN_LTSM_MAP.get(model_id, None) - if model_info is None: - logger.error(f"[Inference][Device-{device}] Model {model_id} not found") - return 0 - - model_type = model_info.model_type - if model_type not in MODEL_MEM_USAGE_MAP: - logger.error(f"[Inference][Device-{device}] Model {model_id} not supported now") - return 0 + if model_info is None or model_info.model_type not in MODEL_MEM_USAGE_MAP: + logger.error( + f"[Inference] Cannot estimate inference pool size on device: {device}, because model: {model_id} is not supported." + ) + raise ModelNotExistError(model_id) - system_res = _evaluate_system_resources(device) + system_res = evaluate_system_resources(device) free_mem = system_res["free_mem"] - mem_usage = MODEL_MEM_USAGE_MAP[model_type] * INFERENCE_EXTRA_MEMORY_RATIO + mem_usage = ( + MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO + ) size = int((free_mem * INFERENCE_MEMORY_USAGE_RATIO) // mem_usage) if size <= 0: logger.error( diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py b/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py index a95af34990b..3b55142350b 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py @@ -30,7 +30,7 @@ from sktime.forecasting.naive import NaiveForecaster from sktime.forecasting.trend import STLForecaster from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.constant import ( +from iotdb.ainode.core.constant import ( MODEL_CONFIG_FILE_IN_JSON, MODEL_WEIGHTS_FILE_IN_SAFETENSORS, AttributeName, @@ -43,21 +43,21 @@ from iotdb.ainode.core.exception import ( StringRangeException, WrongAttributeTypeError, ) -from ainode.core.log import Logger +from iotdb.ainode.core.log import Logger from iotdb.ainode.core.model.model_enums import BuiltInModelType -from ainode.core.model.model_info import TIMER_REPO_ID -from ainode.core.model.sundial import modeling_sundial +from iotdb.ainode.core.model.model_info import TIMER_REPO_ID +from iotdb.ainode.core.model.sundial import modeling_sundial from iotdb.ainode.core.model.timerxl import modeling_timer logger = Logger() -def _download_file_from_hf_if_necessary( - file_path: str, local_dir: str, repo_id: str -) -> bool: - if not os.path.exists(file_path): +def _download_file_from_hf_if_necessary(local_dir: str, repo_id: str) -> bool: + weights_path = os.path.join(local_dir, MODEL_WEIGHTS_FILE_IN_SAFETENSORS) + config_path = os.path.join(local_dir, MODEL_CONFIG_FILE_IN_JSON) + if not os.path.exists(weights_path): logger.info( - f"Model file not found at {file_path}, downloading from HuggingFace..." + f"Model weights file not found at {weights_path}, downloading from HuggingFace..." ) try: hf_hub_download( @@ -65,9 +65,27 @@ def _download_file_from_hf_if_necessary( filename=MODEL_WEIGHTS_FILE_IN_SAFETENSORS, local_dir=local_dir, ) - logger.info(f"Got file to {file_path}") + logger.info(f"Got file to {weights_path}") except Exception as e: - logger.error(f"Failed to download model file to {local_dir} due to {e}") + logger.error( + f"Failed to download model weights file to {local_dir} due to {e}" + ) + return False + if not os.path.exists(config_path): + logger.info( + f"Model config file not found at {config_path}, downloading from HuggingFace..." + ) + try: + hf_hub_download( + repo_id=repo_id, + filename=MODEL_CONFIG_FILE_IN_JSON, + local_dir=local_dir, + ) + logger.info(f"Got file to {config_path}") + except Exception as e: + logger.error( + f"Failed to download model config file to {local_dir} due to {e}" + ) return False return True @@ -82,11 +100,7 @@ def download_built_in_ltsm_from_hf_if_necessary( bool: True if the model is existed or downloaded successfully, False otherwise. """ repo_id = TIMER_REPO_ID[model_type] - weights_path = os.path.join(local_dir, MODEL_WEIGHTS_FILE_IN_SAFETENSORS) - if not _download_file_from_hf_if_necessary(weights_path, local_dir, repo_id): - return False - config_path = os.path.join(local_dir, MODEL_CONFIG_FILE_IN_JSON) - if not _download_file_from_hf_if_necessary(config_path, local_dir, repo_id): + if not _download_file_from_hf_if_necessary(local_dir, repo_id): return False return True diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py index c90222108ad..26d863156f3 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py @@ -22,20 +22,20 @@ from urllib.parse import urljoin import yaml -from ainode.core.constant import ( +from iotdb.ainode.core.constant import ( MODEL_CONFIG_FILE_IN_YAML, MODEL_WEIGHTS_FILE_IN_PT, ) from iotdb.ainode.core.exception import BadConfigValueError, InvalidUriError -from ainode.core.log import Logger +from iotdb.ainode.core.log import Logger from iotdb.ainode.core.model.model_enums import ModelFileType -from ainode.core.model.uri_utils import ( +from iotdb.ainode.core.model.uri_utils import ( UriType, download_file, download_snapshot_from_hf, ) from iotdb.ainode.core.util.serde import get_data_type_byte_from_str -from ainode.thrift.ainode.ttypes import TConfigs +from iotdb.thrift.ainode.ttypes import TConfigs logger = Logger() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py index 7cf15e68225..167bfd76640 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py @@ -18,7 +18,7 @@ import glob import os -from ainode.core.constant import ( +from iotdb.ainode.core.constant import ( MODEL_CONFIG_FILE_IN_JSON, MODEL_CONFIG_FILE_IN_YAML, MODEL_WEIGHTS_FILE_IN_PT, diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py index 1f23c57ff64..e346f569102 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py @@ -27,7 +27,7 @@ import torch from torch import nn from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.constant import ( +from iotdb.ainode.core.constant import ( MODEL_CONFIG_FILE_IN_JSON, MODEL_WEIGHTS_FILE_IN_PT, TSStatusCode, @@ -37,8 +37,8 @@ from iotdb.ainode.core.exception import ( ModelNotExistError, UnsupportedError, ) -from ainode.core.log import Logger -from ainode.core.model.built_in_model_factory import ( +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.model.built_in_model_factory import ( download_built_in_ltsm_from_hf_if_necessary, fetch_built_in_model, ) @@ -49,17 +49,17 @@ from iotdb.ainode.core.model.model_enums import ( ModelStates, ) from iotdb.ainode.core.model.model_factory import fetch_model_by_uri -from ainode.core.model.model_info import ( +from iotdb.ainode.core.model.model_info import ( BUILT_IN_LTSM_MAP, BUILT_IN_MACHINE_LEARNING_MODEL_MAP, ModelInfo, get_built_in_model_type, get_model_file_type, ) -from ainode.core.model.uri_utils import get_model_register_strategy -from ainode.core.util.lock import ModelLockPool -from ainode.thrift.ainode.ttypes import TShowModelsReq, TShowModelsResp -from ainode.thrift.common.ttypes import TSStatus +from iotdb.ainode.core.model.uri_utils import get_model_register_strategy +from iotdb.ainode.core.util.lock import ModelLockPool +from iotdb.thrift.ainode.ttypes import TShowModelsReq, TShowModelsResp +from iotdb.thrift.common.ttypes import TSStatus logger = Logger() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py index 6a2402952db..27ca3f1564e 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py @@ -29,9 +29,9 @@ from transformers.modeling_outputs import ( MoeModelOutputWithPast, ) -from ainode.core.log import Logger -from ainode.core.model.sundial.configuration_sundial import SundialConfig -from ainode.core.model.sundial.flow_loss import FlowLoss +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig +from iotdb.ainode.core.model.sundial.flow_loss import FlowLoss from iotdb.ainode.core.model.sundial.ts_generation_mixin import TSGenerationMixin logger = Logger() @@ -554,7 +554,10 @@ class SundialForPrediction(SundialPreTrainedModel, TSGenerationMixin): loss_masks = loss_masks.reshape(bsz * L).repeat( self.config.diffusion_batch_mul ) - mask_y = mask_y.repeat(L * self.config.diffusion_batch_mul, 1) + if mask_y is not None: + mask_y = mask_y.repeat(L * self.config.diffusion_batch_mul, 1) + else: + mask_y = None loss = self.flow_loss(shift_labels, hidden_states, loss_masks, mask_y) else: diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py index 8dfddd335b7..4d6d288d4d3 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py @@ -29,9 +29,9 @@ from transformers.modeling_outputs import ( MoeModelOutputWithPast, ) -from ainode.core.log import Logger -from ainode.core.model.timerxl.configuration_timer import TimerConfig -from ainode.core.model.timerxl.ts_generation_mixin import TSGenerationMixin +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig +from iotdb.ainode.core.model.timerxl.ts_generation_mixin import TSGenerationMixin logger = Logger() diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py b/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py index 59f03057e85..b2e759e00ce 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py @@ -23,15 +23,15 @@ from huggingface_hub import snapshot_download from requests import Session from requests.adapters import HTTPAdapter -from ainode.core.constant import ( +from iotdb.ainode.core.constant import ( DEFAULT_CHUNK_SIZE, DEFAULT_RECONNECT_TIMEOUT, DEFAULT_RECONNECT_TIMES, ) from iotdb.ainode.core.exception import UnsupportedError -from ainode.core.log import Logger +from iotdb.ainode.core.log import Logger from iotdb.ainode.core.model.model_enums import ModelFileType -from ainode.core.model.model_info import get_model_file_type +from iotdb.ainode.core.model.model_info import get_model_file_type HTTP_PREFIX = "http://"; HTTPS_PREFIX = "https://"; diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py index dc9d84de9ae..0c44ea5d423 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py @@ -19,25 +19,26 @@ import time from thrift.protocol import TBinaryProtocol, TCompactProtocol from thrift.Thrift import TException -from thrift.transport import TSocket, TTransport +from thrift.transport import TSocket, TSSLSocket, TTransport from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.constant import TSStatusCode -from ainode.core.log import Logger -from ainode.core.rpc.status import verify_success -from ainode.core.util.decorator import singleton -from ainode.thrift.common.ttypes import ( +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.rpc.status import verify_success +from iotdb.ainode.core.util.decorator import singleton +from iotdb.thrift.common.ttypes import ( TAINodeConfiguration, TAINodeLocation, TEndPoint, TSStatus, ) -from ainode.thrift.confignode import IConfigNodeRPCService -from ainode.thrift.confignode.ttypes import ( +from iotdb.thrift.confignode import IConfigNodeRPCService +from iotdb.thrift.confignode.ttypes import ( TAINodeRegisterReq, TAINodeRemoveReq, TAINodeRestartReq, TNodeVersionInfo, + TSystemConfigurationResp, TUpdateModelInfoReq, ) @@ -109,9 +110,31 @@ class ConfigNodeClient(object): raise TException(self._MSG_RECONNECTION_FAIL) def _connect(self, target_config_node: TEndPoint) -> None: - transport = TTransport.TFramedTransport( - TSocket.TSocket(target_config_node.ip, target_config_node.port) - ) + if AINodeDescriptor().get_config().get_ain_internal_ssl_enabled(): + import ssl + import sys + + if sys.version_info >= (3, 10): + context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) + else: + context = ssl.SSLContext(ssl.PROTOCOL_TLS) + context.verify_mode = ssl.CERT_REQUIRED + context.check_hostname = False + context.load_verify_locations( + cafile=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file() + ) + context.load_cert_chain( + certfile=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(), + keyfile=AINodeDescriptor().get_config().get_ain_thrift_ssl_key_file(), + ) + socket = TSSLSocket.TSSLSocket( + host=target_config_node.ip, + port=target_config_node.port, + ssl_context=context, + ) + else: + socket = TSocket.TSocket(target_config_node.ip, target_config_node.port) + transport = TTransport.TFramedTransport(socket) if not transport.isOpen(): try: transport.open() @@ -149,6 +172,24 @@ class ConfigNodeClient(object): return True return False + def getSystemConfiguration(self) -> TSystemConfigurationResp: + for _ in range(0, self._RETRY_NUM): + try: + resp = self._client.getSystemConfiguration() + if not self._update_config_node_leader(resp.status): + verify_success( + resp.status, + "An error occurs when calling getSystemConfiguration()", + ) + return resp + except TTransport.TException: + logger.warning( + "Failed to connect to ConfigNode {} from AINode when executing getSystemConfiguration()", + self._config_leader, + ) + self._config_leader = None + self._wait_and_reconnect() + def node_register( self, cluster_name: str, diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py index eaad313903d..58e8d9bd0b1 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py @@ -15,31 +15,54 @@ # specific language governing permissions and limitations # under the License. # -from ainode.core.constant import TSStatusCode -from ainode.core.log import Logger -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.core.rpc.status import get_status -from ainode.thrift.ainode import IAINodeRPCService -from ainode.thrift.ainode.ttypes import ( + +from iotdb.ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.manager.cluster_manager import ClusterManager +from iotdb.ainode.core.manager.inference_manager import InferenceManager +from iotdb.ainode.core.manager.model_manager import ModelManager +from iotdb.ainode.core.model.model_info import ModelCategory, ModelInfo, ModelStates +from iotdb.ainode.core.rpc.status import get_status +from iotdb.ainode.core.util.gpu_mapping import get_available_devices +from iotdb.thrift.ainode import IAINodeRPCService +from iotdb.thrift.ainode.ttypes import ( TAIHeartbeatReq, TAIHeartbeatResp, TDeleteModelReq, TForecastReq, TInferenceReq, TInferenceResp, + TLoadModelReq, TRegisterModelReq, TRegisterModelResp, + TShowAIDevicesResp, + TShowLoadedModelsReq, + TShowLoadedModelsResp, TShowModelsReq, TShowModelsResp, TTrainingReq, + TUnloadModelReq, ) -from ainode.thrift.common.ttypes import TSStatus +from iotdb.thrift.common.ttypes import TSStatus logger = Logger() +def _ensure_device_id_is_available(device_id_list: list[str]) -> TSStatus: + """ + Ensure that the device IDs in the provided list are available. + """ + available_devices = get_available_devices() + for device_id in device_id_list: + if device_id not in available_devices: + return TSStatus( + code=TSStatusCode.INVALID_URI_ERROR.value, + message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.", + ) + return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value) + + class AINodeRPCServiceHandler(IAINodeRPCService.Iface): def __init__(self, ainode): self._ainode = ainode @@ -57,6 +80,24 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface): def registerModel(self, req: TRegisterModelReq) -> TRegisterModelResp: return self._model_manager.register_model(req) + def loadModel(self, req: TLoadModelReq) -> TSStatus: + status = self._ensure_model_is_built_in_or_fine_tuned(req.existingModelId) + if status.code != TSStatusCode.SUCCESS_STATUS.value: + return status + status = _ensure_device_id_is_available(req.deviceIdList) + if status.code != TSStatusCode.SUCCESS_STATUS.value: + return status + return self._inference_manager.load_model(req) + + def unloadModel(self, req: TUnloadModelReq) -> TSStatus: + status = self._ensure_model_is_built_in_or_fine_tuned(req.modelId) + if status.code != TSStatusCode.SUCCESS_STATUS.value: + return status + status = _ensure_device_id_is_available(req.deviceIdList) + if status.code != TSStatusCode.SUCCESS_STATUS.value: + return status + return self._inference_manager.unload_model(req) + def deleteModel(self, req: TDeleteModelReq) -> TSStatus: return self._model_manager.delete_model(req) @@ -72,5 +113,25 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface): def showModels(self, req: TShowModelsReq) -> TShowModelsResp: return self._model_manager.show_models(req) + def showLoadedModels(self, req: TShowLoadedModelsReq) -> TShowLoadedModelsResp: + status = _ensure_device_id_is_available(req.deviceIdList) + if status.code != TSStatusCode.SUCCESS_STATUS.value: + return TShowLoadedModelsResp(status=status, deviceLoadedModelsMap={}) + return self._inference_manager.show_loaded_models(req) + + def showAIDevices(self) -> TShowAIDevicesResp: + return TShowAIDevicesResp( + status=TSStatus(code=TSStatusCode.SUCCESS_STATUS.value), + deviceIdList=get_available_devices(), + ) + def createTrainingTask(self, req: TTrainingReq) -> TSStatus: pass + + def _ensure_model_is_built_in_or_fine_tuned(self, model_id: str) -> TSStatus: + if not self._model_manager.is_built_in_or_fine_tuned(model_id): + return TSStatus( + code=TSStatusCode.MODEL_NOT_FOUND_ERROR.value, + message=f"Model [{model_id}] is not a built-in or fine-tuned model. You can use 'SHOW MODELS' to retrieve the available models.", + ) + return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value) diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py index 2282ae237d8..fc93a0cce7d 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py @@ -19,12 +19,12 @@ import threading from thrift.protocol import TBinaryProtocol, TCompactProtocol from thrift.server import TServer -from thrift.transport import TSocket, TTransport +from thrift.transport import TSocket, TSSLSocket, TTransport from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.log import Logger -from ainode.core.rpc.handler import AINodeRPCServiceHandler -from ainode.thrift.ainode import IAINodeRPCService +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.rpc.handler import AINodeRPCServiceHandler +from iotdb.thrift.ainode import IAINodeRPCService logger = Logger() @@ -73,10 +73,33 @@ class AINodeRPCService(threading.Thread): self._stop_event = threading.Event() self._handler = handler processor = IAINodeRPCService.Processor(handler=self._handler) - transport = TSocket.TServerSocket( - host=AINodeDescriptor().get_config().get_ain_rpc_address(), - port=AINodeDescriptor().get_config().get_ain_rpc_port(), - ) + if AINodeDescriptor().get_config().get_ain_internal_ssl_enabled(): + import ssl + import sys + + if sys.version_info >= (3, 10): + context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) + else: + context = ssl.SSLContext(ssl.PROTOCOL_TLS) + context.verify_mode = ssl.CERT_REQUIRED + context.check_hostname = False + context.load_verify_locations( + cafile=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file() + ) + context.load_cert_chain( + certfile=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(), + keyfile=AINodeDescriptor().get_config().get_ain_thrift_ssl_key_file(), + ) + transport = TSSLSocket.TSSLServerSocket( + host=AINodeDescriptor().get_config().get_ain_rpc_address(), + port=AINodeDescriptor().get_config().get_ain_rpc_port(), + ssl_context=context, + ) + else: + transport = TSocket.TServerSocket( + host=AINodeDescriptor().get_config().get_ain_rpc_address(), + port=AINodeDescriptor().get_config().get_ain_rpc_port(), + ) transport_factory = TTransport.TFramedTransportFactory() if AINodeDescriptor().get_config().get_ain_thrift_compression_enabled(): protocol_factory = TCompactProtocol.TCompactProtocolFactory() diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/status.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/status.py index 37368b0068b..e2bacb53106 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/rpc/status.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/status.py @@ -16,9 +16,9 @@ # under the License. # -from ainode.core.constant import TSStatusCode -from ainode.core.log import Logger -from ainode.thrift.common.ttypes import TSStatus +from iotdb.ainode.core.constant import TSStatusCode +from iotdb.ainode.core.log import Logger +from iotdb.thrift.common.ttypes import TSStatus def get_status(status_code: TSStatusCode, message: str = None) -> TSStatus: diff --git a/iotdb-core/ainode/iotdb/ainode/core/script.py b/iotdb-core/ainode/iotdb/ainode/core/script.py index dd3efa11ac8..6e244e43ab6 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/script.py +++ b/iotdb-core/ainode/iotdb/ainode/core/script.py @@ -15,71 +15,27 @@ # specific language governing permissions and limitations # under the License. # -import os -import shutil +import multiprocessing import sys import torch.multiprocessing as mp from iotdb.ainode.core.ai_node import AINode -from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.constant import TSStatusCode -from iotdb.ainode.core.exception import MissingConfigError -from ainode.core.log import Logger -from ainode.core.rpc.client import ClientManager -from ainode.thrift.common.ttypes import ( - TAINodeLocation, - TEndPoint, -) +from iotdb.ainode.core.log import Logger logger = Logger() -def remove_ainode(arguments): - # Delete the current node - if len(arguments) == 2: - target_ainode_id = AINodeDescriptor().get_config().get_ainode_id() - target_rpc_address = AINodeDescriptor().get_config().get_ain_rpc_address() - target_rpc_port = AINodeDescriptor().get_config().get_ain_rpc_port() - - # Delete the node with a given id - elif len(arguments) == 3: - target_ainode_id = int(arguments[2]) - ainode_configuration_map = ( - ClientManager() - .borrow_config_node_client() - .get_ainode_configuration(target_ainode_id) - ) - - end_point = ainode_configuration_map[target_ainode_id].location.internalEndPoint - target_rpc_address = end_point.ip - target_rpc_port = end_point.port - - if not end_point: - raise MissingConfigError( - "NodeId: {} not found in cluster ".format(target_ainode_id) - ) - - logger.info("Got target AINode id: {}".format(target_ainode_id)) - - else: - raise MissingConfigError("Invalid command") - - location = TAINodeLocation( - target_ainode_id, TEndPoint(target_rpc_address, target_rpc_port) - ) - status = ClientManager().borrow_config_node_client().node_remove(location) - - if status.code == TSStatusCode.SUCCESS_STATUS.get_status_code(): - logger.info("IoTDB-AINode has successfully removed.") - if os.path.exists(AINodeDescriptor().get_config().get_ain_models_dir()): - shutil.rmtree(AINodeDescriptor().get_config().get_ain_models_dir()) - - def main(): + # Handle PyInstaller: filter out Python arguments that might be passed to subprocesses + # These arguments are not needed in frozen executables and cause warnings + # Note: This filtering should happen AFTER freeze_support() has handled child processes + if getattr(sys, "frozen", False): + python_args_to_filter = ["-I", "-B", "-S", "-E", "-O", "-OO"] + sys.argv = [arg for arg in sys.argv if arg not in python_args_to_filter] + + logger.info(f"Starting IoTDB-AINode process with sys argv {sys.argv}.") arguments = sys.argv - # load config - AINodeDescriptor() if len(arguments) == 1: logger.info("Command line argument must be specified.") return @@ -88,23 +44,26 @@ def main(): try: mp.set_start_method("spawn", force=True) logger.info(f"Current multiprocess start method: {mp.get_start_method()}") - logger.info("IoTDB-AINode is starting...") + logger.info("IoTDB AINode is starting...") ai_node = AINode() ai_node.start() except Exception as e: - logger.error("Start AINode failed, because of: {}".format(e)) - sys.exit(1) - # TODO: remove the following function, and add a destroy script - elif command == "remove": - try: - logger.info("Removing AINode...") - remove_ainode(arguments) - except Exception as e: - logger.error("Remove AINode failed, because of: {}".format(e)) + logger.error("Start IoTDB AINode failed, because of: {}".format(e)) sys.exit(1) else: logger.warning("Unknown argument: {}.".format(command)) if __name__ == "__main__": + # PyInstaller multiprocessing support + # freeze_support() is essential for PyInstaller frozen executables on all platforms + # It detects if the current process is a multiprocessing child process + # If it is, it executes the child process target function and exits + # If it's not, it returns immediately and continues with main() execution + # This prevents child processes from executing the main application logic + if getattr(sys, "frozen", False): + # Call freeze_support() for both standard multiprocessing and torch.multiprocessing + multiprocessing.freeze_support() + mp.freeze_support() + main() diff --git a/iotdb-core/ainode/iotdb/ainode/core/util/cache.py b/iotdb-core/ainode/iotdb/ainode/core/util/cache.py index 40e2b8d5771..122873df67b 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/util/cache.py +++ b/iotdb-core/ainode/iotdb/ainode/core/util/cache.py @@ -19,7 +19,7 @@ import sys from collections import OrderedDict from iotdb.ainode.core.config import AINodeDescriptor -from ainode.core.util.decorator import singleton +from iotdb.ainode.core.util.decorator import singleton def _estimate_size_in_byte(obj): diff --git a/iotdb-core/ainode/iotdb/ainode/core/util/decorator.py b/iotdb-core/ainode/iotdb/ainode/core/util/decorator.py index 33b9f4835ac..5a84c3d6bb2 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/util/decorator.py +++ b/iotdb-core/ainode/iotdb/ainode/core/util/decorator.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. # +from functools import wraps + + def singleton(cls): instances = {} @@ -24,3 +27,15 @@ def singleton(cls): return instances[cls] return get_instance + + +def synchronized(lock): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + with lock: + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/iotdb-core/ainode/iotdb/sqlalchemy/IoTDBSQLCompiler.py b/iotdb-core/ainode/iotdb/sqlalchemy/IoTDBSQLCompiler.py index 008a314e683..94e650ba66c 100644 --- a/iotdb-core/ainode/iotdb/sqlalchemy/IoTDBSQLCompiler.py +++ b/iotdb-core/ainode/iotdb/sqlalchemy/IoTDBSQLCompiler.py @@ -16,9 +16,8 @@ # under the License. # -from sqlalchemy.sql.compiler import SQLCompiler -from sqlalchemy.sql.compiler import OPERATORS from sqlalchemy.sql import operators +from sqlalchemy.sql.compiler import OPERATORS, SQLCompiler class IoTDBSQLCompiler(SQLCompiler): diff --git a/iotdb-core/ainode/iotdb/table_session_pool.py b/iotdb-core/ainode/iotdb/table_session_pool.py index f44df3f249b..bd5a84be5c7 100644 --- a/iotdb-core/ainode/iotdb/table_session_pool.py +++ b/iotdb-core/ainode/iotdb/table_session_pool.py @@ -16,7 +16,7 @@ # under the License. # from iotdb.Session import Session -from iotdb.SessionPool import SessionPool, PoolConfig +from iotdb.SessionPool import PoolConfig, SessionPool from iotdb.table_session import TableSession diff --git a/iotdb-core/ainode/iotdb/template/MeasurementNode.py b/iotdb-core/ainode/iotdb/template/MeasurementNode.py index 06cd7c6957e..865e17bdda9 100644 --- a/iotdb-core/ainode/iotdb/template/MeasurementNode.py +++ b/iotdb-core/ainode/iotdb/template/MeasurementNode.py @@ -17,9 +17,10 @@ # import warnings -from iotdb.utils.IoTDBConstants import TSDataType, TSEncoding, Compressor -from .TemplateNode import TemplateNode +from iotdb.utils.IoTDBConstants import Compressor, TSDataType, TSEncoding + from ..tsfile.utils.read_write_io_utils import ReadWriteUtils +from .TemplateNode import TemplateNode warnings.simplefilter("always", DeprecationWarning) diff --git a/iotdb-core/ainode/iotdb/template/Template.py b/iotdb-core/ainode/iotdb/template/Template.py index 98a6082f148..99c97515424 100644 --- a/iotdb-core/ainode/iotdb/template/Template.py +++ b/iotdb-core/ainode/iotdb/template/Template.py @@ -19,10 +19,10 @@ import struct import warnings -from .TemplateNode import TemplateNode from ..tsfile.common.constant.tsfile_constant import TsFileConstant from ..tsfile.utils.pair import Pair from ..tsfile.utils.read_write_io_utils import ReadWriteUtils +from .TemplateNode import TemplateNode warnings.simplefilter("always", DeprecationWarning) diff --git a/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService-remote b/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService-remote index 082e8f84511..4b17e2f82dc 100755 --- a/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService-remote +++ b/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService-remote @@ -7,14 +7,16 @@ # options string: py # -import sys import pprint +import sys + if sys.version_info[0] > 2: from urllib.parse import urlparse else: from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient + from thrift.protocol.TBinaryProtocol import TBinaryProtocol +from thrift.transport import THttpClient, TSocket, TSSLSocket, TTransport from iotdb.thrift.ainode import IAINodeRPCService from iotdb.thrift.ainode.ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService.py b/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService.py index 7db6b3bee65..f2231800d4d 100644 --- a/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService.py +++ b/iotdb-core/ainode/iotdb/thrift/ainode/IAINodeRPCService.py @@ -6,15 +6,23 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import logging +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TProcessor, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys -import logging from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/ainode/constants.py b/iotdb-core/ainode/iotdb/thrift/ainode/constants.py index 69c181ade38..be9fb9f77b3 100644 --- a/iotdb-core/ainode/iotdb/thrift/ainode/constants.py +++ b/iotdb-core/ainode/iotdb/thrift/ainode/constants.py @@ -6,9 +6,16 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.TRecursive import fix_spec -import sys from .ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/ainode/ttypes.py b/iotdb-core/ainode/iotdb/thrift/ainode/ttypes.py index a016b467de1..b1ad583cd64 100644 --- a/iotdb-core/ainode/iotdb/thrift/ainode/ttypes.py +++ b/iotdb-core/ainode/iotdb/thrift/ainode/ttypes.py @@ -6,14 +6,21 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys import iotdb.thrift.common.ttypes -from thrift.transport import TTransport all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/common/constants.py b/iotdb-core/ainode/iotdb/thrift/common/constants.py index 69c181ade38..be9fb9f77b3 100644 --- a/iotdb-core/ainode/iotdb/thrift/common/constants.py +++ b/iotdb-core/ainode/iotdb/thrift/common/constants.py @@ -6,9 +6,16 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.TRecursive import fix_spec -import sys from .ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/common/ttypes.py b/iotdb-core/ainode/iotdb/thrift/common/ttypes.py index 29d58f4dd4b..b55c864d2a0 100644 --- a/iotdb-core/ainode/iotdb/thrift/common/ttypes.py +++ b/iotdb-core/ainode/iotdb/thrift/common/ttypes.py @@ -6,13 +6,19 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException -from thrift.protocol.TProtocol import TProtocolException -from thrift.TRecursive import fix_spec - import sys +from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.transport import TTransport +from thrift.TRecursive import fix_spec + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService-remote b/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService-remote index 014d1635657..74b5a660ecf 100755 --- a/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService-remote +++ b/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService-remote @@ -7,14 +7,16 @@ # options string: py # -import sys import pprint +import sys + if sys.version_info[0] > 2: from urllib.parse import urlparse else: from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient + from thrift.protocol.TBinaryProtocol import TBinaryProtocol +from thrift.transport import THttpClient, TSocket, TSSLSocket, TTransport from iotdb.thrift.confignode import IConfigNodeRPCService from iotdb.thrift.confignode.ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService.py b/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService.py index 439190815c4..a445bbe647a 100644 --- a/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService.py +++ b/iotdb-core/ainode/iotdb/thrift/confignode/IConfigNodeRPCService.py @@ -6,15 +6,23 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import logging +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TProcessor, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys -import logging from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/confignode/constants.py b/iotdb-core/ainode/iotdb/thrift/confignode/constants.py index 69c181ade38..be9fb9f77b3 100644 --- a/iotdb-core/ainode/iotdb/thrift/confignode/constants.py +++ b/iotdb-core/ainode/iotdb/thrift/confignode/constants.py @@ -6,9 +6,16 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.TRecursive import fix_spec -import sys from .ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/confignode/ttypes.py b/iotdb-core/ainode/iotdb/thrift/confignode/ttypes.py index 9a8d9c4b4e4..e5197c428a6 100644 --- a/iotdb-core/ainode/iotdb/thrift/confignode/ttypes.py +++ b/iotdb-core/ainode/iotdb/thrift/confignode/ttypes.py @@ -6,14 +6,21 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys import iotdb.thrift.common.ttypes -from thrift.transport import TTransport all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService-remote b/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService-remote index c37821f4d78..7381e55e0ea 100755 --- a/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService-remote +++ b/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService-remote @@ -7,14 +7,16 @@ # options string: py # -import sys import pprint +import sys + if sys.version_info[0] > 2: from urllib.parse import urlparse else: from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient + from thrift.protocol.TBinaryProtocol import TBinaryProtocol +from thrift.transport import THttpClient, TSocket, TSSLSocket, TTransport from iotdb.thrift.datanode import IDataNodeRPCService from iotdb.thrift.datanode.ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService.py b/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService.py index ca3c29ddd8f..09782524fe1 100644 --- a/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService.py +++ b/iotdb-core/ainode/iotdb/thrift/datanode/IDataNodeRPCService.py @@ -6,15 +6,23 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import logging +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TProcessor, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys -import logging from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService-remote b/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService-remote index b1fbd493625..3fd163e7830 100755 --- a/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService-remote +++ b/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService-remote @@ -7,14 +7,16 @@ # options string: py # -import sys import pprint +import sys + if sys.version_info[0] > 2: from urllib.parse import urlparse else: from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient + from thrift.protocol.TBinaryProtocol import TBinaryProtocol +from thrift.transport import THttpClient, TSocket, TSSLSocket, TTransport from iotdb.thrift.datanode import MPPDataExchangeService from iotdb.thrift.datanode.ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService.py b/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService.py index ce970090b32..cd0248990b9 100644 --- a/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService.py +++ b/iotdb-core/ainode/iotdb/thrift/datanode/MPPDataExchangeService.py @@ -6,15 +6,23 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import logging +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TProcessor, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys -import logging from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/constants.py b/iotdb-core/ainode/iotdb/thrift/datanode/constants.py index 69c181ade38..be9fb9f77b3 100644 --- a/iotdb-core/ainode/iotdb/thrift/datanode/constants.py +++ b/iotdb-core/ainode/iotdb/thrift/datanode/constants.py @@ -6,9 +6,16 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.TRecursive import fix_spec -import sys from .ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/datanode/ttypes.py b/iotdb-core/ainode/iotdb/thrift/datanode/ttypes.py index b1d99ce48e3..53217e2d534 100644 --- a/iotdb-core/ainode/iotdb/thrift/datanode/ttypes.py +++ b/iotdb-core/ainode/iotdb/thrift/datanode/ttypes.py @@ -6,15 +6,22 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys import iotdb.thrift.common.ttypes import iotdb.thrift.rpc.ttypes -from thrift.transport import TTransport all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService-remote b/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService-remote index 45126506e6f..8ea0655ce88 100755 --- a/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService-remote +++ b/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService-remote @@ -7,14 +7,16 @@ # options string: py # -import sys import pprint +import sys + if sys.version_info[0] > 2: from urllib.parse import urlparse else: from urlparse import urlparse -from thrift.transport import TTransport, TSocket, TSSLSocket, THttpClient + from thrift.protocol.TBinaryProtocol import TBinaryProtocol +from thrift.transport import THttpClient, TSocket, TSSLSocket, TTransport from iotdb.thrift.rpc import IClientRPCService from iotdb.thrift.rpc.ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService.py b/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService.py index 647d46a7995..f80b2e33219 100644 --- a/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService.py +++ b/iotdb-core/ainode/iotdb/thrift/rpc/IClientRPCService.py @@ -6,15 +6,23 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import logging +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TProcessor, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys -import logging from .ttypes import * -from thrift.Thrift import TProcessor -from thrift.transport import TTransport + all_structs = [] diff --git a/iotdb-core/ainode/iotdb/thrift/rpc/constants.py b/iotdb-core/ainode/iotdb/thrift/rpc/constants.py index 69c181ade38..be9fb9f77b3 100644 --- a/iotdb-core/ainode/iotdb/thrift/rpc/constants.py +++ b/iotdb-core/ainode/iotdb/thrift/rpc/constants.py @@ -6,9 +6,16 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) from thrift.TRecursive import fix_spec -import sys from .ttypes import * diff --git a/iotdb-core/ainode/iotdb/thrift/rpc/ttypes.py b/iotdb-core/ainode/iotdb/thrift/rpc/ttypes.py index 66ae893696f..df3767c6fdd 100644 --- a/iotdb-core/ainode/iotdb/thrift/rpc/ttypes.py +++ b/iotdb-core/ainode/iotdb/thrift/rpc/ttypes.py @@ -6,14 +6,21 @@ # options string: py # -from thrift.Thrift import TType, TMessageType, TFrozenDict, TException, TApplicationException +import sys + from thrift.protocol.TProtocol import TProtocolException +from thrift.Thrift import ( + TApplicationException, + TException, + TFrozenDict, + TMessageType, + TType, +) +from thrift.transport import TTransport from thrift.TRecursive import fix_spec -import sys import iotdb.thrift.common.ttypes -from thrift.transport import TTransport all_structs = [] diff --git a/iotdb-core/ainode/iotdb/tsfile/utils/read_write_io_utils.py b/iotdb-core/ainode/iotdb/tsfile/utils/read_write_io_utils.py index f6a5c02f235..11782a55807 100644 --- a/iotdb-core/ainode/iotdb/tsfile/utils/read_write_io_utils.py +++ b/iotdb-core/ainode/iotdb/tsfile/utils/read_write_io_utils.py @@ -16,7 +16,7 @@ # under the License. # -from iotdb.utils.IoTDBConstants import TSDataType, TSEncoding, Compressor +from iotdb.utils.IoTDBConstants import Compressor, TSDataType, TSEncoding class ReadWriteUtils: diff --git a/iotdb-core/ainode/iotdb/tsfile/utils/tsblock_serde.py b/iotdb-core/ainode/iotdb/tsfile/utils/tsblock_serde.py index 854be120ffb..40801f96ad7 100644 --- a/iotdb-core/ainode/iotdb/tsfile/utils/tsblock_serde.py +++ b/iotdb-core/ainode/iotdb/tsfile/utils/tsblock_serde.py @@ -18,7 +18,6 @@ import numpy as np - # Serialized tsBlock: # +-------------+---------------+---------+------------+-----------+----------+ # | val col cnt | val col types | pos cnt | encodings | time col | val col | diff --git a/iotdb-core/ainode/iotdb/utils/Field.py b/iotdb-core/ainode/iotdb/utils/Field.py index d9a0ee77776..6f2d9123e52 100644 --- a/iotdb-core/ainode/iotdb/utils/Field.py +++ b/iotdb-core/ainode/iotdb/utils/Field.py @@ -16,12 +16,14 @@ # under the License. # +import numpy as np +import pandas as pd + +from iotdb.tsfile.utils.date_utils import parse_int_to_date + # for package from iotdb.utils.IoTDBConstants import TSDataType -from iotdb.tsfile.utils.date_utils import parse_int_to_date from iotdb.utils.rpc_utils import convert_to_timestamp, isoformat -import numpy as np -import pandas as pd class Field(object): diff --git a/iotdb-core/ainode/iotdb/utils/IoTDBConstants.py b/iotdb-core/ainode/iotdb/utils/IoTDBConstants.py index 4b9082b5353..e51172e0962 100644 --- a/iotdb-core/ainode/iotdb/utils/IoTDBConstants.py +++ b/iotdb-core/ainode/iotdb/utils/IoTDBConstants.py @@ -16,7 +16,8 @@ # under the License. # from datetime import date -from enum import unique, IntEnum +from enum import IntEnum, unique + import numpy as np diff --git a/iotdb-core/ainode/iotdb/utils/NumpyTablet.py b/iotdb-core/ainode/iotdb/utils/NumpyTablet.py index 9a0860d3a43..65b76ffe4aa 100644 --- a/iotdb-core/ainode/iotdb/utils/NumpyTablet.py +++ b/iotdb-core/ainode/iotdb/utils/NumpyTablet.py @@ -17,14 +17,14 @@ # import struct +from typing import List import numpy as np from numpy import ndarray -from typing import List from iotdb.tsfile.utils.date_utils import parse_date_to_int -from iotdb.utils.IoTDBConstants import TSDataType from iotdb.utils.BitMap import BitMap +from iotdb.utils.IoTDBConstants import TSDataType from iotdb.utils.Tablet import ColumnType diff --git a/iotdb-core/ainode/iotdb/utils/SessionDataSet.py b/iotdb-core/ainode/iotdb/utils/SessionDataSet.py index b079ee28c1a..2139e94b9a8 100644 --- a/iotdb-core/ainode/iotdb/utils/SessionDataSet.py +++ b/iotdb-core/ainode/iotdb/utils/SessionDataSet.py @@ -17,15 +17,15 @@ # import logging +import pandas as pd + from iotdb.utils.Field import Field +from iotdb.utils.iotdb_rpc_dataset import IoTDBRpcDataSet # for package from iotdb.utils.IoTDBConstants import TSDataType -from iotdb.utils.iotdb_rpc_dataset import IoTDBRpcDataSet from iotdb.utils.RowRecord import RowRecord -import pandas as pd - logger = logging.getLogger("IoTDB") diff --git a/iotdb-core/ainode/iotdb/utils/Tablet.py b/iotdb-core/ainode/iotdb/utils/Tablet.py index 9b241723fe5..ce2b2c2a5ae 100644 --- a/iotdb-core/ainode/iotdb/utils/Tablet.py +++ b/iotdb-core/ainode/iotdb/utils/Tablet.py @@ -17,7 +17,7 @@ # import struct -from enum import unique, IntEnum +from enum import IntEnum, unique from typing import List, Union from iotdb.tsfile.utils.date_utils import parse_date_to_int diff --git a/iotdb-core/ainode/iotdb/utils/iotdb_rpc_dataset.py b/iotdb-core/ainode/iotdb/utils/iotdb_rpc_dataset.py index dc771368669..cc35c429b3c 100644 --- a/iotdb-core/ainode/iotdb/utils/iotdb_rpc_dataset.py +++ b/iotdb-core/ainode/iotdb/utils/iotdb_rpc_dataset.py @@ -23,12 +23,12 @@ import numpy as np import pandas as pd from thrift.transport import TTransport -from iotdb.thrift.rpc.IClientRPCService import TSFetchResultsReq, TSCloseOperationReq +from iotdb.thrift.rpc.IClientRPCService import TSCloseOperationReq, TSFetchResultsReq from iotdb.tsfile.utils.date_utils import parse_int_to_date from iotdb.tsfile.utils.tsblock_serde import deserialize from iotdb.utils.exception import IoTDBConnectionException from iotdb.utils.IoTDBConstants import TSDataType -from iotdb.utils.rpc_utils import verify_success, convert_to_timestamp +from iotdb.utils.rpc_utils import convert_to_timestamp, verify_success logger = logging.getLogger("IoTDB") TIMESTAMP_STR = "Time" diff --git a/iotdb-core/ainode/pom.xml b/iotdb-core/ainode/pom.xml index c59c719c747..e9353db65a3 100644 --- a/iotdb-core/ainode/pom.xml +++ b/iotdb-core/ainode/pom.xml @@ -24,7 +24,7 @@ <parent> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-core</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> </parent> <artifactId>iotdb-ainode</artifactId> <name>IoTDB: Core: AINode</name> @@ -33,31 +33,31 @@ <dependency> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-thrift</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> <scope>provided</scope> </dependency> <dependency> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-thrift-confignode</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> <scope>provided</scope> </dependency> <dependency> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-thrift-commons</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> <scope>provided</scope> </dependency> <dependency> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-thrift-ainode</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> <scope>provided</scope> </dependency> <dependency> <groupId>org.apache.iotdb</groupId> <artifactId>iotdb-python-api</artifactId> - <version>2.0.6-SNAPSHOT</version> + <version>2.0.6</version> <scope>provided</scope> </dependency> </dependencies> @@ -73,7 +73,7 @@ <directory>dist</directory> </fileset> <fileset> - <directory>ainode</directory> + <directory>iotdb</directory> <includes> <include>ainode/conf/</include> <include>dbapi/</include> diff --git a/iotdb-core/ainode/pyproject.toml b/iotdb-core/ainode/pyproject.toml index 54aacf08768..151eeebb757 100644 --- a/iotdb-core/ainode/pyproject.toml +++ b/iotdb-core/ainode/pyproject.toml @@ -21,7 +21,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "apache-iotdb-ainode" -version = "2.0.6.dev" +version = "2.0.6" description = "Apache IoTDB AINode" readme = "README.md" authors = ["Apache Software Foundation <[email protected]>"] @@ -34,45 +34,49 @@ classifiers = [ "Topic :: Software Development :: Libraries :: Python Modules", ] include = [ - {path = "iotdb/*", format = "wheel"}, - {path = "iotdb/dbapi/*", format = "wheel"}, - {path = "iotdb/sqlalchemy/*", format = "wheel"}, - {path = "iotdb/template/*", format = "wheel"}, - {path = "iotdb/thrift/*", format = "wheel"}, - {path = "iotdb/thrift/ainode/*", format = "wheel"}, - {path = "iotdb/thrift/common/*", format = "wheel"}, - {path = "iotdb/thrift/confignode/*", format = "wheel"}, - {path = "iotdb/thrift/datanode/*", format = "wheel"}, - {path = "iotdb/thrift/rpc/*", format = "wheel"}, - {path = "iotdb/tsfile/*", format = "wheel"}, - {path = "iotdb/tsfile/common/*", format = "wheel"}, - {path = "iotdb/tsfile/common/constant/*", format = "wheel"}, - {path = "iotdb/tsfile/utils/*", format = "wheel"}, - {path = "iotdb/utils/*", format = "wheel"}, + { path = "iotdb/Session.py", format = "wheel" }, + { path = "iotdb/SessionPool.py", format = "wheel" }, + { path = "iotdb/table_session.py", format = "wheel" }, + { path = "iotdb/table_session_pool.py", format = "wheel" }, + { path = "iotdb/dbapi/*", format = "wheel" }, + { path = "iotdb/sqlalchemy/*", format = "wheel" }, + { path = "iotdb/template/*", format = "wheel" }, + { path = "iotdb/thrift/*", format = "wheel" }, + { path = "iotdb/thrift/ainode/*", format = "wheel" }, + { path = "iotdb/thrift/common/*", format = "wheel" }, + { path = "iotdb/thrift/confignode/*", format = "wheel" }, + { path = "iotdb/thrift/datanode/*", format = "wheel" }, + { path = "iotdb/thrift/rpc/*", format = "wheel" }, + { path = "iotdb/tsfile/*", format = "wheel" }, + { path = "iotdb/tsfile/common/*", format = "wheel" }, + { path = "iotdb/tsfile/common/constant/*", format = "wheel" }, + { path = "iotdb/tsfile/utils/*", format = "wheel" }, + { path = "iotdb/utils/*", format = "wheel" }, + { path = "iotdb/__init__.py", format = "wheel" }, + { path = "iotdb/ainode/__init__.py", format = "wheel" }, + { path = "iotdb/ainode/conf/*", format = "wheel" }, ] + packages = [ - { include = "ainode" } + { include = "iotdb/ainode/core" } ] -[tool.poetry.dependencies] -python = ">=3.11,<=3.13.5" - -# Core scientific stack -numpy = [ - { version = "^2.3.2", python = ">=3.10" }, - { version = "^1.26.4", python = ">=3.9,<3.10" } +exclude = [ + "**/__pycache__/**", + "**/*.pyc", "**/*.pyo", + "**/.DS_Store", + "**/*.log", "**/*.log.gz", + "venv/**", ".venv/**", "*/.venv/**", + "iotdb/ainode/core/data/**", + "iotdb/ainode/core/logs/**", + "iotdb/**/.cache/**", + "iotdb/**/download/**", + "iotdb/**/logs/**", + "iotdb/**/log-*.log*" ] -scipy = [ - { version = "^1.12.0", python = ">=3.10" }, - { version = "^1.11.4", python = ">=3.9,<3.10" } -] -pandas = "^2.3.2" -scikit-learn = [ - { version = "^1.7.1", python = ">=3.10" }, - { version = "^1.5.2", python = ">=3.9,<3.10" } -] -statsmodels = "^0.14.5" -sktime = "0.38.5" + +[tool.poetry.dependencies] +python = ">=3.11.0,<3.14.0" # ---- DL / HF stack ---- torch = ">=2.7.0" @@ -83,32 +87,33 @@ huggingface_hub = "^0.34.4" safetensors = "^0.6.2" einops = "^0.8.1" +# ---- Core scientific stack ---- +numpy = "^2.3.2" +scipy = "^1.12.0" +pandas = "^2.3.2" +scikit-learn = "^1.7.1" +statsmodels = "^0.14.5" +sktime = "0.38.5" + # ---- Optimizers / utils ---- optuna = "^4.4.0" psutil = "^7.0.0" requests = "^2.32.5" dynaconf = "^3.2.11" thrift = ">=0.14.0" -dynaconf = "^3.1.11" -requests = "^2.31.0" -optuna = "^3.2.0" -psutil = "^5.9.5" -sktime = "0.38.4" -scikit-learn = "^1.4" -statsmodels = "^0.14" -hmmlearn = "^0.3.0" -einops = "^0.8.1" -safetensors = "^0.5.1" -huggingface_hub = "^0.30.1" -tokenizers = ">=0.19" -hf_xet = ">=1.1.7" +tzlocal = "^5.3.1" + +# ---- HF acceleration ---- +hf_xet = ">=1.1.9" + +# ---- Tooling ---- +pyinstaller = "6.16.0" black = "25.1.0" isort = "6.0.1" -transformers = "==4.40.1" -torchmetrics = ">=1.2.0" [tool.poetry.scripts] -ainode = "ainode.core.script:main" +ainode = "iotdb.ainode.core.script:main" [tool.isort] -profile = "black" \ No newline at end of file +profile = "black" +skip = ["build", "dist", "target"] \ No newline at end of file diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java index 3f5ee055649..3b09a568836 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java @@ -27,7 +27,6 @@ import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation; import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.commons.audit.UserEntity; import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.commons.exception.IoTDBRuntimeException; import org.apache.iotdb.commons.exception.auth.AccessDeniedException;
