Merge branch 'master' into strong_password

Author: jt2594838

integration-test/src/test/java/org/apache/iotdb/db/it/query/IoTDBCaseWhenThenIT.java

@@ -919,4 +919,13 @@ public void testMultipleSatisfyCase() {
         };
     resultSetEqualTest(sql, expectedHeader, retArray);
   }
+
+  @Test
+  public void testThenWithBinarySameConstant() {
+    String sql = "SELECT CASE WHEN true THEN 200 + (s1 - 200) END AS result FROM root.sg.d1";
+    String[] expectedHeader = new String[] {TIMESTAMP_STR, "result"};
+    String[] retArray =
+        new String[] {"0,0.0,", "1000000,11.0,", "20000000,22.0,", "210000000,33.0,"};
+    resultSetEqualTest(sql, expectedHeader, retArray);
+  }
 }

integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/basic/IoTDBPipeExtractorIT.java

@@ -469,7 +469,7 @@ public void testExtractorTimeRangeMatch() throws Exception {
           TestUtils.executeNonQueryWithRetry(receiverEnv, "flush");
         };
 
-    boolean insertResult = true;
+    boolean insertResult;
 
     try (final SyncConfigNodeIServiceClient client =
         (SyncConfigNodeIServiceClient) senderEnv.getLeaderConfigNodeConnection()) {

integration-test/src/test/java/org/apache/iotdb/pipe/it/dual/tablemodel/manual/basic/IoTDBPipePermissionIT.java

@@ -211,7 +211,11 @@ public void testSourcePermission() {
       return;
     }
 
-    TableModelUtils.createDataBaseAndTable(receiverEnv, "test", "test");
+    try {
+      TableModelUtils.createDataBaseAndTable(receiverEnv, "test", "test");
+    } catch (final Exception ignore) {
+      // Ignore because the db/table may be transferred because sender user may see these
+    }
 
     // Exception, block here
     TableModelUtils.assertCountDataAlwaysOnEnv("test", "test", 0, receiverEnv);

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBCaseWhenThenTableIT.java

@@ -673,4 +673,12 @@ public void testKind1Logic() {
     String[] retArray = new String[] {"null,", "just so so~~~,", "null,", "very well~~~,"};
     tableResultSetEqualTest(sql, expectedHeader, retArray, DATABASE);
   }
+
+  @Test
+  public void testThenWithBinarySameConstant() {
+    String sql = "SELECT CASE WHEN true THEN 200 + (s1 - 200) END AS result FROM table1";
+    String[] expectedHeader = new String[] {"result"};
+    String[] retArray = new String[] {"0,", "11,", "22,", "33,"};
+    tableResultSetEqualTest(sql, expectedHeader, retArray, DATABASE);
+  }
 }

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

@@ -134,7 +134,7 @@ def start(self):
                 raise e
 
         # Start the RPC service
-        self._rpc_handler = AINodeRPCServiceHandler(aiNode=self)
+        self._rpc_handler = AINodeRPCServiceHandler(ainode=self)
         self._rpc_service = AINodeRPCService(self._rpc_handler)
         self._rpc_service.start()
         self._rpc_service.join(1)

iotdb-core/ainode/ainode/core/config.py

@@ -31,6 +31,7 @@
     AINODE_CONF_GIT_FILE_NAME,
     AINODE_CONF_POM_FILE_NAME,
     AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
+    AINODE_INFERENCE_MAX_PREDICT_LENGTH,
     AINODE_LOG_DIR,
     AINODE_MODELS_DIR,
     AINODE_ROOT_CONF_DIRECTORY_NAME,
@@ -72,6 +73,9 @@ def __init__(self):
         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
+        )
 
         # log directory
         self._ain_logs_dir: str = AINODE_LOG_DIR
@@ -140,6 +144,14 @@ def set_ain_inference_batch_interval_in_ms(
     ) -> 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_logs_dir(self) -> str:
         return self._ain_logs_dir
 

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

@@ -50,6 +50,7 @@
 
 # AINode inference configuration
 AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
+AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
 
 # AINode folder structure
 AINODE_ROOT_DIR = os.path.dirname(

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

@@ -23,7 +23,7 @@
 import numpy as np
 import torch
 import torch.multiprocessing as mp
-from transformers import PretrainedConfig, PreTrainedModel
+from transformers import PretrainedConfig
 
 from ainode.core.config import AINodeDescriptor
 from ainode.core.inference.inference_request import InferenceRequest
@@ -46,7 +46,7 @@ class InferenceRequestPool(mp.Process):
     def __init__(
         self,
         pool_id: int,
-        model_id: int,
+        model_id: str,
         config: PretrainedConfig,
         request_queue: mp.Queue,
         result_queue: mp.Queue,
@@ -58,6 +58,7 @@ def __init__(
         self.config = config
         self.pool_kwargs = pool_kwargs
         self.model = None
+        self._model_manager = None
         self.device = None
 
         # TODO: A scheduler is necessary for better handling following queues

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

@@ -22,7 +22,7 @@
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 
 
-def _generate_req_id(length=10, charset=string.ascii_letters + string.digits) -> str:
+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.

iotdb-core/ainode/ainode/core/manager/inference_manager.py

@@ -30,6 +30,7 @@
 from ainode.core.exception import (
     InferenceModelInternalError,
     InvalidWindowArgumentError,
+    NumericalRangeException,
     runtime_error_extractor,
 )
 from ainode.core.inference.inference_request import (
@@ -40,7 +41,7 @@
 from ainode.core.inference.strategy.timer_sundial_inference_pipeline import (
     TimerSundialInferencePipeline,
 )
-from ainode.core.inference.utils import _generate_req_id
+from ainode.core.inference.utils import generate_req_id
 from ainode.core.log import Logger
 from ainode.core.manager.model_manager import ModelManager
 from ainode.core.model.sundial.configuration_sundial import SundialConfig
@@ -214,6 +215,20 @@ def _run(
             full_data = deserializer(raw)
             inference_attrs = extract_attrs(req)
 
+            predict_length = inference_attrs.get("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 model_id == self.ACCELERATE_MODEL_ID and self.DEFAULT_POOL_SIZE > 0:
                 # TODO: Logic in this branch shall handle all LTSM inferences
                 # TODO: TSBlock -> Tensor codes should be unified
@@ -223,10 +238,10 @@ def _run(
                 # the inputs should be on CPU before passing to the inference request
                 inputs = torch.tensor(data).unsqueeze(0).float().to("cpu")
                 infer_req = InferenceRequest(
-                    req_id=_generate_req_id(),
+                    req_id=generate_req_id(),
                     inputs=inputs,
                     inference_pipeline=TimerSundialInferencePipeline(SundialConfig()),
-                    max_new_tokens=inference_attrs.get("predict_length", 96),
+                    max_new_tokens=predict_length,
                 )
                 infer_proxy = InferenceRequestProxy(infer_req.req_id)
                 with self._result_wrapper_lock: