Move internal functions to torch.distributed.rpc

Summary: Pull Request resolved: pytorch#27289 Test Plan: Imported from OSS Differential Revision: D17808214 Pulled By: pietern fbshipit-source-id: 4c453028e431c3e951d439784017ef07037ba1a9
kashif · Oct 8, 2019 · a6d26ce · a6d26ce
1 parent 14f1629
commit a6d26ce
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 7 deletions.
diff --git a/test/rpc_test.py b/test/rpc_test.py
@@ -13,7 +13,7 @@
 from common_utils import load_tests
 from dist_utils import INIT_METHOD_TEMPLATE, TEST_CONFIG, dist_init
 from torch.distributed import ProcessGroupAgent
-from torch.distributed.internal_rpc_utils import PythonUDF, _internal_rpc_pickler
+from torch.distributed.rpc.internal import PythonUDF, _internal_rpc_pickler
 from torch.distributed.rpc_api import RpcBackend
 
 

diff --git a/torch/csrc/distributed/rpc/python_rpc_handler.cpp b/torch/csrc/distributed/rpc/python_rpc_handler.cpp
@@ -20,8 +20,7 @@ py::object getFunction(const py::object& module, const char* name) {
 
 PythonRpcHandler::PythonRpcHandler() {
   AutoGIL ag;
-  py::object module =
-      py::module::import("torch.distributed.internal_rpc_utils");
+  py::object module = py::module::import("torch.distributed.rpc.internal");
   pyRunFunction_ = getFunction(module, "_run_function");
   pyLoadReturnValue_ = getFunction(module, "_load_return_value");
   pySerialize_ = getFunction(module, "serialize");

diff --git a/torch/csrc/distributed/rpc/python_rpc_handler.h b/torch/csrc/distributed/rpc/python_rpc_handler.h
@@ -48,13 +48,13 @@ class PYBIND11_EXPORT PythonRpcHandler {
   PythonRpcHandler(PythonRpcHandler&&) = delete;
   PythonRpcHandler& operator=(PythonRpcHandler&&) = delete;
 
-  // Ref to `torch.distributed.internal_rpc_utils._run_function`.
+  // Ref to `torch.distributed.rpc.internal._run_function`.
   py::object pyRunFunction_;
 
-  // Ref to `torch.distributed.internal_rpc_utils._load_return_value`.
+  // Ref to `torch.distributed.rpc.internal._load_return_value`.
   py::object pyLoadReturnValue_;
 
-  // Ref to `torch.distributed.internal_rpc_utils.serialize`.
+  // Ref to `torch.distributed.rpc.internal.serialize`.
   py::object pySerialize_;
 };
 

diff --git a/torch/distributed/internal_rpc_utils.py → torch/distributed/rpc/internal.py b/torch/distributed/internal_rpc_utils.py → torch/distributed/rpc/internal.py
diff --git a/torch/distributed/rpc_api.py b/torch/distributed/rpc_api.py
@@ -4,7 +4,7 @@
 from . import ProcessGroupAgent
 from . import WorkerInfo
 from .rpc import is_backend_registered, init_backend
-from .internal_rpc_utils import _internal_rpc_pickler, PythonUDF
+from .rpc.internal import _internal_rpc_pickler, PythonUDF
 
 import functools
 import sys