refactor: Add metrics collector

python/pylibvw.cc

@@ -327,10 +327,9 @@ py::dict get_learner_metrics(vw_ptr all)
 {
   py::dict dictionary;
 
-  if (all->options->was_supplied("extra_metrics"))
+  if (all->global_metrics.are_metrics_enabled())
   {
-    VW::metric_sink metrics;
-    all->l->persist_metrics(metrics);
+    auto metrics = all->global_metrics.collect_metrics(all->l);
 
     python_dict_writer writer(dictionary);
     metrics.visit(writer);

vowpalwabbit/core/CMakeLists.txt

@@ -56,6 +56,7 @@ set(vw_core_headers
   include/vw/core/loss_functions.h
   include/vw/core/memory.h
   include/vw/core/merge.h
+  include/vw/core/metrics_collector.h
   include/vw/core/metric_sink.h
   include/vw/core/model_utils.h
   include/vw/core/multi_model_utils.h
@@ -228,6 +229,7 @@ set(vw_core_sources
   src/learner.cc
   src/loss_functions.cc
   src/merge.cc
+  src/metrics_collector.cc
   src/metric_sink.cc
   src/multiclass.cc
   src/multilabel.cc

vowpalwabbit/core/include/vw/core/global_data.h

@@ -11,6 +11,7 @@
 #include "vw/core/error_reporting.h"
 #include "vw/core/input_parser.h"
 #include "vw/core/interaction_generation_state.h"
+#include "vw/core/metrics_collector.h"
 #include "vw/core/multi_ex.h"
 #include "vw/core/version.h"
 #include "vw/core/vw_fwd.h"
@@ -151,9 +152,7 @@ class workspace
   std::unique_ptr<VW::parsers::flatbuffer::parser> flat_converter;
 #endif
 
-  // This field is experimental and subject to change.
-  // Used to implement the external binary parser.
-  std::vector<std::function<void(VW::metric_sink&)>> metric_output_hooks;
+  VW::metrics_collector global_metrics;
 
   // Experimental field.
   // Generic parser interface to make it possible to use any external parser.

vowpalwabbit/core/include/vw/core/metrics_collector.h

@@ -0,0 +1,36 @@
+// Copyright (c) by respective owners including Yahoo!, Microsoft, and
+// individual contributors. All rights reserved. Released under a BSD (revised)
+// license as described in the file LICENSE.=
+
+#pragma once
+
+#include "metric_sink.h"
+#include "vw/core/learner_fwd.h"
+
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+namespace VW
+{
+class metrics_collector
+{
+public:
+  metrics_collector(bool enabled = false, std::string filename = "");
+
+  using metrics_callback_fn = std::function<void(VW::metric_sink&)>;
+
+  bool are_metrics_enabled() const;
+  void register_metrics_callback(const metrics_callback_fn& callback);
+  VW::metric_sink collect_metrics(LEARNER::base_learner* l = nullptr) const;
+
+private:
+  bool _are_metrics_enabled;
+  std::vector<metrics_callback_fn> _metrics_callbacks;
+  std::string _metrics_filename;
+};
+}  // namespace VW

vowpalwabbit/core/include/vw/core/reductions/metrics.h

@@ -11,5 +11,6 @@ namespace reductions
 {
 VW::LEARNER::base_learner* metrics_setup(VW::setup_base_i& stack_builder);
 void output_metrics(VW::workspace& all);
+void additional_metrics(VW::workspace& all, VW::metric_sink& sink);
 }  // namespace reductions
 }  // namespace VW

vowpalwabbit/core/src/metrics_collector.cc

@@ -0,0 +1,30 @@
+// Copyright (c) by respective owners including Yahoo!, Microsoft, and
+// individual contributors. All rights reserved. Released under a BSD (revised)
+// license as described in the file LICENSE.
+
+#include "vw/core/metrics_collector.h"
+
+#include "vw/core/learner.h"
+namespace VW
+{
+metrics_collector::metrics_collector(bool enabled, std::string filename)
+    : _are_metrics_enabled(enabled), _metrics_filename(filename)
+{
+}
+
+bool metrics_collector::are_metrics_enabled() const { return _are_metrics_enabled; }
+
+void metrics_collector::register_metrics_callback(const metrics_callback_fn& callback)
+{
+  if (_are_metrics_enabled) { _metrics_callbacks.push_back(callback); }
+}
+
+metric_sink metrics_collector::collect_metrics(LEARNER::base_learner* l) const
+{
+  VW::metric_sink sink;
+  if (!_are_metrics_enabled) { THROW("Metrics must be enabled to call collect_metrics"); }
+  if (l) { l->persist_metrics(sink); }
+  for (const auto& callback : _metrics_callbacks) { callback(sink); }
+  return sink;
+}
+}  // namespace VW

vowpalwabbit/core/src/parse_args.cc

@@ -2143,6 +2143,8 @@ void finish(VW::workspace& all, bool delete_all)
     auto writer = VW::io::open_file_writer(all.json_weights_file_name);
     writer->write(content.c_str(), content.length());
   }
+  all.global_metrics.register_metrics_callback(
+      [&all](VW::metric_sink& sink) -> void { VW::reductions::additional_metrics(all, sink); });
   VW::reductions::output_metrics(all);
   all.logger.log_summary();
 }

vowpalwabbit/core/src/parser.cc

@@ -213,7 +213,7 @@ void set_json_reader(VW::workspace& all, bool dsjson = false)
 
   all.example_parser->decision_service_json = dsjson;
 
-  if (dsjson && all.options->was_supplied("extra_metrics"))
+  if (dsjson && all.global_metrics.are_metrics_enabled())
   {
     all.example_parser->metrics = VW::make_unique<VW::details::dsjson_metrics>();
   }

vowpalwabbit/core/src/reductions/baseline_challenger_cb.cc

@@ -64,12 +64,8 @@ class baseline_challenger_data
   VW::estimators::ChiSquared baseline;
   discounted_expectation policy_expectation;
   float baseline_epsilon;
-  bool emit_metrics;
 
-  baseline_challenger_data(bool emit_metrics, double alpha, double tau)
-      : baseline(alpha, tau), policy_expectation(tau), emit_metrics(emit_metrics)
-  {
-  }
+  baseline_challenger_data(double alpha, double tau) : baseline(alpha, tau), policy_expectation(tau) {}
 
   static int get_chosen_action(const VW::action_scores& action_scores) { return action_scores[0].action; }
 
@@ -176,7 +172,6 @@ void save_load(baseline_challenger_data& data, VW::io_buf& io, bool read, bool t
 
 void persist_metrics(baseline_challenger_data& data, metric_sink& metrics)
 {
-  if (!data.emit_metrics) { return; }
   auto ci = static_cast<float>(data.baseline.lower_bound_and_update());
   auto exp = static_cast<float>(data.policy_expectation.current());
 
@@ -188,6 +183,7 @@ void persist_metrics(baseline_challenger_data& data, metric_sink& metrics)
 VW::LEARNER::base_learner* VW::reductions::baseline_challenger_cb_setup(VW::setup_base_i& stack_builder)
 {
   options_i& options = *stack_builder.get_options();
+
   float alpha;
   float tau;
   bool is_enabled = false;
@@ -215,8 +211,7 @@ VW::LEARNER::base_learner* VW::reductions::baseline_challenger_cb_setup(VW::setu
 
   if (!options.was_supplied("cb_adf")) { THROW("cb_challenger requires cb_explore_adf or cb_adf"); }
 
-  bool emit_metrics = options.was_supplied("extra_metrics");
-  auto data = VW::make_unique<baseline_challenger_data>(emit_metrics, alpha, tau);
+  auto data = VW::make_unique<baseline_challenger_data>(alpha, tau);
 
   auto* l = make_reduction_learner(std::move(data), as_multiline(stack_builder.setup_base_learner()),
       learn_or_predict<true>, learn_or_predict<false>, stack_builder.get_setupfn_name(baseline_challenger_cb_setup))

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_bag.cc

@@ -206,11 +206,9 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_bag_setup(VW::setup_ba
   VW::LEARNER::multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_bag>;
-  auto data = VW::make_unique<explore_type>(
-      with_metrics, epsilon, VW::cast_to_smaller_type<size_t>(bag_size), greedify, first_only, all.get_random_state());
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), epsilon,
+      VW::cast_to_smaller_type<size_t>(bag_size), greedify, first_only, all.get_random_state());
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,
       stack_builder.get_setupfn_name(cb_explore_adf_bag_setup))
                 .set_input_label_type(VW::label_type_t::CB)

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_cover.cc

@@ -323,14 +323,12 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_cover_setup(VW::setup_
     epsilon_decay = true;
   }
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   auto* scorer = VW::LEARNER::as_singleline(base->get_learner_by_name_prefix("scorer"));
 
   using explore_type = cb_explore_adf_base<cb_explore_adf_cover>;
-  auto data =
-      VW::make_unique<explore_type>(with_metrics, VW::cast_to_smaller_type<size_t>(cover_size), psi, nounif, epsilon,
-          epsilon_decay, first_only, as_multiline(all.cost_sensitive), scorer, cb_type, all.model_file_ver, all.logger);
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(),
+      VW::cast_to_smaller_type<size_t>(cover_size), psi, nounif, epsilon, epsilon_decay, first_only,
+      as_multiline(all.cost_sensitive), scorer, cb_type, all.model_file_ver, all.logger);
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,
       stack_builder.get_setupfn_name(cb_explore_adf_cover_setup))
                 .set_input_label_type(VW::label_type_t::CB)

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_first.cc

@@ -122,11 +122,9 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_first_setup(VW::setup_
   multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_first>;
-  auto data =
-      VW::make_unique<explore_type>(with_metrics, VW::cast_to_smaller_type<size_t>(tau), epsilon, all.model_file_ver);
+  auto data = VW::make_unique<explore_type>(
+      all.global_metrics.are_metrics_enabled(), VW::cast_to_smaller_type<size_t>(tau), epsilon, all.model_file_ver);
 
   if (epsilon < 0.0 || epsilon > 1.0) { THROW("The value of epsilon must be in [0,1]"); }
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_greedy.cc

@@ -133,10 +133,8 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_greedy_setup(VW::setup
   VW::LEARNER::multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_greedy>;
-  auto data = VW::make_unique<explore_type>(with_metrics, epsilon, first_only);
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), epsilon, first_only);
 
   if (epsilon < 0.0 || epsilon > 1.0) { THROW("The value of epsilon must be in [0,1]"); }
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_large_action_space.cc

@@ -406,20 +406,19 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_large_action_space_set
   VW::LEARNER::multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   if (use_two_pass_svd_impl)
   {
     auto impl_type = implementation_type::two_pass_svd;
     return make_las_with_impl<two_pass_svd_impl, one_rank_spanner_state>(stack_builder, base, impl_type, all,
-        with_metrics, d, gamma_scale, gamma_exponent, c, apply_shrink_factor, thread_pool_size, block_size,
+        all.global_metrics.are_metrics_enabled(), d, gamma_scale, gamma_exponent, c, apply_shrink_factor,
+        thread_pool_size, block_size,
         /*use_explicit_simd=*/false);
   }
   else
   {
     auto impl_type = implementation_type::one_pass_svd;
     return make_las_with_impl<one_pass_svd_impl, one_rank_spanner_state>(stack_builder, base, impl_type, all,
-        with_metrics, d, gamma_scale, gamma_exponent, c, apply_shrink_factor, thread_pool_size, block_size,
-        use_simd_in_one_pass_svd_impl);
+        all.global_metrics.are_metrics_enabled(), d, gamma_scale, gamma_exponent, c, apply_shrink_factor,
+        thread_pool_size, block_size, use_simd_in_one_pass_svd_impl);
   }
 }

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_regcb.cc

@@ -303,11 +303,9 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_regcb_setup(VW::setup_
   multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_regcb>;
   auto data = VW::make_unique<explore_type>(
-      with_metrics, regcbopt, c0, first_only, min_cb_cost, max_cb_cost, all.model_file_ver);
+      all.global_metrics.are_metrics_enabled(), regcbopt, c0, first_only, min_cb_cost, max_cb_cost, all.model_file_ver);
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,
       stack_builder.get_setupfn_name(cb_explore_adf_regcb_setup))
                 .set_input_label_type(VW::label_type_t::CB)

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_rnd.cc

@@ -304,11 +304,9 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_rnd_setup(VW::setup_ba
   multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_rnd>;
-  auto data = VW::make_unique<explore_type>(
-      with_metrics, epsilon, alpha, invlambda, numrnd, base->increment * problem_multiplier, &all);
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), epsilon, alpha, invlambda, numrnd,
+      base->increment * problem_multiplier, &all);
 
   if (epsilon < 0.0 || epsilon > 1.0) { THROW("The value of epsilon must be in [0,1]"); }
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_softmax.cc

@@ -93,10 +93,8 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_softmax_setup(VW::setu
   VW::LEARNER::multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_softmax>;
-  auto data = VW::make_unique<explore_type>(with_metrics, epsilon, lambda);
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), epsilon, lambda);
 
   if (epsilon < 0.0 || epsilon > 1.0) { THROW("The value of epsilon must be in [0,1]"); }
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_squarecb.cc

@@ -381,13 +381,11 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_squarecb_setup(VW::set
   multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   if (epsilon < 0.0 || epsilon > 1.0) { THROW("The value of epsilon must be in [0,1]"); }
 
   using explore_type = cb_explore_adf_base<cb_explore_adf_squarecb>;
-  auto data = VW::make_unique<explore_type>(
-      with_metrics, gamma_scale, gamma_exponent, elim, c0, min_cb_cost, max_cb_cost, all.model_file_ver, epsilon);
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), gamma_scale, gamma_exponent, elim,
+      c0, min_cb_cost, max_cb_cost, all.model_file_ver, epsilon);
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,
       stack_builder.get_setupfn_name(cb_explore_adf_squarecb_setup))
                 .set_input_label_type(VW::label_type_t::CB)

vowpalwabbit/core/src/reductions/cb/cb_explore_adf_synthcover.cc

@@ -199,10 +199,8 @@ VW::LEARNER::base_learner* VW::reductions::cb_explore_adf_synthcover_setup(VW::s
   VW::LEARNER::multi_learner* base = as_multiline(stack_builder.setup_base_learner());
   all.example_parser->lbl_parser = CB::cb_label;
 
-  bool with_metrics = options.was_supplied("extra_metrics");
-
   using explore_type = cb_explore_adf_base<cb_explore_adf_synthcover>;
-  auto data = VW::make_unique<explore_type>(with_metrics, epsilon, psi,
+  auto data = VW::make_unique<explore_type>(all.global_metrics.are_metrics_enabled(), epsilon, psi,
       VW::cast_to_smaller_type<size_t>(synthcoversize), all.get_random_state(), all.model_file_ver);
   auto* l = make_reduction_learner(std::move(data), base, explore_type::learn, explore_type::predict,
       stack_builder.get_setupfn_name(cb_explore_adf_synthcover_setup))