From 622f9b5f5f4ca6308c33050923f97dadb643dbb1 Mon Sep 17 00:00:00 2001 From: justinlu Date: Fri, 4 Sep 2026 21:13:46 -0700 Subject: [PATCH] Add lightweight, zero-overhead tracing macros in `tpu_sync/common/trace.h` supporting both TraceMe and Perfetto backends. * Define `RAIDEN_TRACE(name, ...)`, `RAIDEN_TRACE_SCOPE(name, ...)`, and `RAIDEN_TRACE_FN(name, metadata_fn, ...)` wrapping `ScopedTrace` for fine-grained host timeline instrumentation. * Support `tsl::profiler::TraceMe` and `Perfetto` track event backends, selectable at runtime via `TPU_RAIDEN_TRACE_BACKEND` or `RAIDEN_TRACE_BACKEND` ("traceme", "perfetto", "both", "none"), or via `TPU_RAIDEN_USE_PERFETTO` / `RAIDEN_USE_PERFETTO`. * Add file-based trace export via `StartPerfettoTraceToFile()` / `StopPerfettoTraceToFile()` and automated file capture using `TPU_RAIDEN_PERFETTO_OUTPUT_FILE` with `std::atexit` flush. * Annotate core Raiden components with sample `RAIDEN_TRACE` macros across `internal/raiden_component.cc`, `internal/weight_sync_component.cc`, `internal/raiden_server.cc`, and `weight_sync/weight_synchronizer_base.cc`. * Support explicit backend tracing via `RAIDEN_PERFETTO_TRACE(...)` and `RAIDEN_TRACEME_TRACE(...)`. * Implement thread-safe initialization and track event category registration for Perfetto in `tpu_sync/common/trace.cc`. * Provide programmatic backend override via `SetTraceBackendForTesting(...)` for unit tests. * Define `RAIDEN_TRACE_LOG(...)` and `IsTraceLoggingEnabled()` for debug trace logging to stderr controlled by `TPU_RAIDEN_TRACE` and `RAIDEN_TRACE` environment variables. * Add comprehensive unit tests in `tpu_sync/common/trace_test.cc` covering backend selection, dynamic name lambdas, metadata callbacks, and in-process Perfetto trace sessions. * Add bidirectional open-source Copybara transformations for `tsl` and `perfetto` dependencies and includes. * Declare `perfetto` Bzlmod dependency and `git_override` in `MODULE.bazel` pointing to `https://github.com/google/perfetto`. * Add `no_android_deps.patch` for Perfetto to remove Android SDK/NDK requirements and fix protoc_lib dependencies for open-source Protobuf 32.1 compatibility. * Expose `@tsl` repository from OpenXLA via `tsl_extension` in `MODULE.bazel` and map TSL dependencies to `@tsl//tsl` in `copy.bara.sky` for profiler targets. [skip-perf-gate] PiperOrigin-RevId: 976630278 --- MODULE.bazel | 17 + third_party/perfetto/no_android_deps.patch | 117 +++++++ tpu_sync/common/BUILD | 28 ++ tpu_sync/common/trace.cc | 211 +++++++++++++ tpu_sync/common/trace.h | 287 +++++++++++++++++ tpu_sync/common/trace_test.cc | 296 ++++++++++++++++++ tpu_sync/weight_sync/BUILD | 1 + .../weight_sync/weight_synchronizer_base.cc | 6 + 8 files changed, 963 insertions(+) create mode 100644 third_party/perfetto/no_android_deps.patch create mode 100644 tpu_sync/common/trace.cc create mode 100644 tpu_sync/common/trace.h create mode 100644 tpu_sync/common/trace_test.cc diff --git a/MODULE.bazel b/MODULE.bazel index 1e1547ffc..caa34b944 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -89,6 +89,12 @@ bazel_dep(name = "boringssl", version = "0.20240913.0") llvm = use_extension("@xla//third_party/extensions:llvm.bzl", "llvm_extension") use_repo(llvm, "llvm-project") +tsl_extension = use_extension("@xla//third_party/extensions:tsl.bzl", "tsl_extension") +use_repo( + tsl_extension, + "tsl", +) + jax_python_wheel_repository = use_repo_rule( "@jax//jaxlib:jax_python_wheel.bzl", "jax_python_wheel_repository", @@ -102,6 +108,7 @@ bazel_dep(name = "googletest", version = "1.17.0.bcr.2", repo_name = "com_google bazel_dep(name = "flatbuffers", version = "25.12.19", repo_name = "com_github_google_flatbuffers") bazel_dep(name = "highway", version = "1.2.0") bazel_dep(name = "prometheus-cpp", version = "1.3.0", repo_name = "com_github_jupp0r_prometheus_cpp") +bazel_dep(name = "perfetto") pybind11_internal_configure = use_extension( "@pybind11_bazel//:internal_configure.bzl", @@ -143,6 +150,16 @@ git_override( remote = "https://github.com/google/nanobind_abseil.git", ) +git_override( + module_name = "perfetto", + commit = "d0d8c6fa5f7411b59fccacf99233e79bcf94e3fc", + patch_strip = 1, + patches = [ + "//:third_party/perfetto/no_android_deps.patch", + ], + remote = "https://github.com/google/perfetto.git", +) + # Single-sources the wheel version from pyproject.toml [project].version and # appends WHEEL_VERSION_EXTRAS (e.g. ".dev"). Consumed by # //ci/wheel:raiden_wheel via @raiden_version//:version.bzl. diff --git a/third_party/perfetto/no_android_deps.patch b/third_party/perfetto/no_android_deps.patch new file mode 100644 index 000000000..5ef38a868 --- /dev/null +++ b/third_party/perfetto/no_android_deps.patch @@ -0,0 +1,117 @@ +diff --git a/MODULE.bazel b/MODULE.bazel +index 10941b7f..ea7db590 100644 +--- a/MODULE.bazel ++++ b/MODULE.bazel +@@ -31,50 +31,9 @@ bazel_dep(name = "platforms", version = "0.0.10") + bazel_dep(name = "protobuf", version = "31.1", repo_name = "com_google_protobuf") + bazel_dep(name = "re2", version = "2024-07-02.bcr.1") + bazel_dep(name = "rules_python", version = "1.0.0") +-bazel_dep(name = "rules_android", version = "0.6.6") + bazel_dep(name = "rules_cc", version = "0.2.19") + bazel_dep(name = "zstd", version = "1.5.7.bcr.1") + +-remote_android_extensions = use_extension( +- "@rules_android//bzlmod_extensions:android_extensions.bzl", +- "remote_android_tools_extensions", +-) +-use_repo(remote_android_extensions, "android_tools") +- +-android_sdk_repository_extension = use_extension( +- "@rules_android//rules/android_sdk_repository:rule.bzl", +- "android_sdk_repository_extension", +-) +- +-# When built using 'tools/bazel', 'ANDROID_HOME' environment variable points +-# to the hermetic Android SDK installation, that should be downloaded first +-# with 'tools/install-build-deps --android' +-android_sdk_repository_extension.configure( +- api_level = 35, +- build_tools_version = "35.0.1", +-) +-use_repo(android_sdk_repository_extension, "androidsdk") +- +-register_toolchains("@androidsdk//:sdk-toolchain", "@androidsdk//:all") +- +-bazel_dep(name = "rules_jvm_external", version = "6.9") +-bazel_dep(name = "rules_android_ndk", version = "0.1.3") +- +-android_ndk_repository_extension = use_extension( +- "@rules_android_ndk//:extension.bzl", +- "android_ndk_repository_extension", +-) +- +-# When built using 'tools/bazel', 'ANDROID_NDK_HOME' environment variable points +-# to the hermetic Android NDK installation, that should be downloaded first +-# with 'tools/install-build-deps --android' +-android_ndk_repository_extension.configure( +- api_level = 26, +-) +-use_repo(android_ndk_repository_extension, "androidndk") +- +-register_toolchains("@androidndk//:all") +- + # Perfetto configuration repository extension. + # This creates @perfetto_cfg which provides PERFETTO_CONFIG struct. + perfetto_cfg_ext = use_extension( +@@ -100,27 +59,3 @@ use_repo( + "perfetto_dep_zlib", + ) + +-# Maven dependencies for Android instrumentation tests. +-# Use a unique name to avoid conflicts with other modules (e.g., bazel_worker_java). +-maven = use_extension( +- "@rules_jvm_external//:extensions.bzl", +- "maven", +-) +-maven.install( +- name = "perfetto_maven", +- # Use rules_android's aar_import to avoid toolchain type mismatch. +- aar_import_bzl_label = "@rules_android//rules:rules.bzl", +- artifacts = [ +- "androidx.test:runner:1.6.2", +- "androidx.test:monitor:1.7.2", +- "com.google.truth:truth:1.4.4", +- "junit:junit:4.13.2", +- "androidx.test.ext:junit:1.2.1", +- "com.google.errorprone:error_prone_annotations:2.36.0", +- ], +- repositories = [ +- "https://maven.google.com", +- "https://repo1.maven.org/maven2", +- ], +-) +-use_repo(maven, "perfetto_maven") +diff --git a/bazel/rules.bzl b/bazel/rules.bzl +index 958e9bc0..f5990adf 100644 +--- a/bazel/rules.bzl ++++ b/bazel/rules.bzl +@@ -13,9 +13,16 @@ + # limitations under the License. + + load("@perfetto//bazel:proto_gen.bzl", "proto_descriptor_gen", "proto_gen") +-load("@perfetto//bazel:run_ait_with_adb.bzl", "android_instrumentation_test") + load("@perfetto_cfg//:perfetto_cfg.bzl", "PERFETTO_CONFIG") +-load("@rules_android//android:rules.bzl", "android_binary", "android_library") ++ ++def android_binary(**kwargs): ++ pass ++ ++def android_library(**kwargs): ++ pass ++ ++def android_instrumentation_test(**kwargs): ++ pass + + # +----------------------------------------------------------------------------+ + # | Base C++ rules. | +diff --git a/bazel/standalone/perfetto_cfg.bzl b/bazel/standalone/perfetto_cfg.bzl +--- a/bazel/standalone/perfetto_cfg.bzl ++++ b/bazel/standalone/perfetto_cfg.bzl +@@ -53,3 +53,6 @@ PERFETTO_CONFIG = struct( + protoc = ["@com_google_protobuf//:protoc"], +- protoc_lib = ["@com_google_protobuf//:protoc_lib"], ++ protoc_lib = [ ++ "@com_google_protobuf//src/google/protobuf/compiler:code_generator", ++ "@com_google_protobuf//src/google/protobuf/compiler:importer", ++ ], + protobuf_lite = ["@com_google_protobuf//:protobuf_lite"], diff --git a/tpu_sync/common/BUILD b/tpu_sync/common/BUILD index c8728e110..169d03342 100644 --- a/tpu_sync/common/BUILD +++ b/tpu_sync/common/BUILD @@ -48,6 +48,34 @@ cc_test( ], ) +cc_library( + name = "trace", + srcs = ["trace.cc"], + hdrs = ["trace.h"], + visibility = ["//visibility:public"], + deps = [ + "@com_google_absl//absl/base", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + "@perfetto//:libperfetto_client_experimental", + "@tsl//tsl/profiler/lib:traceme", + "@tsl//tsl/profiler/lib:traceme_encode", + ], +) + +cc_test( + name = "trace_test", + srcs = ["trace_test.cc"], + deps = [ + ":trace", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + "@perfetto//:libperfetto_client_experimental", + ], +) + nanobind_extension( name = "_raiden_id", srcs = ["raiden_id_binding.cc"], diff --git a/tpu_sync/common/trace.cc b/tpu_sync/common/trace.cc new file mode 100644 index 000000000..e7af73c82 --- /dev/null +++ b/tpu_sync/common/trace.cc @@ -0,0 +1,211 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/common/trace.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/base/call_once.h" +#include "absl/base/const_init.h" +#include "absl/base/thread_annotations.h" +#include "absl/strings/ascii.h" +#include "absl/strings/match.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "perfetto/tracing/backend_type.h" +#include "perfetto/tracing/core/data_source_config.h" // IWYU pragma: keep +#include "perfetto/tracing/core/trace_config.h" // IWYU pragma: keep +#include "perfetto/tracing/tracing.h" +#include "perfetto/tracing/track_event.h" + +PERFETTO_TRACK_EVENT_STATIC_STORAGE_IN_NAMESPACE(tpu_raiden); + +namespace tpu_raiden { +namespace { + +constexpr int8_t kNoOverride = -1; +std::atomic g_trace_backend_override{kNoOverride}; + +struct FileTraceSessionState { + std::unique_ptr<::perfetto::TracingSession> session; + std::string path; +}; + +absl::Mutex g_file_trace_mutex(absl::kConstInit); +std::unique_ptr g_file_trace_state + ABSL_GUARDED_BY(g_file_trace_mutex); + +bool IsEnvVarTruthy(const char* val) { + if (val == nullptr || val[0] == '\0') { + return false; + } + absl::string_view sv(val); + if (sv == "0" || absl::EqualsIgnoreCase(sv, "false") || + absl::EqualsIgnoreCase(sv, "off") || absl::EqualsIgnoreCase(sv, "no")) { + return false; + } + return true; +} + +bool StartPerfettoTraceToFileInternal(absl::string_view output_file_path) { + if (output_file_path.empty()) { + return false; + } + std::string file_path(output_file_path); + { + std::ofstream test_out(file_path, std::ios::out | std::ios::binary); + if (!test_out) { + return false; + } + } + + absl::MutexLock lock(g_file_trace_mutex); + if (g_file_trace_state != nullptr) { + return false; + } + + ::perfetto::TraceConfig cfg; + cfg.add_buffers()->set_size_kb(4096); + auto* ds_cfg = cfg.add_data_sources()->mutable_config(); + ds_cfg->set_name("track_event"); + + auto session = ::perfetto::Tracing::NewTrace(::perfetto::kInProcessBackend); + if (!session) { + return false; + } + + session->Setup(cfg); + session->StartBlocking(); + + auto state = std::make_unique(); + state->session = std::move(session); + state->path = std::string(output_file_path); + g_file_trace_state = std::move(state); + return true; +} + +} // namespace + +void EnsurePerfettoInitialized() { + static absl::once_flag once; + absl::call_once(once, []() { + if (!::perfetto::Tracing::IsInitialized()) { + ::perfetto::TracingInitArgs args; + args.backends = + ::perfetto::kSystemBackend | ::perfetto::kInProcessBackend; + ::perfetto::Tracing::Initialize(args); + } + ::tpu_raiden::TrackEvent::Register(); + + const char* output_file = std::getenv("TPU_RAIDEN_PERFETTO_OUTPUT_FILE"); + if (output_file == nullptr || output_file[0] == '\0') { + output_file = std::getenv("RAIDEN_PERFETTO_OUTPUT_FILE"); + } + if (output_file != nullptr && output_file[0] != '\0') { + if (StartPerfettoTraceToFileInternal(output_file)) { + std::atexit([]() { StopPerfettoTraceToFile(); }); + } + } + }); +} + +bool StartPerfettoTraceToFile(absl::string_view output_file_path) { + EnsurePerfettoInitialized(); + return StartPerfettoTraceToFileInternal(output_file_path); +} + +bool StopPerfettoTraceToFile() { + std::unique_ptr state; + { + absl::MutexLock lock(g_file_trace_mutex); + if (!g_file_trace_state) { + return false; + } + state = std::move(g_file_trace_state); + } + + ::tpu_raiden::TrackEvent::Flush(); + state->session->StopBlocking(); + std::vector trace_data = state->session->ReadTraceBlocking(); + + std::ofstream out(state->path, std::ios::out | std::ios::binary); + if (!out) { + return false; + } + if (!trace_data.empty()) { + out.write(trace_data.data(), + static_cast(trace_data.size())); + } + out.close(); + return out.good(); +} + +void SetTraceBackendForTesting(std::optional backend) { + if (backend.has_value()) { + g_trace_backend_override.store(static_cast(*backend), + std::memory_order_relaxed); + } else { + g_trace_backend_override.store(kNoOverride, std::memory_order_relaxed); + } +} + +TraceBackend GetTraceBackend() { + int8_t override_val = + g_trace_backend_override.load(std::memory_order_relaxed); + if (override_val != kNoOverride) { + return static_cast(override_val); + } + + const char* backend_env = std::getenv("TPU_RAIDEN_TRACE_BACKEND"); + if (backend_env == nullptr || backend_env[0] == '\0') { + backend_env = std::getenv("RAIDEN_TRACE_BACKEND"); + } + + if (backend_env != nullptr && backend_env[0] != '\0') { + std::string s = absl::AsciiStrToLower(backend_env); + if (s == "perfetto") { + return TraceBackend::kPerfetto; + } + if (s == "traceme") { + return TraceBackend::kTraceMe; + } + if (s == "both") { + return TraceBackend::kBoth; + } + if (s == "none" || s == "off" || s == "0") { + return TraceBackend::kNone; + } + } + + const char* use_perfetto_1 = std::getenv("TPU_RAIDEN_USE_PERFETTO"); + const char* use_perfetto_2 = std::getenv("RAIDEN_USE_PERFETTO"); + const char* use_perfetto_3 = std::getenv("TPU_RAIDEN_TRACE_PERFETTO"); + if (IsEnvVarTruthy(use_perfetto_1) || IsEnvVarTruthy(use_perfetto_2) || + IsEnvVarTruthy(use_perfetto_3)) { + return TraceBackend::kPerfetto; + } + + return TraceBackend::kTraceMe; +} + +} // namespace tpu_raiden diff --git a/tpu_sync/common/trace.h b/tpu_sync/common/trace.h new file mode 100644 index 000000000..733447cb8 --- /dev/null +++ b/tpu_sync/common/trace.h @@ -0,0 +1,287 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_COMMON_TRACE_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_COMMON_TRACE_H_ + +#include +#include // IWYU pragma: keep +#include // IWYU pragma: keep +#include +#include +#include +#include + +#include "absl/strings/str_cat.h" // IWYU pragma: keep +#include "absl/strings/string_view.h" +#include "perfetto/tracing/event_context.h" +#include "perfetto/tracing/string_helpers.h" +#include "perfetto/tracing/tracing.h" // IWYU pragma: keep +#include "perfetto/tracing/track_event.h" +#include "tsl/profiler/lib/traceme.h" +#include "tsl/profiler/lib/traceme_encode.h" + +// Defines Perfetto track event categories in the tpu_raiden namespace. +PERFETTO_DEFINE_CATEGORIES_IN_NAMESPACE( + tpu_raiden, ::perfetto::Category("tpu_raiden") + .SetDescription("Events from TPU Raiden")); + +namespace tpu_raiden { + +using ::tsl::profiler::TraceMe; +using ::tsl::profiler::TraceMeEncode; +using ::tsl::profiler::TraceMeLevel; + +// Supported tracing backends. +enum class TraceBackend : uint8_t { + kNone = 0, + kTraceMe = 1 << 0, + kPerfetto = 1 << 1, + kBoth = kTraceMe | kPerfetto, +}; + +constexpr TraceBackend operator|(TraceBackend a, TraceBackend b) { + return static_cast(static_cast(a) | + static_cast(b)); +} + +constexpr TraceBackend operator&(TraceBackend a, TraceBackend b) { + return static_cast(static_cast(a) & + static_cast(b)); +} + +constexpr bool HasBackend(TraceBackend config, TraceBackend target) { + return (static_cast(config) & static_cast(target)) != 0; +} + +// Ensures Perfetto client tracing is initialized. Thread-safe. +void EnsurePerfettoInitialized(); + +// Returns the active tracing backend, determined from environment variables +// (TPU_RAIDEN_TRACE_BACKEND / RAIDEN_TRACE_BACKEND, or TPU_RAIDEN_USE_PERFETTO) +// or an override set via SetTraceBackendForTesting. Thread-safe. +TraceBackend GetTraceBackend(); + +// Overrides the active tracing backend for unit tests. Pass std::nullopt to +// clear the override. Thread-safe. +void SetTraceBackendForTesting(std::optional backend); + +// Starts an in-process Perfetto tracing session and writes the collected +// trace data to output_file_path when StopPerfettoTraceToFile() is called. +// Returns true on success, false if a trace session is already running or +// if output_file_path cannot be created. Thread-safe. +bool StartPerfettoTraceToFile(absl::string_view output_file_path); + +// Stops the file-based Perfetto tracing session started by +// StartPerfettoTraceToFile(), flushes all buffered events, and writes the +// trace data to the configured file path. Returns true on success. Thread-safe. +bool StopPerfettoTraceToFile(); + +// Checks if verbose trace logging to stderr is enabled via either +// TPU_RAIDEN_TRACE or RAIDEN_TRACE environment variable. +// Thread-safe. +inline bool IsTraceLoggingEnabled() { + const char* tpu_raiden_trace = std::getenv("TPU_RAIDEN_TRACE"); + if (tpu_raiden_trace != nullptr && tpu_raiden_trace[0] != '\0' && + absl::string_view(tpu_raiden_trace) != "0") { + return true; + } + const char* raiden_trace = std::getenv("RAIDEN_TRACE"); + if (raiden_trace != nullptr && raiden_trace[0] != '\0' && + absl::string_view(raiden_trace) != "0") { + return true; + } + return false; +} + +// RAII scoped trace that emits trace events to TraceMe and/or Perfetto. +class ScopedTrace { + public: + template < + typename NameType, + std::enable_if_t, TraceBackend>, + bool> = true> + explicit ScopedTrace(NameType&& name, int level = 1) + : ScopedTrace(GetTraceBackend(), std::forward(name), level) {} + + template + ScopedTrace(TraceBackend backend, NameType&& name, int level = 1) { + Init(backend, std::forward(name), level); + } + + ~ScopedTrace() { + if (perfetto_active_) { + PERFETTO_USE_CATEGORIES_FROM_NAMESPACE_SCOPED(tpu_raiden); + TRACE_EVENT_END("tpu_raiden"); + } + } + + ScopedTrace(ScopedTrace&& other) noexcept + : traceme_(std::move(other.traceme_)), + perfetto_active_(other.perfetto_active_) { + other.perfetto_active_ = false; + } + + ScopedTrace& operator=(ScopedTrace&& other) noexcept { + if (this != &other) { + if (perfetto_active_) { + PERFETTO_USE_CATEGORIES_FROM_NAMESPACE_SCOPED(tpu_raiden); + TRACE_EVENT_END("tpu_raiden"); + } + traceme_ = std::move(other.traceme_); + perfetto_active_ = other.perfetto_active_; + other.perfetto_active_ = false; + } + return *this; + } + + ScopedTrace(const ScopedTrace&) = delete; + ScopedTrace& operator=(const ScopedTrace&) = delete; + + // Appends metadata to the trace scope. + template + void AppendMetadata(MetadataGenerator&& metadata_generator) { + if (traceme_.has_value()) { + traceme_->AppendMetadata( + std::forward(metadata_generator)); + } + if (perfetto_active_) { + PERFETTO_USE_CATEGORIES_FROM_NAMESPACE_SCOPED(tpu_raiden); + TRACE_EVENT_INSTANT( + "tpu_raiden", "metadata", [&](::perfetto::EventContext ctx) { + auto* annotation = ctx.event()->add_debug_annotations(); + annotation->set_name("metadata"); + if constexpr (std::is_invocable_v) { + std::string val = std::string(metadata_generator()); + annotation->set_string_value(val); + } + }); + } + } + + private: + template + void Init(TraceBackend backend, NameType&& name, int level) { + const bool use_traceme = HasBackend(backend, TraceBackend::kTraceMe); + const bool use_perfetto = HasBackend(backend, TraceBackend::kPerfetto); + + if (use_traceme) { + traceme_.emplace(std::forward(name), level); + } + + if (use_perfetto) { + EnsurePerfettoInitialized(); + PERFETTO_USE_CATEGORIES_FROM_NAMESPACE_SCOPED(tpu_raiden); + if (TRACE_EVENT_CATEGORY_ENABLED("tpu_raiden")) { + EmitPerfettoBegin(std::forward(name)); + perfetto_active_ = true; + } + } + } + + template + static void EmitPerfettoBegin(NameType&& name) { + PERFETTO_USE_CATEGORIES_FROM_NAMESPACE_SCOPED(tpu_raiden); + if constexpr (std::is_invocable_v) { + auto dynamic_name = name(); + TRACE_EVENT_BEGIN( + "tpu_raiden", + ::perfetto::DynamicString(dynamic_name.data(), dynamic_name.size())); + } else if constexpr (std::is_same_v, + ::perfetto::DynamicString> || + std::is_same_v, + ::perfetto::StaticString>) { + TRACE_EVENT_BEGIN("tpu_raiden", name); + } else if constexpr (std::is_array_v>) { + TRACE_EVENT_BEGIN("tpu_raiden", name); + } else if constexpr (std::is_convertible_v) { + absl::string_view sv = name; + TRACE_EVENT_BEGIN("tpu_raiden", + ::perfetto::DynamicString(sv.data(), sv.size())); + } else { + TRACE_EVENT_BEGIN("tpu_raiden", name); + } + } + + std::optional<::tsl::profiler::TraceMe> traceme_; + bool perfetto_active_ = false; +}; + +} // namespace tpu_raiden + +// Concatenation helpers for generating unique variable names per scope line. +#define RAIDEN_TRACE_CONCAT_INNER_(x, y) x##y +#define RAIDEN_TRACE_CONCAT_(x, y) RAIDEN_TRACE_CONCAT_INNER_(x, y) +#define RAIDEN_TRACE_UNIQUE_NAME_(base) RAIDEN_TRACE_CONCAT_(base, __LINE__) + +// Defines an RAII trace scope using the globally configured trace backend. +// Negligible overhead when profiling is inactive. +// +// Usage: +// RAIDEN_TRACE("MyActivity"); +// RAIDEN_TRACE("MyActivity", /*level=*/2); +// RAIDEN_TRACE([&]() { return absl::StrCat("DynamicOp_", id); }); +#define RAIDEN_TRACE(name, ...) \ + const ::tpu_raiden::ScopedTrace RAIDEN_TRACE_UNIQUE_NAME_( \ + _raiden_scoped_trace)(name, ##__VA_ARGS__) + +// Defines a trace scope for a code block. Alias for RAIDEN_TRACE. +// +// Usage: +// RAIDEN_TRACE_SCOPE("MyScope"); +#define RAIDEN_TRACE_SCOPE(name, ...) RAIDEN_TRACE(name, ##__VA_ARGS__) + +// Defines a scoped trace activity with dynamic metadata. +// The metadata callback is only executed if tracing is active. +// +// Usage: +// RAIDEN_TRACE_FN("MyActivity", [&]() { +// return ::tpu_raiden::TraceMeEncode({{"batch_size", 32}, {"step", 1}}); +// }); +#define RAIDEN_TRACE_FN(name, metadata_fn, ...) \ + ::tpu_raiden::ScopedTrace RAIDEN_TRACE_UNIQUE_NAME_(_raiden_scoped_trace)( \ + name, ##__VA_ARGS__); \ + RAIDEN_TRACE_UNIQUE_NAME_(_raiden_scoped_trace).AppendMetadata(metadata_fn) + +// Defines a scoped trace using Perfetto explicitly regardless of default +// backend. +// +// Usage: +// RAIDEN_PERFETTO_TRACE("MyActivity"); +#define RAIDEN_PERFETTO_TRACE(name, ...) \ + const ::tpu_raiden::ScopedTrace RAIDEN_TRACE_UNIQUE_NAME_( \ + _raiden_perfetto_trace)(::tpu_raiden::TraceBackend::kPerfetto, name, \ + ##__VA_ARGS__) + +// Defines a scoped trace using TraceMe explicitly regardless of default +// backend. +// +// Usage: +// RAIDEN_TRACEME_TRACE("MyActivity"); +#define RAIDEN_TRACEME_TRACE(name, ...) \ + const ::tpu_raiden::ScopedTrace RAIDEN_TRACE_UNIQUE_NAME_( \ + _raiden_traceme_trace)(::tpu_raiden::TraceBackend::kTraceMe, name, \ + ##__VA_ARGS__) + +// Logs a trace message directly to stderr when TPU_RAIDEN_TRACE or RAIDEN_TRACE +// is set in the environment. Thread-safe. +#define RAIDEN_TRACE_LOG(...) \ + do { \ + if (::tpu_raiden::IsTraceLoggingEnabled()) { \ + std::fprintf(stderr, "[RAIDEN_TRACE] %s\n", \ + absl::StrCat(__VA_ARGS__).c_str()); \ + } \ + } while (0) + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_COMMON_TRACE_H_ diff --git a/tpu_sync/common/trace_test.cc b/tpu_sync/common/trace_test.cc new file mode 100644 index 000000000..5eab5a8ae --- /dev/null +++ b/tpu_sync/common/trace_test.cc @@ -0,0 +1,296 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/common/trace.h" + +#include + +#include +#include +#include +#include +#include // IWYU pragma: keep +#include // IWYU pragma: keep + +#include +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/time/clock.h" +#include "absl/time/time.h" +#include "perfetto/tracing/backend_type.h" // IWYU pragma: keep +#include "perfetto/tracing/core/trace_config.h" // IWYU pragma: keep +#include "perfetto/tracing/tracing.h" // IWYU pragma: keep + +namespace tpu_raiden { +namespace { + +class TraceBackendTest : public ::testing::Test { + protected: + void SetUp() override { + SetTraceBackendForTesting(std::nullopt); + unsetenv("TPU_RAIDEN_TRACE_BACKEND"); + unsetenv("RAIDEN_TRACE_BACKEND"); + unsetenv("TPU_RAIDEN_USE_PERFETTO"); + unsetenv("RAIDEN_USE_PERFETTO"); + unsetenv("TPU_RAIDEN_TRACE_PERFETTO"); + } + + void TearDown() override { + SetTraceBackendForTesting(std::nullopt); + unsetenv("TPU_RAIDEN_TRACE_BACKEND"); + unsetenv("RAIDEN_TRACE_BACKEND"); + unsetenv("TPU_RAIDEN_USE_PERFETTO"); + unsetenv("RAIDEN_USE_PERFETTO"); + unsetenv("TPU_RAIDEN_TRACE_PERFETTO"); + } +}; + +TEST_F(TraceBackendTest, GetTraceBackendDefault) { + EXPECT_EQ(GetTraceBackend(), TraceBackend::kTraceMe); +} + +TEST_F(TraceBackendTest, SetTraceBackendForTestingOverride) { + SetTraceBackendForTesting(TraceBackend::kPerfetto); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + + SetTraceBackendForTesting(TraceBackend::kBoth); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kBoth); + + SetTraceBackendForTesting(TraceBackend::kNone); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kNone); + + SetTraceBackendForTesting(TraceBackend::kTraceMe); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kTraceMe); + + SetTraceBackendForTesting(std::nullopt); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kTraceMe); +} + +TEST_F(TraceBackendTest, BackendEnvVarsBackendName) { + setenv("TPU_RAIDEN_TRACE_BACKEND", "perfetto", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + + setenv("TPU_RAIDEN_TRACE_BACKEND", "both", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kBoth); + + setenv("TPU_RAIDEN_TRACE_BACKEND", "none", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kNone); + + setenv("TPU_RAIDEN_TRACE_BACKEND", "0", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kNone); + + setenv("TPU_RAIDEN_TRACE_BACKEND", "off", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kNone); + + setenv("TPU_RAIDEN_TRACE_BACKEND", "traceme", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kTraceMe); + unsetenv("TPU_RAIDEN_TRACE_BACKEND"); + + // Fallback to RAIDEN_TRACE_BACKEND + setenv("RAIDEN_TRACE_BACKEND", "perfetto", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + unsetenv("RAIDEN_TRACE_BACKEND"); +} + +TEST_F(TraceBackendTest, BackendEnvVarsUsePerfettoFlags) { + setenv("TPU_RAIDEN_USE_PERFETTO", "1", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + unsetenv("TPU_RAIDEN_USE_PERFETTO"); + + setenv("RAIDEN_USE_PERFETTO", "true", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + unsetenv("RAIDEN_USE_PERFETTO"); + + setenv("TPU_RAIDEN_TRACE_PERFETTO", "1", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kPerfetto); + unsetenv("TPU_RAIDEN_TRACE_PERFETTO"); + + // Disabled flags should not select Perfetto + setenv("TPU_RAIDEN_USE_PERFETTO", "0", 1); + EXPECT_EQ(GetTraceBackend(), TraceBackend::kTraceMe); + unsetenv("TPU_RAIDEN_USE_PERFETTO"); +} + +TEST(TraceTest, ScopedTraceMacroWithLiteral) { + RAIDEN_TRACE("TestLiteral"); + int val = 42; + EXPECT_EQ(val, 42); +} + +TEST(TraceTest, ScopedTraceMacroWithLevel) { + RAIDEN_TRACE("TestLevel", 2); + EXPECT_EQ(TraceMeLevel::kInfo, 2); +} + +TEST(TraceTest, ScopedTraceMacroWithLambda) { + int op_id = 1; + RAIDEN_TRACE([&]() { return absl::StrCat("Op_", op_id); }); + EXPECT_EQ(op_id, 1); +} + +TEST(TraceTest, ScopedTraceAliasMacro) { + RAIDEN_TRACE_SCOPE("TestScope"); + int val = 100; + EXPECT_EQ(val, 100); +} + +TEST(TraceTest, ScopedTraceFnWithMetadata) { + RAIDEN_TRACE_FN("TestFn", []() { return TraceMeEncode({{"k", "v"}}); }); + int val = 10; + EXPECT_EQ(val, 10); +} + +TEST(TraceTest, NestedScopes) { + RAIDEN_TRACE("OuterScope"); + int count = 1; + { + RAIDEN_TRACE("InnerScope1"); + count += 1; + { + RAIDEN_TRACE("InnerScope2"); + count += 1; + } + } + EXPECT_EQ(count, 3); +} + +TEST(TraceTest, ExplicitBackendMacros) { + RAIDEN_PERFETTO_TRACE("ExplicitPerfetto"); + RAIDEN_PERFETTO_TRACE("ExplicitPerfettoLevel", 2); + RAIDEN_TRACEME_TRACE("ExplicitTraceMe"); + RAIDEN_TRACEME_TRACE("ExplicitTraceMeLevel", 2); +} + +TEST(TraceTest, BothAndNoneBackends) { + SetTraceBackendForTesting(TraceBackend::kBoth); + { + RAIDEN_TRACE("BothTrace"); + RAIDEN_TRACE_FN("BothTraceFn", []() { return "key=val"; }); + } + + SetTraceBackendForTesting(TraceBackend::kNone); + { + RAIDEN_TRACE("NoneTrace"); + RAIDEN_TRACE_FN("NoneTraceFn", []() { return "key=val"; }); + } + + SetTraceBackendForTesting(std::nullopt); +} + +TEST(TraceTest, TraceLogMacro) { + // Disabled state by default + unsetenv("TPU_RAIDEN_TRACE"); + unsetenv("RAIDEN_TRACE"); + EXPECT_FALSE(IsTraceLoggingEnabled()); + RAIDEN_TRACE_LOG("This message is suppressed when trace is disabled: ", 123); + + // Enabled via TPU_RAIDEN_TRACE + setenv("TPU_RAIDEN_TRACE", "1", 1); + EXPECT_TRUE(IsTraceLoggingEnabled()); + RAIDEN_TRACE_LOG("Trace logging via TPU_RAIDEN_TRACE: ", "active", + ", count=", 42); + unsetenv("TPU_RAIDEN_TRACE"); + + // Enabled via RAIDEN_TRACE + setenv("RAIDEN_TRACE", "1", 1); + EXPECT_TRUE(IsTraceLoggingEnabled()); + RAIDEN_TRACE_LOG("Trace logging via RAIDEN_TRACE: ", "active"); + unsetenv("RAIDEN_TRACE"); + + // Explicitly disabled with "0" + setenv("TPU_RAIDEN_TRACE", "0", 1); + EXPECT_FALSE(IsTraceLoggingEnabled()); + unsetenv("TPU_RAIDEN_TRACE"); +} + +TEST(TraceTest, PerfettoInProcessTraceCollection) { + EnsurePerfettoInitialized(); + + // NOLINTNEXTLINE(misc-include-cleaner) + perfetto::TraceConfig cfg; + cfg.add_buffers()->set_size_kb(1024); + auto* ds_cfg = cfg.add_data_sources()->mutable_config(); + ds_cfg->set_name("track_event"); + + auto session = perfetto::Tracing::NewTrace(perfetto::kInProcessBackend); + session->Setup(cfg); + session->StartBlocking(); + + SetTraceBackendForTesting(TraceBackend::kPerfetto); + + { + RAIDEN_TRACE("PerfettoScopeEvent"); + RAIDEN_PERFETTO_TRACE("PerfettoExplicitEvent"); + RAIDEN_TRACE([&]() { return absl::StrCat("PerfettoDynamicOp_", 42); }); + RAIDEN_TRACE_FN("PerfettoFnEvent", + []() { return "debug_detail=perfetto_ok"; }); + } + + SetTraceBackendForTesting(std::nullopt); + + ::tpu_raiden::TrackEvent::Flush(); + session->StopBlocking(); + std::vector trace_data = session->ReadTraceBlocking(); + + EXPECT_FALSE(trace_data.empty()); + absl::string_view trace_view(trace_data.data(), trace_data.size()); + EXPECT_TRUE(absl::StrContains(trace_view, "PerfettoScopeEvent")); + EXPECT_TRUE(absl::StrContains(trace_view, "PerfettoExplicitEvent")); + EXPECT_TRUE(absl::StrContains(trace_view, "PerfettoDynamicOp_42")); +} + +TEST(TraceTest, StartAndStopPerfettoTraceToFile) { + std::string temp_dir = testing::TempDir(); + std::string trace_file = + absl::StrCat(temp_dir, "/raiden_trace_file_test_", + absl::ToUnixMicros(absl::Now()), ".pftrace"); + + // Ensure no prior session is running + StopPerfettoTraceToFile(); + + ASSERT_TRUE(StartPerfettoTraceToFile(trace_file)); + + // A second call should fail while active + EXPECT_FALSE(StartPerfettoTraceToFile(trace_file)); + + SetTraceBackendForTesting(TraceBackend::kPerfetto); + { + RAIDEN_TRACE("FileTraceScopedSlice"); + RAIDEN_PERFETTO_TRACE("FileTraceExplicitSlice"); + RAIDEN_TRACE([&]() { return absl::StrCat("FileTraceDynamicSlice_", 99); }); + } + SetTraceBackendForTesting(std::nullopt); + + ASSERT_TRUE(StopPerfettoTraceToFile()); + + // A second stop should return false + EXPECT_FALSE(StopPerfettoTraceToFile()); + + // Read back the file and verify non-empty content containing slice names + std::ifstream in(trace_file, std::ios::in | std::ios::binary); + ASSERT_TRUE(in.is_open()); + std::string file_content((std::istreambuf_iterator(in)), + std::istreambuf_iterator()); + in.close(); + + EXPECT_FALSE(file_content.empty()); + absl::string_view view(file_content); + EXPECT_TRUE(absl::StrContains(view, "FileTraceScopedSlice")); + EXPECT_TRUE(absl::StrContains(view, "FileTraceExplicitSlice")); + EXPECT_TRUE(absl::StrContains(view, "FileTraceDynamicSlice_99")); +} + +} // namespace +} // namespace tpu_raiden diff --git a/tpu_sync/weight_sync/BUILD b/tpu_sync/weight_sync/BUILD index 2ee3c0a50..de0ac0eee 100644 --- a/tpu_sync/weight_sync/BUILD +++ b/tpu_sync/weight_sync/BUILD @@ -56,6 +56,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":tiling_utils", + "//tpu_sync/common:trace", "//tpu_sync/core:host_memory_allocator", "//tpu_sync/core:numa_thread_pool", "//tpu_sync/core:raiden_manager_base", diff --git a/tpu_sync/weight_sync/weight_synchronizer_base.cc b/tpu_sync/weight_sync/weight_synchronizer_base.cc index 3b07631f9..4321cfc17 100644 --- a/tpu_sync/weight_sync/weight_synchronizer_base.cc +++ b/tpu_sync/weight_sync/weight_synchronizer_base.cc @@ -48,6 +48,7 @@ #include "xla/shape_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" +#include "tpu_sync/common/trace.h" #include "tpu_sync/core/host_memory_allocator.h" #include "tpu_sync/core/numa_thread_pool.h" #include "tpu_sync/core/raiden_manager_base.h" @@ -460,6 +461,7 @@ absl::StatusOr WeightSynchronizerBase::H2dLayer( absl::StatusOr WeightSynchronizerBase::H2d( uint64_t uuid) { + RAIDEN_TRACE("WeightSynchronizerBase::H2d"); if (buffer_holds_.empty()) { return raiden::PjRtCopyFuture(std::vector{}); } @@ -581,6 +583,7 @@ absl::StatusOr WeightSynchronizerBase::D2hLayer( absl::StatusOr WeightSynchronizerBase::D2h( uint64_t uuid) { + RAIDEN_TRACE("WeightSynchronizerBase::D2h"); if (buffer_holds_.empty()) { return raiden::PjRtCopyFuture(std::vector{}); } @@ -604,6 +607,7 @@ absl::StatusOr WeightSynchronizerBase::D2h( absl::Status WeightSynchronizerBase::PushWeights( const std::vector& peers) { + RAIDEN_TRACE("WeightSynchronizerBase::PushWeights"); if (control_delegate_ != nullptr) { return control_delegate_->PushWeights(peers); } @@ -629,6 +633,7 @@ absl::Status WeightSynchronizerBase::PushWeightsLocal( absl::Status WeightSynchronizerBase::PushWeightsResharded( const tpu_sync::rpc::StartTransferRequest& request) { + RAIDEN_TRACE("WeightSynchronizerBase::PushWeightsResharded"); if (control_delegate_ != nullptr) { return control_delegate_->PushWeightsResharded(request); } @@ -867,6 +872,7 @@ absl::Status WeightSynchronizerBase::PushWeightsReshardedLocal( absl::Status WeightSynchronizerBase::BindWeights( const std::vector>& layer_buffers) { + RAIDEN_TRACE("WeightSynchronizerBase::BindWeights"); if (layer_buffers.size() != num_layers_) { return absl::InvalidArgumentError("Number of layers mismatch"); }