Allow turning off individual profiler via environment variable.

PiperOrigin-RevId: 245486114

tensorflow/core/platform/default/device_tracer.cc

@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/step_stats_collector.h"
 #include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/profiler/internal/profiler_interface.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace tensorflow {
 namespace {
@@ -715,7 +717,11 @@ std::unique_ptr<profiler::ProfilerInterface> CreateDeviceTracer(
 }
 
 auto register_device_tracer_factory = [] {
-  RegisterProfilerFactory(&CreateDeviceTracer);
+  bool enable;
+  TF_CHECK_OK(ReadBoolFromEnvVar("TF_ENABLE_OSS_GPU_PROFILER", true, &enable));
+  if (enable) {
+    RegisterProfilerFactory(&CreateDeviceTracer);
+  }
   return 0;
 }();
 

tensorflow/core/profiler/internal/cpu/BUILD

@@ -16,14 +16,14 @@ tf_cuda_library(
         "host_tracer.cc",
     ],
     deps = [
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/strings",
         "//tensorflow/core:core_cpu_lib",
         "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/profiler/internal:profiler_interface",
         "//tensorflow/core/profiler/internal:traceme_recorder",
-        # "//tensorflow/stream_executor/lib",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
     ],
     alwayslink = True,
 )

tensorflow/core/profiler/internal/cpu/host_tracer.cc

@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/profiler/internal/profiler_interface.h"
 #include "tensorflow/core/profiler/internal/traceme_recorder.h"
 #include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -146,7 +147,12 @@ std::unique_ptr<ProfilerInterface> CreateHostTracer(const ProfilerContext*) {
 }
 
 auto register_host_tracer_factory = [] {
-  RegisterProfilerFactory(&CreateHostTracer);
+  bool enable;
+
+  TF_CHECK_OK(ReadBoolFromEnvVar("TF_ENABLE_OSS_CPU_PROFILER", true, &enable));
+  if (enable) {
+    RegisterProfilerFactory(&CreateHostTracer);
+  }
   return 0;
 }();
 

tensorflow/core/profiler/internal/runtime/BUILD

@@ -15,9 +15,10 @@ tf_cuda_library(
         "eager_profiler.cc",
     ],
     deps = [
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core/common_runtime/eager:context",
         "//tensorflow/core/profiler/internal:profiler_interface",
-        # "//tensorflow/stream_executor/lib",
     ],
     alwayslink = True,
 )

tensorflow/core/profiler/internal/runtime/eager_profiler.cc

@@ -13,7 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/profiler/internal/profiler_interface.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace tensorflow {
 namespace profiler {
@@ -89,7 +91,12 @@ std::unique_ptr<ProfilerInterface> CreateEagerProfiler(
 }
 
 auto register_eager_profiler_factory = [] {
-  RegisterProfilerFactory(&CreateEagerProfiler);
+  bool enable;
+  TF_CHECK_OK(
+      ReadBoolFromEnvVar("TF_ENABLE_EAGER_RUNTIME_PROFILER", true, &enable));
+  if (enable) {
+    RegisterProfilerFactory(&CreateEagerProfiler);
+  }
   return 0;
 }();