Add image

Author: rohan-varma

_static/img/rpc_trace_img.png

@@ -44,7 +44,7 @@ two RPC workers on the same host, named "worker0" and "worker1" respectively. Th
 be spawned as subprocesses, and we set some environment variables required for proper
 initialization (see torch.distributed documentation for more details).
 
-.. code:: python
+.. code:: python3
   import torch
   import torch.distributed.rpc as rpc
   import torch.autograd.profiler as profiler
@@ -95,7 +95,7 @@ Now that we have a skeleton setup of our RPC framework, we can move on to
 sending RPCs back and forth and using the profiler to obtain a view of what's
 happening under the hood. Let's add to the above "worker" function:
 
-..code:: python
+..code:: python3
     def worker(rank, world_size):
         # Above code omitted...
         if rank == 0:
@@ -152,7 +152,7 @@ call are prefixed with ::rpc_async#aten::mul(worker0 -> worker1).
 We can also use the profiler gain insight into user-defined functions that are executed over RPC. 
 For example, let's add the following to the above "worker" function:
 
-..code:: python
+..code:: python3
   # Define somewhere outside of worker() func.
   def udf_with_ops():
       import time
@@ -161,7 +161,7 @@ For example, let's add the following to the above "worker" function:
       torch.add(t1, t2)
       torch.mul(t1, t2)
 
-..code::python
+..code::python3
   def worker(rank, world_size):
       # Above code omitted
       with profiler.profile() as p:
@@ -197,7 +197,7 @@ remote operators that have been executed on worker 1 as part of executing this R
 Lastly, we can visualize remote execution using the tracing functionality provided by the profiler.
 Let's add the following code to the above "worker" function:
 
-..code:: python
+..code:: python3
     def worker(rank, world_size):
         # Above code omitted
         # Will generated trace for above profiling output
@@ -215,7 +215,7 @@ in this case, given in the trace column for "node_id: 1".
 
 Putting it all together, we have the following code for this recipe:
 
-..code:: python
+..code:: python3
     import torch
     import torch.distributed.rpc as rpc
     import torch.autograd.profiler as profiler

recipes_source/distributed_rpc_profiling.rst

@@ -0,0 +1,49 @@
+import torch
+import torch.distributed.rpc as rpc
+import torch.autograd.profiler as profiler
+import torch.multiprocessing as mp
+import os
+import logging
+import sys
+
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+logger = logging.getLogger()
+
+def random_tensor():
+    return torch.rand((3, 3), requires_grad=True)
+
+
+def worker(rank, world_size):
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "29500"
+    worker_name = f"worker{rank}"
+
+    # Initialize RPC framework.
+    rpc.init_rpc(
+        name=worker_name,
+        rank=rank,
+        world_size=world_size
+    )
+    logger.debug(f"{worker_name} successfully initialized RPC.")
+
+    pass # to be continued below
+    if rank == 0:  
+      dst_worker_rank = (rank + 1) % world_size
+      dst_worker_name = f"worker{dst_worker_rank}"
+      t1, t2 = random_tensor(), random_tensor() 
+      # Send and wait RPC completion under profiling scope.
+      with profiler.profile() as p:
+        fut1 = rpc.rpc_async(dst_worker_name, torch.add, args=(t1, t2))
+        fut2 = rpc.rpc_async(dst_worker_name, torch.mul, args=(t1, t2))
+        # RPCs must be awaited within profiling scope.
+        fut1.wait()
+        fut2.wait()
+
+      print(p.key_averages().table())
+
+
+
+if __name__ == '__main__':
+    # Run 2 RPC workers.
+    world_size = 2
+    mp.spawn(worker, args=(world_size,), nprocs=world_size)