Merge pull request #204 from anilbey/n_processes

Introducing number of processes argument

Author: wvangeit

bluepymm/main.py

@@ -51,13 +51,16 @@ def run(arg_list):
 
     if args.action == "prepare":
         prepare_combos.prepare_combos(conf_filename=args.conf_filename,
-                                      continu=args.continu)
+                                      continu=args.continu,
+                                      n_processes=args.n_processes)
     elif args.action == "run":
         run_combos.run_combos(conf_filename=args.conf_filename,
                               ipyp=args.ipyp,
-                              ipyp_profile=args.ipyp_profile)
+                              ipyp_profile=args.ipyp_profile,
+                              n_processes=args.n_processes)
     elif args.action == "select":
-        select_combos.select_combos(conf_filename=args.conf_filename)
+        select_combos.select_combos(conf_filename=args.conf_filename,
+                                    n_processes=args.n_processes)
     elif args.action == "validate":
         validate_output.validate_output(conf_filename=args.conf_filename)
 

bluepymm/prepare_combos/main.py

@@ -29,7 +29,7 @@
 from . import create_mm_sqlite
 
 
-def prepare_emodels(conf_dict, continu, scores_db_path):
+def prepare_emodels(conf_dict, continu, scores_db_path, n_processes):
     """Prepare emodels"""
 
     tmp_dir = conf_dict['tmp_dir']
@@ -52,7 +52,7 @@ def prepare_emodels(conf_dict, continu, scores_db_path):
     # Clone the emodels repo and prepare the dirs for all the emodels
     emodel_dirs = prepare_dirs.prepare_emodel_dirs(
         final_dict, emodel_etype_map, emodels_dir, opt_dir, emodels_hoc_dir,
-        emodels_in_repo, continu=continu)
+        emodels_in_repo, continu=continu, n_processes=n_processes)
 
     if not continu:
         print('Creating sqlite db at %s' % scores_db_path)
@@ -97,15 +97,15 @@ def prepare_emodels(conf_dict, continu, scores_db_path):
     return final_dict, emodel_dirs
 
 
-def prepare_combos(conf_filename, continu):
+def prepare_combos(conf_filename, continu, n_processes=None):
     """Prepare combos"""
 
     print('Reading configuration at %s' % conf_filename)
     conf_dict = tools.load_json(conf_filename)
     scores_db_path = os.path.abspath(conf_dict['scores_db'])
 
     final_dict, emodel_dirs = prepare_emodels(
-        conf_dict, continu, scores_db_path)
+        conf_dict, continu, scores_db_path, n_processes)
 
     # Save output
     # TODO: gather all output business here?
@@ -124,3 +124,5 @@ def add_parser(action):
     parser.add_argument('conf_filename')
     parser.add_argument('--continu', action='store_true',
                         help='continue from previous run')
+    parser.add_argument('--n_processes', help='number of processes',
+                        type=int)

bluepymm/prepare_combos/prepare_emodel_dirs.py

@@ -255,13 +255,15 @@ def prepare_emodel_dir(input_args):
 
 
 def prepare_emodel_dirs(
-        final_dict,
-        emodel_etype_map,
-        emodels_dir,
-        opt_dir,
-        emodels_hoc_dir,
-        emodels_in_repo,
-        continu=False):
+    final_dict,
+    emodel_etype_map,
+    emodels_dir,
+    opt_dir,
+    emodels_hoc_dir,
+    emodels_in_repo,
+    continu=False,
+    n_processes=None,
+):
     """Prepare the directories for the emodels.
 
     Args:
@@ -278,6 +280,8 @@ def prepare_emodel_dirs(
             into separate subdirectories.
         continu: True if this BluePyMM run builds on a previous run, False
             otherwise. Default is False.
+        n_processes: the integer number of processes. If `None`,
+        all processes are going to be used.
 
     Return:
         A dict mapping e-models to prepared e-model directories.
@@ -299,10 +303,23 @@ def prepare_emodel_dirs(
              emodels_in_repo,
              continu))
 
-    print('Parallelising preparation of e-model directories')
-    pool = multiprocessing.Pool(maxtasksperchild=1)
     emodel_dirs = {}
-    for emodel_dir_dict in pool.map(prepare_emodel_dir, arg_list, chunksize=1):
+    emodel_dir_dicts = []
+
+    if n_processes == 1:
+        for arg in arg_list:
+            emodel_dir_dict = prepare_emodel_dir(arg)
+            emodel_dir_dicts.append(emodel_dir_dict)
+    else:
+        print('Parallelising preparation of e-model directories')
+        pool = multiprocessing.Pool(processes=n_processes,
+                                    maxtasksperchild=1)
+        for emodel_dir_dict in pool.map(prepare_emodel_dir, arg_list,
+                                        chunksize=1):
+            emodel_dir_dicts.append(emodel_dir_dict)
+
+    for emodel_dir_dict in emodel_dir_dicts:
         for emodel, emodel_dir in emodel_dir_dict.items():
             emodel_dirs[emodel] = emodel_dir
+
     return emodel_dirs

bluepymm/run_combos/calculate_scores.py

@@ -321,7 +321,7 @@ def expand_scores_to_score_values_table(scores_sqlite_filename):
 
 def calculate_scores(final_dict, emodel_dirs, scores_db_filename,
                      use_ipyp=False, ipyp_profile=None, timeout=10,
-                     use_apical_points=True):
+                     use_apical_points=True, n_processes=None):
     """Calculate scores of e-model morphology combinations and update the
     database accordingly.
 
@@ -335,6 +335,8 @@ def calculate_scores(final_dict, emodel_dirs, scores_db_filename,
             False.
         ipyp_profile: path to ipyparallel profile. Default is None.
         use_apical_points: boolean to use apical points or not
+        n_processes: the integer number of processes. If `None`,
+        all processes are going to be used.
     """
 
     print('Creating argument list for parallelisation')
@@ -344,16 +346,15 @@ def calculate_scores(final_dict, emodel_dirs, scores_db_filename,
                                use_apical_points=use_apical_points)
 
     print('Parallelising score evaluation of %d me-combos' % len(arg_list))
-
     if use_ipyp:
         # use ipyparallel
         client = ipyparallel.Client(profile=ipyp_profile, timeout=timeout)
-        lview = client.load_balanced_view()
+        lview = client.load_balanced_view(targets=n_processes)
         results = lview.imap(run_emodel_morph_isolated,
                              arg_list, ordered=False)
     else:
         # use multiprocessing
-        pool = tools.NestedPool()
+        pool = tools.NestedPool(processes=n_processes)
         results = pool.imap_unordered(run_emodel_morph_isolated, arg_list)
 
     # keep track of the number of received results

bluepymm/run_combos/main.py

@@ -41,9 +41,12 @@ def add_parser(action):
                         help='Path to ipyparallel profile')
     parser.add_argument('--timeout',
                         help='Timeout for ipyparallel clients')
+    parser.add_argument('--n_processes', help='number of processes',
+                        type=int)
 
 
-def run_combos_from_conf(conf_dict, ipyp=None, ipyp_profile=None, timeout=10):
+def run_combos_from_conf(conf_dict, ipyp=None, ipyp_profile=None, timeout=10,
+                         n_processes=None):
     """Run combos from conf dictionary"""
     output_dir = conf_dict['output_dir']
     final_dict = tools.load_json(
@@ -69,13 +72,15 @@ def run_combos_from_conf(conf_dict, ipyp=None, ipyp_profile=None, timeout=10):
         use_ipyp=ipyp,
         ipyp_profile=ipyp_profile,
         timeout=timeout,
-        use_apical_points=use_apical_points)
+        use_apical_points=use_apical_points,
+        n_processes=n_processes)
 
 
-def run_combos(conf_filename, ipyp=None, ipyp_profile=None):
+def run_combos(conf_filename, ipyp=None, ipyp_profile=None, n_processes=None):
     """Run combos"""
 
     print('Reading configuration at %s' % conf_filename)
     conf_dict = tools.load_json(conf_filename)
 
-    run_combos_from_conf(conf_dict, ipyp, ipyp_profile)
+    run_combos_from_conf(conf_dict, ipyp, ipyp_profile,
+                         n_processes=n_processes)

bluepymm/select_combos/main.py

@@ -31,20 +31,21 @@
 from . import process_megate_config as proc_config
 
 
-def select_combos(conf_filename):
+def select_combos(conf_filename, n_processes):
     """Parse conf file and run select combos"""
     # Parse configuration file
     conf_dict = tools.load_json(conf_filename)
 
-    select_combos_from_conf(conf_dict)
+    select_combos_from_conf(conf_dict, n_processes)
 
 
-def select_combos_from_conf(conf_dict):
+def select_combos_from_conf(conf_dict, n_processes=None):
     """Compare scores of me-combinations to thresholds, select successful
     combinations, and write results out to file.
 
     Args:
         conf_filename: filename of configuration (.json file)
+        n_processes: integer number of processes, `None` will use all of them
     """
     scores_db_filename = conf_dict['scores_db']
     pdf_filename = conf_dict['pdf_filename']
@@ -85,7 +86,8 @@ def select_combos_from_conf(conf_dict):
             scores, score_values,
             conf_dict.get('plot_emodels_per_morphology', False),
             output_dir,
-            select_perc_best)
+            select_perc_best,
+            n_processes=n_processes)
     print('Wrote pdf to %s' % os.path.abspath(pdf_filename))
 
     # write output files
@@ -112,3 +114,5 @@ def add_parser(action):
     parser = action.add_parser('select',
                                help='Select feasible me-combinations')
     parser.add_argument('conf_filename')
+    parser.add_argument('--n_processes', help='number of processes',
+                        type=int)

bluepymm/select_combos/reporting.py

@@ -330,7 +330,8 @@ def create_final_db_and_write_report(pdf_filename,
                                      score_values,
                                      enable_plot_emodels_per_morphology,
                                      output_dir,
-                                     select_perc_best):
+                                     select_perc_best,
+                                     n_processes=None):
     """Create the final output files and report"""
     ext_neurondb = pandas.DataFrame()
 
@@ -362,7 +363,9 @@ def create_final_db_and_write_report(pdf_filename,
                                                         megate_patterns,
                                                         skip_repaired_exemplar,
                                                         check_opt_scores,
-                                                        select_perc_best)
+                                                        select_perc_best,
+                                                        n_processes=n_processes
+                                                        )
 
         print("All emodels processed, generating output files")
 

bluepymm/select_combos/table_processing.py

@@ -269,7 +269,8 @@ def process_emodels(emodels,
                     megate_patterns,
                     skip_repaired_exemplar,
                     enable_check_opt_scores,
-                    select_perc_best):
+                    select_perc_best,
+                    n_processes=None):
 
     arg_list = [(emodel,
                  scores,
@@ -280,16 +281,22 @@ def process_emodels(emodels,
                  enable_check_opt_scores,
                  select_perc_best) for emodel in emodels]
 
-    print('Parallelising selection processing of e-models')
-    pool = multiprocessing.Pool(maxtasksperchild=1)
     emodel_infos = {}
-    for emodel, emodel_info in pool.imap(process_emodel, arg_list,
-                                         chunksize=1):
-        print('Received processed info from e-model %s' % emodel)
-        emodel_infos[emodel] = emodel_info
 
-    pool.terminate()
-    pool.join()
+    if n_processes == 1:
+        for args in arg_list:
+            emodel, emodel_info = process_emodel(args)
+            emodel_infos[emodel] = emodel_info
+    else:
+        print('Parallelising selection processing of e-models')
+        pool = multiprocessing.Pool(maxtasksperchild=1, processes=n_processes)
+        for emodel, emodel_info in pool.imap(process_emodel, arg_list,
+                                             chunksize=1):
+            print('Received processed info from e-model %s' % emodel)
+            emodel_infos[emodel] = emodel_info
+
+        pool.terminate()
+        pool.join()
 
     return emodel_infos
 

bluepymm/tests/test_calculate_scores.py

@@ -383,7 +383,8 @@ def test_calculate_scores():
                 final_dict,
                 emodel_dirs,
                 test_db_filename,
-                use_ipyp=use_ipyp
+                use_ipyp=use_ipyp,
+                n_processes=1
             )
 
         if use_ipyp:

bluepymm/tests/test_prepare_emodel_dirs.py

@@ -257,9 +257,65 @@ def test_prepare_emodel_dir():
 
 
 @attr('unit')
-def test_prepare_emodel_dirs():
-    """prepare_combos.prepare_emodel_dirs: test prepare_emodel_dirs
-    based on test example 'simple1'.
+def test_prepare_emodel_dirs_single_process():
+    """prepare_combos.prepare_emodel_dirs_single_process:
+    test prepare_emodel_dirs based on test example 'simple1
+    using a single process'.
+    """
+    # create test directory, where output of this test will be created
+    test_dir = os.path.join(TMP_DIR, 'test_prepare_emodel_dirs')
+    tools.makedirs(test_dir)
+
+    # input parameters
+    final_dict = {'emodel1': {'main_path': '.',
+                              'seed': 2,
+                              'rank': 0,
+                              'notes': '',
+                              'branch': 'emodel1',
+                              'params': {'cm': 1.0},
+                              'fitness': {'Step1.SpikeCount': 20.0},
+                              'score': 104.72906197480131,
+                              'morph_path': 'morphologies/morph1.asc'},
+                  'emodel2': {'main_path': '.',
+                              'seed': 2,
+                              'rank': 0,
+                              'notes': '',
+                              'branch': 'emodel2',
+                              'params': {'cm': 0.5},
+                              'fitness': {'Step1.SpikeCount': 20.0},
+                              'score': 104.72906197480131,
+                              'morph_path': 'morphologies/morph2.asc'}}
+    emodel_etype_map = {'emodel1': {'mm_recipe': 'emodel1',
+                                    'etype': 'etype1',
+                                    'layer': ['1', 'str1']},
+                        'emodel2': {'mm_recipe': 'emodel2',
+                                    'etype': 'etype2',
+                                    'layer': ['1', '2']}}
+    emodels_dir = os.path.join(test_dir, 'tmp/emodels/')
+    opt_dir = os.path.join(TEST_DATA_DIR, 'data/emodels_dir/subdir/')
+    emodels_hoc_dir = os.path.join(test_dir, './output/emodels_hoc/')
+    emodels_in_repo = False
+    continu = False
+
+    # run function
+    with tools.cd(TEST_DATA_DIR):
+        ret = prepare_emodel_dirs.prepare_emodel_dirs(
+            final_dict, emodel_etype_map, emodels_dir, opt_dir,
+            emodels_hoc_dir, emodels_in_repo, continu, n_processes=1)
+
+    # verify output
+    expected_ret = {emodel: os.path.join(
+        emodels_dir, emodel) for emodel in final_dict}
+    nt.assert_dict_equal(ret, expected_ret)
+    nt.assert_true(os.path.isdir(emodels_dir))
+    nt.assert_true(os.path.isdir(emodels_hoc_dir))
+
+
+@attr('unit')
+def test_prepare_emodel_dirs_multi_process():
+    """prepare_combos.prepare_emodel_dirs_multi_process:
+    test prepare_emodel_dirs based on test example 'simple1'
+    using multiprocessing.
     """
     # create test directory, where output of this test will be created
     test_dir = os.path.join(TMP_DIR, 'test_prepare_emodel_dirs')
@@ -300,7 +356,7 @@ def test_prepare_emodel_dirs():
     with tools.cd(TEST_DATA_DIR):
         ret = prepare_emodel_dirs.prepare_emodel_dirs(
             final_dict, emodel_etype_map, emodels_dir, opt_dir,
-            emodels_hoc_dir, emodels_in_repo, continu)
+            emodels_hoc_dir, emodels_in_repo, continu, n_processes=None)
 
     # verify output
     expected_ret = {emodel: os.path.join(