Merge branch 'master' of https://github.com/apache/iotdb into fix-audit-logger

Author: Caideyipi

integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentForecastIT.java

@@ -49,7 +49,7 @@ public class AINodeConcurrentForecastIT {
   private static final Logger LOGGER = LoggerFactory.getLogger(AINodeConcurrentForecastIT.class);
 
   private static final String FORECAST_TABLE_FUNCTION_SQL_TEMPLATE =
-      "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time, output_length=>%d)";
+      "SELECT * FROM FORECAST(model_id=>'%s', targets=>(SELECT time,s FROM root.AI) ORDER BY time, output_length=>%d)";
 
   @BeforeClass
   public static void setUp() throws Exception {

integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeForecastIT.java

@@ -38,13 +38,21 @@
 import java.sql.Statement;
 
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.BUILTIN_MODEL_MAP;
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.errorTest;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({AIClusterIT.class})
 public class AINodeForecastIT {
 
   private static final String FORECAST_TABLE_FUNCTION_SQL_TEMPLATE =
-      "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time, s%d FROM db.AI) ORDER BY time)";
+      "SELECT * FROM FORECAST("
+          + "model_id=>'%s', "
+          + "targets=>(SELECT time, s%d FROM db.AI WHERE time<%d ORDER BY time DESC LIMIT %d) ORDER BY time, "
+          + "output_start_time=>%d, "
+          + "output_length=>%d, "
+          + "output_interval=>%d, "
+          + "timecol=>'%s'"
+          + ")";
 
   @BeforeClass
   public static void setUp() throws Exception {
@@ -55,7 +63,7 @@ public static void setUp() throws Exception {
       statement.execute("CREATE DATABASE db");
       statement.execute(
           "CREATE TABLE db.AI (s0 FLOAT FIELD, s1 DOUBLE FIELD, s2 INT32 FIELD, s3 INT64 FIELD)");
-      for (int i = 0; i < 2880; i++) {
+      for (int i = 0; i < 5760; i++) {
         statement.execute(
             String.format(
                 "INSERT INTO db.AI(time,s0,s1,s2,s3) VALUES(%d,%f,%f,%d,%d)",
@@ -81,18 +89,100 @@ public void forecastTableFunctionTest() throws SQLException {
 
   public void forecastTableFunctionTest(
       Statement statement, AINodeTestUtils.FakeModelInfo modelInfo) throws SQLException {
-    // Invoke call inference for specified models, there should exist result.
+    // Invoke forecast table function for specified models, there should exist result.
     for (int i = 0; i < 4; i++) {
       String forecastTableFunctionSQL =
-          String.format(FORECAST_TABLE_FUNCTION_SQL_TEMPLATE, modelInfo.getModelId(), i);
+          String.format(
+              FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
+              modelInfo.getModelId(),
+              i,
+              5760,
+              2880,
+              5760,
+              96,
+              1,
+              "time");
       try (ResultSet resultSet = statement.executeQuery(forecastTableFunctionSQL)) {
         int count = 0;
         while (resultSet.next()) {
           count++;
         }
-        // Ensure the call inference return results
+        // Ensure the forecast sentence return results
         Assert.assertTrue(count > 0);
       }
     }
   }
+
+  @Test
+  public void forecastTableFunctionErrorTest() throws SQLException {
+    for (AINodeTestUtils.FakeModelInfo modelInfo : BUILTIN_MODEL_MAP.values()) {
+      try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+          Statement statement = connection.createStatement()) {
+        forecastTableFunctionErrorTest(statement, modelInfo);
+      }
+    }
+  }
+
+  public void forecastTableFunctionErrorTest(
+      Statement statement, AINodeTestUtils.FakeModelInfo modelInfo) throws SQLException {
+    // OUTPUT_START_TIME error
+    String invalidOutputStartTimeSQL =
+        String.format(
+            FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
+            modelInfo.getModelId(),
+            0,
+            5760,
+            2880,
+            5759,
+            96,
+            1,
+            "time");
+    errorTest(
+        statement,
+        invalidOutputStartTimeSQL,
+        "701: The OUTPUT_START_TIME should be greater than the maximum timestamp of target time series. Expected greater than [5759] but found [5759].");
+
+    // OUTPUT_LENGTH error
+    String invalidOutputLengthSQL =
+        String.format(
+            FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
+            modelInfo.getModelId(),
+            0,
+            5760,
+            2880,
+            5760,
+            0,
+            1,
+            "time");
+    errorTest(statement, invalidOutputLengthSQL, "701: OUTPUT_LENGTH should be greater than 0");
+
+    // OUTPUT_INTERVAL error
+    String invalidOutputIntervalSQL =
+        String.format(
+            FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
+            modelInfo.getModelId(),
+            0,
+            5760,
+            2880,
+            5760,
+            96,
+            -1,
+            "time");
+    errorTest(statement, invalidOutputIntervalSQL, "701: OUTPUT_INTERVAL should be greater than 0");
+
+    // TIMECOL error
+    String invalidTimecolSQL2 =
+        String.format(
+            FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
+            modelInfo.getModelId(),
+            0,
+            5760,
+            2880,
+            5760,
+            96,
+            1,
+            "s0");
+    errorTest(
+        statement, invalidTimecolSQL2, "701: The type of the column [s0] is not as expected.");
+  }
 }

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

@@ -42,7 +42,7 @@ def __init__(
         output_length: int = 96,
         **infer_kwargs,
     ):
-        if inputs.ndim == 1:
+        while inputs.ndim < 3:
             inputs = inputs.unsqueeze(0)
 
         self.req_id = req_id
@@ -54,15 +54,16 @@ def __init__(
         )
 
         self.batch_size = inputs.size(0)
+        self.variable_size = inputs.size(1)
         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(
-            self.batch_size, output_length, device="cpu"
-        )  # shape: [self.batch_size, max_new_steps]
+            self.batch_size, self.variable_size, output_length, device="cpu"
+        )  # shape: [batch_size, target_count, predict_length]
 
     def mark_running(self):
         self.state = InferenceRequestState.RUNNING
@@ -77,26 +78,26 @@ def is_finished(self) -> bool:
         )
 
     def write_step_output(self, step_output: torch.Tensor):
-        if step_output.ndim == 1:
+        while step_output.ndim < 3:
             step_output = step_output.unsqueeze(0)
 
-        batch_size, step_size = step_output.shape
+        batch_size, variable_size, step_size = step_output.shape
         end_idx = self.cur_step_idx + step_size
 
         if end_idx > self.output_length:
-            self.output_tensor[:, self.cur_step_idx :] = step_output[
-                :, : self.output_length - self.cur_step_idx
+            self.output_tensor[:, :, self.cur_step_idx :] = step_output[
+                :, :, : self.output_length - self.cur_step_idx
             ]
             self.cur_step_idx = self.output_length
         else:
-            self.output_tensor[:, self.cur_step_idx : end_idx] = step_output
+            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:
-        return self.output_tensor[:, : self.cur_step_idx]
+        return self.output_tensor[:, :, : self.cur_step_idx]
 
 
 class InferenceRequestProxy:

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

@@ -123,6 +123,7 @@ def _step(self):
             batch_inputs = self._batcher.batch_request(requests).to(
                 "cpu"
             )  # The input data should first load to CPU in current version
+            batch_inputs = self._inference_pipeline.preprocess(batch_inputs)
             if isinstance(self._inference_pipeline, ForecastPipeline):
                 batch_output = self._inference_pipeline.forecast(
                     batch_inputs,
@@ -140,7 +141,10 @@ def _step(self):
                     # more infer kwargs can be added here
                 )
             else:
+                batch_output = None
                 self._logger.error("[Inference] Unsupported pipeline type.")
+            batch_output = self._inference_pipeline.postprocess(batch_output)
+
             offset = 0
             for request in requests:
                 request.output_tensor = request.output_tensor.to(self.device)

iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py

@@ -20,6 +20,7 @@
 
 import torch
 
+from iotdb.ainode.core.exception import InferenceModelInternalException
 from iotdb.ainode.core.model.model_loader import load_model
 
 
@@ -29,59 +30,75 @@ def __init__(self, model_info, **model_kwargs):
         self.device = model_kwargs.get("device", "cpu")
         self.model = load_model(model_info, device_map=self.device, **model_kwargs)
 
-    def _preprocess(self, inputs):
+    @abstractmethod
+    def preprocess(self, inputs):
         """
         Preprocess the input before inference, including shape validation and value transformation.
         """
-        return inputs
+        raise NotImplementedError("preprocess not implemented")
 
-    def _postprocess(self, output: torch.Tensor):
+    @abstractmethod
+    def postprocess(self, outputs: torch.Tensor):
         """
         Post-process the outputs after the entire inference task.
         """
-        return output
+        raise NotImplementedError("postprocess not implemented")
 
 
 class ForecastPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
+        """
+        The inputs should be 3D tensor: [batch_size, target_count, sequence_length].
+        """
+        if len(inputs.shape) != 3:
+            raise InferenceModelInternalException(
+                f"[Inference] Input must be: [batch_size, target_count, sequence_length], but receives {inputs.shape}"
+            )
         return inputs
 
     @abstractmethod
     def forecast(self, inputs, **infer_kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
-        return output
+    def postprocess(self, outputs: torch.Tensor):
+        """
+        The outputs should be 3D tensor: [batch_size, target_count, predict_length].
+        """
+        if len(outputs.shape) != 3:
+            raise InferenceModelInternalException(
+                f"[Inference] Output must be: [batch_size, target_count, predict_length], but receives {outputs.shape}"
+            )
+        return outputs
 
 
 class ClassificationPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     @abstractmethod
     def classify(self, inputs, **kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
-        return output
+    def postprocess(self, outputs: torch.Tensor):
+        return outputs
 
 
 class ChatPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     @abstractmethod
     def chat(self, inputs, **kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
-        return output
+    def postprocess(self, outputs: torch.Tensor):
+        return outputs