[AINode][To rc/2.0.6] cp AINode codes (#16910)

* [AINode] Refactor code base

* [AINode] Implement concurrent inference framework (#16311)

(cherry picked from commit 7b9ec7e)

* [AINode] Fix bugs for SHOW LOADED MODELS (#16410)

(cherry picked from commit 40b2b33)

* [AINode] Add a batcher for inference (#16411)

(cherry picked from commit 7734331)

* [AINode][Bug fix] Concurrent inference (#16518)

* trigger CI

* bug fix 4 show loaded models

(cherry picked from commit b4dde12)

* [AINode] Concurrent inference bug fix (#16595)

(cherry picked from commit 46a0c6a)

* [AINode] Adjust the maximum inference input length (#16640)

(cherry picked from commit 2c9064f)

* [AINode] Fix bug of sundial and forecast udf (#16768)

(cherry picked from commit 2b47be7)

* [AINode] Package AINode via PyInstaller (#16707)

(cherry picked from commit 49c625b)

* [AINode] Enable AINode start as background (-d) (#16762)

(cherry picked from commit 1ebb951)

* [AINode] Update AINodeClient for DataNode to borrow (#16647)

(cherry picked from commit d49d7dd)

* [AINode] Fix bug that AINode cannot compile in Windows (#16767)

(cherry picked from commit cd443ba)

* [AINode] Delete poetry.lock for easier maintain different operating systems (#16793)

(cherry picked from commit 50f92e4)

* [AINode] Fix cp errors

---------

Co-authored-by: Leo <[email protected]>
Co-authored-by: jtmer <[email protected]>
Co-authored-by: Zeyu Zhang <[email protected]>

Author: 4 people

.github/workflows/cluster-it-1c1d1a.yml

@@ -41,9 +41,6 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Build AINode
-        shell: bash
-        run: mvn clean package -DskipTests -P with-ainode
       - name: IT Test
         shell: bash
         run: |
@@ -59,5 +56,5 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: cluster-log-ainode-${{ matrix.os }}
-          path: integration-test/target/ainode-logs
+          path: integration-test/target/*-logs
           retention-days: 30

integration-test/src/assembly/mpp-test.xml

@@ -63,7 +63,7 @@
         </fileSet>
         <fileSet>
             <outputDirectory>lib</outputDirectory>
-            <directory>${project.basedir}/../iotdb-core/ainode/dist/</directory>
+            <directory>${project.basedir}/../iotdb-core/ainode/dist/ainode/</directory>
             <fileMode>0755</fileMode>
         </fileSet>
     </fileSets>

integration-test/src/main/java/org/apache/iotdb/it/env/cluster/node/AINodeWrapper.java

@@ -25,9 +25,9 @@
 import org.apache.commons.io.file.PathUtils;
 import org.slf4j.Logger;
 
-import java.io.BufferedWriter;
 import java.io.File;
-import java.io.FileWriter;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.LinkOption;
@@ -37,6 +37,7 @@
 import java.nio.file.StandardCopyOption;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Properties;
 import java.util.stream.Stream;
 
 import static org.apache.iotdb.it.env.cluster.ClusterConstant.AI_NODE_NAME;
@@ -59,18 +60,22 @@ public class AINodeWrapper extends AbstractNodeWrapper {
   public static final String CONFIG_PATH = "conf";
   public static final String SCRIPT_PATH = "sbin";
   public static final String BUILT_IN_MODEL_PATH = "data/ainode/models/weights";
-  public static final String CACHE_BUILT_IN_MODEL_PATH = "/tmp/data/ainode/models/weights";
+  public static final String CACHE_BUILT_IN_MODEL_PATH = "/data/ainode/models/weights";
 
   private void replaceAttribute(String[] keys, String[] values, String filePath) {
-    try (BufferedWriter writer = new BufferedWriter(new FileWriter(filePath, true))) {
-      for (int i = 0; i < keys.length; i++) {
-        String line = keys[i] + "=" + values[i];
-        writer.newLine();
-        writer.write(line);
-      }
+    Properties props = new Properties();
+    try (FileInputStream in = new FileInputStream(filePath)) {
+      props.load(in);
+    } catch (IOException e) {
+      logger.warn("Failed to load existing AINode properties from {}, because: ", filePath, e);
+    }
+    for (int i = 0; i < keys.length; i++) {
+      props.setProperty(keys[i], values[i]);
+    }
+    try (FileOutputStream out = new FileOutputStream(filePath)) {
+      props.store(out, "Updated by AINode integration-test env");
     } catch (IOException e) {
-      logger.error(
-          "Failed to set attribute for AINode in file: {} because {}", filePath, e.getMessage());
+      logger.error("Failed to save properties to {}, because:", filePath, e);
     }
   }
 

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

@@ -0,0 +1,181 @@
+/*
+ * 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.
+ */
+
+package org.apache.iotdb.ainode.it;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.AIClusterIT;
+import org.apache.iotdb.itbase.env.BaseEnv;
+
+import com.google.common.collect.ImmutableSet;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.concurrentInference;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({AIClusterIT.class})
+public class AINodeConcurrentInferenceIT {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(AINodeConcurrentInferenceIT.class);
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    // Init 1C1D1A cluster environment
+    EnvFactory.getEnv().initClusterEnvironment(1, 1);
+    prepareDataForTreeModel();
+    prepareDataForTableModel();
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  private static void prepareDataForTreeModel() throws SQLException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE root.AI");
+      statement.execute("CREATE TIMESERIES root.AI.s WITH DATATYPE=DOUBLE, ENCODING=RLE");
+      for (int i = 0; i < 2880; i++) {
+        statement.execute(
+            String.format(
+                "INSERT INTO root.AI(timestamp, s) VALUES(%d, %f)",
+                i, Math.sin(i * Math.PI / 1440)));
+      }
+    }
+  }
+
+  private static void prepareDataForTableModel() throws SQLException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      statement.execute("CREATE DATABASE root");
+      statement.execute("CREATE TABLE root.AI (s DOUBLE FIELD)");
+      for (int i = 0; i < 2880; i++) {
+        statement.execute(
+            String.format(
+                "INSERT INTO root.AI(time, s) VALUES(%d, %f)", i, Math.sin(i * Math.PI / 1440)));
+      }
+    }
+  }
+
+  //  @Test
+  public void concurrentGPUCallInferenceTest() throws SQLException, InterruptedException {
+    concurrentGPUCallInferenceTest("timer_xl");
+    concurrentGPUCallInferenceTest("sundial");
+  }
+
+  private void concurrentGPUCallInferenceTest(String modelId)
+      throws SQLException, InterruptedException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      final int threadCnt = 10;
+      final int loop = 100;
+      final int predictLength = 512;
+      final String devices = "0,1";
+      statement.execute(String.format("LOAD MODEL %s TO DEVICES '%s'", modelId, devices));
+      checkModelOnSpecifiedDevice(statement, modelId, devices);
+      concurrentInference(
+          statement,
+          String.format(
+              "CALL INFERENCE(%s, 'SELECT s FROM root.AI', predict_length=%d)",
+              modelId, predictLength),
+          threadCnt,
+          loop,
+          predictLength);
+      statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES '0,1'", modelId));
+    }
+  }
+
+  @Test
+  public void concurrentGPUForecastTest() throws SQLException, InterruptedException {
+    concurrentGPUForecastTest("timer_xl");
+    concurrentGPUForecastTest("sundial");
+  }
+
+  public void concurrentGPUForecastTest(String modelId) throws SQLException, InterruptedException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      final int threadCnt = 10;
+      final int loop = 100;
+      final int predictLength = 512;
+      final String devices = "0,1";
+      statement.execute(String.format("LOAD MODEL %s TO DEVICES '%s'", modelId, devices));
+      checkModelOnSpecifiedDevice(statement, modelId, devices);
+      long startTime = System.currentTimeMillis();
+      concurrentInference(
+          statement,
+          String.format(
+              "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time), predict_length=>%d",
+              modelId, predictLength),
+          threadCnt,
+          loop,
+          predictLength);
+      long endTime = System.currentTimeMillis();
+      LOGGER.info(
+          String.format(
+              "Model %s concurrent inference %d reqs (%d threads, %d loops) in GPU takes time: %dms",
+              modelId, threadCnt * loop, threadCnt, loop, endTime - startTime));
+      statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES '0,1'", modelId));
+    }
+  }
+
+  private void checkModelOnSpecifiedDevice(Statement statement, String modelId, String device)
+      throws SQLException, InterruptedException {
+    Set<String> targetDevices = ImmutableSet.copyOf(device.split(","));
+    LOGGER.info("Checking model: {} on target devices: {}", modelId, targetDevices);
+    for (int retry = 0; retry < 200; retry++) {
+      Set<String> foundDevices = new HashSet<>();
+      try (final ResultSet resultSet =
+          statement.executeQuery(String.format("SHOW LOADED MODELS '%s'", device))) {
+        while (resultSet.next()) {
+          String deviceId = resultSet.getString("DeviceId");
+          String loadedModelId = resultSet.getString("ModelId");
+          int count = resultSet.getInt("Count(instances)");
+          LOGGER.info("Model {} found in device {}, count {}", loadedModelId, deviceId, count);
+          if (loadedModelId.equals(modelId) && targetDevices.contains(deviceId) && count > 0) {
+            foundDevices.add(deviceId);
+            LOGGER.info("Model {} is loaded to device {}", modelId, device);
+          }
+        }
+        if (foundDevices.containsAll(targetDevices)) {
+          LOGGER.info("Model {} is loaded to devices {}, start testing", modelId, targetDevices);
+          return;
+        }
+      }
+      TimeUnit.SECONDS.sleep(3);
+    }
+    Assert.fail("Model " + modelId + " is not loaded on device " + device);
+  }
+}

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

@@ -95,7 +95,7 @@ public static void tearDown() throws Exception {
     EnvFactory.getEnv().cleanClusterEnvironment();
   }
 
-  @Test
+  //  @Test
   public void callInferenceTestInTree() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
         Statement statement = connection.createStatement()) {
@@ -209,7 +209,7 @@ public void callInferenceTest(Statement statement) throws SQLException {
     //    }
   }
 
-  @Test
+  //  @Test
   public void errorCallInferenceTestInTree() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
         Statement statement = connection.createStatement()) {