diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 00d71fff..14bd9d33 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,6 +2,7 @@ name: CI env: CARGO_TERM_COLOR: always + FEATURES: tracel/station,tracel/burn,tracel/schema # For now we execute CI only on PR to save on CI time on: @@ -33,16 +34,16 @@ jobs: run: cargo xtask check format # -------------------------------------------------------------------------------- - name: Lint - run: cargo xtask check lint + run: cargo xtask check --features "$FEATURES" lint # -------------------------------------------------------------------------------- - name: Typos uses: tracel-ai/github-actions/check-typos@v11 # -------------------------------------------------------------------------------- - name: Unit Tests - run: cargo xtask test unit + run: cargo xtask test --features "$FEATURES" unit # -------------------------------------------------------------------------------- - name: Integration Tests - run: cargo xtask test integration + run: cargo xtask test --features "$FEATURES" integration # -------------------------------------------------------------------------------- - name: Documentation Tests - run: cargo xtask doc tests + run: cargo xtask doc --features "$FEATURES" tests diff --git a/Cargo.lock b/Cargo.lock index 6134d814..9e4662e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7226,6 +7226,7 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" name = "tracel" version = "0.12.0" dependencies = [ + "rstest", "serde", "serde_json", "tempfile", @@ -7280,9 +7281,11 @@ dependencies = [ "chrono", "crossbeam", "reqwest", + "rstest", "serde", "serde_json", "serde_with", + "tempfile", "thiserror 2.0.21", "toml", "tracel-artifact", @@ -7312,6 +7315,8 @@ version = "0.12.0" dependencies = [ "burn", "chrono", + "relative-path", + "rstest", "serde", "serde_json", "tempfile", @@ -7355,6 +7360,7 @@ dependencies = [ "clap", "clap_complete", "referencing", + "rstest", "schemars 1.2.2", "serde", "serde_json", @@ -7400,6 +7406,7 @@ dependencies = [ "clap_complete", "json-patch", "reqwest", + "rstest", "schemars 1.2.2", "serde", "serde_json", @@ -7416,6 +7423,7 @@ name = "tracel-station" version = "0.12.0" dependencies = [ "chrono", + "rstest", "serde", "serde_json", "serde_with", diff --git a/Cargo.toml b/Cargo.toml index 14bbd594..6457925e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,8 @@ categories = ["development-tools"] burn = { version = "0.22", default-features = false, features = [ "train" ] } rand = "0.9.2" +relative-path = "1.9.3" +rstest = "0.27.0" tempfile = "3.27.0" url = "2.5.8" diff --git a/README.md b/README.md index 6b8338b4..9e789834 100644 --- a/README.md +++ b/README.md @@ -91,10 +91,8 @@ input, merged onto its default, as its arguments, which the console lists as its The experiments `Target::experiments` builds in a job, on the thread that runs it, are linked to that job, even for runs started on other threads: each one started from them records the job's -number and input, and prints where it records to stderr as it starts, such as -`Recording experiment 42: `, which `tracel run` and a remote job's logs show. Experiments -built another way, such as `ProjectHandle::experiments`, are not linked unless the code links -them, through `Experiments::attribute`, `arguments` and `on_start`. +number and input. Experiments built another way, such as `ProjectHandle::experiments`, are not +linked unless the code links them, through `Experiments::attribute` and `arguments`. The [`cli`](examples/basics/examples/cli.rs) example in [`examples/basics`](examples/basics) is a complete, runnable version. @@ -103,7 +101,7 @@ complete, runnable version. | Variable | Value | Default | | --- | --- | --- | -| `TRACEL_TARGET` | `offline` or `console` | `offline` | +| `TRACEL_TARGET` | `offline`, `console`, or `station` with the `station` feature of `tracel` | `offline` | | `TRACEL_RUNS_DIR` | the directory offline runs are recorded under | `./runs` | | `TRACEL_ENV` | the console to reach | `Production` | | `TRACEL_API_KEY` | an API key or a job token | the stored `tracel login` session | @@ -159,11 +157,12 @@ Options: - A nested field's flag joins the keys with dots and writes `_` as `-`: the field `optimizer.weight_decay` is `--optimizer.weight-decay`. A field that is an array, a map, an optional object or an enum with data takes a JSON literal, such as `--layers '[64, 32]'`. A - boolean flag given no value is `true`. + boolean flag given no value is `true`, and takes a value only as `--flag=value`, such as + `--shuffle=false`. - The input is the default input, then the `--config` file, then the JSON document, then the flags, each merged onto the one before, so a later one wins. Launchers pass the JSON document. -- `help`, `version` and `config` are the command line's names: a field with one of them has no - flag, and is set through `--config` or the JSON document. +- `help` and `config` are the command line's names: a field with one of them has no flag, and is + set through `--config` or the JSON document. - `--completions ` prints the completion script for `bash`, `elvish`, `fish`, `powershell` or `zsh`: `cargo run -- --completions bash > cli.bash`. diff --git a/crates/tracel-artifact/src/bundle/memory/tests.rs b/crates/tracel-artifact/src/bundle/memory/tests.rs index f5ae034e..2bb7707d 100644 --- a/crates/tracel-artifact/src/bundle/memory/tests.rs +++ b/crates/tracel-artifact/src/bundle/memory/tests.rs @@ -96,18 +96,6 @@ impl BundleDecode for ExampleConfig { let config: ExampleConfig = serde_json::from_str(&config_content) .map_err(|e| format!("Failed to parse config JSON: {}", e))?; - // Optionally validate metadata exists (but don't require it for decoding) - if let Ok(mut metadata_reader) = source.open(&settings.metadata_filename) { - let mut metadata_content = String::new(); - if metadata_reader - .read_to_string(&mut metadata_content) - .is_ok() - && let Ok(_metadata) = serde_json::from_str::(&metadata_content) - { - // Could validate version compatibility here - } - } - Ok(config) } } diff --git a/crates/tracel-console/Cargo.toml b/crates/tracel-console/Cargo.toml index aa7ce2c0..a05825d2 100644 --- a/crates/tracel-console/Cargo.toml +++ b/crates/tracel-console/Cargo.toml @@ -29,3 +29,7 @@ tracel-experiment-remote = { workspace = true } tracel-inference = { workspace = true } tracel-models = { workspace = true } url = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } diff --git a/crates/tracel-console/src/console.rs b/crates/tracel-console/src/console.rs index 91b82508..1b2f7ac5 100644 --- a/crates/tracel-console/src/console.rs +++ b/crates/tracel-console/src/console.rs @@ -1,11 +1,16 @@ +use std::ffi::OsString; use std::fmt; +use std::path::Path; use std::sync::{Arc, Mutex, OnceLock, PoisonError}; use std::time::{Duration, SystemTime}; use tracel_artifact::ReqwestTransferClient; use tracel_client::{ ClientError, - console::{Client, Env, TracelCredentials, user::response::UserResponseSchema}, + console::{ + Client, Env, FileSessionStore, SessionStoreError, TracelCredentials, + user::response::UserResponseSchema, + }, }; use tracel_datasets::DatasetRegistry; use tracel_experiment::Experiments; @@ -14,7 +19,7 @@ use tracel_models::ModelRegistry; use url::Url; use crate::datasets::ConsoleDatasetOps; -use crate::env::{CredentialSource, ProjectRef, env_from_environment}; +use crate::env::{CredentialSource, ProjectRef, env_from_environment, env_from_vars}; use crate::experiment::ConsoleExperimentProvider; use crate::inference::ConsoleInferenceProvider; use crate::models::ConsoleModelOps; @@ -92,8 +97,6 @@ pub struct ProjectScope { console: Console, pub owner: String, pub project: String, - /// The kind of namespace that owns the project, once read from the console. - owner_kind: OnceLock, } impl ProjectScope { @@ -111,19 +114,9 @@ impl ProjectScope { pub fn fetch(&self) -> Result { let project = self.client()?.get_project(&self.owner, &self.project)?; let project = Project::try_from(project)?; - let _ = self.owner_kind.set(project.namespace.kind); Ok(project) } - /// Returns the kind of namespace that owns the project, fetching the project on the first - /// call. - pub fn owner_kind(&self) -> Result { - match self.owner_kind.get() { - Some(kind) => Ok(*kind), - None => self.fetch().map(|project| project.namespace.kind), - } - } - /// Returns the client that moves files to and from presigned URLs. pub fn transfer_client(&self) -> &ReqwestTransferClient { &self.console.inner.transfer_client @@ -236,7 +229,6 @@ impl Console { console: self.clone(), owner: owner.into(), project: project.into(), - owner_kind: OnceLock::new(), }), } } @@ -301,11 +293,28 @@ impl ProjectHandle { /// The owner namespace is `TRACEL_NAMESPACE` and the project name `TRACEL_PROJECT`; either /// one unset is read from `namespace` or `project` in `tracel.toml` in the current /// directory. Fails with [`ConsoleError::NoNamespace`] or [`ConsoleError::NoProject`] when - /// neither names it. + /// neither names it, and with [`ConsoleError::InvalidTracelToml`] when `tracel.toml` + /// cannot be read. pub fn from_env() -> Result { - let console = Console::from_env()?; - let project = ProjectRef::from_env()?; - Ok(console.project(project.namespace, project.name)) + Self::from_vars( + |name| std::env::var_os(name), + Path::new("."), + FileSessionStore::for_server, + ) + } + + /// [`from_env`](Self::from_env), reading the variables from `lookup`, `tracel.toml` from + /// `dir`, and the stored `tracel login` session from the store `sessions` gives for the + /// console's URL. + pub fn from_vars( + lookup: impl Fn(&str) -> Option, + dir: &Path, + sessions: impl FnOnce(&Url) -> Result, + ) -> Result { + let env = env_from_vars(&lookup)?; + let credentials = CredentialSource::from_vars(&lookup).resolve_with(&env, sessions)?; + let project = ProjectRef::from_vars(&lookup, dir)?; + Ok(Console::connect(env, credentials).project(project.namespace, project.name)) } /// Returns the console this project is reached through. @@ -385,23 +394,4 @@ mod tests { assert_send_sync::(); assert_send_sync::(); } - - #[test] - fn binding_and_scoping_send_no_request() { - let console = Console::connect( - Env::Development, - TracelCredentials::api_key("tcl_key_not_checked_until_the_first_request"), - ); - let project = console.project("owner", "project"); - let _ = ( - project.datasets(), - project.models(), - project.experiments(), - project.inference(), - ); - - assert!(console.inner.connected.get().is_none()); - assert!(Arc::ptr_eq(&project.console().inner, &console.inner)); - assert_eq!(console.base_url(), &Env::Development.get_url()); - } } diff --git a/crates/tracel-console/src/datasets.rs b/crates/tracel-console/src/datasets.rs index 6065764f..82b5ea50 100644 --- a/crates/tracel-console/src/datasets.rs +++ b/crates/tracel-console/src/datasets.rs @@ -456,143 +456,69 @@ fn map_version_error( #[cfg(test)] mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + use super::*; - #[test] - fn one_batch_of_neighbours_is_one_run() { - assert_eq!(contiguous_runs(&[4, 5, 6]), vec![4..7]); + fn item( + payload: usize, + source_item_id: Option<&str>, + extra: Option, + ) -> DatasetVersionUploadItemRequest { + DatasetVersionUploadItemRequest { + source_item_id: source_item_id.map(str::to_string), + example_payload: vec![7; payload], + annotation: extra.clone(), + metadata: extra, + } } - #[test] - fn a_shuffled_batch_is_sorted_and_coalesced() { - assert_eq!(contiguous_runs(&[9, 1, 8, 0, 2]), vec![0..3, 8..10]); + fn wire_size(value: &impl serde::Serialize) -> usize { + serde_json::to_vec(value).unwrap().len() } - #[test] - fn a_repeated_index_is_asked_for_once() { - assert_eq!(contiguous_runs(&[3, 3, 3]), vec![3..4]); + #[rstest] + #[case::neighbours(&[4, 5, 6], vec![4..7])] + #[case::shuffled(&[9, 1, 8, 0, 2], vec![0..3, 8..10])] + #[case::repeated(&[3, 3, 3], vec![3..4])] + fn indexes_are_asked_for_in_ascending_contiguous_runs( + #[case] indexes: &[u64], + #[case] runs: Vec>, + ) { + assert_eq!(contiguous_runs(indexes), runs); } - #[test] - fn the_estimate_never_undercounts_what_is_actually_sent() { - let payloads = [0usize, 1, 2, 3, 4, 5, 1_000, 1_001]; - let ids = [ - None, - Some(String::new()), - Some("plain".to_string()), - // JSON escapes make the wire form longer than the Rust string. - Some("\"quoted\"".to_string()), - Some("back\\slash".to_string()), - Some("tab\there\nand\nnewlines".to_string()), - Some("\u{1}\u{2}\u{3} control".to_string()), - Some("emoji \u{1F600} and accents \u{e9}\u{e8}".to_string()), - Some("\"".repeat(64)), - ]; - let extras = [ - None, - Some(serde_json::json!(null)), - Some(serde_json::json!({"label": "cat", "nested": [1, 2, 3]})), - Some(serde_json::json!({"quote": "say \"hi\"", "uni": "\u{e9}"})), - ]; - - for payload in payloads { - for id in &ids { - for extra in &extras { - let item = DatasetVersionUploadItemRequest { - source_item_id: id.clone(), - example_payload: vec![7; payload], - annotation: extra.clone(), - metadata: extra.clone(), - }; - - let actual = serde_json::to_vec(&item).expect("an item serializes").len(); - let estimate = encoded_size(&item); - - assert!( - estimate >= actual, - "undercounted by {}: payload={payload} id={id:?} extra={extra:?}", - actual - estimate - ); - } - } - } + #[rstest] + #[case::empty(item(0, None, None))] + #[case::one_byte_past_a_base64_block(item(1, None, None))] + #[case::two_bytes_past_a_base64_block(item(2, None, None))] + #[case::a_whole_base64_block(item(3, None, None))] + #[case::a_large_payload(item(3_000, None, None))] + #[case::an_id_json_escapes(item(0, Some("\"quoted\" back\\slash\ttab\nline \u{1}"), None))] + #[case::an_id_beyond_ascii(item(0, Some("emoji \u{1F600} and accents \u{e9}"), None))] + #[case::null_extras(item(0, None, Some(json!(null))))] + #[case::nested_extras(item(0, None, Some(json!({"quote": "say \"hi\"", "nested": [1, 2, 3]}))))] + fn an_items_estimate_never_undercounts_what_is_sent( + #[case] item: DatasetVersionUploadItemRequest, + ) { + assert!(encoded_size(&item) >= wire_size(&item)); } #[test] fn the_budget_bounds_the_whole_request_body_not_just_the_items() { let items: Vec<_> = (0..500) - .map(|n| DatasetVersionUploadItemRequest { - source_item_id: Some(format!("item-\"{n}\"")), - example_payload: vec![n as u8; n * 7 % 1_000], - annotation: Some(serde_json::json!({ "label": n })), - metadata: None, + .map(|n| { + item( + n * 7 % 1_000, + Some(format!("item-\"{n}\"").as_str()), + Some(json!({ "label": n })), + ) }) .collect(); - let budgeted: usize = items.iter().map(encoded_size).sum(); - let request = AddDatasetVersionUploadItemsRequest { - items: items.clone(), - }; - let actual = serde_json::to_vec(&request) - .expect("a batch serializes") - .len(); - - assert!( - budgeted >= actual, - "a batch of {} items budgeted {budgeted} but sends {actual}", - items.len() - ); - } - #[test] - fn a_payload_is_measured_base64_encoded_not_raw() { - let item = DatasetVersionUploadItemRequest { - source_item_id: None, - example_payload: vec![0; 3_000], - annotation: None, - metadata: None, - }; - - // 3 raw bytes become 4 on the wire, so a batch budgeted on raw sizes overshoots. - assert_eq!(encoded_size(&item) - 96 - 4 - 4 - 4, 4_000); - } - - #[test] - fn a_payload_that_alone_exceeds_the_budget_is_refused() { - let raw = BATCH_BYTES / 4 * 3 + 1; - let item = DatasetVersionUploadItemRequest { - source_item_id: None, - example_payload: vec![0; raw], - annotation: None, - metadata: None, - }; - - assert!(encoded_size(&item) > BATCH_BYTES); - } - - #[test] - fn a_megabyte_item_fills_the_budget_long_before_a_large_count_does() { - let item = DatasetVersionUploadItemRequest { - source_item_id: None, - example_payload: vec![0; 1024 * 1024], - annotation: None, - metadata: None, - }; - - assert!(BATCH_BYTES / encoded_size(&item) < 64); - } - - #[test] - fn small_items_are_not_held_back_by_a_count() { - let item = DatasetVersionUploadItemRequest { - source_item_id: None, - example_payload: vec![0; 1024], - annotation: None, - metadata: None, - }; - - // Bytes alone decide, so a batch holds far more than the 256 it used to. - assert!(BATCH_BYTES / encoded_size(&item) > 10_000); + assert!(budgeted >= wire_size(&AddDatasetVersionUploadItemsRequest { items })); } #[test] @@ -611,22 +537,18 @@ mod tests { ); } - #[test] - fn queries_every_page_needed_to_reach_the_total() { - let mut fetched = Vec::new(); - let mut pages = vec![ - ((0..PAGE_SIZE).collect::>(), u64::from(PAGE_SIZE) + 1), - (vec![PAGE_SIZE], u64::from(PAGE_SIZE) + 1), - ] - .into_iter(); - - let values = collect_pages(|page| { - fetched.push(page); - Ok(pages.next().expect("the test provided this page")) - }) - .unwrap(); + #[rstest] + #[case::one_more_than_a_page(u64::from(PAGE_SIZE) + 1, u64::from(PAGE_SIZE) + 1)] + #[case::fewer_than_the_total_says(u64::from(PAGE_SIZE) + 1, 10 * u64::from(PAGE_SIZE))] + fn a_query_reads_every_item_across_its_pages(#[case] count: u64, #[case] total: u64) { + let all: Vec = (0..count).collect(); + let pages: Vec<&[u64]> = all.chunks(PAGE_SIZE as usize).collect(); + + let read = collect_pages(|page| { + let items = pages.get(page as usize).copied().unwrap_or_default(); + Ok((items.to_vec(), total)) + }); - assert_eq!(fetched, [0, 1]); - assert_eq!(values.len(), PAGE_SIZE as usize + 1); + assert_eq!(read.unwrap(), all); } } diff --git a/crates/tracel-console/src/env.rs b/crates/tracel-console/src/env.rs index fdf454f2..c9cf7700 100644 --- a/crates/tracel-console/src/env.rs +++ b/crates/tracel-console/src/env.rs @@ -1,17 +1,23 @@ //! Reading which console to reach, with which credential, and which project, from the //! environment and `tracel.toml`. +use std::ffi::OsString; +use std::io::ErrorKind; use std::path::Path; use std::sync::Arc; use serde::Deserialize; use tracel_client::ClientError; use tracel_client::console::auth::DeviceAuthClient; -use tracel_client::console::{AppSession, Env, FileSessionStore, SessionStore, TracelCredentials}; +use tracel_client::console::{ + AppSession, Env, FileSessionStore, SessionStore, SessionStoreError, TracelCredentials, +}; +use url::Url; use crate::ConsoleError; const TRACEL_ENV: &str = "TRACEL_ENV"; +const TRACEL_API_KEY: &str = "TRACEL_API_KEY"; const TRACEL_NAMESPACE: &str = "TRACEL_NAMESPACE"; const TRACEL_PROJECT: &str = "TRACEL_PROJECT"; const TRACEL_TOML: &str = "tracel.toml"; @@ -29,7 +35,12 @@ const TRACEL_ENV_VALUES: &str = /// `staging-n` spellings of the `tracel` CLI. Fails with [`ConsoleError::InvalidSetting`] for any /// other value. pub fn env_from_environment() -> Result { - let value = std::env::var_os(TRACEL_ENV).map(|value| value.to_string_lossy().into_owned()); + env_from_vars(|name| std::env::var_os(name)) +} + +/// [`env_from_environment`], reading `TRACEL_ENV` from `lookup`. +pub fn env_from_vars(lookup: impl Fn(&str) -> Option) -> Result { + let value = lookup(TRACEL_ENV).map(|value| value.to_string_lossy().into_owned()); parse_env(value.as_deref()) } @@ -42,6 +53,7 @@ fn parse_env(value: Option<&str>) -> Result { .strip_prefix("Staging(") .and_then(|number| number.strip_suffix(')')) .or_else(|| value.strip_prefix("staging-")) + .filter(|number| number.bytes().all(|byte| byte.is_ascii_digit())) .and_then(|number| number.parse::().ok()) .map(Env::Staging) .ok_or_else(|| ConsoleError::InvalidSetting { @@ -65,7 +77,14 @@ pub enum CredentialSource { impl CredentialSource { /// `TRACEL_API_KEY` when it is set and not empty, else the stored `tracel login` session. pub fn from_env() -> Self { - TracelCredentials::from_env().map_or(Self::StoredLogin, Self::Explicit) + Self::from_vars(|name| std::env::var_os(name)) + } + + /// [`from_env`](Self::from_env), reading `TRACEL_API_KEY` from `lookup`. + pub fn from_vars(lookup: impl Fn(&str) -> Option) -> Self { + non_empty(lookup(TRACEL_API_KEY)).map_or(Self::StoredLogin, |key| { + Self::Explicit(TracelCredentials::api_key(key)) + }) } /// The credential this source names for `env`'s console. @@ -73,16 +92,28 @@ impl CredentialSource { /// Reads a stored session from disk, without network I/O. Fails with /// [`ConsoleError::NoCredentials`] when `tracel login` stored none for that console. pub fn resolve(&self, env: &Env) -> Result { + self.resolve_with(env, FileSessionStore::for_server) + } + + /// [`resolve`](Self::resolve), reading a stored session from the store `sessions` gives for + /// the console's URL. Only [`StoredLogin`](Self::StoredLogin) asks `sessions` for it. + pub fn resolve_with( + &self, + env: &Env, + sessions: impl FnOnce(&Url) -> Result, + ) -> Result { match self { Self::Explicit(credentials) => Ok(credentials.clone()), - Self::StoredLogin => stored_login(env), + Self::StoredLogin => { + let store = sessions(&env.get_url()).map_err(ClientError::from)?; + stored_login(env, store) + } } } } -/// The app session `tracel login` stored for `env`'s console. -fn stored_login(env: &Env) -> Result { - let store = FileSessionStore::for_server(&env.get_url()).map_err(ClientError::from)?; +/// The app session `tracel login` stored in `store` for `env`'s console. +fn stored_login(env: &Env, store: FileSessionStore) -> Result { if store.load().map_err(ClientError::from)?.is_none() { return Err(ConsoleError::NoCredentials); } @@ -117,36 +148,40 @@ impl ProjectRef { /// The owner namespace is `TRACEL_NAMESPACE` and the project name `TRACEL_PROJECT`; either /// one unset is read from `namespace` or `project` in `tracel.toml` in the current /// directory. Fails with [`ConsoleError::NoNamespace`] or [`ConsoleError::NoProject`] when - /// neither names it. + /// neither names it, and with [`ConsoleError::InvalidTracelToml`] when `tracel.toml` + /// cannot be read. pub fn from_env() -> Result { - let namespace = non_empty_var(TRACEL_NAMESPACE); - let project = non_empty_var(TRACEL_PROJECT); + Self::from_vars(|name| std::env::var_os(name), Path::new(".")) + } + + /// [`from_env`](Self::from_env), reading the variables from `lookup` and `tracel.toml` from + /// `dir`. + pub fn from_vars( + lookup: impl Fn(&str) -> Option, + dir: &Path, + ) -> Result { + let namespace = non_empty(lookup(TRACEL_NAMESPACE)); + let project = non_empty(lookup(TRACEL_PROJECT)); let file = if namespace.is_some() && project.is_some() { TracelToml::default() } else { - TracelToml::read(Path::new(TRACEL_TOML)) + TracelToml::read(&dir.join(TRACEL_TOML))? }; - locate_project(namespace, project, file) - } -} + let namespace = namespace + .or(file.namespace) + .ok_or(ConsoleError::NoNamespace)?; + let name = project.or(file.project).ok_or(ConsoleError::NoProject)?; -fn locate_project( - namespace: Option, - project: Option, - file: TracelToml, -) -> Result { - let namespace = namespace - .or(file.namespace) - .ok_or(ConsoleError::NoNamespace)?; - let name = project.or(file.project).ok_or(ConsoleError::NoProject)?; - - Ok(ProjectRef { namespace, name }) + Ok(Self { namespace, name }) + } } -/// A variable's value; an empty one reads as unset. -fn non_empty_var(name: &str) -> Option { - std::env::var(name).ok().filter(|value| !value.is_empty()) +/// A variable's value; an empty or non-UTF-8 one reads as unset. +fn non_empty(value: Option) -> Option { + value + .and_then(|value| value.into_string().ok()) + .filter(|value| !value.is_empty()) } /// The project a `tracel.toml` names with its `namespace` and `project` keys. @@ -157,164 +192,254 @@ struct TracelToml { } impl TracelToml { - /// A missing, unreadable or malformed file names nothing. - fn read(path: &Path) -> Self { - std::fs::read_to_string(path) - .map(|contents| Self::parse(&contents)) - .unwrap_or_default() - } - - fn parse(contents: &str) -> Self { - toml::from_str(contents).unwrap_or_default() + /// The project the file at `path` names; a missing file names nothing. + /// + /// Fails with [`ConsoleError::InvalidTracelToml`] when the file cannot be read or is not + /// valid TOML. + fn read(path: &Path) -> Result { + let invalid = |reason: String| ConsoleError::InvalidTracelToml { + path: path.to_path_buf(), + reason, + }; + match std::fs::read_to_string(path) { + Ok(contents) => toml::from_str(&contents).map_err(|error| invalid(error.to_string())), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(Self::default()), + Err(error) => Err(invalid(error.to_string())), + } } } #[cfg(test)] mod tests { + use std::time::SystemTime; + + use rstest::{fixture, rstest}; + use tempfile::TempDir; + use tracel_client::console::{AccessToken, RefreshToken, StoredAppSession}; + use super::*; + const API_KEY: &str = "tcl_key"; + fn located( namespace: Option<&str>, project: Option<&str>, tracel_toml: &str, ) -> Result { - locate_project( - namespace.map(str::to_string), - project.map(str::to_string), - TracelToml::parse(tracel_toml), - ) + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("tracel.toml"), tracel_toml).unwrap(); + let lookup = |name: &str| match name { + TRACEL_NAMESPACE => namespace.map(OsString::from), + TRACEL_PROJECT => project.map(OsString::from), + _ => None, + }; + ProjectRef::from_vars(lookup, dir.path()) } - #[test] - fn the_environment_names_the_project_over_tracel_toml() { - let file = "namespace = \"file-owner\"\nproject = \"file-project\""; + const FILE_NAMESPACE: &str = "file-owner"; + const FILE_PROJECT: &str = "file-project"; + + #[fixture] + fn tracel_toml() -> String { + format!("namespace = \"{FILE_NAMESPACE}\"\nproject = \"{FILE_PROJECT}\"") + } + #[rstest] + fn the_environment_names_the_project_over_tracel_toml(tracel_toml: String) { assert_eq!( - located(Some("env-owner"), Some("env-project"), file).unwrap(), + located(Some("env-owner"), Some("env-project"), &tracel_toml).unwrap(), ProjectRef::new("env-owner", "env-project") ); } - #[test] - fn tracel_toml_fills_in_what_the_environment_leaves_unset() { - let file = "namespace = \"file-owner\"\nproject = \"file-project\""; + #[rstest] + #[case::namespace_unset( + None, + Some("env-project"), + ProjectRef::new(FILE_NAMESPACE, "env-project") + )] + #[case::namespace_empty( + Some(""), + Some("env-project"), + ProjectRef::new(FILE_NAMESPACE, "env-project") + )] + #[case::project_unset(Some("env-owner"), None, ProjectRef::new("env-owner", FILE_PROJECT))] + #[case::both_unset(None, None, ProjectRef::new(FILE_NAMESPACE, FILE_PROJECT))] + fn tracel_toml_fills_in_what_the_environment_leaves_unset( + tracel_toml: String, + #[case] namespace: Option<&str>, + #[case] project: Option<&str>, + #[case] expected: ProjectRef, + ) { + assert_eq!(located(namespace, project, &tracel_toml).unwrap(), expected); + } - assert_eq!( - located(Some("env-owner"), None, file).unwrap(), - ProjectRef::new("env-owner", "file-project") - ); - assert_eq!( - located(None, Some("env-project"), file).unwrap(), - ProjectRef::new("file-owner", "env-project") - ); - assert_eq!( - located(None, None, file).unwrap(), - ProjectRef::new("file-owner", "file-project") - ); + #[test] + fn no_namespace_in_the_environment_or_tracel_toml_is_an_error() { + assert!(matches!( + located(None, Some("env-project"), ""), + Err(ConsoleError::NoNamespace) + )); } #[test] - fn tracel_toml_names_the_project_only_by_namespace_and_project() { - let file = "name = \"file-project\"\nowner = \"file-owner\""; - - let no_namespace = located(None, None, file).unwrap_err(); - let no_project = located(Some("env-owner"), None, file).unwrap_err(); - - assert!(matches!(no_namespace, ConsoleError::NoNamespace)); - assert!(matches!(no_project, ConsoleError::NoProject)); - for error in [no_namespace, no_project] { - assert!( - error.to_string().contains("`namespace` and `project`"), - "{error}" - ); - } + fn no_project_in_the_environment_or_tracel_toml_is_an_error() { + assert!(matches!( + located(Some("env-owner"), None, ""), + Err(ConsoleError::NoProject) + )); } #[test] - fn a_malformed_tracel_toml_names_nothing() { + fn a_missing_tracel_toml_names_nothing() { + let dir = tempfile::tempdir().unwrap(); + assert!(matches!( - located(None, Some("env-project"), "namespace = "), + ProjectRef::from_vars(|_| None, dir.path()), Err(ConsoleError::NoNamespace) )); } - #[test] - fn a_missing_namespace_or_project_says_where_to_set_it() { - let no_namespace = located(None, Some("env-project"), "").unwrap_err(); - let no_project = located(Some("env-owner"), None, "").unwrap_err(); + #[rstest] + #[case::malformed(|path: &Path| std::fs::write(path, "namespace = ").unwrap())] + #[case::unreadable(|path: &Path| std::fs::create_dir(path).unwrap())] + fn an_unreadable_or_malformed_tracel_toml_is_an_error_naming_it(#[case] write: fn(&Path)) { + let dir = tempfile::tempdir().unwrap(); + let tracel_toml = dir.path().join(TRACEL_TOML); + write(&tracel_toml); - assert!(matches!(no_namespace, ConsoleError::NoNamespace)); - assert!(no_namespace.to_string().contains("TRACEL_NAMESPACE")); - assert!(no_namespace.to_string().contains("tracel.toml")); - assert!(matches!(no_project, ConsoleError::NoProject)); - assert!(no_project.to_string().contains("TRACEL_PROJECT")); - assert!(no_project.to_string().contains("tracel.toml")); + assert!(matches!( + ProjectRef::from_vars(|_| None, dir.path()), + Err(ConsoleError::InvalidTracelToml { path, .. }) if path == tracel_toml + )); } #[test] - fn an_explicit_credential_is_used_as_given() { - let source = CredentialSource::Explicit(TracelCredentials::api_key("tcl_key")); + fn tracel_toml_is_not_read_when_the_environment_names_the_project() { + assert_eq!( + located(Some("env-owner"), Some("env-project"), "namespace = ").unwrap(), + ProjectRef::new("env-owner", "env-project") + ); + } + fn api_key(value: Option<&str>) -> CredentialSource { + CredentialSource::from_vars(|name| match name { + TRACEL_API_KEY => value.map(OsString::from), + _ => None, + }) + } + + #[fixture] + fn sessions() -> TempDir { + tempfile::tempdir().unwrap() + } + + fn store_in(sessions: &TempDir) -> FileSessionStore { + FileSessionStore::new(sessions.path().join("session.json")) + } + + fn no_store(_: &Url) -> Result { + Err(SessionStoreError::new("no home directory")) + } + + #[test] + fn tracel_api_key_is_the_credential() { assert!(matches!( - source.resolve(&Env::Development), - Ok(TracelCredentials::ApiKey(key)) if key == "tcl_key" + api_key(Some(API_KEY)), + CredentialSource::Explicit(TracelCredentials::ApiKey(key)) if key == API_KEY )); } - #[test] - fn tracel_env_unset_or_empty_names_production() { - assert!(matches!(parse_env(None), Ok(Env::Production))); - assert!(matches!(parse_env(Some("")), Ok(Env::Production))); + #[rstest] + #[case::unset(None)] + #[case::empty(Some(""))] + fn without_tracel_api_key_the_credential_is_the_stored_login(#[case] value: Option<&str>) { + assert!(matches!(api_key(value), CredentialSource::StoredLogin)); } #[test] - fn tracel_env_accepts_the_spellings_of_the_tracel_cli() { - for value in ["Production", "production"] { - assert!( - matches!(parse_env(Some(value)), Ok(Env::Production)), - "{value}" - ); - } - for value in ["Development", "development"] { - assert!( - matches!(parse_env(Some(value)), Ok(Env::Development)), - "{value}" - ); - } - for value in ["Staging(2)", "staging-2"] { - assert!( - matches!(parse_env(Some(value)), Ok(Env::Staging(2))), - "{value}" - ); - } + fn a_stored_login_is_read_from_the_store_for_its_console() { + let env = Env::Staging(2); + let mut asked_for = None; + + let _ = CredentialSource::StoredLogin.resolve_with(&env, |url| { + asked_for = Some(url.clone()); + no_store(url) + }); + + assert_eq!(asked_for, Some(env.get_url())); + } + + #[rstest] + fn a_store_without_a_session_gives_no_credential(sessions: TempDir) { + assert!(matches!( + CredentialSource::StoredLogin + .resolve_with(&Env::Production, |_| Ok(store_in(&sessions))), + Err(ConsoleError::NoCredentials) + )); + } + + #[rstest] + fn a_stored_session_is_the_credential(sessions: TempDir) { + store_in(&sessions) + .save(&StoredAppSession { + access_token: AccessToken::new("tcl_at_stored"), + access_token_expires_at: SystemTime::UNIX_EPOCH, + refresh_token: RefreshToken::new("tcl_rt_stored"), + refresh_token_expires_at: SystemTime::UNIX_EPOCH, + }) + .unwrap(); + + assert!(matches!( + CredentialSource::StoredLogin + .resolve_with(&Env::Production, |_| Ok(store_in(&sessions))), + Ok(TracelCredentials::AppSession(_)) + )); } #[test] - fn an_unrecognized_tracel_env_is_an_error_naming_the_variable_and_value() { - for value in [ - "DEVELOPMENT", - "dev", - " production", - "Staging(x)", - "Staging()", - "Staging(1", - "Staging(256)", - "staging-", - "staging--1", - "staging-256", - ] { - let error = parse_env(Some(value)).unwrap_err(); - - assert!( - matches!( - &error, - ConsoleError::InvalidSetting { variable: "TRACEL_ENV", value: held, .. } - if held == value - ), - "{value}: {error:?}" - ); - assert!(error.to_string().contains("TRACEL_ENV"), "{error}"); - assert!(error.to_string().contains(value), "{error}"); - } + fn an_explicit_credential_needs_no_session_store() { + let source = CredentialSource::Explicit(TracelCredentials::api_key(API_KEY)); + + assert!(matches!( + source.resolve_with(&Env::Production, no_store), + Ok(TracelCredentials::ApiKey(key)) if key == API_KEY + )); + } + + #[rstest] + #[case::unset(None)] + #[case::empty(Some(""))] + fn an_unset_or_empty_tracel_env_names_production(#[case] value: Option<&str>) { + assert!(matches!(parse_env(value), Ok(Env::Production))); + } + + #[rstest] + #[case::production("Production", Env::Production)] + #[case::production_as_the_cli_spells_it("production", Env::Production)] + #[case::development("Development", Env::Development)] + #[case::development_as_the_cli_spells_it("development", Env::Development)] + #[case::staging("Staging(2)", Env::Staging(2))] + #[case::staging_as_the_cli_spells_it("staging-2", Env::Staging(2))] + fn tracel_env_accepts_the_spellings_of_the_tracel_cli(#[case] value: &str, #[case] env: Env) { + assert_eq!(parse_env(Some(value)).unwrap().get_url(), env.get_url()); + } + + #[rstest] + #[case::wrong_case("DEVELOPMENT")] + #[case::abbreviated("dev")] + #[case::padded(" production")] + #[case::unclosed("Staging(1")] + #[case::without_a_number("staging-")] + #[case::not_a_number("Staging(x)")] + #[case::signed("staging-+2")] + #[case::signed_in_parentheses("Staging(+2)")] + #[case::past_255("Staging(256)")] + fn an_unrecognized_tracel_env_is_an_invalid_setting_naming_it(#[case] value: &str) { + assert!(matches!( + parse_env(Some(value)), + Err(ConsoleError::InvalidSetting { variable: TRACEL_ENV, value: held, .. }) + if held == value + )); } } diff --git a/crates/tracel-console/src/error.rs b/crates/tracel-console/src/error.rs index 1737adba..f8c10510 100644 --- a/crates/tracel-console/src/error.rs +++ b/crates/tracel-console/src/error.rs @@ -1,3 +1,5 @@ +use std::path::PathBuf; + use tracel_client::error::ClientError; /// Errors produced while configuring, authenticating with or calling the console. @@ -42,6 +44,14 @@ pub enum ConsoleError { (expected keys: `namespace` and `project`)" )] NoProject, + /// The `tracel.toml` that would name the project cannot be read, or is not valid TOML. + #[error("could not read {}: {reason}", path.display())] + InvalidTracelToml { + /// The file. + path: PathBuf, + /// Why it cannot be used. + reason: String, + }, /// The credential cannot be used for this request: an API key reaches project data only. #[error("this credential cannot be used for this request; API keys reach project data only")] CredentialNotAllowed, diff --git a/crates/tracel-console/src/experiment.rs b/crates/tracel-console/src/experiment.rs index 4e505129..bb46087f 100644 --- a/crates/tracel-console/src/experiment.rs +++ b/crates/tracel-console/src/experiment.rs @@ -21,14 +21,13 @@ use tracel_experiment::reader::{ ArtifactRef, ExperimentArtifactReader, ExperimentReaderError, LoadedArtifact, }; use tracel_experiment::{ - ArtifactKind, CancelToken, ExperimentId, ExperimentLocation, ExperimentProvider, ExperimentRun, + ArtifactKind, CancelToken, ExperimentId, ExperimentProvider, ExperimentRun, ExperimentRunControl, }; use tracel_experiment_remote::{ArtifactUploadError, ArtifactUploader, RemoteExperimentSession}; -use url::Url; use crate::console::ProjectScope; -use crate::{Namespace, NamespaceKind}; +use crate::wire::console_timestamp; /// Experiment provider backed by the console's experiment run protocol. pub struct ConsoleExperimentProvider { @@ -93,50 +92,7 @@ fn create_run( let id = ExperimentId::from(format!("{experiment_num}")); let run = ExperimentRun::new_with_control(id, session, reader, control); - Ok(match page_of(scope, experiment_num) { - Some(page) => run.with_location(ExperimentLocation::Url(page.into())), - None => run, - }) -} - -/// The console page of experiment `num` of the project `scope` names, when the project's owner -/// can be read. -fn page_of(scope: &ProjectScope, num: i32) -> Option { - let kind = match scope.owner_kind() { - Ok(kind) => kind, - Err(error) => { - tracing::warn!("Could not find the console page of experiment {num}: {error}"); - return None; - } - }; - let owner = Namespace { - name: scope.owner.clone(), - kind, - }; - experiment_page(scope.console().base_url(), &owner, &scope.project, num) -} - -/// The page of experiment `num` of `project`, owned by `owner`, on the console whose API is at -/// `api_url`. -/// -/// The console serves its API under `api/` of its web address, and an experiment's page at -/// `users//projects//experiments/` of it, or under `orgs/` for a project -/// an organization owns. `None` when the API is not served under `api/`. -fn experiment_page(api_url: &Url, owner: &Namespace, project: &str, num: i32) -> Option { - let mut page = Url::parse(api_url.as_str().strip_suffix("api/")?).ok()?; - let owners = match owner.kind { - NamespaceKind::User => "users", - NamespaceKind::Organization => "orgs", - }; - page.path_segments_mut().ok()?.pop_if_empty().extend([ - owners, - &owner.name, - "projects", - project, - "experiments", - &num.to_string(), - ]); - Some(page) + Ok(run) } /// A scope for artifact operations within a specific experiment. @@ -250,7 +206,7 @@ impl ExperimentArtifactClient { Ok(bundle) } - /// The artifact of the experiment named `name`. + /// The artifact of the experiment named `name`, the newest when several are. fn fetch(&self, name: &str) -> Result { let listed = self.scope.client()?.list_artifacts_by_name( &self.scope.owner, @@ -258,26 +214,24 @@ impl ExperimentArtifactClient { self.experiment_num, name, )?; - named(listed.items, name, |artifact| &artifact.name) + newest_named(listed.items, name) } } -/// The one item of `items` whose name, which `name_of` gives, is `name`. +/// The newest of `artifacts` named exactly `name`. /// /// The console's name filter lists the artifacts whose names contain the name asked for, so -/// asking for `model` can list `model-2` as well. -fn named(items: Vec, name: &str, name_of: impl Fn(&T) -> &str) -> Result { - let mut matching = items.into_iter().filter(|item| name_of(item) == name); - let found = matching - .next() - .ok_or_else(|| ArtifactError::NotFound(name.to_owned()))?; - match matching.count() { - 0 => Ok(found), - others => Err(ArtifactError::Ambiguous { - name: name.to_owned(), - count: others + 1, - }), - } +/// asking for `model` can list `model-2` as well, and an experiment can hold several artifacts of +/// one name. +fn newest_named( + artifacts: Vec, + name: &str, +) -> Result { + artifacts + .into_iter() + .filter(|artifact| artifact.name == name) + .max_by_key(|artifact| console_timestamp(&artifact.created_at)) + .ok_or_else(|| ArtifactError::NotFound(name.to_owned())) } fn artifact_kind_name(kind: ArtifactKind) -> &'static str { @@ -292,8 +246,6 @@ fn artifact_kind_name(kind: ArtifactKind) -> &'static str { enum ArtifactError { #[error("Artifact not found: {0}")] NotFound(String), - #[error("Artifact name is ambiguous: {count} artifacts are named {name}")] - Ambiguous { name: String, count: usize }, #[error(transparent)] Client(#[from] ClientError), #[error(transparent)] @@ -383,104 +335,64 @@ impl ArtifactUploader for ConsoleArtifactUploader { #[cfg(test)] mod tests { - use tracel_client::console::Env; + use rstest::{fixture, rstest}; + use tracel_client::console::artifact::response::ArtifactSourceResponse; use super::*; - /// The id of the artifact `named` picks among `(id, name)` pairs. - fn pick( - artifacts: &[(&'static str, &'static str)], - name: &str, - ) -> Result<&'static str, ArtifactError> { - named(artifacts.to_vec(), name, |(_, name)| name).map(|(id, _)| id) + fn artifact(name: &str, created_at: &str) -> ArtifactResponse { + ArtifactResponse { + id: format!("{name} at {created_at}"), + created_at: created_at.to_string(), + name: name.to_string(), + kind: "model".to_string(), + bucket_id: "bucket".to_string(), + experiment: ArtifactSourceResponse { + id: 1, + experiment_num: 1, + }, + manifest: Value::Null, + } } - #[test] - fn an_artifact_is_picked_by_its_exact_name() { - let listed = [("2", "model-2"), ("1", "model"), ("3", "model.bpk")]; + #[fixture] + fn listed() -> Vec { + vec![ + artifact("model-2", "2026-10-01T00:00:00Z"), + artifact("model", "2026-10-01T00:00:00Z"), + artifact("model.bpk", "2026-10-01T00:00:00Z"), + ] + } - assert_eq!(pick(&listed, "model").unwrap(), "1"); - assert_eq!(pick(&listed, "model-2").unwrap(), "2"); + #[rstest] + #[case::a_name_others_extend("model")] + #[case::a_name_extending_another("model-2")] + fn an_artifact_is_picked_by_its_exact_name(listed: Vec, #[case] name: &str) { + assert_eq!(newest_named(listed, name).unwrap().name, name); } - #[test] - fn a_name_no_artifact_has_exactly_is_not_found() { - let error = pick(&[("2", "model-2")], "model").unwrap_err(); + #[rstest] + #[case::only_longer_names_listed(vec![artifact("model-2", "2026-10-01T00:00:00Z")])] + #[case::nothing_listed(vec![])] + fn a_name_no_artifact_has_exactly_is_not_found(#[case] listed: Vec) { + let asked = "model"; - assert!(matches!(&error, ArtifactError::NotFound(name) if name == "model")); assert!(matches!( - pick(&[], "model"), - Err(ArtifactError::NotFound(_)) + newest_named(listed, asked), + Err(ArtifactError::NotFound(name)) if name == asked )); } #[test] - fn a_name_several_artifacts_have_is_ambiguous() { - let error = pick(&[("1", "model"), ("2", "model-2"), ("3", "model")], "model").unwrap_err(); - - assert!(matches!(&error, ArtifactError::Ambiguous { count: 2, .. })); - assert_eq!( - error.to_string(), - "Artifact name is ambiguous: 2 artifacts are named model" - ); - } - - fn page(api_url: &str, owner: Namespace, project: &str, num: i32) -> Option { - experiment_page(&Url::parse(api_url).unwrap(), &owner, project, num).map(String::from) - } - - #[test] - fn a_user_project_experiment_is_under_users() { - assert_eq!( - page( - Env::Production.get_url().as_str(), - Namespace::user("alice"), - "mnist", - 42 - ) - .as_deref(), - Some("https://console.tracel.ai/users/alice/projects/mnist/experiments/42") - ); - } - - #[test] - fn an_organization_project_experiment_is_under_orgs() { - assert_eq!( - page( - "https://console.example.com/api/", - Namespace::organization("tracel"), - "vision", - 7 - ) - .as_deref(), - Some("https://console.example.com/orgs/tracel/projects/vision/experiments/7") - ); - } - - #[test] - fn names_are_escaped_as_path_segments() { - assert_eq!( - page( - "https://console.example.com/api/", - Namespace::user("a b"), - "x/y", - 1 - ) - .as_deref(), - Some("https://console.example.com/users/a%20b/projects/x%2Fy/experiments/1") - ); - } - - #[test] - fn an_api_not_served_under_api_has_no_pages() { - assert_eq!( - page( - Env::Development.get_url().as_str(), - Namespace::user("alice"), - "mnist", - 1 - ), - None - ); + fn the_newest_of_same_named_artifacts_is_picked() { + let newest = artifact("model", "2026-10-03T00:00:00Z"); + let listed = vec![ + artifact("model", "2026-10-01T00:00:00Z"), + newest.clone(), + artifact("model-2", "2026-10-04T00:00:00Z"), + artifact("model", "2026-10-02T00:00:00Z"), + ]; + + assert_eq!(newest_named(listed, "model").unwrap().id, newest.id); } } diff --git a/crates/tracel-console/src/inference.rs b/crates/tracel-console/src/inference.rs index 76c7d551..109a739f 100644 --- a/crates/tracel-console/src/inference.rs +++ b/crates/tracel-console/src/inference.rs @@ -299,63 +299,28 @@ fn rfc3339(timestamp_ms: i64) -> String { mod tests { use super::*; - #[test] - fn rfc3339_formats_epoch_millis() { - assert_eq!(rfc3339(1000), "1970-01-01T00:00:01.000Z"); + fn descriptor(name: &str) -> TelemetryMsg { + TelemetryMsg::Descriptor(MetricDescriptor { + name: name.to_string(), + kind: MetricKind::Distribution, + unit: Some("ms".to_string()), + description: None, + }) } #[test] - fn converts_metric_sample_to_backend_contract() { - let sample = MetricSample { - name: "inference_duration_ms".to_string(), - timestamp_ms: 1000, - metadata: serde_json::json!({ "request_id": "wordtok/0" }), - data: MetricData::Distribution { value: 5.0 }, - }; - - let json = serde_json::to_value(convert_metric(sample)).unwrap(); - assert_eq!( - json, - serde_json::json!({ - "name": "inference_duration_ms", - "timestamp": "1970-01-01T00:00:01.000Z", - "metadata": { "request_id": "wordtok/0" }, - "kind": "distribution", - "value": 5.0 - }) - ); - } + fn a_metric_descriptor_is_sent_once_per_name() { + let mut batch = Batch::default(); - #[test] - fn batch_dedups_descriptors_and_counts_events() { - let descriptor = || { - TelemetryMsg::Descriptor(MetricDescriptor { - name: "latency".to_string(), - kind: MetricKind::Distribution, - unit: Some("ms".to_string()), - description: None, - }) - }; + batch.push(descriptor("latency")); + batch.push(descriptor("latency")); + batch.push(descriptor("throughput")); - let mut batch = Batch::default(); - batch.push(descriptor()); - batch.push(descriptor()); - batch.push(TelemetryMsg::Log(LogSample { - timestamp_ms: 0, - level: LogLevel::Info, - message: "hi".to_string(), - metadata: serde_json::json!({}), - })); - batch.push(TelemetryMsg::Metric(MetricSample { - name: "latency".to_string(), - timestamp_ms: 0, - metadata: serde_json::json!({}), - data: MetricData::Gauge { value: 1.0 }, - })); - - assert_eq!(batch.descriptors.len(), 1); - assert_eq!(batch.logs.len(), 1); - assert_eq!(batch.metrics.len(), 1); - assert_eq!(batch.len(), 3); + let sent: Vec<_> = batch + .descriptors + .iter() + .map(|descriptor| descriptor.name.as_str()) + .collect(); + assert_eq!(sent, ["latency", "throughput"]); } } diff --git a/crates/tracel-console/src/lib.rs b/crates/tracel-console/src/lib.rs index 1771ee40..a3846832 100644 --- a/crates/tracel-console/src/lib.rs +++ b/crates/tracel-console/src/lib.rs @@ -4,12 +4,15 @@ //! //! [`Console`] owns one connection to the console. Creating it performs no I/O: the first //! request verifies the credential. Project handles are cheap views over that shared client and -//! vend backend-agnostic, project-scoped capabilities without performing I/O when created. +//! vend backend-agnostic, project-scoped registries and services without performing I/O when +//! created. //! //! [`Console::from_env`] and [`ProjectHandle::from_env`] read the console, the credential and the //! project from `TRACEL_ENV`, `TRACEL_API_KEY`, `TRACEL_NAMESPACE` and `TRACEL_PROJECT`, falling //! back to the `tracel login` sign-in and to `tracel.toml`. [`env_from_environment`], //! [`CredentialSource::from_env`] and [`ProjectRef::from_env`] read each of the three on its own. +//! Their `from_vars` variants, and [`ProjectHandle::from_vars`], read the variables from a lookup, +//! `tracel.toml` from a directory and the sign-in from a session store the caller gives. mod console; mod datasets; @@ -24,8 +27,10 @@ mod wire; pub use console::{Console, ProjectHandle}; pub use domain::{Namespace, NamespaceKind, Organization, Project, User, Visibility}; -pub use env::{CredentialSource, ProjectRef, env_from_environment}; +pub use env::{CredentialSource, ProjectRef, env_from_environment, env_from_vars}; pub use error::ConsoleError; pub use login::{DeviceApproval, DeviceLogin, refresh_session, sign_out}; pub use tracel_client::console::auth::IssuedAppSession; -pub use tracel_client::console::{AccessToken, Env, RefreshToken, TracelCredentials}; +pub use tracel_client::console::{ + AccessToken, Env, FileSessionStore, RefreshToken, SessionStoreError, TracelCredentials, +}; diff --git a/crates/tracel-datasets/Cargo.toml b/crates/tracel-datasets/Cargo.toml index ea2bc515..9dcfd726 100644 --- a/crates/tracel-datasets/Cargo.toml +++ b/crates/tracel-datasets/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tracel-datasets" -description = "Backend-agnostic dataset domain and capability for Tracel." +description = "Backend-agnostic dataset domain and registry for Tracel." edition.workspace = true version.workspace = true readme.workspace = true @@ -13,7 +13,7 @@ categories.workspace = true [features] default = [] -# Burn `Dataset` adapters over the capability. Off by default so that backends composing +# Burn `Dataset` adapters over dataset versions. Off by default so that backends composing # this crate stay burn-free. burn = ["dep:burn"] diff --git a/crates/tracel-experiment/Cargo.toml b/crates/tracel-experiment/Cargo.toml index 15a21bfc..9b9947cb 100644 --- a/crates/tracel-experiment/Cargo.toml +++ b/crates/tracel-experiment/Cargo.toml @@ -18,6 +18,7 @@ local = [] [dependencies] chrono.workspace = true +relative-path = { workspace = true } thiserror.workspace = true tracel-artifact.workspace = true serde_json.workspace = true @@ -27,5 +28,6 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] } burn = { workspace = true, optional = true } [dev-dependencies] +rstest = { workspace = true } tempfile.workspace = true tokio = { version = "1", features = ["macros", "rt-multi-thread"] } diff --git a/crates/tracel-experiment/src/activity.rs b/crates/tracel-experiment/src/activity.rs index b6200121..cb39b6b5 100644 --- a/crates/tracel-experiment/src/activity.rs +++ b/crates/tracel-experiment/src/activity.rs @@ -68,7 +68,7 @@ pub struct ActivitySpec { } /// Terminal state for an activity. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum ActivityStatus { /// The activity completed successfully. Success, @@ -644,229 +644,199 @@ impl std::ops::Deref for ActivityGuard { #[cfg(test)] mod tests { - use std::sync::Arc; + use std::convert::Infallible; + use rstest::rstest; use serde_json::json; - use crate::test_support::{MockSession, create_run, create_run_with_control}; + use crate::ExperimentRun; + use crate::context::ExperimentGlobalExt; + use crate::test_support::{ + MockRun, MockSession, create_run, create_run_with_control, mock_run, + }; use super::*; - fn setup_run() -> (Arc, crate::ExperimentRun) { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - (session, run) - } - - #[test] - fn start_reports_configured_activity() { - let (session, run) = setup_run(); - - let _guard = run - .activity("load") - .meter(12, "items") - .attr("split", "train") - .start(); - - let events = session.activity_events(); - let ActivityEvent::Started { activity } = &events[0] else { - panic!("unexpected event: {:?}", events[0]); - }; - assert_eq!(activity.name, "load"); - let meter = activity.meter.as_ref().expect("expected activity meter"); - assert_eq!(meter.total, Some(12)); - assert_eq!(meter.unit.as_deref(), Some("items")); - assert_eq!(activity.attributes.get("split"), Some(&json!("train"))); - } - - #[test] - fn inc_reports_updated_progress() { - let (session, run) = setup_run(); - let guard = run.activity("items").meter(8, "items").start(); - - guard.inc(3); + const REASON: &str = "fold failed"; + const PANIC: &str = "fusion stream out of order"; - let events = session.activity_events(); - let ActivityEvent::Updated { current, .. } = &events[1] else { - panic!("unexpected event: {:?}", events[1]); - }; - assert_eq!(*current, 3); + fn updates(session: &MockSession) -> Vec { + session + .activity_events() + .into_iter() + .filter_map(|event| match event { + ActivityEvent::Updated { current, .. } => Some(current), + _ => None, + }) + .collect() } - #[test] - fn cloned_activity_reports_accumulated_progress_from_another_thread() { - let (session, run) = setup_run(); - let guard = run.activity("items").meter(8, "items").start(); - let activity = guard.share(); - - activity.inc(2); - let worker_activity = activity.clone(); - std::thread::spawn(move || worker_activity.inc(3)) - .join() - .unwrap(); - - let updates = session + fn endings(session: &MockSession) -> Vec<(ActivityStatus, Option)> { + session .activity_events() .into_iter() .filter_map(|event| match event { - ActivityEvent::Updated { current, .. } => Some(current), + ActivityEvent::Finished { + status, message, .. + } => Some((status, message)), _ => None, }) - .collect::>(); - assert_eq!(updates, vec![2, 5]); + .collect() } - #[test] - fn meterless_activity_reports_updates_as_open_ended_counts() { - let (session, run) = setup_run(); - let guard = run.activity("items").start(); + #[rstest] + fn starting_an_activity_reports_how_it_was_built(mock_run: MockRun) { + let _guard = mock_run + .run + .activity("load") + .meter(12, "items") + .attr("split", "train") + .start(); + + let started = mock_run.session.started_activities().remove(0); + let meter = started.meter.unwrap(); + assert_eq!( + ( + started.name.as_str(), + meter.total, + meter.unit.as_deref(), + started.attributes.get("split") + ), + ("load", Some(12), Some("items"), Some(&json!("train"))) + ); + } + + #[rstest] + #[case::metered(|run: &ExperimentRun| run.activity("items").meter(8, "items").start())] + #[case::open_ended(|run: &ExperimentRun| run.activity("items").start())] + fn increments_from_any_thread_report_the_accumulated_progress( + mock_run: MockRun, + #[case] start: fn(&ExperimentRun) -> ActivityGuard, + ) { + let guard = start(&mock_run.run); let activity = guard.share(); - activity.inc(4); + activity.inc(2); + let worker = activity.clone(); + std::thread::spawn(move || worker.inc(3)).join().unwrap(); - let events = session.activity_events(); - assert!(matches!( - &events[0], - ActivityEvent::Started { activity } if activity.meter.is_none() - )); - assert!(matches!( - &events[1], - ActivityEvent::Updated { current: 4, .. } - )); + assert_eq!(updates(&mock_run.session), [2, 5]); } - #[test] - fn message_reports_activity_message() { - let (session, run) = setup_run(); - let guard = run.activity("items").start(); + #[rstest] + fn a_message_is_reported_for_its_activity(mock_run: MockRun) { + let guard = mock_run.run.activity("items").start(); guard.message("halfway"); - let events = session.activity_events(); assert!(matches!( - &events[1], - ActivityEvent::Message { id, message } - if *id == guard.id() && message == "halfway" + mock_run.session.activity_events().last(), + Some(ActivityEvent::Message { id, message }) if *id == guard.id() && message == "halfway" )); } - #[test] - fn finish_reports_one_success_completion() { - let (session, run) = setup_run(); - - run.activity("node").start().finish(); + #[rstest] + #[case::finished(ActivityGuard::finish, ActivityStatus::Success, None)] + #[case::failed(|guard: ActivityGuard| guard.fail(REASON), ActivityStatus::Failed, Some(REASON))] + #[case::abandoned(ActivityGuard::abandon, ActivityStatus::Abandoned, None)] + #[case::dropped(drop, ActivityStatus::Abandoned, None)] + #[case::dropped_once_cancelled( + |guard: ActivityGuard| { + guard.cancel_token().cancel(); + drop(guard); + }, + ActivityStatus::Abandoned, + None + )] + fn an_activity_reports_once_how_it_ended( + mock_run: MockRun, + #[case] end: fn(ActivityGuard), + #[case] status: ActivityStatus, + #[case] message: Option<&str>, + ) { + end(mock_run.run.activity("node").start()); - let events = session.activity_events(); - let finished: Vec<_> = events - .iter() - .filter_map(|event| match event { - ActivityEvent::Finished { status, .. } => Some(status), - _ => None, - }) - .collect(); - assert!(matches!(finished.as_slice(), [ActivityStatus::Success])); + assert_eq!( + endings(&mock_run.session), + [(status, message.map(String::from))] + ); } #[test] - fn fail_reports_failed_completion_with_its_message() { - let (session, run) = setup_run(); + fn a_guard_unwound_by_a_panic_fails_its_activity_with_the_panic() { + let session = Arc::new(MockSession::default()); + let run_session = session.clone(); - run.activity("node").start().fail("kernel panicked"); + std::thread::spawn(move || { + let run = create_run(run_session); + let _stage = run.activity("stage").start(); + panic!("{PANIC}"); + }) + .join() + .unwrap_err(); - let events = session.activity_events(); assert!(matches!( - events.last(), - Some(ActivityEvent::Finished { - status: ActivityStatus::Failed, - message: Some(message), - .. - }) if message == "kernel panicked" + endings(&session).as_slice(), + [(ActivityStatus::Failed, Some(reason))] if reason.ends_with(PANIC) )); } - #[test] - fn abandon_reports_abandoned_completion() { - let (session, run) = setup_run(); - - run.activity("node").start().abandon(); - - let events = session.activity_events(); - assert!(matches!( - events.last(), - Some(ActivityEvent::Finished { - status: ActivityStatus::Abandoned, - .. + #[rstest] + #[case::returned_ok(Ok(()), ActivityStatus::Success, None)] + #[case::returned_an_error(Err(REASON), ActivityStatus::Failed, Some(REASON))] + fn running_an_activity_ends_it_by_what_its_code_returns( + mock_run: MockRun, + #[case] returned: Result<(), &'static str>, + #[case] status: ActivityStatus, + #[case] message: Option<&str>, + ) { + let _ = mock_run.run.activity("fold").run(|_| returned); + + assert_eq!( + endings(&mock_run.session), + [(status, message.map(String::from))] + ); + } + + #[rstest] + fn running_an_activity_makes_it_the_current_scope(mock_run: MockRun) { + let fold = mock_run + .run + .activity("fold") + .run(|activity| { + ExperimentRun::current().unwrap().log_info("in the fold"); + Ok::<_, Infallible>(activity.id()) }) - )); - } - - #[test] - fn drop_reports_abandoned_completion() { - let (session, run) = setup_run(); - - drop(run.activity("node").start()); + .unwrap(); - let events = session.activity_events(); - assert!(matches!( - events.last(), - Some(ActivityEvent::Finished { - status: ActivityStatus::Abandoned, - .. - }) - )); + assert_eq!(mock_run.session.log_activities(), [Some(fold)]); } - #[test] - fn cancellable_activity_reports_control_metadata() { - let (session, run) = setup_run(); - - let _guard = run.activity("node").cancellable().start(); + #[rstest] + fn a_cancellable_activity_says_so_when_it_starts(mock_run: MockRun) { + let _guard = mock_run.run.activity("node").cancellable().start(); - let events = session.activity_events(); - let ActivityEvent::Started { activity } = &events[0] else { - panic!("unexpected event: {:?}", events[0]); - }; - assert!(activity.cancellable); + assert!(mock_run.session.started_activities()[0].cancellable); } #[test] - fn cancellable_activity_registers_with_control() { - let session = Arc::new(MockSession::default()); + fn a_cancellable_activity_is_cancelled_through_its_run_control() { let control = ExperimentRunControl::default(); - let run = create_run_with_control(session, control.clone()); + let run = create_run_with_control(Default::default(), control.clone()); let guard = run.activity("node").cancellable().start(); - assert!(control.cancel_activity(guard.id())); + control.cancel_activity(guard.id()); + assert!(guard.is_cancel_requested()); } - #[test] - fn child_activity_cancel_request_is_linked_to_parent_activity_token() { - let (_session, run) = setup_run(); - let parent = run.activity("parent").start(); + #[rstest] + fn cancelling_an_activity_requests_the_cancel_of_its_children(mock_run: MockRun) { + let parent = mock_run.run.activity("parent").start(); let child = parent.activity("child").start(); parent.cancel_token().cancel(); - assert!(parent.is_cancel_requested()); assert!(child.is_cancel_requested()); } - - #[test] - fn drop_reports_abandoned_even_when_cancel_was_requested() { - let (session, run) = setup_run(); - let guard = run.activity("node").start(); - - guard.cancel_token().cancel(); - drop(guard); - - let events = session.activity_events(); - assert!(matches!( - events.last(), - Some(ActivityEvent::Finished { - status: ActivityStatus::Abandoned, - .. - }) - )); - } } diff --git a/crates/tracel-experiment/src/cancellation.rs b/crates/tracel-experiment/src/cancellation.rs index 38a93610..dc934dd7 100644 --- a/crates/tracel-experiment/src/cancellation.rs +++ b/crates/tracel-experiment/src/cancellation.rs @@ -136,105 +136,82 @@ impl Cancellable for Arc { #[cfg(test)] mod tests { - use super::*; - use std::sync::Arc; + use std::sync::Barrier; use std::sync::atomic::AtomicUsize; use std::thread; - use std::time::Duration; - struct TestCancellable { - cancelled: AtomicBool, - cancel_count: AtomicUsize, - } + use super::*; - impl TestCancellable { - fn new() -> Self { - Self { - cancelled: AtomicBool::new(false), - cancel_count: AtomicUsize::new(0), - } - } + #[derive(Default)] + struct CountsCancels(AtomicUsize); - fn cancel_count(&self) -> usize { - self.cancel_count.load(Ordering::Relaxed) + impl CountsCancels { + fn count(&self) -> usize { + self.0.load(Ordering::Acquire) } } - impl Cancellable for TestCancellable { + impl Cancellable for CountsCancels { fn cancel(&self) { - if !self.cancelled.swap(true, Ordering::AcqRel) { - self.cancel_count.fetch_add(1, Ordering::AcqRel); - } + self.0.fetch_add(1, Ordering::AcqRel); } fn is_cancelled(&self) -> bool { - self.cancelled.load(Ordering::Acquire) + self.count() > 0 } } - #[test] - fn test_cancel_token() { - let token = CancelToken::new(); - - token.cancel(); - - assert!(token.is_cancelled()); + fn linked_child(token: &CancelToken) -> Arc { + let child = Arc::new(CountsCancels::default()); + token.link(child.clone()); + child } #[test] - fn test_cancel_children() { + fn cancelling_a_token_cancels_every_child_linked_to_it() { let token = CancelToken::new(); - let child1 = Arc::new(TestCancellable::new()); - let child2 = Arc::new(TestCancellable::new()); - token.link(child1.clone()); - token.link(child2.clone()); + let children = [linked_child(&token), linked_child(&token)]; token.cancel(); - assert!(child1.is_cancelled()); - assert!(child2.is_cancelled()); + assert!(children.iter().all(|child| child.is_cancelled())); } #[test] - fn test_idempotent_cancel() { + fn cancelling_a_token_twice_cancels_its_child_once() { let token = CancelToken::new(); - let child = Arc::new(TestCancellable::new()); - token.link(child.clone()); - token.cancel(); + let child = linked_child(&token); + token.cancel(); token.cancel(); - assert!(token.is_cancelled()); - assert!(child.is_cancelled()); - assert_eq!(child.cancel_count(), 1); + assert_eq!(child.count(), 1); } #[test] - fn test_concurrent_cancel() { + fn cancelling_a_token_from_two_threads_at_once_cancels_its_child_once() { let token = CancelToken::new(); - let child = Arc::new(TestCancellable::new()); - token.link(child.clone()); - let token_clone = token.clone(); - let handle = thread::spawn(move || { - thread::sleep(Duration::from_millis(50)); - token_clone.cancel(); + let child = linked_child(&token); + let barrier = Barrier::new(2); + + thread::scope(|scope| { + for _ in 0..2 { + scope.spawn(|| { + barrier.wait(); + token.cancel(); + }); + } }); - token.cancel(); - - handle.join().unwrap(); - assert!(token.is_cancelled()); - assert!(child.is_cancelled()); - assert_eq!(child.cancel_count(), 1); // Should only be cancelled once + assert_eq!(child.count(), 1); } #[test] - fn test_link_after_cancel() { + fn a_child_linked_to_a_cancelled_token_is_cancelled() { let token = CancelToken::new(); token.cancel(); - let child = Arc::new(TestCancellable::new()); - token.link(child.clone()); + let child = linked_child(&token); assert!(child.is_cancelled()); } diff --git a/crates/tracel-experiment/src/context.rs b/crates/tracel-experiment/src/context.rs index 33ba3d7d..cc7a7d14 100644 --- a/crates/tracel-experiment/src/context.rs +++ b/crates/tracel-experiment/src/context.rs @@ -214,9 +214,9 @@ impl Drop for CurrentExperimentGuard { #[cfg(test)] mod tests { - use crate::ExperimentRun; use crate::context::ExperimentGlobalExt as _; use crate::test_support::create_run_with_id; + use crate::{ExperimentId, ExperimentRun}; use super::ExperimentInstrument; @@ -224,144 +224,91 @@ mod tests { create_run_with_id(id, Default::default()) } - #[test] - fn nested_experiment_scopes_restore_the_previous_experiment() { - let run_a = create_run("ambient-test-a"); - let run_b = create_run("ambient-test-b"); - - assert!(ExperimentRun::current().is_none()); - { - let _outer = run_a.enter(); - let current = ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id(), run_a.id()); - - { - let _inner = run_b.enter(); - let current = - ExperimentRun::current().expect("inner experiment should override outer scope"); - assert_eq!(current.id(), run_b.id()); - } - - let current = ExperimentRun::current().expect("outer experiment should be restored"); - assert_eq!(current.id(), run_a.id()); - } - assert!(ExperimentRun::current().is_none()); + fn current_id() -> Option { + ExperimentRun::current().map(|current| current.id().clone()) } - #[test] - fn dropping_an_outer_scope_does_not_remove_the_inner_scope() { - let run_a = create_run("ambient-test-out-of-order-a"); - let run_b = create_run("ambient-test-out-of-order-b"); - let outer = run_a.enter(); - let inner = run_b.enter(); - - drop(outer); - - let current = ExperimentRun::current().expect("inner experiment should remain current"); - assert_eq!(current.id(), run_b.id()); - drop(inner); - assert!(ExperimentRun::current().is_none()); + async fn current_id_after_an_await() -> Option { + tokio::task::yield_now().await; + current_id() } #[test] - fn handle_with_current_can_be_used_in_spawned_threads() { - let run = create_run("ambient-test-thread"); - let handle = run.handle(); + fn an_inner_scope_overrides_the_outer_one() { + let outer = create_run("outer"); + let inner = create_run("inner"); + let _outer = outer.enter(); - std::thread::spawn(move || { - handle.in_scope(|| { - let current = ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id().as_str(), "ambient-test-thread"); - }) - }) - .join() - .expect("thread should complete successfully"); + let _inner = inner.enter(); - assert!(ExperimentRun::current().is_none()); + assert_eq!(current_id(), Some(inner.id().clone())); } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn instrumented_tokio_tasks_keep_multiple_experiments_isolated() { - let run_a = create_run("ambient-test-tokio-a"); - let run_b = create_run("ambient-test-tokio-b"); - - let task_a = tokio::spawn( - async { - let current = ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-a"); + #[test] + fn leaving_an_inner_scope_restores_the_outer_one() { + let outer = create_run("outer"); + let inner = create_run("inner"); + let _outer = outer.enter(); - tokio::task::yield_now().await; + drop(inner.enter()); - let current = - ExperimentRun::current().expect("current experiment should still be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-a"); - } - .in_experiment(&run_a), - ); + assert_eq!(current_id(), Some(outer.id().clone())); + } - let task_b = tokio::spawn( - async { - let current = ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-b"); + #[test] + fn leaving_an_outer_scope_first_keeps_the_inner_one() { + let outer = create_run("outer"); + let inner = create_run("inner"); + let outer_scope = outer.enter(); + let _inner_scope = inner.enter(); - tokio::task::yield_now().await; + drop(outer_scope); - let current = - ExperimentRun::current().expect("current experiment should still be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-b"); - } - .in_experiment(&run_b), - ); - - task_a - .await - .expect("first tokio task should complete successfully"); - task_b - .await - .expect("second tokio task should complete successfully"); - assert!(ExperimentRun::current().is_none()); + assert_eq!(current_id(), Some(inner.id().clone())); } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn nested_tokio_spawn_can_use_in_current_experiment() { - let run = create_run("ambient-test-tokio-nested"); - - let outer = tokio::spawn( - async { - let current = ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-nested"); + #[test] + fn a_handle_enters_its_run_on_another_thread() { + let run = create_run("run"); + let handle = run.handle(); - let inner = tokio::spawn( - async { - let current = - ExperimentRun::current().expect("current experiment should be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-nested"); + let current = std::thread::spawn(move || handle.in_scope(current_id)) + .join() + .unwrap(); - tokio::task::yield_now().await; + assert_eq!(current, Some(run.id().clone())); + } - let current = ExperimentRun::current() - .expect("current experiment should still be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-nested"); - } - .in_current_experiment(), - ); + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn instrumented_tasks_each_keep_their_own_experiment_across_an_await() { + let first = create_run("first"); + let second = create_run("second"); + + let tasks = [&first, &second] + .map(|run| tokio::spawn(current_id_after_an_await().in_experiment(run))); + let mut seen = Vec::new(); + for task in tasks { + seen.push(task.await.unwrap()); + } - tokio::task::yield_now().await; + assert_eq!(seen, [Some(first.id().clone()), Some(second.id().clone())]); + } - let current = - ExperimentRun::current().expect("outer current experiment should still be set"); - assert_eq!(current.id().as_str(), "ambient-test-tokio-nested"); + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn a_task_spawned_in_the_current_experiment_keeps_it() { + let run = create_run("run"); - inner + let seen = tokio::spawn( + async { + tokio::spawn(current_id_after_an_await().in_current_experiment()) .await - .expect("nested tokio task should complete successfully"); + .unwrap() } .in_experiment(&run), - ); + ) + .await + .unwrap(); - outer - .await - .expect("outer tokio task should complete successfully"); - assert!(ExperimentRun::current().is_none()); + assert_eq!(seen, Some(run.id().clone())); } } diff --git a/crates/tracel-experiment/src/control.rs b/crates/tracel-experiment/src/control.rs index f3cbd643..cf4195c7 100644 --- a/crates/tracel-experiment/src/control.rs +++ b/crates/tracel-experiment/src/control.rs @@ -86,39 +86,43 @@ impl Default for ExperimentRunControl { mod tests { use std::num::NonZeroU64; + use rstest::rstest; + use super::*; - fn activity_id(value: u64) -> ActivityId { - ActivityId::new(NonZeroU64::new(value).unwrap()) + fn activity_id() -> ActivityId { + ActivityId::new(NonZeroU64::new(1).unwrap()) } #[test] - fn cancel_run_cancels_linked_activity_tokens() { + fn cancelling_the_run_cancels_its_token() { let control = ExperimentRunControl::default(); - let activity_token = control.cancel_token().linked(CancelToken::new()); - control.register_activity_cancellation(activity_id(1), activity_token.clone()); control.cancel_run(); assert!(control.cancel_token().is_cancelled()); - assert!(activity_token.is_cancelled()); } #[test] - fn cancel_activity_cancels_registered_activity_token() { + fn cancelling_a_registered_activity_cancels_its_token() { let control = ExperimentRunControl::default(); - let activity_token = CancelToken::new(); - control.register_activity_cancellation(activity_id(1), activity_token.clone()); + let token = CancelToken::new(); + control.register_activity_cancellation(activity_id(), token.clone()); - assert!(control.cancel_activity(activity_id(1))); + control.cancel_activity(activity_id()); - assert!(activity_token.is_cancelled()); + assert!(token.is_cancelled()); } - #[test] - fn cancel_unknown_activity_returns_false() { + #[rstest] + #[case::registered(true)] + #[case::unknown(false)] + fn cancelling_an_activity_reports_whether_it_was_registered(#[case] registered: bool) { let control = ExperimentRunControl::default(); + if registered { + control.register_activity_cancellation(activity_id(), CancelToken::new()); + } - assert!(!control.cancel_activity(activity_id(99))); + assert_eq!(control.cancel_activity(activity_id()), registered); } } diff --git a/crates/tracel-experiment/src/integration/tracing/layer.rs b/crates/tracel-experiment/src/integration/tracing/layer.rs index a497ab6e..0c609619 100644 --- a/crates/tracel-experiment/src/integration/tracing/layer.rs +++ b/crates/tracel-experiment/src/integration/tracing/layer.rs @@ -224,3 +224,85 @@ impl Visit for JsonFieldVisitor { self.record(field, Value::String(format!("{value:?}"))); } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use serde_json::json; + + use crate::ExperimentGlobalExt; + use crate::test_support::{MockRun, MockSession, create_run_with_id, mock_run, traced}; + + #[rstest] + fn an_event_in_the_current_experiment_is_logged_to_it(mock_run: MockRun) { + traced(|| mock_run.run.in_scope(|| tracing::info!("epoch completed"))); + + assert_eq!(mock_run.session.logs().len(), 1); + } + + #[test] + fn an_event_in_a_span_naming_an_experiment_is_logged_to_it() { + let session = Arc::new(MockSession::default()); + let run = create_run_with_id("layer-span", session.clone()); + + traced(|| { + tracing::info_span!("experiment", experiment_id = %run.id()) + .in_scope(|| tracing::info!("epoch completed")) + }); + + assert_eq!(session.logs().len(), 1); + } + + #[rstest] + fn an_event_outside_any_experiment_is_not_logged(mock_run: MockRun) { + traced(|| tracing::info!("outside any experiment")); + + assert!(mock_run.session.logs().is_empty()); + } + + #[rstest] + fn an_event_field_is_an_attribute_of_its_log(mock_run: MockRun) { + traced(|| { + mock_run + .run + .in_scope(|| tracing::info!(step = 3u64, "epoch completed")) + }); + + assert_eq!( + mock_run.session.logs()[0].attributes.get("step"), + Some(&json!(3)) + ); + } + + #[rstest] + fn a_span_field_is_an_attribute_of_the_logs_inside_it(mock_run: MockRun) { + traced(|| { + mock_run.run.in_scope(|| { + tracing::info_span!("phase", stage = "train") + .in_scope(|| tracing::info!("step done")) + }) + }); + + assert_eq!( + mock_run.session.logs()[0].attributes.get("stage"), + Some(&json!("train")) + ); + } + + #[rstest] + fn an_event_field_wins_over_a_span_field_of_the_same_name(mock_run: MockRun) { + traced(|| { + mock_run.run.in_scope(|| { + tracing::info_span!("phase", shard = 2u64) + .in_scope(|| tracing::info!(shard = 5u64, "step done")) + }) + }); + + assert_eq!( + mock_run.session.logs()[0].attributes.get("shard"), + Some(&json!(5)) + ); + } +} diff --git a/crates/tracel-experiment/src/integration/tracing/mod.rs b/crates/tracel-experiment/src/integration/tracing/mod.rs index d9369bff..05ab674f 100644 --- a/crates/tracel-experiment/src/integration/tracing/mod.rs +++ b/crates/tracel-experiment/src/integration/tracing/mod.rs @@ -125,184 +125,34 @@ impl ActivityTracingExt for Activity { mod tests { use std::sync::Arc; - use tracing_subscriber::layer::SubscriberExt; - - use crate::context::ExperimentGlobalExt; - use crate::session::Event; - use crate::test_support::{MockSession, create_run_with_id}; - use super::*; + use crate::test_support::{MockSession, create_run_with_id, traced}; #[test] - fn tracing_layer_forwards_events_to_current_experiment() { - let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-1", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - - tracing::subscriber::with_default(subscriber, || { - run.in_scope(|| { - tracing::info!(step = 3u64, "epoch completed"); - }) - }); - - let events = session.events.lock().unwrap(); - assert_eq!(events.len(), 1); - match &events[0] { - Event::Log { record, .. } => { - assert!(record.message.contains("epoch completed")); - assert_eq!( - record.attributes.get("step").and_then(|v| v.as_u64()), - Some(3) - ); - } - event => panic!("unexpected event: {event:?}"), - } - } - - #[test] - fn activity_run_installs_its_scope_for_tracing_events() { + fn an_event_in_an_experiment_tracing_span_is_logged_to_that_experiment() { let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-activity", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); + let run = create_run_with_id("tracing-experiment-span", session.clone()); - let activity_id = tracing::subscriber::with_default(subscriber, || { - run.activity("fold") - .run(|activity| { - tracing::info!(fold = 1u64, "training fold"); - Ok::<_, std::convert::Infallible>(activity.id()) - }) - .unwrap() + traced(|| { + run.tracing_span() + .in_scope(|| tracing::info!("epoch completed")) }); - let events = session.events.lock().unwrap(); - let activity = events - .iter() - .find_map(|event| match event { - Event::Log { activity, .. } => *activity, - _ => None, - }) - .expect("tracing log should carry an activity scope"); - assert_eq!(activity, activity_id); + assert_eq!(session.logs().len(), 1); } #[test] - fn activity_span_routes_logs_to_its_activity() { + fn an_event_in_an_activity_tracing_span_is_logged_to_that_activity() { let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-activity-span", session.clone()); + let run = create_run_with_id("tracing-activity-span", session.clone()); let activity = run.activity("fold").start(); - let activity_id = activity.id(); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - tracing::subscriber::with_default(subscriber, || { + traced(|| { activity .tracing_span() - .in_scope(|| tracing::info!("training fold")); - }); - - let events = session.events.lock().unwrap(); - assert!(events.iter().any(|event| matches!( - event, - Event::Log { - activity: Some(id), - .. - } if *id == activity_id - ))); - } - - #[test] - fn tracing_layer_routes_from_span_experiment_id_without_ambient_scope() { - let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-span", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - - tracing::subscriber::with_default(subscriber, || { - let span = tracing::info_span!("experiment", experiment_id = %run.id()); - let _guard = span.enter(); - tracing::info!(step = 7u64, "span-routed event"); - }); - - let events = session.events.lock().unwrap(); - assert_eq!(events.len(), 1); - match &events[0] { - Event::Log { record, .. } => { - assert!(record.message.contains("span-routed event")); - assert_eq!( - record.attributes.get("step").and_then(|v| v.as_u64()), - Some(7) - ); - } - event => panic!("unexpected event: {event:?}"), - } - } - - #[test] - fn tracing_layer_routes_from_experiment_span_helper_without_ambient_scope() { - let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-helper-span", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - - tracing::subscriber::with_default(subscriber, || { - let span = experiment_span(&run); - let _guard = span.enter(); - tracing::info!("helper-span-routed event"); - }); - - let events = session.events.lock().unwrap(); - assert_eq!(events.len(), 1); - match &events[0] { - Event::Log { record, .. } => { - assert!(record.message.contains("helper-span-routed event")); - } - event => panic!("unexpected event: {event:?}"), - } - } - - #[test] - fn tracing_layer_skips_events_without_experiment_scope() { - let session = Arc::new(MockSession::default()); - let _run = create_run_with_id("trace-test-2", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - - tracing::subscriber::with_default(subscriber, || { - tracing::info!("outside experiment scope"); - }); - - let events = session.events.lock().unwrap(); - assert!(events.is_empty()); - } - - #[test] - fn tracing_layer_inherits_span_field_attributes() { - let session = Arc::new(MockSession::default()); - let run = create_run_with_id("trace-test-scope", session.clone()); - let subscriber = tracing_subscriber::registry().with(tracing_log_layer()); - - tracing::subscriber::with_default(subscriber, || { - run.in_scope(|| { - let span = tracing::info_span!("phase", stage = "train", shard = 2u64); - span.in_scope(|| { - // The event's own field overrides the inherited `shard` from the span. - tracing::info!(shard = 5u64, "step done"); - }); - }) + .in_scope(|| tracing::info!("training fold")) }); - let events = session.events.lock().unwrap(); - let log = events - .iter() - .find_map(|event| match event { - Event::Log { record, .. } => Some(record), - _ => None, - }) - .expect("a log event should have been recorded"); - assert_eq!(log.message, "step done"); - assert_eq!( - log.attributes.get("stage").and_then(|v| v.as_str()), - Some("train") - ); - assert_eq!( - log.attributes.get("shard").and_then(|v| v.as_u64()), - Some(5) - ); + assert_eq!(session.log_activities(), [Some(activity.id())]); } } diff --git a/crates/tracel-experiment/src/integration/training/progress.rs b/crates/tracel-experiment/src/integration/training/progress.rs index 2a162cc6..9e05bcd1 100644 --- a/crates/tracel-experiment/src/integration/training/progress.rs +++ b/crates/tracel-experiment/src/integration/training/progress.rs @@ -243,166 +243,178 @@ impl EvaluationProgressLogger for ExperimentEvaluationProgressLogger { #[cfg(test)] mod tests { - use std::sync::Arc; - - use crate::{ - activity::ActivityEvent, - integration::training::ExperimentTrainingExt, - session::Event, - test_support::{MockSession, create_run}, - }; + use rstest::rstest; + use serde_json::json; use super::*; + use crate::integration::training::ExperimentTrainingExt; + use crate::test_support::{MockRun, MockSession, mock_run}; + use crate::{ActivityEvent, ActivitySpec, ActivityStatus}; + + const TRAINING: &str = "Pretraining"; + const EVALUATION: &str = "Holdout"; + const EPOCHS: usize = 2; + const TESTS: usize = 2; + const ITEMS: usize = 10; + const RESUMED_AT: usize = 3; + + fn started(session: &MockSession, name: &str) -> ActivitySpec { + session + .started_activities() + .into_iter() + .find(|activity| activity.name == name) + .unwrap() + } - #[test] - fn training_progress_groups_splits_under_supplied_scope() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let parent = run.activity("Fold 1").start(); - let parent_id = parent.id(); - let mut logger = parent.training_progress_logger(); - - logger.start(2, 1, None, None); - logger.start_split("train", 10); - logger.update_split(4); - logger.end_split(); - logger.start_split("valid", 5); - logger.end_split(); - logger.update_epoch(1); - logger.end(); + fn epoch_of(session: &MockSession, split: &str) -> ActivitySpec { + let split = started(session, split); + session + .started_activities() + .into_iter() + .find(|activity| Some(activity.id) == split.parent) + .unwrap() + } - let events = session.events.lock().unwrap(); - let started = events - .iter() - .filter_map(|event| match event { - Event::Activity(ActivityEvent::Started { activity }) => Some(activity), + fn ending(session: &MockSession, name: &str) -> Option { + let id = started(session, name).id; + session + .activity_events() + .into_iter() + .find_map(|event| match event { + ActivityEvent::Finished { + id: finished, + status, + .. + } if finished == id => Some(status), _ => None, }) - .collect::>(); + } + + fn total(session: &MockSession, name: &str) -> Option { + started(session, name).meter.and_then(|meter| meter.total) + } + + #[rstest] + fn the_training_activity_takes_the_name_of_its_logger(mock_run: MockRun) { + let mut logger = mock_run.run.training_progress_logger().with_name(TRAINING); - assert_eq!(started.len(), 5); - let training = started[1]; - let epoch = started[2]; - let train = started[3]; - let valid = started[4]; + logger.start(EPOCHS, 1, None, None); - assert_eq!(training.name, "Training"); - assert_eq!(training.parent, Some(parent_id)); + assert_eq!(mock_run.session.started_activities()[0].name, TRAINING); + } + + #[rstest] + fn a_split_nests_under_its_epoch_under_the_training_under_its_scope(mock_run: MockRun) { + let fold = mock_run.run.activity("Fold 1").start(); + let mut logger = fold.training_progress_logger().with_name(TRAINING); + + logger.start(EPOCHS, 1, None, None); + logger.start_split("train", ITEMS); + + let training = started(&mock_run.session, TRAINING); + let epoch = epoch_of(&mock_run.session, "train"); assert_eq!( - training.meter.as_ref().unwrap().unit.as_deref(), - Some("epochs") + (training.parent, epoch.parent), + (Some(fold.id()), Some(training.id)) ); - assert_eq!(training.meter.as_ref().unwrap().total, Some(2)); + } + + #[rstest] + fn the_training_is_metered_by_its_epochs_and_a_split_by_its_items(mock_run: MockRun) { + let mut logger = mock_run.run.training_progress_logger().with_name(TRAINING); + + logger.start(EPOCHS, 1, None, None); + logger.start_split("train", ITEMS); - assert_eq!(epoch.name, "Epoch 1"); - assert_eq!(epoch.parent, Some(training.id)); - assert!(epoch.meter.is_none()); assert_eq!( - epoch.attributes.get("activity_type"), - Some(&serde_json::json!("epoch")) + ( + total(&mock_run.session, TRAINING), + total(&mock_run.session, "train") + ), + (Some(EPOCHS as u64), Some(ITEMS as u64)) ); - assert_eq!(epoch.attributes.get("epoch"), Some(&serde_json::json!(1))); + } + + #[rstest] + fn an_epoch_activity_numbers_its_epoch_from_where_the_training_started(mock_run: MockRun) { + let mut logger = mock_run.run.training_progress_logger().with_name(TRAINING); - assert_eq!(train.name, "train"); - assert_eq!(train.parent, Some(epoch.id)); - assert_eq!(train.meter.as_ref().unwrap().unit.as_deref(), Some("steps")); - assert_eq!(train.meter.as_ref().unwrap().total, Some(10)); + logger.start(EPOCHS, RESUMED_AT, None, None); + logger.start_split("train", ITEMS); - assert_eq!(valid.name, "valid"); - assert_eq!(valid.parent, Some(epoch.id)); - assert_eq!(valid.meter.as_ref().unwrap().unit.as_deref(), Some("steps")); - assert_eq!(valid.meter.as_ref().unwrap().total, Some(5)); + assert_eq!( + epoch_of(&mock_run.session, "train").attributes.get("epoch"), + Some(&json!(RESUMED_AT)) + ); } - #[test] - fn training_progress_uses_supplied_name() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let mut logger = run.training_progress_logger().with_name("Pretraining"); + #[rstest] + #[case::through_its_items(ITEMS, ActivityStatus::Success)] + #[case::short_of_its_items(ITEMS - 1, ActivityStatus::Abandoned)] + fn a_split_ends_by_how_far_its_meter_got( + mock_run: MockRun, + #[case] processed: usize, + #[case] status: ActivityStatus, + ) { + let mut logger = mock_run.run.training_progress_logger().with_name(TRAINING); + logger.start(EPOCHS, 1, None, None); + logger.start_split("train", ITEMS); - logger.start(1, 1, None, None); - logger.end(); + logger.update_split(processed); + logger.end_split(); - let names = session - .activity_events() - .into_iter() - .filter_map(|event| match event { - ActivityEvent::Started { activity } => Some(activity.name), - _ => None, - }) - .collect::>(); - assert_eq!(names, ["Pretraining"]); + assert_eq!(ending(&mock_run.session, "train"), Some(status)); } - #[test] - fn endings_follow_the_meter() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let parent = run.activity("Fold 6").start(); - let mut logger = parent.training_progress_logger(); + #[rstest] + #[case::every_epoch(EPOCHS, ActivityStatus::Success)] + #[case::short_of_its_epochs(EPOCHS - 1, ActivityStatus::Abandoned)] + fn the_training_ends_by_how_many_epochs_it_completed( + mock_run: MockRun, + #[case] completed: usize, + #[case] status: ActivityStatus, + ) { + let mut logger = mock_run.run.training_progress_logger().with_name(TRAINING); + logger.start(EPOCHS, 1, None, None); - logger.start(2, 1, None, None); - logger.start_split("train", 10); - logger.update_split(10); - logger.end_split(); - logger.start_split("valid", 10); - logger.update_split(3); - logger.end_split(); + logger.update_epoch(completed); logger.end(); - let completions = session - .activity_events() - .into_iter() - .filter_map(|event| match event { - ActivityEvent::Finished { id, status, .. } => Some((id, status)), - _ => None, - }) - .collect::>(); - - let statuses: Vec<_> = completions.iter().map(|(_, s)| s.clone()).collect(); - assert!( - matches!( - statuses.as_slice(), - [ - crate::ActivityStatus::Success, // the train split ran dry - crate::ActivityStatus::Abandoned, // the valid split stopped at 3 of 10 - crate::ActivityStatus::Abandoned, // the epoch never completed - crate::ActivityStatus::Abandoned, // 0 of 2 epochs is not a finished training - ] - ), - "unexpected endings: {statuses:?}" + assert_eq!(ending(&mock_run.session, TRAINING), Some(status)); + } + + #[rstest] + fn a_test_nests_under_the_evaluation_under_its_scope(mock_run: MockRun) { + let fold = mock_run.run.activity("Fold 1").start(); + let mut logger = fold.evaluation_progress_logger().with_name(EVALUATION); + + logger.start_global_progress(TESTS); + logger.start_test("test", ITEMS); + + let evaluation = started(&mock_run.session, EVALUATION); + let test = started(&mock_run.session, "test"); + assert_eq!( + (evaluation.parent, test.parent), + (Some(fold.id()), Some(evaluation.id)) ); } - #[test] - fn evaluation_progress_groups_tests_under_supplied_scope() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let parent = run.activity("Fold 1").start(); - let mut logger = parent.evaluation_progress_logger().with_name("Holdout"); - - logger.start_global_progress(2); - logger.start_test("test", 4); - logger.update_test_progress(2); - logger.end_test(); - logger.end_global_progress(); - - let events = session.events.lock().unwrap(); - let started = events - .iter() - .filter_map(|event| match event { - Event::Activity(ActivityEvent::Started { activity }) => Some(activity), - _ => None, - }) - .collect::>(); - - assert_eq!(started.len(), 3); - let evaluation = started[1]; - let test = started[2]; - assert_eq!(evaluation.name, "Holdout"); - assert_eq!(evaluation.parent, Some(parent.id())); - assert_eq!(evaluation.meter.as_ref().unwrap().total, Some(2)); - assert_eq!(test.parent, Some(evaluation.id)); - assert_eq!(test.meter.as_ref().unwrap().total, Some(4)); + #[rstest] + fn the_evaluation_is_metered_by_its_tests_and_a_test_by_its_items(mock_run: MockRun) { + let mut logger = mock_run + .run + .evaluation_progress_logger() + .with_name(EVALUATION); + + logger.start_global_progress(TESTS); + logger.start_test("test", ITEMS); + + assert_eq!( + ( + total(&mock_run.session, EVALUATION), + total(&mock_run.session, "test") + ), + (Some(TESTS as u64), Some(ITEMS as u64)) + ); } } diff --git a/crates/tracel-experiment/src/lib.rs b/crates/tracel-experiment/src/lib.rs index 608a55d7..1cb00013 100644 --- a/crates/tracel-experiment/src/lib.rs +++ b/crates/tracel-experiment/src/lib.rs @@ -28,7 +28,6 @@ //! feature, `local::LocalExperiments` records experiments on this machine. use std::fmt; -use std::path::PathBuf; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, Weak}; @@ -133,15 +132,6 @@ impl From for ExperimentId { } } -/// Where an experiment run can be looked at. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ExperimentLocation { - /// The experiment's page on the console. - Url(String), - /// The directory the run is recorded in on this machine. - Dir(PathBuf), -} - /// Artifact category associated with an experiment run. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ArtifactKind { @@ -198,7 +188,6 @@ struct ExperimentMetadata { pub struct ExperimentRun { inner: Arc, handle: ExperimentRunHandle, - location: Option, _tracing_registration: TracingRegistration, } @@ -316,24 +305,10 @@ impl ExperimentRun { Self { inner, handle, - location: None, _tracing_registration: tracing_registration, } } - /// Records where the run can be looked at, such as its page on the console. - /// - /// A provider sets it when it creates the run. - pub fn with_location(mut self, location: ExperimentLocation) -> Self { - self.location = Some(location); - self - } - - /// Where the run can be looked at, when its provider says. - pub fn location(&self) -> Option<&ExperimentLocation> { - self.location.as_ref() - } - /// Stream every panic on any thread into this run's log until the guard /// drops. The failure reason needs no guard: an unwinding run reads the /// panicking thread's record on its own. @@ -878,405 +853,211 @@ impl Drop for ExperimentRun { #[cfg(test)] mod tests { - use crate::activity::ActivityEvent; - use crate::test_support::{MockSession, create_run, create_run_with_control}; + use rstest::rstest; + use serde_json::json; - use super::*; + use crate::test_support::{ + MockRun, MockSession, create_run, create_run_with_control, mock_run, + }; - #[test] - fn panic_unwinding_a_guard_fails_the_activity_with_the_panic() { - let session = Arc::new(MockSession::default()); - let run_session = session.clone(); - - let outcome = std::thread::Builder::new() - .name("stage-runner".into()) - .spawn(move || { - let run = create_run(run_session); - let _stage = run.activity("Train fold 3 of 6").cancellable().start(); - panic!("fusion stream out of order"); - }) - .unwrap() - .join(); - assert!(outcome.is_err()); - - let finished = session - .activity_events() - .into_iter() - .find_map(|event| match event { - ActivityEvent::Finished { - status, message, .. - } => Some((status, message)), - _ => None, - }) - .expect("the stage reported an ending"); - assert!(matches!(finished.0, ActivityStatus::Failed)); - let message = finished.1.expect("the ending carried the panic"); - assert!(message.contains("fusion stream out of order"), "{message}"); - - let completions = session.completions.lock().unwrap(); - match completions.as_slice() { - [ExperimentCompletion::Failed(reason)] => { - assert!(reason.contains("fusion stream out of order"), "{reason}"); - } - other => panic!("unexpected completions: {other:?}"), - } - } - - #[test] - fn panicking_run_reports_the_panic_as_its_failure_reason() { - let session = Arc::new(MockSession::default()); - let run_session = session.clone(); - - let outcome = std::thread::Builder::new() - .name("fold-runner".into()) - .spawn(move || { - let _run = create_run(run_session); - panic!("kernel exploded mid-batch"); - }) - .unwrap() - .join(); - assert!(outcome.is_err()); - - let completions = session.completions.lock().unwrap(); - match completions.as_slice() { - [ExperimentCompletion::Failed(reason)] => { - assert!(reason.contains("kernel exploded mid-batch"), "{reason}"); - assert!(reason.contains("panicked at"), "{reason}"); - assert!(reason.contains("fold-runner"), "{reason}"); - } - other => panic!("unexpected completions: {other:?}"), - } - } - - #[test] - fn watched_run_logs_panics_from_other_threads() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let _watch = run.capture_panics(); - - let outcome = std::thread::Builder::new() - .name("device-runner".into()) - .spawn(|| panic!("unsupported op: plane.elect")) - .unwrap() - .join(); - assert!(outcome.is_err()); - - let events = session.events.lock().unwrap(); - let logged = events.iter().any(|event| match event { - Event::Log { record, .. } => { - record.level == LogLevel::Error - && record.message.contains("unsupported op: plane.elect") - && record.message.contains("device-runner") - } - _ => false, - }); - assert!(logged, "no error log carried the panic: {events:?}"); - drop(events); - - // The hook pushes what it just recorded out of the process while it - // still can; an aborting teardown must not be able to eat it. - assert!( - *session.flushes.lock().unwrap() > 0, - "the hook never flushed" - ); - } - - #[test] - fn run_context_records_events() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); + use super::*; - run.log_info("hello"); + const PANIC: &str = "kernel exploded mid-batch"; + const REASON: &str = "interrupted"; - let events = session.events.lock().unwrap(); - assert_eq!(events.len(), 1); - match &events[0] { - Event::Log { record, .. } => assert_eq!(record.message, "hello"), - event => panic!("unexpected event: {event:?}"), - } + fn metrics() -> Vec { + vec![MetricValue { + name: "mean_score".to_string(), + value: 0.8, + }] } - #[test] - fn telemetry_records_the_emitting_activity_scope() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let activity = run.activity("cross-validation").start(); - let items = || { - vec![MetricValue { - name: "mean_score".to_string(), - value: 0.8, - }] - }; - - activity.log_info("in scope"); - activity.log_metric(1, "train", 1, items()); - activity.log_epoch_summary(1, "train", items()); - activity.log_summary(items()); - run.log_info("out of scope"); - - let events = session.events.lock().unwrap(); - let event_activity = |event: &Event| match event { + fn recorded_activity(event: &Event) -> Option { + match event { Event::Log { activity, .. } | Event::Metrics { activity, .. } | Event::EpochSummary { activity, .. } | Event::Summary { activity, .. } => *activity, - event => panic!("unexpected event: {event:?}"), - }; - assert!( - events[1..=4] - .iter() - .all(|event| event_activity(event) == Some(activity.id())) - ); - assert_eq!(event_activity(&events[5]), None); + event => panic!("{event:?} is not telemetry"), + } } #[test] - fn activity_run_finishes_fails_and_propagates_panics() { + fn a_run_unwound_by_a_panic_fails_with_the_panic() { let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - - let value = run - .activity("ok") - .run(|_| Ok::<_, &'static str>(7)) - .unwrap(); - assert_eq!(value, 7); - - let error = run - .activity("error") - .run(|_| Err::<(), _>("fold failed")) - .unwrap_err(); - assert_eq!(error, "fold failed"); - - let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - let _: Result<(), &'static str> = run.activity("panic").run(|_| panic!("boom")); - })); - assert!(panic.is_err()); - - let events = session.events.lock().unwrap(); - let completions = events - .iter() - .filter_map(|event| match event { - Event::Activity(ActivityEvent::Finished { - status, message, .. - }) => Some((status, message.as_deref())), - _ => None, - }) - .collect::>(); - match completions.as_slice() { - [ - (ActivityStatus::Success, None), - (ActivityStatus::Failed, Some("fold failed")), - (ActivityStatus::Failed, Some(panicked)), - ] => assert!(panicked.contains("boom"), "{panicked}"), - other => panic!("unexpected completions: {other:?}"), - } + let run_session = session.clone(); + + std::thread::spawn(move || { + let _run = create_run(run_session); + panic!("{PANIC}"); + }) + .join() + .unwrap_err(); + + assert!(matches!( + session.completions.lock().unwrap().as_slice(), + [ExperimentCompletion::Failed(reason)] if reason.ends_with(PANIC) + )); + } + + #[rstest] + #[case::log(|activity: &Activity| activity.log_info("in scope"))] + #[case::metric(|activity: &Activity| activity.log_metric(1, "train", 1, metrics()))] + #[case::epoch_summary(|activity: &Activity| activity.log_epoch_summary(1, "train", metrics()))] + #[case::summary(|activity: &Activity| activity.log_summary(metrics()))] + fn telemetry_through_an_activity_records_the_activity( + mock_run: MockRun, + #[case] emit: fn(&Activity), + ) { + let guard = mock_run.run.activity("cross-validation").start(); + + emit(&guard); + + let events = mock_run.session.events.lock().unwrap(); + assert_eq!(recorded_activity(events.last().unwrap()), Some(guard.id())); } - #[test] - fn child_activity_from_reference_uses_reference_as_parent() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let parent = run.activity("parent").start(); - let activity = parent.share(); + #[rstest] + fn an_activity_started_from_another_is_its_child(mock_run: MockRun) { + let parent = mock_run.run.activity("parent").start(); - let _child = activity.activity("child").start(); + let _child = parent.share().activity("child").start(); - let events = session.events.lock().unwrap(); - let Event::Activity(ActivityEvent::Started { activity: child }) = &events[1] else { - panic!("unexpected event: {:?}", events[1]); - }; - assert_eq!(child.parent, Some(parent.id())); + let started = mock_run.session.started_activities(); + assert_eq!(started.last().unwrap().parent, Some(parent.id())); } - #[test] - fn activity_reference_preserves_inherited_log_attributes() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let guard = run + #[rstest] + fn an_activity_keeps_the_attributes_of_the_handle_it_started_from(mock_run: MockRun) { + let guard = mock_run + .run .handle() .with_attr("pipeline", "cv") .activity("fold") .start(); - let activity = guard.share().with_attr("fold", 1i64); - activity.log(LogRecord::info("started").with("fold", 2i64)); + guard.share().log_info("started"); - let events = session.events.lock().unwrap(); - let Event::Log { - record, - activity: Some(id), - } = &events[1] - else { - panic!("unexpected event: {:?}", events[1]); - }; - assert_eq!(*id, guard.id()); assert_eq!( - record.attributes.get("pipeline"), - Some(&serde_json::json!("cv")) + mock_run.session.logs()[0].attributes.get("pipeline"), + Some(&json!("cv")) ); - assert_eq!(record.attributes.get("fold"), Some(&serde_json::json!(2))); } - #[test] - fn activity_reference_discards_telemetry_after_the_run_is_dropped() { - let session = Arc::new(MockSession::default()); - let activity = { - let run = create_run(session.clone()); - let guard = run.activity("fold").start(); - let activity = guard.share(); - guard.finish(); - activity - }; - let recorded = session.events.lock().unwrap().len(); + #[rstest] + fn a_log_through_a_scoped_handle_inherits_its_attributes(mock_run: MockRun) { + let scoped = mock_run.run.handle().with_attr("phase", "train"); - assert!(!activity.is_active()); - activity.log_info("late"); + scoped.log_warn("slow step"); - assert_eq!(session.events.lock().unwrap().len(), recorded); + assert_eq!( + mock_run.session.logs()[0].attributes.get("phase"), + Some(&json!("train")) + ); } - #[test] - fn scoped_handle_inherits_attributes_and_call_site_wins() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); + #[rstest] + #[case::finished(|run: ExperimentRun| run.finish().unwrap())] + #[case::dropped(drop)] + fn a_handle_of_a_run_that_ended_is_inactive(mock_run: MockRun, #[case] end: fn(ExperimentRun)) { + let handle = mock_run.run.handle(); - let scoped = run - .handle() - .with_attr("phase", "train") - .with_attr("shard", 1i64); - scoped.log( - LogRecord::warn("slow step") - .with("shard", 2i64) - .with("elapsed_ms", 900i64), - ); + end(mock_run.run); - let events = session.events.lock().unwrap(); - assert_eq!(events.len(), 1); - match &events[0] { - Event::Log { record, .. } => { - assert_eq!(record.level, LogLevel::Warn); - assert_eq!(record.message, "slow step"); - // Inherited from scope. - assert_eq!( - record.attributes.get("phase").and_then(|v| v.as_str()), - Some("train") - ); - // Call-site value overrides the inherited scope value. - assert_eq!( - record.attributes.get("shard").and_then(|v| v.as_i64()), - Some(2) - ); - assert_eq!( - record.attributes.get("elapsed_ms").and_then(|v| v.as_i64()), - Some(900) - ); - } - event => panic!("unexpected event: {event:?}"), - } + assert!(!handle.is_active()); } - #[test] - fn finish_marks_handle_inactive_and_discards_later_telemetry() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let handle = run.handle(); + #[rstest] + #[case::finished(|run: ExperimentRun| run.finish().unwrap())] + #[case::dropped(drop)] + fn a_log_through_a_handle_of_a_run_that_ended_is_discarded( + mock_run: MockRun, + #[case] end: fn(ExperimentRun), + ) { + let handle = mock_run.run.handle(); + end(mock_run.run); - run.finish().unwrap(); + handle.log_info("late"); - assert!(!handle.is_active()); - handle.log_info("after-finish"); - assert!(session.events.lock().unwrap().is_empty()); + assert!(mock_run.session.logs().is_empty()); } - #[test] - fn a_run_dropped_before_it_finishes_fails() { - let session = Arc::new(MockSession::default()); - let finished = create_run(session.clone()); + #[rstest] + fn an_activity_started_from_a_handle_of_a_dropped_run_records_nothing(mock_run: MockRun) { + let handle = mock_run.run.handle(); + drop(mock_run.run); - drop(create_run(session.clone())); - finished.finish().unwrap(); + let _late = handle.activity("late").meter(1, "items").start(); - let completions = session.completions.lock().unwrap(); - assert_eq!( - completions.as_slice(), - &[ - ExperimentCompletion::Failed(DROPPED_UNFINISHED.to_string()), - ExperimentCompletion::Success, - ] - ); + assert!(mock_run.session.events.lock().unwrap().is_empty()); } - #[test] - fn a_cancelled_run_ends_by_how_it_is_finished() { - let session = Arc::new(MockSession::default()); + #[rstest] + fn a_run_dropped_before_it_finishes_fails(mock_run: MockRun) { + drop(mock_run.run); - let completed = create_run(session.clone()); - completed.cancel().unwrap(); - completed.finish().unwrap(); - let failed = create_run(session.clone()); - failed.cancel().unwrap(); - failed.fail("interrupted").unwrap(); + assert!(matches!( + mock_run.session.completions.lock().unwrap().as_slice(), + [ExperimentCompletion::Failed(_)] + )); + } + + #[rstest] + fn a_finished_run_ends_once_as_completed(mock_run: MockRun) { + mock_run.run.finish().unwrap(); - let completions = session.completions.lock().unwrap(); assert_eq!( - completions.as_slice(), - &[ - ExperimentCompletion::Success, - ExperimentCompletion::Failed("interrupted".to_string()), - ] + *mock_run.session.completions.lock().unwrap(), + [ExperimentCompletion::Success] ); } - #[test] - fn cancel_does_not_prevent_logging_before_drop() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); + #[rstest] + #[case::finished(|run: ExperimentRun| run.finish().unwrap(), ExperimentCompletion::Success)] + #[case::failed( + |run: ExperimentRun| run.fail(REASON).unwrap(), + ExperimentCompletion::Failed(REASON.to_string()) + )] + fn a_cancelled_run_ends_as_it_is_ended( + mock_run: MockRun, + #[case] end: fn(ExperimentRun), + #[case] completion: ExperimentCompletion, + ) { + mock_run.run.cancel().unwrap(); - run.cancel().unwrap(); - run.log_info("still-logging"); + end(mock_run.run); - let events = session.events.lock().unwrap(); - match events.last() { - Some(Event::Log { record, .. }) => assert_eq!(record.message, "still-logging"), - event => panic!("unexpected event: {event:?}"), - } + assert_eq!(*mock_run.session.completions.lock().unwrap(), [completion]); } - /// The level, message, attributes and activity of each log `session` recorded. - fn logs(session: &MockSession) -> Vec<(LogLevel, String, usize, Option)> { - session - .events - .lock() - .unwrap() - .iter() - .filter_map(|event| match event { - Event::Log { record, activity } => Some(( - record.level, - record.message.clone(), - record.attributes.len(), - *activity, - )), - _ => None, - }) - .collect() + #[rstest] + fn a_cancelled_run_keeps_logging_after_the_stop_warning(mock_run: MockRun) { + mock_run.run.cancel().unwrap(); + + mock_run.run.log_info("still logging"); + + assert_eq!( + mock_run.session.log_messages(), + [STOP_REQUESTED, "still logging"] + ); } - #[test] - fn a_stop_request_is_logged_once_while_the_run_is_active() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let activity = run.activity("epoch 1").cancellable().start(); + #[rstest] + fn cancelling_an_activity_logs_no_stop_request(mock_run: MockRun) { + let activity = mock_run.run.activity("epoch 1").cancellable().start(); activity.cancel_token().cancel(); - assert!(logs(&session).is_empty()); - run.cancel_token().cancel(); - run.cancel().unwrap(); - run.cancel_token().cancel(); - assert_eq!( - logs(&session), - [(LogLevel::Warn, STOP_REQUESTED.to_string(), 0, None)] - ); + assert!(mock_run.session.logs().is_empty()); + } + + #[rstest] + fn a_stop_request_is_logged_once_however_often_the_run_is_cancelled(mock_run: MockRun) { + mock_run.run.cancel_token().cancel(); + mock_run.run.cancel().unwrap(); + mock_run.run.cancel_token().cancel(); + + assert_eq!(mock_run.session.log_messages(), [STOP_REQUESTED]); } #[test] @@ -1287,34 +1068,16 @@ mod tests { let _run = create_run_with_control(session.clone(), control); - assert_eq!( - logs(&session), - [(LogLevel::Warn, STOP_REQUESTED.to_string(), 0, None)] - ); + assert_eq!(session.log_messages(), [STOP_REQUESTED]); } - #[test] - fn a_stop_requested_once_the_run_has_finished_is_not_logged() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let cancel_token = run.cancel_token(); + #[rstest] + fn a_stop_requested_once_the_run_has_finished_is_not_logged(mock_run: MockRun) { + let cancel_token = mock_run.run.cancel_token(); + mock_run.run.finish().unwrap(); - run.finish().unwrap(); cancel_token.cancel(); - assert!(logs(&session).is_empty()); - } - - #[test] - fn dropped_run_handle_activity_records_no_event() { - let session = Arc::new(MockSession::default()); - let run = create_run(session.clone()); - let handle = run.handle(); - drop(run); - - let _progress = handle.activity("late").meter(1, "items").start(); - - let events = session.events.lock().unwrap(); - assert!(events.is_empty()); + assert!(mock_run.session.logs().is_empty()); } } diff --git a/crates/tracel-experiment/src/local.rs b/crates/tracel-experiment/src/local.rs deleted file mode 100644 index fae5e008..00000000 --- a/crates/tracel-experiment/src/local.rs +++ /dev/null @@ -1,1013 +0,0 @@ -//! Recording experiments on this machine, without reaching any server. - -use std::collections::HashMap; -use std::fs::{self, File, OpenOptions}; -use std::io::{self, Write}; -use std::path::{Path, PathBuf}; -use std::sync::Mutex; -use std::sync::mpsc::{Receiver, Sender, channel}; -use std::thread::{self, JoinHandle}; - -use chrono::{SecondsFormat, Utc}; -use serde::Serialize; -use serde_json::{Map, Value}; -use tracel_artifact::bundle::FsBundle; - -use crate::error::{ExperimentError, ExperimentErrorKind}; -use crate::reader::{ArtifactRef, ExperimentArtifactReader, ExperimentReaderError, LoadedArtifact}; -use crate::session::{BundleFn, Event, ExperimentCompletion, ExperimentSession}; -use crate::{ - ActivityEvent, ActivitySpec, ActivityStatus, ArtifactKind, CancelToken, ExperimentId, - ExperimentLocation, ExperimentProvider, ExperimentRun, MetricSpec, MetricValue, -}; - -/// The file a run appends its events to, one JSON object per line. -const EVENTS_FILE: &str = "events.jsonl"; -/// The file a run keeps its [`Status`] in. -const STATUS_FILE: &str = "status.json"; -/// The `.gitignore` written in a runs directory this backend creates: everything in it stays out -/// of git and of code packages. -const RUNS_GITIGNORE: &str = "# Created by tracel: experiment runs recorded on this machine.\n*\n"; - -/// Records experiments under a directory on this machine. -/// -/// The runs of the experiment named `name` are numbered from 1 under `/`. Each run -/// keeps its status in `status.json`, written when the run starts and again when it ends, appends -/// its events to `events.jsonl`, and saves its artifacts under `artifacts/`. Creating a -/// `LocalExperiments` performs no I/O; the first run creates the directories. When it creates -/// `dir` itself, it writes a `.gitignore` there that ignores everything under it, so the records -/// stay out of git and of code packages; a `dir` that already exists is left as it is. -/// -/// `status.json` names the experiment and numbers the run, says whether the run is `running`, -/// `completed` or `failed`, when it started and ended, in RFC 3339 UTC to the second, and why it -/// failed: -/// -/// ```json -/// { -/// "name": "mnist", -/// "num": 3, -/// "status": "failed", -/// "started_at": "2026-10-06T14:02:11Z", -/// "finished_at": "2026-10-06T14:31:40Z", -/// "error": "loss is NaN" -/// } -/// ``` -/// -/// Each line of `events.jsonl` is one event, a JSON object whose `type` names it, with the -/// fields the console's experiment API gives it: -/// -/// | `type` | Fields | -/// | --- | --- | -/// | `attribute` | `key`, `value` | -/// | `arguments` | `value`, the run's input, which the console lists as the experiment's `config`; the last one is the run's | -/// | `config` | `name`, `value` | -/// | `log` | `timestamp`, `level`, `message`, `metadata`, `activity` | -/// | `metric` | `name`, `epoch`, `iteration`, `value`, `group`, `activity` | -/// | `metric_definition` | `name`, `description`, `unit`, `higher_is_better` | -/// | `epoch_summary` | `name`, `epoch`, `group`, `value`, `activity` | -/// | `summary` | `name`, `value`, `activity` | -/// | `input_used` | `experiment`, `name`, `artifact_id` | -/// | `activity_started` | `activity`: `id`, `parent`, `name`, `cancellable`, `meter`, `attributes` | -/// | `activity_updated` | `id`, `current` | -/// | `activity_message` | `id`, `message` | -/// | `activity_finished` | `id`, `status`, `message` | -/// -/// `activity` is the ID of the activity an event was recorded in, or `null`. A metric logged with -/// several values is one `metric` line per value. -/// -/// ``` -/// use std::sync::Arc; -/// -/// use tracel_experiment::Experiments; -/// use tracel_experiment::local::LocalExperiments; -/// -/// let experiments = Experiments::new(Arc::new(LocalExperiments::new("./runs"))); -/// ``` -#[derive(Debug, Clone)] -pub struct LocalExperiments { - dir: PathBuf, -} - -impl LocalExperiments { - /// Records experiments under `dir`. - pub fn new(dir: impl Into) -> Self { - Self { dir: dir.into() } - } -} - -impl ExperimentProvider for LocalExperiments { - fn create_experiment( - &self, - name: String, - attributes: HashMap, - ) -> Result { - let internal = |message: &'static str| { - move |error| ExperimentError::with_source(ExperimentErrorKind::Internal, message, error) - }; - create_runs_dir(&self.dir) - .map_err(internal("Failed to create local experiment directory"))?; - let root = self.dir.join(&name); - let root = root - .canonicalize() - .or_else(|_| { - fs::create_dir_all(&root)?; - root.canonicalize() - }) - .map_err(internal("Failed to create local experiment directory"))?; - let (num, run_root) = create_local_run_dir(&root) - .map_err(internal("Failed to create local experiment run directory"))?; - let session = LocalExperimentSession::start(run_root.clone(), name, num, attributes) - .map_err(internal("Failed to initialize local experiment session"))?; - let reader = LocalExperimentReader { root }; - - let run = ExperimentRun::new(num.to_string(), session, reader, CancelToken::new()); - Ok(run.with_location(ExperimentLocation::Dir(run_root))) - } -} - -/// Creates `dir`, the directory runs are recorded under, with a [`RUNS_GITIGNORE`] in it, unless -/// it already exists. -fn create_runs_dir(dir: &Path) -> io::Result<()> { - if let Some(parent) = dir.parent().filter(|parent| !parent.as_os_str().is_empty()) { - fs::create_dir_all(parent)?; - } - match fs::create_dir(dir) { - Ok(()) => fs::write(dir.join(".gitignore"), RUNS_GITIGNORE), - Err(error) if error.kind() == io::ErrorKind::AlreadyExists => Ok(()), - Err(error) => Err(error), - } -} - -struct LocalExperimentSession { - root: PathBuf, - status: Status, - active: Mutex>, -} - -impl LocalExperimentSession { - /// Starts recording the run `num` of the experiment `name` in `root`: writes its status as - /// running and its attributes as its first events. - fn start( - root: PathBuf, - name: String, - num: u64, - attributes: HashMap, - ) -> io::Result { - fs::create_dir_all(root.join("artifacts"))?; - let mut attributes: Vec<(String, Value)> = attributes.into_iter().collect(); - attributes.sort_by(|(a, _), (b, _)| a.cmp(b)); - let attributes = attributes - .into_iter() - .map(|(key, value)| Line::Attribute { key, value }); - let mut events = OpenOptions::new() - .create(true) - .append(true) - .open(root.join(EVENTS_FILE))?; - events.write_all(to_jsonl(attributes)?.as_bytes())?; - let status = Status::running(name, num); - status.write(&root.join(STATUS_FILE))?; - let (sender, receiver) = channel(); - let join = thread::spawn(move || local_worker(receiver, events)); - - Ok(Self { - root, - status, - active: Mutex::new(Some(LocalWorker { sender, join })), - }) - } - - fn sender(&self) -> Result, ExperimentError> { - let guard = self.active.lock().unwrap(); - guard - .as_ref() - .map(|worker| worker.sender.clone()) - .ok_or_else(|| { - ExperimentError::new( - ExperimentErrorKind::AlreadyFinished, - "Local experiment session has already finished", - ) - }) - } - - /// Stops the writer once it has written every event, then writes the run's status. - fn finish_worker(&self, completion: ExperimentCompletion) -> Result<(), ExperimentError> { - let worker = self.active.lock().unwrap().take().ok_or_else(|| { - ExperimentError::new( - ExperimentErrorKind::AlreadyFinished, - "Local experiment session has already finished", - ) - })?; - - let send_result = worker.sender.send(LocalWrite::Finish); - match worker.join.join() { - Ok(Ok(())) if send_result.is_ok() => {} - Ok(Ok(())) => { - return Err(ExperimentError::new( - ExperimentErrorKind::Internal, - "Failed to send local experiment completion", - )); - } - Ok(Err(err)) => { - return Err(ExperimentError::with_source( - ExperimentErrorKind::Internal, - "Local experiment writer failed", - err, - )); - } - Err(_) => { - return Err(ExperimentError::new( - ExperimentErrorKind::Internal, - "Local experiment writer thread panicked", - )); - } - } - - self.status - .ended(completion) - .write(&self.root.join(STATUS_FILE)) - .map_err(|err| { - ExperimentError::with_source( - ExperimentErrorKind::Internal, - "Failed to write the local experiment status", - err, - ) - }) - } -} - -impl ExperimentSession for LocalExperimentSession { - fn record_event(&self, event: Event) -> Result<(), ExperimentError> { - let lines = to_jsonl(lines(event)).map_err(|err| { - ExperimentError::with_source( - ExperimentErrorKind::Internal, - "Failed to serialize local experiment event", - err, - ) - })?; - self.sender()?.send(LocalWrite::Lines(lines)).map_err(|_| { - ExperimentError::new( - ExperimentErrorKind::Internal, - "Failed to queue local experiment event", - ) - }) - } - - fn save_artifact( - &self, - name: &str, - _kind: ArtifactKind, - artifact: Box, - ) -> Result<(), ExperimentError> { - let artifact_root = self.root.join("artifacts").join(name); - if artifact_root.exists() { - fs::remove_dir_all(&artifact_root).map_err(|err| { - ExperimentError::with_source( - ExperimentErrorKind::Artifact, - "Failed to replace existing local artifact", - err, - ) - })?; - } - - let mut bundle = FsBundle::create(artifact_root.clone()).map_err(|err| { - ExperimentError::with_source( - ExperimentErrorKind::Artifact, - "Failed to create local artifact bundle", - err, - ) - })?; - - let res = artifact(&mut bundle); - - if res.is_err() { - _ = bundle.delete(); - } - - res - } - - fn finish(&self, completion: ExperimentCompletion) -> Result<(), ExperimentError> { - self.finish_worker(completion) - } -} - -/// How a run is going, as its `status.json` gives it. -#[derive(Clone, Serialize)] -struct Status { - name: String, - num: u64, - status: StatusKind, - started_at: String, - finished_at: Option, - error: Option, -} - -#[derive(Clone, Copy, Serialize)] -#[serde(rename_all = "snake_case")] -enum StatusKind { - Running, - Completed, - Failed, -} - -impl Status { - /// The status of the run `num` of the experiment `name`, started now. - fn running(name: String, num: u64) -> Self { - Self { - name, - num, - status: StatusKind::Running, - started_at: now(), - finished_at: None, - error: None, - } - } - - /// The status of the run once it ended now with `completion`. - fn ended(&self, completion: ExperimentCompletion) -> Self { - let (status, error) = match completion { - ExperimentCompletion::Success => (StatusKind::Completed, None), - ExperimentCompletion::Failed(reason) => (StatusKind::Failed, Some(reason)), - }; - Self { - status, - finished_at: Some(now()), - error, - ..self.clone() - } - } - - /// Writes the status to `path` as pretty-printed JSON: to `.tmp` first, then renamed to - /// `path`, so it is never read half written. - fn write(&self, path: &Path) -> io::Result<()> { - let mut contents = serde_json::to_vec_pretty(self)?; - contents.push(b'\n'); - let tmp = path.with_extension("json.tmp"); - let written = fs::write(&tmp, contents).and_then(|()| fs::rename(&tmp, path)); - if written.is_err() { - let _ = fs::remove_file(&tmp); - } - written - } -} - -/// The current time, in RFC 3339 UTC to the second. -fn now() -> String { - Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true) -} - -struct LocalWorker { - sender: Sender, - join: JoinHandle>, -} - -enum LocalWrite { - /// Lines to append to `events.jsonl`, each ending with a newline. - Lines(String), - /// Stop once every line sent before is written. - Finish, -} - -/// One line of `events.jsonl`, with the fields the console's experiment API gives the event. -#[derive(Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum Line { - Attribute { - key: String, - value: Value, - }, - Arguments { - value: Value, - }, - Config { - name: String, - value: Value, - }, - Log { - timestamp: String, - level: &'static str, - message: String, - metadata: Map, - activity: Option, - }, - Metric { - name: String, - epoch: usize, - iteration: usize, - value: f64, - group: String, - activity: Option, - }, - MetricDefinition { - name: String, - description: Option, - unit: Option, - higher_is_better: bool, - }, - EpochSummary { - name: String, - epoch: usize, - group: String, - value: f64, - activity: Option, - }, - Summary { - name: String, - value: f64, - activity: Option, - }, - InputUsed { - experiment: String, - name: String, - artifact_id: String, - }, - ActivityStarted { - activity: ActivitySpec, - }, - ActivityUpdated { - id: u64, - current: u64, - }, - ActivityMessage { - id: u64, - message: String, - }, - ActivityFinished { - id: u64, - status: ActivityStatus, - message: Option, - }, -} - -/// The lines `event` is recorded as: one, or one per value of a metric event. -fn lines(event: Event) -> Vec { - match event { - Event::Args(value) => vec![Line::Arguments { value }], - Event::Config { name, value } => vec![Line::Config { name, value }], - Event::Log { record, activity } => vec![Line::Log { - timestamp: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), - level: record.level.as_str(), - message: record.message, - metadata: record.attributes, - activity: activity.map(|id| id.as_u64()), - }], - Event::Metrics { - epoch, - split, - iteration, - items, - activity, - } => items - .into_iter() - .map(|MetricValue { name, value }| Line::Metric { - name, - epoch, - iteration, - value, - group: split.clone(), - activity: activity.map(|id| id.as_u64()), - }) - .collect(), - Event::MetricDefinition(MetricSpec { - name, - description, - unit, - higher_is_better, - }) => vec![Line::MetricDefinition { - name, - description, - unit, - higher_is_better, - }], - Event::EpochSummary { - epoch, - split, - items, - activity, - } => items - .into_iter() - .map(|MetricValue { name, value }| Line::EpochSummary { - name, - epoch, - group: split.clone(), - value, - activity: activity.map(|id| id.as_u64()), - }) - .collect(), - Event::Summary { items, activity } => items - .into_iter() - .map(|MetricValue { name, value }| Line::Summary { - name, - value, - activity: activity.map(|id| id.as_u64()), - }) - .collect(), - Event::ArtifactUsed { - experiment_id, - reference, - } => vec![Line::InputUsed { - experiment: experiment_id.to_string(), - name: reference.name, - artifact_id: reference.id, - }], - Event::Activity(event) => vec![match event { - ActivityEvent::Started { activity } => Line::ActivityStarted { activity }, - ActivityEvent::Updated { id, current } => Line::ActivityUpdated { - id: id.as_u64(), - current, - }, - ActivityEvent::Message { id, message } => Line::ActivityMessage { - id: id.as_u64(), - message, - }, - ActivityEvent::Finished { - id, - status, - message, - } => Line::ActivityFinished { - id: id.as_u64(), - status, - message, - }, - }], - } -} - -/// `lines` as JSON Lines: each one JSON object followed by a newline. -fn to_jsonl(lines: impl IntoIterator) -> serde_json::Result { - let mut jsonl = String::new(); - for line in lines { - jsonl.push_str(&serde_json::to_string(&line)?); - jsonl.push('\n'); - } - Ok(jsonl) -} - -struct LocalExperimentReader { - root: PathBuf, -} - -impl ExperimentArtifactReader for LocalExperimentReader { - fn load_artifact_raw( - &self, - experiment_id: ExperimentId, - name: &str, - ) -> Result { - let experiment_root = parse_local_experiment_root(&self.root, &experiment_id)?; - let artifact_root = experiment_root.join("artifacts").join(name); - - if !artifact_root.is_dir() { - return Err(ExperimentReaderError::new(format!( - "Local artifact not found: {}", - artifact_root.display() - ))); - } - - let files = collect_bundle_files(&artifact_root, &artifact_root).map_err(|err| { - ExperimentReaderError::with_source("Failed to inspect local artifact files", err) - })?; - let bundle = FsBundle::with_files(artifact_root.clone(), files).map_err(|err| { - ExperimentReaderError::with_source("Failed to create local artifact bundle", err) - })?; - - Ok(LoadedArtifact::new( - ArtifactRef { - id: artifact_root.to_string_lossy().to_string(), - name: name.to_string(), - }, - bundle, - )) - } -} - -fn parse_local_experiment_root( - root: &Path, - experiment_id: &ExperimentId, -) -> Result { - let experiment_dir = experiment_id.as_str(); - if experiment_dir.is_empty() || experiment_dir.contains('/') || experiment_dir.contains('\\') { - return Err(ExperimentReaderError::new( - "Invalid local experiment ID format", - )); - } - - Ok(root.join(experiment_dir)) -} - -/// Appends each batch of lines to `events` as it arrives, until told to finish. -fn local_worker(receiver: Receiver, mut events: File) -> Result<(), io::Error> { - while let Ok(message) = receiver.recv() { - match message { - LocalWrite::Lines(lines) => events.write_all(lines.as_bytes())?, - LocalWrite::Finish => break, - } - } - Ok(()) -} - -fn collect_bundle_files(root: &Path, current: &Path) -> Result, std::io::Error> { - let mut files = Vec::new(); - for entry in fs::read_dir(current)? { - let entry = entry?; - let path = entry.path(); - if path.is_dir() { - files.extend(collect_bundle_files(root, &path)?); - continue; - } - - let rel = path - .strip_prefix(root) - .map_err(std::io::Error::other)? - .to_string_lossy() - .to_string(); - files.push(rel); - } - Ok(files) -} - -/// Creates the directory of the next run under `root`, numbered after the highest one there, and -/// returns its number and path. -fn create_local_run_dir(root: &Path) -> Result<(u64, PathBuf), std::io::Error> { - let mut next_id = 1u64; - - for entry in fs::read_dir(root)? { - let entry = entry?; - if !entry.file_type()?.is_dir() { - continue; - } - - if let Some(id) = entry - .file_name() - .to_str() - .and_then(|name| name.parse::().ok()) - { - next_id = next_id.max(id + 1); - } - } - - loop { - let run_root = root.join(next_id.to_string()); - match fs::create_dir(&run_root) { - Ok(()) => return Ok((next_id, run_root)), - Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { - next_id += 1; - } - Err(err) => return Err(err), - } - } -} - -#[cfg(test)] -mod tests { - use std::io::Read; - - use serde_json::json; - use tracel_artifact::bundle::{BundleDecode, BundleEncode, BundleSink, BundleSource}; - - use super::*; - use crate::LogRecord; - - /// A one-file artifact. - #[derive(Debug, PartialEq)] - struct Note(String); - - impl BundleEncode for Note { - type Settings = (); - type Error = String; - - fn encode(self, sink: &mut O, _settings: &()) -> Result<(), String> { - sink.put_bytes("note.txt", self.0.as_bytes()) - } - } - - impl BundleDecode for Note { - type Settings = (); - type Error = String; - - fn decode(source: &I, _settings: &()) -> Result { - let mut text = String::new(); - source - .open("note.txt")? - .read_to_string(&mut text) - .map_err(|error| error.to_string())?; - Ok(Self(text)) - } - } - - fn run(dir: &Path, name: &str) -> ExperimentRun { - LocalExperiments::new(dir) - .create_experiment(name.to_string(), HashMap::new()) - .unwrap() - } - - #[test] - fn a_runs_directory_it_creates_ignores_everything_in_it() { - let dir = tempfile::tempdir().unwrap(); - let runs = dir.path().join("project/runs"); - - let _run = run(&runs, "mnist"); - - assert_eq!( - fs::read_to_string(runs.join(".gitignore")).unwrap(), - RUNS_GITIGNORE - ); - } - - #[test] - fn a_runs_directory_that_exists_is_left_as_it_is() { - let dir = tempfile::tempdir().unwrap(); - - let _run = run(dir.path(), "mnist"); - - assert!(!dir.path().join(".gitignore").exists()); - } - - #[test] - fn runs_are_numbered_from_one_under_the_experiment_name() { - let dir = tempfile::tempdir().unwrap(); - - let first = run(dir.path(), "mnist"); - let second = run(dir.path(), "mnist"); - let other = run(dir.path(), "other"); - - assert_eq!(first.id().as_str(), "1"); - assert_eq!(second.id().as_str(), "2"); - assert_eq!(other.id().as_str(), "1"); - assert!(dir.path().join("mnist/1").is_dir()); - assert!(dir.path().join("mnist/2").is_dir()); - assert!(dir.path().join("other/1").is_dir()); - } - - #[test] - fn numbering_continues_after_the_highest_existing_run() { - let dir = tempfile::tempdir().unwrap(); - fs::create_dir_all(dir.path().join("mnist/7")).unwrap(); - fs::create_dir_all(dir.path().join("mnist/notes")).unwrap(); - - assert_eq!(run(dir.path(), "mnist").id().as_str(), "8"); - } - - fn read_lines(path: &Path) -> Vec { - fs::read_to_string(path) - .unwrap() - .lines() - .map(|line| serde_json::from_str(line).unwrap()) - .collect() - } - - fn read_json(path: &Path) -> Value { - serde_json::from_str(&fs::read_to_string(path).unwrap()).unwrap() - } - - #[test] - fn a_run_records_each_event_as_a_line_of_json() { - let dir = tempfile::tempdir().unwrap(); - let experiment = LocalExperiments::new(dir.path()) - .create_experiment( - "mnist".to_string(), - HashMap::from([ - ("tracel.job_num".to_string(), json!(12)), - ("kind".to_string(), json!("example")), - ]), - ) - .unwrap(); - - experiment.log_args(&json!({"epochs": 2})).unwrap(); - experiment - .log_config("optimizer", &json!({"lr": 0.01})) - .unwrap(); - experiment.log_metric_definition(MetricSpec { - name: "loss".to_string(), - description: Some("training loss".to_string()), - unit: None, - higher_is_better: false, - }); - let epoch = experiment.activity("epoch 1").meter(4, "batch").start(); - epoch.log(LogRecord::warn("slow batch").with("batch", 3)); - epoch.log_metric( - 1, - "train", - 4, - vec![ - MetricValue { - name: "loss".to_string(), - value: 0.5, - }, - MetricValue { - name: "accuracy".to_string(), - value: 0.9, - }, - ], - ); - epoch.inc(1); - epoch.finish_with_message("done"); - experiment.log_epoch_summary( - 1, - "train", - vec![MetricValue { - name: "loss".to_string(), - value: 0.4, - }], - ); - experiment.log_summary(vec![MetricValue { - name: "accuracy".to_string(), - value: 0.95, - }]); - experiment.finish().unwrap(); - - let mut lines = read_lines(&dir.path().join("mnist/1/events.jsonl")); - let timestamp = lines[6]["timestamp"].take(); - assert!( - timestamp - .as_str() - .is_some_and(|t| chrono::DateTime::parse_from_rfc3339(t).is_ok()), - "{timestamp}" - ); - assert_eq!( - lines, - [ - json!({"type": "attribute", "key": "kind", "value": "example"}), - json!({"type": "attribute", "key": "tracel.job_num", "value": 12}), - json!({"type": "arguments", "value": {"epochs": 2}}), - json!({"type": "config", "name": "optimizer", "value": {"lr": 0.01}}), - json!({ - "type": "metric_definition", - "name": "loss", - "description": "training loss", - "unit": null, - "higher_is_better": false - }), - json!({ - "type": "activity_started", - "activity": { - "id": 1, - "parent": null, - "name": "epoch 1", - "cancellable": false, - "meter": {"unit": "batch", "total": 4}, - "attributes": {} - } - }), - json!({ - "type": "log", - "timestamp": null, - "level": "warn", - "message": "slow batch", - "metadata": {"batch": 3}, - "activity": 1 - }), - json!({ - "type": "metric", - "name": "loss", - "epoch": 1, - "iteration": 4, - "value": 0.5, - "group": "train", - "activity": 1 - }), - json!({ - "type": "metric", - "name": "accuracy", - "epoch": 1, - "iteration": 4, - "value": 0.9, - "group": "train", - "activity": 1 - }), - json!({"type": "activity_updated", "id": 1, "current": 1}), - json!({ - "type": "activity_finished", - "id": 1, - "status": "Success", - "message": "done" - }), - json!({ - "type": "epoch_summary", - "name": "loss", - "epoch": 1, - "group": "train", - "value": 0.4, - "activity": null - }), - json!({"type": "summary", "name": "accuracy", "value": 0.95, "activity": null}), - ] - ); - } - - fn is_utc_seconds(timestamp: &Value) -> bool { - timestamp.as_str().is_some_and(|timestamp| { - timestamp.ends_with('Z') - && !timestamp.contains('.') - && chrono::DateTime::parse_from_rfc3339(timestamp).is_ok() - }) - } - - #[test] - fn the_status_follows_the_run_from_the_start_to_the_end() { - let dir = tempfile::tempdir().unwrap(); - let experiment = run(dir.path(), "mnist"); - let run_dir = dir.path().join("mnist/1"); - - let mut running = read_json(&run_dir.join("status.json")); - experiment.finish().unwrap(); - let mut completed = read_json(&run_dir.join("status.json")); - - assert!(is_utc_seconds(&running["started_at"]), "{running}"); - assert!(is_utc_seconds(&completed["finished_at"]), "{completed}"); - assert_eq!(completed["started_at"], running["started_at"]); - running["started_at"] = Value::Null; - completed["started_at"] = Value::Null; - completed["finished_at"] = Value::Null; - assert_eq!( - running, - json!({ - "name": "mnist", - "num": 1, - "status": "running", - "started_at": null, - "finished_at": null, - "error": null - }) - ); - assert_eq!( - completed, - json!({ - "name": "mnist", - "num": 1, - "status": "completed", - "started_at": null, - "finished_at": null, - "error": null - }) - ); - assert!(!run_dir.join("status.json.tmp").exists()); - } - - #[test] - fn a_run_says_how_it_ended_in_its_status() { - let dir = tempfile::tempdir().unwrap(); - run(dir.path(), "mnist").fail("diverged").unwrap(); - let cancelled = run(dir.path(), "mnist"); - cancelled.cancel().unwrap(); - cancelled.finish().unwrap(); - drop(run(dir.path(), "mnist")); - - let status = |num: u32| read_json(&dir.path().join(format!("mnist/{num}/status.json"))); - let (failed, cancelled, dropped) = (status(1), status(2), status(3)); - assert_eq!(failed["status"], "failed"); - assert_eq!(failed["error"], "diverged"); - assert_eq!(cancelled["status"], "completed"); - assert_eq!(cancelled["error"], Value::Null); - assert_eq!(dropped["status"], "failed"); - assert_eq!( - dropped["error"], - "the run was dropped without finish() or fail()" - ); - } - - #[test] - fn a_run_is_located_in_its_directory() { - let dir = tempfile::tempdir().unwrap(); - - let experiment = run(dir.path(), "mnist"); - - assert_eq!( - experiment.location(), - Some(&ExperimentLocation::Dir( - dir.path().canonicalize().unwrap().join("mnist/1") - )) - ); - } - - #[test] - fn a_saved_artifact_loads_back_from_a_later_run() { - let dir = tempfile::tempdir().unwrap(); - let first = run(dir.path(), "mnist"); - first - .save_artifact("note", ArtifactKind::Other, Note("kept".to_string()), &()) - .unwrap(); - let first_id = first.id().clone(); - first.finish().unwrap(); - - let second = run(dir.path(), "mnist"); - let note: Note = second.use_artifact(first_id, "note", &()).unwrap(); - - assert_eq!(note, Note("kept".to_string())); - } - - #[test] - fn an_experiment_id_naming_another_directory_is_refused() { - let root = Path::new("/runs/mnist"); - - for id in ["", "../other", "a/b", "a\\b"] { - assert!( - parse_local_experiment_root(root, &ExperimentId::new(id)).is_err(), - "{id:?}" - ); - } - assert_eq!( - parse_local_experiment_root(root, &ExperimentId::new("3")).unwrap(), - root.join("3") - ); - } -} diff --git a/crates/tracel-experiment/src/local/artifacts.rs b/crates/tracel-experiment/src/local/artifacts.rs new file mode 100644 index 00000000..f1bf3feb --- /dev/null +++ b/crates/tracel-experiment/src/local/artifacts.rs @@ -0,0 +1,394 @@ +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use relative_path::RelativePath; +use tracel_artifact::bundle::FsBundle; + +use super::runs::child_dir; +use crate::ExperimentId; +use crate::error::{ExperimentError, ExperimentErrorKind}; +use crate::reader::{ArtifactRef, ExperimentArtifactReader, ExperimentReaderError, LoadedArtifact}; +use crate::session::BundleFn; + +/// Saves the artifact `name` of the run recorded in `root`, replacing the one saved before. +/// +/// The artifact is written into a temporary directory, which takes the place of the one saved +/// before only once it is complete, so a failed encode leaves that one as it was. +pub fn save_artifact( + root: &Path, + name: &str, + artifact: Box, +) -> Result<(), ExperimentError> { + let artifact_root = artifact_dir(root, name) + .map_err(|message| ExperimentError::new(ExperimentErrorKind::Artifact, message))?; + let mut staged = FsBundle::temp_in(root.join("artifacts"), ".saving-").map_err(|err| { + ExperimentError::with_source( + ExperimentErrorKind::Artifact, + "Failed to create local artifact bundle", + err, + ) + })?; + + artifact(&mut staged)?; + + replace_dir(staged.root(), &artifact_root).map_err(|err| { + ExperimentError::with_source( + ExperimentErrorKind::Artifact, + "Failed to replace existing local artifact", + err, + ) + }) +} + +pub struct LocalExperimentReader { + pub root: PathBuf, +} + +impl ExperimentArtifactReader for LocalExperimentReader { + fn load_artifact_raw( + &self, + experiment_id: ExperimentId, + name: &str, + ) -> Result { + let experiment_root = parse_local_experiment_root(&self.root, &experiment_id)?; + let artifact_root = + artifact_dir(&experiment_root, name).map_err(ExperimentReaderError::new)?; + + if !artifact_root.is_dir() { + return Err(ExperimentReaderError::new(format!( + "Local artifact not found: {}", + artifact_root.display() + ))); + } + + let files = collect_bundle_files(&artifact_root, &artifact_root).map_err(|err| { + ExperimentReaderError::with_source("Failed to inspect local artifact files", err) + })?; + let bundle = FsBundle::with_files(artifact_root.clone(), files).map_err(|err| { + ExperimentReaderError::with_source("Failed to create local artifact bundle", err) + })?; + + Ok(LoadedArtifact::new( + ArtifactRef { + id: artifact_root.to_string_lossy().to_string(), + name: name.to_string(), + }, + bundle, + )) + } +} + +fn parse_local_experiment_root( + root: &Path, + experiment_id: &ExperimentId, +) -> Result { + child_dir(root, experiment_id.as_str()) + .ok_or_else(|| ExperimentReaderError::new("Invalid local experiment ID format")) +} + +/// The directory the artifact `name` of the run recorded in `root` is saved in. +/// +/// `name` is a relative path inside the run's `artifacts` directory, such as `checkpoints/best`, +/// a leading `/` included. Fails when it resolves to that directory itself or to one outside it, +/// as `..` does, or when it holds a `\`, which separates paths on Windows only. +fn artifact_dir(root: &Path, name: &str) -> Result { + let path = RelativePath::new(name).normalize(); + if name.contains('\\') || path.as_str().is_empty() || path.starts_with("..") { + return Err(format!("Invalid local artifact name: {name:?}")); + } + Ok(path.to_path(root.join("artifacts"))) +} + +/// Moves the directory `staged` to `target`, in place of the one there. +/// +/// A directory cannot be renamed over one that is not empty, so the one there is moved aside +/// first, put back if `staged` cannot take its place, and removed once it has. +fn replace_dir(staged: &Path, target: &Path) -> io::Result<()> { + if let Some(parent) = target.parent() { + fs::create_dir_all(parent)?; + } + let mut replaced = staged.as_os_str().to_owned(); + replaced.push(".replaced"); + let replacing = match fs::rename(target, &replaced) { + Ok(()) => true, + Err(error) if error.kind() == io::ErrorKind::NotFound => false, + Err(error) => return Err(error), + }; + if let Err(error) = fs::rename(staged, target) { + if replacing { + let _ = fs::rename(&replaced, target); + } + return Err(error); + } + if replacing { + let _ = fs::remove_dir_all(&replaced); + } + Ok(()) +} + +fn collect_bundle_files(root: &Path, current: &Path) -> Result, std::io::Error> { + let mut files = Vec::new(); + for entry in fs::read_dir(current)? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + files.extend(collect_bundle_files(root, &path)?); + continue; + } + + let rel = path + .strip_prefix(root) + .map_err(std::io::Error::other)? + .to_string_lossy() + .to_string(); + files.push(rel); + } + Ok(files) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::io::Read; + + use rstest::rstest; + use tracel_artifact::bundle::{BundleDecode, BundleEncode, BundleSink, BundleSource}; + + use super::*; + use crate::local::test_support::{EXPERIMENT, LocalRun, local_run, run}; + use crate::{ArtifactKind, ExperimentRun}; + + #[derive(Debug, PartialEq)] + struct Note(String); + + impl BundleEncode for Note { + type Settings = (); + type Error = String; + + fn encode(self, sink: &mut O, _settings: &()) -> Result<(), String> { + sink.put_bytes("note.txt", self.0.as_bytes()) + } + } + + impl BundleDecode for Note { + type Settings = (); + type Error = String; + + fn decode(source: &I, _settings: &()) -> Result { + let mut text = String::new(); + source + .open("note.txt")? + .read_to_string(&mut text) + .map_err(|error| error.to_string())?; + Ok(Self(text)) + } + } + + #[derive(Debug, PartialEq)] + struct Files(BTreeMap); + + impl Files { + fn of(files: &[(&str, &str)]) -> Self { + Self( + files + .iter() + .map(|(path, contents)| (path.to_string(), contents.to_string())) + .collect(), + ) + } + } + + impl BundleEncode for Files { + type Settings = (); + type Error = String; + + fn encode(self, sink: &mut O, _settings: &()) -> Result<(), String> { + for (path, contents) in self.0 { + sink.put_bytes(&path, contents.as_bytes())?; + } + Ok(()) + } + } + + impl BundleDecode for Files { + type Settings = (); + type Error = String; + + fn decode(source: &I, _settings: &()) -> Result { + let mut files = BTreeMap::new(); + for path in source.list()? { + let mut contents = String::new(); + source + .open(&path)? + .read_to_string(&mut contents) + .map_err(|error| error.to_string())?; + files.insert(path, contents); + } + Ok(Self(files)) + } + } + + struct HalfWritten; + + impl BundleEncode for HalfWritten { + type Settings = (); + type Error = String; + + fn encode(self, sink: &mut O, _settings: &()) -> Result<(), String> { + sink.put_bytes("note.txt", b"partial")?; + Err("disk full".to_string()) + } + } + + fn entries(dir: &Path) -> Vec { + let mut names: Vec = fs::read_dir(dir) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + names.sort(); + names + } + + #[rstest] + fn a_saved_artifact_loads_back_from_a_later_run(local_run: LocalRun) { + let saved = || Note("kept".to_string()); + local_run + .run + .save_artifact("note", ArtifactKind::Other, saved(), &()) + .unwrap(); + let saved_by = local_run.run.id().clone(); + local_run.run.finish().unwrap(); + + let later = run(local_run.runs.path(), EXPERIMENT); + + assert_eq!( + later.use_artifact::(saved_by, "note", &()).unwrap(), + saved() + ); + } + + #[rstest] + #[case::empty("")] + #[case::current(".")] + #[case::parent("..")] + #[case::parent_through_a_nested_name("checkpoints/../..")] + #[case::windows_separator("..\\outside")] + fn an_artifact_name_not_inside_the_run_artifacts_is_refused(#[case] name: &str) { + assert!(artifact_dir(Path::new("run"), name).is_err()); + } + + #[rstest] + fn an_absolute_artifact_name_is_saved_inside_its_run(local_run: LocalRun) { + let outside = local_run.runs.path().join("outside"); + + local_run + .run + .save_artifact( + outside.to_str().unwrap(), + ArtifactKind::Other, + Note("saved".to_string()), + &(), + ) + .unwrap(); + + assert!(!outside.exists()); + } + + #[rstest] + fn an_experiment_id_reaching_through_its_parent_is_not_loaded(local_run: LocalRun) { + local_run + .run + .save_artifact("note", ArtifactKind::Other, Note("kept".to_string()), &()) + .unwrap(); + let through_parent = format!("../{EXPERIMENT}/{}", local_run.run.id()); + + let loaded = local_run + .run + .use_artifact::(through_parent, "note", &()); + + assert!(loaded.is_err()); + } + + fn save_checkpoint(run: &ExperimentRun, files: &[(&str, &str)]) { + run.save_artifact("checkpoint", ArtifactKind::Model, Files::of(files), &()) + .unwrap(); + } + + #[rstest] + fn saving_an_artifact_again_replaces_every_file_of_the_one_saved_before(local_run: LocalRun) { + save_checkpoint( + &local_run.run, + &[("model.bin", "v1"), ("optimizer.bin", "v1")], + ); + + save_checkpoint(&local_run.run, &[("model.bin", "v2")]); + + assert_eq!( + local_run + .run + .use_artifact::(local_run.run.id().clone(), "checkpoint", &()) + .unwrap(), + Files::of(&[("model.bin", "v2")]) + ); + } + + #[rstest] + fn replacing_an_artifact_leaves_nothing_else_in_the_run_artifacts(local_run: LocalRun) { + save_checkpoint(&local_run.run, &[("model.bin", "v1")]); + + save_checkpoint(&local_run.run, &[("model.bin", "v2")]); + + assert_eq!(entries(&local_run.dir.join("artifacts")), ["checkpoint"]); + } + + #[rstest] + fn an_artifact_that_fails_to_encode_leaves_the_one_saved_before(local_run: LocalRun) { + let kept = || Note("kept".to_string()); + local_run + .run + .save_artifact("note", ArtifactKind::Other, kept(), &()) + .unwrap(); + + local_run + .run + .save_artifact("note", ArtifactKind::Other, HalfWritten, &()) + .unwrap_err(); + + assert_eq!( + local_run + .run + .use_artifact::(local_run.run.id().clone(), "note", &()) + .unwrap(), + kept() + ); + } + + #[rstest] + fn an_artifact_never_saved_does_not_load(local_run: LocalRun) { + let loaded = + local_run + .run + .use_artifact::(local_run.run.id().clone(), "missing", &()); + + assert!(loaded.is_err()); + } + + #[test] + fn an_artifact_with_a_nested_name_loads_back_with_its_nested_files() { + let dir = tempfile::tempdir().unwrap(); + let experiment = run(dir.path(), EXPERIMENT); + let files = || Files::of(&[("config.json", "{}"), ("weights/layer1.bin", "w1")]); + + experiment + .save_artifact("checkpoints/best", ArtifactKind::Model, files(), &()) + .unwrap(); + + assert_eq!( + experiment + .use_artifact::(experiment.id().clone(), "checkpoints/best", &()) + .unwrap(), + files() + ); + } +} diff --git a/crates/tracel-experiment/src/local/events.rs b/crates/tracel-experiment/src/local/events.rs new file mode 100644 index 00000000..14c87dc0 --- /dev/null +++ b/crates/tracel-experiment/src/local/events.rs @@ -0,0 +1,184 @@ +use chrono::{SecondsFormat, Utc}; +use serde::Serialize; +use serde_json::{Map, Value}; + +use crate::session::Event; +use crate::{ActivityEvent, ActivitySpec, ActivityStatus, MetricSpec, MetricValue}; + +/// The file a run appends its events to, one JSON object per line. +pub const EVENTS_FILE: &str = "events.jsonl"; + +/// One line of `events.jsonl`, with the fields the console's experiment API gives the event. +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum Line { + Attribute { + key: String, + value: Value, + }, + Arguments { + value: Value, + }, + Config { + name: String, + value: Value, + }, + Log { + timestamp: String, + level: &'static str, + message: String, + metadata: Map, + activity: Option, + }, + Metric { + name: String, + epoch: usize, + iteration: usize, + value: f64, + group: String, + activity: Option, + }, + MetricDefinition { + name: String, + description: Option, + unit: Option, + higher_is_better: bool, + }, + EpochSummary { + name: String, + epoch: usize, + group: String, + value: f64, + activity: Option, + }, + Summary { + name: String, + value: f64, + activity: Option, + }, + InputUsed { + experiment: String, + name: String, + artifact_id: String, + }, + ActivityStarted { + activity: ActivitySpec, + }, + ActivityUpdated { + id: u64, + current: u64, + }, + ActivityMessage { + id: u64, + message: String, + }, + ActivityFinished { + id: u64, + status: ActivityStatus, + message: Option, + }, +} + +/// The lines `event` is recorded as: one, or one per value of a metric event. +pub fn lines(event: Event) -> Vec { + match event { + Event::Args(value) => vec![Line::Arguments { value }], + Event::Config { name, value } => vec![Line::Config { name, value }], + Event::Log { record, activity } => vec![Line::Log { + timestamp: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + level: record.level.as_str(), + message: record.message, + metadata: record.attributes, + activity: activity.map(|id| id.as_u64()), + }], + Event::Metrics { + epoch, + split, + iteration, + items, + activity, + } => items + .into_iter() + .map(|MetricValue { name, value }| Line::Metric { + name, + epoch, + iteration, + value, + group: split.clone(), + activity: activity.map(|id| id.as_u64()), + }) + .collect(), + Event::MetricDefinition(MetricSpec { + name, + description, + unit, + higher_is_better, + }) => vec![Line::MetricDefinition { + name, + description, + unit, + higher_is_better, + }], + Event::EpochSummary { + epoch, + split, + items, + activity, + } => items + .into_iter() + .map(|MetricValue { name, value }| Line::EpochSummary { + name, + epoch, + group: split.clone(), + value, + activity: activity.map(|id| id.as_u64()), + }) + .collect(), + Event::Summary { items, activity } => items + .into_iter() + .map(|MetricValue { name, value }| Line::Summary { + name, + value, + activity: activity.map(|id| id.as_u64()), + }) + .collect(), + Event::ArtifactUsed { + experiment_id, + reference, + } => vec![Line::InputUsed { + experiment: experiment_id.to_string(), + name: reference.name, + artifact_id: reference.id, + }], + Event::Activity(event) => vec![match event { + ActivityEvent::Started { activity } => Line::ActivityStarted { activity }, + ActivityEvent::Updated { id, current } => Line::ActivityUpdated { + id: id.as_u64(), + current, + }, + ActivityEvent::Message { id, message } => Line::ActivityMessage { + id: id.as_u64(), + message, + }, + ActivityEvent::Finished { + id, + status, + message, + } => Line::ActivityFinished { + id: id.as_u64(), + status, + message, + }, + }], + } +} + +/// `lines` as JSON Lines: each one JSON object followed by a newline. +pub fn to_jsonl(lines: impl IntoIterator) -> serde_json::Result { + let mut jsonl = String::new(); + for line in lines { + jsonl.push_str(&serde_json::to_string(&line)?); + jsonl.push('\n'); + } + Ok(jsonl) +} diff --git a/crates/tracel-experiment/src/local/mod.rs b/crates/tracel-experiment/src/local/mod.rs new file mode 100644 index 00000000..f1063f8b --- /dev/null +++ b/crates/tracel-experiment/src/local/mod.rs @@ -0,0 +1,400 @@ +//! Recording experiments on this machine, without reaching any server. + +mod artifacts; +mod events; +mod runs; +mod status; +#[cfg(test)] +mod test_support; + +use std::collections::HashMap; +use std::fs::{self, OpenOptions}; +use std::io::{self, BufWriter, Write}; +use std::mem; +use std::path::PathBuf; +use std::sync::Mutex; + +use serde_json::Value; + +use crate::error::{ExperimentError, ExperimentErrorKind}; +use crate::session::{BundleFn, Event, ExperimentCompletion, ExperimentSession}; +use crate::{ArtifactKind, CancelToken, ExperimentProvider, ExperimentRun}; +use artifacts::{LocalExperimentReader, save_artifact}; +use events::{EVENTS_FILE, Line, lines, to_jsonl}; +use runs::{child_dir, create_local_run_dir, create_runs_dir}; +use status::{STATUS_FILE, Status}; + +/// Records experiments under a directory on this machine. +/// +/// The runs of the experiment named `name` are numbered from 1 under `/`, so `name`, +/// read as a relative path, is one directory name: one that is empty, `.`, `..` or a path of +/// several is refused. Each run keeps its status in `status.json`, written when the run starts and +/// again when it ends, appends its events to `events.jsonl`, and saves its artifacts under +/// `artifacts/`, each in the directory its name gives there, read as a relative path, such as +/// `checkpoints/best`; a name that resolves to `artifacts/` itself or outside it, such as `..`, is +/// refused. Creating a `LocalExperiments` performs no I/O; the first run creates the directories. +/// When it creates `dir` itself, it writes a `.gitignore` there that ignores everything under it, +/// so the records stay out of git and of code packages; a `dir` that already exists is left as it +/// is. +/// +/// `status.json` names the experiment and numbers the run, says whether the run is `running`, +/// `completed` or `failed`, when it started and ended, in RFC 3339 UTC to the second, and why it +/// failed: +/// +/// ```json +/// { +/// "name": "mnist", +/// "num": 3, +/// "status": "failed", +/// "started_at": "2026-10-06T14:02:11Z", +/// "finished_at": "2026-10-06T14:31:40Z", +/// "error": "loss is NaN" +/// } +/// ``` +/// +/// Each line of `events.jsonl` is one event, a JSON object whose `type` names it, with the +/// fields the console's experiment API gives it: +/// +/// | `type` | Fields | +/// | --- | --- | +/// | `attribute` | `key`, `value` | +/// | `arguments` | `value`, the run's input, which the console lists as the experiment's `config`; the last one is the run's | +/// | `config` | `name`, `value` | +/// | `log` | `timestamp`, `level`, `message`, `metadata`, `activity` | +/// | `metric` | `name`, `epoch`, `iteration`, `value`, `group`, `activity` | +/// | `metric_definition` | `name`, `description`, `unit`, `higher_is_better` | +/// | `epoch_summary` | `name`, `epoch`, `group`, `value`, `activity` | +/// | `summary` | `name`, `value`, `activity` | +/// | `input_used` | `experiment`, `name`, `artifact_id` | +/// | `activity_started` | `activity`: `id`, `parent`, `name`, `cancellable`, `meter`, `attributes` | +/// | `activity_updated` | `id`, `current` | +/// | `activity_message` | `id`, `message` | +/// | `activity_finished` | `id`, `status`, `message` | +/// +/// `activity` is the ID of the activity an event was recorded in, or `null`. A metric logged with +/// several values is one `metric` line per value. +/// +/// ``` +/// use std::sync::Arc; +/// +/// use tracel_experiment::Experiments; +/// use tracel_experiment::local::LocalExperiments; +/// +/// let experiments = Experiments::new(Arc::new(LocalExperiments::new("./runs"))); +/// ``` +#[derive(Debug, Clone)] +pub struct LocalExperiments { + dir: PathBuf, +} + +impl LocalExperiments { + /// Records experiments under `dir`. + pub fn new(dir: impl Into) -> Self { + Self { dir: dir.into() } + } +} + +impl ExperimentProvider for LocalExperiments { + fn create_experiment( + &self, + name: String, + attributes: HashMap, + ) -> Result { + let root = child_dir(&self.dir, &name).ok_or_else(|| { + ExperimentError::new( + ExperimentErrorKind::Internal, + format!("Invalid local experiment name: {name:?}"), + ) + })?; + let internal = |message: &'static str| { + move |error| ExperimentError::with_source(ExperimentErrorKind::Internal, message, error) + }; + create_runs_dir(&self.dir) + .map_err(internal("Failed to create local experiment directory"))?; + let root = root + .canonicalize() + .or_else(|_| { + fs::create_dir_all(&root)?; + root.canonicalize() + }) + .map_err(internal("Failed to create local experiment directory"))?; + let (num, run_root) = create_local_run_dir(&root) + .map_err(internal("Failed to create local experiment run directory"))?; + let session = LocalExperimentSession::start(run_root, name, num, attributes) + .map_err(internal("Failed to initialize local experiment session"))?; + let reader = LocalExperimentReader { root }; + + let run = ExperimentRun::new(num.to_string(), session, reader, CancelToken::new()); + Ok(run) + } +} + +struct LocalExperimentSession { + root: PathBuf, + status: Status, + events: Mutex, +} + +/// Where a run's events go: appended to its events file, until a write fails or the run finishes. +enum Events { + Recording(BufWriter>), + Failed(io::Error), + Finished, +} + +impl LocalExperimentSession { + /// Starts recording the run `num` of the experiment `name` in `root`: writes its status as + /// running and its attributes as its first events. + fn start( + root: PathBuf, + name: String, + num: u64, + attributes: HashMap, + ) -> io::Result { + fs::create_dir_all(root.join("artifacts"))?; + let file = OpenOptions::new() + .create(true) + .append(true) + .open(root.join(EVENTS_FILE))?; + Self::start_writing(root, Box::new(file), name, num, attributes) + } + + /// Starts recording the run as [`start`](Self::start) does, with its events written to + /// `events`. + fn start_writing( + root: PathBuf, + events: Box, + name: String, + num: u64, + attributes: HashMap, + ) -> io::Result { + let mut attributes: Vec<(String, Value)> = attributes.into_iter().collect(); + attributes.sort_by(|(a, _), (b, _)| a.cmp(b)); + let attributes = attributes + .into_iter() + .map(|(key, value)| Line::Attribute { key, value }); + let mut events = BufWriter::new(events); + events.write_all(to_jsonl(attributes)?.as_bytes())?; + let status = Status::running(name, num); + status.write(&root.join(STATUS_FILE))?; + + Ok(Self { + root, + status, + events: Mutex::new(Events::Recording(events)), + }) + } + + /// Runs `write` on the run's events file, unless the run has finished or a write has failed: + /// the first write that fails ends the recording. + fn write_events( + &self, + write: impl FnOnce(&mut BufWriter>) -> io::Result<()>, + ) -> Result<(), ExperimentError> { + let mut events = self.events.lock().unwrap(); + let file = match &mut *events { + Events::Recording(file) => file, + Events::Failed(error) => return Err(writer_failed(error)), + Events::Finished => return Err(already_finished()), + }; + write(file).map_err(|error| { + let failure = writer_failed(&error); + *events = Events::Failed(error); + failure + }) + } +} + +impl ExperimentSession for LocalExperimentSession { + fn record_event(&self, event: Event) -> Result<(), ExperimentError> { + let lines = to_jsonl(lines(event)).map_err(|err| { + ExperimentError::with_source( + ExperimentErrorKind::Internal, + "Failed to serialize local experiment event", + err, + ) + })?; + self.write_events(|events| events.write_all(lines.as_bytes())) + } + + fn flush(&self) -> Result<(), ExperimentError> { + self.write_events(|events| events.flush()) + } + + fn save_artifact( + &self, + name: &str, + _kind: ArtifactKind, + artifact: Box, + ) -> Result<(), ExperimentError> { + save_artifact(&self.root, name, artifact) + } + + /// Writes every event recorded, then the run's status: failed, with the writer's error, when + /// an event could not be written. + fn finish(&self, completion: ExperimentCompletion) -> Result<(), ExperimentError> { + let events = mem::replace(&mut *self.events.lock().unwrap(), Events::Finished); + let written = match events { + Events::Recording(mut file) => file.flush().map_err(|error| writer_failed(&error)), + Events::Failed(error) => Err(writer_failed(&error)), + Events::Finished => return Err(already_finished()), + }; + let completion = match &written { + Ok(()) => completion, + Err(error) => ExperimentCompletion::Failed(error.message.clone()), + }; + + self.status + .ended(completion) + .write(&self.root.join(STATUS_FILE)) + .map_err(|err| { + ExperimentError::with_source( + ExperimentErrorKind::Internal, + "Failed to write the local experiment status", + err, + ) + })?; + written + } +} + +fn already_finished() -> ExperimentError { + ExperimentError::new( + ExperimentErrorKind::AlreadyFinished, + "Local experiment session has already finished", + ) +} + +fn writer_failed(error: &io::Error) -> ExperimentError { + ExperimentError::new( + ExperimentErrorKind::Internal, + format!("Local experiment writer failed: {error}"), + ) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + use tempfile::TempDir; + + use super::*; + use crate::MetricValue; + use status::StatusKind; + use test_support::{EXPERIMENT, LocalRun, local_run, read_json, read_lines, run_dir}; + + #[test] + fn an_experiment_name_that_is_not_one_directory_creates_nothing() { + let dir = tempfile::tempdir().unwrap(); + + LocalExperiments::new(dir.path().join("runs")) + .create_experiment("..".to_string(), HashMap::new()) + .err() + .expect("the name is refused"); + + assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 0); + } + + #[test] + fn a_run_records_its_attributes_sorted_by_key() { + let dir = tempfile::tempdir().unwrap(); + let experiment = LocalExperiments::new(dir.path()) + .create_experiment( + EXPERIMENT.to_string(), + HashMap::from([ + ("seed".to_string(), json!(42)), + ("kind".to_string(), json!("example")), + ("dataset".to_string(), json!("mnist")), + ]), + ) + .unwrap(); + + experiment.handle().flush(); + + assert_eq!( + read_lines(&run_dir(dir.path(), &experiment).join(EVENTS_FILE)), + [ + json!({"type": "attribute", "key": "dataset", "value": "mnist"}), + json!({"type": "attribute", "key": "kind", "value": "example"}), + json!({"type": "attribute", "key": "seed", "value": 42}), + ] + ); + } + + #[rstest] + fn flush_writes_the_events_recorded_before_it_returns(local_run: LocalRun) { + local_run.run.log_summary(vec![MetricValue { + name: "accuracy".to_string(), + value: 0.95, + }]); + + local_run.run.handle().flush(); + + assert_eq!(read_lines(&local_run.dir.join(EVENTS_FILE)).len(), 1); + } + + #[rstest] + fn a_finished_run_is_completed_in_its_status(local_run: LocalRun) { + local_run.run.finish().unwrap(); + + let status = read_json(&local_run.dir.join(STATUS_FILE)); + assert_eq!(status["status"], json!(StatusKind::Completed)); + } + + #[derive(Default)] + struct FailsFirstWrite { + failed: bool, + } + + impl Write for FailsFirstWrite { + fn write(&mut self, buf: &[u8]) -> io::Result { + if self.failed { + return Ok(buf.len()); + } + self.failed = true; + Err(io::Error::other("no space left on device")) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + struct FailedWrite { + session: LocalExperimentSession, + root: TempDir, + } + + #[fixture] + fn failed_write() -> FailedWrite { + let root = tempfile::tempdir().unwrap(); + let session = LocalExperimentSession::start_writing( + root.path().to_path_buf(), + Box::new(FailsFirstWrite::default()), + EXPERIMENT.to_string(), + 1, + HashMap::new(), + ) + .unwrap(); + session + .record_event(Event::Args(json!({"epochs": 2}))) + .unwrap(); + session.flush().unwrap_err(); + FailedWrite { session, root } + } + + #[rstest] + fn a_run_whose_events_failed_to_write_ends_as_failed(failed_write: FailedWrite) { + let _ = failed_write.session.finish(ExperimentCompletion::Success); + + let status = read_json(&failed_write.root.path().join(STATUS_FILE)); + assert_eq!(status["status"], json!(StatusKind::Failed)); + } + + #[rstest] + fn finishing_a_run_whose_events_failed_to_write_reports_the_failure(failed_write: FailedWrite) { + let finished = failed_write.session.finish(ExperimentCompletion::Success); + + assert_eq!(finished.unwrap_err().kind, ExperimentErrorKind::Internal); + } +} diff --git a/crates/tracel-experiment/src/local/runs.rs b/crates/tracel-experiment/src/local/runs.rs new file mode 100644 index 00000000..9c14343b --- /dev/null +++ b/crates/tracel-experiment/src/local/runs.rs @@ -0,0 +1,167 @@ +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use relative_path::{Component, RelativePath}; + +/// The `.gitignore` written in a runs directory this backend creates: everything in it stays out +/// of git and of code packages. +pub const RUNS_GITIGNORE: &str = + "# Created by tracel: experiment runs recorded on this machine.\n*\n"; + +/// Creates `dir`, the directory runs are recorded under, with a [`RUNS_GITIGNORE`] in it, unless +/// it already exists. +pub fn create_runs_dir(dir: &Path) -> io::Result<()> { + if let Some(parent) = dir.parent().filter(|parent| !parent.as_os_str().is_empty()) { + fs::create_dir_all(parent)?; + } + match fs::create_dir(dir) { + Ok(()) => fs::write(dir.join(".gitignore"), RUNS_GITIGNORE), + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => Ok(()), + Err(error) => Err(error), + } +} + +/// Creates the directory of the next run under `root`, numbered after the highest one there, and +/// returns its number and path. +pub fn create_local_run_dir(root: &Path) -> Result<(u64, PathBuf), std::io::Error> { + let mut next_id = 1u64; + + for entry in fs::read_dir(root)? { + let entry = entry?; + if !entry.file_type()?.is_dir() { + continue; + } + + if let Some(id) = entry + .file_name() + .to_str() + .and_then(|name| name.parse::().ok()) + { + next_id = next_id.max(id + 1); + } + } + + loop { + let run_root = root.join(next_id.to_string()); + match fs::create_dir(&run_root) { + Ok(()) => return Ok((next_id, run_root)), + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { + next_id += 1; + } + Err(err) => return Err(err), + } + } +} + +/// The directory `name` names directly inside `parent`, when `name`, read as a relative path, is +/// one plain component: `None` when it is empty, `.`, `..`, a path of several, or holds a `\`, +/// which separates paths on Windows only. +pub fn child_dir(parent: &Path, name: &str) -> Option { + if name.contains('\\') { + return None; + } + let mut components = RelativePath::new(name).components(); + match (components.next(), components.next()) { + (Some(Component::Normal(dir)), None) => Some(parent.join(dir)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Barrier; + use std::thread; + + use rstest::rstest; + + use super::*; + use crate::local::test_support::{LocalRun, local_run, run}; + + #[test] + fn a_runs_directory_it_creates_ignores_everything_in_it() { + let dir = tempfile::tempdir().unwrap(); + let runs = dir.path().join("project/runs"); + + let _run = run(&runs, "mnist"); + + let gitignore = fs::read_to_string(runs.join(".gitignore")).unwrap(); + assert!( + gitignore.lines().any(|pattern| pattern == "*"), + "{gitignore}" + ); + } + + #[rstest] + fn a_runs_directory_that_exists_is_left_as_it_is(local_run: LocalRun) { + assert!(!local_run.runs.path().join(".gitignore").exists()); + } + + #[rstest] + #[case::plain("mnist")] + #[case::absolute("/mnist")] + fn a_name_of_one_directory_names_that_directory_of_its_parent(#[case] name: &str) { + let parent = Path::new("runs"); + + assert_eq!(child_dir(parent, name), Some(parent.join("mnist"))); + } + + #[rstest] + #[case::empty("")] + #[case::current(".")] + #[case::parent("..")] + #[case::nested("sweep/lr-3")] + #[case::windows_separator("..\\outside")] + fn a_name_that_is_not_exactly_one_directory_is_refused(#[case] name: &str) { + assert_eq!(child_dir(Path::new("runs"), name), None); + } + + #[test] + fn runs_are_numbered_from_one_under_the_experiment_name() { + let dir = tempfile::tempdir().unwrap(); + + let ids = ["mnist", "mnist", "other"].map(|name| run(dir.path(), name).id().to_string()); + + assert_eq!(ids, ["1", "2", "1"]); + } + + #[test] + fn numbering_continues_after_the_highest_existing_run() { + let dir = tempfile::tempdir().unwrap(); + fs::create_dir_all(dir.path().join("mnist/7")).unwrap(); + fs::create_dir_all(dir.path().join("mnist/notes")).unwrap(); + + assert_eq!(run(dir.path(), "mnist").id().as_str(), "8"); + } + + #[test] + fn a_file_named_like_the_next_run_is_numbered_past() { + let dir = tempfile::tempdir().unwrap(); + fs::create_dir(dir.path().join("mnist")).unwrap(); + fs::write(dir.path().join("mnist/1"), "").unwrap(); + + assert_eq!(run(dir.path(), "mnist").id().as_str(), "2"); + } + + #[test] + fn runs_started_at_once_are_numbered_from_one_without_gaps() { + const RUNS: u64 = 8; + let dir = tempfile::tempdir().unwrap(); + let barrier = Barrier::new(RUNS as usize); + + let mut ids: Vec = thread::scope(|scope| { + let runs: Vec<_> = (0..RUNS) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + run(dir.path(), "mnist").id().parse().unwrap() + }) + }) + .collect(); + runs.into_iter().map(|run| run.join().unwrap()).collect() + }); + ids.sort(); + + assert_eq!(ids, (1..=RUNS).collect::>()); + } +} diff --git a/crates/tracel-experiment/src/local/status.rs b/crates/tracel-experiment/src/local/status.rs new file mode 100644 index 00000000..7ff83cd8 --- /dev/null +++ b/crates/tracel-experiment/src/local/status.rs @@ -0,0 +1,119 @@ +use std::fs; +use std::io; +use std::path::Path; + +use chrono::{SecondsFormat, Utc}; +use serde::Serialize; + +use crate::session::ExperimentCompletion; + +/// The file a run keeps its [`Status`] in. +pub const STATUS_FILE: &str = "status.json"; + +/// How a run is going, as its `status.json` gives it. +#[derive(Clone, Serialize)] +pub struct Status { + name: String, + num: u64, + status: StatusKind, + started_at: String, + finished_at: Option, + error: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum StatusKind { + Running, + Completed, + Failed, +} + +impl Status { + /// The status of the run `num` of the experiment `name`, started now. + pub fn running(name: String, num: u64) -> Self { + Self { + name, + num, + status: StatusKind::Running, + started_at: now(), + finished_at: None, + error: None, + } + } + + /// The status of the run once it ended now with `completion`. + pub fn ended(&self, completion: ExperimentCompletion) -> Self { + let (status, error) = match completion { + ExperimentCompletion::Success => (StatusKind::Completed, None), + ExperimentCompletion::Failed(reason) => (StatusKind::Failed, Some(reason)), + }; + Self { + status, + finished_at: Some(now()), + error, + ..self.clone() + } + } + + /// Writes the status to `path` as pretty-printed JSON: to `.tmp` first, then renamed to + /// `path`, so it is never read half written. + pub fn write(&self, path: &Path) -> io::Result<()> { + let mut contents = serde_json::to_vec_pretty(self)?; + contents.push(b'\n'); + let tmp = path.with_extension("json.tmp"); + let written = fs::write(&tmp, contents).and_then(|()| fs::rename(&tmp, path)); + if written.is_err() { + let _ = fs::remove_file(&tmp); + } + written + } +} + +/// The current time, in RFC 3339 UTC to the second. +fn now() -> String { + Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + use crate::local::test_support::EXPERIMENT; + + const REASON: &str = "loss is NaN"; + + fn running() -> Status { + Status::running(EXPERIMENT.to_string(), 1) + } + + #[test] + fn an_ended_status_keeps_when_the_run_started_and_says_when_it_ended() { + let running = running(); + + let ended = running.ended(ExperimentCompletion::Success); + + assert_eq!( + (ended.started_at, ended.finished_at.is_some()), + (running.started_at, true) + ); + } + + #[rstest] + #[case::succeeded(ExperimentCompletion::Success, StatusKind::Completed, None)] + #[case::failed( + ExperimentCompletion::Failed(REASON.to_string()), + StatusKind::Failed, + Some(REASON) + )] + fn an_ended_status_says_how_the_run_ended( + #[case] completion: ExperimentCompletion, + #[case] status: StatusKind, + #[case] error: Option<&str>, + ) { + let ended = running().ended(completion); + + assert_eq!((ended.status, ended.error.as_deref()), (status, error)); + } +} diff --git a/crates/tracel-experiment/src/local/test_support.rs b/crates/tracel-experiment/src/local/test_support.rs new file mode 100644 index 00000000..f0e9283e --- /dev/null +++ b/crates/tracel-experiment/src/local/test_support.rs @@ -0,0 +1,51 @@ +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; + +use rstest::fixture; +use serde_json::Value; +use tempfile::TempDir; + +use super::LocalExperiments; +use crate::{ExperimentProvider, ExperimentRun}; + +pub const EXPERIMENT: &str = "mnist"; + +pub struct LocalRun { + pub run: ExperimentRun, + pub dir: PathBuf, + pub runs: TempDir, +} + +#[fixture] +pub fn local_run() -> LocalRun { + let runs = tempfile::tempdir().unwrap(); + let run = run(runs.path(), EXPERIMENT); + LocalRun { + dir: run_dir(runs.path(), &run), + run, + runs, + } +} + +pub fn run(dir: &Path, name: &str) -> ExperimentRun { + LocalExperiments::new(dir) + .create_experiment(name.to_string(), HashMap::new()) + .unwrap() +} + +pub fn run_dir(runs: &Path, run: &ExperimentRun) -> PathBuf { + runs.join(EXPERIMENT).join(run.id().as_str()) +} + +pub fn read_lines(path: &Path) -> Vec { + fs::read_to_string(path) + .unwrap() + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() +} + +pub fn read_json(path: &Path) -> Value { + serde_json::from_str(&fs::read_to_string(path).unwrap()).unwrap() +} diff --git a/crates/tracel-experiment/src/log.rs b/crates/tracel-experiment/src/log.rs index e661abdf..4f73f038 100644 --- a/crates/tracel-experiment/src/log.rs +++ b/crates/tracel-experiment/src/log.rs @@ -104,3 +104,19 @@ impl LogRecord { } } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn an_attribute_of_the_record_wins_over_an_inherited_one() { + let mut record = LogRecord::warn("slow step").with("shard", 2); + + record.inherit_attrs(&Map::from_iter([("shard".to_string(), json!(1))])); + + assert_eq!(record.attributes.get("shard"), Some(&json!(2))); + } +} diff --git a/crates/tracel-experiment/src/panic_watch.rs b/crates/tracel-experiment/src/panic_watch.rs index bd627c0a..383bceec 100644 --- a/crates/tracel-experiment/src/panic_watch.rs +++ b/crates/tracel-experiment/src/panic_watch.rs @@ -109,3 +109,41 @@ fn describe(info: &PanicHookInfo<'_>) -> String { None => format!("thread '{name}' panicked: {payload}"), } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use crate::LogLevel; + use crate::test_support::{MockRun, mock_run}; + + const PANIC: &str = "unsupported op: plane.elect"; + + fn panic_on_another_thread() { + std::thread::spawn(|| panic!("{PANIC}")).join().unwrap_err(); + } + + #[rstest] + fn a_watched_run_logs_a_panic_on_another_thread_as_an_error(mock_run: MockRun) { + let _watch = mock_run.run.capture_panics(); + + panic_on_another_thread(); + + assert!( + mock_run + .session + .logs() + .iter() + .any(|record| record.level == LogLevel::Error && record.message.ends_with(PANIC)) + ); + } + + #[rstest] + fn a_watched_run_flushes_the_panic_it_logged(mock_run: MockRun) { + let _watch = mock_run.run.capture_panics(); + + panic_on_another_thread(); + + assert!(*mock_run.session.flushes.lock().unwrap() > 0); + } +} diff --git a/crates/tracel-experiment/src/provider.rs b/crates/tracel-experiment/src/provider.rs index 767cdd7d..a7b060e4 100644 --- a/crates/tracel-experiment/src/provider.rs +++ b/crates/tracel-experiment/src/provider.rs @@ -19,11 +19,8 @@ pub trait ExperimentProvider: Send + Sync + 'static { ) -> Result; } -/// Called with the run of each experiment started, once it exists. -type OnStart = Arc; - /// The experiments of a backend, and how each experiment started from them starts: the attributes -/// and arguments it records, and the hooks called with its run. +/// and arguments it records. /// /// [`run`](Self::run) runs code in a new experiment and ends the experiment by what the code /// returns. [`start`](Self::start) starts one and returns its run, for code that ends it itself. @@ -47,7 +44,6 @@ pub struct Experiments { provider: Arc, attributes: HashMap, arguments: Option, - on_start: Vec, } impl Experiments { @@ -57,7 +53,6 @@ impl Experiments { provider, attributes: HashMap::new(), arguments: None, - on_start: Vec::new(), } } @@ -97,20 +92,12 @@ impl Experiments { self } - /// Calls `hook` with the run of each experiment started from the returned experiments, once - /// it exists and before [`start`](Self::start) returns it, after the hooks given before. - pub fn on_start(mut self, hook: impl Fn(&ExperimentRun) + Send + Sync + 'static) -> Self { - self.on_start.push(Arc::new(hook)); - self - } - /// Starts the experiment `name` and returns its run. /// - /// The experiment records the attributes and arguments these experiments give it, and each - /// [`on_start`](Self::on_start) hook is called with the run. The run ends as completed - /// through [`ExperimentRun::finish`], and as failed through [`ExperimentRun::fail`] or when it - /// is dropped before either. Unless the program has installed a `tracing` subscriber, this - /// installs one that forwards events to experiments. + /// The experiment records the attributes and arguments these experiments give it. The run + /// ends as completed through [`ExperimentRun::finish`], and as failed through + /// [`ExperimentRun::fail`] or when it is dropped before either. Unless the program has + /// installed a `tracing` subscriber, this installs one that forwards events to experiments. pub fn start(&self, name: &str) -> Result { let _ = try_init_tracing_subscriber(); @@ -120,9 +107,6 @@ impl Experiments { if let Some(arguments) = &self.arguments { run.record_args(arguments.clone()); } - for hook in &self.on_start { - hook(&run); - } Ok(run) } @@ -144,8 +128,6 @@ impl Experiments { E: From + fmt::Display, { let run = self.start(name)?; - // Worker-thread panics (a kernel compiler, a data loader) land in the - // run's log even when they never unwind the run itself. let _panic_watch = run.capture_panics(); match run.in_scope(|| f(&run)) { Ok(output) => { @@ -163,21 +145,22 @@ impl Experiments { #[cfg(test)] mod tests { use std::sync::Mutex; - use std::sync::atomic::{AtomicBool, Ordering}; + use rstest::{fixture, rstest}; use serde_json::json; use super::*; - use crate::ExperimentId; use crate::session::{Event, ExperimentCompletion}; use crate::test_support::{MockSession, create_run_with_id}; type BoxError = Box; - /// Creates runs numbered 7 on one session, and keeps the attributes each was created with. + const REASON: &str = "diverged"; + + #[derive(Default)] struct MockProvider { session: Arc, - attributes: Arc>>>, + attributes: Mutex>>, } impl ExperimentProvider for MockProvider { @@ -186,39 +169,28 @@ mod tests { _name: String, attributes: HashMap, ) -> Result { - self.attributes.lock().unwrap().push(attributes); - Ok(create_run_with_id("7", self.session.clone())) + let mut created = self.attributes.lock().unwrap(); + created.push(attributes); + Ok(create_run_with_id( + created.len().to_string(), + self.session.clone(), + )) } } struct Fixture { - session: Arc, - attributes: Arc>>>, + provider: Arc, experiments: Experiments, } impl Fixture { - fn new() -> Self { - let session = Arc::new(MockSession::default()); - let attributes = Arc::default(); - let experiments = Experiments::new(Arc::new(MockProvider { - session: session.clone(), - attributes: Arc::clone(&attributes), - })); - Self { - session, - attributes, - experiments, - } - } - fn completions(&self) -> Vec { - self.session.completions.lock().unwrap().clone() + self.provider.session.completions.lock().unwrap().clone() } - /// The arguments the runs recorded, in order. fn arguments(&self) -> Vec { - self.session + self.provider + .session .events .lock() .unwrap() @@ -231,32 +203,30 @@ mod tests { } } - #[test] - fn run_ends_the_experiment_by_what_its_code_returns() { - let fixture = Fixture::new(); - - let output = fixture - .experiments - .run("train", |_run| Ok::<_, BoxError>(3)); - let failed = fixture - .experiments - .run("train", |_run| Err::<(), BoxError>("diverged".into())); - - assert_eq!(output.unwrap(), 3); - assert_eq!(failed.unwrap_err().to_string(), "diverged"); - assert_eq!( - fixture.completions(), - [ - ExperimentCompletion::Success, - ExperimentCompletion::Failed("diverged".to_string()) - ] - ); + #[fixture] + fn fixture() -> Fixture { + let provider = Arc::new(MockProvider::default()); + Fixture { + experiments: Experiments::new(provider.clone()), + provider, + } } - #[test] - fn a_panic_in_run_fails_the_experiment_with_its_message() { - let fixture = Fixture::new(); + #[rstest] + #[case::ok(Ok(()), ExperimentCompletion::Success)] + #[case::err(Err(REASON.into()), ExperimentCompletion::Failed(REASON.to_string()))] + fn run_ends_the_experiment_by_what_its_code_returns( + fixture: Fixture, + #[case] returned: Result<(), BoxError>, + #[case] completion: ExperimentCompletion, + ) { + let _ = fixture.experiments.run("train", |_run| returned); + + assert_eq!(fixture.completions(), [completion]); + } + #[rstest] + fn a_panic_in_run_fails_the_experiment_with_its_message(fixture: Fixture) { let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { fixture .experiments @@ -268,23 +238,20 @@ mod tests { assert!(panicked.is_err()); match fixture.completions().as_slice() { [ExperimentCompletion::Failed(reason)] => { - assert!(reason.contains("loss is NaN"), "{reason}"); + assert!(reason.ends_with("loss is NaN"), "{reason}"); } other => panic!("unexpected completions: {other:?}"), } } - #[test] - fn run_makes_the_run_the_current_experiment() { - let fixture = Fixture::new(); - + #[rstest] + fn run_makes_the_run_the_current_experiment(fixture: Fixture) { let current = fixture.experiments.run("train", |run| { let current = ExperimentRun::current().map(|current| current.id().clone()); Ok::<_, BoxError>(current == Some(run.id().clone())) }); assert!(current.unwrap()); - assert!(ExperimentRun::current().is_none()); } #[test] @@ -303,22 +270,20 @@ mod tests { )) } } - let ran = AtomicBool::new(false); + let mut ran = false; - let error = Experiments::new(Arc::new(Unreachable)) + Experiments::new(Arc::new(Unreachable)) .run("train", |_run| -> Result<(), BoxError> { - ran.store(true, Ordering::SeqCst); + ran = true; Ok(()) }) .unwrap_err(); - assert_eq!(error.to_string(), "the backend is unreachable"); - assert!(!ran.load(Ordering::SeqCst)); + assert!(!ran); } - #[test] - fn a_started_experiment_records_the_attributes() { - let fixture = Fixture::new(); + #[rstest] + fn a_started_experiment_records_the_attributes(fixture: Fixture) { let experiments = fixture .experiments .clone() @@ -329,7 +294,7 @@ mod tests { experiments.start("train").unwrap().finish().unwrap(); assert_eq!( - *fixture.attributes.lock().unwrap(), + *fixture.provider.attributes.lock().unwrap(), [HashMap::from([ ("kind".to_string(), json!("example")), ("seed".to_string(), json!(42)), @@ -337,14 +302,12 @@ mod tests { ); } - #[test] - fn a_started_experiment_records_the_arguments_logging_replaces() { - let fixture = Fixture::new(); + #[rstest] + fn only_experiments_given_arguments_record_them_when_a_run_starts(fixture: Fixture) { let experiments = fixture.experiments.clone().arguments(json!({"epochs": 2})); - let run = experiments.start("train").unwrap(); + experiments.start("train").unwrap().finish().unwrap(); let recorded = fixture.arguments(); - run.log_args(&json!({"epochs": 2, "fold": 1})).unwrap(); fixture .experiments .start("train") @@ -353,40 +316,19 @@ mod tests { .unwrap(); assert_eq!(recorded, [json!({"epochs": 2})]); - assert_eq!( - fixture.arguments(), - [json!({"epochs": 2}), json!({"epochs": 2, "fold": 1})] - ); + assert_eq!(fixture.arguments(), recorded); } - #[test] - fn each_hook_is_called_with_each_run_started() { - let fixture = Fixture::new(); - let called = Arc::new(Mutex::new(Vec::new())); - let hook = |tag: &'static str| { - let called = called.clone(); - move |run: &ExperimentRun| called.lock().unwrap().push((tag, run.id().clone())) - }; - let experiments = fixture - .experiments - .clone() - .on_start(hook("first")) - .on_start(hook("second")); + #[rstest] + fn arguments_logged_on_a_run_are_recorded_after_those_it_started_with(fixture: Fixture) { + let experiments = fixture.experiments.clone().arguments(json!({"epochs": 2})); + let run = experiments.start("train").unwrap(); - experiments.start("train").unwrap().finish().unwrap(); - fixture - .experiments - .start("train") - .unwrap() - .finish() - .unwrap(); + run.log_args(&json!({"epochs": 2, "fold": 1})).unwrap(); assert_eq!( - *called.lock().unwrap(), - [ - ("first", ExperimentId::from("7")), - ("second", ExperimentId::from("7")) - ] + fixture.arguments(), + [json!({"epochs": 2}), json!({"epochs": 2, "fold": 1})] ); } } diff --git a/crates/tracel-experiment/src/test_support.rs b/crates/tracel-experiment/src/test_support.rs index e12628f9..b427ac8e 100644 --- a/crates/tracel-experiment/src/test_support.rs +++ b/crates/tracel-experiment/src/test_support.rs @@ -1,10 +1,14 @@ use std::sync::{Arc, Mutex}; -use crate::activity::ActivityEvent; +use rstest::fixture; +use tracing_subscriber::layer::SubscriberExt; + +use crate::activity::{ActivityEvent, ActivityId, ActivitySpec}; use crate::error::ExperimentError; +use crate::integration::tracing::tracing_log_layer; use crate::reader::{ExperimentArtifactReader, ExperimentReaderError, LoadedArtifact}; use crate::session::{BundleFn, Event, ExperimentCompletion, ExperimentSession}; -use crate::{ArtifactKind, ExperimentId, ExperimentRun, ExperimentRunControl}; +use crate::{ArtifactKind, ExperimentId, ExperimentRun, ExperimentRunControl, LogRecord}; #[derive(Default)] pub struct MockSession { @@ -25,6 +29,47 @@ impl MockSession { }) .collect() } + + pub fn started_activities(&self) -> Vec { + self.activity_events() + .into_iter() + .filter_map(|event| match event { + ActivityEvent::Started { activity } => Some(activity), + _ => None, + }) + .collect() + } + + pub fn logs(&self) -> Vec { + self.events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + Event::Log { record, .. } => Some(record.clone()), + _ => None, + }) + .collect() + } + + pub fn log_activities(&self) -> Vec> { + self.events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + Event::Log { activity, .. } => Some(*activity), + _ => None, + }) + .collect() + } + + pub fn log_messages(&self) -> Vec { + self.logs() + .into_iter() + .map(|record| record.message) + .collect() + } } impl ExperimentSession for MockSession { @@ -90,3 +135,21 @@ pub fn create_run_with_control( control, ) } + +pub struct MockRun { + pub run: ExperimentRun, + pub session: Arc, +} + +#[fixture] +pub fn mock_run() -> MockRun { + let session = Arc::new(MockSession::default()); + MockRun { + run: create_run(session.clone()), + session, + } +} + +pub fn traced(f: impl FnOnce() -> T) -> T { + tracing::subscriber::with_default(tracing_subscriber::registry().with(tracing_log_layer()), f) +} diff --git a/crates/tracel-job/Cargo.toml b/crates/tracel-job/Cargo.toml index b450a567..5241a4b8 100644 --- a/crates/tracel-job/Cargo.toml +++ b/crates/tracel-job/Cargo.toml @@ -24,4 +24,5 @@ referencing = { workspace = true, optional = true } thiserror = { workspace = true, optional = true } [dev-dependencies] +rstest = { workspace = true } schemars.workspace = true diff --git a/crates/tracel-job/src/cli/command.rs b/crates/tracel-job/src/cli/command.rs index 1c25cc7d..40ad4b55 100644 --- a/crates/tracel-job/src/cli/command.rs +++ b/crates/tracel-job/src/cli/command.rs @@ -214,7 +214,10 @@ fn flag_arg(flag: &Flag) -> Arg { arg = arg.required_unless_present_any([INPUT, CONFIG]); } match flag.kind { - Kind::Bool => arg.num_args(0..=1).default_missing_value("true"), + Kind::Bool => arg + .num_args(0..=1) + .require_equals(true) + .default_missing_value("true"), Kind::Integer | Kind::Number | Kind::Json => arg.allow_negative_numbers(true), Kind::String | Kind::Choice(_) => arg, } @@ -259,37 +262,14 @@ fn overlay(input: &mut Value, layer: Value) { } #[cfg(test)] -#[allow(dead_code)] mod tests { - use clap::error::ErrorKind; + use rstest::rstest; use schemars::JsonSchema; use serde::Serialize; use serde_json::json; use super::*; - /// A job's input schema and default input. - struct Definition { - name: &'static str, - schema: Option, - default: Option, - } - - fn jobs(definitions: &[Definition]) -> Vec { - definitions.iter().map(Definition::job).collect() - } - - impl Definition { - fn job(&self) -> JobDefinition { - JobDefinition { - name: self.name.to_string(), - description: Some("Run the job".to_string()), - input_schema: self.schema.clone(), - input_example: self.default.clone(), - } - } - } - #[derive(Serialize, JsonSchema)] struct Training { num_epochs: u32, @@ -297,6 +277,7 @@ mod tests { shuffle: bool, tag: String, layers: Vec, + mode: Mode, resume_from: Option, optimizer: Optimizer, } @@ -312,69 +293,83 @@ mod tests { epoch: u32, } - /// `train`, taking `Training`, with a default input. - fn train() -> Definition { + fn train() -> JobDefinition { let default = Training { num_epochs: 10, batch_size: 64, shuffle: true, tag: "baseline".to_string(), layers: vec![64, 32], + mode: Mode::Fast, resume_from: None, optimizer: Optimizer { lr: 0.001, weight_decay: 5e-5, }, }; - Definition { - name: "train", - schema: Some(schemars::schema_for!(Training).to_value()), - default: Some(serde_json::to_value(default).unwrap()), + JobDefinition { + name: "train".to_string(), + description: Some("Train the model".to_string()), + input_schema: Some(schemars::schema_for!(Training).to_value()), + input_example: Some(serde_json::to_value(default).unwrap()), } } #[derive(JsonSchema)] + #[expect(dead_code)] struct Prompt { - /// The prompt. text: String, mode: Mode, - limit: Option, } - #[derive(JsonSchema)] + #[derive(Serialize, JsonSchema)] #[serde(rename_all = "snake_case")] + #[expect(dead_code)] enum Mode { Fast, Exact, } - /// `prompt`, taking `Prompt`, with no default input. - fn prompt() -> Definition { - Definition { - name: "prompt", - schema: Some(schemars::schema_for!(Prompt).to_value()), - default: None, + fn prompt() -> JobDefinition { + JobDefinition { + name: "prompt".to_string(), + description: Some("Answer a prompt".to_string()), + input_schema: Some(schemars::schema_for!(Prompt).to_value()), + input_example: None, } } - /// `evaluate`, with a default input and no schema. - fn evaluate() -> Definition { - Definition { - name: "evaluate", - schema: None, - default: Some(json!({"split": "test", "limit": 10})), + fn evaluate() -> JobDefinition { + JobDefinition { + name: "evaluate".to_string(), + description: None, + input_schema: None, + input_example: Some(json!({"split": "test", "limit": 10})), + } + } + + fn named_like_flags() -> JobDefinition { + JobDefinition { + name: "train".to_string(), + description: None, + input_schema: Some(json!({ + "type": "object", + "properties": { + "config": {"type": "string"}, + "version": {"type": "integer"} + } + })), + input_example: None, } } - /// The input the job's command line ` ` gives, its `--config` file read from - /// `files`, by name. fn input_from( - definition: &Definition, + job: &JobDefinition, args: &[&str], files: &[(&str, &str)], ) -> Result, clap::Error> { - let matches = job_command(&definition.job()) - .try_get_matches_from([definition.name].iter().chain(args))?; + let matches = + job_command(job).try_get_matches_from([job.name.as_str()].iter().chain(args))?; Ok(job_input(&matches, |path| { files .iter() @@ -384,241 +379,181 @@ mod tests { })) } - /// The input the job's command line ` ` gives, with no `--config` file to read. - fn input(definition: &Definition, args: &[&str]) -> Result { - input_from(definition, args, &[]).map(|input| input.unwrap()) - } - - /// The job's `--help`. - fn help(definition: &Definition) -> String { - job_command(&definition.job()) - .render_long_help() - .to_string() + fn input(job: &JobDefinition, args: &[&str]) -> Result { + input_from(job, args, &[]).map(|input| input.unwrap()) } #[test] - fn the_command_has_a_subcommand_per_job() { - let definitions = [train(), prompt()]; - let mut command = command("program", &jobs(&definitions)); - command.build(); - - let names: Vec<&str> = command.get_subcommands().map(Command::get_name).collect(); - assert_eq!(names, ["train", "prompt"]); - let help = command.render_help().to_string(); - assert!(help.contains("Jobs:"), "{help}"); - assert!(help.contains("Run the job"), "{help}"); - assert!(help.contains("--completions "), "{help}"); + fn each_job_is_a_subcommand_about_its_description() { + let jobs = [train(), prompt(), evaluate()]; + + let command = command("program", &jobs); + + let subcommands: Vec<(&str, Option)> = command + .get_subcommands() + .map(|job| (job.get_name(), job.get_about().map(ToString::to_string))) + .collect(); + let expected: Vec<(&str, Option)> = jobs + .iter() + .map(|job| (job.name.as_str(), job.description.clone())) + .collect(); + assert_eq!(subcommands, expected); + } + + #[rstest] + #[case::at_the_top(&["--num-epochs", "2"], json!({"num_epochs": 2}))] + #[case::in_its_object( + &["--optimizer.weight-decay", "0.5"], + json!({"optimizer": {"weight_decay": 0.5}}) + )] + #[case::beside_another_in_its_object( + &["--optimizer.lr", "0.1", "--optimizer.weight-decay", "0.5"], + json!({"optimizer": {"lr": 0.1, "weight_decay": 0.5}}) + )] + fn a_flag_sets_its_field_under_its_original_keys( + #[case] args: &[&str], + #[case] expected: Value, + ) { + assert_eq!(input(&train(), args).unwrap(), expected); } #[test] - fn flags_set_fields_under_their_original_keys() { - let input = input( - &train(), - &[ - "--num-epochs", - "2", - "--optimizer.weight-decay", - "3e-4", - "--tag", - "42", - ], - ) - .unwrap(); - + fn a_choice_flag_sets_its_field_to_the_value_chosen() { assert_eq!( - input, - json!({"num_epochs": 2, "optimizer": {"weight_decay": 3e-4}, "tag": "42"}) + input(&train(), &["--mode", "exact"]).unwrap(), + json!({"mode": "exact"}) ); } #[test] - fn no_arguments_give_no_input_and_defaults_are_only_shown() { + fn the_defaults_alone_give_no_input() { assert_eq!(input(&train(), &[]).unwrap(), Value::Null); - - let help = help(&train()); - for line in [ - "--num-epochs ", - "[default: 10]", - "--optimizer.weight-decay ", - "[default: 0.00005]", - "--shuffle []", - "--layers ", - "[default: [64,32]]", - "--resume-from ", - "[default: null]", - "-c, --config ", - ] { - assert!(help.contains(line), "no {line} in\n{help}"); - } } - #[test] - fn a_flag_parses_as_its_type() { - let train = train(); + #[rstest] + #[case::integer(&["--batch-size", "-1"], json!({"batch_size": -1}))] + #[case::number(&["--optimizer.lr", "-0.5"], json!({"optimizer": {"lr": -0.5}}))] + #[case::json(&["--layers", "-1"], json!({"layers": -1}))] + fn a_numeric_or_json_flag_takes_a_negative_number( + #[case] args: &[&str], + #[case] expected: Value, + ) { + assert_eq!(input(&train(), args).unwrap(), expected); + } - assert_eq!( - input(&train, &["--shuffle", "--batch-size", "-1"]).unwrap(), - json!({"shuffle": true, "batch_size": -1}) - ); - assert_eq!( - input(&train, &["--shuffle", "false"]).unwrap(), - json!({"shuffle": false}) - ); - for args in [ - &["--num-epochs", "ten"][..], - &["--num-epochs", "1.5"], - &["--optimizer.lr", "fast"], - &["--shuffle", "yes"], - &["--layers", "[1,"], - ] { - let error = input(&train, args).unwrap_err(); - assert!( - matches!( - error.kind(), - ErrorKind::ValueValidation | ErrorKind::InvalidValue - ), - "{args:?}: {error}" - ); - assert_eq!(error.exit_code(), 2); - } - assert_eq!( - input(&train, &["--num-epochs"]).unwrap_err().kind(), - ErrorKind::InvalidValue - ); - assert_eq!( - input(&train, &["--epochs", "2"]).unwrap_err().kind(), - ErrorKind::UnknownArgument - ); + #[rstest] + #[case::integer(&["--num-epochs", "ten"])] + #[case::bool(&["--shuffle=yes"])] + #[case::choice(&["--mode", "slow"])] + fn a_flag_given_a_value_its_field_does_not_take_is_rejected(#[case] args: &[&str]) { + assert!(input(&train(), args).is_err()); } - #[test] - fn json_literal_flags_take_any_json() { - assert_eq!( - input( - &train(), - &["--layers", "[8]", "--resume-from", r#"{"epoch": 3}"#] - ) - .unwrap(), - json!({"layers": [8], "resume_from": {"epoch": 3}}) - ); - assert_eq!( - input(&train(), &["--resume-from", "null"]).unwrap(), - json!({"resume_from": null}) - ); + #[rstest] + #[case::alone(&["--shuffle"], json!({"shuffle": true}))] + #[case::after_an_equals_sign(&["--shuffle=false"], json!({"shuffle": false}))] + #[case::before_the_json_input( + &["--shuffle", r#"{"num_epochs": 2}"#], + json!({"num_epochs": 2, "shuffle": true}) + )] + fn a_bool_flag_takes_a_value_only_after_an_equals_sign( + #[case] args: &[&str], + #[case] expected: Value, + ) { + assert_eq!(input(&train(), args).unwrap(), expected); } - #[test] - fn a_schema_gives_flags_their_values_help_and_requirement() { - let prompt = prompt(); + #[rstest] + #[case::an_object(json!({"epoch": 3}))] + #[case::null(Value::Null)] + fn a_json_flag_sets_its_field_to_the_json_it_is_given(#[case] value: Value) { + let input = input(&train(), &["--resume-from", &value.to_string()]).unwrap(); - assert_eq!( - input( - &prompt, - &["--text", "hi", "--mode", "exact", "--limit", "3"] - ) - .unwrap(), - json!({"text": "hi", "mode": "exact", "limit": 3}) - ); - assert_eq!( - input(&prompt, &["--text", "hi", "--mode", "slow"]) - .unwrap_err() - .kind(), - ErrorKind::InvalidValue - ); - let help = help(&prompt); - assert!(help.contains("The prompt."), "{help}"); - assert!(help.contains("[possible values: fast, exact]"), "{help}"); + assert_eq!(input, json!({"resume_from": value})); } #[test] - fn a_required_flag_can_be_left_to_the_json_input_or_the_config_file() { - let prompt = prompt(); + fn a_required_flag_left_out_is_rejected() { + assert!(input(&prompt(), &["--mode", "fast"]).is_err()); + } + + #[rstest] + #[case::the_json_input(&[r#"{"text": "hi"}"#])] + #[case::the_config_file(&["--config", "config.json"])] + fn a_required_flag_can_be_left_to(#[case] args: &[&str]) { let files = [("config.json", r#"{"text": "from a file"}"#)]; - let missing = input(&prompt, &["--mode", "fast"]).unwrap_err(); + assert!(input_from(&prompt(), args, &files).is_ok()); + } - assert_eq!(missing.kind(), ErrorKind::MissingRequiredArgument); - assert!(missing.to_string().contains("--text "), "{missing}"); - assert_eq!( - input(&prompt, &[r#"{"text": "hi"}"#]).unwrap(), - json!({"text": "hi"}) - ); - assert_eq!( - input_from(&prompt, &["-c", "config.json"], &files) - .unwrap() - .unwrap(), - json!({"text": "from a file"}) - ); + #[rstest] + #[case::an_object(json!({"tag": null, "layers": [], "optimizer": {"lr": 0.1}}))] + #[case::a_number(json!(7))] + #[case::a_negative_number(json!(-7))] + fn the_json_input_alone_is_the_input_as_given(#[case] given: Value) { + assert_eq!(input(&train(), &[&given.to_string()]).unwrap(), given); } #[test] - fn a_job_without_a_schema_takes_its_json_input_and_config_file_only() { - let evaluate = evaluate(); - let files = [("config.json", r#"{"split": "valid"}"#)]; + fn an_input_that_is_not_json_is_rejected() { + assert!(input(&train(), &["{"]).is_err()); + } - let flag = input(&evaluate, &["--split", "valid"]).unwrap_err(); + #[test] + fn the_config_file_then_the_json_input_then_the_flags_win() { + let files = [( + "config.json", + r#"{"num_epochs": 1, "batch_size": 8, "tag": "file"}"#, + )]; + let args = [ + "--num-epochs", + "3", + "--config", + "config.json", + r#"{"num_epochs": 2, "batch_size": 16}"#, + ]; + + let input = input_from(&train(), &args, &files).unwrap().unwrap(); - assert_eq!(flag.kind(), ErrorKind::UnknownArgument); - assert_eq!(flag.exit_code(), 2); - assert_eq!( - input(&evaluate, &[r#"{"split": "valid"}"#]).unwrap(), - json!({"split": "valid"}) - ); assert_eq!( - input_from(&evaluate, &["--config", "config.json"], &files) - .unwrap() - .unwrap(), - json!({"split": "valid"}) + input, + json!({"num_epochs": 3, "batch_size": 16, "tag": "file"}) ); - let help = help(&evaluate); - assert!(help.contains("[INPUT]"), "{help}"); - assert!(help.contains("-c, --config "), "{help}"); - assert!(!help.contains("--split"), "{help}"); } - #[test] - fn the_json_input_alone_is_the_input_as_given() { - let given = json!({"tag": null, "layers": [], "optimizer": {"lr": 0.1}}); + #[rstest] + #[case::a_value_onto_an_object(json!({"num_epochs": 1}), json!(7))] + #[case::an_object_onto_a_value(json!(7), json!({"num_epochs": 3}))] + fn a_layer_replaces_what_is_below_unless_both_are_objects( + #[case] mut below: Value, + #[case] layer: Value, + ) { + overlay(&mut below, layer.clone()); - assert_eq!(input(&train(), &[&given.to_string()]).unwrap(), given); - assert_eq!(input(&train(), &["7"]).unwrap(), json!(7)); - assert_eq!(input(&train(), &["-7"]).unwrap(), json!(-7)); + assert_eq!(below, layer); } #[test] - fn the_config_file_then_the_json_input_then_the_flags_win() { - let files = [( - "config.json", - r#"{"num_epochs": 1, "batch_size": 8, "tag": "file", "optimizer": {"lr": 0.1, "weight_decay": 0.5}}"#, - )]; + fn an_object_layer_is_overlaid_field_by_field() { + let mut below = json!({"num_epochs": 1, "optimizer": {"lr": 0.1, "weight_decay": 0.5}}); - let input = input_from( - &train(), - &[ - "--num-epochs", - "3", - "--config", - "config.json", - r#"{"num_epochs": 2, "batch_size": 16, "tag": null, "optimizer": {"lr": 0.2}}"#, - "--optimizer.lr", - "0.3", - ], - &files, - ) - .unwrap() - .unwrap(); + overlay(&mut below, json!({"optimizer": {"lr": 0.2}})); assert_eq!( - input, - json!({ - "num_epochs": 3, - "batch_size": 16, - "tag": null, - "optimizer": {"lr": 0.3, "weight_decay": 0.5} - }) + below, + json!({"num_epochs": 1, "optimizer": {"lr": 0.2, "weight_decay": 0.5}}) ); } + #[test] + fn a_null_field_is_kept_rather_than_applied() { + let mut below = json!({"num_epochs": 1, "tag": "file"}); + + overlay(&mut below, json!({"tag": null})); + + assert_eq!(below, json!({"num_epochs": 1, "tag": null})); + } + #[test] fn a_repeated_flag_takes_its_last_value() { assert_eq!( @@ -628,117 +563,61 @@ mod tests { } #[test] - fn an_unreadable_or_invalid_config_file_gives_no_input() { - let files = [("config.toml", "epochs = 3")]; - let error = |path| { - input_from(&train(), &["-c", path], &files) - .unwrap() - .unwrap_err() - }; + fn an_unreadable_config_file_is_an_error_naming_it() { + let missing = "missing.json"; - let missing = error("missing.json"); - let not_json = error("config.toml"); + let error = input_from(&train(), &["--config", missing], &[]) + .unwrap() + .unwrap_err(); assert!( - matches!(&missing, ConfigError::Unreadable { path, .. } if path == Path::new("missing.json")), - "{missing:?}" - ); - assert!( - missing - .to_string() - .starts_with("cannot read the --config file missing.json: ") - ); - assert!( - matches!(¬_json, ConfigError::NotJson { .. }), - "{not_json:?}" - ); - assert!( - not_json - .to_string() - .starts_with("the --config file config.toml is not JSON: ") + matches!(error, ConfigError::Unreadable { path, .. } if path == Path::new(missing)) ); } #[test] - fn reserved_names_belong_to_the_command_line() { - let job = Definition { - name: "train", - schema: Some(json!({ - "type": "object", - "properties": { - "config": {"type": "string"}, - "help": {"type": "boolean"}, - "version": {"type": "integer"}, - "epochs": {"type": "integer"} - } - })), - default: None, - }; - let files = [("large.json", r#"{"config": "large"}"#)]; + fn a_config_file_that_is_not_json_is_an_error_naming_it() { + let toml = "config.toml"; - let command = job_command(&job.job()); - let longs: Vec<&str> = command.get_arguments().filter_map(Arg::get_long).collect(); - assert_eq!(longs, ["epochs", "config"]); - assert_eq!( - input_from(&job, &["--config", "large.json"], &files) - .unwrap() - .unwrap(), - json!({"config": "large"}) - ); - assert_eq!( - input(&job, &["--help"]).unwrap_err().kind(), - ErrorKind::DisplayHelp - ); - assert_eq!( - input(&job, &["--version"]).unwrap_err().kind(), - ErrorKind::UnknownArgument - ); + let error = input_from(&train(), &["--config", toml], &[(toml, "epochs = 3")]) + .unwrap() + .unwrap_err(); + + assert!(matches!(error, ConfigError::NotJson { path, .. } if path == Path::new(toml))); } #[test] - fn completions_names_the_shell_the_command_line_asks_for() { - let definitions = [train()]; - let matches = |args: &[&str]| { - command("program", &jobs(&definitions)) - .try_get_matches_from(["program"].iter().chain(args)) - .unwrap() - }; + fn the_config_flag_names_the_config_file_beside_a_field_named_config() { + let files = [("large.json", r#"{"config": "large"}"#)]; - assert_eq!( - completions(&matches(&["--completions", "zsh"])), - Some(Shell::Zsh) - ); - assert_eq!(completions(&matches(&["train"])), None); - assert_eq!( - command("program", &jobs(&definitions)) - .try_get_matches_from(["program", "--completions", "tcsh"]) - .unwrap_err() - .kind(), - ErrorKind::InvalidValue - ); + let input = input_from(&named_like_flags(), &["--config", "large.json"], &files) + .unwrap() + .unwrap(); + + assert_eq!(input, json!({"config": "large"})); } #[test] - fn completions_cover_the_jobs_and_their_flags() { - let definitions = [train(), prompt()]; - let mut script = Vec::new(); - - clap_complete::generate( - Shell::Bash, - &mut command("program", &jobs(&definitions)), - "program", - &mut script, + fn a_field_named_version_is_set_by_its_flag() { + assert_eq!( + input(&named_like_flags(), &["--version", "3"]).unwrap(), + json!({"version": 3}) ); + } - let script = String::from_utf8(script).unwrap(); - for word in [ - "train", - "prompt", - "--optimizer.weight-decay", - "--config", - "--completions", - ] { - assert!(script.contains(word), "no {word} in the script"); - } + #[rstest] + #[case::when_asked(&["--completions", "zsh"], Some(Shell::Zsh))] + #[case::when_a_job_is_named(&["train"], None)] + fn completions_names_the_shell_the_command_line_asks_for( + #[case] args: &[&str], + #[case] shell: Option, + ) { + let jobs = [train()]; + + let matches = command("program", &jobs) + .try_get_matches_from(["program"].iter().chain(args)) + .unwrap(); + + assert_eq!(completions(&matches), shell); } } diff --git a/crates/tracel-job/src/cli/flags.rs b/crates/tracel-job/src/cli/flags.rs index e15a12c4..90a93693 100644 --- a/crates/tracel-job/src/cli/flags.rs +++ b/crates/tracel-job/src/cli/flags.rs @@ -1,17 +1,18 @@ //! The flags of a job's command line, read from its input schema: one per field of its input. +use std::error::Error; use std::ptr; -use referencing::{Registry, Resolver}; +use referencing::{Registry, Resolver, Retrieve, Uri}; use serde_json::{Map, Number, Value}; use crate::JobDefinition; /// Flag names the command line keeps: a field with one of these names has no flag. -const RESERVED: [&str; 3] = ["help", "version", "config"]; +const RESERVED: [&str; 2] = ["help", "config"]; /// The URI the input schema is registered under, which its references resolve against. -const BASE: &str = "urn:tracel:input"; +const BASE: &str = "json-schema:///input"; /// A flag that sets one field of a job's input. #[derive(Debug, Clone, PartialEq)] @@ -163,6 +164,7 @@ pub fn flags(job: &JobDefinition) -> Vec { return Vec::new(); }; let Ok(registry) = Registry::new() + .retriever(Unfetched) .add(BASE, schema) .and_then(|registry| registry.prepare()) else { @@ -188,6 +190,17 @@ pub fn flags(job: &JobDefinition) -> Vec { walk.flags } +/// Answers each schema the input schema refers to outside itself with `true`, the schema of any +/// value, since none is fetched: the registry is prepared whatever the input schema refers to, and +/// a reference into such a schema does not resolve. +struct Unfetched; + +impl Retrieve for Unfetched { + fn retrieve(&self, _: &Uri) -> Result> { + Ok(Value::Bool(true)) + } +} + /// Whether `key` can be part of a flag name. fn is_flag_key(key: &str) -> bool { !key.is_empty() @@ -306,9 +319,18 @@ impl<'r> Walk<'r> { } } + /// The description of `schema`, or else that of the schema it refers to, or else the help of + /// the value it makes optional. + fn help(&self, schema: &'r Value) -> Option { + schema.as_object().and_then(description).or_else(|| { + let target = self.target(schema)?.as_object()?; + description(target).or_else(|| self.help(self.nullable(target)?)) + }) + } + /// Adds the flags of the fields of `object`, the schema of the object at `path`, whose value /// in the default input is `defaults`. The object is required when `required` is. A field's - /// help is its description, or else that of the schema it refers to. + /// help is its [`help`](Self::help). fn object( &mut self, object: &'r Map, @@ -337,15 +359,10 @@ impl<'r> Walk<'r> { match self.shape(property) { Shape::Object(object) => self.object(object, path, default, required), Shape::Leaf(kind) => { - let help = [Some(property), self.target(property)] - .into_iter() - .flatten() - .filter_map(Value::as_object) - .find_map(description); self.flags.push(Flag { path: path.clone(), kind, - help, + help: self.help(property), default: default.cloned(), required: required && default.is_none(), }); @@ -391,17 +408,16 @@ fn description(schema: &Map) -> Option { } #[cfg(test)] -#[allow(dead_code)] mod tests { use std::collections::HashMap; + use rstest::{fixture, rstest}; use schemars::JsonSchema; use serde::Serialize; use serde_json::json; use super::*; - /// The flags of a job whose input schema is `schema` and default input `default`. fn flags_of(schema: Value, default: Option) -> Vec { flags(&JobDefinition { name: "train".to_string(), @@ -411,85 +427,65 @@ mod tests { }) } - /// The flags of a job taking `T`, with no default input. fn flags_for() -> Vec { flags_of(schemars::schema_for!(T).to_value(), None) } - /// The name and kind of each flag. - fn kinds(flags: &[Flag]) -> Vec<(String, Kind)> { - flags - .iter() - .map(|flag| (flag.name(), flag.kind.clone())) - .collect() - } - fn flag<'f>(flags: &'f [Flag], name: &str) -> &'f Flag { flags .iter() .find(|flag| flag.name() == name) - .unwrap_or_else(|| panic!("no flag --{name} in {:?}", kinds(flags))) + .unwrap_or_else(|| panic!("no flag --{name}")) + } + + fn kind_of() -> Kind { + flag(&flags_for::>(), "field").kind.clone() } - fn choices(values: &[(&str, Option<&str>)]) -> Kind { + fn one_of(values: &[Value]) -> Kind { Kind::Choice( values .iter() - .map(|(value, help)| Choice { - text: value.to_string(), - value: json!(value), - help: help.map(str::to_string), - }) + .map(|value| Choice::new(value, None).unwrap()) .collect(), ) } #[derive(JsonSchema)] - #[serde(rename_all = "camelCase")] - struct Shapes { - epochs: u32, - offset: i64, - lr: f64, - decay: f32, - shuffle: bool, - tag: String, - letter: char, - #[serde(rename = "seed_value")] - seed: Option, - label: Option, - optimizer: Optimizer, - resume_from: Option, - layers: Vec, - labels: HashMap, - pair: (u32, String), - anything: serde_json::Value, - nothing: (), - tagged: Tagged, - untagged: Untagged, - mixed: Mixed, - #[serde(flatten)] - flattened: Flattened, + #[expect(dead_code)] + struct Field { + field: T, } - #[derive(Serialize, JsonSchema)] - struct Optimizer { + #[derive(JsonSchema)] + #[expect(dead_code)] + struct Nesting { + optimizer: Nested, + } + + #[derive(JsonSchema)] + #[expect(dead_code)] + struct Nested { lr: f64, weight_decay: f64, schedule: Schedule, } - #[derive(Serialize, JsonSchema)] + #[derive(JsonSchema)] + #[expect(dead_code)] struct Schedule { warmup: u32, } #[derive(JsonSchema)] - struct Resume { + #[expect(dead_code)] + struct Checkpoint { experiment: u32, } #[derive(JsonSchema)] #[serde(tag = "kind")] + #[expect(dead_code)] enum Tagged { Adam { beta: f64 }, Sgd, @@ -497,122 +493,176 @@ mod tests { #[derive(JsonSchema)] #[serde(untagged)] + #[expect(dead_code)] enum Untagged { Count(u32), Name(String), } #[derive(JsonSchema)] + #[expect(dead_code)] enum Mixed { Unit, Data(u32), } - #[derive(JsonSchema)] - struct Flattened { - workers: u32, + #[rstest] + #[case::unsigned_integer(kind_of::(), Kind::Integer)] + #[case::signed_integer(kind_of::(), Kind::Integer)] + #[case::double(kind_of::(), Kind::Number)] + #[case::single(kind_of::(), Kind::Number)] + #[case::boolean(kind_of::(), Kind::Bool)] + #[case::string(kind_of::(), Kind::String)] + #[case::character(kind_of::(), Kind::String)] + fn a_scalar_field_takes_the_values_of_its_type(#[case] kind: Kind, #[case] expected: Kind) { + assert_eq!(kind, expected); + } + + #[rstest] + #[case::integer(kind_of::>(), kind_of::())] + #[case::string(kind_of::>(), kind_of::())] + #[case::unit_enum(kind_of::>(), kind_of::())] + fn an_optional_field_takes_the_values_of_the_type_it_makes_optional( + #[case] optional: Kind, + #[case] plain: Kind, + ) { + assert_eq!(optional, plain); + } + + #[rstest] + #[case::list(kind_of::>())] + #[case::map(kind_of::>())] + #[case::tuple(kind_of::<(u32, String)>())] + #[case::any_json(kind_of::())] + #[case::unit(kind_of::<()>())] + #[case::optional_object(kind_of::>())] + #[case::tagged_enum(kind_of::())] + #[case::untagged_enum(kind_of::())] + #[case::enum_with_data(kind_of::())] + fn a_field_of_any_other_type_takes_json(#[case] kind: Kind) { + assert_eq!(kind, Kind::Json); } #[test] - fn objects_nest_into_dotted_flags_and_other_values_take_json() { - let flags = flags_for::(); + fn a_nested_object_has_a_flag_per_field_named_by_its_keys() { + let mut names: Vec = flags_for::().iter().map(Flag::name).collect(); + names.sort(); assert_eq!( - kinds(&flags), + names, [ - ("anything".to_string(), Kind::Json), - ("decay".to_string(), Kind::Number), - ("epochs".to_string(), Kind::Integer), - ("label".to_string(), Kind::String), - ("labels".to_string(), Kind::Json), - ("layers".to_string(), Kind::Json), - ("letter".to_string(), Kind::String), - ("lr".to_string(), Kind::Number), - ("mixed".to_string(), Kind::Json), - ("nothing".to_string(), Kind::Json), - ("offset".to_string(), Kind::Integer), - ("optimizer.lr".to_string(), Kind::Number), - ("optimizer.schedule.warmup".to_string(), Kind::Integer), - ("optimizer.weight-decay".to_string(), Kind::Number), - ("pair".to_string(), Kind::Json), - ("resumeFrom".to_string(), Kind::Json), - ("seed-value".to_string(), Kind::Integer), - ("shuffle".to_string(), Kind::Bool), - ("tag".to_string(), Kind::String), - ("tagged".to_string(), Kind::Json), - ("untagged".to_string(), Kind::Json), - ("workers".to_string(), Kind::Integer), + "optimizer.lr", + "optimizer.schedule.warmup", + "optimizer.weight-decay" ] ); - let weight_decay = flag(&flags, "optimizer.weight-decay"); - assert_eq!(weight_decay.path, ["optimizer", "weight_decay"]); - assert_eq!(weight_decay.id(), "optimizer.weight_decay"); } - /// How to train. + #[test] + fn an_object_that_may_be_null_takes_json() { + let schema = json!({ + "type": "object", + "properties": { + "resume": { + "type": ["object", "null"], + "properties": {"epoch": {"type": "integer"}} + } + } + }); + + assert_eq!(flag(&flags_of(schema, None), "resume").kind, Kind::Json); + } + #[derive(JsonSchema)] + #[expect(dead_code)] struct Documented { - /// Passes over the data. + #[schemars(description = "Passes over the data.")] epochs: u32, - /// Where to train. - device: Device, backend: Backend, - /// The fallback. + #[schemars(description = "The backend to compare with.")] + baseline: Backend, fallback: Option, - optimizer: Optimizer, + seed: u64, } #[derive(JsonSchema)] - #[serde(rename_all = "snake_case")] + #[schemars(description = "The tensor backend.")] + #[expect(dead_code)] + enum Backend { + Wgpu, + Cuda, + } + + #[derive(JsonSchema)] + #[expect(dead_code)] enum Device { - /// The CPU. + #[schemars(description = "The CPU.")] Cpu, - /// The first GPU. + #[schemars(description = "The first GPU.")] Gpu, } - /// The tensor backend. - #[derive(JsonSchema)] - enum Backend { - Wgpu, - Cuda, + #[rstest] + #[case::its_own_description("epochs", Some("Passes over the data."))] + #[case::the_description_of_its_type("backend", Some("The tensor backend."))] + #[case::its_own_over_its_types("baseline", Some("The backend to compare with."))] + #[case::that_of_the_type_it_makes_optional("fallback", Some("The tensor backend."))] + #[case::none_when_nothing_describes_it("seed", None)] + fn a_flag_is_helped_by_its_fields_description_or_else_its_types( + #[case] name: &str, + #[case] help: Option<&str>, + ) { + assert_eq!(flag(&flags_for::(), name).help.as_deref(), help); + } + + #[rstest] + #[case::described( + kind_of::(), + &[("Cpu", Some("The CPU.")), ("Gpu", Some("The first GPU."))] + )] + #[case::undescribed(kind_of::(), &[("Wgpu", None), ("Cuda", None)])] + fn the_unit_variants_of_an_enum_are_choices_with_their_descriptions( + #[case] kind: Kind, + #[case] expected: &[(&str, Option<&str>)], + ) { + let Kind::Choice(choices) = &kind else { + panic!("{kind:?} is not a choice"); + }; + let choices: Vec<(&str, Option<&str>)> = choices + .iter() + .map(|choice| (choice.text.as_str(), choice.help.as_deref())) + .collect(); + + assert_eq!(choices, expected); + } + + #[rstest] + #[case::leaving_out_null(json!(["adam", null]), &[json!("adam")])] + #[case::keeping_numbers_as_numbers(json!([1, 2.5]), &[json!(1), json!(2.5)])] + fn listed_values_are_the_choices_of_their_flag( + #[case] listed: Value, + #[case] values: &[Value], + ) { + let schema = json!({"type": "object", "properties": {"field": {"enum": listed}}}); + + assert_eq!(flag(&flags_of(schema, None), "field").kind, one_of(values)); } #[test] - fn doc_comments_are_help_and_unit_variants_are_choices() { - let flags = flags_for::(); + fn listed_values_of_which_one_is_not_a_scalar_take_json() { + let schema = json!({"type": "object", "properties": {"field": {"enum": ["adam", [1]]}}}); - assert_eq!( - flag(&flags, "device").kind, - choices(&[("cpu", Some("The CPU.")), ("gpu", Some("The first GPU."))]) - ); - assert_eq!( - flag(&flags, "backend").kind, - choices(&[("Wgpu", None), ("Cuda", None)]) - ); - assert_eq!(flag(&flags, "fallback").kind, flag(&flags, "backend").kind); - for (name, help) in [ - ("epochs", Some("Passes over the data.")), - ("device", Some("Where to train.")), - ("backend", Some("The tensor backend.")), - ("fallback", Some("The fallback.")), - ("optimizer.lr", None), - ] { - assert_eq!(flag(&flags, name).help.as_deref(), help, "--{name}"); - } + assert_eq!(flag(&flags_of(schema, None), "field").kind, Kind::Json); } - #[derive(Serialize, JsonSchema)] + #[derive(JsonSchema)] + #[expect(dead_code)] struct Defaulted { epochs: u32, - #[serde(default)] - workers: u32, #[serde(default = "seed")] seed: u64, #[serde(default)] schedule: Scheduled, - optimizer: Optimizer, - tag: Option, } #[derive(Serialize, JsonSchema)] @@ -635,221 +685,265 @@ mod tests { 42 } - #[test] - fn defaults_come_from_the_default_input_or_else_the_schema() { - let schema = schemars::schema_for!(Defaulted).to_value(); - let default = |flags: &[Flag], name| flag(flags, name).default.clone(); - - let without = flags_of(schema.clone(), None); - let with = flags_of( - schema, - Some(json!({"epochs": 10, "seed": 7, "schedule": {"warmup": 1}, "optimizer": {}})), + #[rstest] + #[case::without_a_schema_default("epochs")] + #[case::over_the_schema_default("seed")] + fn a_flag_defaults_to_its_fields_value_in_the_default_input(#[case] name: &str) { + let default_input = json!({"epochs": 10, "seed": 7}); + + let flags = flags_of( + schemars::schema_for!(Defaulted).to_value(), + Some(default_input.clone()), ); - assert_eq!(default(&without, "epochs"), None); - assert_eq!(default(&without, "workers"), Some(json!(0))); - assert_eq!(default(&without, "seed"), Some(json!(42))); - assert_eq!(default(&without, "schedule.warmup"), Some(json!(5))); - assert_eq!(default(&without, "schedule.cooldown"), Some(json!(2))); - assert_eq!(default(&with, "epochs"), Some(json!(10))); - assert_eq!(default(&with, "seed"), Some(json!(7))); - assert_eq!(default(&with, "schedule.warmup"), Some(json!(1))); - assert_eq!(default(&with, "schedule.cooldown"), Some(json!(2))); + assert_eq!(flag(&flags, name).default.as_ref(), default_input.get(name)); + } + + #[rstest] + #[case::of_a_field("seed", None, json!(seed()))] + #[case::of_a_nested_field("schedule.warmup", None, json!(Scheduled::default().warmup))] + #[case::of_a_nested_field_the_default_input_leaves_out( + "schedule.cooldown", + Some(json!({"schedule": {"warmup": 1}})), + json!(Scheduled::default().cooldown) + )] + fn a_flag_the_default_input_does_not_give_defaults_to_the_schema_default( + #[case] name: &str, + #[case] default_input: Option, + #[case] default: Value, + ) { + let flags = flags_of(schemars::schema_for!(Defaulted).to_value(), default_input); + + assert_eq!(flag(&flags, name).default, Some(default)); + } + + #[fixture] + fn requirements() -> Value { + json!({ + "type": "object", + "required": ["epochs", "seed", "optimizer"], + "properties": { + "epochs": {"type": "integer"}, + "seed": {"type": "integer", "default": 42}, + "tag": {"type": "string"}, + "optimizer": { + "type": "object", + "required": ["lr"], + "properties": {"lr": {"type": "number"}} + }, + "resume": { + "type": "object", + "required": ["epoch"], + "properties": {"epoch": {"type": "integer"}} + } + } + }) + } + + #[rstest] + #[case::required("epochs", true)] + #[case::required_in_a_required_object("optimizer.lr", true)] + #[case::required_with_a_default("seed", false)] + #[case::optional("tag", false)] + #[case::required_in_an_optional_object("resume.epoch", false)] + fn a_flag_is_required_when_the_schema_requires_its_field_and_gives_no_default( + requirements: Value, + #[case] name: &str, + #[case] required: bool, + ) { + assert_eq!(flag(&flags_of(requirements, None), name).required, required); } - #[test] - fn a_field_the_schema_requires_is_required_unless_a_default_gives_it() { - let schema = schemars::schema_for!(Defaulted).to_value(); - let required = |flags: Vec| -> Vec { - flags - .iter() - .filter(|flag| flag.required) - .map(Flag::name) - .collect() - }; + #[rstest] + fn a_required_field_the_default_input_gives_is_not_a_required_flag(requirements: Value) { + let flags = flags_of(requirements, Some(json!({"epochs": 10}))); - assert_eq!( - required(flags_of(schema.clone(), None)), - [ - "epochs", - "optimizer.lr", - "optimizer.schedule.warmup", - "optimizer.weight-decay" - ] - ); - assert_eq!( - required(flags_of( - schema, - Some(json!({"epochs": 10, "optimizer": {"lr": 0.1, "schedule": {"warmup": 1}}})) - )), - ["optimizer.weight-decay"] - ); + assert!(!flag(&flags, "epochs").required); } - #[derive(JsonSchema)] - struct Node { - value: u32, - next: Option>, - children: Vec, + #[test] + fn a_job_without_an_input_schema_has_no_flags_whatever_its_default_input() { + let job = JobDefinition { + name: "train".to_string(), + description: None, + input_schema: None, + input_example: Some(json!({"epochs": 10})), + }; + + assert!(flags(&job).is_empty()); } - #[test] - fn an_input_that_is_not_an_object_with_fields_has_no_flags() { - assert!(flags_for::().is_empty()); - assert!(flags_for::<()>().is_empty()); - assert!(flags_for::>().is_empty()); - assert!(flags_for::().is_empty()); - assert!(flags_for::>().is_empty()); - assert!(flags_for::().is_empty()); - assert!(flags_of(json!({"type": "object", "properties": {}}), None).is_empty()); - assert!( - flags(&JobDefinition { - name: "train".to_string(), - description: None, - input_schema: None, - input_example: Some(json!({"epochs": 10})), - }) - .is_empty() - ); + #[rstest] + #[case::of_an_integer(schemars::schema_for!(u32).to_value())] + #[case::of_a_list(schemars::schema_for!(Vec).to_value())] + #[case::of_a_unit_enum(schemars::schema_for!(Backend).to_value())] + #[case::of_an_optional_object(schemars::schema_for!(Option).to_value())] + #[case::of_any_json(schemars::schema_for!(Value).to_value())] + #[case::of_an_object_without_fields(json!({"type": "object", "properties": {}}))] + fn an_input_that_is_not_an_object_with_fields_has_no_flags(#[case] schema: Value) { + assert!(flags_of(schema, None).is_empty()); } #[test] fn a_field_that_refers_back_to_an_object_it_is_in_takes_json() { - assert_eq!( - kinds(&flags_for::()), - [ - ("children".to_string(), Kind::Json), - ("next".to_string(), Kind::Json), - ("value".to_string(), Kind::Integer), - ] - ); - let flags = flags_of( - json!({ - "type": "object", - "properties": { - "node": {"$ref": "#"}, - "loop": {"$ref": "#/$defs/A"}, - "missing": {"$ref": "#/$defs/Missing"} - }, - "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A"}} - }), - None, - ); - assert_eq!( - kinds(&flags), - [ - ("loop".to_string(), Kind::Json), - ("missing".to_string(), Kind::Json), - ("node".to_string(), Kind::Json), - ] - ); - } + let schema = json!({"type": "object", "properties": {"node": {"$ref": "#"}}}); - #[derive(JsonSchema)] - struct Escaped { - odd: Odd, - slash: Slash, + assert_eq!(flag(&flags_of(schema, None), "node").kind, Kind::Json); } - #[derive(JsonSchema)] - #[schemars(rename = "Odd name%")] - struct Odd { - a: u32, - } + #[rstest] + #[case::missing_from_the_schema("#/$defs/Missing")] + #[case::that_comes_back_round("#/$defs/Loop")] + #[case::into_another_document("https://example.com/input.json#/$defs/Epochs")] + fn a_field_whose_reference_does_not_resolve_takes_json(#[case] reference: &str) { + let schema = json!({ + "type": "object", + "properties": {"field": {"$ref": reference}}, + "$defs": {"Loop": {"$ref": "#/$defs/Loop"}} + }); - #[derive(JsonSchema)] - #[schemars(rename = "a/b~c")] - struct Slash { - b: bool, + assert_eq!(flag(&flags_of(schema, None), "field").kind, Kind::Json); } - #[test] - fn references_resolve_as_schemars_escapes_them() { - let schema = schemars::schema_for!(Escaped).to_value(); - assert_eq!(schema["properties"]["odd"]["$ref"], "#/$defs/Odd%20name%25"); - assert_eq!(schema["properties"]["slash"]["$ref"], "#/$defs/a~1b~0c"); + #[rstest] + #[case::from_a_field(json!({ + "type": "object", + "properties": { + "epochs": {"type": "integer"}, + "remote": {"$ref": "https://example.com/input.json"} + } + }))] + #[case::from_a_definition_to_a_file_beside_it(json!({ + "type": "object", + "properties": {"epochs": {"type": "integer"}}, + "$defs": {"Shared": {"$ref": "shared.json"}} + }))] + #[case::to_its_dialect(json!({ + "$schema": "https://example.com/an-unknown-dialect", + "type": "object", + "properties": {"epochs": {"type": "integer"}} + }))] + fn a_schema_that_refers_outside_itself_keeps_its_flags(#[case] schema: Value) { + assert_eq!(flag(&flags_of(schema, None), "epochs").kind, Kind::Integer); + } + + #[rstest] + #[case::into_defs("$defs")] + #[case::into_definitions("definitions")] + fn a_field_takes_the_values_of_the_schema_its_reference_resolves_to(#[case] definitions: &str) { + let schema = json!({ + "type": "object", + "properties": {"epochs": {"$ref": format!("#/{definitions}/Epochs")}}, + definitions: {"Epochs": {"type": "integer"}} + }); + + assert_eq!(flag(&flags_of(schema, None), "epochs").kind, Kind::Integer); + } + + #[rstest] + #[case::reserved_for_help("help")] + #[case::reserved_for_the_config_file("config")] + #[case::with_a_dot("a.b")] + #[case::with_a_space("a b")] + #[case::starting_with_a_dash("-x")] + #[case::empty("")] + fn a_field_whose_key_cannot_name_a_flag_has_none(#[case] key: &str) { + let schema = json!({"type": "object", "properties": {key: {"type": "integer"}}}); + + assert!(flags_of(schema, None).is_empty()); + } + + #[rstest] + #[case::as_the_object("config", "path")] + #[case::as_the_field("optimizer", "config")] + fn a_reserved_name_is_reserved_only_as_a_whole_flag_name( + #[case] object: &str, + #[case] field: &str, + ) { + let schema = json!({ + "type": "object", + "properties": {object: {"type": "object", "properties": {field: {"type": "integer"}}}} + }); - assert_eq!( - kinds(&flags_of(schema, None)), - [ - ("odd.a".to_string(), Kind::Integer), - ("slash.b".to_string(), Kind::Bool), - ] - ); + let names: Vec = flags_of(schema, None).iter().map(Flag::name).collect(); + + assert_eq!(names, [format!("{object}.{field}")]); } #[test] - fn reserved_unusable_and_taken_names_have_no_flag() { - let integer = json!({"type": "integer"}); - let flags = flags_of( - json!({ - "type": "object", - "properties": { - "help": {"type": "boolean"}, - "version": {"type": "string"}, - "config": {"type": "object", "properties": {"path": {"type": "string"}}}, - "optimizer": {"type": "object", "properties": {"config": integer}}, - "a.b": integer, - "a b": integer, - "-x": integer, - "": integer, - "batch-size": integer, - "batch_size": integer - } - }), - None, - ); + fn a_flag_name_an_earlier_field_took_is_not_given_again() { + let schema = json!({ + "type": "object", + "properties": {"batch-size": {"type": "integer"}, "batch_size": {"type": "integer"}} + }); + + assert_eq!(flags_of(schema, None).len(), 1); + } + + #[rstest] + #[case::a_negative_integer(Kind::Integer, "-3", json!(-3))] + #[case::a_number_in_exponent_form(Kind::Number, "3e-4", json!(3e-4))] + #[case::a_whole_number_kept_whole(Kind::Number, "2", json!(2))] + #[case::a_number_as_rust_reads_it(Kind::Number, ".5", json!(0.5))] + #[case::a_bool(Kind::Bool, "false", json!(false))] + #[case::digits_as_a_string(Kind::String, "12", json!("12"))] + #[case::json(Kind::Json, "[1, 2]", json!([1, 2]))] + #[case::a_choice_as_its_value(one_of(&[json!("adam"), json!(2)]), "2", json!(2))] + fn an_argument_parses_as_the_value_its_kind_reads( + #[case] kind: Kind, + #[case] argument: &str, + #[case] value: Value, + ) { + assert_eq!(kind.parse(argument), Ok(value)); + } + + #[rstest] + #[case::a_fraction_as_an_integer(Kind::Integer, "1.5")] + #[case::a_word_as_an_integer(Kind::Integer, "ten")] + #[case::infinity_as_a_number(Kind::Number, "inf")] + #[case::yes_as_a_bool(Kind::Bool, "yes")] + #[case::json_cut_short(Kind::Json, "[1,")] + #[case::a_value_no_choice_gives(one_of(&[json!("adam")]), "sgd")] + fn an_argument_its_kind_does_not_read_is_rejected(#[case] kind: Kind, #[case] argument: &str) { + assert!(kind.parse(argument).is_err()); + } + + #[rstest] + #[case::a_number(Kind::Number, json!(0.001))] + #[case::a_whole_number(Kind::Number, json!(10))] + #[case::an_integer(Kind::Integer, json!(64))] + #[case::a_string(Kind::String, json!("baseline"))] + #[case::a_bool(Kind::Bool, json!(true))] + #[case::json(Kind::Json, json!({"a": [1]}))] + #[case::a_choice(one_of(&[json!("adam")]), json!("adam"))] + fn a_default_renders_as_an_argument_that_parses_back_to_it( + #[case] kind: Kind, + #[case] default: Value, + ) { + let argument = kind.render(&default).unwrap(); - assert_eq!( - kinds(&flags), - [ - ("batch-size".to_string(), Kind::Integer), - ("config.path".to_string(), Kind::String), - ("optimizer.config".to_string(), Kind::Integer), - ] - ); + assert_eq!(kind.parse(&argument), Ok(default)); } #[test] - fn arguments_parse_as_the_kind_of_their_field() { - assert_eq!(Kind::Integer.parse("-3"), Ok(json!(-3))); - assert!(Kind::Integer.parse("1.5").is_err()); - assert!(Kind::Integer.parse("ten").is_err()); - assert_eq!(Kind::Number.parse("3e-4"), Ok(json!(3e-4))); - assert_eq!(Kind::Number.parse("2"), Ok(json!(2))); - assert_eq!(Kind::Number.parse(".5"), Ok(json!(0.5))); - assert!(Kind::Number.parse("inf").is_err()); - assert_eq!(Kind::Bool.parse("false"), Ok(json!(false))); - assert!(Kind::Bool.parse("yes").is_err()); - assert_eq!(Kind::String.parse("12"), Ok(json!("12"))); - assert_eq!(Kind::Json.parse("[1, 2]"), Ok(json!([1, 2]))); - assert!(Kind::Json.parse("[1,").is_err()); - let choice = Kind::Choice(vec![ - Choice::new(&json!("adam"), None).unwrap(), - Choice::new(&json!(2), None).unwrap(), - ]); - assert_eq!(choice.parse("2"), Ok(json!(2))); - assert!(choice.parse("sgd").is_err()); - } + fn a_default_widened_from_an_f32_renders_as_the_f32() { + let single = 5e-5f32; - #[test] - fn defaults_render_as_arguments_their_flag_takes() { - assert_eq!( - Kind::Number.render(&json!(5e-5f32)).as_deref(), - Some("0.00005") - ); - assert_eq!(Kind::Number.render(&json!(0.001)).as_deref(), Some("0.001")); - assert_eq!(Kind::Number.render(&json!(10)).as_deref(), Some("10")); - assert_eq!(Kind::Integer.render(&json!(64)).as_deref(), Some("64")); assert_eq!( - Kind::Json.render(&json!({"a": [1]})).as_deref(), - Some(r#"{"a":[1]}"#) + Kind::Number.render(&json!(single)), + Some(single.to_string()) ); - assert_eq!(Kind::Integer.render(&json!(1.5)), None); - assert_eq!(Kind::String.render(&json!(3)), None); - assert_eq!(Kind::Number.render(&Value::Null), None); - let choice = Kind::Choice(vec![Choice::new(&json!("adam"), None).unwrap()]); - assert_eq!(choice.render(&json!("adam")).as_deref(), Some("adam")); - assert_eq!(choice.render(&json!("sgd")), None); + } + + #[rstest] + #[case::a_fraction_as_an_integer(Kind::Integer, json!(1.5))] + #[case::a_number_as_a_string(Kind::String, json!(3))] + #[case::null_as_a_number(Kind::Number, Value::Null)] + #[case::a_value_no_choice_gives(one_of(&[json!("adam")]), json!("sgd"))] + fn a_default_its_flag_does_not_take_renders_as_nothing( + #[case] kind: Kind, + #[case] default: Value, + ) { + assert_eq!(kind.render(&default), None); } } diff --git a/crates/tracel-job/src/cli/mod.rs b/crates/tracel-job/src/cli/mod.rs index 04e56780..853fda63 100644 --- a/crates/tracel-job/src/cli/mod.rs +++ b/crates/tracel-job/src/cli/mod.rs @@ -18,7 +18,8 @@ //! - An object's `properties` are its fields, and a nested object's fields are flags too: the //! field `weight_decay` of `optimizer` is `--optimizer.weight-decay`, the keys joined with dots //! and `_` written as `-`. -//! - `$ref` is followed into the schema's `$defs` or `definitions`. +//! - `$ref` is followed within the schema, such as into its `$defs` or `definitions`; a field +//! whose reference leads out of the schema or nowhere takes a JSON literal. //! - `enum`, `const`, and `oneOf` or `anyOf` of constants, such as the unit variants of an enum, //! are the values a flag takes, each with its description. //! - An optional field, `type: [T, "null"]` or `anyOf` of `T` and `null`, takes the values of `T`. @@ -28,13 +29,13 @@ //! A field's `description` is its flag's help, and its value in the job's default input, or else //! its schema `default`, is the flag's shown default. A field the schema requires is a required //! flag unless it has a default or the JSON document or `--config` is given. A boolean flag given -//! no value is `true`. The schema's other constraints, such as `minimum` or `pattern`, are not -//! checked here: the job's input is checked as a whole once it is resolved. +//! no value is `true`, and takes a value only as `--=`, so an argument after it is the +//! JSON document. The schema's other constraints, such as `minimum` or `pattern`, are not checked +//! here: the job's input is checked as a whole once it is resolved. //! -//! `--help`, `--version` and `--config` belong to the command line, so a field named `help`, -//! `version` or `config` has no flag; neither does a field whose key is not made of letters, -//! digits, `_` and `-`, or whose flag name an earlier field took. Such a field is set through the -//! JSON document or `--config`. +//! `--help` and `--config` belong to the command line, so a field named `help` or `config` has no +//! flag; neither does a field whose key is not made of letters, digits, `_` and `-`, or whose flag +//! name an earlier field took. Such a field is set through the JSON document or `--config`. //! //! ## Input //! diff --git a/crates/tracel-job/src/definitions.rs b/crates/tracel-job/src/definitions.rs index f1887dca..3eaeab47 100644 --- a/crates/tracel-job/src/definitions.rs +++ b/crates/tracel-job/src/definitions.rs @@ -20,7 +20,7 @@ pub struct JobDefinition { pub description: Option, /// JSON Schema of the job's input, when the program gives one. pub input_schema: Option, - /// An example input, typically the job's default input, which its input is merged onto. + /// The job's default input, which its input is merged onto. pub input_example: Option, } @@ -36,12 +36,13 @@ pub struct DefinitionsFile { } impl DefinitionsFile { - /// The file of a program that can run `jobs`: following [`PROTOCOL`], and built with this - /// version of the SDK. - pub fn new(jobs: Vec) -> Self { + /// The file of a program built with version `sdk_version` of the Tracel SDK that can run + /// `jobs`, following [`PROTOCOL`], with `jobs` ordered by name. + pub fn new(sdk_version: impl Into, mut jobs: Vec) -> Self { + jobs.sort_by(|a, b| a.name.cmp(&b.name)); Self { protocol: PROTOCOL, - sdk_version: env!("CARGO_PKG_VERSION").to_string(), + sdk_version: sdk_version.into(), jobs, } } @@ -112,86 +113,28 @@ mod tests { use super::*; - fn file() -> DefinitionsFile { - DefinitionsFile::new(vec![ - JobDefinition { - name: "echo".to_string(), - description: None, - input_schema: Some(json!({"type": "string"})), - input_example: None, - }, - JobDefinition { - name: "train".to_string(), - description: Some("Train the model".to_string()), - input_schema: None, - input_example: Some(json!({"epochs": 10, "optimizer": {"lr": 0.001}})), - }, - ]) - } - #[test] - fn the_file_is_pretty_printed_json_ending_with_a_newline() { - let expected = r#"{ - "protocol": 1, - "sdk_version": "SDK_VERSION", - "jobs": [ - { - "name": "echo", - "description": null, - "input_schema": { - "type": "string" - }, - "input_example": null - }, - { - "name": "train", - "description": "Train the model", - "input_schema": null, - "input_example": { - "epochs": 10, - "optimizer": { - "lr": 0.001 - } - } - } - ] -} -"# - .replace("SDK_VERSION", env!("CARGO_PKG_VERSION")); - - assert_eq!(file().to_json(), expected); - assert_eq!( - DefinitionsFile::from_json(expected.as_bytes()).unwrap(), - file() - ); + fn the_jobs_are_ordered_by_name() { + let job = |name: &str| JobDefinition { + name: name.to_string(), + description: None, + input_schema: None, + input_example: None, + }; + + let file = + DefinitionsFile::new("0.10.0", vec![job("train"), job("evaluate"), job("infer")]); + + assert_eq!(file.jobs, [job("evaluate"), job("infer"), job("train")]); } #[test] - fn only_a_definitions_file_of_this_protocol_parses() { - let parse = |json: &str| DefinitionsFile::from_json(json.as_bytes()); - - assert!(matches!( - parse(r#"{"protocol": 1}"#), - Err(ParseError::Invalid(_)) - )); - assert!(matches!( - parse(r#"{"jobs": []}"#), - Err(ParseError::Invalid(_)) - )); - assert!(matches!( - parse(r#"{"protocol": 1, "jo"#), - Err(ParseError::Invalid(_)) - )); - assert!(matches!( - parse(r#"{"protocol": 2, "jobs": "elsewhere"}"#), - Err(ParseError::Protocol { found: 2 }) - )); - } + fn a_file_of_another_protocol_is_refused_whatever_its_other_fields() { + let other = PROTOCOL + 1; + let document = json!({"protocol": other, "jobs": "elsewhere"}); - #[test] - fn a_protocol_error_names_both_versions() { - let error = DefinitionsFile::from_json(br#"{"protocol": 7}"#).unwrap_err(); + let parsed = DefinitionsFile::from_json(document.to_string().as_bytes()); - assert_eq!(error.to_string(), "protocol 7 is not supported, expected 1"); + assert!(matches!(parsed, Err(ParseError::Protocol { found }) if found == other)); } } diff --git a/crates/tracel-job/src/run.rs b/crates/tracel-job/src/run.rs index 5fe4859a..cac2d1bf 100644 --- a/crates/tracel-job/src/run.rs +++ b/crates/tracel-job/src/run.rs @@ -42,20 +42,21 @@ impl From for std::process::ExitCode { #[cfg(test)] mod tests { + use rstest::rstest; + use super::*; - #[test] - fn each_exit_code_reads_back_from_its_code() { - for (exit, code) in [ - (ExitCode::Completed, 0), - (ExitCode::Failed, 1), - (ExitCode::Rejected, 2), - ] { - assert_eq!(exit.code(), code); - assert_eq!(ExitCode::from_code(i32::from(code)), Some(exit)); - } - for code in [-1, 3, 101, 130, 137, 256] { - assert_eq!(ExitCode::from_code(code), None); - } + #[rstest] + fn an_exit_code_reads_back_from_its_code( + #[values(ExitCode::Completed, ExitCode::Failed, ExitCode::Rejected)] exit: ExitCode, + ) { + assert_eq!(ExitCode::from_code(exit.code().into()), Some(exit)); + } + + #[rstest] + fn a_code_the_protocol_does_not_give_is_no_exit_code( + #[values(-1, 3, 101, 130, 137, 256)] code: i32, + ) { + assert_eq!(ExitCode::from_code(code), None); } } diff --git a/crates/tracel-models/Cargo.toml b/crates/tracel-models/Cargo.toml index 0eb5d36d..e3ae71c0 100644 --- a/crates/tracel-models/Cargo.toml +++ b/crates/tracel-models/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tracel-models" -description = "Backend-agnostic model domain and capability for Tracel." +description = "Backend-agnostic model domain and registry for Tracel." edition.workspace = true version.workspace = true readme.workspace = true diff --git a/crates/tracel-runtime/Cargo.toml b/crates/tracel-runtime/Cargo.toml index 7f8561a3..1fde7524 100644 --- a/crates/tracel-runtime/Cargo.toml +++ b/crates/tracel-runtime/Cargo.toml @@ -31,6 +31,7 @@ url = { workspace = true, optional = true } uuid = { workspace = true, optional = true } [dev-dependencies] +rstest = { workspace = true } serde = { workspace = true, features = ["std"] } tempfile.workspace = true uuid = { workspace = true, features = ["v4"] } diff --git a/crates/tracel-runtime/src/execution.rs b/crates/tracel-runtime/src/execution.rs index 0239a6b1..3fd51f0e 100644 --- a/crates/tracel-runtime/src/execution.rs +++ b/crates/tracel-runtime/src/execution.rs @@ -106,155 +106,77 @@ impl Outcome { #[cfg(test)] mod tests { - use std::sync::{Arc, Barrier, Mutex}; + use std::sync::{Arc, OnceLock}; + use rstest::{fixture, rstest}; use serde_json::json; use super::*; - use crate::{Job, Jobs, PrepareError, job_input, job_num, job_running}; + use crate::{Jobs, job_input, job_num}; - /// `train`, which takes a number of epochs, three by default, and fails on zero; and `panic`, - /// which panics. + #[fixture] fn jobs() -> Jobs { Jobs::new() - .job( - "train", - Job::new(|epochs: u32| match epochs { - 0 => Err::<(), BoxError>("no epochs".into()), - _ => Ok(()), - }) - .with_default(3), - ) - .job("panic", |()| -> Result<(), BoxError> { panic!("kaboom") }) - } - - #[test] - fn a_job_is_prepared_with_its_input_merged_onto_its_default() { - let jobs = Jobs::new().job( - "train", - Job::new(|_: Value| Ok::<_, BoxError>(())) - .with_default(json!({"epochs": 3, "lr": 0.1})), - ); - - let prepared = jobs.prepare("train", json!({"epochs": 5})).unwrap(); - - assert_eq!(prepared.name(), "train"); - assert_eq!(prepared.input(), &json!({"epochs": 5, "lr": 0.1})); + .job("completes", |()| Ok::<_, BoxError>(())) + .job("fails", |()| { + Err::<(), BoxError>("no epochs to train".into()) + }) + .job("panics", |()| -> Result<(), BoxError> { + panic!("the kernel exploded") + }) } - #[test] - fn an_unknown_job_or_an_input_that_does_not_decode_is_not_prepared() { - let jobs = jobs(); - - let unknown = jobs.prepare("infer", Value::Null).err().unwrap(); - let invalid = jobs.prepare("train", json!("ten")).err().unwrap(); - - assert!( - matches!( - &unknown, - PrepareError::UnknownJob { name, available } - if name == "infer" && available == &["panic", "train"] - ), - "{unknown:?}" - ); - assert_eq!(unknown.to_string(), "unknown job 'infer'"); - assert!( - matches!(invalid, PrepareError::InvalidInput(_)), - "{invalid:?}" - ); - assert!( - invalid.to_string().starts_with("invalid input: "), - "{invalid}" - ); + fn reading(read: fn() -> T) -> (Jobs, Arc>) + where + T: Send + Sync + 'static, + { + let seen = Arc::new(OnceLock::new()); + let jobs = Jobs::new().job("train", { + let seen = seen.clone(); + move |_: Value| { + let _ = seen.set(read()); + Ok::<_, BoxError>(()) + } + }); + (jobs, seen) } - #[test] - fn a_job_completes_or_fails_by_what_it_returns_and_a_panic_fails_it() { - let run = - |name: &str, input: Value| jobs().prepare(name, input).unwrap().run(Execution::new()); - - let completed = run("train", Value::Null); - let failed = run("train", json!(0)); - let panicked = run("panic", Value::Null); + #[rstest] + #[case::completes("completes", JobStatus::Completed)] + #[case::fails("fails", JobStatus::Failed)] + #[case::panics("panics", JobStatus::Failed)] + fn a_job_ends_as_its_function_does(jobs: Jobs, #[case] name: &str, #[case] status: JobStatus) { + let outcome = jobs + .prepare(name, Value::Null) + .unwrap() + .run(Execution::new()); - assert_eq!(completed.status(), JobStatus::Completed); - assert!(completed.error().is_none()); - assert!(completed.into_result().is_ok()); - assert_eq!(failed.status(), JobStatus::Failed); - assert_eq!(failed.error().unwrap().to_string(), "no epochs"); - assert_eq!(panicked.status(), JobStatus::Failed); - assert_eq!( - panicked.into_result().unwrap_err().to_string(), - "the job panicked: kaboom" - ); + assert_eq!(outcome.status(), status); } - #[test] - fn the_job_reads_its_number_and_input_from_its_scope_while_it_runs() { - let seen = Arc::new(Mutex::new(Vec::new())); - let jobs = Jobs::new().job( - "train", - Job::new({ - let seen = seen.clone(); - move |_: Value| { - seen.lock() - .unwrap() - .push((job_running(), job_num(), job_input())); - Ok::<_, BoxError>(()) - } - }) - .with_default(json!({"epochs": 3})), - ); + #[rstest] + #[case::numbered(Execution::new().with_num("42"), Some("42"))] + #[case::unnumbered(Execution::new(), None)] + fn the_job_reads_the_number_its_execution_gives( + #[case] execution: Execution, + #[case] num: Option<&str>, + ) { + let (jobs, seen) = reading(job_num); - let numbered = jobs.prepare("train", json!({"lr": 0.1})).unwrap(); - numbered.run(Execution::new().with_num("42")); - jobs.prepare("train", Value::Null) - .unwrap() - .run(Execution::new()); + jobs.prepare("train", Value::Null).unwrap().run(execution); - assert_eq!( - *seen.lock().unwrap(), - [ - ( - true, - Some("42".to_string()), - Some(json!({"epochs": 3, "lr": 0.1})) - ), - (true, None, Some(json!({"epochs": 3}))), - ] - ); - assert!(!job_running()); + assert_eq!(seen.get(), Some(&num.map(str::to_string))); } #[test] - fn runs_at_once_on_two_threads_each_see_their_own_execution() { - let both = Arc::new(Barrier::new(2)); - let jobs = Arc::new(Jobs::new().job("train", { - let both = both.clone(); - move |epochs: u32| { - both.wait(); - match (job_num(), job_input()) { - (Some(num), Some(input)) - if num == epochs.to_string() && input == json!(epochs) => - { - Ok(()) - } - seen => Err::<(), BoxError>(format!("job {epochs} saw {seen:?}").into()), - } - } - })); - let run = |epochs: u32| { - let prepared = jobs.prepare("train", json!(epochs)).unwrap(); - let execution = Execution::new().with_num(epochs.to_string()); - std::thread::spawn(move || prepared.run(execution)) - }; + fn the_job_reads_the_input_it_runs_with() { + let (jobs, seen) = reading(job_input); + let prepared = jobs.prepare("train", json!({"lr": 0.1})).unwrap(); + let input = prepared.input().clone(); - let one = run(1); - let two = run(2); + prepared.run(Execution::new()); - let (one, two) = (one.join().unwrap(), two.join().unwrap()); - assert!(one.error().is_none(), "{one:?}"); - assert!(two.error().is_none(), "{two:?}"); + assert_eq!(seen.get(), Some(&Some(input))); } #[test] diff --git a/crates/tracel-runtime/src/job.rs b/crates/tracel-runtime/src/job.rs index 9c04732a..9b3966ee 100644 --- a/crates/tracel-runtime/src/job.rs +++ b/crates/tracel-runtime/src/job.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; +use tracel_job::JobDefinition; /// The error a job fails with. pub type BoxError = Box; @@ -77,21 +78,6 @@ where } } -/// What a transport reads about a registered job: its name and description, and the input it -/// takes. -#[derive(Debug, Clone, PartialEq)] -pub struct JobDefinition { - /// The name that selects the job. - pub name: String, - /// What the job does, as [`Job::with_description`] gives it. - pub description: Option, - /// The JSON Schema of the job's input, when it was given one (requires the `schema` feature). - pub input_schema: Option, - /// The job's default input, which its input is merged onto, as [`Job::with_default`] gives - /// it. - pub input_example: Option, -} - /// What [`Jobs::job`](crate::Jobs::job) registers: a [`Job`], or a function of the job's input. pub trait IntoJob { /// The job. @@ -195,8 +181,10 @@ fn resolve(default: Option, input: Value) -> Value { #[cfg(test)] mod tests { + use std::io; use std::sync::Mutex; + use rstest::rstest; use serde::Deserialize; use serde_json::json; @@ -228,7 +216,6 @@ mod tests { } } - /// A job that keeps the input it runs with. fn recording() -> (Job, Arc>>) where I: DeserializeOwned + Send + 'static, @@ -244,7 +231,6 @@ mod tests { (job, ran_with) } - /// Runs `job`, registered as `train`, with `input`. fn run(job: Job, input: Value) -> Result<(), BoxError> where I: DeserializeOwned + Send + 'static, @@ -256,55 +242,21 @@ mod tests { #[test] fn a_job_is_defined_by_its_name_description_and_default() { + let (name, description) = ("train", "Train the model"); let (job, _) = recording::(); let job = job - .with_description("Train the model") + .with_description(description) .with_default(default_config()); - let (bare, _) = recording::(); assert_eq!( - Registered::new("train".to_string(), job).definition(), + Registered::new(name.to_string(), job).definition(), &JobDefinition { - name: "train".to_string(), - description: Some("Train the model".to_string()), + name: name.to_string(), + description: Some(description.to_string()), input_schema: None, - input_example: Some(json!({ - "epochs": 10, - "optimizer": {"lr": 0.001, "weight_decay": 0.0}, - "tag": "baseline" - })), + input_example: Some(serde_json::to_value(default_config()).unwrap()), } ); - assert_eq!( - Registered::new("evaluate".to_string(), bare).definition(), - &JobDefinition { - name: "evaluate".to_string(), - description: None, - input_schema: None, - input_example: None, - } - ); - } - - #[test] - fn without_a_default_the_input_decodes_as_given() { - let (job, ran_with) = recording::(); - let (no_input, _) = recording::(); - let (unit, ran_with_unit) = recording::<()>(); - - run(job, json!({"lr": 0.1, "weight_decay": 0.01})).unwrap(); - let missing = run(no_input, Value::Null); - run(unit, Value::Null).unwrap(); - - assert_eq!( - ran_with.lock().unwrap().take(), - Some(Optimizer { - lr: 0.1, - weight_decay: 0.01 - }) - ); - assert!(missing.is_err()); - assert_eq!(ran_with_unit.lock().unwrap().take(), Some(())); } #[test] @@ -316,38 +268,6 @@ mod tests { assert_eq!(ran_with.lock().unwrap().take(), Some(default_config())); } - #[test] - fn the_prepared_input_is_the_json_the_job_runs_with() { - let prepare = |input| { - let (job, _) = recording::(); - Registered::new("train".to_string(), job.with_default(default_config())) - .prepare(input) - .unwrap() - .input - }; - let (bare, _) = recording::(); - - assert_eq!( - prepare(json!({"epochs": 2, "optimizer": {"lr": 0.1}})), - json!({ - "epochs": 2, - "optimizer": {"lr": 0.1, "weight_decay": 0.0}, - "tag": "baseline" - }) - ); - assert_eq!( - prepare(Value::Null), - serde_json::to_value(default_config()).unwrap() - ); - assert_eq!( - Registered::new("evaluate".to_string(), bare) - .prepare(json!({"split": "test"})) - .unwrap() - .input, - json!({"split": "test"}) - ); - } - #[test] fn an_input_is_merge_patched_onto_the_default() { let (job, ran_with) = recording(); @@ -364,44 +284,46 @@ mod tests { epochs: 2, optimizer: Optimizer { lr: 0.1, - weight_decay: 0.0, + ..default_config().optimizer }, tag: None, }) ); } - #[test] - fn an_input_of_the_wrong_type_does_not_decode() { - for input in [json!({"epochs": "ten"}), json!([1, 2]), json!("ten")] { - let (job, ran_with) = recording(); + #[rstest] + #[case::a_field_of_the_wrong_type(json!({"epochs": "ten"}))] + #[case::not_an_object(json!("ten"))] + fn an_input_that_does_not_decode_is_refused(#[case] input: Value) { + let (job, _) = recording::(); - let prepared = Registered::new("train".to_string(), job.with_default(default_config())) - .prepare(input.clone()); + let prepared = + Registered::new("train".to_string(), job.with_default(default_config())).prepare(input); - assert!(prepared.is_err(), "{input}"); - assert!(ran_with.lock().unwrap().is_none()); - } + assert!(prepared.is_err()); } #[test] fn a_registered_job_runs_each_time_it_is_prepared() { let (job, ran_with) = recording::(); let registered = Registered::new("train".to_string(), job); + let epochs = [1, 2]; - for epochs in [1, 2] { + let ran = epochs.map(|epochs| { (registered.prepare(json!(epochs)).unwrap().run)().unwrap(); - assert_eq!(ran_with.lock().unwrap().take(), Some(epochs)); - } + ran_with.lock().unwrap().take() + }); + + assert_eq!(ran, epochs.map(Some)); } #[test] fn a_job_fails_with_the_error_it_returns() { - let job = Job::new(|_: ()| Err::<(), _>(std::io::Error::other("disk full"))); + let job = Job::new(|_: ()| Err::<(), _>(io::Error::other("disk full"))); - let failed = run(job, Value::Null); + let failed = run(job, Value::Null).unwrap_err(); - assert_eq!(failed.unwrap_err().to_string(), "disk full"); + assert!(failed.is::()); } #[cfg(feature = "schema")] @@ -411,9 +333,9 @@ mod tests { let registered = Registered::new("train".to_string(), job.with_schema()); - let schema = registered.definition().input_schema.clone().unwrap(); - assert_eq!(schema["title"], "Config"); - assert!(schema["properties"]["epochs"].is_object()); - assert!(schema["properties"]["optimizer"].is_object()); + assert_eq!( + registered.definition().input_schema, + Some(schemars::schema_for!(Config).to_value()) + ); } } diff --git a/crates/tracel-runtime/src/jobs.rs b/crates/tracel-runtime/src/jobs.rs index 092622d2..62a4acc9 100644 --- a/crates/tracel-runtime/src/jobs.rs +++ b/crates/tracel-runtime/src/jobs.rs @@ -4,9 +4,10 @@ use std::collections::BTreeMap; use serde::de::DeserializeOwned; use serde_json::Value; +use tracel_job::JobDefinition; use crate::execution::{self, PreparedJob}; -use crate::job::{IntoJob, JobDefinition, Prepared, Registered}; +use crate::job::{IntoJob, Prepared, Registered}; /// A program's jobs, by name, and how the program runs the one its command line names: `main` /// returns [`run`](Self::run). @@ -174,3 +175,57 @@ pub enum PrepareError { #[error("invalid input: {0}")] InvalidInput(#[source] serde_json::Error), } + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + use crate::BoxError; + + const TRAIN: &str = "train"; + const EVALUATE: &str = "evaluate"; + + #[fixture] + fn jobs() -> Jobs { + Jobs::new() + .job(TRAIN, |_: u32| Ok::<_, BoxError>(())) + .job(EVALUATE, |()| Ok::<_, BoxError>(())) + } + + #[rstest] + fn a_prepared_job_has_its_name_and_the_input_it_runs_with(jobs: Jobs) { + let input = json!(3); + + let prepared = jobs.prepare(TRAIN, input.clone()).unwrap(); + + assert_eq!((prepared.name(), prepared.input()), (TRAIN, &input)); + } + + #[rstest] + fn an_unknown_job_is_not_prepared_and_the_registered_names_are_given_in_order(jobs: Jobs) { + let unknown = "infer"; + + let refused = jobs.prepare(unknown, Value::Null); + + assert!(matches!( + refused, + Err(PrepareError::UnknownJob { name, available }) + if name == unknown && available == [EVALUATE, TRAIN] + )); + } + + #[rstest] + fn an_input_that_does_not_decode_is_not_prepared(jobs: Jobs) { + let refused = jobs.prepare(TRAIN, json!("ten")); + + assert!(matches!(refused, Err(PrepareError::InvalidInput(_)))); + } + + #[rstest] + #[should_panic] + fn registering_a_job_name_twice_panics(jobs: Jobs) { + let _ = jobs.job(TRAIN, |()| Ok::<_, BoxError>(())); + } +} diff --git a/crates/tracel-runtime/src/lib.rs b/crates/tracel-runtime/src/lib.rs index b7a607ea..de638226 100644 --- a/crates/tracel-runtime/src/lib.rs +++ b/crates/tracel-runtime/src/lib.rs @@ -1,3 +1,5 @@ +#![deny(missing_docs)] + //! A program's jobs, and the adapters that run them. //! //! A job is a function of its typed input, `Fn(I) -> Result<(), E>`, registered by name in @@ -71,9 +73,8 @@ //! function, while it runs, so jobs a Station runner runs at once each see their own. The //! functions read it: [`job_running`] says whether there is one, [`job_num`] gives its number, and //! [`job_input`] its input, as JSON. To read it on another thread, hold [`JobScope::current`]. The -//! `tracel` crate links the experiments a job starts to it through its scope. While a thread runs -//! no job, such as in a plain binary or a unit test, there is none: [`job_running`] gives `false`, -//! and [`job_num`] and [`job_input`] give `None`. +//! `tracel` crate links the experiments a job starts to it through its scope. Outside a job there +//! is none: [`job_running`] gives `false`, and [`job_num`] and [`job_input`] give `None`. mod execution; mod job; @@ -85,6 +86,7 @@ mod scope; pub mod station; pub use execution::{Execution, JobStatus, Outcome, PreparedJob}; -pub use job::{BoxError, IntoJob, Job, JobDefinition}; +pub use job::{BoxError, IntoJob, Job}; pub use jobs::{Jobs, PrepareError}; pub use scope::{JobScope, job_input, job_num, job_running}; +pub use tracel_job::JobDefinition; diff --git a/crates/tracel-runtime/src/panics.rs b/crates/tracel-runtime/src/panics.rs index be037f88..04531454 100644 --- a/crates/tracel-runtime/src/panics.rs +++ b/crates/tracel-runtime/src/panics.rs @@ -24,21 +24,30 @@ fn panic_message(panic: &(dyn Any + Send)) -> &str { #[cfg(test)] mod tests { + use std::io; + + use rstest::rstest; + use super::*; + const MESSAGE: &str = "the kernel exploded"; + #[test] - fn a_panic_is_the_failure_the_job_ends_with() { - let failure = |run: fn() -> Result<(), BoxError>| catch_panic(run).unwrap_err().to_string(); - - assert_eq!( - failure(|| panic!("kernel exploded")), - "the job panicked: kernel exploded" - ); - assert_eq!( - failure(|| panic!("{} exploded", "kernel")), - "the job panicked: kernel exploded" - ); - assert_eq!(failure(|| Err("boom".into())), "boom"); - assert!(catch_panic(|| Ok(())).is_ok()); + fn a_panic_fails_the_run() { + assert!(catch_panic(|| panic!("{MESSAGE}")).is_err()); + } + + #[test] + fn an_error_the_run_returns_is_its_failure() { + let failure = catch_panic(|| Err(io::Error::other(MESSAGE).into())).unwrap_err(); + + assert!(failure.is::()); + } + + #[rstest] + #[case::str(Box::new(MESSAGE))] + #[case::string(Box::new(MESSAGE.to_string()))] + fn the_panic_message_is_the_text_the_panic_carries(#[case] payload: Box) { + assert_eq!(panic_message(payload.as_ref()), MESSAGE); } } diff --git a/crates/tracel-runtime/src/process/describe.rs b/crates/tracel-runtime/src/process/describe.rs index d3f8dbea..4809e90f 100644 --- a/crates/tracel-runtime/src/process/describe.rs +++ b/crates/tracel-runtime/src/process/describe.rs @@ -1,30 +1,26 @@ //! Describing the registered jobs instead of running one. -use std::ffi::OsString; use std::io; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; -use tracel_job::{DefinitionsFile, TRACEL_DESCRIBE}; +use tracel_job::DefinitionsFile; -use super::protocol_definition; use crate::JobDefinition; use crate::process::file::write_atomically; -/// Writes the definitions file of the jobs `definitions` gives when `TRACEL_DESCRIBE`, among the -/// environment variables `lookup` gives, names a path. -/// -/// Returns whether it did; the program then runs no job. +/// Writes the definitions file of the jobs `definitions` gives to `path`. pub fn describe<'a>( definitions: impl IntoIterator, - lookup: impl Fn(&str) -> Option, -) -> Result { - let Some(path) = lookup(TRACEL_DESCRIBE).filter(|path| !path.is_empty()) else { - return Ok(false); - }; - let path = PathBuf::from(path); - let file = DefinitionsFile::new(definitions.into_iter().map(protocol_definition).collect()); - write_atomically(&path, &file.to_json()).map_err(|source| DescribeError { path, source })?; - Ok(true) + path: &Path, +) -> Result<(), DescribeError> { + let file = DefinitionsFile::new( + env!("CARGO_PKG_VERSION"), + definitions.into_iter().cloned().collect(), + ); + write_atomically(path, &file.to_json()).map_err(|source| DescribeError { + path: path.to_path_buf(), + source, + }) } /// The definitions file could not be written. @@ -37,9 +33,7 @@ pub struct DescribeError { #[cfg(test)] mod tests { - use std::path::Path; - - use serde_json::{Value, json}; + use serde_json::json; use super::*; @@ -60,66 +54,13 @@ mod tests { ] } - fn describe_to(path: &Path) -> Result { - describe(&definitions(), |name| { - (name == TRACEL_DESCRIBE).then(|| path.as_os_str().to_owned()) - }) - } - - #[test] - fn the_definitions_file_lists_every_job() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("jobs.json"); - - assert!(describe_to(&path).unwrap()); - - let contents = std::fs::read_to_string(&path).unwrap(); - let file: Value = serde_json::from_str(&contents).unwrap(); - assert_eq!( - file, - json!({ - "protocol": 1, - "sdk_version": env!("CARGO_PKG_VERSION"), - "jobs": [ - { - "name": "evaluate", - "description": null, - "input_schema": null, - "input_example": null - }, - { - "name": "train", - "description": "Train the model", - "input_schema": null, - "input_example": {"epochs": 10, "optimizer": {"lr": 0.001}} - } - ] - }) - ); - let read = DefinitionsFile::from_json(contents.as_bytes()).unwrap(); - assert_eq!( - read.jobs, - definitions() - .iter() - .map(protocol_definition) - .collect::>() - ); - assert_eq!(contents, read.to_json()); - } - #[test] fn a_failed_write_names_the_path() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("missing").join("jobs.json"); - let error = describe_to(&path).unwrap_err(); + let error = describe(&definitions(), &path).unwrap_err(); - assert!(error.to_string().contains("jobs.json"), "{error}"); - } - - #[test] - fn an_unset_or_empty_tracel_describe_writes_nothing() { - assert!(!describe(&definitions(), |_| None).unwrap()); - assert!(!describe(&definitions(), |_| Some(OsString::new())).unwrap()); + assert_eq!(error.path, path); } } diff --git a/crates/tracel-runtime/src/process/file.rs b/crates/tracel-runtime/src/process/file.rs index a56c6856..1767b8ed 100644 --- a/crates/tracel-runtime/src/process/file.rs +++ b/crates/tracel-runtime/src/process/file.rs @@ -19,37 +19,59 @@ pub fn write_atomically(path: &Path, contents: &str) -> io::Result<()> { #[cfg(test)] mod tests { + use std::ffi::OsString; + use super::*; + const CONTENTS: &str = "{}\n"; + + fn entries(dir: &Path) -> Vec { + fs::read_dir(dir) + .unwrap() + .map(|entry| entry.unwrap().file_name()) + .collect() + } + #[test] - fn the_file_holds_the_contents_and_no_temporary_file_is_left() { + fn the_file_holds_the_contents_written() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("jobs.json"); - write_atomically(&path, "{}\n").unwrap(); + write_atomically(&path, CONTENTS).unwrap(); - assert_eq!(fs::read_to_string(&path).unwrap(), "{}\n"); - assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1); + assert_eq!(fs::read_to_string(&path).unwrap(), CONTENTS); } #[test] - fn the_file_replaces_an_existing_one_through_a_temporary_file() { + fn the_file_replaces_an_existing_one() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("jobs.json"); fs::write(&path, "stale").unwrap(); - write_atomically(&path, "fresh").unwrap(); + write_atomically(&path, CONTENTS).unwrap(); - assert_eq!(fs::read_to_string(&path).unwrap(), "fresh"); - assert!(!dir.path().join("jobs.json.tmp").exists()); + assert_eq!(fs::read_to_string(&path).unwrap(), CONTENTS); } #[test] - fn a_file_that_cannot_be_written_leaves_no_file() { + fn a_written_file_leaves_no_temporary_file() { let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("missing").join("jobs.json"); + let path = dir.path().join("jobs.json"); + + write_atomically(&path, CONTENTS).unwrap(); + + assert_eq!(entries(dir.path()), [path.file_name().unwrap()]); + } + + #[test] + fn a_failed_rename_leaves_no_temporary_file() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("jobs.json"); + fs::create_dir(&path).unwrap(); + fs::write(path.join("kept"), "").unwrap(); + + write_atomically(&path, "{}").unwrap_err(); - assert!(write_atomically(&path, "{}").is_err()); - assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 0); + assert_eq!(entries(dir.path()), [path.file_name().unwrap()]); } } diff --git a/crates/tracel-runtime/src/process/mod.rs b/crates/tracel-runtime/src/process/mod.rs index 7975c503..53547cf6 100644 --- a/crates/tracel-runtime/src/process/mod.rs +++ b/crates/tracel-runtime/src/process/mod.rs @@ -6,17 +6,17 @@ mod error; mod file; use std::ffi::OsString; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::process::ExitCode; use clap::Command; use clap::error::{ContextKind, ContextValue, ErrorKind}; use clap_complete::Shell; use serde_json::Value; -use tracel_job::TRACEL_JOB_NUM; use tracel_job::cli::{command, completions, job_input}; +use tracel_job::{TRACEL_DESCRIBE, TRACEL_JOB_NUM}; -use crate::{Execution, JobDefinition, Jobs}; +use crate::{Execution, Jobs}; use describe::describe; use error::RunError; @@ -47,7 +47,11 @@ impl Jobs { /// [`job_num`](crate::job_num) gives the number its launcher gave it, and /// [`job_input`](crate::job_input) the input it runs with. pub fn run(self) -> ExitCode { - let exit = match execute(&self) { + let launch = Launch::new( + std::env::var_os(TRACEL_DESCRIBE), + std::env::var_os(TRACEL_JOB_NUM), + ); + let exit = match execute(&self, std::env::args_os(), launch) { Ok(()) => tracel_job::ExitCode::Completed, Err(error) => { error.print(); @@ -58,16 +62,49 @@ impl Jobs { } } -/// Does what the process's environment and command line ask of `jobs`: writes their definitions, -/// prints the help, the version or a completion script, or runs the job named. -fn execute(jobs: &Jobs) -> Result<(), RunError> { - let vars = |name: &str| std::env::var_os(name); - if describe(jobs.definitions(), vars)? { - return Ok(()); +/// What the program's launcher asks of it through [`TRACEL_DESCRIBE`] and [`TRACEL_JOB_NUM`]. +struct Launch { + /// Where to write the jobs' definitions, instead of running a job. + describe: Option, + /// The number of the job the program runs. + num: Option, +} + +impl Launch { + /// The launch the values of `TRACEL_DESCRIBE` and `TRACEL_JOB_NUM` ask for, an empty value + /// counting as unset. + fn new(describe: Option, num: Option) -> Self { + let set = |value: Option| value.filter(|value| !value.is_empty()); + Self { + describe: set(describe).map(PathBuf::from), + num: set(num).map(|num| num.to_string_lossy().into_owned()), + } + } + + /// The execution of the job the program runs: numbered when its launcher numbered it. + fn execution(&self) -> Execution { + match &self.num { + Some(num) => Execution::new().with_num(num.clone()), + None => Execution::new(), + } } - let mut command = command_line(jobs, program_name()); - match read_request(jobs, &mut command, std::env::args_os())? { - Request::Run { job, input } => dispatch(jobs, &job, input, execution(vars)), +} + +/// Does what the command-line arguments `args`, the program's name first, and the `launch` ask +/// of `jobs`: writes their definitions, prints the help, the version or a completion script, or +/// runs the job named. +fn execute( + jobs: &Jobs, + args: impl IntoIterator, + launch: Launch, +) -> Result<(), RunError> { + if let Some(path) = &launch.describe { + return Ok(describe(jobs.definitions(), path)?); + } + let args: Vec = args.into_iter().collect(); + let mut command = command_line(jobs, program_name(args.first())); + match read_request(jobs, &mut command, args)? { + Request::Run { job, input } => dispatch(jobs, &job, input, launch.execution()), Request::Completions(shell) => { let name = command.get_name().to_string(); clap_complete::generate(shell, &mut command, name, &mut std::io::stdout()); @@ -83,25 +120,13 @@ fn execute(jobs: &Jobs) -> Result<(), RunError> { /// The command line of the program `name`: the [`command`] of `jobs`, with `--version` when the /// program gives its version. fn command_line(jobs: &Jobs, name: impl Into) -> Command { - let definitions: Vec<_> = jobs.definitions().map(protocol_definition).collect(); - let command = command(name, &definitions); + let command = command(name, jobs.definitions()); match jobs.get_version() { Some(version) => command.version(version.to_string()), None => command, } } -/// The job `definition` defines, as the protocol gives it: in the definitions file, and to -/// generate its command line from. -fn protocol_definition(definition: &JobDefinition) -> tracel_job::JobDefinition { - tracel_job::JobDefinition { - name: definition.name.clone(), - description: definition.description.clone(), - input_schema: definition.input_schema.clone(), - input_example: definition.input_example.clone(), - } -} - /// The names of the jobs of `jobs`, in order. fn names(jobs: &Jobs) -> Vec { jobs.definitions() @@ -149,15 +174,6 @@ where } } -/// The execution of the job the process runs, as the environment variables `lookup` gives -/// describe it: numbered by [`TRACEL_JOB_NUM`] when its launcher sets it. -fn execution(lookup: impl Fn(&str) -> Option) -> Execution { - match lookup(TRACEL_JOB_NUM).filter(|num| !num.is_empty()) { - Some(num) => Execution::new().with_num(num.to_string_lossy()), - None => Execution::new(), - } -} - /// Runs the job `name` of `jobs` with `input`, launched as `execution` says. fn dispatch(jobs: &Jobs, name: &str, input: Value, execution: Execution) -> Result<(), RunError> { jobs.prepare(name, input)? @@ -166,11 +182,9 @@ fn dispatch(jobs: &Jobs, name: &str, input: Value, execution: Execution) -> Resu .map_err(RunError::JobFailed) } -/// The name the program was run as, without its extension, or `job` when it has none. -fn program_name() -> String { - std::env::args_os() - .next() - .as_deref() +/// The name the program was run as, `program`, without its extension, or `job` when it has none. +fn program_name(program: Option<&OsString>) -> String { + program .map(Path::new) .and_then(Path::file_stem) .map(|name| name.to_string_lossy().into_owned()) @@ -179,567 +193,236 @@ fn program_name() -> String { #[cfg(test)] mod tests { - use std::ffi::OsStr; - use std::sync::{Arc, Mutex}; + use std::sync::{Arc, OnceLock}; + use rstest::{fixture, rstest}; use serde::{Deserialize, Serialize}; - use serde_json::json; + use tracel_job::DefinitionsFile; use super::*; - use crate::{BoxError, IntoJob, Job, job_input, job_num, job_running}; - - #[derive(Clone, Copy)] - enum Outcome { - Succeed, - Fail, - Panic, - } + use crate::{BoxError, job_num}; - /// The input a job ran with, once it ran. - type RanWith = Arc>>; - - /// A job that takes any JSON input, keeps it in `ran_with`, then ends with `outcome`. - fn job( - outcome: Outcome, - ran_with: &RanWith, - ) -> impl Fn(Value) -> Result<(), BoxError> + Send + Sync + use<> { - let ran_with = ran_with.clone(); - move |input| { - *ran_with.lock().unwrap() = Some(input); - match outcome { - Outcome::Succeed => Ok(()), - Outcome::Fail => Err("boom".into()), - Outcome::Panic => panic!("kernel exploded"), - } - } + #[derive(Serialize, Deserialize)] + #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] + struct Epochs { + epochs: u32, } - /// `train`, which runs `train`, and `evaluate`, which succeeds. - fn jobs(train: impl IntoJob) -> Jobs { + #[fixture] + fn jobs() -> Jobs { Jobs::new() - .job("train", train) - .job("evaluate", job(Outcome::Succeed, &RanWith::default())) + .job("train", |_: Epochs| Ok::<_, BoxError>(())) + .job("fail", |()| Err::<(), BoxError>("the job failed".into())) } - /// `train`, which ends with `outcome`, and `evaluate`, and the input `train` ran with. - fn ending(outcome: Outcome) -> (Jobs, RanWith) { - let ran_with = RanWith::default(); - (jobs(job(outcome, &ran_with)), ran_with) + fn parse(jobs: &Jobs, args: &[&str]) -> Result { + let mut command = command_line(jobs, "program"); + read_request(jobs, &mut command, ["program"].iter().chain(args)) } - /// `train`, which succeeds with its input merged onto `default`, and `evaluate`, and the input - /// `train` ran with. - fn with_default(default: Value) -> (Jobs, RanWith) { - let ran_with = RanWith::default(); - let train = Job::new(job(Outcome::Succeed, &ran_with)).with_default(default); - (jobs(train), ran_with) + fn arguments(args: &[&str]) -> impl Iterator { + ["program"].iter().chain(args).map(OsString::from) } - #[derive(Serialize, Deserialize)] - struct Epochs { - epochs: u32, + fn describing(path: &Path) -> Launch { + Launch::new(Some(path.as_os_str().to_owned()), None) } - /// `train`, which takes `Epochs`, ten by default, and `evaluate`, and whether `train` ran. - fn typed() -> (Jobs, Arc>) { - let ran = Arc::new(Mutex::new(false)); - let train = Job::new({ - let ran = ran.clone(); - move |_: Epochs| { - *ran.lock().unwrap() = true; - Ok::<_, BoxError>(()) - } - }) - .with_default(Epochs { epochs: 10 }); - let jobs = Jobs::new() - .job("train", train) - .job("evaluate", job(Outcome::Succeed, &RanWith::default())); - (jobs, ran) + fn run(jobs: Jobs, args: &[&str]) -> Result<(), RunError> { + execute(&jobs, arguments(args), Launch::new(None, None)) } - /// What the command line `program ` asks of `jobs`. - fn parse(jobs: &Jobs, args: &[&str]) -> Result { - let mut command = command_line(jobs, "program"); - read_request(jobs, &mut command, ["program"].iter().chain(args)) - } + #[test] + fn an_empty_launch_variable_counts_as_unset() { + let launch = Launch::new(Some(OsString::new()), Some(OsString::new())); - /// Runs the command line `program ` as [`Jobs::run`] does, printing nothing. - fn run(jobs: Jobs, args: &[&str]) -> Result<(), RunError> { - run_launched(jobs, args, &[]) + assert_eq!((launch.describe, launch.num), (None, None)); } - /// Runs the command line `program ` as [`run`] does, launched with the environment - /// variables `vars`. - fn run_launched(jobs: Jobs, args: &[&str], vars: &[(&str, &OsStr)]) -> Result<(), RunError> { - match parse(&jobs, args)? { - Request::Run { job, input } => { - let execution = execution(|name| { - vars.iter() - .find(|(variable, _)| *variable == name) - .map(|(_, value)| value.to_os_string()) - }); - dispatch(&jobs, &job, input, execution) - } - Request::Completions(_) | Request::Print(_) => Ok(()), - } + #[rstest] + #[case::no_job(&[])] + #[case::a_job_that_fails(&["fail"])] + #[case::an_unusable_command_line(&["--bogus"])] + fn a_launcher_asking_for_the_definitions_gets_them_instead_of_a_run( + jobs: Jobs, + #[case] args: &[&str], + ) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("jobs.json"); + + execute(&jobs, arguments(args), describing(&path)).unwrap(); + + let file = DefinitionsFile::from_json(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!(file.jobs, jobs.definitions().cloned().collect::>()); } - #[test] - fn the_named_job_runs_with_its_json_input() { - let (jobs, ran_with) = ending(Outcome::Succeed); + #[rstest] + fn definitions_that_cannot_be_written_fail_the_run(jobs: Jobs) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("missing").join("jobs.json"); - run(jobs, &["train", r#"{"epochs": 2, "tag": null}"#]).unwrap(); + let failed = execute(&jobs, arguments(&["train"]), describing(&path)).unwrap_err(); - assert_eq!( - *ran_with.lock().unwrap(), - Some(json!({"epochs": 2, "tag": null})) - ); + assert_eq!(failed.exit_code(), tracel_job::ExitCode::Failed); } - #[test] - fn a_missing_input_runs_the_job_with_its_default_or_no_input() { - let (jobs, ran_with) = ending(Outcome::Succeed); - let (defaulted, ran_with_default) = with_default(json!({"epochs": 10})); - let ran = Arc::new(Mutex::new(false)); - let unit = Jobs::new().job("evaluate", { - let ran = ran.clone(); + #[rstest] + #[case::numbered(Some("12"))] + #[case::unnumbered(None)] + fn the_job_is_numbered_as_its_launcher_numbers_it(#[case] num: Option<&str>) { + let seen = Arc::new(OnceLock::new()); + let jobs = Jobs::new().job("train", { + let seen = seen.clone(); move |()| { - *ran.lock().unwrap() = true; + let _ = seen.set(job_num()); Ok::<_, BoxError>(()) } }); + let launch = Launch::new(None, num.map(OsString::from)); - run(jobs, &["train"]).unwrap(); - run(defaulted, &["train"]).unwrap(); - run(unit, &["evaluate"]).unwrap(); + execute(&jobs, arguments(&["train"]), launch).unwrap(); - assert_eq!(*ran_with.lock().unwrap(), Some(Value::Null)); - assert_eq!( - *ran_with_default.lock().unwrap(), - Some(json!({"epochs": 10})) - ); - assert!(*ran.lock().unwrap()); + assert_eq!(seen.get(), Some(&num.map(str::to_string))); } - #[test] - fn a_job_without_an_input_schema_takes_no_flags() { - let (jobs, ran_with) = with_default(json!({"epochs": 10})); - - let flag = run(jobs, &["train", "--epochs", "2"]).unwrap_err(); + #[rstest] + fn the_named_job_runs_with_the_input_its_command_line_gives(jobs: Jobs) { + let given = serde_json::json!({"epochs": 2}); - assert!( - matches!(&flag, RunError::Usage(error) if error.kind() == ErrorKind::UnknownArgument), - "{flag:?}" - ); - assert_eq!(flag.exit_code().code(), 2); - assert!(ran_with.lock().unwrap().is_none()); - } + let request = parse(&jobs, &["train", &given.to_string()]); - #[test] - fn the_json_input_is_merged_onto_the_default() { - let (jobs, ran_with) = with_default(json!({ - "num_epochs": 10, - "batch_size": 64, - "shuffle": true, - "tag": "baseline", - "layers": [64, 32], - "resume_from": null, - "optimizer": {"lr": 0.001, "weight_decay": 5e-5f32} - })); - - run( - jobs, - &[ - "train", - r#"{"num_epochs": 3, "batch_size": 16, "tag": null, "optimizer": {"lr": 0.3, "weight_decay": 0.5}}"#, - ], - ) - .unwrap(); - - // `null` removes the field from the default. - assert_eq!( - *ran_with.lock().unwrap(), - Some(json!({ - "num_epochs": 3, - "batch_size": 16, - "shuffle": true, - "layers": [64, 32], - "resume_from": null, - "optimizer": {"lr": 0.3, "weight_decay": 0.5} - })) - ); + assert!(matches!( + request, + Ok(Request::Run { job, input }) if job == "train" && input == given + )); } - #[test] - fn the_config_file_is_read_and_an_unreadable_one_rejects_the_input() { + #[rstest] + fn a_config_file_gives_the_named_job_its_input(jobs: Jobs) { + let given = serde_json::json!({"epochs": 2}); let dir = tempfile::tempdir().unwrap(); let config = dir.path().join("config.json"); - std::fs::write(&config, r#"{"epochs": 3}"#).unwrap(); - let missing = dir.path().join("missing.json"); - let (jobs, ran_with) = with_default(json!({"epochs": 10, "lr": 0.1})); - let (rejected_jobs, rejected) = with_default(json!({"epochs": 10})); + std::fs::write(&config, given.to_string()).unwrap(); - run(jobs, &["train", "-c", config.to_str().unwrap()]).unwrap(); - let unreadable = run(rejected_jobs, &["train", "-c", missing.to_str().unwrap()]); + let request = parse(&jobs, &["train", "-c", config.to_str().unwrap()]); - assert_eq!( - *ran_with.lock().unwrap(), - Some(json!({"epochs": 3, "lr": 0.1})) - ); - let unreadable = unreadable.unwrap_err(); - assert!(matches!(unreadable, RunError::Config(_)), "{unreadable:?}"); - assert_eq!(unreadable.exit_code().code(), 2); - assert!(rejected.lock().unwrap().is_none()); + assert!(matches!( + request, + Ok(Request::Run { job, input }) if job == "train" && input == given + )); } - #[test] - fn the_default_job_runs_when_no_job_is_named() { - let (jobs, ran_with) = ending(Outcome::Succeed); - - run(jobs.default_job("train"), &[]).unwrap(); - - assert_eq!(*ran_with.lock().unwrap(), Some(Value::Null)); - } + #[rstest] + fn the_default_job_runs_with_no_input_when_no_job_is_named(jobs: Jobs) { + let request = parse(&jobs.default_job("train"), &[]); - #[test] - fn an_unknown_or_missing_job_lists_the_registered_ones() { - let jobs = || ending(Outcome::Succeed).0; - - let unknown = run(jobs(), &["infer"]).unwrap_err(); - let missing = run(jobs(), &[]).unwrap_err(); - let unknown_default = run(jobs().default_job("infer"), &[]).unwrap_err(); - - for unknown in [&unknown, &unknown_default] { - assert!( - matches!( - unknown, - RunError::UnknownJob { name, available } - if name == "infer" && available == &["evaluate", "train"] - ), - "{unknown:?}" - ); - assert!(unknown.to_string().contains("evaluate, train"), "{unknown}"); - } - assert!(matches!(missing, RunError::MissingJob { .. })); - assert!(missing.to_string().contains("evaluate, train"), "{missing}"); + assert!(matches!( + request, + Ok(Request::Run { job, input }) if job == "train" && input.is_null() + )); } - #[test] - fn an_unusable_command_line_or_an_input_that_does_not_decode_is_invalid() { - let rejected = |args: &[&str]| { - let (jobs, ran) = typed(); - let rejected = run(jobs, args); - assert!(!*ran.lock().unwrap(), "{args:?}"); - rejected - }; - - let not_json = rejected(&["train", "{"]); - let unknown_flag = rejected(&["train", "--epochs", "2"]); - let not_decoded = rejected(&["train", "[1, 2]"]); - let wrong_type = rejected(&["train", r#"{"epochs": "ten"}"#]); - - for usage in [not_json, unknown_flag] { - assert!(matches!(usage, Err(RunError::Usage(_))), "{usage:?}"); - } - for invalid in [not_decoded, wrong_type] { - assert!( - matches!(invalid, Err(RunError::InvalidInput(_))), - "{invalid:?}" - ); - } - } + #[rstest] + fn an_unknown_job_is_rejected_with_the_registered_names(jobs: Jobs) { + let rejected = parse(&jobs, &["infer"]); - #[test] - fn help_and_completions_are_printed_instead_of_running_a_job() { - let (jobs, _) = with_default(json!({"epochs": 10})); - - for args in [&["--help"][..], &["train", "--help"], &["train", "-h"]] { - assert!( - matches!(parse(&jobs, args), Ok(Request::Print(help)) if help.exit_code() == 0), - "{args:?}" - ); - } assert!(matches!( - parse(&jobs, &["--completions", "bash"]), - Ok(Request::Completions(Shell::Bash)) - )); - assert!(matches!( - parse(&jobs, &["--completions", "bash", "train"]), - Err(RunError::Usage(_)) + rejected, + Err(RunError::UnknownJob { name, available }) + if name == "infer" && available == names(&jobs) )); } - #[test] - fn the_version_the_program_gives_is_printed_with_its_name() { - let jobs = ending(Outcome::Succeed).0.version("1.2.3"); - - let Ok(Request::Print(version)) = parse(&jobs, &["--version"]) else { - panic!("--version printed nothing"); - }; + #[rstest] + fn an_unknown_default_job_is_rejected_with_the_registered_names(jobs: Jobs) { + let registered = names(&jobs); - assert_eq!(version.kind(), ErrorKind::DisplayVersion); - assert!(!version.use_stderr()); - assert_eq!(version.exit_code(), 0); - assert_eq!(version.to_string(), "program 1.2.3\n"); - } + let rejected = run(jobs.default_job("infer"), &[]); - #[test] - fn without_a_version_there_is_no_version_flag() { - let jobs = ending(Outcome::Succeed).0; - - let unoffered = parse(&jobs, &["--version"]); - - assert!( - matches!( - &unoffered, - Err(RunError::Usage(error)) if error.kind() == ErrorKind::UnknownArgument - ), - "{:?}", - unoffered.as_ref().err() - ); - assert_eq!( - unoffered.err().map(|error| error.exit_code().code()), - Some(2) - ); + assert!(matches!( + rejected, + Err(RunError::UnknownJob { name, available }) + if name == "infer" && available == registered + )); } - #[test] - fn a_failing_or_panicking_job_fails() { - let failing = run(ending(Outcome::Fail).0, &["train"]); - let panicking = run(ending(Outcome::Panic).0, &["train"]); + #[rstest] + fn naming_no_job_without_a_default_is_rejected_with_the_registered_names(jobs: Jobs) { + let rejected = parse(&jobs, &[]); - assert!(matches!(failing, Err(RunError::JobFailed(e)) if e.to_string() == "boom")); assert!(matches!( - panicking, - Err(RunError::JobFailed(e)) if e.to_string().contains("kernel exploded") + rejected, + Err(RunError::MissingJob { available }) if available == names(&jobs) )); } - #[test] - fn each_outcome_has_its_exit_code() { - let exit_code = |args: &[&str]| run(typed().0, args).unwrap_err().exit_code().code(); - - assert_eq!(exit_code(&["infer"]), 2); - assert_eq!(exit_code(&[]), 2); - assert_eq!(exit_code(&["train", "{"]), 2); - assert_eq!(exit_code(&["train", "--epochs", "ten"]), 2); - assert_eq!(exit_code(&["train", "--epochs"]), 2); - assert_eq!(exit_code(&["--completions", "tcsh"]), 2); - assert_eq!(exit_code(&["train", "[1, 2]"]), 2); - for (outcome, code) in [(Outcome::Fail, 1), (Outcome::Panic, 1)] { - let failed = run(ending(outcome).0, &["train"]); - assert_eq!(failed.unwrap_err().exit_code().code(), code); - } - } - - #[test] - fn the_job_is_the_running_job_until_it_ends() { - let given = Arc::new(Mutex::new(Vec::new())); - let numbered = |num: &str| { - let jobs = Jobs::new().job("train", { - let given = given.clone(); - move |()| { - given.lock().unwrap().push((job_running(), job_num())); - Ok::<_, BoxError>(()) - } - }); - run_launched(jobs, &["train"], &[(TRACEL_JOB_NUM, OsStr::new(num))]).unwrap(); - }; + #[rstest] + #[case::an_unknown_job(&["infer"], tracel_job::ExitCode::Rejected)] + #[case::no_job(&[], tracel_job::ExitCode::Rejected)] + #[case::an_unusable_command_line(&["train", "--bogus"], tracel_job::ExitCode::Rejected)] + #[case::an_input_that_does_not_decode( + &["train", r#"{"epochs": "ten"}"#], + tracel_job::ExitCode::Rejected + )] + #[case::a_job_that_fails(&["fail"], tracel_job::ExitCode::Failed)] + fn a_run_that_does_not_complete_exits_with_the_code_that_says_why( + jobs: Jobs, + #[case] args: &[&str], + #[case] exit: tracel_job::ExitCode, + ) { + assert_eq!(run(jobs, args).unwrap_err().exit_code(), exit); + } + + #[rstest] + fn an_unreadable_config_file_is_rejected(jobs: Jobs) { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("missing.json"); - numbered("12"); - numbered(""); - let after = (job_running(), job_num()); + let rejected = run(jobs, &["train", "-c", config.to_str().unwrap()]).unwrap_err(); - assert_eq!( - *given.lock().unwrap(), - [(true, Some("12".to_string())), (true, None)] - ); - assert_eq!(after, (false, None)); + assert_eq!(rejected.exit_code(), tracel_job::ExitCode::Rejected); } - #[test] - fn the_running_job_is_given_the_input_it_runs_with() { - let given = Arc::new(Mutex::new(Vec::new())); - let train = Job::new({ - let given = given.clone(); - move |input: Value| { - given.lock().unwrap().push((job_input(), Some(input))); - Ok::<_, BoxError>(()) - } - }) - .with_default(json!({"epochs": 10, "optimizer": {"lr": 0.001}})); - let evaluate = { - let given = given.clone(); - move |()| { - given.lock().unwrap().push((job_input(), None)); - Ok::<_, BoxError>(()) - } - }; + #[rstest] + #[case::help(&["--help"])] + #[case::version(&["--version"])] + fn help_and_the_version_are_printed_instead_of_running_a_job( + jobs: Jobs, + #[case] args: &[&str], + ) { + let request = parse(&jobs.version("1.2.3"), args); - run( - Jobs::new().job("train", train), - &["train", r#"{"epochs": 2}"#], - ) - .unwrap(); - run(Jobs::new().job("evaluate", evaluate), &["evaluate"]).unwrap(); - - let merged = json!({"epochs": 2, "optimizer": {"lr": 0.001}}); - assert_eq!( - *given.lock().unwrap(), - [ - (Some(merged.clone()), Some(merged)), - (Some(Value::Null), None) - ] - ); + assert!(matches!(request, Ok(Request::Print(_)))); } - #[cfg(feature = "schema")] - #[test] - fn a_documented_input_type_gives_documented_flags() { - /// How to train. - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - struct Config { - /// Passes over the data. - num_epochs: u32, - optimizer: Optimizer, - /// Where to train. - device: Device, - resume_from: Option, - } - - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - struct Optimizer { - /// The learning rate. - lr: f64, - } + #[rstest] + fn the_completion_script_is_printed_for_the_shell_asked_for(jobs: Jobs) { + let request = parse(&jobs, &["--completions", "bash"]); - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - #[serde(rename_all = "snake_case")] - enum Device { - /// The CPU. - Cpu, - /// The first GPU. - Gpu, - } + assert!(matches!(request, Ok(Request::Completions(Shell::Bash)))); + } - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - struct Resume { - experiment: u32, - } + #[rstest] + fn without_a_version_there_is_no_version_flag(jobs: Jobs) { + let request = parse(&jobs, &["--version"]); - let ran_with = Arc::new(Mutex::new(None)); - let train = Job::new({ - let ran_with = ran_with.clone(); - move |config: Config| { - *ran_with.lock().unwrap() = Some(config); - Ok::<_, BoxError>(()) - } - }) - .with_default(Config { - num_epochs: 10, - optimizer: Optimizer { lr: 0.001 }, - device: Device::Cpu, - resume_from: None, - }) - .with_schema(); - let jobs = Jobs::new().job("train", train); - - let help = command_line(&jobs, "program") - .find_subcommand_mut("train") - .unwrap() - .render_long_help() - .to_string(); - for line in [ - "--device ", - "Where to train.", - "cpu: The CPU.", - "--num-epochs ", - "Passes over the data.", - "--optimizer.lr ", - "The learning rate.", - "--resume-from ", - ] { - assert!(help.contains(line), "no {line} in\n{help}"); - } - run( - jobs, - &[ - "train", - "--device", - "gpu", - "--resume-from", - r#"{"experiment": 4}"#, - ], - ) - .unwrap(); - let config = ran_with.lock().unwrap().take().unwrap(); - assert!(matches!(config.device, Device::Gpu)); - assert_eq!(config.resume_from.map(|resume| resume.experiment), Some(4)); - assert_eq!(config.num_epochs, 10); + assert!(matches!(request, Err(RunError::Usage(_)))); } #[cfg(feature = "schema")] #[test] - fn flags_from_the_input_schema_set_fields_after_the_json_input() { - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - struct Training { - epochs: u32, - optimizer: Optimizer, - } - - #[derive(Serialize, Deserialize, schemars::JsonSchema)] - struct Optimizer { - lr: f64, - weight_decay: f64, - } - - let ran_with = RanWith::default(); - let jobs = || { - let train = Job::new({ - let ran_with = ran_with.clone(); - move |config: Training| { - *ran_with.lock().unwrap() = Some(serde_json::to_value(config)?); - Ok::<_, BoxError>(()) - } - }) - .with_default(Training { - epochs: 10, - optimizer: Optimizer { - lr: 0.001, - weight_decay: 0.0, - }, - }) - .with_schema(); - Jobs::new().job("train", train) - }; - - run( - jobs(), - &[ - "train", - "--optimizer.weight-decay", - "0.01", - r#"{"epochs": 2, "optimizer": {"weight_decay": 0.5}}"#, - ], - ) - .unwrap(); - let not_an_integer = run(jobs(), &["train", "--epochs", "ten"]).unwrap_err(); - - assert_eq!( - ran_with.lock().unwrap().take(), - Some(json!({"epochs": 2, "optimizer": {"lr": 0.001, "weight_decay": 0.01}})) + fn a_job_with_an_input_schema_takes_its_fields_as_flags() { + let jobs = Jobs::new().job( + "train", + crate::Job::new(|_: Epochs| Ok::<_, BoxError>(())).with_schema(), ); - assert!( - matches!(not_an_integer, RunError::Usage(_)), - "{not_an_integer:?}" - ); - assert!(ran_with.lock().unwrap().is_none()); - } + let epochs = 2; - #[test] - #[should_panic(expected = "already registered")] - fn a_job_name_registers_once() { - jobs(job(Outcome::Succeed, &RanWith::default())) - .job("evaluate", |()| Ok::<_, BoxError>(())); + let request = parse(&jobs, &["train", "--epochs", &epochs.to_string()]); + + assert!(matches!( + request, + Ok(Request::Run { input, .. }) if input == serde_json::json!({"epochs": epochs}) + )); } } diff --git a/crates/tracel-runtime/src/scope.rs b/crates/tracel-runtime/src/scope.rs index 9a1638ee..2a04bd68 100644 --- a/crates/tracel-runtime/src/scope.rs +++ b/crates/tracel-runtime/src/scope.rs @@ -2,8 +2,8 @@ //! //! Each execution of a job has a scope of its own, current on the thread that runs the job's //! function, while it runs. Code in the job reads it through the functions below, or holds it, -//! through [`JobScope::current`], to read it on another thread. While a thread runs no job, such as -//! in a plain binary or a unit test, there is none, and the functions say so. +//! through [`JobScope::current`], to read it on another thread. Outside a job there is none, and the +//! functions say so. use std::cell::RefCell; use std::sync::Arc; @@ -91,6 +91,9 @@ pub fn job_input() -> Option { #[cfg(test)] mod tests { + use std::sync::Barrier; + use std::thread; + use super::*; fn scope(num: &str) -> JobScope { @@ -98,63 +101,61 @@ mod tests { } #[test] - fn a_scope_is_current_on_its_thread_until_its_guard_drops() { - let entered = enter(scope("12")); - let given = (job_running(), job_num(), job_input()); - drop(entered); + fn the_entered_scope_gives_the_number_and_input_of_its_job() { + let job = scope("12"); + let _entered = enter(job.clone()); assert_eq!( - given, - ( - true, - Some("12".to_string()), - Some(serde_json::json!({"num": "12"})) - ) + (job_num().as_deref(), job_input().as_ref()), + (job.num(), Some(job.input())) ); + } + + #[test] + fn once_its_guard_drops_the_thread_runs_no_job() { + drop(enter(scope("12"))); + assert_eq!((job_running(), job_num(), job_input()), (false, None, None)); } #[test] - fn each_thread_sees_the_scope_of_the_job_it_runs() { - let run = |num: &'static str| { - let scope = scope(num); - std::thread::spawn(move || { - let _entered = enter(scope); - (job_num(), job_input()) + fn a_nested_scope_gives_way_to_the_one_before_when_its_guard_drops() { + let outer = scope("1"); + let _outer = enter(outer.clone()); + + drop(enter(scope("2"))); + + assert_eq!(job_num().as_deref(), outer.num()); + } + + #[test] + fn jobs_running_at_once_each_see_their_own_scope() { + let both_entered = Arc::new(Barrier::new(2)); + let run = |job: JobScope| { + let both_entered = both_entered.clone(); + thread::spawn(move || { + let _entered = enter(job); + both_entered.wait(); + job_num() }) }; + let jobs = [scope("1"), scope("2")]; - let one = run("1"); - let two = run("2"); + let seen = jobs.clone().map(run).map(|running| running.join().unwrap()); - let input = |num: &str| Some(serde_json::json!({ "num": num })); - assert_eq!(one.join().unwrap(), (Some("1".to_string()), input("1"))); - assert_eq!(two.join().unwrap(), (Some("2".to_string()), input("2"))); - assert!(!job_running()); + assert_eq!(seen, jobs.map(|job| job.num().map(str::to_string))); } #[test] fn a_held_scope_reads_the_same_job_on_another_thread() { - let _entered = enter(scope("12")); + let job = scope("12"); + let _entered = enter(job.clone()); let held = JobScope::current().unwrap(); - let seen = std::thread::spawn(move || { - ( - job_running(), - held.num().map(str::to_string), - held.input().clone(), - ) - }) - .join() - .unwrap(); + let read = thread::spawn(move || (held.num().map(str::to_string), held.input().clone())) + .join() + .unwrap(); - assert_eq!( - seen, - ( - false, - Some("12".to_string()), - serde_json::json!({"num": "12"}) - ) - ); + assert_eq!(read, (job.num().map(str::to_string), job.input().clone())); } } diff --git a/crates/tracel-runtime/src/station/error.rs b/crates/tracel-runtime/src/station/error.rs index 631508bd..dfb17bc4 100644 --- a/crates/tracel-runtime/src/station/error.rs +++ b/crates/tracel-runtime/src/station/error.rs @@ -5,13 +5,28 @@ /// permanently rejects its registration. #[derive(Debug, thiserror::Error)] pub enum RunnerError { + /// The station URL does not parse. #[error("invalid station url '{url}': {source}")] InvalidUrl { + /// The URL given. url: String, + /// Why it does not parse. source: url::ParseError, }, + /// The station URL is not an `http` or `https` URL. + #[error("station url '{url}' is not an http or https url")] + NotHttp { + /// The URL given. + url: String, + }, + /// The HTTP client that reaches the station could not be built. + #[error("failed to build the http client: {0}")] + HttpClient(#[source] reqwest::Error), + /// The runner was given no jobs to serve. #[error("no jobs registered")] NoJobs, + /// The station rejected the runner's registration for good, with the status and body it + /// answered with. #[error("station rejected runner registration: {0}")] Registration(String), } diff --git a/crates/tracel-runtime/src/station/infrastructure/mod.rs b/crates/tracel-runtime/src/station/infrastructure/mod.rs index ffb8ea38..e0c0896e 100644 --- a/crates/tracel-runtime/src/station/infrastructure/mod.rs +++ b/crates/tracel-runtime/src/station/infrastructure/mod.rs @@ -12,6 +12,7 @@ use std::time::Duration; use reqwest::Url; use reqwest::blocking::{Client, Response}; +use tracing::warn; use uuid::Uuid; use protocol::{FinishJob, RegisterRunner, RunnerEvent}; @@ -26,10 +27,6 @@ pub enum ClientError { status: reqwest::StatusCode, body: String, }, - #[error("stream io error: {0}")] - Io(#[from] std::io::Error), - #[error("invalid event payload: {0}")] - Decode(#[from] serde_json::Error), } impl ClientError { @@ -54,25 +51,25 @@ pub struct StationRunnerClient { } impl StationRunnerClient { - pub fn new(base_url: Url) -> Self { + /// The client of the station at `base_url`, an `http` or `https` URL. + pub fn new(base_url: Url) -> Result { let events_stream_client = Client::builder() .timeout(None) .tcp_keepalive(Duration::from_secs(60)) - .build() - .expect("failed to build HTTP client"); - let short_call_client = Client::builder() - .timeout(Duration::from_secs(60)) - .build() - .expect("failed to build HTTP client"); - Self { + .build()?; + let short_call_client = Client::builder().timeout(Duration::from_secs(60)).build()?; + Ok(Self { base_url, events_stream_client, short_call_client, - } + }) } /// Register with the station; the successful response IS the runner's event stream. - pub fn open_events(&self, register: &RegisterRunner) -> Result { + pub fn open_events( + &self, + register: &RegisterRunner, + ) -> Result>, ClientError> { let response = self .events_stream_client .post(self.join("runners/events")) @@ -80,7 +77,7 @@ impl StationRunnerClient { .json(register) .send()?; let response = ensure_success(response)?; - Ok(RunnerEventStream::new(response)) + Ok(RunnerEventStream::new(BufReader::new(response))) } pub fn finish_job(&self, job_id: Uuid, finish: &FinishJob) -> Result<(), ClientError> { @@ -110,40 +107,100 @@ fn ensure_success(response: Response) -> Result { Err(ClientError::Api { status, body }) } -/// Blocking iterator over the events of an open runner stream. +/// Blocking iterator over the events of an open runner stream, read from `R`. /// -/// Yields `None` when the station closes the stream; the session is over either way. -pub struct RunnerEventStream { - lines: std::io::Lines>, +/// Skips a frame that does not decode, with a warning, and ends when the station closes the stream +/// or reading it fails; the session is over either way. +pub struct RunnerEventStream { + lines: std::io::Lines, parser: SseParser, } -impl RunnerEventStream { - fn new(response: Response) -> Self { +impl RunnerEventStream { + fn new(reader: R) -> Self { Self { - lines: BufReader::new(response).lines(), + lines: reader.lines(), parser: SseParser::default(), } } } -impl Iterator for RunnerEventStream { - type Item = Result; +impl Iterator for RunnerEventStream { + type Item = RunnerEvent; fn next(&mut self) -> Option { loop { let line = match self.lines.next()? { Ok(line) => line, - Err(e) => return Some(Err(e.into())), + Err(error) => { + warn!(%error, "the event stream failed"); + return None; + } }; let Some(frame) = self.parser.push_line(&line) else { continue; }; match RunnerEvent::decode(&frame.event, &frame.data) { - Ok(Some(event)) => return Some(Ok(event)), - Ok(None) => continue, - Err(e) => return Some(Err(e.into())), + Ok(Some(event)) => return Some(event), + Ok(None) => {} + Err(error) => warn!(event = %frame.event, %error, "skipped a malformed event"), + } + } + } +} + +#[cfg(test)] +mod tests { + use std::io::{self, Read}; + + use serde_json::json; + + use super::*; + + struct FailingOnce(bool); + + impl Read for FailingOnce { + fn read(&mut self, _: &mut [u8]) -> io::Result { + if std::mem::replace(&mut self.0, true) { + return Ok(0); } + Err(io::Error::other("connection reset")) } } + + fn frame(event: &str, data: &str) -> String { + format!("event: {event}\ndata: {data}\n\n") + } + + fn registered(runner_id: Uuid) -> String { + frame("registered", &json!({ "runner_id": runner_id }).to_string()) + } + + fn events(stream: impl Read) -> Vec { + RunnerEventStream::new(BufReader::new(stream)).collect() + } + + #[test] + fn a_frame_whose_data_does_not_decode_is_skipped() { + let runner_id = Uuid::new_v4(); + let stream = frame("job", "not json") + ®istered(runner_id); + + assert_eq!( + events(stream.as_bytes()), + [RunnerEvent::Registered { runner_id }] + ); + } + + #[test] + fn a_failed_read_ends_the_stream() { + let runner_id = Uuid::new_v4(); + let registered = registered(runner_id); + let cancel = frame("cancel", &json!({ "job_id": Uuid::new_v4() }).to_string()); + let stream = registered + .as_bytes() + .chain(FailingOnce(false)) + .chain(cancel.as_bytes()); + + assert_eq!(events(stream), [RunnerEvent::Registered { runner_id }]); + } } diff --git a/crates/tracel-runtime/src/station/infrastructure/protocol.rs b/crates/tracel-runtime/src/station/infrastructure/protocol.rs index 2e5af8e4..e94bcc75 100644 --- a/crates/tracel-runtime/src/station/infrastructure/protocol.rs +++ b/crates/tracel-runtime/src/station/infrastructure/protocol.rs @@ -113,86 +113,7 @@ mod tests { use super::*; #[test] - fn given_known_events_when_decoding_then_yields_them() { - let id = Uuid::nil(); - - let registered = RunnerEvent::decode("registered", &format!("{{\"runner_id\":\"{id}\"}}")) - .unwrap() - .unwrap(); - let job = RunnerEvent::decode( - "job", - &format!("{{\"id\":\"{id}\",\"job_name\":\"train\",\"input\":{{\"epochs\":2}}}}"), - ) - .unwrap() - .unwrap(); - let cancel = RunnerEvent::decode("cancel", &format!("{{\"job_id\":\"{id}\"}}")) - .unwrap() - .unwrap(); - - assert_eq!(registered, RunnerEvent::Registered { runner_id: id }); - assert_eq!(cancel, RunnerEvent::Cancel { job_id: id }); - assert_eq!( - job, - RunnerEvent::Job(DispatchedJob { - id, - job_name: "train".to_string(), - input: serde_json::json!({"epochs": 2}), - }) - ); - } - - #[test] - fn given_unknown_event_when_decoding_then_yields_none() { + fn an_unknown_event_decodes_to_nothing() { assert_eq!(RunnerEvent::decode("heartbeat", "{}").unwrap(), None); } - - #[test] - fn given_malformed_data_when_decoding_then_errors() { - assert!(RunnerEvent::decode("job", "not json").is_err()); - assert!(RunnerEvent::decode("cancel", "{}").is_err()); - } - - #[test] - fn a_registration_gives_each_job_its_name_description_and_input() { - let definition = JobDefinition { - name: "train".to_string(), - description: Some("Train the model".to_string()), - input_schema: Some(serde_json::json!({"type": "integer"})), - input_example: Some(serde_json::json!(3)), - }; - - let register = RegisterRunner { - name: Some("runner".to_string()), - jobs: vec![RegisteredJob::from(&definition)], - }; - - assert_eq!( - serde_json::to_value(register).unwrap(), - serde_json::json!({ - "name": "runner", - "jobs": [{ - "name": "train", - "description": "Train the model", - "input_schema": {"type": "integer"}, - "input_example": 3 - }] - }) - ); - } - - #[test] - fn a_finish_gives_its_status_in_snake_case() { - let finish = |status| { - serde_json::to_value(FinishJob { - runner_id: Uuid::nil(), - status, - reason: None, - }) - .unwrap()["status"] - .clone() - }; - - assert_eq!(finish(FinishStatus::Completed), "completed"); - assert_eq!(finish(FinishStatus::Failed), "failed"); - } } diff --git a/crates/tracel-runtime/src/station/infrastructure/sse.rs b/crates/tracel-runtime/src/station/infrastructure/sse.rs index b5a13c5c..6921f8c3 100644 --- a/crates/tracel-runtime/src/station/infrastructure/sse.rs +++ b/crates/tracel-runtime/src/station/infrastructure/sse.rs @@ -62,7 +62,7 @@ mod tests { } #[test] - fn given_event_and_data_when_blank_line_then_frame_is_complete() { + fn a_blank_line_completes_the_frame_of_its_event_and_data() { let frames = parse(&["event: job", "data: {\"id\":1}", ""]); assert_eq!( @@ -75,40 +75,49 @@ mod tests { } #[test] - fn given_no_event_name_then_defaults_to_message() { + fn a_frame_without_an_event_name_is_a_message() { let frames = parse(&["data: hello", ""]); assert_eq!(frames[0].event, "message"); } #[test] - fn given_multiple_data_lines_then_joined_with_newline() { + fn data_lines_join_with_a_newline() { let frames = parse(&["event: job", "data: a", "data: b", ""]); assert_eq!(frames[0].data, "a\nb"); } #[test] - fn given_comments_and_other_fields_then_ignored() { + fn comments_and_other_fields_are_ignored() { let frames = parse(&[": keep-alive", "id: 42", "retry: 1000", ""]); assert!(frames.is_empty()); } #[test] - fn given_value_without_leading_space_then_parsed() { + fn a_value_needs_no_leading_space() { let frames = parse(&["event:job", "data:x", ""]); - assert_eq!(frames[0].event, "job"); - assert_eq!(frames[0].data, "x"); + assert_eq!( + frames, + [SseFrame { + event: "job".to_string(), + data: "x".to_string(), + }] + ); } #[test] - fn given_consecutive_frames_then_state_resets_between_them() { + fn each_frame_starts_afresh() { let frames = parse(&["event: a", "data: 1", "", "data: 2", ""]); - assert_eq!(frames.len(), 2); - assert_eq!(frames[1].event, "message"); - assert_eq!(frames[1].data, "2"); + assert_eq!( + frames[1], + SseFrame { + event: "message".to_string(), + data: "2".to_string(), + } + ); } } diff --git a/crates/tracel-runtime/src/station/mod.rs b/crates/tracel-runtime/src/station/mod.rs index 704662d4..89bac534 100644 --- a/crates/tracel-runtime/src/station/mod.rs +++ b/crates/tracel-runtime/src/station/mod.rs @@ -88,6 +88,9 @@ impl StationRunner { url: self.url.clone(), source, })?; + if !matches!(url.scheme(), "http" | "https") { + return Err(RunnerError::NotHttp { url: self.url }); + } if jobs.definitions().next().is_none() { return Err(RunnerError::NoJobs); } @@ -97,7 +100,7 @@ impl StationRunner { metadata: self.metadata, jobs: jobs.definitions().map(RegisteredJob::from).collect(), }; - let client = StationRunnerClient::new(url); + let client = StationRunnerClient::new(url).map_err(RunnerError::HttpClient)?; let executor = Executor::new(Arc::new(jobs), Arc::new(client.clone())); Err(serve::serve_forever(client, register, executor)) } @@ -105,6 +108,8 @@ impl StationRunner { #[cfg(test)] mod tests { + use rstest::{fixture, rstest}; + use super::*; use crate::BoxError; @@ -112,17 +117,33 @@ mod tests { Ok(()) } - #[test] - fn given_invalid_url_when_running_then_fails_to_start() { - let result = StationRunner::new("not a url").run(Jobs::new().job("train", train)); + #[fixture] + fn jobs() -> Jobs { + Jobs::new().job("train", train) + } + + #[rstest] + fn a_url_that_does_not_parse_does_not_start_the_runner(jobs: Jobs) { + let given = "not a url"; + + let failed = StationRunner::new(given).run(jobs); + + assert!(matches!(failed, Err(RunnerError::InvalidUrl { url, .. }) if url == given)); + } + + #[rstest] + #[case::a_host_and_port("localhost:8000")] + #[case::of_another_scheme("ftp://station")] + fn a_url_that_is_not_http_does_not_start_the_runner(jobs: Jobs, #[case] given: &str) { + let failed = StationRunner::new(given).run(jobs); - assert!(matches!(result, Err(RunnerError::InvalidUrl { .. }))); + assert!(matches!(failed, Err(RunnerError::NotHttp { url }) if url == given)); } #[test] - fn given_no_jobs_when_running_then_fails_to_start() { - let result = StationRunner::new("http://localhost:9000").run(Jobs::new()); + fn a_runner_with_no_jobs_does_not_start() { + let failed = StationRunner::new("http://localhost:9000").run(Jobs::new()); - assert!(matches!(result, Err(RunnerError::NoJobs))); + assert!(matches!(failed, Err(RunnerError::NoJobs))); } } diff --git a/crates/tracel-runtime/src/station/serve.rs b/crates/tracel-runtime/src/station/serve.rs index bf4917ee..565a4cbf 100644 --- a/crates/tracel-runtime/src/station/serve.rs +++ b/crates/tracel-runtime/src/station/serve.rs @@ -7,22 +7,21 @@ use std::time::Duration; use tracing::{info, warn}; use uuid::Uuid; +use crate::panics::catch_panic; use crate::station::error::RunnerError; +use crate::station::infrastructure::StationRunnerClient; use crate::station::infrastructure::protocol::{ DispatchedJob, FinishJob, FinishStatus, RegisterRunner, RunnerEvent, }; -use crate::station::infrastructure::{ClientError, StationRunnerClient}; use crate::{Execution, Jobs}; -/// Reports job outcomes back to the station. Split out so the loop is testable with a fake. +/// Reports job outcomes back to the station. pub trait FinishSink: Send + Sync + 'static { fn finish_job(&self, job_id: Uuid, finish: &FinishJob); } impl FinishSink for StationRunnerClient { fn finish_job(&self, job_id: Uuid, finish: &FinishJob) { - // A rejected finish is expected after a disconnect: the station has already failed the - // job and this session no longer owns it. if let Err(e) = StationRunnerClient::finish_job(self, job_id, finish) { warn!(%job_id, error = %e, "failed to report the outcome of the job to the station"); } @@ -52,13 +51,13 @@ impl Executor { /// Runs `job` on a thread of its own, and reports how it ended to the station once it ends. pub fn dispatch(&self, runner_id: Uuid, job: DispatchedJob) { - *self.in_flight.0.lock().unwrap() += 1; - let (jobs, sink, in_flight) = - (self.jobs.clone(), self.sink.clone(), self.in_flight.clone()); + let running = Running::start(&self.in_flight); + let (jobs, sink) = (self.jobs.clone(), self.sink.clone()); let job_id = job.id; let spawned = std::thread::Builder::new() .name(format!("job-{job_id}")) .spawn(move || { + let _running = running; let (status, reason) = execute(&jobs, &job); sink.finish_job( job.id, @@ -68,7 +67,6 @@ impl Executor { reason, }, ); - finish_in_flight(&in_flight); }); if let Err(error) = spawned { let reason = format!("the job could not start: {error}"); @@ -80,7 +78,6 @@ impl Executor { reason: Some(reason), }, ); - finish_in_flight(&self.in_flight); } } @@ -94,23 +91,35 @@ impl Executor { } } -fn finish_in_flight(in_flight: &InFlight) { - let (count, finished) = &**in_flight; - *count.lock().unwrap() -= 1; - finished.notify_all(); +/// A dispatched job, counted in flight until this drops, when its thread ends or unwinds. +struct Running(InFlight); + +impl Running { + fn start(in_flight: &InFlight) -> Self { + *in_flight.0.lock().unwrap() += 1; + Self(in_flight.clone()) + } +} + +impl Drop for Running { + fn drop(&mut self) { + let (count, finished) = &*self.0; + *count.lock().unwrap() -= 1; + finished.notify_all(); + } } /// Runs `job` and returns how it ended: completed or failed by what it returned, with its error -/// when it failed. +/// when it failed. A panic fails it, even one while its input decodes. fn execute(jobs: &Jobs, job: &DispatchedJob) -> (FinishStatus, Option) { - let prepared = match jobs.prepare(&job.job_name, job.input.clone()) { - Ok(prepared) => prepared, - // An unknown job is unreachable under the station's strict policy; defend anyway. - Err(error) => return (FinishStatus::Failed, Some(error.to_string())), - }; - info!(job = %job.job_name, job_id = %job.id, "running the job"); - let outcome = prepared.run(Execution::new().with_num(job.id.to_string())); - match outcome.into_result() { + let ran = catch_panic(|| { + let prepared = jobs.prepare(&job.job_name, job.input.clone())?; + info!(job = %job.job_name, job_id = %job.id, "running the job"); + prepared + .run(Execution::new().with_num(job.id.to_string())) + .into_result() + }); + match ran { Ok(()) => (FinishStatus::Completed, None), Err(error) => (FinishStatus::Failed, Some(error.to_string())), } @@ -126,19 +135,15 @@ pub enum StreamOutcome { /// Drive one session: expect the leading `registered` event, then route job dispatches to the /// executor until the stream ends. pub fn serve_stream( - mut events: impl Iterator>, + mut events: impl Iterator, executor: &Executor, ) -> StreamOutcome { let runner_id = match events.next() { - Some(Ok(RunnerEvent::Registered { runner_id })) => runner_id, - Some(Ok(event)) => { + Some(RunnerEvent::Registered { runner_id }) => runner_id, + Some(event) => { warn!(?event, "expected a registered event first"); return StreamOutcome::NeverRegistered; } - Some(Err(e)) => { - warn!(error = %e, "the event stream failed before registration"); - return StreamOutcome::NeverRegistered; - } None => { warn!("the event stream closed before registration"); return StreamOutcome::NeverRegistered; @@ -148,17 +153,13 @@ pub fn serve_stream( for event in events { match event { - Ok(RunnerEvent::Job(job)) => executor.dispatch(runner_id, job), - Ok(RunnerEvent::Cancel { job_id }) => { + RunnerEvent::Job(job) => executor.dispatch(runner_id, job), + RunnerEvent::Cancel { job_id } => { info!(%job_id, "the station cancelled the job; it runs to completion"); } - Ok(RunnerEvent::Registered { .. }) => { + RunnerEvent::Registered { .. } => { warn!("unexpected registered event mid-stream"); } - Err(e) => { - warn!(error = %e, "the event stream failed"); - break; - } } } StreamOutcome::Served @@ -201,8 +202,6 @@ pub fn serve_forever( if let StreamOutcome::Served = serve_stream(stream, &executor) { backoff.reset(); } - // Never re-register while user code still runs: the station failed those jobs when - // the stream dropped, and a fresh session would get new ones. Drain them first. executor.wait_idle(); } Err(e) if e.is_permanent() => return RunnerError::Registration(e.to_string()), @@ -216,10 +215,14 @@ pub fn serve_forever( #[cfg(test)] mod tests { + use rstest::{fixture, rstest}; + use serde::{Deserialize, Deserializer}; use serde_json::{Value, json}; + use std::sync::{Barrier, OnceLock}; + use super::*; - use crate::{BoxError, job_input, job_num}; + use crate::{BoxError, job_num}; #[allow(clippy::type_complexity)] type RecordedFinish = (Uuid, Uuid, FinishStatus, Option); @@ -240,18 +243,25 @@ mod tests { } } - /// `train` succeeds, `fail` fails, `panic` panics, `typed` takes a number, and `num` keeps the - /// job number it is given. - fn jobs(nums: Arc>>>) -> Jobs { + struct Explosive; + + impl<'de> Deserialize<'de> for Explosive { + fn deserialize>(_: D) -> Result { + panic!("the decoder exploded") + } + } + + const FAILURE: &str = "boom"; + + fn jobs() -> Jobs { + let both = Arc::new(Barrier::new(2)); Jobs::new() .job("train", |_: Value| Ok::<_, BoxError>(())) - .job("fail", |_: Value| Err::<(), BoxError>("boom".into())) - .job("panic", |_: Value| -> Result<(), BoxError> { - panic!("kaboom") - }) + .job("fail", |_: Value| Err::<(), BoxError>(FAILURE.into())) .job("typed", |_: u32| Ok::<_, BoxError>(())) - .job("num", move |_: Value| { - nums.lock().unwrap().push(job_num()); + .job("explosive", |_: Explosive| Ok::<_, BoxError>(())) + .job("wait", move |_: Value| { + both.wait(); Ok::<_, BoxError>(()) }) } @@ -260,21 +270,23 @@ mod tests { executor: Executor, sink: Arc, runner_id: Uuid, - nums: Arc>>>, } - fn setup() -> Setup { - let nums = Arc::default(); + fn serving(jobs: Jobs) -> Setup { let sink = Arc::new(RecordingSink::default()); - let executor = Executor::new(Arc::new(jobs(Arc::clone(&nums))), sink.clone()); + let executor = Executor::new(Arc::new(jobs), sink.clone()); Setup { executor, sink, runner_id: Uuid::new_v4(), - nums, } } + #[fixture] + fn setup() -> Setup { + serving(jobs()) + } + fn dispatched(job_name: &str) -> DispatchedJob { DispatchedJob { id: Uuid::new_v4(), @@ -283,136 +295,112 @@ mod tests { } } - fn events( - runner_id: Uuid, - rest: Vec, - ) -> impl Iterator> { - std::iter::once(RunnerEvent::Registered { runner_id }) - .chain(rest) - .map(Ok) + fn events(runner_id: Uuid, rest: Vec) -> impl Iterator { + std::iter::once(RunnerEvent::Registered { runner_id }).chain(rest) } - /// Serves `jobs`, dispatched in order, and returns the finishes reported. fn serve(setup: &Setup, jobs: Vec) -> Vec { let rest = jobs.into_iter().map(RunnerEvent::Job).collect(); serve_stream(events(setup.runner_id, rest), &setup.executor); setup.executor.wait_idle(); - setup.sink.finishes.lock().unwrap().clone() + finishes(setup) } - #[test] - fn given_job_event_when_serving_then_completion_reported() { - let setup = setup(); - let job = dispatched("train"); + fn finishes(setup: &Setup) -> Vec { + let mut finishes = setup.sink.finishes.lock().unwrap().clone(); + finishes.sort_by_key(|(job_id, ..)| *job_id); + finishes + } - let finishes = serve(&setup, vec![job.clone()]); + fn statuses(finishes: &[RecordedFinish]) -> Vec<(Uuid, FinishStatus)> { + finishes + .iter() + .map(|(job_id, _, status, _)| (*job_id, *status)) + .collect() + } - assert_eq!( - finishes, - [(job.id, setup.runner_id, FinishStatus::Completed, None)] - ); + fn completed<'a>( + jobs: impl IntoIterator, + ) -> Vec<(Uuid, FinishStatus)> { + let mut completed: Vec<_> = jobs + .into_iter() + .map(|job| (job.id, FinishStatus::Completed)) + .collect(); + completed.sort_by_key(|(job_id, _)| *job_id); + completed } - #[test] - fn given_failing_job_then_failure_reason_reported() { - let finishes = serve(&setup(), vec![dispatched("fail")]); + #[rstest] + #[case::completed("train", FinishStatus::Completed)] + #[case::failed("fail", FinishStatus::Failed)] + #[case::with_an_input_that_does_not_decode("typed", FinishStatus::Failed)] + #[case::unknown("mystery", FinishStatus::Failed)] + fn a_job_is_reported_as_it_ended( + setup: Setup, + #[case] name: &str, + #[case] status: FinishStatus, + ) { + let job = dispatched(name); + + let finishes = serve(&setup, vec![job.clone()]); - assert_eq!(finishes[0].2, FinishStatus::Failed); - assert_eq!(finishes[0].3.as_deref(), Some("boom")); + assert_eq!(statuses(&finishes), [(job.id, status)]); } - #[test] - fn given_input_the_job_rejects_then_failure_reported_as_invalid_input() { - let finishes = serve(&setup(), vec![dispatched("typed")]); + #[rstest] + fn a_failed_job_is_reported_with_its_error_as_the_reason(setup: Setup) { + let job = dispatched("fail"); + + let finishes = serve(&setup, vec![job.clone()]); - assert_eq!(finishes[0].2, FinishStatus::Failed); assert_eq!( - finishes[0].3.as_deref(), - Some("invalid input: invalid type: map, expected u32") + finishes, + [( + job.id, + setup.runner_id, + FinishStatus::Failed, + Some(FAILURE.to_string()) + )] ); } - #[test] - fn given_panicking_job_then_failure_reported_with_panic_message() { - let finishes = serve(&setup(), vec![dispatched("panic")]); + #[rstest] + fn a_panic_while_the_input_decodes_fails_the_job_and_frees_the_runner(setup: Setup) { + let job = dispatched("explosive"); - assert_eq!(finishes[0].2, FinishStatus::Failed); - assert_eq!(finishes[0].3.as_deref(), Some("the job panicked: kaboom")); - } - - #[test] - fn given_unknown_job_name_then_failure_reported() { - let finishes = serve(&setup(), vec![dispatched("mystery")]); + let finishes = serve(&setup, vec![job.clone()]); - assert_eq!(finishes[0].2, FinishStatus::Failed); - assert_eq!(finishes[0].3.as_deref(), Some("unknown job 'mystery'")); + assert_eq!(statuses(&finishes), [(job.id, FinishStatus::Failed)]); } #[test] fn a_running_job_is_numbered_with_its_station_job_id() { - let setup = setup(); - let job = dispatched("num"); + let seen = Arc::new(OnceLock::new()); + let setup = serving(Jobs::new().job("train", { + let seen = seen.clone(); + move |_: Value| { + let _ = seen.set(job_num()); + Ok::<_, BoxError>(()) + } + })); + let job = dispatched("train"); serve(&setup, vec![job.clone()]); - assert_eq!(*setup.nums.lock().unwrap(), [Some(job.id.to_string())]); + assert_eq!(seen.get(), Some(&Some(job.id.to_string()))); } - #[test] - fn jobs_dispatched_at_once_run_at_once_each_in_its_own_scope() { - let started = Arc::new((Mutex::new(0), Condvar::new())); - let seen = Arc::new(Mutex::new(Vec::new())); - let jobs = Jobs::new().job("wait", { - let (started, seen) = (started.clone(), seen.clone()); - move |_: Value| { - let (count, both) = &*started; - *count.lock().unwrap() += 1; - both.notify_all(); - let (count, _) = both - .wait_timeout_while(count.lock().unwrap(), Duration::from_secs(10), |count| { - *count < 2 - }) - .unwrap(); - seen.lock() - .unwrap() - .push((job_num(), job_input(), *count == 2)); - Ok::<_, BoxError>(()) - } - }); - let sink = Arc::new(RecordingSink::default()); - let executor = Executor::new(Arc::new(jobs), sink.clone()); - let with_input = |n: u32| DispatchedJob { - input: json!({ "n": n }), - ..dispatched("wait") - }; - let (first, second) = (with_input(1), with_input(2)); + #[rstest] + fn jobs_dispatched_at_once_run_at_once(setup: Setup) { + let waiting = [dispatched("wait"), dispatched("wait")]; - serve_stream( - events( - Uuid::new_v4(), - vec![ - RunnerEvent::Job(first.clone()), - RunnerEvent::Job(second.clone()), - ], - ), - &executor, - ); - executor.wait_idle(); - - let seen = seen.lock().unwrap().clone(); - assert_eq!(seen.len(), 2); - for (job, n) in [(&first, 1), (&second, 2)] { - assert!( - seen.contains(&(Some(job.id.to_string()), Some(json!({ "n": n })), true)), - "{seen:?}" - ); - } - assert_eq!(sink.finishes.lock().unwrap().len(), 2); + let finishes = serve(&setup, waiting.to_vec()); + + assert_eq!(statuses(&finishes), completed(&waiting)); } - #[test] - fn a_cancelled_job_runs_to_completion_and_the_stream_goes_on() { - let setup = setup(); + #[rstest] + fn a_cancelled_job_runs_to_completion_and_the_stream_goes_on(setup: Setup) { let (cancelled, next) = (dispatched("train"), dispatched("train")); serve_stream( @@ -430,27 +418,18 @@ mod tests { ); setup.executor.wait_idle(); - let finishes = setup.sink.finishes.lock().unwrap().clone(); - assert_eq!(finishes.len(), 2); - for job in [&cancelled, &next] { - assert!( - finishes.contains(&(job.id, setup.runner_id, FinishStatus::Completed, None)), - "{finishes:?}" - ); - } + assert_eq!(statuses(&finishes(&setup)), completed([&cancelled, &next])); } - #[test] - fn given_stream_without_registered_first_then_nothing_served() { - let setup = setup(); - + #[rstest] + fn a_stream_that_does_not_open_with_its_registration_serves_nothing(setup: Setup) { let outcome = serve_stream( - vec![Ok(RunnerEvent::Job(dispatched("train")))].into_iter(), + vec![RunnerEvent::Job(dispatched("train"))].into_iter(), &setup.executor, ); setup.executor.wait_idle(); assert!(matches!(outcome, StreamOutcome::NeverRegistered)); - assert!(setup.sink.finishes.lock().unwrap().is_empty()); + assert!(finishes(&setup).is_empty()); } } diff --git a/crates/tracel-station/Cargo.toml b/crates/tracel-station/Cargo.toml index fd39cbe0..f59d818d 100644 --- a/crates/tracel-station/Cargo.toml +++ b/crates/tracel-station/Cargo.toml @@ -24,3 +24,6 @@ tracel-experiment = { workspace = true } tracel-experiment-remote = { workspace = true } tracel-models = { workspace = true } url = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } diff --git a/crates/tracel-station/src/experiment/mod.rs b/crates/tracel-station/src/experiment/mod.rs index 16fc7087..2c94fc0d 100644 --- a/crates/tracel-station/src/experiment/mod.rs +++ b/crates/tracel-station/src/experiment/mod.rs @@ -24,6 +24,7 @@ use tracel_experiment_remote::RemoteExperimentSession; use self::artifacts::{StationArtifactReader, StationArtifactUploader}; use crate::station::StationInner; +use crate::wire::station_timestamp; #[derive(Debug, thiserror::Error)] enum RunError { @@ -161,7 +162,7 @@ impl ExperimentArtifactClient { Ok(bundle) } - /// The artifact of the experiment named `name`. + /// The artifact of the experiment named `name`, the newest when several are. pub fn fetch(&self, name: &str) -> Result { let listed = self.client.experiments().list_artifacts( self.exp_path.experiment_num(), @@ -169,26 +170,24 @@ impl ExperimentArtifactClient { name: Some(name.to_string()), }, )?; - named(listed.items, name, |artifact| &artifact.name) + newest_named(listed.items, name) } } -/// The one item of `items` whose name, which `name_of` gives, is `name`. +/// The newest of `artifacts` named exactly `name`. /// /// The Station's name filter lists the artifacts whose names contain the name asked for, so -/// asking for `model` can list `model-2` as well. -fn named(items: Vec, name: &str, name_of: impl Fn(&T) -> &str) -> Result { - let mut matching = items.into_iter().filter(|item| name_of(item) == name); - let found = matching - .next() - .ok_or_else(|| ArtifactError::NotFound(name.to_owned()))?; - match matching.count() { - 0 => Ok(found), - others => Err(ArtifactError::Ambiguous { - name: name.to_owned(), - count: others + 1, - }), - } +/// asking for `model` can list `model-2` as well, and an experiment can hold several artifacts of +/// one name. +fn newest_named( + artifacts: Vec, + name: &str, +) -> Result { + artifacts + .into_iter() + .filter(|artifact| artifact.name == name) + .max_by_key(|artifact| station_timestamp(&artifact.created_at)) + .ok_or_else(|| ArtifactError::NotFound(name.to_owned())) } fn artifact_kind_name(kind: ArtifactKind) -> &'static str { @@ -203,8 +202,6 @@ fn artifact_kind_name(kind: ArtifactKind) -> &'static str { pub enum ArtifactError { #[error("Artifact not found: {0}")] NotFound(String), - #[error("Artifact name is ambiguous: {count} artifacts are named {name}")] - Ambiguous { name: String, count: usize }, #[error(transparent)] Client(#[from] ClientError), #[error(transparent)] @@ -266,43 +263,59 @@ fn create_run( #[cfg(test)] mod tests { + use rstest::{fixture, rstest}; + use super::*; - /// The id of the artifact `named` picks among `(id, name)` pairs. - fn pick( - artifacts: &[(&'static str, &'static str)], - name: &str, - ) -> Result<&'static str, ArtifactError> { - named(artifacts.to_vec(), name, |(_, name)| name).map(|(id, _)| id) + fn artifact(name: &str, created_at: &str) -> ArtifactResponse { + ArtifactResponse { + id: format!("{name} at {created_at}"), + created_at: created_at.to_string(), + name: name.to_string(), + kind: "model".to_string(), + experiment_num: 1, + manifest: Value::Null, + } } - #[test] - fn an_artifact_is_picked_by_its_exact_name() { - let listed = [("2", "model-2"), ("1", "model"), ("3", "model.bpk")]; + #[fixture] + fn listed() -> Vec { + vec![ + artifact("model-2", "2026-10-01T00:00:00Z"), + artifact("model", "2026-10-01T00:00:00Z"), + artifact("model.bpk", "2026-10-01T00:00:00Z"), + ] + } - assert_eq!(pick(&listed, "model").unwrap(), "1"); - assert_eq!(pick(&listed, "model-2").unwrap(), "2"); + #[rstest] + #[case::a_name_others_extend("model")] + #[case::a_name_extending_another("model-2")] + fn an_artifact_is_picked_by_its_exact_name(listed: Vec, #[case] name: &str) { + assert_eq!(newest_named(listed, name).unwrap().name, name); } - #[test] - fn a_name_no_artifact_has_exactly_is_not_found() { - let error = pick(&[("2", "model-2")], "model").unwrap_err(); + #[rstest] + #[case::only_longer_names_listed(vec![artifact("model-2", "2026-10-01T00:00:00Z")])] + #[case::nothing_listed(vec![])] + fn a_name_no_artifact_has_exactly_is_not_found(#[case] listed: Vec) { + let asked = "model"; - assert!(matches!(&error, ArtifactError::NotFound(name) if name == "model")); assert!(matches!( - pick(&[], "model"), - Err(ArtifactError::NotFound(_)) + newest_named(listed, asked), + Err(ArtifactError::NotFound(name)) if name == asked )); } #[test] - fn a_name_several_artifacts_have_is_ambiguous() { - let error = pick(&[("1", "model"), ("2", "model-2"), ("3", "model")], "model").unwrap_err(); - - assert!(matches!(&error, ArtifactError::Ambiguous { count: 2, .. })); - assert_eq!( - error.to_string(), - "Artifact name is ambiguous: 2 artifacts are named model" - ); + fn the_newest_of_same_named_artifacts_is_picked() { + let newest = artifact("model", "2026-10-03T00:00:00Z"); + let listed = vec![ + artifact("model", "2026-10-01T00:00:00Z"), + newest.clone(), + artifact("model-2", "2026-10-04T00:00:00Z"), + artifact("model", "2026-10-02T00:00:00Z"), + ]; + + assert_eq!(newest_named(listed, "model").unwrap().id, newest.id); } } diff --git a/crates/tracel/Cargo.toml b/crates/tracel/Cargo.toml index ffb464f7..652822a1 100644 --- a/crates/tracel/Cargo.toml +++ b/crates/tracel/Cargo.toml @@ -13,7 +13,7 @@ rust-version.workspace = true [features] default = [] -station = ["dep:tracel-station", "dep:url", "tracel-runtime/station"] +station = ["dep:tracel-station", "tracel-runtime/station"] burn = ["tracel-experiment/burn", "tracel-datasets/burn"] schema = ["tracel-runtime/schema"] @@ -28,8 +28,9 @@ tracel-datasets.workspace = true tracel-models.workspace = true serde_json.workspace = true thiserror.workspace = true -url = { workspace = true, optional = true } +url.workspace = true [dev-dependencies] +rstest.workspace = true serde = { workspace = true, features = ["std"] } tempfile.workspace = true diff --git a/crates/tracel/src/lib.rs b/crates/tracel/src/lib.rs index 9ddc70e7..fed9580b 100644 --- a/crates/tracel/src/lib.rs +++ b/crates/tracel/src/lib.rs @@ -83,12 +83,10 @@ //! //! The experiments [`Target::experiments`] builds in a job, on the thread that runs it, are linked //! to that job, even for runs started on other threads: each one started from them records the -//! job's number and input, and prints where it records to stderr as it starts, such as -//! `Recording experiment 42: `, which `tracel run` and a remote job's logs show. Experiments -//! built another way, such as [`console::ProjectHandle::experiments`], are not linked unless the -//! code links them, through [`Experiments::attribute`](experiment::Experiments::attribute), -//! [`arguments`](experiment::Experiments::arguments) and -//! [`on_start`](experiment::Experiments::on_start). +//! job's number and input. Experiments built another way, such as +//! [`console::ProjectHandle::experiments`], are not linked unless the code links them, through +//! [`Experiments::attribute`](experiment::Experiments::attribute) and +//! [`arguments`](experiment::Experiments::arguments). //! //! [`version`](runtime::Jobs::version) gives the binary its `--version`, which prints the binary's //! name and that version; without it, the binary has no `--version`. @@ -125,11 +123,12 @@ //! - A nested field's flag joins the keys with dots and writes `_` as `-`: the field //! `optimizer.weight_decay` is `--optimizer.weight-decay`. A field that is an array, a map, an //! optional object or an enum with data takes a JSON literal, such as `--layers '[64, 32]'`. A -//! boolean flag given no value is `true`. +//! boolean flag given no value is `true`, and takes a value only as `--flag=value`, such as +//! `--shuffle=false`. //! - The input is the default input, then the `--config` file, then the JSON document, then the //! flags, each merged onto the one before, so a later one wins. Launchers pass the JSON document. -//! - `help`, `version` and `config` are the command line's names: a field with one of them has no -//! flag, and is set through `--config` or the JSON document. +//! - `help` and `config` are the command line's names: a field with one of them has no flag, and is +//! set through `--config` or the JSON document. //! - ` --completions ` prints the completion script for `bash`, `elvish`, `fish`, //! `powershell` or `zsh`. //! diff --git a/crates/tracel/src/running_job.rs b/crates/tracel/src/running_job.rs index 5a85b9ee..8de792b4 100644 --- a/crates/tracel/src/running_job.rs +++ b/crates/tracel/src/running_job.rs @@ -3,70 +3,31 @@ use std::collections::HashMap; use serde_json::Value; -use tracel_experiment::{ExperimentLocation, ExperimentRun, Experiments}; +use tracel_experiment::Experiments; use tracel_runtime::JobScope; /// The experiment attribute that links an experiment to the job that started it: the job's /// number, as a number when it is one. const JOB_NUM_ATTRIBUTE: &str = "tracel.job_num"; -/// `experiments`, linked to the job the current thread runs, when it runs one, as [`link`] links -/// them, each run started from them printing where it records to stderr; as they are otherwise. -/// The link holds on any thread a run is started on. +/// `experiments`, linked to the job the current thread runs, when it runs one; as they are +/// otherwise. +/// +/// Each experiment started from linked experiments, on any thread, belongs to the job: it records +/// the job's number as its `tracel.job_num` attribute and the job's input, unless it is `null`, as +/// its arguments. pub fn linked(experiments: Experiments) -> Experiments { - match RunningJob::current() { - Some(job) => link(experiments, job, |line| eprintln!("{line}")), - None => experiments, - } -} - -/// What a job gives the experiments started in it. -struct RunningJob { - /// The number the job's launcher gave it. - num: Option, - /// The input the job runs with, as JSON. - input: Value, -} - -impl RunningJob { - /// The job the current thread runs, as its [`JobScope`] gives it, while it runs one. - fn current() -> Option { - let scope = JobScope::current()?; - Some(Self { - num: scope.num().map(str::to_string), - input: scope.input().clone(), - }) - } -} - -/// `experiments`, each started from which belongs to `job`: it records the job's number as its -/// `tracel.job_num` attribute and the job's input, unless it is `null`, as its arguments, and it -/// gives `announce` the line saying where it records as it starts. -fn link( - experiments: Experiments, - job: RunningJob, - announce: impl Fn(&str) + Send + Sync + 'static, -) -> Experiments { + let Some(job) = JobScope::current() else { + return experiments; + }; let attributes = job - .num - .map(|num| (JOB_NUM_ATTRIBUTE.to_string(), job_num(&num))); - let experiments = match job.input { + .num() + .map(|num| (JOB_NUM_ATTRIBUTE.to_string(), job_num(num))); + let experiments = match job.input() { Value::Null => experiments, - input => experiments.arguments(input), + input => experiments.arguments(input.clone()), }; - experiments - .attributes(HashMap::from_iter(attributes)) - .on_start(move |run| announce(&recording(run))) -} - -/// The line saying where `run` records: its number and its page on the console, or its directory -/// offline. -fn recording(run: &ExperimentRun) -> String { - match run.location() { - Some(ExperimentLocation::Url(url)) => format!("Recording experiment {}: {url}", run.id()), - Some(ExperimentLocation::Dir(dir)) => format!("Recording offline in {}", dir.display()), - None => format!("Recording experiment {}", run.id()), - } + experiments.attributes(HashMap::from_iter(attributes)) } /// The attribute value of the job number `num`: a number when it is one, its text otherwise. @@ -79,139 +40,166 @@ fn job_num(num: &str) -> Value { #[cfg(test)] mod tests { - use std::path::Path; use std::sync::{Arc, Mutex}; + use rstest::{fixture, rstest}; use serde_json::json; - use tracel_experiment::ExperimentProvider; - use tracel_experiment::local::LocalExperiments; + use tracel_experiment::error::ExperimentError; + use tracel_experiment::reader::{ + ExperimentArtifactReader, ExperimentReaderError, LoadedArtifact, + }; + use tracel_experiment::session::{BundleFn, Event, ExperimentCompletion, ExperimentSession}; + use tracel_experiment::{ + ArtifactKind, CancelToken, ExperimentId, ExperimentProvider, ExperimentRun, + }; + use tracel_runtime::{BoxError, Execution, Jobs}; use super::*; - /// What a test keeps of what it is given. - type Kept = Arc>>; + const ID: &str = "7"; - /// A job numbered `num`, running with `input`. - fn job(num: Option<&str>, input: Value) -> RunningJob { - RunningJob { - num: num.map(str::to_string), - input, - } + #[derive(Debug, Default, PartialEq)] + struct Recorded { + attributes: HashMap, + arguments: Option, } - fn offline(dir: &Path) -> Experiments { - Experiments::new(Arc::new(LocalExperiments::new(dir))) - } + #[derive(Clone, Default)] + struct Recorder(Arc>); + + impl Recorder { + fn experiments(&self) -> Experiments { + Experiments::new(Arc::new(self.clone())) + } - /// `experiments` linked to `job`, and the lines its runs announce. - fn linked_to(experiments: Experiments, job: RunningJob) -> (Experiments, Kept) { - let announced = Arc::new(Mutex::new(Vec::new())); - let experiments = link(experiments, job, { - let announced = announced.clone(); - move |line| announced.lock().unwrap().push(line.to_string()) - }); - (experiments, announced) + fn recorded(&self) -> Recorded { + std::mem::take(&mut self.0.lock().unwrap()) + } } - /// The lines of the events file of run `num` of `train`, under `dir`. - fn events(dir: &Path, num: u64) -> Vec { - std::fs::read_to_string(dir.join(format!("train/{num}/events.jsonl"))) - .unwrap() - .lines() - .map(|line| serde_json::from_str(line).unwrap()) - .collect() + impl ExperimentProvider for Recorder { + fn create_experiment( + &self, + _: String, + attributes: HashMap, + ) -> Result { + self.0.lock().unwrap().attributes = attributes; + Ok(ExperimentRun::new( + ID, + self.clone(), + self.clone(), + CancelToken::new(), + )) + } } - #[test] - fn an_experiment_started_in_a_job_belongs_to_it() { - let dir = tempfile::tempdir().unwrap(); - let job = job(Some("12"), json!({"epochs": 2})); - - let (experiments, announced) = linked_to(offline(dir.path()), job); - experiments.start("train").unwrap().finish().unwrap(); - - let runs = dir.path().canonicalize().unwrap(); - assert!(events(dir.path(), 1).contains(&json!({ - "type": "attribute", - "key": "tracel.job_num", - "value": 12 - }))); - assert!(events(dir.path(), 1).contains(&json!({ - "type": "arguments", - "value": {"epochs": 2} - }))); - assert_eq!( - *announced.lock().unwrap(), - [format!( - "Recording offline in {}", - runs.join("train/1").display() - )] - ); + impl ExperimentSession for Recorder { + fn record_event(&self, event: Event) -> Result<(), ExperimentError> { + if let Event::Args(arguments) = event { + self.0.lock().unwrap().arguments = Some(arguments); + } + Ok(()) + } + + fn save_artifact( + &self, + _: &str, + _: ArtifactKind, + _: Box, + ) -> Result<(), ExperimentError> { + Ok(()) + } + + fn finish(&self, _: ExperimentCompletion) -> Result<(), ExperimentError> { + Ok(()) + } } - #[test] - fn an_experiment_started_in_a_job_without_a_number_or_input_records_neither() { - let dir = tempfile::tempdir().unwrap(); - let job = job(None, Value::Null); - let (experiments, announced) = linked_to(offline(dir.path()), job); + impl ExperimentArtifactReader for Recorder { + fn load_artifact_raw( + &self, + _: ExperimentId, + _: &str, + ) -> Result { + Err(ExperimentReaderError::new("a recorder keeps no artifacts")) + } + } - experiments.start("train").unwrap().finish().unwrap(); + type Start = fn(Experiments) -> Result<(), BoxError>; - assert!( - events(dir.path(), 1) - .iter() - .all(|event| event["type"] != "attribute" && event["type"] != "arguments") - ); - assert_eq!(announced.lock().unwrap().len(), 1); + fn start(experiments: Experiments) -> Result<(), BoxError> { + experiments.start("train")?.finish()?; + Ok(()) } - #[test] - fn outside_a_job_experiments_are_not_linked() { - let dir = tempfile::tempdir().unwrap(); + fn start_on_another_thread(experiments: Experiments) -> Result<(), BoxError> { + std::thread::spawn(move || start(experiments)) + .join() + .unwrap() + } - linked(offline(dir.path())) - .start("train") + fn recorded_in_a_job(execution: Execution, input: Value, start: Start) -> Recorded { + let recorder = Recorder::default(); + let train = { + let recorder = recorder.clone(); + move |_: Value| start(linked(recorder.experiments())) + }; + Jobs::new() + .job("train", train) + .prepare("train", input) .unwrap() - .finish() + .run(execution) + .into_result() .unwrap(); - - assert!( - events(dir.path(), 1) - .iter() - .all(|event| event["type"] != "attribute" && event["type"] != "arguments") - ); + recorder.recorded() } - #[test] - fn a_job_number_is_recorded_as_a_number_when_it_is_one() { - assert_eq!(job_num("42"), json!(42)); - assert_eq!(job_num("job-42"), json!("job-42")); + #[fixture] + fn input() -> Value { + json!({"epochs": 2}) } - #[test] - fn an_experiment_announces_its_page_or_its_directory() { - let dir = tempfile::tempdir().unwrap(); - let experiments = LocalExperiments::new(dir.path()); - let page = "https://console.tracel.ai/users/me/projects/mnist/experiments/2"; + #[rstest] + #[case::on_the_jobs_thread(start)] + #[case::on_a_thread_the_job_spawns(start_on_another_thread)] + fn an_experiment_started_in_a_job_belongs_to_it(#[case] start: Start, input: Value) { + let num = "12"; - let offline = experiments - .create_experiment("mnist".to_string(), HashMap::new()) - .unwrap(); - let on_the_console = experiments - .create_experiment("mnist".to_string(), HashMap::new()) - .unwrap() - .with_location(ExperimentLocation::Url(page.to_string())); + let recorded = recorded_in_a_job(Execution::new().with_num(num), input.clone(), start); assert_eq!( - recording(&offline), - format!( - "Recording offline in {}", - dir.path().canonicalize().unwrap().join("mnist/1").display() - ) + recorded, + Recorded { + attributes: HashMap::from([(JOB_NUM_ATTRIBUTE.to_string(), job_num(num))]), + arguments: Some(input), + } ); + } + + #[test] + fn an_experiment_started_in_a_job_without_a_number_or_input_records_neither() { assert_eq!( - recording(&on_the_console), - format!("Recording experiment 2: {page}") + recorded_in_a_job(Execution::new(), Value::Null, start), + Recorded::default() ); } + + #[test] + fn outside_a_job_experiments_are_not_linked() { + let recorder = Recorder::default(); + + start(linked(recorder.experiments())).unwrap(); + + assert_eq!(recorder.recorded(), Recorded::default()); + } + + #[rstest] + #[case::digits("42", json!(42))] + #[case::not_only_digits("job-42", json!("job-42"))] + fn a_job_number_is_recorded_as_a_number_when_it_is_one( + #[case] num: &str, + #[case] recorded: Value, + ) { + assert_eq!(job_num(num), recorded); + } } diff --git a/crates/tracel/src/target.rs b/crates/tracel/src/target.rs index 6194a9e6..a9d4924c 100644 --- a/crates/tracel/src/target.rs +++ b/crates/tracel/src/target.rs @@ -1,8 +1,8 @@ use std::ffi::OsString; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; -use tracel_console::{Console, ConsoleError, CredentialSource, Env, ProjectHandle, ProjectRef}; +use tracel_console::{ConsoleError, FileSessionStore, ProjectHandle, SessionStoreError}; use tracel_datasets::DatasetRegistry; use tracel_experiment::Experiments; use tracel_experiment::local::LocalExperiments; @@ -10,16 +10,12 @@ use tracel_inference::{InferenceModule, NoopInferenceProvider}; use tracel_models::ModelRegistry; #[cfg(feature = "station")] use tracel_station::Station; -#[cfg(feature = "station")] use url::Url; use crate::running_job; const TRACEL_TARGET: &str = "TRACEL_TARGET"; const TRACEL_RUNS_DIR: &str = "TRACEL_RUNS_DIR"; -const TRACEL_API_KEY: &str = "TRACEL_API_KEY"; -const TRACEL_NAMESPACE: &str = "TRACEL_NAMESPACE"; -const TRACEL_PROJECT: &str = "TRACEL_PROJECT"; #[cfg(feature = "station")] const TRACEL_STATION_URL: &str = "TRACEL_STATION_URL"; @@ -38,8 +34,7 @@ const TARGETS: &str = "offline, console or station"; /// /// Build one directly, or read it from the environment with [`Target::from_env`]. Neither /// performs network I/O, and neither do the services a target builds: the first request that -/// needs a server reaches it. Each call to a service method binds a new connection; to share one -/// across services, build them from a [`ProjectHandle`] instead. +/// needs a server reaches it. #[derive(Debug, Clone)] #[non_exhaustive] pub enum Target { @@ -48,20 +43,18 @@ pub enum Target { /// The directory runs are recorded under. dir: PathBuf, }, - /// A project on the Tracel console. + /// A project on the Tracel console. The services the target builds share its connection, so + /// the credential is verified once. Console { - /// The console to reach. - env: Env, - /// Where the credential comes from. - credentials: CredentialSource, /// The project experiments, models and datasets belong to. - project: ProjectRef, + project: ProjectHandle, }, - /// A Tracel Station (requires the `station` feature). + /// A Tracel Station (requires the `station` feature). The services the target builds share + /// its client. #[cfg(feature = "station")] Station { - /// The Station's base URL. - url: Url, + /// The Station experiments, models and datasets belong to. + station: Station, }, } @@ -89,16 +82,26 @@ impl Target { doc = "| `TRACEL_STATION_URL` | the Station's base URL | `http://localhost:8000` |" )] /// - /// Fails with [`TargetError::UnknownTarget`] when `TRACEL_TARGET` names no target, and with - /// [`TargetError::MissingSetting`] or [`TargetError::InvalidSetting`] naming the variable to - /// set when the target cannot be filled in. + /// The console target is the project [`ProjectHandle::from_env`] reads, and fails as it + /// does, with [`TargetError::Console`]. Fails with [`TargetError::UnknownTarget`] when + /// `TRACEL_TARGET` names no target, and with [`TargetError::InvalidSetting`] when + /// `TRACEL_STATION_URL` is not an http or https URL. pub fn from_env() -> Result { - Self::from_vars(|name| std::env::var_os(name)) + Self::from_vars( + |name| std::env::var_os(name), + Path::new("."), + FileSessionStore::for_server, + ) } - /// [`from_env`](Self::from_env) over the values `lookup` gives this crate's variables. The - /// console's own variables are read by [`tracel_console`]. - fn from_vars(lookup: impl Fn(&str) -> Option) -> Result { + /// [`from_env`](Self::from_env), reading the variables from `lookup`, `tracel.toml` from + /// `dir`, and the stored `tracel login` session from the store `sessions` gives for the + /// console's URL. + fn from_vars( + lookup: impl Fn(&str) -> Option, + dir: &Path, + sessions: impl FnOnce(&Url) -> Result, + ) -> Result { let var = |name: &str| lookup(name).filter(|value| !value.is_empty()); let target = var(TRACEL_TARGET).map(|value| value.to_string_lossy().into_owned()); @@ -107,27 +110,12 @@ impl Target { dir: var(TRACEL_RUNS_DIR).map_or_else(|| DEFAULT_RUNS_DIR.into(), PathBuf::from), }), Some("console") => Ok(Target::Console { - env: tracel_console::env_from_environment()?, - credentials: CredentialSource::from_env(), - project: ProjectRef::from_env()?, + project: ProjectHandle::from_vars(&lookup, dir, sessions)?, }), #[cfg(feature = "station")] - Some("station") => { - let url = var(TRACEL_STATION_URL).map_or_else( - || DEFAULT_STATION_URL.to_string(), - |value| value.to_string_lossy().into_owned(), - ); - match Url::parse(&url) { - Ok(url) if matches!(url.scheme(), "http" | "https") => { - Ok(Target::Station { url }) - } - _ => Err(TargetError::InvalidSetting { - variable: TRACEL_STATION_URL, - value: url, - expected: "an http or https URL, such as http://localhost:8000", - }), - } - } + Some("station") => Ok(Target::Station { + station: Station::connect(station_url(var(TRACEL_STATION_URL))?), + }), #[cfg(not(feature = "station"))] Some("station") => Err(TargetError::FeatureRequired { target: "station", @@ -147,22 +135,15 @@ impl Target { /// Built in a job, on the thread its adapter runs it on, the experiments are linked to that /// job, as its [`JobScope`](crate::runtime::JobScope) gives it, wherever runs are started from /// them: each experiment records the job's number, when its launcher gave one, as its - /// `tracel.job_num` attribute and the job's input as its arguments, and prints where it - /// records to stderr as it starts, such as `Recording experiment 42: ` or - /// `Recording offline in `. Experiments built another way, such as - /// [`ProjectHandle::experiments`] or [`Experiments::new`], are not linked unless the code links - /// them through [`Experiments::attribute`], [`Experiments::arguments`] and - /// [`Experiments::on_start`]. + /// `tracel.job_num` attribute and the job's input as its arguments. Experiments built another + /// way, such as [`ProjectHandle::experiments`] or [`Experiments::new`], are not linked unless + /// the code links them through [`Experiments::attribute`] and [`Experiments::arguments`]. pub fn experiments(&self) -> Result { let experiments = match self { Target::Offline { dir } => Experiments::new(Arc::new(LocalExperiments::new(dir))), - Target::Console { - env, - credentials, - project, - } => console_project(env, credentials, project)?.experiments(), + Target::Console { project } => project.experiments(), #[cfg(feature = "station")] - Target::Station { url } => Station::connect(url.clone()).experiments(), + Target::Station { station } => station.experiments(), }; Ok(running_job::linked(experiments)) } @@ -176,13 +157,9 @@ impl Target { service: "models", target: "offline", }), - Target::Console { - env, - credentials, - project, - } => Ok(console_project(env, credentials, project)?.models()), + Target::Console { project } => Ok(project.models()), #[cfg(feature = "station")] - Target::Station { url } => Ok(Station::connect(url.clone()).models()), + Target::Station { station } => Ok(station.models()), } } @@ -195,13 +172,9 @@ impl Target { service: "datasets", target: "offline", }), - Target::Console { - env, - credentials, - project, - } => Ok(console_project(env, credentials, project)?.datasets()), + Target::Console { project } => Ok(project.datasets()), #[cfg(feature = "station")] - Target::Station { url } => Ok(Station::connect(url.clone()).datasets()), + Target::Station { station } => Ok(station.datasets()), } } @@ -213,25 +186,30 @@ impl Target { pub fn inference(&self) -> Result { match self { Target::Offline { .. } => Ok(noop_inference()), - Target::Console { - env, - credentials, - project, - } => Ok(console_project(env, credentials, project)?.inference()), + Target::Console { project } => Ok(project.inference()), #[cfg(feature = "station")] Target::Station { .. } => Ok(noop_inference()), } } } -/// Binds `project` on the console `env` names, with the credential `credentials` provides. -fn console_project( - env: &Env, - credentials: &CredentialSource, - project: &ProjectRef, -) -> Result { - let console = Console::connect(env.clone(), credentials.resolve(env)?); - Ok(console.project(&project.namespace, &project.name)) +/// The Station URL `value` names, `http://localhost:8000` when it is unset. +/// +/// Fails with [`TargetError::InvalidSetting`] when it is not an http or https URL. +#[cfg(feature = "station")] +fn station_url(value: Option) -> Result { + let url = value.map_or_else( + || DEFAULT_STATION_URL.to_string(), + |value| value.to_string_lossy().into_owned(), + ); + match Url::parse(&url) { + Ok(url) if matches!(url.scheme(), "http" | "https") => Ok(url), + _ => Err(TargetError::InvalidSetting { + variable: TRACEL_STATION_URL, + value: url, + expected: "an http or https URL, such as http://localhost:8000", + }), + } } fn noop_inference() -> InferenceModule { @@ -258,13 +236,7 @@ pub enum TargetError { /// The `tracel` feature that adds the target. feature: &'static str, }, - /// A setting the target needs is not set. - #[error("{variable} is not set{}", fallback(.variable))] - MissingSetting { - /// The environment variable to set. - variable: &'static str, - }, - /// An environment variable holds a value the SDK does not accept. + /// `TRACEL_STATION_URL` holds a value that is not an http or https URL. #[error("invalid {variable} value `{value}`: expected {expected}")] InvalidSetting { /// The environment variable. @@ -282,252 +254,166 @@ pub enum TargetError { /// The kind of target that does not serve it. target: &'static str, }, - /// The console settings could not be read for another reason, such as a stored sign-in - /// that cannot be read. + /// The console target cannot be filled in or reached, as the console says. #[error(transparent)] - Console(ConsoleError), -} - -/// What else provides a setting whose variable is unset, as [`TargetError::MissingSetting`] says. -fn fallback(variable: &str) -> &'static str { - match variable { - TRACEL_API_KEY => ", and there is no `tracel login` sign-in for this console", - TRACEL_NAMESPACE => { - ", and tracel.toml sets no `namespace` (expected keys: `namespace` and `project`)" - } - TRACEL_PROJECT => { - ", and tracel.toml sets no `project` (expected keys: `namespace` and `project`)" - } - _ => "", - } -} - -impl From for TargetError { - fn from(error: ConsoleError) -> Self { - match error { - ConsoleError::NoCredentials => Self::MissingSetting { - variable: TRACEL_API_KEY, - }, - ConsoleError::NoNamespace => Self::MissingSetting { - variable: TRACEL_NAMESPACE, - }, - ConsoleError::NoProject => Self::MissingSetting { - variable: TRACEL_PROJECT, - }, - ConsoleError::InvalidSetting { - variable, - value, - expected, - } => Self::InvalidSetting { - variable, - value, - expected, - }, - error => Self::Console(error), - } - } + Console(#[from] ConsoleError), } #[cfg(test)] mod tests { - use std::path::Path; - - use tracel_console::TracelCredentials; + use rstest::rstest; + use tracel_console::ProjectRef; use super::*; fn target(vars: &[(&str, &str)]) -> Result { - Target::from_vars(|name| { - vars.iter() - .find(|(variable, _)| *variable == name) - .map(|(_, value)| OsString::from(value)) - }) - } - - fn offline_dir(target: Result) -> PathBuf { - match target { - Ok(Target::Offline { dir }) => dir, - other => panic!("expected an offline target, got {other:?}"), - } + let dir = tempfile::tempdir().unwrap(); + Target::from_vars( + |name| { + vars.iter() + .find(|(variable, _)| *variable == name) + .map(|(_, value)| OsString::from(value)) + }, + dir.path(), + |_| Ok(FileSessionStore::new(dir.path().join("session.json"))), + ) } - #[test] - fn an_unset_or_empty_target_records_offline_under_runs() { - assert_eq!(offline_dir(target(&[])), Path::new("./runs")); - assert_eq!( - offline_dir(target(&[(TRACEL_TARGET, "")])), - Path::new("./runs") - ); - assert_eq!( - offline_dir(target(&[(TRACEL_TARGET, "offline")])), - Path::new("./runs") - ); + #[rstest] + #[case::unset(&[])] + #[case::empty(&[(TRACEL_TARGET, "")])] + #[case::offline(&[(TRACEL_TARGET, "offline")])] + #[case::runs_dir_empty(&[(TRACEL_RUNS_DIR, "")])] + fn an_offline_target_records_under_the_default_runs_dir(#[case] vars: &[(&str, &str)]) { + assert!(matches!( + target(vars), + Ok(Target::Offline { dir }) if dir == Path::new(DEFAULT_RUNS_DIR) + )); } #[test] fn tracel_runs_dir_moves_offline_runs() { - assert_eq!( - offline_dir(target(&[(TRACEL_RUNS_DIR, "/data/runs")])), - Path::new("/data/runs") - ); - assert_eq!( - offline_dir(target(&[(TRACEL_RUNS_DIR, "")])), - Path::new("./runs") - ); + let runs = "/data/runs"; + + assert!(matches!( + target(&[(TRACEL_RUNS_DIR, runs)]), + Ok(Target::Offline { dir }) if dir == Path::new(runs) + )); } - #[test] - fn an_unknown_target_names_the_value_and_the_accepted_ones() { - for value in ["cloud", "Offline", "CONSOLE", " console"] { - let error = target(&[(TRACEL_TARGET, value)]).unwrap_err(); - - assert!( - matches!( - &error, - TargetError::UnknownTarget { value: held, .. } if held == value - ), - "{value}: {error:?}" - ); - let message = error.to_string(); - assert!(message.contains("TRACEL_TARGET"), "{message}"); - assert!(message.contains(value), "{message}"); - assert!(message.contains("offline"), "{message}"); - assert!(message.contains("console"), "{message}"); - } + #[rstest] + #[case::unknown("cloud")] + #[case::capitalized("Offline")] + #[case::upper_case("CONSOLE")] + #[case::padded(" console")] + fn an_unknown_target_is_an_error_naming_it(#[case] value: &str) { + assert!(matches!( + target(&[(TRACEL_TARGET, value)]), + Err(TargetError::UnknownTarget { value: held, .. }) if held == value + )); } #[cfg(not(feature = "station"))] #[test] - fn a_station_target_says_it_needs_the_station_feature() { - let error = target(&[(TRACEL_TARGET, "station")]).unwrap_err(); - + fn a_station_target_needs_the_station_feature() { assert!(matches!( - error, - TargetError::FeatureRequired { + target(&[(TRACEL_TARGET, "station")]), + Err(TargetError::FeatureRequired { feature: "station", .. - } + }) )); - assert!(error.to_string().contains("`station` feature"), "{error}"); } #[cfg(feature = "station")] #[test] - fn a_station_target_reads_tracel_station_url() { - let station_url = |vars: &[(&str, &str)]| match target(vars) { - Ok(Target::Station { url }) => url.to_string(), - other => panic!("expected a station target, got {other:?}"), - }; + fn tracel_target_station_reaches_a_station() { + assert!(matches!( + target(&[(TRACEL_TARGET, "station")]), + Ok(Target::Station { .. }) + )); + } + #[cfg(feature = "station")] + #[rstest] + #[case::unset(None, DEFAULT_STATION_URL)] + #[case::set(Some("https://station.example.com"), "https://station.example.com")] + fn tracel_station_url_names_the_station(#[case] value: Option<&str>, #[case] url: &str) { assert_eq!( - station_url(&[(TRACEL_TARGET, "station")]), - "http://localhost:8000/" - ); - assert_eq!( - station_url(&[ - (TRACEL_TARGET, "station"), - (TRACEL_STATION_URL, "https://station.example.com") - ]), - "https://station.example.com/" + station_url(value.map(OsString::from)).unwrap(), + Url::parse(url).unwrap() ); } #[cfg(feature = "station")] - #[test] - fn an_invalid_station_url_names_the_variable() { - for value in ["not a url", "localhost:8000", "ftp://localhost:8000"] { - let error = - target(&[(TRACEL_TARGET, "station"), (TRACEL_STATION_URL, value)]).unwrap_err(); - - assert!( - matches!( - &error, - TargetError::InvalidSetting { variable: "TRACEL_STATION_URL", value: held, .. } - if held == value - ), - "{value}: {error:?}" - ); - assert!(error.to_string().contains("TRACEL_STATION_URL"), "{error}"); - } + #[rstest] + #[case::not_a_url("not a url")] + #[case::not_http("ftp://localhost:8000")] + fn a_station_url_that_is_not_http_is_an_invalid_setting_naming_it(#[case] value: &str) { + assert!(matches!( + target(&[(TRACEL_TARGET, "station"), (TRACEL_STATION_URL, value)]), + Err(TargetError::InvalidSetting { variable: TRACEL_STATION_URL, value: held, .. }) + if held == value + )); } #[test] - fn an_offline_target_serves_experiments_and_inference_without_touching_the_disk() { + fn building_offline_experiments_and_inference_creates_no_directory() { let root = tempfile::tempdir().unwrap(); let dir = root.path().join("runs"); let offline = Target::Offline { dir: dir.clone() }; - assert!(offline.experiments().is_ok()); - assert!(offline.inference().is_ok()); + let _ = (offline.experiments(), offline.inference()); + assert!(!dir.exists()); } - #[test] - fn an_offline_target_does_not_serve_models_or_datasets() { + #[rstest] + #[case::models("models", |offline: &Target| offline.models().err())] + #[case::datasets("datasets", |offline: &Target| offline.datasets().err())] + fn an_offline_target_serves_no_registry( + #[case] service: &str, + #[case] build: fn(&Target) -> Option, + ) { let offline = Target::Offline { - dir: PathBuf::from("./runs"), + dir: PathBuf::from(DEFAULT_RUNS_DIR), }; assert!(matches!( - offline.models(), - Err(TargetError::Unsupported { - service: "models", - target: "offline" - }) - )); - assert!(matches!( - offline.datasets(), - Err(TargetError::Unsupported { - service: "datasets", - target: "offline" - }) + build(&offline), + Some(TargetError::Unsupported { service: held, target: "offline" }) if held == service )); } #[test] - fn a_console_target_builds_its_services_without_a_request() { - let console = Target::Console { - env: Env::Development, - credentials: CredentialSource::Explicit(TracelCredentials::api_key( - "tcl_key_not_checked_until_the_first_request", - )), - project: ProjectRef::new("owner", "project"), - }; + fn a_console_target_is_the_project_its_variables_name() { + let mnist = ProjectRef::new("alice", "mnist"); - assert!(console.experiments().is_ok()); - assert!(console.models().is_ok()); - assert!(console.datasets().is_ok()); - assert!(console.inference().is_ok()); - } + let console = target(&[ + (TRACEL_TARGET, "console"), + ("TRACEL_API_KEY", "tcl_key"), + ("TRACEL_NAMESPACE", mnist.namespace.as_str()), + ("TRACEL_PROJECT", mnist.name.as_str()), + ]); - #[test] - fn a_missing_console_setting_names_the_variable_to_set() { - for (error, variable, fallback) in [ - (ConsoleError::NoCredentials, TRACEL_API_KEY, "tracel login"), - (ConsoleError::NoNamespace, TRACEL_NAMESPACE, "tracel.toml"), - (ConsoleError::NoProject, TRACEL_PROJECT, "tracel.toml"), - ] { - let error = TargetError::from(error); - - assert!( - matches!(error, TargetError::MissingSetting { variable: held } if held == variable), - "{error:?}" - ); - assert!(error.to_string().contains(variable), "{error}"); - assert!(error.to_string().contains(fallback), "{error}"); - } + assert!(matches!( + console, + Ok(Target::Console { project }) + if ProjectRef::new(project.owner(), project.name()) == mnist + )); } #[test] - fn a_missing_namespace_or_project_names_the_keys_of_tracel_toml() { - for (error, key) in [ - (ConsoleError::NoNamespace, "no `namespace`"), - (ConsoleError::NoProject, "no `project`"), - ] { - let message = TargetError::from(error).to_string(); - - assert!(message.contains(key), "{message}"); - assert!(message.contains("`namespace` and `project`"), "{message}"); - } + fn a_console_target_without_a_credential_fails_as_the_console_does() { + let console = target(&[ + (TRACEL_TARGET, "console"), + ("TRACEL_NAMESPACE", "alice"), + ("TRACEL_PROJECT", "mnist"), + ]); + + assert!(matches!( + console, + Err(TargetError::Console(ConsoleError::NoCredentials)) + )); } }