diff --git a/Cargo.lock b/Cargo.lock index 1ad77cf..a75fe51 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3091,8 +3091,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracel-artifact" -version = "0.11.0" -source = "git+https://github.com/tracel-ai/tracel?rev=f52ee9e5c76153abf2ca750c31427d44dd86cc0b#f52ee9e5c76153abf2ca750c31427d44dd86cc0b" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5242b42fceb3cef1141e3094a7a26b24b9310d8ae51df19c3b42ce5f11b70fd6" dependencies = [ "crossbeam", "reqwest", @@ -3104,7 +3105,7 @@ dependencies = [ [[package]] name = "tracel-cli" -version = "0.6.0" +version = "0.6.1" dependencies = [ "anyhow", "cargo_metadata", @@ -3136,9 +3137,8 @@ dependencies = [ [[package]] name = "tracel-client" -version = "0.22.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "456275c8a4a1fd65e147aac5e39a34de32070ad8efe03929cb7cf6104fe2f407" +version = "0.23.0" +source = "git+https://github.com/tracel-ai/tracel-client?rev=36edfd1bf8191724bb9e07ae7f1ed3c6f55194dd#36edfd1bf8191724bb9e07ae7f1ed3c6f55194dd" dependencies = [ "directories", "reqwest", @@ -3155,8 +3155,9 @@ dependencies = [ [[package]] name = "tracel-job" -version = "0.11.0" -source = "git+https://github.com/tracel-ai/tracel?rev=f52ee9e5c76153abf2ca750c31427d44dd86cc0b#f52ee9e5c76153abf2ca750c31427d44dd86cc0b" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d78edfde9f0a780eadca19ca21913519271d2b0e5a6673107c0e20cffab5c55" dependencies = [ "clap", "clap_complete", @@ -3757,7 +3758,7 @@ checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" [[package]] name = "xtask" -version = "0.6.0" +version = "0.6.1" dependencies = [ "strum 0.27.2", "tracel-xtask", diff --git a/Cargo.toml b/Cargo.toml index 08f8dee..9acf2d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,7 @@ members = ["crates/*", "xtask"] [workspace.package] edition = "2024" -version = "0.6.0" +version = "0.6.1" readme = "README.md" license = "MIT OR Apache-2.0" rust-version = "1.97.0" @@ -48,12 +48,12 @@ rstest = { version = "0.27.0", default-features = false } ## External Crate -tracel-client = "0.22.0" +tracel-client = { git = "https://github.com/tracel-ai/tracel-client", rev = "36edfd1bf8191724bb9e07ae7f1ed3c6f55194dd" } tracel-artifact = "0.12.0" tracel-job = "0.12.0" ## Crate -tracel-cli = { version = "0.6.0", path = "crates/tracel-cli" } +tracel-cli = { version = "0.6.1", path = "crates/tracel-cli" } ### For xtask crate ### tracel-xtask = { version = "=5.0.0" } diff --git a/crates/tracel-cli/src/commands/jobs.rs b/crates/tracel-cli/src/commands/jobs.rs index b2e96d5..27d6235 100644 --- a/crates/tracel-cli/src/commands/jobs.rs +++ b/crates/tracel-cli/src/commands/jobs.rs @@ -263,7 +263,7 @@ pub fn handle_command(args: JobsArgs, context: CliContext) -> anyhow::Result Ok(read_job(&client()?, &project, args.num)?.into()), JobsCommands::Logs(args) if args.follow => { let start = args.start.unwrap_or(0); - let job = follow_logs(&client()?, context.output(), &project, args.num, start)?; + let job = follow_logs(&client()?, &context, &project, args.num, start)?; context.terminal().outro(&state(&job)); Ok(Outcome::streamed()) } diff --git a/crates/tracel-cli/src/commands/run/remote.rs b/crates/tracel-cli/src/commands/run/remote.rs index 798698e..72b6bc2 100644 --- a/crates/tracel-cli/src/commands/run/remote.rs +++ b/crates/tracel-cli/src/commands/run/remote.rs @@ -163,7 +163,7 @@ pub fn queue(run: RemoteRun, context: &CliContext) -> anyhow::Result { let queued = format!("Queued job {num} with code version {digest}."); if run.follow { terminal.success(&queued); - follow_job(&client, context.output(), linked, num)?; + follow_job(&client, context, linked, num)?; terminal.outro(&format!("Job {num} completed.")); return Ok(Outcome::streamed()); } diff --git a/crates/tracel-cli/src/helpers/job_logs.rs b/crates/tracel-cli/src/helpers/job_logs.rs new file mode 100644 index 0000000..0f06d79 --- /dev/null +++ b/crates/tracel-cli/src/helpers/job_logs.rs @@ -0,0 +1,277 @@ +//! Following a job's logs until it finishes: live, as its compute provider streams them, +//! while it runs, and from the log stored once it has ended. + +use std::io::{self, Write}; +use std::time::{Duration, Instant}; + +use serde::Serialize; +use tracel_client::WebSocketClient; +use tracel_client::console::Client; +use tracel_client::console::job::response::{JobLogResponse, JobLogStreamMessage, JobResponse}; +use tracel_client::websocket::WebSocketError; + +use super::jobs::{final_status, missing_job, read_job, require_completed}; +use crate::context::CliContext; +use crate::tools::tracel_config::TracelProject; +use crate::ui::{Human, Output, Render}; + +/// How often a followed job is read again once its logs are caught up. +const FOLLOW_INTERVAL: Duration = Duration::from_secs(2); + +/// How long to wait before reading a live log that had nothing new again. +const PAUSE_BETWEEN_EMPTY_LIVE_LOG_READS: Duration = Duration::from_millis(100); + +/// An event of following a job's logs. As JSON, a page of logs gains `"type": "log"`. +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum LogEvent<'a> { + Log(&'a JobLogResponse), + LiveLog { + logs: &'a str, + }, + End { + status: &'a str, + /// Whether the logs so far end inside a line, which the text then ends. + #[serde(skip)] + mid_line: bool, + }, +} + +impl Render for LogEvent<'_> { + fn render(&self, out: &mut Human<'_>) -> io::Result<()> { + match self { + Self::Log(page) => out.write_all(page.logs.as_bytes()), + Self::LiveLog { logs } => out.write_all(logs.as_bytes()), + Self::End { mid_line: true, .. } => writeln!(out), + Self::End { .. } => Ok(()), + } + } +} + +/// Find where the next read of the logs starts after `page`, read from `start`, and +/// whether more is already there to read. +fn next_read(start: u64, page: &JobLogResponse) -> anyhow::Result<(u64, bool)> { + let next = page.end.max(start); + if page.has_more && next == start { + anyhow::bail!("The Tracel Console has more logs to send, but not from where."); + } + Ok((next, page.has_more)) +} + +/// The log a compute provider uploads once the job has ended, written from a byte on. +struct StoredLogFollower { + start: u64, + has_written_logs: bool, + ends_mid_line: bool, +} + +impl StoredLogFollower { + fn write_what_is_new( + &mut self, + client: &Client, + output: &Output, + project: &TracelProject, + num: i32, + ) -> anyhow::Result<()> { + loop { + let page = client + .get_job_logs(&project.namespace, &project.name, num, Some(self.start)) + .map_err(|error| missing_job(error, project, num))?; + if !page.logs.is_empty() { + output.event(&LogEvent::Log(&page))?; + self.has_written_logs = true; + self.ends_mid_line = !page.logs.ends_with('\n'); + } + let more; + (self.start, more) = next_read(self.start, &page)?; + if !more { + return Ok(()); + } + } + } +} + +enum LiveStream { + NotOpened, + Open(Box), + Unavailable, +} + +/// The lines a running job's compute provider streams, written as they come. +struct LiveLogFollower { + stream: LiveStream, + has_written_lines: bool, + ends_mid_line: bool, +} + +impl LiveLogFollower { + /// Write the lines that come until `deadline`, opening the stream first when it isn't + /// open. A message that isn't a line is skipped. A stream that fails is opened again + /// the next time; one that can't be opened is given up with a warning. + fn write_lines_until( + &mut self, + deadline: Instant, + client: &Client, + context: &CliContext, + project: &TracelProject, + num: i32, + ) -> anyhow::Result<()> { + if let LiveStream::NotOpened = self.stream { + self.stream = open_or_warn(client, context, project, num); + } + let LiveStream::Open(stream) = &mut self.stream else { + return Ok(()); + }; + while Instant::now() < deadline { + match stream.receive::() { + Ok(Some(JobLogStreamMessage::Log(line))) => { + context.output().event(&LogEvent::LiveLog { logs: &line })?; + self.has_written_lines = true; + self.ends_mid_line = !line.ends_with('\n'); + } + Ok(None) => std::thread::sleep( + PAUSE_BETWEEN_EMPTY_LIVE_LOG_READS + .min(deadline.saturating_duration_since(Instant::now())), + ), + Err(WebSocketError::SerializationError(_)) => {} + Err(_) => { + self.stream = LiveStream::NotOpened; + return Ok(()); + } + } + } + Ok(()) + } +} + +fn open_or_warn( + client: &Client, + context: &CliContext, + project: &TracelProject, + num: i32, +) -> LiveStream { + match client.open_job_log_websocket(&project.namespace, &project.name, num) { + Ok(stream) => LiveStream::Open(Box::new(stream)), + Err(error) => { + context.terminal().warning(&format!( + "Can't follow job {num}'s log live: {error}. It shows once the job has ended." + )); + LiveStream::Unavailable + } + } +} + +/// Write job `num`'s logs from byte `start` until the job has finished, and return it as +/// last read. Its compute provider uploads the stored log once the job has ended, so while +/// the stored log is empty, the lines the provider streams are written live instead, and +/// the stored log isn't written after them. The job is read before the stored log, so the +/// stored log read after a final status is complete. +pub fn follow_logs( + client: &Client, + context: &CliContext, + project: &TracelProject, + num: i32, + start: u64, +) -> anyhow::Result { + let output = context.output(); + let mut stored_log = StoredLogFollower { + start, + has_written_logs: false, + ends_mid_line: false, + }; + let mut live_log = LiveLogFollower { + stream: LiveStream::NotOpened, + has_written_lines: false, + ends_mid_line: false, + }; + let job = loop { + let job = read_job(client, project, num)?; + if !live_log.has_written_lines { + stored_log.write_what_is_new(client, output, project, num)?; + } + if final_status(&job.status).is_some() { + break job; + } + let next_read_at = Instant::now() + FOLLOW_INTERVAL; + if !stored_log.has_written_logs { + live_log.write_lines_until(next_read_at, client, context, project, num)?; + } + std::thread::sleep(next_read_at.saturating_duration_since(Instant::now())); + }; + let ends_mid_line = match live_log.has_written_lines { + true => live_log.ends_mid_line, + false => stored_log.ends_mid_line, + }; + output.event(&LogEvent::End { + status: &job.status, + mid_line: ends_mid_line, + })?; + if live_log.has_written_lines { + context.terminal().info(&format!( + "The whole log, with what came before following it: `tracel jobs logs {num}`." + )); + } + Ok(job) +} + +/// Follow job `num`'s logs from the start until it finishes, and fail unless it +/// completed. +pub fn follow_job( + client: &Client, + context: &CliContext, + project: &TracelProject, + num: i32, +) -> anyhow::Result { + let job = follow_logs(client, context, project, num, 0)?; + require_completed(&job)?; + Ok(job) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::plain; + + fn page(end: u64, has_more: bool) -> JobLogResponse { + JobLogResponse { + logs: "epoch 1\n".into(), + start: 0, + end, + total_size: 4096, + has_more, + } + } + + #[test] + fn given_a_page_when_reading_on_then_the_next_read_starts_where_it_ends() { + let next = next_read(0, &page(8, true)).unwrap(); + + assert_eq!(next, (8, true)); + } + + #[test] + fn given_a_page_that_ends_before_the_start_when_reading_on_then_the_start_is_kept() { + let next = next_read(16, &page(8, false)).unwrap(); + + assert_eq!(next, (16, false)); + } + + #[test] + fn given_more_logs_but_no_progress_when_reading_on_then_it_fails() { + let next = next_read(8, &page(8, true)); + + assert!(next.is_err()); + } + + #[test] + fn given_logs_that_end_mid_line_when_the_job_ends_then_the_line_is_ended() { + let end = LogEvent::End { + status: "completed", + mid_line: true, + }; + + let text = plain(&end); + + assert_eq!(text, "\n"); + } +} diff --git a/crates/tracel-cli/src/helpers/jobs.rs b/crates/tracel-cli/src/helpers/jobs.rs index 686ec98..9122098 100644 --- a/crates/tracel-cli/src/helpers/jobs.rs +++ b/crates/tracel-cli/src/helpers/jobs.rs @@ -2,22 +2,16 @@ //! reading a job until it finishes. use std::borrow::Cow; -use std::io::{self, Write}; use std::time::{Duration, Instant}; -use serde::Serialize; use serde_json::Value; use tracel_client::ClientError; use tracel_client::console::Client; -use tracel_client::console::job::response::{JobLogResponse, JobResponse}; +use tracel_client::console::job::response::JobResponse; use super::{Resource, map_resource_error}; use crate::error::{CliError, ErrorKind}; use crate::tools::tracel_config::TracelProject; -use crate::ui::{Human, Output, Render}; - -/// How often a followed job is read again once its logs are caught up. -const FOLLOW_INTERVAL: Duration = Duration::from_secs(2); /// Write the job command that gives the program `arguments` unchanged. /// @@ -144,94 +138,12 @@ pub fn poll_job( } } -/// An event of following a job's logs. As JSON, a page of logs gains `"type": "log"`. -#[derive(Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum LogEvent<'a> { - Log(&'a JobLogResponse), - End { - status: &'a str, - /// Whether the logs so far end inside a line, which the text then ends. - #[serde(skip)] - mid_line: bool, - }, -} - -impl Render for LogEvent<'_> { - fn render(&self, out: &mut Human<'_>) -> io::Result<()> { - match self { - Self::Log(page) => out.write_all(page.logs.as_bytes()), - Self::End { mid_line: true, .. } => writeln!(out), - Self::End { .. } => Ok(()), - } - } -} - -/// Find where the next read of the logs starts after `page`, read from `start`, and -/// whether more is already there to read. -fn next_read(start: u64, page: &JobLogResponse) -> anyhow::Result<(u64, bool)> { - let next = page.end.max(start); - if page.has_more && next == start { - anyhow::bail!("The Tracel Console has more logs to send, but not from where."); - } - Ok((next, page.has_more)) -} - -/// Write job `num`'s logs from byte `start` until the job has finished, and return it as -/// last read. The job is read before its logs, so the logs read after a final status are -/// complete. -pub fn follow_logs( - client: &Client, - output: &Output, - project: &TracelProject, - num: i32, - start: u64, -) -> anyhow::Result { - let mut start = start; - let mut mid_line = false; - let job = poll_job(client, project, num, FOLLOW_INTERVAL, None, |_| { - loop { - let page = client - .get_job_logs(&project.namespace, &project.name, num, Some(start)) - .map_err(|error| missing_job(error, project, num))?; - if !page.logs.is_empty() { - output.event(&LogEvent::Log(&page))?; - mid_line = !page.logs.ends_with('\n'); - } - let more; - (start, more) = next_read(start, &page)?; - if !more { - return Ok(()); - } - } - })?; - output.event(&LogEvent::End { - status: &job.status, - mid_line, - })?; - Ok(job) -} - -/// Follow job `num`'s logs from the start until it finishes, and fail unless it -/// completed. -pub fn follow_job( - client: &Client, - output: &Output, - project: &TracelProject, - num: i32, -) -> anyhow::Result { - let job = follow_logs(client, output, project, num, 0)?; - require_completed(&job)?; - Ok(job) -} - #[cfg(test)] mod tests { use rstest::rstest; use serde_json::json; use super::*; - use crate::test_support::plain; fn command(arguments: &[&str]) -> String { let arguments: Vec = arguments @@ -241,16 +153,6 @@ mod tests { job_command(&arguments) } - fn page(end: u64, has_more: bool) -> JobLogResponse { - JobLogResponse { - logs: "epoch 1\n".into(), - start: 0, - end, - total_size: 4096, - has_more, - } - } - #[test] fn given_plain_arguments_when_writing_a_job_command_then_they_are_joined_unchanged() { let written = command(&["train", "--epochs", "10", "--lr=0.01", "$HOME", "*.rs"]); @@ -314,37 +216,4 @@ mod tests { assert_eq!(invocation, None); } - - #[test] - fn given_a_page_when_reading_on_then_the_next_read_starts_where_it_ends() { - let next = next_read(0, &page(8, true)).unwrap(); - - assert_eq!(next, (8, true)); - } - - #[test] - fn given_a_page_that_ends_before_the_start_when_reading_on_then_the_start_is_kept() { - let next = next_read(16, &page(8, false)).unwrap(); - - assert_eq!(next, (16, false)); - } - - #[test] - fn given_more_logs_but_no_progress_when_reading_on_then_it_fails() { - let next = next_read(8, &page(8, true)); - - assert!(next.is_err()); - } - - #[test] - fn given_logs_that_end_mid_line_when_the_job_ends_then_the_line_is_ended() { - let end = LogEvent::End { - status: "completed", - mid_line: true, - }; - - let text = plain(&end); - - assert_eq!(text, "\n"); - } } diff --git a/crates/tracel-cli/src/helpers/mod.rs b/crates/tracel-cli/src/helpers/mod.rs index 3ab5797..757af64 100644 --- a/crates/tracel-cli/src/helpers/mod.rs +++ b/crates/tracel-cli/src/helpers/mod.rs @@ -1,4 +1,5 @@ mod download; +mod job_logs; mod jobs; mod model; mod project; @@ -6,9 +7,10 @@ mod resources; mod transfer; pub use download::{DownloadFile, DownloadResult, download_files, validate_rel_path}; +pub use job_logs::{follow_job, follow_logs}; pub use jobs::{ - FinalStatus, final_status, follow_job, follow_logs, job_command, job_invocation, missing_job, - poll_job, read_job, require_completed, + FinalStatus, final_status, job_command, job_invocation, missing_job, poll_job, read_job, + require_completed, }; pub use model::ensure_model_exists; pub use project::{ diff --git a/docs/cli.md b/docs/cli.md index 0f138b2..89d79f2 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -50,7 +50,10 @@ and the instructions a login asks you to follow. instead: each log entry, with `"type":"log"` added, then `{"type":"end","running":false}` when the experiment has finished. `tracel jobs logs --follow --json` does the same with each page of logs, then -`{"type":"end","status":""}`, and `tracel run --remote --follow --json` +`{"type":"end","status":""}`. While the job runs with nothing in its +log file, each line its compute provider streams comes as +`{"type":"live_log","logs":""}` instead, and no page follows them. +`tracel run --remote --follow --json` starts with the queued job, with `"type":"queued"` added, then writes the job's log events. An error still writes the failure line. @@ -216,7 +219,12 @@ time left. - `jobs logs ` reads one page of the job's log file, and shows it as it is, then the command that reads the next page when there is more. `--start ` reads from a byte offset, and a page's `end` is the next - page's `start`. `--follow` reads until the job has finished. + page's `start`. `--follow` reads until the job has finished. The compute + provider uploads the log file once the job has ended, so while it is empty, + `--follow` shows the lines the provider streams as they come, the way the web + console does, and doesn't show the file after them. Those lines start from + when following began; `jobs logs ` reads the whole file once the job + has ended. - `jobs cancel ` cancels a new, queued or running job, after asking; `--yes` skips the question. A running job stays `pending_cancellation` until its compute provider stops it.