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: 3 16message IterationKeyDef { 17 string name = 1; 18 int64 iteration = 2; 19} 20 21// Next tag: 14 22message TaskDef { 23 reserved 6; 24 // The dataset to iterate over. 25 oneof dataset { 26 DatasetDef dataset_def = 1; 27 string path = 2; 28 } 29 string dataset_id = 3; 30 int64 task_id = 4; 31 int64 iteration_id = 5; 32 // In distributed epoch processing mode, we use one split provider for each 33 // source that feeds into the dataset. In parallel_epochs mode, 34 // `num_split_providers` is always zero. 35 int64 num_split_providers = 9; 36 // Address of the worker that the task is assigned to. 37 string worker_address = 8; 38 ProcessingModeDef processing_mode_def = 10; 39 // Optional number of consumers. If set, the results of the task will be 40 // provided to consumers round-robin. 41 oneof optional_num_consumers { 42 int64 num_consumers = 7; 43 } 44 // Number of workers and the worker index. These are only populated when the 45 // `processing_mode_def` specifies a static sharding policy. 46 int64 num_workers = 11; 47 int64 worker_index = 12; 48 // True if cross-trainer cache is enabled. 49 bool use_cross_trainer_cache = 13; 50} 51 52// Next tag: 8 53message TaskInfo { 54 // The address of the worker processing the task. 55 string worker_address = 1; 56 // The transfer address of the worker processing the task. 57 string transfer_address = 4; 58 // Tags attached to the worker. This allows reading from selected workers. 59 // For example, by applying a "COLOCATED" tag, tf.data service is able to read 60 // from the local tf.data worker if one exists, then from off-TF-host workers, 61 // to avoid cross-TF-host reads. 62 repeated string worker_tags = 6; 63 // The task id. 64 int64 task_id = 2; 65 // The id of the iteration that the task is part of. 66 int64 iteration_id = 3; 67 // The UID of the worker Borg job, used for telemetry. 68 int64 worker_uid = 7; 69 // The round to start reading from the task in. For non-round-robin reads, 70 // this is always 0. 71 int64 starting_round = 5; 72} 73 74// Specifies which tf.data service workers to read from. 75enum TargetWorkers { 76 TARGET_WORKERS_UNSPECIFIED = 0; 77 // tf.data service runtime decides which workers to read from. 78 TARGET_WORKERS_AUTO = 1; 79 // Reads from any available worker. 80 TARGET_WORKERS_ANY = 2; 81 // Only reads from local workers. If no local worker is found, it is an error. 82 TARGET_WORKERS_LOCAL = 3; 83} 84