syntax = "proto3";

import "google/protobuf/empty.proto";
import "elasticdl/proto/tensor_dtype.proto";

package master;

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

enum MethodType {
    MINIMUM = 0;
    FIXED = 1;
}

// 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;
}

message Tensor {
    // Dimensions of the tensor. The first entry in "dim" is the outermost
    // dimension used to layout the values, the last entry is the innermost
    // dimension.
    repeated int32 dim = 1;

    // ndarray's buffer dump. Each element must be a 32 bit float value.
    bytes content = 2;

    // Indices will be tf.IndexedSlices.indices if the tensor is in the form
    // of tf.IndexedSlices. Ohterwise indices will be None.
    repeated int32 indices = 3;

    // Dtype of the tensor, e.g. "int64", "float32".
    TensorDtype dtype = 4;
}

message Model {
    int32 version = 1;
    map<string, Tensor> param = 2;
}

message GetTaskRequest {
    int32 worker_id = 1;
}

message GetModelRequest {
    MethodType method = 1;
    int32 version = 2;
}

message ReportVariableRequest {
    map<string, Tensor> variable = 1;
}

message ReportGradientRequest {
    // Used by worker to identify the report. Not used by master. It will be
    // copied directly into ReportGradientResponse
    int32 gradient_id = 1;
    
    // Model version used to compute gradients.
    int32 model_version = 2;
    map<string, Tensor> gradient = 3;
}

message ReportGradientResponse {
    // If the gradient is accepted.
    bool accepted = 1;
    // Current model version on master.
    int32 model_version = 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;
}

message ReportEvaluationMetricsResponse {
    // If the evaluation metric is accepted.
    bool accepted = 1;
    // Current model version on master.
    int32 model_version = 2;
}

message ReportEvaluationMetricsRequest {
    // Model version used to compute evaluation metrics.
    int32 model_version = 1;
    map<string, Tensor> evaluation_metrics = 2;
}

service Master {
    rpc GetTask(GetTaskRequest) returns (Task);
    rpc GetModel(GetModelRequest) returns (Model);
    rpc ReportVariable(ReportVariableRequest) returns (google.protobuf.Empty);
    rpc ReportGradient(ReportGradientRequest) returns (ReportGradientResponse);
    rpc ReportEvaluationMetrics(ReportEvaluationMetricsRequest) returns (ReportEvaluationMetricsResponse);
    rpc ReportTaskResult(ReportTaskResultRequest) returns (google.protobuf.Empty);
}
