HPC Job Executor

HPC Job Executor#

In contrast to the HPC Cluster Executor which submits individual Python functions to HPC job schedulers, the HPC Job Executors take a given job allocation of the HPC job scheduler and executes Python functions with the resources available in this job allocation. In this regard it is similar to the Single Node Executor as it communicates with the individual Python processes using the zero message queue, still it is more advanced as it can access the computational resources of all compute nodes of the given HPC job allocation and also provides the option to assign GPUs as accelerators for parallel execution.

Available Functionality:

Submit Python functions with the submit() function or the map() function.
Support for parallel execution, either using the message passing interface (MPI), thread based parallelism or by assigning dedicated GPUs to selected Python functions. All these resources assignments are handled via the resource dictionary parameter resource_dict.
Performance optimization features, like block allocation, dependency resolution and caching.

The only parameter the user has to change is the backend parameter.

SLURM#

With the Simple Linux Utility for Resource Management (SLURM) currently being the most commonly used job scheduler, executorlib provides an interface to submit Python functions to SLURM. Internally, this is based on the srun command of the SLURM scheduler, which creates job steps in a given allocation. Given that all resource requests in SLURM are communicated via a central database a large number of submitted Python functions and resulting job steps can slow down the performance of SLURM. To address this limitation it is recommended to install the hierarchical job scheduler flux in addition to SLURM, to use flux for distributing the resources within a given allocation. This configuration is discussed in more detail below in the section SLURM with flux.

from executorlib import SlurmJobExecutor

with SlurmJobExecutor() as exe:
    future = exe.submit(sum, [1, 1])
    print(future.result())

SLURM with Flux#

As discussed in the installation section it is important to select the flux version compatible to the installation of a given HPC cluster. Which GPUs are available? Who manufactured these GPUs? Does the HPC use mpich or OpenMPI or one of their commercial counter parts like cray MPI or intel MPI? Depending on the configuration different installation options can be choosen, as explained in the installation section.

Afterwards flux can be started in an sbatch submission script using:

srun flux start python <script.py>

In this Python script <script.py> the "flux_allocation" backend can be used.

When the Python functions are submitted from the login node with the SlurmClusterExecutor, a flux instance can also be started automatically for each submitted job, rather than wrapping the script manually. This is achieved by wrapping the {{command}} of the submission template with srun flux start, so each sbatch job boots its own flux instance:

submission_template = """\
#!/bin/bash
#SBATCH --output=time.out
#SBATCH --job-name={{job_name}}
#SBATCH --chdir={{working_directory}}
#SBATCH --get-user-env=L
#SBATCH --partition={{partition}}
{%- if run_time_max %}
#SBATCH --time={{ [1, run_time_max // 60]|max }}
{%- endif %}
#SBATCH --cpus-per-task={{cores}}

srun --cpus-per-task={{cores}} flux start {{command}}
"""

Inside the function which is submitted to the SlurmClusterExecutor a nested FluxJobExecutor can then distribute many small tasks within the flux instance of the allocation. This combines the robustness of the file based submission of the HPC Cluster Executor with the low overhead of the socket based execution of the HPC Job Executor, and is the recommended setup to distribute a large number of tasks within a single SLURM allocation - see the installation section.

Resource Assignment#

Independent of the selected Executor Single Node Executor, HPC Cluster Executor or HPC job executor the assignment of the computational resources remains the same. They can either be specified in the submit() function by adding the resource dictionary parameter resource_dict or alternatively during the initialization of the Executor class by adding the resource dictionary parameter resource_dict there.

This functionality of executorlib is commonly used to rewrite individual Python functions to use MPI while the rest of the Python program remains serial.

def calc_mpi(i):
    from mpi4py import MPI

    size = MPI.COMM_WORLD.Get_size()
    rank = MPI.COMM_WORLD.Get_rank()
    return i, size, rank

Depending on the choice of MPI version, it is recommended to specify the pmi standard which flux should use internally for the resource assignment. For example for OpenMPI >=5 "pmix" is the recommended pmi standard.

from executorlib import FluxJobExecutor

with FluxJobExecutor(pmi_mode="pmix") as exe:
    fs = exe.submit(calc_mpi, 3, resource_dict={"cores": 2})
    print(fs.result())

[(3, 2, 0), (3, 2, 1)]

Block Allocation#

The block allocation for the HPC allocation mode follows the same implementation as the block allocation for the Single Node Executor. It starts by defining the initialization function init_function() which returns a dictionary which is internally used to look up input parameters for Python functions submitted to the FluxJobExecutor class. Commonly this functionality is used to store large data objects inside the Python process created for the block allocation, rather than reloading these Python objects for each submitted function.

def init_function():
    return {"j": 4, "k": 3, "l": 2}

def calc_with_preload(i, j, k):
    return i + j + k

with FluxJobExecutor(
    pmi_mode="pmix",
    max_workers=2,
    init_function=init_function,
    block_allocation=True,
) as exe:
    fs = exe.submit(calc_with_preload, 2, j=5)
    print(fs.result())

In this example the parameter k is used from the dataset created by the initialization function while the parameters i and j are specified by the call of the submit() function.

When using the block allocation mode, it is recommended to set either the maxium number of workers using the max_workers parameter or the maximum number of CPU cores using the max_cores parameter to prevent oversubscribing the available resources.

Dependencies#

Python functions with rather different computational resource requirements should not be merged into a single function. So to able to execute a series of Python functions which each depend on the output of the previous Python function executorlib internally handles the dependencies based on the concurrent futures future objects from the Python standard library. This implementation is independent of the selected backend and works for HPC allocation mode just like explained in the Single Node Executor section.

def add_funct(a, b):
    return a + b

with FluxJobExecutor(pmi_mode="pmix") as exe:
    future = 0
    for i in range(1, 4):
        future = exe.submit(add_funct, i, future)
    print(future.result())

Caching#

Finally, also the caching is available for HPC allocation mode, in analogy to the Single Node Executor. Again this functionality is not designed to identify function calls with the same parameters, but rather provides the option to reload previously cached results even after the Python processes which contained the executorlib Executor class is closed. As the cache is stored on the file system, this option can decrease the performance of executorlib. Consequently the caching option should primarily be used during the prototyping phase.

with FluxJobExecutor(pmi_mode="pmix", cache_directory="./file") as exe:
    future_lst = [exe.submit(sum, [i, i]) for i in range(1, 4)]
    print([f.result() for f in future_lst])

[2, 4, 6]

import os
import shutil

cache_dir = "./file"
if os.path.exists(cache_dir):
    print(os.listdir(cache_dir))
    try:
        shutil.rmtree(cache_dir)
    except OSError:
        pass

['sum89afbdf9da5eb1794f6976a3f01697c2_o.h5', 'sum0f7710227cda6456e5d07187702313f3_o.h5', 'sumf5ad27b855231a293ddd735a8554c9ea_o.h5']

Nested executors#

The hierarchical nature of the flux job scheduler allows the creation of additional executorlib Executors inside the functions submitted to the Executor. This hierarchy can be beneficial to separate the logic to saturate the available computational resources.

def calc_nested():
    from executorlib import FluxJobExecutor

    with FluxJobExecutor(pmi_mode="pmix") as exe:
        fs = exe.submit(sum, [1, 1])
        return fs.result()

with FluxJobExecutor(pmi_mode="pmix", flux_executor_nesting=True) as exe:
    fs = exe.submit(calc_nested)
    print(fs.result())

Executor from Flux#

The flux framework provides its own FluxExecutor which can be used to submit shell scripts to the flux framework for execution. The FluxExecutor returns its own representation of future objects which is incompatible with the concurrent.futures.Future which is used by executorlib. Combining both provides the opportunity to link Python fucntions and external executables. For this purpose executorlib provides the option to use a FluxExecutor as an input for the FluxJobExecutor:

from executorlib import FluxJobExecutor
import flux.job

with flux.job.FluxExecutor() as flux_executor:
    with FluxJobExecutor(flux_executor=flux_executor) as exe:
        future = exe.submit(sum, [1, 1])
        print(future.result())

Resource Monitoring#

For debugging it is commonly helpful to keep track of the computational resources. flux provides a number of features to analyse the resource utilization, so here only the two most commonly used ones are introduced. Starting with the option to list all the resources available in a given allocation with the flux resource list command:

! flux resource list

     STATE NNODES   NCORES    NGPUS NODELIST
      free      1       24        0 jupyter-pyiron-executorlib-wx8wv67z
 allocated      0        0        0 
      down      0        0        0 

Followed by the list of jobs which were executed in a given flux session. This can be retrieved using the flux jobs -a command:

! flux jobs -a

       JOBID USER     NAME       ST NTASKS NNODES     TIME INFO
    ƒ66TjsQs jovyan   python     CD      1      1   0.149s jupyter-pyiron-executorlib-wx8wv67z
    ƒ4R3m4Sj jovyan   flux       CD      1      1   3.509s jupyter-pyiron-executorlib-wx8wv67z
    ƒ3N4Qc3y jovyan   python     CD      1      1   1.922s jupyter-pyiron-executorlib-wx8wv67z
    ƒ3DuUZ9y jovyan   python     CD      1      1   2.291s jupyter-pyiron-executorlib-wx8wv67z
    ƒ3DrWabH jovyan   python     CD      1      1   2.204s jupyter-pyiron-executorlib-wx8wv67z
    ƒ2z9sDYT jovyan   python     CD      1      1   0.271s jupyter-pyiron-executorlib-wx8wv67z
    ƒ2m9FX6w jovyan   python     CD      1      1   0.404s jupyter-pyiron-executorlib-wx8wv67z
    ƒ2dGdLJj jovyan   python     CD      1      1   0.346s jupyter-pyiron-executorlib-wx8wv67z
    ƒ29qrcvj jovyan   python     CD      1      1   0.848s jupyter-pyiron-executorlib-wx8wv67z
    ƒ29tpbVR jovyan   python     CD      1      1   0.539s jupyter-pyiron-executorlib-wx8wv67z
     ƒZsZ5QT jovyan   python     CD      2      1   0.966s jupyter-pyiron-executorlib-wx8wv67z

Flux#

While the number of HPC clusters which use flux as primary job scheduler is currently still limited the setup and functionality provided by executorlib for running SLURM with flux also applies to HPCs which use flux as primary job scheduler.