From 4b9147f9dd2747fa817e6f0f31f65a674b930fa6 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Fri, 12 Jan 2024 17:47:05 -0600 Subject: [PATCH] cache the delayed heavy object so we always use the same key in the task graph --- src/coffea/lookup_tools/lookup_base.py | 5 +++-- src/coffea/ml_tools/helper.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/coffea/lookup_tools/lookup_base.py b/src/coffea/lookup_tools/lookup_base.py index a34021247..5f681ac92 100644 --- a/src/coffea/lookup_tools/lookup_base.py +++ b/src/coffea/lookup_tools/lookup_base.py @@ -107,11 +107,12 @@ def __call__(self, *args, **kwargs): out_meta = tomap(self, *tuple([arg._meta for arg in actual_args])) - delayed_corr = dask.delayed(self) + if not hasattr(self, "_delayed_corr"): + setattr(self, "_delayed_corr", dask.delayed(self)) return dask_awkward.map_partitions( tomap, - delayed_corr, + self._delayed_corr, *actual_args, label=dask_label, meta=out_meta, diff --git a/src/coffea/ml_tools/helper.py b/src/coffea/ml_tools/helper.py index 4d00de7b5..2ecedbe6e 100644 --- a/src/coffea/ml_tools/helper.py +++ b/src/coffea/ml_tools/helper.py @@ -352,10 +352,11 @@ def __call__(self, wrapper, *args): packed_metas = repack(flattened_metas) wrap_meta = wrap(self, *packed_metas) - delayed_wrapper = dask.delayed(self) + if not hasattr(self, "_delayed_wrapper"): + setattr(self, "_delayed_wrapper", dask.delayed(self)) arr = dask_awkward.lib.core.map_partitions( wrap, - delayed_wrapper, + self._delayed_wrapper, *packed_args, label=f"numpy_call_{self.__class__.__name__}_", meta=wrap_meta,