def dask_setup(worker):
import os
def get_classads():
fname = os.getenv("_CONDOR_JOB_AD")
if not fname:
return {}
d = {}
with open(fname) as fh:
for line in fh:
if "=" not in line:
continue
k, v = line.split("=", 1)
d[k.strip()] = v.strip().lstrip('"').strip('"')
return d
worker.array_cache = None
try:
import uproot
worker.array_cache = uproot.ArrayCache("8 GB")
except ImportError as e:
print(e, "so we can't make a global ArrayCache")
except AttributeError as e:
print(e,
" Maybe this is an older version of uproot without ArrayCache?")
worker.classads = get_classads()
def read_root(tree, branches, start, stop):
global uproot_cache
read_root_lock.acquire()
if uproot_cache is None:
uproot_cache = uproot.ArrayCache("2 GB")
data = tree.arrays(branches=branches, outputtype=pandas.DataFrame, namedecode='utf-8',
entrystart=start, entrystop=stop, cache=uproot_cache)
read_root_lock.release()
return data
def dask_setup(worker):
import os
def get_classads():
fname = os.getenv("_CONDOR_JOB_AD")
if not fname:
return {}
d = {}
with open(fname) as fh:
for line in fh:
if "=" not in line:
continue
k, v = line.split("=", 1)
d[k.strip()] = v.strip().lstrip('"').strip('"')
return d
worker.classads = get_classads()
worker.array_cache = None
worker.tree_cache = {}
worker.nevents = 0
try:
import uproot
worker.array_cache = uproot.ArrayCache("6 GB")
except ImportError as e:
print(e, "so we can't make a global ArrayCache")
except AttributeError as e:
print(e,
" Maybe this is an older version of uproot without ArrayCache?")
def cachesize_metric(worker):
if hasattr(worker, "array_cache"):
return "{:.2f}GB".format(worker.array_cache._cache.currsize / 1e9)
return 0
def numtreescached_metric(worker):
if hasattr(worker, "tree_cache"):
return len(list(worker.tree_cache.keys()))
return 0
def nevents_metric(worker):
if hasattr(worker, "nevents"):
return "{:.2f}M".format(worker.nevents / 1e6)
return 0
worker.metrics["cachesize"] = cachesize_metric
worker.metrics["numtreescached"] = numtreescached_metric
worker.metrics["eventsprocessed"] = nevents_metric
try:
# Load some imports initially
import coffea.processor
import coffea.executor
except:
pass
def Lazy(file_list,branches):
cache = uproot.ArrayCache("2 GB")
tree = uproot.lazyarrays(
file_list,
"Delphes",
branches,
cache=cache,
)
print(len(tree))
Electron = ak.zip(
{
"PT": tree["Electron.PT"],
"Eta": tree["Electron.Eta"],
"Phi": tree["Electron.Phi"],
"Charge": tree["Electron.Charge"],
})
Muon = ak.zip(
{
"PT": tree["MuonLoose.PT"],
"Eta": tree["MuonLoose.Eta"],
"Phi": tree["MuonLoose.Phi"],
"Charge": tree["MuonLoose.Charge"],
})
Photon = ak.zip(
{
"PT": tree["PhotonLoose.PT"],
"Eta": tree["PhotonLoose.Eta"],
"Phi": tree["PhotonLoose.Phi"],
})
MET = ak.zip(
{
"PT": tree["PuppiMissingET.MET"],
"Phi": tree["PuppiMissingET.Phi"],
})
def __init__(self, name):
self.name_ = name
self.cache = up.ArrayCache("1 GB")
self.worker_meta["total_read_bytes"] += read_bytes
self.worker_meta["total_write_bytes"] += write_bytes
self.worker_meta["total_tasks"] += 1
self.worker_meta["total_time_elapsed"] += t1 - t0
self.worker_meta["busy"] = False
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"url", help="redis url. e.g., redis://[:password]@localhost:6379")
args = parser.parse_args()
w = Worker(args.url)
w.cache = None
try:
import uproot
ARRAY_CACHE = uproot.ArrayCache("8 GB")
w.cache = ARRAY_CACHE
except ImportError as e:
print(e, "so we can't make a global ArrayCache")
except AttributeError as e:
print(e,
" Maybe this is an older version of uproot without ArrayCache?")
w.start_pubsub()
w.run()