Extract make_gpu_plan and it's users from #1207

[sk1p]

For discussion - as this changes the way resources are allocated in a quite significant way, we might need to iterate a bit on the interface.

Edit @uellue, discussion with @sk1p:

We should use the information of lstopo and os.sched_getaffinity(0) to get the physical cores we are actually allowed to use, as opposed to the ones available in the system. Only with this information we can pin the PipelinedExecutor to the correct cores and start the correct number of workers. This is relevant for running within a container, in particular in CI.

• Query os.sched_getaffinity(0) and topology to determine physical cores to run on.

Contributor Checklist:

• I have added or updated my entry in the creators.json file
• I have added a changelog entry for my contribution
• I have added/updated documentation for all user-facing changes
• I have added/updated test cases
• I have included the rebuilt production build of the client (only if changes were made to the GUI)

Reviewer Checklist:

• /azp run libertem.libertem-data passed
• No import of GPL code from MIT code

[sk1p]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[sk1p]

I think there are two main topics to discuss:

1. The change of resource allocation - after this PR, a CUDA worker "consumes" a CPU core, where previously the CUDA workers were counted on top of the CPU workers

2. The interface itself - having to pass in cuda_info for spawning a cluster is a bit inconvenient (and backwards-incompatible); also it's convenient to have round-robin assignment for integer arguments as implemented in Allow integer arguments for cluster_spec #1336.

We might want to take a step back and enumerate the common use-cases and re-think how we can implement them in a convenient and backwards-compatible way.

[sk1p]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[matbryan52]

I think there are two main topics to discuss:

1. The change of resource allocation - after this PR, a CUDA worker "consumes" a CPU core, where previously the CUDA workers were counted on top of the CPU workers

2. The interface itself - having to pass in `cuda_info` for spawning a cluster is a bit inconvenient (and backwards-incompatible); also it's convenient to have round-robin assignment for integer arguments as implemented in [Allow integer arguments for `cluster_spec` #1336](https://github.com/LiberTEM/LiberTEM/pull/1336).

We might want to take a step back and enumerate the common use-cases and re-think how we can implement them in a convenient and backwards-compatible way.

For summary, the semantics were / are:

Originally:

• ...dask.cluster_spec => cpus[Iterable] create n workers equal to length of iterable
• ...pipelined._make_spec => cpus[Iterable] pin 1 worker each to exactly these cores, where supported
• for both => , cudas[Iterable] pin 1 worker each to exactly these GPU ids (allow repeat pinning)

After #1336 for both dask and pipelined:

• cpus[int] => create this many workers exactly
• cudas[int] => assign N workers round robin across all available GPUs

With this PR as-is:

• Require cuda_info dict if cudas is nonzero, else make_gpu_plan will fall over
• Automatic extra workers per GPU within RAM limits (by default)
• Maximum N workers per CUDA param
• Reserve one CPU per CUDA worker (prioritise spreading workers over the GPUs before many workers on one GPU if CPU-limited)
• cudas[int] => converted to range(int) i.e. pin to n GPU ids up to value (might not be available on a system)

[matbryan52]

How about the following interface:

cluster_spec(
    cpus: int | Iterable[int],
    cudas: int | Iterable[int],
    has_cupy: bool,
    cuda_scaling: Optional[Dict[str, Any]] = None,
):
    ...

• cpus: assign either n_workers [int] or len(cpus) [iterable], pin when possible, raise if pinning and num_cores < max(cpus)
• cudas: assign n_workers across available GPUs (round_robin), or pin workers to GPU ids in iterable (allow repeats), raise if iterable contains GPU ids not visible, or no GPUs are visible at all and cudas > 0
• has_cupy, unchanged
• cuda_scaling, when provided require cudas is an Iterable specifying which CUDA device ids to scale on, and require cuda_scaling to contain these ids at least, insert scaling params (ram-per-worker, max-per-device) here also. Use this info to scale the workers on the assigned cudas as in the current make_gpu_plan

For reserving CPUs for each CUDA worker - I feel like it's quite analysis dependent? Some tasks take both CPU and GPU power, others might be almost purely GPU. If we could guarantee pinning we could compromise and assign 2 GPU workers to a shared single CPU core, for example, but this is not even an option with Dask or some OSs. It could be an option within cuda_scaling, though.

[uellue]

The change is supposed to address three issues:

First, when running a set of UDFs, CUDA workers are performing not only GPU but also a lot of CPU work in case not all UDFs support CUDA. When doing CUDA work they also spin a core at full speed. That means CUDA workers do, in fact, consume a CPU resource on the system and having too many CUDA workers in parallel to saturating the cores with CPU workers will lead to oversubscription of the CPU. Second, @matbryan52, you had observed that multiple GPU workers improve throughput. And third, the number of partitions should be a multiple of the worker count to reduce stragglers with smaller datasets, meaning the number of workers that actually work on a UDF run should be predictable and constant independent on the UDFs that are executed.

I'll try to rework the interface in such a way that it is backwards-compatible!

[uellue]

I've tested the current state of the PR with ApplyMasksUDF on the large_raw fixture on ptycho, showing how much it can improve performance. Definitely worthwhile IMO. :-)

Baseline

------------------------------------------- benchmark: 1 tests ------------------------------------------
Name (time in s)        Min     Max    Mean  StdDev  Median     IQR  Outliers     OPS  Rounds  Iterations
---------------------------------------------------------------------------------------------------------
test_masks_udf       6.1322  6.4635  6.2467  0.1361  6.2004  0.1905       1;0  0.1601       5           1
---------------------------------------------------------------------------------------------------------

Current PR

---------------------------------------- benchmark 'udf': 1 tests ---------------------------------------
Name (time in s)        Min     Max    Mean  StdDev  Median     IQR  Outliers     OPS  Rounds  Iterations
---------------------------------------------------------------------------------------------------------
test_masks_udf       3.8892  3.9851  3.9467  0.0366  3.9516  0.0465       2;0  0.2534       5           1
---------------------------------------------------------------------------------------------------------

[uellue]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[uellue]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[uellue]

@sk1p can you have a look if the changes address your concerns? :-)

[uellue]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[uellue]

/azp run libertem.libertem-data

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[uellue]

Call with @sk1p and @matbryan52:

ctx = Context.make_with(cpus=10, gpus=10)

ctx = Context.make_with(cpus=10, gpus=(0, 0, 1, 1, 2))

# does this pin to CPU cores?
# Yes, if supported by the executor
# The first one missing means "whatever the Context thinks is the default executor". Currently "dask" (not sure if implemented already)
ctx = Context.make_with(cpus=(2, 4, 6, 8), gpus=(0, 0, 1, 1, 2))

ctx = Context.make_with('pipelined', cpus=(2, 4, 6, 8), gpus=(0, 0, 1, 1, 2))

# this will probably warn that the inline executor can't
# create multiple workers:
ctx = Context.make_with('inline', cpus=(2, 4, 6, 8), gpus=(0, 0, 1, 1, 2))

# just create a dask-distributed thingy:
ctx = Context.make_with('dask')

# this warns/raises that the resources can't be "created" as specified:
ctx = Context.make_with('dask-integration', cpus=(1, 2, 3))

def make_with(
    self,
    executor_type: Optional[str] = None,
    cpus: Optional[Union[int, Iterable[int]]] = None,
    gpus: Optional[Union[int, Iterable[int]]] = None,
    plot_class: Optional[str] = None,
):
    ...

TODO

• specify details of a "resource specifier"
• hybrid worker?

[uellue]

Collection of links on topology:

• https://github.com/gbitzes/cpu-topology/blob/master/cpu-topology.py
• https://man7.org/linux/man-pages/man1/nproc.1.html
• https://psutil.readthedocs.io/en/latest/#psutil.Process.cpu_affinity
• https://docs.rs/atopology/latest/atopology/ (thx @sk1p!)
• /sys/devices/system/cpu/

[sk1p]

Collection of links on topology:

• https://github.com/gbitzes/cpu-topology/blob/master/cpu-topology.py
• https://man7.org/linux/man-pages/man1/nproc.1.html
• https://psutil.readthedocs.io/en/latest/#psutil.Process.cpu_affinity
• https://docs.rs/atopology/latest/atopology/ (thx @sk1p!)
• /sys/devices/system/cpu/

And a way to check how this works when cgroup limits are present:

$ docker run --cpuset-cpus 0-2 --rm -it python python -c 'import os; print(os.sched_getaffinity(0))'
{0, 1, 2}
$ docker run --cpuset-cpus 0-2 --rm -it debian:stable nproc
3