fix(cluster): size the worker pool by what the job was granted, not by the machine (#616)

CHANGELOG.md

@@ -198,6 +198,27 @@ All notable changes to PyBNF are documented below. This project adheres to
   says what it has not bisected.
 
 ### Fixed
+- **A multi-machine fit sizes its worker pool by what the job was granted, not by how big the
+  machine is (#616).** PyBNF decided how many worker processes to start on each node by calling
+  `multiprocessing.cpu_count()`, which reports every processor the machine has whatever the job
+  scheduler granted. On a cluster where this was measured, a job that asked for **4** CPUs was
+  told the node had **128**: PyBNF would have started one worker per processor and overshot the
+  job's real capacity **32-fold**. Every worker is a separate process, so that multiplies memory
+  use and leaves the workers competing for the same four CPUs — a fit that runs slower than it
+  would have on the share it was given, or that runs out of memory. The defect could hide because
+  a job that asks for *whole* nodes gets the right answer by coincidence: there the two numbers
+  are equal.
+  Both launchers now take the count from `Cluster.cpus_per_node`, the one place that decides it,
+  which prefers **`$SLURM_CPUS_ON_NODE`** — what the allocation granted, and the only one of the
+  three numbers that describes the *allocation* rather than the process asking, so it is still
+  right for a worker started on another machine — then **`dask.system.CPU_COUNT`**, the machine's
+  processors narrowed by CPU affinity and by any cgroup quota, which is what a single-machine run
+  already sizes itself by, and only then the machine's whole processor count, which is correct
+  only when nothing is limiting the job. The count and **which of the three it came from** are
+  written to the log at the start of the run, so an unexpected number of workers can be traced to
+  the number PyBNF believed; setting `parallel_count` still overrides all of it, and the log then
+  names that key as the source. `-t slurm-srun`, which already read what SLURM granted, is
+  unchanged apart from logging the source.
 - **Multi-machine fits run the command dask actually installs, so a cluster run gets past
   its first second (#615).** PyBNF started remote workers by running **`dask-ssh`** — one of
   three standalone scripts (with `dask-scheduler` and `dask-worker`) that distributed stopped

docs/adr/0089-every-locally-spawned-dask-worker-is-single-threaded-so-a-fits-thread-policy-no-longer-depends-on-whether-parallel-count-is-set.md

@@ -66,6 +66,12 @@ CPU affinity and cgroup quotas, so a run confined to 4 cores of a 64-core host g
 64. (`setup_cluster` still uses `multiprocessing.cpu_count()` for `dask-ssh`, where the number
 being computed is a remote node's core count anyway.)
 
+**Superseded by issue #616:** the parenthesis above was the defect. A remote node's core count is
+not what a *job* holds on that node -- a job granted 4 CPUs of a 128-processor node was told 128 --
+so `setup_cluster` now takes its default from `Cluster.cpus_per_node`, which prefers what the
+scheduler granted (`$SLURM_CPUS_ON_NODE`) and falls back to `dask.system.CPU_COUNT` before ever
+reaching `multiprocessing.cpu_count()`. Both launchers now decide this in that one place.
+
 **Total concurrency is unchanged.** Measured on a 6-core machine, before and after:
 
 ```text

docs/adr/0122-a-cluster-that-cannot-be-logged-into-is-still-a-cluster-so-pybnf-starts-its-workers-inside-the-allocation-slurm-already-granted-rather-than-over-ssh.md

@@ -120,6 +120,16 @@ that number, it asks SLURM for that many CPUs, and a number taken from the whole
 refused. The SSH launcher still uses `multiprocessing.cpu_count()`; that it does so is issue #616,
 and it has to be fixed there rather than here.
 
+**Superseded by issue #616:** the launchers no longer differ. `Cluster.cpus_per_node` is now the one
+place either of them decides how many workers a node gets, and it returns the count together with a
+phrase naming where it came from, which both launchers log. The precedence is `$SLURM_CPUS_ON_NODE`,
+then `dask.system.CPU_COUNT`, then `multiprocessing.cpu_count()`: the scheduler's number is
+preferred because it describes the *allocation* rather than the process asking, so it remains the
+right number for a worker the SSH launcher starts on some other machine; the affinity- and
+cgroup-aware count is next because it is what the operating system will actually permit here, and is
+what a local run already sizes itself by; the whole machine is last, since it is correct only when
+nothing is limiting the job at all.
+
 ### `scheduler_file` names an output under this launcher
 
 Everywhere else, `scheduler_file` means *attach to a cluster someone else brought up* -- PyBNF starts

docs/cluster.rst

@@ -73,10 +73,25 @@ An example batch script, the ``-t slurm`` one with a single word changed::
 
     pybnf -c tcr-ss.conf -t slurm-srun -o
 
-By default, each node runs one single-threaded worker process per CPU the job was granted on that node. Setting ``parallel_count`` overrides that with a total number of worker processes over all nodes, divided evenly among them.
-
 Two log files are written to the output directory: ``dask_scheduler.log`` and ``dask_workers.log``. The second is where ``srun`` reports anything that went wrong with placing the workers, and PyBNF quotes from it in the error message if no worker ever registers.
 
+.. _workercount:
+
+How many workers run on each node
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+By default -- with either launcher -- each node runs one single-threaded worker process per CPU **the job was granted** on that node. That is the number your allocation asked for, not the number of processors the machine happens to have: a job given 4 CPUs of a 128-processor node runs 4 workers on it, not 128. Each worker is a separate process, so a pool sized to the machine rather than to the job would multiply memory use and leave the workers competing for the same few CPUs.
+
+PyBNF takes that number from the first of these that is available:
+
+* ``$SLURM_CPUS_ON_NODE``, which is what SLURM granted the job on a node;
+* the CPU count dask derives for this process, which is the machine's processors narrowed by CPU affinity and by any cgroup CPU quota -- the same number a single-machine PyBNF run sizes itself by; or
+* the machine's whole processor count, which is correct only when nothing is limiting the job.
+
+Which number was used, and which of the three it came from, is written to the log at the start of the run, so an unexpected worker count can be traced to the number PyBNF believed.
+
+Setting ``parallel_count`` overrides all of this with a total number of worker processes over all nodes, divided evenly among them; the log then names ``parallel_count`` as the source. Nodes of different sizes still get equal shares.
+
 
 TORQUE/PBS
 ----------
@@ -102,7 +117,7 @@ PyBNF uses `Dask.distributed <http://distributed.readthedocs.io/en/latest/index.
 
 For a local (single-machine) run, PyBNF builds a Dask ``LocalCluster`` with one thread per worker process, and as many worker processes as there are available cores (or ``parallel_count`` of them, if that key is set). One thread per worker is not tunable: the simulation backends hold process-wide state that is not thread-safe, so two jobs running concurrently in one process can interfere with each other.
 
-In the automatic PyBNF setup, the command ``dask ssh`` is run on one of the available nodes (which becomes the scheduler node), with all available nodes as arguments (which become the worker nodes). ``dask ssh`` is run with ``--nthreads 1`` and ``--nworkers`` equal to the number of available cores per node. The default number of available processes per core is the value returned by ``multiprocessing.cpu_count()``; this default can be overridden by specifying the ``parallel_count`` key equal to the total number of processes over all nodes. This entire automatic setup with ``dask ssh`` can be overridden as described below. If overriding the automatic setup, it is recommended to keep ``nthreads`` equal to 1 for SBML models because the SBML simulator is not thread safe.
+In the automatic PyBNF setup, the command ``dask ssh`` is run on one of the available nodes (which becomes the scheduler node), with all available nodes as arguments (which become the worker nodes). ``dask ssh`` is run with ``--nthreads 1`` and ``--nworkers`` equal to the number of CPUs the job was granted on a node, as described under `How many workers run on each node`_; this default can be overridden by specifying the ``parallel_count`` key equal to the total number of processes over all nodes. This entire automatic setup with ``dask ssh`` can be overridden as described below. If overriding the automatic setup, it is recommended to keep ``nthreads`` equal to 1 for SBML models because the SBML simulator is not thread safe.
 
 For manual configuration, you will need to run the series of commands described below. All of these commands must remain running during the entire PyBNF run. Utilites such as ``nohup`` or ``screen`` are helpful for keeping multiple commands running at once. 
 

docs/config_keys.rst

@@ -971,7 +971,7 @@ Parallel Computing
 
   Each parallel job runs in its own **single-threaded worker process**, whether or not this key is set: the simulation backends hold process-wide state that is not thread-safe, so PyBNF never places two concurrently running jobs in one process. This key therefore sets a process count, not a thread count. Lowering it is the way to reduce the memory a run uses, since each worker process holds its own copy of the models.
 
-  Default: Use all available cores -- one single-threaded worker per core. Locally, the core count comes from Dask, which honors CPU affinity and cgroup quotas (so a run confined to 4 cores gets 4 workers, not the host's full count). On a cluster, the number of available cores per node is determined by running ``multiprocessing.cpu_count()`` from the scheduler node; with ``cluster_type = slurm-srun`` it is instead the number of cores SLURM granted the job on a node (``$SLURM_CPUS_ON_NODE``).
+  Default: Use all available cores -- one single-threaded worker per core. Locally, the core count comes from Dask, which honors CPU affinity and cgroup quotas (so a run confined to 4 cores gets 4 workers, not the host's full count). On a cluster, with either ``cluster_type``, it is the number of CPUs the job was granted on a node (``$SLURM_CPUS_ON_NODE``), falling back to the affinity- and cgroup-aware count Dask derives and then to the machine's whole processor count; PyBNF logs the number it used and which of the three it came from. See :ref:`How many workers run on each node <workercount>`.
 
   Example:
 

pybnf/cluster.py

@@ -15,6 +15,10 @@
 host-based support and dask never enables its GSSAPI support -- so on such a cluster the
 login fails no matter what the user configures, and no amount of ``ssh-keygen`` helps. See
 docs/adr/0122 for the full argument.
+
+Both launchers size their default worker pool from what the *job* was granted rather than
+from how big the machine is, and record which number they used and where it came from
+(#616); ``Cluster.cpus_per_node`` is the one place that decides it.
 """
 
 
@@ -35,6 +39,11 @@
 from multiprocessing import cpu_count
 from distributed import Client, LocalCluster
 from dask import __version__ as daskv
+# What the operating system will actually let this process run on: the machine's
+# processors narrowed by CPU affinity and by any cgroup CPU quota. Bound to a module
+# global (rather than read through ``dask.system``) both because dask computes it once at
+# import time and because that makes it substitutable in tests, the way ``cpu_count`` is.
+from dask.system import CPU_COUNT as DASK_CPU_COUNT
 from distributed import __version__ as distributedv
 from .config import init_logging, reinit_logging
 
@@ -342,8 +351,9 @@ def setup_cluster(node_string, out_dir, parallel_count=None):
 
         :param node_string: A string composed of a list of compute nodes
         :param out_dir: A directory for cluster logging output
-        :param parallel_count: Total number of parallel threads to use over all nodes. If None, use all available threads
-            (the dask ssh default)
+        :param parallel_count: Total number of single-threaded worker processes over all
+            nodes, divided evenly among them. If None, one worker per CPU the job was
+            granted on a node (``cpus_per_node``)
         :return: subprocess.Popen
         """
         # Ask before launching, so a dask that cannot do this reads as a configuration
@@ -362,11 +372,21 @@ def setup_cluster(node_string, out_dir, parallel_count=None):
         # (pyproject pins dask/distributed >=2024.1.0).
         nodes = node_string.split()
         if parallel_count is None:
+            # One worker per CPU the *job* holds on a node, not per processor the machine
+            # has (#616). The two differ by more than an order of magnitude on a job that
+            # asked for a small share of a large node, and the machine's count is the one
+            # that oversubscribes it. The source is logged because a user who sees an
+            # unexpected worker count needs to know which number PyBNF believed.
+            n_per_node, source = Cluster.cpus_per_node()
+            logger.info('Starting %i worker process(es) on each of %i node(s), one per CPU, '
+                        'from %s' % (n_per_node, len(nodes), source))
             dask_ssh_cmd = [*DASK_CLI, 'ssh', *nodes,
-                            '--log-directory', out_dir, '--nthreads', '1', '--nworkers', str(cpu_count())]
+                            '--log-directory', out_dir, '--nthreads', '1', '--nworkers', str(n_per_node)]
         else:
             n_per_node = int(np.ceil(parallel_count/len(nodes)))
-            logger.info('Manually setting %i workers per node' % n_per_node)
+            logger.info('Manually setting %i worker process(es) on each of %i node(s), from the '
+                        'parallel_count key (%i over all nodes)'
+                        % (n_per_node, len(nodes), parallel_count))
             dask_ssh_cmd = [*DASK_CLI, 'ssh', *nodes,
                             '--log-directory', out_dir, '--nworkers', str(n_per_node), '--nthreads', '1']
         # Capture stderr to a temp file rather than a PIPE: dask ssh stays
@@ -442,23 +462,46 @@ def srun_scheduler_file(config):
     @staticmethod
     def cpus_per_node():
         """
-        The number of CPUs the running job was granted on a node.
-
-        ``$SLURM_CPUS_ON_NODE`` is what the allocation actually granted; ``cpu_count()`` is
-        the size of the whole machine, which is only the same number when whole nodes were
-        allocated. The srun launcher reads the former because it does not merely count
-        workers with it -- it also asks SLURM for that many CPUs per task, and a request
-        larger than the allocation is refused outright. (The SSH launcher still uses
-        ``cpu_count()``; correcting that is issue #616, and it has to be corrected there
-        too rather than here.)
-
-        :return: CPUs granted per node, falling back to the machine's core count
-        :rtype: int
+        The number of CPUs the running job was granted on a node, and where that came from.
+
+        Both launchers size their default worker pool with this, because the number that
+        decides how many processes to start has to describe what the *job* holds, not what
+        the machine has (#616). ``multiprocessing.cpu_count()`` answers the second question:
+        it reports every processor on the machine whatever the scheduler granted, so a job
+        given 4 CPUs of a 128-processor node is told 128, and one worker process per
+        processor oversubscribes it 32-fold -- 32 times the memory, and workers competing
+        for time rather than a fit that runs faster. The two numbers agree only when whole
+        nodes were allocated, which is why the defect stayed hidden.
+
+        Three sources are consulted, best first:
+
+        * ``$SLURM_CPUS_ON_NODE`` -- what the allocation granted on a node. Preferred
+          because it is the only one that describes the *allocation* rather than the process
+          doing the asking, so it is still the right number for a worker the SSH launcher
+          starts on some other machine. (When nodes differ in size it describes this node;
+          per-node counts are issue #617.)
+        * ``dask.system.CPU_COUNT`` -- what the operating system will let this process run
+          on: the machine's processors narrowed by CPU affinity and by any cgroup CPU quota.
+          This is the number a local run already sizes itself by, and it is the right one
+          whenever the job is confined on the machine PyBNF is running on but no scheduler
+          published a count.
+        * ``multiprocessing.cpu_count()`` -- the whole machine, correct only when nothing is
+          limiting the job at all, and reached only if neither number above is usable.
+
+        The srun launcher does not merely count workers with this: it also asks SLURM for
+        that many CPUs per task, and a request larger than the allocation is refused
+        outright.
+
+        :return: CPUs granted per node, and a phrase naming where that number came from
+        :rtype: tuple
         """
         granted = os.environ.get('SLURM_CPUS_ON_NODE', '').strip()
         if granted.isdigit() and int(granted) > 0:
-            return int(granted)
-        return cpu_count()
+            return int(granted), 'what SLURM granted the job ($SLURM_CPUS_ON_NODE)'
+        if DASK_CPU_COUNT > 0:
+            return DASK_CPU_COUNT, ("this process's CPU affinity and cgroup limits "
+                                    '(dask.system.CPU_COUNT)')
+        return cpu_count(), "this machine's whole processor count (multiprocessing.cpu_count)"
 
     @staticmethod
     def srun_worker_command(scheduler_file, node_count, parallel_count=None):
@@ -475,7 +518,7 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None):
         :return: the srun argument list
         :rtype: list
         """
-        granted = Cluster.cpus_per_node()
+        granted, source = Cluster.cpus_per_node()
         if parallel_count is None:
             n_per_node = granted
         else:
@@ -490,8 +533,9 @@ def srun_worker_command(scheduler_file, node_count, parallel_count=None):
         # a deliberately oversubscribed parallel_count still runs (SLURM refuses a request
         # for more CPUs than the job holds) rather than failing the run.
         cpus_per_task = max(1, min(n_per_node, granted))
-        logger.info('Starting %i worker process(es) per node on %i node(s), %i CPU(s) per node'
-                    % (n_per_node, node_count, cpus_per_task))
+        logger.info('Starting %i worker process(es) per node on %i node(s), asking SLURM for %i '
+                    'CPU(s) per node; %i available per node, from %s'
+                    % (n_per_node, node_count, cpus_per_task, granted, source))
         return ['srun',
                 '--nodes', str(node_count), '--ntasks', str(node_count),
                 '--ntasks-per-node', '1', '--cpus-per-task', str(cpus_per_task),