[AINode] Integrate device manager framework (#16998)

Author: CRZbulabula

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

@@ -83,13 +83,15 @@ private static void prepareDataForTableModel() throws SQLException {
   }
 
   @Test
-  public void concurrentGPUForecastTest() throws SQLException, InterruptedException {
+  public void concurrentForecastTest() throws SQLException, InterruptedException {
     for (AINodeTestUtils.FakeModelInfo modelInfo : MODEL_LIST) {
-      concurrentGPUForecastTest(modelInfo);
+      concurrentGPUForecastTest(modelInfo, "0,1");
+      // TODO: Enable cpu test after optimize memory consumption
+      //      concurrentGPUForecastTest(modelInfo, "cpu");
     }
   }
 
-  public void concurrentGPUForecastTest(AINodeTestUtils.FakeModelInfo modelInfo)
+  public void concurrentGPUForecastTest(AINodeTestUtils.FakeModelInfo modelInfo, String devices)
       throws SQLException, InterruptedException {
     final int forecastLength = 512;
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
@@ -100,7 +102,6 @@ public void concurrentGPUForecastTest(AINodeTestUtils.FakeModelInfo modelInfo)
               FORECAST_TABLE_FUNCTION_SQL_TEMPLATE, modelInfo.getModelId(), forecastLength);
       final int threadCnt = 10;
       final int loop = 100;
-      final String devices = "0,1";
       statement.execute(
           String.format("LOAD MODEL %s TO DEVICES '%s'", modelInfo.getModelId(), devices));
       checkModelOnSpecifiedDevice(statement, modelInfo.getModelId(), devices);

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

@@ -0,0 +1,96 @@
+/*
+ * 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.LinkedList;
+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 = new LinkedList<>(Arrays.asList("0", "1", "cpu"));
+    final List<String> expectedDeviceTypeList =
+        new LinkedList<>(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/constant.py

@@ -56,7 +56,7 @@
     "timer": 856 * 1024**2,  # 856 MiB
 }  # the memory usage of each model in bytes
 
-AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.4  # the device space allocated for inference
+AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.2  # the device space allocated for inference
 AINODE_INFERENCE_EXTRA_MEMORY_RATIO = (
     1.2  # the overhead ratio for inference, used to estimate the pool size
 )

iotdb-core/ainode/iotdb/ainode/core/device/__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/iotdb/ainode/core/device/backend/__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/iotdb/ainode/core/device/backend/base.py

@@ -0,0 +1,42 @@
+# 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.
+#
+
+from enum import Enum
+from typing import ContextManager, Optional, Protocol
+
+import torch
+
+
+class BackendType(Enum):
+    """
+    Different types of supported computation backends.
+    AINode will automatically select the available backend according to the order defined here.
+    """
+
+    CUDA = "cuda"
+    CPU = "cpu"
+
+
+class BackendAdapter(Protocol):
+    type: BackendType
+
+    # device basics
+    def is_available(self) -> bool: ...
+    def device_count(self) -> int: ...
+    def make_device(self, index: Optional[int]) -> torch.device: ...
+    def set_device(self, index: int) -> None: ...

iotdb-core/ainode/iotdb/ainode/core/device/backend/cpu_backend.py

@@ -0,0 +1,37 @@
+# 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 torch
+
+from iotdb.ainode.core.device.backend.base import BackendAdapter, BackendType
+
+
+class CPUBackend(BackendAdapter):
+    type = BackendType.CPU
+
+    def is_available(self) -> bool:
+        return True
+
+    def device_count(self) -> int:
+        return 1
+
+    def make_device(self, index: int | None) -> torch.device:
+        return torch.device("cpu")
+
+    def set_device(self, index: int) -> None:
+        return None

iotdb-core/ainode/iotdb/ainode/core/device/backend/cuda_backend.py

@@ -0,0 +1,39 @@
+# 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 torch
+
+from iotdb.ainode.core.device.backend.base import BackendAdapter, BackendType
+
+
+class CUDABackend(BackendAdapter):
+    type = BackendType.CUDA
+
+    def is_available(self) -> bool:
+        return torch.cuda.is_available()
+
+    def device_count(self) -> int:
+        return torch.cuda.device_count()
+
+    def make_device(self, index: int | None) -> torch.device:
+        if index is None:
+            raise ValueError("CUDA backend requires a valid device index")
+        return torch.device(f"cuda:{index}")
+
+    def set_device(self, index: int) -> None:
+        torch.cuda.set_device(index)

iotdb-core/ainode/iotdb/ainode/core/device/device_utils.py

@@ -0,0 +1,49 @@
+# 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.
+#
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+DeviceLike = Union[torch.device, str, int]
+
+
+@dataclass(frozen=True)
+class DeviceSpec:
+    type: str
+    index: Optional[int]
+
+
+def parse_device_like(x: DeviceLike) -> DeviceSpec:
+    if isinstance(x, int):
+        return DeviceSpec("index", x)
+
+    if isinstance(x, str):
+        try:
+            return DeviceSpec("index", int(x))
+        except ValueError:
+            s = x.strip().lower()
+            if ":" in s:
+                t, idx = s.split(":", 1)
+                return DeviceSpec(t, int(idx))
+            return DeviceSpec(s, None)
+
+    if isinstance(x, torch.device):
+        return DeviceSpec(x.type, x.index)
+
+    raise TypeError(f"Unsupported device: {x!r}")

iotdb-core/ainode/iotdb/ainode/core/device/env.py

@@ -0,0 +1,39 @@
+# 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 os
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DistEnv:
+    rank: int
+    local_rank: int
+    world_size: int
+
+
+def read_dist_env() -> DistEnv:
+    # torchrun:
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+
+    # torchrun provides LOCAL_RANK; slurm often provides SLURM_LOCALID
+    local_rank = os.environ.get("LOCAL_RANK", os.environ.get("SLURM_LOCALID", "0"))
+    local_rank = int(local_rank)
+
+    return DistEnv(rank=rank, local_rank=local_rank, world_size=world_size)