117 lines
2.3 KiB
Protocol Buffer
117 lines
2.3 KiB
Protocol Buffer
syntax = "proto3";
|
|
|
|
package node_service;
|
|
|
|
service NodeService {
|
|
rpc SendPrompt (PromptRequest) returns (Tensor) {}
|
|
rpc SendTensor (TensorRequest) returns (Tensor) {}
|
|
rpc SendExample (ExampleRequest) returns (Loss) {}
|
|
rpc CollectTopology (CollectTopologyRequest) returns (Topology) {}
|
|
rpc SendResult (SendResultRequest) returns (Empty) {}
|
|
rpc SendOpaqueStatus (SendOpaqueStatusRequest) returns (Empty) {}
|
|
rpc HealthCheck (HealthCheckRequest) returns (HealthCheckResponse) {}
|
|
}
|
|
|
|
message Shard {
|
|
string model_id = 1;
|
|
int32 start_layer = 2;
|
|
int32 end_layer = 3;
|
|
int32 n_layers = 4;
|
|
}
|
|
|
|
message PromptRequest {
|
|
Shard shard = 1;
|
|
string prompt = 2;
|
|
optional string request_id = 3;
|
|
optional InferenceState inference_state = 4;
|
|
}
|
|
|
|
message TensorRequest {
|
|
Shard shard = 1;
|
|
Tensor tensor = 2;
|
|
optional string request_id = 3;
|
|
optional InferenceState inference_state = 4;
|
|
}
|
|
|
|
message ExampleRequest {
|
|
Shard shard = 1;
|
|
Tensor example = 2;
|
|
Tensor target = 3;
|
|
Tensor length = 4;
|
|
bool train = 5;
|
|
optional string request_id = 6;
|
|
}
|
|
|
|
message Loss {
|
|
float loss = 1;
|
|
optional Tensor grads = 2;
|
|
}
|
|
|
|
message Tensor {
|
|
bytes tensor_data = 1;
|
|
repeated int32 shape = 2;
|
|
string dtype = 3;
|
|
}
|
|
|
|
message TensorList {
|
|
repeated Tensor tensors = 1;
|
|
}
|
|
|
|
message InferenceState {
|
|
map<string, Tensor> tensor_data = 1;
|
|
map<string, TensorList> tensor_list_data = 2;
|
|
string other_data_json = 3;
|
|
}
|
|
|
|
message CollectTopologyRequest {
|
|
repeated string visited = 1;
|
|
int32 max_depth = 2;
|
|
}
|
|
|
|
message Topology {
|
|
map<string, DeviceCapabilities> nodes = 1;
|
|
map<string, PeerConnections> peer_graph = 2;
|
|
}
|
|
|
|
message PeerConnection {
|
|
string to_id = 1;
|
|
optional string description = 2;
|
|
}
|
|
|
|
message PeerConnections {
|
|
repeated PeerConnection connections = 1;
|
|
}
|
|
|
|
message DeviceFlops {
|
|
double fp32 = 1;
|
|
double fp16 = 2;
|
|
double int8 = 3;
|
|
}
|
|
|
|
message DeviceCapabilities {
|
|
string model = 1;
|
|
string chip = 2;
|
|
int32 memory = 3;
|
|
DeviceFlops flops = 4;
|
|
}
|
|
|
|
message SendResultRequest {
|
|
string request_id = 1;
|
|
repeated int32 result = 2;
|
|
optional Tensor tensor = 3;
|
|
bool is_finished = 4;
|
|
}
|
|
|
|
message SendOpaqueStatusRequest {
|
|
string request_id = 1;
|
|
string status = 2;
|
|
}
|
|
|
|
message HealthCheckRequest {}
|
|
|
|
message HealthCheckResponse {
|
|
bool is_healthy = 1;
|
|
}
|
|
|
|
message Empty {}
|