syntax = "proto3";
package tensorflow;

import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/protobuf/config.proto";
import "tensorflow/contrib/tpu/profiler/op_profile.proto";

// The TPUProfiler service retrieves performance information about
// the programs running on connected TPUs over a period of time.
service TPUProfiler {
  // Starts a profiling session, blocks until it completes, and returns data.
  rpc Profile(ProfileRequest) returns (ProfileResponse) {
  }
}

message ProfileRequest {
  // In future, the caller will be able to customize when profiling starts and
  // stops. For now, it collects `duration_ms` milliseconds worth of data.
  uint64 duration_ms = 1;

  // The maximum number of events to return. By default (value 0), return all
  // events.
  uint64 max_events = 2;

  // required profiling tools name such as "input_pipeline_analyzer" etc
  repeated string tools = 3;

  // In future, the caller will indicate which TF session is being profiled, and
  // only data relating to that program will be returned. For now, we assume
  // all activity during the profiling period is relevant.
  // next-field: 4
}

message ProfileToolData {
  // The tool's name which this data is associated. (e.g. "input_pipeline".)
  string name = 1;

  // The data payload (likely json) for the specific tool.
  bytes data = 2;
}

message ProfileResponse {
  uint64 xprof_response_size = 1;  // Placeholder: return something meaningful.
  // Graphs of programs executed on TPUs during the profiling period.
  repeated GraphDef computation_graph = 2;

  // Performance profile that can be used to annotate HLO operations in the
  // computation graph.
  RunMetadata hlo_metadata = 5;

  // Encoded Trace proto message that contains metadata about the trace captured
  // during the profiling period. Describes the devices and resources that
  // 'trace_events' refers to.
  bytes encoded_trace = 3;

  // Assembles a hierarchical performance profile based on HLOs in trace events.
  // If the trace covers multiple programs, the longest-running one is analyzed.
  // See op_profile.proto for the detailed semantics of the returned profile.
  tpu.op_profile.Profile op_profile = 4;

  // Data payload for each required tools.
  repeated ProfileToolData tool_data = 6;
  // next-field: 7
}
