materialize_branches.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180f4d32-e5c4-42f9-84e8-430034493a55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import glob\n",
    "import json\n",
    "import os\n",
    "from collections import defaultdict\n",
    "from pathlib import Path\n",
    "import traceback\n",
    "import time\n",
    "import warnings\n",
    "import copy\n",
    "import pathlib\n",
    "\n",
    "import awkward as ak\n",
    "import dask\n",
    "import dask_awkward as dak\n",
    "import hist.dask\n",
    "import coffea\n",
    "import numpy as np\n",
    "import uproot\n",
    "from dask.distributed import Client\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "mpl.style.use(\"ggplot\")\n",
    "\n",
    "from coffea.nanoevents import NanoEventsFactory, PHYSLITESchema\n",
    "from coffea import dataset_tools\n",
    "\n",
    "import utils\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import input_files.utils\n",
    "\n",
    "from dask.distributed import LocalCluster, Client, progress, performance_report\n",
    "\n",
    "# local: single thread, single worker\n",
    "# cluster = LocalCluster(n_workers=1, processes=False, threads_per_worker=1)\n",
    "# client = Client(cluster)\n",
    "\n",
    "# for UChicago\n",
    "# update this to point to your own client!\n",
    "client = Client(\"tcp://dask-alheld-fa8ea830-2.af-jupyter:8786\")\n",
    "\n",
    "# create a folder for output tracking of uproot.open setup\n",
    "MEASUREMENT_PATH = pathlib.Path(datetime.datetime.now().strftime(\"measurements/%Y-%m-%d_%H-%M-%S\"))\n",
    "os.makedirs(MEASUREMENT_PATH)\n",
    "\n",
    "print(f\"awkward: {ak.__version__}\")\n",
    "print(f\"dask-awkward: {dak.__version__}\")\n",
    "print(f\"uproot: {uproot.__version__}\")\n",
    "print(f\"hist: {hist.__version__}\")\n",
    "print(f\"coffea: {coffea.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef7dd8d-68e5-43fa-b8bf-5d8f972520c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -------------------\n",
    "# INPUT CONFIGURATION\n",
    "# -------------------\n",
    "# modify this to change how many files are being processed\n",
    "# top-level processes determine containers/DSIDs, which each have some number of files\n",
    "# full list is list(find_containers.container_dict.keys()) + [\"data15_13TeV\", \"data16_13TeV\", \"data17_13TeV\", \"data18_13TeV\"]\n",
    "\n",
    "PROCESSES_TO_USE = [\"ttbar\"]  # 6.7 TB\n",
    "# PROCESSES_TO_USE = [\"db\", \"zjets\", \"wjets\", \"ttV\", \"othertop\", \"ttbar\"]  # all simulation, 47 TB\n",
    "# PROCESSES_TO_USE = [\"db\", \"zjets\", \"wjets\", \"ttV\", \"othertop\", \"ttbar\", \"data15_13TeV\", \"data16_13TeV\", \"data17_13TeV\", \"data18_13TeV\"]  # 190 TB\n",
    "\n",
    "fileset = input_files.utils.get_fileset(PROCESSES_TO_USE, max_files_per_container=1, max_containers_per_dsid=None, max_dsid_per_process=None)\n",
    "\n",
    "# example for how to veto files\n",
    "# files_to_veto = [(\"root://192.170.240.148//root://fax.mwt2.org:1094//pnfs/uchicago.edu/atlaslocalgroupdisk/rucio/mc20_13TeV/f5/99/DAOD_PHYSLITE.37230013._001196.pool.root.1\", \"CollectionTree\")]\n",
    "# fileset = dataset_tools.filter_files(fileset, lambda x: x in files_to_veto)\n",
    "\n",
    "# duplicate each entry in the fileset\n",
    "# print(\"DUPLICATING FILESET CONTENT\")\n",
    "# fileset_with_duplicates = copy.deepcopy(fileset)\n",
    "# for k, v in fileset.items():\n",
    "#     fileset_with_duplicates.update({f\"{k}-duplicate\": dict(v)})\n",
    "# fileset = fileset_with_duplicates\n",
    "\n",
    "utils.save_fileset(fileset, MEASUREMENT_PATH)\n",
    "print(f\"total number of files (including duplicates): {sum([len(v['files']) for v in fileset.values()])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff312282-a87b-4394-aff4-6acc12d27c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check for files not yet replicated to MWT2\n",
    "files_at_mwt2 = 0\n",
    "files_elsewhere = 0\n",
    "for process in fileset.keys():\n",
    "    for file in fileset[process][\"files\"]:\n",
    "        if \"mwt2\" in file:\n",
    "            files_at_mwt2 += 1\n",
    "        else:\n",
    "            files_elsewhere += 1\n",
    "\n",
    "print(f\"files at MWT2: {files_at_mwt2}, elsewhere: {files_elsewhere}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c438a9f-372a-4968-9758-466b52ad30e3",
   "metadata": {},
   "source": [
    "## Dask distributing `uproot.open`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd0ec454-7620-4c63-930a-2c28ffa6efda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn fileset into simple list of files to run over\n",
    "all_files = []\n",
    "for process in fileset:\n",
    "    all_files += fileset[process][\"files\"]\n",
    "\n",
    "# define work to be done\n",
    "def uproot_open_materialize(fname):\n",
    "    # ~9%, around 400 Mbps single core, ~150 Mbps with 100 workers\n",
    "    # BRANCH_LIST = [\n",
    "    #     'PrimaryVerticesAuxDyn.z',\n",
    "    #     'PrimaryVerticesAuxDyn.x',\n",
    "    #     'PrimaryVerticesAuxDyn.y',\n",
    "    #     'AnalysisJetsAuxDyn.Timing',\n",
    "    #     'AnalysisJetsAuxDyn.JetConstitScaleMomentum_phi',\n",
    "    #     'AnalysisJetsAuxDyn.DetectorEta',\n",
    "    #     'AnalysisJetsAuxDyn.ActiveArea4vec_eta',\n",
    "    #     'AnalysisJetsAuxDyn.JetConstitScaleMomentum_eta',\n",
    "    #     'AnalysisJetsAuxDyn.phi',\n",
    "    #     'AnalysisJetsAuxDyn.m',\n",
    "    #     'AnalysisJetsAuxDyn.JetConstitScaleMomentum_pt',\n",
    "    #     'AnalysisJetsAuxDyn.ActiveArea4vec_phi',\n",
    "    #     'AnalysisJetsAuxDyn.JetConstitScaleMomentum_m',\n",
    "    #     'AnalysisJetsAuxDyn.ActiveArea4vec_m',\n",
    "    #     'AnalysisJetsAuxDyn.pt',\n",
    "    #     'AnalysisJetsAuxDyn.Width',\n",
    "    #     'AnalysisJetsAuxDyn.EMFrac',\n",
    "    #     'AnalysisJetsAuxDyn.ActiveArea4vec_pt',\n",
    "    #     'AnalysisJetsAuxDyn.PSFrac'\n",
    "    # ]\n",
    "\n",
    "    # ~15%, around 300 Mbps single core, ~130 Mbps with 100 workers\n",
    "    BRANCH_LIST = [\n",
    "        'PrimaryVerticesAuxDyn.z',\n",
    "        'PrimaryVerticesAuxDyn.x',\n",
    "        'PrimaryVerticesAuxDyn.y',\n",
    "        'AnalysisJetsAuxDyn.Timing',\n",
    "        'AnalysisJetsAuxDyn.JetConstitScaleMomentum_phi',\n",
    "        'AnalysisJetsAuxDyn.DetectorEta',\n",
    "        'AnalysisJetsAuxDyn.ActiveArea4vec_eta',\n",
    "        'AnalysisJetsAuxDyn.JetConstitScaleMomentum_eta',\n",
    "        'AnalysisJetsAuxDyn.phi',\n",
    "        'AnalysisJetsAuxDyn.m',\n",
    "        'AnalysisJetsAuxDyn.JetConstitScaleMomentum_pt',\n",
    "        'AnalysisJetsAuxDyn.ActiveArea4vec_phi',\n",
    "        'AnalysisJetsAuxDyn.JetConstitScaleMomentum_m',\n",
    "        'AnalysisJetsAuxDyn.ActiveArea4vec_m',\n",
    "        'AnalysisJetsAuxDyn.pt',\n",
    "        'AnalysisJetsAuxDyn.Width',\n",
    "        'AnalysisJetsAuxDyn.EMFrac',\n",
    "        'AnalysisJetsAuxDyn.ActiveArea4vec_pt',\n",
    "        'AnalysisJetsAuxDyn.PSFrac',\n",
    "        'AnalysisJetsAuxDyn.JVFCorr',\n",
    "        'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_TracksC1',\n",
    "        'AnalysisJetsAuxDyn.eta',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone40_CloseByCorr',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone40',\n",
    "        'AnalysisPhotonsAuxDyn.eta',\n",
    "        'AnalysisJetsAuxDyn.DFCommonJets_fJvt',\n",
    "        'AnalysisPhotonsAuxDyn.phi',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone20_CloseByCorr',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone40ptCorrection',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone20ptCorrection',\n",
    "        'AnalysisPhotonsAuxDyn.pt',\n",
    "        'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_NTracks',\n",
    "        'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_TracksWidth',\n",
    "        'AnalysisJetsAuxDyn.GhostMuonSegmentCount',\n",
    "        'AnalysisPhotonsAuxDyn.topoetcone20',\n",
    "        'AnalysisPhotonsAuxDyn.f1',\n",
    "        'AnalysisPhotonsAuxDyn.DFCommonPhotonsIsEMTightIsEMValue',\n",
    "        'AnalysisPhotonsAuxDyn.ptcone20_CloseByCorr',\n",
    "        'AnalysisPhotonsAuxDyn.OQ',\n",
    "        'AnalysisPhotonsAuxDyn.ptcone20',\n",
    "        'AnalysisTauJetsAuxDyn.RNNJetScore',\n",
    "        'AnalysisTauJetsAuxDyn.JetDeepSetScore',\n",
    "        'AnalysisTauJetsAuxDyn.etaTauEnergyScale',\n",
    "        'AnalysisTauJetsAuxDyn.etaFinalCalib',\n",
    "        'AnalysisTauJetsAuxDyn.RNNEleScoreSigTrans_v1'\n",
    "    ]\n",
    "\n",
    "    filter_name = lambda x: x in BRANCH_LIST\n",
    "\n",
    "    size_uncompressed = 0\n",
    "    t0 = time.perf_counter()\n",
    "    try:\n",
    "        with uproot.open(fname, filter_name=filter_name) as f:\n",
    "            num_entries = f[\"CollectionTree\"].num_entries\n",
    "\n",
    "            # iterate approach\n",
    "            # for _ in f[\"CollectionTree\"].iterate(expressions=BRANCH_LIST):\n",
    "            #     pass\n",
    "\n",
    "            # branch loop approach\n",
    "            for b in BRANCH_LIST:\n",
    "                f[\"CollectionTree\"][b].array()\n",
    "                size_uncompressed += f[\"CollectionTree\"][b].uncompressed_bytes\n",
    "\n",
    "            size_read = f.file.source.num_requested_bytes\n",
    "        exception = None\n",
    "\n",
    "    except:\n",
    "        num_entries = 0\n",
    "        size_read = 0\n",
    "        size_uncompressed = 0\n",
    "        exception = traceback.format_exc()\n",
    "\n",
    "    t1 = time.perf_counter()\n",
    "    time_finished = datetime.datetime.now()\n",
    "    return {\"fname\": fname, \"read\": size_read, \"uncompressed\": size_uncompressed, \"num_entries\": num_entries,\n",
    "            \"runtime\": t1-t0, \"time_finished\": time_finished, \"exception\": exception}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bc64663-8ca0-46f8-8ca1-54c1aac51979",
   "metadata": {},
   "outputs": [],
   "source": [
    "# perform computation\n",
    "print(f\"running with {len(all_files)} files\")\n",
    "# scattered_data = client.scatter([f for f in all_files])  # instead of submitting (possibly big) object directly\n",
    "\n",
    "utils.start_tracking_workers(client, MEASUREMENT_PATH)  # track worker count in background\n",
    "with performance_report(filename=MEASUREMENT_PATH/\"dask-report-plain-uproot.html\"):\n",
    "    # futures = client.map(uproot_open_materialize, scattered_data)\n",
    "    # out = ak.Array([r for r in client.gather(iter(futures))])\n",
    "\n",
    "    tasks = [dask.delayed(uproot_open_materialize)(f) for f in all_files]  # create tasks\n",
    "    t0 = time.perf_counter()\n",
    "    out = ak.Array(dask.compute(*tasks))  # perform computations\n",
    "    t1 = time.perf_counter()\n",
    "\n",
    "utils.stop_tracking_workers()\n",
    "\n",
    "print(f\"wall clock time: {t1-t0:.2f}s\")\n",
    "utils.save_measurement(out, t0, t1, MEASUREMENT_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5ea00ca-43f9-4e47-bc2c-2ffbad32fe95",
   "metadata": {},
   "source": [
    "while waiting, check out out the XCache output: https://grafana.mwt2.org/d/EKefjM-Sz/af-network-200gbps-challenge?orgId=1&viewPanel=205&from=now-30m&to=now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e786f1a-2df7-4e35-95f9-2438512a17f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# when repeating plots for old measurements, set appropriate MEASUREMENT_PATH here\n",
    "# MEASUREMENT_PATH = pathlib.Path(\"measurements/2024-05-09_20-39-29\")\n",
    "\n",
    "# load measurements from file again\n",
    "timestamps, nworkers, avg_num_workers = utils.get_timestamps_and_counts(MEASUREMENT_PATH)  # worker count info\n",
    "out, t0, t1 = utils.load_measurement(MEASUREMENT_PATH)\n",
    "\n",
    "# summary of performance\n",
    "read_GB = sum(out['read']) / 1000**3\n",
    "print(f\"total read (compressed): {read_GB:.2f} GB\")\n",
    "print(f\"total read (uncompressed): {sum(out['uncompressed']) / 1000**3:.2f} GB\")\n",
    "\n",
    "rate_Gbps = read_GB*8/(t1-t0)\n",
    "print(f\"average data rate: {rate_Gbps:.2f} Gbps (need to scale by x{200/rate_Gbps:.1f} to reach 200 Gbps)\")\n",
    "\n",
    "n_evts = sum(out[\"num_entries\"])\n",
    "print(f\"total event rate (wall clock time): {n_evts / (t1-t0) / 1000:.2f} kHz (processed {n_evts} events total)\")\n",
    "\n",
    "total_runtime = sum(out[\"runtime\"])\n",
    "print(f\"total aggregated runtime in function: {total_runtime:.2f} s\")\n",
    "print(f\"ratio total runtime / wall clock time: {total_runtime / (t1-t0):.2f} \"\\\n",
    "      \"(should match # cores without overhead / scheduling issues)\")\n",
    "print(f\"time-averaged number of workers: {avg_num_workers:.1f}\")\n",
    "print(f\"\\\"efficiency\\\" (ratio of two numbers above): {total_runtime / (t1-t0) / avg_num_workers:.1%}\")\n",
    "print(f\"event rate (aggregated time spent in function): {n_evts / total_runtime / 1000:.2f} kHz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "698ae2c8-4c83-43b2-9f46-3c8163add495",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get arrays for starting time, runtime and end time of all tasks\n",
    "runtimes = np.asarray([datetime.timedelta(seconds=t) for t in out[\"runtime\"]], dtype=np.timedelta64)\n",
    "ends = out[\"time_finished\"].to_numpy()\n",
    "starts = ends - runtimes\n",
    "\n",
    "# calculate instantaneous rates for given timestamp\n",
    "times_for_rates = []\n",
    "instantaneous_rates = []\n",
    "for t in timestamps[::30]:  # only calculate every 30 seconds\n",
    "    mask = np.logical_and((starts <= t), (t <= ends))  # mask for tasks running at given timestamp\n",
    "    rate_Gbps_at_timestamp = sum(out[mask]['read']*8 / 1000**3 / out[mask][\"runtime\"])\n",
    "    times_for_rates.append(t)\n",
    "    instantaneous_rates.append(rate_Gbps_at_timestamp)\n",
    "\n",
    "utils.plot_worker_count(timestamps, nworkers, avg_num_workers, times_for_rates, instantaneous_rates, MEASUREMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "136da759-b7ea-49c2-afab-fe4e12c5422e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{sum(o is not None for o in out['exception'])} files failed\\n\")\n",
    "\n",
    "# use below to get full list with details\n",
    "# for report in out:\n",
    "#     if report[\"exception\"] is not None:\n",
    "#         print(f\"{report['fname']} failed in {report['runtime']:.2f} s\\n{report['exception']}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d51de7a-722e-4bf8-8a3e-f39861b0261f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# runtime distribution for all files\n",
    "fig, ax = plt.subplots()\n",
    "bins = np.linspace(0, max(out[\"runtime\"])*1.01, 100)\n",
    "ax.hist(out[\"runtime\"], bins=bins)\n",
    "ax.set_xlabel(\"runtime [s]\")\n",
    "ax.set_xlim([0, ax.get_xlim()[1]])\n",
    "ax.set_ylabel(\"count\")\n",
    "ax.semilogy()\n",
    "fig.savefig(MEASUREMENT_PATH / \"runtime_distribution.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af23ee49-a5fe-4fd3-8c01-8ff76944d95f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# runtime vs number of events in file\n",
    "fig, ax = plt.subplots()\n",
    "ax.scatter(out[\"num_entries\"], out[\"runtime\"], marker=\"x\")\n",
    "ax.set_xlabel(\"number of events\")\n",
    "ax.set_ylabel(\"runtime [s]\")\n",
    "fig.savefig(MEASUREMENT_PATH / \"runtime_vs_nevts.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36bcc374-834f-4aa9-8975-0118a98ddd35",
   "metadata": {},
   "source": [
    "## Using coffea 2024"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a38efe4-8024-422c-ba42-12bd3a3b44cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BRANCH_LIST = [\n",
    "    'PrimaryVerticesAuxDyn.z',\n",
    "    'PrimaryVerticesAuxDyn.x',\n",
    "    'PrimaryVerticesAuxDyn.y',\n",
    "    'AnalysisJetsAuxDyn.Timing',\n",
    "    'AnalysisJetsAuxDyn.JetConstitScaleMomentum_phi',\n",
    "    'AnalysisJetsAuxDyn.DetectorEta',\n",
    "    'AnalysisJetsAuxDyn.ActiveArea4vec_eta',\n",
    "    'AnalysisJetsAuxDyn.JetConstitScaleMomentum_eta',\n",
    "    'AnalysisJetsAuxDyn.phi',\n",
    "    'AnalysisJetsAuxDyn.m',\n",
    "    'AnalysisJetsAuxDyn.JetConstitScaleMomentum_pt',\n",
    "    'AnalysisJetsAuxDyn.ActiveArea4vec_phi',\n",
    "    'AnalysisJetsAuxDyn.JetConstitScaleMomentum_m',\n",
    "    'AnalysisJetsAuxDyn.ActiveArea4vec_m',\n",
    "    'AnalysisJetsAuxDyn.pt',\n",
    "    'AnalysisJetsAuxDyn.Width',\n",
    "    'AnalysisJetsAuxDyn.EMFrac',\n",
    "    'AnalysisJetsAuxDyn.ActiveArea4vec_pt',\n",
    "    'AnalysisJetsAuxDyn.PSFrac',\n",
    "    'AnalysisJetsAuxDyn.JVFCorr',\n",
    "    'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_TracksC1',\n",
    "    'AnalysisJetsAuxDyn.eta',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone40_CloseByCorr',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone40',\n",
    "    'AnalysisPhotonsAuxDyn.eta',\n",
    "    'AnalysisJetsAuxDyn.DFCommonJets_fJvt',\n",
    "    'AnalysisPhotonsAuxDyn.phi',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone20_CloseByCorr',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone40ptCorrection',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone20ptCorrection',\n",
    "    'AnalysisPhotonsAuxDyn.pt',\n",
    "    'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_NTracks',\n",
    "    'AnalysisJetsAuxDyn.DFCommonJets_QGTagger_TracksWidth',\n",
    "    'AnalysisJetsAuxDyn.GhostMuonSegmentCount',\n",
    "    'AnalysisPhotonsAuxDyn.topoetcone20',\n",
    "    'AnalysisPhotonsAuxDyn.f1',\n",
    "    'AnalysisPhotonsAuxDyn.DFCommonPhotonsIsEMTightIsEMValue',\n",
    "    'AnalysisPhotonsAuxDyn.ptcone20_CloseByCorr',\n",
    "    'AnalysisPhotonsAuxDyn.OQ',\n",
    "    'AnalysisPhotonsAuxDyn.ptcone20',\n",
    "    'AnalysisTauJetsAuxDyn.RNNJetScore',\n",
    "    'AnalysisTauJetsAuxDyn.JetDeepSetScore',\n",
    "    'AnalysisTauJetsAuxDyn.etaTauEnergyScale',\n",
    "    'AnalysisTauJetsAuxDyn.etaFinalCalib',\n",
    "    'AnalysisTauJetsAuxDyn.RNNEleScoreSigTrans_v1'\n",
    "]\n",
    "\n",
    "\n",
    "def materialize_branches(events):\n",
    "    num_events = ak.num(events, axis=0)  # track number of events\n",
    "\n",
    "    # this will read around 25% of data files\n",
    "    # materialize branches, just derive integers from them that will be aggregated to avoid memory issues\n",
    "    _counter = 0\n",
    "    for branch in BRANCH_LIST:\n",
    "        obj_name, obj_prop = branch.split(\".\")\n",
    "        obj_name = obj_name.replace(\"Analysis\", \"\").replace(\"AuxDyn\", \"\")\n",
    "        if \"Link\" not in obj_prop:\n",
    "            branch_data = events[obj_name, obj_prop]\n",
    "        else:\n",
    "            branch_data = events[obj_name, obj_prop][\"m_persIndex\"]\n",
    "\n",
    "        _counter_to_add = ak.count_nonzero(branch_data, axis=-1)  # reduce innermost\n",
    "\n",
    "        # reduce >2-dimensional (per event) branches further\n",
    "        for _ in range(_counter_to_add.ndim - 1):\n",
    "            _counter_to_add = ak.count_nonzero(_counter_to_add, axis=-1)\n",
    "\n",
    "        _counter = _counter + _counter_to_add  # sum 1-dim array built from new branch\n",
    "\n",
    "    _counter = ak.count_nonzero(_counter, axis=0)  # reduce to int\n",
    "\n",
    "    return {\"nevts\": num_events, \"_counter\": _counter}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed1ecc41-0a72-44ec-a8b8-654286a02196",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# pre-process\n",
    "with performance_report(filename=MEASUREMENT_PATH/\"dask-report-preprocess.html\"):\n",
    "    samples, report = dataset_tools.preprocess(fileset, skip_bad_files=True, uproot_options={\"allow_read_errors_with_report\": True})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b56b6c-dd28-4979-a26a-1cd7cdba24cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# find issues where access did not work\n",
    "for process in report:\n",
    "    for k, v in report[process][\"files\"].items():\n",
    "        if v[\"steps\"] is None:\n",
    "            print(f\"could not read {k}\")\n",
    "\n",
    "samples = dict((k, dict(v.items())) for k, v in samples.items())  # convert to dict (no defaultdict left)\n",
    "\n",
    "# save pre-processing output to json\n",
    "with open(\"preprocessed_files.json\", \"w\") as f:\n",
    "    f.write(json.dumps(samples, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "463c7d3d-c906-4bd9-99f7-232a1efe6514",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the pre-processing results: useful to skip repeated pre-processing\n",
    "with open(\"preprocessed_files.json\") as f:\n",
    "    samples = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15835a57-1182-4efb-8306-07f36af7a5b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# create the task graph\n",
    "# filter_name seems to not do anything here in terms of performance\n",
    "filter_name = lambda name: name in BRANCH_LIST\n",
    "tasks = dataset_tools.apply_to_fileset(materialize_branches,\n",
    "                                       samples,\n",
    "                                       uproot_options={\"allow_read_errors_with_report\": (OSError, TypeError, KeyError), \"filter_name\": filter_name},\n",
    "                                       schemaclass=PHYSLITESchema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b6dfb17-6806-46c8-aa59-cd475e40cf64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# execute task graph\n",
    "utils.start_tracking_workers(client, MEASUREMENT_PATH)  # track worker count in background\n",
    "t0 = time.perf_counter()\n",
    "with performance_report(filename=MEASUREMENT_PATH/\"dask-report-compute.html\"):\n",
    "    ((out, report),) = dask.compute(tasks)  # feels strange that this is a tuple-of-tuple\n",
    "t1 = time.perf_counter()\n",
    "utils.stop_tracking_workers()\n",
    "\n",
    "time_uproot = ak.sum([v['duration'] for v in report.values()])\n",
    "print(f\"total time spent in uproot reading data: {time_uproot:.2f} s\")\n",
    "print(f\"wall time: {t1-t0:.2f}s\")\n",
    "\n",
    "timestamps, nworkers, avg_num_workers = utils.get_timestamps_and_counts(MEASUREMENT_PATH)  # worker count info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb47e9e5-57ea-4385-b8f8-b2a22ded030e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"output: {out}\")\n",
    "\n",
    "print(\"\\nperformance metrics:\")\n",
    "event_rate = sum([out[process][\"nevts\"] for process in out.keys()]) / (t1-t0)\n",
    "print(f\" - event rate: {event_rate / 1_000:.2f} kHz\")\n",
    "\n",
    "num_bytes = ak.sum([report[process][\"performance_counters\"][\"num_requested_bytes\"] for process in out.keys()])\n",
    "read_MB = num_bytes / 1_000**2\n",
    "rate_Mbs = read_MB / (t1-t0)\n",
    "print(f\" - read {read_MB:.2f} MB in {t1-t0:.2f} s -> {rate_Mbs/1000*8:.2f} Gbps (need to scale by x{200/8/rate_Mbs*1000:.1f} to reach 200 Gbps)\")\n",
    "print(f\" - time-averaged number of workers: {avg_num_workers:.1f}\")\n",
    "print(f\" - spent {time_uproot:.1f} s reading data with wall time {t1-t0:.2f} and {avg_num_workers:.1f} cores on average -> \\\"efficiency\\\": {time_uproot / (t1-t0) / avg_num_workers:.1%}\")\n",
    "\n",
    "utils.plot_worker_count(timestamps, nworkers, avg_num_workers, [], [], pathlib.Path(\".\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8da90743-8045-42f3-a59c-12c7ff9f8b7b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# report problematic files that caused exceptions\n",
    "for process in report.keys():\n",
    "    for i_file in range(len(report[process].exception)):\n",
    "        file_report = report[process][i_file]\n",
    "        if file_report.exception is not None:\n",
    "            print(file_report.args[0].strip(\"\\'\"))\n",
    "            print(file_report.message + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51fb1ab6-0be4-4b99-942a-c1b53eda5639",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sanity check that the right colums are being touched\n",
    "# dak.report_necessary_columns(tasks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25225440-d57c-49e7-b48e-43fe5362327d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if issues with files exist, paste in path and reproduce\n",
    "# fname = \"root://192.170.240.148//root://fax.mwt2.org:1094//pnfs/uchicago.edu/atlaslocalgroupdisk/rucio/mc20_13TeV/f5/99/DAOD_PHYSLITE.37230013._001196.pool.root.1\"\n",
    "# treename = \"CollectionTree\"\n",
    "# events = NanoEventsFactory.from_root({fname: treename}, schemaclass=PHYSLITESchema).events()\n",
    "# task = materialize_branches(events)\n",
    "# task[\"_counter\"].compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5f767bc-495a-4eba-a809-cd4768d65497",
   "metadata": {},
   "source": [
    "## Dask distributing `xrdcp`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a2d0a00-e696-41df-b4ec-c9e183ee818f",
   "metadata": {},
   "outputs": [],
   "source": [
    "RUN_XRDCP = False  # off by default\n",
    "\n",
    "# distribute `xrdcp` with Dask\n",
    "def run_xrdcp(fname):\n",
    "    t0 = time.perf_counter()\n",
    "    os.system(f\"xrdcp {fname} /dev/null -f\")\n",
    "    t1 = time.perf_counter()\n",
    "    time_finished = datetime.datetime.now()\n",
    "    return {\"runtime\": t1-t0, \"time_finished\": time_finished}\n",
    "\n",
    "all_files = []\n",
    "for process in fileset:\n",
    "    all_files += fileset[process][\"files\"]\n",
    "\n",
    "if RUN_XRDCP:\n",
    "    # perform computation\n",
    "    print(f\"running with {len(all_files)} files\")\n",
    "    \n",
    "    t0 = time.perf_counter()\n",
    "    with performance_report(filename=\"dask-report-xrdcp.html\"):\n",
    "        futures = client.map(run_xrdcp, all_files)\n",
    "        out = ak.Array([r for r in client.gather(iter(futures))])\n",
    "    t1 = time.perf_counter()\n",
    "    \n",
    "    print(f\"wall clock time: {t1-t0:.2f}s\")\n",
    "    print(f\"current number of workers: {len(client.scheduler_info()['workers'])}\")\n",
    "    \n",
    "    total_runtime = sum(out[\"runtime\"])\n",
    "    print(f\"total aggregated runtime in function: {total_runtime:.2f} s\")\n",
    "    print(f\"ratio total runtime / wall clock time: {total_runtime / (t1-t0):.2f} \"\\\n",
    "          \"(should match # cores without overhead / scheduling issues)\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "notebook_metadata_filter": "all,-jupytext.text_representation.jupytext_version,-language_info.version"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}