diff --git a/tensorflow/core/platform/default/device_tracer.cc b/tensorflow/core/platform/default/device_tracer.cc index 6761472a043..0836ef5b5e5 100644 --- a/tensorflow/core/platform/default/device_tracer.cc +++ b/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 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; }(); diff --git a/tensorflow/core/profiler/internal/cpu/BUILD b/tensorflow/core/profiler/internal/cpu/BUILD index 9a326c2ebba..a07e51fe003 100644 --- a/tensorflow/core/profiler/internal/cpu/BUILD +++ b/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, ) diff --git a/tensorflow/core/profiler/internal/cpu/host_tracer.cc b/tensorflow/core/profiler/internal/cpu/host_tracer.cc index 039b5a06570..5fde361fd2c 100644 --- a/tensorflow/core/profiler/internal/cpu/host_tracer.cc +++ b/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 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; }(); diff --git a/tensorflow/core/profiler/internal/runtime/BUILD b/tensorflow/core/profiler/internal/runtime/BUILD index aa670f88b4f..085fed81578 100644 --- a/tensorflow/core/profiler/internal/runtime/BUILD +++ b/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, ) diff --git a/tensorflow/core/profiler/internal/runtime/eager_profiler.cc b/tensorflow/core/profiler/internal/runtime/eager_profiler.cc index 8f09438671f..30182da5db1 100644 --- a/tensorflow/core/profiler/internal/runtime/eager_profiler.cc +++ b/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 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; }();