parallel processing of files from a directory

Question

This code does parallel processing of files read from a directory. It divides the directory into 'core' number of file chunks and process those chunks in parallel. 'cores' is the number of cores in the linux system.

from multiprocessing import Process
import os
from time import sleep
import multiprocessing

pros = []

def getFiles(directory):
    '''returns the files from a directory'''
    for dirpath,_,filenames in os.walk(directory):
        for f in filenames:
            yield os.path.abspath(os.path.join(dirpath, f))

def countFiles(directory):
    return len([name for name in os.listdir(directory) if os.path.isfile(os.path.join(directory, name))])

def foo(fileList):
    print(fileList)

def isCommon(a, b):
    aSet = set(a)
    bSet = set(b)
    if len(aSet.intersection(bSet)) > 0:
        return(True)
    return(False)

if __name__ == "__main__":
  '''get count of files in directory and split it in based on number of cores'''
  directory = ""
  noCores = multiprocessing.cpu_count()
  totalFilesCount = countFiles(directory)
  chunkSize = totalFilesCount/noCores
  totalChunks = noCores
  print("total files", totalFilesCount, "total chunks", totalChunks, "chunk size", chunkSize)
  filesProcessed = 0
  currentChunk = 0
  fileObj = getFiles(directory)

  listOFFiles = []
  while filesProcessed < totalFilesCount:
      filesList = []

      # if it is last chunk and totalFilesCount can't be divided equally then put leftover files in last core to get processed
      if currentChunk == totalChunks - 1 and totalFilesCount%noCores:
          chunkSize += totalFilesCount%noCores

      reached = 0
      for f in fileObj:
          filesList.append(f)
          if chunkSize == reached:
              break
          reached += 1

      listOFFiles.append(filesList)
      p = Process(target=foo, args=(filesList,))
      pros.append(p)
      p.start()

      currentChunk += 1
      filesProcessed += chunkSize

  for t in pros:
     t.join()

  for a, b in zip(listOFFiles, listOFFiles[1:]):
     assert isCommon(a, b) == False

Answer 1

The python standard library offers two process pools to simplify what you're trying to achieve: multiprocessing.Pool (as mentioned by @RootTwo) and concurrent.futures.ProcessPoolExecutor.

I personally like the second one (despite its more limited API) because you can replace it with a ThreadPoolExecutor without changing the rest of the code. Threads are better suited if your processing is I/O bound, because inter-process communication is slow and processes are slower to start (especially on Windows). With the concurrent.futures API, you can easily swap the two and test which one works better.

There are some other ways you can improve the code, most notably:

1. If you're on Python 3.x as you should be, I suggest the use of pathlib instead of os.path.
2. Put code from the if __name__ == "__main__": into a "main" function to make automated testing easier.
3. Use snake_case instead of camelCase as that's the accepted standard for Python.
4. (Opinion) If you're on python 3.x, you could also use type annotations to improve the readability of the code.

With this in mind, I would rewrite your code as:

from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Iterable

def get_files(directory: Path) -> Iterable[Path]:
    # Using glob simplifies the code in these cases.
    return (file for file in directory.glob("**/*") if file.is_file())

def foo(file: Path):
    ...

# This can be easily called from a test
def process_files(directory: Path):
    # I am using sum, so that the result is computed lazily and we
    # do not need to build a list of all files. If the directory is
    # very large, this could save a lot of memory.
    # Since get_files returns a one-shot generator, we cannot
    # save it to a variable and reuse it here and below.
    file_count = sum(1 for _ in get_files(directory))
    
    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = executor.map(foo, get_files(directory))
        # Reading the value returned by the executor is very
        # important, because exceptions happening in the `foo`
        # function will be captured and not be raised until the
        # value is read, thus obfuscating the errors.
        #        
        # Another nice solution for the progress would be
        # using the library tqdm.
        for i, _ in enumerate(as_completed(futures)):
            print(f"Processed file {i+1} / {file_count}")

if __name__ == "__main__":
    process_files(Path(""))