allow rpc path to override ProfilerOptions.
To allow backward compatible with old client, a version control field is added.
PiperOrigin-RevId: 303618399
Change-Id: If305c6923f93daaa67633e63e6ac2ede2d9d653f
diff --git a/tensorflow/core/profiler/internal/profiler_interface.h b/tensorflow/core/profiler/internal/profiler_interface.h
index 9e7819d..74aa565 100644
--- a/tensorflow/core/profiler/internal/profiler_interface.h
+++ b/tensorflow/core/profiler/internal/profiler_interface.h
@@ -36,8 +36,22 @@
// DeviceType::kTpu: only CPU/TPU will be profiled.
DeviceType device_type = DeviceType::kUnspecified;
- // Inexpensive ops are not traced by default.
- int host_tracer_level = 2;
+ // Levels of host tracing:
+ // - Level 0 is used to disable host traces.
+ // - Level 1 enables tracing of only user instrumented (or default) TraceMe.
+ // - Level 2 enables tracing of all level 1 TraceMe(s) and instrumented high
+ // level program execution details (expensive TF ops, XLA ops, etc).
+ // This is the default.
+ // - Level 3 enables tracing of all level 2 TraceMe(s) and more verbose
+ // (low-level) program execution details (cheap TF ops, etc).
+ uint32 host_tracer_level = 2;
+
+ // Levels of device tracing:
+ // - Level 0 is used to disable device traces.
+ // - Level 1 is used to enable device traces.
+ // - More levels might be defined for specific device for controlling the
+ // verbosity of the trace.
+ uint32 device_tracer_level = 1;
// Whether to enable python function calls tracer.
bool enable_python_tracer = false;
diff --git a/tensorflow/core/profiler/profiler_service.proto b/tensorflow/core/profiler/profiler_service.proto
index ab7ae49..007b68e 100644
--- a/tensorflow/core/profiler/profiler_service.proto
+++ b/tensorflow/core/profiler/profiler_service.proto
@@ -14,11 +14,37 @@
}
message ProfileOptions {
+ // Some default value of option are not proto3 default value. Use this version
+ // to determine if we should use default option value instead of proto3
+ // default value.
+ uint32 version = 5;
+
// We don't collect the dataset ops by default for better trace-viewer
// scalability. The caller can mannually set this field to include the ops.
bool include_dataset_ops = 1;
- // next-field: 2
+ // Levels of host tracing: (version >= 1)
+ // - Level 0 is used to disable host traces.
+ // - Level 1 enables tracing of only user instrumented (or default) TraceMe.
+ // - Level 2 enables tracing of all level 1 TraceMe(s) and instrumented high
+ // level program execution details (expensive TF ops, XLA ops, etc).
+ // This is the default.
+ // - Level 3 enables tracing of all level 2 TraceMe(s) and more verbose
+ // (low-level) program execution details (cheap TF ops, etc).
+ uint32 host_tracer_level = 2;
+
+ // Levels of device tracing: (version >= 1)
+ // - Level 0 is used to disable device traces.
+ // - Level 1 is used to enable device traces.
+ // - More levels might be defined for specific device for controlling the
+ // verbosity of the trace.
+ uint32 device_tracer_level = 3;
+
+ // Whether enable python function calls tracing. Runtime overhead ensues if
+ // enabled. Default off. (version >= 1)
+ uint32 python_tracer_level = 4;
+
+ // next-field: 6
}
message ToolRequestOptions {
diff --git a/tensorflow/core/profiler/rpc/profiler_service_impl.cc b/tensorflow/core/profiler/rpc/profiler_service_impl.cc
index 01cd35a..407cd0a 100644
--- a/tensorflow/core/profiler/rpc/profiler_service_impl.cc
+++ b/tensorflow/core/profiler/rpc/profiler_service_impl.cc
@@ -24,6 +24,7 @@
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/env_time.h"
#include "tensorflow/core/profiler/convert/xplane_to_profile_response.h"
+#include "tensorflow/core/profiler/internal/profiler_interface.h"
#include "tensorflow/core/profiler/lib/profiler_session.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
#include "tensorflow/core/util/ptr_util.h"
@@ -51,7 +52,8 @@
::grpc::Status Profile(::grpc::ServerContext* ctx, const ProfileRequest* req,
ProfileResponse* response) override {
VLOG(1) << "Received a profile request: " << req->DebugString();
- std::unique_ptr<ProfilerSession> profiler = ProfilerSession::Create();
+ std::unique_ptr<ProfilerSession> profiler =
+ ProfilerSession::Create(GetOptions(req->opts()));
Status status = profiler->Status();
if (!status.ok()) {
return ::grpc::Status(::grpc::StatusCode::INTERNAL,
@@ -74,6 +76,19 @@
return ::grpc::Status::OK;
}
+
+ private:
+ profiler::ProfilerOptions GetOptions(const tensorflow::ProfileOptions& opts) {
+ profiler::ProfilerOptions options;
+ if (opts.version()) {
+ options.host_tracer_level = opts.host_tracer_level();
+ options.device_tracer_level = opts.device_tracer_level();
+ options.enable_python_tracer = opts.python_tracer_level() > 0;
+ } else {
+ // use default options value;
+ }
+ return options;
+ }
};
} // namespace
