1use anyhow::{Context, Result};
2use serde::{Deserialize, Serialize};
3use serde_json::{json, Value};
4use std::fs;
5use std::path::PathBuf;
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "lowercase")]
10pub enum PythonBackend {
11 Uv,
13 Python,
15}
16
17impl Default for PythonBackend {
18 fn default() -> Self {
19 if crate::platform::environment::command_on_path("uv") {
21 PythonBackend::Uv
22 } else {
23 PythonBackend::Python
24 }
25 }
26}
27
28fn python_backend_from_env() -> PythonBackend {
36 match std::env::var("MLSTACK_PYTHON_BACKEND")
37 .ok()
38 .map(|v| v.trim().to_ascii_lowercase())
39 .as_deref()
40 {
41 Some("uv") => PythonBackend::Uv,
42 Some("python") => PythonBackend::Python,
43 _ => PythonBackend::default(),
44 }
45}
46
47fn default_vllm_version_lag() -> u32 {
50 1
51}
52
53fn default_onnx_install_method() -> String {
55 "migraphx".to_string()
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct InstallerConfig {
60 pub scripts_dir: String,
61 pub log_dir: String,
62 pub install_path: String,
63 pub batch_mode: bool,
64 pub install_method: String,
65 pub auto_confirm: bool,
66 pub star_repos: bool,
67 pub force_reinstall: bool,
68 pub theme: String,
69 pub performance_profile: String,
70 #[serde(default = "default_vllm_version_lag")]
73 pub vllm_version_lag: u32,
74 #[serde(default)]
77 pub vllm_version_min_age_days: u32,
78 #[serde(default = "default_onnx_install_method")]
83 pub onnx_install_method: String,
84 #[serde(default)]
87 pub onnx_version: Option<String>,
88 #[serde(default)]
89 pub python_backend: PythonBackend,
90 #[serde(skip)]
91 pub config_path: PathBuf,
92}
93
94impl InstallerConfig {
95 pub fn load_or_default(scripts_dir: &str) -> Result<Self> {
96 let config_path = config_file_path()?;
97 let log_dir = config_path
98 .parent()
99 .and_then(|p| p.parent())
100 .map(|p| p.join("logs"))
101 .unwrap_or_else(|| PathBuf::from("/tmp/mlstack/logs"));
102
103 let mut config = InstallerConfig::default_with_paths(
104 scripts_dir,
105 log_dir.to_string_lossy().to_string(),
106 config_path.clone(),
107 );
108
109 if config_path.exists() {
110 let raw = fs::read_to_string(&config_path).context("Failed to read config file")?;
111 let value: Value = serde_json::from_str(&raw).unwrap_or_else(|_| json!({}));
112 config.apply_from_value(&value);
113 } else {
114 let _ = config.save(None);
116 }
117
118 Ok(config)
119 }
120
121 pub fn save(&self, existing: Option<Value>) -> Result<()> {
122 if let Some(parent) = self.config_path.parent() {
123 fs::create_dir_all(parent).context("Failed to create config directory")?;
124 }
125 let mut value = existing.unwrap_or_else(|| json!({}));
126
127 value["scripts_dir"] = json!(self.scripts_dir);
128 value["log_dir"] = json!(self.log_dir);
129 value["batch_mode"] = json!(self.batch_mode);
130 value["install_method"] = json!(self.install_method);
131
132 let user_prefs = value
133 .get("user_preferences")
134 .cloned()
135 .unwrap_or_else(|| json!({}));
136 let mut user_prefs = user_prefs;
137 user_prefs["installation_path"] = json!(self.install_path);
138 user_prefs["install_method"] = json!(self.install_method);
139 user_prefs["auto_confirm"] = json!(self.auto_confirm);
140 user_prefs["star_repos"] = json!(self.star_repos);
141 user_prefs["force_reinstall"] = json!(self.force_reinstall);
142 user_prefs["theme"] = json!(self.theme);
143 user_prefs["performance_profile"] = json!(self.performance_profile);
144 user_prefs["python_backend"] = json!(self.python_backend);
145 user_prefs["vllm_version_lag"] = json!(self.vllm_version_lag);
146 user_prefs["vllm_version_min_age_days"] = json!(self.vllm_version_min_age_days);
147 user_prefs["onnx_install_method"] = json!(self.onnx_install_method);
148 user_prefs["onnx_version"] = json!(self.onnx_version);
149 value["user_preferences"] = user_prefs;
150
151 fs::write(&self.config_path, serde_json::to_string_pretty(&value)?)
152 .context("Failed to save config")?;
153
154 Ok(())
155 }
156
157 fn apply_from_value(&mut self, value: &Value) {
158 if let Some(dir) = value.get("scripts_dir").and_then(|v| v.as_str()) {
159 self.scripts_dir = dir.to_string();
160 }
161 if let Some(dir) = value.get("log_dir").and_then(|v| v.as_str()) {
162 self.log_dir = dir.to_string();
163 }
164 if let Some(batch) = value.get("batch_mode").and_then(|v| v.as_bool()) {
165 self.batch_mode = batch;
166 }
167 if let Some(method) = value.get("install_method").and_then(|v| v.as_str()) {
168 self.install_method = normalize_install_method(method);
169 }
170
171 if let Some(prefs) = value.get("user_preferences") {
172 if let Some(path) = prefs.get("installation_path").and_then(|v| v.as_str()) {
173 self.install_path = path.to_string();
174 }
175 if let Some(method) = prefs.get("install_method").and_then(|v| v.as_str()) {
176 self.install_method = normalize_install_method(method);
177 }
178 if let Some(auto) = prefs.get("auto_confirm").and_then(|v| v.as_bool()) {
179 self.auto_confirm = auto;
180 }
181 if let Some(star) = prefs.get("star_repos").and_then(|v| v.as_bool()) {
182 self.star_repos = star;
183 }
184 if let Some(force) = prefs.get("force_reinstall").and_then(|v| v.as_bool()) {
185 self.force_reinstall = force;
186 }
187 if let Some(theme) = prefs.get("theme").and_then(|v| v.as_str()) {
188 self.theme = theme.to_string();
189 }
190 if let Some(profile) = prefs.get("performance_profile").and_then(|v| v.as_str()) {
191 self.performance_profile = profile.to_string();
192 }
193 if let Some(backend) = prefs.get("python_backend").and_then(|v| v.as_str()) {
194 self.python_backend = match backend.to_lowercase().as_str() {
195 "uv" => PythonBackend::Uv,
196 "python" => PythonBackend::Python,
197 _ => PythonBackend::default(),
198 };
199 }
200 if let Some(lag) = prefs.get("vllm_version_lag").and_then(|v| v.as_u64()) {
201 self.vllm_version_lag = lag as u32;
202 }
203 if let Some(age) = prefs
204 .get("vllm_version_min_age_days")
205 .and_then(|v| v.as_u64())
206 {
207 self.vllm_version_min_age_days = age as u32;
208 }
209 if let Some(method) = prefs.get("onnx_install_method").and_then(|v| v.as_str()) {
210 self.onnx_install_method = method.to_string();
211 }
212 if let Some(v) = prefs.get("onnx_version").and_then(|v| v.as_str()) {
213 let v = v.trim();
214 self.onnx_version = if v.is_empty() {
215 None
216 } else {
217 Some(v.to_string())
218 };
219 }
220 }
221
222 if std::env::var("MLSTACK_PYTHON_BACKEND").is_ok() {
224 self.python_backend = python_backend_from_env();
225 }
226 }
227
228 pub fn default_with_paths(scripts_dir: &str, log_dir: String, config_path: PathBuf) -> Self {
229 let python_backend = if std::env::var("MLSTACK_PYTHON_BACKEND").is_ok() {
231 python_backend_from_env()
232 } else {
233 PythonBackend::default()
234 };
235
236 Self {
237 scripts_dir: scripts_dir.to_string(),
238 log_dir,
239 install_path: "/opt/rocm".into(),
240 batch_mode: std::env::var("MLSTACK_BATCH_MODE").unwrap_or_default() == "1",
241 install_method: normalize_install_method(
242 &std::env::var("MLSTACK_INSTALL_METHOD").unwrap_or_else(|_| "auto".into()),
243 ),
244 auto_confirm: false,
245 star_repos: true,
246 force_reinstall: false,
247 theme: "dark".into(),
248 performance_profile: "balanced".into(),
249 vllm_version_lag: 1,
250 vllm_version_min_age_days: 0,
251 onnx_install_method: std::env::var("MLSTACK_ONNX_INSTALL_METHOD")
252 .unwrap_or_else(|_| "migraphx".into()),
253 onnx_version: std::env::var("MLSTACK_ONNX_VERSION")
254 .ok()
255 .map(|s| s.trim().to_string())
256 .filter(|s| !s.is_empty()),
257 python_backend,
258 config_path,
259 }
260 }
261}
262
263fn normalize_install_method(input: &str) -> String {
264 match input.trim().to_ascii_lowercase().as_str() {
265 "global" => "global".into(),
266 "venv" => "venv".into(),
267 _ => "auto".into(),
268 }
269}
270
271pub fn config_file_path() -> Result<PathBuf> {
272 let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
273 let config_dir = PathBuf::from(home).join(".mlstack").join("config");
274 fs::create_dir_all(&config_dir).context("Failed to create config directory")?;
275 Ok(config_dir.join("config.json"))
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281
282 #[test]
283 fn onnx_config_defaults_and_apply() {
284 let _env = crate::test_support::lock_env();
285 std::env::remove_var("MLSTACK_ONNX_INSTALL_METHOD");
286 std::env::remove_var("MLSTACK_ONNX_VERSION");
287
288 let mut cfg = InstallerConfig::default_with_paths(
289 "scripts",
290 "logs".into(),
291 PathBuf::from("/tmp/rusty-config-test.json"),
292 );
293 assert_eq!(cfg.onnx_install_method, "migraphx");
295 assert_eq!(cfg.onnx_version, None);
296
297 let raw = serde_json::json!({
299 "user_preferences": {
300 "onnx_install_method": "source",
301 "onnx_version": "1.27.1",
302 }
303 });
304 cfg.apply_from_value(&raw);
305 assert_eq!(cfg.onnx_install_method, "source");
306 assert_eq!(cfg.onnx_version.as_deref(), Some("1.27.1"));
307
308 let raw = serde_json::json!({ "user_preferences": { "onnx_version": " " } });
310 cfg.apply_from_value(&raw);
311 assert_eq!(cfg.onnx_version, None);
312 }
313}