Merge pull request #1065 from opencobra/update/multiprocessing

refactor: multiprocessing.Pool usage

Author: synchon

src/cobra/flux_analysis/deletion.py

@@ -1,10 +1,10 @@
 # -*- coding: utf-8 -*-
 
 import logging
-import multiprocessing
 from builtins import dict, map
 from functools import partial
 from itertools import product
+from multiprocessing import Pool
 from typing import List, Set, Union
 
 import pandas as pd
@@ -149,14 +149,11 @@ def extract_knockout_results(result_iter):
                 gene=_gene_deletion_worker, reaction=_reaction_deletion_worker
             )[entity]
             chunk_size = len(args) // processes
-            pool = multiprocessing.Pool(
-                processes, initializer=_init_worker, initargs=(model,)
-            )
-            results = extract_knockout_results(
-                pool.imap_unordered(worker, args, chunksize=chunk_size)
-            )
-            pool.close()
-            pool.join()
+
+            with Pool(processes, initializer=_init_worker, initargs=(model,)) as pool:
+                results = extract_knockout_results(
+                    pool.imap_unordered(worker, args, chunksize=chunk_size)
+                )
         else:
             worker = dict(gene=_gene_deletion, reaction=_reaction_deletion)[entity]
             results = extract_knockout_results(map(partial(worker, model), args))

src/cobra/flux_analysis/variability.py

@@ -3,8 +3,8 @@
 from __future__ import absolute_import
 
 import logging
-import multiprocessing
 from builtins import map
+from multiprocessing import Pool
 from warnings import warn
 
 from numpy import zeros
@@ -208,17 +208,16 @@ def flux_variability_analysis(
                 # objective direction for all reactions. This creates a
                 # slight overhead but seems the most clean.
                 chunk_size = len(reaction_ids) // processes
-                pool = multiprocessing.Pool(
+
+                with Pool(
                     processes,
                     initializer=_init_worker,
                     initargs=(model, loopless, what[:3]),
-                )
-                for rxn_id, value in pool.imap_unordered(
-                    _fva_step, reaction_ids, chunksize=chunk_size
-                ):
-                    fva_result.at[rxn_id, what] = value
-                pool.close()
-                pool.join()
+                ) as pool:
+                    for rxn_id, value in pool.imap_unordered(
+                        _fva_step, reaction_ids, chunksize=chunk_size
+                    ):
+                        fva_result.at[rxn_id, what] = value
             else:
                 _init_worker(model, loopless, what[:3])
                 for rxn_id, value in map(_fva_step, reaction_ids):

src/cobra/sampling/optgp.py

@@ -214,12 +214,8 @@ def sample(self, n: int, fluxes: bool = True) -> pd.DataFrame:
             # limit errors, something weird going on with multiprocessing
             args = list(zip([n_process] * self.processes, range(self.processes)))
 
-            # No with statement or starmap here since Python 2.x
-            # does not support it :(
-            mp = Pool(self.processes, initializer=mp_init, initargs=(self,))
-            results = mp.map(_sample_chain, args, chunksize=1)
-            mp.close()
-            mp.join()
+            with Pool(self.processes, initializer=mp_init, initargs=(self,)) as pool:
+                results = pool.map(_sample_chain, args, chunksize=1)
 
             chains = np.vstack([r[1] for r in results])
             self.retries += sum(r[0] for r in results)