-
Notifications
You must be signed in to change notification settings - Fork 7
perf: add experimental cityhash option for bandit shuffling #394
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,11 @@ | ||
| --- | ||
| "python-sdk": patch | ||
| "eppo_core": patch | ||
| "ruby-sdk": patch | ||
| "elixir-sdk": patch | ||
| "rust-sdk": patch | ||
| --- | ||
|
|
||
| Added experimental support for CityHash-based hashing in bandit evaluation via the `EPPO_EXPERIMENTAL_BANDITS_CITYHASH` environment variable (set to `"1"`, `"true"`, or `"TRUE"` to enable). This provides significant performance improvements over the default MD5 implementation, especially when evaluating bandits with many actions. | ||
|
|
||
| **Warning**: This feature is experimental and unstable. Enabling CityHash will produce different bandit evaluation results compared to the default MD5 implementation and other Eppo SDKs. Do not enable this if you need consistent results across multiple SDKs, services, or for historical data comparisons. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,6 +8,13 @@ use crate::{ | |
| Str, | ||
| }; | ||
|
|
||
| /// Hashing algorithm to use for bandit evaluation. | ||
| #[derive(Debug, Clone, Copy)] | ||
| pub(crate) enum BanditHashingAlgorithm { | ||
| Md5, | ||
| CityHash, | ||
| } | ||
|
|
||
| /// Remote configuration for the eppo client. It's a central piece that defines client behavior. | ||
| #[derive(Debug)] | ||
| pub struct Configuration { | ||
|
|
@@ -17,6 +24,8 @@ pub struct Configuration { | |
| pub(crate) flags: UniversalFlagConfig, | ||
| /// Bandits configuration. | ||
| pub(crate) bandits: Option<BanditResponse>, | ||
| /// Hashing algorithm for bandit evaluation. | ||
| pub(crate) bandit_hashing_algorithm: BanditHashingAlgorithm, | ||
| } | ||
|
|
||
| impl Configuration { | ||
|
|
@@ -27,10 +36,23 @@ impl Configuration { | |
| ) -> Configuration { | ||
| let now = Utc::now(); | ||
|
|
||
| // Check environment variable for experimental CityHash support | ||
| let bandit_hashing_algorithm = std::env::var("EPPO_EXPERIMENTAL_BANDITS_CITYHASH") | ||
| .ok() | ||
| .and_then(|val| { | ||
| if val == "1" || val == "true" || val == "TRUE" { | ||
| Some(BanditHashingAlgorithm::CityHash) | ||
| } else { | ||
| None | ||
| } | ||
| }) | ||
| .unwrap_or(BanditHashingAlgorithm::Md5); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nice use of strategy pattern! |
||
|
|
||
| Configuration { | ||
| fetched_at: now, | ||
| flags: config, | ||
| bandits, | ||
| bandit_hashing_algorithm, | ||
| } | ||
| } | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,12 +4,16 @@ use std::sync::Arc; | |
| use chrono::{DateTime, Utc}; | ||
| use serde::Serialize; | ||
|
|
||
| use cityhasher::CityHasher; | ||
| use md5::{Digest, Md5}; | ||
| use std::hash::Hasher; | ||
|
|
||
| use crate::bandits::{ | ||
| BanditCategoricalAttributeCoefficient, BanditModelData, BanditNumericAttributeCoefficient, | ||
| }; | ||
| use crate::configuration::BanditHashingAlgorithm; | ||
| use crate::error::EvaluationFailure; | ||
| use crate::events::{AssignmentEvent, BanditEvent}; | ||
| use crate::sharder::PreSaltedSharder; | ||
| use crate::ufc::{Assignment, AssignmentValue, VariationType}; | ||
| use crate::{Configuration, EvaluationError, Str}; | ||
| use crate::{ContextAttributes, SdkMetadata}; | ||
|
|
@@ -187,27 +191,29 @@ fn get_bandit_action_with_visitor<V: EvalBanditVisitor>( | |
| return result; | ||
| }; | ||
|
|
||
| let evaluation = | ||
| match bandit | ||
| .model_data | ||
| .evaluate(flag_key, subject_key, subject_attributes, actions.iter()) | ||
| { | ||
| Ok(evaluation) => evaluation, | ||
| Err(err) => { | ||
| // We've evaluated a flag but now bandit evaluation failed. (Likely to user supplying | ||
| // empty actions, or NaN attributes.) | ||
| // | ||
| // Abort evaluation and return default variant. | ||
| let result = BanditResult { | ||
| variation, | ||
| action: None, | ||
| assignment_event: assignment.event, | ||
| bandit_event: None, | ||
| }; | ||
| visitor.on_result(Err(err), &result); | ||
| return result; | ||
| } | ||
| }; | ||
| let evaluation = match bandit.model_data.evaluate( | ||
| flag_key, | ||
| subject_key, | ||
| subject_attributes, | ||
| actions.iter(), | ||
| configuration.bandit_hashing_algorithm, | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ^ new parameter (caused a bit of reformatting in this section) |
||
| ) { | ||
| Ok(evaluation) => evaluation, | ||
| Err(err) => { | ||
| // We've evaluated a flag but now bandit evaluation failed. (Likely to user supplying | ||
| // empty actions, or NaN attributes.) | ||
| // | ||
| // Abort evaluation and return default variant. | ||
| let result = BanditResult { | ||
| variation, | ||
| action: None, | ||
| assignment_event: assignment.event, | ||
| bandit_event: None, | ||
| }; | ||
| visitor.on_result(Err(err), &result); | ||
| return result; | ||
| } | ||
| }; | ||
|
|
||
| let action_attributes = &actions[&evaluation.action_key]; | ||
| let bandit_event = BanditEvent { | ||
|
|
@@ -236,6 +242,108 @@ fn get_bandit_action_with_visitor<V: EvalBanditVisitor>( | |
| return result; | ||
| } | ||
|
|
||
| /// Trait for hashing in bandit evaluation. | ||
| /// | ||
| /// This trait abstracts the hashing logic for bandit evaluation, allowing different | ||
| /// implementations (MD5, CityHash) to be used interchangeably. | ||
| trait BanditHasher: Clone { | ||
| /// Create a new hasher pre-initialized with flag_key + "-" + subject_key | ||
| fn new(flag_key: &str, subject_key: &str) -> Self; | ||
|
|
||
| /// Get the selection hash (0.0..1.0) for choosing action based on weights | ||
| fn selection_hash(&self) -> f64; | ||
|
|
||
| /// Compute hash for shuffling a specific action | ||
| fn action_shuffle_hash(&self, action_key: &str) -> u64; | ||
| } | ||
|
Comment on lines
+249
to
+258
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nice job distilling the universal hashing operations needed for bandits here |
||
|
|
||
| /// MD5-based bandit hasher (10k shards, compatible with existing SDKs) | ||
| #[derive(Clone)] | ||
| struct Md5BanditHasher { | ||
| selection_hash: f64, | ||
| shuffle_ctx: Md5, // flag_key + "-" + subject_key + "-" | ||
| } | ||
|
|
||
| impl BanditHasher for Md5BanditHasher { | ||
| fn new(flag_key: &str, subject_key: &str) -> Self { | ||
| const TOTAL_SHARDS: u32 = 10_000; | ||
| let mut base_ctx = Md5::new(); | ||
| base_ctx.update(flag_key.as_bytes()); | ||
| base_ctx.update(b"-"); | ||
| base_ctx.update(subject_key.as_bytes()); | ||
|
|
||
| // Compute selection hash once | ||
| let selection_hash = { | ||
| let hash = base_ctx.clone().finalize(); | ||
| let value = u32::from_be_bytes(hash[0..4].try_into().unwrap()); | ||
| (value % TOTAL_SHARDS) as f64 / TOTAL_SHARDS as f64 | ||
| }; | ||
|
|
||
| // Prepare context for shuffling | ||
| let mut shuffle_ctx = base_ctx; | ||
| shuffle_ctx.update(b"-"); | ||
|
|
||
| Md5BanditHasher { | ||
| selection_hash, | ||
| shuffle_ctx, | ||
| } | ||
| } | ||
|
|
||
| fn selection_hash(&self) -> f64 { | ||
| self.selection_hash | ||
| } | ||
|
|
||
| fn action_shuffle_hash(&self, action_key: &str) -> u64 { | ||
| const TOTAL_SHARDS: u32 = 10_000; | ||
| let mut ctx = self.shuffle_ctx.clone(); | ||
| ctx.update(action_key.as_bytes()); | ||
| let hash = ctx.finalize(); | ||
| let value = u32::from_be_bytes(hash[0..4].try_into().unwrap()); | ||
| (value % TOTAL_SHARDS) as u64 | ||
| } | ||
| } | ||
|
|
||
| /// CityHash-based bandit hasher (experimental, better performance) | ||
| #[derive(Clone)] | ||
| struct CityHashBanditHasher { | ||
| selection_hash: f64, | ||
| shuffle_ctx: CityHasher, | ||
| } | ||
|
|
||
| impl BanditHasher for CityHashBanditHasher { | ||
| fn new(flag_key: &str, subject_key: &str) -> Self { | ||
| let mut base_ctx = CityHasher::new(); | ||
| base_ctx.write(flag_key.as_bytes()); | ||
| base_ctx.write(b"-"); | ||
| base_ctx.write(subject_key.as_bytes()); | ||
|
|
||
| // Compute selection hash once | ||
| let selection_hash = { | ||
| let hash = base_ctx.clone().finish(); | ||
| hash as u32 as f64 / u32::MAX as f64 | ||
|
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. "sharding" by 10k was arbitrary and somewhat useless, so using truncation to 32-bits in the new algorithm (which should be a tad faster and less biased) |
||
| }; | ||
|
|
||
| // Prepare context for shuffling | ||
| let mut shuffle_ctx = base_ctx; | ||
| shuffle_ctx.write(b"-"); | ||
|
|
||
| CityHashBanditHasher { | ||
| selection_hash, | ||
| shuffle_ctx, | ||
| } | ||
| } | ||
|
|
||
| fn selection_hash(&self) -> f64 { | ||
| self.selection_hash | ||
| } | ||
|
|
||
| fn action_shuffle_hash(&self, action_key: &str) -> u64 { | ||
| let mut ctx = self.shuffle_ctx.clone(); | ||
| ctx.write(action_key.as_bytes()); | ||
| ctx.finish() as u32 as u64 | ||
| } | ||
| } | ||
|
|
||
| impl BanditModelData { | ||
| // Exported to super, so we can use it in precomputed evaluation. | ||
| pub(super) fn evaluate<'a>( | ||
|
|
@@ -244,11 +352,32 @@ impl BanditModelData { | |
| subject_key: &str, | ||
| subject_attributes: &ContextAttributes, | ||
| actions: impl Iterator<Item = (&'a Str, &'a ContextAttributes)>, | ||
| hashing_algorithm: BanditHashingAlgorithm, | ||
| ) -> Result<BanditEvaluationDetails, EvaluationFailure> { | ||
| // total_shards is not configurable at the moment. | ||
| const TOTAL_SHARDS: u32 = 10_000; | ||
| match hashing_algorithm { | ||
| BanditHashingAlgorithm::Md5 => self.evaluate_with_hasher::<Md5BanditHasher>( | ||
| flag_key, | ||
| subject_key, | ||
| subject_attributes, | ||
| actions, | ||
| ), | ||
| BanditHashingAlgorithm::CityHash => self.evaluate_with_hasher::<CityHashBanditHasher>( | ||
| flag_key, | ||
| subject_key, | ||
| subject_attributes, | ||
| actions, | ||
| ), | ||
| } | ||
| } | ||
|
|
||
| let sharder = PreSaltedSharder::new(&[flag_key, "-", subject_key], TOTAL_SHARDS); | ||
| fn evaluate_with_hasher<'a, H: BanditHasher>( | ||
| &self, | ||
| flag_key: &str, | ||
| subject_key: &str, | ||
| subject_attributes: &ContextAttributes, | ||
| actions: impl Iterator<Item = (&'a Str, &'a ContextAttributes)>, | ||
| ) -> Result<BanditEvaluationDetails, EvaluationFailure> { | ||
| let hasher = H::new(flag_key, subject_key); | ||
|
|
||
| // Pseudo-random deterministic shuffle of actions. Shuffling is unique per subject, so when | ||
| // weights change slightly, large swatches of subjects are not reassigned from one action to | ||
|
|
@@ -260,7 +389,7 @@ impl BanditModelData { | |
| .collect::<Vec<_>>(); | ||
| // Sort actions by their shard value. Use action key as tie breaker. | ||
| shuffled_actions.sort_by_cached_key(|action| { | ||
| let hash = sharder.shard(&["-", action.key]); | ||
| let hash = hasher.action_shuffle_hash(action.key); | ||
| (hash, action.key) | ||
| }); | ||
| shuffled_actions | ||
|
|
@@ -298,7 +427,7 @@ impl BanditModelData { | |
| let weights = self.weigh_actions(&scores, best); | ||
| debug_assert_eq!(shuffled_actions.len(), weights.len()); | ||
|
|
||
| let selection_hash = sharder.shard(&[] as &[&str; 0]) as f64 / TOTAL_SHARDS as f64; | ||
| let selection_hash = hasher.selection_hash(); | ||
|
|
||
| let selected_action = { | ||
| let mut cumulative_weight = 0.0; | ||
|
|
@@ -315,7 +444,6 @@ impl BanditModelData { | |
|
|
||
| Ok(BanditEvaluationDetails { | ||
| action_key: shuffled_actions[selected_action].key.to_owned(), | ||
| // action_attributes: actions[selected_action].to_owned(), | ||
| action_weight: weights[selected_action], | ||
| optimality_gap, | ||
| }) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Don't want to read env variable on the hot evaluation path, so reading it here (which usually happens in background thread)