From 2026b4dd0ecc38bc806b45855c1eeaa593c540c3 Mon Sep 17 00:00:00 2001
From: vodkaslime <646329483@qq.com>
Date: Tue, 15 Aug 2023 23:22:03 +0800
Subject: [PATCH] feat: add architecture/cpu info to health api response
 [TAB-162] (#355)

* feat: add architecture, cpu and gpu info to health command

* chore: fix

* chore: fix

* chore: fix

* chore: fix lint

* chore: fix lint

* chore: remove gpu

* chore: resolve comments

* chore: resolve comments

* Update health.rs

---------

Co-authored-by: Meng Zhang <meng@tabbyml.com>
---
 Cargo.lock                       | 33 ++++++++++++++++++++++++++++----
 crates/tabby/Cargo.toml          |  1 +
 crates/tabby/src/serve/health.rs | 19 +++++++++++++++++-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1935bc0..b3dd739 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1740,6 +1740,15 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "ntapi"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "nu-ansi-term"
 version = "0.46.0"
@@ -2774,6 +2783,21 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160"
 
+[[package]]
+name = "sysinfo"
+version = "0.29.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d10ed79c22663a35a255d289a7fdcb43559fc77ff15df5ce6c341809e7867528"
+dependencies = [
+ "cfg-if",
+ "core-foundation-sys",
+ "libc",
+ "ntapi",
+ "once_cell",
+ "rayon",
+ "winapi",
+]
+
 [[package]]
 name = "tabby"
 version = "0.1.0"
@@ -2794,6 +2818,7 @@ dependencies = [
  "serdeconv",
  "strfmt",
  "strum",
+ "sysinfo",
  "tabby-common",
  "tabby-download",
  "tabby-inference",
@@ -3002,18 +3027,18 @@ dependencies = [
 
 [[package]]
 name = "thiserror"
-version = "1.0.40"
+version = "1.0.45"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "978c9a314bd8dc99be594bc3c175faaa9794be04a5a5e153caba6915336cebac"
+checksum = "dedd246497092a89beedfe2c9f176d44c1b672ea6090edc20544ade01fbb7ea0"
 dependencies = [
  "thiserror-impl",
 ]
 
 [[package]]
 name = "thiserror-impl"
-version = "1.0.40"
+version = "1.0.45"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f9456a42c5b0d803c8cd86e73dd7cc9edd429499f37a3550d286d5e86720569f"
+checksum = "7d7b1fadccbbc7e19ea64708629f9d8dccd007c260d66485f20a6d41bc1cf4b3"
 dependencies = [
  "proc-macro2",
  "quote",
diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml
index b614939..3b3b1f3 100644
--- a/crates/tabby/Cargo.toml
+++ b/crates/tabby/Cargo.toml
@@ -33,6 +33,7 @@ axum-tracing-opentelemetry = "0.10.0"
 tracing-opentelemetry = "0.18.0"
 tantivy = { workspace = true }
 anyhow = { workspace = true }
+sysinfo = "0.29.8"
 
 
 [dependencies.uuid]
diff --git a/crates/tabby/src/serve/health.rs b/crates/tabby/src/serve/health.rs
index 283eb60..a36cac5 100644
--- a/crates/tabby/src/serve/health.rs
+++ b/crates/tabby/src/serve/health.rs
@@ -1,7 +1,8 @@
-use std::sync::Arc;
+use std::{env::consts::ARCH, sync::Arc};
 
 use axum::{extract::State, Json};
 use serde::{Deserialize, Serialize};
+use sysinfo::{CpuExt, System, SystemExt};
 use utoipa::ToSchema;
 
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@@ -9,14 +10,30 @@ pub struct HealthState {
     model: String,
     device: String,
     compute_type: String,
+    arch: String,
+    cpu_info: String,
+    cpu_count: usize,
 }
 
 impl HealthState {
     pub fn new(args: &super::ServeArgs) -> Self {
+        let mut sys = System::new_all();
+        sys.refresh_cpu();
+        let cpus = sys.cpus();
+        let cpu_info = if cpus.len() > 0 {
+            let cpu = &cpus[0];
+            cpu.brand().to_string()
+        } else {
+            "unknown".to_string()
+        };
+
         Self {
             model: args.model.clone(),
             device: args.device.to_string(),
             compute_type: args.compute_type.to_string(),
+            arch: ARCH.to_string(),
+            cpu_info,
+            cpu_count: cpus.len(),
         }
     }
 }