Add helper function to compute concurrency params

Author: dimkarakostas

tests/test_helper.py

@@ -1,4 +1,5 @@
 import tokenomics_decentralization.helper as hlp
+from collections import namedtuple
 import pathlib
 import os
 import datetime
@@ -485,3 +486,28 @@ def test_get_clusters(mocker):
     assert clusters['entity1'] == clusters['entity3']
     assert clusters['entity4'] == clusters['entity5']
     assert 'entity7' not in clusters.keys()
+
+
+def test_get_concurrency_per_ledger(mocker):
+    psutil_memory_mock = mocker.patch('psutil.virtual_memory')
+    psutil_memory_mock.return_value = namedtuple('VM', 'total')(10*10**9)
+
+    get_input_directories_mock = mocker.patch('tokenomics_decentralization.helper.get_input_directories')
+    get_input_directories_mock.return_value = [pathlib.Path('/').resolve()]
+
+    get_ledgers_mock = mocker.patch('tokenomics_decentralization.helper.get_ledgers')
+    get_ledgers_mock.return_value = ['bitcoin', 'ethereum']
+
+    os_walk_mock = mocker.patch('os.walk')
+    os_walk_mock.return_value = [('/', 'foo', ['bitcoin_2010-01-01_raw_data.csv'])]
+
+    os_stat_mock = mocker.patch('os.stat')
+    os_stat_mock.return_value = namedtuple('ST', 'st_size')(10*10**8)
+
+    concurrency = hlp.get_concurrency_per_ledger()
+    assert concurrency == {'bitcoin': 3, 'ethereum': 1}
+
+    os_stat_mock.return_value = namedtuple('ST', 'st_size')(5*10**9)
+
+    with pytest.raises(ValueError):
+        hlp.get_concurrency_per_ledger()

tokenomics_decentralization/helper.py

@@ -6,6 +6,7 @@
 import os
 import datetime
 import calendar
+import psutil
 import json
 from collections import defaultdict
 import logging
@@ -606,3 +607,41 @@ def get_clusters(ledger):
                 cluster_mapping[item[0]] = cluster_name
 
     return cluster_mapping
+
+
+def get_concurrency_per_ledger():
+    """
+    Computes the maximum number of parallel processes that can run per ledger,
+    based on the system's available memory.
+    :returns: a dictionary where the keys are ledger names and values are integers
+    """
+    system_memory_total = psutil.virtual_memory().total  # Get the system's total memory
+    system_memory_total -= 10**9  # Leave 1GB of memory to be used by other processes
+
+    concurrency = {}
+    too_large_ledgers = set()
+    input_dirs = get_input_directories()
+    for ledger in get_ledgers():
+        # Find the size of the largest input file per ledger
+        max_file_size = 0
+        for input_dir in input_dirs:
+            for folder, _, files in os.walk(input_dir):
+                for file in files:
+                    if file.startswith(ledger):
+                        max_file_size = max(max_file_size, os.stat(os.path.join(folder, file)).st_size)
+        # Compute the max number of processes that can open the largest ledger file
+        # and run in parallel without exhausting the system's memory.
+        if max_file_size > 0:
+            # When loaded in (a dict in) memory, each file consumes approx. 2.5x space compared to storage.
+            concurrency[ledger] = int(system_memory_total / (2.5 * max_file_size))
+            # Find if some ledger files are too large to fit in the system's available memory.
+            if concurrency[ledger] == 0:
+                too_large_ledgers.add(ledger)
+        else:
+            concurrency[ledger] = 1
+
+    if too_large_ledgers:
+        raise ValueError('The max input files of the following ledgers are too'
+                         'large to load in memory' + ','.join(too_large_ledgers))
+
+    return concurrency