You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ja...@apache.org on 2025/12/16 08:28:22 UTC
(iotdb) branch force_ci/object_type updated (ef5cef5c9d5 -> 3bc6101a93d)
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a change to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
from ef5cef5c9d5 Two Implementation of Object File Path (#16858)
new 8f1235835ae Throw exception while readObject column is not OBJECT type
new bc9418a8af3 [AINode] More accurate exception for model management (#16895)
new 2d3c19a5656 Fix the config 'query_cost_stat_window' cannot be hot loaded #16896
new 75e8579e600 Hotfix concurrent connect error (#16900)
new 191edf13094 [AINode] More strict concurrent inference IT (#16898)
new a5c3b4387bb Delete audit implementation (#16899)
new 900ff029b1f Avoid CN OOM by Pulling User/Roles to DN When Cache Misses. (#16888)
new 36cb21f4fd6 fix showQueries (#16855)
new bd677fb9954 Add state predicate push down for system table current_queries (#16904)
new 3bc6101a93d Optimized the error message & prevented stack printing of information_schema selection when the cluster is closing (#16905)
The 10 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../iotdb/ainode/it/AINodeClusterConfigIT.java | 3 -
.../ainode/it/AINodeConcurrentForecastIT.java | 2 +-
.../iotdb/ainode/it/AINodeModelManageIT.java | 4 +-
.../apache/iotdb/ainode/utils/AINodeTestUtils.java | 17 +-
.../iotdb/db/it/audit/IoTDBAuditLogBasicIT.java | 2 +
.../informationschema/IoTDBCurrentQueriesIT.java | 11 +-
.../iotdb/relational/it/schema/IoTDBTableIT.java | 6 +
.../java/org/apache/iotdb/rpc/TSStatusCode.java | 21 +-
iotdb-core/ainode/iotdb/ainode/core/config.py | 8 +-
iotdb-core/ainode/iotdb/ainode/core/constant.py | 33 +-
iotdb-core/ainode/iotdb/ainode/core/exception.py | 108 +++---
.../core/inference/dispatcher/basic_dispatcher.py | 6 +-
.../core/inference/inference_request_pool.py | 4 +-
.../iotdb/ainode/core/inference/pool_controller.py | 4 +-
.../iotdb/ainode/core/inference/pool_group.py | 6 +-
.../pool_scheduler/basic_pool_scheduler.py | 4 +-
.../iotdb/ainode/core/manager/inference_manager.py | 4 +-
.../iotdb/ainode/core/manager/model_manager.py | 43 +--
.../ainode/iotdb/ainode/core/manager/utils.py | 4 +-
.../iotdb/ainode/core/model/model_constants.py | 7 -
.../ainode/iotdb/ainode/core/model/model_loader.py | 4 +-
.../iotdb/ainode/core/model/model_storage.py | 158 +++-----
.../core/model/sktime/configuration_sktime.py | 36 +-
.../ainode/core/model/sktime/modeling_sktime.py | 12 +-
.../ainode/core/model/sundial/pipeline_sundial.py | 4 +-
.../ainode/core/model/timer_xl/pipeline_timer.py | 4 +-
iotdb-core/ainode/iotdb/ainode/core/model/utils.py | 53 ++-
iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py | 48 +--
iotdb-core/ainode/iotdb/ainode/core/util/serde.py | 6 +-
.../iotdb/confignode/audit/CNAuditLogger.java | 50 +--
.../org/apache/iotdb/db/audit/DNAuditLogger.java | 398 +--------------------
.../iotdb/db/auth/ClusterAuthorityFetcher.java | 59 ++-
.../org/apache/iotdb/db/conf/IoTDBDescriptor.java | 6 +
.../db/exception/ainode/GetModelInfoException.java | 2 +-
.../exception/ainode/ModelNotFoundException.java | 28 --
.../iotdb/db/protocol/client/an/AINodeClient.java | 6 +-
.../iotdb/db/protocol/session/IClientSession.java | 3 +-
.../iotdb/db/protocol/session/SessionManager.java | 4 +-
.../operator/process/function/partition/Slice.java | 3 +
.../InformationSchemaContentSupplierFactory.java | 30 +-
.../relational/aggregation/RecordIterator.java | 3 +
.../iotdb/db/queryengine/plan/Coordinator.java | 33 ++
.../plan/execution/config/ConfigExecution.java | 1 +
.../execution/config/TableConfigTaskVisitor.java | 4 +-
.../execution/config/TreeConfigTaskVisitor.java | 4 +-
.../config/executor/ClusterConfigTaskExecutor.java | 10 +-
.../config/executor/IConfigTaskExecutor.java | 2 +-
...eateTrainingTask.java => CreateTuningTask.java} | 8 +-
.../db/queryengine/plan/parser/ASTVisitor.java | 22 +-
.../plan/planner/TableOperatorGenerator.java | 1 +
.../planner/distribution/ExchangeNodeAdder.java | 6 +-
.../planner/plan/node/source/ShowQueriesNode.java | 2 +-
.../distribute/TableDistributedPlanGenerator.java | 8 +-
.../optimizations/PushPredicateIntoTableScan.java | 100 ++++++
.../plan/statement/component/OrderByComponent.java | 6 +
.../FragmentInstanceStatisticsDrawer.java | 2 +-
.../plan/parser/StatementGeneratorTest.java | 33 ++
...NodeSerdeTest.java => SourceNodeSerdeTest.java} | 20 +-
.../planner/assertions/ColumnReference.java | 10 +-
.../planner/assertions/PlanMatchPattern.java | 20 +-
.../informationschema/CurrentQueriesTest.java | 77 ++++
.../iotdb/commons/audit/AbstractAuditLogger.java | 77 +---
.../thrift-ainode/src/main/thrift/ainode.thrift | 17 +-
63 files changed, 708 insertions(+), 969 deletions(-)
delete mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java
rename iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/{CreateTrainingTask.java => CreateTuningTask.java} (93%)
rename iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/{LastQueryScanNodeSerdeTest.java => SourceNodeSerdeTest.java} (76%)
(iotdb) 08/10: fix showQueries (#16855)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 36cb21f4fd60f61b9c41ea1161b02d5930bb7a86
Author: shuwenwei <55...@users.noreply.github.com>
AuthorDate: Mon Dec 15 17:56:47 2025 +0800
fix showQueries (#16855)
(cherry picked from commit b71401a42271cc45162ac0ffb244fa86dd1a9d8e)
---
.../db/queryengine/plan/parser/ASTVisitor.java | 22 +++++++++++----
.../planner/distribution/ExchangeNodeAdder.java | 6 ++--
.../planner/plan/node/source/ShowQueriesNode.java | 2 +-
.../plan/statement/component/OrderByComponent.java | 6 ++++
.../FragmentInstanceStatisticsDrawer.java | 2 +-
.../plan/parser/StatementGeneratorTest.java | 33 ++++++++++++++++++++++
...NodeSerdeTest.java => SourceNodeSerdeTest.java} | 20 +++++++++++--
7 files changed, 80 insertions(+), 11 deletions(-)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
index 7f3936af511..fafa0995db1 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java
@@ -1634,7 +1634,8 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
queryStatement.setOrderByComponent(
parseOrderByClause(
ctx.orderByClause(),
- ImmutableSet.of(OrderByKey.TIME, OrderByKey.DEVICE, OrderByKey.TIMESERIES)));
+ ImmutableSet.of(OrderByKey.TIME, OrderByKey.DEVICE, OrderByKey.TIMESERIES),
+ true));
}
// parse FILL
@@ -1937,7 +1938,9 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
// ---- Order By Clause
// all SortKeys should be contained by limitSet
private OrderByComponent parseOrderByClause(
- IoTDBSqlParser.OrderByClauseContext ctx, ImmutableSet<String> limitSet) {
+ IoTDBSqlParser.OrderByClauseContext ctx,
+ ImmutableSet<String> limitSet,
+ boolean allowExpression) {
OrderByComponent orderByComponent = new OrderByComponent();
Set<String> sortKeySet = new HashSet<>();
for (IoTDBSqlParser.OrderByAttributeClauseContext orderByAttributeClauseContext :
@@ -1946,7 +1949,8 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
if (orderByComponent.isUnique()) {
break;
}
- SortItem sortItem = parseOrderByAttributeClause(orderByAttributeClauseContext, limitSet);
+ SortItem sortItem =
+ parseOrderByAttributeClause(orderByAttributeClauseContext, limitSet, allowExpression);
String sortKey = sortItem.getSortKey();
if (sortKeySet.contains(sortKey)) {
@@ -1965,7 +1969,9 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
}
private SortItem parseOrderByAttributeClause(
- IoTDBSqlParser.OrderByAttributeClauseContext ctx, ImmutableSet<String> limitSet) {
+ IoTDBSqlParser.OrderByAttributeClauseContext ctx,
+ ImmutableSet<String> limitSet,
+ boolean allowExpression) {
if (ctx.sortKey() != null) {
String sortKey = ctx.sortKey().getText().toUpperCase();
if (!limitSet.contains(sortKey)) {
@@ -1974,6 +1980,11 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
}
return new SortItem(sortKey, ctx.DESC() != null ? Ordering.DESC : Ordering.ASC);
} else {
+ if (!allowExpression) {
+ throw new SemanticException(
+ "ORDER BY expression is not supported for current statement, supported sort key: "
+ + limitSet.toString());
+ }
Expression sortExpression = parseExpression(ctx.expression(), true);
return new SortItem(
sortExpression,
@@ -3719,7 +3730,8 @@ public class ASTVisitor extends IoTDBSqlParserBaseVisitor<Statement> {
OrderByKey.QUERYID,
OrderByKey.DATANODEID,
OrderByKey.ELAPSEDTIME,
- OrderByKey.STATEMENT)));
+ OrderByKey.STATEMENT),
+ false));
}
// parse LIMIT & OFFSET
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
index 8f36e44361a..69105a094d1 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/distribution/ExchangeNodeAdder.java
@@ -567,10 +567,12 @@ public class ExchangeNodeAdder extends PlanVisitor<PlanNode, NodeGroupContext> {
PlanNode newNode = node.clone();
PlanNode child = visit(node.getChildren().get(0), context);
newNode.addChild(child);
- TRegionReplicaSet dataRegion = context.getNodeDistribution(child.getPlanNodeId()).getRegion();
+ NodeDistribution nodeDistribution = context.getNodeDistribution(child.getPlanNodeId());
context.putNodeDistribution(
newNode.getPlanNodeId(),
- new NodeDistribution(NodeDistributionType.SAME_WITH_ALL_CHILDREN, dataRegion));
+ new NodeDistribution(
+ NodeDistributionType.SAME_WITH_ALL_CHILDREN,
+ nodeDistribution == null ? null : nodeDistribution.getRegion()));
return newNode;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowQueriesNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowQueriesNode.java
index c5925f93480..25e94c0da5c 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowQueriesNode.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/source/ShowQueriesNode.java
@@ -105,8 +105,8 @@ public class ShowQueriesNode extends VirtualSourceNode {
}
public static ShowQueriesNode deserialize(ByteBuffer byteBuffer) {
- PlanNodeId planNodeId = PlanNodeId.deserialize(byteBuffer);
String allowedUsername = ReadWriteIOUtils.readString(byteBuffer);
+ PlanNodeId planNodeId = PlanNodeId.deserialize(byteBuffer);
return new ShowQueriesNode(planNodeId, null, allowedUsername);
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/OrderByComponent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/OrderByComponent.java
index 8774db378f5..b6aa8c67c0a 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/OrderByComponent.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/component/OrderByComponent.java
@@ -63,6 +63,12 @@ public class OrderByComponent extends StatementNode {
orderByDevice = true;
deviceOrderPriority = sortItemList.size() - 1;
break;
+ case OrderByKey.QUERYID:
+ case OrderByKey.DATANODEID:
+ case OrderByKey.ELAPSEDTIME:
+ case OrderByKey.STATEMENT:
+ // show queries statement
+ break;
default:
throw new IllegalArgumentException(
String.format("Unknown sort key %s", sortItem.getSortKey()));
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/statistics/FragmentInstanceStatisticsDrawer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/statistics/FragmentInstanceStatisticsDrawer.java
index b9f1a55b5bd..7a8b7097eb9 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/statistics/FragmentInstanceStatisticsDrawer.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/statistics/FragmentInstanceStatisticsDrawer.java
@@ -484,7 +484,7 @@ public class FragmentInstanceStatisticsDrawer {
addLineWithValueCheck(
singleFragmentInstanceArea,
indentNum + 2,
- "Estimated Memory Size: ",
+ "Estimated Memory Size",
operatorStatistic.getMemoryUsage());
if (operatorStatistic.getSpecifiedInfoSize() != 0) {
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
index 9da2cdcf8a4..4d27792e532 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/StatementGeneratorTest.java
@@ -39,7 +39,10 @@ import org.apache.iotdb.db.queryengine.plan.relational.type.InternalTypeManager;
import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.StatementTestUtils;
import org.apache.iotdb.db.queryengine.plan.statement.StatementType;
+import org.apache.iotdb.db.queryengine.plan.statement.component.OrderByKey;
+import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
import org.apache.iotdb.db.queryengine.plan.statement.component.ResultColumn;
+import org.apache.iotdb.db.queryengine.plan.statement.component.SortItem;
import org.apache.iotdb.db.queryengine.plan.statement.crud.DeleteDataStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
@@ -60,6 +63,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.template.ShowNode
import org.apache.iotdb.db.queryengine.plan.statement.metadata.template.UnsetSchemaTemplateStatement;
import org.apache.iotdb.db.queryengine.plan.statement.metadata.view.CreateLogicalViewStatement;
import org.apache.iotdb.db.queryengine.plan.statement.sys.AuthorStatement;
+import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowQueriesStatement;
import org.apache.iotdb.isession.template.TemplateNode;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.service.rpc.thrift.TSAggregationQueryReq;
@@ -116,6 +120,35 @@ import static org.junit.Assert.assertTrue;
public class StatementGeneratorTest {
+ @Test
+ public void testShowQueries() {
+ Statement showQueries =
+ StatementGenerator.createStatement(
+ "show queries order by time, queryid, datanodeid, elapsedtime, statement",
+ ZonedDateTime.now().getOffset());
+ Assert.assertTrue(showQueries instanceof ShowQueriesStatement);
+ Assert.assertEquals(
+ ((ShowQueriesStatement) showQueries).getSortItemList().get(0),
+ new SortItem(OrderByKey.TIME, Ordering.ASC));
+ Assert.assertEquals(
+ ((ShowQueriesStatement) showQueries).getSortItemList().get(1),
+ new SortItem(OrderByKey.QUERYID, Ordering.ASC));
+ Assert.assertEquals(
+ ((ShowQueriesStatement) showQueries).getSortItemList().get(2),
+ new SortItem(OrderByKey.DATANODEID, Ordering.ASC));
+ Assert.assertEquals(
+ ((ShowQueriesStatement) showQueries).getSortItemList().get(3),
+ new SortItem(OrderByKey.ELAPSEDTIME, Ordering.ASC));
+ Assert.assertEquals(
+ ((ShowQueriesStatement) showQueries).getSortItemList().get(4),
+ new SortItem(OrderByKey.STATEMENT, Ordering.ASC));
+ Assert.assertThrows(
+ SemanticException.class,
+ () ->
+ StatementGenerator.createStatement(
+ "show queries order by a", ZonedDateTime.now().getOffset()));
+ }
+
@Test
public void testRawDataQuery() throws IllegalPathException {
TSRawDataQueryReq req =
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/LastQueryScanNodeSerdeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java
similarity index 76%
rename from iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/LastQueryScanNodeSerdeTest.java
rename to iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java
index 68a7fa0df0d..068b2bbc4b6 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/LastQueryScanNodeSerdeTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/node/source/SourceNodeSerdeTest.java
@@ -24,6 +24,7 @@ import org.apache.iotdb.commons.path.PartialPath;
import org.apache.iotdb.db.queryengine.plan.planner.node.PlanNodeDeserializeHelper;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.LastQueryScanNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.source.ShowQueriesNode;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.write.schema.MeasurementSchema;
@@ -34,9 +35,9 @@ import java.util.Arrays;
import static org.junit.Assert.assertEquals;
-public class LastQueryScanNodeSerdeTest {
+public class SourceNodeSerdeTest {
@Test
- public void test() throws IllegalPathException {
+ public void testLastQueryScanNode() throws IllegalPathException {
LastQueryScanNode node =
new LastQueryScanNode(
new PlanNodeId("test"),
@@ -69,4 +70,19 @@ public class LastQueryScanNodeSerdeTest {
byteBuffer.flip();
assertEquals(PlanNodeDeserializeHelper.deserialize(byteBuffer), node);
}
+
+ @Test
+ public void testShowQueriesNode() throws IllegalPathException {
+ ShowQueriesNode node = new ShowQueriesNode(new PlanNodeId("test"), null, "root");
+
+ ByteBuffer byteBuffer = ByteBuffer.allocate(2048);
+ node.serialize(byteBuffer);
+ byteBuffer.flip();
+ assertEquals(PlanNodeDeserializeHelper.deserialize(byteBuffer), node);
+ node = new ShowQueriesNode(new PlanNodeId("test"), null, "root");
+ byteBuffer = ByteBuffer.allocate(2048);
+ node.serialize(byteBuffer);
+ byteBuffer.flip();
+ assertEquals(PlanNodeDeserializeHelper.deserialize(byteBuffer), node);
+ }
}
(iotdb) 05/10: [AINode] More strict concurrent inference IT (#16898)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 191edf13094927c16f0f9216864e5f75d64781de
Author: Yongzao <yo...@apache.org>
AuthorDate: Fri Dec 12 16:43:52 2025 +0800
[AINode] More strict concurrent inference IT (#16898)
(cherry picked from commit 7c7b2a6c091f6a2f9fabb6b4fbfe2e737202eaec)
---
.../iotdb/ainode/it/AINodeConcurrentForecastIT.java | 2 +-
.../org/apache/iotdb/ainode/utils/AINodeTestUtils.java | 17 ++++++++++++++++-
.../ainode/core/inference/inference_request_pool.py | 4 +++-
3 files changed, 20 insertions(+), 3 deletions(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentForecastIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentForecastIT.java
index 64029c1e34b..844ec1d8223 100644
--- a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentForecastIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentForecastIT.java
@@ -49,7 +49,7 @@ public class AINodeConcurrentForecastIT {
private static final Logger LOGGER = LoggerFactory.getLogger(AINodeConcurrentForecastIT.class);
private static final String FORECAST_TABLE_FUNCTION_SQL_TEMPLATE =
- "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time, forecast_length=>%d)";
+ "SELECT * FROM FORECAST(model_id=>'%s', input=>(SELECT time,s FROM root.AI) ORDER BY time, output_length=>%d)";
@BeforeClass
public static void setUp() throws Exception {
diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java b/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java
index 0de90c42925..1d21a4d90f0 100644
--- a/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java
+++ b/integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java
@@ -20,6 +20,7 @@
package org.apache.iotdb.ainode.utils;
import com.google.common.collect.ImmutableSet;
+import org.junit.Assert;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -34,6 +35,7 @@ import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import java.util.stream.Stream;
@@ -101,6 +103,7 @@ public class AINodeTestUtils {
public static void concurrentInference(
Statement statement, String sql, int threadCnt, int loop, int expectedOutputLength)
throws InterruptedException {
+ AtomicBoolean allPass = new AtomicBoolean(true);
Thread[] threads = new Thread[threadCnt];
for (int i = 0; i < threadCnt; i++) {
threads[i] =
@@ -113,12 +116,23 @@ public class AINodeTestUtils {
while (resultSet.next()) {
outputCnt++;
}
- assertEquals(expectedOutputLength, outputCnt);
+ if (expectedOutputLength != outputCnt) {
+ allPass.set(false);
+ fail(
+ "Output count mismatch for SQL: "
+ + sql
+ + ". Expected: "
+ + expectedOutputLength
+ + ", but got: "
+ + outputCnt);
+ }
} catch (SQLException e) {
+ allPass.set(false);
fail(e.getMessage());
}
}
} catch (Exception e) {
+ allPass.set(false);
fail(e.getMessage());
}
});
@@ -130,6 +144,7 @@ public class AINodeTestUtils {
fail("Thread timeout after 10 minutes");
}
}
+ Assert.assertTrue(allPass.get());
}
public static void checkModelOnSpecifiedDevice(Statement statement, String modelId, String device)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
index a6c415a6c84..c31bcd3d762 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
@@ -120,7 +120,9 @@ class InferenceRequestPool(mp.Process):
grouped_requests = list(grouped_requests.values())
for requests in grouped_requests:
- batch_inputs = self._batcher.batch_request(requests).to(self.device)
+ batch_inputs = self._batcher.batch_request(requests).to(
+ "cpu"
+ ) # The input data should first load to CPU in current version
if isinstance(self._inference_pipeline, ForecastPipeline):
batch_output = self._inference_pipeline.forecast(
batch_inputs,
(iotdb) 02/10: [AINode] More accurate exception for model management (#16895)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit bc9418a8af3b48b9c8e8da294d8515faa891d35c
Author: Yongzao <yo...@apache.org>
AuthorDate: Thu Dec 11 16:32:56 2025 +0800
[AINode] More accurate exception for model management (#16895)
(cherry picked from commit c722aaf6c8d3bb3b040d537ec8696efecc806156)
---
.../iotdb/ainode/it/AINodeClusterConfigIT.java | 3 -
.../iotdb/ainode/it/AINodeModelManageIT.java | 4 +-
.../java/org/apache/iotdb/rpc/TSStatusCode.java | 21 +--
iotdb-core/ainode/iotdb/ainode/core/config.py | 8 +-
iotdb-core/ainode/iotdb/ainode/core/constant.py | 33 ++---
iotdb-core/ainode/iotdb/ainode/core/exception.py | 108 ++++++--------
.../core/inference/dispatcher/basic_dispatcher.py | 6 +-
.../iotdb/ainode/core/inference/pool_controller.py | 4 +-
.../iotdb/ainode/core/inference/pool_group.py | 6 +-
.../pool_scheduler/basic_pool_scheduler.py | 4 +-
.../iotdb/ainode/core/manager/inference_manager.py | 4 +-
.../iotdb/ainode/core/manager/model_manager.py | 43 +++---
.../ainode/iotdb/ainode/core/manager/utils.py | 4 +-
.../iotdb/ainode/core/model/model_constants.py | 7 -
.../ainode/iotdb/ainode/core/model/model_loader.py | 4 +-
.../iotdb/ainode/core/model/model_storage.py | 158 +++++++--------------
.../core/model/sktime/configuration_sktime.py | 36 ++---
.../ainode/core/model/sktime/modeling_sktime.py | 12 +-
.../ainode/core/model/sundial/pipeline_sundial.py | 4 +-
.../ainode/core/model/timer_xl/pipeline_timer.py | 4 +-
iotdb-core/ainode/iotdb/ainode/core/model/utils.py | 53 ++++++-
iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py | 48 ++++---
iotdb-core/ainode/iotdb/ainode/core/util/serde.py | 6 +-
.../db/exception/ainode/GetModelInfoException.java | 2 +-
.../exception/ainode/ModelNotFoundException.java | 28 ----
.../iotdb/db/protocol/client/an/AINodeClient.java | 6 +-
.../execution/config/TableConfigTaskVisitor.java | 4 +-
.../execution/config/TreeConfigTaskVisitor.java | 4 +-
.../config/executor/ClusterConfigTaskExecutor.java | 10 +-
.../config/executor/IConfigTaskExecutor.java | 2 +-
...eateTrainingTask.java => CreateTuningTask.java} | 8 +-
.../thrift-ainode/src/main/thrift/ainode.thrift | 17 +--
32 files changed, 285 insertions(+), 376 deletions(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeClusterConfigIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeClusterConfigIT.java
index 2e62f618091..e148be6b20a 100644
--- a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeClusterConfigIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeClusterConfigIT.java
@@ -107,7 +107,4 @@ public class AINodeClusterConfigIT {
}
Assert.fail("The target AINode is not removed successfully after all retries.");
}
-
- // TODO: We might need to add remove unknown test in the future, but current infrastructure is too
- // hard to implement it.
}
diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
index b92b80aecf3..3315617e7fd 100644
--- a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
@@ -131,7 +131,7 @@ public class AINodeModelManageIT {
public void dropBuiltInModelErrorTestInTree() throws SQLException {
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
Statement statement = connection.createStatement()) {
- errorTest(statement, "drop model sundial", "1510: Cannot delete built-in model: sundial");
+ errorTest(statement, "drop model sundial", "1506: Cannot delete built-in model: sundial");
}
}
@@ -139,7 +139,7 @@ public class AINodeModelManageIT {
public void dropBuiltInModelErrorTestInTable() throws SQLException {
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
Statement statement = connection.createStatement()) {
- errorTest(statement, "drop model sundial", "1510: Cannot delete built-in model: sundial");
+ errorTest(statement, "drop model sundial", "1506: Cannot delete built-in model: sundial");
}
}
diff --git a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
index 527e6eae3f4..bb533ddbbd5 100644
--- a/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
+++ b/iotdb-client/service-rpc/src/main/java/org/apache/iotdb/rpc/TSStatusCode.java
@@ -249,16 +249,17 @@ public enum TSStatusCode {
CQ_UPDATE_LAST_EXEC_TIME_ERROR(1403),
// AI
- CREATE_MODEL_ERROR(1500),
- DROP_MODEL_ERROR(1501),
- MODEL_EXIST_ERROR(1502),
- GET_MODEL_INFO_ERROR(1503),
- NO_REGISTERED_AI_NODE_ERROR(1504),
- MODEL_NOT_FOUND_ERROR(1505),
- REGISTER_AI_NODE_ERROR(1506),
- UNAVAILABLE_AI_DEVICE_ERROR(1507),
- AI_NODE_INTERNAL_ERROR(1510),
- REMOVE_AI_NODE_ERROR(1511),
+ NO_REGISTERED_AI_NODE_ERROR(1500),
+ REGISTER_AI_NODE_ERROR(1501),
+ REMOVE_AI_NODE_ERROR(1502),
+ MODEL_EXISTED_ERROR(1503),
+ MODEL_NOT_EXIST_ERROR(1504),
+ CREATE_MODEL_ERROR(1505),
+ DROP_BUILTIN_MODEL_ERROR(1506),
+ DROP_MODEL_ERROR(1507),
+ UNAVAILABLE_AI_DEVICE_ERROR(1508),
+
+ AINODE_INTERNAL_ERROR(1599), // In case somebody too lazy to add a new error code
// Pipe Plugin
CREATE_PIPE_PLUGIN_ERROR(1600),
diff --git a/iotdb-core/ainode/iotdb/ainode/core/config.py b/iotdb-core/ainode/iotdb/ainode/core/config.py
index e465df7e36d..b14efa3bedf 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/config.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/config.py
@@ -46,7 +46,7 @@ from iotdb.ainode.core.constant import (
AINODE_THRIFT_COMPRESSION_ENABLED,
AINODE_VERSION_INFO,
)
-from iotdb.ainode.core.exception import BadNodeUrlError
+from iotdb.ainode.core.exception import BadNodeUrlException
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.util.decorator import singleton
from iotdb.thrift.common.ttypes import TEndPoint
@@ -437,7 +437,7 @@ class AINodeDescriptor(object):
file_configs["ain_cluster_ingress_time_zone"]
)
- except BadNodeUrlError:
+ except BadNodeUrlException:
logger.warning("Cannot load AINode conf file, use default configuration.")
except Exception as e:
@@ -489,7 +489,7 @@ def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
"""
split = endpoint_url.split(":")
if len(split) != 2:
- raise BadNodeUrlError(endpoint_url)
+ raise BadNodeUrlException(endpoint_url)
ip = split[0]
try:
@@ -497,4 +497,4 @@ def parse_endpoint_url(endpoint_url: str) -> TEndPoint:
result = TEndPoint(ip, port)
return result
except ValueError:
- raise BadNodeUrlError(endpoint_url)
+ raise BadNodeUrlException(endpoint_url)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/constant.py b/iotdb-core/ainode/iotdb/ainode/core/constant.py
index d8f730c829c..44e76840f73 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/constant.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/constant.py
@@ -81,33 +81,18 @@ DEFAULT_CHUNK_SIZE = 8192
class TSStatusCode(Enum):
SUCCESS_STATUS = 200
REDIRECTION_RECOMMEND = 400
- MODEL_EXIST_ERROR = 1502
- MODEL_NOT_FOUND_ERROR = 1505
- UNAVAILABLE_AI_DEVICE_ERROR = 1507
- AINODE_INTERNAL_ERROR = 1510
+ MODEL_EXISTED_ERROR = 1503
+ MODEL_NOT_EXIST_ERROR = 1504
+ CREATE_MODEL_ERROR = 1505
+ DROP_BUILTIN_MODEL_ERROR = 1506
+ DROP_MODEL_ERROR = 1507
+ UNAVAILABLE_AI_DEVICE_ERROR = 1508
+
INVALID_URI_ERROR = 1511
INVALID_INFERENCE_CONFIG = 1512
INFERENCE_INTERNAL_ERROR = 1520
- def get_status_code(self) -> int:
- return self.value
-
+ AINODE_INTERNAL_ERROR = 1599 # In case somebody too lazy to add a new error code
-class HyperparameterName(Enum):
- # Training hyperparameter
- LEARNING_RATE = "learning_rate"
- EPOCHS = "epochs"
- BATCH_SIZE = "batch_size"
- USE_GPU = "use_gpu"
- NUM_WORKERS = "num_workers"
-
- # Structure hyperparameter
- KERNEL_SIZE = "kernel_size"
- INPUT_VARS = "input_vars"
- BLOCK_TYPE = "block_type"
- D_MODEL = "d_model"
- INNER_LAYERS = "inner_layer"
- OUTER_LAYERS = "outer_layer"
-
- def name(self):
+ def get_status_code(self) -> int:
return self.value
diff --git a/iotdb-core/ainode/iotdb/ainode/core/exception.py b/iotdb-core/ainode/iotdb/ainode/core/exception.py
index 30b9d54dcc7..b007ee58c48 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/exception.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/exception.py
@@ -23,7 +23,7 @@ from iotdb.ainode.core.model.model_constants import (
)
-class _BaseError(Exception):
+class _BaseException(Exception):
"""Base class for exceptions in this module."""
def __init__(self):
@@ -33,86 +33,74 @@ class _BaseError(Exception):
return self.message
-class BadNodeUrlError(_BaseError):
+class BadNodeUrlException(_BaseException):
def __init__(self, node_url: str):
+ super().__init__()
self.message = "Bad node url: {}".format(node_url)
-class ModelNotExistError(_BaseError):
- def __init__(self, msg: str):
- self.message = "Model is not exists: {} ".format(msg)
-
-
-class BadConfigValueError(_BaseError):
- def __init__(self, config_name: str, config_value, hint: str = ""):
- self.message = "Bad value [{0}] for config {1}. {2}".format(
- config_value, config_name, hint
- )
-
+# ==================== Model Management ====================
-class MissingConfigError(_BaseError):
- def __init__(self, config_name: str):
- self.message = "Missing config: {}".format(config_name)
-
-class MissingOptionError(_BaseError):
- def __init__(self, config_name: str):
- self.message = "Missing task option: {}".format(config_name)
+class ModelExistedException(_BaseException):
+ def __init__(self, model_id: str):
+ super().__init__()
+ self.message = "Model {} already exists".format(model_id)
-class RedundantOptionError(_BaseError):
- def __init__(self, option_name: str):
- self.message = "Redundant task option: {}".format(option_name)
+class ModelNotExistException(_BaseException):
+ def __init__(self, model_id: str):
+ super().__init__()
+ self.message = "Model {} does not exist".format(model_id)
-class WrongTypeConfigError(_BaseError):
- def __init__(self, config_name: str, expected_type: str):
- self.message = "Wrong type for config: {0}, expected: {1}".format(
- config_name, expected_type
+class InvalidModelUriException(_BaseException):
+ def __init__(self, msg: str):
+ super().__init__()
+ self.message = (
+ "Model registration failed because the specified uri is invalid: {}".format(
+ msg
+ )
)
-class UnsupportedError(_BaseError):
- def __init__(self, msg: str):
- self.message = "{0} is not supported in current version".format(msg)
+class BuiltInModelDeletionException(_BaseException):
+ def __init__(self, model_id: str):
+ super().__init__()
+ self.message = "Cannot delete built-in model: {}".format(model_id)
-class InvalidUriError(_BaseError):
- def __init__(self, uri: str):
- self.message = "Invalid uri: {}, there are no {} or {} under this uri.".format(
- uri, MODEL_WEIGHTS_FILE_IN_PT, MODEL_CONFIG_FILE_IN_YAML
+class BadConfigValueException(_BaseException):
+ def __init__(self, config_name: str, config_value, hint: str = ""):
+ super().__init__()
+ self.message = "Bad value [{0}] for config {1}. {2}".format(
+ config_value, config_name, hint
)
-class InvalidWindowArgumentError(_BaseError):
- def __init__(self, window_interval, window_step, dataset_length):
- self.message = f"Invalid inference input: window_interval {window_interval}, window_step {window_step}, dataset_length {dataset_length}"
-
-
-class InferenceModelInternalError(_BaseError):
+class InferenceModelInternalException(_BaseException):
def __init__(self, msg: str):
+ super().__init__()
self.message = "Inference model internal error: {0}".format(msg)
-class BuiltInModelNotSupportError(_BaseError):
+class BuiltInModelNotSupportException(_BaseException):
def __init__(self, msg: str):
+ super().__init__()
self.message = "Built-in model not support: {0}".format(msg)
-class BuiltInModelDeletionError(_BaseError):
- def __init__(self, model_id: str):
- self.message = "Cannot delete built-in model: {0}".format(model_id)
-
-
-class WrongAttributeTypeError(_BaseError):
+class WrongAttributeTypeException(_BaseException):
def __init__(self, attribute_name: str, expected_type: str):
+ super().__init__()
self.message = "Wrong type for attribute: {0}, expected: {1}".format(
attribute_name, expected_type
)
-class NumericalRangeException(_BaseError):
+class NumericalRangeException(_BaseException):
def __init__(self, attribute_name: str, value, min_value, max_value):
+ super().__init__()
self.message = (
"Attribute {0} expect value between {1} and {2}, got {3} instead.".format(
attribute_name, min_value, max_value, value
@@ -120,35 +108,19 @@ class NumericalRangeException(_BaseError):
)
-class StringRangeException(_BaseError):
+class StringRangeException(_BaseException):
def __init__(self, attribute_name: str, value: str, expect_value):
+ super().__init__()
self.message = "Attribute {0} expect value in {1}, got {2} instead.".format(
attribute_name, expect_value, value
)
-class ListRangeException(_BaseError):
+class ListRangeException(_BaseException):
def __init__(self, attribute_name: str, value: list, expected_type: str):
+ super().__init__()
self.message = (
"Attribute {0} expect value type list[{1}], got {2} instead.".format(
attribute_name, expected_type, value
)
)
-
-
-class AttributeNotSupportError(_BaseError):
- def __init__(self, model_name: str, attribute_name: str):
- self.message = "Attribute {0} is not supported in model {1}".format(
- attribute_name, model_name
- )
-
-
-# This is used to extract the key message in RuntimeError instead of the traceback message
-def runtime_error_extractor(error_message):
- pattern = re.compile(r"RuntimeError: (.+)")
- match = pattern.search(error_message)
-
- if match:
- return match.group(1)
- else:
- return ""
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
index d06ba55546f..d9a2d8663e5 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/dispatcher/basic_dispatcher.py
@@ -16,7 +16,7 @@
# under the License.
#
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.dispatcher.abstract_dispatcher import (
AbstractDispatcher,
)
@@ -41,7 +41,7 @@ class BasicDispatcher(AbstractDispatcher):
"""
model_id = req.model_id
if not pool_ids:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"No available pools for model {model_id}"
)
start_idx = hash(req.req_id) % len(pool_ids)
@@ -51,7 +51,7 @@ class BasicDispatcher(AbstractDispatcher):
state = self.pool_states[pool_id]
if state == PoolState.RUNNING:
return pool_id
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"No RUNNING pools available for model {model_id}"
)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
index 8ffa89ffd67..c580a89916d 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
@@ -24,7 +24,7 @@ from typing import Dict, Optional
import torch.multiprocessing as mp
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.inference_request import (
InferenceRequest,
InferenceRequestProxy,
@@ -374,7 +374,7 @@ class PoolController:
if not self.has_request_pools(model_id):
logger.error(f"[Inference] No pools found for model {model_id}.")
infer_proxy.set_result(None)
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
"Dispatch request failed, because no inference pools are init."
)
# TODO: Implement adaptive scaling based on requests.(e.g. lazy initialization)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
index a700dcee473..b85f64d42cc 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_group.py
@@ -19,7 +19,7 @@ from typing import Dict, Tuple
import torch.multiprocessing as mp
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.dispatcher.basic_dispatcher import BasicDispatcher
from iotdb.ainode.core.inference.inference_request import (
InferenceRequest,
@@ -90,14 +90,14 @@ class PoolGroup:
def get_request_pool(self, pool_id) -> InferenceRequestPool:
if pool_id not in self.pool_group:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
)
return self.pool_group[pool_id][0]
def get_request_queue(self, pool_id) -> mp.Queue:
if pool_id not in self.pool_group:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"[Inference][Pool-{pool_id}] Pool not found for model {self.model_id}"
)
return self.pool_group[pool_id][1]
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
index d2e7292ecd8..21140cafb1f 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
@@ -20,7 +20,7 @@ from typing import Dict, List, Optional
import torch
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.pool_group import PoolGroup
from iotdb.ainode.core.inference.pool_scheduler.abstract_pool_scheduler import (
AbstractPoolScheduler,
@@ -113,7 +113,7 @@ class BasicPoolScheduler(AbstractPoolScheduler):
if model_id not in self._request_pool_map:
pool_num = estimate_pool_size(self.DEFAULT_DEVICE, model_id)
if pool_num <= 0:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"Not enough memory to run model {model_id}."
)
return [ScaleAction(ScaleActionType.SCALE_UP, pool_num, model_id)]
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
index 1ce2e84e059..34c315274f5 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
@@ -27,7 +27,7 @@ import torch.multiprocessing as mp
from iotdb.ainode.core.config import AINodeDescriptor
from iotdb.ainode.core.constant import TSStatusCode
from iotdb.ainode.core.exception import (
- InferenceModelInternalError,
+ InferenceModelInternalException,
NumericalRangeException,
)
from iotdb.ainode.core.inference.inference_request import (
@@ -161,7 +161,7 @@ class InferenceManager:
return outputs
except Exception as e:
logger.error(e)
- raise InferenceModelInternalError(str(e))
+ raise InferenceModelInternalException(str(e))
finally:
with self._result_wrapper_lock:
del self._result_wrapper_map[req_id]
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
index 8ffb33d91e2..ef0846c3d78 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
@@ -16,12 +16,16 @@
# under the License.
#
-from typing import Any, List, Optional
+from typing import Optional
from iotdb.ainode.core.constant import TSStatusCode
-from iotdb.ainode.core.exception import BuiltInModelDeletionError
+from iotdb.ainode.core.exception import (
+ BuiltInModelDeletionException,
+ InvalidModelUriException,
+ ModelExistedException,
+ ModelNotExistException,
+)
from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.model.model_loader import load_model
from iotdb.ainode.core.model.model_storage import ModelCategory, ModelInfo, ModelStorage
from iotdb.ainode.core.rpc.status import get_status
from iotdb.ainode.core.util.decorator import singleton
@@ -47,16 +51,15 @@ class ModelManager:
req: TRegisterModelReq,
) -> TRegisterModelResp:
try:
- if self._model_storage.register_model(model_id=req.modelId, uri=req.uri):
- return TRegisterModelResp(get_status(TSStatusCode.SUCCESS_STATUS))
- return TRegisterModelResp(get_status(TSStatusCode.AINODE_INTERNAL_ERROR))
- except ValueError as e:
+ self._model_storage.register_model(model_id=req.modelId, uri=req.uri)
+ return TRegisterModelResp(get_status(TSStatusCode.SUCCESS_STATUS))
+ except ModelExistedException as e:
return TRegisterModelResp(
- get_status(TSStatusCode.INVALID_URI_ERROR, str(e))
+ get_status(TSStatusCode.MODEL_EXISTED_ERROR, str(e))
)
- except Exception as e:
+ except InvalidModelUriException as e:
return TRegisterModelResp(
- get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+ get_status(TSStatusCode.CREATE_MODEL_ERROR, str(e))
)
def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
@@ -67,12 +70,12 @@ class ModelManager:
try:
self._model_storage.delete_model(req.modelId)
return get_status(TSStatusCode.SUCCESS_STATUS)
- except BuiltInModelDeletionError as e:
- logger.warning(e)
- return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+ except ModelNotExistException as e:
+ return get_status(TSStatusCode.MODEL_NOT_EXIST_ERROR, str(e))
+ except BuiltInModelDeletionException as e:
+ return get_status(TSStatusCode.DROP_BUILTIN_MODEL_ERROR, str(e))
except Exception as e:
- logger.warning(e)
- return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+ return get_status(TSStatusCode.DROP_MODEL_ERROR, str(e))
def get_model_info(
self,
@@ -81,19 +84,9 @@ class ModelManager:
) -> Optional[ModelInfo]:
return self._model_storage.get_model_info(model_id, category)
- def get_model_infos(
- self,
- category: Optional[ModelCategory] = None,
- model_type: Optional[str] = None,
- ) -> List[ModelInfo]:
- return self._model_storage.get_model_infos(category, model_type)
-
def _refresh(self):
"""Refresh the model list (re-scan the file system)"""
self._model_storage.discover_all_models()
- def get_registered_models(self) -> List[str]:
- return self._model_storage.get_registered_models()
-
def is_model_registered(self, model_id: str) -> bool:
return self._model_storage.is_model_registered(model_id)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
index 23a98f26bbf..17516876201 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
@@ -22,7 +22,7 @@ import psutil
import torch
from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.exception import ModelNotExistError
+from iotdb.ainode.core.exception import ModelNotExistException
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.manager.model_manager import ModelManager
from iotdb.ainode.core.model.model_loader import load_model
@@ -86,7 +86,7 @@ def estimate_pool_size(device: torch.device, model_id: str) -> int:
logger.error(
f"[Inference] Cannot estimate inference pool size on device: {device}, because model: {model_id} is not supported."
)
- raise ModelNotExistError(model_id)
+ raise ModelNotExistException(model_id)
system_res = evaluate_system_resources(device)
free_mem = system_res["free_mem"]
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_constants.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_constants.py
index c42ec98551b..9f1801b5073 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/model_constants.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_constants.py
@@ -24,13 +24,6 @@ MODEL_WEIGHTS_FILE_IN_PT = "model.pt"
MODEL_CONFIG_FILE_IN_YAML = "config.yaml"
-# Model file constants
-MODEL_WEIGHTS_FILE_IN_SAFETENSORS = "model.safetensors"
-MODEL_CONFIG_FILE_IN_JSON = "config.json"
-MODEL_WEIGHTS_FILE_IN_PT = "model.pt"
-MODEL_CONFIG_FILE_IN_YAML = "config.yaml"
-
-
class ModelCategory(Enum):
BUILTIN = "builtin"
USER_DEFINED = "user_defined"
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py
index a6e3b1f7b5e..29a7c14c972 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py
@@ -32,7 +32,7 @@ from transformers import (
)
from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.exception import ModelNotExistError
+from iotdb.ainode.core.exception import ModelNotExistException
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.model.model_constants import ModelCategory
from iotdb.ainode.core.model.model_info import ModelInfo
@@ -131,7 +131,7 @@ def load_model_from_pt(model_info: ModelInfo, **kwargs):
model_file = os.path.join(model_path, "model.pt")
if not os.path.exists(model_file):
logger.error(f"Model file not found at {model_file}.")
- raise ModelNotExistError(model_file)
+ raise ModelNotExistException(model_file)
model = torch.jit.load(model_file)
if isinstance(model, torch._dynamo.eval_frame.OptimizedModule) or not acceleration:
return model
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
index 5194ed4df1b..a79371d2e79 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
@@ -23,12 +23,16 @@ import shutil
from pathlib import Path
from typing import Dict, List, Optional
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForCausalLM
from iotdb.ainode.core.config import AINodeDescriptor
from iotdb.ainode.core.constant import TSStatusCode
-from iotdb.ainode.core.exception import BuiltInModelDeletionError
+from iotdb.ainode.core.exception import (
+ BuiltInModelDeletionException,
+ ModelExistedException,
+ ModelNotExistException,
+)
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.model.model_constants import (
MODEL_CONFIG_FILE_IN_JSON,
@@ -43,6 +47,8 @@ from iotdb.ainode.core.model.model_info import (
ModelInfo,
)
from iotdb.ainode.core.model.utils import (
+ _fetch_model_from_hf_repo,
+ _fetch_model_from_local,
ensure_init_file,
get_parsed_uri,
import_class_from_path,
@@ -233,12 +239,23 @@ class ModelStorage:
# ==================== Registration Methods ====================
- def register_model(self, model_id: str, uri: str) -> bool:
+ def register_model(self, model_id: str, uri: str):
"""
- Supported URI formats:
- - repo://<huggingface_repo_id> (Maybe in the future)
- - file://<local_path>
+ Register a user-defined model from a given URI.
+ Args:
+ model_id (str): Unique identifier for the model.
+ uri (str): URI to fetch the model from.
+ Supported URI formats:
+ - file://<local_path>
+ - repo://<huggingface_repo_id> (Maybe in the future)
+ Raises:
+ ModelExistedException: If the model_id already exists.
+ InvalidModelUriException: If the URI format is invalid.
"""
+
+ if self.is_model_registered(model_id):
+ raise ModelExistedException(model_id)
+
uri_type = parse_uri_type(uri)
parsed_uri = get_parsed_uri(uri)
@@ -249,9 +266,9 @@ class ModelStorage:
ensure_init_file(model_dir)
if uri_type == UriType.REPO:
- self._fetch_model_from_hf_repo(parsed_uri, model_dir)
+ _fetch_model_from_hf_repo(parsed_uri, model_dir)
else:
- self._fetch_model_from_local(os.path.expanduser(parsed_uri), model_dir)
+ _fetch_model_from_local(os.path.expanduser(parsed_uri), model_dir)
config_path, _ = validate_model_files(model_dir)
config = load_model_config_in_json(config_path)
@@ -272,7 +289,7 @@ class ModelStorage:
self._models[ModelCategory.USER_DEFINED.value][model_id] = model_info
if auto_map:
- # Transformers model: immediately register to Transformers auto-loading mechanism
+ # Transformers model: immediately register to Transformers autoloading mechanism
success = self._register_transformers_model(model_info)
if success:
with self._lock_pool.get_lock(model_id).write_lock():
@@ -281,46 +298,15 @@ class ModelStorage:
with self._lock_pool.get_lock(model_id).write_lock():
model_info.state = ModelStates.INACTIVE
logger.error(f"Failed to register Transformers model {model_id}")
- return False
else:
# Other type models: only log
self._register_other_model(model_info)
logger.info(f"Successfully registered model {model_id} from URI: {uri}")
- return True
- def _fetch_model_from_hf_repo(self, repo_id: str, storage_path: str):
- logger.info(
- f"Downloading model from HuggingFace repository: {repo_id} -> {storage_path}"
- )
- # Use snapshot_download to download entire repository (including config.json and model.safetensors)
- try:
- snapshot_download(
- repo_id=repo_id,
- local_dir=storage_path,
- local_dir_use_symlinks=False,
- )
- except Exception as e:
- logger.error(f"Failed to download model from HuggingFace: {e}")
- raise
-
- def _fetch_model_from_local(self, source_path: str, storage_path: str):
- logger.info(f"Copying model from local path: {source_path} -> {storage_path}")
- source_dir = Path(source_path)
- if not source_dir.is_dir():
- raise ValueError(
- f"Source path does not exist or is not a directory: {source_path}"
- )
-
- storage_dir = Path(storage_path)
- for file in source_dir.iterdir():
- if file.is_file():
- shutil.copy2(file, storage_dir / file.name)
- return
-
- def _register_transformers_model(self, model_info: ModelInfo) -> bool:
+ def _register_transformers_model(self, model_info: ModelInfo):
"""
- Register Transformers model to auto-loading mechanism (internal method)
+ Register Transformers model to autoloading mechanism (internal method)
"""
auto_map = model_info.auto_map
if not auto_map:
@@ -350,7 +336,6 @@ class ModelStorage:
logger.info(
f"Registered AutoModelForCausalLM: {config_class.__name__} -> {auto_model_path}"
)
-
return True
except Exception as e:
logger.warning(
@@ -471,8 +456,16 @@ class ModelStorage:
stateMap=state_map,
)
- def delete_model(self, model_id: str) -> None:
- # Use write lock to protect entire deletion process
+ def delete_model(self, model_id: str):
+ """
+ Delete a user-defined model by model_id.
+ Args:
+ model_id (str): Unique identifier for the model to be deleted.
+ Raises:
+ ModelNotExistException: If the model_id does not exist.
+ BuiltInModelDeletionException: If attempting to delete a built-in model.
+ Others: Any exceptions raised during file deletion.
+ """
with self._lock_pool.get_lock(model_id).write_lock():
model_info = None
category_value = None
@@ -481,30 +474,25 @@ class ModelStorage:
model_info = category_dict[model_id]
category_value = cat_value
break
-
if not model_info:
logger.warning(f"Model {model_id} does not exist, cannot delete")
- return
-
+ raise ModelNotExistException(model_id)
if model_info.category == ModelCategory.BUILTIN:
- raise BuiltInModelDeletionError(model_id)
+ logger.warning(f"Model {model_id} is builtin, cannot delete")
+ raise BuiltInModelDeletionException(model_id)
model_info.state = ModelStates.DROPPING
model_path = os.path.join(
self._models_dir, model_info.category.value, model_id
)
- if model_path.exists():
+ if os.path.exists(model_path):
try:
shutil.rmtree(model_path)
- logger.info(f"Deleted model directory: {model_path}")
+ logger.info(f"Model directory is deleted: {model_path}")
except Exception as e:
logger.error(f"Failed to delete model directory {model_path}: {e}")
- raise
-
- if category_value and model_id in self._models[category_value]:
- del self._models[category_value][model_id]
- logger.info(f"Model {model_id} has been removed from storage")
-
- return
+ raise e
+ del self._models[category_value][model_id]
+ logger.info(f"Model {model_id} has been removed from model storage")
# ==================== Query Methods ====================
@@ -512,10 +500,14 @@ class ModelStorage:
self, model_id: str, category: Optional[ModelCategory] = None
) -> Optional[ModelInfo]:
"""
- Get single model information
-
- If category is specified, use model_id's lock
- If category is not specified, need to traverse all dictionaries, use global lock
+ Get specified model information.
+ Args:
+ model_id (str): Unique identifier for the model.
+ category (Optional[ModelCategory]): Category of the model (if known).
+ Returns:
+ ModelInfo: Information of the specified model.
+ Raises:
+ ModelNotExistException: If the model_id does not exist.
"""
if category:
# Category specified, only need to access specific dictionary, use model_id's lock
@@ -527,39 +519,7 @@ class ModelStorage:
for category_dict in self._models.values():
if model_id in category_dict:
return category_dict[model_id]
- return None
-
- def get_model_infos(
- self, category: Optional[ModelCategory] = None, model_type: Optional[str] = None
- ) -> List[ModelInfo]:
- """
- Get model information list
-
- Note: Since we need to traverse all models, use a global lock to protect the entire dictionary structure
- For single model access, using model_id-based lock would be more efficient
- """
- matching_models = []
-
- # For traversal operations, we need to protect the entire dictionary structure
- # Use a special lock (using empty string as key) to protect the entire dictionary
- with self._lock_pool.get_lock("").read_lock():
- if category and model_type:
- for model_info in self._models[category.value].values():
- if model_info.model_type == model_type:
- matching_models.append(model_info)
- return matching_models
- elif category:
- return list(self._models[category.value].values())
- elif model_type:
- for category_dict in self._models.values():
- for model_info in category_dict.values():
- if model_info.model_type == model_type:
- matching_models.append(model_info)
- return matching_models
- else:
- for category_dict in self._models.values():
- matching_models.extend(category_dict.values())
- return matching_models
+ raise ModelNotExistException(model_id)
def is_model_registered(self, model_id: str) -> bool:
"""Check if model is registered (search in _models)"""
@@ -572,11 +532,3 @@ class ModelStorage:
if model_id in category_dict:
return True
return False
-
- def get_registered_models(self) -> List[str]:
- """Get list of all registered model IDs"""
- with self._lock_pool.get_lock("").read_lock():
- model_ids = []
- for category_dict in self._models.values():
- model_ids.extend(category_dict.keys())
- return model_ids
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py
index 261de3c9abe..d9d20545af6 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py
@@ -20,11 +20,11 @@ from dataclasses import dataclass, field
from typing import Any, Dict, List, Union
from iotdb.ainode.core.exception import (
- BuiltInModelNotSupportError,
+ BuiltInModelNotSupportException,
ListRangeException,
NumericalRangeException,
StringRangeException,
- WrongAttributeTypeError,
+ WrongAttributeTypeException,
)
from iotdb.ainode.core.log import Logger
@@ -49,7 +49,7 @@ class AttributeConfig:
if value is None:
return True # Allow None for optional int parameters
if not isinstance(value, int):
- raise WrongAttributeTypeError(self.name, "int")
+ raise WrongAttributeTypeException(self.name, "int")
if self.low is not None and self.high is not None:
if not (self.low <= value <= self.high):
raise NumericalRangeException(self.name, value, self.low, self.high)
@@ -57,7 +57,7 @@ class AttributeConfig:
if value is None:
return True # Allow None for optional float parameters
if not isinstance(value, (int, float)):
- raise WrongAttributeTypeError(self.name, "float")
+ raise WrongAttributeTypeException(self.name, "float")
value = float(value)
if self.low is not None and self.high is not None:
if not (self.low <= value <= self.high):
@@ -66,26 +66,26 @@ class AttributeConfig:
if value is None:
return True # Allow None for optional str parameters
if not isinstance(value, str):
- raise WrongAttributeTypeError(self.name, "str")
+ raise WrongAttributeTypeException(self.name, "str")
if self.choices and value not in self.choices:
raise StringRangeException(self.name, value, self.choices)
elif self.type == "bool":
if value is None:
return True # Allow None for optional bool parameters
if not isinstance(value, bool):
- raise WrongAttributeTypeError(self.name, "bool")
+ raise WrongAttributeTypeException(self.name, "bool")
elif self.type == "list":
if not isinstance(value, list):
- raise WrongAttributeTypeError(self.name, "list")
+ raise WrongAttributeTypeException(self.name, "list")
for item in value:
if not isinstance(item, self.value_type):
- raise WrongAttributeTypeError(self.name, self.value_type)
+ raise WrongAttributeTypeException(self.name, self.value_type)
elif self.type == "tuple":
if not isinstance(value, tuple):
- raise WrongAttributeTypeError(self.name, "tuple")
+ raise WrongAttributeTypeException(self.name, "tuple")
for item in value:
if not isinstance(item, self.value_type):
- raise WrongAttributeTypeError(self.name, self.value_type)
+ raise WrongAttributeTypeException(self.name, self.value_type)
return True
def parse(self, string_value: str):
@@ -96,14 +96,14 @@ class AttributeConfig:
try:
return int(string_value)
except:
- raise WrongAttributeTypeError(self.name, "int")
+ raise WrongAttributeTypeException(self.name, "int")
elif self.type == "float":
if string_value.lower() == "none" or string_value.strip() == "":
return None
try:
return float(string_value)
except:
- raise WrongAttributeTypeError(self.name, "float")
+ raise WrongAttributeTypeException(self.name, "float")
elif self.type == "str":
if string_value.lower() == "none" or string_value.strip() == "":
return None
@@ -116,14 +116,14 @@ class AttributeConfig:
elif string_value.lower() == "none" or string_value.strip() == "":
return None
else:
- raise WrongAttributeTypeError(self.name, "bool")
+ raise WrongAttributeTypeException(self.name, "bool")
elif self.type == "list":
try:
list_value = eval(string_value)
except:
- raise WrongAttributeTypeError(self.name, "list")
+ raise WrongAttributeTypeException(self.name, "list")
if not isinstance(list_value, list):
- raise WrongAttributeTypeError(self.name, "list")
+ raise WrongAttributeTypeException(self.name, "list")
for i in range(len(list_value)):
try:
list_value[i] = self.value_type(list_value[i])
@@ -136,9 +136,9 @@ class AttributeConfig:
try:
tuple_value = eval(string_value)
except:
- raise WrongAttributeTypeError(self.name, "tuple")
+ raise WrongAttributeTypeException(self.name, "tuple")
if not isinstance(tuple_value, tuple):
- raise WrongAttributeTypeError(self.name, "tuple")
+ raise WrongAttributeTypeException(self.name, "tuple")
list_value = list(tuple_value)
for i in range(len(list_value)):
try:
@@ -390,7 +390,7 @@ def get_attributes(model_id: str) -> Dict[str, AttributeConfig]:
"""Get attribute configuration for Sktime model"""
model_id = "EXPONENTIAL_SMOOTHING" if model_id == "HOLTWINTERS" else model_id
if model_id not in MODEL_CONFIGS:
- raise BuiltInModelNotSupportError(model_id)
+ raise BuiltInModelNotSupportException(model_id)
return MODEL_CONFIGS[model_id]
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py
index eca812d35ec..9ddbcab286f 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py
@@ -30,8 +30,8 @@ from sktime.forecasting.naive import NaiveForecaster
from sktime.forecasting.trend import STLForecaster
from iotdb.ainode.core.exception import (
- BuiltInModelNotSupportError,
- InferenceModelInternalError,
+ BuiltInModelNotSupportException,
+ InferenceModelInternalException,
)
from iotdb.ainode.core.log import Logger
@@ -66,7 +66,7 @@ class ForecastingModel(SktimeModel):
output = self._model.predict(fh=range(predict_length))
return np.array(output, dtype=np.float64)
except Exception as e:
- raise InferenceModelInternalError(str(e))
+ raise InferenceModelInternalException(str(e))
class DetectionModel(SktimeModel):
@@ -82,7 +82,7 @@ class DetectionModel(SktimeModel):
else:
return np.array(output, dtype=np.int32)
except Exception as e:
- raise InferenceModelInternalError(str(e))
+ raise InferenceModelInternalException(str(e))
class ArimaModel(ForecastingModel):
@@ -155,7 +155,7 @@ class STRAYModel(DetectionModel):
scaled_data = pd.Series(scaled_data.flatten())
return super().generate(scaled_data, **kwargs)
except Exception as e:
- raise InferenceModelInternalError(str(e))
+ raise InferenceModelInternalException(str(e))
# Model factory mapping
@@ -176,5 +176,5 @@ def create_sktime_model(model_id: str, **kwargs) -> SktimeModel:
attributes = update_attribute({**kwargs}, get_attributes(model_id.upper()))
model_class = _MODEL_FACTORY.get(model_id.upper())
if model_class is None:
- raise BuiltInModelNotSupportError(model_id)
+ raise BuiltInModelNotSupportException(model_id)
return model_class(attributes)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/pipeline_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/pipeline_sundial.py
index 85b6f7db2ff..ee128802d24 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/pipeline_sundial.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/pipeline_sundial.py
@@ -18,7 +18,7 @@
import torch
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
@@ -28,7 +28,7 @@ class SundialPipeline(ForecastPipeline):
def _preprocess(self, inputs):
if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
)
return inputs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timer_xl/pipeline_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timer_xl/pipeline_timer.py
index c0f00b1f5ca..65c6cdd74cd 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/timer_xl/pipeline_timer.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/timer_xl/pipeline_timer.py
@@ -18,7 +18,7 @@
import torch
-from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.exception import InferenceModelInternalException
from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
@@ -28,7 +28,7 @@ class TimerPipeline(ForecastPipeline):
def _preprocess(self, inputs):
if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
+ raise InferenceModelInternalException(
f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
)
return inputs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/utils.py b/iotdb-core/ainode/iotdb/ainode/core/model/utils.py
index 1cd0ee44912..815232c52b0 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/utils.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/utils.py
@@ -19,25 +19,31 @@
import importlib
import json
import os.path
+import shutil
import sys
from contextlib import contextmanager
+from pathlib import Path
from typing import Dict, Tuple
+from huggingface_hub import snapshot_download
+
+from iotdb.ainode.core.exception import InvalidModelUriException
+from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.model.model_constants import (
MODEL_CONFIG_FILE_IN_JSON,
MODEL_WEIGHTS_FILE_IN_SAFETENSORS,
UriType,
)
+logger = Logger()
+
def parse_uri_type(uri: str) -> UriType:
- if uri.startswith("repo://"):
- return UriType.REPO
- elif uri.startswith("file://"):
+ if uri.startswith("file://"):
return UriType.FILE
else:
- raise ValueError(
- f"Unsupported URI type: {uri}. Supported formats: repo:// or file://"
+ raise InvalidModelUriException(
+ f"Unknown uri type {uri}, currently supporting formats: file://"
)
@@ -70,9 +76,13 @@ def validate_model_files(model_dir: str) -> Tuple[str, str]:
weights_path = os.path.join(model_dir, MODEL_WEIGHTS_FILE_IN_SAFETENSORS)
if not os.path.exists(config_path):
- raise ValueError(f"Model config file does not exist: {config_path}")
+ raise InvalidModelUriException(
+ f"Model config file does not exist: {config_path}"
+ )
if not os.path.exists(weights_path):
- raise ValueError(f"Model weights file does not exist: {weights_path}")
+ raise InvalidModelUriException(
+ f"Model weights file does not exist: {weights_path}"
+ )
# Create __init__.py file to ensure model directory can be imported as a module
init_file = os.path.join(model_dir, "__init__.py")
@@ -96,3 +106,32 @@ def ensure_init_file(dir_path: str):
if not os.path.exists(init_file):
with open(init_file, "w"):
pass
+
+
+def _fetch_model_from_local(source_path: str, storage_path: str):
+ logger.info(f"Copying model from local path: {source_path} -> {storage_path}")
+ source_dir = Path(source_path)
+ if not source_dir.exists():
+ raise InvalidModelUriException(f"Source path does not exist: {source_path}")
+ if not source_dir.is_dir():
+ raise InvalidModelUriException(f"Source path is not a directory: {source_path}")
+ storage_dir = Path(storage_path)
+ for file in source_dir.iterdir():
+ if file.is_file():
+ shutil.copy2(file, storage_dir / file.name)
+
+
+def _fetch_model_from_hf_repo(repo_id: str, storage_path: str):
+ logger.info(
+ f"Downloading model from HuggingFace repository: {repo_id} -> {storage_path}"
+ )
+ # Use snapshot_download to download entire repository (including config.json and model.safetensors)
+ try:
+ snapshot_download(
+ repo_id=repo_id,
+ local_dir=storage_path,
+ local_dir_use_symlinks=False,
+ )
+ except Exception as e:
+ logger.error(f"Failed to download model from HuggingFace: {e}")
+ raise
diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
index 6c4eedeb99f..492802fc060 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
@@ -40,7 +40,7 @@ from iotdb.thrift.ainode.ttypes import (
TShowLoadedModelsResp,
TShowModelsReq,
TShowModelsResp,
- TTrainingReq,
+ TTuningReq,
TUnloadModelReq,
)
from iotdb.thrift.common.ttypes import TSStatus
@@ -56,8 +56,8 @@ def _ensure_device_id_is_available(device_id_list: list[str]) -> TSStatus:
for device_id in device_id_list:
if device_id not in available_devices:
return TSStatus(
- code=TSStatusCode.INVALID_URI_ERROR.value,
- message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
+ code=TSStatusCode.UNAVAILABLE_AI_DEVICE_ERROR.value,
+ message=f"AIDevice ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
)
return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value)
@@ -68,7 +68,9 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
self._model_manager = ModelManager()
self._inference_manager = InferenceManager()
- def stop(self) -> None:
+ # ==================== Cluster Management ====================
+
+ def stop(self):
logger.info("Stopping the RPC service handler of IoTDB-AINode...")
self._inference_manager.stop()
@@ -76,6 +78,17 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
self._ainode.stop()
return get_status(TSStatusCode.SUCCESS_STATUS, "AINode stopped successfully.")
+ def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp:
+ return ClusterManager.get_heart_beat(req)
+
+ def showAIDevices(self) -> TShowAIDevicesResp:
+ return TShowAIDevicesResp(
+ status=TSStatus(code=TSStatusCode.SUCCESS_STATUS.value),
+ deviceIdList=get_available_devices(),
+ )
+
+ # ==================== Model Management ====================
+
def registerModel(self, req: TRegisterModelReq) -> TRegisterModelResp:
return self._model_manager.register_model(req)
@@ -109,11 +122,15 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
return TShowLoadedModelsResp(status=status, deviceLoadedModelsMap={})
return self._inference_manager.show_loaded_models(req)
- def showAIDevices(self) -> TShowAIDevicesResp:
- return TShowAIDevicesResp(
- status=TSStatus(code=TSStatusCode.SUCCESS_STATUS.value),
- deviceIdList=get_available_devices(),
- )
+ def _ensure_model_is_registered(self, model_id: str) -> TSStatus:
+ if not self._model_manager.is_model_registered(model_id):
+ return TSStatus(
+ code=TSStatusCode.MODEL_NOT_EXIST_ERROR.value,
+ message=f"Model [{model_id}] is not registered yet. You can use 'SHOW MODELS' to retrieve the available models.",
+ )
+ return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value)
+
+ # ==================== Inference ====================
def inference(self, req: TInferenceReq) -> TInferenceResp:
status = self._ensure_model_is_registered(req.modelId)
@@ -127,16 +144,7 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
return TForecastResp(status, [])
return self._inference_manager.forecast(req)
- def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp:
- return ClusterManager.get_heart_beat(req)
+ # ==================== Tuning ====================
- def createTrainingTask(self, req: TTrainingReq) -> TSStatus:
+ def createTuningTask(self, req: TTuningReq) -> TSStatus:
pass
-
- def _ensure_model_is_registered(self, model_id: str) -> TSStatus:
- if not self._model_manager.is_model_registered(model_id):
- return TSStatus(
- code=TSStatusCode.MODEL_NOT_FOUND_ERROR.value,
- message=f"Model [{model_id}] is not registered yet. You can use 'SHOW MODELS' to retrieve the available models.",
- )
- return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/util/serde.py b/iotdb-core/ainode/iotdb/ainode/core/util/serde.py
index f8188209a37..9c6020019fc 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/util/serde.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/util/serde.py
@@ -21,7 +21,7 @@ from enum import Enum
import numpy as np
import pandas as pd
-from iotdb.ainode.core.exception import BadConfigValueError
+from iotdb.ainode.core.exception import BadConfigValueException
class TSDataType(Enum):
@@ -122,7 +122,7 @@ def _get_type_in_byte(data_type: pd.Series):
elif data_type == "text":
return b"\x05"
else:
- raise BadConfigValueError(
+ raise BadConfigValueException(
"data_type",
data_type,
"data_type should be in ['bool', 'int32', 'int64', 'float32', 'float64', 'text']",
@@ -138,7 +138,7 @@ def get_data_type_byte_from_str(value):
byte: corresponding data type in [b'\x00', b'\x01', b'\x02', b'\x03', b'\x04', b'\x05']
"""
if value not in ["bool", "int32", "int64", "float32", "float64", "text"]:
- raise BadConfigValueError(
+ raise BadConfigValueException(
"data_type",
value,
"data_type should be in ['bool', 'int32', 'int64', 'float32', 'float64', 'text']",
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java
index 03402d30c64..b8d98c9c1d7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/GetModelInfoException.java
@@ -23,6 +23,6 @@ import org.apache.iotdb.rpc.TSStatusCode;
public class GetModelInfoException extends ModelException {
public GetModelInfoException(String message) {
- super(message, TSStatusCode.GET_MODEL_INFO_ERROR);
+ super(message, TSStatusCode.AINODE_INTERNAL_ERROR);
}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java
deleted file mode 100644
index 38a5105cded..00000000000
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/ainode/ModelNotFoundException.java
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * 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.db.exception.ainode;
-
-import org.apache.iotdb.rpc.TSStatusCode;
-
-public class ModelNotFoundException extends ModelException {
- public ModelNotFoundException(String message) {
- super(message, TSStatusCode.MODEL_NOT_FOUND_ERROR);
- }
-}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/an/AINodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/an/AINodeClient.java
index 5eaffc40af9..ffad889a223 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/an/AINodeClient.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/an/AINodeClient.java
@@ -35,7 +35,7 @@ import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsReq;
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsResp;
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsResp;
-import org.apache.iotdb.ainode.rpc.thrift.TTrainingReq;
+import org.apache.iotdb.ainode.rpc.thrift.TTuningReq;
import org.apache.iotdb.ainode.rpc.thrift.TUnloadModelReq;
import org.apache.iotdb.common.rpc.thrift.TAINodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
@@ -129,8 +129,8 @@ public class AINodeClient implements IAINodeRPCService.Iface, AutoCloseable, Thr
}
@Override
- public TSStatus createTrainingTask(TTrainingReq req) throws TException {
- return executeRemoteCallWithRetry(() -> client.createTrainingTask(req));
+ public TSStatus createTuningTask(TTuningReq req) throws TException {
+ return executeRemoteCallWithRetry(() -> client.createTuningTask(req));
}
@Override
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java
index f9d668944e0..d659f30fc8f 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TableConfigTaskVisitor.java
@@ -63,7 +63,7 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowPipePl
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowRegionTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowVariablesTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateModelTask;
-import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateTrainingTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateTuningTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.DropModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.LoadModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.ShowAIDevicesTask;
@@ -1530,7 +1530,7 @@ public class TableConfigTaskVisitor extends AstVisitor<IConfigTask, MPPQueryCont
protected IConfigTask visitCreateTraining(CreateTraining node, MPPQueryContext context) {
context.setQueryType(QueryType.WRITE);
accessControl.checkUserGlobalSysPrivilege(context);
- return new CreateTrainingTask(
+ return new CreateTuningTask(
node.getModelId(), node.getParameters(), node.getExistingModelId(), node.getTargetSql());
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java
index cb4d05f1b99..ed5e2e434f4 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java
@@ -68,7 +68,7 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowTrigge
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ShowVariablesTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.UnSetTTLTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateModelTask;
-import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateTrainingTask;
+import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.CreateTuningTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.DropModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.LoadModelTask;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai.ShowAIDevicesTask;
@@ -926,7 +926,7 @@ public class TreeConfigTaskVisitor extends StatementVisitor<IConfigTask, MPPQuer
for (PartialPath partialPath : partialPathList) {
targetPathPatterns.add(partialPath.getFullPath());
}
- return new CreateTrainingTask(
+ return new CreateTuningTask(
createTrainingStatement.getModelId(),
createTrainingStatement.getParameters(),
createTrainingStatement.getTargetTimeRanges(),
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
index 5f27b3feb8c..e33cc615407 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
@@ -28,7 +28,7 @@ import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsReq;
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsResp;
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsResp;
-import org.apache.iotdb.ainode.rpc.thrift.TTrainingReq;
+import org.apache.iotdb.ainode.rpc.thrift.TTuningReq;
import org.apache.iotdb.ainode.rpc.thrift.TUnloadModelReq;
import org.apache.iotdb.common.rpc.thrift.FunctionType;
import org.apache.iotdb.common.rpc.thrift.Model;
@@ -3748,7 +3748,7 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
}
@Override
- public SettableFuture<ConfigTaskResult> createTraining(
+ public SettableFuture<ConfigTaskResult> createTuningTask(
String modelId,
boolean isTableModel,
Map<String, String> parameters,
@@ -3757,9 +3757,9 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
@Nullable String targetSql,
@Nullable List<String> pathList) {
final SettableFuture<ConfigTaskResult> future = SettableFuture.create();
- try (final AINodeClient ai =
+ try (final AINodeClient aiNodeClient =
AINodeClientManager.getInstance().borrowClient(AINodeClientManager.AINODE_ID_PLACEHOLDER)) {
- final TTrainingReq req = new TTrainingReq();
+ final TTuningReq req = new TTuningReq();
req.setModelId(modelId);
req.setParameters(parameters);
if (existingModelId != null) {
@@ -3768,7 +3768,7 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor {
if (existingModelId != null) {
req.setExistingModelId(existingModelId);
}
- final TSStatus status = ai.createTrainingTask(req);
+ final TSStatus status = aiNodeClient.createTuningTask(req);
if (TSStatusCode.SUCCESS_STATUS.getStatusCode() != status.getCode()) {
future.setException(new IoTDBException(status));
} else {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
index 85455756b05..2d7c0f6d1f3 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java
@@ -443,7 +443,7 @@ public interface IConfigTaskExecutor {
SettableFuture<ConfigTaskResult> unloadModel(String existingModelId, List<String> deviceIdList);
- SettableFuture<ConfigTaskResult> createTraining(
+ SettableFuture<ConfigTaskResult> createTuningTask(
String modelId,
boolean isTableModel,
Map<String, String> parameters,
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTrainingTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTuningTask.java
similarity index 93%
rename from iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTrainingTask.java
rename to iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTuningTask.java
index 9c93c5b7577..a66b7e700b7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTrainingTask.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/CreateTuningTask.java
@@ -28,7 +28,7 @@ import com.google.common.util.concurrent.ListenableFuture;
import java.util.List;
import java.util.Map;
-public class CreateTrainingTask implements IConfigTask {
+public class CreateTuningTask implements IConfigTask {
private final String modelId;
private final boolean isTableModel;
@@ -43,7 +43,7 @@ public class CreateTrainingTask implements IConfigTask {
private List<List<Long>> timeRanges;
// For table model
- public CreateTrainingTask(
+ public CreateTuningTask(
String modelId, Map<String, String> parameters, String existingModelId, String targetSql) {
this.modelId = modelId;
this.parameters = parameters;
@@ -53,7 +53,7 @@ public class CreateTrainingTask implements IConfigTask {
}
// For tree model
- public CreateTrainingTask(
+ public CreateTuningTask(
String modelId,
Map<String, String> parameters,
List<List<Long>> timeRanges,
@@ -71,7 +71,7 @@ public class CreateTrainingTask implements IConfigTask {
@Override
public ListenableFuture<ConfigTaskResult> execute(IConfigTaskExecutor configTaskExecutor)
throws InterruptedException {
- return configTaskExecutor.createTraining(
+ return configTaskExecutor.createTuningTask(
modelId, isTableModel, parameters, timeRanges, existingModelId, targetSql, targetPaths);
}
}
diff --git a/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift b/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift
index 1dc2f025f5c..cda356a948e 100644
--- a/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift
+++ b/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift
@@ -73,7 +73,7 @@ struct IDataSchema {
2: optional list<i64> timeRange
}
-struct TTrainingReq {
+struct TTuningReq {
1: required string dbType
2: required string modelId
3: required string existingModelId
@@ -131,30 +131,27 @@ struct TUnloadModelReq {
service IAINodeRPCService {
- // -------------- For Config Node --------------
common.TSStatus stopAINode()
+ TAIHeartbeatResp getAIHeartbeat(TAIHeartbeatReq req)
+
+ TShowAIDevicesResp showAIDevices()
+
TShowModelsResp showModels(TShowModelsReq req)
TShowLoadedModelsResp showLoadedModels(TShowLoadedModelsReq req)
- TShowAIDevicesResp showAIDevices()
-
common.TSStatus deleteModel(TDeleteModelReq req)
TRegisterModelResp registerModel(TRegisterModelReq req)
- TAIHeartbeatResp getAIHeartbeat(TAIHeartbeatReq req)
-
- common.TSStatus createTrainingTask(TTrainingReq req)
-
common.TSStatus loadModel(TLoadModelReq req)
common.TSStatus unloadModel(TUnloadModelReq req)
- // -------------- For Data Node --------------
-
TInferenceResp inference(TInferenceReq req)
TForecastResp forecast(TForecastReq req)
+
+ common.TSStatus createTuningTask(TTuningReq req)
}
\ No newline at end of file
(iotdb) 01/10: Throw exception while readObject column is not OBJECT type
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 8f1235835ae17ec620ecb9a48ae0a2716fed0dd7
Author: JackieTien97 <ja...@gmail.com>
AuthorDate: Mon Dec 15 18:10:14 2025 +0800
Throw exception while readObject column is not OBJECT type
---
.../execution/operator/process/function/partition/Slice.java | 3 +++
.../operator/source/relational/aggregation/RecordIterator.java | 3 +++
2 files changed, 6 insertions(+)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/function/partition/Slice.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/function/partition/Slice.java
index a34a0eb650d..57f04be9e00 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/function/partition/Slice.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/process/function/partition/Slice.java
@@ -221,6 +221,9 @@ public class Slice {
@Override
public Binary readObject(int columnIndex, long offset, int length) {
+ if (getDataType(columnIndex) != Type.OBJECT) {
+ throw new UnsupportedOperationException("current column is not object column");
+ }
Binary binary = getBinarySafely(columnIndex);
return new Binary(ObjectTypeUtils.readObjectContent(binary, offset, length, true).array());
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/aggregation/RecordIterator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/aggregation/RecordIterator.java
index ae66a11729d..234723b229b 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/aggregation/RecordIterator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/aggregation/RecordIterator.java
@@ -166,6 +166,9 @@ public class RecordIterator implements Iterator<Record> {
@Override
public Binary readObject(int columnIndex, long offset, int length) {
+ if (getDataType(columnIndex) != Type.OBJECT) {
+ throw new UnsupportedOperationException("current column is not object column");
+ }
Binary binary = getBinarySafely(columnIndex);
return new Binary(ObjectTypeUtils.readObjectContent(binary, offset, length, true).array());
}
(iotdb) 10/10: Optimized the error message & prevented stack printing of information_schema selection when the cluster is closing (#16905)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 3bc6101a93d5ad0cc438991b7fde37194bdc434e
Author: Caideyipi <87...@users.noreply.github.com>
AuthorDate: Tue Dec 16 09:16:11 2025 +0800
Optimized the error message & prevented stack printing of information_schema selection when the cluster is closing (#16905)
* DL
* ancient-IT
(cherry picked from commit c77bb57cb2e9e99f2511c7bb59894773f873b720)
---
.../java/org/apache/iotdb/relational/it/schema/IoTDBTableIT.java | 6 ++++++
.../planner/distribute/TableDistributedPlanGenerator.java | 8 +++++++-
2 files changed, 13 insertions(+), 1 deletion(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBTableIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBTableIT.java
index 748a6937bcf..5eaf230b600 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBTableIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBTableIT.java
@@ -1229,6 +1229,12 @@ public class IoTDBTableIT {
} catch (final SQLException e) {
assertEquals("701: The system view does not support show create.", e.getMessage());
}
+ try {
+ statement.execute("show create table information_schema.tables");
+ fail();
+ } catch (final SQLException e) {
+ assertEquals("701: The system view does not support show create.", e.getMessage());
+ }
try {
statement.execute("create or replace view a () as root.b.**");
fail();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
index 069d1df7469..e13d09de4a5 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/distribute/TableDistributedPlanGenerator.java
@@ -23,6 +23,7 @@ import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot;
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.exception.IoTDBRuntimeException;
import org.apache.iotdb.commons.partition.DataPartition;
import org.apache.iotdb.commons.partition.SchemaPartition;
import org.apache.iotdb.commons.schema.table.TsTable;
@@ -100,6 +101,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
import org.apache.iotdb.db.queryengine.plan.statement.component.Ordering;
import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache;
import org.apache.iotdb.db.schemaengine.table.DataNodeTreeViewSchemaUtils;
+import org.apache.iotdb.rpc.TSStatusCode;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
@@ -1004,7 +1006,11 @@ public class TableDistributedPlanGenerator
List<TDataNodeLocation> dataNodeLocations =
dataNodeLocationSupplier.getDataNodeLocations(
node.getQualifiedObjectName().getObjectName());
- checkArgument(!dataNodeLocations.isEmpty(), "DataNodeLocations shouldn't be empty");
+ if (dataNodeLocations.isEmpty()) {
+ throw new IoTDBRuntimeException(
+ "No available dataNodes, may be the cluster is closing",
+ TSStatusCode.NO_AVAILABLE_REPLICA.getStatusCode());
+ }
List<PlanNode> resultTableScanNodeList = new ArrayList<>();
dataNodeLocations.forEach(
(iotdb) 04/10: Hotfix concurrent connect error (#16900)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 75e8579e6008a5a10a921a59515dfef61a313e9a
Author: libo <ze...@163.com>
AuthorDate: Fri Dec 12 11:18:42 2025 +0800
Hotfix concurrent connect error (#16900)
(cherry picked from commit fad9bdedada706c521393bc4b926c81763d7bdf4)
---
.../java/org/apache/iotdb/db/protocol/session/IClientSession.java | 3 ++-
.../java/org/apache/iotdb/db/protocol/session/SessionManager.java | 4 ++--
2 files changed, 4 insertions(+), 3 deletions(-)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
index 97585673e82..16cd1dc2f75 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
@@ -20,6 +20,7 @@
package org.apache.iotdb.db.protocol.session;
import org.apache.iotdb.commons.conf.IoTDBConstant.ClientVersion;
+import org.apache.iotdb.commons.utils.CommonDateTimeUtils;
import org.apache.iotdb.db.queryengine.common.ConnectionInfo;
import org.apache.iotdb.service.rpc.thrift.TSConnectionInfo;
import org.apache.iotdb.service.rpc.thrift.TSConnectionType;
@@ -57,7 +58,7 @@ public abstract class IClientSession {
@Nullable private String databaseName;
- private long lastActiveTime;
+ private long lastActiveTime = CommonDateTimeUtils.currentTime();
public abstract String getClientAddress();
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
index e5e95d4f82c..db6955f6dc7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/SessionManager.java
@@ -497,8 +497,8 @@ public class SessionManager implements SessionManagerMBean {
session.setUsername(username);
session.setZoneId(zoneId);
session.setClientVersion(clientVersion);
- session.setLogin(true);
session.setLogInTime(System.currentTimeMillis());
+ session.setLogin(true);
}
public void closeDataset(
@@ -581,7 +581,7 @@ public class SessionManager implements SessionManagerMBean {
public List<ConnectionInfo> getAllSessionConnectionInfo() {
return sessions.keySet().stream()
- .filter(s -> StringUtils.isNotEmpty(s.getUsername()))
+ .filter(s -> StringUtils.isNotEmpty(s.getUsername()) && s.isLogin())
.map(IClientSession::convertToConnectionInfo)
.sorted(Comparator.comparingLong(ConnectionInfo::getLastActiveTime))
.collect(Collectors.toList());
(iotdb) 09/10: Add state predicate push down for system table current_queries (#16904)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit bd677fb99541d4d8eec3529323ec9239c0445f7c
Author: Weihao Li <60...@users.noreply.github.com>
AuthorDate: Mon Dec 15 18:05:43 2025 +0800
Add state predicate push down for system table current_queries (#16904)
(cherry picked from commit 2710cb3db653d8c674495647bd34cd5afdbc3544)
---
.../informationschema/IoTDBCurrentQueriesIT.java | 9 +-
.../InformationSchemaContentSupplierFactory.java | 30 ++++++-
.../iotdb/db/queryengine/plan/Coordinator.java | 33 +++++++
.../plan/planner/TableOperatorGenerator.java | 1 +
.../optimizations/PushPredicateIntoTableScan.java | 100 +++++++++++++++++++++
.../planner/assertions/ColumnReference.java | 10 +--
.../planner/assertions/PlanMatchPattern.java | 20 +++--
.../informationschema/CurrentQueriesTest.java | 77 ++++++++++++++++
8 files changed, 259 insertions(+), 21 deletions(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
index a68bd92b50b..4f9c8cab794 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
@@ -83,7 +83,7 @@ public class IoTDBCurrentQueriesIT {
statement.execute("set configuration \"query_cost_stat_window\"='1'");
// 1. query current_queries table
- String sql = "SELECT * FROM current_queries";
+ String sql = "SELECT * FROM current_queries WHERE state='RUNNING'";
ResultSet resultSet = statement.executeQuery(sql);
ResultSetMetaData metaData = resultSet.getMetaData();
Assert.assertEquals(CURRENT_QUERIES_COLUMN_NUM, metaData.getColumnCount());
@@ -116,7 +116,7 @@ public class IoTDBCurrentQueriesIT {
Assert.assertEquals(61, rowNum);
// 3. requery current_queries table
- sql = "SELECT * FROM current_queries";
+ sql = "SELECT * FROM current_queries WHERE state='FINISHED'";
resultSet = statement.executeQuery(sql);
metaData = resultSet.getMetaData();
Assert.assertEquals(CURRENT_QUERIES_COLUMN_NUM, metaData.getColumnCount());
@@ -128,13 +128,14 @@ public class IoTDBCurrentQueriesIT {
}
rowNum++;
}
- // three rows in the result, 2 FINISHED and 1 RUNNING
- Assert.assertEquals(3, rowNum);
+ // two rows in the result, 2 FINISHED
+ Assert.assertEquals(2, rowNum);
Assert.assertEquals(2, finishedQueries);
resultSet.close();
// 4. test the expired QueryInfo was evicted
Thread.sleep(61_001);
+ sql = "SELECT * FROM current_queries";
resultSet = statement.executeQuery(sql);
rowNum = 0;
while (resultSet.next()) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java
index 76d84c741de..fe0727ec3f2 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/relational/InformationSchemaContentSupplierFactory.java
@@ -69,11 +69,15 @@ import org.apache.iotdb.db.protocol.session.IClientSession;
import org.apache.iotdb.db.protocol.session.SessionManager;
import org.apache.iotdb.db.queryengine.common.ConnectionInfo;
import org.apache.iotdb.db.queryengine.common.QueryId;
+import org.apache.iotdb.db.queryengine.execution.QueryState;
import org.apache.iotdb.db.queryengine.plan.Coordinator;
import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
import org.apache.iotdb.db.queryengine.plan.execution.config.metadata.relational.ShowCreateViewTask;
import org.apache.iotdb.db.queryengine.plan.relational.function.TableBuiltinTableFunction;
import org.apache.iotdb.db.queryengine.plan.relational.security.AccessControl;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Expression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.StringLiteral;
import org.apache.iotdb.db.queryengine.plan.relational.sql.util.ReservedIdentifiers;
import org.apache.iotdb.db.relational.grammar.sql.RelationalSqlKeywords;
import org.apache.iotdb.db.schemaengine.table.InformationSchemaUtils;
@@ -130,7 +134,10 @@ public class InformationSchemaContentSupplierFactory {
private InformationSchemaContentSupplierFactory() {}
public static Iterator<TsBlock> getSupplier(
- final String tableName, final List<TSDataType> dataTypes, final UserEntity userEntity) {
+ final String tableName,
+ final List<TSDataType> dataTypes,
+ Expression predicate,
+ final UserEntity userEntity) {
try {
switch (tableName) {
case InformationSchema.QUERIES:
@@ -168,7 +175,7 @@ public class InformationSchemaContentSupplierFactory {
case InformationSchema.CONNECTIONS:
return new ConnectionsSupplier(dataTypes, userEntity);
case InformationSchema.CURRENT_QUERIES:
- return new CurrentQueriesSupplier(dataTypes, userEntity);
+ return new CurrentQueriesSupplier(dataTypes, predicate, userEntity);
case InformationSchema.QUERIES_COSTS_HISTOGRAM:
return new QueriesCostsHistogramSupplier(dataTypes, userEntity);
default:
@@ -1188,9 +1195,24 @@ public class InformationSchemaContentSupplierFactory {
private int nextConsumedIndex;
private List<Coordinator.StatedQueriesInfo> queriesInfo;
- private CurrentQueriesSupplier(final List<TSDataType> dataTypes, final UserEntity userEntity) {
+ private CurrentQueriesSupplier(
+ final List<TSDataType> dataTypes, Expression predicate, final UserEntity userEntity) {
super(dataTypes);
- queriesInfo = Coordinator.getInstance().getCurrentQueriesInfo();
+
+ if (predicate == null) {
+ queriesInfo = Coordinator.getInstance().getCurrentQueriesInfo();
+ } else if (QueryState.RUNNING
+ .toString()
+ .equals(((StringLiteral) ((ComparisonExpression) predicate).getRight()).getValue())) {
+ queriesInfo = Coordinator.getInstance().getRunningQueriesInfos();
+ } else if (QueryState.FINISHED
+ .toString()
+ .equals(((StringLiteral) ((ComparisonExpression) predicate).getRight()).getValue())) {
+ queriesInfo = Coordinator.getInstance().getFinishedQueriesInfos();
+ } else {
+ queriesInfo = Collections.emptyList();
+ }
+
try {
accessControl.checkUserGlobalSysPrivilege(userEntity);
} catch (final AccessDeniedException e) {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
index 78cdee720da..57e3d071074 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java
@@ -865,6 +865,39 @@ public class Coordinator {
}
}
+ public List<StatedQueriesInfo> getRunningQueriesInfos() {
+ long currentTime = System.currentTimeMillis();
+ return getAllQueryExecutions().stream()
+ .map(
+ queryExecution ->
+ new StatedQueriesInfo(
+ QueryState.RUNNING,
+ queryExecution.getQueryId(),
+ queryExecution.getStartExecutionTime(),
+ DEFAULT_END_TIME,
+ (currentTime - queryExecution.getStartExecutionTime()) / 1000,
+ queryExecution.getExecuteSQL().orElse("UNKNOWN"),
+ queryExecution.getUser(),
+ queryExecution.getClientHostname()))
+ .collect(Collectors.toList());
+ }
+
+ public List<StatedQueriesInfo> getFinishedQueriesInfos() {
+ long currentTime = System.currentTimeMillis();
+ List<StatedQueriesInfo> result = new ArrayList<>();
+ Iterator<QueryInfo> historyQueriesIterator = currentQueriesInfo.iterator();
+ long needRecordTime = currentTime - CONFIG.getQueryCostStatWindow() * 60 * 1_000L;
+ while (historyQueriesIterator.hasNext()) {
+ QueryInfo queryInfo = historyQueriesIterator.next();
+ if (queryInfo.endTime < needRecordTime) {
+ // out of time window, ignore it
+ } else {
+ result.add(new StatedQueriesInfo(QueryState.FINISHED, queryInfo));
+ }
+ }
+ return result;
+ }
+
public List<StatedQueriesInfo> getCurrentQueriesInfo() {
List<IQueryExecution> runningQueries = getAllQueryExecutions();
Set<String> runningQueryIdSet =
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TableOperatorGenerator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TableOperatorGenerator.java
index c29ceaff04f..68036f9d514 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TableOperatorGenerator.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/TableOperatorGenerator.java
@@ -1296,6 +1296,7 @@ public class TableOperatorGenerator extends PlanVisitor<Operator, LocalExecution
getSupplier(
node.getQualifiedObjectName().getObjectName(),
dataTypes,
+ node.getPushDownPredicate(),
context
.getDriverContext()
.getFragmentInstanceContext()
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PushPredicateIntoTableScan.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PushPredicateIntoTableScan.java
index fab09a2e3f5..55905f67682 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PushPredicateIntoTableScan.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/optimizations/PushPredicateIntoTableScan.java
@@ -56,6 +56,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationN
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AssignUniqueId;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.FilterNode;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.node.InformationSchemaTableScanNode;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.JoinNode;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.ProjectNode;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.SemiJoinNode;
@@ -68,6 +69,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.FunctionCall;
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LogicalExpression;
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Node;
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.NullLiteral;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.StringLiteral;
import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
import org.apache.iotdb.db.utils.TimestampPrecisionUtils;
@@ -99,6 +101,8 @@ import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static java.util.Objects.requireNonNull;
+import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.STATE_TABLE_MODEL;
+import static org.apache.iotdb.commons.schema.table.InformationSchema.CURRENT_QUERIES;
import static org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory.ATTRIBUTE;
import static org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory.FIELD;
import static org.apache.iotdb.commons.schema.table.column.TsTableColumnCategory.TIME;
@@ -440,6 +444,102 @@ public class PushPredicateIntoTableScan implements PlanOptimizer {
return output;
}
+ @Override
+ public PlanNode visitInformationSchemaTableScan(
+ InformationSchemaTableScanNode node, RewriteContext context) {
+ // no predicate or table is not current_queries, just return dierectly
+ if (TRUE_LITERAL.equals(context.inheritedPredicate)) {
+ return node;
+ }
+
+ // push-down for CURRENT_QUERIES
+ if (CURRENT_QUERIES.equals(node.getQualifiedObjectName().getObjectName())) {
+ SplitExpression splitExpression = splitCurrentQueriesPredicate(context.inheritedPredicate);
+ // exist expressions can push down to scan operator
+ if (!splitExpression.getExpressionsCanPushDown().isEmpty()) {
+ List<Expression> expressions = splitExpression.getExpressionsCanPushDown();
+ checkState(expressions.size() == 1, "Unexpected number of expressions in table scan");
+ node.setPushDownPredicate(expressions.get(0));
+ }
+
+ // exist expressions cannot push down
+ if (!splitExpression.getExpressionsCannotPushDown().isEmpty()) {
+ List<Expression> expressions = splitExpression.getExpressionsCannotPushDown();
+ return new FilterNode(
+ queryId.genPlanNodeId(),
+ node,
+ expressions.size() == 1
+ ? expressions.get(0)
+ : new LogicalExpression(LogicalExpression.Operator.AND, expressions));
+ }
+ return node;
+ }
+
+ FilterNode filterNode =
+ new FilterNode(queryId.genPlanNodeId(), node, context.inheritedPredicate);
+ context.inheritedPredicate = TRUE_LITERAL;
+ return filterNode;
+ }
+
+ private SplitExpression splitCurrentQueriesPredicate(Expression predicate) {
+ List<Expression> expressionsCanPushDown = new ArrayList<>();
+ List<Expression> expressionsCannotPushDown = new ArrayList<>();
+
+ if (predicate instanceof LogicalExpression
+ && ((LogicalExpression) predicate).getOperator() == LogicalExpression.Operator.AND) {
+
+ // predicate like state = 'xxx' can be push down
+ // Note: the optimizer CanonicalizeExpressionRewriter will ensure the predicate like 'xxx' =
+ // state will be canonicalized to state = 'xxx'
+ boolean hasExpressionPushDown = false;
+ for (Expression expression : ((LogicalExpression) predicate).getTerms()) {
+ if (isStateComparedWithConstant(expression) && !hasExpressionPushDown) {
+ // if there are more than one state = 'xxx' terms, only add first to push-down candidate
+ expressionsCanPushDown.add(expression);
+ hasExpressionPushDown = true;
+ } else {
+ expressionsCannotPushDown.add(expression);
+ }
+ }
+
+ return new SplitExpression(
+ Collections.emptyList(), expressionsCanPushDown, expressionsCannotPushDown, null);
+ }
+
+ if (isStateComparedWithConstant(predicate)) {
+ expressionsCanPushDown.add(predicate);
+ } else {
+ expressionsCannotPushDown.add(predicate);
+ }
+
+ return new SplitExpression(
+ Collections.emptyList(), expressionsCanPushDown, expressionsCannotPushDown, null);
+ }
+
+ private boolean isStateComparedWithConstant(Expression expression) {
+ if (!(expression instanceof ComparisonExpression)) {
+ return false;
+ }
+
+ ComparisonExpression comparisonExpression = (ComparisonExpression) expression;
+
+ if (ComparisonExpression.Operator.EQUAL != comparisonExpression.getOperator()) {
+ return false;
+ }
+
+ if (!(comparisonExpression.getLeft() instanceof SymbolReference)
+ || !STATE_TABLE_MODEL.equals(
+ ((SymbolReference) comparisonExpression.getLeft()).getName())) {
+ return false;
+ }
+
+ if (!(comparisonExpression.getRight() instanceof StringLiteral)) {
+ return false;
+ }
+
+ return true;
+ }
+
@Override
public PlanNode visitDeviceTableScan(
DeviceTableScanNode tableScanNode, RewriteContext context) {
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/ColumnReference.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/ColumnReference.java
index 6dc2e6d9f96..a93176c19a8 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/ColumnReference.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/ColumnReference.java
@@ -25,7 +25,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.metadata.ColumnSchema;
import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
import org.apache.iotdb.db.queryengine.plan.relational.planner.Symbol;
import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationTableScanNode;
-import org.apache.iotdb.db.queryengine.plan.relational.planner.node.DeviceTableScanNode;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.node.TableScanNode;
import java.util.Map;
import java.util.Optional;
@@ -49,10 +49,10 @@ public class ColumnReference implements RvalueMatcher {
String actualTableName;
Map<Symbol, ColumnSchema> assignments;
- if (node instanceof DeviceTableScanNode) {
- DeviceTableScanNode deviceTableScanNode = (DeviceTableScanNode) node;
- actualTableName = deviceTableScanNode.getQualifiedObjectName().toString();
- assignments = deviceTableScanNode.getAssignments();
+ if (node instanceof TableScanNode) {
+ TableScanNode tableScanNode = (TableScanNode) node;
+ actualTableName = tableScanNode.getQualifiedObjectName().toString();
+ assignments = tableScanNode.getAssignments();
}
/*else if (node instanceof IndexSourceNode indexSourceNode) {
tableHandle = indexSourceNode.getTableHandle();
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/PlanMatchPattern.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/PlanMatchPattern.java
index 9146f46eeca..3ec60536aec 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/PlanMatchPattern.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/assertions/PlanMatchPattern.java
@@ -149,14 +149,18 @@ public final class PlanMatchPattern {
public static PlanMatchPattern infoSchemaTableScan(
String expectedTableName, Optional<Integer> dataNodeId, List<String> outputSymbols) {
- return node(InformationSchemaTableScanNode.class)
- .with(
- new InformationSchemaTableScanMatcher(
- expectedTableName,
- Optional.empty(),
- outputSymbols,
- Collections.emptySet(),
- dataNodeId));
+ PlanMatchPattern pattern =
+ node(InformationSchemaTableScanNode.class)
+ .with(
+ new InformationSchemaTableScanMatcher(
+ expectedTableName,
+ Optional.empty(),
+ outputSymbols,
+ Collections.emptySet(),
+ dataNodeId));
+ outputSymbols.forEach(
+ symbol -> pattern.withAlias(symbol, new ColumnReference(expectedTableName, symbol)));
+ return pattern;
}
public static PlanMatchPattern treeDeviceViewTableScan(
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/CurrentQueriesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/CurrentQueriesTest.java
index 1e3163321ec..c2f4eb68399 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/CurrentQueriesTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/informationschema/CurrentQueriesTest.java
@@ -20,6 +20,9 @@ package org.apache.iotdb.db.queryengine.plan.relational.planner.informationschem
import org.apache.iotdb.db.queryengine.plan.planner.plan.LogicalQueryPlan;
import org.apache.iotdb.db.queryengine.plan.relational.planner.PlanTester;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.StringLiteral;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
import com.google.common.collect.ImmutableList;
import org.junit.Test;
@@ -40,8 +43,10 @@ import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.USER_T
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanAssert.assertPlan;
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.collect;
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.exchange;
+import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.filter;
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.infoSchemaTableScan;
import static org.apache.iotdb.db.queryengine.plan.relational.planner.assertions.PlanMatchPattern.output;
+import static org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression.Operator.EQUAL;
public class CurrentQueriesTest {
private final PlanTester planTester = new PlanTester();
@@ -80,6 +85,78 @@ public class CurrentQueriesTest {
infoSchemaTableScan("information_schema.current_queries", Optional.of(2)));
}
+ @Test
+ public void testCurrentQueriesFilterPushDown() {
+ // Normal case
+ LogicalQueryPlan logicalQueryPlan =
+ planTester.createPlan(
+ "select * from information_schema.current_queries where state='RUNNING'");
+ assertPlan(
+ logicalQueryPlan,
+ output(
+ infoSchemaTableScan(
+ "information_schema.current_queries",
+ Optional.empty(),
+ ImmutableList.of(
+ QUERY_ID_TABLE_MODEL,
+ STATE_TABLE_MODEL,
+ START_TIME_TABLE_MODEL,
+ END_TIME_TABLE_MODEL,
+ DATA_NODE_ID_TABLE_MODEL,
+ COST_TIME,
+ STATEMENT_TABLE_MODEL,
+ USER_TABLE_MODEL,
+ CLIENT_IP))));
+
+ // mixed push down and cannot push down term
+ logicalQueryPlan =
+ planTester.createPlan(
+ "select * from information_schema.current_queries where state='RUNNING' and query_id='1'");
+ assertPlan(
+ logicalQueryPlan,
+ output(
+ filter(
+ new ComparisonExpression(
+ EQUAL, new SymbolReference(QUERY_ID_TABLE_MODEL), new StringLiteral("1")),
+ infoSchemaTableScan(
+ "information_schema.current_queries",
+ Optional.empty(),
+ ImmutableList.of(
+ QUERY_ID_TABLE_MODEL,
+ STATE_TABLE_MODEL,
+ START_TIME_TABLE_MODEL,
+ END_TIME_TABLE_MODEL,
+ DATA_NODE_ID_TABLE_MODEL,
+ COST_TIME,
+ STATEMENT_TABLE_MODEL,
+ USER_TABLE_MODEL,
+ CLIENT_IP)))));
+
+ // More than one state='xxx' terms
+ logicalQueryPlan =
+ planTester.createPlan(
+ "select * from information_schema.current_queries where state='RUNNING' and state='xx'");
+ assertPlan(
+ logicalQueryPlan,
+ output(
+ filter(
+ new ComparisonExpression(
+ EQUAL, new SymbolReference(STATE_TABLE_MODEL), new StringLiteral("xx")),
+ infoSchemaTableScan(
+ "information_schema.current_queries",
+ Optional.empty(),
+ ImmutableList.of(
+ QUERY_ID_TABLE_MODEL,
+ STATE_TABLE_MODEL,
+ START_TIME_TABLE_MODEL,
+ END_TIME_TABLE_MODEL,
+ DATA_NODE_ID_TABLE_MODEL,
+ COST_TIME,
+ STATEMENT_TABLE_MODEL,
+ USER_TABLE_MODEL,
+ CLIENT_IP)))));
+ }
+
@Test
public void testQueriesCostsHistogram() {
LogicalQueryPlan logicalQueryPlan =
(iotdb) 03/10: Fix the config 'query_cost_stat_window' cannot be hot loaded #16896
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 2d3c19a565651c2f6bbed39abf512fa27b01ab95
Author: Weihao Li <60...@users.noreply.github.com>
AuthorDate: Thu Dec 11 19:45:26 2025 +0800
Fix the config 'query_cost_stat_window' cannot be hot loaded #16896
(cherry picked from commit 9abac5c2d76da6f1ade9ca3e958d194fd6a2acfc)
---
.../it/query/recent/informationschema/IoTDBCurrentQueriesIT.java | 2 +-
.../src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java | 6 ++++++
.../iotdb/db/queryengine/plan/execution/config/ConfigExecution.java | 1 +
3 files changed, 8 insertions(+), 1 deletion(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
index 66941ec784f..a68bd92b50b 100644
--- a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/informationschema/IoTDBCurrentQueriesIT.java
@@ -64,7 +64,6 @@ public class IoTDBCurrentQueriesIT {
@BeforeClass
public static void setUp() throws Exception {
- EnvFactory.getEnv().getConfig().getDataNodeConfig().setQueryCostStatWindow(1);
EnvFactory.getEnv().initClusterEnvironment();
createUser("test", "test123123456");
}
@@ -81,6 +80,7 @@ public class IoTDBCurrentQueriesIT {
EnvFactory.getEnv().getConnection(ADMIN_NAME, ADMIN_PWD, BaseEnv.TABLE_SQL_DIALECT);
Statement statement = connection.createStatement();
statement.execute("USE information_schema");
+ statement.execute("set configuration \"query_cost_stat_window\"='1'");
// 1. query current_queries table
String sql = "SELECT * FROM current_queries";
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
index 6a1e8dbcccb..a039fc61774 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java
@@ -2041,6 +2041,12 @@ public class IoTDBDescriptor {
properties.getProperty(
"slow_query_threshold",
ConfigurationFileUtils.getConfigurationDefaultValue("slow_query_threshold"))));
+ // update query_cost_stat_window
+ conf.setQueryCostStatWindow(
+ Integer.parseInt(
+ properties.getProperty(
+ "query_cost_stat_window",
+ ConfigurationFileUtils.getConfigurationDefaultValue("query_cost_stat_window"))));
// update select into operation max buffer size
conf.setIntoOperationBufferSizeInByte(
Long.parseLong(
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
index 823a620820f..7868460f387 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/ConfigExecution.java
@@ -354,6 +354,7 @@ public class ConfigExecution implements IQueryExecution {
return context.getSession().getUserName();
}
+ @Override
public String getClientHostname() {
return context.getCliHostname();
}
(iotdb) 06/10: Delete audit implementation (#16899)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit a5c3b4387bb057ece36b27e89a59b4c3c3d64818
Author: jintao zhu <10...@users.noreply.github.com>
AuthorDate: Fri Dec 12 16:44:18 2025 +0800
Delete audit implementation (#16899)
(cherry picked from commit 48bd352fb00a0785156ae381c199a38468dd5e15)
---
.../iotdb/db/it/audit/IoTDBAuditLogBasicIT.java | 2 +
.../iotdb/confignode/audit/CNAuditLogger.java | 50 +--
.../org/apache/iotdb/db/audit/DNAuditLogger.java | 398 +--------------------
.../iotdb/commons/audit/AbstractAuditLogger.java | 77 +---
4 files changed, 8 insertions(+), 519 deletions(-)
diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/audit/IoTDBAuditLogBasicIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/audit/IoTDBAuditLogBasicIT.java
index f12dd4ca66e..83aa06dc61d 100644
--- a/integration-test/src/test/java/org/apache/iotdb/db/it/audit/IoTDBAuditLogBasicIT.java
+++ b/integration-test/src/test/java/org/apache/iotdb/db/it/audit/IoTDBAuditLogBasicIT.java
@@ -32,6 +32,7 @@ import org.apache.iotdb.itbase.runtime.NodeConnection;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
+import org.junit.Ignore;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;
@@ -52,6 +53,7 @@ import java.util.stream.Stream;
* This test class ensures the audit log behave exactly the same as we expected, including the
* number, sequence and content of the audit logs.
*/
+@Ignore
@RunWith(IoTDBTestRunner.class)
@Category({LocalStandaloneIT.class})
public class IoTDBAuditLogBasicIT {
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/audit/CNAuditLogger.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/audit/CNAuditLogger.java
index 9b14d7b8423..ccc1008eec3 100644
--- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/audit/CNAuditLogger.java
+++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/audit/CNAuditLogger.java
@@ -19,24 +19,16 @@
package org.apache.iotdb.confignode.audit;
-import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
-import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
-import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
-import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
import org.apache.iotdb.commons.audit.AbstractAuditLogger;
import org.apache.iotdb.commons.audit.IAuditEntity;
-import org.apache.iotdb.confignode.client.async.AsyncDataNodeHeartbeatClientPool;
-import org.apache.iotdb.confignode.client.async.handlers.audit.DataNodeWriteAuditLogHandler;
import org.apache.iotdb.confignode.conf.ConfigNodeConfig;
import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
import org.apache.iotdb.confignode.manager.ConfigManager;
import org.apache.iotdb.confignode.manager.IManager;
-import org.apache.iotdb.mpp.rpc.thrift.TAuditLogReq;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.util.List;
import java.util.function.Supplier;
public class CNAuditLogger extends AbstractAuditLogger {
@@ -50,45 +42,5 @@ public class CNAuditLogger extends AbstractAuditLogger {
}
@Override
- public void log(IAuditEntity auditLogFields, Supplier<String> log) {
- if (!IS_AUDIT_LOG_ENABLED) {
- return;
- }
- if (noNeedInsertAuditLog(auditLogFields)) {
- return;
- }
- // find database "__audit"'s data_region
- List<TRegionReplicaSet> auditReplicaSets =
- configManager
- .getPartitionManager()
- .getAllReplicaSets("root.__audit", TConsensusGroupType.DataRegion);
- if (auditReplicaSets.isEmpty()) {
- logger.warn("Database {} does not exist.", "root.__audit");
- return;
- }
- TConsensusGroupId regionId = auditReplicaSets.get(0).getRegionId();
- // use ConfigManager.getLoadManager().getLoadCache().getRegionLeaderMap() to get regionLeaderId
- TDataNodeLocation regionLeader = configManager.getRegionLeaderLocation(regionId);
- TAuditLogReq req =
- new TAuditLogReq(
- auditLogFields.getUsername(),
- auditLogFields.getUserId(),
- auditLogFields.getCliHostname(),
- auditLogFields.getAuditEventType().toString(),
- auditLogFields.getAuditLogOperation().toString(),
- auditLogFields.getPrivilegeTypeString(),
- auditLogFields.getResult(),
- auditLogFields.getDatabase(),
- auditLogFields.getSqlString(),
- log.get(),
- CONF.getConfigNodeId());
- // refer the implementation of HeartbeatService.pingRegisteredDataNode(). By appending a new
- // writeAudtiLog() interface in AsyncDataNodeHeartbeatClientPool, the main thread is not
- // required to wait until the write audit log request to be complete.
- AsyncDataNodeHeartbeatClientPool.getInstance()
- .writeAuditLog(
- regionLeader.getInternalEndPoint(),
- req,
- new DataNodeWriteAuditLogHandler(regionLeader.getDataNodeId()));
- }
+ public void log(IAuditEntity auditLogFields, Supplier<String> log) {}
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/audit/DNAuditLogger.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/audit/DNAuditLogger.java
index 713d79f77d9..3a07ac9aa69 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/audit/DNAuditLogger.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/audit/DNAuditLogger.java
@@ -19,99 +19,20 @@
package org.apache.iotdb.db.audit;
-import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.audit.AbstractAuditLogger;
-import org.apache.iotdb.commons.audit.AuditEventType;
import org.apache.iotdb.commons.audit.AuditLogFields;
-import org.apache.iotdb.commons.audit.AuditLogOperation;
import org.apache.iotdb.commons.audit.IAuditEntity;
-import org.apache.iotdb.commons.audit.PrivilegeLevel;
-import org.apache.iotdb.commons.audit.UserEntity;
-import org.apache.iotdb.commons.auth.entity.PrivilegeType;
-import org.apache.iotdb.commons.client.IClientManager;
-import org.apache.iotdb.commons.client.exception.ClientManagerException;
-import org.apache.iotdb.commons.conf.IoTDBConstant;
-import org.apache.iotdb.commons.consensus.ConfigRegionId;
import org.apache.iotdb.commons.exception.IllegalPathException;
import org.apache.iotdb.commons.path.PartialPath;
-import org.apache.iotdb.commons.pipe.config.constant.SystemConstant;
-import org.apache.iotdb.commons.utils.CommonDateTimeUtils;
-import org.apache.iotdb.confignode.rpc.thrift.TGetDatabaseReq;
-import org.apache.iotdb.confignode.rpc.thrift.TShowDatabaseResp;
-import org.apache.iotdb.db.auth.AuthorityChecker;
-import org.apache.iotdb.db.conf.IoTDBConfig;
-import org.apache.iotdb.db.conf.IoTDBDescriptor;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClient;
-import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager;
-import org.apache.iotdb.db.protocol.client.ConfigNodeInfo;
-import org.apache.iotdb.db.protocol.session.IClientSession;
-import org.apache.iotdb.db.protocol.session.InternalClientSession;
-import org.apache.iotdb.db.protocol.session.SessionManager;
-import org.apache.iotdb.db.queryengine.common.SessionInfo;
import org.apache.iotdb.db.queryengine.plan.Coordinator;
-import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher;
-import org.apache.iotdb.db.queryengine.plan.analyze.cache.schema.DataNodeDevicePathCache;
-import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
-import org.apache.iotdb.db.queryengine.plan.parser.StatementGenerator;
-import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
-import org.apache.iotdb.db.queryengine.plan.relational.metadata.Metadata;
-import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
-import org.apache.iotdb.db.queryengine.plan.statement.Statement;
import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement;
-import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowDatabaseStatement;
-import org.apache.iotdb.rpc.TSStatusCode;
-
-import org.apache.thrift.TException;
-import org.apache.tsfile.common.conf.TSFileConfig;
-import org.apache.tsfile.enums.TSDataType;
-import org.apache.tsfile.utils.Binary;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
import javax.validation.constraints.NotNull;
-import java.io.IOException;
-import java.time.ZoneId;
-import java.util.Arrays;
-import java.util.List;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-import static org.apache.iotdb.db.pipe.receiver.protocol.legacy.loader.ILoader.SCHEMA_FETCHER;
public class DNAuditLogger extends AbstractAuditLogger {
public static final String PREFIX_PASSWORD_HISTORY = "root.__audit.password_history";
- private static final Logger logger = LoggerFactory.getLogger(DNAuditLogger.class);
-
- // TODO: @zhujt20 Optimize the following stupid intervals
- private static final int INSERT_RETRY_COUNT = 5;
- private static final int INSERT_RETRY_INTERVAL_MS = 2000;
- private static final int INSERT_INTERVAL_MS = 50;
-
- private static final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
-
- private static final String AUDIT_LOG_DEVICE = "root.__audit.log.node_%s.u_%s";
- private static final String AUDIT_CN_LOG_DEVICE = "root.__audit.log.node_%s.u_all";
- private static final SessionInfo sessionInfo =
- new SessionInfo(
- 0,
- new UserEntity(
- AuthorityChecker.INTERNAL_AUDIT_USER_ID,
- AuthorityChecker.INTERNAL_AUDIT_USER,
- IoTDBDescriptor.getInstance().getConfig().getInternalAddress()),
- ZoneId.systemDefault());
-
- private static final SessionManager SESSION_MANAGER = SessionManager.getInstance();
-
- private static final IClientManager<ConfigRegionId, ConfigNodeClient> CONFIG_NODE_CLIENT_MANAGER =
- ConfigNodeClientManager.getInstance();
-
- private static final DataNodeDevicePathCache DEVICE_PATH_CACHE =
- DataNodeDevicePathCache.getInstance();
- private static final AtomicBoolean tableViewIsInitialized = new AtomicBoolean(false);
private Coordinator coordinator;
@@ -130,327 +51,16 @@ public class DNAuditLogger extends AbstractAuditLogger {
@NotNull
private static InsertRowStatement generateInsertStatement(
IAuditEntity auditLogFields, String log, PartialPath logDevice) {
- String username = auditLogFields.getUsername();
- String address = auditLogFields.getCliHostname();
- AuditEventType type = auditLogFields.getAuditEventType();
- AuditLogOperation operation = auditLogFields.getAuditLogOperation();
- PrivilegeLevel privilegeLevel = null;
- if (auditLogFields.getPrivilegeTypes() != null) {
- for (PrivilegeType privilegeType : auditLogFields.getPrivilegeTypes()) {
- privilegeLevel = judgePrivilegeLevel(privilegeType);
- if (privilegeLevel.equals(PrivilegeLevel.GLOBAL)) {
- break;
- }
- }
- } else {
- privilegeLevel = PrivilegeLevel.GLOBAL;
- }
- InsertRowStatement insertStatement = new InsertRowStatement();
- insertStatement.setDevicePath(logDevice);
- insertStatement.setTime(CommonDateTimeUtils.currentTime());
- insertStatement.setMeasurements(
- new String[] {
- AUDIT_LOG_USERNAME,
- AUDIT_LOG_CLI_HOSTNAME,
- AUDIT_LOG_AUDIT_EVENT_TYPE,
- AUDIT_LOG_OPERATION_TYPE,
- AUDIT_LOG_PRIVILEGE_TYPE,
- AUDIT_LOG_PRIVILEGE_LEVEL,
- AUDIT_LOG_RESULT,
- AUDIT_LOG_DATABASE,
- AUDIT_LOG_SQL_STRING,
- AUDIT_LOG_LOG
- });
- insertStatement.setAligned(false);
- String sqlString = auditLogFields.getSqlString();
- if (sqlString != null) {
- if (sqlString.toUpperCase().startsWith("CREATE USER")) {
- sqlString = String.join(" ", Arrays.asList(sqlString.split(" ")).subList(0, 3)) + " ...";
- }
- Pattern pattern = Pattern.compile("(?i)(values)\\([^)]*\\)");
- Matcher matcher = pattern.matcher(sqlString);
- StringBuffer sb = new StringBuffer();
- while (matcher.find()) {
- matcher.appendReplacement(sb, matcher.group(1) + "(...)");
- }
- matcher.appendTail(sb);
- sqlString = sb.toString();
- }
- insertStatement.setValues(
- new Object[] {
- new Binary(username == null ? "null" : username, TSFileConfig.STRING_CHARSET),
- new Binary(address == null ? "null" : address, TSFileConfig.STRING_CHARSET),
- new Binary(type == null ? "null" : type.toString(), TSFileConfig.STRING_CHARSET),
- new Binary(
- operation == null ? "null" : operation.toString(), TSFileConfig.STRING_CHARSET),
- new Binary(
- auditLogFields.getPrivilegeTypes() == null
- ? "null"
- : auditLogFields.getPrivilegeTypeString(),
- TSFileConfig.STRING_CHARSET),
- new Binary(
- privilegeLevel == null ? "null" : privilegeLevel.toString(),
- TSFileConfig.STRING_CHARSET),
- auditLogFields.getResult(),
- new Binary(
- auditLogFields.getDatabase() == null ? "null" : auditLogFields.getDatabase(),
- TSFileConfig.STRING_CHARSET),
- new Binary(sqlString == null ? "null" : sqlString, TSFileConfig.STRING_CHARSET),
- new Binary(log == null ? "null" : log, TSFileConfig.STRING_CHARSET)
- });
- insertStatement.setDataTypes(
- new TSDataType[] {
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.BOOLEAN,
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.STRING,
- });
- return insertStatement;
+ return null;
}
- public void createViewIfNecessary() {
- if (!tableViewIsInitialized.get()) {
- synchronized (this) {
- if (tableViewIsInitialized.get()) {
- return;
- }
- Statement statement =
- StatementGenerator.createStatement(
- "SHOW DATABASES " + SystemConstant.AUDIT_DATABASE, ZoneId.systemDefault());
- try (final ConfigNodeClient client =
- CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
- ShowDatabaseStatement showStatement = (ShowDatabaseStatement) statement;
- final List<String> databasePathPattern =
- Arrays.asList(showStatement.getPathPattern().getNodes());
- final TGetDatabaseReq req =
- new TGetDatabaseReq(
- databasePathPattern, showStatement.getAuthorityScope().serialize())
- .setIsTableModel(false);
- final TShowDatabaseResp resp = client.showDatabase(req);
- if (resp.getDatabaseInfoMapSize() > 0) {
- tableViewIsInitialized.set(true);
- return;
- }
- } catch (ClientManagerException | TException | IOException e) {
- logger.error(
- "Failed to show database before creating database {} for audit log",
- SystemConstant.AUDIT_DATABASE);
- }
-
- statement =
- StatementGenerator.createStatement(
- "CREATE DATABASE "
- + SystemConstant.AUDIT_DATABASE
- + " WITH SCHEMA_REGION_GROUP_NUM=1, DATA_REGION_GROUP_NUM=1",
- ZoneId.systemDefault());
- ExecutionResult result =
- coordinator.executeForTreeModel(
- statement,
- SESSION_MANAGER.requestQueryId(),
- sessionInfo,
- "",
- ClusterPartitionFetcher.getInstance(),
- SCHEMA_FETCHER);
- if (result.status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
- || result.status.getCode() == TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
-
- SqlParser relationSqlParser = new SqlParser();
- IClientSession session =
- new InternalClientSession(
- String.format(
- "%s_%s", DNAuditLogger.class.getSimpleName(), SystemConstant.AUDIT_DATABASE));
- session.setUsername(AuthorityChecker.INTERNAL_AUDIT_USER);
- session.setZoneId(ZoneId.systemDefault());
- session.setClientVersion(IoTDBConstant.ClientVersion.V_1_0);
- session.setDatabaseName(SystemConstant.AUDIT_DATABASE);
- session.setSqlDialect(IClientSession.SqlDialect.TABLE);
- SESSION_MANAGER.registerSession(session);
- Metadata metadata = LocalExecutionPlanner.getInstance().metadata;
-
- org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement stmt =
- relationSqlParser.createStatement(
- "CREATE DATABASE " + SystemConstant.AUDIT_PREFIX_KEY,
- ZoneId.systemDefault(),
- session);
- TSStatus status =
- coordinator.executeForTableModel(
- stmt,
- relationSqlParser,
- session,
- SESSION_MANAGER.requestQueryId(),
- SESSION_MANAGER.getSessionInfoOfTableModel(session),
- "",
- metadata,
- config.getQueryTimeoutThreshold(),
- false)
- .status;
- if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
- && status.getCode() != TSStatusCode.DATABASE_ALREADY_EXISTS.getStatusCode()) {
- logger.error(
- "Failed to create database in table model for audit log, because {}",
- status.getMessage());
- }
- stmt =
- relationSqlParser.createStatement(
- String.format(
- "CREATE VIEW __audit.audit_log (\n"
- + " %s STRING TAG,\n"
- + " %s STRING TAG,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s BOOLEAN FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD,\n"
- + " %s STRING FIELD\n"
- + ") AS root.__audit.log.**",
- AUDIT_LOG_NODE_ID,
- AUDIT_LOG_USER_ID,
- AUDIT_LOG_USERNAME,
- AUDIT_LOG_CLI_HOSTNAME,
- AUDIT_LOG_AUDIT_EVENT_TYPE,
- AUDIT_LOG_OPERATION_TYPE,
- AUDIT_LOG_PRIVILEGE_TYPE,
- AUDIT_LOG_PRIVILEGE_LEVEL,
- AUDIT_LOG_RESULT,
- AUDIT_LOG_DATABASE,
- AUDIT_LOG_SQL_STRING,
- AUDIT_LOG_LOG),
- ZoneId.systemDefault(),
- session);
- status =
- coordinator.executeForTableModel(
- stmt,
- relationSqlParser,
- session,
- SESSION_MANAGER.requestQueryId(),
- SESSION_MANAGER.getSessionInfoOfTableModel(session),
- "",
- metadata,
- config.getQueryTimeoutThreshold(),
- false)
- .status;
- if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()
- && status.getCode()
- != TSStatusCode.MEASUREMENT_ALREADY_EXISTS_IN_TEMPLATE.getStatusCode()) {
- logger.error("Failed to create view for audit log, because {}", status.getMessage());
- } else {
- logger.info("Create view for audit log successfully");
- tableViewIsInitialized.set(true);
- }
- } else {
- logger.error("Failed to create database {} for audit log", SystemConstant.AUDIT_DATABASE);
- }
- }
- }
- }
+ public void createViewIfNecessary() {}
@Override
- public synchronized void log(IAuditEntity auditLogFields, Supplier<String> log) {
- if (!IS_AUDIT_LOG_ENABLED) {
- return;
- }
- try {
- createViewIfNecessary();
- if (noNeedInsertAuditLog(auditLogFields)) {
- return;
- }
- long userId = auditLogFields.getUserId();
- String user = String.valueOf(userId);
- if (userId == -1) {
- user = "none";
- }
- String dataNodeId = String.valueOf(config.getDataNodeId());
- InsertRowStatement statement =
- generateInsertStatement(
- auditLogFields,
- log.get(),
- DEVICE_PATH_CACHE.getPartialPath(String.format(AUDIT_LOG_DEVICE, dataNodeId, user)));
- for (int retry = 0; retry < INSERT_RETRY_COUNT; retry++) {
- ExecutionResult insertResult =
- coordinator.executeForTreeModel(
- statement,
- SESSION_MANAGER.requestQueryId(),
- sessionInfo,
- "",
- ClusterPartitionFetcher.getInstance(),
- SCHEMA_FETCHER);
- if (insertResult.status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- TimeUnit.MILLISECONDS.sleep(INSERT_INTERVAL_MS);
- return;
- }
- TimeUnit.MILLISECONDS.sleep(INSERT_RETRY_INTERVAL_MS);
- }
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- logger.warn("[AUDIT] Audit log insertion retry sleep was interrupted because", e);
- } catch (Exception e) {
- logger.warn("[AUDIT] Failed to log audit events because", e);
- }
- AuditEventType type = auditLogFields.getAuditEventType();
- if (isLoginEvent(type)) {
- // TODO: @wenyanshi-123 Reactivate the following codes in the future
- // try {
- // statement.setDevicePath(
- // DEVICE_PATH_CACHE.getPartialPath(
- // String.format(AUDIT_LOGIN_LOG_DEVICE, dataNodeId, user)));
- // } catch (IllegalPathException e) {
- // logger.error("Failed to log audit login events because ", e);
- // return;
- // }
- // coordinator.executeForTreeModel(
- // statement,
- // SESSION_MANAGER.requestQueryId(),
- // sessionInfo,
- // "",
- // ClusterPartitionFetcher.getInstance(),
- // SCHEMA_FETCHER);
- }
- }
+ public synchronized void log(IAuditEntity auditLogFields, Supplier<String> log) {}
public void logFromCN(AuditLogFields auditLogFields, String log, int nodeId)
- throws IllegalPathException {
- if (!IS_AUDIT_LOG_ENABLED) {
- return;
- }
- createViewIfNecessary();
- if (noNeedInsertAuditLog(auditLogFields)) {
- return;
- }
- InsertRowStatement statement =
- generateInsertStatement(
- auditLogFields,
- log,
- DEVICE_PATH_CACHE.getPartialPath(String.format(AUDIT_CN_LOG_DEVICE, nodeId)));
- for (int retry = 0; retry < INSERT_RETRY_COUNT; retry++) {
- ExecutionResult insertResult =
- coordinator.executeForTreeModel(
- statement,
- SESSION_MANAGER.requestQueryId(),
- sessionInfo,
- "",
- ClusterPartitionFetcher.getInstance(),
- SCHEMA_FETCHER);
- if (insertResult.status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
- return;
- }
- try {
- TimeUnit.MILLISECONDS.sleep(INSERT_RETRY_INTERVAL_MS);
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- logger.error("Audit log insertion retry sleep was interrupted", e);
- }
- }
- }
+ throws IllegalPathException {}
private static class DNAuditLoggerHolder {
diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/AbstractAuditLogger.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/AbstractAuditLogger.java
index ca659a532cd..cec03366321 100644
--- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/AbstractAuditLogger.java
+++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/AbstractAuditLogger.java
@@ -19,11 +19,9 @@
package org.apache.iotdb.commons.audit;
-import org.apache.iotdb.commons.auth.entity.PrivilegeType;
import org.apache.iotdb.commons.conf.CommonConfig;
import org.apache.iotdb.commons.conf.CommonDescriptor;
-import java.util.List;
import java.util.function.Supplier;
public abstract class AbstractAuditLogger {
@@ -43,83 +41,10 @@ public abstract class AbstractAuditLogger {
private static final CommonConfig CONFIG = CommonDescriptor.getInstance().getConfig();
protected static final boolean IS_AUDIT_LOG_ENABLED = CONFIG.isEnableAuditLog();
- private static final List<AuditLogOperation> AUDITABLE_OPERATION_TYPE =
- CONFIG.getAuditableOperationType();
- private static final PrivilegeLevel AUDITABLE_OPERATION_LEVEL =
- CONFIG.getAuditableOperationLevel();
- private static final String AUDITABLE_OPERATION_RESULT = CONFIG.getAuditableOperationResult();
public abstract void log(IAuditEntity auditLogFields, Supplier<String> log);
public boolean noNeedInsertAuditLog(IAuditEntity auditLogFields) {
- AuditLogOperation operation = auditLogFields.getAuditLogOperation();
- boolean result = auditLogFields.getResult();
-
- // to do: check whether this event should be logged.
- // if whitelist or blacklist is used, only ip on the whitelist or blacklist can be logged
-
- if (AUDITABLE_OPERATION_TYPE == null || !AUDITABLE_OPERATION_TYPE.contains(operation)) {
- return true;
- }
- if (auditLogFields.getPrivilegeTypes() != null) {
- for (PrivilegeType privilegeType : auditLogFields.getPrivilegeTypes()) {
- PrivilegeLevel privilegeLevel = judgePrivilegeLevel(privilegeType);
- if (AUDITABLE_OPERATION_LEVEL == PrivilegeLevel.OBJECT
- && privilegeLevel == PrivilegeLevel.GLOBAL) {
- return true;
- }
- }
- }
- if (result && !AUDITABLE_OPERATION_RESULT.contains("SUCCESS")) {
- return true;
- }
- return !result && !AUDITABLE_OPERATION_RESULT.contains("FAIL");
- }
-
- public static PrivilegeLevel judgePrivilegeLevel(PrivilegeType type) {
- if (type == null) {
- return PrivilegeLevel.GLOBAL;
- }
- switch (type) {
- case READ_DATA:
- case DROP:
- case ALTER:
- case CREATE:
- case DELETE:
- case INSERT:
- case SELECT:
- case MANAGE_DATABASE:
- case WRITE_DATA:
- case READ_SCHEMA:
- case WRITE_SCHEMA:
- return PrivilegeLevel.OBJECT;
- case USE_CQ:
- case USE_UDF:
- case USE_PIPE:
- case USE_MODEL:
- case MAINTAIN:
- case MANAGE_ROLE:
- case MANAGE_USER:
- case USE_TRIGGER:
- case EXTEND_TEMPLATE:
- default:
- return PrivilegeLevel.GLOBAL;
- }
- }
-
- public static Boolean isLoginEvent(AuditEventType type) {
- switch (type) {
- case LOGIN:
- case LOGIN_FINAL:
- case MODIFY_PASSWD:
- case LOGIN_EXCEED_LIMIT:
- case LOGIN_FAILED_TRIES:
- case LOGIN_REJECT_IP:
- case LOGIN_FAIL_MAX_TIMES:
- case LOGIN_RESOURCE_RESTRICT:
- return true;
- default:
- return false;
- }
+ return true;
}
}
(iotdb) 07/10: Avoid CN OOM by Pulling User/Roles to DN When Cache Misses. (#16888)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jackietien pushed a commit to branch force_ci/object_type
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 900ff029b1f3b6ea95cf62cc358c8aa71815b1e9
Author: wenyanshi-123 <34...@qq.com>
AuthorDate: Mon Dec 15 09:47:17 2025 +0800
Avoid CN OOM by Pulling User/Roles to DN When Cache Misses. (#16888)
* Fix OOM problem for PIPE.
* Fix IT.
* Fix potential NPE.
---------
Co-authored-by: shiwenyan <we...@timecho.com>
(cherry picked from commit 44409934426cf9d786669b98996a9cc1edf5c4ab)
---
.../iotdb/db/auth/ClusterAuthorityFetcher.java | 59 +++++++++-------------
1 file changed, 25 insertions(+), 34 deletions(-)
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/ClusterAuthorityFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/ClusterAuthorityFetcher.java
index b7b3815caca..2894ce8bb47 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/ClusterAuthorityFetcher.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/auth/ClusterAuthorityFetcher.java
@@ -167,37 +167,37 @@ public class ClusterAuthorityFetcher implements IAuthorityFetcher {
@Override
public List<Integer> checkUserPathPrivileges(
String username, List<? extends PartialPath> allPath, PrivilegeType permission) {
- checkCacheAvailable();
List<Integer> posList = new ArrayList<>();
- User user = iAuthorCache.getUserCache(username);
- if (user != null) {
- if (user.isOpenIdUser()) {
- return posList;
- }
- int pos = 0;
- for (PartialPath path : allPath) {
- if (!user.checkPathPrivilege(path, permission)) {
- boolean checkFromRole = false;
- for (String rolename : user.getRoleSet()) {
- Role cachedRole = iAuthorCache.getRoleCache(rolename);
- if (cachedRole == null) {
- return checkPathFromConfigNode(username, allPath, permission);
- }
- if (cachedRole.checkPathPrivilege(path, permission)) {
- checkFromRole = true;
- break;
- }
+ if (username.equals(AuthorityChecker.INTERNAL_AUDIT_USER)) {
+ return posList;
+ }
+ checkCacheAvailable();
+ User user = getUser(username);
+ if (user.isOpenIdUser()) {
+ return posList;
+ }
+ int pos = 0;
+ for (PartialPath path : allPath) {
+ if (!user.checkPathPrivilege(path, permission)) {
+ boolean checkFromRole = false;
+ for (String rolename : user.getRoleSet()) {
+ Role cachedRole = iAuthorCache.getRoleCache(rolename);
+ if (cachedRole == null) {
+ checkRoleFromConfigNode(username, rolename);
+ cachedRole = iAuthorCache.getRoleCache(rolename);
}
- if (!checkFromRole) {
- posList.add(pos);
+ if (cachedRole != null && cachedRole.checkPathPrivilege(path, permission)) {
+ checkFromRole = true;
+ break;
}
}
- pos++;
+ if (!checkFromRole) {
+ posList.add(pos);
+ }
}
- return posList;
- } else {
- return checkPathFromConfigNode(username, allPath, permission);
+ pos++;
}
+ return posList;
}
@Override
@@ -642,15 +642,6 @@ public class ClusterAuthorityFetcher implements IAuthorityFetcher {
return permissionInfoResp;
}
- private List<Integer> checkPathFromConfigNode(
- String username, List<? extends PartialPath> allPath, PrivilegeType permission) {
- TCheckUserPrivilegesReq req =
- new TCheckUserPrivilegesReq(
- username, PrivilegeModelType.TREE.ordinal(), permission.ordinal(), false);
- req.setPaths(AuthUtils.serializePartialPathList(allPath));
- return checkPrivilegeFromConfigNode(req).getFailPos();
- }
-
private boolean checkRoleFromConfigNode(String username, String rolename) {
TAuthorizerReq req = new TAuthorizerReq();
// just reuse authorizer request. only need username and rolename field.