1syntax = "proto3"; 2 3package tensorflow.data; 4 5import "tensorflow/core/framework/graph.proto"; 6import "tensorflow/core/protobuf/data_service.proto"; 7 8// Next tag: 2 9message DatasetDef { 10 // We represent datasets as tensorflow GraphDefs which define the operations 11 // needed to create a tf.data dataset. 12 GraphDef graph = 1; 13} 14 15// Next tag: 13 16message TaskDef { 17 reserved 6; 18 // The dataset to iterate over. 19 oneof dataset { 20 DatasetDef dataset_def = 1; 21 string path = 2; 22 } 23 int64 dataset_id = 3; 24 int64 task_id = 4; 25 int64 job_id = 5; 26 // In distributed epoch processing mode, we use one split provider for each 27 // source that feeds into the dataset. In parallel_epochs mode, 28 // `num_split_providers` is always zero. 29 int64 num_split_providers = 9; 30 // Address of the worker that the task is assigned to. 31 string worker_address = 8; 32 ProcessingModeDef processing_mode_def = 10; 33 // Optional number of consumers. If set, the results of the task will be 34 // provided to consumers round-robin. 35 oneof optional_num_consumers { 36 int64 num_consumers = 7; 37 } 38 // Number of workers and the worker index. These are only populated when the 39 // `processing_mode_def` specifies a static sharding policy. 40 int64 num_workers = 11; 41 int64 worker_index = 12; 42} 43 44// Next tag: 6 45message TaskInfo { 46 // The address of the worker processing the task. 47 string worker_address = 1; 48 // The transfer address of the worker processing the task. 49 string transfer_address = 4; 50 // The task id. 51 int64 task_id = 2; 52 // The id of the job that the task is part of. 53 int64 job_id = 3; 54 // The round to start reading from the task in. For non-round-robin reads, 55 // this is always 0. 56 int64 starting_round = 5; 57} 58 59// Specifies which tf.data service workers to read from. 60enum TargetWorkers { 61 TARGET_WORKERS_UNSPECIFIED = 0; 62 // tf.data service runtime decides which workers to read from. 63 TARGET_WORKERS_AUTO = 1; 64 // Reads from any available worker. 65 TARGET_WORKERS_ANY = 2; 66 // Only reads from local workers. If no local worker is found, it is an error. 67 TARGET_WORKERS_LOCAL = 3; 68} 69