sqlite_graphrag/commands/
slots.rs1use clap::{Args, Subcommand};
10use serde::Serialize;
11
12use crate::errors::AppError;
13use crate::llm_slots::{slot_path, slots_dir};
14use crate::output::emit_json_compact;
15use crate::output::OutputFormat;
16
17#[derive(Debug, Args)]
20pub struct SlotsArgs {
21 #[command(subcommand)]
22 pub cmd: SlotsCmd,
23}
24
25#[derive(Debug, Subcommand)]
26pub enum SlotsCmd {
27 Status(SlotsStatusArgs),
29 Release {
31 #[arg(long)]
33 slot_id: u32,
34 #[arg(long)]
36 yes: bool,
37 },
38 Cleanup {
40 #[arg(long, default_value_t = 3600)]
42 stale_after: u64,
43 #[arg(long)]
45 yes: bool,
46 #[arg(long)]
48 dry_run: bool,
49 },
50}
51
52#[derive(Debug, clap::Args)]
53pub struct SlotsStatusArgs {
54 #[arg(long, value_enum, default_value_t = OutputFormat::Json)]
56 pub format: OutputFormat,
57}
58
59#[derive(Serialize)]
60struct SlotEntry {
61 slot_id: u32,
62 path: String,
63 age_secs: u64,
64 pid_hint: Option<u32>,
65}
66
67#[derive(Serialize)]
68struct SlotsStatusOutput {
69 action: &'static str,
70 max_concurrency: u32,
71 active: usize,
72 free: usize,
73 slots: Vec<SlotEntry>,
74 elapsed_ms: u64,
75}
76
77pub fn run(args: SlotsArgs) -> Result<(), AppError> {
78 run_cmd(args.cmd)
79}
80
81fn run_cmd(cmd: SlotsCmd) -> Result<(), AppError> {
82 match cmd {
83 SlotsCmd::Status(args) => run_status(args),
84 SlotsCmd::Release { slot_id, yes } => run_release(slot_id, yes),
85 SlotsCmd::Cleanup {
86 stale_after,
87 yes,
88 dry_run,
89 } => run_cleanup(stale_after, yes, dry_run),
90 }
91}
92
93fn run_status(args: SlotsStatusArgs) -> Result<(), AppError> {
94 let start = std::time::Instant::now();
95 let max = crate::llm_slots::default_max_concurrency();
96 let dir = slots_dir();
97 let mut entries: Vec<SlotEntry> = Vec::new();
98
99 if dir.is_dir() {
100 for slot_id in 0..max {
101 let path = slot_path(slot_id);
102 if path.is_file() {
103 let age_secs = path
104 .metadata()
105 .and_then(|m| m.modified())
106 .ok()
107 .and_then(|t| t.elapsed().ok())
108 .map(|d| d.as_secs())
109 .unwrap_or(0);
110 let pid_hint = std::fs::read_to_string(&path)
111 .ok()
112 .and_then(|s| s.trim().parse::<u32>().ok());
113 entries.push(SlotEntry {
114 slot_id,
115 path: path.to_string_lossy().into_owned(),
116 age_secs,
117 pid_hint,
118 });
119 }
120 }
121 }
122
123 let output = SlotsStatusOutput {
124 action: "slots_status",
125 max_concurrency: max,
126 active: entries.len(),
127 free: (max as usize).saturating_sub(entries.len()),
128 slots: entries,
129 elapsed_ms: start.elapsed().as_millis() as u64,
130 };
131
132 if matches!(args.format, OutputFormat::Json) {
133 let json = serde_json::to_string_pretty(&output).map_err(AppError::Json)?;
134 println!("{json}");
136 } else {
137 tracing::info!(target: "slots", max_concurrency = output.max_concurrency, "slot status");
142 tracing::info!(
143 target: "slots",
144 active = output.active,
145 free = output.free,
146 "slot occupancy"
147 );
148 for s in &output.slots {
149 let pid = s.pid_hint.map(|p| p.to_string()).unwrap_or_default();
150 tracing::info!(
151 target: "slots",
152 slot_id = s.slot_id,
153 age_secs = s.age_secs,
154 pid = %pid,
155 path = %s.path,
156 "slot entry"
157 );
158 }
159 }
160 Ok(())
161}
162
163fn run_release(slot_id: u32, yes: bool) -> Result<(), AppError> {
164 let path = slot_path(slot_id);
165 if !path.is_file() {
166 return Err(AppError::NotFound(format!(
167 "slot {slot_id} is not held (no file at {})",
168 path.display()
169 )));
170 }
171 if !yes {
172 eprintln!(
173 "About to release slot {slot_id} at {}. Pass --yes to skip confirmation.",
174 path.display()
175 );
176 }
177 std::fs::remove_file(&path).map_err(AppError::Io)?;
178 let out = serde_json::json!({
179 "action": "slot_released",
180 "slot_id": slot_id,
181 "path": path.to_string_lossy(),
182 });
183 let _ = emit_json_compact(&out);
184 Ok(())
185}
186
187fn run_cleanup(stale_after: u64, yes: bool, dry_run: bool) -> Result<(), AppError> {
188 let start = std::time::Instant::now();
189 let max = crate::llm_slots::default_max_concurrency();
190 let mut removed: Vec<u32> = Vec::new();
191 for slot_id in 0..max {
192 let path = slot_path(slot_id);
193 if !path.is_file() {
194 continue;
195 }
196 let age = path
197 .metadata()
198 .and_then(|m| m.modified())
199 .ok()
200 .and_then(|t| t.elapsed().ok())
201 .map(|d| d.as_secs())
202 .unwrap_or(0);
203 if age >= stale_after {
204 if !dry_run {
205 if let Err(e) = std::fs::remove_file(&path) {
206 tracing::warn!(target: "slots", slot_id, error = %e, "stale slot removal failed");
207 continue;
208 }
209 }
210 removed.push(slot_id);
211 }
212 }
213 let out = serde_json::json!({
214 "action": if dry_run { "slots_cleanup_dry_run" } else { "slots_cleanup" },
215 "stale_after_secs": stale_after,
216 "removed": removed,
217 "removed_count": removed.len(),
218 "elapsed_ms": start.elapsed().as_millis() as u64,
219 "yes": yes,
220 });
221 let _ = emit_json_compact(&out);
222 Ok(())
223}
224
225#[cfg(test)]
229mod tests {
230 use super::*;
231 use crate::llm_slots::acquire_llm_slot;
232
233 #[test]
234 fn acquire_then_drop_releases_slot() {
235 let _ = std::fs::remove_dir_all(crate::llm_slots::slots_dir());
236 let guard = acquire_llm_slot(2, 5).expect("acquire");
237 let path = slot_path(guard.slot_id());
238 assert!(path.is_file(), "slot file must exist after acquire");
239 drop(guard);
240 assert!(
241 !path.is_file(),
242 "slot file must be removed after Drop (RAII guarantee)"
243 );
244 }
245}