syntax = "proto3";

package tensorflow.data;

import "tensorflow/core/data/service/common.proto";
import "tensorflow/core/framework/tensor.proto";

message TaskProgress {
  // The task that this message is about.
  int64 task_id = 1;
  // Whether the task has completed.
  bool completed = 2;
}

message WorkerHeartbeatRequest {
  string worker_address = 1;
  repeated int64 current_tasks = 2;
}

message WorkerHeartbeatResponse {
  repeated TaskDef new_tasks = 1;
  repeated int64 tasks_to_delete = 2;
}

message WorkerUpdateRequest {
  string worker_address = 1;
  repeated TaskProgress updates = 2;
}

message WorkerUpdateResponse {}

message GetDatasetDefRequest {
  int64 dataset_id = 1;
}

message GetDatasetDefResponse {
  DatasetDef dataset_def = 1;
}

message GetSplitRequest {
  int64 job_id = 1;
  int64 repetition = 2;
}

message GetSplitResponse {
  TensorProto split = 1;
  bool end_of_splits = 2;
}

message GetOrRegisterDatasetRequest {
  // The dataset to register.
  DatasetDef dataset = 1;
}

message GetOrRegisterDatasetResponse {
  // The id for the registered dataset.
  int64 dataset_id = 1;
}

message CreateJobRequest {
  // The id of the dataset to create a job for.
  int64 dataset_id = 1;
  // A mode controlling how the tf.data service produces data for the job.
  ProcessingModeDef processing_mode = 2;
}

message CreateJobResponse {
  // An id for the client that will read from the job. When the client is done
  // with the job, they should call ReleaseJobClient with this id.
  int64 job_client_id = 1;
}

message GetOrCreateJobRequest {
  // The id of the dataset to create a job for.
  int64 dataset_id = 1;
  // A mode controlling how the tf.data service produces data for the job.
  ProcessingModeDef processing_mode = 2;
  // A name for the job.
  string job_name = 3;
  // An index for the job. Multiple jobs can be created for the same name, if
  // they have different indices.
  int64 job_name_index = 4;
}

message GetOrCreateJobResponse {
  // An id for the client that will read from the job. When the client is done
  // with the job, they should call ReleaseJobClient with this id.
  int64 job_client_id = 1;
}

message ReleaseJobClientRequest {
  int64 job_client_id = 1;
}

message ReleaseJobClientResponse {}

message GetTasksRequest {
  // The job client id to look up tasks for.
  int64 job_client_id = 1;
}

message GetTasksResponse {
  // A list of all tasks for a job.
  repeated TaskInfo task_info = 1;
  // Whether the job has finished. An empty `task_info` list could either mean
  // that no tasks have been started yet, or that all tasks have finished. This
  // field gives us a way to tell the difference.
  bool job_finished = 2;
}

message WorkerInfo {
  string address = 1;
  int64 id = 2;
}

message GetWorkersRequest {}

message GetWorkersResponse {
  // A list of all workers.
  repeated WorkerInfo workers = 1;
}

service DispatcherService {
  // Performs a periodic worker heartbeat.
  rpc WorkerHeartbeat(WorkerHeartbeatRequest) returns (WorkerHeartbeatResponse);

  // Updates the dispatcher with information about the worker's state.
  rpc WorkerUpdate(WorkerUpdateRequest) returns (WorkerUpdateResponse);

  // Gets a dataset defintion.
  rpc GetDatasetDef(GetDatasetDefRequest) returns (GetDatasetDefResponse);

  // Gets the next split for a given job.
  rpc GetSplit(GetSplitRequest) returns (GetSplitResponse);

  // Registers a dataset with the server, or returns its id if it is already
  // registered.
  //
  // The dataset is constructed in a new graph, so it must not refer to
  // external resources or variables.
  rpc GetOrRegisterDataset(GetOrRegisterDatasetRequest)
      returns (GetOrRegisterDatasetResponse);

  // Gets a job if it already exists, otherwise creates it.
  rpc GetOrCreateJob(GetOrCreateJobRequest) returns (GetOrCreateJobResponse);

  // Creates a job for reading from the tf.data service.
  rpc CreateJob(CreateJobRequest) returns (CreateJobResponse);

  // Releases a job client so that a job may eventually be cleaned up.
  rpc ReleaseJobClient(ReleaseJobClientRequest)
      returns (ReleaseJobClientResponse);

  // Reports a list of all tasks for a job.
  rpc GetTasks(GetTasksRequest) returns (GetTasksResponse);

  // Reports a list of all workers registered with the dispatcher.
  rpc GetWorkers(GetWorkersRequest) returns (GetWorkersResponse);
}
