This is an automated email from the ASF dual-hosted git repository. liyuheng pushed a commit to branch lyh/ainode/logger-refactor in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 70a17fc5b56f143e2a0343fe91eba005c88430f7 Author: liyuheng <[email protected]> AuthorDate: Wed Jul 30 19:17:51 2025 +0800 save --- .../ainode/ainode/core/{ainode.py => ai_node.py} | 0 .../ainode/core/inference/inference_request.py | 4 ++++ .../ainode/core/manager/inference_manager.py | 25 ++++++++++++++++++++++ iotdb-core/ainode/ainode/core/script.py | 2 +- .../ainode/resources/conf/iotdb-ainode.properties | 2 +- 5 files changed, 31 insertions(+), 2 deletions(-) diff --git a/iotdb-core/ainode/ainode/core/ainode.py b/iotdb-core/ainode/ainode/core/ai_node.py similarity index 100% rename from iotdb-core/ainode/ainode/core/ainode.py rename to iotdb-core/ainode/ainode/core/ai_node.py diff --git a/iotdb-core/ainode/ainode/core/inference/inference_request.py b/iotdb-core/ainode/ainode/core/inference/inference_request.py index 2c45826fd26..714bde05bf0 100644 --- a/iotdb-core/ainode/ainode/core/inference/inference_request.py +++ b/iotdb-core/ainode/ainode/core/inference/inference_request.py @@ -63,11 +63,15 @@ class InferenceRequest: self.batch_size, max_new_tokens, device="cpu" ) # shape: [self.batch_size, max_new_steps] + self.logger = Logger() + def mark_running(self): self.state = InferenceRequestState.RUNNING + logger.info("mark_running") def mark_finished(self): self.state = InferenceRequestState.FINISHED + logger.info("mark_finished") def is_finished(self) -> bool: return ( diff --git a/iotdb-core/ainode/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/ainode/core/manager/inference_manager.py index 5a853ac4e72..f5751010e0a 100644 --- a/iotdb-core/ainode/ainode/core/manager/inference_manager.py +++ b/iotdb-core/ainode/ainode/core/manager/inference_manager.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. # +import os import threading import time from abc import ABC, abstractmethod @@ -72,6 +73,29 @@ class InferenceStrategy(ABC): # we only get valueList currently. class TimerXLStrategy(InferenceStrategy): def infer(self, full_data, predict_length=96, **_): + if torch.cuda.is_available(): + device = next(self.model.parameters()).device + else: + device = torch.device("cpu") + + # Get possible rank + if torch.distributed.is_initialized(): + global_rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + else: + # Not distribution, default rank=0, world_size=1 + global_rank, world_size = 0, 1 + + if device.type == "cuda": + gpu_name = torch.cuda.get_device_name(device.index) + logger.info( + f"[rank {global_rank}/{world_size}] " + f"Running on GPU {device.index} ({gpu_name})" + ) + else: + logger.info(f"[rank {global_rank}/{world_size}] Running on CPU") + + logger.info("Start inference") data = full_data[1][0] if data.dtype.byteorder not in ("=", "|"): data = data.byteswap().newbyteorder() @@ -79,6 +103,7 @@ class TimerXLStrategy(InferenceStrategy): # TODO: unify model inference input output = self.model.generate(seqs, max_new_tokens=predict_length, revin=True) df = pd.DataFrame(output[0]) + logger.info("Complete inference") return convert_to_binary(df) diff --git a/iotdb-core/ainode/ainode/core/script.py b/iotdb-core/ainode/ainode/core/script.py index 8cccfec5d17..a01ce691dc7 100644 --- a/iotdb-core/ainode/ainode/core/script.py +++ b/iotdb-core/ainode/ainode/core/script.py @@ -21,7 +21,7 @@ import sys import torch.multiprocessing as mp -from ainode.core.ainode import AINode +from ainode.core.ai_node import AINode from ainode.core.config import AINodeDescriptor from ainode.core.constant import TSStatusCode from ainode.core.exception import MissingConfigError diff --git a/iotdb-core/ainode/resources/conf/iotdb-ainode.properties b/iotdb-core/ainode/resources/conf/iotdb-ainode.properties index b31079ce00d..6c8c28e3f26 100644 --- a/iotdb-core/ainode/resources/conf/iotdb-ainode.properties +++ b/iotdb-core/ainode/resources/conf/iotdb-ainode.properties @@ -19,7 +19,7 @@ # Used for indicate cluster name and distinguish different cluster. # Datatype: string -cluster_name=defaultCluster +cluster_name=???? # ConfigNode address registered at AINode startup. # Allow modifications only before starting the service for the first time.
