This is an automated email from the ASF dual-hosted git repository. yongzao pushed a commit to branch cp-ain-to-206 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit e77da3951e62c45de901c7826217673a66a7eb3f Author: Yongzao <[email protected]> AuthorDate: Tue Dec 16 13:29:41 2025 +0800 [AINode] Refactor code base --- iotdb-core/ainode/.gitignore | 2 +- .../core/inference/inference_request_pool.py | 197 ------------------ iotdb-core/ainode/{ => iotdb}/ainode/__init__.py | 0 iotdb-core/ainode/iotdb/ainode/conf/git.properties | 3 - .../ainode/{ => iotdb}/ainode/core/__init__.py | 0 .../ainode/{ => iotdb}/ainode/core/ai_node.py | 2 +- .../ainode/{ => iotdb}/ainode/core/config.py | 2 +- .../ainode/{ => iotdb}/ainode/core/constant.py | 2 +- .../ainode/{ => iotdb}/ainode/core/exception.py | 0 .../{ => iotdb}/ainode/core/inference/__init__.py | 0 .../ainode/core/inference/batcher}/__init__.py | 0 .../inference/batcher/abstract_batcher.py} | 29 ++- .../ainode/core/inference/batcher/basic_batcher.py | 63 ++++++ .../ainode/core/inference/dispatcher}/__init__.py | 0 .../inference/dispatcher/abstract_dispatcher.py} | 32 ++- .../core/inference/dispatcher/basic_dispatcher.py | 60 ++++++ .../ainode/core/inference/inference_request.py | 18 +- .../core/inference/inference_request_pool.py | 223 +++++++++++++++++++++ .../core/inference/pool_scheduler}/__init__.py | 0 .../pool_scheduler/abstract_pool_scheduler.py | 86 ++++++++ .../pool_scheduler/basic_pool_scheduler.py | 193 ++++++++++++++++++ .../core/inference/request_scheduler}/__init__.py | 0 .../abstract_request_scheduler.py} | 2 +- .../request_scheduler/basic_request_scheduler.py} | 11 +- .../ainode/core/inference/strategy}/__init__.py | 0 .../strategy/abstract_inference_pipeline.py | 0 .../strategy/timer_sundial_inference_pipeline.py | 6 +- .../strategy/timerxl_inference_pipeline.py | 12 +- .../{ => iotdb}/ainode/core/inference/utils.py | 0 .../ainode/core/ingress}/__init__.py | 0 .../{ => iotdb}/ainode/core/ingress/dataset.py | 0 .../{ => iotdb}/ainode/core/ingress/iotdb.py | 4 +- iotdb-core/ainode/{ => iotdb}/ainode/core/log.py | 0 .../ainode/core/manager}/__init__.py | 0 .../ainode/core/manager/cluster_manager.py | 0 .../ainode/core/manager/inference_manager.py | 16 +- .../ainode/core/manager/model_manager.py | 4 +- .../{ => iotdb}/ainode/core/manager/utils.py | 2 +- .../rpc => iotdb/ainode/core/model}/__init__.py | 0 .../ainode/core/model/built_in_model_factory.py | 8 +- .../{ => iotdb}/ainode/core/model/model_enums.py | 0 .../{ => iotdb}/ainode/core/model/model_factory.py | 6 +- .../{ => iotdb}/ainode/core/model/model_info.py | 2 +- .../{ => iotdb}/ainode/core/model/model_storage.py | 8 +- .../ainode/core/model/sundial}/__init__.py | 0 .../core/model/sundial/configuration_sundial.py | 0 .../ainode/core/model/sundial/flow_loss.py | 0 .../ainode/core/model/sundial/modeling_sundial.py | 5 +- .../core/model/sundial/ts_generation_mixin.py | 0 .../ainode/core/model/timerxl}/__init__.py | 0 .../core/model/timerxl/configuration_timer.py | 0 .../ainode/core/model/timerxl/modeling_timer.py | 0 .../core/model/timerxl/ts_generation_mixin.py | 0 .../{ => iotdb}/ainode/core/model/uri_utils.py | 4 +- .../core => iotdb/ainode/core/rpc}/__init__.py | 0 .../ainode/{ => iotdb}/ainode/core/rpc/client.py | 2 +- .../ainode/{ => iotdb}/ainode/core/rpc/handler.py | 0 .../ainode/{ => iotdb}/ainode/core/rpc/service.py | 2 +- .../ainode/{ => iotdb}/ainode/core/rpc/status.py | 0 .../ainode/{ => iotdb}/ainode/core/script.py | 6 +- .../core => iotdb/ainode/core/util}/__init__.py | 0 .../{ => iotdb}/ainode/core/util/activation.py | 0 .../ainode/{ => iotdb}/ainode/core/util/cache.py | 2 +- .../{ => iotdb}/ainode/core/util/decorator.py | 0 .../ainode/core/util/huggingface_cache.py | 0 .../ainode/{ => iotdb}/ainode/core/util/lock.py | 0 .../ainode/{ => iotdb}/ainode/core/util/masking.py | 0 .../ainode/{ => iotdb}/ainode/core/util/serde.py | 2 +- 68 files changed, 740 insertions(+), 276 deletions(-) diff --git a/iotdb-core/ainode/.gitignore b/iotdb-core/ainode/.gitignore index c30efa5f662..8cc2098c3fd 100644 --- a/iotdb-core/ainode/.gitignore +++ b/iotdb-core/ainode/.gitignore @@ -12,7 +12,7 @@ /iotdb/table_session_pool.py # generated by maven -/ainode/conf/ +/iotdb/ainode/conf/ # .whl of ainode, generated by Poetry /dist/ diff --git a/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py b/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py deleted file mode 100644 index 4bf594860f7..00000000000 --- a/iotdb-core/ainode/ainode/core/inference/inference_request_pool.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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 random -import threading -import time - -import numpy as np -import torch -import torch.multiprocessing as mp -from transformers import PretrainedConfig - -from ainode.core.config import AINodeDescriptor -from ainode.core.constant import INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE -from ainode.core.inference.scheduler.basic_scheduler import BasicScheduler -from ainode.core.log import Logger -from ainode.core.manager.model_manager import ModelManager - - -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, - 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.model = None - self._model_manager = None - self.device = None - self.ready_event = ready_event - - 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._scheduler = BasicScheduler( - self._waiting_queue, self._running_queue, self._finished_queue, self.pool_id - ) - self._stop_event = mp.Event() - - # Fix inference seed - random.seed(self.FIX_SEED) - torch.manual_seed(self.FIX_SEED) - np.random.seed(self.FIX_SEED) - - def _warm_up_and_estimate_memory(self): - # TODO: Test per token memory usage, add support for cpu in the future - torch.cuda.empty_cache() - gc.collect() - dummy_input = torch.zeros( - (1, self.config.input_token_len), dtype=torch.float32 - ).to(self.device) - - # force cuda synchronization to avoid any asynchronous memory allocation issues - torch.cuda.reset_peak_memory_stats(self.device) - torch.cuda.synchronize(self.device) - memory_before_warmup = torch.cuda.memory_allocated(self.device) - self.logger.info( - f"[Inference][Device-{self.device}][Pool-{self.pool_id}] Before warm-up, peak memory usage: {memory_before_warmup:.2f} bytes" - ) - - # warm-up - with torch.no_grad(): - self.model.generate(dummy_input, max_new_tokens=1) - torch.cuda.synchronize(self.device) - peak_memory_1_token = torch.cuda.max_memory_allocated(self.device) - self.logger.info( - f"[Inference][Device-{self.device}][Pool-{self.pool_id}] Baseline memory usage for 1 token: {peak_memory_1_token:.2f} bytes" - ) - self.logger.info( - f"[Inference][Device-{self.device}][Pool-{self.pool_id}] Differentiation : {peak_memory_1_token-memory_before_warmup:.2f} bytes" - ) - - def _activate_requests(self): - requests = self._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}][ID-{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): - requests = self._scheduler.schedule_step() - # TODO: We need a batcher to accelerate the concurrent inference - for request in requests: - if self.model_id == "sundial": - request.inputs = request.inputs.to(self.device) - output = self.model.generate( - request.inputs, - max_new_tokens=request.max_new_tokens, - num_samples=10, - revin=True, - ) - request.output_tensor = request.output_tensor.to(self.device) - request.write_step_output(output[0].mean(dim=0)) - elif self.model_id == "timer_xl": - request.inputs = request.inputs.to(self.device) - output = self.model.generate( - request.inputs, - max_new_tokens=request.max_new_tokens, - revin=True, - ) - request.output_tensor = request.output_tensor.to(self.device) - request.write_step_output(output[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) - - 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._model_manager = ModelManager() - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - self.logger = Logger( - INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE.format(self.device) - ) - self._scheduler.device = self.device - self.model = self._model_manager.load_model(self.model_id, {}).to(self.device) - self.ready_event.set() - - # self._warm_up_and_estimate_memory() - - 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 exited cleanly." - ) - - def stop(self): - self._stop_event.set() diff --git a/iotdb-core/ainode/ainode/__init__.py b/iotdb-core/ainode/iotdb/ainode/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/__init__.py rename to iotdb-core/ainode/iotdb/ainode/__init__.py diff --git a/iotdb-core/ainode/iotdb/ainode/conf/git.properties b/iotdb-core/ainode/iotdb/ainode/conf/git.properties deleted file mode 100644 index a298fc3dc61..00000000000 --- a/iotdb-core/ainode/iotdb/ainode/conf/git.properties +++ /dev/null @@ -1,3 +0,0 @@ -#Generated by Git-Commit-Id-Plugin -git.commit.id.abbrev=391429b -git.dirty=false diff --git a/iotdb-core/ainode/ainode/core/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/__init__.py similarity index 100% copy from iotdb-core/ainode/ainode/core/__init__.py copy to iotdb-core/ainode/iotdb/ainode/core/__init__.py diff --git a/iotdb-core/ainode/ainode/core/ai_node.py b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py similarity index 99% rename from iotdb-core/ainode/ainode/core/ai_node.py rename to iotdb-core/ainode/iotdb/ainode/core/ai_node.py index 71536344391..4fe2a9b97dc 100644 --- a/iotdb-core/ainode/ainode/core/ai_node.py +++ b/iotdb-core/ainode/iotdb/ainode/core/ai_node.py @@ -22,7 +22,7 @@ from datetime import datetime import psutil -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import AINODE_SYSTEM_FILE_NAME from ainode.core.log import Logger from ainode.core.rpc.client import ClientManager diff --git a/iotdb-core/ainode/ainode/core/config.py b/iotdb-core/ainode/iotdb/ainode/core/config.py similarity index 99% rename from iotdb-core/ainode/ainode/core/config.py rename to iotdb-core/ainode/iotdb/ainode/core/config.py index d0244b3784e..4e9843c143e 100644 --- a/iotdb-core/ainode/ainode/core/config.py +++ b/iotdb-core/ainode/iotdb/ainode/core/config.py @@ -47,7 +47,7 @@ from ainode.core.constant import ( AINODE_THRIFT_COMPRESSION_ENABLED, AINODE_VERSION_INFO, ) -from ainode.core.exception import BadNodeUrlError +from iotdb.ainode.core.exception import BadNodeUrlError from ainode.core.log import Logger from ainode.core.util.decorator import singleton from ainode.thrift.common.ttypes import TEndPoint diff --git a/iotdb-core/ainode/ainode/core/constant.py b/iotdb-core/ainode/iotdb/ainode/core/constant.py similarity index 99% rename from iotdb-core/ainode/ainode/core/constant.py rename to iotdb-core/ainode/iotdb/ainode/core/constant.py index 9db8aec1d5b..954380b225f 100644 --- a/iotdb-core/ainode/ainode/core/constant.py +++ b/iotdb-core/ainode/iotdb/ainode/core/constant.py @@ -21,7 +21,7 @@ import os from enum import Enum from typing import List -from ainode.core.model.model_enums import BuiltInModelType +from iotdb.ainode.core.model.model_enums import BuiltInModelType from ainode.thrift.common.ttypes import TEndPoint AINODE_VERSION_INFO = "UNKNOWN" diff --git a/iotdb-core/ainode/ainode/core/exception.py b/iotdb-core/ainode/iotdb/ainode/core/exception.py similarity index 100% rename from iotdb-core/ainode/ainode/core/exception.py rename to iotdb-core/ainode/iotdb/ainode/core/exception.py diff --git a/iotdb-core/ainode/ainode/core/inference/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/inference/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/__init__.py diff --git a/iotdb-core/ainode/ainode/core/inference/scheduler/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/inference/scheduler/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/batcher/__init__.py diff --git a/iotdb-core/ainode/iotdb/ainode/conf/pom.properties b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py similarity index 58% rename from iotdb-core/ainode/iotdb/ainode/conf/pom.properties rename to iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py index 96a8d57baa0..7b668e93a8f 100644 --- a/iotdb-core/ainode/iotdb/ainode/conf/pom.properties +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/abstract_batcher.py @@ -1,4 +1,3 @@ -# # 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 @@ -15,7 +14,29 @@ # 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 -artifactId=iotdb-ainode -groupId=org.apache.iotdb -version=2.0.6-SNAPSHOT \ No newline at end of file + @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 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 new file mode 100644 index 00000000000..591a0d7c1dd --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/batcher/basic_batcher.py @@ -0,0 +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 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 diff --git a/iotdb-core/ainode/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/inference/strategy/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/__init__.py diff --git a/iotdb-core/ainode/ainode/core/rpc/status.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py similarity index 50% copy from iotdb-core/ainode/ainode/core/rpc/status.py copy to iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py index 37368b0068b..3abc1477a97 100644 --- a/iotdb-core/ainode/ainode/core/rpc/status.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/abstract_dispatcher.py @@ -16,18 +16,28 @@ # under the License. # -from ainode.core.constant import TSStatusCode -from ainode.core.log import Logger -from ainode.thrift.common.ttypes import TSStatus +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 -def get_status(status_code: TSStatusCode, message: str = None) -> TSStatus: - status = TSStatus(status_code.get_status_code()) - status.message = message - return status +class AbstractDispatcher(ABC): + """ + Abstract base class for dispatchers that handle inference requests. + """ -def verify_success(status: TSStatus, err_msg: str) -> None: - if status.code != TSStatusCode.SUCCESS_STATUS.get_status_code(): - Logger().warning(err_msg + ", error status is ", status) - raise RuntimeError(str(status.code) + ": " + status.message) + 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 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 new file mode 100644 index 00000000000..d06ba55546f --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py @@ -0,0 +1,60 @@ +# 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 diff --git a/iotdb-core/ainode/ainode/core/inference/inference_request.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py similarity index 84% rename from iotdb-core/ainode/ainode/core/inference/inference_request.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py index 2c45826fd26..82c72cc37ab 100644 --- a/iotdb-core/ainode/ainode/core/inference/inference_request.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py @@ -20,10 +20,11 @@ from typing import Any import torch -from ainode.core.inference.strategy.abstract_inference_pipeline import ( +from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import ( AbstractInferencePipeline, ) -from ainode.core.log import Logger +from iotdb.ainode.core.log import Logger +from iotdb.ainode.core.util.atmoic_int import AtomicInt logger = Logger() @@ -38,6 +39,7 @@ class InferenceRequest: def __init__( self, req_id: str, + model_id: str, inputs: torch.Tensor, inference_pipeline: AbstractInferencePipeline, max_new_tokens: int = 96, @@ -47,6 +49,7 @@ class InferenceRequest: 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 @@ -57,6 +60,8 @@ class InferenceRequest: 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( @@ -106,15 +111,22 @@ class InferenceRequestProxy: 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 wait_for_completion(self) -> Any: + 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 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 new file mode 100644 index 00000000000..6b054c91fe3 --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py @@ -0,0 +1,223 @@ +# 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.model.model_enums import BuiltInModelType +from iotdb.ainode.core.model.model_info import ModelInfo +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_info: ModelInfo, + 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_info = model_info + 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_info.model_type == BuiltInModelType.SUNDIAL.value: + 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 + request.write_step_output(cur_output.mean(dim=1)) + + 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_info.model_type == BuiltInModelType.TIMER_XL.value: + 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_info.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() + self._logger.info( + f"[Inference][Device-{self.device}][Pool-{self.pool_id}] InferenceRequestPool for model {self.model_info.model_id} is activated." + ) + for thread in self._threads: + thread.join() + self._logger.info( + f"[Inference][Device-{self.device}][Pool-{self.pool_id}] InferenceRequestPool for model {self.model_info.model_id} exited cleanly." + ) + + def stop(self): + self._stop_event.set() diff --git a/iotdb-core/ainode/ainode/core/ingress/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/ingress/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/__init__.py 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 new file mode 100644 index 00000000000..19d21f5822d --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/abstract_pool_scheduler.py @@ -0,0 +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 +from iotdb.ainode.core.model.model_info import ModelInfo + + +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_info: ModelInfo, device_id: str + ) -> List[ScaleAction]: + """ + Schedule a series of actions to load the model to the device. + Args: + model_info: 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_info: ModelInfo, device_id: str + ) -> List[ScaleAction]: + """ + Schedule a series of actions to unload the model from the device. + Args: + model_info: The model to be unloaded. + device_id: The device to unload the model from. + Returns: + A list of ScaleAction to be performed. + """ + pass 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 new file mode 100644 index 00000000000..6a2bd2b619a --- /dev/null +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py @@ -0,0 +1,193 @@ +# 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.model_manager import ModelManager +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, ModelInfo +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_infos: List[ModelInfo], + new_model_info: Optional[ModelInfo] = 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_infos + ( + [new_model_info] if new_model_info is not None else [] + ) + + # Seize memory usage for each model + mem_usages: Dict[str, float] = {} + for model_info in all_models: + mem_usages[model_info.model_id] = ( + MODEL_MEM_USAGE_MAP[model_info.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_info in all_models: + pool_num = int(per_model_share // mem_usages[model_info.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_info.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_info.model_id] / 1024 ** 2:.2f} MB" + ) + allocation[model_info.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) + self._model_manager = ModelManager() + + 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_info: ModelInfo, device_id: str + ) -> List[ScaleAction]: + existing_model_infos = [ + self._model_manager.get_model_info(existing_model_id) + for existing_model_id, pool_group_map in self._request_pool_map.items() + if existing_model_id != model_info.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_infos=existing_model_infos, + new_model_info=model_info, + ) + return self._convert_allocation_result_to_scale_actions( + allocation_result, device_id + ) + + def schedule_unload_model_from_device( + self, model_info: ModelInfo, device_id: str + ) -> List[ScaleAction]: + existing_model_infos = [ + self._model_manager.get_model_info(existing_model_id) + for existing_model_id, pool_group_map in self._request_pool_map.items() + if existing_model_id != model_info.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_infos=existing_model_infos, + new_model_info=None, + ) + if len(existing_model_infos) > 0 + else {model_info.model_id: 0} + ) + if len(existing_model_infos) > 0: + allocation_result[model_info.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 diff --git a/iotdb-core/ainode/ainode/core/manager/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/manager/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/__init__.py diff --git a/iotdb-core/ainode/ainode/core/inference/scheduler/abstract_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py similarity index 98% rename from iotdb-core/ainode/ainode/core/inference/scheduler/abstract_scheduler.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py index 8bc34e529c7..a6c2fe53cae 100644 --- a/iotdb-core/ainode/ainode/core/inference/scheduler/abstract_scheduler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/abstract_request_scheduler.py @@ -19,7 +19,7 @@ from abc import ABC, abstractmethod -class AbstractScheduler(ABC): +class AbstractRequestScheduler(ABC): """ Abstract base class for inference scheduling strategies. diff --git a/iotdb-core/ainode/ainode/core/inference/scheduler/basic_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py similarity index 89% rename from iotdb-core/ainode/ainode/core/inference/scheduler/basic_scheduler.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py index 65dee81ca1c..ef5d37a18a1 100644 --- a/iotdb-core/ainode/ainode/core/inference/scheduler/basic_scheduler.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/request_scheduler/basic_request_scheduler.py @@ -21,16 +21,17 @@ import os import psutil import torch -from ainode.core.inference.inference_request import InferenceRequest -from ainode.core.inference.scheduler.abstract_scheduler import AbstractScheduler -from ainode.core.log import Logger +from iotdb.ainode.core.inference.request_scheduler.abstract_request_scheduler import ( + AbstractRequestScheduler, +) +from iotdb.ainode.core.log import Logger logger = Logger() -class BasicScheduler(AbstractScheduler): +class BasicRequestScheduler(AbstractRequestScheduler): """ - A simple FIFO scheduler that selects requests based on memory availability and activation/step size. + A simple FIFO request scheduler that selects requests based on memory availability and activation/step size. """ def __init__( diff --git a/iotdb-core/ainode/ainode/core/model/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py diff --git a/iotdb-core/ainode/ainode/core/inference/strategy/abstract_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py similarity index 100% rename from iotdb-core/ainode/ainode/core/inference/strategy/abstract_inference_pipeline.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py diff --git a/iotdb-core/ainode/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py similarity index 88% rename from iotdb-core/ainode/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py index ffa76751713..17c88e32fb5 100644 --- a/iotdb-core/ainode/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py @@ -18,11 +18,11 @@ import torch -from ainode.core.exception import InferenceModelInternalError -from ainode.core.inference.strategy.abstract_inference_pipeline import ( +from iotdb.ainode.core.exception import InferenceModelInternalError +from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import ( AbstractInferencePipeline, ) -from ainode.core.model.sundial.configuration_sundial import SundialConfig +from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig class TimerSundialInferencePipeline(AbstractInferencePipeline): diff --git a/iotdb-core/ainode/ainode/core/inference/strategy/timerxl_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py similarity index 71% rename from iotdb-core/ainode/ainode/core/inference/strategy/timerxl_inference_pipeline.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py index 4add928cb38..dc1dd304f68 100644 --- a/iotdb-core/ainode/ainode/core/inference/strategy/timerxl_inference_pipeline.py +++ b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py @@ -18,11 +18,11 @@ import torch -from ainode.core.exception import InferenceModelInternalError -from ainode.core.inference.strategy.abstract_inference_pipeline import ( +from iotdb.ainode.core.exception import InferenceModelInternalError +from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import ( AbstractInferencePipeline, ) -from ainode.core.model.timerxl.configuration_timer import TimerConfig +from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig class TimerXLInferencePipeline(AbstractInferencePipeline): @@ -39,13 +39,13 @@ class TimerXLInferencePipeline(AbstractInferencePipeline): raise InferenceModelInternalError( f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}" ) - # TODO: Disassemble and adapt with TimerXL's ts_generation_mixin.py + # Considering that we are currently using the generate function interface, it seems that no pre-processing is required return inputs def post_decode(self): - # TODO: Disassemble and adapt with TimerXL's ts_generation_mixin.py + # Considering that we are currently using the generate function interface, it seems that no post-processing is required pass def post_inference(self): - # TODO: Disassemble and adapt with TimerXL's ts_generation_mixin.py + # Considering that we are currently using the generate function interface, it seems that no post-processing is required pass diff --git a/iotdb-core/ainode/ainode/core/inference/utils.py b/iotdb-core/ainode/iotdb/ainode/core/inference/utils.py similarity index 100% rename from iotdb-core/ainode/ainode/core/inference/utils.py rename to iotdb-core/ainode/iotdb/ainode/core/inference/utils.py diff --git a/iotdb-core/ainode/ainode/core/model/sundial/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/sundial/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/ingress/__init__.py diff --git a/iotdb-core/ainode/ainode/core/ingress/dataset.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py similarity index 100% rename from iotdb-core/ainode/ainode/core/ingress/dataset.py rename to iotdb-core/ainode/iotdb/ainode/core/ingress/dataset.py diff --git a/iotdb-core/ainode/ainode/core/ingress/iotdb.py b/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py similarity index 99% rename from iotdb-core/ainode/ainode/core/ingress/iotdb.py rename to iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py index 2d1bb74016f..792d64b2980 100644 --- a/iotdb-core/ainode/ainode/core/ingress/iotdb.py +++ b/iotdb-core/ainode/iotdb/ainode/core/ingress/iotdb.py @@ -23,8 +23,8 @@ from iotdb.utils.Field import Field from iotdb.utils.IoTDBConstants import TSDataType from torch.utils.data import Dataset -from ainode.core.config import AINodeDescriptor -from ainode.core.ingress.dataset import BasicDatabaseForecastDataset +from iotdb.ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.ingress import BasicDatabaseForecastDataset from ainode.core.log import Logger from ainode.core.util.cache import MemoryLRUCache from ainode.core.util.decorator import singleton diff --git a/iotdb-core/ainode/ainode/core/log.py b/iotdb-core/ainode/iotdb/ainode/core/log.py similarity index 100% rename from iotdb-core/ainode/ainode/core/log.py rename to iotdb-core/ainode/iotdb/ainode/core/log.py diff --git a/iotdb-core/ainode/ainode/core/model/timerxl/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/timerxl/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/manager/__init__.py diff --git a/iotdb-core/ainode/ainode/core/manager/cluster_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py similarity index 100% rename from iotdb-core/ainode/ainode/core/manager/cluster_manager.py rename to iotdb-core/ainode/iotdb/ainode/core/manager/cluster_manager.py diff --git a/iotdb-core/ainode/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py similarity index 97% rename from iotdb-core/ainode/ainode/core/manager/inference_manager.py rename to iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py index fe7adbc2396..65f4abacfca 100644 --- a/iotdb-core/ainode/ainode/core/manager/inference_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py @@ -15,21 +15,19 @@ # specific language governing permissions and limitations # under the License. # -import gc import threading import time from abc import ABC, abstractmethod from typing import Dict, List import pandas as pd -import psutil import torch import torch.multiprocessing as mp from iotdb.tsfile.utils.tsblock_serde import deserialize -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import TSStatusCode -from ainode.core.exception import ( +from iotdb.ainode.core.exception import ( InferenceModelInternalError, InvalidWindowArgumentError, NumericalRangeException, @@ -39,11 +37,11 @@ from ainode.core.inference.inference_request import ( InferenceRequest, InferenceRequestProxy, ) -from ainode.core.inference.inference_request_pool import InferenceRequestPool -from ainode.core.inference.strategy.timer_sundial_inference_pipeline import ( +from iotdb.ainode.core.inference.inference_request_pool import InferenceRequestPool +from iotdb.ainode.core.inference.strategy import ( TimerSundialInferencePipeline, ) -from ainode.core.inference.strategy.timerxl_inference_pipeline import ( +from iotdb.ainode.core.inference.strategy import ( TimerXLInferencePipeline, ) from ainode.core.inference.utils import generate_req_id @@ -56,9 +54,9 @@ from ainode.core.manager.utils import ( from ainode.core.model.sundial.configuration_sundial import SundialConfig from ainode.core.model.sundial.modeling_sundial import SundialForPrediction from ainode.core.model.timerxl.configuration_timer import TimerConfig -from ainode.core.model.timerxl.modeling_timer import TimerForPrediction +from iotdb.ainode.core.model.timerxl.modeling_timer import TimerForPrediction from ainode.core.rpc.status import get_status -from ainode.core.util.serde import convert_to_binary +from iotdb.ainode.core.util.serde import convert_to_binary from ainode.thrift.ainode.ttypes import ( TForecastReq, TForecastResp, diff --git a/iotdb-core/ainode/ainode/core/manager/model_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py similarity index 97% rename from iotdb-core/ainode/ainode/core/manager/model_manager.py rename to iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py index 69b163afd97..818311ed57f 100644 --- a/iotdb-core/ainode/ainode/core/manager/model_manager.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py @@ -21,12 +21,12 @@ from torch import nn from yaml import YAMLError from ainode.core.constant import TSStatusCode -from ainode.core.exception import ( +from iotdb.ainode.core.exception import ( BadConfigValueError, InvalidUriError, ) from ainode.core.log import Logger -from ainode.core.model.model_enums import BuiltInModelType, ModelStates +from iotdb.ainode.core.model.model_enums import BuiltInModelType, ModelStates from ainode.core.model.model_info import ModelInfo from ainode.core.model.model_storage import ModelStorage from ainode.core.rpc.status import get_status diff --git a/iotdb-core/ainode/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py similarity index 98% rename from iotdb-core/ainode/ainode/core/manager/utils.py rename to iotdb-core/ainode/iotdb/ainode/core/manager/utils.py index 5fbd444c38c..65106e08e08 100644 --- a/iotdb-core/ainode/ainode/core/manager/utils.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py @@ -21,7 +21,7 @@ import gc import psutil import torch -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.log import Logger from ainode.core.manager.model_manager import ModelManager from ainode.core.model.model_info import BUILT_IN_LTSM_MAP diff --git a/iotdb-core/ainode/ainode/core/rpc/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/rpc/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/model/__init__.py diff --git a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py b/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py similarity index 99% rename from iotdb-core/ainode/ainode/core/model/built_in_model_factory.py rename to iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py index f1bf4952fe2..a95af34990b 100644 --- a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/built_in_model_factory.py @@ -29,13 +29,13 @@ from sktime.forecasting.exp_smoothing import ExponentialSmoothing from sktime.forecasting.naive import NaiveForecaster from sktime.forecasting.trend import STLForecaster -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import ( MODEL_CONFIG_FILE_IN_JSON, MODEL_WEIGHTS_FILE_IN_SAFETENSORS, AttributeName, ) -from ainode.core.exception import ( +from iotdb.ainode.core.exception import ( BuiltInModelNotSupportError, InferenceModelInternalError, ListRangeException, @@ -44,10 +44,10 @@ from ainode.core.exception import ( WrongAttributeTypeError, ) from ainode.core.log import Logger -from ainode.core.model.model_enums import BuiltInModelType +from iotdb.ainode.core.model.model_enums import BuiltInModelType from ainode.core.model.model_info import TIMER_REPO_ID from ainode.core.model.sundial import modeling_sundial -from ainode.core.model.timerxl import modeling_timer +from iotdb.ainode.core.model.timerxl import modeling_timer logger = Logger() diff --git a/iotdb-core/ainode/ainode/core/model/model_enums.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_enums.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/model_enums.py rename to iotdb-core/ainode/iotdb/ainode/core/model/model_enums.py diff --git a/iotdb-core/ainode/ainode/core/model/model_factory.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py similarity index 98% rename from iotdb-core/ainode/ainode/core/model/model_factory.py rename to iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py index 5bbcf321644..c90222108ad 100644 --- a/iotdb-core/ainode/ainode/core/model/model_factory.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_factory.py @@ -26,15 +26,15 @@ from ainode.core.constant import ( MODEL_CONFIG_FILE_IN_YAML, MODEL_WEIGHTS_FILE_IN_PT, ) -from ainode.core.exception import BadConfigValueError, InvalidUriError +from iotdb.ainode.core.exception import BadConfigValueError, InvalidUriError from ainode.core.log import Logger -from ainode.core.model.model_enums import ModelFileType +from iotdb.ainode.core.model.model_enums import ModelFileType from ainode.core.model.uri_utils import ( UriType, download_file, download_snapshot_from_hf, ) -from ainode.core.util.serde import get_data_type_byte_from_str +from iotdb.ainode.core.util.serde import get_data_type_byte_from_str from ainode.thrift.ainode.ttypes import TConfigs logger = Logger() diff --git a/iotdb-core/ainode/ainode/core/model/model_info.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py similarity index 99% rename from iotdb-core/ainode/ainode/core/model/model_info.py rename to iotdb-core/ainode/iotdb/ainode/core/model/model_info.py index 6226179dbcb..7cf15e68225 100644 --- a/iotdb-core/ainode/ainode/core/model/model_info.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py @@ -24,7 +24,7 @@ from ainode.core.constant import ( MODEL_WEIGHTS_FILE_IN_PT, MODEL_WEIGHTS_FILE_IN_SAFETENSORS, ) -from ainode.core.model.model_enums import ( +from iotdb.ainode.core.model.model_enums import ( BuiltInModelType, ModelCategory, ModelFileType, diff --git a/iotdb-core/ainode/ainode/core/model/model_storage.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py similarity index 98% rename from iotdb-core/ainode/ainode/core/model/model_storage.py rename to iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py index ebfa9211811..93f695805dd 100644 --- a/iotdb-core/ainode/ainode/core/model/model_storage.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py @@ -26,13 +26,13 @@ from typing import Dict import torch from torch import nn -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import ( MODEL_CONFIG_FILE_IN_JSON, MODEL_WEIGHTS_FILE_IN_PT, TSStatusCode, ) -from ainode.core.exception import ( +from iotdb.ainode.core.exception import ( BuiltInModelDeletionError, ModelNotExistError, UnsupportedError, @@ -42,13 +42,13 @@ from ainode.core.model.built_in_model_factory import ( download_built_in_ltsm_from_hf_if_necessary, fetch_built_in_model, ) -from ainode.core.model.model_enums import ( +from iotdb.ainode.core.model.model_enums import ( BuiltInModelType, ModelCategory, ModelFileType, ModelStates, ) -from ainode.core.model.model_factory import fetch_model_by_uri +from iotdb.ainode.core.model.model_factory import fetch_model_by_uri from ainode.core.model.model_info import ( BUILT_IN_LTSM_MAP, BUILT_IN_MACHINE_LEARNING_MODEL_MAP, diff --git a/iotdb-core/ainode/ainode/core/util/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/model/sundial/__init__.py diff --git a/iotdb-core/ainode/ainode/core/model/sundial/configuration_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/configuration_sundial.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/sundial/configuration_sundial.py rename to iotdb-core/ainode/iotdb/ainode/core/model/sundial/configuration_sundial.py diff --git a/iotdb-core/ainode/ainode/core/model/sundial/flow_loss.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/sundial/flow_loss.py rename to iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py diff --git a/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py similarity index 99% rename from iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py rename to iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py index 2e6b436fbda..f66b4b13f22 100644 --- a/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py @@ -16,13 +16,10 @@ # under the License. # -import os from typing import List, Optional, Tuple, Union import torch import torch.nn.functional as F -from huggingface_hub import hf_hub_download -from safetensors.torch import load_file as load_safetensors from torch import nn from transformers import Cache, DynamicCache, PreTrainedModel from transformers.activations import ACT2FN @@ -35,7 +32,7 @@ from transformers.modeling_outputs import ( from ainode.core.log import Logger from ainode.core.model.sundial.configuration_sundial import SundialConfig from ainode.core.model.sundial.flow_loss import FlowLoss -from ainode.core.model.sundial.ts_generation_mixin import TSGenerationMixin +from iotdb.ainode.core.model.sundial.ts_generation_mixin import TSGenerationMixin logger = Logger() diff --git a/iotdb-core/ainode/ainode/core/model/sundial/ts_generation_mixin.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/ts_generation_mixin.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/sundial/ts_generation_mixin.py rename to iotdb-core/ainode/iotdb/ainode/core/model/sundial/ts_generation_mixin.py diff --git a/iotdb-core/ainode/ainode/core/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/__init__.py similarity index 100% copy from iotdb-core/ainode/ainode/core/__init__.py copy to iotdb-core/ainode/iotdb/ainode/core/model/timerxl/__init__.py diff --git a/iotdb-core/ainode/ainode/core/model/timerxl/configuration_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/configuration_timer.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/timerxl/configuration_timer.py rename to iotdb-core/ainode/iotdb/ainode/core/model/timerxl/configuration_timer.py diff --git a/iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py rename to iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py diff --git a/iotdb-core/ainode/ainode/core/model/timerxl/ts_generation_mixin.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/ts_generation_mixin.py similarity index 100% rename from iotdb-core/ainode/ainode/core/model/timerxl/ts_generation_mixin.py rename to iotdb-core/ainode/iotdb/ainode/core/model/timerxl/ts_generation_mixin.py diff --git a/iotdb-core/ainode/ainode/core/model/uri_utils.py b/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py similarity index 97% rename from iotdb-core/ainode/ainode/core/model/uri_utils.py rename to iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py index de2fb2a1848..59f03057e85 100644 --- a/iotdb-core/ainode/ainode/core/model/uri_utils.py +++ b/iotdb-core/ainode/iotdb/ainode/core/model/uri_utils.py @@ -28,9 +28,9 @@ from ainode.core.constant import ( DEFAULT_RECONNECT_TIMEOUT, DEFAULT_RECONNECT_TIMES, ) -from ainode.core.exception import UnsupportedError +from iotdb.ainode.core.exception import UnsupportedError from ainode.core.log import Logger -from ainode.core.model.model_enums import ModelFileType +from iotdb.ainode.core.model.model_enums import ModelFileType from ainode.core.model.model_info import get_model_file_type HTTP_PREFIX = "http://"; diff --git a/iotdb-core/ainode/ainode/core/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/__init__.py similarity index 100% copy from iotdb-core/ainode/ainode/core/__init__.py copy to iotdb-core/ainode/iotdb/ainode/core/rpc/__init__.py diff --git a/iotdb-core/ainode/ainode/core/rpc/client.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py similarity index 99% rename from iotdb-core/ainode/ainode/core/rpc/client.py rename to iotdb-core/ainode/iotdb/ainode/core/rpc/client.py index c595000bb0f..dc9d84de9ae 100644 --- a/iotdb-core/ainode/ainode/core/rpc/client.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/client.py @@ -21,7 +21,7 @@ from thrift.protocol import TBinaryProtocol, TCompactProtocol from thrift.Thrift import TException from thrift.transport import TSocket, TTransport -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import TSStatusCode from ainode.core.log import Logger from ainode.core.rpc.status import verify_success diff --git a/iotdb-core/ainode/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py similarity index 100% rename from iotdb-core/ainode/ainode/core/rpc/handler.py rename to iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py diff --git a/iotdb-core/ainode/ainode/core/rpc/service.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py similarity index 98% rename from iotdb-core/ainode/ainode/core/rpc/service.py rename to iotdb-core/ainode/iotdb/ainode/core/rpc/service.py index b97abeee65b..8642f62fc20 100644 --- a/iotdb-core/ainode/ainode/core/rpc/service.py +++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/service.py @@ -21,7 +21,7 @@ from thrift.protocol import TBinaryProtocol, TCompactProtocol from thrift.server import TServer from thrift.transport import TSocket, TTransport -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.log import Logger from ainode.core.rpc.handler import AINodeRPCServiceHandler from ainode.thrift.ainode import IAINodeRPCService diff --git a/iotdb-core/ainode/ainode/core/rpc/status.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/status.py similarity index 100% rename from iotdb-core/ainode/ainode/core/rpc/status.py rename to iotdb-core/ainode/iotdb/ainode/core/rpc/status.py diff --git a/iotdb-core/ainode/ainode/core/script.py b/iotdb-core/ainode/iotdb/ainode/core/script.py similarity index 95% rename from iotdb-core/ainode/ainode/core/script.py rename to iotdb-core/ainode/iotdb/ainode/core/script.py index a01ce691dc7..dd3efa11ac8 100644 --- a/iotdb-core/ainode/ainode/core/script.py +++ b/iotdb-core/ainode/iotdb/ainode/core/script.py @@ -21,10 +21,10 @@ import sys import torch.multiprocessing as mp -from ainode.core.ai_node import AINode -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.ai_node import AINode +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.constant import TSStatusCode -from ainode.core.exception import MissingConfigError +from iotdb.ainode.core.exception import MissingConfigError from ainode.core.log import Logger from ainode.core.rpc.client import ClientManager from ainode.thrift.common.ttypes import ( diff --git a/iotdb-core/ainode/ainode/core/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/util/__init__.py similarity index 100% rename from iotdb-core/ainode/ainode/core/__init__.py rename to iotdb-core/ainode/iotdb/ainode/core/util/__init__.py diff --git a/iotdb-core/ainode/ainode/core/util/activation.py b/iotdb-core/ainode/iotdb/ainode/core/util/activation.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/activation.py rename to iotdb-core/ainode/iotdb/ainode/core/util/activation.py diff --git a/iotdb-core/ainode/ainode/core/util/cache.py b/iotdb-core/ainode/iotdb/ainode/core/util/cache.py similarity index 98% rename from iotdb-core/ainode/ainode/core/util/cache.py rename to iotdb-core/ainode/iotdb/ainode/core/util/cache.py index e0cfd588316..40e2b8d5771 100644 --- a/iotdb-core/ainode/ainode/core/util/cache.py +++ b/iotdb-core/ainode/iotdb/ainode/core/util/cache.py @@ -18,7 +18,7 @@ import sys from collections import OrderedDict -from ainode.core.config import AINodeDescriptor +from iotdb.ainode.core.config import AINodeDescriptor from ainode.core.util.decorator import singleton diff --git a/iotdb-core/ainode/ainode/core/util/decorator.py b/iotdb-core/ainode/iotdb/ainode/core/util/decorator.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/decorator.py rename to iotdb-core/ainode/iotdb/ainode/core/util/decorator.py diff --git a/iotdb-core/ainode/ainode/core/util/huggingface_cache.py b/iotdb-core/ainode/iotdb/ainode/core/util/huggingface_cache.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/huggingface_cache.py rename to iotdb-core/ainode/iotdb/ainode/core/util/huggingface_cache.py diff --git a/iotdb-core/ainode/ainode/core/util/lock.py b/iotdb-core/ainode/iotdb/ainode/core/util/lock.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/lock.py rename to iotdb-core/ainode/iotdb/ainode/core/util/lock.py diff --git a/iotdb-core/ainode/ainode/core/util/masking.py b/iotdb-core/ainode/iotdb/ainode/core/util/masking.py similarity index 100% rename from iotdb-core/ainode/ainode/core/util/masking.py rename to iotdb-core/ainode/iotdb/ainode/core/util/masking.py diff --git a/iotdb-core/ainode/ainode/core/util/serde.py b/iotdb-core/ainode/iotdb/ainode/core/util/serde.py similarity index 98% rename from iotdb-core/ainode/ainode/core/util/serde.py rename to iotdb-core/ainode/iotdb/ainode/core/util/serde.py index affd96992ee..f8188209a37 100644 --- a/iotdb-core/ainode/ainode/core/util/serde.py +++ b/iotdb-core/ainode/iotdb/ainode/core/util/serde.py @@ -21,7 +21,7 @@ from enum import Enum import numpy as np import pandas as pd -from ainode.core.exception import BadConfigValueError +from iotdb.ainode.core.exception import BadConfigValueError class TSDataType(Enum):
