Skip to content

Commit fd34d93

Browse files
CRZbulabulaCRZbulabula
committed
[AINode] Concurrent inference bug fix (#16595)
(cherry picked from commit 46a0c6a)
1 parent cfaa0b6 commit fd34d93

File tree

14 files changed

+90
-126
lines changed

14 files changed

+90
-126
lines changed

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

Lines changed: 36 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,17 @@
2020
package org.apache.iotdb.ainode.it;
2121

2222
import org.apache.iotdb.it.env.EnvFactory;
23+
import org.apache.iotdb.it.framework.IoTDBTestRunner;
24+
import org.apache.iotdb.itbase.category.AIClusterIT;
2325
import org.apache.iotdb.itbase.env.BaseEnv;
2426

25-
import com.google.common.collect.ImmutableMap;
2627
import com.google.common.collect.ImmutableSet;
2728
import org.junit.AfterClass;
2829
import org.junit.Assert;
2930
import org.junit.BeforeClass;
3031
import org.junit.Test;
32+
import org.junit.experimental.categories.Category;
33+
import org.junit.runner.RunWith;
3134
import org.slf4j.Logger;
3235
import org.slf4j.LoggerFactory;
3336

@@ -36,21 +39,17 @@
3639
import java.sql.SQLException;
3740
import java.sql.Statement;
3841
import java.util.HashSet;
39-
import java.util.Map;
4042
import java.util.Set;
4143
import java.util.concurrent.TimeUnit;
4244

4345
import static org.apache.iotdb.ainode.utils.AINodeTestUtils.concurrentInference;
4446

47+
@RunWith(IoTDBTestRunner.class)
48+
@Category({AIClusterIT.class})
4549
public class AINodeConcurrentInferenceIT {
4650

4751
private static final Logger LOGGER = LoggerFactory.getLogger(AINodeConcurrentInferenceIT.class);
4852

49-
private static final Map<String, String> MODEL_ID_TO_TYPE_MAP =
50-
ImmutableMap.of(
51-
"timer_xl", "Timer-XL",
52-
"sundial", "Timer-Sundial");
53-
5453
@BeforeClass
5554
public static void setUp() throws Exception {
5655
// Init 1C1D1A cluster environment
@@ -86,13 +85,12 @@ private static void prepareDataForTableModel() throws SQLException {
8685
for (int i = 0; i < 2880; i++) {
8786
statement.execute(
8887
String.format(
89-
"INSERT INTO root.AI(timestamp, s) VALUES(%d, %f)",
90-
i, Math.sin(i * Math.PI / 1440)));
88+
"INSERT INTO root.AI(time, s) VALUES(%d, %f)", i, Math.sin(i * Math.PI / 1440)));
9189
}
9290
}
9391
}
9492

95-
@Test
93+
// @Test
9694
public void concurrentCPUCallInferenceTest() throws SQLException, InterruptedException {
9795
concurrentCPUCallInferenceTest("timer_xl");
9896
concurrentCPUCallInferenceTest("sundial");
@@ -105,21 +103,21 @@ private void concurrentCPUCallInferenceTest(String modelId)
105103
final int threadCnt = 4;
106104
final int loop = 10;
107105
final int predictLength = 96;
108-
statement.execute(String.format("LOAD MODEL %s TO DEVICES \"cpu\"", modelId));
109-
checkModelOnSpecifiedDevice(statement, MODEL_ID_TO_TYPE_MAP.get(modelId), "cpu");
106+
statement.execute(String.format("LOAD MODEL %s TO DEVICES 'cpu'", modelId));
107+
checkModelOnSpecifiedDevice(statement, modelId, "cpu");
110108
concurrentInference(
111109
statement,
112110
String.format(
113-
"CALL INFERENCE(%s, \"SELECT s FROM root.AI\", predict_length=%d)",
111+
"CALL INFERENCE(%s, 'SELECT s FROM root.AI', predict_length=%d)",
114112
modelId, predictLength),
115113
threadCnt,
116114
loop,
117115
predictLength);
118-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES \"cpu\"", modelId));
116+
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES 'cpu'", modelId));
119117
}
120118
}
121119

122-
@Test
120+
// @Test
123121
public void concurrentGPUCallInferenceTest() throws SQLException, InterruptedException {
124122
concurrentGPUCallInferenceTest("timer_xl");
125123
concurrentGPUCallInferenceTest("sundial");
@@ -133,17 +131,17 @@ private void concurrentGPUCallInferenceTest(String modelId)
133131
final int loop = 100;
134132
final int predictLength = 512;
135133
final String devices = "0,1";
136-
statement.execute(String.format("LOAD MODEL %s TO DEVICES \"%s\"", modelId, devices));
137-
checkModelOnSpecifiedDevice(statement, MODEL_ID_TO_TYPE_MAP.get(modelId), devices);
134+
statement.execute(String.format("LOAD MODEL %s TO DEVICES '%s'", modelId, devices));
135+
checkModelOnSpecifiedDevice(statement, modelId, devices);
138136
concurrentInference(
139137
statement,
140138
String.format(
141-
"CALL INFERENCE(%s, \"SELECT s FROM root.AI\", predict_length=%d)",
139+
"CALL INFERENCE(%s, 'SELECT s FROM root.AI', predict_length=%d)",
142140
modelId, predictLength),
143141
threadCnt,
144142
loop,
145143
predictLength);
146-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES \"0,1\"", modelId));
144+
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES '0,1'", modelId));
147145
}
148146
}
149147

@@ -159,8 +157,8 @@ private void concurrentCPUForecastTest(String modelId) throws SQLException, Inte
159157
final int threadCnt = 4;
160158
final int loop = 10;
161159
final int predictLength = 96;
162-
statement.execute(String.format("LOAD MODEL %s TO DEVICES \"cpu\"", modelId));
163-
checkModelOnSpecifiedDevice(statement, MODEL_ID_TO_TYPE_MAP.get(modelId), "cpu");
160+
statement.execute(String.format("LOAD MODEL %s TO DEVICES 'cpu'", modelId));
161+
checkModelOnSpecifiedDevice(statement, modelId, "cpu");
164162
long startTime = System.currentTimeMillis();
165163
concurrentInference(
166164
statement,
@@ -175,7 +173,7 @@ private void concurrentCPUForecastTest(String modelId) throws SQLException, Inte
175173
String.format(
176174
"Model %s concurrent inference %d reqs (%d threads, %d loops) in CPU takes time: %dms",
177175
modelId, threadCnt * loop, threadCnt, loop, endTime - startTime));
178-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES \"cpu\"", modelId));
176+
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES 'cpu'", modelId));
179177
}
180178
}
181179

@@ -192,8 +190,8 @@ public void concurrentGPUForecastTest(String modelId) throws SQLException, Inter
192190
final int loop = 100;
193191
final int predictLength = 512;
194192
final String devices = "0,1";
195-
statement.execute(String.format("LOAD MODEL %s TO DEVICES \"%s\"", modelId, devices));
196-
checkModelOnSpecifiedDevice(statement, MODEL_ID_TO_TYPE_MAP.get(modelId), devices);
193+
statement.execute(String.format("LOAD MODEL %s TO DEVICES '%s'", modelId, devices));
194+
checkModelOnSpecifiedDevice(statement, modelId, devices);
197195
long startTime = System.currentTimeMillis();
198196
concurrentInference(
199197
statement,
@@ -208,32 +206,35 @@ public void concurrentGPUForecastTest(String modelId) throws SQLException, Inter
208206
String.format(
209207
"Model %s concurrent inference %d reqs (%d threads, %d loops) in GPU takes time: %dms",
210208
modelId, threadCnt * loop, threadCnt, loop, endTime - startTime));
211-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES \"0,1\"", modelId));
209+
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES '0,1'", modelId));
212210
}
213211
}
214212

215-
private void checkModelOnSpecifiedDevice(Statement statement, String modelType, String device)
213+
private void checkModelOnSpecifiedDevice(Statement statement, String modelId, String device)
216214
throws SQLException, InterruptedException {
217-
for (int retry = 0; retry < 10; retry++) {
218-
Set<String> targetDevices = ImmutableSet.copyOf(device.split(","));
215+
Set<String> targetDevices = ImmutableSet.copyOf(device.split(","));
216+
LOGGER.info("Checking model: {} on target devices: {}", modelId, targetDevices);
217+
for (int retry = 0; retry < 20; retry++) {
219218
Set<String> foundDevices = new HashSet<>();
220219
try (final ResultSet resultSet =
221-
statement.executeQuery(String.format("SHOW LOADED MODELS %s", device))) {
220+
statement.executeQuery(String.format("SHOW LOADED MODELS '%s'", device))) {
222221
while (resultSet.next()) {
223-
String deviceId = resultSet.getString(1);
224-
String loadedModelType = resultSet.getString(2);
225-
int count = resultSet.getInt(3);
226-
if (loadedModelType.equals(modelType) && targetDevices.contains(deviceId)) {
227-
Assert.assertTrue(count > 1);
222+
String deviceId = resultSet.getString("DeviceId");
223+
String loadedModelId = resultSet.getString("ModelId");
224+
int count = resultSet.getInt("Count(instances)");
225+
LOGGER.info("Model {} found in device {}, count {}", loadedModelId, deviceId, count);
226+
if (loadedModelId.equals(modelId) && targetDevices.contains(deviceId) && count > 0) {
228227
foundDevices.add(deviceId);
228+
LOGGER.info("Model {} is loaded to device {}", modelId, device);
229229
}
230230
}
231231
if (foundDevices.containsAll(targetDevices)) {
232+
LOGGER.info("Model {} is loaded to devices {}, start testing", modelId, targetDevices);
232233
return;
233234
}
234235
}
235236
TimeUnit.SECONDS.sleep(3);
236237
}
237-
Assert.fail("Model " + modelType + " is not loaded on device " + device);
238+
Assert.fail("Model " + modelId + " is not loaded on device " + device);
238239
}
239240
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ public static void tearDown() throws Exception {
9595
EnvFactory.getEnv().cleanClusterEnvironment();
9696
}
9797

98-
@Test
98+
// @Test
9999
public void callInferenceTestInTree() throws SQLException {
100100
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
101101
Statement statement = connection.createStatement()) {
@@ -209,7 +209,7 @@ public void callInferenceTest(Statement statement) throws SQLException {
209209
// }
210210
}
211211

212-
@Test
212+
// @Test
213213
public void errorCallInferenceTestInTree() throws SQLException {
214214
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
215215
Statement statement = connection.createStatement()) {

‎iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
ScaleActionType,
4141
)
4242
from iotdb.ainode.core.log import Logger
43+
from iotdb.ainode.core.manager.model_manager import ModelManager
44+
from iotdb.ainode.core.model.model_enums import BuiltInModelType
4345
from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
4446
from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
4547
from iotdb.ainode.core.util.atmoic_int import AtomicInt
@@ -48,6 +50,7 @@
4850
from iotdb.ainode.core.util.thread_name import ThreadName
4951

5052
logger = Logger()
53+
MODEL_MANAGER = ModelManager()
5154

5255

5356
class PoolController:
@@ -169,7 +172,7 @@ def show_loaded_models(
169172
for model_id, device_map in self._request_pool_map.items():
170173
if device_id in device_map:
171174
pool_group = device_map[device_id]
172-
device_models[model_id] = pool_group.get_pool_count()
175+
device_models[model_id] = pool_group.get_running_pool_count()
173176
result[device_id] = device_models
174177
return result
175178

@@ -191,7 +194,7 @@ def _load_model_task(self, model_id: str, device_id_list: list[str]):
191194
def _load_model_on_device_task(device_id: str):
192195
if not self.has_request_pools(model_id, device_id):
193196
actions = self._pool_scheduler.schedule_load_model_to_device(
194-
model_id, device_id
197+
MODEL_MANAGER.get_model_info(model_id), device_id
195198
)
196199
for action in actions:
197200
if action.action == ScaleActionType.SCALE_UP:
@@ -218,7 +221,7 @@ def _unload_model_task(self, model_id: str, device_id_list: list[str]):
218221
def _unload_model_on_device_task(device_id: str):
219222
if self.has_request_pools(model_id, device_id):
220223
actions = self._pool_scheduler.schedule_unload_model_from_device(
221-
model_id, device_id
224+
MODEL_MANAGER.get_model_info(model_id), device_id
222225
)
223226
for action in actions:
224227
if action.action == ScaleActionType.SCALE_DOWN:
@@ -253,13 +256,19 @@ def _expand_pools_on_device(self, model_id: str, device_id: str, count: int):
253256
def _expand_pool_on_device(*_):
254257
result_queue = mp.Queue()
255258
pool_id = self._new_pool_id.get_and_increment()
256-
if model_id == "sundial":
259+
model_info = MODEL_MANAGER.get_model_info(model_id)
260+
model_type = model_info.model_type
261+
if model_type == BuiltInModelType.SUNDIAL.value:
257262
config = SundialConfig()
258-
elif model_id == "timer_xl":
263+
elif model_type == BuiltInModelType.TIMER_XL.value:
259264
config = TimerConfig()
265+
else:
266+
raise InferenceModelInternalError(
267+
f"Unsupported model type {model_type} for loading model {model_id}"
268+
)
260269
pool = InferenceRequestPool(
261270
pool_id=pool_id,
262-
model_id=model_id,
271+
model_info=model_info,
263272
device=device_id,
264273
config=config,
265274
request_queue=result_queue,

‎iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,12 @@ def get_pool_ids(self) -> list[int]:
7070
def get_pool_count(self) -> int:
7171
return len(self.pool_group)
7272

73+
def get_running_pool_count(self) -> int:
74+
count = 0
75+
for _, state in self.pool_states.items():
76+
count += 1 if state == PoolState.RUNNING else 0
77+
return count
78+
7379
def dispatch_request(
7480
self, req: InferenceRequest, infer_proxy: InferenceRequestProxy
7581
):

‎iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,9 @@ def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
147147
def register_built_in_model(self, model_info: ModelInfo):
148148
self.model_storage.register_built_in_model(model_info)
149149

150+
def get_model_info(self, model_id: str) -> ModelInfo:
151+
return self.model_storage.get_model_info(model_id)
152+
150153
def update_model_state(self, model_id: str, state: ModelStates):
151154
self.model_storage.update_model_state(model_id, state)
152155

‎iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,13 @@ def register_built_in_model(self, model_info: ModelInfo):
423423
with self._lock_pool.get_lock(model_info.model_id).write_lock():
424424
self._model_info_map[model_info.model_id] = model_info
425425

426+
def get_model_info(self, model_id: str) -> ModelInfo:
427+
with self._lock_pool.get_lock(model_id).read_lock():
428+
if model_id in self._model_info_map:
429+
return self._model_info_map[model_id]
430+
else:
431+
raise ValueError(f"Model {model_id} does not exist.")
432+
426433
def update_model_state(self, model_id: str, state: ModelStates):
427434
with self._lock_pool.get_lock(model_id).write_lock():
428435
if model_id in self._model_info_map:

‎iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/GetModelInfoResp.java‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,9 @@
2525
import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
2626
import org.apache.iotdb.consensus.common.DataSet;
2727

28-
import java.nio.ByteBuffer;
29-
3028
public class GetModelInfoResp implements DataSet {
3129

3230
private final TSStatus status;
33-
private ByteBuffer serializedModelInformation;
3431

3532
private int targetAINodeId;
3633
private TEndPoint targetAINodeAddress;
@@ -43,10 +40,6 @@ public GetModelInfoResp(TSStatus status) {
4340
this.status = status;
4441
}
4542

46-
public void setModelInfo(ByteBuffer serializedModelInformation) {
47-
this.serializedModelInformation = serializedModelInformation;
48-
}
49-
5043
public int getTargetAINodeId() {
5144
return targetAINodeId;
5245
}
@@ -64,7 +57,6 @@ public void setTargetAINodeAddress(TAINodeConfiguration aiNodeConfiguration) {
6457

6558
public TGetModelInfoResp convertToThriftResponse() {
6659
TGetModelInfoResp resp = new TGetModelInfoResp(status);
67-
resp.setModelInfo(serializedModelInformation);
6860
resp.setAiNodeAddress(targetAINodeAddress);
6961
return resp;
7062
}

‎iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java‎

Lines changed: 9 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,8 @@
3131
import org.apache.iotdb.commons.model.ModelInformation;
3232
import org.apache.iotdb.commons.model.ModelStatus;
3333
import org.apache.iotdb.commons.model.ModelType;
34-
import org.apache.iotdb.confignode.consensus.request.read.model.GetModelInfoPlan;
3534
import org.apache.iotdb.confignode.consensus.request.write.model.CreateModelPlan;
3635
import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelInfoPlan;
37-
import org.apache.iotdb.confignode.consensus.response.model.GetModelInfoResp;
3836
import org.apache.iotdb.confignode.exception.NoAvailableAINodeException;
3937
import org.apache.iotdb.confignode.persistence.ModelInfo;
4038
import org.apache.iotdb.confignode.rpc.thrift.TAINodeInfo;
@@ -186,33 +184,15 @@ public TShowAIDevicesResp showAIDevices() {
186184
}
187185

188186
public TGetModelInfoResp getModelInfo(TGetModelInfoReq req) {
189-
try {
190-
GetModelInfoResp response =
191-
(GetModelInfoResp) configManager.getConsensusManager().read(new GetModelInfoPlan(req));
192-
if (response.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
193-
return new TGetModelInfoResp(response.getStatus());
194-
}
195-
int aiNodeId = response.getTargetAINodeId();
196-
if (aiNodeId != 0) {
197-
response.setTargetAINodeAddress(
198-
configManager.getNodeManager().getRegisteredAINode(aiNodeId));
199-
} else {
200-
if (configManager.getNodeManager().getRegisteredAINodes().isEmpty()) {
201-
return new TGetModelInfoResp(
202-
new TSStatus(TSStatusCode.GET_MODEL_INFO_ERROR.getStatusCode())
203-
.setMessage("There is no AINode available"));
204-
}
205-
response.setTargetAINodeAddress(
206-
configManager.getNodeManager().getRegisteredAINodes().get(0));
207-
}
208-
return response.convertToThriftResponse();
209-
} catch (ConsensusException e) {
210-
LOGGER.warn("Unexpected error happened while getting model: ", e);
211-
// consensus layer related errors
212-
TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode());
213-
res.setMessage(e.getMessage());
214-
return new TGetModelInfoResp(res);
215-
}
187+
return new TGetModelInfoResp()
188+
.setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()))
189+
.setAiNodeAddress(
190+
configManager
191+
.getNodeManager()
192+
.getRegisteredAINodes()
193+
.get(0)
194+
.getLocation()
195+
.getInternalEndPoint());
216196
}
217197

218198
// Currently this method is only used by built-in timer_xl

0 commit comments

Comments
 (0)