diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b3427537..e2f52566d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -246,7 +246,7 @@ jobs: release: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') runs-on: ubuntu-latest - needs: [test, test-no-dask, test-vine] + needs: [test, test-dask-client, test-no-dask, test-vine] permissions: id-token: write attestations: write @@ -278,7 +278,10 @@ jobs: password: ${{ secrets.PYPI_TOKEN }} pass: + if: always() needs: [test, test-dask-client, test-no-dask, test-vine] runs-on: ubuntu-latest steps: + - if: contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') || contains(needs.*.result, 'skipped') + run: exit 1 - run: echo "All jobs passed" diff --git a/codecov.yml b/codecov.yml index 3dbf99e0d..9e35b5072 100644 --- a/codecov.yml +++ b/codecov.yml @@ -1,7 +1,6 @@ comment: off codecov: - token: 045255bb-e0d8-4c5d-b413-9c33128a03a4 notify: require_ci_to_pass: yes diff --git a/src/coffea/analysis_tools.py b/src/coffea/analysis_tools.py index 5f5b9436c..8abdf44ea 100644 --- a/src/coffea/analysis_tools.py +++ b/src/coffea/analysis_tools.py @@ -55,28 +55,6 @@ def _get_hist_class(delayed_mode): return hist.Hist -def _generate_slices(array_length, max_elements=128): - """Generate slices to split an array into chunks of at most `max_elements` elements - - Parameters - ---------- - array_length : int - The length of the array to split - max_elements : int, optional - The maximum number of elements in each chunk. Default is 128. - - Returns - ------- - slices : list of slice objects - A list of slice objects to iterate over and split the array into chunks with at most `max_elements` elements per slice - """ - slices = [] - for start in range(0, array_length, max_elements): - end = min(start + max_elements, array_length) - slices.append(slice(start, end)) - return slices - - def boolean_masks_to_categorical_integers( masks, insert_unmasked_as_zeros=False, @@ -187,10 +165,12 @@ def add(self, other): def __add__(self, other): temp = WeightStatistics(self.sumw, self.sumw2, self.minw, self.maxw, self.n) - return temp.add(other) + temp.add(other) + return temp def __iadd__(self, other): - return self.add(other) + self.add(other) + return self class Weights: @@ -394,7 +374,7 @@ def __add_multivariation_delayed( """Add a new weight with multiple variations in delayed mode""" dask_awkward = _import_dask_awkward() - if isinstance(weight, awkward.types.OptionType): + if isinstance(dask_awkward.type(weight), awkward.types.OptionType): # TODO what to do with option-type? is it representative of unknown weight # and we default to one or is it an invalid weight and we should never use this # event in the first place (0) ? @@ -567,8 +547,8 @@ def weight(self, modifier=None): """ if modifier is None: return self._weight - elif "Down" in modifier and modifier not in self._modifiers: - return self._weight / self._modifiers[modifier.replace("Down", "Up")] + elif modifier.endswith("Down") and modifier not in self._modifiers: + return self._weight / self._modifiers[modifier[:-4] + "Up"] return self._weight * self._modifiers[modifier] def partial_weight(self, include=[], exclude=[], modifier=None): @@ -628,12 +608,17 @@ def _partial_weight(self, include, exclude, modifier=None): if modifier is None: return w - elif modifier.replace("Down", "").replace("Up", "") not in names: + base = ( + modifier[:-4] + if modifier.endswith("Down") + else modifier[:-2] if modifier.endswith("Up") else modifier + ) + if not any(base == n or base.startswith(n + "_") for n in names): raise ValueError( f"Modifier {modifier} is not in the list of included weights" ) - elif "Down" in modifier and modifier not in self._modifiers: - return w / self._modifiers[modifier.replace("Down", "Up")] + if modifier.endswith("Down") and modifier not in self._modifiers: + return w / self._modifiers[modifier[:-4] + "Up"] return w * self._modifiers[modifier] @property @@ -642,7 +627,8 @@ def variations(self): keys = set(self._modifiers.keys()) # add any missing 'Down' variation for k in self._modifiers.keys(): - keys.add(k.replace("Up", "Down")) + if k.endswith("Up"): + keys.add(k[:-2] + "Down") return keys @@ -1188,10 +1174,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None): if do_weighted: axes.append(hist.storage.Weight()) if not self._delayed_mode and not do_categorical: - if categorical is not None: - raise NotImplementedError( - "yieldhist is not implemented for non-delayed mode (v1) with categorical" - ) h = hist.Hist(*axes) weighttofill = self._wgtev if do_weighted else self._nev if do_scaled: @@ -1200,10 +1182,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None): elif self._delayed_mode and not do_categorical: dask_awkward = _import_dask_awkward() - if categorical is not None: - raise NotImplementedError( - "yieldhist is not implemented for non-delayed mode (v1) with categorical" - ) h = Hist(*axes) for i, mask in enumerate(self._masks, 1): @@ -1772,10 +1750,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None): if do_weighted: axes.append(hist.storage.Weight()) if not self._delayed_mode and not do_categorical: - if categorical is not None: - raise NotImplementedError( - "yieldhist is not implemented for non-delayed mode (v1) with categorical" - ) honecut = hist.Hist(*axes) hcutflow = honecut.copy() hcutflow.axes.name = ("cutflow",) @@ -1789,10 +1763,6 @@ def yieldhist(self, weighted=None, scale=None, categorical=None): elif self._delayed_mode and not do_categorical: dask_awkward = _import_dask_awkward() - if categorical is not None: - raise NotImplementedError( - "yieldhist is not implemented for non-delayed mode (v1) with categorical" - ) honecut = Hist(*axes) hcutflow = honecut.copy() hcutflow.axes.name = ("cutflow",) @@ -2291,7 +2261,6 @@ def add_multiple(self, selections, fill_value=False): for name, selection in selections.items(): self.add(name, selection, fill_value) - @lru_cache def require(self, **names): """Return a mask vector corresponding to specific requirements @@ -2318,6 +2287,12 @@ def require(self, **names): returns a boolean array where an entry is True if the corresponding entries ``cut1 == True``, ``cut2 == False``, and ``cut3`` arbitrary. """ + # copy so a caller mutating the returned mask cannot corrupt the shared cache + result = self._require(**names) + return result.copy() if isinstance(result, numpy.ndarray) else result + + @lru_cache + def _require(self, **names): for cut, v in names.items(): if not isinstance(cut, str) or cut not in self._names: raise ValueError( diff --git a/src/coffea/btag_tools/btagscalefactor.py b/src/coffea/btag_tools/btagscalefactor.py index 1d57525d9..2baf0a5f8 100644 --- a/src/coffea/btag_tools/btagscalefactor.py +++ b/src/coffea/btag_tools/btagscalefactor.py @@ -147,31 +147,40 @@ def __init__(self, filename, workingpoint, methods="comb,comb,incl", keep_df=Fal ) mapping = numpy.full(bin_low_edges[0].shape, -1) - def findbin(flavor, eta, pt, discr): - btvflavor = self._flavor2btvflavor[flavor] - for i, (fbin, ebin, pbin, dbin) in enumerate(allbins): - if ( - btvflavor == fbin - and ebin[0] <= eta < ebin[1] - and pbin[0] <= pt < pbin[1] - and dbin[0] <= discr < dbin[1] - ): - return i - if eta < 0: - # maybe in this region we have only abseta - for i, (fbin, ebin, pbin, dbin) in enumerate(allbins): - if ( - btvflavor == fbin - and -ebin[1] <= eta < -ebin[0] - and pbin[0] <= pt < pbin[1] - and dbin[0] <= discr < dbin[1] - ): - return i - return -1 - - for idx, _ in numpy.ndenumerate(mapping): - flavor, eta, pt, discr = (x[idx] for x in bin_low_edges) - mapping[idx] = findbin(flavor, eta, pt, discr) + fbins = numpy.array([b[0] for b in allbins]) + eta_lo = numpy.array([b[1][0] for b in allbins]) + eta_hi = numpy.array([b[1][1] for b in allbins]) + pt_lo = numpy.array([b[2][0] for b in allbins]) + pt_hi = numpy.array([b[2][1] for b in allbins]) + discr_lo = numpy.array([b[3][0] for b in allbins]) + discr_hi = numpy.array([b[3][1] for b in allbins]) + + eta_cell = bin_low_edges[1].reshape(-1, 1) + pt_cell = bin_low_edges[2].reshape(-1, 1) + discr_cell = bin_low_edges[3].reshape(-1, 1) + btvflavor_cell = numpy.empty(bin_low_edges[0].size, dtype=fbins.dtype) + for flav, btv in self._flavor2btvflavor.items(): + btvflavor_cell[bin_low_edges[0].reshape(-1) == flav] = btv + btvflavor_cell = btvflavor_cell.reshape(-1, 1) + + base_match = ( + (btvflavor_cell == fbins) + & (pt_lo <= pt_cell) + & (pt_cell < pt_hi) + & (discr_lo <= discr_cell) + & (discr_cell < discr_hi) + ) + match = base_match & (eta_lo <= eta_cell) & (eta_cell < eta_hi) + found = match.any(axis=1) + result = numpy.where(found, match.argmax(axis=1), -1) + + # maybe in this region we have only abseta + abseta_match = base_match & (-eta_hi <= eta_cell) & (eta_cell < -eta_lo) + fallback = (~found) & (eta_cell[:, 0] < 0) + fallback &= abseta_match.any(axis=1) + result[fallback] = abseta_match.argmax(axis=1)[fallback] + + mapping = result.reshape(mapping.shape) if self.workingpoint == BTagScaleFactor.RESHAPE: self._corrections[syst] = dense_mapped_lookup( diff --git a/src/coffea/dataset_tools/manipulations.py b/src/coffea/dataset_tools/manipulations.py index b744475ec..b0879cb32 100644 --- a/src/coffea/dataset_tools/manipulations.py +++ b/src/coffea/dataset_tools/manipulations.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ast import copy import sys from collections.abc import Callable @@ -321,7 +322,7 @@ def get_failed_steps_for_dataset( ) for failure in failures: - args_as_types = tuple(eval(arg) for arg in failure.args) + args_as_types = tuple(ast.literal_eval(arg) for arg in failure.args) fname, object_path, start, stop, is_step = args_as_types diff --git a/src/coffea/dataset_tools/rucio_utils.py b/src/coffea/dataset_tools/rucio_utils.py index 53e100854..1635c7ef1 100644 --- a/src/coffea/dataset_tools/rucio_utils.py +++ b/src/coffea/dataset_tools/rucio_utils.py @@ -274,8 +274,9 @@ def get_dataset_files_replicas( outfiles.append(outfile) outsites.append(outsite) elif mode == "first": - outfiles.append(outfile[0]) - outsites.append(outsite[0]) + if outfile: + outfiles.append(outfile[0]) + outsites.append(outsite[0]) else: raise NotImplementedError(f"Mode {mode} not yet implemented!") @@ -287,7 +288,7 @@ def get_dataset_files_replicas( sites_counts[site] += 1 elif mode == "first": for site_by_file in outsites: - sites_counts[site] += 1 + sites_counts[site_by_file] += 1 return outfiles, outsites, sites_counts diff --git a/src/coffea/jetmet_tools/CorrectedJetsFactory.py b/src/coffea/jetmet_tools/CorrectedJetsFactory.py index 76aa8b1bd..7b51372fa 100644 --- a/src/coffea/jetmet_tools/CorrectedJetsFactory.py +++ b/src/coffea/jetmet_tools/CorrectedJetsFactory.py @@ -34,8 +34,11 @@ def __call__(self, array, like_what): def rand_gauss(item): + np_item = awkward.typetracer.length_one_if_typetracer(item).to_numpy() seeds = ( - awkward.typetracer.length_one_if_typetracer(item).to_numpy()[[0, -1]].view("i4") + np_item[[0, -1]].view("i4") + if len(np_item) + else numpy.zeros(2, dtype=numpy.int32) ) randomstate = numpy.random.Generator(numpy.random.PCG64(seeds)) @@ -148,20 +151,20 @@ def __init__(self, name_map, jec_stack): # from PhysicsTools/PatUtils/interface/SmearedJetProducerT.h#L283 self.forceStochastic = False - if "ptRaw" not in name_map or name_map["ptRaw"] is None: + self.treat_pt_as_raw = "ptRaw" not in name_map or name_map["ptRaw"] is None + if self.treat_pt_as_raw: warnings.warn( "There is no name mapping for ptRaw," " CorrectedJets will assume that .pt is raw pt!" ) name_map["ptRaw"] = name_map["JetPt"] + "_raw" - self.treat_pt_as_raw = "ptRaw" not in name_map if "massRaw" not in name_map or name_map["massRaw"] is None: warnings.warn( "There is no name mapping for massRaw," - " CorrectedJets will assume that .mass is raw pt!" + " CorrectedJets will assume that .mass is raw mass!" ) - name_map["ptRaw"] = name_map["JetMass"] + "_raw" + name_map["massRaw"] = name_map["JetMass"] + "_raw" total_signature = set() for part in _stack_parts: @@ -226,7 +229,7 @@ def build(self, injets): fields = awkward.fields(jets) if len(fields) == 0: raise Exception( - "Empty record, please pass a jet object with at least {self.real_sig} defined!" + f"Empty record, please pass a jet object with at least {self.real_sig} defined!" ) out = awkward.flatten(jets) wrap = partial( @@ -484,7 +487,7 @@ def build_variant(unc, template, jetpt, jetpt_orig, jetmass, jetmass_orig): label=f"{name}", ) - out_parms = out.layout.parameters + out_parms = dict(out.layout.parameters) out_parms["corrected"] = True out = awkward.zip( out_dict, depth_limit=1, parameters=out_parms, behavior=out.behavior diff --git a/src/coffea/jetmet_tools/FactorizedJetCorrector.py b/src/coffea/jetmet_tools/FactorizedJetCorrector.py index fee8d6673..0e41db008 100644 --- a/src/coffea/jetmet_tools/FactorizedJetCorrector.py +++ b/src/coffea/jetmet_tools/FactorizedJetCorrector.py @@ -13,8 +13,7 @@ def _checkConsistency(against, tocheck): else: if against != tocheck: raise Exception( - "Corrector for {} is mixed" - "with correctors for {}!".format(tocheck, against) + f"Corrector for {tocheck} is mixed with correctors for {against}!" ) return tocheck diff --git a/src/coffea/lookup_tools/evaluator.py b/src/coffea/lookup_tools/evaluator.py index 9472ad2d5..1079f24cd 100644 --- a/src/coffea/lookup_tools/evaluator.py +++ b/src/coffea/lookup_tools/evaluator.py @@ -4,6 +4,7 @@ from coffea.lookup_tools.jec_uncertainty_lookup import jec_uncertainty_lookup from coffea.lookup_tools.jersf_lookup import jersf_lookup from coffea.lookup_tools.jme_standard_function import jme_standard_function +from coffea.lookup_tools.json_lookup import json_lookup from coffea.lookup_tools.rochester_lookup import rochester_lookup lookup_types = { @@ -14,6 +15,7 @@ "jec_uncertainty_lookup": jec_uncertainty_lookup, "rochester_lookup": rochester_lookup, "correctionlib_wrapper": correctionlib_wrapper, + "json_lookup": json_lookup, } diff --git a/src/coffea/lookup_tools/extractor.py b/src/coffea/lookup_tools/extractor.py index 0155faa75..c4895b0c0 100644 --- a/src/coffea/lookup_tools/extractor.py +++ b/src/coffea/lookup_tools/extractor.py @@ -130,8 +130,6 @@ def add_weight_sets(self, weightsdescs): else: weights, thetype = self.extract_from_file(thefile, name) self.add_weight_set(local_name, thetype, weights) - if thetype == "json_lookup": - self._names[local_name] = 0 def import_file(self, thefile): """ diff --git a/src/coffea/lookup_tools/jme_standard_function.py b/src/coffea/lookup_tools/jme_standard_function.py index c83256888..cdb459781 100644 --- a/src/coffea/lookup_tools/jme_standard_function.py +++ b/src/coffea/lookup_tools/jme_standard_function.py @@ -46,37 +46,6 @@ def masked_bin_eval(dim1_indices, dimN_bins, dimN_vals): return dimN_indices, dimN_overflows -# idx_in is a tuple of indices in increasing jaggedness -# idx_out is a list of flat indices -def flatten_idxs(idx_in, jaggedarray): - """ - This provides a faster way to convert between tuples of - jagged indices and flat indices in a jagged array's contents - """ - if len(idx_in) == 0: - return numpy.array([], dtype=numpy.int) - idx_out = jaggedarray.starts[idx_in[0]] - if len(idx_in) == 1: - pass - elif len(idx_in) == 2: - idx_out += idx_in[1] - else: - raise Exception("jme_standard_function only works for two binning dimensions!") - - flattened = awkward.flatten(jaggedarray) - good_idx = idx_out < len(flattened) - if (~good_idx).any(): - input_idxs = tuple( - [idx_out[~good_idx]] + [idx_in[i][~good_idx] for i in range(len(idx_in))] - ) - raise Exception( - "Calculated invalid index {} for" - " array with length {}".format(numpy.vstack(input_idxs), len(flattened)) - ) - - return idx_out - - class jme_standard_function(lookup_base): """ This class defines a lookup table for jet energy corrections and resolutions. diff --git a/src/coffea/lookup_tools/json_lookup.py b/src/coffea/lookup_tools/json_lookup.py new file mode 100644 index 000000000..9afbf3411 --- /dev/null +++ b/src/coffea/lookup_tools/json_lookup.py @@ -0,0 +1,10 @@ +class json_lookup: + def __init__(self, wrapped_values): + self.values = wrapped_values + + def __call__(self, run, lumi): + out = [] + for r, ls in zip(run, lumi): + table = self.values.get(str(r)) + out.append(None if table is None else table.get(ls)) + return out diff --git a/src/coffea/lookup_tools/txt_converters.py b/src/coffea/lookup_tools/txt_converters.py index 4cf789b60..dfbe11a2f 100644 --- a/src/coffea/lookup_tools/txt_converters.py +++ b/src/coffea/lookup_tools/txt_converters.py @@ -497,7 +497,7 @@ def convert_effective_area_file(eaFilePath): binMaxs = numpy.unique(pars[columns[1]]) bins[layout[i + offset_name]] = numpy.union1d(binMins, binMaxs) else: - counts = numpy.zeros(0, dtype=numpy.int) + counts = numpy.zeros(0, dtype=numpy.int64) allBins = numpy.zeros(0, dtype=numpy.double) for binMin in bins[bin_order[0]][:-1]: binMins = numpy.unique( diff --git a/src/coffea/lumi_tools/lumi_tools.py b/src/coffea/lumi_tools/lumi_tools.py index 3b95c8027..ce33e7ac4 100644 --- a/src/coffea/lumi_tools/lumi_tools.py +++ b/src/coffea/lumi_tools/lumi_tools.py @@ -119,19 +119,23 @@ def get_lumi(self, runlumis): self.index = _make_lumi_data_dict() runs = self._lumidata[:, 0].astype("u4") lumis = self._lumidata[:, 1].astype("u4") - # fill self.index LumiData._build_lumi_table_kernel( runs, lumis, self._lumidata[:, 2], self.index ) - # delayed object cache - if _isinstance(runlumis, "dask_awkward.lib.core.Array"): - dask = _import_dask() - self.index_delayed = dask.delayed( - tuple([runs, lumis, self._lumidata[:, 2]]) - ) tot_lumi = numpy.zeros((1,), dtype=numpy.dtype("float64")) if _isinstance(runlumis, "dask_awkward.lib.core.Array"): + if self.index_delayed is None: + dask = _import_dask() + self.index_delayed = dask.delayed( + tuple( + [ + self._lumidata[:, 0].astype("u4"), + self._lumidata[:, 1].astype("u4"), + self._lumidata[:, 2], + ] + ) + ) dask_awkward = _import_dask_awkward() lumi_meta = wrap_get_lumi(runlumis._meta, self.index) lumi_per_partition = dask_awkward.map_partitions( @@ -144,7 +148,10 @@ def get_lumi(self, runlumis): tot_lumi = awkward.sum(lumi_per_partition, keepdims=True) else: LumiData._get_lumi_kernel( - runlumis[:, 0], runlumis[:, 1], self.index, tot_lumi + runlumis[:, 0].astype(numpy.uint32), + runlumis[:, 1].astype(numpy.uint32), + self.index, + tot_lumi, ) return ( tot_lumi[0] * self.seconds_per_lumi_LHC diff --git a/src/coffea/ml_tools/tf_wrapper.py b/src/coffea/ml_tools/tf_wrapper.py index 5e8af73d9..2196a6eec 100644 --- a/src/coffea/ml_tools/tf_wrapper.py +++ b/src/coffea/ml_tools/tf_wrapper.py @@ -114,7 +114,7 @@ def numpy_call(self, *args: numpy.array, **kwargs: numpy.array) -> numpy.array: kwargs = { key: ( tensorflow.convert_to_tensor(arr) - if arr.flags["WRITABLE"] + if arr.flags["WRITEABLE"] else tensorflow.convert_to_tensor(numpy.copy(arr)) ) for key, arr in kwargs.items() diff --git a/src/coffea/ml_tools/triton_wrapper.py b/src/coffea/ml_tools/triton_wrapper.py index 2b5136b8c..1ae0a1bf2 100644 --- a/src/coffea/ml_tools/triton_wrapper.py +++ b/src/coffea/ml_tools/triton_wrapper.py @@ -289,8 +289,8 @@ def _get_infer_shape(name): self.pmod.InferRequestedOutput(output) for output in output_list ] - # Setting up container for storing output. - output = None + # Collect per-batch outputs, concatenate once at the end. + output_batches = {o: [] for o in output_list} # Padding the outermost dimension to a multiple of of the batch size orig_len = list(input_dict.values())[0].shape[0] # saving original length @@ -312,25 +312,21 @@ def _get_infer_shape(name): # Running the request with fall back request = self.run_infer(infer_inputs, infer_outputs) - if output is None: - output = { - o: request.as_numpy(o)[start_idx:stop_idx] for o in output_list - } - else: - for o in output_list: - output[o] = numpy.concatenate( - (output[o], request.as_numpy(o)), axis=0 - ) + for o in output_list: + output_batches[o].append(request.as_numpy(o)[: stop_idx - start_idx]) if ( - output is None + orig_len == 0 ): # Input was a length-0, so we should generate the length-0 outputs with correct dimension return { o: numpy.zeros(shape=(0, *self.model_outputs[o]["shape"][1:])) for o in output_list } - return {k: v[:orig_len] for k, v in output.items()} + return { + o: numpy.concatenate(batches, axis=0)[:orig_len] + for o, batches in output_batches.items() + } def run_infer(self, inputs, outputs, attempt=0): """Thin wrapper around tritonclient.infer to automatic retry with backoff+jitter on inference server failures""" @@ -345,6 +341,8 @@ def run_infer(self, inputs, outputs, attempt=0): # Retry backoff + full jitter: # https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/ time.sleep( - numpy.random.rand() * self.retry_jitter_base_ms * (2**attempt) + numpy.random.rand() + * (self.retry_jitter_base_ms / 1000.0) + * (2**attempt) ) return self.run_infer(inputs, outputs, attempt + 1) diff --git a/src/coffea/nanoevents/factory.py b/src/coffea/nanoevents/factory.py index 5388e0c4f..33566ecdf 100644 --- a/src/coffea/nanoevents/factory.py +++ b/src/coffea/nanoevents/factory.py @@ -253,12 +253,14 @@ def __getstate__(self): "schema": self._schema, "mapping": self._mapping, "partition_key": self._partition_key, + "mode": self._mode, } def __setstate__(self, state): self._schema = state["schema"] self._mapping = state["mapping"] self._partition_key = state["partition_key"] + self._mode = state.get("mode", "eager") self._events = lambda: None @classmethod @@ -685,7 +687,11 @@ def from_preloaded( ) uuidpfn = {uuid: array_source} mapping = PreloadedSourceMapping( - PreloadedOpener(uuidpfn), entry_start, entry_stop, access_log=access_log + PreloadedOpener(uuidpfn), + entry_start, + entry_stop, + access_log=access_log, + buffer_cache=buffer_cache, ) mapping.preload_column_source(partition_key[0], partition_key[1], array_source) diff --git a/src/coffea/nanoevents/mapping/base.py b/src/coffea/nanoevents/mapping/base.py index 06b995dd7..cdba9abe4 100644 --- a/src/coffea/nanoevents/mapping/base.py +++ b/src/coffea/nanoevents/mapping/base.py @@ -9,6 +9,8 @@ from coffea.nanoevents import transforms from coffea.nanoevents.util import key_to_tuple, tuple_to_key +_MISSING = object() + class Accessed(NamedTuple): branch: str @@ -95,8 +97,10 @@ def interpret_key(cls, key): def __getitem__(self, key): def _getitem(key): - if self._buffer_cache is not None and key in self._buffer_cache: - return self._buffer_cache[key] + if self._buffer_cache is not None: + cached = self._buffer_cache.get(key, _MISSING) + if cached is not _MISSING: + return cached uuid, treepath, start, stop, partition, nodes = self.interpret_key(key) if self._debug: print( diff --git a/src/coffea/nanoevents/mapping/buffer_cache.py b/src/coffea/nanoevents/mapping/buffer_cache.py index f19cb868e..4224361de 100644 --- a/src/coffea/nanoevents/mapping/buffer_cache.py +++ b/src/coffea/nanoevents/mapping/buffer_cache.py @@ -24,6 +24,7 @@ def decode(self, buffer: ByteBuffer, struct: ShapeDTypeStruct) -> np.ndarray: .. class NoCompressionCodec(Codec): def encode(self, arr: np.ndarray) -> tuple[ByteBuffer, ShapeDTypeStruct]: + arr = np.ascontiguousarray(arr) struct = ShapeDTypeStruct( dtype=arr.dtype, shape=arr.shape, strides=arr.strides, compressed=False ) @@ -41,6 +42,7 @@ def __init__(self, codec: tp.Any) -> None: self._codec = codec def encode(self, arr: np.ndarray) -> tuple[ByteBuffer, ShapeDTypeStruct]: + arr = np.ascontiguousarray(arr) if arr.nbytes > 16: encoded = self._codec.encode(arr.tobytes()) struct = ShapeDTypeStruct( @@ -208,21 +210,23 @@ def BufferCache( if codec is not None: try: import numcodecs - except ModuleNotFoundError as err: - raise ModuleNotFoundError("""to use BufferCache, you must install numcodecs: + except ModuleNotFoundError: + numcodecs = None -pip install numcodecs + if numcodecs is not None and isinstance(codec, numcodecs.abc.Codec): + codec = NumCodecsWrapper(codec=codec) -or + if not isinstance(codec, Codec): + if numcodecs is None: + raise ModuleNotFoundError( + """to use BufferCache, you must install numcodecs: -conda install -c conda-forge numcodecs""") from err +pip install numcodecs - # auto-wrap for numcodecs.abc.Codec - if isinstance(codec, numcodecs.abc.Codec): - codec = NumCodecsWrapper(codec=codec) +or - # at this point we expect a proper Codec instance - if not isinstance(codec, Codec): +conda install -c conda-forge numcodecs""" + ) raise TypeError(f"codec must be an instance of Codec, got {type(codec)}") return CodecAwareCache(cache=cache, codec=codec) diff --git a/src/coffea/nanoevents/mapping/preloaded.py b/src/coffea/nanoevents/mapping/preloaded.py index b1784ba99..d4933334f 100644 --- a/src/coffea/nanoevents/mapping/preloaded.py +++ b/src/coffea/nanoevents/mapping/preloaded.py @@ -36,7 +36,14 @@ class PreloadedSourceMapping(BaseSourceMapping): def __init__( self, array_source, start, stop, cache=None, access_log=None, buffer_cache=None ): - super().__init__(array_source, start, stop, cache, access_log, buffer_cache) + super().__init__( + array_source, + start, + stop, + cache, + access_log, + buffer_cache=buffer_cache, + ) @classmethod def _extract_base_form(cls, column_source, force_to_i64=False): diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index afd8cef37..5b5f0b25a 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -35,6 +35,21 @@ def add(self, other): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.subtract, {"Candidate"}) + def subtract(self, other): + """Subtract a candidate from another elementwise using ``x``, ``y``, ``z``, ``t``, and ``charge`` components""" + return awkward.zip( + { + "x": self.x - other.x, + "y": self.y - other.y, + "z": self.z - other.z, + "t": self.t - other.t, + "charge": self.charge - other.charge, + }, + with_name="Candidate", + behavior=self.behavior, + ) + def sum(self, axis=-1): """Sum an array of vectors elementwise using ``x``, ``y``, ``z``, ``t``, and ``charge`` components""" return awkward.zip( diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index e62ed7216..8451ec65f 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -833,7 +833,7 @@ def jet(self): @jet.dask def jet(self, dask_array): collection = self.collection_map[self._collection_name()][0] - return dask_array.events()[collection]._apply_global_index(dask_array.jetIdxG) + return dask_array._events()[collection]._apply_global_index(dask_array.jetIdxG) @dask_property def pf(self): @@ -865,12 +865,12 @@ class AssociatedSV(base.NanoCollection): @dask_property def jet(self): - collection = self._events()[self.collection_map[self._collection_name()][0]] + collection = self.collection_map[self._collection_name()][0] return self._events()[collection]._apply_global_index(self.jetIdxG) @jet.dask def jet(self, dask_array): - collection = self._events()[self.collection_map[self._collection_name()][0]] + collection = self.collection_map[self._collection_name()][0] return dask_array._events()[collection]._apply_global_index(dask_array.jetIdxG) @dask_property diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 7de84c73e..e6938e925 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -590,13 +590,24 @@ def multiply(self, other): In reality, this directly adjusts ``pt``, ``eta``, ``phi`` and ``mass`` for performance """ - absother = abs(other) + if other < 0: + # (pt, eta, phi, mass) cannot represent t < 0, so match cartesian in x, y, z, t + return awkward.zip( + { + "x": self.x * other, + "y": self.y * other, + "z": self.z * other, + "t": self.t * other, + }, + with_name="LorentzVector", + behavior=self.behavior, + ) return awkward.zip( { - "pt": self.pt * absother, - "eta": self.eta * numpy.sign(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - "mass": self.mass * absother, + "pt": self.pt * other, + "eta": self.eta, + "phi": self.phi, + "mass": self.mass * other, }, with_name="PtEtaPhiMLorentzVector", behavior=self.behavior, @@ -605,16 +616,7 @@ def multiply(self, other): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - { - "pt": self.pt, - "eta": -self.eta, - "phi": self.phi % (2 * numpy.pi) - numpy.pi, - "mass": self.mass, - }, - with_name="PtEtaPhiMLorentzVector", - behavior=self.behavior, - ) + return self.multiply(-1) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): diff --git a/src/coffea/nanoevents/schemas/edm4hep.py b/src/coffea/nanoevents/schemas/edm4hep.py index b23bd3a7c..522d9b9d9 100755 --- a/src/coffea/nanoevents/schemas/edm4hep.py +++ b/src/coffea/nanoevents/schemas/edm4hep.py @@ -1,6 +1,7 @@ import copy import re import warnings +from functools import lru_cache from coffea.nanoevents import transforms from coffea.nanoevents.assets import edm4hep_ver @@ -98,6 +99,17 @@ def sort_dict(d): return {k: d[k] for k in sorted(d)} +@lru_cache(maxsize=None) +def load_edm4hep(version): + """Load and parse the edm4hep yaml for a version, caching the result. + + The returned ``(raw, parsed)`` dicts are treated as read-only by the schema, + so a single parse is shared across all schema builds for a given version. + """ + raw = edm4hep_ver[version]() + return raw, parse_yaml(raw, copy.deepcopy(raw)) + + class EDM4HEPSchema(BaseSchema): """Schema-builder for EDM4HEP root file structure. EDM4HEPSchema for edm4hep version 00.99.01 @@ -156,8 +168,7 @@ def __init__(self, base_form, *args, **kwargs): super().__init__(base_form) # Detect Collection Datatypes and create a datatype mixin - self.edm4hep = edm4hep_ver[self.edm4hep_version]() - self.parsed_edm4hep = parse_yaml(self.edm4hep, copy.deepcopy(self.edm4hep)) + self.edm4hep, self.parsed_edm4hep = load_edm4hep(self.edm4hep_version) self._create_mixin(base_form) self._form["fields"], self._form["contents"] = self._build_collections( @@ -218,6 +229,7 @@ def _create_mixin(self, base_form): for collection_name in self._form["fields"] if _all_collections.match(collection_name) } + self._all_collections = all_collections collections = { collection_name for collection_name in all_collections @@ -304,7 +316,7 @@ def _lookup_branch(self, collection_name, branch_name, key=None): Members = collection_edm4hep.get("Members", {}) VectorMembers = collection_edm4hep.get("VectorMembers", {}) OneToOneRelations = collection_edm4hep.get("OneToOneRelations", {}) - OneToManyRelations = collection_edm4hep.get("OneToOneRelations", {}) + OneToManyRelations = collection_edm4hep.get("OneToManyRelations", {}) composite_dict = { **Members, **VectorMembers, @@ -1045,13 +1057,9 @@ def _build_collections(self, field_names, input_contents): Builds all the collections with the necessary behaviors defined in the mixins dictionary """ branch_forms = {k: v for k, v in zip(field_names, input_contents)} - # All collection names + # All collection names (computed once in _create_mixin from the same fields) # Example: ReconstructedParticles or _ReconstructedParticle_clusters, etc - all_collections = { - collection_name.split("/")[0] - for collection_name in field_names - if _all_collections.match(collection_name) - } + all_collections = self._all_collections output = {} branch_forms = self._doc_strings(branch_forms, all_collections) diff --git a/src/coffea/nanoevents/schemas/fcc.py b/src/coffea/nanoevents/schemas/fcc.py index b73bc37e6..f43661bb5 100644 --- a/src/coffea/nanoevents/schemas/fcc.py +++ b/src/coffea/nanoevents/schemas/fcc.py @@ -404,7 +404,7 @@ def _unknown_collections(self, output, branch_forms, all_collections): output[record_name] = zip_forms( sort_dict(contents), record_name, - self._datatype_mixins.get(record_name, "NanoCollection"), + self.mixins_dictionary.get(record_name, "NanoCollection"), ) # If a branch is non-empty and is one of its kind (i.e. has no other associated branch) # call it a singleton and assign it directly to the output @@ -712,4 +712,7 @@ def get_schema(cls, version="latest"): elif version == "edm4hep1": return FCCSchema_edm4hep1 else: - pass + raise ValueError( + f"Invalid FCC schema version {version!r}. " + "Valid versions: 'latest', 'pre-edm4hep1', 'edm4hep1'." + ) diff --git a/src/coffea/nanoevents/transforms.py b/src/coffea/nanoevents/transforms.py index 10c97f5bd..16550d444 100644 --- a/src/coffea/nanoevents/transforms.py +++ b/src/coffea/nanoevents/transforms.py @@ -778,7 +778,7 @@ def get_index_ranges(begin, end): ) ranges = get_index_ranges_kernel(begin_end, awkward.ArrayBuilder()).snapshot() - if awkward.sum(ranges) == 0: # empty ranges, return a twice nested empty array + if awkward.count(ranges, axis=None) == 0: # no entries, twice nested empty array ranges = begin_end[begin_end < 0] return ranges @@ -923,48 +923,6 @@ def begin_end_mapping_nested_target_form(begin_form, end_form, target_form): # begin_end_mapping_with_xyzrecord -@numba.njit -def get_array_from_indices_xyzrecord_target_kernel(indices, target, builder): - for ev in range(len(indices)): - builder.begin_list() - for j in range(len(indices[ev])): - builder.begin_list() - for k in indices[ev][j]: - builder.begin_record() - builder.field("x").real(target[ev][k]["x"]) - builder.field("y").real(target[ev][k]["y"]) - builder.field("z").real(target[ev][k]["z"]) - builder.end_record() - builder.end_list() - builder.end_list() - return builder - - -def get_array_from_indices_xyzrecord_target(indices, target): - return get_array_from_indices_xyzrecord_target_kernel( - indices, target, awkward.ArrayBuilder() - ).snapshot() - - -def begin_end_mapping_with_xyzrecord(stack): - target = stack.pop() - end = stack.pop() - begin = stack.pop() - indices, o1, o2 = get_index_ranges(begin, end) - - if len(target.fields) == 0: # Target is a ListOffset type - raise RuntimeError("Target is a ListOffset.") - else: # Target is a Record type - if awkward.sum(awkward.num(target, axis=1)) == 0: # Empty Target - out = indices[indices < 0] # return an empty array - else: - if awkward.sum(awkward.num(indices, axis=1)) == 0: # Empty Indices - out = indices[indices < 0] # return an empty array - else: # The usual case when both of the indices and target are non-empty - out = get_array_from_indices_xyzrecord_target(indices, target) - stack.append(out) - - def begin_end_mapping_with_xyzrecord_form(begin_form, end_form, target_form): if not begin_form["class"].startswith("ListOffset"): raise RuntimeError diff --git a/src/coffea/processor/dask/__init__.py b/src/coffea/processor/dask/__init__.py index f68aa7664..8f1843cbb 100644 --- a/src/coffea/processor/dask/__init__.py +++ b/src/coffea/processor/dask/__init__.py @@ -89,4 +89,7 @@ def register_columncache(client): for p in client.run(lambda: set(get_worker().plugins)).values(): plugins |= p if ColumnCache.name not in plugins: - client.register_worker_plugin(ColumnCache()) + if hasattr(client, "register_plugin"): + client.register_plugin(ColumnCache()) + else: + client.register_worker_plugin(ColumnCache()) diff --git a/src/coffea/processor/executor.py b/src/coffea/processor/executor.py index 4cb039eec..12a256a59 100644 --- a/src/coffea/processor/executor.py +++ b/src/coffea/processor/executor.py @@ -16,7 +16,7 @@ Mapping, MutableMapping, ) -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field from functools import partial from io import BytesIO @@ -496,17 +496,18 @@ def _watcher( ) else: # Merge within process batch = FH.fetch(len(FH.completed)) - merged = _compress( - accumulate( - progress.track( - map(_decompress, (c for c in batch)), - task_id=p_idm, - total=progress._tasks[p_idm].total + len(batch), + if batch: + merged = _compress( + accumulate( + progress.track( + map(_decompress, (c for c in batch)), + task_id=p_idm, + total=progress._tasks[p_idm].total + len(batch), + ), + _decompress(merged), ), - _decompress(merged), - ), - executor.compression, - ) + executor.compression, + ) # Add checkpointing if executor.merging: @@ -652,8 +653,7 @@ class FuturesExecutor(ExecutorBase): Supply an additional executor to process merge jobs independently. An ``int`` will be interpreted as ``ProcessPoolExecutor(max_workers=int)``. tailtimeout : int, optional - Timeout requirement on job tails. Cancel all remaining jobs if none have finished - in the timeout window. + Deprecated and ignored; it never had an effect and will be removed in a future release. retries : int, optional Number of retries for failed tasks (default: 3) @@ -678,6 +678,12 @@ class FuturesExecutor(ExecutorBase): retries: int = 3 def __post_init__(self): + if self.tailtimeout is not None: + warnings.warn( + "tailtimeout has never had an effect and is deprecated; it will be removed in a future release", + DeprecationWarning, + stacklevel=2, + ) if not ( isinstance(self.merging, bool) or (isinstance(self.merging, tuple) and len(self.merging) == 3) @@ -732,21 +738,32 @@ def _processwith(pool, mergepool): raise e from None if isinstance(self.pool, concurrent.futures.Executor): - return _processwith(pool=self.pool, mergepool=self.mergepool) + with ExitStack() as stack: + mergepoolinstance = self._resolve_mergepool(stack) + return _processwith(pool=self.pool, mergepool=mergepoolinstance) else: # assume its a class then with ExitStack() as stack: poolinstance = stack.enter_context(self.pool(max_workers=self.workers)) - if self.mergepool is not None: - if isinstance(self.mergepool, int): - self.mergepool = concurrent.futures.ProcessPoolExecutor( - max_workers=self.mergepool - ) - mergepoolinstance = stack.enter_context(self.mergepool) - else: - mergepoolinstance = None + mergepoolinstance = self._resolve_mergepool(stack) return _processwith(pool=poolinstance, mergepool=mergepoolinstance) + def _resolve_mergepool(self, stack): + mp = self.mergepool + if mp is None or mp is False: + return None + if mp is True: + return stack.enter_context( + concurrent.futures.ProcessPoolExecutor(max_workers=self.workers) + ) + if isinstance(mp, int): + return stack.enter_context( + concurrent.futures.ProcessPoolExecutor(max_workers=mp) + ) + if isinstance(mp, concurrent.futures.Executor): + return mp + return stack.enter_context(mp(max_workers=self.workers)) + @dataclass class DaskExecutor(ExecutorBase): @@ -972,8 +989,7 @@ class ParslExecutor(ExecutorBase): Labels of the executors (from dfk.config.executors) that will process main jobs. Default is 'all'. Recommended is ``['merges']``, while passing ``label='merges'`` to the executor dedicated towards merge jobs. tailtimeout : int, optional - Timeout requirement on job tails. Cancel all remaining jobs if none have finished - in the timeout window. + Deprecated and ignored; it never had an effect and will be removed in a future release. retries : int, optional Number of retries for failed tasks (default: 3) @@ -994,6 +1010,12 @@ class ParslExecutor(ExecutorBase): retries: int = 3 def __post_init__(self): + if self.tailtimeout is not None: + warnings.warn( + "tailtimeout has never had an effect and is deprecated; it will be removed in a future release", + DeprecationWarning, + stacklevel=2, + ) if not ( isinstance(self.merging, bool) or (isinstance(self.merging, tuple) and len(self.merging) == 3) @@ -1275,25 +1297,21 @@ def automatic_retries( try: return func(*args, **kwargs) except Exception as e: + chain = _exception_chain(e) + if retries == retry_count: + if skipbadfiles and any(isinstance(c, skipbadfiles) for c in chain): + if use_result_type: + # surface the exception instead of silently skipping + # so the Runner can wrap it as Err + raise e + warnings.warn( + f"Skipping bad file after {retry_count + 1} attempts. The last exception was: {str(e)}" + ) + break + raise e warnings.warn( f"Performed attempt {retry_count + 1} out of {retries + 1}" ) - chain = _exception_chain(e) - if ( - skipbadfiles - and (retries == retry_count) - and any(isinstance(c, skipbadfiles) for c in chain) - ): - if use_result_type: - # surface the exception instead of silently skipping - # so the Runner can wrap it as Err - raise e - warnings.warn( - f"Skipping bad file after {retry_count + 1} attempts. The last exception was: {str(e)}" - ) - break - if not skipbadfiles or (retries == retry_count): - raise e retry_count += 1 @staticmethod @@ -1357,6 +1375,8 @@ def metadata_fetcher_root( uproot_options: dict, item: FileMeta, ) -> Accumulatable: + uproot_options = dict(uproot_options) + xrootdtimeout = uproot_options.pop("timeout", xrootdtimeout) with uproot.open( {item.filename: None}, timeout=xrootdtimeout, **uproot_options ) as file: @@ -1414,8 +1434,6 @@ def _preprocess_fileset_root(self, fileset: dict, uproot_options: dict) -> None: "unit": "file", "compression": None, } - if isinstance(self.pre_executor, (FuturesExecutor, ParslExecutor)): - pre_arg_override.update({"tailtimeout": None}) if isinstance(self.pre_executor, (DaskExecutor)): self.pre_executor.heavy_input = None pre_arg_override.update({"worker_affinity": False}) @@ -1454,8 +1472,6 @@ def _preprocess_fileset_parquet(self, fileset: dict) -> None: "unit": "file", "compression": None, } - if isinstance(self.pre_executor, (FuturesExecutor, ParslExecutor)): - pre_arg_override.update({"tailtimeout": None}) if isinstance(self.pre_executor, (DaskExecutor)): self.pre_executor.heavy_input = None pre_arg_override.update({"worker_affinity": False}) @@ -1485,6 +1501,20 @@ def _filter_badfiles(self, fileset: dict) -> list: ) return final_fileset + def _limit_preprocess_files(self, fileset): + if self.maxchunks is None: + return fileset + # each file yields at least one chunk, so at most maxchunks files per + # dataset need opening to satisfy maxchunks chunks + seen = defaultdict(int) + limited = [] + for filemeta in fileset: + if seen[filemeta.dataset] >= self.maxchunks: + continue + limited.append(filemeta) + seen[filemeta.dataset] += 1 + return limited + def _trace_preload( self, fileset: list[FileMeta], @@ -1585,11 +1615,11 @@ def _work_function( checkpointer: CheckpointerABC, cache_function: Callable[[], MutableMapping], ) -> dict: - if "timeout" in uproot_options: - xrootdtimeout = uproot_options["timeout"] + uproot_options = dict(uproot_options) + xrootdtimeout = uproot_options.pop("timeout", xrootdtimeout) if processor_instance == "heavy": item, processor_instance = item - if not isinstance(processor_instance, ProcessorABC) or not callable( + if not isinstance(processor_instance, ProcessorABC) and not callable( processor_instance ): processor_instance = cloudpickle.loads(lz4f.decompress(processor_instance)) @@ -1803,6 +1833,9 @@ def __call__( ) if self.use_dataframes: return wrapped_out # not wrapped anymore + exception = wrapped_out.get("exception", 0) + if exception != 0: + raise exception if self.savemetrics: return wrapped_out["out"], wrapped_out["metrics"] return wrapped_out["out"] @@ -1871,7 +1904,9 @@ def preprocess( if uproot_options is None: uproot_options = {} if self.format == "root": - fileset = list(self._normalize_fileset(fileset, treename)) + fileset = self._limit_preprocess_files( + list(self._normalize_fileset(fileset, treename)) + ) for filemeta in fileset: filemeta.maybe_populate(self.metadata_cache) @@ -1885,7 +1920,9 @@ def preprocess( process_fn = processor_instance self._trace_preload(fileset, trace, process_fn, uproot_options) elif self.format == "parquet": - fileset = list(self._normalize_fileset(fileset, treename)) + fileset = self._limit_preprocess_files( + list(self._normalize_fileset(fileset, treename)) + ) if any(filemeta.preload is not None for filemeta in fileset): raise NotImplementedError( "fileset-level 'preload' is not supported for parquet input" @@ -1982,6 +2019,30 @@ def run( return Err(e) raise + @contextmanager + def _auto_dask_client(self): + auto = [] + if isinstance(self.executor, DaskExecutor) and self.executor.client is None: + auto.append(self.executor) + if ( + self.pre_executor is not self.executor + and isinstance(self.pre_executor, DaskExecutor) + and self.pre_executor.client is None + ): + auto.append(self.pre_executor) + if not auto: + yield None + return + client = _import_distributed().client.Client(threads_per_worker=1) + for ex in auto: + ex.client = client + try: + yield client + finally: + for ex in auto: + ex.client = None + client.close() + def _run( self, fileset: dict | str | list[WorkItem] | Generator, @@ -1991,6 +2052,26 @@ def _run( uproot_options: dict | None = {}, iteritems_options: dict | None = {}, trace: Callable | None = None, + ) -> Accumulatable: + with self._auto_dask_client(): + return self._run_impl( + fileset, + processor_instance, + treename=treename, + uproot_options=uproot_options, + iteritems_options=iteritems_options, + trace=trace, + ) + + def _run_impl( + self, + fileset: dict | str | list[WorkItem] | Generator, + processor_instance: ProcessorABC | Callable[[awkward.highlevel.Array], Any], + *, + treename: str | None = None, + uproot_options: dict | None = {}, + iteritems_options: dict | None = {}, + trace: Callable | None = None, ) -> Accumulatable: if uproot_options is None: uproot_options = {} diff --git a/src/coffea/processor/parsl/detail.py b/src/coffea/processor/parsl/detail.py index e34ae78d6..e8c4e7241 100644 --- a/src/coffea/processor/parsl/detail.py +++ b/src/coffea/processor/parsl/detail.py @@ -1,12 +1,8 @@ import parsl -from parsl.app.app import python_app from parsl.config import Config from parsl.executors import HighThroughputExecutor from parsl.providers import LocalProvider -from ..executor import _futures_handler -from .timeout import timeout - _default_cfg = Config( executors=[ HighThroughputExecutor( @@ -30,58 +26,3 @@ def _parsl_initialize(config=None): def _parsl_stop(): parsl.dfk().cleanup() parsl.clear() - - -@timeout -@python_app -def derive_chunks(filename, treename, chunksize, ds, timeout=10): - from collections.abc import Sequence - - import uproot - - uproot.XRootDSource.defaults["parallel"] = False - - a_file = uproot.open({filename: None}) - - tree = None - if isinstance(treename, str): - tree = a_file[treename] - elif isinstance(treename, Sequence): - for name in reversed(treename): - if name in a_file: - tree = a_file[name] - else: - raise Exception( - "treename must be a str or Sequence but is a %s!" % repr(type(treename)) - ) - - if tree is None: - raise Exception( - "No tree found, out of possible tree names: %s" % repr(treename) - ) - - nentries = tree.numentries - return ( - ds, - treename, - [(filename, chunksize, index) for index in range(nentries // chunksize + 1)], - ) - - -def _parsl_get_chunking(filelist, chunksize, status=True, timeout=10): - futures = { - derive_chunks(fn, tn, chunksize, ds, timeout=timeout) for ds, fn, tn in filelist - } - - items = [] - - def chunk_accumulator(total, result): - ds, treename, chunks = result - for chunk in chunks: - total.append((ds, chunk[0], treename, chunk[1], chunk[2])) - - _futures_handler( - futures, items, status, "files", "Preprocessing", chunk_accumulator, None - ) - - return items diff --git a/src/coffea/processor/taskvine_executor.py b/src/coffea/processor/taskvine_executor.py index 74a00109a..08d24f5c9 100644 --- a/src/coffea/processor/taskvine_executor.py +++ b/src/coffea/processor/taskvine_executor.py @@ -1,4 +1,5 @@ import collections +import functools import math import os import re @@ -408,8 +409,11 @@ def _processing(self, items, function, accumulator): function = _compression_wrapper(self.executor.compression, function) accumulate_fn = _compression_wrapper( self.executor.compression, - accumulate_result_files, - self.executor.concurrent_reads, + functools.partial( + accumulate_result_files, + concurrent_reads=self.executor.concurrent_reads, + ), + name="accumulate_result_files", ) sc = self.stats_coffea @@ -973,7 +977,6 @@ def debug_info(self): def _handle_early_terminate(signum, frame, raise_on_repeat=True): global early_terminate - raise KeyboardInterrupt if early_terminate and raise_on_repeat: raise KeyboardInterrupt diff --git a/src/coffea/util.py b/src/coffea/util.py index 14672a562..a94948405 100644 --- a/src/coffea/util.py +++ b/src/coffea/util.py @@ -11,7 +11,6 @@ import cloudpickle import fsspec import hist -import numba import numpy import uproot from rich.console import Console @@ -28,7 +27,6 @@ ak = awkward np = numpy -nb = numba __all__ = [ @@ -75,20 +73,6 @@ def save(output, filename, compression="lz4"): cloudpickle.dump(output, fout) -def _hex(string): - try: - return string.hex() - except AttributeError: - return "".join(f"{ord(c):02x}" for c in string) - - -def _ascii(maybebytes): - try: - return maybebytes.decode("ascii") - except AttributeError: - return maybebytes - - def _hash(items): # python 3.3 salts hash(), we want it to persist across processes x = hashlib.md5(bytes(";".join(str(x) for x in items), "ascii")) diff --git a/tests/samples/photon_id_2d.ea.txt b/tests/samples/photon_id_2d.ea.txt new file mode 100644 index 000000000..d43041712 --- /dev/null +++ b/tests/samples/photon_id_2d.ea.txt @@ -0,0 +1,5 @@ +{2 AbsEta Pt 1 EA_Pho} +0.0 1.0 0.0 20.0 0.10 +0.0 1.0 20.0 50.0 0.12 +1.0 2.0 0.0 20.0 0.20 +1.0 2.0 20.0 50.0 0.22 diff --git a/tests/samples/testpu.pileup.json b/tests/samples/testpu.pileup.json new file mode 100644 index 000000000..71731190f --- /dev/null +++ b/tests/samples/testpu.pileup.json @@ -0,0 +1 @@ +{"1": [[1, 0.0, 0.0, 25.0], [2, 0.0, 0.0, 30.0]], "2": [[1, 0.0, 0.0, 12.0]]} diff --git a/tests/test_analysis_tools.py b/tests/test_analysis_tools.py index 03d65812a..3da27ae37 100644 --- a/tests/test_analysis_tools.py +++ b/tests/test_analysis_tools.py @@ -2275,3 +2275,109 @@ def test_packed_selection_cutflow_dak_uproot_only(optimization_enabled): counts[-2] += counts[-1] c, e = np.histogram(dak.flatten(array[truth]).compute(), bins=edges) assert np.all(np.isclose(counts[1:-1], c)) + + +def test_weight_statistics_add_returns_object(): + from coffea.analysis_tools import WeightStatistics + + a = WeightStatistics(sumw=1.0, sumw2=1.0, minw=0.5, maxw=2.0, n=3) + b = WeightStatistics(sumw=2.0, sumw2=4.0, minw=0.1, maxw=3.0, n=5) + + c = a + b + assert isinstance(c, WeightStatistics) + assert (c.sumw, c.sumw2, c.n, c.minw, c.maxw) == (3.0, 5.0, 8, 0.1, 3.0) + assert a.sumw == 1.0 and a.n == 3 # operand untouched by __add__ + + acc = WeightStatistics() + for _ in range(3): + acc += b # must not rebind to None + assert isinstance(acc, WeightStatistics) + assert acc.n == 15 and acc.sumw == 6.0 + + +def test_weights_multivariation_dak_option_fill(): + dak = pytest.importorskip("dask_awkward") + import awkward as ak + + from coffea.analysis_tools import Weights + + weight = dak.from_awkward(ak.Array([1.0, None, 2.0, 3.0]), npartitions=1) + weights = Weights(None) + weights.add_multivariation( + "test", weight, modifierNames=[], weightsUp=[], weightsDown=[] + ) + result = weights.weight().compute() + assert not ak.any(ak.is_none(result)) # option filled with 1.0, like eager path + assert ak.to_list(result) == [1.0, 1.0, 2.0, 3.0] + + +@pytest.mark.parametrize("mode", ["eager", "delayed"]) +def test_weights_modifier_suffix_and_multivariation(mode): + from coffea.analysis_tools import Weights + + if mode == "delayed": + dak = pytest.importorskip("dask_awkward") + import dask.array as da + + def arr(x): + return dak.from_dask_array(da.from_array(np.asarray(x, dtype=float))) + + def get(x): + return np.asarray(x.compute()) + + size = None + else: + + def arr(x): + return np.asarray(x, dtype=float) + + def get(x): + return np.asarray(x) + + size = 4 + + central = arr([1.0, 1.0, 1.0, 1.0]) + up = arr([1.25, 1.25, 1.25, 1.25]) + down = arr([0.8, 0.8, 0.8, 0.8]) + + weights = Weights(size, storeIndividual=True) + weights.add_multivariation( + "test", central, modifierNames=["A"], weightsUp=[up], weightsDown=[down] + ) + # weight name carrying 'Up'/'Down' mid-string must survive suffix handling + weights.add("myUpdate", central, weightUp=up) + + # bug 41: multivariation modifiers must be recognized by partial_weight + assert np.allclose( + get(weights.partial_weight(include=["test"], modifier="test_AUp")), 1.25 + ) + assert np.allclose( + get(weights.partial_weight(include=["test"], modifier="test_ADown")), 0.8 + ) + + # bug 46: 'Up' inside a weight name must not be stripped mid-string + assert np.allclose( + get(weights.partial_weight(include=["myUpdate"], modifier="myUpdateUp")), 1.25 + ) + + # bug 46: auto-derived Down for a name containing 'Down'-like fragments + assert np.allclose(get(weights.weight("myUpdateDown")), 0.8) + + # modifier of an excluded weight is still rejected + with pytest.raises(ValueError, match="not in the list of included weights"): + weights.partial_weight(include=["myUpdate"], modifier="test_AUp") + + +def test_packed_selection_require_returns_independent_copy(): + from coffea.analysis_tools import PackedSelection + + selection = PackedSelection() + selection.add("c1", np.array([True, True, False, False])) + selection.add("c2", np.array([True, False, True, False])) + + mask = selection.all("c1") + original = mask.copy() + mask &= np.array([False, False, False, False]) # in-place mutation by caller + + assert np.array_equal(selection.all("c1"), original) # cache not corrupted + assert selection.all("c1") is not selection.all("c1") # independent copies diff --git a/tests/test_buffer_cache.py b/tests/test_buffer_cache.py index 89ab6acfa..294fd4e1b 100644 --- a/tests/test_buffer_cache.py +++ b/tests/test_buffer_cache.py @@ -135,6 +135,85 @@ def test_buffer_cache_hierarchical(tests_directory): os.rmdir(f"{tests_directory}/mycache") +@pytest.mark.parametrize("codec_name", ["none", "numcodecs"]) +@pytest.mark.parametrize("layout", ["contiguous", "sliced", "fortran"]) +def test_buffer_cache_roundtrip_layouts(codec_name, layout): + from coffea.nanoevents.mapping.buffer_cache import ( + CodecAwareCache, + NoCompressionCodec, + ) + + if codec_name == "none": + codec = NoCompressionCodec() + else: + pytest.importorskip("numcodecs") + from numcodecs import Blosc + + from coffea.nanoevents.mapping.buffer_cache import NumCodecsWrapper + + codec = NumCodecsWrapper(Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE)) + + base = np.arange(64, dtype=np.int64) + if layout == "contiguous": + arr = base.copy() + elif layout == "sliced": + arr = base[::2] + else: + arr = np.asfortranarray(base.reshape(8, 8)) + assert layout == "contiguous" or not arr.flags["C_CONTIGUOUS"] + + cache = CodecAwareCache(cache={}, codec=codec) + cache["key"] = arr + np.testing.assert_array_equal(cache["key"], arr) + + +def test_no_compression_codec_without_numcodecs(monkeypatch): + import builtins + + from coffea.nanoevents.mapping import buffer_cache as bc_mod + from coffea.nanoevents.mapping.buffer_cache import ( + CodecAwareCache, + NoCompressionCodec, + ) + + real_import = builtins.__import__ + + def _no_numcodecs(name, *args, **kwargs): + if name == "numcodecs": + raise ModuleNotFoundError("numcodecs is not available") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _no_numcodecs) + + cache = bc_mod.BufferCache(cache=None, codec=NoCompressionCodec()) + assert isinstance(cache, CodecAwareCache) + + +def test_buffer_cache_decodes_once_per_hit(): + from coffea.nanoevents.mapping.buffer_cache import ( + CodecAwareCache, + NoCompressionCodec, + ) + from coffea.nanoevents.mapping.preloaded import PreloadedSourceMapping + + class CountingCodec(NoCompressionCodec): + def __init__(self): + self.decodes = 0 + + def decode(self, buffer, struct): + self.decodes += 1 + return super().decode(buffer, struct) + + codec = CountingCodec() + buffer_cache = CodecAwareCache(cache={}, codec=codec) + mapping = PreloadedSourceMapping(object(), 0, 5, buffer_cache=buffer_cache) + buffer_cache["mykey"] = np.arange(5, dtype=np.int64) + + out = mapping["mykey"] + np.testing.assert_array_equal(out, np.arange(5, dtype=np.int64)) + assert codec.decodes == 1 + + def test_buffer_cache_small_and_empty_array_compression(): pytest.importorskip("numcodecs") pytest.importorskip("zict") diff --git a/tests/test_dataset_tools_rucio.py b/tests/test_dataset_tools_rucio.py new file mode 100644 index 000000000..d4f5b62ff --- /dev/null +++ b/tests/test_dataset_tools_rucio.py @@ -0,0 +1,77 @@ +import sys +import types as _types + + +def _install_fake_rucio(): + if "rucio" in sys.modules: + return + rucio = _types.ModuleType("rucio") + client_mod = _types.ModuleType("rucio.client") + + class Client: + pass + + client_mod.Client = Client + rucio.client = client_mod + sys.modules["rucio"] = rucio + sys.modules["rucio.client"] = client_mod + + +_install_fake_rucio() + +from coffea.dataset_tools import rucio_utils # noqa: E402 + + +class FakeClient: + def __init__(self, replicas): + self._replicas = replicas + + def list_replicas(self, dids): + return self._replicas + + +def _filedata(name, sites): + rses = {} + pfns = {} + states = {} + for site, available in sites.items(): + key = f"root://{site}/{name}" + rses[site] = [key] + pfns[key] = {"type": "DISK", "volatile": False} + states[site] = "AVAILABLE" if available else "UNAVAILABLE" + return {"name": name, "rses": rses, "pfns": pfns, "states": states} + + +_PREFIXES = {"T2_A": "root://a/", "T2_B": "root://b/"} + + +def test_first_partial_allowed_skips_unviable(monkeypatch): + monkeypatch.setattr(rucio_utils, "get_xrootd_sites_map", lambda: _PREFIXES) + client = FakeClient( + [ + _filedata("/f1.root", {"T2_A": True}), + _filedata("/f2.root", {"T2_B": False}), + ] + ) + files, sites, counts = rucio_utils.get_dataset_files_replicas( + "ds", client=client, mode="first", partial_allowed=True + ) + assert files == ["root://a//f1.root"] + assert sites == ["T2_A"] + assert dict(counts) == {"T2_A": 1} + + +def test_first_sites_counts_not_stale(monkeypatch): + monkeypatch.setattr(rucio_utils, "get_xrootd_sites_map", lambda: _PREFIXES) + client = FakeClient( + [ + _filedata("/f1.root", {"T2_A": True}), + _filedata("/f2.root", {"T2_B": True}), + ] + ) + files, sites, counts = rucio_utils.get_dataset_files_replicas( + "ds", client=client, mode="first" + ) + assert files == ["root://a//f1.root", "root://b//f2.root"] + assert sites == ["T2_A", "T2_B"] + assert dict(counts) == {"T2_A": 1, "T2_B": 1} diff --git a/tests/test_jetmet_tools.py b/tests/test_jetmet_tools.py index 4453358f0..f079dde52 100644 --- a/tests/test_jetmet_tools.py +++ b/tests/test_jetmet_tools.py @@ -1220,3 +1220,95 @@ def smear_factor(jetPt, pt_gen, jersf): print("build all met variations =", toc - tic) print(prof.output_text(unicode=True, color=True, show_all=True)) + + +_JEC_ONLY_NAMES = [ + "Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi", + "Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi", + "Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi", + "Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi", +] + + +def _jec_only_setup(with_raw): + from coffea.jetmet_tools import JECStack + + counts = np.array([2, 0, 1]) + fields = { + "pt": ak.unflatten(np.array([50.0, 100.0, 30.0]), counts), + "mass": ak.unflatten(np.array([10.0, 20.0, 5.0]), counts), + "eta": ak.unflatten(np.array([0.5, -1.2, 2.0]), counts), + "area": ak.unflatten(np.array([0.5, 0.5, 0.5]), counts), + "Rho": ak.unflatten(np.array([15.0, 15.0, 15.0]), counts), + } + if with_raw: + fields["pt_raw"] = fields["pt"] * 0.95 + fields["mass_raw"] = fields["mass"] * 0.95 + jets = ak.zip(fields) + + jec_stack = JECStack({name: evaluator[name] for name in _JEC_ONLY_NAMES}) + name_map = jec_stack.blank_name_map + name_map["JetPt"] = "pt" + name_map["JetMass"] = "mass" + name_map["JetEta"] = "eta" + name_map["JetA"] = "area" + name_map["Rho"] = "Rho" + if with_raw: + name_map["ptRaw"] = "pt_raw" + name_map["massRaw"] = "mass_raw" + return jets, jec_stack, name_map + + +def test_corrected_jets_factory_no_raw_name_map(): + from coffea.jetmet_tools import CorrectedJetsFactory, FactorizedJetCorrector + + jets, jec_stack, name_map = _jec_only_setup(with_raw=False) + with pytest.warns(UserWarning): + jet_factory = CorrectedJetsFactory(name_map, jec_stack) + assert jet_factory.treat_pt_as_raw + assert jet_factory.name_map["ptRaw"] == "pt_raw" + assert jet_factory.name_map["massRaw"] == "mass_raw" + + corrected_jets = jet_factory.build(jets) + + corrector = FactorizedJetCorrector( + **{name: evaluator[name] for name in _JEC_ONLY_NAMES} + ) + check_corrs = corrector.getCorrection( + JetEta=jets.eta, Rho=jets.Rho, JetPt=jets.pt, JetA=jets.area + ) + assert ak.all(np.abs(corrected_jets.pt - check_corrs * jets.pt) < 1e-6) + assert ak.all(np.abs(corrected_jets.mass - check_corrs * jets.mass) < 1e-6) + assert ak.all(corrected_jets.pt_raw == jets.pt) + assert ak.all(corrected_jets.mass_raw == jets.mass) + + +def test_rand_gauss_empty(): + from coffea.jetmet_tools.CorrectedJetsFactory import rand_gauss + + out = rand_gauss(ak.Array(np.array([], dtype=np.float32))) + assert len(out) == 0 + assert out.to_numpy().dtype == np.dtype("float32") + + +def test_corrected_jets_factory_does_not_mutate_input(): + from coffea.jetmet_tools import CorrectedJetsFactory + + jets, jec_stack, name_map = _jec_only_setup(with_raw=True) + with pytest.warns(UserWarning): + jet_factory = CorrectedJetsFactory(name_map, jec_stack) + + corrected_jets = jet_factory.build(jets) + + assert corrected_jets.layout.content.parameters.get("corrected") is True + assert "corrected" not in jets.layout.content.parameters + + +def test_corrected_jets_factory_empty_record_message(): + from coffea.jetmet_tools import CorrectedJetsFactory + + _, jec_stack, name_map = _jec_only_setup(with_raw=True) + with pytest.warns(UserWarning): + jet_factory = CorrectedJetsFactory(name_map, jec_stack) + with pytest.raises(Exception, match="'pt'"): + jet_factory.build(ak.Array([[{}], []])) diff --git a/tests/test_local_executors.py b/tests/test_local_executors.py index cc5ec8b83..285881da3 100644 --- a/tests/test_local_executors.py +++ b/tests/test_local_executors.py @@ -464,3 +464,204 @@ def test_use_result_type_requires_skipbadfiles(): use_result_type=True, skipbadfiles=(FileNotFoundError,), ) + + +def _one(x): + return {"n": 1} + + +def test_processor_compression_none_runs(): + # Bug 11: processor_compression=None must not feed raw processor to lz4f.decompress + run = processor.Runner( + executor=processor.IterativeExecutor(), + schema=schemas.NanoAODSchema, + processor_compression=None, + ) + out = run( + _good_fileset, + processor_instance=NanoEventsProcessor(mode="eager"), + treename="Events", + ) + assert out["cutflow"]["ZJets_pt"] == 18 + assert out["cutflow"]["ZJets_mass"] == 6 + + +@pytest.mark.parametrize("seam", ["metadata", "run"]) +def test_uproot_options_timeout(seam): + # Bug 21: a timeout inside uproot_options must not be passed twice to uproot.open + from coffea.processor.executor import FileMeta, Runner + + if seam == "metadata": + item = FileMeta("ZJets", osp.abspath("tests/samples/nano_dy.root"), "Events") + out = Runner.metadata_fetcher_root(60, False, {"timeout": 30}, item) + (fetched,) = out + assert fetched.metadata["numentries"] == 40 + else: + run = processor.Runner( + executor=processor.IterativeExecutor(), + schema=schemas.NanoAODSchema, + ) + out = run( + _good_fileset, + processor_instance=NanoEventsProcessor(mode="eager"), + treename="Events", + uproot_options={"timeout": 30}, + ) + assert out["cutflow"]["ZJets_pt"] == 18 + + +@pytest.mark.parametrize("skipbadfiles", [False, (ValueError,)]) +def test_automatic_retries_retry_without_skipbadfiles(skipbadfiles): + # Bug 22: retries must happen regardless of skipbadfiles + calls = {"n": 0} + + def flaky(): + calls["n"] += 1 + if calls["n"] < 3: + raise ValueError("transient") + return "ok" + + out = processor.Runner.automatic_retries(3, skipbadfiles, flaky) + assert out == "ok" + assert calls["n"] == 3 + + +@pytest.mark.parametrize("form", ["int", "class", "instance", "bool"]) +def test_futures_mergepool_forms(form): + # Bug 25: every documented mergepool form must work and not mutate self.mergepool + import concurrent.futures as cf + + mp = { + "int": 2, + "class": cf.ProcessPoolExecutor, + "instance": cf.ProcessPoolExecutor(max_workers=2), + "bool": True, + }[form] + ex = processor.FuturesExecutor( + workers=2, merging=True, mergepool=mp, compression=None, status=False + ) + out, code = ex(list(range(6)), _one, None) + assert out == {"n": 6} + assert code == 0 + assert ex.mergepool is mp + if isinstance(mp, cf.Executor): + mp.shutdown() + + +def test_recoverable_exception_raised_in_default_path(): + # Bug 47: use_result_type=False must not silently drop a captured exception + run = processor.Runner( + executor=processor.IterativeExecutor(), + schema=schemas.NanoAODSchema, + ) + boom = RuntimeError("recovered failure") + + def fake_run(**kwargs): + return {"out": {"cutflow": {"events": 1}}, "exception": boom} + + run.run = fake_run + with pytest.raises(RuntimeError, match="recovered failure"): + run( + {"x": {"files": {"f.root": "Events"}}}, + processor_instance=NanoEventsProcessor(mode="eager"), + ) + + +def test_maxchunks_limits_preprocess_file_opening(): + # Bug 56: maxchunks=1 must not open every file during preprocessing + fileset = { + "ZJets": { + "treename": "Events", + "files": [ + osp.abspath("tests/samples/nano_dy.root"), + osp.abspath("tests/samples/non_existent.root"), + ], + } + } + run = processor.Runner( + executor=processor.IterativeExecutor(retries=0), + schema=schemas.NanoAODSchema, + maxchunks=1, + ) + chunks = list(run.preprocess(fileset)) + assert len(chunks) == 1 + assert chunks[0].filename.endswith("nano_dy.root") + + +def test_watcher_recompresses_only_on_change(monkeypatch): + # Bug 55: the running accumulator must not be recompressed on idle poll cycles + from coffea.processor import executor as ex_mod + + item = ex_mod._compress({"n": 1}, 1) + calls = {"n": 0} + real_compress = ex_mod._compress + + def counting(value, compression): + calls["n"] += 1 + return real_compress(value, compression) + + monkeypatch.setattr(ex_mod, "_compress", counting) + + class StubExec: + desc = "Processing" + unit = "items" + merging = False + compression = 1 + + class StubFH: + def __init__(self): + self.running = 1 + self.done = {"futures": 0, "merges": 0} + self.completed = [] + self.futures = [] + self.merges = [] + self._polls = 0 + + def update(self): + self._polls += 1 + if self._polls >= 3: + self.completed = [item] + self.running = 0 + + def fetch(self, n): + batch = self.completed[:n] + self.completed = self.completed[n:] + return batch + + out = ex_mod._watcher(StubFH(), StubExec(), lambda b: b, None) + assert calls["n"] == 1 + assert ex_mod._decompress(out) == {"n": 1} + + +def test_auto_dask_client_single_and_cleanup(monkeypatch): + # Bug 24: at most one auto-created dask Client, assigned to both executors and + # cleaned up afterwards + from coffea.processor import executor as ex_mod + + created = [] + + class FakeClient: + def __init__(self, **kwargs): + self.closed = False + created.append(self) + + def close(self): + self.closed = True + + class FakeDistributed: + class client: + Client = FakeClient + + monkeypatch.setattr(ex_mod, "_import_distributed", lambda: FakeDistributed) + + executor = processor.DaskExecutor() + runner = processor.Runner(executor=executor) + assert executor.client is None + with runner._auto_dask_client() as client: + assert client is created[0] + assert executor.client is client + assert runner.pre_executor.client is client + assert len(created) == 1 + assert len(created) == 1 + assert executor.client is None + assert created[0].closed is True diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index d3495ec7a..e9a6b8ca6 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -385,6 +385,36 @@ def test_jec_txt_effareas(): print(evaluator["photon_id_EA_Pho"]) +def test_effective_area_2d_binning(): + from coffea.lookup_tools.txt_converters import convert_effective_area_file + + wrapped = convert_effective_area_file("tests/samples/photon_id_2d.ea.txt") + assert any(name.startswith("photon_id_2d") for (name, _t) in wrapped.keys()) + + +def test_pileup_json_wildcard(): + ext = lookup_tools.extractor() + ext.add_weight_sets(["* * tests/samples/testpu.pileup.json"]) + ext.finalize() + ev = ext.make_evaluator() + out = ev["pileup"](np.array([2]), np.array([1])) + assert list(out) == [12.0] + + +def test_pileup_json_named_no_nametable_corruption(): + ext = lookup_tools.extractor() + ext.add_weight_sets( + [ + "photon_id_EA_Pho photon_id_EA_Pho tests/samples/photon_id.ea.txt", + "pu pileup tests/samples/testpu.pileup.json", + ] + ) + ext.finalize() + ev = ext.make_evaluator() + out = ev["pu"](np.array([1, 1, 2]), np.array([1, 2, 1])) + assert list(out) == [25.0, 30.0, 12.0] + + def test_rochester(tests_directory): dak = pytest.importorskip("dask_awkward") rochester_data = lookup_tools.txt_converters.convert_rochester_file( diff --git a/tests/test_lumi_tools.py b/tests/test_lumi_tools.py index 799b022c3..1064b03b7 100644 --- a/tests/test_lumi_tools.py +++ b/tests/test_lumi_tools.py @@ -58,6 +58,40 @@ def test_lumidata(): assert len(results["index"][lumidata]) == len(results["index"][lumidata_pickle]) +@pytest.mark.parametrize("order", ["eager_first", "dask_first"]) +@pytest.mark.parametrize("dtype", ["u4", "i8", "f8"]) +def test_lumidata_eager_dask_order_and_dtype(order, dtype): + dak = pytest.importorskip("dask_awkward") + import dask + + lumidata = LumiData("tests/samples/lumi_small.csv") + + base = np.zeros((10, 2), dtype="u4") + base[:, 0] = lumidata._lumidata[0:10, 0] + base[:, 1] = lumidata._lumidata[0:10, 1] + + baseline = LumiData("tests/samples/lumi_small.csv").get_lumi(base) + + eager_input = base.astype(dtype) + dask_input = dak.from_awkward(ak.Array(base), 3) + + def run_eager(): + return lumidata.get_lumi(eager_input) + + def run_dask(): + return dask.compute(lumidata.get_lumi(dask_input))[0] + + if order == "eager_first": + eager_result = run_eager() + dask_result = run_dask() + else: + dask_result = run_dask() + eager_result = run_eager() + + assert eager_result == baseline + assert dask_result == baseline + + @pytest.mark.dask_client @pytest.mark.network @pytest.mark.parametrize( diff --git a/tests/test_ml_tools.py b/tests/test_ml_tools.py index e08abd30d..d695e9420 100644 --- a/tests/test_ml_tools.py +++ b/tests/test_ml_tools.py @@ -1,3 +1,6 @@ +import importlib +import types + import awkward as ak import numpy as np import pytest @@ -5,6 +8,116 @@ dak = pytest.importorskip("dask_awkward") +def test_tf_wrapper_kwargs_writeable_flag(monkeypatch): + tfmod = importlib.import_module("coffea.ml_tools.tf_wrapper") + + monkeypatch.setattr( + tfmod, + "tensorflow", + types.SimpleNamespace(convert_to_tensor=lambda a: a), + raising=False, + ) + arr = np.ones((3, 2)) + expected = np.arange(6).reshape(3, 2) + fake_self = types.SimpleNamespace( + skip_length_zero=False, + model=lambda *a, **k: types.SimpleNamespace(numpy=lambda: expected), + ) + out = tfmod.tf_wrapper.numpy_call(fake_self, **{"feat": arr}) + assert np.array_equal(out, expected) + + +def test_triton_run_infer_backoff_seconds(monkeypatch): + twmod = importlib.import_module("coffea.ml_tools.triton_wrapper") + + class FakeInferException(Exception): + pass + + monkeypatch.setattr( + twmod, + "tritonclient", + types.SimpleNamespace( + utils=types.SimpleNamespace(InferenceServerException=FakeInferException) + ), + raising=False, + ) + + sleeps = [] + monkeypatch.setattr(twmod.time, "sleep", lambda s: sleeps.append(s)) + monkeypatch.setattr(twmod.numpy.random, "rand", lambda: 1.0) + + calls = {"n": 0} + + def infer(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise FakeInferException() + return "ok" + + class _tw(twmod.triton_wrapper): + def prepare_awkward(self, *a, **k): + return [], {} + + obj = _tw.__new__(_tw) + obj.client = types.SimpleNamespace(infer=infer) + obj.model, obj.version = "m", "1" + + assert obj.run_infer(inputs=None, outputs=None) == "ok" + # retry_jitter_base_ms=100 -> attempt 0 with max jitter is 0.1 s, not 100 s + assert sleeps == [pytest.approx(0.1)] + + +def test_triton_numpy_call_single_concatenate(monkeypatch): + twmod = importlib.import_module("coffea.ml_tools.triton_wrapper") + + monkeypatch.setattr( + twmod, + "tritonclient", + types.SimpleNamespace( + utils=types.SimpleNamespace(triton_to_np_dtype=lambda dt: np.float32) + ), + raising=False, + ) + + real_concat = twmod.numpy.concatenate + concat_calls = {"n": 0} + + def counting_concat(seq, *a, **k): + if len(tuple(seq)) >= 2: + concat_calls["n"] += 1 + return real_concat(seq, *a, **k) + + monkeypatch.setattr(twmod.numpy, "concatenate", counting_concat) + + class FakeInferInput: + def __init__(self, name, shape, dtype): + self.name = name + + def set_data_from_numpy(self, arr): + self.arr = arr + + class _tw(twmod.triton_wrapper): + def prepare_awkward(self, *a, **k): + return [], {} + + _tw.pmod = types.SimpleNamespace( + InferInput=FakeInferInput, InferRequestedOutput=lambda o: o + ) + obj = _tw.__new__(_tw) + obj._batch_size = 2 + obj.model_inputs = {"x": {"shape": (-1, 3), "datatype": "FP32"}} + obj.model_outputs = {"y": {"shape": (-1, 3)}} + obj.run_infer = lambda inputs, outputs: types.SimpleNamespace( + as_numpy=lambda o: inputs[0].arr.copy() + ) + + x = np.arange(6 * 3).reshape(6, 3).astype(float) # 3 full batches of size 2 + out = obj.numpy_call(["y"], {"x": x}) + + assert np.array_equal(out["y"], x) + assert concat_calls["n"] == 1 + + def prepare_jets_array(njets, tmp_path): # Creating jagged Jet-with-constituent array, returning both awkward and lazy # dask_awkward arrays diff --git a/tests/test_nanoevents.py b/tests/test_nanoevents.py index cf48c561e..141489497 100644 --- a/tests/test_nanoevents.py +++ b/tests/test_nanoevents.py @@ -272,6 +272,22 @@ def test_file_handle_from_path(tests_directory, mode): assert factory.file_handle is not None +@pytest.mark.parametrize("mode", ["eager", "virtual"]) +def test_factory_pickle_preserves_mode(tests_directory, mode): + import pickle + + path = f"{tests_directory}/samples/nano_dy.root:Events" + factory = NanoEventsFactory.from_root( + path, + schemaclass=NanoAODSchema, + mode=mode, + ) + + unpickled = pickle.loads(pickle.dumps(factory)) + assert unpickled._mode == mode + assert unpickled.events() is not None + + @pytest.mark.parametrize("mode", ["eager", "virtual"]) def test_file_handle_from_directory(tests_directory, mode): """Test that file_handle is available when passing ReadOnlyDirectory.""" diff --git a/tests/test_nanoevents_edm4hep.py b/tests/test_nanoevents_edm4hep.py index ff72250d7..66d84d9ea 100644 --- a/tests/test_nanoevents_edm4hep.py +++ b/tests/test_nanoevents_edm4hep.py @@ -327,3 +327,47 @@ def test_Relations(eager_events, delayed_events, field): elif d_fin.layout.branch_depth[1] == 3: mixin = d_fin.layout.content.content.content.parameter("__record__") assert target_name.startswith(mixin) + + +def test_edm4hep_lookup_one_to_many_relation(): + import copy + + from coffea.nanoevents.assets import edm4hep_ver + from coffea.nanoevents.schemas.edm4hep import parse_yaml + + schema = EDM4HEPSchema.__new__(EDM4HEPSchema) + schema.edm4hep = edm4hep_ver["00-99-01"]() + schema.parsed_edm4hep = parse_yaml(schema.edm4hep, copy.deepcopy(schema.edm4hep)) + schema._datatype_mixins = {"MCParticleCollection": "MCParticle"} + assert ( + schema._lookup_branch("MCParticleCollection", "daughters", key="type") + == "edm4hep::MCParticle" + ) + + +def test_edm4hep_yaml_cache_is_readonly(): + # The parsed edm4hep yaml is loaded once and shared across all schema + # builds; a build must therefore treat it as read-only. Guards against + # reintroducing per-build mutation of the shared cache. + import copy + + from coffea.nanoevents.schemas import edm4hep as edm4hep_module + + version = EDM4HEPSchema.edm4hep_version + raw, parsed = edm4hep_module.load_edm4hep(version) + raw_snapshot = copy.deepcopy(raw) + parsed_snapshot = copy.deepcopy(parsed) + + # A full schema build exercises every path that reads the cached dicts. + _events( + mode="eager", + iteritems_options={"filter_name": "/^(?!.*(PARAMETERS|_.*Map))/"}, + ) + + raw_after, parsed_after = edm4hep_module.load_edm4hep(version) + # Caching is active: the same objects are handed to every build ... + assert raw_after is raw + assert parsed_after is parsed + # ... and the build did not mutate them. + assert raw_after == raw_snapshot + assert parsed_after == parsed_snapshot diff --git a/tests/test_nanoevents_fcc_spring2021.py b/tests/test_nanoevents_fcc_spring2021.py index e947f4e0f..7b04da1f1 100644 --- a/tests/test_nanoevents_fcc_spring2021.py +++ b/tests/test_nanoevents_fcc_spring2021.py @@ -296,3 +296,27 @@ def test_KaonParent_to_PionDaughters_Loop(eager_events): nested_bool_parent = p.PDG == PDG_IDs["K(S)0"] daughters_have_K_S0_parent = awkward.all(awkward.ravel(nested_bool_parent)) assert daughters_have_K_S0_parent + + +def test_fcc_unknown_collections_record_array(): + from coffea.nanoevents.schemas.fcc import FCCSchema + + schema = FCCSchema.__new__(FCCSchema) + schema.mixins_dictionary = {} + branch_forms = { + "Foo/Foo.bar": { + "class": "RecordArray", + "fields": ["x"], + "contents": [ + {"class": "NumpyArray", "primitive": "float64", "form_key": "k"} + ], + "form_key": "fk", + } + } + output, _ = schema._unknown_collections({}, branch_forms, set()) + assert output["Foo"]["parameters"]["__record__"] == "NanoCollection" + + +def test_fcc_get_schema_bad_version(): + with pytest.raises(ValueError): + FCC.get_schema("bad-version") diff --git a/tests/test_nanoevents_pfnano.py b/tests/test_nanoevents_pfnano.py index f065208fb..e89e371ca 100644 --- a/tests/test_nanoevents_pfnano.py +++ b/tests/test_nanoevents_pfnano.py @@ -52,6 +52,28 @@ def check_fields_recursive(coll, field): check_fields_recursive(events, field) +@pytest.mark.parametrize("mode", ["eager", "dask"]) +def test_jet_associations(tests_directory, mode): + if mode == "dask": + pytest.importorskip("dask_awkward") + path = os.path.join(tests_directory, "samples/pfnano.root") + events = NanoEventsFactory.from_root( + {path: "Events"}, schemaclass=PFNanoAODSchema, mode=mode + ).events() + + svs, pfc = events.JetSVs, events.JetPFCands + checks = [ + (svs.jet.pt, events.Jet[svs.jetIdx].pt), + (svs.sv.pt, events.SV[svs.sVIdx].pt), + (pfc.jet.pt, events.Jet[pfc.jetIdx].pt), + (pfc.pf.pt, events.PFCands[pfc.pFCandsIdx].pt), + ] + for linked, expected in checks: + if mode == "dask": + linked, expected = linked.compute(), expected.compute() + assert ak.all(linked == expected) + + def test_uproot_write(tmp_path): path = os.path.abspath("tests/samples/pfnano.root") orig_events = NanoEventsFactory.from_root( diff --git a/tests/test_nanoevents_transforms.py b/tests/test_nanoevents_transforms.py new file mode 100644 index 000000000..637706cf9 --- /dev/null +++ b/tests/test_nanoevents_transforms.py @@ -0,0 +1,13 @@ +import awkward as ak + +from coffea.nanoevents import transforms + + +def test_get_index_ranges_zero_valued_indices(): + # begin=0, end=1 yields the single index [0]; the all-zero sum must not be + # mistaken for an empty range and replaced with [[]]. + ranges = transforms.get_index_ranges(ak.Array([[0]]), ak.Array([[1]])) + assert ranges.tolist() == [[[0]]] + + empty = transforms.get_index_ranges(ak.Array([[]]), ak.Array([[]])) + assert empty.tolist() == [[]] diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 9f30f00a5..0e51c3ad9 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -1225,3 +1225,86 @@ def test_genvistau_addition_propagates_charge(): common = (ak.num(gvt) >= 1) & (ak.num(mu) >= 1) gm = gvt[common][:, 0] + mu[common][:, 0] assert "charge" in gm.fields + + +@pytest.mark.parametrize( + "name,kin1,kin2", + [ + ( + "Candidate", + {"x": 1.0, "y": 2.0, "z": 3.0, "t": 10.0}, + {"x": 0.5, "y": 1.0, "z": 1.5, "t": 4.0}, + ), + ( + "PtEtaPhiMCandidate", + {"pt": 10.0, "eta": 0.5, "phi": 0.1, "mass": 1.0}, + {"pt": 5.0, "eta": -0.2, "phi": 1.0, "mass": 0.5}, + ), + ( + "PtEtaPhiECandidate", + {"pt": 10.0, "eta": 0.5, "phi": 0.1, "energy": 20.0}, + {"pt": 5.0, "eta": -0.2, "phi": 1.0, "energy": 8.0}, + ), + ("Muon", None, None), + ], +) +def test_candidate_subtraction_differences_charge(name, kin1, kin2): + """Candidate subtraction keeps and negates the ``charge`` field.""" + from coffea.nanoevents.methods import candidate + + if name == "Muon": + from coffea.nanoevents import NanoAODSchema, NanoEventsFactory + + NanoAODSchema.warn_missing_crossrefs = False + events = NanoEventsFactory.from_root( + {"tests/samples/nano_dy.root": "Events"}, + schemaclass=NanoAODSchema, + mode="eager", + ).events() + mu = events.Muon[ak.num(events.Muon) >= 2] + assert len(mu) > 0 + a, b = mu[:, 0], mu[:, 1] + else: + a = ak.zip( + {**{k: [v] for k, v in kin1.items()}, "charge": [1]}, + with_name=name, + behavior=candidate.behavior, + ) + b = ak.zip( + {**{k: [v] for k, v in kin2.items()}, "charge": [-1]}, + with_name=name, + behavior=candidate.behavior, + ) + diff = a - b + assert "charge" in diff.fields + assert ak.all(diff.charge == a.charge - b.charge) + for c in ("x", "y", "z", "t"): + assert_allclose( + ak.to_list(getattr(diff, c)), + ak.to_list(getattr(a, c) - getattr(b, c)), + atol=ATOL, + ) + + +@pytest.mark.parametrize( + "name,temporal", + [("PtEtaPhiMLorentzVector", "mass"), ("PtEtaPhiELorentzVector", "energy")], +) +def test_polar_lorentz_negative_scalar_matches_cartesian(name, temporal): + """The time component transforms consistently with the Cartesian components + under negative scaling.""" + a = ak.zip( + {"pt": [1.0, 2.0], "eta": [1.2, -0.8], "phi": [0.3, 2.5], temporal: [3.0, 4.0]}, + with_name=name, + behavior=vector.behavior, + ) + cart = ak.zip( + {"x": a.x, "y": a.y, "z": a.z, "t": a.t}, + with_name="LorentzVector", + behavior=vector.behavior, + ) + for scaled, ref in ((a * (-2), cart * (-2)), (-a, -cart), (a / (-2), cart / (-2))): + for c in ("x", "y", "z", "t"): + assert_allclose( + ak.to_list(getattr(scaled, c)), ak.to_list(getattr(ref, c)), atol=ATOL + ) diff --git a/tests/test_preloaded.py b/tests/test_preloaded.py index 5fe6f5ac0..d09968057 100644 --- a/tests/test_preloaded.py +++ b/tests/test_preloaded.py @@ -1,13 +1,34 @@ import os +import awkward as ak import pytest import uproot from coffea.nanoevents import NanoEventsFactory -from coffea.nanoevents.mapping import SimplePreloadedColumnSource +from coffea.nanoevents.mapping import BufferCache, SimplePreloadedColumnSource +from coffea.nanoevents.schemas import BaseSchema from coffea.processor.test_items import NanoEventsProcessor +def test_from_preloaded_honors_buffer_cache(tests_directory): + rootdir = uproot.open(f"{tests_directory}/samples/nano_dy.root") + tree = rootdir["Events"] + arrays = tree.arrays(["nMuon", "Muon_pt"], how=dict) + src = SimplePreloadedColumnSource( + arrays, rootdir.file.uuid, tree.num_entries, object_path="/Events" + ) + + cache = BufferCache(cache=None, codec=None) + factory = NanoEventsFactory.from_preloaded( + src, buffer_cache=cache, schemaclass=BaseSchema + ) + assert factory.buffer_cache is cache + + events = factory.events() + ak.materialize(events.Muon_pt) + assert len(cache) > 0 + + def test_preloaded_nanoevents(): pytest.xfail("preloaded nanoevents doesn't support dask yet") diff --git a/tests/test_taskvine_executor.py b/tests/test_taskvine_executor.py new file mode 100644 index 000000000..fca0a20ed --- /dev/null +++ b/tests/test_taskvine_executor.py @@ -0,0 +1,75 @@ +import signal +import types + +import pytest + +from coffea.processor import taskvine_executor as tv +from coffea.processor.executor import _compress + + +def test_handle_early_terminate_soft_then_hard(monkeypatch): + cancelled = [] + + class FakeManager: + class console: + @staticmethod + def printf(*args, **kwargs): + pass + + @staticmethod + def cancel_by_category(category): + cancelled.append(category) + + monkeypatch.setattr(tv, "manager", FakeManager) + monkeypatch.setattr(tv, "early_terminate", False) + + # first interrupt: soft-terminate the run, no exception, results preserved + tv._handle_early_terminate(signal.SIGINT, None) + assert tv.early_terminate is True + assert cancelled == ["processing", "accumulating"] + + # second interrupt: hard-terminate + with pytest.raises(KeyboardInterrupt): + tv._handle_early_terminate(signal.SIGINT, None) + + +def test_processing_binds_concurrent_reads(monkeypatch, tmp_path): + captured = {} + + monkeypatch.setattr(tv.signal, "signal", lambda *a, **k: None) + monkeypatch.setattr( + tv.CoffeaVine, "_make_process_bars", lambda self: None, raising=False + ) + monkeypatch.setattr( + tv.CoffeaVine, + "_process_events", + lambda self, proc_fn, accum_fn, items: captured.__setitem__( + "accum_fn", accum_fn + ), + raising=False, + ) + monkeypatch.setattr( + tv.CoffeaVine, "_final_accumulation", lambda self, acc: acc, raising=False + ) + monkeypatch.setattr( + tv.CoffeaVine, "_update_bars", lambda self, **k: None, raising=False + ) + + sizes = [] + real_pool = tv.ThreadPool + monkeypatch.setattr(tv, "ThreadPool", lambda n: sizes.append(n) or real_pool(n)) + + m = tv.CoffeaVine.__new__(tv.CoffeaVine) + m.executor = types.SimpleNamespace(compression=1, concurrent_reads=5) + m.stats_coffea = {} + + m._processing([1, 2, 3], lambda x: x, None) + + accum_fn = captured["accum_fn"] + # buggy code passes concurrent_reads as the wrapper's name + assert str(accum_fn) == "accumulate_result_files" + + f = tmp_path / "file.0" + f.write_bytes(_compress({"x": 1}, 1)) + accum_fn([str(f)]) + assert sizes[-1] == 5