feat: per-file size tracking for ntuples and monitoring improvements (#31)

* track per-file sizes for ntuples
* use per-file sizes for analysis workload
* track rate in xrdcp notebook

Author: alexander-held

atlas/ntuple_production/cache_everything.ipynb

@@ -82,19 +82,22 @@ def rucio_file_paths(name, num_files_expected):
     """file paths from rucio list-file-replicas call"""
     cmd = f"rucio list-file-replicas --protocols root {name}"
     output = subprocess.check_output(cmd, shell=True)
-    rses_and_paths = re.findall(r"(\w+): (root:\/\/.*?)\s", output.decode())
+    size_unit_rse_path = re.findall(r"(\d+\.\d+)\s(\wB).+?([\w-]+): (root:\/\/.*?)\s", output.decode())
 
     # select a single RSE for each file
-    filenames = sorted(set([rp[1].split("/")[-1] for rp in rses_and_paths]))
+    filenames = sorted(set([rp[-1].split("/")[-1] for rp in size_unit_rse_path]))
     unique_paths = []
+    sizes_GB = []
     for filename in filenames:
-        fpaths = [rp for rp in rses_and_paths if filename in rp[1]]
+        matches = [m for m in size_unit_rse_path if filename in m[-1]]
         # pick MWT2_UC_LOCALGROUPDISK match by default, otherwise first in the list
-        fpath = next((fp for fp in fpaths if fp[0] == "MWT2_UC_LOCALGROUPDISK"), fpaths[0])[1]
-        unique_paths.append(fpath)
+        match = next((m for m in matches if m[2] == "MWT2_UC_LOCALGROUPDISK"), matches[0])
+        unique_paths.append(match[3])
+        size_to_GB = lambda num, unit: float(num) * {"kB": 1e-6, "MB": 1e-3, "GB": 1}[unit]
+        sizes_GB.append(size_to_GB(*match[:2]))
 
     assert len(unique_paths) == num_files_expected
-    return unique_paths
+    return unique_paths, sizes_GB
 
 
 def process_one_category(category, container_list, production_map):
@@ -132,8 +135,10 @@ def process_one_category(category, container_list, production_map):
             metadata[container]["size_output_GB"] = info_output["size_GB"]
 
             # add xrootd file paths
-            paths = rucio_file_paths(production_map[container]["output"], info_output["nfiles"])
+            paths, sizes = rucio_file_paths(production_map[container]["output"], info_output["nfiles"])
             metadata[container]["files_output"] = paths
+            metadata[container]["sizes_output_GB"] = sizes
+            assert abs(sum(sizes) - info_output["size_GB"]) < 0.01  # agree within 10 MB
 
     return {category: metadata}
 
@@ -168,4 +173,4 @@ def save_full_metadata(production_map, fname, max_workers=8):
     production_map = parse_job_json(fname_bigpanda)
 
     fname_full = "file_metadata.json.gz"
-    metadata = save_full_metadata(production_map, fname_full)
+    metadata = save_full_metadata(production_map, fname_full, max_workers=8)

atlas/ntuple_production/collect_file_metadata.py

@@ -0,0 +1,141 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff6ca0fb-e358-4dbe-9951-e5167758d030",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datetime\n",
+    "import gzip\n",
+    "import json\n",
+    "import time\n",
+    "import os\n",
+    "\n",
+    "import dask\n",
+    "from dask.distributed import Client\n",
+    "import numpy as np\n",
+    "import matplotlib.dates as mdates\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "client = Client(\"tls://localhost:8786\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d55ec891-e0c2-4f43-ba35-2077ecafcc07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_input(max_size_GB = None):\n",
+    "    with gzip.open(\"file_metadata.json.gz\") as f:\n",
+    "        dataset_info = json.loads(f.read().decode())\n",
+    "\n",
+    "    all_files = []\n",
+    "    all_sizes_GB = []\n",
+    "    for containers_for_category in dataset_info.values():\n",
+    "        for container, metadata in containers_for_category.items():\n",
+    "            if metadata[\"files_output\"] is None:\n",
+    "                continue\n",
+    "            for fname, size in zip(metadata[\"files_output\"], metadata[\"sizes_output_GB\"]):\n",
+    "                all_files.append(fname)\n",
+    "                all_sizes_GB.append(size)\n",
+    "                if max_size_GB and sum(all_sizes_GB) > max_size_GB:\n",
+    "                    return all_files, all_sizes_GB\n",
+    "    return all_files, all_sizes_GB\n",
+    "\n",
+    "all_files, all_sizes_GB = get_input(max_size_GB = None)  # limit list to specific total size\n",
+    "print(f\"list of {len(all_files)} files with total size {sum(all_sizes_GB):.2f} GB\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c5cb8c1-ebf6-41f7-a12d-feff9a0456b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def run_xrdcp(fname, size):\n",
+    "    t0 = time.time()\n",
+    "    os.system(f\"xrdcp {fname} /dev/null -f\")\n",
+    "    t1 = time.time()\n",
+    "    return {\"t0\": t0, \"t1\": t1, \"GBread\": size}\n",
+    "\n",
+    "t0 = time.time()\n",
+    "tasks = [dask.delayed(run_xrdcp)(fname, size) for fname, size in zip(all_files, all_sizes_GB)]\n",
+    "res = dask.compute(*tasks)\n",
+    "t1 = time.time()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c04faacc-8136-4ef9-910a-98b8660cb4d6",
+   "metadata": {},
+   "source": [
+    "track egress: [link](https://grafana.mwt2.org/d/EKefjM-Sz/af-network-200gbps-challenge?orgId=1&from=now-1h&to=now&viewPanel=panel-205&refresh=5s)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd03f327-857a-4d80-a013-3a13549de37f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "total_runtime_sum = sum(r[\"t1\"] - r[\"t0\"] for r in res)\n",
+    "\n",
+    "print(f\"processtime: {total_runtime_sum:.2f} s\")\n",
+    "print(f\" -> data rate per worker: {sum(all_sizes_GB) * 8 / total_runtime_sum:.2f} Gbps\")\n",
+    "\n",
+    "print(f\"walltime: {t1-t0:.2f} s\")\n",
+    "print(f\" -> total data rate: {sum(all_sizes_GB) * 8 / (t1-t0):.2f} Gbps\")\n",
+    "\n",
+    "starts = np.asarray([r[\"t0\"] for r in res])\n",
+    "ends = np.asarray([r[\"t1\"] for r in res])\n",
+    "GBread = [r[\"GBread\"] for r in res]\n",
+    "rates_per_chunk = GBread / (ends - starts)\n",
+    "\n",
+    "t_samples = np.linspace(t0, t1, 100)\n",
+    "rate_samples = []\n",
+    "for t in t_samples:\n",
+    "    mask = np.logical_and((starts <= t), (t < ends))\n",
+    "    rate_samples.append(float(sum(rates_per_chunk[mask]) * 8))\n",
+    "\n",
+    "print(f\"total data read from data rate integral: {sum((t_samples[1] - t_samples[0]) * np.asarray(rate_samples)) / 8:.2f} GB\")\n",
+    "t_samples = [datetime.datetime.fromtimestamp(t) for t in t_samples.tolist()]\n",
+    "\n",
+    "fig, ax = plt.subplots()\n",
+    "ax.plot(t_samples, rate_samples, marker=\"v\", linewidth=0)\n",
+    "ax.set_xlabel(\"time\")\n",
+    "ax.tick_params(axis=\"x\", labelrotation=45)\n",
+    "ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))\n",
+    "ax.set_ylabel(\"data rate [Gbps]\")\n",
+    "ax.set_ylim([0, ax.get_ylim()[1] * 1.1])\n",
+    "fig.savefig(\"xrdcp_rate.png\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "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.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

atlas/ntuple_production/distributed_xrdcp.ipynb

@@ -238,9 +238,6 @@ def get_fileset(campaign_filter: list | None = None, dsid_filter: list | None =
     with gzip.open(fname) as f:
         dataset_info = json.loads(f.read().decode())
 
-    if max_files_per_sample is not None:
-        print(f"[WARNING] limiting files per sample to {max_files_per_sample}, input size estimate is invalid")
-
     # construct fileset
     fileset = {}
     input_size_GB = 0
@@ -266,7 +263,7 @@ def get_fileset(campaign_filter: list | None = None, dsid_filter: list | None =
                 "files": dict((path, "reco") for path in metadata["files_output"][:num_files]),
                 "metadata": {"dsid": dsid, "campaign": campaign, "category": category, "weight_xs": weight_xs, "lumi": lumi}
             }
-            input_size_GB += metadata["size_output_GB"]
+            input_size_GB += sum(metadata["sizes_output_GB"][:num_files])
 
     print(f"[INFO] fileset has {len(fileset)} categories with {sum([len(f["files"]) for f in fileset.values()])} files total, size is {input_size_GB:.2f} GB")
     return fileset, input_size_GB