1 | #include "taichi/python/memory_usage_monitor.h" |
2 | |
3 | #include "pybind11/embed.h" |
4 | #include "pybind11/pybind11.h" |
5 | #include "taichi/common/core.h" |
6 | #include "taichi/common/task.h" |
7 | #include "taichi/math/scalar.h" |
8 | #include "taichi/system/threading.h" |
9 | #include "taichi/system/timer.h" |
10 | |
11 | namespace taichi { |
12 | |
13 | namespace py = pybind11; |
14 | using namespace py::literals; |
15 | |
16 | constexpr size_t VirtualMemoryAllocator::page_size; |
17 | |
18 | float64 bytes_to_GB(float64 bytes) { |
19 | return float64(bytes) * pow<3>(1.0_f64 / 1024.0_f64); |
20 | } |
21 | |
22 | float64 get_memory_usage_gb(int pid) { |
23 | return bytes_to_GB(get_memory_usage(pid)); |
24 | } |
25 | |
26 | uint64 get_memory_usage(int pid) { |
27 | if (pid == -1) { |
28 | pid = PID::get_pid(); |
29 | } |
30 | |
31 | auto locals = py::dict("pid"_a = pid); |
32 | py::exec(R"( |
33 | import os, psutil |
34 | process = psutil.Process(pid) |
35 | mem = process.memory_info().rss)" , |
36 | py::globals(), locals); |
37 | |
38 | return locals["mem" ].cast<int64>(); |
39 | } |
40 | |
41 | MemoryMonitor::MemoryMonitor(int pid, std::string output_fn) { |
42 | log_.open(output_fn, std::ios_base::out); |
43 | locals_ = new py::dict; |
44 | (*reinterpret_cast<py::dict *>(locals_))["pid" ] = pid; |
45 | py::exec(R"( |
46 | import os, psutil |
47 | process = psutil.Process(pid))" , |
48 | py::globals(), *reinterpret_cast<py::dict *>(locals_)); |
49 | } |
50 | |
51 | MemoryMonitor::~MemoryMonitor() { |
52 | delete reinterpret_cast<py::dict *>(locals_); |
53 | } |
54 | |
55 | uint64 MemoryMonitor::get_usage() const { |
56 | py::gil_scoped_acquire acquire; |
57 | py::exec(R"( |
58 | try: |
59 | mem = process.memory_info().rss |
60 | except: |
61 | mem = -1)" , |
62 | py::globals(), *reinterpret_cast<py::dict *>(locals_)); |
63 | return (*reinterpret_cast<py::dict *>(locals_))["mem" ].cast<uint64>(); |
64 | } |
65 | |
66 | void MemoryMonitor::append_sample() { |
67 | auto t = std::chrono::system_clock::now(); |
68 | log_ << fmt::format( |
69 | "{:.5f} {}\n" , |
70 | (t.time_since_epoch() / std::chrono::nanoseconds(1)) / 1e9_f64, |
71 | get_usage()); |
72 | log_.flush(); |
73 | } |
74 | |
75 | void start_memory_monitoring(std::string output_fn, int pid, real interval) { |
76 | if (pid == -1) { |
77 | pid = PID::get_pid(); |
78 | } |
79 | TI_P(pid); |
80 | std::thread th([=]() { |
81 | MemoryMonitor monitor(pid, output_fn); |
82 | while (true) { |
83 | monitor.append_sample(); |
84 | Time::sleep(interval); |
85 | } |
86 | }); |
87 | th.detach(); |
88 | } |
89 | |
90 | class MemoryTest : public Task { |
91 | public: |
92 | std::string run(const std::vector<std::string> ¶meters) override { |
93 | TI_P(get_memory_usage()); |
94 | Time::sleep(3); |
95 | std::vector<uint8> a(1024ul * 1024 * 1024 * 10, 3); |
96 | TI_P(get_memory_usage()); |
97 | Time::sleep(3); |
98 | return "" ; |
99 | } |
100 | }; |
101 | |
102 | class MemoryTest2 : public Task { |
103 | public: |
104 | std::string run(const std::vector<std::string> ¶meters) override { |
105 | start_memory_monitoring("test.txt" ); |
106 | std::vector<uint8> a; |
107 | for (int i = 0; i < 10; i++) { |
108 | a.resize(1024ul * 1024 * 1024 * i / 2); |
109 | std::fill(std::begin(a), std::end(a), 3); |
110 | Time::sleep(0.5); |
111 | } |
112 | return "" ; |
113 | } |
114 | }; |
115 | |
116 | TI_IMPLEMENTATION(Task, MemoryTest, "mem_test" ); |
117 | TI_IMPLEMENTATION(Task, MemoryTest2, "mem_test2" ); |
118 | |
119 | } // namespace taichi |
120 | |