Memory Issue using Dask_histogram/Coffea/Dask_awkward

[JuanDuarte2003]

I've been experiencing the next problem when trying to build some histograms using Dask_histogram along a Dask cluster.

I have a YAML file with paths to root trees in CMS EOS storage

fileset:
  "ee_ttbarsignalviatau_2016preVFP":
    "files":
      "/eos/purdue/store/user/jduarteq/2016preVFP/spinCorrInput_2016preVFP_January2023/Nominal/ee/ee_ttbarsignalviatau_fromDilepton_2016ULpreVFP_11.root": "ttBar_treeVariables_step8"
.
.
.

In total, they're 660 root files divided into 24 datasets (not necessarily with the same amount of files). Each root tree has fields with weights used to construct the histograms, they have one "nominal" weight and a bunch of systematic weights. I load this YAML file in Python and then preprocess the full dataset:

with open(path_to_fileset, 'r') as f:
    cfg = yaml.full_load(f)
    fileset = cfg["fileset"]
    
dataset_runnable, dataset_updated = preprocess(
    fileset,
    align_clusters=True,
    step_size=100_000,
    files_per_batch=1,
    skip_bad_files=True,
    save_form=False
)

Then, I have a list of variables and bin edges to create the histogram that I also load from a YAML file like this

variables:
  top_pt:
    - [0.0, 30.0, 40.0, 50.0, 100.0, 600.0]
  top_scatteringangle_ttbarframe:
    - [-1.0,-0.6,-0.5,-0.4,0.4,0.5,0.6,1.0]
  ttbar_mass:
    - [0.0, 360.0, 380.0, 400.0, 700.0, 800.0, 900.0, 1100.0, 1200.0]
  beta:
    - [0.0, 0.60, 0.70, 0.80, 0.90, 0.95, 1.0]
  ll_cHel:
    - [-1.    , -0.9375, -0.875 , -0.8125, -0.75  , -0.6875, -0.625 ,
       -0.5625, -0.5   , -0.4375, -0.375 , -0.3125, -0.25  , -0.1875,
       -0.125 , -0.0625,  0.    ,  0.0625,  0.125 ,  0.1875,  0.25  ,
        0.3125,  0.375 ,  0.4375,  0.5   ,  0.5625,  0.625 ,  0.6875,
        0.75  ,  0.8125,  0.875 ,  0.9375,  1.    ]

The variables can be stored in the root files or be computed from other variables. Then, I pass this to a Coffea Processor Processor = HistogramProcessor(process, var_list, bin_edges_list). The processors looks like this:

class HistogramProcessor(processor.ProcessorABC):

    def __init__(self, process, var_list, bin_edges_list, systematic="Nominal",weight_tag="eventWeight"):
        # Parameters
        self.process_ = process
        self.var_list = var_list
        self.axis_labels = var_list
        self.systematic = systematic
        self.bin_edges_list = bin_edges_list
        self.weight_tag = weight_tag

    def process(self, events):
        # Get weight-based systematics
        systematics = [self.systematic]
        for _syst in constants.weight_based_systematics:
            if f'{self.weight_tag}_{_syst}' in events.fields:
                systematics.append(_syst)

        # Lumi-Weights 
        dataset_info = events.metadata['dataset']
        era = dataset_info.split('_')[-1]
        proc_name = '_'.join(dataset_info.split('_')[1:-1])
        cross_section = constants.get_cross_section(dataset_info)
        integrated_luminosity = constants.INTEGRATED_LUMI[era]
        totalWeighted = constants.get_weighted_events(era, self.systematic, proc_name)
        
        integrated_luminosity_scalefactor = (integrated_luminosity*cross_section / totalWeighted)
        weight_factor = integrated_luminosity_scalefactor*dak.ones_like(events['eventWeight']) 

        # Store variables to fill histograms
        fill_variables=[]
        for var in self.var_list:
            if var not in events.fields:
                var_val = calc_external_var(var,events)
            else:    
                var_val = events[var]
            fill_variables.append(
                var_val
                )

        # Create histogram
        axes = []
        ## Systematics axes
        axes.extend([hist.axis.StrCategory(
                systematics, name='weight_dict'
            )])

        # ## Var axes
        axes.extend([
            hist.axis.Variable(
                self.bin_edges_list[j], name=self.var_list[j],
                underflow=True, overflow=True
            )
            for j in range(len(self.bin_edges_list))
        ])
        histogram = hda.Hist(*axes, storage=hist.storage.Weight())

        # Fill histogram
        for _syst in systematics:
           syst_label = f'eventWeight_{_syst}' if _syst != self.systematic else 'eventWeight'
           match _syst:
               case _:
                   histogram.fill(_syst, *fill_variables, weight=events[syst_label]*weight_factor)

        # Correcting the weights to sqrt(N) from, sqrt(sum(weigths^2))
        histogram[...] = (np.concatenate((histogram.values()[..., None], 
                                        histogram.variances()[..., None]),
                                        axis=-1))

        return {"histogram":histogram}

    def postprocess(self, accumulator):
        pass

Here constants is just an imported script to load some numbers. As you can see, aside from the 5 axes associated with the histograms, I also added a 6th axis for the various systematic weights mentioned above. Then, this processor should return hist_dask when applied over the datasets like this:

to_compute = apply_to_fileset(
    Processor,
    dataset,
    schemaclass=BaseSchema # type: ignore
)

Now, to convert the hist_dask into futures and save them in disk after computation, I used the following function

saveHist = lambda hh, fn: save(hh, fn)

def compute_and_save(delay_hist, out_name) -> None:
    hist_fut = client.compute(delay_hist)
    save_fut = client.submit(saveHist,hist_fut,out_name)
    fire_and_forget(save_fut)
    return

Since the dataset has various keys, I do a for loop to apply the function.

keys = list(dataset.keys())
    for i in range(len(keys)):
        compute_and_save(to_compute[keys[i]]["histogram"], f'{outputDir}/subHistograms/Histogram_{keys[i]}_{process}.coffea')

Initially, it works fine, looking at the dashboard it shows that it's computing. At some point, the cluster gets killed. When sending the histograms to the cluster, dask shows some warnings about the graph size, which could indicate a memory problem. Therefore, I talked with someone from the IT team who works directly with the cluster, and he didn't find any problems on that side.

Maybe is there some way to build the hist.dask directly in the cluster so they don't have to go through the scheduler or something like that? That would be a solution to my main problem, still, it looks weird how dask dies without further messages.