use std::collections::{BTreeMap, HashMap};
use std::io;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use tokio::sync::broadcast::Receiver;
use tokio::sync::broadcast::error::RecvError;
use tokio::sync::mpsc::UnboundedSender;
use tokio::task::{AbortHandle, JoinHandle, JoinSet};
use tokio_util::sync::CancellationToken;
use super::catalog::{Listed, Query, fetch};
use super::compare::CompareListing;
use super::cost::history_cost;
use super::editor::Edited;
use super::pod_logs::{PodLogRead, read_pod_log};
use super::project::ProjectConfig;
use super::project_edit::{SaveRefusal, save_config};
use super::start::{AutoPlan, Catalog, StartPlan, look_up, prepare, prepare_auto, prepare_resume};
use super::training::{Listing, list_runs, read_series};
use crate::cli::compare::Preview;
use crate::cli::data::{Command, Load};
use crate::cli::front::{Frontend, Report};
use crate::cli::push::{self, PushOptions, PushPlan};
use crate::cli::{CompareArgs, StageArgs, TrainArgs, TrainCommand};
use crate::config::{DotenvKeys, EnvSource, ReloadError, Source, Stamp, reload, stamp};
use crate::dataset::{Counts, DataFiles, Dataset, Deletion};
use crate::events::{Event, EventBus, Observer, Stage};
use crate::history::{self, Cost, Entry, Total};
use crate::pipeline::{Ctx, SplitReport};
use crate::prompts::Prompts;
use crate::runs::Runs;
use crate::train::sizing::VramFloor;
use crate::train::{Phases, TrainMetric};
pub(super) const BUS_CAPACITY: usize = 32_768;
const FORWARD_GRACE: Duration = Duration::from_secs(5);
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(super) struct TaskId(pub(super) u64);
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum Task {
Load,
Edit(Edit),
Pipeline(Command),
Train(TrainJob),
Runs,
Series(String),
PodLog {
run: String,
offset: u64,
},
Prepare,
PrepareResume(String),
PrepareAuto,
PreparePush(String),
Push(String),
Compares,
PrepareCompare(Option<String>),
Compare(CompareArgs),
StartCatalog(u32),
Catalog(Query),
CheckConfig {
seen: Stamp,
dotenv: DotenvKeys,
},
SaveConfig {
text: String,
base: String,
env: EnvSource,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum TrainJob {
Start {
vram_floor: VramFloor,
},
Attach(String),
Cancel(String),
Stop(String),
Resume(String),
}
impl TrainJob {
fn args(&self) -> TrainArgs {
let command = match self {
Self::Start { vram_floor } => {
return TrainArgs {
vram_floor: *vram_floor,
..TrainArgs::default()
};
},
Self::Resume(run_id) => {
return TrainArgs {
resume_from: Some(run_id.clone()),
..TrainArgs::default()
};
},
Self::Stop(run_id) => TrainCommand::Stop {
run_id: run_id.clone(),
},
Self::Attach(run_id) => TrainCommand::Attach {
run_id: run_id.clone(),
},
Self::Cancel(run_id) => TrainCommand::Cancel {
run_id: run_id.clone(),
},
};
TrainArgs {
command: Some(command),
..TrainArgs::default()
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum Edit {
Change(Edited),
Delete {
deletion: Deletion,
counts: Counts,
},
}
#[derive(Debug, Clone, Default, PartialEq)]
pub(super) struct History {
pub(super) cost: Option<Cost>,
pub(super) stages: BTreeMap<Stage, Total>,
pub(super) all: Total,
pub(super) models: BTreeMap<String, Total>,
}
impl History {
pub(super) fn of(entries: &[Entry]) -> Self {
let (stages, all) = history::totals(entries);
Self {
cost: history_cost(&all),
stages,
all,
models: history::per_model(entries),
}
}
#[cfg(test)]
pub(super) fn costing(cost: Cost) -> Self {
Self {
cost: Some(cost),
..Self::default()
}
}
}
#[derive(Debug)]
pub(super) enum Done {
Loaded(Result<Dataset, String>, Result<History, String>),
Saved(Result<Saved, String>),
Pipeline(Result<(), String>),
Trained(Result<(), String>),
Runs(Listing),
Series {
run: String,
series: Option<(Vec<TrainMetric>, Phases)>,
},
PodLog(PodLogRead),
Prepared(Result<StartPlan, String>),
PreparedAuto(Result<AutoPlan, String>),
PushPrepared(Result<PushPlan, String>),
Pushed(Result<String, String>),
Compares(CompareListing),
ComparePrepared(Result<Preview, String>),
Compared(Result<(), String>),
StartCatalog(Catalog),
Catalog(Result<Listed, String>),
ConfigSaved(Result<Box<ProjectConfig>, SaveRefusal>),
ConfigChecked(Box<Checked>),
}
#[derive(Debug)]
pub(super) struct Checked {
pub(super) stamp: Stamp,
pub(super) read: Option<Result<Reread, ReloadError>>,
}
#[derive(Debug)]
pub(super) struct Reread {
pub(super) config: ProjectConfig,
pub(super) env: EnvSource,
}
pub(super) fn check_config(dir: &Path, seen: Stamp, dotenv: &DotenvKeys) -> Checked {
let now = stamp(dir);
if now == seen {
return Checked {
stamp: now,
read: None,
};
}
let read = reload(dir, dotenv).and_then(|reloaded| {
let mut config = ProjectConfig::new(&reloaded.text, &reloaded.env)?;
config.stamp = Some(now);
Ok(Reread {
config,
env: reloaded.env,
})
});
Checked {
stamp: now,
read: Some(read),
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct Saved {
pub(super) message: String,
pub(super) split: Result<SplitReport, String>,
}
#[derive(Debug)]
pub(super) enum Msg {
EditorExited(io::Result<ExitStatus>),
BrowserFailed(String),
LogsExported(Result<(String, usize), String>),
NewerRelease(String),
Event(TaskId, Event),
Lagged(TaskId, u64),
Report(TaskId, Report),
}
pub(super) fn forward(
id: TaskId,
mut events: Receiver<Event>,
messages: UnboundedSender<Msg>,
) -> JoinHandle<()> {
tokio::spawn(async move {
loop {
let message = match events.recv().await {
Ok(event) => Msg::Event(id, event),
Err(RecvError::Lagged(skipped)) => Msg::Lagged(id, skipped),
Err(RecvError::Closed) => break,
};
if messages.send(message).is_err() {
break;
}
}
})
}
async fn forwarded(front: Frontend, forwarder: JoinHandle<()>, what: &str) {
drop(front);
if tokio::time::timeout(FORWARD_GRACE, forwarder)
.await
.is_err()
{
tracing::warn!("the events of the {what} were not all forwarded");
}
}
fn front_end(
id: TaskId,
messages: &UnboundedSender<Msg>,
detach: &CancellationToken,
abandon: &Arc<AtomicBool>,
observer: Option<Arc<dyn Observer>>,
) -> (Frontend, JoinHandle<()>) {
let bus = EventBus::observed(BUS_CAPACITY, observer);
let forwarder = forward(id, bus.subscribe(), messages.clone());
let sink = messages.clone();
let front = Frontend::Tui {
bus,
detach: detach.clone(),
abandon: Arc::clone(abandon),
report: Arc::new(move |report| {
sink.send(Msg::Report(id, report)).ok();
}),
};
(front, forwarder)
}
pub(super) fn save(project_dir: &Path, edit: &Edit, source: &Source) -> Result<Saved, String> {
let error = |error: anyhow::Error| format!("{error:#}");
let settings = source.load(project_dir).map_err(|e| error(e.into()))?;
let files = DataFiles::new(project_dir);
let mut data = Dataset::read(&files).map_err(|e| error(e.into()))?;
let change = match edit {
Edit::Change(Edited::Question { id, before, text }) => data.edit_question(id, before, text),
Edit::Change(Edited::Answer { id, before, after }) => {
data.edit_answer(id, before, after.clone())
},
Edit::Change(Edited::Subtopic { id, before, name }) => {
data.rename_subtopic(id, before, name)
},
Edit::Delete { deletion, counts } => data.delete(deletion, *counts),
}
.map_err(|e| error(e.into()))?;
data.save(&files, &change).map_err(|e| error(e.into()))?;
let ctx = Ctx {
settings: &settings,
files: &files,
prompts: &Prompts::empty(),
bus: &EventBus::new(),
topic: None,
force: false,
};
let split = crate::pipeline::split(&ctx).map_err(|e| error(e.into()));
Ok(Saved {
message: change.message,
split,
})
}
pub(super) struct Tasks {
project_dir: PathBuf,
messages: UnboundedSender<Msg>,
set: JoinSet<Done>,
ids: HashMap<tokio::task::Id, TaskId>,
tokens: HashMap<TaskId, CancellationToken>,
abandons: HashMap<TaskId, Arc<AtomicBool>>,
lookups: HashMap<TaskId, AbortHandle>,
source: Source,
observer: Option<Arc<dyn Observer>>,
}
impl Tasks {
pub(super) fn new(project_dir: &Path, messages: UnboundedSender<Msg>) -> Self {
Self {
project_dir: project_dir.to_path_buf(),
messages,
set: JoinSet::new(),
ids: HashMap::new(),
tokens: HashMap::new(),
abandons: HashMap::new(),
lookups: HashMap::new(),
source: Source::from(EnvSource::Process),
observer: None,
}
}
pub(super) fn use_config(&mut self, source: Source) {
self.source = source;
}
pub(super) fn observe(&mut self, observer: Option<Arc<dyn Observer>>) {
self.observer = observer;
}
pub(super) fn project_dir(&self) -> &Path {
&self.project_dir
}
pub(super) fn abort_lookups(&mut self) {
for (_, lookup) in self.lookups.drain() {
lookup.abort();
}
}
pub(super) fn abandon(&self, id: TaskId) {
if let Some(abandon) = self.abandons.get(&id) {
abandon.store(true, std::sync::atomic::Ordering::SeqCst);
}
self.cancel(id);
}
pub(super) fn cancel(&self, id: TaskId) {
if let Some(token) = self.tokens.get(&id) {
token.cancel();
}
}
pub(super) fn is_empty(&self) -> bool {
self.set.is_empty()
}
fn spawn_prepare(&mut self, auto: bool) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.set.spawn(async move {
if auto {
let plan = tokio::task::spawn_blocking(move || prepare_auto(&dir, &source)).await;
return Done::PreparedAuto(match plan {
Ok(plan) => plan,
Err(error) => Err(format!("cannot prepare auto mode: {error}")),
});
}
let plan = tokio::task::spawn_blocking(move || prepare(&dir, &source)).await;
Done::Prepared(match plan {
Ok(plan) => plan,
Err(error) => Err(format!("cannot prepare the run: {error}")),
})
})
}
fn spawn_prepare_resume(&mut self, run: String) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.set.spawn(async move {
let plan =
tokio::task::spawn_blocking(move || prepare_resume(&dir, &source, &run)).await;
Done::Prepared(match plan {
Ok(plan) => plan,
Err(error) => Err(format!("cannot prepare the run: {error}")),
})
})
}
pub(super) fn spawn(&mut self, id: TaskId, task: Task) {
let files = DataFiles::new(&self.project_dir);
let handle = match task {
Task::Load => {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let read = tokio::task::spawn_blocking(move || {
let history = history::read(&dir)
.map(|entries| History::of(&entries))
.map_err(|error| {
format!("cannot read {}: {error}", history::path(&dir).display())
});
(Dataset::read(&files), history)
})
.await;
match read {
Ok((Ok(data), history)) => Done::Loaded(Ok(data), history),
Ok((Err(error), history)) => {
Done::Loaded(Err(format!("{:#}", anyhow::Error::from(error))), history)
},
Err(error) => Done::Loaded(
Err(format!("cannot read the data files: {error}")),
Err(format!("cannot read the history: {error}")),
),
}
})
},
Task::Edit(edit) => {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.set.spawn(async move {
let saved =
tokio::task::spawn_blocking(move || save(&dir, &edit, &source)).await;
Done::Saved(match saved {
Ok(saved) => saved,
Err(error) => Err(format!("the edit failed: {error}")),
})
})
},
Task::CheckConfig { seen, dotenv } => self.spawn_check(seen, dotenv),
Task::SaveConfig { text, base, env } => self.spawn_save(text, base, env),
Task::Pipeline(command) => self.spawn_pipeline(id, command),
Task::Train(job) => self.spawn_train(id, job),
Task::Runs => self.spawn_runs(),
Task::Series(run) => {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let id = run.clone();
let series = tokio::task::spawn_blocking(move || read_series(&dir, &id)).await;
Done::Series {
run,
series: series.ok().flatten(),
}
})
},
Task::PodLog { run, offset } => {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let name = run.clone();
let read =
tokio::task::spawn_blocking(move || read_pod_log(&dir, &name, offset))
.await;
Done::PodLog(read.unwrap_or_else(|error| PodLogRead {
run,
offset,
restarted: false,
lines: Vec::new(),
error: Some(format!("cannot read the pod log: {error}")),
}))
})
},
Task::Prepare => self.spawn_prepare(false),
Task::PrepareResume(run) => self.spawn_prepare_resume(run),
Task::PrepareAuto => self.spawn_prepare(true),
Task::PreparePush(run) => self.spawn_prepare_push(run),
Task::Push(run) => self.spawn_push(id, run),
Task::Compares => self.spawn_compares(),
Task::PrepareCompare(run) => self.spawn_prepare_compare(run),
Task::Compare(args) => self.spawn_compare(id, args),
Task::StartCatalog(gpu_count) => {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.spawn_lookup(id, async move {
Done::StartCatalog(look_up(&dir, source, gpu_count).await)
})
},
Task::Catalog(query) => {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.spawn_lookup(id, async move {
Done::Catalog(fetch(&dir, source, query).await)
})
},
};
self.ids.insert(handle.id(), id);
}
fn spawn_save(&mut self, text: String, base: String, env: EnvSource) -> AbortHandle {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let saved = tokio::task::spawn_blocking(move || {
save_config(&dir, &text, &base, &env).map(Box::new)
})
.await;
Done::ConfigSaved(saved.unwrap_or_else(|error| {
Err(SaveRefusal::Failed(format!("the save failed: {error}")))
}))
})
}
fn spawn_check(&mut self, seen: Stamp, dotenv: DotenvKeys) -> AbortHandle {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let checked =
tokio::task::spawn_blocking(move || check_config(&dir, seen, &dotenv)).await;
Done::ConfigChecked(Box::new(checked.unwrap_or_else(|error| {
tracing::warn!("cannot check the configuration files: {error}");
Checked {
stamp: seen,
read: None,
}
})))
})
}
fn spawn_lookup(
&mut self,
id: TaskId,
read: impl Future<Output = Done> + Send + 'static,
) -> AbortHandle {
let handle = self.set.spawn(read);
self.lookups.insert(id, handle.clone());
handle
}
fn spawn_pipeline(&mut self, id: TaskId, command: Command) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) =
front_end(id, &self.messages, &token, &abandon, self.observer.clone());
self.tokens.insert(id, token.clone());
self.set.spawn(async move {
let args = StageArgs::default();
let outcome =
crate::cli::data::run(&dir, command, &args, &front, Load::Source(&source))
.await
.map_err(|error| format!("{error:#}"));
forwarded(front, forwarder, "stage").await;
Done::Pipeline(outcome)
})
}
fn spawn_train(&mut self, id: TaskId, job: TrainJob) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) =
front_end(id, &self.messages, &token, &abandon, self.observer.clone());
self.tokens.insert(id, token);
self.abandons.insert(id, abandon);
self.set.spawn(async move {
let result = crate::cli::train::run(&dir, &job.args(), &front, &source).await;
forwarded(front, forwarder, "training").await;
Done::Trained(result.map_err(|error| format!("{error:#}")))
})
}
fn spawn_prepare_push(&mut self, run: String) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.set.spawn(async move {
let plan = tokio::task::spawn_blocking(move || {
let settings = source.load(&dir)?;
push::plan(&Runs::new(&dir), &settings, &run)
})
.await;
Done::PushPrepared(match plan {
Ok(plan) => plan.map_err(|error| format!("{error:#}")),
Err(error) => Err(format!("cannot prepare the push: {error}")),
})
})
}
fn spawn_push(&mut self, id: TaskId, run: String) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) = front_end(id, &self.messages, &token, &abandon, None);
self.tokens.insert(id, token.clone());
self.set.spawn(async move {
let push = async {
let settings = source.load(&dir)?;
let opts = PushOptions::default();
push::push_with_token(&Runs::new(&dir), &settings, &run, &opts, &front).await
};
let pushed = tokio::select! {
biased;
pushed = push => pushed
.map(|pushed| pushed.url.unwrap_or_else(|| pushed.repo.to_string()))
.map_err(|error| format!("{error:#}")),
() = token.cancelled() => Err(push::PUSH_CANCELLED.to_string()),
};
forwarded(front, forwarder, "push").await;
Done::Pushed(pushed)
})
}
fn spawn_runs(&mut self) -> AbortHandle {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let listing = tokio::task::spawn_blocking(move || list_runs(&dir)).await;
Done::Runs(listing.unwrap_or_else(|error| Listing {
runs: Err(format!("cannot list the runs: {error}")),
skipped: Vec::new(),
}))
})
}
fn spawn_compares(&mut self) -> AbortHandle {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let listing = tokio::task::spawn_blocking(move || super::compare::list(&dir)).await;
Done::Compares(listing.unwrap_or_else(|error| CompareListing {
rows: Err(format!("cannot list the compares: {error}")),
skipped: Vec::new(),
}))
})
}
fn spawn_prepare_compare(&mut self, run: Option<String>) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
self.set.spawn(async move {
let preview = crate::cli::compare::preview(&dir, &source, run.as_deref())
.await
.map_err(|error| format!("{error:#}"));
Done::ComparePrepared(preview)
})
}
fn spawn_compare(&mut self, id: TaskId, args: CompareArgs) -> AbortHandle {
let dir = self.project_dir.clone();
let source = self.source.clone();
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) =
front_end(id, &self.messages, &token, &abandon, self.observer.clone());
self.tokens.insert(id, token);
self.abandons.insert(id, abandon);
self.set.spawn(async move {
let result = crate::cli::compare::run(&dir, &args, &front, &source).await;
forwarded(front, forwarder, "compare").await;
Done::Compared(result.map_err(|error| format!("{error:#}")))
})
}
pub(super) async fn next(&mut self) -> Option<(TaskId, Result<Done, String>)> {
loop {
let joined = self.set.join_next_with_id().await?;
let (task, result) = match joined {
Ok((task, done)) => (task, Ok(done)),
Err(error) => (
error.id(),
Err(format!("a background task failed: {error}")),
),
};
if let Some(id) = self.ids.remove(&task) {
self.tokens.remove(&id);
self.abandons.remove(&id);
self.lookups.remove(&id);
return Some((id, result));
}
}
}
}
#[cfg(test)]
impl Tasks {
pub(super) fn park(&mut self, id: TaskId, release: CancellationToken) -> Arc<AtomicBool> {
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
self.tokens.insert(id, token.clone());
self.abandons.insert(id, Arc::clone(&abandon));
let handle = self.set.spawn(async move {
tokio::select! {
() = release.cancelled() => {},
() = token.cancelled() => {},
}
Done::Trained(Err("parked".to_string()))
});
self.ids.insert(handle.id(), id);
abandon
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dataset::{Example, Id, Question, rewrite};
use crate::tui::snapshots::{MOVED, dataset, project};
type TestResult = Result<(), Box<dyn std::error::Error>>;
const LIMIT: Duration = Duration::from_secs(30);
#[tokio::test]
async fn a_save_task_writes_the_configuration_and_ends_with_it() -> TestResult {
use crate::tui::snapshots::{PROJECT_CONFIG, project_env};
let dir = tempfile::tempdir()?;
std::fs::write(dir.path().join("overbrainer.toml"), PROJECT_CONFIG)?;
let text = PROJECT_CONFIG.replace("rust_expert", "rust_pro");
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.spawn(
TaskId(5),
Task::SaveConfig {
text: text.clone(),
base: PROJECT_CONFIG.to_string(),
env: project_env(),
},
);
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(5), Ok(Done::ConfigSaved(Ok(config))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(config.settings.project.name, "rust_pro");
assert_eq!(config.text, text);
assert_eq!(
std::fs::read_to_string(dir.path().join("overbrainer.toml"))?,
text
);
Ok(())
}
#[tokio::test]
async fn a_task_started_after_a_reload_reads_the_settings_with_its_env() -> TestResult {
let dir = crate::tui::snapshots::project()?;
let config = format!(
"{}\n[training]\ntarget = \"homelab\"\nbase_model = \"Qwen/Qwen3-4B\"\n\
adapter = \"qlora\"\nhub_model_id = \"me/model\"\n\n[targets.homelab]\n\
kind = \"ssh\"\nruntime = \"docker\"\n",
crate::tui::snapshots::CONFIG
);
std::fs::write(dir.path().join("overbrainer.toml"), config)?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
let mut warned = Vec::new();
for (id, env) in [
(TaskId(1), Vec::new()),
(
TaskId(2),
vec![("OVERBRAINER_HF_TOKEN".to_string(), "hf_x".to_string())],
),
] {
tasks.use_config(Source::from(EnvSource::Vars(env)));
tasks.spawn(id, Task::Prepare);
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((ended, Ok(Done::Prepared(Ok(plan))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(ended, id);
warned.push(
plan.warnings
.iter()
.any(|w| w.contains("OVERBRAINER_HF_TOKEN")),
);
}
assert_eq!(warned, [true, false], "the token set by the new env");
Ok(())
}
#[test]
fn a_look_reads_the_files_again_only_once_they_changed() -> TestResult {
let dir = crate::tui::snapshots::project()?;
let seen = stamp(dir.path());
let none = check_config(dir.path(), seen, &DotenvKeys::default());
assert!(none.read.is_none());
assert_eq!(none.stamp, seen);
std::fs::write(
dir.path().join(crate::config::DOTENV_FILE),
"OVERBRAINER_PROJECT__NAME=renamed\n",
)?;
let checked = check_config(dir.path(), seen, &DotenvKeys::default());
assert_ne!(checked.stamp, seen);
let Some(Ok(reread)) = checked.read else {
return Err(format!("not read again: {:?}", checked.read).into());
};
assert_eq!(reread.config.settings.project.name, "renamed");
assert_eq!(reread.config.stamp, Some(checked.stamp));
assert!(reread.config.env.contains("project.name"));
Ok(())
}
#[tokio::test]
async fn a_load_reads_the_data_files_and_the_history_cost_and_ends_with_its_id() -> TestResult {
let dir = tempfile::tempdir()?;
let files = DataFiles::new(dir.path());
rewrite(&files.subtopics, &dataset().subtopics)?;
let stats = crate::events::StageStats {
cost: Some(0.5),
..crate::events::StageStats::default()
};
let span = history::Span {
started_at: "2026-09-27T10:00:00Z".into(),
ended_at: "2026-09-27T10:01:00Z".into(),
};
let entry = history::Entry::from_stats(
crate::events::Stage::Answers,
span,
history::Status::Ok,
None,
&stats,
);
history::append(dir.path(), &entry)?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
assert!(tasks.is_empty());
tasks.spawn(TaskId(7), Task::Load);
assert!(!tasks.is_empty());
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(7), Ok(Done::Loaded(Ok(data), Ok(history))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
let Some(Cost::Known(cost)) = history.cost else {
return Err(format!("unexpected history: {history:?}").into());
};
assert_eq!(data.subtopics.len(), 3);
let answers = history.stages.get(&crate::events::Stage::Answers);
assert_eq!(answers.map(|total| total.cost), Some(Cost::Known(0.5)));
assert!(history.models.is_empty(), "no model recorded");
assert!((cost - 0.5).abs() < f64::EPSILON, "{cost}");
assert!(tasks.is_empty());
assert!(tasks.next().await.is_none());
Ok(())
}
#[tokio::test]
async fn a_load_of_a_broken_file_ends_with_its_error() -> TestResult {
let dir = tempfile::tempdir()?;
let files = DataFiles::new(dir.path());
if let Some(parent) = files.answers.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(&files.answers, "not json\n")?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.spawn(TaskId(1), Task::Load);
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(1), Ok(Done::Loaded(Err(error), Ok(history))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert!(error.contains("answers.jsonl"), "{error}");
assert_eq!(history, History::default());
Ok(())
}
#[test]
fn a_saved_edit_rewrites_the_files_then_train_and_eval() -> TestResult {
let dir = project()?;
let files = DataFiles::new(dir.path());
let subtopic = Id::subtopic("ownership", "Borrowing");
let edit = Edit::Change(Edited::Question {
id: Id::question(&subtopic, MOVED),
before: MOVED.into(),
text: "What happens to a borrow after a move?".into(),
});
let saved = save(
dir.path(),
&edit,
&Source::from(EnvSource::Vars(Vec::new())),
)?;
assert_eq!(
saved.message,
"question saved; its old answer was deleted: 1 question unanswered, run answers to answer it"
);
let report = saved.split?;
assert_eq!((report.train, report.eval, report.orphaned), (1, 1, 1));
let questions: Vec<Question> = crate::dataset::read(&files.questions)?;
assert!(
questions
.iter()
.any(|q| q.text == "What happens to a borrow after a move?")
);
let train: Vec<Example> = crate::dataset::read(&files.train)?;
assert_eq!(train.len(), 1);
Ok(())
}
#[test]
fn a_broken_config_refuses_the_edit_before_writing() -> TestResult {
let dir = project()?;
let files = DataFiles::new(dir.path());
let before = std::fs::read(&files.answers)?;
std::fs::write(dir.path().join("overbrainer.toml"), "[project\n")?;
let edit = Edit::Delete {
deletion: Deletion::Answer(Id::question(
&Id::subtopic("ownership", "Borrowing"),
MOVED,
)),
counts: Counts {
questions: 0,
answers: 1,
},
};
assert!(
save(
dir.path(),
&edit,
&Source::from(EnvSource::Vars(Vec::new()))
)
.is_err()
);
assert_eq!(std::fs::read(&files.answers)?, before);
assert!(!files.train.exists());
Ok(())
}
#[tokio::test]
async fn a_cancel_task_runs_the_command_line_flow() -> TestResult {
let dir = project()?;
let run = "20260921-133200-a1b2";
let mut record = crate::tui::snapshots::run(run, "homelab", crate::runs::RunState::Failed);
record.job = None;
crate::runs::Runs::new(dir.path()).save(&record)?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.spawn(TaskId(3), Task::Train(TrainJob::Cancel(run.into())));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(3), Ok(Done::Trained(Err(error))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(error, format!("run {run} has not started"));
assert!(tasks.is_empty());
Ok(())
}
#[tokio::test]
async fn push_tasks_refuse_as_the_command_does() -> TestResult {
let dir = project()?;
let run = "20260921-133200-a1b2";
crate::runs::Runs::new(dir.path()).save(&crate::tui::snapshots::run(
run,
"homelab",
crate::runs::RunState::Running,
))?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.use_config(Source::from(EnvSource::Vars(Vec::new())));
tasks.spawn(TaskId(1), Task::PreparePush(run.into()));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(1), Ok(Done::PushPrepared(Err(error))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(
error,
format!("run {run} is running: only a succeeded or stopped run can be pushed")
);
tasks.spawn(TaskId(2), Task::Push(run.into()));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(2), Ok(Done::Pushed(Err(error))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(
error,
"OVERBRAINER_HF_TOKEN is not set: a Hugging Face token with write access is needed \
to push"
);
assert!(tasks.is_empty());
Ok(())
}
#[tokio::test]
async fn runs_and_metrics_are_read_in_tasks() -> TestResult {
let dir = tempfile::tempdir()?;
let run = "20260921-133200-a1b2";
let runs = crate::runs::Runs::new(dir.path());
runs.save(&crate::tui::snapshots::run(
run,
"homelab",
crate::runs::RunState::Running,
))?;
std::fs::write(
runs.run_dir(run)?.join(crate::train::METRICS_FILE),
"{\"event\": \"log\", \"time\": 2, \"step\": 1, \"loss\": 1.5}\n",
)?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.spawn(TaskId(1), Task::Runs);
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(1), Ok(Done::Runs(Listing { runs: Ok(rows), .. })))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(rows.len(), 1);
tasks.spawn(TaskId(2), Task::Series(run.into()));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(2), Ok(Done::Series { series, .. }))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(series.map(|(s, _)| s.len()), Some(1));
Ok(())
}
#[test]
fn abandon_sets_the_flag_then_cancels_the_token() {
let mut tasks = Tasks::new(
Path::new("/nonexistent"),
tokio::sync::mpsc::unbounded_channel().0,
);
let token = CancellationToken::new();
let flag = Arc::new(AtomicBool::new(false));
tasks.tokens.insert(TaskId(1), token.clone());
tasks.abandons.insert(TaskId(1), Arc::clone(&flag));
tasks.abandon(TaskId(2));
assert!(!flag.load(std::sync::atomic::Ordering::SeqCst));
tasks.abandon(TaskId(1));
assert!(flag.load(std::sync::atomic::Ordering::SeqCst));
assert!(token.is_cancelled());
}
fn copy(message: &Msg) -> Option<Msg> {
match message {
Msg::Event(id, event) => Some(Msg::Event(*id, event.clone())),
Msg::Lagged(id, skipped) => Some(Msg::Lagged(*id, *skipped)),
Msg::Report(id, report) => Some(Msg::Report(*id, report.clone())),
Msg::EditorExited(_)
| Msg::BrowserFailed(_)
| Msg::LogsExported(_)
| Msg::NewerRelease(_) => None,
}
}
#[derive(Default)]
struct Seen(std::sync::Mutex<Vec<Event>>);
impl Observer for Seen {
fn event(&self, _bus: usize, event: &Event) {
if let Ok(mut seen) = self.0.lock() {
seen.push(event.clone());
}
}
fn closed(&self, _bus: usize) {}
}
#[tokio::test]
async fn the_observer_sees_tasks_started_with_a_reloaded_configuration() -> TestResult {
let dir = project()?;
let seen = Arc::new(Seen::default());
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
tasks.observe(Some(Arc::clone(&seen) as Arc<dyn Observer>));
tasks.use_config(Source::from(EnvSource::Vars(Vec::new())));
tasks.spawn(TaskId(1), Task::Pipeline(Command::Split));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(1), Ok(Done::Pipeline(Ok(()))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
let seen = seen.0.lock().map_err(|e| e.to_string())?;
assert!(
seen.iter().any(|event| matches!(
event,
Event::StageFinished {
stage: Stage::Split,
..
}
)),
"{seen:?}"
);
Ok(())
}
#[tokio::test]
async fn a_split_task_sends_every_message_before_it_ends() -> TestResult {
use crossterm::event::KeyCode;
use crate::events::Stage;
use crate::tui::app::Effect;
use crate::tui::pipeline::StageState;
use crate::tui::snapshots::{app, key};
let dir = project()?;
let (messages, mut inbox) = tokio::sync::mpsc::unbounded_channel();
let mut tasks = Tasks::new(dir.path(), messages);
tasks.spawn(TaskId(1), Task::Pipeline(Command::Split));
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
let Some((TaskId(1), Ok(Done::Pipeline(Ok(()))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
let mut received = Vec::new();
while let Ok(message) = inbox.try_recv() {
received.push(message);
}
let started = received.iter().any(|message| {
matches!(
message,
Msg::Event(
TaskId(1),
Event::StageStarted {
stage: Stage::Split,
..
}
)
)
});
let finished = received.iter().any(|message| {
matches!(
message,
Msg::Event(
TaskId(1),
Event::StageFinished {
stage: Stage::Split,
..
}
)
)
});
assert!(started && finished, "{received:?}");
let line = received.iter().find_map(|message| match message {
Msg::Report(TaskId(1), Report::Line(line)) => Some(line.clone()),
_ => None,
});
let Some(line) = line else {
return Err(format!("no split line: {received:?}").into());
};
assert!(line.starts_with("split: "), "{line}");
for done_first in [false, true] {
let mut app = app();
let mut effects = Vec::new();
for code in [
KeyCode::Char('r'),
KeyCode::Down,
KeyCode::Down,
KeyCode::Down,
KeyCode::Down,
KeyCode::Enter,
] {
effects.extend(app.on_input(&key(code)));
}
assert_eq!(
effects,
[Effect::Spawn(TaskId(1), Task::Pipeline(Command::Split))]
);
if done_first {
app.on_done(TaskId(1), Ok(Done::Pipeline(Ok(()))));
}
for message in received.iter().filter_map(copy) {
app.on_message(message);
}
if !done_first {
app.on_done(TaskId(1), Ok(Done::Pipeline(Ok(()))));
}
assert_eq!(
app.pipeline.results,
std::slice::from_ref(&line),
"done first: {done_first}"
);
let row = app.pipeline.row(Stage::Split);
assert_eq!(row.state, StageState::Done, "done first: {done_first}");
assert_eq!(row.finished, row.total, "done first: {done_first}");
assert!(row.total > 0);
}
Ok(())
}
const FAKE_AXOLOTL: &str = r#"#!/bin/sh
here="$(dirname "$0")/.."
[ "$1" = train ] || exit 0
i=0
until [ -f "$here/release" ]; do
i=$((i + 1))
[ "$i" -gt 600 ] && exit 3
sleep 0.1
done
printf '{"event": "begin", "time": 1, "max_steps": 2}\n' >> "$OVERBRAINER_METRICS"
printf '{"event": "log", "time": 2, "step": 1, "epoch": 0.5, "max_steps": 2, "loss": 1.5}\n' >> "$OVERBRAINER_METRICS"
printf '{"event": "log", "time": 3, "step": 2, "epoch": 1.0, "max_steps": 2, "eval_loss": 1.25}\n' >> "$OVERBRAINER_METRICS"
mkdir -p output && echo adapter > output/adapter_model.safetensors
exit 0
"#;
fn local_project() -> Result<tempfile::TempDir, Box<dyn std::error::Error>> {
use std::os::unix::fs::PermissionsExt;
let dir = project()?;
let venv = dir.path().join("venv");
std::fs::create_dir_all(venv.join("bin"))?;
std::fs::write(venv.join("bin/axolotl"), FAKE_AXOLOTL)?;
std::fs::set_permissions(
venv.join("bin/axolotl"),
std::fs::Permissions::from_mode(0o755),
)?;
let files = DataFiles::new(dir.path());
std::fs::write(&files.train, "{\"id\": 1}\n")?;
std::fs::write(&files.eval, "{\"id\": 2}\n")?;
let config = format!(
"{}\n[training]\ntarget = \"here\"\nbase_model = \"Qwen/Qwen3-4B\"\n\
adapter = \"qlora\"\n\n[targets.here]\nkind = \"local\"\nruntime = \"native\"\n\
venv = \"{}\"\n",
crate::tui::snapshots::CONFIG,
venv.display()
);
std::fs::write(dir.path().join("overbrainer.toml"), config)?;
Ok(dir)
}
async fn next_message(
inbox: &mut tokio::sync::mpsc::UnboundedReceiver<Msg>,
) -> Result<Option<Msg>, tokio::time::error::Elapsed> {
tokio::time::timeout(Duration::from_secs(60), inbox.recv()).await
}
#[tokio::test]
async fn a_training_task_reports_its_run_and_detaches_without_cancelling_it() -> TestResult {
let dir = local_project()?;
let (messages, mut inbox) = tokio::sync::mpsc::unbounded_channel();
let mut tasks = Tasks::new(dir.path(), messages);
tasks.spawn(
TaskId(1),
Task::Train(TrainJob::Start {
vram_floor: VramFloor::ToEstimate,
}),
);
let mut created = None;
loop {
match next_message(&mut inbox).await? {
Some(Msg::Report(TaskId(1), Report::RunCreated(id))) => created = Some(id),
Some(Msg::Event(TaskId(1), Event::JobStatus(_))) => break,
Some(_) => {},
None => return Err("the task ended before its job started".into()),
}
}
let run = created.ok_or("no RunCreated report before the job started")?;
tasks.cancel(TaskId(1));
let ended = tokio::time::timeout(Duration::from_secs(60), tasks.next()).await?;
let Some((TaskId(1), Ok(Done::Trained(Err(error))))) = ended else {
return Err(format!("{ended:?}").into());
};
assert_eq!(
error,
format!(
"interrupted: run {run} keeps running on target `here`; follow it again with \
`overbrainer train attach {run}`"
)
);
let record = crate::runs::Runs::new(dir.path()).load(&run)?;
assert_eq!(
record.state,
crate::runs::RunState::Running,
"detached, not cancelled"
);
std::fs::write(dir.path().join("venv/release"), "")?;
tasks.spawn(TaskId(2), Task::Train(TrainJob::Attach(run.clone())));
let ended = tokio::time::timeout(Duration::from_secs(60), tasks.next()).await?;
assert!(
matches!(ended, Some((TaskId(2), Ok(Done::Trained(Ok(())))))),
"{ended:?}"
);
drop(tasks);
let (mut metrics, mut lines) = (0, Vec::new());
while let Some(message) = next_message(&mut inbox).await? {
match message {
Msg::Event(TaskId(2), Event::Metric(_)) => metrics += 1,
Msg::Report(TaskId(2), Report::Line(line)) => lines.push(line),
_ => {},
}
}
assert_eq!(metrics, 2, "every metric of the run, from its first");
assert_eq!(lines.len(), 2, "{lines:?}");
assert!(
lines[0].starts_with(&format!("train: run {run} succeeded; ")),
"the command line's own outcome line: {lines:?}"
);
assert!(
lines[1].ends_with(&format!(" in runs/{run}/output")),
"then where the model is: {lines:?}"
);
Ok(())
}
#[tokio::test]
async fn a_catalog_lookup_is_aborted_when_the_tui_ends() -> TestResult {
let dir = tempfile::tempdir()?;
let mut tasks = Tasks::new(dir.path(), tokio::sync::mpsc::unbounded_channel().0);
let handle = tasks.set.spawn(std::future::pending::<Done>());
tasks.ids.insert(handle.id(), TaskId(9));
tasks.lookups.insert(TaskId(9), handle);
tasks.abort_lookups();
let next = tokio::time::timeout(LIMIT, tasks.next()).await?;
assert!(matches!(next, Some((TaskId(9), Err(_)))), "{next:?}");
Ok(())
}
}