Rename variables and add comments (#27286)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/27286
The name `runUDFFunction` stutters because the F in UDF also stands
for function. Renamed these variables to be identical to their Python
equivalents. Renamed those to share a prefix and drop `internal`,
because internal functions can use an underscore prefix.
Test Plan: Imported from OSS
Differential Revision: D17808208
Pulled By: pietern
fbshipit-source-id: 7619f07fc8215203dfb1da1eb281845edcd2bb99
diff --git a/torch/csrc/distributed/rpc/python_rpc_handler.cpp b/torch/csrc/distributed/rpc/python_rpc_handler.cpp
index 34290b2..7331ba4 100644
--- a/torch/csrc/distributed/rpc/python_rpc_handler.cpp
+++ b/torch/csrc/distributed/rpc/python_rpc_handler.cpp
@@ -22,9 +22,9 @@
AutoGIL ag;
py::object module =
py::module::import("torch.distributed.internal_rpc_utils");
- runUDFFunction_ = getFunction(module, "run_python_udf_internal");
- loadResultFunction_ = getFunction(module, "load_python_udf_result_internal");
- serializeFunction_ = getFunction(module, "serialize");
+ pyRunFunction_ = getFunction(module, "_run_function");
+ pyLoadReturnValue_ = getFunction(module, "_load_return_value");
+ pySerialize_ = getFunction(module, "serialize");
}
PythonRpcHandler& PythonRpcHandler::getInstance() {
@@ -38,8 +38,7 @@
std::vector<torch::Tensor>& responseTensorTable) {
AutoGIL ag;
auto pargs = py::bytes(pickledPayload.data(), pickledPayload.size());
- py::tuple pres =
- serializeFunction_(runUDFFunction_(pargs, requestTensorTable));
+ py::tuple pres = pySerialize_(pyRunFunction_(pargs, requestTensorTable));
const auto& presStr = pres[0].cast<std::string>();
responseTensorTable = pres[1].cast<std::vector<torch::Tensor>>();
std::vector<char> payload(presStr.begin(), presStr.end());
@@ -51,26 +50,26 @@
const std::vector<torch::Tensor>& tensorTable) {
AutoGIL ag;
auto pargs = py::bytes(pickledPayload.data(), pickledPayload.size());
- return loadResultFunction_(pargs, tensorTable);
+ return pyLoadReturnValue_(pargs, tensorTable);
}
py::object PythonRpcHandler::runPythonUDF(
const SerializedPyObj& serializedObj) {
AutoGIL ag;
- return runUDFFunction_(
+ return pyRunFunction_(
py::bytes(serializedObj.payload_), serializedObj.tensors_);
}
SerializedPyObj PythonRpcHandler::serialize(const py::object& obj) {
AutoGIL ag;
- py::tuple t = serializeFunction_(obj);
+ py::tuple t = pySerialize_(obj);
return SerializedPyObj(
t[0].cast<std::string>(), t[1].cast<std::vector<torch::Tensor>>());
}
py::object PythonRpcHandler::deserialize(const SerializedPyObj& serializedObj) {
AutoGIL ag;
- return loadResultFunction_(
+ return pyLoadReturnValue_(
py::bytes(serializedObj.payload_), serializedObj.tensors_);
}
diff --git a/torch/csrc/distributed/rpc/python_rpc_handler.h b/torch/csrc/distributed/rpc/python_rpc_handler.h
index ce55115..612de07 100644
--- a/torch/csrc/distributed/rpc/python_rpc_handler.h
+++ b/torch/csrc/distributed/rpc/python_rpc_handler.h
@@ -17,20 +17,25 @@
class PYBIND11_EXPORT PythonRpcHandler {
public:
static PythonRpcHandler& getInstance();
- // Execute python UDF, result is pickled to binary string
+
+ // Deserialize Python function, run it, and serialize its return value.
std::vector<char> generatePythonUDFResult(
const std::vector<char>& pickledPayload,
const std::vector<torch::Tensor>& requestTensorTable,
std::vector<torch::Tensor>& responseTensorTable);
+
// Returned python UDF result is pickled binary string, so run python
// function to unpickle the python UDF result and return py::object to user
py::object loadPythonUDFResult(
const std::vector<char>& pickledPayload,
const std::vector<torch::Tensor>& tensorTable);
+
// Run a pickled Python UDF and return the result py::object
py::object runPythonUDF(const SerializedPyObj& serializedObj);
+
// Serialized a py::object into a string
SerializedPyObj serialize(const py::object& obj);
+
// Deserialize a string into a py::object
py::object deserialize(const SerializedPyObj& serializedObj);
@@ -43,9 +48,14 @@
PythonRpcHandler(PythonRpcHandler&&) = delete;
PythonRpcHandler& operator=(PythonRpcHandler&&) = delete;
- py::object runUDFFunction_;
- py::object loadResultFunction_;
- py::object serializeFunction_;
+ // Ref to `torch.distributed.internal_rpc_utils._run_function`.
+ py::object pyRunFunction_;
+
+ // Ref to `torch.distributed.internal_rpc_utils._load_return_value`.
+ py::object pyLoadReturnValue_;
+
+ // Ref to `torch.distributed.internal_rpc_utils.serialize`.
+ py::object pySerialize_;
};
} // namespace rpc
diff --git a/torch/distributed/internal_rpc_utils.py b/torch/distributed/internal_rpc_utils.py
index 3a9f4b3..e0e9e0a 100644
--- a/torch/distributed/internal_rpc_utils.py
+++ b/torch/distributed/internal_rpc_utils.py
@@ -97,32 +97,38 @@
# Create _internal_rpc_pickler only once to initialize _dispatch_table only once
_internal_rpc_pickler = _InternalRPCPickler()
+
def serialize(obj):
return _internal_rpc_pickler.serialize(obj)
-def run_python_udf_internal(pickled_python_udf, tensors):
+
+def _run_function(binary_data, tensor_table):
r"""
- Internal python function will be imported and executed in C++ land
- it unpickles pickled python udf strings and tensors and run the python
- udf, return serialized result and tensor tables
+ This function is exclusively called from C++.
+ See ``torch/csrc/distributed/rpc/python_rpc_handler.cpp``.
+
+ Runs a Python UDF and returns its return value.
+ Wraps any exception in ``RemoteException`` if the function raises.
"""
- python_udf = _internal_rpc_pickler.deserialize(pickled_python_udf, tensors)
+ python_udf = _internal_rpc_pickler.deserialize(binary_data, tensor_table)
try:
result = python_udf.func(*python_udf.args, **python_udf.kwargs)
except Exception as e:
# except str = exception info + traceback string
except_str = "{}\n{}".format(repr(e), traceback.format_exc())
result = RemoteException(except_str)
- # return _internal_rpc_pickler.serialize(result)
return result
-def load_python_udf_result_internal(pickled_python_result, tensors):
+def _load_return_value(binary_data, tensor_table):
r"""
- Internal python function will be imported and executed in C++ land
- it unpickled pickled python udf result and tensor tables, return python object
+ This function is exclusively called from C++.
+ See ``torch/csrc/distributed/rpc/python_rpc_handler.cpp``.
+
+ Processes the return value of a Python function.
+ Raises exception if the return value is a wrapped exception.
"""
- result = _internal_rpc_pickler.deserialize(pickled_python_result, tensors)
+ result = _internal_rpc_pickler.deserialize(binary_data, tensor_table)
if isinstance(result, RemoteException):
raise Exception(result.msg)
return result
