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
11namespace taichi {
12
13namespace py = pybind11;
14using namespace py::literals;
15
16constexpr size_t VirtualMemoryAllocator::page_size;
17
18float64 bytes_to_GB(float64 bytes) {
19 return float64(bytes) * pow<3>(1.0_f64 / 1024.0_f64);
20}
21
22float64 get_memory_usage_gb(int pid) {
23 return bytes_to_GB(get_memory_usage(pid));
24}
25
26uint64 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
41MemoryMonitor::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
51MemoryMonitor::~MemoryMonitor() {
52 delete reinterpret_cast<py::dict *>(locals_);
53}
54
55uint64 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
66void 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
75void 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
90class MemoryTest : public Task {
91 public:
92 std::string run(const std::vector<std::string> &parameters) 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
102class MemoryTest2 : public Task {
103 public:
104 std::string run(const std::vector<std::string> &parameters) 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
116TI_IMPLEMENTATION(Task, MemoryTest, "mem_test");
117TI_IMPLEMENTATION(Task, MemoryTest2, "mem_test2");
118
119} // namespace taichi
120