use std::collections::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::editor::Edited;
use super::start::{Prices, StartPlan, list_prices, prepare};
use super::training::{Listing, list_runs, read_series};
use crate::cli::data::Command;
use crate::cli::front::{Frontend, Report};
use crate::cli::{StageArgs, TrainArgs, TrainCommand};
use crate::config::EnvSource;
use crate::dataset::{Counts, DataFiles, Dataset, Deletion};
use crate::events::{Event, EventBus};
use crate::pipeline::{Ctx, SplitReport};
use crate::prompts::Prompts;
use crate::train::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),
Prepare,
Prices(Vec<String>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum TrainJob {
Start,
Attach(String),
Cancel(String),
}
impl TrainJob {
fn args(&self) -> TrainArgs {
let command = match self {
Self::Start => return TrainArgs::default(),
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)]
pub(super) enum Done {
Loaded(Result<Dataset, String>),
Saved(Result<Saved, String>),
Pipeline(Result<(), String>),
Trained(Result<(), String>),
Runs(Listing),
Series {
run: String,
series: Option<Vec<TrainMetric>>,
},
Prepared(Result<StartPlan, String>),
Prices(Prices),
}
#[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>),
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>,
) -> (Frontend, JoinHandle<()>) {
let bus = EventBus::with_capacity(BUS_CAPACITY);
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, env: EnvSource) -> Result<Saved, String> {
let error = |error: anyhow::Error| format!("{error:#}");
let settings = crate::config::load(project_dir, env).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>,
}
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(),
}
}
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()
}
pub(super) fn spawn(&mut self, id: TaskId, task: Task) {
let files = DataFiles::new(&self.project_dir);
let handle = match task {
Task::Load => self.set.spawn(async move {
let read = tokio::task::spawn_blocking(move || Dataset::read(&files)).await;
Done::Loaded(match read {
Ok(Ok(data)) => Ok(data),
Ok(Err(error)) => Err(format!("{:#}", anyhow::Error::from(error))),
Err(error) => Err(format!("cannot read the data files: {error}")),
})
}),
Task::Edit(edit) => {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let saved =
tokio::task::spawn_blocking(move || save(&dir, &edit, EnvSource::Process))
.await;
Done::Saved(match saved {
Ok(saved) => saved,
Err(error) => Err(format!("the edit failed: {error}")),
})
})
},
Task::Pipeline(command) => self.spawn_pipeline(id, command),
Task::Train(job) => self.spawn_train(id, job),
Task::Runs => {
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(),
}))
})
},
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::Prepare => {
let dir = self.project_dir.clone();
self.set.spawn(async move {
let plan =
tokio::task::spawn_blocking(move || prepare(&dir, EnvSource::Process))
.await;
Done::Prepared(match plan {
Ok(plan) => plan,
Err(error) => Err(format!("cannot prepare the run: {error}")),
})
})
},
Task::Prices(gpu_types) => {
let dir = self.project_dir.clone();
let handle = self.set.spawn(async move {
Done::Prices(list_prices(&dir, EnvSource::Process, gpu_types).await)
});
self.lookups.insert(id, handle.clone());
handle
},
};
self.ids.insert(handle.id(), id);
}
fn spawn_pipeline(&mut self, id: TaskId, command: Command) -> AbortHandle {
let dir = self.project_dir.clone();
let token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) = front_end(id, &self.messages, &token, &abandon);
self.tokens.insert(id, token.clone());
self.set.spawn(async move {
let args = StageArgs::default();
let outcome = {
let run = crate::cli::data::run(&dir, command, &args, &front);
tokio::select! {
result = run => result.map_err(|error| format!("{error:#}")),
() = token.cancelled() => {
Err("interrupted: the stage resumes on its next run".to_string())
},
}
};
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 token = CancellationToken::new();
let abandon = Arc::new(AtomicBool::new(false));
let (front, forwarder) = front_end(id, &self.messages, &token, &abandon);
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).await;
forwarded(front, forwarder, "training").await;
Done::Trained(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_load_reads_the_data_files_and_ends_with_its_id() -> TestResult {
let dir = tempfile::tempdir()?;
let files = DataFiles::new(dir.path());
rewrite(&files.subtopics, &dataset().subtopics)?;
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))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert_eq!(data.subtopics.len(), 3);
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))))) = next else {
return Err(format!("unexpected end: {next:?}").into());
};
assert!(error.contains("answers.jsonl"), "{error}");
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, 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, 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 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(_) => None,
}
}
#[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::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));
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(), 1);
assert!(
lines[0].starts_with(&format!("train: run {run} succeeded; ")),
"the command line's own outcome line: {lines:?}"
);
Ok(())
}
#[tokio::test]
async fn a_price_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(())
}
}