You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by yo...@apache.org on 2025/10/16 02:07:55 UTC
(iotdb) 01/02: [AINode] Code format
This is an automated email from the ASF dual-hosted git repository.
yongzao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit a1c34887c3b732767d66f2784f450cd50780bd94
Author: Yongzao <53...@qq.com>
AuthorDate: Thu Oct 16 09:54:43 2025 +0800
[AINode] Code format
---
iotdb-core/ainode/iotdb/ainode/__init__.py | 35 +-
iotdb-core/ainode/iotdb/ainode/core/__init__.py | 35 +-
iotdb-core/ainode/iotdb/ainode/core/ai_node.py | 339 ++--
iotdb-core/ainode/iotdb/ainode/core/config.py | 983 +++++-----
iotdb-core/ainode/iotdb/ainode/core/constant.py | 549 +++---
iotdb-core/ainode/iotdb/ainode/core/exception.py | 309 +--
.../ainode/iotdb/ainode/core/inference/__init__.py | 35 +-
.../{ => core/inference/batcher}/__init__.py | 2 +-
.../core/inference/batcher/abstract_batcher.py | 85 +-
.../ainode/core/inference/batcher/basic_batcher.py | 127 +-
.../ainode/core/inference/dispatcher/__init__.py | 35 +-
.../inference/dispatcher/abstract_dispatcher.py | 87 +-
.../core/inference/dispatcher/basic_dispatcher.py | 121 +-
.../ainode/core/inference/inference_request.py | 265 +--
.../core/inference/inference_request_pool.py | 435 ++---
.../iotdb/ainode/core/inference/pool_controller.py | 1019 +++++-----
.../iotdb/ainode/core/inference/pool_group.py | 231 +--
.../core/inference/pool_scheduler/__init__.py | 35 +-
.../pool_scheduler/abstract_pool_scheduler.py | 171 +-
.../pool_scheduler/basic_pool_scheduler.py | 383 ++--
.../core/inference/request_scheduler/__init__.py | 35 +-
.../abstract_request_scheduler.py | 125 +-
.../request_scheduler/basic_request_scheduler.py | 177 +-
.../ainode/core/inference/strategy/__init__.py | 35 +-
.../strategy/abstract_inference_pipeline.py | 121 +-
.../strategy/timer_sundial_inference_pipeline.py | 103 +-
.../strategy/timerxl_inference_pipeline.py | 103 +-
.../ainode/iotdb/ainode/core/inference/utils.py | 161 +-
.../ainode/iotdb/ainode/core/ingress/__init__.py | 35 +-
.../ainode/iotdb/ainode/core/ingress/dataset.py | 93 +-
.../ainode/iotdb/ainode/core/ingress/iotdb.py | 803 ++++----
iotdb-core/ainode/iotdb/ainode/core/log.py | 323 ++--
.../ainode/iotdb/ainode/core/manager/__init__.py | 35 +-
.../iotdb/ainode/core/manager/cluster_manager.py | 93 +-
.../iotdb/ainode/core/manager/inference_manager.py | 745 ++++----
.../iotdb/ainode/core/manager/model_manager.py | 333 ++--
.../ainode/iotdb/ainode/core/manager/utils.py | 217 +--
.../ainode/iotdb/ainode/core/model/__init__.py | 35 +-
.../ainode/core/model/built_in_model_factory.py | 2000 ++++++++++----------
.../ainode/iotdb/ainode/core/model/model_enums.py | 141 +-
.../iotdb/ainode/core/model/model_factory.py | 583 +++---
.../ainode/iotdb/ainode/core/model/model_info.py | 309 +--
.../iotdb/ainode/core/model/model_storage.py | 899 ++++-----
.../iotdb/ainode/core/model/sundial/__init__.py | 35 +-
.../core/model/sundial/configuration_sundial.py | 131 +-
.../iotdb/ainode/core/model/sundial/flow_loss.py | 511 ++---
.../ainode/core/model/sundial/modeling_sundial.py | 1303 ++++++-------
.../core/model/sundial/ts_generation_mixin.py | 767 ++++----
.../iotdb/ainode/core/model/timerxl/__init__.py | 35 +-
.../core/model/timerxl/configuration_timer.py | 119 +-
.../ainode/core/model/timerxl/modeling_timer.py | 1281 ++++++-------
.../core/model/timerxl/ts_generation_mixin.py | 741 ++++----
.../ainode/iotdb/ainode/core/model/uri_utils.py | 275 +--
.../ainode/iotdb/ainode/core/rpc/__init__.py | 35 +-
iotdb-core/ainode/iotdb/ainode/core/rpc/client.py | 613 +++---
iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py | 271 +--
iotdb-core/ainode/iotdb/ainode/core/rpc/service.py | 255 +--
iotdb-core/ainode/iotdb/ainode/core/rpc/status.py | 67 +-
iotdb-core/ainode/iotdb/ainode/core/script.py | 215 +--
.../ainode/iotdb/ainode/core/util/__init__.py | 35 +-
.../ainode/iotdb/ainode/core/util/activation.py | 111 +-
.../ainode/iotdb/ainode/core/util/atmoic_int.py | 133 +-
.../iotdb/ainode/core/util/batch_executor.py | 267 +--
iotdb-core/ainode/iotdb/ainode/core/util/cache.py | 177 +-
.../ainode/iotdb/ainode/core/util/decorator.py | 83 +-
.../ainode/iotdb/ainode/core/util/gpu_mapping.py | 187 +-
.../iotdb/ainode/core/util/huggingface_cache.py | 455 ++---
iotdb-core/ainode/iotdb/ainode/core/util/lock.py | 169 +-
.../ainode/iotdb/ainode/core/util/masking.py | 187 +-
iotdb-core/ainode/iotdb/ainode/core/util/serde.py | 315 +--
.../ainode/iotdb/ainode/core/util/thread_name.py | 45 +-
71 files changed, 10836 insertions(+), 10767 deletions(-)
diff --git a/iotdb-core/ainode/iotdb/ainode/__init__.py b/iotdb-core/ainode/iotdb/ainode/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py
index 587ae4930de..ba9b269cdb4 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/ai_node.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py
@@ -1,169 +1,170 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import os
-import signal
-import threading
-from datetime import datetime
-
-import psutil
-
-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.constant import AINODE_SYSTEM_FILE_NAME
-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 iotdb.thrift.confignode.ttypes import TNodeVersionInfo
-
-logger = Logger()
-
-
-def _generate_configuration() -> TAINodeConfiguration:
- location = TAINodeLocation(
- AINodeDescriptor().get_config().get_ainode_id(),
- TEndPoint(
- AINodeDescriptor().get_config().get_ain_rpc_address(),
- AINodeDescriptor().get_config().get_ain_rpc_port(),
- ),
- )
- resource = TNodeResource(int(psutil.cpu_count()), int(psutil.virtual_memory()[0]))
-
- return TAINodeConfiguration(location, resource)
-
-
-def _generate_version_info() -> TNodeVersionInfo:
- return TNodeVersionInfo(
- AINodeDescriptor().get_config().get_version_info(),
- AINodeDescriptor().get_config().get_build_info(),
- )
-
-
-def _check_path_permission():
- system_path = AINodeDescriptor().get_config().get_ain_system_dir()
- if not os.path.exists(system_path):
- try:
- os.makedirs(system_path)
- except PermissionError as e:
- logger.error(e)
- raise e
-
-
-def _generate_system_properties(ainode_id: int):
- return {
- "ainode_id": ainode_id,
- "cluster_name": AINodeDescriptor().get_config().get_cluster_name(),
- "iotdb_version": AINodeDescriptor().get_config().get_version_info(),
- "commit_id": AINodeDescriptor().get_config().get_build_info(),
- "ain_rpc_address": AINodeDescriptor().get_config().get_ain_rpc_address(),
- "ain_rpc_port": AINodeDescriptor().get_config().get_ain_rpc_port(),
- "config_node_list": AINodeDescriptor()
- .get_config()
- .get_ain_target_config_node_list(),
- }
-
-
-class AINode:
- def __init__(self):
- self._rpc_service = None
- self._rpc_handler = None
- self._stop_event = None
-
- def start(self):
- _check_path_permission()
- system_properties_file = os.path.join(
- AINodeDescriptor().get_config().get_ain_system_dir(),
- AINODE_SYSTEM_FILE_NAME,
- )
- if not os.path.exists(system_properties_file):
- # If the system.properties file does not exist, the AINode will register to IoTDB cluster.
- try:
- logger.info("IoTDB-AINode is registering to IoTDB cluster...")
- ainode_id = (
- ClientManager()
- .borrow_config_node_client()
- .node_register(
- AINodeDescriptor().get_config().get_cluster_name(),
- _generate_configuration(),
- _generate_version_info(),
- )
- )
- AINodeDescriptor().get_config().set_ainode_id(ainode_id)
- system_properties = _generate_system_properties(ainode_id)
- with open(system_properties_file, "w") as f:
- f.write("#" + str(datetime.now()) + "\n")
- for key, value in system_properties.items():
- f.write(key + "=" + str(value) + "\n")
- except Exception as e:
- logger.error(
- "IoTDB-AINode failed to register to IoTDB cluster: {}".format(e)
- )
- raise e
- else:
- # If the system.properties file does exist, the AINode will just restart.
- try:
- logger.info("IoTDB-AINode is restarting...")
- ClientManager().borrow_config_node_client().node_restart(
- AINodeDescriptor().get_config().get_cluster_name(),
- _generate_configuration(),
- _generate_version_info(),
- )
- except Exception as e:
- logger.error("IoTDB-AINode failed to restart: {}".format(e))
- raise e
-
- # Start the RPC service
- self._rpc_handler = AINodeRPCServiceHandler(ainode=self)
- self._rpc_service = AINodeRPCService(self._rpc_handler)
- self._rpc_service.start()
- self._rpc_service.join(1)
- if self._rpc_service.exit_code != 0:
- logger.info("IoTDB-AINode failed to start, please check previous logs.")
- return
-
- logger.info("IoTDB-AINode has successfully started.")
-
- # Register stop hook
- self._stop_event = threading.Event()
- signal.signal(signal.SIGTERM, self._handle_signal)
-
- self._rpc_service.join()
-
- def _handle_signal(self, signum, frame):
- signal_name = {signal.SIGTERM: "SIGTERM", signal.SIGINT: "SIGINT"}.get(
- signum, f"SIGNAL {signum}"
- )
-
- logger.info(f"IoTDB-AINode receives {signal_name}, initiating graceful stop...")
- self.stop()
-
- def stop(self):
- if not self._stop_event.is_set():
- self._stop_event.set()
- self._rpc_handler.stop()
- if self._rpc_service:
- self._rpc_service.stop()
- self._rpc_service.join(1)
- if self._rpc_service.is_alive():
- logger.warning("RPC service thread failed to stop in time.")
- logger.info("IoTDB-AINode has successfully stopped.")
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import os
+1import signal
+1import threading
+1from datetime import datetime
+1
+1import psutil
+1
+1from iotdb.ainode.core.config import AINodeDescriptor
+1from iotdb.ainode.core.constant import AINODE_SYSTEM_FILE_NAME
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.rpc.client import ClientManager
+1from iotdb.ainode.core.rpc.handler import AINodeRPCServiceHandler
+1from iotdb.ainode.core.rpc.service import AINodeRPCService
+1from iotdb.thrift.common.ttypes import (
+1 TAINodeConfiguration,
+1 TAINodeLocation,
+1 TEndPoint,
+1 TNodeResource,
+1)
+1from iotdb.thrift.confignode.ttypes import TNodeVersionInfo
+1
+1logger = Logger()
+1
+1
+1def _generate_configuration() -> TAINodeConfiguration:
+1 location = TAINodeLocation(
+1 AINodeDescriptor().get_config().get_ainode_id(),
+1 TEndPoint(
+1 AINodeDescriptor().get_config().get_ain_rpc_address(),
+1 AINodeDescriptor().get_config().get_ain_rpc_port(),
+1 ),
+1 )
+1 resource = TNodeResource(int(psutil.cpu_count()), int(psutil.virtual_memory()[0]))
+1
+1 return TAINodeConfiguration(location, resource)
+1
+1
+1def _generate_version_info() -> TNodeVersionInfo:
+1 return TNodeVersionInfo(
+1 AINodeDescriptor().get_config().get_version_info(),
+1 AINodeDescriptor().get_config().get_build_info(),
+1 )
+1
+1
+1def _check_path_permission():
+1 system_path = AINodeDescriptor().get_config().get_ain_system_dir()
+1 if not os.path.exists(system_path):
+1 try:
+1 os.makedirs(system_path)
+1 except PermissionError as e:
+1 logger.error(e)
+1 raise e
+1
+1
+1def _generate_system_properties(ainode_id: int):
+1 return {
+1 "ainode_id": ainode_id,
+1 "cluster_name": AINodeDescriptor().get_config().get_cluster_name(),
+1 "iotdb_version": AINodeDescriptor().get_config().get_version_info(),
+1 "commit_id": AINodeDescriptor().get_config().get_build_info(),
+1 "ain_rpc_address": AINodeDescriptor().get_config().get_ain_rpc_address(),
+1 "ain_rpc_port": AINodeDescriptor().get_config().get_ain_rpc_port(),
+1 "config_node_list": AINodeDescriptor()
+1 .get_config()
+1 .get_ain_target_config_node_list(),
+1 }
+1
+1
+1class AINode:
+1 def __init__(self):
+1 self._rpc_service = None
+1 self._rpc_handler = None
+1 self._stop_event = None
+1
+1 def start(self):
+1 _check_path_permission()
+1 system_properties_file = os.path.join(
+1 AINodeDescriptor().get_config().get_ain_system_dir(),
+1 AINODE_SYSTEM_FILE_NAME,
+1 )
+1 if not os.path.exists(system_properties_file):
+1 # If the system.properties file does not exist, the AINode will register to IoTDB cluster.
+1 try:
+1 logger.info("IoTDB-AINode is registering to IoTDB cluster...")
+1 ainode_id = (
+1 ClientManager()
+1 .borrow_config_node_client()
+1 .node_register(
+1 AINodeDescriptor().get_config().get_cluster_name(),
+1 _generate_configuration(),
+1 _generate_version_info(),
+1 )
+1 )
+1 AINodeDescriptor().get_config().set_ainode_id(ainode_id)
+1 system_properties = _generate_system_properties(ainode_id)
+1 with open(system_properties_file, "w") as f:
+1 f.write("#" + str(datetime.now()) + "\n")
+1 for key, value in system_properties.items():
+1 f.write(key + "=" + str(value) + "\n")
+1 except Exception as e:
+1 logger.error(
+1 "IoTDB-AINode failed to register to IoTDB cluster: {}".format(e)
+1 )
+1 raise e
+1 else:
+1 # If the system.properties file does exist, the AINode will just restart.
+1 try:
+1 logger.info("IoTDB-AINode is restarting...")
+1 ClientManager().borrow_config_node_client().node_restart(
+1 AINodeDescriptor().get_config().get_cluster_name(),
+1 _generate_configuration(),
+1 _generate_version_info(),
+1 )
+1 except Exception as e:
+1 logger.error("IoTDB-AINode failed to restart: {}".format(e))
+1 raise e
+1
+1 # Start the RPC service
+1 self._rpc_handler = AINodeRPCServiceHandler(ainode=self)
+1 self._rpc_service = AINodeRPCService(self._rpc_handler)
+1 self._rpc_service.start()
+1 self._rpc_service.join(1)
+1 if self._rpc_service.exit_code != 0:
+1 logger.info("IoTDB-AINode failed to start, please check previous logs.")
+1 return
+1
+1 logger.info("IoTDB-AINode has successfully started.")
+1
+1 # Register stop hook
+1 self._stop_event = threading.Event()
+1 signal.signal(signal.SIGTERM, self._handle_signal)
+1
+1 self._rpc_service.join()
+1
+1 def _handle_signal(self, signum, frame):
+1 signal_name = {signal.SIGTERM: "SIGTERM", signal.SIGINT: "SIGINT"}.get(
+1 signum, f"SIGNAL {signum}"
+1 )
+1
+1 logger.info(f"IoTDB-AINode receives {signal_name}, initiating graceful stop...")
+1 self.stop()
+1
+1 def stop(self):
+1 if not self._stop_event.is_set():
+1 self._stop_event.set()
+1 self._rpc_handler.stop()
+1 if self._rpc_service:
+1 self._rpc_service.stop()
+1 self._rpc_service.join(1)
+1 if self._rpc_service.is_alive():
+1 logger.warning("RPC service thread failed to stop in time.")
+1 logger.info("IoTDB-AINode has successfully stopped.")
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/config.py b/iotdb-core/ainode/iotdb/ainode/core/config.py
index 328f4a5faa6..66f2321e438 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/config.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/config.py
@@ -1,491 +1,492 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import os
-
-from iotdb.ainode.core.constant import (
- AINODE_BUILD_INFO,
- AINODE_BUILTIN_MODELS_DIR,
- AINODE_CLUSTER_INGRESS_ADDRESS,
- AINODE_CLUSTER_INGRESS_PASSWORD,
- AINODE_CLUSTER_INGRESS_PORT,
- AINODE_CLUSTER_INGRESS_TIME_ZONE,
- AINODE_CLUSTER_INGRESS_USERNAME,
- AINODE_CLUSTER_NAME,
- AINODE_CONF_DIRECTORY_NAME,
- AINODE_CONF_FILE_NAME,
- AINODE_CONF_GIT_FILE_NAME,
- AINODE_CONF_POM_FILE_NAME,
- AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
- AINODE_INFERENCE_EXTRA_MEMORY_RATIO,
- AINODE_INFERENCE_MAX_PREDICT_LENGTH,
- AINODE_INFERENCE_MEMORY_USAGE_RATIO,
- AINODE_INFERENCE_MODEL_MEM_USAGE_MAP,
- AINODE_LOG_DIR,
- AINODE_MODELS_DIR,
- AINODE_ROOT_CONF_DIRECTORY_NAME,
- AINODE_ROOT_DIR,
- AINODE_RPC_ADDRESS,
- AINODE_RPC_PORT,
- AINODE_SYSTEM_DIR,
- AINODE_SYSTEM_FILE_NAME,
- AINODE_TARGET_CONFIG_NODE_LIST,
- AINODE_THRIFT_COMPRESSION_ENABLED,
- AINODE_VERSION_INFO,
-)
-from iotdb.ainode.core.exception import BadNodeUrlError
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.util.decorator import singleton
-from iotdb.thrift.common.ttypes import TEndPoint
-
-logger = Logger()
-
-
-class AINodeConfig(object):
- def __init__(self):
- self._version_info = AINODE_VERSION_INFO
- self._build_info = AINODE_BUILD_INFO
-
- # Cluster configuration
- self._ainode_id = 0
- self._cluster_name = AINODE_CLUSTER_NAME
- self._ain_target_config_node_list: TEndPoint = AINODE_TARGET_CONFIG_NODE_LIST
- self._ain_rpc_address: str = AINODE_RPC_ADDRESS
- self._ain_rpc_port: int = AINODE_RPC_PORT
- self._ain_cluster_ingress_address = AINODE_CLUSTER_INGRESS_ADDRESS
- self._ain_cluster_ingress_port = AINODE_CLUSTER_INGRESS_PORT
- self._ain_cluster_ingress_username = AINODE_CLUSTER_INGRESS_USERNAME
- self._ain_cluster_ingress_password = AINODE_CLUSTER_INGRESS_PASSWORD
- self._ain_cluster_ingress_time_zone = AINODE_CLUSTER_INGRESS_TIME_ZONE
-
- # Inference configuration
- self._ain_inference_batch_interval_in_ms: int = (
- AINODE_INFERENCE_BATCH_INTERVAL_IN_MS
- )
- self._ain_inference_max_predict_length: int = (
- AINODE_INFERENCE_MAX_PREDICT_LENGTH
- )
- self._ain_inference_model_mem_usage_map: dict[str, int] = (
- AINODE_INFERENCE_MODEL_MEM_USAGE_MAP
- )
- self._ain_inference_memory_usage_ratio: float = (
- AINODE_INFERENCE_MEMORY_USAGE_RATIO
- )
- self._ain_inference_extra_memory_ratio: float = (
- AINODE_INFERENCE_EXTRA_MEMORY_RATIO
- )
- # log directory
- self._ain_logs_dir: str = AINODE_LOG_DIR
-
- # cache size for ingress dataloader (MB)
- self._ain_data_cache_size = 50
-
- # Directory to save models
- self._ain_models_dir = AINODE_MODELS_DIR
- self._ain_builtin_models_dir = AINODE_BUILTIN_MODELS_DIR
- self._ain_system_dir = AINODE_SYSTEM_DIR
-
- # 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
-
- 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 get_version_info(self) -> str:
- return self._version_info
-
- def get_ainode_id(self) -> int:
- return self._ainode_id
-
- def set_ainode_id(self, ainode_id: int) -> None:
- self._ainode_id = ainode_id
-
- def get_build_info(self) -> str:
- return self._build_info
-
- def set_build_info(self, build_info: str) -> None:
- self._build_info = build_info
-
- def get_ain_data_storage_cache_size(self) -> int:
- return self._ain_data_cache_size
-
- def set_ain_data_cache_size(self, ain_data_cache_size: int) -> None:
- self._ain_data_cache_size = ain_data_cache_size
-
- def set_version_info(self, version_info: str) -> None:
- self._version_info = version_info
-
- def get_ain_rpc_address(self) -> str:
- return self._ain_rpc_address
-
- def set_ain_rpc_address(self, ain_rpc_address: str) -> None:
- self._ain_rpc_address = ain_rpc_address
-
- def get_ain_rpc_port(self) -> int:
- return self._ain_rpc_port
-
- def set_ain_rpc_port(self, ain_rpc_port: int) -> None:
- self._ain_rpc_port = ain_rpc_port
-
- def get_ain_inference_batch_interval_in_ms(self) -> int:
- return self._ain_inference_batch_interval_in_ms
-
- def set_ain_inference_batch_interval_in_ms(
- self, ain_inference_batch_interval_in_ms: int
- ) -> None:
- self._ain_inference_batch_interval_in_ms = ain_inference_batch_interval_in_ms
-
- def get_ain_inference_max_predict_length(self) -> int:
- return self._ain_inference_max_predict_length
-
- def set_ain_inference_max_predict_length(
- self, ain_inference_max_predict_length: int
- ) -> None:
- self._ain_inference_max_predict_length = ain_inference_max_predict_length
-
- def get_ain_inference_model_mem_usage_map(self) -> dict[str, int]:
- return self._ain_inference_model_mem_usage_map
-
- def set_ain_inference_model_mem_usage_map(
- self, ain_inference_model_mem_usage_map: dict[str, int]
- ) -> None:
- self._ain_inference_model_mem_usage_map = ain_inference_model_mem_usage_map
-
- def get_ain_inference_memory_usage_ratio(self) -> float:
- return self._ain_inference_memory_usage_ratio
-
- def set_ain_inference_memory_usage_ratio(
- self, ain_inference_memory_usage_ratio: float
- ) -> None:
- self._ain_inference_memory_usage_ratio = ain_inference_memory_usage_ratio
-
- def get_ain_inference_extra_memory_ratio(self) -> float:
- return self._ain_inference_extra_memory_ratio
-
- def set_ain_inference_extra_memory_ratio(
- self, ain_inference_extra_memory_ratio: float
- ) -> None:
- self._ain_inference_extra_memory_ratio = ain_inference_extra_memory_ratio
-
- def get_ain_logs_dir(self) -> str:
- return self._ain_logs_dir
-
- def set_ain_logs_dir(self, ain_logs_dir: str) -> None:
- self._ain_logs_dir = ain_logs_dir
-
- def get_ain_models_dir(self) -> str:
- return self._ain_models_dir
-
- def set_ain_models_dir(self, ain_models_dir: str) -> None:
- self._ain_models_dir = ain_models_dir
-
- def get_ain_builtin_models_dir(self) -> str:
- return self._ain_builtin_models_dir
-
- def set_ain_builtin_models_dir(self, ain_builtin_models_dir: str) -> None:
- self._ain_builtin_models_dir = ain_builtin_models_dir
-
- def get_ain_system_dir(self) -> str:
- return self._ain_system_dir
-
- def set_ain_system_dir(self, ain_system_dir: str) -> None:
- self._ain_system_dir = ain_system_dir
-
- def get_ain_thrift_compression_enabled(self) -> bool:
- return self._ain_thrift_compression_enabled
-
- def set_ain_thrift_compression_enabled(
- self, ain_thrift_compression_enabled: int
- ) -> 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
-
- def get_ain_target_config_node_list(self) -> TEndPoint:
- return self._ain_target_config_node_list
-
- def set_ain_target_config_node_list(self, ain_target_config_node_list: str) -> None:
- self._ain_target_config_node_list = parse_endpoint_url(
- ain_target_config_node_list
- )
-
- def get_ain_cluster_ingress_address(self) -> str:
- return self._ain_cluster_ingress_address
-
- def set_ain_cluster_ingress_address(self, ain_cluster_ingress_address: str) -> None:
- self._ain_cluster_ingress_address = ain_cluster_ingress_address
-
- def get_ain_cluster_ingress_port(self) -> int:
- return self._ain_cluster_ingress_port
-
- def set_ain_cluster_ingress_port(self, ain_cluster_ingress_port: int) -> None:
- self._ain_cluster_ingress_port = ain_cluster_ingress_port
-
- def get_ain_cluster_ingress_username(self) -> str:
- return self._ain_cluster_ingress_username
-
- def set_ain_cluster_ingress_username(
- self, ain_cluster_ingress_username: str
- ) -> None:
- self._ain_cluster_ingress_username = ain_cluster_ingress_username
-
- def get_ain_cluster_ingress_password(self) -> str:
- return self._ain_cluster_ingress_password
-
- def set_ain_cluster_ingress_password(
- self, ain_cluster_ingress_password: str
- ) -> None:
- self._ain_cluster_ingress_password = ain_cluster_ingress_password
-
- def get_ain_cluster_ingress_time_zone(self) -> str:
- return self._ain_cluster_ingress_time_zone
-
- def set_ain_cluster_ingress_time_zone(
- self, ain_cluster_ingress_time_zone: str
- ) -> None:
- self._ain_cluster_ingress_time_zone = ain_cluster_ingress_time_zone
-
-
-@singleton
-class AINodeDescriptor(object):
- def __init__(self):
- self._config = AINodeConfig()
- self._load_config_from_file()
- logger.info("AINodeDescriptor is init successfully.")
-
- def _load_config_from_file(self) -> None:
- system_properties_file = os.path.join(
- self._config.get_ain_system_dir(), AINODE_SYSTEM_FILE_NAME
- )
- if os.path.exists(system_properties_file):
- system_configs = load_properties(system_properties_file)
- if "ainode_id" in system_configs:
- self._config.set_ainode_id(int(system_configs["ainode_id"]))
-
- git_file = os.path.join(
- AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_GIT_FILE_NAME
- )
- if os.path.exists(git_file):
- git_configs = load_properties(git_file)
- if "git.commit.id.abbrev" in git_configs:
- build_info = git_configs["git.commit.id.abbrev"]
- if "git.dirty" in git_configs:
- if git_configs["git.dirty"] == "true":
- build_info += "-dev"
- self._config.set_build_info(build_info)
-
- pom_file = os.path.join(
- AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_POM_FILE_NAME
- )
- if os.path.exists(pom_file):
- pom_configs = load_properties(pom_file)
- if "version" in pom_configs:
- self._config.set_version_info(pom_configs["version"])
-
- conf_file = os.path.join(AINODE_CONF_DIRECTORY_NAME, AINODE_CONF_FILE_NAME)
- if not os.path.exists(conf_file):
- logger.info(
- "Cannot find AINode config file '{}', use default configuration.".format(
- conf_file
- )
- )
- return
-
- # noinspection PyBroadException
- try:
- file_configs = load_properties(conf_file)
-
- config_keys = file_configs.keys()
-
- if "ain_rpc_address" in config_keys:
- self._config.set_ain_rpc_address(file_configs["ain_rpc_address"])
-
- if "ain_rpc_port" in config_keys:
- self._config.set_ain_rpc_port(int(file_configs["ain_rpc_port"]))
-
- if "ain_inference_batch_interval_in_ms" in config_keys:
- self._config.set_ain_inference_batch_interval_in_ms(
- int(file_configs["ain_inference_batch_interval_in_ms"])
- )
-
- if "ain_inference_model_mem_usage_map" in config_keys:
- self._config.set_ain_inference_model_mem_usage_map(
- eval(file_configs["ain_inference_model_mem_usage_map"])
- )
-
- if "ain_inference_memory_usage_ratio" in config_keys:
- self._config.set_ain_inference_memory_usage_ratio(
- float(file_configs["ain_inference_memory_usage_ratio"])
- )
-
- if "ain_inference_extra_memory_ratio" in config_keys:
- self._config.set_ain_inference_extra_memory_ratio(
- float(file_configs["ain_inference_extra_memory_ratio"])
- )
-
- if "ain_models_dir" in config_keys:
- self._config.set_ain_models_dir(file_configs["ain_models_dir"])
-
- if "ain_system_dir" in config_keys:
- self._config.set_ain_system_dir(file_configs["ain_system_dir"])
-
- if "ain_seed_config_node" in config_keys:
- self._config.set_ain_target_config_node_list(
- file_configs["ain_seed_config_node"]
- )
-
- if "cluster_name" in config_keys:
- self._config.set_cluster_name(file_configs["cluster_name"])
-
- if "ain_thrift_compression_enabled" in config_keys:
- self._config.set_ain_thrift_compression_enabled(
- 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_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)
-
- if "ain_cluster_ingress_address" in config_keys:
- self._config.set_ain_cluster_ingress_address(
- file_configs["ain_cluster_ingress_address"]
- )
-
- if "ain_cluster_ingress_port" in config_keys:
- self._config.set_ain_cluster_ingress_port(
- int(file_configs["ain_cluster_ingress_port"])
- )
-
- if "ain_cluster_ingress_username" in config_keys:
- self._config.set_ain_cluster_ingress_username(
- file_configs["ain_cluster_ingress_username"]
- )
-
- if "ain_cluster_ingress_password" in config_keys:
- self._config.set_ain_cluster_ingress_password(
- file_configs["ain_cluster_ingress_password"]
- )
-
- if "ain_cluster_ingress_time_zone" in config_keys:
- self._config.set_ain_cluster_ingress_time_zone(
- file_configs["ain_cluster_ingress_time_zone"]
- )
-
- except BadNodeUrlError:
- logger.warning("Cannot load AINode conf file, use default configuration.")
-
- except Exception as e:
- logger.warning(
- "Cannot load AINode conf file caused by: {}, use default configuration. ".format(
- e
- )
- )
-
- def get_config(self) -> AINodeConfig:
- return self._config
-
-
-def load_properties(filepath, sep="=", comment_char="#"):
- """
- Read the file passed as parameter as a properties file.
- """
- props = {}
- with open(filepath, "rt") as f:
- for line in f:
- l = line.strip()
- if l and not l.startswith(comment_char):
- key_value = l.split(sep)
- key = key_value[0].strip()
- value = sep.join(key_value[1:]).strip().strip('"')
- props[key] = value
- return props
-
-
-def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
- """Parse TEndPoint from a given endpoint url.
- Args:
- endpoint_url: an endpoint url, format: ip:port
- Returns:
- TEndPoint
- Raises:
- BadNodeUrlError
- """
- split = endpoint_url.split(":")
- if len(split) != 2:
- raise BadNodeUrlError(endpoint_url)
-
- ip = split[0]
- try:
- port = int(split[1])
- result = TEndPoint(ip, port)
- return result
- except ValueError:
- raise BadNodeUrlError(endpoint_url)
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import os
+1
+1from iotdb.ainode.core.constant import (
+1 AINODE_BUILD_INFO,
+1 AINODE_BUILTIN_MODELS_DIR,
+1 AINODE_CLUSTER_INGRESS_ADDRESS,
+1 AINODE_CLUSTER_INGRESS_PASSWORD,
+1 AINODE_CLUSTER_INGRESS_PORT,
+1 AINODE_CLUSTER_INGRESS_TIME_ZONE,
+1 AINODE_CLUSTER_INGRESS_USERNAME,
+1 AINODE_CLUSTER_NAME,
+1 AINODE_CONF_DIRECTORY_NAME,
+1 AINODE_CONF_FILE_NAME,
+1 AINODE_CONF_GIT_FILE_NAME,
+1 AINODE_CONF_POM_FILE_NAME,
+1 AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
+1 AINODE_INFERENCE_EXTRA_MEMORY_RATIO,
+1 AINODE_INFERENCE_MAX_PREDICT_LENGTH,
+1 AINODE_INFERENCE_MEMORY_USAGE_RATIO,
+1 AINODE_INFERENCE_MODEL_MEM_USAGE_MAP,
+1 AINODE_LOG_DIR,
+1 AINODE_MODELS_DIR,
+1 AINODE_ROOT_CONF_DIRECTORY_NAME,
+1 AINODE_ROOT_DIR,
+1 AINODE_RPC_ADDRESS,
+1 AINODE_RPC_PORT,
+1 AINODE_SYSTEM_DIR,
+1 AINODE_SYSTEM_FILE_NAME,
+1 AINODE_TARGET_CONFIG_NODE_LIST,
+1 AINODE_THRIFT_COMPRESSION_ENABLED,
+1 AINODE_VERSION_INFO,
+1)
+1from iotdb.ainode.core.exception import BadNodeUrlError
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.util.decorator import singleton
+1from iotdb.thrift.common.ttypes import TEndPoint
+1
+1logger = Logger()
+1
+1
+1class AINodeConfig(object):
+1 def __init__(self):
+1 self._version_info = AINODE_VERSION_INFO
+1 self._build_info = AINODE_BUILD_INFO
+1
+1 # Cluster configuration
+1 self._ainode_id = 0
+1 self._cluster_name = AINODE_CLUSTER_NAME
+1 self._ain_target_config_node_list: TEndPoint = AINODE_TARGET_CONFIG_NODE_LIST
+1 self._ain_rpc_address: str = AINODE_RPC_ADDRESS
+1 self._ain_rpc_port: int = AINODE_RPC_PORT
+1 self._ain_cluster_ingress_address = AINODE_CLUSTER_INGRESS_ADDRESS
+1 self._ain_cluster_ingress_port = AINODE_CLUSTER_INGRESS_PORT
+1 self._ain_cluster_ingress_username = AINODE_CLUSTER_INGRESS_USERNAME
+1 self._ain_cluster_ingress_password = AINODE_CLUSTER_INGRESS_PASSWORD
+1 self._ain_cluster_ingress_time_zone = AINODE_CLUSTER_INGRESS_TIME_ZONE
+1
+1 # Inference configuration
+1 self._ain_inference_batch_interval_in_ms: int = (
+1 AINODE_INFERENCE_BATCH_INTERVAL_IN_MS
+1 )
+1 self._ain_inference_max_predict_length: int = (
+1 AINODE_INFERENCE_MAX_PREDICT_LENGTH
+1 )
+1 self._ain_inference_model_mem_usage_map: dict[str, int] = (
+1 AINODE_INFERENCE_MODEL_MEM_USAGE_MAP
+1 )
+1 self._ain_inference_memory_usage_ratio: float = (
+1 AINODE_INFERENCE_MEMORY_USAGE_RATIO
+1 )
+1 self._ain_inference_extra_memory_ratio: float = (
+1 AINODE_INFERENCE_EXTRA_MEMORY_RATIO
+1 )
+1 # log directory
+1 self._ain_logs_dir: str = AINODE_LOG_DIR
+1
+1 # cache size for ingress dataloader (MB)
+1 self._ain_data_cache_size = 50
+1
+1 # Directory to save models
+1 self._ain_models_dir = AINODE_MODELS_DIR
+1 self._ain_builtin_models_dir = AINODE_BUILTIN_MODELS_DIR
+1 self._ain_system_dir = AINODE_SYSTEM_DIR
+1
+1 # Whether to enable compression for thrift
+1 self._ain_thrift_compression_enabled = AINODE_THRIFT_COMPRESSION_ENABLED
+1
+1 # use for ssl
+1 self._ain_cluster_ingress_ssl_enabled = False
+1 self._ain_internal_ssl_enabled = False
+1 self._ain_thrift_ssl_cert_file = None
+1 self._ain_thrift_ssl_key_file = None
+1
+1 # Cache number of model storage to avoid repeated loading
+1 self._ain_model_storage_cache_size = 30
+1
+1 def get_cluster_name(self) -> str:
+1 return self._cluster_name
+1
+1 def set_cluster_name(self, cluster_name: str) -> None:
+1 self._cluster_name = cluster_name
+1
+1 def get_version_info(self) -> str:
+1 return self._version_info
+1
+1 def get_ainode_id(self) -> int:
+1 return self._ainode_id
+1
+1 def set_ainode_id(self, ainode_id: int) -> None:
+1 self._ainode_id = ainode_id
+1
+1 def get_build_info(self) -> str:
+1 return self._build_info
+1
+1 def set_build_info(self, build_info: str) -> None:
+1 self._build_info = build_info
+1
+1 def get_ain_data_storage_cache_size(self) -> int:
+1 return self._ain_data_cache_size
+1
+1 def set_ain_data_cache_size(self, ain_data_cache_size: int) -> None:
+1 self._ain_data_cache_size = ain_data_cache_size
+1
+1 def set_version_info(self, version_info: str) -> None:
+1 self._version_info = version_info
+1
+1 def get_ain_rpc_address(self) -> str:
+1 return self._ain_rpc_address
+1
+1 def set_ain_rpc_address(self, ain_rpc_address: str) -> None:
+1 self._ain_rpc_address = ain_rpc_address
+1
+1 def get_ain_rpc_port(self) -> int:
+1 return self._ain_rpc_port
+1
+1 def set_ain_rpc_port(self, ain_rpc_port: int) -> None:
+1 self._ain_rpc_port = ain_rpc_port
+1
+1 def get_ain_inference_batch_interval_in_ms(self) -> int:
+1 return self._ain_inference_batch_interval_in_ms
+1
+1 def set_ain_inference_batch_interval_in_ms(
+1 self, ain_inference_batch_interval_in_ms: int
+1 ) -> None:
+1 self._ain_inference_batch_interval_in_ms = ain_inference_batch_interval_in_ms
+1
+1 def get_ain_inference_max_predict_length(self) -> int:
+1 return self._ain_inference_max_predict_length
+1
+1 def set_ain_inference_max_predict_length(
+1 self, ain_inference_max_predict_length: int
+1 ) -> None:
+1 self._ain_inference_max_predict_length = ain_inference_max_predict_length
+1
+1 def get_ain_inference_model_mem_usage_map(self) -> dict[str, int]:
+1 return self._ain_inference_model_mem_usage_map
+1
+1 def set_ain_inference_model_mem_usage_map(
+1 self, ain_inference_model_mem_usage_map: dict[str, int]
+1 ) -> None:
+1 self._ain_inference_model_mem_usage_map = ain_inference_model_mem_usage_map
+1
+1 def get_ain_inference_memory_usage_ratio(self) -> float:
+1 return self._ain_inference_memory_usage_ratio
+1
+1 def set_ain_inference_memory_usage_ratio(
+1 self, ain_inference_memory_usage_ratio: float
+1 ) -> None:
+1 self._ain_inference_memory_usage_ratio = ain_inference_memory_usage_ratio
+1
+1 def get_ain_inference_extra_memory_ratio(self) -> float:
+1 return self._ain_inference_extra_memory_ratio
+1
+1 def set_ain_inference_extra_memory_ratio(
+1 self, ain_inference_extra_memory_ratio: float
+1 ) -> None:
+1 self._ain_inference_extra_memory_ratio = ain_inference_extra_memory_ratio
+1
+1 def get_ain_logs_dir(self) -> str:
+1 return self._ain_logs_dir
+1
+1 def set_ain_logs_dir(self, ain_logs_dir: str) -> None:
+1 self._ain_logs_dir = ain_logs_dir
+1
+1 def get_ain_models_dir(self) -> str:
+1 return self._ain_models_dir
+1
+1 def set_ain_models_dir(self, ain_models_dir: str) -> None:
+1 self._ain_models_dir = ain_models_dir
+1
+1 def get_ain_builtin_models_dir(self) -> str:
+1 return self._ain_builtin_models_dir
+1
+1 def set_ain_builtin_models_dir(self, ain_builtin_models_dir: str) -> None:
+1 self._ain_builtin_models_dir = ain_builtin_models_dir
+1
+1 def get_ain_system_dir(self) -> str:
+1 return self._ain_system_dir
+1
+1 def set_ain_system_dir(self, ain_system_dir: str) -> None:
+1 self._ain_system_dir = ain_system_dir
+1
+1 def get_ain_thrift_compression_enabled(self) -> bool:
+1 return self._ain_thrift_compression_enabled
+1
+1 def set_ain_thrift_compression_enabled(
+1 self, ain_thrift_compression_enabled: int
+1 ) -> None:
+1 self._ain_thrift_compression_enabled = ain_thrift_compression_enabled
+1
+1 def get_ain_cluster_ingress_ssl_enabled(self) -> bool:
+1 return self._ain_cluster_ingress_ssl_enabled
+1
+1 def set_ain_cluster_ingress_ssl_enabled(
+1 self, ain_cluster_ingress_ssl_enabled: int
+1 ) -> None:
+1 self._ain_cluster_ingress_ssl_enabled = ain_cluster_ingress_ssl_enabled
+1
+1 def get_ain_internal_ssl_enabled(self) -> bool:
+1 return self._ain_internal_ssl_enabled
+1
+1 def set_ain_internal_ssl_enabled(self, ain_internal_ssl_enabled: int) -> None:
+1 self._ain_internal_ssl_enabled = ain_internal_ssl_enabled
+1
+1 def get_ain_thrift_ssl_cert_file(self) -> str:
+1 return self._ain_thrift_ssl_cert_file
+1
+1 def set_ain_thrift_ssl_cert_file(self, ain_thrift_ssl_cert_file: str) -> None:
+1 self._ain_thrift_ssl_cert_file = ain_thrift_ssl_cert_file
+1
+1 def get_ain_thrift_ssl_key_file(self) -> str:
+1 return self._ain_thrift_ssl_key_file
+1
+1 def set_ain_thrift_ssl_key_file(self, ain_thrift_ssl_key_file: str) -> None:
+1 self._ain_thrift_ssl_key_file = ain_thrift_ssl_key_file
+1
+1 def get_ain_model_storage_cache_size(self) -> int:
+1 return self._ain_model_storage_cache_size
+1
+1 def get_ain_target_config_node_list(self) -> TEndPoint:
+1 return self._ain_target_config_node_list
+1
+1 def set_ain_target_config_node_list(self, ain_target_config_node_list: str) -> None:
+1 self._ain_target_config_node_list = parse_endpoint_url(
+1 ain_target_config_node_list
+1 )
+1
+1 def get_ain_cluster_ingress_address(self) -> str:
+1 return self._ain_cluster_ingress_address
+1
+1 def set_ain_cluster_ingress_address(self, ain_cluster_ingress_address: str) -> None:
+1 self._ain_cluster_ingress_address = ain_cluster_ingress_address
+1
+1 def get_ain_cluster_ingress_port(self) -> int:
+1 return self._ain_cluster_ingress_port
+1
+1 def set_ain_cluster_ingress_port(self, ain_cluster_ingress_port: int) -> None:
+1 self._ain_cluster_ingress_port = ain_cluster_ingress_port
+1
+1 def get_ain_cluster_ingress_username(self) -> str:
+1 return self._ain_cluster_ingress_username
+1
+1 def set_ain_cluster_ingress_username(
+1 self, ain_cluster_ingress_username: str
+1 ) -> None:
+1 self._ain_cluster_ingress_username = ain_cluster_ingress_username
+1
+1 def get_ain_cluster_ingress_password(self) -> str:
+1 return self._ain_cluster_ingress_password
+1
+1 def set_ain_cluster_ingress_password(
+1 self, ain_cluster_ingress_password: str
+1 ) -> None:
+1 self._ain_cluster_ingress_password = ain_cluster_ingress_password
+1
+1 def get_ain_cluster_ingress_time_zone(self) -> str:
+1 return self._ain_cluster_ingress_time_zone
+1
+1 def set_ain_cluster_ingress_time_zone(
+1 self, ain_cluster_ingress_time_zone: str
+1 ) -> None:
+1 self._ain_cluster_ingress_time_zone = ain_cluster_ingress_time_zone
+1
+1
+1@singleton
+1class AINodeDescriptor(object):
+1 def __init__(self):
+1 self._config = AINodeConfig()
+1 self._load_config_from_file()
+1 logger.info("AINodeDescriptor is init successfully.")
+1
+1 def _load_config_from_file(self) -> None:
+1 system_properties_file = os.path.join(
+1 self._config.get_ain_system_dir(), AINODE_SYSTEM_FILE_NAME
+1 )
+1 if os.path.exists(system_properties_file):
+1 system_configs = load_properties(system_properties_file)
+1 if "ainode_id" in system_configs:
+1 self._config.set_ainode_id(int(system_configs["ainode_id"]))
+1
+1 git_file = os.path.join(
+1 AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_GIT_FILE_NAME
+1 )
+1 if os.path.exists(git_file):
+1 git_configs = load_properties(git_file)
+1 if "git.commit.id.abbrev" in git_configs:
+1 build_info = git_configs["git.commit.id.abbrev"]
+1 if "git.dirty" in git_configs:
+1 if git_configs["git.dirty"] == "true":
+1 build_info += "-dev"
+1 self._config.set_build_info(build_info)
+1
+1 pom_file = os.path.join(
+1 AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_POM_FILE_NAME
+1 )
+1 if os.path.exists(pom_file):
+1 pom_configs = load_properties(pom_file)
+1 if "version" in pom_configs:
+1 self._config.set_version_info(pom_configs["version"])
+1
+1 conf_file = os.path.join(AINODE_CONF_DIRECTORY_NAME, AINODE_CONF_FILE_NAME)
+1 if not os.path.exists(conf_file):
+1 logger.info(
+1 "Cannot find AINode config file '{}', use default configuration.".format(
+1 conf_file
+1 )
+1 )
+1 return
+1
+1 # noinspection PyBroadException
+1 try:
+1 file_configs = load_properties(conf_file)
+1
+1 config_keys = file_configs.keys()
+1
+1 if "ain_rpc_address" in config_keys:
+1 self._config.set_ain_rpc_address(file_configs["ain_rpc_address"])
+1
+1 if "ain_rpc_port" in config_keys:
+1 self._config.set_ain_rpc_port(int(file_configs["ain_rpc_port"]))
+1
+1 if "ain_inference_batch_interval_in_ms" in config_keys:
+1 self._config.set_ain_inference_batch_interval_in_ms(
+1 int(file_configs["ain_inference_batch_interval_in_ms"])
+1 )
+1
+1 if "ain_inference_model_mem_usage_map" in config_keys:
+1 self._config.set_ain_inference_model_mem_usage_map(
+1 eval(file_configs["ain_inference_model_mem_usage_map"])
+1 )
+1
+1 if "ain_inference_memory_usage_ratio" in config_keys:
+1 self._config.set_ain_inference_memory_usage_ratio(
+1 float(file_configs["ain_inference_memory_usage_ratio"])
+1 )
+1
+1 if "ain_inference_extra_memory_ratio" in config_keys:
+1 self._config.set_ain_inference_extra_memory_ratio(
+1 float(file_configs["ain_inference_extra_memory_ratio"])
+1 )
+1
+1 if "ain_models_dir" in config_keys:
+1 self._config.set_ain_models_dir(file_configs["ain_models_dir"])
+1
+1 if "ain_system_dir" in config_keys:
+1 self._config.set_ain_system_dir(file_configs["ain_system_dir"])
+1
+1 if "ain_seed_config_node" in config_keys:
+1 self._config.set_ain_target_config_node_list(
+1 file_configs["ain_seed_config_node"]
+1 )
+1
+1 if "cluster_name" in config_keys:
+1 self._config.set_cluster_name(file_configs["cluster_name"])
+1
+1 if "ain_thrift_compression_enabled" in config_keys:
+1 self._config.set_ain_thrift_compression_enabled(
+1 int(file_configs["ain_thrift_compression_enabled"])
+1 )
+1
+1 if "ain_cluster_ingress_ssl_enabled" in config_keys:
+1 self._config.set_ain_cluster_ingress_ssl_enabled(
+1 int(file_configs["ain_cluster_ingress_ssl_enabled"])
+1 )
+1
+1 if "ain_thrift_ssl_cert_file" in config_keys:
+1 self._config.set_ain_thrift_ssl_cert_file(
+1 file_configs["ain_thrift_ssl_cert_file"]
+1 )
+1
+1 if "ain_thrift_ssl_key_file" in config_keys:
+1 self._config.set_ain_thrift_ssl_key_file(
+1 file_configs["ain_thrift_ssl_key_file"]
+1 )
+1
+1 if "ain_logs_dir" in config_keys:
+1 log_dir = file_configs["ain_logs_dir"]
+1 self._config.set_ain_logs_dir(log_dir)
+1
+1 if "ain_cluster_ingress_address" in config_keys:
+1 self._config.set_ain_cluster_ingress_address(
+1 file_configs["ain_cluster_ingress_address"]
+1 )
+1
+1 if "ain_cluster_ingress_port" in config_keys:
+1 self._config.set_ain_cluster_ingress_port(
+1 int(file_configs["ain_cluster_ingress_port"])
+1 )
+1
+1 if "ain_cluster_ingress_username" in config_keys:
+1 self._config.set_ain_cluster_ingress_username(
+1 file_configs["ain_cluster_ingress_username"]
+1 )
+1
+1 if "ain_cluster_ingress_password" in config_keys:
+1 self._config.set_ain_cluster_ingress_password(
+1 file_configs["ain_cluster_ingress_password"]
+1 )
+1
+1 if "ain_cluster_ingress_time_zone" in config_keys:
+1 self._config.set_ain_cluster_ingress_time_zone(
+1 file_configs["ain_cluster_ingress_time_zone"]
+1 )
+1
+1 except BadNodeUrlError:
+1 logger.warning("Cannot load AINode conf file, use default configuration.")
+1
+1 except Exception as e:
+1 logger.warning(
+1 "Cannot load AINode conf file caused by: {}, use default configuration. ".format(
+1 e
+1 )
+1 )
+1
+1 def get_config(self) -> AINodeConfig:
+1 return self._config
+1
+1
+1def load_properties(filepath, sep="=", comment_char="#"):
+1 """
+1 Read the file passed as parameter as a properties file.
+1 """
+1 props = {}
+1 with open(filepath, "rt") as f:
+1 for line in f:
+1 l = line.strip()
+1 if l and not l.startswith(comment_char):
+1 key_value = l.split(sep)
+1 key = key_value[0].strip()
+1 value = sep.join(key_value[1:]).strip().strip('"')
+1 props[key] = value
+1 return props
+1
+1
+1def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
+1 """Parse TEndPoint from a given endpoint url.
+1 Args:
+1 endpoint_url: an endpoint url, format: ip:port
+1 Returns:
+1 TEndPoint
+1 Raises:
+1 BadNodeUrlError
+1 """
+1 split = endpoint_url.split(":")
+1 if len(split) != 2:
+1 raise BadNodeUrlError(endpoint_url)
+1
+1 ip = split[0]
+1 try:
+1 port = int(split[1])
+1 result = TEndPoint(ip, port)
+1 return result
+1 except ValueError:
+1 raise BadNodeUrlError(endpoint_url)
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/constant.py b/iotdb-core/ainode/iotdb/ainode/core/constant.py
index f4547e99803..87735a14b25 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/constant.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/constant.py
@@ -1,274 +1,275 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import inspect
-import logging
-import os
-from enum import Enum
-from typing import List
-
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
-from iotdb.thrift.common.ttypes import TEndPoint
-
-AINODE_VERSION_INFO = "UNKNOWN"
-AINODE_BUILD_INFO = "UNKNOWN"
-AINODE_CONF_DIRECTORY_NAME = "conf"
-AINODE_ROOT_CONF_DIRECTORY_NAME = "conf"
-AINODE_CONF_FILE_NAME = "iotdb-ainode.properties"
-AINODE_CONF_GIT_FILE_NAME = "git.properties"
-AINODE_CONF_POM_FILE_NAME = "pom.properties"
-AINODE_SYSTEM_FILE_NAME = "system.properties"
-
-# AINode cluster configuration
-AINODE_CLUSTER_NAME = "defaultCluster"
-AINODE_TARGET_CONFIG_NODE_LIST = TEndPoint("127.0.0.1", 10710)
-AINODE_RPC_ADDRESS = "127.0.0.1"
-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_TIME_ZONE = "UTC+8"
-
-# RPC config
-AINODE_THRIFT_COMPRESSION_ENABLED = False
-DEFAULT_RECONNECT_TIMEOUT = 20
-DEFAULT_RECONNECT_TIMES = 3
-
-# AINode inference configuration
-AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
-AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
-AINODE_INFERENCE_MODEL_MEM_USAGE_MAP = {
- BuiltInModelType.SUNDIAL.value: 1036 * 1024**2, # 1036 MiB
- BuiltInModelType.TIMER_XL.value: 856 * 1024**2, # 856 MiB
-} # the memory usage of each model in bytes
-AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.4 # the device space allocated for inference
-AINODE_INFERENCE_EXTRA_MEMORY_RATIO = (
- 1.2 # the overhead ratio for inference, used to estimate the pool size
-)
-
-# AINode folder structure
-AINODE_ROOT_DIR = os.path.dirname(
- os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
-)
-AINODE_MODELS_DIR = "data/ainode/models"
-AINODE_BUILTIN_MODELS_DIR = "data/ainode/models/weights" # For built-in models, we only need to store their weights and config.
-AINODE_SYSTEM_DIR = "data/ainode/system"
-AINODE_LOG_DIR = "logs"
-
-# AINode log
-LOG_FILE_TYPE = ["all", "info", "warn", "error"]
-AINODE_LOG_FILE_NAME_PREFIX = "log_ainode_"
-AINODE_LOG_FILE_LEVELS = [logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR]
-DEFAULT_LOG_LEVEL = logging.INFO
-INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE = (
- "log_inference_rank_{}_" # example: log_inference_rank_0_all.log
-)
-
-# AINode model management
-MODEL_WEIGHTS_FILE_IN_SAFETENSORS = "model.safetensors"
-MODEL_CONFIG_FILE_IN_JSON = "config.json"
-MODEL_WEIGHTS_FILE_IN_PT = "model.pt"
-MODEL_CONFIG_FILE_IN_YAML = "config.yaml"
-DEFAULT_CHUNK_SIZE = 8192
-
-
-class TSStatusCode(Enum):
- SUCCESS_STATUS = 200
- REDIRECTION_RECOMMEND = 400
- MODEL_EXIST_ERROR = 1502
- MODEL_NOT_FOUND_ERROR = 1505
- UNAVAILABLE_AI_DEVICE_ERROR = 1507
- AINODE_INTERNAL_ERROR = 1510
- INVALID_URI_ERROR = 1511
- INVALID_INFERENCE_CONFIG = 1512
- INFERENCE_INTERNAL_ERROR = 1520
-
- def get_status_code(self) -> int:
- return self.value
-
-
-class TaskType(Enum):
- FORECAST = "forecast"
-
-
-class OptionsKey(Enum):
- # common
- TASK_TYPE = "task_type"
- MODEL_TYPE = "model_type"
- AUTO_TUNING = "auto_tuning"
- INPUT_VARS = "input_vars"
-
- # forecast
- INPUT_LENGTH = "input_length"
- PREDICT_LENGTH = "predict_length"
- PREDICT_INDEX_LIST = "predict_index_list"
- INPUT_TYPE_LIST = "input_type_list"
-
- def name(self) -> str:
- return self.value
-
-
-class HyperparameterName(Enum):
- # Training hyperparameter
- LEARNING_RATE = "learning_rate"
- EPOCHS = "epochs"
- BATCH_SIZE = "batch_size"
- USE_GPU = "use_gpu"
- NUM_WORKERS = "num_workers"
-
- # Structure hyperparameter
- KERNEL_SIZE = "kernel_size"
- INPUT_VARS = "input_vars"
- BLOCK_TYPE = "block_type"
- D_MODEL = "d_model"
- INNER_LAYERS = "inner_layer"
- OUTER_LAYERS = "outer_layer"
-
- def name(self):
- return self.value
-
-
-class ForecastModelType(Enum):
- DLINEAR = "dlinear"
- DLINEAR_INDIVIDUAL = "dlinear_individual"
- NBEATS = "nbeats"
-
- @classmethod
- def values(cls) -> List[str]:
- values = []
- for item in list(cls):
- values.append(item.value)
- return values
-
-
-class ModelInputName(Enum):
- DATA_X = "data_x"
- TIME_STAMP_X = "time_stamp_x"
- TIME_STAMP_Y = "time_stamp_y"
- DEC_INP = "dec_inp"
-
-
-class AttributeName(Enum):
- # forecast Attribute
- PREDICT_LENGTH = "predict_length"
-
- # NaiveForecaster
- STRATEGY = "strategy"
- SP = "sp"
-
- # STLForecaster
- # SP = 'sp'
- SEASONAL = "seasonal"
- SEASONAL_DEG = "seasonal_deg"
- TREND_DEG = "trend_deg"
- LOW_PASS_DEG = "low_pass_deg"
- SEASONAL_JUMP = "seasonal_jump"
- TREND_JUMP = "trend_jump"
- LOSS_PASS_JUMP = "low_pass_jump"
-
- # ExponentialSmoothing
- DAMPED_TREND = "damped_trend"
- INITIALIZATION_METHOD = "initialization_method"
- OPTIMIZED = "optimized"
- REMOVE_BIAS = "remove_bias"
- USE_BRUTE = "use_brute"
-
- # Arima
- ORDER = "order"
- SEASONAL_ORDER = "seasonal_order"
- METHOD = "method"
- MAXITER = "maxiter"
- SUPPRESS_WARNINGS = "suppress_warnings"
- OUT_OF_SAMPLE_SIZE = "out_of_sample_size"
- SCORING = "scoring"
- WITH_INTERCEPT = "with_intercept"
- TIME_VARYING_REGRESSION = "time_varying_regression"
- ENFORCE_STATIONARITY = "enforce_stationarity"
- ENFORCE_INVERTIBILITY = "enforce_invertibility"
- SIMPLE_DIFFERENCING = "simple_differencing"
- MEASUREMENT_ERROR = "measurement_error"
- MLE_REGRESSION = "mle_regression"
- HAMILTON_REPRESENTATION = "hamilton_representation"
- CONCENTRATE_SCALE = "concentrate_scale"
-
- # GAUSSIAN_HMM
- N_COMPONENTS = "n_components"
- COVARIANCE_TYPE = "covariance_type"
- MIN_COVAR = "min_covar"
- STARTPROB_PRIOR = "startprob_prior"
- TRANSMAT_PRIOR = "transmat_prior"
- MEANS_PRIOR = "means_prior"
- MEANS_WEIGHT = "means_weight"
- COVARS_PRIOR = "covars_prior"
- COVARS_WEIGHT = "covars_weight"
- ALGORITHM = "algorithm"
- N_ITER = "n_iter"
- TOL = "tol"
- PARAMS = "params"
- INIT_PARAMS = "init_params"
- IMPLEMENTATION = "implementation"
-
- # GMMHMM
- # N_COMPONENTS = "n_components"
- N_MIX = "n_mix"
- # MIN_COVAR = "min_covar"
- # STARTPROB_PRIOR = "startprob_prior"
- # TRANSMAT_PRIOR = "transmat_prior"
- WEIGHTS_PRIOR = "weights_prior"
-
- # MEANS_PRIOR = "means_prior"
- # MEANS_WEIGHT = "means_weight"
- # ALGORITHM = "algorithm"
- # COVARIANCE_TYPE = "covariance_type"
- # N_ITER = "n_iter"
- # TOL = "tol"
- # INIT_PARAMS = "init_params"
- # PARAMS = "params"
- # IMPLEMENTATION = "implementation"
-
- # STRAY
- ALPHA = "alpha"
- K = "k"
- KNN_ALGORITHM = "knn_algorithm"
- P = "p"
- SIZE_THRESHOLD = "size_threshold"
- OUTLIER_TAIL = "outlier_tail"
-
- # timerxl
- INPUT_TOKEN_LEN = "input_token_len"
- HIDDEN_SIZE = "hidden_size"
- INTERMEDIATE_SIZE = "intermediate_size"
- OUTPUT_TOKEN_LENS = "output_token_lens"
- NUM_HIDDEN_LAYERS = "num_hidden_layers"
- NUM_ATTENTION_HEADS = "num_attention_heads"
- HIDDEN_ACT = "hidden_act"
- USE_CACHE = "use_cache"
- ROPE_THETA = "rope_theta"
- ATTENTION_DROPOUT = "attention_dropout"
- INITIALIZER_RANGE = "initializer_range"
- MAX_POSITION_EMBEDDINGS = "max_position_embeddings"
- CKPT_PATH = "ckpt_path"
-
- # sundial
- DROPOUT_RATE = "dropout_rate"
- FLOW_LOSS_DEPTH = "flow_loss_depth"
- NUM_SAMPLING_STEPS = "num_sampling_steps"
- DIFFUSION_BATCH_MUL = "diffusion_batch_mul"
-
- def name(self) -> str:
- return self.value
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import inspect
+1import logging
+1import os
+1from enum import Enum
+1from typing import List
+1
+1from iotdb.ainode.core.model.model_enums import BuiltInModelType
+1from iotdb.thrift.common.ttypes import TEndPoint
+1
+1AINODE_VERSION_INFO = "UNKNOWN"
+1AINODE_BUILD_INFO = "UNKNOWN"
+1AINODE_CONF_DIRECTORY_NAME = "conf"
+1AINODE_ROOT_CONF_DIRECTORY_NAME = "conf"
+1AINODE_CONF_FILE_NAME = "iotdb-ainode.properties"
+1AINODE_CONF_GIT_FILE_NAME = "git.properties"
+1AINODE_CONF_POM_FILE_NAME = "pom.properties"
+1AINODE_SYSTEM_FILE_NAME = "system.properties"
+1
+1# AINode cluster configuration
+1AINODE_CLUSTER_NAME = "defaultCluster"
+1AINODE_TARGET_CONFIG_NODE_LIST = TEndPoint("127.0.0.1", 10710)
+1AINODE_RPC_ADDRESS = "127.0.0.1"
+1AINODE_RPC_PORT = 10810
+1AINODE_CLUSTER_INGRESS_ADDRESS = "127.0.0.1"
+1AINODE_CLUSTER_INGRESS_PORT = 6667
+1AINODE_CLUSTER_INGRESS_USERNAME = "root"
+1AINODE_CLUSTER_INGRESS_PASSWORD = "root"
+1AINODE_CLUSTER_INGRESS_TIME_ZONE = "UTC+8"
+1
+1# RPC config
+1AINODE_THRIFT_COMPRESSION_ENABLED = False
+1DEFAULT_RECONNECT_TIMEOUT = 20
+1DEFAULT_RECONNECT_TIMES = 3
+1
+1# AINode inference configuration
+1AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
+1AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
+1AINODE_INFERENCE_MODEL_MEM_USAGE_MAP = {
+1 BuiltInModelType.SUNDIAL.value: 1036 * 1024**2, # 1036 MiB
+1 BuiltInModelType.TIMER_XL.value: 856 * 1024**2, # 856 MiB
+1} # the memory usage of each model in bytes
+1AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.4 # the device space allocated for inference
+1AINODE_INFERENCE_EXTRA_MEMORY_RATIO = (
+1 1.2 # the overhead ratio for inference, used to estimate the pool size
+1)
+1
+1# AINode folder structure
+1AINODE_ROOT_DIR = os.path.dirname(
+1 os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+1)
+1AINODE_MODELS_DIR = "data/ainode/models"
+1AINODE_BUILTIN_MODELS_DIR = "data/ainode/models/weights" # For built-in models, we only need to store their weights and config.
+1AINODE_SYSTEM_DIR = "data/ainode/system"
+1AINODE_LOG_DIR = "logs"
+1
+1# AINode log
+1LOG_FILE_TYPE = ["all", "info", "warn", "error"]
+1AINODE_LOG_FILE_NAME_PREFIX = "log_ainode_"
+1AINODE_LOG_FILE_LEVELS = [logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR]
+1DEFAULT_LOG_LEVEL = logging.INFO
+1INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE = (
+1 "log_inference_rank_{}_" # example: log_inference_rank_0_all.log
+1)
+1
+1# AINode model management
+1MODEL_WEIGHTS_FILE_IN_SAFETENSORS = "model.safetensors"
+1MODEL_CONFIG_FILE_IN_JSON = "config.json"
+1MODEL_WEIGHTS_FILE_IN_PT = "model.pt"
+1MODEL_CONFIG_FILE_IN_YAML = "config.yaml"
+1DEFAULT_CHUNK_SIZE = 8192
+1
+1
+1class TSStatusCode(Enum):
+1 SUCCESS_STATUS = 200
+1 REDIRECTION_RECOMMEND = 400
+1 MODEL_EXIST_ERROR = 1502
+1 MODEL_NOT_FOUND_ERROR = 1505
+1 UNAVAILABLE_AI_DEVICE_ERROR = 1507
+1 AINODE_INTERNAL_ERROR = 1510
+1 INVALID_URI_ERROR = 1511
+1 INVALID_INFERENCE_CONFIG = 1512
+1 INFERENCE_INTERNAL_ERROR = 1520
+1
+1 def get_status_code(self) -> int:
+1 return self.value
+1
+1
+1class TaskType(Enum):
+1 FORECAST = "forecast"
+1
+1
+1class OptionsKey(Enum):
+1 # common
+1 TASK_TYPE = "task_type"
+1 MODEL_TYPE = "model_type"
+1 AUTO_TUNING = "auto_tuning"
+1 INPUT_VARS = "input_vars"
+1
+1 # forecast
+1 INPUT_LENGTH = "input_length"
+1 PREDICT_LENGTH = "predict_length"
+1 PREDICT_INDEX_LIST = "predict_index_list"
+1 INPUT_TYPE_LIST = "input_type_list"
+1
+1 def name(self) -> str:
+1 return self.value
+1
+1
+1class HyperparameterName(Enum):
+1 # Training hyperparameter
+1 LEARNING_RATE = "learning_rate"
+1 EPOCHS = "epochs"
+1 BATCH_SIZE = "batch_size"
+1 USE_GPU = "use_gpu"
+1 NUM_WORKERS = "num_workers"
+1
+1 # Structure hyperparameter
+1 KERNEL_SIZE = "kernel_size"
+1 INPUT_VARS = "input_vars"
+1 BLOCK_TYPE = "block_type"
+1 D_MODEL = "d_model"
+1 INNER_LAYERS = "inner_layer"
+1 OUTER_LAYERS = "outer_layer"
+1
+1 def name(self):
+1 return self.value
+1
+1
+1class ForecastModelType(Enum):
+1 DLINEAR = "dlinear"
+1 DLINEAR_INDIVIDUAL = "dlinear_individual"
+1 NBEATS = "nbeats"
+1
+1 @classmethod
+1 def values(cls) -> List[str]:
+1 values = []
+1 for item in list(cls):
+1 values.append(item.value)
+1 return values
+1
+1
+1class ModelInputName(Enum):
+1 DATA_X = "data_x"
+1 TIME_STAMP_X = "time_stamp_x"
+1 TIME_STAMP_Y = "time_stamp_y"
+1 DEC_INP = "dec_inp"
+1
+1
+1class AttributeName(Enum):
+1 # forecast Attribute
+1 PREDICT_LENGTH = "predict_length"
+1
+1 # NaiveForecaster
+1 STRATEGY = "strategy"
+1 SP = "sp"
+1
+1 # STLForecaster
+1 # SP = 'sp'
+1 SEASONAL = "seasonal"
+1 SEASONAL_DEG = "seasonal_deg"
+1 TREND_DEG = "trend_deg"
+1 LOW_PASS_DEG = "low_pass_deg"
+1 SEASONAL_JUMP = "seasonal_jump"
+1 TREND_JUMP = "trend_jump"
+1 LOSS_PASS_JUMP = "low_pass_jump"
+1
+1 # ExponentialSmoothing
+1 DAMPED_TREND = "damped_trend"
+1 INITIALIZATION_METHOD = "initialization_method"
+1 OPTIMIZED = "optimized"
+1 REMOVE_BIAS = "remove_bias"
+1 USE_BRUTE = "use_brute"
+1
+1 # Arima
+1 ORDER = "order"
+1 SEASONAL_ORDER = "seasonal_order"
+1 METHOD = "method"
+1 MAXITER = "maxiter"
+1 SUPPRESS_WARNINGS = "suppress_warnings"
+1 OUT_OF_SAMPLE_SIZE = "out_of_sample_size"
+1 SCORING = "scoring"
+1 WITH_INTERCEPT = "with_intercept"
+1 TIME_VARYING_REGRESSION = "time_varying_regression"
+1 ENFORCE_STATIONARITY = "enforce_stationarity"
+1 ENFORCE_INVERTIBILITY = "enforce_invertibility"
+1 SIMPLE_DIFFERENCING = "simple_differencing"
+1 MEASUREMENT_ERROR = "measurement_error"
+1 MLE_REGRESSION = "mle_regression"
+1 HAMILTON_REPRESENTATION = "hamilton_representation"
+1 CONCENTRATE_SCALE = "concentrate_scale"
+1
+1 # GAUSSIAN_HMM
+1 N_COMPONENTS = "n_components"
+1 COVARIANCE_TYPE = "covariance_type"
+1 MIN_COVAR = "min_covar"
+1 STARTPROB_PRIOR = "startprob_prior"
+1 TRANSMAT_PRIOR = "transmat_prior"
+1 MEANS_PRIOR = "means_prior"
+1 MEANS_WEIGHT = "means_weight"
+1 COVARS_PRIOR = "covars_prior"
+1 COVARS_WEIGHT = "covars_weight"
+1 ALGORITHM = "algorithm"
+1 N_ITER = "n_iter"
+1 TOL = "tol"
+1 PARAMS = "params"
+1 INIT_PARAMS = "init_params"
+1 IMPLEMENTATION = "implementation"
+1
+1 # GMMHMM
+1 # N_COMPONENTS = "n_components"
+1 N_MIX = "n_mix"
+1 # MIN_COVAR = "min_covar"
+1 # STARTPROB_PRIOR = "startprob_prior"
+1 # TRANSMAT_PRIOR = "transmat_prior"
+1 WEIGHTS_PRIOR = "weights_prior"
+1
+1 # MEANS_PRIOR = "means_prior"
+1 # MEANS_WEIGHT = "means_weight"
+1 # ALGORITHM = "algorithm"
+1 # COVARIANCE_TYPE = "covariance_type"
+1 # N_ITER = "n_iter"
+1 # TOL = "tol"
+1 # INIT_PARAMS = "init_params"
+1 # PARAMS = "params"
+1 # IMPLEMENTATION = "implementation"
+1
+1 # STRAY
+1 ALPHA = "alpha"
+1 K = "k"
+1 KNN_ALGORITHM = "knn_algorithm"
+1 P = "p"
+1 SIZE_THRESHOLD = "size_threshold"
+1 OUTLIER_TAIL = "outlier_tail"
+1
+1 # timerxl
+1 INPUT_TOKEN_LEN = "input_token_len"
+1 HIDDEN_SIZE = "hidden_size"
+1 INTERMEDIATE_SIZE = "intermediate_size"
+1 OUTPUT_TOKEN_LENS = "output_token_lens"
+1 NUM_HIDDEN_LAYERS = "num_hidden_layers"
+1 NUM_ATTENTION_HEADS = "num_attention_heads"
+1 HIDDEN_ACT = "hidden_act"
+1 USE_CACHE = "use_cache"
+1 ROPE_THETA = "rope_theta"
+1 ATTENTION_DROPOUT = "attention_dropout"
+1 INITIALIZER_RANGE = "initializer_range"
+1 MAX_POSITION_EMBEDDINGS = "max_position_embeddings"
+1 CKPT_PATH = "ckpt_path"
+1
+1 # sundial
+1 DROPOUT_RATE = "dropout_rate"
+1 FLOW_LOSS_DEPTH = "flow_loss_depth"
+1 NUM_SAMPLING_STEPS = "num_sampling_steps"
+1 DIFFUSION_BATCH_MUL = "diffusion_batch_mul"
+1
+1 def name(self) -> str:
+1 return self.value
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/exception.py b/iotdb-core/ainode/iotdb/ainode/core/exception.py
index bc89cdc3066..5e1a2774bc1 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/exception.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/exception.py
@@ -1,154 +1,155 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import re
-
-from iotdb.ainode.core.constant import (
- MODEL_CONFIG_FILE_IN_YAML,
- MODEL_WEIGHTS_FILE_IN_PT,
-)
-
-
-class _BaseError(Exception):
- """Base class for exceptions in this module."""
-
- def __init__(self):
- self.message = None
-
- def __str__(self) -> str:
- return self.message
-
-
-class BadNodeUrlError(_BaseError):
- def __init__(self, node_url: str):
- self.message = "Bad node url: {}".format(node_url)
-
-
-class ModelNotExistError(_BaseError):
- def __init__(self, msg: str):
- self.message = "Model is not exists: {} ".format(msg)
-
-
-class BadConfigValueError(_BaseError):
- def __init__(self, config_name: str, config_value, hint: str = ""):
- self.message = "Bad value [{0}] for config {1}. {2}".format(
- config_value, config_name, hint
- )
-
-
-class MissingConfigError(_BaseError):
- def __init__(self, config_name: str):
- self.message = "Missing config: {}".format(config_name)
-
-
-class MissingOptionError(_BaseError):
- def __init__(self, config_name: str):
- self.message = "Missing task option: {}".format(config_name)
-
-
-class RedundantOptionError(_BaseError):
- def __init__(self, option_name: str):
- self.message = "Redundant task option: {}".format(option_name)
-
-
-class WrongTypeConfigError(_BaseError):
- def __init__(self, config_name: str, expected_type: str):
- self.message = "Wrong type for config: {0}, expected: {1}".format(
- config_name, expected_type
- )
-
-
-class UnsupportedError(_BaseError):
- def __init__(self, msg: str):
- self.message = "{0} is not supported in current version".format(msg)
-
-
-class InvalidUriError(_BaseError):
- def __init__(self, uri: str):
- self.message = "Invalid uri: {}, there are no {} or {} under this uri.".format(
- uri, MODEL_WEIGHTS_FILE_IN_PT, MODEL_CONFIG_FILE_IN_YAML
- )
-
-
-class InvalidWindowArgumentError(_BaseError):
- def __init__(self, window_interval, window_step, dataset_length):
- self.message = f"Invalid inference input: window_interval {window_interval}, window_step {window_step}, dataset_length {dataset_length}"
-
-
-class InferenceModelInternalError(_BaseError):
- def __init__(self, msg: str):
- self.message = "Inference model internal error: {0}".format(msg)
-
-
-class BuiltInModelNotSupportError(_BaseError):
- def __init__(self, msg: str):
- self.message = "Built-in model not support: {0}".format(msg)
-
-
-class BuiltInModelDeletionError(_BaseError):
- def __init__(self, model_id: str):
- self.message = "Cannot delete built-in model: {0}".format(model_id)
-
-
-class WrongAttributeTypeError(_BaseError):
- def __init__(self, attribute_name: str, expected_type: str):
- self.message = "Wrong type for attribute: {0}, expected: {1}".format(
- attribute_name, expected_type
- )
-
-
-class NumericalRangeException(_BaseError):
- def __init__(self, attribute_name: str, value, min_value, max_value):
- self.message = (
- "Attribute {0} expect value between {1} and {2}, got {3} instead.".format(
- attribute_name, min_value, max_value, value
- )
- )
-
-
-class StringRangeException(_BaseError):
- def __init__(self, attribute_name: str, value: str, expect_value):
- self.message = "Attribute {0} expect value in {1}, got {2} instead.".format(
- attribute_name, expect_value, value
- )
-
-
-class ListRangeException(_BaseError):
- def __init__(self, attribute_name: str, value: list, expected_type: str):
- self.message = (
- "Attribute {0} expect value type list[{1}], got {2} instead.".format(
- attribute_name, expected_type, value
- )
- )
-
-
-class AttributeNotSupportError(_BaseError):
- def __init__(self, model_name: str, attribute_name: str):
- self.message = "Attribute {0} is not supported in model {1}".format(
- attribute_name, model_name
- )
-
-
-# This is used to extract the key message in RuntimeError instead of the traceback message
-def runtime_error_extractor(error_message):
- pattern = re.compile(r"RuntimeError: (.+)")
- match = pattern.search(error_message)
-
- if match:
- return match.group(1)
- else:
- return ""
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import re
+1
+1from iotdb.ainode.core.constant import (
+1 MODEL_CONFIG_FILE_IN_YAML,
+1 MODEL_WEIGHTS_FILE_IN_PT,
+1)
+1
+1
+1class _BaseError(Exception):
+1 """Base class for exceptions in this module."""
+1
+1 def __init__(self):
+1 self.message = None
+1
+1 def __str__(self) -> str:
+1 return self.message
+1
+1
+1class BadNodeUrlError(_BaseError):
+1 def __init__(self, node_url: str):
+1 self.message = "Bad node url: {}".format(node_url)
+1
+1
+1class ModelNotExistError(_BaseError):
+1 def __init__(self, msg: str):
+1 self.message = "Model is not exists: {} ".format(msg)
+1
+1
+1class BadConfigValueError(_BaseError):
+1 def __init__(self, config_name: str, config_value, hint: str = ""):
+1 self.message = "Bad value [{0}] for config {1}. {2}".format(
+1 config_value, config_name, hint
+1 )
+1
+1
+1class MissingConfigError(_BaseError):
+1 def __init__(self, config_name: str):
+1 self.message = "Missing config: {}".format(config_name)
+1
+1
+1class MissingOptionError(_BaseError):
+1 def __init__(self, config_name: str):
+1 self.message = "Missing task option: {}".format(config_name)
+1
+1
+1class RedundantOptionError(_BaseError):
+1 def __init__(self, option_name: str):
+1 self.message = "Redundant task option: {}".format(option_name)
+1
+1
+1class WrongTypeConfigError(_BaseError):
+1 def __init__(self, config_name: str, expected_type: str):
+1 self.message = "Wrong type for config: {0}, expected: {1}".format(
+1 config_name, expected_type
+1 )
+1
+1
+1class UnsupportedError(_BaseError):
+1 def __init__(self, msg: str):
+1 self.message = "{0} is not supported in current version".format(msg)
+1
+1
+1class InvalidUriError(_BaseError):
+1 def __init__(self, uri: str):
+1 self.message = "Invalid uri: {}, there are no {} or {} under this uri.".format(
+1 uri, MODEL_WEIGHTS_FILE_IN_PT, MODEL_CONFIG_FILE_IN_YAML
+1 )
+1
+1
+1class InvalidWindowArgumentError(_BaseError):
+1 def __init__(self, window_interval, window_step, dataset_length):
+1 self.message = f"Invalid inference input: window_interval {window_interval}, window_step {window_step}, dataset_length {dataset_length}"
+1
+1
+1class InferenceModelInternalError(_BaseError):
+1 def __init__(self, msg: str):
+1 self.message = "Inference model internal error: {0}".format(msg)
+1
+1
+1class BuiltInModelNotSupportError(_BaseError):
+1 def __init__(self, msg: str):
+1 self.message = "Built-in model not support: {0}".format(msg)
+1
+1
+1class BuiltInModelDeletionError(_BaseError):
+1 def __init__(self, model_id: str):
+1 self.message = "Cannot delete built-in model: {0}".format(model_id)
+1
+1
+1class WrongAttributeTypeError(_BaseError):
+1 def __init__(self, attribute_name: str, expected_type: str):
+1 self.message = "Wrong type for attribute: {0}, expected: {1}".format(
+1 attribute_name, expected_type
+1 )
+1
+1
+1class NumericalRangeException(_BaseError):
+1 def __init__(self, attribute_name: str, value, min_value, max_value):
+1 self.message = (
+1 "Attribute {0} expect value between {1} and {2}, got {3} instead.".format(
+1 attribute_name, min_value, max_value, value
+1 )
+1 )
+1
+1
+1class StringRangeException(_BaseError):
+1 def __init__(self, attribute_name: str, value: str, expect_value):
+1 self.message = "Attribute {0} expect value in {1}, got {2} instead.".format(
+1 attribute_name, expect_value, value
+1 )
+1
+1
+1class ListRangeException(_BaseError):
+1 def __init__(self, attribute_name: str, value: list, expected_type: str):
+1 self.message = (
+1 "Attribute {0} expect value type list[{1}], got {2} instead.".format(
+1 attribute_name, expected_type, value
+1 )
+1 )
+1
+1
+1class AttributeNotSupportError(_BaseError):
+1 def __init__(self, model_name: str, attribute_name: str):
+1 self.message = "Attribute {0} is not supported in model {1}".format(
+1 attribute_name, model_name
+1 )
+1
+1
+1# This is used to extract the key message in RuntimeError instead of the traceback message
+1def runtime_error_extractor(error_message):
+1 pattern = re.compile(r"RuntimeError: (.+)")
+1 match = pattern.search(error_message)
+1
+1 if match:
+1 return match.group(1)
+1 else:
+1 return ""
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/__init__.py
similarity index 99%
copy from iotdb-core/ainode/iotdb/ainode/__init__.py
copy to iotdb-core/ainode/iotdb/ainode/core/inference/batcher/__init__.py
index 2a1e720805f..4b8ee97fad2 100644
--- a/iotdb-core/ainode/iotdb/ainode/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/__init__.py
@@ -14,4 +14,4 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-#
+#
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py
index 7b668e93a8f..b12b9cd2d69 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py
@@ -1,42 +1,43 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-from typing import List
-
-from iotdb.ainode.core.inference.inference_request import InferenceRequest
-
-
-class AbstractBatcher(ABC):
- """
- Abstract base class for batchers that batch inference requests.
- """
-
- def __init__(self):
- """
- Args:
-
- """
- pass
-
- @abstractmethod
- def batch_request(self, reqs: List[InferenceRequest]):
- """
- batch given requests, such that they can be delivered to the model and be executed concurrently.
- """
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from abc import ABC, abstractmethod
+1from typing import List
+1
+1from iotdb.ainode.core.inference.inference_request import InferenceRequest
+1
+1
+1class AbstractBatcher(ABC):
+1 """
+1 Abstract base class for batchers that batch inference requests.
+1 """
+1
+1 def __init__(self):
+1 """
+1 Args:
+1
+1 """
+1 pass
+1
+1 @abstractmethod
+1 def batch_request(self, reqs: List[InferenceRequest]):
+1 """
+1 batch given requests, such that they can be delivered to the model and be executed concurrently.
+1 """
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py
index 591a0d7c1dd..7a12ea2ba57 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py
@@ -1,63 +1,64 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from typing import List
-
-import torch
-
-from iotdb.ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
-from iotdb.ainode.core.inference.inference_request import InferenceRequest
-
-
-class BasicBatcher(AbstractBatcher):
- """
- Basic batcher for inference requests.
- """
-
- def __init__(self):
- """
- Args:
-
- """
-
- def batch_request(self, reqs: List[InferenceRequest]) -> torch.Tensor:
- """
- Batch given requests by simply concatenating their inputs, only requests with uniformed output length can be batched.
-
- - Considering the current implementation of AINode, we might merely be piecing together the input for now.
-
- Args:
- reqs (List[InferenceRequest]): List of inference requests.
-
- Returns:
- torch.Tensor: Concatenated input tensor of shape
- [sum(req.batch_size), length].
- """
- if not reqs:
- raise ValueError("No requests provided to batch_request.")
-
- # Ensure length consistency
- length_set = {req.inputs.shape[1] for req in reqs}
- if len(length_set) != 1:
- raise ValueError(
- f"All requests must have the same length, " f"but got {length_set}"
- )
-
- batch_inputs = torch.cat([req.inputs for req in reqs], dim=0)
-
- return batch_inputs
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from typing import List
+1
+1import torch
+1
+1from iotdb.ainode.core.inference.batcher.abstract_batcher import AbstractBatcher
+1from iotdb.ainode.core.inference.inference_request import InferenceRequest
+1
+1
+1class BasicBatcher(AbstractBatcher):
+1 """
+1 Basic batcher for inference requests.
+1 """
+1
+1 def __init__(self):
+1 """
+1 Args:
+1
+1 """
+1
+1 def batch_request(self, reqs: List[InferenceRequest]) -> torch.Tensor:
+1 """
+1 Batch given requests by simply concatenating their inputs, only requests with uniformed output length can be batched.
+1
+1 - Considering the current implementation of AINode, we might merely be piecing together the input for now.
+1
+1 Args:
+1 reqs (List[InferenceRequest]): List of inference requests.
+1
+1 Returns:
+1 torch.Tensor: Concatenated input tensor of shape
+1 [sum(req.batch_size), length].
+1 """
+1 if not reqs:
+1 raise ValueError("No requests provided to batch_request.")
+1
+1 # Ensure length consistency
+1 length_set = {req.inputs.shape[1] for req in reqs}
+1 if len(length_set) != 1:
+1 raise ValueError(
+1 f"All requests must have the same length, " f"but got {length_set}"
+1 )
+1
+1 batch_inputs = torch.cat([req.inputs for req in reqs], dim=0)
+1
+1 return batch_inputs
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py
index 3abc1477a97..e2fca5c188a 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py
@@ -1,43 +1,44 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-from typing import Dict
-
-from iotdb.ainode.core.inference.inference_request import InferenceRequest
-from iotdb.ainode.core.inference.inference_request_pool import PoolState
-
-
-class AbstractDispatcher(ABC):
- """
- Abstract base class for dispatchers that handle inference requests.
- """
-
- def __init__(self, pool_states: Dict[int, PoolState]):
- """
- Args:
- pool_states: Dictionary containing the states of inference request pools in the same pool group.
- """
- self.pool_states = pool_states
-
- @abstractmethod
- def dispatch_request(self, req: InferenceRequest, pool_ids: list[int]) -> int:
- """
- Dispatch an inference request to the appropriate pool.
- """
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from abc import ABC, abstractmethod
+1from typing import Dict
+1
+1from iotdb.ainode.core.inference.inference_request import InferenceRequest
+1from iotdb.ainode.core.inference.inference_request_pool import PoolState
+1
+1
+1class AbstractDispatcher(ABC):
+1 """
+1 Abstract base class for dispatchers that handle inference requests.
+1 """
+1
+1 def __init__(self, pool_states: Dict[int, PoolState]):
+1 """
+1 Args:
+1 pool_states: Dictionary containing the states of inference request pools in the same pool group.
+1 """
+1 self.pool_states = pool_states
+1
+1 @abstractmethod
+1 def dispatch_request(self, req: InferenceRequest, pool_ids: list[int]) -> int:
+1 """
+1 Dispatch an inference request to the appropriate pool.
+1 """
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
index d06ba55546f..6e14a974c73 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
@@ -1,60 +1,61 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.dispatcher.abstract_dispatcher import (
- AbstractDispatcher,
-)
-from iotdb.ainode.core.inference.inference_request import InferenceRequest
-from iotdb.ainode.core.inference.inference_request_pool import PoolState
-from iotdb.ainode.core.log import Logger
-
-logger = Logger()
-
-
-class BasicDispatcher(AbstractDispatcher):
- """
- Basic dispatcher for inference requests.
- """
-
- def __init__(self, pool_states):
- super().__init__(pool_states)
-
- def _select_pool_by_hash(self, req, pool_ids) -> int:
- """
- Select a pool for the given request using a hash-based approach.
- """
- model_id = req.model_id
- if not pool_ids:
- raise InferenceModelInternalError(
- f"No available pools for model {model_id}"
- )
- start_idx = hash(req.req_id) % len(pool_ids)
- n = len(pool_ids)
- for i in range(n):
- pool_id = pool_ids[(start_idx + i) % n]
- state = self.pool_states[pool_id]
- if state == PoolState.RUNNING:
- return pool_id
- raise InferenceModelInternalError(
- f"No RUNNING pools available for model {model_id}"
- )
-
- def dispatch_request(self, req: InferenceRequest, pool_ids: list[int]) -> int:
- pool_idx = self._select_pool_by_hash(req, pool_ids)
- return pool_idx
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.dispatcher.abstract_dispatcher import (
+1 AbstractDispatcher,
+1)
+1from iotdb.ainode.core.inference.inference_request import InferenceRequest
+1from iotdb.ainode.core.inference.inference_request_pool import PoolState
+1from iotdb.ainode.core.log import Logger
+1
+1logger = Logger()
+1
+1
+1class BasicDispatcher(AbstractDispatcher):
+1 """
+1 Basic dispatcher for inference requests.
+1 """
+1
+1 def __init__(self, pool_states):
+1 super().__init__(pool_states)
+1
+1 def _select_pool_by_hash(self, req, pool_ids) -> int:
+1 """
+1 Select a pool for the given request using a hash-based approach.
+1 """
+1 model_id = req.model_id
+1 if not pool_ids:
+1 raise InferenceModelInternalError(
+1 f"No available pools for model {model_id}"
+1 )
+1 start_idx = hash(req.req_id) % len(pool_ids)
+1 n = len(pool_ids)
+1 for i in range(n):
+1 pool_id = pool_ids[(start_idx + i) % n]
+1 state = self.pool_states[pool_id]
+1 if state == PoolState.RUNNING:
+1 return pool_id
+1 raise InferenceModelInternalError(
+1 f"No RUNNING pools available for model {model_id}"
+1 )
+1
+1 def dispatch_request(self, req: InferenceRequest, pool_ids: list[int]) -> int:
+1 pool_idx = self._select_pool_by_hash(req, pool_ids)
+1 return pool_idx
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
index 82c72cc37ab..a26d0074a51 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
@@ -1,132 +1,133 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import threading
-from typing import Any
-
-import torch
-
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.util.atmoic_int import AtomicInt
-
-logger = Logger()
-
-
-class InferenceRequestState:
- WAITING = "waiting"
- RUNNING = "running"
- FINISHED = "finished"
-
-
-class InferenceRequest:
- def __init__(
- self,
- req_id: str,
- model_id: str,
- inputs: torch.Tensor,
- inference_pipeline: AbstractInferencePipeline,
- max_new_tokens: int = 96,
- **infer_kwargs,
- ):
- if inputs.ndim == 1:
- inputs = inputs.unsqueeze(0)
-
- self.req_id = req_id
- self.model_id = model_id
- self.inputs = inputs
- self.infer_kwargs = infer_kwargs
- self.inference_pipeline = inference_pipeline
- self.max_new_tokens = (
- max_new_tokens # Number of time series data points to generate
- )
-
- self.batch_size = inputs.size(0)
- self.state = InferenceRequestState.WAITING
- self.cur_step_idx = 0 # Current write position in the output step index
- self.assigned_pool_id = -1 # The pool handling this request
- self.assigned_device_id = -1 # The device handling this request
-
- # Preallocate output buffer [batch_size, max_new_tokens]
- self.output_tensor = torch.zeros(
- self.batch_size, max_new_tokens, device="cpu"
- ) # shape: [self.batch_size, max_new_steps]
-
- def mark_running(self):
- self.state = InferenceRequestState.RUNNING
-
- def mark_finished(self):
- self.state = InferenceRequestState.FINISHED
-
- def is_finished(self) -> bool:
- return (
- self.state == InferenceRequestState.FINISHED
- or self.cur_step_idx >= self.max_new_tokens
- )
-
- def write_step_output(self, step_output: torch.Tensor):
- if step_output.ndim == 1:
- step_output = step_output.unsqueeze(0)
-
- batch_size, step_size = step_output.shape
- end_idx = self.cur_step_idx + step_size
-
- if end_idx > self.max_new_tokens:
- self.output_tensor[:, self.cur_step_idx :] = step_output[
- :, : self.max_new_tokens - self.cur_step_idx
- ]
- self.cur_step_idx = self.max_new_tokens
- else:
- self.output_tensor[:, self.cur_step_idx : end_idx] = step_output
- self.cur_step_idx = end_idx
-
- if self.is_finished():
- self.mark_finished()
-
- def get_final_output(self) -> torch.Tensor:
- return self.output_tensor[:, : self.cur_step_idx]
-
-
-class InferenceRequestProxy:
- """
- Wrap the raw request for handling multiprocess processing.
- """
-
- def __init__(self, req_id: str):
- self.req_id = req_id
- self.result = None
- self._counter: AtomicInt = None
- self._lock = threading.Lock()
- self._condition = threading.Condition(self._lock)
-
- def set_result(self, result: Any):
- with self._lock:
- self.result = result
- if self._counter is not None:
- self._counter.decrement_and_get()
- self._condition.notify_all()
-
- def set_counter(self, counter: AtomicInt):
- with self._lock:
- self._counter = counter
-
- def wait_for_result(self) -> Any:
- with self._lock:
- self._condition.wait()
- return self.result
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import threading
+1from typing import Any
+1
+1import torch
+1
+1from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
+1 AbstractInferencePipeline,
+1)
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.util.atmoic_int import AtomicInt
+1
+1logger = Logger()
+1
+1
+1class InferenceRequestState:
+1 WAITING = "waiting"
+1 RUNNING = "running"
+1 FINISHED = "finished"
+1
+1
+1class InferenceRequest:
+1 def __init__(
+1 self,
+1 req_id: str,
+1 model_id: str,
+1 inputs: torch.Tensor,
+1 inference_pipeline: AbstractInferencePipeline,
+1 max_new_tokens: int = 96,
+1 **infer_kwargs,
+1 ):
+1 if inputs.ndim == 1:
+1 inputs = inputs.unsqueeze(0)
+1
+1 self.req_id = req_id
+1 self.model_id = model_id
+1 self.inputs = inputs
+1 self.infer_kwargs = infer_kwargs
+1 self.inference_pipeline = inference_pipeline
+1 self.max_new_tokens = (
+1 max_new_tokens # Number of time series data points to generate
+1 )
+1
+1 self.batch_size = inputs.size(0)
+1 self.state = InferenceRequestState.WAITING
+1 self.cur_step_idx = 0 # Current write position in the output step index
+1 self.assigned_pool_id = -1 # The pool handling this request
+1 self.assigned_device_id = -1 # The device handling this request
+1
+1 # Preallocate output buffer [batch_size, max_new_tokens]
+1 self.output_tensor = torch.zeros(
+1 self.batch_size, max_new_tokens, device="cpu"
+1 ) # shape: [self.batch_size, max_new_steps]
+1
+1 def mark_running(self):
+1 self.state = InferenceRequestState.RUNNING
+1
+1 def mark_finished(self):
+1 self.state = InferenceRequestState.FINISHED
+1
+1 def is_finished(self) -> bool:
+1 return (
+1 self.state == InferenceRequestState.FINISHED
+1 or self.cur_step_idx >= self.max_new_tokens
+1 )
+1
+1 def write_step_output(self, step_output: torch.Tensor):
+1 if step_output.ndim == 1:
+1 step_output = step_output.unsqueeze(0)
+1
+1 batch_size, step_size = step_output.shape
+1 end_idx = self.cur_step_idx + step_size
+1
+1 if end_idx > self.max_new_tokens:
+1 self.output_tensor[:, self.cur_step_idx :] = step_output[
+1 :, : self.max_new_tokens - self.cur_step_idx
+1 ]
+1 self.cur_step_idx = self.max_new_tokens
+1 else:
+1 self.output_tensor[:, self.cur_step_idx : end_idx] = step_output
+1 self.cur_step_idx = end_idx
+1
+1 if self.is_finished():
+1 self.mark_finished()
+1
+1 def get_final_output(self) -> torch.Tensor:
+1 return self.output_tensor[:, : self.cur_step_idx]
+1
+1
+1class InferenceRequestProxy:
+1 """
+1 Wrap the raw request for handling multiprocess processing.
+1 """
+1
+1 def __init__(self, req_id: str):
+1 self.req_id = req_id
+1 self.result = None
+1 self._counter: AtomicInt = None
+1 self._lock = threading.Lock()
+1 self._condition = threading.Condition(self._lock)
+1
+1 def set_result(self, result: Any):
+1 with self._lock:
+1 self.result = result
+1 if self._counter is not None:
+1 self._counter.decrement_and_get()
+1 self._condition.notify_all()
+1
+1 def set_counter(self, counter: AtomicInt):
+1 with self._lock:
+1 self._counter = counter
+1
+1 def wait_for_result(self) -> Any:
+1 with self._lock:
+1 self._condition.wait()
+1 return self.result
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
index b5f26c98358..9e92aceafcc 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
@@ -1,217 +1,218 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import random
-import threading
-import time
-from collections import defaultdict
-from enum import Enum
-
-import numpy as np
-import torch
-import torch.multiprocessing as mp
-from transformers import PretrainedConfig
-
-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.constant import INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE
-from iotdb.ainode.core.inference.batcher.basic_batcher import BasicBatcher
-from iotdb.ainode.core.inference.inference_request import InferenceRequest
-from iotdb.ainode.core.inference.request_scheduler.basic_request_scheduler import (
- BasicRequestScheduler,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.model_manager import ModelManager
-from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
-
-
-class PoolState(Enum):
- INITIALIZING = "INITIALIZING"
- RUNNING = "RUNNING"
- STOPPING = "STOPPING"
-
-
-class InferenceRequestPool(mp.Process):
- """
- The request pool to handle inference for a specific model.
- """
-
- FIX_SEED = 2021
- WAITING_INTERVAL_IN_MS = (
- AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms()
- ) # How often to check for requests in the waiting/running queue
-
- def __init__(
- self,
- pool_id: int,
- model_id: str,
- device: str,
- config: PretrainedConfig,
- request_queue: mp.Queue,
- result_queue: mp.Queue,
- ready_event,
- **pool_kwargs,
- ):
- super().__init__()
- self.pool_id = pool_id
- self.model_id = model_id
- self.config = config
- self.pool_kwargs = pool_kwargs
- self.ready_event = ready_event
- self.device = convert_device_id_to_torch_device(device)
-
- self._threads = []
- self._waiting_queue = request_queue # Requests that are waiting to be processed
- self._running_queue = mp.Queue() # Requests that are currently being processed
- self._finished_queue = result_queue # Requests that are finished
- self._request_scheduler = BasicRequestScheduler(
- self._waiting_queue, self._running_queue, self._finished_queue, self.pool_id
- )
- self._batcher = BasicBatcher()
- self._stop_event = mp.Event()
-
- self._model = None
- self._model_manager = None
- self._logger = None
-
- # Fix inference seed
- random.seed(self.FIX_SEED)
- torch.manual_seed(self.FIX_SEED)
- np.random.seed(self.FIX_SEED)
-
- def _activate_requests(self):
- requests = self._request_scheduler.schedule_activate()
- for request in requests:
- request.inputs = request.inference_pipeline.preprocess_inputs(
- request.inputs
- )
- request.mark_running()
- self._running_queue.put(request)
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][Req-{request.req_id}] Request is activated with inputs shape {request.inputs.shape}"
- )
-
- def _requests_activate_loop(self):
- while not self._stop_event.is_set():
- time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
- self._activate_requests()
-
- def _step(self):
- all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
-
- grouped_requests = defaultdict(list)
- for req in all_requests:
- key = (req.inputs.shape[1], req.max_new_tokens)
- grouped_requests[key].append(req)
- grouped_requests = list(grouped_requests.values())
-
- for requests in grouped_requests:
- batch_inputs = self._batcher.batch_request(requests).to(self.device)
- if self.model_id == "sundial":
- batch_output = self._model.generate(
- batch_inputs,
- max_new_tokens=requests[0].max_new_tokens,
- num_samples=10,
- revin=True,
- )
-
- offset = 0
- for request in requests:
- request.output_tensor = request.output_tensor.to(self.device)
- cur_batch_size = request.batch_size
- cur_output = batch_output[offset : offset + cur_batch_size]
- offset += cur_batch_size
- # TODO Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here, such as: `cur_output[0]` maybe not true in multi-variable scene
- request.write_step_output(cur_output[0].mean(dim=0))
-
- request.inference_pipeline.post_decode()
- if request.is_finished():
- request.inference_pipeline.post_inference()
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
- )
- # ensure the output tensor is on CPU before sending to result queue
- request.output_tensor = request.output_tensor.cpu()
- self._finished_queue.put(request)
- else:
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
- )
- self._waiting_queue.put(request)
-
- elif self.model_id == "timer_xl":
- batch_output = self._model.generate(
- batch_inputs,
- max_new_tokens=requests[0].max_new_tokens,
- revin=True,
- )
-
- offset = 0
- for request in requests:
- request.output_tensor = request.output_tensor.to(self.device)
- cur_batch_size = request.batch_size
- cur_output = batch_output[offset : offset + cur_batch_size]
- offset += cur_batch_size
- request.write_step_output(cur_output)
-
- request.inference_pipeline.post_decode()
- if request.is_finished():
- request.inference_pipeline.post_inference()
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
- )
- # ensure the output tensor is on CPU before sending to result queue
- request.output_tensor = request.output_tensor.cpu()
- self._finished_queue.put(request)
- else:
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
- )
- self._waiting_queue.put(request)
-
- def _requests_execute_loop(self):
- while not self._stop_event.is_set():
- time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
- self._step()
-
- def run(self):
- self._logger = Logger(
- INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE.format(self.device)
- )
- self._model_manager = ModelManager()
- self._request_scheduler.device = self.device
- self._model = self._model_manager.load_model(self.model_id, {}).to(self.device)
- self.ready_event.set()
-
- activate_daemon = threading.Thread(
- target=self._requests_activate_loop, daemon=True
- )
- self._threads.append(activate_daemon)
- activate_daemon.start()
- execute_daemon = threading.Thread(
- target=self._requests_execute_loop, daemon=True
- )
- self._threads.append(execute_daemon)
- execute_daemon.start()
- for thread in self._threads:
- thread.join()
- self._logger.info(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}] InferenceRequestPool for model {self.model_id} exited cleanly."
- )
-
- def stop(self):
- self._stop_event.set()
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import random
+1import threading
+1import time
+1from collections import defaultdict
+1from enum import Enum
+1
+1import numpy as np
+1import torch
+1import torch.multiprocessing as mp
+1from transformers import PretrainedConfig
+1
+1from iotdb.ainode.core.config import AINodeDescriptor
+1from iotdb.ainode.core.constant import INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE
+1from iotdb.ainode.core.inference.batcher.basic_batcher import BasicBatcher
+1from iotdb.ainode.core.inference.inference_request import InferenceRequest
+1from iotdb.ainode.core.inference.request_scheduler.basic_request_scheduler import (
+1 BasicRequestScheduler,
+1)
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.manager.model_manager import ModelManager
+1from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
+1
+1
+1class PoolState(Enum):
+1 INITIALIZING = "INITIALIZING"
+1 RUNNING = "RUNNING"
+1 STOPPING = "STOPPING"
+1
+1
+1class InferenceRequestPool(mp.Process):
+1 """
+1 The request pool to handle inference for a specific model.
+1 """
+1
+1 FIX_SEED = 2021
+1 WAITING_INTERVAL_IN_MS = (
+1 AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms()
+1 ) # How often to check for requests in the waiting/running queue
+1
+1 def __init__(
+1 self,
+1 pool_id: int,
+1 model_id: str,
+1 device: str,
+1 config: PretrainedConfig,
+1 request_queue: mp.Queue,
+1 result_queue: mp.Queue,
+1 ready_event,
+1 **pool_kwargs,
+1 ):
+1 super().__init__()
+1 self.pool_id = pool_id
+1 self.model_id = model_id
+1 self.config = config
+1 self.pool_kwargs = pool_kwargs
+1 self.ready_event = ready_event
+1 self.device = convert_device_id_to_torch_device(device)
+1
+1 self._threads = []
+1 self._waiting_queue = request_queue # Requests that are waiting to be processed
+1 self._running_queue = mp.Queue() # Requests that are currently being processed
+1 self._finished_queue = result_queue # Requests that are finished
+1 self._request_scheduler = BasicRequestScheduler(
+1 self._waiting_queue, self._running_queue, self._finished_queue, self.pool_id
+1 )
+1 self._batcher = BasicBatcher()
+1 self._stop_event = mp.Event()
+1
+1 self._model = None
+1 self._model_manager = None
+1 self._logger = None
+1
+1 # Fix inference seed
+1 random.seed(self.FIX_SEED)
+1 torch.manual_seed(self.FIX_SEED)
+1 np.random.seed(self.FIX_SEED)
+1
+1 def _activate_requests(self):
+1 requests = self._request_scheduler.schedule_activate()
+1 for request in requests:
+1 request.inputs = request.inference_pipeline.preprocess_inputs(
+1 request.inputs
+1 )
+1 request.mark_running()
+1 self._running_queue.put(request)
+1 self._logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][Req-{request.req_id}] Request is activated with inputs shape {request.inputs.shape}"
+1 )
+1
+1 def _requests_activate_loop(self):
+1 while not self._stop_event.is_set():
+1 time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
+1 self._activate_requests()
+1
+1 def _step(self):
+1 all_requests: list[InferenceRequest] = self._request_scheduler.schedule_step()
+1
+1 grouped_requests = defaultdict(list)
+1 for req in all_requests:
+1 key = (req.inputs.shape[1], req.max_new_tokens)
+1 grouped_requests[key].append(req)
+1 grouped_requests = list(grouped_requests.values())
+1
+1 for requests in grouped_requests:
+1 batch_inputs = self._batcher.batch_request(requests).to(self.device)
+1 if self.model_id == "sundial":
+1 batch_output = self._model.generate(
+1 batch_inputs,
+1 max_new_tokens=requests[0].max_new_tokens,
+1 num_samples=10,
+1 revin=True,
+1 )
+1
+1 offset = 0
+1 for request in requests:
+1 request.output_tensor = request.output_tensor.to(self.device)
+1 cur_batch_size = request.batch_size
+1 cur_output = batch_output[offset : offset + cur_batch_size]
+1 offset += cur_batch_size
+1 # TODO Here we only considered the case where batchsize=1 in one request. If multi-variable adaptation is required in the future, modifications may be needed here, such as: `cur_output[0]` maybe not true in multi-variable scene
+1 request.write_step_output(cur_output[0].mean(dim=0))
+1
+1 request.inference_pipeline.post_decode()
+1 if request.is_finished():
+1 request.inference_pipeline.post_inference()
+1 self._logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+1 )
+1 # ensure the output tensor is on CPU before sending to result queue
+1 request.output_tensor = request.output_tensor.cpu()
+1 self._finished_queue.put(request)
+1 else:
+1 self._logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+1 )
+1 self._waiting_queue.put(request)
+1
+1 elif self.model_id == "timer_xl":
+1 batch_output = self._model.generate(
+1 batch_inputs,
+1 max_new_tokens=requests[0].max_new_tokens,
+1 revin=True,
+1 )
+1
+1 offset = 0
+1 for request in requests:
+1 request.output_tensor = request.output_tensor.to(self.device)
+1 cur_batch_size = request.batch_size
+1 cur_output = batch_output[offset : offset + cur_batch_size]
+1 offset += cur_batch_size
+1 request.write_step_output(cur_output)
+1
+1 request.inference_pipeline.post_decode()
+1 if request.is_finished():
+1 request.inference_pipeline.post_inference()
+1 self._logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+1 )
+1 # ensure the output tensor is on CPU before sending to result queue
+1 request.output_tensor = request.output_tensor.cpu()
+1 self._finished_queue.put(request)
+1 else:
+1 self._logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+1 )
+1 self._waiting_queue.put(request)
+1
+1 def _requests_execute_loop(self):
+1 while not self._stop_event.is_set():
+1 time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
+1 self._step()
+1
+1 def run(self):
+1 self._logger = Logger(
+1 INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE.format(self.device)
+1 )
+1 self._model_manager = ModelManager()
+1 self._request_scheduler.device = self.device
+1 self._model = self._model_manager.load_model(self.model_id, {}).to(self.device)
+1 self.ready_event.set()
+1
+1 activate_daemon = threading.Thread(
+1 target=self._requests_activate_loop, daemon=True
+1 )
+1 self._threads.append(activate_daemon)
+1 activate_daemon.start()
+1 execute_daemon = threading.Thread(
+1 target=self._requests_execute_loop, daemon=True
+1 )
+1 self._threads.append(execute_daemon)
+1 execute_daemon.start()
+1 for thread in self._threads:
+1 thread.join()
+1 self._logger.info(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}] InferenceRequestPool for model {self.model_id} exited cleanly."
+1 )
+1
+1 def stop(self):
+1 self._stop_event.set()
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
index ce04120474e..2a06fd0c8cf 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
@@ -1,509 +1,510 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import concurrent
-import queue
-import random
-import threading
-from concurrent.futures import wait
-from typing import Dict, Optional
-
-import torch
-import torch.multiprocessing as mp
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.inference_request import (
- InferenceRequest,
- InferenceRequestProxy,
-)
-from iotdb.ainode.core.inference.inference_request_pool import (
- InferenceRequestPool,
- PoolState,
-)
-from iotdb.ainode.core.inference.pool_group import PoolGroup
-from iotdb.ainode.core.inference.pool_scheduler.basic_pool_scheduler import (
- BasicPoolScheduler,
- ScaleActionType,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
-from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
-from iotdb.ainode.core.util.atmoic_int import AtomicInt
-from iotdb.ainode.core.util.batch_executor import BatchExecutor
-from iotdb.ainode.core.util.decorator import synchronized
-from iotdb.ainode.core.util.thread_name import ThreadName
-
-logger = Logger()
-
-
-class PoolController:
- """
- A controller for handling inference request pools.
- """
-
- def __init__(self, result_queue: mp.Queue):
- # structure: {model_id: {device_id: PoolGroup}}
- self._request_pool_map: Dict[str, Dict[str, PoolGroup]] = {}
- self._new_pool_id = AtomicInt()
- self._result_queue = result_queue
- self._pool_scheduler = BasicPoolScheduler(self._request_pool_map)
- self._stop_event = threading.Event()
-
- # for pool instances control
- self._task_queue = queue.Queue()
- self._pool_control_worker_thread = threading.Thread(
- target=self._worker_loop, daemon=True
- )
- self._pool_control_worker_thread.start()
- self._executor = BatchExecutor(
- thread_name_prefix=ThreadName.INFERENCE_POOL_CONTROLLER.value
- )
-
- # =============== Pool Management ===============
- @synchronized(threading.Lock())
- def first_req_init(self, model_id: str):
- """
- Initialize the pools when the first request for the given model_id arrives.
- """
- if not self.has_request_pools(model_id, device.index):
- # TODO: choose a device based on some strategy
- device = self.DEFAULT_DEVICE
- actions = self._pool_scheduler.schedule(model_id, device)
- for action in actions:
- if action.action == ScaleActionType.SCALE_UP:
- # initialize the first pool
- self._first_pool_init(action.model_id, str(device))
- # start a background thread to expand pools
- expand_thread = threading.Thread(
- target=self._expand_pools_on_device,
- args=(action.model_id, str(device), action.amount - 1),
- daemon=True,
- )
- expand_thread.start()
- elif action.action == ScaleActionType.SCALE_DOWN:
- # TODO: implement scale down logic
- pass
-
- def _first_pool_init(self, model_id: str, device_str: str):
- """
- Initialize the first pool for the given model_id.
- Ensure the pool is ready before returning.
- """
- device = torch.device(device_str)
- device_id = device.index
-
- 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,
- device=device_str,
- config=config,
- request_queue=first_queue,
- result_queue=self._result_queue,
- ready_event=ready_event,
- )
- first_pool.start()
- self._register_pool(model_id, device_str, 0, first_pool, first_queue)
-
- if not ready_event.wait(timeout=30):
- self._erase_pool(model_id, device_id, 0)
- logger.error(
- f"[Inference][Device-{device}][Pool-0] Pool failed to be ready in time"
- )
- else:
- self.set_state(model_id, device_id, 0, PoolState.RUNNING)
- logger.info(
- f"[Inference][Device-{device}][Pool-0] Pool started running for model {model_id}"
- )
-
- def load_model(self, model_id: str, device_id_list: list[str]):
- """
- Load the model to the specified devices asynchronously.
- Args:
- model_id (str): The ID of the model to be loaded.
- device_id_list (list[str]): List of device_ids where the model should be loaded.
- """
- self._task_queue.put((self._load_model_task, (model_id, device_id_list), {}))
-
- def unload_model(self, model_id: str, device_id_list: list[str]):
- """
- Unload the model from the specified devices asynchronously.
- Args:
- model_id (str): The ID of the model to be unloaded.
- device_id_list (list[str]): List of device_ids where the model should be unloaded.
- """
- self._task_queue.put((self._unload_model_task, (model_id, device_id_list), {}))
-
- def show_loaded_models(
- self, device_id_list: list[str]
- ) -> Dict[str, Dict[str, int]]:
- """
- Show loaded model instances on the specified devices.
- Args:
- device_id_list (list[str]): List of device_ids where to examine loaded instances.
- Return:
- Dict[str, Dict[str, int]]: Dict[device_id, Dict[model_id, Count(instances)]].
- """
- result = {}
- for device_id in device_id_list:
- device_models = {}
- for model_id, device_map in self._request_pool_map.items():
- if device_id in device_map:
- pool_group = device_map[device_id]
- device_models[model_id] = pool_group.get_pool_count()
- result[device_id] = device_models
- return result
-
- def _worker_loop(self):
- while not self._stop_event.is_set():
- task = self._task_queue.get()
- if task is None:
- self._task_queue.task_done()
- break
- task_fn, args, kwargs = task
- try:
- task_fn(*args, **kwargs)
- except Exception as e:
- logger.error(f"Error executing task: {e}")
- finally:
- self._task_queue.task_done()
-
- def _load_model_task(self, model_id: str, device_id_list: list[str]):
- def _load_model_on_device_task(device_id: str):
- if not self.has_request_pools(model_id, device_id):
- actions = self._pool_scheduler.schedule_load_model_to_device(
- model_id, device_id
- )
- for action in actions:
- if action.action == ScaleActionType.SCALE_UP:
- self._expand_pools_on_device(
- action.model_id, device_id, action.amount
- )
- elif action.action == ScaleActionType.SCALE_DOWN:
- self._shrink_pools_on_device(
- action.model_id, device_id, action.amount
- )
- else:
- logger.info(
- f"[Inference][Device-{device_id}] Model {model_id} is already installed."
- )
-
- load_model_futures = self._executor.submit_batch(
- device_id_list, _load_model_on_device_task
- )
- concurrent.futures.wait(
- load_model_futures, return_when=concurrent.futures.ALL_COMPLETED
- )
-
- def _unload_model_task(self, model_id: str, device_id_list: list[str]):
- def _unload_model_on_device_task(device_id: str):
- if self.has_request_pools(model_id, device_id):
- actions = self._pool_scheduler.schedule_unload_model_from_device(
- model_id, device_id
- )
- for action in actions:
- if action.action == ScaleActionType.SCALE_DOWN:
- self._shrink_pools_on_device(
- action.model_id, device_id, action.amount
- )
- elif action.action == ScaleActionType.SCALE_UP:
- self._expand_pools_on_device(
- action.model_id, device_id, action.amount
- )
- else:
- logger.info(
- f"[Inference][Device-{device_id}] Model {model_id} is not installed."
- )
-
- unload_model_futures = self._executor.submit_batch(
- device_id_list, _unload_model_on_device_task
- )
- concurrent.futures.wait(
- unload_model_futures, return_when=concurrent.futures.ALL_COMPLETED
- )
-
- def _expand_pools_on_device(self, model_id: str, device_id: str, count: int):
- """
- Expand the pools for the given model_id and device_id sequentially.
- Args:
- model_id (str): The ID of the model.
- device_id (str): The ID of the device.
- count (int): The number of pools to be expanded.
- """
-
- def _expand_pool_on_device(*_):
- result_queue = mp.Queue()
- pool_id = self._new_pool_id.get_and_increment()
- if model_id == "sundial":
- config = SundialConfig()
- elif model_id == "timer_xl":
- config = TimerConfig()
- pool = InferenceRequestPool(
- pool_id=pool_id,
- model_id=model_id,
- device=device_id,
- config=config,
- request_queue=result_queue,
- result_queue=self._result_queue,
- ready_event=mp.Event(),
- )
- pool.start()
- self._register_pool(model_id, device_id, pool_id, pool, result_queue)
- if not pool.ready_event.wait(timeout=30):
- logger.error(
- f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool failed to be ready in time"
- )
- # TODO: retry or decrease the count? this error should be better handled
- self._erase_pool(model_id, device_id, pool_id)
- else:
- self.set_state(model_id, device_id, pool_id, PoolState.RUNNING)
- logger.info(
- f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool started running for model {model_id}"
- )
-
- expand_pool_futures = self._executor.submit_batch(
- [None for _ in range(count)], _expand_pool_on_device
- )
- concurrent.futures.wait(
- expand_pool_futures, return_when=concurrent.futures.ALL_COMPLETED
- )
-
- def _shrink_pools_on_device(self, model_id: str, device_id: str, count):
- """
- Shrink the pools for the given model_id by count sequentially.
- TODO: shrink pools in parallel
- """
-
- pool_ids = self.get_pool_ids(model_id, device_id)
- current_count = len(pool_ids)
- if current_count < count:
- logger.warning(
- f"[Inference] Cannot shrink {count} pools for model {model_id}, only {current_count} pools available."
- )
-
- # 1. Find pools with zero load and set them to STOP
- zero_load_pools = [
- pid for pid in pool_ids if self.get_load(model_id, device_id, pid) == 0
- ]
- stop_count = min(len(zero_load_pools), count)
- pools_to_stop = zero_load_pools[:stop_count]
-
- for pid in pools_to_stop:
- self.set_state(
- model_id, device_id, pid, PoolState.STOPPING
- ) # Set to STOPPING state
-
- remaining_to_close = count - len(pools_to_stop)
- if remaining_to_close <= 0:
- # close the STOP pools whose load is 0
- for pid in pools_to_stop:
- self.get_request_pool(model_id, device_id, pid).stop()
- self._erase_pool(model_id, device_id, pid)
- return
-
- # 2. If not enough, randomly select additional pools to set to STOP
- candidates = [pid for pid in pool_ids if pid not in pools_to_stop]
- additional_stop = random.sample(
- candidates, min(remaining_to_close, len(candidates))
- ) # TODO: implement better selection strategy
- for pid in additional_stop:
- self.set_state(
- model_id, device_id, pid, PoolState.STOPPING
- ) # Set to STOPPING state
-
- # 3. Wait for all STOPPING pools to have zero load, then close them
- final_stop_pools = pools_to_stop + additional_stop
- while final_stop_pools and not self._stop_event.is_set():
- for pid in final_stop_pools[:]: # iterate over a copy of the list
- if self.get_load(model_id, device_id, pid) == 0:
- self.get_request_pool(model_id, device_id, pid).stop()
- self._erase_pool(model_id, device_id, pid)
- final_stop_pools.remove(pid)
-
- def _register_pool(
- self,
- model_id: str,
- device_id: str,
- pool_id: int,
- request_pool: InferenceRequestPool,
- request_queue: mp.Queue,
- ):
- """
- Register a new inference request pool for the given model_id, device and pool_id.
- """
- self.set_request_pool_map(
- model_id, device_id, pool_id, request_pool, request_queue
- )
- pool_group: PoolGroup = self.get_request_pools_group(model_id, device_id)
- pool_group.set_state(pool_id, PoolState.INITIALIZING)
- logger.info(
- f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool initializing for model {model_id}"
- )
-
- def _erase_pool(self, model_id: str, device_id: str, pool_id: int):
- """
- Erase the specified inference request pool for the given model_id, device_id and pool_id.
- """
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- pool_group.remove_pool(pool_id)
- logger.info(
- f"[Inference][Device-{device_id}][Pool-{pool_id}] Erase pool for model {model_id}"
- )
- # Clean up empty structures
- if pool_group and not pool_group.get_pool_ids():
- self._request_pool_map[model_id].pop(device_id, None)
- if model_id in self._request_pool_map and not self._request_pool_map[model_id]:
- self._request_pool_map.pop(model_id, None)
-
- # =============== Request Management ===============
- def add_request(self, req: InferenceRequest, infer_proxy: InferenceRequestProxy):
- model_id = req.model_id
- if not self.has_request_pools(model_id):
- logger.error(f"[Inference] No pools found for model {model_id}.")
- infer_proxy.set_result(None)
- raise InferenceModelInternalError(
- "Dispatch request failed, because no inference pools are init."
- )
- # TODO: Implement adaptive scaling based on requests.(e.g. lazy initialization)
- # self.first_req_init(model_id)
-
- device_ids = self.get_device_ids(model_id)
- device_id = random.choice(
- device_ids
- ) # TODO: Implement better device selection strategy
- self._request_pool_map[model_id][device_id].dispatch_request(req, infer_proxy)
-
- # =============== Getters / Setters ===============
- def get_state(self, model_id, device_id, pool_id) -> Optional[PoolState]:
- """
- Get the state of the specified pool based on model_id, device_id, and pool_id.
- """
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- return pool_group.get_state(pool_id)
- return None
-
- def set_state(self, model_id, device_id, pool_id, state):
- """
- Set the state of the specified pool based on model_id, device_id, and pool_id.
- """
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- pool_group.set_state(pool_id, state)
-
- def get_device_ids(self, model_id) -> list[str]:
- """
- Get the list of device IDs for the given model_id, where the corresponding instances are loaded.
- """
- if model_id in self._request_pool_map:
- return list(self._request_pool_map[model_id].keys())
- return []
-
- def get_pool_ids(self, model_id: str, device_id: str) -> list[int]:
- """
- Get the list of pool IDs for the given model_id and device_id.
- """
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- return pool_group.get_pool_ids()
- return []
-
- def has_request_pools(self, model_id: str, device_id: Optional[str] = None) -> bool:
- """
- Check if there are request pools for the given model_id and device_id (optional).
- """
- if model_id not in self._request_pool_map:
- return False
- if device_id is not None:
- return device_id in self._request_pool_map[model_id]
- return True
-
- def get_request_pools_group(
- self, model_id: str, device_id: str
- ) -> Optional[PoolGroup]:
- if (
- model_id in self._request_pool_map
- and device_id in self._request_pool_map[model_id]
- ):
- return self._request_pool_map[model_id][device_id]
- else:
- return None
-
- def get_request_pool(
- self, model_id, device_id, pool_id
- ) -> Optional[InferenceRequestPool]:
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- return pool_group.get_request_pool(pool_id)
- return None
-
- def get_request_queue(self, model_id, device_id, pool_id) -> Optional[mp.Queue]:
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- return pool_group.get_request_queue(pool_id)
- return None
-
- def set_request_pool_map(
- self,
- model_id: str,
- device_id: str,
- pool_id: int,
- request_pool: InferenceRequestPool,
- request_queue: mp.Queue,
- ):
- """
- Set the request pool map for the given model_id, device and pool_id.
- """
- if model_id not in self._request_pool_map:
- self._request_pool_map[model_id] = {}
- if device_id not in self._request_pool_map[model_id]:
- self._request_pool_map[model_id][device_id] = PoolGroup(model_id)
- self._request_pool_map[model_id][device_id].add_pool(
- pool_id, request_pool, request_queue
- )
- logger.info(
- f"[Inference][Device-{device_id}][Pool-{pool_id}] Registered pool for model {model_id}"
- )
-
- def get_load(self, model_id, device_id, pool_id) -> int:
- """
- Get the current load of the specified pool.
- """
- pool_group = self.get_request_pools_group(model_id, device_id)
- if pool_group:
- return pool_group.get_load(pool_id)
- return -1
-
- def shutdown(self):
- self._stop_event.set()
-
- # shutdown pool controller
- self._task_queue.put(None)
- self._pool_control_worker_thread.join()
- self._executor.close()
-
- # shutdown pool instances
- # TODO: pool instances can be shutdown in parallel
- for inner in self._request_pool_map.values():
- for group in inner.values():
- group.shutdown()
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import concurrent
+1import queue
+1import random
+1import threading
+1from concurrent.futures import wait
+1from typing import Dict, Optional
+1
+1import torch
+1import torch.multiprocessing as mp
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.inference_request import (
+1 InferenceRequest,
+1 InferenceRequestProxy,
+1)
+1from iotdb.ainode.core.inference.inference_request_pool import (
+1 InferenceRequestPool,
+1 PoolState,
+1)
+1from iotdb.ainode.core.inference.pool_group import PoolGroup
+1from iotdb.ainode.core.inference.pool_scheduler.basic_pool_scheduler import (
+1 BasicPoolScheduler,
+1 ScaleActionType,
+1)
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
+1from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
+1from iotdb.ainode.core.util.atmoic_int import AtomicInt
+1from iotdb.ainode.core.util.batch_executor import BatchExecutor
+1from iotdb.ainode.core.util.decorator import synchronized
+1from iotdb.ainode.core.util.thread_name import ThreadName
+1
+1logger = Logger()
+1
+1
+1class PoolController:
+1 """
+1 A controller for handling inference request pools.
+1 """
+1
+1 def __init__(self, result_queue: mp.Queue):
+1 # structure: {model_id: {device_id: PoolGroup}}
+1 self._request_pool_map: Dict[str, Dict[str, PoolGroup]] = {}
+1 self._new_pool_id = AtomicInt()
+1 self._result_queue = result_queue
+1 self._pool_scheduler = BasicPoolScheduler(self._request_pool_map)
+1 self._stop_event = threading.Event()
+1
+1 # for pool instances control
+1 self._task_queue = queue.Queue()
+1 self._pool_control_worker_thread = threading.Thread(
+1 target=self._worker_loop, daemon=True
+1 )
+1 self._pool_control_worker_thread.start()
+1 self._executor = BatchExecutor(
+1 thread_name_prefix=ThreadName.INFERENCE_POOL_CONTROLLER.value
+1 )
+1
+1 # =============== Pool Management ===============
+1 @synchronized(threading.Lock())
+1 def first_req_init(self, model_id: str):
+1 """
+1 Initialize the pools when the first request for the given model_id arrives.
+1 """
+1 if not self.has_request_pools(model_id, device.index):
+1 # TODO: choose a device based on some strategy
+1 device = self.DEFAULT_DEVICE
+1 actions = self._pool_scheduler.schedule(model_id, device)
+1 for action in actions:
+1 if action.action == ScaleActionType.SCALE_UP:
+1 # initialize the first pool
+1 self._first_pool_init(action.model_id, str(device))
+1 # start a background thread to expand pools
+1 expand_thread = threading.Thread(
+1 target=self._expand_pools_on_device,
+1 args=(action.model_id, str(device), action.amount - 1),
+1 daemon=True,
+1 )
+1 expand_thread.start()
+1 elif action.action == ScaleActionType.SCALE_DOWN:
+1 # TODO: implement scale down logic
+1 pass
+1
+1 def _first_pool_init(self, model_id: str, device_str: str):
+1 """
+1 Initialize the first pool for the given model_id.
+1 Ensure the pool is ready before returning.
+1 """
+1 device = torch.device(device_str)
+1 device_id = device.index
+1
+1 if model_id == "sundial":
+1 config = SundialConfig()
+1 elif model_id == "timer_xl":
+1 config = TimerConfig()
+1 first_queue = mp.Queue()
+1 ready_event = mp.Event()
+1 first_pool = InferenceRequestPool(
+1 pool_id=0,
+1 model_id=model_id,
+1 device=device_str,
+1 config=config,
+1 request_queue=first_queue,
+1 result_queue=self._result_queue,
+1 ready_event=ready_event,
+1 )
+1 first_pool.start()
+1 self._register_pool(model_id, device_str, 0, first_pool, first_queue)
+1
+1 if not ready_event.wait(timeout=30):
+1 self._erase_pool(model_id, device_id, 0)
+1 logger.error(
+1 f"[Inference][Device-{device}][Pool-0] Pool failed to be ready in time"
+1 )
+1 else:
+1 self.set_state(model_id, device_id, 0, PoolState.RUNNING)
+1 logger.info(
+1 f"[Inference][Device-{device}][Pool-0] Pool started running for model {model_id}"
+1 )
+1
+1 def load_model(self, model_id: str, device_id_list: list[str]):
+1 """
+1 Load the model to the specified devices asynchronously.
+1 Args:
+1 model_id (str): The ID of the model to be loaded.
+1 device_id_list (list[str]): List of device_ids where the model should be loaded.
+1 """
+1 self._task_queue.put((self._load_model_task, (model_id, device_id_list), {}))
+1
+1 def unload_model(self, model_id: str, device_id_list: list[str]):
+1 """
+1 Unload the model from the specified devices asynchronously.
+1 Args:
+1 model_id (str): The ID of the model to be unloaded.
+1 device_id_list (list[str]): List of device_ids where the model should be unloaded.
+1 """
+1 self._task_queue.put((self._unload_model_task, (model_id, device_id_list), {}))
+1
+1 def show_loaded_models(
+1 self, device_id_list: list[str]
+1 ) -> Dict[str, Dict[str, int]]:
+1 """
+1 Show loaded model instances on the specified devices.
+1 Args:
+1 device_id_list (list[str]): List of device_ids where to examine loaded instances.
+1 Return:
+1 Dict[str, Dict[str, int]]: Dict[device_id, Dict[model_id, Count(instances)]].
+1 """
+1 result = {}
+1 for device_id in device_id_list:
+1 device_models = {}
+1 for model_id, device_map in self._request_pool_map.items():
+1 if device_id in device_map:
+1 pool_group = device_map[device_id]
+1 device_models[model_id] = pool_group.get_pool_count()
+1 result[device_id] = device_models
+1 return result
+1
+1 def _worker_loop(self):
+1 while not self._stop_event.is_set():
+1 task = self._task_queue.get()
+1 if task is None:
+1 self._task_queue.task_done()
+1 break
+1 task_fn, args, kwargs = task
+1 try:
+1 task_fn(*args, **kwargs)
+1 except Exception as e:
+1 logger.error(f"Error executing task: {e}")
+1 finally:
+1 self._task_queue.task_done()
+1
+1 def _load_model_task(self, model_id: str, device_id_list: list[str]):
+1 def _load_model_on_device_task(device_id: str):
+1 if not self.has_request_pools(model_id, device_id):
+1 actions = self._pool_scheduler.schedule_load_model_to_device(
+1 model_id, device_id
+1 )
+1 for action in actions:
+1 if action.action == ScaleActionType.SCALE_UP:
+1 self._expand_pools_on_device(
+1 action.model_id, device_id, action.amount
+1 )
+1 elif action.action == ScaleActionType.SCALE_DOWN:
+1 self._shrink_pools_on_device(
+1 action.model_id, device_id, action.amount
+1 )
+1 else:
+1 logger.info(
+1 f"[Inference][Device-{device_id}] Model {model_id} is already installed."
+1 )
+1
+1 load_model_futures = self._executor.submit_batch(
+1 device_id_list, _load_model_on_device_task
+1 )
+1 concurrent.futures.wait(
+1 load_model_futures, return_when=concurrent.futures.ALL_COMPLETED
+1 )
+1
+1 def _unload_model_task(self, model_id: str, device_id_list: list[str]):
+1 def _unload_model_on_device_task(device_id: str):
+1 if self.has_request_pools(model_id, device_id):
+1 actions = self._pool_scheduler.schedule_unload_model_from_device(
+1 model_id, device_id
+1 )
+1 for action in actions:
+1 if action.action == ScaleActionType.SCALE_DOWN:
+1 self._shrink_pools_on_device(
+1 action.model_id, device_id, action.amount
+1 )
+1 elif action.action == ScaleActionType.SCALE_UP:
+1 self._expand_pools_on_device(
+1 action.model_id, device_id, action.amount
+1 )
+1 else:
+1 logger.info(
+1 f"[Inference][Device-{device_id}] Model {model_id} is not installed."
+1 )
+1
+1 unload_model_futures = self._executor.submit_batch(
+1 device_id_list, _unload_model_on_device_task
+1 )
+1 concurrent.futures.wait(
+1 unload_model_futures, return_when=concurrent.futures.ALL_COMPLETED
+1 )
+1
+1 def _expand_pools_on_device(self, model_id: str, device_id: str, count: int):
+1 """
+1 Expand the pools for the given model_id and device_id sequentially.
+1 Args:
+1 model_id (str): The ID of the model.
+1 device_id (str): The ID of the device.
+1 count (int): The number of pools to be expanded.
+1 """
+1
+1 def _expand_pool_on_device(*_):
+1 result_queue = mp.Queue()
+1 pool_id = self._new_pool_id.get_and_increment()
+1 if model_id == "sundial":
+1 config = SundialConfig()
+1 elif model_id == "timer_xl":
+1 config = TimerConfig()
+1 pool = InferenceRequestPool(
+1 pool_id=pool_id,
+1 model_id=model_id,
+1 device=device_id,
+1 config=config,
+1 request_queue=result_queue,
+1 result_queue=self._result_queue,
+1 ready_event=mp.Event(),
+1 )
+1 pool.start()
+1 self._register_pool(model_id, device_id, pool_id, pool, result_queue)
+1 if not pool.ready_event.wait(timeout=30):
+1 logger.error(
+1 f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool failed to be ready in time"
+1 )
+1 # TODO: retry or decrease the count? this error should be better handled
+1 self._erase_pool(model_id, device_id, pool_id)
+1 else:
+1 self.set_state(model_id, device_id, pool_id, PoolState.RUNNING)
+1 logger.info(
+1 f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool started running for model {model_id}"
+1 )
+1
+1 expand_pool_futures = self._executor.submit_batch(
+1 [None for _ in range(count)], _expand_pool_on_device
+1 )
+1 concurrent.futures.wait(
+1 expand_pool_futures, return_when=concurrent.futures.ALL_COMPLETED
+1 )
+1
+1 def _shrink_pools_on_device(self, model_id: str, device_id: str, count):
+1 """
+1 Shrink the pools for the given model_id by count sequentially.
+1 TODO: shrink pools in parallel
+1 """
+1
+1 pool_ids = self.get_pool_ids(model_id, device_id)
+1 current_count = len(pool_ids)
+1 if current_count < count:
+1 logger.warning(
+1 f"[Inference] Cannot shrink {count} pools for model {model_id}, only {current_count} pools available."
+1 )
+1
+1 # 1. Find pools with zero load and set them to STOP
+1 zero_load_pools = [
+1 pid for pid in pool_ids if self.get_load(model_id, device_id, pid) == 0
+1 ]
+1 stop_count = min(len(zero_load_pools), count)
+1 pools_to_stop = zero_load_pools[:stop_count]
+1
+1 for pid in pools_to_stop:
+1 self.set_state(
+1 model_id, device_id, pid, PoolState.STOPPING
+1 ) # Set to STOPPING state
+1
+1 remaining_to_close = count - len(pools_to_stop)
+1 if remaining_to_close <= 0:
+1 # close the STOP pools whose load is 0
+1 for pid in pools_to_stop:
+1 self.get_request_pool(model_id, device_id, pid).stop()
+1 self._erase_pool(model_id, device_id, pid)
+1 return
+1
+1 # 2. If not enough, randomly select additional pools to set to STOP
+1 candidates = [pid for pid in pool_ids if pid not in pools_to_stop]
+1 additional_stop = random.sample(
+1 candidates, min(remaining_to_close, len(candidates))
+1 ) # TODO: implement better selection strategy
+1 for pid in additional_stop:
+1 self.set_state(
+1 model_id, device_id, pid, PoolState.STOPPING
+1 ) # Set to STOPPING state
+1
+1 # 3. Wait for all STOPPING pools to have zero load, then close them
+1 final_stop_pools = pools_to_stop + additional_stop
+1 while final_stop_pools and not self._stop_event.is_set():
+1 for pid in final_stop_pools[:]: # iterate over a copy of the list
+1 if self.get_load(model_id, device_id, pid) == 0:
+1 self.get_request_pool(model_id, device_id, pid).stop()
+1 self._erase_pool(model_id, device_id, pid)
+1 final_stop_pools.remove(pid)
+1
+1 def _register_pool(
+1 self,
+1 model_id: str,
+1 device_id: str,
+1 pool_id: int,
+1 request_pool: InferenceRequestPool,
+1 request_queue: mp.Queue,
+1 ):
+1 """
+1 Register a new inference request pool for the given model_id, device and pool_id.
+1 """
+1 self.set_request_pool_map(
+1 model_id, device_id, pool_id, request_pool, request_queue
+1 )
+1 pool_group: PoolGroup = self.get_request_pools_group(model_id, device_id)
+1 pool_group.set_state(pool_id, PoolState.INITIALIZING)
+1 logger.info(
+1 f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool initializing for model {model_id}"
+1 )
+1
+1 def _erase_pool(self, model_id: str, device_id: str, pool_id: int):
+1 """
+1 Erase the specified inference request pool for the given model_id, device_id and pool_id.
+1 """
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 pool_group.remove_pool(pool_id)
+1 logger.info(
+1 f"[Inference][Device-{device_id}][Pool-{pool_id}] Erase pool for model {model_id}"
+1 )
+1 # Clean up empty structures
+1 if pool_group and not pool_group.get_pool_ids():
+1 self._request_pool_map[model_id].pop(device_id, None)
+1 if model_id in self._request_pool_map and not self._request_pool_map[model_id]:
+1 self._request_pool_map.pop(model_id, None)
+1
+1 # =============== Request Management ===============
+1 def add_request(self, req: InferenceRequest, infer_proxy: InferenceRequestProxy):
+1 model_id = req.model_id
+1 if not self.has_request_pools(model_id):
+1 logger.error(f"[Inference] No pools found for model {model_id}.")
+1 infer_proxy.set_result(None)
+1 raise InferenceModelInternalError(
+1 "Dispatch request failed, because no inference pools are init."
+1 )
+1 # TODO: Implement adaptive scaling based on requests.(e.g. lazy initialization)
+1 # self.first_req_init(model_id)
+1
+1 device_ids = self.get_device_ids(model_id)
+1 device_id = random.choice(
+1 device_ids
+1 ) # TODO: Implement better device selection strategy
+1 self._request_pool_map[model_id][device_id].dispatch_request(req, infer_proxy)
+1
+1 # =============== Getters / Setters ===============
+1 def get_state(self, model_id, device_id, pool_id) -> Optional[PoolState]:
+1 """
+1 Get the state of the specified pool based on model_id, device_id, and pool_id.
+1 """
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 return pool_group.get_state(pool_id)
+1 return None
+1
+1 def set_state(self, model_id, device_id, pool_id, state):
+1 """
+1 Set the state of the specified pool based on model_id, device_id, and pool_id.
+1 """
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 pool_group.set_state(pool_id, state)
+1
+1 def get_device_ids(self, model_id) -> list[str]:
+1 """
+1 Get the list of device IDs for the given model_id, where the corresponding instances are loaded.
+1 """
+1 if model_id in self._request_pool_map:
+1 return list(self._request_pool_map[model_id].keys())
+1 return []
+1
+1 def get_pool_ids(self, model_id: str, device_id: str) -> list[int]:
+1 """
+1 Get the list of pool IDs for the given model_id and device_id.
+1 """
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 return pool_group.get_pool_ids()
+1 return []
+1
+1 def has_request_pools(self, model_id: str, device_id: Optional[str] = None) -> bool:
+1 """
+1 Check if there are request pools for the given model_id and device_id (optional).
+1 """
+1 if model_id not in self._request_pool_map:
+1 return False
+1 if device_id is not None:
+1 return device_id in self._request_pool_map[model_id]
+1 return True
+1
+1 def get_request_pools_group(
+1 self, model_id: str, device_id: str
+1 ) -> Optional[PoolGroup]:
+1 if (
+1 model_id in self._request_pool_map
+1 and device_id in self._request_pool_map[model_id]
+1 ):
+1 return self._request_pool_map[model_id][device_id]
+1 else:
+1 return None
+1
+1 def get_request_pool(
+1 self, model_id, device_id, pool_id
+1 ) -> Optional[InferenceRequestPool]:
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 return pool_group.get_request_pool(pool_id)
+1 return None
+1
+1 def get_request_queue(self, model_id, device_id, pool_id) -> Optional[mp.Queue]:
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 return pool_group.get_request_queue(pool_id)
+1 return None
+1
+1 def set_request_pool_map(
+1 self,
+1 model_id: str,
+1 device_id: str,
+1 pool_id: int,
+1 request_pool: InferenceRequestPool,
+1 request_queue: mp.Queue,
+1 ):
+1 """
+1 Set the request pool map for the given model_id, device and pool_id.
+1 """
+1 if model_id not in self._request_pool_map:
+1 self._request_pool_map[model_id] = {}
+1 if device_id not in self._request_pool_map[model_id]:
+1 self._request_pool_map[model_id][device_id] = PoolGroup(model_id)
+1 self._request_pool_map[model_id][device_id].add_pool(
+1 pool_id, request_pool, request_queue
+1 )
+1 logger.info(
+1 f"[Inference][Device-{device_id}][Pool-{pool_id}] Registered pool for model {model_id}"
+1 )
+1
+1 def get_load(self, model_id, device_id, pool_id) -> int:
+1 """
+1 Get the current load of the specified pool.
+1 """
+1 pool_group = self.get_request_pools_group(model_id, device_id)
+1 if pool_group:
+1 return pool_group.get_load(pool_id)
+1 return -1
+1
+1 def shutdown(self):
+1 self._stop_event.set()
+1
+1 # shutdown pool controller
+1 self._task_queue.put(None)
+1 self._pool_control_worker_thread.join()
+1 self._executor.close()
+1
+1 # shutdown pool instances
+1 # TODO: pool instances can be shutdown in parallel
+1 for inner in self._request_pool_map.values():
+1 for group in inner.values():
+1 group.shutdown()
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
index 96dce845585..0fa5ca1c7f9 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
@@ -1,115 +1,116 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-from typing import Dict, Tuple
-
-import torch.multiprocessing as mp
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.dispatcher.basic_dispatcher import BasicDispatcher
-from iotdb.ainode.core.inference.inference_request import (
- InferenceRequest,
- InferenceRequestProxy,
-)
-from iotdb.ainode.core.inference.inference_request_pool import (
- InferenceRequestPool,
- PoolState,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.util.atmoic_int import AtomicInt
-
-logger = Logger()
-
-
-class PoolGroup:
- """
- A group of inference request pools for a specific model.
- """
-
- def __init__(self, model_id):
- # structure: {pool_id: (InferenceRequestPool, waitingQueue)}
- self.pool_group: Dict[int, Tuple[InferenceRequestPool, mp.Queue]] = {}
- # structure: {pool_id: PoolState}
- self.pool_states: Dict[int, PoolState] = {}
- # structure: {pool_id: remaining_reqs_cnt}
- self.pool_remaining_reqs: Dict[int, AtomicInt] = {}
- self.model_id = model_id
- self.request_dispatcher = BasicDispatcher(self.pool_states)
-
- def get_pool_group(self) -> Dict[int, Tuple[InferenceRequestPool, mp.Queue]]:
- return self.pool_group
-
- def add_pool(
- self, pool_id: int, request_pool: InferenceRequestPool, request_queue: mp.Queue
- ):
- self.pool_group[pool_id] = (request_pool, request_queue)
- self.pool_remaining_reqs[pool_id] = AtomicInt()
-
- def remove_pool(self, pool_id: int):
- self.pool_group.pop(pool_id, None)
- self.pool_states.pop(pool_id, None)
- self.pool_remaining_reqs.pop(pool_id, None)
-
- def get_pool_ids(self) -> list[int]:
- return list(self.pool_group.keys())
-
- def get_pool_count(self) -> int:
- return len(self.pool_group)
-
- def dispatch_request(
- self, req: InferenceRequest, infer_proxy: InferenceRequestProxy
- ):
- pool_id = self.request_dispatcher.dispatch_request(req, self.get_pool_ids())
- req_q = self.pool_group[pool_id][1]
- self.pool_remaining_reqs[pool_id].increment_and_get()
- infer_proxy.set_counter(self.pool_remaining_reqs[pool_id])
- req_q.put(req)
- logger.debug(
- f"[Inference][Pool-{pool_id}][Req-{req.req_id}] Request is queued for inference"
- )
-
- def get_request_pool(self, pool_id) -> InferenceRequestPool:
- if pool_id not in self.pool_group:
- raise InferenceModelInternalError(
- f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
- )
- return self.pool_group[pool_id][0]
-
- def get_request_queue(self, pool_id) -> mp.Queue:
- if pool_id not in self.pool_group:
- raise InferenceModelInternalError(
- f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
- )
- return self.pool_group[pool_id][1]
-
- def get_state(self, pool_id) -> PoolState:
- return self.pool_states[pool_id]
-
- def set_state(self, pool_id, state):
- self.pool_states[pool_id] = state
-
- def get_load(self, pool_id) -> int:
- """
- Currently, we use the number of remaining requests in the pool as the load estimation.
- TODO: we will implement better load estimation in the future.
- """
- return self.pool_remaining_reqs[pool_id].get()
-
- def shutdown(self):
- for pool_id, (request_pool, _) in self.pool_group.items():
- logger.info(f"[Inference][Pool-{pool_id}] Shutting down the pool...")
- request_pool.stop()
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1from typing import Dict, Tuple
+1
+1import torch.multiprocessing as mp
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.dispatcher.basic_dispatcher import BasicDispatcher
+1from iotdb.ainode.core.inference.inference_request import (
+1 InferenceRequest,
+1 InferenceRequestProxy,
+1)
+1from iotdb.ainode.core.inference.inference_request_pool import (
+1 InferenceRequestPool,
+1 PoolState,
+1)
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.util.atmoic_int import AtomicInt
+1
+1logger = Logger()
+1
+1
+1class PoolGroup:
+1 """
+1 A group of inference request pools for a specific model.
+1 """
+1
+1 def __init__(self, model_id):
+1 # structure: {pool_id: (InferenceRequestPool, waitingQueue)}
+1 self.pool_group: Dict[int, Tuple[InferenceRequestPool, mp.Queue]] = {}
+1 # structure: {pool_id: PoolState}
+1 self.pool_states: Dict[int, PoolState] = {}
+1 # structure: {pool_id: remaining_reqs_cnt}
+1 self.pool_remaining_reqs: Dict[int, AtomicInt] = {}
+1 self.model_id = model_id
+1 self.request_dispatcher = BasicDispatcher(self.pool_states)
+1
+1 def get_pool_group(self) -> Dict[int, Tuple[InferenceRequestPool, mp.Queue]]:
+1 return self.pool_group
+1
+1 def add_pool(
+1 self, pool_id: int, request_pool: InferenceRequestPool, request_queue: mp.Queue
+1 ):
+1 self.pool_group[pool_id] = (request_pool, request_queue)
+1 self.pool_remaining_reqs[pool_id] = AtomicInt()
+1
+1 def remove_pool(self, pool_id: int):
+1 self.pool_group.pop(pool_id, None)
+1 self.pool_states.pop(pool_id, None)
+1 self.pool_remaining_reqs.pop(pool_id, None)
+1
+1 def get_pool_ids(self) -> list[int]:
+1 return list(self.pool_group.keys())
+1
+1 def get_pool_count(self) -> int:
+1 return len(self.pool_group)
+1
+1 def dispatch_request(
+1 self, req: InferenceRequest, infer_proxy: InferenceRequestProxy
+1 ):
+1 pool_id = self.request_dispatcher.dispatch_request(req, self.get_pool_ids())
+1 req_q = self.pool_group[pool_id][1]
+1 self.pool_remaining_reqs[pool_id].increment_and_get()
+1 infer_proxy.set_counter(self.pool_remaining_reqs[pool_id])
+1 req_q.put(req)
+1 logger.debug(
+1 f"[Inference][Pool-{pool_id}][Req-{req.req_id}] Request is queued for inference"
+1 )
+1
+1 def get_request_pool(self, pool_id) -> InferenceRequestPool:
+1 if pool_id not in self.pool_group:
+1 raise InferenceModelInternalError(
+1 f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
+1 )
+1 return self.pool_group[pool_id][0]
+1
+1 def get_request_queue(self, pool_id) -> mp.Queue:
+1 if pool_id not in self.pool_group:
+1 raise InferenceModelInternalError(
+1 f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
+1 )
+1 return self.pool_group[pool_id][1]
+1
+1 def get_state(self, pool_id) -> PoolState:
+1 return self.pool_states[pool_id]
+1
+1 def set_state(self, pool_id, state):
+1 self.pool_states[pool_id] = state
+1
+1 def get_load(self, pool_id) -> int:
+1 """
+1 Currently, we use the number of remaining requests in the pool as the load estimation.
+1 TODO: we will implement better load estimation in the future.
+1 """
+1 return self.pool_remaining_reqs[pool_id].get()
+1
+1 def shutdown(self):
+1 for pool_id, (request_pool, _) in self.pool_group.items():
+1 logger.info(f"[Inference][Pool-{pool_id}] Shutting down the pool...")
+1 request_pool.stop()
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/abstract_pool_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/abstract_pool_scheduler.py
index 6a26d1fe15b..941db0bfdbb 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/abstract_pool_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/abstract_pool_scheduler.py
@@ -1,85 +1,86 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from enum import Enum
-from typing import Dict, List
-
-from iotdb.ainode.core.inference.pool_group import PoolGroup
-
-
-class ScaleActionType(Enum):
- SCALE_UP = "Scale Up"
- SCALE_DOWN = "Scale Down"
-
-
-@dataclass(frozen=True)
-class ScaleAction:
- action: ScaleActionType
- amount: int
- model_id: str
- device_id: str
-
-
-class AbstractPoolScheduler(ABC):
- """
- Abstract base class for pool scheduling strategies.
- """
-
- def __init__(self, request_pool_map: Dict[str, Dict[str, PoolGroup]]):
- """
- Args:
- request_pool_map: Dict["model_id", Dict["device_id", PoolGroup]].
- """
- self._request_pool_map = request_pool_map
-
- @abstractmethod
- def schedule(self, model_id: str) -> List[ScaleAction]:
- """
- Schedule a scaling action for the given model_id.
- """
- pass
-
- @abstractmethod
- def schedule_load_model_to_device(
- self, model_id: str, device_id: str
- ) -> List[ScaleAction]:
- """
- Schedule a series of actions to load the model to the device.
- Args:
- model_id: The model to be loaded.
- device_id: The device to load the model to.
- Returns:
- A list of ScaleAction to be performed.
- """
- pass
-
- @abstractmethod
- def schedule_unload_model_from_device(
- self, model_id: str, device_id: str
- ) -> List[ScaleAction]:
- """
- Schedule a series of actions to unload the model from the device.
- Args:
- model_id: The model to be unloaded.
- device_id: The device to unload the model from.
- Returns:
- A list of ScaleAction to be performed.
- """
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from abc import ABC, abstractmethod
+1from dataclasses import dataclass
+1from enum import Enum
+1from typing import Dict, List
+1
+1from iotdb.ainode.core.inference.pool_group import PoolGroup
+1
+1
+1class ScaleActionType(Enum):
+1 SCALE_UP = "Scale Up"
+1 SCALE_DOWN = "Scale Down"
+1
+1
+1@dataclass(frozen=True)
+1class ScaleAction:
+1 action: ScaleActionType
+1 amount: int
+1 model_id: str
+1 device_id: str
+1
+1
+1class AbstractPoolScheduler(ABC):
+1 """
+1 Abstract base class for pool scheduling strategies.
+1 """
+1
+1 def __init__(self, request_pool_map: Dict[str, Dict[str, PoolGroup]]):
+1 """
+1 Args:
+1 request_pool_map: Dict["model_id", Dict["device_id", PoolGroup]].
+1 """
+1 self._request_pool_map = request_pool_map
+1
+1 @abstractmethod
+1 def schedule(self, model_id: str) -> List[ScaleAction]:
+1 """
+1 Schedule a scaling action for the given model_id.
+1 """
+1 pass
+1
+1 @abstractmethod
+1 def schedule_load_model_to_device(
+1 self, model_id: str, device_id: str
+1 ) -> List[ScaleAction]:
+1 """
+1 Schedule a series of actions to load the model to the device.
+1 Args:
+1 model_id: The model to be loaded.
+1 device_id: The device to load the model to.
+1 Returns:
+1 A list of ScaleAction to be performed.
+1 """
+1 pass
+1
+1 @abstractmethod
+1 def schedule_unload_model_from_device(
+1 self, model_id: str, device_id: str
+1 ) -> List[ScaleAction]:
+1 """
+1 Schedule a series of actions to unload the model from the device.
+1 Args:
+1 model_id: The model to be unloaded.
+1 device_id: The device to unload the model from.
+1 Returns:
+1 A list of ScaleAction to be performed.
+1 """
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
index 9aefd236730..f41e5f87fdf 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
@@ -1,191 +1,192 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from typing import Dict, List, Optional
-
-import torch
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.pool_group import PoolGroup
-from iotdb.ainode.core.inference.pool_scheduler.abstract_pool_scheduler import (
- AbstractPoolScheduler,
- ScaleAction,
- ScaleActionType,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.utils import (
- INFERENCE_EXTRA_MEMORY_RATIO,
- INFERENCE_MEMORY_USAGE_RATIO,
- MODEL_MEM_USAGE_MAP,
- estimate_pool_size,
- evaluate_system_resources,
-)
-from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP
-from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
-
-logger = Logger()
-
-
-def _estimate_shared_pool_size_by_total_mem(
- device: torch.device,
- existing_model_ids: List[str],
- new_model_id: Optional[str] = None,
-) -> Dict[str, int]:
- """
- Estimate pool counts for (existing_model_ids + new_model_id) by equally
- splitting the device's TOTAL memory among models.
-
- Returns:
- mapping {model_id: pool_num}
- """
- # Extract unique model IDs
- all_models = existing_model_ids + (
- [new_model_id] if new_model_id is not None else []
- )
-
- # Seize memory usage for each model
- mem_usages: Dict[str, float] = {}
- for model_id in all_models:
- model_info = BUILT_IN_LTSM_MAP.get(model_id)
- model_type = model_info.model_type
- mem_usages[model_id] = (
- MODEL_MEM_USAGE_MAP[model_type] * INFERENCE_EXTRA_MEMORY_RATIO
- )
-
- # Evaluate system resources and get TOTAL memory
- system_res = evaluate_system_resources(device)
- # TODO: Its better to consider free memory, but we need to track the memory usage of existing pools
- total_mem = system_res.get("total_mem")
-
- usable_mem = total_mem * INFERENCE_MEMORY_USAGE_RATIO
- if usable_mem <= 0:
- logger.error(
- f"[Inference][Device-{device}] No usable memory on device. total={total_mem / 1024 ** 2:.2f} MB, usable={usable_mem / 1024 ** 2:.2f} MB"
- )
-
- # Each model gets an equal share of the TOTAL memory
- num_models = len(all_models)
- per_model_share = usable_mem / num_models # TODO: Implement more strategies later
-
- # Calculate pool allocation for each model
- allocation: Dict[str, int] = {}
- for model_id in all_models:
- pool_num = int(per_model_share // mem_usages[model_id])
- if pool_num <= 0:
- logger.warning(
- f"[Inference][Device-{device}] Not enough TOTAL memory to guarantee at least 1 pool for model {model_id}, no pool will be scheduled for this model. "
- f"Per-model share={per_model_share / 1024 ** 2:.2f} MB, need>={mem_usages[model_id] / 1024 ** 2:.2f} MB"
- )
- allocation[model_id] = pool_num
- logger.info(
- f"[Inference][Device-{device}] Shared pool allocation (by TOTAL memory): {allocation}"
- )
- return allocation
-
-
-class BasicPoolScheduler(AbstractPoolScheduler):
- """
- A basic scheduler to init the request pools. In short, different kind of models will equally share the available resource of the located device, and scale down actions are always ahead of scale up.
- """
-
- def __init__(self, request_pool_map: Dict[str, Dict[str, PoolGroup]]):
- super().__init__(request_pool_map)
-
- def schedule(self, model_id: str) -> List[ScaleAction]:
- """
- Schedule a scaling action for the given model_id.
- """
- 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}."
- )
- return [ScaleAction(ScaleActionType.SCALE_UP, pool_num, model_id)]
-
- def schedule_load_model_to_device(
- self, model_id: str, device_id: str
- ) -> List[ScaleAction]:
- existing_model_ids = [
- existing_model_id
- for existing_model_id, pool_group_map in self._request_pool_map.items()
- if existing_model_id != model_id and device_id in pool_group_map
- ]
- allocation_result = _estimate_shared_pool_size_by_total_mem(
- device=convert_device_id_to_torch_device(device_id),
- existing_model_ids=existing_model_ids,
- new_model_id=model_id,
- )
- return self._convert_allocation_result_to_scale_actions(
- allocation_result, device_id
- )
-
- def schedule_unload_model_from_device(
- self, model_id: str, device_id: str
- ) -> List[ScaleAction]:
- existing_model_ids = [
- existing_model_id
- for existing_model_id, pool_group_map in self._request_pool_map.items()
- if existing_model_id != model_id and device_id in pool_group_map
- ]
- allocation_result = (
- _estimate_shared_pool_size_by_total_mem(
- device=convert_device_id_to_torch_device(device_id),
- existing_model_ids=existing_model_ids,
- new_model_id=None,
- )
- if len(existing_model_ids) > 0
- else {model_id: 0}
- )
- return self._convert_allocation_result_to_scale_actions(
- allocation_result, device_id
- )
-
- def _convert_allocation_result_to_scale_actions(
- self, allocation_result: Dict[str, int], device_id: str
- ) -> List[ScaleAction]:
- """
- Convert the model allocation result to List[ScaleAction], where the scale down actions are always ahead of the scale up.
- """
- actions = []
- for model_id, target_num in allocation_result.items():
- current_num = self._request_pool_map.get(model_id, {}).get(device_id, None)
- current_num = current_num.get_pool_count() if current_num else 0
- diff = target_num - current_num
- if diff > 0:
- actions.append(
- ScaleAction(
- action=ScaleActionType.SCALE_UP,
- amount=diff,
- model_id=model_id,
- device_id=device_id,
- )
- )
- elif diff < 0:
- actions.append(
- ScaleAction(
- action=ScaleActionType.SCALE_DOWN,
- amount=-diff,
- model_id=model_id,
- device_id=device_id,
- )
- )
- sorted_actions = sorted(
- actions, key=lambda a: (0 if a.action == ScaleActionType.SCALE_DOWN else 1)
- )
- return sorted_actions
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from typing import Dict, List, Optional
+1
+1import torch
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.pool_group import PoolGroup
+1from iotdb.ainode.core.inference.pool_scheduler.abstract_pool_scheduler import (
+1 AbstractPoolScheduler,
+1 ScaleAction,
+1 ScaleActionType,
+1)
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.manager.utils import (
+1 INFERENCE_EXTRA_MEMORY_RATIO,
+1 INFERENCE_MEMORY_USAGE_RATIO,
+1 MODEL_MEM_USAGE_MAP,
+1 estimate_pool_size,
+1 evaluate_system_resources,
+1)
+1from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP
+1from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
+1
+1logger = Logger()
+1
+1
+1def _estimate_shared_pool_size_by_total_mem(
+1 device: torch.device,
+1 existing_model_ids: List[str],
+1 new_model_id: Optional[str] = None,
+1) -> Dict[str, int]:
+1 """
+1 Estimate pool counts for (existing_model_ids + new_model_id) by equally
+1 splitting the device's TOTAL memory among models.
+1
+1 Returns:
+1 mapping {model_id: pool_num}
+1 """
+1 # Extract unique model IDs
+1 all_models = existing_model_ids + (
+1 [new_model_id] if new_model_id is not None else []
+1 )
+1
+1 # Seize memory usage for each model
+1 mem_usages: Dict[str, float] = {}
+1 for model_id in all_models:
+1 model_info = BUILT_IN_LTSM_MAP.get(model_id)
+1 model_type = model_info.model_type
+1 mem_usages[model_id] = (
+1 MODEL_MEM_USAGE_MAP[model_type] * INFERENCE_EXTRA_MEMORY_RATIO
+1 )
+1
+1 # Evaluate system resources and get TOTAL memory
+1 system_res = evaluate_system_resources(device)
+1 # TODO: Its better to consider free memory, but we need to track the memory usage of existing pools
+1 total_mem = system_res.get("total_mem")
+1
+1 usable_mem = total_mem * INFERENCE_MEMORY_USAGE_RATIO
+1 if usable_mem <= 0:
+1 logger.error(
+1 f"[Inference][Device-{device}] No usable memory on device. total={total_mem / 1024 ** 2:.2f} MB, usable={usable_mem / 1024 ** 2:.2f} MB"
+1 )
+1
+1 # Each model gets an equal share of the TOTAL memory
+1 num_models = len(all_models)
+1 per_model_share = usable_mem / num_models # TODO: Implement more strategies later
+1
+1 # Calculate pool allocation for each model
+1 allocation: Dict[str, int] = {}
+1 for model_id in all_models:
+1 pool_num = int(per_model_share // mem_usages[model_id])
+1 if pool_num <= 0:
+1 logger.warning(
+1 f"[Inference][Device-{device}] Not enough TOTAL memory to guarantee at least 1 pool for model {model_id}, no pool will be scheduled for this model. "
+1 f"Per-model share={per_model_share / 1024 ** 2:.2f} MB, need>={mem_usages[model_id] / 1024 ** 2:.2f} MB"
+1 )
+1 allocation[model_id] = pool_num
+1 logger.info(
+1 f"[Inference][Device-{device}] Shared pool allocation (by TOTAL memory): {allocation}"
+1 )
+1 return allocation
+1
+1
+1class BasicPoolScheduler(AbstractPoolScheduler):
+1 """
+1 A basic scheduler to init the request pools. In short, different kind of models will equally share the available resource of the located device, and scale down actions are always ahead of scale up.
+1 """
+1
+1 def __init__(self, request_pool_map: Dict[str, Dict[str, PoolGroup]]):
+1 super().__init__(request_pool_map)
+1
+1 def schedule(self, model_id: str) -> List[ScaleAction]:
+1 """
+1 Schedule a scaling action for the given model_id.
+1 """
+1 if model_id not in self._request_pool_map:
+1 pool_num = estimate_pool_size(self.DEFAULT_DEVICE, model_id)
+1 if pool_num <= 0:
+1 raise InferenceModelInternalError(
+1 f"Not enough memory to run model {model_id}."
+1 )
+1 return [ScaleAction(ScaleActionType.SCALE_UP, pool_num, model_id)]
+1
+1 def schedule_load_model_to_device(
+1 self, model_id: str, device_id: str
+1 ) -> List[ScaleAction]:
+1 existing_model_ids = [
+1 existing_model_id
+1 for existing_model_id, pool_group_map in self._request_pool_map.items()
+1 if existing_model_id != model_id and device_id in pool_group_map
+1 ]
+1 allocation_result = _estimate_shared_pool_size_by_total_mem(
+1 device=convert_device_id_to_torch_device(device_id),
+1 existing_model_ids=existing_model_ids,
+1 new_model_id=model_id,
+1 )
+1 return self._convert_allocation_result_to_scale_actions(
+1 allocation_result, device_id
+1 )
+1
+1 def schedule_unload_model_from_device(
+1 self, model_id: str, device_id: str
+1 ) -> List[ScaleAction]:
+1 existing_model_ids = [
+1 existing_model_id
+1 for existing_model_id, pool_group_map in self._request_pool_map.items()
+1 if existing_model_id != model_id and device_id in pool_group_map
+1 ]
+1 allocation_result = (
+1 _estimate_shared_pool_size_by_total_mem(
+1 device=convert_device_id_to_torch_device(device_id),
+1 existing_model_ids=existing_model_ids,
+1 new_model_id=None,
+1 )
+1 if len(existing_model_ids) > 0
+1 else {model_id: 0}
+1 )
+1 return self._convert_allocation_result_to_scale_actions(
+1 allocation_result, device_id
+1 )
+1
+1 def _convert_allocation_result_to_scale_actions(
+1 self, allocation_result: Dict[str, int], device_id: str
+1 ) -> List[ScaleAction]:
+1 """
+1 Convert the model allocation result to List[ScaleAction], where the scale down actions are always ahead of the scale up.
+1 """
+1 actions = []
+1 for model_id, target_num in allocation_result.items():
+1 current_num = self._request_pool_map.get(model_id, {}).get(device_id, None)
+1 current_num = current_num.get_pool_count() if current_num else 0
+1 diff = target_num - current_num
+1 if diff > 0:
+1 actions.append(
+1 ScaleAction(
+1 action=ScaleActionType.SCALE_UP,
+1 amount=diff,
+1 model_id=model_id,
+1 device_id=device_id,
+1 )
+1 )
+1 elif diff < 0:
+1 actions.append(
+1 ScaleAction(
+1 action=ScaleActionType.SCALE_DOWN,
+1 amount=-diff,
+1 model_id=model_id,
+1 device_id=device_id,
+1 )
+1 )
+1 sorted_actions = sorted(
+1 actions, key=lambda a: (0 if a.action == ScaleActionType.SCALE_DOWN else 1)
+1 )
+1 return sorted_actions
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py
index a6c2fe53cae..5688f0effc1 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py
@@ -1,62 +1,63 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-
-
-class AbstractRequestScheduler(ABC):
- """
- Abstract base class for inference scheduling strategies.
-
- This class defines the high-level interface for scheduling inference requests.
- A scheduler is responsible for managing the execution order of inference tasks across different
- stages: waiting, running, and finished.
-
- Subclasses should implement specific scheduling logic.
- """
-
- def __init__(self, waiting_queue, running_queue, finished_queue):
- """
- Args:
- waiting_queue: Queue containing inference requests that are waiting to be executed.
- running_queue: Queue containing currently running inference tasks.
- finished_queue: Queue containing completed inference tasks.
- """
- self.waiting_queue = waiting_queue
- self.running_queue = running_queue
- self.finished_queue = finished_queue
-
- @abstractmethod
- def schedule_activate(self) -> list:
- """
- Select one or more inference requests from the waiting queue that are ready to be activated and processed.
-
- Returns:
- List: A list of inference requests that will be moved to the running queue.
- """
- pass
-
- @abstractmethod
- def schedule_step(self) -> list:
- """
- Select one or more inference requests from the running queue that are ready to perform the next inference step.
-
- Returns:
- List: A list of inference requests that are scheduled to run an inference step.
- """
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from abc import ABC, abstractmethod
+1
+1
+1class AbstractRequestScheduler(ABC):
+1 """
+1 Abstract base class for inference scheduling strategies.
+1
+1 This class defines the high-level interface for scheduling inference requests.
+1 A scheduler is responsible for managing the execution order of inference tasks across different
+1 stages: waiting, running, and finished.
+1
+1 Subclasses should implement specific scheduling logic.
+1 """
+1
+1 def __init__(self, waiting_queue, running_queue, finished_queue):
+1 """
+1 Args:
+1 waiting_queue: Queue containing inference requests that are waiting to be executed.
+1 running_queue: Queue containing currently running inference tasks.
+1 finished_queue: Queue containing completed inference tasks.
+1 """
+1 self.waiting_queue = waiting_queue
+1 self.running_queue = running_queue
+1 self.finished_queue = finished_queue
+1
+1 @abstractmethod
+1 def schedule_activate(self) -> list:
+1 """
+1 Select one or more inference requests from the waiting queue that are ready to be activated and processed.
+1
+1 Returns:
+1 List: A list of inference requests that will be moved to the running queue.
+1 """
+1 pass
+1
+1 @abstractmethod
+1 def schedule_step(self) -> list:
+1 """
+1 Select one or more inference requests from the running queue that are ready to perform the next inference step.
+1
+1 Returns:
+1 List: A list of inference requests that are scheduled to run an inference step.
+1 """
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py
index ef5d37a18a1..60f75b7bd97 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py
@@ -1,88 +1,89 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import os
-
-import psutil
-import torch
-
-from iotdb.ainode.core.inference.request_scheduler.abstract_request_scheduler import (
- AbstractRequestScheduler,
-)
-from iotdb.ainode.core.log import Logger
-
-logger = Logger()
-
-
-class BasicRequestScheduler(AbstractRequestScheduler):
- """
- A simple FIFO request scheduler that selects requests based on memory availability and activation/step size.
- """
-
- def __init__(
- self,
- waiting_queue,
- running_queue,
- finished_queue,
- pool_id,
- max_memory_bytes=1 << 30,
- max_activate_size=10,
- max_step_size=10,
- ):
- super().__init__(waiting_queue, running_queue, finished_queue)
- self.max_memory_bytes = max_memory_bytes
- self.max_activate_size = max_activate_size
- self.max_step_size = max_step_size
- self.pool_id = pool_id
- self.device = None
-
- def memory_is_available(self):
- if "cuda" in self.device.type:
- used = torch.cuda.memory_allocated(self.device)
- reserved = torch.cuda.memory_reserved(self.device)
- elif "cpu" in self.device.type:
- process = psutil.Process(os.getpid())
- used = process.memory_info().rss
- reserved = used
- else:
- used = 0
- reserved = 0
- logger.warning(
- f"[Inference] Unsupported device type: {self.device.type}. Memory checks will not be performed."
- )
- logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}] "
- f"Memory used: {used/1024**2:.2f} MB, Max memory: {self.max_memory_bytes/1024**2:.2f} MB"
- )
- return used < self.max_memory_bytes
-
- def schedule_activate(self) -> list:
- requests = []
- while not self.waiting_queue.empty() and len(requests) < self.max_activate_size:
- if not self.memory_is_available():
- break
- requests.append(self.waiting_queue.get())
- return requests
-
- def schedule_step(self) -> list:
- requests = []
- while not self.running_queue.empty() and len(requests) < self.max_step_size:
- if not self.memory_is_available():
- break
- requests.append(self.running_queue.get())
- return requests
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import os
+1
+1import psutil
+1import torch
+1
+1from iotdb.ainode.core.inference.request_scheduler.abstract_request_scheduler import (
+1 AbstractRequestScheduler,
+1)
+1from iotdb.ainode.core.log import Logger
+1
+1logger = Logger()
+1
+1
+1class BasicRequestScheduler(AbstractRequestScheduler):
+1 """
+1 A simple FIFO request scheduler that selects requests based on memory availability and activation/step size.
+1 """
+1
+1 def __init__(
+1 self,
+1 waiting_queue,
+1 running_queue,
+1 finished_queue,
+1 pool_id,
+1 max_memory_bytes=1 << 30,
+1 max_activate_size=10,
+1 max_step_size=10,
+1 ):
+1 super().__init__(waiting_queue, running_queue, finished_queue)
+1 self.max_memory_bytes = max_memory_bytes
+1 self.max_activate_size = max_activate_size
+1 self.max_step_size = max_step_size
+1 self.pool_id = pool_id
+1 self.device = None
+1
+1 def memory_is_available(self):
+1 if "cuda" in self.device.type:
+1 used = torch.cuda.memory_allocated(self.device)
+1 reserved = torch.cuda.memory_reserved(self.device)
+1 elif "cpu" in self.device.type:
+1 process = psutil.Process(os.getpid())
+1 used = process.memory_info().rss
+1 reserved = used
+1 else:
+1 used = 0
+1 reserved = 0
+1 logger.warning(
+1 f"[Inference] Unsupported device type: {self.device.type}. Memory checks will not be performed."
+1 )
+1 logger.debug(
+1 f"[Inference][Device-{self.device}][Pool-{self.pool_id}] "
+1 f"Memory used: {used/1024**2:.2f} MB, Max memory: {self.max_memory_bytes/1024**2:.2f} MB"
+1 )
+1 return used < self.max_memory_bytes
+1
+1 def schedule_activate(self) -> list:
+1 requests = []
+1 while not self.waiting_queue.empty() and len(requests) < self.max_activate_size:
+1 if not self.memory_is_available():
+1 break
+1 requests.append(self.waiting_queue.get())
+1 return requests
+1
+1 def schedule_step(self) -> list:
+1 requests = []
+1 while not self.running_queue.empty() and len(requests) < self.max_step_size:
+1 if not self.memory_is_available():
+1 break
+1 requests.append(self.running_queue.get())
+1 return requests
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py
index 2300169a6ee..aff8a3361e6 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py
@@ -1,60 +1,61 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-
-import torch
-
-
-class AbstractInferencePipeline(ABC):
- """
- Abstract assistance strategy class for model inference.
- This class shall define the interface process for specific model.
- """
-
- def __init__(self, model_config, **infer_kwargs):
- self.model_config = model_config
- self.infer_kwargs = infer_kwargs
-
- @abstractmethod
- def preprocess_inputs(self, inputs: torch.Tensor):
- """
- Preprocess the inputs before inference, including shape validation and value transformation.
-
- Args:
- inputs (torch.Tensor): The input tensor to be preprocessed.
-
- Returns:
- torch.Tensor: The preprocessed input tensor.
- """
- # TODO: Integrate with the data processing pipeline operators
- pass
-
- @abstractmethod
- def post_decode(self):
- """
- Post-process the outputs after each decode step.
- """
- pass
-
- @abstractmethod
- def post_inference(self):
- """
- Post-process the outputs after the entire inference task.
- """
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from abc import ABC, abstractmethod
+1
+1import torch
+1
+1
+1class AbstractInferencePipeline(ABC):
+1 """
+1 Abstract assistance strategy class for model inference.
+1 This class shall define the interface process for specific model.
+1 """
+1
+1 def __init__(self, model_config, **infer_kwargs):
+1 self.model_config = model_config
+1 self.infer_kwargs = infer_kwargs
+1
+1 @abstractmethod
+1 def preprocess_inputs(self, inputs: torch.Tensor):
+1 """
+1 Preprocess the inputs before inference, including shape validation and value transformation.
+1
+1 Args:
+1 inputs (torch.Tensor): The input tensor to be preprocessed.
+1
+1 Returns:
+1 torch.Tensor: The preprocessed input tensor.
+1 """
+1 # TODO: Integrate with the data processing pipeline operators
+1 pass
+1
+1 @abstractmethod
+1 def post_decode(self):
+1 """
+1 Post-process the outputs after each decode step.
+1 """
+1 pass
+1
+1 @abstractmethod
+1 def post_inference(self):
+1 """
+1 Post-process the outputs after the entire inference task.
+1 """
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py
index 17c88e32fb5..5bcc2018a7d 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py
@@ -1,51 +1,52 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import torch
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
-from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
-
-
-class TimerSundialInferencePipeline(AbstractInferencePipeline):
- """
- Strategy for Timer-Sundial model inference.
- """
-
- def __init__(self, model_config: SundialConfig, **infer_kwargs):
- super().__init__(model_config, infer_kwargs=infer_kwargs)
-
- def preprocess_inputs(self, inputs: torch.Tensor):
- super().preprocess_inputs(inputs)
- if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
- f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
- )
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- return inputs
-
- def post_decode(self):
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- pass
-
- def post_inference(self):
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import torch
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
+1 AbstractInferencePipeline,
+1)
+1from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
+1
+1
+1class TimerSundialInferencePipeline(AbstractInferencePipeline):
+1 """
+1 Strategy for Timer-Sundial model inference.
+1 """
+1
+1 def __init__(self, model_config: SundialConfig, **infer_kwargs):
+1 super().__init__(model_config, infer_kwargs=infer_kwargs)
+1
+1 def preprocess_inputs(self, inputs: torch.Tensor):
+1 super().preprocess_inputs(inputs)
+1 if len(inputs.shape) != 2:
+1 raise InferenceModelInternalError(
+1 f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
+1 )
+1 # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
+1 return inputs
+1
+1 def post_decode(self):
+1 # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
+1 pass
+1
+1 def post_inference(self):
+1 # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py
index dc1dd304f68..0ed8fd2b14b 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py
@@ -1,51 +1,52 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import torch
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
-from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
-
-
-class TimerXLInferencePipeline(AbstractInferencePipeline):
- """
- Strategy for Timer-XL model inference.
- """
-
- def __init__(self, model_config: TimerConfig, **infer_kwargs):
- super().__init__(model_config, infer_kwargs=infer_kwargs)
-
- def preprocess_inputs(self, inputs: torch.Tensor):
- super().preprocess_inputs(inputs)
- if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
- f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
- )
- # Considering that we are currently using the generate function interface, it seems that no pre-processing is required
- return inputs
-
- def post_decode(self):
- # Considering that we are currently using the generate function interface, it seems that no post-processing is required
- pass
-
- def post_inference(self):
- # Considering that we are currently using the generate function interface, it seems that no post-processing is required
- pass
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import torch
+1
+1from iotdb.ainode.core.exception import InferenceModelInternalError
+1from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
+1 AbstractInferencePipeline,
+1)
+1from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
+1
+1
+1class TimerXLInferencePipeline(AbstractInferencePipeline):
+1 """
+1 Strategy for Timer-XL model inference.
+1 """
+1
+1 def __init__(self, model_config: TimerConfig, **infer_kwargs):
+1 super().__init__(model_config, infer_kwargs=infer_kwargs)
+1
+1 def preprocess_inputs(self, inputs: torch.Tensor):
+1 super().preprocess_inputs(inputs)
+1 if len(inputs.shape) != 2:
+1 raise InferenceModelInternalError(
+1 f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
+1 )
+1 # Considering that we are currently using the generate function interface, it seems that no pre-processing is required
+1 return inputs
+1
+1 def post_decode(self):
+1 # Considering that we are currently using the generate function interface, it seems that no post-processing is required
+1 pass
+1
+1 def post_inference(self):
+1 # Considering that we are currently using the generate function interface, it seems that no post-processing is required
+1 pass
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/utils.py b/iotdb-core/ainode/iotdb/ainode/core/inference/utils.py
index cf10b5b2cd4..9d26a6b1130 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/utils.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/utils.py
@@ -1,80 +1,81 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import secrets
-import string
-
-import torch
-from transformers.modeling_outputs import MoeCausalLMOutputWithPast
-
-
-def generate_req_id(length=10, charset=string.ascii_letters + string.digits) -> str:
- """
- Generate a random req_id string of specified length.
- The length is 10 by default, with 10^{17} possible combinations.
- """
- return "".join(secrets.choice(charset) for _ in range(length))
-
-
-def _slice_tensor(t, s, e):
- return None if t is None else t[s:e]
-
-
-def _slice_tuple_of_tensors(tup, s, e):
- """
- hidden_states / attentions: Tuple[layer0, layer1, ...]
- every layer maybe Tensor or None。
- """
- if tup is None:
- return None
- sliced = []
- for x in tup:
- sliced.append(_slice_tensor(x, s, e) if torch.is_tensor(x) else x)
- return tuple(sliced)
-
-
-def _slice_pkv(pkv, s, e):
- if pkv is None:
- return None
- out = []
- for layer in pkv: # layer: Tuple[key, value, ...]
- out.append(tuple(x[s:e] for x in layer))
- return out
-
-
-def split_moe_output(batch_out: MoeCausalLMOutputWithPast, split_sizes):
- """
- split batch_out with type: MoeCausalLMOutputWithPast into len(split_sizes)
- split_sizes[i] = ith request's batch_size。
- """
- outs = []
- start = 0
- for bsz in split_sizes:
- end = start + bsz
- outs.append(
- MoeCausalLMOutputWithPast(
- loss=_slice_tensor(batch_out.loss, start, end),
- logits=batch_out.logits[start:end],
- past_key_values=_slice_pkv(batch_out.past_key_values, start, end),
- hidden_states=_slice_tuple_of_tensors(
- batch_out.hidden_states, start, end
- ),
- attentions=_slice_tuple_of_tensors(batch_out.attentions, start, end),
- )
- )
- start = end
- return outs
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import secrets
+1import string
+1
+1import torch
+1from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+1
+1
+1def generate_req_id(length=10, charset=string.ascii_letters + string.digits) -> str:
+1 """
+1 Generate a random req_id string of specified length.
+1 The length is 10 by default, with 10^{17} possible combinations.
+1 """
+1 return "".join(secrets.choice(charset) for _ in range(length))
+1
+1
+1def _slice_tensor(t, s, e):
+1 return None if t is None else t[s:e]
+1
+1
+1def _slice_tuple_of_tensors(tup, s, e):
+1 """
+1 hidden_states / attentions: Tuple[layer0, layer1, ...]
+1 every layer maybe Tensor or None。
+1 """
+1 if tup is None:
+1 return None
+1 sliced = []
+1 for x in tup:
+1 sliced.append(_slice_tensor(x, s, e) if torch.is_tensor(x) else x)
+1 return tuple(sliced)
+1
+1
+1def _slice_pkv(pkv, s, e):
+1 if pkv is None:
+1 return None
+1 out = []
+1 for layer in pkv: # layer: Tuple[key, value, ...]
+1 out.append(tuple(x[s:e] for x in layer))
+1 return out
+1
+1
+1def split_moe_output(batch_out: MoeCausalLMOutputWithPast, split_sizes):
+1 """
+1 split batch_out with type: MoeCausalLMOutputWithPast into len(split_sizes)
+1 split_sizes[i] = ith request's batch_size。
+1 """
+1 outs = []
+1 start = 0
+1 for bsz in split_sizes:
+1 end = start + bsz
+1 outs.append(
+1 MoeCausalLMOutputWithPast(
+1 loss=_slice_tensor(batch_out.loss, start, end),
+1 logits=batch_out.logits[start:end],
+1 past_key_values=_slice_pkv(batch_out.past_key_values, start, end),
+1 hidden_states=_slice_tuple_of_tensors(
+1 batch_out.hidden_states, start, end
+1 ),
+1 attentions=_slice_tuple_of_tensors(batch_out.attentions, start, end),
+1 )
+1 )
+1 start = end
+1 return outs
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py
index 4e3b5293c16..1b9e95ba0a5 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py
@@ -1,46 +1,47 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from torch.utils.data import Dataset
-
-
-class BasicDatabaseDataset(Dataset):
- def __init__(self, ip: str, port: int):
- self.ip = ip
- self.port = port
-
-
-class BasicDatabaseForecastDataset(BasicDatabaseDataset):
- def __init__(
- self,
- ip: str,
- port: int,
- seq_len: int,
- input_token_len: int,
- output_token_len: int,
- window_step: int,
- ):
- super().__init__(ip, port)
- # The number of the time series data points of the model input
- self.seq_len = seq_len
- # The number of the time series data points of each model token
- self.input_token_len = input_token_len
- # The number of the time series data points of the model output
- self.output_token_len = output_token_len
- self.token_num = self.seq_len // self.input_token_len
- self.window_step = window_step
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1from torch.utils.data import Dataset
+1
+1
+1class BasicDatabaseDataset(Dataset):
+1 def __init__(self, ip: str, port: int):
+1 self.ip = ip
+1 self.port = port
+1
+1
+1class BasicDatabaseForecastDataset(BasicDatabaseDataset):
+1 def __init__(
+1 self,
+1 ip: str,
+1 port: int,
+1 seq_len: int,
+1 input_token_len: int,
+1 output_token_len: int,
+1 window_step: int,
+1 ):
+1 super().__init__(ip, port)
+1 # The number of the time series data points of the model input
+1 self.seq_len = seq_len
+1 # The number of the time series data points of each model token
+1 self.input_token_len = input_token_len
+1 # The number of the time series data points of the model output
+1 self.output_token_len = output_token_len
+1 self.token_num = self.seq_len // self.input_token_len
+1 self.window_step = window_step
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py
index 13c56ca9d2d..426e98833b4 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py
@@ -1,401 +1,402 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-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
-
-logger = Logger()
-
-
-def get_field_value(field: Field):
- data_type = field.get_data_type()
- if data_type == TSDataType.INT32:
- return field.get_int_value()
- elif data_type == TSDataType.INT64:
- return field.get_long_value()
- elif data_type == TSDataType.FLOAT:
- return field.get_float_value()
- elif data_type == TSDataType.DOUBLE:
- return field.get_double_value()
- else:
- return field.get_string_value()
-
-
-def _cache_enable() -> bool:
- return AINodeDescriptor().get_config().get_ain_data_storage_cache_size() > 0
-
-
-class IoTDBTreeModelDataset(BasicDatabaseForecastDataset):
- cache = MemoryLRUCache()
-
- def __init__(
- self,
- model_id: str,
- seq_len: int,
- input_token_len: int,
- output_token_len: int,
- window_step: int,
- data_schema_list: list,
- ip: str = AINodeDescriptor().get_config().get_ain_cluster_ingress_address(),
- port: int = AINodeDescriptor().get_config().get_ain_cluster_ingress_port(),
- username: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_username(),
- password: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_password(),
- time_zone: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_time_zone(),
- use_rate: float = 1.0,
- offset_rate: float = 0.0,
- ):
- super().__init__(
- ip, port, seq_len, input_token_len, output_token_len, window_step
- )
-
- self.SHOW_TIMESERIES = "show timeseries %s%s"
- self.COUNT_SERIES_SQL = "select count(%s) from %s%s"
- self.FETCH_SERIES_SQL = "select %s from %s%s"
- self.FETCH_SERIES_RANGE_SQL = "select %s from %s offset %s limit %s%s"
-
- self.TIME_CONDITION = " where time>=%s and time<%s"
-
- self.session = Session.init_from_node_urls(
- node_urls=[f"{ip}:{port}"],
- 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
- self.offset_rate = offset_rate
- self._fetch_schema(data_schema_list)
- self.cache_enable = _cache_enable()
- self.cache_key_prefix = model_id + "_"
-
- def _fetch_schema(self, data_schema_list: list):
- series_to_length = {}
- for schema in data_schema_list:
- path_pattern = schema.schemaName
- series_list = []
- time_condition = (
- self.TIME_CONDITION % (schema.timeRange[0], schema.timeRange[1])
- if schema.timeRange
- else ""
- )
- with self.session.execute_query_statement(
- self.SHOW_TIMESERIES % (path_pattern, time_condition)
- ) as show_timeseries_result:
- while show_timeseries_result.has_next():
- # parse the name of time series according to data_schema_list
- series_list.append(
- get_field_value(show_timeseries_result.next().get_fields()[0])
- )
-
- for series in series_list:
- split_series = series.split(".")
- with self.session.execute_query_statement(
- self.COUNT_SERIES_SQL
- % (split_series[-1], ".".join(split_series[:-1]), time_condition)
- ) as count_series_result:
- while count_series_result.has_next():
- length = get_field_value(
- count_series_result.next().get_fields()[0]
- )
- # count the number (length) of time series data points
- # structure: {series_name: (split_series_name, length, time_condition)}
- series_to_length[series] = (
- split_series,
- length,
- time_condition,
- )
-
- sorted_series = sorted(series_to_length.items(), key=lambda x: x[1][1])
- # structure: [(split_series_name, the number of windows of this series, prefix sum of window number, window start offset, time_condition of this series), ...]
- sorted_series_with_prefix_sum = []
- window_sum = 0
- for seq_name, seq_value in sorted_series:
- # calculate and sum the number of training data windows for each time series
- window_count = (
- seq_value[1] - self.seq_len - self.output_token_len + 1
- ) // self.window_step
- if window_count <= 1:
- continue
- use_window_count = int(window_count * self.use_rate)
- window_sum += use_window_count
- sorted_series_with_prefix_sum.append(
- (
- seq_value[0],
- use_window_count,
- window_sum,
- int(window_count * self.offset_rate),
- seq_value[2],
- )
- )
-
- self.total_window_count = window_sum
- self.sorted_series = sorted_series_with_prefix_sum
-
- def __getitem__(self, index):
- window_index = index
- # locate the series to be queried
- series_index = 0
- while self.sorted_series[series_index][2] < window_index:
- series_index += 1
- # locate the window of this series to be queried
- if series_index != 0:
- window_index -= self.sorted_series[series_index - 1][2]
- series = self.sorted_series[series_index][0]
- window_index += self.sorted_series[series_index][3]
- time_condition = self.sorted_series[series_index][4]
- if self.cache_enable:
- cache_key = self.cache_key_prefix + ".".join(series) + time_condition
- series_data = self.cache.get(cache_key)
- if series_data is not None:
- # try to get the training data window from cache first
- series_data = torch.tensor(series_data)
- result = series_data[
- window_index * self.window_step : window_index * self.window_step
- + self.seq_len
- + self.output_token_len
- ]
- return (
- result[0 : self.seq_len],
- result[self.input_token_len : self.seq_len + self.output_token_len],
- np.ones(self.token_num, dtype=np.int32),
- )
- series_data = []
- sql = ""
- try:
- if self.cache_enable:
- # fetch all data points of the specific series when cache is enabled
- sql = self.FETCH_SERIES_SQL % (
- series[-1],
- ".".join(series[0:-1]),
- time_condition,
- )
- else:
- # TODO: fetch only the specific data window, current logic is wrong
- sql = self.FETCH_SERIES_RANGE_SQL % (
- series[-1],
- ".".join(series[0:-1]),
- window_index,
- self.seq_len,
- time_condition,
- )
- with self.session.execute_query_statement(sql) as query_result:
- while query_result.has_next():
- series_data.append(
- get_field_value(query_result.next().get_fields()[0])
- )
- except Exception as e:
- logger.error("Executing sql: {} with exception: {}".format(sql, e))
- if self.cache_enable:
- self.cache.put(cache_key, series_data)
- result = series_data[
- window_index * self.window_step : window_index * self.window_step
- + self.seq_len
- + self.output_token_len
- ]
- result = torch.tensor(result)
- return (
- result[0 : self.seq_len],
- result[self.input_token_len : self.seq_len + self.output_token_len],
- np.ones(self.token_num, dtype=np.int32),
- )
-
- def __len__(self):
- return self.total_window_count
-
-
-class IoTDBTableModelDataset(BasicDatabaseForecastDataset):
-
- DEFAULT_TAG = "__DEFAULT_TAG__"
-
- def __init__(
- self,
- model_id: str,
- seq_len: int,
- input_token_len: int,
- output_token_len: int,
- window_step: int,
- data_schema_list: list,
- ip: str = AINodeDescriptor().get_config().get_ain_cluster_ingress_address(),
- port: int = AINodeDescriptor().get_config().get_ain_cluster_ingress_port(),
- username: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_username(),
- password: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_password(),
- time_zone: str = AINodeDescriptor()
- .get_config()
- .get_ain_cluster_ingress_time_zone(),
- use_rate: float = 1.0,
- offset_rate: float = 0.0,
- ):
- super().__init__(
- ip, port, seq_len, input_token_len, output_token_len, window_step
- )
-
- table_session_config = TableSessionConfig(
- node_urls=[f"{ip}:{port}"],
- 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
- self.offset_rate = offset_rate
-
- # used for caching data
- self._fetch_schema(data_schema_list)
-
- def _fetch_schema(self, data_schema_list: list):
- series_map = {}
- for target_sql in data_schema_list:
- target_sql = target_sql.schemaName
- with self.session.execute_query_statement(target_sql) as target_data:
- while target_data.has_next():
- cur_data = target_data.next()
- # TODO: currently, we only support the following simple table form
- time_col, value_col, tag_col = -1, -1, -1
- for i, field in enumerate(cur_data.get_fields()):
- if field.get_data_type() == TSDataType.TIMESTAMP:
- time_col = i
- elif field.get_data_type() in (
- TSDataType.INT32,
- TSDataType.INT64,
- TSDataType.FLOAT,
- TSDataType.DOUBLE,
- ):
- value_col = i
- elif field.get_data_type() == TSDataType.TEXT:
- tag_col = i
- if time_col == -1 or value_col == -1:
- raise ValueError(
- "The training cannot start due to invalid data schema"
- )
- if tag_col == -1:
- tag = self.DEFAULT_TAG
- else:
- tag = cur_data.get_fields()[tag_col].get_string_value()
- if tag not in series_map:
- series_map[tag] = []
- series_list = series_map[tag]
- series_list.append(
- get_field_value(cur_data.get_fields()[value_col])
- )
-
- # TODO: Unify the following implementation
- # structure: [(series_name, the number of windows of this series, prefix sum of window number, window start offset, series_data), ...]
- series_with_prefix_sum = []
- window_sum = 0
- for seq_name, seq_values in series_map.items():
- # calculate and sum the number of training data windows for each time series
- window_count = (
- len(seq_values) - self.seq_len - self.output_token_len + 1
- ) // self.window_step
- if window_count <= 1:
- continue
- use_window_count = int(window_count * self.use_rate)
- window_sum += use_window_count
- series_with_prefix_sum.append(
- (
- seq_name,
- use_window_count,
- window_sum,
- int(window_count * self.offset_rate),
- seq_values,
- )
- )
-
- self.total_window_count = window_sum
- self.series_with_prefix_sum = series_with_prefix_sum
-
- def __getitem__(self, index):
- window_index = index
- # locate the series to be queried
- series_index = 0
- while self.series_with_prefix_sum[series_index][1] < window_index:
- series_index += 1
- # locate the window of this series to be queried
- if series_index != 0:
- window_index -= self.series_with_prefix_sum[series_index - 1][2]
- window_index += self.series_with_prefix_sum[series_index][3]
- result = self.series_with_prefix_sum[series_index][4][
- window_index * self.window_step : window_index * self.window_step
- + self.seq_len
- + self.output_token_len
- ]
- result = torch.tensor(result)
- return (
- result[0 : self.seq_len],
- result[self.input_token_len : self.seq_len + self.output_token_len],
- np.ones(self.token_num, dtype=np.int32),
- )
-
- def __len__(self):
- return self.total_window_count
-
-
-def register_dataset(key: str, dataset: Dataset):
- DatasetFactory().register(key, dataset)
-
-
-@singleton
-class DatasetFactory(object):
-
- def __init__(self):
- self.dataset_list = {
- "iotdb.table": IoTDBTableModelDataset,
- "iotdb.tree": IoTDBTreeModelDataset,
- }
-
- def register(self, key: str, dataset: Dataset):
- if key not in self.dataset_list:
- self.dataset_list[key] = dataset
- else:
- raise KeyError(f"Dataset {key} already exists")
-
- def deregister(self, key: str):
- del self.dataset_list[key]
-
- def get_dataset(self, key: str):
- if key not in self.dataset_list.keys():
- raise KeyError(f"Dataset {key} does not exist")
- return self.dataset_list[key]
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import numpy as np
+1import torch
+1from torch.utils.data import Dataset
+1
+1from iotdb.ainode.core.config import AINodeDescriptor
+1from iotdb.ainode.core.ingress.dataset import BasicDatabaseForecastDataset
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.util.cache import MemoryLRUCache
+1from iotdb.ainode.core.util.decorator import singleton
+1from iotdb.Session import Session
+1from iotdb.table_session import TableSession, TableSessionConfig
+1from iotdb.utils.Field import Field
+1from iotdb.utils.IoTDBConstants import TSDataType
+1
+1logger = Logger()
+1
+1
+1def get_field_value(field: Field):
+1 data_type = field.get_data_type()
+1 if data_type == TSDataType.INT32:
+1 return field.get_int_value()
+1 elif data_type == TSDataType.INT64:
+1 return field.get_long_value()
+1 elif data_type == TSDataType.FLOAT:
+1 return field.get_float_value()
+1 elif data_type == TSDataType.DOUBLE:
+1 return field.get_double_value()
+1 else:
+1 return field.get_string_value()
+1
+1
+1def _cache_enable() -> bool:
+1 return AINodeDescriptor().get_config().get_ain_data_storage_cache_size() > 0
+1
+1
+1class IoTDBTreeModelDataset(BasicDatabaseForecastDataset):
+1 cache = MemoryLRUCache()
+1
+1 def __init__(
+1 self,
+1 model_id: str,
+1 seq_len: int,
+1 input_token_len: int,
+1 output_token_len: int,
+1 window_step: int,
+1 data_schema_list: list,
+1 ip: str = AINodeDescriptor().get_config().get_ain_cluster_ingress_address(),
+1 port: int = AINodeDescriptor().get_config().get_ain_cluster_ingress_port(),
+1 username: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_username(),
+1 password: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_password(),
+1 time_zone: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_time_zone(),
+1 use_rate: float = 1.0,
+1 offset_rate: float = 0.0,
+1 ):
+1 super().__init__(
+1 ip, port, seq_len, input_token_len, output_token_len, window_step
+1 )
+1
+1 self.SHOW_TIMESERIES = "show timeseries %s%s"
+1 self.COUNT_SERIES_SQL = "select count(%s) from %s%s"
+1 self.FETCH_SERIES_SQL = "select %s from %s%s"
+1 self.FETCH_SERIES_RANGE_SQL = "select %s from %s offset %s limit %s%s"
+1
+1 self.TIME_CONDITION = " where time>=%s and time<%s"
+1
+1 self.session = Session.init_from_node_urls(
+1 node_urls=[f"{ip}:{port}"],
+1 user=username,
+1 password=password,
+1 zone_id=time_zone,
+1 use_ssl=AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_ssl_enabled(),
+1 ca_certs=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(),
+1 )
+1 self.session.open(False)
+1 self.use_rate = use_rate
+1 self.offset_rate = offset_rate
+1 self._fetch_schema(data_schema_list)
+1 self.cache_enable = _cache_enable()
+1 self.cache_key_prefix = model_id + "_"
+1
+1 def _fetch_schema(self, data_schema_list: list):
+1 series_to_length = {}
+1 for schema in data_schema_list:
+1 path_pattern = schema.schemaName
+1 series_list = []
+1 time_condition = (
+1 self.TIME_CONDITION % (schema.timeRange[0], schema.timeRange[1])
+1 if schema.timeRange
+1 else ""
+1 )
+1 with self.session.execute_query_statement(
+1 self.SHOW_TIMESERIES % (path_pattern, time_condition)
+1 ) as show_timeseries_result:
+1 while show_timeseries_result.has_next():
+1 # parse the name of time series according to data_schema_list
+1 series_list.append(
+1 get_field_value(show_timeseries_result.next().get_fields()[0])
+1 )
+1
+1 for series in series_list:
+1 split_series = series.split(".")
+1 with self.session.execute_query_statement(
+1 self.COUNT_SERIES_SQL
+1 % (split_series[-1], ".".join(split_series[:-1]), time_condition)
+1 ) as count_series_result:
+1 while count_series_result.has_next():
+1 length = get_field_value(
+1 count_series_result.next().get_fields()[0]
+1 )
+1 # count the number (length) of time series data points
+1 # structure: {series_name: (split_series_name, length, time_condition)}
+1 series_to_length[series] = (
+1 split_series,
+1 length,
+1 time_condition,
+1 )
+1
+1 sorted_series = sorted(series_to_length.items(), key=lambda x: x[1][1])
+1 # structure: [(split_series_name, the number of windows of this series, prefix sum of window number, window start offset, time_condition of this series), ...]
+1 sorted_series_with_prefix_sum = []
+1 window_sum = 0
+1 for seq_name, seq_value in sorted_series:
+1 # calculate and sum the number of training data windows for each time series
+1 window_count = (
+1 seq_value[1] - self.seq_len - self.output_token_len + 1
+1 ) // self.window_step
+1 if window_count <= 1:
+1 continue
+1 use_window_count = int(window_count * self.use_rate)
+1 window_sum += use_window_count
+1 sorted_series_with_prefix_sum.append(
+1 (
+1 seq_value[0],
+1 use_window_count,
+1 window_sum,
+1 int(window_count * self.offset_rate),
+1 seq_value[2],
+1 )
+1 )
+1
+1 self.total_window_count = window_sum
+1 self.sorted_series = sorted_series_with_prefix_sum
+1
+1 def __getitem__(self, index):
+1 window_index = index
+1 # locate the series to be queried
+1 series_index = 0
+1 while self.sorted_series[series_index][2] < window_index:
+1 series_index += 1
+1 # locate the window of this series to be queried
+1 if series_index != 0:
+1 window_index -= self.sorted_series[series_index - 1][2]
+1 series = self.sorted_series[series_index][0]
+1 window_index += self.sorted_series[series_index][3]
+1 time_condition = self.sorted_series[series_index][4]
+1 if self.cache_enable:
+1 cache_key = self.cache_key_prefix + ".".join(series) + time_condition
+1 series_data = self.cache.get(cache_key)
+1 if series_data is not None:
+1 # try to get the training data window from cache first
+1 series_data = torch.tensor(series_data)
+1 result = series_data[
+1 window_index * self.window_step : window_index * self.window_step
+1 + self.seq_len
+1 + self.output_token_len
+1 ]
+1 return (
+1 result[0 : self.seq_len],
+1 result[self.input_token_len : self.seq_len + self.output_token_len],
+1 np.ones(self.token_num, dtype=np.int32),
+1 )
+1 series_data = []
+1 sql = ""
+1 try:
+1 if self.cache_enable:
+1 # fetch all data points of the specific series when cache is enabled
+1 sql = self.FETCH_SERIES_SQL % (
+1 series[-1],
+1 ".".join(series[0:-1]),
+1 time_condition,
+1 )
+1 else:
+1 # TODO: fetch only the specific data window, current logic is wrong
+1 sql = self.FETCH_SERIES_RANGE_SQL % (
+1 series[-1],
+1 ".".join(series[0:-1]),
+1 window_index,
+1 self.seq_len,
+1 time_condition,
+1 )
+1 with self.session.execute_query_statement(sql) as query_result:
+1 while query_result.has_next():
+1 series_data.append(
+1 get_field_value(query_result.next().get_fields()[0])
+1 )
+1 except Exception as e:
+1 logger.error("Executing sql: {} with exception: {}".format(sql, e))
+1 if self.cache_enable:
+1 self.cache.put(cache_key, series_data)
+1 result = series_data[
+1 window_index * self.window_step : window_index * self.window_step
+1 + self.seq_len
+1 + self.output_token_len
+1 ]
+1 result = torch.tensor(result)
+1 return (
+1 result[0 : self.seq_len],
+1 result[self.input_token_len : self.seq_len + self.output_token_len],
+1 np.ones(self.token_num, dtype=np.int32),
+1 )
+1
+1 def __len__(self):
+1 return self.total_window_count
+1
+1
+1class IoTDBTableModelDataset(BasicDatabaseForecastDataset):
+1
+1 DEFAULT_TAG = "__DEFAULT_TAG__"
+1
+1 def __init__(
+1 self,
+1 model_id: str,
+1 seq_len: int,
+1 input_token_len: int,
+1 output_token_len: int,
+1 window_step: int,
+1 data_schema_list: list,
+1 ip: str = AINodeDescriptor().get_config().get_ain_cluster_ingress_address(),
+1 port: int = AINodeDescriptor().get_config().get_ain_cluster_ingress_port(),
+1 username: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_username(),
+1 password: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_password(),
+1 time_zone: str = AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_time_zone(),
+1 use_rate: float = 1.0,
+1 offset_rate: float = 0.0,
+1 ):
+1 super().__init__(
+1 ip, port, seq_len, input_token_len, output_token_len, window_step
+1 )
+1
+1 table_session_config = TableSessionConfig(
+1 node_urls=[f"{ip}:{port}"],
+1 username=username,
+1 password=password,
+1 time_zone=time_zone,
+1 use_ssl=AINodeDescriptor()
+1 .get_config()
+1 .get_ain_cluster_ingress_ssl_enabled(),
+1 ca_certs=AINodeDescriptor().get_config().get_ain_thrift_ssl_cert_file(),
+1 )
+1 self.session = TableSession(table_session_config)
+1 self.use_rate = use_rate
+1 self.offset_rate = offset_rate
+1
+1 # used for caching data
+1 self._fetch_schema(data_schema_list)
+1
+1 def _fetch_schema(self, data_schema_list: list):
+1 series_map = {}
+1 for target_sql in data_schema_list:
+1 target_sql = target_sql.schemaName
+1 with self.session.execute_query_statement(target_sql) as target_data:
+1 while target_data.has_next():
+1 cur_data = target_data.next()
+1 # TODO: currently, we only support the following simple table form
+1 time_col, value_col, tag_col = -1, -1, -1
+1 for i, field in enumerate(cur_data.get_fields()):
+1 if field.get_data_type() == TSDataType.TIMESTAMP:
+1 time_col = i
+1 elif field.get_data_type() in (
+1 TSDataType.INT32,
+1 TSDataType.INT64,
+1 TSDataType.FLOAT,
+1 TSDataType.DOUBLE,
+1 ):
+1 value_col = i
+1 elif field.get_data_type() == TSDataType.TEXT:
+1 tag_col = i
+1 if time_col == -1 or value_col == -1:
+1 raise ValueError(
+1 "The training cannot start due to invalid data schema"
+1 )
+1 if tag_col == -1:
+1 tag = self.DEFAULT_TAG
+1 else:
+1 tag = cur_data.get_fields()[tag_col].get_string_value()
+1 if tag not in series_map:
+1 series_map[tag] = []
+1 series_list = series_map[tag]
+1 series_list.append(
+1 get_field_value(cur_data.get_fields()[value_col])
+1 )
+1
+1 # TODO: Unify the following implementation
+1 # structure: [(series_name, the number of windows of this series, prefix sum of window number, window start offset, series_data), ...]
+1 series_with_prefix_sum = []
+1 window_sum = 0
+1 for seq_name, seq_values in series_map.items():
+1 # calculate and sum the number of training data windows for each time series
+1 window_count = (
+1 len(seq_values) - self.seq_len - self.output_token_len + 1
+1 ) // self.window_step
+1 if window_count <= 1:
+1 continue
+1 use_window_count = int(window_count * self.use_rate)
+1 window_sum += use_window_count
+1 series_with_prefix_sum.append(
+1 (
+1 seq_name,
+1 use_window_count,
+1 window_sum,
+1 int(window_count * self.offset_rate),
+1 seq_values,
+1 )
+1 )
+1
+1 self.total_window_count = window_sum
+1 self.series_with_prefix_sum = series_with_prefix_sum
+1
+1 def __getitem__(self, index):
+1 window_index = index
+1 # locate the series to be queried
+1 series_index = 0
+1 while self.series_with_prefix_sum[series_index][1] < window_index:
+1 series_index += 1
+1 # locate the window of this series to be queried
+1 if series_index != 0:
+1 window_index -= self.series_with_prefix_sum[series_index - 1][2]
+1 window_index += self.series_with_prefix_sum[series_index][3]
+1 result = self.series_with_prefix_sum[series_index][4][
+1 window_index * self.window_step : window_index * self.window_step
+1 + self.seq_len
+1 + self.output_token_len
+1 ]
+1 result = torch.tensor(result)
+1 return (
+1 result[0 : self.seq_len],
+1 result[self.input_token_len : self.seq_len + self.output_token_len],
+1 np.ones(self.token_num, dtype=np.int32),
+1 )
+1
+1 def __len__(self):
+1 return self.total_window_count
+1
+1
+1def register_dataset(key: str, dataset: Dataset):
+1 DatasetFactory().register(key, dataset)
+1
+1
+1@singleton
+1class DatasetFactory(object):
+1
+1 def __init__(self):
+1 self.dataset_list = {
+1 "iotdb.table": IoTDBTableModelDataset,
+1 "iotdb.tree": IoTDBTreeModelDataset,
+1 }
+1
+1 def register(self, key: str, dataset: Dataset):
+1 if key not in self.dataset_list:
+1 self.dataset_list[key] = dataset
+1 else:
+1 raise KeyError(f"Dataset {key} already exists")
+1
+1 def deregister(self, key: str):
+1 del self.dataset_list[key]
+1
+1 def get_dataset(self, key: str):
+1 if key not in self.dataset_list.keys():
+1 raise KeyError(f"Dataset {key} does not exist")
+1 return self.dataset_list[key]
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/log.py b/iotdb-core/ainode/iotdb/ainode/core/log.py
index a7a05d0ea92..e23966f2bde 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/log.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/log.py
@@ -1,161 +1,162 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import gzip
-import logging
-import os
-import random
-import re
-import shutil
-import sys
-import threading
-from logging.handlers import TimedRotatingFileHandler
-
-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 iotdb.ainode.core.util.decorator import singleton
-
-
-class BaseLogger:
- """
- BaseLogger is a base class for logging, which implements daily compress, log files management, custom format, etc.
-
- Args:
- log_file_name_prefix: the prefix of the log file name, it is used to distinguish different processes.
- log_dir: log directory, default is AINODE_HOME/logs.
-
- Attributes:
- logger_format: log format
- logger: global logger with custom format and level
- console_handler: console handler for stdout
- _lock: process lock for logger. This is just a precaution, we currently do not have multiprocessing
- """
-
- def __init__(self, log_file_name_prefix: str, log_dir=AINODE_LOG_DIR):
- self.logger_format = logging.Formatter(
- fmt="%(asctime)s %(levelname)s [%(process)d:%(processName)s] "
- "%(filename)s:%(funcName)s:%(lineno)d - %(message)s",
- datefmt="%Y-%m-%d %H:%M:%S",
- )
-
- self._lock = threading.Lock()
-
- self.logger = logging.getLogger(str(random.random()))
- self.logger.handlers.clear()
- self.logger.setLevel(DEFAULT_LOG_LEVEL)
- self.console_handler = logging.StreamHandler(sys.stdout)
- self.console_handler.setLevel(DEFAULT_LOG_LEVEL)
- self.console_handler.setFormatter(self.logger_format)
- self.logger.addHandler(self.console_handler)
-
- # Set log file handler
- for i in range(len(LOG_FILE_TYPE)):
- file_name = log_file_name_prefix + LOG_FILE_TYPE[i] + ".log"
- # create log file if not exist
- os.makedirs(log_dir, exist_ok=True)
- file_path = os.path.join(log_dir, f"{file_name}")
- if not os.path.exists(file_path):
- with open(file_path, "w", encoding="utf-8") as f:
- f.write("")
- # create handler
- file_level = AINODE_LOG_FILE_LEVELS[i]
- file_handler = TimedRotatingFileHandler(
- filename=file_path,
- when="MIDNIGHT",
- interval=1,
- encoding="utf-8",
- )
- if LOG_FILE_TYPE[i] != "all":
- file_handler.setLevel(logging.NOTSET)
- file_handler.addFilter(LogLevelFilter(file_level))
-
- # set renamer
- def universal_namer(
- default_name: str, internal_file_name: str = file_name
- ) -> str:
- # to avoid outer variable late binding
- base, ext = os.path.splitext(internal_file_name)
- base = base.replace("_", "-")
- base_log_dir = os.path.dirname(default_name)
- suffix = default_name.rsplit(".", 1)[-1] # e.g. 2025-08-01_13-45
- digits = re.sub(r"[-_]", "", suffix)
- return os.path.join(base_log_dir, f"{base}-{digits}{ext}.gz")
-
- file_handler.namer = universal_namer
-
- # set gzip
- def gzip_rotator(src: str, dst: str):
- with open(src, "rb") as f_in, gzip.open(dst, "wb") as f_out:
- shutil.copyfileobj(f_in, f_out)
- # delete the old .log file
- os.remove(src)
-
- file_handler.rotator = gzip_rotator
- # other settings
- file_handler.setLevel(file_level)
- file_handler.setFormatter(self.logger_format)
- self.logger.addHandler(file_handler)
-
- self.info(f"Logger init successfully.")
-
- # interfaces
- def debug(self, *msg):
- self._write(self.logger.debug, *msg)
-
- def info(self, *msg):
- self._write(self.logger.info, *msg)
-
- def warning(self, *msg):
- self._write(self.logger.warning, *msg)
-
- def error(self, *msg):
- self._write(self.logger.error, *msg)
-
- def _write(self, function, *msg):
- with self._lock:
- # The stack level of caller is 3, because:
- # caller -> BaseLogger.info -> BaseLogger._write
- function(" ".join(map(str, msg)), stacklevel=3)
-
-
-@singleton
-class Logger(BaseLogger):
- """
- Logger is a singleton, just use Logger() to get it in the specified process.
-
- Args:
- log_file_name_prefix: the prefix of the log file name, it is used to distinguish inference, training and the main process.
- """
-
- def __init__(self, log_file_name_prefix: str = AINODE_LOG_FILE_NAME_PREFIX):
- super().__init__(log_file_name_prefix=log_file_name_prefix)
-
-
-class LogLevelFilter(logging.Filter):
-
- def __init__(self, level: int):
- super().__init__()
- self.level = level
-
- def filter(self, record: logging.LogRecord) -> bool:
- return record.levelno == self.level
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import gzip
+1import logging
+1import os
+1import random
+1import re
+1import shutil
+1import sys
+1import threading
+1from logging.handlers import TimedRotatingFileHandler
+1
+1from iotdb.ainode.core.constant import (
+1 AINODE_LOG_DIR,
+1 AINODE_LOG_FILE_LEVELS,
+1 AINODE_LOG_FILE_NAME_PREFIX,
+1 DEFAULT_LOG_LEVEL,
+1 LOG_FILE_TYPE,
+1)
+1from iotdb.ainode.core.util.decorator import singleton
+1
+1
+1class BaseLogger:
+1 """
+1 BaseLogger is a base class for logging, which implements daily compress, log files management, custom format, etc.
+1
+1 Args:
+1 log_file_name_prefix: the prefix of the log file name, it is used to distinguish different processes.
+1 log_dir: log directory, default is AINODE_HOME/logs.
+1
+1 Attributes:
+1 logger_format: log format
+1 logger: global logger with custom format and level
+1 console_handler: console handler for stdout
+1 _lock: process lock for logger. This is just a precaution, we currently do not have multiprocessing
+1 """
+1
+1 def __init__(self, log_file_name_prefix: str, log_dir=AINODE_LOG_DIR):
+1 self.logger_format = logging.Formatter(
+1 fmt="%(asctime)s %(levelname)s [%(process)d:%(processName)s] "
+1 "%(filename)s:%(funcName)s:%(lineno)d - %(message)s",
+1 datefmt="%Y-%m-%d %H:%M:%S",
+1 )
+1
+1 self._lock = threading.Lock()
+1
+1 self.logger = logging.getLogger(str(random.random()))
+1 self.logger.handlers.clear()
+1 self.logger.setLevel(DEFAULT_LOG_LEVEL)
+1 self.console_handler = logging.StreamHandler(sys.stdout)
+1 self.console_handler.setLevel(DEFAULT_LOG_LEVEL)
+1 self.console_handler.setFormatter(self.logger_format)
+1 self.logger.addHandler(self.console_handler)
+1
+1 # Set log file handler
+1 for i in range(len(LOG_FILE_TYPE)):
+1 file_name = log_file_name_prefix + LOG_FILE_TYPE[i] + ".log"
+1 # create log file if not exist
+1 os.makedirs(log_dir, exist_ok=True)
+1 file_path = os.path.join(log_dir, f"{file_name}")
+1 if not os.path.exists(file_path):
+1 with open(file_path, "w", encoding="utf-8") as f:
+1 f.write("")
+1 # create handler
+1 file_level = AINODE_LOG_FILE_LEVELS[i]
+1 file_handler = TimedRotatingFileHandler(
+1 filename=file_path,
+1 when="MIDNIGHT",
+1 interval=1,
+1 encoding="utf-8",
+1 )
+1 if LOG_FILE_TYPE[i] != "all":
+1 file_handler.setLevel(logging.NOTSET)
+1 file_handler.addFilter(LogLevelFilter(file_level))
+1
+1 # set renamer
+1 def universal_namer(
+1 default_name: str, internal_file_name: str = file_name
+1 ) -> str:
+1 # to avoid outer variable late binding
+1 base, ext = os.path.splitext(internal_file_name)
+1 base = base.replace("_", "-")
+1 base_log_dir = os.path.dirname(default_name)
+1 suffix = default_name.rsplit(".", 1)[-1] # e.g. 2025-08-01_13-45
+1 digits = re.sub(r"[-_]", "", suffix)
+1 return os.path.join(base_log_dir, f"{base}-{digits}{ext}.gz")
+1
+1 file_handler.namer = universal_namer
+1
+1 # set gzip
+1 def gzip_rotator(src: str, dst: str):
+1 with open(src, "rb") as f_in, gzip.open(dst, "wb") as f_out:
+1 shutil.copyfileobj(f_in, f_out)
+1 # delete the old .log file
+1 os.remove(src)
+1
+1 file_handler.rotator = gzip_rotator
+1 # other settings
+1 file_handler.setLevel(file_level)
+1 file_handler.setFormatter(self.logger_format)
+1 self.logger.addHandler(file_handler)
+1
+1 self.info(f"Logger init successfully.")
+1
+1 # interfaces
+1 def debug(self, *msg):
+1 self._write(self.logger.debug, *msg)
+1
+1 def info(self, *msg):
+1 self._write(self.logger.info, *msg)
+1
+1 def warning(self, *msg):
+1 self._write(self.logger.warning, *msg)
+1
+1 def error(self, *msg):
+1 self._write(self.logger.error, *msg)
+1
+1 def _write(self, function, *msg):
+1 with self._lock:
+1 # The stack level of caller is 3, because:
+1 # caller -> BaseLogger.info -> BaseLogger._write
+1 function(" ".join(map(str, msg)), stacklevel=3)
+1
+1
+1@singleton
+1class Logger(BaseLogger):
+1 """
+1 Logger is a singleton, just use Logger() to get it in the specified process.
+1
+1 Args:
+1 log_file_name_prefix: the prefix of the log file name, it is used to distinguish inference, training and the main process.
+1 """
+1
+1 def __init__(self, log_file_name_prefix: str = AINODE_LOG_FILE_NAME_PREFIX):
+1 super().__init__(log_file_name_prefix=log_file_name_prefix)
+1
+1
+1class LogLevelFilter(logging.Filter):
+1
+1 def __init__(self, level: int):
+1 super().__init__()
+1 self.level = level
+1
+1 def filter(self, record: logging.LogRecord) -> bool:
+1 return record.levelno == self.level
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
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 47fce599ab2..c024efb93dc 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py
@@ -1,46 +1,47 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import psutil
-
-from iotdb.thrift.ainode.ttypes import TAIHeartbeatReq, TAIHeartbeatResp
-from iotdb.thrift.common.ttypes import TLoadSample
-
-
-class ClusterManager:
- @staticmethod
- def get_heart_beat(req: TAIHeartbeatReq) -> TAIHeartbeatResp:
- if req.needSamplingLoad:
- cpu_percent = psutil.cpu_percent(interval=1)
- memory_percent = psutil.virtual_memory().percent
- disk_usage = psutil.disk_usage("/")
- disk_free = disk_usage.free
- load_sample = TLoadSample(
- cpuUsageRate=cpu_percent,
- memoryUsageRate=memory_percent,
- diskUsageRate=disk_usage.percent,
- freeDiskSpace=disk_free / 1024 / 1024 / 1024,
- )
- return TAIHeartbeatResp(
- heartbeatTimestamp=req.heartbeatTimestamp,
- status="Running",
- loadSample=load_sample,
- )
- else:
- return TAIHeartbeatResp(
- heartbeatTimestamp=req.heartbeatTimestamp, status="Running"
- )
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1import psutil
+1
+1from iotdb.thrift.ainode.ttypes import TAIHeartbeatReq, TAIHeartbeatResp
+1from iotdb.thrift.common.ttypes import TLoadSample
+1
+1
+1class ClusterManager:
+1 @staticmethod
+1 def get_heart_beat(req: TAIHeartbeatReq) -> TAIHeartbeatResp:
+1 if req.needSamplingLoad:
+1 cpu_percent = psutil.cpu_percent(interval=1)
+1 memory_percent = psutil.virtual_memory().percent
+1 disk_usage = psutil.disk_usage("/")
+1 disk_free = disk_usage.free
+1 load_sample = TLoadSample(
+1 cpuUsageRate=cpu_percent,
+1 memoryUsageRate=memory_percent,
+1 diskUsageRate=disk_usage.percent,
+1 freeDiskSpace=disk_free / 1024 / 1024 / 1024,
+1 )
+1 return TAIHeartbeatResp(
+1 heartbeatTimestamp=req.heartbeatTimestamp,
+1 status="Running",
+1 loadSample=load_sample,
+1 )
+1 else:
+1 return TAIHeartbeatResp(
+1 heartbeatTimestamp=req.heartbeatTimestamp, status="Running"
+1 )
+1
\ No newline at end of file
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 6f14036c8dc..d90a5204936 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
@@ -1,372 +1,373 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import threading
-import time
-from abc import ABC, abstractmethod
-from typing import Dict
-
-import pandas as pd
-import torch
-import torch.multiprocessing as mp
-
-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.constant import TSStatusCode
-from iotdb.ainode.core.exception import (
- InferenceModelInternalError,
- InvalidWindowArgumentError,
- NumericalRangeException,
- runtime_error_extractor,
-)
-from iotdb.ainode.core.inference.inference_request import (
- InferenceRequest,
- InferenceRequestProxy,
-)
-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.timerxl_inference_pipeline import (
- TimerXLInferencePipeline,
-)
-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.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 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 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()
-
-
-class InferenceStrategy(ABC):
- def __init__(self, model):
- self.model = model
-
- @abstractmethod
- def infer(self, full_data, **kwargs):
- pass
-
-
-# [IoTDB] full data deserialized from iotdb is composed of [timestampList, valueList, length],
-# we only get valueList currently.
-class TimerXLStrategy(InferenceStrategy):
- def infer(self, full_data, predict_length=96, **_):
- data = full_data[1][0]
- if data.dtype.byteorder not in ("=", "|"):
- np_data = data.byteswap()
- data = np_data.view(np_data.dtype.newbyteorder())
- seqs = torch.tensor(data).unsqueeze(0).float()
- # TODO: unify model inference input
- output = self.model.generate(seqs, max_new_tokens=predict_length, revin=True)
- df = pd.DataFrame(output[0])
- return convert_to_binary(df)
-
-
-class SundialStrategy(InferenceStrategy):
- def infer(self, full_data, predict_length=96, **_):
- data = full_data[1][0]
- if data.dtype.byteorder not in ("=", "|"):
- np_data = data.byteswap()
- data = np_data.view(np_data.dtype.newbyteorder())
- seqs = torch.tensor(data).unsqueeze(0).float()
- # TODO: unify model inference input
- output = self.model.generate(
- seqs, max_new_tokens=predict_length, num_samples=10, revin=True
- )
- df = pd.DataFrame(output[0].mean(dim=0))
- return convert_to_binary(df)
-
-
-class BuiltInStrategy(InferenceStrategy):
- def infer(self, full_data, **_):
- data = pd.DataFrame(full_data[1]).T
- output = self.model.inference(data)
- df = pd.DataFrame(output)
- return convert_to_binary(df)
-
-
-class RegisteredStrategy(InferenceStrategy):
- def infer(self, full_data, window_interval=None, window_step=None, **_):
- _, dataset, _, length = full_data
- if window_interval is None or window_step is None:
- window_interval = length
- window_step = float("inf")
-
- if window_interval <= 0 or window_step <= 0 or window_interval > length:
- raise InvalidWindowArgumentError(window_interval, window_step, length)
-
- data = torch.tensor(dataset, dtype=torch.float32).unsqueeze(0).permute(0, 2, 1)
-
- times = int((length - window_interval) // window_step + 1)
- results = []
- try:
- for i in range(times):
- start = 0 if window_step == float("inf") else i * window_step
- end = start + window_interval
- window = data[:, start:end, :]
- out = self.model(window)
- df = pd.DataFrame(out.squeeze(0).detach().numpy())
- results.append(df)
- except Exception as e:
- msg = runtime_error_extractor(str(e)) or str(e)
- raise InferenceModelInternalError(msg)
-
- # concatenate or return first window for forecast
- return [convert_to_binary(df) for df in results]
-
-
-class InferenceManager:
- 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()
-
- self._stop_event = mp.Event()
- self._result_handler_thread = threading.Thread(
- target=self._handle_results, daemon=True
- )
- self._result_handler_thread.start()
- 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
- )
- 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.',
- )
-
- 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:
- 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():
- if self._result_queue.empty():
- time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
- continue
- infer_req: InferenceRequest = self._result_queue.get()
- with self._result_wrapper_lock:
- self._result_wrapper_map[infer_req.req_id].set_result(
- 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):
- return BuiltInStrategy(model)
- return RegisteredStrategy(model)
-
- def _run(
- self,
- req,
- data_getter,
- deserializer,
- extract_attrs,
- resp_cls,
- single_output: bool,
- ):
- model_id = req.modelId
- try:
- raw = data_getter(req)
- full_data = deserializer(raw)
- inference_attrs = extract_attrs(req)
-
- predict_length = int(inference_attrs.pop("predict_length", 96))
- if (
- predict_length
- > AINodeDescriptor().get_config().get_ain_inference_max_predict_length()
- ):
- raise NumericalRangeException(
- "output_length",
- 1,
- AINodeDescriptor()
- .get_config()
- .get_ain_inference_max_predict_length(),
- predict_length,
- )
-
- 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 ("=", "|"):
- np_data = data.byteswap()
- 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":
- inference_pipeline = TimerSundialInferencePipeline(SundialConfig())
- elif model_id == "timer_xl":
- 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,
- )
- outputs = self._process_request(infer_req)
- outputs = convert_to_binary(pd.DataFrame(outputs[0]))
- else:
- # load model
- accel = str(inference_attrs.get("acceleration", "")).lower() == "true"
- model = self._model_manager.load_model(model_id, inference_attrs, accel)
- # inference by strategy
- strategy = self._get_strategy(model_id, model)
- outputs = strategy.infer(
- full_data, predict_length=predict_length, **inference_attrs
- )
-
- # construct response
- status = get_status(TSStatusCode.SUCCESS_STATUS)
-
- if isinstance(outputs, list):
- return resp_cls(status, outputs[0] if single_output else outputs)
- return resp_cls(status, outputs if single_output else [outputs])
-
- except Exception as e:
- logger.error(e)
- status = get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
- empty = b"" if single_output else []
- return resp_cls(status, empty)
-
- def forecast(self, req: TForecastReq):
- return self._run(
- req,
- data_getter=lambda r: r.inputData,
- deserializer=deserialize,
- extract_attrs=lambda r: {
- "predict_length": r.outputLength,
- **(r.options or {}),
- },
- resp_cls=TForecastResp,
- single_output=True,
- )
-
- def inference(self, req: TInferenceReq):
- return self._run(
- req,
- data_getter=lambda r: r.dataset,
- deserializer=deserialize,
- extract_attrs=lambda r: {
- "window_interval": getattr(r.windowParams, "windowInterval", None),
- "window_step": getattr(r.windowParams, "windowStep", None),
- **(r.inferenceAttributes or {}),
- },
- resp_cls=TInferenceResp,
- single_output=False,
- )
-
- def shutdown(self):
- self._stop_event.set()
- self._pool_controller.shutdown()
- while not self._result_queue.empty():
- self._result_queue.get_nowait()
- self._result_queue.close()
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import threading
+1import time
+1from abc import ABC, abstractmethod
+1from typing import Dict
+1
+1import pandas as pd
+1import torch
+1import torch.multiprocessing as mp
+1
+1from iotdb.ainode.core.config import AINodeDescriptor
+1from iotdb.ainode.core.constant import TSStatusCode
+1from iotdb.ainode.core.exception import (
+1 InferenceModelInternalError,
+1 InvalidWindowArgumentError,
+1 NumericalRangeException,
+1 runtime_error_extractor,
+1)
+1from iotdb.ainode.core.inference.inference_request import (
+1 InferenceRequest,
+1 InferenceRequestProxy,
+1)
+1from iotdb.ainode.core.inference.pool_controller import PoolController
+1from iotdb.ainode.core.inference.strategy.timer_sundial_inference_pipeline import (
+1 TimerSundialInferencePipeline,
+1)
+1from iotdb.ainode.core.inference.strategy.timerxl_inference_pipeline import (
+1 TimerXLInferencePipeline,
+1)
+1from iotdb.ainode.core.inference.utils import generate_req_id
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.manager.model_manager import ModelManager
+1from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
+1from iotdb.ainode.core.model.sundial.modeling_sundial import SundialForPrediction
+1from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
+1from iotdb.ainode.core.model.timerxl.modeling_timer import TimerForPrediction
+1from iotdb.ainode.core.rpc.status import get_status
+1from iotdb.ainode.core.util.gpu_mapping import get_available_devices
+1from iotdb.ainode.core.util.serde import convert_to_binary
+1from iotdb.thrift.ainode.ttypes import (
+1 TForecastReq,
+1 TForecastResp,
+1 TInferenceReq,
+1 TInferenceResp,
+1 TLoadModelReq,
+1 TShowLoadedModelsReq,
+1 TShowLoadedModelsResp,
+1 TUnloadModelReq,
+1)
+1from iotdb.thrift.common.ttypes import TSStatus
+1from iotdb.tsfile.utils.tsblock_serde import deserialize
+1
+1logger = Logger()
+1
+1
+1class InferenceStrategy(ABC):
+1 def __init__(self, model):
+1 self.model = model
+1
+1 @abstractmethod
+1 def infer(self, full_data, **kwargs):
+1 pass
+1
+1
+1# [IoTDB] full data deserialized from iotdb is composed of [timestampList, valueList, length],
+1# we only get valueList currently.
+1class TimerXLStrategy(InferenceStrategy):
+1 def infer(self, full_data, predict_length=96, **_):
+1 data = full_data[1][0]
+1 if data.dtype.byteorder not in ("=", "|"):
+1 np_data = data.byteswap()
+1 data = np_data.view(np_data.dtype.newbyteorder())
+1 seqs = torch.tensor(data).unsqueeze(0).float()
+1 # TODO: unify model inference input
+1 output = self.model.generate(seqs, max_new_tokens=predict_length, revin=True)
+1 df = pd.DataFrame(output[0])
+1 return convert_to_binary(df)
+1
+1
+1class SundialStrategy(InferenceStrategy):
+1 def infer(self, full_data, predict_length=96, **_):
+1 data = full_data[1][0]
+1 if data.dtype.byteorder not in ("=", "|"):
+1 np_data = data.byteswap()
+1 data = np_data.view(np_data.dtype.newbyteorder())
+1 seqs = torch.tensor(data).unsqueeze(0).float()
+1 # TODO: unify model inference input
+1 output = self.model.generate(
+1 seqs, max_new_tokens=predict_length, num_samples=10, revin=True
+1 )
+1 df = pd.DataFrame(output[0].mean(dim=0))
+1 return convert_to_binary(df)
+1
+1
+1class BuiltInStrategy(InferenceStrategy):
+1 def infer(self, full_data, **_):
+1 data = pd.DataFrame(full_data[1]).T
+1 output = self.model.inference(data)
+1 df = pd.DataFrame(output)
+1 return convert_to_binary(df)
+1
+1
+1class RegisteredStrategy(InferenceStrategy):
+1 def infer(self, full_data, window_interval=None, window_step=None, **_):
+1 _, dataset, _, length = full_data
+1 if window_interval is None or window_step is None:
+1 window_interval = length
+1 window_step = float("inf")
+1
+1 if window_interval <= 0 or window_step <= 0 or window_interval > length:
+1 raise InvalidWindowArgumentError(window_interval, window_step, length)
+1
+1 data = torch.tensor(dataset, dtype=torch.float32).unsqueeze(0).permute(0, 2, 1)
+1
+1 times = int((length - window_interval) // window_step + 1)
+1 results = []
+1 try:
+1 for i in range(times):
+1 start = 0 if window_step == float("inf") else i * window_step
+1 end = start + window_interval
+1 window = data[:, start:end, :]
+1 out = self.model(window)
+1 df = pd.DataFrame(out.squeeze(0).detach().numpy())
+1 results.append(df)
+1 except Exception as e:
+1 msg = runtime_error_extractor(str(e)) or str(e)
+1 raise InferenceModelInternalError(msg)
+1
+1 # concatenate or return first window for forecast
+1 return [convert_to_binary(df) for df in results]
+1
+1
+1class InferenceManager:
+1 WAITING_INTERVAL_IN_MS = (
+1 AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms()
+1 ) # How often to check for requests in the result queue
+1
+1 def __init__(self):
+1 self._model_manager = ModelManager()
+1 self._model_mem_usage_map: Dict[str, int] = (
+1 {}
+1 ) # store model memory usage for each model
+1 self._result_queue = mp.Queue()
+1 self._result_wrapper_map = {}
+1 self._result_wrapper_lock = threading.RLock()
+1
+1 self._stop_event = mp.Event()
+1 self._result_handler_thread = threading.Thread(
+1 target=self._handle_results, daemon=True
+1 )
+1 self._result_handler_thread.start()
+1 self._pool_controller = PoolController(self._result_queue)
+1
+1 def load_model(self, req: TLoadModelReq) -> TSStatus:
+1 devices_to_be_processed = []
+1 devices_not_to_be_processed = []
+1 for device_id in req.deviceIdList:
+1 if self._pool_controller.has_request_pools(
+1 model_id=req.existingModelId, device_id=device_id
+1 ):
+1 devices_not_to_be_processed.append(device_id)
+1 else:
+1 devices_to_be_processed.append(device_id)
+1 if len(devices_to_be_processed) > 0:
+1 self._pool_controller.load_model(
+1 model_id=req.existingModelId, device_id_list=devices_to_be_processed
+1 )
+1 logger.info(
+1 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."
+1 )
+1 return TSStatus(
+1 code=TSStatusCode.SUCCESS_STATUS.value,
+1 message='Successfully submitted load model task, please use "SHOW LOADED MODELS" to check progress.',
+1 )
+1
+1 def unload_model(self, req: TUnloadModelReq) -> TSStatus:
+1 devices_to_be_processed = []
+1 devices_not_to_be_processed = []
+1 for device_id in req.deviceIdList:
+1 if self._pool_controller.has_request_pools(
+1 model_id=req.modelId, device_id=device_id
+1 ):
+1 devices_to_be_processed.append(device_id)
+1 else:
+1 devices_not_to_be_processed.append(device_id)
+1 if len(devices_to_be_processed) > 0:
+1 self._pool_controller.unload_model(
+1 model_id=req.modelId, device_id_list=req.deviceIdList
+1 )
+1 logger.info(
+1 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."
+1 )
+1 return TSStatus(
+1 code=TSStatusCode.SUCCESS_STATUS.value,
+1 message='Successfully submitted unload model task, please use "SHOW LOADED MODELS" to check progress.',
+1 )
+1
+1 def show_loaded_models(self, req: TShowLoadedModelsReq) -> TShowLoadedModelsResp:
+1 return TShowLoadedModelsResp(
+1 status=get_status(TSStatusCode.SUCCESS_STATUS),
+1 deviceLoadedModelsMap=self._pool_controller.show_loaded_models(
+1 req.deviceIdList
+1 if len(req.deviceIdList) > 0
+1 else get_available_devices()
+1 ),
+1 )
+1
+1 def _handle_results(self):
+1 while not self._stop_event.is_set():
+1 if self._result_queue.empty():
+1 time.sleep(self.WAITING_INTERVAL_IN_MS / 1000)
+1 continue
+1 infer_req: InferenceRequest = self._result_queue.get()
+1 with self._result_wrapper_lock:
+1 self._result_wrapper_map[infer_req.req_id].set_result(
+1 infer_req.get_final_output()
+1 )
+1
+1 def _process_request(self, req):
+1 req_id = req.req_id
+1 infer_proxy = InferenceRequestProxy(req_id)
+1 with self._result_wrapper_lock:
+1 self._result_wrapper_map[req_id] = infer_proxy
+1 try:
+1 # dispatch request to the pool
+1 self._pool_controller.add_request(req, infer_proxy)
+1 outputs = infer_proxy.wait_for_result()
+1 return outputs
+1 except Exception as e:
+1 logger.error(e)
+1 raise InferenceModelInternalError(str(e))
+1 finally:
+1 with self._result_wrapper_lock:
+1 del self._result_wrapper_map[req_id]
+1
+1 def _get_strategy(self, model_id, model):
+1 if isinstance(model, TimerForPrediction):
+1 return TimerXLStrategy(model)
+1 if isinstance(model, SundialForPrediction):
+1 return SundialStrategy(model)
+1 if self._model_manager.model_storage.is_built_in_or_fine_tuned(model_id):
+1 return BuiltInStrategy(model)
+1 return RegisteredStrategy(model)
+1
+1 def _run(
+1 self,
+1 req,
+1 data_getter,
+1 deserializer,
+1 extract_attrs,
+1 resp_cls,
+1 single_output: bool,
+1 ):
+1 model_id = req.modelId
+1 try:
+1 raw = data_getter(req)
+1 full_data = deserializer(raw)
+1 inference_attrs = extract_attrs(req)
+1
+1 predict_length = int(inference_attrs.pop("predict_length", 96))
+1 if (
+1 predict_length
+1 > AINodeDescriptor().get_config().get_ain_inference_max_predict_length()
+1 ):
+1 raise NumericalRangeException(
+1 "output_length",
+1 1,
+1 AINodeDescriptor()
+1 .get_config()
+1 .get_ain_inference_max_predict_length(),
+1 predict_length,
+1 )
+1
+1 if self._pool_controller.has_request_pools(model_id):
+1 # use request pool to accelerate inference when the model instance is already loaded.
+1 # TODO: TSBlock -> Tensor codes should be unified
+1 data = full_data[1][0]
+1 if data.dtype.byteorder not in ("=", "|"):
+1 np_data = data.byteswap()
+1 data = np_data.view(np_data.dtype.newbyteorder())
+1 # the inputs should be on CPU before passing to the inference request
+1 inputs = torch.tensor(data).unsqueeze(0).float().to("cpu")
+1 if model_id == "sundial":
+1 inference_pipeline = TimerSundialInferencePipeline(SundialConfig())
+1 elif model_id == "timer_xl":
+1 inference_pipeline = TimerXLInferencePipeline(TimerConfig())
+1 else:
+1 raise InferenceModelInternalError(
+1 f"Unsupported model_id: {model_id}"
+1 )
+1 infer_req = InferenceRequest(
+1 req_id=generate_req_id(),
+1 model_id=model_id,
+1 inputs=inputs,
+1 inference_pipeline=inference_pipeline,
+1 max_new_tokens=predict_length,
+1 )
+1 outputs = self._process_request(infer_req)
+1 outputs = convert_to_binary(pd.DataFrame(outputs[0]))
+1 else:
+1 # load model
+1 accel = str(inference_attrs.get("acceleration", "")).lower() == "true"
+1 model = self._model_manager.load_model(model_id, inference_attrs, accel)
+1 # inference by strategy
+1 strategy = self._get_strategy(model_id, model)
+1 outputs = strategy.infer(
+1 full_data, predict_length=predict_length, **inference_attrs
+1 )
+1
+1 # construct response
+1 status = get_status(TSStatusCode.SUCCESS_STATUS)
+1
+1 if isinstance(outputs, list):
+1 return resp_cls(status, outputs[0] if single_output else outputs)
+1 return resp_cls(status, outputs if single_output else [outputs])
+1
+1 except Exception as e:
+1 logger.error(e)
+1 status = get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+1 empty = b"" if single_output else []
+1 return resp_cls(status, empty)
+1
+1 def forecast(self, req: TForecastReq):
+1 return self._run(
+1 req,
+1 data_getter=lambda r: r.inputData,
+1 deserializer=deserialize,
+1 extract_attrs=lambda r: {
+1 "predict_length": r.outputLength,
+1 **(r.options or {}),
+1 },
+1 resp_cls=TForecastResp,
+1 single_output=True,
+1 )
+1
+1 def inference(self, req: TInferenceReq):
+1 return self._run(
+1 req,
+1 data_getter=lambda r: r.dataset,
+1 deserializer=deserialize,
+1 extract_attrs=lambda r: {
+1 "window_interval": getattr(r.windowParams, "windowInterval", None),
+1 "window_step": getattr(r.windowParams, "windowStep", None),
+1 **(r.inferenceAttributes or {}),
+1 },
+1 resp_cls=TInferenceResp,
+1 single_output=False,
+1 )
+1
+1 def shutdown(self):
+1 self._stop_event.set()
+1 self._pool_controller.shutdown()
+1 while not self._result_queue.empty():
+1 self._result_queue.get_nowait()
+1 self._result_queue.close()
+1
\ No newline at end of file
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 d6b0f17d5c0..6b4bb4fea2d 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
@@ -1,166 +1,167 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-from typing import Callable, Dict
-
-from torch import nn
-from yaml import YAMLError
-
-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 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 iotdb.thrift.common.ttypes import TSStatus
-
-logger = Logger()
-
-
-@singleton
-class ModelManager:
- def __init__(self):
- self.model_storage = ModelStorage()
-
- def register_model(self, req: TRegisterModelReq) -> TRegisterModelResp:
- logger.info(f"register model {req.modelId} from {req.uri}")
- try:
- configs, attributes = self.model_storage.register_model(
- req.modelId, req.uri
- )
- return TRegisterModelResp(
- get_status(TSStatusCode.SUCCESS_STATUS), configs, attributes
- )
- except InvalidUriError as e:
- logger.warning(e)
- return TRegisterModelResp(
- get_status(TSStatusCode.INVALID_URI_ERROR, e.message)
- )
- except BadConfigValueError as e:
- logger.warning(e)
- return TRegisterModelResp(
- get_status(TSStatusCode.INVALID_INFERENCE_CONFIG, e.message)
- )
- except YAMLError as e:
- logger.warning(e)
- if hasattr(e, "problem_mark"):
- mark = e.problem_mark
- return TRegisterModelResp(
- get_status(
- TSStatusCode.INVALID_INFERENCE_CONFIG,
- f"An error occurred while parsing the yaml file, "
- f"at line {mark.line + 1} column {mark.column + 1}.",
- )
- )
- return TRegisterModelResp(
- get_status(
- TSStatusCode.INVALID_INFERENCE_CONFIG,
- f"An error occurred while parsing the yaml file",
- )
- )
- except Exception as e:
- logger.warning(e)
- return TRegisterModelResp(get_status(TSStatusCode.AINODE_INTERNAL_ERROR))
-
- def delete_model(self, req: TDeleteModelReq) -> TSStatus:
- logger.info(f"delete model {req.modelId}")
- try:
- self.model_storage.delete_model(req.modelId)
- return get_status(TSStatusCode.SUCCESS_STATUS)
- except Exception as e:
- logger.warning(e)
- return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
-
- def load_model(
- self, model_id: str, inference_attrs: Dict[str, str], acceleration: bool = False
- ) -> Callable:
- """
- Load the model with the given model_id.
- """
- logger.info(f"Load model {model_id}")
- try:
- model = self.model_storage.load_model(
- model_id, inference_attrs, acceleration
- )
- logger.info(f"Model {model_id} loaded")
- return model
- except Exception as e:
- logger.error(f"Failed to load model {model_id}: {e}")
- raise
-
- def save_model(self, model_id: str, model: nn.Module) -> TSStatus:
- """
- Save the model using save_pretrained
- """
- logger.info(f"Saving model {model_id}")
- try:
- self.model_storage.save_model(model_id, model)
- logger.info(f"Saving model {model_id} successfully")
- return get_status(
- TSStatusCode.SUCCESS_STATUS, f"Model {model_id} saved successfully"
- )
- except Exception as e:
- logger.error(f"Save model failed: {e}")
- return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
-
- def get_ckpt_path(self, model_id: str) -> str:
- """
- Get the checkpoint path for a given model ID.
-
- Args:
- model_id (str): The ID of the model.
-
- Returns:
- str: The path to the checkpoint file for the model.
- """
- return self.model_storage.get_ckpt_path(model_id)
-
- def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
- return self.model_storage.show_models(req)
-
- def register_built_in_model(self, model_info: ModelInfo):
- self.model_storage.register_built_in_model(model_info)
-
- def update_model_state(self, model_id: str, state: ModelStates):
- self.model_storage.update_model_state(model_id, state)
-
- def get_built_in_model_type(self, model_id: str) -> BuiltInModelType:
- """
- Get the type of the model with the given model_id.
- """
- return self.model_storage.get_built_in_model_type(model_id)
-
- def is_built_in_or_fine_tuned(self, model_id: str) -> bool:
- """
- Check if the model_id corresponds to a built-in or fine-tuned model.
-
- Args:
- model_id (str): The ID of the model.
-
- Returns:
- bool: True if the model is built-in or fine_tuned, False otherwise.
- """
- return self.model_storage.is_built_in_or_fine_tuned(model_id)
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1from typing import Callable, Dict
+1
+1from torch import nn
+1from yaml import YAMLError
+1
+1from iotdb.ainode.core.constant import TSStatusCode
+1from iotdb.ainode.core.exception import BadConfigValueError, InvalidUriError
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.model.model_enums import BuiltInModelType, ModelStates
+1from iotdb.ainode.core.model.model_info import ModelInfo
+1from iotdb.ainode.core.model.model_storage import ModelStorage
+1from iotdb.ainode.core.rpc.status import get_status
+1from iotdb.ainode.core.util.decorator import singleton
+1from iotdb.thrift.ainode.ttypes import (
+1 TDeleteModelReq,
+1 TRegisterModelReq,
+1 TRegisterModelResp,
+1 TShowModelsReq,
+1 TShowModelsResp,
+1)
+1from iotdb.thrift.common.ttypes import TSStatus
+1
+1logger = Logger()
+1
+1
+1@singleton
+1class ModelManager:
+1 def __init__(self):
+1 self.model_storage = ModelStorage()
+1
+1 def register_model(self, req: TRegisterModelReq) -> TRegisterModelResp:
+1 logger.info(f"register model {req.modelId} from {req.uri}")
+1 try:
+1 configs, attributes = self.model_storage.register_model(
+1 req.modelId, req.uri
+1 )
+1 return TRegisterModelResp(
+1 get_status(TSStatusCode.SUCCESS_STATUS), configs, attributes
+1 )
+1 except InvalidUriError as e:
+1 logger.warning(e)
+1 return TRegisterModelResp(
+1 get_status(TSStatusCode.INVALID_URI_ERROR, e.message)
+1 )
+1 except BadConfigValueError as e:
+1 logger.warning(e)
+1 return TRegisterModelResp(
+1 get_status(TSStatusCode.INVALID_INFERENCE_CONFIG, e.message)
+1 )
+1 except YAMLError as e:
+1 logger.warning(e)
+1 if hasattr(e, "problem_mark"):
+1 mark = e.problem_mark
+1 return TRegisterModelResp(
+1 get_status(
+1 TSStatusCode.INVALID_INFERENCE_CONFIG,
+1 f"An error occurred while parsing the yaml file, "
+1 f"at line {mark.line + 1} column {mark.column + 1}.",
+1 )
+1 )
+1 return TRegisterModelResp(
+1 get_status(
+1 TSStatusCode.INVALID_INFERENCE_CONFIG,
+1 f"An error occurred while parsing the yaml file",
+1 )
+1 )
+1 except Exception as e:
+1 logger.warning(e)
+1 return TRegisterModelResp(get_status(TSStatusCode.AINODE_INTERNAL_ERROR))
+1
+1 def delete_model(self, req: TDeleteModelReq) -> TSStatus:
+1 logger.info(f"delete model {req.modelId}")
+1 try:
+1 self.model_storage.delete_model(req.modelId)
+1 return get_status(TSStatusCode.SUCCESS_STATUS)
+1 except Exception as e:
+1 logger.warning(e)
+1 return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+1
+1 def load_model(
+1 self, model_id: str, inference_attrs: Dict[str, str], acceleration: bool = False
+1 ) -> Callable:
+1 """
+1 Load the model with the given model_id.
+1 """
+1 logger.info(f"Load model {model_id}")
+1 try:
+1 model = self.model_storage.load_model(
+1 model_id, inference_attrs, acceleration
+1 )
+1 logger.info(f"Model {model_id} loaded")
+1 return model
+1 except Exception as e:
+1 logger.error(f"Failed to load model {model_id}: {e}")
+1 raise
+1
+1 def save_model(self, model_id: str, model: nn.Module) -> TSStatus:
+1 """
+1 Save the model using save_pretrained
+1 """
+1 logger.info(f"Saving model {model_id}")
+1 try:
+1 self.model_storage.save_model(model_id, model)
+1 logger.info(f"Saving model {model_id} successfully")
+1 return get_status(
+1 TSStatusCode.SUCCESS_STATUS, f"Model {model_id} saved successfully"
+1 )
+1 except Exception as e:
+1 logger.error(f"Save model failed: {e}")
+1 return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+1
+1 def get_ckpt_path(self, model_id: str) -> str:
+1 """
+1 Get the checkpoint path for a given model ID.
+1
+1 Args:
+1 model_id (str): The ID of the model.
+1
+1 Returns:
+1 str: The path to the checkpoint file for the model.
+1 """
+1 return self.model_storage.get_ckpt_path(model_id)
+1
+1 def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
+1 return self.model_storage.show_models(req)
+1
+1 def register_built_in_model(self, model_info: ModelInfo):
+1 self.model_storage.register_built_in_model(model_info)
+1
+1 def update_model_state(self, model_id: str, state: ModelStates):
+1 self.model_storage.update_model_state(model_id, state)
+1
+1 def get_built_in_model_type(self, model_id: str) -> BuiltInModelType:
+1 """
+1 Get the type of the model with the given model_id.
+1 """
+1 return self.model_storage.get_built_in_model_type(model_id)
+1
+1 def is_built_in_or_fine_tuned(self, model_id: str) -> bool:
+1 """
+1 Check if the model_id corresponds to a built-in or fine-tuned model.
+1
+1 Args:
+1 model_id (str): The ID of the model.
+1
+1 Returns:
+1 bool: True if the model is built-in or fine_tuned, False otherwise.
+1 """
+1 return self.model_storage.is_built_in_or_fine_tuned(model_id)
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
index 0264e27331a..0ed1bf1f140 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
@@ -1,108 +1,109 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import gc
-
-import psutil
-import torch
-
-from iotdb.ainode.core.config import AINodeDescriptor
-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()
-
-MODEL_MEM_USAGE_MAP = (
- AINodeDescriptor().get_config().get_ain_inference_model_mem_usage_map()
-) # the memory usage of each model in bytes
-INFERENCE_MEMORY_USAGE_RATIO = (
- AINodeDescriptor().get_config().get_ain_inference_memory_usage_ratio()
-) # the device space allocated for inference
-INFERENCE_EXTRA_MEMORY_RATIO = (
- AINodeDescriptor().get_config().get_ain_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:
- # TODO: support CPU in the future
- # TODO: we can estimate the memory usage by running a dummy inference
- torch.cuda.empty_cache()
- torch.cuda.synchronize(device)
- start = torch.cuda.memory_reserved(device)
-
- model = ModelManager().load_model(model_id, {}).to(device)
- torch.cuda.synchronize(device)
- end = torch.cuda.memory_reserved(device)
- usage = end - start
-
- # delete model to free memory
- del model
- torch.cuda.empty_cache()
- gc.collect()
-
- # add inference factor and cuda context overhead
- overhead = 500 * 1024**2 # 500 MiB
- final = int(max(usage, 1) * INFERENCE_EXTRA_MEMORY_RATIO + overhead)
- return final
-
-
-def evaluate_system_resources(device: torch.device) -> dict:
- if torch.cuda.is_available():
- free_mem, total_mem = torch.cuda.mem_get_info()
- logger.info(
- f"[Inference][Device-{device}] CUDA device memory: free={free_mem/1024**2:.2f} MB, total={total_mem/1024**2:.2f} MB"
- )
- return {"device": "cuda", "free_mem": free_mem, "total_mem": total_mem}
- else:
- free_mem = psutil.virtual_memory().available
- total_mem = psutil.virtual_memory().total
- logger.info(
- f"[Inference][Device-{device}] CPU memory: free={free_mem/1024**2:.2f} MB, total={total_mem/1024**2:.2f} MB"
- )
- return {"device": "cpu", "free_mem": free_mem, "total_mem": total_mem}
-
-
-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 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)
- free_mem = system_res["free_mem"]
-
- 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(
- f"[Inference][Device-{device}] Not enough memory to run model {model_id}. free={free_mem/1024**2:.2f} MB, need>={mem_usage/1024**2:.2f} MB"
- )
- return 0
-
- logger.info(
- f"[Inference][Device-{device}] "
- f"model={model_id}, mem_usage={mem_usage/1024**2:.2f} MB, "
- f"pool_num={size}"
- )
- return size
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
+1import gc
+1
+1import psutil
+1import torch
+1
+1from iotdb.ainode.core.config import AINodeDescriptor
+1from iotdb.ainode.core.exception import ModelNotExistError
+1from iotdb.ainode.core.log import Logger
+1from iotdb.ainode.core.manager.model_manager import ModelManager
+1from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP
+1
+1logger = Logger()
+1
+1MODEL_MEM_USAGE_MAP = (
+1 AINodeDescriptor().get_config().get_ain_inference_model_mem_usage_map()
+1) # the memory usage of each model in bytes
+1INFERENCE_MEMORY_USAGE_RATIO = (
+1 AINodeDescriptor().get_config().get_ain_inference_memory_usage_ratio()
+1) # the device space allocated for inference
+1INFERENCE_EXTRA_MEMORY_RATIO = (
+1 AINodeDescriptor().get_config().get_ain_inference_extra_memory_ratio()
+1) # the overhead ratio for inference, used to estimate the pool size
+1
+1
+1def measure_model_memory(device: torch.device, model_id: str) -> int:
+1 # TODO: support CPU in the future
+1 # TODO: we can estimate the memory usage by running a dummy inference
+1 torch.cuda.empty_cache()
+1 torch.cuda.synchronize(device)
+1 start = torch.cuda.memory_reserved(device)
+1
+1 model = ModelManager().load_model(model_id, {}).to(device)
+1 torch.cuda.synchronize(device)
+1 end = torch.cuda.memory_reserved(device)
+1 usage = end - start
+1
+1 # delete model to free memory
+1 del model
+1 torch.cuda.empty_cache()
+1 gc.collect()
+1
+1 # add inference factor and cuda context overhead
+1 overhead = 500 * 1024**2 # 500 MiB
+1 final = int(max(usage, 1) * INFERENCE_EXTRA_MEMORY_RATIO + overhead)
+1 return final
+1
+1
+1def evaluate_system_resources(device: torch.device) -> dict:
+1 if torch.cuda.is_available():
+1 free_mem, total_mem = torch.cuda.mem_get_info()
+1 logger.info(
+1 f"[Inference][Device-{device}] CUDA device memory: free={free_mem/1024**2:.2f} MB, total={total_mem/1024**2:.2f} MB"
+1 )
+1 return {"device": "cuda", "free_mem": free_mem, "total_mem": total_mem}
+1 else:
+1 free_mem = psutil.virtual_memory().available
+1 total_mem = psutil.virtual_memory().total
+1 logger.info(
+1 f"[Inference][Device-{device}] CPU memory: free={free_mem/1024**2:.2f} MB, total={total_mem/1024**2:.2f} MB"
+1 )
+1 return {"device": "cpu", "free_mem": free_mem, "total_mem": total_mem}
+1
+1
+1def estimate_pool_size(device: torch.device, model_id: str) -> int:
+1 model_info = BUILT_IN_LTSM_MAP.get(model_id, None)
+1 if model_info is None or model_info.model_type not in MODEL_MEM_USAGE_MAP:
+1 logger.error(
+1 f"[Inference] Cannot estimate inference pool size on device: {device}, because model: {model_id} is not supported."
+1 )
+1 raise ModelNotExistError(model_id)
+1
+1 system_res = evaluate_system_resources(device)
+1 free_mem = system_res["free_mem"]
+1
+1 mem_usage = (
+1 MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO
+1 )
+1 size = int((free_mem * INFERENCE_MEMORY_USAGE_RATIO) // mem_usage)
+1 if size <= 0:
+1 logger.error(
+1 f"[Inference][Device-{device}] Not enough memory to run model {model_id}. free={free_mem/1024**2:.2f} MB, need>={mem_usage/1024**2:.2f} MB"
+1 )
+1 return 0
+1
+1 logger.info(
+1 f"[Inference][Device-{device}] "
+1 f"model={model_id}, mem_usage={mem_usage/1024**2:.2f} MB, "
+1 f"pool_num={size}"
+1 )
+1 return size
+1
\ No newline at end of file
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py
index 2a1e720805f..3ea5029913c 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py
@@ -1,17 +1,18 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
+1# Licensed to the Apache Software Foundation (ASF) under one
+1# or more contributor license agreements. See the NOTICE file
+1# distributed with this work for additional information
+1# regarding copyright ownership. The ASF licenses this file
+1# to you under the Apache License, Version 2.0 (the
+1# "License"); you may not use this file except in compliance
+1# with the License. You may obtain a copy of the License at
+1#
+1# http://www.apache.org/licenses/LICENSE-2.0
+1#
+1# Unless required by applicable law or agreed to in writing,
+1# software distributed under the License is distributed on an
+1# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+1# KIND, either express or implied. See the License for the
+1# specific language governing permissions and limitations
+1# under the License.
+1#
+1
\ No newline at end of file
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 3b55142350b..dcf6741e739 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
@@ -1,1003 +1,1003 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-import os
-from abc import abstractmethod
-from typing import Callable, Dict, List
-
-import numpy as np
-from huggingface_hub import hf_hub_download
-from sklearn.preprocessing import MinMaxScaler
-from sktime.detection.hmm_learn import GMMHMM, GaussianHMM
-from sktime.detection.stray import STRAY
-from sktime.forecasting.arima import ARIMA
-from sktime.forecasting.exp_smoothing import ExponentialSmoothing
-from sktime.forecasting.naive import NaiveForecaster
-from sktime.forecasting.trend import STLForecaster
-
-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.constant import (
- MODEL_CONFIG_FILE_IN_JSON,
- MODEL_WEIGHTS_FILE_IN_SAFETENSORS,
- AttributeName,
-)
-from iotdb.ainode.core.exception import (
- BuiltInModelNotSupportError,
- InferenceModelInternalError,
- ListRangeException,
- NumericalRangeException,
- StringRangeException,
- WrongAttributeTypeError,
-)
-from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
-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(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 weights file not found at {weights_path}, downloading from HuggingFace..."
- )
- try:
- hf_hub_download(
- repo_id=repo_id,
- filename=MODEL_WEIGHTS_FILE_IN_SAFETENSORS,
- local_dir=local_dir,
- )
- logger.info(f"Got file to {weights_path}")
- except Exception as 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
-
-
-def download_built_in_ltsm_from_hf_if_necessary(
- model_type: BuiltInModelType, local_dir: str
-) -> bool:
- """
- Download the built-in ltsm from HuggingFace repository when necessary.
-
- Return:
- bool: True if the model is existed or downloaded successfully, False otherwise.
- """
- repo_id = TIMER_REPO_ID[model_type]
- if not _download_file_from_hf_if_necessary(local_dir, repo_id):
- return False
- return True
-
-
-def get_model_attributes(model_type: BuiltInModelType):
- if model_type == BuiltInModelType.ARIMA:
- attribute_map = arima_attribute_map
- elif model_type == BuiltInModelType.NAIVE_FORECASTER:
- attribute_map = naive_forecaster_attribute_map
- elif (
- model_type == BuiltInModelType.EXPONENTIAL_SMOOTHING
- or model_type == BuiltInModelType.HOLTWINTERS
- ):
- attribute_map = exponential_smoothing_attribute_map
- elif model_type == BuiltInModelType.STL_FORECASTER:
- attribute_map = stl_forecaster_attribute_map
- elif model_type == BuiltInModelType.GMM_HMM:
- attribute_map = gmmhmm_attribute_map
- elif model_type == BuiltInModelType.GAUSSIAN_HMM:
- attribute_map = gaussian_hmm_attribute_map
- elif model_type == BuiltInModelType.STRAY:
- attribute_map = stray_attribute_map
- elif model_type == BuiltInModelType.TIMER_XL:
- attribute_map = timerxl_attribute_map
- elif model_type == BuiltInModelType.SUNDIAL:
- attribute_map = sundial_attribute_map
- else:
- raise BuiltInModelNotSupportError(model_type.value)
- return attribute_map
-
-
-def fetch_built_in_model(
- model_type: BuiltInModelType, model_dir, inference_attrs: Dict[str, str]
-) -> Callable:
- """
- Fetch the built-in model according to its id and directory, not that this directory only contains model weights and config.
- Args:
- model_type: the type of the built-in model
- model_dir: for huggingface models only, the directory where the model is stored
- Returns:
- model: the built-in model
- """
- default_attributes = get_model_attributes(model_type)
- # parse the attributes from inference_attrs
- attributes = parse_attribute(inference_attrs, default_attributes)
-
- # build the built-in model
- if model_type == BuiltInModelType.ARIMA:
- model = ArimaModel(attributes)
- elif (
- model_type == BuiltInModelType.EXPONENTIAL_SMOOTHING
- or model_type == BuiltInModelType.HOLTWINTERS
- ):
- model = ExponentialSmoothingModel(attributes)
- elif model_type == BuiltInModelType.NAIVE_FORECASTER:
- model = NaiveForecasterModel(attributes)
- elif model_type == BuiltInModelType.STL_FORECASTER:
- model = STLForecasterModel(attributes)
- elif model_type == BuiltInModelType.GMM_HMM:
- model = GMMHMMModel(attributes)
- elif model_type == BuiltInModelType.GAUSSIAN_HMM:
- model = GaussianHmmModel(attributes)
- elif model_type == BuiltInModelType.STRAY:
- model = STRAYModel(attributes)
- elif model_type == BuiltInModelType.TIMER_XL:
- model = modeling_timer.TimerForPrediction.from_pretrained(model_dir)
- elif model_type == BuiltInModelType.SUNDIAL:
- model = modeling_sundial.SundialForPrediction.from_pretrained(model_dir)
- else:
- raise BuiltInModelNotSupportError(model_type.value)
-
- return model
-
-
-class Attribute(object):
- def __init__(self, name: str):
- """
- Args:
- name: the name of the attribute
- """
- self._name = name
-
- @abstractmethod
- def get_default_value(self):
- raise NotImplementedError
-
- @abstractmethod
- def validate_value(self, value):
- raise NotImplementedError
-
- @abstractmethod
- def parse(self, string_value: str):
- raise NotImplementedError
-
-
-class IntAttribute(Attribute):
- def __init__(
- self,
- name: str,
- default_value: int,
- default_low: int,
- default_high: int,
- ):
- super(IntAttribute, self).__init__(name)
- self.__default_value = default_value
- self.__default_low = default_low
- self.__default_high = default_high
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if self.__default_low <= value <= self.__default_high:
- return True
- raise NumericalRangeException(
- self._name, value, self.__default_low, self.__default_high
- )
-
- def parse(self, string_value: str):
- try:
- int_value = int(string_value)
- except:
- raise WrongAttributeTypeError(self._name, "int")
- return int_value
-
-
-class FloatAttribute(Attribute):
- def __init__(
- self,
- name: str,
- default_value: float,
- default_low: float,
- default_high: float,
- ):
- super(FloatAttribute, self).__init__(name)
- self.__default_value = default_value
- self.__default_low = default_low
- self.__default_high = default_high
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if self.__default_low <= value <= self.__default_high:
- return True
- raise NumericalRangeException(
- self._name, value, self.__default_low, self.__default_high
- )
-
- def parse(self, string_value: str):
- try:
- float_value = float(string_value)
- except:
- raise WrongAttributeTypeError(self._name, "float")
- return float_value
-
-
-class StringAttribute(Attribute):
- def __init__(self, name: str, default_value: str, value_choices: List[str]):
- super(StringAttribute, self).__init__(name)
- self.__default_value = default_value
- self.__value_choices = value_choices
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if value in self.__value_choices:
- return True
- raise StringRangeException(self._name, value, self.__value_choices)
-
- def parse(self, string_value: str):
- return string_value
-
-
-class BooleanAttribute(Attribute):
- def __init__(self, name: str, default_value: bool):
- super(BooleanAttribute, self).__init__(name)
- self.__default_value = default_value
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if isinstance(value, bool):
- return True
- raise WrongAttributeTypeError(self._name, "bool")
-
- def parse(self, string_value: str):
- if string_value.lower() == "true":
- return True
- elif string_value.lower() == "false":
- return False
- else:
- raise WrongAttributeTypeError(self._name, "bool")
-
-
-class ListAttribute(Attribute):
- def __init__(self, name: str, default_value: List, value_type):
- """
- value_type is the type of the elements in the list, e.g. int, float, str
- """
- super(ListAttribute, self).__init__(name)
- self.__default_value = default_value
- self.__value_type = value_type
- self.__type_to_str = {str: "str", int: "int", float: "float"}
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if not isinstance(value, list):
- raise WrongAttributeTypeError(self._name, "list")
- for value_item in value:
- if not isinstance(value_item, self.__value_type):
- raise WrongAttributeTypeError(self._name, self.__value_type)
- return True
-
- def parse(self, string_value: str):
- try:
- list_value = eval(string_value)
- except:
- raise WrongAttributeTypeError(self._name, "list")
- if not isinstance(list_value, list):
- raise WrongAttributeTypeError(self._name, "list")
- for i in range(len(list_value)):
- try:
- list_value[i] = self.__value_type(list_value[i])
- except:
- raise ListRangeException(
- self._name, list_value, self.__type_to_str[self.__value_type]
- )
- return list_value
-
-
-class TupleAttribute(Attribute):
- def __init__(self, name: str, default_value: tuple, value_type):
- """
- value_type is the type of the elements in the list, e.g. int, float, str
- """
- super(TupleAttribute, self).__init__(name)
- self.__default_value = default_value
- self.__value_type = value_type
- self.__type_to_str = {str: "str", int: "int", float: "float"}
-
- def get_default_value(self):
- return self.__default_value
-
- def validate_value(self, value):
- if not isinstance(value, tuple):
- raise WrongAttributeTypeError(self._name, "tuple")
- for value_item in value:
- if not isinstance(value_item, self.__value_type):
- raise WrongAttributeTypeError(self._name, self.__value_type)
- return True
-
- def parse(self, string_value: str):
- try:
- tuple_value = eval(string_value)
- except:
- raise WrongAttributeTypeError(self._name, "tuple")
- if not isinstance(tuple_value, tuple):
- raise WrongAttributeTypeError(self._name, "tuple")
- list_value = list(tuple_value)
- for i in range(len(list_value)):
- try:
- list_value[i] = self.__value_type(list_value[i])
- except:
- raise ListRangeException(
- self._name, list_value, self.__type_to_str[self.__value_type]
- )
- tuple_value = tuple(list_value)
- return tuple_value
-
-
-def parse_attribute(
- input_attributes: Dict[str, str], attribute_map: Dict[str, Attribute]
-):
- """
- Args:
- input_attributes: a dict of attributes, where the key is the attribute name, the value is the string value of
- the attribute
- attribute_map: a dict of hyperparameters, where the key is the attribute name, the value is the Attribute
- object
- Returns:
- a dict of attributes, where the key is the attribute name, the value is the parsed value of the attribute
- """
- attributes = {}
- for attribute_name in attribute_map:
- # user specified the attribute
- if attribute_name in input_attributes:
- attribute = attribute_map[attribute_name]
- value = attribute.parse(input_attributes[attribute_name])
- attribute.validate_value(value)
- attributes[attribute_name] = value
- # user did not specify the attribute, use the default value
- else:
- try:
- attributes[attribute_name] = attribute_map[
- attribute_name
- ].get_default_value()
- except NotImplementedError as e:
- logger.error(f"attribute {attribute_name} is not implemented.")
- raise e
- return attributes
-
-
-sundial_attribute_map = {
- AttributeName.INPUT_TOKEN_LEN.value: IntAttribute(
- name=AttributeName.INPUT_TOKEN_LEN.value,
- default_value=16,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.HIDDEN_SIZE.value: IntAttribute(
- name=AttributeName.HIDDEN_SIZE.value,
- default_value=768,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.INTERMEDIATE_SIZE.value: IntAttribute(
- name=AttributeName.INTERMEDIATE_SIZE.value,
- default_value=3072,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.OUTPUT_TOKEN_LENS.value: ListAttribute(
- name=AttributeName.OUTPUT_TOKEN_LENS.value, default_value=[720], value_type=int
- ),
- AttributeName.NUM_HIDDEN_LAYERS.value: IntAttribute(
- name=AttributeName.NUM_HIDDEN_LAYERS.value,
- default_value=12,
- default_low=1,
- default_high=16,
- ),
- AttributeName.NUM_ATTENTION_HEADS.value: IntAttribute(
- name=AttributeName.NUM_ATTENTION_HEADS.value,
- default_value=12,
- default_low=1,
- default_high=192,
- ),
- AttributeName.HIDDEN_ACT.value: StringAttribute(
- name=AttributeName.HIDDEN_ACT.value,
- default_value="silu",
- value_choices=["relu", "gelu", "silu", "tanh"],
- ),
- AttributeName.USE_CACHE.value: BooleanAttribute(
- name=AttributeName.USE_CACHE.value,
- default_value=True,
- ),
- AttributeName.ROPE_THETA.value: IntAttribute(
- name=AttributeName.ROPE_THETA.value,
- default_value=10000,
- default_low=1000,
- default_high=50000,
- ),
- AttributeName.DROPOUT_RATE.value: FloatAttribute(
- name=AttributeName.DROPOUT_RATE.value,
- default_value=0.1,
- default_low=0.0,
- default_high=1.0,
- ),
- AttributeName.INITIALIZER_RANGE.value: FloatAttribute(
- name=AttributeName.INITIALIZER_RANGE.value,
- default_value=0.02,
- default_low=0.0,
- default_high=1.0,
- ),
- AttributeName.MAX_POSITION_EMBEDDINGS.value: IntAttribute(
- name=AttributeName.MAX_POSITION_EMBEDDINGS.value,
- default_value=10000,
- default_low=1,
- default_high=50000,
- ),
- AttributeName.FLOW_LOSS_DEPTH.value: IntAttribute(
- name=AttributeName.FLOW_LOSS_DEPTH.value,
- default_value=3,
- default_low=1,
- default_high=50,
- ),
- AttributeName.NUM_SAMPLING_STEPS.value: IntAttribute(
- name=AttributeName.NUM_SAMPLING_STEPS.value,
- default_value=50,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.DIFFUSION_BATCH_MUL.value: IntAttribute(
- name=AttributeName.DIFFUSION_BATCH_MUL.value,
- default_value=4,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.CKPT_PATH.value: StringAttribute(
- name=AttributeName.CKPT_PATH.value,
- default_value=os.path.join(
- os.getcwd(),
- AINodeDescriptor().get_config().get_ain_models_dir(),
- "weights",
- "sundial",
- ),
- value_choices=[""],
- ),
-}
-
-timerxl_attribute_map = {
- AttributeName.INPUT_TOKEN_LEN.value: IntAttribute(
- name=AttributeName.INPUT_TOKEN_LEN.value,
- default_value=96,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.HIDDEN_SIZE.value: IntAttribute(
- name=AttributeName.HIDDEN_SIZE.value,
- default_value=1024,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.INTERMEDIATE_SIZE.value: IntAttribute(
- name=AttributeName.INTERMEDIATE_SIZE.value,
- default_value=2048,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.OUTPUT_TOKEN_LENS.value: ListAttribute(
- name=AttributeName.OUTPUT_TOKEN_LENS.value, default_value=[96], value_type=int
- ),
- AttributeName.NUM_HIDDEN_LAYERS.value: IntAttribute(
- name=AttributeName.NUM_HIDDEN_LAYERS.value,
- default_value=8,
- default_low=1,
- default_high=16,
- ),
- AttributeName.NUM_ATTENTION_HEADS.value: IntAttribute(
- name=AttributeName.NUM_ATTENTION_HEADS.value,
- default_value=8,
- default_low=1,
- default_high=192,
- ),
- AttributeName.HIDDEN_ACT.value: StringAttribute(
- name=AttributeName.HIDDEN_ACT.value,
- default_value="silu",
- value_choices=["relu", "gelu", "silu", "tanh"],
- ),
- AttributeName.USE_CACHE.value: BooleanAttribute(
- name=AttributeName.USE_CACHE.value,
- default_value=True,
- ),
- AttributeName.ROPE_THETA.value: IntAttribute(
- name=AttributeName.ROPE_THETA.value,
- default_value=10000,
- default_low=1000,
- default_high=50000,
- ),
- AttributeName.ATTENTION_DROPOUT.value: FloatAttribute(
- name=AttributeName.ATTENTION_DROPOUT.value,
- default_value=0.0,
- default_low=0.0,
- default_high=1.0,
- ),
- AttributeName.INITIALIZER_RANGE.value: FloatAttribute(
- name=AttributeName.INITIALIZER_RANGE.value,
- default_value=0.02,
- default_low=0.0,
- default_high=1.0,
- ),
- AttributeName.MAX_POSITION_EMBEDDINGS.value: IntAttribute(
- name=AttributeName.MAX_POSITION_EMBEDDINGS.value,
- default_value=10000,
- default_low=1,
- default_high=50000,
- ),
- AttributeName.CKPT_PATH.value: StringAttribute(
- name=AttributeName.CKPT_PATH.value,
- default_value=os.path.join(
- os.getcwd(),
- AINodeDescriptor().get_config().get_ain_models_dir(),
- "weights",
- "timerxl",
- "model.safetensors",
- ),
- value_choices=[""],
- ),
-}
-
-# built-in sktime model attributes
-# NaiveForecaster
-naive_forecaster_attribute_map = {
- AttributeName.PREDICT_LENGTH.value: IntAttribute(
- name=AttributeName.PREDICT_LENGTH.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.STRATEGY.value: StringAttribute(
- name=AttributeName.STRATEGY.value,
- default_value="last",
- value_choices=["last", "mean"],
- ),
- AttributeName.SP.value: IntAttribute(
- name=AttributeName.SP.value, default_value=1, default_low=1, default_high=5000
- ),
-}
-# ExponentialSmoothing
-exponential_smoothing_attribute_map = {
- AttributeName.PREDICT_LENGTH.value: IntAttribute(
- name=AttributeName.PREDICT_LENGTH.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.DAMPED_TREND.value: BooleanAttribute(
- name=AttributeName.DAMPED_TREND.value,
- default_value=False,
- ),
- AttributeName.INITIALIZATION_METHOD.value: StringAttribute(
- name=AttributeName.INITIALIZATION_METHOD.value,
- default_value="estimated",
- value_choices=["estimated", "heuristic", "legacy-heuristic", "known"],
- ),
- AttributeName.OPTIMIZED.value: BooleanAttribute(
- name=AttributeName.OPTIMIZED.value,
- default_value=True,
- ),
- AttributeName.REMOVE_BIAS.value: BooleanAttribute(
- name=AttributeName.REMOVE_BIAS.value,
- default_value=False,
- ),
- AttributeName.USE_BRUTE.value: BooleanAttribute(
- name=AttributeName.USE_BRUTE.value,
- default_value=False,
- ),
-}
-# Arima
-arima_attribute_map = {
- AttributeName.PREDICT_LENGTH.value: IntAttribute(
- name=AttributeName.PREDICT_LENGTH.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.ORDER.value: TupleAttribute(
- name=AttributeName.ORDER.value, default_value=(1, 0, 0), value_type=int
- ),
- AttributeName.SEASONAL_ORDER.value: TupleAttribute(
- name=AttributeName.SEASONAL_ORDER.value,
- default_value=(0, 0, 0, 0),
- value_type=int,
- ),
- AttributeName.METHOD.value: StringAttribute(
- name=AttributeName.METHOD.value,
- default_value="lbfgs",
- value_choices=["lbfgs", "bfgs", "newton", "nm", "cg", "ncg", "powell"],
- ),
- AttributeName.MAXITER.value: IntAttribute(
- name=AttributeName.MAXITER.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.SUPPRESS_WARNINGS.value: BooleanAttribute(
- name=AttributeName.SUPPRESS_WARNINGS.value,
- default_value=True,
- ),
- AttributeName.OUT_OF_SAMPLE_SIZE.value: IntAttribute(
- name=AttributeName.OUT_OF_SAMPLE_SIZE.value,
- default_value=0,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.SCORING.value: StringAttribute(
- name=AttributeName.SCORING.value,
- default_value="mse",
- value_choices=["mse", "mae", "rmse", "mape", "smape", "rmsle", "r2"],
- ),
- AttributeName.WITH_INTERCEPT.value: BooleanAttribute(
- name=AttributeName.WITH_INTERCEPT.value,
- default_value=True,
- ),
- AttributeName.TIME_VARYING_REGRESSION.value: BooleanAttribute(
- name=AttributeName.TIME_VARYING_REGRESSION.value,
- default_value=False,
- ),
- AttributeName.ENFORCE_STATIONARITY.value: BooleanAttribute(
- name=AttributeName.ENFORCE_STATIONARITY.value,
- default_value=True,
- ),
- AttributeName.ENFORCE_INVERTIBILITY.value: BooleanAttribute(
- name=AttributeName.ENFORCE_INVERTIBILITY.value,
- default_value=True,
- ),
- AttributeName.SIMPLE_DIFFERENCING.value: BooleanAttribute(
- name=AttributeName.SIMPLE_DIFFERENCING.value,
- default_value=False,
- ),
- AttributeName.MEASUREMENT_ERROR.value: BooleanAttribute(
- name=AttributeName.MEASUREMENT_ERROR.value,
- default_value=False,
- ),
- AttributeName.MLE_REGRESSION.value: BooleanAttribute(
- name=AttributeName.MLE_REGRESSION.value,
- default_value=True,
- ),
- AttributeName.HAMILTON_REPRESENTATION.value: BooleanAttribute(
- name=AttributeName.HAMILTON_REPRESENTATION.value,
- default_value=False,
- ),
- AttributeName.CONCENTRATE_SCALE.value: BooleanAttribute(
- name=AttributeName.CONCENTRATE_SCALE.value,
- default_value=False,
- ),
-}
-# STLForecaster
-stl_forecaster_attribute_map = {
- AttributeName.PREDICT_LENGTH.value: IntAttribute(
- name=AttributeName.PREDICT_LENGTH.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.SP.value: IntAttribute(
- name=AttributeName.SP.value, default_value=2, default_low=1, default_high=5000
- ),
- AttributeName.SEASONAL.value: IntAttribute(
- name=AttributeName.SEASONAL.value,
- default_value=7,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.SEASONAL_DEG.value: IntAttribute(
- name=AttributeName.SEASONAL_DEG.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.TREND_DEG.value: IntAttribute(
- name=AttributeName.TREND_DEG.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.LOW_PASS_DEG.value: IntAttribute(
- name=AttributeName.LOW_PASS_DEG.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.SEASONAL_JUMP.value: IntAttribute(
- name=AttributeName.SEASONAL_JUMP.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.TREND_JUMP.value: IntAttribute(
- name=AttributeName.TREND_JUMP.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
- AttributeName.LOSS_PASS_JUMP.value: IntAttribute(
- name=AttributeName.LOSS_PASS_JUMP.value,
- default_value=1,
- default_low=0,
- default_high=5000,
- ),
-}
-
-# GAUSSIAN_HMM
-gaussian_hmm_attribute_map = {
- AttributeName.N_COMPONENTS.value: IntAttribute(
- name=AttributeName.N_COMPONENTS.value,
- default_value=1,
- default_low=1,
- default_high=5000,
- ),
- AttributeName.COVARIANCE_TYPE.value: StringAttribute(
- name=AttributeName.COVARIANCE_TYPE.value,
- default_value="diag",
- value_choices=["spherical", "diag", "full", "tied"],
- ),
- AttributeName.MIN_COVAR.value: FloatAttribute(
- name=AttributeName.MIN_COVAR.value,
- default_value=1e-3,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.STARTPROB_PRIOR.value: FloatAttribute(
- name=AttributeName.STARTPROB_PRIOR.value,
- default_value=1.0,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.TRANSMAT_PRIOR.value: FloatAttribute(
- name=AttributeName.TRANSMAT_PRIOR.value,
- default_value=1.0,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.MEANS_PRIOR.value: FloatAttribute(
- name=AttributeName.MEANS_PRIOR.value,
- default_value=0.0,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.MEANS_WEIGHT.value: FloatAttribute(
- name=AttributeName.MEANS_WEIGHT.value,
- default_value=0.0,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.COVARS_PRIOR.value: FloatAttribute(
- name=AttributeName.COVARS_PRIOR.value,
- default_value=1e-2,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.COVARS_WEIGHT.value: FloatAttribute(
- name=AttributeName.COVARS_WEIGHT.value,
- default_value=1.0,
- default_low=-1e10,
- default_high=1e10,
- ),
- AttributeName.ALGORITHM.value: StringAttribute(
- name=AttributeName.ALGORITHM.value,
- default_value="viterbi",
- value_choices=["viterbi", "map"],
- ),
... 12116 lines suppressed ...