Browse Source

[ML] Start a new trace context before loading a trained model (#103124)

Each distinct task is a different span in APM tracing, so needs
a new trace context.
David Roberts 1 year ago
parent
commit
d8bbc62c6e

+ 5 - 0
docs/changelog/103124.yaml

@@ -0,0 +1,5 @@
+pr: 103124
+summary: Start a new trace context before loading a trained model
+area: Machine Learning
+type: bug
+issues: []

+ 7 - 2
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentNodeService.java

@@ -369,9 +369,14 @@ public class TrainedModelAssignmentNodeService implements ClusterStateListener {
                     }
 
                     if (shouldLoadModel(routingInfo, trainedModelAssignment.getDeploymentId(), isResetMode)) {
-                        prepareModelToLoad(
-                            createStartTrainedModelDeploymentTaskParams(trainedModelAssignment, routingInfo.getCurrentAllocations())
+                        StartTrainedModelDeploymentAction.TaskParams params = createStartTrainedModelDeploymentTaskParams(
+                            trainedModelAssignment,
+                            routingInfo.getCurrentAllocations()
                         );
+                        // Loading the model is done by a separate task, so needs a new trace context
+                        try (var ignored = threadPool.getThreadContext().newTraceContext()) {
+                            prepareModelToLoad(params);
+                        }
                     }
                 }