Skip to content

Commit

Permalink
Retrieve remote model id from registration response in IT to avoid fl…
Browse files Browse the repository at this point in the history
…aky (opensearch-project#3244)

Signed-off-by: zane-neo <[email protected]>
  • Loading branch information
zane-neo authored Dec 3, 2024
1 parent 78a304a commit 1d30671
Showing 1 changed file with 1 addition and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1002,15 +1002,11 @@ public String registerRemoteModel(String createConnectorInput, String modelName,
String connectorId = (String) responseMap.get("connector_id");
response = RestMLRemoteInferenceIT.registerRemoteModel(modelName, modelName, connectorId);
responseMap = parseResponseToMap(response);
String taskId = (String) responseMap.get("task_id");
waitForTask(taskId, MLTaskState.COMPLETED);
response = RestMLRemoteInferenceIT.getTask(taskId);
responseMap = parseResponseToMap(response);
String modelId = (String) responseMap.get("model_id");
if (deploy) {
response = RestMLRemoteInferenceIT.deployRemoteModel(modelId);
responseMap = parseResponseToMap(response);
taskId = (String) responseMap.get("task_id");
String taskId = (String) responseMap.get("task_id");
waitForTask(taskId, MLTaskState.COMPLETED);
}
return modelId;
Expand Down

0 comments on commit 1d30671

Please sign in to comment.