Use MPI version and SLURM mulitnode allocation in notebook for powder diffraction

ac9e28a1 · Egor Sobolev · ff7cd862 · ac9e28a1
Commit ac9e28a1 authored 6 months ago by Egor Sobolev
--- a/nb/powder.ipynb
+++ b/nb/powder.ipynb
@@ -10,7 +10,14 @@
    "runno = 189\n",
    "propno = 900463\n",
    "\n",
-    "ref_geom_fn = \"/gpfs/exfel/exp/SPB/202430/p900429/usr/Shared/geom/p5681_r64_67_refined.geom\"  # initial geometry file name"
+    "ref_geom_fn = \"/gpfs/exfel/exp/SPB/202430/p900429/usr/Shared/geom/p5681_r64_67_refined.geom\"  # initial geometry file name\n",
+    "\n",
+    "# image averager parameters\n",
+    "mask_fn = \"/gpfs/exfel/exp/SPB/202431/p900463/usr/Shared/mask/mask_hvoff_20240807.h5\"\n",
+    "\n",
+    "# slurm paramenters\n",
+    "num_nodes = 16\n",
+    "slurm_partition = \"exfel\""
   ]
  },
  {
@@ -54,7 +61,13 @@
   "source": [
    "fn = f\"powder_sum_p{propno:06d}_r{runno:04d}.h5\"\n",
    "if not os.path.exists(fn):\n",
-    "    !srun -p exfel geomtools-powder-sum -p {propno} -r {runno} -D ^scratch/pycal/proc -m ^usr/Shared/mask/mask_hvoff_20240807.h5 -l agipd1m SPB_DET_AGIPD1M-1"
+    "    sopts = f\"-N {num_nodes} -p {slurm_partition} --ntasks-per-node=1\"\n",
+    "    mpiopts = (\"--map-by NODES -x LD_LIBRARY_PATH -x PATH --bind-to none \"\n",
+    "               \"--mca btl ^smcuda --mca rcache ^gpusm,rgpusm --mca accelerator ^cuda\")\n",
+    "    powder_sum_args = (\n",
+    "        f\"-p {propno} -r {runno} -m {mask_fn} --legacy -l agipd1m SPB_DET_AGIPD1M-1\"\n",
+    "    )\n",
+    "    !salloc {sopts} mpirun {mpiopts} geomtools-powder-sum {powder_sum_args}"
   ]
  },
  {

 %% Cell type:code id:c134ff02-0eba-4146-803b-2b69f9ccaef9 tags:

 ``` python
 runno = 189
 propno = 900463

 ref_geom_fn = "/gpfs/exfel/exp/SPB/202430/p900429/usr/Shared/geom/p5681_r64_67_refined.geom"  # initial geometry file name
+
+# image averager parameters
+mask_fn = "/gpfs/exfel/exp/SPB/202431/p900463/usr/Shared/mask/mask_hvoff_20240807.h5"
+
+# slurm paramenters
+num_nodes = 16
+slurm_partition = "exfel"
 ```

 %% Cell type:code id:0bdf665e-9ad7-4a40-a59e-23b428312f4f tags:

 ``` python
 # !geomtools-powder-center -g {ref_geom_fn} -c AgBh ../../../nb/powder_sum_p{propno:06d}_r{runno:04d}.h5
 ```

 %% Cell type:code id:504eca33-7f3d-4b7e-9ce1-9edcadff6387 tags:

 ``` python
 import os
 import numpy as np
 import matplotlib.pyplot as plt
 import h5py

 from extra_data import open_run
 from extra_geom import agipd_asic_seams, AGIPD_1MGeometry
 from extra_geom.motors import AGIPD_1MMotors
 from extra.components import AGIPD1MQuadrantMotors

 from geomtools.powder.calibrants import get_calibrant
 from geomtools.powder.misc import agipd_borders
 from geomtools.powder import PowderDiffraction
 ```

 %% Cell type:code id:33f4f527-d61e-45c1-8301-1be7c6aa8edf tags:

 ``` python
 fn = f"powder_sum_p{propno:06d}_r{runno:04d}.h5"
 if not os.path.exists(fn):
-    !srun -p exfel geomtools-powder-sum -p {propno} -r {runno} -D ^scratch/pycal/proc -m ^usr/Shared/mask/mask_hvoff_20240807.h5 -l agipd1m SPB_DET_AGIPD1M-1
+    sopts = f"-N {num_nodes} -p {slurm_partition} --ntasks-per-node=1"
+    mpiopts = ("--map-by NODES -x LD_LIBRARY_PATH -x PATH --bind-to none "
+               "--mca btl ^smcuda --mca rcache ^gpusm,rgpusm --mca accelerator ^cuda")
+    powder_sum_args = (
+        f"-p {propno} -r {runno} -m {mask_fn} --legacy -l agipd1m SPB_DET_AGIPD1M-1"
+    )
+    !salloc {sopts} mpirun {mpiopts} geomtools-powder-sum {powder_sum_args}
 ```

 %% Cell type:code id:e9c147f0-e011-46cb-bf0f-e0ef01a6d8a4 tags:

 ``` python
 with h5py.File(fn, "r") as f:
    # conditions
    clen = f["powderSum/conditions/cameraLen"][()]
    photon_en = f["powderSum/conditions/photonEnergy"][()]
    lmd = f["powderSum/conditions/waveLength"][()]
    detector_id = f["powderSum/conditions/detectorId"][()].decode()

    # average image
    img_mean = f["powderSum/image/mean"][:]
    img_std = f["powderSum/image/std"][:]
    img_count = f["powderSum/image/count"][:]
    num_frames = f["powderSum/image/numFrames"][:]
    mask = f["powderSum/image/mask"][:].astype(bool)
    modules = f["powderSum/image/modules"][:]

 print("Photon energy", photon_en, "(keV)")
 print("Wave length", lmd, "(nm)")
 print("Camera length", clen, "(m)")
 ```

 %% Cell type:code id:8fdf43d7-db13-4493-a030-3132337e1512 tags:

 ``` python
 # geometry
 run = open_run(propno, runno, data="raw")
 ref_geom = AGIPD_1MGeometry.from_crystfel_geom(ref_geom_fn)
 tracker = AGIPD_1MMotors(ref_geom)
 motors = AGIPD1MQuadrantMotors(run)
 geom = tracker.geom_at(motors.most_frequent_positions())
 ```

 %% Cell type:code id:332a77f6-b3b7-4e13-b6e4-5838391949df tags:

 ``` python
 img = np.copy(img_mean)
 geom.plot_data(img, norm="log", interpolation="none")
 plt.show()
 ```

 %% Cell type:code id:83dd3cb1-00d5-4337-80c5-e83d9eeec0b1 tags:

 ``` python
 rings = get_calibrant("AgBh")
 borders = agipd_borders()
 pw = PowderDiffraction(geom, rings, clen, lmd,
                       border_mask=~borders, make_shadow_mask=True)
 pw.fit(img_mean, mask, num_frames)
 pw.refinement_info()
 pw.draw()
 plt.show()
 ```

 %% Cell type:code id:1bd08493-2562-4f92-82aa-eb866d5af791 tags:

 ``` python
 pw.draw_radial_profile(img_mean, mask)
 plt.show()
 ```