scikit-hep / cabinetry

design and steer profile likelihood fits
https://cabinetry.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
26 stars 19 forks source link

feat: enable template histogram creation and collection to accept list of tasks #421

Open alexander-held opened 1 year ago

alexander-held commented 1 year ago

This updates the templates.build, templates.collect, and templates.postprocess implementations to accept a list of templates to produce via a new template_list keyword argument. The full list can be obtained via route.required_templates. By default, all templates will still be processed.

This allows users to parallelize the template histogram production / collection, partially addressing #401. An example of this would be the following:

from dask.distributed import Client, LocalCluster, wait

def produce_single_template(template):
    cabinetry.templates.build(config, template_list=[template])

template_list = cabinetry.route.required_templates(config)

with LocalCluster(n_workers=2) as cluster:
    client = Client(cluster)
    wait(client.map(produce_single_template, template_list))

compared to the non-parallelized version, which is unchanged:

cabinetry.templates.build(config)

breaking change:

* templates.build, templates.collect, and templates.postprocess now accept a list of tasks
* breaking change: route. apply_to_all_templates renamed to route.apply_to_templates
* breaking change: route.apply_to_templates now requires a list of tasks as argument
alexander-held commented 8 months ago

This will require a different approach to handle remote workers without access to a shared filesystem. To handle this, histograms would have to be passed back and saved at the head node. The current API does not surface the call to utils._name_and_save, which is quite deep inside internal APIs.

Currently route.apply_to_templates does generic function execution. This might have to be changed to support handling return values and passing those back to the higher level, e.g. templates.build for saving. The downside of this is that saving would only happen (in the non-parallel version) after all templates are constructed (and passed back as histograms).

to-do / think about: