Allow locking the profiler via environment variable

Alternative to https://github.com/tensorflow/tensorflow/pull/46544

Use case: Allow an outer profiler (e.g., a command-line tool) to set an environment variable to prevent interference from TF Profiler. We already have mutual exclusion of profiler clients via ProfilerLock to guarantee at most one profiler client can actually profile. The environment variable TF_DISABLE_PROFILING, when set to 1, disables all profiler clients. We assume that clients already know how to gracefully handle failure to profile.

PiperOrigin-RevId: 357003057
Change-Id: I79ceccd289114d6eac3b9b7b4b14d0c5f1950c4d

tensorflow/core/profiler/lib/BUILD

@@ -227,6 +227,10 @@ cc_library(
     hdrs = ["profiler_lock.h"],
     copts = tf_profiler_copts(),
     visibility = ["//tensorflow/core/profiler:internal"],
+    deps = [
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+    ],
 )
 
 filegroup(

tensorflow/core/profiler/lib/profiler_lock.cc

@@ -16,6 +16,9 @@ limitations under the License.
 
 #include <atomic>
 
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/util/env_var.h"
+
 namespace tensorflow {
 namespace profiler {
 
@@ -23,7 +26,22 @@ namespace profiler {
 // Prevents another profiler session from creating ProfilerInterface(s).
 std::atomic<bool> session_active = ATOMIC_VAR_INIT(false);
 
-bool AcquireProfilerLock() { return !session_active.exchange(true); }
+bool AcquireProfilerLock() {
+  // Use environment variable to permanently lock the profiler.
+  // This allows running TensorFlow under an external profiling tool with all
+  // built-in profiling disabled.
+  static bool tf_profiler_disabled = [] {
+    bool disabled = false;
+    ReadBoolFromEnvVar("TF_DISABLE_PROFILING", false, &disabled).IgnoreError();
+    return disabled;
+  }();
+  if (TF_PREDICT_FALSE(tf_profiler_disabled)) {
+    LOG(WARNING) << "TensorFlow Profiler is permanently disabled by env var "
+                    "TF_DISABLE_PROFILING.";
+    return false;
+  }
+  return !session_active.exchange(true);
+}
 
 void ReleaseProfilerLock() { session_active.store(false); }