syntax = "proto3";

package tensorflow.data;

import "tensorflow/core/framework/graph.proto";

// Next tag: 2
message DatasetDef {
  // We represent datasets as tensorflow GraphDefs which define the operations
  // needed to create a tf.data dataset.
  GraphDef graph = 1;
}

// Next tag: 10
message TaskDef {
  // The dataset to iterate over.
  oneof dataset {
    DatasetDef dataset_def = 1;
    string path = 2;
  }
  int64 dataset_id = 3;
  int64 task_id = 4;
  int64 job_id = 5;
  // In distributed epoch processing mode, we use one split provider for each
  // source that feeds into the dataset. In parallel_epochs mode,
  // `num_split_providers` is always zero.
  int64 num_split_providers = 9;
  // Address of the worker that the task is assigned to.
  string worker_address = 8;
  ProcessingModeDef processing_mode = 6;
  // Optional number of consumers. If set, the results of the task will be
  // provided to consumers round-robin.
  oneof optional_num_consumers {
    int64 num_consumers = 7;
  }
}

// Next tag: 6
message TaskInfo {
  // The address of the worker processing the task.
  string worker_address = 1;
  // The transfer address of the worker processing the task.
  string transfer_address = 4;
  // The task id.
  int64 task_id = 2;
  // The id of the job that the task is part of.
  int64 job_id = 3;
  // The round to start reading from the task in. For non-round-robin reads,
  // this is always 0.
  int64 starting_round = 5;
}

// Next tag: 3
enum ProcessingModeDef {
  INVALID = 0;
  // Each tf.data worker processes an entire epoch.
  PARALLEL_EPOCHS = 1;
  // Processing of an epoch is distributed across all tf.data workers.
  DISTRIBUTED_EPOCH = 2;
}
