bug fix & append it

Author: CRZbulabula

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

@@ -0,0 +1,94 @@
+/*
+ * 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 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 java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkHeader;
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.prepareDataInTable;
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.prepareDataInTree;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({AIClusterIT.class})
+public class AINodeDeviceManageIT {
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    // Init 1C1D1A cluster environment
+    EnvFactory.getEnv().initClusterEnvironment(1, 1);
+    prepareDataInTree();
+    prepareDataInTable();
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  @Test
+  public void showAIDeviceTestInTree() throws SQLException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      showAIDevicesTest(statement);
+    }
+  }
+
+  @Test
+  public void showAIDeviceTestInTable() throws SQLException {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      showAIDevicesTest(statement);
+    }
+  }
+
+  private void showAIDevicesTest(Statement statement) throws SQLException {
+    final String showSql = "SHOW AI_DEVICES";
+    final List<String> expectedDeviceIdList = Arrays.asList("0", "1", "cpu");
+    final List<String> expectedDeviceTypeList = Arrays.asList("cuda", "cuda", "cpu");
+    try (ResultSet resultSet = statement.executeQuery(showSql)) {
+      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+      checkHeader(resultSetMetaData, "DeviceId,DeviceType");
+      while (resultSet.next()) {
+        String deviceId = resultSet.getString(1);
+        String deviceType = resultSet.getString(2);
+        Assert.assertEquals(expectedDeviceIdList.remove(0), deviceId);
+        Assert.assertEquals(expectedDeviceTypeList.remove(0), deviceType);
+      }
+    }
+  }
+}

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

@@ -77,8 +77,6 @@ def __init__(
         self.ready_event = ready_event
         self.device = device
 
-        self._backend = DeviceManager()
-
         self._threads = []
         self._waiting_queue = request_queue  # Requests that are waiting to be processed
         self._running_queue = mp.Queue()  # Requests that are currently being processed
@@ -89,8 +87,8 @@ def __init__(
         self._batcher = BasicBatcher()
         self._stop_event = mp.Event()
 
+        self._backend = None
         self._inference_pipeline = None
-
         self._logger = None
 
         # Fix inference seed
@@ -186,6 +184,7 @@ def run(self):
         self._logger = Logger(
             INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE.format(self.device)
         )
+        self._backend = DeviceManager()
         self._request_scheduler.device = self.device
         self._inference_pipeline = load_pipeline(self.model_info, self.device)
         self.ready_event.set()

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/ShowAIDevicesTask.java

@@ -54,12 +54,15 @@ public static void buildTsBlock(
             .map(ColumnHeader::getColumnType)
             .collect(Collectors.toList());
     TsBlockBuilder builder = new TsBlockBuilder(outputDataTypes);
-    for (Map.Entry<String, String> deviceEntry : resp.getDeviceIdMap().entrySet()) {
-      builder.getTimeColumnBuilder().writeLong(0L);
-      builder.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(deviceEntry.getKey()));
-      builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(deviceEntry.getValue()));
-      builder.declarePosition();
-    }
+    resp.getDeviceIdMap().entrySet().stream()
+        .sorted(Map.Entry.comparingByKey())
+        .forEach(
+            deviceEntry -> {
+              builder.getTimeColumnBuilder().writeLong(0L);
+              builder.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(deviceEntry.getKey()));
+              builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(deviceEntry.getValue()));
+              builder.declarePosition();
+            });
     DatasetHeader datasetHeader = DatasetHeaderFactory.getShowAIDevicesHeader();
     future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS, builder.build(), datasetHeader));
   }