• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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