package tpu_driver

Get desktop application:
View/edit binary Protocol Buffers messages

• rpc Open (OpenRequest, OpenResponse)

  tpu_service.proto:172

  Open the driver. If the driver is already open, return an error.

  message OpenRequest

  tpu_service.proto:139

  • optional int32 client_version = 1

    The version number for this client. Versions are bumped in case of backwards incompatible client-server protocol changes. Servers will reject clients with an unsupported version.

  message OpenResponse

  tpu_service.proto:146

  • required uint32 client_id = 1

  • optional int32 max_idle_time_seconds = 2

    Maximum time this client can be idle before it is GC'ed and all resources released.

• rpc Close (CloseRequest, CloseResponse)

  tpu_service.proto:175

  Close the driver. Any outstanding requests will be terminated.

  message CloseRequest

  tpu_service.proto:154

  • required fixed32 client_id = 1

  message CloseResponse

  tpu_service.proto:158

  (message has no fields)

• rpc Reset (ResetRequest, ResetResponse)

  tpu_service.proto:178

  Reset the driver. All connected clients will be disconnected.

  message ResetRequest

  tpu_service.proto:160

  (message has no fields)

  message ResetResponse

  tpu_service.proto:162

  (message has no fields)

• rpc QuerySystemInfo (QuerySystemInfoRequest, QuerySystemInfoResponse)

  tpu_service.proto:181

  Query the driver for current system performance information.

  message QuerySystemInfoRequest

  tpu_service.proto:164

  (message has no fields)

  message QuerySystemInfoResponse

  tpu_service.proto:166

  • required SystemInfo system_info = 1

• rpc StreamExecute (stream StreamRequest, stream StreamResponse)

  tpu_service.proto:184

  Enqueue an operation to be executed when its dependencies are satisfied.

  message StreamRequest

  tpu_service.proto:97

  • repeated StreamRequest.Entry entry = 30

  message StreamResponse

  tpu_service.proto:125

  • repeated StreamResponse.Entry entry = 20

Used in: StreamRequest.Entry

• required int32 core_id = 1

• required MemoryRegion region = 2

• oneof size

  • int64 num_bytes = 3

  • xla.ShapeProto shape = 4

Used in: StreamRequest.Entry

• required int32 core_id = 1

• required MemoryRegion region = 2

• repeated int64 children = 3

Used in: TpuChipInfo

• required int32 x = 1

• required int32 y = 2

• required int32 z = 3

Used in: StreamRequest.Entry

• required xla.HloProto hlo_program = 1

• optional int64 num_replicas = 2

Used in: StreamResponse.Entry

• required CompiledProgramMetadata metadata = 1

Used in: CompileResponse

• required xla.ProgramShapeProto program_shape = 1

Used in: SystemInfo

• required int32 num_cpu_cores = 1

• required float cpu_load_average_1min = 2

• required int64 ram_bytes_total = 100

• required int64 ram_bytes_available = 101

Used in: StreamRequest.Entry

• required int64 handle = 1

Used in: StreamRequest.Entry

• required int64 loaded_program_handle = 1

• repeated int64 input_handle = 2

• repeated int64 output_handle = 3

• optional xla.DeviceAssignmentProto device_assignment = 4

Used in: StreamRequest.Entry

• required int32 core_id = 1

• required int64 compiled_program_handle = 2

Used in: AllocateRequest, AllocateTupleRequest

• HBM = 1

Used in: StreamResponse.Entry

• required int32 code = 1

• optional string message = 2

Used in: StreamRequest

• oneof request

  • AllocateRequest alloc = 1

  • AllocateTupleRequest alloc_tuple = 2

  • DeallocateRequest dealloc = 3

  • TransferToDeviceRequest transfer_to = 4

  • TransferFromDeviceRequest transfer_from = 5

  • TransferFromDeviceToDeviceRequest transfer_from_to = 10

  • CompileRequest compile = 6

  • LoadProgramRequest load = 7

  • UnloadProgramRequest unload = 8

  • ExecuteRequest execute = 9

• repeated int64 wait_for_id = 20

  If specified, a list of encoded EventId values.

• required int64 operation_id = 21

  A unique, encoded EventId value. For Allocate, Compile, and Load, this also defines the result handle.

• optional int64 thread_id = 22

  A unique identifier for the thread that issued this request. Currently for debugging purposes only.

Used in: StreamResponse

• oneof response

  • TransferFromDeviceResponse transfer_from = 3

  • CompileResponse compile = 4

• required StatusMessage status = 10

• required int64 operation_id = 11

  Echos the given encoded EventId value.

Used in: QuerySystemInfoResponse

• repeated TpuChipInfo tpu_chip = 1

• required CpuInfo cpu = 2

• repeated TpuCoreInfo local_core = 3

• optional int32 host_id = 4

• optional int32 host_count = 5

• optional int32 chip_count = 6

• optional int32 core_count = 7

Used in: SystemInfo

• repeated TpuCoreInfo core = 1

• optional int32 host_id = 2

• optional ChipCoordinate chip_coord = 3

Used in: SystemInfo, TpuChipInfo

• required int32 id = 1

• optional int32 core_on_chip_index = 2

• optional int32 core_on_host_index = 3

• optional int64 hbm_bytes_available = 100

• optional int64 hbm_bytes_allocatable = 101

• optional string worker = 1

• optional TpuDriverConfig.GrpcConfig grpc = 2

Used in: TpuDriverConfig

• optional int64 connection_timeout_secs = 1

  Time in seconds before the initial connection to the server will timeout.

• optional int64 keepalive_timeout_secs = 2

  Time in seconds the server may be unresponsive before terminating the connection.

Used in: StreamRequest.Entry

• required int64 source_handle = 1

Used in: StreamResponse.Entry

• required bytes data = 2

Used in: StreamRequest.Entry

• required int64 source_handle = 1

• required int64 target_handle = 2

Used in: StreamRequest.Entry

• required int64 target_handle = 1

• required bytes data = 2

Used in: StreamRequest.Entry

• required int64 loaded_program_handle = 1