mp_arima.py

# how to pool processes for arima

# given chunked data from data_chunker.py
# https://gist.github.com/justinhchae/13d246e8e2e2d521a8d2cce20eb09a09

# given arima function from run_arima.py
# https://gist.github.com/justinhchae/d2a2dc8b71b5f5fbbb0f7eabf68b6850

# dependencies
from tqdm import tqdm
import torch
import torch.multiprocessing as mp
from functools import partial

if __name__ == '__main__':
    
    # given a dataframe, df 
    chunked_data = chunk_data(df, price_col='c', time_col='t', n_prediction_units=1)

    model = run_arima

    p = mp.Pool(8)
    
    # pass the model and its params to a new partial object
    model_ = partial(model, n_prediction_units=1)
    
    # iterate over the partial object and the data
    # wrap the object inside tqdm to get a progress bar
    results = list(tqdm(p.imap(model_, chunked_data)))
    
    # close out the pool
    p.close()
    p.join()