Skip to content

Commit e849a1c

Browse files
Fix TraceMe context propagation for pathways inference requests by replacing kPathwaysExecutor and kTfrtExecutor with a new enumkGenericModelExecutor and add back some tracing
PiperOrigin-RevId: 846455577
1 parent 90cb856 commit e849a1c

File tree

2 files changed

+5
-2
lines changed

2 files changed

+5
-2
lines changed

tsl/profiler/lib/context_types.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,8 @@ const char* GetContextTypeString(ContextType context_type) {
5454
return "jax_serving";
5555
case ContextType::kScOffload:
5656
return "sparsecore_offload";
57+
case ContextType::kGenericModelExecutor:
58+
return "generic_model_executor";
5759
}
5860
}
5961

tsl/profiler/lib/context_types.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ enum class ContextType : int {
2525
kGeneric = 0,
2626
kLegacy,
2727
kTfExecutor,
28-
kTfrtExecutor,
28+
kTfrtExecutor [[deprecated]],
2929
kSharedBatchScheduler,
3030
kPjRt,
3131
kAdaptiveSharedBatchScheduler,
@@ -35,12 +35,13 @@ enum class ContextType : int {
3535
kBatcher,
3636
kTpuStream,
3737
kTpuLaunch,
38-
kPathwaysExecutor,
38+
kPathwaysExecutor [[deprecated]],
3939
kPjrtLibraryCall,
4040
kThreadpoolEvent,
4141
kJaxServingExecutor,
4242
kScOffload,
4343
kLastContextType = ContextType::kScOffload,
44+
kGenericModelExecutor,
4445
};
4546

4647
// In XFlow we encode context type as flow category as 6 bits.

0 commit comments

Comments
 (0)