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
20enum MemoryRegion { HBM = 1; }
21
22message ChipCoordinate {
23  required int32 x = 1;
24  required int32 y = 2;
25  required int32 z = 3;
26}
27
28message TpuCoreInfo {
29  required int32 id = 1;
30  optional int32 core_on_chip_index = 2;
31  optional int32 core_on_host_index = 3;
32  optional int64 hbm_bytes_available = 100;
33  optional int64 hbm_bytes_allocatable = 101;
34}
35
36message TpuChipInfo {
37  repeated TpuCoreInfo core = 1;
38  optional int32 host_id = 2;
39  optional ChipCoordinate chip_coord = 3;
40}
41
42message CpuInfo {
43  required int32 num_cpu_cores = 1;
44  required float cpu_load_average_1min = 2;
45  required int64 ram_bytes_total = 100;
46  required int64 ram_bytes_available = 101;
47}
48
49message SystemInfo {
50  repeated TpuChipInfo tpu_chip = 1;
51  required CpuInfo cpu = 2;
52  repeated TpuCoreInfo local_core = 3;
53  optional int32 host_id = 4;
54  optional int32 host_count = 5;
55  optional int32 chip_count = 6;
56  optional int32 core_count = 7;
57}
58
59message TpuDriverConfig {
60  optional string worker = 1;
61
62  message GrpcConfig {
63    // Time in seconds before the initial connection to the server will timeout.
64    optional int64 connection_timeout_secs = 1 [default = 30];
65
66    // Time in seconds the server may be unresponsive before terminating the
67    // connection.
68    optional int64 keepalive_timeout_secs = 2 [default = 30];
69  }
70
71  optional GrpcConfig grpc = 2;
72}
73