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