mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
95dbbf713e
Follows #125102 Pull Request resolved: https://github.com/pytorch/pytorch/pull/130109 Approved by: https://github.com/ezyang
39 lines
1.4 KiB
C++
39 lines
1.4 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/distributed/rpc/rpc_command_base.h>
|
|
#include <torch/csrc/distributed/rpc/types.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
|
|
namespace torch::distributed::rpc {
|
|
|
|
// This class converts the content in a PythonCall into py::object. This is a
|
|
// helper class to make sure that all arguments deserialization is done before
|
|
// entering RequestCallbackImpl::processRpc(...), so that the deserialization
|
|
// related logic can be carried out in one spot instead of scattered in multiple
|
|
// places for different message types.
|
|
// NB: The reason for not consolidating class into PythonCall is because
|
|
// PythonCall is a libtorch type which should not depend on Python types.
|
|
class TORCH_API UnpickledPythonCall : public RpcCommandBase {
|
|
public:
|
|
UnpickledPythonCall(
|
|
const SerializedPyObj& serializedPyObj,
|
|
bool isAsyncExecution);
|
|
~UnpickledPythonCall() override;
|
|
|
|
// toMessage() method is not implemented, as objects of this class should
|
|
// never be directly converted into a Message object.
|
|
c10::intrusive_ptr<Message> toMessageImpl() && override;
|
|
const py::object& pythonUdf() const;
|
|
|
|
inline bool isAsyncExecution() const {
|
|
return isAsyncExecution_;
|
|
}
|
|
|
|
private:
|
|
py::object pythonUdf_;
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members)
|
|
const bool isAsyncExecution_;
|
|
};
|
|
|
|
} // namespace torch::distributed::rpc
|