[WIP] Coffea2023

.github/workflows/main.yml

@@ -87,7 +87,7 @@ jobs:
         run: |
           mkdir dir_for_topcoffea
           cd dir_for_topcoffea
-          git clone https://github.com/TopEFT/topcoffea.git
+          git clone -b coffea2023 https://github.com/TopEFT/topcoffea.git
           cd topcoffea
           conda run -n coffea-env pip install -e .
           cd ../..
@@ -105,7 +105,7 @@ jobs:
 
       - name: Download root files
         run: |
-          wget --no-verbose http://uaf-10.t2.ucsd.edu/~kmohrman/for_ci/for_wwz/WWZJetsTo4L2Nu_4F_TuneCP5_13TeV-amcatnlo-pythia8_RunIISummer20UL17NanoAODv9-106X_mc2017_realistic_v9-v2_NANOAODSIM_3LepTau_4Lep/output_1.root
+          wget --no-verbose http://uaf-8.t2.ucsd.edu/~kmohrman/for_ci/for_wwz/WWZJetsTo4L2Nu_4F_TuneCP5_13TeV-amcatnlo-pythia8_RunIISummer20UL17NanoAODv9-106X_mc2017_realistic_v9-v2_NANOAODSIM_3LepTau_4Lep/output_1.root
 
       - name: Pytest setup
         run: |

analysis/wwz/get_wwz_counts.py

@@ -47,18 +47,32 @@ def get_counts(histos_dict):
     return out_dict
 
 
+# Restructure hists so samples are in category axis
+def restructure_hist(in_dict):
+
+    dense_axes = list(list(in_dict.values())[0].keys())
+    restructured_hist = {}
+    for dense_axis in dense_axes:
+        for dataset_name in in_dict.keys():
+            if dense_axis not in restructured_hist:
+                restructured_hist[dense_axis] = in_dict[dataset_name][dense_axis]
+            else:
+                restructured_hist[dense_axis] += in_dict[dataset_name][dense_axis]
+    return restructured_hist
+
+
 def main():
 
     # Set up the command line parser
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--pkl-file-path", default="histos/plotsTopEFT.pkl.gz", help = "The path to the pkl file")
+    parser.add_argument("pkl_file_path", help = "The path to the pkl file")
     parser.add_argument("-o", "--output-path", default=".", help = "The path the output files should be saved to")
     parser.add_argument("-n", "--output-name", default="counts_wwz_sync", help = "A name for the output directory")
     args = parser.parse_args()
 
     # Get the counts from the input hiso
     histo_dict = pickle.load(gzip.open(args.pkl_file_path))
-
+    histo_dict = restructure_hist(histo_dict)
 
     counts_dict = get_counts(histo_dict)
 

analysis/wwz/run_wrapper.sh

@@ -0,0 +1,8 @@
+# Get the file the CI uses, and move it to the directory the JSON expects
+printf "\nDownloading root file...\n"
+wget -nc http://uaf-10.t2.ucsd.edu/~kmohrman/for_ci/for_wwz/WWZJetsTo4L2Nu_4F_TuneCP5_13TeV-amcatnlo-pythia8_RunIISummer20UL17NanoAODv9-106X_mc2017_realistic_v9-v2_NANOAODSIM_3LepTau_4Lep/output_1.root
+
+# Run the processor via the run_wwz4l.py script
+# Note the -x executor argument does not do anything anymore, though should keep it in for now since without it the code tries to default to wq so tries to start packaging up env
+
+time python run_wwz4l.py ../../input_samples/sample_jsons/test_samples/UL17_WWZJetsTo4L2Nu_forCI.json,../../input_samples/sample_jsons/test_samples/UL17_WWZJetsTo4L2Nu_forCI_extra.json -x iterative

analysis/wwz/run_wwz4l.py

@@ -6,11 +6,20 @@
 import cloudpickle
 import gzip
 import os
-from coffea import processor
+import dask
+import dask_awkward as dak
+from distributed import Client
+
 from coffea.nanoevents import NanoAODSchema
+from coffea.nanoevents import NanoEventsFactory
+
+from coffea.dataset_tools import preprocess
+from coffea.dataset_tools import apply_to_fileset
 
 import topcoffea.modules.remote_environment as remote_environment
 
+from ndcctools.taskvine import DaskVine
+
 import wwz4l
 
 LST_OF_KNOWN_EXECUTORS = ["futures","work_queue","iterative"]
@@ -165,11 +174,11 @@ def LoadJsonToSampleName(jsonFile, prefix):
                         else:
                             LoadJsonToSampleName(l, prefix)
 
-    flist = {}
+    fdict = {}
     nevts_total = 0
     for sname in samplesdict.keys():
         redirector = samplesdict[sname]['redirector']
-        flist[sname] = [(redirector+f) for f in samplesdict[sname]['files']]
+        fdict[sname] = [(redirector+f) for f in samplesdict[sname]['files']]
         samplesdict[sname]['year'] = samplesdict[sname]['year']
         samplesdict[sname]['xsec'] = float(samplesdict[sname]['xsec'])
         samplesdict[sname]['nEvents'] = int(samplesdict[sname]['nEvents'])
@@ -306,34 +315,96 @@ def LoadJsonToSampleName(jsonFile, prefix):
     # Run the processor and get the output
     tstart = time.time()
 
-    if executor == "futures":
-        exec_instance = processor.FuturesExecutor(workers=nworkers)
-        runner = processor.Runner(exec_instance, schema=NanoAODSchema, chunksize=chunksize, maxchunks=nchunks)
-    elif executor == "iterative":
-        exec_instance = processor.IterativeExecutor()
-        runner = processor.Runner(exec_instance, schema=NanoAODSchema, chunksize=chunksize, maxchunks=nchunks)
-    elif executor ==  "work_queue":
-        executor = processor.WorkQueueExecutor(**executor_args)
-        runner = processor.Runner(executor, schema=NanoAODSchema, chunksize=chunksize, maxchunks=nchunks, skipbadfiles=False, xrootdtimeout=180)
+    ####################################3
+    ### coffea2023 ###
+
+    # Get fileset
+    fileset = {}
+    for name, fpaths in fdict.items():
+        fileset[name] = {}
+        fileset[name]["files"] = {}
+        for fpath in fpaths:
+            fileset[name]["files"][fpath] = {"object_path": "Events"}
+            fileset[name]["metadata"] = {"dataset": name}
+    print(fileset)
+    print("Number of datasets:",len(fdict))
+
+
+    #### Try with distributed Client ####
+
+    #with dask.config.set({"scheduler": "sync"}): # Single thread
+    #with Client() as _: # distributed Client scheduler
+    with Client() as client:
+
+        # Run preprocess
+        print("\nRunning preprocess...")
+        dataset_runnable, dataset_updated = preprocess(
+            fileset,
+            maybe_step_size=50_000,
+            align_clusters=False,
+            files_per_batch=1,
+            #skip_bad_files=True,
+            #calculate_form=True,
+        )
+
+        # Run apply_to_fileset
+        print("\nRunning apply_to_fileset...")
+        histos_to_compute, reports = apply_to_fileset(
+            processor_instance,
+            dataset_runnable,
+            uproot_options={"allow_read_errors_with_report": True}
+        )
+
+        # Check columns to be read
+        print("\nRunning necessary_columns...")
+        columns_read = dak.necessary_columns(histos_to_compute[list(histos_to_compute.keys())[0]])
+        print(columns_read)
+
+        # Compute
+        print("\nRunning compute...")
+        output_futures, report_futures = {}, {}
+        for key in histos_to_compute:
+            output_futures[key], report_futures[key] = client.compute((histos_to_compute[key], reports[key],)) # , scheduler=taskvine
+
+        coutputs, creports = client.gather((output_futures, report_futures,))
+
+
+
+    ### Task vine testing ###
+    do_tv = 0
+    if do_tv:
+
+        fdict = {"UL17_WWZJetsTo4L2Nu_forCI": ["/home/k.mohrman/coffea_dir/migrate_to_coffea2023_repo/ewkcoffea/analysis/wwz/output_1.root"]}
+
+        # Create dict of events objects
+        print("Number of datasets:",len(fdict))
+        events_dict = {}
+        for name, fpaths in fdict.items():
+            events_dict[name] = NanoEventsFactory.from_root(
+                {fpath: "/Events" for fpath in fpaths},
+                schemaclass=NanoAODSchema,
+                metadata={"dataset": name},
+            ).events()
+
+        # Get and compute the histograms
+        histos_to_compute = {}
+        for json_name in fdict.keys():
+            print(f"Getting histos for {json_name}")
+            histos_to_compute[json_name] = processor_instance.process(events_dict[json_name])
+
+        m = DaskVine([9123,9128], name=f"coffea-vine-{os.environ['USER']}")
+
+        print("Compute histos")
+        #output = dask.compute(histos_to_compute)[0] # Output of dask.compute is a tuple
+        coutputs = dask.compute(histos_to_compute, scheduler=m.get, resources={"cores": 1}, resources_mode=None, lazy_transfers=True)
 
-    output = runner(flist, treename, processor_instance)
 
     dt = time.time() - tstart
 
-    if executor == "work_queue":
-        print('Processed {} events in {} seconds ({:.2f} evts/sec).'.format(nevts_total,dt,nevts_total/dt))
-
-    #nbins = sum(sum(arr.size for arr in h._sumw.values()) for h in output.values() if isinstance(h, hist.Hist))
-    #nfilled = sum(sum(np.sum(arr > 0) for arr in h._sumw.values()) for h in output.values() if isinstance(h, hist.Hist))
-    #print("Filled %.0f bins, nonzero bins: %1.1f %%" % (nbins, 100*nfilled/nbins,))
-
-    if executor == "futures":
-        print("Processing time: %1.2f s with %i workers (%.2f s cpu overall)" % (dt, nworkers, dt*nworkers, ))
-
     # Save the output
     if not os.path.isdir(outpath): os.system("mkdir -p %s"%outpath)
     out_pkl_file = os.path.join(outpath,outname+".pkl.gz")
     print(f"\nSaving output in {out_pkl_file}...")
     with gzip.open(out_pkl_file, "wb") as fout:
-        cloudpickle.dump(output, fout)
+        cloudpickle.dump(coutputs, fout)
     print("Done!")