Skip to content

Commit f5470e8

Browse files
yunbow30944yunbow30944
authored
[AINode] More accurate exceptions (#17061)
1 parent 46a4ba2 commit f5470e8

10 files changed

Lines changed: 113 additions & 28 deletions

File tree

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ public static void forecastTableFunctionErrorTest(
133133
"701: The OUTPUT_START_TIME should be greater than the maximum timestamp of target time series. Expected greater than [5759] but found [5759].");
134134

135135
// OUTPUT_LENGTH error
136-
String invalidOutputLengthSQL =
136+
String invalidOutputLengthSQLWithZero =
137137
String.format(
138138
FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
139139
modelInfo.getModelId(),
@@ -144,7 +144,24 @@ public static void forecastTableFunctionErrorTest(
144144
0,
145145
1,
146146
"time");
147-
errorTest(statement, invalidOutputLengthSQL, "701: OUTPUT_LENGTH should be greater than 0");
147+
errorTest(
148+
statement, invalidOutputLengthSQLWithZero, "701: OUTPUT_LENGTH should be greater than 0");
149+
150+
String invalidOutputLengthSQLWithOutOfRange =
151+
String.format(
152+
FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
153+
modelInfo.getModelId(),
154+
0,
155+
5760,
156+
2880,
157+
5760,
158+
2881,
159+
1,
160+
"time");
161+
errorTest(
162+
statement,
163+
invalidOutputLengthSQLWithOutOfRange,
164+
"1599: Error occurred while executing forecast:[Attribute output_length expect value between 1 and 2880, got 2881 instead.]");
148165

149166
// OUTPUT_INTERVAL error
150167
String invalidOutputIntervalSQL =

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,5 +155,13 @@ private void failTest(Statement statement) {
155155
statement,
156156
"UNLOAD MODEL sundial FROM DEVICES \"unknown\"",
157157
"1507: Device ID [unknown] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.");
158+
errorTest(
159+
statement,
160+
"LOAD MODEL sundial TO DEVICES \"0,0\"",
161+
"1509: Device ID list contains duplicate entries.");
162+
errorTest(
163+
statement,
164+
"UNLOAD MODEL sundial FROM DEVICES \"0,0\"",
165+
"1510: Device ID list contains duplicate entries.");
158166
}
159167
}

‎iotdb-core/ainode/iotdb/ainode/core/constant.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,8 @@ class TSStatusCode(Enum):
8787
DROP_BUILTIN_MODEL_ERROR = 1506
8888
DROP_MODEL_ERROR = 1507
8989
UNAVAILABLE_AI_DEVICE_ERROR = 1508
90+
LOAD_MODEL_ERROR = 1509
91+
UNLOAD_MODEL_ERROR = 1510
9092

9193
INVALID_URI_ERROR = 1511
9294
INVALID_INFERENCE_CONFIG = 1512

‎iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,11 @@ def loadModel(self, req: TLoadModelReq) -> TSStatus:
9292
status = self._ensure_model_is_registered(req.existingModelId)
9393
if status.code != TSStatusCode.SUCCESS_STATUS.value:
9494
return status
95+
if len(set(req.deviceIdList)) != len(req.deviceIdList):
96+
return TSStatus(
97+
code=TSStatusCode.LOAD_MODEL_ERROR.value,
98+
message="Device ID list contains duplicate entries.",
99+
)
95100
status = self._ensure_device_id_is_available(req.deviceIdList)
96101
if status.code != TSStatusCode.SUCCESS_STATUS.value:
97102
return status
@@ -104,6 +109,11 @@ def unloadModel(self, req: TUnloadModelReq) -> TSStatus:
104109
status = self._ensure_model_is_registered(req.modelId)
105110
if status.code != TSStatusCode.SUCCESS_STATUS.value:
106111
return status
112+
if len(set(req.deviceIdList)) != len(req.deviceIdList):
113+
return TSStatus(
114+
code=TSStatusCode.UNLOAD_MODEL_ERROR.value,
115+
message="Device ID list contains duplicate entries.",
116+
)
107117
status = self._ensure_device_id_is_available(req.deviceIdList)
108118
if status.code != TSStatusCode.SUCCESS_STATUS.value:
109119
return status

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/AINodeConnectionException.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iotdb.db.exception.ainode;
21+
22+
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
23+
24+
import static org.apache.iotdb.rpc.TSStatusCode.CAN_NOT_CONNECT_AINODE;
25+
26+
public class AINodeConnectionException extends IoTDBRuntimeException {
27+
28+
private static final String message =
29+
"Failed to connect to AINode because [%s], please check the status of your AINode.";
30+
31+
public AINodeConnectionException(Exception e) {
32+
super(String.format(message, e.toString()), CAN_NOT_CONNECT_AINODE.getStatusCode());
33+
}
34+
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/ai/InferenceOperator.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121

2222
import org.apache.iotdb.ainode.rpc.thrift.TInferenceReq;
2323
import org.apache.iotdb.ainode.rpc.thrift.TInferenceResp;
24+
import org.apache.iotdb.commons.client.exception.ClientManagerException;
25+
import org.apache.iotdb.db.exception.ainode.AINodeConnectionException;
2426
import org.apache.iotdb.db.exception.runtime.ModelInferenceProcessException;
2527
import org.apache.iotdb.db.protocol.client.an.AINodeClient;
2628
import org.apache.iotdb.db.protocol.client.an.AINodeClientManager;
@@ -33,6 +35,7 @@
3335

3436
import com.google.common.util.concurrent.Futures;
3537
import com.google.common.util.concurrent.ListenableFuture;
38+
import org.apache.thrift.TException;
3639
import org.apache.tsfile.block.column.Column;
3740
import org.apache.tsfile.block.column.ColumnBuilder;
3841
import org.apache.tsfile.read.common.block.TsBlock;
@@ -247,8 +250,8 @@ private void submitInferenceTask() {
247250
new TInferenceReq(
248251
modelInferenceDescriptor.getModelId(), serde.serialize(inputTsBlock))
249252
.setInferenceAttributes(modelInferenceDescriptor.getInferenceAttributes()));
250-
} catch (Exception e) {
251-
throw new ModelInferenceProcessException(e.getMessage());
253+
} catch (ClientManagerException | TException e) {
254+
throw new AINodeConnectionException(e);
252255
}
253256
},
254257
modelInferenceExecutor);

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@
175175
import org.apache.iotdb.db.conf.IoTDBDescriptor;
176176
import org.apache.iotdb.db.exception.BatchProcessException;
177177
import org.apache.iotdb.db.exception.StorageEngineException;
178+
import org.apache.iotdb.db.exception.ainode.AINodeConnectionException;
178179
import org.apache.iotdb.db.exception.metadata.PathNotExistException;
179180
import org.apache.iotdb.db.exception.metadata.SchemaQuotaExceededException;
180181
import org.apache.iotdb.db.exception.sql.SemanticException;
@@ -3705,8 +3706,8 @@ public SettableFuture<ConfigTaskResult> createModel(String modelId, String uri)
37053706
} else {
37063707
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
37073708
}
3708-
} catch (final TException | ClientManagerException e) {
3709-
future.setException(e);
3709+
} catch (final ClientManagerException | TException e) {
3710+
future.setException(new AINodeConnectionException(e));
37103711
}
37113712
return future;
37123713
}
@@ -3723,7 +3724,7 @@ public SettableFuture<ConfigTaskResult> dropModel(final String modelId) {
37233724
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
37243725
}
37253726
} catch (final ClientManagerException | TException e) {
3726-
future.setException(e);
3727+
future.setException(new AINodeConnectionException(e));
37273728
}
37283729
return future;
37293730
}
@@ -3743,8 +3744,8 @@ public SettableFuture<ConfigTaskResult> showModels(final String modelId) {
37433744
return future;
37443745
}
37453746
ShowModelsTask.buildTsBlock(resp, future);
3746-
} catch (final Exception e) {
3747-
future.setException(e);
3747+
} catch (final ClientManagerException | TException e) {
3748+
future.setException(new AINodeConnectionException(e));
37483749
}
37493750
return future;
37503751
}
@@ -3762,8 +3763,8 @@ public SettableFuture<ConfigTaskResult> showLoadedModels(List<String> deviceIdLi
37623763
return future;
37633764
}
37643765
ShowLoadedModelsTask.buildTsBlock(resp, future);
3765-
} catch (final Exception e) {
3766-
future.setException(e);
3766+
} catch (final ClientManagerException | TException e) {
3767+
future.setException(new AINodeConnectionException(e));
37673768
}
37683769
return future;
37693770
}
@@ -3779,8 +3780,8 @@ public SettableFuture<ConfigTaskResult> showAIDevices() {
37793780
return future;
37803781
}
37813782
ShowAIDevicesTask.buildTsBlock(resp, future);
3782-
} catch (final Exception e) {
3783-
future.setException(e);
3783+
} catch (final ClientManagerException | TException e) {
3784+
future.setException(new AINodeConnectionException(e));
37843785
}
37853786
return future;
37863787
}
@@ -3798,8 +3799,8 @@ public SettableFuture<ConfigTaskResult> loadModel(
37983799
} else {
37993800
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
38003801
}
3801-
} catch (final Exception e) {
3802-
future.setException(e);
3802+
} catch (final ClientManagerException | TException e) {
3803+
future.setException(new AINodeConnectionException(e));
38033804
}
38043805
return future;
38053806
}
@@ -3817,8 +3818,8 @@ public SettableFuture<ConfigTaskResult> unloadModel(
38173818
} else {
38183819
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
38193820
}
3820-
} catch (final Exception e) {
3821-
future.setException(e);
3821+
} catch (final ClientManagerException | TException e) {
3822+
future.setException(new AINodeConnectionException(e));
38223823
}
38233824
return future;
38243825
}
@@ -3850,8 +3851,8 @@ public SettableFuture<ConfigTaskResult> createTuningTask(
38503851
} else {
38513852
future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS));
38523853
}
3853-
} catch (final Exception e) {
3854-
future.setException(e);
3854+
} catch (final ClientManagerException | TException e) {
3855+
future.setException(new AINodeConnectionException(e));
38553856
}
38563857
return future;
38573858
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/function/tvf/ClassifyTableFunction.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import org.apache.iotdb.ainode.rpc.thrift.TForecastReq;
2323
import org.apache.iotdb.ainode.rpc.thrift.TForecastResp;
2424
import org.apache.iotdb.commons.client.IClientManager;
25+
import org.apache.iotdb.commons.client.exception.ClientManagerException;
2526
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
27+
import org.apache.iotdb.db.exception.ainode.AINodeConnectionException;
2628
import org.apache.iotdb.db.exception.sql.SemanticException;
2729
import org.apache.iotdb.db.protocol.client.an.AINodeClient;
2830
import org.apache.iotdb.db.protocol.client.an.AINodeClientManager;
@@ -44,6 +46,7 @@
4446
import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification;
4547
import org.apache.iotdb.udf.api.type.Type;
4648

49+
import org.apache.thrift.TException;
4750
import org.apache.tsfile.block.column.Column;
4851
import org.apache.tsfile.block.column.ColumnBuilder;
4952
import org.apache.tsfile.enums.TSDataType;
@@ -71,7 +74,6 @@
7174

7275
import static org.apache.iotdb.commons.udf.builtin.relational.tvf.WindowTVFUtils.findColumnIndex;
7376
import static org.apache.iotdb.db.queryengine.plan.relational.utils.ResultColumnAppender.createResultColumnAppender;
74-
import static org.apache.iotdb.rpc.TSStatusCode.CAN_NOT_CONNECT_AINODE;
7577

7678
public class ClassifyTableFunction implements TableFunction {
7779

@@ -367,8 +369,10 @@ private TsBlock classify() {
367369
try (AINodeClient client =
368370
CLIENT_MANAGER.borrowClient(AINodeClientManager.AINODE_ID_PLACEHOLDER)) {
369371
resp = client.forecast(new TForecastReq(modelId, SERDE.serialize(inputData), outputLength));
370-
} catch (Exception e) {
371-
throw new IoTDBRuntimeException(e.getMessage(), CAN_NOT_CONNECT_AINODE.getStatusCode());
372+
} catch (ClientManagerException | TException e) {
373+
throw new AINodeConnectionException(e);
374+
} catch (IOException e) {
375+
throw new RuntimeException(e);
372376
}
373377

374378
if (resp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/function/tvf/ForecastTableFunction.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import org.apache.iotdb.ainode.rpc.thrift.TForecastReq;
2323
import org.apache.iotdb.ainode.rpc.thrift.TForecastResp;
2424
import org.apache.iotdb.commons.client.IClientManager;
25+
import org.apache.iotdb.commons.client.exception.ClientManagerException;
2526
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
27+
import org.apache.iotdb.db.exception.ainode.AINodeConnectionException;
2628
import org.apache.iotdb.db.exception.sql.SemanticException;
2729
import org.apache.iotdb.db.protocol.client.an.AINodeClient;
2830
import org.apache.iotdb.db.protocol.client.an.AINodeClientManager;
@@ -43,6 +45,7 @@
4345
import org.apache.iotdb.udf.api.relational.table.specification.TableParameterSpecification;
4446
import org.apache.iotdb.udf.api.type.Type;
4547

48+
import org.apache.thrift.TException;
4649
import org.apache.tsfile.block.column.Column;
4750
import org.apache.tsfile.block.column.ColumnBuilder;
4851
import org.apache.tsfile.enums.TSDataType;
@@ -72,7 +75,6 @@
7275

7376
import static org.apache.iotdb.commons.udf.builtin.relational.tvf.WindowTVFUtils.findColumnIndex;
7477
import static org.apache.iotdb.db.queryengine.plan.relational.utils.ResultColumnAppender.createResultColumnAppender;
75-
import static org.apache.iotdb.rpc.TSStatusCode.CAN_NOT_CONNECT_AINODE;
7678

7779
public class ForecastTableFunction implements TableFunction {
7880

@@ -563,8 +565,10 @@ protected TsBlock forecast() {
563565
client.forecast(
564566
new TForecastReq(modelId, SERDE.serialize(inputData), outputLength)
565567
.setOptions(options));
566-
} catch (Exception e) {
567-
throw new IoTDBRuntimeException(e.getMessage(), CAN_NOT_CONNECT_AINODE.getStatusCode());
568+
} catch (ClientManagerException | TException e) {
569+
throw new AINodeConnectionException(e);
570+
} catch (IOException e) {
571+
throw new RuntimeException(e);
568572
}
569573

570574
if (resp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/udf/UDTFForecast.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import org.apache.iotdb.ainode.rpc.thrift.TForecastReq;
2323
import org.apache.iotdb.ainode.rpc.thrift.TForecastResp;
2424
import org.apache.iotdb.commons.client.IClientManager;
25+
import org.apache.iotdb.commons.client.exception.ClientManagerException;
2526
import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
27+
import org.apache.iotdb.db.exception.ainode.AINodeConnectionException;
2628
import org.apache.iotdb.db.protocol.client.an.AINodeClient;
2729
import org.apache.iotdb.db.protocol.client.an.AINodeClientManager;
2830
import org.apache.iotdb.rpc.TSStatusCode;
@@ -34,6 +36,7 @@
3436
import org.apache.iotdb.udf.api.customizer.strategy.RowByRowAccessStrategy;
3537
import org.apache.iotdb.udf.api.type.Type;
3638

39+
import org.apache.thrift.TException;
3740
import org.apache.tsfile.enums.TSDataType;
3841
import org.apache.tsfile.read.common.block.TsBlock;
3942
import org.apache.tsfile.read.common.block.TsBlockBuilder;
@@ -213,9 +216,8 @@ private TsBlock forecast() throws Exception {
213216
client.forecast(
214217
new TForecastReq(model_id, serde.serialize(inputData), outputLength)
215218
.setOptions(options));
216-
} catch (Exception e) {
217-
throw new IoTDBRuntimeException(
218-
e.getMessage(), TSStatusCode.CAN_NOT_CONNECT_AINODE.getStatusCode());
219+
} catch (ClientManagerException | TException e) {
220+
throw new AINodeConnectionException(e);
219221
}
220222

221223
if (resp.getStatus().getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {

0 commit comments

Comments
 (0)