materialize_w_manual_graph.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ddaf7e-ce36-4a83-9b5e-a3dd19d5191a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import datetime\n",
    "import traceback\n",
    "\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",
    "\n",
    "from coffea.nanoevents import NanoEventsFactory, NanoAODSchema\n",
    "from coffea.analysis_tools import PackedSelection\n",
    "from coffea import dataset_tools\n",
    "\n",
    "from functools import partial\n",
    "import cloudpickle\n",
    "from collections import defaultdict\n",
    "import os\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "from ndcctools.taskvine import DaskVine, PythonTask, FunctionCall\n",
    "    \n",
    "warnings.filterwarnings(\"ignore\")\n",
    "NanoAODSchema.warn_missing_crossrefs = False # silences warnings about branches we will not use here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0e5e52-b15a-4484-804a-7d88e9b77b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "BRANCH_LIST = [\n",
    "        \"run\", \"luminosityBlock\", \"event\",\n",
    "        \"GenPart_pt\", \"GenPart_eta\", \"GenPart_phi\", \"CorrT1METJet_phi\",\n",
    "        \"GenJet_pt\", \"CorrT1METJet_eta\", \"SoftActivityJet_pt\",\n",
    "        \"Jet_eta\", \"Jet_phi\", \"SoftActivityJet_eta\", \"SoftActivityJet_phi\",\n",
    "        \"CorrT1METJet_rawPt\", \"Jet_btagDeepFlavB\", \"GenJet_eta\", \n",
    "        \"GenPart_mass\", \"GenJet_phi\",\n",
    "        \"Jet_puIdDisc\", \"CorrT1METJet_muonSubtrFactor\", \"Jet_btagDeepFlavCvL\",\n",
    "        \"Jet_btagDeepFlavQG\", \"Jet_mass\", \"Jet_pt\", \"GenPart_pdgId\",\n",
    "        \"Jet_btagDeepFlavCvB\", \"Jet_cRegCorr\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04fac84-abf6-481e-bbb3-1514ca8a3bdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def report_stats(out, t0=None, t1=None, failed=False):\n",
    "    if not out:\n",
    "        return\n",
    "\n",
    "    if not isinstance(out, dict):\n",
    "        # something is weird with that output\n",
    "        print(out)\n",
    "        return\n",
    "\n",
    "    if t0 is None:\n",
    "        t0 = out[\"start\"]\n",
    "        \n",
    "    if t1 is None:\n",
    "        t1 = out[\"end\"]\n",
    "\n",
    "    if failed:\n",
    "        print(f\"{out['failed']}\")\n",
    "\n",
    "    print(f\"chunks: {out['chunks']}\")\n",
    "    print(f\"events: {out['num_events']}\")\n",
    "    # summary of performance\n",
    "    read_GB = out['read'] / 1000**3\n",
    "    print(f\"total read (compressed): {read_GB:.2f} GB\")\n",
    "    print(f\"total read (uncompressed): {out['uncompressed'] / 1000**3:.2f} GB\")\n",
    "\n",
    "    rate_Gbps = read_GB*8/(t1-t0)\n",
    "    if rate_Gbps == 0:\n",
    "        rate_Gbps = 0.000000001\n",
    "    print(f\"average data rate: {rate_Gbps:.2f} Gbps (need to scale by x{200/rate_Gbps:.0f} to reach 200 Gbps)\")\n",
    "\n",
    "    n_evts = 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 = 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\"event rate (aggregated time spent in function): {n_evts / total_runtime / 1000:.2f} kHz\")               \n",
    "    print(f\"failed: {len(out['failed'])}/{out['chunks']}\")\n",
    "    print(f\"not dict: {out['not_dict']}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "673aa975-9565-42a1-b883-9a857e95529b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# wrappers for debugging\n",
    "def get_bytes(key, fn, *args):\n",
    "    # get the transfer stats out of the dask function\n",
    "    r = fn(*args)\n",
    "    if key[0] == \"accum\":\n",
    "        return (None, r)\n",
    "    else:\n",
    "        return ((r['read'], r['uncompressed'], r['num_entries'], r['start'], r['end'], r['chunks'], len(r['failed'])), r)\n",
    "\n",
    "def avg_bytes():\n",
    "    output_file = open(\"stats.csv\", \"w\")\n",
    "    print(f\"read,uncompressed,num_entries,start,end,chunks,nfailed\", file=output_file)\n",
    "\n",
    "    stats = defaultdict(lambda: 0)\n",
    "    stats['start'] = None\n",
    "\n",
    "    def avg(info, stats=stats, output_file=output_file):\n",
    "        if not info:\n",
    "            return\n",
    "        read, uncompressed, num_entries, start, end, chunks, nfailed = info\n",
    "\n",
    "        print(f\"{read},{uncompressed},{num_entries},{start},{end},{chunks},{nfailed}\", file=output_file)\n",
    "\n",
    "        if not stats['start']:\n",
    "            stats['start'] = start\n",
    "\n",
    "        if stats['start'] > start:\n",
    "            stats['start'] = start\n",
    "\n",
    "        if stats['end'] < start:\n",
    "            stats['end'] = end\n",
    "\n",
    "        stats['total'] += read\n",
    "\n",
    "        if stats['counter'] % 1000 == 0:\n",
    "            bps = (8*stats['total'])/(stats['end'] - stats['start'])\n",
    "            if bps > stats['max_rate_seen']:\n",
    "                stats['max_rate_seen'] = bps\n",
    "            print(f\"rate: {bps/1000**3:.2f} Gbps, read: {stats['total']/1000**3:.2f} GB,  max seen: {stats['max_rate_seen']/1000**3:.2f} Gbps\")\n",
    "            output_file.flush()\n",
    "\n",
    "        stats['counter'] += 1\n",
    "    return avg\n",
    "\n",
    "\n",
    "def trace_memory(key, fn, *args):\n",
    "    import tracemalloc\n",
    "    tracemalloc.start()\n",
    "    before = tracemalloc.get_traced_memory()\n",
    "    b = fn(*args)\n",
    "    after = tracemalloc.get_traced_memory()\n",
    "    tracemalloc.stop()\n",
    "\n",
    "    memory_used = (after[1] - before[1])/1000**3\n",
    "    memory_used = after[1]/1000**3\n",
    "    a = (key, memory_used)\n",
    "\n",
    "    return (a, b)\n",
    "\n",
    "def trace_memory_peak(arg):\n",
    "    stats = {'max_memory': 0}\n",
    "    def mem(arg, stats=stats):\n",
    "        (key, memory) = arg\n",
    "        if memory > stats['max_memory']:\n",
    "            stats['max_memory'] = memory\n",
    "            print(\"new max\", key, memory)\n",
    "    return mem\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dbebe2a-491e-4ad8-a9b2-7de656ab3edd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# work for coffea\n",
    "def do_stuff_real(events):\n",
    "    # track number of events\n",
    "    num_entries = ak.num(events, axis=0)\n",
    "\n",
    "    # read out all other branches into integers to avoid memory issues\n",
    "    _counter = 0\n",
    "    for obj_to_add in [\n",
    "        events.GenPart.pt,\n",
    "        events.GenPart.eta,\n",
    "        events.GenPart.phi,\n",
    "        events.CorrT1METJet.phi,\n",
    "        events.GenJet.pt,\n",
    "        events.CorrT1METJet.eta,\n",
    "        events.SoftActivityJet.pt,\n",
    "        events.Jet.eta,\n",
    "        events.Jet.phi,\n",
    "        events.SoftActivityJet.eta,\n",
    "        events.SoftActivityJet.phi,\n",
    "        events.LHEPart.eta,\n",
    "        events.LHEPart.phi,\n",
    "        events.CorrT1METJet.rawPt,\n",
    "        events.Jet.btagDeepFlavB,\n",
    "        events.GenJet.eta,\n",
    "        events.GenPart.mass,\n",
    "        events.GenJet.phi,\n",
    "        events.Jet.puIdDisc,\n",
    "        events.CorrT1METJet.muonSubtrFactor,\n",
    "        events.Jet.btagDeepFlavCvL,\n",
    "        events.LHEPart.mass,\n",
    "        events.LHEPart.pt,\n",
    "        events.Jet.btagDeepFlavQG,\n",
    "        events.Jet.mass,\n",
    "        events.Jet.pt,\n",
    "        events.GenPart.pdgId,\n",
    "        events.Jet.btagDeepFlavCvB,\n",
    "        events.Jet.cRegCorr,\n",
    "        events.LHEPart.incomingpz\n",
    "    ]:\n",
    "        _counter_to_add = ak.count_nonzero(obj_to_add, axis=1)\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 {\"num_entries\": num_entries, \"_counter\": _counter}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb9b3d0-367c-4e1e-874e-39796e8e92c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# work for uproot.dask\n",
    "# specs: [(filename, obj_path, entry_start, entry_stop), ...]\n",
    "def do_uproot_read_dask(specs, compute=False):\n",
    "    import dask\n",
    "    import awkward as ak\n",
    "    import time\n",
    "    import gc\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    size_read = 0\n",
    "    size_uncompressed = 0\n",
    "    num_entries = 0\n",
    "    ccounter_all = 0\n",
    "    failed = []\n",
    "\n",
    "    try:\n",
    "        for (filename, path, entry_start, entry_stop) in specs:\n",
    "            try:\n",
    "                spec = {f\"{filename}\": {\"object_path\": \"Events\", \"steps\": [entry_start, entry_stop]}}\n",
    "                pre_events, pre_report = uproot.dask(spec, filter_name=BRANCH_LIST, allow_read_errors_with_report=True, entry_start=entry_start, entry_stop=entry_stop)\n",
    "                events, report = dask.compute(pre_events, pre_report, num_workers=1)\n",
    "\n",
    "                if report.exception and report.exception[0]:\n",
    "                    raise Exception(f\"{filename}: {report.exception} {report.message}\")\n",
    "\n",
    "                num_entries += ak.num(events, axis=0)\n",
    "\n",
    "                if compute:\n",
    "                    ccounter_all = 0\n",
    "                    for b in BRANCH_LIST:\n",
    "                        try:\n",
    "                            ccounter_all += ak.count_nonzero(events[b], axis=None)\n",
    "                        except Exception:\n",
    "                            pass\n",
    "\n",
    "                size_uncompressed = events.nbytes\n",
    "                size_read = sum(report['performance_counters']['num_requested_bytes'])\n",
    "            except Exception as e:\n",
    "                #failed.append({filename: traceback.format_exc()})\n",
    "                failed.append({filename: e})\n",
    "                print(f\"{filename}: {e}\")\n",
    "            finally:\n",
    "                events = None\n",
    "                report = None\n",
    "                # gc.collect()\n",
    "    except Exception as e:\n",
    "        size_read = 0\n",
    "        size_uncompressed = 0\n",
    "        num_entries = 0\n",
    "        ccounter_all = 0\n",
    "\n",
    "    t1 = time.time()\n",
    "\n",
    "    return {\"chunks\": len(specs), \"read\": size_read, \"uncompressed\":\n",
    "            size_uncompressed, \"num_entries\": num_entries, \"compute_counter\":\n",
    "            ccounter_all, \"runtime\": t1-t0, \"start\": t0, \"end\": t1, \"failed\": failed}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95a5a574-9e4e-4c83-92e7-f1f1685f6628",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the one used for the workshop results\n",
    "def do_uproot_read_open(specs, compute=False):\n",
    "    import uproot\n",
    "    import time\n",
    "    import gc\n",
    "    import signal\n",
    "\n",
    "    def handler(num, stack):\n",
    "        raise TimeoutError(f\"{specs}\")\n",
    "\n",
    "    signal.alarm(300)\n",
    "    signal.signal(signal.SIGALRM, handler)\n",
    "\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    size_read = 0\n",
    "    size_uncompressed = 0\n",
    "    num_entries = 0\n",
    "    ccounter_all = 0\n",
    "    failed = []\n",
    "    try:\n",
    "        for (filename, path, entry_start, entry_stop) in specs:\n",
    "            try:\n",
    "                with uproot.open(f\"{filename}\") as froot:\n",
    "                    if not entry_start:\n",
    "                        entry_start = 0\n",
    "                    if not entry_stop:\n",
    "                        entry_stop = froot[path].num_entries\n",
    "                    num_entries += entry_stop - entry_start\n",
    "                    for b in BRANCH_LIST:\n",
    "                        try:\n",
    "                            froot[\"Events\"][b].array(entry_start=entry_start, entry_stop=entry_stop)\n",
    "                            size_uncompressed += froot[\"Events\"][b].uncompressed_bytes\n",
    "                        except uproot.exceptions.KeyInFileError:\n",
    "                            pass\n",
    "\n",
    "                    size_read += froot.file.source.num_requested_bytes\n",
    "            except Exception as e:\n",
    "                #failed.append({filename: traceback.format_exc()})\n",
    "                failed.append({filename: e})\n",
    "                print(f\"{filename}: {e}\")\n",
    "                raise\n",
    "            finally:\n",
    "                # gc.collect()\n",
    "                pass\n",
    "    except Exception as e:\n",
    "        size_read = 0\n",
    "        size_uncompressed = 0\n",
    "        num_entries = 0\n",
    "        ccounter_all = 0\n",
    "        #raise\n",
    "\n",
    "    t1 = time.time()\n",
    "\n",
    "    return {\"chunks\": len(specs), \"read\": size_read, \"uncompressed\":\n",
    "            size_uncompressed, \"num_entries\": num_entries, \"compute_counter\":\n",
    "            ccounter_all, \"runtime\": t1-t0, \"start\": t0, \"end\": t1, \"failed\": failed}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908c9c2d-8d5e-4d40-83a5-43268c984668",
   "metadata": {},
   "outputs": [],
   "source": [
    "# accumulation task for manual graph\n",
    "# adds numbers and extends lists values from partial dictionaries\n",
    "def accum_dict(*partials):\n",
    "    from collections import defaultdict\n",
    "    out = defaultdict(lambda: 0)\n",
    "    out[\"failed\"] = []\n",
    "\n",
    "    start = None\n",
    "    end = None\n",
    "    for p in partials:\n",
    "        if not p:\n",
    "            continue\n",
    "\n",
    "        if not isinstance(p, dict):\n",
    "            out[\"not_dict\"] += 1\n",
    "            continue\n",
    "        try:\n",
    "            pstart = p[\"start\"]\n",
    "            if not start or pstart < start:\n",
    "                start = pstart\n",
    "        except KeyError:\n",
    "            pass\n",
    "\n",
    "        try:\n",
    "            pend = p[\"end\"]\n",
    "            if not end or pend > end:\n",
    "                end = pend\n",
    "        except KeyError:\n",
    "            pass\n",
    "\n",
    "        if \"failed\" in p and p[\"failed\"]:\n",
    "            print(p[\"failed\"])\n",
    "\n",
    "        for k, v in p.items():\n",
    "            if isinstance(v, list):\n",
    "                out[k].extend(v)\n",
    "            else:\n",
    "                try:\n",
    "                    out[k] += v\n",
    "                except KeyError:\n",
    "                    out[k] = v\n",
    "                except TypeError:\n",
    "                    print(k, v)\n",
    "    out[\"start\"] = start\n",
    "    out[\"end\"] = end\n",
    "\n",
    "    return out\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50308091-6b6b-45dd-9d83-3b937886b2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct manual graph\n",
    "# samples: from coffea preprocess\n",
    "# fn: function to apply to samples (e.g., do_uproot_read_dask)\n",
    "# times: how many times to repeat each sample\n",
    "# files_per_task: number of files to assign to a single read task\n",
    "# compute: Whether to make a simple computation on the data read (i.e., count nonzero entries)\n",
    "# accumulate: Whether to accumulate per dataset (True), or return each read result indiviudually (False)\n",
    "# accumulate_all: Whether to accumulate the individual datasets into a single result.\n",
    "# accum_fn: Function to accumulate results (default is to accumulate as dictionaries)\n",
    "# accum_chunk: Number of resutls per accumualtion task\n",
    "def manual_graph(samples, fn, times=1, files_per_task=1, compute=False, accumulate=True, accumulate_all=False, accum_fn=accum_dict, accum_chunk=10):\n",
    "    \n",
    "    def add_vertex(graph, process, on_queue, files_for_task):\n",
    "        # if len(graph) > 1000:\n",
    "        #    return\n",
    "        key = (f\"uproot:{process}\", len(graph))\n",
    "        graph[key] = (fn, files_for_task, compute)\n",
    "        on_queue.append(key)\n",
    "\n",
    "    targets = []\n",
    "    graph = {}\n",
    "    for _ in range(times):\n",
    "        files_for_task = []\n",
    "\n",
    "        for process in samples:\n",
    "            on_queue = []\n",
    "\n",
    "            for name, info in samples[process][\"files\"].items():\n",
    "                obj_path = info.get('object_path', 'Events')\n",
    "                steps = info.get('steps', [[None, None]])\n",
    "\n",
    "                for (start, stop) in steps:\n",
    "                    files_for_task.append([name, obj_path, start, stop])\n",
    "                    if len(files_for_task) >= files_per_task:\n",
    "                        add_vertex(graph, process, on_queue, files_for_task)\n",
    "                        files_for_task = []\n",
    "\n",
    "            if len(files_for_task) > 0:\n",
    "                add_vertex(graph, process, on_queue, files_for_task)\n",
    "                files_for_task = []\n",
    "\n",
    "            if accumulate:\n",
    "                while on_queue:\n",
    "                    args = on_queue[0:accum_chunk]\n",
    "                    on_queue = on_queue[accum_chunk:]\n",
    "                    key = (\"accum\", process, len(graph))\n",
    "                    graph[key] = (accum_fn, *args)\n",
    "                    if on_queue:\n",
    "                        on_queue.append(key)\n",
    "                    else:\n",
    "                        targets.append(key)\n",
    "            else:\n",
    "                targets.extend(on_queue)\n",
    "    if accumulate and accumulate_all:\n",
    "        on_queue = targets\n",
    "        targets = []\n",
    "        while on_queue:\n",
    "            args = on_queue[0:accum_chunk]\n",
    "            on_queue = on_queue[accum_chunk:]\n",
    "            key = (\"accum\", len(graph))\n",
    "            graph[key] = (accum_fn, *args)\n",
    "            if on_queue:\n",
    "                on_queue.append(key)\n",
    "            else:\n",
    "                targets.append(key)\n",
    "    return (graph, targets)\n",
    " \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e275ea61-191e-449e-920a-da009f170c89",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"uproot: {uproot.__version__}\")\n",
    "print(f\"hist: {hist.__version__}\")\n",
    "print(f\"coffea: {coffea.__version__}\")\n",
    "\n",
    "manager = DaskVine(port=8786, ssl=True, name=f\"{os.environ.get('USER', 'noname')}-coffea-casa\",run_info_path=\"/mnt/data/btovar-logs/\",)\n",
    "\n",
    "extra_files = {}\n",
    "env_vars = {}\n",
    "    \n",
    "token_acc_path = \"/etc/cmsaf-secrets-chown/access_token\"\n",
    "token_xch_path = \"/etc/cmsaf-secrets-chown/xcache_token\"\n",
    "\n",
    "if Path(token_acc_path).is_file():\n",
    "    extra_files[manager.declare_file(token_acc_path, cache=True)] = \"access_token\"\n",
    "    env_vars[\"BEARER_TOKEN_FILE\"] = \"access_token\"\n",
    "if Path(token_xch_path).is_file():\n",
    "    extra_files[manager.declare_file(token_xch_path, cache=True)] = \"xcache_token\"\n",
    "    env_vars[\"XCACHE_FILE\"] = \"xcache_token\"\n",
    "\n",
    "\n",
    "# bring back accumulation task results for better disk garbage collection\n",
    "def checkpoint_accum(dag, key):\n",
    "    if key[0] == \"accum\":\n",
    "        return True\n",
    "\n",
    "# usually we put all these options directly into dask calls,\n",
    "# but coffea preprocessing only allows one argument to set the scheduler,\n",
    "# thus we create a partial of manager.get, which is the function that takes\n",
    "# a dask graph and executes it.\n",
    "vine_scheduler = partial(manager.get,\n",
    "                         resources={\"cores\": 1},  #  max 1 core, 5GB of disk per task\n",
    "                         #resources_mode='fixed',\n",
    "                         resources_mode=None,   # set to \"fixed\" to kill tasks on resources\n",
    "                         extra_files=extra_files,\n",
    "                         checkpoint_fn=checkpoint_accum,\n",
    "                         env_vars=env_vars,\n",
    "                         submit_per_cycle=1000,  # throttle submission to keep memory usage low,\n",
    "                         max_pending=6000,       # and start doing work faster\n",
    "                         worker_transfers=True,  # keep partials at workers\n",
    "                         task_mode=\"function-calls\", # use one interpreter per worker\n",
    "                         lib_resources={\"cores\": 8, \"slots\": 8}, # resources and functions a single interpreter can run\n",
    "                         # environment=\"env.tar.gz\",   # needed for task_mode=\"tasks\" if taskvine version at worker \n",
    "                                                       # is behind: poncho_package_create $CONDA_PREFIX env.tar.gz,\n",
    "                                                       # or if more modules are needed at the execution site.\n",
    "                        )\n",
    "\n",
    "# given to coffea and dask functions as **scheduler_options to make taskvine the scheduler\n",
    "scheduler_options = {}\n",
    "scheduler_options['scheduler'] = vine_scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2217faf-b4db-47bb-8950-c1fdf224f838",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read datasets\n",
    "import json\n",
    "fname = \"zstd_files.json\"\n",
    "\n",
    "fileset = {}\n",
    "with open(fname,'r') as fp:\n",
    "    for i,(dataset_name,file_list) in enumerate(json.load(fp).items()):\n",
    "        fileset[dataset_name] = {\"files\": {}}\n",
    "        for j,dataset_fpath in enumerate(file_list):\n",
    "            xrd_fpath = f\"root://xcache.cmsaf-dev.flatiron.hollandhpc.org:1094/{dataset_fpath}\"\n",
    "            fileset[dataset_name][\"files\"][xrd_fpath] = \"Events\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25ecaa10-f3c0-4f82-9d39-837e49d95c5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess\n",
    "try:\n",
    "    # step_size = 50_000\n",
    "    # step_size = 100_000\n",
    "    # step_size = 500_000\n",
    "    step_size = 5_000_000\n",
    "\n",
    "    with open(f\"preprocessed_{step_size}.pkl\", \"rb\") as f:\n",
    "        # do not re preprocess if we don't have too...\n",
    "        samples = cloudpickle.load(f)\n",
    "except Exception:\n",
    "    scheduler_options['scheduler'] = vine_scheduler\n",
    "    samples, report = dataset_tools.preprocess(fileset, step_size=step_size, skip_bad_files=True, uproot_options={\"allow_read_errors_with_report\": True}, **scheduler_options)\n",
    "    with open(f\"preprocessed_{step_size}.pkl\", \"wb\") as f:\n",
    "        cloudpickle.dump(samples, f)\n",
    "\n",
    "    total_files  = sum([len(p[\"files\"]) for p in samples.values()])\n",
    "    total_chunks = sum(sum(len(f[\"steps\"]) for f in p[\"files\"].values()) for p in samples.values())\n",
    "    print(f\"nfiles: {total_files} chunks: {total_chunks}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b475b0-80aa-4aea-8137-13b21ad294c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run with manual graph\n",
    "if True:\n",
    "    # create the graph\n",
    "    tasks, targets = manual_graph(samples, times=2, files_per_task=1, fn=do_uproot_read_open, compute=False, accumulate=True, accumulate_all=True, accum_fn=accum_dict, accum_chunk=10)\n",
    "\n",
    "    print(f\"tasks: {len(tasks)}\")\n",
    "    \n",
    "    # send twice the functions that can fit to a worker. With this, a worker can start working in the next load\n",
    "    # while waiting for the manager to retrieve completed results.\n",
    "    manager.tune(\"resource-submit-multiplier\", 2)\n",
    "    \n",
    "    # keep a second copy of each results in the workers. More relevant when there is eviction, but it helps a little\n",
    "    # when accumulating results, as workers can more easily find available sources for transfers.\n",
    "    manager.tune(\"temp-replica-count\", 2)\n",
    "\n",
    "    # execute task graph\n",
    "    t0 = time.time()\n",
    "    outs = vine_scheduler(tasks, targets, progress_label='[green]process',\n",
    "                          wrapper=get_bytes, wrapper_proc=avg_bytes())\n",
    "    #outs = vine_scheduler(tasks, targets, progress_label='[green]process', wrapper=trace_memory, wrapper_proc=trace_memory_peak())\n",
    "\n",
    "    ### or use dask client\n",
    "    # from dask.distributed import Client, performance_report\n",
    "    # client = Client(\"tls://localhost:8786\")\n",
    "    # outs = client.get(tasks, targets)\n",
    "\n",
    "    t1 = time.time()\n",
    "\n",
    "    print(\"stats with IO time\")\n",
    "    report_stats(outs[0], failed=True)\n",
    "\n",
    "    print(\"stats IO + accumulation\")\n",
    "    report_stats(outs[0], t0, t1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845ecefd-880b-479a-b720-683502589424",
   "metadata": {
    "lines_to_next_cell": 3
   },
   "outputs": [],
   "source": [
    "# regular coffea\n",
    "if False:\n",
    "    t0 = time.perf_counter()\n",
    "\n",
    "    # change default scheduler\n",
    "    tasks = dataset_tools.apply_to_fileset(do_stuff_real, samples, uproot_options={\"allow_read_errors_with_report\": (OSError, TypeError, KeyError), \"filter_name\": BRANCH_LIST})\n",
    "\n",
    "    #(out, report) = dask.compute(tasks, **scheduler_options, progress_label=\"[green]process\")\n",
    "    ((out,report),) = dask.compute(tasks, **scheduler_options)\n",
    "    t1 = time.perf_counter()\n",
    "\n",
    "    try:\n",
    "        print(f\"total time spent in uproot reading data: {ak.sum([v['duration'] for v in report.values()]):.2f} s\")\n",
    "        print(f\"wall time: {t1-t0:.2f}s\")\n",
    "        print(f\"events: {sum(out[k]['num_entries'] for k in out)}\")\n",
    "        event_rate = sum(out[k][\"num_entries\"] for k in out)\n",
    "\n",
    "        event_rate = event_rate / (t1-t0)\n",
    "        print(f\"event rate: {event_rate / 1_000:.2f} kHz\")\n",
    "\n",
    "        read_GB = ak.sum([v['performance_counters']['num_requested_bytes'] for v in report.values()]) / 1_000**3\n",
    "        rate_Gbs = read_GB / (t1-t0)\n",
    "        print(f\" - read {read_GB:.2f} GB in {t1-t0:.2f} s -> {rate_Gbs:.2f} GBps (need to scale by x{200/8/rate_Gbs*1000:.0f} to reach 200 Gbps)\")\n",
    "    except Exception:\n",
    "        pass\n",
    "\n",
    "    with open(\"outs.pkl\", \"wb\") as f:\n",
    "        cloudpickle.dump((out, report), f)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "notebook_metadata_filter": "all,-jupytext.text_representation.jupytext_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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}