syntax = "proto3";
package ek.worker.v1;

import "ek/object/v1/object.proto";

message ForwardReq {
  message SequenceInfo { repeated string experts = 2; }
  string instance_id = 1;
  repeated SequenceInfo sequences = 2;
  // dimension: [seq,dim]
  bytes tensor = 3;
}

message ForwardResp {
  // dimension: [seq,expert,dim]
  bytes output_tensor = 1;
}

service ComputationService {
  rpc Forward(ForwardReq) returns (ForwardResp);
  // TODO: Batch Forward
}

message ExpertState {
  enum Stage {
    STAGE_UNSPECIFIED = 0;
    STAGE_ACTIVE = 1;
    STAGE_LOADING = 2;
    STAGE_EVICTING = 3;
  }
  Stage stage = 1;
}

message ExchangeReq {
  string id = 1;
  string addr = 2;
  string channel = 3;
  string device = 4;
  bool last_will = 5;
  uint32 rdma_tcp_port = 6;
  // List of currently loaded expert IDs on this worker (for progressive startup)
  // Workers report their full list of loaded experts in every heartbeat
  repeated string loaded_experts = 7;
  // gRPC address of this worker's LocalWeightManager (e.g., "http://10.0.0.5:5004")
  // Used by peer workers to fetch expert weights directly
  string wm_addr = 8;
  // Static memory capacity of this worker in megabytes (announced once at startup)
  uint64 mem_capacity_mb = 9;
  // Per-expert request counts since the last heartbeat; reset after each snapshot.
  // Key: expert_id (e.g. "qwen3/l0-e3"), Value: request count in this interval.
  map<string, uint64> expert_request_counts = 10;
}

message ExchangeResp {
  message ExpertWithState { ek.object.v1.ExpertSlice target = 1; }
  ExpertWithState state = 1;
  // Controller signals that proactive expert migration is complete and the
  // worker may proceed with shutdown.
  bool preemption_complete = 2;
}

service StateService {
  rpc Exchange(stream ExchangeReq) returns (stream ExchangeResp);
}