6d8877a6d4* feat(search): per-profile simeon bandit with training-free self-tuning * fix(zyp): free parallel buffers correctly * fix: code cleanup and simeon version bump * fix(search): guard Simeon bandit state
125 lines
3.9 KiB
C++
125 lines
3.9 KiB
C++
// SPDX-License-Identifier: GPL-3.0-or-later
|
||
// Copyright 2025 YAMS Contributors
|
||
//
|
||
// Micro-benchmarks for simeon bandit overhead.
|
||
// Measures: arm selection, reward recording, init, and dispatch latency.
|
||
|
||
#include <catch2/benchmark/catch_benchmark.hpp>
|
||
#include <catch2/catch_test_macros.hpp>
|
||
|
||
#include <chrono>
|
||
#include <string>
|
||
#include <unordered_map>
|
||
#include <vector>
|
||
|
||
#include <yams/search/search_engine_config.h>
|
||
#include <yams/search/tuner_mab.h>
|
||
|
||
using namespace yams::search;
|
||
|
||
namespace {
|
||
|
||
std::vector<TunerMAB::Arm> standardArms() {
|
||
std::vector<TunerMAB::Arm> arms;
|
||
arms.push_back({"sab_smooth", 0.0, {}});
|
||
arms.push_back({"sab_smooth_rm3_adaptive", 0.0, {}});
|
||
arms.push_back({"sab_smooth_rm3_diverse", 0.0, {}});
|
||
arms.push_back({"keyphrase", 0.0, {}});
|
||
arms.push_back({"lead_field", 0.0, {}});
|
||
return arms;
|
||
}
|
||
|
||
} // namespace
|
||
|
||
TEST_CASE("Bench: bandit arm selection on warm bandit", "[.bench][search][bandit][bench]") {
|
||
TunerMAB mab;
|
||
mab.setArms(standardArms());
|
||
|
||
// Warm up: pull each arm once.
|
||
for (size_t i = 0; i < mab.arms().size(); ++i) {
|
||
auto idx = mab.selectArm();
|
||
mab.recordReward(*idx, 0.5, TunerMAB::RewardSource::Proxy);
|
||
}
|
||
|
||
BENCHMARK("selectArm warm") {
|
||
return mab.selectArm();
|
||
};
|
||
}
|
||
|
||
TEST_CASE("Bench: bandit reward recording", "[.bench][search][bandit][bench]") {
|
||
TunerMAB mab;
|
||
mab.setArms(standardArms());
|
||
|
||
// Ensure at least one arm has been pulled.
|
||
auto idx = mab.selectArm();
|
||
REQUIRE(idx.has_value());
|
||
|
||
BENCHMARK("recordReward") {
|
||
mab.recordReward(0, 0.5, TunerMAB::RewardSource::Proxy);
|
||
};
|
||
}
|
||
|
||
TEST_CASE("Bench: bandit init with 5 profiles × 5 arms", "[.bench][search][bandit][bench]") {
|
||
using ProfileKey = std::string;
|
||
|
||
BENCHMARK("init 5x5 bandits") {
|
||
std::unordered_map<ProfileKey, TunerMAB> bandits;
|
||
using CP = SearchEngineConfig::CorpusProfile;
|
||
for (auto profile : {CP::CODE, CP::PROSE, CP::DOCS, CP::MIXED, CP::CUSTOM}) {
|
||
bandits[SearchEngineConfig::corpusProfileToString(profile)].setArms(standardArms());
|
||
}
|
||
return bandits.size();
|
||
};
|
||
}
|
||
|
||
TEST_CASE("Bench: per-profile arm selection + reward", "[.bench][search][bandit][bench]") {
|
||
using ProfileKey = std::string;
|
||
std::unordered_map<ProfileKey, TunerMAB> bandits;
|
||
using CP = SearchEngineConfig::CorpusProfile;
|
||
for (auto profile : {CP::CODE, CP::PROSE, CP::DOCS, CP::MIXED, CP::CUSTOM}) {
|
||
bandits[SearchEngineConfig::corpusProfileToString(profile)].setArms(standardArms());
|
||
}
|
||
|
||
// Warm all bandits.
|
||
for (auto& [_, mab] : bandits) {
|
||
for (size_t i = 0; i < mab.arms().size(); ++i) {
|
||
auto idx = mab.selectArm();
|
||
mab.recordReward(*idx, 0.5, TunerMAB::RewardSource::Proxy);
|
||
}
|
||
}
|
||
|
||
BENCHMARK("select+record per-profile") {
|
||
std::string profileKey = "MIXED";
|
||
auto it = bandits.find(profileKey);
|
||
auto idx = it->second.selectArm();
|
||
if (idx.has_value()) {
|
||
it->second.recordReward(*idx, 0.5, TunerMAB::RewardSource::Proxy);
|
||
}
|
||
return idx.has_value();
|
||
};
|
||
}
|
||
|
||
TEST_CASE("Bench: corpusProfileToString lookup", "[.bench][search][bandit][bench]") {
|
||
BENCHMARK("corpusProfileToString") {
|
||
return SearchEngineConfig::corpusProfileToString(SearchEngineConfig::CorpusProfile::MIXED);
|
||
};
|
||
}
|
||
|
||
TEST_CASE("Bench: bandit UCB1 score computation (150 iterations)",
|
||
"[.bench][search][bandit][bench]") {
|
||
TunerMAB mab;
|
||
mab.setArms(standardArms());
|
||
|
||
// Warm up: 10 pulls per arm with varying rewards.
|
||
for (int round = 0; round < 10; ++round) {
|
||
for (size_t i = 0; i < mab.arms().size(); ++i) {
|
||
auto idx = mab.selectArm();
|
||
double reward = (i == 0) ? 0.9 : 0.5; // arm 0 is best
|
||
mab.recordReward(*idx, reward, TunerMAB::RewardSource::Proxy);
|
||
}
|
||
}
|
||
|
||
BENCHMARK("bestArmId after convergence") {
|
||
return mab.bestArmId();
|
||
};
|
||
}
|