<p dir="auto">Steps to reproduce the behavior:</p>
<p dir="auto">Call torch.ceil on a sparse tensor<br>
Below is a code sample:</p>
<p dir="auto">import torch<br>
i = torch.cuda.LongTensor([[0, 1, 1],<br>
[2, 0, 2],<br>
[2, 0, 2]])</p>
<p dir="auto">v  = torch.cuda.FloatTensor([3.2, -44.5, -5.1])</p>
<p dir="auto">s = torch.sparse_coo_tensor(i, v, torch.Size([2, 4, 4]), device=torch.device('cuda'))<br>
print(torch.ceil(s))</p>
<p dir="auto">Actual behavior<br>
Traceback (most recent call last):<br>
File "sparse_floor.py", line 9, in <br>
print(torch.ceil(s))<br>
RuntimeError: Could not run 'aten::ceil.out' with arguments from the 'SparseCUDA' backend. 'aten::ceil.out' is only available for these backends: [CPU, CUDA, Named, Autograd, Profiler, Tracer].</p>
<p dir="auto">Expected behavior<br>
ceil of the sparse tensor should get calculated and output as a sparse tensor</p>
<p dir="auto">Environment<br>
Versions of relevant libraries:<br>
[pip3] numpy==1.19.1<br>
[pip3] torch==1.8.0a0<br>
[pip3] torchtext==0.8.0a0+7e267d2<br>
[pip3] torchvision==0.8.0a0+be8192e<br>
[conda] blas 1.0 mkl<br>
[conda] cudatoolkit 10.2.89 hfd86e86_1<br>
[conda] magma-cuda102 2.5.2 1 pytorch<br>
[conda] mkl 2020.1 217<br>
[conda] mkl-include 2020.1 217<br>
[conda] mkl-service 2.3.0 py38he904b0f_0<br>
[conda] mkl_fft 1.1.0 py38h23d657b_0<br>
[conda] mkl_random 1.1.1 py38h0573a6f_0<br>
[conda] numpy 1.19.1 py38hbc911f0_0<br>
[conda] numpy-base 1.19.1 py38hfa32c7d_0<br>
[conda] torch 1.8.0a0 dev_0<br>
[conda] torchtext 0.8.0a0+7e267d2 pypi_0 pypi<br>
[conda] torchvision 0.8.0a0+be8192e pypi_0 pypi</p>
<p dir="auto">PyTorch Version: 1.8.0a0<br>
OS: Linux<br>
How you installed PyTorch (conda, pip, source): source<br>
Build command you used (if compiling from source): python setup.py develop<br>
Python version: 3.8<br>
CUDA/cuDNN version: 10.2</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/aocsa/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/aocsa">@aocsa</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/nikitaved/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/nikitaved">@nikitaved</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pearu/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pearu">@pearu</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mruberry/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mruberry">@mruberry</a></p>
<p dir="auto">-- with <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/Tierex/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/Tierex">@Tierex</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mrshenli/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mrshenli">@mrshenli</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/kiukchung/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/kiukchung">@kiukchung</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pritamdamania/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pritamdamania">@pritamdamania</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/lw/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/lw">@lw</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/agolynski/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/agolynski">@agolynski</a> , <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/osalpekar/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/osalpekar">@osalpekar</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zhaojuanmao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zhaojuanmao">@zhaojuanmao</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/rohan-varma/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rohan-varma">@rohan-varma</a></p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature</h2>
<p dir="auto">Concept and a new set of rpc APIs in <code class="notranslate">torch.distributed.rpc</code> aiming to help users construct distributed applications more easily, apis maily comprise of:</p>
<ol dir="auto">
<li>Higher level torch rpc APIs, including group creation, scoped service registering and easier value sharing among processes.</li>
<li>Succinct init methods to for apps using both rpc and process groups with a heterogeneous topology (roles),</li>
<li>Ensure all the above works well when launching these apps using TorchElastic, or with a simple launch script<br>
such as <code class="notranslate">torch.distributed.launch</code>.</li>
</ol>
<h2 dir="auto">Motivation</h2>
<p dir="auto">This RFC is a refined and extended version of <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="656856394" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/41425" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/41425/hovercard" href="https://github.com/pytorch/pytorch/issues/41425">#41425</a></p>
<h3 dir="auto">Homogeneous/Heterogeneous Torch RPC Apps</h3>
<p dir="auto">This section has been described in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="656856394" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/41425" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/41425/hovercard" href="https://github.com/pytorch/pytorch/issues/41425">#41425</a> with details, and therefore it will be omitted here.</p>
<h3 dir="auto">Service registering</h3>
<p dir="auto">Currently <code class="notranslate">torch.distributed.rpc</code> has only defined three basic point to point style APIs: <code class="notranslate">rpc_sync</code>, <code class="notranslate">rpc_async</code>, and <code class="notranslate">rpc_remote</code>. They require users to manually specify a specific action: <code class="notranslate">func</code> to perform, and a specific destination: <code class="notranslate">to</code> where this action will be performed on:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)"><pre class="notranslate"><code class="notranslate">torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)
</code></pre></div>
<p dir="auto">However, in most real application scenarios, a single process does not provide sufficient computing power to provide a complete service, therefore users are usually dealing with a complex of different processes serving different roles, such as a parameter server architecture diagram below:<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/17508457/87754941-96af0e00-c838-11ea-8798-02ff8dd2cb2f.png"><img src="https://user-images.githubusercontent.com/17508457/87754941-96af0e00-c838-11ea-8798-02ff8dd2cb2f.png" alt="PS" style="max-width: 100%;"></a></p>
<p dir="auto">If users want to invoke <code class="notranslate">push()</code> or the <code class="notranslate">pull()</code> method, it would be most sensible to provide a mechanism to register these two methods along with the method provider (i.e: job queue, model distributors) in the scope where they would like to use them, and let processes query for the registered method.<br>
Compared to the mechanism mentioned in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="656856394" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/41425" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/41425/hovercard" href="https://github.com/pytorch/pytorch/issues/41425">#41425</a></p>
<blockquote>
<p dir="auto">group-rpc calls (e.g. call the same function on all trainers)</p>
</blockquote>
<p dir="auto">which can also provide a similar function as "Service registering", "Service registering" enables users to call different processes with different roles, as well as hiding the actual service provider and make detail implementations transparent to users. This is especially important if users wants to maintain modularity or write complex applications.</p>
<h3 dir="auto">Resource sharing</h3>
<p dir="auto">Although RPC is an "action" based paradigm, compared to the "resource" based paradigm used by http &amp; restful apis, sometimes it could be useful to borrow some ideas from there. For example, to implement a simple synchronization mechanism using a shared value in a group of processes:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def process1():
    rpc.init_rpc(...)
    some_signal = Switch(False)           # a switch class
    rpc.pair(&quot;signal&quot;, some_signal)   # expose this variable in the implicit global rpc group
    sleep(10)                              # do some initialization
    some_signal.set(True)
    # synchronized
   
def process2():
    rpc.init_rpc(...)
    some_signal = Switch(False)           # a switch class
    rpc.pair(&quot;signal2&quot;, some_signal) # expose this variable in the implicit global rpc group
    sleep(5)                              # do some initialization
    some_signal.set(True)
    while not rpc.get_paired(&quot;signal&quot;):
        sleep(1)
    # synchronized

def process2():
    rpc.init_rpc(...)
    while not rpc.get_paired(&quot;signal&quot;) or not rpc.get_paired(&quot;signal2&quot;):
        sleep(1)
    # synchronized"><pre class="notranslate"><code class="notranslate">def process1():
    rpc.init_rpc(...)
    some_signal = Switch(False)           # a switch class
    rpc.pair("signal", some_signal)   # expose this variable in the implicit global rpc group
    sleep(10)                              # do some initialization
    some_signal.set(True)
    # synchronized
   
def process2():
    rpc.init_rpc(...)
    some_signal = Switch(False)           # a switch class
    rpc.pair("signal2", some_signal) # expose this variable in the implicit global rpc group
    sleep(5)                              # do some initialization
    some_signal.set(True)
    while not rpc.get_paired("signal"):
        sleep(1)
    # synchronized

def process2():
    rpc.init_rpc(...)
    while not rpc.get_paired("signal") or not rpc.get_paired("signal2"):
        sleep(1)
    # synchronized
</code></pre></div>
<p dir="auto">compared to</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="signal1 = False
signal2 = False

def get_signal1():
    return signal1

def get_signal2():
    return signal2

def process1():
    rpc.init_rpc(...)
    sleep(10)                              # do some initialization
    some_signal = True
    
    # synchronized
   
def process2():
    rpc.init_rpc(...)
    sleep(5)                              # do some initialization
    while not rpc.rpc_sync(&quot;proc1&quot;, get_signal1):
        sleep(1)
    # synchronized

def process2():
    rpc.init_rpc(...)
    while (not rpc.rpc_sync(&quot;proc1&quot;, get_signal1) or
           not rpc.rpc_sync(&quot;proc1&quot;, get_signal2)):
        sleep(1)
    # synchronized"><pre class="notranslate"><code class="notranslate">signal1 = False
signal2 = False

def get_signal1():
    return signal1

def get_signal2():
    return signal2

def process1():
    rpc.init_rpc(...)
    sleep(10)                              # do some initialization
    some_signal = True
    
    # synchronized
   
def process2():
    rpc.init_rpc(...)
    sleep(5)                              # do some initialization
    while not rpc.rpc_sync("proc1", get_signal1):
        sleep(1)
    # synchronized

def process2():
    rpc.init_rpc(...)
    while (not rpc.rpc_sync("proc1", get_signal1) or
           not rpc.rpc_sync("proc1", get_signal2)):
        sleep(1)
    # synchronized
</code></pre></div>
<p dir="auto">is simpler, because in the second example, users will have to roll out a new function for each signal they would like to look up, or<br>
they have to implement a custom lookup function taking a signal name and return a signal value, like below:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="signals = {}
def lookup(signal_name: str):
   return signals.get(signal_name, None)"><pre class="notranslate"><code class="notranslate">signals = {}
def lookup(signal_name: str):
   return signals.get(signal_name, None)
</code></pre></div>
<p dir="auto">Therefore, it is reasonable to provide a simple api like <code class="notranslate">rpc.pair</code> and <code class="notranslate">rpc.get_paired</code> to expose the value to the respective scope, so that processes can easily look up other values.</p>
<h3 dir="auto">Groups</h3>
<p dir="auto">If we want to either share a resource, or access a service, by RPC, it is crutial to limit their effect scope, so we can define services / pair values in a controllable, safe range, without overwriting other paired values and creating hidden conflicts. RPC has a implicit global group, as stated in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="656856394" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/41425" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/41425/hovercard" href="https://github.com/pytorch/pytorch/issues/41425">#41425</a> :</p>
<blockquote>
<p dir="auto">rpc not really needing a group (other than the implicit global one).</p>
</blockquote>
<p dir="auto">Again, we must be clear that <code class="notranslate">rpc_sync</code>, <code class="notranslate">rpc_async</code>, and <code class="notranslate">rpc_remote</code> are all operating on a global scope, and only having one scope is a severe disadvantage in OOP designs.</p>
<p dir="auto">There is one difference in the group conception between this RFC and <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="656856394" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/41425" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/41425/hovercard" href="https://github.com/pytorch/pytorch/issues/41425">#41425</a>, that is group may consists of different roles, for example, in the parameter server architecture above, users may implement each sub-module and group them by their kind, or consider all submodules as members of the parameter server RPC group:<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/17508457/87755322-5ef49600-c839-11ea-8397-ef1af1f96f90.png"><img src="https://user-images.githubusercontent.com/17508457/87755322-5ef49600-c839-11ea-8397-ef1af1f96f90.png" alt="PS2" style="max-width: 100%;"></a></p>
<h2 dir="auto">Pitch</h2>
<h3 dir="auto">Role based auto initialization</h3>
<p dir="auto">Ideally we'd like the application code to look like the following:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def master_main():
    # do something
    pass

def trainer_main():
    # do something
    pass

if __name__ == &quot;__main__&quot;:
    rpc.init_role(backend=backend, backend_options)
    if rpc.get_worker_info().role == &quot;master&quot;:
        master_main()
    elif rpc.get_worker_info().role == &quot;trainer&quot;:
        trainer_main()
    rpc.shutdown()"><pre class="notranslate"><code class="notranslate">def master_main():
    # do something
    pass

def trainer_main():
    # do something
    pass

if __name__ == "__main__":
    rpc.init_role(backend=backend, backend_options)
    if rpc.get_worker_info().role == "master":
        master_main()
    elif rpc.get_worker_info().role == "trainer":
        trainer_main()
    rpc.shutdown()
</code></pre></div>
<p dir="auto">So users do not specify world size and process rank by hand, instead they should pass how many roles to launch to torchelastic launcher or a simpler launcher scipt similar to <code class="notranslate">torch.distributed.launch</code> if they don't want to configure torchelastic with etcd:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; python -m torchelastic.distributed.launch
    --nnodes=$NUM_NODES
    --nproc_per_node=$NUM_TRAINERS
    --role=trainer:4         # launch 4 trainer processes
    --role=master:1        # launch 1 master process
    --rdzv_id=$JOB_ID
    --rdzv_backend=etcd
    --rdzv_endpoint=$ETCD_HOST:$ETCD_PORT
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)"><pre class="notranslate"><code class="notranslate">&gt;&gt;&gt; python -m torchelastic.distributed.launch
    --nnodes=$NUM_NODES
    --nproc_per_node=$NUM_TRAINERS
    --role=trainer:4         # launch 4 trainer processes
    --role=master:1        # launch 1 master process
    --rdzv_id=$JOB_ID
    --rdzv_backend=etcd
    --rdzv_endpoint=$ETCD_HOST:$ETCD_PORT
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
</code></pre></div>
<p dir="auto">And torchelastic or the launching script will set <code class="notranslate">RANK</code>, <code class="notranslate">WORLD_SIZE</code>, <code class="notranslate">NODE_RANK</code>, <code class="notranslate">NODE_DIMS</code>, <code class="notranslate">LOCAL_RANK</code>, <code class="notranslate">ROLE</code>, <code class="notranslate">ROLE_NUM</code>, <code class="notranslate">ROLE_NUMS</code>, <code class="notranslate">ROLE_RANK</code>, <code class="notranslate">NAME</code>, <code class="notranslate">MASTER_ADDR</code>, <code class="notranslate">MASTER_PORT</code> in the env vars, and <code class="notranslate">rpc.init_role</code> will pick these env vars up.</p>
<p dir="auto">After the application is initialized the following hold true:</p>
<ol dir="auto">
<li>
<p dir="auto">Every process (regardless of role) can invoke rpc APIs on one another.</p>
</li>
<li>
<p dir="auto">Every process has a global rank (e.g. one that was passed as RANK env var) that can be obtained via dist.get_rank()</p>
</li>
<li>
<p dir="auto">Each node is assigned a rank. (this is equivalent to the elastic agent’s rank - a.k.a GROUP_RANK in torchelastic)</p>
</li>
<li>
<p dir="auto">The global rank can be derived from the following formula:<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/43595115/87468522-a75c4b80-c5ce-11ea-9994-72b39363fd97.png"><img src="https://user-images.githubusercontent.com/43595115/87468522-a75c4b80-c5ce-11ea-9994-72b39363fd97.png" alt="eq" style="max-width: 100%;"></a><br>
Where <code class="notranslate">node_dims[i]</code> returns the number of local processes on node rank i.<br>
But cannot be derived from <code class="notranslate">ROLE_NUMS</code> since there is no explicit garantee on the order of roles.</p>
</li>
<li>
<p dir="auto">the launching utility (torchelastic or a simple script launcher) will set <code class="notranslate">RANK</code>, <code class="notranslate">WORLD_SIZE</code>, <code class="notranslate">NODE_RANK</code>(<code class="notranslate">GROUP_RANK</code>), <code class="notranslate">NODE_DIMS</code>, <code class="notranslate">LOCAL_RANK</code>, <code class="notranslate">ROLE</code>, <code class="notranslate">ROLE_NUM</code>, <code class="notranslate">ROLE_NUMS</code>,  <code class="notranslate">ROLE_RANK</code>, <code class="notranslate">NAME</code>, <code class="notranslate">MASTER_ADDR</code>, <code class="notranslate">MASTER_PORT</code> in the env vars.<br>
<code class="notranslate">ROLE_NUM</code> is the number of processes in current role, and <code class="notranslate">ROLE_NUMS</code> is the number of processes of each role.<br>
<code class="notranslate">NODE_DIMS</code> is a comma seperated string of &gt;0 integers: 2,3,4,5<br>
<code class="notranslate">ROLE_NUMS</code> is a comma sepeated string of role names and &gt;0 integers: trainer,4,master,1<br>
There is no trailling comma in them.<br>
Role names must not contain commas.</p>
</li>
<li>
<p dir="auto">Every process has a unique name in form <code class="notranslate">name = {role}:{role_id}</code>, <code class="notranslate">{role_id}</code> is a number in [0, num_workers_in_role),</p>
</li>
<li>
<p dir="auto">The <code class="notranslate">WorkerInfo</code> object returned by <code class="notranslate">torch.distributed.rpc.get_worker_info()</code> should contain 4 properties: <code class="notranslate">name</code>, <code class="notranslate">id</code>, <code class="notranslate">role</code>, <code class="notranslate">role_id</code>. <code class="notranslate">id</code> is equal to the global rank, users can look up <code class="notranslate">role</code> and <code class="notranslate">role_id</code> of current process by passing <code class="notranslate">None</code> to <code class="notranslate">get_worker_info()</code>. If users are calling <code class="notranslate">rpc.init_rpc</code> instead of <code class="notranslate">rpc.init_role</code> to initialize, <code class="notranslate">role</code> and <code class="notranslate">role_id</code> should be <code class="notranslate">None</code>.</p>
</li>
</ol>
<p dir="auto">Clarification or different ranks:<br>
! <code class="notranslate">RANK</code> depicts the global rank of the current process among all launched processes, it is an int in <code class="notranslate">[0, num_total_processes)</code><br>
! <code class="notranslate">NODE_RANK</code> depicts the rank of container / physical server, it is an int in <code class="notranslate">[0, num_container_or_server)</code><br>
! <code class="notranslate">LOCAL_RANK</code> depicts the rank of the current process among processes on this container/server, it is an int in <code class="notranslate">[0, NODE_DIMS[NODE_RANK])</code><br>
! <code class="notranslate">ROLE_RANK</code> depicts the rank of current process among processes in this role,  it is an int in <code class="notranslate">[0, ROLE_NUMS[ROLE])</code></p>
<h3 dir="auto">Role-based rpc APIs</h3>
<p dir="auto">Since we are working with the concept of roles this is really a syntactic sugar that allows users to perform group rpc calls on all workers belonging to a role rather than having to for-loop around them. It can trivially be implemented by first getting all the worker names for the role (or derive it by using role_world_size = ROLE_NUMS[ROLE] = ROLE_NUM)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="ret_futures = {}
role = os.environ['ROLE']
role_num = int(os.environ['ROLE_NUM'])
for idx in range(0, role_num):
    name=f&quot;trainer:{idx}&quot;
    ret_futures[name] = rpc.rpc_async(to=name, func=foobar, args=(...)))

# versus
ret_futures = rpc.rpc_async_role(role=&quot;trainer&quot;, selector=lambda _: True, func=foobar, args=(...))

# selector is a function which takes in a role rank and returns True for selecting it, False for deselecting it.
# could be None, and selecting all processes with this role"><pre class="notranslate"><code class="notranslate">ret_futures = {}
role = os.environ['ROLE']
role_num = int(os.environ['ROLE_NUM'])
for idx in range(0, role_num):
    name=f"trainer:{idx}"
    ret_futures[name] = rpc.rpc_async(to=name, func=foobar, args=(...)))

# versus
ret_futures = rpc.rpc_async_role(role="trainer", selector=lambda _: True, func=foobar, args=(...))

# selector is a function which takes in a role rank and returns True for selecting it, False for deselecting it.
# could be None, and selecting all processes with this role
</code></pre></div>
<p dir="auto">Note The original design also included an rpc.wait_all() API for completeness. This has already been implemented in the form of torch.futures.collect_all (<a href="https://github.com/pytorch/pytorch/blob/master/torch/futures/__init__.py#L88">https://github.com/pytorch/pytorch/blob/master/torch/futures/__init__.py#L88</a>)</p>
<h3 dir="auto">Create groups</h3>
<p dir="auto">In this section:</p>
<ol dir="auto">
<li>MPI style APIs will be refered to as "collective communication primitives" or "coll comm" for short.</li>
<li>"sub" means a group whose members is a subset of a larger group of same kind.</li>
</ol>
<h4 dir="auto">Collective group</h4>
<p dir="auto">In current distributed apis, we can create a collective group for coll comm using:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800), backend=None)"><pre class="notranslate"><code class="notranslate">torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800), backend=None)
</code></pre></div>
<p dir="auto">this function takes in a list of raw ranks, is blocking, and requires all mentioned processes in <code class="notranslate">ranks</code> to enter. However, many users find coll comm extremely useful in many scenarios, therefore, we decide to import all of its functions into the scope of rpc,<br>
so that instead of using a list of raw ranks, users may now create a collective group with:</p>
<ol dir="auto">
<li>a list of names, such as ["worker:0", "worker:1", "manager:0"]</li>
<li>a role name, with a index selector: role="worker", selector=lambda id: id in (0, 1)</li>
</ol>
<p dir="auto">api has the following signature:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="rpc.new_collective_group(names, role: str = None, selector: Callable = None, timeout=datetime.timedelta(0, 1800), backend=None) -&gt; CollectiveGroup"><pre class="notranslate"><code class="notranslate">rpc.new_collective_group(names, role: str = None, selector: Callable = None, timeout=datetime.timedelta(0, 1800), backend=None) -&gt; CollectiveGroup
</code></pre></div>
<p dir="auto"><code class="notranslate">CollectiveGroup</code> has the same coll comm apis as <code class="notranslate">torch.distributed</code>, except the <code class="notranslate">group</code> keyword argument is the created group itself, and <code class="notranslate">dst</code> and <code class="notranslate">src</code> are process names instead of global ranks:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def send(self, tensor, dst, tag=0)
def recv(self, tensor, src=None, tag=0)
def isend(self, tensor, dst, tag=0)
def irecv(self, tensor, src, tag=0)
..."><pre class="notranslate"><code class="notranslate">def send(self, tensor, dst, tag=0)
def recv(self, tensor, src=None, tag=0)
def isend(self, tensor, dst, tag=0)
def irecv(self, tensor, src, tag=0)
...
</code></pre></div>
<h4 dir="auto">RPC group</h4>
<p dir="auto">RPC group is a name scope, each process may join multiple RPC groups, and each RPC group may contain multiple processes, RPC group is a place where users may:</p>
<ol dir="auto">
<li>share variables only among members by pairing their value with a key to the group</li>
<li>expose services so that only its member processes may access.</li>
</ol>
<p dir="auto">The default RPC group is the rpc world, that is the <code class="notranslate">rpc</code> module itself:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)"><pre class="notranslate"><code class="notranslate">torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)
</code></pre></div>
<p dir="auto">are all operating on this implicit global group.</p>
<p dir="auto">In order to create a sub rpc group:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="rpc.new_rpc_group(self, names, role: str = None, selector: Callable = None, group_name: str) -&gt; RpcGroup"><pre class="notranslate"><code class="notranslate">rpc.new_rpc_group(self, names, role: str = None, selector: Callable = None, group_name: str) -&gt; RpcGroup
</code></pre></div>
<p dir="auto">compared to the collective group, <code class="notranslate">group_name</code> is necessary here because it will identify the scope of shared variables and exposed services.</p>
<p dir="auto"><code class="notranslate">RpcGroup</code> class, as well as the implicit global group <code class="notranslate">rpc</code>, provides the following APIs:<br>
<strong>This part is open to discussion</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="# pair a value to this group, the same key cannot be used twice, otherwise an exception must be thrown
pair(key: str, Value: Any) -&gt; None

# unpair a value from this group, the key must exist, otherwise an exception must be thrown
unpair(key: str) -&gt; None

# query whether a key has been used to pair
is_paired(key: str) -&gt; bool

# get a paired value, if it exist, otherwise an exception must be thrown
get_paired(key: str) -&gt; Any

# expose a service, the service function will be executed when called
# could be a global scope function, a class method, a local function returned by a decorator, or lambda, 
# or torch script, etc.
register(key: str, func: Callable) -&gt; None

# deregister an exposed service, the key must exist, otherwise an exception must be thrown
deregister(key: str) -&gt; None

# query whether a service has been registered:
is_registered(key: str) -&gt; bool

# call a service
registered_sync(key: str, args=None, kwargs=None)
registered_async(key: str, args=None, kwargs=None)
registered_remote(key: str, args=None, kwargs=None)

# destroy the group, global rpc does not have this function, since it has &quot;shutdown&quot;
# also deregister the group along with values, services from the global discovery manager
destroy()

# info accessors
# returns a list of member names
get_members() -&gt; List[str]

# judge whether the target process is a member of this group, by default it is the current process
is_member(name: str=None) -&gt; bool

# group name
get_group_name() -&gt; str

# number of group members
size() -&gt; int"><pre class="notranslate"><code class="notranslate"># pair a value to this group, the same key cannot be used twice, otherwise an exception must be thrown
pair(key: str, Value: Any) -&gt; None

# unpair a value from this group, the key must exist, otherwise an exception must be thrown
unpair(key: str) -&gt; None

# query whether a key has been used to pair
is_paired(key: str) -&gt; bool

# get a paired value, if it exist, otherwise an exception must be thrown
get_paired(key: str) -&gt; Any

# expose a service, the service function will be executed when called
# could be a global scope function, a class method, a local function returned by a decorator, or lambda, 
# or torch script, etc.
register(key: str, func: Callable) -&gt; None

# deregister an exposed service, the key must exist, otherwise an exception must be thrown
deregister(key: str) -&gt; None

# query whether a service has been registered:
is_registered(key: str) -&gt; bool

# call a service
registered_sync(key: str, args=None, kwargs=None)
registered_async(key: str, args=None, kwargs=None)
registered_remote(key: str, args=None, kwargs=None)

# destroy the group, global rpc does not have this function, since it has "shutdown"
# also deregister the group along with values, services from the global discovery manager
destroy()

# info accessors
# returns a list of member names
get_members() -&gt; List[str]

# judge whether the target process is a member of this group, by default it is the current process
is_member(name: str=None) -&gt; bool

# group name
get_group_name() -&gt; str

# number of group members
size() -&gt; int
</code></pre></div>
<p dir="auto">the pair/register process happens with two stages:</p>
<ol dir="auto">
<li>register to local map</li>
<li>push the result to a global discovery manager</li>
</ol>
<p dir="auto">the unpair/deregister process happens with two stages:</p>
<ol dir="auto">
<li>degister from local map</li>
<li>push the result to a global discovery manager</li>
</ol>
<p dir="auto">the get_paired/registered call process happens with two stages:</p>
<ol dir="auto">
<li>look up the local discovery cache (every process has it)</li>
<li>if cache does not contain the address: (group, process name, key, query_type), query the global discovery manager</li>
<li>otherwise go to the target address</li>
<li>if the target group has been destroyed / target value deregistered, return look up failed exception</li>
<li>otherwise go on as normal, return the result or execute the service</li>
</ol>
<p dir="auto">In order to discover a shared value / an exposed service, a register thread / process must be started once <code class="notranslate">rpc.init_rpc</code> or <code class="notranslate">rpc.init_auto</code> has been called, since torchelastic tears down and restart all processes when at least one processes becomes unhealthy, it is reasonable to regard the period in between as stable.<br>
Therefore, we can create a register service on a fixed process such as the process with rank=0, or use a more sophisticated service management infrastructure such as zookeeper, we should allow users to choose the discovery backend by specifying it in the <code class="notranslate">ProcessGroupRpcBackendOptions</code> passed to <code class="notranslate">rpc.init_rpc</code> and <code class="notranslate">rpc.init_auto</code>.<br>
To prevent the explosion of discovery queries, a local query cache should be placed on each process and store looked up values, outdated address values (caused by <code class="notranslate">unpair</code> or <code class="notranslate">deregister</code>) will be cleared when <code class="notranslate">get_paired</code> or <code class="notranslate">registered_sync</code> etc. returns a lookup-failed signal.</p>
<h3 dir="auto">RPC Info Accessors</h3>
<p dir="auto">Some times users may want to dive deeper and see the actual mapping relation. <code class="notranslate">rpc.get_role_map()</code> will return a dictionary mapping role name to global process rank.<br>
<strong>This part is open to discussion</strong></p>
<h3 dir="auto">Launching utilities</h3>
<p dir="auto">Since torchelastic currently is not included in the default pytorch installation, a new launching utility similar to <code class="notranslate">torch.distributed.launch</code> should be included, its name could be <code class="notranslate">torch.distributed.launch_roles</code>, or it could be merged into <code class="notranslate">torch.distributed.launch</code>. this utility should be able to takein a list of roles and spawn subprocesses, allocating roles to global ranks like:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="worker:0 -&gt; Process 0, Node 0
worker:1 -&gt; Process 1, Node 0
worker:2 -&gt; Process 2, Node 0
worker:3 -&gt; Process 0, Node 1
manager:0 -&gt; Process 1, Node 1
manager:1 -&gt; Process 2, Node 1"><pre class="notranslate"><code class="notranslate">worker:0 -&gt; Process 0, Node 0
worker:1 -&gt; Process 1, Node 0
worker:2 -&gt; Process 2, Node 0
worker:3 -&gt; Process 0, Node 1
manager:0 -&gt; Process 1, Node 1
manager:1 -&gt; Process 2, Node 1
</code></pre></div>
<p dir="auto">when invoked with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Node 0:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4  --role=manager:2
               --nnodes=2 --node_rank=0 --master_addr=&quot;192.168.1.1&quot;
               --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
               and all other arguments of your training script)

Node 1:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4  --role=manager:2
               --nnodes=2 --node_rank=1 --master_addr=&quot;192.168.1.1&quot;
               --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
               and all other arguments of your training script)
"><pre class="notranslate"><code class="notranslate">Node 0:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4  --role=manager:2
               --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
               --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
               and all other arguments of your training script)

Node 1:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4  --role=manager:2
               --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
               --master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
               and all other arguments of your training script)

</code></pre></div>
<p dir="auto">and <code class="notranslate">nproc_per_node * nnodes &gt;= num of all roles</code>, otherwise an exception must be thrown.</p>
<h3 dir="auto">Failure Handling Behaviors (naive launching tool)</h3>
<p dir="auto">No failure could be dealt with, one process fails, all processes fail.</p>
<h3 dir="auto">Failure Handling Behaviors (torchelastic)</h3>
<p dir="auto">This section describes the types of failures and who/how those failures are handled, when users are launching the application with torchelastic.</p>
<p dir="auto">We focus on worker process(es) failure since a node failure can be viewed as multiple worker failures. In general TorchElastic views failures as two scaling events: scale-down + scale-up. In TorchElastic the world size is fixed between "rounds" of quorum where "rounds" is defined as the state of the world between two rendezvous versions. Between rounds ALL existing workers are shut down and restarted, hence TorchElastic follows a "all-or-nothing" model when it comes to dealing with faults.</p>
<p dir="auto">There are three types of torch applications that we care about:</p>
<p dir="auto">RPC only (no process groups)<br>
RPC + DDP<br>
RPC + DDP + Pipelining<br>
In all three cases the following behavior on failures is guaranteed by TorchElastic:</p>
<ol dir="auto">
<li>Worker process(es) fail</li>
<li>elastic agent detects the failure</li>
<li>elastic agent tears down all other workers on the same host</li>
<li>elastic agent enters a re-rendezvous</li>
<li>other elastic agents are notified of the re-rendezvous event</li>
<li>other elastic agents tear down all their respective local workers</li>
<li>all elastic agents re-rendezvous (e.g. next round) and get assigned a rank</li>
<li>each elastic agent computes the worker ranks based on the agent ranks</li>
<li>all elastic agents restart the worker processes<br>
Note: This implies that work between checkpoints are lost.</li>
</ol>
<p dir="auto">Note on node failure: TorchElastic is NOT a scheduler hence it cannot deal with replacement of NODES. We rely on the scheduler to replace failed nodes (or containers). Assuming the scheduler replaces the failed node the following behavior is observed on node failures:</p>
<ol dir="auto">
<li>Node(s) fail and are replaced by the scheduler</li>
<li>Elastic agents are started on the replaced nodes and attempt to join the rendezvous</li>
<li>The surviving agents are notified of such event and they all tear down their local workers and re-rendezvous</li>
<li>A new version of rendezvous (e.g. round) is created and the worker processes are started again.</li>
</ol>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pietern/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pietern">@pietern</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mrshenli/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mrshenli">@mrshenli</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/kiukchung/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/kiukchung">@kiukchung</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pritamdamania87/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pritamdamania87">@pritamdamania87</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zhaojuanmao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zhaojuanmao">@zhaojuanmao</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/satgera/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/satgera">@satgera</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gqchen/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gqchen">@gqchen</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/aazzolini/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/aazzolini">@aazzolini</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/rohan-varma/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rohan-varma">@rohan-varma</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/xush6528/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/xush6528">@xush6528</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jjlilley/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jjlilley">@jjlilley</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/osalpekar/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/osalpekar">@osalpekar</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jiayisuse/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jiayisuse">@jiayisuse</a></p>