Remote start stop control
This commit is contained in:
Jaakko Vanhala
@@ -5,6 +5,7 @@ use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
mod inference;
|
||||
mod tui;
|
||||
|
||||
/// GPU-tietorakenne — yhtenäinen kaikille valmistajille
|
||||
struct GpuInfo {
|
||||
@@ -354,74 +355,115 @@ async fn main() {
|
||||
continue;
|
||||
}
|
||||
|
||||
while let Some(Ok(msg)) = read.next().await {
|
||||
if let Message::Text(text) = msg {
|
||||
// LLM-promptit
|
||||
if text.contains("llm_prompt") {
|
||||
if let Ok(task) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
let prompt = task.get("prompt").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let task_id = task.get("task_id").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let msg_model = task.get("model").and_then(|v| v.as_str()).unwrap_or("");
|
||||
|
||||
if !prompt.is_empty() && (msg_model.starts_with("qwen-coder") || msg_model.starts_with("qwen2.5-coder")) {
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
let mut stdin_lines = tokio::io::BufReader::new(tokio::io::stdin()).lines();
|
||||
|
||||
if let Some(ref engine) = llm {
|
||||
let max_tokens = task.get("max_tokens").and_then(|v| v.as_u64()).unwrap_or(1024) as usize;
|
||||
let prompt_lines = prompt.lines().count();
|
||||
let prompt_last: String = prompt.lines().last().unwrap_or("").chars().take(60).collect();
|
||||
tracing::info!("→ task_id:{} | {}r prompti | \"{}...\"", task_id, prompt_lines, prompt_last);
|
||||
|
||||
let model_name = engine.model_name();
|
||||
match engine.generate(prompt, max_tokens).await {
|
||||
Ok(result) => {
|
||||
tracing::info!(
|
||||
"✓ {} | {} tok | {:.0}ms | {:.1} tok/s",
|
||||
model_name,
|
||||
result.tokens_generated,
|
||||
result.duration_ms,
|
||||
result.tokens_per_sec,
|
||||
);
|
||||
|
||||
// Lähetetään vain lyhyt prompti-esikatselu (ei koko kontekstia)
|
||||
let prompt_short: String = prompt.lines().last().unwrap_or("").chars().take(100).collect();
|
||||
let done = json!({
|
||||
"type": "llm_done",
|
||||
"prompt": prompt_short,
|
||||
"model": format!("{} (Ollama)", model_name),
|
||||
"response": result.text,
|
||||
"tokens_generated": result.tokens_generated,
|
||||
"duration_ms": result.duration_ms,
|
||||
"tokens_per_sec": (result.tokens_per_sec * 10.0).round() / 10.0,
|
||||
"load_time_ms": 0,
|
||||
"task_id": task_id,
|
||||
});
|
||||
let _ = write.send(Message::Text(done.to_string())).await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Inferenssivirhe: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
loop {
|
||||
tokio::select! {
|
||||
line = stdin_lines.next_line() => {
|
||||
if let Ok(Some(text)) = line {
|
||||
let t = text.trim();
|
||||
if t == "p" || t == "pause" {
|
||||
tracing::info!("Tauotetaan solmun suoritus (Hub ei lähetä tehtäviä)...");
|
||||
let req = json!({"type": "status_update", "status": "paused"});
|
||||
let _ = write.send(Message::Text(req.to_string())).await;
|
||||
} else if t == "r" || t == "resume" || t == "s" {
|
||||
tracing::info!("Jatketaan solmun suoritusta...");
|
||||
let req = json!({"type": "status_update", "status": "active"});
|
||||
let _ = write.send(Message::Text(req.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Mallin vaihto lennossa
|
||||
if text.contains("change_model") {
|
||||
if let Ok(task) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
if let Some(new_model) = task.get("model").and_then(|v| v.as_str()) {
|
||||
if let Some(ref engine) = llm {
|
||||
tracing::info!("Vaihdetaan malli: {}", new_model);
|
||||
engine.set_model(new_model.to_string());
|
||||
match engine.ensure_model().await {
|
||||
Ok(()) => tracing::info!("Malli {} valmis!", new_model),
|
||||
Err(e) => tracing::error!("Mallin lataus epäonnistui: {}", e),
|
||||
ws_msg = read.next() => {
|
||||
match ws_msg {
|
||||
Some(Ok(Message::Text(text))) => {
|
||||
// Hubin control-viestit
|
||||
if text.contains(r#""type":"control""#) {
|
||||
if let Ok(task) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
if let Some(action) = task.get("action").and_then(|v| v.as_str()) {
|
||||
if action == "pause" {
|
||||
tracing::info!("Hub pakotti solmun tauolle (Pause)");
|
||||
let req = json!({"type": "status_update", "status": "paused"});
|
||||
let _ = write.send(Message::Text(req.to_string())).await;
|
||||
} else if action == "resume" {
|
||||
tracing::info!("Hub aktivoi solmun suorituksen (Resume)");
|
||||
let req = json!({"type": "status_update", "status": "active"});
|
||||
let _ = write.send(Message::Text(req.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// LLM-promptit
|
||||
if text.contains("llm_prompt") {
|
||||
if let Ok(task) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
let prompt = task.get("prompt").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let task_id = task.get("task_id").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let msg_model = task.get("model").and_then(|v| v.as_str()).unwrap_or("");
|
||||
|
||||
if !prompt.is_empty() && (msg_model.starts_with("qwen-coder") || msg_model.starts_with("qwen2.5-coder") || msg_model.starts_with("phi")) {
|
||||
if let Some(ref engine) = llm {
|
||||
let max_tokens = task.get("max_tokens").and_then(|v| v.as_u64()).unwrap_or(1024) as usize;
|
||||
let prompt_lines = prompt.lines().count();
|
||||
let prompt_last: String = prompt.lines().last().unwrap_or("").chars().take(60).collect();
|
||||
tracing::info!("→ task_id:{} | {}r prompti | \"{}...\"", task_id, prompt_lines, prompt_last);
|
||||
|
||||
let model_name = engine.model_name();
|
||||
match engine.generate(prompt, max_tokens).await {
|
||||
Ok(result) => {
|
||||
tracing::info!(
|
||||
"✓ {} | {} tok | {:.0}ms | {:.1} tok/s",
|
||||
model_name,
|
||||
result.tokens_generated,
|
||||
result.duration_ms,
|
||||
result.tokens_per_sec,
|
||||
);
|
||||
let prompt_short: String = prompt.lines().last().unwrap_or("").chars().take(100).collect();
|
||||
let done = json!({
|
||||
"type": "llm_done",
|
||||
"prompt": prompt_short,
|
||||
"model": format!("{} (Ollama)", model_name),
|
||||
"response": result.text,
|
||||
"tokens_generated": result.tokens_generated,
|
||||
"duration_ms": result.duration_ms,
|
||||
"tokens_per_sec": (result.tokens_per_sec * 10.0).round() / 10.0,
|
||||
"load_time_ms": 0,
|
||||
"task_id": task_id,
|
||||
});
|
||||
let _ = write.send(Message::Text(done.to_string())).await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Inferenssivirhe: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mallin vaihto lennossa
|
||||
if text.contains("change_model") {
|
||||
if let Ok(task) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
if let Some(new_model) = task.get("model").and_then(|v| v.as_str()) {
|
||||
if let Some(ref engine) = llm {
|
||||
tracing::info!("Vaihdetaan malli: {}", new_model);
|
||||
engine.set_model(new_model.to_string());
|
||||
match engine.ensure_model().await {
|
||||
Ok(()) => tracing::info!("Malli {} valmis!", new_model),
|
||||
Err(e) => tracing::error!("Mallin lataus epäonnistui: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Ok(_)) => {} // Muut viestityypit (binary/ping)
|
||||
Some(Err(_)) | None => break, // Yhteys poikki
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::warn!("Yhteys hubiin katkesi — yritetään uudelleen 5s...");
|
||||
}
|
||||
Err(e) => {
|
||||
|
||||
Reference in New Issue
Block a user