From 4c9d7122f91834ed0c7bd26f27c23181f3731377 Mon Sep 17 00:00:00 2001 From: MatthewMckee4 Date: Wed, 6 May 2026 08:02:51 +0100 Subject: [PATCH] Add --partition slice:M/N for round-robin test slicing Closes #587. Slices the collected test set after sorting by qualified name, keeping only tests at positions matching slice M of N. Validation on the M/N values lives in PartitionSelection's FromStr. --- crates/karva/src/commands/test/mod.rs | 2 + crates/karva/tests/it/main.rs | 1 + crates/karva/tests/it/partition.rs | 138 ++++++++++++++++++++++ crates/karva_benchmark/src/lib.rs | 1 + crates/karva_cli/src/lib.rs | 2 + crates/karva_cli/src/partition.rs | 142 +++++++++++++++++++++++ crates/karva_cli/src/test.rs | 16 +++ crates/karva_runner/src/orchestration.rs | 5 +- crates/karva_runner/src/partition.rs | 16 +++ docs/reference/cli.md | 5 +- docs/usage/running-tests/parallel.md | 12 ++ 11 files changed, 338 insertions(+), 2 deletions(-) create mode 100644 crates/karva/tests/it/partition.rs create mode 100644 crates/karva_cli/src/partition.rs diff --git a/crates/karva/src/commands/test/mod.rs b/crates/karva/src/commands/test/mod.rs index 45ec6016..717b64aa 100644 --- a/crates/karva/src/commands/test/mod.rs +++ b/crates/karva/src/commands/test/mod.rs @@ -46,6 +46,7 @@ pub fn test(args: TestCommand) -> Result { let watch = args.watch; let durations = args.durations; let last_failed = args.last_failed; + let partition = args.partition; let no_cache = args.no_cache.unwrap_or(false); let num_workers = if args.no_parallel.unwrap_or(false) || args.no_capture { 1 @@ -76,6 +77,7 @@ pub fn test(args: TestCommand) -> Result { create_ctrlc_handler: true, last_failed, profile, + partition, }; if watch { diff --git a/crates/karva/tests/it/main.rs b/crates/karva/tests/it/main.rs index c267a0a2..6e923fd2 100644 --- a/crates/karva/tests/it/main.rs +++ b/crates/karva/tests/it/main.rs @@ -10,6 +10,7 @@ mod durations; mod extensions; mod filterset; mod last_failed; +mod partition; mod run_ignored; mod version; mod watch; diff --git a/crates/karva/tests/it/partition.rs b/crates/karva/tests/it/partition.rs new file mode 100644 index 00000000..bb2fcbb6 --- /dev/null +++ b/crates/karva/tests/it/partition.rs @@ -0,0 +1,138 @@ +use insta_cmd::assert_cmd_snapshot; + +use crate::common::TestContext; + +const SIX_TESTS: &str = " +def test_a(): pass +def test_b(): pass +def test_c(): pass +def test_d(): pass +def test_e(): pass +def test_f(): pass +"; + +#[test] +fn slice_first_of_three() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=slice:1/3"), + @r" + success: true + exit_code: 0 + ----- stdout ----- + Starting 6 tests across 1 worker + PASS [TIME] test_mod::test_a + PASS [TIME] test_mod::test_d + ──────────── + Summary [TIME] 2 tests run: 2 passed, 0 skipped + + ----- stderr ----- + " + ); +} + +#[test] +fn slice_second_of_three() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=slice:2/3"), + @r" + success: true + exit_code: 0 + ----- stdout ----- + Starting 6 tests across 1 worker + PASS [TIME] test_mod::test_b + PASS [TIME] test_mod::test_e + ──────────── + Summary [TIME] 2 tests run: 2 passed, 0 skipped + + ----- stderr ----- + " + ); +} + +#[test] +fn slice_third_of_three() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=slice:3/3"), + @r" + success: true + exit_code: 0 + ----- stdout ----- + Starting 6 tests across 1 worker + PASS [TIME] test_mod::test_c + PASS [TIME] test_mod::test_f + ──────────── + Summary [TIME] 2 tests run: 2 passed, 0 skipped + + ----- stderr ----- + " + ); +} + +#[test] +fn slice_one_of_one_runs_everything() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=slice:1/1"), + @r" + success: true + exit_code: 0 + ----- stdout ----- + Starting 6 tests across 1 worker + PASS [TIME] test_mod::test_a + PASS [TIME] test_mod::test_b + PASS [TIME] test_mod::test_c + PASS [TIME] test_mod::test_d + PASS [TIME] test_mod::test_e + PASS [TIME] test_mod::test_f + ──────────── + Summary [TIME] 6 tests run: 6 passed, 0 skipped + + ----- stderr ----- + " + ); +} + +#[test] +fn invalid_partition_index_above_total_errors() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=slice:4/3"), + @r" + success: false + exit_code: 2 + ----- stdout ----- + + ----- stderr ----- + error: invalid value 'slice:4/3' for '--partition ': partition index `M` (4) must not exceed partition count `N` (3) + + For more information, try '--help'. + " + ); +} + +#[test] +fn invalid_partition_strategy_errors() { + let context = TestContext::with_file("test_mod.py", SIX_TESTS); + + assert_cmd_snapshot!( + context.command_no_parallel().arg("--partition=hash:1/3"), + @r" + success: false + exit_code: 2 + ----- stdout ----- + + ----- stderr ----- + error: invalid value 'hash:1/3' for '--partition ': unknown partition strategy `hash`; supported strategies: `slice` + + For more information, try '--help'. + " + ); +} diff --git a/crates/karva_benchmark/src/lib.rs b/crates/karva_benchmark/src/lib.rs index c767ed93..2dc36db4 100644 --- a/crates/karva_benchmark/src/lib.rs +++ b/crates/karva_benchmark/src/lib.rs @@ -58,6 +58,7 @@ pub fn run_karva(project: &Project) { create_ctrlc_handler: false, last_failed: false, profile: None, + partition: None, }; let args = SubTestCommand { diff --git a/crates/karva_cli/src/lib.rs b/crates/karva_cli/src/lib.rs index f7826cb9..94181683 100644 --- a/crates/karva_cli/src/lib.rs +++ b/crates/karva_cli/src/lib.rs @@ -4,12 +4,14 @@ use clap::builder::styling::{AnsiColor, Effects}; mod cache; mod enums; +mod partition; mod snapshot; mod test; mod verbosity; pub use cache::{CacheAction, CacheCommand}; pub use enums::{CovReport, NoTests, OutputFormat, RunIgnored}; +pub use partition::PartitionSelection; pub use snapshot::{ SnapshotAction, SnapshotCommand, SnapshotDeleteArgs, SnapshotFilterArgs, SnapshotPruneArgs, }; diff --git a/crates/karva_cli/src/partition.rs b/crates/karva_cli/src/partition.rs new file mode 100644 index 00000000..5086a6aa --- /dev/null +++ b/crates/karva_cli/src/partition.rs @@ -0,0 +1,142 @@ +use std::fmt; +use std::str::FromStr; + +/// Selection of a single partition (slice) from the collected tests. +/// +/// Used by `--partition slice:M/N` to run only the tests assigned to slice +/// `M` of `N`. Slice indices are 1-indexed: `slice:1/3`, `slice:2/3`, +/// `slice:3/3` together cover every collected test exactly once. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PartitionSelection { + pub index: u32, + pub total: u32, +} + +impl PartitionSelection { + /// Returns true if the test at `position` (0-indexed, in the deterministic + /// post-filter ordering) belongs to this slice. + #[must_use] + pub fn contains(self, position: usize) -> bool { + // 1-indexed input -> 0-indexed modulo target. + let target = (self.index - 1) as usize; + let total = self.total as usize; + position % total == target + } +} + +impl fmt::Display for PartitionSelection { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "slice:{}/{}", self.index, self.total) + } +} + +impl FromStr for PartitionSelection { + type Err = String; + + fn from_str(raw: &str) -> Result { + let (kind, body) = raw.split_once(':').ok_or_else(|| { + format!("expected `:/` (e.g. `slice:1/3`), got `{raw}`") + })?; + + if kind != "slice" { + return Err(format!( + "unknown partition strategy `{kind}`; supported strategies: `slice`" + )); + } + + let (m, n) = body + .split_once('/') + .ok_or_else(|| format!("expected `slice:/`, got `slice:{body}`"))?; + + let index: u32 = m + .parse() + .map_err(|err| format!("`{m}` is not a valid partition index: {err}"))?; + let total: u32 = n + .parse() + .map_err(|err| format!("`{n}` is not a valid partition count: {err}"))?; + + if total == 0 { + return Err("partition count `N` must be at least 1".to_string()); + } + if index == 0 { + return Err("partition index `M` must be at least 1".to_string()); + } + if index > total { + return Err(format!( + "partition index `M` ({index}) must not exceed partition count `N` ({total})" + )); + } + + Ok(Self { index, total }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_valid_slice() { + assert_eq!( + "slice:1/3".parse::().unwrap(), + PartitionSelection { index: 1, total: 3 }, + ); + assert_eq!( + "slice:3/3".parse::().unwrap(), + PartitionSelection { index: 3, total: 3 }, + ); + assert_eq!( + "slice:1/1".parse::().unwrap(), + PartitionSelection { index: 1, total: 1 }, + ); + } + + #[test] + fn rejects_zero_total() { + assert!("slice:1/0".parse::().is_err()); + } + + #[test] + fn rejects_zero_index() { + assert!("slice:0/3".parse::().is_err()); + } + + #[test] + fn rejects_index_above_total() { + assert!("slice:4/3".parse::().is_err()); + } + + #[test] + fn rejects_unknown_strategy() { + assert!("hash:1/3".parse::().is_err()); + } + + #[test] + fn rejects_missing_separators() { + assert!("slice".parse::().is_err()); + assert!("slice:13".parse::().is_err()); + assert!("1/3".parse::().is_err()); + } + + #[test] + fn contains_round_robin() { + let p = PartitionSelection { index: 1, total: 3 }; + assert!(p.contains(0)); + assert!(!p.contains(1)); + assert!(!p.contains(2)); + assert!(p.contains(3)); + + let q = PartitionSelection { index: 3, total: 3 }; + assert!(!q.contains(0)); + assert!(!q.contains(1)); + assert!(q.contains(2)); + assert!(q.contains(5)); + } + + #[test] + fn display_round_trip() { + let p = PartitionSelection { index: 2, total: 5 }; + assert_eq!(p.to_string(), "slice:2/5"); + assert_eq!(p.to_string().parse::().unwrap(), p); + } +} diff --git a/crates/karva_cli/src/test.rs b/crates/karva_cli/src/test.rs index 188a58da..12c3bc11 100644 --- a/crates/karva_cli/src/test.rs +++ b/crates/karva_cli/src/test.rs @@ -9,6 +9,7 @@ use karva_metadata::{ }; use crate::enums::{CovReport, NoTests, OutputFormat, RunIgnored}; +use crate::partition::PartitionSelection; use crate::verbosity::Verbosity; /// Shared test execution options that can be used by both main CLI and worker processes @@ -219,6 +220,21 @@ pub struct TestCommand { #[clap(long, alias = "lf", help_heading = "Filter options")] pub last_failed: bool, + /// Run only a slice of the collected tests, distributed round-robin. + /// + /// Accepts `slice:M/N` where this run executes slice `M` of `N` total + /// slices (1-indexed). Tests are sorted by qualified name and then + /// distributed by cycling through slices: test 1 to slice 1, test 2 to + /// slice 2, ..., test N+1 to slice 1, and so on. Running every + /// `slice:1/N` through `slice:N/N` together covers every collected test + /// exactly once. + /// + /// Useful for splitting a test run across CI jobs. Slice membership + /// shifts when tests are added or removed, so it gives less stable + /// per-test placement than a hash-based scheme. + #[clap(long, value_name = "STRATEGY:M/N", help_heading = "Filter options")] + pub partition: Option, + /// Number of parallel workers (default: number of CPU cores) #[clap(short = 'n', long, help_heading = "Runner options")] pub num_workers: Option, diff --git a/crates/karva_runner/src/orchestration.rs b/crates/karva_runner/src/orchestration.rs index 593603dd..c1506264 100644 --- a/crates/karva_runner/src/orchestration.rs +++ b/crates/karva_runner/src/orchestration.rs @@ -13,7 +13,7 @@ use karva_cache::{ AggregatedResults, CACHE_DIR, RunCache, RunHash, read_last_failed, read_recent_durations, write_last_failed, }; -use karva_cli::SubTestCommand; +use karva_cli::{PartitionSelection, SubTestCommand}; use karva_collector::{CollectedPackage, CollectionSettings}; use karva_logging::Printer; use karva_logging::time::format_duration; @@ -154,6 +154,8 @@ pub struct ParallelTestConfig { /// Active configuration profile name. Propagated to workers as /// `KARVA_PROFILE`; falls back to `"default"` when `None`. pub profile: Option, + /// When set, restrict the run to the selected slice of collected tests. + pub partition: Option, } /// Spawn worker processes for each partition @@ -303,6 +305,7 @@ pub fn run_parallel_tests( num_workers, &previous_durations, &last_failed_set, + config.partition, ); let run_hash = RunHash::current_time(); diff --git a/crates/karva_runner/src/partition.rs b/crates/karva_runner/src/partition.rs index f6ac6844..ab5eced1 100644 --- a/crates/karva_runner/src/partition.rs +++ b/crates/karva_runner/src/partition.rs @@ -1,6 +1,8 @@ use std::collections::{HashMap, HashSet}; use std::time::Duration; +use karva_cli::PartitionSelection; + /// Test metadata used for partitioning decisions #[derive(Debug, Clone)] struct TestInfo { @@ -104,6 +106,7 @@ pub fn partition_collected_tests( num_workers: usize, previous_durations: &HashMap, last_failed: &HashSet, + partition_selection: Option, ) -> Vec { let mut test_infos = Vec::new(); collect_test_paths_recursive(package, &mut test_infos, previous_durations); @@ -112,6 +115,19 @@ pub fn partition_collected_tests( test_infos.retain(|info| last_failed.contains(&info.qualified_name)); } + // Slice partitioning runs on a deterministic ordering of the post-filter + // test set so that `slice:M/N` is stable across runs and machines (modulo + // changes to the test set itself). + if let Some(selection) = partition_selection { + test_infos.sort_by(|a, b| a.qualified_name.cmp(&b.qualified_name)); + let mut position = 0usize; + test_infos.retain(|_| { + let keep = selection.contains(position); + position += 1; + keep + }); + } + // Shuffle tests without durations so they distribute randomly across partitions shuffle_tests_without_durations(&mut test_infos); diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 6253b571..edb79a36 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -101,7 +101,10 @@ karva test [OPTIONS] [PATH]...
  • full: Print diagnostics verbosely, with context and helpful hints (default)
  • concise: Print diagnostics concisely, one per line
  • -
--profile, -P name

Configuration profile to use.

+
--partition strategy:m/n

Run only a slice of the collected tests, distributed round-robin.

+

Accepts slice:M/N where this run executes slice M of N total slices (1-indexed). Tests are sorted by qualified name and then distributed by cycling through slices: test 1 to slice 1, test 2 to slice 2, ..., test N+1 to slice 1, and so on. Running every slice:1/N through slice:N/N together covers every collected test exactly once.

+

Useful for splitting a test run across CI jobs. Slice membership shifts when tests are added or removed, so it gives less stable per-test placement than a hash-based scheme.

+
--profile, -P name

Configuration profile to use.

Profiles are defined as [profile.<name>] sections in karva.toml (or [tool.karva.profile.<name>] in pyproject.toml) and may override any of the [src], [terminal], and [test] settings. The selected profile is layered on top of any [profile.default] overrides, which themselves layer on top of the top-level options.

Defaults to default.

May also be set with the KARVA_PROFILE environment variable.

--retry retry

When set, the test will retry failed tests up to this number of times

diff --git a/docs/usage/running-tests/parallel.md b/docs/usage/running-tests/parallel.md index 601e1c11..bb52c846 100644 --- a/docs/usage/running-tests/parallel.md +++ b/docs/usage/running-tests/parallel.md @@ -42,3 +42,15 @@ karva test --no-capture ``` Reach for `--no-capture` when debugging with `print` statements or attaching `pdb`. For ad-hoc inspection without giving up parallelism, prefer `-s` / `--show-output`, which keeps capture on but prints the captured output for every test. + +## Splitting a run across CI jobs + +`--partition slice:M/N` runs only slice `M` of `N` total slices. Tests are sorted by qualified name and distributed round-robin: test 1 to slice 1, test 2 to slice 2, ..., test `N+1` to slice 1, and so on. Running every `slice:1/N` through `slice:N/N` together covers every collected test exactly once. + +```bash +karva test --partition slice:1/3 +karva test --partition slice:2/3 +karva test --partition slice:3/3 +``` + +Slices are computed deterministically from the current test set, so the same revision splits the same way on every machine. Adding or removing tests can shift which slice a given test falls into, so this is less stable per-test than a hash-based scheme but does not need any historical data.