From d8bbc62c6e228774f103da9f879a3bae5a676149 Mon Sep 17 00:00:00 2001
From: David Roberts <dave.roberts@elastic.co>
Date: Thu, 7 Dec 2023 16:42:56 +0000
Subject: [PATCH] [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.
---
 docs/changelog/103124.yaml                               | 5 +++++
 .../assignment/TrainedModelAssignmentNodeService.java    | 9 +++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)
 create mode 100644 docs/changelog/103124.yaml

diff --git a/docs/changelog/103124.yaml b/docs/changelog/103124.yaml
new file mode 100644
index 0000000000000..078c8249bbf5d
--- /dev/null
+++ b/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: []
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentNodeService.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentNodeService.java
index 3efae0ed58bf6..fdb007862cfdc 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentNodeService.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentNodeService.java
@@ -369,9 +369,14 @@ public void clusterChanged(ClusterChangedEvent event) {
                     }
 
                     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);
+                        }
                     }
                 }