ITP map logic prevents paralellization

[JoranAngevaare]

Describe the bug
When loading plugins requiring an interpolation map, funny business ensures. I think some things are not handelled properly as would be required for proper multithreading in this function, see the traceback below.

To Reproduce
Insert the MWE of how to reproduce the error

st.make(run_list, 'event_pattern_fit', progress_bar=True, max_workers=20)

Versions
Please add the output of:

  | module | version | path | git
-- | -- | -- | -- | --
python | 3.8.13 | /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/b... | None
strax | 1.2.2 | /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/l... | None
straxen | 1.7.0 | /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/l... | None
cutax | 1.9.0 | /dali/lgrandi/xenonnt/software/cutax/v1.9.0/cutax | None
wfsim | 0.6.1 | /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/l... | None
pema | 0.4.3 | /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/l... | None

Trackback

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/context.py:1403, in Context.make(self, run_id, targets, save, max_workers, _skip_if_built, **kwargs)
   1401     raise ValueError("Cannot build empty list of runs")
   1402 if len(run_ids) > 1:
-> 1403     return strax.multi_run(
   1404         self.get_array, run_ids, targets=targets,
   1405         throw_away_result=True, log=self.log,
   1406         save=save, max_workers=max_workers, **kwargs)
   1408 if _skip_if_built and self.is_stored(run_id, targets):
   1409     return

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/utils.py:541, in multi_run(exec_function, run_ids, max_workers, throw_away_result, multi_run_progress_bar, ignore_errors, log, *args, **kwargs)
    539         failures.append(_run_id)
    540         continue
--> 541     raise f.exception()
    543 if throw_away_result:
    544     continue

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/concurrent/futures/thread.py:57, in _WorkItem.run(self)
     54     return
     56 try:
---> 57     result = self.fn(*self.args, **self.kwargs)
     58 except BaseException as exc:
     59     self.future.set_exception(exc)

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/context.py:1440, in Context.get_array(self, run_id, targets, save, max_workers, **kwargs)
   1433 else:
   1434     source = self.get_iter(
   1435         run_ids[0],
   1436         targets,
   1437         save=save,
   1438         max_workers=max_workers,
   1439         **kwargs)
-> 1440     results = [x.data for x in source]
   1442 results = np.concatenate(results)
   1443 return results

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/context.py:1440, in <listcomp>(.0)
   1433 else:
   1434     source = self.get_iter(
   1435         run_ids[0],
   1436         targets,
   1437         save=save,
   1438         max_workers=max_workers,
   1439         **kwargs)
-> 1440     results = [x.data for x in source]
   1442 results = np.concatenate(results)
   1443 return results

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/context.py:1263, in Context.get_iter(self, run_id, targets, save, max_workers, time_range, seconds_range, time_within, time_selection, selection_str, keep_columns, drop_columns, allow_multiple, progress_bar, _chunk_number, **kwargs)
   1254     elif (self.context_config['timeout'] > 7200 or (
   1255             self.context_config['allow_lazy'] and
   1256             not self.context_config['allow_multiprocess'])):
   1257         # For allow_multiple we don't want allow this when in lazy mode
   1258         # with long timeouts (lazy-mode is disabled if multiprocessing
   1259         # so if that is activated, we can also continue)
   1260         raise RuntimeError(f'Cannot allow_multiple in lazy mode or '
   1261                            f'with long timeouts.')
-> 1263 components = self.get_components(run_id,
   1264                                  targets=targets,
   1265                                  save=save,
   1266                                  time_range=time_range,
   1267                                  chunk_number=_chunk_number)
   1269 # Cleanup the temp plugins
   1270 for k in list(self._plugin_class_registry.keys()):

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/strax/context.py:1024, in Context.get_components(self, run_id, targets, save, time_range, chunk_number)
   1022 for d in plugins.values():
   1023     self._set_plugin_config(d, run_id, tolerant=False)
-> 1024     d.setup()
   1025 return strax.ProcessorComponents(
   1026     plugins=plugins,
   1027     loaders=loaders,
   1028     loader_plugins=loader_plugins,
   1029     savers=savers,
   1030     targets=strax.to_str_tuple(final_plugin))

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/straxen/plugins/event_patternfit.py:127, in EventPatternFit.setup(self)
    124 self.mean_pe_photon = self.config['mean_pe_per_photon']
    126 # Getting optical maps
--> 127 self.s1_pattern_map = straxen.InterpolatingMap(
    128     straxen.get_resource(
    129         self.config['s1_optical_map'],
    130         fmt=self._infer_map_format(self.config['s1_optical_map'])))
    131 self.s2_pattern_map = straxen.InterpolatingMap(
    132     straxen.get_resource(
    133         self.config['s2_optical_map'],
    134         fmt=self._infer_map_format(self.config['s2_optical_map'])))
    136 # Getting S2 data-driven tensorflow models

File /opt/XENONnT/anaconda/envs/XENONnT_2022.05.2/lib/python3.8/site-packages/straxen/itp_map.py:113, in InterpolatingMap.__init__(self, data, method, **kwargs)
    109     compressor, dtype, shape = self.data['compressed']
    110     self.data['map'] = np.frombuffer(
    111         strax.io.COMPRESSORS[compressor]['decompress'](self.data['map']),
    112         dtype=dtype).reshape(*shape)
--> 113     del self.data['compressed']
    114 if 'quantized' in self.data:
    115     self.data['map'] = self.data['quantized'] * self.data['map'].astype(np.float32)

KeyError: 'compressed'