Allow different number of cores for the worker and the simulation tool

[bennoschoenstein]

Description of the Issue:

Currently, the num_procs setting on a QUEENS scheduler controls two things at once:

1. How many CPUs the workload manager (e.g. SLURM) gives to a worker.
2. How many MPI processes the simulation tool is started with (e.g. mpirun -n N 4C ...).

Because both come from the same value, a worker that gets 4 CPUs from the cluster is forced to run the simulation with exactly 4 ranks - and vice versa. There is no way to:

• give a worker some headroom for Python pre- or post-processing while running the simulation with fewer ranks, or
• run many short-lived jobs in small worker slots whose simulations still need several ranks.

This couples two things that are conceptually independent: how big the slot on the cluster is and how many cores the simulation actually uses.

Proposed Solution:

Let drivers carry their own optional num_procs value:

• If the driver's num_procs is set, it is used for the simulation call.
• If it is not set (default), the driver falls back to the scheduler's num_procs — exactly today's behaviour.

This is fully backward compatible (default None = unchanged behaviour) and adds one optional argument on the driver side. The scheduler keeps its current role of requesting the worker resources from the workload manager; the driver gets to say how many MPI ranks the simulation tool actually uses.

Out of scope (future work): Different core counts for each step of a single driver pipeline (e.g. meshing on 1 core, simulation on N cores, plotting on 1 core in the same job). That would require a deeper architectural change. This issue is a prerequisite but does not attempt to solve it.

Action Items:

• Add optional num_procs argument to the driver base class.
• Use it for the jobscript rendering when set; fall back to the scheduler value otherwise.
• Add a small regression test with two drivers using different values on the same scheduler.
• One-paragraph note in the docs.

Related Issues:

No response

Interested Parties:

No response

[gilrrei]

@bennoschoenstein, this is an interesting use case that has constantly come up in the past :)

I want to motivate you to think outside existing QUEENS structures to enable use cases you consider out of scope. In the past, I tried to find a structure to enable the chain you mentioned on meshing -> simulation -> something else, just to discover that this ends up introducing a restrictive structure and will never be general-purpose.

With #237, you can pass callables to be evaluated embarrassingly parallel using the schedulers. Here an example

mesher = Simulation(
    parameters_1, driver_1, data_processor_1
)


def sequential_execution(
    sample: np.ndarray,
    job_id: int,
    job_dir: Path,
    num_procs: int,
    experiment_dir: Path,
    experiment_name: str,
):

    metadata = SimulationMetadata(job_id, sample, job_dir, file_name="job_metadata")

    # Make job dir
    job_dir.mkdir(exist_ok=True, parents=True)

    # Meshing
    with metadata.time_code("meshing"):
        results_model_1 = mesher(
            sample,
            job_id,
            job_dir / "mesh",
            1, # single core
            experiment_dir,
            experiment_name,
        )

    # Run 4C
    with metadata.time_code("run_4C"):
        driver_2.run(
            results_model_1,
            job_id,
            job_dir / "fourc_output",
            num_procs,
            experiment_dir=experiment_dir,
            experiment_name=experiment_dir,
        )

    # Extract the data if needed
    data = {
        "model_1": results_model_1,
        "results_model_2": data_processor_model_2(job_dir / "fourc_output"),
    }

    return data

which can then be used with a Dask scheduler. I think this encapsulates your use case, is general-purpose, and does not require any infrastructure change. QUEENS has evolved into a library with a wide range of heterogeneous scheduling use cases. I think yours is one of them, you are always free to build your custom tools without having to introduce the changes to main if they don't fit the scope :)

PS: cool to see other people with complex sequential model calls :)

[bennoschoenstein]

Thanks @gilrrei - that's a much cleaner solution than what I proposed, and #237 is exactly the mechanism I missed. I'll implement this pattern in my project repo. Closing this for now.

[mayrmt]

Is this documented in a tutorial, yet?

[bennoschoenstein]

I will make a tutorial out of this, to showcase how it can be done.

[sbrandstaeter]

If I understand this correctly, one limitation of this approach remains: the entire job runs on a single scheduler allocation, which means that the resources allocated are determined by the sub-step with the highest resource demand. Consequently, during the execution of other sub-steps that require fewer resources, the remaining resources (in most cases, CPU cores) would remain idle.