Add a trace type for Jax serving

SiqiaoWu1993 · Google-ML-Automation · commit 37b580d0dd7b · 2025-05-13T11:35:00.000-07:00
PiperOrigin-RevId: 758303660
diff --git a/third_party/tsl/tsl/profiler/lib/context_types.cc b/third_party/tsl/tsl/profiler/lib/context_types.cc
@@ -50,6 +50,8 @@ const char* GetContextTypeString(ContextType context_type) {
       return "pjrt_library_call";
     case ContextType::kThreadpoolEvent:
       return "threadpool_event";
+    case ContextType::kJaxServingExecutor:
+      return "jax_serving";
   }
 }
 
diff --git a/third_party/tsl/tsl/profiler/lib/context_types.h b/third_party/tsl/tsl/profiler/lib/context_types.h
@@ -38,7 +38,8 @@ enum class ContextType : int {
   kPathwaysExecutor,
   kPjrtLibraryCall,
   kThreadpoolEvent,
-  kLastContextType = ContextType::kTpuLaunch,
+  kJaxServingExecutor,
+  kLastContextType = ContextType::kJaxServingExecutor,
 };
 
 // In XFlow we encode context type as flow category as 6 bits.

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,8 @@ const char* GetContextTypeString(ContextType context_type) {`
`50`	`50`	`return "pjrt_library_call";`
`51`	`51`	`case ContextType::kThreadpoolEvent:`
`52`	`52`	`return "threadpool_event";`
	`53`	`+ case ContextType::kJaxServingExecutor:`
	`54`	`+ return "jax_serving";`
`53`	`55`	`}`
`54`	`56`	`}`
`55`	`57`