stash changes

Author: CRZbulabula

iotdb-core/ainode/ainode/core/constant.py

@@ -31,23 +31,23 @@
 AINODE_SYSTEM_FILE_NAME = "system.properties"
 # inference_rpc_address
 AINODE_INFERENCE_RPC_ADDRESS = "127.0.0.1"
-AINODE_INFERENCE_RPC_PORT = 10810
+AINODE_INFERENCE_RPC_PORT = 11810
 # AINode folder structure
 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"
 AINODE_THRIFT_COMPRESSION_ENABLED = False
 # use for node management
-AINODE_CLUSTER_NAME = "defaultCluster"
+AINODE_CLUSTER_NAME = "yongzaoCluster"
 AINODE_VERSION_INFO = "UNKNOWN"
 AINODE_BUILD_INFO = "UNKNOWN"
 AINODE_ROOT_DIR = os.path.dirname(
     os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
 )
 # connect IoTDB cluster
 AINODE_CLUSTER_INGRESS_ADDRESS = "127.0.0.1"
-AINODE_CLUSTER_INGRESS_PORT = 6667
+AINODE_CLUSTER_INGRESS_PORT = 7667
 AINODE_CLUSTER_INGRESS_USERNAME = "root"
 AINODE_CLUSTER_INGRESS_PASSWORD = "root"
 AINODE_CLUSTER_INGRESS_TIME_ZONE = "UTC+8"

iotdb-core/ainode/ainode/core/inference/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+#

iotdb-core/ainode/ainode/core/inference/inference_request.py

@@ -0,0 +1,108 @@
+# 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 ainode.core.inference.strategy.abstract_strategy import AbstractStrategy
+
+
+class InferenceRequestState:
+    WAITING = "waiting"
+    RUNNING = "running"
+    FINISHED = "finished"
+
+
+class InferenceRequest:
+    def __init__(
+        self,
+        req_id: int,
+        inputs: torch.Tensor,
+        strategy: AbstractStrategy,
+        max_new_tokens: int = 96,
+        **infer_kwargs,
+    ):
+        if inputs.ndim == 1:
+            inputs = inputs.unsqueeze(0)
+
+        self.id = req_id
+        self.inputs = inputs
+        self.infer_kwargs = infer_kwargs
+        self.strategy = strategy
+        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
+
+        # Preallocate output buffer [batch_size, max_new_tokens]
+        device = inputs.device
+        self.output_tensor = torch.zeros(
+            self.batch_size, max_new_tokens, device=device
+        )  # shape: [self.batch_size, max_new_steps]
+
+        self._lock = threading.Lock()
+        self._condition = threading.Condition(self._lock)
+
+    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):
+        with self._lock:
+            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:
+        with self._lock:
+            return self.output_tensor[:, : self.cur_step_idx]
+
+    def notify_completion(self):
+        with self._lock:
+            self._condition.notify_all()
+
+    def wait_for_completion(self) -> Any:
+        with self._lock:
+            while self.state != InferenceRequestState.FINISHED:
+                self._condition.wait()