{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Parallel 3D RM-synthesis and RM-CLEAN with Dask\n", "\n", "The `rm_lite.tools_3d` tools build **lazy** `dask.array` graphs: `rmsynth_3d`\n", "and `rmclean_3d` map their per-spectrum math over spatial chunks with\n", "`dask.array.map_blocks`/`dask.delayed`, one task per chunk. Nothing runs until\n", "you call `.compute()` (or write to zarr). That means the *same* code runs\n", "serially, across threads/processes on one machine, or across a\n", "`dask.distributed` cluster on many nodes; only the scheduler changes.\n", "\n", "This notebook shows how to pick that scheduler, and the one built-in config\n", "that matters for it: `nufft_nthreads`.\n", "\n", "A stage's inner work is either compiled C or a per-pixel Python loop, and the\n", "two want opposite schedulers; a full run mixes both:\n", "\n", "| Stage | Inner work | Releases the GIL? | Best scheduler |\n", "|---|---|---|---|\n", "| RM-synth / RMSF (`rmsynth_3d`) | `finufft` NUFFT (C) | yes | **threads** |\n", "| Stokes I fit (`rmsynth_3d` with `stokes_i`) | per-pixel `scipy` fit (Python) | no | **processes** |\n", "| RM-CLEAN (`rmclean_3d`) | per-pixel Hogbom loop (Python) | no | **processes** |\n", "| All stages together (synth + fit + CLEAN) | mix of C and Python | partly | **distributed** (processes, a few threads each) |\n", "\n", "The distinction is the Python GIL:\n", "\n", "- **Threads** share one process and its memory, so there is no data copying and\n", " startup is instant. But the GIL lets only one thread run Python at a time, so\n", " threads only speed up work that *releases* the GIL. The NUFFT does (it is\n", " compiled C), so RM-synth scales on threads.\n", "- **Processes** each run a separate Python interpreter, so they run pure-Python\n", " code truly in parallel, at the cost of copying inputs/outputs between them and\n", " slower startup. The per-pixel Stokes I fit and the RM-CLEAN minor loop are\n", " pure Python (they hold the GIL), so they only scale on processes.\n", "\n", "Rule of thumb: **C/NUFFT work wants threads, per-pixel Python loops want\n", "processes.** The rest of this notebook shows each in turn.\n", "\n", "## Where to run it: picking a cluster\n", "\n", "The stage above decides threads vs processes; *where* you run decides the\n", "cluster. It is the same lazy graph either way, so start simple and scale up:\n", "\n", "| Where you are running | Cluster / launcher | Notes |\n", "|---|---|---|\n", "| One machine, quick run | none: just `.compute(scheduler=\"threads\"` or `\"processes\")` | simplest, no `Client` |\n", "| One machine, want a dashboard + memory management | `dask.distributed.LocalCluster` | `processes=` and `threads_per_worker=` set the layout |\n", "| HPC batch scheduler | `dask_jobqueue.SLURMCluster` (or `PBSCluster`, `LSFCluster`) | one job per node, `cluster.scale(jobs=N)` |\n", "| A few machines over SSH | `dask.distributed.SSHCluster` | no batch system needed |\n", "| Kubernetes / cloud | `dask_kubernetes`, `dask-gateway`, or Coiled | beyond this notebook |\n", "\n", "Whichever cluster you pick, a `Client` makes it the default scheduler, so the\n", "rm-lite calls do not change. The per-stage threads/processes choice and the\n", "thread counts in the final summary stay the same.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "import time\n", "\n", "import dask\n", "import dask.array as da\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "from dask.diagnostics import ProgressBar\n", "from rm_lite.utils.synthesis import faraday_simple_spectrum\n", "\n", "plt.rcParams[\"figure.dpi\"] = 150\n", "rng = np.random.default_rng(42)" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "## A small chunked cube\n", "\n", "Simulate a compact Stokes Q/U cube with a per-pixel RM, then wrap it in a\n", "`dask.array` chunked **spatially only** (the frequency axis is always kept\n", "whole, since RM-synthesis needs every channel per pixel). A real cube would\n", "come from `read_fits_cube_dask` or `dask.array.from_zarr`; here we build one\n", "in memory so the notebook is self-contained.\n", "\n", "The single most important tuning rule: **make more chunks than you have\n", "workers**, so every worker stays busy. We use small chunks here to get many\n", "blocks out of a tiny demo cube." ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "n_freq, ny, nx = 48, 96, 96\n", "freq_hz = np.linspace(800e6, 1088e6, n_freq)\n", "\n", "y_grid, x_grid = np.mgrid[0:ny, 0:nx]\n", "rm_map = 60.0 * (x_grid / nx - 0.5) * 2\n", "frac_pol = 0.6\n", "\n", "stokes_q = np.empty((n_freq, ny, nx), dtype=np.float64)\n", "stokes_u = np.empty((n_freq, ny, nx), dtype=np.float64)\n", "for j in range(ny):\n", " for i in range(nx):\n", " spec = faraday_simple_spectrum(\n", " freq_hz, frac_pol=frac_pol, psi0_deg=30.0, rm_radm2=rm_map[j, i]\n", " )\n", " stokes_q[:, j, i] = spec.real + rng.normal(0, 0.02, n_freq)\n", " stokes_u[:, j, i] = spec.imag + rng.normal(0, 0.02, n_freq)\n", "\n", "# freq axis whole (-1), 16x16 spatial chunks -> 6x6 = 36 blocks\n", "q_dask = da.from_array(stokes_q, chunks=(-1, 16, 16))\n", "u_dask = da.from_array(stokes_u, chunks=(-1, 16, 16))\n", "print(f\"chunks: freq={q_dask.chunks[0]}, spatial blocks={q_dask.numblocks}\")\n", "assert len(q_dask.chunks[0]) == 1 # freq axis must be a single chunk" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "## The built-in config: `nufft_nthreads`\n", "\n", "`finufft` is itself OpenMP-threaded and by default grabs *all* cores for a\n", "single transform. When Dask already runs many chunks in parallel, that nests\n", "two layers of threading and oversubscribes the cores, which is *slower* than\n", "either layer alone.\n", "\n", "So `rmsynth_3d` defaults to **`nufft_nthreads=1`**: one thread per NUFFT call,\n", "and let the Dask scheduler parallelise across chunks instead. This is the fast\n", "config whenever you have many chunks (the usual case). It is a performance\n", "knob only; it does not change the result. The cell below asserts exactly that." ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [], "source": [ "from rm_lite.tools_3d.rmsynth import rmsynth_3d\n", "\n", "synth_kw = {\"phi_max_radm2\": 120.0, \"d_phi_radm2\": 8.0}\n", "\n", "# default: one finufft thread per chunk, dask parallelises across chunks\n", "res = rmsynth_3d(q_dask, u_dask, freq_hz, **synth_kw)\n", "# nufft_nthreads=0 hands each call back to finufft's own OpenMP (all cores)\n", "res_omp = rmsynth_3d(q_dask, u_dask, freq_hz, nufft_nthreads=0, **synth_kw)\n", "\n", "# ProgressBar (from dask.diagnostics) shows a live bar for any local-scheduler\n", "# compute (synchronous/threads/processes).\n", "with ProgressBar():\n", " fdf_default = res.fdf_dirty_cube.compute()\n", " fdf_omp = res_omp.fdf_dirty_cube.compute()\n", "np.testing.assert_allclose(fdf_default, fdf_omp, atol=1e-10)\n", "print(\"nufft_nthreads changes speed, not the result: OK\")" ] }, { "cell_type": "markdown", "id": "6", "metadata": {}, "source": [ "## Choosing a scheduler\n", "\n", "The lazy cube computes under whatever scheduler you ask for. For the NUFFT\n", "stage the threaded scheduler is the right default (the NUFFT releases the\n", "GIL). Pass it explicitly with `scheduler=` on `.compute()`, or set it\n", "globally with `dask.config.set`.\n", "\n", "The result is identical across schedulers; only the wall-clock differs. We\n", "assert the invariance and print the timings (timings are machine-dependent,\n", "so they are not asserted)." ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "def timed(label, arr, **kw):\n", " t = time.time()\n", " with ProgressBar():\n", " out = arr.compute(**kw)\n", " print(f\"{label:24s} {time.time() - t:6.3f}s\")\n", " return out\n", "\n", "\n", "serial = timed(\"synchronous\", res.fdf_dirty_cube, scheduler=\"synchronous\")\n", "threaded = timed(\"threads\", res.fdf_dirty_cube, scheduler=\"threads\")\n", "np.testing.assert_allclose(serial, threaded, atol=1e-10)\n", "print(\"same result under both schedulers: OK\")" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "peak_pi = np.max(np.abs(threaded), axis=0)\n", "peak_rm = res.phi_arr_radm2[np.argmax(np.abs(threaded), axis=0)]\n", "\n", "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.5))\n", "im1 = ax1.imshow(peak_pi, origin=\"lower\")\n", "fig.colorbar(im1, ax=ax1, label=\"Peak polarised intensity\")\n", "ax1.set(title=\"Recovered peak PI\")\n", "im2 = ax2.imshow(peak_rm, origin=\"lower\", cmap=\"RdBu_r\", vmin=-60, vmax=60)\n", "fig.colorbar(im2, ax=ax2, label=\"RM [rad m^-2]\")\n", "ax2.set(title=\"Recovered RM\")\n", "# the recovered RM gradient should track the input gradient in x\n", "assert peak_rm[:, -1].mean() > peak_rm[:, 0].mean()" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "## The GIL-bound stages: use processes\n", "\n", "RM-CLEAN (`rmclean_3d`) and the optional per-pixel Stokes I fit are pure-Python\n", "per-pixel loops. The threaded scheduler cannot overlap them (the GIL lets only\n", "one thread run Python at a time), so here we compute with\n", "`scheduler=\"processes\"`. Contrast the RM-synth cells above, which used threads\n", "because the NUFFT is C: same lazy cubes, opposite scheduler, and the choice is\n", "driven entirely by whether the inner work is C or Python.\n", "\n", "For these stages `nufft_nthreads` is irrelevant (no NUFFT in the CLEAN loop),\n", "and the tuning flips from *one process, many threads* to *many processes, one\n", "thread each*. For anything beyond one machine, use a `dask.distributed`\n", "cluster instead (next section) for the dashboard, memory management, and\n", "multi-node reach." ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "from rm_lite.tools_3d.rmclean import rmclean_3d_from_synth\n", "\n", "# The processes scheduler spawns workers. Use the 'fork' start method so it\n", "# runs in a plain script/notebook without an `if __name__ == '__main__'`\n", "# guard (the default on Linux; set explicitly here for macOS too).\n", "dask.config.set({\"multiprocessing.context\": \"fork\"})\n", "\n", "# CLEAN needs the dirty FDF + RMSF cubes; keep them chunked spatially.\n", "clean = rmclean_3d_from_synth(res, auto_mask=6, auto_threshold=1, max_iter=200)\n", "# lazy, like the synth results; compute across processes (GIL-bound loop).\n", "# num_workers=4 keeps this tiny demo light; use n_cores in practice.\n", "with ProgressBar():\n", " clean_fdf = clean.clean_fdf_cube.compute(scheduler=\"processes\", num_workers=4)\n", "print(f\"clean FDF cube: {clean_fdf.shape}, {clean_fdf.dtype}\")\n", "assert clean_fdf.shape == threaded.shape" ] }, { "cell_type": "markdown", "id": "11", "metadata": {}, "source": [ "## Dask distributed (single node and multi-node)\n", "\n", "A `dask.distributed` `Client` becomes the default scheduler for every\n", "subsequent `.compute()`, so **no rm-lite code changes**. Its workers are\n", "processes, each with a thread pool, which serves both stage classes at once:\n", "the NUFFT parallelises over each worker's threads, while the GIL-bound fit and\n", "CLEAN parallelise over the worker processes.\n", "\n", "`distributed` is a `dev`/`docs` extra of rm-lite (`pip install distributed`),\n", "so it is available whenever these docs run. The runnable cell below uses\n", "**in-process threaded workers** (`processes=False`) so it executes inside the\n", "docs build. A real deployment uses process workers, as in the\n", "`LocalCluster`/`SLURMCluster` snippets below.\n", "\n", "Progress with a `Client` works differently from `ProgressBar` (see\n", "[dask/distributed#926](https://github.com/dask/distributed/issues/926)):\n", "submit the graph as a **future** with `client.compute(...)` (it runs in the\n", "background), then call `progress(future)` to draw the bar and `wait(future)`\n", "to block until it finishes. The live **dashboard** at `client.dashboard_link`\n", "is the richer view.\n", "\n", "```python\n", "# Single fat node: processes for the GIL-bound stages, a few threads each\n", "# for the NUFFT. Pin OMP so finufft does not re-oversubscribe.\n", "import os\n", "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n", "from dask.distributed import Client, LocalCluster\n", "cluster = LocalCluster(n_workers=6, threads_per_worker=2)\n", "client = Client(cluster)\n", "\n", "# Many nodes via a batch scheduler (SLURM shown; PBS/LSF analogous):\n", "from dask_jobqueue import SLURMCluster\n", "cluster = SLURMCluster(\n", " cores=12, processes=6, memory=\"64GB\",\n", " job_extra_directives=[\"--export=OMP_NUM_THREADS=1\"],\n", ")\n", "cluster.scale(jobs=10) # 10 nodes\n", "client = Client(cluster)\n", "```\n", "\n", "For multi-node runs the input cube must be reachable from every worker: put it\n", "on a shared filesystem, or (better) store it as **zarr** and read with\n", "`da.from_zarr`, so each worker reads only its own chunks. Write outputs the\n", "same way with `write_zarr_group`." ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "from dask.distributed import Client, LocalCluster, progress\n", "\n", "with (\n", " LocalCluster(\n", " n_workers=1, threads_per_worker=4, processes=True, dashboard_address=None\n", " ) as cluster,\n", " Client(cluster) as client,\n", "):\n", " future = client.compute(clean.clean_fdf_cube)\n", " progress(future, notebook=False)\n", " clean_dist = future.result()\n", "\n", "np.testing.assert_allclose(clean_dist, clean_fdf, atol=1e-8)\n", "print(\"\\ndistributed Client: same CLEAN result, OK\")" ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "## Summary: what to run\n", "\n", "Two settings are always the same, so set them once and forget them:\n", "\n", "- **`nufft_nthreads=1`** (the default, an argument to `rmsynth_3d`) and\n", " **`OMP_NUM_THREADS=1`** (an environment variable, set *before* rm-lite or\n", " finufft is imported, on every worker). Both say the same thing: let Dask do\n", " the parallelism, not finufft's internal OpenMP. Leaving finufft's OpenMP on\n", " is the main way to make this *slower*, by oversubscribing the cores.\n", " The one time to change `nufft_nthreads` is when Dask is **not** spreading\n", " chunks across the cores: the `synchronous` scheduler, or a cube with fewer\n", " chunks than cores. Then set `nufft_nthreads=0` so finufft's OpenMP fills the\n", " otherwise-idle cores. The invariant is `nufft_nthreads x dask_workers`\n", " `≈ n_cores`: the default is the `1 x n_cores` split, `nufft_nthreads=0` is the\n", " opposite `n_cores x 1` extreme, only sensible at about one chunk. Never run\n", " `nufft_nthreads > 1` *and* many workers.\n", "- **More chunks than workers.** Pick `target_chunk_mb` (on the FITS reader) so\n", " there are several chunks per worker; that balances load and caps per-task\n", " memory.\n", "\n", "The only real choice is the scheduler and its thread/worker counts, and it\n", "follows from what you are running. Counts below are for one 12-core node;\n", "scale by your core count, and keep total threads across all workers equal to\n", "the cores (Dask threads x processes = cores, with `OMP_NUM_THREADS=1`):\n", "\n", "| What you are computing | Scheduler / cluster | Threads to set, and where | Why |\n", "|---|---|---|---|\n", "| RM-synth only (`rmsynth_3d`, no Stokes I fit) | threaded | `.compute(scheduler=\"threads\", num_workers=12)` (= n_cores) | the NUFFT releases the GIL, so threads parallelise it |\n", "| RM-CLEAN, or RM-synth with the Stokes I fit | processes | `.compute(scheduler=\"processes\", num_workers=12)` (= n_cores) | per-pixel Python loops are GIL-bound; only processes overlap them |\n", "| All stages in one run (synth + fit + CLEAN) | `dask.distributed` | `LocalCluster(n_workers=6, threads_per_worker=2)` (6 x 2 = 12), `OMP_NUM_THREADS=1` | one `Client` for the whole pipeline: processes for the Python loops, threads for the NUFFT |\n", "| A cube too big for one machine, or many nodes | `dask.distributed` | `SLURMCluster(cores=12, processes=6, ...)` per node, `OMP_NUM_THREADS=1` | the same split, scaled across nodes |\n", "\n", "For a mixed synth-then-CLEAN pipeline, do not switch schedulers between\n", "stages: use one distributed `Client` of **several processes, a few threads\n", "each** (the 6 x 2 = 12 layout above). Processes cover the GIL-bound CLEAN,\n", "threads cover the NUFFT, and the cheap NUFFT stage loses less to the process\n", "layout than CLEAN would lose to a threads-only one. On a batch scheduler the\n", "same split is `SLURMCluster(cores=12, processes=6, ...)` (6 workers x 2\n", "threads per job), with `OMP_NUM_THREADS=1` exported to the job.\n" ] } ], "metadata": { "kernelspec": { "display_name": "rm-lite (3.11.13)", "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.11.13" } }, "nbformat": 4, "nbformat_minor": 5 }