feat: ops & UX batch — jobs survive restarts, logs/system in the console, media queue, gallery, chat, exact stream tokens
Fifteen changes aimed at the first day of real use: - Import and media jobs persist to <data>/jobs-*.json and reload on start (running ones settle as errors); the Register button no longer vanishes because the gateway restarted. - Docker ships the cortiq binary (built in the builder stage) — local .cmf chat models and image/video generation now work in containers, where auto-install never could (no cargo in the runtime image). libvulkan1 added for optional GPU offload. - Open admin console on 0.0.0.0 without a token now warns in the log and shows a dismissible banner — the VPS footgun is at least labeled. - "Restart gateway" button in Settings: stops managed model servers, spawns a fresh copy (bind retried while the port frees), Docker relies on the container restart policy instead. - Dashboard "System" panel: CPU sparkline + load, RAM (with the gateway's own RSS), free disk at the models dir — sampled by sysinfo every 2 s. - Featured model cards warn "not enough disk" / "may not fit in RAM" before the download, from the same live snapshot. - Log viewer in the sidebar: a 500-line in-memory ring fed by a second ANSI-free tracing layer; filter, level, pause, copy. - Update check: crates.io polled twice a day, the console shows an "update available" notice. - Media generations queue up (with position shown) instead of being rejected; progress gained an honest ETA from measured throughput. - Generation gallery under the image/video tabs: lazy PPM thumbnails, click to reopen with seed/steps, delete. - Playground chat is a real dialogue now: history is resent, bubbles, Clear, Cmd/Ctrl+Enter. - "Connect your app" snippets (curl / Python / JS) on the keys page, with the real key inlined right after creation. - One-click benchmark on the Models page via `cortiq bench --json`. - Config export now downloads the actual gateway.toml (comments intact); import accepts TOML (validated, hot-reloaded) and the legacy JSON. - Streaming token counts are exact when the model server provides usage (cortiq-server patch, shipped separately); the gateway fallback now counts content chunks (1/token on cortiq) instead of chars/4. Wizard's Finish lands on the Playground. 37 new strings in all 7 languages.
This commit is contained in:
+102
-3
@@ -184,6 +184,12 @@ pub fn api_routes(admin_token: String) -> Router<SharedState> {
|
||||
.route("/admin/api/requests", get(get_requests))
|
||||
.route("/admin/api/test", post(run_test))
|
||||
.route("/admin/api/test/stream", post(run_test_stream))
|
||||
.route("/admin/api/logs", get(get_logs))
|
||||
.route("/admin/api/system", get(get_system))
|
||||
.route("/admin/api/restart", post(restart_gateway))
|
||||
.route("/admin/api/bench", get(bench_list).post(bench_start))
|
||||
.route("/admin/api/config/export", get(config_export))
|
||||
.route("/admin/api/config/import", post(config_import))
|
||||
.route("/admin/api/media/models", get(media_models))
|
||||
.route("/admin/api/media", get(media_jobs).post(media_generate))
|
||||
.route(
|
||||
@@ -591,9 +597,13 @@ async fn provider_models(
|
||||
ok(resp)
|
||||
}
|
||||
|
||||
async fn meta() -> ApiResult {
|
||||
async fn meta(State(state): State<SharedState>) -> ApiResult {
|
||||
ok(json!({
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
"update_available": state.gateway_update.lock().unwrap().clone(),
|
||||
"admin_open_public": state
|
||||
.admin_open_public
|
||||
.load(std::sync::atomic::Ordering::Relaxed),
|
||||
"providers": ["openai", "anthropic", "openrouter", "lmstudio", "ollama", "http"],
|
||||
"provider_impl": { "openai": true, "anthropic": true, "openrouter": true, "lmstudio": true, "ollama": true, "http": true },
|
||||
"cost_tiers": ["cheap", "mid", "expensive", "local"],
|
||||
@@ -2029,6 +2039,89 @@ async fn get_requests(State(state): State<SharedState>, Query(q): Query<RecentQu
|
||||
ok(json!({ "requests": items }))
|
||||
}
|
||||
|
||||
// ───────────────────────── system / logs / lifecycle ─────────────────────────
|
||||
|
||||
/// Recent gateway log lines (in-memory ring) — for users who never see the
|
||||
/// terminal: Docker, autostart, a double-clicked exe.
|
||||
async fn get_logs(Query(q): Query<RecentQuery>) -> ApiResult {
|
||||
let limit = q.limit.unwrap_or(300).min(500);
|
||||
ok(json!({ "lines": crate::logbuf::snapshot(limit) }))
|
||||
}
|
||||
|
||||
/// Host metrics snapshot: CPU, RAM, free disk at the models dir, gateway RSS.
|
||||
async fn get_system(State(state): State<SharedState>) -> ApiResult {
|
||||
let snap = state.sysmon.snapshot();
|
||||
ok(serde_json::to_value(snap).unwrap_or_else(|_| json!({})))
|
||||
}
|
||||
|
||||
/// Restart the gateway in place: stop managed model servers, spawn a fresh
|
||||
/// copy of ourselves (it retries the bind while our port frees), then exit.
|
||||
/// In Docker we only exit — we are PID 1, a spawned child would die with the
|
||||
/// container anyway; the container's restart policy brings the gateway back.
|
||||
async fn restart_gateway(State(state): State<SharedState>) -> ApiResult {
|
||||
state.cmf.stop().await;
|
||||
let in_docker = std::path::Path::new("/.dockerenv").exists();
|
||||
let exe = std::env::current_exe().map_err(|e| ApiError::bad(e.to_string()))?;
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
tokio::spawn(async move {
|
||||
// let the 200 reach the browser first
|
||||
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
|
||||
if !in_docker {
|
||||
let _ = std::process::Command::new(exe)
|
||||
.args(args)
|
||||
.env("CORTIQ_BIND_RETRY", "1")
|
||||
.spawn();
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
}
|
||||
std::process::exit(0);
|
||||
});
|
||||
ok(json!({ "ok": true }))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BenchBody {
|
||||
model: String,
|
||||
}
|
||||
|
||||
/// Benchmark a local .cmf with `cortiq bench --json` — real tok/s on this box.
|
||||
async fn bench_start(State(state): State<SharedState>, Json(b): Json<BenchBody>) -> ApiResult {
|
||||
let live = state.live();
|
||||
crate::bench::start(state.bench.clone(), &live.cfg.cmf, b.model).map_err(ApiError::bad)?;
|
||||
ok(json!({ "ok": true }))
|
||||
}
|
||||
|
||||
async fn bench_list(State(state): State<SharedState>) -> ApiResult {
|
||||
ok(json!({ "runs": state.bench.list() }))
|
||||
}
|
||||
|
||||
/// The config file as-is (comments included) for download. Secrets live in
|
||||
/// secrets.toml and are deliberately NOT part of the export.
|
||||
async fn config_export(State(state): State<SharedState>) -> Response {
|
||||
match tokio::fs::read_to_string(&state.config_path).await {
|
||||
Ok(text) => (
|
||||
[
|
||||
(header::CONTENT_TYPE, "application/toml".to_string()),
|
||||
(
|
||||
header::CONTENT_DISPOSITION,
|
||||
"attachment; filename=\"gateway.toml\"".to_string(),
|
||||
),
|
||||
],
|
||||
text,
|
||||
)
|
||||
.into_response(),
|
||||
Err(e) => ApiError::bad(format!("cannot read config: {e}")).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace the whole config from an uploaded TOML. Validated and hot-reloaded
|
||||
/// exactly like a panel edit; on any error the current config stays untouched.
|
||||
async fn config_import(State(state): State<SharedState>, body: String) -> ApiResult {
|
||||
let cfg: crate::config::Config =
|
||||
toml::from_str(&body).map_err(|e| ApiError::bad(format!("invalid TOML: {e}")))?;
|
||||
state.reload(cfg)?;
|
||||
ok(json!({ "ok": true }))
|
||||
}
|
||||
|
||||
// ───────────────────────────── media generation ─────────────────────────────
|
||||
|
||||
/// Local `.cmf` files with their modality (chat / image / video) — the
|
||||
@@ -2056,7 +2149,13 @@ async fn media_generate(
|
||||
|
||||
async fn media_job_status(State(state): State<SharedState>, Path(job): Path<String>) -> ApiResult {
|
||||
match state.media.get(&job) {
|
||||
Some(j) => ok(serde_json::to_value(j).unwrap_or_else(|_| json!({}))),
|
||||
Some(j) => {
|
||||
let mut v = serde_json::to_value(&j).unwrap_or_else(|_| json!({}));
|
||||
if j.state == "queued" {
|
||||
v["queue_position"] = json!(state.media.queue_position(&job));
|
||||
}
|
||||
ok(v)
|
||||
}
|
||||
None => Err(ApiError::not_found(format!("media job '{job}' not found"))),
|
||||
}
|
||||
}
|
||||
@@ -2065,7 +2164,7 @@ async fn media_cancel(State(state): State<SharedState>, Path(job): Path<String>)
|
||||
if state.media.get(&job).is_none() {
|
||||
return Err(ApiError::not_found(format!("media job '{job}' not found")));
|
||||
}
|
||||
ok(json!({ "ok": state.media.cancel(&job) }))
|
||||
ok(json!({ "ok": state.media.cancel(&job, &crate::media::media_dir()) }))
|
||||
}
|
||||
|
||||
async fn media_delete(State(state): State<SharedState>, Path(job): Path<String>) -> ApiResult {
|
||||
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
//! One-click local model benchmark: run `cortiq bench <model> --json` and keep
|
||||
//! the parsed result per model, so the Models page can show real tok/s numbers
|
||||
//! measured on this machine.
|
||||
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[derive(Clone, Serialize)]
|
||||
pub struct BenchRun {
|
||||
pub model: String, // file stem under models_dir
|
||||
pub state: String, // running | done | error
|
||||
pub log: Vec<String>,
|
||||
pub tok_s: Option<f64>,
|
||||
/// The full JSON object the benchmark printed, passed through untouched.
|
||||
pub raw: Option<serde_json::Value>,
|
||||
pub started: u64,
|
||||
pub finished: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct BenchStore {
|
||||
runs: Mutex<HashMap<String, BenchRun>>,
|
||||
busy: Mutex<bool>,
|
||||
}
|
||||
|
||||
fn now() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
/// tok/s out of the benchmark's JSON: first numeric field whose key mentions
|
||||
/// tokens-per-second (the exact name has varied across cortiq releases).
|
||||
fn extract_tok_s(v: &serde_json::Value) -> Option<f64> {
|
||||
let obj = v.as_object()?;
|
||||
for want in ["tok_s", "tok_per_s", "tokens_per_sec", "decode_tok_s"] {
|
||||
if let Some(n) = obj.get(want).and_then(|x| x.as_f64()) {
|
||||
return Some(n);
|
||||
}
|
||||
}
|
||||
obj.iter()
|
||||
.find(|(k, val)| {
|
||||
let k = k.to_ascii_lowercase();
|
||||
(k.contains("tok") && (k.contains("/s") || k.contains("_s") || k.contains("per")))
|
||||
&& val.as_f64().is_some()
|
||||
})
|
||||
.and_then(|(_, val)| val.as_f64())
|
||||
}
|
||||
|
||||
impl BenchStore {
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self::default())
|
||||
}
|
||||
pub fn get(&self, model: &str) -> Option<BenchRun> {
|
||||
self.runs.lock().unwrap().get(model).cloned()
|
||||
}
|
||||
pub fn list(&self) -> Vec<BenchRun> {
|
||||
self.runs.lock().unwrap().values().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Kick off a benchmark for `<models_dir>/<model>.cmf`. One at a time — a
|
||||
/// benchmark loads the whole model and owns the machine while it runs.
|
||||
pub fn start(
|
||||
store: Arc<BenchStore>,
|
||||
cmf: &crate::config::CmfCfg,
|
||||
model: String,
|
||||
) -> Result<(), String> {
|
||||
let path = std::path::Path::new(&cmf.models_dir).join(format!("{model}.cmf"));
|
||||
if !path.exists() {
|
||||
return Err(format!("model file not found: {}", path.display()));
|
||||
}
|
||||
{
|
||||
let mut busy = store.busy.lock().unwrap();
|
||||
if *busy {
|
||||
return Err("a benchmark is already running — wait for it to finish".into());
|
||||
}
|
||||
*busy = true;
|
||||
}
|
||||
store.runs.lock().unwrap().insert(
|
||||
model.clone(),
|
||||
BenchRun {
|
||||
model: model.clone(),
|
||||
state: "running".into(),
|
||||
log: Vec::new(),
|
||||
tok_s: None,
|
||||
raw: None,
|
||||
started: now(),
|
||||
finished: None,
|
||||
},
|
||||
);
|
||||
let bin = cmf.cortiq_bin.clone();
|
||||
let gpu = cmf.gpu;
|
||||
tokio::spawn(async move {
|
||||
let out = tokio::process::Command::new(&bin)
|
||||
.arg("bench")
|
||||
.arg(&path)
|
||||
.arg("--json")
|
||||
.env("CMF_GPU", if gpu { "1" } else { "0" })
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
let mut g = store.runs.lock().unwrap();
|
||||
let run = g.get_mut(&model).unwrap();
|
||||
run.finished = Some(now());
|
||||
match out {
|
||||
Ok(o) => {
|
||||
let text = format!(
|
||||
"{}\n{}",
|
||||
String::from_utf8_lossy(&o.stdout),
|
||||
String::from_utf8_lossy(&o.stderr)
|
||||
);
|
||||
// keep a readable tail; the JSON line is parsed separately
|
||||
run.log = text
|
||||
.lines()
|
||||
.map(|l| l.trim())
|
||||
.filter(|l| !l.is_empty() && !l.contains("Metal GPU path"))
|
||||
.map(String::from)
|
||||
.collect::<Vec<_>>();
|
||||
let n = run.log.len();
|
||||
if n > 30 {
|
||||
run.log.drain(0..n - 30);
|
||||
}
|
||||
// the --json contract: one JSON object on its own line
|
||||
let parsed = text
|
||||
.lines()
|
||||
.rev()
|
||||
.map(str::trim)
|
||||
.find(|l| l.starts_with('{') && l.ends_with('}'))
|
||||
.and_then(|l| serde_json::from_str::<serde_json::Value>(l).ok());
|
||||
match parsed {
|
||||
Some(v) if o.status.success() => {
|
||||
run.tok_s = extract_tok_s(&v);
|
||||
run.raw = Some(v);
|
||||
run.state = "done".into();
|
||||
}
|
||||
_ => run.state = "error".into(),
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
run.log.push(format!("✗ failed to start cortiq bench: {e}"));
|
||||
run.state = "error".into();
|
||||
}
|
||||
}
|
||||
drop(g);
|
||||
*store.busy.lock().unwrap() = false;
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn tok_s_from_common_keys() {
|
||||
let v = serde_json::json!({"tok_s": 42.5, "allocs_per_token": 0});
|
||||
assert_eq!(extract_tok_s(&v), Some(42.5));
|
||||
let v = serde_json::json!({"decode_tokens_per_sec": 31.0});
|
||||
assert_eq!(extract_tok_s(&v), Some(31.0));
|
||||
}
|
||||
}
|
||||
+13
-3
@@ -197,14 +197,14 @@ impl Default for CmfRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
/// crates.io `newest_version` for `cortiq-cli`.
|
||||
async fn latest_crates_version() -> Option<String> {
|
||||
/// crates.io `newest_version` for any crate.
|
||||
async fn latest_version_of(krate: &str) -> Option<String> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.ok()?;
|
||||
let body = client
|
||||
.get("https://crates.io/api/v1/crates/cortiq-cli")
|
||||
.get(format!("https://crates.io/api/v1/crates/{krate}"))
|
||||
.header("User-Agent", "cortiq-gateway")
|
||||
.send()
|
||||
.await
|
||||
@@ -216,6 +216,16 @@ async fn latest_crates_version() -> Option<String> {
|
||||
v["crate"]["newest_version"].as_str().map(|s| s.to_string())
|
||||
}
|
||||
|
||||
/// crates.io `newest_version` for `cortiq-cli`.
|
||||
async fn latest_crates_version() -> Option<String> {
|
||||
latest_version_of("cortiq-cli").await
|
||||
}
|
||||
|
||||
/// Newest published gateway version — the "update available" footer check.
|
||||
pub(crate) async fn latest_gateway_version() -> Option<String> {
|
||||
latest_version_of("cortiq-gateway").await
|
||||
}
|
||||
|
||||
/// `<bin> --version` → the trailing version token (e.g. "cortiq 0.1.2" → "0.1.2").
|
||||
pub(crate) fn installed_version(bin: &str) -> Option<String> {
|
||||
let out = std::process::Command::new(bin)
|
||||
|
||||
+51
-1
@@ -132,7 +132,7 @@ fn default_quant() -> String {
|
||||
}
|
||||
|
||||
/// Live state of one conversion job.
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct Job {
|
||||
pub id: String,
|
||||
pub repo: String,
|
||||
@@ -162,14 +162,60 @@ pub struct JobStore {
|
||||
/// Cancel handles, keyed by job id (not serialized). Notifying aborts the
|
||||
/// converter process and cleans up its partial output.
|
||||
cancels: Mutex<HashMap<String, Arc<Notify>>>,
|
||||
/// When set, every state change lands in this JSON file, so the job list —
|
||||
/// and its Register buttons — survives a gateway restart.
|
||||
persist: Mutex<Option<std::path::PathBuf>>,
|
||||
last_save: Mutex<Option<std::time::Instant>>,
|
||||
}
|
||||
|
||||
impl JobStore {
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self::default())
|
||||
}
|
||||
/// Load previously saved jobs and start persisting to `path`. A job that
|
||||
/// was `running` when the gateway died is settled as an error — its process
|
||||
/// is gone; the file it produced (if complete) is still on disk.
|
||||
pub fn attach_persistence(&self, path: std::path::PathBuf) {
|
||||
if let Ok(text) = std::fs::read_to_string(&path) {
|
||||
if let Ok(mut jobs) = serde_json::from_str::<Vec<Job>>(&text) {
|
||||
let mut g = self.jobs.lock().unwrap();
|
||||
for j in jobs.drain(..) {
|
||||
let mut j = j;
|
||||
if j.state == "running" {
|
||||
j.state = "error".into();
|
||||
j.log.push("✗ interrupted by a gateway restart".into());
|
||||
j.finished = Some(now());
|
||||
}
|
||||
g.insert(j.id.clone(), j);
|
||||
}
|
||||
}
|
||||
}
|
||||
*self.persist.lock().unwrap() = Some(path);
|
||||
}
|
||||
fn save(&self, force: bool) {
|
||||
let Some(path) = self.persist.lock().unwrap().clone() else {
|
||||
return;
|
||||
};
|
||||
if !force {
|
||||
let mut last = self.last_save.lock().unwrap();
|
||||
if let Some(t) = *last {
|
||||
if t.elapsed() < std::time::Duration::from_secs(2) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
*last = Some(std::time::Instant::now());
|
||||
}
|
||||
let jobs: Vec<Job> = self.jobs.lock().unwrap().values().cloned().collect();
|
||||
if let Ok(json) = serde_json::to_vec(&jobs) {
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
if std::fs::write(&tmp, json).is_ok() {
|
||||
let _ = std::fs::rename(&tmp, &path);
|
||||
}
|
||||
}
|
||||
}
|
||||
fn insert(&self, job: Job) {
|
||||
self.jobs.lock().unwrap().insert(job.id.clone(), job);
|
||||
self.save(true);
|
||||
}
|
||||
fn set_cancel(&self, id: &str, n: Arc<Notify>) {
|
||||
self.cancels.lock().unwrap().insert(id.to_string(), n);
|
||||
@@ -215,6 +261,7 @@ impl JobStore {
|
||||
j.phase = Some(phase);
|
||||
}
|
||||
}
|
||||
self.save(false);
|
||||
}
|
||||
/// Expected total download size, once the HF tree is known.
|
||||
fn set_total(&self, id: &str, total: u64) {
|
||||
@@ -263,6 +310,7 @@ impl JobStore {
|
||||
j.state = state.into();
|
||||
j.finished = Some(now());
|
||||
}
|
||||
self.save(true);
|
||||
}
|
||||
fn finish(&self, id: &str, ok: bool) {
|
||||
if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
|
||||
@@ -280,6 +328,7 @@ impl JobStore {
|
||||
j.total_bytes = Some(j.total_bytes.unwrap_or(on_disk).max(on_disk));
|
||||
}
|
||||
}
|
||||
self.save(true);
|
||||
}
|
||||
pub fn get(&self, id: &str) -> Option<Job> {
|
||||
self.jobs.lock().unwrap().get(id).cloned()
|
||||
@@ -295,6 +344,7 @@ impl JobStore {
|
||||
cleanup_output(&job.output);
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
self.jobs.lock().unwrap().remove(id);
|
||||
self.save(true);
|
||||
Ok(true)
|
||||
}
|
||||
pub fn list(&self) -> Vec<Job> {
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
//! In-memory ring of recent log lines, so the admin console can show logs to
|
||||
//! users who never see the terminal (Docker, autostart, double-clicked exe).
|
||||
//! Fed by a second `tracing` layer writing plain text (no ANSI) into the ring
|
||||
//! while the normal stdout layer stays colored.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::io;
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
|
||||
const CAP: usize = 500;
|
||||
|
||||
fn ring() -> &'static Arc<Mutex<VecDeque<String>>> {
|
||||
static RING: OnceLock<Arc<Mutex<VecDeque<String>>>> = OnceLock::new();
|
||||
RING.get_or_init(|| Arc::new(Mutex::new(VecDeque::with_capacity(CAP))))
|
||||
}
|
||||
|
||||
/// Last `limit` lines, oldest first.
|
||||
pub fn snapshot(limit: usize) -> Vec<String> {
|
||||
let g = ring().lock().unwrap();
|
||||
g.iter()
|
||||
.skip(g.len().saturating_sub(limit))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn push_line(line: &str) {
|
||||
let line = line.trim_end();
|
||||
if line.is_empty() {
|
||||
return;
|
||||
}
|
||||
let mut g = ring().lock().unwrap();
|
||||
if g.len() >= CAP {
|
||||
g.pop_front();
|
||||
}
|
||||
g.push_back(line.to_string());
|
||||
}
|
||||
|
||||
/// `MakeWriter` for the ring layer: buffers partial writes until a newline.
|
||||
#[derive(Clone, Default)]
|
||||
pub struct RingMakeWriter;
|
||||
|
||||
pub struct RingWriter {
|
||||
buf: String,
|
||||
}
|
||||
|
||||
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for RingMakeWriter {
|
||||
type Writer = RingWriter;
|
||||
fn make_writer(&'a self) -> Self::Writer {
|
||||
RingWriter { buf: String::new() }
|
||||
}
|
||||
}
|
||||
|
||||
impl io::Write for RingWriter {
|
||||
fn write(&mut self, data: &[u8]) -> io::Result<usize> {
|
||||
self.buf.push_str(&String::from_utf8_lossy(data));
|
||||
while let Some(idx) = self.buf.find('\n') {
|
||||
let line: String = self.buf.drain(..=idx).collect();
|
||||
push_line(&line);
|
||||
}
|
||||
Ok(data.len())
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RingWriter {
|
||||
fn drop(&mut self) {
|
||||
if !self.buf.is_empty() {
|
||||
push_line(&self.buf.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
+51
-7
@@ -9,11 +9,13 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
mod admin;
|
||||
mod bench;
|
||||
mod cache;
|
||||
mod cmf_runtime;
|
||||
mod config;
|
||||
mod error;
|
||||
mod import;
|
||||
mod logbuf;
|
||||
mod media;
|
||||
mod model;
|
||||
mod pipeline;
|
||||
@@ -26,6 +28,7 @@ mod routing;
|
||||
mod secrets;
|
||||
mod state;
|
||||
mod stats;
|
||||
mod sysmon;
|
||||
|
||||
use axum::{routing::get, Router};
|
||||
use clap::Parser;
|
||||
@@ -204,12 +207,23 @@ async fn main() -> anyhow::Result<()> {
|
||||
if listen != original_listen {
|
||||
let _ = cfg.save(&config_path);
|
||||
}
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| cfg.log.level.clone().into()),
|
||||
)
|
||||
.init();
|
||||
{
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| cfg.log.level.clone().into());
|
||||
tracing_subscriber::registry()
|
||||
.with(filter)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
// second copy of every line, without ANSI, into the in-memory ring
|
||||
// that backs the admin console's Логи view
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_ansi(false)
|
||||
.with_writer(logbuf::RingMakeWriter),
|
||||
)
|
||||
.init();
|
||||
}
|
||||
|
||||
let admin_enabled = cfg.admin.enabled;
|
||||
let admin_token = resolve_admin_token(&cfg, &args);
|
||||
@@ -217,6 +231,18 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
let state = SharedState::build(cfg, config_path.clone())?;
|
||||
|
||||
// Open admin console reachable from other machines: fine on a laptop,
|
||||
// a footgun on a VPS. Flag it for the UI banner and say it in the log.
|
||||
let host_public = listen.starts_with("0.0.0.0") || listen.starts_with("[::]");
|
||||
if admin_enabled && admin_token.is_empty() && host_public && !is_docker() {
|
||||
state
|
||||
.admin_open_public
|
||||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
tracing::warn!(
|
||||
"admin console is OPEN (no token) and listening on {listen} — anyone on your network can manage this gateway; set a token via --admin-token or [admin].token if this machine is reachable from outside"
|
||||
);
|
||||
}
|
||||
|
||||
// Close the loop with the CMF format (github.com/infosave2007/cmf): install
|
||||
// / update cortiq-cli and run a local `cortiq serve` if configured. Runs in
|
||||
// the background so startup is never blocked by a cargo install or a model
|
||||
@@ -247,7 +273,25 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
let app = app.with_state(state).layer(cors);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(&listen).await?;
|
||||
let listener = {
|
||||
let mut attempt = 0u32;
|
||||
loop {
|
||||
match tokio::net::TcpListener::bind(&listen).await {
|
||||
Ok(l) => break l,
|
||||
Err(e) => {
|
||||
// during self-restart the old process is still letting go of
|
||||
// the port — retry for a few seconds before giving up
|
||||
let retrying = std::env::var("CORTIQ_BIND_RETRY").is_ok();
|
||||
if retrying && attempt < 50 {
|
||||
attempt += 1;
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
continue;
|
||||
}
|
||||
return Err(e.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
tracing::info!("cortiq-gateway listening on {listen}");
|
||||
if admin_enabled {
|
||||
let host_for_url = if listen.starts_with("0.0.0.0:") {
|
||||
|
||||
+151
-61
@@ -184,7 +184,7 @@ pub struct MediaParams {
|
||||
}
|
||||
|
||||
/// Live state of one generation.
|
||||
#[derive(Clone, Serialize)]
|
||||
#[derive(Clone, Serialize, serde::Deserialize)]
|
||||
pub struct MediaJob {
|
||||
pub id: String,
|
||||
pub kind: String, // image | video
|
||||
@@ -195,7 +195,7 @@ pub struct MediaJob {
|
||||
pub steps: u32,
|
||||
pub seed: u64,
|
||||
pub frames: Option<u32>,
|
||||
pub state: String, // running | done | error | cancelled
|
||||
pub state: String, // queued | running | done | error | cancelled
|
||||
pub progress: Option<f32>,
|
||||
pub phase: Option<String>,
|
||||
pub log: Vec<String>,
|
||||
@@ -206,18 +206,67 @@ pub struct MediaJob {
|
||||
pub out_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
/// Everything needed to actually launch a queued job when its turn comes.
|
||||
struct Prepared {
|
||||
bin: String,
|
||||
args: Vec<String>,
|
||||
gpu: bool,
|
||||
dir: PathBuf,
|
||||
out_path: PathBuf,
|
||||
kind: String,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct MediaStore {
|
||||
jobs: Mutex<HashMap<String, MediaJob>>,
|
||||
cancels: Mutex<HashMap<String, Arc<Notify>>>,
|
||||
/// One generation at a time — these models take the whole GPU/RAM budget.
|
||||
busy: Mutex<Option<String>>,
|
||||
/// One generation runs at a time — these models take the whole GPU/RAM
|
||||
/// budget — but new requests queue up instead of being rejected.
|
||||
queue: Mutex<std::collections::VecDeque<String>>,
|
||||
current: Mutex<Option<String>>,
|
||||
prepared: Mutex<HashMap<String, Prepared>>,
|
||||
persist: Mutex<Option<PathBuf>>,
|
||||
}
|
||||
|
||||
impl MediaStore {
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self::default())
|
||||
}
|
||||
/// Load saved generations (their files are still under media_dir) and start
|
||||
/// persisting. Jobs that were queued/running when the gateway died settle
|
||||
/// as errors — the process is gone.
|
||||
pub fn attach_persistence(&self, path: PathBuf) {
|
||||
if let Ok(text) = std::fs::read_to_string(&path) {
|
||||
if let Ok(mut jobs) = serde_json::from_str::<Vec<MediaJob>>(&text) {
|
||||
let mut g = self.jobs.lock().unwrap();
|
||||
for mut j in jobs.drain(..) {
|
||||
if j.state == "running" || j.state == "queued" {
|
||||
j.state = "error".into();
|
||||
j.log.push("✗ interrupted by a gateway restart".into());
|
||||
j.finished = Some(now());
|
||||
}
|
||||
g.insert(j.id.clone(), j);
|
||||
}
|
||||
}
|
||||
}
|
||||
*self.persist.lock().unwrap() = Some(path);
|
||||
}
|
||||
fn save(&self) {
|
||||
let Some(path) = self.persist.lock().unwrap().clone() else {
|
||||
return;
|
||||
};
|
||||
let jobs: Vec<MediaJob> = self.jobs.lock().unwrap().values().cloned().collect();
|
||||
if let Ok(json) = serde_json::to_vec(&jobs) {
|
||||
let tmp = path.with_extension("json.tmp");
|
||||
if std::fs::write(&tmp, json).is_ok() {
|
||||
let _ = std::fs::rename(&tmp, &path);
|
||||
}
|
||||
}
|
||||
}
|
||||
/// How many jobs sit in the queue ahead of `id` (0 = next up).
|
||||
pub fn queue_position(&self, id: &str) -> Option<usize> {
|
||||
self.queue.lock().unwrap().iter().position(|x| x == id)
|
||||
}
|
||||
pub fn list(&self) -> Vec<MediaJob> {
|
||||
let mut v: Vec<MediaJob> = self.jobs.lock().unwrap().values().cloned().collect();
|
||||
v.sort_by_key(|j| std::cmp::Reverse(j.started));
|
||||
@@ -227,30 +276,43 @@ impl MediaStore {
|
||||
pub fn get(&self, id: &str) -> Option<MediaJob> {
|
||||
self.jobs.lock().unwrap().get(id).cloned()
|
||||
}
|
||||
pub fn cancel(&self, id: &str) -> bool {
|
||||
let running = matches!(
|
||||
self.jobs.lock().unwrap().get(id).map(|j| j.state.as_str()),
|
||||
Some("running")
|
||||
);
|
||||
if !running {
|
||||
return false;
|
||||
}
|
||||
if let Some(n) = self.cancels.lock().unwrap().get(id) {
|
||||
n.notify_one();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
pub fn cancel(self: &Arc<Self>, id: &str, media_dir: &Path) -> bool {
|
||||
let state = self
|
||||
.jobs
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(id)
|
||||
.map(|j| j.state.clone())
|
||||
.unwrap_or_default();
|
||||
match state.as_str() {
|
||||
"queued" => {
|
||||
self.queue.lock().unwrap().retain(|x| x != id);
|
||||
self.prepared.lock().unwrap().remove(id);
|
||||
let _ = std::fs::remove_dir_all(media_dir.join(id));
|
||||
self.set_cancelled(id);
|
||||
true
|
||||
}
|
||||
"running" => {
|
||||
if let Some(n) = self.cancels.lock().unwrap().get(id) {
|
||||
n.notify_one();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
/// Remove a settled job and its output directory. Err = still running.
|
||||
pub fn delete(&self, id: &str, media_dir: &Path) -> Result<bool, String> {
|
||||
let job = self.jobs.lock().unwrap().get(id).cloned();
|
||||
let Some(job) = job else { return Ok(false) };
|
||||
if job.state == "running" {
|
||||
if job.state == "running" || job.state == "queued" {
|
||||
return Err("cancel the running generation first".into());
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(media_dir.join(id));
|
||||
self.jobs.lock().unwrap().remove(id);
|
||||
self.save();
|
||||
Ok(true)
|
||||
}
|
||||
fn push_line(&self, id: &str, line: String) {
|
||||
@@ -273,10 +335,12 @@ impl MediaStore {
|
||||
}
|
||||
}
|
||||
}
|
||||
fn finish(&self, id: &str, ok: bool, outputs: Vec<String>, out_bytes: u64) {
|
||||
fn finish(self: &Arc<Self>, id: &str, ok: bool, outputs: Vec<String>, out_bytes: u64) {
|
||||
let mut g = self.jobs.lock().unwrap();
|
||||
if let Some(j) = g.get_mut(id) {
|
||||
if j.state == "cancelled" {
|
||||
drop(g);
|
||||
self.after_slot_freed(id);
|
||||
return;
|
||||
}
|
||||
j.finished = Some(now());
|
||||
@@ -288,25 +352,67 @@ impl MediaStore {
|
||||
}
|
||||
}
|
||||
drop(g);
|
||||
self.clear_busy(id);
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
self.after_slot_freed(id);
|
||||
}
|
||||
fn set_cancelled(&self, id: &str) {
|
||||
fn set_cancelled(self: &Arc<Self>, id: &str) {
|
||||
if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
|
||||
j.state = "cancelled".into();
|
||||
j.finished = Some(now());
|
||||
}
|
||||
self.clear_busy(id);
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
self.after_slot_freed(id);
|
||||
}
|
||||
fn clear_busy(&self, id: &str) {
|
||||
let mut b = self.busy.lock().unwrap();
|
||||
if b.as_deref() == Some(id) {
|
||||
*b = None;
|
||||
fn after_slot_freed(self: &Arc<Self>, id: &str) {
|
||||
{
|
||||
let mut cur = self.current.lock().unwrap();
|
||||
if cur.as_deref() == Some(id) {
|
||||
*cur = None;
|
||||
}
|
||||
}
|
||||
self.cancels.lock().unwrap().remove(id);
|
||||
self.save();
|
||||
promote_next(self.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the next queued generation, if the slot is free.
|
||||
fn promote_next(store: Arc<MediaStore>) {
|
||||
let next = {
|
||||
let mut cur = store.current.lock().unwrap();
|
||||
if cur.is_some() {
|
||||
return;
|
||||
}
|
||||
let Some(id) = store.queue.lock().unwrap().pop_front() else {
|
||||
return;
|
||||
};
|
||||
*cur = Some(id.clone());
|
||||
id
|
||||
};
|
||||
let Some(p) = store.prepared.lock().unwrap().remove(&next) else {
|
||||
// job was cancelled while queued and slipped through — free the slot
|
||||
*store.current.lock().unwrap() = None;
|
||||
return;
|
||||
};
|
||||
if let Some(j) = store.jobs.lock().unwrap().get_mut(&next) {
|
||||
j.state = "running".into();
|
||||
j.phase = Some("starting".into());
|
||||
j.started = now();
|
||||
}
|
||||
store.save();
|
||||
let cancel = Arc::new(Notify::new());
|
||||
store
|
||||
.cancels
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(next.clone(), cancel.clone());
|
||||
let s = store.clone();
|
||||
tokio::spawn(async move {
|
||||
run_generation(
|
||||
s, next, p.bin, p.args, p.gpu, p.dir, p.out_path, p.kind, cancel,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
/// Progress markers in `cortiq imagine`/`animate` stdout. The exact wording has
|
||||
/// varied between releases, so accept anything that carries an X/Y pair near a
|
||||
/// step/frame word, a bare percentage, or the converter-style `@PROGRESS`.
|
||||
@@ -432,22 +538,10 @@ pub fn start_generation(
|
||||
if !model_path.exists() {
|
||||
return Err(format!("model file not found: {}", model_path.display()));
|
||||
}
|
||||
{
|
||||
let mut busy = store.busy.lock().unwrap();
|
||||
if let Some(running) = busy.as_ref() {
|
||||
return Err(format!(
|
||||
"a generation is already running (job {running}) — these models need the whole machine; wait or cancel it"
|
||||
));
|
||||
}
|
||||
*busy = Some(String::new()); // reserved; real id set below
|
||||
}
|
||||
|
||||
let id = gen_media_id();
|
||||
*store.busy.lock().unwrap() = Some(id.clone());
|
||||
|
||||
let dir = media_dir().join(&id);
|
||||
if let Err(e) = std::fs::create_dir_all(&dir) {
|
||||
store.clear_busy(&id);
|
||||
return Err(format!("cannot create media dir {}: {e}", dir.display()));
|
||||
}
|
||||
|
||||
@@ -553,9 +647,9 @@ pub fn start_generation(
|
||||
steps,
|
||||
seed,
|
||||
frames: (kind == "video").then_some(frames),
|
||||
state: "running".into(),
|
||||
state: "queued".into(),
|
||||
progress: None,
|
||||
phase: Some("starting".into()),
|
||||
phase: None,
|
||||
log: vec![format!("→ cortiq {}", args.join(" "))],
|
||||
started: now(),
|
||||
finished: None,
|
||||
@@ -563,25 +657,21 @@ pub fn start_generation(
|
||||
out_bytes: None,
|
||||
},
|
||||
);
|
||||
let cancel = Arc::new(Notify::new());
|
||||
store
|
||||
.cancels
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(id.clone(), cancel.clone());
|
||||
|
||||
let bin = cmf.cortiq_bin.clone();
|
||||
let gpu = cmf.gpu;
|
||||
let ret = id.clone();
|
||||
let dir_clone = dir.clone();
|
||||
let kind_owned = kind.to_string();
|
||||
tokio::spawn(async move {
|
||||
run_generation(
|
||||
store, id, bin, args, gpu, dir_clone, out_path, kind_owned, cancel,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
Ok(ret)
|
||||
store.prepared.lock().unwrap().insert(
|
||||
id.clone(),
|
||||
Prepared {
|
||||
bin: cmf.cortiq_bin.clone(),
|
||||
args,
|
||||
gpu: cmf.gpu,
|
||||
dir,
|
||||
out_path,
|
||||
kind: kind.to_string(),
|
||||
},
|
||||
);
|
||||
store.queue.lock().unwrap().push_back(id.clone());
|
||||
store.save();
|
||||
promote_next(store.clone());
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
|
||||
+12
-2
@@ -738,6 +738,7 @@ fn tap_stream(
|
||||
let mut completion_tokens = 0u32;
|
||||
let mut got_usage = false;
|
||||
let mut delta_chars = 0usize;
|
||||
let mut delta_chunks = 0u32;
|
||||
let mut errored: Option<String> = None;
|
||||
|
||||
while let Some(item) = inner.next().await {
|
||||
@@ -774,6 +775,7 @@ fn tap_stream(
|
||||
v["choices"][0]["delta"]["content"].as_str()
|
||||
{
|
||||
delta_chars += c.chars().count();
|
||||
delta_chunks += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -789,7 +791,11 @@ fn tap_stream(
|
||||
// estimate (~4 chars/token), flagged `cortiq_estimated` so nobody
|
||||
// mistakes it for the tokenizer's truth.
|
||||
if done_marker && !got_usage {
|
||||
let est_completion = (delta_chars / 4) as u32;
|
||||
let est_completion = if delta_chunks > 0 {
|
||||
delta_chunks
|
||||
} else {
|
||||
(delta_chars / 4) as u32
|
||||
};
|
||||
let est_prompt = (prompt_chars / 4) as u32;
|
||||
let evt = serde_json::json!({
|
||||
"object": "chat.completion.chunk",
|
||||
@@ -812,7 +818,11 @@ fn tap_stream(
|
||||
|
||||
if !got_usage && completion_tokens == 0 {
|
||||
// stream ended without a [DONE] marker — estimate for the stats
|
||||
completion_tokens = (delta_chars / 4) as u32;
|
||||
completion_tokens = if delta_chunks > 0 {
|
||||
delta_chunks
|
||||
} else {
|
||||
(delta_chars / 4) as u32
|
||||
};
|
||||
prompt_tokens = (prompt_chars / 4) as u32;
|
||||
}
|
||||
let cost = provider.price(prompt_tokens, completion_tokens);
|
||||
|
||||
@@ -73,6 +73,15 @@ pub struct AppState {
|
||||
pub imports: Arc<crate::import::JobStore>,
|
||||
/// Image/video generation jobs (cortiq imagine / animate).
|
||||
pub media: Arc<crate::media::MediaStore>,
|
||||
/// One-click local model benchmarks (cortiq bench --json).
|
||||
pub bench: Arc<crate::bench::BenchStore>,
|
||||
/// Host metrics sampler (CPU / RAM / disk) for the dashboard.
|
||||
pub sysmon: Arc<crate::sysmon::SysMon>,
|
||||
/// Newest cortiq-gateway version on crates.io, when newer than ours.
|
||||
pub gateway_update: Arc<std::sync::Mutex<Option<String>>>,
|
||||
/// Admin console reachable from other machines without a token — the
|
||||
/// footgun banner flag (set at startup, shown by the SPA).
|
||||
pub admin_open_public: std::sync::atomic::AtomicBool,
|
||||
/// Managed local CMF model server (install/update/spawn lifecycle + status).
|
||||
pub cmf: Arc<crate::cmf_runtime::CmfRuntime>,
|
||||
pub cache: Arc<SemanticCache>,
|
||||
@@ -116,6 +125,29 @@ impl SharedState {
|
||||
let promotion = crate::promotion::Promotion::new(promotion_cfg(&cfg.shadow));
|
||||
let imports = crate::import::JobStore::new();
|
||||
let media = crate::media::MediaStore::new();
|
||||
let bench = crate::bench::BenchStore::new();
|
||||
// jobs survive restarts: files under <data>/, reconciled on load
|
||||
let data = crate::config::data_dir();
|
||||
let _ = std::fs::create_dir_all(&data);
|
||||
imports.attach_persistence(data.join("jobs-import.json"));
|
||||
media.attach_persistence(data.join("jobs-media.json"));
|
||||
let sysmon = crate::sysmon::spawn(cfg.cmf.models_dir.clone());
|
||||
let gateway_update: Arc<std::sync::Mutex<Option<String>>> =
|
||||
Arc::new(std::sync::Mutex::new(None));
|
||||
{
|
||||
// check crates.io for a newer gateway once at start and twice a day
|
||||
let slot = gateway_update.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if let Some(latest) = crate::cmf_runtime::latest_gateway_version().await {
|
||||
let newer =
|
||||
crate::cmf_runtime::version_lt(env!("CARGO_PKG_VERSION"), &latest);
|
||||
*slot.lock().unwrap() = newer.then_some(latest);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(12 * 3600)).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
let cmf = crate::cmf_runtime::CmfRuntime::new();
|
||||
let cache = SemanticCache::new(&cfg.cache);
|
||||
let router_status = RouterLastStatus::default();
|
||||
@@ -126,6 +158,10 @@ impl SharedState {
|
||||
promotion,
|
||||
imports,
|
||||
media,
|
||||
bench,
|
||||
sysmon,
|
||||
gateway_update,
|
||||
admin_open_public: std::sync::atomic::AtomicBool::new(false),
|
||||
cmf,
|
||||
cache,
|
||||
secrets,
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
//! Host metrics sampler for the admin dashboard: CPU load, RAM, free disk at
|
||||
//! the models dir, and the gateway's own RSS. One background task refreshes a
|
||||
//! snapshot every 2 s; the API endpoint only reads it.
|
||||
|
||||
use serde::Serialize;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
#[derive(Clone, Default, Serialize)]
|
||||
pub struct Snapshot {
|
||||
pub ts: u64,
|
||||
/// Whole-machine CPU usage 0..100 (average over all cores).
|
||||
pub cpu_percent: f32,
|
||||
pub cores: usize,
|
||||
pub load1: f64,
|
||||
pub ram_total: u64,
|
||||
pub ram_used: u64,
|
||||
/// Disk that holds models_dir.
|
||||
pub disk_total: u64,
|
||||
pub disk_free: u64,
|
||||
/// Resident memory of the gateway process itself.
|
||||
pub process_rss: u64,
|
||||
}
|
||||
|
||||
pub struct SysMon {
|
||||
snap: Mutex<Snapshot>,
|
||||
}
|
||||
|
||||
impl SysMon {
|
||||
pub fn snapshot(&self) -> Snapshot {
|
||||
self.snap.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn now() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
/// Start the sampler. `models_dir` picks which disk's free space is reported.
|
||||
pub fn spawn(models_dir: String) -> Arc<SysMon> {
|
||||
let mon = Arc::new(SysMon {
|
||||
snap: Mutex::new(Snapshot::default()),
|
||||
});
|
||||
let ret = mon.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
use sysinfo::{Disks, ProcessesToUpdate, System};
|
||||
let mut sys = System::new();
|
||||
let pid = sysinfo::Pid::from_u32(std::process::id());
|
||||
let models_abs = std::fs::canonicalize(&models_dir)
|
||||
.unwrap_or_else(|_| std::path::PathBuf::from(&models_dir));
|
||||
loop {
|
||||
sys.refresh_cpu_usage();
|
||||
sys.refresh_memory();
|
||||
let _ = sys.refresh_processes(ProcessesToUpdate::Some(&[pid]), true);
|
||||
let disks = Disks::new_with_refreshed_list();
|
||||
// the disk whose mount point is the longest prefix of models_dir
|
||||
let mut best: Option<(usize, u64, u64)> = None;
|
||||
for d in disks.list() {
|
||||
let mp = d.mount_point();
|
||||
if models_abs.starts_with(mp) {
|
||||
let len = mp.as_os_str().len();
|
||||
if best.map(|(l, _, _)| len > l).unwrap_or(true) {
|
||||
best = Some((len, d.total_space(), d.available_space()));
|
||||
}
|
||||
}
|
||||
}
|
||||
let (disk_total, disk_free) = best.map(|(_, t, f)| (t, f)).unwrap_or((0, 0));
|
||||
let snap = Snapshot {
|
||||
ts: now(),
|
||||
cpu_percent: sys.global_cpu_usage(),
|
||||
cores: sys.cpus().len(),
|
||||
load1: System::load_average().one,
|
||||
ram_total: sys.total_memory(),
|
||||
ram_used: sys.used_memory(),
|
||||
disk_total,
|
||||
disk_free,
|
||||
process_rss: sys.process(pid).map(|p| p.memory()).unwrap_or(0),
|
||||
};
|
||||
*mon.snap.lock().unwrap() = snap;
|
||||
std::thread::sleep(std::time::Duration::from_secs(2));
|
||||
}
|
||||
});
|
||||
ret
|
||||
}
|
||||
Reference in New Issue
Block a user