From dc18f59d828f1b77a29a9121d1586a482d4954d6 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 22 Mar 2022 06:53:47 -0500 Subject: [PATCH] Allow user to override num processes used for depletion --- docs/source/pythonapi/deplete.rst | 6 ++++++ openmc/deplete/pool.py | 6 +++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/source/pythonapi/deplete.rst b/docs/source/pythonapi/deplete.rst index d4a0f128f..383268f8b 100644 --- a/docs/source/pythonapi/deplete.rst +++ b/docs/source/pythonapi/deplete.rst @@ -166,6 +166,12 @@ with :func:`cram.CRAM48` being the default. :type: bool +.. data:: pool.NUM_PROCESSES + + Number of worker processes used for depletion calculations, which rely on the + :class:`multiprocessing.pool.Pool` class. If set to ``None`` (default), the + number returned by :func:`os.cpu_count` is used. + The following classes are used to help the :class:`openmc.deplete.Operator` compute quantities like effective fission yields, reaction rates, and total system energy. diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py index 8360ef602..a3d37f79f 100644 --- a/openmc/deplete/pool.py +++ b/openmc/deplete/pool.py @@ -10,6 +10,10 @@ from multiprocessing import Pool # multiprocessing routines during depletion USE_MULTIPROCESSING = True +# Allow user to override the number of worker processes to use for depletion +# calculations +NUM_PROCESSES = None + def deplete(func, chain, x, rates, dt, matrix_func=None): """Deplete materials using given reaction rates for a specified time @@ -58,7 +62,7 @@ def deplete(func, chain, x, rates, dt, matrix_func=None): inputs = zip(matrices, x, repeat(dt)) if USE_MULTIPROCESSING: - with Pool() as pool: + with Pool(NUM_PROCESSES) as pool: x_result = list(pool.starmap(func, inputs)) else: x_result = list(starmap(func, inputs))