GlobalCallbackManager Class — pytorch Architecture

Source Code

aten/src/ATen/record_function.cpp lines 102–125

class GlobalCallbackManager {
 public:
  static GlobalCallbackManager& get(); // Singleton

 private:
  GlobalCallbackManager() = default;

 public:
  static constexpr size_t NoVersion = 0;
  using snapshot_t = std::pair<size_t, RecordFunctionCallbacks>;

  //                                                                Locking?
  size_t version() const; //                                     No
  snapshot_t getSnapshot() const; //                                Yes
  CallbackHandle addCallback(RecordFunctionCallback cb); //         Yes
  void setCallbackEnabled(CallbackHandle handle, bool enabled); //  Yes
  void removeCallback(CallbackHandle handle); //                    Yes
  void clearCallbacks(); //                                         Yes

 private:
  std::atomic<size_t> version_{NoVersion + 1};
  RecordFunctionCallbacks global_callbacks_; // Source of truth.
  mutable std::mutex update_mutex_;
};