Get Tokens Per Second metric

backend/backend.proto

@@ -19,6 +19,8 @@ service Backend {
   rpc SoundGeneration(SoundGenerationRequest) returns (Result) {}
   rpc TokenizeString(PredictOptions) returns (TokenizationResponse) {}
   rpc Status(HealthMessage) returns (StatusResponse) {}
+  // expose tokens_per_second metric
+  rpc GetMetrics(HealthMessage) returns (MetricsResponse);
 
   rpc StoresSet(StoresSetOptions) returns (Result) {}
   rpc StoresDelete(StoresDeleteOptions) returns (Result) {}
@@ -28,6 +30,12 @@ service Backend {
   rpc Rerank(RerankRequest) returns (RerankResult) {}
 }
 
+message MetricsResponse {
+  float tokens_per_second = 1;
+  int32 tokens_generated = 2;
+  int32 prompt_tokens_processed = 3;
+}
+
 message RerankRequest {
   string query = 1;
   repeated string documents = 2;

backend/cpp/llama/grpc-server.cpp

@@ -441,6 +441,16 @@ struct llama_server_context
         }
     }
 
+    llama_client_slot* get_active_slot() {
+        for (llama_client_slot& slot : slots) {
+            // Check if the slot is currently processing
+            if (slot.is_processing()) {
+                return &slot;  // Return the active slot
+            }
+        }
+        return nullptr;  // No active slot found
+    }
+
     bool load_model(const gpt_params &params_)
     {
         params = params_;
@@ -2031,9 +2041,10 @@ inline void signal_handler(int signal) { shutdown_handler(signal); }
 bool loaded_model; // TODO: add a mutex for this, but happens only once loading the model
 
 // The class has a llama instance that is shared across all RPCs
-llama_server_context llama;
+static llama_server_context llama;
 
 static void start_llama_server() {
+
     // Wait for model to be loaded first
     while (!loaded_model) {
         std::this_thread::sleep_for(std::chrono::milliseconds(100));
@@ -2406,6 +2417,27 @@ class BackendServiceImpl final : public backend::Backend::Service {
 
         return grpc::Status::OK;
     }
+
+    grpc::Status GetMetrics(ServerContext* context, const backend::HealthMessage* request, backend::MetricsResponse* response) {
+        llama_client_slot* active_slot = llama.get_active_slot();
+
+        if (active_slot != nullptr) {
+            // Calculate the tokens per second using existing logic
+            double tokens_per_second = 1e3 / active_slot->t_token_generation * active_slot->n_decoded;
+
+            // Populate the response with metrics
+            response->set_tokens_per_second(tokens_per_second);
+            response->set_tokens_generated(active_slot->n_decoded);
+            response->set_prompt_tokens_processed(active_slot->num_prompt_tokens_processed);
+        } else {
+            // Handle case when no active slot exists
+            response->set_tokens_per_second(0);
+            response->set_tokens_generated(0);
+            response->set_prompt_tokens_processed(0);
+        }
+
+        return grpc::Status::OK;
+    } 
 };
 
 void RunServer(const std::string& server_address) {

pkg/grpc/backend.go

@@ -51,4 +51,5 @@ type Backend interface {
 	StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts ...grpc.CallOption) (*pb.StoresFindResult, error)
 
 	Rerank(ctx context.Context, in *pb.RerankRequest, opts ...grpc.CallOption) (*pb.RerankResult, error)
+	GetTokenMetrics(ctx context.Context) (*pb.MetricsResponse, error)
 }

pkg/grpc/client.go

@@ -50,6 +50,32 @@ func (c *Client) wdUnMark() {
 	}
 }
 
+func (c *Client) GetTokenMetrics(ctx context.Context) (*pb.MetricsResponse, error) {
+	if !c.parallel {
+		c.opMutex.Lock()
+		defer c.opMutex.Unlock()
+	}
+	c.setBusy(true)
+	defer c.setBusy(false)
+	conn, err := grpc.Dial(c.address, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close()
+	client := pb.NewBackendClient(conn)
+
+	// The healthcheck call shouldn't take long time
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	res, err := client.GetMetrics(ctx, &pb.HealthMessage{})
+	if err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
 func (c *Client) HealthCheck(ctx context.Context) (bool, error) {
 	if !c.parallel {
 		c.opMutex.Lock()

pkg/grpc/embed.go

@@ -23,6 +23,10 @@ func (e *embedBackend) HealthCheck(ctx context.Context) (bool, error) {
 	return true, nil
 }
 
+func (e *embedBackend) GetTokenMetrics(ctx context.Context) (*pb.MetricsResponse, error) {
+	return e.s.GetMetrics(ctx, &pb.HealthMessage{})
+}
+
 func (e *embedBackend) Embeddings(ctx context.Context, in *pb.PredictOptions, opts ...grpc.CallOption) (*pb.EmbeddingResult, error) {
 	return e.s.Embedding(ctx, in)
 }