1 | #include <torch/csrc/monitor/events.h> |
---|---|
2 | |
3 | #include <algorithm> |
4 | #include <mutex> |
5 | #include <sstream> |
6 | #include <unordered_set> |
7 | #include <vector> |
8 | |
9 | namespace torch { |
10 | namespace monitor { |
11 | |
12 | namespace { |
13 | class EventHandlers { |
14 | public: |
15 | void registerEventHandler(std::shared_ptr<EventHandler> handler) noexcept { |
16 | std::unique_lock<std::mutex> lock(mu_); |
17 | |
18 | handlers_.emplace_back(std::move(handler)); |
19 | } |
20 | |
21 | void unregisterEventHandler( |
22 | const std::shared_ptr<EventHandler>& handler) noexcept { |
23 | std::unique_lock<std::mutex> lock(mu_); |
24 | |
25 | auto it = std::find(handlers_.begin(), handlers_.end(), handler); |
26 | handlers_.erase(it); |
27 | } |
28 | |
29 | void logEvent(const Event& e) { |
30 | std::unique_lock<std::mutex> lock(mu_); |
31 | |
32 | for (auto& handler : handlers_) { |
33 | handler->handle(e); |
34 | } |
35 | } |
36 | |
37 | static EventHandlers& get() noexcept { |
38 | static EventHandlers ehs; |
39 | return ehs; |
40 | } |
41 | |
42 | private: |
43 | std::mutex mu_{}; |
44 | std::vector<std::shared_ptr<EventHandler>> handlers_{}; |
45 | }; |
46 | } // namespace |
47 | |
48 | void logEvent(const Event& e) { |
49 | EventHandlers::get().logEvent(e); |
50 | } |
51 | |
52 | void registerEventHandler(std::shared_ptr<EventHandler> p) { |
53 | EventHandlers::get().registerEventHandler(std::move(p)); |
54 | } |
55 | |
56 | void unregisterEventHandler(const std::shared_ptr<EventHandler>& p) { |
57 | EventHandlers::get().unregisterEventHandler(p); |
58 | } |
59 | |
60 | } // namespace monitor |
61 | } // namespace torch |
62 |