research_helpers.sweep

This module provides tools for planning and running parameter sweeps as scheduler job arrays and collecting the results. Planning, running, and status checks need only the standard library. collect requires the sweep extra. See Parameter sweeps.

research_helpers.sweep.collect_results(run_dir, *, write=True, leading=())

Merge the part files into a table de-duplicating by combination id.

Duplicates are legitimate: a requeued task can recompute a combination whose result was already flushed. The last record written wins.

Parameters:

• run_dir (Path | str) – the run directory.

• write (bool, default: True) – also write ‘results.csv’, and ‘results.parquet’ where pyarrow allows it.

• leading (Iterable[str], default: ()) – columns to move to the front, e.g. the swept parameters, so the table reads configuration first and measurements second.

Return type:

DataFrame

Returns:

One row per combination.

research_helpers.sweep.read_parts(run_dir)

Read every result recorded by this run’s tasks.

Parameters:

run_dir (Path | str) – the run directory.

Return type:

list[dict[str, Any]]

Returns:

One dict per recorded result, in task order. A partial final line, left by a task killed mid-write, is skipped.

Raises:

FileNotFoundError – if the array has not run yet.

research_helpers.sweep.run_status(run_dir)

Report how much of a planned sweep has finished.

Parameters:

run_dir (Path | str) – the run directory.

Return type:

dict[str, float]

Returns:

Counts, percentage complete, and the runtime seen so far.

class research_helpers.sweep.Manifest(grid, combinations, n_tasks=1, metadata=<factory>, notes='')

Bases: object

The full description of a planned sweep, written once and read by every task.

Parameters:

• grid (dict[str, list[Any]])

• combinations (list[dict[str, Any]])

• n_tasks (int)

• metadata (dict[str, Any])

• notes (str)

property n_combinations: int

Return how many parameter combinations the sweep covers.

save(run_dir)

Write the manifest into ‘run_dir’ and return the path.

Return type:

Path

Parameters:

run_dir (Path | str)

classmethod load(run_dir)

Read the manifest from ‘run_dir’.

Raises:

FileNotFoundError – if the sweep has not been planned.

Return type:

Manifest

Parameters:

run_dir (Path | str)

research_helpers.sweep.combination_id(params)

Return a short, stable id for a parameter combination.

Derived from the values alone, so the same combination keeps the same id however the grid was expanded and on whichever machine.

Parameters:

params (dict[str, Any]) – the parameter values, excluding any id already assigned.

Return type:

str

Returns:

Twelve hex characters.

research_helpers.sweep.expand_grid(grid, constants=None)

Expand a parameter grid into the full Cartesian product of combinations.

Parameters:

• grid (dict[str, list[Any]]) – parameter name to the list of values to sweep.

• constants (dict[str, Any] | None, default: None) – fixed values added to every combination before it is hashed, e.g. the label of the dataset being swept. Part of the identity, so two datasets do not collide.

Return type:

list[dict[str, Any]]

Returns:

One dict per combination, each carrying the constants and a ‘combination_id’. Ordering is deterministic: parameter names sorted, values in the order given.

Raises:

ValueError – if the grid is empty.

research_helpers.sweep.slice_bounds(n_items, n_tasks, task_index)

Return the [start, end) slice of items belonging to a 1-based array task.

Parameters:

• n_items (int) – total number of combinations.

• n_tasks (int) – number of array tasks the sweep was planned for.

• task_index (int) – 1-based task index.

Return type:

tuple[int, int]

Returns:

Start and end indices into the combination list.

Raises:

ValueError – if ‘task_index’ is outside 1..n_tasks.

research_helpers.sweep.completed_ids(run_dir)

Return the combination ids already recorded in this run’s parts.

A task killed mid-write leaves a partial final line, which is skipped rather than treated as corruption.

Parameters:

run_dir (Path | str) – the run directory.

Return type:

set[str]

Returns:

The ids found.

research_helpers.sweep.run_slice(run_dir, evaluate, *, task_index, n_tasks=None, context=None, resume=True, quiet=False)

Run this task’s slice of the sweep, appending each result to its own part file.

Parameters:

• run_dir (Path | str) – the run directory, holding the manifest.

• evaluate (Callable[[dict[str, Any], Any], dict[str, Any]]) – called as ‘evaluate(params, context)’ for each combination, returning the measurements to record. The parameters are recorded alongside them automatically.

• task_index (int) – 1-based array task index.

• n_tasks (int | None, default: None) – total array tasks. Defaults to the manifest’s value, which is what it should be; pass one only to deliberately re-partition the grid.

• context (Any, default: None) – whatever the evaluation needs, built once for the task.

• resume (bool, default: True) – skip combinations already recorded in this run.

• quiet (bool, default: False) – suppress per-combination progress.

Return type:

Path

Returns:

The part file written.

research_helpers.sweep.task_index_from_env(explicit=None)

Return the array task index, from the argument, the scheduler, or 1 for a local run.

Parameters:

explicit (int | None, default: None) – an index given on the command line, which wins.

Return type:

int

Returns:

The 1-based task index.

class research_helpers.sweep.Sweep(evaluate=None, context=None)

Bases: object

A parameter sweep: an evaluation, an optional per-task context, and a command line.

Parameters:

• evaluate (Callable[[dict[str, Any], Any], dict[str, Any]] | None)

• context (Callable[[Manifest, Path], Any] | None)

evaluate(function)

Register the evaluation, called as ‘function(params, context)’ for each combination.

Return type:

Callable[[dict[str, Any], Any], dict[str, Any]]

Parameters:

function (Callable[[dict[str, Any], Any], dict[str, Any]])

context(function)

Register what to build once per task, called as ‘function(manifest, run_dir)’.

Return type:

Callable[[Manifest, Path], Any]

Parameters:

function (Callable[[Manifest, Path], Any])

property evaluation: Callable[[dict[str, Any], Any], dict[str, Any]]

Return the registered evaluation.

Raises:

RuntimeError – if none was registered.

build_context(manifest, run_dir)

Return the per-task context, or None if the sweep registered no context factory.

Return type:

Any

Parameters:

• manifest (Manifest)

• run_dir (Path)

runs_dir()

Return where sweeps live.

Return type:

Path

plan(run_dir, grid, *, constants=None, metadata=None, notes='', n_tasks=None, max_tasks=LARGE_ARRAY)

Expand a grid and write the manifest that every task will read.

Parameters:

• run_dir (Path | str) – directory to hold the manifest, parts and results.

• grid (dict[str, list[Any]]) – parameter name to the values to sweep.

• constants (dict[str, Any] | None, default: None) – fixed values on every combination, included in its identity.

• metadata (dict[str, Any] | None, default: None) – anything a task needs, e.g. a dataset name or a seed count.

• notes (str, default: '') – free text recorded with the sweep.

• n_tasks (int | None, default: None) – array width. Defaults to the smaller of ‘max_tasks’ and the grid size.

• max_tasks (int, default: LARGE_ARRAY) – the cap applied when ‘n_tasks’ is not given.

Return type:

Manifest

Returns:

The manifest, already written.

main(argv=None)

Run the sweep command line.

Parameters:

argv (list[str] | None, default: None) – command line arguments, or None to read ‘sys.argv’.

Return type:

int

Returns:

A process exit status.

research_helpers.sweep.estimate_runtime(manifest, evaluate, context=None, *, samples=DEFAULT_SAMPLES, seed=0)

Time a random sample of combinations.

The sample is random, not strided. ‘research_helpers.sweep.grid.expand_grid’ orders combinations by the Cartesian product over sorted parameter names, so a stride lands on a systematically correlated subset.

Parameters:

• manifest (Manifest) – the planned sweep.

• evaluate (Callable[[dict[str, Any], Any], dict[str, Any]]) – the evaluation, called as for a real task.

• context (Any, default: None) – whatever the evaluation needs.

• samples (int, default: DEFAULT_SAMPLES) – how many combinations to time.

• seed (int, default: 0) – seed for choosing the sample.

Return type:

list[float]

Returns:

The timings, sorted ascending.

research_helpers.sweep.read_config(path)

Read a sweep config from TOML or JSON.

Both are standard library, so planning a sweep needs nothing installed. A project preferring another format can parse it itself and call ‘Sweep.plan’ directly.

Parameters:

path (Path | str) – the config file. ‘.json’ is read as JSON, anything else as TOML.

Return type:

dict[str, Any]

Returns:

The parsed config – ‘grid’, and optionally ‘constants’, ‘metadata’, ‘notes’, ‘max_tasks’.