#!/usr/bin/env python

import multiprocessing 
import os # For reading the amount of CPUs requested. 
import time # For clocking the calculation. 

def double(data):
    return data ** 2

if __name__ == '__main__':
    begin = time.time()
    inputs = list(range(10000)) # Makes an array from 0 to 10
    poolSize = int(os.environ['SLURM_JOB_CPUS_PER_NODE']) # Amount of CPUs requested.
    pool = multiprocessing.Pool(processes=poolSize,)
    poolResults = pool.map(double, inputs) # Do the calculation.
    pool.close() # Stop pool accordingly.
    pool.join()  # Wrap up data from the workers in the pool.
    print ('Pool output:', poolResults) # Results.
    elapsedTime = time.time() - begin
    print ('Time elapsed for ' , poolSize, ' workers: ', elapsedTime, ' seconds')