sokuto_cli/commands/
serve.rs

1//! # HTTP API サーバー起動コマンド
2//!
3//! TypeSafe AI Jev 互換 HTTP API サーバーを初期化し、
4//! 指定されたモデルと Execution Provider でエンドポイントを提供する。
5
6use std::fs;
7use std::net::SocketAddr;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11use sokuto_core::contract::calibration::CalibrationConfig;
12use sokuto_runtime::engine::{
13    CacheCompressionMode, DualSessionEngine, ExecutionProvider, InferenceEngine, SessionConfig,
14    StateCacheManager,
15};
16use sokuto_runtime::tokenizer::JevTokenizer;
17use sokuto_server::metrics::setup_metrics_recorder;
18use sokuto_server::state::AppState;
19
20/// サーバー起動オプション構造体。
21#[derive(Debug, Clone)]
22pub struct ServeArgs {
23    /// バインド先ホストアドレス。
24    pub host: String,
25    /// バインド先ポート番号。
26    pub port: u16,
27    /// モデルファイル格納ディレクトリ。
28    pub model_dir: PathBuf,
29    /// セッションプールサイズ。
30    pub pool_size: usize,
31    /// 優先 Execution Provider (auto, coreml, cuda, tensorrt, cpu)。
32    pub provider: Option<String>,
33    /// 単一オペレータ内の並列スレッド数 (intra-op)。
34    pub intra_threads: Option<usize>,
35    /// 複数オペレータ間の並列スレッド数 (inter-op)。
36    pub inter_threads: Option<usize>,
37}
38
39/// 文字列から Execution Provider をパースする。
40pub fn parse_provider(name: &str) -> Result<ExecutionProvider, String> {
41    match name.to_lowercase().as_str() {
42        "auto" => Ok(ExecutionProvider::Auto),
43        "coreml" => Ok(ExecutionProvider::CoreML),
44        "cuda" => Ok(ExecutionProvider::CUDA),
45        "tensorrt" => Ok(ExecutionProvider::TensorRT),
46        "cpu" => Ok(ExecutionProvider::CPU),
47        other => Err(format!(
48            "未知の Execution Provider です: '{other}'。利用可能: auto, coreml, cuda, tensorrt, cpu。"
49        )),
50    }
51}
52
53/// モデルディレクトリの完全性を検証する。
54pub fn validate_model_dir(dir: &Path) -> Result<(), String> {
55    if !dir.exists() {
56        return Err(format!(
57            "モデルディレクトリが存在しません: {}。",
58            dir.display()
59        ));
60    }
61
62    let onnx_file = dir.join("model.onnx");
63    if !onnx_file.exists() {
64        return Err(format!(
65            "必須モデルファイルが見つかりません: {}。",
66            onnx_file.display()
67        ));
68    }
69
70    let tok_file = dir.join("tokenizer.json");
71    if !tok_file.exists() {
72        return Err(format!(
73            "必須トークナイザー定義ファイルが見つかりません: {}。",
74            tok_file.display()
75        ));
76    }
77
78    Ok(())
79}
80
81/// HTTP サーバー起動処理を実行する。
82///
83/// # 引数
84/// - `args`: サーバー起動構成パラメータ。
85pub async fn run_serve(args: ServeArgs) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
86    tracing_subscriber::fmt()
87        .with_env_filter(
88            tracing_subscriber::EnvFilter::try_from_default_env()
89                .unwrap_or_else(|_| "sokuto_server=info,tower_http=info".into()),
90        )
91        .init();
92
93    // 1. モデルディレクトリの検証 (Fast-Fail)
94    if let Err(err) = validate_model_dir(&args.model_dir) {
95        tracing::error!("{err}");
96        return Err(err.into());
97    }
98
99    // 2. セッション設定の構築
100    let mut session_config = SessionConfig {
101        pool_size: args.pool_size.max(1),
102        intra_threads: args.intra_threads,
103        inter_threads: args.inter_threads.or(Some(1)),
104        ..Default::default()
105    };
106
107    if let Some(ref provider_str) = args.provider {
108        let provider = parse_provider(provider_str)?;
109        session_config.preferred_providers = match provider {
110            ExecutionProvider::Auto => SessionConfig::default().preferred_providers,
111            other => vec![other, ExecutionProvider::CPU],
112        };
113    }
114
115    let onnx_path = args.model_dir.join("model.onnx");
116    let tok_path = args.model_dir.join("tokenizer.json");
117    let calib_path = args.model_dir.join("calibration.json");
118
119    tracing::info!(
120        "sokuto サーバーを起動します (モデル: {}, アドレス: {}:{}, セッションプール数: {})...",
121        args.model_dir.display(),
122        args.host,
123        args.port,
124        args.pool_size
125    );
126
127    // 3. Prometheus レコーダーの初期化
128    let prometheus_handle = match setup_metrics_recorder() {
129        Ok(handle) => Some(handle),
130        Err(err) => {
131            tracing::warn!("Prometheus メトリクス初期化をスキップしました: {err}");
132            None
133        }
134    };
135
136    // 4. 推論エンジンとトークナイザーのロード
137    let engine = Arc::new(InferenceEngine::new(&onnx_path, session_config.clone())?);
138    let tokenizer = Arc::new(JevTokenizer::from_file(&tok_path)?);
139
140    let calib_config = if calib_path.exists() {
141        let calib_json = fs::read_to_string(&calib_path)?;
142        tracing::info!("較正温度設定をロードしました: {}", calib_path.display());
143        Arc::new(CalibrationConfig::from_json_str(&calib_json)?)
144    } else {
145        tracing::info!("較正設定ファイルが存在しないため、デフォルト温度設定を使用します。");
146        Arc::new(CalibrationConfig::default())
147    };
148
149    let mut state = AppState::new(engine, tokenizer, calib_config)
150        .load_hierarchical_mapping_from_dir(&args.model_dir);
151
152    // デュアルセッションモデル (state_encoder.onnx, query_evaluator.onnx) の検出と初期化
153    let state_model_path = args.model_dir.join("state_encoder.onnx");
154    let query_model_path = args.model_dir.join("query_evaluator.onnx");
155    if state_model_path.exists() && query_model_path.exists() {
156        tracing::info!(
157            "2 セッション分割型モデルを検出しました: State={}, Query={}",
158            state_model_path.display(),
159            query_model_path.display()
160        );
161
162        let cache_max_bytes = std::env::var("SOKUTO_STATE_CACHE_MAX_BYTES")
163            .ok()
164            .and_then(|s| s.parse::<u64>().ok())
165            .unwrap_or(512 * 1024 * 1024);
166
167        let cache_ttl_secs = std::env::var("SOKUTO_STATE_CACHE_TTL_SECS")
168            .ok()
169            .and_then(|s| s.parse::<u64>().ok())
170            .unwrap_or(600);
171
172        let compression_mode = match std::env::var("SOKUTO_STATE_CACHE_COMPRESSION")
173            .unwrap_or_else(|_| "fp16".to_string())
174            .to_lowercase()
175            .as_str()
176        {
177            "none" | "fp32" => CacheCompressionMode::None,
178            "int8" => CacheCompressionMode::Int8Symmetric,
179            _ => CacheCompressionMode::Fp16,
180        };
181
182        let cache_manager = Arc::new(StateCacheManager::new(
183            cache_max_bytes,
184            cache_ttl_secs,
185            compression_mode,
186        ));
187
188        match DualSessionEngine::new(
189            &state_model_path,
190            &query_model_path,
191            session_config,
192            Arc::clone(&cache_manager),
193        ) {
194            Ok(dual_engine) => {
195                tracing::info!(
196                    "DualSessionEngine を初期化しました (キャッシュ上限: {}MB, 圧縮: {:?})。",
197                    cache_max_bytes / (1024 * 1024),
198                    compression_mode
199                );
200                state = state.with_dual_engine(Arc::new(dual_engine), cache_manager);
201            }
202            Err(err) => {
203                tracing::warn!(
204                    "DualSessionEngine の構築に失敗したため、単一モデルモードで運用します: {err}"
205                );
206            }
207        }
208    }
209
210    let state = Arc::new(state);
211    let addr: SocketAddr = format!("{}:{}", args.host, args.port).parse()?;
212
213    sokuto_server::run_server(addr, state, prometheus_handle).await?;
214
215    Ok(())
216}
sokuto_cli/commands/serve.rs

sokuto_cli/commands/
serve.rs
