syntax = "proto3";

package node_service;

service NodeService {
  rpc SendPrompt (PromptRequest) returns (Tensor) {}
  rpc SendTensor (TensorRequest) returns (Tensor) {}
  rpc GetInferenceResult (GetInferenceResultRequest) returns (InferenceResult) {}
  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 image_str = 3;
  optional string request_id = 4;
  optional string inference_state = 5;
}

message TensorRequest {
  Shard shard = 1;
  Tensor tensor = 2;
  optional string request_id = 3;
  optional string inference_state = 4;
}

message GetInferenceResultRequest {
  string request_id = 1;
}

message InferenceResult {
  optional Tensor tensor = 1;
  bool is_finished = 2;
}

message Tensor {
  bytes tensor_data = 1;
  repeated int32 shape = 2;
  string dtype = 3;
}

message CollectTopologyRequest {
  repeated string visited = 1;
  int32 max_depth = 2;
}

message Topology {
  map<string, DeviceCapabilities> nodes = 1;
  map<string, Peers> peer_graph = 2;
}

message Peers {
    repeated string peer_ids = 1;
}

message DeviceFlops {
  float fp32 = 1;
  float fp16 = 2;
  float 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;
  bool is_finished = 3;
}

message SendOpaqueStatusRequest {
  string request_id = 1;
  string status = 2;
}

message HealthCheckRequest {}

message HealthCheckResponse {
  bool is_healthy = 1;
}

message Empty {}