Use parallel processes when analyzing ledger snapshots

Author: dimkarakostas

tests/test_analyze.py

@@ -1,5 +1,5 @@
-from tokenomics_decentralization.analyze import analyze_snapshot, analyze, get_entries
-from unittest.mock import call
+from tokenomics_decentralization.analyze import analyze_snapshot, analyze, get_entries, analyze_ledger_snapshot
+from unittest.mock import call, Mock
 import pathlib
 
 
@@ -129,20 +129,25 @@ def test_get_entries(mocker):
 
 
 def test_analyze(mocker):
+    get_concurrency_mock = mocker.patch('tokenomics_decentralization.helper.get_concurrency_per_ledger')
+    get_concurrency_mock.return_value = {'bitcoin': 2, 'ethereum': 2}
+
+    write_csv_output_mock = mocker.patch('tokenomics_decentralization.helper.write_csv_output')
+
+    analyze(['bitcoin'], ['2010-01-01'])
+    assert len(write_csv_output_mock.call_args_list) == 1
+
+
+def test_analyze_ledger_snapshot(mocker):
     get_input_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_input_directories')
     get_input_directories_mock.return_value = [pathlib.Path('/').resolve()]
 
     is_file_mock = mocker.patch('os.path.isfile')
     is_file_mock.side_effect = {
         pathlib.Path('/bitcoin_2010-01-01_raw_data.csv').resolve(): True,
-        pathlib.Path('/bitcoin_2011-01-01_raw_data.csv').resolve(): False,
         pathlib.Path('/ethereum_2010-01-01_raw_data.csv').resolve(): False,
-        pathlib.Path('/ethereum_2011-01-01_raw_data.csv').resolve(): True,
     }.get
 
-    get_db_connector_mock = mocker.patch('tokenomics_decentralization.db_helper.get_connector')
-    get_db_connector_mock.return_value = 'connector'
-
     get_entries_mock = mocker.patch('tokenomics_decentralization.analyze.get_entries')
     entries = [1, 2]
     get_entries_mock.return_value = entries
@@ -153,36 +158,26 @@ def test_analyze(mocker):
     get_output_row_mock = mocker.patch('tokenomics_decentralization.helper.get_output_row')
     get_output_row_mock.return_value = 'row'
 
-    write_csv_output_mock = mocker.patch('tokenomics_decentralization.helper.write_csv_output')
+    sema = Mock()
 
-    get_input_dirs_calls = []
     get_entries_calls = []
     analyze_snapshot_calls = []
     get_row_calls = []
-    write_output_calls = []
+    sema_release_calls = []
 
-    analyze(['bitcoin'], ['2010-01-01'])
-    get_input_dirs_calls.append(call())
-    assert get_input_directories_mock.call_args_list == get_input_dirs_calls
+    analyze_ledger_snapshot('bitcoin', '2010-01-01', [], sema)
     get_entries_calls.append(call('bitcoin', '2010-01-01', pathlib.Path('/bitcoin_2010-01-01_raw_data.csv').resolve()))
     assert get_entries_mock.call_args_list == get_entries_calls
     analyze_snapshot_calls.append(call(entries))
     assert analyze_snapshot_mock.call_args_list == analyze_snapshot_calls
     get_row_calls.append(call('bitcoin', '2010-01-01', {'hhi': 1}))
     assert get_output_row_mock.call_args_list == get_row_calls
-    write_output_calls.append(call(['row']))
-    assert write_csv_output_mock.call_args_list == write_output_calls
+    sema_release_calls.append(call())
+    assert sema.release.call_args_list == sema_release_calls
 
-    analyze(['bitcoin', 'ethereum'], ['2010-01-01', '2011-01-01'])
-    get_input_dirs_calls += 4 * [call()]
-    assert get_input_directories_mock.call_args_list == get_input_dirs_calls
-    get_entries_calls.append(call('bitcoin', '2010-01-01', pathlib.Path('/bitcoin_2010-01-01_raw_data.csv').resolve()))
-    get_entries_calls.append(call('ethereum', '2011-01-01', pathlib.Path('/ethereum_2011-01-01_raw_data.csv').resolve()))
+    analyze_ledger_snapshot('ethereum', '2010-01-01', [], sema)
     assert get_entries_mock.call_args_list == get_entries_calls
-    analyze_snapshot_calls += 2 * [call(entries)]
     assert analyze_snapshot_mock.call_args_list == analyze_snapshot_calls
-    get_row_calls.append(call('bitcoin', '2010-01-01', {'hhi': 1}))
-    get_row_calls.append(call('ethereum', '2011-01-01', {'hhi': 1}))
     assert get_output_row_mock.call_args_list == get_row_calls
-    write_output_calls.append(call(['row', 'row']))
-    assert write_csv_output_mock.call_args_list == write_output_calls
+    sema_release_calls.append(call())  # Test that semaphore release is called even if file does not exist
+    assert sema.release.call_args_list == sema_release_calls

tokenomics_decentralization/analyze.py

@@ -1,4 +1,5 @@
 import csv
+import multiprocessing
 import os.path
 import tokenomics_decentralization.helper as hlp
 import tokenomics_decentralization.db_helper as db_hlp
@@ -113,35 +114,53 @@ def get_entries(ledger, date, filename):
     return entries
 
 
+def analyze_ledger_snapshot(ledger, date, output_rows, sema):
+    """
+    Executes the analysis of a given ledgers and snapshot date.
+    :param ledger: a ledger name
+    :param date: a string in YYYY-MM-DD format
+    :param output_rows: a list of strings in the form of csv output rows
+    :param sema: a multiprocessing semaphore
+    """
+    input_filename = None
+    input_paths = [input_dir / f'{ledger}_{date}_raw_data.csv' for input_dir in hlp.get_input_directories()]
+    for filename in input_paths:
+        if os.path.isfile(filename):
+            input_filename = filename
+            break
+    if input_filename:
+        logging.info(f'[*] {ledger} - {date}')
+
+        entries = get_entries(ledger, date, filename)
+        metrics_values = analyze_snapshot(entries)
+        del entries
+
+        row = hlp.get_output_row(ledger, date, metrics_values)
+        output_rows.append(row)
+
+    sema.release()  # Release the semaphore s.t. the loop in analyze() can continue
+
+
 def analyze(ledgers, snapshot_dates):
     """
     Executes the analysis of the given ledgers for the snapshot dates and writes the output
     to csv files.
     :param ledgers: a list of ledger names
     :param snapshot_dates: a list of strings in YYYY-MM-DD format
     """
-    output_rows = []
+    manager = multiprocessing.Manager()
+    output_rows = manager.list()  # output_rows is a shared list across all parallel processes
+
+    concurrency = hlp.get_concurrency_per_ledger()
     for ledger in ledgers:
-        logging.info(f'[*] {ledger} - Analyzing')
+        sema = multiprocessing.Semaphore(concurrency[ledger])
+        jobs = []
         for date in snapshot_dates:
-            logging.info(f'[*] {ledger} - {date}')
-
-            input_filename = None
-            input_paths = [input_dir / f'{ledger}_{date}_raw_data.csv' for input_dir in hlp.get_input_directories()]
-            for filename in input_paths:
-                if os.path.isfile(filename):
-                    input_filename = filename
-                    break
-            if not input_filename:
-                logging.error(f'{ledger} input data for {date} do not exist')
-                continue
-
-            entries = get_entries(ledger, date, filename)
-            metrics_values = analyze_snapshot(entries)
-            del entries
-
-            output_rows.append(hlp.get_output_row(ledger, date, metrics_values))
-            for metric, value in metrics_values.items():
-                logging.info(f'{metric}: {value}')
-
-    hlp.write_csv_output(output_rows)
+            sema.acquire()  # Loop blocks here while the active processes are as many as the semaphore's limit
+            p = multiprocessing.Process(target=analyze_ledger_snapshot, args=(ledger, date, output_rows, sema))
+            jobs.append(p)
+            p.start()
+        for proc in jobs:
+            proc.join()
+
+    hlp.write_csv_output(sorted(output_rows, key=lambda x: (x[0], x[1])))  # Csv rows ordered by ledger and date