Merge pull request #9 from tilesprivacy/feat/cli-server-commands

Added server commands

Author: madclaws

.gitignore

@@ -1 +1,2 @@
 /target
+.tiles_dev

scripts/install.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-ENV="dev" # prod is another env, try taking it from github env
+ENV="prod" # prod is another env, try taking it from github env
 REPO="tilesprivacy/tilekit" 
 # VERSION="${TILES_VERSION:-latest}"       
 VERSION="0.1.0"       
@@ -39,7 +39,6 @@ fi
 log "⬇️  Downloading Tiles (${VERSION}) for ${ARCH}-${OS}..."
 
 
-
 if [[ "$ENV" == "prod" ]]; then
   TAR_URL="https://github.com/${REPO}/releases/download/${VERSION}/tiles-v${VERSION}-${ARCH}-${OS}.tar.gz"
   curl -fsSL -o "${TMPDIR}/tiles.tar.gz" "$TAR_URL"

src/commands/mod.rs

@@ -20,3 +20,11 @@ pub async fn run(modelfile: &str) {
 pub fn check_health() {
     health::check_health();
 }
+
+pub fn start_server() {
+    let _ = mlx::start_server_daemon();
+}
+
+pub fn stop_server() {
+    let _ = mlx::stop_server_daemon();
+}

src/main.rs

@@ -1,9 +1,9 @@
 use std::error::Error;
 
-use clap::{Parser, Subcommand};
+use clap::{Args, Parser, Subcommand};
 mod commands;
 #[derive(Debug, Parser)]
-#[command(name = "tilekit")]
+#[command(name = "tiles")]
 #[command(version, about = "Run, fine-tune models locally with Modelfile", long_about = None)]
 struct Cli {
     #[command(subcommand)]
@@ -17,8 +17,27 @@ enum Commands {
 
     /// Checks the status of dependencies
     Health,
+
+    /// start or stop the daemon server
+    Server(ServerArgs),
+}
+
+#[derive(Debug, Args)]
+#[command(args_conflicts_with_subcommands = true)]
+#[command(flatten_help = true)]
+struct ServerArgs {
+    #[command(subcommand)]
+    command: Option<ServerCommands>,
 }
 
+#[derive(Debug, Subcommand)]
+enum ServerCommands {
+    /// Start the py server as a daemon
+    Start,
+
+    /// Stops the daemon py server
+    Stop,
+}
 #[tokio::main]
 pub async fn main() -> Result<(), Box<dyn Error>> {
     let cli = Cli::parse();
@@ -29,6 +48,11 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
         Commands::Health => {
             commands::check_health();
         }
+        Commands::Server(server) => match server.command {
+            Some(ServerCommands::Start) => commands::start_server(),
+            Some(ServerCommands::Stop) => commands::stop_server(),
+            _ => println!("Expected start or stop"),
+        },
     }
     Ok(())
 }

src/runner/mlx.rs

@@ -3,6 +3,7 @@ use reqwest::Client;
 use serde_json::{Value, json};
 use std::io::Write;
 use std::path::PathBuf;
+use std::process::Stdio;
 use std::{env, fs};
 use std::{io, process::Command};
 
@@ -72,6 +73,52 @@ fn run_model_by_sub_process(modelfile: Modelfile) {
     }
 }
 
+#[allow(clippy::zombie_processes)]
+pub fn start_server_daemon() -> Result<()> {
+    // check if the server is running
+    // start server as a child process
+    // save the pid in a file under ~/.config/tiles/server_pid
+    let config_dir = get_config_dir()?;
+    let server_dir = get_server_dir()?;
+    let pid_file = config_dir.join("server.pid");
+    if pid_file.exists() {
+        eprintln!("Server is already running");
+        return Ok(());
+    }
+
+    let child = Command::new("uv")
+        .args([
+            "run",
+            "--project",
+            server_dir.to_str().unwrap(),
+            "python",
+            "-m",
+            "server.main",
+        ])
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()
+        .expect("failed to start server");
+    fs::create_dir_all(&config_dir).context("Failed to create config directory")?;
+    std::fs::write(pid_file, child.id().to_string()).unwrap();
+    println!("Server started with PID {}", child.id());
+    Ok(())
+}
+
+pub fn stop_server_daemon() -> Result<()> {
+    let pid_file = get_config_dir()?.join("server.pid");
+
+    if !pid_file.exists() {
+        eprintln!("Server is not running");
+        return Ok(());
+    }
+
+    let pid = std::fs::read_to_string(&pid_file).unwrap();
+    Command::new("kill").arg(pid.trim()).status().unwrap();
+    std::fs::remove_file(pid_file).unwrap();
+    println!("Server stopped.");
+    Ok(())
+}
 async fn run_model_with_server(modelfile: Modelfile) -> reqwest::Result<()> {
     let stdin = io::stdin();
     let mut stdout = io::stdout();
@@ -157,19 +204,8 @@ async fn chat(input: &str, model_name: &str) -> Result<String, String> {
 }
 
 fn get_memory_path() -> Result<String> {
-    let home_dir = env::home_dir().context("Failed to fetch $HOME")?;
-    let config_dir = match env::var("XDG_CONFIG_HOME") {
-        Ok(val) => PathBuf::from(val),
-        Err(_err) => home_dir.join(".config"),
-    };
-
-    let data_dir = match env::var("XDG_DATA_HOME") {
-        Ok(val) => PathBuf::from(val),
-        Err(_err) => home_dir.join(".local/share"),
-    };
-
-    let tiles_config_dir = config_dir.join("tiles");
-    let tiles_data_dir = data_dir.join("tiles");
+    let tiles_config_dir = get_config_dir()?;
+    let tiles_data_dir = get_data_dir()?;
     let mut is_memory_path_found: bool = false;
     let mut memory_path: String = String::from("");
     if tiles_config_dir.is_dir()
@@ -193,3 +229,44 @@ fn get_memory_path() -> Result<String> {
         Ok(memory_path.to_string_lossy().to_string())
     }
 }
+
+fn get_server_dir() -> Result<PathBuf> {
+    if cfg!(debug_assertions) {
+        let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?;
+        Ok(base_dir.join("server"))
+    } else {
+        let home_dir = env::home_dir().context("Failed to fetch $HOME")?;
+        let data_dir = match env::var("XDG_DATA_HOME") {
+            Ok(val) => PathBuf::from(val),
+            Err(_err) => home_dir.join(".local/share"),
+        };
+        Ok(data_dir.join("tiles/server"))
+    }
+}
+fn get_config_dir() -> Result<PathBuf> {
+    if cfg!(debug_assertions) {
+        let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?;
+        Ok(base_dir.join(".tiles_dev/tiles"))
+    } else {
+        let home_dir = env::home_dir().context("Failed to fetch $HOME")?;
+        let config_dir = match env::var("XDG_CONFIG_HOME") {
+            Ok(val) => PathBuf::from(val),
+            Err(_err) => home_dir.join(".config"),
+        };
+        Ok(config_dir.join("tiles"))
+    }
+}
+
+fn get_data_dir() -> Result<PathBuf> {
+    if cfg!(debug_assertions) {
+        let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?;
+        Ok(base_dir.join(".tiles_dev/tiles"))
+    } else {
+        let home_dir = env::home_dir().context("Failed to fetch $HOME")?;
+        let data_dir = match env::var("XDG_DATA_HOME") {
+            Ok(val) => PathBuf::from(val),
+            Err(_err) => home_dir.join(".local/share"),
+        };
+        Ok(data_dir.join("tiles"))
+    }
+}