multi-user-exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "787c34c1-52ef-485c-bd63-062305fc1b75",
   "metadata": {},
   "source": [
    "# Multi-user exercise at 2024 IRIS-HEP retreat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e578b0c-a7cf-429d-94e3-b0641970b2ca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import datetime\n",
    "import traceback\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, performance_report\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "mpl.style.use(\"ggplot\")\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 os\n",
    "import time\n",
    "import warnings\n",
    "import pathlib\n",
    "\n",
    "import utils  # worker count tracking\n",
    "\n",
    "executor = \"dask\"   # \"dask\" or \"taskvine\" or \"dask_gateway\"\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "NanoAODSchema.warn_missing_crossrefs = False # silences warnings about branches we will not use here\n",
    "\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__}\")\n",
    "\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f891e515-6bef-4ac3-a72f-9fec2b13c25c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scheduler_options = {}\n",
    "\n",
    "if executor == \"dask_gateway\":\n",
    "    num_workers = 100   #number of workers desired\n",
    "    from dask.distributed import LocalCluster, Client, progress\n",
    "    from dask_gateway import Gateway\n",
    "    import pathlib\n",
    "    \n",
    "    gateway = Gateway()\n",
    "    clusters=gateway.list_clusters()\n",
    "    cluster = gateway.connect(clusters[0].name)\n",
    "    client = cluster.get_client()\n",
    "    cluster.scale(num_workers)\n",
    "    # %%\n",
    "    def set_env(dask_worker):\n",
    "        path = str(pathlib.Path(dask_worker.local_directory) / 'access_token')\n",
    "        os.environ[\"BEARER_TOKEN_FILE\"] = path\n",
    "        os.chmod(path, 0o600)\n",
    "        os.chmod(\"/etc/grid-security/certificates\", 0o755)\n",
    "\n",
    "    client.wait_for_workers(num_workers)\n",
    "    client.upload_file(\"/etc/cmsaf-secrets/access_token\")\n",
    "    client.run(set_env)\n",
    "        \n",
    "else:\n",
    "    # by default use dask   \n",
    "    # local: single thread, single worker\n",
    "    from dask.distributed import LocalCluster, Client, progress\n",
    "    \n",
    "    # cluster = LocalCluster(n_workers=1, processes=False, threads_per_worker=1)\n",
    "    # client = Client(cluster)\n",
    "    client = Client(\"tls://localhost:8786\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "684f1240-a754-4dd1-b861-1bfac65ec288",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "fname = \"zstd_files.json\"\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": "05e87e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply optional filtering to limit number of input files\n",
    "\n",
    "# limite to the first N files per container, None if no limit\n",
    "LIMIT_NUM_FILES = 100\n",
    "\n",
    "# limit to the first N containers, None if no limit\n",
    "LIMIT_NUM_CONTAINERS = 30\n",
    "\n",
    "fileset = coffea.dataset_tools.max_files(fileset, LIMIT_NUM_FILES)\n",
    "\n",
    "if LIMIT_NUM_CONTAINERS is not None:\n",
    "    fileset = dict((k,v) for i, (k,v) in enumerate(fileset.items()) if i <LIMIT_NUM_CONTAINERS)\n",
    "\n",
    "print(f\"number of input files after filter: {sum([len(f['files']) for f in fileset.values()])}\")\n",
    "utils.worker_tracking.save_fileset(fileset, MEASUREMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "633689aa-419d-4c79-a0b8-b36dad9abd79",
   "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",
    "    BRANCH_LIST = [\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",
    "        ]\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[\"Events\"].num_entries\n",
    "            for b in BRANCH_LIST:\n",
    "                size_uncompressed += f[\"Events\"][b].uncompressed_bytes\n",
    "\n",
    "            for _ in f[\"Events\"].iterate(expressions=BRANCH_LIST):\n",
    "                pass\n",
    "\n",
    "            size_read = f.file.source.num_requested_bytes\n",
    "        exception = None\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": "35833a73-9c4e-4b74-8029-8b138b0f8c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# perform computation\n",
    "print(f\"running with {len(all_files)} files\")\n",
    "\n",
    "utils.worker_tracking.start_tracking_workers(client, MEASUREMENT_PATH)  # track worker count in background\n",
    "with performance_report(filename=MEASUREMENT_PATH/\"dask-report-plain-uproot.html\"):\n",
    "    tasks = [dask.delayed(uproot_open_materialize)(f) for f in all_files]\n",
    "    t0 = time.perf_counter()\n",
    "    out = ak.Array(dask.compute(*tasks))\n",
    "    t1 = time.perf_counter()\n",
    "\n",
    "utils.worker_tracking.stop_tracking_workers()\n",
    "\n",
    "print(f\"wall clock time: {t1-t0:.2f}s\")\n",
    "utils.worker_tracking.save_measurement(out, t0, t1, MEASUREMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e9943a-e242-4b36-9dba-ef892096a232",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load measurements from file again\n",
    "timestamps, nworkers, avg_num_workers = utils.worker_tracking.get_timestamps_and_counts(MEASUREMENT_PATH)  # worker count info\n",
    "out, t0, t1 = utils.worker_tracking.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": "db3663b9-40b9-4b68-919f-95798b9676de",
   "metadata": {
    "tags": []
   },
   "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[::10]:  # 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.worker_tracking.plot_worker_count(timestamps, nworkers, avg_num_workers, times_for_rates, instantaneous_rates, MEASUREMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "def5675b-6a0b-40c8-bccb-260f5d6f4938",
   "metadata": {
    "tags": []
   },
   "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": "32a93978-161a-49d3-b0e2-70e55929bec6",
   "metadata": {
    "tags": []
   },
   "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": "9bcf3496-c845-4431-9bf7-e50055d8e75f",
   "metadata": {
    "tags": []
   },
   "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",
    "\n",
    "xlim = ax.get_xlim()\n",
    "ylim = ax.get_ylim()\n",
    "xvals = np.linspace(*xlim, 100)\n",
    "ax.plot(xvals, xvals/(25*1_000), label=\"25 kHz\", linestyle=\"-\", c=\"C1\")\n",
    "ax.plot(xvals, xvals/(50*1_000), label=\"50 kHz\", linestyle=\"--\", c=\"C2\")\n",
    "ax.plot(xvals, xvals/(100*1_000), label=\"100 kHz\", linestyle=\":\", c=\"C3\")\n",
    "ax.set_xlim([0, xlim[1]])\n",
    "ax.set_ylim([0, ylim[1]])\n",
    "ax.legend()\n",
    "\n",
    "fig.savefig(MEASUREMENT_PATH / \"runtime_vs_nevts.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38f998ea-4052-4587-92e1-806c3e466452",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}