Add PID metadata to XPlanes and TraceMeConsumers to enable producer/consumer differentiation across processes profiled together.

PiperOrigin-RevId: 824554472

Author: bmass02

tsl/profiler/lib/BUILD

@@ -217,6 +217,7 @@ cc_library(
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
+        "@xla//xla/tsl/platform:env",
         "@xla//xla/tsl/platform:errors",
         "@xla//xla/tsl/platform:logging",
         "@xla//xla/tsl/platform:types",
@@ -357,8 +358,6 @@ cc_library(
         ":context_types_hdrs",
         ":traceme",
         ":traceme_encode",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:optional",
         "@xla//xla/tsl/platform:types",
     ],
 )

tsl/profiler/lib/connected_traceme.h

@@ -16,11 +16,8 @@ limitations under the License.
 #define TENSORFLOW_TSL_PROFILER_LIB_CONNECTED_TRACEME_H_
 
 #include <optional>
-#include <string>
 #include <utility>
 
-#include "absl/strings/string_view.h"
-#include "absl/types/optional.h"
 #include "xla/tsl/platform/types.h"
 #include "tsl/profiler/lib/context_types.h"
 #include "tsl/profiler/lib/traceme.h"

tsl/profiler/lib/profiler_session.cc

@@ -30,6 +30,7 @@ limitations under the License.
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
 #if !defined(IS_MOBILE_PLATFORM)
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/profiler/convert/post_process_single_host_xplane.h"
 #include "xla/tsl/profiler/utils/time_utils.h"
 #include "tsl/platform/host_info.h"
@@ -95,6 +96,7 @@ absl::Status ProfilerSession::CollectData(XSpace* space) {
 #if !defined(IS_MOBILE_PLATFORM)
   space->add_hostnames(port::Hostname());
   TF_RETURN_IF_ERROR(CollectDataInternal(space));
+  profiler::SetXSpacePidIfNotSet(*space, tsl::Env::Default()->GetProcessId());
   profiler::PostProcessSingleHostXSpace(space, start_time_ns_, stop_time_ns_);
 #endif
   SetProfileOptionsIntoSpace(options_, space);