Orchestrate Processes

Dragon provides its own native API to finely and programmatically control where and how processes get executed. Below, we work through the use of dragon.native to execute a combination of user applications, specify their placement on hardware, and how to manage their output

ProcessGroup

Anytime you have some number of processes you want to execute, ProcessGroup is where you want to begin. In fact, ProcessGroup is so powerful that Dragon uses it as the backbone for its implementation of Pool().

Hello World!

We’ll begin by just executing the classic “Hello world!” example. In the snippet, we begin by creating a ProcessGroup object that contains all the API for managing the processes we’ll assign to it. We’ll assign processes to the group by defining a ProcessTemplate.

A ProcessGroup can contain as many templates as we’d like, and we can also tell ProcessGroup how many instances of a given template we want to execute. In this example, we’ll launch 4 instances of our "Hello World!" template. After all that setup is complete, we’ll initialize the infrastructure for the the ProcessGroup object and start execution of the 4 "Hello World!" instances. We then tell our ProcessGroup object to join on the completion of those 4 instances and then close all the ProcessGroup infrastructure

Listing 54 Execute a group of “Hello world!” processes with ProcessGroup
 1import socket
 2from dragon import ProcessGroup, ProcessTemplate
 3
 4
 5def hello_world():
 6    print(f'hello from process {socket.gethostname()}!')
 7
 8
 9def run_hello_world_group():
10
11    pg = ProcessGroup()
12    hello_world_template = ProcessTemplate(target=hello_world)
13    pg.add_process(nproc=4, template=hello_world_template)
14
15    pg.init()
16    pg.start()
17
18    pg.join()
19    pg.close()
20
21
22if __name__ == '__main__':
23
24    run_hello_world_group()

Defining Multiple Templates

Say you’d like to run different applications but have them be part of the same ProcessGroup. That is easily done by providing multiple templates to your ProcessGroup object. In the following example, we’ll create a data generator app and a consumer of that data that will be connected to each other via a Queue. The Queue will passed as input to each of the processes.

Listing 55 Run ProcessGroup with a process generating data passed to consumer via a Queue
 1import random
 2
 3from dragon import ProcessGroup, ProcessTemplate
 4from dragon.native.queue import Queue
 5
 6
 7def data_generator(q_out, n_outputs):
 8
 9    for _ in range(n_outputs):
10        output_data = int(100 * random.random())
11        print(f'generator feeding {output_data} to consumer', flush=True)
12        q_out.put(output_data)
13
14
15def data_consumer(q_in, n_inputs):
16
17    for _ in range(n_inputs):
18        input_data = q_in.get()
19        result = input_data * 2
20        print(f'consumer computed result {result} from input {input_data}', flush=True)
21
22
23def run_group():
24
25    q = Queue()
26    pg = ProcessGroup()
27
28    generator_template = ProcessTemplate(target=data_generator,
29                                         args=(q, 5))
30    consumer_template = ProcessTemplate(target=data_consumer,
31                                        args=(q, 5))
32
33    pg.add_process(nproc=1, template=generator_template)
34    pg.add_process(nproc=1, template=consumer_template)
35
36    pg.init()
37    pg.start()
38
39    pg.join()
40    pg.close()
41
42
43if __name__ == '__main__':
44
45    run_group()

Managing Output/stdout

In the above example, we had a bit of redundant output. We get the input via the generator process printed to stdout and then that value is echoed in the consumer process:

Listing 56 Output from execution of consumer/generator example without piping generator output to `/dev/null`
 1(_env) user@hostname:~/dragon_example> dragon generator_consumer_example.py
 2consumer computed result 140 from input 70
 3generator feeding 70 to consumer
 4consumer computed result 160 from input 80
 5generator feeding 80 to consumer
 6consumer computed result 14 from input 7
 7generator feeding 7 to consumer
 8consumer computed result 28 from input 14
 9generator feeding 14 to consumer
10generator feeding 72 to consumer
11consumer computed result 144 from input 72

Since the generator information is redundant, let’s send it to /dev/null by modifying the driver function in above example:

Listing 57 Sending generator stdout to `/dev/null`
 1 from dragon import ProcessGroup, ProcessTemplate
 2 from dragon.native.process import Popen
 3 from dragon.native.queue import Queue
 4
 5 def run_group():
 6
 7     q = Queue()
 8     pg = ProcessGroup()
 9
10     # Tell the dragon to get rid of the geneator's stdout
11     generator_template = ProcessTemplate(target=data_generator,
12                                          args=(q, 5),
13                                          stdout=Popen.DEVNULL)
14
15     consumer_template = ProcessTemplate(target=data_consumer,
16                                         args=(q, 5))
17
18     pg.add_process(nproc=1, template=generator_template)
19     pg.add_process(nproc=1, template=consumer_template)
20
21     pg.init()
22     pg.start()
23
24     pg.join()
25     pg.close()

The end result is an easier to parse stream of output:

(_env) user@hostname:~/dragon_example> dragon generator_consumer_sanitized_output.py
consumer computed result 50 from input 25
consumer computed result 30 from input 15
consumer computed result 80 from input 40
consumer computed result 44 from input 22
consumer computed result 12 from input 6

Placement of ProcessGroup Processes via Policy

Commonly, a user wants to have one process run on a particular hardware resource (eg: GPU) while other processes are perhaps agnostic about their compute resources. In Dragon, this is done via the Policy API. To illustrate this, we’ll take the basic template of the consumer-generator example above and replace it with some simple PyTorch code While we’re not doing anything complicated or exercising this paradigm as you might in reality (e.g., generating model data on a CPU and feeding training inputs to a GPU), it provides a template of how you might do somethign more complicated. We’ll replace the data generator function from above with initialization of PyTorch model parameters. We’ll pass these to the consumer process which will use a GPU to train the data. And lastly, we’ll use the Policy API to specify the PyTorch model is trained on a compute node we know has a GPU present.

Listing 58 Generating training input data on the CPU and passing to a GPU process for PyTorch training
  1from dragon import ProcessGroup, ProcessTemplate, Policy
  2from dragon.native.queue import Queue
  3
  4import torch
  5import math
  6
  7
  8def data_generate(q_out):
  9
 10    torch.set_default_device("cpu")
 11
 12    dtype = torch.float
 13
 14    # Create Tensors to hold input and outputs.
 15    # By default, requires_grad=False, which indicates that we do not need to
 16    # compute gradients with respect to these Tensors during the backward pass.
 17    x = torch.linspace(-math.pi, math.pi, 2000, dtype=dtype)
 18    y = torch.sin(x)
 19
 20    # Create random Tensors for weights. For a third order polynomial, we need
 21    # 4 weights: y = a + b x + c x^2 + d x^3
 22    # Setting requires_grad=True indicates that we want to compute gradients with
 23    # respect to these Tensors during the backward pass.
 24    a = torch.randn((), dtype=dtype, requires_grad=True)
 25    b = torch.randn((), dtype=dtype, requires_grad=True)
 26    c = torch.randn((), dtype=dtype, requires_grad=True)
 27    d = torch.randn((), dtype=dtype, requires_grad=True)
 28
 29    q_out.put((x, y, a, b, c, d))
 30
 31
 32def pytorch_train(q_in):
 33
 34    torch.set_default_device("cuda")
 35
 36    x, y, a, b, c, d = q_in.get()
 37
 38    x.to('cuda')
 39    y.to('cuda')
 40    a.to('cuda')
 41    b.to('cuda')
 42    c.to('cuda')
 43    d.to('cuda')
 44
 45    learning_rate = 1e-6
 46    for t in range(2000):
 47        # Forward pass: compute predicted y using operations on Tensors.
 48        y_pred = a + b * x + c * x ** 2 + d * x ** 3
 49
 50        # Compute and print loss using operations on Tensors.
 51        loss = (y_pred - y).pow(2).sum()
 52        if t % 100 == 99:
 53            print(t, loss.item())
 54
 55        # Use autograd to compute the backward pass.
 56        loss.backward()
 57
 58        # Manually update weights using gradient descent.
 59        with torch.no_grad():
 60            a -= learning_rate * a.grad
 61            b -= learning_rate * b.grad
 62            c -= learning_rate * c.grad
 63            d -= learning_rate * d.grad
 64
 65            # Manually zero the gradients after updating weights
 66            a.grad = None
 67            b.grad = None
 68            c.grad = None
 69            d.grad = None
 70
 71    print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')
 72
 73
 74def run_group():
 75
 76    q = Queue()
 77    pg = ProcessGroup()
 78
 79    # Since we don't care where the data gets generated, we let
 80    # Dragon determine the placement by leaving the placement kwarg blank
 81    generator_template = ProcessTemplate(target=data_generate,
 82                                         args=(q,))
 83
 84    # node 'pinoak0033' is the hostname for a node with NVIDIA A100 GPUs.
 85    # We tell Dragon to use it for this process via the policy kwarg.
 86    train_template = ProcessTemplate(target=pytorch_train,
 87                                     args=(q,),
 88                                     policy=Policy(placement=Policy.Placement.HOST_NAME,
 89                                                   host_name='pinoak0033'))
 90
 91    pg.add_process(nproc=1, template=generator_template)
 92    pg.add_process(nproc=1, template=train_template)
 93
 94    pg.init()
 95    pg.start()
 96
 97    pg.join()
 98    pg.close()
 99
100
101if __name__ == '__main__':
102
103    run_group()