offline_intelligence/model_runtime/

safetensors_runtime.rs

//! Safetensors Runtime Adapter
//!
//! Adapter for Safetensors format models using Candle framework.

use async_trait::async_trait;
use super::runtime_trait::*;
use std::process::{Child, Command, Stdio};
use std::time::Duration;
use tracing::{info, warn};
use tokio::time::sleep;

pub struct SafetensorsRuntime {
    config: Option<RuntimeConfig>,
    server_process: Option<Child>,
    http_client: reqwest::Client,
    base_url: String,
}

impl SafetensorsRuntime {
    pub fn new() -> Self {
        Self {
            config: None,
            server_process: None,
            http_client: reqwest::Client::builder()
                .timeout(Duration::from_secs(600))
                .build()
                .unwrap_or_default(),
            base_url: String::new(),
        }
    }

    async fn start_server(&mut self, config: &RuntimeConfig) -> anyhow::Result<()> {
        let binary_path = config.runtime_binary.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Safetensors runtime requires runtime_binary path (e.g., candle-server)"))?;

        if !binary_path.exists() {
            return Err(anyhow::anyhow!(
                "Safetensors server binary not found at: {}. Use Candle or similar framework.",
                binary_path.display()
            ));
        }

        info!("Starting Safetensors server for model: {}", config.model_path.display());
        
        let mut cmd = Command::new(binary_path);
        cmd.arg("--model").arg(&config.model_path)
            .arg("--host").arg(&config.host)
            .arg("--port").arg(config.port.to_string())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let child = cmd.spawn()
            .map_err(|e| anyhow::anyhow!("Failed to spawn Safetensors server: {}", e))?;

        self.server_process = Some(child);
        self.base_url = format!("http://{}:{}", config.host, config.port);

        // Wait for server to be ready (up to 120 seconds) with exponential backoff.
        let _start = std::time::Instant::now();
        let mut delay_ms: u64 = 100;
        let mut last_log_secs: u64 = 0;
        loop {
            sleep(Duration::from_millis(delay_ms)).await;
            if self.is_ready().await {
                info!("✅ Safetensors runtime ready after {:.1}s", _start.elapsed().as_secs_f64());
                return Ok(());
            }
            let elapsed_secs = _start.elapsed().as_secs();
            if elapsed_secs >= 120 {
                break;
            }
            if elapsed_secs >= last_log_secs + 10 {
                info!("Still waiting for Safetensors server... ({}/120s)", elapsed_secs);
                last_log_secs = elapsed_secs;
            }
            delay_ms = (delay_ms * 2).min(2_000);
        }

        Err(anyhow::anyhow!("Safetensors server failed to start within 120 seconds"))
    }
}

impl Default for SafetensorsRuntime {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl ModelRuntime for SafetensorsRuntime {
    fn supported_format(&self) -> ModelFormat {
        ModelFormat::Safetensors
    }

    async fn initialize(&mut self, config: RuntimeConfig) -> anyhow::Result<()> {
        info!("Initializing Safetensors runtime");
        
        if config.format != ModelFormat::Safetensors {
            return Err(anyhow::anyhow!("Safetensors runtime received wrong format: {:?}", config.format));
        }

        self.config = Some(config.clone());
        self.start_server(&config).await?;
        Ok(())
    }

    async fn is_ready(&self) -> bool {
        if self.base_url.is_empty() {
            return false;
        }

        let health_url = format!("{}/health", self.base_url);
        match self.http_client.get(&health_url).send().await {
            Ok(resp) => resp.status().is_success(),
            Err(_) => false,
        }
    }

    async fn health_check(&self) -> anyhow::Result<String> {
        if self.base_url.is_empty() {
            return Err(anyhow::anyhow!("Runtime not initialized"));
        }

        let health_url = format!("{}/health", self.base_url);
        let resp = self.http_client.get(&health_url).send().await
            .map_err(|e| anyhow::anyhow!("Health check failed: {}", e))?;

        if resp.status().is_success() {
            Ok("healthy".to_string())
        } else {
            Err(anyhow::anyhow!("Health check returned: {}", resp.status()))
        }
    }

    fn base_url(&self) -> String {
        self.base_url.clone()
    }

    async fn generate(&self, request: InferenceRequest) -> anyhow::Result<InferenceResponse> {
        let url = self.completions_url();
        
        let payload = serde_json::json!({
            "model": "safetensors-llm",
            "messages": request.messages,
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
            "stream": false,
        });

        let resp = self.http_client.post(&url).json(&payload).send().await
            .map_err(|e| anyhow::anyhow!("Inference request failed: {}", e))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(anyhow::anyhow!("Inference failed ({}): {}", status, body));
        }

        let response: serde_json::Value = resp.json().await
            .map_err(|e| anyhow::anyhow!("Failed to parse response: {}", e))?;

        let content = response["choices"][0]["message"]["content"].as_str().unwrap_or("").to_string();
        let finish_reason = response["choices"][0]["finish_reason"].as_str().map(|s| s.to_string());

        Ok(InferenceResponse { content, finish_reason })
    }

    async fn generate_stream(
        &self,
        request: InferenceRequest,
    ) -> anyhow::Result<Box<dyn futures_util::Stream<Item = Result<String, anyhow::Error>> + Send + Unpin>> {
        use futures_util::StreamExt;
        
        let url = self.completions_url();
        let payload = serde_json::json!({
            "model": "safetensors-llm",
            "messages": request.messages,
            "max_tokens": request.max_tokens,
            "temperature": request.temperature,
            "stream": true,
        });

        let resp = self.http_client.post(&url).json(&payload).send().await
            .map_err(|e| anyhow::anyhow!("Stream request failed: {}", e))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(anyhow::anyhow!("Stream failed ({}): {}", status, body));
        }

        let byte_stream = resp.bytes_stream();
        let sse_stream = async_stream::try_stream! {
            let mut buffer = String::new();
            futures_util::pin_mut!(byte_stream);

            while let Some(chunk_result) = byte_stream.next().await {
                let chunk = chunk_result.map_err(|e| anyhow::anyhow!("Stream read error: {}", e))?;
                buffer.push_str(&String::from_utf8_lossy(&chunk));

                while let Some(newline_pos) = buffer.find('\n') {
                    let line = buffer[..newline_pos].trim().to_string();
                    buffer = buffer[newline_pos + 1..].to_string();

                    if line.is_empty() || !line.starts_with("data: ") {
                        continue;
                    }

                    let data = &line[6..];
                    if data == "[DONE]" {
                        return;
                    }

                    yield format!("data: {}\n\n", data);
                }
            }
        };

        Ok(Box::new(Box::pin(sse_stream)))
    }

    async fn shutdown(&mut self) -> anyhow::Result<()> {
        info!("Shutting down Safetensors runtime");
        
        if let Some(mut child) = self.server_process.take() {
            match child.kill() {
                Ok(_) => {
                    info!("Safetensors server process killed successfully");
                    let _ = child.wait();
                }
                Err(e) => {
                    warn!("Failed to kill Safetensors server: {}", e);
                }
            }
        }

        self.config = None;
        self.base_url.clear();
        Ok(())
    }

    fn metadata(&self) -> RuntimeMetadata {
        RuntimeMetadata {
            format: ModelFormat::Safetensors,
            runtime_name: "Safetensors (Candle)".to_string(),
            version: "latest".to_string(),
            supports_gpu: true,
            supports_streaming: true,
        }
    }
}

impl Drop for SafetensorsRuntime {
    fn drop(&mut self) {
        if let Some(mut child) = self.server_process.take() {
            let _ = child.kill();
            let _ = child.wait();
        }
    }
}