experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allow turning off individual profiler via environment variable.

Browse Source 
PiperOrigin-RevId: 245486114
This commit is contained in:
Xiao Yu 2019-04-26 14:18:03 -07:00 committed by TensorFlower Gardener
parent a2b7b5d5a0
commit 8cd28ec6d0
5 changed files with 27 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/core/platform/default/device_tracer.cc
View File

@ -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;
}();

6
tensorflow/core/profiler/internal/cpu/BUILD
View File

@ -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,
)

8
tensorflow/core/profiler/internal/cpu/host_tracer.cc
View File

@ -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;
}();

3
tensorflow/core/profiler/internal/runtime/BUILD
View File

@ -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,
)

9
tensorflow/core/profiler/internal/runtime/eager_profiler.cc
View File

@ -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;
}();