Skip to content

Commit 71505c4

Browse files
committed
used sysinfo as optional dependency
1 parent ad14fee commit 71505c4

File tree

2 files changed

+113
-62
lines changed

2 files changed

+113
-62
lines changed

workflows/Cargo.toml

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,15 @@ rand.workspace = true
3030
log.workspace = true
3131
eyre.workspace = true
3232

33-
# system info
34-
sysinfo = "0.32.0"
33+
sysinfo = { version = "0.32.0", optional = true }
34+
35+
[features]
36+
profiling = ["sysinfo"]
3537

3638
[dev-dependencies]
3739
# only used for tests
3840
dotenvy.workspace = true
3941

4042
[[bin]]
4143
name = "tps"
42-
path = "src/bin/tps.rs"
44+
path = "src/bin/tps.rs"

workflows/src/bin/tps.rs

Lines changed: 108 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,69 +1,116 @@
1+
use std::vec;
2+
13
use dkn_workflows::{DriaWorkflowsConfig, OllamaConfig};
24
use ollama_workflows::ollama_rs::{
35
generation::{completion::request::GenerationRequest, options::GenerationOptions},
46
Ollama,
57
};
6-
use sysinfo::{CpuRefreshKind, RefreshKind, System};
8+
use ollama_workflows::Model;
9+
10+
#[cfg(feature = "profiling")]
11+
use sysinfo::{CpuRefreshKind, RefreshKind, System, MINIMUM_CPU_UPDATE_INTERVAL};
712

813
#[tokio::main]
914
async fn main() {
10-
// initialize logger
11-
env_logger::init();
15+
#[cfg(feature = "profiling")]
16+
{
17+
// initialize logger
18+
env_logger::init();
1219

13-
let cfg = DriaWorkflowsConfig::new_from_csv("finalend/hermes-3-llama-3.1:8b-q8_0,phi3:14b-medium-4k-instruct-q4_1,phi3:14b-medium-128k-instruct-q4_1,phi3.5:3.8b,phi3.5:3.8b-mini-instruct-fp16,gemma2:9b-instruct-q8_0,gemma2:9b-instruct-fp16,llama3.1:latest,llama3.1:8b-instruct-q8_0,llama3.1:8b-instruct-fp16,llama3.1:70b-instruct-q4_0,llama3.1:70b-instruct-q8_0,llama3.2:1b,llama3.2:3b,qwen2.5:7b-instruct-q5_0,qwen2.5:7b-instruct-fp16,qwen2.5:32b-instruct-fp16,qwen2.5-coder:1.5b,qwen2.5-coder:7b-instruct,llama3.2:3b,qwen2.5-coder:7b-instruct-q8_0,qwen2.5-coder:7b-instruct-fp16,deepseek-coder:6.7b,mixtral:8x7b");
14-
let config = OllamaConfig::default();
15-
let ollama = Ollama::new(config.host, config.port);
20+
let models = vec![
21+
Model::NousTheta,
22+
Model::Phi3Medium,
23+
Model::Phi3Medium128k,
24+
Model::Phi3_5Mini,
25+
Model::Phi3_5MiniFp16,
26+
Model::Gemma2_9B,
27+
Model::Gemma2_9BFp16,
28+
Model::Llama3_1_8B,
29+
Model::Llama3_1_8Bq8,
30+
Model::Llama3_1_8Bf16,
31+
Model::Llama3_1_8BTextQ4KM,
32+
Model::Llama3_1_8BTextQ8,
33+
Model::Llama3_1_70B,
34+
Model::Llama3_1_70Bq8,
35+
Model::Llama3_1_70BTextQ4KM,
36+
Model::Llama3_2_1B,
37+
Model::Llama3_2_3B,
38+
Model::Llama3_2_1BTextQ4KM,
39+
Model::Qwen2_5_7B,
40+
Model::Qwen2_5_7Bf16,
41+
Model::Qwen2_5_32Bf16,
42+
Model::Qwen2_5Coder1_5B,
43+
Model::Qwen2_5coder7B,
44+
Model::Qwen2_5oder7Bq8,
45+
Model::Qwen2_5coder7Bf16,
46+
Model::DeepSeekCoder6_7B,
47+
Model::Mixtral8_7b,
48+
Model::GPT4Turbo,
49+
Model::GPT4o,
50+
Model::GPT4oMini,
51+
Model::O1Preview,
52+
Model::O1Mini,
53+
Model::Gemini15ProExp0827,
54+
Model::Gemini15Pro,
55+
Model::Gemini15Flash,
56+
Model::Gemini10Pro,
57+
Model::Gemma2_2bIt,
58+
Model::Gemma2_27bIt,
59+
];
1660

17-
log::info!("Starting...");
18-
// ensure that all lists of CPUs and processes are filled
19-
let mut system = System::new_all();
20-
// update all information of the system
21-
system.refresh_all();
61+
let cfg = DriaWorkflowsConfig::new(models);
62+
let config = OllamaConfig::default();
63+
let ollama = Ollama::new(config.host, config.port);
64+
log::info!("Starting...");
65+
// ensure that all lists of CPUs and processes are filled
66+
let mut system = System::new_all();
67+
// update all information of the system
68+
system.refresh_all();
2269

23-
log::debug!("Getting system information...");
24-
let brand = system.cpus()[0].brand().to_string();
25-
let os_name = System::name().unwrap_or_else(|| "Unknown".to_string());
26-
let os_version = System::long_os_version().unwrap_or_else(|| "Unknown".to_string());
27-
let cpu_usage = system.global_cpu_usage();
28-
let total_memory = system.total_memory();
29-
let used_memory = system.used_memory();
70+
log::debug!("Getting system information...");
71+
let brand = system.cpus()[0].brand().to_string();
72+
let os_name = System::name().unwrap_or_else(|| "Unknown".to_string());
73+
let os_version = System::long_os_version().unwrap_or_else(|| "Unknown".to_string());
74+
let cpu_usage = system.global_cpu_usage();
75+
let total_memory = system.total_memory();
76+
let used_memory = system.used_memory();
3077

31-
for (_, model) in cfg.models {
32-
log::info!("Pulling model: {}", model);
78+
for (_, model) in cfg.models {
79+
log::info!("Pulling model: {}", model);
3380

34-
// pull model
35-
match ollama.pull_model(model.to_string(), false).await {
36-
Ok(status) => log::info!("Status: {}", status.message),
37-
Err(err) => {
38-
log::error!("Failed to pull model {}: {:?}", model, err);
81+
// pull model
82+
match ollama.pull_model(model.to_string(), false).await {
83+
Ok(status) => log::info!("Status: {}", status.message),
84+
Err(err) => {
85+
log::error!("Failed to pull model {}: {:?}", model, err);
86+
}
3987
}
40-
}
4188

42-
log::debug!("Creating request...");
43-
// create dummy request
44-
let mut generation_request =
45-
GenerationRequest::new(model.to_string(), "compute 6780 * 1200".to_string());
89+
log::debug!("Creating request...");
90+
// create dummy request
91+
let mut generation_request =
92+
GenerationRequest::new(model.to_string(), "compute 6780 * 1200".to_string());
4693

47-
if let Ok(num_thread) = std::env::var("OLLAMA_NUM_THREAD") {
48-
generation_request = generation_request.options(
49-
GenerationOptions::default().num_thread(
50-
num_thread
51-
.parse()
52-
.expect("num threads should be a positive integer"),
53-
),
54-
);
55-
}
94+
if let Ok(num_thread) = std::env::var("OLLAMA_NUM_THREAD") {
95+
generation_request = generation_request.options(
96+
GenerationOptions::default().num_thread(
97+
num_thread
98+
.parse()
99+
.expect("num threads should be a positive integer"),
100+
),
101+
);
102+
}
56103

57-
// generate response
58-
match ollama.generate(generation_request).await {
59-
Ok(response) => {
60-
log::debug!("Got response for model {}", model);
61-
// compute TPS
62-
let tps = (response.eval_count.unwrap_or_default() as f64)
63-
/ (response.eval_duration.unwrap_or(1) as f64)
64-
* 1_000_000_000f64;
65-
// report machine info
66-
log::info!(
104+
// generate response
105+
match ollama.generate(generation_request).await {
106+
Ok(response) => {
107+
log::debug!("Got response for model {}", model);
108+
// compute TPS
109+
let tps = (response.eval_count.unwrap_or_default() as f64)
110+
/ (response.eval_duration.unwrap_or(1) as f64)
111+
* 1_000_000_000f64;
112+
// report machine info
113+
log::info!(
67114
"\n Model: {} \n TPS: {} \n OS: {} {} \n Version: {} \n CPU Usage: % {} \n Total Memory: {} KB \n Used Memory: {} KB ",
68115
model,
69116
tps,
@@ -74,18 +121,20 @@ async fn main() {
74121
total_memory,
75122
used_memory,
76123
);
124+
}
125+
Err(e) => {
126+
log::warn!("Ignoring model {}: Workflow failed with error {}", model, e);
127+
}
77128
}
78-
Err(e) => {
79-
log::warn!("Ignoring model {}: Workflow failed with error {}", model, e);
80-
}
129+
// refresh CPU usage (https://docs.rs/sysinfo/latest/sysinfo/struct.Cpu.html#method.cpu_usage)
130+
system = System::new_with_specifics(
131+
RefreshKind::new().with_cpu(CpuRefreshKind::everything()),
132+
);
133+
// wait a bit because CPU usage is based on diff
134+
std::thread::sleep(MINIMUM_CPU_UPDATE_INTERVAL);
135+
// refresh CPUs again to get actual value
136+
system.refresh_cpu_usage();
81137
}
82-
// refresh CPU usage (https://docs.rs/sysinfo/latest/sysinfo/struct.Cpu.html#method.cpu_usage)
83-
system =
84-
System::new_with_specifics(RefreshKind::new().with_cpu(CpuRefreshKind::everything()));
85-
// wait a bit because CPU usage is based on diff
86-
std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
87-
// refresh CPUs again to get actual value
88-
system.refresh_cpu_usage();
89138
}
90139
log::info!("Finished");
91140
}

0 commit comments

Comments
 (0)