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 {}