Luiz Paulo Carvalho преди 3 месеца
ревизия
7824387779
променени са 3 файла, в които са добавени 284 реда и са изтрити 0 реда
  1. 8 0
      Cargo.toml
  2. 64 0
      README.md
  3. 212 0
      src/main.rs

+ 8 - 0
Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "estagiario"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+clap = { version = "4", features = ["derive"] }
+home = "0.5"

+ 64 - 0
README.md

@@ -0,0 +1,64 @@
+# estagiário
+
+A Rust tool that binary-searches for the optimal GPU layer count (-ngl) when running llama-server with a GGUF model, then launches the server with that configuration.
+
+## How it works
+
+- Tests different -ngl values (0..255) using binary search.
+- For each candidate:
+  - Spawns llama-server with that -ngl.
+  - Checks if it outputs the expected response.
+- Narrows down the highest valid -ngl.
+- Spawns the final llama-server with best_fit - 1 to stay safe.
+
+## Requirements
+
+- Rust (edition 2024)
+- llama-server in $PATH
+- GGUF model(s) under ~/guff_models
+
+## Build and run
+
+- Build:
+  - cargo build --release
+- Run:
+  - ./target/release/estagiario
+
+## CLI arguments
+
+Run with --help for full list:
+
+- -c, --ctx <CTX>
+  - Context size (in 1024-byte units). Default: 32.
+
+- -t, --threads <THREADS>
+  - Threads to pass to llama-server. Default: 32.
+- --port <PORT>
+  - Port for llama-server. Default: 11434.
+- --host <HOST>
+  - Host to bind. Default: 0.0.0.0.
+- -b, --batch <BATCH>
+  - CPU offload batch size. Default: 512.
+- -u, --ubatch <UBATCH>
+  - GPU offload batch size. Default: 256.
+- -v, --verbose
+  - Print llama-server logs instead of discarding them.
+
+## Model selection
+
+- The tool scans ~/guff_models for .gguf files.
+- Lists them as:
+  - a) model1.gguf
+  - b) model2.gguf
+  - ...
+- Prompts for a letter to select.
+
+## Behavior notes
+
+- During binary search:
+  - Each test run’s stdout/stderr is discarded.
+  - Results are printed to stderr as:
+    - <ngl> [ok] or <ngl> [fail]
+- Final server:
+  - Runs until interrupted (Ctrl+C).
+  - Use --verbose to see its logs.

+ 212 - 0
src/main.rs

@@ -0,0 +1,212 @@
+use std::{
+    io::{self, BufRead, Write},
+    process,
+    thread::sleep,
+    time::Duration,
+};
+
+use clap::Parser;
+
+#[derive(Parser)]
+#[command(about = "Binary search for best NGL and start llama-server")]
+struct Args {
+    #[arg(short, long, default_value_t = 32)]
+    ctx: u32,
+
+ 
+    #[arg(short, long, default_value_t = 32)]
+    threads: u32,
+
+    #[arg(short, long, default_value_t = 11434)]
+    port: u16,
+
+    #[arg(long, default_value = "0.0.0.0")]
+    host: String,
+
+    #[arg(short, long, default_value_t = 512)]
+    batch: u32,
+
+    #[arg(short, long, default_value_t = 256)]
+    ubatch: u32,
+
+    #[arg(short)]
+    verbose: bool,
+}
+
+fn main() {
+    let args = Args::parse();
+    let ctx_size = args.ctx * 1024;
+
+    let model_path = select_model();
+
+    let best_fit = (u8::MIN..u8::MAX)
+        .collect::<Vec<_>>()
+        .binary_search_by(|x| {
+            if test_ngl(*x, &args, ctx_size, &model_path) {
+                std::cmp::Ordering::Less
+            } else {
+                std::cmp::Ordering::Greater
+            }
+        })
+        .unwrap_or_else(|x| x - 1);
+
+    eprintln!(
+        "final spawn with ngl={} ctx={}x1024 ({} bytes)",
+        best_fit,
+        args.ctx,
+        ctx_size
+    );
+
+    sleep(Duration::from_secs(1));
+
+    let mut child = spawn_final_server(best_fit as u8 - 1, &args, ctx_size, &model_path);
+
+    child.wait().expect("Sempre vou matar com control+c");
+}
+
+fn select_model() -> String {
+    let home_dir = home::home_dir().expect("home dir not found");
+    let model_dir = home_dir.join("guff_models");
+
+    let entries = match std::fs::read_dir(&model_dir) {
+        Ok(entries) => entries,
+        Err(e) => {
+            eprintln!("Error: cannot access '{}': {}", model_dir.display(), e);
+            std::process::exit(1);
+        }
+    };
+
+    let mut models: Vec<std::path::PathBuf> = entries
+        .filter_map(|e| e.ok())
+        .filter(|e| {
+            e.file_name()
+                .to_string_lossy()
+                .to_lowercase()
+                .ends_with(".gguf")
+        })
+        .map(|e| e.path())
+        .collect();
+
+    models.sort_by(|a, b| {
+        let a_name = a.file_name().unwrap_or_default();
+        let b_name = b.file_name().unwrap_or_default();
+        a_name.cmp(b_name)
+    });
+
+    if models.is_empty() {
+        eprintln!("No .gguf files found in {}", model_dir.display());
+        std::process::exit(1);
+    }
+
+    println!("\nAvailable models:");
+    for (i, path) in models.iter().enumerate() {
+        let label = (b'a' + i as u8) as char;
+        let name = path
+            .file_name()
+            .unwrap_or_default()
+            .to_string_lossy();
+        println!("  {label}) {name}");
+    }
+
+    loop {
+        let upper = (b'a' + models.len() as u8 - 1) as char;
+        print!("\nSelect model [a-{upper}]: ");
+        io::stdout().flush().unwrap();
+
+        let mut input = String::new();
+        io::stdin().read_line(&mut input).unwrap();
+
+        let c = input.trim().chars().next().unwrap_or('a');
+        let idx = match c.to_ascii_lowercase() {
+            c if c as u8 >= b'a' && (c as u8) < { b'a' + models.len() as u8 } => {
+                c as u8 - b'a'
+            }
+            _ => continue,
+        };
+
+        return models[idx as usize]
+            .to_string_lossy()
+            .into_owned();
+    }
+}
+
+fn test_ngl(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> bool {
+    let mut child = spawn_llama_server(ngl, args, ctx_size, model_path);
+
+    let found = {
+        let handle = child.stderr.take().expect("stderr not available");
+        let reader = io::BufReader::new(handle);
+        reader
+            .lines()
+            .any(|line| {
+                if let Ok(l) = line {
+                    if l.contains("You are a helpful assistant") {
+                        return true;
+                    }
+                }
+                false
+            })
+    };
+
+    let _ = child.kill();
+    let _ = child.wait();
+
+    sleep(Duration::from_secs(1));
+
+    if found {
+        eprintln!("{ngl} [ok]");
+        true
+    } else {
+        eprintln!("{ngl} [fail]");
+        false
+    }
+}
+
+fn spawn_llama_server(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> process::Child {
+    process::Command::new("llama-server")
+        .arg("-m")
+        .arg(model_path)
+        .arg("-t")
+        .arg(args.threads.to_string())
+        .arg("-c")
+        .arg(ctx_size.to_string())
+        .arg("--port")
+        .arg(args.port.to_string())
+        .arg("--host")
+        .arg(&args.host)
+        .arg("-ngl")
+        .arg(&ngl.to_string())
+        .arg("-b")
+        .arg(args.batch.to_string())
+        .arg("-ub")
+        .arg(args.ubatch.to_string())
+        .arg("--flash-attn")
+        .arg("on")
+        .stdout(std::process::Stdio::piped())
+        .stderr(std::process::Stdio::piped())
+        .spawn()
+        .expect("O spawn precisa sempre acontecer")
+}
+
+ fn spawn_final_server(ngl: u8, args: &Args, ctx_size: u32, model_path: &str) -> process::Child {
+    let mut cmd = process::Command::new("llama-server");
+    cmd.arg("-m").arg(model_path)
+        .arg("-t").arg(args.threads.to_string())
+        .arg("-c").arg(ctx_size.to_string())
+        .arg("--port").arg(args.port.to_string())
+        .arg("--host").arg(&args.host)
+        .arg("-ngl").arg(&ngl.to_string())
+        .arg("-b").arg(args.batch.to_string())
+        .arg("-ub").arg(args.ubatch.to_string())
+        .arg("--flash-attn").arg("on");
+
+    if args.verbose {
+        cmd.stdout(std::process::Stdio::inherit())
+            .stderr(std::process::Stdio::inherit());
+    } else {
+        cmd.stdout(std::process::Stdio::null())
+            .stderr(std::process::Stdio::null());
+    }
+
+    cmd.spawn().expect("O spawn precisa sempre acontecer")
+}