From 5de004e0803b93c18e41c6896892589ec166ef29 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 18 Dec 2025 16:19:33 -0800
Subject: [PATCH] Fix TraceMe context propagation for some inference requests
 by replacing `kPathwaysExecutor` and `kTfrtExecutor` with a new
 enum`kGenericModelExecutor` and add back some tracing

PiperOrigin-RevId: 846455577
---
 tsl/profiler/lib/context_types.cc | 2 ++
 tsl/profiler/lib/context_types.h  | 7 ++++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/tsl/profiler/lib/context_types.cc b/tsl/profiler/lib/context_types.cc
index 58e63b17d..2f54dc18a 100644
--- a/tsl/profiler/lib/context_types.cc
+++ b/tsl/profiler/lib/context_types.cc
@@ -54,6 +54,8 @@ const char* GetContextTypeString(ContextType context_type) {
       return "jax_serving";
     case ContextType::kScOffload:
       return "sparsecore_offload";
+    case ContextType::kGenericModelExecutor:
+      return "generic_model_executor";
   }
 }
 
diff --git a/tsl/profiler/lib/context_types.h b/tsl/profiler/lib/context_types.h
index 42389c1db..a986fd014 100644
--- a/tsl/profiler/lib/context_types.h
+++ b/tsl/profiler/lib/context_types.h
@@ -25,7 +25,7 @@ enum class ContextType : int {
   kGeneric = 0,
   kLegacy,
   kTfExecutor,
-  kTfrtExecutor,
+  kTfrtExecutor [[deprecated]],
   kSharedBatchScheduler,
   kPjRt,
   kAdaptiveSharedBatchScheduler,
@@ -35,12 +35,13 @@ enum class ContextType : int {
   kBatcher,
   kTpuStream,
   kTpuLaunch,
-  kPathwaysExecutor,
+  kPathwaysExecutor [[deprecated]],
   kPjrtLibraryCall,
   kThreadpoolEvent,
   kJaxServingExecutor,
   kScOffload,
-  kLastContextType = ContextType::kScOffload,
+  kGenericModelExecutor,
+  kLastContextType = ContextType::kGenericModelExecutor,
 };
 
 // In XFlow we encode context type as flow category as 6 bits.