Make rpc and tensorlist part of PBG

Reviewed By: adamlerer

Differential Revision: D17742505

fbshipit-source-id: 6791a0a57127989f5af3ef0b63cee7b05f1b111f

test/test_entitylist.py

@@ -10,9 +10,9 @@
 from unittest import TestCase, main
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.entitylist import EntityList
+from torchbiggraph.tensorlist import TensorList
 
 
 def tensor_list_from_lists(lists: Sequence[Sequence[int]]) -> TensorList:

test/test_graph_storages.py

@@ -12,9 +12,9 @@
 import h5py
 import numpy as np
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.graph_storages import FileEdgeAppender
+from torchbiggraph.tensorlist import TensorList
 
 
 class TestFileEdgeAppender(TestCase):

test/test_model.py

@@ -14,7 +14,6 @@
 from unittest import TestCase, main
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.entitylist import EntityList
 from torchbiggraph.model import (
@@ -39,6 +38,7 @@
     TranslationOperator,
     match_shape,
 )
+from torchbiggraph.tensorlist import TensorList
 
 
 class TensorTestCase(TestCase):

torchbiggraph/bucket_scheduling.py

@@ -11,10 +11,9 @@
 from abc import ABC, abstractmethod
 from typing import Dict, List, NamedTuple, Optional, Set, Tuple
 
-from torch_extensions.rpc.rpc import Client, Server
-
 from torchbiggraph.config import BucketOrder
 from torchbiggraph.distributed import Startable
+from torchbiggraph.rpc import Client, Server
 from torchbiggraph.stats import Stats
 from torchbiggraph.types import Bucket, EntityName, Partition, Rank, Side
 

torchbiggraph/entitylist.py

@@ -9,8 +9,8 @@
 from typing import Any, Sequence, Union
 
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
+from torchbiggraph.tensorlist import TensorList
 from torchbiggraph.types import LongTensorType
 
 

torchbiggraph/graph_storages.py

@@ -18,11 +18,11 @@
 import h5py
 import numpy as np
 import torch
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.edgelist import EdgeList
 from torchbiggraph.entitylist import EntityList
 from torchbiggraph.plugin import URLPluginRegistry
+from torchbiggraph.tensorlist import TensorList
 from torchbiggraph.util import CouldNotLoadData
 
 

torchbiggraph/model.py

@@ -23,7 +23,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_extensions.tensorlist.tensorlist import TensorList
 
 from torchbiggraph.config import (
     ConfigSchema,
@@ -34,6 +33,7 @@
 from torchbiggraph.entitylist import EntityList
 from torchbiggraph.graph_storages import RELATION_TYPE_STORAGES
 from torchbiggraph.plugin import PluginRegistry
+from torchbiggraph.tensorlist import TensorList
 from torchbiggraph.types import FloatTensorType, LongTensorType, Side
 from torchbiggraph.util import CouldNotLoadData
 

torchbiggraph/parameter_sharing.py

@@ -63,7 +63,7 @@ class ParameterServer(Startable):
     get tensors by string key. Operations on the parameter server are globally
     synchronous.
 
-    FIXME: torch_extensions.rpc should be fixed to not require torch.serialization,
+    FIXME: torchbiggraph.rpc should be fixed to not require torch.serialization,
     then most of this code can be removed.
     FIXME: torch.distributed.recv should not require you to provide the
     tensor to write to; the type and size should be sent in the header.