-
Notifications
You must be signed in to change notification settings - Fork 1.1k
[AINode] Update AINodeClient for DataNode to borrow #16647
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PTAL.
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PTAL.
iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
Outdated
Show resolved
Hide resolved
iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java
Outdated
Show resolved
Hide resolved
iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/client/ainode/AINodeClient.java
Outdated
Show resolved
Hide resolved
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PTAL.
...rg/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
Outdated
Show resolved
Hide resolved
iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
Outdated
Show resolved
Hide resolved
iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
# Conflicts: # iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
cd3b39a to
6ba848f
Compare
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PTAL.
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/AINodeClientFactory.java
Show resolved
Hide resolved
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PTAL.
integration-test/src/main/java/org/apache/iotdb/it/env/cluster/env/AbstractEnv.java
Outdated
Show resolved
Hide resolved
integration-test/src/main/java/org/apache/iotdb/it/env/remote/env/RemoteServerEnv.java
Outdated
Show resolved
Hide resolved
.../src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java
Outdated
Show resolved
Hide resolved
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ainode/AINodeClient.java
Outdated
Show resolved
Hide resolved
iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ainode/AINodeClient.java
Outdated
Show resolved
Hide resolved
...rg/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
Outdated
Show resolved
Hide resolved
...n/java/org/apache/iotdb/commons/consensus/iotv2/container/IoTV2GlobalComponentContainer.java
Outdated
Show resolved
Hide resolved
...rg/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java
Outdated
Show resolved
Hide resolved
|
LGTM~ |
CRZbulabula
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM!!!
Description
This PR refactors the AINode client infrastructure to support direct communication between DataNode and AINode, removing the dependency on ConfigNode for AI-related operations such as model loading and inference. Should be reviewed by @CRZbulabula
Contents
AINodeClientAdded a new executeRemoteCallWithRetry() method for automatic retry and reconnection on Thrift transport failures, following the same design pattern as ConfigNodeClient.
Updated the loadModel(TLoadModelReq req) API to use this retry wrapper for improved resilience.
Simplified connection lifecycle management (init(), close()) to ensure stable client reuse via AINodeClientManager.
ClusterConfigTaskExecutorReplaced indirect ConfigNode RPCs with direct calls to AINodeClientManager.borrowClient(TEndPoint) for model operations (currently loadModel as an example).
Ensured the DataNode→AINode invocation flow mirrors the ConfigNode client style while maintaining compatibility with existing client pooling.
Updated Thrift imports to use org.apache.iotdb.ainode.rpc.thrift.* instead of org.apache.iotdb.confignode.rpc.thrift.*.
AINodeClientManagerImpact
DataNode can now directly send AI-related requests (e.g., model load/unload, inference) to AINode without routing through ConfigNode.
Next Steps
Extend the same direct invocation pattern (AINodeClientManager.borrowClient()) to other AI APIs:
unloadModel, showModel, showLoadedModel, showAIDevices, createTraining, and getModelInfo.
This PR has:
for an unfamiliar reader.
for code coverage.
Key changed/added classes (or packages if there are too many classes) in this PR
As former.