<?xml version="1.0" encoding="utf-8"?><feed xmlns="http://www.w3.org/2005/Atom" xml:lang="en"><generator uri="https://jekyllrb.com/" version="4.3.4">Jekyll</generator><link href="https://ltmeyer.github.io/feed.xml" rel="self" type="application/atom+xml"/><link href="https://ltmeyer.github.io/" rel="alternate" type="text/html" hreflang="en"/><updated>2024-10-23T10:43:11+00:00</updated><id>https://ltmeyer.github.io/feed.xml</id><title type="html">blank</title><entry><title type="html">Online Data Generation for Better, Faster, and Cheaper Training</title><link href="https://ltmeyer.github.io/2024/10/15/online_data_generation.html" rel="alternate" type="text/html" title="Online Data Generation for Better, Faster, and Cheaper Training"/><published>2024-10-15T00:00:00+00:00</published><updated>2024-10-15T00:00:00+00:00</updated><id>https://ltmeyer.github.io/2024/10/15/online_data_generation</id><content type="html" xml:base="https://ltmeyer.github.io/2024/10/15/online_data_generation.html"><![CDATA[<div class="intro"> This blog post discusses the design of an online deep learning framework that generates synthetic data simultaneously with the training. It summarizes the ideas of papers presented at <a href="https://dl.acm.org/doi/abs/10.1145/3581784.3607083">SuperComputing</a> and <a href="https://proceedings.mlr.press/v202/meyer23b.html">ICML</a>. <br/><br/> </div> <h1 id="bigger-models-need-bigger-datasets">Bigger Models Need Bigger Datasets</h1> <p>The current scaling laws of deep learning indicate that not only bigger models tend to perform better, but also that they require bigger datasets to do so. Nowadays, it is common to see training on hundreds of gigabytes if not terabytes of data. To such an extent that easily accessible big datasets are now hardly enough to train large models. There are some application domains, however, for which training data can be generated synthetically, and thus dataset are virtually unlimited. For instance, one can think about using a smaller generative model to produce synthetic texts in order to train a larger one. For this post, we focus on AI for Science applications. In this context, training data can be generated by executing an external program: typically a numerical solver that takes as input physical parameters and simulates relevant quantities ruled by partial differential equations.</p> <h1 id="data-limitations-with-the-typical-training-pipeline">Data Limitations with the Typical Training Pipeline</h1> <p>The training on synthetically generated data is a two-step process. First, training data are produced using an oracle and stored to disk. Second, the actual training phase loads these data back from disk and performs the forward and backward passes to update model’s weights.</p> <div class="row mt-3"> <div class="col-sm mt-3 mt-md-0"> <figure> <picture> <source class="responsive-img-srcset" srcset="/assets/img/traditional_training.png" sizes="95vw"/> <img src="/assets/img/traditional_training.png" class="img-fluid rounded z-depth-1" width="100%" height="auto" loading="eager" onerror="this.onerror=null; $('.responsive-img-srcset').remove();"/> </picture> </figure> </div> </div> <div class="caption"> <b>Typical Training Pipeline</b>: The dataset is generated beforehand, then loaded from disk. Data loading and model updates overlap. </div> <p>This process presents several disadvantages:</p> <ul> <li>Because I/O is much slower than computing, <strong>data generation is slow</strong>. Nonetheless, this cost is paid only once. The same dataset can be used for several trainings;</li> <li>For the same reason, because data are loaded from disk, the overall <strong>training phase may be slow</strong>. This is generally mitigated by overlapping data loading and model updates;</li> <li>Because disk storage is expensive and limited, <strong>the training dataset is generally reduced to hundreds of GBs</strong>;</li> <li>Because the dataset is generated beforehand, it is rare to generate only the data that are the most useful for training. <strong>Training does not benefit from techniques like active learning</strong>, curriculum learning, or Bayesian inference.</li> </ul> <blockquote> <p>How can we alleviate these limitations and improve large model training on synthetically generated data?</p> </blockquote> <h1 id="leveraging-network-speed">Leveraging Network Speed</h1> <p>One possible solution appears while considering the characteristics of hardware. On the clusters that are used both for the training of deep learning models and the running of numerical solvers, <strong>network speed is generally two orders faster than I/O</strong>. For instance, if we consider an SSD device with a random read speed of 200,000 IOPS and a block size of 4 KB, we will get a throughput of 0.8 GB/s. In comparison, regarding the network, InfiniBand individual signal rate can provide a throughput up to 25 GB/s. Besides, on these clusters, <strong>CPU hours are cheaper than storage</strong>.</p> <p>It is thus promising to move from a traditionally offline training pipeline to an online configuration, where synthetic data are generated along the training. Multiple instances of the numerical solver are executed in parallel. Generated data are not stored on disk anymore, which circumvents I/O slowness. Instead, they are directly streamed through network for training. Training thus becomes faster. Because it needs less storage it is also cheaper. And because it allows using more data for the same training budget, it is also better.</p> <div class="row mt-3"> <div class="col-sm mt-3 mt-md-0"> <figure> <picture> <source class="responsive-img-srcset" srcset="/assets/img/online_training.png" sizes="95vw"/> <img src="/assets/img/online_training.png" class="img-fluid rounded z-depth-1" width="100%" height="auto" loading="eager" onerror="this.onerror=null; $('.responsive-img-srcset').remove();"/> </picture> </figure> </div> </div> <div class="caption"> <b>Online Training Pipeline</b>: Dataset is generated on-the-fly. Data are streamed for training as soon as they are generated. It circumvents expensive I/O, allowing for faster training on more diverse dataset. </div> <h1 id="an-online-training-framework">An Online Training Framework</h1> <p>To expose the design of such an online training framework, we introduce the following elements:</p> <ul> <li><strong>Clients</strong>: Each client runs an instance of the oracle the user provides for a different set of parameters. The oracle is a program that can be written, but not necessarily, in Python. For instance, numerical solvers are often C or Fortran programs that run on multiple processes using MPI.</li> <li><strong>Runner</strong>: The runner manages the executions of the clients. It triggers and monitors oracle executions on the clients. The runner serves as an orchestrator or a job scheduler. However, we avoid these terms to prevent any confusion with job schedulers commonly found on clusters (e.g. Slurm). The runner can nonetheless rely on such schedulers.</li> <li><strong>Server</strong>: A server serves data to the training loop. It receives and consumes the data generated by the clients. It typically corresponds to the dataset of the deep learning training pipeline. There can be multiple servers in case of data distributed parallelism.</li> </ul> <div class="row mt-3"> <div class="col-sm mt-3 mt-md-0"> <figure> <picture> <source class="responsive-img-srcset" srcset="/assets/img/framework_elements.png" sizes="95vw"/> <img src="/assets/img/framework_elements.png" class="img-fluid rounded z-depth-1" width="100%" height="auto" loading="eager" onerror="this.onerror=null; $('.responsive-img-srcset').remove();"/> </picture> </figure> </div> </div> <div class="caption"> <b>Interaction between Framework's Elements</b> </div> <p>There already exist some attempts to implement this kind of online training framework (e.g. <a href="https://joss.theoj.org/papers/10.21105/joss.05291">Melissa</a>). However, these frameworks require advance HPC knowledge and adds a significant overlay to PyTorch, which limits their generic use and thus adoption.</p> <h1 id="for-a-smooth-integration-to-pytorch">For a Smooth Integration to Pytorch</h1> <p>The online framework can be reduced to a data generation tool. The actual training loop of the model should be delegated to common deep learning libraries specifically designed for that (e.g. PyTorch or JAX). The framework should incur only minimal changes to the training loop. A good way to do so is by providing a dataset class that interfaces smoothly with these libraries. For instance, in the case of Pytorch, the framework should offer a class inheriting from the <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset"><code class="language-plaintext highlighter-rouge">IterableDataset</code></a>.</p> <p>In fact, in Pytorch, we can already obtain something similar with a <code class="language-plaintext highlighter-rouge">DataLoader</code> of multiple workers in conjunction of an <code class="language-plaintext highlighter-rouge">IterableDataset</code> object that would execute the oracle instances in the <code class="language-plaintext highlighter-rouge">__iter__</code> method. This is nonetheless limited. There will be as many instances of the oracle running in parallel as the number of workers of the data loader. This number will be bounded by the number of CPUs associated to each GPU node. It does not take advantage of CPUs that would be available on other nodes of the cluster.</p> <div class="language-python highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
6
7
8
9
10
</pre></td><td class="rouge-code"><pre><span class="k">class</span> <span class="nc">Dataset</span><span class="p">(</span><span class="n">IterableDataset</span><span class="p">):</span>

    <span class="k">def</span> <span class="nf">__iter__</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="c1"># Select parameters depending on the worker info
</span>        <span class="c1"># Call the oracle with some parameters
</span>        <span class="k">yield</span> <span class="k">from</span> <span class="nf">oracle</span><span class="p">(</span><span class="n">parameters</span><span class="p">)</span>

<span class="n">dataset</span> <span class="o">=</span> <span class="nc">Dataset</span><span class="p">()</span>
<span class="n">n_cpus</span> <span class="o">=</span> <span class="p">...</span> <span class="c1"># Number of CPUs associated to the node
</span><span class="n">dataloader</span> <span class="o">=</span> <span class="nc">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="n">n_cpus</span><span class="p">)</span>
</pre></td></tr></tbody></table></code></pre></div></div> <p>Moreover, whenever we want to guarantee <em>fault tolerance</em>, as we expect some clients to fail due to hardware failure which regularly occurs on clusters, or provide <em>elasticity</em> for evolving resource availability, we will need finer control on the execution of the oracle. This can only be achieved by relying on a dedicated <em>Runner</em>. Several packages like <a href="https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool"><code class="language-plaintext highlighter-rouge">multiprocessing.Pool</code></a>, <a href="https://docs.dask.org/en/stable/"><code class="language-plaintext highlighter-rouge">dask</code></a>, or even <a href="https://github.com/facebookincubator/submitit/"><code class="language-plaintext highlighter-rouge">submitit</code></a> can then be used for the Runner.</p> <h1 id="an-event-driven-architecture">An Event-driven Architecture</h1> <p>By designing the framework as a set of clients, servers, and runner interacting altogether asynchronously, we are actually defining <strong>an event-driven application</strong>. Indeed, upon submission of new parameters by the server to the runner, the later must trigger the execution of a new oracle instance given the computational resources allocated. Upon generation of the data by the different oracle instances, they are streamed to the server, which in turns accumulate these data and yield batches. All the dynamics of the framework can be expressed as occurring events that must trigger a response. Even <em>fault tolerance</em> feature can be expressed as a monitoring of failing oracle instances that informs the server of the issue.</p> <p>Even though the event-driven architecture pattern seems fairly common and good libraries with probably similar approach exist (e.g. <code class="language-plaintext highlighter-rouge">dask</code> which uses a <a href="https://www.tornadoweb.org/en/stable/ioloop.html"><code class="language-plaintext highlighter-rouge">tornado.IOLoop</code></a>), I didn’t find a clear and minimal example on how to build such architectures in Python. This kind of examples would help to avoid thread deadlocks and difficult maintenance, which are common issues for asynchronous applications.</p> <blockquote> <p>What is a good approach for event-driven applications in Python that are not prone to deadlocks and easy to maintain?</p> </blockquote> <h1 id="a-minimal-example">A Minimal Example</h1> <p>In this section we present a minimal example that work locally. It only needs few modifications for running at scale on a cluster. Only the content of the <code class="language-plaintext highlighter-rouge">__main__</code> functions below would require modification in a real application.</p> <p>To set up a Python environment to run the example, install the following packages. The dependencies include:</p> <ul> <li><code class="language-plaintext highlighter-rouge">numpy</code> and <code class="language-plaintext highlighter-rouge">torch</code>;</li> <li><a href="https://pyzmq.readthedocs.io/en/latest/"><code class="language-plaintext highlighter-rouge">zmq</code></a> for communication between runner, servers, and clients;</li> <li><code class="language-plaintext highlighter-rouge">tornado</code> to handle asynchronous events;</li> <li><code class="language-plaintext highlighter-rouge">submitit</code> to submit jobs to the clients.</li> </ul> <div class="language-bash highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
</pre></td><td class="rouge-code"><pre>conda create <span class="nt">-n</span> demo <span class="nv">python</span><span class="o">=</span>3.12
conda activate demo
pip <span class="nb">install </span>numpy torch zmq tornado submitit
<span class="c"># MPI Python bindings must be installed along the proper MPI binaries</span>
<span class="nv">CC</span><span class="o">=</span>mpicc pip <span class="nb">install</span> <span class="nt">--no-cache</span> mpi4py 
</pre></td></tr></tbody></table></code></pre></div></div> <h2 id="utility-functions">Utility Functions</h2> <p>First, we define some utility functions to format commands passed from the server to the runner and encapsulate signals sent between the different components.</p> <details> <summary>utils.py<br/><br/></summary> <div class="language-python highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
</pre></td><td class="rouge-code"><pre><span class="kn">import</span> <span class="n">os</span>
<span class="kn">import</span> <span class="n">pickle</span>
<span class="kn">import</span> <span class="n">shlex</span>
<span class="kn">from</span> <span class="n">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span>
<span class="kn">from</span> <span class="n">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="kn">from</span> <span class="n">functools</span> <span class="kn">import</span> <span class="n">wraps</span>
<span class="kn">from</span> <span class="n">typing</span> <span class="kn">import</span> <span class="n">List</span>

<span class="kn">import</span> <span class="n">numpy</span> <span class="k">as</span> <span class="n">np</span>


<span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">Task</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">A task to execute by a client. Generated data must be send to the server address.</span><span class="sh">"""</span>

    <span class="n">command</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>
    <span class="n">server_address</span><span class="p">:</span> <span class="nb">str</span>


<span class="k">def</span> <span class="nf">format_command</span><span class="p">(</span><span class="n">command</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
    <span class="sh">"""</span><span class="s">Format command for execution by the client.</span><span class="sh">"""</span>
    <span class="n">full_command</span> <span class="o">=</span> <span class="p">(</span>
        <span class="n">command</span>
        <span class="o">+</span> <span class="sh">"</span><span class="s"> </span><span class="sh">"</span><span class="p">.</span><span class="nf">join</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>
        <span class="o">+</span> <span class="sh">"</span><span class="s"> </span><span class="sh">"</span><span class="p">.</span><span class="nf">join</span><span class="p">([</span><span class="sa">f</span><span class="sh">"</span><span class="si">{</span><span class="n">key</span><span class="si">}</span><span class="s">=</span><span class="si">{</span><span class="n">val</span><span class="si">}</span><span class="sh">"</span> <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="p">.</span><span class="nf">items</span><span class="p">()])</span>
    <span class="p">)</span>
    <span class="n">split_command</span> <span class="o">=</span> <span class="n">shlex</span><span class="p">.</span><span class="nf">split</span><span class="p">(</span><span class="n">full_command</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">split_command</span>


<span class="k">def</span> <span class="nf">deserialize</span><span class="p">(</span><span class="n">msg</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">bytes</span><span class="p">]):</span>
    <span class="k">return</span> <span class="n">pickle</span><span class="p">.</span><span class="nf">loads</span><span class="p">(</span><span class="n">msg</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>


<span class="k">class</span> <span class="nc">Status</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
    <span class="sh">"""</span><span class="s">Status of running process.</span><span class="sh">"""</span>

    <span class="n">START</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
    <span class="n">SUCCESS</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="n">FAIL</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="n">READY</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="n">FINISH</span> <span class="o">=</span> <span class="mi">3</span>


<span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">Signal</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Signal sent about clients to communicate about their status.</span><span class="sh">"""</span>

    <span class="n">status</span><span class="p">:</span> <span class="n">Status</span>
    <span class="n">client_id</span><span class="p">:</span> <span class="nb">int</span>


<span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">ClientData</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Numerical solvers used to generate the data
    typically produce data for different time steps.</span><span class="sh">"""</span>

    <span class="n">data</span><span class="p">:</span> <span class="n">np</span><span class="p">.</span><span class="n">ndarray</span>
    <span class="n">job_id</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="bp">None</span>
    <span class="n">step_id</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="bp">None</span>

    <span class="k">def</span> <span class="nf">__repr__</span><span class="p">(</span><span class="n">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
        <span class="k">return</span> <span class="sa">f</span><span class="sh">"</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">__class__</span><span class="p">.</span><span class="n">__name__</span><span class="si">}</span><span class="s">(job_id=</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">job_id</span><span class="si">}</span><span class="s">, step_id=</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">step_id</span><span class="si">}</span><span class="s">, data_shape=</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">data</span><span class="p">.</span><span class="n">shape</span><span class="si">}</span><span class="s">)</span><span class="sh">"</span>


<span class="k">def</span> <span class="nf">get_rank</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Get the local rank for the client.
    A client may run on several processes.</span><span class="sh">"""</span>
    <span class="c1"># For Slurm
</span>    <span class="c1"># Get rank from environment variables
</span>    <span class="n">cpus_per_task</span> <span class="o">=</span> <span class="nf">int</span><span class="p">(</span><span class="n">os</span><span class="p">.</span><span class="n">environ</span><span class="p">.</span><span class="nf">get</span><span class="p">(</span><span class="sh">"</span><span class="s">SLURM_CPUS_PER_TASK</span><span class="sh">"</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
    <span class="k">if</span> <span class="n">cpus_per_task</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">return</span> <span class="mi">0</span>

    <span class="n">rank_keys</span> <span class="o">=</span> <span class="p">(</span><span class="sh">"</span><span class="s">RANK</span><span class="sh">"</span><span class="p">,</span> <span class="sh">"</span><span class="s">LOCAL_RANK</span><span class="sh">"</span><span class="p">,</span> <span class="sh">"</span><span class="s">SLURM_PROCID</span><span class="sh">"</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">rank_keys</span><span class="p">:</span>
        <span class="n">rank</span> <span class="o">=</span> <span class="n">os</span><span class="p">.</span><span class="n">environ</span><span class="p">.</span><span class="nf">get</span><span class="p">(</span><span class="n">key</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">rank</span> <span class="ow">is</span> <span class="ow">not</span> <span class="bp">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="nf">int</span><span class="p">(</span><span class="n">rank</span><span class="p">)</span>

    <span class="c1"># For MPI
</span>    <span class="kn">from</span> <span class="n">mpi4py</span> <span class="kn">import</span> <span class="n">MPI</span>

    <span class="n">comm</span> <span class="o">=</span> <span class="n">MPI</span><span class="p">.</span><span class="n">COMM_WORLD</span>
    <span class="n">rank</span> <span class="o">=</span> <span class="n">comm</span><span class="p">.</span><span class="nc">Get_rank</span><span class="p">()</span>
    <span class="k">return</span> <span class="nf">int</span><span class="p">(</span><span class="n">rank</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">on_rank_zero_only</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
    <span class="nd">@wraps</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
    <span class="k">def</span> <span class="nf">wrapper</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="n">rank</span> <span class="o">=</span> <span class="nf">get_rank</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">rank</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">return</span> <span class="nf">fn</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">wrapper</span>

</pre></td></tr></tbody></table></code></pre></div> </div> </details> <h2 id="clients">Clients</h2> <p>For the client, we create a dummy example that generates random arrays for 20 time steps. In practice, the client can be any numerical simulation program. To work along the framework, the program must be edited to send the generated data over the network instead of saving them on disk, as it would be generally done. To do so, the program can use an API that is managed by the <code class="language-plaintext highlighter-rouge">ClientCommunicator</code> in the example below. This API does three things:</p> <ul> <li>initiate the communication between the client and the targeted servers (<code class="language-plaintext highlighter-rouge">__init__</code> method of the communicator);</li> <li>send data as they are produced through ZMQ socket (<code class="language-plaintext highlighter-rouge">send_array</code> method);</li> <li>close the communication and signal the runner the client has terminated successfully (<code class="language-plaintext highlighter-rouge">terminate</code> method).</li> </ul> <p>Bindings to the API can be easily provided for Fortan, C, and C++ code that are generally used to write parallel numerical solvers.</p> <details> <summary>client.py<br/><br/></summary> <div class="language-python highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
</pre></td><td class="rouge-code"><pre><span class="kn">import</span> <span class="n">logging</span>
<span class="kn">import</span> <span class="n">os</span>
<span class="kn">import</span> <span class="n">random</span>
<span class="kn">import</span> <span class="n">socket</span>
<span class="kn">import</span> <span class="n">time</span>
<span class="kn">from</span> <span class="n">typing</span> <span class="kn">import</span> <span class="n">Optional</span>

<span class="kn">import</span> <span class="n">numpy</span> <span class="k">as</span> <span class="n">np</span>
<span class="kn">import</span> <span class="n">zmq</span>

<span class="kn">from</span> <span class="n">utils</span> <span class="kn">import</span> <span class="n">ClientData</span><span class="p">,</span> <span class="n">Signal</span><span class="p">,</span> <span class="n">Status</span><span class="p">,</span> <span class="n">get_rank</span><span class="p">,</span> <span class="n">on_rank_zero_only</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="p">.</span><span class="nf">getLogger</span><span class="p">(</span><span class="n">__name__</span><span class="p">)</span>
<span class="n">logging</span><span class="p">.</span><span class="nf">basicConfig</span><span class="p">(</span><span class="n">level</span><span class="o">=</span><span class="n">logging</span><span class="p">.</span><span class="n">DEBUG</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">send_array</span><span class="p">(</span><span class="n">socket</span><span class="p">:</span> <span class="n">zmq</span><span class="p">.</span><span class="n">Socket</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="n">ClientData</span><span class="p">):</span>
    <span class="c1"># TODO: Check if serializing is needed for performances
</span>    <span class="c1"># C.f. PyZMQ doc:
</span>    <span class="c1"># https://pyzmq.readthedocs.io/en/latest/howto/serialization.html#serializing-messages-with-pyzmq
</span>    <span class="n">socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">ClientCommunicator</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Class to be used by the client to send data and termination signal.</span><span class="sh">"""</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="n">self</span><span class="p">.</span><span class="n">hostname</span> <span class="o">=</span> <span class="n">socket</span><span class="p">.</span><span class="nf">gethostbyname</span><span class="p">(</span><span class="n">socket</span><span class="p">.</span><span class="nf">gethostname</span><span class="p">())</span>
        <span class="n">self</span><span class="p">.</span><span class="n">rank</span> <span class="o">=</span> <span class="nf">get_rank</span><span class="p">()</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">info</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Start </span><span class="si">{</span><span class="n">self</span><span class="si">}</span><span class="s"> on </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">hostname</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>

        <span class="n">self</span><span class="p">.</span><span class="n">job_id</span> <span class="o">=</span> <span class="n">os</span><span class="p">.</span><span class="n">environ</span><span class="p">.</span><span class="nf">get</span><span class="p">(</span><span class="sh">"</span><span class="s">JOB_ID</span><span class="sh">"</span><span class="p">,</span> <span class="bp">None</span><span class="p">)</span>
        <span class="n">socket_addr</span> <span class="o">=</span> <span class="n">os</span><span class="p">.</span><span class="n">environ</span><span class="p">.</span><span class="nf">get</span><span class="p">(</span><span class="sh">"</span><span class="s">DATA_ADDR</span><span class="sh">"</span><span class="p">)</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span>
            <span class="sa">f</span><span class="sh">"</span><span class="s">Instantiate communicator on rank </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">rank</span><span class="si">}</span><span class="s"> of client </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">job_id</span><span class="si">}</span><span class="s">.</span><span class="sh">"</span>
        <span class="p">)</span>
        <span class="k">if</span> <span class="n">self</span><span class="p">.</span><span class="n">rank</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">context</span> <span class="o">=</span> <span class="n">zmq</span><span class="p">.</span><span class="n">Context</span><span class="p">.</span><span class="nf">instance</span><span class="p">()</span>
            <span class="n">self</span><span class="p">.</span><span class="n">socket</span> <span class="o">=</span> <span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">PUSH</span><span class="p">)</span>
            <span class="n">self</span><span class="p">.</span><span class="n">socket</span><span class="p">.</span><span class="nf">connect</span><span class="p">(</span><span class="n">socket_addr</span><span class="p">)</span>

    <span class="nd">@on_rank_zero_only</span>
    <span class="k">def</span> <span class="nf">send_array</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="n">np</span><span class="p">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">step_id</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="bp">None</span><span class="p">):</span>
        <span class="n">client_data</span> <span class="o">=</span> <span class="nc">ClientData</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">self</span><span class="p">.</span><span class="n">job_id</span><span class="p">,</span> <span class="n">step_id</span><span class="p">)</span>
        <span class="nf">send_array</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">socket</span><span class="p">,</span> <span class="n">client_data</span><span class="p">)</span>

    <span class="nd">@on_rank_zero_only</span>
    <span class="k">def</span> <span class="nf">termniate</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="n">signal</span> <span class="o">=</span> <span class="nc">Signal</span><span class="p">(</span><span class="n">Status</span><span class="p">.</span><span class="n">SUCCESS</span><span class="p">,</span> <span class="n">self</span><span class="p">.</span><span class="n">job_id</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">signal</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">socket</span><span class="p">.</span><span class="nf">close</span><span class="p">()</span>


<span class="k">def</span> <span class="nf">main</span><span class="p">(</span><span class="n">steps</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">20</span><span class="p">):</span>
    <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Client initiates communicator.</span><span class="sh">"</span><span class="p">)</span>
    <span class="n">communicator</span> <span class="o">=</span> <span class="nc">ClientCommunicator</span><span class="p">()</span>
    <span class="k">for</span> <span class="n">step</span> <span class="ow">in</span> <span class="nf">range</span><span class="p">(</span><span class="n">steps</span><span class="p">):</span>
        <span class="n">time</span><span class="p">.</span><span class="nf">sleep</span><span class="p">(</span><span class="n">random</span><span class="p">.</span><span class="nf">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">500</span><span class="p">)</span> <span class="o">/</span> <span class="mi">1_000</span><span class="p">)</span>
        <span class="n">data</span> <span class="o">=</span> <span class="n">np</span><span class="p">.</span><span class="n">random</span><span class="p">.</span><span class="nf">rand</span><span class="p">(</span><span class="mi">256</span><span class="p">,</span> <span class="mi">256</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="n">communicator</span><span class="p">.</span><span class="nf">send_array</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">step</span><span class="p">)</span>
    <span class="n">communicator</span><span class="p">.</span><span class="nf">termniate</span><span class="p">()</span>
    <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Client terminates.</span><span class="sh">"</span><span class="p">)</span>


<span class="k">if</span> <span class="n">__name__</span> <span class="o">==</span> <span class="sh">"</span><span class="s">__main__</span><span class="sh">"</span><span class="p">:</span>
    <span class="nf">main</span><span class="p">()</span>

</pre></td></tr></tbody></table></code></pre></div> </div> </details> <h2 id="runner">Runner</h2> <p>The runner starts by waiting a signal from the server for synchronization. Then it launches commands it receives from the synchronized server.</p> <p>In the example, the runner relies on <code class="language-plaintext highlighter-rouge">submitit</code> to launch clients locally. <code class="language-plaintext highlighter-rouge">submitit</code> can normally be used to submit Slurm jobs. Other options exist to launch clients within an already granted Slurm allocation.</p> <details> <summary>runner.py<br/><br/></summary> <div class="language-python highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
</pre></td><td class="rouge-code"><pre>
<span class="kn">import</span> <span class="n">logging</span>
<span class="kn">import</span> <span class="n">os</span>
<span class="kn">import</span> <span class="n">socket</span>
<span class="kn">import</span> <span class="n">threading</span>
<span class="kn">from</span> <span class="n">queue</span> <span class="kn">import</span> <span class="n">Empty</span><span class="p">,</span> <span class="n">Queue</span>
<span class="kn">from</span> <span class="n">typing</span> <span class="kn">import</span> <span class="n">Iterator</span>

<span class="kn">import</span> <span class="n">submitit</span>
<span class="kn">import</span> <span class="n">zmq</span>
<span class="kn">import</span> <span class="n">zmq.asyncio</span>
<span class="kn">from</span> <span class="n">tornado.ioloop</span> <span class="kn">import</span> <span class="n">IOLoop</span>

<span class="kn">from</span> <span class="n">utils</span> <span class="kn">import</span> <span class="n">Status</span><span class="p">,</span> <span class="n">Task</span><span class="p">,</span> <span class="n">get_rank</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="p">.</span><span class="nf">getLogger</span><span class="p">(</span><span class="n">__name__</span><span class="p">)</span>
<span class="n">logging</span><span class="p">.</span><span class="nf">basicConfig</span><span class="p">(</span><span class="n">level</span><span class="o">=</span><span class="n">logging</span><span class="p">.</span><span class="n">DEBUG</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">RunnerCommunicator</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Class to be used by the runner to send signals about client status.</span><span class="sh">"""</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="n">self</span><span class="p">.</span><span class="n">hostname</span> <span class="o">=</span> <span class="n">socket</span><span class="p">.</span><span class="nf">gethostbyname</span><span class="p">(</span><span class="n">socket</span><span class="p">.</span><span class="nf">gethostname</span><span class="p">())</span>
        <span class="n">self</span><span class="p">.</span><span class="n">rank</span> <span class="o">=</span> <span class="nf">get_rank</span><span class="p">()</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">info</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Start </span><span class="si">{</span><span class="n">self</span><span class="si">}</span><span class="s"> on </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">hostname</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">servers</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>

        <span class="c1"># Communication variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">context</span> <span class="o">=</span> <span class="n">zmq</span><span class="p">.</span><span class="n">asyncio</span><span class="p">.</span><span class="n">Context</span><span class="p">.</span><span class="nf">instance</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">synchronization_socket</span> <span class="o">=</span> <span class="n">self</span><span class="p">.</span><span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">REP</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">synchronization_socket</span><span class="p">.</span><span class="nf">bind</span><span class="p">(</span><span class="sh">"</span><span class="s">tcp://*:5558</span><span class="sh">"</span><span class="p">)</span>
        <span class="c1"># Event loop variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span> <span class="o">=</span> <span class="n">IOLoop</span><span class="p">.</span><span class="nf">instance</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">add_callback</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">_handle_task</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">add_callback</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">_synchronize</span><span class="p">)</span>
        <span class="c1"># Thread variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">lock</span> <span class="o">=</span> <span class="n">threading</span><span class="p">.</span><span class="nc">Lock</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span> <span class="o">=</span> <span class="bp">True</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_tasks</span> <span class="o">=</span> <span class="nc">Queue</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span> <span class="o">=</span> <span class="n">threading</span><span class="p">.</span><span class="nc">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">self</span><span class="p">.</span><span class="n">_run_ioloop</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span><span class="p">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="bp">True</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span><span class="p">.</span><span class="nf">start</span><span class="p">()</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">is_receiving</span><span class="p">(</span><span class="n">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="k">with</span> <span class="n">self</span><span class="p">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span>

    <span class="nd">@is_receiving.setter</span>
    <span class="k">def</span> <span class="nf">is_receiving</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">status</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
        <span class="k">with</span> <span class="n">self</span><span class="p">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span> <span class="o">=</span> <span class="n">status</span>

    <span class="k">def</span> <span class="nf">_run_ioloop</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">start</span><span class="p">()</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_synchronize</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="sh">"""</span><span class="s">Synchronize with a Server</span><span class="sh">"""</span>
        <span class="k">while</span> <span class="bp">True</span><span class="p">:</span>
            <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Waiting for synchronization...</span><span class="sh">"</span><span class="p">)</span>
            <span class="p">(</span>
                <span class="n">server_hostname</span><span class="p">,</span>
                <span class="n">server_data_port</span><span class="p">,</span>
            <span class="p">)</span> <span class="o">=</span> <span class="k">await</span> <span class="n">self</span><span class="p">.</span><span class="n">synchronization_socket</span><span class="p">.</span><span class="nf">recv_pyobj</span><span class="p">()</span>
            <span class="n">self</span><span class="p">.</span><span class="n">servers</span><span class="p">.</span><span class="nf">append</span><span class="p">((</span><span class="n">server_hostname</span><span class="p">,</span> <span class="n">server_data_port</span><span class="p">))</span>
            <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span> <span class="o">=</span> <span class="bp">True</span>
            <span class="k">await</span> <span class="n">self</span><span class="p">.</span><span class="n">synchronization_socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">Status</span><span class="p">.</span><span class="n">SUCCESS</span><span class="p">)</span>
            <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Synchronized with </span><span class="si">{</span><span class="n">server_hostname</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_handle_task</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="n">task_socket</span> <span class="o">=</span> <span class="n">self</span><span class="p">.</span><span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">PULL</span><span class="p">)</span>
        <span class="n">task_socket</span><span class="p">.</span><span class="nf">bind</span><span class="p">(</span><span class="sh">"</span><span class="s">tcp://*:5559</span><span class="sh">"</span><span class="p">)</span>
        <span class="n">poller</span> <span class="o">=</span> <span class="n">zmq</span><span class="p">.</span><span class="n">asyncio</span><span class="p">.</span><span class="nc">Poller</span><span class="p">()</span>
        <span class="n">poller</span><span class="p">.</span><span class="nf">register</span><span class="p">(</span><span class="n">task_socket</span><span class="p">)</span>

        <span class="n">timeout</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
        <span class="k">while</span> <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span><span class="p">:</span>
            <span class="n">events</span> <span class="o">=</span> <span class="k">await</span> <span class="n">poller</span><span class="p">.</span><span class="nf">poll</span><span class="p">(</span><span class="n">timeout</span><span class="p">)</span>
            <span class="n">events</span> <span class="o">=</span> <span class="nf">dict</span><span class="p">(</span><span class="n">events</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">task_socket</span> <span class="ow">in</span> <span class="n">events</span><span class="p">:</span>
                <span class="n">msg</span> <span class="o">=</span> <span class="k">await</span> <span class="n">task_socket</span><span class="p">.</span><span class="nf">recv_pyobj</span><span class="p">()</span>
                <span class="k">if</span> <span class="nf">isinstance</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="n">Task</span><span class="p">):</span>
                    <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Receive task </span><span class="si">{</span><span class="n">msg</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
                    <span class="n">self</span><span class="p">.</span><span class="n">_tasks</span><span class="p">.</span><span class="nf">put_nowait</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
                <span class="c1"># TODO: Put message in queue to avoid setting self.is_receiving
</span>                <span class="k">elif</span> <span class="nf">isinstance</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="n">Status</span><span class="p">)</span> <span class="ow">and</span> <span class="n">msg</span> <span class="o">==</span> <span class="n">Status</span><span class="p">.</span><span class="n">FINISH</span><span class="p">:</span>
                    <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Receive submission termination signal</span><span class="sh">"</span><span class="p">)</span>
                    <span class="n">timeout</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span> <span class="o">=</span> <span class="bp">False</span>
                <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">stop</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">tasks</span><span class="p">(</span><span class="n">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Iterator</span><span class="p">[</span><span class="n">Task</span><span class="p">]:</span>
        <span class="sh">"""</span><span class="s">Iterate over the tasks received from a server submission.</span><span class="sh">"""</span>
        <span class="k">while</span> <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span> <span class="ow">or</span> <span class="n">self</span><span class="p">.</span><span class="n">_tasks</span><span class="p">.</span><span class="nf">qsize</span><span class="p">():</span>
            <span class="c1"># TODO: Retrieve termination signal from queue
</span>            <span class="k">try</span><span class="p">:</span>
                <span class="k">yield</span> <span class="n">self</span><span class="p">.</span><span class="n">_tasks</span><span class="p">.</span><span class="nf">get_nowait</span><span class="p">()</span>
            <span class="k">except</span> <span class="n">Empty</span><span class="p">:</span>
                <span class="k">continue</span>

    <span class="k">def</span> <span class="nf">terminate</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="k">with</span> <span class="n">zmq</span><span class="p">.</span><span class="nc">Context</span><span class="p">()</span> <span class="k">as</span> <span class="n">ctx</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">server</span> <span class="ow">in</span> <span class="n">self</span><span class="p">.</span><span class="n">servers</span><span class="p">:</span>
                <span class="n">hostname</span><span class="p">,</span> <span class="n">port</span> <span class="o">=</span> <span class="n">server</span>
                <span class="k">with</span> <span class="n">ctx</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">PUSH</span><span class="p">)</span> <span class="k">as</span> <span class="n">socket</span><span class="p">:</span>
                    <span class="n">socket</span><span class="p">.</span><span class="nf">connect</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">tcp://</span><span class="si">{</span><span class="n">hostname</span><span class="si">}</span><span class="s">:</span><span class="si">{</span><span class="n">port</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
                    <span class="n">socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">Status</span><span class="p">.</span><span class="n">SUCCESS</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">stop</span><span class="p">()</span>


<span class="k">if</span> <span class="n">__name__</span> <span class="o">==</span> <span class="sh">"</span><span class="s">__main__</span><span class="sh">"</span><span class="p">:</span>
    <span class="n">executor</span> <span class="o">=</span> <span class="n">submitit</span><span class="p">.</span><span class="nc">AutoExecutor</span><span class="p">(</span><span class="n">folder</span><span class="o">=</span><span class="sh">"</span><span class="s">log_test</span><span class="sh">"</span><span class="p">)</span>
    <span class="n">jobs</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="n">communicator</span> <span class="o">=</span> <span class="nc">RunnerCommunicator</span><span class="p">()</span>
    <span class="k">for</span> <span class="n">task_id</span><span class="p">,</span> <span class="n">task</span> <span class="ow">in</span> <span class="nf">enumerate</span><span class="p">(</span><span class="n">communicator</span><span class="p">.</span><span class="nf">tasks</span><span class="p">()):</span>
        <span class="n">env</span> <span class="o">=</span> <span class="n">os</span><span class="p">.</span><span class="n">environ</span><span class="p">.</span><span class="nf">copy</span><span class="p">()</span>
        <span class="n">env</span><span class="p">.</span><span class="nf">update</span><span class="p">({</span><span class="sh">"</span><span class="s">DATA_ADDR</span><span class="sh">"</span><span class="p">:</span> <span class="n">task</span><span class="p">.</span><span class="n">server_address</span><span class="p">,</span> <span class="sh">"</span><span class="s">JOB_ID</span><span class="sh">"</span><span class="p">:</span> <span class="nf">str</span><span class="p">(</span><span class="n">task_id</span><span class="p">)})</span>
        <span class="n">function</span> <span class="o">=</span> <span class="n">submitit</span><span class="p">.</span><span class="n">helpers</span><span class="p">.</span><span class="nc">CommandFunction</span><span class="p">(</span><span class="n">task</span><span class="p">.</span><span class="n">command</span><span class="p">,</span> <span class="n">env</span><span class="o">=</span><span class="n">env</span><span class="p">)</span>
        <span class="n">job</span> <span class="o">=</span> <span class="n">executor</span><span class="p">.</span><span class="nf">submit</span><span class="p">(</span><span class="n">function</span><span class="p">)</span>
        <span class="n">jobs</span><span class="p">.</span><span class="nf">append</span><span class="p">(</span><span class="n">job</span><span class="p">)</span>
    <span class="nf">print</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Results for </span><span class="si">{</span><span class="nf">len</span><span class="p">(</span><span class="n">jobs</span><span class="p">)</span><span class="si">}</span><span class="s"> jobs:</span><span class="sh">"</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">job</span> <span class="ow">in</span> <span class="n">submitit</span><span class="p">.</span><span class="n">helpers</span><span class="p">.</span><span class="nf">as_completed</span><span class="p">(</span><span class="n">jobs</span><span class="p">):</span>
        <span class="nf">print</span><span class="p">(</span><span class="n">job</span><span class="p">.</span><span class="nf">result</span><span class="p">())</span>
    <span class="n">communicator</span><span class="p">.</span><span class="nf">terminate</span><span class="p">()</span>
</pre></td></tr></tbody></table></code></pre></div> </div> </details> <h2 id="server">Server</h2> <p>The server first synchronizes with the runner. Then, it submits command to be run by the clients to generate the synthetic data. Finally, it receives these data from the clients. Whenever the clients have completed, the runner signals it to the server for it to know no more data are to expect.</p> <details> <summary>server.py<br/><br/></summary> <div class="language-python highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
</pre></td><td class="rouge-code"><pre><span class="kn">import</span> <span class="n">argparse</span>
<span class="kn">import</span> <span class="n">asyncio</span>
<span class="kn">import</span> <span class="n">logging</span>
<span class="kn">import</span> <span class="n">os</span>
<span class="kn">import</span> <span class="n">socket</span>
<span class="kn">import</span> <span class="n">threading</span>
<span class="kn">from</span> <span class="n">queue</span> <span class="kn">import</span> <span class="n">Empty</span><span class="p">,</span> <span class="n">Queue</span>
<span class="kn">from</span> <span class="n">typing</span> <span class="kn">import</span> <span class="n">List</span>

<span class="kn">import</span> <span class="n">zmq</span>
<span class="kn">import</span> <span class="n">zmq.asyncio</span>
<span class="kn">from</span> <span class="n">tornado.ioloop</span> <span class="kn">import</span> <span class="n">IOLoop</span>
<span class="kn">from</span> <span class="n">torch.utils.data</span> <span class="kn">import</span> <span class="n">IterableDataset</span><span class="p">,</span> <span class="n">DataLoader</span>

<span class="kn">from</span> <span class="n">utils</span> <span class="kn">import</span> <span class="n">ClientData</span><span class="p">,</span> <span class="n">Status</span><span class="p">,</span> <span class="n">Task</span><span class="p">,</span> <span class="n">format_command</span><span class="p">,</span> <span class="n">get_rank</span>
<span class="kn">import</span> <span class="n">client</span> <span class="k">as</span> <span class="n">dummy_client</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="p">.</span><span class="nf">getLogger</span><span class="p">(</span><span class="n">__name__</span><span class="p">)</span>

<span class="c1"># Set static port number for communication purposes
</span><span class="n">SYNC_PORT</span> <span class="o">=</span> <span class="mi">5558</span>
<span class="n">TASK_PORT</span> <span class="o">=</span> <span class="mi">5559</span>


<span class="k">class</span> <span class="nc">ServerCommunicator</span><span class="p">:</span>
    <span class="sh">"""</span><span class="s">Class to be used by the server to receive and distribute client data.</span><span class="sh">"""</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">runner_hostname</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">self</span><span class="p">.</span><span class="n">hostname</span> <span class="o">=</span> <span class="n">socket</span><span class="p">.</span><span class="nf">gethostbyname</span><span class="p">(</span><span class="n">socket</span><span class="p">.</span><span class="nf">gethostname</span><span class="p">())</span>
        <span class="n">self</span><span class="p">.</span><span class="n">rank</span> <span class="o">=</span> <span class="nf">get_rank</span><span class="p">()</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">info</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Start </span><span class="si">{</span><span class="n">self</span><span class="si">}</span><span class="s"> on </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">hostname</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="c1"># Communication variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">runner_hostname</span> <span class="o">=</span> <span class="n">runner_hostname</span>
        <span class="n">self</span><span class="p">.</span><span class="n">context</span> <span class="o">=</span> <span class="n">zmq</span><span class="p">.</span><span class="n">asyncio</span><span class="p">.</span><span class="n">Context</span><span class="p">.</span><span class="nf">instance</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">data_port</span> <span class="o">=</span> <span class="nf">str</span><span class="p">(</span><span class="mi">5560</span> <span class="o">+</span> <span class="n">self</span><span class="p">.</span><span class="n">rank</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">task_socket</span> <span class="o">=</span> <span class="n">self</span><span class="p">.</span><span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">PUSH</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">task_socket</span><span class="p">.</span><span class="nf">connect</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">tcp://</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">runner_hostname</span><span class="si">}</span><span class="s">:</span><span class="si">{</span><span class="n">TASK_PORT</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="c1"># Event loop variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">_synchronized</span> <span class="o">=</span> <span class="n">asyncio</span><span class="p">.</span><span class="nc">Event</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span> <span class="o">=</span> <span class="n">IOLoop</span><span class="p">.</span><span class="nf">current</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">add_callback</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">_synchronize</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">add_callback</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">_receive</span><span class="p">)</span>
        <span class="c1"># Thread variables
</span>        <span class="n">self</span><span class="p">.</span><span class="n">_data</span> <span class="o">=</span> <span class="nc">Queue</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span> <span class="o">=</span> <span class="bp">True</span>
        <span class="n">self</span><span class="p">.</span><span class="n">lock</span> <span class="o">=</span> <span class="n">threading</span><span class="p">.</span><span class="nc">Lock</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span> <span class="o">=</span> <span class="n">threading</span><span class="p">.</span><span class="nc">Thread</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">self</span><span class="p">.</span><span class="n">_run_ioloop</span><span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span><span class="p">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="bp">True</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_thread</span><span class="p">.</span><span class="nf">start</span><span class="p">()</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">is_receiving</span><span class="p">(</span><span class="n">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="k">with</span> <span class="n">self</span><span class="p">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span>

    <span class="nd">@is_receiving.setter</span>
    <span class="k">def</span> <span class="nf">is_receiving</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">status</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
        <span class="k">with</span> <span class="n">self</span><span class="p">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="n">self</span><span class="p">.</span><span class="n">_is_receiving</span> <span class="o">=</span> <span class="n">status</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_synchronize</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="sh">"""</span><span class="s">Synchronize with the runner.</span><span class="sh">"""</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Waiting for synchronization...</span><span class="sh">"</span><span class="p">)</span>
        <span class="k">with</span> <span class="n">self</span><span class="p">.</span><span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">REQ</span><span class="p">)</span> <span class="k">as</span> <span class="n">runner_socket</span><span class="p">:</span>
            <span class="n">runner_socket</span><span class="p">.</span><span class="nf">connect</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">tcp://</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">runner_hostname</span><span class="si">}</span><span class="s">:</span><span class="si">{</span><span class="n">SYNC_PORT</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
            <span class="k">await</span> <span class="n">runner_socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">((</span><span class="n">self</span><span class="p">.</span><span class="n">hostname</span><span class="p">,</span> <span class="n">self</span><span class="p">.</span><span class="n">data_port</span><span class="p">))</span>
            <span class="n">synchronization_status</span><span class="p">:</span> <span class="n">Status</span> <span class="o">=</span> <span class="k">await</span> <span class="n">runner_socket</span><span class="p">.</span><span class="nf">recv_pyobj</span><span class="p">()</span>
            <span class="k">assert</span> <span class="n">synchronization_status</span> <span class="o">==</span> <span class="n">Status</span><span class="p">.</span><span class="n">SUCCESS</span>
            <span class="n">self</span><span class="p">.</span><span class="n">_synchronized</span><span class="p">.</span><span class="nf">set</span><span class="p">()</span>
            <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span> <span class="o">=</span> <span class="bp">True</span>
            <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Syncrhonized with </span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">runner_hostname</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_run_ioloop</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="sh">"""</span><span class="s">Method to start the tornado.IOLoop in a thread.</span><span class="sh">"""</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">start</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">submit</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">command</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]):</span>
        <span class="sh">"""</span><span class="s">Submit the command as a task to the runner.</span><span class="sh">"""</span>
        <span class="n">task</span> <span class="o">=</span> <span class="nc">Task</span><span class="p">(</span>
            <span class="n">command</span><span class="o">=</span><span class="n">command</span><span class="p">,</span> <span class="n">server_address</span><span class="o">=</span><span class="sa">f</span><span class="sh">"</span><span class="s">tcp://</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">hostname</span><span class="si">}</span><span class="s">:</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">data_port</span><span class="si">}</span><span class="sh">"</span>
        <span class="p">)</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">add_callback</span><span class="p">(</span><span class="n">self</span><span class="p">.</span><span class="n">_submit</span><span class="p">,</span> <span class="n">task</span><span class="p">)</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_submit</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">task</span><span class="p">:</span> <span class="n">Task</span><span class="p">):</span>
        <span class="k">await</span> <span class="n">self</span><span class="p">.</span><span class="n">_synchronized</span><span class="p">.</span><span class="nf">wait</span><span class="p">()</span>
        <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Submit task </span><span class="si">{</span><span class="n">task</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="k">await</span> <span class="n">self</span><span class="p">.</span><span class="n">task_socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">task</span><span class="p">)</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_terminate</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="k">await</span> <span class="n">self</span><span class="p">.</span><span class="n">task_socket</span><span class="p">.</span><span class="nf">send_pyobj</span><span class="p">(</span><span class="n">Status</span><span class="p">.</span><span class="n">FINISH</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">terminate</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="sh">"""</span><span class="s">Send signal to runner no more task will be submitted for the round.
        Delay the call to not shortcut current submissions.

        </span><span class="sh">"""</span>
        <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">call_later</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="n">self</span><span class="p">.</span><span class="n">_terminate</span><span class="p">)</span>

    <span class="k">async</span> <span class="k">def</span> <span class="nf">_receive</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="sh">"""</span><span class="s">Reception to be run in a thread.
        Listen to the sockets connected to the clients.

        </span><span class="sh">"""</span>
        <span class="n">client_data_socket</span> <span class="o">=</span> <span class="n">self</span><span class="p">.</span><span class="n">context</span><span class="p">.</span><span class="nf">socket</span><span class="p">(</span><span class="n">zmq</span><span class="p">.</span><span class="n">PULL</span><span class="p">)</span>
        <span class="n">client_data_socket</span><span class="p">.</span><span class="nf">bind</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">tcp://*:</span><span class="si">{</span><span class="n">self</span><span class="p">.</span><span class="n">data_port</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="n">poller</span> <span class="o">=</span> <span class="n">zmq</span><span class="p">.</span><span class="n">asyncio</span><span class="p">.</span><span class="nc">Poller</span><span class="p">()</span>
        <span class="n">poller</span><span class="p">.</span><span class="nf">register</span><span class="p">(</span><span class="n">client_data_socket</span><span class="p">)</span>

        <span class="n">timeout</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
        <span class="k">while</span> <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span><span class="p">:</span>
            <span class="n">events</span> <span class="o">=</span> <span class="k">await</span> <span class="n">poller</span><span class="p">.</span><span class="nf">poll</span><span class="p">(</span><span class="n">timeout</span><span class="p">)</span>
            <span class="n">events</span> <span class="o">=</span> <span class="nf">dict</span><span class="p">(</span><span class="n">events</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">client_data_socket</span> <span class="ow">in</span> <span class="n">events</span><span class="p">:</span>
                <span class="n">msg</span> <span class="o">=</span> <span class="k">await</span> <span class="n">client_data_socket</span><span class="p">.</span><span class="nf">recv_pyobj</span><span class="p">()</span>
                <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">Received client data: </span><span class="si">{</span><span class="n">msg</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
                <span class="k">if</span> <span class="nf">isinstance</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="n">ClientData</span><span class="p">):</span>
                    <span class="n">self</span><span class="p">.</span><span class="n">_data</span><span class="p">.</span><span class="nf">put_nowait</span><span class="p">(</span><span class="n">msg</span><span class="p">.</span><span class="n">data</span><span class="p">)</span>
                <span class="k">elif</span> <span class="nf">isinstance</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="n">Status</span><span class="p">):</span>
                    <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Wait for last messages</span><span class="sh">"</span><span class="p">)</span>
                    <span class="n">timeout</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">logger</span><span class="p">.</span><span class="nf">debug</span><span class="p">(</span><span class="sh">"</span><span class="s">Done receiving</span><span class="sh">"</span><span class="p">)</span>
                <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span> <span class="o">=</span> <span class="bp">False</span>
                <span class="n">self</span><span class="p">.</span><span class="n">_loop</span><span class="p">.</span><span class="nf">stop</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">__iter__</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="k">while</span> <span class="n">self</span><span class="p">.</span><span class="n">is_receiving</span><span class="p">:</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="k">yield</span> <span class="n">self</span><span class="p">.</span><span class="n">_data</span><span class="p">.</span><span class="nf">get_nowait</span><span class="p">()</span>
            <span class="k">except</span> <span class="n">Empty</span><span class="p">:</span>
                <span class="k">continue</span>


<span class="k">class</span> <span class="nc">Dataset</span><span class="p">(</span><span class="n">IterableDataset</span><span class="p">):</span>
    <span class="sh">"""</span><span class="s">Online dataset.</span><span class="sh">"""</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="n">self</span><span class="p">,</span> <span class="n">communicator</span><span class="p">:</span> <span class="n">ServerCommunicator</span><span class="p">):</span>
        <span class="nf">super</span><span class="p">().</span><span class="nf">__init__</span><span class="p">()</span>
        <span class="n">self</span><span class="p">.</span><span class="n">communicator</span> <span class="o">=</span> <span class="n">communicator</span>

    <span class="k">def</span> <span class="nf">__iter__</span><span class="p">(</span><span class="n">self</span><span class="p">):</span>
        <span class="k">yield</span> <span class="k">from</span> <span class="n">self</span><span class="p">.</span><span class="n">communicator</span>


<span class="k">def</span> <span class="nf">main</span><span class="p">(</span><span class="n">runner_hostname</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">nsim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="n">communicator</span> <span class="o">=</span> <span class="nc">ServerCommunicator</span><span class="p">(</span><span class="n">runner_hostname</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nf">range</span><span class="p">(</span><span class="n">nsim</span><span class="p">):</span>
        <span class="c1"># Format the command to send to the runner for execution
</span>        <span class="c1"># The command could include specific parameters
</span>        <span class="n">command</span> <span class="o">=</span> <span class="nf">format_command</span><span class="p">(</span><span class="sa">f</span><span class="sh">"</span><span class="s">python </span><span class="si">{</span><span class="n">os</span><span class="p">.</span><span class="n">path</span><span class="p">.</span><span class="nf">abspath</span><span class="p">(</span><span class="n">dummy_client</span><span class="p">.</span><span class="n">__file__</span><span class="p">)</span><span class="si">}</span><span class="sh">"</span><span class="p">)</span>
        <span class="n">communicator</span><span class="p">.</span><span class="nf">submit</span><span class="p">(</span><span class="n">command</span><span class="p">)</span>
    <span class="c1"># Signal no more simulation will be submitted
</span>    <span class="n">communicator</span><span class="p">.</span><span class="nf">terminate</span><span class="p">()</span>

    <span class="c1"># Use the data as in a regular PyTorch training loop
</span>    <span class="n">dataset</span> <span class="o">=</span> <span class="nc">Dataset</span><span class="p">(</span><span class="n">communicator</span><span class="p">)</span>
    <span class="n">dataloader</span> <span class="o">=</span> <span class="nc">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">batch</span><span class="p">,</span> <span class="n">data</span> <span class="ow">in</span> <span class="nf">enumerate</span><span class="p">(</span><span class="n">dataloader</span><span class="p">):</span>
        <span class="nf">print</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="n">data</span><span class="p">.</span><span class="n">shape</span><span class="p">)</span>
        <span class="c1"># Perform the training forward and backward passes
</span>        <span class="c1"># ...
</span>

<span class="k">if</span> <span class="n">__name__</span> <span class="o">==</span> <span class="sh">"</span><span class="s">__main__</span><span class="sh">"</span><span class="p">:</span>
    <span class="n">parser</span> <span class="o">=</span> <span class="n">argparse</span><span class="p">.</span><span class="nc">ArgumentParser</span><span class="p">(</span><span class="sh">"</span><span class="s">Multiprocessing example</span><span class="sh">"</span><span class="p">)</span>
    <span class="n">parser</span><span class="p">.</span><span class="nf">add_argument</span><span class="p">(</span><span class="sh">"</span><span class="s">hostname</span><span class="sh">"</span><span class="p">,</span> <span class="nb">type</span><span class="o">=</span><span class="nb">str</span><span class="p">)</span>
    <span class="n">parser</span><span class="p">.</span><span class="nf">add_argument</span><span class="p">(</span>
        <span class="sh">"</span><span class="s">--nsim</span><span class="sh">"</span><span class="p">,</span> <span class="nb">type</span><span class="o">=</span><span class="nb">int</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">help</span><span class="o">=</span><span class="sh">"</span><span class="s">Number of simulations to run.</span><span class="sh">"</span>
    <span class="p">)</span>
    <span class="n">parser</span><span class="p">.</span><span class="nf">add_argument</span><span class="p">(</span><span class="sh">"</span><span class="s">--batch_size</span><span class="sh">"</span><span class="p">,</span> <span class="nb">type</span><span class="o">=</span><span class="nb">int</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">args</span> <span class="o">=</span> <span class="n">parser</span><span class="p">.</span><span class="nf">parse_args</span><span class="p">()</span>
    <span class="n">runner_hostname</span> <span class="o">=</span> <span class="n">args</span><span class="p">.</span><span class="n">hostname</span>
    <span class="n">nsim</span> <span class="o">=</span> <span class="n">args</span><span class="p">.</span><span class="n">nsim</span>
    <span class="n">batch_size</span> <span class="o">=</span> <span class="n">args</span><span class="p">.</span><span class="n">batch_size</span>
    <span class="nf">main</span><span class="p">(</span><span class="n">runner_hostname</span><span class="p">,</span> <span class="n">nsim</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">)</span>
</pre></td></tr></tbody></table></code></pre></div> </div> </details> <h2 id="run-the-example">Run the Example</h2> <p>Copy the python files above in the same folder. In one terminal run the following command to start the runner. The output will display the runner’s hostname.</p> <div class="language-bash highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
</pre></td><td class="rouge-code"><pre>python runner.py
</pre></td></tr></tbody></table></code></pre></div></div> <p>In another terminal, starts the server by executing the following command. For synchronization the runner’s hostname must be specified (for instance <code class="language-plaintext highlighter-rouge">127.0.1.1</code>).</p> <div class="language-bash highlighter-rouge"><div class="highlight"><pre class="highlight"><code><table class="rouge-table"><tbody><tr><td class="rouge-gutter gl"><pre class="lineno">1
</pre></td><td class="rouge-code"><pre>python server.py &lt;runner<span class="s1">'s hostname&gt; --nsim 20 --batch_size 4
</span></pre></td></tr></tbody></table></code></pre></div></div> <p>The server will ask the runner to submit jobs to the clients. The clients will then produce data and stream them to the server, These data will be available at the server level for the training loop as a regular iterable dataset.</p> <h1 id="could-it-be-implemented-with-ray">Could it be Implemented with Ray?</h1> <p>Instead of reinventing the wheel and implement the framework from scratch, we may first wonder whether there is not already existing libraries that would do the job. Indeed, such online learning settings are not new and even common in domains like <em>reinforcement learning</em>. The submodule of the library <a href="https://docs.ray.io/en/latest/index.html">Ray</a> for reinforcement learning, <a href="https://docs.ray.io/en/latest/rllib/index.html">RLlib</a>, allows orchestrating <em>Actors</em> that will explore the environment given a policy and send <em>state</em> and <em>rewards</em> to <em>Learners</em> that will update the policy. Both <em>Actors</em> and <em>Learners</em> run on different resources of a cluster. We can draw a clear parallel between the use case of RLlib and our framework design by identifying <em>Actors</em> with <em>Clients</em>, <em>Learners</em> with <em>Servers</em>, and <em>rewards</em> with <em>simulated data</em> by the oracle. The comparison is nonetheless limited. With RLlib, <em>Actors</em> are generally Python executables that run on a single process, whereas in our case the oracle is most likely to be a non-Python program that runs on multiple processes using MPI. It is not clear yet how RLlib can be tweaked to support the proposed framework.</p> <h1 id="conclusion">Conclusion</h1> <ul> <li>We presented an online deep learning framework tailored for AI for Science applications;</li> <li>We provided a starting example of such framework;</li> <li>Wait a minute! You said better, cheaper, and faster? By alleviating IO, the framework accelerates data generation thus training. By reducing disk storage use it also makes it cheaper. Enabling greater data diversity improves training quality. Check the <a href="../../../assets/pdf/High_Throughput_Training_of_Deep_Surrogates_from_Large_Ensemble_Runs.pdf">paper</a> for more details.</li> </ul> <h1 id="correspondence">Correspondence</h1> <p>Any question or suggestion to improve the framework design is welcome. You can contact me through <a href="https://www.linkedin.com/in/lucas-meyer-a7983b103/">LinkedIn</a>.</p>]]></content><author><name></name></author><summary type="html"><![CDATA[An online training framework design for AI for Science applications.]]></summary></entry></feed>