Skip to content

Commit 9a4b680

Browse files
committed
Merge branch 'master' of https://github.com/apache/iotdb into fix-audit-logger
2 parents 90dcbd3 + 2b47be7 commit 9a4b680

File tree

38 files changed

+543
-482
lines changed

38 files changed

+543
-482
lines changed

iotdb-core/ainode/build_binary.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ def install_dependencies(venv_python, venv_dir, script_dir):
324324
shutil.rmtree(poetry_venv_path, ignore_errors=True)
325325
poetry_venv_path.parent.mkdir(parents=True, exist_ok=True)
326326
poetry_venv_path.symlink_to(venv_dir)
327-
print(f"Symlink created successfully")
327+
print(f"Symlink created successfully")
328328
except Exception as e:
329329
print(f"WARNING: Failed to create symlink: {e}")
330330
print("Will try to use poetry install directly with VIRTUAL_ENV set")
@@ -386,9 +386,7 @@ def install_dependencies(venv_python, venv_dir, script_dir):
386386
print("The symlink approach may not have worked. Please check the symlink.")
387387
sys.exit(1)
388388
else:
389-
print(
390-
f"✓ Poetry is correctly using virtual environment: {poetry_venv_path}"
391-
)
389+
print(f"Poetry is correctly using virtual environment: {poetry_venv_path}")
392390
else:
393391
print("Warning: Could not verify poetry virtual environment path")
394392
print(
@@ -465,12 +463,12 @@ def verify_poetry_env():
465463
)
466464
if test_result.returncode == 0:
467465
version = test_result.stdout.strip()
468-
print(f"{package} {version} installed")
466+
print(f"{package} {version} installed")
469467
else:
470468
error_msg = (
471469
test_result.stderr.strip() if test_result.stderr else "Unknown error"
472470
)
473-
print(f"{package} NOT found in virtual environment: {error_msg}")
471+
print(f"{package} NOT found in virtual environment: {error_msg}")
474472
missing_packages.append(package)
475473

476474
if missing_packages:

iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -616,7 +616,11 @@ def prepare_inputs_for_generation(
616616
if attention_mask is not None and attention_mask.shape[1] > (
617617
input_ids.shape[1] // self.config.input_token_len
618618
):
619-
input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
619+
input_ids = input_ids[
620+
:,
621+
-(attention_mask.shape[1] - past_length)
622+
* self.config.input_token_len :,
623+
]
620624
# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
621625
# input_ids based on the past_length.
622626
elif past_length < (input_ids.shape[1] // self.config.input_token_len):
@@ -629,9 +633,10 @@ def prepare_inputs_for_generation(
629633
position_ids = attention_mask.long().cumsum(-1) - 1
630634
position_ids.masked_fill_(attention_mask == 0, 1)
631635
if past_key_values:
632-
position_ids = position_ids[
633-
:, -(input_ids.shape[1] // self.config.input_token_len) :
634-
]
636+
token_num = (
637+
input_ids.shape[1] + self.config.input_token_len - 1
638+
) // self.config.input_token_len
639+
position_ids = position_ids[:, -token_num:]
635640

636641
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
637642
if inputs_embeds is not None and past_key_values is None:

iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -606,7 +606,11 @@ def prepare_inputs_for_generation(
606606
if attention_mask is not None and attention_mask.shape[1] > (
607607
input_ids.shape[1] // self.config.input_token_len
608608
):
609-
input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
609+
input_ids = input_ids[
610+
:,
611+
-(attention_mask.shape[1] - past_length)
612+
* self.config.input_token_len :,
613+
]
610614
# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
611615
# input_ids based on the past_length.
612616
elif past_length < (input_ids.shape[1] // self.config.input_token_len):

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/AsyncAINodeHeartbeatClientPool.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@
2121

2222
import org.apache.iotdb.ainode.rpc.thrift.TAIHeartbeatReq;
2323
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
24-
import org.apache.iotdb.commons.client.ClientPoolFactory;
2524
import org.apache.iotdb.commons.client.IClientManager;
26-
import org.apache.iotdb.commons.client.ainode.AsyncAINodeServiceClient;
2725
import org.apache.iotdb.confignode.client.async.handlers.heartbeat.AINodeHeartbeatHandler;
26+
import org.apache.iotdb.db.protocol.client.AINodeClientFactory;
27+
import org.apache.iotdb.db.protocol.client.ainode.AsyncAINodeServiceClient;
2828

2929
public class AsyncAINodeHeartbeatClientPool {
3030

@@ -33,8 +33,7 @@ public class AsyncAINodeHeartbeatClientPool {
3333
private AsyncAINodeHeartbeatClientPool() {
3434
clientManager =
3535
new IClientManager.Factory<TEndPoint, AsyncAINodeServiceClient>()
36-
.createClientManager(
37-
new ClientPoolFactory.AsyncAINodeHeartbeatServiceClientPoolFactory());
36+
.createClientManager(new AINodeClientFactory.AINodeHeartbeatClientPoolFactory());
3837
}
3938

4039
public void getAINodeHeartBeat(

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/read/model/ShowModelPlan.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@
1919

2020
package org.apache.iotdb.confignode.consensus.request.read.model;
2121

22+
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
2223
import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
2324
import org.apache.iotdb.confignode.consensus.request.read.ConfigPhysicalReadPlan;
24-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
2525

2626
import java.util.Objects;
2727

@@ -33,7 +33,7 @@ public ShowModelPlan() {
3333
super(ConfigPhysicalPlanType.ShowModel);
3434
}
3535

36-
public ShowModelPlan(final TShowModelReq showModelReq) {
36+
public ShowModelPlan(final TShowModelsReq showModelReq) {
3737
super(ConfigPhysicalPlanType.ShowModel);
3838
if (showModelReq.isSetModelId()) {
3939
this.modelName = showModelReq.getModelId();

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

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,14 @@
2020
package org.apache.iotdb.confignode.manager;
2121

2222
import org.apache.iotdb.ainode.rpc.thrift.IDataSchema;
23+
import org.apache.iotdb.ainode.rpc.thrift.TLoadModelReq;
24+
import org.apache.iotdb.ainode.rpc.thrift.TShowAIDevicesResp;
25+
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsReq;
26+
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsResp;
27+
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
28+
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsResp;
2329
import org.apache.iotdb.ainode.rpc.thrift.TTrainingReq;
30+
import org.apache.iotdb.ainode.rpc.thrift.TUnloadModelReq;
2431
import org.apache.iotdb.common.rpc.thrift.TAINodeConfiguration;
2532
import org.apache.iotdb.common.rpc.thrift.TAINodeLocation;
2633
import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation;
@@ -42,8 +49,6 @@
4249
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
4350
import org.apache.iotdb.commons.auth.AuthException;
4451
import org.apache.iotdb.commons.auth.entity.PrivilegeUnion;
45-
import org.apache.iotdb.commons.client.ainode.AINodeClient;
46-
import org.apache.iotdb.commons.client.ainode.AINodeClientManager;
4752
import org.apache.iotdb.commons.cluster.NodeStatus;
4853
import org.apache.iotdb.commons.cluster.NodeType;
4954
import org.apache.iotdb.commons.conf.CommonConfig;
@@ -214,7 +219,6 @@
214219
import org.apache.iotdb.confignode.rpc.thrift.TGetTriggerTableResp;
215220
import org.apache.iotdb.confignode.rpc.thrift.TGetUDFTableResp;
216221
import org.apache.iotdb.confignode.rpc.thrift.TGetUdfTableReq;
217-
import org.apache.iotdb.confignode.rpc.thrift.TLoadModelReq;
218222
import org.apache.iotdb.confignode.rpc.thrift.TMigrateRegionReq;
219223
import org.apache.iotdb.confignode.rpc.thrift.TNodeVersionInfo;
220224
import org.apache.iotdb.confignode.rpc.thrift.TPermissionInfoResp;
@@ -227,7 +231,6 @@
227231
import org.apache.iotdb.confignode.rpc.thrift.TSchemaPartitionTableResp;
228232
import org.apache.iotdb.confignode.rpc.thrift.TSetDataNodeStatusReq;
229233
import org.apache.iotdb.confignode.rpc.thrift.TSetSchemaTemplateReq;
230-
import org.apache.iotdb.confignode.rpc.thrift.TShowAIDevicesResp;
231234
import org.apache.iotdb.confignode.rpc.thrift.TShowAINodesResp;
232235
import org.apache.iotdb.confignode.rpc.thrift.TShowCQResp;
233236
import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
@@ -236,10 +239,6 @@
236239
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodes4InformationSchemaResp;
237240
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp;
238241
import org.apache.iotdb.confignode.rpc.thrift.TShowDatabaseResp;
239-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelReq;
240-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelResp;
241-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
242-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
243242
import org.apache.iotdb.confignode.rpc.thrift.TShowPipePluginReq;
244243
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
245244
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
@@ -257,12 +256,13 @@
257256
import org.apache.iotdb.confignode.rpc.thrift.TSubscribeReq;
258257
import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp;
259258
import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList;
260-
import org.apache.iotdb.confignode.rpc.thrift.TUnloadModelReq;
261259
import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq;
262260
import org.apache.iotdb.confignode.rpc.thrift.TUnsubscribeReq;
263261
import org.apache.iotdb.confignode.rpc.thrift.TUpdateModelInfoReq;
264262
import org.apache.iotdb.consensus.common.DataSet;
265263
import org.apache.iotdb.consensus.exception.ConsensusException;
264+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClient;
265+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClientManager;
266266
import org.apache.iotdb.db.schemaengine.template.TemplateAlterOperationType;
267267
import org.apache.iotdb.db.schemaengine.template.alter.TemplateAlterOperationUtil;
268268
import org.apache.iotdb.rpc.RpcUtils;
@@ -2863,19 +2863,19 @@ public TSStatus unloadModel(TUnloadModelReq req) {
28632863
}
28642864

28652865
@Override
2866-
public TShowModelResp showModel(TShowModelReq req) {
2866+
public TShowModelsResp showModel(TShowModelsReq req) {
28672867
TSStatus status = confirmLeader();
28682868
return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
28692869
? modelManager.showModel(req)
2870-
: new TShowModelResp(status);
2870+
: new TShowModelsResp(status);
28712871
}
28722872

28732873
@Override
2874-
public TShowLoadedModelResp showLoadedModel(TShowLoadedModelReq req) {
2874+
public TShowLoadedModelsResp showLoadedModel(TShowLoadedModelsReq req) {
28752875
TSStatus status = confirmLeader();
28762876
return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()
28772877
? modelManager.showLoadedModel(req)
2878-
: new TShowLoadedModelResp(status, Collections.emptyMap());
2878+
: new TShowLoadedModelsResp(status, Collections.emptyMap());
28792879
}
28802880

28812881
@Override

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,13 @@
1919

2020
package org.apache.iotdb.confignode.manager;
2121

22+
import org.apache.iotdb.ainode.rpc.thrift.TLoadModelReq;
23+
import org.apache.iotdb.ainode.rpc.thrift.TShowAIDevicesResp;
24+
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsReq;
25+
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsResp;
26+
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
27+
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsResp;
28+
import org.apache.iotdb.ainode.rpc.thrift.TUnloadModelReq;
2229
import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation;
2330
import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
2431
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
@@ -128,7 +135,6 @@
128135
import org.apache.iotdb.confignode.rpc.thrift.TGetTriggerTableResp;
129136
import org.apache.iotdb.confignode.rpc.thrift.TGetUDFTableResp;
130137
import org.apache.iotdb.confignode.rpc.thrift.TGetUdfTableReq;
131-
import org.apache.iotdb.confignode.rpc.thrift.TLoadModelReq;
132138
import org.apache.iotdb.confignode.rpc.thrift.TMigrateRegionReq;
133139
import org.apache.iotdb.confignode.rpc.thrift.TPermissionInfoResp;
134140
import org.apache.iotdb.confignode.rpc.thrift.TPipeConfigTransferReq;
@@ -140,7 +146,6 @@
140146
import org.apache.iotdb.confignode.rpc.thrift.TSchemaPartitionTableResp;
141147
import org.apache.iotdb.confignode.rpc.thrift.TSetDataNodeStatusReq;
142148
import org.apache.iotdb.confignode.rpc.thrift.TSetSchemaTemplateReq;
143-
import org.apache.iotdb.confignode.rpc.thrift.TShowAIDevicesResp;
144149
import org.apache.iotdb.confignode.rpc.thrift.TShowAINodesResp;
145150
import org.apache.iotdb.confignode.rpc.thrift.TShowCQResp;
146151
import org.apache.iotdb.confignode.rpc.thrift.TShowClusterResp;
@@ -149,10 +154,6 @@
149154
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodes4InformationSchemaResp;
150155
import org.apache.iotdb.confignode.rpc.thrift.TShowDataNodesResp;
151156
import org.apache.iotdb.confignode.rpc.thrift.TShowDatabaseResp;
152-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelReq;
153-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelResp;
154-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
155-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
156157
import org.apache.iotdb.confignode.rpc.thrift.TShowPipePluginReq;
157158
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
158159
import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
@@ -166,7 +167,6 @@
166167
import org.apache.iotdb.confignode.rpc.thrift.TStartPipeReq;
167168
import org.apache.iotdb.confignode.rpc.thrift.TStopPipeReq;
168169
import org.apache.iotdb.confignode.rpc.thrift.TSubscribeReq;
169-
import org.apache.iotdb.confignode.rpc.thrift.TUnloadModelReq;
170170
import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq;
171171
import org.apache.iotdb.confignode.rpc.thrift.TUnsubscribeReq;
172172
import org.apache.iotdb.consensus.common.DataSet;
@@ -893,10 +893,10 @@ TDataPartitionTableResp getOrCreateDataPartition(
893893
TSStatus unloadModel(TUnloadModelReq req);
894894

895895
/** Return the model table. */
896-
TShowModelResp showModel(TShowModelReq req);
896+
TShowModelsResp showModel(TShowModelsReq req);
897897

898898
/** Return the loaded model instances. */
899-
TShowLoadedModelResp showLoadedModel(TShowLoadedModelReq req);
899+
TShowLoadedModelsResp showLoadedModel(TShowLoadedModelsReq req);
900900

901901
/** Return all available AI devices. */
902902
TShowAIDevicesResp showAIDevices();

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

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,15 @@
1919

2020
package org.apache.iotdb.confignode.manager;
2121

22+
import org.apache.iotdb.ainode.rpc.thrift.TLoadModelReq;
23+
import org.apache.iotdb.ainode.rpc.thrift.TShowAIDevicesResp;
2224
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsReq;
2325
import org.apache.iotdb.ainode.rpc.thrift.TShowLoadedModelsResp;
2426
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsReq;
2527
import org.apache.iotdb.ainode.rpc.thrift.TShowModelsResp;
28+
import org.apache.iotdb.ainode.rpc.thrift.TUnloadModelReq;
2629
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
2730
import org.apache.iotdb.common.rpc.thrift.TSStatus;
28-
import org.apache.iotdb.commons.client.ainode.AINodeClient;
29-
import org.apache.iotdb.commons.client.ainode.AINodeClientManager;
3031
import org.apache.iotdb.commons.client.exception.ClientManagerException;
3132
import org.apache.iotdb.commons.model.ModelInformation;
3233
import org.apache.iotdb.commons.model.ModelStatus;
@@ -40,15 +41,10 @@
4041
import org.apache.iotdb.confignode.rpc.thrift.TDropModelReq;
4142
import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoReq;
4243
import org.apache.iotdb.confignode.rpc.thrift.TGetModelInfoResp;
43-
import org.apache.iotdb.confignode.rpc.thrift.TLoadModelReq;
44-
import org.apache.iotdb.confignode.rpc.thrift.TShowAIDevicesResp;
45-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelReq;
46-
import org.apache.iotdb.confignode.rpc.thrift.TShowLoadedModelResp;
47-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelReq;
48-
import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
49-
import org.apache.iotdb.confignode.rpc.thrift.TUnloadModelReq;
5044
import org.apache.iotdb.confignode.rpc.thrift.TUpdateModelInfoReq;
5145
import org.apache.iotdb.consensus.exception.ConsensusException;
46+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClient;
47+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClientManager;
5248
import org.apache.iotdb.rpc.TSStatusCode;
5349

5450
import org.slf4j.Logger;
@@ -124,42 +120,43 @@ public TSStatus unloadModel(TUnloadModelReq req) {
124120
}
125121
}
126122

127-
public TShowModelResp showModel(final TShowModelReq req) {
123+
public TShowModelsResp showModel(final TShowModelsReq req) {
128124
try (AINodeClient client = getAINodeClient()) {
129125
TShowModelsReq showModelsReq = new TShowModelsReq();
130126
if (req.isSetModelId()) {
131127
showModelsReq.setModelId(req.getModelId());
132128
}
133129
TShowModelsResp resp = client.showModels(showModelsReq);
134-
TShowModelResp res =
135-
new TShowModelResp().setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
130+
TShowModelsResp res =
131+
new TShowModelsResp()
132+
.setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
136133
res.setModelIdList(resp.getModelIdList());
137134
res.setModelTypeMap(resp.getModelTypeMap());
138135
res.setCategoryMap(resp.getCategoryMap());
139136
res.setStateMap(resp.getStateMap());
140137
return res;
141138
} catch (Exception e) {
142139
LOGGER.warn("Failed to show models due to", e);
143-
return new TShowModelResp()
140+
return new TShowModelsResp()
144141
.setStatus(
145142
new TSStatus(TSStatusCode.AI_NODE_INTERNAL_ERROR.getStatusCode())
146143
.setMessage(e.getMessage()));
147144
}
148145
}
149146

150-
public TShowLoadedModelResp showLoadedModel(final TShowLoadedModelReq req) {
147+
public TShowLoadedModelsResp showLoadedModel(final TShowLoadedModelsReq req) {
151148
try (AINodeClient client = getAINodeClient()) {
152149
TShowLoadedModelsReq showModelsReq =
153150
new TShowLoadedModelsReq().setDeviceIdList(req.getDeviceIdList());
154151
TShowLoadedModelsResp resp = client.showLoadedModels(showModelsReq);
155-
TShowLoadedModelResp res =
156-
new TShowLoadedModelResp()
152+
TShowLoadedModelsResp res =
153+
new TShowLoadedModelsResp()
157154
.setStatus(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
158155
res.setDeviceLoadedModelsMap(resp.getDeviceLoadedModelsMap());
159156
return res;
160157
} catch (Exception e) {
161158
LOGGER.warn("Failed to show loaded models due to", e);
162-
return new TShowLoadedModelResp()
159+
return new TShowLoadedModelsResp()
163160
.setStatus(
164161
new TSStatus(TSStatusCode.AI_NODE_INTERNAL_ERROR.getStatusCode())
165162
.setMessage(e.getMessage()));
@@ -235,7 +232,11 @@ private AINodeClient getAINodeClient() throws NoAvailableAINodeException, Client
235232
}
236233
TEndPoint targetAINodeEndPoint =
237234
new TEndPoint(aiNodeInfo.get(0).getInternalAddress(), aiNodeInfo.get(0).getInternalPort());
238-
return AINodeClientManager.getInstance().borrowClient(targetAINodeEndPoint);
235+
try {
236+
return AINodeClientManager.getInstance().borrowClient(targetAINodeEndPoint);
237+
} catch (Exception e) {
238+
throw new RuntimeException(e);
239+
}
239240
}
240241

241242
public List<Integer> getModelDistributions(String modelName) {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/model/CreateModelProcedure.java

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

2222
import org.apache.iotdb.common.rpc.thrift.TAINodeConfiguration;
2323
import org.apache.iotdb.common.rpc.thrift.TSStatus;
24-
import org.apache.iotdb.commons.client.ainode.AINodeClient;
25-
import org.apache.iotdb.commons.client.ainode.AINodeClientManager;
2624
import org.apache.iotdb.commons.exception.ainode.LoadModelException;
2725
import org.apache.iotdb.commons.model.ModelInformation;
2826
import org.apache.iotdb.commons.model.ModelStatus;
@@ -36,6 +34,8 @@
3634
import org.apache.iotdb.confignode.procedure.state.model.CreateModelState;
3735
import org.apache.iotdb.confignode.procedure.store.ProcedureType;
3836
import org.apache.iotdb.consensus.exception.ConsensusException;
37+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClient;
38+
import org.apache.iotdb.db.protocol.client.ainode.AINodeClientManager;
3939
import org.apache.iotdb.rpc.TSStatusCode;
4040

4141
import org.apache.tsfile.utils.ReadWriteIOUtils;

0 commit comments

Comments
 (0)