<jupyter_start><jupyter_text>Core> Core functions for MRL, mostly low level plumbing and parallel processing<jupyter_code>#hide
from nbdev.showdoc import *
%load_ext autoreload
%autoreload 2
# export
from mrl.imports import *
from multiprocessing import get_context
import requests
import zipfile<jupyter_output><empty_output><jupyter_text>Miscellaneous FunctionsLow level helper functions<jupyter_code># export
def is_container(x):
    "check if `x` is a container (used for parallel processing)"
    if isinstance(x, (list, tuple, np.ndarray)):
        return True
    else:
        return False

def flatten_recursive(list_of_lists):
    "Recursively flattel list of lists"
    flat_list = []
    for item in list_of_lists:
        if type(item) == list:
            flat_list += flatten_recursive(item)
        else:
            flat_list.append(item)
            
    return flat_list

def flatten_list_of_lists(list_of_lists):
    "Flattens list of lists (not recursive)"
    return [item for sublist in list_of_lists for item in sublist]

def deduplicate_list(l):
    "Deduplicates list l"
    return list(set(l))

def chunk_list(input_list, chunksize):
    'Breaks `input_list` into chunks of size `chunksize`, ragged on last list'
    return [input_list[i:i+chunksize] for i in range(0, len(input_list), chunksize)]

def filter_passing(inputs, bools):
    'Subsets `inputs` (list) by `bools` (list of bools)'
    assert len(inputs)==len(bools), '`inputs` and `bools` must have the same length'
    return [inputs[i] for i in range(len(inputs)) if bools[i]]
assert flatten_list_of_lists([[1],[2],[3]]) == [1,2,3]
assert flatten_recursive([[1],[2],[3, [4,5, [6,7,8]]]]) == [1,2,3,4,5,6,7,8]<jupyter_output><empty_output><jupyter_text>Parallel ProcessingMRL tries to build in parallel processing at every level. This can make a huge difference when you're processing millions of molecules`new_pool_parallel` and `maybe_parallel` are convenient wrappers for parallel processing. The given `func` is wrapped with `**kwargs` and used to process the `iterable`. If `iterable` is a `list` or `np.ndarray`, the elements in `iterable` are run in parallel by `func`. Parallel processing tradeoffsParallel processing can significantly speed up a process. There are however some trade-offs.In Python, parallel processing is creating using a `Pool`. A pool maps instances of a function over an iterable.``` uses 5 processes to map `my_iterable` to `my_func`with Pool(processes=5) as p:    outputs = p.map(my_func, my_iterable)```Using the above code creates a new `Pool` with 5 processes, and uses those 5 processes to process the function arguments. The code incurs some i/o overhead creating the `Pool`. This means that if the time required to process the function calls is less than the Pool overhead, using parallel processing will actually be slower than serial processing. There are two ways around this:1. Use process pools for bulk processing (ie function time much greater than i/o time)2. Maintain an open process pool to avoid repeated pool creation overheadMRL uses different functions to control the method of parallel processing used.The `new_pool_parallel` function implements parallel processing using a new pool for every function call, similar to the above code. This function is best used to process large numbers of inputs infrequently. Parallel processing is controlled by the `cpus` argument. If `cpus=None`, the `ncpus` environment variable is used (ie `os.environ['ncpus'] = '8'`)The `maybe_parallel` function allows for repeated use of a stateful process Pool, defined by the `GLOBAL_POOL` variable. By default, `GLOBAL_POOL=None`. To create a global pool, use the `set_global_pool` function.```set_global_pool(cpus=8)```If the `cpus=None`, `maybe_parallel` will run processes using `GLOBAL_POOL` if it exists, or serial processing if it does not. If `cpus` is not None, `maybe_parallel` defaults back to using `new_pool_parallel`If you need to frequently use parallel processing on small batches of inputs (ie batches from a model), set a global pool and use `maybe_parallel` Global Pool GotchasUsing a global pool allows us to take advantage of parallel processing on small batches without having the overhead of creating process pools over and over again. However, process pools left open accumulate memory. If memory usage builds up, use `refresh_global_pool` to release the memory and create a new global pool, or use `close_global_pool` to delete the global pool and reset it to `None`<jupyter_code># export

GLOBAL_POOL = None
os.environ['max_global_threads'] = '2000'

def set_global_pool(cpus=None):
    global GLOBAL_POOL
    if GLOBAL_POOL is not None:
        close_global_pool()
    
    if cpus is None:
        GLOBAL_POOL = None
    else:   
        GLOBAL_POOL = Pool(processes=cpus)
        GLOBAL_POOL.uses = 0
        
def close_global_pool():
    global GLOBAL_POOL
    if GLOBAL_POOL is not None:
        GLOBAL_POOL.close()
        del GLOBAL_POOL
        GLOBAL_POOL = None
        gc.collect()
    
def refresh_global_pool():
    global GLOBAL_POOL
    if GLOBAL_POOL is not None:
        cpus = GLOBAL_POOL._processes
        close_global_pool()
        set_global_pool(cpus=cpus)
        
def new_pool_parallel(func, iterable, cpus=None, **kwargs):
    p_func = partial(func, **kwargs)
    if is_container(iterable):
    
        if cpus is None:
            if 'ncpus' in os.environ.keys():
                cpus = int(os.environ['ncpus'])
            else:
                cpus = 0

        processes = min(cpus, len(iterable))

        if processes == 1:
            # spinning up a single pool has more overhead
            processes = 0

        if processes == 0:
            output = [p_func(i) for i in iterable]

        else:
            with Pool(processes=cpus) as p:
                output = p.map(p_func, iterable)
            
    else:
        output = p_func(iterable)
        
    return output

def maybe_parallel(func, iterable, cpus=None, **kwargs):
    global GLOBAL_POOL
    
    p_func = partial(func, **kwargs)
    
    if is_container(iterable):
        if cpus is not None:
            
            output = new_pool_parallel(func, iterable, cpus, **kwargs)
                    
        elif GLOBAL_POOL is not None:
            output = GLOBAL_POOL.map(p_func, iterable)
            GLOBAL_POOL.uses += 1
            if GLOBAL_POOL.uses > int(os.environ['max_global_threads']):
                refresh_global_pool()
                gc.collect()
            
        else:
            output = [p_func(i) for i in iterable]
            
    else:
        output = p_func(iterable)
        
    return output

def test_func(x):
    time.sleep(0.5)
    return x
start = time.time()
_ = [test_func(i) for i in range(10)]
t1 = time.time()
_ = new_pool_parallel(test_func, list(range(10))) # serial processing
t2 = time.time()
_ = new_pool_parallel(test_func, list(range(10)), cpus=4) # 4 cpus manually defined
t3 = time.time()
os.environ['ncpus'] = '4'
_ = new_pool_parallel(test_func, list(range(10))) # 4 cpus defined by environ variable
t4 = time.time()

print_str = f'''
Serial time: {t1-start:.2f}\n
new_pool_parallel, 0 cpus time: {t2-t1:.2f}\n
new_pool_parallel, 4 cpus (arg defined) time: {t3-t2:.2f}\n
new_pool_parallel, 4 cpus (environ defined) time: {t4-t3:.2f}\n
'''

print(print_str)
print(type(GLOBAL_POOL))
set_global_pool(5)
print(type(GLOBAL_POOL))

start = time.time()
_ = maybe_parallel(test_func, list(range(10)))
t1 = time.time()
_ = maybe_parallel(test_func, list(range(10)), cpus=2)
t2 = time.time()

print_str = f'''
maybe_parallel Global Pool (5 cpus) time: {t1-start:.2f}\n
maybe_parallel arg override 2 cpus time: {t2-t1:.2f}\n
'''

print(print_str)
start = time.time()
_ = maybe_parallel(test_func, list(range(10)))
t1 = time.time()
_ = maybe_parallel(test_func, list(range(10)), cpus=2)
t2 = time.time()

print_str = f'''
maybe_parallel Global Pool (5 cpus) time: {t1-start:.2f}\n
maybe_parallel arg override 2 cpus time: {t2-t1:.2f}\n
'''

print(print_str)<jupyter_output>maybe_parallel Global Pool (5 cpus) time: 1.00

maybe_parallel arg override 2 cpus time: 3.02<jupyter_text>Examples of different pools<jupyter_code>start = time.time()

for i in range(10):
    _ = new_pool_parallel(test_func, list(range(10)))
    
end = time.time() - start
print(f'{end:.2f} elapsed')
start = time.time()

for i in range(10):
    _ = maybe_parallel(test_func, list(range(10)))
    
end = time.time() - start
print(f'{end:.2f} elapsed')<jupyter_output>10.03 elapsed<jupyter_text>In the above example, `new_pool_parallel` takes 5 seconds longer to execute compared to `maybe_parallel`. The time difference is driven by the overhead generated by creating new pools  Debugging Parallel ProcessingErrors in parallel processing can be difficult to debug because the true error and stack trace are obscured by the parallel processing stack trace. If you have errors in parallel processing, first try setting `os.environ['ncpus'] = '0'` and running `close_global_pool` to disable python multiprocessing. This should reveal the true error.If everything works fine when multiprocessing is disabled, it is likely one of your functions is failing to pickle.  Download FilesDownload the `files` directory from the repo. Used for Collab notebooks<jupyter_code># export

def download_files():
    if not os.path.exists('files'):
        r = requests.get('https://dmai-mrl.s3.us-west-2.amazonaws.com/mrl_public/files.zip')
        
        with open('files.zip', 'wb') as f:
            f.write(r.content)
            
        with zipfile.ZipFile('files.zip', 'r') as zip_ref:
            zip_ref.extractall('.')
            
        os.remove('files.zip')
# hide
from nbdev.export import notebook2script; notebook2script()<jupyter_output>Converted 00_core.ipynb.
Converted 01_chem.ipynb.
Converted 02_template.filters.ipynb.
Converted 03_template.template.ipynb.
Converted 04_template.blocks.ipynb.
Converted index.ipynb.
{'fraction_non_alphanumeric': 0.096048024012006, 'fraction_numerical': 0.014107053526763383, 'mean_word_length': 3.6753975678203927, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 20, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '32669', 'license': 'MIT', 'n_tokens_mistral': 3198, 'n_tokens_neox': 3122, 'n_words': 1041, 'path': 'nbs/00_core.ipynb', 'repo_name': 'DarkMatterAI/mrl', 'size': '9995'}