Fix n_jobs=-1 oversubscribing cores under CPU affinity restrictions - #4724
Fix n_jobs=-1 oversubscribing cores under CPU affinity restrictions#4724lochhh wants to merge 6 commits into
n_jobs=-1 oversubscribing cores under CPU affinity restrictions#4724Conversation
- `fix_job_kwargs` caps `n_jobs` using `get_usable_cpu_count` and tests need to compute their expected boundary with the same function
- `set_global_job_kwargs` merges via `dict.update()` and doesn't remove `chunk_memory` key
| def get_usable_cpu_count(): | ||
| """ | ||
| Number of CPUs usable by this process. | ||
| """ | ||
| if hasattr(os, "process_cpu_count"): | ||
| # Python >= 3.13; respects both Windows job object and Linux cgroup/cpuset restrictions | ||
| return os.process_cpu_count() | ||
| if hasattr(os, "sched_getaffinity"): | ||
| # Respects Linux cgroup/cpuset restrictions | ||
| return len(os.sched_getaffinity(0)) | ||
| return os.cpu_count() |
There was a problem hiding this comment.
Would be good to test this function out on as many real life systems as possible to try and find any edge cases. I've tested on mac, local linux desktop Edinburgh supercomputer.
@lochhh have you tested this on the HPC?
@zm711 could you quickly test on windows?
@alejoe91 could you check in your set up?
There was a problem hiding this comment.
requested 4 cpus on HPC: srun --cpus-per-task=4 --pty bash -i
| System | Python | process_cpu_count | sched_getaffinity | cpu_count |
|---|---|---|---|---|
| HPC Linux 6.8 (Ubuntu) | 3.11.5 | None | 4 | 24 |
| HPC Linux 6.8 (Ubuntu) | 3.13.2 | 4 | 4 | 24 |
| Laptop Windows 11 | 3.14.4 | 16 | AttributeError | 16 |
There was a problem hiding this comment.
I also just want to make it clear in general that even working locally on a Windows 10/11 station setting n_jobs to be 7 (for 8 cores) would lead to all 8 cores being used at various stages because:
- n_jobs is ultimately setting the number of processes while the OS chooses to schedule across cores so even if you say 1 job per core that doesn't mean that you are telling the computer to use 1 process/core and all processes could go to one core or setting 2 processes could utilize all cores
- we don't control scipy, numpy, torch, sklearn allocation of resources so they can also can affect this behavior and lead to
n_jobs =/= n_cores.
That being said I have no problem with the logic for -1 being changed to not lead to too many processes for the number of cores (ie trying to keep the semantic that -1 = n_cpus available). But I just want us to remember not to overpromise what n_jobs is actually setting.
|
@lochhh thanks for this! I agree it was an hassle! Please let's give us a couple of weeks as many of us are off for vacation at the moment. I agree with @chrishalcrow that we should test on a bunch of platforms to make sure it works as expected everywhere |
alejoe91
added
core
4 participants
Closes #4723
This PR
get_usable_cpu_count()injob_tools.py, preferringos.process_cpu_count()(Python >= 3.13, cross-platform affinity-aware) oros.sched_getaffinity()(older Python on Linux) and falls back toos.cpu_count()when neither is availableget_usable_cpu_count()get_usable_cpu_count()everywheren_jobswas previously resolved fromos.cpu_count()get_usable_cpu_count()test_fix_job_kwargs:set_global_job_kwargsmerges viadict.update()and can't remove keys, so the manual "restore" step that usesset_global_job_kwargsleftchunk_memorykey inglobal_job_kwargs. The fix usesreset_global_job_kwargs()following the same pattern intest_globals.py