syntax = "proto3";

import "google/protobuf/empty.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/framework/tensor.proto";

package proto;
option go_package = "elasticdl/go/pkg/proto";

enum TaskType {
  TRAINING = 0;
  EVALUATION = 1;
  PREDICTION = 2;
  WAIT = 3;
  TRAIN_END_CALLBACK = 4;
}

// A task is a unit of work for ElasticDL training workers, assigned by master.
// Worker divides a task into multiple minibatches and compute a gradient for
// each minibatch. For now, only RecordIO file format is supported.
message Task {
  // Unique id assigned by master.
  int32 task_id = 1;

  int32 minibatch_size = 2;

  // Name for the shard. If RecordIO file format is used, this should be the
  // filename for a RecordIO shard. An empty shard name signifies that the
  // master has no pending tasks to assign to the requesting worker.
  string shard_name = 3;

  // Starting and ending (non-inclusive) record number.
  int64 start = 4;
  int64 end = 5;

  // Current model version on master
  int32 model_version = 6;

  // Whether this is training or evaluation task.
  TaskType type = 7;

  // Extended task configuration as a list of key-value pair.
  // The task message is mainly targeted to data and model.
  // What's more, we also leverage task mechanism to complete
  // some work such as saving the model. To keep the message
  // fields clean, we put the additional configurations of these
  // task type into this field.
  // For example:
  // SAVE_MODEL task: saved_model_path
  map<string, string> extended_config = 8;
}

message IndexedSlicesProto {
  tensorflow.TensorProto concat_tensors = 1;
  repeated int64 ids = 2;
}

message EmbeddingTableInfo {
  string name = 1;
  int64 dim = 2;
  string initializer = 3;
  tensorflow.DataType dtype = 4;
}

message Model {
  int32 version = 1;
  repeated EmbeddingTableInfo embedding_table_infos = 2;
  map<string, tensorflow.TensorProto> dense_parameters = 3;
  map<string, IndexedSlicesProto> embedding_tables = 4;
}

message GetTaskRequest {
  int32 worker_id = 1;
  TaskType task_type = 2;
}

message ReportTaskResultRequest {
  // Task id assigned by master.
  int32 task_id = 1;

  // When error occurred, err_message contains error message in plain text.
  string err_message = 2;
  // statistics of the task being executed.
  map<string, int32> exec_counters = 3;
}

message ReportEvaluationMetricsRequest {
  map<string, tensorflow.TensorProto> model_outputs = 1;
  tensorflow.TensorProto labels = 2;
  int32 worker_id = 3;
}

message ReportVersionRequest {
  int32 model_version = 1;
}

service Master {
  rpc get_task(GetTaskRequest) returns (Task);
  rpc report_evaluation_metrics(ReportEvaluationMetricsRequest)
      returns (google.protobuf.Empty);
  rpc report_task_result(ReportTaskResultRequest)
      returns (google.protobuf.Empty);
  rpc report_version(ReportVersionRequest) returns (google.protobuf.Empty);
}

message PullEmbeddingVectorRequest {
  string name = 1;
  repeated int64 ids = 2;
}

message PullDenseParametersRequest {
  int32 version = 1;
}

message PullDenseParametersResponse {
  bool initialized = 1;
  int32 version = 2;
  map<string, tensorflow.TensorProto> dense_parameters = 3;
}

message PullEmbeddingVectorsRequest {
  string name = 1;
  repeated int64 ids = 2;
}

message PushGradientsRequest {
  Model gradients = 1;
  float learning_rate = 2;
}

message PushGradientsResponse {
  bool accepted = 1;
  int32 version = 2;
}

// PS service
service Pserver {
  rpc push_model(Model) returns (google.protobuf.Empty);
  rpc push_embedding_table_infos(Model) returns (google.protobuf.Empty);
  rpc pull_dense_parameters(PullDenseParametersRequest)
      returns (PullDenseParametersResponse);
  rpc pull_embedding_vectors(PullEmbeddingVectorsRequest)
      returns (tensorflow.TensorProto);
  rpc push_gradients(PushGradientsRequest) returns (PushGradientsResponse);
}
