# coding: utf-8

# In[1]:

#ENVIRONMENT SETUP

#Import the required libraries
#- pandas for data management
import pandas as pd
#- os and subprocess for system calls
import os
import subprocess as sp
#- time for execution time computation
import time

#Helper function for running piped Linux commands easily
def executePipedCommands(cmd):
    procs = []
    for n in range(len(cmd)):
        subprocessCmd = cmd[n].split(" " )
        if n == 0:
            procs.append(sp.Popen(subprocessCmd, stdout=sp.PIPE))
        elif n == len(cmd)-1:
            out = sp.check_output(subprocessCmd, stdin=procs[-1].stdout)
        else:
            procs.append(sp.Popen(subprocessCmd, stdin=procs[-1].stdout, stdout=sp.PIPE))
    return out

#processFile is the function used to simulate the processing of a single file
#Here, processFile does not actually process the input file, it simply runs a Linux sleep command
#whose length in proportional to the file's size
#Technically, processFile relies on the subprocess Python package which allows to run system commands
#in async or sync mode (the latter requiring an additional call to the wait() function that waits for
#the command's completion)
#We will use processFile in sync mode for the serial case and in async mode for the parallel case
#since this will allow us to start several concurrent processing tasks without any additional
#parallel processing framework
#Inputs:
#filename: name of the file to process
#sleepFactor: factor to apply to the file size for determining the sleep command length
def processFile(filename, sleepFactor, async=False):
    p = sp.Popen(["sleep", str(os.path.getsize(filename) * sleepFactor)])
    if not async:
        p.wait()
        p.communicate()
        if p.returncode != 0:
            raise Exception("Processing task errored" )
    return p

#From now on, the functions below are specifically dedicated to parallel computations

#startNewProcessingTask starts a new processing task identified by a task index on a worker
#identified by a worker index
#Inputs:
#fileTable: table of files to process
#sleepFactor: factor to apply to the file size for determining the sleep command length
#worker: worker index, belongs to the [0, (#workers-1)] range
#task: task index, belongs to the [0, (#files-1)] range
def startNewProcessingTask(fileTable, sleepFactor, worker, task, verbose=False):
    #This is for the parallel case, processFile is therefore launched in async mode
    p = processFile(fileTable.at[task, "name"], sleepFactor, async=True)
    if verbose:
        print("Worker " + str(worker) + " starting task " + str(task))
    return p

#getNextAssignedTask returns the index of the next task a worker has to run
#If there are no tasks left to run for the worker, the function returns None
#Inputs:
#currentTask: index of the current task
#tasksAssigned: list of all tasks assigned to the worker
def getNextAssignedTask(currentTask, tasksAssigned):
    #Find the index of currentTask in tasksAssigned
    idx = tasksAssigned.index(currentTask)
    if idx == len(tasksAssigned)-1:
        #If currentTask was the last element, return None
        return None
    else:
        #Otherwise return the next value
        return tasksAssigned[idx+1]

#parallelProcessing is the function that distributes and handles processing tasks on the workers
#It mainly relies on the input taskAssignment variable which has to be computed beforehand and
#and describes the distribution schemes (it contains for each worker the list of assigned tasks)
#Inputs:
#numFiles: number of files to process
#numWorkers: number of workers to use
#fileTable: table of files to process
#sleepFactor: factor to apply to the file size for determining the sleep command length
#taskAssignment: list of tasks assigned to each worker (the i-th element of taskAssignment is the
#list of tasks assigned to the i-th worker)
def parallelProcessing(numFiles, numWorkers, fileTable, sleepFactor, taskAssignment, verbose=False):
    #Start a timer
    start = time.time()
    #Keep track of completed tasks, globally and for each worker
    numTasksCompleted = 0
    tasksCompleted = [[] for _ in range(numWorkers)]
    #Keep track of the execution time on each worker
    execTimeWorkers = [None] * numWorkers
    #The following lists store the task that is currently executing on each worker,
    #more precisely the task's index in currentTask and the associated process in currentProcess
    #Initialize currentTask with the first element found in taskAssignment
    currentTask = [elem[0] for elem in taskAssignment]
    currentProcess = [None] * numWorkers
    #Start processing and run as long as not all files have been processed
    while numTasksCompleted < numFiles:
        #Loop through workers
        for idx in range(0, numWorkers):
            #The following test checks if the worker has already completed all its tasks,
            #in which case its currentTask would have a None value
            if currentTask[idx] is not None:
                if not currentProcess[idx]:
                    #If the current worker does not have any task assigned yet, start a new one
                    if verbose:
                        print("Worker " + str(idx) + ": no task assigned yet" )
                    currentProcess[idx] = startNewProcessingTask(fileTable, sleepFactor, idx, currentTask[idx])
                else:
                    #Otherwise investigate the running process
                    p = currentProcess[idx]
                    #The following tests checks if the task has terminated
                    if not(p.poll() == None):
                        #If the task has not terminated, there is nothing to do, just move on to the next worker
                        #If it has terminated, we will increase the completed tasks counter and start a new task
                        #Before that, we check the return code to make sure the process terminated successfully
                        p.communicate()
                        if p.returncode != 0:
                            raise Exception("Processing task errored" )
                        #Increase the completed tasks counter and add the index of the completed task to the
                        #worker's list of completed tasks
                        numTasksCompleted += 1
                        tasksCompleted[idx].append(currentTask[idx])
                        if verbose:
                            print("Worker " + str(idx) + ": task " + str(currentTask[idx]) + " completed!" )
                        #Get the next task the worker has to run
                        currentTask[idx] = getNextAssignedTask(currentTask[idx], taskAssignment[idx])
                        if currentTask[idx] is not None:
                            #If there is a task left to run, start it
                            currentProcess[idx] = startNewProcessingTask(fileTable, sleepFactor, idx, currentTask[idx])
                        else:
                            #Otherwise it means that the worker has completed all its tasks
                            #Save the execution time
                            execTimeWorkers[idx] = time.time() - start
    #Before exiting, we check that execution went well by comparing the taskAssignment and tasksCompleted variables
    #which are supposed to be equal
    for idx in range(0, numWorkers):
        t1 = taskAssignment[idx]
        t2 = tasksCompleted[idx]
        if len(t1) != len(t2) or not all([t1[elem] == t2[elem] for elem in range(0, len(t1))]):
            raise Exception("Something went wrong on worker " + str(idx))
    #End the timer and display the execution time
    end = time.time()
    execTimeOverall = end - start
    print("Execution time in PARALLEL mode: %.2f seconds" % execTimeOverall)
    return execTimeOverall, execTimeWorkers

# In[2]:

#BUILD THE LIST OF FILES TO PROCESS

#In  this case, we will use log files of jobs that have run on our server
#We randomly select files larger than 1 kB and smaller than 10 MB
numFiles = 100;
rootPath = "/some/path/"
cmd = ["find " + rootPath + " -type f -name *log* -size -10000k -size +1k -printf '%s\t%p\n", "shuf -n " + str(numFiles)]
out = executePipedCommands(cmd)

#Get the file list from the commands output
fileList = out.split("\n" )

#At the end of the list is an extra empty element that we need to discard
fileList = fileList[0:-1]

#Get the size of each file: the size is the first field in each element of fileList
#The file size is also preceded by a single quote character that we need to discard
fileSize = [int(elem[1:len(elem)].split("\t" )[0]) for elem in fileList]

#Compute a 'sleep factor' so that the serial execution runs in around 1 minute
sleepFactor = 60.0 / sum(fileSize)

#The file name is the second field in each element of fileList
fileName = [elem.split("\t" )[1] for elem in fileList]

#Create a single table concatenating name and size data
fileTable = pd.DataFrame({"name":fileName, "size":fileSize})

print(fileTable)

# In[3]:

#SERIAL FILE PROCESSING

#We run the processing a first time in serial mode with the aim of getting the execution time
#in serial mode so we can compute the parallel speed-up

start = time.time()
for idx in range(0, len(fileTable.index)):
    processFile(fileTable.at[idx, "name"], sleepFactor, async=False)
end = time.time()
execTimeSerial = end - start;

print("Execution time in SERIAL mode: %.2f seconds" % execTimeSerial)

# In[4]:

#PARALLEL FILE PROCESSING - Naive version

#We start with the 'naive' distribution of tasks, which consists in the following:
#Assume there are M files and N workers
#File indices range from 0 to M-1, worker indices from 0 to N-1
#The i-th worker processes all files for which mod(file index, M) == i

#Use 4 workers
numWorkers = 4

#Generate the task distribution scheme
taskAssignmentNaive = [[] for _ in range(numWorkers)]
for idx in range(0, numWorkers):
    taskAssignmentNaive[idx] = [elem for elem in range(0, numFiles) if elem%numWorkers == idx]
    print("Tasks assigned to worker " + str(idx) + ":" )
    print(taskAssignmentNaive[idx])
print("\n" )

#Run processing tasks
execTimeParallelNaive, execTimeWorkersNaive = parallelProcessing(numFiles, numWorkers, fileTable, sleepFactor, taskAssignmentNaive)
print("\nParallel speed up: %.2fx" % (execTimeSerial / execTimeParallelNaive))
print("Execution time on the individual workers: " + ", ".join(["%.2f" % elem for elem in execTimeWorkersNaive]))

#As we can see below, the parallel speed up is suboptimal and there is a significant skew
#among workers' individual execution times
#The overall execution time actually is the execution time seen on the worker which took the most time
#to complete its tasks

# In[5]:

#PARALLEL FILE PROCESSING - Optimized version

#The optimized version focuses on building a more elaborate distribution scheme
#In pseudo code, it reads like this:
#    - sort the list of files to process according to file size in descending order
#    - keep track of the expected workload on each worker, initializing it to zero
#    - loop over the sorted list of files to process
#         - at each step, look for the worker which has the least work to do
#         - assign the next task to this worker
#         - increment this worker's expected workload with the file size
#This will end up with a different task assignment with a lower skew

#Sort the table of files to process as indicated above: according to file size in descending order
fileTableSorted = fileTable.sort_values(by=["size"], ascending=False, inplace=False)
fileTableSorted = fileTableSorted.reset_index(drop=True)

#Build the new distribution scheme, keeping track of the expected workload on each worker
taskAssignmentOptimized = [[] for _ in range(numWorkers)]
expectedWorkload = [0] * numWorkers
for idxFile in range(0, numFiles):
    #Look for worker with the minimum workload
    idxWorker = expectedWorkload.index(min(expectedWorkload))
    taskAssignmentOptimized[idxWorker].append(idxFile)
    expectedWorkload[idxWorker] += fileTableSorted.at[idxFile, "size"]

#Display the optimized task assignment
for idx in range(0, numWorkers):
    print("Tasks assigned to worker " + str(idx) + ":" )
    print(taskAssignmentOptimized[idx])
print("\n" )

#Display the expected workload
print("Expected workload on each worker: " + ", ".join(["%.2f" % elem for elem in expectedWorkload]) + "\n" )

#Run processing tasks
execTimeParallelOptimized, execTimeWorkersOptimized = parallelProcessing(numFiles, numWorkers, fileTableSorted, sleepFactor, taskAssignmentOptimized)
print("\nParallel speed up: %.2fx" % (execTimeSerial / execTimeParallelOptimized))
print("Execution time on the individual workers: " + ", ".join(["%.2f" % elem for elem in execTimeWorkersOptimized]))

#As we can see below, the parallel speed up is now much better and the skew has been nearly eliminated