This is an automated email from the ASF dual-hosted git repository. liyuheng pushed a commit to branch lyh/show-models-arch in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit a9e4ba92cc9f99ceef5a2184fec75b40b06e6db5 Author: liyuheng <[email protected]> AuthorDate: Thu Jun 26 17:57:50 2025 +0800 seems done --- iotdb-core/ainode/ainode/core/handler.py | 6 ++++++ .../consensus/response/model/ModelTableResp.java | 8 +++++++- .../apache/iotdb/confignode/manager/ModelManager.java | 18 ++++++++++++++++++ .../config/executor/ClusterConfigTaskExecutor.java | 2 +- .../execution/config/metadata/ai/ShowModelsTask.java | 16 +++++++++++----- .../iotdb/commons/client/ainode/AINodeClient.java | 13 +++++++++++++ .../commons/schema/column/ColumnHeaderConstant.java | 2 ++ .../thrift-ainode/src/main/thrift/ainode.thrift | 5 +++++ .../src/main/thrift/confignode.thrift | 1 + 9 files changed, 64 insertions(+), 7 deletions(-) diff --git a/iotdb-core/ainode/ainode/core/handler.py b/iotdb-core/ainode/ainode/core/handler.py index b338286e97b..401a57b870f 100644 --- a/iotdb-core/ainode/ainode/core/handler.py +++ b/iotdb-core/ainode/ainode/core/handler.py @@ -32,6 +32,7 @@ from ainode.thrift.ainode.ttypes import ( TRegisterModelReq, TRegisterModelResp, TTrainingReq, + TModelsArchResp, ) from ainode.thrift.common.ttypes import TSStatus @@ -58,5 +59,10 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface): def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp: return ClusterManager.get_heart_beat(req) + def getArch(self) -> TModelsArchResp: + x = TModelsArchResp(self) + x.archMap = {"_sundial":"good"} + return x + def createTrainingTask(self, req: TTrainingReq) -> TSStatus: pass diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java index 9a23d9ed713..3c456b41269 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/response/model/ModelTableResp.java @@ -28,11 +28,13 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; +import java.util.Map; public class ModelTableResp implements DataSet { private final TSStatus status; private final List<ByteBuffer> serializedAllModelInformation; + private Map<String, String> archMap; public ModelTableResp(TSStatus status) { this.status = status; @@ -50,6 +52,10 @@ public class ModelTableResp implements DataSet { } public TShowModelResp convertToThriftResponse() throws IOException { - return new TShowModelResp(status, serializedAllModelInformation); + return new TShowModelResp(status, serializedAllModelInformation).setArchMap(archMap); + } + + public void setArchMap(Map<String, String> archMap) { + this.archMap = archMap; } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java index 2ebc4d71dfc..76f4f42fa1c 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ModelManager.java @@ -19,7 +19,10 @@ package org.apache.iotdb.confignode.manager; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.ainode.AINodeClient; +import org.apache.iotdb.commons.client.ainode.AINodeClientManager; import org.apache.iotdb.commons.model.ModelInformation; import org.apache.iotdb.commons.model.ModelStatus; import org.apache.iotdb.commons.model.ModelType; @@ -30,6 +33,7 @@ import org.apache.iotdb.confignode.consensus.request.write.model.UpdateModelInfo import org.apache.iotdb.confignode.consensus.response.model.GetModelInfoResp; import org.apache.iotdb.confignode.consensus.response.model.ModelTableResp; import org.apache.iotdb.confignode.persistence.ModelInfo; +import org.apache.iotdb.confignode.rpc.thrift.TAINodeInfo; import org.apache.iotdb.confignode.rpc.thrift.TCreateModelReq; import org.apache.iotdb.confignode.rpc.thrift.TDropModelReq; import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq; @@ -47,6 +51,7 @@ import org.slf4j.LoggerFactory; import java.io.IOException; import java.util.Collections; import java.util.List; +import java.util.Map; public class ModelManager { @@ -92,8 +97,21 @@ public class ModelManager { } public TShowModelResp showModel(final TShowModelReq req) { + TAINodeInfo registeredAINode = + configManager.getNodeManager().getRegisteredAINodeInfoList().get(0); + TEndPoint targetAINodeEndPoint = + new TEndPoint(registeredAINode.getInternalAddress(), registeredAINode.getInternalPort()); + Map<String, String> archMap; + try (AINodeClient client = + AINodeClientManager.getInstance().borrowClient(targetAINodeEndPoint)) { + archMap = client.getArch().getArchMap(); + } catch (Exception e) { + LOGGER.warn("Failed to get archMap", e); + archMap = Collections.emptyMap(); + } try { final DataSet response = configManager.getConsensusManager().read(new ShowModelPlan(req)); + ((ModelTableResp) response).setArchMap(archMap); return ((ModelTableResp) response).convertToThriftResponse(); } catch (final ConsensusException e) { LOGGER.warn( 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 39c2bda7298..eef11e684f8 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 @@ -3321,7 +3321,7 @@ public class ClusterConfigTaskExecutor implements IConfigTaskExecutor { return future; } // convert model info list and buildTsBlock - ShowModelsTask.buildTsBlock(showModelResp.getModelInfoList(), future); + ShowModelsTask.buildTsBlock(showModelResp, future); } catch (final ClientManagerException | TException e) { future.setException(e); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/ShowModelsTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/ShowModelsTask.java index 78ded3da6e9..d4aac9dc91c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/ShowModelsTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/metadata/ai/ShowModelsTask.java @@ -22,6 +22,7 @@ package org.apache.iotdb.db.queryengine.plan.execution.config.metadata.ai; import org.apache.iotdb.commons.model.ModelType; import org.apache.iotdb.commons.schema.column.ColumnHeader; import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; @@ -38,6 +39,7 @@ import org.apache.tsfile.utils.ReadWriteIOUtils; import java.nio.ByteBuffer; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; @@ -62,7 +64,9 @@ public class ShowModelsTask implements IConfigTask { } public static void buildTsBlock( - List<ByteBuffer> modelInfoList, SettableFuture<ConfigTaskResult> future) { + TShowModelResp resp, SettableFuture<ConfigTaskResult> future) { + List<ByteBuffer> modelInfoList = resp.getModelInfoList(); + Map<String, String> archMap = resp.getArchMap(); List<TSDataType> outputDataTypes = ColumnHeaderConstant.showModelsColumnHeaders.stream() .map(ColumnHeader::getColumnType) @@ -71,6 +75,7 @@ public class ShowModelsTask implements IConfigTask { for (ByteBuffer modelInfo : modelInfoList) { String modelId = ReadWriteIOUtils.readString(modelInfo); String modelType = ReadWriteIOUtils.readString(modelInfo); + String modelArch = archMap.getOrDefault(modelId, "Unknown"); String state = ReadWriteIOUtils.readString(modelInfo); String note; String config; @@ -97,12 +102,13 @@ public class ShowModelsTask implements IConfigTask { builder.getTimeColumnBuilder().writeLong(0L); builder.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(modelId)); builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(modelType)); - builder.getColumnBuilder(2).writeBinary(BytesUtils.valueOf(state)); - builder.getColumnBuilder(3).writeBinary(BytesUtils.valueOf(config)); + builder.getColumnBuilder(2).writeBinary(BytesUtils.valueOf(modelArch)); + builder.getColumnBuilder(3).writeBinary(BytesUtils.valueOf(state)); + builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf(config)); if (note != null) { - builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf(note)); + builder.getColumnBuilder(5).writeBinary(BytesUtils.valueOf(note)); } else { - builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf("")); + builder.getColumnBuilder(5).writeBinary(BytesUtils.valueOf("")); } builder.declarePosition(); } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ainode/AINodeClient.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ainode/AINodeClient.java index 346a459136a..a7c10bbd92e 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ainode/AINodeClient.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ainode/AINodeClient.java @@ -26,6 +26,7 @@ import org.apache.iotdb.ainode.rpc.thrift.TForecastReq; import org.apache.iotdb.ainode.rpc.thrift.TForecastResp; import org.apache.iotdb.ainode.rpc.thrift.TInferenceReq; import org.apache.iotdb.ainode.rpc.thrift.TInferenceResp; +import org.apache.iotdb.ainode.rpc.thrift.TModelsArchResp; import org.apache.iotdb.ainode.rpc.thrift.TRegisterModelReq; import org.apache.iotdb.ainode.rpc.thrift.TRegisterModelResp; import org.apache.iotdb.ainode.rpc.thrift.TTrainingReq; @@ -158,6 +159,18 @@ public class AINodeClient implements AutoCloseable, ThriftClient { } } + public TModelsArchResp getArch() throws TException { + try { + return client.getArch(); + } catch (TException e) { + logger.warn( + "Failed to connect to AINode from ConfigNode when executing {}: {}", + Thread.currentThread().getStackTrace()[1].getMethodName(), + e.getMessage()); + throw new TException(MSG_CONNECTION_FAIL); + } + } + public TInferenceResp inference( String modelId, TsBlock inputTsBlock, diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java index 72fe1ce0236..325127ab8b5 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java @@ -278,6 +278,7 @@ public class ColumnHeaderConstant { // column names for show models/trials public static final String MODEL_ID = "ModelId"; + public static final String ARCHITECTURE = "Architecture"; // column names for views (e.g. logical view) public static final String VIEW_TYPE = "ViewType"; @@ -612,6 +613,7 @@ public class ColumnHeaderConstant { ImmutableList.of( new ColumnHeader(MODEL_ID, TSDataType.TEXT), new ColumnHeader(MODEL_TYPE, TSDataType.TEXT), + new ColumnHeader(ARCHITECTURE, TSDataType.TEXT), new ColumnHeader(STATE, TSDataType.TEXT), new ColumnHeader(CONFIGS, TSDataType.TEXT), new ColumnHeader(NOTES, TSDataType.TEXT)); diff --git a/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift b/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift index db1a15c2460..dd1ad4afd0b 100644 --- a/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift +++ b/iotdb-protocol/thrift-ainode/src/main/thrift/ainode.thrift @@ -98,9 +98,14 @@ struct TForecastResp { 2: required binary forecastResult } +struct TModelsArchResp { + 1: required map<string, string> archMap +} + service IAINodeRPCService { // -------------- For Config Node -------------- + TModelsArchResp getArch() common.TSStatus deleteModel(TDeleteModelReq req) diff --git a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift index f767f35d67c..2e4327d653a 100644 --- a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift +++ b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift @@ -1064,6 +1064,7 @@ struct TShowModelReq { struct TShowModelResp { 1: required common.TSStatus status 2: required list<binary> modelInfoList + 3: optional map<string, string> archMap } struct TGetModelInfoReq {
