1 | #pragma once |
2 | |
3 | #include <chrono> |
4 | #include <string> |
5 | #include <unordered_map> |
6 | |
7 | #include <c10/macros/Macros.h> |
8 | #include <c10/util/variant.h> |
9 | |
10 | namespace torch { |
11 | namespace monitor { |
12 | |
13 | // data_value_t is the type for Event data values. |
14 | using data_value_t = c10::variant<std::string, double, int64_t, bool>; |
15 | |
16 | // Event represents a single event that can be logged out to an external |
17 | // tracker. This does acquire a lock on logging so should be used relatively |
18 | // infrequently to avoid performance issues. |
19 | struct TORCH_API Event { |
20 | // name is the name of the event. This is a static string that's used to |
21 | // differentiate between event types for programmatic access. The type should |
22 | // be in the format of a fully qualified Python-style class name. |
23 | // Ex: torch.monitor.MonitorEvent |
24 | std::string name; |
25 | |
26 | // timestamp is a timestamp relative to the Unix epoch time. |
27 | std::chrono::system_clock::time_point timestamp; |
28 | |
29 | // data contains rich information about the event. The contents are event |
30 | // specific so you should check the type to ensure it's what you expect before |
31 | // accessing the data. |
32 | // |
33 | // NOTE: these events are not versioned and it's up to the consumer of the |
34 | // events to check the fields to ensure backwards compatibility. |
35 | std::unordered_map<std::string, data_value_t> data; |
36 | }; |
37 | |
38 | TORCH_API inline bool operator==(const Event& lhs, const Event& rhs) { |
39 | return lhs.name == rhs.name && lhs.timestamp == rhs.timestamp && |
40 | lhs.data == rhs.data; |
41 | } |
42 | |
43 | // EventHandler represents an abstract event handler that can be registered to |
44 | // capture events. Every time an event is logged every handler will be called |
45 | // with the events contents. |
46 | // |
47 | // NOTE: The handlers should avoid any IO, blocking calls or heavy computation |
48 | // as this may block the main thread and cause performance issues. |
49 | class TORCH_API EventHandler { |
50 | public: |
51 | virtual ~EventHandler() = default; |
52 | |
53 | // handle needs to be implemented to handle the events. This may be called |
54 | // from multiple threads so needs to be thread safe. |
55 | virtual void handle(const Event& e) = 0; |
56 | }; |
57 | |
58 | // logEvent calls each registered event handler with the event. This method can |
59 | // be called from concurrently from multiple threads. |
60 | TORCH_API void logEvent(const Event& e); |
61 | |
62 | // registerEventHandler registers an EventHandler so it receives any logged |
63 | // events. Typically an EventHandler will be registered during program |
64 | // setup and unregistered at the end. |
65 | TORCH_API void registerEventHandler(std::shared_ptr<EventHandler> p); |
66 | |
67 | // unregisterEventHandler unregisters the event handler pointed to by the |
68 | // shared_ptr. |
69 | TORCH_API void unregisterEventHandler(const std::shared_ptr<EventHandler>& p); |
70 | |
71 | } // namespace monitor |
72 | } // namespace torch |
73 | |