You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by li...@apache.org on 2025/07/30 11:19:12 UTC

(iotdb) branch lyh/ainode/logger-refactor created (now 70a17fc5b56)

This is an automated email from the ASF dual-hosted git repository.

liyuheng pushed a change to branch lyh/ainode/logger-refactor
in repository https://gitbox.apache.org/repos/asf/iotdb.git


      at 70a17fc5b56 save

This branch includes the following new commits:

     new 70a17fc5b56 save

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.



(iotdb) 01/01: save

Posted by li...@apache.org.
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 <li...@126.com>
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.