1// Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14// ============================================================================== 15 16syntax = "proto2"; 17 18package tpu_driver; 19 20import "tensorflow/compiler/xla/python/tpu_driver/tpu_driver.proto"; 21import "tensorflow/compiler/xla/service/hlo.proto"; 22import "tensorflow/compiler/xla/xla_data.proto"; 23 24option optimize_for = SPEED; 25 26message StatusMessage { 27 required int32 code = 1; 28 optional string message = 2; 29} 30 31message AllocateRequest { 32 required int32 core_id = 1; 33 required MemoryRegion region = 2; 34 oneof size { 35 int64 num_bytes = 3; 36 xla.ShapeProto shape = 4; 37 } 38} 39 40message AllocateTupleRequest { 41 required int32 core_id = 1; 42 required MemoryRegion region = 2; 43 repeated int64 children = 3; 44} 45 46message DeallocateRequest { 47 required int64 handle = 1; 48} 49 50message TransferToDeviceRequest { 51 required int64 target_handle = 1; 52 required bytes data = 2; 53} 54 55message TransferFromDeviceRequest { 56 required int64 source_handle = 1; 57} 58 59message TransferFromDeviceResponse { 60 required bytes data = 2; 61} 62 63message TransferFromDeviceToDeviceRequest { 64 required int64 source_handle = 1; 65 required int64 target_handle = 2; 66} 67 68message CompileRequest { 69 required xla.HloProto hlo_program = 1; 70 optional int64 num_replicas = 2; 71} 72 73message CompiledProgramMetadata { 74 required xla.ProgramShapeProto program_shape = 1; 75} 76 77message CompileResponse { 78 required CompiledProgramMetadata metadata = 1; 79} 80 81message LoadProgramRequest { 82 required int32 core_id = 1; 83 required int64 compiled_program_handle = 2; 84} 85 86message UnloadProgramRequest { 87 required int64 loaded_program_handle = 1; 88} 89 90message ExecuteRequest { 91 required int64 loaded_program_handle = 1; 92 repeated int64 input_handle = 2; 93 repeated int64 output_handle = 3; 94 optional xla.DeviceAssignmentProto device_assignment = 4; 95} 96 97message StreamRequest { 98 message Entry { 99 oneof request { 100 AllocateRequest alloc = 1; 101 AllocateTupleRequest alloc_tuple = 2; 102 DeallocateRequest dealloc = 3; 103 TransferToDeviceRequest transfer_to = 4; 104 TransferFromDeviceRequest transfer_from = 5; 105 TransferFromDeviceToDeviceRequest transfer_from_to = 10; 106 CompileRequest compile = 6; 107 LoadProgramRequest load = 7; 108 UnloadProgramRequest unload = 8; 109 ExecuteRequest execute = 9; 110 } 111 // If specified, a list of encoded EventId values. 112 repeated int64 wait_for_id = 20; 113 // A unique, encoded EventId value. 114 // For Allocate, Compile, and Load, this also defines the result handle. 115 required int64 operation_id = 21; 116 117 // A unique identifier for the thread that issued this request. Currently 118 // for debugging purposes only. 119 optional int64 thread_id = 22; 120 } 121 122 repeated Entry entry = 30; 123} 124 125message StreamResponse { 126 message Entry { 127 oneof response { 128 TransferFromDeviceResponse transfer_from = 3; 129 CompileResponse compile = 4; 130 } 131 required StatusMessage status = 10; 132 // Echos the given encoded EventId value. 133 required int64 operation_id = 11; 134 } 135 136 repeated Entry entry = 20; 137} 138 139message OpenRequest { 140 // The version number for this client. Versions are bumped in case of 141 // backwards incompatible client-server protocol changes. Servers will reject 142 // clients with an unsupported version. 143 optional int32 client_version = 1 [default = 0]; 144} 145 146message OpenResponse { 147 required uint32 client_id = 1; 148 149 // Maximum time this client can be idle before it is GC'ed and all resources 150 // released. 151 optional int32 max_idle_time_seconds = 2 [default = 3600]; 152} 153 154message CloseRequest { 155 required fixed32 client_id = 1; 156} 157 158message CloseResponse {} 159 160message ResetRequest {} 161 162message ResetResponse {} 163 164message QuerySystemInfoRequest {} 165 166message QuerySystemInfoResponse { 167 required SystemInfo system_info = 1; 168} 169 170service CloudTpuDriver { 171 // Open the driver. If the driver is already open, return an error. 172 rpc Open(OpenRequest) returns (OpenResponse); 173 174 // Close the driver. Any outstanding requests will be terminated. 175 rpc Close(CloseRequest) returns (CloseResponse); 176 177 // Reset the driver. All connected clients will be disconnected. 178 rpc Reset(ResetRequest) returns (ResetResponse); 179 180 // Query the driver for current system performance information. 181 rpc QuerySystemInfo(QuerySystemInfoRequest) returns (QuerySystemInfoResponse); 182 183 // Enqueue an operation to be executed when its dependencies are satisfied. 184 rpc StreamExecute(stream StreamRequest) returns (stream StreamResponse); 185} 186