From b5320a2c78b82d9c9b5a14c36fe44d16fdda9148 Mon Sep 17 00:00:00 2001 From: Bill Hlavacek Date: Fri, 21 Aug 2026 12:46:04 -0600 Subject: [PATCH] fix(cluster): size the worker pool by what the job was granted, not by the machine (#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 it. On the 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. `Cluster.cpus_per_node` is now the one place either launcher decides this, and it returns the count together with a phrase naming where it came from. The precedence is: 1. `$SLURM_CPUS_ON_NODE` -- what the allocation granted. Preferred because it is the only one of the three that describes the *allocation* rather than the process asking, so it is still the right number for a worker the SSH launcher starts on some other machine. 2. `dask.system.CPU_COUNT` -- the machine's processors narrowed by CPU affinity and by any cgroup quota, i.e. what the operating system will actually permit here. This is the number a single-machine run already sizes itself by. 3. `multiprocessing.cpu_count()` -- the whole machine, correct only when nothing is limiting the job at all. Both launchers log the count, the node count and the source, so a user who sees an unexpected number of workers can trace it to the number PyBNF believed. Setting `parallel_count` still overrides all of it, and its branch now names that key as the source rather than logging a bare "Manually setting N workers per node". `-t slurm-srun`, which already read `$SLURM_CPUS_ON_NODE`, is unchanged apart from the added provenance in its log line. The tests give the three sources three different numbers, so each one pins *which* source PyBNF consulted rather than merely a plausible count; the reported case (4 granted of a 128-processor node) is asserted directly, and all four new setup_cluster assertions go red against the old code. ADR-0089 and ADR-0122 both recorded the old split -- 0089's parenthesis that a remote node's core count "is a remote node's core count anyway" was the defect itself -- and now carry superseding notes. Verified by running the real SSH bring-up against an unreachable host under `SLURM_CPUS_ON_NODE=4`: dask's CLI accepts the constructed argv (`--nthreads 1 --nworkers 4`) and the process is still running after the bring-up wait. The srun launcher was re-exercised end to end with stand-ins for srun/scontrol on PATH -- a real scheduler, two real workers, a real task through the client, no orphans after teardown. Full suite green (4472 passed, 23 skipped); docs build clean under -W --keep-going. --- CHANGELOG.md | 21 +++ ...epends-on-whether-parallel-count-is-set.md | 6 + ...rm-already-granted-rather-than-over-ssh.md | 10 ++ docs/cluster.rst | 21 ++- docs/config_keys.rst | 2 +- pybnf/cluster.py | 86 ++++++++--- tests/test_cluster.py | 145 +++++++++++++++--- 7 files changed, 241 insertions(+), 50 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9efc46b2b..efffcbf30 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/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 b/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 index 00582caab..f7ba3d835 100644 --- a/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 +++ b/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 diff --git a/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 b/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 index 643c2143a..fd53a510f 100644 --- a/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 +++ b/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 diff --git a/docs/cluster.rst b/docs/cluster.rst index 4fc869671..3069d72de 100644 --- a/docs/cluster.rst +++ b/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 `. Example: diff --git a/pybnf/cluster.py b/pybnf/cluster.py index ad2b92e3d..6529974b1 100644 --- a/pybnf/cluster.py +++ b/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), diff --git a/tests/test_cluster.py b/tests/test_cluster.py index 3a470a243..ba97b2000 100644 --- a/tests/test_cluster.py +++ b/tests/test_cluster.py @@ -21,8 +21,11 @@ * ``cluster.run`` (subprocess) — **mock**: a recorder returning a fake proc with canned ``.stdout`` bytes, or raising ``TimeoutExpired`` / ``CalledProcessError``. - * ``cluster.Popen`` / ``cluster.time.sleep`` / ``cluster.cpu_count`` — - **mock**: capture the command; stub the 10s sleep so the test is instant. + * ``cluster.Popen`` / ``cluster.time.sleep`` — **mock**: capture the command; + stub the 10s sleep so the test is instant. + * the worker-count sources — ``$SLURM_CPUS_ON_NODE``, ``cluster.DASK_CPU_COUNT`` + and ``cluster.cpu_count`` — **stubbed to three different numbers**, so a test + pins which source PyBNF consulted and not merely a plausible count (#616). * ``cluster.Client`` / ``cluster.LocalCluster`` / ``cluster.init_logging`` / ``cluster.reinit_logging`` and ``Cluster.read_node_names`` / ``Cluster.setup_cluster`` — **fakes** recording their call args, so the @@ -263,9 +266,15 @@ def test_the_group_queried_is_the_one_dask_itself_reads(self, monkeypatch): class TestSetupCluster: - def _patch(self, monkeypatch, cpu=4, returncode=None, stderr_bytes=b''): - """Patch the three externals setup_cluster touches: Popen (capture the - command), time.sleep (don't actually wait 10s), cpu_count (deterministic). + def _patch(self, monkeypatch, granted=None, affinity=4, cpu=64, + returncode=None, stderr_bytes=b''): + """Patch what setup_cluster touches: Popen (capture the command), + time.sleep (don't actually wait 10s), and every source the default worker + count can come from, so the count is deterministic *and* it is visible which + source produced it -- ``$SLURM_CPUS_ON_NODE`` (what the job was granted; + removed from the environment unless ``granted`` is passed), the + affinity/cgroup count dask derives, and the whole machine's ``cpu_count()``. + The three defaults are deliberately three different numbers. The fake proc's ``poll()`` returns ``returncode`` (None = still running, the healthy default); if ``stderr_bytes`` is given the fake writes it to the stderr file setup_cluster handed to Popen, so the early-exit error @@ -280,15 +289,21 @@ def fake_popen(*args, **kwargs): monkeypatch.setattr(cluster, 'Popen', fake_popen) monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) + monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) + if granted is not None: + monkeypatch.setenv('SLURM_CPUS_ON_NODE', str(granted)) + monkeypatch.setattr(cluster, 'DASK_CPU_COUNT', affinity) monkeypatch.setattr(cluster, 'cpu_count', lambda: cpu) return popen_calls - def test_default_parallel_count_uses_cpu_count(self, monkeypatch): - """parallel_count=None ⇒ dask ssh's own default of one worker per CPU: - ``--nthreads 1 --nworkers {cpu_count()}`` (note this branch's flag order is - --nthreads then --nworkers). Oracle: the exact argument list (ROB-3: an argv - list launched with no shell, each node its own entry) with cpu_count()=7.""" - popen_calls = self._patch(monkeypatch, cpu=7) + def test_default_worker_count_is_what_the_job_was_granted(self, monkeypatch): + """parallel_count=None ⇒ one single-threaded worker per CPU **the job holds + on a node**: ``--nthreads 1 --nworkers {cpus_per_node()}`` (note this + branch's flag order is --nthreads then --nworkers). Oracle: the exact + argument list (ROB-3: an argv list launched with no shell, each node its own + entry) with SLURM having granted 7, over an affinity count of 4 and a + machine of 64.""" + popen_calls = self._patch(monkeypatch, granted=7, affinity=4, cpu=64) proc = cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=None) assert proc.poll() is None @@ -302,6 +317,55 @@ def test_default_parallel_count_uses_cpu_count(self, monkeypatch): assert kwargs['stderr'] is not cluster.DEVNULL assert hasattr(kwargs['stderr'], 'read') + def test_default_worker_count_is_not_the_size_of_the_machine(self, monkeypatch): + """#616, stated as the reported case: a job granted 4 CPUs of a + 128-processor node must start 4 workers per node, not 128. This is the + assertion the old code failed -- it passed ``multiprocessing.cpu_count()``, + which reports every processor the machine has whatever the scheduler + granted, so the pool overshot the job by 32x. The two numbers agree only + when whole nodes were allocated, which is why the defect could hide.""" + popen_calls = self._patch(monkeypatch, granted=4, affinity=4, cpu=128) + cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=None) + + (args, _), = popen_calls + assert args[0][args[0].index('--nworkers') + 1] == '4' + assert '128' not in args[0] + + def test_default_worker_count_falls_back_to_affinity_not_the_machine(self, monkeypatch): + """With no scheduler count published, the next-best number is what the + operating system will let this process run on -- CPU affinity narrowed by + any cgroup quota, which is what dask derives and what a local PyBNF run + already sizes itself by -- rather than the whole machine.""" + popen_calls = self._patch(monkeypatch, granted=None, affinity=6, cpu=64) + cluster.Cluster.setup_cluster('n1', '/out', parallel_count=None) + + (args, _), = popen_calls + assert args[0][args[0].index('--nworkers') + 1] == '6' + + def test_default_worker_count_and_its_source_are_logged(self, monkeypatch, caplog): + """A user who sees an unexpected number of workers has to be able to find + out which number PyBNF believed and where it read it, so the count, the node + count and the source are all logged (#616).""" + self._patch(monkeypatch, granted=7) + with caplog.at_level('INFO', logger='pybnf.cluster'): + cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=None) + + line, = [r.message for r in caplog.records if 'worker process' in r.message] + assert '7' in line and '2 node' in line + assert 'SLURM_CPUS_ON_NODE' in line + + def test_explicit_parallel_count_is_logged_as_its_own_source(self, monkeypatch, caplog): + """When parallel_count decides the count, the log says so: the number did + not come from the job's CPUs, and a user comparing the two needs to know + which one is in force.""" + self._patch(monkeypatch, granted=7) + with caplog.at_level('INFO', logger='pybnf.cluster'): + cluster.Cluster.setup_cluster('n1 n2', '/out', parallel_count=6) + + line, = [r.message for r in caplog.records if 'worker process' in r.message] + assert 'parallel_count' in line + assert 'SLURM_CPUS_ON_NODE' not in line + def test_the_launcher_is_dask_ssh_through_this_interpreter(self, monkeypatch): """#615: the command is the ``dask ssh`` *subcommand*, run through the interpreter running PyBNF -- not the standalone ``dask-ssh`` script, which @@ -616,7 +680,7 @@ def test_local_and_dask_ssh_defaults_agree_on_one_thread(self, monkeypatch): monkeypatch.setattr(cluster, 'Popen', lambda *a, **k: popen_calls.append((a, k)) or _FakeDaskProc()) monkeypatch.setattr(cluster.time, 'sleep', lambda *_: None) - monkeypatch.setattr(cluster, 'cpu_count', lambda: 7) + monkeypatch.setenv('SLURM_CPUS_ON_NODE', '7') cluster.Cluster.setup_cluster('n1', '/out', parallel_count=None) (ssh_args, _), = popen_calls cmd = ssh_args[0] @@ -828,23 +892,54 @@ def test_scheduler_file_chooses_where_it_is_written(self): class TestCpusPerNode: + """The one place either launcher decides how many workers a node gets (#616). + + The three sources are given three different numbers throughout, so each test + pins *which* one was consulted rather than merely a plausible count. + """ + + def _sources(self, monkeypatch, granted=None, affinity=6, cpu=64): + monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) + if granted is not None: + monkeypatch.setenv('SLURM_CPUS_ON_NODE', str(granted)) + monkeypatch.setattr(cluster, 'DASK_CPU_COUNT', affinity) + monkeypatch.setattr(cluster, 'cpu_count', lambda: cpu) def test_reads_what_slurm_granted(self, monkeypatch): - """$SLURM_CPUS_ON_NODE is what the *allocation* granted, which is the - number the launcher can actually ask SLURM for.""" - monkeypatch.setenv('SLURM_CPUS_ON_NODE', '12') - monkeypatch.setattr(cluster, 'cpu_count', lambda: 64) - assert cluster.Cluster.cpus_per_node() == 12 + """$SLURM_CPUS_ON_NODE is what the *allocation* granted. It is preferred + over both local numbers because it describes the allocation rather than the + process asking, so it is still right for a worker started on another + machine -- and because it is the number the srun launcher can actually ask + SLURM for.""" + self._sources(monkeypatch, granted=12, affinity=6, cpu=64) + count, source = cluster.Cluster.cpus_per_node() + assert count == 12 + assert 'SLURM_CPUS_ON_NODE' in source @pytest.mark.parametrize('value', [None, '', 'many', '0', '-4']) - def test_falls_back_to_the_machine_core_count(self, monkeypatch, value): - """With nothing usable in the environment, fall back to the machine's own - core count (the number the SSH launcher uses unconditionally).""" - monkeypatch.delenv('SLURM_CPUS_ON_NODE', raising=False) - if value is not None: - monkeypatch.setenv('SLURM_CPUS_ON_NODE', value) - monkeypatch.setattr(cluster, 'cpu_count', lambda: 64) - assert cluster.Cluster.cpus_per_node() == 64 + def test_an_unusable_slurm_value_falls_through_to_the_affinity_count(self, monkeypatch, value): + """With no usable scheduler count, the next-best number is what the OS will + let this process run on -- CPU affinity narrowed by any cgroup quota, the + number dask derives and a local PyBNF run already uses -- not the machine.""" + self._sources(monkeypatch, granted=value, affinity=6, cpu=64) + count, source = cluster.Cluster.cpus_per_node() + assert count == 6 + assert 'affinity' in source + + def test_the_whole_machine_is_the_last_resort(self, monkeypatch): + """The machine's own processor count is right only when nothing is limiting + the job at all, so it is reached only when neither better number exists.""" + self._sources(monkeypatch, granted=None, affinity=0, cpu=64) + count, source = cluster.Cluster.cpus_per_node() + assert count == 64 + assert 'machine' in source + + def test_a_granted_count_never_reports_the_machine_count(self, monkeypatch): + """The #616 oracle, stated once for both launchers: a job granted a small + share of a large node is sized by the share. The old SSH-launcher code + returned 128 here, oversubscribing the job 32-fold.""" + self._sources(monkeypatch, granted=4, affinity=4, cpu=128) + assert cluster.Cluster.cpus_per_node()[0] == 4 class TestSrunWorkerCommand: