Merge pull request #904 from SciPioneer/patch-1

Clang-format

Author: brianjo

distributed/rpc/ddp_rpc/main.py

@@ -1,28 +1,28 @@
-from functools import wraps
 import os
 import random
+from functools import wraps
 
 import torch
 import torch.distributed as dist
 import torch.distributed.autograd as dist_autograd
-from torch.distributed.optim import DistributedOptimizer
 import torch.distributed.rpc as rpc
+import torch.multiprocessing as mp
+import torch.optim as optim
+from torch.distributed.optim import DistributedOptimizer
 from torch.distributed.rpc import RRef
 from torch.distributed.rpc import TensorPipeRpcBackendOptions
-import torch.multiprocessing as mp
 from torch.nn.parallel import DistributedDataParallel as DDP
-import torch.optim as optim
 
 NUM_EMBEDDINGS = 100
 EMBEDDING_DIM = 16
 
+
 class HybridModel(torch.nn.Module):
     r"""
     The model consists of a sparse part and a dense part. The dense part is an
     nn.Linear module that is replicated across all trainers using
     DistributedDataParallel. The sparse part is an nn.EmbeddingBag that is
     stored on the parameter server.
-
     The model holds a Remote Reference to the embedding table on the parameter
     server.
     """
@@ -37,6 +37,7 @@ def forward(self, indices, offsets):
         emb_lookup = self.emb_rref.rpc_sync().forward(indices, offsets)
         return self.fc(emb_lookup.cuda(self.device))
 
+
 def _retrieve_embedding_parameters(emb_rref):
     return [RRef(p) for p in emb_rref.local_value().parameters()]
 
@@ -57,7 +58,8 @@ def _run_trainer(emb_rref, rank):
 
     # Retrieve parameters for embedding table.
     model_parameter_rrefs = rpc.rpc_sync(
-            "ps", _retrieve_embedding_parameters, args=(emb_rref,))
+        "ps", _retrieve_embedding_parameters, args=(emb_rref,)
+    )
 
     # model.parameters() only includes local parameters.
     for param in model.parameters():
@@ -118,29 +120,30 @@ def run_worker(rank, world_size):
     # We need to use different port numbers in TCP init_method for init_rpc and
     # init_process_group to avoid port conflicts.
     rpc_backend_options = TensorPipeRpcBackendOptions()
-    rpc_backend_options.init_method='tcp://localhost:29501'
+    rpc_backend_options.init_method = "tcp://localhost:29501"
 
     # Rank 2 is master, 3 is ps and 0 and 1 are trainers.
     if rank == 2:
         rpc.init_rpc(
-                "master",
-                rank=rank,
-                world_size=world_size,
-                rpc_backend_options=rpc_backend_options)
+            "master",
+            rank=rank,
+            world_size=world_size,
+            rpc_backend_options=rpc_backend_options,
+        )
 
         # Build the embedding table on the ps.
         emb_rref = rpc.remote(
-                "ps",
-                torch.nn.EmbeddingBag,
-                args=(NUM_EMBEDDINGS, EMBEDDING_DIM),
-                kwargs={"mode": "sum"})
+            "ps",
+            torch.nn.EmbeddingBag,
+            args=(NUM_EMBEDDINGS, EMBEDDING_DIM),
+            kwargs={"mode": "sum"},
+        )
 
         # Run the training loop on trainers.
         futs = []
         for trainer_rank in [0, 1]:
             trainer_name = "trainer{}".format(trainer_rank)
-            fut = rpc.rpc_async(
-                    trainer_name, _run_trainer, args=(emb_rref, rank))
+            fut = rpc.rpc_async(trainer_name, _run_trainer, args=(emb_rref, rank))
             futs.append(fut)
 
         # Wait for all training to finish.
@@ -149,32 +152,34 @@ def run_worker(rank, world_size):
     elif rank <= 1:
         # Initialize process group for Distributed DataParallel on trainers.
         dist.init_process_group(
-                backend="gloo", rank=rank, world_size=2,
-                init_method='tcp://localhost:29500')
+            backend="gloo", rank=rank, world_size=2, init_method="tcp://localhost:29500"
+        )
 
         # Initialize RPC.
         trainer_name = "trainer{}".format(rank)
         rpc.init_rpc(
-                trainer_name,
-                rank=rank,
-                world_size=world_size,
-                rpc_backend_options=rpc_backend_options)
+            trainer_name,
+            rank=rank,
+            world_size=world_size,
+            rpc_backend_options=rpc_backend_options,
+        )
 
         # Trainer just waits for RPCs from master.
     else:
         rpc.init_rpc(
-                "ps",
-                rank=rank,
-                world_size=world_size,
-                rpc_backend_options=rpc_backend_options)
+            "ps",
+            rank=rank,
+            world_size=world_size,
+            rpc_backend_options=rpc_backend_options,
+        )
         # parameter server do nothing
         pass
 
     # block until all rpcs finish
     rpc.shutdown()
 
 
-if __name__=="__main__":
+if __name__ == "__main__":
     # 2 trainers, 1 parameter server, 1 master.
     world_size = 4
-    mp.spawn(run_worker, args=(world_size, ), nprocs=world_size, join=True)
+    mp.spawn(run_worker, args=(world_size,), nprocs=world_size, join=True)