Caffe2 - C++ API
A deep learning, cross platform ML framework
init.cpp
1 #include <torch/csrc/python_headers.h>
2 #include <torch/csrc/utils/object_ptr.h>
3 #include <torch/csrc/utils/pybind.h>
4 
5 #include <stdexcept>
6 
7 #if defined(__linux__)
8 #include <sys/prctl.h>
9 #endif
10 
11 #define SYSASSERT(rv, ...) \
12  if ((rv) < 0) { \
13  throw std::system_error(errno, std::system_category(), ##__VA_ARGS__); \
14  }
15 
16 namespace torch {
17 namespace multiprocessing {
18 
19 namespace {
20 
21 PyObject* multiprocessing_init(PyObject* _unused) {
22  auto multiprocessing_module =
23  THPObjectPtr(PyImport_ImportModule("torch.multiprocessing"));
24  if (!multiprocessing_module) {
25  throw python_error();
26  }
27 
28  auto module = py::handle(multiprocessing_module).cast<py::module>();
29 
30  module.def("_prctl_pr_set_pdeathsig", [](int signal) {
31 #if defined(__linux__)
32  auto rv = prctl(PR_SET_PDEATHSIG, signal);
33  SYSASSERT(rv, "prctl");
34 #endif
35  });
36 
37  Py_RETURN_TRUE;
38 }
39 
40 } // namespace
41 
42 // multiprocessing methods on torch._C
43 static PyMethodDef methods[] = {
44  {
45  "_multiprocessing_init",
46  (PyCFunction)multiprocessing_init,
47  METH_NOARGS,
48  nullptr,
49  },
50  {nullptr, nullptr, 0, nullptr},
51 };
52 
53 PyMethodDef* python_functions() {
54  return methods;
55 }
56 
57 } // namespace multiprocessing
58 } // namespace torch
Definition: jit_type.h:17