PrincetonUniversity / PsyNeuLink / 28132814695

84%

master: 85%

Build Type

push

github

Committed by web-flow

Commit Message

Distributed (Dask) fitting for ParameterEstimationComposition (#3561)

## Summary

Adds distributed evaluation to `ParameterEstimationComposition` fitting.
Candidate
parameterizations are evaluated in parallel across a Dask cluster
instead of
serially, controlled by a single constructor flag. Serial behavior is
unchanged
when the flag is off, and Dask is imported only when distributed fitting
is
requested.

## User-facing API

`PECOptimizationFunction` and `ParameterEstimationComposition` gain two
arguments:

- `distributed: bool = False` — evaluate candidates across a Dask
cluster.
- `distributed_options: Mapping` — configuration with the following
keys:
- `pec_factory` — required; a picklable callable that takes the observed
data and
    returns a fresh `(pec, inputs)` for a worker to score.
- `worker_cores` — LLVM threads per worker; defaults to
`$SLURM_CPUS_PER_TASK`.
- `max_concurrent_evaluations` — candidates per ask/tell round; defaults
to the
    live worker count.
  - `client` — a Dask client to use instead of an auto-resolved cluster.
  - `n_workers` — size of the auto-created single-node `LocalCluster`.

Both the optuna samplers and `differential_evolution` are supported.
LLVM
execution and a single scalar objective are required.

## Launch models

- Single node: `python study.py` forms a `LocalCluster` automatically.
- Multiple nodes: `srun -n N python -m psyneulink.dask_run study.py`. A
new
launcher module forms the scheduler, driver, and workers from the SLURM
ranks,
so the study script contains no scheduler addresses or cluster setup.
The worker
  count is `N - 2`.

A worked DDM example and an SBATCH template are under
`Scripts/Debug/pec_dask/`.

## How it works

The driver owns the optimizer. Each ask/tell round dispatches a batch of
candidates to workers; a worker rebuilds one PEC locally from
`pec_factory`,
caches it, and returns `float(pec.log_likelihood(...))`. A live PEC is
never
shipped across the wire, since Compositions are not... (continued)

Coverage Stats

11056 of 14436 branches covered (76.59%)

Branch coverage included in aggregate %.

158 of 272 new or added lines in 3 files covered. (58.09%)

1 existing line in 1 file now uncovered.

37804 of 43795 relevant lines covered (86.32%)

0.86 hits per line