9
\$\begingroup\$

I'm building a script where I pass in a text file of hostnames, and a text file of commands and each command will then be run against each host.

To prevent this getting out of hand I'm also passing in a maximum thread count and trying to manage a maximum number of workers at any one time.

The code I would like feedback on is as follows:

main.py

#!/usr/bin/python3
import sys
from lib.core.input import InputParser, InputHelper
from lib.core.output import OutputHelper, Level
from lib.core.threader import Pool


def build_queue(arguments, output):
    queue = list()
    for target in InputHelper.process_targets(arguments):
        for command in InputHelper.process_commands(arguments):
            output.terminal(Level.VERBOSE, target, command, "Added to Queue")
            queue.append(command)
    return queue


def main():
    parser = InputParser()
    arguments = parser.parse(sys.argv[1:])

    output = OutputHelper(arguments)

    output.print_banner()

    pool = Pool(arguments.threads, build_queue(arguments, output), arguments.timeout, output)
    pool.run()


if __name__ == "__main__":
    main()

threader.py

import threading
import os


class Worker(object):
    def __init__(self, pool):
        self.pool = pool

    def __call__(self, task, output, timeout):
        self.run_task(task)
        self.pool.workers.append(self)

    @staticmethod
    def run_task(task):
        os.system(task)


class Pool(object):
    def __init__(self, max_workers, queue, timeout, output):
        self.queue = queue
        self.workers = [Worker(self) for w in range(max_workers)]
        self.timeout = timeout
        self.output = output

    def run(self):
        while True:

            # make sure resources are available
            if not self.workers:
                continue

            # check if the queue is empty
            if not self.queue:
                break

            # get a worker
            worker = self.workers.pop(0)

            # get task from queue
            task = self.queue.pop(0)

            # run
            thread = threading.Thread(target=worker, args=(task, self.output, self.timeout))
            thread.start()

Essentially, I don't know what I don't know. There's likely some basic gaps in my knowledge here, and I'd love feedback and examples that help me to bridge those gaps. Any and all help is very welcome, and very appreciated.

\$\endgroup\$
5
  • \$\begingroup\$ Are you doing this for the fun of learning things or are you planning on using this on real systems to automate executing the same commands on several hosts? \$\endgroup\$ Oct 23 '18 at 15:46
  • \$\begingroup\$ @MathiasEttinger to learn, but I'm planning to use it to manage some components in my home lab also \$\endgroup\$
    – Michael A
    Oct 23 '18 at 22:59
  • \$\begingroup\$ In such case, you might be better off learning to use mature automation tools such as Ansible for instance. You’ll be able to focus on what you want to do, not how. \$\endgroup\$ Oct 24 '18 at 12:42
  • \$\begingroup\$ @MathiasEttinger I'm familiar with a few other approaches (even in this case I can acheive much with xargs) but I'm trying to use this project as a way to better understand how to do such things in Python. \$\endgroup\$
    – Michael A
    Oct 25 '18 at 1:08
  • 1
    \$\begingroup\$ For learning purposes, I'd say good for you. For actual use, you want something like GNU parallel. It likely can achieve exactly what your code does (and then some) with a single command line invocation. \$\endgroup\$ Oct 29 '18 at 5:22
9
+100
\$\begingroup\$

1. Introduction

From a practical point of view, the most important points are that Python has batteries included:

  1. To run tasks in a pool of worker threads, use concurrent.futures.ThreadPoolExecutor.

  2. But if what you really want to do is run external programs via the shell, as suggested by the use of os.system, then you don't need threads at all! Use subprocess.Popen instead.

    (Note that on some operating systems, notably macOS, you can only have one call to os.system running at a time. Subsequent calls, even from other threads, have to wait for the first call to finish. On these systems you have to use subprocess.Popen if you want the subprocesses to run in parallel.)

Nonetheless I think it's a good exercise to have a go at writing your own thread pool, and so I'm go to look at how you might improve the code in the post. There's quite a lot of ground to cover, so I'm going to take it in stages, starting by reviewing improving the code in the post, and then identifying and fixing various problems with the design.

2. Initial review

  1. There are no docstrings. How do I use this module? Which classes are public and which are private? Do I create my own Worker objects or do I let the Pool create them for me? What kind of object do I pass for queue? How should I specify the timeout?

  2. The Pool constructor takes timeout and output arguments but these have no effect: they are passed to Worker.__call__ but then they are ignored.

  3. The Pool class has only one method (other than __init__). When you have a class with just one method, then what you need is a function. Not everything needs to be a class! It would simplify the code if it were turned into a function with a specification like this:

    def run_shell_commands_in_parallel(commands, max_threads):
        """Run shell commands in parallel. max_threads is the maximum
        number of threads that may run simultaneously.
    
        """
    
  4. Running a shell command is not very general. It would make sense to generalize the code to do arbitrary function calls in parallel:

    def apply_in_parallel(fun, args, max_threads=4):
        """Apply fun to each of the arguments in parallel. max_threads is the
        maximum number of threads that may run simultaneously.
    
        """
    

    To try to run a bunch of shell commands in parallel, you'd call it like this:

    apply_in_parallel(os.system, ['command1', 'command2', ...])
    

    but now you can use it for other parallel tasks:

    apply_in_parallel(print, range(10))
    
  5. The main thread calls workers.pop(0) and the worker threads call workers.append(self), but the workers data structure is just an ordinary Python list, which is not thread-safe. Whenever you have a data structure that's shared between threads, access to the data structure needs to be synchronized, for example by using a lock to ensure that only one thread can be updating it at a time.

  6. All the Worker objects are identical — their only attribute is self.pool, and that's the same for every worker, so instead of a list of workers all we really need is a count of how many workers are idle.

3. Revised code

Fixing the issues in §1 above yields the following, which does essentially the same thing as the code in the post, but in a simpler way:

def apply_in_parallel(fun, args, max_threads=4):
    """Apply fun to each of the arguments in parallel. max_threads is the
    maximum number of threads to run simultaneously.

    """
    available_workers = max_threads
    available_workers_lock = Lock()

    def worker(arg):
        nonlocal available_workers
        fun(arg)
        with available_workers_lock:
            available_workers += 1

    for arg in args:
        while True:
            with available_workers_lock:
                if available_workers:
                    available_workers -= 1
                    break
        Thread(target=worker, args=(arg,)).start()

4. Problem: busy-waiting

The while True: loop is busy-waiting. If it takes a long time for a worker to become available, the main thread will be wasting effort by repeatedly taking the lock and testing the condition if available_workers: which remains false.

What we'd like instead is to be able to suspend the main thread until a worker becomes available. For this use case what we need is a semaphore, and Python provides us with threading.Semaphore:

from threading import Semaphore, Thread

def apply_in_parallel(fun, args, max_threads=4):
    """Apply fun to each of the arguments in parallel. max_threads is the
    maximum number of threads that may run simultaneously.

    """
    available = Semaphore(max_threads)

    def worker(arg):
        fun(arg)
        available.release()

    for arg in args:
        available.acquire()
        Thread(target=worker, args=(arg,)).start()

5. Problem: can't tell when threads are done

The code in §4 starts the shell commands but does not wait for them all to complete. But in many use cases it is important to wait for workers to finish running before proceding.

To ensure that all threads have finished, we must join them before returning:

from threading import Semaphore, Thread

def apply_in_parallel(fun, args, max_threads=4):
    """Apply fun to each of the arguments in parallel and wait until all
    calls have completed. max_threads is the maximum number of threads
    that may run simultaneously.

    """
    available = Semaphore(max_threads)

    def worker(arg):
        fun(arg)
        available.release()

    threads = []
    for arg in args:
        available.acquire()
        thread = Thread(target=worker, args=(arg,))
        threads.append(thread)
        thread.start()

    for thread in threads:
        thread.join()

6. Problem: too many threads

The code in §5 creates a new thread for every task, and every thread exits after completing its task. This defeats one of the purposes of having a pool of threads, namely that it avoids some of the overhead of creating new threads by reusing the threads in the pool.

The usual approach to implementing a pool of threads is for each thread to execute another task when it is finished (so long as there are more tasks to run). That means that the main thread is going to need a safe way of sending the tasks to the running threads, and Python provides queue.Queue which is exactly what we want:

from queue import Queue
from threading import Thread

def apply_in_parallel(fun, args, max_threads=4):
    """Apply fun to each of the arguments in parallel and wait until all
    calls have completed. max_threads is the maximum number of threads
    that may run simultaneously.

    """
    queue = Queue()

    def worker():
        while True:
            arg = queue.get()
            fun(arg)
            queue.task_done()

    for _ in range(max_threads):
        Thread(target=worker).start()
    for arg in args:
        queue.put(arg)
    queue.join()

7. Problem: threads are left dangling

The problem with the implementation in §6 is that the worker threads never finish! When all the tasks are done, each worker thread is blocked in queue.get waiting for the next task, but it will never arrive. This is a dangerous resource leak because the number of threads that we can create is limited by the operating system, and so to ensure that we can continue to create new threads later in the program, we must clean up all the threads that we started.

So when all tasks are done, we must tell the threads to exit. A convenient way to do that is to pass in a special sentinel argument that is different from any argument that the caller could pass in:

from queue import Queue
from threading import Thread

def apply_in_parallel(fun, args, max_threads=4):
    """Apply fun to each of the arguments in parallel and wait until all
    calls have completed. max_threads is the maximum number of threads
    that may run simultaneously.

    """
    queue = Queue()
    sentinel = object()

    def worker():
        while True:
            arg = queue.get()
            if arg is sentinel:
                break
            fun(arg)

    threads = []
    for _ in range(max_threads):
        thread = Thread(target=worker)
        threads.append(thread)
        thread.start()
    for arg in args:
        queue.put(arg)
    for _ in range(max_threads):
        queue.put(sentinel)
    for thread in threads:
        thread.join()
\$\endgroup\$
1
  • 1
    \$\begingroup\$ This is awesome, thank-you for going into such detail! \$\endgroup\$
    – Michael A
    Oct 30 '18 at 0:48
5
+50
\$\begingroup\$

Prefer generators instead of lists

build_queue creates a list of tasks up front. It's unnecessary to store all task details in memory up front. You could use a generator instead, and yield the task parameters. That will minimize the memory use, by generating the task details just before the execution of the individual tasks.

Avoid busy loops

This piece of code results in busy loop / busy waiting:

while True:
    # make sure resources are available
    if not self.workers:
        continue

    # ...

That is, when no workers are available and the queue is not empty, the main thread will be spinning until a worker becomes available. Such spinning is useless activity for the CPU, wasting resources. Some amount of sleeping would help, but there's a much better solution (keep reading).

Use the libraries, Luke...

Instead of implementing a worker pool and job queue yourself, it's better to use what is provided by the Python Standard Library for concurrent execution. Notably, the Pool class if you go with process-based parallelism (recommended for your example), or the Queue class if you go with thread-based parallelism.

You will be able to get rid of your manual queue management code, get something more robust and performant in return, and be able to focus on the implementation of the worker.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.