Parallelism among multi-model jobs

[ryants]

If a job has more than one independent model evaluation, then allow dask to schedule to multiple workers. Similar to JobGroup but with multiple models instead of smoothing

[emitra17]

Now I've also encountered a case where I would want the opposite: in a job with smoothing, don't use JobGroup and instead run all replicates on the same worker.

[emitra17]

We have a model of interest where this feature would be useful... removing "low priority" tag.

[emitra17]

Partially implemented with #256 .
Still need to support simultaneous use of multi-model parallelism and smoothing.

[wshlavacek]

Current state of the implementation

PR #256 added parallelize_models as a separate feature alongside the existing smoothing. The two features share a common interface (JobGroup / MultimodelJobGroup) but are implemented as mutually exclusive branches.

How smoothing works (pre-existing)

When smoothing > 1, make_job() creates N identical replicate jobs, each running the full model list. Results are collected into a JobGroup, and average_results() computes element-wise np.mean across the replicate Data objects.

Job names: sim_1_rep0, sim_1_rep1, ...

How parallelize_models works (PR #256)

When parallelize_models > 1, make_job() partitions the model list into M slices and creates one job per slice. Results are collected into a MultimodelJobGroup, and average_results() merges the disjoint model dictionaries into a single Result via Result.add_result().

Job names: sim_1_part0, sim_1_part1, ...

Why they can't be used together

1. make_job() uses if/elif/else — smoothing and parallelize_models are separate branches that can't both execute.

2. Different aggregation strategies — smoothing averages identical results; parallelize_models merges disjoint model results. Combined use needs two-level aggregation: merge models first, then average across replicas.

3. Hard validation error at config.py:461:

   if self.config['smoothing'] > 1 and self.config['parallelize_models'] > 1:
       raise PybnfError('Simultaneous use of "smoothing" and "parallelize_models" is not supported')

4. Postprocessing mismatch — smoothing jobs get the user's postprocessing settings; parallelize_models jobs get an empty dict (postprocessing must happen after model merging).

What already works well

The main run loop (lines 898-905) is generic — it just calls group.job_finished(res) and group.average_results() on whatever group type is in job_group_dir. A new group class implementing the same interface would plug in without changes to the run loop.

[wshlavacek]

Implementation plan

Step 1: Add HybridJobGroup class

New class in algorithms.py alongside JobGroup and MultimodelJobGroup. Manages N×M subjobs (N smoothing replicas × M model partitions).

Internal structure: N MultimodelJobGroup instances (one per replica). When a subjob arrives:

1. Route it to the correct replica's MultimodelJobGroup based on the job name
2. When a replica's MultimodelJobGroup is complete, merge its models via add_result()
3. When all N replicas are merged, average across them via Data.average()

job_finished(res) → True only when all N×M subjobs have arrived
average_results() → merge models within each replica, then average across replicas

Step 2: Update make_job() branching

Replace the if/elif/else with nested logic:

smoothing = self.config.config['smoothing']
parallelize_models = self.config.config['parallelize_models']

if smoothing > 1 and parallelize_models > 1:
    # Create smoothing × parallelize_models jobs
    # Names: sim_1_rep0_part0, sim_1_rep0_part1, sim_1_rep1_part0, ...
    # Group: HybridJobGroup
elif smoothing > 1:
    # Existing smoothing logic (unchanged)
elif parallelize_models > 1:
    # Existing parallelize_models logic (unchanged)
else:
    # Single job (unchanged)

Step 3: Remove validation error

Delete the check at config.py:461-462 that raises PybnfError when both are > 1.

Step 4: Handle postprocessing order

In the combined case, postprocessing should be applied after model merging but before replica averaging. Each partial job gets an empty postprocessing dict (as parallelize_models currently does). HybridJobGroup.average_results() applies postprocessing to each merged replica before averaging.

Alternatively, if postprocessing is always safe to apply to the final averaged result, apply it only once at the end. This needs checking — if postprocessing is nonlinear, postprocess(average(x)) ≠ average(postprocess(x)).

Step 5: Objective function evaluation

Keep local evaluation (no calc_future) when either feature is active. The existing condition at line 860-861 already handles this correctly — it only uses distributed evaluation when both smoothing == 1 and parallelize_models == 1.

Step 6: Tests

• Unit tests for HybridJobGroup.job_finished() and average_results() with mock Result objects
• Unit test for make_job() in combined mode: verify correct number of jobs, correct model partitioning, correct group type
• Integration test with a config using both smoothing > 1 and parallelize_models > 1

Step 7: Documentation

Update config key docs to note that smoothing and parallelize_models can now be used together.

---

Pre-implementation unit tests needed

There are currently zero unit tests for JobGroup, MultimodelJobGroup, make_job(), job_finished(), average_results(), or Result.add_result(). These features are only exercised by integration tests (full fitting runs with BNG).

Before implementing the combined feature, we should add baseline tests for the existing behavior:

1. JobGroup.job_finished() — feed N results one at a time, verify returns False for the first N-1, True for the Nth. Test with FailedSimulation results too (should return True immediately on failure).

2. JobGroup.average_results() — feed results with known Data values, verify the output is the element-wise mean. Verify a FailedSimulation is returned when any subjob fails.

3. MultimodelJobGroup.average_results() — feed results with disjoint model keys, verify the output contains all models merged.

4. Result.add_result() — verify simdata dictionaries are merged correctly.

5. make_job() branches — verify smoothing creates N full-model jobs + a JobGroup, parallelize_models creates M partial-model jobs + a MultimodelJobGroup, and default creates a single job. This requires mocking the model list and config.

These tests protect the existing behavior and make it safe to refactor make_job() and add HybridJobGroup.

[wshlavacek]

Pre-implementation tests added

tests/test_job_groups.py — 19 tests covering JobGroup, MultimodelJobGroup, Result.add_result(), and the make_job() branching logic. Run with:

cd tests
python3 -m pytest test_job_groups.py -v

No BioNetGen installation required; BNG simulator detection and model initialization are mocked.

What the tests cover

TestResultAddResult (3 tests) — Verifies Result.add_result() correctly merges simdata dictionaries (disjoint keys, duplicate keys, empty target).

TestJobGroup (6 tests) — Verifies smoothing group behavior:

• job_finished() returns True only when all replicas arrive
• job_finished() returns True immediately on FailedSimulation, ignores subsequent results
• average_results() computes element-wise mean across replicate Data arrays
• average_results() handles multiple models/suffixes
• average_results() returns FailedSimulation when any subjob fails

TestMultimodelJobGroup (4 tests) — Verifies parallel-model group behavior:

• job_finished() completion tracking
• average_results() merges disjoint model results into a single Result
• Three-partition merge
• Failure handling

TestMakeJob (6 tests) — Verifies make_job() branching:

• Default (both disabled): returns 1 job
• smoothing > 1: returns N replicate jobs with full model list, creates JobGroup
• parallelize_models > 1: returns M jobs with partitioned model list, creates MultimodelJobGroup
• Auto-incrementing job IDs

These tests protect existing behavior and make it safe to add the combined smoothing + parallelize_models branch and new HybridJobGroup class.