diff --git a/README.md b/README.md index ca947bf9..c70eba36 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,9 @@ This repository contains the core code that determines which posts a viewer sees ## Table of Contents -- [Latest update](#latest-update--august-13th-2026) +- [Notable Updates](#notable-updates) + - [August 14th, 2026](#august-14th-2026) + - [August 13th, 2026](#august-13th-2026) - [Overview](#overview) - [System Architecture](#system-architecture) - [Request Path](#request-path) @@ -23,7 +25,16 @@ This repository contains the core code that determines which posts a viewer sees -## Latest update — August 13th, 2026 +## Notable Updates + +### August 14th, 2026 + +Notable updates: + +- **How weights work.** There's a common misconception about how weights related to actions (e.g. Like, Share, Block, Report, etc) work in ranking. The weights scale the predicted probabilities of such actions (or predicted continuous values, e.g. dwell time) — they do *not* scale the raw engagement counts, so e.g. it'd be incorrect to see that a report has 468 times higher weight than a like and conclude that e.g. "1 report cancels out 468 likes". The weights are a multiple on your own predicted probability of Liking, Reporting, etc, which is substantially driven by your own behavior. We've [added comments](home-mixer/params/param.rs) [to the code](home-mixer/scorers/ranking_scorer.rs) so that LLMs or people reading it are more likely to understand it correctly. +- **Brazil 2026 Elections.** As [announced by X](https://x.com/XBR/status/2088341967864320507?s=20), in accordance with Brazilian electoral law, For You now runs `Brazil2026ElectionFilter`, which removes posts from accounts reported to Brazil's Electoral Court for the 2026 election, unless the viewer explicitly follows the account. *(Account list updated August 27, 2026.)* A benefit of open-source is that you can see that changes like this exist, and exactly how they work — take a [look at the code](home-mixer/filters/brazil_2026_election_filter.rs). + +### August 13th, 2026 This release: @@ -329,6 +340,8 @@ Final Score = Σ (weight_i × P(action_i)) Positive actions carry positive weights, negative actions negative ones. The weights are in [`home-mixer/params/param.rs`](home-mixer/params/param.rs); the arithmetic is in [`home-mixer/scorers/ranking_scorer.rs`](home-mixer/scorers/ranking_scorer.rs). +There is a common misconception to be aware of about the weights: they scale the predicted probabilities (or predicted continuous values, e.g. dwell time) — they do *not* scale the raw engagement counts, so e.g. it'd be incorrect to see that a report has 468 times higher weight than a like and conclude that e.g. "1 report cancels out 468 likes". The weights are a multiple on your own predicted probability of Liking, Reporting, etc, which is substantially driven by your own behavior. + Three adjustments follow: - **Author Diversity**: each post after an author's first is multiplied by a decaying factor, down to a floor. diff --git a/abuse-enforcement-service/service-lib/rules/enforcement_post.yaml b/abuse-enforcement-service/service-lib/rules/enforcement_post.yaml index bf4de586..f16fc220 100644 --- a/abuse-enforcement-service/service-lib/rules/enforcement_post.yaml +++ b/abuse-enforcement-service/service-lib/rules/enforcement_post.yaml @@ -1,4 +1,4 @@ -# mirrored from GrowthBook dynamic config; last sync 2026-08-06T16:20:00Z +# mirrored from GrowthBook dynamic config; last sync 2026-08-25T16:15:48Z for_entity: post @@ -75,6 +75,11 @@ rules: perm: true policy: Cse + - id: act_requested_actions + when: "size(score.requested_actions) > 0" + then: + kind: act_requested_actions + - id: post_no_actionable_label when: "true" then: diff --git a/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml b/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml index 96327641..b58d5a31 100644 --- a/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml +++ b/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml @@ -1,4 +1,4 @@ -# mirrored from GrowthBook dynamic config; last sync 2026-08-12T16:22:21Z +# mirrored from GrowthBook dynamic config; last sync 2026-08-26T16:31:18Z for_entity: user @@ -15,6 +15,13 @@ rules: kind: skip reason: user_not_found + - id: very_high_follower_count + # Prod uses a different follower count floor; this is a mock value to reduce gaming. + when: cred.follower_count >= 12.34 + then: + kind: skip + reason: very_high_follower_count + - id: high_follower_count # Prod uses a different follower count floor; this is a mock value to reduce gaming. when: cred.follower_count >= 12.34 @@ -42,6 +49,12 @@ rules: when: '"anchor_campaign_suspend_cse" in score.labels' then: { kind: act_suspend_user, perm: true, policy: "Cse" } + - id: panda_reports_embedding_v10_rough_spam + when: '"panda_reports_embedding_v10_rough_spam" in score.labels' + then: { kind: act_suspend_user, perm: false, policy: "PlatformManipulation" } + + + - id: already_spam_high_recall_labeled_llm_slop when: '"llm_slop_user" in score.labels && "SpamHighRecall" in user.labels' then: @@ -201,6 +214,17 @@ rules: labels: ["SpamHighRecall"] ttl_msec: 2592000000 + - id: platform_row_without_requested_actions + when: '"abuse_platform_requested_actions" in score.labels && size(score.requested_actions) == 0' + then: + kind: skip + reason: platform_row_without_requested_actions + + - id: act_requested_actions + when: "size(score.requested_actions) > 0" + then: + kind: act_requested_actions + - id: act_suspend when: "true" then: diff --git a/abuse-enforcement-service/service-lib/src/ais_client.rs b/abuse-enforcement-service/service-lib/src/ais_client.rs index 63d84295..2bc79047 100644 --- a/abuse-enforcement-service/service-lib/src/ais_client.rs +++ b/abuse-enforcement-service/service-lib/src/ais_client.rs @@ -471,8 +471,10 @@ mod tests { "additionalInfo": [] }}), json!({"bounceViaSelection": { - "userId": 1, - "uncheckedTags": ["FAKE"] + "target": {"user": {"userId": 1}}, + "tags": [], + "uncheckedTags": ["FAKE"], + "bounceActor": {"simpleService": {"serviceName": "xai-abuse-enforcement-service"}} }}), ], "RtpPlugin", diff --git a/abuse-enforcement-service/service-lib/src/config.rs b/abuse-enforcement-service/service-lib/src/config.rs index 2825f2fc..ac1d2714 100644 --- a/abuse-enforcement-service/service-lib/src/config.rs +++ b/abuse-enforcement-service/service-lib/src/config.rs @@ -38,25 +38,14 @@ pub struct Config { #[arg(long, env = "SOCKS_PROXY")] pub socks_proxy: Option, - #[arg( - long, - default_value = "/s/kafka/phoenix-kafka-scram-bootstrap", - env = "KAFKA_DEST" - )] - pub kafka_dest: String, + #[arg(long, default_value_t = false, env = "KAFKA_CONSUMER_ENABLED")] + pub kafka_consumer_enabled: bool, - #[arg(long, default_value = "SCRAM-SHA-512", env = "KAFKA_SASL_MECHANISM")] - pub kafka_sasl_mechanism: String, + #[arg(long, default_value = "phoenix", env = "KAFKA_CONSUMER_MTLS_CLUSTER")] + pub kafka_consumer_mtls_cluster: String, - #[arg( - long, - default_value = "scram-client-phoenix", - env = "KAFKA_SASL_USERNAME" - )] - pub kafka_sasl_username: String, - - #[arg(long, env = "KAFKA_SASL_PASSWORD")] - pub kafka_sasl_password: Option, + #[arg(long, default_value = "atla", env = "KAFKA_CONSUMER_MTLS_ZONE")] + pub kafka_consumer_mtls_zone: String, #[arg( long, @@ -68,39 +57,12 @@ pub struct Config { #[arg(long, env = "KAFKA_CONSUMER_GROUP_ID")] pub kafka_consumer_group_id: Option, - #[arg(long, env = "KAFKA_CONSUMER_DEST")] - pub kafka_consumer_dest: Option, - - #[arg(long, env = "KAFKA_CONSUMER_SASL_MECHANISM")] - pub kafka_consumer_sasl_mechanism: Option, - - #[arg(long, env = "KAFKA_CONSUMER_SASL_USERNAME")] - pub kafka_consumer_sasl_username: Option, - - #[arg(long, env = "KAFKA_CONSUMER_SASL_PASSWORD")] - pub kafka_consumer_sasl_password: Option, - #[arg(long, default_value_t = 256, env = "KAFKA_MAX_MESSAGES_PER_POLL")] pub kafka_max_messages_per_poll: usize, #[arg(long, default_value_t = 64, env = "KAFKA_MAX_IN_FLIGHT")] pub kafka_max_in_flight: usize, - #[arg(long, env = "KAFKA_PRODUCER_DEST")] - pub kafka_producer_dest: Option, - - #[arg(long, env = "KAFKA_PRODUCER_SASL_MECHANISM")] - pub kafka_producer_sasl_mechanism: Option, - - #[arg(long, env = "KAFKA_PRODUCER_SASL_USERNAME")] - pub kafka_producer_sasl_username: Option, - - #[arg(long, env = "KAFKA_PRODUCER_SASL_PASSWORD")] - pub kafka_producer_sasl_password: Option, - - #[arg(long, default_value_t = false, env = "KAFKA_PRODUCER_MTLS_ENABLED")] - pub kafka_producer_mtls_enabled: bool, - #[arg(long, default_value = "coredata", env = "KAFKA_PRODUCER_MTLS_CLUSTER")] pub kafka_producer_mtls_cluster: String, @@ -257,61 +219,7 @@ pub struct Config { pub api_keys_json: Option, } -#[derive(Debug, Clone)] -pub struct KafkaConnConfig { - pub dest: String, - pub sasl_mechanism: String, - pub sasl_username: String, - pub sasl_password: Option, -} - -#[derive(Debug, Default)] -struct KafkaConnOverride { - dest: Option, - sasl_mechanism: Option, - sasl_username: Option, - sasl_password: Option, -} - -impl KafkaConnConfig { - fn with_overrides(self, overrides: KafkaConnOverride) -> KafkaConnConfig { - KafkaConnConfig { - dest: overrides.dest.unwrap_or(self.dest), - sasl_mechanism: overrides.sasl_mechanism.unwrap_or(self.sasl_mechanism), - sasl_username: overrides.sasl_username.unwrap_or(self.sasl_username), - sasl_password: overrides.sasl_password.or(self.sasl_password), - } - } -} - impl Config { - fn kafka_common(&self) -> KafkaConnConfig { - KafkaConnConfig { - dest: self.kafka_dest.clone(), - sasl_mechanism: self.kafka_sasl_mechanism.clone(), - sasl_username: self.kafka_sasl_username.clone(), - sasl_password: self.kafka_sasl_password.clone(), - } - } - - pub fn kafka_consumer(&self) -> KafkaConnConfig { - self.kafka_common().with_overrides(KafkaConnOverride { - dest: self.kafka_consumer_dest.clone(), - sasl_mechanism: self.kafka_consumer_sasl_mechanism.clone(), - sasl_username: self.kafka_consumer_sasl_username.clone(), - sasl_password: self.kafka_consumer_sasl_password.clone(), - }) - } - - pub fn kafka_producer(&self) -> KafkaConnConfig { - self.kafka_common().with_overrides(KafkaConnOverride { - dest: self.kafka_producer_dest.clone(), - sasl_mechanism: self.kafka_producer_sasl_mechanism.clone(), - sasl_username: self.kafka_producer_sasl_username.clone(), - sasl_password: self.kafka_producer_sasl_password.clone(), - }) - } - pub fn kafka_group_id(&self) -> String { self.kafka_consumer_group_id .clone() @@ -326,10 +234,27 @@ mod tests { use super::Config; #[test] - fn coredata_producer_mtls_defaults_pin_atla() { + fn kafka_defaults_preserve_consumer_and_watchdog_behavior() { + let config = Config::parse_from(["xai-abuse-enforcement-service"]); + + assert!(!config.kafka_consumer_enabled); + assert_eq!(config.kafka_consumer_mtls_cluster, "phoenix"); + assert_eq!(config.kafka_consumer_mtls_zone, "atla"); + assert_eq!(config.kafka_group_id(), "xai-abuse-enforcement-service"); + assert_eq!(config.kafka_max_messages_per_poll, 256); + assert_eq!(config.kafka_max_in_flight, 64); + assert!(config.kafka_self_delete_enabled); + assert_eq!(config.kafka_watchdog_interval_secs, 15); + assert_eq!(config.kafka_watchdog_stale_secs, 120); + assert_eq!(config.kafka_watchdog_self_delete_secs, 240); + assert_eq!(config.kafka_watchdog_error_rate_per_sec, 2.0); + assert_eq!(config.kafka_watchdog_error_self_delete_secs, 60); + } + + #[test] + fn coredata_producer_defaults_use_mtls() { let config = Config::parse_from(["xai-abuse-enforcement-service"]); - assert!(!config.kafka_producer_mtls_enabled); assert_eq!(config.kafka_producer_mtls_cluster, "coredata"); assert_eq!(config.kafka_producer_mtls_zone, "atla"); } diff --git a/abuse-enforcement-service/service-lib/src/decision.rs b/abuse-enforcement-service/service-lib/src/decision.rs index a64a3c84..f2473f6d 100644 --- a/abuse-enforcement-service/service-lib/src/decision.rs +++ b/abuse-enforcement-service/service-lib/src/decision.rs @@ -2,6 +2,7 @@ pub enum Decision { Skip(String), Act(Vec), + ActRequestedActions, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/abuse-enforcement-service/service-lib/src/facts.rs b/abuse-enforcement-service/service-lib/src/facts.rs index f06de079..6b6c25cc 100644 --- a/abuse-enforcement-service/service-lib/src/facts.rs +++ b/abuse-enforcement-service/service-lib/src/facts.rs @@ -139,6 +139,29 @@ impl Facts { } } +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RequestedActionFacts { + pub kind: String, + pub perm: bool, + pub policy: String, + pub labels: Vec, + pub ttl_msec: i64, + pub head: String, +} + +impl RequestedActionFacts { + fn from_proto(a: &xai_abuse_proto::enforcement::RequestedAction) -> Self { + Self { + kind: a.kind.clone(), + perm: a.perm, + policy: a.policy.clone(), + labels: a.labels.clone(), + ttl_msec: a.ttl_msec, + head: a.head.clone(), + } + } +} + #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ScoreFacts { pub model_version: String, @@ -147,6 +170,10 @@ pub struct ScoreFacts { pub labels: Vec, #[serde(default)] pub skip_author_credibility_prechecks: bool, + #[serde(default)] + pub requested_actions: Vec, + #[serde(default)] + pub policy_version: String, } impl ScoreFacts { @@ -170,6 +197,12 @@ impl ScoreFacts { decoded_actions: take_string("decoded_actions"), labels, skip_author_credibility_prechecks: score.skip_author_credibility_prechecks, + requested_actions: score + .requested_actions + .iter() + .map(RequestedActionFacts::from_proto) + .collect(), + policy_version: score.policy_version.clone(), } } } @@ -497,6 +530,73 @@ mod tests { assert_eq!(f.enforcement_note, "test note"); } + #[test] + fn score_facts_projects_requested_actions_and_policy_version() { + use xai_abuse_proto::enforcement::RequestedAction; + let s = ScoreResult { + requested_actions: vec![ + RequestedAction { + kind: "suspend".into(), + perm: true, + policy: "PlatformManipulation".into(), + labels: vec![], + ttl_msec: 0, + head: "IsSpammer".into(), + }, + RequestedAction { + kind: "label".into(), + perm: false, + policy: String::new(), + labels: vec!["SpamHighRecall".into()], + ttl_msec: 86_400_000, + head: "IsLabelHead".into(), + }, + ], + policy_version: "7".into(), + ..Default::default() + }; + let f = ScoreFacts::from_score(&s); + assert_eq!(f.policy_version, "7"); + assert_eq!(f.requested_actions.len(), 2); + assert_eq!( + f.requested_actions[0], + RequestedActionFacts { + kind: "suspend".into(), + perm: true, + policy: "PlatformManipulation".into(), + labels: vec![], + ttl_msec: 0, + head: "IsSpammer".into(), + } + ); + assert_eq!(f.requested_actions[1].kind, "label"); + assert_eq!( + f.requested_actions[1].labels, + vec!["SpamHighRecall".to_owned()] + ); + assert_eq!(f.requested_actions[1].ttl_msec, 86_400_000); + } + + #[test] + fn score_facts_without_requested_actions_projects_empty() { + let f = ScoreFacts::from_score(&ScoreResult::default()); + assert!(f.requested_actions.is_empty()); + assert_eq!(f.policy_version, ""); + } + + #[test] + fn score_facts_pre_requested_actions_json_still_deserializes() { + let legacy = r#"{ + "model_version": "v1", + "enforcement_note": "", + "decoded_actions": "[]", + "labels": [] + }"#; + let f: ScoreFacts = serde_json::from_str(legacy).expect("legacy ScoreFacts must parse"); + assert!(f.requested_actions.is_empty()); + assert_eq!(f.policy_version, ""); + } + #[test] fn score_facts_labels_empty_when_summary_absent() { let s = ScoreResult { diff --git a/abuse-enforcement-service/service-lib/src/fixtures/mock_rules_user.yaml b/abuse-enforcement-service/service-lib/src/fixtures/mock_rules_user.yaml index 1d6d8fa4..cea4a20a 100644 --- a/abuse-enforcement-service/service-lib/src/fixtures/mock_rules_user.yaml +++ b/abuse-enforcement-service/service-lib/src/fixtures/mock_rules_user.yaml @@ -7,7 +7,7 @@ rules: when: "!user.present" then: { kind: skip, reason: not_present } - id: cred_skip - when: cred.is_high || cred.follower_count >= 1234 || cred.score >= 7.5 + when: cred.is_high || cred.follower_count >= 12.34 || cred.score >= 7.5 then: { kind: skip, reason: cred_skip } - id: composite when: '"composite_label" in score.labels' diff --git a/abuse-enforcement-service/service-lib/src/generic_actions.rs b/abuse-enforcement-service/service-lib/src/generic_actions.rs new file mode 100644 index 00000000..44c770d9 --- /dev/null +++ b/abuse-enforcement-service/service-lib/src/generic_actions.rs @@ -0,0 +1,663 @@ +use std::collections::HashSet; + +use serde::Serialize; +use tracing::warn; + +use crate::decision::ActionSpec; +use crate::facts::{EntityType, RequestedActionFacts}; + +pub const USER_KINDS: &[&str] = &[ + "suspend", + "label", + "bounce_captcha", + "bounce_arkose", + "spam_liveness_check", +]; + +pub const POST_KINDS: &[&str] = &["post_label", "suspend_author"]; + +pub const MAX_REQUESTED_ACTIONS_PER_MESSAGE: usize = 16; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct GenericActionAllowlist { + pub kinds: HashSet, + pub suspend_policies: HashSet, + pub labels: HashSet, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct SkippedRequestedAction { + pub kind: String, + pub head: String, + pub reason: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub count: Option, +} + +#[derive(Debug, Default)] +pub struct ResolvedRequestedActions { + pub specs: Vec, + pub skipped: Vec, +} + +impl ResolvedRequestedActions { + pub fn skipped_info_json(&self) -> Option { + if self.skipped.is_empty() { + return None; + } + serde_json::to_string(&self.skipped).ok() + } +} + +fn hardcoded_kinds(entity_type: EntityType) -> &'static [&'static str] { + match entity_type { + EntityType::User => USER_KINDS, + EntityType::Post => POST_KINDS, + } +} + +fn metric_kind(kind: &str) -> &str { + if USER_KINDS.contains(&kind) || POST_KINDS.contains(&kind) { + kind + } else { + "unknown" + } +} + +fn sanitize_for_log(s: &str) -> String { + const MAX: usize = 64; + let mut out: String = s + .chars() + .take(MAX) + .map(|c| if c.is_control() { '?' } else { c }) + .collect(); + if s.chars().count() > MAX { + out.push('…'); + } + out +} + +fn action_to_spec( + entity_type: EntityType, + action: &RequestedActionFacts, + allowlist: &GenericActionAllowlist, +) -> Result { + let kind = action.kind.as_str(); + if !hardcoded_kinds(entity_type).contains(&kind) { + let other = match entity_type { + EntityType::User => EntityType::Post, + EntityType::Post => EntityType::User, + }; + return Err(if hardcoded_kinds(other).contains(&kind) { + "entity_type_mismatch" + } else { + "unknown_kind" + }); + } + if !allowlist.kinds.contains(kind) { + return Err("kind_not_allowlisted"); + } + match kind { + "suspend" | "suspend_author" => { + if !allowlist.suspend_policies.contains(&action.policy) { + return Err("policy_not_allowlisted"); + } + Ok(ActionSpec::SuspendUser { + perm: action.perm, + policy: action.policy.clone(), + }) + } + "label" | "post_label" => { + if action.labels.is_empty() { + return Err("no_labels"); + } + if action.labels.iter().any(|l| !allowlist.labels.contains(l)) { + return Err("label_not_allowlisted"); + } + if action.ttl_msec < 0 { + return Err("invalid_ttl"); + } + let ttl_msec = (action.ttl_msec > 0).then_some(action.ttl_msec); + Ok(if kind == "label" { + ActionSpec::AddLabelsV2 { + labels: action.labels.clone(), + ttl_msec, + } + } else { + ActionSpec::AddPostLabelsV2 { + labels: action.labels.clone(), + ttl_msec, + } + }) + } + "bounce_captcha" => Ok(ActionSpec::Captcha), + "bounce_arkose" => Ok(ActionSpec::Arkose), + "spam_liveness_check" => Ok(ActionSpec::SpamLivenessCheck), + _ => Err("unknown_kind"), + } +} + +pub fn resolve_requested_actions( + entity_type: EntityType, + requested: &[RequestedActionFacts], + allowlist: &GenericActionAllowlist, +) -> ResolvedRequestedActions { + let mut out = ResolvedRequestedActions::default(); + let considered = &requested[..requested.len().min(MAX_REQUESTED_ACTIONS_PER_MESSAGE)]; + let overflow = requested.len() - considered.len(); + for action in considered { + let resolved = action_to_spec(entity_type, action, allowlist).and_then(|spec| { + if out.specs.contains(&spec) { + Err("duplicate_action") + } else { + Ok(spec) + } + }); + match resolved { + Ok(spec) => { + crate::metrics::GENERIC_ACTION_TOTAL + .with_label_values(&[ + entity_type.as_str(), + metric_kind(&action.kind), + "resolved", + "", + ]) + .inc(); + out.specs.push(spec); + } + Err(reason) => { + warn!( + entity_type = entity_type.as_str(), + kind = %sanitize_for_log(&action.kind), + kind_len = action.kind.len(), + head = %sanitize_for_log(&action.head), + head_len = action.head.len(), + reason, + "requested action skipped (fail-closed)" + ); + crate::metrics::GENERIC_ACTION_TOTAL + .with_label_values(&[ + entity_type.as_str(), + metric_kind(&action.kind), + "skipped", + reason, + ]) + .inc(); + out.skipped.push(SkippedRequestedAction { + kind: action.kind.clone(), + head: action.head.clone(), + reason, + count: None, + }); + } + } + } + if overflow > 0 { + warn!( + entity_type = entity_type.as_str(), + overflow, + cap = MAX_REQUESTED_ACTIONS_PER_MESSAGE, + reason = "too_many_actions", + "requested actions past the per-message cap skipped (fail-closed, aggregated)" + ); + crate::metrics::GENERIC_ACTION_TOTAL + .with_label_values(&[ + entity_type.as_str(), + "aggregate", + "skipped", + "too_many_actions", + ]) + .inc(); + out.skipped.push(SkippedRequestedAction { + kind: String::new(), + head: String::new(), + reason: "too_many_actions", + count: Some(overflow), + }); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn full_user_allowlist() -> GenericActionAllowlist { + GenericActionAllowlist { + kinds: USER_KINDS.iter().map(|s| (*s).to_owned()).collect(), + suspend_policies: ["PlatformManipulation"] + .iter() + .map(|s| (*s).to_owned()) + .collect(), + labels: ["SpamHighRecall"].iter().map(|s| (*s).to_owned()).collect(), + } + } + + fn full_post_allowlist() -> GenericActionAllowlist { + GenericActionAllowlist { + kinds: POST_KINDS.iter().map(|s| (*s).to_owned()).collect(), + suspend_policies: ["Cse"].iter().map(|s| (*s).to_owned()).collect(), + labels: ["SpamHighRecall"].iter().map(|s| (*s).to_owned()).collect(), + } + } + + fn req(kind: &str) -> RequestedActionFacts { + RequestedActionFacts { + kind: kind.into(), + head: "SomeHead".into(), + ..Default::default() + } + } + + fn suspend_req(kind: &str, policy: &str, perm: bool) -> RequestedActionFacts { + RequestedActionFacts { + kind: kind.into(), + perm, + policy: policy.into(), + head: "SomeHead".into(), + ..Default::default() + } + } + + fn label_req(kind: &str, labels: &[&str], ttl_msec: i64) -> RequestedActionFacts { + RequestedActionFacts { + kind: kind.into(), + labels: labels.iter().map(|s| (*s).to_owned()).collect(), + ttl_msec, + head: "SomeHead".into(), + ..Default::default() + } + } + + #[test] + fn user_suspend_maps_to_suspend_user_spec() { + let r = resolve_requested_actions( + EntityType::User, + &[suspend_req("suspend", "PlatformManipulation", true)], + &full_user_allowlist(), + ); + assert!(r.skipped.is_empty()); + assert_eq!( + r.specs, + vec![ActionSpec::SuspendUser { + perm: true, + policy: "PlatformManipulation".into(), + }] + ); + } + + #[test] + fn user_label_maps_to_add_labels_v2_with_ttl() { + let r = resolve_requested_actions( + EntityType::User, + &[label_req("label", &["SpamHighRecall"], 86_400_000)], + &full_user_allowlist(), + ); + assert!(r.skipped.is_empty()); + assert_eq!( + r.specs, + vec![ActionSpec::AddLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: Some(86_400_000), + }] + ); + } + + #[test] + fn label_ttl_zero_maps_to_no_expiry() { + let r = resolve_requested_actions( + EntityType::User, + &[label_req("label", &["SpamHighRecall"], 0)], + &full_user_allowlist(), + ); + assert_eq!( + r.specs, + vec![ActionSpec::AddLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: None, + }] + ); + } + + #[test] + fn user_challenge_kinds_map_to_their_specs() { + let r = resolve_requested_actions( + EntityType::User, + &[ + req("bounce_captcha"), + req("bounce_arkose"), + req("spam_liveness_check"), + ], + &full_user_allowlist(), + ); + assert!(r.skipped.is_empty()); + assert_eq!( + r.specs, + vec![ + ActionSpec::Captcha, + ActionSpec::Arkose, + ActionSpec::SpamLivenessCheck, + ] + ); + } + + #[test] + fn post_label_maps_to_add_post_labels_v2() { + let r = resolve_requested_actions( + EntityType::Post, + &[label_req("post_label", &["SpamHighRecall"], 86_400_000)], + &full_post_allowlist(), + ); + assert!(r.skipped.is_empty()); + assert_eq!( + r.specs, + vec![ActionSpec::AddPostLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: Some(86_400_000), + }] + ); + } + + #[test] + fn post_suspend_author_maps_to_suspend_user_spec() { + let r = resolve_requested_actions( + EntityType::Post, + &[suspend_req("suspend_author", "Cse", true)], + &full_post_allowlist(), + ); + assert!(r.skipped.is_empty()); + assert_eq!( + r.specs, + vec![ActionSpec::SuspendUser { + perm: true, + policy: "Cse".into(), + }] + ); + } + + #[test] + fn post_kinds_on_user_entity_are_refused_even_when_allowlisted() { + let mut allowlist = full_user_allowlist(); + allowlist.kinds.insert("post_label".into()); + allowlist.kinds.insert("suspend_author".into()); + for kind in POST_KINDS { + let r = resolve_requested_actions( + EntityType::User, + &[suspend_req(kind, "PlatformManipulation", false)], + &allowlist, + ); + assert!(r.specs.is_empty(), "{kind} must not dispatch on user"); + assert_eq!(r.skipped[0].reason, "entity_type_mismatch"); + } + } + + #[test] + fn user_kinds_on_post_entity_are_refused_even_when_allowlisted() { + let mut allowlist = full_post_allowlist(); + for kind in USER_KINDS { + allowlist.kinds.insert((*kind).to_owned()); + } + allowlist + .suspend_policies + .insert("PlatformManipulation".into()); + for kind in USER_KINDS { + let r = resolve_requested_actions( + EntityType::Post, + &[RequestedActionFacts { + kind: (*kind).to_owned(), + policy: "PlatformManipulation".into(), + labels: vec!["SpamHighRecall".into()], + head: "SomeHead".into(), + ..Default::default() + }], + &allowlist, + ); + assert!(r.specs.is_empty(), "{kind} must not dispatch on post"); + assert_eq!(r.skipped[0].reason, "entity_type_mismatch"); + } + } + + #[test] + fn unknown_kind_is_refused() { + for entity in [EntityType::User, EntityType::Post] { + let allowlist = match entity { + EntityType::User => full_user_allowlist(), + EntityType::Post => full_post_allowlist(), + }; + let r = resolve_requested_actions(entity, &[req("bounce")], &allowlist); + assert!(r.specs.is_empty()); + assert_eq!(r.skipped[0].reason, "unknown_kind"); + } + } + + #[test] + fn empty_allowlist_refuses_everything() { + let empty = GenericActionAllowlist::default(); + let r = resolve_requested_actions( + EntityType::User, + &[ + suspend_req("suspend", "PlatformManipulation", false), + label_req("label", &["SpamHighRecall"], 0), + req("bounce_captcha"), + req("bounce_arkose"), + req("spam_liveness_check"), + ], + &empty, + ); + assert!(r.specs.is_empty()); + assert_eq!(r.skipped.len(), 5); + assert!(r.skipped.iter().all(|s| s.reason == "kind_not_allowlisted")); + + let r = resolve_requested_actions( + EntityType::Post, + &[ + label_req("post_label", &["SpamHighRecall"], 0), + suspend_req("suspend_author", "Cse", true), + ], + &empty, + ); + assert!(r.specs.is_empty()); + assert_eq!(r.skipped.len(), 2); + } + + #[test] + fn non_allowlisted_kind_is_refused() { + let mut allowlist = full_user_allowlist(); + allowlist.kinds.remove("suspend"); + let r = resolve_requested_actions( + EntityType::User, + &[suspend_req("suspend", "PlatformManipulation", false)], + &allowlist, + ); + assert!(r.specs.is_empty()); + assert_eq!(r.skipped[0].reason, "kind_not_allowlisted"); + } + + #[test] + fn non_allowlisted_suspend_policy_is_refused() { + for (entity, kind, allowlist) in [ + (EntityType::User, "suspend", full_user_allowlist()), + (EntityType::Post, "suspend_author", full_post_allowlist()), + ] { + let r = resolve_requested_actions( + entity, + &[suspend_req(kind, "SomeOtherPolicy", false)], + &allowlist, + ); + assert!(r.specs.is_empty(), "{kind} with foreign policy"); + assert_eq!(r.skipped[0].reason, "policy_not_allowlisted"); + + let r = resolve_requested_actions(entity, &[suspend_req(kind, "", false)], &allowlist); + assert_eq!(r.skipped[0].reason, "policy_not_allowlisted"); + } + } + + #[test] + fn non_allowlisted_label_is_refused() { + for (entity, kind, allowlist) in [ + (EntityType::User, "label", full_user_allowlist()), + (EntityType::Post, "post_label", full_post_allowlist()), + ] { + let r = resolve_requested_actions( + entity, + &[label_req(kind, &["SpamHighRecall", "SomethingElse"], 0)], + &allowlist, + ); + assert!(r.specs.is_empty(), "{kind} with foreign label"); + assert_eq!(r.skipped[0].reason, "label_not_allowlisted"); + } + } + + #[test] + fn label_kind_with_no_labels_is_refused() { + let r = resolve_requested_actions( + EntityType::User, + &[label_req("label", &[], 0)], + &full_user_allowlist(), + ); + assert!(r.specs.is_empty()); + assert_eq!(r.skipped[0].reason, "no_labels"); + } + + #[test] + fn negative_ttl_is_refused() { + for (entity, kind, allowlist) in [ + (EntityType::User, "label", full_user_allowlist()), + (EntityType::Post, "post_label", full_post_allowlist()), + ] { + let r = resolve_requested_actions( + entity, + &[label_req(kind, &["SpamHighRecall"], -1)], + &allowlist, + ); + assert!(r.specs.is_empty(), "{kind} with negative ttl"); + assert_eq!(r.skipped[0].reason, "invalid_ttl"); + } + } + + #[test] + fn entries_past_the_hardcoded_cap_are_refused_as_one_aggregate() { + let mut requested: Vec = (0..MAX_REQUESTED_ACTIONS_PER_MESSAGE) + .map(|i| label_req("label", &["SpamHighRecall"], (i as i64 + 1) * 1000)) + .collect(); + requested.push(req("bounce_captcha")); + requested.push(req("bounce_arkose")); + let r = resolve_requested_actions(EntityType::User, &requested, &full_user_allowlist()); + assert_eq!(r.specs.len(), MAX_REQUESTED_ACTIONS_PER_MESSAGE); + assert_eq!( + r.skipped, + vec![SkippedRequestedAction { + kind: String::new(), + head: String::new(), + reason: "too_many_actions", + count: Some(2), + }] + ); + let json = r.skipped_info_json().expect("aggregate serializes"); + assert!(json.contains("\"count\":2"), "{json}"); + } + + #[test] + fn per_entry_refusals_omit_the_count_field() { + let r = resolve_requested_actions( + EntityType::User, + &[req("post_label")], + &full_user_allowlist(), + ); + assert_eq!(r.skipped[0].count, None); + let json = r.skipped_info_json().expect("skipped list serializes"); + assert!(!json.contains("count"), "{json}"); + } + + #[test] + fn identical_resolved_specs_are_deduplicated() { + let r = resolve_requested_actions( + EntityType::Post, + &[ + label_req("post_label", &["SpamHighRecall"], 0), + suspend_req("suspend_author", "Cse", true), + label_req("post_label", &["SpamHighRecall"], 0), + suspend_req("suspend_author", "Cse", true), + ], + &full_post_allowlist(), + ); + assert_eq!( + r.specs, + vec![ + ActionSpec::AddPostLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: None, + }, + ActionSpec::SuspendUser { + perm: true, + policy: "Cse".into(), + }, + ] + ); + assert_eq!(r.skipped.len(), 2); + assert!(r.skipped.iter().all(|s| s.reason == "duplicate_action")); + } + + #[test] + fn mixed_list_dispatches_allowed_and_skips_refused_in_order() { + let r = resolve_requested_actions( + EntityType::User, + &[ + label_req("label", &["SpamHighRecall"], 1000), + req("post_label"), + suspend_req("suspend", "PlatformManipulation", false), + ], + &full_user_allowlist(), + ); + assert_eq!( + r.specs, + vec![ + ActionSpec::AddLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: Some(1000), + }, + ActionSpec::SuspendUser { + perm: false, + policy: "PlatformManipulation".into(), + }, + ] + ); + assert_eq!(r.skipped.len(), 1); + assert_eq!(r.skipped[0].kind, "post_label"); + + let json = r.skipped_info_json().expect("skipped list serializes"); + assert!(json.contains("\"post_label\"")); + assert!(json.contains("entity_type_mismatch")); + } + + #[test] + fn skipped_info_json_none_when_nothing_skipped() { + let r = resolve_requested_actions( + EntityType::User, + &[req("bounce_captcha")], + &full_user_allowlist(), + ); + assert!(r.skipped_info_json().is_none()); + } + + #[test] + fn metric_kind_collapses_unknown_kinds() { + assert_eq!(metric_kind("suspend"), "suspend"); + assert_eq!(metric_kind("post_label"), "post_label"); + assert_eq!(metric_kind("totally-made-up"), "unknown"); + assert_eq!(metric_kind(""), "unknown"); + } + + #[test] + fn sanitize_for_log_strips_control_chars_and_truncates() { + assert_eq!(sanitize_for_log("suspend"), "suspend"); + assert_eq!( + sanitize_for_log("evil\nline\x1b[31mred"), + "evil?line?[31mred" + ); + let long = "a".repeat(200); + let sanitized = sanitize_for_log(&long); + assert_eq!(sanitized.chars().count(), 65); + assert!(sanitized.ends_with('…')); + } +} diff --git a/abuse-enforcement-service/service-lib/src/growthbook.rs b/abuse-enforcement-service/service-lib/src/growthbook.rs index 8a7a9e25..f0dc9a88 100644 --- a/abuse-enforcement-service/service-lib/src/growthbook.rs +++ b/abuse-enforcement-service/service-lib/src/growthbook.rs @@ -32,6 +32,11 @@ const DEFAULT_DEDUP_HEAD_ACTION_CLASS: &[(&str, u8)] = &[ ), ]; +const GENERIC_ACTIONS_KEY: &str = "generic_actions"; +const GENERIC_KINDS_KEY: &str = "kinds"; +const GENERIC_SUSPEND_POLICIES_KEY: &str = "suspend_policies"; +const GENERIC_LABELS_KEY: &str = "labels"; + const KAFKA_KEY: &str = "kafka"; const CONSUMER_KEY: &str = "consumer"; const PRODUCER_KEY: &str = "producer"; @@ -187,6 +192,13 @@ impl DynamicConfig { dedup_action_class_from_config(self.config().as_ref(), head) } + pub fn generic_action_allowlist( + &self, + entity_type: EntityType, + ) -> crate::generic_actions::GenericActionAllowlist { + generic_actions_from_config(self.config().as_ref(), entity_type) + } + pub fn config(&self) -> Option { let c = self.client.as_ref()?; Some(c.feature_result(GROWTHBOOK_CONFIG_KEY, None).value) @@ -226,6 +238,30 @@ fn parse_bool(config: Option<&Value>, key: &str, default: bool) -> bool { .unwrap_or(default) } +fn string_set(v: Option<&Value>) -> std::collections::HashSet { + v.and_then(Value::as_array) + .map(|arr| { + arr.iter() + .filter_map(Value::as_str) + .filter(|s| !s.trim().is_empty()) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default() +} + +fn generic_actions_from_config( + config: Option<&Value>, + entity_type: EntityType, +) -> crate::generic_actions::GenericActionAllowlist { + let entity = get_path(config, &[GENERIC_ACTIONS_KEY, entity_type.as_str()]); + crate::generic_actions::GenericActionAllowlist { + kinds: string_set(entity.and_then(|e| e.get(GENERIC_KINDS_KEY))), + suspend_policies: string_set(entity.and_then(|e| e.get(GENERIC_SUSPEND_POLICIES_KEY))), + labels: string_set(entity.and_then(|e| e.get(GENERIC_LABELS_KEY))), + } +} + fn dedup_action_class_from_config(config: Option<&Value>, head: &str) -> u8 { use crate::dedup_cache::{parse_action_class, CLASS_NONE, CLASS_SUSPEND}; if head.is_empty() { @@ -417,6 +453,96 @@ mod tests { assert_eq!(pick_u32(Some(&cfg), &[MAX_POST_ENFORCEMENTS_KEY], 7), 42); } + #[test] + fn generic_actions_parses_per_entity_allowlists() { + let cfg = json!({"generic_actions": { + "user": { + "kinds": ["suspend", "label", "bounce_captcha"], + "suspend_policies": ["PlatformManipulation"], + "labels": ["SpamHighRecall"], + }, + "post": { + "kinds": ["post_label", "suspend_author"], + "suspend_policies": ["Cse"], + "labels": ["SpamHighRecall", "RiskyHighVizReply"], + }, + }}); + let user = generic_actions_from_config(Some(&cfg), EntityType::User); + assert!(user.kinds.contains("suspend")); + assert!(user.kinds.contains("bounce_captcha")); + assert!(!user.kinds.contains("post_label"), "no cross-entity bleed"); + assert!(user.suspend_policies.contains("PlatformManipulation")); + assert!(!user.suspend_policies.contains("Cse")); + assert_eq!(user.labels.len(), 1); + + let post = generic_actions_from_config(Some(&cfg), EntityType::Post); + assert!(post.kinds.contains("suspend_author")); + assert!(!post.kinds.contains("suspend")); + assert!(post.suspend_policies.contains("Cse")); + assert!(post.labels.contains("RiskyHighVizReply")); + } + + #[test] + fn generic_actions_absent_or_partial_config_fails_closed() { + let empty = generic_actions_from_config(None, EntityType::User); + assert_eq!( + empty, + crate::generic_actions::GenericActionAllowlist::default() + ); + + let cfg = json!({"dry_run": true}); + let a = generic_actions_from_config(Some(&cfg), EntityType::User); + assert!(a.kinds.is_empty() && a.suspend_policies.is_empty() && a.labels.is_empty()); + + let cfg = json!({"generic_actions": {"post": {"kinds": ["post_label"]}}}); + let user = generic_actions_from_config(Some(&cfg), EntityType::User); + assert!(user.kinds.is_empty()); + let post = generic_actions_from_config(Some(&cfg), EntityType::Post); + assert!(post.kinds.contains("post_label")); + assert!(post.suspend_policies.is_empty()); + assert!(post.labels.is_empty()); + } + + #[test] + fn generic_actions_malformed_values_fail_closed() { + let cfg = json!({"generic_actions": { + "user": { + "kinds": "suspend", + "suspend_policies": [1, true, null], + "labels": ["SpamHighRecall", 42], + }, + }}); + let user = generic_actions_from_config(Some(&cfg), EntityType::User); + assert!(user.kinds.is_empty()); + assert!(user.suspend_policies.is_empty()); + assert_eq!(user.labels.len(), 1); + assert!(user.labels.contains("SpamHighRecall")); + + let cfg = json!({"generic_actions": ["suspend"]}); + let user = generic_actions_from_config(Some(&cfg), EntityType::User); + assert_eq!( + user, + crate::generic_actions::GenericActionAllowlist::default() + ); + } + + #[test] + fn generic_actions_empty_string_entries_are_dropped() { + let cfg = json!({"generic_actions": { + "user": { + "kinds": ["suspend", ""], + "suspend_policies": ["", " ", "PlatformManipulation"], + "labels": ["\t", "SpamHighRecall"], + }, + }}); + let user = generic_actions_from_config(Some(&cfg), EntityType::User); + assert_eq!(user.kinds.len(), 1); + assert_eq!(user.suspend_policies.len(), 1); + assert!(user.suspend_policies.contains("PlatformManipulation")); + assert_eq!(user.labels.len(), 1); + assert!(user.labels.contains("SpamHighRecall")); + } + #[test] fn topic_labels_prefers_nested_then_falls_back_to_flat() { let cfg = json!({ diff --git a/abuse-enforcement-service/service-lib/src/lib.rs b/abuse-enforcement-service/service-lib/src/lib.rs index 244ab59e..89cd1f11 100644 --- a/abuse-enforcement-service/service-lib/src/lib.rs +++ b/abuse-enforcement-service/service-lib/src/lib.rs @@ -7,6 +7,7 @@ pub mod decision; pub mod dedup_cache; pub mod entities; pub mod facts; +pub mod generic_actions; pub mod gizmoduck; pub mod gizmoduck_labels; pub mod growthbook; @@ -41,16 +42,15 @@ use tokio::sync::Semaphore; use tokio::task::JoinHandle; use tracing::{Instrument, error, info, warn}; use xai_kafka::{ - BatchConsumerConfig, BatchResult, CancellationToken, KafkaBatchProcessor, KafkaConsumerBuilder, - KafkaConsumerConfig, KafkaMessage, KafkaProducer, KafkaProducerConfigBuilder, SslConfig, + BatchConsumerConfig, BatchResult, CancellationToken, KafkaBatchProcessor, KafkaConsumerConfig, + KafkaConsumerConfigBuilder, KafkaMessage, KafkaProducer, KafkaProducerConfigBuilder, apply_auth_config, resolve_kafka_brokers, run_batch_consumer, self_delete_pod, }; use xai_service_runner::{ServerBuilder, ServerInfo}; -use xai_wily::WilyConfig; use xai_strato::{Strato, StratoClientConfig}; -use crate::config::{Config, KafkaConnConfig}; +use crate::config::Config; use crate::decision::Decision; use crate::facts::{EntityFacts, EntityType, Facts, PostFacts, ScoreFacts, UserFacts}; use crate::growthbook::DynamicConfig; @@ -132,6 +132,7 @@ struct EnforcementCtx { rules_cache: Arc, allowlist: Option, kafka_producer_decisions: Option>, + kafka_producer_decisions_json: Option>, } struct ScoreResultProcessor { @@ -170,6 +171,35 @@ fn decision_outcome( } } +const REQUESTED_ACTIONS_DENIED: &str = "requested_actions_denied"; + +#[derive(Debug, PartialEq)] +enum ExpandedDecision { + Skip(String), + Act(Vec), +} + +fn expand_requested_actions_decision( + entity_type: EntityType, + score_facts: &ScoreFacts, + allowlist: &generic_actions::GenericActionAllowlist, +) -> (ExpandedDecision, Option) { + let resolved = generic_actions::resolve_requested_actions( + entity_type, + &score_facts.requested_actions, + allowlist, + ); + let skipped_json = resolved.skipped_info_json(); + if resolved.specs.is_empty() { + ( + ExpandedDecision::Skip(REQUESTED_ACTIONS_DENIED.into()), + skipped_json, + ) + } else { + (ExpandedDecision::Act(resolved.specs), skipped_json) + } +} + #[tracing::instrument(skip_all, fields(dry_run, uas))] async fn run_enforcement_inner( ctx: &EnforcementCtx, @@ -341,20 +371,41 @@ async fn run_enforcement_inner( } }; + let (decision, mut requested_actions_skipped) = match decision { + Decision::ActRequestedActions => expand_requested_actions_decision( + facts.entity_type, + &facts.score, + &ctx.dynamic_config + .generic_action_allowlist(facts.entity_type), + ), + Decision::Skip(reason) => (ExpandedDecision::Skip(reason), None), + Decision::Act(specs) => (ExpandedDecision::Act(specs), None), + }; + match decision { - Decision::Skip(reason) => Ok(skip_outcome(reason, dry_run, &facts, score)), - Decision::Act(specs) => { + ExpandedDecision::Skip(reason) => { + let mut outcome = skip_outcome(reason, dry_run, &facts, score); + if let Some(json) = requested_actions_skipped { + outcome + .info + .insert("requested_actions_skipped".into(), json); + } + Ok(outcome) + } + ExpandedDecision::Act(specs) => { if !ctx.dynamic_config.try_enforce(facts.entity_type).await { warn!( entity_type = facts.entity_type.as_str(), "max enforcement has been reached; skipping" ); - return Ok(skip_outcome( - "max_enforcement_reached".into(), - dry_run, - &facts, - score, - )); + let mut outcome = + skip_outcome("max_enforcement_reached".into(), dry_run, &facts, score); + if let Some(json) = requested_actions_skipped.take() { + outcome + .info + .insert("requested_actions_skipped".into(), json); + } + return Ok(outcome); } let user_id = facts.user_id; @@ -391,6 +442,15 @@ async fn run_enforcement_inner( .map(|a| (a.name(), a.metric_label())) .collect(); + additional_info_map.insert( + "action_kinds".into(), + serde_json::to_string(&uas_action_dims.iter().map(|(a, _)| *a).collect::>()) + .unwrap_or_else(|_| "[]".into()), + ); + if let Some(json) = requested_actions_skipped { + additional_info_map.insert("requested_actions_skipped".into(), json); + } + enforce_actions( &ctx.ais_client, dry_run, @@ -493,6 +553,8 @@ async fn write_dedup_outcome( const DECISIONS_PRODUCER: &str = "decisions"; +const DECISIONS_JSON_PRODUCER: &str = "decisions_json"; + pub(crate) const ADMIN_ACTIONS_PRODUCER: &str = "admin_actions"; #[derive(Clone, Default)] @@ -509,6 +571,10 @@ impl KafkaProducers { self.get(DECISIONS_PRODUCER) } + pub fn decisions_json(&self) -> Option<&Arc> { + self.get(DECISIONS_JSON_PRODUCER) + } + pub fn admin_actions(&self) -> Option<&Arc> { self.get(ADMIN_ACTIONS_PRODUCER) } @@ -531,9 +597,32 @@ pub(crate) fn spawn_publish( let Some(producer) = producer else { return; }; - let key = key(); - let bytes = record.encode_to_vec(); - let producer = producer.clone(); + spawn_send(producer.clone(), sink, key(), record.encode_to_vec()); +} + +pub(crate) fn spawn_publish_json( + producer: Option<&Arc>, + sink: &'static str, + key: impl FnOnce() -> Vec, + record: &M, +) { + let Some(producer) = producer else { + return; + }; + let bytes = match serde_json::to_vec(record) { + Ok(bytes) => bytes, + Err(e) => { + warn!("{sink} JSON encode failed, record dropped: {e}"); + metrics::KAFKA_PUBLISH_TOTAL + .with_label_values(&[sink, "encode_error"]) + .inc(); + return; + } + }; + spawn_send(producer.clone(), sink, key(), bytes); +} + +fn spawn_send(producer: Arc, sink: &'static str, key: Vec, bytes: Vec) { tokio::spawn(async move { let result = match producer.send_with_key(Some(key.as_slice()), &bytes).await { Ok(_) => "ok", @@ -548,16 +637,19 @@ pub(crate) fn spawn_publish( }); } -fn publish_decision_outcome( - producer: Option<&Arc>, - outcome: &abuse_proto::DecisionOutcome, -) { +fn publish_decision_outcome(ctx: &EnforcementCtx, outcome: &abuse_proto::DecisionOutcome) { spawn_publish( - producer, + ctx.kafka_producer_decisions.as_ref(), DECISIONS_PRODUCER, || outcome.entity_id.to_string().into_bytes(), outcome, ); + spawn_publish_json( + ctx.kafka_producer_decisions_json.as_ref(), + DECISIONS_JSON_PRODUCER, + || outcome.entity_id.to_string().into_bytes(), + outcome, + ); } async fn run_enforcement( @@ -566,29 +658,13 @@ async fn run_enforcement( source_topic: &str, ) -> Result { let outcome = run_enforcement_inner(ctx, score, source_topic).await?; - publish_decision_outcome(ctx.kafka_producer_decisions.as_ref(), &outcome); + publish_decision_outcome(ctx, &outcome); Ok(outcome) } -fn sasl_ssl_config(conn: &KafkaConnConfig, password: &str) -> SslConfig { - SslConfig { - security_protocol: "SASL_SSL".to_string(), - sasl_mechanism: Some(conn.sasl_mechanism.clone()), - sasl_username: Some(conn.sasl_username.clone()), - sasl_password: Some(password.to_owned()), - } -} - async fn build_kafka_producers(cfg: &Config, dynamic_config: &DynamicConfig) -> KafkaProducers { let mut producers = KafkaProducers::default(); - let conn = cfg.kafka_producer(); - let sasl_password = conn.sasl_password.clone(); - - if !cfg.kafka_producer_mtls_enabled && sasl_password.is_none() { - return producers; - } - for (name, spec) in dynamic_config.kafka_producers() { if !spec.enabled { warn!("kafka producer '{name}' configured but disabled (enabled=false)"); @@ -598,28 +674,16 @@ async fn build_kafka_producers(cfg: &Config, dynamic_config: &DynamicConfig) -> warn!("kafka producer '{name}' enabled but no topic set; skipping"); continue; }; - let producer_config = if cfg.kafka_producer_mtls_enabled { - match KafkaProducerConfigBuilder::for_cluster_mtls_auto( - &cfg.kafka_producer_mtls_cluster, - topic.clone(), - Some(&cfg.kafka_producer_mtls_zone), - ) { - Ok(builder) => builder.build(), - Err(e) => { - error!("kafka producer '{name}' mTLS config failed; skipping: {e}"); - continue; - } + let producer_config = match KafkaProducerConfigBuilder::for_cluster_mtls_auto( + &cfg.kafka_producer_mtls_cluster, + topic.clone(), + Some(&cfg.kafka_producer_mtls_zone), + ) { + Ok(builder) => builder.build(), + Err(e) => { + error!("kafka producer '{name}' mTLS config failed; skipping: {e}"); + continue; } - } else { - KafkaProducerConfigBuilder::new(conn.dest.clone(), topic.clone()) - .with_wily_config(WilyConfig::default()) - .with_ssl(sasl_ssl_config( - &conn, - sasl_password - .as_deref() - .expect("SASL password checked before producer construction"), - )) - .build() }; let mut producer = KafkaProducer::new(producer_config); match producer.start().await { @@ -663,7 +727,7 @@ impl ScoreResultProcessor { ]) .inc(); publish_decision_outcome( - self.ctx.kafka_producer_decisions.as_ref(), + &self.ctx, &decision_outcome( score, topic, @@ -716,7 +780,7 @@ impl ScoreResultProcessor { ]) .inc(); publish_decision_outcome( - self.ctx.kafka_producer_decisions.as_ref(), + &self.ctx, &decision_outcome(score, topic, "dedup_skipped".into(), false, BTreeMap::new()), ); return Ok("dedup_skipped".into()); @@ -1348,6 +1412,10 @@ async fn kafka_health_watchdog( } } +const KAFKA_PREFLIGHT_ATTEMPTS: u32 = 3; +const KAFKA_PREFLIGHT_TIMEOUT: Duration = Duration::from_secs(10); +const KAFKA_PREFLIGHT_BACKOFF: Duration = Duration::from_secs(2); + async fn probe_broker_reachability(config: KafkaConsumerConfig, timeout: Duration) -> Result<()> { use rdkafka::ClientConfig; use rdkafka::consumer::{BaseConsumer, Consumer}; @@ -1386,14 +1454,11 @@ pub async fn start_kafka_consumers( state: Arc, cfg: &Config, ) -> Result> { - let consumer_conn = cfg.kafka_consumer(); - let Some(sasl_password) = consumer_conn.sasl_password.clone() else { - info!( - "no Kafka consumer credential (KAFKA_SASL_PASSWORD / KAFKA_CONSUMER_SASL_PASSWORD unset) — Kafka consumer disabled" - ); + if !cfg.kafka_consumer_enabled { + info!("KAFKA_CONSUMER_ENABLED=false — Kafka consumer disabled"); state.kafka_ready.store(true, Ordering::Relaxed); return Ok(tokio::spawn(async {})); - }; + } let growthbook_enabled = cfg.growthbook_url.is_some() && cfg.growthbook_key.is_some(); let topic_labels_json: serde_json::Value = @@ -1446,25 +1511,24 @@ pub async fn start_kafka_consumers( } if cfg.kafka_self_delete_enabled && !topics.is_empty() { - const PREFLIGHT_ATTEMPTS: u32 = 3; - const PREFLIGHT_TIMEOUT: Duration = Duration::from_secs(10); - const PREFLIGHT_BACKOFF: Duration = Duration::from_secs(2); - let probe_topic = topics.keys().min().expect("topics non-empty").clone(); - let probe_config: KafkaConsumerConfig = KafkaConsumerBuilder::new( - consumer_conn.dest.clone(), + let probe_config = KafkaConsumerConfigBuilder::for_cluster_mtls_auto( + &cfg.kafka_consumer_mtls_cluster, probe_topic.clone(), format!("{}-preflight", cfg.kafka_group_id()), + Some(&cfg.kafka_consumer_mtls_zone), ) - .with_wily_config(WilyConfig::default()) - .with_ssl(sasl_ssl_config(&consumer_conn, &sasl_password)) - .into(); + .context("failed to configure Phoenix mTLS consumer preflight")? + .with_enable_auto_offset_store(false) + .with_enable_auto_commit(false) + .with_fetch_timeout_ms(10000) + .build(); let mut attempt = 0u32; loop { attempt += 1; let start = std::time::Instant::now(); - match probe_broker_reachability(probe_config.clone(), PREFLIGHT_TIMEOUT).await { + match probe_broker_reachability(probe_config.clone(), KAFKA_PREFLIGHT_TIMEOUT).await { Ok(()) => { info!( attempt, @@ -1474,18 +1538,18 @@ pub async fn start_kafka_consumers( ); break; } - Err(e) if attempt < PREFLIGHT_ATTEMPTS => { + Err(e) if attempt < KAFKA_PREFLIGHT_ATTEMPTS => { warn!( topic = %probe_topic, - "Kafka broker preflight attempt {attempt}/{PREFLIGHT_ATTEMPTS} failed \ - (retrying in {PREFLIGHT_BACKOFF:?}): {e:#}" + "Kafka broker preflight attempt {attempt}/{KAFKA_PREFLIGHT_ATTEMPTS} failed \ + (retrying in {KAFKA_PREFLIGHT_BACKOFF:?}): {e:#}" ); - tokio::time::sleep(PREFLIGHT_BACKOFF).await; + tokio::time::sleep(KAFKA_PREFLIGHT_BACKOFF).await; } Err(e) => { error!( topic = %probe_topic, - "Kafka broker preflight failed after {PREFLIGHT_ATTEMPTS} attempts — \ + "Kafka broker preflight failed after {KAFKA_PREFLIGHT_ATTEMPTS} attempts — \ likely a bad node; self-deleting to reschedule: {e:#}" ); metrics::KAFKA_SELF_DELETE_TOTAL @@ -1493,7 +1557,7 @@ pub async fn start_kafka_consumers( .inc(); self_delete_pod().await; return Err(anyhow::anyhow!( - "Kafka broker preflight failed after {PREFLIGHT_ATTEMPTS} attempts: {e:#}" + "Kafka broker preflight failed after {KAFKA_PREFLIGHT_ATTEMPTS} attempts: {e:#}" )); } } @@ -1570,6 +1634,7 @@ pub async fn start_kafka_consumers( rules_cache: state.rules_cache.clone(), allowlist: state.allowlist.clone(), kafka_producer_decisions: state.kafka_producers.decisions().cloned(), + kafka_producer_decisions_json: state.kafka_producers.decisions_json().cloned(), }; { @@ -1696,22 +1761,25 @@ pub async fn start_kafka_consumers( max_in_flight, "Kafka consumer batch limits" ); - let make_batch_config = |topic: &str| { - let kafka = KafkaConsumerBuilder::new( - consumer_conn.dest.clone(), + let make_batch_config = |topic: &str| -> Result { + let kafka = KafkaConsumerConfigBuilder::for_cluster_mtls_auto( + &cfg.kafka_consumer_mtls_cluster, topic.to_owned(), - format!("{}-{}", kafka_group_id, topic), + topic_consumer_group_id(&kafka_group_id, topic), + Some(&cfg.kafka_consumer_mtls_zone), ) - .with_wily_config(WilyConfig::default()) - .with_ssl(sasl_ssl_config(&consumer_conn, &sasl_password)) - .with_fetch_timeout_ms(10000); - - BatchConsumerConfig::new(kafka, SERVICE_NAME) - .with_max_messages_per_poll(max_messages_per_poll) + .with_context(|| format!("failed to configure Phoenix mTLS consumer for {topic}"))? + .with_enable_auto_offset_store(false) + .with_enable_auto_commit(false) + .with_fetch_timeout_ms(10000) + .build(); + + Ok(BatchConsumerConfig::new(kafka, SERVICE_NAME) + .with_max_messages_per_poll(max_messages_per_poll)) }; for (topic, topic_cfg) in topics { - let batch_config = make_batch_config(&topic); + let batch_config = make_batch_config(&topic)?; let cancel = cancel.clone(); match topic_cfg { @@ -1781,6 +1849,10 @@ pub async fn start_kafka_consumers( })) } +fn topic_consumer_group_id(base_group_id: &str, topic: &str) -> String { + format!("{base_group_id}-{topic}") +} + pub async fn serve(router: Router, cfg: &Config) -> Result<()> { let server = ServerBuilder::new(cfg.port) .merge(router) @@ -1852,6 +1924,36 @@ mod router_split_tests { } } +#[cfg(test)] +mod kafka_topic_config_tests { + use super::{ + KAFKA_PREFLIGHT_ATTEMPTS, KAFKA_PREFLIGHT_BACKOFF, KAFKA_PREFLIGHT_TIMEOUT, + topic_consumer_group_id, + }; + use std::time::Duration; + + #[test] + fn dynamic_topics_keep_distinct_existing_group_ids() { + let base_group_id = "xai-abuse-enforcement-service"; + + assert_eq!( + topic_consumer_group_id(base_group_id, "scores.primary"), + "xai-abuse-enforcement-service-scores.primary" + ); + assert_eq!( + topic_consumer_group_id(base_group_id, "scores.secondary"), + "xai-abuse-enforcement-service-scores.secondary" + ); + } + + #[test] + fn preflight_retry_budget_stays_bounded() { + assert_eq!(KAFKA_PREFLIGHT_ATTEMPTS, 3); + assert_eq!(KAFKA_PREFLIGHT_TIMEOUT, Duration::from_secs(10)); + assert_eq!(KAFKA_PREFLIGHT_BACKOFF, Duration::from_secs(2)); + } +} + #[cfg(test)] mod dedup_retention_tests { use super::outcome_holds_full_dedup; @@ -1862,6 +1964,7 @@ mod dedup_retention_tests { for skip in [ "dry_run", "dedup_skipped", + "very_high_follower_count", "high_follower_count", "pagerank_skipped", "gizmoduck_skipped", @@ -1872,12 +1975,175 @@ mod dedup_retention_tests { "rule_eval_error", "max_enforcement_reached", "invalid_entity_id", + "requested_actions_denied", + "platform_row_without_requested_actions", ] { assert!(!outcome_holds_full_dedup(skip), "{skip} must not hold 24h"); } } } +#[cfg(test)] +mod generic_dispatch_tests { + use super::*; + use crate::decision::ActionSpec; + use crate::facts::RequestedActionFacts; + use crate::generic_actions::GenericActionAllowlist; + + fn score_facts_with(requested: Vec) -> ScoreFacts { + let mut f = ScoreFacts::from_score(&abuse_proto::ScoreResult::default()); + f.requested_actions = requested; + f + } + + fn user_allowlist() -> GenericActionAllowlist { + GenericActionAllowlist { + kinds: ["suspend", "label"] + .iter() + .map(|s| (*s).to_owned()) + .collect(), + suspend_policies: ["PlatformManipulation"] + .iter() + .map(|s| (*s).to_owned()) + .collect(), + labels: ["SpamHighRecall"].iter().map(|s| (*s).to_owned()).collect(), + } + } + + #[test] + fn expand_allowed_requested_actions_yields_plain_act_decision() { + let facts = score_facts_with(vec![RequestedActionFacts { + kind: "suspend".into(), + perm: false, + policy: "PlatformManipulation".into(), + head: "IsSpammer".into(), + ..Default::default() + }]); + let (decision, skipped) = + expand_requested_actions_decision(EntityType::User, &facts, &user_allowlist()); + assert_eq!( + decision, + ExpandedDecision::Act(vec![ActionSpec::SuspendUser { + perm: false, + policy: "PlatformManipulation".into(), + }]) + ); + assert!(skipped.is_none()); + } + + #[test] + fn expand_denied_requested_actions_yields_skip_with_reason() { + let facts = score_facts_with(vec![RequestedActionFacts { + kind: "suspend".into(), + policy: "PlatformManipulation".into(), + head: "IsSpammer".into(), + ..Default::default() + }]); + let (decision, skipped) = expand_requested_actions_decision( + EntityType::User, + &facts, + &GenericActionAllowlist::default(), + ); + assert_eq!( + decision, + ExpandedDecision::Skip(REQUESTED_ACTIONS_DENIED.into()) + ); + let skipped = skipped.expect("refused entries must be reported"); + assert!(skipped.contains("kind_not_allowlisted"), "{skipped}"); + } + + #[test] + fn expand_empty_requested_actions_yields_skip() { + let (decision, skipped) = expand_requested_actions_decision( + EntityType::User, + &score_facts_with(vec![]), + &user_allowlist(), + ); + assert_eq!( + decision, + ExpandedDecision::Skip(REQUESTED_ACTIONS_DENIED.into()) + ); + assert!(skipped.is_none()); + } + + #[test] + fn expand_partial_allowlist_dispatches_allowed_and_reports_skipped() { + let facts = score_facts_with(vec![ + RequestedActionFacts { + kind: "label".into(), + labels: vec!["SpamHighRecall".into()], + ttl_msec: 1000, + head: "IsLabelHead".into(), + ..Default::default() + }, + RequestedActionFacts { + kind: "bounce_captcha".into(), + head: "IsCuspHead".into(), + ..Default::default() + }, + ]); + let (decision, skipped) = + expand_requested_actions_decision(EntityType::User, &facts, &user_allowlist()); + assert_eq!( + decision, + ExpandedDecision::Act(vec![ActionSpec::AddLabelsV2 { + labels: vec!["SpamHighRecall".into()], + ttl_msec: Some(1000), + }]) + ); + let skipped = skipped.expect("refused entry must be reported"); + assert!(skipped.contains("bounce_captcha"), "{skipped}"); + assert!(skipped.contains("kind_not_allowlisted"), "{skipped}"); + } +} + +#[cfg(test)] +mod decision_outcome_json_tests { + use super::*; + use xai_abuse_proto::enforcement::{ + EntityType as ProtoEntityType, FiredHead, ScoreResult, SummaryCounters, + }; + + #[test] + fn decision_outcome_json_mirror_carries_funnel_fields() { + let score = ScoreResult { + user_id: 100, + model_version: "my_model@1".into(), + entity_type: ProtoEntityType::Post as i32, + entity_id: 555, + summary: Some(SummaryCounters { + labels: vec!["my_model_threshold_reached".into()], + fired_heads: vec![FiredHead { + name: "IsSpamPost".into(), + score: 0.99, + threshold: 0.9, + }], + ..Default::default() + }), + ..Default::default() + }; + let mut info = BTreeMap::new(); + info.insert( + "action_kinds".to_owned(), + r#"["addPostLabelsV2"]"#.to_owned(), + ); + let outcome = decision_outcome(&score, "some.topic", "success".into(), true, info); + let v = serde_json::to_value(&outcome).expect("DecisionOutcome serializes to JSON"); + + assert!(v["decided_at_ms"].as_i64().unwrap() > 0); + assert_eq!(v["source_topic"], "some.topic"); + assert_eq!(v["entity_type"], "post"); + assert_eq!(v["entity_id"], 555); + assert_eq!(v["model_version"], "my_model@1"); + assert_eq!(v["status"], "success"); + assert_eq!(v["dry_run"], true); + assert_eq!(v["head"], "IsSpamPost"); + assert_eq!(v["fired_heads"][0]["name"], "IsSpamPost"); + assert_eq!(v["labels"][0], "my_model_threshold_reached"); + assert_eq!(v["info"]["action_kinds"], r#"["addPostLabelsV2"]"#); + } +} + #[cfg(test)] mod health_tests { use super::health_decision; diff --git a/abuse-enforcement-service/service-lib/src/metrics.rs b/abuse-enforcement-service/service-lib/src/metrics.rs index 96e3a539..fda5fe6a 100644 --- a/abuse-enforcement-service/service-lib/src/metrics.rs +++ b/abuse-enforcement-service/service-lib/src/metrics.rs @@ -175,6 +175,14 @@ lazy_static! { .unwrap(); + pub static ref GENERIC_ACTION_TOTAL: IntCounterVec = + register_int_counter_vec!( + "abuse_enforcement_generic_action_total", + "Requested actions handled by the generic executor, by entity_type, kind, result, and fail-closed skip reason.", + &["entity_type", "kind", "result", "reason"]) + .unwrap(); + + pub static ref HTTP_LATENCY: HistogramVec = register_histogram_vec!( "abuse_enforcement_http_latency_seconds", @@ -191,7 +199,7 @@ lazy_static! { .unwrap(); - pub static ref KAFKA_PUBLISH_TOTAL: IntCounterVec = + pub static ref KAFKA_PUBLISH_TOTAL: IntCounterVec = register_int_counter_vec!( "abuse_enforcement_kafka_publish_total", "Records published to a Kafka publish sink, by sink and produce result.", @@ -258,6 +266,7 @@ pub fn init() { let _ = &*MANHATTAN_ERRORS_TOTAL; let _ = &*KAFKA_SELF_DELETE_TOTAL; let _ = &*RULES_YAML_COMPILED; + let _ = &*GENERIC_ACTION_TOTAL; let _ = &*HTTP_LATENCY; let _ = &*HTTP_INFLIGHT; let _ = &*KAFKA_PUBLISH_TOTAL; diff --git a/abuse-enforcement-service/service-lib/src/rules.rs b/abuse-enforcement-service/service-lib/src/rules.rs index 3ebcdebf..ef081aa5 100644 --- a/abuse-enforcement-service/service-lib/src/rules.rs +++ b/abuse-enforcement-service/service-lib/src/rules.rs @@ -51,6 +51,7 @@ enum ActionStep { enum SpecialOutcome { Skip { reason: String }, ActAll { actions: Vec }, + ActRequestedActions, } #[derive(Clone, Debug, PartialEq, Eq, Deserialize)] @@ -172,6 +173,7 @@ fn outcome_to_decision(o: &Outcome) -> Decision { Outcome::Special(SpecialOutcome::ActAll { actions }) => { Decision::Act(actions.iter().map(step_to_spec).collect()) } + Outcome::Special(SpecialOutcome::ActRequestedActions) => Decision::ActRequestedActions, Outcome::Action(step) => Decision::Act(vec![step_to_spec(step)]), } } @@ -182,6 +184,18 @@ struct ScoreCel<'a> { labels: &'a [String], model_version: &'a str, skip_author_credibility_prechecks: bool, + requested_actions: Vec>, + policy_version: &'a str, +} + +#[derive(Serialize)] +struct RequestedActionCel<'a> { + kind: &'a str, + perm: bool, + policy: &'a str, + labels: &'a [String], + ttl_msec: i64, + head: &'a str, } #[derive(Serialize)] @@ -217,6 +231,20 @@ fn project_score(facts: &Facts) -> ScoreCel<'_> { labels: &facts.score.labels, model_version: &facts.score.model_version, skip_author_credibility_prechecks: facts.score.skip_author_credibility_prechecks, + requested_actions: facts + .score + .requested_actions + .iter() + .map(|a| RequestedActionCel { + kind: &a.kind, + perm: a.perm, + policy: &a.policy, + labels: &a.labels, + ttl_msec: a.ttl_msec, + head: &a.head, + }) + .collect(), + policy_version: &facts.score.policy_version, } } @@ -642,6 +670,7 @@ mod tests { for id in [ "user_in_allowlist", "user_not_found", + "very_high_follower_count", "high_follower_count", "pagerank_skipped", ] { @@ -651,6 +680,17 @@ mod tests { user.rule_ids ); } + let pos = |id: &str| { + user.rule_ids + .iter() + .position(|r| r == id) + .unwrap_or_else(|| panic!("user pipeline missing {id:?}")) + }; + assert!( + pos("very_high_follower_count") < pos("high_follower_count"), + "very_high_follower_count must precede high_follower_count; got {:?}", + user.rule_ids + ); for id in [ "post_in_allowlist", "user_in_allowlist", @@ -1069,6 +1109,249 @@ rules: } + fn facts_with_requested_action() -> Facts { + let mut f = base_facts(); + f.score.requested_actions = vec![crate::facts::RequestedActionFacts { + kind: "suspend".into(), + perm: false, + policy: "PlatformManipulation".into(), + labels: vec!["SpamHighRecall".into()], + ttl_msec: 1000, + head: "IsSpammer".into(), + }]; + f.score.policy_version = "3".into(); + f + } + + #[test] + fn act_requested_actions_outcome_parses_and_maps_to_decision() { + let rules = one_rule_pipeline( + r#" +rules: + - id: generic_test + when: "size(score.requested_actions) > 0" + then: + kind: act_requested_actions + - id: fallthrough + when: "true" + then: {kind: skip, reason: no_requested_actions} +"#, + ); + assert_eq!( + decide_with(&rules, &base_facts()).expect("eval should succeed"), + Decision::Skip("no_requested_actions".into()) + ); + assert_eq!( + decide_with(&rules, &facts_with_requested_action()).expect("eval should succeed"), + Decision::ActRequestedActions + ); + } + + #[test] + fn requested_actions_entries_and_policy_version_are_cel_readable() { + let rules = one_rule_pipeline( + r#" +rules: + - id: entry_fields + when: > + score.policy_version == "3" + && score.requested_actions.exists(a, + a.kind == "suspend" + && a.policy == "PlatformManipulation" + && !a.perm + && a.ttl_msec == 1000 + && a.head == "IsSpammer" + && "SpamHighRecall" in a.labels) + then: {kind: skip, reason: matched} + - id: fallthrough + when: "true" + then: {kind: skip, reason: unmatched} +"#, + ); + assert_eq!( + decide_with(&rules, &facts_with_requested_action()).expect("eval should succeed"), + Decision::Skip("matched".into()) + ); + assert_eq!( + decide_with(&rules, &base_facts()).expect("eval should succeed"), + Decision::Skip("unmatched".into()) + ); + } + + #[test] + fn act_requested_actions_is_rejected_inside_act_all() { + let yaml = "rules:\n - id: x\n when: \"true\"\n then: {kind: act_all, actions: [{kind: act_requested_actions}]}\n"; + let err = CompiledRules::from_yaml(yaml).unwrap_err(); + assert!(matches!(err, RuleCompileError::Yaml(_)), "{err:?}"); + } + + + fn rule_index(ids: &[String], id: &str) -> usize { + ids.iter() + .position(|r| r == id) + .unwrap_or_else(|| panic!("rule {id:?} missing; got {ids:?}")) + } + + #[test] + fn baked_in_generic_rule_sits_directly_above_each_terminal_rule() { + let cache = RulesCache::new(); + let user_ids = cache.status(EntityType::User, None).rule_ids; + let generic = rule_index(&user_ids, "act_requested_actions"); + let guard = rule_index(&user_ids, "platform_row_without_requested_actions"); + let terminal = rule_index(&user_ids, "act_suspend"); + assert_eq!(terminal, user_ids.len() - 1, "act_suspend must be last"); + assert_eq!( + generic, + terminal - 1, + "generic rule must sit directly above act_suspend" + ); + assert_eq!( + guard, + generic - 1, + "empty-list guard must sit directly above the generic rule" + ); + assert!(generic > rule_index(&user_ids, "pagerank_skipped")); + + let post_ids = cache.status(EntityType::Post, None).rule_ids; + let generic = rule_index(&post_ids, "act_requested_actions"); + let terminal = rule_index(&post_ids, "post_no_actionable_label"); + assert_eq!( + terminal, + post_ids.len() - 1, + "post terminal skip must be last" + ); + assert_eq!( + generic, + terminal - 1, + "generic rule must sit directly above the post terminal skip" + ); + assert!(generic > rule_index(&post_ids, "pagerank_skipped")); + } + + #[test] + fn baked_in_user_pipeline_routes_requested_actions_to_generic_dispatch() { + let rules = RulesCache::new().resolve(EntityType::User, None); + assert_eq!( + decide_with(&rules, &facts_with_requested_action()).expect("eval should succeed"), + Decision::ActRequestedActions + ); + assert_eq!( + decide_with(&rules, &base_facts()).expect("eval should succeed"), + Decision::Act(vec![ActionSpec::SuspendUser { + perm: false, + policy: "PlatformManipulation".into(), + }]) + ); + } + + #[test] + fn baked_in_user_platform_label_with_empty_list_skips_not_suspends() { + let rules = RulesCache::new().resolve(EntityType::User, None); + let mut f = base_facts(); + f.score + .labels + .push("abuse_platform_requested_actions".into()); + assert_eq!( + decide_with(&rules, &f).expect("eval should succeed"), + Decision::Skip("platform_row_without_requested_actions".into()) + ); + let mut f = facts_with_requested_action(); + f.score + .labels + .push("abuse_platform_requested_actions".into()); + assert_eq!( + decide_with(&rules, &f).expect("eval should succeed"), + Decision::ActRequestedActions + ); + } + + #[test] + fn baked_in_user_guardrails_still_win_over_generic_dispatch() { + let rules = RulesCache::new().resolve(EntityType::User, None); + let mut f = facts_with_requested_action(); + f.user_allowlist_mut().is_allowlisted = true; + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("user_in_allowlist".into()) + ); + let mut f = facts_with_requested_action(); + f.cred_mut().follower_count = Some(5_000); + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("high_follower_count".into()) + ); + } + + #[test] + fn baked_in_user_pipeline_splits_follower_bands_at_the_ceiling() { + let rules = RulesCache::new().resolve(EntityType::User, None); + let mut f = facts_with_requested_action(); + f.cred_mut().follower_count = Some(250_000); + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("very_high_follower_count".into()) + ); + f.cred_mut().follower_count = Some(50_000); + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("high_follower_count".into()) + ); + f.cred_mut().follower_count = Some(500); + assert_ne!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("high_follower_count".into()) + ); + assert_ne!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("very_high_follower_count".into()) + ); + } + + #[test] + fn baked_in_user_ceiling_ignores_skip_author_credibility_prechecks() { + let rules = RulesCache::new().resolve(EntityType::User, None); + let mut f = facts_with_requested_action(); + f.cred_mut().follower_count = Some(250_000); + f.score.skip_author_credibility_prechecks = true; + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("very_high_follower_count".into()) + ); + } + + #[test] + fn baked_in_post_pipeline_routes_requested_actions_to_generic_dispatch() { + let rules = RulesCache::new().resolve(EntityType::Post, None); + let mut f = post_facts(safe_author(), vec![]); + f.score.requested_actions = vec![crate::facts::RequestedActionFacts { + kind: "post_label".into(), + labels: vec!["SpamHighRecall".into()], + head: "IsSpamPost".into(), + ..Default::default() + }]; + assert_eq!( + decide_with(&rules, &f).expect("eval should succeed"), + Decision::ActRequestedActions + ); + assert_eq!( + decide_with(&rules, &post_facts(safe_author(), vec![])).unwrap(), + Decision::Skip("post_no_actionable_label".into()) + ); + let mut author = safe_author(); + author.cred.is_high = Some(true); + let mut f = post_facts(author, vec![]); + f.score.requested_actions = vec![crate::facts::RequestedActionFacts { + kind: "post_label".into(), + labels: vec!["SpamHighRecall".into()], + ..Default::default() + }]; + assert_eq!( + decide_with(&rules, &f).unwrap(), + Decision::Skip("pagerank_skipped".into()) + ); + } + + #[test] fn mock_user_allowlisted_skips() { let mut f = base_facts(); diff --git a/abuse-enforcement-service/service-lib/src/service.rs b/abuse-enforcement-service/service-lib/src/service.rs index fd63db9d..a0ea8f29 100644 --- a/abuse-enforcement-service/service-lib/src/service.rs +++ b/abuse-enforcement-service/service-lib/src/service.rs @@ -2100,6 +2100,64 @@ mod tests { const SPAM_HIGH_RECALL_LABEL: &str = "SpamHighRecall"; const LABEL_EXPIRY_MSEC: i64 = 30 * 24 * 60 * 60 * 1000; + #[test] + fn every_action_shape_encodes_against_the_ais_schema() { + use xai_abuse_thrift_codec::{SCHEMA_AIS, Transcoder}; + + let transcoder = Transcoder::from_schema_bytes(SCHEMA_AIS).expect("AIS schema loads"); + let method = transcoder + .lookup_method( + "com.twitter.agenttools.ais.thriftscala.ActionIntakeService", + "intakeAction", + ) + .expect("intakeAction in schema"); + + let actions = [ + EnforcementAction::suspend_user(1, false, POLICY_IN_VIOLATION, vec!["n".to_owned()]), + EnforcementAction::suspend_user(1, true, POLICY_IN_VIOLATION, vec![]), + EnforcementAction::add_user_label( + 1, + vec![SPAM_HIGH_RECALL_LABEL.to_owned()], + Some(LABEL_EXPIRY_MSEC), + vec!["n".to_owned()], + ), + EnforcementAction::add_user_label( + 1, + vec![SPAM_HIGH_RECALL_LABEL.to_owned()], + None, + vec![], + ), + EnforcementAction::add_post_label( + 2, + vec![SPAM_HIGH_RECALL_LABEL.to_owned()], + Some(LABEL_EXPIRY_MSEC), + vec![], + ), + EnforcementAction::bounce_arkose(1, vec!["n".to_owned()]), + EnforcementAction::bounce_captcha(1, vec![]), + EnforcementAction::spam_liveness_check(1, vec![]), + ]; + + for action in actions { + let request = crate::ais_client::build_intake_request( + action.actor_header(), + action.to_json(), + AIS_TOOL, + Some(action.audit_note()), + ); + transcoder + .encode_call( + method.thrift_method_name, + method.arg_field_name, + method.arg_field_id, + method.request_type, + &request, + 0, + ) + .unwrap_or_else(|e| panic!("{} must encode: {e}", action.name())); + } + } + #[test] fn ais_outcome_failure_is_detected_as_permanent() { let rejection = json!({"success": {"outcome": {"failure": {"businessLogic": { diff --git a/bdsm/README.md b/bdsm/README.md index ec45c545..7b928949 100644 --- a/bdsm/README.md +++ b/bdsm/README.md @@ -1,4 +1,4 @@ -# Behavioral Inauthentic-Account Detection +# Behavioral Detection Sequence Model (BDSM) A sequence-of-actions transformer that detects inauthentic (bot / spam / coordinated) accounts from their behavioral event streams, together with the @@ -106,9 +106,21 @@ The scorer publishes an 8-wide row in `heads.HEAD_NAMES` order. in `score_results_sink_focal.py`) — are **redacted** in this public release: they ship as an out-of-range `9.99` sentinel (the fields are probabilities in `[0, 1]`, so `9.99` never fires and is plainly a placeholder, not a real - value). Publishing exact operating points would hand adversaries the - detector's evasion boundary. The policy *structure*, head names, and gate - logic are real and unredacted; only the tuned numbers are withheld. Supply - your own via `--policy-file` / `BDSM_SINK_POLICY`. -- Per-head **appeal-note templates** that described the tripping features in - prose are withheld. The `enforcement_note` proto field remains and is empty. + value). The min-actions enforcement gate (a count, not a probability) is + redacted the same way with an impossible `999999` sentinel — far longer + than any scoreable sequence. Publishing exact operating points would hand + adversaries the detector's evasion boundary — including the account-size + floor below which scoring never fires. The policy *structure*, head names, + and gate logic are real and unredacted; only the tuned numbers are + withheld. Supply your own via `--policy-file` / `BDSM_SINK_POLICY`. +- Per-head **appeal-note templates**: the production sink interpolates a + short prose paragraph from the dominant bot head and selected histogram + counts (`build_enforcement_note` in `runtime/score_results_sink_focal.py`). + The public package keeps the **gates** (the min-actions gate — value + redacted, wired from the policy — and the dominant-head pick) and the + `enforcement_note` proto field. The template *strings* and the + per-head `key_actions` interpolator are the sentinel `""` — + same idea as the `9.99` operating points. When a note would have fired + it carries that sentinel plus the model-head suffix, not the internal + appeal paragraph or the action types each head keys off. ActionName + proto enums are unchanged. diff --git a/bdsm/proto/abuse_inference.proto b/bdsm/proto/abuse_inference.proto index 7c52d3e9..62937bb1 100644 --- a/bdsm/proto/abuse_inference.proto +++ b/bdsm/proto/abuse_inference.proto @@ -98,6 +98,9 @@ message ScoreResult { repeated DecodedAction decoded_actions = 18; + // Optional human-readable note. In this release templates and + // key_actions are the sentinel "" (see README); production + // fills the real appeal paragraph. ActionName enums are unchanged. string enforcement_note = 19; } diff --git a/bdsm/runtime/score_results_sink_focal.py b/bdsm/runtime/score_results_sink_focal.py index 04f44134..520807d6 100644 --- a/bdsm/runtime/score_results_sink_focal.py +++ b/bdsm/runtime/score_results_sink_focal.py @@ -78,8 +78,57 @@ def _resolve_head_names(_model_version): return HEAD_NAMES -def build_enforcement_note(*_args, **_kwargs): - return None +# Public-release note: per-head appeal-note templates are REDACTED +# (the prose AND the key_actions each head interpolates from the +# histogram). What ships: MIN_ACTIONS gate, dominant bot-head pick, +# and the sentinel "" plus a [model: Head=score] suffix. +# Production fills placeholders from an internal table that is not +# part of this release. ActionName proto enums are untouched. +_REDACTED_TEMPLATE = "" + +_ENFORCEMENT_TEMPLATES = { + "FollowBot": _REDACTED_TEMPLATE, + "LikeBot": _REDACTED_TEMPLATE, + "EngagementAmplifier": _REDACTED_TEMPLATE, + "ReplySpamBot": _REDACTED_TEMPLATE, + "TweetSpamBot": _REDACTED_TEMPLATE, + "RTBot": _REDACTED_TEMPLATE, + "MultiActionBot": _REDACTED_TEMPLATE, +} + + +def build_enforcement_note(head_scores_list, action_hist_list, min_actions): + """Build an enforcement note from head scores + action histogram. + + Returns None when no bot head is above 0.5 or the sequence is shorter + than `min_actions` (the policy min-actions gate). In this public release + the note is the sentinel "" plus a model-head suffix; production + interpolates from a private template table (prose and key_actions). + """ + if not head_scores_list: + return None + + total = sum(h["cnt"] for h in (action_hist_list or ())) + if total < min_actions: + return None + + bot_heads = [ + (h["head_name"], h["score"]) + for h in head_scores_list + if h["head_name"] != "LegitimateUser" and h["score"] > 0.5 + ] + if not bot_heads: + return None + + bot_heads.sort(key=lambda x: x[1], reverse=True) + dominant_name, dominant_score = bot_heads[0] + note = _ENFORCEMENT_TEMPLATES.get(dominant_name, _REDACTED_TEMPLATE) + note += f" [model: {dominant_name}={dominant_score:.2f}" + if len(bot_heads) > 1: + secondary = ", ".join(f"{n}={s:.2f}" for n, s in bot_heads[1:3]) + note += f", also: {secondary}" + note += "]" + return note def _build_flag_config(head_names): @@ -171,7 +220,7 @@ def client_dwell_dropout(r, uid: int, args, metrics) -> bool: class SinkPolicy: version: str = "baked-in-defaults" source: str = "defaults" - min_actions_for_enforcement: int = 30 + min_actions_for_enforcement: int = 999999 thresholds: dict = field( default_factory=lambda: { "FollowBot": (9.99, 9.99), @@ -309,9 +358,9 @@ def _parse_args(): parser.add_argument( "--min-actions-for-long-cooldown", type=int, - default=30, + default=None, help="Action-count threshold separating short vs long cooldown buckets. " - "Default matches MIN_ACTIONS_FOR_ENFORCEMENT (30).", + "Defaults to the policy min-actions enforcement gate.", ) parser.add_argument( "--bq-project", default="your-gcp-project", help="GCP project for the scores table." @@ -1089,7 +1138,15 @@ def _spam_bounce_lane( def _build_bq_row( - uid_int, score_id, now_ts, user_head_scores, user_action_hist, user_labels, bsummary, args + uid_int, + score_id, + now_ts, + user_head_scores, + user_action_hist, + user_labels, + bsummary, + args, + min_note_actions, ): total_actions = None if user_action_hist: @@ -1116,7 +1173,9 @@ def _build_bq_row( "head_scores": user_head_scores, "action_histogram": user_action_hist, "total_actions": total_actions, - "enforcement_note": build_enforcement_note(), + "enforcement_note": build_enforcement_note( + user_head_scores, user_action_hist, min_note_actions + ), "labels": user_labels, "model_version": args.model_version, "pipeline_version": "gpu_scorer_kafka_v3", @@ -1442,7 +1501,11 @@ def main(): r = redis_lib.Redis(host=args.redis_host, port=args.redis_port) cooldown_low_sec = int(args.cooldown_low_hours * 3600) cooldown_high_sec = int(args.cooldown_high_hours * 3600) - cooldown_threshold = args.min_actions_for_long_cooldown + cooldown_threshold = ( + args.min_actions_for_long_cooldown + if args.min_actions_for_long_cooldown is not None + else pol.min_actions_for_enforcement + ) try: r.ping() log.info( @@ -1588,6 +1651,7 @@ def main(): user_labels, bsummary, args, + pol.min_actions_for_enforcement, ) bq_buffer.append(row) diff --git a/bdsm/runtime/sink_policy.yaml b/bdsm/runtime/sink_policy.yaml index eda36890..4ec9eaea 100644 --- a/bdsm/runtime/sink_policy.yaml +++ b/bdsm/runtime/sink_policy.yaml @@ -1,9 +1,18 @@ -version: "2026-08-13" +version: "2026-08-14" -# Public-release note: the per-head operating points below (the 2-D +# Public-release note: the operating points below (the 2-D +# [tau, lambda] pairs, cusp_delta, reply_spam_hard_suspend_tau, and +# the min_actions_for_enforcement count gate) are REDACTED in the +# open-source export. Probabilities are replaced with an out-of-range +# 9.99 sentinel (these fields live in [0, 1], so a 9.99 threshold +# never fires and is plainly a stub, not a real value); the action +# count with an impossible 999999 (far longer than any scoreable +# sequence, so the gate can never pass). The production values are +# configured internally and are not part of this release. Structure, +# head names, and the action map are real. -min_actions_for_enforcement: 30 +min_actions_for_enforcement: 999999 thresholds: FollowBot: [9.99, 9.99] @@ -31,5 +40,6 @@ official_client_app_ids: [3033300, 129032, 258901, 191841] reply_spam_hard_suspend_tau: 9.99 notes: > - 2026-08-13: enforcement policy for the 8-head behavioral model. - Operating points are redacted in this export (9.99 sentinel). + 2026-08-14: enforcement policy for the 8-head behavioral model. + Operating points are redacted in this export (out-of-range sentinels: + 9.99 for probabilities, 999999 for the min-actions count gate). diff --git a/bdsm/tests/test_sink_policy.py b/bdsm/tests/test_sink_policy.py index 89db1c5c..0b09c544 100644 --- a/bdsm/tests/test_sink_policy.py +++ b/bdsm/tests/test_sink_policy.py @@ -30,6 +30,7 @@ def test_shipped_policy_matches_baked_in_defaults(): got.pop(k) want.pop(k) assert got == want, "sink_policy.yaml drifted from the baked-in defaults" + assert pol.min_actions_for_enforcement > 10_000 def test_missing_policy_file_falls_back_to_defaults(): @@ -43,3 +44,45 @@ def test_invalid_policy_key_fails_loudly(tmp_path): bad.write_text("version: x\nnot_a_policy_key: 1\n") with pytest.raises(ValueError): m._load_policy(str(bad)) + + +def _hist(*pairs): + return [{"action_type": a, "cnt": n} for a, n in pairs] + + +def _heads(**scores): + return [ + {"head_index": i, "head_name": n, "score": s} for i, (n, s) in enumerate(scores.items()) + ] + + +def test_enforcement_note_gates_and_redacted_prose(): + m = _load_sink_module() + gate = 25 + heads = _heads(FollowBot=0.91, LegitimateUser=0.1) + short = _hist(("SERVER_PROFILE_FOLLOW", 10)) + long = _hist(("SERVER_PROFILE_FOLLOW", 40), ("SERVER_PROFILE_UNFOLLOW", 5)) + + assert m.build_enforcement_note(heads, short, gate) is None + assert m.build_enforcement_note(heads, None, gate) is None + assert m.build_enforcement_note(heads, [], gate) is None + assert m.build_enforcement_note(_heads(FollowBot=0.4, LegitimateUser=0.8), long, gate) is None + + note = m.build_enforcement_note(heads, long, gate) + assert note == f"{m._REDACTED_TEMPLATE} [model: FollowBot=0.91]" + + +def test_enforcement_templates_are_the_redacted_sentinel(): + m = _load_sink_module() + assert set(m._ENFORCEMENT_TEMPLATES) == { + "FollowBot", + "LikeBot", + "EngagementAmplifier", + "ReplySpamBot", + "TweetSpamBot", + "RTBot", + "MultiActionBot", + } + for tmpl in m._ENFORCEMENT_TEMPLATES.values(): + assert tmpl == m._REDACTED_TEMPLATE + assert not hasattr(m, "_ACTION_FRIENDLY") diff --git a/grox/config/config.py b/grox/config/config.py index 87ad75da..8a030c81 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -86,9 +86,11 @@ class ModelName: GROK_4_MINI_CRITICAL_SAFETY = "critical-safety" EAPI_GROK_420_REASONING_X_ALGO = "eapi-grok-420-reasoning-x-algo" EAPI_GROK_420_REASONING_INTERNAL = "eapi-grok-420-reasoning-internal" + EAPI_GROK_4_1_FAST_X_ALGO = "eapi-grok-4-1-fast-x-algo" EAPI_GROK_4_3_INTERNAL = "eapi-grok-4-3-internal" EAPI_GROK_4_3_X_ALGO = "eapi-grok-4-3-x-algo" - EAPI_GROK_4_5_INTERNAL = "eapi-grok-4-5-internal" + EAPI_GROK_4_5_X_ALGO = "eapi-grok-4-5-x-algo" + EAPI_GROK_4_6_INTERNAL = "eapi-grok-4-6-internal" class NightOwlConfig(BaseModel): @@ -109,6 +111,7 @@ class MediaHydrationConfig(BaseModel): image_tile_size: int = 448 enable_light_dark_enhancement: bool = False enable_clahe_enhancement: bool = False + enable_motion_reveal: bool = False deluxe_fav_count_threshold: int = 64 deluxe_video_max_frames: int = 30 deluxe_video_tile_size: int = 600 @@ -120,6 +123,7 @@ class GroxKafkaLoaderConfig(BaseModel): prefetching_threshold: int = 256 prefetching_batch_size: int = 1024 + max_qps_per_partition: int | None = None class GrpcServerConfig(BaseModel): diff --git a/grox/core/data_loaders/data_types.py b/grox/core/data_loaders/data_types.py index a441fb43..893a2247 100644 --- a/grox/core/data_loaders/data_types.py +++ b/grox/core/data_loaders/data_types.py @@ -39,6 +39,7 @@ class User(BaseModel): urls: list[str] | None = None affiliated_business: AffiliatedBusiness | None = None recent_posts: list["Post"] | None = None + profile_image: "Image | None" = None @classmethod def from_thrift_model(cls, author_metadata: t.AuthorMetadata) -> "User": @@ -86,6 +87,9 @@ def from_thrift_model(cls, author_metadata: t.AuthorMetadata) -> "User": ) if author_metadata.affiliatedBusinessMetadata else None, + profile_image=Image(url=author_metadata.profileImageUrl) + if getattr(author_metadata, "profileImageUrl", None) + else None, ) @@ -367,30 +371,47 @@ def from_thrift_model( ) def to_convo(self) -> list[str | ConvoImage]: - res: list[str | ConvoImage] = ["\n\n[Card] ", " "] + body: list[str | ConvoImage] = [] if self.title: - res.append(f"\n\nTitle: {self.title}") + body.append(f"\n\n[Card Title] {self.title}") if self.description: - res.append(f"\n\nDescription: {self.description}") + body.append(f"\n\nDescription: {self.description}") if self.domain: - res.append(f"\n\nDomain: {self.domain}") + body.append(f"\n\nDomain: {self.domain}") if self.thumbnail_image and self.thumbnail_image.convo_image: - res.extend(["\n\n[Card Image] ", self.thumbnail_image.convo_image, " "]) + body.extend(["\n\n[Card Image] ", self.thumbnail_image.convo_image, " "]) if self.poll_cards: - res.append( + body.append( "\n\nThis post includes a poll where user can vote their choices. The choices are as rendered below:" ) for poll_card in self.poll_cards: - res.extend(poll_card.to_convo()) - if self.grok_share_cards: - share_parts: list[str | ConvoImage] = [] - for grok_share_card in self.grok_share_cards: - share_parts.extend(grok_share_card.to_convo()) - if share_parts: - res.append( - "\n\nThis post includes a shared Grok conversation. The conversation messages are as rendered below:" - ) - res.extend(share_parts) + body.extend(poll_card.to_convo()) + if not body: + return [] + return ["\n\n[Card] ", " "] + body + + +class GrokShare(BaseModel): + sender: str | None = None + message: str | None = None + + @classmethod + def from_thrift_model(cls, metadata: t.GrokShareMetadata) -> "GrokShare": + sender = ( + t.GrokShareConversationSender._VALUES_TO_NAMES.get(metadata.sender) + if metadata.sender is not None + else None + ) + return cls(sender=sender, message=metadata.message) + + def to_convo(self) -> list[str | ConvoImage]: + if not self.sender and not self.message: + return [] + res: list[str | ConvoImage] = ["\n\n[Grok Share Message] ", " "] + if self.sender: + res.append(f"\n\nSender: {self.sender}") + if self.message: + res.append(f"\n\nMessage: {self.message}") return res @@ -424,7 +445,7 @@ def from_thrift_model( def to_convo(self) -> list[str | ConvoImage | ConvoVideo]: res: list[str | ConvoImage | ConvoVideo] = ["\n\n[Card] ", " "] if self.title: - res.append(f"\n\nTitle: {self.title}") + res.append(f"\n\n[Card Title] {self.title}") if self.description: res.append(f"\n\nDescription: {self.description}") if self.url: @@ -578,6 +599,22 @@ def to_convo(self) -> list[str | ConvoImage]: return res +class SpaceMetadata(BaseModel): + title: str | None = None + + @classmethod + def from_thrift_model(cls, metadata: t.SpaceMetadata) -> "SpaceMetadata": + return cls(title=metadata.title) + + def to_convo(self) -> list[str]: + if not self.title: + return [] + return [ + "\n\nThis post has the following audio space card attached:", + f"\nAudio Space title: {self.title}", + ] + + class ArticleMetadata(BaseModel): title: str | None = None media: list[Image | Video] | None = None @@ -648,6 +685,7 @@ class Post(BaseModel): screenshot: Image | None = None reply: Reply | None = None cardsV2: list[CardV2] | None = None + grok_share_metadatas: list[GrokShare] | None = None article_metadata: ArticleMetadata | None = None descendants: list["Post"] | None = None user_agent: str | None = None @@ -657,6 +695,7 @@ class Post(BaseModel): safety_labels: list[str] | None = None list_metadata: ListMetadata | None = None chat_group_metadata: ChatGroupMetadata | None = None + space_metadata: SpaceMetadata | None = None def get_images(self) -> list[Image]: images: list[Image] = [] @@ -744,6 +783,12 @@ def from_post_metadata(cls, post_metadata: t.PostMetadataV2) -> "Post": ] else: cardsV2 = None + if post_metadata.grokShareMetadatas: + grok_share_metadatas = [ + GrokShare.from_thrift_model(m) for m in post_metadata.grokShareMetadatas + ] + else: + grok_share_metadatas = None if post_metadata.articleMetadata: article_metadata = ArticleMetadata.from_thrift_model( post_metadata.articleMetadata @@ -760,6 +805,12 @@ def from_post_metadata(cls, post_metadata: t.PostMetadataV2) -> "Post": ) else: chat_group_metadata = None + if post_metadata.spaceMetadata: + space_metadata = SpaceMetadata.from_thrift_model( + post_metadata.spaceMetadata + ) + else: + space_metadata = None return cls( id=str(post_metadata.postId), @@ -778,6 +829,7 @@ def from_post_metadata(cls, post_metadata: t.PostMetadataV2) -> "Post": ancestors=[], screenshot=None, cardsV2=cardsV2, + grok_share_metadatas=grok_share_metadatas, article_metadata=article_metadata, user_agent=post_metadata.userAgent if hasattr(post_metadata, "userAgent") @@ -793,6 +845,7 @@ def from_post_metadata(cls, post_metadata: t.PostMetadataV2) -> "Post": else None, list_metadata=list_metadata, chat_group_metadata=chat_group_metadata, + space_metadata=space_metadata, ) @classmethod diff --git a/grox/core/data_loaders/kafka_loader.py b/grox/core/data_loaders/kafka_loader.py index 60ed5276..f04a451b 100644 --- a/grox/core/data_loaders/kafka_loader.py +++ b/grox/core/data_loaders/kafka_loader.py @@ -1,4 +1,3 @@ -import os import struct import uuid import asyncio @@ -13,6 +12,7 @@ from grox.config.config import grox_config from kafka_cli.consumer import KafkaConsumer from kafka_cli.multi_region_consumer import MultiRegionKafkaConsumer +from limits import RateLimitItemPerSecond, storage, strategies from grox.core.data_loaders.data_types import Post from grox.core.data_loaders.message_queue_loader import ( MessageQueueLoader, @@ -27,6 +27,7 @@ logger = logging.getLogger(__name__) MAX_WORKING_THREADS = 12 +_limiter = strategies.FixedWindowRateLimiter(storage.MemoryStorage()) def parse_native_encoding(data: bytes) -> int: @@ -51,24 +52,11 @@ def __init__(self, topic_name: str): self.consumer = MultiRegionKafkaConsumer(self.consumer_config) else: self.consumer_config = grox_config.get_kafka_consumer_topic(topic_name) - self._maybe_inject_scram_password(self.consumer_config) self.consumer = KafkaConsumer(self.consumer_config) self.queue: asyncio.Queue[MessageQueuePayload] = asyncio.Queue() self._prefetcher_task: asyncio.Task | None = None - - @staticmethod - def _maybe_inject_scram_password(consumer_config) -> None: - ssl_conf = consumer_config.ssl - if ssl_conf is None or not ssl_conf.sasl_mechanism.startswith("SCRAM"): - return - if ssl_conf.sasl_plain_password: - return - password = os.environ.get("KAFKA_RECSYS_PASSWORD") - if not password: - raise RuntimeError( - "KAFKA_RECSYS_PASSWORD env var is required for SCRAM auth but is not set." - ) - ssl_conf.sasl_plain_password = password + self._max_qps_per_partition = self.loader_config.max_qps_per_partition + self._limit_item: RateLimitItemPerSecond | None = None def _is_shutdown(self) -> bool: try: @@ -100,20 +88,44 @@ async def stop(self): await self.consumer.stop() logger.warning(f"KafkaLoader stopped, topic: {self.topic_name}") + def _current_limit(self) -> RateLimitItemPerSecond | None: + if not self._max_qps_per_partition: + return None + partitions = self.consumer.assigned_partitions + if partitions <= 0: + return self._limit_item + amount = self._max_qps_per_partition * partitions + if self._limit_item is None or self._limit_item.amount != amount: + self._limit_item = RateLimitItemPerSecond(amount, 1) + return self._limit_item + async def poll(self) -> AsyncGenerator[MessageQueuePayload | None, None]: while not self._shutdown_event.is_set() or not self.queue.empty(): + limit = self._current_limit() + if ( + limit + and not self._shutdown_event.is_set() + and not _limiter.test(limit, self.topic_name) + ): + yield None + continue try: - yield self.queue.get_nowait() + payload = self.queue.get_nowait() except asyncio.QueueEmpty: logger.debug( f"Queue is empty, waiting for prefetcher to fill, topic: {self.topic_name}" ) yield None + continue except Exception: logger.error( f"Error polling from kafka, topic: {self.topic_name}, error: {traceback.format_exc()}" ) yield None + continue + if limit: + _limiter.hit(limit, self.topic_name) + yield payload async def ack(self, mid: str, success: bool = True): pass diff --git a/grox/core/data_loaders/media_loader.py b/grox/core/data_loaders/media_loader.py index e954aef6..9c75d986 100644 --- a/grox/core/data_loaders/media_loader.py +++ b/grox/core/data_loaders/media_loader.py @@ -227,6 +227,11 @@ async def hydrate_media( grox_config.media_hydration.enable_light_dark_enhancement and is_main_post ) + should_enable_motion_reveal = ( + grox_config.media_hydration.enable_motion_reveal + and is_main_post + and is_high_fav + ) for medium in post.media: if isinstance(medium, Image): tasks.append( @@ -244,6 +249,7 @@ async def hydrate_media( is_main_post, is_high_fav, should_enable_clahe_enhancement, + should_enable_motion_reveal, ) ) if post.broadcast_metadata and post.broadcast_metadata.thumbnail_image: @@ -377,6 +383,7 @@ async def hydrate_video( is_main_post: bool = False, is_high_fav: bool = False, enable_clahe_enhancement: bool = False, + enable_motion_reveal: bool = False, ) -> None: url = None if video.videoInfo and video.videoInfo.durationMillis: @@ -458,6 +465,7 @@ async def hydrate_video( is_main_post, is_high_fav, enable_clahe_enhancement, + enable_motion_reveal, ) Metrics.counter("media_loader.hydrate_video_success.count").add( 1, attributes=cls._metrics_attributes() @@ -527,6 +535,7 @@ async def construct_convo_video( is_main_post: bool = False, is_high_fav: bool = False, enable_clahe_enhancement: bool = False, + enable_motion_reveal: bool = False, ) -> ConvoVideo: video_max_frames = grox_config.media_hydration.video_max_frames_light video_tile_size = grox_config.media_hydration.video_tile_size @@ -545,6 +554,7 @@ async def construct_convo_video( video_tile_size, enable_clahe=enable_clahe_enhancement, include_combined_video_bytes=False, + enable_motion_reveal=enable_motion_reveal, ) times = [frame.time_sec for frame in video_data.frames] frames = [frame.frame for frame in video_data.frames] @@ -574,10 +584,20 @@ async def construct_convo_video( 1, attributes=cls._metrics_attributes() ) + if video_data.motion_reveal_frames: + Metrics.counter("media_loader.motion_reveal_built.count").add( + 1, attributes=cls._metrics_attributes() + ) + elif enable_motion_reveal: + Metrics.counter("media_loader.motion_reveal_skipped.count").add( + 1, attributes=cls._metrics_attributes() + ) + return ConvoVideo( frames=frames, subtitles=subtitles, duration=duration, total_duration=total_duration, is_deluxe_target=is_main_post and is_high_fav, + motion_reveal_frames=video_data.motion_reveal_frames, ) diff --git a/grox/core/data_loaders/post_mapper.py b/grox/core/data_loaders/post_mapper.py index ff432bf6..aef93d11 100644 --- a/grox/core/data_loaders/post_mapper.py +++ b/grox/core/data_loaders/post_mapper.py @@ -11,9 +11,11 @@ CardMetadataV2 as StratoCardMetadataV2, PollCardMetadata as StratoPollCardMetadata, GrokShareCardMetadata as StratoGrokShareCardMetadata, + GrokShareMetadata as StratoGrokShareMetadata, ArticleMetadata as StratoArticleMetadata, ListMetadata as StratoListMetadata, ChatGroupMetadata as StratoChatGroupMetadata, + SpaceMetadata as StratoSpaceMetadata, ) from grox.core.data_loaders.data_types import ( Post, @@ -30,9 +32,11 @@ BroadcastMetadata, PollCard, GrokShareCard, + GrokShare, ArticleMetadata, ListMetadata, ChatGroupMetadata, + SpaceMetadata, AffiliatedBusiness, ) @@ -117,6 +121,12 @@ def from_post_metadata_strato(cls, post_metadata: StratoPostMetadata) -> Post: cls._from_strato_cardmetadataV2_to_cardV2(cardV2) for cardV2 in post_metadata.cardMetadatasV2 ] + grok_share_metadatas = None + if post_metadata.grokShareMetadatas: + grok_share_metadatas = [ + cls._from_strato_grok_share_metadata(m) + for m in post_metadata.grokShareMetadatas + ] article_metadata = None if post_metadata.articleMetadata: article_metadata = cls._from_strato_article_metadata_to_article_metadata( @@ -134,6 +144,11 @@ def from_post_metadata_strato(cls, post_metadata: StratoPostMetadata) -> Post: post_metadata.chatGroupMetadata ) ) + space_metadata = None + if post_metadata.spaceMetadata: + space_metadata = cls._from_strato_space_metadata_to_space_metadata( + post_metadata.spaceMetadata + ) return Post( id=str(post_metadata.postId), user=user, @@ -151,9 +166,11 @@ def from_post_metadata_strato(cls, post_metadata: StratoPostMetadata) -> Post: ancestors=[], screenshot=None, cardsV2=cardsV2, + grok_share_metadatas=grok_share_metadatas, article_metadata=article_metadata, list_metadata=list_metadata, chat_group_metadata=chat_group_metadata, + space_metadata=space_metadata, ) @classmethod @@ -286,6 +303,12 @@ def _from_strato_chat_group_metadata_to_chat_group_metadata( else None, ) + @classmethod + def _from_strato_space_metadata_to_space_metadata( + cls, metadata: StratoSpaceMetadata + ) -> SpaceMetadata: + return SpaceMetadata(title=metadata.title) + @classmethod def _from_strato_poll_card_metadata_to_poll_card( cls, poll_card_metadata: StratoPollCardMetadata @@ -303,6 +326,13 @@ def _from_strato_grok_share_card_metadata_to_grok_share_card( ) -> GrokShareCard: return GrokShareCard(sender=metadata.sender, message=metadata.message) + @classmethod + def _from_strato_grok_share_metadata( + cls, metadata: StratoGrokShareMetadata + ) -> GrokShare: + sender = metadata.sender.name if metadata.sender is not None else None + return GrokShare(sender=sender, message=metadata.message) + @classmethod def _from_strato_user_metadata_to_user( cls, user_metadata: StratoUserMetadata @@ -331,6 +361,9 @@ def _from_strato_user_metadata_to_user( ) if user_metadata.affiliatedBusinessMetadata else None, + profile_image=Image(url=user_metadata.profileImageUrl) + if user_metadata.profileImageUrl + else None, ) @classmethod diff --git a/grox/core/lm/convo.py b/grox/core/lm/convo.py index 1f8a05f3..43b47201 100644 --- a/grox/core/lm/convo.py +++ b/grox/core/lm/convo.py @@ -20,6 +20,13 @@ STORYBOARD_TILE_SIZE = 448 NO_THINKING_PROMPT = grox_config.prompt_tokens.no_thinking_prompt +MOTION_REVEAL_DESCRIPTION = ( + "Motion-reveal stills for this video follow (static layer subtracted, faint moving layers amplified); " + "they distort colors and motion, and can make opaque clothing look like bare skin. People or acts " + "absent from the frames above but visible here are the video's actual overlay-hidden content; judge " + "anyone already visible above solely from the original frames." +) + class Role(str, Enum): USER = "User" @@ -50,12 +57,13 @@ class Video(BaseModel): duration: float total_duration: float is_deluxe_target: bool = False + motion_reveal_frames: list[bytes] = Field(default_factory=lambda: []) - @field_serializer("frames", when_used="json") + @field_serializer("frames", "motion_reveal_frames", when_used="json") def serialize_frames(self, value: list[bytes]) -> list[str]: return [b64encode(frame).decode("utf-8") for frame in value] - @field_validator("frames", mode="before") + @field_validator("frames", "motion_reveal_frames", mode="before") @classmethod def decode_frames(cls, value: list[str] | list[bytes]) -> list[bytes]: is_list_str = len(value) > 0 and isinstance(value[0], str) @@ -155,6 +163,11 @@ def interleave_storyboard( res.append("\n") + if self.motion_reveal_frames: + res.append(f"{MOTION_REVEAL_DESCRIPTION}\n") + for idx, reveal_bytes in enumerate(self.motion_reveal_frames): + res.extend([" ", reveal_bytes, f" Motion-reveal still {idx + 1}\n"]) + return res @@ -371,6 +384,18 @@ def _image_part(image_bytes: bytes) -> dict: } ) parts.append(_image_part(frame)) + if c.motion_reveal_frames: + parts.append( + {"type": "text", "text": MOTION_REVEAL_DESCRIPTION} + ) + for idx, reveal in enumerate(c.motion_reveal_frames): + parts.append( + { + "type": "text", + "text": f"Motion-reveal still {idx + 1}", + } + ) + parts.append(_image_part(reveal)) if not parts: continue diff --git a/grox/core/lm/post.py b/grox/core/lm/post.py index 4df57298..1aa51b6a 100644 --- a/grox/core/lm/post.py +++ b/grox/core/lm/post.py @@ -40,12 +40,19 @@ def render( max_media: int | None = None, include_reply_to: bool = False, cards_note_override: str | None = None, + include_follower_count: bool = False, + include_bio: bool = False, ) -> list[Content]: indent_str = ">" + (" " * indent) if indent > 0 else "" res = [] if not post.user: post.user = User(name="unknown", handle="unknown") - res.append(f"\n{indent_str}Metadata: {post.user.name} @{post.user.handle}") + res.append(f"\n{indent_str}User Handle: @{post.user.handle}") + if include_follower_count and post.user.follower_count is not None: + res.append(f"\n{indent_str}User Follower Count: {post.user.follower_count}") + res.append(f'\n{indent_str}User Name: "{post.user.name}"') + if include_bio and post.user.bio: + res.append(f'\n{indent_str}User Bio: "{post.user.bio}"') all_media = list(post.media or []) + list(post.url_videos or []) if all_media: res.append(f"\n{indent_str}Media:") @@ -67,10 +74,9 @@ def render( reply_handles = get_replies_handle_string(post) res.append(f"\n{indent_str}Replying to: {reply_handles}") res.append(f"\n{indent_str}Text: {formatted_text}") - if post.urls: - res.append( - f"\n{indent_str}The Post contains these URLs: {', '.join(post.urls)}" - ) + urls = [url for url in post.urls or [] if url] + if urls: + res.append(f"\n{indent_str}The Post contains these URLs: {', '.join(urls)}") if post.broadcast_metadata: res.extend(post.broadcast_metadata.to_convo()) if cards_note_override is not None: @@ -79,6 +85,12 @@ def render( res.append(f"\n{indent_str}This post has the following cards attached: ") for cardV2 in post.cardsV2: res.extend(cardV2.to_convo()) + if post.grok_share_metadatas: + res.append( + f"\n{indent_str}This post includes a shared Grok conversation. The conversation messages are as rendered below:" + ) + for grok_share in post.grok_share_metadatas: + res.extend(grok_share.to_convo()) if post.article_metadata: res.append(f"\n{indent_str}[Article Post] This post is an Article.") res.extend(post.article_metadata.to_convo()) @@ -86,12 +98,20 @@ def render( res.extend(post.list_metadata.to_convo()) if post.chat_group_metadata: res.extend(post.chat_group_metadata.to_convo()) + if post.space_metadata: + res.extend(post.space_metadata.to_convo()) if post.quoted_post: res.append( f"\n\n{indent_str}This Post quotes Post {post.quoted_post.id}\n\n" ) res.extend( - cls.render(post.quoted_post, indent=indent + 4, max_media=max_media) + cls.render( + post.quoted_post, + indent=indent + 4, + max_media=max_media, + include_follower_count=include_follower_count, + include_bio=include_bio, + ) ) if post.descendants: res.append( diff --git a/grox/core/lm/thread.py b/grox/core/lm/thread.py index 12999465..b5b68092 100644 --- a/grox/core/lm/thread.py +++ b/grox/core/lm/thread.py @@ -18,12 +18,18 @@ def render( role: Role = Role.USER, separator=SEPARATOR, include_signals: bool = False, + include_follower_count: bool = False, ) -> Message: message = Message(role=role, content=[], separator=separator) + message.content.append( + "\n# Reply Author Info\n\nThe user profile and any signals below describe the author of the reply being evaluated (the final post of the thread).\n" + ) message.content.extend(UserRenderer.render(post.user)) if include_signals: message.content.extend(cls._render_signals(post)) - message.content.extend(cls._render_thread(post)) + message.content.extend( + cls._render_thread(post, include_follower_count=include_follower_count) + ) return message @classmethod @@ -60,7 +66,9 @@ def _render_signals(cls, post: Post) -> list[Content]: return ["\n# Additional Signals\n" + "\n".join(lines) + "\n"] @classmethod - def _render_thread(cls, post: Post) -> list[Content]: + def _render_thread( + cls, post: Post, include_follower_count: bool = False + ) -> list[Content]: res = [] res.append( "\n\n# Thread \n\nListed below is the X post thread before the reply. Post 0 is the original post. \n\n" @@ -76,11 +84,24 @@ def _render_thread(cls, post: Post) -> list[Content]: if p is None: res.append("\n<- A post has been deleted ->\n") else: - res.extend(PostRenderer.render(p, max_media=MAX_MEDIA_PER_POST)) + res.extend( + PostRenderer.render( + p, + max_media=MAX_MEDIA_PER_POST, + include_follower_count=include_follower_count, + include_bio=include_follower_count and post_idx == 0, + ) + ) res.append("\n\n------\n\n") res.append("\n\n# Reply \n\nBelow is the reply that you need to evaluate. \n\n") - res.extend(PostRenderer.render(post, max_media=MAX_MEDIA_PER_POST)) + res.extend( + PostRenderer.render( + post, + max_media=MAX_MEDIA_PER_POST, + include_follower_count=include_follower_count, + ) + ) res.append("\n\n------\n\n") return res diff --git a/grox/flows/mm_emb/task_embedding_pub.py b/grox/flows/mm_emb/task_embedding_pub.py index 19aaf1d6..ecd9bf5a 100644 --- a/grox/flows/mm_emb/task_embedding_pub.py +++ b/grox/flows/mm_emb/task_embedding_pub.py @@ -1,14 +1,12 @@ +import asyncio import logging -import os -from functools import cache - from thrifts.gen.twitter.strato.columns.content_understanding.content_understanding.ttypes import ( SimpleTweetEmbedding, ) from thrifts.serdes import Serializer from grox.core.data_loaders.data_types import Post from grox.config.config import grox_config -from kafka_cli.producer import ScramKafkaProducer +from kafka_cli.multi_region_producer import MultiRegionKafkaProducer from grox.flows.mm_emb.constants import ( TOPIC_EMBEDDING_V5, TOPIC_EMBEDDING_V5_ALL, @@ -18,39 +16,44 @@ logger = logging.getLogger(__name__) -class TaskPublishEmbeddingKafka: +def _serialize(post: Post, embedding: list[float]) -> bytes: + return Serializer.serialize( + SimpleTweetEmbedding(tweetId=int(post.id), embedding1=embedding) + ) + + +class TaskPublishEmbeddingMultiRegionKafka: KAFKA_TOPIC_NAME: str + _producer: MultiRegionKafkaProducer | None = None + _producer_lock = asyncio.Lock() + @classmethod async def _publish_to_kafka(cls, post: Post, embedding: list[float]) -> None: - tweet_embedding = SimpleTweetEmbedding( - tweetId=int(post.id), - embedding1=embedding, - ) - serialized_bytes = Serializer.serialize(tweet_embedding) - await cls._get_kafka_producer().send(id=post.id, value=serialized_bytes) + producer = await cls._get_kafka_producer() + await producer.send(id=post.id, value=_serialize(post, embedding)) logger.info(f"Published embedding for post {post.id} to {cls.KAFKA_TOPIC_NAME}") @classmethod - @cache - def _get_kafka_producer(cls) -> ScramKafkaProducer: - producer_config = grox_config.get_kafka_producer_topic(cls.KAFKA_TOPIC_NAME) - password = os.environ.get("KAFKA_RECSYS_PASSWORD") - if not password: - raise RuntimeError( - "KAFKA_RECSYS_PASSWORD env var is required for SCRAM auth but is not set." - ) - producer_config.ssl.sasl_plain_password = password - return ScramKafkaProducer(producer_config) - - -class TaskPublishEmbeddingV5Kafka(TaskPublishEmbeddingKafka): + async def _get_kafka_producer(cls) -> MultiRegionKafkaProducer: + if cls._producer is None: + async with cls._producer_lock: + if cls._producer is None: + producer = MultiRegionKafkaProducer( + grox_config.get_kafka_producer(cls.KAFKA_TOPIC_NAME) + ) + await producer.start() + cls._producer = producer + return cls._producer + + +class TaskPublishEmbeddingV5Kafka(TaskPublishEmbeddingMultiRegionKafka): KAFKA_TOPIC_NAME = TOPIC_EMBEDDING_V5 -class TaskPublishEmbeddingV5AllKafka(TaskPublishEmbeddingKafka): +class TaskPublishEmbeddingV5AllKafka(TaskPublishEmbeddingMultiRegionKafka): KAFKA_TOPIC_NAME = TOPIC_EMBEDDING_V5_ALL -class TaskPublishEmbeddingV82Kafka(TaskPublishEmbeddingKafka): +class TaskPublishEmbeddingV82Kafka(TaskPublishEmbeddingMultiRegionKafka): KAFKA_TOPIC_NAME = TOPIC_EMBEDDING_V8_2 diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index 46805c11..554d5964 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -74,7 +74,15 @@ def _fav_bucket(fav_count: int) -> str: half_open_max_calls=5, excluded_exceptions=(asyncio.CancelledError,), ) -_EAPI_4_5_INTERNAL_BREAKER_CONFIG = CircuitBreakerConfig( +_EAPI_4_6_INTERNAL_BREAKER_CONFIG = CircuitBreakerConfig( + failure_rate_threshold=0.5, + window_size=600.0, + min_calls_in_window=10, + recovery_timeout=600.0, + half_open_max_calls=5, + excluded_exceptions=(asyncio.CancelledError,), +) +_EAPI_4_5_X_ALGO_BREAKER_CONFIG = CircuitBreakerConfig( failure_rate_threshold=0.5, window_size=600.0, min_calls_in_window=10, @@ -85,8 +93,11 @@ def _fav_bucket(fav_count: int) -> str: _eapi_4_3_x_algo_breaker = CircuitBreaker( ModelName.EAPI_GROK_4_3_X_ALGO, _EAPI_4_3_X_ALGO_BREAKER_CONFIG ) -_eapi_4_5_internal_breaker = CircuitBreaker( - ModelName.EAPI_GROK_4_5_INTERNAL, _EAPI_4_5_INTERNAL_BREAKER_CONFIG +_eapi_4_5_x_algo_breaker = CircuitBreaker( + ModelName.EAPI_GROK_4_5_X_ALGO, _EAPI_4_5_X_ALGO_BREAKER_CONFIG +) +_eapi_4_6_internal_breaker = CircuitBreaker( + ModelName.EAPI_GROK_4_6_INTERNAL, _EAPI_4_6_INTERNAL_BREAKER_CONFIG ) @@ -171,13 +182,232 @@ async def _sample(self, convo: Conversation) -> str: ) +def _policy_no_violation(reason: str) -> SafetyPolicy: + return SafetyPolicy(policyType=SafetyPolicyType.NoViolation, reason=reason) + + +def _policy_with_type( + policy_type: SafetyPolicyType, source: SafetyPolicy, reason: str +) -> SafetyPolicy: + return SafetyPolicy( + policyType=policy_type, confidenceScore=source.confidenceScore, reason=reason + ) + + +class SafetyPtosPolicyCrossValidator: + result_pattern = re.compile(r"(.*)(.*)", re.DOTALL) + + def __init__(self): + eapi_4_5 = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_5_X_ALGO) + self.eapi_4_5_x_algo = EapiSampler(EapiModelConfig(**eapi_4_5.model_dump())) + eapi_4_1 = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_1_FAST_X_ALGO) + self.eapi_4_1_x_algo = EapiSampler(EapiModelConfig(**eapi_4_1.model_dump())) + + def _parse_policy(self, raw: str) -> SafetyPolicy | None: + match = self.result_pattern.search(raw) + if not match: + return None + try: + return SafetyPolicy.model_validate_json(match.group(2).strip()) + except Exception: + return None + + def _build_policy_convo( + self, post: Post, category: SafetyPolicyCategory, system_prompt: str + ) -> Conversation: + content = _strip_thinking_restrictions(system_prompt) + convo = Conversation(conversation_id=uuid.uuid4().hex) + convo.messages.append(Message(role=Role.SYSTEM, content=[content])) + user_msg = Message(role=Role.USER, content=[]) + user_msg.content.extend(UserRenderer.render(post.user)) + user_msg.content.extend(PostRenderer.render(post, include_reply_to=True)) + user_msg.content.append( + f"\n\nAnalyze the post {post.id} for the specific safety policy violation category: {category.value}" + ) + user_msg.content.append( + f"\n\nProvide the requested JSON object for the specific safety policy type.{THINKING_CONTROL_START}" + ) + convo.messages.append(user_msg) + convo.messages.append(Message(role=Role.ASSISTANT, content=[])) + return convo + + async def validate( + self, category: SafetyPolicyCategory, post: Post, policy: SafetyPolicy | None + ) -> SafetyPolicy | None: + if policy is None or policy.policyType == SafetyPolicyType.NoViolation: + return policy + if category == SafetyPolicyCategory.ChildSafety: + return await self._validate_child_safety(post, policy) + if category == SafetyPolicyCategory.ViolentMedia: + return await self._validate_violent_media(post, policy) + if category == SafetyPolicyCategory.IllegalAndRegulatedBehaviors: + return await self._validate_illegal_and_regulated_behaviors(post, policy) + return policy + + async def _validate_child_safety( + self, post: Post, policy: SafetyPolicy + ) -> SafetyPolicy: + metric = "safety_ptos.child_safety_cross_model_validate_with_grok_4_5" + post_creation_time = ( + post.created_at if post.created_at else datetime.now() + ).strftime("%Y-%m-%d") + convo = self._build_policy_convo( + post, + SafetyPolicyCategory.ChildSafety, + child_safety_policy_prompt(post_creation_time), + ) + try: + async with _eapi_4_5_x_algo_breaker.guard(): + raw = await self.eapi_4_5_x_algo.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) + confirm = self._parse_policy(raw) + if confirm is None: + logger.error( + f"child_safety cross_validate unparseable post={post.id} primary={policy.policyType.value} " + f"conversation_id={convo.conversation_id} raw={raw[:500]!r}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "unparseable"}) + return _policy_no_violation("child_safety_grok_4_5_parse_error") + if confirm.policyType == policy.policyType: + logger.info( + f"child_safety cross_validate agreed post={post.id} policy={policy.policyType.value} " + f"conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "agreed"}) + return policy + logger.info( + f"child_safety cross_validate disagreed post={post.id} primary={policy.policyType.value} " + f"confirm={confirm.policyType.value} conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "disagreed"}) + return _policy_no_violation("child_safety_grok_4_5_disagreed") + except Exception: + logger.error( + f"child_safety cross_validate failed post={post.id} primary={policy.policyType.value} " + f"conversation_id={convo.conversation_id}: {traceback.format_exc()}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "error"}) + return _policy_no_violation("child_safety_grok_4_5_sample_error") + + async def _validate_violent_media( + self, post: Post, policy: SafetyPolicy + ) -> SafetyPolicy: + metric = "safety_ptos.violent_media_cross_model_validate_with_grok_4_5" + convo = self._build_policy_convo( + post, SafetyPolicyCategory.ViolentMedia, violent_media_policy_prompt() + ) + try: + async with _eapi_4_5_x_algo_breaker.guard(): + raw = await self.eapi_4_5_x_algo.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) + confirm = self._parse_policy(raw) + if confirm is None: + logger.error( + f"violent_media cross_validate unparseable post={post.id} primary={policy.policyType.value} " + f"conversation_id={convo.conversation_id} raw={raw[:500]!r}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "unparseable"}) + return _policy_with_type( + SafetyPolicyType.ViolentMediaGraphicMedia, + policy, + f"violent_media_grok_4_5_parse_error: downgraded from {policy.policyType.value}", + ) + if confirm.policyType == policy.policyType: + logger.info( + f"violent_media cross_validate agreed post={post.id} policy={policy.policyType.value} " + f"conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "agreed"}) + return policy + if confirm.policyType == SafetyPolicyType.NoViolation: + logger.info( + f"violent_media cross_validate no_violation post={post.id} primary={policy.policyType.value} " + f"conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "no_violation"}) + return _policy_no_violation("violent_media_grok_4_5_no_violation") + logger.info( + f"violent_media cross_validate type_mismatch post={post.id} primary={policy.policyType.value} " + f"confirm={confirm.policyType.value} conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "type_mismatch"}) + return _policy_with_type( + SafetyPolicyType.ViolentMediaGraphicMedia, + policy, + f"violent_media_grok_4_5_type_mismatch: downgraded from {policy.policyType.value}, cv={confirm.policyType.value}", + ) + except Exception: + logger.error( + f"violent_media cross_validate failed post={post.id} primary={policy.policyType.value} " + f"conversation_id={convo.conversation_id}: {traceback.format_exc()}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "error"}) + return _policy_with_type( + SafetyPolicyType.ViolentMediaGraphicMedia, + policy, + f"violent_media_grok_4_5_sample_error: downgraded from {policy.policyType.value}", + ) + + async def _validate_illegal_and_regulated_behaviors( + self, post: Post, policy: SafetyPolicy + ) -> SafetyPolicy: + metric = "safety_ptos.illegal_and_regulated_behaviors_cross_model_validate_with_grok_4_1" + convo = self._build_policy_convo( + post, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, + illegal_and_regulated_behaviors_policy_prompt(), + ) + try: + raw = await self.eapi_4_1_x_algo.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) + confirm = self._parse_policy(raw) + if confirm is None: + logger.error( + f"illegal_and_regulated_behaviors cross_validate unparseable post={post.id} " + f"primary={policy.policyType.value} conversation_id={convo.conversation_id} raw={raw[:500]!r}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "unparseable"}) + return _policy_no_violation( + "illegal_and_regulated_behaviors_grok_4_1_parse_error" + ) + if confirm.policyType == policy.policyType: + logger.info( + f"illegal_and_regulated_behaviors cross_validate agreed post={post.id} " + f"policy={policy.policyType.value} conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "agreed"}) + return _policy_with_type( + policy.policyType, policy, confirm.reason or policy.reason + ) + logger.info( + f"illegal_and_regulated_behaviors cross_validate disagreed post={post.id} " + f"primary={policy.policyType.value} confirm={confirm.policyType.value} " + f"conversation_id={convo.conversation_id}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "disagreed"}) + return _policy_no_violation( + confirm.reason + or f"illegal_and_regulated_behaviors_grok_4_1_disagreed: cv={confirm.policyType.value}" + ) + except Exception: + logger.error( + f"illegal_and_regulated_behaviors cross_validate failed post={post.id} " + f"primary={policy.policyType.value} conversation_id={convo.conversation_id}: {traceback.format_exc()}" + ) + Metrics.counter(metric).add(1, attributes={"outcome": "error"}) + return _policy_no_violation( + "illegal_and_regulated_behaviors_grok_4_1_sample_error" + ) + + class SafetyPtosChildSafetyPolicyClassifier: result_pattern = re.compile(r"(.*)(.*)", re.DOTALL) def __init__(self, gemma_model_name: str = GEMMA): self.oai_gemma4 = OaiSampler(grox_config.get_oai_model(gemma_model_name)) - eapi_cfg = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_5_INTERNAL) - self.eapi_4_5_internal = EapiSampler(EapiModelConfig(**eapi_cfg.model_dump())) def build_convo(self, post: Post) -> Conversation: post_creation_time = ( @@ -203,10 +433,6 @@ def build_convo(self, post: Post) -> Conversation: convo.messages.append(Message(role=Role.ASSISTANT, content=[])) return convo - @staticmethod - def _no_violation(reason: str) -> SafetyPolicy: - return SafetyPolicy(policyType=SafetyPolicyType.NoViolation, reason=reason) - def _parse_policy(self, raw: str) -> SafetyPolicy | None: match = self.result_pattern.search(raw) if not match: @@ -229,7 +455,7 @@ async def classify_policy(self, post: Post) -> SafetyPolicy: Metrics.counter("safety_ptos.child_safety_policy.count").add( 1, attributes={"outcome": "gemma_error"} ) - return self._no_violation("child_safety_gemma_sample_error") + return _policy_no_violation("child_safety_gemma_sample_error") policy = self._parse_policy(raw) if policy is None: @@ -239,49 +465,8 @@ async def classify_policy(self, post: Post) -> SafetyPolicy: Metrics.counter("safety_ptos.child_safety_policy.count").add( 1, attributes={"outcome": "gemma_unparseable"} ) - return self._no_violation("child_safety_gemma_parse_error") - - if policy.policyType == SafetyPolicyType.NoViolation: - return policy - return await self._cross_model_validate_with_4_5(convo, policy, post_id=post.id) - - async def _cross_model_validate_with_4_5( - self, convo: Conversation, policy: SafetyPolicy, post_id: str - ) -> SafetyPolicy: - metric = "safety_ptos.cross_model_validate_with_grok_4_5" - try: - async with _eapi_4_5_internal_breaker.guard(): - raw = await self.eapi_4_5_internal.sample( - convo.interleaveToEapi(), conversation_id=convo.conversation_id - ) - confirm = self._parse_policy(raw) - if confirm is None: - logger.error( - f"child_safety cross_validate_4_5 unparseable post={post_id} primary={policy.policyType.value} " - f"conversation_id={convo.conversation_id} raw={raw[:500]!r}" - ) - Metrics.counter(metric).add(1, attributes={"outcome": "unparseable"}) - return self._no_violation("child_safety_grok_4_5_parse_error") - if confirm.policyType == policy.policyType: - logger.info( - f"child_safety cross_validate_4_5 agreed post={post_id} policy={policy.policyType.value} " - f"conversation_id={convo.conversation_id}" - ) - Metrics.counter(metric).add(1, attributes={"outcome": "agreed"}) - return policy - logger.info( - f"child_safety cross_validate_4_5 disagreed post={post_id} primary={policy.policyType.value} " - f"confirm={confirm.policyType.value} conversation_id={convo.conversation_id}, clearing to NoViolation" - ) - Metrics.counter(metric).add(1, attributes={"outcome": "disagreed"}) - return self._no_violation("child_safety_grok_4_5_disagreed") - except Exception: - logger.error( - f"child_safety cross_validate_4_5 failed post={post_id} primary={policy.policyType.value} " - f"conversation_id={convo.conversation_id}: {traceback.format_exc()}" - ) - Metrics.counter(metric).add(1, attributes={"outcome": "error"}) - return self._no_violation("child_safety_grok_4_5_sample_error") + return _policy_no_violation("child_safety_gemma_parse_error") + return policy class SafetyPtosPolicyClassifier: @@ -303,7 +488,6 @@ def __init__( oai_config = grox_config.get_oai_model(gemma_model_name) self.oai_gemma4 = OaiSampler(oai_config) - self.use_oai_gemma4_dial = 1.0 if self.deluxe: eapi_config_4_3_x_algo = grox_config.get_eapi_model( @@ -313,11 +497,11 @@ def __init__( EapiModelConfig(**eapi_config_4_3_x_algo.model_dump()) ) - eapi_config_4_5_internal = grox_config.get_eapi_model( - ModelName.EAPI_GROK_4_5_INTERNAL + eapi_config_4_6_internal = grox_config.get_eapi_model( + ModelName.EAPI_GROK_4_6_INTERNAL ) - self.eapi_4_5_internal = EapiSampler( - EapiModelConfig(**eapi_config_4_5_internal.model_dump()) + self.eapi_4_6_internal = EapiSampler( + EapiModelConfig(**eapi_config_4_6_internal.model_dump()) ) @staticmethod @@ -402,6 +586,7 @@ def build_convo( USE_GEMMA_CATEGORIES = { SafetyPolicyCategory.Spam, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, } USE_THREAD_RENDERER_CATEGORIES = { @@ -433,8 +618,8 @@ async def classify_policy_for_violation( and violation.category in self.DELUXE_4_3_CATEGORIES ): if fav_count >= 1024: - mode = "deluxe-4.5-internal" - result = await self._sample_4_5_internal(convo) + mode = "deluxe-4.6-internal" + result = await self._sample_4_6_internal(convo) else: mode = "deluxe-4.3" result = await self._sample_4_3(convo) @@ -498,46 +683,84 @@ async def _sample_4_3(self, convo: Conversation) -> str: convo.interleave(), conversation_id=convo.conversation_id ) - async def _sample_4_5_internal(self, convo: Conversation) -> str: - breaker, sampler = _eapi_4_5_internal_breaker, self.eapi_4_5_internal + async def _sample_4_6_internal(self, convo: Conversation) -> str: + breaker, sampler = _eapi_4_6_internal_breaker, self.eapi_4_6_internal try: async with breaker.guard(): return await sampler.sample( convo.interleaveToEapi(), conversation_id=convo.conversation_id ) except CircuitBreakerOpen as e: - Metrics.counter("safety_ptos.eapi_4_5_fallback.count").add( + Metrics.counter("safety_ptos.eapi_4_6_fallback.count").add( 1, attributes={"endpoint": breaker.name, "reason": "breaker_open"} ) logger.warning( - f"4.5 circuit breaker '{e.name}' open (recovery in {e.remaining_seconds:.0f}s), falling back to 4.1" + f"4.6 circuit breaker '{e.name}' open (recovery in {e.remaining_seconds:.0f}s), falling back to 4.1" ) except Exception: - Metrics.counter("safety_ptos.eapi_4_5_fallback.count").add( + Metrics.counter("safety_ptos.eapi_4_6_fallback.count").add( 1, attributes={"endpoint": breaker.name, "reason": "error"} ) logger.error( - f"Failed to call 4.5-internal reasoning, conversation_id={convo.conversation_id}, error: {traceback.format_exc()}" + f"Failed to call 4.6-internal reasoning, conversation_id={convo.conversation_id}, error: {traceback.format_exc()}" ) return await self.llm.sample( convo.interleave(), conversation_id=convo.conversation_id ) async def _sample(self, convo: Conversation, sample_for_gemma: bool = False) -> str: - if ( - sample_for_gemma - and not self.deluxe - and self.use_gemma - and random.random() < self.use_oai_gemma4_dial - ): + if sample_for_gemma: try: return await self.oai_gemma4.sample( convo.to_openai_messages(), conversation_id=convo.conversation_id ) except Exception: logger.error( - f"OaiSampler (gemma4) failed for spam policy, falling back to grok: {traceback.format_exc()}" + f"OaiSampler (gemma4) failed for policy, falling back to grok: {traceback.format_exc()}" ) return await self.llm.sample( convo.interleave(), conversation_id=convo.conversation_id ) + + +class SafetyPtosAdultContentCrossValidationJudge: + result_pattern = re.compile(r"(.*)(.*)", re.DOTALL) + + def __init__(self): + eapi_cfg = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_5_X_ALGO) + self.llm = EapiSampler(EapiModelConfig(**eapi_cfg.model_dump())) + + def build_convo(self, post: Post) -> Conversation: + convo = Conversation(conversation_id=uuid.uuid4().hex) + convo.messages.append( + Message( + role=Role.SYSTEM, + content=[_strip_thinking_restrictions(adult_content_policy_prompt())], + ) + ) + + user_msg = Message(role=Role.USER, content=[]) + user_msg.content.extend(UserRenderer.render(post.user)) + user_msg.content.extend(PostRenderer.render(post, include_reply_to=True)) + user_msg.content.append( + f"\n\nAnalyze the post {post.id} for the specific safety policy violation category: {SafetyPolicyCategory.AdultContent.value}" + ) + user_msg.content.append( + f"\n\nProvide the requested JSON object for the specific safety policy type.{THINKING_CONTROL_START}" + ) + convo.messages.append(user_msg) + convo.messages.append(Message(role=Role.ASSISTANT, content=[])) + return convo + + async def judge(self, post: Post) -> SafetyPolicy: + convo = self.build_convo(post) + async with _eapi_4_5_x_algo_breaker.guard(): + raw = await self.llm.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) + match = self.result_pattern.search(raw) + if not match: + raise ValueError( + f"Invalid output for adult content cross validation judge: {raw[:500]!r}" + ) + return SafetyPolicy.model_validate_json(match.group(2).strip()) diff --git a/grox/flows/ptos/constants.py b/grox/flows/ptos/constants.py index 26e0a486..95b6cb28 100644 --- a/grox/flows/ptos/constants.py +++ b/grox/flows/ptos/constants.py @@ -23,3 +23,7 @@ GEMMA_PTOS_REALTIME = "oai-gemma4-26b-ptos-realtime" HIGH_FAV_THRESHOLD = 128 + +SAFETY_PTOS_SPECIAL_VIDEO = "safety_ptos_special_video" +TOPIC_SPECIAL_VIDEO = "safety-ptos-special-video" +DELUXE_TIER_TASK_TYPES = frozenset({SAFETY_PTOS_DELUXE, SAFETY_PTOS_SPECIAL_VIDEO}) diff --git a/grox/flows/ptos/generators.py b/grox/flows/ptos/generators.py index e769a2fb..cc7c886d 100644 --- a/grox/flows/ptos/generators.py +++ b/grox/flows/ptos/generators.py @@ -3,6 +3,8 @@ from grox.core.registry import register from grox.flows.ptos.constants import ( POST_MIN_IMPRESSION_STREAM_FOR_GROX_PTOS, + SAFETY_PTOS_SPECIAL_VIDEO, + TOPIC_SPECIAL_VIDEO, POST_MIN_TRACTION_STREAM_FOR_GROX_PTOS, SAFETY_PTOS_BACKFILL, SAFETY_PTOS_DELUXE, @@ -19,6 +21,7 @@ ) from grox.flows.ptos.kafka_loader import KafkaLiveClusterAnchorLoader from grox.flows.ptos.plan_safety_ptos import PlanSafetyPtos +from grox.flows.ptos.plan_safety_ptos_special_video import PlanSafetyPtosSpecialVideo from grox.flows.ptos.plan_safety_ptos_live_cluster_anchors import ( PlanSafetyPtosLiveClusterAnchors, ) @@ -88,3 +91,12 @@ class SafetyPtosLiveClusterAnchorsStreamTaskGenerator(StreamTaskGenerator): def _get_loader(self): return KafkaLiveClusterAnchorLoader(TOPIC_LIVE_CLUSTER_ANCHORS) + + +@register +class SafetyPtosSpecialVideoStreamTaskGenerator(StreamTaskGenerator): + TASK_GENERATOR_TYPE = SAFETY_PTOS_SPECIAL_VIDEO + PLANS_TO_INJECT = {PlanSafetyPtosSpecialVideo.KEY} + + def _get_loader(self): + return KafkaPostLoader(TOPIC_SPECIAL_VIDEO) diff --git a/grox/flows/ptos/mode.py b/grox/flows/ptos/mode.py index eb894e80..2381e30b 100644 --- a/grox/flows/ptos/mode.py +++ b/grox/flows/ptos/mode.py @@ -12,7 +12,7 @@ class SafetyPtosMode(str, Enum): @classmethod def from_task_type(cls, task_type: str | None) -> "SafetyPtosMode": match task_type: - case constants.SAFETY_PTOS_DELUXE: + case t if t in constants.DELUXE_TIER_TASK_TYPES: return cls.DELUXE case constants.SAFETY_PTOS_RECOVERY: return cls.RECOVERY diff --git a/grox/flows/ptos/plan_safety_ptos.py b/grox/flows/ptos/plan_safety_ptos.py index 92215599..9f96d0e9 100644 --- a/grox/flows/ptos/plan_safety_ptos.py +++ b/grox/flows/ptos/plan_safety_ptos.py @@ -3,6 +3,9 @@ from grox.core.tasks.task_media import TaskMediaHydration from grox.flows.ptos.task_filter import TaskSafetyPtosFilter from grox.flows.ptos.task_safety_ptos_category import TaskSafetyPtosCategoryDetection +from grox.flows.ptos.task_safety_ptos_adult_content_cross_validation import ( + TaskSafetyPtosAdultContentCrossValidation, +) from grox.flows.ptos.task_safety_ptos_policy import TaskSafetyPtosPolicyDetection from grox.flows.ptos.task_safety_ptos_safemodel_sex_nudity import ( TaskSafetyPtosSafemodelSexNudity, @@ -28,6 +31,7 @@ class PlanSafetyPtos(Plan): "task_safety_ptos_category_detection": TaskSafetyPtosCategoryDetection, "task_safety_ptos_policy_detection": TaskSafetyPtosPolicyDetection, "task_safety_ptos_safemodel_sex_nudity": TaskSafetyPtosSafemodelSexNudity, + "task_safety_ptos_adult_content_cross_validation": TaskSafetyPtosAdultContentCrossValidation, "task_write_safety_post_annotations_result_sink": TaskWriteSafetyPostAnnotationsResultSink, } @@ -41,8 +45,11 @@ class PlanSafetyPtos(Plan): "task_safety_ptos_category_detection": {"task_media_hydration"}, "task_safety_ptos_policy_detection": {"task_safety_ptos_category_detection"}, "task_safety_ptos_safemodel_sex_nudity": {"task_safety_ptos_policy_detection"}, - "task_write_safety_post_annotations_result_sink": { + "task_safety_ptos_adult_content_cross_validation": { "task_safety_ptos_policy_detection", "task_safety_ptos_safemodel_sex_nudity", }, + "task_write_safety_post_annotations_result_sink": { + "task_safety_ptos_adult_content_cross_validation" + }, } diff --git a/grox/flows/ptos/plan_safety_ptos_special_video.py b/grox/flows/ptos/plan_safety_ptos_special_video.py new file mode 100644 index 00000000..dcd2f653 --- /dev/null +++ b/grox/flows/ptos/plan_safety_ptos_special_video.py @@ -0,0 +1,55 @@ +from grox.core.plans.plan import Plan +from grox.core.registry import register +from grox.core.tasks.task_media import TaskMediaHydration +from grox.flows.ptos.task_special_video_screen import TaskSpecialVideoScreen +from grox.flows.ptos.task_safety_ptos_adult_content_cross_validation import ( + TaskSafetyPtosAdultContentCrossValidation, +) +from grox.flows.ptos.task_rate_limit import TaskRateLimitSafetyPtosAnnotationWithPost +from grox.flows.ptos.task_safety_ptos_category import TaskSafetyPtosCategoryDetection +from grox.flows.ptos.task_safety_ptos_policy import TaskSafetyPtosPolicyDetection +from grox.flows.ptos.task_safety_ptos_safemodel_sex_nudity import ( + TaskSafetyPtosSafemodelSexNudity, +) +from grox.flows.ptos.task_safety_ptos_special_video_filter import ( + TaskSafetyPtosSpecialVideoFilter, +) +from grox.flows.ptos.task_write_safety_post_annotations_result_sink import ( + TaskWriteSafetyPostAnnotationsResultSink, +) + + +@register +class PlanSafetyPtosSpecialVideo(Plan): + KEY = "safety_ptos_special_video" + + TASKS = { + "task_safety_ptos_special_video_filter": TaskSafetyPtosSpecialVideoFilter, + "task_safety_ptos_annotation_rate_limit": TaskRateLimitSafetyPtosAnnotationWithPost, + "task_media_hydration": TaskMediaHydration, + "task_special_video_screen": TaskSpecialVideoScreen, + "task_safety_ptos_category_detection": TaskSafetyPtosCategoryDetection, + "task_safety_ptos_policy_detection": TaskSafetyPtosPolicyDetection, + "task_safety_ptos_safemodel_sex_nudity": TaskSafetyPtosSafemodelSexNudity, + "task_safety_ptos_adult_content_cross_validation": TaskSafetyPtosAdultContentCrossValidation, + "task_write_safety_post_annotations_result_sink": TaskWriteSafetyPostAnnotationsResultSink, + } + + TASK_DEPENDENCIES = { + "task_safety_ptos_special_video_filter": {}, + "task_safety_ptos_annotation_rate_limit": { + "task_safety_ptos_special_video_filter" + }, + "task_media_hydration": {"task_safety_ptos_annotation_rate_limit"}, + "task_special_video_screen": {"task_media_hydration"}, + "task_safety_ptos_category_detection": {"task_special_video_screen"}, + "task_safety_ptos_policy_detection": {"task_safety_ptos_category_detection"}, + "task_safety_ptos_safemodel_sex_nudity": {"task_safety_ptos_policy_detection"}, + "task_safety_ptos_adult_content_cross_validation": { + "task_safety_ptos_policy_detection", + "task_safety_ptos_safemodel_sex_nudity", + }, + "task_write_safety_post_annotations_result_sink": { + "task_safety_ptos_adult_content_cross_validation" + }, + } diff --git a/grox/flows/ptos/prior_nsfw.py b/grox/flows/ptos/prior_nsfw.py new file mode 100644 index 00000000..b86242ea --- /dev/null +++ b/grox/flows/ptos/prior_nsfw.py @@ -0,0 +1,34 @@ +import logging + +from grox.core.data_loaders.data_types import Post +from grox.core.schedules.types import TaskContext +from grox.flows.ptos.state import SafetyPtosState +from monitor.metrics import Metrics +from strato_http.queries.safety_post_annotations_result import ( + StratoSafetyPostAnnotationsResultDirectMh, +) + +logger = logging.getLogger(__name__) + +_result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() + + +async def post_is_already_flagged_nsfw(ctx: TaskContext, post: Post) -> bool: + state = ctx.state(SafetyPtosState) + if state.prior_nsfw is None: + state.prior_nsfw = await _fetch(post) + return state.prior_nsfw + + +async def _fetch(post: Post) -> bool: + try: + result = await _result_direct_mh.fetch(int(post.id)) + except Exception as e: + Metrics.counter("safety_ptos.prior_nsfw_lookup_error.count").add(1) + logger.warning( + f"Post {post.id}: NSFW MH lookup failed, treating as not flagged: {e}" + ) + return False + return bool( + result and result.safetyBoolMetadata and result.safetyBoolMetadata.isNsfw + ) diff --git a/grox/flows/ptos/state.py b/grox/flows/ptos/state.py index 9be56b7e..cdaa80ea 100644 --- a/grox/flows/ptos/state.py +++ b/grox/flows/ptos/state.py @@ -88,6 +88,7 @@ class SafetyPostAnnotations(BaseModel): class SafemodelResult(BaseModel): + scored: bool = False positive: bool = False confidence: float = 0.0 @@ -123,3 +124,4 @@ class LiveClusterAnchorKafkaVerdict(BaseModel): class SafetyPtosState: annotations: SafetyPostAnnotations | None = None safemodel_sex_nudity: SafemodelResult = field(default_factory=SafemodelResult) + prior_nsfw: bool | None = None diff --git a/grox/flows/ptos/task_rate_limit.py b/grox/flows/ptos/task_rate_limit.py index 2cf8c396..4b9f88c0 100644 --- a/grox/flows/ptos/task_rate_limit.py +++ b/grox/flows/ptos/task_rate_limit.py @@ -5,7 +5,7 @@ from grox.core.data_loaders.data_types import Post from grox.core.schedules.types import TaskContext from grox.core.tasks.task_rate_limit import TaskTTLDedupeWithPost -from grox.flows.ptos.constants import SAFETY_PTOS_DELUXE +from grox.flows.ptos.constants import DELUXE_TIER_TASK_TYPES class TaskRateLimitSafetyPtosAnnotationWithPost(TaskTTLDedupeWithPost): @@ -15,7 +15,7 @@ class TaskRateLimitSafetyPtosAnnotationWithPost(TaskTTLDedupeWithPost): @override @classmethod async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: - is_deluxe = ctx.payload.task_type == SAFETY_PTOS_DELUXE + is_deluxe = ctx.payload.task_type in DELUXE_TIER_TASK_TYPES cache = ( cls.POST_CACHE_FOR_SAFETY_PTOS_DELUXE if is_deluxe diff --git a/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py new file mode 100644 index 00000000..648167d9 --- /dev/null +++ b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py @@ -0,0 +1,152 @@ +import logging +from enum import Enum + +from grox.core.data_loaders.data_types import Post +from grox.core.schedules.types import TaskContext +from grox.core.tasks.task import Task, TaskResultCategory, TaskWithPost +from grox.flows.ptos.classifier import SafetyPtosAdultContentCrossValidationJudge +from grox.flows.ptos.constants import DELUXE_TIER_TASK_TYPES, SAFETY_PTOS_SPECIAL_VIDEO +from grox.flows.ptos.state import ( + SafetyPolicy, + SafetyPolicyCategory, + SafetyPolicyType, + SafetyPtosState, + SafetyPtosViolatedPolicy, +) +from monitor.metrics import Metrics +from grox.flows.ptos.prior_nsfw import post_is_already_flagged_nsfw + +logger = logging.getLogger(__name__) + +_METRIC_PREFIX = "task.safety_ptos_adult_content_cross_validation" +_CROSS_VALIDATION_REASON = "Grok 4.5 Cross Validation disagreed" +_CROSS_VALIDATION_FAILED_REASON = "Grok 4.5 Cross Validation failed" + + +class CompareOutcome(str, Enum): + BOTH_POSITIVE = "both_positive" + BOTH_NEGATIVE = "both_negative" + SAFEMODEL_ONLY_POSITIVE = "safemodel_only_positive" + PTOS_ONLY_POSITIVE = "ptos_only_positive" + + @property + def is_disagreement(self) -> bool: + return self in ( + CompareOutcome.SAFEMODEL_ONLY_POSITIVE, + CompareOutcome.PTOS_ONLY_POSITIVE, + ) + + +class TaskSafetyPtosAdultContentCrossValidation(TaskWithPost): + _judge = SafetyPtosAdultContentCrossValidationJudge() + + @classmethod + async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: + try: + await cls._run(ctx, post) + except Exception as e: + Metrics.counter(f"{_METRIC_PREFIX}.error.count").add(1) + logger.warning( + f"Post {post.id}: cross validation failed, keeping original decision: {e}" + ) + + @classmethod + async def _run(cls, ctx: TaskContext, post: Post) -> None: + state = ctx.state(SafetyPtosState) + if not state.safemodel_sex_nudity.scored: + return + + is_deluxe = ctx.payload.task_type in DELUXE_TIER_TASK_TYPES + flow = "deluxe" if is_deluxe else "standard" + safemodel_positive = state.safemodel_sex_nudity.positive + violations = ( + (state.annotations.violatedPolicies or []) if state.annotations else [] + ) + ptos_positive = any( + v.category == SafetyPolicyCategory.AdultContent + and v.safetyPolicy is not None + and v.safetyPolicy.policyType == SafetyPolicyType.AdultContentSexualHard + for v in violations + ) + + outcome = cls._compare_outcome(safemodel_positive, ptos_positive) + Metrics.counter(f"{_METRIC_PREFIX}.compare.count").add( + 1, attributes={"outcome": outcome.value, "flow": flow} + ) + logger.info( + f"Post {post.id} ({flow}): safemodel={'positive' if safemodel_positive else 'negative'} " + f"ptos={'positive' if ptos_positive else 'negative'} outcome={outcome.value}" + ) + + if is_deluxe and outcome.is_disagreement: + if await post_is_already_flagged_nsfw(ctx, post): + Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( + 1, attributes={"reason": "prior_nsfw"} + ) + return + await cls._cross_validate(ctx, post) + + @staticmethod + def _compare_outcome( + safemodel_positive: bool, ptos_positive: bool + ) -> CompareOutcome: + if safemodel_positive and ptos_positive: + return CompareOutcome.BOTH_POSITIVE + if not safemodel_positive and not ptos_positive: + return CompareOutcome.BOTH_NEGATIVE + if safemodel_positive: + return CompareOutcome.SAFEMODEL_ONLY_POSITIVE + return CompareOutcome.PTOS_ONLY_POSITIVE + + @classmethod + async def _cross_validate(cls, ctx: TaskContext, post: Post) -> None: + metric = f"{_METRIC_PREFIX}.judged.count" + try: + judged = await cls._judge.judge(post) + except Exception as e: + Metrics.counter(metric).add(1, attributes={"outcome": "error"}) + if ctx.payload.task_type == SAFETY_PTOS_SPECIAL_VIDEO: + logger.warning( + f"Post {post.id}: grok 4.5 cross validation failed, failing closed to Soft: {e}" + ) + cls._apply_soft_verdict(ctx, _CROSS_VALIDATION_FAILED_REASON) + else: + logger.warning( + f"Post {post.id}: grok 4.5 cross validation failed, keeping original decision: {e}" + ) + return + + is_hard = judged.policyType == SafetyPolicyType.AdultContentSexualHard + Metrics.counter(metric).add(1, attributes={"outcome": judged.policyType.value}) + logger.info( + f"Post {post.id}: grok 4.5 cross validation judged {judged.policyType.value}" + ) + if is_hard: + return + + cls._apply_soft_verdict(ctx, _CROSS_VALIDATION_REASON) + + @classmethod + def _apply_soft_verdict(cls, ctx: TaskContext, reason: str) -> None: + state = ctx.state(SafetyPtosState) + violations = state.annotations.violatedPolicies or [] + adult_violations = [ + v for v in violations if v.category == SafetyPolicyCategory.AdultContent + ] + if not adult_violations: + adult_violations = [ + SafetyPtosViolatedPolicy( + category=SafetyPolicyCategory.AdultContent, reason=reason + ) + ] + violations.append(adult_violations[0]) + for violation in adult_violations: + violation.safetyPolicy = SafetyPolicy( + policyType=SafetyPolicyType.AdultContentSexualSoft, reason=reason + ) + state.annotations.violatedPolicies = violations + state.safemodel_sex_nudity.positive = False + + @classmethod + async def exec(cls, ctx: TaskContext) -> TaskResultCategory: + return await Task.exec.__wrapped__(cls, ctx) diff --git a/grox/flows/ptos/task_safety_ptos_policy.py b/grox/flows/ptos/task_safety_ptos_policy.py index d3343acd..d0d7ea13 100644 --- a/grox/flows/ptos/task_safety_ptos_policy.py +++ b/grox/flows/ptos/task_safety_ptos_policy.py @@ -15,13 +15,12 @@ from grox.flows.ptos.classifier import ( SafetyPtosChildSafetyPolicyClassifier, SafetyPtosPolicyClassifier, + SafetyPtosPolicyCrossValidator, ) from grox.config.config import ModelName from grox.flows.ptos.mode import SafetyPtosMode from grox.flows.ptos.constants import GEMMA, HIGH_FAV_THRESHOLD -from strato_http.queries.safety_post_annotations_result import ( - StratoSafetyPostAnnotationsResultDirectMh, -) +from grox.flows.ptos.prior_nsfw import post_is_already_flagged_nsfw logger = logging.getLogger(__name__) @@ -30,6 +29,7 @@ class TaskSafetyPtosPolicyDetection(TaskWithPost): child_safety_classifier = SafetyPtosChildSafetyPolicyClassifier( gemma_model_name=GEMMA ) + cross_validator = SafetyPtosPolicyCrossValidator() classifiers = { SafetyPtosMode.STANDARD: SafetyPtosPolicyClassifier( @@ -64,24 +64,6 @@ class TaskSafetyPtosPolicyDetection(TaskWithPost): ), } - _result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() - - @classmethod - async def _post_is_already_flagged_nsfw(cls, post: Post) -> bool: - try: - result = await cls._result_direct_mh.fetch(int(post.id)) - except Exception as e: - Metrics.counter( - "task.safety_ptos_deluxe_policy.nsfw_lookup_error.count" - ).add(1) - logger.warning( - f"Post {post.id}: NSFW MH lookup failed, treating as not flagged: {e}" - ) - return False - return bool( - result and result.safetyBoolMetadata and result.safetyBoolMetadata.isNsfw - ) - @classmethod async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: annotations = ctx.state(SafetyPtosState).annotations @@ -98,26 +80,10 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: violations = list(annotations.violatedPolicies or []) - injected_recheck = None - if mode.is_deluxe and post.get_fav_count() >= HIGH_FAV_THRESHOLD: - if ( - not any( - v.category == SafetyPolicyCategory.AdultContent for v in violations - ) - and PostRenderer.has_media(post) - and not await cls._post_is_already_flagged_nsfw(post) - ): - injected_recheck = SafetyPtosViolatedPolicy( - category=SafetyPolicyCategory.AdultContent, - reason="high-fav adult content recheck", - score=50, - ) - violations.append(injected_recheck) - for violation in violations: if ( violation.category == SafetyPolicyCategory.AdultContent - and await cls._post_is_already_flagged_nsfw(post) + and await post_is_already_flagged_nsfw(ctx, post) ): Metrics.counter( f"{metric_prefix}.skipped_adult_content_already_nsfw.count" @@ -127,8 +93,19 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: reason="post was previously flagged nsfw by ptos", ) elif violation.category == SafetyPolicyCategory.ChildSafety: - violation.safetyPolicy = ( - await cls.child_safety_classifier.classify_policy(post) + policy = await cls.child_safety_classifier.classify_policy(post) + violation.safetyPolicy = await cls.cross_validator.validate( + violation.category, post, policy + ) + elif violation.category in ( + SafetyPolicyCategory.ViolentMedia, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, + ): + policy = await active_classifier.classify_policy_for_violation( + post, violation + ) + violation.safetyPolicy = await cls.cross_validator.validate( + violation.category, post, policy ) else: violation.safetyPolicy = ( @@ -139,13 +116,47 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: if violation.safetyPolicy: cls._record_policy_metrics(metric_prefix, violation) - if injected_recheck is not None: - policy = injected_recheck.safetyPolicy - if not policy or policy.policyType == SafetyPolicyType.NoViolation: - violations.remove(injected_recheck) + recheck = await cls._high_fav_adult_recheck( + ctx, post, mode, active_classifier, violations, metric_prefix + ) + if recheck is not None: + violations.append(recheck) annotations.violatedPolicies = violations + @classmethod + async def _high_fav_adult_recheck( + cls, + ctx: TaskContext, + post: Post, + mode: SafetyPtosMode, + active_classifier: SafetyPtosPolicyClassifier, + violations: list[SafetyPtosViolatedPolicy], + metric_prefix: str, + ) -> SafetyPtosViolatedPolicy | None: + if not mode.is_deluxe or post.get_fav_count() < HIGH_FAV_THRESHOLD: + return None + if any(v.category == SafetyPolicyCategory.AdultContent for v in violations): + return None + if not PostRenderer.has_media(post) or await post_is_already_flagged_nsfw( + ctx, post + ): + return None + + recheck = SafetyPtosViolatedPolicy( + category=SafetyPolicyCategory.AdultContent, + reason="high-fav adult content recheck", + score=50, + ) + policy = await active_classifier.classify_policy_for_violation(post, recheck) + if policy is None: + return None + recheck.safetyPolicy = policy + cls._record_policy_metrics(metric_prefix, recheck) + if policy.policyType == SafetyPolicyType.NoViolation: + return None + return recheck + @classmethod def _record_policy_metrics( cls, metric_prefix: str, violation: SafetyPtosViolatedPolicy diff --git a/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py b/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py index a7709b61..aafea867 100644 --- a/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py +++ b/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py @@ -9,15 +9,12 @@ from grox.core.lib.utils import detect_image_content_type from grox.core.data_loaders.data_types import Image, Post, Video -from grox.flows.ptos.state import ( - SafetyPolicyCategory, - SafetyPolicyType, - SafetyPtosState, -) +from grox.flows.ptos.state import SafetyPolicyCategory, SafetyPtosState from grox.core.schedules.types import TaskContext from grox.core.tasks.task import Task, TaskWithPost, TaskResultCategory from monitor.metrics import Metrics -from grox.flows.ptos.constants import SAFETY_PTOS_DELUXE +from grox.flows.ptos.constants import DELUXE_TIER_TASK_TYPES +from grox.flows.ptos.prior_nsfw import post_is_already_flagged_nsfw logger = logging.getLogger(__name__) @@ -62,7 +59,7 @@ def _has_adult_content_suspicion(cls, ctx: TaskContext) -> bool: @classmethod async def _run(cls, ctx: TaskContext, post: Post) -> None: - is_deluxe = ctx.payload.task_type == SAFETY_PTOS_DELUXE + is_deluxe = ctx.payload.task_type in DELUXE_TIER_TASK_TYPES flow = "deluxe" if is_deluxe else "standard" if not is_deluxe and not cls._has_adult_content_suspicion(ctx): @@ -71,6 +68,12 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: ) return + if is_deluxe and await post_is_already_flagged_nsfw(ctx, post): + Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( + 1, attributes={"reason": "prior_nsfw", "flow": flow} + ) + return + payloads = cls._collect_payloads(post) if not payloads: Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( @@ -123,43 +126,18 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: ) return - annotations = ctx.state(SafetyPtosState).annotations - violations = (annotations.violatedPolicies or []) if annotations else [] - ptos_positive = any( - v.category == SafetyPolicyCategory.AdultContent - and v.safetyPolicy is not None - and v.safetyPolicy.policyType == SafetyPolicyType.AdultContentSexualHard - for v in violations - ) - - outcome = cls._compare_outcome(safemodel_positive, ptos_positive) - Metrics.counter(f"{_METRIC_PREFIX}.compare.count").add( - 1, - attributes={"outcome": outcome, "has_video": has_video_attr, "flow": flow}, - ) - logger.info( f"Post {post.id} ({flow}): safemodel={'positive' if safemodel_positive else 'negative'} " - f"(buckets={buckets_seen}, n_errors={n_errors}, n_images={n_images}, n_video_frames={n_video_frames}) " - f"ptos={'positive' if ptos_positive else 'negative'} outcome={outcome}" + f"(buckets={buckets_seen}, n_errors={n_errors}, n_images={n_images}, n_video_frames={n_video_frames})" ) + ctx.state(SafetyPtosState).safemodel_sex_nudity.scored = True if safemodel_positive: ctx.state(SafetyPtosState).safemodel_sex_nudity.positive = True ctx.state( SafetyPtosState ).safemodel_sex_nudity.confidence = max_positive_confidence - @staticmethod - def _compare_outcome(safemodel_positive: bool, ptos_positive: bool) -> str: - if safemodel_positive and ptos_positive: - return "both_positive" - if not safemodel_positive and not ptos_positive: - return "both_negative" - if safemodel_positive: - return "safemodel_only_positive" - return "ptos_only_positive" - @classmethod def _collect_payloads(cls, post: Post) -> list[tuple[str, bytes]]: payloads: list[tuple[str, bytes]] = [] diff --git a/grox/flows/ptos/task_safety_ptos_special_video_filter.py b/grox/flows/ptos/task_safety_ptos_special_video_filter.py new file mode 100644 index 00000000..fa9c4aa6 --- /dev/null +++ b/grox/flows/ptos/task_safety_ptos_special_video_filter.py @@ -0,0 +1,40 @@ +from typing import override + +from monitor.metrics import Metrics + +from grox.config.config import grox_config +from grox.core.data_loaders.data_types import Post, Video +from grox.core.schedules.types import TaskContext +from grox.core.tasks.task_filters import TaskFilterWithPost + +_METRIC_PREFIX = "task.safety_ptos_special_video_filter" + + +class TaskSafetyPtosSpecialVideoFilter(TaskFilterWithPost): + @override + @classmethod + async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: + reason = cls._skip_reason(post) + if reason is not None: + Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( + 1, attributes={"reason": reason} + ) + return False + Metrics.counter(f"{_METRIC_PREFIX}.eligible.count").add(1) + return True + + @classmethod + def _skip_reason(cls, post: Post) -> str | None: + if not post.user: + return "no_user" + media = list(post.media or []) + if post.quoted_post and post.quoted_post.media: + media.extend(post.quoted_post.media) + if not any(isinstance(medium, Video) for medium in media): + return "no_video" + if ( + post.get_fav_count() + < grox_config.media_hydration.deluxe_fav_count_threshold + ): + return "not_high_fav" + return None diff --git a/grox/flows/ptos/task_special_video_screen.py b/grox/flows/ptos/task_special_video_screen.py new file mode 100644 index 00000000..2e44f9f0 --- /dev/null +++ b/grox/flows/ptos/task_special_video_screen.py @@ -0,0 +1,49 @@ +from typing import override + +from monitor.metrics import Metrics + +from grox.core.data_loaders.data_types import Post, Video +from grox.core.lm.convo import Video as ConvoVideo +from grox.core.schedules.types import TaskContext +from grox.core.tasks.task_filters import TaskFilterWithPost + +_METRIC_PREFIX = "task.safety_ptos_special_video_screen" + + +class TaskSpecialVideoScreen(TaskFilterWithPost): + @classmethod + def _technique_signals(cls, convo_video: ConvoVideo) -> dict[str, bool]: + return { + "motion_reveal": bool(convo_video.motion_reveal_frames), + } + + @override + @classmethod + async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: + screened = 0 + hits: set[str] = set() + media = list(post.media or []) + if post.quoted_post and post.quoted_post.media: + media.extend(post.quoted_post.media) + for medium in media: + if not isinstance(medium, Video) or not medium.convo_video: + continue + screened += 1 + hits.update( + technique + for technique, hit in cls._technique_signals(medium.convo_video).items() + if hit + ) + if screened == 0: + Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( + 1, attributes={"reason": "no_hydrated_video"} + ) + return False + Metrics.counter(f"{_METRIC_PREFIX}.screened.count").add( + 1, + attributes={ + "hit": str(bool(hits)).lower(), + "techniques": ",".join(sorted(hits)) or "none", + }, + ) + return bool(hits) diff --git a/grox/flows/ptos/task_write_safety_post_annotations_result_sink.py b/grox/flows/ptos/task_write_safety_post_annotations_result_sink.py index 335bb28b..85b93b2b 100644 --- a/grox/flows/ptos/task_write_safety_post_annotations_result_sink.py +++ b/grox/flows/ptos/task_write_safety_post_annotations_result_sink.py @@ -235,6 +235,7 @@ async def _is_test_user() -> bool: if policy_type in ( SafetyPolicyType.SpamEngagementBaiting, SafetyPolicyType.SpamEngagementFarming, + SafetyPolicyType.SpamCardManipulation, ): should_label_for_spam = True elif policy_type != SafetyPolicyType.NoViolation: diff --git a/grox/flows/reply_spam/classifier_coordinated_spam.py b/grox/flows/reply_spam/classifier_coordinated_spam.py index ca97e56f..362e6453 100644 --- a/grox/flows/reply_spam/classifier_coordinated_spam.py +++ b/grox/flows/reply_spam/classifier_coordinated_spam.py @@ -123,7 +123,11 @@ async def _to_convo(self, post: Post) -> Conversation: else: tag = "" content.append(f"\n\n#### Post {i}{tag}\n\n") - content.extend(PostRenderer.render(p, max_media=MAX_MEDIA_PER_POST)) + content.extend( + PostRenderer.render( + p, max_media=MAX_MEDIA_PER_POST, include_follower_count=True + ) + ) content.append("\n\n------\n\n") convo.messages.append(Message(role=Role.HUMAN, content=content)) return convo diff --git a/grox/flows/reply_spam/classifier_multi_step_reply_spam.py b/grox/flows/reply_spam/classifier_multi_step_reply_spam.py new file mode 100644 index 00000000..32f717d5 --- /dev/null +++ b/grox/flows/reply_spam/classifier_multi_step_reply_spam.py @@ -0,0 +1,226 @@ +import re +import json +import uuid +import logging + +import json_repair + +from grok_sampler.oai_sampler import OaiSampler +from grox.core.lm.post import PostRenderer +from grox.core.lm.user import UserRenderer +from grox.core.lm.convo import Role, Message, Conversation +from grox.config.config import grox_config +from grox.flows.reply_spam.prompts import multi_step_reply_spam_system_prompt +from grox.core.data_loaders.data_types import Post +from grox.flows.reply_spam.state_multi_step_reply_spam import MultiStepReplySpamResult +from monitor.metrics import Metrics +from grox.flows.reply_spam.constants import GEMMA_2 + +logger = logging.getLogger(__name__) + +MAX_MEDIA_PER_POST = 2 +MAX_THREAD_SIZE = 10 + + +class MultiStepReplySpamScorer: + model_name = GEMMA_2 + + def __init__(self): + oai_config = grox_config.get_oai_model(GEMMA_2) + self.oai_gemma4 = OaiSampler(oai_config) + + async def score(self, post: Post) -> MultiStepReplySpamResult: + convo = await self._to_convo(post) + output = await self._sample(convo) + result = await self._parse(post, output) + result = self._drop_if_newest_reply_not_flagged(post, result) + result = self._drop_other_author_posts(post, result) + result = self._drop_singleton(post, result) + Metrics.counter("multi_step_reply_spam.scored.count").add( + 1, attributes={"is_spam": bool(result.spam_post_ids)} + ) + return result + + @classmethod + def _drop_if_newest_reply_not_flagged( + cls, post: Post, result: MultiStepReplySpamResult + ) -> MultiStepReplySpamResult: + if not result.spam_post_ids or post.id in result.spam_post_ids: + return result + logger.info( + f"Dropping multi-step-reply-spam flags for post {post.id}: newest reply not in flagged set " + f"{result.spam_post_ids} reason={result.reason!r}" + ) + Metrics.counter("multi_step_reply_spam.newest_not_flagged_dropped.count").add(1) + return MultiStepReplySpamResult(reason=result.reason) + + @classmethod + def _drop_other_author_posts( + cls, post: Post, result: MultiStepReplySpamResult + ) -> MultiStepReplySpamResult: + if not result.spam_post_ids or not post.user: + return result + author_by_id = { + p.id: (p.user.id if p.user else None) for p in cls._thread_posts(post) + } + kept = [ + pid for pid in result.spam_post_ids if author_by_id.get(pid) == post.user.id + ] + if len(kept) == len(result.spam_post_ids): + return result + logger.info( + f"Dropping other-author posts from multi-step-reply-spam flags for post {post.id}: " + f"{[pid for pid in result.spam_post_ids if pid not in kept]} kept={kept} reason={result.reason!r}" + ) + Metrics.counter("multi_step_reply_spam.other_author_dropped.count").add(1) + return MultiStepReplySpamResult(spam_post_ids=kept, reason=result.reason) + + @classmethod + def _drop_singleton( + cls, post: Post, result: MultiStepReplySpamResult + ) -> MultiStepReplySpamResult: + if len(result.spam_post_ids) != 1: + return result + logger.info( + f"Dropping singleton multi-step-reply-spam flag for post {post.id}: {result.spam_post_ids[0]} reason={result.reason!r}" + ) + Metrics.counter("multi_step_reply_spam.singleton_dropped.count").add(1) + return MultiStepReplySpamResult(reason=result.reason) + + @staticmethod + def _thread_posts(post: Post) -> list[Post]: + ancestors = post.ancestors or [] + if len(ancestors) > MAX_THREAD_SIZE: + ancestors = ( + ancestors[: MAX_THREAD_SIZE // 2] + ancestors[-(MAX_THREAD_SIZE // 2) :] + ) + return ancestors + [post] + + async def _to_convo(self, post: Post) -> Conversation: + convo = Conversation(conversation_id=uuid.uuid4().hex) + convo.messages.append( + Message(role=Role.SYSTEM, content=[multi_step_reply_spam_system_prompt()]) + ) + + posts = self._thread_posts(post) + last = len(posts) - 1 + content: list = [] + content.append( + "\n# Reply Author Info\n\n" + "The user profile below describes the author of the reply being evaluated " + "(the final post of the thread).\n" + ) + content.extend(UserRenderer.render(post.user)) + content.append( + "\n\n# Thread\n\nThe reply thread is listed below; each post is labeled with a `#### Post N` index. " + "Post 0 is the original (root) post and the highest-numbered post is the reply being evaluated. " + "Refer to posts only by their integer index.\n\n" + ) + for i, p in enumerate(posts): + if i == 0: + tag = " (root post — never flag)" + elif i == last: + tag = " (the reply being evaluated)" + else: + tag = "" + content.append(f"\n\n#### Post {i}{tag}\n\n") + content.extend( + PostRenderer.render( + p, + max_media=MAX_MEDIA_PER_POST, + include_follower_count=True, + include_bio=i == 0, + ) + ) + content.append("\n\n------\n\n") + convo.messages.append(Message(role=Role.HUMAN, content=content)) + return convo + + async def _sample(self, convo: Conversation) -> str: + return await self.oai_gemma4.sample( + convo.to_openai_messages(), conversation_id=convo.conversation_id + ) + + async def _clean_output(self, output: str) -> str: + if output.endswith("<|eos|>"): + output = output.removesuffix("<|eos|>") + output = output.strip() + if output.startswith("```json"): + output = output[7:] + elif output.startswith("```"): + output = output[3:] + if output.endswith("```"): + output = output[:-3] + output = output.strip() + return output + + @staticmethod + def _resolve_post_ids(posts: list[Post], raw_values: list) -> list[str]: + id_by_index = {i: p.id for i, p in enumerate(posts)} + candidate_ids = {p.id for p in posts[1:]} + out: list[str] = [] + for v in raw_values: + s = str(v).strip() + if not s: + continue + if s.lstrip("-").isdigit(): + n = int(s) + if 0 < n < len(posts): + pid = id_by_index[n] + if pid not in out: + out.append(pid) + continue + if s in candidate_ids and s not in out: + out.append(s) + return out + + async def _parse(self, post: Post, output: str) -> MultiStepReplySpamResult: + posts = self._thread_posts(post) + match = re.search(r"\{.*\}", output, re.DOTALL) + raw_json = ( + match.group(0).strip() + if (match and "spam_post" in match.group(0)) + else output + ) + cleaned = await self._clean_output(raw_json) + + data = None + try: + data = json.loads(cleaned) + except Exception: + try: + repaired = json_repair.repair_json(cleaned, return_objects=True) + if isinstance(repaired, dict): + data = repaired + Metrics.counter("multi_step_reply_spam.json_repaired.count").add(1) + except Exception: + data = None + + if isinstance(data, dict): + raw_values = data.get("spam_post_indexes") + if raw_values is None: + raw_values = data.get("spam_post_ids") + if raw_values is None: + raw_values = [] + if not isinstance(raw_values, list): + raw_values = [raw_values] + reason = data.get("reason") or "" + return MultiStepReplySpamResult( + spam_post_ids=self._resolve_post_ids(posts, raw_values), reason=reason + ) + + arr_match = re.search( + r'"spam_post_(?:indexes|ids)"\s*:\s*\[([^\]]*)\]', cleaned, re.DOTALL + ) + if arr_match is None: + logger.error(f"Invalid multi-step reply spam output format: {output}") + Metrics.counter("multi_step_reply_spam.invalid.count").add(1) + raise ValueError(f"Invalid output: {output}") + raw_values = re.findall(r"\d+", arr_match.group(1)) + reason = "" + reason_match = re.search(r'"reason":\s*"((?:[^"\\]|\\.)*)"', cleaned, re.DOTALL) + if reason_match: + reason = reason_match.group(1) + return MultiStepReplySpamResult( + spam_post_ids=self._resolve_post_ids(posts, raw_values), reason=reason + ) diff --git a/grox/flows/reply_spam/classifier_reply_ranking.py b/grox/flows/reply_spam/classifier_reply_ranking.py index 1166a127..a4d08486 100644 --- a/grox/flows/reply_spam/classifier_reply_ranking.py +++ b/grox/flows/reply_spam/classifier_reply_ranking.py @@ -64,6 +64,7 @@ async def _to_convo(self, post: Post, non_reasoning: bool = False) -> Conversati role=Role.HUMAN, separator=THINKING_CONTROL_START + SEPARATOR, include_signals=True, + include_follower_count=True, ) ) if non_reasoning: diff --git a/grox/flows/reply_spam/classifier_simple_reply_scorer.py b/grox/flows/reply_spam/classifier_simple_reply_scorer.py index bf3b1798..b224dee8 100644 --- a/grox/flows/reply_spam/classifier_simple_reply_scorer.py +++ b/grox/flows/reply_spam/classifier_simple_reply_scorer.py @@ -20,7 +20,9 @@ async def _to_convo(self, post: Post) -> Conversation: system_prompt = reply_scoring_system_simple_prompt(100_000) convo.messages.append(Message(role=Role.SYSTEM, content=[system_prompt])) convo.messages.append( - ThreadRenderer.render(post, role=Role.HUMAN, include_signals=True) + ThreadRenderer.render( + post, role=Role.HUMAN, include_signals=True, include_follower_count=True + ) ) return convo diff --git a/grox/flows/reply_spam/constants.py b/grox/flows/reply_spam/constants.py index 5b60b6f6..1a74a667 100644 --- a/grox/flows/reply_spam/constants.py +++ b/grox/flows/reply_spam/constants.py @@ -1,5 +1,5 @@ -POST_STREAM = "post_stream" +REPLY_RANKING = "reply_ranking" REPLY_RANKING_RECOVERY = "reply_ranking_recovery" -TOPIC_UNIFIED_POSTS = "content-understanding-realtime-unified-posts" +TOPIC_UNIFIED_POSTS_V3 = "content-understanding-realtime-unified-posts-v3" TOPIC_REPLY_RANKING_RECOVERY = "reply_ranking_annotation_recovery_v2" GEMMA_2 = "oai-gemma4-26b-2" diff --git a/grox/flows/reply_spam/generators.py b/grox/flows/reply_spam/generators.py index 20a1d1e7..a8fa683a 100644 --- a/grox/flows/reply_spam/generators.py +++ b/grox/flows/reply_spam/generators.py @@ -3,32 +3,38 @@ from grox.flows.reply_spam.plan_spam_comment import PlanSpamComment from grox.flows.reply_spam.plan_reply_ranking import PlanReplyRanking from grox.flows.reply_spam.plan_coordinated_spam import PlanCoordinatedSpam +from grox.flows.reply_spam.plan_multi_step_reply_spam import PlanMultiStepReplySpam from grox.core.registry import register from grox.flows.reply_spam.constants import ( - POST_STREAM, + REPLY_RANKING, REPLY_RANKING_RECOVERY, TOPIC_REPLY_RANKING_RECOVERY, - TOPIC_UNIFIED_POSTS, + TOPIC_UNIFIED_POSTS_V3, ) @register -class PostStreamTaskGenerator(StreamTaskGenerator): - TASK_GENERATOR_TYPE = POST_STREAM +class ReplyRankingTaskGenerator(StreamTaskGenerator): + TASK_GENERATOR_TYPE = REPLY_RANKING PLANS_TO_INJECT = { - PlanSpamComment.KEY, PlanReplyRanking.KEY, + PlanSpamComment.KEY, PlanCoordinatedSpam.KEY, + PlanMultiStepReplySpam.KEY, } def _get_loader(self): - return KafkaPostLoader(TOPIC_UNIFIED_POSTS) + return KafkaPostLoader(TOPIC_UNIFIED_POSTS_V3) @register class ReplyRankingRecoveryTaskGenerator(StreamTaskGenerator): TASK_GENERATOR_TYPE = REPLY_RANKING_RECOVERY - PLANS_TO_INJECT = {PlanReplyRanking.KEY} + PLANS_TO_INJECT = { + PlanReplyRanking.KEY, + PlanSpamComment.KEY, + PlanMultiStepReplySpam.KEY, + } def _get_loader(self): return KafkaPostLoader(TOPIC_REPLY_RANKING_RECOVERY) diff --git a/grox/flows/reply_spam/plan_multi_step_reply_spam.py b/grox/flows/reply_spam/plan_multi_step_reply_spam.py new file mode 100644 index 00000000..beae75e4 --- /dev/null +++ b/grox/flows/reply_spam/plan_multi_step_reply_spam.py @@ -0,0 +1,29 @@ +from grox.core.plans.plan import Plan +from grox.core.registry import register +from grox.core.tasks.task_media import TaskMediaHydration +from grox.flows.reply_spam.task_multi_step_reply_spam import ( + TaskMultiStepReplySpamDetection, + TaskMultiStepReplySpamFilter, + TaskWriteMultiStepReplySpamReplyRanking, +) + + +@register +class PlanMultiStepReplySpam(Plan): + KEY = "multi_step_reply_spam" + + TASKS = { + "task_multi_step_reply_spam_filter": TaskMultiStepReplySpamFilter, + "task_media_hydration": TaskMediaHydration, + "task_multi_step_reply_spam_detection": TaskMultiStepReplySpamDetection, + "task_write_multi_step_reply_spam_reply_ranking": TaskWriteMultiStepReplySpamReplyRanking, + } + + TASK_DEPENDENCIES = { + "task_multi_step_reply_spam_filter": set(), + "task_media_hydration": {"task_multi_step_reply_spam_filter"}, + "task_multi_step_reply_spam_detection": {"task_media_hydration"}, + "task_write_multi_step_reply_spam_reply_ranking": { + "task_multi_step_reply_spam_detection" + }, + } diff --git a/grox/flows/reply_spam/plan_reply_ranking.py b/grox/flows/reply_spam/plan_reply_ranking.py index c539a1e4..81a3c1ba 100644 --- a/grox/flows/reply_spam/plan_reply_ranking.py +++ b/grox/flows/reply_spam/plan_reply_ranking.py @@ -4,9 +4,6 @@ from grox.core.tasks.task_media import TaskMediaHydration from grox.flows.reply_spam.task_filter import TaskReplyRankingFilter from grox.flows.reply_spam.task_rank_replies import TaskRankReplies -from grox.flows.reply_spam.task_rate_limit import ( - TaskRateLimitReplyRankingAnnotationWithPost, -) @register @@ -15,7 +12,6 @@ class PlanReplyRanking(Plan): TASKS = { "task_reply_ranking_filter": TaskReplyRankingFilter, - "task_reply_ranking_annotation_rate_limit": TaskRateLimitReplyRankingAnnotationWithPost, "task_media_hydration": TaskMediaHydration, "task_rank_replies": TaskRankReplies, "task_write_reply_ranking_manhattan": TaskWriteReplyRankingManhattan, @@ -23,8 +19,7 @@ class PlanReplyRanking(Plan): TASK_DEPENDENCIES = { "task_reply_ranking_filter": set(), - "task_reply_ranking_annotation_rate_limit": {"task_reply_ranking_filter"}, - "task_media_hydration": {"task_reply_ranking_annotation_rate_limit"}, + "task_media_hydration": {"task_reply_ranking_filter"}, "task_rank_replies": {"task_media_hydration"}, "task_write_reply_ranking_manhattan": {"task_rank_replies"}, } diff --git a/grox/flows/reply_spam/plan_spam_comment.py b/grox/flows/reply_spam/plan_spam_comment.py index df5aac10..63650225 100644 --- a/grox/flows/reply_spam/plan_spam_comment.py +++ b/grox/flows/reply_spam/plan_spam_comment.py @@ -4,9 +4,6 @@ from grox.core.tasks.task_media import TaskMediaHydration from grox.flows.reply_spam.task_filter import TaskSpamFilter from grox.flows.reply_spam.task_spam_detection import TaskSpamDetection -from grox.flows.reply_spam.task_rate_limit import ( - TaskRateLimitReplySpamAnnotationWithPost, -) @register @@ -15,7 +12,6 @@ class PlanSpamComment(Plan): TASKS = { "task_spam_filter": TaskSpamFilter, - "task_reply_spam_annotation_rate_limit": TaskRateLimitReplySpamAnnotationWithPost, "task_media_hydration": TaskMediaHydration, "task_spam_detection": TaskSpamDetection, "task_write_reply_ranking_manhattan": TaskWriteReplyRankingManhattan, @@ -23,8 +19,7 @@ class PlanSpamComment(Plan): TASK_DEPENDENCIES = { "task_spam_filter": set(), - "task_reply_spam_annotation_rate_limit": {"task_spam_filter"}, - "task_media_hydration": {"task_reply_spam_annotation_rate_limit"}, + "task_media_hydration": {"task_spam_filter"}, "task_spam_detection": {"task_media_hydration"}, "task_write_reply_ranking_manhattan": {"task_spam_detection"}, } diff --git a/grox/flows/reply_spam/prompts.py b/grox/flows/reply_spam/prompts.py index b4c91459..169c30e2 100644 --- a/grox/flows/reply_spam/prompts.py +++ b/grox/flows/reply_spam/prompts.py @@ -28,3 +28,8 @@ def reply_scoring_system_simple_prompt(large_account_follower_threshold: int) -> @cache def coordinated_spam_system_prompt() -> str: return _env.get_template("coordinated_spam_system.j2").render() + + +@cache +def multi_step_reply_spam_system_prompt() -> str: + return _env.get_template("multi_step_reply_spam_system.j2").render() diff --git a/grox/flows/reply_spam/state_multi_step_reply_spam.py b/grox/flows/reply_spam/state_multi_step_reply_spam.py new file mode 100644 index 00000000..f459c129 --- /dev/null +++ b/grox/flows/reply_spam/state_multi_step_reply_spam.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + +from pydantic import BaseModel, Field + + +class MultiStepReplySpamResult(BaseModel): + spam_post_ids: list[str] = Field(default_factory=list) + reason: str = "" + + +@dataclass +class MultiStepReplySpamState: + result: MultiStepReplySpamResult | None = None diff --git a/grox/flows/reply_spam/strato_loader.py b/grox/flows/reply_spam/strato_loader.py index f882d985..dae31294 100644 --- a/grox/flows/reply_spam/strato_loader.py +++ b/grox/flows/reply_spam/strato_loader.py @@ -1,3 +1,4 @@ +import asyncio import logging from strato_http.queries.data_types import ( @@ -5,6 +6,10 @@ ReplyRankingScoreKafka, ) from strato_http.queries.reply_ranking_score import StratoReplyRankingScore +from strato_http.queries.reply_ranking_score_cache import ( + StratoReplyRankingScoreCacheAtla, + StratoReplyRankingScoreCachePdxa, +) from strato_http.queries.reply_ranking_score_kafka_v2 import ( StratoReplyRankingScoreV2Kafka, ) @@ -12,16 +17,32 @@ logger = logging.getLogger(__name__) +_CACHE_FANOUT_SCORE_MAX = 1.0 + class ReplyRankingScoreStratoLoader: strato = StratoReplyRankingScore() + strato_cache_atla = StratoReplyRankingScoreCacheAtla() + strato_cache_pdxa = StratoReplyRankingScoreCachePdxa() reply_ranking_v2_kafka_strato = StratoReplyRankingScoreV2Kafka() + @classmethod + async def fetch_reply_ranking_score(cls, post_id: str) -> ReplyRankingScore | None: + return await cls.strato.fetch(int(post_id)) + @classmethod async def save_reply_ranking_score( cls, post_id: str, reply_ranking_score: ReplyRankingScore ): await cls.strato.put(int(post_id), reply_ranking_score) + if ( + reply_ranking_score.score is not None + and reply_ranking_score.score <= _CACHE_FANOUT_SCORE_MAX + ): + await asyncio.gather( + cls.strato_cache_atla.put(int(post_id), reply_ranking_score), + cls.strato_cache_pdxa.put(int(post_id), reply_ranking_score), + ) @classmethod async def save_reply_ranking_kafka_v2( diff --git a/grox/flows/reply_spam/task_filter.py b/grox/flows/reply_spam/task_filter.py index 7b39d7d2..91467f38 100644 --- a/grox/flows/reply_spam/task_filter.py +++ b/grox/flows/reply_spam/task_filter.py @@ -14,7 +14,7 @@ class TaskSpamFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 15000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 120000 @override @classmethod @@ -94,7 +94,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskCoordinatedSpamFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 1000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 5000 FILTER_NAME = "coordinated_spam" @override @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 15000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 120000 @override @classmethod diff --git a/grox/flows/reply_spam/task_multi_step_reply_spam.py b/grox/flows/reply_spam/task_multi_step_reply_spam.py new file mode 100644 index 00000000..aa72bdd2 --- /dev/null +++ b/grox/flows/reply_spam/task_multi_step_reply_spam.py @@ -0,0 +1,206 @@ +import asyncio +import logging +from typing import override + + +from grox.core.constants import GORK_USER_ID, GROK_USER_ID +from grox.core.data_loaders.data_types import Post +from grox.core.data_loaders.strato_loader import UserStratoLoader +from grox.core.schedules.types import TaskContext +from grox.core.tasks.disable_rules import DisableTaskForNonProd +from grox.core.tasks.task import Task, TaskResultCategory, TaskWithPost +from grox.core.tasks.task_filters import TaskFilterWithPost +from grox.flows.reply_spam.classifier_multi_step_reply_spam import ( + MultiStepReplySpamScorer, +) +from grox.flows.reply_spam.state_multi_step_reply_spam import MultiStepReplySpamState +from grox.flows.reply_spam.strato_loader import ReplyRankingScoreStratoLoader +from grox.flows.reply_spam.task_write import _apply_reply_spam_label +from monitor.metrics import Metrics +from strato_http.queries.data_types import ReplyRankingScore, ReplyRankingScoreKafka + +logger = logging.getLogger(__name__) + + +class TaskMultiStepReplySpamFilter(TaskFilterWithPost): + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 1000 + FILTER_NAME = "multi_step_reply_spam" + + @override + @classmethod + async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: + if not post.ancestors: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "not_reply"} + ) + return False + if any(a is None for a in post.ancestors): + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "deleted_ancestor"} + ) + return False + if not post.user: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "no_user"} + ) + return False + if post.user.id == GROK_USER_ID: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "is_grok_reply"} + ) + return False + if post.user.id == GORK_USER_ID: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "is_gork_reply"} + ) + return False + if not post.ancestors[-1].user: + Metrics.counter("task.filter.skipped.count").add( + 1, + attributes={ + "filter": cls.FILTER_NAME, + "reason": "previous_post_no_user", + }, + ) + return False + if not post.ancestors[0].user: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "root_post_no_user"} + ) + return False + if post.user.id == post.ancestors[0].user.id: + logger.info( + f"Skipping multi-step reply spam since the replier is same as reply root post {post.id}" + ) + Metrics.counter("task.filter.skipped.count").add( + 1, + attributes={ + "filter": cls.FILTER_NAME, + "reason": "same_user_reply_as_root", + }, + ) + return False + if len(post.ancestors) < 2: + logger.info( + f"Skipping multi-step reply spam since the reply thread is not more than two level deep {post.id}" + ) + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "one_level_deep"} + ) + return False + if post.ancestors[-1].user.id != post.user.id: + Metrics.counter("task.filter.skipped.count").add( + 1, + attributes={ + "filter": cls.FILTER_NAME, + "reason": "parent_not_same_author", + }, + ) + return False + root_user_follower_count = post.ancestors[0].user.follower_count or 0 + if root_user_follower_count < cls.FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION: + Metrics.counter("task.filter.skipped.count").add( + 1, attributes={"filter": cls.FILTER_NAME, "reason": "low_blast_radius"} + ) + return False + is_high_page_rank, is_grey_badge = await asyncio.gather( + UserStratoLoader.is_high_page_rank_v2_user(post.user.id), + UserStratoLoader.is_grey_badge_user(post.user.id), + ) + if is_high_page_rank or is_grey_badge: + logger.info( + f"Skipping multi-step reply spam for post {post.id} user {post.user.id} " + f"(high_page_rank_v2={is_high_page_rank}, grey_badge={is_grey_badge})" + ) + Metrics.counter("task.filter.skipped.count").add( + 1, + attributes={ + "filter": cls.FILTER_NAME, + "reason": "high_page_rank_or_grey_badge", + }, + ) + return False + Metrics.counter("task.multi_step_reply_spam_filter.eligible.count").add(1) + return True + + +class TaskMultiStepReplySpamDetection(TaskWithPost): + scorer = MultiStepReplySpamScorer() + + @classmethod + async def exec(cls, ctx: TaskContext) -> TaskResultCategory: + return await Task.exec.__wrapped__(cls, ctx) + + @classmethod + async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: + result = await cls.scorer.score(post) + ctx.state(MultiStepReplySpamState).result = result + if result.spam_post_ids: + logger.info( + f"Multi-step reply spam found for post {post.id}: spam_post_ids={result.spam_post_ids} reason={result.reason!r}" + ) + Metrics.counter("task.multi_step_reply_spam_detection.positive.count").add( + 1 + ) + else: + Metrics.counter("task.multi_step_reply_spam_detection.negative.count").add( + 1 + ) + + +class TaskWriteMultiStepReplySpamReplyRanking(Task): + DISABLE_RULES = [DisableTaskForNonProd] + + @classmethod + async def _exec(cls, ctx: TaskContext) -> None: + post = ctx.payload.post + if not post: + return + result = ctx.state(MultiStepReplySpamState).result + if result is None: + Metrics.counter("task.write_multi_step_reply_spam.skipped.count").add( + 1, attributes={"reason": "no_results"} + ) + return + if not result.spam_post_ids: + Metrics.counter("task.write_multi_step_reply_spam.skipped.count").add( + 1, attributes={"reason": "no_spam"} + ) + return + + thread_posts = (post.ancestors or []) + [post] + spam_ids = set(result.spam_post_ids) + reasoning = (result.reason or "")[-500:] + + Metrics.counter("task.write_multi_step_reply_spam.intaken.count").add(1) + for p in thread_posts: + if p.id not in spam_ids: + continue + if not p.user: + Metrics.counter( + "task.write_multi_step_reply_spam.skipped_post.count" + ).add(1, attributes={"reason": "no_author"}) + continue + await cls._mark_spam(p.id, p.user.id, reasoning) + Metrics.counter("task.write_multi_step_reply_spam.success.count").add(1) + + @classmethod + async def _mark_spam(cls, post_id: str, author_id: int, reasoning: str) -> None: + await _apply_reply_spam_label(post_id, author_id) + + await ReplyRankingScoreStratoLoader.save_reply_ranking_score( + post_id=post_id, + reply_ranking_score=ReplyRankingScore( + score=0.0, reasoning=reasoning[-500:] + ), + ) + + await ReplyRankingScoreStratoLoader.save_reply_ranking_kafka_v2( + post_id=post_id, + reply_ranking_score_kafka=ReplyRankingScoreKafka( + postId=int(post_id), score=0.0, reasoning=reasoning + ), + ) + logger.info( + f"Published multi-step reply spam reply ranking score 0 for post {post_id}" + ) diff --git a/grox/flows/reply_spam/task_rate_limit.py b/grox/flows/reply_spam/task_rate_limit.py index 636bd427..9caa8065 100644 --- a/grox/flows/reply_spam/task_rate_limit.py +++ b/grox/flows/reply_spam/task_rate_limit.py @@ -3,16 +3,6 @@ from grox.core.tasks.task_rate_limit import TaskTTLDedupeWithPost -class TaskRateLimitReplySpamAnnotationWithPost(TaskTTLDedupeWithPost): - DEDUPE_CACHE = TTLCache(maxsize=10_000, ttl=60) - DEDUPE_NAME = "reply spam" - - -class TaskRateLimitReplyRankingAnnotationWithPost(TaskTTLDedupeWithPost): - DEDUPE_CACHE = TTLCache(maxsize=10_000, ttl=60) - DEDUPE_NAME = "reply ranking" - - class TaskRateLimitCoordinatedSpamAnnotationWithPost(TaskTTLDedupeWithPost): DEDUPE_CACHE = TTLCache(maxsize=10_000, ttl=60) DEDUPE_NAME = "coordinated spam" diff --git a/grox/flows/reply_spam/task_write.py b/grox/flows/reply_spam/task_write.py index 1149af26..a76d4e53 100644 --- a/grox/flows/reply_spam/task_write.py +++ b/grox/flows/reply_spam/task_write.py @@ -86,9 +86,6 @@ async def _exec(cls, ctx: TaskContext) -> None: Metrics.counter("task.write_reply_ranking_manhattan.intaken.count").add(1) try: await cls._publish_to_reply_ranking_manhattan(post, result) - logger.info( - f"Published reply ranking post to manhattan: {post.id=} user_id={post.user.id if post.user else None}" - ) except Exception: Metrics.counter("task.write_reply_ranking_manhattan.failed.count").add(1) logger.error( @@ -113,26 +110,45 @@ async def _publish_to_reply_ranking_manhattan( f"Missing user id [_publish_to_reply_ranking_manhattan] {reasoning=} {post.id=} {score=}" ) + existing = await ReplyRankingScoreStratoLoader.fetch_reply_ranking_score( + post.id + ) + if ( + existing is not None + and existing.score is not None + and score >= existing.score + ): + Metrics.counter("task.write_reply_ranking_manhattan.skipped.count").add( + 1, attributes={"reason": "not_lower_than_existing"} + ) + logger.info( + f"[_publish_to_reply_ranking_manhattan] skipping write: new {score=} >= existing={existing.score} {post.id=}" + ) + return + if score == 0.0: await _apply_reply_spam_label(post.id, post.user.id if post.user else None) - await ReplyRankingScoreStratoLoader.save_reply_ranking_score( + await ReplyRankingScoreStratoLoader.save_reply_ranking_kafka_v2( post_id=post.id, - reply_ranking_score=ReplyRankingScore( - score=score, reasoning=reasoning[-500:] + reply_ranking_score_kafka=ReplyRankingScoreKafka( + postId=int(post.id), score=score, reasoning=reasoning[-500:] ), ) - await ReplyRankingScoreStratoLoader.save_reply_ranking_kafka_v2( + await ReplyRankingScoreStratoLoader.save_reply_ranking_score( post_id=post.id, - reply_ranking_score_kafka=ReplyRankingScoreKafka( - postId=int(post.id), score=score, reasoning=reasoning[-500:] + reply_ranking_score=ReplyRankingScore( + score=score, reasoning=reasoning[-500:] ), ) Metrics.counter("task.write_reply_ranking_manhattan.success.count").add( 1, attributes={"column": "reply_ranking"} ) + logger.info( + f"Published reply ranking post to manhattan: {post.id=} user_id={post.user.id if post.user else None} {score=}" + ) class TaskWriteCoordinatedSpamReplyRanking(Task): diff --git a/grox/flows/upa/constants.py b/grox/flows/upa/constants.py index a3281ec1..513bf459 100644 --- a/grox/flows/upa/constants.py +++ b/grox/flows/upa/constants.py @@ -1,9 +1,11 @@ POST_MIN_TRACTION_STREAM_FOR_GROX = "post_min_traction_stream_for_grox" POST_STREAM_RECOVERY = "post_stream_recovery" POST_SAFETY_STREAM = "post_safety_stream" +POST_PRIORITY_STREAM = "post_priority_stream" TOPIC_MIN_TRACTION = ( "content-understanding-realtime-unified-posts-min-traction-for-grox" ) TOPIC_RECOVERY = "content_understanding_realtime_unified_posts_recovery_v2" TOPIC_POPULAR = "content-understanding-realtime-unified-posts-popular" +TOPIC_PRIORITY = "content_understanding_realtime_unified_posts_priority" GEMMA_UPA = "oai-gemma4-26b-upa" diff --git a/grox/flows/upa/generators.py b/grox/flows/upa/generators.py index b9d3190d..92a9ca85 100644 --- a/grox/flows/upa/generators.py +++ b/grox/flows/upa/generators.py @@ -5,10 +5,12 @@ from grox.core.registry import register from grox.flows.upa.constants import ( POST_MIN_TRACTION_STREAM_FOR_GROX, + POST_PRIORITY_STREAM, POST_SAFETY_STREAM, POST_STREAM_RECOVERY, TOPIC_MIN_TRACTION, TOPIC_POPULAR, + TOPIC_PRIORITY, TOPIC_RECOVERY, ) @@ -38,3 +40,12 @@ class PostSafetyStreamTaskGenerator(StreamTaskGenerator): def _get_loader(self): return KafkaPostLoader(TOPIC_POPULAR) + + +@register +class PostPriorityStreamTaskGenerator(StreamTaskGenerator): + TASK_GENERATOR_TYPE = POST_PRIORITY_STREAM + PLANS_TO_INJECT = {PlanInitialBanger.KEY} + + def _get_loader(self): + return KafkaPostLoader(TOPIC_PRIORITY) diff --git a/grox/libs/grok_sampler/llm.py b/grox/libs/grok_sampler/llm.py index 3fca7609..57ae1173 100644 --- a/grox/libs/grok_sampler/llm.py +++ b/grox/libs/grok_sampler/llm.py @@ -2,7 +2,7 @@ import logging import traceback from abc import ABC, abstractmethod -from typing import Generic, TypeVar, AsyncGenerator +from typing import Callable, Generic, TypeVar, AsyncGenerator from contextlib import aclosing from monitor.logging import Logging @@ -69,6 +69,7 @@ async def _sample_streaming(self, query: T, **kwargs) -> AsyncGenerator[str, Non keep_separator = kwargs.get("keep_separator", False) separator = kwargs.get("separator", SEPARATOR) log_prompt = kwargs.get("log_prompt", False) + stop_predicate: Callable[[str], bool] | None = kwargs.get("stop_predicate") request = await self._get_sample_request(query, separator, **kwargs) prompt_len = self._prompt_len(request) logger.info(f"Started sampling request, prompt_len={prompt_len}") @@ -84,19 +85,32 @@ async def _sample_streaming(self, query: T, **kwargs) -> AsyncGenerator[str, Non tokens_received = 0 start = time.perf_counter() try: - async for tok in self._sample_streaming_raw(request, **kwargs): - if tokens_received == 0: - ttft = time.perf_counter() - start - Metrics.histogram("llm.sample.ttft").record( - ttft, attributes=attributes + async with aclosing( + self._sample_streaming_raw(request, **kwargs) + ) as raw_stream: + async for tok in raw_stream: + if tokens_received == 0: + ttft = time.perf_counter() - start + Metrics.histogram("llm.sample.ttft").record( + ttft, attributes=attributes + ) + logger.info(f"Time to first token: {ttft:.3f}") + tokens_received += 1 + Metrics.counter("llm.sample.token.count").add( + 1, attributes=attributes ) - logger.info(f"Time to first token: {ttft:.3f}") - tokens_received += 1 - Metrics.counter("llm.sample.token.count").add(1, attributes=attributes) - if not keep_separator: - tok.text = tok.text.replace(separator, "") - resp += tok.text - yield tok.text + if not keep_separator: + tok.text = tok.text.replace(separator, "") + resp += tok.text + yield tok.text + if stop_predicate is not None and stop_predicate(resp): + logger.info( + f"Response complete after {tokens_received} tokens, cancelling the stream" + ) + Metrics.counter("llm.sample.early_stop.count").add( + 1, attributes=attributes + ) + break logger.info( f"Finished sampling request in {time.perf_counter() - start:.2f} seconds" ) diff --git a/grox/libs/kafka_cli/consumer.py b/grox/libs/kafka_cli/consumer.py index cfdef888..f8108f89 100644 --- a/grox/libs/kafka_cli/consumer.py +++ b/grox/libs/kafka_cli/consumer.py @@ -24,6 +24,12 @@ def __init__(self, config: KafkaConsumerConfig): self.group_id: str = config.group_id self._consumer: AIOKafkaConsumer | None = None + @property + def assigned_partitions(self) -> int: + if self._consumer is None: + return 0 + return len(self._consumer.assignment() or ()) + def _auth_mode(self) -> str: if not self.config.ssl: return "PLAINTEXT" diff --git a/grox/libs/kafka_cli/multi_region_consumer.py b/grox/libs/kafka_cli/multi_region_consumer.py index 6e99b1bb..06a569a4 100644 --- a/grox/libs/kafka_cli/multi_region_consumer.py +++ b/grox/libs/kafka_cli/multi_region_consumer.py @@ -116,6 +116,12 @@ async def _start_region(self, region: str, brokers: list[str]) -> AIOKafkaConsum ) return consumer + @property + def assigned_partitions(self) -> int: + return sum( + len(consumer.assignment() or ()) for consumer in self._consumers.values() + ) + async def stop(self): if self._region_retry_task is not None: self._region_retry_task.cancel() diff --git a/grox/libs/video_tools/image.py b/grox/libs/video_tools/image.py index 03d41365..4033ecf5 100644 --- a/grox/libs/video_tools/image.py +++ b/grox/libs/video_tools/image.py @@ -1,4 +1,5 @@ import io +import math from dataclasses import dataclass import logging import cv2 @@ -152,3 +153,127 @@ def pad_image(image_bytes: bytes) -> bytes: with Image.open(io.BytesIO(image_bytes)) as img: with img.convert("RGB") as image: return _padded_image(image) + + +_MOTION_REVEAL_MIN_FRAMES = 4 +_MOTION_REVEAL_MAX_INPUT_FRAMES = 32 +_MOTION_REVEAL_WORK_MAX_DIM = 640 +_MOTION_REVEAL_COVER_WEIGHT = 0.7 +_MOTION_REVEAL_GRID_FRAMES = 4 +_MOTION_REVEAL_GRID_COLUMNS = 2 +_MOTION_REVEAL_JPEG_QUALITY = 90 +_MOTION_REVEAL_MIN_RESIDUAL_P99 = 3.0 +_MOTION_REVEAL_MAX_RESIDUAL_P99 = 75.0 + + +def _decode_reveal_frame(frame_bytes: bytes) -> np.ndarray | None: + img = cv2.imdecode(np.frombuffer(frame_bytes, dtype=np.uint8), cv2.IMREAD_COLOR) + if img is None: + return None + h, w = img.shape[:2] + scale = _MOTION_REVEAL_WORK_MAX_DIM / max(h, w) + if scale < 1.0: + img = cv2.resize( + img, + (max(1, int(w * scale)), max(1, int(h * scale))), + interpolation=cv2.INTER_AREA, + ) + return img.astype(np.float32) + + +def _stretch_to_u8(arr: np.ndarray) -> np.ndarray: + lo = float(np.percentile(arr, 2)) + hi = float(np.percentile(arr, 98)) + if hi <= lo + 1e-3: + return np.zeros(arr.shape, dtype=np.uint8) + return np.clip((arr - lo) / (hi - lo) * 255.0, 0, 255).astype(np.uint8) + + +def _encode_reveal(bgr: np.ndarray) -> bytes | None: + ok, jpeg = cv2.imencode( + ".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), _MOTION_REVEAL_JPEG_QUALITY] + ) + return jpeg.tobytes() if ok else None + + +def _skin_fraction(bgr_u8: np.ndarray) -> float: + ycrcb = cv2.cvtColor(bgr_u8, cv2.COLOR_BGR2YCrCb) + mask = ( + (ycrcb[:, :, 0] > 60) + & (ycrcb[:, :, 1] > 135) + & (ycrcb[:, :, 1] < 175) + & (ycrcb[:, :, 2] > 80) + & (ycrcb[:, :, 2] < 130) + ) + return float(mask.mean()) + + +def build_motion_reveal_images(frame_jpegs: list[bytes]) -> list[bytes]: + if len(frame_jpegs) < _MOTION_REVEAL_MIN_FRAMES: + return [] + if len(frame_jpegs) > _MOTION_REVEAL_MAX_INPUT_FRAMES: + indices = np.linspace( + 0, len(frame_jpegs) - 1, _MOTION_REVEAL_MAX_INPUT_FRAMES + ).astype(int) + frame_jpegs = [frame_jpegs[i] for i in indices] + + decoded: list[np.ndarray] = [] + target_hw: tuple[int, int] | None = None + for frame_bytes in frame_jpegs: + img = _decode_reveal_frame(frame_bytes) + if img is None: + continue + if target_hw is None: + target_hw = (img.shape[0], img.shape[1]) + elif (img.shape[0], img.shape[1]) != target_hw: + img = cv2.resize( + img, (target_hw[1], target_hw[0]), interpolation=cv2.INTER_AREA + ) + decoded.append(img) + if len(decoded) < _MOTION_REVEAL_MIN_FRAMES: + return [] + + stack = np.stack(decoded, axis=0) + median = np.median(stack, axis=0) + residual = np.abs(stack - median[None]).mean(axis=-1) + per_frame_p99 = np.percentile(residual.reshape(len(decoded), -1), 99, axis=1) + if ( + not _MOTION_REVEAL_MIN_RESIDUAL_P99 + <= float(np.median(per_frame_p99)) + <= _MOTION_REVEAL_MAX_RESIDUAL_P99 + ): + return [] + motion_energy = residual.mean(axis=(1, 2)) + + unmixed = [ + _stretch_to_u8(frame - _MOTION_REVEAL_COVER_WEIGHT * median) for frame in stack + ] + skin = np.array([_skin_fraction(still) for still in unmixed]) + order = np.argsort(-(motion_energy * (0.5 + skin))) + + out: list[bytes] = [] + + best = int(order[0]) + pair = np.concatenate([stack[best].astype(np.uint8), unmixed[best]], axis=1) + encoded = _encode_reveal(pair) + if encoded: + out.append(encoded) + + encoded = _encode_reveal(unmixed[best]) + if encoded: + out.append(encoded) + + grid_indices = sorted(int(i) for i in order[:_MOTION_REVEAL_GRID_FRAMES]) + tiles = [unmixed[i] for i in grid_indices] + h, w = tiles[0].shape[:2] + cols = _MOTION_REVEAL_GRID_COLUMNS + rows = math.ceil(len(tiles) / cols) + canvas = np.zeros((rows * h, cols * w, 3), dtype=np.uint8) + for i, tile in enumerate(tiles): + row, col = divmod(i, cols) + canvas[row * h : (row + 1) * h, col * w : (col + 1) * w] = tile + encoded = _encode_reveal(canvas) + if encoded: + out.append(encoded) + + return out diff --git a/grox/libs/video_tools/video_frames.py b/grox/libs/video_tools/video_frames.py index bfb1b554..37d86e92 100644 --- a/grox/libs/video_tools/video_frames.py +++ b/grox/libs/video_tools/video_frames.py @@ -8,14 +8,21 @@ import cv2 import numpy as np from PIL import Image -from pydantic import BaseModel +from pydantic import Field, BaseModel from av.stream import Stream from av.container import InputContainer -from video_tools.image import resize_tile, enhance_image_with_clahe +from video_tools.image import ( + resize_tile, + enhance_image_with_clahe, + build_motion_reveal_images, +) logger = logging.getLogger(__name__) +_MOTION_REVEAL_DENSE_THRESHOLD = 12 +_MOTION_REVEAL_DENSE_SAMPLES = 16 + class VideoFrame(BaseModel): time_sec: float @@ -26,6 +33,7 @@ class VideoData(BaseModel): frames: list[VideoFrame] combined_bytes: bytes | None = None total_duration: float | None = None + motion_reveal_frames: list[bytes] = Field(default_factory=list) class VideoFramesExtractor: @@ -37,6 +45,7 @@ async def extract_frames( tile_size: int | None = None, enable_clahe: bool = False, include_combined_video_bytes: bool = True, + enable_motion_reveal: bool = False, ) -> VideoData: loop = asyncio.get_running_loop() return await loop.run_in_executor( @@ -47,6 +56,7 @@ async def extract_frames( tile_size, enable_clahe, include_combined_video_bytes, + enable_motion_reveal, ) @classmethod @@ -57,6 +67,7 @@ def _extract_frames( tile_size: int | None, enable_clahe: bool = False, include_combined_video_bytes: bool = True, + enable_motion_reveal: bool = False, ) -> VideoData: logger.info(f"Extracting maximum {max_frames} frames from video") @@ -72,6 +83,39 @@ def _extract_frames( total_duration = float(c_duration / av.time_base) sample_times = cls._sample_frames(total_duration, max_frames) frames = cls._extract_frames_at_times(container, sample_times) + + reveal_input = frames + if ( + enable_motion_reveal + and 0 < len(frames) < _MOTION_REVEAL_DENSE_THRESHOLD + and total_duration > 0 + ): + dense_times = [ + i * total_duration / _MOTION_REVEAL_DENSE_SAMPLES + for i in range(_MOTION_REVEAL_DENSE_SAMPLES) + ] + try: + dense_frames = cls._extract_frames_at_times(container, dense_times) + if len(dense_frames) > len(frames): + reveal_input = dense_frames + except Exception: + logger.warning( + "Failed to extract dense motion-reveal frames; using sampled frames", + exc_info=True, + ) + + motion_reveal_frames: list[bytes] = [] + if enable_motion_reveal: + try: + motion_reveal_frames = build_motion_reveal_images( + [frame.frame for frame in reveal_input] + ) + except Exception: + logger.warning( + "Failed to build motion-reveal stills; continuing with RGB frames only", + exc_info=True, + ) + for frame in frames: frame.frame = cls._process_frame(frame.frame, tile_size, enable_clahe) logger.info(f"Extracted {len(frames)} frames") @@ -81,7 +125,10 @@ def _extract_frames( else None ) return VideoData( - frames=frames, combined_bytes=combined_bytes, total_duration=total_duration + frames=frames, + combined_bytes=combined_bytes, + total_duration=total_duration, + motion_reveal_frames=motion_reveal_frames, ) @classmethod diff --git a/grox/main.py b/grox/main.py new file mode 100644 index 00000000..c9dc8892 --- /dev/null +++ b/grox/main.py @@ -0,0 +1,73 @@ +import signal +import asyncio +import logging + +from kerberos_cli.kerberos import KerberosRenewer + +from grox.core.engine import Engine +from grox.core.services.service import GrpcServer +from grox.core.dispatcher import Dispatcher +from grox.config.config import grox_config +from grox.core.schedules.init import init_metrics, init_proc +from grox.core.schedules.context import ( + cleanup, + new_context, + shutdown_context, + queue_connection_shutdown_context, +) + +logger = logging.getLogger(__name__) +shutdown = asyncio.Event() + + +def init_kerberos_renewer() -> KerberosRenewer | None: + keytab_path = grox_config.kerberos.keytab_path + principal = grox_config.kerberos.principal + if keytab_path is None or principal is None: + logger.warning("Kerberos is not enabled, skipping") + return None + return KerberosRenewer(keytab_path=keytab_path, principal=principal) + + +async def serve(): + await init_proc("main", defer_metrics=True) + logger.info("Starting grox server...") + context = new_context() + kerberos_renewer = init_kerberos_renewer() + engine = Engine(context) + dispatcher = Dispatcher(context) + grpc_server = GrpcServer() + + if kerberos_renewer is not None: + await kerberos_renewer.renew() + kerberos_renewer.start() + + await engine.start() + await dispatcher.start() + init_metrics("main") + await grpc_server.start() + + logger.info("Grox server started") + event_loop = asyncio.get_running_loop() + event_loop.add_signal_handler(signal.SIGINT, lambda: shutdown.set()) + event_loop.add_signal_handler(signal.SIGTERM, lambda: shutdown.set()) + + await shutdown.wait() + logger.warning("Grox server shutting down...") + queue_connection_shutdown_context(context) + await asyncio.sleep(grox_config.shutdown_drain_timeout) + + shutdown_context(context) + await asyncio.gather( + grpc_server.stop(), + dispatcher.stop(), + engine.stop(), + ) + if kerberos_renewer is not None: + kerberos_renewer.stop() + cleanup() + logger.warning("Grox server stopped") + + +if __name__ == "__main__": + asyncio.run(serve()) diff --git a/home-mixer/ads/mod.rs b/home-mixer/ads/mod.rs index 1727f777..3b9e4701 100644 --- a/home-mixer/ads/mod.rs +++ b/home-mixer/ads/mod.rs @@ -1,4 +1,5 @@ mod following_ad_blender; +mod multi_risk_blender; mod partition_organic_blender; mod safe_gap_blender; #[cfg(test)] @@ -7,6 +8,7 @@ mod time_gap_blender; pub(crate) mod util; pub use following_ad_blender::FollowingAdBlender; +pub use multi_risk_blender::MultiRiskAdsBlender; pub use partition_organic_blender::PartitionOrganicAdsBlender; pub use safe_gap_blender::SafeGapAdsBlender; pub use time_gap_blender::{TimeGapAdsBlender, TimeGapConfig}; diff --git a/home-mixer/ads/multi_risk_blender.rs b/home-mixer/ads/multi_risk_blender.rs new file mode 100644 index 00000000..b7d32649 --- /dev/null +++ b/home-mixer/ads/multi_risk_blender.rs @@ -0,0 +1,332 @@ +use super::util::*; +use super::AdsBlender; +use crate::params::RESULT_SIZE; +use xai_home_mixer_proto::{feed_item, FeedItem, ScoredPost}; +use xai_post_text::TokenSequence; +use xai_recsys_proto::AdIndexInfo; +use xai_stats_receiver::global_stats_receiver; + +const ENFORCEMENT_METRIC: &str = "MultiRisk.enforcement"; +const SLOT_OUTCOME_METRIC: &str = "MultiRisk.slot_outcome"; +const SERVING_LIMITATION_METRIC: &str = "MultiRisk.serving_limitation"; + +pub struct MultiRiskAdsBlender; + +impl AdsBlender for MultiRiskAdsBlender { + fn blend_inner(&self, scored_posts: Vec, ads: Vec) -> Vec { + blend_impl(scored_posts, ads, MIN_POSTS_FOR_ADS) + } +} + +pub(crate) fn blend_impl( + scored_posts: Vec, + ads: Vec, + min_posts: usize, +) -> Vec { + let n = scored_posts.len(); + + if ads.is_empty() || n < min_posts { + emit_serving_limitation(if ads.is_empty() { + "no_ads" + } else { + "too_few_posts" + }); + return posts_to_feed_items(scored_posts); + } + + let spacing = compute_spacing(&ads); + let spacing_cap = n + .saturating_sub(1) + .checked_div(spacing.requested) + .unwrap_or(0); + + let safe_count = scored_posts + .iter() + .filter(|p| !is_medium_risk(p) && !is_high_risk(p)) + .count(); + let medium_count = scored_posts.iter().filter(|p| is_medium_risk(p)).count(); + let max_safe_slots = safe_count / 2; + let max_medium_slots = medium_count / 2; + let num_ads = ads.len(); + let safe_budget = num_ads.min(spacing_cap).min(max_safe_slots); + emit_serving_limitation(serving_limitation(num_ads, spacing_cap, max_safe_slots)); + + let any_bsr_high = ads.iter().any(is_bsr_high_ad); + if safe_budget == 0 && !(any_bsr_high && max_medium_slots > 0 && spacing_cap > 0) { + return posts_to_feed_items(scored_posts); + } + + let mut safe: Vec = Vec::new(); + let mut medium: Vec = Vec::new(); + let mut high: Vec = Vec::new(); + for post in scored_posts { + if is_high_risk(&post) { + high.push(post); + } else if is_medium_risk(&post) { + medium.push(post); + } else { + safe.push(post); + } + } + + let num_safe = safe.len(); + let group_size = if safe_budget > 0 { + num_safe / safe_budget + } else { + 0 + }; + + let mut safe_opts: Vec> = safe.into_iter().map(Some).collect(); + let mut medium_opts: Vec> = medium.into_iter().map(Some).collect(); + let mut triples: Vec<(AdIndexInfo, ScoredPost, ScoredPost)> = Vec::new(); + + let mut slot_tokens: Option<[Option; 2]> = None; + let mut medium_slot_tokens: Option<[Option; 2]> = None; + + let mut bsr_drop: u64 = 0; + let mut bsr_ok: u64 = 0; + let mut handle_drop: u64 = 0; + let mut keyword_drop: u64 = 0; + + let mut slot_rejections = SlotRejections::default(); + + let mut safe_group_idx = 0; + let mut medium_pair_idx = 0; + + for ad in ads { + if triples.len() >= spacing_cap { + break; + } + + if is_bsr_high_ad(&ad) && medium_pair_idx < max_medium_slots { + let start = medium_pair_idx * 2; + let above_ref = medium_opts[start].as_ref(); + let below_ref = medium_opts[start + 1].as_ref(); + if should_drop_handle(&ad, above_ref, below_ref) { + if safe_group_idx >= safe_budget || group_size == 0 { + handle_drop += 1; + slot_rejections.handle += 1; + continue; + } + } else if keyword_matches(&ad, above_ref, below_ref, &mut medium_slot_tokens) { + if safe_group_idx >= safe_budget || group_size == 0 { + keyword_drop += 1; + slot_rejections.keyword += 1; + continue; + } + } else { + let above = medium_opts[start].take().unwrap(); + let below = medium_opts[start + 1].take().unwrap(); + triples.push((ad, above, below)); + medium_pair_idx += 1; + medium_slot_tokens = None; + slot_rejections = SlotRejections::default(); + continue; + } + } + + if safe_group_idx >= safe_budget || group_size == 0 { + continue; + } + let group_start = safe_group_idx * group_size; + let above_ref = safe_opts[group_start].as_ref(); + let below_ref = safe_opts[group_start + 1].as_ref(); + + if should_drop_bsr_low(&ad, above_ref, below_ref) { + bsr_drop += 1; + slot_rejections.bsr_low += 1; + continue; + } + if is_bsr_low_ad(&ad) { + bsr_ok += 1; + } + + if should_drop_handle(&ad, above_ref, below_ref) { + handle_drop += 1; + slot_rejections.handle += 1; + continue; + } + + if keyword_matches(&ad, above_ref, below_ref, &mut slot_tokens) { + keyword_drop += 1; + slot_rejections.keyword += 1; + continue; + } + + let above = safe_opts[group_start].take().unwrap(); + let below = safe_opts[group_start + 1].take().unwrap(); + triples.push((ad, above, below)); + safe_group_idx += 1; + slot_tokens = None; + slot_rejections = SlotRejections::default(); + } + + let placed_ads = triples.len(); + emit_enforcement_metrics(bsr_drop, bsr_ok, handle_drop, keyword_drop); + let offered_slots = num_ads + .min(spacing_cap) + .min(max_safe_slots.saturating_add(medium_pair_idx)); + emit_slot_outcome_metrics( + placed_ads as u64, + offered_slots.saturating_sub(placed_ads) as u64, + &slot_rejections, + ); + + if placed_ads == 0 { + let mut all_posts: Vec = safe_opts.into_iter().flatten().collect(); + all_posts.extend(medium_opts.into_iter().flatten()); + all_posts.extend(high); + all_posts.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + return posts_to_feed_items(all_posts); + } + + let mut filler: Vec = + Vec::with_capacity(num_safe + medium_count + high.len() - 2 * placed_ads); + filler.extend(safe_opts.into_iter().flatten()); + filler.extend(medium_opts.into_iter().flatten()); + filler.extend(high); + filler.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + let inter_ad_gaps = placed_ads; + let filler_per_gap = filler.len() / inter_ad_gaps; + let remainder = filler.len() % inter_ad_gaps; + let mut filler_iter = filler.into_iter(); + + let mut items: Vec = Vec::with_capacity(n + placed_ads); + + for (i, (ad, above, below)) in triples.into_iter().enumerate() { + items.push(FeedItem { + position: 0, + item: Some(feed_item::Item::Post(above)), + }); + items.push(FeedItem { + position: 0, + item: Some(feed_item::Item::Ad(ad)), + }); + items.push(FeedItem { + position: 0, + item: Some(feed_item::Item::Post(below)), + }); + + let count = filler_per_gap + if i >= inter_ad_gaps - remainder { 1 } else { 0 }; + for _ in 0..count { + if let Some(post) = filler_iter.next() { + items.push(FeedItem { + position: 0, + item: Some(feed_item::Item::Post(post)), + }); + } + } + } + + items.truncate(RESULT_SIZE); + if matches!(items.last(), Some(item) if matches!(item.item, Some(feed_item::Item::Ad(_)))) { + items.pop(); + } + for (i, item) in items.iter_mut().enumerate() { + item.position = i as i32; + } + + items +} + +pub(crate) fn serving_limitation( + ads_supply: usize, + from_spacing: usize, + from_safe: usize, +) -> &'static str { + let budget = ads_supply.min(from_spacing).min(from_safe); + if budget == ads_supply { + "ads_supply" + } else if budget == from_spacing { + "spacing" + } else { + "safe_posts" + } +} + +fn emit_serving_limitation(factor: &'static str) { + let Some(receiver) = global_stats_receiver() else { + return; + }; + receiver.incr(SERVING_LIMITATION_METRIC, &[("factor", factor)], 1); +} + +#[derive(Default)] +pub(crate) struct SlotRejections { + pub(crate) bsr_low: u64, + pub(crate) handle: u64, + pub(crate) keyword: u64, +} + +impl SlotRejections { + pub(crate) fn stuck_outcome(&self) -> &'static str { + if self.bsr_low == 0 && self.handle == 0 && self.keyword == 0 { + "unfilled_no_ads" + } else if self.bsr_low >= self.keyword && self.bsr_low >= self.handle { + "unfilled_bsr_low" + } else if self.keyword >= self.handle { + "unfilled_keyword" + } else { + "unfilled_handle" + } + } +} + +fn emit_slot_outcome_metrics(filled: u64, unfilled: u64, stuck: &SlotRejections) { + let Some(receiver) = global_stats_receiver() else { + return; + }; + if filled > 0 { + receiver.incr(SLOT_OUTCOME_METRIC, &[("outcome", "filled")], filled); + } + if unfilled == 0 { + return; + } + receiver.incr( + SLOT_OUTCOME_METRIC, + &[("outcome", stuck.stuck_outcome())], + 1, + ); + if unfilled > 1 { + receiver.incr( + SLOT_OUTCOME_METRIC, + &[("outcome", "unfilled_no_ads")], + unfilled - 1, + ); + } +} + +fn emit_enforcement_metrics(bsr_drop: u64, bsr_ok: u64, handle_drop: u64, keyword_drop: u64) { + let Some(receiver) = global_stats_receiver() else { + return; + }; + if bsr_drop > 0 { + receiver.incr(ENFORCEMENT_METRIC, &[("action", "drop")], bsr_drop); + } + if bsr_ok > 0 { + receiver.incr(ENFORCEMENT_METRIC, &[("action", "ok")], bsr_ok); + } + if handle_drop > 0 { + receiver.incr( + ENFORCEMENT_METRIC, + &[("action", "handle_drop")], + handle_drop, + ); + } + if keyword_drop > 0 { + receiver.incr( + ENFORCEMENT_METRIC, + &[("action", "keyword_drop")], + keyword_drop, + ); + } +} diff --git a/home-mixer/ads/tests/blender_comparison_test.rs b/home-mixer/ads/tests/blender_comparison_test.rs index b4026816..8aca0230 100644 --- a/home-mixer/ads/tests/blender_comparison_test.rs +++ b/home-mixer/ads/tests/blender_comparison_test.rs @@ -1,5 +1,5 @@ use crate::ads::{ - partition_organic_blender, safe_gap_blender, + multi_risk_blender, partition_organic_blender, safe_gap_blender, util::{should_drop_handle, should_drop_keyword}, }; use crate::params::RESULT_SIZE; @@ -9,13 +9,17 @@ use std::collections::HashMap; use xai_home_mixer_proto::{feed_item, BrandSafetyVerdict, FeedItem, ScoredPost}; use xai_recsys_proto::{AdAdjacencyControl, AdIndexInfo, BrandSafetyRiskLevel}; -const UNSAFE_POST_RATE: f64 = 0.52; +const HIGH_RISK_LABEL_POST_RATE: f64 = 0.045; +const NSFW_AUTHOR_POST_RATE: f64 = 0.117; +const MEDIUM_RISK_POST_RATE: f64 = 0.411; +const LOW_RISK_POST_RATE: f64 = 0.093; -const LOW_RISK_POST_RATE: f64 = 0.06; +const NSFW_AUTHOR_POOL: &[u64] = &[2001, 2002, 2003, 2004, 2005]; -const BSR_LOW_AD_RATE: f64 = 0.145; -const BSR_IAS_AD_RATE: f64 = 0.008; -const BSR_NO_RISK_AD_RATE: f64 = 0.029; +const BSR_LOW_AD_RATE: f64 = 0.214; +const BSR_IAS_AD_RATE: f64 = 0.003; +const BSR_NO_RISK_AD_RATE: f64 = 0.013; +const BSR_HIGH_AD_RATE: f64 = 0.10; const KNOWN_HANDLE_POST_RATE: f64 = 0.10; const AD_EXCLUDED_HANDLES_RATE: f64 = 0.10; @@ -37,7 +41,9 @@ const SAFE_TEXT_POOL: &[&str] = &[ "happy birthday to my best friend", ]; -const NUM_SCENARIOS: usize = 200; +const IDEAL_GAPS: &[usize] = &[1, 5, 9, 13, 17, 21, 25]; + +const NUM_SCENARIOS: usize = 1000; const RNG_SEED: u64 = 412; @@ -46,22 +52,63 @@ struct Scenario { ads: Vec, } +fn expected_high_risk_post_rate() -> f64 { + 1.0 - (1.0 - HIGH_RISK_LABEL_POST_RATE) * (1.0 - NSFW_AUTHOR_POST_RATE) +} + +fn is_avoid_verdict(verdict: BrandSafetyVerdict) -> bool { + matches!( + verdict, + BrandSafetyVerdict::MediumRisk | BrandSafetyVerdict::HighRisk + ) +} + +fn block_ideal_gaps(posts: &mut [ScoredPost]) { + let n = posts.len(); + for &gap in IDEAL_GAPS { + if gap >= n { + continue; + } + let left = gap - 1; + let right = gap; + if is_avoid_verdict(posts[left].brand_safety_verdict()) + || is_avoid_verdict(posts[right].brand_safety_verdict()) + { + continue; + } + if let Some(j) = (right + 1..n).find(|&j| is_avoid_verdict(posts[j].brand_safety_verdict())) + { + posts.swap(right, j); + } + } +} + fn generate_scenarios(rng: &mut StdRng) -> Vec { (0..NUM_SCENARIOS) .map(|_| { - let num_posts = rng.random_range(10..=35); + let num_posts = rng.random_range(31..=35); let num_ads = rng.random_range(3..=15); - let posts: Vec = (1..=num_posts) + let mut posts: Vec = (1..=num_posts) .map(|id| { let roll: f64 = rng.random(); - let verdict = if roll < UNSAFE_POST_RATE { + let tweet_verdict = if roll < HIGH_RISK_LABEL_POST_RATE { + BrandSafetyVerdict::HighRisk + } else if roll < HIGH_RISK_LABEL_POST_RATE + MEDIUM_RISK_POST_RATE { BrandSafetyVerdict::MediumRisk - } else if roll < UNSAFE_POST_RATE + LOW_RISK_POST_RATE { + } else if roll + < HIGH_RISK_LABEL_POST_RATE + MEDIUM_RISK_POST_RATE + LOW_RISK_POST_RATE + { BrandSafetyVerdict::LowRisk } else { BrandSafetyVerdict::Safe }; + let nsfw_author = rng.random_bool(NSFW_AUTHOR_POST_RATE); + let verdict = if nsfw_author { + BrandSafetyVerdict::HighRisk + } else { + tweet_verdict + }; let (author_id, screen_names) = if rng.random_bool(KNOWN_HANDLE_POST_RATE) { let idx = rng.random_range(0..HANDLE_POOL.len()); @@ -69,6 +116,9 @@ fn generate_scenarios(rng: &mut StdRng) -> Vec { let mut map = HashMap::new(); map.insert(aid, SCREEN_NAME_POOL[idx].to_string()); (aid, map) + } else if nsfw_author { + let aid = NSFW_AUTHOR_POOL[rng.random_range(0..NSFW_AUTHOR_POOL.len())]; + (aid, HashMap::new()) } else { (id as u64, HashMap::new()) }; @@ -90,10 +140,11 @@ fn generate_scenarios(rng: &mut StdRng) -> Vec { } }) .collect(); + block_ideal_gaps(&mut posts); let ads: Vec = (0..num_ads) .map(|i| { - let insert_position = i * 3 + 1; + let insert_position = i * 8 + 1; let roll: f64 = rng.random(); let risk = if roll < BSR_LOW_AD_RATE { BrandSafetyRiskLevel::BsrLow @@ -101,6 +152,10 @@ fn generate_scenarios(rng: &mut StdRng) -> Vec { BrandSafetyRiskLevel::BsrIas } else if roll < BSR_LOW_AD_RATE + BSR_IAS_AD_RATE + BSR_NO_RISK_AD_RATE { BrandSafetyRiskLevel::BsrNoRisk + } else if roll + < BSR_LOW_AD_RATE + BSR_IAS_AD_RATE + BSR_NO_RISK_AD_RATE + BSR_HIGH_AD_RATE + { + BrandSafetyRiskLevel::BsrHigh } else { BrandSafetyRiskLevel::BsrNormal }; @@ -165,11 +220,32 @@ fn ad_positions(items: &[FeedItem]) -> Vec { .collect() } +fn ad_bsr(ad: &AdIndexInfo) -> BrandSafetyRiskLevel { + ad.ad_adjacency_control + .as_ref() + .map(|c| c.brand_safety_risk()) + .unwrap_or(BrandSafetyRiskLevel::BsrUnknown) +} + +fn is_bsr_low_or_ias(risk: BrandSafetyRiskLevel) -> bool { + matches!( + risk, + BrandSafetyRiskLevel::BsrLow | BrandSafetyRiskLevel::BsrIas + ) +} + struct AdjacencyViolations { med_risk_adj: usize, + non_high_med_risk_adj: usize, bsr_low_low_risk_adj: usize, bsr_low_placed: usize, bsr_low_available: usize, + bsr_high_placed: usize, + bsr_high_available: usize, + bsr_low_only_placed: usize, + bsr_ias_placed: usize, + bsr_normal_placed: usize, + bsr_no_risk_placed: usize, handle_adj: usize, keyword_adj: usize, } @@ -180,23 +256,23 @@ fn count_adjacency_violations( ) -> AdjacencyViolations { let mut v = AdjacencyViolations { med_risk_adj: 0, + non_high_med_risk_adj: 0, bsr_low_low_risk_adj: 0, bsr_low_placed: 0, handle_adj: 0, keyword_adj: 0, + bsr_high_placed: 0, + bsr_low_only_placed: 0, + bsr_ias_placed: 0, + bsr_normal_placed: 0, + bsr_no_risk_placed: 0, bsr_low_available: input_ads .iter() - .filter(|a| { - a.ad_adjacency_control - .as_ref() - .map(|c| { - matches!( - c.brand_safety_risk(), - BrandSafetyRiskLevel::BsrLow | BrandSafetyRiskLevel::BsrIas - ) - }) - .unwrap_or(false) - }) + .filter(|a| is_bsr_low_or_ias(ad_bsr(a))) + .count(), + bsr_high_available: input_ads + .iter() + .filter(|a| ad_bsr(a) == BrandSafetyRiskLevel::BsrHigh) .count(), }; @@ -206,19 +282,21 @@ fn count_adjacency_violations( _ => continue, }; - let is_sensitive = ad - .ad_adjacency_control - .as_ref() - .map(|c| { - matches!( - c.brand_safety_risk(), - BrandSafetyRiskLevel::BsrLow | BrandSafetyRiskLevel::BsrIas - ) - }) - .unwrap_or(false); + let risk = ad_bsr(ad); + let is_sensitive = is_bsr_low_or_ias(risk); if is_sensitive { v.bsr_low_placed += 1; } + if risk == BrandSafetyRiskLevel::BsrHigh { + v.bsr_high_placed += 1; + } + match risk { + BrandSafetyRiskLevel::BsrLow => v.bsr_low_only_placed += 1, + BrandSafetyRiskLevel::BsrIas => v.bsr_ias_placed += 1, + BrandSafetyRiskLevel::BsrNormal => v.bsr_normal_placed += 1, + BrandSafetyRiskLevel::BsrNoRisk => v.bsr_no_risk_placed += 1, + _ => {} + } let above_post = if i > 0 { match &items[i - 1].item { @@ -244,6 +322,9 @@ fn count_adjacency_violations( || below_verdict == Some(BrandSafetyVerdict::MediumRisk) { v.med_risk_adj += 1; + if risk != BrandSafetyRiskLevel::BsrHigh { + v.non_high_med_risk_adj += 1; + } } if is_sensitive @@ -307,6 +388,13 @@ struct BlenderStats { bsr_low_with_lowrisk_adj: usize, bsr_low_available: usize, bsr_low_placed: usize, + non_high_med_risk_adj: usize, + bsr_high_available: usize, + bsr_high_placed: usize, + bsr_low_only_placed: usize, + bsr_ias_placed: usize, + bsr_normal_placed: usize, + bsr_no_risk_placed: usize, handle_adj: usize, keyword_adj: usize, } @@ -326,6 +414,13 @@ impl BlenderStats { bsr_low_with_lowrisk_adj: 0, bsr_low_available: 0, bsr_low_placed: 0, + non_high_med_risk_adj: 0, + bsr_high_available: 0, + bsr_high_placed: 0, + bsr_low_only_placed: 0, + bsr_ias_placed: 0, + bsr_normal_placed: 0, + bsr_no_risk_placed: 0, handle_adj: 0, keyword_adj: 0, } @@ -372,6 +467,13 @@ impl BlenderStats { } (1.0 - self.bsr_low_placed as f64 / self.bsr_low_available as f64) * 100.0 } + + fn bsr_high_drop_pct(&self) -> f64 { + if self.bsr_high_available == 0 { + return 0.0; + } + (1.0 - self.bsr_high_placed as f64 / self.bsr_high_available as f64) * 100.0 + } } type BlendFn = fn(Vec, Vec, usize) -> Vec; @@ -392,6 +494,7 @@ fn blender_comparison_report() { let blenders: Vec<(&str, BlendFn)> = vec![ ("safe_gap", safe_gap_blend), ("partition_organic", partition_organic_blender::blend_impl), + ("multi_risk", multi_risk_blender::blend_impl), ]; let mut all_stats: Vec = blenders @@ -423,9 +526,16 @@ fn blender_comparison_report() { let v = count_adjacency_violations(&result, &scenario.ads); all_stats[blender_idx].ads_with_medrisk_adj += v.med_risk_adj; + all_stats[blender_idx].non_high_med_risk_adj += v.non_high_med_risk_adj; all_stats[blender_idx].bsr_low_with_lowrisk_adj += v.bsr_low_low_risk_adj; all_stats[blender_idx].bsr_low_available += v.bsr_low_available; all_stats[blender_idx].bsr_low_placed += v.bsr_low_placed; + all_stats[blender_idx].bsr_high_available += v.bsr_high_available; + all_stats[blender_idx].bsr_high_placed += v.bsr_high_placed; + all_stats[blender_idx].bsr_low_only_placed += v.bsr_low_only_placed; + all_stats[blender_idx].bsr_ias_placed += v.bsr_ias_placed; + all_stats[blender_idx].bsr_normal_placed += v.bsr_normal_placed; + all_stats[blender_idx].bsr_no_risk_placed += v.bsr_no_risk_placed; all_stats[blender_idx].handle_adj += v.handle_adj; all_stats[blender_idx].keyword_adj += v.keyword_adj; @@ -444,26 +554,32 @@ fn blender_comparison_report() { "║ ADS BLENDER COMPARISON REPORT ║" ); eprintln!( - "║ {NUM_SCENARIOS} scenarios — posts: ~{:.0}% MedRisk, ~{:.0}% LowRisk, ~{:.0}% Safe — ads: ~{:.0}% BSR_LOW, ~{:.0}% BSR_IAS, ~{:.0}% BSR_NORMAL ║", - UNSAFE_POST_RATE * 100.0, - LOW_RISK_POST_RATE * 100.0, - (1.0 - UNSAFE_POST_RATE - LOW_RISK_POST_RATE) * 100.0, + "║ {NUM_SCENARIOS} scenarios — posts: ~{:.0}% HighRisk ({:.0}% labels + {:.0}% nsfw_author), ~{:.0}% MedRisk, ~{:.0}% LowRisk, ~{:.0}% Safe — ads: ~{:.0}% BSR_LOW, ~{:.0}% BSR_IAS, ~{:.0}% BSR_HIGH, ~{:.0}% BSR_NORMAL ║", + expected_high_risk_post_rate() * 100.0, + HIGH_RISK_LABEL_POST_RATE * 100.0, + NSFW_AUTHOR_POST_RATE * 100.0, + MEDIUM_RISK_POST_RATE * (1.0 - NSFW_AUTHOR_POST_RATE) * 100.0, + LOW_RISK_POST_RATE * (1.0 - NSFW_AUTHOR_POST_RATE) * 100.0, + (1.0 - HIGH_RISK_LABEL_POST_RATE - MEDIUM_RISK_POST_RATE - LOW_RISK_POST_RATE) + * (1.0 - NSFW_AUTHOR_POST_RATE) + * 100.0, BSR_LOW_AD_RATE * 100.0, BSR_IAS_AD_RATE * 100.0, - (1.0 - BSR_LOW_AD_RATE - BSR_IAS_AD_RATE - BSR_NO_RISK_AD_RATE) * 100.0, + BSR_HIGH_AD_RATE * 100.0, + (1.0 - BSR_LOW_AD_RATE - BSR_IAS_AD_RATE - BSR_NO_RISK_AD_RATE - BSR_HIGH_AD_RATE) * 100.0, ); eprintln!( - "╠══════════════════════╦═════════╦═════════╦═════════╦═════════╦═════════╦══════════╦══════════╦══════════╦══════════╦══════════╦══════════╣" + "╠══════════════════════╦═════════╦═════════╦═════════╦═════════╦═════════╦══════════╦══════════╦══════════╦══════════╦══════════╦══════════╦══════════╣" ); eprintln!( - "║ Blender ║ AdPlace ║ AdAvail ║ Fill% ║ Load% ║ AvgSpc ║ MedAdj% ║ LR→BSR% ║ BSRDrop% ║ BSRPlcd ║ HndlAdj ║ KwAdj ║" + "║ Blender ║ AdPlace ║ AdAvail ║ Fill% ║ Load% ║ AvgSpc ║ MedAdj% ║ LR→BSR% ║ BSRDrop% ║ HighDrop%║ HighPlcd ║ HndlAdj ║ KwAdj ║" ); eprintln!( - "╠══════════════════════╬═════════╬═════════╬═════════╬═════════╬═════════╬══════════╬══════════╬══════════╬══════════╬══════════╬══════════╣" + "╠══════════════════════╬═════════╬═════════╬═════════╬═════════╬═════════╬══════════╬══════════╬══════════╬══════════╬══════════╬══════════╬══════════╣" ); for stats in &all_stats { eprintln!( - "║ {:<20} ║ {:>7} ║ {:>7} ║ {:>6.1}% ║ {:>6.1}% ║ {:>7.2} ║ {:>7.1}% ║ {:>7.1}% ║ {:>7.1}% ║ {:>4}/{:<4}║ {:>8} ║ {:>8} ║", + "║ {:<20} ║ {:>7} ║ {:>7} ║ {:>6.1}% ║ {:>6.1}% ║ {:>7.2} ║ {:>7.1}% ║ {:>7.1}% ║ {:>7.1}% ║ {:>7.1}% ║ {:>4}/{:<4}║ {:>8} ║ {:>8} ║", stats.name, stats.total_ads_placed, stats.total_ads_available, @@ -473,16 +589,45 @@ fn blender_comparison_report() { stats.medrisk_adj_pct(), stats.bsr_low_lowrisk_adj_pct(), stats.bsr_low_drop_pct(), - stats.bsr_low_placed, - stats.bsr_low_available, + stats.bsr_high_drop_pct(), + stats.bsr_high_placed, + stats.bsr_high_available, stats.handle_adj, stats.keyword_adj, ); } eprintln!( - "╚══════════════════════╩═════════╩═════════╩═════════╩═════════╩═════════╩══════════╩══════════╩══════════╩══════════╩══════════╩══════════╝" + "╚══════════════════════╩═════════╩═════════╩═════════╩═════════╩═════════╩══════════╩══════════╩══════════╩══════════╩══════════╩══════════╩══════════╝" ); eprintln!(); + eprintln!("Placed mix (share of AdPlace; Safe = BSR_NORMAL):"); + for stats in &all_stats { + let n = stats.total_ads_placed as f64; + let pct = |c: usize| if n == 0.0 { 0.0 } else { c as f64 / n * 100.0 }; + eprintln!( + " {:<20} LOW {:>4} ({:>4.1}%) IAS {:>4} ({:>4.1}%) HIGH {:>4} ({:>4.1}%) Safe {:>4} ({:>4.1}%) other {:>4} ({:>4.1}%)", + stats.name, + stats.bsr_low_only_placed, + pct(stats.bsr_low_only_placed), + stats.bsr_ias_placed, + pct(stats.bsr_ias_placed), + stats.bsr_high_placed, + pct(stats.bsr_high_placed), + stats.bsr_normal_placed, + pct(stats.bsr_normal_placed), + stats.total_ads_placed + - stats.bsr_low_only_placed + - stats.bsr_ias_placed + - stats.bsr_high_placed + - stats.bsr_normal_placed, + pct(stats.total_ads_placed + - stats.bsr_low_only_placed + - stats.bsr_ias_placed + - stats.bsr_high_placed + - stats.bsr_normal_placed), + ); + } + eprintln!(); eprintln!("Legend: Fill% = ads placed / ads available"); eprintln!(" Load% = ads placed / total items (ad load in timeline)"); eprintln!(" AvgSpc = average number of posts between consecutive ads"); @@ -491,7 +636,8 @@ fn blender_comparison_report() { " LR→BSR% = % of placed BSR_LOW/IAS ads adjacent to LowRisk post (should be 0 for partition_organic)" ); eprintln!(" BSRDrop% = % of BSR_LOW/IAS ads dropped (not placed)"); - eprintln!(" BSRPlcd = BSR_LOW/IAS ads placed / available"); + eprintln!(" HighDrop% = % of BSR_HIGH ads dropped (not placed)"); + eprintln!(" HighPlcd = BSR_HIGH ads placed / available"); eprintln!( " HndlAdj = placed ads adjacent to an excluded handle (should be 0 for partition_organic)" ); @@ -527,4 +673,34 @@ fn blender_comparison_report() { "partition_organic had {} ads adjacent to excluded keyword matches!", po_stats.keyword_adj ); + assert_eq!( + po_stats.ads_with_medrisk_adj, 0, + "partition_organic had {} ads adjacent to MediumRisk posts!", + po_stats.ads_with_medrisk_adj + ); + + let npo_stats = all_stats + .iter() + .find(|s| s.name == "multi_risk") + .expect("multi_risk blender not found in stats"); + assert_eq!( + npo_stats.bsr_low_with_lowrisk_adj, 0, + "multi_risk had {} BSR_LOW/IAS ads adjacent to LowRisk posts!", + npo_stats.bsr_low_with_lowrisk_adj + ); + assert_eq!( + npo_stats.handle_adj, 0, + "multi_risk had {} ads adjacent to excluded handles!", + npo_stats.handle_adj + ); + assert_eq!( + npo_stats.keyword_adj, 0, + "multi_risk had {} ads adjacent to excluded keyword matches!", + npo_stats.keyword_adj + ); + assert_eq!( + npo_stats.non_high_med_risk_adj, 0, + "multi_risk had {} non-BSR_HIGH ads adjacent to MediumRisk posts!", + npo_stats.non_high_med_risk_adj + ); } diff --git a/home-mixer/ads/tests/mod.rs b/home-mixer/ads/tests/mod.rs index c082296e..3b877a9c 100644 --- a/home-mixer/ads/tests/mod.rs +++ b/home-mixer/ads/tests/mod.rs @@ -1,5 +1,6 @@ mod blender_comparison_test; mod following_ad_blender_tests; +mod multi_risk_blender_tests; mod partition_organic_blender_tests; mod safe_gap_blender_tests; mod time_gap_blender_tests; diff --git a/home-mixer/ads/tests/multi_risk_blender_tests.rs b/home-mixer/ads/tests/multi_risk_blender_tests.rs new file mode 100644 index 00000000..82a0ca8a --- /dev/null +++ b/home-mixer/ads/tests/multi_risk_blender_tests.rs @@ -0,0 +1,215 @@ +use crate::ads::multi_risk_blender::*; +use xai_home_mixer_proto::{feed_item, BrandSafetyVerdict, FeedItem, ScoredPost}; +use xai_recsys_proto::{AdAdjacencyControl, AdIndexInfo, BrandSafetyRiskLevel}; + +fn make_post(tweet_id: u64) -> ScoredPost { + ScoredPost { + tweet_id, + score: 1.0 - (tweet_id as f32 * 0.01), + ..Default::default() + } +} + +fn make_avoid_post(tweet_id: u64) -> ScoredPost { + ScoredPost { + tweet_id, + brand_safety_verdict: BrandSafetyVerdict::MediumRisk.into(), + score: 1.0 - (tweet_id as f32 * 0.01), + ..Default::default() + } +} + +fn make_high_risk_post(tweet_id: u64) -> ScoredPost { + ScoredPost { + tweet_id, + brand_safety_verdict: BrandSafetyVerdict::HighRisk.into(), + score: 1.0 - (tweet_id as f32 * 0.01), + ..Default::default() + } +} + +fn make_normal_ad(post_id: i64) -> AdIndexInfo { + AdIndexInfo { + post_id, + ad_adjacency_control: Some(AdAdjacencyControl { + brand_safety_risk: BrandSafetyRiskLevel::BsrNormal.into(), + ..Default::default() + }), + ..Default::default() + } +} + +fn make_bsr_high_ad(post_id: i64) -> AdIndexInfo { + AdIndexInfo { + post_id, + ad_adjacency_control: Some(AdAdjacencyControl { + brand_safety_risk: BrandSafetyRiskLevel::BsrHigh.into(), + ..Default::default() + }), + ..Default::default() + } +} + +fn make_bsr_high_ad_with_handles(post_id: i64, handles: &[i64]) -> AdIndexInfo { + AdIndexInfo { + post_id, + ad_adjacency_control: Some(AdAdjacencyControl { + brand_safety_risk: BrandSafetyRiskLevel::BsrHigh.into(), + handles: handles.to_vec(), + ..Default::default() + }), + ..Default::default() + } +} + +fn ad_count(items: &[FeedItem]) -> usize { + items + .iter() + .filter(|i| matches!(i.item, Some(feed_item::Item::Ad(_)))) + .count() +} + +fn ad_neighbour_verdicts(items: &[FeedItem]) -> Vec<(BrandSafetyVerdict, BrandSafetyVerdict)> { + let mut result = Vec::new(); + for (i, item) in items.iter().enumerate() { + if matches!(item.item, Some(feed_item::Item::Ad(_))) { + let above = if i > 0 { + match &items[i - 1].item { + Some(feed_item::Item::Post(p)) => p.brand_safety_verdict(), + _ => BrandSafetyVerdict::VerdictUnspecified, + } + } else { + BrandSafetyVerdict::VerdictUnspecified + }; + let below = if i + 1 < items.len() { + match &items[i + 1].item { + Some(feed_item::Item::Post(p)) => p.brand_safety_verdict(), + _ => BrandSafetyVerdict::VerdictUnspecified, + } + } else { + BrandSafetyVerdict::VerdictUnspecified + }; + result.push((above, below)); + } + } + result +} + +fn ad_bsr_levels(items: &[FeedItem]) -> Vec { + items + .iter() + .filter_map(|i| match &i.item { + Some(feed_item::Item::Ad(ad)) => Some( + ad.ad_adjacency_control + .as_ref() + .map(|c| c.brand_safety_risk()) + .unwrap_or(BrandSafetyRiskLevel::BsrUnknown), + ), + _ => None, + }) + .collect() +} + +#[test] +fn bsr_high_ad_adjacent_to_medium_risk() { + let mut posts: Vec<_> = (1..=6).map(make_avoid_post).collect(); + posts.extend((7..=12).map(make_post)); + let result = blend_impl(posts, vec![make_bsr_high_ad(100)], 5); + assert_eq!(ad_count(&result), 1); + let (above, below) = ad_neighbour_verdicts(&result)[0]; + assert_eq!(above, BrandSafetyVerdict::MediumRisk); + assert_eq!(below, BrandSafetyVerdict::MediumRisk); +} + +#[test] +fn normal_ad_not_adjacent_to_medium_risk_when_bsr_high_fills() { + let mut posts: Vec<_> = (1..=6).map(make_avoid_post).collect(); + posts.extend((7..=12).map(make_post)); + let result = blend_impl(posts, vec![make_normal_ad(100), make_bsr_high_ad(200)], 5); + let bsr_levels = ad_bsr_levels(&result); + let neighbours = ad_neighbour_verdicts(&result); + assert!(bsr_levels.contains(&BrandSafetyRiskLevel::BsrHigh)); + for (i, bsr) in bsr_levels.iter().enumerate() { + let (above, below) = neighbours[i]; + if *bsr == BrandSafetyRiskLevel::BsrNormal { + assert_ne!(above, BrandSafetyVerdict::MediumRisk); + assert_ne!(below, BrandSafetyVerdict::MediumRisk); + } + if *bsr == BrandSafetyRiskLevel::BsrHigh { + assert_eq!(above, BrandSafetyVerdict::MediumRisk); + assert_eq!(below, BrandSafetyVerdict::MediumRisk); + } + assert_ne!(above, BrandSafetyVerdict::HighRisk); + assert_ne!(below, BrandSafetyVerdict::HighRisk); + } +} + +#[test] +fn high_risk_never_adjacent_including_bsr_high() { + let mut posts: Vec<_> = (1..=6).map(make_high_risk_post).collect(); + posts.extend((7..=12).map(make_post)); + let result = blend_impl(posts, vec![make_bsr_high_ad(100), make_normal_ad(200)], 5); + assert!(ad_count(&result) > 0); + for (above, below) in ad_neighbour_verdicts(&result) { + assert_ne!(above, BrandSafetyVerdict::HighRisk); + assert_ne!(below, BrandSafetyVerdict::HighRisk); + assert_ne!(above, BrandSafetyVerdict::MediumRisk); + assert_ne!(below, BrandSafetyVerdict::MediumRisk); + } +} + +#[test] +fn all_medium_risk_places_bsr_high_drops_normal() { + let posts: Vec<_> = (1..=8).map(make_avoid_post).collect(); + let with_high = blend_impl(posts.clone(), vec![make_bsr_high_ad(100)], 5); + assert_eq!(ad_count(&with_high), 1); + let (above, below) = ad_neighbour_verdicts(&with_high)[0]; + assert_eq!(above, BrandSafetyVerdict::MediumRisk); + assert_eq!(below, BrandSafetyVerdict::MediumRisk); + + let with_normal = blend_impl(posts, vec![make_normal_ad(100)], 5); + assert_eq!(ad_count(&with_normal), 0); +} + +#[test] +fn excluded_high_does_not_burn_medium_pair() { + let mut medium: Vec<_> = (1..=4).map(make_avoid_post).collect(); + medium[0].author_id = 9999; + medium[1].author_id = 9999; + let mut posts = medium; + posts.extend((5..=12).map(make_post)); + + let result = blend_impl( + posts, + vec![ + make_bsr_high_ad_with_handles(100, &[9999]), + make_bsr_high_ad(200), + ], + 5, + ); + assert_eq!(ad_count(&result), 2); + + let mut neighbours_by_ad = std::collections::HashMap::new(); + for (i, item) in result.iter().enumerate() { + let Some(feed_item::Item::Ad(ad)) = &item.item else { + continue; + }; + let above = match &result[i - 1].item { + Some(feed_item::Item::Post(p)) => p.brand_safety_verdict(), + _ => BrandSafetyVerdict::VerdictUnspecified, + }; + let below = match &result[i + 1].item { + Some(feed_item::Item::Post(p)) => p.brand_safety_verdict(), + _ => BrandSafetyVerdict::VerdictUnspecified, + }; + neighbours_by_ad.insert(ad.post_id, (above, below)); + } + + let (above_100, below_100) = neighbours_by_ad[&100]; + assert_ne!(above_100, BrandSafetyVerdict::MediumRisk); + assert_ne!(below_100, BrandSafetyVerdict::MediumRisk); + + let (above_200, below_200) = neighbours_by_ad[&200]; + assert_eq!(above_200, BrandSafetyVerdict::MediumRisk); + assert_eq!(below_200, BrandSafetyVerdict::MediumRisk); +} diff --git a/home-mixer/ads/tests/partition_organic_blender_tests.rs b/home-mixer/ads/tests/partition_organic_blender_tests.rs index aeab6fd9..e54ef528 100644 --- a/home-mixer/ads/tests/partition_organic_blender_tests.rs +++ b/home-mixer/ads/tests/partition_organic_blender_tests.rs @@ -51,6 +51,17 @@ fn make_sensitive_ad(post_id: i64) -> AdIndexInfo { } } +fn make_bsr_high_ad(post_id: i64) -> AdIndexInfo { + AdIndexInfo { + post_id, + ad_adjacency_control: Some(AdAdjacencyControl { + brand_safety_risk: BrandSafetyRiskLevel::BsrHigh.into(), + ..Default::default() + }), + ..Default::default() + } +} + fn ad_count(items: &[FeedItem]) -> usize { items .iter() @@ -1106,3 +1117,22 @@ fn serving_limitation_ties_break_supply_then_spacing() { assert_eq!(serving_limitation(9, 5, 5), "spacing"); assert_eq!(serving_limitation(5, 9, 5), "ads_supply"); } + +#[test] +fn bsr_high_ad_sits_next_to_safe_not_medium() { + let mut posts: Vec<_> = (1..=6).map(make_avoid_post).collect(); + posts.extend((7..=12).map(make_post)); + let result = blend_impl(posts, vec![make_bsr_high_ad(100)], 5); + assert_eq!(ad_count(&result), 1); + let (above, below) = ad_neighbour_verdicts(&result)[0]; + assert_ne!(above, BrandSafetyVerdict::MediumRisk); + assert_ne!(below, BrandSafetyVerdict::MediumRisk); + assert_eq!(ad_bsr_levels(&result)[0], BrandSafetyRiskLevel::BsrHigh); +} + +#[test] +fn bsr_high_ad_dropped_when_only_medium_posts() { + let posts: Vec<_> = (1..=8).map(make_avoid_post).collect(); + let result = blend_impl(posts, vec![make_bsr_high_ad(100)], 5); + assert_eq!(ad_count(&result), 0); +} diff --git a/home-mixer/ads/util.rs b/home-mixer/ads/util.rs index cc40aa7c..26025ab3 100644 --- a/home-mixer/ads/util.rs +++ b/home-mixer/ads/util.rs @@ -22,6 +22,17 @@ pub(crate) struct AdSpacing { } pub(crate) fn has_avoid(post: &ScoredPost) -> bool { + matches!( + post.brand_safety_verdict(), + BrandSafetyVerdict::MediumRisk | BrandSafetyVerdict::HighRisk + ) +} + +pub(crate) fn is_high_risk(post: &ScoredPost) -> bool { + post.brand_safety_verdict() == BrandSafetyVerdict::HighRisk +} + +pub(crate) fn is_medium_risk(post: &ScoredPost) -> bool { post.brand_safety_verdict() == BrandSafetyVerdict::MediumRisk } @@ -75,6 +86,14 @@ pub(crate) fn is_bsr_low_ad(ad: &AdIndexInfo) -> bool { ) } +pub(crate) fn is_bsr_high_ad(ad: &AdIndexInfo) -> bool { + ad.ad_adjacency_control + .as_ref() + .map(|c| c.brand_safety_risk()) + .unwrap_or(BrandSafetyRiskLevel::BsrUnknown) + == BrandSafetyRiskLevel::BsrHigh +} + pub(crate) fn should_drop_bsr_low( ad: &AdIndexInfo, above: Option<&ScoredPost>, @@ -130,6 +149,25 @@ pub(crate) fn should_drop_keyword( above.map(text_matches).unwrap_or(false) || below.map(text_matches).unwrap_or(false) } +pub(crate) fn keyword_matches( + ad: &AdIndexInfo, + above: Option<&ScoredPost>, + below: Option<&ScoredPost>, + slot_tokens: &mut Option<[Option; 2]>, +) -> bool { + let Some(keywords) = tokenize_ad_keywords(ad) else { + return false; + }; + let tokens = slot_tokens.get_or_insert_with(|| { + [above, below].map(|post| post.map(|p| tokenize_tweet_text(&p.tweet_text))) + }); + let matches = |t: &Option| { + t.as_ref() + .is_some_and(|t| tokens_match_any_keyword(t, &keywords)) + }; + matches(&tokens[0]) || matches(&tokens[1]) +} + pub(crate) fn tokenize_tweet_text(text: &str) -> TokenSequence { TWEET_TOKENIZER.tokenize(text) } diff --git a/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs b/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs index 2d271fa7..3ac4e132 100644 --- a/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs +++ b/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs @@ -1,9 +1,10 @@ use crate::models::brand_safety::{ - botmaker_rule_category, botmaker_rule_id_from, compute_verdict, truncate_description, - worst_verdict, BrandSafetyVerdict, + botmaker_rule_category, botmaker_rule_id_from, compute_verdict, compute_verdict_v2, + truncate_description, worst_verdict, BrandSafetyVerdict, }; use crate::models::candidate::{PostCandidate, SafetyLabelInfo}; use crate::models::query::ScoredPostsQuery; +use crate::params::EnableAdsBrandSafetyVerdictV2; use std::collections::{HashMap, HashSet}; use std::sync::Arc; use tonic::async_trait; @@ -30,7 +31,7 @@ fn to_safety_label_infos(labels: &SafetyLabelMap) -> impl Iterator for AdsBrandSafetyVfHydrator { async fn hydrate( &self, - _query: &ScoredPostsQuery, + query: &ScoredPostsQuery, candidates: &[PostCandidate], ) -> Vec> { let mut all_ids: HashSet = HashSet::new(); @@ -54,6 +55,12 @@ impl Hydrator for AdsBrandSafetyVfHydrator { let failed_ids: HashSet = batch.failures.keys().copied().collect(); let label_map = batch.labels; + let compute = if query.params.get(EnableAdsBrandSafetyVerdictV2) { + compute_verdict_v2 + } else { + compute_verdict + }; + let mut nsfw_author_seen: u64 = 0; let mut nsfw_author_dropped: u64 = 0; @@ -68,7 +75,7 @@ impl Hydrator for AdsBrandSafetyVfHydrator { let empty = HashMap::new(); let primary_labels = label_map.get(&primary_id).unwrap_or(&empty); - let mut verdict = compute_verdict(primary_labels, primary_id); + let mut verdict = compute(primary_labels, primary_id); let mut safety_labels: Vec = to_safety_label_infos(primary_labels).collect(); @@ -77,7 +84,7 @@ impl Hydrator for AdsBrandSafetyVfHydrator { verdict = worst_verdict(&verdict, &BrandSafetyVerdict::MediumRisk); } else { let qt_labels = label_map.get(&qt_id).unwrap_or(&empty); - verdict = worst_verdict(&verdict, &compute_verdict(qt_labels, qt_id)); + verdict = worst_verdict(&verdict, &compute(qt_labels, qt_id)); safety_labels.extend(to_safety_label_infos(qt_labels)); } } @@ -87,8 +94,7 @@ impl Hydrator for AdsBrandSafetyVfHydrator { verdict = worst_verdict(&verdict, &BrandSafetyVerdict::MediumRisk); } else { let ancestor_labels = label_map.get(&ancestor_id).unwrap_or(&empty); - verdict = - worst_verdict(&verdict, &compute_verdict(ancestor_labels, ancestor_id)); + verdict = worst_verdict(&verdict, &compute(ancestor_labels, ancestor_id)); safety_labels.extend(to_safety_label_infos(ancestor_labels)); } } @@ -96,7 +102,7 @@ impl Hydrator for AdsBrandSafetyVfHydrator { if c.nsfw_author_ads == Some(true) { nsfw_author_seen += 1; let before = verdict; - verdict = worst_verdict(&verdict, &BrandSafetyVerdict::MediumRisk); + verdict = worst_verdict(&verdict, &BrandSafetyVerdict::HighRisk); if verdict != before { nsfw_author_dropped += 1; } @@ -138,6 +144,7 @@ impl Hydrator for AdsBrandSafetyVfHydrator { #[cfg(test)] mod tests { use super::*; + use crate::models::query::ScoredPostsQuery; use xai_safety_label_store::types::SafetyLabelMap; use xai_visibility_filtering::tweet_safety_label::{SafetyLabelFailure, SafetyLabelsBatch}; use xai_x_thrift::tweet_safety_label::{SafetyLabel, SafetyLabelType}; @@ -314,7 +321,7 @@ mod tests { } #[tokio::test] - async fn ancestor_medium_risk_escalates_verdict() { + async fn ancestor_high_risk_escalates_verdict() { let mut safe_labels: SafetyLabelMap = HashMap::new(); safe_labels.insert(SafetyLabelType::GROK_SFA, SafetyLabel::default()); let mut risky_labels: SafetyLabelMap = HashMap::new(); @@ -340,7 +347,7 @@ mod tests { let hydrated = results[0].as_ref().unwrap(); assert_eq!( hydrated.brand_safety_verdict, - Some(BrandSafetyVerdict::MediumRisk) + Some(BrandSafetyVerdict::HighRisk) ); assert!(hydrated .safety_labels @@ -474,7 +481,7 @@ mod tests { } #[tokio::test] - async fn nsfw_author_escalates_to_medium_risk() { + async fn nsfw_author_escalates_to_high_risk() { let mut labels: SafetyLabelMap = HashMap::new(); labels.insert(SafetyLabelType::GROK_SFA, SafetyLabel::default()); let client = Arc::new(FakeVfClient { @@ -497,7 +504,7 @@ mod tests { let hydrated = results[0].as_ref().unwrap(); assert_eq!( hydrated.brand_safety_verdict, - Some(BrandSafetyVerdict::MediumRisk) + Some(BrandSafetyVerdict::HighRisk) ); } diff --git a/home-mixer/candidate_hydrators/ai_trend_feedback_context_hydrator.rs b/home-mixer/candidate_hydrators/ai_trend_feedback_context_hydrator.rs new file mode 100644 index 00000000..795849bb --- /dev/null +++ b/home-mixer/candidate_hydrators/ai_trend_feedback_context_hydrator.rs @@ -0,0 +1,157 @@ +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use crate::params::EnableAiTrendFeedbackContext; +use rand::seq::IndexedRandom; +use std::collections::HashMap; +use std::sync::Arc; +use tonic::async_trait; +use tracing::warn; +use xai_candidate_pipeline::component_library::clients::StratoClient; +use xai_candidate_pipeline::hydrator::Hydrator; + +const TOP_TREND_COUNT: usize = 2; + +pub struct AiTrendFeedbackContextHydrator { + pub strato_client: Arc, +} + +impl AiTrendFeedbackContextHydrator { + fn is_eligible_original_post(candidate: &PostCandidate) -> bool { + candidate.in_reply_to_tweet_id.is_none() + && candidate.retweeted_tweet_id.is_none() + && candidate.ancestors.is_empty() + && candidate.following_replied_user_ids.is_empty() + } + + fn select_feedback_targets(candidate_trends: &[(usize, i64)]) -> HashMap { + let mut frequency: HashMap = HashMap::new(); + for (_, trend_id) in candidate_trends { + *frequency.entry(*trend_id).or_default() += 1; + } + + let mut ranked: Vec<(i64, usize)> = frequency.into_iter().collect(); + ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + ranked.truncate(TOP_TREND_COUNT); + + let mut rng = rand::rng(); + let mut selected: HashMap = HashMap::new(); + let mut used_indices = std::collections::HashSet::new(); + + for (trend_id, _) in ranked { + let candidates_for_trend: Vec = candidate_trends + .iter() + .filter(|(idx, id)| !used_indices.contains(idx) && *id == trend_id) + .map(|(idx, _)| *idx) + .collect(); + let Some(&idx) = candidates_for_trend.choose(&mut rng) else { + continue; + }; + used_indices.insert(idx); + selected.insert(idx, trend_id); + } + + selected + } +} + +#[async_trait] +impl Hydrator for AiTrendFeedbackContextHydrator { + fn enable(&self, query: &ScoredPostsQuery) -> bool { + query.params.get(EnableAiTrendFeedbackContext) + && !query.is_topic_request() + && !query.in_network_only + } + + async fn hydrate( + &self, + _query: &ScoredPostsQuery, + candidates: &[PostCandidate], + ) -> Vec> { + let eligible: Vec<(usize, u64)> = candidates + .iter() + .enumerate() + .filter(|(_, c)| Self::is_eligible_original_post(c)) + .map(|(i, c)| (i, c.tweet_id)) + .collect(); + + let selected = if eligible.is_empty() { + HashMap::new() + } else { + let tweet_ids: Vec = eligible.iter().map(|(_, id)| *id).collect(); + let results = self + .strato_client + .batch_get_tweet_ai_trend(&tweet_ids) + .await; + + let mut candidate_trends: Vec<(usize, i64)> = Vec::new(); + for ((idx, _), result) in eligible.into_iter().zip(results) { + match result { + Ok(Some(trend_id)) if trend_id != 0 => { + candidate_trends.push((idx, trend_id)); + } + Ok(_) => {} + Err(e) => { + warn!("AiTrendFeedbackContextHydrator: strato fetch error: {}", e); + } + } + } + let selected_ids = Self::select_feedback_targets(&candidate_trends); + if selected_ids.is_empty() { + HashMap::new() + } else { + let mut unique_ids: Vec = selected_ids.values().copied().collect(); + unique_ids.sort_unstable(); + unique_ids.dedup(); + let names = self + .strato_client + .batch_get_ai_trend_name(&unique_ids) + .await; + let mut name_by_id: HashMap = HashMap::new(); + for (trend_id, result) in unique_ids.into_iter().zip(names) { + match result { + Ok(Some(name)) if !name.is_empty() => { + name_by_id.insert(trend_id, name); + } + Ok(_) => {} + Err(e) => { + warn!( + "AiTrendFeedbackContextHydrator: trend name fetch error: {}", + e + ); + } + } + } + selected_ids + .into_iter() + .filter_map(|(idx, trend_id)| { + name_by_id + .get(&trend_id) + .cloned() + .map(|name| (idx, (name, trend_id.to_string()))) + }) + .collect() + } + }; + + candidates + .iter() + .enumerate() + .map(|(i, _)| { + let (name, trend_id) = match selected.get(&i) { + Some((n, id)) => (Some(n.clone()), Some(id.clone())), + None => (None, None), + }; + Ok(PostCandidate { + ai_trend_name: name, + ai_trend_id: trend_id, + ..Default::default() + }) + }) + .collect() + } + + fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { + candidate.ai_trend_name = hydrated.ai_trend_name; + candidate.ai_trend_id = hydrated.ai_trend_id; + } +} diff --git a/home-mixer/candidate_hydrators/conversation_gap_ancestor_hydrator.rs b/home-mixer/candidate_hydrators/conversation_gap_ancestor_hydrator.rs index ce0569fe..bf369b34 100644 --- a/home-mixer/candidate_hydrators/conversation_gap_ancestor_hydrator.rs +++ b/home-mixer/candidate_hydrators/conversation_gap_ancestor_hydrator.rs @@ -40,12 +40,13 @@ impl Hydrator for ConversationGapAncestorHydrat candidates .iter() .map(|candidate| { - expand_ancestors_for_gap(&candidate.ancestors, &ancestor_core).map(|ancestors| { - PostCandidate { - tombstone_ancestor_ids: tombstone_ancestor_ids(&ancestors, &ancestor_core), - ancestors, - ..Default::default() - } + let ancestors = expand_ancestors_for_gap(&candidate.ancestors, &ancestor_core) + .unwrap_or_else(|_| candidate.ancestors.to_vec()); + Ok(PostCandidate { + tombstone_ancestor_ids: tombstone_ancestor_ids(&ancestors, &ancestor_core), + ancestor_texts: ancestor_texts(&candidate.ancestors, &ancestor_core), + ancestors, + ..Default::default() }) }) .collect() @@ -54,6 +55,7 @@ impl Hydrator for ConversationGapAncestorHydrat fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { candidate.ancestors = hydrated.ancestors; candidate.tombstone_ancestor_ids = hydrated.tombstone_ancestor_ids; + candidate.ancestor_texts = hydrated.ancestor_texts; } } @@ -89,6 +91,20 @@ fn tombstone_ancestor_ids( .collect() } +fn ancestor_texts( + ancestors: &[u64], + ancestor_core: &HashMap>>, +) -> HashMap { + ancestors + .iter() + .copied() + .filter_map(|id| match ancestor_core.get(&id) { + Some(Ok(Some(data))) if !data.text.is_empty() => Some((id, data.text.clone())), + _ => None, + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; diff --git a/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs b/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs index 0ef172d3..e74280d8 100644 --- a/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs +++ b/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs @@ -1,7 +1,7 @@ use crate::clients::engagement_counts_client::EngagementCountsClient; use crate::models::candidate::{CandidateHelpers, PostCandidate}; use crate::models::query::ScoredPostsQuery; -use crate::params::{ColdStartFollowerCap, EnableEngagementCountsHydration, EnableViewerColdStart}; +use crate::params::{ColdStartFollowerCap, EnableEngagementCountsHydration}; use crate::scorers::author_cold_start::cold_start_base_eligible; use std::collections::HashMap; use std::sync::Arc; @@ -21,6 +21,7 @@ pub struct CachedCounts { repost_count: Option, quote_count: Option, view_count: Option, + view_count_on_home: Option, bookmark_count: Option, } @@ -32,6 +33,7 @@ impl CachedCounts { repost_count: Some(c.retweet_count as i64), quote_count: Some(c.quote_count as i64), view_count: Some(c.view_count), + view_count_on_home: Some(c.view_count_on_home), bookmark_count: Some(c.bookmark_count as i64), } } @@ -43,6 +45,7 @@ impl CachedCounts { repost_count: self.repost_count, quote_count: self.quote_count, view_count: self.view_count, + view_count_on_home: self.view_count_on_home, bookmark_count: self.bookmark_count, ..Default::default() } @@ -56,6 +59,7 @@ fn preserve_counts(c: &PostCandidate) -> PostCandidate { repost_count: c.repost_count, quote_count: c.quote_count, view_count: c.view_count, + view_count_on_home: c.view_count_on_home, bookmark_count: c.bookmark_count, ..Default::default() } @@ -82,9 +86,7 @@ impl CachedHydrator for EngagementCountsHydrato type CacheValue = CachedCounts; fn enable(&self, query: &ScoredPostsQuery) -> bool { - query.params.get(EnableViewerColdStart) - || (!query.has_cached_posts - && (query.params.get(EnableEngagementCountsHydration) || query.is_shadow_traffic)) + query.params.get(EnableEngagementCountsHydration) || query.is_shadow_traffic } fn cache_store(&self) -> &dyn CacheStore { @@ -102,6 +104,7 @@ impl CachedHydrator for EngagementCountsHydrato repost_count: hydrated.repost_count, quote_count: hydrated.quote_count, view_count: hydrated.view_count, + view_count_on_home: hydrated.view_count_on_home, bookmark_count: hydrated.bookmark_count, } } @@ -160,6 +163,7 @@ impl CachedHydrator for EngagementCountsHydrato candidate.repost_count = hydrated.repost_count; candidate.quote_count = hydrated.quote_count; candidate.view_count = hydrated.view_count; + candidate.view_count_on_home = hydrated.view_count_on_home; candidate.bookmark_count = hydrated.bookmark_count; } } @@ -212,16 +216,25 @@ mod tests { async fn enable_matrix() { let h = hydrator(HashMap::new()).await; assert!(h.enable(&query(false, &[(COUNTS, "true")]))); - assert!(h.enable(&query(false, &[(COLD_START, "true"), (COUNTS, "false")]))); - assert!(h.enable(&query(true, &[(COLD_START, "true")]))); - assert!(!h.enable(&query(true, &[(COUNTS, "true"), (COLD_START, "false")]))); - assert!(!h.enable(&query(false, &[(COUNTS, "false"), (COLD_START, "false")]))); - assert!(!h.enable(&query(true, &[(COUNTS, "false"), (COLD_START, "false")]))); + assert!(h.enable(&query(true, &[(COUNTS, "true")]))); + assert!(!h.enable(&query(false, &[(COUNTS, "false")]))); + assert!(!h.enable(&query(true, &[(COUNTS, "false")]))); + let mut shadow = query(true, &[(COUNTS, "false")]); + shadow.is_shadow_traffic = true; + assert!(h.enable(&shadow)); } #[tokio::test] async fn no_cached_posts_hydrates_all() { - let h = hydrator(view_counts(&[(20, 7)])).await; + let h = hydrator(HashMap::from([( + 20, + EngagementCounts { + view_count: 7, + view_count_on_home: 3, + ..Default::default() + }, + )])) + .await; let candidates = vec![PostCandidate { tweet_id: 20, author_id: 2, @@ -232,6 +245,7 @@ mod tests { let q = query(false, &[(COUNTS, "true"), (CAP, "1000")]); let result = h.hydrate_from_client(&q, &candidates).await; assert_eq!(result[0].as_ref().unwrap().view_count, Some(7)); + assert_eq!(result[0].as_ref().unwrap().view_count_on_home, Some(3)); } #[tokio::test] diff --git a/home-mixer/candidate_hydrators/following_blocked_by_hydrator.rs b/home-mixer/candidate_hydrators/following_blocked_by_hydrator.rs new file mode 100644 index 00000000..d714cbd9 --- /dev/null +++ b/home-mixer/candidate_hydrators/following_blocked_by_hydrator.rs @@ -0,0 +1,65 @@ +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use std::sync::Arc; +use tonic::async_trait; +use xai_candidate_pipeline::component_library::clients::SocialGraphClientOps; +use xai_candidate_pipeline::hydrator::Hydrator; + +pub struct FollowingBlockedByHydrator { + socialgraph_client: Arc, +} + +impl FollowingBlockedByHydrator { + pub async fn new(socialgraph_client: Arc) -> Self { + Self { socialgraph_client } + } +} + +#[async_trait] +impl Hydrator for FollowingBlockedByHydrator { + async fn hydrate( + &self, + query: &ScoredPostsQuery, + candidates: &[PostCandidate], + ) -> Vec> { + let user_ids: Vec = candidates + .iter() + .flat_map(|c| c.quoted_user_id.into_iter().chain(c.retweeted_user_id)) + .collect(); + + let blocked_by_user_ids = match self + .socialgraph_client + .check_blocked_by(query.user_id, &user_ids) + .await + { + Ok(ids) => ids, + Err(e) => { + let err_msg = e.to_string(); + return candidates.iter().map(|_| Err(err_msg.clone())).collect(); + } + }; + candidates + .iter() + .map(|candidate| { + let author_blocks_viewer = candidate + .retweeted_user_id + .is_some_and(|uid| blocked_by_user_ids.contains(&uid)); + let quoted_author_blocks_viewer = candidate + .quoted_user_id + .map(|uid| blocked_by_user_ids.contains(&uid)); + Ok(PostCandidate { + author_blocks_viewer: Some(author_blocks_viewer), + quoted_author_blocks_viewer, + ..Default::default() + }) + }) + .collect() + } + + fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { + candidate.author_blocks_viewer = hydrated.author_blocks_viewer; + if hydrated.quoted_author_blocks_viewer.is_some() { + candidate.quoted_author_blocks_viewer = hydrated.quoted_author_blocks_viewer; + } + } +} diff --git a/home-mixer/candidate_hydrators/gizmoduck_hydrator.rs b/home-mixer/candidate_hydrators/gizmoduck_hydrator.rs index 15ff2bda..c2eeb9e9 100644 --- a/home-mixer/candidate_hydrators/gizmoduck_hydrator.rs +++ b/home-mixer/candidate_hydrators/gizmoduck_hydrator.rs @@ -49,6 +49,7 @@ impl CachedHydrator for GizmoduckCandidateHydra retweeted_screen_name: hydrated.retweeted_screen_name.clone(), nsfw_author: hydrated.nsfw_author, nsfw_author_ads: hydrated.nsfw_author_ads, + nsfw_author_phoenix: hydrated.nsfw_author_phoenix, } } @@ -59,6 +60,7 @@ impl CachedHydrator for GizmoduckCandidateHydra retweeted_screen_name: value.retweeted_screen_name, nsfw_author: value.nsfw_author, nsfw_author_ads: value.nsfw_author_ads, + nsfw_author_phoenix: value.nsfw_author_phoenix, ..Default::default() } } @@ -134,6 +136,17 @@ impl CachedHydrator for GizmoduckCandidateHydra label.label_value == LabelValue::POSSIBLY_NSFW_ACCOUNT.0 }) }); + let nsfw_author_phoenix: Option = author.map(|u| { + u.safety.nsfw_user + || u.safety.nsfw_admin + || u.labels.labels.iter().any(|l| { + matches!( + LabelValue(l.label_value), + LabelValue::NSFW_HIGH_PRECISION + | LabelValue::POSSIBLY_NSFW_ACCOUNT + ) + }) + }); Ok(PostCandidate { author_followers_count, @@ -141,6 +154,7 @@ impl CachedHydrator for GizmoduckCandidateHydra retweeted_screen_name, nsfw_author, nsfw_author_ads, + nsfw_author_phoenix, ..Default::default() }) } @@ -158,6 +172,7 @@ impl CachedHydrator for GizmoduckCandidateHydra candidate.retweeted_screen_name = hydrated.retweeted_screen_name; candidate.nsfw_author = hydrated.nsfw_author; candidate.nsfw_author_ads = hydrated.nsfw_author_ads; + candidate.nsfw_author_phoenix = hydrated.nsfw_author_phoenix; } } @@ -174,4 +189,5 @@ pub struct GizmoduckCacheValue { pub retweeted_screen_name: Option, pub nsfw_author: Option, pub nsfw_author_ads: Option, + pub nsfw_author_phoenix: Option, } diff --git a/home-mixer/candidate_hydrators/mod.rs b/home-mixer/candidate_hydrators/mod.rs index 6527acf1..27a4f3db 100644 --- a/home-mixer/candidate_hydrators/mod.rs +++ b/home-mixer/candidate_hydrators/mod.rs @@ -1,10 +1,12 @@ pub mod ads_brand_safety_vf_hydrator; +pub mod ai_trend_feedback_context_hydrator; pub mod bidirectional_follow_hydrator; pub mod blocked_by_hydrator; pub mod conversation_gap_ancestor_hydrator; pub mod core_data_candidate_hydrator; pub mod engagement_counts_hydrator; pub mod filtered_topics_hydrator; +pub mod following_blocked_by_hydrator; pub mod following_replied_users_hydrator; pub mod gizmoduck_hydrator; pub mod in_network_candidate_hydrator; @@ -12,8 +14,10 @@ pub mod language_code_hydrator; pub mod media_info_hydrator; pub mod mutual_follow_jaccard_hydrator; pub mod quote_hydrator; +pub mod quoted_post_text_hydrator; pub mod semantic_id_hydrator; pub mod subscription_hydrator; pub mod topic_feedback_context_hydrator; pub mod tweet_type_metrics_hydrator; pub mod vf_candidate_hydrator; +pub mod vf_following_candidate_hydrator; diff --git a/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs b/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs new file mode 100644 index 00000000..7c6041bd --- /dev/null +++ b/home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs @@ -0,0 +1,106 @@ +use crate::clients::tweet_entity_service_client::TESClient; +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use tonic::async_trait; +use xai_candidate_pipeline::hydrator::Hydrator; + +pub struct QuotedPostTextHydrator { + pub tes_client: Arc, +} + +impl QuotedPostTextHydrator { + pub fn new(tes_client: Arc) -> Self { + Self { tes_client } + } +} + +#[async_trait] +impl Hydrator for QuotedPostTextHydrator { + async fn hydrate( + &self, + _query: &ScoredPostsQuery, + candidates: &[PostCandidate], + ) -> Vec> { + let quoted_ids: Vec = candidates + .iter() + .filter_map(|c| c.quoted_tweet_id) + .collect::>() + .into_iter() + .collect(); + + let quoted_core = if quoted_ids.is_empty() { + HashMap::new() + } else { + self.tes_client.get_tweet_core_datas(quoted_ids).await + }; + + candidates + .iter() + .map(|candidate| { + Ok(PostCandidate { + quoted_tweet_text: candidate.quoted_tweet_id.and_then(|id| { + match quoted_core.get(&id) { + Some(Ok(Some(data))) if !data.text.is_empty() => { + Some(data.text.clone()) + } + _ => None, + } + }), + ..Default::default() + }) + }) + .collect() + } + + fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { + candidate.quoted_tweet_text = hydrated.quoted_tweet_text; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::clients::tweet_entity_service_client::MockTESClient; + use xai_core_entities::entities::PureCoreData; + + #[tokio::test] + async fn fills_quoted_text_and_leaves_non_quotes_empty() { + let mut core_data = HashMap::new(); + core_data.insert( + 99, + Some(PureCoreData { + text: "quoted text".to_string(), + ..Default::default() + }), + ); + let client = Arc::new(MockTESClient { + core_data, + ..Default::default() + }); + let hydrator = QuotedPostTextHydrator::new(client as Arc); + + let mut with_quote = PostCandidate { + tweet_id: 1, + quoted_tweet_id: Some(99), + ..Default::default() + }; + let mut without_quote = PostCandidate { + tweet_id: 2, + ..Default::default() + }; + + let hydrated = hydrator + .hydrate( + &ScoredPostsQuery::default(), + &[with_quote.clone(), without_quote.clone()], + ) + .await; + hydrator.update(&mut with_quote, hydrated[0].clone().unwrap()); + hydrator.update(&mut without_quote, hydrated[1].clone().unwrap()); + + assert_eq!(with_quote.quoted_tweet_text.as_deref(), Some("quoted text")); + assert_eq!(without_quote.quoted_tweet_text, None); + } +} diff --git a/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs index 24cca235..b6554ed6 100644 --- a/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs +++ b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs @@ -12,7 +12,7 @@ use xai_twittercontext_proto::TwitterContextViewer; use xai_visibility_filtering::models::{Action, FilteredReason}; use xai_visibility_filtering::vf_client::SafetyLevel; use xai_visibility_filtering::vf_client::SafetyLevel::{TimelineHome, TimelineHomeRecommendations}; -use xai_visibility_filtering::vf_client::VfClient; +use xai_visibility_filtering::vf_client::{TweetVisibility, VfClient}; pub struct VFCandidateHydrator { pub strato_vf_client: Arc, @@ -36,7 +36,7 @@ impl VFCandidateHydrator { safety_level: SafetyLevel, for_user_id: u64, context: Option, - ) -> HashMap>> { + ) -> HashMap> { if tweet_ids.is_empty() { return HashMap::new(); } @@ -83,6 +83,8 @@ impl Hydrator for VFCandidateHydrator { } } + in_network_ids.sort_unstable(); + in_network_ids.dedup(); oon_ids.sort_unstable(); oon_ids.dedup(); @@ -104,8 +106,12 @@ impl Hydrator for VFCandidateHydrator { let (in_network_result, oon_result) = join(in_network_future, oon_future).await; let mut all_results: HashMap>> = HashMap::new(); - all_results.extend(in_network_result); - all_results.extend(oon_result); + all_results.extend( + oon_result + .into_iter() + .chain(in_network_result) + .map(|(id, r)| (id, r.map(|t| t.reason))), + ); let mut hydrated_candidates = Vec::with_capacity(candidates.len()); for candidate in candidates { @@ -136,7 +142,7 @@ impl Hydrator for VFCandidateHydrator { } } -fn should_drop_ancillary( +pub(crate) fn should_drop_ancillary( candidate: &PostCandidate, vf_results: &HashMap>>, ) -> bool { diff --git a/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs new file mode 100644 index 00000000..19273fc0 --- /dev/null +++ b/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs @@ -0,0 +1,98 @@ +use crate::candidate_hydrators::vf_candidate_hydrator::should_drop_ancillary; +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use crate::params::EnableXaiVfClient; +use anyhow::Result; +use std::collections::HashMap; +use std::sync::Arc; +use tonic::async_trait; +use xai_candidate_pipeline::hydrator::Hydrator; +use xai_twittercontext_proto::GetTwitterContextViewer; +use xai_visibility_filtering::models::FilteredReason; +use xai_visibility_filtering::vf_client::SafetyLevel::TimelineHome; +use xai_visibility_filtering::vf_client::VfClient; + +pub struct VFFollowingCandidateHydrator { + pub strato_vf_client: Arc, + pub xai_vf_client: Arc, +} + +impl VFFollowingCandidateHydrator { + pub fn new( + strato_vf_client: Arc, + xai_vf_client: Arc, + ) -> Self { + Self { + strato_vf_client, + xai_vf_client, + } + } +} + +#[async_trait] +impl Hydrator for VFFollowingCandidateHydrator { + async fn hydrate( + &self, + query: &ScoredPostsQuery, + candidates: &[PostCandidate], + ) -> Vec> { + let context = query.get_viewer(); + let client = if query.params.get(EnableXaiVfClient) { + &self.xai_vf_client + } else { + &self.strato_vf_client + }; + + let mut post_ids: Vec = Vec::new(); + for candidate in candidates { + post_ids.push(candidate.tweet_id); + post_ids.extend(candidate.ancestors.iter().copied()); + if let Some(quoted_post_id) = candidate.quoted_tweet_id { + post_ids.push(quoted_post_id); + } + if let Some(reposted_post_id) = candidate.retweeted_tweet_id { + post_ids.push(reposted_post_id); + } + } + post_ids.sort_unstable(); + post_ids.dedup(); + + let all_results: HashMap>> = if post_ids.is_empty() { + HashMap::new() + } else { + client + .get_result(post_ids, TimelineHome, query.user_id, context) + .await + .into_iter() + .map(|(id, r)| (id, r.map(|t| t.reason))) + .collect() + }; + + let mut hydrated_candidates = Vec::with_capacity(candidates.len()); + for candidate in candidates { + let primary_result = all_results.get(&candidate.tweet_id); + let visibility_reason = match primary_result { + Some(Ok(Some(reason))) => Some(reason.clone()), + _ => None, + }; + + let drop_ancillary = should_drop_ancillary(candidate, &all_results); + + let hydrated = match primary_result { + Some(Err(err)) => Err(err.to_string()), + _ => Ok(PostCandidate { + visibility_reason, + drop_ancillary_posts: Some(drop_ancillary), + ..Default::default() + }), + }; + hydrated_candidates.push(hydrated); + } + hydrated_candidates + } + + fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { + candidate.visibility_reason = hydrated.visibility_reason; + candidate.drop_ancillary_posts = hydrated.drop_ancillary_posts; + } +} diff --git a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index dd38792c..e8c9c397 100644 --- a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs +++ b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs @@ -1,4 +1,5 @@ use crate::candidate_hydrators::ads_brand_safety_vf_hydrator::AdsBrandSafetyVfHydrator; +use crate::candidate_hydrators::ai_trend_feedback_context_hydrator::AiTrendFeedbackContextHydrator; use crate::candidate_hydrators::bidirectional_follow_hydrator::BidirectionalFollowHydrator; use crate::candidate_hydrators::blocked_by_hydrator::BlockedByHydrator; use crate::candidate_hydrators::core_data_candidate_hydrator::CoreDataCandidateHydrator; @@ -37,12 +38,12 @@ use crate::clients::vm_ranker_client::{MockVMRankerClient, ProdVMRankerClient, V use crate::filters::age_filter::AgeFilter; use crate::filters::ancillary_vf_filter::AncillaryVFFilter; use crate::filters::author_socialgraph_filter::AuthorSocialgraphFilter; +use crate::filters::brazil_2026_election_filter::Brazil2026ElectionFilter; use crate::filters::core_data_hydration_filter::CoreDataHydrationFilter; use crate::filters::dedup_conversation_filter::DedupConversationFilter; use crate::filters::drop_duplicates_filter::DropDuplicatesFilter; use crate::filters::ineligible_subscription_filter::IneligibleSubscriptionFilter; use crate::filters::inventory_holdout_filter::InventoryHoldoutFilter; -use crate::filters::muted_keyword_filter::MutedKeywordFilter; use crate::filters::new_user_min_engagement_filter::NewUserMinEngagementFilter; use crate::filters::oon_nsfw_simclusters_filter::OONNsfwSimclustersFilter; use crate::filters::oon_retweet_reply_filter::OONRetweetReplyFilter; @@ -54,6 +55,7 @@ use crate::filters::self_tweet_filter::SelfTweetFilter; use crate::filters::topic_ids_filter::TopicIdsFilter; use crate::filters::vf_filter::VFFilter; use crate::filters::video_filter::VideoFilter; +use crate::filters::viewer_muted_keyword_filter::ViewerMutedKeywordFilter; use crate::models::candidate::PostCandidate; use crate::models::query::ScoredPostsQuery; use crate::params; @@ -353,8 +355,18 @@ impl PhoenixCandidatePipeline { Box::new(PreviouslySeenPostsFilter), Box::new(PreviouslySeenPostsBackupFilter), Box::new(PreviouslyServedPostsFilter), - Box::new(MutedKeywordFilter::new()), + Box::new(ViewerMutedKeywordFilter::new()), Box::new(AuthorSocialgraphFilter), + // Brazil 2026 election filter + + // Application providers that use a recommendation system for users must exclude from the + // results the channels and profiles reported to the Electoral Court under the terms of + // § 1º of this article and, except in cases of paid boosting, the content posted on them. + + // https://dadosabertos.tse.jus.br/dataset/candidatos-2026 + + // OmarAzizSenador deleted his account at the time this code was written. + Box::new(Brazil2026ElectionFilter), Box::new(VideoFilter), Box::new(TopicIdsFilter), Box::new(NewUserMinEngagementFilter), @@ -383,14 +395,11 @@ impl PhoenixCandidatePipeline { feature_switches, )); let author_cold_start = crate::scorers::author_cold_start::AuthorColdStart { author_rules }; - let ranking_scorer = Box::new(RankingScorer { - author_cold_start: author_cold_start.clone(), - }); + let ranking_scorer = Box::new(RankingScorer { author_cold_start }); let xds_vm_ranker_client = super::build_vm_ranker_xds_client(vm_ranker_xds).await; let vm_ranker = Box::new(VMRanker { client: vm_ranker_client, xds_client: xds_vm_ranker_client, - author_cold_start, }); let scorers: Vec>> = vec![phoenix_scorer, ranking_scorer, vm_ranker]; @@ -412,6 +421,9 @@ impl PhoenixCandidatePipeline { Box::new(TopicFeedbackContextHydrator { strato_client: strato_client.clone(), }), + Box::new(AiTrendFeedbackContextHydrator { + strato_client: strato_client.clone(), + }), ]; let post_selection_filters: Vec>> = vec![ diff --git a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs index 90cac898..1ea8adf6 100644 --- a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs +++ b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs @@ -1,13 +1,17 @@ use crate::candidate_hydrators::ads_brand_safety_vf_hydrator::AdsBrandSafetyVfHydrator; use crate::candidate_hydrators::conversation_gap_ancestor_hydrator::ConversationGapAncestorHydrator; use crate::candidate_hydrators::core_data_candidate_hydrator::CoreDataCandidateHydrator; +use crate::candidate_hydrators::following_blocked_by_hydrator::FollowingBlockedByHydrator; +use crate::candidate_hydrators::quoted_post_text_hydrator::QuotedPostTextHydrator; use crate::candidate_hydrators::tweet_type_metrics_hydrator::TweetTypeMetricsHydrator; -use crate::candidate_hydrators::vf_candidate_hydrator::VFCandidateHydrator; +use crate::candidate_hydrators::vf_following_candidate_hydrator::VFFollowingCandidateHydrator; use crate::clients::night_owl_client::{MockNightOwlClient, NightOwlClient, ProdNightOwlClient}; use crate::clients::s2s::{S2S_CHAIN_PATH, S2S_CRT_PATH, S2S_KEY_PATH}; use crate::clients::tweet_entity_service_client::{MockTESClient, ProdTESClient, TESClient}; use crate::filters::ancillary_vf_filter::AncillaryVFFilter; +use crate::filters::author_socialgraph_filter::AuthorSocialgraphFilter; use crate::filters::following_retweet_deduplication_filter::FollowingRetweetDeduplicationFilter; +use crate::filters::following_viewer_muted_keyword_filter::FollowingViewerMutedKeywordFilter; use crate::filters::self_reply_chain_filter::SelfReplyChainFilter; use crate::filters::vf_filter::VFFilter; use crate::models::candidate::PostCandidate; @@ -18,6 +22,9 @@ use crate::sources::following_night_owl_source::FollowingNightOwlSource; use std::sync::Arc; use tonic::async_trait; use xai_candidate_pipeline::candidate_pipeline::CandidatePipeline; +use xai_candidate_pipeline::component_library::clients::{ + MockSocialGraphClient, SocialGraphClient, SocialGraphClientOps, +}; use xai_candidate_pipeline::filter::Filter; use xai_candidate_pipeline::hydrator::Hydrator; use xai_candidate_pipeline::query_hydrator::QueryHydrator; @@ -48,6 +55,7 @@ impl ReverseChronPostsPipeline { strato_vf_client, xai_vf_client, vf_safety_labels_client, + socialgraph_client, ) = tokio::join!( async { Arc::new( @@ -92,6 +100,18 @@ impl ReverseChronPostsPipeline { .with_max_batch_size(50), ) as Arc }, + async { + Arc::new( + SocialGraphClient::new( + datacenter, + &S2S_CHAIN_PATH, + &S2S_CRT_PATH, + &S2S_KEY_PATH, + ) + .await + .expect("Failed to create flock SocialGraphClient"), + ) as Arc + }, ); Self::build( @@ -100,6 +120,7 @@ impl ReverseChronPostsPipeline { strato_vf_client, xai_vf_client, vf_safety_labels_client, + socialgraph_client, ) .await } @@ -111,6 +132,7 @@ impl ReverseChronPostsPipeline { Arc::new(MockVfClient) as Arc, Arc::new(MockVfClient) as Arc, Arc::new(MockTweetSafetyLabelClient) as Arc, + Arc::new(MockSocialGraphClient) as Arc, ) .await } @@ -121,6 +143,7 @@ impl ReverseChronPostsPipeline { strato_vf_client: Arc, xai_vf_client: Arc, vf_safety_labels_client: Arc, + socialgraph_client: Arc, ) -> Self { let sources: Vec>> = vec![Box::new(FollowingNightOwlSource { @@ -129,24 +152,35 @@ impl ReverseChronPostsPipeline { let hydrators: Vec>> = vec![ Box::new(CoreDataCandidateHydrator::new(Arc::clone(&tes_client)).await), - Box::new(ConversationGapAncestorHydrator::new(tes_client)), + Box::new(ConversationGapAncestorHydrator::new(Arc::clone( + &tes_client, + ))), + Box::new(QuotedPostTextHydrator::new(tes_client)), ]; let filters: Vec>> = vec![ Box::new(FollowingRetweetDeduplicationFilter), + Box::new(FollowingViewerMutedKeywordFilter::new()), Box::new(SelfReplyChainFilter), ]; let post_selection_hydrators: Vec>> = vec![ - Box::new(VFCandidateHydrator::new(strato_vf_client, xai_vf_client).await), + Box::new(FollowingBlockedByHydrator::new(socialgraph_client).await), + Box::new(VFFollowingCandidateHydrator::new( + strato_vf_client, + xai_vf_client, + )), Box::new(AdsBrandSafetyVfHydrator { client: vf_safety_labels_client, }), Box::new(TweetTypeMetricsHydrator::new()), ]; - let post_selection_filters: Vec>> = - vec![Box::new(VFFilter), Box::new(AncillaryVFFilter)]; + let post_selection_filters: Vec>> = vec![ + Box::new(AuthorSocialgraphFilter), + Box::new(VFFilter), + Box::new(AncillaryVFFilter), + ]; Self { sources, diff --git a/home-mixer/filters/brazil_2026_election_filter.rs b/home-mixer/filters/brazil_2026_election_filter.rs new file mode 100644 index 00000000..04bb9a09 --- /dev/null +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -0,0 +1,5416 @@ +use rustc_hash::FxHashSet; +use std::sync::LazyLock; + +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use xai_candidate_pipeline::filter::{Filter, FilterResult}; + +// Brazil 2026 election filter + +// Art. 28 § 1º-A of Electoral Resolution No. 23.610: Application providers that use a +// recommendation system for users must exclude from the results the channels and +// profiles reported to the Electoral Court under the terms of § 1º of this article +// and, except in cases of paid boosting, the content posted on them. + +// https://dadosabertos.tse.jus.br/dataset/candidatos-2026 + +// User ids below are obfuscated; usernames are included for transparency. + +// We believe the account @ABR reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @ACMNETO_ reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @ALESILVAOFICIAL reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @BENICIODIAS reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @CARLOSSAMPAIO reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @CARLOSVIANA reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @CLECIACARVALHO1 reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @DAYSE reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @PAULODIMELO reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @PSTUPE reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @RICAR reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @SERGINHOCAXIAS reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// We believe the account @ZANAAMANDA reported by the candidate is not the candidate's actual account, so we are not currently filtering it. +// @ADRIANASOUSAPIAUI no live account found. +// @AGOLDBACH no live account found. +// @AHELIXO no live account found. +// @ALCEU_ALCEUMOREIRA no live account found. +// @ALEXROSETI no live account found. +// @ALKORAP1 no live account found. +// @ARAFETHNASREDDINE no live account found. +// @BETORICHAOFICIAL no live account found. +// @BRUNOPORTODEALMEIDA no live account found. +// @CAIOZMENDONCA no live account found. +// @CHARLES067277 no live account found. +// @CRISTINAGRAEM no live account found. +// @DANIELBRSOARES no live account found. +// @DANILOBALASOFICIAL no live account found. +// @DANILOTORRES100 no live account found. +// @DELEGADOEGUCHI no live account found. +// @DEMAOLIVEIRA70 no live account found. +// @DEPCELSOSABINO no live account found. +// @DEPLUANAREGIA no live account found. +// @DEPUTADOALTAIRSILVA no live account found. +// @DUARTEJR70 no live account found. +// @DUDUSIVINSKI no live account found. +// @EDSONSANTOSRJ no live account found. +// @EUANGELAGARCIALINKTREE no live account found. +// @EXPEDITOFUCAP no live account found. +// @FADAPSICANALISE no live account found. +// @FEDERALFELICIO no live account found. +// @GERSONBURMANNIV no live account found. +// @GSMA1986 no live account found. +// @IGORRAYARN no live account found. +// @JAIZAMETODIO no live account found. +// @LEOMASCARENHASP no live account found. +// @LUCIANALIPPI30 no live account found. +// @LUCIANAOROZIMBO no live account found. +// @LUIZPAULOCURVELO no live account found. +// @MARCELOSILVACAMPINAS no live account found. +// @MEUCANA669499 no live account found. +// @MIRCOCORONETTI no live account found. +// @NADIAGERHARD no live account found. +// @NETOFEITOSA6891 no live account found. +// @PATRICIACRIZANTO2 no live account found. +// @PAULOMOURAOTO no live account found. +// @PEDROPONCIOBE no live account found. +// @PELUCIO_LUCA no live account found. +// @POLICIALPAULOBASTOS no live account found. +// @PRADOCORONEL no live account found. +// @SUSANNAPFEDERAL no live account found. +// @TIAKEYLAECIA no live account found. +// @TWITTERADRIANAACCORSI no live account found. +// @XIGORPORTO no live account found. +// @_ANDREDOPRADO no live account found. +// @_EDUARDOMANTOAN no live account found. + +/// User ids reported to the Electoral Court for the Brazil 2026 election. +static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(|| { + FxHashSet::from_iter([ + // @madeleinelacsko + 9179462, + // @nanacachen + 13258012, + // @renildo + 14160928, + // @renatoroseno + 14492205, + // @pedro_lupion + 15022409, + // @erasintetica + 15585094, + // @soninhafrancine + 15768105, + // @AecioNeves + 15828014, + // @tatyanavaleria + 15908023, + // @ClariceChacon + 16524287, + // @rosilenedf + 16526497, + // @raulchristiano + 16976843, + // @Rafael_Parente + 16979783, + // @RicardoFabrizio + 17084658, + // @marlonluz + 17244270, + // @eduardopinheiro + 17525670, + // @falcaopatos + 18950196, + // @diegotavares + 19546911, + // @RubinhoDivi + 19617146, + // @depmariomotta + 19665407, + // @depchinaglia + 19723670, + // @Sen_Cristovam + 20242549, + // @marcelvanhattem + 21069302, + // @lgmbrasilia + 21438601, + // @nilvanferreira + 21571098, + // @RafaellMilas + 21857669, + // @chagasvieira + 22035789, + // @GabrielSouza_RS + 22480147, + // @Pimenta13Br + 22864100, + // @fabriziomeller + 22908605, + // @rigotto + 23443097, + // @gustavotutuca + 24055827, + // @camasao50 + 24103292, + // @radiovaldo + 24588898, + // @euserafimcorrea + 24761475, + // @rafangeli + 25164669, + // @pauloteixeira13 + 25562342, + // @crismonteirosp + 25841252, + // @ManuelaDavila + 25858078, + // @tomioyano + 26214420, + // @CintyaMuniz + 26284058, + // @rodrigolimamdh + 26498680, + // @Biango + 26560559, + // @samiabomfim + 27703690, + // @profsta + 27883076, + // @Jfelippeneto + 28177866, + // @depguibismarck + 29026134, + // @RicardoFerraco + 29434532, + // @mendoncafilho + 29803424, + // @hamiltonassis + 30247132, + // @MarcoMartinsSP + 30495186, + // @maragabrilli + 30978177, + // @gleisi + 31139434, + // @Miwky + 31401730, + // @ronaldocaiado + 32101740, + // @lunazarattini + 32160446, + // @cirogomes + 33374761, + // @gabrielmagno_13 + 33717372, + // @ticokuzma + 33759666, + // @CristinaMel + 33911227, + // @aldenorlima + 33973761, + // @_sidney_ + 34083807, + // @SharleneAZ + 34286009, + // @vimarchese + 34374755, + // @Peter_Costa + 34430921, + // @RafaCupertino + 34618485, + // @BetoRicha + 34665220, + // @marconiperillo + 34790309, + // @Donato_PT + 34795040, + // @jooliveirapb + 34909888, + // @denilsonsoares + 34955073, + // @julianapt + 35105584, + // @LeonelRadde + 35268237, + // @caduxavier + 35470350, + // @marcofeliciano + 35805725, + // @aavasantiago + 35827365, + // @NetoAM + 36145775, + // @rfalcao13 + 36248739, + // @marciokieller + 36403221, + // @ThiagoMassuda + 36694689, + // @PauloMartins10 + 36714942, + // @marceloverly + 36733456, + // @RenanCeschin + 36753162, + // @AlicePortugal + 36971658, + // @edinhosilva + 37055387, + // @PompeodeMattos + 37320286, + // @eniobritodesa + 37654300, + // @pauloserra_sp + 37700244, + // @dep_geraldo + 37711911, + // @SouCrisCaiado + 37745764, + // @ruialves10_ + 37949658, + // @mauricioscalco + 38212742, + // @MarcioNakashima + 38585953, + // @inacioarruda + 38659819, + // @LucasCalilGo + 38929561, + // @jandira_feghali + 39116717, + // @tourinhopedro + 39818820, + // @murilogaldinopb + 40040727, + // @FlavioBolsonaro + 40053694, + // @BetoAlbuquerque + 40463380, + // @alexandrekalil + 40721032, + // @RomanelliPR + 41271043, + // @jessicamichels + 41355867, + // @nidepp + 41403744, + // @EGCARREIRA + 41410587, + // @pauloabiackel + 41590536, + // @profcassiano + 42106436, + // @KleberfreireAdv + 42202762, + // @RaoniMendes + 42284334, + // @lcbusato + 42300711, + // @Alice_Portugal + 42304398, + // @DepArthurMaia + 42454330, + // @marciofrancasp + 42455446, + // @kruke1 + 42487937, + // @Domingos_Neto + 42626592, + // @ArlenSantiago + 42630237, + // @Francischini_ + 42736936, + // @onyxlorenzoni + 43041690, + // @pedropaulo + 43189774, + // @axelgrael + 43326346, + // @owagnerpro + 43364776, + // @edurodrigues_25 + 43856097, + // @profleomatos + 44070650, + // @marcelorangel1 + 44107208, + // @Fufaazevedo + 44153555, + // @renatasouzario + 44195955, + // @duarte_nogueira + 44455876, + // @rafaelcampelo + 44460118, + // @rafamacris + 44690902, + // @perpetua_acre + 44693900, + // @pedrosuplicy30 + 44739698, + // @FelixMendoncaJr + 45235841, + // @PedroTaquesMT + 45448098, + // @CarlosZarattini + 45473463, + // @MarceloFreixo + 45870897, + // @alicedrummond + 46204905, + // @adriano_viana + 46229863, + // @RodrigoSchroder + 46313017, + // @Igorbaimapsol + 46647000, + // @SolangeFreitas + 46891928, + // @betinhogomes + 46958570, + // @DanielVilela15 + 47371488, + // @RicMarques_RM + 47387995, + // @chrispuppi + 47418811, + // @zeca_dirceu + 47461491, + // @augustocury + 47529845, + // @policarpodf + 47851105, + // @Altineu + 47991805, + // @fernandojordao + 48062025, + // @marciojerry + 48122425, + // @eduardopaes + 48298703, + // @gilbertokassab + 49051293, + // @senadorcidgomes + 49089646, + // @Camilo_Capi + 49339134, + // @paulopinheirorj + 49437070, + // @nelioaguiarstm + 49597151, + // @guilhermepasin + 49616087, + // @cicerolucena + 49632512, + // @bethsahao + 49790504, + // @FernandoLadeia + 49818264, + // @ZimbaldiRafa + 49848446, + // @acmnetoba + 49966078, + // @ReginaldoLopes + 50088692, + // @RicardoBarrosPP + 50101324, + // @zecadopt + 50133997, + // @danielxdonizet + 50360881, + // @anyortiz + 50430144, + // @ludiocabral + 50514890, + // @deppauloguedes + 50713061, + // @priscilakrause + 51066167, + // @advcorrea + 51169114, + // @noelnit + 51178338, + // @VandinhoLeite + 51483562, + // @VanildaBordieri + 51498674, + // @afranioboppre50 + 51735736, + // @blogdogarotinho + 51834513, + // @AlexandreCuri + 51882661, + // @celinaleao + 51935259, + // @DarcioVix + 51946581, + // @jilmartatto + 52045368, + // @thigagliasso + 52297557, + // @DeputadoBacelar + 52364013, + // @amastha2026 + 52371426, + // @JuniorMochi + 52483984, + // @ADALBERTO_MDB + 52540070, + // @duiliodecastro + 52564555, + // @aknoploch + 52690350, + // @OsmarTerra + 52722451, + // @lidicedamata + 52724814, + // @DepTraiano + 52736824, + // @GuilhermePaz + 52750366, + // @rosanedopv + 52842057, + // @RodrigoGuedesam + 52954632, + // @joliveirambl + 52971495, + // @RollembergPSB + 53050115, + // @jorginhomello + 53073647, + // @neyleprevost + 53163365, + // @DayseHansa + 53473894, + // @renanroto + 53544192, + // @advandrebarros + 53681220, + // @AngeloAlmeidaBA + 53719510, + // @paulosalimmaluf + 53776600, + // @tatiroque + 53911544, + // @kikosilveira + 53998138, + // @VINIANZILIERO + 54316589, + // @nelsinhotrad + 54412355, + // @RubensOtoni + 54539654, + // @DeputadoWelter + 54545619, + // @vanderloubet + 54547924, + // @xuxudalmolin + 54585047, + // @marcosdoval + 54600557, + // @Reimont + 54612905, + // @flavinhocn + 54735703, + // @mineiroptrn + 54897084, + // @luizaerundina + 55022037, + // @RodLago + 55152459, + // @simaopedro_SP + 55366483, + // @silvionavarro + 55366770, + // @HomeroMarchese + 55370663, + // @giriboni + 55385927, + // @RobertoCarlosTB + 55450890, + // @BiologoHenrique + 55544977, + // @cristianostpr + 55558067, + // @stephanesjunior + 55607796, + // @eduardo_vidal + 55686825, + // @vitor_bicca + 56081072, + // @Jrsantosrosa + 56400752, + // @EduPinheiro7022 + 56409767, + // @MauricioPeixer + 56444778, + // @Iran_Barbosa + 56466859, + // @dep_acoutinho + 56480030, + // @joaonatel + 56610250, + // @FredCostaDep + 56689975, + // @venezianovital + 56734024, + // @marlonreis + 56829466, + // @romerojuca + 56843641, + // @rosedefreitas + 56864413, + // @sandesjunior + 56922936, + // @fernandapsol + 57044549, + // @josefortunati + 57106033, + // @DelegadoMeneses + 57108807, + // @RobinsonFaria + 57151676, + // @HarifeViegas + 57163107, + // @Rafaelpicciani + 57529926, + // @CanzianiAlex + 57641073, + // @depHugoLeal + 57771926, + // @gabriellimambl + 58029195, + // @KarenGregoris + 58044789, + // @marinapassadore + 58247896, + // @AlmeidaMarcus + 58293305, + // @AlineMariano_pe + 58314645, + // @fredlinhares + 58328715, + // @yurimourarj + 58340662, + // @depluizfernando + 58483633, + // @Dyamondharper + 58517952, + // @carloselula + 58538668, + // @welbert__pedro + 58661156, + // @GugaJP + 58878180, + // @LeonoraPerico + 58911503, + // @adrianofritz + 59135814, + // @JutayMeneses + 59243227, + // @fernandofilhope + 59323812, + // @louiselima_md + 59346999, + // @CarlosBezerraJr + 59483409, + // @fernandabarth + 59534429, + // @brunotopete + 59855833, + // @subgonzagamg + 59868993, + // @nandopinheiro22 + 60150569, + // @marcio_motta + 60469505, + // @anapaulagold + 60708760, + // @claudioapolina + 60731692, + // @Isquierdorio + 60805457, + // @andrerodini + 60879681, + // @luisbenambl + 60983130, + // @CovattiFilho + 60994156, + // @FelicioRamuth + 61119320, + // @franzepiaui + 61190865, + // @mariocaixa + 61196446, + // @SigaPepeVargas + 61208942, + // @ale_campelo + 61325857, + // @honoratopvh + 61472552, + // @pauloabarbosa + 61579803, + // @carlaayres + 62286707, + // @juliacasamasso + 62289167, + // @jaqueswagner + 62501888, + // @HenriqueFontana + 62804559, + // @fabiotokarski + 62906951, + // @mariadorosario + 63118359, + // @BohnGass + 63127680, + // @maurotramonte + 63158278, + // @luizsarraf + 63164895, + // @FilipeSabara + 63474515, + // @orlandopesoti + 63494090, + // @leandrograss + 63507573, + // @marceloaro + 63510130, + // @vereadorsamuel1 + 63817204, + // @JoseAirtonPT + 63868587, + // @Bernarditv + 64297102, + // @Bobadra + 64302060, + // @antonionetopdt + 64391577, + // @Marco_Brasil + 64434503, + // @wladmesquita + 64460310, + // @josenunes_ARI + 64482750, + // @hmecabo + 64493005, + // @helencabral13 + 64493500, + // @romulorippa + 64605196, + // @JulioLopesRio + 64755437, + // @livioluciano + 65059970, + // @tenentemelo + 65194252, + // @FatimaCleidePT + 65491379, + // @Glauber_Braga + 65720380, + // @ladyfontenelle + 65763684, + // @marcoslatino + 65972973, + // @fabio_novo + 66413485, + // @OgierBuchi + 66459281, + // @miguelcoelhope + 66525428, + // @Brigadeiromma + 66529851, + // @VerGuilherme + 66704302, + // @cantojocelito + 66719730, + // @betodoisaum + 66753261, + // @mateus_simoesmg + 66789220, + // @iriny_13 + 66810575, + // @carlosgiannazi + 66814653, + // @Priscilaromano + 66900651, + // @Assis_Gurgacz_N + 66982240, + // @monica_benicio + 67002637, + // @UlissesMaia + 67009496, + // @anaperugini + 67061352, + // @ciro_nogueira + 67098726, + // @LarissaPucca + 67184191, + // @PattyParente + 67234932, + // @acirgurgacz + 67601773, + // @gersongabrielli + 67727566, + // @gleisonpego + 67737645, + // @eduardoreiner + 67947379, + // @celamericosp + 67978929, + // @lucasmortimer + 68154178, + // @fabiofeijo + 68181274, + // @andrekamai + 68404402, + // @costa_rui + 68466700, + // @AloisioParana + 68486960, + // @MarcosSantanaSC + 68554516, + // @clarianabarao + 68578439, + // @joaopaulorillo + 68694603, + // @CarlosBolsonaro + 68712576, + // @ProfIsrael + 68719944, + // @DatenaOficial + 68722955, + // @StelaFarias + 68759158, + // @DenisAndia + 68763092, + // @vanbrandao + 68891838, + // @acrodriguessp + 69016521, + // @noraldinojunior + 69020170, + // @juliophilbert + 69073488, + // @MonicaSeixas + 69370960, + // @esuplicy + 69373037, + // @juniorsinforma + 69874322, + // @Gyselle_Soares + 70131422, + // @ernanipolo + 70284501, + // @pedroribeiropdt + 70301797, + // @janetepieta + 70361625, + // @lucarminatti + 70405671, + // @luis_fabuloso + 70443471, + // @deputadoismael + 70453020, + // @jonasdonizette_ + 70956209, + // @rdlorenzoni + 71056246, + // @sen_wellington + 71065638, + // @charlesribeiro_ + 71098452, + // @lindberghfarias + 71310152, + // @FabioFayad + 71345547, + // @Daniel_PCdoB + 71545154, + // @FaissalCalil + 71602909, + // @NabilBondukiSP + 71613541, + // @lohannaf + 71917109, + // @Maristeladutra + 71953093, + // @SANDRAFARAJ + 72296368, + // @maricarvalhoro + 72556466, + // @aaluisfernando + 72575848, + // @DepPatriciaAlba + 72949846, + // @marcoalba + 73034114, + // @EmidioDeSouza_ + 73217377, + // @pedrolimasjc + 73288051, + // @SamuelMalafaia + 73442547, + // @rowennabrito + 73478714, + // @pauderney + 73803827, + // @geraldoalckmin + 74215006, + // @WaldeckCarneiro + 74234605, + // @AlencarBraga13 + 74243131, + // @eduardobismarck + 74361905, + // @LeurLomantoJr + 74538721, + // @RogerPerestelo + 74583511, + // @Casagrande_ES + 74722174, + // @luizcoutopt + 74738674, + // @VICENTINHOPT + 74762633, + // @MCrivella + 74860967, + // @agenorsantospa + 74867163, + // @caetraven + 75035844, + // @profpaulamarisa + 75058892, + // @dasilvabenedita + 75060031, + // @cassioafsoares + 75128422, + // @kaiofeitosa + 75163315, + // @santinroveda + 75849275, + // @Gisele_Nasc + 75977555, + // @fabriciopref + 76039110, + // @drwilsonbatista + 76043437, + // @senadorhumberto + 76049312, + // @carloschiodini + 76093489, + // @paulofiorilo + 76206825, + // @leitaothales + 76224437, + // @sofiacavedonPT + 76329038, + // @DepMajorAraujo + 76383384, + // @guerinocolatina + 76698584, + // @DepAfonsoHamm + 76741399, + // @Patrus_Ananias + 77210725, + // @MarcelAlexandre + 77266759, + // @malafaia_d + 77297183, + // @denisespessoa + 77658135, + // @ClesioSalvaro + 77860393, + // @rodrivaladares + 78154723, + // @AlexManente23 + 78673777, + // @leo_picciani + 78690962, + // @raquelferreirar + 78707944, + // @joslene65 + 78714361, + // @profdorinha + 79174387, + // @colattodeputado + 79252957, + // @gilbertoabramo + 80095058, + // @michelschlemper + 80123403, + // @capitaotadeu + 80214491, + // @LuisTibeOficial + 80307396, + // @MarcioMacedoPT + 80548664, + // @aldairrizzi + 80559557, + // @heldersalomao + 80575628, + // @aniellefranco + 80582734, + // @carlosmatosce + 80597262, + // @EvairdeMelo + 80626542, + // @depdarcidematos + 80712669, + // @Adrianageronim + 80733107, + // @HermanoMorais + 80995814, + // @fabiofelixdf + 81384580, + // @marciusmachado + 81445552, + // @liviaduartepsol + 81742151, + // @DaviSacer + 81836518, + // @laurasito + 82143155, + // @IzalciLucas + 82144764, + // @carlosjordy + 82271629, + // @Joyceeduca + 82412969, + // @CassianoCaron + 82415360, + // @Manato_es + 82433790, + // @leonelquerino + 82760526, + // @DepMauricioRS + 82776436, + // @darypagung + 82936046, + // @rodrigocruzz + 83489126, + // @MUNIQUEBUSSON + 83582018, + // @uczai + 83722886, + // @FRANBONI + 83730940, + // @alberto_fraga + 83844236, + // @_LuizEmanuel + 84127211, + // @joaoromaneto + 84141432, + // @andrerochamg + 84196661, + // @ProfPicler + 84325242, + // @MaguinhaMalta + 84396300, + // @oliviasantana65 + 84608790, + // @filhomarcio + 84639445, + // @lubloureiro + 85150664, + // @RobertoPSOL + 85293890, + // @randolfeap + 85327394, + // @emanuelcacho + 85461555, + // @danielter + 85585608, + // @maxlemos + 85613796, + // @PCBpartidao + 85647830, + // @wevertonrocha + 85859074, + // @realrcoutinho + 86065271, + // @vozdaenf + 86348991, + // @Baleia_Rossi + 86373674, + // @JeanVolpato + 86825438, + // @DomingosSavioMG + 87473436, + // @romeropelapb + 88303403, + // @Doutorluizinhot + 89419656, + // @DepJeferson + 89787624, + // @luizinhopatria + 90083628, + // @carladuwe + 90213796, + // @MariliaArraes + 90309661, + // @MGRodrigoCastro + 90478644, + // @arthurgurgeladv + 90630534, + // @FernandoCFAR + 90687549, + // @ANDREAESPANHA + 90898890, + // @wilsonlimaAM + 91109801, + // @RaffieDellon + 91840828, + // @f_trad + 92509126, + // @felipeaugusto01 + 92528438, + // @LucasRedecker + 92739280, + // @dpmarciomarinho + 92743190, + // @nancythame + 93065930, + // @marcelodeputado + 93116173, + // @bupessoa + 93424680, + // @danielsoranz + 93426187, + // @DuarteBechir + 93894534, + // @tomazteixeira + 93958167, + // @DrRafaelFavatto + 94115118, + // @Obsevador + 94167900, + // @coroneldavidms + 94378206, + // @ruycarneiropb + 94428241, + // @netemoura + 94450303, + // @samanthacavalca + 94806323, + // @EdicarlosVieira + 94943878, + // @requiaooficial + 95253000, + // @cabogilberto + 95526088, + // @overissimo + 95621286, + // @deciolimasc + 95699837, + // @anapaulalimapt + 95939603, + // @PROFTULIO + 96570084, + // @cassioandradepa + 96991211, + // @TrzeciakDaniel + 97587760, + // @evagoncalvesmg + 97677980, + // @maxmacieldf + 98786988, + // @matheusquintal + 98892109, + // @alexandrenau + 99001454, + // @ailtonlopespsol + 99389654, + // @michelyfarina + 99617723, + // @CacaLeao + 100521721, + // @julioarcoverde + 101016613, + // @fernando_it + 101892949, + // @PabloValenteDF + 102257530, + // @EduardoGomesTO + 102444270, + // @falcon + 103423503, + // @taliriapetrone + 103704608, + // @eubrasileiro_ + 103723639, + // @RafaelGreca + 104819806, + // @MarinaSilva + 105155795, + // @gustavopetta + 105236036, + // @nubiapassos + 105799989, + // @AmaliaTortato + 105872129, + // @NicolasTrancho + 106110720, + // @ThabattaPimenta + 106126597, + // @brendamars + 107012654, + // @DivaneidePT + 107976868, + // @tenenteromulo + 108322618, + // @mayradiasam + 108527180, + // @gui_pugliese + 108937956, + // @HugoMottaPB + 108988113, + // @rosenvergreis + 109006854, + // @DiogoForjaz + 109147422, + // @deplucasdelima + 109318072, + // @adjutoafonso + 109332428, + // @joaopaulodopt + 109657263, + // @paulaschild + 110187886, + // @depchicoalencar + 110522807, + // @thomeprefeito + 110560697, + // @Cleiton1Pereira + 110876570, + // @jeanwyllys_real + 111123176, + // @AndreMouraSE + 111148190, + // @talitagalhardo + 111695906, + // @HendersonPinto + 111717686, + // @marialuciaamary + 112476794, + // @dep_padrejoao + 113885149, + // @alielmachado + 115519533, + // @Silvio_CFilho + 115657305, + // @deputadomarcon + 115676753, + // @luizaogoulart + 116541810, + // @mahomsi + 116736901, + // @adelmosoaress + 116751115, + // @EdegarPretto + 116829646, + // @acrisiosena + 117425043, + // @fredprocopioofc + 117490508, + // @vereadorjulio + 117594353, + // @ronaldornrn + 117801614, + // @patybenzaquem + 118492432, + // @ZeRicardoAM + 119079224, + // @julio_cesar_pi + 119115954, + // @capitaosamuelof + 119416561, + // @adrianogaldino + 119586707, + // @AndreQuintaoPT + 119818761, + // @sibellebarros + 120535860, + // @aryvanazzi + 121422504, + // @walteralvesrn + 121425571, + // @netoevangelista + 121594926, + // @NelsonHossri + 121954870, + // @depulyssesgomes + 122123728, + // @felipecarreras + 122184686, + // @delegadowaldir + 123689660, + // @Valdeci13RS + 124157459, + // @FlavioCoutinho_ + 124360967, + // @Thais_Margarido + 124531510, + // @drgeorgemorais + 124753094, + // @edmundosouza7 + 124879050, + // @OlenoMatos + 125433649, + // @dario4e20 + 125727969, + // @annemouraam + 125816456, + // @MarceloCastroPI + 125822264, + // @susanevidal + 125851531, + // @anaffonso13 + 126091135, + // @jardelinacioam + 126323536, + // @deputadopezenti + 127708617, + // @lucasneves_sc + 127952305, + // @ricardozguidi + 127975238, + // @ortizjuniorm + 128550627, + // @GodriJunior + 128906180, + // @apjunqueira + 129055364, + // @ninamarinabraga + 129070751, + // @Mersinho_Lucena + 129681031, + // @depleomonteiro + 129837652, + // @alexandrebaldy + 130620293, + // @paulaodopt + 131004593, + // @minc_rj + 131286287, + // @Miriampetrone + 132187525, + // @depdelmasso + 132190299, + // @PedrodoOvo + 132220469, + // @Larissafgaspar + 132480664, + // @betofantinel + 132735215, + // @40rodrigofarias + 133685785, + // @glaustindafokus + 133692726, + // @drthiagopeixoto + 134284069, + // @sidneyleite_ + 135349877, + // @LucianoDucci + 136004062, + // @matheuscadorin + 136381913, + // @paulomansur_ + 136710714, + // @andreia_zito + 137369968, + // @BiradoPindare + 137548563, + // @tadeuveneri + 137919701, + // @sigageraldoluis + 138172915, + // @Eduardo_Cury + 138504441, + // @Katiadiasjf + 139059131, + // @rodrigominotto + 139809677, + // @MerlongSolano + 140413929, + // @raphaelsebba + 140485170, + // @wiliantonezi + 141023529, + // @DiogoPBotelho + 141087783, + // @lucianogenesio + 142068227, + // @gilmarribeirojr + 142309506, + // @silvioantonioma + 142501634, + // @moemagramacho + 142707910, + // @NelsiWelter + 142895907, + // @Carloshbfavaro + 143529694, + // @joaorodriguessc + 143924396, + // @SenadorRogerio + 144372753, + // @oacelio + 144599667, + // @fredpachecorj + 144681516, + // @majorpalumbo + 147744504, + // @dep_fatimanunes + 148355994, + // @depzanchin + 148424655, + // @vicentinhojr + 149013361, + // @sanchilispe + 149518311, + // @marcospereira04 + 149746462, + // @YvisEvelynn + 150019112, + // @narciakelly + 150448458, + // @DrLindoso + 151034428, + // @helderbarbalho + 151653693, + // @pedrogomatos + 152900822, + // @CezinhaNunes + 153131779, + // @AntonioTelesJr + 153268918, + // @RicardoCappelli + 153563550, + // @JarbasFilho_ + 155230905, + // @FabiolaMansur_ + 155564771, + // @dacassia1 + 155768056, + // @sgtalexandre + 156144539, + // @profcanguru + 156326262, + // @Sergio_Turra + 156355487, + // @AzambujaReinald + 157184730, + // @JoseMedeirosMT + 157226645, + // @RenataBuenoITA + 158067709, + // @Marivaldo4P + 159643822, + // @CoronelMedina14 + 160541377, + // @sandroalexpr + 160554168, + // @MarcosRogerio + 160895960, + // @MarcoVamosaLuta + 161318399, + // @marinorpsol + 161404367, + // @LianaCirne + 161416659, + // @maneco_hassen + 162794224, + // @Vanderlan_VC + 162807255, + // @TJMFernandes + 163251815, + // @emidinhomadeira + 163546697, + // @MariliaPFerrari + 163606193, + // @Ruicpimenta29 + 163831440, + // @DepZeMilton + 163935204, + // @wandnogueira + 164012285, + // @dimasgadelha13 + 164058788, + // @DepNeilando + 164416842, + // @EduardoBraga_AM + 164439493, + // @profjosemarpsol + 164478764, + // @TiagoCado + 164816642, + // @crisrbritto + 165329128, + // @_Heloisa_Helena + 165499618, + // @Jeronimoba13 + 165754961, + // @DeputadoRoberio + 165914831, + // @andrepdt12 + 165922194, + // @geovaniadesa + 165936805, + // @caiopontess + 165949065, + // @NelterQueiroz + 166164275, + // @deputadotoninho + 166417133, + // @rafacastro_40 + 166612108, + // @artagaojunior + 166645637, + // @depsamueljunior + 167761418, + // @serafinipsol + 168155564, + // @WladGarotinho + 168360473, + // @marcocastilhoss + 168408841, + // @drthiagoduarte + 168504384, + // @liberato1125 + 168510620, + // @RachelSherazade + 168520768, + // @gustavohaguera + 168653875, + // @TiberioLimeira + 168657354, + // @edbernasil + 168814204, + // @pepecollaco + 169158984, + // @DepEzequielRN + 169180911, + // @matheusmanholer + 170176086, + // @pedrofrancez + 170188044, + // @depJeferson10 + 170505916, + // @walterlfcaval + 170638771, + // @wilsonsousajr + 171158811, + // @giovaniculau + 171278562, + // @dudualfaia + 172117779, + // @JanineLucena + 172579844, + // @tayannystefany + 172581180, + // @rdrgcassollima + 172883136, + // @arthurboavista + 173527362, + // @toinharocha2 + 173823248, + // @rodrigomoraes65 + 174272995, + // @raquellyra + 174370381, + // @deputadoalisson + 174507960, + // @CinthiaCRibeiro + 175438359, + // @ThiagoBagatin + 175898388, + // @queirozmfilho + 175901537, + // @julianecviana + 176245917, + // @RenanFilho_ + 176502945, + // @MiltonVieiraOfc + 176920334, + // @bordalopt + 177595081, + // @alxlindenmeyer + 179749078, + // @leatricebez + 179938194, + // @JulyverModesto + 180256758, + // @VictorCoelhoES + 182050588, + // @eupugina + 184098414, + // @BrunoEnglerDM + 184210354, + // @rogeriobarrapa + 185053556, + // @FaveroNeto + 185482131, + // @rodfvale + 186331377, + // @Andiara1400 + 187726597, + // @acarlosmendes + 190308328, + // @pablomarcal + 191223319, + // @lorran_rebeldia + 192010152, + // @esperidiaoamin_ + 192658045, + // @pretagilsa + 193045733, + // @joeldaharpa + 195423732, + // @MarciaTaschetti + 198338329, + // @JacksonAndre7 + 198350574, + // @DepEduardoCunha + 198535390, + // @EdsonSilvaCotia + 199102127, + // @MariaSeffair + 199526953, + // @Anderson_Prego + 201100473, + // @silvyealves + 202856731, + // @natbonavides + 203332874, + // @reisptsp + 203784995, + // @profagraciele + 203948587, + // @profhenriquejr + 204042807, + // @Rafaeik + 204136918, + // @toninhobondade + 204759112, + // @lucasrfamaral + 205959233, + // @luboiteux + 207739610, + // @AntonioCoelhoPe + 208636149, + // @ivaniserotta + 209071918, + // @MeuEuPolitico + 209674215, + // @saullovianna + 210299199, + // @fatimacnagitos + 212950503, + // @Odarlone + 213105137, + // @betaopt + 213382232, + // @rafaelprimo + 215727649, + // @Lucinildo + 217347443, + // @celmarcosantos + 219262636, + // @LaizPerrut + 220062558, + // @arilsonchiorato + 220353446, + // @boscosaraiva + 221399756, + // @capitaowelton + 222230527, + // @DATENAREAL + 223852535, + // @LuisMirandaUSA + 225426734, + // @filipebarrost + 225925013, + // @gardelrolim + 226984087, + // @emersonbacil + 229028677, + // @gislenemoura + 229229129, + // @osmarfilhoma + 230495542, + // @marisol_santoss + 230769498, + // @luisa_canziani + 230950274, + // @caroldecaxias + 231379712, + // @kelitalks + 235765762, + // @euberlucas + 237071521, + // @lyneoliver + 237497430, + // @elzefacchinetti + 237562432, + // @wildermorais + 237687306, + // @deltapericles + 239036359, + // @_lucavalcante + 239059356, + // @silascamara_ + 243018634, + // @UrsulaVidalPA + 243429150, + // @gutoschiavetto + 243429651, + // @moarasaboia + 244491558, + // @depgurgel + 245320150, + // @RogerioCorreia_ + 245392082, + // @gelizabetesp + 247745869, + // @DeAssisDiniz + 247906787, + // @iginomarcos13 + 250070980, + // @prgilsondesouza + 250214822, + // @ranypaulino + 252553750, + // @RachelMaroja + 253219200, + // @rdnarede + 254172269, + // @deprsantos + 254201392, + // @SargentoFAHUR + 255300173, + // @wellington_luiz + 255637975, + // @fredantunes_rs + 256499222, + // @DelegadoJacovos + 257247983, + // @TamyresFilgueir + 258957259, + // @Rodrigopreis13 + 259022842, + // @Marceloalvaroan + 261007129, + // @Nenebabyvianna + 261446056, + // @macielchris + 261472149, + // @paulobregolin + 262375578, + // @coronelfrota + 263153198, + // @victorantoun + 264298649, + // @catulejr + 264391266, + // @deboratruck7 + 266856351, + // @AureoRibeiroRJ + 267467458, + // @alinegurgel_ap + 267981392, + // @depsebastiao + 268318846, + // @pastorflamarion + 269618056, + // @ayres_jr + 269938072, + // @RafaelMottaRN + 271212223, + // @alexgalvaodf + 272477039, + // @diegogarciapr + 273616279, + // @JairMiotto + 273974254, + // @doutorgutemberg + 274515672, + // @BiaCerqueira_ + 274688443, + // @netocoelhoo + 276794985, + // @DafneOrion + 277586067, + // @danealencar + 278126758, + // @CostaMarinara + 278549268, + // @AllanPombopdt + 279092433, + // @InspetorRobison + 280306544, + // @mi_andrews + 284903425, + // @victorsalatiel + 285614672, + // @maicolmed + 286427700, + // @gabrielaorttiz + 287219048, + // @marcelomaranata + 287876093, + // @anaacarlinha + 288430204, + // @humberto130 + 288512944, + // @delegado_waldir + 288735634, + // @laualencar_ + 288761063, + // @ZeRobertoLula + 289318056, + // @waltercamargo40 + 289521136, + // @coelho_rodrigo + 290204659, + // @fefrancischini + 290475286, + // @maitebrusman + 294065810, + // @alexdapiata + 295697900, + // @gamanaiton + 295949375, + // @RinaldoJunior40 + 295964934, + // @cleniltonsc + 297025171, + // @DepSostenes + 298308683, + // @CarlaPrataReal + 299254693, + // @titotorresmg + 299923575, + // @depkleberRN + 302056921, + // @LopesCancadoAdv + 302235657, + // @ricardopinaffi + 302725028, + // @LelinhoLopes + 303942993, + // @orlandosilva + 304092926, + // @marcosjorgebv + 306253241, + // @fabiarichter + 307557586, + // @gleicejanems + 308780059, + // @guto_schiavetto + 309004600, + // @SamiraDaud + 310033093, + // @GildeteAlves + 310502840, + // @TercioTinoco + 311616488, + // @paulobilynskyj1 + 311720533, + // @katiabacelar + 312252804, + // @depandresoares + 313466411, + // @WedersonLopes + 313684933, + // @franciane_bayer + 317784577, + // @_sergiosouza + 317973567, + // @danielvalencapt + 318019551, + // @depmaracaseiro + 321414691, + // @NonatoSampaio + 321556855, + // @loureirocris + 322406139, + // @jorgeviana + 325131009, + // @AnibelliNeto + 327327238, + // @MartaVitorino + 327423765, + // @OtoniDepFederal + 330770436, + // @patriciamelobr + 331174978, + // @MajorVitorHugo + 332324517, + // @RequiaoFilho + 333720455, + // @Fabinho_Gaspar + 334978230, + // @simboramudar22 + 337269106, + // @_soedi_ + 339854006, + // @marisa_lobo + 340331807, + // @danniellibrelon + 340748015, + // @MariliaFreireAM + 343512877, + // @MarceloBelinati + 345512946, + // @SertanejoUFC + 346580194, + // @edusantosdf + 348089005, + // @leilafonsecapb + 349053302, + // @paulinhoramosap + 349163261, + // @DenianCouto + 350980280, + // @JanderBrum + 351008004, + // @paulapreta50 + 351272601, + // @RenatoLoffi + 351817280, + // @Haddad_Fernando + 354095556, + // @Biakicis + 357030742, + // @wanderleyporto + 358996581, + // @andrefernm + 360231929, + // @PROFAVILETE + 360649695, + // @milkleileite + 365210575, + // @Thiagofreitaz + 365552309, + // @andresalineiro + 367519089, + // @Ana_claudiapb + 370942708, + // @RONIESILVA15 + 373465468, + // @coroneljunior + 381045128, + // @ieda_chaves + 389660780, + // @RichardSelvagem + 394195842, + // @depbosco + 396265009, + // @JorgeFrederico2 + 398245257, + // @fabigoulart30 + 399159303, + // @FeFrancoOficial + 399589593, + // @hugosilva63 + 402050446, + // @WMello69 + 402958613, + // @brena_dianna + 412508127, + // @faf_freitas + 412669574, + // @_robertoluiz + 415073795, + // @zeliomota + 420614002, + // @AMonicaFacio + 422814320, + // @Carlos_Gaguim + 423706952, + // @Edianefolle + 423874716, + // @vitordeangelo + 424455116, + // @BalbinottiFilho + 428247512, + // @RealNabor + 429039681, + // @celaocordeiro + 429940274, + // @joicealvarengab + 433833555, + // @instrutormarcio + 437397340, + // @RenataAbreu2020 + 441108081, + // @wilsinhodatabu + 442693955, + // @AllysonBezerra_ + 444925483, + // @LorruanaM + 456789752, + // @CarlosMoises + 456795842, + // @drpaulocruz + 459276565, + // @soudaniellapb + 459320302, + // @JulioCesarRib + 459681629, + // @marcosbrazrio + 460343731, + // @advemanuelbueno + 465613391, + // @bellagoncalvs + 469918968, + // @crisbrasilreal + 472033751, + // @andreawerner_ + 475996406, + // @leopratesba + 480604517, + // @IanBlois + 480755252, + // @xerifedoconsum + 483269816, + // @chicorafa_s + 487108193, + // @pedroluislongo + 487622592, + // @LPescinelli + 492334281, + // @brunosouzasc + 494268633, + // @brauliolaranovo + 505339844, + // @caiocordeirotv + 521514005, + // @glauberbastos_ + 527053144, + // @MoisesSantosAc + 538269241, + // @JulianaBRedivo + 545696318, + // @mariarosassp + 554803332, + // @alcyvania + 556852639, + // @limmajr168 + 565889975, + // @ThammyReal + 578495086, + // @mariohildebrand + 580176611, + // @OthelinoNeto + 583377940, + // @HelioWirbiski + 589375704, + // @adautosoutoms + 596446640, + // @AfonsoFlorence + 599427558, + // @CiceroSimplicio + 604095777, + // @lucaspavanato + 608452346, + // @Jasson_Goulart + 608730390, + // @MarioEsteves2 + 610309440, + // @pastorhenriquev + 617468983, + // @MateusWesp + 618506366, + // @elmanooficial + 626602522, + // @fabiogov55 + 630758039, + // @rosemodestoms + 632251844, + // @helenaduailibe_ + 632737286, + // @andrewleal2 + 636590052, + // @WesleyCasaForte + 637247323, + // @charlessantosmg + 709095821, + // @Cidinho_Santos + 745333897, + // @mitchellemeira + 746090918, + // @FKrelling + 753495397, + // @nikolas_dm + 758264276, + // @suelenmarques06 + 796882578, + // @bneydavid + 797062418, + // @MottaTarcisio + 799260530, + // @JenirNeves + 813802178, + // @PedroDeyrot + 857054846, + // @carlaopelobem + 893975196, + // @Isoldadantaspt + 999393290, + // @dr_nesio + 1008361322, + // @MarcelloPaula + 1008742886, + // @SHEILAKLENER + 1038802238, + // @rodrigostm10 + 1038849745, + // @gabrieldiedrich + 1059637784, + // @xambinhoes + 1068257582, + // @depjpassarinho + 1071798366, + // @MacaeEvaristo + 1075036110, + // @tatianehelena81 + 1075326786, + // @moisesselerges + 1081998169, + // @deproosevelt + 1084884007, + // @depjanetedesa + 1094959356, + // @NeumannJarbas + 1124876318, + // @McSmithOriginal + 1127018335, + // @cacateixeira45 + 1130957935, + // @daltonlueders + 1162799238, + // @prof_juniorgeo + 1226451780, + // @marcoswesleymw + 1308817182, + // @alexceoficial + 1314840228, + // @ThaysBieberbach + 1316758495, + // @zecarlospt + 1325494376, + // @boettchermarcos + 1339028060, + // @CARLOSVALADARE7 + 1356677952, + // @D_GoretePereira + 1362596354, + // @RobertoRocha_MA + 1436541721, + // @mickasevalho + 1461892208, + // @reginetebispo + 1485556436, + // @KimKataguiri + 1494658207, + // @leticiamattossc + 1498645730, + // @deyvidbacelar + 1526183576, + // @bublitz_a + 1532130170, + // @Luiza_RibeiroG + 1544201047, + // @pedrocampospe + 1546304018, + // @natthpaccola + 1570635661, + // @NegrahLima + 1571332381, + // @oficialdmarques + 1604525460, + // @carlosfportinho + 1612753909, + // @marciopachecopf + 1613063005, + // @renancalheiros + 1650330319, + // @GilvanMaximoOfc + 1651852124, + // @carlosedrsantos + 1685646356, + // @alcidesfernm + 1687046911, + // @Brandaveneno + 1710385291, + // @chaficlays + 1725112411, + // @bombeirorafa + 1733073672, + // @dr_furlan + 1844887189, + // @josaqueirozpt + 1848767107, + // @margabuzetti + 1854699607, + // @ReginaldoVeras + 1864033711, + // @lucianomattosmp + 1892291491, + // @valmirdesergipe + 1960681207, + // @Alexandresantan + 1964899736, + // @helinhocastro + 1977089148, + // @DepAntoniaLucia + 2161527577, + // @JackRochaes + 2174078260, + // @juniortunao + 2178147073, + // @talitavazbh + 2208038935, + // @lidiamourac + 2210601883, + // @LUANARUIZSILVA + 2216639570, + // @DepFederalMoses + 2217650233, + // @andredopradosp + 2250697531, + // @ZeniteRosa + 2289590857, + // @BahMatteuss + 2293776768, + // @deltanmd + 2296138146, + // @BrunoCarianha + 2299610603, + // @MatSchilling + 2310607019, + // @delmartharocha + 2312487756, + // @ProfClaudioBran + 2352495202, + // @DavidAlmeidaAM + 2353403137, + // @capitaoassis10 + 2359974485, + // @marciolabre + 2401917402, + // @sheikhrodrigo + 2429206439, + // @OmarAzizAm_ + 2445921702, + // @Francadf_ + 2445938091, + // @Alfredoficial22 + 2474258532, + // @GenPeternelli + 2491980048, + // @yuriarrudam + 2495584641, + // @simonetebetbr + 2508415207, + // @prof_rosaneide + 2523520542, + // @rosecipriano_ + 2523530742, + // @tiagossimon + 2525023980, + // @DelegadoOlim + 2540255982, + // @gleidept13 + 2544622051, + // @leila_0ficial + 2560238086, + // @viagensdaiw + 2572998767, + // @Mata4Adrianada + 2580412784, + // @anadogasoficial + 2583179096, + // @depheliolopes + 2585621855, + // @profsoniameire + 2604536294, + // @depjorgesolla + 2605560932, + // @mvictoriabb + 2609799432, + // @barbosinhams + 2612421128, + // @_newtoncardoso + 2613918878, + // @EderMauroPA + 2632802395, + // @paulo_litro + 2647646724, + // @LuizianneLinsCE + 2666819714, + // @LulaOficial + 2670726740, + // @NiltoTatto + 2674564802, + // @ReginaFortunati + 2674769546, + // @depjoaodanielpt + 2690172127, + // @marcoaurelioITZ + 2691147288, + // @cafabrini + 2691507383, + // @UalidRabah + 2712571453, + // @LucasVergilioGO + 2717062461, + // @netto_expedito + 2717313965, + // @tiaomedeiros + 2732504341, + // @luthrebeloPA + 2741279361, + // @crismachadotity + 2762117711, + // @DepSanderson + 2767167039, + // @RenanPaesSP + 2784935916, + // @carteiroreaca + 2797010717, + // @MarleideCunhaRN + 2820378598, + // @mqueiroga22 + 2823869072, + // @HigaWagner + 2830971208, + // @AlbertoMaiaf + 2834745257, + // @RubinhoNunes + 2838953716, + // @daianasantospoa + 2858823694, + // @yagorVieira + 2868228117, + // @_akalicia_ + 2879108776, + // @paulolemosap + 2881293743, + // @manascimentogo + 2892653500, + // @LizianyM + 2894163148, + // @RUBENSCANTUARIO + 2903026858, + // @andreicastroba + 2925491427, + // @KuhlmannJean + 2927136550, + // @NoletoMBL + 2938978041, + // @DepJuscelino + 2970617333, + // @neioluciofp + 2977449981, + // @cabobonadiman + 2977624732, + // @guipaoficial + 2979670457, + // @depjosenelto + 2995610423, + // @isaakalmeida93 + 2997861520, + // @carlosveraspt + 3010698441, + // @miguelcosta46 + 3015267897, + // @jbittencourtjr + 3018188571, + // @mendanhagustavo + 3018780587, + // @VillaMarcovilla + 3021556996, + // @lvaroDomingues1 + 3025314889, + // @karlosbernardoA + 3029604164, + // @tatianeruas + 3030593279, + // @depniltonfranco + 3044931389, + // @AdrillesRJorge + 3060235071, + // @jackson56786523 + 3084537202, + // @DepVitorLippi + 3092560931, + // @rosangelawm + 3096479489, + // @ericgustavodf + 3111776739, + // @TeonilioBarba + 3119378914, + // @otavio_camp + 3125324049, + // @jarbassoaresjr + 3125532669, + // @zuleidequeirozf + 3130842411, + // @doutor_vicente + 3130887358, + // @DrLeonardomt + 3131609429, + // @senadorasoraya + 3167874665, + // @meire_cruvinel + 3205786257, + // @andre_fufuca + 3254270357, + // @Marcio_Honaiser + 3294107902, + // @deputadopriante + 3305760711, + // @marcos_reategui + 3306282352, + // @elikatakimoto + 3316550554, + // @ClaudiomirCast1 + 3335733075, + // @drjeanfreire + 3342622547, + // @castrothulio + 3354383374, + // @PbnConcursos + 3357932231, + // @brasleiroluc + 3366080079, + // @moisesbrazpt + 3373574517, + // @caformigoni + 3385089700, + // @kleybe_morais + 3512854216, + // @carmelonetobr + 3662771592, + // @OPaiakan + 3674682197, + // @delegadanadine + 3744465381, + // @ninasouzarn + 3904595243, + // @FelipeMichelRJ + 4026638189, + // @geivissonvieira + 4028627062, + // @andersonrib_rj + 4031541093, + // @paulo_mavignier + 4043244676, + // @depsorayasantos + 4052741685, + // @zeaugustonalin + 4056653428, + // @GeraldoStocco + 4057530885, + // @Vinylobianco + 4077498658, + // @depzegeraldopt + 4204580127, + // @_AlessandroSE + 4250596815, + // @IrataAbreu + 4309564575, + // @DepLucianoMDB + 4350613047, + // @josuelsantosbch + 4493572241, + // @CarolDeToni + 4566967516, + // @brunocunhablu + 4648157621, + // @FofaBorges + 4775264669, + // @andrebrandao_rj + 4820124735, + // @SandroFilhoBA + 4836857632, + // @ederborgesbr + 4871172143, + // @emersonosasco + 4892851287, + // @AdvogadaDoPovo + 4895047809, + // @jumildemberg + 4899470602, + // @joaoazevedolins + 698129524106653697, + // @AndreaOficial55 + 703604819538407424, + // @PaulinhaQuint + 703685122265116673, + // @juvircostella + 707214837584105472, + // @JuninhoSinono + 711521612219158528, + // @MARTACLERIALIMA + 713343399898820608, + // @drlaudicerio + 713734821747560448, + // @OctavioSampaio_ + 714887920931442690, + // @MichellePSOL + 714891302723264512, + // @MiguelSRossetto + 714958583239151616, + // @BentoLeiteML + 717529076974620672, + // @SauanRockenbach + 719724675908124672, + // @alvaroporto_pe + 719998490831634432, + // @LiaGomesCE + 724003686863740933, + // @ViniciusFerroC + 726243481669242880, + // @emersonjarude + 726777982115799040, + // @lpbragancabr + 728281672731471873, + // @brisabracchi13 + 728292397130592256, + // @zoemartinez_05 + 730865469054435328, + // @djalmaneryneto + 732299122242359296, + // @ErikakHilton + 738143559920934912, + // @dreltonjr + 739270494629842946, + // @cabo_senna + 744609688415789056, + // @rjdouglasgomes + 745611833248186368, + // @jarir_pereira + 746804180430487554, + // @doorgalandrada + 748233031337545728, + // @paulobrant_ + 750386404417495040, + // @TamirFelipe + 750744436271964165, + // @coronelrochase + 752300463693914113, + // @femirandapsol + 753276824743018497, + // @samuelsalazarPE + 753589568314703872, + // @jeffersonlimapa + 757671198125858816, + // @JanainaDoBrasil + 759001618884939776, + // @gauchodageral + 760848309892161536, + // @ronaldodimasto + 761550012652322816, + // @mazer_pg + 762337048858595328, + // @toinhocarolino + 762649334504689664, + // @lindabrasilse + 764288327713452032, + // @angelaamin11 + 766624608338477056, + // @LeonardoBalbi14 + 767499420384497664, + // @kauan_poubel + 770098743412752385, + // @vanessapstubh + 770298023582765057, + // @fernandogorgen + 771705262671556608, + // @pereira_lenilda + 779061594970025984, + // @Euallephhillz + 780616664601604096, + // @MariaConstantin + 782293493364457473, + // @taguayork + 785545895748173826, + // @deputadodocarmo + 793261830072311812, + // @RafaelMarcal33 + 793424665431666688, + // @DaniloBalas + 797527121879068672, + // @iyagomedeiros + 799344884197064704, + // @SocorroLac + 803633234705645568, + // @CatharinaDon + 806302110899994636, + // @DragaAlana + 806523898917515266, + // @Jorgepinheiroof + 807903184576532480, + // @ruanmartinsjp + 811160203483893765, + // @fabianodaluzsc + 813543058457423872, + // @mariana_psol + 818583425149902849, + // @amorimvivian_ + 820867290749161472, + // @goura_nataraj + 821001145359409152, + // @VictorRuizSP + 823586870864932866, + // @ProfessorEuler + 824285052590845955, + // @luizfernandopt + 826789430425960450, + // @angelocoronel_ + 829391158208032768, + // @AdelitaMonteiro + 830144764360192002, + // @ranallipf + 832322309436403712, + // @OseiasVarao + 832603702737379329, + // @depeniotatto + 834065907743854592, + // @peumendonca23 + 839127759955841024, + // @DaSilvaBelford + 841440085035896834, + // @GeneralGirao + 841700087143288832, + // @viniciusaithsp + 845046382910197760, + // @_LeandroBello + 850113810681802756, + // @Guischleder + 850728101868863490, + // @WillaceSouza + 850775334362505216, + // @vivitobiasms + 852960064159842305, + // @Marcio_Canella + 860275250830999554, + // @rodrigodiasbsb + 863812402680397826, + // @elke_pimentel + 864274647583469568, + // @MiqueiasS0ares + 868677911674597376, + // @DanielaCarneiro + 869919798767091712, + // @BuchiOgier + 874408561572552704, + // @renato_battista + 877159534430650369, + // @deputadohalley + 879906917371564033, + // @pedroalmeidace + 882539177757298689, + // @AndreCeciliano + 885883830489554944, + // @eucricielle + 887083474104049664, + // @paulosafederal + 889464426222563328, + // @brunopedralva + 889497158491271169, + // @jonesmanoel_PE + 889628930436730881, + // @MayconRobertoPR + 893177406705553410, + // @izadutrabr + 893178307591770112, + // @CatiaColombo2 + 894671613576269826, + // @NiccSan + 898362443717627904, + // @meuamigojoao + 899602419302240257, + // @v_medioli + 902634025566732288, + // @Diego_Bentim + 903169689202954241, + // @todandara + 904842960931565568, + // @natalimamt + 905965883952164864, + // @AlcyPinheiroCE + 909049597091356673, + // @coroneljonildo + 913597378002866177, + // @deborapsol + 915203945479458818, + // @leonidio_boucas + 917373635605786624, + // @cabojunioamaral + 917727546615193601, + // @NoemiaTiago + 918524953561100288, + // @wanderley_vieir + 921969720546480128, + // @GuajajaraSonia + 924978508224516096, + // @valentinarrocha + 927752925031620608, + // @MajorMecca + 931022287066804224, + // @TorminCassiana + 931087564064329728, + // @ptalfredinho + 931140401293070337, + // @julianopsol + 937683426882326529, + // @rigoni_felipe + 939423395376136192, + // @sergiokruke + 940210949943910401, + // @fabiofreitaspa + 940623667322609664, + // @pedrocheoficial + 947255546062745600, + // @jusmarioficial + 948577496177561600, + // @DpRicardoArruda + 953055428124045313, + // @Carlos_cabral81 + 956932819086913536, + // @WilliamSiriRJ + 958357145761845248, + // @delegadopiquet + 958833209961254912, + // @DepJorgeEverton + 959184931552464896, + // @ClaudioLCaivano + 962035539133177858, + // @Fbgg40 + 963928182368980993, + // @RafaelLustoza_ + 965096425313972229, + // @XandePessoaPE + 966833467391725569, + // @Pastorellux + 968686315326799872, + // @rsallesmma + 971822131024748545, + // @RNBolsonaro22 + 973888139637993472, + // @dep_paulinha + 973982865997385728, + // @ORicardoDaKarol + 976888734267445248, + // @Diegoferdfc + 977195483507707904, + // @MoRosenbergSP + 977579494344142848, + // @AndreTrindadeRP + 978076973300879360, + // @ZeCocaOficial + 978602905690427392, + // @jairfarias_to + 979750671728685056, + // @julialucydf + 980846559444291585, + // @Indianarae1 + 981429809413873664, + // @Fausto_Pinato + 981500250602041344, + // @oficialigortimo + 981604712238780432, + // @ChristianProf_ + 981899518244544514, + // @Renatoafjr + 982217430297554945, + // @fernandomaximoX + 983360917227364353, + // @RenilceOficial + 984221723544444928, + // @eng_angelo44 + 984522534623301632, + // @israelsantosap + 985962825909723141, + // @Dep_GilPereira + 986284866533842944, + // @IndiaraNOVO + 986396254287646721, + // @leaolais_ + 986659343436255232, + // @ptmatiaspedro + 986903953014128641, + // @leiladovolei + 986978033159626754, + // @DougjButzke + 988809130256273408, + // @paulambelmonte + 989854043370606592, + // @CaboDaciolo + 989899804200325121, + // @GutoZacariasMBL + 991090809578708992, + // @RJRogerioAmorim + 992223431906287616, + // @KerexuOficial + 992801222926196736, + // @danimontpsol + 993146341294596096, + // @kekabagno + 993288307625943040, + // @andreiadejesuus + 994268335486504963, + // @Alessandrojesu_ + 996849375635820544, + // @prdinhosouza + 999747749364076545, + // @thiagoavilabr + 1000525959257378816, + // @Izalourenca + 1000724258740473856, + // @tabataamaralsp + 1001251931812220928, + // @NatanSperafico + 1001476076466593792, + // @zuccors + 1002182052341534720, + // @romulobuldrini + 1002646368496865282, + // @paparicobacchi + 1003614554394415104, + // @jmtavaresz + 1004078050521305088, + // @DudaSalabert + 1004511711251099653, + // @FreireJocimar + 1005181143627522049, + // @ferreirinharj + 1006202251990437888, + // @ZeDirceu_ + 1007663803948027904, + // @annakarinapsol + 1009608950176796677, + // @CapitaoContar + 1011667518728089602, + // @f_francischini + 1012066500373557253, + // @EricaGorga + 1012086799328522246, + // @gabino + 1012469182330355713, + // @Trom_Petista + 1013494732834639874, + // @danicunhario + 1013945423676010497, + // @limmapiaui + 1014568888078651392, + // @costa_major + 1015188576584298497, + // @acaroldartora + 1016697971843502080, + // @CarlosBurigo + 1019197414136320000, + // @danielibalbi + 1020645096226795520, + // @RomeuZema + 1020798087449776128, + // @EduGiraoOficial + 1024403315164160000, + // @Laurez_Moreira + 1025352673292378113, + // @juliadecastrobr + 1025519861592666113, + // @zenaidern + 1025870063461720064, + // @deppimentel + 1026636587298443265, + // @EuclydesPetter + 1027187280438546433, + // @Pataielloo + 1027646933110804480, + // @MarcosFelipiPoa + 1028075839269888001, + // @MaisaMitidieri + 1030455388678889472, + // @professorjoziel + 1030551918127529990, + // @KenjiNohama + 1030577196598001664, + // @BocaAbertaOf + 1031914738933018624, + // @leomartins45rs + 1032017342937661441, + // @vereadornetuno + 1032654017418211330, + // @SalvinoOliveir1 + 1034050960262418433, + // @victorhugoforte + 1034054725107363840, + // @RafagninLuciana + 1034079788162535424, + // @erikaamorimce + 1034138653147242496, + // @AuricchioThiago + 1034858280231809024, + // @victoriogallimt + 1036259095458795526, + // @AmauriRibeiroGO + 1036681658760658945, + // @NeymarPesadao + 1037281805118910466, + // @FaleiroAirton + 1037384667836637185, + // @phbarroso45 + 1039361882845528064, + // @AlanaPassosRJ + 1039881241301016576, + // @DouglasGarcia + 1040704983358955526, + // @RenanSantosMBL + 1042601099566436352, + // @Daysehansa1 + 1044658383922638849, + // @CavalarEmanuel + 1047079027696177152, + // @gutopfonseca + 1048016722433900550, + // @tome_abduch + 1049634819028717569, + // @raphaelbarrabr + 1049870508643233793, + // @ContaratoSenado + 1050121324436307970, + // @EuJorgeMiranda + 1051329992662142977, + // @Cristian0Engel + 1051904895106908161, + // @manumirella_ + 1052932498181840898, + // @depmgualberto + 1053123917785808901, + // @RafaelDemarchi5 + 1053334763858214912, + // @deputadatalita + 1055640991200411648, + // @cleitinhotmj + 1057231251743170562, + // @gilson__marques + 1057602315664924673, + // @delegadasheila + 1058010509256126464, + // @mauricio_lindol + 1058341993372360705, + // @evandro_stacruz + 1059551999258255361, + // @veronicalima_ve + 1059554967600685058, + // @pinheirinhomg + 1060134845043666945, + // @RuthVenceremos + 1061207070412886016, + // @giordanmes + 1062407588724252673, + // @pluviapt + 1062505159824150530, + // @capitaocarpe + 1062678892216020992, + // @FlaviaHellen_13 + 1062871543007662080, + // @rafaelsimoesmg + 1064322572261699585, + // @clarissatercio + 1065379080084865024, + // @professorabebel + 1065745997391949824, + // @GIBERTOPINHEIRO + 1066799808172695552, + // @evandroaraujodf + 1069259896892329984, + // @DelHelioBressan + 1070740683386957825, + // @wilkerleaods + 1073662759009771520, + // @marxbeltrao + 1074651587140902913, + // @capalbertoneto + 1076135802969772032, + // @tarcisiogdf + 1078618844007157761, + // @EduardoBraide + 1080814654195204096, + // @MarcosA_Sampaio + 1080911369447399430, + // @capitao_alden + 1081245039920144384, + // @deputadomoraes + 1081341581314150400, + // @DerriteSP + 1081517956964802561, + // @proanalucia + 1081700266410426369, + // @ingrapsol + 1081957744725381120, + // @matheuspggomes + 1083017623837777921, + // @Khalill_gui + 1083098735675084800, + // @PEDROCO13904182 + 1083745378489589760, + // @DamaresAlves + 1083774484123975680, + // @BennyBriolly + 1084272914202091520, + // @majorfabianadep + 1084712593292443648, + // @waltinhofoguete + 1084870539666247682, + // @wilkerbarretoam + 1085537170276958208, + // @ArthurLira_ + 1086390169970884613, + // @IsaacRicalde1 + 1086404913314385920, + // @emanuelzinhomt + 1087326506559389696, + // @DaviBenevides + 1090741557383311362, + // @issam_saado + 1090980245098979328, + // @DelegadoFurtado + 1091410188110831616, + // @PlinioValerio45 + 1091691201844207617, + // @joaorenatobr + 1092217577026318337, + // @doriel_barros + 1092446595889745921, + // @joaoluizam + 1092549114762539009, + // @karensantospoa + 1092781921560641543, + // @betopereirams + 1092816222012534784, + // @biologiagabriel + 1092879030465032192, + // @doutordanielpa + 1093086549825208321, + // @ZequinhaMarinho + 1093223195258298368, + // @enfermeiranaza + 1093488527583690753, + // @igortavaresmg + 1093510541404971008, + // @DepGuiLandim + 1095422600854032390, + // @MondardoGiovana + 1095645491419852800, + // @JuniorManoDep + 1095696286974652422, + // @IdilvanAlencar + 1095754303367720960, + // @depsargentolima + 1095990374550695938, + // @DiogoTalento + 1096490764677320704, + // @ReMesquita1977 + 1097042489813409792, + // @DFDanielFreitas + 1097498693719199744, + // @dilvandafaroPT + 1097822283387809793, + // @joaocardosobr + 1097888011650502656, + // @thalesdacostaal + 1098031305818812416, + // @drmarcelmuscat + 1098089617209913344, + // @CelAlfredo + 1098328459825287187, + // @helio_ver + 1098911814824411137, + // @DelegadoFreitas + 1098981858321276929, + // @benesleocadiorn + 1099003656064643073, + // @BlogdoSavio + 1102579349939724288, + // @FBarcellosSC + 1103370630827855872, + // @AcilonGoncalves + 1103616211311570944, + // @depPedroLucasF + 1104120809482866688, + // @jurailtonsantos + 1104747705434288128, + // @PollonMarcos + 1105081277135417345, + // @GabrielDosAnima + 1105170906987593728, + // @CoelhoAntonioPE + 1106161642486796288, + // @rodrigomoraisof + 1106200795098292229, + // @SargentoBetania + 1107097837194694663, + // @robsonjs_ofc + 1107238951990059008, + // @DelegadoZucco + 1107381354680012806, + // @drfranciscopi + 1108052690783928322, + // @DelegadoCaveira + 1109546175982698496, + // @leolupi_rio + 1109869057153683462, + // @DrCesarMello + 1110177096373100545, + // @ToniettoChris + 1110626741314306049, + // @keilapereirasp + 1111421096501432320, + // @JoaoCampos + 1112804630453501953, + // @SF_Moro + 1113094855281008641, + // @deputadothiago + 1113811245378015232, + // @keitlimasp + 1114595200113029120, + // @CaldeiraJos1 + 1114734016811474946, + // @DeputadoRatinho + 1115425047135571974, + // @faustojr_am + 1116362076459601920, + // @DeputadoLemos + 1116429709049573376, + // @10ronaldomartin + 1118164353121968129, + // @guitodeschini + 1118455651901083648, + // @deputadaniella + 1118481835376431104, + // @flavionpi + 1118489062866849793, + // @NajaraCosta_ + 1118899291983175680, + // @italomoreirasp + 1119731927991443462, + // @NelPiaui + 1120084261917351936, + // @henriquecesarr + 1121117485900730368, + // @guicolombosc + 1123714030391132161, + // @emanuelvenzo17 + 1124016358893785089, + // @AltairMoraessp + 1124045706883473408, + // @eunatashapoa + 1124140833547141121, + // @GilbertoCattan1 + 1125062220302385153, + // @Jeffers67876956 + 1125132737491480576, + // @RobertoJustusPR + 1125404580374822912, + // @FernandaLoubac2 + 1125416190199967744, + // @doutorfran + 1126519869053251584, + // @eduardo_borgo + 1126708527274315776, + // @diretoreliasjr + 1130259444326109184, + // @fabiosilvadep + 1130509226281975814, + // @GersonSPires + 1130572020901699585, + // @gilmarpetrolina + 1131947049690255360, + // @brunolessarj + 1131984902390460416, + // @rosamst_ + 1132357062703439873, + // @DeputadoPoubel + 1135964153309478912, + // @Fabianohorta_ + 1137809351387865088, + // @vinicios_betiol + 1137885440768450561, + // @vilmareisfem + 1137930985109110785, + // @PapoDoCappa + 1140263037049475072, + // @FernandaSixel + 1140840515124047873, + // @DirleteP + 1141306501019119616, + // @georgebastos30 + 1141394024860966912, + // @delegadotayah + 1141737559816585217, + // @CarboniRogerio + 1143607219302387712, + // @lesinhaduarte + 1144077937228034048, + // @felipebecari + 1145863684511674369, + // @luisedulelis + 1148639803908472832, + // @pedrorochafilho + 1151964529896644620, + // @ShirleyCruz22 + 1152873734157393922, + // @washingtonban + 1154733508088270849, + // @DacunhaDelegado + 1155962982486171648, + // @AlanLopesRio + 1157984524116156416, + // @torinomarques + 1158806983925030913, + // @BertolucciGab + 1159503401463484417, + // @karlasarney_ + 1160505413697253376, + // @erickdenil65 + 1161802414921527296, + // @JalserRenier + 1164616428944801793, + // @jofariasm_ + 1167172974451073024, + // @RobertoCidadeAm + 1168549394830020608, + // @umJovemPaulo + 1169327756544479232, + // @monicacunhario + 1171205026724929537, + // @jadearomero + 1171411853282557952, + // @prjuniortercio + 1171959927771975680, + // @alice_psol + 1174767854459195393, + // @uaivittor + 1174847040406376449, + // @carlosnovofsa + 1176225764200591361, + // @joaobettega_ + 1177185451536465920, + // @marcimeirelles + 1177246955942154241, + // @GalbaSheyla + 1179012777366736899, + // @delegadopalumbo + 1179437585275465729, + // @by_dzn + 1184622911174389762, + // @rickazzevedo + 1184969015023800320, + // @PauloSussumu + 1188038521501749248, + // @FranciscodoPT13 + 1188191474636206082, + // @samueljesusjb + 1188281892229074951, + // @atenabr051 + 1188827192999972865, + // @BrunoZambelli3 + 1191788412833075201, + // @rebecaromerobr + 1192983475663638529, + // @ProtetorAle + 1193938478456868865, + // @thiagomedinamd + 1196894264334127107, + // @leninhamoc13 + 1197135895469711362, + // @JulianaBenicio_ + 1197281720694956032, + // @PettersNeto + 1198325133368340480, + // @DeputadoCarlosH + 1199293398194233344, + // @mariliacamposmg + 1199714008217063424, + // @SandraLimadeVa1 + 1199740816018853890, + // @tukuma_pataxo + 1200840202194952194, + // @Ronygabriel_ofc + 1208038091962888194, + // @LeoSuricate + 1208544960032727040, + // @ManuVieiraSC + 1210296676520513536, + // @CruzOrleans + 1211689081861623808, + // @michelbeckerbr + 1211812977579421699, + // @mello_bandeira + 1213217041739419648, + // @mfriasoficial + 1213876635331465216, + // @OficialNenemAl + 1216746154626625536, + // @JairSoutoAM + 1218904655591243777, + // @JohnRobertPA + 1219003418347483136, + // @cortezpsol + 1219445803854548994, + // @mariamarighella + 1220387710462021633, + // @BernardinoNOVO + 1223018829393145856, + // @GuiBianco65 + 1223960824403939328, + // @VicMello16 + 1224035025882140673, + // @Vilsondafetaemg + 1224332353927073792, + // @DrVictorAmoras + 1224505061558161408, + // @luanlennonbr + 1224592244465905666, + // @KlesleyGarcia + 1224615230195609601, + // @PortaldoJose + 1224767212101283840, + // @joaobmaresguia + 1225055364124762112, + // @profterezinhaPT + 1225158156545904642, + // @brenogaribalde + 1225774842735202306, + // @gerlanebaccarin + 1225861584972664833, + // @deputadoelizeu + 1226865113635926021, + // @dramayraoficial + 1226874382938787841, + // @marinahelenabr + 1227202002066837504, + // @VaniceMatos + 1227388838592614400, + // @JulioCLamim + 1227572253522526208, + // @Alfredogaspar_ + 1228008951813492739, + // @izadoradias29 + 1231267069762691072, + // @BabaTupinamba + 1232082766071767045, + // @felipegomespsol + 1232252941681221632, + // @amandasalecosta + 1232788071218843648, + // @queciareismbl + 1233041635480678401, + // @carolnunesof + 1233941640441663488, + // @bolsonaronegona + 1234168067191648259, + // @MassamiMiki + 1234278834519920642, + // @juizcubas + 1234488841479884801, + // @JacqueMoraes_es + 1235371041712726022, + // @AndreKubitschek + 1238100206488629248, + // @julinhodeputado + 1238202050745446401, + // @kikoceleguim + 1238511273752563712, + // @NaldyBianca + 1238835947661377536, + // @eusouamom + 1238868377675980803, + // @eduardomenegol_ + 1240359378081001474, + // @juliana_macieel + 1240694913420926976, + // @fabriciochaves_ + 1241145962082467851, + // @barcllay + 1241746085820932098, + // @fabiolopes_38 + 1241821464111853570, + // @cop_santana + 1241842077979283466, + // @elianexunakalo + 1241878280841674752, + // @EiflerPam + 1241960263886286854, + // @depvermelho + 1242093483151839238, + // @claudiaguerramg + 1242834650155941888, + // @gabrielcarvalce + 1243211951863463941, + // @samaramartinsup + 1244783574676627460, + // @docporto + 1244954874049114113, + // @allinyserrao + 1245187621007089666, + // @vitimporto + 1246125973755609088, + // @wmf_oficial + 1246641106391052288, + // @GallinatiRaquel + 1246945909881044994, + // @RodrigoLivra + 1247224255311499264, + // @drluizovando + 1247287375803355136, + // @robertlemosss + 1247578937510723585, + // @enricolopescba + 1247916183392837633, + // @leodacostas + 1248042609731387392, + // @victorjansendf + 1248682133943717889, + // @a_jessicao + 1248830800281378817, + // @CoronelFernand9 + 1249769581725573121, + // @lucaspoleseES + 1250037140647481344, + // @tuannepsol + 1250165291558068230, + // @VandaMonteiro_ + 1251221027822211073, + // @jaziel_dr + 1251558672821600259, + // @cirineu_costa + 1251676296842805257, + // @Victorinogustav + 1252237633008238593, + // @fcordeirosc + 1252561439686037506, + // @ProfDrMauroRosa + 1252709404975169538, + // @fabio_schiochet + 1253698500401008641, + // @ivoneidecaetano + 1253699087209127936, + // @tyago_hoffmann + 1254019417148661760, + // @peritatirotti + 1254073103661088768, + // @tanielmacedo + 1254084527422689282, + // @deputadaleandre + 1254103540387254274, + // @israelcosta_rs + 1255134132889214984, + // @MoraisDino + 1255145793507348482, + // @beckhausersc + 1257020913893195779, + // @DraSilvana2 + 1257476064219140097, + // @GiovaniMattoss + 1258511929146040320, + // @najuliaribeiro + 1259624231765184514, + // @adailfilhoam + 1260696151029940225, + // @camilajarams + 1261423487585005575, + // @brunoseccobr + 1261452346028105728, + // @ze_haroldo + 1262408615564099596, + // @luannasantos_13 + 1262645801936879617, + // @nise_dra + 1262944997118181385, + // @mikaelgringoPE + 1265422039260762113, + // @hana_ghassan + 1266449427985829888, + // @professorarita_ + 1266477339413749761, + // @mariadoscamelos + 1267127677468753920, + // @Junacamara + 1267549843586768896, + // @aanaelisast + 1267822695195930630, + // @ChirleyPankara + 1268292556967874566, + // @maiarafelicioo + 1269374979222765568, + // @manupeloES + 1269632201760690176, + // @DRodrigueiro + 1269706255347650565, + // @ThiagoResiste + 1270774020808626176, + // @juliermesenav + 1271216272148176896, + // @GilvanDaFederal + 1273299541786349569, + // @carladicksonrn + 1273322623380983815, + // @MucioBotelho + 1275549634757365760, + // @LucasCaregnato + 1275636866075803648, + // @GiorgiaPratesMP + 1278774539942584322, + // @CrBeraldo + 1278784862313472005, + // @victorcarvamt + 1280120303818072064, + // @JuniorGeraldo_ + 1280224335307907077, + // @ChinaoRuiz + 1280315144338313217, + // @FeCurti13 + 1280517974449893377, + // @wagnertavaresrj + 1281590900834078720, + // @jonasreispt + 1281598785081180162, + // @KodamaThiago + 1281980054730407936, + // @daniel_sucupira + 1283586848019951621, + // @adrianaraujomg + 1284456064432377856, + // @GlauberPoubel + 1284849642895740930, + // @DepChrisostomo + 1285226176953360384, + // @eugenialima_pe + 1285282863471001601, + // @daniportelape + 1285295088525082631, + // @Sonaira_sp + 1285299871340167173, + // @euLiviaNoronha + 1285729712795525125, + // @ladisouzams + 1286438113896796170, + // @ericodonovo + 1287396715331452932, + // @sorriso_elisa + 1287490158510723072, + // @prjuniortrovao + 1288544254856507392, + // @majorvitorsa + 1290713462285500417, + // @MatheusLaiola + 1292497353166004227, + // @DepCoronel + 1295725787790942208, + // @rzampieri22 + 1295750020554203136, + // @joaoccoser + 1295830662721736711, + // @KuertenRoberto + 1297864783841103877, + // @gustavosefer + 1298369792169127939, + // @junynhomartinss + 1298413776715358209, + // @dinhodowsley + 1298603065457680385, + // @luladafonte + 1300279136620085249, + // @telma_rodolpho + 1302310926482317313, + // @marcelinhoguima + 1303443223663325185, + // @DimasCostaRS + 1304120510821945346, + // @chaves_hildon + 1304134002547331072, + // @pepeliberdadefm + 1304145702122147840, + // @marioleonypsol + 1304884724603772931, + // @draraissasoares + 1304911931036307456, + // @chris4patinhas + 1305541363644104706, + // @deppedrokemp + 1305583653934727175, + // @damiresrinarlly + 1307532712794873856, + // @GuguSeba + 1310253107310460933, + // @CoronelRomualdo + 1310744586139119618, + // @EduardoRiedel_ + 1310977778574032900, + // @PabloSilvaLira + 1311620519935045632, + // @OdairTramontin + 1311753908549750784, + // @gedalvaumbauba + 1312771507601498112, + // @cassymonteiro + 1312841615111843848, + // @delegadkatarina + 1314575606487670784, + // @andersonlimaadm + 1318194625400737792, + // @delegadalia + 1318672895573479424, + // @ElianaBayer_ + 1319073837485576197, + // @anaportelams + 1319825203313184768, + // @DeoliAnderson + 1322213572953493505, + // @PamelaGiedre + 1325859971905572869, + // @bfeministapsol + 1329213002051149824, + // @paulomelo_sa + 1330155785951866881, + // @cabomeireles + 1330393555290951681, + // @patisborges + 1330690808090124289, + // @_PastorDiego + 1334311453592064000, + // @bocalomoficial + 1335685798746845190, + // @RoseanaSarneyM + 1335942718040793091, + // @matheuscampinas + 1338827932849090563, + // @RenatoG05313289 + 1339709453738926081, + // @Danny57037067 + 1341344192874864640, + // @MorandoOrlando + 1341392073799380992, + // @RenatoTaroco + 1343385618559197188, + // @RicattoEduardo + 1344480762180022272, + // @BrandelJunior + 1345054306614042626, + // @ProfessoraNadir + 1345965375582851072, + // @DepDianaBelo + 1346196281832693768, + // @elmovazoficial + 1346929915040571399, + // @CrisWainer + 1348639812497125377, + // @FlavioNJunior + 1349056893235458051, + // @adeildoreisofc + 1349210221604962306, + // @DrGeorgeLins + 1349425371146547204, + // @annacarolinaadv + 1349690681464385540, + // @GuiUchoaJr + 1350924815759269888, + // @rosanadasaude_ + 1351505730323558400, + // @Paulatitan1 + 1351686391265193987, + // @silenoguedes + 1352236209800671232, + // @dimasfabianomg + 1355216660068761606, + // @AnaPimentelmg + 1356434859934298115, + // @faustinorn01 + 1356762322489004032, + // @AndrePiresDF + 1358426024250343427, + // @RafaelFonteles_ + 1358739792163389440, + // @deputadodrPaulo + 1358880298407243778, + // @miroteixeira + 1359157620897153024, + // @FadelMoacyr + 1359162686228078593, + // @robertoclaudio + 1359194831415828488, + // @Nilvo17 + 1362188334924242949, + // @del_guilhermed + 1362442787988381696, + // @boulosnat + 1366796662614749187, + // @NeidinhaSurui + 1368009955912130561, + // @_mapeoficial + 1372519763818217477, + // @NilmaSanches + 1373064659817926661, + // @rodrigoestacho + 1374044251475091457, + // @PiauienseO + 1374436642375659530, + // @VitorioJrMG + 1374502893533884421, + // @FelipeAlecrimPE + 1375498458132537348, + // @rodolfoms + 1376603305879736320, + // @CarolineKalil4 + 1376746550593003525, + // @celiaxakriaba + 1379025327314329608, + // @DrLeviMelo1 + 1379424037068234752, + // @fabiosilveirarn + 1379583456812924935, + // @marleipr + 1380366998488645636, + // @depprofcleiton + 1382460589625204742, + // @apropriajulia + 1382550118314999812, + // @PrefeitoCesar + 1383362690584768514, + // @santanna_cs + 1385208512955957250, + // @SylvioMenicucci + 1385213373332213778, + // @PCO29antonio + 1387843668053200897, + // @depmarcionunes + 1390475281316646917, + // @davibrandaobac1 + 1392451077656698881, + // @juarezcostamt + 1394317065750724608, + // @lucasbovesp + 1394364586346885120, + // @MassaelB + 1395784254269841408, + // @DelegadoMarcus + 1396055754793230341, + // @renataseneofc + 1399347605776314368, + // @jusoaresft + 1400452986178981891, + // @eduacostario + 1400996132982038532, + // @antidio_lunelli + 1402602168838938625, + // @gu_camillo + 1405334928926162946, + // @Sandromarttinss + 1405740717595529225, + // @rubensuchoa_ + 1406043032437264386, + // @lelecopimentel_ + 1408810159376420865, + // @ananiasnauar + 1409194659872653317, + // @sauloportolivei + 1409385758838906882, + // @jrfreitasofc_ + 1409935535443951616, + // @juniorferrari55 + 1409990621536919555, + // @rafaelsaraivasp + 1410633602354794508, + // @ProfeBonatto + 1411854540182327299, + // @yasminvsh + 1412747070511984641, + // @deputadodrhugo + 1413128451154866185, + // @cozzolino_RJ + 1417893820876967943, + // @dudasanchesba + 1422326852539015168, + // @rjandremonteiro + 1423628533696540673, + // @depclaudiac + 1423660842239791115, + // @coronelbonates + 1427816945735319554, + // @CapitaoMartim + 1430162160320188433, + // @MauricioNeves_ + 1430622957102145543, + // @amandagentiI + 1432351330379644929, + // @RicardoAbrao_RJ + 1433065898064191490, + // @robertadahorta + 1434501105241792512, + // @NFgoes + 1435915606541410305, + // @0xcarlosh + 1437029285387243526, + // @PeKelmon + 1437437148769226757, + // @ScalcoDarlan + 1437758990172184577, + // @nisia_trindade + 1438501569301991424, + // @Jadielmoraes20 + 1439280377936367621, + // @OperadoraManu + 1439763846772703232, + // @edilenxavier + 1441394605338021888, + // @ivanirdosantos + 1442471110927290379, + // @mariiluse + 1445155253880594432, + // @renatmirandarj + 1446451214292553733, + // @DayaneCruzPE + 1446618743350730754, + // @DanielaadvAP + 1446856314215337989, + // @joelrodriguespi + 1447921024117493761, + // @atilaliraof + 1448414094441291777, + // @LuizEduardo_RN + 1452666475844706306, + // @PretoniDacio + 1452992283616366607, + // @victordiaspa + 1453677155485896708, + // @RenamTassio + 1454041097395744776, + // @dep_dayany + 1456344779625799688, + // @TremeaMarcio + 1456999668961906690, + // @ChicoVieiramg + 1457348048334594049, + // @NilceBregalda + 1457433413028356098, + // @fala_mafia + 1458089891795976202, + // @sargentoportuga + 1458495272716222465, + // @josecam01577970 + 1459681826700673030, + // @schumarker7 + 1461714774996230145, + // @RicardoArrruda + 1462770465068437504, + // @BrunoBrazKart + 1463552472052502528, + // @KAUMAGN0 + 1463557962589429767, + // @Alexandregonrs + 1463574028468330497, + // @padovanidep + 1464329410115514369, + // @michelkourimbl + 1464589158035468306, + // @marinadomst + 1465395158103597065, + // @amarianalescano + 1468292579775348737, + // @FernandoManso20 + 1468757430632865793, + // @Maubmarcon + 1469007240279597067, + // @orenatomachado + 1471286780737503232, + // @joaquimroriznet + 1472329069027119107, + // @LuizinhoMinas + 1477024586764009478, + // @yurydoparedao + 1478380311532736512, + // @Anderson_ma123 + 1478686290270896128, + // @CapitaBrasil22 + 1478840583481446400, + // @diegocastroba + 1479821013009551371, + // @firmo_oficial + 1481351815669112840, + // @Johnatanmaravi2 + 1481369125159067656, + // @erickmonteiropa + 1481981850159652865, + // @foliveirapr + 1483509817620709381, + // @laisjordy + 1485332978976886786, + // @HungaroIgor + 1487071314410188807, + // @alexsousaam + 1488595085361041408, + // @lucascaculajp + 1488612280241696775, + // @deparimateia + 1488691561185619969, + // @depjaqueline + 1488866652037029889, + // @GayerGus + 1489108473027739654, + // @valdiroliveira_ + 1489559619118805001, + // @AdaoPrettoFilho + 1489648399334969349, + // @DuCazellato + 1489744612558356481, + // @JooDePaulaDosS2 + 1490037307176636420, + // @AndersGimenes + 1491165950422495244, + // @eribertomfilho + 1491176697445683205, + // @SdMadalhano + 1491215424675004421, + // @RosaRezendeGO + 1491387105913810944, + // @joao_herzer + 1492198027620143113, + // @neyamorimac + 1492921263299379207, + // @cleoniceback1 + 1493565805304455174, + // @depchicao + 1496537874237468679, + // @Delboni_Isis + 1496817240049606661, + // @deposcargutz + 1497201650196430861, + // @profangelapsol + 1497605679229644803, + // @fsdecastilho + 1499529892039434251, + // @EnfBrunoFarias + 1500910030593445888, + // @annasebbaj + 1501892430949453827, + // @sgtgoncalves22 + 1504105286876991489, + // @EdianeMariaMTST + 1504228964319080449, + // @edsonferrazsc + 1505881785762304002, + // @AlvaroJeronymo + 1506029258967240715, + // @SKaripuna + 1506313161921777676, + // @rhdeverdade + 1508246534370082821, + // @bm_oliveira14 + 1508783941037268993, + // @Thais_ProfeChef + 1509245802815922177, + // @MaalouliMari + 1509589806325579784, + // @orleansbrandao_ + 1510790014111784967, + // @diegoandrademg + 1510954548348784645, + // @rodrigomarcial_ + 1511399052675633156, + // @WebaNatassia + 1512098373570113538, + // @DelRodrigoSa + 1512105284558282765, + // @MarciaHuculak + 1512756653199872006, + // @reillerlopes + 1513329845551390720, + // @franze_carneiro + 1513903701483790343, + // @mariaarraespe + 1515390927958908928, + // @SocorroWaquiim + 1516019212145287168, + // @PenhaBernardes + 1516033327941181447, + // @DanielleDVale + 1516121393028644869, + // @thiagorodripa + 1516560691410481156, + // @andrebuenoofc + 1517514873105797120, + // @ivanilsonrn + 1517576201338040321, + // @leandrobasson + 1518652449371922433, + // @drpaulomedinamg + 1518728271176871936, + // @EDILSONLIMA70 + 1518759485694779392, + // @EdRaposo_ + 1518965841261441024, + // @luizgastaoce + 1519011181574434816, + // @PazuelloGeneral + 1519440468949573632, + // @drviniciuspi + 1519744542173536259, + // @Ricardinhofoz + 1520051894319788033, + // @Julianoomarti14 + 1520087295822598146, + // @OtacilioDeSous3 + 1520962489931939840, + // @LuizZacarias22 + 1521138334256480259, + // @Alessandrogb74 + 1521179786915233795, + // @pricarrijo28 + 1521558274230919171, + // @clementecampo15 + 1522266010837041153, + // @AnaPaulaGoffi + 1522529565733707777, + // @coronelamadeu + 1523726346232414208, + // @SERGIOD84944206 + 1525599750854266880, + // @SindoleyMorais + 1525834800233336832, + // @MariaClaraMarra + 1526909238622175232, + // @misspretapt + 1526912162194653184, + // @IrineuCruz4 + 1527267836250537984, + // @lucianovieirarj + 1527317724979810304, + // @ArielBrandao7 + 1527695566184009731, + // @Leonegreiross + 1528462133477916678, + // @RicardoEndrigo5 + 1529074828933746688, + // @Marcelo00945011 + 1529506446009827329, + // @CriaCranio + 1530904551582220288, + // @JadyelAlencarr + 1531271393219837956, + // @Jaimenunes_ap + 1531655976008531968, + // @ReginaluciaSi15 + 1531933026208436231, + // @eudonaneuma + 1533265972659990528, + // @melcafariaspb + 1534926600743051266, + // @helenadaasatur + 1536471955209076736, + // @obrenocop + 1536700734892298241, + // @FaoDoBolsonaro + 1537386416963039233, + // @eliassantiagopt + 1538940264173211648, + // @rejanepstu + 1539038049937547269, + // @AnneMarques_AP + 1539992525980815364, + // @eudociasenadora + 1542212008082382855, + // @marussago + 1542233934435696646, + // @a33lucas + 1542497617363492864, + // @AAdautooficial + 1542590670359207941, + // @soududaramos + 1542878687120482306, + // @victorlinhalis + 1543943524189720576, + // @rsodrec + 1544490645636751360, + // @LuisaCela87 + 1545437399446077440, + // @AnibalLins + 1545770498486898689, + // @leandrosoaressp + 1546610777519521794, + // @GuaracyJunior4 + 1546615155009667072, + // @iossef_hassan + 1547001236163100672, + // @AdaXavier9 + 1547257445243830275, + // @DrHenriquePaes + 1547672436933419017, + // @drjosericardo92 + 1548433495097151489, + // @abraogodois + 1548671250679173126, + // @ThiagoManzoniDF + 1549429329494478848, + // @williambarros30 + 1550186795333296128, + // @julio_kuller + 1551652560951562240, + // @DepAndreAbdon + 1552328294665781249, + // @Gi_MonteiroRJ + 1553042102908502021, + // @MissiasDias + 1553167499901943808, + // @draalehaber + 1554462338715197446, + // @pedroaiharamg + 1555552838008414210, + // @catanhopt + 1555559206635315201, + // @PQueirozadv + 1555860014383988737, + // @LaironCarlos + 1556676552674394114, + // @KleberRosa50 + 1556777705387032577, + // @josyjacob_ + 1556827882516893700, + // @VitormoreiraCG + 1557086339740405761, + // @ScaranteRenato + 1557393886431027202, + // @Raphael89700864 + 1558100800664178688, + // @ArianeRSAssis + 1559215037746741251, + // @samuelvianamg + 1559598025420464134, + // @Clecio_Luis + 1560635958994812928, + // @BartonCutler + 1561212563836264449, + // @StellaGaio1224 + 1561545109715517440, + // @carlosmoraestv + 1563342678011842560, + // @vinivenades + 1564049349000183809, + // @DeboraMenezes22 + 1564101654496133120, + // @daiane_daiane18 + 1564304740850221058, + // @jaquepetrovik_ + 1565401487836061698, + // @NMousquer + 1569113308078170115, + // @pitmagrin + 1569391051160313858, + // @oAllaxSiqueira + 1569761874471972865, + // @regisethur + 1574862745325146112, + // @weibetapeba + 1575549757900132379, + // @14luizfranca + 1575672365094440960, + // @PauloEchebarria + 1576793388845760512, + // @ProfAlexandreS + 1577470375272878083, + // @Arnaldodeputado + 1578127734362046489, + // @Mariada67416358 + 1578375608878391296, + // @TamirisPeixoto1 + 1578450870127263757, + // @CarlosValdevin6 + 1579466590512418817, + // @eusouolimpio + 1580519722877128708, + // @quinhoprefeito + 1581818536137302021, + // @marlipaulinopr + 1583101243987206145, + // @renatodepaularj + 1583183342979190785, + // @WesleyCosta_GO + 1584241246817624064, + // @LucasLasmarMG + 1584270172952694784, + // @Joellobatoc + 1584299141550702593, + // @andreadantasc + 1584527735686397952, + // @tcoronel_maia + 1586016154262315009, + // @rafandradembl + 1586192715137667072, + // @CombatPatriota + 1586416104632721416, + // @LucasBackes14pr + 1586733661113794561, + // @marcosfonsecapi + 1587131451593658369, + // @antoniodoidoofc + 1588344483766321154, + // @munique_busson + 1588523690760732672, + // @RONALDO90199231 + 1588685187629678592, + // @gianninogueira2 + 1589336373101723649, + // @JUNIORCESARLEI9 + 1589412695664742401, + // @LuanNun43052492 + 1589645603616817154, + // @pitypaguiar + 1589741307630690304, + // @Marisaloboreal + 1589768586096148483, + // @enzosamuelthe + 1589963570720264193, + // @soldadoarruda + 1590877665837436928, + // @ricadeFreitasG1 + 1591548255313305607, + // @MartaPadovani73 + 1592147414500122624, + // @marcelomb1993 + 1593005387296514048, + // @deputadomatheus + 1593032863066251265, + // @izaarrudape + 1593200340777803776, + // @davivalencambl + 1594709250835709957, + // @FelipeBolsonare + 1595399593012887552, + // @DraLuxmonteiro + 1595451834730168321, + // @shanvirmond + 1597101228181397508, + // @Josecar84649656 + 1598498909608988673, + // @DartanhanCampos + 1599722975997067264, + // @camposoptica + 1599886697096945699, + // @davipqdtrj + 1600721241534504960, + // @vandrofamilia + 1603020469145374720, + // @AlexanderBrasil + 1604310468906176512, + // @Geraldomendesof + 1605547416366841856, + // @MarinaCallega13 + 1609042267095941123, + // @gutembergfrios + 1610745043358187520, + // @WickRyanAM + 1612178912947183620, + // @SchiavoMaurilio + 1612183554787614722, + // @deputadairacema + 1612229099551858688, + // @aveiltonsouza + 1612484459034550273, + // @MalconMazzucato + 1612520249735135232, + // @depgilbertinho + 1612589448256016385, + // @gersonzocchi1 + 1612930692932837377, + // @oGabrielAzevedo + 1612935067487145985, + // @WaldenorPereira + 1613183981230366728, + // @DiegoQuaqua + 1613315565203906562, + // @ToGomes5 + 1613522130368356352, + // @gmoraisdai + 1617357506459598848, + // @realthiagonunes + 1617967890107432961, + // @DepLucianoTo + 1618625038176980995, + // @tadeudesouzaam + 1618970134005121025, + // @DanielSantanaES + 1619000816714629122, + // @delegadagabi + 1619448856009138182, + // @Priscil34438752 + 1620207638397976577, + // @IndiaArmelau + 1622610845002870784, + // @gracinhamaosant + 1624490961236631553, + // @francisconace01 + 1631055193675644932, + // @EdsonKambeba + 1633100825403826177, + // @Jaderbfilho + 1633282883682009091, + // @PrEzequielbueno + 1633928858180239363, + // @DeivisOss12 + 1634553908402987010, + // @Dep_CoronelNeil + 1635697153979891712, + // @LucasAl29111516 + 1637240612620587013, + // @AdrianaLeal2507 + 1637963109573828608, + // @adrianalmeidapt + 1638181881127612416, + // @leticiaaguiarsp + 1638274776413134849, + // @Rogerio_Ulysses + 1641171155540123655, + // @edercostarj + 1642933716279328768, + // @giuargolo + 1645831168586137600, + // @fialhooooo + 1646851212329754626, + // @Edy_Tocantins + 1647755717158264835, + // @wistongomess + 1648046380106104833, + // @Alannleal_ + 1651297457374887936, + // @felixsantosrn + 1653188395324014593, + // @CamilaGodoiSP + 1653466465855471617, + // @MatiasSamuka + 1655587131136307204, + // @DeputadoQuirino + 1656674946163277826, + // @DepComanDANte + 1657817923728166916, + // @drbrunoresende_ + 1661373663256408066, + // @CosttaIvony + 1662406778263420929, + // @wellersind + 1663242435420356609, + // @DouglasRuas_RJ + 1663563614954086401, + // @bemoreiradf + 1666457401816391682, + // @TathianaGuzella + 1668266165922144262, + // @carrarajh23 + 1669323757700169729, + // @vanessarosajlle + 1675490861440802817, + // @feschmittvet + 1676911232614293505, + // @GiFreitas1982 + 1684192295459905536, + // @fellipe1971 + 1690338078827728896, + // @babatupinamba_ + 1691097420166254592, + // @tonytretanews + 1691834974301773824, + // @JacoLulaDaSilva + 1693307464354041856, + // @souricardojusto + 1693956921449996289, + // @MarianaNaime + 1700652030501519360, + // @PauloAssun68233 + 1702440269138866176, + // @mateuspepicepr + 1703957971380604928, + // @opropriokogos + 1707814589209944064, + // @rosianepolitica + 1708485431405293568, + // @marcelodino_rj + 1708844117914972162, + // @willianrochapr + 1709986969277624320, + // @osenhorsaldanha + 1710320865517162496, + // @francisco_arten + 1710430261932908545, + // @hebertcsgyn + 1713936872370532352, + // @oengenheiroleo + 1714288040334831616, + // @TitoBarichello + 1715811403708166144, + // @karisantospt + 1716415063471316993, + // @moreiramissao + 1721592893511483392, + // @ThomazSJC + 1721921193588973568, + // @jorge6050384252 + 1723066872403214336, + // @guikilter + 1725596039263006720, + // @Gloria_Vale72 + 1725943187250831360, + // @eduardowilliamm + 1728271638104358912, + // @ManuellaTyler + 1731034177175171072, + // @profRodneyRJ + 1731046357425668096, + // @aureacarolinax + 1734643690751078400, + // @Aludydias + 1738591794403655682, + // @MalluAlmeida12 + 1739812685187784704, + // @ysanireal + 1740577335395414017, + // @drfabriciok + 1742246045835300864, + // @bellaccarmelo + 1743284668764463105, + // @bittencourt_rg + 1744512067682316289, + // @jotinhapiaui + 1744704300616417280, + // @miriamguzella_ + 1750522589548879872, + // @andrabianchessi + 1752837451175907328, + // @rmpicoli + 1757491174871363584, + // @prwellingtonst + 1758502910638329856, + // @samuel_al_silva + 1760365813456904192, + // @RicardoAlv32716 + 1761085576227213312, + // @FePresidente + 1762888285796343808, + // @dulce_lmendes + 1764371188489330689, + // @JeffreyChiquini + 1767171124629037056, + // @a_caroliveira + 1767282969708822528, + // @Lenesilllva + 1771346241449865216, + // @rafaelsatiebr + 1771625536772571136, + // @MayaraKeiko + 1773171382886576128, + // @jorgemaciel05 + 1775273264774045697, + // @gisvaldopsol + 1776773968726298625, + // @Ap52467Jovelino + 1777297081134178304, + // @BrunoAlvesSFFS + 1777373652708683786, + // @fabiocarneirojp + 1777538059891777537, + // @MarcoRo43166598 + 1777608369236271105, + // @igorrayanrn + 1777621309335183360, + // @Dhe_paiv + 1777901547629764608, + // @juhliasantost + 1778077250908274688, + // @OficialVanucci + 1779252601269141504, + // @asargentolorena + 1779345802629914624, + // @carlosiranrs + 1779913435825795072, + // @BiaCoimbraSP + 1780537023859761153, + // @PedroNassifrj + 1780805773603258368, + // @paulacoutinho65 + 1781108419853701120, + // @CamillaGonda + 1781485235047174144, + // @brunnomattospt + 1782511140561367040, + // @MicheleMarxsc + 1782916900923543552, + // @aucilene_a75929 + 1784209048671322112, + // @eliabecamposs + 1784275708606320641, + // @NivaldoNoga + 1785161935685562368, + // @oescobarpro + 1786921594390016001, + // @LeandroGol15952 + 1789440263187771392, + // @GuimaraesAlpha + 1790692252840243201, + // @EloyNhandewa + 1791120964731744256, + // @coronellrosses + 1793428861348167680, + // @akarinaclaro + 1793823661369044992, + // @nataliademesmao + 1797413322385485825, + // @Soldado_Sampaio + 1798332723989295104, + // @Lindenbergbra + 1799431977625366528, + // @maykondelfinomg + 1800001219932409856, + // @mqueiroz_rio11 + 1800605526683754496, + // @GilsonMachado22 + 1800983500541030400, + // @severoeulalio + 1801749578938564608, + // @leorondonn13 + 1803570674914426880, + // @igor_ooliveira1 + 1809003112658767872, + // @gabi_bvnt + 1809421300433317892, + // @ustramarcelo22 + 1809603567520718849, + // @Adrianadasilvax + 1810502940416991232, + // @passinhoisa + 1812967972358569984, + // @TenenteNilton + 1813622960579588096, + // @RGracie79575 + 1814770780871532545, + // @Fernand51081218 + 1814824898956828672, + // @VolmirGordo + 1815772052063694848, + // @taniadacreche + 1815835183834390528, + // @SandroOmar48237 + 1816185635578978304, + // @JoaoRochaFranca + 1816267815063531520, + // @ofcjusantana + 1820574539312508929, + // @007Douggomes + 1821931832008392704, + // @DaSilva69157 + 1822302684935790593, + // @joaogcandidoadv + 1826732949376479232, + // @betsMartins + 1827352420541763584, + // @MarcusLopesPsol + 1827593492387745792, + // @matheussimoespr + 1836487149769596928, + // @AGoldbach60024 + 1844153830201589760, + // @gabrielpiauhysp + 1844436268706418699, + // @cironogueirapi + 1844442355421614089, + // @ledapco + 1845277420854837249, + // @ProfMarcio58657 + 1845590017445400576, + // @henriquemgjf + 1845839091515981833, + // @GoulartVla23085 + 1846568278950346752, + // @DragUrbana + 1847431729415397376, + // @LuizaDoClezao + 1848153540012945408, + // @abreu_de63729 + 1855074530613403648, + // @luladobemofici + 1855768833115541504, + // @debora__romani + 1855793756332519424, + // @LucasReis13_ + 1856394921189494784, + // @sandsonmenezes + 1856816774517268480, + // @Airtonjose26 + 1857173636689362944, + // @rubensnascimepr + 1857825522765766656, + // @CristianeN74380 + 1858882502707810304, + // @depcabomacielam + 1859357204278542336, + // @LusFelipeV68715 + 1861380355875373056, + // @lucylulv + 1863493161240133632, + // @_AnaHering + 1864125622928187392, + // @denistaveiradn + 1865917852424798208, + // @joaopaulo_tprs + 1870514350164705280, + // @Pcbcastelo + 1872641473021087744, + // @saulofreitas22 + 1873333201445289984, + // @Lukaovereador + 1874897810753208320, + // @johnysantos_sp + 1875239088028299264, + // @oiurecastro + 1878758512639033344, + // @fsantanapsd55 + 1878782691375759360, + // @FlavioMant5441 + 1878922692860006401, + // @grazimacedo_ + 1879635567639777280, + // @ProfRonaldo13 + 1882045370613600257, + // @_marquinhostrad + 1883938805306052608, + // @bolsonaro__jr + 1884423273490108416, + // @digportella + 1884620263549001728, + // @yasminsarrafsp + 1884758507272216576, + // @ravengaar + 1885085802645880832, + // @ameliocayresdep + 1887508211445702656, + // @FelipeVasquesce + 1889410539639517184, + // @celprincipebr + 1890456925147734016, + // @Herculano1011 + 1892340577796313092, + // @julianafideliis + 1896723491371581440, + // @gualbertoap + 1898078097452548097, + // @FRomulo13 + 1898340153896136704, + // @Rpachecopinho + 1898684990428151808, + // @viscontioficial + 1898913625697292288, + // @GersonClaroMS + 1899113980544897024, + // @galvaomicheles + 1899797969400180737, + // @Moanavaladares_ + 1899875749907357697, + // @caixeta_oficial + 1900205421123817472, + // @falaCosenza + 1900729566542655488, + // @RafaMinatoSP + 1901680009430892544, + // @MacAntonioRJ + 1902731327461318656, + // @LenirOficial + 1904187270342574082, + // @profeVinicius + 1906160472354873344, + // @ErikaGuerrieri + 1906823280281174016, + // @rodrigospada_ + 1908175828417949697, + // @MarinhoGui65411 + 1909642092424368128, + // @glenioseixas + 1909726014202171392, + // @yolandasilva_sc + 1911181511837011968, + // @betioldebochado + 1911652608915193856, + // @DrDiogoFranco + 1915836821826646016, + // @RenatoAngraRJ + 1916570073663475712, + // @aparecidobian + 1920469901136764928, + // @ericafmissao26 + 1922067675074727936, + // @joselaviola + 1922378485412331522, + // @daniele_carva + 1925558068206518272, + // @MuriloM50730958 + 1928141892065316864, + // @paulomeloparana + 1929325555234803712, + // @manoela__peres + 1932149081620484096, + // @souadibelias + 1932784165545533440, + // @leo_grandini13 + 1932924340917710848, + // @Camargo1Amanda + 1937005496398962688, + // @JulianaMontrios + 1937336667356205056, + // @prof_elson_sc + 1937980838299504644, + // @marcioalvinosp + 1940410346411597824, + // @ArturdeFariasV + 1941922310568443904, + // @beto_vaz_ + 1944257271942287360, + // @jotabrandaoam + 1949664156606459905, + // @semeadorborges + 1952556730463707137, + // @DrSergioNeves + 1953095258784579586, + // @IgorMarquesBRA + 1957556152117653505, + // @lucasvenditeMS + 1958412111752699904, + // @NoronhaPro7153 + 1958622002719182848, + // @RubensAngiolett + 1959013340661035008, + // @sebasticoelho + 1960040602877190144, + // @isabellaGedeon + 1960414219074998272, + // @DelAmirSalmen + 1961811802271952896, + // @blogmmedeiros + 1962570898310840320, + // @brunoboaretto + 1963656656342192129, + // @josearrudadf + 1964022075540312064, + // @juliemilkreal + 1965782186403012608, + // @delegadali47533 + 1966489240566444033, + // @edilsondamiaorr + 1970177461070528512, + // @joaoadolfo_14 + 1970451742430035968, + // @Drfilipecm + 1972344363654037504, + // @delsmontanari + 1973824020907700224, + // @romulobraz_13 + 1975026318975840257, + // @DrManuelMarcos + 1975317845962858501, + // @VereadorFox + 1975709711870906369, + // @CarlosCostaPE10 + 1975873882654441472, + // @RangelJucy18837 + 1977431822557700096, + // @isapaixaosp + 1980463624649945088, + // @EliCorreaFilhoo + 1980794408141258752, + // @drjuliostobbe + 1981365553173299200, + // @Kiko_Caputo + 1982938414002647045, + // @mazzei_fel47870 + 1983280886067195904, + // @marciaabrahaodf + 1985760866109964288, + // @alepaivaoficial + 1986993348008419328, + // @GustavoHenryR + 1987721882838462464, + // @fig35508 + 1988116025670594563, + // @DaClaudio33805 + 1988208205575712772, + // @RenatoBolsonar0 + 1988322911367950337, + // @vandawitoto + 1988323381792681984, + // @SamaraMadureira + 1990473255539621889, + // @YanProf + 1992656375382704128, + // @lemuel_SV + 1995470250738122752, + // @MarianaServente + 1996216267678875648, + // @jorgerrosario + 1997124812662362113, + // @oleandroissa + 1998066982999212032, + // @Euberlucas14 + 1998418256953233408, + // @profakellsilva + 1998797409712009216, + // @joaopiresrj + 2000205916009074690, + // @nelsongrasselli + 2002006532184281088, + // @rafoliver_pa + 2005281365072449536, + // @moisesbarboza + 2006965910309896192, + // @alvarenga35289 + 2007441310102523904, + // @coronelbusnello + 2009066265156202496, + // @RogerioChimi + 2010017755194589185, + // @AntoniadeJessp + 2010525266565857280, + // @MichelPadao + 2010532868271816705, + // @profraydf + 2010751427623497728, + // @PolianadoPsol + 2010869406650245120, + // @edinhosouzaaa + 2011825328793296896, + // @Efreu_Quintana + 2013273184640860160, + // @meunomeejhonebr + 2013398295859609601, + // @brenomacedopi + 2013660202873032707, + // @esthermoraessp + 2014421538418642944, + // @robson_cacau + 2015964533689323520, + // @carolontiveros0 + 2016185817392099336, + // @anderalvesgo + 2016851339624538112, + // @Landovoigttt + 2017294158558339072, + // @ingridcardososp + 2017575066381238272, + // @sicchar1389 + 2018445464908042240, + // @isabeldesouza01 + 2020307564903165953, + // @vanesckaessusp + 2021065723573829633, + // @leninhavalente_ + 2021229327594000386, + // @tamirisgodoysp + 2021276913012969472, + // @MatheusCambuiBa + 2021559145728454656, + // @ruandutrace + 2021639549625991168, + // @henriqueduraesd + 2021718062261506049, + // @Jordambritosc + 2022354011035222016, + // @sophiafechinece + 2022429125399478273, + // @mateusquadrosms + 2024092564014669824, + // @tulioteconta + 2024178238990536704, + // @NetoFeitos68916 + 2026461466728292352, + // @catarinanevespb + 2026549782278529024, + // @jessicathis_mbl + 2026754488128909313, + // @GreguyLoooban + 2027415714953396224, + // @AraceliLemosOF + 2028870149546373120, + // @rfurtado22 + 2028873608903401472, + // @mendes__babi + 2029985728059559938, + // @helen_vitaRJ + 2033587304795963392, + // @jeaninepresente + 2034019355629846528, + // @CidaCarvalhoSP + 2034086382705274880, + // @EdneyBatalha + 2034242236678848513, + // @WickRyaAM + 2034644862679547904, + // @Fabio_x86 + 2036151026986696704, + // @MarcaoVivacqua + 2036163287587467265, + // @isamamedi027 + 2036512323020500992, + // @missaogabi + 2036817223558234112, + // @marinanamissao + 2036894494973440000, + // @isabeldesouzasp + 2036944910713081856, + // @beatrizoPR + 2037168439836512257, + // @RicardoBec31758 + 2037169927618961408, + // @stellabragasp + 2040479116068081664, + // @maurodeAL1930 + 2041451935685906432, + // @DepDrFlavio + 2041621533454450688, + // @oenzozibellini + 2041763427488608256, + // @Grazypasqualeto + 2042058510536482816, + // @vivianegomesrn + 2042331395725430784, + // @CrisNavarroMIDH + 2043108953177899008, + // @vanessacfortes + 2043516611831701505, + // @owilsonmartins + 2043714770591735809, + // @CarolineSa5644 + 2043839578394533890, + // @viniciusdiaspi + 2045135821695574016, + // @CMolinariBR + 2046074615957540864, + // @VivianeSan11286 + 2046810859695947776, + // @emersonrrosa + 2047035719500103680, + // @Barbarabbotega + 2047324463612485632, + // @JoaoPaulo_2026 + 2047395841715920897, + // @DeyconSP + 2048379106517987328, + // @VasconcellosCel + 2048607450593468416, + // @PablodoMST + 2048698490906169344, + // @CGasparin11011 + 2048961988999454720, + // @livianoronha_PA + 2049127918794670080, + // @MarcioRezendeRJ + 2049539394701262848, + // @bia_pedagoga + 2049646600088117249, + // @enzodangelo_mbl + 2049839952624492544, + // @oalanmansurrj + 2049927296442658816, + // @14DanielAguiar + 2050843854325067777, + // @jaimeverruckms + 2052755712644730887, + // @Aminjhannouche + 2053565789421047808, + // @ThefiAmancio21 + 2054040940885532672, + // @flaviadovalmg + 2054361311832616960, + // @marianasartor50 + 2054619154762674176, + // @onenencoelho + 2054729517864771584, + // @JulianaBrizola + 2054954070083813376, + // @mahmoudamer_rs + 2054974955603861504, + // @luizgamabsb + 2055147146618302464, + // @DelEduardoK + 2055437449740873728, + // @glaucelima12 + 2056410704031121408, + // @edmartresoitao + 2056559848750276608, + // @celmarcioasouza + 2056815030935416832, + // @CidaVillasBoas + 2056824758201708544, + // @RodolfoFiorucci + 2057456621865881600, + // @brenobarcelos_ + 2057550738427887616, + // @PL22Al + 2057808888297062400, + // @ProfNelsiWelter + 2058186878449278976, + // @obrasilcomlula + 2058595300873289728, + // @andreiamoura_df + 2059328849125490688, + // @eleusespaivaf + 2059680529746644993, + // @vivianeluizams + 2059726216119107587, + // @ProfWiterNaves + 2059790203091341316, + // @profcfabian + 2060000313843572736, + // @mayaelliz + 2060371890124857347, + // @cmidf_oficial + 2060549322207375368, + // @nicolasravipsol + 2061280727711289345, + // @Renatofonsecape + 2061591502136983552, + // @DrCrisVeloso + 2061889166779019264, + // @DaversonMatos + 2062692988606717952, + // @jmonteirosc + 2062866727252250624, + // @CoronelMenezes_ + 2062867137497059328, + // @profjoaohs + 2063783489275600896, + // @viniciusrodsp + 2064052780407336960, + // @aBiaAlcantara + 2064088166588440577, + // @DelmaPSOL + 2064731236023541760, + // @pr_itamar_paim + 2064819106545565696, + // @MalluCortes + 2065526840198832128, + // @FelipeGambaroP + 2066731525400330240, + // @drcassiohprado + 2066997147971497984, + // @manoelsseverino + 2067256474024173568, + // @nenemalbuquerqu + 2067271467000045568, + // @SofiaFavero_ + 2067775247542095872, + // @HVilelaporGoias + 2069828777517973504, + // @verapinheirosc + 2069869784972369920, + // @Waltinhogo + 2070185289100775424, + // @candidotelesdr + 2070216303479066624, + // @Deborahzanchi_ + 2070650915871240192, + // @Delmariacorsato + 2070667303931256832, + // @transformarpsol + 2070866245293932545, + // @avictoriagallo + 2071779601815134208, + // @angelaabukce + 2071949683736354816, + // @wilsonamigao + 2071988583334862849, + // @Drrafujr + 2074177828606652416, + // @rogeriozabdalla + 2074246655759552513, + // @MariRibeiro42 + 2074876656561393664, + // @WiliansDouglaas + 2074899795404103681, + // @drthalescoelho + 2075213887213797376, + // @ArndaAcademia + 2075242788568899584, + // @RafaelPerlasca + 2076673278274383872, + // @AlinemunizRJ + 2076717583307255808, + // @zealexandrerj + 2076749095092256768, + // @CarineTAdv + 2077757274961911808, + // @LarianeTellMend + 2077832017736011777, + // @DepCarlaMachado + 2078177903271956480, + // @Nayladasilva0 + 2078543961333850112, + // @drbrenoaraujo + 2079197725384454144, + // @raimundomce + 2079247986442248192, + // @evaldo_gomespi + 2079281807556587520, + // @colet_unidade + 2079306862097215488, + // @FernandoEsporte + 2079569340517462016, + // @MaedjaCampos + 2079640052393472000, + // @PauloMassettiMS + 2079658058842591232, + // @AlineDinizadv + 2079710939197177856, + // @PaiRoblez + 2080002260814233600, + // @miss__Paulinha + 2080010486397988864, + // @colemulheressp + 2080282215653441537, + // @helderdelegado + 2080292155315138560, + // @JoseMoitamxx7 + 2080321827289677824, + // @matteus_hnrq13 + 2080378112114593792, + // @pituxosergipe + 2080630684989730816, + // @Irmamabelmelo + 2080632627887734784, + // @thiagocampeloof + 2081154349267345408, + // @erigreyce__ + 2081487029443940352, + // @DanielSantanamc + 2081774169167974415, + // @luismartarj + 2081890844735426560, + // @drjoaomartins26 + 2082193801267888129, + // @Julio_Neto_SP + 2082200815763148800, + // @BibianoRN + 2082203221980778496, + // @DorvaniloNilo1 + 2082221380297236480, + // @SubTenSergio + 2082541411111493635, + // @anamariaa_of + 2082561982377406464, + // @Fabiani_vasco + 2082567723423342592, + // @Clarianabr + 2082581908005806081, + // @joaogusmao_pt + 2082766963659403264, + // @karimeefayad + 2082897339065233409, + // @barbararesende0 + 2082924803145551872, + // @guihenriquesc + 2082932165466058752, + // @odrlucasoficial + 2083260564294295552, + // @JHONNSOM70 + 2083921524394700800, + // @eliethdefatima + 2084021684495904769, + // @Profsamuelsiebr + 2084048343609540608, + // @KelllenGuerra + 2084209199404191745, + // @profalumatias + 2084258501887406080, + // @Suely7033 + 2084272983443374080, + // @capitaodaviof + 2084329694938251264, + // @drgiovanimendes + 2084387783674564608, + // @tiocarlosrio + 2084426199242022912, + // @ProfessorBezerr + 2084640378666295296, + // @AuriJuniorr13 + 2084644771927076864, + // @Rosangelanegaro + 2084724540194643968, + // @CarmemOliverofc + 2084756136004083712, + // @PauloLealMissao + 2085083851395620864, + // @panayotisdolula + 2085133776741437441, + // @diegojejees + 2085201272139886592, + // @profyoshida + 2085344308480049152, + // @lucianoleitoa_ + 2085736572914159616, + // @Virginiaenfer + 2085739235093446656, + // @profandersonfig + 2085812265509388288, + // @nandovianapsol + 2086122224126136320, + // @boracomMarcel + 2086812805034815488, + // @roneymariachi + 2086873322587881473, + // @RaphaelaSpadi + 2087175409540497408, + // @capitaocesario + 2087749550719086592, + // @DarbideJesusrr + 2088266727519862784, + // @tatibarrapa + 2088334004445487104, + // @Mariguedes1406 + 2088832528383639552, + // @PedroAbib15 + 2089024802669334528, + // @brunoscheid222 + 2089064009643220993, + // @RafaGonzagaSP + 2089736144921444352, + // @marciacandidata + 2089837151844212736, + // @SimonePimehb + 2090082297370312704, + // @RadmanGadiel + 2090252697354067968, + // @cabodaciolo33 + 2091177221310345216, + // @danilosoaresce + 2091948027967676416, + // @IuriMarque60kt + 2092800629278126081, + ]) +}); + +pub struct Brazil2026ElectionFilter; + +impl Brazil2026ElectionFilter { + fn is_excluded_author(user_id: u64, followed_user_ids: &FxHashSet) -> bool { + BRAZIL_2026_ELECTION_USER_IDS.contains(&user_id) && !followed_user_ids.contains(&user_id) + } + + fn should_remove(candidate: &PostCandidate, followed_user_ids: &FxHashSet) -> bool { + let is_excluded = |user_id: u64| Self::is_excluded_author(user_id, followed_user_ids); + if is_excluded(candidate.author_id) { + return true; + } + if candidate.retweeted_user_id.is_some_and(is_excluded) { + return true; + } + if candidate.quoted_user_id.is_some_and(is_excluded) { + return true; + } + // Drop replies whose conversation context would surface listed authors + // (ancestors are returned on the scored post for For You thread UI). + candidate.ancestor_users.iter().copied().any(is_excluded) + } +} + +impl Filter for Brazil2026ElectionFilter { + fn filter( + &self, + query: &ScoredPostsQuery, + candidates: Vec, + ) -> FilterResult { + let followed_user_ids: FxHashSet = query + .user_features + .followed_user_ids + .iter() + .map(|&id| id as u64) + .collect(); + + let (removed, kept): (Vec<_>, Vec<_>) = candidates + .into_iter() + .partition(|candidate| Self::should_remove(candidate, &followed_user_ids)); + + FilterResult { kept, removed } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Stable sample ids from the set (sorted source list order). + const SAMPLE_LISTED: [u64; 4] = [14160928, 14492205, 15022409, 20242549]; + + fn make_candidate(tweet_id: u64, author_id: u64) -> PostCandidate { + PostCandidate { + tweet_id, + author_id, + ..Default::default() + } + } + + #[test] + fn keeps_candidates_from_non_listed_authors() { + let filter = Brazil2026ElectionFilter; + let candidates = vec![ + make_candidate(1, 1), + make_candidate(2, 2), + make_candidate(3, 3), + ]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 3); + assert_eq!(result.removed.len(), 0); + } + + #[test] + fn removes_candidates_from_listed_authors() { + let filter = Brazil2026ElectionFilter; + let listed = SAMPLE_LISTED[0]; + let candidates = vec![ + make_candidate(1, 1), + make_candidate(2, listed), + make_candidate(3, 3), + ]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 2); + assert_eq!(result.kept[0].tweet_id, 1); + assert_eq!(result.kept[1].tweet_id, 3); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].author_id, listed); + } + + #[test] + fn removes_retweets_of_listed_authors() { + let filter = Brazil2026ElectionFilter; + let listed = SAMPLE_LISTED[1]; + let mut retweet = make_candidate(10, 999); + retweet.retweeted_user_id = Some(listed); + let candidates = vec![make_candidate(1, 1), retweet]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 1); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 10); + } + + #[test] + fn removes_quotes_of_listed_authors() { + let filter = Brazil2026ElectionFilter; + let listed = SAMPLE_LISTED[2]; + let mut quote = make_candidate(20, 888); + quote.quoted_user_id = Some(listed); + let candidates = vec![make_candidate(1, 1), quote]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 1); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 20); + } + + #[test] + fn removes_replies_with_listed_ancestor_users() { + let filter = Brazil2026ElectionFilter; + let listed = SAMPLE_LISTED[3]; + let mut reply = make_candidate(30, 777); + reply.in_reply_to_tweet_id = Some(29); + reply.ancestors = vec![29, 28]; + reply.ancestor_users = vec![listed, 42]; + let candidates = vec![make_candidate(1, 1), reply]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 1); + assert_eq!(result.kept[0].tweet_id, 1); + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 30); + } + + #[test] + fn keeps_replies_with_only_non_listed_ancestor_users() { + let filter = Brazil2026ElectionFilter; + let mut reply = make_candidate(40, 777); + reply.in_reply_to_tweet_id = Some(39); + reply.ancestors = vec![39, 38]; + reply.ancestor_users = vec![100, 200]; + let candidates = vec![make_candidate(1, 1), reply]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert_eq!(result.kept.len(), 2); + assert_eq!(result.removed.len(), 0); + } + + #[test] + fn empty_candidates_list() { + let filter = Brazil2026ElectionFilter; + let result = filter.filter(&ScoredPostsQuery::default(), vec![]); + + assert!(result.kept.is_empty()); + assert!(result.removed.is_empty()); + } + + #[test] + fn all_candidates_listed() { + let filter = Brazil2026ElectionFilter; + let listed_a = SAMPLE_LISTED[0]; + let listed_b = SAMPLE_LISTED[1]; + let candidates = vec![make_candidate(1, listed_a), make_candidate(2, listed_b)]; + + let result = filter.filter(&ScoredPostsQuery::default(), candidates); + + assert!(result.kept.is_empty()); + assert_eq!(result.removed.len(), 2); + } + + #[test] + fn hardcoded_list_is_non_empty() { + assert!(!BRAZIL_2026_ELECTION_USER_IDS.is_empty()); + assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2554); + } + + #[test] + fn contains_known_listed_ids() { + let no_follows = FxHashSet::default(); + assert!(Brazil2026ElectionFilter::is_excluded_author( + 40053694, + &no_follows + )); + assert!(Brazil2026ElectionFilter::is_excluded_author( + 21069302, + &no_follows + )); + assert!(Brazil2026ElectionFilter::is_excluded_author( + 1355216660068761606, + &no_follows + )); + assert!(!Brazil2026ElectionFilter::is_excluded_author( + 0, + &no_follows + )); + assert!(!Brazil2026ElectionFilter::is_excluded_author( + 1, + &no_follows + )); + for id in SAMPLE_LISTED { + assert!(Brazil2026ElectionFilter::is_excluded_author( + id, + &no_follows + )); + } + } +} diff --git a/home-mixer/filters/following_viewer_muted_keyword_filter.rs b/home-mixer/filters/following_viewer_muted_keyword_filter.rs new file mode 100644 index 00000000..b5ad3769 --- /dev/null +++ b/home-mixer/filters/following_viewer_muted_keyword_filter.rs @@ -0,0 +1,65 @@ +use crate::models::candidate::PostCandidate; +use crate::models::query::ScoredPostsQuery; +use std::sync::Arc; +use xai_candidate_pipeline::filter::{Filter, FilterResult}; +use xai_post_text::{MatchTweetGroup, TokenSequence, TweetTokenizer, UserMutes}; + +pub struct FollowingViewerMutedKeywordFilter { + pub tokenizer: Arc, +} + +impl FollowingViewerMutedKeywordFilter { + pub fn new() -> Self { + Self { + tokenizer: Arc::new(TweetTokenizer::new()), + } + } +} + +impl Filter for FollowingViewerMutedKeywordFilter { + fn filter( + &self, + query: &ScoredPostsQuery, + candidates: Vec, + ) -> FilterResult { + let muted_keywords = &query.user_features.muted_keywords; + if muted_keywords.is_empty() { + return FilterResult { + kept: candidates, + removed: vec![], + }; + } + + let tokenizer = Arc::clone(&self.tokenizer); + tokio::task::block_in_place(|| { + let token_sequences: Vec = muted_keywords + .iter() + .map(|k| tokenizer.tokenize(k)) + .collect(); + let matcher = MatchTweetGroup::new(UserMutes::new(token_sequences)); + + let mut kept = Vec::new(); + let mut removed = Vec::new(); + for candidate in candidates { + if candidate_matches(&candidate, &tokenizer, &matcher) { + removed.push(candidate); + } else { + kept.push(candidate); + } + } + FilterResult { kept, removed } + }) + } +} + +fn candidate_matches( + candidate: &PostCandidate, + tokenizer: &TweetTokenizer, + matcher: &MatchTweetGroup, +) -> bool { + std::iter::once(candidate.tweet_text.as_str()) + .chain(candidate.quoted_tweet_text.as_deref()) + .chain(candidate.ancestor_texts.values().map(String::as_str)) + .filter(|text| !text.is_empty()) + .any(|text| matcher.matches(&tokenizer.tokenize(text))) +} diff --git a/home-mixer/filters/mod.rs b/home-mixer/filters/mod.rs index b9f09b9a..ce6231cc 100644 --- a/home-mixer/filters/mod.rs +++ b/home-mixer/filters/mod.rs @@ -2,15 +2,16 @@ pub mod ad_adjacent_served_filter; pub mod age_filter; pub mod ancillary_vf_filter; pub mod author_socialgraph_filter; +pub mod brazil_2026_election_filter; pub mod core_data_hydration_filter; pub mod dedup_conversation_filter; pub mod drop_duplicates_filter; pub mod following_retweet_deduplication_filter; +pub mod following_viewer_muted_keyword_filter; pub mod ineligible_subscription_filter; pub mod invalid_conversation_module_filter; pub mod inventory_holdout_filter; -pub mod muted_keyword_filter; pub mod new_user_min_engagement_filter; pub mod oon_nsfw_simclusters_filter; pub mod oon_retweet_reply_filter; @@ -25,3 +26,4 @@ pub mod self_tweet_filter; pub mod topic_ids_filter; pub mod vf_filter; pub mod video_filter; +pub mod viewer_muted_keyword_filter; diff --git a/home-mixer/filters/muted_keyword_filter.rs b/home-mixer/filters/viewer_muted_keyword_filter.rs similarity index 86% rename from home-mixer/filters/muted_keyword_filter.rs rename to home-mixer/filters/viewer_muted_keyword_filter.rs index 463e50eb..482a8182 100644 --- a/home-mixer/filters/muted_keyword_filter.rs +++ b/home-mixer/filters/viewer_muted_keyword_filter.rs @@ -4,11 +4,11 @@ use std::sync::Arc; use xai_candidate_pipeline::filter::{Filter, FilterResult}; use xai_post_text::{MatchTweetGroup, TokenSequence, TweetTokenizer, UserMutes}; -pub struct MutedKeywordFilter { +pub struct ViewerMutedKeywordFilter { pub tokenizer: Arc, } -impl MutedKeywordFilter { +impl ViewerMutedKeywordFilter { pub fn new() -> Self { let tokenizer = TweetTokenizer::new(); Self { @@ -17,7 +17,7 @@ impl MutedKeywordFilter { } } -impl Filter for MutedKeywordFilter { +impl Filter for ViewerMutedKeywordFilter { fn filter( &self, query: &ScoredPostsQuery, @@ -84,7 +84,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_no_muted_keywords() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec![]); let candidates = vec![ @@ -100,7 +100,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_simple_keyword_match() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["spam".to_string()]); let candidates = vec![ @@ -118,7 +118,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_hashtag_keyword_without_hash() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["widget".to_string()]); let candidates = vec![ @@ -138,7 +138,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_hashtag_keyword_with_hash() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["#launch".to_string()]); let candidates = vec![ @@ -156,7 +156,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_mention_keyword() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["exampleuser".to_string()]); let candidates = vec![ @@ -174,7 +174,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_multi_word_phrase() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["crypto scam".to_string()]); let candidates = vec![ @@ -192,7 +192,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_multiple_muted_keywords() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec![ "spam".to_string(), "scam".to_string(), @@ -215,7 +215,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_case_insensitive_matching() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["SPAM".to_string()]); let candidates = vec![ @@ -234,7 +234,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_unicode_and_accents() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["café".to_string()]); let candidates = vec![ @@ -252,7 +252,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_empty_candidates() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["spam".to_string()]); let candidates = vec![]; @@ -265,7 +265,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_all_candidates_removed() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["spam".to_string()]); let candidates = vec![ @@ -279,9 +279,27 @@ mod tests { assert_eq!(result.removed.len(), 2); } + #[tokio::test(flavor = "multi_thread")] + async fn test_cjk_keyword_whole_token_only() { + let filter = ViewerMutedKeywordFilter::new(); + let query = create_test_query(vec!["京都".to_string()]); + + let candidates = vec![ + create_test_candidate(1, "東京都に行くのが楽しみ"), + create_test_candidate(2, "I visited 京都 last week"), + create_test_candidate(3, "unrelated content"), + ]; + + let result = filter.filter(&query, candidates); + + assert_eq!(result.removed.len(), 1); + assert_eq!(result.removed[0].tweet_id, 2); + assert_eq!(result.kept.len(), 2); + } + #[tokio::test(flavor = "multi_thread")] async fn test_punctuation_handling() { - let filter = MutedKeywordFilter::new(); + let filter = ViewerMutedKeywordFilter::new(); let query = create_test_query(vec!["bitcoin".to_string()]); let candidates = vec![ diff --git a/home-mixer/main.rs b/home-mixer/main.rs index 511cd25a..fdcc3dc3 100644 --- a/home-mixer/main.rs +++ b/home-mixer/main.rs @@ -137,7 +137,7 @@ async fn main() -> anyhow::Result<()> { .metrics_port(args.metrics_port) .datacenter(args.datacenter) .otel_endpoint(args.otel_endpoint) - .with_featureswitches(params::FS_PATH, true) + .with_featureswitches_experiment_logging(params::FS_PATH) .with_decider(params::decider_path(), None) .with_tls(TlsMode::server_mtls_from_env()?) .with_max_connection_age(Duration::from_secs(300)) diff --git a/home-mixer/models/brand_safety.rs b/home-mixer/models/brand_safety.rs index b225d8c9..38cee469 100644 --- a/home-mixer/models/brand_safety.rs +++ b/home-mixer/models/brand_safety.rs @@ -9,22 +9,27 @@ pub enum BrandSafetyVerdict { Safe = 1, LowRisk = 2, MediumRisk = 3, + HighRisk = 4, } -pub(crate) const MEDIUM_RISK_LABELS: &[SafetyLabelType] = &[ +pub(crate) const HIGH_RISK_LABELS: &[SafetyLabelType] = &[ SafetyLabelType::NSFW_HIGH_PRECISION, SafetyLabelType::NSFW_HIGH_RECALL, + SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, + SafetyLabelType::PDNA, + SafetyLabelType::EGREGIOUS_NSFW, + SafetyLabelType::SOFT_NSFW, +]; + +pub(crate) const MEDIUM_RISK_LABELS: &[SafetyLabelType] = &[ SafetyLabelType::NSFA_HIGH_PRECISION, SafetyLabelType::NSFA_KEYWORDS_HIGH_PRECISION, - SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, SafetyLabelType::NSFW_REPORTED_HEURISTICS, SafetyLabelType::GORE_AND_VIOLENCE_REPORTED_HEURISTICS, SafetyLabelType::NSFW_CARD_IMAGE, SafetyLabelType::DO_NOT_AMPLIFY, SafetyLabelType::MALICIOUS_URL, SafetyLabelType::NSFA_COMMUNITY_NOTE, - SafetyLabelType::PDNA, - SafetyLabelType::EGREGIOUS_NSFW, SafetyLabelType::GROK_NSFA, SafetyLabelType::NSFW_TEXT, ]; @@ -41,6 +46,10 @@ pub fn compute_verdict( labels: &HashMap, tweet_id: u64, ) -> BrandSafetyVerdict { + if HIGH_RISK_LABELS.iter().any(|l| labels.contains_key(l)) { + return BrandSafetyVerdict::HighRisk; + } + if MEDIUM_RISK_LABELS.iter().any(|l| labels.contains_key(l)) { return BrandSafetyVerdict::MediumRisk; } @@ -62,6 +71,62 @@ pub fn compute_verdict( BrandSafetyVerdict::Safe } +pub(crate) const MEDIUM_RISK_LABELS_V2: &[SafetyLabelType] = &[ + SafetyLabelType::NSFA_KEYWORDS_HIGH_PRECISION, + SafetyLabelType::NSFW_REPORTED_HEURISTICS, + SafetyLabelType::GORE_AND_VIOLENCE_REPORTED_HEURISTICS, + SafetyLabelType::NSFW_CARD_IMAGE, + SafetyLabelType::DO_NOT_AMPLIFY, + SafetyLabelType::MALICIOUS_URL, + SafetyLabelType::NSFA_COMMUNITY_NOTE, + SafetyLabelType::GROK_NSFA_V2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2, + SafetyLabelType::NSFW_TEXT, +]; + +pub(crate) const LOW_RISK_LABELS_V2: &[SafetyLabelType] = &[ + SafetyLabelType::GROK_NSFA_LIMITED_V2, + SafetyLabelType::NSFA_HIGH_RECALL, +]; + +const V2_WRITTEN_LABELS: &[SafetyLabelType] = &[ + SafetyLabelType::GROK_SFA_V2, + SafetyLabelType::GROK_NSFA_V2, + SafetyLabelType::GROK_NSFA_LIMITED_V2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2, +]; + +pub(crate) fn compute_verdict_v2( + labels: &HashMap, + tweet_id: u64, +) -> BrandSafetyVerdict { + if !V2_WRITTEN_LABELS.iter().any(|l| labels.contains_key(l)) { + return compute_verdict(labels, tweet_id); + } + if HIGH_RISK_LABELS.iter().any(|l| labels.contains_key(l)) { + return BrandSafetyVerdict::HighRisk; + } + if MEDIUM_RISK_LABELS_V2.iter().any(|l| labels.contains_key(l)) { + return BrandSafetyVerdict::MediumRisk; + } + + let scored_by_grok = labels.contains_key(&SafetyLabelType::GROK_SFA_V2) + || labels.contains_key(&SafetyLabelType::GROK_NSFA_LIMITED_V2); + if !scored_by_grok { + return BrandSafetyVerdict::MediumRisk; + } + + if tweet_id >= PTOS_CUTOFF_TWEET_ID && !labels.contains_key(&SafetyLabelType::PTOS_REVIEWED) { + return BrandSafetyVerdict::MediumRisk; + } + + if LOW_RISK_LABELS_V2.iter().any(|l| labels.contains_key(l)) { + return BrandSafetyVerdict::LowRisk; + } + + BrandSafetyVerdict::Safe +} + pub fn worst_verdict(a: &BrandSafetyVerdict, b: &BrandSafetyVerdict) -> BrandSafetyVerdict { if *a as i32 >= *b as i32 { *a @@ -117,14 +182,14 @@ mod tests { } #[test] - fn medium_risk_with_nsfw_label() { + fn high_risk_with_nsfw_label() { let labels = labels_with(&[ SafetyLabelType::GROK_SFA, SafetyLabelType::NSFW_HIGH_PRECISION, ]); assert_eq!( compute_verdict(&labels, PRE_CUTOFF_ID), - BrandSafetyVerdict::MediumRisk + BrandSafetyVerdict::HighRisk ); } @@ -174,7 +239,7 @@ mod tests { } #[test] - fn medium_risk_trumps_low_risk() { + fn high_risk_trumps_low_risk() { let labels = labels_with(&[ SafetyLabelType::GROK_SFA, SafetyLabelType::NSFA_LIMITED_INVENTORY, @@ -182,7 +247,7 @@ mod tests { ]); assert_eq!( compute_verdict(&labels, PRE_CUTOFF_ID), - BrandSafetyVerdict::MediumRisk + BrandSafetyVerdict::HighRisk ); } @@ -213,6 +278,146 @@ mod tests { ); } + #[test] + fn v2_defers_to_v1_when_v2_has_not_ruled() { + let v1_safe = labels_with(&[SafetyLabelType::GROK_SFA]); + assert_eq!( + compute_verdict_v2(&v1_safe, PRE_CUTOFF_ID), + compute_verdict(&v1_safe, PRE_CUTOFF_ID) + ); + assert_eq!( + compute_verdict_v2(&v1_safe, PRE_CUTOFF_ID), + BrandSafetyVerdict::Safe + ); + + let v1_nsfa = labels_with(&[SafetyLabelType::GROK_NSFA]); + assert_eq!( + compute_verdict_v2(&v1_nsfa, PRE_CUTOFF_ID), + BrandSafetyVerdict::MediumRisk + ); + let v1_limited = labels_with(&[ + SafetyLabelType::NSFA_LIMITED_INVENTORY, + SafetyLabelType::GROK_NSFA_LIMITED, + ]); + assert_eq!( + compute_verdict_v2(&v1_limited, PRE_CUTOFF_ID), + BrandSafetyVerdict::LowRisk + ); + + assert_eq!( + compute_verdict_v2(&labels_with(&[]), PRE_CUTOFF_ID), + BrandSafetyVerdict::MediumRisk + ); + + let disagreement = labels_with(&[SafetyLabelType::GROK_SFA, SafetyLabelType::GROK_NSFA_V2]); + assert_eq!( + compute_verdict_v2(&disagreement, PRE_CUTOFF_ID), + BrandSafetyVerdict::MediumRisk + ); + + let freed = labels_with(&[ + SafetyLabelType::NSFA_HIGH_PRECISION, + SafetyLabelType::GROK_NSFA, + SafetyLabelType::GROK_SFA_V2, + ]); + assert_eq!( + compute_verdict(&freed, PRE_CUTOFF_ID), + BrandSafetyVerdict::MediumRisk + ); + assert_eq!( + compute_verdict_v2(&freed, PRE_CUTOFF_ID), + BrandSafetyVerdict::Safe + ); + } + + #[test] + fn v2_mirrors_v1_across_tier_matrix() { + fn to_v2(v1_set: &[SafetyLabelType]) -> Vec { + v1_set + .iter() + .filter_map(|l| match *l { + SafetyLabelType::GROK_SFA => Some(SafetyLabelType::GROK_SFA_V2), + SafetyLabelType::GROK_NSFA => Some(SafetyLabelType::GROK_NSFA_V2), + SafetyLabelType::GROK_NSFA_LIMITED => { + Some(SafetyLabelType::GROK_NSFA_LIMITED_V2) + } + SafetyLabelType::NSFA_HIGH_PRECISION + | SafetyLabelType::NSFA_LIMITED_INVENTORY => None, + other => Some(other), + }) + .collect() + } + + let matrix: &[&[SafetyLabelType]] = &[ + &[], + &[SafetyLabelType::GROK_SFA], + &[SafetyLabelType::GROK_SFA, SafetyLabelType::PTOS_REVIEWED], + &[ + SafetyLabelType::NSFA_LIMITED_INVENTORY, + SafetyLabelType::GROK_NSFA_LIMITED, + ], + &[ + SafetyLabelType::NSFA_LIMITED_INVENTORY, + SafetyLabelType::GROK_NSFA_LIMITED, + SafetyLabelType::PTOS_REVIEWED, + ], + &[ + SafetyLabelType::GROK_SFA, + SafetyLabelType::NSFA_HIGH_PRECISION, + SafetyLabelType::GROK_NSFA, + ], + ]; + + for v1_set in matrix { + let v2_set = to_v2(v1_set); + for tweet_id in [PRE_CUTOFF_ID, POST_CUTOFF_ID] { + assert_eq!( + compute_verdict(&labels_with(v1_set), tweet_id), + compute_verdict_v2(&labels_with(&v2_set), tweet_id), + "v1 {v1_set:?} vs v2 {v2_set:?} at tweet_id {tweet_id}" + ); + } + } + + let labels = labels_with(&[ + SafetyLabelType::GROK_SFA_V2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2, + ]); + assert_eq!( + compute_verdict_v2(&labels, PRE_CUTOFF_ID), + BrandSafetyVerdict::MediumRisk + ); + + use std::collections::HashSet; + let medium: HashSet<_> = MEDIUM_RISK_LABELS.iter().copied().collect(); + let medium_v2: HashSet<_> = MEDIUM_RISK_LABELS_V2.iter().copied().collect(); + let expected_medium_v2: HashSet<_> = medium + .iter() + .copied() + .filter(|l| { + *l != SafetyLabelType::GROK_NSFA && *l != SafetyLabelType::NSFA_HIGH_PRECISION + }) + .chain([ + SafetyLabelType::GROK_NSFA_V2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2, + ]) + .collect(); + assert_eq!(medium_v2, expected_medium_v2); + + let low: HashSet<_> = LOW_RISK_LABELS.iter().copied().collect(); + let low_v2: HashSet<_> = LOW_RISK_LABELS_V2.iter().copied().collect(); + let expected_low_v2: HashSet<_> = low + .iter() + .copied() + .filter(|l| { + *l != SafetyLabelType::GROK_NSFA_LIMITED + && *l != SafetyLabelType::NSFA_LIMITED_INVENTORY + }) + .chain([SafetyLabelType::GROK_NSFA_LIMITED_V2]) + .collect(); + assert_eq!(low_v2, expected_low_v2); + } + #[test] fn worst_verdict_ordering() { assert_eq!( diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index b72582d3..a9d747d6 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; pub use xai_candidate_pipeline::component_library::models::PhoenixScores; use xai_home_mixer_proto as pb; +use xai_recsys_proto::SAFETY_BIT_AUTHOR_NSFW; use xai_visibility_filtering::models as vf; #[derive(Clone, Debug, Default, Serialize, Deserialize)] @@ -22,7 +23,11 @@ pub struct PostCandidate { pub score: Option, pub slate_context: Option, #[serde(default)] - pub mpn_parts: Option, + pub served_slate_context: Option, + #[serde(default)] + pub reranker_head_tag: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub backbone_scores: Option, #[serde( serialize_with = "serialize_served_type", deserialize_with = "deserialize_served_type" @@ -32,6 +37,8 @@ pub struct PostCandidate { pub ancestors: Vec, pub tombstone_ancestor_ids: Vec, pub ancestor_users: Vec, + pub ancestor_texts: HashMap, + pub quoted_tweet_text: Option, pub min_video_duration_ms: Option, pub quoted_video_duration_ms: Option, pub author_followers_count: Option, @@ -55,6 +62,8 @@ pub struct PostCandidate { pub repost_count: Option, pub quote_count: Option, pub view_count: Option, + #[serde(default)] + pub view_count_on_home: Option, pub bookmark_count: Option, pub mutual_follow_jaccard: Option, pub is_mutual_follow_author: Option, @@ -62,6 +71,7 @@ pub struct PostCandidate { pub brand_safety_verdict: Option, pub nsfw_author: Option, pub nsfw_author_ads: Option, + pub nsfw_author_phoenix: Option, #[serde(default)] pub safety_labels: Vec, #[serde(default)] @@ -69,22 +79,59 @@ pub struct PostCandidate { pub topic_feedback_topic: Option, pub topic_feedback_topic_id: Option, pub grok_topics: Option>, + pub ai_trend_name: Option, + pub ai_trend_id: Option, } #[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] pub struct SlateContext { pub k: u32, pub pool_rank: u32, pub pool_rank_gap: Option, pub fatigue: f64, pub pre_diversity_score: f64, + pub sid_known: bool, + pub sid_k_l1: u32, + pub sid_k_l2: u32, + pub sid_k_l3: u32, + pub sid_gap_l1: Option, + pub sid_gap_l2: Option, + pub sid_gap_l3: Option, + #[serde(default)] + pub recon_cos_milli: Option, + #[serde(default)] + pub recon_count_above: Option, + #[serde(default)] + pub recon_gap_above: Option, + #[serde(default)] + pub exact_k: Option, + #[serde(default)] + pub exact_gap: Option, } -#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct MpnParts { - pub pos: f64, - pub neg: f64, - pub scalar_multiplier: f64, +impl From for SlateContext { + fn from(c: xai_recsys_proto::SlateContext) -> Self { + Self { + k: c.k, + pool_rank: c.pool_rank, + pool_rank_gap: c.pool_rank_gap, + fatigue: c.fatigue, + pre_diversity_score: c.pre_diversity_score, + sid_known: c.sid_known, + sid_k_l1: c.sid_k1, + sid_k_l2: c.sid_k2, + sid_k_l3: c.sid_k3, + sid_gap_l1: c.sid_gap1, + sid_gap_l2: c.sid_gap2, + sid_gap_l3: c.sid_gap3, + recon_cos_milli: c.recon_cos_milli, + recon_count_above: c.recon_count_above, + recon_gap_above: c.recon_gap_above, + exact_k: c.exact_k, + exact_gap: c.exact_gap, + } + } } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -123,7 +170,7 @@ pub trait CandidateHelpers { fn get_original_tweet_id(&self) -> u64; fn get_original_author_id(&self) -> u64; fn as_tweet_info(&self, is_followed_by_viewer: bool) -> xai_recsys_proto::TweetInfo; - fn as_score_info(&self) -> xai_recsys_proto::ScoreInfo; + fn as_score_info_no_prediction_scores(&self) -> xai_recsys_proto::ScoreInfo; } impl CandidateHelpers for PostCandidate { @@ -148,9 +195,9 @@ impl CandidateHelpers for PostCandidate { self.retweeted_user_id.unwrap_or(self.author_id) } - fn as_score_info(&self) -> xai_recsys_proto::ScoreInfo { + fn as_score_info_no_prediction_scores(&self) -> xai_recsys_proto::ScoreInfo { xai_recsys_proto::ScoreInfo { - prediction_scores: self.phoenix_scores.as_prediction_scores(), + prediction_scores: Default::default(), weighted_score: self.weighted_score, final_score: self.score, slate_context: self.slate_context.map(|c| xai_recsys_proto::SlateContext { @@ -159,8 +206,22 @@ impl CandidateHelpers for PostCandidate { pool_rank_gap: c.pool_rank_gap, fatigue: c.fatigue, pre_diversity_score: c.pre_diversity_score, + sid_known: c.sid_known, + sid_k1: c.sid_k_l1, + sid_k2: c.sid_k_l2, + sid_k3: c.sid_k_l3, + sid_gap1: c.sid_gap_l1, + sid_gap2: c.sid_gap_l2, + sid_gap3: c.sid_gap_l3, + recon_cos_milli: c.recon_cos_milli, + recon_count_above: c.recon_count_above, + recon_gap_above: c.recon_gap_above, + exact_k: c.exact_k, + exact_gap: c.exact_gap, }), reward_rerank_slot_prob: None, + page_decode_slot_prob: None, + reranker_head_tag: self.reranker_head_tag, } } @@ -182,6 +243,13 @@ impl CandidateHelpers for PostCandidate { quoted_author_id: self.quoted_user_id.unwrap_or(0), in_reply_to_tweet_id: self.in_reply_to_tweet_id.unwrap_or(0), is_author_followed_by_user: is_followed_by_viewer, + safety_label_mask: if self.retweeted_user_id.is_none() + && self.nsfw_author_phoenix.unwrap_or(false) + { + SAFETY_BIT_AUTHOR_NSFW + } else { + 0 + }, min_video_duration_ms: self.min_video_duration_ms.map(|ms| ms as u64).unwrap_or(0), fav_count: self.fav_count.unwrap_or(0) as u64, retweet_count: self.repost_count.unwrap_or(0) as u64, @@ -207,7 +275,11 @@ impl CandidateHelpers for PostCandidate { } else { None }, - ..Default::default() + followers: if self.retweeted_user_id.is_none() { + self.author_followers_count.map(|c| c.max(0) as u64) + } else { + None + }, }), semantic_ids: self.semantic_ids.clone().unwrap_or_default(), ..Default::default() diff --git a/home-mixer/models/query.rs b/home-mixer/models/query.rs index d0856c55..de645475 100644 --- a/home-mixer/models/query.rs +++ b/home-mixer/models/query.rs @@ -71,6 +71,7 @@ pub struct ScoredPostsQuery { pub request_time_ms: i64, pub cached_posts: Vec, pub has_cached_posts: bool, + pub return_backbone_scores: bool, pub topic_ids: Vec, pub excluded_topic_ids: Vec, pub exclude_videos: bool, @@ -187,6 +188,7 @@ impl ScoredPostsQuery { request_time_ms: current_time_ms(), cached_posts: vec![], has_cached_posts: false, + return_backbone_scores: false, topic_ids, excluded_topic_ids, exclude_videos, diff --git a/home-mixer/params/config.rs b/home-mixer/params/config.rs index 6048fa6f..0bdbe354 100644 --- a/home-mixer/params/config.rs +++ b/home-mixer/params/config.rs @@ -21,7 +21,7 @@ pub const MAX_JETFUEL_FRAMES_PER_RESPONSE: usize = 8; pub const FOR_YOU_MAX_RESULT_SIZE: usize = RESULT_SIZE + FEED_MODULE_SLOTS + MAX_JETFUEL_FRAMES_PER_RESPONSE; pub const RANKED_FOLLOWING_MAX_RESULT_SIZE: usize = 38; -pub const FOLLOWING_MAX_RESULT_SIZE: usize = 100; +pub const FOLLOWING_MAX_RESULT_SIZE: usize = 110; pub const FOLLOWING_ADS_TOP_K: usize = 9; pub const FOLLOWING_POST_FETCH_SIZE: usize = FOLLOWING_MAX_RESULT_SIZE - FOLLOWING_ADS_TOP_K; pub const FOLLOWING_PIPELINE_RESULT_SIZE: usize = FOLLOWING_MAX_RESULT_SIZE + 2; diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index dc3a12e4..4df2ffa9 100644 --- a/home-mixer/params/param.rs +++ b/home-mixer/params/param.rs @@ -1,4 +1,4 @@ -// mirrored from config feature-switch defaults; last sync 2026-08-12T04:09:22Z +// mirrored from config feature-switch defaults; last sync 2026-09-04T16:22:24Z use xai_feature_switches::param; param!( @@ -80,7 +80,7 @@ param!( ShadowTrafficPhoenixClusterRates, Vec, "rust_home_mixer_shadow_traffic_phoenix_cluster_rates", - vec!["Experiment6Fou:1.5".to_string()] + vec![] ); param!( ShadowTrafficDefaultPercent, @@ -94,6 +94,12 @@ param!( "rust_home_mixer_enable_phoenix_retrieval_stats_experiment_bucket", false ); +param!( + EnablePhoenixScoreStatsExperimentBucket, + bool, + "rust_home_mixer_enable_phoenix_score_stats_experiment_bucket", + false +); param!( PhoenixRetrievalTopicInferenceClusterId, String, @@ -122,13 +128,13 @@ param!( PhoenixAggregationType, String, "rust_home_mixer_phoenix_aggregation_type", - "DENSE_WITH_SHORT_DWELL" + "DENSE_WITH_LONG_DWELL" ); param!( PhoenixRetrievalAggregationType, String, "rust_home_mixer_phoenix_retrieval_aggregation_type", - "DENSE_WITH_SHORT_DWELL" + "DENSE_WITH_LONG_DWELL" ); param!( @@ -243,17 +249,17 @@ param!( "rust_home_mixer_log_slate_context", false ); +param!(RerankerHeadTag, i64, "rust_home_mixer_reranker_head_tag", 0); param!( OonWeightFactor, f64, "rust_home_mixer_oon_weight_factor", 0.75 ); - param!( - EnableMpnScoring, + MultiplierPreOffset, bool, - "rust_home_mixer_enable_mpn_scoring", + "rust_home_mixer_multiplier_pre_offset", false ); @@ -279,6 +285,32 @@ param!( // These weights reflect a combination of how much an action is // valued in ranking and typical propensities of these actions // across the X network (e.g. negative feedback is overall rare). + +// Each weight multiplies the *predicted* probability of that +// action (P(favorite), P(repost), …) or a continuous value e.g. +// watch time -- the weights do not multiply raw engagement counts. +// One common misinterpretation is that you can read these weight +// ratios as count equivalences, e.g. the incorrect statement that +// "one report cancels 468 likes" -- this is incorrect because the +// weights apply to the predicted probabilities rather than raw counts. + +// And the baseline probability of a Report is more than 1000x lower +// than a Like, so it’s weighted more to allow the prediction to affect +// the final ranking at all. + +// Related to the above is a misunderstanding that bad actors engaging +// in mass blocking/reporting will significantly suppress reach. There +// are multiple things inhibiting this: +// 1. It’s predicting your likelihood of the action, not summing up +// raw weights on counts. Also, recommendations are personalized, so +// reports from bad actors will primarily affect recommendations for +// users who are similar to the bad actors, rather than having the same +// effect on the post's ranking to everyone. +// 2. For an account to count in the algorithms recommendation system, +// it must take place on a post served in Home Timeline. Directly +// navigating to a post (i.e., coordinating via groupchat) has no +// ranking impact. And users cannot manufacture a post to show up in +// their Timeline in any consistently reproducible way. param!(FavoriteWeight, f64, "rust_home_mixer_favorite_weight", 0.5); param!(ReplyWeight, f64, "rust_home_mixer_reply_weight", 5.0); param!( @@ -304,7 +336,7 @@ param!( VideoOpenWeight, f64, "rust_home_mixer_video_open_weight", - 0.05 + 0.07 ); param!(ClickWeight, f64, "rust_home_mixer_click_weight", 0.4); param!(OpenLinkWeight, f64, "rust_home_mixer_open_link_weight", 0.2); @@ -314,7 +346,7 @@ param!( "rust_home_mixer_profile_click_weight", 0.0 ); -param!(VqvWeight, f64, "rust_home_mixer_vqv_weight", 0.05); +param!(VqvWeight, f64, "rust_home_mixer_vqv_weight", 0.0); param!(ShareWeight, f64, "rust_home_mixer_share_weight", 2.0); param!( ShareViaDmWeight, @@ -328,7 +360,7 @@ param!( "rust_home_mixer_share_via_copy_link_weight", 20.0 ); -param!(DwellWeight, f64, "rust_home_mixer_dwell_weight", 0.0); +param!(DwellWeight, f64, "rust_home_mixer_dwell_weight", 0.05); param!(QuoteWeight, f64, "rust_home_mixer_quote_weight", 5.0); param!( QuotedClickWeight, @@ -594,16 +626,10 @@ param!( "Experiment3" ); param!( - VMRankerValueModelId, + PhoenixExperimentOverrides, String, - "rust_home_mixer_vm_ranker_value_model_id", - "dpp" -); -param!( - VMRankerSendHeadWeights, - bool, - "rust_home_mixer_vm_ranker_send_head_weights", - false + "rust_home_mixer_phoenix_experiment_overrides", + "" ); param!( VMRankerDppTheta, @@ -661,6 +687,42 @@ param!( "rust_home_mixer_enable_viewer_cold_start_boost", true ); +param!( + EnableColdStartThompsonSampling, + bool, + "rust_home_mixer_enable_cold_start_thompson_sampling", + false +); +param!( + ColdStartBetaAlpha0, + f64, + "rust_home_mixer_cold_start_beta_alpha0", + 0.75 +); +param!( + ColdStartBetaBeta0, + f64, + "rust_home_mixer_cold_start_beta_beta0", + 49.25 +); +param!( + ColdStartTsTopK, + u32, + "rust_home_mixer_cold_start_ts_top_k", + 2 +); +param!( + ColdStartImpressionScale, + f64, + "rust_home_mixer_cold_start_impression_scale", + 1.0 +); +param!( + ColdStartTrackedIds, + String, + "rust_home_mixer_cold_start_tracked_ids", + "" +); param!( EnableCachedPosts, @@ -853,7 +915,14 @@ param!( AdsBlenderType, String, "rust_home_mixer_ads_blender", - "partition_organic_low_risk" + "multi_risk" +); + +param!( + EnableAdsBrandSafetyVerdictV2, + bool, + "rust_home_mixer_ads_bs_v2_exp_enabled", + true ); param!( AdsTimeGapTSec, @@ -951,6 +1020,13 @@ param!( false ); +param!( + EnableAiTrendFeedbackContext, + bool, + "rust_home_mixer_enable_ai_trend_feedback_context", + false +); + param!( EnablePhoenixRequestCacheSideEffect, bool, diff --git a/home-mixer/scored_posts_server.rs b/home-mixer/scored_posts_server.rs index 0bec6ab8..ff84fb1e 100644 --- a/home-mixer/scored_posts_server.rs +++ b/home-mixer/scored_posts_server.rs @@ -117,6 +117,8 @@ fn candidates_to_scored_posts(candidates: &[PostCandidate]) -> Vec { .map(|d| d as f32), topic_feedback_topic: candidate.topic_feedback_topic.clone(), topic_feedback_topic_id: candidate.topic_feedback_topic_id.clone(), + ai_trend_name: candidate.ai_trend_name.clone(), + ai_trend_id: candidate.ai_trend_id.clone(), } }) .collect() @@ -220,6 +222,11 @@ fn safety_label_to_proto(label: SafetyLabelType) -> Option { SafetyLabelType::NSFA_HIGH_RECALL => HM::NsfaHighRecall, SafetyLabelType::GROK_SFA => HM::GrokSfa, SafetyLabelType::MALICIOUS_URL => HM::MaliciousUrl, + SafetyLabelType::GROK_SFA_V2 => HM::GrokSfaV2, + SafetyLabelType::GROK_NSFA_LIMITED_V2 => HM::GrokNsfaLimitedV2, + SafetyLabelType::GROK_NSFA_V2 => HM::GrokNsfaV2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2 => HM::GrokNsfaExpandedV2, + SafetyLabelType::PTOS_REVIEWED => HM::PtosReviewed, _ => return None, }; Some(v.into()) diff --git a/home-mixer/scorers/author_cold_start.rs b/home-mixer/scorers/author_cold_start.rs index 246d26ce..a00202ca 100644 --- a/home-mixer/scorers/author_cold_start.rs +++ b/home-mixer/scorers/author_cold_start.rs @@ -1,13 +1,17 @@ use crate::models::candidate::PostCandidate; use crate::models::query::ScoredPostsQuery; use crate::params::{ - AuthorIsControl, AuthorIsTreatment, ColdStartFollowerCap, ColdStartImpressionThreshold, - ColdStartMaxPostAgeSecs, ColdStartSlotMax, ColdStartSlotMin, EnableViewerColdStart, + AuthorIsControl, AuthorIsTreatment, ColdStartBetaAlpha0, ColdStartBetaBeta0, + ColdStartFollowerCap, ColdStartImpressionScale, ColdStartImpressionThreshold, + ColdStartMaxPostAgeSecs, ColdStartSlotMax, ColdStartSlotMin, ColdStartTrackedIds, + ColdStartTsTopK, EnableColdStartThompsonSampling, EnableViewerColdStart, LowImpressionsMaxPositionRatio, PhoenixMoeCodivertViewerIsControl, PhoenixMoeCodivertViewerIsTreatment, }; use crate::util::author_rules::AuthorRulesEvaluator; use rand::Rng; +use rand_distr::{Beta, Distribution}; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::Duration; use xai_candidate_pipeline::component_library::utils::duration_since_creation_opt; @@ -37,14 +41,52 @@ enum AuthorCorpus { Treatment, } -fn viewer_arm(query: &ScoredPostsQuery) -> ViewerArm { - if query.params.get(PhoenixMoeCodivertViewerIsTreatment) { - return ViewerArm::Treatment; +struct ColdStartParams { + enabled: bool, + tracked_ids: String, + arm: ViewerArm, + follower_cap: i64, + impression_threshold: u64, + max_post_age: Duration, + max_position_ratio: f64, + slot_min: usize, + slot_max: usize, + use_thompson_sampling: bool, + beta_alpha0: f64, + beta_beta0: f64, + impression_scale: f64, + ts_top_k: usize, +} + +impl ColdStartParams { + fn read(query: &ScoredPostsQuery) -> Self { + Self { + enabled: query.params.get(EnableViewerColdStart), + tracked_ids: query.params.get(ColdStartTrackedIds), + arm: viewer_arm(query), + follower_cap: query.params.get(ColdStartFollowerCap), + impression_threshold: query.params.get(ColdStartImpressionThreshold) as u64, + max_post_age: Duration::from_secs(query.params.get(ColdStartMaxPostAgeSecs)), + max_position_ratio: query.params.get(LowImpressionsMaxPositionRatio), + slot_min: query.params.get(ColdStartSlotMin) as usize, + slot_max: query.params.get(ColdStartSlotMax) as usize, + use_thompson_sampling: query.params.get(EnableColdStartThompsonSampling), + beta_alpha0: query.params.get(ColdStartBetaAlpha0), + beta_beta0: query.params.get(ColdStartBetaBeta0), + impression_scale: query.params.get(ColdStartImpressionScale), + ts_top_k: query.params.get(ColdStartTsTopK) as usize, + } } - if query.params.get(PhoenixMoeCodivertViewerIsControl) { - return ViewerArm::Control; +} + +fn viewer_arm(query: &ScoredPostsQuery) -> ViewerArm { + let is_treatment = query.params.get(PhoenixMoeCodivertViewerIsTreatment); + let is_control = query.params.get(PhoenixMoeCodivertViewerIsControl); + match (is_treatment, is_control) { + (true, _) => ViewerArm::Treatment, + (false, true) => ViewerArm::Control, + (false, false) => ViewerArm::Holdout, } - ViewerArm::Holdout } fn positions_among_nonzero(scores: &[f64]) -> (Vec, usize) { @@ -79,6 +121,49 @@ fn record_cold_started_posts(is_moe: bool, viewer_arm: &str, count: u64) { } } +fn parse_tracked_ids(raw: &str) -> HashSet { + raw.split(',') + .filter_map(|s| s.trim().parse().ok()) + .collect() +} + +fn count_tracked_ids( + candidates: &[PostCandidate], + tracked: &HashSet, +) -> HashMap<(u64, i32), u64> { + let mut counts = HashMap::new(); + if tracked.is_empty() { + return counts; + } + for c in candidates { + if !tracked.contains(&c.tweet_id) && !tracked.contains(&c.author_id) { + continue; + } + let source = c.served_type.map(|t| t as i32).unwrap_or(0); + *counts.entry((c.tweet_id, source)).or_insert(0) += 1; + } + counts +} + +fn record_tracked_ids(candidates: &[PostCandidate], raw: &str) { + let tracked = parse_tracked_ids(raw); + if tracked.is_empty() { + return; + } + let Some(receiver) = xai_stats_receiver::global_stats_receiver() else { + return; + }; + for ((id, source), count) in count_tracked_ids(candidates, &tracked) { + let id_str = id.to_string(); + let source_str = source.to_string(); + receiver.incr( + "home_mixer.cold_start_tracked_ids_total", + &[("id", id_str.as_str()), ("source", source_str.as_str())], + count, + ); + } +} + fn is_phoenix_moe(c: &PostCandidate) -> bool { c.served_type == Some(pb::ServedType::ForYouPhoenixRetrievalMoe) } @@ -97,12 +182,12 @@ fn author_corpus( candidates .iter() .map(|c| { - if author_rules.get(c.author_id, AuthorIsTreatment) { - AuthorCorpus::Treatment - } else if author_rules.get(c.author_id, AuthorIsControl) { - AuthorCorpus::Control - } else { - AuthorCorpus::NotBucketed + let is_treatment = author_rules.get(c.author_id, AuthorIsTreatment); + let is_control = author_rules.get(c.author_id, AuthorIsControl); + match (is_treatment, is_control) { + (true, _) => AuthorCorpus::Treatment, + (false, true) => AuthorCorpus::Control, + (false, false) => AuthorCorpus::NotBucketed, } }) .collect() @@ -127,11 +212,11 @@ fn apply_moe_ranking_policy( out } -fn cold_start_target(query: &ScoredPostsQuery, scores: &[f64]) -> Option { +fn cold_start_target(params: &ColdStartParams, scores: &[f64]) -> Option { let mut ranked = scores.to_vec(); ranked.sort_by(|a, b| b.total_cmp(a)); - let hi = (query.params.get(ColdStartSlotMax) as usize).min(ranked.len()); - let lo = (query.params.get(ColdStartSlotMin) as usize).min(hi); + let hi = params.slot_max.min(ranked.len()); + let lo = params.slot_min.min(hi); if lo >= hi { return None; } @@ -142,7 +227,7 @@ fn cold_start_corpus_eligible(arm: ViewerArm, c: &PostCandidate, corpus: AuthorC match arm { ViewerArm::Holdout => !is_phoenix_moe(c), ViewerArm::Control => corpus == AuthorCorpus::Control && !is_phoenix_moe(c), - ViewerArm::Treatment => corpus == AuthorCorpus::Treatment, + ViewerArm::Treatment => corpus == AuthorCorpus::Treatment && is_phoenix_moe(c), } } @@ -153,35 +238,96 @@ fn cold_start_freshness_eligible(arm: ViewerArm, c: &PostCandidate, max_age: Dur duration_since_creation_opt(c.tweet_id).is_some_and(|age| age <= max_age) } +fn pick_by_score(eligible: &[usize], scores: &[f64]) -> Option { + eligible + .iter() + .copied() + .max_by(|&i, &j| scores[i].total_cmp(&scores[j]).then(i.cmp(&j))) +} + +fn sample_reward( + candidate: &PostCandidate, + alpha0: f64, + beta0: f64, + scale: f64, + rng: &mut R, +) -> f64 { + let n = scale * candidate.view_count_on_home.unwrap_or(0) as f64; + let x = (candidate.fav_count.unwrap_or(0).max(0) as f64).min(n); + let alpha = alpha0 + x; + let beta = beta0 + (n - x).max(0.0); + Beta::new(alpha, beta).map(|d| d.sample(rng)).unwrap_or(0.5) +} + +fn pick_thompson( + eligible: &[usize], + candidates: &[PostCandidate], + scores: &[f64], + alpha0: f64, + beta0: f64, + scale: f64, + top_k: usize, + rng: &mut R, +) -> Option { + if eligible.is_empty() { + return None; + } + if top_k == 0 || alpha0 <= 0.0 || beta0 <= 0.0 { + return pick_by_score(eligible, scores); + } + let mut sampled: Vec<(usize, f64)> = eligible + .iter() + .map(|&i| (i, sample_reward(&candidates[i], alpha0, beta0, scale, rng))) + .collect(); + sampled.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); + let k = top_k.min(sampled.len()); + sampled[..k] + .iter() + .map(|(i, _)| *i) + .max_by(|&i, &j| scores[i].total_cmp(&scores[j]).then(i.cmp(&j))) +} + fn apply_cold_start( - query: &ScoredPostsQuery, + params: &ColdStartParams, candidates: &[PostCandidate], scores: &[f64], corpus: &[AuthorCorpus], - arm: ViewerArm, target: f64, ) -> Vec { - let follower_cap = query.params.get(ColdStartFollowerCap); - let threshold = query.params.get(ColdStartImpressionThreshold) as u64; - let max_post_age = Duration::from_secs(query.params.get(ColdStartMaxPostAgeSecs)); + let arm = params.arm; let (positions, nonzero) = positions_among_nonzero(scores); - let max_cold_start_slot = - (query.params.get(LowImpressionsMaxPositionRatio) * nonzero as f64) as usize; + let max_cold_start_slot = (params.max_position_ratio * nonzero as f64) as usize; - let best = candidates + let eligible: Vec = candidates .iter() .enumerate() .filter(|(i, c)| { - cold_start_base_eligible(c, follower_cap) + cold_start_base_eligible(c, params.follower_cap) && cold_start_corpus_eligible(arm, c, corpus[*i]) - && cold_start_freshness_eligible(arm, c, max_post_age) + && cold_start_freshness_eligible(arm, c, params.max_post_age) && positions[*i] < max_cold_start_slot - && c.view_count.is_some_and(|imp| imp < threshold) + && c.view_count_on_home + .is_some_and(|imp| imp < params.impression_threshold) }) .map(|(i, _)| i) - .max_by(|&i, &j| scores[i].total_cmp(&scores[j])); + .collect(); + + let best_idx = if params.use_thompson_sampling { + pick_thompson( + &eligible, + candidates, + scores, + params.beta_alpha0, + params.beta_beta0, + params.impression_scale, + params.ts_top_k, + &mut rand::rng(), + ) + } else { + pick_by_score(&eligible, scores) + }; - let Some(best_idx) = best else { + let Some(best_idx) = best_idx else { return scores.to_vec(); }; @@ -203,21 +349,24 @@ impl AuthorColdStart { candidates: &[PostCandidate], scores: &[f64], ) -> Vec { - if !query.params.get(EnableViewerColdStart) { - return scores.to_vec(); - } - - let arm = viewer_arm(query); - let corpus = match arm { + let params = ColdStartParams::read(query); + record_tracked_ids(candidates, ¶ms.tracked_ids); + let corpus = match params.arm { ViewerArm::Holdout => vec![AuthorCorpus::NotBucketed; candidates.len()], ViewerArm::Control | ViewerArm::Treatment => { author_corpus(&self.author_rules, candidates) } }; + if !params.enabled { + return scores.to_vec(); + } + + let arm = params.arm; + let mut effective = apply_moe_ranking_policy(arm, candidates, &corpus, scores); - if let Some(target) = cold_start_target(query, scores) { - effective = apply_cold_start(query, candidates, &effective, &corpus, arm, target); + if let Some(target) = cold_start_target(¶ms, scores) { + effective = apply_cold_start(¶ms, candidates, &effective, &corpus, target); } effective } @@ -226,10 +375,12 @@ impl AuthorColdStart { #[cfg(test)] mod tests { use super::*; + use rand::SeedableRng; use xai_candidate_pipeline::component_library::utils::current_time_to_id; use xai_feature_switches::{ BucketMembership, ExperimentBucket, ExperimentBucketsChooser, FeatureSwitches, - NullBucketImpressor, Recipient, + MockExperimentBucketsChooser, NullBucketImpressor, Recipient, RecipientBuilder, + SpyingBucketImpressor, }; #[derive(Debug)] @@ -264,6 +415,14 @@ mod tests { } fn cold_start_with_arms(treatment: Vec, control: Vec) -> AuthorColdStart { + cold_start_with_author_impressor(treatment, control, Arc::new(NullBucketImpressor::new())) + } + + fn cold_start_with_author_impressor( + treatment: Vec, + control: Vec, + impressor: Arc, + ) -> AuthorColdStart { let yaml = r#" rust_home_mixer: description: "x" @@ -281,10 +440,12 @@ rust_home_mixer: [moe_exp author_bucket_membership treatment] values: rust_home_mixer_author_is_treatment: true + rust_home_mixer_author_is_control: false - description: "moe control corpus" query: > [moe_exp author_bucket_membership control] values: + rust_home_mixer_author_is_treatment: false rust_home_mixer_author_is_control: true "#; let features = xai_feature_switches::load_yaml_string(yaml).unwrap(); @@ -292,10 +453,9 @@ rust_home_mixer: FeatureSwitches::with_options( features, Arc::new(ArmChooser { treatment, control }), - Arc::new(NullBucketImpressor::new()), + impressor, None, false, - None, ) .unwrap(), ); @@ -312,24 +472,43 @@ rust_home_mixer: current_time_to_id() as u64 - ((age.as_millis() as u64) << 22) } - fn cold_start_candidate(author_id: u64, age: Duration, view_count: u64) -> PostCandidate { + fn cold_start_candidate( + author_id: u64, + age: Duration, + view_count_on_home: u64, + ) -> PostCandidate { + cold_start_candidate_with_favs(author_id, age, view_count_on_home, 0) + } + + fn cold_start_candidate_with_favs( + author_id: u64, + age: Duration, + view_count_on_home: u64, + fav_count: i64, + ) -> PostCandidate { PostCandidate { author_id, tweet_id: tweet_id_with_age(age), author_followers_count: Some(100), - view_count: Some(view_count), + view_count_on_home: Some(view_count_on_home), + fav_count: Some(fav_count), ..Default::default() } } - fn moe_candidate(author_id: u64, age: Duration, view_count: u64) -> PostCandidate { + fn moe_candidate(author_id: u64, age: Duration, view_count_on_home: u64) -> PostCandidate { + moe_candidate_with_favs(author_id, age, view_count_on_home, 0) + } + + fn moe_candidate_with_favs( + author_id: u64, + age: Duration, + view_count_on_home: u64, + fav_count: i64, + ) -> PostCandidate { PostCandidate { - author_id, - tweet_id: tweet_id_with_age(age), - author_followers_count: Some(100), - view_count: Some(view_count), served_type: Some(pb::ServedType::ForYouPhoenixRetrievalMoe), - ..Default::default() + ..cold_start_candidate_with_favs(author_id, age, view_count_on_home, fav_count) } } @@ -360,6 +539,32 @@ rust_home_mixer: "rust_home_mixer_cold_start_max_post_age_secs".to_string(), "7200", ); + results.override_fs( + "rust_home_mixer_enable_cold_start_thompson_sampling".to_string(), + "false", + ); + results.override_fs("rust_home_mixer_cold_start_beta_alpha0".to_string(), "0.75"); + results.override_fs("rust_home_mixer_cold_start_beta_beta0".to_string(), "49.25"); + results.override_fs("rust_home_mixer_cold_start_ts_top_k".to_string(), "5"); + results.override_fs( + "rust_home_mixer_cold_start_impression_scale".to_string(), + "1.0", + ); + query.params = results.into(); + query + } + + fn ts_query(viewer_treatment: bool, top_k: u32) -> ScoredPostsQuery { + let mut query = codivert_query(!viewer_treatment, viewer_treatment); + let mut results = query.params.0.expect("params set"); + results.override_fs( + "rust_home_mixer_enable_cold_start_thompson_sampling".to_string(), + "true", + ); + results.override_fs( + "rust_home_mixer_cold_start_ts_top_k".to_string(), + &top_k.to_string(), + ); query.params = results.into(); query } @@ -423,7 +628,7 @@ rust_home_mixer: } #[test] - fn treatment_viewer_keeps_treatment_moe_and_cold_starts_treatment_corpus() { + fn treatment_viewer_cold_starts_treatment_moe_only() { let author_cold_start = cold_start_with_arms(vec![1, 2], vec![3]); let candidates = vec![ moe_candidate(1, minutes(10), 3), @@ -433,19 +638,19 @@ rust_home_mixer: let result = author_cold_start.apply( &codivert_query(false, true), &candidates, - &[80.0, 50.0, 90.0], + &[50.0, 80.0, 90.0], ); assert_eq!(result[0], 90.0); - assert_eq!(result[1], 50.0); + assert_eq!(result[1], 80.0); assert_eq!(result[2], 90.0); } #[test] fn treatment_skips_post_older_than_max_post_age() { - let author_cold_start = cold_start_with_arms(vec![1], vec![]); + let author_cold_start = cold_start_with_arms(vec![1, 2], vec![]); let candidates = vec![ - cold_start_candidate(1, minutes(180), 3), - cold_start_candidate(2, minutes(30), 1000), + moe_candidate(1, minutes(180), 3), + moe_candidate(2, minutes(30), 1000), ]; let result = author_cold_start.apply( &query_with_max_post_age(true, 7200), @@ -491,4 +696,220 @@ rust_home_mixer: let result = author_cold_start.apply(&query, &candidates, &[10.0, 40.0, 30.0, 20.0]); assert_eq!(result, vec![10.0, 40.0, 30.0, 20.0]); } + + #[test] + fn ts_top_k_zero_falls_back_to_argmax_score() { + let author_cold_start = cold_start_with_arms(vec![1, 2], vec![]); + let candidates = vec![ + moe_candidate_with_favs(1, minutes(10), 0, 0), + moe_candidate_with_favs(2, minutes(20), 3, 0), + ]; + let result = author_cold_start.apply(&ts_query(true, 0), &candidates, &[10.0, 90.0]); + assert_eq!(result, vec![10.0, 90.0]); + } + + #[test] + fn treatment_ts_among_top_k_picks_highest_score() { + let author_cold_start = cold_start_with_arms(vec![1, 2], vec![]); + let candidates = vec![ + moe_candidate_with_favs(1, minutes(10), 0, 0), + moe_candidate_with_favs(2, minutes(20), 0, 0), + ]; + let result = author_cold_start.apply(&ts_query(true, 10), &candidates, &[10.0, 90.0]); + assert_eq!(result, vec![10.0, 90.0]); + } + + #[test] + fn control_ts_among_top_k_picks_highest_score() { + let author_cold_start = cold_start_with_arms(vec![], vec![1, 2]); + let candidates = vec![ + cold_start_candidate_with_favs(1, minutes(10), 0, 0), + cold_start_candidate_with_favs(2, minutes(20), 0, 0), + ]; + let result = author_cold_start.apply(&ts_query(false, 10), &candidates, &[10.0, 90.0]); + assert_eq!(result, vec![10.0, 90.0]); + } + + const CODIVERT_EXPERIMENT: &str = "moe_codivert_viewer"; + + const CODIVERT_YAML: &str = r#" +rust_home_mixer: + description: "x" + owner: "t@example.com" + parameters: + rust_home_mixer_phoenix_moe_codivert_viewer_is_control: + type: boolean + default: false + rust_home_mixer_phoenix_moe_codivert_viewer_is_treatment: + type: boolean + default: false + rust_home_mixer_enable_viewer_cold_start_boost: + type: boolean + default: true + rules: + - description: "co-divert viewer treatment" + query: > + [moe_codivert_viewer bucket_membership treatment] + values: + rust_home_mixer_phoenix_moe_codivert_viewer_is_treatment: true + rust_home_mixer_phoenix_moe_codivert_viewer_is_control: false + rust_home_mixer_enable_viewer_cold_start_boost: true + - description: "co-divert viewer control" + query: > + [moe_codivert_viewer bucket_membership control] + values: + rust_home_mixer_phoenix_moe_codivert_viewer_is_treatment: false + rust_home_mixer_phoenix_moe_codivert_viewer_is_control: true + rust_home_mixer_enable_viewer_cold_start_boost: false +"#; + + fn codivert_viewer(arm: &str) -> (ScoredPostsQuery, Arc) { + let mut buckets = BucketMembership::new(); + buckets.add(ExperimentBucket::new(CODIVERT_EXPERIMENT, arm).with_version(1)); + let impressor = Arc::new(SpyingBucketImpressor::new()); + let fs = FeatureSwitches::with_options( + xai_feature_switches::load_yaml_string(CODIVERT_YAML).unwrap(), + Arc::new(MockExperimentBucketsChooser::with_buckets(buckets)), + impressor.clone(), + None, + false, + ) + .unwrap(); + + let query = ScoredPostsQuery { + params: fs + .match_recipient(&RecipientBuilder::new().user_id(7).build()) + .into(), + ..Default::default() + }; + (query, impressor) + } + + #[test] + fn both_arms_log_the_same_codivert_impressions() { + let author_cold_start = cold_start_with_arms(vec![2], vec![1]); + let candidates = vec![ + cold_start_candidate(1, minutes(10), 3), + moe_candidate(2, minutes(20), 3), + ]; + + let (control_query, control_impressor) = codivert_viewer("control"); + author_cold_start.apply(&control_query, &candidates, &[5.0, 40.0]); + + let (treatment_query, treatment_impressor) = codivert_viewer("treatment"); + author_cold_start.apply(&treatment_query, &candidates, &[5.0, 40.0]); + + assert_eq!(control_impressor.impression_count(CODIVERT_EXPERIMENT), 3); + assert_eq!(treatment_impressor.impression_count(CODIVERT_EXPERIMENT), 3); + } + + fn viewer_query(control: bool, treatment: bool, cold_start: bool) -> ScoredPostsQuery { + let mut query = codivert_query(control, treatment); + let mut results = query.params.0.expect("params set"); + results.override_fs( + "rust_home_mixer_enable_viewer_cold_start_boost".to_string(), + if cold_start { "true" } else { "false" }, + ); + query.params = results.into(); + query + } + + #[test] + fn both_viewer_arms_log_the_same_author_impressions() { + let candidates = vec![ + cold_start_candidate(1, minutes(10), 3), + moe_candidate(2, minutes(20), 3), + ]; + let scores = [5.0, 40.0]; + + let cases = [ + viewer_query(true, false, false), + viewer_query(false, true, true), + ]; + for query in cases { + let impressor = Arc::new(SpyingBucketImpressor::new()); + let author_cold_start = + cold_start_with_author_impressor(vec![2], vec![1], impressor.clone()); + author_cold_start.apply(&query, &candidates, &scores); + assert_eq!(impressor.impression_count("moe_exp"), 2); + } + } + + #[test] + fn parse_tracked_ids_splits_and_skips_junk() { + let ids = parse_tracked_ids("10, 20, x, 10"); + assert_eq!(ids, HashSet::from([10, 20])); + assert!(parse_tracked_ids("").is_empty()); + assert!(parse_tracked_ids(" , , ").is_empty()); + } + + #[test] + fn count_tracked_ids_matches_tweet_or_author() { + let tracked = parse_tracked_ids("10,20"); + let candidates = vec![ + PostCandidate { + author_id: 10, + tweet_id: 1, + served_type: Some(pb::ServedType::ForYouPhoenixRetrieval), + ..Default::default() + }, + PostCandidate { + author_id: 10, + tweet_id: 2, + served_type: Some(pb::ServedType::ForYouPhoenixRetrievalMoe), + ..Default::default() + }, + PostCandidate { + author_id: 99, + tweet_id: 20, + served_type: Some(pb::ServedType::ForYouInNetwork), + ..Default::default() + }, + PostCandidate { + author_id: 30, + tweet_id: 3, + ..Default::default() + }, + ]; + let counts = count_tracked_ids(&candidates, &tracked); + assert_eq!( + counts.get(&(1, pb::ServedType::ForYouPhoenixRetrieval as i32)), + Some(&1) + ); + assert_eq!( + counts.get(&(2, pb::ServedType::ForYouPhoenixRetrievalMoe as i32)), + Some(&1) + ); + assert_eq!( + counts.get(&(20, pb::ServedType::ForYouInNetwork as i32)), + Some(&1) + ); + assert_eq!( + counts.get(&(10, pb::ServedType::ForYouPhoenixRetrieval as i32)), + None + ); + assert_eq!(counts.get(&(3, 0)), None); + } + + #[test] + fn pick_thompson_top_k_one_prefers_uncertain_over_peaked_zero() { + let candidates = vec![ + cold_start_candidate_with_favs(1, minutes(10), 10_000, 0), + cold_start_candidate_with_favs(2, minutes(20), 0, 0), + ]; + let scores = [100.0, 10.0]; + let eligible = vec![0, 1]; + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let picked = pick_thompson( + &eligible, + &candidates, + &scores, + 0.75, + 49.25, + 1.0, + 1, + &mut rng, + ); + assert_eq!(picked, Some(1)); + } } diff --git a/home-mixer/scorers/phoenix_scorer.rs b/home-mixer/scorers/phoenix_scorer.rs index 4ff7b1c8..e7c5aa51 100644 --- a/home-mixer/scorers/phoenix_scorer.rs +++ b/home-mixer/scorers/phoenix_scorer.rs @@ -3,7 +3,7 @@ use crate::models::candidate::PostCandidate; use crate::models::query::ScoredPostsQuery; use crate::params::{ PhoenixInferenceClusterId, PhoenixRankerNewUserHistoryThreshold, - PhoenixRankerNewUserInferenceClusterId, + PhoenixRankerNewUserInferenceClusterId, RerankerHeadTag, }; use crate::util::egress::PredictionDispatch; use crate::util::phoenix_request::build_prediction_request; @@ -107,8 +107,13 @@ impl Scorer for PhoenixScorer { .iter() .map(|c| PostCandidate { phoenix_scores: predictions.candidate_scores(&c.get_original_tweet_id()), + backbone_scores: predictions.candidate_backbone_scores(&c.get_original_tweet_id()), + served_slate_context: predictions + .candidate_slate_context(&c.get_original_tweet_id()) + .map(Into::into), prediction_request_id: Some(query.prediction_id), last_scored_at_ms, + reranker_head_tag: Some(query.params.get(RerankerHeadTag) as u32), ..Default::default() }) .map(Ok) @@ -117,7 +122,10 @@ impl Scorer for PhoenixScorer { fn update(&self, candidate: &mut PostCandidate, scored: PostCandidate) { candidate.phoenix_scores = scored.phoenix_scores; + candidate.backbone_scores = scored.backbone_scores; + candidate.served_slate_context = scored.served_slate_context; candidate.prediction_request_id = scored.prediction_request_id; candidate.last_scored_at_ms = scored.last_scored_at_ms; + candidate.reranker_head_tag = scored.reranker_head_tag; } } diff --git a/home-mixer/scorers/ranking_scorer.rs b/home-mixer/scorers/ranking_scorer.rs index 6dfcb7aa..facbe8d2 100644 --- a/home-mixer/scorers/ranking_scorer.rs +++ b/home-mixer/scorers/ranking_scorer.rs @@ -1,4 +1,4 @@ -use crate::models::candidate::{MpnParts, PhoenixScores, PostCandidate, SlateContext}; +use crate::models::candidate::{PhoenixScores, PostCandidate, SlateContext}; use crate::models::query::ScoredPostsQuery; use crate::params::*; use crate::scorers::author_cold_start::AuthorColdStart; @@ -221,88 +221,6 @@ impl ScoringWeights { self.dwell } - pub(crate) fn effective_head_weights( - &self, - query: &ScoredPostsQuery, - candidate: &PostCandidate, - ) -> xai_vm_ranker_proto::HeadWeights { - let scores = &candidate.phoenix_scores; - let vqv = crate::util::candidates_util::vqv_weight( - query, - candidate, - self.min_video_duration_ms, - self.vqv, - ); - let dwell_time = match scores.post_unexplored_score { - Some(post_unexplored) - if self.enable_multiplicative_post_unexplored - && self.post_unexplored_active_for(candidate) => - { - self.cont_dwell_time - * (1.0 + post_unexplored * self.multiplicative_post_unexplored_alpha) - } - _ => self.cont_dwell_time, - }; - let click_dwell_time = if self.enable_click_dwell_low_fav_rate_penalty { - match (scores.click_dwell_time, scores.favorite_score) { - (Some(_), Some(fav)) => { - let baseline = self - .click_dwell_low_fav_rate_penalty_baseline - .max(f64::EPSILON); - let multiplier = (fav / baseline) - .powf(self.click_dwell_low_fav_rate_penalty_alpha) - .max(self.click_dwell_low_fav_rate_penalty_floor) - .min(self.click_dwell_low_fav_rate_penalty_cap); - self.cont_click_dwell_time * multiplier - } - _ => self.cont_click_dwell_time, - } - } else { - self.cont_click_dwell_time - }; - let quoted_vqv = crate::util::candidates_util::quoted_vqv_weight( - candidate, - self.min_video_duration_ms, - self.quoted_vqv, - self.enable_quoted_vqv_duration_check, - ); - let post_unexplored = if !self.enable_multiplicative_post_unexplored - && self.post_unexplored_active_for(candidate) - { - self.post_unexplored - } else { - 0.0 - }; - xai_vm_ranker_proto::HeadWeights { - favorite: Some(self.favorite), - reply: Some(self.reply_weight_for(candidate)), - retweet: Some(self.retweet), - photo_expand: Some(self.photo_expand), - click: Some(self.click), - profile_click: Some(self.profile_click), - vqv: Some(vqv), - share: Some(self.share), - share_via_dm: Some(self.share_via_dm), - share_via_copy_link: Some(self.share_via_copy_link), - dwell: Some(self.dwell_weight_for(candidate)), - quote: Some(self.quote), - quoted_click: Some(self.quoted_click), - follow_author: Some(self.follow_author), - not_interested: Some(self.not_interested), - block_author: Some(self.block_author), - mute_author: Some(self.mute_author), - report: Some(self.report), - dwell_time: Some(dwell_time), - click_dwell_time: Some(click_dwell_time), - not_dwelled: Some(self.not_dwelled), - video_open: Some(self.video_open), - open_link: Some(self.open_link), - quoted_vqv: Some(quoted_vqv), - post_unexplored: Some(post_unexplored), - active_secs_5m_residual_norm: Some(self.cont_active_secs_5m_residual_norm), - } - } - pub(crate) fn applied_weights_map(&self) -> HashMap { HashMap::from( [ @@ -415,6 +333,35 @@ pub struct RankingScorer { } impl RankingScorer { + // These weights reflect a combination of how much an action is + // valued in ranking and typical propensities of these actions + // across the X network (e.g. negative feedback is overall rare). + + // Each weight multiplies the *predicted* probability of that + // action (P(favorite), P(repost), …) or a continuous value e.g. + // watch time -- the weights do not multiply raw engagement counts. + // One common misinterpretation is that you can read these weight + // ratios as count equivalences, e.g. the incorrect statement that + // "one report cancels 468 likes" -- this is incorrect because the + // weights apply to the predicted probabilities rather than raw counts. + + // And the baseline probability of a Report is more than 1000x lower + // than a Like, so it’s weighted more to allow the prediction to affect + // the final ranking at all. + + // Related to the above is a misunderstanding that bad actors engaging + // in mass blocking/reporting will significantly suppress reach. There + // are multiple things inhibiting this: + // 1. It’s predicting your likelihood of the action, not summing up + // raw weights on counts. Also, recommendations are personalized, so + // reports from bad actors will primarily affect recommendations for + // users who are similar to the bad actors, rather than having the same + // effect on the post's ranking to everyone. + // 2. For an account to count in the algorithms recommendation system, + // it must take place on a post served in Home Timeline. Directly + // navigating to a post (i.e., coordinating via groupchat) has no + // ranking impact. And users cannot manufacture a post to show up in + // their Timeline in any consistently reproducible way. fn apply(score: Option, weight: f64) -> f64 { score.unwrap_or(0.0) * weight } @@ -615,10 +562,7 @@ impl RankingScorer { (1.0 - floor) * decay_factor.powf(exponent) + floor } - fn compute_slate_contexts( - candidates: &[PostCandidate], - pre_diversity_scores: &[f64], - ) -> Vec { + fn author_pool_counts(candidates: &[PostCandidate], pre_diversity_scores: &[f64]) -> Vec { let mut indexed: Vec<(usize, f64)> = pre_diversity_scores .iter() .enumerate() @@ -626,51 +570,42 @@ impl RankingScorer { .collect(); indexed.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(Ordering::Equal)); - let mut contexts = vec![SlateContext::default(); candidates.len()]; + let mut counts = vec![0u32; candidates.len()]; let mut author_counts: FxHashMap = FxHashMap::default(); - let mut last_author_rank: FxHashMap = FxHashMap::default(); - for (rank, (idx, score)) in indexed.into_iter().enumerate() { - let rank = rank as u32; + for (idx, _) in indexed { let author_id = candidates[idx].author_id; let k = author_counts.get(&author_id).copied().unwrap_or(0); - let rank_gap = last_author_rank.get(&author_id).map(|last| rank - last); - contexts[idx] = SlateContext { - k, - pool_rank: rank, - pool_rank_gap: rank_gap, - fatigue: 0.0, - pre_diversity_score: score, - }; + counts[idx] = k; author_counts.insert(author_id, k + 1); - last_author_rank.insert(author_id, rank); } + counts + } - contexts + fn served_slate_contexts(candidates: &[PostCandidate]) -> Option> { + candidates.iter().map(|c| c.served_slate_context).collect() } fn stored_slate_contexts(candidates: &[PostCandidate]) -> Option> { candidates.iter().map(|c| c.slate_context).collect() } - fn author_diversity_multipliers( - query: &ScoredPostsQuery, - contexts: &[SlateContext], - ) -> Vec { + fn author_diversity_multipliers(query: &ScoredPostsQuery, counts: &[u32]) -> Vec { let decay_factor = query.params.get(AuthorDiversityDecay); let floor = query.params.get(AuthorDiversityFloor); - contexts + counts .iter() - .map(|context| Self::diversity_multiplier(decay_factor, floor, f64::from(context.k))) + .map(|&k| Self::diversity_multiplier(decay_factor, floor, f64::from(k))) .collect() } fn apply_author_diversity( query: &ScoredPostsQuery, - contexts: &[SlateContext], + candidates: &[PostCandidate], pre_diversity_scores: &[f64], ) -> Vec { - let multipliers = Self::author_diversity_multipliers(query, contexts); + let counts = Self::author_pool_counts(candidates, pre_diversity_scores); + let multipliers = Self::author_diversity_multipliers(query, &counts); pre_diversity_scores .iter() .zip(multipliers) @@ -738,8 +673,6 @@ impl Scorer for RankingScorer { .collect() }; - let mpn_scoring = query.params.get(EnableMpnScoring) && !use_dwell_regret; - let effective_oon = Self::effective_oon_weight(query); let deboost_in_network_replies_retweets = query .params @@ -753,52 +686,35 @@ impl Scorer for RankingScorer { None => false, }; - if mpn_scoring { - let persisted_contexts: Option> = if query.has_cached_posts { - Self::stored_slate_contexts(candidates) - } else { - Some(Self::compute_slate_contexts(candidates, &weighted_scores)) - }; + let persisted_contexts: Option> = Self::served_slate_contexts(candidates) + .or_else(|| { + query + .has_cached_posts + .then(|| Self::stored_slate_contexts(candidates)) + .flatten() + }); + if !use_dwell_regret && query.params.get(MultiplierPreOffset) { let diversity_multipliers: Vec = if enable_author_diversity { - let recomputed_contexts; - let scoring_contexts: &[SlateContext] = match &persisted_contexts { - Some(contexts) if !query.has_cached_posts => contexts, - _ => { - recomputed_contexts = - Self::compute_slate_contexts(candidates, &weighted_scores); - &recomputed_contexts - } - }; - Self::author_diversity_multipliers(query, scoring_contexts) + let counts = Self::author_pool_counts(candidates, &weighted_scores); + Self::author_diversity_multipliers(query, &counts) } else { vec![1.0; candidates.len()] }; - - let scalar_multipliers: Vec = candidates + let scores: Vec = weighted_parts .iter() .enumerate() - .map(|(i, c)| { + .map(|(i, &(pos, neg))| { let mut m = diversity_multipliers[i]; - if oon_applies(c) { + if oon_applies(&candidates[i]) { m *= effective_oon; } - m - }) - .collect(); - - let mpn_scores: Vec = weighted_parts - .iter() - .zip(&scalar_multipliers) - .map(|(&(pos, neg), &m)| { let net = pos - neg; let scaled = if net >= 0.0 { m * net } else { net }; Self::offset_score(scaled, &weights) }) .collect(); - - let final_scores = self.author_cold_start.apply(query, candidates, &mpn_scores); - + let final_scores = self.author_cold_start.apply(query, candidates, &scores); return weighted_scores .iter() .zip(final_scores) @@ -808,11 +724,6 @@ impl Scorer for RankingScorer { weighted_score: Some(weighted), score: Some(score), slate_context: persisted_contexts.as_ref().map(|contexts| contexts[i]), - mpn_parts: Some(MpnParts { - pos: weighted_parts[i].0, - neg: weighted_parts[i].1, - scalar_multiplier: scalar_multipliers[i], - }), ..Default::default() }) }) @@ -823,23 +734,8 @@ impl Scorer for RankingScorer { .author_cold_start .apply(query, candidates, &weighted_scores); - let persisted_contexts: Option> = if query.has_cached_posts { - Self::stored_slate_contexts(candidates) - } else { - Some(Self::compute_slate_contexts(candidates, &adjusted_scores)) - }; - let diversity_adjusted = if enable_author_diversity { - let recomputed_contexts; - let scoring_contexts: &[SlateContext] = match &persisted_contexts { - Some(contexts) if !query.has_cached_posts => contexts, - _ => { - recomputed_contexts = - Self::compute_slate_contexts(candidates, &adjusted_scores); - &recomputed_contexts - } - }; - Self::apply_author_diversity(query, scoring_contexts, &adjusted_scores) + Self::apply_author_diversity(query, candidates, &adjusted_scores) } else { adjusted_scores.clone() }; @@ -876,7 +772,6 @@ impl Scorer for RankingScorer { candidate.weighted_score = scored.weighted_score; candidate.score = scored.score; candidate.slate_context = scored.slate_context; - candidate.mpn_parts = scored.mpn_parts; } } @@ -949,7 +844,6 @@ mod tests { ("rust_home_mixer_author_diversity_decay", "0.5"), ("rust_home_mixer_author_diversity_floor", "0.25"), ("rust_home_mixer_value_model_mode", "weighted"), - ("rust_home_mixer_enable_mpn_scoring", "false"), ]); let scored = scorer.score(&query, &candidates).await; @@ -974,6 +868,7 @@ mod tests { pool_rank_gap: Some(3), fatigue: 0.0, pre_diversity_score: 0.5, + ..Default::default() }; let stored_repeat = PostCandidate { slate_context: Some(stored_context), @@ -1013,7 +908,6 @@ mod tests { ("rust_home_mixer_author_diversity_decay", "0.5"), ("rust_home_mixer_author_diversity_floor", "0.25"), ("rust_home_mixer_value_model_mode", "weighted"), - ("rust_home_mixer_enable_mpn_scoring", "false"), ]); query.has_cached_posts = true; @@ -1034,7 +928,6 @@ mod tests { let query = query_with_flags(&[ ("rust_home_mixer_oon_weight_factor", "0.75"), ("rust_home_mixer_value_model_mode", "weighted"), - ("rust_home_mixer_enable_mpn_scoring", "false"), ]); let scored = scorer.score(&query, &candidates).await; @@ -1348,7 +1241,6 @@ mod tests { ), ("rust_home_mixer_oon_weight_factor", "0.75"), ("rust_home_mixer_value_model_mode", "weighted"), - ("rust_home_mixer_enable_mpn_scoring", "false"), ]); let scored = scorer.score(&query, &candidates).await; diff --git a/home-mixer/scorers/vm_ranker.rs b/home-mixer/scorers/vm_ranker.rs index 624264b0..bc40c229 100644 --- a/home-mixer/scorers/vm_ranker.rs +++ b/home-mixer/scorers/vm_ranker.rs @@ -2,21 +2,17 @@ use crate::clients::vm_ranker_client::{VMRankerClient, VMRankerCluster}; use crate::models::candidate::PostCandidate; use crate::models::query::ScoredPostsQuery; use crate::params::*; -use crate::scorers::author_cold_start::AuthorColdStart; -use crate::util::candidates_util; use rustc_hash::FxHashMap; use std::sync::Arc; use tonic::async_trait; use xai_candidate_pipeline::scorer::Scorer; -use xai_vm_ranker_proto::{DppParams, PhoenixScores, RankCandidate, RankRequest, SlateContext}; +use xai_vm_ranker_proto::{DppParams, RankCandidate, RankRequest}; -const AUTHOR_DIVERSITY_VALUE_MODEL_ID: &str = "author_diversity"; -const MPN_FOLD_MULTIPLIER_MAX: f64 = 10.0; +const DPP_VALUE_MODEL_ID: &str = "dpp"; pub struct VMRanker { pub client: Arc, pub xds_client: Option>, - pub author_cold_start: AuthorColdStart, } #[async_trait] @@ -72,48 +68,11 @@ impl Scorer for VMRanker { .map(|sc| (sc.tweet_id, sc.score)) .collect(); - let fold_weights = (query.params.get(EnableMpnScoring) - && query.params.get(VMRankerValueModelId) == AUTHOR_DIVERSITY_VALUE_MODEL_ID) - .then(|| crate::scorers::ranking_scorer::ScoringWeights::from_params(&query.params)); - - let mut scored: Vec> = candidates + candidates .iter() .map(|c| { - let returned = score_map.get(&c.tweet_id).copied(); - match (&fold_weights, c.mpn_parts, returned, c.score) { - (Some(weights), Some(parts), Some(resp), Some(sent)) if sent > 0.0 => { - let multiplier = (resp / sent).clamp(0.0, MPN_FOLD_MULTIPLIER_MAX); - let model_net = multiplier * parts.pos - parts.neg; - let scaled = if model_net >= 0.0 { - parts.scalar_multiplier * model_net - } else { - model_net - }; - Some(crate::scorers::ranking_scorer::RankingScorer::offset_score( - scaled, weights, - )) - } - _ if fold_weights.is_some() => c.score, - _ => returned.or(c.score), - } - }) - .collect(); - - if fold_weights.is_some() { - let scores: Vec = scored.iter().map(|s| s.unwrap_or(0.0)).collect(); - let effective = self.author_cold_start.apply(query, candidates, &scores); - for (slot, value) in scored.iter_mut().zip(effective) { - if slot.is_some() { - *slot = Some(value); - } - } - } - - scored - .into_iter() - .map(|score| { Ok(PostCandidate { - score, + score: score_map.get(&c.tweet_id).copied().or(c.score), ..Default::default() }) }) @@ -126,80 +85,13 @@ impl Scorer for VMRanker { } fn build_request(query: &ScoredPostsQuery, candidates: &[PostCandidate]) -> RankRequest { - let min_video_duration_ms = query.params.get(MinVideoDurationMs); - let vqv_weight_value = query.params.get(VqvWeight); - let request_timestamp_ms = query.request_time_ms as u64; - let scoring_weights = query - .params - .get(VMRankerSendHeadWeights) - .then(|| crate::scorers::ranking_scorer::ScoringWeights::from_params(&query.params)); - let proto_candidates: Vec = candidates .iter() - .map(|c| { - let phoenix_scores = Some(PhoenixScores { - favorite_score: c.phoenix_scores.favorite_score, - reply_score: c.phoenix_scores.reply_score, - retweet_score: c.phoenix_scores.retweet_score, - photo_expand_score: c.phoenix_scores.photo_expand_score, - click_score: c.phoenix_scores.click_score, - profile_click_score: c.phoenix_scores.profile_click_score, - vqv_score: c.phoenix_scores.vqv_score, - share_score: c.phoenix_scores.share_score, - share_via_dm_score: c.phoenix_scores.share_via_dm_score, - share_via_copy_link_score: c.phoenix_scores.share_via_copy_link_score, - dwell_score: c.phoenix_scores.dwell_score, - quote_score: c.phoenix_scores.quote_score, - quoted_click_score: c.phoenix_scores.quoted_click_score, - follow_author_score: c.phoenix_scores.follow_author_score, - not_interested_score: c.phoenix_scores.not_interested_score, - block_author_score: c.phoenix_scores.block_author_score, - mute_author_score: c.phoenix_scores.mute_author_score, - report_score: c.phoenix_scores.report_score, - not_dwelled_score: c.phoenix_scores.not_dwelled_score, - dwell_time: c.phoenix_scores.dwell_time, - click_dwell_time: c.phoenix_scores.click_dwell_time, - video_open_score: scoring_weights - .as_ref() - .and(c.phoenix_scores.video_open_score), - open_link_score: scoring_weights - .as_ref() - .and(c.phoenix_scores.open_link_score), - quoted_vqv_score: scoring_weights - .as_ref() - .and(c.phoenix_scores.quoted_vqv_score), - post_unexplored_score: scoring_weights - .as_ref() - .and(c.phoenix_scores.post_unexplored_score), - active_secs_5m_residual_norm: scoring_weights - .as_ref() - .and(c.phoenix_scores.active_secs_5m_residual_norm), - }); - - let vqv_weight = - candidates_util::vqv_weight(query, c, min_video_duration_ms, vqv_weight_value); - - RankCandidate { - tweet_id: c.tweet_id, - author_id: c.author_id, - in_network: c.in_network.unwrap_or(false), - is_retweet: c.retweeted_tweet_id.is_some(), - is_reply: c.in_reply_to_tweet_id.is_some(), - author_followers_count: c.author_followers_count.unwrap_or(0), - vqv_ineligible: vqv_weight == 0.0, - retweeted_tweet_id: c.retweeted_tweet_id.unwrap_or(0), - score: c.score, - phoenix_scores, - slate_context: c.slate_context.as_ref().map(|s| SlateContext { - k: s.k, - pool_rank: s.pool_rank, - pool_rank_gap: s.pool_rank_gap, - }), - head_weights: scoring_weights - .as_ref() - .map(|w| w.effective_head_weights(query, c)), - weighted_score: scoring_weights.as_ref().and(c.weighted_score), - } + .map(|c| RankCandidate { + tweet_id: c.tweet_id, + retweeted_tweet_id: c.retweeted_tweet_id.unwrap_or(0), + score: c.score, + ..Default::default() }) .collect(); @@ -217,11 +109,9 @@ fn build_request(query: &ScoredPostsQuery, candidates: &[PostCandidate]) -> Rank RankRequest { viewer_id: query.user_id, - request_timestamp_ms, candidates: proto_candidates, - value_model_id: query.params.get(VMRankerValueModelId), - viewer_following_count: query.user_features.followed_user_ids.len() as u32, + value_model_id: DPP_VALUE_MODEL_ID.to_string(), dpp_params, - new_user_age_threshold_secs: Some(query.params.get(NewUserAgeThresholdSecs)), + ..Default::default() } } diff --git a/home-mixer/selectors/blender_selector.rs b/home-mixer/selectors/blender_selector.rs index 4d98df6e..754b7de7 100644 --- a/home-mixer/selectors/blender_selector.rs +++ b/home-mixer/selectors/blender_selector.rs @@ -1,5 +1,6 @@ use crate::ads::{ - AdsBlender, PartitionOrganicAdsBlender, SafeGapAdsBlender, TimeGapAdsBlender, TimeGapConfig, + AdsBlender, MultiRiskAdsBlender, PartitionOrganicAdsBlender, SafeGapAdsBlender, + TimeGapAdsBlender, TimeGapConfig, }; use crate::frames; use crate::models::query::ScoredPostsQuery; @@ -16,6 +17,7 @@ use xai_recsys_proto::AdIndexInfo; pub struct BlenderSelector { safe_gap_blender: SafeGapAdsBlender, partition_organic_blender: PartitionOrganicAdsBlender, + multi_risk_blender: MultiRiskAdsBlender, } impl BlenderSelector { @@ -23,6 +25,7 @@ impl BlenderSelector { Self { safe_gap_blender: SafeGapAdsBlender::default(), partition_organic_blender: PartitionOrganicAdsBlender, + multi_risk_blender: MultiRiskAdsBlender, } } } @@ -48,7 +51,9 @@ impl Selector for BlenderSelector { let time_gap_blender; let blender: &dyn AdsBlender = match query.params.get(AdsBlenderType).as_str() { + "partition_organic_low_risk" => &self.partition_organic_blender, "safe_gap" => &self.safe_gap_blender, + "multi_risk" => &self.multi_risk_blender, "time_gap" => { time_gap_blender = TimeGapAdsBlender { config: TimeGapConfig { diff --git a/home-mixer/server.rs b/home-mixer/server.rs index fde0054d..d8f59d18 100644 --- a/home-mixer/server.rs +++ b/home-mixer/server.rs @@ -18,6 +18,7 @@ use crate::params; use crate::phoenix_scores_server::{build_query_builder_input, PhoenixScoresServer}; use crate::ranked_following_feed_server::RankedFollowingFeedServer; use crate::scored_posts_server::{build_debug_json, ScoredPostsServer}; +use crate::util::strato_context; use std::collections::HashMap; use std::sync::Arc; use tonic::codec::CompressionEncoding; @@ -328,9 +329,10 @@ impl pb::scored_posts_service_server::ScoredPostsService for ScoredPostsServer { .await?; let RequestContext { b3_info, - query, + mut query, root_span, } = ctx; + query.return_backbone_scores = true; let output = self.run_pipeline(query).instrument(root_span).await?; let debug_json = build_debug_json(&output.pipeline_result); @@ -385,6 +387,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { + let strato_ctx = strato_context::parse(request.metadata()).unwrap_or_default(); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -392,7 +395,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { .ok_or_else(|| Status::invalid_argument("query must be specified"))?; let cursor_str = proto_query.cursor.clone(); let request_context = proto_query.request_context.clone(); - let is_polling = proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -411,6 +414,8 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { query.request_context = request_context; query.is_polling = is_polling; + query.mobile_device_id = strato_ctx.mobile_device_id; + query.mobile_device_ad_id = strato_ctx.ad_id; if !cursor_str.is_empty() { match cursor_utils::decode_ordered_cursor(&cursor_str) { Ok(Some(c)) => { @@ -459,9 +464,10 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { .await?; let RequestContext { b3_info, - query, + mut query, root_span, } = ctx; + query.return_backbone_scores = true; let output = self.get_for_you_feed(query).instrument(root_span).await?; let mut response = Response::new(ForYouFeedResponse { @@ -476,6 +482,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { + let strato_ctx = strato_context::parse(request.metadata()).unwrap_or_default(); let mut b3_info = extract_b3_info(request.metadata()); b3_info.force_sample(); @@ -486,7 +493,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { .ok_or_else(|| Status::invalid_argument("query must be specified"))?; let cursor_str = proto_query.cursor.clone(); let request_context = proto_query.request_context.clone(); - let is_polling = proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -503,8 +510,11 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { root_span, } = ctx; + query.return_backbone_scores = true; query.request_context = request_context; query.is_polling = is_polling; + query.mobile_device_id = strato_ctx.mobile_device_id; + query.mobile_device_ad_id = strato_ctx.ad_id; if !cursor_str.is_empty() { match cursor_utils::decode_ordered_cursor(&cursor_str) { Ok(Some(c)) => { @@ -538,6 +548,7 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService &self, request: Request, ) -> Result, Status> { + let strato_ctx = strato_context::parse(request.metadata()).unwrap_or_default(); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -545,7 +556,7 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService .ok_or_else(|| Status::invalid_argument("query must be specified"))?; let cursor_str = proto_query.cursor.clone(); let request_context = proto_query.request_context.clone(); - let is_polling = proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -565,6 +576,8 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService query.in_network_only = true; query.request_context = request_context; query.is_polling = is_polling; + query.mobile_device_id = strato_ctx.mobile_device_id; + query.mobile_device_ad_id = strato_ctx.ad_id; if !cursor_str.is_empty() { match cursor_utils::decode_ordered_cursor(&cursor_str) { Ok(Some(c)) => { @@ -616,6 +629,7 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService mut query, root_span, } = ctx; + query.return_backbone_scores = true; query.in_network_only = true; let output = self .get_ranked_following_feed(query) @@ -637,6 +651,7 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe &self, request: Request, ) -> Result, Status> { + let strato_ctx = strato_context::parse(request.metadata()).unwrap_or_default(); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -644,7 +659,7 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe .ok_or_else(|| Status::invalid_argument("query must be specified"))?; let cursor_str = proto_query.cursor.clone(); let request_context = proto_query.request_context.clone(); - let is_polling = proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -664,6 +679,8 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe query.in_network_only = true; query.request_context = request_context; query.is_polling = is_polling; + query.mobile_device_id = strato_ctx.mobile_device_id; + query.mobile_device_ad_id = strato_ctx.ad_id; if !cursor_str.is_empty() { match cursor_utils::decode_ordered_cursor(&cursor_str) { Ok(Some(c)) => { @@ -715,6 +732,7 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe mut query, root_span, } = ctx; + query.return_backbone_scores = true; query.in_network_only = true; let output = self.get_following_feed(query).instrument(root_span).await?; diff --git a/home-mixer/side_effects/ads_injection_logging_side_effect.rs b/home-mixer/side_effects/ads_injection_logging_side_effect.rs index 37cdc76d..1a7147ee 100644 --- a/home-mixer/side_effects/ads_injection_logging_side_effect.rs +++ b/home-mixer/side_effects/ads_injection_logging_side_effect.rs @@ -1,5 +1,5 @@ use crate::models::query::{RequestType, ScoredPostsQuery}; -use crate::params::EnableAdsInjectionLogging; +use crate::params::{EnableAdsBrandSafetyVerdictV2, EnableAdsInjectionLogging}; use prost::Message; use std::sync::Arc; use tonic::async_trait; @@ -95,6 +95,12 @@ impl SideEffect for AdsInjectionLoggingSideEffect { ip_address: query.ip_address.clone(), user_agent: query.user_agent.clone(), ads_injection_experiment_bucket: String::new(), + ddg_experiment_bucket: query + .params + .experiment_buckets(EnableAdsBrandSafetyVerdictV2) + .first() + .map(|b| format!("{}:{}", b.experiment, b.bucket)) + .unwrap_or_default(), product_surface: product_surface.into(), }; diff --git a/home-mixer/side_effects/client_events_kafka_side_effect.rs b/home-mixer/side_effects/client_events_kafka_side_effect.rs index 0540eab9..e5765170 100644 --- a/home-mixer/side_effects/client_events_kafka_side_effect.rs +++ b/home-mixer/side_effects/client_events_kafka_side_effect.rs @@ -1,19 +1,19 @@ use crate::models::query::{RequestType, ScoredPostsQuery}; use crate::params::EnableUrtMigrationComponents; -use crate::util::tweet_type_metrics::{bitset_get, TWEET_TYPE_PREDICATES, VIDEO}; +use crate::util::tweet_type_metrics::{TWEET_TYPE_PREDICATES, VIDEO, bitset_get}; use std::collections::HashMap; use std::sync::Arc; use tonic::async_trait; use xai_candidate_pipeline::component_library::clients::kafka_publisher_client::{ - KafkaCluster, KafkaPublisherClient, ProdKafkaPublisherClient, CLIENT_EVENT_TOPIC, + CLIENT_EVENT_TOPIC, KafkaCluster, KafkaPublisherClient, ProdKafkaPublisherClient, }; use xai_candidate_pipeline::component_library::utils::client_utils::{ ClientPlatform, RequestContext, }; use xai_candidate_pipeline::component_library::utils::is_prod; use xai_candidate_pipeline::side_effect::{SideEffect, SideEffectInput}; -use xai_home_mixer_proto::{feed_item, FeedItem, ScoredPost, ServedType}; +use xai_home_mixer_proto::{FeedItem, ScoredPost, ServedType, feed_item}; use xai_x_thrift::log_event::{EventNamespace, LogBase, LogEvent}; use xai_x_thrift::serialize_binary; @@ -61,7 +61,7 @@ impl SideEffect for ClientEventsKafkaSideEffect { .iter() .filter(|i| matches!(i.item, Some(feed_item::Item::WhoToFollow(_)))) .count() as i64; - let post_count = posts.len() as i64; + let post_count = posts.iter().map(|p| conversation_post_count(p)).sum(); let base = ClientEventParams { query, @@ -101,6 +101,22 @@ impl SideEffect for ClientEventsKafkaSideEffect { } } +fn conversation_post_count(post: &ScoredPost) -> i64 { + let mut count = 1; + if let Some(&root) = post.ancestors.iter().min() { + if !post.tombstone_ancestor_ids.contains(&root) { + count += 1; + } + if let Some(&parent) = post.ancestors.iter().max() + && parent != root + && !post.tombstone_ancestor_ids.contains(&parent) + { + count += 1; + } + } + count +} + fn section_for(request_type: RequestType) -> &'static str { match request_type { RequestType::RankedFollowing => "ranked_following", diff --git a/home-mixer/side_effects/mutual_follow_stats_side_effect.rs b/home-mixer/side_effects/mutual_follow_stats_side_effect.rs index da5cb3c6..26b8135b 100644 --- a/home-mixer/side_effects/mutual_follow_stats_side_effect.rs +++ b/home-mixer/side_effects/mutual_follow_stats_side_effect.rs @@ -34,14 +34,13 @@ impl SideEffect for MutualFollowStatsSideEffect let retrieval_cluster: String = input.query.params.get(PhoenixRetrievalInferenceClusterId); let moe_enabled: bool = input.query.params.get(EnablePhoenixMOESource); + let moe_cluster: String = input + .query + .params + .get(PhoenixRetrievalMOEInferenceClusterId); let mut scope: Vec<(&str, &str)> = vec![("retrieval_cluster", &retrieval_cluster)]; - let moe_cluster: String; if moe_enabled { - moe_cluster = input - .query - .params - .get(PhoenixRetrievalMOEInferenceClusterId); scope.push(("moe_cluster", &moe_cluster)); } diff --git a/home-mixer/side_effects/phoenix_request_cache_side_effect.rs b/home-mixer/side_effects/phoenix_request_cache_side_effect.rs index 632affdb..da6761af 100644 --- a/home-mixer/side_effects/phoenix_request_cache_side_effect.rs +++ b/home-mixer/side_effects/phoenix_request_cache_side_effect.rs @@ -94,7 +94,7 @@ impl SideEffect for PhoenixRequestCacheSideEffe let mut tweet_infos = build_tweet_infos(query, &input.selected_candidates); if log_slate_context { for (info, candidate) in tweet_infos.iter_mut().zip(input.selected_candidates.iter()) { - info.score_info = Some(candidate.as_score_info()); + info.score_info = Some(candidate.as_score_info_no_prediction_scores()); } } for candidate in tweet_infos { @@ -230,6 +230,7 @@ mod tests { pool_rank_gap: Some(3), fatigue: 1.5, pre_diversity_score: 0.42, + ..Default::default() }), ..candidate(1, 100) }; @@ -250,8 +251,7 @@ mod tests { ) .await; let score_info = t1.score_info.expect("score_info should be set"); - assert_eq!(score_info.prediction_scores["favorite"], 0.25); - assert_eq!(score_info.prediction_scores["vqv"], 0.0); + assert!(score_info.prediction_scores.is_empty()); assert_eq!(score_info.weighted_score, Some(0.9)); assert_eq!(score_info.final_score, Some(0.6)); let slate_context = score_info diff --git a/home-mixer/side_effects/response_stats_side_effect.rs b/home-mixer/side_effects/response_stats_side_effect.rs index 96a6128c..509d3f85 100644 --- a/home-mixer/side_effects/response_stats_side_effect.rs +++ b/home-mixer/side_effects/response_stats_side_effect.rs @@ -4,6 +4,10 @@ use crate::util::country_codes::bucket_country; use crate::util::feed_log::{count_ads, count_posts, feed_name}; use std::sync::Arc; use tonic::async_trait; +use xai_candidate_pipeline::component_library::utils::client_utils::{ + ClientPlatform, RequestContext, +}; +use xai_candidate_pipeline::component_library::utils::client_version; use xai_candidate_pipeline::side_effect::{SideEffect, SideEffectInput}; use xai_core_entities::entities::SubscriptionLevel; use xai_home_mixer_proto::FeedItem; @@ -40,6 +44,7 @@ impl SideEffect for ResponseStatsSideEffect { }; let feed = feed_name(query.request_type); + stat_request_context(feed, query); stat_response( feed, &input.selected_candidates, @@ -118,3 +123,32 @@ fn stat_response( } } } + +fn stat_request_context(feed: &str, query: &ScoredPostsQuery) { + let Some(receiver) = global_stats_receiver() else { + return; + }; + let platform = ClientPlatform::from_app_id(query.client_app_id); + let client = ("client", platform.stats_client()); + let client_version = ( + "client_version", + client_version::stats_label(client.1, &query.client_version), + ); + let polling = if query.is_polling { "true" } else { "false" }; + let request_context = RequestContext::parse(&query.request_context); + receiver.incr( + &format!("{feed}.request.polling"), + &[("polling", polling), client], + 1, + ); + receiver.incr( + &format!("{feed}.request.request_context"), + &[("request_context", request_context.as_ref()), client], + 1, + ); + receiver.incr( + &format!("{feed}.request.client"), + &[client, client_version], + 1, + ); +} diff --git a/home-mixer/side_effects/scored_stats_side_effect.rs b/home-mixer/side_effects/scored_stats_side_effect.rs index 3be99356..adfb877b 100644 --- a/home-mixer/side_effects/scored_stats_side_effect.rs +++ b/home-mixer/side_effects/scored_stats_side_effect.rs @@ -1,8 +1,8 @@ use crate::models::candidate::PostCandidate; use crate::models::query::{RequestType, ScoredPostsQuery}; use crate::params::{ - EnablePhoenixRetrievalStatsExperimentBucket, PhoenixRetrievalInferenceClusterId, - PhoenixRetrievalMOEInferenceClusterId, TRACE_USER_IDS, + EnablePhoenixRetrievalStatsExperimentBucket, EnablePhoenixScoreStatsExperimentBucket, + PhoenixRetrievalInferenceClusterId, PhoenixRetrievalMOEInferenceClusterId, TRACE_USER_IDS, }; use rand::random; @@ -69,9 +69,21 @@ impl SideEffect for ScoredStatsSideEffect { .query .params .experiment_buckets(EnablePhoenixRetrievalStatsExperimentBucket); + let score_buckets = input + .query + .params + .experiment_buckets(EnablePhoenixScoreStatsExperimentBucket); - if !experiment_buckets.is_empty() || random::() < DEFAULT_SAMPLING_RATE { - record_score_distributions(receiver.as_ref(), "score", candidates.iter()); + let sampled = random::() < DEFAULT_SAMPLING_RATE; + if !score_buckets.is_empty() || sampled { + record_score_distributions( + receiver.as_ref(), + "score", + candidates.iter(), + &score_buckets, + ); + } + if !experiment_buckets.is_empty() || sampled { record_phoenix_retrieval_stats( receiver.as_ref(), &input.selected_candidates, @@ -89,7 +101,7 @@ impl SideEffect for ScoredStatsSideEffect { } } else { if random::() < DEFAULT_SAMPLING_RATE { - record_score_distributions(receiver.as_ref(), "score", candidates.iter()); + record_score_distributions(receiver.as_ref(), "score", candidates.iter(), &[]); } } @@ -122,8 +134,16 @@ fn record_head( metric: &str, name: &str, scores: impl Iterator>, + experiment_buckets: &[&ExperimentBucket], ) { - record_head_with_buckets(receiver, metric, name, scores, HistogramBuckets::Bucket0To1); + record_head_with_buckets( + receiver, + metric, + name, + scores, + HistogramBuckets::Bucket0To1, + experiment_buckets, + ); } fn record_head_with_buckets( @@ -132,9 +152,12 @@ fn record_head_with_buckets( name: &str, scores: impl Iterator>, buckets: HistogramBuckets, + experiment_buckets: &[&ExperimentBucket], ) { let distribution_key = format!("{METRIC_PREFIX}.{metric}Distribution.{name}"); let missing_key = format!("{METRIC_PREFIX}.{metric}Missing.{name}"); + let by_bucket_key = (!experiment_buckets.is_empty()) + .then(|| format!("{METRIC_PREFIX}.{metric}DistributionByBucket.{name}")); let mut present = 0u64; let mut missing = 0u64; for score in scores { @@ -142,6 +165,16 @@ fn record_head_with_buckets( Some(value) => { present += 1; receiver.observe(&distribution_key, &[], value, buckets); + if let Some(key) = &by_bucket_key { + for b in experiment_buckets { + receiver.observe( + key, + &[("ddg", &b.experiment), ("bucket", &b.bucket)], + value, + buckets, + ); + } + } } None => { missing += 1; @@ -156,42 +189,49 @@ fn record_score_distributions<'a>( receiver: &dyn StatsReceiverExt, metric: &str, candidates: impl Iterator + Clone, + experiment_buckets: &[&ExperimentBucket], ) { record_head( receiver, metric, "favorite", candidates.clone().map(|c| c.phoenix_scores.favorite_score), + experiment_buckets, ); record_head( receiver, metric, "reply", candidates.clone().map(|c| c.phoenix_scores.reply_score), + experiment_buckets, ); record_head( receiver, metric, "retweet", candidates.clone().map(|c| c.phoenix_scores.retweet_score), + experiment_buckets, ); record_head( receiver, metric, "click", candidates.clone().map(|c| c.phoenix_scores.click_score), + experiment_buckets, ); record_head( receiver, metric, "vqv", candidates.clone().map(|c| c.phoenix_scores.vqv_score), + experiment_buckets, ); record_head( receiver, metric, "share", candidates.clone().map(|c| c.phoenix_scores.share_score), + experiment_buckets, ); record_head( receiver, @@ -200,6 +240,16 @@ fn record_score_distributions<'a>( candidates .clone() .map(|c| c.phoenix_scores.not_interested_score), + experiment_buckets, + ); + record_head( + receiver, + metric, + "not_dwelled", + candidates + .clone() + .map(|c| c.phoenix_scores.not_dwelled_score), + experiment_buckets, ); record_head_with_buckets( receiver, @@ -207,14 +257,22 @@ fn record_score_distributions<'a>( "dwellTime", candidates.clone().map(|c| c.phoenix_scores.dwell_time), HistogramBuckets::Bucket0To50, + experiment_buckets, ); record_head( receiver, metric, "weightedScore", candidates.clone().map(|c| c.weighted_score), + experiment_buckets, + ); + record_head( + receiver, + metric, + "finalScore", + candidates.map(|c| c.score), + experiment_buckets, ); - record_head(receiver, metric, "finalScore", candidates.map(|c| c.score)); } fn record_trace_author_score_distributions( @@ -249,8 +307,8 @@ fn record_trace_author_score_distributions( .filter(is_original) .partition(|c| c.author_id == author_id); - record_score_distributions(receiver, "traceAuthorScore", author.iter().copied()); - record_score_distributions(receiver, "traceOtherScore", others.iter().copied()); + record_score_distributions(receiver, "traceAuthorScore", author.iter().copied(), &[]); + record_score_distributions(receiver, "traceOtherScore", others.iter().copied(), &[]); } fn post_type(candidate: &PostCandidate) -> &'static str { diff --git a/home-mixer/sources/following_night_owl_source.rs b/home-mixer/sources/following_night_owl_source.rs index 98ce0e82..0195c758 100644 --- a/home-mixer/sources/following_night_owl_source.rs +++ b/home-mixer/sources/following_night_owl_source.rs @@ -178,6 +178,10 @@ fn hit_to_post_candidate(hit: night_owl::SearchHit) -> PostCandidate { .and_then(|d| d.retweet_of_status_id) .and_then(|id| u64::try_from(id).ok()) .unwrap_or(0); + let retweeted_user_id = doc + .and_then(|d| d.retweeted_user_id) + .and_then(|id| u64::try_from(id).ok()) + .unwrap_or(0); let quoted_tweet_id = doc .and_then(|d| d.quoted_tweet_id) .and_then(|id| u64::try_from(id).ok()) @@ -202,6 +206,7 @@ fn hit_to_post_candidate(hit: night_owl::SearchHit) -> PostCandidate { tweet_id: hit.doc_id, author_id, retweeted_tweet_id: (retweeted_tweet_id != 0).then_some(retweeted_tweet_id), + retweeted_user_id: (retweeted_user_id != 0).then_some(retweeted_user_id), in_reply_to_tweet_id: (in_reply_to_tweet_id != 0).then_some(in_reply_to_tweet_id), quoted_tweet_id: (quoted_tweet_id != 0).then_some(quoted_tweet_id), quoted_user_id: (quoted_user_id != 0).then_some(quoted_user_id), diff --git a/home-mixer/util/author_rules.rs b/home-mixer/util/author_rules.rs index ca41f54d..da90f943 100644 --- a/home-mixer/util/author_rules.rs +++ b/home-mixer/util/author_rules.rs @@ -77,7 +77,6 @@ mod tests { impressor, None, false, - None, ) .unwrap(), ) diff --git a/home-mixer/util/mod.rs b/home-mixer/util/mod.rs index 74c06f8c..75d041ec 100644 --- a/home-mixer/util/mod.rs +++ b/home-mixer/util/mod.rs @@ -7,6 +7,7 @@ pub mod feed_log; pub mod phoenix_request; pub mod rescore; pub mod shadow; +pub mod strato_context; pub mod string_case; pub mod tweet_type_metrics; pub mod url; diff --git a/home-mixer/util/phoenix_request.rs b/home-mixer/util/phoenix_request.rs index 96635b5a..222a6ad7 100644 --- a/home-mixer/util/phoenix_request.rs +++ b/home-mixer/util/phoenix_request.rs @@ -1,5 +1,6 @@ use crate::models::candidate::{CandidateHelpers, PostCandidate}; use crate::models::query::ScoredPostsQuery; +use crate::params::{PhoenixExperimentOverrides, RerankerHeadTag}; use rustc_hash::FxHashSet; use xai_candidate_pipeline::component_library::clients::phoenix_prediction_client::TOP_LOG_PROBS_NUM; use xai_geo_ip::zip_to_dma_code; @@ -131,13 +132,27 @@ pub fn build_request_without_sequence_and_candidates( candidate_sets: vec![candidate_set], return_logprob: true, top_logprobs_num: TOP_LOG_PROBS_NUM, + return_backbone_scores: query.return_backbone_scores, client_context: build_client_context(query), user_context: build_user_context(query), metadata: query.request_id.to_string(), + experiment_overrides: parse_experiment_overrides( + &query.params.get(PhoenixExperimentOverrides), + ), ..Default::default() } } +pub fn parse_experiment_overrides(spec: &str) -> std::collections::HashMap { + spec.split(';') + .filter_map(|kv| { + let (k, v) = kv.split_once('=')?; + let (k, v) = (k.trim(), v.trim()); + (!k.is_empty()).then(|| (k.to_string(), v.to_string())) + }) + .collect() +} + pub fn build_prediction_request( query: &ScoredPostsQuery, candidates: &[PostCandidate], @@ -145,7 +160,9 @@ pub fn build_prediction_request( ) -> PredictNextActionsRequest { let mut request = build_request_without_sequence_and_candidates(query, product_surface); request.candidate_sets[0].candidates = build_tweet_infos(query, candidates); - request.sequences = vec![query.scoring_sequence.clone().unwrap_or_default()]; + let mut sequence = query.scoring_sequence.clone().unwrap_or_default(); + sequence.reranker_head_tag = Some(query.params.get(RerankerHeadTag) as u32); + request.sequences = vec![sequence]; request.columnar_sequences = query.columnar_scoring_sequence.iter().cloned().collect(); request } diff --git a/home-mixer/util/strato_context.rs b/home-mixer/util/strato_context.rs new file mode 100644 index 00000000..c7a85b39 --- /dev/null +++ b/home-mixer/util/strato_context.rs @@ -0,0 +1,92 @@ +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use prost::Message; +use tonic::metadata::MetadataMap; + +const STRATO_CONTEXT_KEY: &str = "stratocontext"; +const STRATO_CONTEXT_BIN_KEY: &str = "stratocontext-bin"; + +#[derive(Clone, PartialEq, prost::Message)] +pub struct StratoContext { + #[prost(string, tag = "8")] + pub ad_id: String, + #[prost(bool, tag = "11")] + pub is_polling: bool, + #[prost(string, tag = "12")] + pub mobile_device_id: String, +} + +pub fn parse(metadata: &MetadataMap) -> Option { + if let Some(value) = metadata.get(STRATO_CONTEXT_KEY) + && let Ok(s) = value.to_str() + && let Ok(bytes) = STANDARD.decode(s.trim()) + && let Ok(ctx) = StratoContext::decode(bytes.as_slice()) + { + return Some(ctx); + } + if let Some(value) = metadata.get_bin(STRATO_CONTEXT_BIN_KEY) + && let Ok(bytes) = value.to_bytes() + && let Ok(ctx) = StratoContext::decode(bytes.as_ref()) + { + return Some(ctx); + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + use tonic::metadata::MetadataValue; + + fn encode_ascii(ctx: &StratoContext) -> MetadataMap { + let mut map = MetadataMap::new(); + let encoded = STANDARD.encode(ctx.encode_to_vec()); + map.insert(STRATO_CONTEXT_KEY, encoded.parse().unwrap()); + map + } + + fn polling(ctx: StratoContext) -> StratoContext { + StratoContext { + is_polling: true, + ..ctx + } + } + + #[test] + fn polling_true() { + let ctx = parse(&encode_ascii(&polling(StratoContext::default()))).unwrap(); + assert!(ctx.is_polling); + } + + #[test] + fn polling_false() { + let ctx = parse(&encode_ascii(&StratoContext::default())).unwrap(); + assert!(!ctx.is_polling); + } + + #[test] + fn binary_metadata() { + let mut map = MetadataMap::new(); + map.insert_bin( + STRATO_CONTEXT_BIN_KEY, + MetadataValue::from_bytes(&polling(StratoContext::default()).encode_to_vec()), + ); + assert!(parse(&map).unwrap().is_polling); + } + + #[test] + fn missing_context() { + assert!(parse(&MetadataMap::new()).is_none()); + } + + #[test] + fn decodes_device_ids() { + let ctx = parse(&encode_ascii(&StratoContext { + ad_id: "ad-1".into(), + is_polling: false, + mobile_device_id: "dev-1".into(), + })) + .unwrap(); + assert_eq!(ctx.ad_id, "ad-1"); + assert_eq!(ctx.mobile_device_id, "dev-1"); + } +} diff --git a/home-mixer/util/urt/ad_marshaller.rs b/home-mixer/util/urt/ad_marshaller.rs index ba533753..d3f5bcb4 100644 --- a/home-mixer/util/urt/ad_marshaller.rs +++ b/home-mixer/util/urt/ad_marshaller.rs @@ -1,4 +1,5 @@ use super::client_event::ad_client_event_info; +use crate::models::query::RequestType; const ENTRY_NAMESPACE_PROMOTED_TWEET: &str = "promoted-tweet"; use super::post_marshaller::{make_tweet, make_tweet_item}; @@ -25,13 +26,22 @@ const EMPTY_LANDING_URL: &str = "emptyLandingUrl"; const DSP_IMPRESSION_PREAMBLE: &str = "IS1:"; const DEFAULT_DSP_CREATIVE_ID: &str = "default-dsp-creative-id"; -pub(super) fn marshal_ad(ad: &AdIndexInfo, sort_index: i64) -> TimelineEntry { +pub(super) fn marshal_ad( + ad: &AdIndexInfo, + sort_index: i64, + request_type: RequestType, +) -> TimelineEntry { let impression_string = format!("{:x}", ad.impression_id); if ad.post_id == 0 && ad.rtb_ad_metadata.is_some() { return marshal_ssp_ad(ad, sort_index, &impression_string); } + let entry_id = format!( + "{}-{}-{}", + ENTRY_NAMESPACE_PROMOTED_TWEET, ad.post_id, impression_string + ); + let url_params: BTreeMap = ad .url_params .iter() @@ -183,11 +193,11 @@ pub(super) fn marshal_ad(ad: &AdIndexInfo, sort_index: i64) -> TimelineEntry { tweet.contextual_tweet_ref = Some(contextual_ref(ad.post_id)); TimelineEntry { - entry_id: format!("{}-{}", ENTRY_NAMESPACE_PROMOTED_TWEET, ad.post_id), + entry_id, sort_index, content: TimelineEntryContent::Item(make_tweet_item( tweet, - Some(ad_client_event_info(ad)), + Some(ad_client_event_info(ad, request_type)), None, )), expiry_time: None, @@ -305,7 +315,7 @@ mod tests { use xai_urt_thrift::item::TimelineItemContent; fn promoted_metadata_for(ad: AdIndexInfo) -> PromotedMetadata { - let entry = marshal_ad(&ad, 0); + let entry = marshal_ad(&ad, 0, RequestType::ForYou); let item = match entry.content { TimelineEntryContent::Item(item) => item, _ => panic!("expected an item entry"), @@ -391,7 +401,7 @@ mod tests { #[test] fn ssp_ad_emits_rtb_image_ad_item() { - let entry = marshal_ad(&ssp_ad(), 7); + let entry = marshal_ad(&ssp_ad(), 7, RequestType::ForYou); assert_eq!(entry.entry_id, "rtb-image-ad-abc123"); assert_eq!(entry.sort_index, 7); @@ -440,8 +450,8 @@ mod tests { }), ..Default::default() }; - let entry = marshal_ad(&ad, 0); - assert_eq!(entry.entry_id, "promoted-tweet-42"); + let entry = marshal_ad(&ad, 0, RequestType::ForYou); + assert_eq!(entry.entry_id, "promoted-tweet-42-0"); let item = match entry.content { TimelineEntryContent::Item(item) => item, _ => panic!("expected an item entry"), @@ -455,7 +465,7 @@ mod tests { post_id: 0, ..Default::default() }; - let entry = marshal_ad(&ad, 0); + let entry = marshal_ad(&ad, 0, RequestType::ForYou); let item = match entry.content { TimelineEntryContent::Item(item) => item, _ => panic!("expected an item entry"), diff --git a/home-mixer/util/urt/client_event.rs b/home-mixer/util/urt/client_event.rs index 1d9b6cd4..431b7fc4 100644 --- a/home-mixer/util/urt/client_event.rs +++ b/home-mixer/util/urt/client_event.rs @@ -1,4 +1,5 @@ use super::controller_data; +use crate::models::query::RequestType; use crate::util::string_case::upper_snake_to_pascal; use xai_home_mixer_proto::{ScoredPost, ServedType}; use xai_recsys_proto::AdIndexInfo; @@ -6,9 +7,13 @@ use xai_urt_thrift::metadata::{ClientEventDetails, ClientEventInfo, TimelinesDet pub(super) const ELEMENT_TWEET: &str = "tweet"; pub(super) const ELEMENT_USER: &str = "user"; -pub(super) const COMPONENT_ADS: &str = "for_you_promoted"; pub(super) const COMPONENT_WTF: &str = "suggest_who_to_follow"; -const ADS_INJECTION_TYPE: &str = "ForYouPromoted"; +const COMPONENT_ADS_FOR_YOU: &str = "for_you_promoted"; +const ADS_INJECTION_TYPE_FOR_YOU: &str = "ForYouPromoted"; +const COMPONENT_ADS_FOLLOWING: &str = "following_promoted"; +const ADS_INJECTION_TYPE_FOLLOWING: &str = "FollowingPromoted"; +const COMPONENT_ADS_RANKED_FOLLOWING: &str = "ranked_following_promoted"; +const ADS_INJECTION_TYPE_RANKED_FOLLOWING: &str = "RankedFollowingPromoted"; const WTF_INJECTION_TYPE: &str = "WhoToFollow"; pub(super) fn served_type_component(st: i32) -> String { @@ -68,17 +73,25 @@ pub(super) fn post_client_event_info( } } -pub(super) fn ad_client_event_info(ad: &AdIndexInfo) -> ClientEventInfo { +pub(super) fn ad_client_event_info(ad: &AdIndexInfo, request_type: RequestType) -> ClientEventInfo { + let (component, injection_type) = match request_type { + RequestType::Following => (COMPONENT_ADS_FOLLOWING, ADS_INJECTION_TYPE_FOLLOWING), + RequestType::RankedFollowing => ( + COMPONENT_ADS_RANKED_FOLLOWING, + ADS_INJECTION_TYPE_RANKED_FOLLOWING, + ), + _ => (COMPONENT_ADS_FOR_YOU, ADS_INJECTION_TYPE_FOR_YOU), + }; let mut details = empty_details(); details.timelines_details = Some(TimelinesDetails { - injection_type: Some(ADS_INJECTION_TYPE.to_string()), + injection_type: Some(injection_type.to_string()), controller_data: controller_data::ad_item_controller_data(), source_data: None, }); details.adindex_details = controller_data::ad_index_controller_data(ad); ClientEventInfo { - component: Some(COMPONENT_ADS.to_string()), + component: Some(component.to_string()), element: Some(ELEMENT_TWEET.to_string()), details: Some(details), action: None, @@ -103,11 +116,18 @@ pub(super) fn wtf_item_client_event_info(tracking_token: Option<&str>) -> Client } } -pub(super) fn wtf_module_client_event_info() -> ClientEventInfo { +pub(super) fn wtf_module_client_event_info(tracking_token: Option<&str>) -> ClientEventInfo { + let mut details = empty_details(); + details.timelines_details = Some(TimelinesDetails { + injection_type: Some(WTF_INJECTION_TYPE.to_string()), + controller_data: tracking_token.and_then(controller_data::wtf_controller_data), + source_data: tracking_token.map(|s| s.to_string()), + }); + ClientEventInfo { component: Some(COMPONENT_WTF.to_string()), element: None, - details: None, + details: Some(details), action: None, entity_token: None, } diff --git a/home-mixer/util/urt/mod.rs b/home-mixer/util/urt/mod.rs index 9b8e4924..3f646f37 100644 --- a/home-mixer/util/urt/mod.rs +++ b/home-mixer/util/urt/mod.rs @@ -51,6 +51,7 @@ pub(crate) fn make_urt_timeline( cursor, language_code, country_code, + true, )); let mut entries: Vec = items @@ -87,12 +88,17 @@ pub(crate) fn make_urt_timeline( Some(FeedItemKind::PushToHome(post)) => Some( push_to_home_marshaller::marshal_push_to_home(post, feed_item.position as i64), ), - Some(FeedItemKind::Ad(ad)) => { - Some(ad_marshaller::marshal_ad(ad, feed_item.position as i64)) - } - Some(FeedItemKind::WhoToFollow(wtf)) => { - wtf_marshaller::marshal_wtf(wtf, feed_item.position as i64, client_app_id) - } + Some(FeedItemKind::Ad(ad)) => Some(ad_marshaller::marshal_ad( + ad, + feed_item.position as i64, + request_type, + )), + Some(FeedItemKind::WhoToFollow(wtf)) => wtf_marshaller::marshal_wtf( + wtf, + feed_item.position as i64, + client_app_id, + initial_sort_index, + ), Some(FeedItemKind::Prompt(module)) => { match prompt_marshaller::marshal_prompt(module, feed_item.position as i64) { Some(prompt_marshaller::PromptUrtResult::Entry(entry)) => Some(entry), @@ -258,7 +264,7 @@ mod tests { fn without_nonce_entry_ids_are_unchanged() { let ids = rendered_entry_ids(&[post_item(100), ad_item(200)], None); assert!(ids.contains(&"tweet-100".to_string()), "{ids:?}"); - assert!(ids.contains(&"promoted-tweet-200".to_string()), "{ids:?}"); + assert!(ids.contains(&"promoted-tweet-200-0".to_string()), "{ids:?}"); } #[test] @@ -266,7 +272,7 @@ mod tests { let ids = rendered_entry_ids(&[post_item(100), ad_item(200)], Some(42)); assert!(ids.contains(&"tweet-100-42".to_string()), "{ids:?}"); assert!(!ids.contains(&"tweet-100".to_string()), "{ids:?}"); - assert!(ids.contains(&"promoted-tweet-200".to_string()), "{ids:?}"); + assert!(ids.contains(&"promoted-tweet-200-0".to_string()), "{ids:?}"); assert!( ids.iter().any(|id| id.starts_with("cursor-top-")), "{ids:?}" diff --git a/home-mixer/util/urt/new_tweets_pill.rs b/home-mixer/util/urt/new_tweets_pill.rs index e86c9d0e..237df4e8 100644 --- a/home-mixer/util/urt/new_tweets_pill.rs +++ b/home-mixer/util/urt/new_tweets_pill.rs @@ -24,6 +24,7 @@ pub(super) fn build_new_tweets_pill_instruction( cursor: Option<&UrtOrderedCursor>, language: &str, country: Option<&str>, + require_full_facepile: bool, ) -> Option { let has_top_cursor = cursor .and_then(|c| c.cursor_type.as_ref()) @@ -37,8 +38,8 @@ pub(super) fn build_new_tweets_pill_instruction( return None; } - let user_ids = extract_pill_user_ids(items, viewer_id); - if user_ids.as_ref().is_none_or(|ids| ids.len() < NUM_AVATARS) { + let user_ids = extract_pill_user_ids(items, viewer_id).filter(|ids| ids.len() >= NUM_AVATARS); + if require_full_facepile && user_ids.is_none() { return None; } diff --git a/home-mixer/util/urt/post_marshaller.rs b/home-mixer/util/urt/post_marshaller.rs index 4852a00d..80b95348 100644 --- a/home-mixer/util/urt/post_marshaller.rs +++ b/home-mixer/util/urt/post_marshaller.rs @@ -4,7 +4,7 @@ use xai_urt_thrift::entry::{TimelineEntry, TimelineEntryContent}; use xai_urt_thrift::item::{TimelineItem, TimelineItemContent}; use xai_urt_thrift::metadata::{ ClientEventInfo, ContextType, FeedbackInfo, TopicFeedbackContext, TweetContext, - TweetContextDetails, + TweetContextDetails, Url, UrlType, }; use xai_urt_thrift::tweet::{Tweet, TweetDisplayType, TweetFacepile}; @@ -72,17 +72,39 @@ pub(super) fn tweet_context_from_post(post: &ScoredPost) -> Option if !is_standalone_original_post(post) { return None; } - let topic = post.topic_feedback_topic.as_ref()?; + if let Some(topic) = post.topic_feedback_topic.as_ref().filter(|t| !t.is_empty()) { + return Some(TweetContext { + context_type: ContextType::TOPIC, + text: topic.clone(), + context_image_urls: None, + landing_url: None, + context: Some(TweetContextDetails::TopicFeedbackContext( + TopicFeedbackContext { + topic: Some(topic.clone()), + url: None, + topic_id: post.topic_feedback_topic_id.clone(), + }, + )), + icon: None, + }); + } + let trend = post.ai_trend_name.as_ref().filter(|t| !t.is_empty())?; + let trend_id = post.ai_trend_id.as_ref().filter(|id| !id.is_empty())?; + let landing_url = Url { + url_type: UrlType::EXTERNAL_URL, + url: format!("https://x.com/i/trending/{trend_id}"), + urt_endpoint_options: None, + }; Some(TweetContext { context_type: ContextType::TOPIC, - text: topic.clone(), + text: trend.clone(), context_image_urls: None, - landing_url: None, + landing_url: Some(landing_url.clone()), context: Some(TweetContextDetails::TopicFeedbackContext( TopicFeedbackContext { - topic: Some(topic.clone()), - url: None, - topic_id: post.topic_feedback_topic_id.clone(), + topic: Some(trend.clone()), + url: Some(landing_url), + topic_id: None, }, )), icon: None, diff --git a/home-mixer/util/urt/reverse_chron_following/mod.rs b/home-mixer/util/urt/reverse_chron_following/mod.rs index 9b0b257e..6967dbc2 100644 --- a/home-mixer/util/urt/reverse_chron_following/mod.rs +++ b/home-mixer/util/urt/reverse_chron_following/mod.rs @@ -6,7 +6,7 @@ use super::new_tweets_pill::build_new_tweets_pill_instruction; use super::post_marshaller::marshal_post; use super::prompt_marshaller::{marshal_prompt, PromptUrtResult}; use super::wtf_marshaller::marshal_wtf; -use crate::models::query::FollowingPaginationMeta; +use crate::models::query::{FollowingPaginationMeta, RequestType}; use cursors::{build_cursors, initial_sort_index as compute_initial_sort_index}; use xai_home_mixer_proto::feed_item::Item as FeedItemKind; use xai_home_mixer_proto::FeedItem; @@ -35,6 +35,7 @@ pub(crate) fn make_urt_timeline( cursor, language_code, country_code, + false, ) .into_iter() .collect(); @@ -53,10 +54,17 @@ pub(crate) fn make_urt_timeline( )) } } - Some(FeedItemKind::Ad(ad)) => Some(marshal_ad(ad, feed_item.position as i64)), - Some(FeedItemKind::WhoToFollow(wtf)) => { - marshal_wtf(wtf, feed_item.position as i64, client_app_id) - } + Some(FeedItemKind::Ad(ad)) => Some(marshal_ad( + ad, + feed_item.position as i64, + RequestType::Following, + )), + Some(FeedItemKind::WhoToFollow(wtf)) => marshal_wtf( + wtf, + feed_item.position as i64, + client_app_id, + initial_sort_index, + ), Some(FeedItemKind::Prompt(module)) => { match marshal_prompt(module, feed_item.position as i64) { Some(PromptUrtResult::Entry(entry)) => Some(entry), diff --git a/home-mixer/util/urt/wtf_marshaller.rs b/home-mixer/util/urt/wtf_marshaller.rs index 135efcb7..8dc38d50 100644 --- a/home-mixer/util/urt/wtf_marshaller.rs +++ b/home-mixer/util/urt/wtf_marshaller.rs @@ -20,6 +20,7 @@ pub(super) fn marshal_wtf( wtf: &WhoToFollowModule, sort_index: i64, client_app_id: i32, + module_id: i64, ) -> Option { let is_android = ClientPlatform::is_android(client_app_id); let user_display_type = if is_android { @@ -34,6 +35,8 @@ pub(super) fn marshal_wtf( }; let resp = wtf.who_to_follow_response.as_ref()?; + let module_entry_id = format!("{}-{}", ENTRY_NAMESPACE_WTF, module_id); + let header = resp.header.as_ref().and_then(|h| { h.title.as_ref().map(|title| ModuleHeader { text: title.clone(), @@ -110,7 +113,10 @@ pub(super) fn marshal_wtf( awards_given: None, }; ModuleItem { - entry_id: format!("{}-{}", ENTRY_NAMESPACE_USER, rec.user_id), + entry_id: format!( + "{}-{}-{}", + module_entry_id, ENTRY_NAMESPACE_USER, rec.user_id + ), item: TimelineItem { content: TimelineItemContent::User(user), client_event_info: Some(wtf_item_client_event_info( @@ -127,15 +133,20 @@ pub(super) fn marshal_wtf( }) .collect(); + let first_tracking_token = resp + .user_recommendations + .first() + .and_then(|rec| rec.tracking_token.as_deref()); + Some(TimelineEntry { - entry_id: format!("{}-0", ENTRY_NAMESPACE_WTF), + entry_id: module_entry_id, sort_index, content: TimelineEntryContent::TimelineModule(TimelineModule { items: module_items, display_type: module_display_type, header, footer, - client_event_info: Some(wtf_module_client_event_info()), + client_event_info: Some(wtf_module_client_event_info(first_tracking_token)), feedback_info: None, metadata: None, show_more_behavior: None, diff --git a/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato b/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato index 7b92f8d2..552e4fd2 100644 --- a/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato +++ b/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato @@ -76,19 +76,16 @@ def build32FavIndex(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiv } def build1FavIndex(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiver): Unit = { - candidate.engagementCount match { - case Some(ec) => - if(ec.favoriteCount.getOrElse(0) >= 1) { - val object = { - postId = candidate.postId, - authorId = candidate.authorId, - indexName = "1fav" - } - #.insert((), object) - stats.counter("1fav_indexing_event_created").incr(1) - } - case None => () + if (candidate.engagementCount.flatMap { ec => ec.favoriteCount }.getOrElse(0) == 0) { + stats.counter("1fav_uec_count_race").incr(1) } + val object = { + postId = candidate.postId, + authorId = candidate.authorId, + indexName = "1fav" + } + #.insert((), object) + stats.counter("1fav_indexing_event_created").incr(1) } def build1FavIndexBackup(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiver): Unit = { @@ -263,6 +260,20 @@ def build1FavTopicIndex(candidate: PhoenixRankAllCandidate, stats: Stats.StatsRe buildTopicIndexWithTopics(candidate, stats, 5, postBased50PctTopics) } +def build1FavVideoIndex(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiver): Unit = { + candidate.tweetMedadata match { + case Some(tweet) if eventProcessing.hasValidImmersiveVideo(tweet) => + val object = { + postId = candidate.postId, + authorId = candidate.authorId, + indexName = "1fav_video" + } + #.insert((), object) + stats.counter("1fav_video_indexing_event_created").incr(1) + case _ => () + } +} + def buildPostCreationIndex(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiver): Unit = { val object = { postId = candidate.postId, @@ -340,39 +351,33 @@ def buildMetadataDump(candidate: PhoenixRankAllCandidate, stats: Stats.StatsRece } def buildMMEmbMetadataDump(candidate: PhoenixRankAllCandidate, stats: Stats.StatsReceiver): Unit = { - candidate.engagementCount match { - case Some(ec) => - if(ec.favoriteCount.getOrElse(0) >= 1) { - candidate.tweetMedadata match { - case Some(tweet) => - val hasImage = tweetUtil.hasImage(tweet) - val hasVideo = tweetUtil.hasVideo(tweet) - try { - val emb3Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v3"), ()) - val emb5Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v5_1"), ()) - val emb6Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v6_dev"), ()) - val mmEmbs = { - mmEmbV3 = emb3Result.v.flatMap { emb => emb.embedding1 }, - mmEmbV5 = emb5Result.v.flatMap { emb => emb.embedding1 }, - mmEmbV6 = emb6Result.v.flatMap { emb => emb.embedding1 } - } - val flatObject = { - postId = candidate.postId, - authorId = candidate.authorId, - hasVideo = hasVideo, - hasImage = hasImage, - authorFollowersCount = candidate.authorFollowers, - mmEmbs = Some(mmEmbs), - indexName = Some("mm_emb_metadata") - } - #.insert((), flatObject) - stats.counter("metadata_dump_mm_emb_created").incr(1) - } catch { - case _ => - stats.counter("metadata_dump_mm_emb_fetch_error").incr(1) - } - case None => () + candidate.tweetMedadata match { + case Some(tweet) => + val hasImage = tweetUtil.hasImage(tweet) + val hasVideo = tweetUtil.hasVideo(tweet) + try { + val emb3Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v3"), ()) + val emb5Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v5_1"), ()) + val emb6Result = postMultimodalEmbeddingMhColumnV2.fetch((candidate.postId, "v6_dev"), ()) + val mmEmbs = { + mmEmbV3 = emb3Result.v.flatMap { emb => emb.embedding1 }, + mmEmbV5 = emb5Result.v.flatMap { emb => emb.embedding1 }, + mmEmbV6 = emb6Result.v.flatMap { emb => emb.embedding1 } + } + val flatObject = { + postId = candidate.postId, + authorId = candidate.authorId, + hasVideo = hasVideo, + hasImage = hasImage, + authorFollowersCount = candidate.authorFollowers, + mmEmbs = Some(mmEmbs), + indexName = Some("mm_emb_metadata") } + #.insert((), flatObject) + stats.counter("metadata_dump_mm_emb_created").incr(1) + } catch { + case _ => + stats.counter("metadata_dump_mm_emb_fetch_error").incr(1) } case None => () } @@ -461,6 +466,7 @@ val executeOp = Op.execute({ idempotent = true })[PhoenixRankAllIndexingRequest, build32FavIndex(candidate, receiverStats) if(hasImmersiveVideo) { buildVideoIndex(candidate, receiverStats) + build1FavVideoIndex(candidate, receiverStats) buildImagineIndex(candidate, receiverStats) } buildMetadataDump(candidate, receiverStats) diff --git a/phoenix-rankall-strato/columns/phoenix_rank_all/postCreationEventProcessor.strato b/phoenix-rankall-strato/columns/phoenix_rank_all/postCreationEventProcessor.strato index 831e1816..1848c374 100644 --- a/phoenix-rankall-strato/columns/phoenix_rank_all/postCreationEventProcessor.strato +++ b/phoenix-rankall-strato/columns/phoenix_rank_all/postCreationEventProcessor.strato @@ -4,9 +4,10 @@ type Event = com.twitter.strato.columns.content_understanding.content_understand val processor = # -val executeOp = Op.execute({ idempotent = true })[Event, Unit] { ctx => +val executeOp = Op.execute({ idempotent = true })[(Long, Event), Unit] { ctx => + val (postId, _) = ctx.arg val request = { - postId = ctx.arg.postMetadata.post.postId, + postId = postId, eventSource = PostCreation } processor.execute(request) diff --git a/phoenix-rankall-strato/stream_forwarders/postCreationEventForwarder.strato b/phoenix-rankall-strato/stream_forwarders/postCreationEventForwarder.strato index bb4855dc..6d3726ff 100644 --- a/phoenix-rankall-strato/stream_forwarders/postCreationEventForwarder.strato +++ b/phoenix-rankall-strato/stream_forwarders/postCreationEventForwarder.strato @@ -1,14 +1,16 @@ -import +import StreamForwarder({ source = KafkaTopic({ - topic = "content_understanding_realtime_unified_posts", - dest = "/s/kafka/main-2:kafka-tls", - consumerGroupId = "post_creation_event", - develConsumerGroupId = None, - mapping = ValueOnly({ + topic = "content_understanding_realtime_unified_posts_v3", + dest = "/s/kafka/phoenix-kafka-external-bootstrap", + consumerGroupId = "post_creation_event_v3", + develConsumerGroupId = Some("post_creation_event_v3_devel"), + mapping = KeyValue({ + key = Type(Long), + keyEncoding = NativeEncoding, value = Type(com.twitter.strato.columns.content_understanding.content_understanding.ContentUnderstandingMetadataV2), - encoding = ThriftEncoding + valueEncoding = ThriftEncoding }) }), forwardToOp = Execute, diff --git a/phoenix-rankall/src/config/mod.rs b/phoenix-rankall/src/config/mod.rs index bbf582d6..7aa6b0d4 100644 --- a/phoenix-rankall/src/config/mod.rs +++ b/phoenix-rankall/src/config/mod.rs @@ -142,6 +142,10 @@ impl PipelineKind { WindowConfig::new("post_creation", 24), WindowConfig::new("1fav", 24), WindowConfig::new("1fav", 48), + WindowConfig::new("1fav_video", 48), + WindowConfig::bounded("1fav_video", 24 * 2, 24 * 4), + WindowConfig::bounded("1fav_video", 24 * 4, 24 * 14), + WindowConfig::bounded("1fav_video", 24 * 4, 24 * 30), WindowConfig::new("32fav", 24), WindowConfig::new("video", 48), WindowConfig::new("video", 96), @@ -152,7 +156,6 @@ impl PipelineKind { WindowConfig::new("nsfw_video", 168), WindowConfig::new("evergreen_video", 24 * 365 * 5), WindowConfig::new("evergreen_nsfw_video", 24 * 365 * 5), - WindowConfig::new("evergreen_video_grok", 24 * 30), ], Self::Topic => vec![ WindowConfig::new("1fav", 24), @@ -175,14 +178,20 @@ impl PipelineKind { ], Self::Sid => vec![ WindowConfig::new("1fav", 24), + WindowConfig::new("1fav_video", 48), + WindowConfig::bounded("1fav_video", 24 * 2, 24 * 4), + WindowConfig::bounded("1fav_video", 24 * 4, 24 * 14), + WindowConfig::bounded("1fav_video", 24 * 4, 24 * 30), WindowConfig::new("video", 48), WindowConfig::new("video", 96), + WindowConfig::new("video", 24 * 14), WindowConfig::new("nsfw_video", 48), WindowConfig::new("nsfw_video", 168), WindowConfig::new("nsfw_video", 24 * 14), WindowConfig::new("evergreen_video", 24 * 365 * 5), WindowConfig::new("imagine", 96), - WindowConfig::new("evergreen_video_grok", 24 * 30), + WindowConfig::bounded("video", 24 * 4, 24 * 14), + WindowConfig::bounded("nsfw_video", 24 * 4, 24 * 14), ], Self::SidTail => vec![WindowConfig::new("tail", 24)], Self::Analysis | Self::Ads => vec![], @@ -209,6 +218,7 @@ impl fmt::Display for PipelineKind { pub struct WindowConfig { pub name: String, pub retention: Duration, + pub min_age: Option, } impl WindowConfig { @@ -216,12 +226,25 @@ impl WindowConfig { Self { name: name.into(), retention: Duration::from_secs(retention_hours * 3600), + min_age: None, + } + } + + pub fn bounded(name: impl Into, min_age_hours: u64, retention_hours: u64) -> Self { + assert!(min_age_hours < retention_hours); + Self { + name: name.into(), + retention: Duration::from_secs(retention_hours * 3600), + min_age: Some(Duration::from_secs(min_age_hours * 3600)), } } pub fn window_name(&self) -> String { let days = self.retention.as_secs() / 86400; - format!("{}_{days}day", self.name) + match self.min_age { + Some(min) => format!("{}_{}to{days}day", self.name, min.as_secs() / 86400), + None => format!("{}_{days}day", self.name), + } } } @@ -239,8 +262,12 @@ mod tests { "evergreen_video_1825day" ); assert_eq!( - WindowConfig::new("evergreen_video_grok", 24 * 30).window_name(), - "evergreen_video_grok_30day" + WindowConfig::bounded("video", 24 * 4, 24 * 14).window_name(), + "video_4to14day" + ); + assert_eq!( + WindowConfig::bounded("nsfw_video", 24 * 4, 24 * 14).window_name(), + "nsfw_video_4to14day" ); } @@ -253,8 +280,11 @@ mod tests { assert!(names.contains(&"video_2day".to_string())); assert!(names.contains(&"post_creation_1day".to_string())); assert!(names.contains(&"evergreen_video_1825day".to_string())); - assert!(names.contains(&"evergreen_video_grok_30day".to_string())); - assert_eq!(configs.len(), 14); + assert!(names.contains(&"1fav_video_4to14day".to_string())); + assert!(names.contains(&"1fav_video_4to30day".to_string())); + assert!(names.contains(&"1fav_video_2day".to_string())); + assert!(names.contains(&"1fav_video_2to4day".to_string())); + assert_eq!(configs.len(), 17); } #[test] @@ -333,16 +363,26 @@ mod tests { } #[test] - fn sid_pipeline_includes_evergreen_video_grok_window() { + fn sid_pipeline_includes_bounded_4to14_windows() { let names: Vec = PipelineKind::Sid .window_configs() .iter() .map(|w| w.window_name()) .collect(); - assert!( - names.contains(&"evergreen_video_grok_30day".to_string()), - "Sid window list missing evergreen_video_grok_30day: {names:?}", - ); + for w in [ + "video_4to14day", + "nsfw_video_4to14day", + "1fav_video_4to14day", + "1fav_video_4to30day", + "1fav_video_2day", + "1fav_video_2to4day", + ] { + assert!( + names.contains(&w.to_string()), + "Sid window list missing {w}: {names:?}", + ); + } + assert!(!names.iter().any(|n| n.contains("evergreen_video_grok"))); } #[test] @@ -352,7 +392,7 @@ mod tests { .iter() .map(|w| w.window_name()) .collect(); - for w in ["nsfw_video_7day", "nsfw_video_14day"] { + for w in ["nsfw_video_7day", "nsfw_video_14day", "video_14day"] { assert!( names.contains(&w.to_string()), "Sid window list missing {w}: {names:?}", diff --git a/phoenix-rankall/src/store/base.rs b/phoenix-rankall/src/store/base.rs index 67b7b3f1..13f7d750 100644 --- a/phoenix-rankall/src/store/base.rs +++ b/phoenix-rankall/src/store/base.rs @@ -90,6 +90,8 @@ pub struct BaseSnapshotStore { dump_completed: Arc, } +const PREFLOOR_SLACK_SECS: f64 = 6.0 * 3600.0; + impl BaseSnapshotStore { pub fn new(config: StoreConfig, router: WindowRouter, metrics: Arc) -> Self { std::fs::create_dir_all(&config.output_dir).ok(); @@ -119,42 +121,99 @@ impl BaseSnapshotStore { .output_dir .join(format!("{window_name}.parquet")); - let actual_path = if symlink_path.exists() { + let public_path = if symlink_path.exists() { Some(symlink_path.clone()) } else { find_latest_versioned_file(&self.config.output_dir, &window_name) }; - let Some(actual_path) = actual_path else { - info!("no existing data for window '{window_name}', creating empty parquet"); - state.main_data.insert(window_name.clone(), HashMap::new()); - if let Err(e) = write_empty_parquet( - &self.config.output_dir, - &window_name, - self.config.versions_to_keep, - &self.config.pipeline, - ) { - warn!("failed to create empty parquet for {window_name}: {e}"); + let mut window_data = HashMap::new(); + if let Some(path) = &public_path { + match read_parquet_file(path) { + Ok(batches) => { + for batch in &batches { + load_post_ids_from_batch( + batch, + &mut window_data, + &self.config.pipeline, + ); + } + } + Err(e) => warn!("failed to load {}: {e}", path.display()), } - continue; - }; - - match read_parquet_file(&actual_path) { - Ok(batches) => { - let mut window_data = HashMap::new(); - for batch in &batches { - load_post_ids_from_batch(batch, &mut window_data, &self.config.pipeline); + } + if wc.min_age.is_some() { + let pf_name = format!("prefloor_{window_name}"); + let pf_link = self.config.output_dir.join(format!("{pf_name}.parquet")); + let pf_path = if pf_link.exists() { + Some(pf_link) + } else { + find_latest_versioned_file(&self.config.output_dir, &pf_name) + }; + if let Some(pf_path) = pf_path { + match read_parquet_file(&pf_path) { + Ok(pf_batches) => { + for batch in &pf_batches { + load_post_ids_from_batch( + batch, + &mut window_data, + &self.config.pipeline, + ); + } + } + Err(e) => warn!("failed to load {}: {e}", pf_path.display()), } - let count = window_data.len(); - total_loaded += count; - info!("loaded {count} records into window '{window_name}'"); - state.main_data.insert(window_name, window_data); } - Err(e) => { - warn!("failed to load {}: {e}", actual_path.display()); - state.main_data.insert(window_name, HashMap::new()); + } + + if public_path.is_none() { + let now = chrono::Utc::now().timestamp() as f64; + let retention_cutoff = + timestamp_secs_to_snowflake(now - wc.retention.as_secs() as f64); + let min_age_cutoff = wc + .min_age + .map(|d| timestamp_secs_to_snowflake(now - d.as_secs() as f64)); + let mut entries: Vec<(i64, &StoredValue)> = window_data + .iter() + .map(|(&pid, val)| (pid, val)) + .filter(|&(pid, _)| { + pid >= retention_cutoff && min_age_cutoff.is_none_or(|c| pid < c) + }) + .collect(); + entries.sort_by_key(|(pid, _)| *pid); + info!( + "public file missing for window '{window_name}', regenerating with {} records", + entries.len() + ); + let regen = if entries.is_empty() { + Ok(empty_batch_for_window(&window_name, &self.config.pipeline)) + } else { + stored_to_batch( + &entries, + &window_name, + &self.config.pipeline, + self.config.sid_num_levels, + ) + }; + match regen { + Ok(batch) => { + if let Err(e) = atomic_write_parquet( + &self.config.output_dir, + &window_name, + now as i64, + &batch, + self.config.versions_to_keep, + ) { + warn!("failed to regenerate public parquet for {window_name}: {e}"); + } + } + Err(e) => warn!("failed to build regen batch for {window_name}: {e}"), } } + let count = window_data.len(); + total_loaded += count; + info!("loaded {count} records into window '{window_name}'"); + state.main_data.insert(window_name.clone(), window_data); } if total_loaded > 0 { @@ -259,6 +318,12 @@ impl BaseSnapshotStore { for wc in &self.config.windows { let window_name = wc.window_name(); let retention_secs = wc.retention.as_secs() as f64; + let min_age_cutoff = wc + .min_age + .map(|d| timestamp_secs_to_snowflake(now - d.as_secs() as f64)); + let prefloor_cutoff = wc.min_age.map(|d| { + timestamp_secs_to_snowflake(now - d.as_secs() as f64 - PREFLOOR_SLACK_SECS) + }); if let Some(pending) = pending_snapshot.get(&window_name) && !pending.is_empty() @@ -286,19 +351,48 @@ impl BaseSnapshotStore { let pipeline = &self.config.pipeline; let sid_num_levels = self.config.sid_num_levels; - let batch = if main.is_empty() { + let in_memory = main.len(); + let mut entries: Vec<(i64, &StoredValue)> = main + .iter() + .map(|(&pid, val)| (pid, val)) + .filter(|&(pid, _)| min_age_cutoff.is_none_or(|c| pid < c)) + .collect(); + entries.sort_by_key(|(pid, _)| *pid); + let batch = if entries.is_empty() { empty_batch_for_window(&window_name, pipeline) } else { - let mut entries: Vec<(i64, &StoredValue)> = - main.iter().map(|(&pid, val)| (pid, val)).collect(); - entries.sort_by_key(|(pid, _)| *pid); stored_to_batch(&entries, &window_name, pipeline, sid_num_levels)? }; + let prefloor_batch = match prefloor_cutoff { + Some(c) => { + let mut young: Vec<(i64, &StoredValue)> = main + .iter() + .map(|(&pid, val)| (pid, val)) + .filter(|&(pid, _)| pid >= c) + .collect(); + young.sort_by_key(|(pid, _)| *pid); + Some(if young.is_empty() { + empty_batch_for_window(&window_name, pipeline) + } else { + stored_to_batch(&young, &window_name, pipeline, sid_num_levels)? + }) + } + None => None, + }; let output_dir = self.config.output_dir.clone(); let wn = window_name.clone(); let keep = self.config.versions_to_keep; let written = tokio::task::spawn_blocking(move || { + if let Some(pb) = prefloor_batch { + atomic_write_parquet( + &output_dir, + &format!("prefloor_{wn}"), + epoch_secs, + &pb, + keep, + )?; + } atomic_write_parquet(&output_dir, &wn, epoch_secs, &batch, keep) }) .await??; @@ -320,7 +414,7 @@ impl BaseSnapshotStore { self.metrics .store_buffer_size .with_label_values(&[index_name, window_type]) - .set(written as f64); + .set(in_memory as f64); } self.metrics.store_dump_completed.inc(); @@ -385,9 +479,16 @@ impl SnapshotStore for BaseSnapshotStore { snapshot }; + let mut all_ok = true; for wc in &config.windows { let window_name = wc.window_name(); let retention_secs = wc.retention.as_secs() as f64; + let min_age_cutoff = wc + .min_age + .map(|d| timestamp_secs_to_snowflake(now - d.as_secs() as f64)); + let prefloor_cutoff = wc.min_age.map(|d| { + timestamp_secs_to_snowflake(now - d.as_secs() as f64 - PREFLOOR_SLACK_SECS) + }); if let Some(pending) = pending_snapshot.get(&window_name) && !pending.is_empty() @@ -413,12 +514,16 @@ impl SnapshotStore for BaseSnapshotStore { main.shrink_to_fit(); } - let batch_result = if main.is_empty() { + let in_memory = main.len(); + let mut entries: Vec<(i64, &StoredValue)> = main + .iter() + .map(|(&pid, val)| (pid, val)) + .filter(|&(pid, _)| min_age_cutoff.is_none_or(|c| pid < c)) + .collect(); + entries.sort_by_key(|(pid, _)| *pid); + let batch_result = if entries.is_empty() { Ok(empty_batch_for_window(&window_name, &config.pipeline)) } else { - let mut entries: Vec<(i64, &StoredValue)> = - main.iter().map(|(&pid, val)| (pid, val)).collect(); - entries.sort_by_key(|(pid, _)| *pid); stored_to_batch( &entries, &window_name, @@ -426,12 +531,53 @@ impl SnapshotStore for BaseSnapshotStore { config.sid_num_levels, ) }; + let prefloor_result = match prefloor_cutoff { + Some(c) => { + let mut young: Vec<(i64, &StoredValue)> = main + .iter() + .map(|(&pid, val)| (pid, val)) + .filter(|&(pid, _)| pid >= c) + .collect(); + young.sort_by_key(|(pid, _)| *pid); + if young.is_empty() { + Some(Ok(empty_batch_for_window(&window_name, &config.pipeline))) + } else { + Some(stored_to_batch( + &young, + &window_name, + &config.pipeline, + config.sid_num_levels, + )) + } + } + None => None, + }; match batch_result { Ok(batch) => { let dir = config.output_dir.clone(); let wn = window_name.clone(); let keep = config.versions_to_keep; + let prefloor_batch = match prefloor_result { + Some(Ok(pb)) => Some(pb), + Some(Err(e)) => { + warn!( + "prefloor batch conversion failed for {window_name}: {e}" + ); + all_ok = false; + continue; + } + None => None, + }; match tokio::task::spawn_blocking(move || { + if let Some(pb) = prefloor_batch { + atomic_write_parquet( + &dir, + &format!("prefloor_{wn}"), + epoch_secs, + &pb, + keep, + )?; + } atomic_write_parquet(&dir, &wn, epoch_secs, &batch, keep) }) .await @@ -454,20 +600,33 @@ impl SnapshotStore for BaseSnapshotStore { metrics .store_buffer_size .with_label_values(&[idx, wt]) - .set(written as f64); + .set(in_memory as f64); + } + Ok(Err(e)) => { + warn!("dump failed for {window_name}: {e}"); + all_ok = false; + } + Err(e) => { + warn!("dump task panicked for {window_name}: {e}"); + all_ok = false; } - Ok(Err(e)) => warn!("dump failed for {window_name}: {e}"), - Err(e) => warn!("dump task panicked for {window_name}: {e}"), } } - Err(e) => warn!("batch conversion failed for {window_name}: {e}"), + Err(e) => { + warn!("batch conversion failed for {window_name}: {e}"); + all_ok = false; + } } } - metrics.store_dump_completed.inc(); - metrics.store_dump_last_success_ts.set(epoch_secs); - store_dump_completed.store(true, std::sync::atomic::Ordering::Release); - info!("periodic dump complete"); + if all_ok { + metrics.store_dump_completed.inc(); + metrics.store_dump_last_success_ts.set(epoch_secs); + store_dump_completed.store(true, std::sync::atomic::Ordering::Release); + info!("periodic dump complete"); + } else { + warn!("periodic dump had failures; not signaling completion (offsets held)"); + } } }); @@ -961,24 +1120,6 @@ fn empty_metadata_batch() -> RecordBatch { .expect("empty metadata batch creation cannot fail") } -fn write_empty_parquet( - output_dir: &std::path::Path, - window_name: &str, - versions_to_keep: usize, - pipeline: &str, -) -> anyhow::Result<()> { - let batch = empty_batch_for_window(window_name, pipeline); - let epoch_secs = chrono::Utc::now().timestamp(); - atomic_write_parquet( - output_dir, - window_name, - epoch_secs, - &batch, - versions_to_keep, - )?; - Ok(()) -} - pub fn prefix_window_router() -> WindowRouter { Box::new(|record: &IndexRecord, windows: &[WindowConfig]| { diff --git a/phoenix-rankall/tests/integration_test.rs b/phoenix-rankall/tests/integration_test.rs index b21c130e..f2a4e30a 100644 --- a/phoenix-rankall/tests/integration_test.rs +++ b/phoenix-rankall/tests/integration_test.rs @@ -278,3 +278,95 @@ async fn pipeline_shutdown_flushes_data() { let batches = read_parquet_file(&symlink).unwrap(); assert_eq!(batches[0].num_rows(), 1); } + +#[tokio::test] +async fn bounded_window_young_records_survive_restart() { + let dir = TempDir::new().unwrap(); + let now = chrono::Utc::now().timestamp() as f64; + let young_id = timestamp_secs_to_snowflake(now - 3600.0); + let recent_id = timestamp_secs_to_snowflake(now - 2.0 * 86400.0 - 1800.0); + let mature_id = timestamp_secs_to_snowflake(now - 3.0 * 86400.0); + let expired_id = timestamp_secs_to_snowflake(now - 5.0 * 86400.0); + + let windows = vec![xai_recsys_rankall::config::WindowConfig::bounded( + "1fav_video", + 24 * 2, + 24 * 4, + )]; + let store_config = || StoreConfig { + output_dir: dir.path().to_path_buf(), + windows: windows.clone(), + compaction_interval_secs: 999, + versions_to_keep: 3, + pipeline: "main".to_string(), + sid_num_levels: 6, + }; + let run = |msgs: Vec>>| { + let store = BaseSnapshotStore::new( + store_config(), + prefix_window_router(), + Arc::new(PipelineMetrics::new(prometheus::Registry::new()).unwrap()), + ); + Pipeline::new( + Box::new(MockConsumer::new(msgs)), + Box::new(MainProcessor::new()), + Box::new(store), + PipelineMetrics::new(prometheus::Registry::new()).unwrap(), + ) + .run(CancellationToken::new()) + }; + let read_ids = |path: &std::path::Path| -> Vec { + read_parquet_file(path) + .unwrap() + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() + }; + + run(vec![vec![ + make_core_thrift(young_id, 10, "1fav_video"), + make_core_thrift(recent_id, 15, "1fav_video"), + make_core_thrift(mature_id, 20, "1fav_video"), + make_core_thrift(expired_id, 30, "1fav_video"), + ]]) + .await + .unwrap(); + + let public = dir.path().join("1fav_video_2to4day.parquet"); + let sidecar = dir.path().join("prefloor_1fav_video_2to4day.parquet"); + assert_eq!( + read_ids(&public), + vec![mature_id, recent_id], + "public: matured records only, expired dropped" + ); + assert_eq!( + read_ids(&sidecar), + vec![recent_id, young_id], + "sidecar: young plus slack-band overlap of recently matured" + ); + + run(vec![]).await.unwrap(); + assert_eq!(read_ids(&sidecar), vec![recent_id, young_id]); + + for entry in std::fs::read_dir(dir.path()).unwrap() { + let path = entry.unwrap().path(); + let name = path.file_name().unwrap().to_string_lossy().to_string(); + if name.starts_with("1fav_video_2to4day") { + std::fs::remove_file(&path).unwrap(); + } + } + run(vec![]).await.unwrap(); + assert_eq!( + read_ids(&public), + vec![recent_id], + "public regenerated from sidecar-recovered memory" + ); + assert_eq!(read_ids(&sidecar), vec![recent_id, young_id]); +} diff --git a/phoenix/Cargo.lock b/phoenix/Cargo.lock index eac26746..27d2211a 100644 --- a/phoenix/Cargo.lock +++ b/phoenix/Cargo.lock @@ -2868,6 +2868,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64 0.22.1", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -3468,6 +3478,19 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rcgen" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "yasna", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3941,6 +3964,31 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "serial_test" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "699f4197115b8a7e7ff19c9a315a4bd6fffec26cc4626ef45ecaea389e081c6d" +dependencies = [ + "futures-executor", + "futures-util", + "log", + "once_cell", + "parking_lot", + "serial_test_derive", +] + +[[package]] +name = "serial_test_derive" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e153fc76e1c6a068703d6d29c508a0b15c061c4b7e43da59cc097bc342673c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "sha1" version = "0.10.7" @@ -5048,6 +5096,15 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "xai-o2" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-trait", + "bytes", +] + [[package]] name = "xai-recsys" version = "0.1.0" @@ -5099,10 +5156,13 @@ dependencies = [ "pyo3", "rand 0.9.5", "rayon", + "rcgen", "scopeguard", "serde_json", + "serial_test", "simd-adler32", "static_assertions", + "tempfile", "tokio", "tokio-stream", "tokio-util", @@ -5111,6 +5171,7 @@ dependencies = [ "tonic-reflection", "tower", "url", + "xai-o2", "xai-recsys", "xai-recsys-mm-server", "xai-recsys-proto", @@ -5187,6 +5248,15 @@ dependencies = [ "tonic-prost-build", ] +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + [[package]] name = "yoke" version = "0.8.3" diff --git a/phoenix/Cargo.toml b/phoenix/Cargo.toml index 9cbc645e..dd9cfd42 100644 --- a/phoenix/Cargo.toml +++ b/phoenix/Cargo.toml @@ -9,6 +9,7 @@ members = [ "crates/serving/xai-recsys-proto", "crates/serving/xai-recsys-server", "crates/serving/xai-recsys-sid-proto", + "crates/storage/xai-o2", ] [workspace.package] @@ -57,9 +58,11 @@ prost-build = "0.14" pyo3 = { version = "0.26", features = ["bytes"] } rand = "0.9" rayon = "1" +rcgen = { version = "0.13", features = ["pem"] } rustls = { version = "0.23", features = ["aws-lc-rs"] } scopeguard = "1" serde_json = "1" +serial_test = "3" simd-adler32 = "0.3.7" static_assertions = "1.1.0" tempfile = "3" @@ -76,6 +79,7 @@ tonic-reflection = "0.14" tower = "0.5" tracing = "0.1" url = "2" +xai-o2 = { path = "crates/storage/xai-o2" } xai-recsys = { path = "crates/common/xai-recsys" } xai-recsys-mm-server = { path = "crates/serving/xai-recsys-mm-server" } xai-recsys-proto = { path = "crates/serving/xai-recsys-proto" } diff --git a/phoenix/README.md b/phoenix/README.md index 1f9e7c39..f96de574 100644 --- a/phoenix/README.md +++ b/phoenix/README.md @@ -14,9 +14,11 @@ predicted engagement). > infrastructure (production data feeds, cluster orchestration, internal > telemetry) — every such seam is replaced by a documented local equivalent, > and synthetic data generators are included so the whole system runs end to -> end with nothing external. One training-recipe exception is disclosed in -> [TRAINING.md](TRAINING.md): the dense-optimizer slot ships as standard -> AdamW rather than production's tuned internal variant. +> end with nothing external. One training-recipe substitution is disclosed in +> [TRAINING.md](TRAINING.md): for configs on the legacy dense-optimizer +> slot, the export ships standard AdamW rather than production's tuned +> internal variant. The flagship ranking configs and the nano twin train +> the production Muon recipe, which ships in full. ## Table of Contents @@ -305,7 +307,7 @@ toolchain, `cmake`, `pkg-config`, RDMA verbs headers, bindgen's `libclang`, and `numa_num_possible_nodes` warning) — on Debian/Ubuntu: ```shell -apt update && apt install build-essential cmake pkg-config unzip \ +apt update && apt install build-essential ca-certificates cmake curl pkg-config unzip \ libibverbs-dev libnl-3-dev libnl-route-3-dev libclang-dev libnuma-dev ``` diff --git a/phoenix/TRAINING.md b/phoenix/TRAINING.md index 4675101b..398d96fe 100644 --- a/phoenix/TRAINING.md +++ b/phoenix/TRAINING.md @@ -11,18 +11,28 @@ production-scale recipe. Supply those parts for your own deployment. ## Optimizers and training step -The dense-parameter optimizer in this export is **standard Optax AdamW**. The -internal deployment uses a tuned RMS-normalized-Adam derivative in that -optimizer slot. AdamW is the validated equivalent for the released recipes; -every shipped config trains with it end to end. - -Embedding tables use a separate sparse rowwise AdaGrad optimizer. At a high +Two dense-parameter optimizer families ship in this export: + +- **Muon** (`xrex/optimizers/recsys/muon.py`): the production home-ranker + recipe — consistent-RMS scaling with decoupled weight decay on the matrix + and embedding partitions. The flagship ranking configs and + `home_direct_packed_nano` select it (`optim="muon"`), so the nano trains + with the same dense-optimizer recipe production runs. +- **Standard Optax AdamW**: the slot used by the remaining shipped configs + (two-tower retrieval, gen-recs, and the legacy ranking presets). The + internal deployment uses a tuned RMS-normalized-Adam derivative in that + slot; AdamW is the validated equivalent for those released recipes, and + every config on that slot trains with it end to end. + +Embedding tables use a separate sparse rowwise AdaGrad optimizer; the ranking +flagship recipe (and the nano) additionally runs it with accumulator +half-life decay, lazy per-row decay, and decoupled weight decay. At a high level, each step: 1. looks up and deduplicates the embedding rows used by the batch; 2. computes the loss and gradients for dense parameters and embeddings; -3. applies AdamW to dense parameters and rowwise AdaGrad to the referenced - embedding rows; and +3. applies the config's dense optimizer to dense parameters and rowwise + AdaGrad to the referenced embedding rows; and 4. skips updates when gradients are non-finite. The runnable reference implementation is diff --git a/phoenix/crates/common/xai-recsys/Cargo.toml b/phoenix/crates/common/xai-recsys/Cargo.toml index 88cd403e..af58bbb5 100644 --- a/phoenix/crates/common/xai-recsys/Cargo.toml +++ b/phoenix/crates/common/xai-recsys/Cargo.toml @@ -10,6 +10,7 @@ arrow = { workspace = true } crc32fast = { workspace = true } half = { workspace = true } lazy_static = { workspace = true } +log = { workspace = true } prometheus = { workspace = true } xai-recsys-proto = { workspace = true } diff --git a/phoenix/crates/common/xai-recsys/src/model_config.rs b/phoenix/crates/common/xai-recsys/src/model_config.rs index 853aa05d..a1a94eb7 100644 --- a/phoenix/crates/common/xai-recsys/src/model_config.rs +++ b/phoenix/crates/common/xai-recsys/src/model_config.rs @@ -1,5 +1,73 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 X.AI Corp. +use std::fmt; +use std::str::FromStr; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum MultimodalEmbeddingType { + #[default] + None, + V1, + V3, + V5, + V6, + V8, +} + +impl MultimodalEmbeddingType { + pub const fn dim(self) -> usize { + match self { + Self::None => 0, + Self::V1 => 1536, + Self::V3 | Self::V5 | Self::V6 | Self::V8 => 1024, + } + } + + pub const fn as_str(self) -> &'static str { + match self { + Self::None => "none", + Self::V1 => "v1", + Self::V3 => "v3", + Self::V5 => "v5", + Self::V6 => "v6", + Self::V8 => "v8", + } + } + + pub const fn is_enabled(self) -> bool { + !matches!(self, Self::None) + } + + pub fn trainer_override(self) -> Option { + self.is_enabled() + .then(|| format!("multimodal_embedding_type={}", self.as_str())) + } +} + +impl fmt::Display for MultimodalEmbeddingType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for MultimodalEmbeddingType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.trim().to_ascii_lowercase().as_str() { + "" | "none" | "null" => Ok(Self::None), + "v1" => Ok(Self::V1), + "v3" => Ok(Self::V3), + "v5" => Ok(Self::V5), + "v6" => Ok(Self::V6), + "v8" => Ok(Self::V8), + other => Err(format!( + "unknown multimodal_embedding_type {other:?} (expected none|v1|v3|v5|v6|v8)" + )), + } + } +} + #[derive(Debug, Clone)] pub struct HashTableConfig { pub user_id_table_size: usize, @@ -25,7 +93,6 @@ pub struct HashTableConfig { pub output_vocab_size: usize, pub num_continuous_actions: usize, - pub search_query_embedding_dim: usize, pub num_user_categorical_features: usize, pub num_user_bool_features: usize, @@ -37,6 +104,8 @@ pub struct HashTableConfig { pub num_post_bool_features: usize, pub num_post_float_features: usize, pub num_post_int64_features: usize, + + pub enable_stale_post: bool, } impl HashTableConfig { @@ -249,6 +318,7 @@ impl ModelConfig { num_post_bool_features: usize, num_post_float_features: usize, num_post_int64_features: usize, + enable_stale_post: bool, ) -> Self { ModelConfig { hash_table: HashTableConfig { @@ -271,7 +341,6 @@ impl ModelConfig { ip_modulus, output_vocab_size, num_continuous_actions, - search_query_embedding_dim, num_user_categorical_features, num_user_bool_features, num_user_float_features, @@ -281,6 +350,7 @@ impl ModelConfig { num_post_bool_features, num_post_float_features, num_post_int64_features, + enable_stale_post, }, history_seq_len, candidate_seq_len, @@ -291,3 +361,31 @@ impl ModelConfig { } } } + +#[cfg(test)] +mod tests { + use super::MultimodalEmbeddingType; + + #[test] + fn type_dim_matches_python_embedding_config() { + assert_eq!(MultimodalEmbeddingType::None.dim(), 0); + assert_eq!(MultimodalEmbeddingType::V1.dim(), 1536); + assert_eq!(MultimodalEmbeddingType::V3.dim(), 1024); + assert_eq!(MultimodalEmbeddingType::V5.dim(), 1024); + assert_eq!(MultimodalEmbeddingType::V6.dim(), 1024); + assert_eq!(MultimodalEmbeddingType::V8.dim(), 1024); + } + + #[test] + fn type_parses_cli_and_override_spellings() { + assert_eq!( + "none".parse::().unwrap(), + MultimodalEmbeddingType::None + ); + assert_eq!( + "V8".parse::().unwrap(), + MultimodalEmbeddingType::V8 + ); + assert!("v2".parse::().is_err()); + } +} diff --git a/phoenix/crates/common/xai-recsys/src/util.rs b/phoenix/crates/common/xai-recsys/src/util.rs index 2aeeed6f..35f7de7f 100644 --- a/phoenix/crates/common/xai-recsys/src/util.rs +++ b/phoenix/crates/common/xai-recsys/src/util.rs @@ -1,11 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 X.AI Corp. use crate::model_config::ModelConfig; -use arrow::array::{Array, AsArray, BooleanArray, Float32Array, Int32Array, Int64Array}; +use arrow::array::{ + Array, AsArray, BooleanArray, FixedSizeListArray, Float32Array, Int32Array, Int64Array, +}; use arrow::ipc::reader::StreamReader; use half::f16; use lazy_static::lazy_static; use prometheus::{IntCounterVec, register_int_counter_vec}; +use std::collections::HashMap; use std::io::Cursor; use std::time::{SystemTime, UNIX_EPOCH}; use xai_recsys_proto as pb; @@ -26,7 +29,76 @@ lazy_static! { .unwrap(); } -fn record_sid_coverage(sequence: &str, present: u64, count: u64) { +pub const WEB_CONV_FAKE_TWEET_ID: i64 = 4; + +pub fn conv_asset_map(ids: Option<&pb::ConvAssetIds>) -> HashMap<(i64, i64), i64> { + ids.filter(|ids| { + ids.asset_id.len() == ids.author_id.len() + && ids.asset_id.len() == ids.impressed_time_ms.len() + }) + .map(|ids| { + (0..ids.asset_id.len()) + .filter(|&i| ids.asset_id[i] > 0) + .map(|i| { + ( + (ids.author_id[i], ids.impressed_time_ms[i]), + ids.asset_id[i], + ) + }) + .collect() + }) + .unwrap_or_default() +} + +pub fn is_web_conv_row(tweet_id: i64, has_conv_bit: bool) -> bool { + tweet_id == WEB_CONV_FAKE_TWEET_ID || has_conv_bit +} + +pub fn conv_asset_ids_for_batch( + batch: &arrow::record_batch::RecordBatch, + map: &HashMap<(i64, i64), i64>, +) -> Vec { + let n = batch.num_rows(); + let mut out = vec![0i64; n]; + if map.is_empty() { + return out; + } + let col_i64 = |name: &str| { + batch + .column_by_name(name) + .and_then(|c| c.as_any().downcast_ref::()) + }; + let (Some(tweet_ids), Some(author_ids), Some(impressed_ms)) = ( + col_i64("tweetId"), + col_i64("authorId"), + col_i64("impressedTimeMs"), + ) else { + return out; + }; + let conv_bit = pb::ActionName::AdsWebConversion as usize; + let multi_hot = batch + .column_by_name("actionNameMultiHot") + .and_then(|c| c.as_any().downcast_ref::()) + .filter(|fsl| conv_bit < fsl.value_length() as usize); + let (bools, vocab) = match multi_hot { + Some(fsl) => ( + fsl.values().as_any().downcast_ref::(), + fsl.value_length() as usize, + ), + None => (None, 0), + }; + for (row, slot) in out.iter_mut().enumerate() { + let has_conv_bit = bools.is_some_and(|b| b.value(row * vocab + conv_bit)); + if is_web_conv_row(tweet_ids.value(row), has_conv_bit) + && let Some(&asset_id) = map.get(&(author_ids.value(row), impressed_ms.value(row))) + { + *slot = asset_id; + } + } + out +} + +pub fn record_sid_coverage(sequence: &str, present: u64, count: u64) { SID_COVERAGE_TOTAL .with_label_values(&[sequence, "present"]) .inc_by(present); @@ -35,6 +107,22 @@ fn record_sid_coverage(sequence: &str, present: u64, count: u64) { .inc_by(count.saturating_sub(present)); } +pub fn stamp_semantic_ids(dst: &mut [u16], entry_idx: usize, sid_num_levels: usize, codes: &[i32]) { + if sid_num_levels > 0 && codes.len() == sid_num_levels { + let base = entry_idx * sid_num_levels; + for (d, &c) in dst[base..base + sid_num_levels] + .iter_mut() + .zip(codes.iter()) + { + *d = (c + 1) as u16; + } + } +} + +use crate::feature_config::bool_feature::{ + IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ, IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ_COLUMN, + IS_AUTHOR_FOLLOWING_VIEWER_SEQ, IS_AUTHOR_FOLLOWING_VIEWER_SEQ_COLUMN, IS_STALE_POST14D, +}; use crate::feature_config::categorical_feature::{ AUTHOR_IS_NSFW_SEQ, LOCAL_DAY_OF_WEEK_SEQ, LOCAL_HOUR_OF_DAY_SEQ, PRODUCT_SURFACE_SEQ, PRODUCT_SURFACE_SEQ_COLUMN, TIMEZONE_SEQ, @@ -43,6 +131,7 @@ use crate::feature_config::categorical_feature::{ use crate::feature_config::constants::{ ADS_PRODUCT_KEY_HASH_BIAS, ADS_PRODUCT_KEY_HASH_BIAS_2, ADS_PRODUCT_KEY_HASH_MODULUS, ADS_PRODUCT_KEY_HASH_SCALE, ADS_PRODUCT_KEY_HASH_SCALE_2, ADS_PRODUCT_KEY_TABLE_SIZE, + STALE_POST_14D_TTL_SEC, }; use crate::feature_config::int64_feature::{ FAV_COUNT_SEQ, FAV_COUNT_SEQ_COLUMN, QUOTE_COUNT_SEQ, QUOTE_COUNT_SEQ_COLUMN, REPLY_COUNT_SEQ, @@ -133,6 +222,14 @@ pub fn stamp_i32_as_categorical( } } +pub fn stamp_bool_seq(source: &[bool], dest: &mut [bool], num_features: usize, feature_idx: usize) { + if num_features > feature_idx { + for (i, &val) in source.iter().enumerate() { + dest[i * num_features + feature_idx] = val; + } + } +} + pub fn stamp_local_time_features( impr_ts_sec: &[i32], tz_enums: &[i16], @@ -306,6 +403,30 @@ pub struct InputBuffer { pub num_history: usize, } +pub fn repeat_query_into(dest: &mut [f32], query: &[f32], n_slots: usize) { + let dim = query.len(); + if dim == 0 || dest.is_empty() || n_slots == 0 { + return; + } + let max_slots = dest.len() / dim; + let n = n_slots.min(max_slots); + for i in 0..n { + let off = i * dim; + dest[off..off + dim].copy_from_slice(query); + } +} + +impl InputBuffer { + pub fn num_real_candidates(&self, num_item_hashes: usize, cap: usize) -> usize { + let w = num_item_hashes.max(1); + self.candidate_post_hashes + .chunks(w) + .take(cap) + .take_while(|ch| ch.first().is_some_and(|&h| h != 0)) + .count() + } +} + struct CandidateData { post_hashes: Vec, auth_hashes: Vec, @@ -333,7 +454,7 @@ impl InputBuffer { let num_author_hashes = model_config.hash_table.num_author_hashes(); let num_item_hashes = model_config.hash_table.num_item_hashes(); let candidate_seq_len = model_config.candidate_seq_len; - let search_query_embedding_dim = model_config.hash_table.search_query_embedding_dim; + let search_query_embedding_dim = model_config.search_query_embedding_dim; let mut candidate_post_hashes = vec![0i32; candidate_seq_len * num_item_hashes]; let mut candidate_auth_hashes = vec![0i32; candidate_seq_len * num_author_hashes]; @@ -384,15 +505,12 @@ impl InputBuffer { candidate_funding_instrument_ids[j] = ad.funding_instrument_id; } - if sid_num_levels > 0 && candidate.semantic_ids.len() == sid_num_levels { - let base = j * sid_num_levels; - for (d, &c) in candidate_semantic_ids[base..base + sid_num_levels] - .iter_mut() - .zip(candidate.semantic_ids.iter()) - { - *d = (c + 1) as u16; - } - } + stamp_semantic_ids( + &mut candidate_semantic_ids, + j, + sid_num_levels, + &candidate.semantic_ids, + ); } if sid_num_levels > 0 { @@ -402,19 +520,19 @@ impl InputBuffer { record_sid_coverage("candidate", present, candidates_to_process as u64); } - let mut candidate_search_query_embeddings = - vec![0.0f32; candidate_seq_len * search_query_embedding_dim]; - if search_query_embedding_dim > 0 && !candidate_set.search_query_embedding.is_empty() { - let provided_dim = candidate_set.search_query_embedding.len(); - if provided_dim == search_query_embedding_dim { - for j in 0..candidates_to_process { - let base_idx = j * search_query_embedding_dim; - candidate_search_query_embeddings - [base_idx..base_idx + search_query_embedding_dim] - .copy_from_slice(&candidate_set.search_query_embedding); - } + let candidate_search_query_embeddings = if search_query_embedding_dim > 0 + && candidate_set.search_query_embedding.len() == search_query_embedding_dim + { + candidate_set.search_query_embedding.clone() + } else { + if search_query_embedding_dim > 0 && !candidate_set.search_query_embedding.is_empty() { + log::error!( + "search_query_embedding dim {} != model {search_query_embedding_dim}; leaving empty", + candidate_set.search_query_embedding.len() + ); } - } + Vec::new() + }; let n_post_cat = model_config.hash_table.num_post_categorical_features; let n_post_bool = model_config.hash_table.num_post_bool_features; @@ -443,22 +561,46 @@ impl InputBuffer { ); let mut candidate_int64_features = vec![0i64; candidate_seq_len * n_post_int64]; + let mut candidate_is_author_followed = vec![false; candidate_seq_len]; + let mut candidate_is_author_following = vec![false; candidate_seq_len]; + let mut candidate_is_stale_post = vec![false; candidate_seq_len]; + let mut candidate_bool_features = vec![false; candidate_seq_len * n_post_bool]; + + let stale_post_enabled = model_config.hash_table.enable_stale_post; for (j, candidate) in candidate_set .candidates .iter() .take(candidates_to_process) .enumerate() { - stamp_engagement_counts( - &mut candidate_int64_features, - n_post_int64, - j, - candidate.fav_count, - candidate.reply_count, - candidate.retweet_count, - candidate.quote_count, - candidate.view_count, - ); + let creation_valid = candidate_post_creation_ts_sec[j] > 0; + let original_age_sec = now_sec as i64 - candidate_post_creation_ts_sec[j] as i64; + let is_stale = + stale_post_enabled && creation_valid && original_age_sec > STALE_POST_14D_TTL_SEC; + candidate_is_stale_post[j] = is_stale; + if is_stale { + stamp_engagement_counts( + &mut candidate_int64_features, + n_post_int64, + j, + 0, + 0, + 0, + 0, + 0, + ); + } else { + stamp_engagement_counts( + &mut candidate_int64_features, + n_post_int64, + j, + candidate.fav_count, + candidate.reply_count, + candidate.retweet_count, + candidate.quote_count, + candidate.view_count, + ); + } #[cfg(recsys_ads_dpa)] { @@ -476,6 +618,12 @@ impl InputBuffer { hash_dpa_product_key_2(raw_key); } } + candidate_is_author_followed[j] = candidate.is_author_followed_by_user; + candidate_is_author_following[j] = candidate + .author_info + .as_ref() + .and_then(|ai| ai.is_following_user) + .unwrap_or(false); } let candidate_author_is_nsfw: Vec = candidate_set @@ -491,6 +639,25 @@ impl InputBuffer { AUTHOR_IS_NSFW_CATEGORICAL_IDX, ); + stamp_bool_seq( + &candidate_is_stale_post, + &mut candidate_bool_features, + n_post_bool, + IS_STALE_POST14D, + ); + stamp_bool_seq( + &candidate_is_author_followed, + &mut candidate_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ, + ); + stamp_bool_seq( + &candidate_is_author_following, + &mut candidate_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWING_VIEWER_SEQ, + ); + CandidateData { post_hashes: candidate_post_hashes, auth_hashes: candidate_auth_hashes, @@ -498,7 +665,7 @@ impl InputBuffer { embeddings: mm_embeddings_opt.unwrap_or_default(), search_query_embeddings: candidate_search_query_embeddings, categorical_features, - bool_features: vec![false; candidate_seq_len * n_post_bool], + bool_features: candidate_bool_features, float_features: vec![0.0f32; candidate_seq_len * n_post_float], int64_features: candidate_int64_features, impr_ts: candidate_impr_ts, @@ -666,8 +833,12 @@ impl InputBuffer { let mut history_post_creation_ts_sec = vec![0i32; history_seq_len]; let mut history_tz_enums = vec![0i16; history_seq_len]; let mut history_author_is_nsfw = vec![0i32; history_seq_len]; + let mut history_is_author_followed = vec![false; history_seq_len]; + let mut history_is_author_following = vec![false; history_seq_len]; let mut history_post_ids = vec![0i64; history_seq_len]; let mut history_int64_features = vec![0i64; history_seq_len * n_post_int64]; + let sid_num_levels = model_config.sid_num_levels; + let mut history_semantic_ids = vec![0u16; history_seq_len * sid_num_levels]; let _empty = Vec::new(); let agg_user_actions = match sequence { @@ -709,6 +880,13 @@ impl InputBuffer { history_post_ids[valid_entry_count] = tweet_id; + stamp_semantic_ids( + &mut history_semantic_ids, + valid_entry_count, + sid_num_levels, + &tweet_info.semantic_ids, + ); + let base_idx = valid_entry_count * output_vocab_size; let continuous_base_idx = valid_entry_count * num_continuous_actions; let mut total_dwell_time: f64 = 0.0; @@ -792,10 +970,25 @@ impl InputBuffer { tweet_info.view_count, ); + history_is_author_followed[valid_entry_count] = + tweet_info.is_author_followed_by_user; + history_is_author_following[valid_entry_count] = tweet_info + .author_info + .as_ref() + .and_then(|ai| ai.is_following_user) + .unwrap_or(false); + valid_entry_count += 1; } } + if sid_num_levels > 0 { + let present = (0..valid_entry_count) + .filter(|&i| history_semantic_ids[i * sid_num_levels] != 0) + .count() as u64; + record_sid_coverage("history", present, valid_entry_count as u64); + } + let user_features = Self::extract_user_features(&model_config.hash_table, client_context, user_context); @@ -823,6 +1016,20 @@ impl InputBuffer { AUTHOR_IS_NSFW_CATEGORICAL_IDX, ); + let mut history_bool_features = vec![false; history_seq_len * n_post_bool]; + stamp_bool_seq( + &history_is_author_followed, + &mut history_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ, + ); + stamp_bool_seq( + &history_is_author_following, + &mut history_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWING_VIEWER_SEQ, + ); + let request_ip = candidate_set .device_feature .as_ref() @@ -883,7 +1090,7 @@ impl InputBuffer { user_int64_features: user_features.int64_features, user_installed_apps: user_features.installed_apps, history_categorical_features, - history_bool_features: vec![false; history_seq_len * n_post_bool], + history_bool_features, history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features, candidate_categorical_features, @@ -908,7 +1115,7 @@ impl InputBuffer { user_conversion_history_hashes, candidate_account_hashes, history_post_ids, - history_semantic_ids: vec![0u16; history_seq_len * model_config.sid_num_levels], + history_semantic_ids, candidate_semantic_ids, num_history: valid_entry_count, } @@ -965,6 +1172,8 @@ impl InputBuffer { let mut history_tz_enums = vec![0i16; history_seq_len]; let mut history_post_ids = vec![0i64; history_seq_len]; let mut history_int64_features = vec![0i64; history_seq_len * n_post_int64]; + let mut history_is_author_followed = vec![false; history_seq_len]; + let mut history_is_author_following = vec![false; history_seq_len]; let sid_num_levels = model_config.sid_num_levels; let mut history_semantic_ids = vec![0u16; history_seq_len * sid_num_levels]; @@ -984,6 +1193,7 @@ impl InputBuffer { let n_post_bool = model_config.hash_table.num_post_bool_features; let n_post_float = model_config.hash_table.num_post_float_features; let n_post_int64 = model_config.hash_table.num_post_int64_features; + let history_bool_features = vec![false; history_seq_len * n_post_bool]; let request_ip = candidate_set .device_feature .as_ref() @@ -1044,7 +1254,7 @@ impl InputBuffer { user_int64_features: user_features.int64_features, user_installed_apps: user_features.installed_apps, history_categorical_features: vec![0i16; history_seq_len * n_post_cat], - history_bool_features: vec![false; history_seq_len * n_post_bool], + history_bool_features, history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features: vec![0i64; history_seq_len * n_post_int64], candidate_categorical_features, @@ -1114,6 +1324,12 @@ impl InputBuffer { let col_view_count = batch .column_by_name(VIEW_COUNT_SEQ_COLUMN) .and_then(|c| c.as_any().downcast_ref::()); + let col_is_author_followed_by_viewer = batch + .column_by_name(IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); + let col_is_author_following_viewer = batch + .column_by_name(IS_AUTHOR_FOLLOWING_VIEWER_SEQ_COLUMN) + .and_then(|c| c.as_any().downcast_ref::()); let start_row = num_rows.saturating_sub(history_seq_len); let mut valid_entry_count = 0; @@ -1141,19 +1357,17 @@ impl InputBuffer { history_post_ids[valid_entry_count] = tweet_id; - if sid_num_levels > 0 - && let Some(sid_col) = col_semantic_id - { + if let Some(sid_col) = col_semantic_id { let fsl = sid_col.as_fixed_size_list(); if !fsl.is_null(row_idx) { let inner = fsl.value(row_idx); - if let Some(codes) = inner.as_any().downcast_ref::() - && codes.len() == sid_num_levels - { - let base = valid_entry_count * sid_num_levels; - for d in 0..sid_num_levels { - history_semantic_ids[base + d] = (codes.value(d) + 1) as u16; - } + if let Some(codes) = inner.as_any().downcast_ref::() { + stamp_semantic_ids( + &mut history_semantic_ids, + valid_entry_count, + sid_num_levels, + codes.values(), + ); } } } @@ -1208,6 +1422,11 @@ impl InputBuffer { col_view_count.map_or(0, |arr| arr.value(row_idx)) as u64, ); + history_is_author_followed[valid_entry_count] = col_is_author_followed_by_viewer + .is_some_and(|arr| !arr.is_null(row_idx) && arr.value(row_idx)); + history_is_author_following[valid_entry_count] = col_is_author_following_viewer + .is_some_and(|arr| !arr.is_null(row_idx) && arr.value(row_idx)); + valid_entry_count += 1; } @@ -1239,6 +1458,20 @@ impl InputBuffer { n_post_cat, ); + let mut history_bool_features = vec![false; history_seq_len * n_post_bool]; + stamp_bool_seq( + &history_is_author_followed, + &mut history_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWED_BY_VIEWER_SEQ, + ); + stamp_bool_seq( + &history_is_author_following, + &mut history_bool_features, + n_post_bool, + IS_AUTHOR_FOLLOWING_VIEWER_SEQ, + ); + let request_ip = candidate_set .device_feature .as_ref() @@ -1299,7 +1532,7 @@ impl InputBuffer { user_int64_features: user_features.int64_features, user_installed_apps: user_features.installed_apps, history_categorical_features, - history_bool_features: vec![false; history_seq_len * n_post_bool], + history_bool_features, history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features, candidate_categorical_features, @@ -1385,6 +1618,78 @@ impl InputBuffer { #[cfg(test)] mod tests { + + #[test] + fn conv_asset_ids_for_batch_predicate_and_keys() { + use arrow::array::{ArrayRef, BooleanArray, FixedSizeListArray, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + let rows = [ + (11, 21, 111_000, false), + (WEB_CONV_FAKE_TWEET_ID, 7, 1_000, false), + (555, 8, 2_000, true), + (WEB_CONV_FAKE_TWEET_ID, 7, 3_000, false), + (WEB_CONV_FAKE_TWEET_ID, 7, 4_000, false), + ]; + let vocab = 256usize; + let conv_bit = pb::ActionName::AdsWebConversion as usize; + assert!(conv_bit < vocab); + let mut bits = vec![false; rows.len() * vocab]; + for (i, r) in rows.iter().enumerate() { + bits[i * vocab + conv_bit] = r.3; + } + let multi_hot = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Boolean, true)), + vocab as i32, + Arc::new(BooleanArray::from(bits)), + None, + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("tweetId", DataType::Int64, false), + Field::new("authorId", DataType::Int64, false), + Field::new("impressedTimeMs", DataType::Int64, false), + Field::new("actionNameMultiHot", multi_hot.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from( + rows.iter().map(|r| r.0).collect::>(), + )) as ArrayRef, + Arc::new(Int64Array::from( + rows.iter().map(|r| r.1).collect::>(), + )), + Arc::new(Int64Array::from( + rows.iter().map(|r| r.2).collect::>(), + )), + Arc::new(multi_hot), + ], + ) + .unwrap(); + + let ids = pb::ConvAssetIds { + impressed_time_ms: vec![111_000, 1_000, 2_000, 4_000], + author_id: vec![21, 7, 8, 7], + asset_id: vec![99, 42, 9, 0], + }; + let map = conv_asset_map(Some(&ids)); + assert_eq!(map.len(), 3, "asset_id 0 windows are dropped from the map"); + assert_eq!(conv_asset_ids_for_batch(&batch, &map), vec![0, 42, 9, 0, 0]); + + let ragged = pb::ConvAssetIds { + impressed_time_ms: vec![1_000], + author_id: vec![7, 8], + asset_id: vec![42], + }; + assert!(conv_asset_map(Some(&ragged)).is_empty()); + assert!(conv_asset_map(None).is_empty()); + assert_eq!( + conv_asset_ids_for_batch(&batch, &HashMap::new()), + vec![0; 5] + ); + } use super::*; use crate::feature_config::categorical_feature::{ PRODUCT_SURFACE_SEQ, PRODUCT_SURFACE_SEQ_COLUMN, PRODUCT_SURFACE_SEQ_NAME, @@ -1413,7 +1718,6 @@ mod tests { ip_modulus: 1_073_741_789, output_vocab_size: 64, num_continuous_actions: 2, - search_query_embedding_dim: 0, num_user_categorical_features: 0, num_user_bool_features: 0, num_user_float_features: 0, @@ -1423,6 +1727,7 @@ mod tests { num_post_bool_features: 0, num_post_float_features: 0, num_post_int64_features: 0, + enable_stale_post: false, }, history_seq_len: 4, candidate_seq_len: 3, @@ -1591,6 +1896,40 @@ mod tests { } } + #[test] + fn search_query_stays_one_vector() { + let mut model_config = test_model_config(1); + model_config.search_query_embedding_dim = 4; + let query = vec![1.0f32, 2.0, 3.0, 4.0]; + let candidate_set = pb::CandidateSet { + search_query_embedding: query.clone(), + candidates: vec![ + pb::TweetInfo { + tweet_id: 1000, + author_id: 2000, + ..Default::default() + }, + pb::TweetInfo { + tweet_id: 1001, + author_id: 2001, + ..Default::default() + }, + ], + ..Default::default() + }; + let cand = InputBuffer::new_with_candidates(&model_config, &candidate_set, None); + assert_eq!(cand.search_query_embeddings, query); + assert_eq!(cand.search_query_embeddings.len(), 4); + } + + #[test] + fn repeat_query_into_writes_prefix_and_keeps_tail() { + let query = [1.0f32, 2.0]; + let mut dest = vec![9.0f32; 8]; + repeat_query_into(&mut dest, &query, 2); + assert_eq!(dest, vec![1.0, 2.0, 1.0, 2.0, 9.0, 9.0, 9.0, 9.0]); + } + fn test_user_action_sequence(actions: &[(u64, u64)]) -> pb::UserActionSequence { use pb::user_action_sequence_data_container::Data; let aggregated_user_actions = actions @@ -1665,6 +2004,123 @@ mod tests { assert_eq!(0, buf.num_history); } + #[test] + fn proto_history_semantic_ids_use_ranking_plus_one_shift() { + let mut model_config = test_history_model_config(); + model_config.sid_num_levels = 3; + let sequence = Some(pb::UserActionSequence { + user_actions_data: Some(pb::UserActionSequenceDataContainer { + data: Some( + pb::user_action_sequence_data_container::Data::OrderedAggregatedUserActionsList( + pb::AggregatedUserActionList { + aggregated_user_actions: vec![ + pb::AggregatedUserAction { + tweet_info: Some(pb::TweetInfo { + tweet_id: 11, + author_id: 21, + semantic_ids: vec![0, 7, -1], + ..Default::default() + }), + ..Default::default() + }, + pb::AggregatedUserAction { + tweet_info: Some(pb::TweetInfo { + tweet_id: 12, + author_id: 22, + semantic_ids: vec![3, 4], + ..Default::default() + }), + ..Default::default() + }, + ], + ..Default::default() + }, + ), + ), + }), + ..Default::default() + }); + + let buf = InputBuffer::compute_for_item( + &model_config, + &sequence, + &pb::CandidateSet::default(), + None, + None, + None, + None, + ); + + assert_eq!( + buf.history_semantic_ids, + vec![1, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + ); + assert_eq!(buf.history_post_ids, vec![11, 12, 0, 0]); + } + + #[test] + fn stamp_semantic_ids_shifts_and_skips_wrong_arity() { + let mut dst = vec![0u16; 6]; + stamp_semantic_ids(&mut dst, 0, 3, &[0, 7, -1]); + stamp_semantic_ids(&mut dst, 1, 3, &[3, 4]); + assert_eq!(dst, vec![1, 8, 0, 0, 0, 0]); + } + + #[test] + fn columnar_history_semantic_ids_read_semantic_id_column() { + use arrow::array::{FixedSizeListArray, Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::ipc::writer::StreamWriter; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + let item = Arc::new(Field::new("item", DataType::Int32, true)); + let schema = Schema::new(vec![ + Field::new("tweetId", DataType::Int64, false), + Field::new("authorId", DataType::Int64, false), + Field::new("semanticId", DataType::FixedSizeList(item.clone(), 3), true), + ]); + let sid = FixedSizeListArray::new( + item, + 3, + Arc::new(Int32Array::from(vec![0, 7, -1, 3, 4, 5])), + None, + ); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int64Array::from(vec![11i64, 12])), + Arc::new(Int64Array::from(vec![21i64, 22])), + Arc::new(sid), + ], + ) + .unwrap(); + let mut bytes = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut bytes, batch.schema().as_ref()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + let mut model_config = test_history_model_config(); + model_config.sid_num_levels = 3; + let buf = InputBuffer::compute_from_columnar_bytes( + &model_config, + &bytes, + &pb::CandidateSet::default(), + None, + None, + None, + None, + ) + .unwrap(); + + assert_eq!( + buf.history_semantic_ids, + vec![1, 8, 0, 4, 5, 6, 0, 0, 0, 0, 0, 0] + ); + } + #[test] fn stamp_engagement_counts_writes_correct_indices() { let num_features = VIEW_COUNT_SEQ + 1; @@ -1846,4 +2302,57 @@ mod tests { assert_eq!(0, cand.int64_features[2 * n_post_int64 + slot]); } } + + #[test] + fn stale_post_14d_zeroes_and_flags_old_candidate() { + let n_post_cat = 1; + let mut model_config = test_model_config(n_post_cat); + let n_post_int64 = VIEW_COUNT_SEQ + 1; + model_config.hash_table.num_post_int64_features = n_post_int64; + model_config.hash_table.num_post_bool_features = IS_STALE_POST14D + 1; + model_config.hash_table.enable_stale_post = true; + let n_post_bool = model_config.hash_table.num_post_bool_features; + + let now_ms = now_epoch_sec() as i64 * 1000; + let tweet_id_for_age_h = |age_h: i64| -> u64 { + let creation_ms = now_ms - age_h * 3600 * 1000; + (((creation_ms - TWITTER_EPOCH_MS) << 22) as u64) & !((1u64 << 22) - 1) + }; + + let mut candidate_set = pb::CandidateSet::default(); + candidate_set.candidates.push(pb::TweetInfo { + tweet_id: tweet_id_for_age_h(400), + author_id: 2000, + fav_count: 7, + reply_count: 8, + retweet_count: 9, + quote_count: 10, + view_count: 11, + ..Default::default() + }); + candidate_set.candidates.push(pb::TweetInfo { + tweet_id: tweet_id_for_age_h(1), + author_id: 2001, + fav_count: 1, + reply_count: 2, + retweet_count: 3, + quote_count: 4, + view_count: 5, + ..Default::default() + }); + + let cand = InputBuffer::new_with_candidates(&model_config, &candidate_set, None); + + assert_eq!(0, cand.int64_features[FAV_COUNT_SEQ]); + assert_eq!(0, cand.int64_features[REPLY_COUNT_SEQ]); + assert_eq!(0, cand.int64_features[REPOST_COUNT_SEQ]); + assert_eq!(0, cand.int64_features[QUOTE_COUNT_SEQ]); + assert_eq!(0, cand.int64_features[VIEW_COUNT_SEQ]); + assert!(cand.bool_features[IS_STALE_POST14D]); + + let base1 = n_post_int64; + assert_eq!(1, cand.int64_features[base1 + FAV_COUNT_SEQ]); + assert_eq!(5, cand.int64_features[base1 + VIEW_COUNT_SEQ]); + assert!(!cand.bool_features[n_post_bool + IS_STALE_POST14D]); + } } diff --git a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml index 6e0cecf4..a2a993ad 100644 --- a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml +++ b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml @@ -33,9 +33,9 @@ url = { workspace = true } http-body = { workspace = true } http-body-util = { workspace = true } lazy_static = { workspace = true } -libc = { workspace = true } log = { workspace = true } md5 = { workspace = true } +libc = { workspace = true } memmap2 = { workspace = true } murmur3 = { workspace = true } nix = { workspace = true, features = ["fs", "uio"] } @@ -62,13 +62,16 @@ tonic-reflection = { workspace = true } xai-recsys = { workspace = true } xai-recsys-proto = { workspace = true } xai-recsys-server = { workspace = true } -xai-recsys-sid-proto = { workspace = true } xai-recsys-mm-server = { workspace = true } +xai-o2 = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] ibverbs = { workspace = true } [dev-dependencies] +rcgen = { workspace = true } +serial_test = { workspace = true } +tempfile = { workspace = true } tokio-stream = { workspace = true } [lints] diff --git a/phoenix/crates/serving/xai-recsys-engine/pyproject.toml b/phoenix/crates/serving/xai-recsys-engine/pyproject.toml index 459cb7c0..1c4d1ca0 100644 --- a/phoenix/crates/serving/xai-recsys-engine/pyproject.toml +++ b/phoenix/crates/serving/xai-recsys-engine/pyproject.toml @@ -21,5 +21,4 @@ cache-keys = [ { file = "../../../crates/serving/xai-recsys-mm-server/**" }, { file = "../../../crates/serving/xai-recsys-proto/**" }, { file = "../../../crates/serving/xai-recsys-server/**" }, - { file = "../../../crates/serving/xai-recsys-sid-proto/**" }, ] diff --git a/phoenix/crates/serving/xai-recsys-engine/src/admission.rs b/phoenix/crates/serving/xai-recsys-engine/src/admission.rs index b50d1bf5..28d5fe17 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/admission.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/admission.rs @@ -98,6 +98,11 @@ impl AdmissionController { } } + pub fn reset_estimates(&self) { + self.service_time_us.store(0, Ordering::Relaxed); + self.pipeline_time_us.store(0, Ordering::Relaxed); + } + pub fn set_enabled(&self, on: bool) { self.enabled.store(on, Ordering::Relaxed); } @@ -372,6 +377,20 @@ mod tests { assert_eq!(c.predicted_eta_pipeline_us(0), None); } + #[test] + fn reset_estimates_clears_s_and_p() { + let c = AdmissionController::new(cfg_pipeline(true)); + c.record_service_time_us(100_000); + c.record_sojourn_us(80_000, 0); + assert!(c.service_time_us() > 0); + assert!(c.pipeline_time_us() > 0); + assert!(c.should_reject(Some(1), 32)); + c.reset_estimates(); + assert_eq!(c.service_time_us(), 0); + assert_eq!(c.pipeline_time_us(), 0); + assert!(!c.should_reject(Some(1), 32)); + } + #[test] fn eta_model_parses_aliases() { assert_eq!( diff --git a/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs index 6e33ef71..225d58fe 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs @@ -3,7 +3,7 @@ use crate::{ checkpoint_proxy_metrics::*, r#const::ROOT_DIR, - emb_table::{list_entries, send_entries}, + emb_table::{apply_copy_port_http2, list_entries, send_entries}, }; use log::{info, warn}; use memmap2::MmapMut; @@ -86,7 +86,9 @@ impl CheckpointProxy { } } - async fn resolve_and_connect(&self) -> Result, io::Error> { + async fn resolve_and_connect( + &self, + ) -> Result, io::Error> { let (scheme, hostport) = if let Some(rest) = self.copy_url.strip_prefix("https://") { ("https", rest) } else { @@ -131,7 +133,7 @@ impl CheckpointProxy { futures::future::join_all(addrs.iter().zip(endpoints).map(|(addr, endpoint)| { let timeout = connect_timeout; async move { - let result = endpoint + let result = apply_copy_port_http2(endpoint) .connect_timeout(timeout) .timeout(request_timeout) .http2_keep_alive_interval(Duration::from_secs(30)) @@ -180,6 +182,50 @@ impl CheckpointProxy { ) } + async fn extra_trainer_channels(&self, addr: SocketAddr, n: usize) -> Vec { + if n == 0 { + return Vec::new(); + } + let scheme = if self.copy_url.starts_with("https://") { + "https" + } else { + "http" + }; + let tls = match crate::tls::ClientTlsOptions::from_env() { + Ok(tls) => tls, + Err(e) => { + warn!("extra trainer channels: tls config: {e}"); + return Vec::new(); + } + }; + let connect_timeout = Duration::from_secs( + std::env::var("COPY_PORT_CONNECT_TIMEOUT_SECS") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(10), + ); + let request_timeout = self.request_timeout; + let url = format!("{scheme}://{addr}"); + let results = futures::future::join_all((0..n).map(|_| { + let tls = tls.as_ref(); + let url = url.clone(); + async move { + let endpoint = crate::tls::endpoint_for(&url, tls).ok()?; + apply_copy_port_http2(endpoint) + .connect_timeout(connect_timeout) + .timeout(request_timeout) + .http2_keep_alive_interval(Duration::from_secs(30)) + .keep_alive_timeout(Duration::from_secs(10)) + .keep_alive_while_idle(true) + .connect() + .await + .ok() + } + })) + .await; + results.into_iter().flatten().collect() + } + pub async fn poll_and_download( &self, ) -> Result> { @@ -197,8 +243,9 @@ impl CheckpointProxy { self.downloading.store(false, Ordering::SeqCst); }); - let channels = self.resolve_and_connect().await?; - TRAINER_URL_COUNT.set(channels.len() as f64); + let assigned = self.resolve_and_connect().await?; + TRAINER_URL_COUNT.set(assigned.len() as f64); + let channels: Vec = assigned.iter().map(|(_, ch)| ch.clone()).collect(); let entries = list_entries(&channels, "").await?; @@ -230,7 +277,7 @@ impl CheckpointProxy { } match self - .download_checkpoint(&paths, &channels, &checkpoint_entries) + .download_checkpoint(&paths, &assigned, &checkpoint_entries) .await { Ok(()) => {} @@ -277,10 +324,10 @@ impl CheckpointProxy { async fn download_checkpoint( &self, paths: &CheckpointPaths, - channels: &[transport::Channel], + trainers: &[(SocketAddr, transport::Channel)], entries: &[Vec<(String, usize)>], ) -> Result<(), Box> { - let concurrency = self.download_concurrency; + let concurrency = self.download_concurrency.max(1); let mut seen_files = std::collections::HashSet::new(); let mut deduped_entries: Vec> = @@ -316,16 +363,34 @@ impl CheckpointProxy { ); } - let mut handles = Vec::with_capacity(channels.len()); + let mut handles = Vec::with_capacity(trainers.len()); - for (ch_idx, (channel, ch_entries)) in - channels.iter().zip(deduped_entries.into_iter()).enumerate() + for (ch_idx, ((addr, channel), ch_entries)) in + trainers.iter().zip(deduped_entries).enumerate() { if ch_entries.is_empty() { continue; } - let channel = channel.clone(); + let extras = concurrency.saturating_sub(1); + let mut pool = vec![channel.clone()]; + if extras > 0 { + let extra = self.extra_trainer_channels(*addr, extras).await; + if extra.len() < extras { + warn!( + "trainer {addr}: opened {}/{} extra connections; downloading on {}", + extra.len(), + extras, + extra.len() + 1 + ); + } + pool.extend(extra); + } + info!( + "trainer {addr}: downloading on {} HTTP/2 connection(s) (concurrency={concurrency})", + pool.len(), + ); + let staging_dir = paths.staging_dir.clone(); let grpc_prefix = paths.grpc_prefix.clone(); let verify = self.verify_checksums; @@ -334,11 +399,10 @@ impl CheckpointProxy { download_channel_files( &staging_dir, &grpc_prefix, - channel, + pool, &ch_entries, ch_idx, verify, - concurrency, ) .await }); @@ -483,11 +547,10 @@ impl CheckpointProxy { async fn download_channel_files( staging_dir: &Path, grpc_prefix: &str, - channel: transport::Channel, + channels: Vec, entries: &[(String, usize)], ch_idx: usize, _verify_checksums: bool, - concurrency: usize, ) -> Result<(), Box> { let t0 = Instant::now(); let mut total_bytes: usize = 0; @@ -503,8 +566,10 @@ async fn download_channel_files( } } - let sem = Arc::new(tokio::sync::Semaphore::new(concurrency)); + let n_conns = channels.len().max(1); + let sem = Arc::new(tokio::sync::Semaphore::new(n_conns)); let mut handles = Vec::new(); + let mut chunk_idx = 0usize; for (name, size) in entries { if *size == 0 { @@ -548,7 +613,8 @@ async fn download_channel_files( let mut offset = 0usize; while offset < *size { let this_chunk = cmp::min(chunk_size, *size - offset); - let channel = channel.clone(); + let channel = channels[chunk_idx % n_conns].clone(); + chunk_idx += 1; let name = name.clone(); let sem = sem.clone(); let mmap = mmap.clone(); @@ -580,7 +646,7 @@ async fn download_channel_files( CHANNEL_DOWNLOAD_RATE_BYTES_PER_SEC.observe(total_bytes as f64 / elapsed.as_secs_f64()); } info!( - "channel {}: downloaded {} files, {:.2} GB in {:.1}s ({:.2} GB/s) [concurrency={}]", + "channel {}: downloaded {} files, {:.2} GB in {:.1}s ({:.2} GB/s) [connections={}]", ch_idx, entries.len(), total_bytes as f64 * 1e-9, @@ -590,7 +656,7 @@ async fn download_channel_files( } else { 0.0 }, - concurrency, + n_conns, ); Ok(()) @@ -647,7 +713,7 @@ fn assign_trainer_channels( num_trainers: Option, server_index: usize, num_servers: usize, -) -> Result, io::Error> { +) -> Result, io::Error> { let total_trainers = if let Some(expected) = num_trainers { if connected.len() < expected { return Err(io::Error::new( @@ -685,7 +751,6 @@ fn assign_trainer_channels( .collect(); let my_ips: Vec<_> = my_slice.iter().map(|(addr, _)| addr.to_string()).collect(); - let my_channels: Vec = my_slice.into_iter().map(|(_, ch)| ch).collect(); info!( "{} trainers reachable (sorted by IP), this server (index {}) mirrors {}-{}: {:?}", @@ -696,7 +761,7 @@ fn assign_trainer_channels( my_ips, ); - Ok(my_channels) + Ok(my_slice) } fn free_bytes(path: &str) -> io::Result { diff --git a/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_store.rs b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_store.rs index 3cc6ea99..c7b0bf4b 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_store.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_store.rs @@ -2,7 +2,7 @@ // Copyright 2026 X.AI Corp. use async_trait::async_trait; use bytes::Bytes; -use log::{error, info}; +use log::{error, info, warn}; use object_store::aws::AmazonS3Builder; use object_store::chunked::ChunkedStore; use object_store::gcp::{GcpCredential, GoogleCloudStorageBuilder}; @@ -18,13 +18,33 @@ use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; use tokio::sync::Semaphore; +use xai_o2::{O2ClientBuilder, base_client::BaseO2Client}; + const MAX_CONCURRENT: usize = 16; const MAX_RETRIES: usize = 6; const MULTIPART_THRESHOLD: usize = 100 * 1024 * 1024; -const MULTIPART_PART_SIZE: usize = 16 * 1024 * 1024; +const MULTIPART_PART_SIZE: usize = 32 * 1024 * 1024; + +const MAX_CONCURRENT_WRITE_REQUESTS: usize = 1; + +fn env_usize(name: &str, default: usize) -> usize { + match env::var(name) { + Ok(raw) => match raw.trim().parse::() { + Ok(v) => v, + Err(_) => { + warn!( + "[checkpoint-store] invalid {}={:?}, using default {}", + name, raw, default + ); + default + } + }, + Err(_) => default, + } +} pub struct CheckpointStore { client: Arc, @@ -153,6 +173,34 @@ impl CheckpointStore { } } +pub struct XaiO2Store { + client: Arc, +} + +impl XaiO2Store { + pub async fn put(&self, path: &ObjectPath, data: Bytes, label: &str) -> Result<(), String> { + self.client + .put(path.as_ref(), data) + .await + .map_err(|e| format!("PUT {} failed: {}", label, e)) + } +} + +#[derive(Clone)] +pub enum Store { + ObjectStore(Arc), + XaiO2(Arc), +} + +impl Store { + pub async fn put(&self, path: &ObjectPath, data: Bytes, label: &str) -> Result<(), String> { + match self { + Store::ObjectStore(s) => s.put(path, data, label).await, + Store::XaiO2(s) => s.put(path, data, label).await, + } + } +} + #[derive(Debug)] struct WorkloadIdentityCredentialProvider { inner: google_cloud_auth::credentials::AccessTokenCredentials, @@ -222,6 +270,8 @@ fn is_external_account_credential(path: &str) -> bool { static STORES: OnceLock>>> = OnceLock::new(); +static XAI_O2_STORES: OnceLock>>> = OnceLock::new(); + pub fn get_checkpoint_store( url: &str, runtime: Option<&tokio::runtime::Runtime>, @@ -292,8 +342,9 @@ pub fn get_checkpoint_store( ..Default::default() }; - let endpoint = - env::var("O2_ENDPOINT_URL").unwrap_or_else(|_| "http://o2.example.invalid".into()); + let endpoint = env::var("O2_ENDPOINT_URL").map_err(|_| { + "O2 endpoint not set: set XAI_O2_ENDPOINT_URL or O2_ENDPOINT_URL".to_string() + })?; let access_key = env::var("O2_ACCESS_KEY_ID").unwrap_or_else(|_| "anonymous".to_string()); let secret_key = env::var("O2_SECRET_ACCESS_KEY").unwrap_or_else(|_| "anonymous".to_string()); @@ -329,3 +380,123 @@ pub fn get_checkpoint_store( Ok((store, ObjectPath::from(path))) } + +fn is_o2_scheme(url: &str) -> bool { + url.starts_with("s3://") || url.starts_with("s3-o2://") +} + +fn use_xai_o2_for(url: &str) -> bool { + if !is_o2_scheme(url) { + return false; + } + match env::var("XAI_RECSYS_S3_BACKEND") { + Ok(raw) => { + let v = raw.trim(); + if v.eq_ignore_ascii_case("xai-o2") { + true + } else { + if !v.is_empty() { + warn!( + "[checkpoint-store] unknown XAI_RECSYS_S3_BACKEND={:?}, \ + using object_store (set to \"xai-o2\" for the xai-o2 backend)", + raw + ); + } + false + } + } + Err(_) => false, + } +} + +pub fn get_store( + url: &str, + runtime: Option<&tokio::runtime::Runtime>, +) -> Result<(Store, ObjectPath), Box> { + if use_xai_o2_for(url) { + let (store, path) = get_xai_o2_store(url)?; + return Ok((Store::XaiO2(store), path)); + } + let (store, path) = get_checkpoint_store(url, runtime)?; + Ok((Store::ObjectStore(store), path)) +} + +fn get_xai_o2_store( + url: &str, +) -> Result<(Arc, ObjectPath), Box> { + let parsed = url::Url::parse(url)?; + let bucket = parsed + .host_str() + .ok_or_else(|| format!("No bucket in URL: {}", url))? + .to_string(); + let path = parsed.path().trim_start_matches('/'); + + let cache_key = format!("s3://{}", bucket); + let stores = XAI_O2_STORES.get_or_init(|| Mutex::new(HashMap::new())); + let mut map = stores.lock().map_err(|e| format!("lock poisoned: {}", e))?; + if let Some(store) = map.get(&cache_key) { + return Ok((store.clone(), ObjectPath::from(path))); + } + + let store = Arc::new(XaiO2Store { + client: build_xai_o2_client(&bucket)?, + }); + map.insert(cache_key, store.clone()); + Ok((store, ObjectPath::from(path))) +} + +fn build_xai_o2_client( + bucket: &str, +) -> Result, Box> { + let endpoint = env::var("O2_ENDPOINT_URL").map_err(|_| { + "O2 endpoint not set: set XAI_O2_ENDPOINT_URL or O2_ENDPOINT_URL".to_string() + })?; + let access_key = env::var("O2_ACCESS_KEY_ID").unwrap_or_else(|_| "anonymous".to_string()); + let secret_key = env::var("O2_SECRET_ACCESS_KEY").unwrap_or_else(|_| "anonymous".to_string()); + let skip_signature = access_key == "anonymous" && secret_key == "anonymous"; + + let max_write_reqs = env_usize( + "XAI_RECSYS_O2_MAX_CONCURRENT_WRITE_REQUESTS", + MAX_CONCURRENT_WRITE_REQUESTS, + ) + .max(1); + let write_chunk_concurrency = + env_usize("XAI_RECSYS_O2_WRITE_CHUNK_CONCURRENCY", MAX_CONCURRENT).max(1); + let write_chunk_size_mib = env_usize( + "XAI_RECSYS_O2_WRITE_CHUNK_SIZE_MIB", + MULTIPART_PART_SIZE / (1024 * 1024), + ) + .max(1); + + info!( + "[checkpoint-store] Creating xai-o2 client bucket={} \ + max_concurrent_write_requests={} write_chunk_concurrency={} \ + write_chunk_size_mib={}", + bucket, max_write_reqs, write_chunk_concurrency, write_chunk_size_mib + ); + + let mut builder = O2ClientBuilder::default(); + builder + .endpoint(endpoint) + .bucket(bucket) + .prefix("") + .write_chunk_size_mib(write_chunk_size_mib) + .write_chunk_concurrency(write_chunk_concurrency) + .max_concurrent_write_requests(max_write_reqs) + .retry_max_times(MAX_RETRIES) + .retry_min_delay(Duration::from_secs(1)) + .retry_max_delay(Duration::from_secs(32)) + .timeout(Duration::from_secs(300)) + .io_timeout(Duration::from_secs(300)) + .trace_always(false); + if skip_signature { + builder.allow_anonymous(true); + } else { + builder + .access_key_id(access_key) + .secret_access_key(secret_key); + } + builder + .build() + .map_err(|e| -> Box { e.into() }) +} diff --git a/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs index 4ae210e2..27b75aa6 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs @@ -2,6 +2,7 @@ // Copyright 2026 X.AI Corp. use std::cmp; use std::collections::{BTreeMap, HashMap, HashSet}; +use std::hash::{Hash, Hasher}; use std::mem; use std::slice; #[cfg(target_os = "linux")] @@ -233,6 +234,16 @@ async fn connect_and_list( type TransferFuture = BoxFuture<'static, (usize, u32)>; type DenseDownloadPlan = (Vec, Vec, Vec); +async fn join_transfers( + futures: impl IntoIterator, +) -> Result, CopyPortError> { + join_all(futures.into_iter().map(tokio::task::spawn)) + .await + .into_iter() + .collect::, _>>() + .map_err(|e| CopyPortError::Other(format!("copy_port download task join: {e}"))) +} + async fn run_downloads( futures: Vec, rate_limit_bytes_per_sec: Option, @@ -248,7 +259,7 @@ async fn run_downloads( .max(1); join_rate_limited(futures, limit, max_c).await } - _ => Ok(join_all(futures).await), + _ => join_transfers(futures).await, } } @@ -268,7 +279,7 @@ async fn join_rate_limited( break; } let batch_size = batch.len(); - let batch_results = join_all(batch).await; + let batch_results = join_transfers(batch).await?; let failed = batch_results .iter() .filter(|r| r.0 == TRANSFER_FAILED_SENTINEL) @@ -440,10 +451,6 @@ pub async fn download_dense_and_embeddings( let index = index_dense_listing(&prefix, &entries[channel_idx]); let (futures, sizes, checksums_buf) = build_dense_downloads(&channels[channel_idx], &index, tensors)?; - log::info!( - "copy_port: downloading dense weights prefix={prefix} futures={}", - futures.len() - ); let results = run_downloads(futures, rate_limit_bytes_per_sec, max_concurrent_downloads).await?; sfence_after_download(); @@ -454,9 +461,17 @@ pub async fn download_dense_and_embeddings( .ok() .and_then(|v| v.get("created_timestamp")?.as_f64()) .unwrap_or(0.0); + let bytes: u64 = tensors.iter().map(|t| t.buf.len() as u64).sum(); + let secs = t_all.elapsed().as_secs_f64(); + let gbs = if secs > 0.0 { + bytes as f64 / secs / 1e9 + } else { + 0.0 + }; log::info!( - "copy_port: dense weights loaded prefix={prefix} in {:.2}s", - t_all.elapsed().as_secs_f64() + "copy_port: dense weights loaded bytes={bytes} in {:.2}s ({:.2} GB/s)", + secs, + gbs ); let prefix_slash = if prefix.ends_with('/') { @@ -477,10 +492,6 @@ pub async fn download_dense_and_embeddings( }) .collect(); - log::info!( - "copy_port: downloading emb_table prefix={prefix_slash} bytes={}", - emb.len() - ); let t_emb = Instant::now(); let emb_ck = download_sharded_with_channels( &channels, @@ -492,10 +503,17 @@ pub async fn download_dense_and_embeddings( max_concurrent_downloads, ) .await?; + let secs = t_emb.elapsed().as_secs_f64(); + let gbs = if secs > 0.0 { + emb.len() as f64 / secs / 1e9 + } else { + 0.0 + }; log::info!( - "copy_port: emb_table loaded bytes={} in {:.2}s", + "copy_port: emb_table loaded bytes={} in {:.2}s ({:.2} GB/s)", emb.len(), - t_emb.elapsed().as_secs_f64() + secs, + gbs ); let pe_ck = if let Some(pe_buf) = pe.filter(|b| !b.is_empty()) { @@ -541,10 +559,12 @@ async fn download_sharded_with_channels( ) -> Result { let layout = ShardedLayout::from_listing(name, entries)?; layout.check_buffer_size(name, buf.len())?; - let (futures, expected) = spawn_sharded_downloads(&layout, prefix, name, channels, buf).await?; + let (futures, expected, schedule) = + spawn_sharded_downloads(&layout, prefix, name, channels, buf).await?; let concurrent = max_concurrent_downloads.or(Some((channels.len() / 2).max(1))); let results = run_downloads(futures, rate_limit_bytes_per_sec, concurrent).await?; sfence_after_download(); + let results = results_in_piece_order(results, &schedule)?; combine_transfer_checksums(&results, &expected) } @@ -592,11 +612,12 @@ async fn download_embedding_table_with_conns( .await; } - let (futures, expected) = + let (futures, expected, schedule) = spawn_sharded_downloads(&layout, &prefix, name, &channels, buf).await?; let concurrent = max_concurrent_downloads.or(Some((channels.len() / 2).max(1))); let results = run_downloads(futures, rate_limit_bytes_per_sec, concurrent).await?; sfence_after_download(); + let results = results_in_piece_order(results, &schedule)?; combine_transfer_checksums(&results, &expected) } @@ -789,6 +810,52 @@ pub(crate) fn classify_shard_ownership( ))) } +pub(crate) fn shuffle_sharded_schedule(n: usize, name: &str) -> Vec { + let mut schedule: Vec = (0..n).collect(); + if n <= 1 { + return schedule; + } + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + if let Ok(id) = std::env::var("POD_NAME").or_else(|_| std::env::var("HOSTNAME")) { + id.hash(&mut hasher); + } + name.hash(&mut hasher); + let mut rng = rand::rngs::StdRng::seed_from_u64(hasher.finish()); + schedule.shuffle(&mut rng); + schedule +} + +pub(crate) fn restore_piece_order(shuffled: Vec, schedule: &[usize]) -> Option> { + if shuffled.len() != schedule.len() { + return None; + } + let mut out: Vec> = (0..schedule.len()).map(|_| None).collect(); + for (item, &orig) in shuffled.into_iter().zip(schedule) { + if orig >= out.len() || out[orig].is_some() { + return None; + } + out[orig] = Some(item); + } + out.into_iter().collect() +} + +fn results_in_piece_order( + results: Vec<(usize, u32)>, + schedule: &[usize], +) -> Result, CopyPortError> { + if results + .iter() + .any(|(sent, _)| *sent == TRANSFER_FAILED_SENTINEL) + { + return Err(CopyPortError::TransferFailed( + "gRPC/RDMA transfer failed (check Rust logs)".into(), + )); + } + restore_piece_order(results, schedule).ok_or_else(|| { + CopyPortError::TransferFailed("shuffled download result count/order mismatch".into()) + }) +} + pub(crate) fn combine_transfer_checksums( results: &[(usize, u32)], expected: &[usize], @@ -876,7 +943,7 @@ async fn spawn_sharded_downloads( name: &str, channels: &[Channel], buf: &mut [u8], -) -> Result<(Vec, Vec), CopyPortError> { +) -> Result<(Vec, Vec, Vec), CopyPortError> { let name_prefix = format!("{name}/c/"); let piece = layout.piece_bytes; @@ -911,6 +978,18 @@ async fn spawn_sharded_downloads( } }; + let schedule = match &layout.ownership { + ShardOwnership::Sharded { .. } => shuffle_sharded_schedule(ranges.len(), name), + ShardOwnership::Replicated { .. } => (0..ranges.len()).collect(), + }; + if matches!(layout.ownership, ShardOwnership::Sharded { .. }) && ranges.len() > 1 { + log::info!( + "copy_port: shard schedule shuffled name={name} n={} first={:?}", + schedule.len(), + &schedule[..schedule.len().min(8)] + ); + } + #[cfg(target_os = "linux")] let (contexts, devicez, mrx) = { use crate::emb_table::register_tensor_for_download; @@ -952,11 +1031,10 @@ async fn spawn_sharded_downloads( (contexts, devicez, mrx) }; + let expected: Vec = ranges.iter().map(|&(_, a, b)| (b - a) * piece).collect(); let mut futures = Vec::with_capacity(ranges.len()); - let mut expected = Vec::with_capacity(ranges.len()); - for (i, &(idx, a, b)) in ranges.iter().enumerate() { - #[cfg(not(target_os = "linux"))] - let _ = i; + for &i in &schedule { + let (idx, a, b) = ranges[i]; let n = b - a; let slice = &mut buf[a * piece..b * piece]; let slice: &'static mut [u8] = @@ -973,9 +1051,8 @@ async fn spawn_sharded_downloads( (devicez[i].clone(), contexts.clone(), mrx[i].clone()), )); futures.push(fut); - expected.push(n * piece); } - Ok((futures, expected)) + Ok((futures, expected, schedule)) } #[cfg(test)] @@ -1119,6 +1196,49 @@ mod tests { assert!(combine_transfer_checksums(&[(TRANSFER_FAILED_SENTINEL, 1)], &[64]).is_err()); } + #[test] + fn shuffle_schedule_is_stable_permutation() { + let a = shuffle_sharded_schedule(32, "emb_table"); + let b = shuffle_sharded_schedule(32, "emb_table"); + assert_eq!(a, b); + let mut sorted = a.clone(); + sorted.sort_unstable(); + assert_eq!(sorted, (0..32).collect::>()); + let c = shuffle_sharded_schedule(32, "post_embeddings"); + assert_ne!(a, c); + assert_ne!(a, (0..32).collect::>()); + } + + #[test] + fn restore_piece_order_inverts_schedule() { + let schedule = vec![3, 0, 2, 1]; + let shuffled = vec!['d', 'a', 'c', 'b']; + assert_eq!( + restore_piece_order(shuffled, &schedule).unwrap(), + vec!['a', 'b', 'c', 'd'] + ); + assert!(restore_piece_order(vec![1, 2], &[0, 1, 2]).is_none()); + assert!(restore_piece_order(vec![1, 2, 3], &[0, 0, 1]).is_none()); + } + + #[test] + fn combine_after_restore_matches_piece_order() { + let piece_order = vec![(10, 11u32), (10, 22), (10, 33)]; + let expected = vec![10usize, 10, 10]; + let direct = combine_transfer_checksums(&piece_order, &expected).unwrap(); + let schedule = vec![2, 0, 1]; + let shuffled = vec![piece_order[2], piece_order[0], piece_order[1]]; + let restored = restore_piece_order(shuffled.clone(), &schedule).unwrap(); + assert_eq!( + combine_transfer_checksums(&restored, &expected).unwrap(), + direct + ); + assert_ne!( + combine_transfer_checksums(&shuffled, &expected).unwrap(), + direct + ); + } + #[test] fn replicated_send_ranges_chunk_within_blocks() { assert_eq!( @@ -1206,6 +1326,150 @@ mod tests { let partial = vec![entries[0].clone(), vec![]]; assert!(choose_prefix(Some("elapsed_samples_1/run"), &partial).is_err()); } + + fn tracking_downloads( + n: usize, + bytes: usize, + hold: Duration, + in_flight: std::sync::Arc, + peak: std::sync::Arc, + starts: std::sync::Arc>>>, + ) -> Vec { + (0..n) + .map(|i| { + let in_flight = in_flight.clone(); + let peak = peak.clone(); + let starts = starts.clone(); + Box::pin(async move { + starts.lock().unwrap()[i] = Some(Instant::now()); + let cur = in_flight.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1; + peak.fetch_max(cur, std::sync::atomic::Ordering::SeqCst); + tokio::time::sleep(hold).await; + in_flight.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + (bytes, i as u32) + }) as TransferFuture + }) + .collect() + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rate_limit_caps_in_flight() { + let in_flight = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let peak = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let starts = std::sync::Arc::new(std::sync::Mutex::new(vec![None; 6])); + let futures = tracking_downloads( + 6, + 1, + Duration::from_millis(80), + in_flight, + peak.clone(), + starts, + ); + let results = join_rate_limited(futures, 1 << 40, 2).await.unwrap(); + assert_eq!( + results.iter().map(|r| r.1).collect::>(), + vec![0, 1, 2, 3, 4, 5] + ); + assert_eq!(peak.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rate_limit_paces_between_batches() { + let in_flight = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let peak = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let starts = std::sync::Arc::new(std::sync::Mutex::new(vec![None; 4])); + let bytes = 200 * 1024; + let rate = 400 * 1024; + let t0 = Instant::now(); + let results = run_downloads( + tracking_downloads( + 4, + bytes, + Duration::from_millis(20), + in_flight, + peak.clone(), + starts.clone(), + ), + Some(rate), + Some(2), + ) + .await + .unwrap(); + let elapsed = t0.elapsed(); + assert_eq!(results.len(), 4); + assert_eq!(peak.load(std::sync::atomic::Ordering::SeqCst), 2); + + let starts = starts.lock().unwrap(); + let s: Vec = starts.iter().map(|t| t.expect("started")).collect(); + let first_batch_start = s[0].min(s[1]); + let second_batch_start = s[2].min(s[3]); + let between = second_batch_start.saturating_duration_since(first_batch_start); + assert!( + between >= Duration::from_millis(700), + "second batch started {between:?} after the first" + ); + + let total_bytes = (4 * bytes) as f64; + let min_elapsed = Duration::from_secs_f64(total_bytes / rate as f64 * 0.85); + assert!( + elapsed >= min_elapsed, + "elapsed {elapsed:?} is below the {min_elapsed:?} floor for {total_bytes} bytes at {rate} B/s" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rate_limit_skips_remaining_on_sentinel() { + let started = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let futures: Vec = (0..4) + .map(|i| { + let started = started.clone(); + Box::pin(async move { + started.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if i == 1 { + (TRANSFER_FAILED_SENTINEL, 0) + } else { + (1_000_000, i as u32) + } + }) as TransferFuture + }) + .collect(); + let t0 = Instant::now(); + let results = join_rate_limited(futures, 1, 2).await.unwrap(); + assert!( + t0.elapsed() < Duration::from_secs(2), + "must not pace a failed batch" + ); + assert_eq!(results.len(), 2); + assert!(results.iter().any(|r| r.0 == TRANSFER_FAILED_SENTINEL)); + assert_eq!(started.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn unlimited_path_spawns_all() { + let in_flight = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let peak = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let starts = std::sync::Arc::new(std::sync::Mutex::new(vec![None; 4])); + let results = run_downloads( + tracking_downloads( + 4, + 1, + Duration::from_millis(80), + in_flight, + peak.clone(), + starts, + ), + Some(0), + Some(1), + ) + .await + .unwrap(); + assert_eq!(results.len(), 4); + assert_eq!( + peak.load(std::sync::atomic::Ordering::SeqCst), + 4, + "rate_limit=0 must ignore max_concurrent and spawn every future" + ); + } } #[cfg(all(test, target_os = "linux", feature = "rdma-tests"))] diff --git a/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs index 64a73542..002c53d0 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs @@ -469,6 +469,16 @@ fn parse_window_mib(v: Option<&str>, default_mib: u32) -> Option { (mib > 0).then(|| mib.min(1024) << 20) } +pub(crate) fn apply_copy_port_http2(endpoint: transport::Endpoint) -> transport::Endpoint { + let mut endpoint = endpoint + .initial_connection_window_size(*H2_CONN_WINDOW) + .initial_stream_window_size(*H2_STREAM_WINDOW); + if *H2_ADAPTIVE_WINDOW { + endpoint = endpoint.http2_adaptive_window(true); + } + endpoint +} + pub(crate) async fn get_channels(target: String) -> Result, Status> { if target.is_empty() { return Ok(Vec::new()); @@ -491,13 +501,8 @@ pub(crate) async fn get_channels(target: String) -> Result { + handle: Option>, +} + +impl JoinOnDrop { + fn new(handle: tokio::task::JoinHandle) -> Self { + Self { + handle: Some(handle), + } + } + + async fn join(mut self) -> Result { + self.handle.take().expect("join").await + } +} + +impl Drop for JoinOnDrop { + fn drop(&mut self) { + let Some(handle) = self.handle.take() else { + return; + }; + let _ = tokio::task::block_in_place(|| tokio::runtime::Handle::current().block_on(handle)); + } +} + +fn install_grpc_chunk(dest: &mut [u8], pos: &mut usize, checksum: &mut u32, chunk: &[u8]) { + let n = chunk.len(); + if *pos + n <= dest.len() { + #[cfg(target_arch = "x86_64")] + unsafe { + copy_nontemporal(&mut dest[*pos..*pos + n], chunk); + } + #[cfg(not(target_arch = "x86_64"))] + dest[*pos..*pos + n].copy_from_slice(chunk); + } + *pos += n; + adler32_combine(checksum, adler32(&chunk), n); +} + pub(crate) async fn send_entries( mut channel: transport::Channel, names: Vec>, @@ -636,23 +680,33 @@ pub(crate) async fn send_entries( }; let body = freeze(bytes); - let mut pos = 0; - let mut checksum = 1; + const GRPC_COPY_PIECE: usize = 4 << 20; + let dest_ptr = buf.as_mut_ptr() as usize; + let dest_len = buf.len(); + let (tx, rx) = std::sync::mpsc::sync_channel::(2); + let copy_join = JoinOnDrop::new(tokio::task::spawn_blocking(move || { + let dest = unsafe { std::slice::from_raw_parts_mut(dest_ptr as *mut u8, dest_len) }; + let mut pos = 0; + let mut checksum = 1u32; + while let Ok(chunk) = rx.recv() { + install_grpc_chunk(dest, &mut pos, &mut checksum, &chunk); + } + (pos, checksum) + })); let mut endpoints = Vec::new(); let mut rmrs = Vec::new(); let mut use_rdma = Vec::new(); let proto = proto![ - (1, |_, data: &[u8]| { - if pos + data.len() <= buf.len() { - #[cfg(target_arch = "x86_64")] - unsafe { - copy_nontemporal(&mut buf[pos..pos + data.len()], data); + (1, move |_, data: &[u8]| { + let mut off = 0; + while off < data.len() { + let n = (data.len() - off).min(GRPC_COPY_PIECE); + let piece = bytes::Bytes::copy_from_slice(&data[off..off + n]); + if tx.send(piece).is_err() { + break; } - #[cfg(not(target_arch = "x86_64"))] - buf[pos..pos + data.len()].copy_from_slice(data); + off += n; } - pos += data.len(); - adler32_combine(&mut checksum, adler32(&data), data.len()); }), (2, repeated_bytes(&mut endpoints)), (3, repeated_bytes(&mut rmrs)), @@ -660,8 +714,13 @@ pub(crate) async fn send_entries( ]; if let Err(e) = ready_call_parse(SEND, proto, body, &mut channel).await { log::error!("gRPC error: {e}"); + let _ = copy_join.join().await; return (TRANSFER_FAILED_SENTINEL, 0); } + let Ok((pos, checksum)) = copy_join.join().await else { + return (TRANSFER_FAILED_SENTINEL, 0); + }; + use_rdma.resize(sizes.len(), 0); let f = |acc, (&x, &y)| if y != 0 { acc + x } else { acc }; let size = sizes.iter().zip(&use_rdma).fold(0, f); @@ -673,7 +732,20 @@ pub(crate) async fn send_entries( .await { Ok(_) => { - adler32_combine(&mut checksum, adler32(&&buf[pos..]), size); + let Ok(checksum) = JoinOnDrop::new(tokio::task::spawn_blocking(move || { + if size == 0 || pos >= dest_len { + return checksum; + } + let dest = unsafe { std::slice::from_raw_parts(dest_ptr as *const u8, dest_len) }; + let mut c = checksum; + adler32_combine(&mut c, adler32(&&dest[pos..]), size); + c + })) + .join() + .await + else { + return (TRANSFER_FAILED_SENTINEL, 0); + }; (pos + size, checksum) } Err(e) => { @@ -976,6 +1048,7 @@ pub fn load_tensor_no_resharding<'py>( ) -> PyResult> { use crate::copy_port_client::{ ShardOwnership, classify_shard_ownership, peer_send_max_pieces, replicated_send_ranges, + restore_piece_order, shuffle_sharded_schedule, }; let runtime = runtime::Builder::new_multi_thread() @@ -1057,6 +1130,18 @@ pub fn load_tensor_no_resharding<'py>( ), }; + let schedule = match &ownership { + ShardOwnership::Sharded { .. } => shuffle_sharded_schedule(ranges.len(), &tensor_name), + ShardOwnership::Replicated { .. } => (0..ranges.len()).collect(), + }; + if matches!(ownership, ShardOwnership::Sharded { .. }) && ranges.len() > 1 { + log::info!( + "copy_port: shard schedule shuffled name={tensor_name} n={} first={:?}", + schedule.len(), + &schedule[..schedule.len().min(8)] + ); + } + #[cfg(target_os = "linux")] let (contexts, devicez, mrx) = { let contexts = Arc::new(if matches!(ownership, ShardOwnership::Sharded { .. }) { @@ -1101,11 +1186,13 @@ pub fn load_tensor_no_resharding<'py>( (contexts, devicez, mrx) }; + let expected: Vec = ranges + .iter() + .map(|&(_, a, b)| (b - a) * shard_size) + .collect(); let mut futures = Vec::>::with_capacity(ranges.len()); - let mut expected = Vec::with_capacity(ranges.len()); - for (i, &(idx, a, b)) in ranges.iter().enumerate() { - #[cfg(not(target_os = "linux"))] - let _ = i; + for &i in &schedule { + let (idx, a, b) = ranges[i]; let n = b - a; let slice = &mut tensor_slice[a * shard_size..b * shard_size]; let slice: &'static mut [u8] = @@ -1121,7 +1208,6 @@ pub fn load_tensor_no_resharding<'py>( #[cfg(target_os = "linux")] (devicez[i].clone(), contexts.clone(), mrx[i].clone()), ))); - expected.push(n * shard_size); } let mut checksum = 1; @@ -1135,6 +1221,9 @@ pub fn load_tensor_no_resharding<'py>( _ => block_on(&runtime, futures, None), }) .ok_or_else(|| PyOSError::new_err("copy_port timed out waiting for tensor shards"))?; + let results = restore_piece_order(results, &schedule).ok_or_else(|| { + PyOSError::new_err("copy_port shuffled download result count/order mismatch") + })?; #[cfg(target_arch = "x86_64")] unsafe { std::arch::x86_64::_mm_sfence(); @@ -1161,63 +1250,59 @@ pub fn load_tensor_no_resharding<'py>( Ok((checksum,).into_pyobject(py)?.into()) } -#[pyfunction] -pub fn load_tensor<'py>( - py: Python<'py>, - path: Bound<'py, PyString>, - urls: Bound<'py, PyString>, - shard_sources: Bound<'py, PyList>, - mut tensor: PyReadwriteArray1<'py, u8>, +pub fn load_tensor_into( + path: &str, + urls: &str, + shard_sources: &[(String, String, usize, usize)], + tensor: &mut [u8], row_size: usize, num_row_segments: usize, -) -> PyResult<()> { +) -> Result<(), String> { #[cfg(target_os = "linux")] let oflags = OFlag::O_RDONLY | OFlag::O_DIRECT; #[cfg(not(target_os = "linux"))] let oflags = OFlag::O_RDONLY; - let tensor_slice = tensor.as_slice_mut()?; - let tensor_size = tensor_slice.len(); - if num_row_segments == 0 || tensor_size % num_row_segments != 0 { - return Err(PyValueError::new_err(format!( + let tensor_size = tensor.len(); + if num_row_segments == 0 || !tensor_size.is_multiple_of(num_row_segments) { + return Err(format!( "tensor size {tensor_size} must be divisible by num_row_segments {num_row_segments}" - ))); + )); } let row_segment_size = tensor_size / num_row_segments; if row_size == 0 || !row_segment_size.is_multiple_of(row_size) { - return Err(PyValueError::new_err(format!( - "row_segment_size {row_segment_size} must be divisible by row_size {row_size}", - ))); + return Err(format!( + "row_segment_size {row_segment_size} must be divisible by row_size {row_size}" + )); } - - if shard_sources.len() % num_row_segments != 0 { - return Err(PyValueError::new_err(format!( + if !shard_sources.len().is_multiple_of(num_row_segments) { + return Err(format!( "shard sources size {} must be divisible by num_row_segments {num_row_segments}", shard_sources.len() - ))); + )); } let num_shards = shard_sources.len() / num_row_segments; if !PAGE_SIZE.is_multiple_of(row_size) { - return Err(PyValueError::new_err(format!( + return Err(format!( "page size {PAGE_SIZE} must be divisible by row_size {row_size}" - ))); + )); } if num_shards == 0 || !row_size.is_multiple_of(num_shards) { - return Err(PyValueError::new_err(format!( - "row_size {row_size} must be divisible by num_shards {num_shards}", - ))); + return Err(format!( + "row_size {row_size} must be divisible by num_shards {num_shards}" + )); } let w_s = row_size / num_shards; if w_s < WIDTH && !WIDTH.is_multiple_of(w_s) { - return Err(PyValueError::new_err(format!( + return Err(format!( "{WIDTH} must be divisible by (row_size / num_shards) {w_s}" - ))); + )); } let delta_shards = cmp::max(1, WIDTH / w_s); if !num_shards.is_multiple_of(delta_shards) { - return Err(PyValueError::new_err(format!( + return Err(format!( "num_shards {num_shards} must be divisible by delta_shards {delta_shards}" - ))); + )); } let num_blocks = num_shards / delta_shards; let shard_size = row_segment_size / num_shards; @@ -1234,14 +1319,16 @@ pub fn load_tensor<'py>( let runtime = runtime::Builder::new_multi_thread() .enable_all() .build() - .unwrap(); + .map_err(|e| format!("tokio runtime: {e}"))?; - let path = path.to_string(); let prefix: Vec<_> = path.split('/').rev().take(3).collect(); + if prefix.len() < 3 { + return Err(format!("checkpoint path too short: {path}")); + } let prefix = format!("{}/{}/", prefix[2], prefix[1]); let (channels, entries) = runtime .block_on(get_channels_list_entries(urls.to_string(), prefix.clone())) - .map_err(|s| PyOSError::new_err(format!("gRPC error: {}", s.message())))?; + .map_err(|s| format!("gRPC error: {}", s.message()))?; let mut channel_indexes = HashMap::new(); for (channel_idx, inner) in entries.into_iter().enumerate() { @@ -1249,171 +1336,193 @@ pub fn load_tensor<'py>( channel_indexes.insert(key, (channel_idx, value)); } } - for t in shard_sources.iter() { - let err0 = - || PyTypeError::new_err("shard_sources entry must be (name, fname, offset, size)"); - let tuple = t.downcast::().map_err(|_| err0())?; - if tuple.len() != 4 { - return Err(err0()); + for (key, fname, offset, size) in shard_sources { + if *size != shard_size { + return Err(format!( + "bad shard size of name {key}: wanted {shard_size}, got {size}" + )); } - - let k = tuple.get_item(0)?; - let key: String = k.extract().map_err(|_| { - PyTypeError::new_err(format!( - "shard_sources name {} must be a string", - k.repr() - .map(|s| s.to_string()) - .unwrap_or("".to_string()) - )) - })?; - let err = || { - PyTypeError::new_err(format!( - "shard_sources entry must be (name, fname, offset, size) for name {key}" - )) - }; - - let fname: String = tuple.get_item(1)?.extract().map_err(|_| err())?; let empty_fname = fname.is_empty(); let fname = format!("{path}/{fname}"); - let offset: usize = tuple.get_item(2)?.extract().map_err(|_| err())?; - let size: usize = tuple.get_item(3)?.extract().map_err(|_| err())?; - if size != shard_size { - return Err(PyValueError::new_err(format!( - "bad shard size of name {key}: wanted {shard_size}, got {size}" - ))); - } - - let idx_or_fd: Result = if let Some(&(idx, sz)) = channel_indexes.get(&key) - { - if sz != size { - return Err(PyValueError::new_err(format!( + let idx_or_fd: Result = if let Some(&(idx, sz)) = channel_indexes.get(key) { + if sz != *size { + return Err(format!( "bad shard size of name {key}: wanted {size}, got {sz}" - ))); + )); } Ok(idx) } else { if empty_fname { - return Err(PyOSError::new_err(format!("cannot load shard {key}"))); + return Err(format!("cannot load shard {key}")); } Err(open(fname.as_str(), oflags, Mode::empty()) - .map_err(|e| PyOSError::new_err(format!("failed to open file {}: {}", fname, e)))?) + .map_err(|e| format!("failed to open file {fname}: {e}"))?) }; - shards.push((idx_or_fd, offset, format!("{}{}", prefix, key))); + shards.push((idx_or_fd, *offset, format!("{prefix}{key}"))); } let ok = AtomicBool::new(true); - py.detach(|| { - let counters = (AtomicUsize::new(0), AtomicUsize::new(0)); - for row_segment_idx in 0..num_row_segments { - let shard_idx = row_segment_idx * num_shards; - let chunk_size = cmp::max( - READ_SIZE * num_shards, - row_segment_size / num_shards / PAGE_SIZE / CONCURRENCY_OUTER - * num_shards - * PAGE_SIZE, - ); - let row_segment_slice = &mut tensor_slice - [row_segment_size * row_segment_idx..row_segment_size * (row_segment_idx + 1)]; - let base_ptr = row_segment_slice.as_ptr() as usize; - row_segment_slice - .par_chunks_mut(chunk_size) - .for_each(|chunk| { - let file_offset = (chunk.as_ptr() as usize - base_ptr) / num_shards; - let min_count = chunk.len() / num_shards; - let buf_size = min_count.div_ceil(PAGE_SIZE) * PAGE_SIZE; - - let mut vecs = vec![vec![0; buf_size + 2 * PAGE_SIZE]; delta_shards]; - let mut bufs = Vec::with_capacity(delta_shards); - for v in vecs.iter_mut() { - let _ = madvise_hugepage_internal(v); - let o = (!(v.as_ptr() as usize) + 1) & (PAGE_SIZE - 1); - bufs.push(&mut v[o..o + buf_size + PAGE_SIZE]); - } - let mut pads = vec![0; delta_shards]; - let mut counts = (0usize, 0usize); - - for block_idx in 0..num_blocks { - let mut futures = Vec::>::with_capacity(delta_shards); - - for idx in 0..delta_shards { - let shard = &shards[shard_idx + block_idx * delta_shards + idx]; - let offset = shard.1 + file_offset; - pads[idx] = offset & (PAGE_SIZE - 1); - let o = (offset - pads[idx]) as i64; - match &shard.0 { - Ok(channel_idx) => { - let b = bufs[idx].as_mut_ptr(); - let b: &'static mut [u8] = unsafe { - mem::transmute(slice::from_raw_parts_mut( - b.add(pads[idx]), - min_count, - )) - }; - futures.push(Box::pin(send_entries( - channels[*channel_idx].clone(), - vec![shard.2.as_bytes().to_vec()], - vec![file_offset], - vec![min_count], - b, - #[cfg(target_os = "linux")] - (Vec::new(), Arc::new(Vec::new()), Arc::new(Vec::new())), - ))); - } - Err(fd) => { - let count = pread(fd.as_fd(), bufs[idx], o).unwrap_or(0); - if count < min_count + pads[idx] { - ok.store(false, Ordering::Relaxed); - } - counts.0 += count; - } + let counters = (AtomicUsize::new(0), AtomicUsize::new(0)); + for row_segment_idx in 0..num_row_segments { + let shard_idx = row_segment_idx * num_shards; + let chunk_size = cmp::max( + READ_SIZE * num_shards, + row_segment_size / num_shards / PAGE_SIZE / CONCURRENCY_OUTER * num_shards * PAGE_SIZE, + ); + let row_segment_slice = &mut tensor + [row_segment_size * row_segment_idx..row_segment_size * (row_segment_idx + 1)]; + let base_ptr = row_segment_slice.as_ptr() as usize; + row_segment_slice + .par_chunks_mut(chunk_size) + .for_each(|chunk| { + let file_offset = (chunk.as_ptr() as usize - base_ptr) / num_shards; + let min_count = chunk.len() / num_shards; + let buf_size = min_count.div_ceil(PAGE_SIZE) * PAGE_SIZE; + + let mut vecs = vec![vec![0; buf_size + 2 * PAGE_SIZE]; delta_shards]; + let mut bufs = Vec::with_capacity(delta_shards); + for v in vecs.iter_mut() { + let _ = madvise_hugepage_internal(v); + let o = (!(v.as_ptr() as usize) + 1) & (PAGE_SIZE - 1); + bufs.push(&mut v[o..o + buf_size + PAGE_SIZE]); + } + let mut pads = vec![0; delta_shards]; + let mut counts = (0usize, 0usize); + + for block_idx in 0..num_blocks { + let mut futures = Vec::>::with_capacity(delta_shards); + + for idx in 0..delta_shards { + let shard = &shards[shard_idx + block_idx * delta_shards + idx]; + let offset = shard.1 + file_offset; + pads[idx] = offset & (PAGE_SIZE - 1); + let o = (offset - pads[idx]) as i64; + match &shard.0 { + Ok(channel_idx) => { + let b = bufs[idx].as_mut_ptr(); + let b: &'static mut [u8] = unsafe { + mem::transmute(slice::from_raw_parts_mut( + b.add(pads[idx]), + min_count, + )) + }; + futures.push(Box::pin(send_entries( + channels[*channel_idx].clone(), + vec![shard.2.as_bytes().to_vec()], + vec![file_offset], + vec![min_count], + b, + #[cfg(target_os = "linux")] + (Vec::new(), Arc::new(Vec::new()), Arc::new(Vec::new())), + ))); } - } - - if let Some(results) = block_on(&runtime, futures, None) { - for (count, _) in results { - if count != min_count { + Err(fd) => { + let count = pread(fd.as_fd(), bufs[idx], o).unwrap_or(0); + if count < min_count + pads[idx] { ok.store(false, Ordering::Relaxed); } - counts.1 += count; + counts.0 += count; } - } else { - ok.store(false, Ordering::Relaxed); } + } - let y0 = block_idx * delta_shards * w_s; - let slice_size = chunk.len() / CONCURRENCY_INNER / row_size * row_size; - let base_ptr = chunk.as_ptr() as usize; - chunk.par_chunks_mut(slice_size).for_each(|slice| { - let row_idx_base = (slice.as_ptr() as usize - base_ptr) / row_size; - for row_idx in 0..slice.len() / row_size { - let x1 = (row_idx_base + row_idx) * w_s; - let y1 = y0 + row_idx * row_size; - for idx in 0..delta_shards { - let x2 = x1 + pads[idx]; - let y2 = y1 + idx * w_s; - slice[y2..y2 + w_s].copy_from_slice(&bufs[idx][x2..x2 + w_s]); - } + if let Some(results) = block_on(&runtime, futures, None) { + for (count, _) in results { + if count != min_count { + ok.store(false, Ordering::Relaxed); } - }); + counts.1 += count; + } + } else { + ok.store(false, Ordering::Relaxed); } - counters.0.fetch_add(counts.0, Ordering::Relaxed); - counters.1.fetch_add(counts.1, Ordering::Relaxed); - }); - } - log::info!( - "load_tensor stats: {} via file system, {} via gRPC, {} total", - counters.0.load(Ordering::Relaxed), - counters.1.load(Ordering::Relaxed), - tensor_size, - ); - }); + let y0 = block_idx * delta_shards * w_s; + let slice_size = chunk.len() / CONCURRENCY_INNER / row_size * row_size; + let base_ptr = chunk.as_ptr() as usize; + chunk.par_chunks_mut(slice_size).for_each(|slice| { + let row_idx_base = (slice.as_ptr() as usize - base_ptr) / row_size; + for row_idx in 0..slice.len() / row_size { + let x1 = (row_idx_base + row_idx) * w_s; + let y1 = y0 + row_idx * row_size; + for idx in 0..delta_shards { + let x2 = x1 + pads[idx]; + let y2 = y1 + idx * w_s; + slice[y2..y2 + w_s].copy_from_slice(&bufs[idx][x2..x2 + w_s]); + } + } + }); + } + counters.0.fetch_add(counts.0, Ordering::Relaxed); + counters.1.fetch_add(counts.1, Ordering::Relaxed); + }); + } + + log::info!( + "load_tensor stats: {} via file system, {} via gRPC, {} total", + counters.0.load(Ordering::Relaxed), + counters.1.load(Ordering::Relaxed), + tensor_size, + ); if !ok.load(Ordering::Relaxed) { - return Err(PyOSError::new_err("could not read files")); + return Err("could not read files".into()); } Ok(()) } +#[pyfunction] +pub fn load_tensor<'py>( + py: Python<'py>, + path: Bound<'py, PyString>, + urls: Bound<'py, PyString>, + shard_sources: Bound<'py, PyList>, + mut tensor: PyReadwriteArray1<'py, u8>, + row_size: usize, + num_row_segments: usize, +) -> PyResult<()> { + let mut parsed = Vec::with_capacity(shard_sources.len()); + for t in shard_sources.iter() { + let err0 = + || PyTypeError::new_err("shard_sources entry must be (name, fname, offset, size)"); + let tuple = t.downcast::().map_err(|_| err0())?; + if tuple.len() != 4 { + return Err(err0()); + } + let k = tuple.get_item(0)?; + let key: String = k.extract().map_err(|_| { + PyTypeError::new_err(format!( + "shard_sources name {} must be a string", + k.repr() + .map(|s| s.to_string()) + .unwrap_or("".to_string()) + )) + })?; + let err = || { + PyTypeError::new_err(format!( + "shard_sources entry must be (name, fname, offset, size) for name {key}" + )) + }; + let fname: String = tuple.get_item(1)?.extract().map_err(|_| err())?; + let offset: usize = tuple.get_item(2)?.extract().map_err(|_| err())?; + let size: usize = tuple.get_item(3)?.extract().map_err(|_| err())?; + parsed.push((key, fname, offset, size)); + } + let path = path.to_string(); + let urls = urls.to_string(); + let tensor_slice = tensor.as_slice_mut()?; + py.detach(|| { + load_tensor_into( + &path, + &urls, + &parsed, + tensor_slice, + row_size, + num_row_segments, + ) + .map_err(PyOSError::new_err) + }) +} + #[cfg(test)] mod copy_rows_tests { use super::*; diff --git a/phoenix/crates/serving/xai-recsys-engine/src/grpc_compression.rs b/phoenix/crates/serving/xai-recsys-engine/src/grpc_compression.rs index c16fb450..53dd2567 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/grpc_compression.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/grpc_compression.rs @@ -3,13 +3,14 @@ use axum::http; use axum::response::IntoResponse; use bytes::{BufMut, Bytes, BytesMut}; -use http_body_util::BodyExt; +use http_body_util::{BodyExt, Full}; use lazy_static::lazy_static; use prometheus::{HistogramVec, exponential_buckets, register_histogram_vec}; use std::convert::Infallible; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; +use std::time::Instant; use tonic::server::NamedService; use tower::Service; @@ -24,6 +25,13 @@ lazy_static! { exponential_buckets(512.0, 2.0, 15).unwrap() ) .unwrap(); + static ref RESPONSE_COMPRESS_MS: HistogramVec = register_histogram_vec!( + "recsys_engine_response_compress_ms", + "Outbound zstd time in ms (spawn_blocking queue + encode), by method.", + &["method"], + crate::request_metrics::latency_buckets_ms() + ) + .unwrap(); } fn method_label(path: &str) -> &str { @@ -97,17 +105,19 @@ where return Ok(http::Response::from_parts(parts, axum::body::Body::empty())); } }; + let trailers = collected.trailers().cloned(); let data = collected.to_bytes(); if data.len() <= GRPC_FRAME_HEADER_SIZE { - return Ok(http::Response::from_parts( - parts, - axum::body::Body::from(data), - )); + return Ok(rebuild_response(parts, data, trailers)); } let original = data.clone(); + let started = Instant::now(); let result = tokio::task::spawn_blocking(move || compress_grpc_frame(&data)).await; + RESPONSE_COMPRESS_MS + .with_label_values(&[&method]) + .observe(started.elapsed().as_secs_f64() * 1000.0); match result { Ok(compressed) if compressed[0] == 1 => { @@ -117,20 +127,29 @@ where parts .headers .insert("grpc-encoding", http::HeaderValue::from_static("zstd")); - Ok(http::Response::from_parts( - parts, - axum::body::Body::from(compressed), - )) + Ok(rebuild_response(parts, compressed, trailers)) } - _ => Ok(http::Response::from_parts( - parts, - axum::body::Body::from(original), - )), + _ => Ok(rebuild_response(parts, original, trailers)), } }) } } +fn rebuild_response( + parts: http::response::Parts, + data: Bytes, + trailers: Option, +) -> http::Response { + let full = Full::new(data); + let body = match trailers { + Some(tr) => { + axum::body::Body::new(full.with_trailers(async move { Some(Ok::<_, Infallible>(tr)) })) + } + None => axum::body::Body::new(full), + }; + http::Response::from_parts(parts, body) +} + fn compress_grpc_frame(data: &Bytes) -> Bytes { if data.len() < GRPC_FRAME_HEADER_SIZE { return data.clone(); @@ -194,4 +213,81 @@ mod tests { assert!(out_len < payload.len()); assert_eq!(out.len(), GRPC_FRAME_HEADER_SIZE + out_len); } + + #[test] + fn compress_ms_accepts_method_label() { + RESPONSE_COMPRESS_MS + .with_label_values(&["PredictNextActions"]) + .observe(1.0); + } + + #[derive(Clone)] + struct FakeSvc { + payload: Vec, + } + + impl NamedService for FakeSvc { + const NAME: &'static str = "test"; + } + + impl tower::Service> for FakeSvc { + type Response = http::Response; + type Error = Infallible; + type Future = std::future::Ready>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: http::Request<()>) -> Self::Future { + let mut frame = BytesMut::with_capacity(GRPC_FRAME_HEADER_SIZE + self.payload.len()); + frame.put_u8(0); + frame.put_u32(self.payload.len() as u32); + frame.put_slice(&self.payload); + let full = Full::new(frame.freeze()); + let body = axum::body::Body::new(full.with_trailers(async { + let mut t = http::HeaderMap::new(); + t.insert("grpc-status", http::HeaderValue::from_static("0")); + Some(Ok::<_, Infallible>(t)) + })); + std::future::ready(Ok(http::Response::new(body))) + } + } + + async fn call_zstd(payload: Vec) -> http::Response { + let mut svc = GrpcCompressionService::new(FakeSvc { payload }); + let req = http::Request::builder() + .uri("/xai_recsys.RecsysPredictor/PredictNextActions") + .header("grpc-accept-encoding", "zstd") + .body(()) + .unwrap(); + svc.call(req).await.unwrap() + } + + #[tokio::test] + async fn rebuild_keeps_grpc_status_trailers_when_compressed() { + let resp = call_zstd(vec![0u8; 4096]).await; + assert_eq!( + resp.headers().get("grpc-encoding").map(|v| v.as_bytes()), + Some(&b"zstd"[..]) + ); + let collected = resp.into_body().collect().await.unwrap(); + let trailers = collected.trailers().expect("grpc-status trailers"); + assert_eq!( + trailers.get("grpc-status").map(|v| v.as_bytes()), + Some(&b"0"[..]) + ); + assert_eq!(collected.to_bytes()[0], 1); + } + + #[tokio::test] + async fn rebuild_keeps_grpc_status_trailers_when_too_small_to_compress() { + let resp = call_zstd(vec![]).await; + let collected = resp.into_body().collect().await.unwrap(); + let trailers = collected.trailers().expect("grpc-status trailers"); + assert_eq!( + trailers.get("grpc-status").map(|v| v.as_bytes()), + Some(&b"0"[..]) + ); + } } diff --git a/phoenix/crates/serving/xai-recsys-engine/src/lib.rs b/phoenix/crates/serving/xai-recsys-engine/src/lib.rs index 01515e96..d7aec504 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/lib.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/lib.rs @@ -23,7 +23,6 @@ pub mod multimodal_retrieval; pub mod proto_parser; pub mod request_metrics; pub mod request_queue; -pub mod sid_client; pub mod storage_util; pub mod tls; pub mod util; diff --git a/phoenix/crates/serving/xai-recsys-engine/src/proto_parser.rs b/phoenix/crates/serving/xai-recsys-engine/src/proto_parser.rs index 4b0935da..13b5bfaa 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/proto_parser.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/proto_parser.rs @@ -248,6 +248,7 @@ type FuncU32<'a> = Box; type FuncU64<'a> = Box; type FuncInt<'a> = Box; type FuncStr<'a> = Box; +type FuncBytes<'a> = Box; pub enum Func<'a> { Msg(Vec<(u64, Func<'a>)>), @@ -256,6 +257,7 @@ pub enum Func<'a> { U64(FuncU64<'a>), Int(FuncInt<'a>), Str(FuncStr<'a>), + Bytes(FuncBytes<'a>), } pub trait FuncConvertible<'a, Args> { @@ -304,6 +306,13 @@ impl<'a, F: FnMut(usize, &[u8]) + Send + 'a> FuncConvertible<'a, (usize, &[u8])> } } +impl<'a, F: FnMut(usize, Bytes) + Send + 'a> FuncConvertible<'a, (usize, Bytes)> for F { + fn into(self, x: u64) -> (u64, Func<'a>) { + assert_ne!(x, 0); + ((x << 3) | 2, Func::Bytes(Box::new(self))) + } +} + pub fn proto<'a, F, Args>(x: u64, arg: F) -> (u64, Func<'a>) where F: FuncConvertible<'a, Args>, @@ -501,6 +510,8 @@ pub async fn parse<'a, T: Body + Send + Unpin>( let len = min(left, chunk.len() - pos); if let Func::Str(f) = func { f(left, &chunk[pos..pos + len]); + } else if let Func::Bytes(f) = func { + f(left, chunk.slice(pos..pos + len)); } pos += len; left -= len; diff --git a/phoenix/crates/serving/xai-recsys-engine/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index 0639fb14..a7c02436 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -895,8 +895,6 @@ struct RecsysPredictorImpl { reload_directive: Arc>>, enqueue_timeout_ms: u64, mm_embeddings_client: Option, - #[allow(dead_code)] - sid_client: Option>, admission: Arc, #[allow(dead_code)] prefetch_mm_query_config: Option>, @@ -1171,36 +1169,6 @@ impl pb::recsys_predictor_server::RecsysPredictor for RecsysPredictorImpl { } } -fn fill_semantic_ids( - sids: &std::collections::HashMap>, - post_ids: &[i64], - dst: &mut [u16], - sid_dim: usize, -) { - for (j, &post_id) in post_ids.iter().enumerate() { - if post_id == 0 { - continue; - } - if let Some(codes) = sids.get(&post_id) { - assert_eq!( - codes.len(), - sid_dim, - "SID codes length ({}) != sid_num_levels ({}) for post_id {}", - codes.len(), - sid_dim, - post_id, - ); - let dst_start = j * sid_dim; - for (d, &c) in dst[dst_start..dst_start + sid_dim] - .iter_mut() - .zip(codes.iter()) - { - *d = (c + 1) as u16; - } - } - } -} - async fn handle_request( candidate_set: CandidateSet, sequence: Option, @@ -1907,7 +1875,7 @@ impl PrepareBatch for RankingBatchPrep { let output_vocab_size = model_config.hash_table.output_vocab_size; let num_continuous_actions = model_config.hash_table.num_continuous_actions; let embedding_dim = model_config.multimodal_embedding_dim; - let search_query_embedding_dim = model_config.hash_table.search_query_embedding_dim; + let search_query_embedding_dim = model_config.search_query_embedding_dim; let sid_num_levels = model_config.sid_num_levels; py.detach(|| { @@ -2054,6 +2022,7 @@ impl PrepareBatch for RankingBatchPrep { .par_chunks(chunk_size) .zip(candidate_embeddings_slices.par_iter_mut()) .for_each(|(items_chunk, shard_slice)| { + shard_slice.fill(f16::ZERO); shard_slice .par_chunks_exact_mut(row_size) .zip(items_chunk.par_iter()) @@ -2061,7 +2030,10 @@ impl PrepareBatch for RankingBatchPrep { if let Some(ref input_buffer) = item.input_buffer { let src = &input_buffer.candidate_embeddings; let copy_len = embedding_row.len().min(src.len()); - embedding_row[..copy_len].copy_from_slice(&src[..copy_len]); + if copy_len > 0 { + embedding_row[..copy_len] + .copy_from_slice(&src[..copy_len]); + } } }); }); @@ -2070,14 +2042,23 @@ impl PrepareBatch for RankingBatchPrep { if search_query_embedding_dim > 0 && let Some(sq_slice) = candidate_search_query_embeddings_slice { + sq_slice.fill(0.0); sq_slice .par_chunks_exact_mut(candidate_seq_len * search_query_embedding_dim) .take(length_of_input) .zip(request.items.par_iter()) .for_each(|(search_query_emb_row, item)| { if let Some(ref input_buffer) = item.input_buffer { - search_query_emb_row - .copy_from_slice(&input_buffer.candidate_search_query_embeddings); + let src = &input_buffer.candidate_search_query_embeddings; + if src.len() == search_query_embedding_dim { + let n_rep = input_buffer + .num_real_candidates(num_item_hashes, candidate_seq_len); + xai_recsys::util::repeat_query_into( + search_query_emb_row, + src, + n_rep, + ); + } } }); } @@ -2751,7 +2732,6 @@ struct RecsysRetrievalPredictorImpl { reload_directive: Arc>>, enqueue_timeout_ms: u64, mm_embeddings_client: Option, - sid_client: Option>, admission: Arc, prefetch_mm_query_config: Option>, } @@ -2870,7 +2850,6 @@ impl RecsysRetrievalPredictorImpl { &self.model_config, &cancel_token, self.enqueue_timeout_ms, - &self.sid_client, request.client_context, request.user_context, deadline, @@ -2975,7 +2954,6 @@ async fn handle_retrieval_request( model_config: &ModelConfig, cancellation_token: &CancellationToken, enqueue_timeout_ms: u64, - sid_client: &Option>, client_context: Option, user_context: Option, deadline: Option, @@ -2990,7 +2968,7 @@ async fn handle_retrieval_request( .unwrap_or(""); let start = std::time::Instant::now(); - let mut input_buffer = if let Some(ref columnar_bytes) = columnar_sequence_bytes { + let input_buffer = if let Some(ref columnar_bytes) = columnar_sequence_bytes { log::debug!( "Computing retrieval input buffer from columnar bytes ({} bytes)", columnar_bytes.len() @@ -3029,25 +3007,6 @@ async fn handle_retrieval_request( ) }; - let sid_num_levels = model_config.sid_num_levels; - if sid_num_levels > 0 - && let Some(client) = sid_client.as_ref() - { - let history_ids: Vec = input_buffer - .history_post_ids - .iter() - .copied() - .filter(|&id| id != 0) - .collect(); - let sids = client.lookup(&history_ids).await; - fill_semantic_ids( - &sids, - &input_buffer.history_post_ids, - &mut input_buffer.history_semantic_ids, - sid_num_levels, - ); - } - let duration = start.elapsed(); INPUT_BUFFER_COMPUTATION_TIME.observe(duration.as_secs_f64()); @@ -3134,7 +3093,6 @@ macro_rules! server_impl { enqueue_timeout_ms = ENQUEUE_TIMEOUT_MS, queue_max_staleness_ms = QUEUE_MAX_STALENESS_MS, mm_client = None, - sid_client = None, user_id_table_size = 100_000, user_hash_scales = vec![196742702, 1852108266], user_biases = vec![1935840681, 167407236], @@ -3168,6 +3126,7 @@ macro_rules! server_impl { num_post_bool_features = 0, num_post_float_features = 0, num_post_int64_features = 0, + enable_stale_post = false, enable_async_response_compression = false, sid_num_levels = 0, copy_max_entries = DEFAULT_MAX_ENTRIES, @@ -3198,7 +3157,6 @@ macro_rules! server_impl { enqueue_timeout_ms: u64, queue_max_staleness_ms: u64, mm_client: Option<&PyMmEmbeddingsClient>, - sid_client: Option<&crate::sid_client::PySemanticIdClient>, user_id_table_size: usize, user_hash_scales: Vec, user_biases: Vec, @@ -3232,6 +3190,7 @@ macro_rules! server_impl { num_post_bool_features: usize, num_post_float_features: usize, num_post_int64_features: usize, + enable_stale_post: bool, enable_async_response_compression: bool, sid_num_levels: usize, #[allow(unused_variables)] @@ -3317,6 +3276,7 @@ macro_rules! server_impl { num_post_bool_features, num_post_float_features, num_post_int64_features, + enable_stale_post, ); internal_model_config.sid_num_levels = sid_num_levels; let reload_request = Arc::new(AtomicBool::new(false)); @@ -3345,7 +3305,6 @@ macro_rules! server_impl { let _guard = runtime.enter(); let mm_embeddings_client = mm_client.map(|c| c.client().clone()); - let sid_client_arc = sid_client.map(|c| c.build()); let prefetch_mm_query_config: Option> = if prefetch_mm_query_for_retrieval && mm_embeddings_client.is_some() { @@ -3384,7 +3343,6 @@ macro_rules! server_impl { reload_directive: reload_directive.clone(), enqueue_timeout_ms, mm_embeddings_client, - sid_client: sid_client_arc, admission: admission.clone(), prefetch_mm_query_config, }; @@ -3723,7 +3681,6 @@ pub fn xai_recsys_engine(_py: Python<'_>, m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs b/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs index bb7a9646..853d368f 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs @@ -172,8 +172,9 @@ lazy_static! { "recsys_engine_client_server_network_delay_ms", "Approximate client->server network delay (ms): wall-clock delta \ between the caller's x-request-send-time-ms (stamped just before \ - the request hits the wire) and server handler entry.", - &["client"], + the request hits the wire) and server handler entry. Labels: \ + client, src_dc.", + &["client", "src_dc"], buckets ) .unwrap() @@ -365,7 +366,7 @@ pub fn export() { let _ = TOKENS_PER_ROW.with_label_values(&["history"]); let _ = TOKENS_PER_ROW.with_label_values(&["candidate"]); let _ = NOT_READY_ARRIVAL_SECONDS.with_label_values(&["unknown"]); - let _ = CLIENT_SERVER_NETWORK_DELAY.with_label_values(&["unknown"]); + let _ = CLIENT_SERVER_NETWORK_DELAY.with_label_values(&["unknown", "unknown"]); let _ = ADMISSION_BUDGET_REMAINING.with_label_values(&["unknown"]); NUM_REQUESTS_REJECTED .with_label_values(&["deadline_admission", "unknown"]) @@ -385,8 +386,9 @@ pub fn record_client_server_network_delay(request: &Request, client: &str) .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis() as u64) .unwrap_or(0); + let src_dc = src_dc_label(request); CLIENT_SERVER_NETWORK_DELAY - .with_label_values(&[client]) + .with_label_values(&[client, &src_dc]) .observe(now_ms.saturating_sub(sent_ms) as f64); } @@ -432,9 +434,7 @@ pub fn observe_admission( return Ok(()); } DEADLINE_SHED.with_label_values(&[client]).inc(); - NUM_REQUESTS_REJECTED - .with_label_values(&["deadline_admission", client]) - .inc(); + guard.record_reject("deadline_admission"); let budget_source = if grpc_timeout.is_some() { "grpc-timeout" } else { @@ -454,6 +454,12 @@ pub fn record_admission_service_sample(admission: &AdmissionController, sample_u ADMISSION_SERVICE_TIME_MS.set(admission.service_time_us() as f64 / 1000.0); } +pub fn reset_admission_estimates(admission: &AdmissionController) { + admission.reset_estimates(); + ADMISSION_SERVICE_TIME_MS.set(0.0); + ADMISSION_PIPELINE_TIME_MS.set(0.0); +} + pub fn record_admission_sojourn_sample( admission: &AdmissionController, sample_us: u64, @@ -560,6 +566,7 @@ pub struct RequestMetricsGuard { admit_backlog: i64, admission: Option>, finished: bool, + recorded_reject: bool, } impl RequestMetricsGuard { @@ -571,9 +578,17 @@ impl RequestMetricsGuard { admit_backlog: 0, admission: None, finished: false, + recorded_reject: false, } } + pub fn record_reject(&mut self, reason: &str) { + NUM_REQUESTS_REJECTED + .with_label_values(&[reason, self.labels.client.as_str()]) + .inc(); + self.recorded_reject = true; + } + pub fn client(&self) -> &str { &self.labels.client } @@ -600,7 +615,7 @@ impl RequestMetricsGuard { pub fn finish(mut self, result: tonic::Result>) -> tonic::Result> { self.record_sojourn(); self.finished = true; - on_request_finished(&self.labels, &result); + on_request_finished(&self.labels, &result, self.recorded_reject); result } } @@ -615,7 +630,11 @@ impl Drop for RequestMetricsGuard { } } -pub fn on_request_finished(labels: &RequestLabels, result: &tonic::Result>) { +pub fn on_request_finished( + labels: &RequestLabels, + result: &tonic::Result>, + already_rejected: bool, +) { let code = status_code(result); emit_request_metrics( labels.start, @@ -628,14 +647,18 @@ pub fn on_request_finished(labels: &RequestLabels, result: &tonic::Result NUM_REQUESTS_SUCCEEDED.inc(), tonic::Code::Cancelled => NUM_REQUESTS_CANCELED.inc(), tonic::Code::Unavailable => { - NUM_REQUESTS_REJECTED - .with_label_values(&["checkpoint_loading", labels.client.as_str()]) - .inc(); + if !already_rejected { + NUM_REQUESTS_REJECTED + .with_label_values(&["checkpoint_loading", labels.client.as_str()]) + .inc(); + } } tonic::Code::ResourceExhausted => { - NUM_REQUESTS_REJECTED - .with_label_values(&["inflight_cap", labels.client.as_str()]) - .inc(); + if !already_rejected { + NUM_REQUESTS_REJECTED + .with_label_values(&["inflight_cap", labels.client.as_str()]) + .inc(); + } } tonic::Code::InvalidArgument => {} _ => NUM_REQUESTS_FAILED.inc(), @@ -645,6 +668,7 @@ pub fn on_request_finished(labels: &RequestLabels, result: &tonic::Result u64 { + NUM_REQUESTS_REJECTED + .with_label_values(&[reason, client]) + .get() + } + + #[test] + fn reset_admission_estimates_clears_gauges_and_p() { + let admission = AdmissionController::new(AdmissionConfig { + enabled: true, + batch_size: 4, + margin_us: 0, + post_us: 0, + fallback_budget_us: 0, + ewma_alpha: 1.0, + pipeline_depth: 0, + eta_model: AdmissionEtaModel::Pipeline, + }); + admission.record_service_time_us(50_000); + admission.record_sojourn_us(80_000, 0); + record_admission_service_sample(&admission, 50_000); + record_admission_sojourn_sample(&admission, 80_000, 0); + assert!(admission.pipeline_time_us() > 0); + reset_admission_estimates(&admission); + assert_eq!(admission.service_time_us(), 0); + assert_eq!(admission.pipeline_time_us(), 0); + assert_eq!(ADMISSION_SERVICE_TIME_MS.get(), 0.0); + assert_eq!(ADMISSION_PIPELINE_TIME_MS.get(), 0.0); + } + + #[test] + fn deadline_shed_does_not_double_count_as_inflight_cap() { + let client = "test-deadline-shed-no-double"; + let admission = AdmissionController::new(AdmissionConfig { + enabled: true, + batch_size: 1, + margin_us: 0, + post_us: 0, + fallback_budget_us: 0, + ewma_alpha: 1.0, + pipeline_depth: 0, + eta_model: AdmissionEtaModel::Loop, + }); + admission.record_service_time_us(100_000); + + let (tx, _rx) = tokio::sync::mpsc::channel::<()>(2); + tx.try_send(()).unwrap(); + + let mut req = Request::new(()); + req.metadata_mut() + .insert("x-client-name", client.parse().unwrap()); + req.metadata_mut() + .insert("grpc-timeout", "1m".parse().unwrap()); + + let shed_before = DEADLINE_SHED.with_label_values(&[client]).get(); + let deadline_before = rejected("deadline_admission", client); + let inflight_before = rejected("inflight_cap", client); + + let mut guard = RequestMetricsGuard::begin(&req); + let err = observe_admission(&req, client, 0, &tx, &admission, &mut guard) + .expect_err("deadline admission must shed"); + assert_eq!(err.code(), tonic::Code::ResourceExhausted); + let _ = guard.finish::<()>(Err(err)); + + assert_eq!( + DEADLINE_SHED.with_label_values(&[client]).get(), + shed_before + 1 + ); + assert_eq!(rejected("deadline_admission", client), deadline_before + 1); + assert_eq!(rejected("inflight_cap", client), inflight_before); + } + + #[test] + fn resource_exhausted_without_prior_reject_counts_inflight_cap() { + let client = "test-inflight-cap-from-finish"; + let mut req = Request::new(()); + req.metadata_mut() + .insert("x-client-name", client.parse().unwrap()); + let before = rejected("inflight_cap", client); + let guard = RequestMetricsGuard::begin(&req); + let _ = guard.finish::<()>(Err(Status::resource_exhausted("queue full"))); + assert_eq!(rejected("inflight_cap", client), before + 1); + } } diff --git a/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs deleted file mode 100644 index 50b92704..00000000 --- a/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs +++ /dev/null @@ -1,144 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright 2026 X.AI Corp. -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::{Duration, Instant}; - -use lazy_static::lazy_static; -use log::{info, warn}; -use prometheus::{ - CounterVec, HistogramVec, exponential_buckets, register_counter_vec, register_histogram_vec, -}; -use pyo3::prelude::*; -use tonic::transport::Endpoint; -use xai_recsys_sid_proto::{LookupSidsRequest, sid_lookup_service_client::SidLookupServiceClient}; - -const REQUEST_TIMEOUT: Duration = Duration::from_millis(100); -const NUM_CHANNELS: usize = 2; - -lazy_static! { - static ref SID_LOOKUP_LATENCY_SECS: HistogramVec = register_histogram_vec!( - "recsys_engine_sid_lookup_latency_seconds", - "Latency of semantic ID lookup gRPC calls", - &["status"], - exponential_buckets(0.001, 2.0, 15).unwrap(), - ) - .unwrap(); - static ref SID_LOOKUP_COUNT: CounterVec = register_counter_vec!( - "recsys_engine_sid_lookup_count", - "Number of semantic ID lookups by status", - &["status"], - ) - .unwrap(); -} - -fn build_endpoint(endpoint: &str) -> Endpoint { - let host_port = endpoint.strip_prefix("http://").unwrap_or(endpoint); - let url = format!("http://{host_port}"); - Endpoint::from_shared(url) - .expect("invalid SID endpoint URL") - .timeout(REQUEST_TIMEOUT) - .tcp_nodelay(true) - .http2_keep_alive_interval(Duration::from_secs(30)) - .keep_alive_timeout(Duration::from_secs(20)) - .keep_alive_while_idle(true) -} - -pub struct SemanticIdClient { - clients: Vec>, - next: AtomicUsize, - sid_num_levels: usize, -} - -impl SemanticIdClient { - pub fn new(endpoint: &str, sid_num_levels: usize) -> Self { - let clients = (0..NUM_CHANNELS) - .map(|_| { - let channel = build_endpoint(endpoint).connect_lazy(); - SidLookupServiceClient::new(channel) - }) - .collect(); - Self { - clients, - next: AtomicUsize::new(0), - sid_num_levels, - } - } - - pub async fn lookup(&self, post_ids: &[i64]) -> HashMap> { - if post_ids.is_empty() { - return HashMap::new(); - } - - let start = Instant::now(); - let request = LookupSidsRequest { - post_ids: post_ids.to_vec(), - }; - - let idx = self.next.fetch_add(1, Ordering::Relaxed) % self.clients.len(); - let mut client = self.clients[idx].clone(); - match client.lookup_sids(request).await { - Ok(response) => { - let elapsed = start.elapsed().as_secs_f64(); - SID_LOOKUP_LATENCY_SECS - .with_label_values(&["success"]) - .observe(elapsed); - - let resp = response.into_inner(); - let mut result = HashMap::with_capacity(post_ids.len()); - let missing = vec![-1i32; self.sid_num_levels]; - for (post_id, post_sids) in post_ids.iter().zip(resp.results.iter()) { - if post_sids.codes.is_empty() { - SID_LOOKUP_COUNT.with_label_values(&["miss"]).inc(); - result.insert(*post_id, missing.clone()); - } else { - SID_LOOKUP_COUNT.with_label_values(&["hit"]).inc(); - result.insert(*post_id, post_sids.codes.clone()); - } - } - result - } - Err(e) => { - let elapsed = start.elapsed().as_secs_f64(); - SID_LOOKUP_LATENCY_SECS - .with_label_values(&["error"]) - .observe(elapsed); - SID_LOOKUP_COUNT - .with_label_values(&["error"]) - .inc_by(post_ids.len() as f64); - warn!("SID lookup gRPC call failed: {}", e); - let missing = vec![-1i32; self.sid_num_levels]; - post_ids.iter().map(|id| (*id, missing.clone())).collect() - } - } - } -} - -#[pyclass(module = "xai_recsys_engine")] -pub struct PySemanticIdClient { - endpoint: String, - sid_num_levels: usize, -} - -impl PySemanticIdClient { - pub fn build(&self) -> Arc { - Arc::new(SemanticIdClient::new(&self.endpoint, self.sid_num_levels)) - } -} - -#[pymethods] -impl PySemanticIdClient { - #[new] - #[pyo3(signature = (endpoint, sid_num_levels))] - fn new(endpoint: String, sid_num_levels: usize) -> PyResult { - info!( - "Semantic ID client configured: endpoint={}, sid_num_levels={}", - endpoint, sid_num_levels - ); - Ok(Self { - endpoint, - sid_num_levels, - }) - } -} diff --git a/phoenix/crates/serving/xai-recsys-engine/src/storage_util.rs b/phoenix/crates/serving/xai-recsys-engine/src/storage_util.rs index 6f89c472..e14d075e 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/storage_util.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/storage_util.rs @@ -5,7 +5,7 @@ use futures::future::join_all; use futures::{StreamExt, TryStreamExt}; use google_cloud_storage::client::{Storage, StorageControl}; -use crate::checkpoint_store::get_checkpoint_store; +use crate::checkpoint_store::{get_checkpoint_store, get_store}; use log::{debug, error, info}; use numpy::PyReadwriteArray1; use object_store::aws::AmazonS3Builder; @@ -76,8 +76,9 @@ fn o2_store_from_url( ..Default::default() }; - let endpoint = - env::var("O2_ENDPOINT_URL").unwrap_or_else(|_| "http://o2.example.invalid".into()); + let endpoint = env::var("O2_ENDPOINT_URL").map_err(|_| { + "O2 endpoint not set: set XAI_O2_ENDPOINT_URL or O2_ENDPOINT_URL".to_string() + })?; let access_key = env::var("O2_ACCESS_KEY_ID").unwrap_or_else(|_| "anonymous".to_string()); let secret_key = env::var("O2_SECRET_ACCESS_KEY").unwrap_or_else(|_| "anonymous".to_string()); @@ -223,7 +224,7 @@ pub fn upload_files_to_storage( num_files, storage_url, ); let upload_start = std::time::Instant::now(); - let store_result = get_checkpoint_store(&storage_url, Some(upload_runtime)); + let store_result = get_store(&storage_url, Some(upload_runtime)); match store_result { Ok((store, base_path)) => { let base = base_path.as_ref().to_string(); diff --git a/phoenix/crates/serving/xai-recsys-engine/src/util.rs b/phoenix/crates/serving/xai-recsys-engine/src/util.rs index 99f325fe..e239110e 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/util.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/util.rs @@ -11,7 +11,8 @@ use xai_recsys::{ model_config::ModelConfig, util::{ InputBuffer, PRODUCT_SURFACE_CATEGORICAL_IDX, compute_user_ip_hashes, parse_ipv4_to_u32, - stamp_i32_as_categorical, stamp_local_time_features, + record_sid_coverage, stamp_i32_as_categorical, stamp_local_time_features, + stamp_semantic_ids, }, }; use xai_recsys_proto as pb; @@ -63,6 +64,8 @@ impl RetrievalInputBuffer { let mut history_impr_ts = vec![0i32; history_seq_len]; let mut history_tz_enums = vec![0i16; history_seq_len]; let mut history_post_ids = vec![0i64; history_seq_len]; + let sid_num_levels = model_config.sid_num_levels; + let mut history_semantic_ids = vec![0u16; history_seq_len * sid_num_levels]; let user_id = match sequence { Some(sequence) => sequence.user_id as i64, @@ -118,6 +121,12 @@ impl RetrievalInputBuffer { } history_post_ids[valid_entry_count] = tweet_id; + stamp_semantic_ids( + &mut history_semantic_ids, + valid_entry_count, + sid_num_levels, + &tweet_info.semantic_ids, + ); let base_idx = valid_entry_count * output_vocab_size; let continuous_base_idx = valid_entry_count * num_continuous_actions; @@ -173,6 +182,13 @@ impl RetrievalInputBuffer { } } + if sid_num_levels > 0 { + let present = (0..valid_entry_count) + .filter(|&i| history_semantic_ids[i * sid_num_levels] != 0) + .count() as u64; + record_sid_coverage("history", present, valid_entry_count as u64); + } + let n_post_cat = model_config.hash_table.num_post_categorical_features; let n_post_bool = model_config.hash_table.num_post_bool_features; let n_post_float = model_config.hash_table.num_post_float_features; @@ -216,7 +232,7 @@ impl RetrievalInputBuffer { history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features: vec![0i64; history_seq_len * n_post_int64], history_post_ids, - history_semantic_ids: vec![0u16; history_seq_len * model_config.sid_num_levels], + history_semantic_ids, } } @@ -256,6 +272,8 @@ impl RetrievalInputBuffer { let mut history_impr_ts = vec![0i32; history_seq_len]; let mut history_tz_enums = vec![0i16; history_seq_len]; let mut history_post_ids = vec![0i64; history_seq_len]; + let sid_num_levels = model_config.sid_num_levels; + let mut history_semantic_ids = vec![0u16; history_seq_len * sid_num_levels]; let cursor = Cursor::new(columnar_bytes); let mut reader = StreamReader::try_new(cursor, None)?; @@ -291,7 +309,7 @@ impl RetrievalInputBuffer { history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features: vec![0i64; history_seq_len * n_post_int64], history_post_ids: vec![0i64; history_seq_len], - history_semantic_ids: vec![0u16; history_seq_len * model_config.sid_num_levels], + history_semantic_ids, }); } }; @@ -315,6 +333,7 @@ impl RetrievalInputBuffer { let col_timezone_id = batch .column_by_name("timezoneId") .and_then(|c| c.as_any().downcast_ref::()); + let col_semantic_id = batch.column_by_name("semanticId"); let start_row = num_rows.saturating_sub(history_seq_len); let mut valid_entry_count = 0; @@ -341,6 +360,20 @@ impl RetrievalInputBuffer { } history_post_ids[valid_entry_count] = tweet_id; + if let Some(sid_col) = col_semantic_id { + let fsl = sid_col.as_fixed_size_list(); + if !fsl.is_null(row_idx) { + let inner = fsl.value(row_idx); + if let Some(codes) = inner.as_any().downcast_ref::() { + stamp_semantic_ids( + &mut history_semantic_ids, + valid_entry_count, + sid_num_levels, + codes.values(), + ); + } + } + } if let Some(ps) = product_surfaces { history_product_surfaces[valid_entry_count] = ps.value(row_idx); @@ -381,6 +414,13 @@ impl RetrievalInputBuffer { valid_entry_count += 1; } + if sid_num_levels > 0 { + let present = (0..valid_entry_count) + .filter(|&i| history_semantic_ids[i * sid_num_levels] != 0) + .count() as u64; + record_sid_coverage("history", present, valid_entry_count as u64); + } + let n_post_cat = model_config.hash_table.num_post_categorical_features; let n_post_bool = model_config.hash_table.num_post_bool_features; let n_post_float = model_config.hash_table.num_post_float_features; @@ -424,7 +464,7 @@ impl RetrievalInputBuffer { history_float_features: vec![0.0f32; history_seq_len * n_post_float], history_int64_features: vec![0i64; history_seq_len * n_post_int64], history_post_ids, - history_semantic_ids: vec![0u16; history_seq_len * model_config.sid_num_levels], + history_semantic_ids, }) } } @@ -480,3 +520,80 @@ pub struct RetrieveRequestItem { pub mm_query_embeddings: Option>, pub mm_query_seed_post_ids: Option>, } + +#[cfg(test)] +mod tests { + use super::*; + + fn sid_test_config() -> ModelConfig { + let mut mc = ModelConfig::from_params( + 8, + vec![1], + vec![0], + 1009, + 8, + 0, + vec![1], + vec![0], + 1009, + 8, + vec![1], + vec![0], + 1009, + 4, + 2, + 4, + 1, + 0, + 0, + 0, + vec![], + vec![], + 1009, + 0, + 7, + 0, + 4, + 0, + 0, + 1, + 0, + 0, + 0, + false, + ); + mc.sid_num_levels = 3; + mc + } + + #[test] + fn proto_history_semantic_ids_are_read_from_tweet_info() { + let sequence = Some(pb::UserActionSequence { + user_actions_data: Some(pb::UserActionSequenceDataContainer { + data: Some( + pb::user_action_sequence_data_container::Data::OrderedAggregatedUserActionsList( + pb::AggregatedUserActionList { + aggregated_user_actions: vec![pb::AggregatedUserAction { + tweet_info: Some(pb::TweetInfo { + tweet_id: 11, + author_id: 21, + semantic_ids: vec![0, 7, -1], + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }, + ), + ), + }), + ..Default::default() + }); + let buf = + RetrievalInputBuffer::compute_for_item(&sid_test_config(), &sequence, "", None, None); + assert_eq!( + buf.history_semantic_ids, + vec![1, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] + ); + } +} diff --git a/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi b/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi index cafdcf25..7aabeb4e 100644 --- a/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi +++ b/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi @@ -19,9 +19,6 @@ class PyMmEmbeddingsClient: ) -> npt.NDArray[np.float16]: ... def wait_until_ready(self) -> None: ... -class PySemanticIdClient: - def __init__(self, endpoint: str, sid_num_levels: int) -> None: ... - class RankingBatchPrep: def __init__( self, @@ -178,7 +175,6 @@ class RecsysPredictorServer: enqueue_timeout_ms: int = 1000, queue_max_staleness_ms: int = 0, mm_client: PyMmEmbeddingsClient | None = None, - sid_client: PySemanticIdClient | None = None, user_id_table_size: int = 100_000, user_hash_scales: list[int] = ..., user_biases: list[int] = ..., @@ -211,6 +207,7 @@ class RecsysPredictorServer: num_post_bool_features: int = 0, num_post_float_features: int = 0, num_post_int64_features: int = 0, + enable_stale_post: bool = False, enable_async_response_compression: bool = False, sid_num_levels: int = 0, enable_deadline_admission: bool = False, @@ -278,7 +275,6 @@ class RecsysRetrievalPredictorServer: enqueue_timeout_ms: int = 1000, queue_max_staleness_ms: int = 0, mm_client: PyMmEmbeddingsClient | None = None, - sid_client: PySemanticIdClient | None = None, user_id_table_size: int = 100_000, user_hash_scales: list[int] = ..., user_biases: list[int] = ..., @@ -311,6 +307,7 @@ class RecsysRetrievalPredictorServer: num_post_bool_features: int = 0, num_post_float_features: int = 0, num_post_int64_features: int = 0, + enable_stale_post: bool = False, enable_async_response_compression: bool = False, sid_num_levels: int = 0, enable_deadline_admission: bool = False, diff --git a/phoenix/crates/serving/xai-recsys-mm-server/Cargo.toml b/phoenix/crates/serving/xai-recsys-mm-server/Cargo.toml index 49c987e8..f43407da 100644 --- a/phoenix/crates/serving/xai-recsys-mm-server/Cargo.toml +++ b/phoenix/crates/serving/xai-recsys-mm-server/Cargo.toml @@ -11,6 +11,7 @@ arrow = { workspace = true } chrono = { workspace = true } bytes = { workspace = true } lazy_static = { workspace = true } +rayon = { workspace = true } log = { workspace = true } parquet = { workspace = true, features = ["arrow"] } prometheus = { workspace = true } diff --git a/phoenix/crates/serving/xai-recsys-mm-server/src/mm_embedding_client.rs b/phoenix/crates/serving/xai-recsys-mm-server/src/mm_embedding_client.rs index c5313efc..3536b87b 100644 --- a/phoenix/crates/serving/xai-recsys-mm-server/src/mm_embedding_client.rs +++ b/phoenix/crates/serving/xai-recsys-mm-server/src/mm_embedding_client.rs @@ -120,82 +120,133 @@ impl MmEmbeddingsClient { self.shards.iter().map(|s| s.len()).sum() } - pub async fn preload_embeddings(&self, embeddings: Vec, parallel: bool) { - let num_shards = self.num_shards; + fn shard_records( + embeddings: Vec, + num_shards: usize, + ) -> Vec)>> { let mut buckets: Vec)>> = (0..num_shards).map(|_| Vec::new()).collect(); for record in embeddings { let idx = (record.post_id % num_shards as u64) as usize; buckets[idx].push((record.post_id, record.embedding)); } + buckets + } + + fn insert_bucket(shard: &CacheShard, bucket: Vec<(u64, Vec)>) { + match shard { + CacheShard::InProcess(c) => { + c.bulk_insert(bucket.into_iter().map(|(id, emb)| (id, Arc::new(emb)))); + } + CacheShard::Lmdb(c) => { + c.insert_many(bucket); + } + } + } + pub fn insert_sync(&self, embeddings: Vec) { + let buckets = Self::shard_records(embeddings, self.num_shards); + for (shard_idx, bucket) in buckets.into_iter().enumerate() { + Self::insert_bucket(&self.shards[shard_idx], bucket); + } + MM_EMBEDDING_CACHE_SIZE.set(self.total_len() as f64); + } + + pub fn ingest_parquet_bytes_sync(&self, data: bytes::Bytes, file_name: &str) -> Result { + let embeddings = read_embeddings_from_parquet(data, file_name)?; + let n = embeddings.len(); + self.insert_sync(embeddings); + Ok(n) + } + + pub async fn preload_embeddings(&self, embeddings: Vec, parallel: bool) { + let num_shards = self.num_shards; + let shards = self.shards.clone(); if parallel { + let buckets = match tokio::task::spawn_blocking(move || { + Self::shard_records(embeddings, num_shards) + }) + .await + { + Ok(b) => b, + Err(e) => { + log::error!("Shard-split task panicked: {e:#}"); + return; + } + }; let mut handles = Vec::with_capacity(num_shards); for (shard_idx, bucket) in buckets.into_iter().enumerate() { - let shard = self.shards[shard_idx].clone(); - handles.push(tokio::task::spawn_blocking(move || match shard { - CacheShard::InProcess(c) => { - c.bulk_insert(bucket.into_iter().map(|(id, emb)| (id, Arc::new(emb)))); - } - CacheShard::Lmdb(c) => { - c.insert_many(bucket); - } + let shard = shards[shard_idx].clone(); + handles.push(tokio::task::spawn_blocking(move || { + Self::insert_bucket(&shard, bucket); })); } for handle in handles { if let Err(e) = handle.await { - log::error!("Shard insertion task panicked: {:#}", e); + log::error!("Shard insertion task panicked: {e:#}"); } } + MM_EMBEDDING_CACHE_SIZE.set(self.total_len() as f64); } else { - let shards = self.shards.clone(); - if let Err(e) = tokio::task::spawn_blocking(move || { - for (shard_idx, bucket) in buckets.into_iter().enumerate() { - match &shards[shard_idx] { - CacheShard::InProcess(c) => { - c.bulk_insert(bucket.into_iter().map(|(id, emb)| (id, Arc::new(emb)))); - } - CacheShard::Lmdb(c) => { - c.insert_many(bucket); - } - } - } - }) - .await + let client = self.clone(); + if let Err(e) = + tokio::task::spawn_blocking(move || client.insert_sync(embeddings)).await { - log::error!("Shard insertion task panicked: {:#}", e); + log::error!("Cache insert task panicked: {e:#}"); } } - - MM_EMBEDDING_CACHE_SIZE.set(self.total_len() as f64); } pub fn get(&self, post_id: u64) -> Option>> { self.shard(post_id).get(post_id) } - pub async fn fetch_mm_embeddings( + pub fn in_process() -> Self { + let shard_capacity = MAX_EMBEDDING_CACHE_SIZE / NUM_SHARDS; + let ttl = embedding_ttl(); + log::info!( + "Creating in-process MM cache ({} shards, capacity={})", + NUM_SHARDS, + MAX_EMBEDDING_CACHE_SIZE + ); + Self { + shards: (0..NUM_SHARDS) + .map(|_| CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, ttl)))) + .collect(), + num_shards: NUM_SHARDS, + } + } + + pub fn fetch_mm_embeddings_into( &self, - post_ids: Vec, + post_ids: &[u64], emb_dim: usize, - candidate_seq_len: usize, - ) -> Result> { + dest: &mut [f16], + ) -> Result<()> { + use rayon::prelude::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + anyhow::ensure!(emb_dim > 0, "emb_dim must be > 0"); + anyhow::ensure!( + dest.len().is_multiple_of(emb_dim), + "dest len {} is not a multiple of emb_dim {}", + dest.len(), + emb_dim + ); let start_time = std::time::Instant::now(); - assert!(emb_dim > 0); let total_posts = post_ids.len(); - let processed = candidate_seq_len.min(total_posts); - let mut results = vec![f16::ZERO; emb_dim * processed]; - let mut missing = 0usize; - - for (idx, &post_id) in post_ids.iter().take(processed).enumerate() { - let start = idx * emb_dim; - if !self - .shard(post_id) - .copy_to(post_id, &mut results[start..start + emb_dim]) - { - missing += 1; - } - } + let processed = total_posts.min(dest.len() / emb_dim); + let missing = AtomicUsize::new(0); + + dest[..processed * emb_dim] + .par_chunks_mut(emb_dim) + .zip(post_ids.par_iter().take(processed)) + .for_each(|(dst, &post_id)| { + if !self.shard(post_id).copy_to(post_id, dst) { + missing.fetch_add(1, Ordering::Relaxed); + } + }); + let missing = missing.load(Ordering::Relaxed); let found = processed - missing; MM_EMBEDDING_TWEET_LOOKUP .with_label_values(&["found"]) @@ -209,8 +260,29 @@ impl MmEmbeddingsClient { } FETCH_MM_EMBEDDINGS_DURATION.observe(start_time.elapsed().as_secs_f64()); + Ok(()) + } + + pub fn fetch_mm_embeddings_sync( + &self, + post_ids: Vec, + emb_dim: usize, + candidate_seq_len: usize, + ) -> Result> { + let processed = candidate_seq_len.min(post_ids.len()); + let mut results = vec![f16::ZERO; emb_dim.saturating_mul(processed)]; + self.fetch_mm_embeddings_into(&post_ids, emb_dim, &mut results)?; Ok(results) } + + pub async fn fetch_mm_embeddings( + &self, + post_ids: Vec, + emb_dim: usize, + candidate_seq_len: usize, + ) -> Result> { + self.fetch_mm_embeddings_sync(post_ids, emb_dim, candidate_seq_len) + } } fn o2_env_vars() -> (String, String) { @@ -228,7 +300,7 @@ pub fn get_mm_client( let shard_capacity = MAX_EMBEDDING_CACHE_SIZE / NUM_SHARDS; let ttl = embedding_ttl(); - let shards: Vec = if shared { + let client = if shared { let mut shards = Vec::with_capacity(NUM_SHARDS); std::thread::scope(|s| { let handles: Vec<_> = (0..NUM_SHARDS) @@ -259,21 +331,12 @@ pub fn get_mm_client( shards.push(handle.join().expect("shard creation thread panicked")); } }); - shards + MmEmbeddingsClient { + shards, + num_shards: NUM_SHARDS, + } } else { - log::info!( - "Creating in-process MM cache ({} shards, capacity={})", - NUM_SHARDS, - MAX_EMBEDDING_CACHE_SIZE - ); - (0..NUM_SHARDS) - .map(|_| CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, ttl)))) - .collect() - }; - - let client = MmEmbeddingsClient { - shards, - num_shards: NUM_SHARDS, + MmEmbeddingsClient::in_process() }; if !is_writer { @@ -346,6 +409,41 @@ mod tests { use crate::snapshot::EmbeddingRecord; use std::time::Instant; + #[tokio::test] + async fn in_process_fetch_hits_and_misses() { + const TWEPOCH_MS: u64 = 1_288_834_974_657; + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64; + let hit_id = ((now_ms - TWEPOCH_MS) << 22) | 1; + let miss_id = ((now_ms - TWEPOCH_MS) << 22) | 2; + let client = MmEmbeddingsClient::in_process(); + client + .preload_embeddings( + vec![EmbeddingRecord { + post_id: hit_id, + embedding: vec![f16::from_f32(0.5); 4], + }], + false, + ) + .await; + let hit = client + .fetch_mm_embeddings_sync(vec![hit_id, miss_id], 4, 2) + .unwrap(); + assert_eq!(hit.len(), 8); + assert_eq!(f32::from(hit[0]), 0.5); + assert_eq!(f32::from(hit[4]), 0.0); + + let mut dest = vec![f16::from_f32(9.0); 12]; + client + .fetch_mm_embeddings_into(&[hit_id, miss_id], 4, &mut dest) + .unwrap(); + assert_eq!(f32::from(dest[0]), 0.5); + assert_eq!(f32::from(dest[4]), 9.0); + assert_eq!(f32::from(dest[8]), 9.0); + } + #[tokio::test] #[ignore] async fn bench_serving() { diff --git a/phoenix/crates/serving/xai-recsys-mm-server/src/snapshot.rs b/phoenix/crates/serving/xai-recsys-mm-server/src/snapshot.rs index ddb97469..6968d905 100644 --- a/phoenix/crates/serving/xai-recsys-mm-server/src/snapshot.rs +++ b/phoenix/crates/serving/xai-recsys-mm-server/src/snapshot.rs @@ -505,10 +505,7 @@ pub async fn fetch_recent_embeddings_from_o2( tokio::spawn(async move { let poll_interval = Duration::from_secs(30); - log::info!( - "O2 watcher task started: polling every {:?} for new embedding files", - poll_interval - ); + log::info!("O2 watcher task started: polling every {:?}", poll_interval); loop { tokio::time::sleep(poll_interval).await; @@ -532,17 +529,42 @@ pub async fn fetch_recent_embeddings_from_o2( snap_time ); - match read_embeddings_from_o2(&mm_fetcher, &obj.location).await { - Ok(embeddings) => { - let count = embeddings.len(); - log::info!("Loaded {} new embeddings from {}", count, path_str); - SNAPSHOT_READ_SUCCESS.inc(); - SNAPSHOT_READ_EMBEDDINGS_TOTAL.inc_by(count as u64); - - mm_client.preload_embeddings(embeddings, false).await; + match mm_fetcher.get_raw(&obj.location).await { + Ok(data) => { + let client = mm_client.clone(); + let parse_path = path_str.clone(); + match tokio::task::spawn_blocking(move || { + client.ingest_parquet_bytes_sync(data, &parse_path) + }) + .await + { + Ok(Ok(count)) => { + log::info!( + "Loaded {} new embeddings from {}", + count, + path_str + ); + SNAPSHOT_READ_SUCCESS.inc(); + SNAPSHOT_READ_EMBEDDINGS_TOTAL.inc_by(count as u64); + } + Ok(Err(e)) => { + log::error!( + "Failed to parse/insert embeddings from {}: {e:#}", + path_str + ); + SNAPSHOT_READ_FAILURE.inc(); + } + Err(e) => { + log::error!( + "Blocking ingest panicked for {}: {e:#}", + path_str + ); + SNAPSHOT_READ_FAILURE.inc(); + } + } } Err(e) => { - log::error!("Failed to read embeddings from {}: {:#}", path_str, e); + log::error!("Failed to read embeddings from {}: {e:#}", path_str); SNAPSHOT_READ_FAILURE.inc(); } } @@ -551,7 +573,7 @@ pub async fn fetch_recent_embeddings_from_o2( } } Err(e) => { - log::error!("Failed to list O2 objects: {:#}", e); + log::error!("Failed to list O2 objects: {e:#}"); } } diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index ca21eecf..2505571e 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -110,6 +110,36 @@ message PredictNextActionsRequest { optional uint64 seedTweetId = 16; repeated uint32 requestedContinuousActionIndices = 17; + + repeated ConvAssetIds conv_asset_ids = 18; + + PageDecodeParams pageDecode = 19; + + bool returnBackboneScores = 20; + + map experiment_overrides = 21; +} + +message PageDecodeParams { + uint32 poolN = 1; + + uint32 steps = 2; + + float temperature = 3; +} + +message PageDecodeResult { + repeated uint32 pickedIndices = 1; + + repeated float pickedProbs = 2; + + repeated float pickedScores = 3; +} + +message ConvAssetIds { + repeated int64 impressed_time_ms = 1; + repeated int64 author_id = 2; + repeated int64 asset_id = 3; } message PredictNextActionsResponse { @@ -192,6 +222,8 @@ message CandidateDistributionSet { string checkpointPath = 3; string intermediateTensorsPath = 4; + + PageDecodeResult pageDecodeResult = 5; } message ScoredCandidates { @@ -224,6 +256,19 @@ message NextActionDistribution { repeated float continuousActionsValues = 6; map indexToContinuousValues = 7; + + RewardOutputs rewardOutputs = 8; + + SlateContext slateContext = 9; + + repeated float backboneTopLogProbs = 10; + repeated float backboneContinuousValues = 11; +} + +message RewardOutputs { + repeated float values = 1; + + repeated float policyLogits = 2; } enum ActionName { @@ -311,7 +356,7 @@ enum ActionName { CLIENT_TWEET_VIDEO_MRC_VIEW = 81; CLIENT_TWEET_VIDEO_VIEW = 82; CLIENT_TWEET_VIDEO_VIEW_THRESHOLD = 83; - PLACE_HOLDER_84 = 84; + CLIENT_TWEET_VIDEO_6SEC_VIEW = 84; CLIENT_TWEET_VIDEO_WATCH_TIME = 85; CLIENT_EXTERNAL_LINK_SESSION_LESS_THAN_3_SEC = 86; CLIENT_EXTERNAL_LINK_SESSION_LESS_THAN_5_SEC = 87; @@ -332,6 +377,8 @@ enum ActionName { BRAND_ENGAGEMENT = 102; CLIENT_TWEET_DETAIL_DWELL = 103; CLIENT_TWEET_SEARCH_QUERY_REFORMULATED = 104; + CLIENT_TWEET_RELEVANT_TO_SEARCH = 105; + CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH = 106; ADS_WEB_CONVERSION = 128; ADS_VIDEO_VIEW = 129; @@ -366,71 +413,56 @@ enum ActionName { CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_15_SEC = 158; CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_20_SEC = 159; CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_25_SEC = 160; - CLIENT_TWEET_RELEVANT_TO_SEARCH = 161; - CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH = 162; + CLIENT_TWEET_RELEVANT_TO_SEARCH_DEPRECATED = 161; + CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH_DEPRECATED = 162; CLIENT_TWEET_REPLY_DOWNVOTE = 163; ADS_LONG_DWELL_10_SEC_AND_ATTRIBUTED_CONVERSION = 164; ADS_LONG_DWELL_15_SEC_AND_ATTRIBUTED_CONVERSION = 165; ADS_LONG_SITE_DWELL_CONVERSION = 166; - PLACE_HOLDER_167 = 167; - PLACE_HOLDER_168 = 168; - PLACE_HOLDER_169 = 169; - PLACE_HOLDER_170 = 170; + P_OPEN_LINK_P10 = 167; + P_OPEN_LINK_P25 = 168; + P_OPEN_LINK_P50 = 169; + P_OPEN_LINK_P75 = 170; + ADS_ATTRIBUTED_MACT_PURCHASE = 171; + ADS_ATTRIBUTED_MACT_ADD_TO_CART = 172; + ADS_ATTRIBUTED_MACT_LEVEL_ACHIEVED = 173; + ADS_ATTRIBUTED_MACT_TUTORIAL_COMPLETE = 174; + ADS_ATTRIBUTED_MACT_SIGN_UP = 175; + ADS_ATTRIBUTED_MACT_CUSTOM = 176; + ADS_ATTRIBUTED_MACT_VIEW_PURCHASE = 177; + ADS_ATTRIBUTED_MACT_VIEW_ADD_TO_CART = 178; + ADS_ATTRIBUTED_MACT_VIEW_LEVEL_ACHIEVED = 179; + ADS_ATTRIBUTED_MACT_VIEW_TUTORIAL_COMPLETE = 180; + ADS_ATTRIBUTED_MACT_VIEW_SIGN_UP = 181; + ADS_ATTRIBUTED_MACT_VIEW_CUSTOM = 182; + ADS_PURCHASE_CONVERSION_VIEW_THROUGH = 183; + ADS_ADD_TO_CART_CONVERSION_VIEW_THROUGH = 184; + ADS_CHECKOUT_INITIATED_CONVERSION_VIEW_THROUGH = 185; + ADS_SIGN_UP_CONVERSION_VIEW_THROUGH = 186; + ADS_SITE_VISIT_CONVERSION_VIEW_THROUGH = 187; + ADS_SESSION_CONVERSION_VIEW_THROUGH = 188; + ADS_LANDING_PAGE_VIEW_CONVERSION_VIEW_THROUGH = 189; + ADS_UPPER_FUNNEL_CONVERSION_VIEW_THROUGH = 190; + ADS_PIXEL_FIRE = 191; + ADS_LONG_DWELL_AND_PIXEL_FIRE = 192; + P_LONG_DWELL_AND_PIXEL_FIRE = 193; ADS_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; ADS_UPPER_FUNNEL_CONVERSION = 203; - ADS_RAW_ENGAGEMENT_TYPE_DISPLAYED = 204; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_MRC_VIEW = 205; - ADS_RAW_ENGAGEMENT_TYPE_CAROUSEL_SWIPE_NEXT = 206; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_VIEW_V2 = 207; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_6SEC_VIEW = 208; - ADS_RAW_ENGAGEMENT_TYPE_CARD_URL_CLICK = 209; - ADS_RAW_ENGAGEMENT_TYPE_CAROUSEL_SWIPE_PREVIOUS = 210; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_SHORT_FORM_COMPLETE = 211; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_PLAYBACK_START = 212; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_MRC_VIEW = 213; - ADS_RAW_ENGAGEMENT_TYPE_DISMISS_WITHOUT_REASON = 214; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_VIEW = 215; - ADS_RAW_ENGAGEMENT_TYPE_CLOSE_WEBVIEW = 216; - ADS_RAW_ENGAGEMENT_TYPE_CARD_APP_INSTALL_ATTEMPT = 217; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_6SEC_VIEW = 218; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_SHORT_FORM_COMPLETE = 219; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_SHORT = 220; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_MEDIUM = 221; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_LONG = 222; - ADS_RAW_ENGAGEMENT_TYPE_BLOCK = 223; - ADS_RAW_ENGAGEMENT_TYPE_CARD_CLICK = 224; - ADS_RAW_ENGAGEMENT_TYPE_CARD_OPEN_APP = 225; - ADS_RAW_ENGAGEMENT_TYPE_DETAIL = 226; - ADS_RAW_ENGAGEMENT_TYPE_EMBEDDED_MEDIA = 227; - ADS_RAW_ENGAGEMENT_TYPE_FAV = 228; - ADS_RAW_ENGAGEMENT_TYPE_FOLLOW = 229; - ADS_RAW_ENGAGEMENT_TYPE_HASHTAG = 230; - ADS_RAW_ENGAGEMENT_TYPE_MUTE = 231; - ADS_RAW_ENGAGEMENT_TYPE_POLL_CARD_VOTE = 232; - ADS_RAW_ENGAGEMENT_TYPE_PROFILE_PIC = 233; - ADS_RAW_ENGAGEMENT_TYPE_REPLY = 234; - ADS_RAW_ENGAGEMENT_TYPE_REPORT = 235; - ADS_RAW_ENGAGEMENT_TYPE_RETWEET = 236; - ADS_RAW_ENGAGEMENT_TYPE_SCREEN_NAME = 237; - ADS_RAW_ENGAGEMENT_TYPE_SCROLL_WEBVIEW = 238; - ADS_RAW_ENGAGEMENT_TYPE_SEND = 239; - ADS_RAW_ENGAGEMENT_TYPE_SPOTLIGHT_CLICK = 240; - ADS_RAW_ENGAGEMENT_TYPE_SPOTLIGHT_VIEW = 241; - ADS_RAW_ENGAGEMENT_TYPE_UNBLOCK = 242; - ADS_RAW_ENGAGEMENT_TYPE_UNFAV = 243; - ADS_RAW_ENGAGEMENT_TYPE_UNFOLLOW = 244; - ADS_RAW_ENGAGEMENT_TYPE_UNIFIED_CARD = 245; - ADS_RAW_ENGAGEMENT_TYPE_URL = 246; - ADS_RAW_ENGAGEMENT_TYPE_USER_NAME = 247; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_CTA_URL_CLICK = 248; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_CTA_WATCH_CLICK = 249; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_PLAY_FROM_TAP_V2 = 250; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CTA_URL_CLICK = 251; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CTA_WATCH_CLICK = 252; - PLACE_HOLDER_253 = 253; - PLACE_HOLDER_254 = 254; + ADS_ATTRIBUTED_KEY_CLICK_CONVERSION_DELAYED = 204; + ADS_ATTRIBUTED_CLICK_CONVERSION_DELAYED = 205; + ADS_PURCHASE_CONVERSION_DELAYED = 206; + ADS_MID_FUNNEL_CONVERSION_DELAYED = 207; + ADS_ADD_TO_CART_CONVERSION_DELAYED = 208; + ADS_UPPER_FUNNEL_CONVERSION_DELAYED = 209; + ADS_WEB_CONVERSION_DELAYED = 210; + ADS_SEARCH_CONVERSION_DELAYED = 211; + ADS_SIGN_UP_CONVERSION_DELAYED = 212; + ADS_CHECKOUT_INITIATED_CONVERSION_DELAYED = 213; + reserved 214 to 252; + P_OPEN_LINK_P90 = 253; + ADS_MMP_CLICK = 254; PLACE_HOLDER_255 = 255; } @@ -1192,6 +1224,9 @@ message ScoreInfo { optional double weightedScore = 2; optional double finalScore = 3; SlateContext slateContext = 4; + optional double rewardRerankSlotProb = 5; + optional uint32 rerankerHeadTag = 6; + optional double pageDecodeSlotProb = 7; } message SlateContext { @@ -1200,6 +1235,18 @@ message SlateContext { optional uint32 poolRankGap = 3; double fatigue = 4; double preDiversityScore = 5; + bool sidKnown = 6; + uint32 sidK1 = 7; + uint32 sidK2 = 8; + uint32 sidK3 = 9; + optional uint32 sidGap1 = 10; + optional uint32 sidGap2 = 11; + optional uint32 sidGap3 = 12; + optional uint32 reconCosMilli = 13; + optional uint32 reconCountAbove = 14; + optional uint32 reconGapAbove = 15; + optional uint32 exactK = 16; + optional uint32 exactGap = 17; } message ActionInfo { @@ -1257,6 +1304,8 @@ message UserActionSequence { AdUserInfo adUserInfo = 7; repeated Mask masks = 8; + + optional uint32 rerankerHeadTag = 9; } message UserActionList { @@ -1320,6 +1369,7 @@ message DpaProductInfo { int32 gpc_l1_id = 5; int32 gpc_l2_id = 6; int32 gpc_l3_id = 7; + repeated int32 sid_codes = 8; } message AdInfo { @@ -1549,7 +1599,7 @@ enum UserActionAggregationType { DENSE_WITH_NOT_INTERESTED_IN = 7; DENSE_WITH_LIMITED_QUOTED_ACTION = 8; UPSAMPLED_ADS = 9 [deprecated = true]; - DENSE_WITH_PROFILE_SHORT_DWELL = 10; + DENSE_WITH_PROFILE_LONG_DWELL = 10; NOTIFICATIONS = 11; DENSE_WITH_CONTEXT_FEATURES = 12; NOTIFICATIONS_MORE = 13; @@ -1635,6 +1685,7 @@ message FetchAggregatedUserActionRequests { string client_id = 2; UaasModelType model_type = 3; UaasProductSurface product_surface = 4; + map feature_switch_overrides = 5; } message ChunkMetadata { diff --git a/phoenix/crates/serving/xai-recsys-proto/src/lib.rs b/phoenix/crates/serving/xai-recsys-proto/src/lib.rs index 1b4bb0ef..90fb5cb7 100644 --- a/phoenix/crates/serving/xai-recsys-proto/src/lib.rs +++ b/phoenix/crates/serving/xai-recsys-proto/src/lib.rs @@ -12,6 +12,8 @@ pub mod grok_topics; pub mod installed_apps; pub mod starter_packs; +pub const SAFETY_BIT_AUTHOR_NSFW: u64 = 1 << 2; + pub fn timezone_string_to_enum(tz: &str) -> Timezone { match tz { "" | "unknown" => Timezone::Unknown, diff --git a/phoenix/crates/storage/xai-o2/Cargo.toml b/phoenix/crates/storage/xai-o2/Cargo.toml new file mode 100644 index 00000000..2c305f10 --- /dev/null +++ b/phoenix/crates/storage/xai-o2/Cargo.toml @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +[package] +name = "xai-o2" +version = "0.1.0" +edition.workspace = true +description = "Client interface for the xai-o2 (S3-compatible) object storage backend" +license = "MIT" + +[dependencies] +anyhow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } + +[lints] +workspace = true diff --git a/phoenix/crates/storage/xai-o2/src/base_client.rs b/phoenix/crates/storage/xai-o2/src/base_client.rs new file mode 100644 index 00000000..394c44d6 --- /dev/null +++ b/phoenix/crates/storage/xai-o2/src/base_client.rs @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 X.AI Corp. +use std::fmt::Debug; + +use anyhow::Result; +use async_trait::async_trait; +use bytes::Bytes; + +#[async_trait] +pub trait BaseO2Client: Debug + Send + Sync { + async fn put(&self, key: &str, data: Bytes) -> Result<()>; +} diff --git a/phoenix/crates/storage/xai-o2/src/lib.rs b/phoenix/crates/storage/xai-o2/src/lib.rs new file mode 100644 index 00000000..cbf8497a --- /dev/null +++ b/phoenix/crates/storage/xai-o2/src/lib.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 X.AI Corp. +pub mod base_client; +mod o2_client_builder; + +pub use base_client::BaseO2Client; +pub use o2_client_builder::O2ClientBuilder; diff --git a/phoenix/crates/storage/xai-o2/src/o2_client_builder.rs b/phoenix/crates/storage/xai-o2/src/o2_client_builder.rs new file mode 100644 index 00000000..73f79aa0 --- /dev/null +++ b/phoenix/crates/storage/xai-o2/src/o2_client_builder.rs @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 X.AI Corp. +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; + +use crate::base_client::BaseO2Client; + +pub const BACKEND_UNAVAILABLE: &str = + "the xai-o2 backend is not included in this build; use the default object_store backend"; + +#[derive(Debug, Default, Clone)] +pub struct O2ClientBuilder {} + +impl O2ClientBuilder { + pub fn endpoint(&mut self, _value: impl Into) -> &mut Self { + self + } + + pub fn bucket(&mut self, _value: impl Into) -> &mut Self { + self + } + + pub fn prefix(&mut self, _value: impl Into) -> &mut Self { + self + } + + pub fn access_key_id(&mut self, _value: impl Into) -> &mut Self { + self + } + + pub fn secret_access_key(&mut self, _value: impl Into) -> &mut Self { + self + } + + pub fn allow_anonymous(&mut self, _value: bool) -> &mut Self { + self + } + + pub fn timeout(&mut self, _value: Duration) -> &mut Self { + self + } + + pub fn io_timeout(&mut self, _value: Duration) -> &mut Self { + self + } + + pub fn retry_max_times(&mut self, _value: usize) -> &mut Self { + self + } + + pub fn retry_min_delay(&mut self, _value: Duration) -> &mut Self { + self + } + + pub fn retry_max_delay(&mut self, _value: Duration) -> &mut Self { + self + } + + pub fn write_chunk_size_mib(&mut self, _value: usize) -> &mut Self { + self + } + + pub fn write_chunk_concurrency(&mut self, _value: usize) -> &mut Self { + self + } + + pub fn max_concurrent_write_requests(&mut self, _value: usize) -> &mut Self { + self + } + + pub fn trace_always(&mut self, _value: bool) -> &mut Self { + self + } + + pub fn build(&self) -> Result> { + Err(anyhow::anyhow!(BACKEND_UNAVAILABLE)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn build_fails_loudly_after_full_configuration() { + let mut builder = O2ClientBuilder::default(); + builder + .endpoint("http://localhost:9000") + .bucket("some-bucket") + .prefix("") + .write_chunk_size_mib(8) + .write_chunk_concurrency(16) + .max_concurrent_write_requests(64) + .retry_max_times(6) + .retry_min_delay(Duration::from_secs(1)) + .retry_max_delay(Duration::from_secs(32)) + .timeout(Duration::from_secs(300)) + .io_timeout(Duration::from_secs(300)) + .trace_always(false); + builder.allow_anonymous(true); + builder.access_key_id("key").secret_access_key("secret"); + + let err = builder.build().expect_err("build must fail in this build"); + assert_eq!(err.to_string(), BACKEND_UNAVAILABLE); + } + + #[test] + fn build_error_converts_to_boxed_std_error() { + let err: Box = + O2ClientBuilder::default().build().unwrap_err().into(); + assert!(err.to_string().contains("object_store")); + } +} diff --git a/phoenix/pyproject.toml b/phoenix/pyproject.toml index c85ef087..eadb289d 100644 --- a/phoenix/pyproject.toml +++ b/phoenix/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "optax>=0.2.4", "chex>=0.1.86", "numba>=0.60.0", + "flatbuffers", "grpcio>=1.66.0", "protobuf>=5.28.0", diff --git a/phoenix/python/common/xai-configlib/src/xai_configlib/__init__.py b/phoenix/python/common/xai-configlib/src/xai_configlib/__init__.py index 10b9de19..13543352 100644 --- a/phoenix/python/common/xai-configlib/src/xai_configlib/__init__.py +++ b/phoenix/python/common/xai-configlib/src/xai_configlib/__init__.py @@ -4,6 +4,7 @@ import dataclasses import datetime import enum +import functools import importlib import inspect import json @@ -288,20 +289,21 @@ def to_dict(obj: Any, container: str = "???") -> Any: return obj -def resolve_type_hints(dcls: Any) -> dict[str, Type[Any]]: - if isinstance(dcls, Config): - resolved_hints = type(dcls).get_type_hints() - elif isinstance(dcls, type): - if dataclasses.is_dataclass(dcls): - resolved_hints = get_type_hints(dcls) - field_names = [field.name for field in dataclasses.fields(dcls)] - return {name: resolved_hints[name] for name in field_names} - else: - return {} +@functools.cache +def _resolve_type_hints(cls: type) -> dict[str, Type[Any]]: + if issubclass(cls, Config): + resolved_hints = cls.get_type_hints() + elif dataclasses.is_dataclass(cls): + resolved_hints = get_type_hints(cls) else: - resolved_hints = get_type_hints(type(dcls)) - field_names = [field.name for field in dataclasses.fields(dcls)] - return {name: resolved_hints[name] for name in field_names} + return {} + return {field.name: resolved_hints[field.name] for field in dataclasses.fields(cls)} + + +def resolve_type_hints(dcls: Any) -> dict[str, Type[Any]]: + if isinstance(dcls, type): + return _resolve_type_hints(dcls) + return _resolve_type_hints(type(dcls)) def replace_cli_subs( diff --git a/phoenix/python/common/xai-proto/hatch_build.py b/phoenix/python/common/xai-proto/hatch_build.py index 5e348744..26b84d87 100644 --- a/phoenix/python/common/xai-proto/hatch_build.py +++ b/phoenix/python/common/xai-proto/hatch_build.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import os import pathlib import subprocess import sys @@ -31,11 +32,18 @@ def initialize(self, version, build_data): def run_protoc(out_dir: Path) -> None: + skip_missing = os.environ.get("XAI_PROTO_SKIP_MISSING", "") == "1" xai_root = pathlib.Path(__file__).parent / PARENT_DIR includes = [] protos = [] for proto in PROTOS: p = xai_root / proto + if skip_missing and not p.exists(): + print( + f"xai-proto hatch_build: SKIPPING {proto} (absent; XAI_PROTO_SKIP_MISSING=1)", + file=sys.stderr, + ) + continue includes.append(f"-I{p.parent}") protos.append(str(p)) diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index ca21eecf..2505571e 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -110,6 +110,36 @@ message PredictNextActionsRequest { optional uint64 seedTweetId = 16; repeated uint32 requestedContinuousActionIndices = 17; + + repeated ConvAssetIds conv_asset_ids = 18; + + PageDecodeParams pageDecode = 19; + + bool returnBackboneScores = 20; + + map experiment_overrides = 21; +} + +message PageDecodeParams { + uint32 poolN = 1; + + uint32 steps = 2; + + float temperature = 3; +} + +message PageDecodeResult { + repeated uint32 pickedIndices = 1; + + repeated float pickedProbs = 2; + + repeated float pickedScores = 3; +} + +message ConvAssetIds { + repeated int64 impressed_time_ms = 1; + repeated int64 author_id = 2; + repeated int64 asset_id = 3; } message PredictNextActionsResponse { @@ -192,6 +222,8 @@ message CandidateDistributionSet { string checkpointPath = 3; string intermediateTensorsPath = 4; + + PageDecodeResult pageDecodeResult = 5; } message ScoredCandidates { @@ -224,6 +256,19 @@ message NextActionDistribution { repeated float continuousActionsValues = 6; map indexToContinuousValues = 7; + + RewardOutputs rewardOutputs = 8; + + SlateContext slateContext = 9; + + repeated float backboneTopLogProbs = 10; + repeated float backboneContinuousValues = 11; +} + +message RewardOutputs { + repeated float values = 1; + + repeated float policyLogits = 2; } enum ActionName { @@ -311,7 +356,7 @@ enum ActionName { CLIENT_TWEET_VIDEO_MRC_VIEW = 81; CLIENT_TWEET_VIDEO_VIEW = 82; CLIENT_TWEET_VIDEO_VIEW_THRESHOLD = 83; - PLACE_HOLDER_84 = 84; + CLIENT_TWEET_VIDEO_6SEC_VIEW = 84; CLIENT_TWEET_VIDEO_WATCH_TIME = 85; CLIENT_EXTERNAL_LINK_SESSION_LESS_THAN_3_SEC = 86; CLIENT_EXTERNAL_LINK_SESSION_LESS_THAN_5_SEC = 87; @@ -332,6 +377,8 @@ enum ActionName { BRAND_ENGAGEMENT = 102; CLIENT_TWEET_DETAIL_DWELL = 103; CLIENT_TWEET_SEARCH_QUERY_REFORMULATED = 104; + CLIENT_TWEET_RELEVANT_TO_SEARCH = 105; + CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH = 106; ADS_WEB_CONVERSION = 128; ADS_VIDEO_VIEW = 129; @@ -366,71 +413,56 @@ enum ActionName { CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_15_SEC = 158; CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_20_SEC = 159; CLIENT_EXTERNAL_LINK_SESSION_MORE_THAN_25_SEC = 160; - CLIENT_TWEET_RELEVANT_TO_SEARCH = 161; - CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH = 162; + CLIENT_TWEET_RELEVANT_TO_SEARCH_DEPRECATED = 161; + CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH_DEPRECATED = 162; CLIENT_TWEET_REPLY_DOWNVOTE = 163; ADS_LONG_DWELL_10_SEC_AND_ATTRIBUTED_CONVERSION = 164; ADS_LONG_DWELL_15_SEC_AND_ATTRIBUTED_CONVERSION = 165; ADS_LONG_SITE_DWELL_CONVERSION = 166; - PLACE_HOLDER_167 = 167; - PLACE_HOLDER_168 = 168; - PLACE_HOLDER_169 = 169; - PLACE_HOLDER_170 = 170; + P_OPEN_LINK_P10 = 167; + P_OPEN_LINK_P25 = 168; + P_OPEN_LINK_P50 = 169; + P_OPEN_LINK_P75 = 170; + ADS_ATTRIBUTED_MACT_PURCHASE = 171; + ADS_ATTRIBUTED_MACT_ADD_TO_CART = 172; + ADS_ATTRIBUTED_MACT_LEVEL_ACHIEVED = 173; + ADS_ATTRIBUTED_MACT_TUTORIAL_COMPLETE = 174; + ADS_ATTRIBUTED_MACT_SIGN_UP = 175; + ADS_ATTRIBUTED_MACT_CUSTOM = 176; + ADS_ATTRIBUTED_MACT_VIEW_PURCHASE = 177; + ADS_ATTRIBUTED_MACT_VIEW_ADD_TO_CART = 178; + ADS_ATTRIBUTED_MACT_VIEW_LEVEL_ACHIEVED = 179; + ADS_ATTRIBUTED_MACT_VIEW_TUTORIAL_COMPLETE = 180; + ADS_ATTRIBUTED_MACT_VIEW_SIGN_UP = 181; + ADS_ATTRIBUTED_MACT_VIEW_CUSTOM = 182; + ADS_PURCHASE_CONVERSION_VIEW_THROUGH = 183; + ADS_ADD_TO_CART_CONVERSION_VIEW_THROUGH = 184; + ADS_CHECKOUT_INITIATED_CONVERSION_VIEW_THROUGH = 185; + ADS_SIGN_UP_CONVERSION_VIEW_THROUGH = 186; + ADS_SITE_VISIT_CONVERSION_VIEW_THROUGH = 187; + ADS_SESSION_CONVERSION_VIEW_THROUGH = 188; + ADS_LANDING_PAGE_VIEW_CONVERSION_VIEW_THROUGH = 189; + ADS_UPPER_FUNNEL_CONVERSION_VIEW_THROUGH = 190; + ADS_PIXEL_FIRE = 191; + ADS_LONG_DWELL_AND_PIXEL_FIRE = 192; + P_LONG_DWELL_AND_PIXEL_FIRE = 193; ADS_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; ADS_UPPER_FUNNEL_CONVERSION = 203; - ADS_RAW_ENGAGEMENT_TYPE_DISPLAYED = 204; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_MRC_VIEW = 205; - ADS_RAW_ENGAGEMENT_TYPE_CAROUSEL_SWIPE_NEXT = 206; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_VIEW_V2 = 207; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_6SEC_VIEW = 208; - ADS_RAW_ENGAGEMENT_TYPE_CARD_URL_CLICK = 209; - ADS_RAW_ENGAGEMENT_TYPE_CAROUSEL_SWIPE_PREVIOUS = 210; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_SHORT_FORM_COMPLETE = 211; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_PLAYBACK_START = 212; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_MRC_VIEW = 213; - ADS_RAW_ENGAGEMENT_TYPE_DISMISS_WITHOUT_REASON = 214; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_VIEW = 215; - ADS_RAW_ENGAGEMENT_TYPE_CLOSE_WEBVIEW = 216; - ADS_RAW_ENGAGEMENT_TYPE_CARD_APP_INSTALL_ATTEMPT = 217; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_6SEC_VIEW = 218; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_SHORT_FORM_COMPLETE = 219; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_SHORT = 220; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_MEDIUM = 221; - ADS_RAW_ENGAGEMENT_TYPE_DWELL_LONG = 222; - ADS_RAW_ENGAGEMENT_TYPE_BLOCK = 223; - ADS_RAW_ENGAGEMENT_TYPE_CARD_CLICK = 224; - ADS_RAW_ENGAGEMENT_TYPE_CARD_OPEN_APP = 225; - ADS_RAW_ENGAGEMENT_TYPE_DETAIL = 226; - ADS_RAW_ENGAGEMENT_TYPE_EMBEDDED_MEDIA = 227; - ADS_RAW_ENGAGEMENT_TYPE_FAV = 228; - ADS_RAW_ENGAGEMENT_TYPE_FOLLOW = 229; - ADS_RAW_ENGAGEMENT_TYPE_HASHTAG = 230; - ADS_RAW_ENGAGEMENT_TYPE_MUTE = 231; - ADS_RAW_ENGAGEMENT_TYPE_POLL_CARD_VOTE = 232; - ADS_RAW_ENGAGEMENT_TYPE_PROFILE_PIC = 233; - ADS_RAW_ENGAGEMENT_TYPE_REPLY = 234; - ADS_RAW_ENGAGEMENT_TYPE_REPORT = 235; - ADS_RAW_ENGAGEMENT_TYPE_RETWEET = 236; - ADS_RAW_ENGAGEMENT_TYPE_SCREEN_NAME = 237; - ADS_RAW_ENGAGEMENT_TYPE_SCROLL_WEBVIEW = 238; - ADS_RAW_ENGAGEMENT_TYPE_SEND = 239; - ADS_RAW_ENGAGEMENT_TYPE_SPOTLIGHT_CLICK = 240; - ADS_RAW_ENGAGEMENT_TYPE_SPOTLIGHT_VIEW = 241; - ADS_RAW_ENGAGEMENT_TYPE_UNBLOCK = 242; - ADS_RAW_ENGAGEMENT_TYPE_UNFAV = 243; - ADS_RAW_ENGAGEMENT_TYPE_UNFOLLOW = 244; - ADS_RAW_ENGAGEMENT_TYPE_UNIFIED_CARD = 245; - ADS_RAW_ENGAGEMENT_TYPE_URL = 246; - ADS_RAW_ENGAGEMENT_TYPE_USER_NAME = 247; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_CTA_URL_CLICK = 248; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_AD_CTA_WATCH_CLICK = 249; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CONTENT_PLAY_FROM_TAP_V2 = 250; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CTA_URL_CLICK = 251; - ADS_RAW_ENGAGEMENT_TYPE_VIDEO_CTA_WATCH_CLICK = 252; - PLACE_HOLDER_253 = 253; - PLACE_HOLDER_254 = 254; + ADS_ATTRIBUTED_KEY_CLICK_CONVERSION_DELAYED = 204; + ADS_ATTRIBUTED_CLICK_CONVERSION_DELAYED = 205; + ADS_PURCHASE_CONVERSION_DELAYED = 206; + ADS_MID_FUNNEL_CONVERSION_DELAYED = 207; + ADS_ADD_TO_CART_CONVERSION_DELAYED = 208; + ADS_UPPER_FUNNEL_CONVERSION_DELAYED = 209; + ADS_WEB_CONVERSION_DELAYED = 210; + ADS_SEARCH_CONVERSION_DELAYED = 211; + ADS_SIGN_UP_CONVERSION_DELAYED = 212; + ADS_CHECKOUT_INITIATED_CONVERSION_DELAYED = 213; + reserved 214 to 252; + P_OPEN_LINK_P90 = 253; + ADS_MMP_CLICK = 254; PLACE_HOLDER_255 = 255; } @@ -1192,6 +1224,9 @@ message ScoreInfo { optional double weightedScore = 2; optional double finalScore = 3; SlateContext slateContext = 4; + optional double rewardRerankSlotProb = 5; + optional uint32 rerankerHeadTag = 6; + optional double pageDecodeSlotProb = 7; } message SlateContext { @@ -1200,6 +1235,18 @@ message SlateContext { optional uint32 poolRankGap = 3; double fatigue = 4; double preDiversityScore = 5; + bool sidKnown = 6; + uint32 sidK1 = 7; + uint32 sidK2 = 8; + uint32 sidK3 = 9; + optional uint32 sidGap1 = 10; + optional uint32 sidGap2 = 11; + optional uint32 sidGap3 = 12; + optional uint32 reconCosMilli = 13; + optional uint32 reconCountAbove = 14; + optional uint32 reconGapAbove = 15; + optional uint32 exactK = 16; + optional uint32 exactGap = 17; } message ActionInfo { @@ -1257,6 +1304,8 @@ message UserActionSequence { AdUserInfo adUserInfo = 7; repeated Mask masks = 8; + + optional uint32 rerankerHeadTag = 9; } message UserActionList { @@ -1320,6 +1369,7 @@ message DpaProductInfo { int32 gpc_l1_id = 5; int32 gpc_l2_id = 6; int32 gpc_l3_id = 7; + repeated int32 sid_codes = 8; } message AdInfo { @@ -1549,7 +1599,7 @@ enum UserActionAggregationType { DENSE_WITH_NOT_INTERESTED_IN = 7; DENSE_WITH_LIMITED_QUOTED_ACTION = 8; UPSAMPLED_ADS = 9 [deprecated = true]; - DENSE_WITH_PROFILE_SHORT_DWELL = 10; + DENSE_WITH_PROFILE_LONG_DWELL = 10; NOTIFICATIONS = 11; DENSE_WITH_CONTEXT_FEATURES = 12; NOTIFICATIONS_MORE = 13; @@ -1635,6 +1685,7 @@ message FetchAggregatedUserActionRequests { string client_id = 2; UaasModelType model_type = 3; UaasProductSurface product_surface = 4; + map feature_switch_overrides = 5; } message ChunkMetadata { diff --git a/phoenix/python/common/xai-proto/pyproject.toml b/phoenix/python/common/xai-proto/pyproject.toml index 9367d42c..7abf73d8 100644 --- a/phoenix/python/common/xai-proto/pyproject.toml +++ b/phoenix/python/common/xai-proto/pyproject.toml @@ -1,14 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. [project] -dependencies = ["grpcio-tools>=1.80,<1.82", "protobuf>=6.31.1,<7"] +dependencies = ["grpcio-tools==1.83.0", "protobuf>=7.36.0"] name = "xai-proto" version = "0.0.2" [build-system] backend-path = ["."] build-backend = "backend_proxy" -requires = ["hatchling>=1.20", "grpcio-tools>=1.80,<1.82"] +requires = ["hatchling>=1.20", "grpcio-tools==1.83.0"] [tool.hatch.build.hooks.custom] path = "hatch_build.py" diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py new file mode 100644 index 00000000..de12f3fc --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import json +import logging +import os +import pathlib +import random +import tempfile +import time + +rank_logger = logging.getLogger("rank") + +TREE_DEK_NAME = "_DEK" +TREE_DEK_CLAIM_NAME = "_DEK.claim" +PUBLISH_TIMEOUT_SECS = 120.0 + +_TRANSIENT_KMS_MARKERS = ("429", "502", "503", "504", "KMS transport") +_UNWRAP_ATTEMPTS = 4 +_UNWRAP_BACKOFF_CAP_SECS = 2.0 + + +def _unwrap_with_backoff(kms_client, dek_path) -> bytes: + import xai_kms + + slept = 0.0 + for attempt in range(_UNWRAP_ATTEMPTS): + try: + return bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, dek_path)) + except Exception as error: + transient = any(m in str(error) for m in _TRANSIENT_KMS_MARKERS) + remaining = _UNWRAP_BACKOFF_CAP_SECS - slept + if not transient or attempt == _UNWRAP_ATTEMPTS - 1 or remaining <= 0: + raise + delay = min(0.2 * 2**attempt, 1.0, remaining) * (0.5 + random.random()) + delay = min(delay, remaining) + rank_logger.warning( + "KMS unwrap of %s failed transiently (%s); retrying in %.1fs (attempt %d/%d)", + dek_path, + error, + delay, + attempt + 1, + _UNWRAP_ATTEMPTS, + ) + time.sleep(delay) + slept += delay + raise AssertionError("unreachable") + + +def _read_wrapped_dek(dek_path: pathlib.Path) -> dict | None: + try: + entry = json.loads(dek_path.read_text()) + return entry if entry.get("wrapped") else None + except (OSError, ValueError): + return None + + +def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict, dict]: + import xai_kms + + dek_path = path / TREE_DEK_NAME + claim = path / TREE_DEK_CLAIM_NAME + + deadline = time.monotonic() + PUBLISH_TIMEOUT_SECS + if _read_wrapped_dek(dek_path) is not None: + rank_logger.info("adopted existing _DEK at %s", path) + while (entry := _read_wrapped_dek(dek_path)) is None: + try: + claim.open("x").close() + won_claim = True + except FileExistsError: + won_claim = False + if won_claim: + try: + if _read_wrapped_dek(dek_path) is None: + xai_kms.nfs.write_shared_dek(kms_client, dek_path) + rank_logger.info( + "minted and KMS-wrapped new DEK at %s (key_id=%s)", + path, + (_read_wrapped_dek(dek_path) or {}).get("key_id"), + ) + else: + rank_logger.info("adopted existing _DEK at %s", path) + except BaseException: + if _read_wrapped_dek(dek_path) is None: + claim.unlink(missing_ok=True) + raise + continue + if time.monotonic() >= deadline: + raise RuntimeError( + f"timed out waiting for the wrapped DEK at {dek_path}; its minter " + "(the rank holding the .claim marker) likely died before publishing" + ) + rank_logger.debug("waiting for wrapped DEK at %s", dek_path) + time.sleep(0.1) + raw = _unwrap_with_backoff(kms_client, dek_path) + return raw, entry["wrapped"], entry.get("context") or {}, entry + + +_ENVELOPE_HEADER_SIZE = 4096 +_ENVELOPE_FIXED_LEN = 36 +_DERIVED_PREFIX = "xai-dek1:" + + +def adopt_tree_dek(path: pathlib.Path, kms_client) -> bytes: + dek_path = path / TREE_DEK_NAME + if dek_path.exists(): + entry = _read_wrapped_dek(dek_path) or {} + raw = _unwrap_with_backoff(kms_client, dek_path) + key_id = entry.get("key_id") + else: + raw = _unwrap_header_master(kms_client, path) + key_id = "header" + rank_logger.info("KMS unwrap OK for %s (key_id=%s)", path, key_id) + return raw + + +def _unwrap_header_master(kms_client, path: pathlib.Path) -> bytes: + wrapped, context = _header_wrapped_master(path) + entry = {"key_id": "header", "context": context, "wrapped": wrapped} + fd, tmp = tempfile.mkstemp(prefix="xai-ckpt-dek-", suffix=".json") + try: + with os.fdopen(fd, "w") as f: + json.dump(entry, f) + return _unwrap_with_backoff(kms_client, tmp) + finally: + os.unlink(tmp) + + +def _header_wrapped_master(path: pathlib.Path) -> tuple[str, dict]: + envelope = _first_envelope(path) + head = envelope.read_bytes()[:_ENVELOPE_HEADER_SIZE] + if len(head) < _ENVELOPE_FIXED_LEN or head[:8] != b"XAIENC01": + raise ValueError(f"not an envelope: {envelope}") + wrapped_len = int.from_bytes(head[28:32], "big") + context_len = int.from_bytes(head[32:36], "big") + start = _ENVELOPE_FIXED_LEN + field = head[start : start + wrapped_len].decode() + context = json.loads(head[start + wrapped_len : start + wrapped_len + context_len]) + if not field.startswith(_DERIVED_PREFIX): + raise ValueError( + f"unsupported wrapped-data-key format at {envelope}: expected " + f"'{_DERIVED_PREFIX}:'; convert the tree to the " + "master-DEK format to load it" + ) + _salt, master = field[len(_DERIVED_PREFIX) :].split(":", 1) + if not master: + raise ValueError(f"malformed xai-dek1 header at {envelope}") + return master, context + + +def _first_envelope(path: pathlib.Path) -> pathlib.Path: + meta = path / "_METADATA" + if meta.is_file(): + with meta.open("rb") as f: + if f.read(8) == b"XAIENC01": + return meta + for dirpath, dirnames, filenames in os.walk(path): + dirnames.sort() + for name in sorted(filenames): + candidate = pathlib.Path(dirpath) / name + try: + with candidate.open("rb") as f: + if f.read(8) == b"XAIENC01": + return candidate + except OSError: + continue + raise FileNotFoundError(f"no envelope in {path}") diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py new file mode 100644 index 00000000..526edbc8 --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import base64 +import logging +import pathlib + +from xai_checkpointing.dek import publish_tree_dek + +rank_logger = logging.getLogger("rank") + + +def at_dir(kvstore: dict, path: pathlib.Path | str) -> dict: + inner = kvstore.get("base") + root = inner.get("path", "/") if isinstance(inner, dict) else "/" + rel = pathlib.Path(path).relative_to(root) + scoped = dict(kvstore) + if rel != pathlib.Path("."): + scoped["path"] = scoped.get("path", "") + rel.as_posix() + "/" + return scoped + + +def use_encrypted_kvstore(ocdbt_config: dict, kvstore: dict) -> dict: + local = ocdbt_config["base"] + assert local["driver"] == "file", local + return {**ocdbt_config, "base": at_dir(kvstore, local["path"])} + + +def _envelope_spec( + checkpoint_root: pathlib.Path | str, + dek_b64: str, + wrapped_dek: str, + chunk_size: int, + encryption_context: dict, +) -> dict: + base = { + "driver": "xai_encrypted", + "base": {"driver": "file", "path": pathlib.Path(checkpoint_root).as_posix() + "/"}, + "dek_b64": dek_b64, + "wrapped_dek": wrapped_dek, + "chunk_size": chunk_size, + } + if encryption_context: + base["encryption_context"] = encryption_context + return base + + +def encrypted_kvstore(path: pathlib.Path, kms_client, encryption_chunk_size: int) -> dict: + path.mkdir(parents=True, exist_ok=True) + raw, wrapped, context, entry = publish_tree_dek(path, kms_client) + rank_logger.info( + "KMS unwrap OK for %s (key_id=%s)", + path, + entry.get("key_id"), + ) + return _envelope_spec( + path, base64.b64encode(raw).decode(), wrapped, encryption_chunk_size, context + ) diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py index 62fba013..99de47bb 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py @@ -1,12 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import base64 +import contextlib +import ctypes +import fcntl import functools +import gc import json import logging +import math import os import pathlib import re import time +from collections.abc import Iterable from typing import Any, Callable import jax @@ -17,15 +24,222 @@ common, fix_jax, ) +from xai_checkpointing.dek import adopt_tree_dek +from xai_checkpointing.encrypted_kvstore import at_dir, use_encrypted_kvstore from xai_checkpointing.tree_util import has_subtree, tree_to_dict import orbax.checkpoint as ocp -logger = logging.getLogger("checkpointing") rank_logger = logging.getLogger("rank") PyTree = common.PyTree +_NODE_SERIALIZE_ENV = "XAI_RESTORE_NODE_SERIALIZE" +_NODE_LOCK_FILE_ENV = "XAI_RESTORE_NODE_LOCK_FILE" +_ENCRYPTION_MAGIC = b"XAIENC01" + + +def _is_encrypted_tree(tree: pathlib.Path) -> bool: + try: + with (tree / "_METADATA").open("rb") as f: + return f.read(len(_ENCRYPTION_MAGIC)) == _ENCRYPTION_MAGIC + except FileNotFoundError: + return False + + +def _names_from_tree_metadata(metadata_json: dict[str, Any]) -> list[str]: + return [ + ".".join(str(k["key"]) for k in v["key_metadata"]) + for v in metadata_json["tree_metadata"].values() + if not v["value_metadata"]["skip_deserialize"] + ] + + +def _read_encrypted_file( + path: pathlib.Path, key: str, kvstore_base: dict[str, Any], ts_context: ts.Context +) -> bytes: + kv = ts.KvStore.open(at_dir(kvstore_base, path), context=ts_context).result() + return kv.read(key).result().value + + +def _encrypted_tspec_transform(kvstore_base: dict[str, Any]) -> Callable: + def transform(tspec: dict[str, Any]) -> dict[str, Any]: + return {**tspec, "kvstore": use_encrypted_kvstore(tspec["kvstore"], kvstore_base)} + + return transform + + +def _prepare_checkpoint_read( + path: pathlib.Path, kms_client: object | None +) -> tuple[dict[str, Any] | None, dict[str, Any], list[str], ts.Context]: + ts_context = ts.Context( + { + "file_io_concurrency": {"limit": 128}, + "cache_pool#ocdbt": {"total_bytes_limit": 100000000}, + } + ) + if not _is_encrypted_tree(path): + with (path / "_METADATA").open() as f: + metadata_json = json.load(f) + checkpoint_names = list( + tree_to_dict(ocp.StandardCheckpointer().metadata(path), keep_none=False).keys() + ) + return None, metadata_json, checkpoint_names, ts_context + + import xai_kms + + if kms_client is None: + kms_client = xai_kms.KmsClient.from_cluster_env() + raw = adopt_tree_dek(path, kms_client) + kvstore_base = { + "driver": "xai_encrypted", + "base": {"driver": "file", "path": path.as_posix() + "/"}, + "dek_b64": base64.b64encode(raw).decode(), + } + metadata_json = json.loads(_read_encrypted_file(path, "_METADATA", kvstore_base, ts_context)) + return kvstore_base, metadata_json, _names_from_tree_metadata(metadata_json), ts_context + + +def _restore_node_serialize_enabled() -> bool: + return os.getenv(_NODE_SERIALIZE_ENV, "1").lower() not in ("0", "", "false") + + +def _node_lock_path() -> str: + path = os.getenv(_NODE_LOCK_FILE_ENV) + if path: + return path + if os.path.isdir("/dev/shm"): + return "/dev/shm/xai_restore_node_lock" + return "/tmp/xai_restore_node_lock" + + +class _NodeBatchLock: + def __init__(self, path: str): + self._path = path + self._fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o666) + + def __enter__(self): + t0 = time.time() + fcntl.flock(self._fd, fcntl.LOCK_EX) + waited = time.time() - t0 + if waited > 1.0: + rank_logger.info( + "restore node-serialize: waited %.1fs for node lock %s", waited, self._path + ) + return self + + def __exit__(self, *exc): + fcntl.flock(self._fd, fcntl.LOCK_UN) + + def close(self): + try: + os.close(self._fd) + except OSError: + pass + + +def _release_batch_memory(): + gc.collect() + try: + ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception: + pass + + +def _build_read_plan( + checkpoint_names: Iterable[str], + host_state: dict[str, jax.Array], + load_mask: dict[str, jax.Array], + rename: Callable[[str], str] | None, +) -> list[tuple[str, str, list[bool], int]]: + plan: list[tuple[str, str, list[bool], int]] = [] + for checkpoint_name in checkpoint_names: + name = checkpoint_name + if rename is not None: + name = rename(checkpoint_name) + + if host_state.get(name) is None: + if not has_subtree(name, host_state): + rank_logger.warning( + "Not loading %r from checkpoint because it's not in the initialized state", name + ) + continue + + if not load_mask: + mask = [True for _ in host_state[name].addressable_shards] + else: + mask = [shard.data.item() for shard in load_mask[name].addressable_shards] + if any(mask): + array = host_state[name] + shard_nbytes = array.dtype.itemsize * math.prod(array.sharding.shard_shape(array.shape)) + plan.append((checkpoint_name, name, mask, shard_nbytes * sum(mask))) + + return plan + + +def _pack_read_batches( + plan: list[tuple[str, str, list[bool], int]], + concurrent_bytes: int | None, +) -> list[list[tuple[str, str, list[bool], int]]]: + batches: list[list[tuple[str, str, list[bool], int]]] = [] + if concurrent_bytes: + cur: list[tuple[str, str, list[bool], int]] = [] + cur_bytes = 0 + for item in plan: + nbytes = item[3] + if cur and cur_bytes + nbytes > concurrent_bytes: + batches.append(cur) + cur, cur_bytes = [], 0 + cur.append(item) + cur_bytes += nbytes + if cur: + batches.append(cur) + elif plan: + batches.append(plan) + return batches + + +def _convert_domains(domains: dict[str, Any], host_state: dict[str, jax.Array]) -> dict[str, Any]: + for name, domain in domains.items(): + assert host_state.get(name) is not None, f"Cannot restrict domain of skipped tensor {name}" + if isinstance(domain, dict): + domains[name] = ts.IndexDomain(json=domain) + elif isinstance(domain, ts.DimExpression): + domains[name] = ts.IndexDomain(shape=host_state[name].shape)[domain] + elif not isinstance(domain, ts.IndexDomain): + raise ValueError(f"Unknown domain: {domain} for {name!r}") + return domains + + +def _open_tensor( + checkpoint_name: str, + name: str, + path: pathlib.Path, + use_zarr3: bool, + ts_context: ts.Context, + dest: jax.Array, + has_domain: bool, + tspec_transform: Callable[[dict[str, Any]], dict[str, Any]] | None, +) -> ts.TensorStore: + info = ocp.type_handlers.ParamInfo( + name=checkpoint_name, + path=path / checkpoint_name, + parent_dir=path, + is_ocdbt_checkpoint=True, + use_zarr3=use_zarr3, + ) + tspec = ocp.type_handlers.get_json_tspec_read(info, use_ocdbt=True) + if tspec_transform is not None: + tspec = tspec_transform(tspec) + t = ts.open(ts.Spec(tspec), open=True, context=ts_context).result() + if not has_domain and tuple(t.shape) != tuple(dest.shape): + raise ValueError( + f"Tensor {name!r}: checkpoint has shape {tuple(t.shape)}, " + f"but initialized state has shape {tuple(dest.shape)}. " + f"Use 'no_loading' to skip this tensor or 'domains' to load a partial slice." + ) + return t + def _read_into_shards( t: ts.TensorStore, @@ -63,6 +277,43 @@ def _read_into_shards( yield dest.write(src).commit +def _drain_read_futures( + futures: dict[Any, tuple[str, str, int]], + arrays: dict[str, jax.Array], + path, + timeout: float, + *, + log_loaded: bool = False, +) -> None: + for future in common._ready(list(futures), timeout=timeout): + try: + future.result() + except Exception: + checkpoint_name, name, _ = futures.pop(future, (None, None, None)) + tensor = arrays.get(name) if name is not None else None + extra = ( + f" (in checkpoint: {checkpoint_name})" + if checkpoint_name is not None and checkpoint_name != name + else "" + ) + detail = f"{tensor.shape=}, {tensor.dtype=}" if tensor is not None else "missing tensor" + rank_logger.exception( + "Checkpoint error from loading %s. Error loading from %s into %s.%s", + path, + name, + detail, + extra, + ) + raise + + checkpoint_name, name, _ = futures.pop(future, (None, None, None)) + if log_loaded and name is not None: + if checkpoint_name != name: + rank_logger.debug("Loaded %s (name in checkpoint: %s)", name, checkpoint_name) + else: + rank_logger.debug("Loaded %s", name) + + def load_checkpoint( path: str, host_state: PyTree[jax.Array], @@ -71,100 +322,282 @@ def load_checkpoint( domains: dict[str, Any] | None = None, tag: str | None = None, timeout: float = 900.0, + concurrent_gb: float | None = None, + kms_client: object | None = None, ): if tag is None: tag = "orbax-ckpt" host_state = tree_to_dict(host_state) load_mask = tree_to_dict(load_mask) - domains = tree_to_dict(domains) - - for name, domain in domains.items(): - assert host_state.get(name) is not None, f"Cannot restrict domain of skipped tensor {name}" - if isinstance(domain, dict): - domains[name] = ts.IndexDomain(json=domain) - elif isinstance(domain, ts.DimExpression): - domains[name] = ts.IndexDomain(shape=host_state[name].shape)[domain] - elif not isinstance(domain, ts.IndexDomain): - raise ValueError(f"Unknown domain: {domain} for {name!r}") + domains = _convert_domains(tree_to_dict(domains), host_state) start = time.time() rank_logger.info("Restoring checkpoint from %s", path) path = pathlib.Path(path) / tag - metadata = ocp.StandardCheckpointer().metadata(path) - with (path / "_METADATA").open() as f: - use_zarr3 = json.load(f)["use_zarr3"] - ts_context = ts.Context( - { - "file_io_concurrency": {"limit": 128}, - "cache_pool#ocdbt": {"total_bytes_limit": 100000000}, - } + kvstore_base, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + path, kms_client ) + use_zarr3 = metadata_json["use_zarr3"] + + plan = _build_read_plan(checkpoint_names, host_state, load_mask, rename) + + concurrent_bytes = int(concurrent_gb * 10**9) if concurrent_gb else None + total_bytes = sum(nbytes for *_, nbytes in plan) + max_leaf = max((nbytes for *_, nbytes in plan), default=0) + + batches = _pack_read_batches(plan, concurrent_bytes) + num_batches = len(batches) + + if concurrent_bytes: + rank_logger.info( + "load_checkpoint: restore read throttle ACTIVE concurrent_gb=%s " + "(batch limit %.2fGiB): total=%.2fGiB in %d tensors, max_leaf=%.2fGiB, " + "num_batches=%d (issue+drain one read batch at a time to bound host transient)", + concurrent_gb, + concurrent_bytes / (1 << 30), + total_bytes / (1 << 30), + len(plan), + max_leaf / (1 << 30), + num_batches, + ) + if max_leaf > concurrent_bytes: + rank_logger.warning( + "load_checkpoint: largest tensor loads %.2fGiB > concurrent_gb %.2fGiB; " + "it is issued as its own batch but still reads in one shot.", + max_leaf / (1 << 30), + concurrent_bytes / (1 << 30), + ) + else: + rank_logger.info( + "load_checkpoint: no restore read throttle (concurrent_gb=None); issuing " + "reads for all %d tensors (%.2fGiB) at once", + len(plan), + total_bytes / (1 << 30), + ) - unloaded_state = host_state.copy() + futures: dict[Any, tuple[str, str, int]] = {} - futures = {} - for _i, checkpoint_name in enumerate(tree_to_dict(metadata, keep_none=False).keys()): - name = checkpoint_name - if rename is not None: - name = rename(checkpoint_name) + node_lock: _NodeBatchLock | None = None + if _restore_node_serialize_enabled(): + node_lock = _NodeBatchLock(_node_lock_path()) + rank_logger.info( + "restore node-serialize ACTIVE (flock per batch) lock=%s pid=%d num_batches=%d", + node_lock._path, + os.getpid(), + num_batches, + ) - if host_state.get(name) is None: - if not has_subtree(name, host_state): - rank_logger.warning( - "Not loading %r from checkpoint because it's not in the initialized state", name + with contextlib.ExitStack() as stack: + if node_lock is not None: + stack.callback(node_lock.close) + tspec_transform = None + if kvstore_base is not None: + tspec_transform = _encrypted_tspec_transform(kvstore_base) + + for batch_index, batch in enumerate(batches, start=1): + with node_lock if node_lock is not None else contextlib.nullcontext(): + stores = [] + for checkpoint_name, name, mask, _nbytes in batch: + t = _open_tensor( + checkpoint_name, + name, + path, + use_zarr3, + ts_context, + host_state[name], + has_domain=domains.get(name) is not None, + tspec_transform=tspec_transform, + ) + for s, future in enumerate( + _read_into_shards(t, host_state[name], mask, domains.get(name)) + ): + futures[future] = (checkpoint_name, name, s) + stores.append(t) + + inflight_bytes = sum(nbytes for *_, nbytes in batch) + rank_logger.info( + "load_checkpoint: draining read batch %d (%.2fGiB in %d futures)", + batch_index, + inflight_bytes / (1 << 30), + len(futures), ) - continue + _drain_read_futures(futures, host_state, path, timeout, log_loaded=True) + del stores + _release_batch_memory() - if not load_mask: - mask = [True for _ in host_state[name].addressable_shards] - else: - mask = [shard.data.item() for shard in load_mask[name].addressable_shards] - if any(mask): - info = ocp.type_handlers.ParamInfo( - name=checkpoint_name, - path=path / checkpoint_name, - parent_dir=path, - is_ocdbt_checkpoint=True, - use_zarr3=use_zarr3, - ) - tspec = ocp.type_handlers.get_json_tspec_read(info, use_ocdbt=True) - t = ts.open(ts.Spec(tspec), open=True, context=ts_context).result() - if domains.get(name) is None and tuple(t.shape) != tuple(host_state[name].shape): - raise ValueError( - f"Tensor {name!r}: checkpoint has shape {tuple(t.shape)}, " - f"but initialized state has shape {tuple(host_state[name].shape)}. " - f"Use 'no_loading' to skip this tensor or 'domains' to load a partial slice." - ) - for s, future in enumerate( - _read_into_shards(t, host_state[name], mask, domains.get(name)) - ): - futures[future] = (checkpoint_name, name, s) + rank_logger.info("Loading checkpoint took %.2f sec", time.time() - start) - del unloaded_state[name] - for future in common._ready(list(futures), timeout=timeout): - try: - future.result() - except Exception as e: - logger.exception(e) - checkpoint_name, name, _ = futures.pop(future, None) - tensor = host_state[name] - err = f"Checkpoint error from loading {path}. Error loading from {name} into {tensor.shape=}, {tensor.dtype=}." - if checkpoint_name != name: - err = f"{err} (in checkpoint: {checkpoint_name})" - logger.error(err) - raise +_AUTO_WINDOW_TARGET_BATCHES = 16 +_AUTO_WINDOW_MAX_BYTES = 32 * 10**9 - checkpoint_name, name, _ = futures.pop(future, None) - if checkpoint_name != name: - rank_logger.debug("Loaded %s (name in checkpoint: %s)", name, checkpoint_name) - else: - rank_logger.debug("Loaded %s", name) - rank_logger.info("Loading checkpoint took %.2f sec", time.time() - start) +def _auto_window_bytes(plan: list[tuple[str, str, list[bool], int]]) -> int: + total = sum(nbytes for *_, nbytes in plan) + max_leaf = max((nbytes for *_, nbytes in plan), default=0) + return int(min(max(max_leaf, total / _AUTO_WINDOW_TARGET_BATCHES), _AUTO_WINDOW_MAX_BYTES)) + + +def _stage_to_host(arrays: dict[str, jax.Array]) -> dict[str, jax.Array]: + staged = {} + for name, array in arrays.items(): + is_cpu = all(d.platform == "cpu" for d in array.sharding.addressable_devices) + kind = "unpinned_host" if is_cpu else "pinned_host" + staged[name] = jax.device_put(array, array.sharding.with_memory_kind(kind)) + jax.block_until_ready(list(staged.values())) + return staged + + +def copy_aliased_arrays(tree): + seen: set[int] = set() + + def _copy(x): + if not isinstance(x, jax.Array): + return x + if id(x) in seen: + copied = jax.device_put(x, x.sharding, may_alias=False) + assert copied is not x + return copied + seen.add(id(x)) + return x + + return jax.tree.map(_copy, tree) + + +def load_checkpoint_streamed( + path: str, + device_state: dict[str, jax.Array], + load_mask: PyTree[jax.Array] | None, + rename: Callable[[str], str] | None = None, + domains: dict[str, Any] | None = None, + tag: str | None = None, + timeout: float = 900.0, + window_gb: float | None = None, + window_cap_gb: float | None = None, + on_replaced: Callable[[list[tuple[jax.Array, jax.Array]]], None] | None = None, + kms_client: Any | None = None, +) -> list[tuple[jax.Array, jax.Array]]: + if tag is None: + tag = "orbax-ckpt" + + load_mask = tree_to_dict(load_mask) + domains = _convert_domains(tree_to_dict(domains), device_state) + + src_path = path + start = time.time() + + path = pathlib.Path(path) / tag + kvstore_base, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + path, kms_client + ) + use_zarr3 = metadata_json["use_zarr3"] + + plan = _build_read_plan(checkpoint_names, device_state, load_mask, rename) + plan.sort(key=lambda item: item[3], reverse=True) + + auto_window = window_gb is None + window_bytes = _auto_window_bytes(plan) if auto_window else int(window_gb * 10**9) + total_bytes = sum(nbytes for *_, nbytes in plan) + max_leaf = max((nbytes for *_, nbytes in plan), default=0) + if window_cap_gb is not None: + window_bytes = min(window_bytes, max(max_leaf, int(window_cap_gb * 10**9))) + + batches = _pack_read_batches(plan, window_bytes) + num_batches = len(batches) + + rank_logger.info( + "Restoring %s (streamed): window=%.2fGiB%s total=%.2fGiB tensors=%d " + "max_leaf=%.2fGiB batches=%d", + src_path, + window_bytes / (1 << 30), + " (auto)" if auto_window else "", + total_bytes / (1 << 30), + len(plan), + max_leaf / (1 << 30), + num_batches, + ) + if num_batches <= 1 and len(plan) > 1: + rank_logger.warning( + "load_checkpoint_streamed: window %.2fGiB covers the full %.2fGiB state " + "(one batch); set restore_window_gb between %.2fGiB and the total to bound RSS", + window_bytes / (1 << 30), + total_bytes / (1 << 30), + max_leaf / (1 << 30), + ) + elif max_leaf > window_bytes: + rank_logger.warning( + "load_checkpoint_streamed: largest tensor %.2fGiB > window %.2fGiB; " + "that leaf still needs its full size in pinned memory", + max_leaf / (1 << 30), + window_bytes / (1 << 30), + ) + + replaced: list[tuple[jax.Array, jax.Array]] = [] + futures: dict[Any, tuple[str, str, int]] = {} + + node_lock: _NodeBatchLock | None = None + if _restore_node_serialize_enabled(): + node_lock = _NodeBatchLock(_node_lock_path()) + rank_logger.info( + "restore node-serialize ACTIVE (flock per batch) lock=%s pid=%d num_batches=%d", + node_lock._path, + os.getpid(), + num_batches, + ) + + with contextlib.ExitStack() as stack: + if node_lock is not None: + stack.callback(node_lock.close) + tspec_transform = None + if kvstore_base is not None: + tspec_transform = _encrypted_tspec_transform(kvstore_base) + + for batch in batches: + with node_lock if node_lock is not None else contextlib.nullcontext(): + staging = _stage_to_host({name: device_state[name] for _, name, _, _ in batch}) + + stores = [] + for checkpoint_name, name, mask, _nbytes in batch: + t = _open_tensor( + checkpoint_name, + name, + path, + use_zarr3, + ts_context, + staging[name], + has_domain=domains.get(name) is not None, + tspec_transform=tspec_transform, + ) + for s, future in enumerate( + _read_into_shards(t, staging[name], mask, domains.get(name)) + ): + futures[future] = (checkpoint_name, name, s) + stores.append(t) + + _drain_read_futures(futures, staging, path, timeout) + + new_arrays = { + name: jax.device_put(staging[name], device_state[name].sharding) + for _, name, _, _ in batch + } + jax.block_until_ready(list(new_arrays.values())) + batch_replaced = [] + for _, name, _, _ in batch: + batch_replaced.append((device_state[name], new_arrays[name])) + device_state[name] = new_arrays[name] + if on_replaced is not None: + on_replaced(batch_replaced) + else: + replaced.extend(batch_replaced) + + del staging, new_arrays, stores, batch_replaced + _release_batch_memory() + + rank_logger.info("Loading checkpoint (streamed) took %.2f sec", time.time() - start) + return replaced def broadcast_replicated( diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py new file mode 100644 index 00000000..be367773 --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import dataclasses +import json +import logging +import os +import pathlib +import threading + +import jax +import orbax.checkpoint as ocp +import tensorstore as ts +from orbax.checkpoint._src.handlers import pytree_checkpoint_handler as _pytree +from orbax.checkpoint._src.handlers.base_pytree_checkpoint_handler import ( + BasePyTreeCheckpointHandler, +) +from orbax.checkpoint._src.metadata import checkpoint as _step_metadata +from orbax.checkpoint._src.metadata import sharding as sharding_metadata +from orbax.checkpoint._src.multihost import multihost + +from xai_checkpointing import save as checkpointing_save +from xai_checkpointing.encrypted_kvstore import at_dir, encrypted_kvstore, use_encrypted_kvstore + +rank_logger = logging.getLogger("rank") + + +def encrypt_write(base, directory, name: str, data: bytes) -> None: + kv = ts.KvStore.open(at_dir(base, pathlib.Path(str(directory)))).result() + kv.write(name, data).result() + + +def decrypt_read(base, directory, name: str) -> bytes: + kv = ts.KvStore.open(at_dir(base, pathlib.Path(str(directory)))).result() + result = kv.read(name).result() + if result.state == "missing": + raise FileNotFoundError(f"{name} does not exist at {directory}") + return result.value + + +def _require_array_leaves(state) -> None: + bad = [ + f"{jax.tree_util.keystr(key_path)}: {type(leaf).__name__}" + for key_path, leaf in jax.tree_util.tree_leaves_with_path(state) + if not isinstance(leaf, jax.Array) + ] + if bad: + raise ValueError( + "encrypted orbax saves support only jax.Array leaves (add support in " + f"xai_checkpointing/orbax_encrypted.py if needed); offending leaves: {bad}" + ) + + +class EncryptedPyTreeCheckpointHandler(BasePyTreeCheckpointHandler): + def __init__(self, kms_client, encryption_chunk_size: int, **kwargs): + super().__init__(**kwargs) + self._kms_client = kms_client + self._encryption_chunk_size = encryption_chunk_size + self._base_lock = threading.Lock() + self._base_dir: str | None = None + self._base: dict | None = None + + def _base_for(self, directory) -> dict: + with self._base_lock: + key = str(directory) + if self._base_dir != key: + self._base_dir = key + self._base = encrypted_kvstore( + pathlib.Path(key), self._kms_client, self._encryption_chunk_size + ) + return self._base + + def set_kms_client(self, kms_client) -> None: + self._kms_client = kms_client + + def _write_metadata_file(self, directory, param_infos, save_args, use_zarr3=False): + def _save_fn(): + if multihost.is_primary_host(self._primary_host): + metadata = ocp._src.metadata.tree.InternalTreeMetadata.build( + param_infos, save_args=save_args, use_zarr3=use_zarr3 + ) + encrypt_write( + self._base_for(directory), + directory, + "_METADATA", + json.dumps(metadata.to_json()).encode(), + ) + return 0 + + return self._thread_pool.submit(_save_fn) + + def _read_metadata_file(self, directory): + return ocp._src.metadata.tree.InternalTreeMetadata.from_json( + json.loads(decrypt_read(self._base_for(directory), directory, "_METADATA")) + ) + + def finalize(self, directory): + path = pathlib.Path(str(directory)) + checkpointing_save.finalize_ts( + path, + world_size=jax.process_count(), + ts_context=checkpointing_save._get_ts_context(), + encrypted_base=self._base_for(directory), + ) + assert_all_enveloped(path) + rank_logger.info("Encrypted orbax finalize done (graft + sweep) at %s", path) + + +class EncryptedCheckpointMetadataStore: + def __init__(self, base_for): + self._base_for = base_for + + def is_blocking_writer(self) -> bool: + return True + + def write( + self, + checkpoint_path: str | os.PathLike, + checkpoint_metadata: _step_metadata.StepMetadata, + ) -> None: + directory = pathlib.Path(str(checkpoint_path)) + encrypt_write( + self._base_for(directory), + directory, + "_CHECKPOINT_METADATA", + json.dumps(dataclasses.asdict(checkpoint_metadata)).encode(), + ) + + def read(self, checkpoint_path: str | os.PathLike) -> _step_metadata.StepMetadata | None: + directory = pathlib.Path(str(checkpoint_path)) + kv = ts.KvStore.open(at_dir(self._base_for(directory), directory)).result() + result = kv.read("_CHECKPOINT_METADATA").result() + if result.state == "missing": + return None + return _step_metadata.StepMetadata.from_dict(json.loads(result.value.decode())) + + def update(self, checkpoint_path: str | os.PathLike, **kwargs) -> None: + self.write( + checkpoint_path, + dataclasses.replace( + self.read(checkpoint_path) or _step_metadata.StepMetadata(), **kwargs + ), + ) + + def wait_until_finished(self) -> None: + return None + + def close(self) -> None: + return None + + +def _encrypted_array_handler(array_handler_cls, base_for, array_handler_args=()): + class EncryptedArrayHandler(array_handler_cls): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._base_for = base_for + + def _get_json_tspec_write(self, info, *args, **kwargs): + spec = super()._get_json_tspec_write(info, *args, **kwargs) + spec["kvstore"] = use_encrypted_kvstore( + spec["kvstore"], self._base_for(info.parent_dir) + ) + return spec + + def _get_json_tspec_read(self, info, *args, **kwargs): + spec = super()._get_json_tspec_read(info, *args, **kwargs) + spec["kvstore"] = use_encrypted_kvstore( + spec["kvstore"], self._base_for(info.parent_dir) + ) + return spec + + async def _serialize_sharding(self, sharding, info, sharding_metadata_txn): + if info.parent_dir is None: + raise ValueError("parent_dir cannot be None") + tspec = ocp.type_handlers.get_sharding_tensorstore_spec( + info.parent_dir.as_posix(), info.name + ) + scoped = at_dir(self._base_for(info.parent_dir), info.parent_dir) + tspec["kvstore"] = {**scoped, "path": scoped.get("path", "") + "_sharding"} + if multihost.is_primary_host(self._primary_host): + t = await ts.open(tspec, open=True, context=info.ts_context) + serialized_sharding = None + sharding_metadata_value = sharding_metadata.from_jax_sharding(sharding) + if sharding_metadata_value is not None: + serialized_sharding = sharding_metadata_value.to_serialized_string() + if serialized_sharding is not None: + await t.with_transaction(sharding_metadata_txn).write(serialized_sharding) + + return EncryptedArrayHandler(*array_handler_args) + + +def encrypted_checkpointer( + kms_client, + encryption_chunk_size: int, + timeout_secs: int, + array_handler_cls, + array_handler_args=(), + *, + checkpoint_handler_kwargs, +): + impl_ref: list[EncryptedPyTreeCheckpointHandler] = [] + + def base_for(directory) -> dict: + return impl_ref[0]._base_for(directory) + + array_handler = _encrypted_array_handler(array_handler_cls, base_for, array_handler_args) + registry = _ArraysOnlyRegistry( + ocp.type_handlers.create_type_handler_registry((jax.Array, array_handler)) + ) + impl = EncryptedPyTreeCheckpointHandler( + kms_client, + encryption_chunk_size, + use_ocdbt=True, + use_zarr3=checkpoint_handler_kwargs.get("use_zarr3", False), + save_concurrent_bytes=_pytree._concurrent_bytes( + checkpoint_handler_kwargs.get("save_concurrent_gb") + ), + restore_concurrent_bytes=_pytree._concurrent_bytes( + checkpoint_handler_kwargs.get("restore_concurrent_gb") + ), + type_handler_registry=registry, + ) + impl_ref.append(impl) + return ocp.AsyncCheckpointer( + ocp.PyTreeCheckpointHandler( + handler_impl=impl, type_handler_registry=registry, **checkpoint_handler_kwargs + ), + timeout_secs, + checkpoint_metadata_store=EncryptedCheckpointMetadataStore(impl._base_for), + ) + + +class _ArraysOnlyRegistry: + def __init__(self, inner): + self._inner = inner + + def get(self, ty): + if not (isinstance(ty, type) and issubclass(ty, jax.Array)): + raise ValueError( + "encrypted orbax saves support only jax.Array leaves (add support in " + f"xai_checkpointing/orbax_encrypted.py if needed); offending type: {ty}" + ) + return self._inner.get(ty) + + def has(self, ty): + return self._inner.has(ty) + + def add(self, *args, **kwargs): + return self._inner.add(*args, **kwargs) + + +_ENCRYPTED_CHECKPOINTER = None +_ENCRYPTED_SAVE_CONCURRENT_GB: int | None = None +_ENCRYPTED_TIMEOUT_SECS: int | None = None +_ENCRYPTED_CHUNK_SIZE: int | None = None +_ENCRYPTED_ARRAY_HANDLER_CLS = None + + +def wait_until_finished() -> None: + if _ENCRYPTED_CHECKPOINTER is not None: + _ENCRYPTED_CHECKPOINTER.wait_until_finished() + + +def get_encrypted_checkpointer( + kms_client, + encryption_chunk_size: int, + timeout_secs: int, + save_concurrent_gb: int | None, + array_handler_cls, +): + global \ + _ENCRYPTED_CHECKPOINTER, \ + _ENCRYPTED_SAVE_CONCURRENT_GB, \ + _ENCRYPTED_TIMEOUT_SECS, \ + _ENCRYPTED_CHUNK_SIZE, \ + _ENCRYPTED_ARRAY_HANDLER_CLS + if _ENCRYPTED_CHECKPOINTER is None: + checkpoint_handler_kwargs: dict = {"use_zarr3": True} + if save_concurrent_gb is not None: + checkpoint_handler_kwargs["save_concurrent_gb"] = save_concurrent_gb + checkpoint_handler_kwargs["restore_concurrent_gb"] = save_concurrent_gb + _ENCRYPTED_SAVE_CONCURRENT_GB = save_concurrent_gb + _ENCRYPTED_TIMEOUT_SECS = timeout_secs + array_handler_args = ( + (int(save_concurrent_gb) * 10**9,) if save_concurrent_gb is not None else () + ) + _ENCRYPTED_CHUNK_SIZE = encryption_chunk_size + _ENCRYPTED_ARRAY_HANDLER_CLS = array_handler_cls + _ENCRYPTED_CHECKPOINTER = encrypted_checkpointer( + kms_client, + encryption_chunk_size, + timeout_secs, + array_handler_cls, + array_handler_args, + checkpoint_handler_kwargs=checkpoint_handler_kwargs, + ) + else: + _ENCRYPTED_CHECKPOINTER._handler._handler_impl.set_kms_client(kms_client) + if save_concurrent_gb is not None and _ENCRYPTED_SAVE_CONCURRENT_GB != save_concurrent_gb: + rank_logger.warning( + "get_encrypted_checkpointer(save_concurrent_gb=%s) ignored; checkpointer already " + "created with save_concurrent_gb=%s.", + save_concurrent_gb, + _ENCRYPTED_SAVE_CONCURRENT_GB, + ) + if _ENCRYPTED_TIMEOUT_SECS is not None and timeout_secs != _ENCRYPTED_TIMEOUT_SECS: + rank_logger.warning( + "get_encrypted_checkpointer(timeout_secs=%s) ignored; checkpointer already " + "created with timeout_secs=%s.", + timeout_secs, + _ENCRYPTED_TIMEOUT_SECS, + ) + if _ENCRYPTED_CHUNK_SIZE is not None and encryption_chunk_size != _ENCRYPTED_CHUNK_SIZE: + rank_logger.warning( + "get_encrypted_checkpointer(encryption_chunk_size=%s) ignored; checkpointer already " + "created with encryption_chunk_size=%s.", + encryption_chunk_size, + _ENCRYPTED_CHUNK_SIZE, + ) + if ( + _ENCRYPTED_ARRAY_HANDLER_CLS is not None + and array_handler_cls is not _ENCRYPTED_ARRAY_HANDLER_CLS + ): + rank_logger.warning( + "get_encrypted_checkpointer(array_handler_cls=%s) ignored; checkpointer already " + "created with array_handler_cls=%s.", + array_handler_cls, + _ENCRYPTED_ARRAY_HANDLER_CLS, + ) + return _ENCRYPTED_CHECKPOINTER + + +def assert_all_enveloped(path: pathlib.Path) -> None: + if not (path / "_DEK").is_file(): + raise RuntimeError(f"encrypted save left no master _DEK at {path}; not committing") + offenders = [] + for f in sorted(path.rglob("*")): + if not f.is_file() or f.name in ("_DEK", "_DEK.claim"): + continue + with f.open("rb") as fh: + if fh.read(8) != b"XAIENC01": + offenders.append(f) + if offenders: + raise RuntimeError( + f"encrypted save left plaintext at rest; not committing: {list(map(str, offenders))}" + ) + + +def extend_base(path: str, kms_client, encryption_chunk_size: int) -> dict: + return encrypted_kvstore(pathlib.Path(path) / "orbax-ckpt", kms_client, encryption_chunk_size) diff --git a/phoenix/reference/README.md b/phoenix/reference/README.md index b213185b..4891e1a1 100644 --- a/phoenix/reference/README.md +++ b/phoenix/reference/README.md @@ -51,7 +51,7 @@ comments or docstrings, so no in-file banner survives to say it there). | File | What it is | | --- | --- | -| `train_step.py` | The canonical single-device training step — the released composition (dense AdamW + sparse rowwise-AdaGrad) with the sharding infrastructure removed. See [`TRAINING.md`](../TRAINING.md). | +| `train_step.py` | The canonical single-device training step — the released composition (dense optimizer + sparse rowwise-AdaGrad; this reference composes the AdamW arm, while the ranking flagship/nano configs train the shipped Muon recipe) with the sharding infrastructure removed. See [`TRAINING.md`](../TRAINING.md). | | `train_synth.py` | The public launcher: points the shipped trainer config at a `dump_gen.py` dump and trains the nano model end to end. | | `repack_checkpoint.py` | Repacks a trained checkpoint into a publishable artifact: keeps load/infer tensors, drops optimizer state, scrubs internal metadata, regenerates checksums. | | `retrieve_then_rank.py` | The QUICKSTART §5 driver: sends real dump sessions through the two live servers — `RetrieveTopKCandidates` on retrieval, then `PredictNextActions` on ranking — over the production gRPC contract. | diff --git a/phoenix/xrex/configs/data_feeds.py b/phoenix/xrex/configs/data_feeds.py index d25c81d3..1eeb95af 100644 --- a/phoenix/xrex/configs/data_feeds.py +++ b/phoenix/xrex/configs/data_feeds.py @@ -56,6 +56,7 @@ def phoenix_kafka_kwargs() -> dict: GROUP_IDS: dict[str, str] = { "xrecsys_seqpack": "user_action_sequence_xrecsys", + "xrecsys_search": "user_action_sequence_xrecsys", "home_direct_packed": "user_action_sequence_home_direct_packed", "home_direct_packed_gb300": "user_action_sequence_home_direct_packed", "home_direct_packed_nano": "user_action_sequence_home_direct_packed_nano", @@ -112,6 +113,7 @@ def _ranking_aggregated_kafka(mparams, hash_table, use_post_sid, sid_num_levels, history_seq_len=mparams["history_seq_len"], candidate_seq_len=mparams["candidate_seq_len"], input_vocab_size=mparams["input_vocab_size"], + output_vocab_size=mparams["output_vocab_size"], num_continuous_actions=mparams["num_continuous_actions"], pad_token=PAD_TOKEN, num_negatives_per_example=mparams.get("num_negatives_per_example", 1), @@ -120,6 +122,8 @@ def _ranking_aggregated_kafka(mparams, hash_table, use_post_sid, sid_num_levels, use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, compute_post_unexplored_label=mparams.get("compute_post_unexplored_label", False), + enable_stale_post=mparams.get("enable_stale_post", False), + exclude_required_columns=mparams.get("exclude_required_columns", ""), ) @@ -144,6 +148,8 @@ def _ranking_rust_kafka(mparams, hash_table, use_post_sid, sid_num_levels, confi use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, compute_post_unexplored_label=mparams.get("compute_post_unexplored_label", False), + enable_stale_post=mparams.get("enable_stale_post", False), + exclude_required_columns=mparams.get("exclude_required_columns", ""), ) @@ -173,6 +179,7 @@ def _ranking_kafka_dispatcher(mparams, hash_table, use_post_sid, sid_num_levels, use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, compute_post_unexplored_label=mparams.get("compute_post_unexplored_label", False), + enable_stale_post=mparams.get("enable_stale_post", False), ) @@ -192,6 +199,7 @@ def _ranking_offline_kafka_dump(mparams, hash_table, use_post_sid, sid_num_level use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, compute_post_unexplored_label=mparams.get("compute_post_unexplored_label", False), + enable_stale_post=mparams.get("enable_stale_post", False), ) @@ -211,11 +219,16 @@ def _ranking_rust_parquet(mparams, hash_table, use_post_sid, sid_num_levels, con use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, compute_post_unexplored_label=mparams.get("compute_post_unexplored_label", False), + enable_stale_post=mparams.get("enable_stale_post", False), ) def _ranking_grpc_recsys(mparams, hash_table, use_post_sid, sid_num_levels, config_name): del config_name + if mparams.get("enable_stale_post", False): + raise ValueError( + "enable_stale_post is unsupported on grpc_recsys, which does not transport per-post feature arrays (int64/bool/categorical/float)" + ) return PhoenixGrpcDataset( hash_table=hash_table, history_seq_len=mparams["history_seq_len"], @@ -279,6 +292,7 @@ def _retrieval_aggregated_kafka( output_vocab_size=mparams["output_vocab_size"], use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, + exclude_required_columns=mparams.get("exclude_required_columns", ""), **extra_kwargs, ) @@ -313,6 +327,7 @@ def _retrieval_rust_kafka( output_vocab_size=mparams["output_vocab_size"], use_post_sid=use_post_sid, sid_num_levels=sid_num_levels, + exclude_required_columns=mparams.get("exclude_required_columns", ""), **extra_kwargs, ) diff --git a/phoenix/xrex/configs/xrecsys.py b/phoenix/xrex/configs/xrecsys.py index 787061dd..cbaac72e 100644 --- a/phoenix/xrex/configs/xrecsys.py +++ b/phoenix/xrex/configs/xrecsys.py @@ -29,13 +29,12 @@ ) from xrex.models.scaling import ScaleConfig from xrex.models.transformer import FeedForwardConfig, RematType, TransformerConfig -from xrex.optimizers.optim import OptimConfig from xrex.optimizers.recsys.config import RecsysEmbeddingOptimConfig +from xrex.optimizers.recsys.dense_optim import RecsysDenseOptimConfig from xrex.optimizers.recsys.rowwise_adagrad import RecsysRowwiseAdagradConfig from xrex.optimizers.schedule import ConstantSampleSchedule from xrex.train.parallel_config import ParallelConfig -from xrex.train.trainer import CheckpointConfig -from xrex.train.trainer_recsys import RecsysTrainer +from xrex.train.trainer_recsys import RecsysCheckpointConfig, RecsysTrainer PAD_TOKEN = 0 INPUT_VOCAB_K = 512 @@ -87,6 +86,7 @@ def _make_feature_prep_config(mparams: dict, scale_config: ScaleConfig) -> Featu sid_codebook_size=mparams.get("sid_codebook_size", 1024), sid_hash_level=mparams.get("sid_hash_level", False), sid_cross_attn=mparams.get("sid_cross_attn", False), + enable_stale_post=mparams.get("enable_stale_post", False), ) if "feature_prep" in mparams: @@ -142,7 +142,9 @@ def _make_cfg( ) input_vocab_size = _round_up_to_multiple(input_vocab_size, INPUT_VOCAB_K) - output_vocab_size = _round_up_to_multiple(ACTION_TYPE_MAP_LEN, OUTPUT_VOCAB_K) + output_vocab_size = cfg.get("output_vocab_size") or _round_up_to_multiple( + ACTION_TYPE_MAP_LEN, OUTPUT_VOCAB_K + ) num_continuous_actions = _round_up_to_multiple( len(continuous_action_type_map), CONTINUOUS_ACTION_TYPE_MAP_LEN ) @@ -193,6 +195,7 @@ def _make_dataset( _use_post_sid = mparams.get("use_post_sid", False) _sid_num_levels = mparams.get("sid_num_levels", 6) + _enable_stale_post = mparams.get("enable_stale_post", False) match dataset_type: case "aggregated_kafka": @@ -209,6 +212,7 @@ def _make_dataset( multimodal_embedding_type=mparams.get("multimodal_embedding_type"), use_post_sid=_use_post_sid, sid_num_levels=_sid_num_levels, + enable_stale_post=_enable_stale_post, ) case "toy_dataset": return PhoenixToyDataset( @@ -222,6 +226,7 @@ def _make_dataset( multimodal_embedding_type=mparams.get("multimodal_embedding_type"), use_post_sid=_use_post_sid, sid_num_levels=_sid_num_levels, + enable_stale_post=_enable_stale_post, ) case _: raise ValueError(f"Uknown {dataset_type=}, must be one of {DATASET_TYPES}") @@ -248,7 +253,6 @@ def _home_direct_packed_base() -> dict: "base_batch_size": 32, "dp": 1, "total_samples": 1e11, - "emb_learning_rate": 0.2, "log_q_correction": True, "continuous_metrics_mae_mean": True, "post_age_granularity_mins": 60, @@ -265,6 +269,7 @@ def _home_direct_packed_base() -> dict: "sid_codebook_size": 256, "sid_hash_level": True, "sid_cross_attn": False, + "enable_stale_post": True, "feature_prep_enabled": True, "feature_prep": FeaturePrepConfig( enable_post_sid=True, @@ -281,6 +286,9 @@ def _home_direct_packed_base() -> dict: enable_dwell_time=True, enable_time_of_day=False, enable_hour_of_day=True, + enable_is_author_followed_by_viewer=True, + enable_is_author_following_viewer=True, + enable_engagement_counts=True, hour_of_day_dither_fraction=0.1, ), "seqpack_distribution": BetaLengthDistribution( @@ -298,8 +306,25 @@ def _home_direct_packed_base() -> dict: "bs_per_device": 256, "ep": 256, "attn_impl": "pallas_ranker_varlen_attn", - "learning_rate": 1e-3, + "learning_rate": 7.1e-4, "checkpoint_every_n": 150, + "optim_config": RecsysDenseOptimConfig( + optim="muon", + muon_consistent_rms=0.2, + muon_matrix_weight_decay=0.014, + muon_split_fused="qkv:128", + adam_embedding_weight_decay=0.014, + b1=0.95, + b2=0.98, + ), + "emb_optim_config": RecsysEmbeddingOptimConfig( + rowwise_adagrad=RecsysRowwiseAdagradConfig( + learning_rate=0.28, + half_life_steps=2500, + lazy_decay=True, + weight_decay=2.8e-4, + ), + ), } _GB300_OVERRIDES = { @@ -310,6 +335,23 @@ def _home_direct_packed_base() -> dict: "unroll_layer_stack": True, "learning_rate": 5e-4, "checkpoint_every_n": 300, + "optim_config": RecsysDenseOptimConfig( + optim="muon", + muon_consistent_rms=0.2, + muon_matrix_weight_decay=0.01, + muon_split_fused="qkv:128", + adam_embedding_weight_decay=0.01, + b1=0.95, + b2=0.98, + ), + "emb_optim_config": RecsysEmbeddingOptimConfig( + rowwise_adagrad=RecsysRowwiseAdagradConfig( + learning_rate=0.28, + half_life_steps=5000, + lazy_decay=True, + weight_decay=1.4e-4, + ), + ), } @@ -374,7 +416,9 @@ def _home_direct_packed_base() -> dict: "home_direct_packed_nano": _make_cfg( { **_home_direct_packed_base(), - "learning_rate": 2e-3, + "learning_rate": _GB300_OVERRIDES["learning_rate"], + "optim_config": _GB300_OVERRIDES["optim_config"], + "emb_optim_config": _GB300_OVERRIDES["emb_optim_config"], "bs_per_device": 64, "ep": 1, "dp": 1, @@ -393,6 +437,43 @@ def _home_direct_packed_base() -> dict: author_vocab_size=30_000, ip_vocab_size=10_000, ), + "xrecsys_search": _make_cfg( + { + "history_seq_len": 1022, + "candidate_seq_len": 64, + "enable_user_country_feature": True, + "enable_user_language_feature": True, + "enable_user_location_feature": False, + "enable_user_gender_feature": False, + "enable_user_age_feature": False, + "num_layers": 8, + "emb_size": 2560, + "emb_table_width": 1024, + "query_heads": 20, + "kv_heads": 4, + "base_batch_size": 32, + "bs_per_device": 128, + "tp": 1, + "ep": 512, + "fsdp": 1, + "dp": 2, + "total_samples": 1e11, + "group_id": "user_action_sequence_xrecsys", + "learning_rate": 2e-3, + "attn_impl": "pallas_ranker_attn", + "log_q_correction": True, + "use_product_surface": True, + "post_age_granularity_mins": 60, + "output_vocab_size": 128, + "metric_group": "search", + "use_dense_action_table": True, + "condition_search_relevance_on_prompt": True, + }, + user_vocab_size=100_000_000, + item_vocab_size=100_000_000, + author_vocab_size=30_000_000, + ip_vocab_size=10_000_000, + ), } for _build_model_cfgs in config_registry.RANKING_MODEL_CFG_BUILDERS: @@ -448,9 +529,13 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | for config in configs: config_name, mparams = config["config_name__mparams"] + assert isinstance(mparams, dict) dataset_type = config["dataset_type"] config_name_gen = f"{config_name}_{dataset_type}" + if dataset_type == "grpc_recsys" and mparams.get("enable_stale_post", False): + mparams = {**mparams, "enable_stale_post": False} + hash_table = HashTable( hash_keys=HashKeys( user_id_table_size=mparams["user_vocab_size"], @@ -529,6 +614,7 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | reuse_run_id=False, evals=evals, model_config=RecsysAggregatedModelConfig( + use_dense_action_table=mparams.get("use_dense_action_table", False), multimodal_embedding_type=mparams.get("multimodal_embedding_type"), search_query_embedding_dim=mparams.get("search_query_embedding_dim", 0), use_ip_address=use_ip_address, @@ -542,6 +628,8 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | sid_codebook_size=mparams.get("sid_codebook_size", 1024), sid_hash_level=mparams.get("sid_hash_level", False), sid_cross_attn=mparams.get("sid_cross_attn", False), + sid_embedding_mode=mparams.get("sid_embedding_mode", "learned"), + sid_decoder_path=mparams.get("sid_decoder_path", ""), use_seqpack=mparams["use_seqpack"], right_anchored_rope=mparams.get("right_anchored_rope", False), user_features=user_features, @@ -554,7 +642,15 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | mask_candidate_positive_when_negative_action_present=mparams.get( "mask_candidate_positive_when_negative_action_present", False ), + train_view_through_heads=mparams.get("train_view_through_heads", False), + concat_history_bridge_prob=mparams.get("concat_history_bridge_prob", False), + mact_in_app_loss_weight=mparams.get("mact_in_app_loss_weight", 1.0), + split_head_training_by_source=mparams.get("split_head_training_by_source", False), + condition_search_relevance_on_prompt=mparams.get( + "condition_search_relevance_on_prompt", False + ), metric_group=mparams.get("metric_group", "default"), + metric_mask_keys=mparams.get("metric_mask_keys"), continuous_metrics_mae_mean=mparams.get("continuous_metrics_mae_mean", False), emb_table_width=mparams["emb_table_width"], history_seq_len=mparams["history_seq_len"], @@ -705,23 +801,29 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | ep=mparams["ep"], dp=mparams["dp"], ), - optim_config=OptimConfig( - optim="adam", - weight_decay=1e-3, - b1=0.95, - b2=0.98, + optim_config=mparams.get( + "optim_config", + RecsysDenseOptimConfig( + optim="adam", + weight_decay=1e-3, + b1=0.95, + b2=0.98, + ), ), lr_schedule_in_samples_config=ConstantSampleSchedule( learning_rate=mparams["learning_rate"], ), - emb_optim_config=RecsysEmbeddingOptimConfig( - rowwise_adagrad=RecsysRowwiseAdagradConfig( - learning_rate=mparams.get("emb_learning_rate", 0.1), + emb_optim_config=mparams.get( + "emb_optim_config", + RecsysEmbeddingOptimConfig( + rowwise_adagrad=RecsysRowwiseAdagradConfig( + learning_rate=mparams.get("emb_learning_rate", 0.1), + ), ), ), max_steps=int(mparams["total_samples"] / mparams["base_batch_size"]) - 100, max_samples=None, - checkpoint_config=CheckpointConfig( + checkpoint_config=RecsysCheckpointConfig( from_checkpoint=True, checkpoint_every_n=mparams.get("checkpoint_every_n", 100), shm_max_entries=3, diff --git a/phoenix/xrex/configs/xrecsys_gen_recs.py b/phoenix/xrex/configs/xrecsys_gen_recs.py index 96a5310f..e3cbafb4 100644 --- a/phoenix/xrex/configs/xrecsys_gen_recs.py +++ b/phoenix/xrex/configs/xrecsys_gen_recs.py @@ -15,11 +15,11 @@ from xrex.models.recsys_gen_recs_model import RecsysGenRecsModelConfig from xrex.models.scaling import ScaleConfig from xrex.models.transformer import FeedForwardConfig, RematType, TransformerConfig -from xrex.optimizers.optim import OptimConfig +from xrex.optimizers.recsys.dense_optim import RecsysDenseOptimConfig from xrex.optimizers.schedule import ConstantSampleSchedule from xrex.train.parallel_config import ParallelConfig -from xrex.train.trainer import CheckpointConfig from xrex.train.trainer_gen_recs import GenRecsTrainer +from xrex.train.trainer_recsys import RecsysCheckpointConfig PAD_TOKEN = 0 INPUT_VOCAB_K = 512 @@ -307,7 +307,7 @@ def _make_dataset( ep=mparams["ep"], dp=mparams["dp"], ), - optim_config=OptimConfig( + optim_config=RecsysDenseOptimConfig( optim="adam", weight_decay=1e-3, b1=0.95, @@ -318,7 +318,7 @@ def _make_dataset( ), max_steps=int(mparams["total_samples"] / mparams["base_batch_size"]) - 100, max_samples=None, - checkpoint_config=CheckpointConfig( + checkpoint_config=RecsysCheckpointConfig( from_checkpoint=True, checkpoint_every_n=500, checkpoint_keep_every_nth=100, diff --git a/phoenix/xrex/configs/xrecsys_sid_retrieval.py b/phoenix/xrex/configs/xrecsys_sid_retrieval.py index a56bfa9a..18e8c790 100644 --- a/phoenix/xrex/configs/xrecsys_sid_retrieval.py +++ b/phoenix/xrex/configs/xrecsys_sid_retrieval.py @@ -14,10 +14,10 @@ from xrex.models.recsys_sid_retrieval_model import RecsysSIDRetrievalConfig from xrex.models.scaling import ScaleConfig from xrex.models.transformer import FeedForwardConfig, RematType, TransformerConfig -from xrex.optimizers.optim import OptimConfig +from xrex.optimizers.recsys.dense_optim import RecsysDenseOptimConfig from xrex.optimizers.schedule import ConstantSampleSchedule from xrex.train.parallel_config import ParallelConfig -from xrex.train.trainer import CheckpointConfig +from xrex.train.trainer_recsys import RecsysCheckpointConfig from xrex.train.trainer_sid_retrieval import SidRetrievalTrainer PAD_TOKEN = 0 @@ -226,7 +226,7 @@ def make_trainer( ep=128, dp=1, ), - optim_config=OptimConfig( + optim_config=RecsysDenseOptimConfig( optim="adam", weight_decay=1e-3, b1=0.95, @@ -237,7 +237,7 @@ def make_trainer( ), max_steps=int(1e11 / 32) - 100, max_samples=None, - checkpoint_config=CheckpointConfig( + checkpoint_config=RecsysCheckpointConfig( from_checkpoint=True, checkpoint_every_n=100, checkpoint_keep_every_nth=1000, diff --git a/phoenix/xrex/configs/xrecsys_two_tower.py b/phoenix/xrex/configs/xrecsys_two_tower.py index 135bf80d..6fd9cb37 100644 --- a/phoenix/xrex/configs/xrecsys_two_tower.py +++ b/phoenix/xrex/configs/xrecsys_two_tower.py @@ -31,13 +31,12 @@ ) from xrex.models.scaling import ScaleConfig from xrex.models.transformer import FeedForwardConfig, RematType, TransformerConfig -from xrex.optimizers.optim import OptimConfig from xrex.optimizers.recsys.config import RecsysEmbeddingOptimConfig +from xrex.optimizers.recsys.dense_optim import RecsysDenseOptimConfig from xrex.optimizers.recsys.rowwise_adagrad import RecsysRowwiseAdagradConfig from xrex.optimizers.schedule import ConstantSampleSchedule from xrex.train.parallel_config import ParallelConfig -from xrex.train.trainer import CheckpointConfig -from xrex.train.trainer_recsys import RecsysTrainer +from xrex.train.trainer_recsys import RecsysCheckpointConfig, RecsysTrainer PAD_TOKEN = 0 INPUT_VOCAB_K = 512 @@ -271,8 +270,8 @@ def _xrecsys_two_tower_combined_base() -> dict: "empty_history_user_dropout_rate": 0.1, "learning_rate": 2e-3, "emb_learning_rate": 0.1, - "qk_norm": False, - "attn_logit_cap": 80.0, + "qk_norm": True, + "attn_logit_cap": -1, "primer_norm": True, "feature_prep_enabled": True, "enable_candidate_tower_linear_proj": False, @@ -334,13 +333,7 @@ def _xrecsys_two_tower_combined_base() -> dict: "attn_impl": "pallas_ranker_varlen_attn", } -_GB300_OVERRIDES = { - "bs_per_device": 768, - "ep": 32, - "attn_impl": "cutedsl_ranker_varlen_attn", - "remat_policy": RematType.SAVE_GB300_RECSYS, - "unroll_layer_stack": True, -} +_GB300_OVERRIDES = {"bs_per_device": 960, "ep": 64, "attn_impl": "cutedsl_ranker_varlen_attn"} MODEL_CFGS = { @@ -427,6 +420,9 @@ def _xrecsys_two_tower_combined_base() -> dict: "dp": 1, "total_samples": 1e11, "learning_rate": 2e-3, + "qk_norm": True, + "attn_logit_cap": -1, + "right_anchored_rope": True, "attn_impl": "pallas_ranker_attn", "enable_candidate_tower_linear_proj": True, "apply_u2u_and_i2i_loss": False, @@ -589,6 +585,7 @@ def _xrecsys_two_tower_combined_base() -> dict: num_global_negatives_per_example=mparams["num_global_negatives_per_example"], debug_mode=False, apply_u2u_and_i2i_loss=mparams.get("apply_u2u_and_i2i_loss", False), + use_history_segment_ids=mparams.get("use_history_segment_ids", False), positive_actions=positive_actions, hard_negative_actions=hard_negative_actions, soft_negative_actions=soft_negative_actions, @@ -717,7 +714,7 @@ def _xrecsys_two_tower_combined_base() -> dict: ep=mparams["ep"], dp=mparams["dp"], ), - optim_config=OptimConfig( + optim_config=RecsysDenseOptimConfig( optim="adam", weight_decay=1e-3, b1=0.95, @@ -733,7 +730,7 @@ def _xrecsys_two_tower_combined_base() -> dict: ), max_steps=int(mparams["total_samples"] / mparams["base_batch_size"]) - 100, max_samples=None, - checkpoint_config=CheckpointConfig( + checkpoint_config=RecsysCheckpointConfig( from_checkpoint=True, checkpoint_every_n=300, checkpoint_keep_every_nth=100, diff --git a/phoenix/xrex/cuda/async_emb/async_emb.py b/phoenix/xrex/cuda/async_emb/async_emb.py index 2fecd7d6..c2fc0f95 100644 --- a/phoenix/xrex/cuda/async_emb/async_emb.py +++ b/phoenix/xrex/cuda/async_emb/async_emb.py @@ -42,6 +42,14 @@ }, platform="CUDA", ) + jax.ffi.register_ffi_target( + "xrex_async_emb_rowwise_adagrad_lazy_update_start", + fn={ + "initialize": async_emb_api.rowwise_adagrad_lazy_update_start_init(), + "execute": async_emb_api.rowwise_adagrad_lazy_update_start(), + }, + platform="CUDA", + ) jax.ffi.register_ffi_target( "xrex_async_emb_rowwise_adagrad_update_done", fn=async_emb_api.rowwise_adagrad_update_done(), @@ -186,6 +194,7 @@ def rowwise_adagrad_update_start( learning_rate: float, eps: float, decay_factor: float, + weight_decay_factor: float = 1.0, ): assert grads.shape == (ctx.tokens_per_rank, ctx.emb_width) and grads.dtype == jnp.bfloat16 assert math.prod(segment_ids.shape) == ctx.tokens_per_rank @@ -215,6 +224,61 @@ def rowwise_adagrad_update_start( learning_rate=float(learning_rate), eps=float(eps), decay_factor=float(decay_factor), + weight_decay_factor=float(weight_decay_factor), + ) + return tuple(outs) + + +def rowwise_adagrad_lazy_update_start( + grads: jax.Array, + segment_ids: jax.Array, + unique_tokens: jax.Array, + table: jax.Array, + row_state: jax.Array, + last_step: jax.Array, + step: jax.Array, + pending: jax.Array, + gate: jax.Array, + ctx: AsyncEmbContextHandle, + *, + learning_rate: float, + eps: float, + accum_decay_rate: float, + weight_decay_rate: float, +): + assert grads.shape == (ctx.tokens_per_rank, ctx.emb_width) and grads.dtype == jnp.bfloat16 + assert math.prod(segment_ids.shape) == ctx.tokens_per_rank + assert math.prod(unique_tokens.shape) == ctx.num_unique + assert last_step.shape == row_state.shape and last_step.dtype == jnp.int32 + flat_segment_ids = segment_ids.reshape(-1).astype(jnp.int32) + flat_unique_tokens = unique_tokens.reshape(-1).astype(jnp.int32) + with jax.named_scope("async_emb.rowwise_adagrad_lazy_update_start"): + outs = jax.ffi.ffi_call( + "xrex_async_emb_rowwise_adagrad_lazy_update_start", + [ + jax.ShapeDtypeStruct(grads.shape, grads.dtype), + jax.ShapeDtypeStruct(flat_segment_ids.shape, jnp.int32), + jax.ShapeDtypeStruct(flat_unique_tokens.shape, jnp.int32), + jax.ShapeDtypeStruct(table.shape, table.dtype), + jax.ShapeDtypeStruct(row_state.shape, row_state.dtype), + jax.ShapeDtypeStruct(last_step.shape, jnp.int32), + ], + input_output_aliases={0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5}, + )( + grads, + flat_segment_ids, + flat_unique_tokens, + table, + row_state, + last_step, + step.astype(jnp.int32).reshape(1), + pending.astype(jnp.int32).reshape(1), + gate, + **ctx.attrs(), + learning_rate=float(learning_rate), + eps=float(eps), + accum_decay_rate=float(accum_decay_rate), + weight_decay_rate=float(weight_decay_rate), ) return tuple(outs) diff --git a/phoenix/xrex/cuda/async_emb/src/async_emb_api.cc b/phoenix/xrex/cuda/async_emb/src/async_emb_api.cc index a0c0acfd..7589f621 100644 --- a/phoenix/xrex/cuda/async_emb/src/async_emb_api.cc +++ b/phoenix/xrex/cuda/async_emb/src/async_emb_api.cc @@ -215,14 +215,18 @@ ffi::Error LookupDone( }); } -ApplyUpdateRule makeRowwiseAdagradRule(AdagradParams params, float* row_state) { +ApplyUpdateRule makeRowwiseAdagradRule( + AdagradParams params, float* row_state, int32_t* last_step = nullptr +) { return [params, - row_state](const ReducedGradients& reduced, const UpdateJob& job, cudaStream_t stream) { + row_state, + last_step](const ReducedGradients& reduced, const UpdateJob& job, cudaStream_t stream) { launch_rowwise_adagrad_apply( reduced.row_grads, reduced.row_sq_sums, job.unique_tokens, row_state, + last_step, job.table, job.vocab_rows, reduced.scalars, @@ -263,12 +267,57 @@ ffi::Error RowwiseAdagradUpdateStartInit( return invalid("rowwise_adagrad_update_start: context not initialized"); } std::call_once(rowwise_adagrad_warmup_once, [&] { - ctx->warmupUpdateRule(makeRowwiseAdagradRule(AdagradParams{0.f, 1.f, 1.f, 1.f}, nullptr)); + ctx->warmupUpdateRule( + makeRowwiseAdagradRule(AdagradParams{0.f, 1.f, 1.f, 1.f, 1.f, 0.f, 0.f}, nullptr) + ); }); return ffi::Error::Success(); }); } +ffi::Error validateRowwiseAdagradBuffers( + AsyncEmbContext* ctx, + const ffi::AnyBuffer& grads, + const ffi::AnyBuffer& segment_ids, + const ffi::AnyBuffer& unique_tokens, + const ffi::AnyBuffer& table, + const ffi::AnyBuffer& row_state, + const ffi::AnyBuffer& pending +) { + const auto& spec = ctx->spec(); + if (32 % spec.shard_width != 0) { + return invalid("rowwise_adagrad_update_start: shard_width must divide 32 (warp broadcast)"); + } + if (grads.element_type() != ffi::DataType::BF16 || + grads.element_count() != size_t(spec.tokens_per_rank * spec.shard_width * ctx->worldSize())) { + return invalid("rowwise_adagrad_update_start: bad grads"); + } + if (segment_ids.element_type() != ffi::DataType::S32 || + segment_ids.element_count() != size_t(spec.tokens_per_rank)) { + return invalid("rowwise_adagrad_update_start: bad indices"); + } + if (unique_tokens.element_type() != ffi::DataType::S32 || + unique_tokens.element_count() != size_t(spec.num_unique)) { + return invalid("rowwise_adagrad_update_start: bad unique tokens"); + } + if (table.element_type() != ffi::DataType::BF16 || table.dimensions().size() != 2 || + table.dimensions()[1] != spec.shard_width) { + return invalid("rowwise_adagrad_update_start: bad table"); + } + if (row_state.element_type() != ffi::DataType::F32 || + row_state.element_count() != size_t(table.dimensions()[0])) { + return invalid("rowwise_adagrad_update_start: bad row state"); + } + if (pending.element_type() != ffi::DataType::S32 || pending.element_count() != 1) { + return invalid("rowwise_adagrad_update_start: pending must be s32[1]"); + } + return ffi::Error::Success(); +} + +bool bufferAliased(const ffi::AnyBuffer& in, ffi::Result& out) { + return out->untyped_data() == in.untyped_data(); +} + ffi::Error RowwiseAdagradUpdateStart( cudaStream_t stream, ffi::CollectiveParamsPartial, @@ -298,55 +347,160 @@ ffi::Error RowwiseAdagradUpdateStart( auto learning_rate = attrs.get("learning_rate"); auto eps = attrs.get("eps"); auto decay_factor = attrs.get("decay_factor"); - if (!learning_rate.has_value() || !eps.has_value() || !decay_factor.has_value()) { + auto weight_decay_factor = attrs.get("weight_decay_factor"); + if (!learning_rate.has_value() || !eps.has_value() || !decay_factor.has_value() || + !weight_decay_factor.has_value()) { return invalid("rowwise_adagrad_update_start: missing optimizer attrs"); } const float lr_f = float(*learning_rate); const float eps_f = float(*eps); const float decay_f = float(*decay_factor); + const float wd_f = float(*weight_decay_factor); if (!std::isfinite(lr_f) || !std::isfinite(eps_f) || !std::isfinite(decay_f) || eps_f <= 0.f || - decay_f < 0.f || decay_f > 1.f) { + decay_f < 0.f || decay_f > 1.f || !std::isfinite(wd_f) || wd_f <= 0.f || wd_f > 1.f) { return invalid("rowwise_adagrad_update_start: invalid optimizer parameters"); } + ffi::Error buffers_error = validateRowwiseAdagradBuffers( + ctx.get(), grads, segment_ids, unique_tokens, table, row_state, pending + ); + if (buffers_error.failure()) { + return buffers_error; + } + if (!bufferAliased(grads, grads_pin) || !bufferAliased(segment_ids, segment_ids_pin) || + !bufferAliased(unique_tokens, unique_tokens_pin) || !bufferAliased(table, table_out) || + !bufferAliased(row_state, state_out)) { + return invalid( + "rowwise_adagrad_update_start: outputs must alias inputs (check input_output_aliases)" + ); + } const auto& spec = ctx->spec(); - if (32 % spec.shard_width != 0) { - return invalid("rowwise_adagrad_update_start: shard_width must divide 32 (warp broadcast)"); + const AdagradParams adagrad{lr_f, eps_f, decay_f, 1.f / float(spec.emb_width), wd_f, 0.f, 0.f}; + UpdateJob job{ + static_cast(grads.untyped_data()), + static_cast(segment_ids.untyped_data()), + static_cast(unique_tokens.untyped_data()), + static_cast<__nv_bfloat16*>(table.untyped_data()), + table.dimensions()[0] + }; + ctx->armUpdate( + job, + static_cast(pending.untyped_data()), + makeRowwiseAdagradRule(adagrad, static_cast(row_state.untyped_data())), + stream + ); + return ffi::Error::Success(); + }); +} + +ffi::Error RowwiseAdagradLazyUpdateStartInit( + cudaStream_t, + ffi::CollectiveParamsPartial params, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::AnyBuffer, + ffi::Dictionary attrs, + ffi::Result, + ffi::Result, + ffi::Result, + ffi::Result, + ffi::Result, + ffi::Result +) { + return ffiGuard([&]() -> ffi::Error { + ffi::Error error = initializeContext(attrs, params); + if (error.failure()) { + return error; + } + auto ctx = readyContext(*attrs.get("context_id")); + if (ctx == nullptr) { + return invalid("rowwise_adagrad_lazy_update_start: context not initialized"); + } + std::call_once(rowwise_adagrad_warmup_once, [&] { + ctx->warmupUpdateRule( + makeRowwiseAdagradRule(AdagradParams{0.f, 1.f, 1.f, 1.f, 1.f, 0.f, 0.f}, nullptr) + ); + }); + return ffi::Error::Success(); + }); +} + +ffi::Error RowwiseAdagradLazyUpdateStart( + cudaStream_t stream, + ffi::CollectiveParamsPartial, + ffi::AnyBuffer grads, + ffi::AnyBuffer segment_ids, + ffi::AnyBuffer unique_tokens, + ffi::AnyBuffer table, + ffi::AnyBuffer row_state, + ffi::AnyBuffer last_step, + ffi::AnyBuffer logical_step, + ffi::AnyBuffer pending, + ffi::AnyBuffer , + ffi::Dictionary attrs, + ffi::Result grads_pin, + ffi::Result segment_ids_pin, + ffi::Result unique_tokens_pin, + ffi::Result table_out, + ffi::Result state_out, + ffi::Result last_step_out +) { + return ffiGuard([&]() -> ffi::Error { + auto context_id = attrs.get("context_id"); + if (!context_id.has_value()) { + return invalid("rowwise_adagrad_lazy_update_start: missing context_id"); } - if (grads.element_type() != ffi::DataType::BF16 || - grads.element_count() != - size_t(spec.tokens_per_rank * spec.shard_width * ctx->worldSize())) { - return invalid("rowwise_adagrad_update_start: bad grads"); + auto ctx = readyContext(*context_id); + if (ctx == nullptr) { + return invalid("rowwise_adagrad_lazy_update_start: context not initialized"); } - if (segment_ids.element_type() != ffi::DataType::S32 || - segment_ids.element_count() != size_t(spec.tokens_per_rank)) { - return invalid("rowwise_adagrad_update_start: bad indices"); + auto learning_rate = attrs.get("learning_rate"); + auto eps = attrs.get("eps"); + auto accum_decay_rate = attrs.get("accum_decay_rate"); + auto weight_decay_rate = attrs.get("weight_decay_rate"); + if (!learning_rate.has_value() || !eps.has_value() || !accum_decay_rate.has_value() || + !weight_decay_rate.has_value()) { + return invalid("rowwise_adagrad_lazy_update_start: missing optimizer attrs"); } - if (unique_tokens.element_type() != ffi::DataType::S32 || - unique_tokens.element_count() != size_t(spec.num_unique)) { - return invalid("rowwise_adagrad_update_start: bad unique tokens"); + const float lr_f = float(*learning_rate); + const float eps_f = float(*eps); + const float accum_rate_f = float(*accum_decay_rate); + const float wd_rate_f = float(*weight_decay_rate); + if (!std::isfinite(lr_f) || !std::isfinite(eps_f) || eps_f <= 0.f || + !std::isfinite(accum_rate_f) || accum_rate_f < 0.f || !std::isfinite(wd_rate_f) || + wd_rate_f < 0.f) { + return invalid("rowwise_adagrad_lazy_update_start: invalid optimizer parameters"); } - if (table.element_type() != ffi::DataType::BF16 || table.dimensions().size() != 2 || - table.dimensions()[1] != spec.shard_width) { - return invalid("rowwise_adagrad_update_start: bad table"); + ffi::Error buffers_error = validateRowwiseAdagradBuffers( + ctx.get(), grads, segment_ids, unique_tokens, table, row_state, pending + ); + if (buffers_error.failure()) { + return buffers_error; } - if (row_state.element_type() != ffi::DataType::F32 || - row_state.element_count() != size_t(table.dimensions()[0])) { - return invalid("rowwise_adagrad_update_start: bad row state"); + if (last_step.element_type() != ffi::DataType::S32 || + last_step.element_count() != size_t(table.dimensions()[0])) { + return invalid("rowwise_adagrad_lazy_update_start: bad last_step"); } - if (pending.element_type() != ffi::DataType::S32 || pending.element_count() != 1) { - return invalid("rowwise_adagrad_update_start: pending must be s32[1]"); + if (logical_step.element_type() != ffi::DataType::S32 || logical_step.element_count() != 1) { + return invalid("rowwise_adagrad_lazy_update_start: step must be s32[1]"); } - auto aliased = [](const ffi::AnyBuffer& in, ffi::Result& out) -> bool { - return out->untyped_data() == in.untyped_data(); - }; - if (!aliased(grads, grads_pin) || !aliased(segment_ids, segment_ids_pin) || - !aliased(unique_tokens, unique_tokens_pin) || !aliased(table, table_out) || - !aliased(row_state, state_out)) { + if (!bufferAliased(grads, grads_pin) || !bufferAliased(segment_ids, segment_ids_pin) || + !bufferAliased(unique_tokens, unique_tokens_pin) || !bufferAliased(table, table_out) || + !bufferAliased(row_state, state_out) || !bufferAliased(last_step, last_step_out)) { return invalid( - "rowwise_adagrad_update_start: outputs must alias inputs (check input_output_aliases)" + "rowwise_adagrad_lazy_update_start: outputs must alias inputs (check " + "input_output_aliases)" ); } - const AdagradParams adagrad{lr_f, eps_f, decay_f, 1.f / float(spec.emb_width)}; + const auto& spec = ctx->spec(); + const AdagradParams adagrad{ + lr_f, eps_f, 1.f, 1.f / float(spec.emb_width), 1.f, accum_rate_f, wd_rate_f + }; UpdateJob job{ static_cast(grads.untyped_data()), static_cast(segment_ids.untyped_data()), @@ -357,8 +511,13 @@ ffi::Error RowwiseAdagradUpdateStart( ctx->armUpdate( job, static_cast(pending.untyped_data()), - makeRowwiseAdagradRule(adagrad, static_cast(row_state.untyped_data())), - stream + makeRowwiseAdagradRule( + adagrad, + static_cast(row_state.untyped_data()), + static_cast(last_step.untyped_data()) + ), + stream, + static_cast(logical_step.untyped_data()) ); return ffi::Error::Success(); }); @@ -484,6 +643,54 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL( .Ret() ); +XLA_FFI_DEFINE_HANDLER_SYMBOL( + kRowwiseAdagradLazyUpdateStartInit, + RowwiseAdagradLazyUpdateStartInit, + ffi::Ffi::Bind() + .Ctx>() + .Ctx() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Attrs() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() +); + +XLA_FFI_DEFINE_HANDLER_SYMBOL( + kRowwiseAdagradLazyUpdateStart, + RowwiseAdagradLazyUpdateStart, + ffi::Ffi::Bind() + .Ctx>() + .Ctx() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Attrs() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() +); + XLA_FFI_DEFINE_HANDLER_SYMBOL( kRowwiseAdagradUpdateDone, RowwiseAdagradUpdateDone, @@ -575,6 +782,12 @@ NB_MODULE(async_emb_api, m) { return encapsulate(kRowwiseAdagradUpdateStartInit); }); m.def("rowwise_adagrad_update_start", [] { return encapsulate(kRowwiseAdagradUpdateStart); }); + m.def("rowwise_adagrad_lazy_update_start_init", [] { + return encapsulate(kRowwiseAdagradLazyUpdateStartInit); + }); + m.def("rowwise_adagrad_lazy_update_start", [] { + return encapsulate(kRowwiseAdagradLazyUpdateStart); + }); m.def("rowwise_adagrad_update_done", [] { return encapsulate(kRowwiseAdagradUpdateDone); }); m.def( "wait_step_ready", diff --git a/phoenix/xrex/cuda/async_emb/src/async_emb_comm.cc b/phoenix/xrex/cuda/async_emb/src/async_emb_comm.cc index d8b455c1..ea64d7fe 100644 --- a/phoenix/xrex/cuda/async_emb/src/async_emb_comm.cc +++ b/phoenix/xrex/cuda/async_emb/src/async_emb_comm.cc @@ -705,7 +705,11 @@ void AsyncEmbContext::armLookup(LookupJob job, cudaStream_t main_stream) { } void AsyncEmbContext::armUpdate( - UpdateJob job, const int32_t* pending, ApplyUpdateRule apply, cudaStream_t main_stream + UpdateJob job, + const int32_t* pending, + ApplyUpdateRule apply, + cudaStream_t main_stream, + const int32_t* logical_step ) { std::lock_guard arm_lock(arm_mu_); ensureHealthy(); @@ -727,6 +731,15 @@ void AsyncEmbContext::armUpdate( cudaMemcpyDeviceToDevice, main_stream )); + if (logical_step != nullptr) { + XAI_CUDA_CHECK(cudaMemcpyAsync( + arena() + layout_.scalars + offsetof(UpdateScalars, step), + logical_step, + sizeof(int32_t), + cudaMemcpyDeviceToDevice, + main_stream + )); + } recordInputReady(update_, main_stream); enqueueUpdate(job, apply); { diff --git a/phoenix/xrex/cuda/async_emb/src/async_emb_comm.hpp b/phoenix/xrex/cuda/async_emb/src/async_emb_comm.hpp index 9bad8033..67c127ba 100644 --- a/phoenix/xrex/cuda/async_emb/src/async_emb_comm.hpp +++ b/phoenix/xrex/cuda/async_emb/src/async_emb_comm.hpp @@ -131,7 +131,11 @@ class AsyncEmbContext : public common::CommContext { void armLookup(LookupJob job, cudaStream_t main_stream); void armUpdate( - UpdateJob job, const int32_t* pending, ApplyUpdateRule apply, cudaStream_t main_stream + UpdateJob job, + const int32_t* pending, + ApplyUpdateRule apply, + cudaStream_t main_stream, + const int32_t* logical_step = nullptr ); void warmupUpdateRule(const ApplyUpdateRule& apply); void resetTableBinding(); diff --git a/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.cu b/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.cu index 8b4c913d..4b48a38f 100644 --- a/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.cu +++ b/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.cu @@ -302,6 +302,7 @@ __global__ void rowwise_adagrad_apply( const float* row_sq_sums, const int32_t* unique_tokens, float* row_state, + int32_t* last_step, __nv_bfloat16* table_shard, UpdateScalars* scalars, AdagradParams adagrad, @@ -311,6 +312,8 @@ __global__ void rowwise_adagrad_apply( ) { const float norm = sqrtf(__ldg(row_sq_sums + num_unique)); const bool apply = isfinite(norm) && __ldg(&scalars->pending) != 0; + const bool lazy = last_step != nullptr; + const int32_t clock = lazy ? __ldg(&scalars->step) : 0; const int64_t thread_id = blockIdx.x * int64_t(blockDim.x) + threadIdx.x; if (thread_id == 0) { scalars->norm = norm; @@ -330,16 +333,28 @@ __global__ void rowwise_adagrad_apply( } const bool ok = in_range && apply && token >= 0 && token < vocab_rows; float accum = 0.f; + float wd = 1.f; if (ok && col == 0) { - accum = row_state[token] * adagrad.decay + __ldg(row_sq_sums + row) * adagrad.inv_emb_width; + float accum_factor = adagrad.decay; + if (lazy) { + const float delta = fmaxf(0.f, float(clock - last_step[token])); + accum_factor = expf(-adagrad.accum_decay_rate * delta); + wd = __ldg(row_sq_sums + row) > 0.f ? expf(-adagrad.weight_decay_rate * delta) : 1.f; + last_step[token] = clock; + } else { + wd = __ldg(row_sq_sums + row) > 0.f ? adagrad.weight_decay_factor : 1.f; + } + accum = row_state[token] * accum_factor + __ldg(row_sq_sums + row) * adagrad.inv_emb_width; } accum = __shfl_sync(0xffffffffu, accum, lane - col); + wd = __shfl_sync(0xffffffffu, wd, lane - col); if (ok) { const float step = (-adagrad.lr) * (1.f / (sqrtf(accum) + adagrad.eps)); const float grad = __bfloat162float(__float2bfloat16(grad_accum[elem])); const float delta = __bfloat162float(__float2bfloat16(step * grad)); __nv_bfloat16* cell = table_shard + token * shard_width + col; - *cell = __float2bfloat16(__bfloat162float(*cell) + delta); + const float decayed = __bfloat162float(__float2bfloat16(__bfloat162float(*cell) * wd)); + *cell = __float2bfloat16(decayed + delta); if (col == 0) { row_state[token] = accum; } @@ -355,6 +370,7 @@ void launch_rowwise_adagrad_apply( const float* row_sq_sums, const int32_t* unique_tokens, float* row_state, + int32_t* last_step, __nv_bfloat16* table_shard, int64_t vocab_rows, UpdateScalars* scalars, @@ -368,6 +384,7 @@ void launch_rowwise_adagrad_apply( row_sq_sums, unique_tokens, row_state, + last_step, table_shard, scalars, adagrad, diff --git a/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.hpp b/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.hpp index a9ae6808..d6247521 100644 --- a/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.hpp +++ b/phoenix/xrex/cuda/async_emb/src/async_emb_kernel.hpp @@ -18,15 +18,19 @@ struct AdagradParams { float eps; float decay; float inv_emb_width; + float weight_decay_factor; + float accum_decay_rate; + float weight_decay_rate; }; struct UpdateScalars { float norm; int32_t valid; int32_t pending; + int32_t step; }; -static_assert(sizeof(UpdateScalars) == 3 * sizeof(int32_t)); +static_assert(sizeof(UpdateScalars) == 4 * sizeof(int32_t)); void launch_lookup_dispatch( const int32_t* token_ids_all, @@ -66,6 +70,7 @@ void launch_rowwise_adagrad_apply( const float* row_sq_sums, const int32_t* unique_tokens, float* row_state, + int32_t* last_step, __nv_bfloat16* table_shard, int64_t vocab_rows, UpdateScalars* scalars, diff --git a/phoenix/xrex/cutedsl/ranker_attention_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_fa4.py index 5240f25f..47d0fe68 100644 --- a/phoenix/xrex/cutedsl/ranker_attention_fa4.py +++ b/phoenix/xrex/cutedsl/ranker_attention_fa4.py @@ -6,74 +6,99 @@ import jax import jax.numpy as jnp -from jax.ad_checkpoint import checkpoint_name _FA4_KERNEL_CACHE = {} -_DENSE_BS_LAYOUT_CACHE = {} - -def build_dense_block_sparse_layout(batch_size, seq_len, hist_len, num_q_heads): - key = (batch_size, seq_len, hist_len, num_q_heads) - if key not in _DENSE_BS_LAYOUT_CACHE: - import numpy as np - - cand_len = seq_len - hist_len - assert hist_len % 128 == 0 and cand_len % 128 == 0, ( - f"dense block-sparse needs tile-aligned lengths (hist={hist_len}, cand={cand_len})" - ) - h_blocks = hist_len // 128 - num_blocks = seq_len // 128 - blocks = np.arange(num_blocks, dtype=np.int32) - is_cand = blocks >= h_blocks - - fwd_mask_cnt = np.zeros(num_blocks, np.int32) - fwd_mask_idx = np.zeros((num_blocks, 1), np.int32) - fwd_full_cnt = np.full(num_blocks, h_blocks, np.int32) - fwd_full_idx = np.tile(blocks[:h_blocks], (num_blocks, 1)) - fwd_diag_cnt = is_cand.astype(np.int32) - fwd_diag_idx = (blocks * is_cand).astype(np.int32)[:, None] - - bwd_mask_cnt = np.zeros(num_blocks, np.int32) - bwd_mask_idx = np.zeros((num_blocks, 1), np.int32) - bwd_full_cnt = np.where(is_cand, 0, num_blocks).astype(np.int32) - bwd_full_idx = np.tile(blocks, (num_blocks, 1)) - bwd_full_idx[is_cand] = 0 - bwd_diag_cnt = is_cand.astype(np.int32) - bwd_diag_idx = (blocks * is_cand).astype(np.int32)[:, None] - - def _bcast(arr): - return np.ascontiguousarray(np.broadcast_to(arr, (batch_size, num_q_heads) + arr.shape)) - - _DENSE_BS_LAYOUT_CACHE[key] = ( - tuple( - _bcast(a) - for a in ( - fwd_mask_cnt, - fwd_mask_idx, - fwd_full_cnt, - fwd_full_idx, - fwd_diag_cnt, - fwd_diag_idx, - ) - ), - tuple( - _bcast(a) - for a in ( - bwd_mask_cnt, - bwd_mask_idx, - bwd_full_cnt, - bwd_full_idx, - bwd_diag_cnt, - bwd_diag_idx, - ) - ), - ) - return _DENSE_BS_LAYOUT_CACHE[key] - - -def ranker_attention_fa4(q, k, v, sm_scale, block_sparse_layout): +def build_dense_block_sparse_layout(seq_len, hist_len, num_q_heads, hist_valid_len): + block = 128 + cand_len = seq_len - hist_len + assert hist_len % block == 0 and cand_len % block == 0, ( + f"dense block-sparse needs tile-aligned lengths (hist={hist_len}, cand={cand_len})" + ) + h_blocks = hist_len // block + num_blocks = seq_len // block + cand_blocks = num_blocks - h_blocks + idx_w = max(h_blocks, 1) + + valid_len = jnp.asarray(hist_valid_len, jnp.int32).reshape(-1)[:, None] + batch_size = valid_len.shape[0] + n_full = valid_len // block + has_partial = (valid_len % block != 0).astype(jnp.int32) + t = jnp.arange(num_blocks, dtype=jnp.int32)[None, :] + is_hist = t < h_blocks + is_cand_i32 = (~is_hist).astype(jnp.int32) + is_full = ~is_hist | (t < n_full) + is_partial = is_hist & (t == n_full) & (has_partial == 1) + + j_h = jnp.arange(idx_w, dtype=jnp.int32)[None, None, :] + j_n = jnp.arange(num_blocks, dtype=jnp.int32)[None, None, :] + n_full3 = n_full[:, :, None] + + fwd_mask_cnt = jnp.where(is_partial, n_full + 1, jnp.where(is_full, has_partial, 0)) + fwd_mask_idx = jnp.where(is_partial[:, :, None], j_h, n_full3) + fwd_full_cnt = jnp.where(is_full, n_full, 0) + fwd_full_idx = jnp.broadcast_to(j_h, (batch_size, num_blocks, idx_w)) + fwd_diag_cnt = jnp.broadcast_to(is_cand_i32, (batch_size, num_blocks)) + fwd_diag_idx = jnp.broadcast_to((t * is_cand_i32)[:, :, None], (batch_size, num_blocks, 1)) + + full_q_seq = jnp.clip( + jnp.where(j_n < n_full3, j_n, h_blocks + j_n - n_full3), 0, num_blocks - 1 + ) + mask_q_seq = jnp.clip( + jnp.where(j_n <= n_full3, j_n, h_blocks + j_n - n_full3 - 1), 0, num_blocks - 1 + ) + bwd_mask_cnt = jnp.where( + is_partial, n_full + cand_blocks + 1, jnp.where(t < n_full, has_partial, 0) + ) + bwd_mask_idx = jnp.where( + is_partial[:, :, None], + jnp.broadcast_to(mask_q_seq, (batch_size, num_blocks, num_blocks)), + n_full3, + ) + bwd_full_cnt = jnp.where(t < n_full, n_full + cand_blocks, 0) + bwd_full_idx = jnp.broadcast_to(full_q_seq, (batch_size, num_blocks, num_blocks)) + bwd_diag_cnt = fwd_diag_cnt + bwd_diag_idx = fwd_diag_idx + + valid_block_upper = jnp.where(is_hist, jnp.clip(valid_len - block * t, 0, block), 0) + valid_block_lower = jnp.where(is_hist, block, 0) + + def _bcast(arr): + arr = arr.astype(jnp.int32) + return jnp.broadcast_to(arr[:, None], (batch_size, num_q_heads) + arr.shape[1:]) + + fwd_arrays = ( + fwd_mask_cnt, + fwd_mask_idx, + fwd_full_cnt, + fwd_full_idx, + fwd_diag_cnt, + fwd_diag_idx, + ) + bwd_arrays = ( + bwd_mask_cnt, + bwd_mask_idx, + bwd_full_cnt, + bwd_full_idx, + bwd_diag_cnt, + bwd_diag_idx, + ) + fwd_bs = tuple(_bcast(a) for a in fwd_arrays) + bwd_bs = tuple(_bcast(a) for a in bwd_arrays) + return fwd_bs, bwd_bs, _bcast(valid_block_upper), _bcast(valid_block_lower) + + +def ranker_attention_fa4( + q, + k, + v, + sm_scale, + block_sparse_layout, + valid_block_upper=None, + valid_block_lower=None, +): import cuda.bindings.driver as cuda_driver import cutlass import cutlass.cute as cute @@ -95,6 +120,13 @@ def ranker_attention_fa4(q, k, v, sm_scale, block_sparse_layout): dKV_postprocess = True fwd_bs, bwd_bs = block_sparse_layout + if valid_block_upper is None or valid_block_lower is None: + if valid_block_upper is not None or valid_block_lower is not None: + raise ValueError("valid_block_upper and valid_block_lower must be provided together") + valid_block_upper = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32) + valid_block_lower = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32) + valid_block_upper = jnp.broadcast_to(valid_block_upper, fwd_bs[2].shape) + valid_block_lower = jnp.broadcast_to(valid_block_lower, fwd_bs[2].shape) bs_num_blocks = int(fwd_bs[3].shape[-2]) bs_max_hist_blocks = int(fwd_bs[3].shape[-1]) if bs_num_blocks != (seq_len + m_block - 1) // m_block: @@ -112,6 +144,8 @@ def ranker_attention_fa4(q, k, v, sm_scale, block_sparse_layout): seq_len, bs_num_blocks, bs_max_hist_blocks, + int(fwd_bs[1].shape[-1]), + int(bwd_bs[1].shape[-1]), use_pack_gqa, ) @@ -144,6 +178,8 @@ def launch_fwd( mFullIdx: cute.Tensor, mDiagCnt: cute.Tensor, mDiagIdx: cute.Tensor, + mValidUpper: cute.Tensor, + mValidLower: cute.Tensor, mO: cute.Tensor, mLSE: cute.Tensor, softmax_scale: cutlass.Float32, @@ -160,6 +196,8 @@ def launch_fwd( diag_block_cnt=mDiagCnt, diag_block_idx=mDiagIdx, dq_write_order_diag=None, + valid_block_upper=mValidUpper, + valid_block_lower=mValidLower, ) fa_fwd( mQ, @@ -220,6 +258,8 @@ def launch_bwd( mFullIdx: cute.Tensor, mDiagCnt: cute.Tensor, mDiagIdx: cute.Tensor, + mValidUpper: cute.Tensor, + mValidLower: cute.Tensor, mdQa: cute.Tensor, mdKa: cute.Tensor, mdVa: cute.Tensor, @@ -237,6 +277,8 @@ def launch_bwd( diag_block_cnt=mDiagCnt, diag_block_idx=mDiagIdx, dq_write_order_diag=None, + valid_block_upper=mValidUpper, + valid_block_lower=mValidLower, ) fa_bwd( mQ, @@ -270,7 +312,7 @@ def launch_bwd( jax.ShapeDtypeStruct(dk_accum_shape, jnp.float32), jax.ShapeDtypeStruct(dk_accum_shape, jnp.float32), ], - input_output_aliases={12: 0, 13: 1, 14: 2}, + input_output_aliases={14: 0, 15: 1, 16: 2}, use_static_tensors=False, softmax_scale=cutlass.Float32(sm_scale), ) @@ -352,22 +394,23 @@ def launch_post_dv( c = _FA4_KERNEL_CACHE[cache_key] - _fwd_extra = tuple(fwd_bs) - _bwd_extra = tuple(bwd_bs) - @jax.custom_vjp - def _attention(q, k, v): - out, _lse = c["fwd_call"](q, k, v, *_fwd_extra) + def _attention(q, k, v, *bs_args): + fbs = bs_args[:6] + valid_bounds = bs_args[12:] + out, _lse = c["fwd_call"](q, k, v, *fbs, *valid_bounds) return out - def _attention_fwd(q, k, v): - out, lse = c["fwd_call"](q, k, v, *_fwd_extra) - out = checkpoint_name(out, "attn_outputs") - lse = checkpoint_name(lse, "attn_outputs") - return out, (q, k, v, out, lse) + def _attention_fwd(q, k, v, *bs_args): + fbs = bs_args[:6] + valid_bounds = bs_args[12:] + out, lse = c["fwd_call"](q, k, v, *fbs, *valid_bounds) + return out, (q, k, v, out, lse, bs_args) def _attention_bwd(res, g): - q, k, v, out, lse = res + q, k, v, out, lse, bs_args = res + bbs = bs_args[6:12] + valid_bounds = bs_args[12:] dpsum = jnp.sum(out.astype(jnp.float32) * g.astype(jnp.float32), axis=-1).transpose(0, 2, 1) if dpsum.shape[-1] < sr_q: @@ -381,13 +424,31 @@ def _attention_bwd(res, g): dk_accum_init = jnp.zeros((batch_size, num_kv_heads, sr_k * hdr), dtype=jnp.float32) dv_accum_init = jnp.zeros_like(dk_accum_init) dq_accum, dk_accum, dv_accum = c["bwd_call"]( - q, k, v, g, lse_log2, dpsum, *_bwd_extra, dq_accum_init, dk_accum_init, dv_accum_init + q, + k, + v, + g, + lse_log2, + dpsum, + *bbs, + *valid_bounds, + dq_accum_init, + dk_accum_init, + dv_accum_init, ) (dq,) = c["post_dq_call"](dq_accum) (dk,) = c["post_dk_call"](dk_accum) (dv,) = c["post_dv_call"](dv_accum) - return dq, dk, dv + return (dq, dk, dv) + (None,) * len(bs_args) _attention.defvjp(_attention_fwd, _attention_bwd) - return _attention(q, k, v) + return _attention( + q, + k, + v, + *fwd_bs, + *bwd_bs, + valid_block_upper, + valid_block_lower, + ) diff --git a/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py index c0dc49e1..928d2072 100644 --- a/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py +++ b/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py @@ -10,7 +10,6 @@ import jax import jax.numpy as jnp import numpy as np -from jax.ad_checkpoint import checkpoint_name @dataclasses.dataclass(frozen=True) @@ -33,12 +32,25 @@ def _build_packed_block_sparse( block_size=128, max_hist_blocks=None, num_blocks=None, + real_history_starts=None, + num_user_prefix_tokens=0, ): - assert all(h % block_size == 0 for h in hist_sizes), ( + hist_sizes = np.asarray(hist_sizes, dtype=np.int32) + assert np.all(hist_sizes % block_size == 0), ( f"hist_sizes must be multiples of block_size={block_size}, got {hist_sizes}" ) assert transformer_candidate_seq_len % block_size == 0 - h_blocks = [h // block_size for h in hist_sizes] + assert 0 <= num_user_prefix_tokens <= block_size + if real_history_starts is None: + real_history_starts = np.zeros_like(hist_sizes) + num_user_prefix_tokens = 0 + else: + real_history_starts = np.asarray(real_history_starts, dtype=np.int32) + assert real_history_starts.shape == hist_sizes.shape + assert np.all(real_history_starts >= num_user_prefix_tokens) + assert np.all(real_history_starts <= hist_sizes) + + h_blocks = (hist_sizes // block_size).tolist() cand_blocks_per_user = transformer_candidate_seq_len // block_size actual_num_blocks = sum(hb + cand_blocks_per_user for hb in h_blocks) if num_blocks is None: @@ -48,21 +60,6 @@ def _build_packed_block_sparse( f"num_blocks={num_blocks} too small for hist_sizes (need {actual_num_blocks})" ) - example_id_per_block = np.full(num_blocks, -1, dtype=np.int32) - is_cand_per_block = np.zeros(num_blocks, dtype=bool) - hist_block_start_per_user: list[int] = [] - cand_block_start_per_user: list[int] = [] - cur = 0 - for i, hb in enumerate(h_blocks): - hist_block_start_per_user.append(cur) - example_id_per_block[cur : cur + hb] = i - cur += hb - cand_block_start_per_user.append(cur) - for j in range(cand_blocks_per_user): - example_id_per_block[cur + j] = i - is_cand_per_block[cur + j] = True - cur += cand_blocks_per_user - actual_max_hist_blocks = max(h_blocks) if h_blocks else 1 if max_hist_blocks is None: max_hist_blocks = actual_max_hist_blocks @@ -70,58 +67,109 @@ def _build_packed_block_sparse( assert max_hist_blocks >= actual_max_hist_blocks, ( f"max_hist_blocks={max_hist_blocks} too small (max h={actual_max_hist_blocks})" ) + max_hist_blocks = max(max_hist_blocks, 1) max_q_per_n = max_hist_blocks + cand_blocks_per_user - fwd_mask_cnt = np.zeros((num_blocks,), dtype=np.int32) - fwd_mask_idx = np.zeros((num_blocks, 1), dtype=np.int32) - fwd_full_cnt = np.zeros((num_blocks,), dtype=np.int32) + valid_block_upper = np.zeros(num_blocks, dtype=np.int32) + valid_block_lower = np.full(num_blocks, block_size, dtype=np.int32) + is_partial_per_block = np.zeros(num_blocks, dtype=bool) + is_empty_per_block = np.ones(num_blocks, dtype=bool) + example_id_per_block = np.full(num_blocks, -1, dtype=np.int32) + is_cand_per_block = np.zeros(num_blocks, dtype=bool) + hist_block_start_per_user: list[int] = [] + + cur = 0 + for user_idx, hb in enumerate(h_blocks): + hist_block_start_per_user.append(cur) + example_id_per_block[cur : cur + hb] = user_idx + real_start = int(real_history_starts[user_idx]) + for local_block in range(hb): + physical_block = cur + local_block + tile_start = local_block * block_size + upper = int(np.clip(num_user_prefix_tokens - tile_start, 0, block_size)) + lower = int(np.clip(real_start - tile_start, 0, block_size)) + if upper >= lower: + upper = lower = 0 + is_empty = False + is_partial = False + elif upper == 0 and lower == block_size: + is_empty = True + is_partial = False + else: + is_empty = False + is_partial = True + valid_block_upper[physical_block] = upper + valid_block_lower[physical_block] = lower + is_empty_per_block[physical_block] = is_empty + is_partial_per_block[physical_block] = is_partial + cur += hb + + for j in range(cand_blocks_per_user): + physical_block = cur + j + example_id_per_block[physical_block] = user_idx + is_cand_per_block[physical_block] = True + is_empty_per_block[physical_block] = False + valid_block_upper[physical_block] = 0 + valid_block_lower[physical_block] = 0 + cur += cand_blocks_per_user + + fwd_mask_cnt = np.zeros(num_blocks, dtype=np.int32) + fwd_mask_idx = np.zeros((num_blocks, max_hist_blocks), dtype=np.int32) + fwd_full_cnt = np.zeros(num_blocks, dtype=np.int32) fwd_full_idx = np.zeros((num_blocks, max_hist_blocks), dtype=np.int32) - fwd_diag_cnt = np.zeros((num_blocks,), dtype=np.int32) + fwd_diag_cnt = np.zeros(num_blocks, dtype=np.int32) fwd_diag_idx = np.zeros((num_blocks, 1), dtype=np.int32) - for m in range(num_blocks): - ex_id = int(example_id_per_block[m]) - if ex_id < 0: + + for m_block in range(num_blocks): + user_idx = int(example_id_per_block[m_block]) + if user_idx < 0 or (not is_cand_per_block[m_block] and is_empty_per_block[m_block]): continue - hb = h_blocks[ex_id] - hstart = hist_block_start_per_user[ex_id] - fwd_full_cnt[m] = hb - for j in range(hb): - fwd_full_idx[m, j] = hstart + j - if is_cand_per_block[m]: - fwd_diag_cnt[m] = 1 - fwd_diag_idx[m, 0] = m - - bwd_mask_cnt = np.zeros((num_blocks,), dtype=np.int32) - bwd_mask_idx = np.zeros((num_blocks, 1), dtype=np.int32) - bwd_full_cnt = np.zeros((num_blocks,), dtype=np.int32) + hstart = hist_block_start_per_user[user_idx] + for n_block in range(hstart, hstart + h_blocks[user_idx]): + if is_empty_per_block[n_block]: + continue + is_mask_edge = is_partial_per_block[m_block] or is_partial_per_block[n_block] + if is_mask_edge: + idx = int(fwd_mask_cnt[m_block]) + fwd_mask_idx[m_block, idx] = n_block + fwd_mask_cnt[m_block] += 1 + else: + idx = int(fwd_full_cnt[m_block]) + fwd_full_idx[m_block, idx] = n_block + fwd_full_cnt[m_block] += 1 + if is_cand_per_block[m_block]: + fwd_diag_cnt[m_block] = 1 + fwd_diag_idx[m_block, 0] = m_block + + bwd_mask_lists: list[list[int]] = [[] for _ in range(num_blocks)] + bwd_full_lists: list[list[int]] = [[] for _ in range(num_blocks)] + bwd_diag_lists: list[list[int]] = [[] for _ in range(num_blocks)] + for m_block in range(num_blocks): + for i in range(int(fwd_mask_cnt[m_block])): + bwd_mask_lists[int(fwd_mask_idx[m_block, i])].append(m_block) + for i in range(int(fwd_full_cnt[m_block])): + bwd_full_lists[int(fwd_full_idx[m_block, i])].append(m_block) + for i in range(int(fwd_diag_cnt[m_block])): + bwd_diag_lists[int(fwd_diag_idx[m_block, i])].append(m_block) + + bwd_mask_cnt = np.asarray([len(xs) for xs in bwd_mask_lists], dtype=np.int32) + bwd_mask_idx = np.zeros((num_blocks, max_q_per_n), dtype=np.int32) + bwd_full_cnt = np.asarray([len(xs) for xs in bwd_full_lists], dtype=np.int32) bwd_full_idx = np.zeros((num_blocks, max_q_per_n), dtype=np.int32) - bwd_diag_cnt = np.zeros((num_blocks,), dtype=np.int32) + bwd_diag_cnt = np.asarray([len(xs) for xs in bwd_diag_lists], dtype=np.int32) bwd_diag_idx = np.zeros((num_blocks, 1), dtype=np.int32) - for n_blk in range(num_blocks): - ex_id = int(example_id_per_block[n_blk]) - if ex_id < 0: - continue - if is_cand_per_block[n_blk]: - bwd_diag_cnt[n_blk] = 1 - bwd_diag_idx[n_blk, 0] = n_blk - else: - hb = h_blocks[ex_id] - hstart = hist_block_start_per_user[ex_id] - cand_blk_start = cand_block_start_per_user[ex_id] - cnt = 0 - for j in range(hb): - bwd_full_idx[n_blk, cnt] = hstart + j - cnt += 1 - for j in range(cand_blocks_per_user): - bwd_full_idx[n_blk, cnt] = cand_blk_start + j - cnt += 1 - bwd_full_cnt[n_blk] = cnt + for n_block in range(num_blocks): + bwd_mask_idx[n_block, : len(bwd_mask_lists[n_block])] = bwd_mask_lists[n_block] + bwd_full_idx[n_block, : len(bwd_full_lists[n_block])] = bwd_full_lists[n_block] + bwd_diag_idx[n_block, : len(bwd_diag_lists[n_block])] = bwd_diag_lists[n_block] return dict( total_seq_len=num_blocks * block_size, num_blocks=num_blocks, fwd=(fwd_mask_cnt, fwd_mask_idx, fwd_full_cnt, fwd_full_idx, fwd_diag_cnt, fwd_diag_idx), bwd=(bwd_mask_cnt, bwd_mask_idx, bwd_full_cnt, bwd_full_idx, bwd_diag_cnt, bwd_diag_idx), + valid_block_upper=valid_block_upper, + valid_block_lower=valid_block_lower, ) @@ -140,6 +188,8 @@ class BlockSparseLayout: bwd_full_idx: np.ndarray bwd_diag_cnt: np.ndarray bwd_diag_idx: np.ndarray + valid_block_upper: np.ndarray + valid_block_lower: np.ndarray def build_block_sparse_layout( @@ -147,6 +197,8 @@ def build_block_sparse_layout( transformer_candidate_seq_len: int, max_history_seq_len: int, packed_seq_len: int | None = None, + padding_mask: np.ndarray | None = None, + num_user_prefix_tokens: int = 0, ) -> BlockSparseLayout: per_user_lens = np.diff(cu_seqlens, axis=-1) block_size = get_device_tuning_config().block_q @@ -154,8 +206,15 @@ def build_block_sparse_layout( assert transformer_candidate_seq_len % block_size == 0, ( "Candidate sequence length must be block-aligned" ) + assert 0 <= num_user_prefix_tokens <= block_size + + num_devices, bs_per_device = per_user_lens.shape + if padding_mask is not None: + padding_mask = np.asarray(padding_mask, dtype=np.bool_) + assert padding_mask.ndim == 2 and padding_mask.shape[0] == num_devices + if packed_seq_len is not None: + assert padding_mask.shape[1] == packed_seq_len - num_devices, _ = per_user_lens.shape if packed_seq_len is not None: assert packed_seq_len % block_size == 0, ( f"packed_seq_len={packed_seq_len} must be a multiple of block_size={block_size}" @@ -167,22 +226,48 @@ def build_block_sparse_layout( total_blocks_per_device = packed_seq_len // block_size else: total_blocks_per_device = int(per_user_lens[0].sum()) // block_size - max_hist_blocks = max_history_seq_len // block_size + max_hist_blocks = (max_history_seq_len + block_size - 1) // block_size fwd_arrays: list[list[np.ndarray]] = [[] for _ in range(6)] bwd_arrays: list[list[np.ndarray]] = [[] for _ in range(6)] + valid_upper_arrays: list[np.ndarray] = [] + valid_lower_arrays: list[np.ndarray] = [] for d in range(num_devices): hist_sizes = per_user_lens[d] - transformer_candidate_seq_len + if padding_mask is None: + real_history_starts = None + else: + real_history_starts = np.zeros(bs_per_device, dtype=np.int32) + for user_idx in range(bs_per_device): + seq_start = int(cu_seqlens[d, user_idx]) + hist_size = int(hist_sizes[user_idx]) + hist_valid = padding_mask[ + d, + seq_start + num_user_prefix_tokens : seq_start + hist_size, + ] + valid_offsets = np.flatnonzero(hist_valid) + real_start = ( + num_user_prefix_tokens + int(valid_offsets[0]) + if valid_offsets.size + else hist_size + ) + assert np.all(hist_valid[: real_start - num_user_prefix_tokens] == 0) + assert np.all(hist_valid[real_start - num_user_prefix_tokens :] == 1) + real_history_starts[user_idx] = real_start info = _build_packed_block_sparse( hist_sizes, transformer_candidate_seq_len=transformer_candidate_seq_len, max_hist_blocks=max_hist_blocks, num_blocks=total_blocks_per_device, + real_history_starts=real_history_starts, + num_user_prefix_tokens=num_user_prefix_tokens, ) for i, arr in enumerate(info["fwd"]): fwd_arrays[i].append(arr) for i, arr in enumerate(info["bwd"]): bwd_arrays[i].append(arr) + valid_upper_arrays.append(info["valid_block_upper"]) + valid_lower_arrays.append(info["valid_block_lower"]) fwd = [np.stack(arrs, axis=0) for arrs in fwd_arrays] bwd = [np.stack(arrs, axis=0) for arrs in bwd_arrays] @@ -199,13 +284,23 @@ def build_block_sparse_layout( bwd_full_idx=bwd[3], bwd_diag_cnt=bwd[4], bwd_diag_idx=bwd[5], + valid_block_upper=np.stack(valid_upper_arrays, axis=0), + valid_block_lower=np.stack(valid_lower_arrays, axis=0), ) _FA4_PACKED_CACHE = {} -def ranker_attention_varlen_fa4(q, k, v, sm_scale, block_sparse_layout): +def ranker_attention_varlen_fa4( + q, + k, + v, + sm_scale, + block_sparse_layout, + valid_block_upper=None, + valid_block_lower=None, +): import cuda.bindings.driver as cuda_driver import cutlass import cutlass.cute as cute @@ -227,6 +322,13 @@ def ranker_attention_varlen_fa4(q, k, v, sm_scale, block_sparse_layout): use_pack_gqa = qpk > 1 and (block_size % qpk == 0) fwd_bs, _ = block_sparse_layout + if valid_block_upper is None or valid_block_lower is None: + if valid_block_upper is not None or valid_block_lower is not None: + raise ValueError("valid_block_upper and valid_block_lower must be provided together") + valid_block_upper = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32) + valid_block_lower = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32) + valid_block_upper = jnp.broadcast_to(valid_block_upper, fwd_bs[2].shape) + valid_block_lower = jnp.broadcast_to(valid_block_lower, fwd_bs[2].shape) bs_max_hist_blocks = int(fwd_bs[3].shape[-1]) bs_num_blocks = int(fwd_bs[3].shape[-2]) @@ -279,6 +381,8 @@ def launch_fwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mO, mLSE, softmax_scale: cutlass.Float32, @@ -295,6 +399,8 @@ def launch_fwd( diag_block_cnt=mDiagCnt, diag_block_idx=mDiagIdx, dq_write_order_diag=None, + valid_block_upper=mValidUpper, + valid_block_lower=mValidLower, ) fa_fwd( mQ, @@ -355,6 +461,8 @@ def launch_bwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mdQa, mdKa, mdVa, @@ -372,6 +480,8 @@ def launch_bwd( diag_block_cnt=mDiagCnt, diag_block_idx=mDiagIdx, dq_write_order_diag=None, + valid_block_upper=mValidUpper, + valid_block_lower=mValidLower, ) fa_bwd( mQ, @@ -405,7 +515,7 @@ def launch_bwd( jax.ShapeDtypeStruct(dk_accum_shape, jnp.float32), jax.ShapeDtypeStruct(dk_accum_shape, jnp.float32), ], - input_output_aliases={12: 0, 13: 1, 14: 2}, + input_output_aliases={14: 0, 15: 1, 16: 2}, use_static_tensors=False, softmax_scale=cutlass.Float32(sm_scale), ) @@ -426,6 +536,8 @@ def launch_bwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mdQa, mdK, mdV, @@ -443,6 +555,8 @@ def launch_bwd( diag_block_cnt=mDiagCnt, diag_block_idx=mDiagIdx, dq_write_order_diag=None, + valid_block_upper=mValidUpper, + valid_block_lower=mValidLower, ) fa_bwd( mQ, @@ -476,7 +590,7 @@ def launch_bwd( jax.ShapeDtypeStruct(k_shape, jnp.float32), jax.ShapeDtypeStruct(k_shape, jnp.float32), ], - input_output_aliases={12: 0}, + input_output_aliases={14: 0}, use_static_tensors=False, softmax_scale=cutlass.Float32(sm_scale), ) @@ -547,19 +661,20 @@ def launch_post_dv(stream, mAccum, mOut, scale: cutlass.Float32): @jax.custom_vjp def _attention(q, k, v, *bs_args): fbs = bs_args[:6] - out, _lse = c["fwd_call"](q, k, v, *fbs) + valid_bounds = bs_args[12:] + out, _lse = c["fwd_call"](q, k, v, *fbs, *valid_bounds) return out def _attention_fwd(q, k, v, *bs_args): fbs = bs_args[:6] - out, lse = c["fwd_call"](q, k, v, *fbs) - out = checkpoint_name(out, "attn_outputs") - lse = checkpoint_name(lse, "attn_outputs") + valid_bounds = bs_args[12:] + out, lse = c["fwd_call"](q, k, v, *fbs, *valid_bounds) return out, (q, k, v, out, lse, bs_args) def _attention_bwd(res, g): q, k, v, out, lse, bs_args = res - bbs = bs_args[6:] + bbs = bs_args[6:12] + valid_bounds = bs_args[12:] dpsum = jnp.sum(out.astype(jnp.float32) * g.astype(jnp.float32), axis=-1).transpose(0, 2, 1) if dpsum.shape[-1] < c["sr_q"]: dpsum = jnp.pad(dpsum, ((0, 0), (0, 0), (0, c["sr_q"] - dpsum.shape[-1]))) @@ -580,6 +695,7 @@ def _attention_bwd(res, g): lse_log2, dpsum, *bbs, + *valid_bounds, dq_accum_init, dk_accum_init, dv_accum_init, @@ -596,6 +712,7 @@ def _attention_bwd(res, g): lse_log2, dpsum, *bbs, + *valid_bounds, dq_accum_init, ) (dq,) = c["post_dq_call"](dq_accum) @@ -605,4 +722,12 @@ def _attention_bwd(res, g): return (dq, dk, dv) + (None,) * len(bs_args) _attention.defvjp(_attention_fwd, _attention_bwd) - return _attention(q, k, v, *fwd_bs, *bwd_bs) + return _attention( + q, + k, + v, + *fwd_bs, + *bwd_bs, + valid_block_upper, + valid_block_lower, + ) diff --git a/phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py b/phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py index db9df4d7..8df0bc0b 100644 --- a/phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py +++ b/phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py @@ -55,6 +55,8 @@ class BlockSparseTensors(NamedTuple): diag_block_cnt: cute.Tensor | None = None diag_block_idx: cute.Tensor | None = None dq_write_order_diag: cute.Tensor | None = None + valid_block_upper: cute.Tensor | None = None + valid_block_lower: cute.Tensor | None = None def __new_from_mlir_values__(self, values): new_fields = [] @@ -82,6 +84,8 @@ class BlockSparseTensorsTorch(NamedTuple): diag_block_cnt: torch.Tensor | None = None diag_block_idx: torch.Tensor | None = None dq_write_order_diag: torch.Tensor | None = None + valid_block_upper: torch.Tensor | None = None + valid_block_lower: torch.Tensor | None = None def _ordered_to_dense_simple( @@ -454,6 +458,27 @@ def normalize_block_sparse_tensors( hint, mask_cnt.device, ) + metadata_block_shape = expected_count_shape + valid_block_upper = _check_and_expand_metadata_tensor( + "valid_block_upper", + tensors.valid_block_upper, + metadata_block_shape, + context, + hint, + mask_cnt.device, + ) + valid_block_lower = _check_and_expand_metadata_tensor( + "valid_block_lower", + tensors.valid_block_lower, + metadata_block_shape, + context, + hint, + mask_cnt.device, + ) + if (valid_block_upper is None) != (valid_block_lower is None): + raise ValueError( + "valid_block_upper and valid_block_lower must both be provided or both be None" + ) spt = tensors.spt if spt is not None and not isinstance(spt, bool): raise ValueError("spt must be a bool when provided") @@ -474,6 +499,8 @@ def normalize_block_sparse_tensors( diag_block_cnt=diag_cnt, diag_block_idx=diag_idx, dq_write_order_diag=dq_write_order_diag, + valid_block_upper=valid_block_upper, + valid_block_lower=valid_block_lower, ) @@ -501,6 +528,8 @@ def get_block_sparse_broadcast_pattern( tensors.diag_block_cnt, tensors.diag_block_idx, tensors.dq_write_order_diag, + tensors.valid_block_upper, + tensors.valid_block_lower, ): if tensor is not None: patterns.append(get_broadcast_dims(tensor)) @@ -658,6 +687,12 @@ def to_cute_block_sparse_tensors( if tensors.dq_write_order_diag is not None else None ) + valid_block_upper_tensor, valid_block_lower_tensor = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + if t is not None + else None + for t in (tensors.valid_block_upper, tensors.valid_block_lower) + ] return BlockSparseTensors( mask_block_cnt_tensor, @@ -671,6 +706,8 @@ def to_cute_block_sparse_tensors( diag_block_cnt_tensor, diag_block_idx_tensor, dq_write_order_diag_tensor, + valid_block_upper_tensor, + valid_block_lower_tensor, ) diff --git a/phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py b/phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py index 3a0d2eba..c442fc24 100644 --- a/phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py +++ b/phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py @@ -2721,6 +2721,11 @@ def compute_loop( ) mask = AttentionMaskCls(seqlen) n_block_for_cluster = n_block // self.cta_group_size + mask_block_union = const_expr( + blocksparse_tensors is not None + and blocksparse_tensors.valid_block_upper is not None + and blocksparse_tensors.valid_block_lower is not None + ) mask_fn = partial( mask.apply_mask_sm100_transposed, tScS_t2r=tScS_t2r, @@ -2730,6 +2735,13 @@ def compute_loop( mask_causal=self.is_causal, mask_local=self.is_local, mask_mod=self.mask_mod, + mask_block_union=mask_block_union, + valid_block_upper=( + blocksparse_tensors.valid_block_upper if mask_block_union else None + ), + valid_block_lower=( + blocksparse_tensors.valid_block_lower if mask_block_union else None + ), batch_idx=batch_idx, head_idx=head_idx, aux_tensors=aux_tensors, @@ -2857,8 +2869,16 @@ def compute_loop( (softmax_scale_log2, softmax_scale_log2), (-lse_pair[0], -lse_pair[1]), ) - tSrS_cur[2 * v] = cute.math.exp2(tSrS_cur[2 * v], fastmath=True) - tSrS_cur[2 * v + 1] = cute.math.exp2(tSrS_cur[2 * v + 1], fastmath=True) + tSrS_cur[2 * v] = ( + Float32(0.0) + if lse_pair[0] == -Float32.inf + else cute.math.exp2(tSrS_cur[2 * v], fastmath=True) + ) + tSrS_cur[2 * v + 1] = ( + Float32(0.0) + if lse_pair[1] == -Float32.inf + else cute.math.exp2(tSrS_cur[2 * v + 1], fastmath=True) + ) utils.cvt_f16(tSrS_cur, tSrP_r2t[None, stage, 0, 0]) if const_expr(stage == 0): cute.arch.fence_view_async_tmem_load() diff --git a/phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py b/phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py index ac25d8d8..ceb56785 100644 --- a/phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py +++ b/phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py @@ -1877,9 +1877,21 @@ def softmax_loop( ) mask_mod = self.mask_mod if const_expr(self.mask_mod is not None) else None + mask_block_union = const_expr( + blocksparse_tensors is not None + and blocksparse_tensors.valid_block_upper is not None + and blocksparse_tensors.valid_block_lower is not None + ) mask_fn = partial( mask.apply_mask_sm100, mask_mod=mask_mod, + mask_block_union=mask_block_union, + valid_block_upper=( + blocksparse_tensors.valid_block_upper if mask_block_union else None + ), + valid_block_lower=( + blocksparse_tensors.valid_block_lower if mask_block_union else None + ), fastdiv_mods=fastdiv_mods, head_divmod=head_divmod, **shared_mask_kwargs, diff --git a/phoenix/xrex/cutedsl/ranker_fa4/mask.py b/phoenix/xrex/cutedsl/ranker_fa4/mask.py index 2dcfcc70..202de122 100644 --- a/phoenix/xrex/cutedsl/ranker_fa4/mask.py +++ b/phoenix/xrex/cutedsl/ranker_fa4/mask.py @@ -354,6 +354,9 @@ def apply_mask_sm100( mask_local: cutlass.Constexpr[bool] = False, mask_mod: cutlass.Constexpr[Optional[Callable]] = None, mask_diagonal: cutlass.Constexpr[bool] = False, + mask_block_union: cutlass.Constexpr[bool] = False, + valid_block_upper: Optional[cute.Tensor] = None, + valid_block_lower: Optional[cute.Tensor] = None, batch_idx: Int32 = None, head_idx: Int32 = None, aux_tensors: Optional[list] = None, @@ -364,8 +367,12 @@ def apply_mask_sm100( rBitmask: Optional[cute.Tensor] = None, ) -> None: assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" - assert not (mask_diagonal and (mask_causal or mask_local or mask_mod is not None)), ( - "mask_diagonal is mutually exclusive with mask_causal/mask_local/mask_mod" + assert not ( + mask_diagonal + and (mask_causal or mask_local or mask_mod is not None or mask_block_union) + ), "mask_diagonal is mutually exclusive with all other masks" + assert not (mask_block_union and (mask_causal or mask_local or mask_mod is not None)), ( + "mask_block_union is mutually exclusive with causal/local/mask_mod" ) acc_shape = (self.tile_m, self.tile_n) cS = cute.make_identity_tensor(acc_shape if not self.swap_AB else acc_shape[::-1]) @@ -414,6 +421,41 @@ def apply_mask_sm100( mask_row = global_row acc_S[i] = -Float32.inf if mask_row >= self.seqlen_q else acc_S[i] + elif const_expr(mask_block_union): + assert valid_block_upper is not None and valid_block_lower is not None + qpk = const_expr(self.qhead_per_kvhead_packgqa) + if const_expr(qpk != 1): + q_block = m_block // qpk + else: + q_block = m_block + q_upper = valid_block_upper[batch_idx, head_idx, q_block] + q_lower = valid_block_lower[batch_idx, head_idx, q_block] + kv_upper = valid_block_upper[batch_idx, head_idx, n_block] + kv_lower = valid_block_lower[batch_idx, head_idx, n_block] + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][0] + col_coord = tScS_t2r[i][1] + if const_expr(qpk != 1): + q_coord = ((m_block % qpk) * self.tile_m + row_coord) // qpk + else: + q_coord = row_coord + q_valid = (q_coord < q_upper) or (q_coord >= q_lower) + kv_valid = (col_coord < kv_upper) or (col_coord >= kv_lower) + valid = q_valid and kv_valid + if const_expr(mask_seqlen): + global_col = col_coord + n_block * self.tile_n + valid = valid and global_col < self.seqlen_k + if check_q_boundary: + global_row = row_coord + m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert head_divmod is not None + mask_row, _ = divmod(global_row, head_divmod) + else: + mask_row = global_row + valid = valid and mask_row < self.seqlen_q + acc_S[i] = acc_S[i] if valid else -Float32.inf + elif const_expr(not mask_causal and not mask_local and mask_mod is None): if const_expr(mask_seqlen): if const_expr(not r2p): @@ -551,6 +593,9 @@ def apply_mask_sm100_transposed( fastdiv_mods=(None, None), is_full_block: bool = False, is_diagonal_block: bool = False, + mask_block_union: cutlass.Constexpr[bool] = False, + valid_block_upper: Optional[cute.Tensor] = None, + valid_block_lower: Optional[cute.Tensor] = None, check_m_boundary: bool = True, ) -> None: assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" @@ -581,6 +626,44 @@ def apply_mask_sm100_transposed( out_of_bounds = q_out_of_bounds or kv_out_of_bounds acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + elif const_expr(mask_block_union): + assert valid_block_upper is not None and valid_block_lower is not None + if is_full_block: + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = -cutlass.Float32.inf + elif check_m_boundary: + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + out_of_bounds = (global_q >= self.seqlen_q) or ( + global_kv >= self.seqlen_k + ) + acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + else: + q_upper = valid_block_upper[batch_idx, head_idx, m_block] + q_lower = valid_block_lower[batch_idx, head_idx, m_block] + kv_upper = valid_block_upper[batch_idx, head_idx, n_block] + kv_lower = valid_block_lower[batch_idx, head_idx, n_block] + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + q_valid = (row_coord < q_upper) or (row_coord >= q_lower) + kv_valid = (col_coord < kv_upper) or (col_coord >= kv_lower) + valid = q_valid and kv_valid + if const_expr(mask_seqlen): + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + q_out_of_bounds = check_m_boundary and (global_q >= self.seqlen_q) + kv_out_of_bounds = global_kv >= self.seqlen_k + valid = valid and not (q_out_of_bounds or kv_out_of_bounds) + acc_S[i] = acc_S[i] if valid else -cutlass.Float32.inf + elif const_expr(not mask_causal and not mask_local and mask_mod is not None): if is_full_block: if const_expr(mask_seqlen): diff --git a/phoenix/xrex/data/grpc_recsys.py b/phoenix/xrex/data/grpc_recsys.py index 05016047..1c26bc45 100644 --- a/phoenix/xrex/data/grpc_recsys.py +++ b/phoenix/xrex/data/grpc_recsys.py @@ -90,6 +90,7 @@ def make( output_vocab_size=self.hash_table.output_vocab_size, history_seq_len=self.history_seq_len, candidate_seq_len=self.candidate_seq_len, + enable_stale_post=self.enable_stale_post, ) prefetch_queue = Queue(maxsize=2) @@ -252,6 +253,9 @@ def _create_recsys_features_batch(self, batch_size: int) -> RecsysFeaturesBatch: ), product_surface=np.zeros((batch_size, self.candidate_seq_len), dtype=np.int32), client_app_id=np.zeros((batch_size, self.candidate_seq_len), dtype=np.int32), + trained_candidate_mask=np.ones( + (batch_size, self.candidate_seq_len), dtype=np.bool_ + ), post_creation_ts_sec=np.zeros((batch_size, self.candidate_seq_len), dtype=np.int32), post_ids=None, promoted_ids=np.zeros((batch_size, self.candidate_seq_len), dtype=np.int64), @@ -368,6 +372,7 @@ def example_data( post_ids=None, product_surface=np.zeros((batch_size, candidate_seq_len), dtype=np.int32), client_app_id=np.zeros((batch_size, candidate_seq_len), dtype=np.int32), + trained_candidate_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), post_creation_ts_sec=np.zeros((batch_size, candidate_seq_len), dtype=np.int32), continuous_actions=np.zeros((batch_size, candidate_seq_len, 2), dtype=np.float32), promoted_ids=np.zeros((batch_size, candidate_seq_len), dtype=np.int64), diff --git a/phoenix/xrex/data/parquet_recsys.py b/phoenix/xrex/data/parquet_recsys.py index d35abe31..f27721a4 100644 --- a/phoenix/xrex/data/parquet_recsys.py +++ b/phoenix/xrex/data/parquet_recsys.py @@ -767,6 +767,14 @@ def pad_array(arr: np.ndarray) -> np.ndarray: ) def pad_post_seq(post_seq: PostSeq) -> PostSeq: + padded = _pad_post_seq_fields(post_seq) + if (_tcm := post_seq.get("trained_candidate_mask")) is not None: + padded["trained_candidate_mask"] = np.pad( + _tcm, ((0, batch_size - num_rows), (0, 0)), constant_values=True + ) + return padded + + def _pad_post_seq_fields(post_seq: PostSeq) -> PostSeq: return PostSeq( impr_ts=pad_array(post_seq["impr_ts"]) if post_seq["impr_ts"] is not None else None, actions=pad_array(post_seq["actions"]) if post_seq["actions"] is not None else None, @@ -818,6 +826,9 @@ def pad_post_seq(post_seq: PostSeq) -> PostSeq: "sample_weights": pad_array(sw) if (sw := batch_unpadded.get("sample_weights")) is not None else None, + "sample_source": pad_array(ss) + if (ss := batch_unpadded.get("sample_source")) is not None + else None, } extras = cast(dict[str, np.ndarray], batch_unpadded) @@ -868,6 +879,7 @@ class PhoenixDataset(Dataset): sid_num_levels: int = 0 compute_post_unexplored_label: bool = False + enable_stale_post: bool = False multimodal_embedding_type: EmbeddingType | None = None @@ -1153,6 +1165,7 @@ def producer() -> None: global_post_sids=global_post_sids, sid_num_levels=self.sid_num_levels if self.use_post_sid else 0, compute_post_unexplored_label=self.compute_post_unexplored_label, + zero_stale_post_14d_candidate_counts=self.enable_stale_post, ) if self.use_conversion_labels and self.emit_conversion_label_keys: @@ -1273,6 +1286,9 @@ def example_data_shape(self, batch_size: int) -> Any: "sample_weights": jax.ShapeDtypeStruct(sw.shape, sw.dtype) if (sw := example_data.get("sample_weights")) is not None else None, + "sample_source": jax.ShapeDtypeStruct(ss.shape, ss.dtype) + if (ss := example_data.get("sample_source")) is not None + else None, } return batch_shape @@ -1345,6 +1361,7 @@ def example_data( ), product_surface=np.zeros((batch_size, candidate_seq_len), dtype=np.int32), client_app_id=np.zeros((batch_size, candidate_seq_len), dtype=np.int32), + trained_candidate_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), post_ids=np.zeros((batch_size, candidate_seq_len), dtype=np.int64) if self.include_candidate_post_ids else None, @@ -1379,6 +1396,7 @@ def example_data( and self.candidate_negative_filter != CandidateNegativeFilter.NONE else None, sample_weights=np.ones((batch_size, 1), dtype=np.float32), + sample_source=np.zeros((batch_size, 1), dtype=np.bool_), ) return batch @@ -1474,6 +1492,7 @@ def make_recsys_features_batch(self, batch_size: int) -> RecsysFeaturesBatch: auth_hashes=self.hash_table.get_author_hash(candidate_author_ids), product_surface=candidate_product_surface, client_app_id=np.zeros((batch_size, self.candidate_seq_len), dtype=np.int32), + trained_candidate_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), post_ids=candidate_tweet_ids.astype(np.int64) if self.include_candidate_post_ids else None, diff --git a/phoenix/xrex/data/recsys/constants.py b/phoenix/xrex/data/recsys/constants.py index 7d3e6868..ff350157 100644 --- a/phoenix/xrex/data/recsys/constants.py +++ b/phoenix/xrex/data/recsys/constants.py @@ -179,6 +179,22 @@ def to_pascal_case(s): ], } +SEARCH_RELEVANCE_ACTION_INDICES = [ + recsys_pb2.ActionName.CLIENT_TWEET_RELEVANT_TO_SEARCH, + recsys_pb2.ActionName.CLIENT_TWEET_NOT_RELEVANT_TO_SEARCH, +] + +search_engagement_to_action_types = { + **primary_engagement_to_action_types, + "IsRelevantToSearch": [ + "ClientTweetRelevantToSearch", + ], + "IsNotRelevantToSearch": [ + "ClientTweetNotRelevantToSearch", + ], + "IsSearchQueryReformulated": ["ClientTweetSearchQueryReformulated"], +} + ads_conversion_engagement_to_action_types = { **primary_engagement_to_action_types, "IsPurchaseConversion": [ @@ -209,6 +225,42 @@ def to_pascal_case(s): "IsAttributedMactClickInstall": [ "AdsAttributedMactClickInstall", ], + "IsAttributedMactPurchase": [ + "AdsAttributedMactPurchase", + ], + "IsAttributedMactAddToCart": [ + "AdsAttributedMactAddToCart", + ], + "IsAttributedMactLevelAchieved": [ + "AdsAttributedMactLevelAchieved", + ], + "IsAttributedMactTutorialComplete": [ + "AdsAttributedMactTutorialComplete", + ], + "IsAttributedMactSignUp": [ + "AdsAttributedMactSignUp", + ], + "IsAttributedMactCustom": [ + "AdsAttributedMactCustom", + ], + "IsAttributedMactViewPurchase": [ + "AdsAttributedMactViewPurchase", + ], + "IsAttributedMactViewAddToCart": [ + "AdsAttributedMactViewAddToCart", + ], + "IsAttributedMactViewLevelAchieved": [ + "AdsAttributedMactViewLevelAchieved", + ], + "IsAttributedMactViewTutorialComplete": [ + "AdsAttributedMactViewTutorialComplete", + ], + "IsAttributedMactViewSignUp": [ + "AdsAttributedMactViewSignUp", + ], + "IsAttributedMactViewCustom": [ + "AdsAttributedMactViewCustom", + ], "IsPurchaseConversion": [ "AdsPurchaseConversion", ], @@ -236,6 +288,59 @@ def to_pascal_case(s): "IsAttributedViewConversion": [ "AdsAttributedViewConversion", ], + "IsPurchaseConversionViewThrough": [ + "AdsPurchaseConversionViewThrough", + ], + "IsAddToCartConversionViewThrough": [ + "AdsAddToCartConversionViewThrough", + ], + "IsCheckoutInitiatedConversionViewThrough": [ + "AdsCheckoutInitiatedConversionViewThrough", + ], + "IsSignUpConversionViewThrough": [ + "AdsSignUpConversionViewThrough", + ], + "IsSiteVisitConversionViewThrough": [ + "AdsSiteVisitConversionViewThrough", + ], + "IsSessionConversionViewThrough": [ + "AdsSessionConversionViewThrough", + ], + "IsLandingPageViewConversionViewThrough": [ + "AdsLandingPageViewConversionViewThrough", + ], + "IsUpperFunnelConversionViewThrough": [ + "AdsUpperFunnelConversionViewThrough", + ], +} + +ads_slim_engagement_to_action_types = { + "IsOpenLink": [ + "ClientTweetOpenLink", + ], + "IsExternalLinkLongDwelled": [ + "ClientTweetExternalLinkLongDwelled", + ], + "IsVideoQualityViewed": [ + "ClientTweetVideoQualityView", + ], + "IsNotInterestedIn": [ + "ClientTweetNotInterestedIn", + ], + "IsBlockAuthor": [ + "ClientTweetBlockAuthor", + ], + "IsReported": [ + "ClientTweetReport", + ], + "IsMuteAuthor": [ + "ClientTweetMuteAuthor", + ], + **{ + eng: actions + for eng, actions in ads_p_conv_click_engagement_to_action_types.items() + if eng not in primary_engagement_to_action_types + }, } @@ -243,8 +348,10 @@ def to_pascal_case(s): "default": primary_engagement_to_action_types, "all": engagement_to_action_types, "notifications": notification_engagement_to_action_types, + "search": search_engagement_to_action_types, "ads_conversion": ads_conversion_engagement_to_action_types, "ads_p_conv_click": ads_p_conv_click_engagement_to_action_types, + "ads_slim": ads_slim_engagement_to_action_types, "none": {}, } @@ -269,8 +376,52 @@ def engagement_to_ids(metric_group): recsys_pb2.ActionName.ADS_SIGN_UP_CONVERSION, recsys_pb2.ActionName.ADS_CHECKOUT_INITIATED_CONVERSION, recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_CLICK_INSTALL, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_PURCHASE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_ADD_TO_CART, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_LEVEL_ACHIEVED, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_TUTORIAL_COMPLETE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_SIGN_UP, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_CUSTOM, +] + +MACT_IN_APP_LOSS_ACTION_INDICES = [ + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_PURCHASE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_ADD_TO_CART, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_LEVEL_ACHIEVED, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_TUTORIAL_COMPLETE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_SIGN_UP, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_CUSTOM, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_PURCHASE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_ADD_TO_CART, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_LEVEL_ACHIEVED, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_TUTORIAL_COMPLETE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_SIGN_UP, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_CUSTOM, ] +VIEW_THROUGH_ACTION_INDICES = [ + recsys_pb2.ActionName.ADS_ATTRIBUTED_KEY_VIEW_CONVERSION, + recsys_pb2.ActionName.ADS_ATTRIBUTED_VIEW_CONVERSION, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_PURCHASE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_ADD_TO_CART, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_LEVEL_ACHIEVED, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_TUTORIAL_COMPLETE, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_SIGN_UP, + recsys_pb2.ActionName.ADS_ATTRIBUTED_MACT_VIEW_CUSTOM, + recsys_pb2.ActionName.ADS_PURCHASE_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_ADD_TO_CART_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_CHECKOUT_INITIATED_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_SIGN_UP_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_SITE_VISIT_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_SESSION_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_LANDING_PAGE_VIEW_CONVERSION_VIEW_THROUGH, + recsys_pb2.ActionName.ADS_UPPER_FUNNEL_CONVERSION_VIEW_THROUGH, +] + +SOURCE_SPLIT_CONVERSION_HEAD_INDICES = ( + CLICK_CONDITIONED_ACTION_INDICES + VIEW_THROUGH_ACTION_INDICES +) + NEGATIVE_FEEDBACK_HEAD_INDICES = [ recsys_pb2.ActionName.CLIENT_TWEET_REPORT, recsys_pb2.ActionName.CLIENT_TWEET_NOT_INTERESTED_IN, @@ -280,4 +431,5 @@ def engagement_to_ids(metric_group): recsys_pb2.ActionName.CLIENT_TWEET_NOT_RELEVANT, recsys_pb2.ActionName.CLIENT_TWEET_MUTE_CONVERSATION, recsys_pb2.ActionName.CLIENT_NOTIFICATION_SEE_LESS_OFTEN, + recsys_pb2.ActionName.CLIENT_TWEET_SEARCH_QUERY_REFORMULATED, ] diff --git a/phoenix/xrex/data/recsys/feature_config.py b/phoenix/xrex/data/recsys/feature_config.py index dd4d7bb3..1c42d1f0 100644 --- a/phoenix/xrex/data/recsys/feature_config.py +++ b/phoenix/xrex/data/recsys/feature_config.py @@ -33,11 +33,15 @@ class CategoricalFeature(enum.IntEnum): } ) +COMPUTED_BOOL_FEATURE_NAMES: frozenset[str] = frozenset({"isStalePost14d"}) + AUTHOR_NSFW_BIT = 2 class BoolFeature(enum.IntEnum): - pass + isStalePost14d = 0 + isAuthorFollowedByViewerSeq = 1 + isAuthorFollowingViewerSeq = 2 class FloatFeature(enum.IntEnum): @@ -53,6 +57,21 @@ class Int64Feature(enum.IntEnum): quoteCountSeq = 8 viewCountSeq = 9 ipAddressSeq = 10 + firstDpaProductKey = 11 + firstDpaProductKeyHash2 = 12 + + +COMPUTED_INT64_FEATURE_NAMES: frozenset[str] = frozenset({"firstDpaProductKeyHash2"}) + + +ADS_PRODUCT_KEY_TABLE_SIZE = 10_000_000 +ADS_PRODUCT_KEY_HASH_SCALE = 1_566_083_941 +ADS_PRODUCT_KEY_HASH_BIAS = 774_047_449 +ADS_PRODUCT_KEY_HASH_SCALE_2 = 1_204_318_477 +ADS_PRODUCT_KEY_HASH_BIAS_2 = 393_342_739 +ADS_PRODUCT_KEY_HASH_MODULUS = 2_147_483_647 + +STALE_POST_14D_TTL_SEC = 1_213_200 CATEGORICAL_FEATURES: list[str] = [f.name for f in CategoricalFeature] @@ -115,7 +134,27 @@ class UserInt64Feature(enum.IntEnum): ] +OPTIONAL_COLUMNS: list[str] = [ + "searchQueryEmbeddingSeq", + "lineItemObjectiveSeq", + "safetyLabelMaskSeq", + "semanticIdSeq", + "sampleWeight", + "firstDpaProductKey", + "authorFollowerCountSeq", + "inReplyToPostIdSeq", + "is_delayed_feedback", + "conversionKeepMask", +] + + def _build_required_columns() -> list[str]: + excluded_names = ( + COMPUTED_CATEGORICAL_FEATURE_NAMES + | COMPUTED_INT64_FEATURE_NAMES + | COMPUTED_BOOL_FEATURE_NAMES + | set(OPTIONAL_COLUMNS) + ) seen: set[str] = set() result: list[str] = [] for col in BASE_REQUIRED_COLUMNS: @@ -133,20 +172,10 @@ def _build_required_columns() -> list[str]: USER_INT64_FEATURES, ): for col in feature_list: - if col not in seen and col not in COMPUTED_CATEGORICAL_FEATURE_NAMES: + if col not in seen and col not in excluded_names: result.append(col) seen.add(col) return result REQUIRED_COLUMNS: list[str] = _build_required_columns() - -OPTIONAL_COLUMNS: list[str] = [ - "searchQueryEmbeddingSeq", - "lineItemObjectiveSeq", - "safetyLabelMaskSeq", - "semanticIdSeq", - "sampleWeight", - "authorFollowerCountSeq", - "inReplyToPostIdSeq", -] diff --git a/phoenix/xrex/data/recsys/recsys_batch.py b/phoenix/xrex/data/recsys/recsys_batch.py index 6a66a409..efb2106a 100644 --- a/phoenix/xrex/data/recsys/recsys_batch.py +++ b/phoenix/xrex/data/recsys/recsys_batch.py @@ -16,11 +16,18 @@ from xrex.data.recsys.constants import action_type_map from xrex.data.recsys.feature_config import ( + ADS_PRODUCT_KEY_HASH_BIAS, + ADS_PRODUCT_KEY_HASH_BIAS_2, + ADS_PRODUCT_KEY_HASH_MODULUS, + ADS_PRODUCT_KEY_HASH_SCALE, + ADS_PRODUCT_KEY_HASH_SCALE_2, + ADS_PRODUCT_KEY_TABLE_SIZE, AUTHOR_NSFW_BIT, BOOL_FEATURES, CATEGORICAL_FEATURES, FLOAT_FEATURES, INT64_FEATURES, + STALE_POST_14D_TTL_SEC, BoolFeature, CategoricalFeature, FloatFeature, @@ -174,7 +181,10 @@ def _col( batch_size: int, t: type[_T], ) -> npt.NDArray[_T]: - arr = (rb.column(col).values.to_numpy(zero_copy_only=False).reshape(batch_size, -1)).astype(t) + values = rb.column(col).values + if values.null_count: + values = values.fill_null(False if pa.types.is_boolean(values.type) else 0) + arr = values.to_numpy(zero_copy_only=False).reshape(batch_size, -1).astype(t) return typing.cast(npt.NDArray[_T], arr) @@ -203,6 +213,7 @@ class PostSeq(TypedDict): continuous_actions: npt.NDArray[np.float32] promoted_ids: npt.NDArray[np.int64] | None line_item_objective: npt.NDArray[np.int16] | None + trained_candidate_mask: NotRequired[npt.NDArray[np.bool_] | None] safety_label_mask: npt.NDArray[np.int64] | None embedding: npt.NDArray[np.float32] | jax.Array | None search_query_embeddings: npt.NDArray[np.float32] | None @@ -296,6 +307,7 @@ class RecsysFeaturesBatch(TypedDict): user_installed_apps_multihot: npt.NDArray[np.bool_] num_positive_candidates: npt.NDArray[np.int32] | None sample_weights: NotRequired[npt.NDArray[np.float32] | None] + sample_source: NotRequired[npt.NDArray[np.bool_] | None] packing_layout: NotRequired[SequencePackedLayout | None] @@ -347,6 +359,7 @@ def from_record_batch( global_post_sids: np.ndarray | None = None, sid_num_levels: int = 0, compute_post_unexplored_label: bool = False, + zero_stale_post_14d_candidate_counts: bool = False, ) -> RecsysFeaturesBatch: start = time.time() batch_size = record_batch.num_rows @@ -396,6 +409,14 @@ def from_record_batch( else: seq_len = actions.shape[1] client_app_id = np.zeros((batch_size, seq_len), dtype=np.int32) + if "conversionKeepMask" in record_batch.schema.names: + conversion_keep = _col(record_batch, "conversionKeepMask", batch_size, np.bool_) + null_rows = ~record_batch.column("conversionKeepMask").is_valid().to_numpy( + zero_copy_only=False + ) + conversion_keep[null_rows] = True + else: + conversion_keep = None post_creation_ts_sec = (((tweet_ids >> 22) + TWITTER_EPOCH_MS) // 1000).astype(np.int32) post_creation_ts_sec = np.where(tweet_ids == 0, 0, post_creation_ts_sec) @@ -560,6 +581,7 @@ def from_record_batch( candidate_impr_ts = np.zeros(cand_shape_2d, dtype=np.int32) candidate_product_surface = np.zeros(cand_shape_2d, dtype=np.int32) candidate_client_app_id = np.zeros(cand_shape_2d, dtype=np.int32) + candidate_trained_mask = np.ones(cand_shape_2d, dtype=np.bool_) candidate_post_creation_ts_sec = np.zeros(cand_shape_2d, dtype=np.int32) candidate_actions = np.zeros(cand_shape_3d, dtype=actions.dtype) candidate_continuous_actions = np.zeros( @@ -719,6 +741,8 @@ def from_record_batch( candidate_impr_ts[*cslice] = ts_values.astype(np.int32) candidate_product_surface[*cslice] = product_surface[*dslice] candidate_client_app_id[*cslice] = client_app_id[*dslice] + if conversion_keep is not None: + candidate_trained_mask[*cslice] = conversion_keep[*dslice] candidate_post_creation_ts_sec[*cslice] = post_creation_ts_sec[*dslice] candidate_actions[*cslice] = actions[*dslice, :] candidate_continuous_actions[*cslice] = continuous_actions[*dslice, :] @@ -832,6 +856,41 @@ def from_record_batch( out=cand_int64_features[:, :, _ec_idx], ) + _dpa_idx = Int64Feature.firstDpaProductKey.value + + def _hash_dpa_keys(keys: npt.NDArray[np.int64], scale: int, bias: int) -> npt.NDArray: + hashed = ( + (keys * scale + bias) % ADS_PRODUCT_KEY_HASH_MODULUS % (ADS_PRODUCT_KEY_TABLE_SIZE - 1) + ) + 1 + return np.where(keys == 0, 0, hashed) + + _dpa_idx2 = Int64Feature.firstDpaProductKeyHash2.value + for _dpa_arr in (hist_int64_features, cand_int64_features): + _dpa_keys = _dpa_arr[:, :, _dpa_idx].copy() + _dpa_arr[:, :, _dpa_idx2] = _hash_dpa_keys( + _dpa_keys, ADS_PRODUCT_KEY_HASH_SCALE_2, ADS_PRODUCT_KEY_HASH_BIAS_2 + ) + _dpa_arr[:, :, _dpa_idx] = _hash_dpa_keys( + _dpa_keys, ADS_PRODUCT_KEY_HASH_SCALE, ADS_PRODUCT_KEY_HASH_BIAS + ) + + if zero_stale_post_14d_candidate_counts: + ttl_sec = np.int64(STALE_POST_14D_TTL_SEC) + original_age_sec = candidate_impr_ts.astype( + np.int64 + ) - candidate_post_creation_ts_sec.astype(np.int64) + stale_post_14d = ( + (candidate_impr_ts > 0) + & (candidate_post_creation_ts_sec > 0) + & (original_age_sec > ttl_sec) + ) + for _ec_idx in _ENGAGEMENT_COUNT_INT64_INDICES: + cand_int64_features[:, :, _ec_idx] = np.where( + stale_post_14d, 0, cand_int64_features[:, :, _ec_idx] + ) + if cand_bool_features.shape[2] > BoolFeature.isStalePost14d.value: + cand_bool_features[:, :, BoolFeature.isStalePost14d.value] = stale_post_14d + history_seq = PostSeq( impr_ts=history_impr_ts, actions=history_actions, @@ -863,6 +922,7 @@ def from_record_batch( product_surface=candidate_product_surface, client_app_id=candidate_client_app_id, post_ids=candidate_post_ids if include_candidate_post_ids else None, + trained_candidate_mask=candidate_trained_mask, continuous_actions=candidate_continuous_actions, promoted_ids=candidate_promoted_ids, line_item_objective=candidate_line_item_objective, @@ -920,6 +980,15 @@ def from_record_batch( if "sampleWeight" in record_batch.schema.names else np.ones((batch_size, 1), dtype=np.float32) ), + sample_source=( + record_batch.column("is_delayed_feedback") + .fill_null(False) + .to_numpy(zero_copy_only=False) + .astype(np.bool_) + .reshape(-1, 1) + if "is_delayed_feedback" in record_batch.schema.names + else np.zeros((batch_size, 1), dtype=np.bool_) + ), ) batch["num_positive_candidates"] = ( num_positive_per_user.reshape(-1, 1) @@ -948,6 +1017,7 @@ def apply_negative_sampling( post_ids = post_seq["post_ids"] product_surface = post_seq["product_surface"] client_app_id = post_seq["client_app_id"] + trained_candidate_mask = post_seq.get("trained_candidate_mask") post_creation_ts_sec = post_seq["post_creation_ts_sec"] continuous_actions = post_seq["continuous_actions"] promoted_ids = post_seq["promoted_ids"] @@ -997,6 +1067,7 @@ def apply_negative_sampling( ) new_product_surface = np.zeros((batch_size, total_candidate_slots), dtype=product_surface.dtype) new_client_app_id = np.zeros((batch_size, total_candidate_slots), dtype=client_app_id.dtype) + new_trained_mask = np.ones((batch_size, total_candidate_slots), dtype=np.bool_) new_post_creation_ts_sec = np.zeros( (batch_size, total_candidate_slots), dtype=post_creation_ts_sec.dtype ) @@ -1027,6 +1098,8 @@ def apply_negative_sampling( new_ip_hashes[:, positive_slice, :] = ip_hashes new_product_surface[:, positive_slice] = product_surface new_client_app_id[:, positive_slice] = client_app_id + if trained_candidate_mask is not None: + new_trained_mask[:, positive_slice] = trained_candidate_mask new_post_creation_ts_sec[:, positive_slice] = post_creation_ts_sec if new_post_sids is not None and _post_sids is not None: new_post_sids[:, positive_slice, :] = _post_sids @@ -1193,6 +1266,7 @@ def _copy_neg_features(curr_user_idx, start_slot, end_slot, post_src, query_src) product_surface=new_product_surface, client_app_id=new_client_app_id, post_ids=new_post_ids, + trained_candidate_mask=new_trained_mask, continuous_actions=new_continuous_actions, promoted_ids=new_promoted_ids, line_item_objective=new_line_item_objective, @@ -1264,6 +1338,8 @@ def apply_global_negative_sampling( (batch_size, expanded_candidate_slots), dtype=client_app_id.dtype, ) + _gn_trained_mask = post_seq.get("trained_candidate_mask") + new_gn_trained_mask = np.ones((batch_size, expanded_candidate_slots), dtype=np.bool_) new_post_creation_ts_sec = np.zeros( (batch_size, expanded_candidate_slots), dtype=post_creation_ts_sec.dtype, @@ -1297,6 +1373,8 @@ def apply_global_negative_sampling( new_ip_hashes[:, original_slice, :] = ip_hashes new_product_surface[:, original_slice] = product_surface new_client_app_id[:, original_slice] = client_app_id + if _gn_trained_mask is not None: + new_gn_trained_mask[:, original_slice] = _gn_trained_mask new_post_creation_ts_sec[:, original_slice] = post_creation_ts_sec new_continuous_actions[:, original_slice, :] = continuous_actions if categorical_features.shape[2] > 0: @@ -1420,6 +1498,7 @@ def apply_global_negative_sampling( product_surface=new_product_surface, client_app_id=new_client_app_id, post_ids=new_post_ids, + trained_candidate_mask=new_gn_trained_mask, continuous_actions=new_continuous_actions, promoted_ids=new_promoted_ids, line_item_objective=new_line_item_objective, diff --git a/phoenix/xrex/data/recsys/sequence_packing.py b/phoenix/xrex/data/recsys/sequence_packing.py index b868c896..9c1eb6e0 100644 --- a/phoenix/xrex/data/recsys/sequence_packing.py +++ b/phoenix/xrex/data/recsys/sequence_packing.py @@ -166,6 +166,11 @@ def pack_batch( rng: np.random.Generator | None, block_size: int = 128, ) -> RecsysFeaturesBatch: + _tcm = batch["candidate_seq"].get("trained_candidate_mask") + assert _tcm is None or bool(np.asarray(_tcm).all()), ( + "trained_candidate_mask with masked candidates is not supported with sequence packing" + ) + read_bsz_per_process = batch["user_hashes"].shape[0] history_seq_len = batch["history_seq"]["post_hashes"].shape[1] candidate_seq_len = batch["candidate_seq"]["post_hashes"].shape[1] @@ -252,8 +257,13 @@ def _make_batch( else None ), sample_weights=( - batch["sample_weights"].reshape(D, B, *batch["sample_weights"].shape[1:]) - if batch.get("sample_weights") is not None + sw.reshape(D, B, *sw.shape[1:]) + if (sw := batch.get("sample_weights")) is not None + else None + ), + sample_source=( + ss.reshape(D, B, *ss.shape[1:]) + if (ss := batch.get("sample_source")) is not None else None ), packing_layout=layout, diff --git a/phoenix/xrex/data/retrieval_dataset.py b/phoenix/xrex/data/retrieval_dataset.py index 6a4af404..1e381dce 100644 --- a/phoenix/xrex/data/retrieval_dataset.py +++ b/phoenix/xrex/data/retrieval_dataset.py @@ -18,7 +18,7 @@ _O2_SCHEME = "o2://" -_O2_DEFAULT_ENDPOINT = settings.OBJECT_STORE_ENDPOINT +O2_DEFAULT_ENDPOINT = settings.OBJECT_STORE_ENDPOINT _O2_CREDENTIALS_DIR = Path(settings.OBJECT_STORE_CREDENTIALS_DIR) _GCS_CREDENTIALS_DIR = Path(settings.GCS_CREDENTIALS_DIR) @@ -30,12 +30,12 @@ PHOENIX_INDEX_BASE = Path(settings.PHOENIX_INDEX_BASE) -def _bridge_ads_s3_env() -> None: - if not settings.ADS_S3_ENV_PREFIX: +def _bridge_o2_env() -> None: + if not settings.O2_ENV_PREFIX: return for suffix in ("ENDPOINT", "ACCESS_KEY", "SECRET_KEY"): - prefixed = f"{settings.ADS_S3_ENV_PREFIX}ADS_S3_{suffix}" - generic = f"ADS_S3_{suffix}" + prefixed = f"{settings.O2_ENV_PREFIX}O2_PROD_{suffix}" + generic = f"O2_PROD_{suffix}" if os.environ.get(generic): continue val = os.environ.get(prefixed) @@ -47,11 +47,20 @@ def _bridge_ads_s3_env() -> None: os.environ[generic] = val -_bridge_ads_s3_env() +_bridge_o2_env() + + +_SID_SNAPSHOTS = "post_sid_v8_256x6_snapshots" + + +def _sid_window(filename: str) -> tuple[str, str]: + return ( + str(PHOENIX_INDEX_BASE / _SID_SNAPSHOTS / filename), + str(PHOENIX_INDEX_BASE / f"{_SID_SNAPSHOTS}_backup" / filename), + ) def _idx(sub: str) -> str: - sub = sub.replace("post_sid_v5_256x6_snapshots", "post_sid_v8_256x6_snapshots") return str(PHOENIX_INDEX_BASE / sub) @@ -96,11 +105,15 @@ def _parse_snapshot_timestamp(key: str) -> int | None: def _download_from_o2(name: str, bucket: str, prefix: str) -> tuple[bytes, int, str]: import boto3 + endpoint = os.environ.get("O2_PROD_ENDPOINT", O2_DEFAULT_ENDPOINT) + access_key = _read_credential("O2_PROD_ACCESS_KEY") + secret_key = _read_credential("O2_PROD_SECRET_KEY") + logger.info("%s: O2 endpoint: %s", name, endpoint) s3 = boto3.client( "s3", - endpoint_url=os.environ.get("ADS_S3_ENDPOINT", _O2_DEFAULT_ENDPOINT), - aws_access_key_id=_read_credential("ADS_S3_ACCESS_KEY"), - aws_secret_access_key=_read_credential("ADS_S3_SECRET_KEY"), + endpoint_url=endpoint, + aws_access_key_id=access_key, + aws_secret_access_key=secret_key, ) objects = [ (obj["Key"], obj["LastModified"].timestamp()) @@ -228,16 +241,8 @@ def _load_from_o2( class RetrievalDataset(Enum): PAD = (0, None, None) - HOME = ( - 1, - _idx("post_sid_v5_256x6_snapshots/1fav_1day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/1fav_1day.parquet"), - ) - IMMERSIVE2Day = ( - 2, - _idx("post_sid_v5_256x6_snapshots/video_2day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/video_2day.parquet"), - ) + HOME = (1, *_sid_window("1fav_1day.parquet")) + IMMERSIVE2Day = (2, *_sid_window("video_2day.parquet")) RELEVANT_ADS = ( 3, _idx("relevant_ads/v6/post_id_author_id_pair.parquet"), @@ -248,31 +253,15 @@ class RetrievalDataset(Enum): _idx("relevant_ads/carousel/post_id_author_id_pair.parquet"), _idx("relevant_ads/carousel/post_id_author_id_pair.parquet"), ) - EVERGREEN = ( - 5, - _idx("post_sid_v5_256x6_snapshots/evergreen_video_1825day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/evergreen_video_1825day.parquet"), - ) - IMMERSIVENSFW = ( - 6, - _idx("post_sid_v5_256x6_snapshots/nsfw_video_2day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/nsfw_video_2day.parquet"), - ) + EVERGREEN = (5, *_sid_window("video_4to14day.parquet")) + IMMERSIVENSFW = (6, *_sid_window("nsfw_video_2day.parquet")) ACTIVE_ADS = ( 7, settings.ADS_INDEX_URI, settings.ADS_INDEX_URI, ) - IMMERSIVE4Day = ( - 8, - _idx("post_sid_v5_256x6_snapshots/video_4day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/video_4day.parquet"), - ) - IMAGINE = ( - 9, - _idx("post_sid_v5_256x6_snapshots/imagine_4day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/imagine_4day.parquet"), - ) + IMMERSIVE4Day = (8, *_sid_window("video_4day.parquet")) + IMAGINE = (9, *_sid_window("imagine_4day.parquet")) TAIL = ( 10, _idx("post_sid_v5_256x6_tail_snapshots/tail_1day.parquet"), diff --git a/phoenix/xrex/data/rust_kafka_recsys.py b/phoenix/xrex/data/rust_kafka_recsys.py index 2aa5ccbd..d5832110 100644 --- a/phoenix/xrex/data/rust_kafka_recsys.py +++ b/phoenix/xrex/data/rust_kafka_recsys.py @@ -15,7 +15,9 @@ from xrex.data.streaming.kafkaconsumer import ConsumerMode, _check_reset_sentinel from xrex.data.streaming.kafkaloader import ( PhoenixKafkaDataset, + _rearm_catchup_request, _resolve_sasl_password, + _take_catchup_request, ) @@ -135,7 +137,7 @@ def _async_kafka_consumer_arrow( RecordBatchProvider = rust_ext.load("xai_recsys_kafka_reader").RecordBatchProvider provider: RecordBatchProvider | None = None - reset_watcher: threading.Thread | None = None + control_watcher: threading.Thread | None = None try: sasl_password = _resolve_sasl_password(self.bootstrap_servers) @@ -178,7 +180,7 @@ def _async_kafka_consumer_arrow( ) in_rebaseline = threading.Event() - reset_watcher = self._start_reset_watcher( + control_watcher = self._start_control_watcher( provider, stop_event, reset_event, in_rebaseline ) @@ -240,8 +242,8 @@ def _async_kafka_consumer_arrow( line += f" | {drop_status}" rank_logger.info("%s", line) finally: - if reset_watcher is not None: - reset_watcher.join(timeout=2.0) + if control_watcher is not None: + control_watcher.join(timeout=2.0) if reload_thread is not None: reload_thread.join(timeout=5.0) @@ -295,7 +297,7 @@ def reload_loop() -> None: thread.start() return thread - def _start_reset_watcher( + def _start_control_watcher( self, provider, stop_event: threading.Event | None, @@ -312,6 +314,25 @@ def watcher() -> None: else: http_hit = False time.sleep(1.0) + catchup_request = _take_catchup_request() + if catchup_request is not None: + catchup_generation, catchup_window_secs = catchup_request + try: + deadline_epoch_secs = provider.request_catchup(catchup_window_secs) + rank_logger.info( + "RustKafkaDataset: live catch-up via HTTP endpoint — " + "temporarily engaging drop mode with T=%.0fs " + "(request expires at epoch %.0f)", + catchup_window_secs, + deadline_epoch_secs, + ) + except Exception as e: + _rearm_catchup_request(catchup_generation) + rank_logger.error( + "Failed to forward live catch-up request to Rust " + "reader (will retry next tick): %s", + e, + ) sentinel_mtime = _check_reset_sentinel(last_reset_ts) if sentinel_mtime is not None: last_reset_ts = sentinel_mtime @@ -338,7 +359,7 @@ def watcher() -> None: e, ) - thread = threading.Thread(target=watcher, name="rust-kafka-reset-watcher", daemon=True) + thread = threading.Thread(target=watcher, name="rust-kafka-control-watcher", daemon=True) thread.start() return thread diff --git a/phoenix/xrex/data/rust_parquet_recsys.py b/phoenix/xrex/data/rust_parquet_recsys.py index 4375e48a..718fcf70 100644 --- a/phoenix/xrex/data/rust_parquet_recsys.py +++ b/phoenix/xrex/data/rust_parquet_recsys.py @@ -216,6 +216,7 @@ def producer() -> None: search_query_embedding_dim=dataset.search_query_embedding_dim, sid_num_levels=(dataset.sid_num_levels if dataset.use_post_sid else 0), compute_post_unexplored_label=dataset.compute_post_unexplored_label, + zero_stale_post_14d_candidate_counts=dataset.enable_stale_post, ) if record_batch.num_rows < batch_size: diff --git a/phoenix/xrex/data/streaming/kafkaloader.py b/phoenix/xrex/data/streaming/kafkaloader.py index 6176d006..3443fe56 100644 --- a/phoenix/xrex/data/streaming/kafkaloader.py +++ b/phoenix/xrex/data/streaming/kafkaloader.py @@ -6,8 +6,10 @@ import logging import os import queue +import re import threading import time +import urllib.parse from functools import partial from threading import Event from typing import Any, Iterator, Mapping, Optional, cast @@ -67,6 +69,70 @@ _kafka_reset_event: Optional[threading.Event] = None +_kafka_catchup_event: threading.Event = threading.Event() +_kafka_catchup_window_secs: float = 7200.0 +_kafka_catchup_lock: threading.Lock = threading.Lock() +_kafka_catchup_generation: int = 0 + +_CATCHUP_MIN_SECS = 60.0 +_CATCHUP_MAX_SECS = 24.0 * 3600.0 +_CATCHUP_DEFAULT_SECS = 7200.0 + +_CATCHUP_DURATION_RE = re.compile( + r"^\s*(\d+(?:\.\d+)?)\s*(s|sec|secs|m|min|mins|h|hr|hrs|hour|hours)?\s*$", + re.IGNORECASE, +) + + +def _parse_catchup_duration(raw: str | None) -> float: + if raw is None or not raw.strip(): + return _CATCHUP_DEFAULT_SECS + m = _CATCHUP_DURATION_RE.match(raw) + if m is None: + raise ValueError( + f"unparseable duration {raw!r} (examples: '30min', '2hr', '90m', " + "'7200s'; a bare number means minutes)" + ) + value = float(m.group(1)) + unit = (m.group(2) or "m").lower() + if unit.startswith("s"): + secs = value + elif unit.startswith("m"): + secs = value * 60.0 + else: + secs = value * 3600.0 + if not _CATCHUP_MIN_SECS <= secs <= _CATCHUP_MAX_SECS: + raise ValueError( + f"duration {raw!r} = {secs:.0f}s is outside " + f"[{_CATCHUP_MIN_SECS:.0f}s, {_CATCHUP_MAX_SECS:.0f}s] " + "(below 1 min use /reset-kafka instead; above 24 h dropping " + "is not worth it)" + ) + return secs + + +def _post_catchup_request(window_secs: float) -> None: + global _kafka_catchup_window_secs, _kafka_catchup_generation + with _kafka_catchup_lock: + _kafka_catchup_window_secs = window_secs + _kafka_catchup_generation += 1 + _kafka_catchup_event.set() + + +def _take_catchup_request() -> tuple[int, float] | None: + with _kafka_catchup_lock: + if _kafka_catchup_event.is_set(): + _kafka_catchup_event.clear() + return _kafka_catchup_generation, _kafka_catchup_window_secs + return None + + +def _rearm_catchup_request(generation: int) -> None: + with _kafka_catchup_lock: + if _kafka_catchup_generation == generation: + _kafka_catchup_event.set() + + _meter = None _kafka_batches_processed = None _kafka_messages_processed = None @@ -95,34 +161,72 @@ def _start_reset_http_server(base_port: int, num_ports: int) -> int | None: from http.server import BaseHTTPRequestHandler, HTTPServer class _ResetHandler(BaseHTTPRequestHandler): + timeout = 10 + + def _respond(self, status: int, body: str) -> None: + self.send_response(status) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(body.encode()) + def do_POST(self): - if self.path == "/reset-kafka": + parsed = urllib.parse.urlsplit(self.path) + if parsed.path == "/reset-kafka": if _kafka_reset_event is not None: _kafka_reset_event.set() - self.send_response(200) - self.send_header("Content-Type", "text/plain") - self.end_headers() - self.wfile.write(b"Kafka reset signal sent\n") + self._respond(200, "Kafka reset signal sent\n") rank_logger.info( "Kafka reset signal received via HTTP POST /reset-kafka on port %d", self.server.server_address[1], ) else: - self.send_response(503) - self.send_header("Content-Type", "text/plain") - self.end_headers() - self.wfile.write(b"Consumer not started yet\n") + self._respond(503, "Consumer not started yet\n") + elif parsed.path == "/catchup-kafka": + if _kafka_reset_event is None: + self._respond(503, "Consumer not started yet\n") + return + raw = (urllib.parse.parse_qs(parsed.query).get("duration") or [None])[0] + if raw is None: + try: + length = int(self.headers.get("Content-Length") or 0) + except ValueError: + self._respond(400, "Invalid Content-Length\n") + return + if length > 0: + raw = self.rfile.read(min(length, 64)).decode("utf-8", "replace") + try: + window_secs = _parse_catchup_duration(raw) + except ValueError as e: + self._respond(400, f"{e}\n") + return + _post_catchup_request(window_secs) + self._respond( + 200, + f"Kafka catch-up signal sent: window={window_secs:.0f}s " + "(acted on by the Rust consumer path only)\n", + ) + rank_logger.info( + "Kafka catch-up signal received via HTTP POST /catchup-kafka " + "on port %d (window=%.0fs)", + self.server.server_address[1], + window_secs, + ) else: self.send_response(404) self.end_headers() def do_GET(self): - if self.path == "/reset-kafka": + parsed = urllib.parse.urlsplit(self.path) + if parsed.path == "/reset-kafka": is_pending = _kafka_reset_event is not None and _kafka_reset_event.is_set() - self.send_response(200) - self.send_header("Content-Type", "text/plain") - self.end_headers() - self.wfile.write(f"reset_pending={is_pending}\n".encode()) + self._respond(200, f"reset_pending={is_pending}\n") + elif parsed.path == "/catchup-kafka": + self._respond( + 200, + f"catchup_pending={_kafka_catchup_event.is_set()} " + f"window_secs={_kafka_catchup_window_secs:.0f} " + "(consumed by the Rust consumer path only)\n", + ) else: self.send_response(404) self.end_headers() @@ -140,10 +244,15 @@ def log_message(self, format, *args): return None t = threading.Thread(target=server.serve_forever, daemon=True) t.start() - rank_logger.info("Kafka reset HTTP server listening on port %d (POST /reset-kafka)", port) + rank_logger.info( + "Kafka control HTTP server listening on port %d " + "(POST /reset-kafka, POST /catchup-kafka)", + port, + ) return port rank_logger.warning( - "Kafka reset HTTP server: no free port in %d-%d; live /reset-kafka disabled for this worker", + "Kafka control HTTP server: no free port in %d-%d; live /reset-kafka " + "and /catchup-kafka disabled for this worker", base_port, base_port + num_ports - 1, ) @@ -650,7 +759,8 @@ def _unify_record_batch_schemas( raise new_columns.append(col) else: - logger.warning(f"Column {name} not found in batch") + if name in required_columns: + logger.warning(f"Column {name} not found in batch") new_columns.append(pa.nulls(num_rows, type=target_type)) result.append(pa.RecordBatch.from_arrays(new_columns, names=col_names)) @@ -757,17 +867,23 @@ class PhoenixKafkaDataset(PhoenixDataset): compute_post_unexplored_label: bool = False + exclude_required_columns: str = "" + def kafka_to_training_batch( self, batch: list[pa.RecordBatch], batch_size: int, shard_index: int = 0 ) -> RecsysFeaturesBatch: del batch_size start_time = time.time() + required_columns: list[str] = list(REQUIRED_COLUMNS) + if self.exclude_required_columns: + excluded = {c.strip() for c in self.exclude_required_columns.split(",") if c.strip()} + required_columns = [c for c in required_columns if c not in excluded] + if not self._logged_schema and batch: first_rb = batch[0] available_cols = first_rb.schema.names - expected_cols = REQUIRED_COLUMNS - missing_cols = [col for col in expected_cols if col not in available_cols] + missing_cols = [col for col in required_columns if col not in available_cols] if missing_cols: rank_logger.error( f"Schema mismatch! Missing required columns: {missing_cols}. " @@ -777,7 +893,6 @@ def kafka_to_training_batch( rank_logger.info(f"Schema validated. Available columns: {available_cols}") self._logged_schema = True - required_columns: list[str] = list(REQUIRED_COLUMNS) if self.multimodal_embedding_type is not None: embedding_col_name, _ = EMBEDDING_CONFIG[self.multimodal_embedding_type] if batch and embedding_col_name in batch[0].schema.names: @@ -825,6 +940,7 @@ def kafka_to_training_batch( global_post_sids=self.global_post_sids, sid_num_levels=self.sid_num_levels if self.use_post_sid else 0, compute_post_unexplored_label=self.compute_post_unexplored_label, + zero_stale_post_14d_candidate_counts=self.enable_stale_post, ) elapsed = time.time() - t diff --git a/phoenix/xrex/driver/config_factory.py b/phoenix/xrex/driver/config_factory.py index cb9a7f29..c1070b9d 100644 --- a/phoenix/xrex/driver/config_factory.py +++ b/phoenix/xrex/driver/config_factory.py @@ -5,7 +5,7 @@ from xai_configlib import Config from xrex.configs.config_loader import replace_cli_subs -from xrex.driver.driver_local import LocalDriverConfig +from xrex.driver.local import LocalDriverConfig logger = logging.getLogger(__name__) diff --git a/phoenix/xrex/driver/driver.py b/phoenix/xrex/driver/core.py similarity index 100% rename from phoenix/xrex/driver/driver.py rename to phoenix/xrex/driver/core.py diff --git a/phoenix/xrex/driver/hooks.py b/phoenix/xrex/driver/hooks.py index 7e05f432..a99c2594 100644 --- a/phoenix/xrex/driver/hooks.py +++ b/phoenix/xrex/driver/hooks.py @@ -57,6 +57,10 @@ def _filter_wandb_metrics(metrics: Mapping[str, Any]) -> dict[str, Any]: return {k: v for k, v in metrics.items() if not _should_omit_wandb_metric(k)} +def _write_filtered_wandb_metrics(metrics: Mapping[str, Any]) -> None: + write_wandb_log(_filter_wandb_metrics(metrics)) + + _CLICKHOUSE_HOST = settings.CLICKHOUSE_HOST _CLICKHOUSE_RUN_URL = settings.CLICKHOUSE_RUN_URL @@ -405,10 +409,7 @@ def _do_init(wandb_id: str, resume: str) -> None: @log_elapsed_time def on_step(self, metrics: Mapping[str, Any]): self.future.result() - - metrics = _filter_wandb_metrics(metrics) - - self.future = self.threadpool.submit(write_wandb_log, metrics=metrics) + self.future = self.threadpool.submit(_write_filtered_wandb_metrics, metrics) def on_train_rollback(self): wandb = _import_wandb() @@ -584,8 +585,10 @@ class CheckpointMonitoringState(Enum): _should_keep_cache = {} -def _should_keep(checkpoint_dir: str, run_id: str, keep_every_n: int) -> bool: - cache_key = (checkpoint_dir, run_id, keep_every_n) +def _should_keep( + checkpoint_dir: str, run_id: str, keep_every_n: int, checkpoint_every_n: int = 0 +) -> bool: + cache_key = (checkpoint_dir, run_id, keep_every_n, checkpoint_every_n) if cache_key in _should_keep_cache: return _should_keep_cache[cache_key] @@ -595,13 +598,17 @@ def _should_keep(checkpoint_dir: str, run_id: str, keep_every_n: int) -> bool: checkpoint_path = Path(os.path.join(checkpoint_dir, run_id)) completion_data = read_metadata_file(checkpoint_path) - if completion_data is None or completion_data[1] is None: - logger.info(f"No checkpoint index found in {checkpoint_dir}. Keep it.") + if completion_data is None: return True - - global_ckpt_index = completion_data[1] - if keep_every_n > 0 and global_ckpt_index % keep_every_n == 0: - logger.info(f"Keep candidate (ckpt_index: {global_ckpt_index}): {checkpoint_path}") + if completion_data.step is not None: + every_n = checkpoint_every_n if checkpoint_every_n > 0 else 1 + ckpt_index = completion_data.step // every_n + else: + ckpt_index = completion_data.checkpoint_index + if ckpt_index is None: + return True + if ckpt_index % keep_every_n == 0: + logger.info(f"Keep candidate (ckpt_index: {ckpt_index}): {checkpoint_path}") _should_keep_cache[cache_key] = True return True @@ -631,7 +638,13 @@ def _rm_rf_batch(paths: list[str]): ) -def _checkpoint_cleaner(checkpoint_dir: str, keep_last_n: int, keep_every_n: int, run_id: str): +def _checkpoint_cleaner( + checkpoint_dir: str, + keep_last_n: int, + keep_every_n: int, + run_id: str, + checkpoint_every_n: int = 0, +): if not os.path.exists(checkpoint_dir): logger.debug(f"Checkpoint {checkpoint_dir} does not exist.") return @@ -663,7 +676,7 @@ def _checkpoint_cleaner(checkpoint_dir: str, keep_last_n: int, keep_every_n: int for i, candidate in enumerate(entries): if i >= len(entries) - keep_last_n: break - if _should_keep(candidate.path, run_id, keep_every_n): + if _should_keep(candidate.path, run_id, keep_every_n, checkpoint_every_n): continue to_remove.append(candidate) @@ -729,6 +742,7 @@ def _checkpoint_monitoring(self, trainer_config: Jsonable, run_id: str): ) keep_last_n = ckpt_config.checkpoint_keep_last_n keep_every_n = ckpt_config.checkpoint_keep_every_nth + checkpoint_every_n = ckpt_config.checkpoint_every_n while True: try: @@ -739,7 +753,7 @@ def _checkpoint_monitoring(self, trainer_config: Jsonable, run_id: str): continue if signal == CheckpointMonitoringState.RUN_CLEANUP: - _checkpoint_cleaner(ckpt_dir, keep_last_n, keep_every_n, run_id) + _checkpoint_cleaner(ckpt_dir, keep_last_n, keep_every_n, run_id, checkpoint_every_n) elif signal == CheckpointMonitoringState.SHUTDOWN: logger.info("Checkpoint monitoring thread got shutdown signal. Shutting down.") return diff --git a/phoenix/xrex/driver/driver_local.py b/phoenix/xrex/driver/local.py similarity index 98% rename from phoenix/xrex/driver/driver_local.py rename to phoenix/xrex/driver/local.py index 29e9e309..931e1042 100644 --- a/phoenix/xrex/driver/driver_local.py +++ b/phoenix/xrex/driver/local.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Any, Optional -from xrex.driver.driver import ( +from xrex.driver.core import ( NUM_PHYSICAL_DEVICES_PER_NODE, DriverConfig, MultiRunDriverState, diff --git a/phoenix/xrex/eval/eval_utils.py b/phoenix/xrex/eval/eval_utils.py index b8ad117e..d6a809aa 100644 --- a/phoenix/xrex/eval/eval_utils.py +++ b/phoenix/xrex/eval/eval_utils.py @@ -3,7 +3,6 @@ import logging import traceback from collections.abc import Iterator -from dataclasses import dataclass from functools import partial from typing import Callable, Dict, List, Optional, Tuple @@ -39,7 +38,7 @@ def initialize(self): raise ValueError("All fields must be specified") -@dataclass +@configclass class EvaluationTaskNew(Config): max_steps: int = -1 eval_step_timeout: int = 300 diff --git a/phoenix/xrex/eval/metrics.py b/phoenix/xrex/eval/metrics.py index 07efb0a4..3cb592c9 100644 --- a/phoenix/xrex/eval/metrics.py +++ b/phoenix/xrex/eval/metrics.py @@ -8,7 +8,7 @@ from xai_configlib import Config, configclass -@configclass +@configclass(kw_only=True) class ForwardMetrics(Config): @abstractmethod def run_batch(self, logits, inputs, targets, mask=None): diff --git a/phoenix/xrex/inference/launch_inference.py b/phoenix/xrex/inference/launch_inference.py index a97a506a..bd109b5e 100644 --- a/phoenix/xrex/inference/launch_inference.py +++ b/phoenix/xrex/inference/launch_inference.py @@ -376,6 +376,7 @@ def run( runner.hotswap_download_max_concurrent = args.hotswap_download_max_concurrent runner.hotswap_stage_rate_limit_gbps = args.hotswap_stage_rate_limit_gib_per_s runner.hotswap_stage_chunk_mib = args.hotswap_stage_chunk_mib + runner.hotswap_malloc_trim = args.hotswap_malloc_trim if args.worker_id is not None: runner.worker_id = args.worker_id if args.num_workers is not None: @@ -383,8 +384,6 @@ def run( runner.log_rotate = args.log_rotate runner.log_rotate_max_bytes = args.log_rotate_max_bytes runner.log_rotate_backup_count = args.log_rotate_backup_count - if args.sid_endpoint is not None: - runner.sid_endpoint = args.sid_endpoint if hasattr(runner, "beam_width") and args.beam_width != 1: runner.beam_width = args.beam_width if hasattr(runner, "decode_levels") and args.decode_levels != 0: @@ -463,7 +462,16 @@ def run( "--copy_url", type=str, default="", - help="address of a weight source, typically trainer", + help="address of a weight source, typically trainer; an https:// prefix opts into TLS", + ) + parser.add_argument("--copy_tls_ca", type=str, default="", help="CA PEM verifying the server") + parser.add_argument("--copy_tls_cert", type=str, default="", help="client cert PEM (mTLS)") + parser.add_argument("--copy_tls_key", type=str, default="", help="client key PEM (mTLS)") + parser.add_argument( + "--copy_tls_server_name", + type=str, + default="", + help="hostname-verification override matching a server cert SAN", ) parser.add_argument( "--grpc_port", @@ -639,17 +647,17 @@ def run( parser.add_argument( "--use_pinned_d2h", type=str2bool, - default=False, + default=True, help="Use CUDA pinned host memory for D2H transfer (~50 GB/s vs JAX's ~3 GB/s). " - "Saves ~29ms/inference.", + "Saves ~29ms/inference. Default True; pass False to opt out.", ) parser.add_argument( "--pinned_d2h_num_buffers", type=int, default=3, - help="Number of pinned host buffers per output shape when --use_pinned_d2h is enabled. " - "Only needs to cover the Python pipeline depth (not Rust reply threads) since " - "reply_request() copies into heap memory before donating to Rust.", + help="Floor for the number of pinned host buffers per output shape when " + "--use_pinned_d2h is enabled; the runner auto-raises it to " + "len(retrieval_dataset_types)+1 to cover all transfers within one batch.", ) parser.add_argument( "--embedding_gather_threads", @@ -706,6 +714,14 @@ def run( "does not starve per-request embedding H2D. None = unlimited / single " "device_put per leaf (default). Only used by the live swap path.", ) + parser.add_argument( + "--hotswap_malloc_trim", + type=str2bool, + default=True, + help="glibc malloc_trim(0) on the coordinator thread after each applied " + "hotswap; returns the cycle's freed pages (glibc slots otherwise " + "ratchet ~196 MiB/cycle). No-op under jemalloc.", + ) parser.add_argument( "--hotswap_stage_chunk_mib", type=int, @@ -786,16 +802,6 @@ def run( "but makes each beam SID a prefix matching many corpus posts." ), ) - parser.add_argument( - "--sid_endpoint", - type=str, - default=None, - help=( - "gRPC endpoint of the SID service for SID-aware retrieval inference. " - "Required when the model uses use_post_sid=True. The endpoint must match " - "the model's trained sid_codebook_size." - ), - ) parser.add_argument( "--jax_compilation_cache_dir", type=str, @@ -845,5 +851,13 @@ def run( ) args = parser.parse_args() + for flag, env in ( + ("copy_tls_ca", "COPY_PORT_TLS_CA"), + ("copy_tls_cert", "COPY_PORT_TLS_CERT"), + ("copy_tls_key", "COPY_PORT_TLS_KEY"), + ("copy_tls_server_name", "COPY_PORT_TLS_SERVER_NAME"), + ): + if value := getattr(args, flag): + os.environ[env] = value pin_visible_devices(args.num_devices_per_process) run(args=args) diff --git a/phoenix/xrex/inference/model_runner.py b/phoenix/xrex/inference/model_runner.py index 3455b865..ed432442 100644 --- a/phoenix/xrex/inference/model_runner.py +++ b/phoenix/xrex/inference/model_runner.py @@ -120,9 +120,12 @@ def _classify_copy_port_error(error: Exception) -> str: def _resolve_copy_url_to_http(copy_url: str) -> str: + scheme = "http" + if "://" in copy_url: + scheme, copy_url = copy_url.split("://", 1) name, port = copy_url.split(":") addr_info = socket.getaddrinfo(name, None, family=socket.AF_INET, type=socket.SOCK_STREAM) - return ",".join(f"http://{x[4][0]}:{port}" for x in addr_info) + return ",".join(f"{scheme}://{x[4][0]}:{port}" for x in addr_info) _SHM_WEIGHTS_PATH = "/dev/shm/model_weights.bin" @@ -163,6 +166,29 @@ def _process_start_time_epoch() -> float | None: return None +_malloc_trim_fn: Any = None + + +def _glibc_malloc_trim() -> tuple[bool, float] | None: + global _malloc_trim_fn + if _malloc_trim_fn is False: + return None + try: + if _malloc_trim_fn is None: + import ctypes + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + _malloc_trim_fn = libc.malloc_trim + _malloc_trim_fn.argtypes = [ctypes.c_size_t] + _malloc_trim_fn.restype = ctypes.c_int + t0 = time.time() + released = _malloc_trim_fn(0) + return bool(released), time.time() - t0 + except Exception: + _malloc_trim_fn = False + return None + + def _save_tensors_to_shm( tensors: list[tuple[str, np.ndarray]], prefix: str, @@ -475,8 +501,6 @@ class BaseModelRunner(RecsysTrainer, Generic[RequestBatch, ModelConfig], ABC): readiness_port: int | None = None max_inflight_requests: int = 4096 - sid_endpoint: str | None = None - channel_size: int = 2048 enqueue_timeout_ms: int = 1000 queue_max_staleness_ms: int = 1200 @@ -489,7 +513,7 @@ class BaseModelRunner(RecsysTrainer, Generic[RequestBatch, ModelConfig], ABC): _service_timer: Timer | None = field(default=None, init=False) use_pipelining: bool = True embedding_gather_threads: int = 16 - use_pinned_d2h: bool = False + use_pinned_d2h: bool = True pinned_d2h_num_buffers: int = 3 log_rotate: bool = False @@ -631,6 +655,7 @@ def _write_reload_done(self) -> None: hotswap_stage_rate_limit_gbps: float | None = None hotswap_stage_chunk_mib: int = 32 + hotswap_malloc_trim: bool = True _emb_table_slots: list[np.ndarray] = field(default_factory=list, init=False) _active_emb_slot: int = field(default=0, init=False) @@ -638,6 +663,7 @@ def _write_reload_done(self) -> None: _host_state_slots: list[Any] = field(default_factory=list, init=False) _active_host_slot: int = field(default=0, init=False) _reload_requested: threading.Event = field(default_factory=threading.Event, init=False) + _hotswap_cycle_inflight: threading.Event = field(default_factory=threading.Event, init=False) _swap_ready: threading.Event = field(default_factory=threading.Event, init=False) _swap_complete: threading.Event = field(default_factory=threading.Event, init=False) _hotswap_stop: threading.Event = field(default_factory=threading.Event, init=False) @@ -859,17 +885,19 @@ def _init_hotswap_buffers(self) -> None: self.h2d_states[rid] = self.h2d_states[rid]._replace(emb_table=slot0) del old_heap_emb + t_slot1 = time.time() slot1 = create_memmap_emb_table( "/dev/shm/emb_table_standby.dat", emb_table_shape, self.model_config.embedding_dtype, - prefault=False, + prefault=True, ) - self._standby_emb_prefaulted = False + self._standby_emb_prefaulted = True logger.info( - "[hotswap] Standby emb_table created sparsely (%.2f GiB); page " - "prefault deferred to the coordinator thread", + "[hotswap] Standby emb_table created and prefaulted (%.2f GiB) " + "in %.1fs, before readiness", slot1.nbytes / (1 << 30), + time.time() - t_slot1, ) else: logger.info( @@ -1493,6 +1521,24 @@ def _respawn_loader_subprocess(self) -> None: pe_key = "post_embeddings.embeddings" self._spawn_loader_subprocess(emb_table_shape, pe_key) + def _hotswap_cycle_busy(self) -> bool: + return ( + self._reload_requested.is_set() + or self._hotswap_cycle_inflight.is_set() + or self._swap_ready.is_set() + or self._live_swap_ready.is_set() + ) + + def _reset_server_reload_request(self) -> None: + server = self._server_for_hotswap + reset = getattr(server, "reset_reload_request", None) if server is not None else None + if reset is None: + return + try: + reset() + except Exception as e: + logger.warning("[hotswap] reset_reload_request failed: %s", e) + def _abort_hotswap_cycle(self, stage: str, reason: str) -> None: logger.warning( "[hotswap] Reload cycle aborted at stage=%s (%s); keeping current " @@ -1511,6 +1557,22 @@ def _abort_hotswap_cycle(self, stage: str, reason: str) -> None: if self.is_multi_worker: self._write_reload_done() self._hotswap_aborted.set() + self._reset_server_reload_request() + + def _malloc_trim_after_swap(self) -> None: + if not self.hotswap_malloc_trim: + return + result = _glibc_malloc_trim() + if result is None: + self.hotswap_malloc_trim = False + logger.warning("[hotswap] malloc_trim unavailable; per-cycle trim disabled") + return + released, secs = result + if self.metrics_publisher is not None: + self.metrics_publisher.checkpoint_reload_step_seconds.labels( + step="malloc_trim" + ).observe(secs) + logger.info("[hotswap] malloc_trim(0) released_any=%s in %.3fs", released, secs) def _coordinator_loop(self) -> None: logger.info("[hotswap] Coordinator thread started (subprocess loader).") @@ -1535,97 +1597,105 @@ def _coordinator_loop(self) -> None: if not triggered: continue - self._reload_requested.clear() - logger.info("[hotswap] Coordinator: reload triggered, sending to subprocess.") - bg_start = time.time() - + self._hotswap_cycle_inflight.set() try: - assert self._request_pipe is not None - assert self._result_pipe is not None - verify_in_subprocess = ( - self.checkpoint_config.verify_checksums and self._live_swap_enabled - ) - self._request_pipe.send( - (self.elapsed_samples, self._active_emb_slot, verify_in_subprocess) - ) - result = self._result_pipe.recv() - except (BrokenPipeError, EOFError, OSError) as e: - self._abort_hotswap_cycle("subprocess_crash", f"subprocess pipe broken: {e}") - self._respawn_loader_subprocess() - continue - - bg_elapsed = time.time() - bg_start - - if self.metrics_publisher is not None: - self.metrics_publisher.hotswap_background_load_seconds.observe(bg_elapsed) - self.metrics_publisher.checkpoint_reload_step_seconds.labels( - step="hotswap_background_load" - ).observe(bg_elapsed) + self._reload_requested.clear() + logger.info("[hotswap] Coordinator: reload triggered, sending to subprocess.") + bg_start = time.time() - if result[0] == "error": - error_msg = result[1] - if "no greater than old prefix" in error_msg: - logger.info("[hotswap] Subprocess: no new checkpoint (%.2fs)", bg_elapsed) - self._reload_noop.set() - if self.is_multi_worker: - self._write_reload_done() - if self._hotswap_standby_meta_file is not None: - self._publish_standby_metadata({"status": "noop"}) - else: - self._abort_hotswap_cycle( - "background_load", - f"subprocess download failed after {bg_elapsed:.2f}s: {error_msg}", + try: + assert self._request_pipe is not None + assert self._result_pipe is not None + verify_in_subprocess = ( + self.checkpoint_config.verify_checksums and self._live_swap_enabled ) - if self._hotswap_standby_meta_file is not None: - self._publish_standby_metadata({"status": "error", "error": error_msg}) - continue + self._request_pipe.send( + (self.elapsed_samples, self._active_emb_slot, verify_in_subprocess) + ) + result = self._result_pipe.recv() + except (BrokenPipeError, EOFError, OSError) as e: + self._abort_hotswap_cycle("subprocess_crash", f"subprocess pipe broken: {e}") + self._respawn_loader_subprocess() + continue - _, prefix, checksums, created_ts, emb_checksum, pe_checksum = result - self._pending_prefix = prefix - self._pending_checksums = checksums - self._pending_emb_checksum = emb_checksum - self._pending_pe_checksum = pe_checksum - self._pending_elapsed_samples = int(prefix.split("_")[2].split("/")[0]) - self._pending_checkpoint_timestamp = created_ts if created_ts > 0 else time.time() + bg_elapsed = time.time() - bg_start - logger.info( - "[hotswap] Subprocess download complete (%.2fs), prefix=%s", - bg_elapsed, - prefix, - ) + if self.metrics_publisher is not None: + self.metrics_publisher.hotswap_background_load_seconds.observe(bg_elapsed) + self.metrics_publisher.checkpoint_reload_step_seconds.labels( + step="hotswap_background_load" + ).observe(bg_elapsed) + + if result[0] == "error": + error_msg = result[1] + if "no greater than old prefix" in error_msg: + logger.info("[hotswap] Subprocess: no new checkpoint (%.2fs)", bg_elapsed) + self._reload_noop.set() + if self.is_multi_worker: + self._write_reload_done() + if self._hotswap_standby_meta_file is not None: + self._publish_standby_metadata({"status": "noop"}) + self._reset_server_reload_request() + else: + self._abort_hotswap_cycle( + "background_load", + f"subprocess download failed after {bg_elapsed:.2f}s: {error_msg}", + ) + if self._hotswap_standby_meta_file is not None: + self._publish_standby_metadata({"status": "error", "error": error_msg}) + continue - if self._live_swap_enabled: - staged = self._stage_live_params() - if staged and self._live_pe_enabled: - staged = self._stage_live_post_embeddings() - if staged: - t_stage = time.time() - bg_start - if self.metrics_publisher is not None: - self.metrics_publisher.checkpoint_reload_step_seconds.labels( - step="hotswap_live_stage" - ).observe(t_stage) - logger.info("[hotswap-live] Standby GPU state staged (%.2fs)", t_stage) - self._live_swap_ready.set() - self._swap_complete.wait() - self._swap_complete.clear() - continue + _, prefix, checksums, created_ts, emb_checksum, pe_checksum = result + self._pending_prefix = prefix + self._pending_checksums = checksums + self._pending_emb_checksum = emb_checksum + self._pending_pe_checksum = pe_checksum + self._pending_elapsed_samples = int(prefix.split("_")[2].split("/")[0]) + self._pending_checkpoint_timestamp = created_ts if created_ts > 0 else time.time() - if self._hotswap_standby_meta_file is not None: - self._publish_standby_metadata( - { - "status": "success", - "prefix": prefix, - "checksums": checksums, - "emb_checksum": emb_checksum, - "pe_checksum": pe_checksum, - "created_ts": created_ts, - "elapsed_samples": self._pending_elapsed_samples, - } + logger.info( + "[hotswap] Subprocess download complete (%.2fs), prefix=%s", + bg_elapsed, + prefix, ) - self._swap_ready.set() - self._swap_complete.wait() - self._swap_complete.clear() + if self._live_swap_enabled: + staged = self._stage_live_params() + if staged and self._live_pe_enabled: + staged = self._stage_live_post_embeddings() + if staged: + t_stage = time.time() - bg_start + if self.metrics_publisher is not None: + self.metrics_publisher.checkpoint_reload_step_seconds.labels( + step="hotswap_live_stage" + ).observe(t_stage) + logger.info("[hotswap-live] Standby GPU state staged (%.2fs)", t_stage) + self._live_swap_ready.set() + self._swap_complete.wait() + self._swap_complete.clear() + self._malloc_trim_after_swap() + continue + + if self._hotswap_standby_meta_file is not None: + self._publish_standby_metadata( + { + "status": "success", + "prefix": prefix, + "checksums": checksums, + "emb_checksum": emb_checksum, + "pe_checksum": pe_checksum, + "created_ts": created_ts, + "elapsed_samples": self._pending_elapsed_samples, + } + ) + + self._swap_ready.set() + self._swap_complete.wait() + self._swap_complete.clear() + self._malloc_trim_after_swap() + + finally: + self._hotswap_cycle_inflight.clear() logger.info("[hotswap] Coordinator thread exiting.") @@ -1676,68 +1746,77 @@ def _follower_coordinator_loop(self) -> None: break if not triggered: continue - self._reload_requested.clear() + self._hotswap_cycle_inflight.set() + try: + self._reload_requested.clear() + + t_wait = time.time() + while not self._hotswap_stop.is_set(): + try: + with open(self._hotswap_standby_version_file) as f: + cur = int(f.read().strip() or "0") + except (FileNotFoundError, ValueError): + cur = last_version + if cur > last_version: + last_version = cur + break + if time.time() - t_wait > 1800: + logger.warning( + "[hotswap] Worker %d: timed out waiting for leader's " + "standby payload after 30 min, giving up this cycle", + wid, + ) + break + time.sleep(0.5) + else: + continue - t_wait = time.time() - while not self._hotswap_stop.is_set(): try: - with open(self._hotswap_standby_version_file) as f: - cur = int(f.read().strip() or "0") - except (FileNotFoundError, ValueError): - cur = last_version - if cur > last_version: - last_version = cur - break - if time.time() - t_wait > 1800: + with open(self._hotswap_standby_meta_file) as f: + meta = _json.load(f) + except (FileNotFoundError, _json.JSONDecodeError) as e: logger.warning( - "[hotswap] Worker %d: timed out waiting for leader's " - "standby payload after 30 min, giving up this cycle", + "[hotswap] Worker %d: failed to read leader metadata: %s", wid, + e, ) - break - time.sleep(0.5) - else: - continue + continue - try: - with open(self._hotswap_standby_meta_file) as f: - meta = _json.load(f) - except (FileNotFoundError, _json.JSONDecodeError) as e: - logger.warning( - "[hotswap] Worker %d: failed to read leader metadata: %s", - wid, - e, - ) - continue + status = meta.get("status") + if status == "success": + self._pending_prefix = meta["prefix"] + self._pending_checksums = meta["checksums"] + self._pending_emb_checksum = meta.get("emb_checksum") + self._pending_pe_checksum = meta.get("pe_checksum") + self._pending_elapsed_samples = meta["elapsed_samples"] + created_ts = meta.get("created_ts", 0) + self._pending_checkpoint_timestamp = ( + created_ts if created_ts > 0 else time.time() + ) + logger.info( + "[hotswap] Worker %d: leader standby ready (version=%d, prefix=%s)", + wid, + last_version, + self._pending_prefix, + ) + self._swap_ready.set() + self._swap_complete.wait() + self._swap_complete.clear() + self._malloc_trim_after_swap() + elif status == "noop": + logger.info("[hotswap] Worker %d: leader reported noop (no new ckpt)", wid) + self._reload_noop.set() + if self.is_multi_worker: + self._write_reload_done() + self._reset_server_reload_request() + else: + self._abort_hotswap_cycle( + "leader_error", + f"worker {wid}: leader reported error/unknown status {status!r}", + ) - status = meta.get("status") - if status == "success": - self._pending_prefix = meta["prefix"] - self._pending_checksums = meta["checksums"] - self._pending_emb_checksum = meta.get("emb_checksum") - self._pending_pe_checksum = meta.get("pe_checksum") - self._pending_elapsed_samples = meta["elapsed_samples"] - created_ts = meta.get("created_ts", 0) - self._pending_checkpoint_timestamp = created_ts if created_ts > 0 else time.time() - logger.info( - "[hotswap] Worker %d: leader standby ready (version=%d, prefix=%s)", - wid, - last_version, - self._pending_prefix, - ) - self._swap_ready.set() - self._swap_complete.wait() - self._swap_complete.clear() - elif status == "noop": - logger.info("[hotswap] Worker %d: leader reported noop (no new ckpt)", wid) - self._reload_noop.set() - if self.is_multi_worker: - self._write_reload_done() - else: - self._abort_hotswap_cycle( - "leader_error", - f"worker {wid}: leader reported error/unknown status {status!r}", - ) + finally: + self._hotswap_cycle_inflight.clear() logger.info("[hotswap] Follower coordinator (worker %d) exiting.", wid) @@ -3065,12 +3144,15 @@ def as_np_array(x: jax.Array) -> np.ndarray: buffer_key = (x.shape, np.dtype(x.dtype)) if buffer_key not in self._pinned_buffers: n_shards = self.parallel_config.num_devices_per_process - self._pinned_buffers[buffer_key] = PinnedD2HBuffer( + min_buffers = len(self.retrieval_dataset_types) + 1 + buf = PinnedD2HBuffer( shape=x.shape, dtype=np.dtype(x.dtype), num_shards=n_shards, - num_buffers=self.pinned_d2h_num_buffers, + num_buffers=max(self.pinned_d2h_num_buffers, min_buffers), ) + buf.begin_cycle() + self._pinned_buffers[buffer_key] = buf return self._pinned_buffers[buffer_key].transfer(x) if x.dtype == jnp.bfloat16: return np.asarray(x, dtype=np.float32) @@ -3116,6 +3198,10 @@ def maybe_process_and_reply( batch_id=last_batch_id, approx_output_bytes=approx_output_bytes, ): + if self.use_pinned_d2h: + for pinned_buf in self._pinned_buffers.values(): + pinned_buf.begin_cycle() + if isinstance(last_out, jax.Array): with jax_profiler.TraceAnnotation( "block_until_ready", @@ -3165,7 +3251,7 @@ def maybe_process_and_reply( ) *model_out_jax, has_nan_jax = last_out output_array = tuple(as_np_array(x) for x in model_out_jax) - has_nan = bool(np.asarray(has_nan_jax)) + has_nan = bool(np.any(np.asarray(has_nan_jax))) logger.info( f"[batch={last_batch_id}] as_np_array took {as_np_array_t.elapsed() * 1000:.2f}ms" ) @@ -3229,11 +3315,7 @@ def maybe_process_and_reply( break reload_triggered = False - hotswap_busy = self.enable_hotswap and ( - self._reload_requested.is_set() - or self._swap_ready.is_set() - or self._live_swap_ready.is_set() - ) + hotswap_busy = self.enable_hotswap and self._hotswap_cycle_busy() if check_reload and not hotswap_busy: if self.is_multi_worker and self._observe_peer_reload_trigger(): reload_triggered = True @@ -3269,11 +3351,7 @@ def maybe_process_and_reply( if self.enable_hotswap and restart_timeout is not None: if time.time() - hotswap_last_reload_trigger >= restart_timeout: - if ( - not self._reload_requested.is_set() - and not self._swap_ready.is_set() - and not self._live_swap_ready.is_set() - ): + if not self._hotswap_cycle_busy(): logger.info( "[hotswap] Periodic reload trigger (every %.0fs)", restart_timeout, @@ -4043,6 +4121,11 @@ def create_server( num_user_installed_apps=recsys_batch.NUM_USER_INSTALLED_APPS, num_post_categorical_features=recsys_batch.POST_CATEGORICAL_FEATURE_SIZE, num_post_bool_features=recsys_batch.POST_BOOL_FEATURE_SIZE, + enable_stale_post=bool( + getattr( + getattr(self.model_config, "feature_prep", None), "enable_stale_post", False + ) + ), num_post_float_features=recsys_batch.POST_FLOAT_FEATURE_SIZE, num_post_int64_features=recsys_batch.POST_INT64_FEATURE_SIZE, enable_async_response_compression=self.enable_async_response_compression, @@ -4191,10 +4274,13 @@ def forward_fn( logits, candidate_continuous_predictions = model.forward(batch, recsys_embeddings) log_probs = jax.nn.log_sigmoid(logits).astype(jnp.float32) cont_preds = candidate_continuous_predictions.astype(jnp.float32) - has_nan = jnp.any(jnp.isnan(log_probs)) + has_nan = jnp.any(jnp.isnan(log_probs), axis=tuple(range(1, log_probs.ndim))) return log_probs, cont_preds, has_nan self.forward_fn = forward_fn + nan_sharding = jax.sharding.NamedSharding( + self.data_sharding.mesh, P(self.data_sharding.spec[0]) + ) forward_jit = JittedOrCompiled( jax.jit( @@ -4205,7 +4291,7 @@ def forward_fn( self.data_sharding, self.data_sharding, ), - out_shardings=(self.data_sharding, self.data_sharding, None), + out_shardings=(self.data_sharding, self.data_sharding, nan_sharding), ), name=f"forward_fn_bs{bs}", ) @@ -4659,21 +4745,8 @@ def create_server( assert isinstance(self.model_config, RecsysTwoTowerModelConfig) hash_keys = self.dataset.hash_table.hash_keys - sid_client = None - _use_post_sid = self.model_config.user_tower_config.use_post_sid - sid_num_levels = self.model_config.user_tower_config.sid_num_levels - if _use_post_sid and self.sid_endpoint and sid_num_levels > 0: - sid_client = xai_recsys_engine.PySemanticIdClient( - self.sid_endpoint, - sid_num_levels, - ) - logger.info( - "SID client connected: endpoint=%s, sid_num_levels=%d", - self.sid_endpoint, - sid_num_levels, - ) - elif _use_post_sid and not self.sid_endpoint: - raise ValueError("use_post_sid=True but no sid_endpoint configured") + user_tower = self.model_config.user_tower_config + sid_num_levels = user_tower.sid_num_levels if user_tower.use_post_sid else 0 return xai_recsys_engine.RecsysRetrievalPredictorServer( self.grpc_port, @@ -4690,7 +4763,6 @@ def create_server( service_time_ewma_alpha=self.service_time_ewma_alpha, pipeline_depth=1 if self.use_pipelining else 0, mm_client=mm_client, - sid_client=sid_client, user_id_table_size=hash_keys.user_id_table_size, user_hash_scales=hash_keys.user_hash_scales, user_biases=hash_keys.user_biases, @@ -4705,7 +4777,7 @@ def create_server( author_biases=hash_keys.author_biases, author_modulus=hash_keys.author_modulus, output_vocab_size=self.dataset.hash_table.output_vocab_size, - num_continuous_actions=self.model_config.num_continuous_actions, + num_continuous_actions=self.dataset.num_continuous_actions, history_seq_len=self.history_seq_len, candidate_seq_len=self.candidate_seq_len, multimodal_embedding_dim=getattr(self.model_config, "multimodal_embedding_dim", 0), @@ -4720,10 +4792,15 @@ def create_server( num_user_installed_apps=recsys_batch.NUM_USER_INSTALLED_APPS, num_post_categorical_features=recsys_batch.POST_CATEGORICAL_FEATURE_SIZE, num_post_bool_features=recsys_batch.POST_BOOL_FEATURE_SIZE, + enable_stale_post=bool( + getattr( + getattr(self.model_config, "feature_prep", None), "enable_stale_post", False + ) + ), num_post_float_features=recsys_batch.POST_FLOAT_FEATURE_SIZE, num_post_int64_features=recsys_batch.POST_INT64_FEATURE_SIZE, enable_async_response_compression=self.enable_async_response_compression, - sid_num_levels=sid_num_levels if sid_client is not None else 0, + sid_num_levels=sid_num_levels, ) def create_executable(self) -> None: diff --git a/phoenix/xrex/inference/pinned_d2h.py b/phoenix/xrex/inference/pinned_d2h.py index 5afaaafa..fbab640b 100644 --- a/phoenix/xrex/inference/pinned_d2h.py +++ b/phoenix/xrex/inference/pinned_d2h.py @@ -41,6 +41,8 @@ def __init__( self._buf_idx = 0 + self._transfers_since_cycle: int | None = None + logger.info( f"PinnedD2HBuffer: allocated {num_buffers * self._nbytes / 1024 / 1024:.1f} MB " f"pinned memory ({num_buffers}x{self._nbytes / 1024 / 1024:.1f} MB ring) " @@ -65,6 +67,14 @@ def transfer(self, arr: jax.Array) -> np.ndarray: f"Expected {self._num_shards} shards, got {len(shards)}" ) + if self._transfers_since_cycle is not None: + self._transfers_since_cycle += 1 + if self._transfers_since_cycle > self._num_buffers: + raise RuntimeError( + f"PinnedD2HBuffer ring overflow: more than {self._num_buffers} transfers " + f"in one cycle for shape {self._shape} — would overwrite a live view." + ) + buf_idx = self._buf_idx self._buf_idx = (self._buf_idx + 1) % self._num_buffers @@ -93,5 +103,8 @@ def transfer(self, arr: jax.Array) -> np.ndarray: return self._outs[buf_idx] + def begin_cycle(self) -> None: + self._transfers_since_cycle = 0 + def __del__(self) -> None: pass diff --git a/phoenix/xrex/inference/sid_retrieval_runner.py b/phoenix/xrex/inference/sid_retrieval_runner.py index 82d48afc..095ea3d5 100644 --- a/phoenix/xrex/inference/sid_retrieval_runner.py +++ b/phoenix/xrex/inference/sid_retrieval_runner.py @@ -377,22 +377,7 @@ def create_server( assert isinstance(self.model_config, RecsysSIDRetrievalConfig) hash_keys = self.dataset.hash_table.hash_keys - sid_client = None - if self.model_config.use_post_sid and self.model_config.sid_num_levels > 0: - if not self.sid_endpoint: - raise ValueError( - "v4 SID retrieval has use_post_sid=True; " - "must pass --sid_endpoint to launch_inference.py" - ) - sid_client = xai_recsys_engine.PySemanticIdClient( - self.sid_endpoint, - self.model_config.sid_num_levels, - ) - logger.info( - "SID client connected: endpoint=%s, sid_num_levels=%d", - self.sid_endpoint, - self.model_config.sid_num_levels, - ) + sid_num_levels = self.model_config.sid_num_levels if self.model_config.use_post_sid else 0 return xai_recsys_engine.RecsysRetrievalPredictorServer( self.grpc_port, @@ -409,7 +394,6 @@ def create_server( service_time_ewma_alpha=self.service_time_ewma_alpha, pipeline_depth=1 if self.use_pipelining else 0, mm_client=mm_client, - sid_client=sid_client, user_id_table_size=hash_keys.user_id_table_size, user_hash_scales=hash_keys.user_hash_scales, user_biases=hash_keys.user_biases, @@ -442,5 +426,5 @@ def create_server( num_post_float_features=recsys_batch.POST_FLOAT_FEATURE_SIZE, num_post_int64_features=recsys_batch.POST_INT64_FEATURE_SIZE, enable_async_response_compression=self.enable_async_response_compression, - sid_num_levels=(self.model_config.sid_num_levels if sid_client is not None else 0), + sid_num_levels=sid_num_levels, ) diff --git a/phoenix/xrex/models/recsys_attention.py b/phoenix/xrex/models/recsys_attention.py index 59a0619e..4d2660aa 100644 --- a/phoenix/xrex/models/recsys_attention.py +++ b/phoenix/xrex/models/recsys_attention.py @@ -218,11 +218,25 @@ def sharded_custom_op_with_extra_args(self): ) def sharded_mha(q, k, v, segment_ids, segment_ids_k, temp): - del segment_ids, segment_ids_k, temp - batch_size, seq_len, num_q_heads, _ = q.shape + del segment_ids_k, temp + _, seq_len, num_q_heads, _ = q.shape hist_len = config.num_user_prefix_tokens + config.history_seq_len - bs_layout = build_dense_block_sparse_layout(batch_size, seq_len, hist_len, num_q_heads) - return ranker_attention_fa4(q, k, v, sm_scale, bs_layout), None + hist_valid_len = jnp.sum(segment_ids[:, :hist_len] == 1, axis=-1, dtype=jnp.int32) + fwd_bs, bwd_bs, valid_upper, valid_lower = build_dense_block_sparse_layout( + seq_len, hist_len, num_q_heads, hist_valid_len + ) + return ( + ranker_attention_fa4( + q, + k, + v, + sm_scale, + (fwd_bs, bwd_bs), + valid_block_upper=valid_upper, + valid_block_lower=valid_lower, + ), + None, + ) return sharded_mha, () @@ -265,6 +279,8 @@ def sharded_mha( bsp_bwd_full_idx, bsp_bwd_diag_cnt, bsp_bwd_diag_idx, + bsp_valid_block_upper, + bsp_valid_block_lower, ): del segment_ids, segment_ids_k, temp fwd_bs = ( @@ -290,6 +306,8 @@ def sharded_mha( v, sm_scale, block_sparse_layout=(fwd_bs, bwd_bs), + valid_block_upper=bsp_valid_block_upper, + valid_block_lower=bsp_valid_block_lower, ), None, ) @@ -310,6 +328,8 @@ def _rep(s): "bsp_bwd_full_idx", "bsp_bwd_diag_cnt", "bsp_bwd_diag_idx", + "bsp_valid_block_upper", + "bsp_valid_block_lower", ) extras = tuple((name, _rep) for name in bsp_names) return sharded_mha, extras diff --git a/phoenix/xrex/models/recsys_feature_prep.py b/phoenix/xrex/models/recsys_feature_prep.py index 4b300010..3158b4a5 100644 --- a/phoenix/xrex/models/recsys_feature_prep.py +++ b/phoenix/xrex/models/recsys_feature_prep.py @@ -13,7 +13,9 @@ from xai_configlib import Config, configclass from xrex.data.recsys.feature_config import ( + BoolFeature, CategoricalFeature, + Int64Feature, UserCategoricalFeature, UserFloatFeature, ) @@ -22,6 +24,15 @@ from xrex.models.recsys_embedding import HashKeys, RecsysEmbeddings from xrex.models.scaling import ScaleConfig +ENGAGEMENT_COUNT_ORDER: tuple[Int64Feature, ...] = ( + Int64Feature.favCountSeq, + Int64Feature.replyCountSeq, + Int64Feature.repostCountSeq, + Int64Feature.quoteCountSeq, + Int64Feature.viewCountSeq, +) +ENGAGEMENT_COUNT_LOG2_SCALE: float = 32.0 + class FeaturePrepStreams(NamedTuple): user_tokens: jax.Array @@ -166,6 +177,8 @@ class FeaturePrepConfig(Config): enable_bridge_prob: bool = False product_surface_cardinality: int = 16 timezone_cardinality: int = 32 + enable_engagement_counts: bool = False + engagement_count_mlp_hidden_dim: int = 64 enable_time_of_day: bool = False time_of_day_kernel: Literal["box", "triangle", "cosine"] = "cosine" @@ -189,6 +202,9 @@ class FeaturePrepConfig(Config): enable_day_of_week: bool = False day_of_week_cardinality: int = 8 + enable_is_author_followed_by_viewer: bool = False + enable_is_author_following_viewer: bool = False + enable_post_sid: bool = False sid_embed_dim: int = 1024 sid_num_levels: int = 6 @@ -198,6 +214,7 @@ class FeaturePrepConfig(Config): multimodal_embedding_dim: int = 0 search_query_embedding_dim: int = 0 + enable_stale_post: bool = False @property def has_user_features(self) -> bool: @@ -589,6 +606,65 @@ def _add_context_features( dow, config.day_of_week_cardinality, f"{prefix}_day_of_week_emb", config ).astype(fprop_dtype) + if config.enable_is_author_followed_by_viewer or config.enable_is_author_following_viewer: + bool_features = batch_seq.get("bool_features") + if bool_features is not None: + bool_features = _cast_jax(bool_features) + if config.enable_is_author_followed_by_viewer: + followed = bool_features[:, :, BoolFeature.isAuthorFollowedByViewerSeq].astype( + jnp.int32 + ) + result = result + _embed_categorical( + followed, 2, f"{prefix}_is_author_followed_by_viewer_emb", config + ).astype(fprop_dtype) + if config.enable_is_author_following_viewer: + following = bool_features[:, :, BoolFeature.isAuthorFollowingViewerSeq].astype( + jnp.int32 + ) + result = result + _embed_categorical( + following, 2, f"{prefix}_is_author_following_viewer_emb", config + ).astype(fprop_dtype) + + if config.enable_engagement_counts: + i64 = batch_seq.get("int64_features") + if i64 is not None: + i64 = _cast_jax(i64) + x = jnp.stack( + [ + jnp.log2(jnp.maximum(i64[:, :, f.value].astype(jnp.float32), 0.0) + 1.0) + for f in ENGAGEMENT_COUNT_ORDER + ], + axis=-1, + ) + v = x / ENGAGEMENT_COUNT_LOG2_SCALE + v = jnp.concatenate([v, jnp.ones_like(v[..., :1])], axis=-1) + h_proj = _get_proj( + f"{prefix}_count_lograw_mlp_in", + v.shape[-1], + config.engagement_count_mlp_hidden_dim, + config, + role="input_proj", + ) + out_proj = _get_proj( + f"{prefix}_count_lograw_mlp_out", + config.engagement_count_mlp_hidden_dim, + config.emb_size, + config, + role="input_proj", + ) + hidden = jax.nn.gelu(jnp.dot(v.astype(h_proj.dtype), h_proj)) + result = result + jnp.dot(hidden, out_proj).astype(fprop_dtype) + + if config.enable_stale_post and prefix != "hist": + bools = batch_seq.get("bool_features") + if bools is not None and bools.shape[-1] > BoolFeature.isStalePost14d.value: + is_stale = _cast_jax(bools)[:, :, BoolFeature.isStalePost14d.value].astype(jnp.int32) + emb_table = _get_emb_table( + f"{prefix}_is_stale_post_14d_emb", 2, config.emb_size, config + ) + one_hot = jax.nn.one_hot(jnp.clip(is_stale, 0, 1), 2) + result = result + jnp.dot(one_hot, emb_table).astype(fprop_dtype) + return result @@ -754,9 +830,14 @@ def _add_history_features( ).astype(fprop_dtype) if config.enable_bridge_prob: + from xai_proto import recsys_pb2 + + _bridge_idx = recsys_pb2.ContinuousActionName.BRIDGE_PROBABILITY cont_actions = batch["history_seq"].get("continuous_actions") - if cont_actions is not None: - bridge_p = jnp.clip(_cast_jax(cont_actions)[:, :, 0].astype(jnp.float32), 0.0, 1.0) + if cont_actions is not None and cont_actions.shape[-1] > _bridge_idx: + bridge_p = jnp.clip( + _cast_jax(cont_actions)[:, :, _bridge_idx].astype(jnp.float32), 0.0, 1.0 + ) result = result + _embed_scalar_times_vector( bridge_p, "hist_bridge_prob_vec", config ).astype(fprop_dtype) diff --git a/phoenix/xrex/models/recsys_model.py b/phoenix/xrex/models/recsys_model.py index fabf4f6f..e23c2053 100644 --- a/phoenix/xrex/models/recsys_model.py +++ b/phoenix/xrex/models/recsys_model.py @@ -21,11 +21,21 @@ from xrex.data.recsys.constants import ( CLICK_ACTION_INDEX, CLICK_CONDITIONED_ACTION_INDICES, + MACT_IN_APP_LOSS_ACTION_INDICES, NEGATIVE_FEEDBACK_HEAD_INDICES, + SEARCH_RELEVANCE_ACTION_INDICES, + SOURCE_SPLIT_CONVERSION_HEAD_INDICES, + VIEW_THROUGH_ACTION_INDICES, action_type_map, engagement_to_ids, ) -from xrex.data.recsys.feature_config import ENGAGEMENT_COUNT_BUCKET_MAP, CategoricalFeature +from xrex.data.recsys.feature_config import ( + ADS_PRODUCT_KEY_TABLE_SIZE, + ENGAGEMENT_COUNT_BUCKET_MAP, + BoolFeature, + CategoricalFeature, + Int64Feature, +) from xrex.data.recsys.recsys_batch import EMBEDDING_CONFIG, EmbeddingType, RecsysFeaturesBatch from xrex.data.recsys.safety_filter import apply_safety_filter, safety_filter_stats from xrex.data.recsys.sequence_packing import SequencePackedLayout @@ -49,6 +59,7 @@ FeaturePrepConfig, build_feature_prep_inputs, ) +from xrex.models.recsys_sid import reconstruct_entity_sid from xrex.models.recsys_user_features import ( UserFeaturesConfig, build_user_feature_parts, @@ -90,6 +101,11 @@ ) +_DPA_PRODUCT_KEY_SLOTS: list[int] = [ + Int64Feature.firstDpaProductKey.value, + Int64Feature.firstDpaProductKeyHash2.value, +] + POST_AGE_MAX_MINUTES = 4800 @@ -498,6 +514,7 @@ class RecsysAggregatedModelConfig(Config): effective_sequence_len: int | None = None transformer_output_only: bool = False + use_dense_action_table: bool = False use_product_surface: bool = False post_age_granularity_mins: int = 60 @@ -512,6 +529,8 @@ class RecsysAggregatedModelConfig(Config): metric_group: str = "default" + metric_mask_keys: list[str] | None = None + multimodal_embedding_type: EmbeddingType | None = None search_query_embedding_dim: int = 0 @@ -524,6 +543,9 @@ class RecsysAggregatedModelConfig(Config): sid_hash_level: bool = False sid_cross_attn: bool = False + sid_embedding_mode: Literal["learned", "recon"] = "learned" + sid_decoder_path: str = "" + use_user_embedding: bool = True use_ip_address: bool = False @@ -537,8 +559,22 @@ class RecsysAggregatedModelConfig(Config): condition_conversion_on_click: bool = False + train_view_through_heads: bool = False + + split_head_training_by_source: bool = False + + mact_in_app_loss_weight: float = 1.0 + + condition_search_relevance_on_prompt: bool = False + enable_platform_metrics: bool = False + dpa_product_embed_dim: int = 32 + + dpa_product_table_size: int = ADS_PRODUCT_KEY_TABLE_SIZE + + enable_dpa_input_embedding: bool = False + user_features: UserFeaturesConfig = UserFeaturesConfig() context_features: ContextFeaturesConfig = ContextFeaturesConfig() @@ -582,6 +618,20 @@ def make(self, sharding_context: ShardingContext) -> RecsysAggregatedModel: "build feature_prep via _make_feature_prep_config or set them consistently." ) + if self.sid_embedding_mode == "recon": + assert not self.feature_prep_enabled, ( + "sid_embedding_mode='recon' is only implemented on the legacy build_inputs " + "path; the feature_prep path embeds SIDs via its own learned tables " + "(FeaturePrepConfig.enable_post_sid) and would silently ignore recon." + ) + assert not self.sid_hash_level and not self.sid_cross_attn, ( + "sid_embedding_mode='recon' replaces the learned SID tables outright; " + "sid_hash_level / sid_cross_attn only apply to mode='learned'." + ) + assert self.sid_decoder_path, ( + "sid_embedding_mode='recon' requires sid_decoder_path " + "(codebook.safetensors with `stages` + `dec_*` tensors)." + ) if self.use_seqpack: attn_config = self.model_config.attn_config assert self.right_anchored_rope @@ -1117,6 +1167,7 @@ def pad_to_next_128_multiple( client_app_id: jax.Array | None = None, line_item_objective: jax.Array | None = None, safety_label_mask: jax.Array | None = None, + dpa_product_key: jax.Array | None = None, ) -> tuple[ jax.Array, jax.Array, @@ -1129,6 +1180,7 @@ def pad_to_next_128_multiple( jax.Array | None, jax.Array | None, jax.Array | None, + jax.Array | None, ]: batch_size, seq_len, emb_dim = embeddings.shape @@ -1200,6 +1252,14 @@ def pad_to_next_128_multiple( else: padded_safety_label_mask = None + if dpa_product_key is not None: + pad_dpa_product_key = jnp.zeros( + (batch_size, pad_length) + dpa_product_key.shape[2:], dtype=dpa_product_key.dtype + ) + padded_dpa_product_key = jnp.concatenate([dpa_product_key, pad_dpa_product_key], axis=1) + else: + padded_dpa_product_key = None + return ( padded_embeddings, padded_mask, @@ -1212,6 +1272,7 @@ def pad_to_next_128_multiple( padded_client_app_id, padded_line_item_objective, padded_safety_label_mask, + padded_dpa_product_key, ) @@ -1229,9 +1290,14 @@ def build_metric_masks( new_user_mask: jax.Array | None = None, line_item_objective: jax.Array | None = None, no_history_mask: jax.Array | None = None, + dpa_product_key: jax.Array | None = None, + delayed_sample_mask: jax.Array | None = None, *, condition_conversion_on_click: bool = False, + condition_search_relevance_on_prompt: bool = False, enable_platform_metrics: bool = False, + split_head_training_by_source: bool = False, + metric_mask_keys: list[str] | None = None, ) -> dict[str, jax.Array]: promoted_mask = mask * (promoted_ids != 0) if promoted_ids is not None else jnp.zeros_like(mask) @@ -1307,17 +1373,58 @@ def build_metric_masks( non_negative_mask * promoted_mask * home_timeline_mask * website_clicks_objective ) + dpa_mask = ( + (dpa_product_key != 0).astype(mask.dtype) + if dpa_product_key is not None + else jnp.zeros_like(mask) + ) + masks["dpa"] = mask * dpa_mask + masks["non_negative_dpa"] = non_negative_mask * dpa_mask + if condition_conversion_on_click: click_mask = raw_targets[:, :, CLICK_ACTION_INDEX].astype(mask.dtype) masks["clicked"] = mask * click_mask masks["non_negative_clicked"] = mask * (1 - negative_sample_mask) * click_mask + if split_head_training_by_source: + assert condition_conversion_on_click, ( + "split_head_training_by_source metric slices require condition_conversion_on_click" + ) + delayed = ( + jnp.zeros_like(mask) + if delayed_sample_mask is None + else delayed_sample_mask.astype(mask.dtype) + ) + click_mask = raw_targets[:, :, CLICK_ACTION_INDEX].astype(mask.dtype) + masks["fresh"] = mask * (1 - delayed) + masks["delayed_clicked"] = mask * delayed * click_mask + masks["delayed_non_clicked"] = mask * delayed * (1 - click_mask) + masks["fresh_home_website_clicks"] = ( + masks["fresh"] * home_timeline_mask * website_clicks_objective + ) + + if condition_search_relevance_on_prompt: + prompt_mask = jnp.any( + raw_targets[:, :, jnp.array(SEARCH_RELEVANCE_ACTION_INDICES)] == 1, axis=-1 + ).astype(mask.dtype) + masks["prompted"] = mask * prompt_mask + masks["non_negative_prompted"] = non_negative_mask * prompt_mask + if enable_platform_metrics and ios_mask is not None and android_mask is not None: masks["ios"] = ios_mask masks["android"] = android_mask masks["ios_non_negative"] = ios_mask * (1 - negative_sample_mask) masks["android_non_negative"] = android_mask * (1 - negative_sample_mask) + if metric_mask_keys: + keep = set(metric_mask_keys) + unknown = keep - masks.keys() + if unknown: + raise ValueError( + f"metric_mask_keys entries {sorted(unknown)} not in available masks {sorted(masks)}" + ) + masks = {k: v for k, v in masks.items() if k in keep} + return masks @@ -1381,7 +1488,12 @@ def _compute_metrics_after_masks( pos_sum = stats[f"{eng_name}_{mask_key}_num_tokens"] batch_stat = jnp.stack([ce_sum, pos_sum, total_count]) - old = jnp.stack([rce_ema[f"{base_key}/{ws}"] for ws in smoothing_windows]) + old = jnp.stack( + [ + rce_ema.get(f"{base_key}/{ws}", jnp.zeros((3,), dtype=jnp.float32)) + for ws in smoothing_windows + ] + ) raw_updated = (1.0 - a) * old + a * batch_stat[None, :] updated = jnp.where( jnp.isnan(raw_updated), @@ -1416,7 +1528,12 @@ def _compute_metrics_after_masks( pos_sum_batch = jnp.sum(y * mask_val) batch_stat = jnp.stack([pred_sum, pos_sum_batch]) - old = jnp.stack([calib_ema[f"{base_key}/{ws}"] for ws in smoothing_windows]) + old = jnp.stack( + [ + calib_ema.get(f"{base_key}/{ws}", jnp.zeros((2,), dtype=jnp.float32)) + for ws in smoothing_windows + ] + ) raw_updated = (1.0 - a) * old + a * batch_stat[None, :] updated = jnp.where( jnp.isnan(raw_updated), @@ -1495,6 +1612,8 @@ def _build_metric_masks( new_user_mask: jax.Array | None = None, line_item_objective: jax.Array | None = None, no_history_mask: jax.Array | None = None, + dpa_product_key: jax.Array | None = None, + delayed_sample_mask: jax.Array | None = None, ) -> dict[str, jax.Array]: return build_metric_masks( mask, @@ -1506,8 +1625,13 @@ def _build_metric_masks( new_user_mask, line_item_objective, no_history_mask, + dpa_product_key, + delayed_sample_mask, condition_conversion_on_click=self.config.condition_conversion_on_click, + condition_search_relevance_on_prompt=self.config.condition_search_relevance_on_prompt, enable_platform_metrics=self.config.enable_platform_metrics, + split_head_training_by_source=self.config.split_head_training_by_source, + metric_mask_keys=self.config.metric_mask_keys, ) def compute_recsys_metrics( @@ -1522,6 +1646,8 @@ def compute_recsys_metrics( new_user_mask: jax.Array | None = None, line_item_objective: jax.Array | None = None, no_history_mask: jax.Array | None = None, + dpa_product_key: jax.Array | None = None, + delayed_sample_mask: jax.Array | None = None, stats: dict | None = None, rce_ema: dict[str, jax.Array] | None = None, rce_alpha: jax.Array | None = None, @@ -1542,6 +1668,8 @@ def compute_recsys_metrics( new_user_mask, line_item_objective, no_history_mask, + dpa_product_key, + delayed_sample_mask, ) return self._compute_metrics_after_masks( @@ -1612,29 +1740,114 @@ def compute_engagement_count_metrics( ) -> dict: if stats is None: stats = {} - ctx_config = self.config.context_features - if not ctx_config.enabled or not ctx_config.enable_engagement_counts or batch is None: + fp_on = ( + self.config.feature_prep_enabled and self.config.feature_prep.enable_engagement_counts + ) + ctx_on = ( + self.config.context_features.enabled + and self.config.context_features.enable_engagement_counts + ) + stale_on = self.config.feature_prep_enabled and self.config.feature_prep.enable_stale_post + if batch is None or not (fp_on or ctx_on or stale_on): return stats - for cat_feat_enum, int64_feat_enum in ENGAGEMENT_COUNT_BUCKET_MAP: - name = cat_feat_enum.name.removesuffix("CountBucketSeq").lower() - max_bucket = ENGAGEMENT_COUNT_MAX_BUCKET.get( - cat_feat_enum, ENGAGEMENT_COUNT_NUM_BUCKETS - 1 - ) - for seq_name, side in (("history_seq", "history"), ("candidate_seq", "candidate")): - seq = batch[seq_name] - raw_i64 = seq.get("int64_features") - post_hashes = seq.get("post_hashes") - if raw_i64 is None or post_hashes is None: - continue - raw_counts = cast_jax(raw_i64)[:, :, int64_feat_enum.value] + def _add_scalar_stats( + name: str, + side: str, + raw_counts: jax.Array, + valid_float: jax.Array, + ) -> None: + raw_non_negative = jnp.maximum(raw_counts.astype(jnp.float32), 0.0) + log2_counts = jnp.log2(raw_non_negative + 1.0) + n_valid = jnp.maximum(jnp.sum(valid_float), 1.0) + positive = (raw_non_negative > 0.0).astype(jnp.float32) + + stats[f"ec/{name}/{side}/coverage"] = jnp.sum(positive * valid_float) / n_valid + stats[f"ec/{name}/{side}/mean_raw"] = jnp.sum(raw_non_negative * valid_float) / n_valid + stats[f"ec/{name}/{side}/mean_log2"] = jnp.sum(log2_counts * valid_float) / n_valid + + def _add_bucket_percentiles( + name: str, + side: str, + bucket_values: jax.Array, + valid_float: jax.Array, + nbins: int, + ) -> None: + n_valid = jnp.maximum(jnp.sum(valid_float), 1.0) + binned = jnp.clip(bucket_values.astype(jnp.int32), 0, nbins - 1) + hist = ( + jnp.zeros((nbins,), jnp.float32) + .at[jnp.reshape(binned, (-1,))] + .add(jnp.reshape(valid_float, (-1,))) + ) / n_valid + cdf = jnp.cumsum(hist) + for q_name, q in (("p50", 0.50), ("p90", 0.90), ("p99", 0.99)): + stats[f"ec/{name}/{side}/{q_name}_bucket"] = jnp.argmax(cdf >= q).astype( + jnp.float32 + ) + + for seq_name, side in (("history_seq", "history"), ("candidate_seq", "candidate")): + seq = batch[seq_name] + raw_i64 = seq.get("int64_features") + post_hashes = seq.get("post_hashes") + if raw_i64 is None or post_hashes is None: + continue + + i64 = cast_jax(raw_i64) + valid = (cast_jax(post_hashes)[..., 0] > 0).astype(jnp.float32) + n_valid = jnp.maximum(jnp.sum(valid), 1.0) + raw_by_name: dict[str, jax.Array] = {} + + for cat_feat_enum, int64_feat_enum in ENGAGEMENT_COUNT_BUCKET_MAP: + name = cat_feat_enum.name.removesuffix("CountBucketSeq").lower() + max_bucket = ENGAGEMENT_COUNT_MAX_BUCKET.get( + cat_feat_enum, ENGAGEMENT_COUNT_NUM_BUCKETS - 1 + ) + raw_counts = i64[:, :, int64_feat_enum.value] + raw_by_name[name] = jnp.maximum(raw_counts.astype(jnp.float32), 0.0) buckets = compute_engagement_count_bucket(raw_counts, max_bucket) - valid = (cast_jax(post_hashes)[..., 0] > 0).astype(jnp.float32) - n_valid = jnp.maximum(jnp.sum(valid), 1.0) - coverage = jnp.sum((buckets > 0).astype(jnp.float32) * valid) / n_valid - mean_bucket = jnp.sum(buckets.astype(jnp.float32) * valid) / n_valid - stats[f"ec/{name}/{side}/coverage"] = coverage - stats[f"ec/{name}/{side}/mean_bucket"] = mean_bucket + + _add_scalar_stats(name, side, raw_counts, valid) + stats[f"ec/{name}/{side}/mean_bucket"] = ( + jnp.sum(buckets.astype(jnp.float32) * valid) / n_valid + ) + _add_bucket_percentiles(name, side, buckets, valid, max_bucket + 1) + + non_view_engagement = ( + raw_by_name["fav"] + + raw_by_name["reply"] + + raw_by_name["repost"] + + raw_by_name["quote"] + ) + view = raw_by_name["view"] + _add_scalar_stats("engagement_no_view", side, non_view_engagement, valid) + _add_bucket_percentiles( + "engagement_no_view", + side, + compute_engagement_count_bucket( + non_view_engagement, ENGAGEMENT_COUNT_NUM_BUCKETS - 1 + ), + valid, + ENGAGEMENT_COUNT_NUM_BUCKETS, + ) + + view_missing_like = ((view <= 0.0) & (non_view_engagement > 0.0)).astype(jnp.float32) + stats[f"ec/view/{side}/missing_like_rate"] = ( + jnp.sum(view_missing_like * valid) / n_valid + ) + + bools = seq.get("bool_features") + if ( + side == "candidate" + and stale_on + and bools is not None + and bools.shape[-1] > BoolFeature.isStalePost14d.value + ): + is_stale = cast_jax(bools)[:, :, BoolFeature.isStalePost14d.value].astype( + jnp.float32 + ) + stats[f"ec/stale_post_14d/{side}/zeroed_frac"] = jnp.sum(is_stale * valid) / n_valid + return stats def compute_author_nsfw_metrics( @@ -1678,8 +1891,21 @@ def multi_hot_to_embeddings( emb_size: int, embed_init_scale: float, name: str, + dense: bool = False, ) -> tuple[jax.Array, jax.Array]: embed_init = hk.initializers.VarianceScaling(embed_init_scale, mode="fan_out") + if dense: + embedding_table = get_parameter( + name, + shape=[output_vocab_size, emb_size], + init=embed_init, + dtype=jnp.float32, + pspec=P(), + rms_clip_axes=(-2, -1), + ) + output = jnp.dot(input.astype(embedding_table.dtype), embedding_table) + output = output.astype(DTYPE_BY_NAME[self.config.fprop_dtype]) + return output, embedding_table embedding_table = get_parameter( name, shape=[ @@ -1907,6 +2133,19 @@ def _get_unembedding(self) -> jax.Array: self.unembed_mat: jax.Array = with_sharding_constraint(unembed_mat, out_pspec) return unembed_mat + @hk.transparent + def _get_dpa_product_embedding_table(self) -> jax.Array: + _config = self.config + dim = _config.dpa_product_embed_dim + return get_parameter( + "dpa_product_embedding_table", + [_config.dpa_product_table_size, dim], + dtype=jnp.float32, + init=hk.initializers.Constant(0.0), + pspec=P(), + rms_clip_axes=(-2, -1), + ) + @hk.transparent def decode(self, inputs: jax.Array) -> jax.Array: unembeddings = self._get_unembedding() @@ -2067,6 +2306,41 @@ def _build_user_embedding( ) return user_embeddings, user_padding_mask + @hk.transparent + def _maybe_add_dpa_input_embedding( + self, + candidate_embeddings: jax.Array, + recsys_features_batch: RecsysFeaturesBatch, + ) -> jax.Array: + _config = self.config + if not _config.enable_dpa_input_embedding: + return candidate_embeddings + table = self._get_dpa_product_embedding_table() + dpa_dim = _config.dpa_product_embed_dim + embed_init = hk.initializers.VarianceScaling(_config.embed_init_scale, mode="fan_out") + proj = get_parameter( + "dpa_input_proj", + [dpa_dim, _config.emb_table_width], + dtype=jnp.float32, + init=lambda shape, dtype: embed_init(list(reversed(shape)), dtype).T, + pspec=P(None, None), + lr_multiplier=_config.model_config.scale_config.emb_lr_multiplier(dpa_dim), + ) + raw_i64 = recsys_features_batch["candidate_seq"].get("int64_features") + if raw_i64 is None: + return candidate_embeddings + keys = cast_jax(raw_i64)[:, :, _DPA_PRODUCT_KEY_SLOTS].astype(jnp.int32) + assert keys.shape[:2] == candidate_embeddings.shape[:2], ( + f"dpa keys shape {keys.shape} must match candidate embeddings " + f"{candidate_embeddings.shape[:2]}" + ) + valid = keys != 0 + ids = jnp.where(valid, jnp.clip(keys, 1, _config.dpa_product_table_size - 1), 0) + rows = jnp.take(table, ids, axis=0) + product_emb = jnp.where(valid[..., None], rows, 0.0).sum(axis=2) + emb = jnp.dot(product_emb.astype(proj.dtype), proj) + return candidate_embeddings + emb.astype(candidate_embeddings.dtype) + @hk.transparent def build_inputs( self, @@ -2078,6 +2352,9 @@ def build_inputs( assert _config.model_config.output_vocab_size is not None, "output_vocab_size is required" if _config.feature_prep_enabled: + assert not _config.enable_dpa_input_embedding, ( + "enable_dpa_input_embedding is not implemented for the feature_prep path" + ) fp = _config.feature_prep scale_multiplier = fp.scale_config.input_scale(fp.emb_size) tokens, padding_mask, candidate_start_offset = build_feature_prep_inputs( @@ -2147,8 +2424,9 @@ def build_inputs( ca = cast_jax(history_continuous_actions) if ca.shape[-1] > 1: history_dwell_time = ca[:, :, 1] - if _config.concat_history_bridge_prob and ca.shape[-1] > 0: - history_bridge_prob = ca[:, :, 0] + _bridge_idx = recsys_pb2.ContinuousActionName.BRIDGE_PROBABILITY + if _config.concat_history_bridge_prob and ca.shape[-1] > _bridge_idx: + history_bridge_prob = ca[:, :, _bridge_idx] if ctx_config.enable_engagement_counts: for cat_feat_enum, int64_feat_enum in ENGAGEMENT_COUNT_BUCKET_MAP: @@ -2204,6 +2482,7 @@ def build_inputs( _config.emb_table_width, _config.embed_init_scale, "action_embedding_table", + dense=_config.use_dense_action_table, ) multimodal_embeddings: jax.Array | None = None @@ -2290,6 +2569,9 @@ def build_inputs( search_query_embedding_dim=self.config.search_query_embedding_dim, fprop_dtype=DTYPE_BY_NAME[self.config.fprop_dtype], ) + candidate_embeddings = self._maybe_add_dpa_input_embedding( + candidate_embeddings, recsys_features_batch + ) user_features_token = None if _config.user_features.has_user_features: @@ -2378,6 +2660,16 @@ def _embed_post_sid(seq_name: str) -> jnp.ndarray | None: ) else: sids_jax = cast_jax(sids_in) + if _config.sid_embedding_mode == "recon": + return reconstruct_entity_sid( + sids_jax, + _config.emb_table_width, + _config.sid_decoder_path, + self.config.model_config.scale_config.emb_lr_multiplier, + self.config.embed_init_scale, + DTYPE_BY_NAME[_config.fprop_dtype], + "post", + ) needs_hashes = _config.sid_hash_level return embed_entity_sid( sids_jax, @@ -2396,6 +2688,7 @@ def _embed_post_sid(seq_name: str) -> jnp.ndarray | None: _sid_post_emb_h = _embed_post_sid("history_seq") if _config.use_post_sid else None _sid_post_emb_c = _embed_post_sid("candidate_seq") if _config.use_post_sid else None + _sid_recon = _config.use_post_sid and _config.sid_embedding_mode == "recon" history_embeddings, history_padding_mask = block_history_reduce( history_post_hashes, @@ -2408,9 +2701,11 @@ def _embed_post_sid(seq_name: str) -> jnp.ndarray | None: self.config.model_config.scale_config.emb_lr_multiplier, self.config.embed_init_scale, history_unified_context is not None, - sid_post_embeddings=_sid_post_emb_h, + sid_post_embeddings=None if _sid_recon else _sid_post_emb_h, history_bridge_prob=history_bridge_prob, ) + if _sid_recon and _sid_post_emb_h is not None: + history_embeddings += _sid_post_emb_h.astype(history_embeddings.dtype) if ( self.config.safety_filter_mode == "hard" @@ -2440,7 +2735,12 @@ def _embed_post_sid(seq_name: str) -> jnp.ndarray | None: candidate_search_query_embeddings=candidate_search_query_embeddings, search_query_embedding_dim=self.config.search_query_embedding_dim, fprop_dtype=DTYPE_BY_NAME[self.config.fprop_dtype], - sid_post_embeddings=_sid_post_emb_c, + sid_post_embeddings=None if _sid_recon else _sid_post_emb_c, + ) + if _sid_recon and _sid_post_emb_c is not None: + candidate_embeddings += _sid_post_emb_c.astype(candidate_embeddings.dtype) + candidate_embeddings = self._maybe_add_dpa_input_embedding( + candidate_embeddings, recsys_features_batch ) if _config.user_features.has_user_features: @@ -2596,6 +2896,12 @@ def loss( line_item_objective = ( cast_jax(raw_line_item_objective) if raw_line_item_objective is not None else None ) + raw_candidate_int64 = batch["candidate_seq"].get("int64_features") + dpa_product_key = ( + cast_jax(raw_candidate_int64)[:, :, _DPA_PRODUCT_KEY_SLOTS] + if raw_candidate_int64 is not None + else None + ) candidate_safety_mask = ( cast_jax(candidate_safety_mask) if candidate_safety_mask is not None else None ) @@ -2623,6 +2929,23 @@ def loss( else: raw_weights = jnp.broadcast_to(sample_weights, targets.shape[:2]) + delayed_mask: jax.Array | None = None + sample_source = batch.get("sample_source") + if sample_source is not None: + sample_source = cast_jax(sample_source).astype(jnp.float32) + if self.config.use_seqpack: + delayed_mask = jnp.repeat( + sample_source.squeeze(-1), packed_candidate_seq_len, axis=1 + ) + else: + delayed_mask = sample_source + if self.config.split_head_training_by_source: + assert delayed_mask is not None, ( + "split_head_training_by_source=True requires the sample_source batch " + "field (from the is_delayed_feedback column); training unsplit " + "silently would defeat the flag" + ) + if self.config.log_q_correction: tweet_counts = get_candidate_tweet_counts( batch, self.config.log_q_num_bins, negative_sample_mask @@ -2708,6 +3031,7 @@ def loss( client_app_id, line_item_objective, candidate_safety_mask, + dpa_product_key, ) = pad_to_next_128_multiple( input_embeddings, padding_mask, @@ -2720,6 +3044,7 @@ def loss( client_app_id, line_item_objective, candidate_safety_mask, + dpa_product_key, ) idx = jnp.arange(padding_mask.shape[1], dtype=jnp.int32)[None, :] @@ -2765,6 +3090,11 @@ def loss( ) target_padding_mask = padding_mask[:, candidate_start_offset:] + trained_candidate_mask = batch["candidate_seq"].get("trained_candidate_mask") + if trained_candidate_mask is not None: + keep = cast_jax(trained_candidate_mask) + pad_len = target_padding_mask.shape[1] - keep.shape[1] + target_padding_mask = target_padding_mask & jnp.pad(keep, ((0, 0), (0, pad_len))) history_padding = padding_mask[ :, @@ -2831,6 +3161,39 @@ def loss( conv_zero_mask = no_click[:, :, None] * conv_head_mask loss_mask = loss_mask * (1 - conv_zero_mask) + vt_head_mask = jnp.zeros(num_actions).at[jnp.array(VIEW_THROUGH_ACTION_INDICES)].set(1.0) + if self.config.train_view_through_heads: + vt_zero = targets[:, :, CLICK_ACTION_INDEX][:, :, None] * vt_head_mask + else: + vt_zero = vt_head_mask + loss_mask = loss_mask * (1 - vt_zero) + + if self.config.mact_in_app_loss_weight != 1.0: + mact_w = ( + jnp.ones(num_actions) + .at[jnp.array(MACT_IN_APP_LOSS_ACTION_INDICES)] + .set(self.config.mact_in_app_loss_weight) + ) + loss_mask = loss_mask * mact_w + + if self.config.condition_search_relevance_on_prompt: + prompt_shown = jnp.any(targets[:, :, SEARCH_RELEVANCE_ACTION_INDICES], axis=-1) + no_prompt = 1 - prompt_shown + search_head_mask = ( + jnp.zeros(num_actions).at[jnp.array(SEARCH_RELEVANCE_ACTION_INDICES)].set(1.0) + ) + search_zero_mask = no_prompt[:, :, None] * search_head_mask + loss_mask = loss_mask * (1 - search_zero_mask) + + if self.config.split_head_training_by_source and delayed_mask is not None: + conv_head_split_mask = ( + jnp.zeros(num_actions).at[jnp.array(SOURCE_SPLIT_CONVERSION_HEAD_INDICES)].set(1.0) + ) + eng_head_mask = 1.0 - conv_head_split_mask + is_delayed = delayed_mask[:, :, None] + loss_mask = loss_mask * (1 - is_delayed * eng_head_mask) + loss_mask = loss_mask * (1 - (1 - is_delayed) * conv_head_split_mask) + safety_stats = safety_filter_stats( candidate_safety_mask, target_padding_mask, @@ -2865,6 +3228,7 @@ def loss( 1.0 if self.config.safety_filter_mode != "off" else 0.0 ) stats["origin-loss"] = loss + stats["mact-in-app-loss-weight"] = jnp.float32(self.config.mact_in_app_loss_weight) stats = self.compute_recsys_metrics( raw_targets=targets, @@ -2877,6 +3241,8 @@ def loss( new_user_mask=new_user_mask, line_item_objective=line_item_objective, no_history_mask=no_history_mask, + dpa_product_key=dpa_product_key[..., 0] if dpa_product_key is not None else None, + delayed_sample_mask=delayed_mask, stats=stats, rce_ema=rce_ema, rce_alpha=rce_alpha, @@ -2911,14 +3277,22 @@ def loss( product_surface=product_surface, new_user_mask=new_user_mask, no_history_mask=no_history_mask, + dpa_product_key=dpa_product_key[..., 0] if dpa_product_key is not None else None, + delayed_sample_mask=delayed_mask, ) + continuous_base_mask = target_padding_mask + if self.config.split_head_training_by_source and delayed_mask is not None: + continuous_base_mask = continuous_base_mask * ( + 1 - delayed_mask.astype(continuous_base_mask.dtype) + ) + for loss_config in self.config.continuous_action_losses: if loss_config.loss_weight > 0 and loss_config.action_index < data_num_continuous: gt_raw = candidate_continuous_actions[:, :, loss_config.action_index] pred_raw = candidate_continuous_preds[:, :, loss_config.action_index] - head_valid_mask = target_padding_mask + head_valid_mask = continuous_base_mask if loss_config.product_surfaces or loss_config.exclude_product_surfaces: head_surface_mask = _get_surface_mask(loss_config, product_surface) head_valid_mask = head_valid_mask * head_surface_mask.astype( diff --git a/phoenix/xrex/models/recsys_sid.py b/phoenix/xrex/models/recsys_sid.py new file mode 100644 index 00000000..7dbc5b17 --- /dev/null +++ b/phoenix/xrex/models/recsys_sid.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +from __future__ import annotations + +import functools +from typing import Callable + +import haiku as hk +import jax +import jax.numpy as jnp +from jax.sharding import PartitionSpec as P +from numpy import typing as npt + +from xrex.models.layers import get_parameter + + +@functools.lru_cache(maxsize=2) +def _load_sid_decoder( + path: str, +) -> tuple[npt.NDArray, tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray]]: + import safetensors.numpy + + tensors = safetensors.numpy.load_file(path) + return tensors["stages"], ( + tensors["dec_w0"], + tensors["dec_b0"], + tensors["dec_w1"], + tensors["dec_b1"], + ) + + +def reconstruct_entity_sid( + sids: jnp.ndarray, + target_dim: int, + decoder_path: str, + lr_multiplier_func: Callable[[int], float], + embed_init_scale: float, + fprop_dtype: jnp.dtype, + name_prefix: str, +) -> jnp.ndarray: + stages_np, (w0_np, b0_np, w1_np, b1_np) = _load_sid_decoder(decoder_path) + num_levels, codebook_size, _ = stages_np.shape + assert sids.shape[-1] == num_levels, ( + f"post_sids have {sids.shape[-1]} levels but the decoder at {decoder_path!r} " + f"expects {num_levels}; set sid_num_levels to match the artifact." + ) + + stages = jnp.asarray(stages_np) + w0, b0 = jnp.asarray(w0_np), jnp.asarray(b0_np) + w1, b1 = jnp.asarray(w1_np), jnp.asarray(b1_np) + + codes = sids.astype(jnp.int32) - 1 + missing = sids[..., 0] == 0 + codes = jnp.clip(codes, 0, codebook_size - 1) + quant = stages[jnp.arange(num_levels), codes].sum(-2) + h = jax.nn.relu(quant @ w0 + b0) + recon = h @ w1 + b1 + recon = recon * jax.lax.rsqrt((recon**2).sum(-1, keepdims=True) + 1e-12) + recon = jnp.where(missing[..., None], 0.0, recon) + recon = jax.lax.stop_gradient(recon) + + recon_dim = recon.shape[-1] + embed_init = hk.initializers.VarianceScaling(1.0, mode="fan_out") + proj = get_parameter( + f"{name_prefix}_sid_recon_proj", + [recon_dim, target_dim], + dtype=jnp.float32, + init=lambda shape, dtype: embed_init(list(reversed(shape)), dtype).T, + pspec=P(None, None), + lr_multiplier=lr_multiplier_func(recon_dim), + ) + return jnp.dot(recon, proj).astype(fprop_dtype) diff --git a/phoenix/xrex/models/recsys_two_tower_model.py b/phoenix/xrex/models/recsys_two_tower_model.py index e142e9c2..dcf274cd 100644 --- a/phoenix/xrex/models/recsys_two_tower_model.py +++ b/phoenix/xrex/models/recsys_two_tower_model.py @@ -57,6 +57,7 @@ ) from xrex.models.scaling import ScaleConfig from xrex.models.sharding_context import ShardingContext +from xrex.pallas.ranker_attention_utils import HISTORY_SEGMENT_ID from xrex.train.misc import PostEmbeddings from xrex.utils.utils import Summary @@ -66,6 +67,12 @@ INF = 1e12 +def user_tower_segment_ids(batch: int, seq_len: int, *, use_history_segment_ids: bool) -> jax.Array: + if use_history_segment_ids: + return jnp.full((batch, seq_len), HISTORY_SEGMENT_ID, dtype=jnp.int32) + return jnp.zeros((batch, seq_len), dtype=jnp.int32) + + def _l2_normalize_candidates(embeddings: jax.Array) -> jax.Array: norm_sq = jnp.sum(embeddings**2, axis=-1, keepdims=True) norm = jnp.sqrt(jnp.maximum(norm_sq, EPS)) @@ -1245,23 +1252,23 @@ def _pool(outputs, mask, cu_seqlens): user_representation, P(self.data_axis, None) ) else: - user_embeddings, user_padding_mask, _, _, _, _, _, _, _, _, _ = ( - pad_to_next_128_multiple( - user_embeddings, - user_padding_mask, - jnp.zeros_like(user_padding_mask), - None, - None, - None, - None, - None, - None, - None, - ) + user_embeddings, user_padding_mask, *_ = pad_to_next_128_multiple( + user_embeddings, + user_padding_mask, + jnp.zeros_like(user_padding_mask), + None, + None, + None, + None, + None, + None, + None, ) B, T = user_padding_mask.shape - user_segment_ids = jnp.zeros((B, T), dtype=jnp.int32) + user_segment_ids = user_tower_segment_ids( + B, T, use_history_segment_ids=self.config.use_history_segment_ids + ) if self.config.user_tower_config.right_anchored_rope: user_positions = right_anchored_rope_positions( @@ -1584,7 +1591,7 @@ class RecsysTwoTowerModelConfig(Config): ads_only_candidates: bool = False - num_continuous_actions: int = 0 + use_history_segment_ids: bool = False multimodal_embedding_type: EmbeddingType | None = None diff --git a/phoenix/xrex/models/remat.py b/phoenix/xrex/models/remat.py index bd4be3be..a321610e 100644 --- a/phoenix/xrex/models/remat.py +++ b/phoenix/xrex/models/remat.py @@ -8,6 +8,7 @@ class RematType(enum.IntEnum): WHOLE = 0 SAVE_GB300_RECSYS = 29 + SAVE_H100_RECSYS = 32 def custom_remat_policy(policy: RematType): @@ -28,5 +29,21 @@ def custom_remat_policy(policy: RematType): "value_heads", "scalar_stats", ) + elif policy == RematType.SAVE_H100_RECSYS: + return jax.checkpoint_policies.save_only_these_names( + "attn_outputs", + "dense_outputs", + "dense_outputs_individual", + "attn", + "gate_up_proj", + "dense_up_proj", + "query_heads_rope", + "key_heads_rope", + "cutedsl_attn_outputs", + "query_heads", + "key_heads", + "value_heads", + "scalar_stats", + ) else: raise NotImplementedError(f"Unknown remat policy: {policy}") diff --git a/phoenix/xrex/models/transformer.py b/phoenix/xrex/models/transformer.py index e29d6578..d69517bd 100644 --- a/phoenix/xrex/models/transformer.py +++ b/phoenix/xrex/models/transformer.py @@ -57,7 +57,6 @@ class FeedForwardConfig(Config): with_bias: bool = False force_ffn_base_size: bool = True base_ffn_size: int = 32 - merge_gate_up: bool = False diff --git a/phoenix/xrex/optimizers/recsys/__init__.py b/phoenix/xrex/optimizers/recsys/__init__.py index c054b2a7..e919699e 100644 --- a/phoenix/xrex/optimizers/recsys/__init__.py +++ b/phoenix/xrex/optimizers/recsys/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. from xrex.optimizers.recsys.config import RecsysEmbeddingOptimConfig +from xrex.optimizers.recsys.dense_optim import RecsysDenseOptimConfig from xrex.optimizers.recsys.protocol import ( RecsysEmbeddingOptimizer, _lookup, diff --git a/phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py b/phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py index f7fa9d94..39e8171f 100644 --- a/phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py +++ b/phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py @@ -28,7 +28,7 @@ def gradient_update_start( table: jax.Array, state: Any, gate: jax.Array, - ) -> tuple[tuple[jax.Array, ...], jax.Array, Any]: + ) -> tuple[tuple[jax.Array, ...], jax.Array, Any, dict[str, jax.Array]]: raise NotImplementedError( f"use_async_emb requires an embedding optimizer implementing " f"AsyncEmbOptimizer, and {type(self).__name__} has no fused table update" diff --git a/phoenix/xrex/optimizers/recsys/dense_optim.py b/phoenix/xrex/optimizers/recsys/dense_optim.py new file mode 100644 index 00000000..8386026a --- /dev/null +++ b/phoenix/xrex/optimizers/recsys/dense_optim.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +from __future__ import annotations + +from typing import Optional + +import optax + +from xai_configlib import configclass +from xrex.optimizers.optim import ( + OptimConfig, + _instantiate, + inject_hyperparams, + scale_by_lr_multiplier, +) +from xrex.optimizers.recsys.muon import make_muon_optimizer + + +@configclass +class RecsysDenseOptimConfig(OptimConfig): + muon_beta: float = 0.95 + muon_ns_steps: int = 5 + muon_adam_patterns: str = "embed,_emb,logits,vocab,norm,table" + muon_consistent_rms: Optional[float] = None + muon_ns_dtype: str = "bfloat16" + muon_preconditioning: str = "frobenius" + muon_split_fused: str = "" + muon_matrix_weight_decay: float = 0.0 + adam_embedding_weight_decay: float = 0.0 + adam_embedding_decay_patterns: str = "embed,_emb,logits,vocab,table" + + def make(self): + if self.optim != "muon": + return super().make() + + optimizer = make_muon_optimizer(self) + + @inject_hyperparams + def schedule_optim(learning_rate, b1, b2, weight_decay): + return optax.chain( + optax.clip_by_global_norm(self.clip_by_global_norm), + optimizer(learning_rate, b1=b1, b2=b2, weight_decay=weight_decay), + scale_by_lr_multiplier(), + ) + + schedule_args = map(_instantiate, (self.learning_rate, self.b1, self.b2, self.weight_decay)) + return schedule_optim(*schedule_args) diff --git a/phoenix/xrex/optimizers/recsys/muon.py b/phoenix/xrex/optimizers/recsys/muon.py new file mode 100644 index 00000000..91bd5795 --- /dev/null +++ b/phoenix/xrex/optimizers/recsys/muon.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import functools + +import jax +import jax.numpy as jnp +import optax +from optax import contrib as optax_contrib +from optax._src import numerics as _optax_numerics + +from optax.contrib._muon import _DEFAULT_NS_COEFFS, orthogonalize_via_newton_schulz + +from xrex.models.model_utils import Parameter + +_VECTOR_LEAF_NAMES = frozenset({"b", "bias"}) + + +def _leaf_name(path) -> str: + if not path: + return "" + last = path[-1] + key = getattr(last, "key", None) or getattr(last, "name", None) + return str(key).lower() if key is not None else "" + + +def _is_matrix_group(path, x, adam_patterns: tuple[str, ...]) -> bool: + if not hasattr(x, "ndim") or x.ndim < 2: + return False + if _leaf_name(path) in _VECTOR_LEAF_NAMES: + return False + name = jax.tree_util.keystr(path).lower() + return not any(s in name for s in adam_patterns) + + +def muon_dimension_numbers( + params, adam_patterns: tuple[str, ...], split_patterns: tuple[str, ...] = () +): + assignments: list[tuple[str, str, tuple[int, ...]]] = [] + saw_masked = False + + def _label(path, p): + nonlocal saw_masked + name = jax.tree_util.keystr(path).lower() + x = p.x if isinstance(p, Parameter) else p + if not hasattr(x, "ndim"): + saw_masked = True + return None + if isinstance(p, Parameter): + assert p.weight_decay_mask == 0.0, ( + f"optim='muon' ignores weight_decay_mask, but {name} has " + f"weight_decay_mask={p.weight_decay_mask}; use " + "muon_matrix_weight_decay / adam_embedding_weight_decay." + ) + is_muon = _is_matrix_group(path, x, adam_patterns) + assignments.append(("muon" if is_muon else "adam", name, tuple(x.shape))) + if not is_muon: + return None + return optax_contrib.MuonDimensionNumbers(reduction_axis=-2, output_axis=-1) + + dim_nums = jax.tree.map_with_path(_label, params, is_leaf=lambda v: isinstance(v, Parameter)) + if not saw_masked and any(g == "muon" for g, _, _ in assignments): + for pattern in split_patterns: + assert any(g == "muon" and pattern in name for g, name, _ in assignments), ( + "muon_split_fused pattern %r matched no muon matrix" % pattern + ) + return dim_nums + + +def _cast_muon_group(updates, adam_patterns: tuple[str, ...], dtype): + def _cast(path, u): + x = u.x if isinstance(u, Parameter) else u + if not _is_matrix_group(path, x, adam_patterns): + return u + return jax.tree.map(lambda a: a.astype(dtype), u) + + return jax.tree.map_with_path(_cast, updates, is_leaf=lambda v: isinstance(v, Parameter)) + + +def embedding_decay_mask(emb_patterns: tuple[str, ...]): + def mask_fn(params): + def _m(path, p): + x = p.x if isinstance(p, Parameter) else p + if not hasattr(x, "ndim"): + return False + name = jax.tree_util.keystr(path).lower() + if "norm" in name or _leaf_name(path) in _VECTOR_LEAF_NAMES: + return False + return any(s in name for s in emb_patterns) + + return jax.tree.map_with_path(_m, params, is_leaf=lambda v: isinstance(v, Parameter)) + + return mask_fn + + +def parse_split_fused(spec: str) -> dict[str, int]: + out: dict[str, int] = {} + for item in spec.split(","): + item = item.strip() + if not item: + continue + pattern, _, dim = item.partition(":") + out[pattern.strip().lower()] = int(dim) + return out + + +def _split_block_dim(path, split_spec: dict[str, int]): + name = jax.tree_util.keystr(path).lower() + for pattern, dim in split_spec.items(): + if pattern in name: + return dim + return None + + +def scale_by_muon_shaped( + cfg, + dim_nums_for_partition, + adam_patterns: tuple[str, ...], + split_spec: dict[str, int], +) -> optax.GradientTransformation: + beta = cfg.muon_beta + ns_steps = cfg.muon_ns_steps + preconditioning = cfg.muon_preconditioning + consistent_rms = cfg.muon_consistent_rms + ns_dtype = jnp.dtype(cfg.muon_ns_dtype) if cfg.muon_ns_dtype else None + mu_dtype = ns_dtype or cfg.mu_dtype + ns_coeffs = jnp.asarray(_DEFAULT_NS_COEFFS) + + def _orthogonalize(path, x, dim_num): + block = _split_block_dim(path, split_spec) + if block is not None: + assert x.shape[-1] % block == 0, ( + f"muon_split_fused: block {block} does not divide output axis " + f"{x.shape[-1]} of {jax.tree_util.keystr(path)}" + ) + if block is not None and x.shape[-1] > block: + n_blocks = x.shape[-1] // block + xs = x.reshape(x.shape[:-1] + (n_blocks, block)) + dn = optax_contrib.MuonDimensionNumbers(reduction_axis=-3, output_axis=-1) + o = orthogonalize_via_newton_schulz(xs, ns_coeffs, ns_steps, preconditioning, 1e-8, dn) + fan_in, fan_out = x.shape[-2], block + o = o.reshape(x.shape) + else: + o = orthogonalize_via_newton_schulz( + x, ns_coeffs, ns_steps, preconditioning, 1e-8, dim_num + ) + fan_in, fan_out = x.shape[-2], x.shape[-1] + if consistent_rms is not None: + return o * (jnp.sqrt(jnp.maximum(fan_in, fan_out)) * consistent_rms) + return o * jnp.sqrt(jnp.maximum(1, fan_out / fan_in)) + + def init_fn(params): + mu = optax.tree.zeros_like(params, dtype=mu_dtype) + return optax_contrib.MuonState(count=jnp.zeros([], jnp.int32), mu=mu, ns_coeffs=ns_coeffs) + + def update_fn(updates, state, params=None): + del params + dim_nums = dim_nums_for_partition(updates) + mu = optax.tree.update_moment(updates, state.mu, beta, 1) + count_inc = _optax_numerics.safe_increment(state.count) + mu_hat = jax.tree.map( + lambda m, g: beta * m + (1 - beta) * g, + optax.tree.bias_correction(mu, beta, _optax_numerics.safe_increment(count_inc)), + optax.tree.bias_correction(updates, beta, count_inc), + ) + new_updates = jax.tree_util.tree_map_with_path( + lambda path, x, dn: _orthogonalize(path, x, dn), + mu_hat, + dim_nums, + is_leaf=lambda v: isinstance(v, optax_contrib.MuonDimensionNumbers), + ) + mu = optax.tree.cast(mu, mu_dtype) + return new_updates, optax_contrib.MuonState( + count=count_inc, mu=mu, ns_coeffs=state.ns_coeffs + ) + + return optax.GradientTransformation(init_fn, update_fn) + + +def make_muon_optimizer(cfg): + adam_patterns = tuple(s.strip().lower() for s in cfg.muon_adam_patterns.split(",") if s.strip()) + emb_patterns = tuple( + s.strip().lower() for s in cfg.adam_embedding_decay_patterns.split(",") if s.strip() + ) + ns_dtype = jnp.dtype(cfg.muon_ns_dtype) if cfg.muon_ns_dtype else None + split_spec = parse_split_fused(getattr(cfg, "muon_split_fused", "")) + dim_nums_fn = functools.partial( + muon_dimension_numbers, + adam_patterns=adam_patterns, + split_patterns=tuple(split_spec), + ) + + _is_dim_nums = lambda x: isinstance(x, optax_contrib.MuonDimensionNumbers) + + def param_labels(params): + dim_nums = dim_nums_fn(params) + populate = lambda dn, x: jax.tree.map(lambda _: "muon" if dn is not None else "adam", x) + return jax.tree.map( + populate, dim_nums, params, is_leaf=lambda x: x is None or _is_dim_nums(x) + ) + + def dim_nums_for_partition(params): + dim_nums = dim_nums_fn(params) + mask = jax.tree.map(lambda label: label == "muon", param_labels(params)) + is_leaf = lambda x: ( + x is None or _is_dim_nums(x) or isinstance(x, optax.transforms.MaskedNode) + ) + populate = lambda dn, submask: jax.tree.map( + lambda m: dn if m else optax.transforms.MaskedNode(), submask + ) + return jax.tree.map(populate, dim_nums, mask, is_leaf=is_leaf) + + def optimizer(learning_rate, b1, b2, weight_decay): + del weight_decay + mu_dtype = ns_dtype or cfg.mu_dtype + muon_chain = optax.chain( + scale_by_muon_shaped(cfg, dim_nums_for_partition, adam_patterns, split_spec), + optax.add_decayed_weights(cfg.muon_matrix_weight_decay), + optax.scale_by_learning_rate(learning_rate), + ) + adam = optax.adamw( + learning_rate=learning_rate, + b1=b1, + b2=b2, + eps=1e-8, + eps_root=0.0, + mu_dtype=mu_dtype, + weight_decay=cfg.adam_embedding_weight_decay, + mask=embedding_decay_mask(emb_patterns), + nesterov=False, + ) + base = optax.transforms.partition({"muon": muon_chain, "adam": adam}, param_labels) + + def wrapped_update(updates, state, params=None): + if ns_dtype is not None: + updates = _cast_muon_group(updates, adam_patterns, ns_dtype) + return base.update(updates, state, params) + + return optax.GradientTransformation(base.init, wrapped_update) + + return optimizer diff --git a/phoenix/xrex/optimizers/recsys/rowwise_adagrad.py b/phoenix/xrex/optimizers/recsys/rowwise_adagrad.py index 86ba9440..0e4ef23b 100644 --- a/phoenix/xrex/optimizers/recsys/rowwise_adagrad.py +++ b/phoenix/xrex/optimizers/recsys/rowwise_adagrad.py @@ -99,10 +99,8 @@ def gradient_update_start( state: RecsysRowwiseAdagradState, gate: jax.Array, ) -> tuple[tuple[jax.Array, ...], jax.Array, RecsysRowwiseAdagradState]: - if self._lazy_decay: - raise ValueError("the fused rowwise Adagrad update does not support lazy decay") - if self._weight_decay > 0.0: - raise ValueError("the fused rowwise Adagrad update does not support weight decay") + if self._lazy_decay and (state.step is None or state.last_step is None): + raise ValueError("fused lazy decay needs timestamped state (step/last_step)") if 32 % context.shard_width != 0: raise ValueError( @@ -112,6 +110,87 @@ def gradient_update_start( from xrex.cuda.async_emb import async_emb + metrics: dict[str, jax.Array] = {} + + if self._lazy_decay: + + @shard_map( + mesh=context.mesh, + in_specs=( + P(context.data_axis, None), + P(context.data_axis), + P(), + P(None, context.table_axis), + P(), + P(), + P(), + P(), + P(context.data_axis, None), + ), + out_specs=( + P(context.data_axis, None), + P(context.data_axis), + P(), + P(None, context.table_axis), + P(), + P(), + ), + check_vma=False, + ) + def start_lazy( + grads: jax.Array, + segment_ids: jax.Array, + unique_tokens: jax.Array, + table: jax.Array, + accum: jax.Array, + last_step: jax.Array, + step: jax.Array, + pending: jax.Array, + gate: jax.Array, + ) -> tuple[jax.Array, ...]: + return tuple( + async_emb.rowwise_adagrad_lazy_update_start( + grads, + segment_ids, + unique_tokens, + table, + accum, + last_step, + step, + pending, + gate, + context, + learning_rate=self._learning_rate, + eps=self._eps, + accum_decay_rate=self._decay_rate or 0.0, + weight_decay_rate=self._weight_decay, + ) + ) + + assert state.step is not None and state.last_step is not None + grads_pin, segments_pin, tokens_pin, table_out, accum_out, last_step_out = start_lazy( + update.grads, + update.segment_ids, + update.unique_tokens, + table, + state.row_sum_sq["table"], + state.last_step["table"], + state.step, + update.pending, + gate, + ) + new_step = state.step + update.pending.astype(jnp.int32).reshape(()) + return ( + (grads_pin, segments_pin, tokens_pin), + table_out, + RecsysRowwiseAdagradState( + row_sum_sq={**state.row_sum_sq, "table": accum_out}, + step=new_step, + last_step={**state.last_step, "table": last_step_out}, + ), + metrics, + ) + @shard_map( mesh=context.mesh, in_specs=( @@ -154,6 +233,7 @@ def start( learning_rate=self._learning_rate, eps=self._eps, decay_factor=self.decay_factor, + weight_decay_factor=math.exp(-self._weight_decay), ) ) return grads_pin, segments_pin, tokens_pin, table_out, accum_out @@ -171,6 +251,7 @@ def start( (grads_pin, segments_pin, tokens_pin), table_out, state._replace(row_sum_sq={**state.row_sum_sq, "table": accum_out}), + metrics, ) def gradient_update_done( diff --git a/phoenix/xrex/settings.py b/phoenix/xrex/settings.py index 5e9873f9..8a9dd92e 100644 --- a/phoenix/xrex/settings.py +++ b/phoenix/xrex/settings.py @@ -70,7 +70,7 @@ GCS_MIRROR_BUCKET: str = os.environ.get("XREX_GCS_MIRROR_BUCKET", "") GCS_MIRROR_PREFIX: str = os.environ.get("GCS_PREFIX", "") -ADS_S3_ENV_PREFIX: str = os.environ.get("XREX_ADS_S3_ENV_PREFIX", "") +O2_ENV_PREFIX: str = os.environ.get("XREX_O2_ENV_PREFIX", "") PHOENIX_INDEX_BASE: str = os.environ.get("PHOENIX_INDEX_BASE", "phoenix_index") diff --git a/phoenix/xrex/train/checkpoint_write.py b/phoenix/xrex/train/checkpoint_write.py index 94e98152..4e4232b4 100644 --- a/phoenix/xrex/train/checkpoint_write.py +++ b/phoenix/xrex/train/checkpoint_write.py @@ -66,6 +66,35 @@ def save_checkpoint( path = self.get_checkpoint_path(self.ctx) + encrypt_checkpoint = self.checkpoint_config.encrypt + if getattr(self, "_restored_encrypted", False) and not encrypt_checkpoint: + raise ValueError( + "restored from an ENCRYPTED checkpoint but checkpoint_config.encrypt is " + "false — refusing to downgrade to plaintext saves; set " + "checkpoint_config.encrypt=true with encryption_key_id/encryption_context" + ) + kms_client = None + if encrypt_checkpoint: + if self.checkpoint_config.encryption_key_id is None: + raise ValueError( + "checkpoint_config.encrypt requires checkpoint_config.encryption_key_id" + ) + missing_context = {"domain", "run"} - self.checkpoint_config.encryption_context.keys() + if missing_context: + raise ValueError( + "encrypted checkpoint context is missing " + ", ".join(sorted(missing_context)) + ) + import hashlib + + import xai_kms + + encryption_context = dict(self.checkpoint_config.encryption_context) + encryption_context["mint_nonce"] = hashlib.sha256(os.fsencode(path)).hexdigest() + kms_client = xai_kms.KmsClient.from_cluster_env( + key_id=self.checkpoint_config.encryption_key_id, + encryption_context=encryption_context, + ) + with tracer.start_as_current_span("write_checksum"): if self.checkpoint_config.write_checksums: if checksum_dict is None: @@ -84,6 +113,7 @@ def save_checkpoint( def _callback( elapsed_samples: int = self.elapsed_samples, elapsed_tokens: int = self.elapsed_tokens, + trainer_step: int = self.step, ) -> None: if checkpoint_index is not None: self.current_ckpt_index = checkpoint_index @@ -96,6 +126,7 @@ def _callback( elapsed_samples, self.current_ckpt_index, elapsed_tokens, + trainer_step, self.checkpoint_config.checkpoint_ttl, ) @@ -124,6 +155,9 @@ def _callback( compressed=self.checkpoint_config.checkpoint_compressed, chunk_byte_size=self.checkpoint_config.checkpoint_chunk_size_bytes, tracer=tracer, + save_concurrent_gb=self.checkpoint_config.save_concurrent_gb, + kms_client=kms_client, + encryption_chunk_size=self.checkpoint_config.encryption_chunk_size_bytes, ) diff --git a/phoenix/xrex/train/misc.py b/phoenix/xrex/train/misc.py index fba69fb8..0243c24e 100644 --- a/phoenix/xrex/train/misc.py +++ b/phoenix/xrex/train/misc.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. import datetime +from dataclasses import field from typing import Literal, NamedTuple, Protocol import haiku as hk @@ -20,6 +21,14 @@ class CheckpointConfig(Config): copy_port: int = 0 + copy_port_tls_cert: str = "" + copy_port_tls_key: str = "" + copy_port_tls_client_ca: str = "" + copy_port_tls_ca: str = "" + copy_port_tls_server_name: str = "" + copy_port_tls_client_cert: str = "" + copy_port_tls_client_key: str = "" + checkpoint_disk_every_s: int = 0 shm_max_entries: int = 5 @@ -44,6 +53,18 @@ class CheckpointConfig(Config): checkpoint_chunked: bool = True checkpoint_compressed: bool = True checkpoint_chunk_size_bytes: int = 1024 * 1024 * 4 + encrypt: bool = False + encryption_key_id: str | None = None + encryption_context: dict[str, str] = field(default_factory=dict) + encryption_chunk_size_bytes: int = 8 * 1024 * 1024 + + save_concurrent_gb: int | None = None + + restore_concurrent_gb: int | None = 32 + + restore_streamed: bool = False + + restore_window_gb: float | None = None checkpoint_ttl: int = datetime.timedelta(weeks=2).total_seconds() diff --git a/phoenix/xrex/train/recsys_bundle_export.py b/phoenix/xrex/train/recsys_bundle_export.py new file mode 100644 index 00000000..c7993d0e --- /dev/null +++ b/phoenix/xrex/train/recsys_bundle_export.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +from __future__ import annotations + +import dataclasses +import json +import logging +import re +import time +import typing +import zlib +from typing import Any, NamedTuple + +import jax +import numpy as np + +from xai_checkpointing.tree_util import tree_to_dict +from xrex.models.model_utils import unwrap_tree +from xrex.models.recsys_embedding import RecsysEmbeddings +from xrex.models.sharding_context import make_legacy_sharding_context + +if typing.TYPE_CHECKING: + from xrex.train.trainer_recsys import RecsysTrainer + +logger = logging.getLogger(__name__) + +BUNDLE_SCHEMA_VERSION = 2 +BUNDLE_DIR = "export" +MANIFEST_NAME = f"{BUNDLE_DIR}/MANIFEST.json" + + +def restamp_manifest(data: bytes) -> bytes: + manifest = json.loads(data) + manifest["created_timestamp"] = time.time() + return json.dumps(manifest, indent=2).encode() + + +class EmbeddingSlices(NamedTuple): + hist_post_end: int + hist_auth_end: int + cand_post_end: int + cand_auth_end: int + user_end: int + user_ip_end: int + + +class PackedGeometry(NamedTuple): + packed_history_len: int + packed_candidate_len: int + bs_per_device: int + merged_batch: int + + +@dataclasses.dataclass(frozen=True) +class BundleFile: + name: str + data: bytes + + +_parameter_serialization_registered = False + + +def _pspec_to_json(pspec: Any) -> list[Any]: + return [None if e is None else (e if isinstance(e, str) else list(e)) for e in tuple(pspec)] + + +def _pspec_from_json(entries: list[Any]) -> Any: + from jax.sharding import PartitionSpec + + return PartitionSpec( + *[None if e is None else (e if isinstance(e, str) else tuple(e)) for e in entries] + ) + + +def ensure_parameter_serialization_registered() -> None: + global _parameter_serialization_registered + if _parameter_serialization_registered: + return + from jax import export as jax_export + + from xrex.models.model_utils import Parameter + + def serialize_auxdata(aux: dict[str, Any]) -> bytes: + aux = dict(aux) + aux["pspec"] = _pspec_to_json(aux["pspec"]) + return json.dumps(aux, default=list).encode() + + def deserialize_auxdata(data: bytes) -> dict[str, Any]: + aux = json.loads(data) + aux["pspec"] = _pspec_from_json(aux["pspec"]) + if isinstance(aux.get("rms_clip_axes"), list): + aux["rms_clip_axes"] = tuple(aux["rms_clip_axes"]) + return aux + + jax_export.register_pytree_node_serialization( + Parameter, + serialized_name="xrex.models.model_utils.Parameter", + serialize_auxdata=serialize_auxdata, + deserialize_auxdata=deserialize_auxdata, + ) + _parameter_serialization_registered = True + + +_packing_layout_serialization_registered = False + + +def ensure_packing_layout_serialization_registered() -> None: + global _packing_layout_serialization_registered + if _packing_layout_serialization_registered: + return + from jax import export as jax_export + + from xrex.data.recsys.sequence_packing import SequencePackedLayout + + jax_export.register_pytree_node_serialization( + SequencePackedLayout, + serialized_name="xrex.data.recsys.sequence_packing.SequencePackedLayout", + serialize_auxdata=lambda aux: json.dumps(list(aux)).encode(), + deserialize_auxdata=lambda data: tuple(json.loads(data)), + ) + _packing_layout_serialization_registered = True + + +def _dtype_name(dtype: Any) -> str: + return np.dtype(dtype).name + + +def _aval_entry(leaf: Any) -> dict[str, Any]: + return {"shape": [int(d) for d in leaf.shape], "dtype": _dtype_name(leaf.dtype)} + + +def _to_shape_dtype_struct(tree: Any) -> Any: + return jax.tree.map( + lambda leaf: jax.ShapeDtypeStruct(leaf.shape, jax.dtypes.canonicalize_dtype(leaf.dtype)), + tree, + ) + + +def _make_export_config(trainer: RecsysTrainer, history_seq_len: int, candidate_seq_len: int): + from xrex.configs.config_loader import replace_cli_subs + + init_params = trainer.to_dict() + init_params.pop("__class") + export_cfg = type(trainer).from_dict(init_params, ensure_class=type(trainer)) + + overrides = [ + "num_devices_per_process=1", + "ep=1", + "dp=1", + "num_negatives_per_example=0", + "num_global_negatives_per_example=0", + f"history_seq_len={history_seq_len}", + f"candidate_seq_len={candidate_seq_len}", + ] + export_cfg, used = replace_cli_subs(export_cfg, overrides) + unused = [k for k, v in used.items() if not v] + if unused: + raise ValueError(f"StableHLO bundle export overrides did not match config: {unused}") + return export_cfg + + +def _batch_avals(export_cfg: Any, bs: int, *, packed: bool) -> Any: + model_config = export_cfg.model_config + batch = export_cfg.dataset.example_data(bs) + + if packed: + from xrex.data.recsys.sequence_packing import pack_batch + + batch = pack_batch( + batch=batch, + num_devices_per_process=export_cfg.parallel_config.num_devices_per_process, + num_user_prefix_tokens=model_config.num_user_prefix_tokens, + dist=None, + rng=None, + block_size=export_cfg._seqpack_block_size, + ) + + batch = _to_shape_dtype_struct(batch) + + if model_config.multimodal_embedding_type is not None: + cand_post = batch["candidate_seq"]["post_hashes"] + batch["candidate_seq"]["embedding"] = jax.ShapeDtypeStruct( + (cand_post.shape[0], cand_post.shape[1], model_config.multimodal_embedding_dim), + np.float32, + ) + + if model_config.use_post_sid and model_config.sid_num_levels > 0: + for seq_name in ("history_seq", "candidate_seq"): + post_hashes = batch[seq_name]["post_hashes"] + batch[seq_name]["post_sids"] = jax.ShapeDtypeStruct( + (post_hashes.shape[0], post_hashes.shape[1], model_config.sid_num_levels), + np.uint16, + ) + + return batch + + +def _embedding_slices(export_cfg: Any, history_seq_len: int, candidate_seq_len: int): + ht = export_cfg.dataset.hash_table + hist_post_seq = ht.num_item_hashes * history_seq_len + hist_auth_seq = ht.num_author_hashes * history_seq_len + cand_post_seq = ht.num_item_hashes * candidate_seq_len + cand_auth_seq = ht.num_author_hashes * candidate_seq_len + user_seq = ht.num_user_hashes + ip_seq = ht.num_ip_hashes if export_cfg.model_config.use_ip_address else 0 + + user_end = hist_post_seq + hist_auth_seq + cand_post_seq + cand_auth_seq + user_seq + return EmbeddingSlices( + hist_post_end=hist_post_seq, + hist_auth_end=hist_post_seq + hist_auth_seq, + cand_post_end=hist_post_seq + hist_auth_seq + cand_post_seq, + cand_auth_end=hist_post_seq + hist_auth_seq + cand_post_seq + cand_auth_seq, + user_end=user_end, + user_ip_end=user_end + ip_seq, + ) + + +def _packed_embedding_slices( + export_cfg: Any, batch_avals: Any +) -> tuple[EmbeddingSlices, PackedGeometry]: + hist = batch_avals["history_seq"] + cand = batch_avals["candidate_seq"] + hist_post_seq = int(np.prod(hist["post_hashes"].shape[1:])) + hist_auth_seq = int(np.prod(hist["auth_hashes"].shape[1:])) + cand_post_seq = int(np.prod(cand["post_hashes"].shape[1:])) + cand_auth_seq = int(np.prod(cand["auth_hashes"].shape[1:])) + user_seq = int(np.prod(batch_avals["user_hashes"].shape[1:])) + ip_seq = ( + int(np.prod(batch_avals["user_ip_hashes"].shape[1:])) + if export_cfg.model_config.use_ip_address + else 0 + ) + + user_end = hist_post_seq + hist_auth_seq + cand_post_seq + cand_auth_seq + user_seq + slices = EmbeddingSlices( + hist_post_end=hist_post_seq, + hist_auth_end=hist_post_seq + hist_auth_seq, + cand_post_end=hist_post_seq + hist_auth_seq + cand_post_seq, + cand_auth_end=hist_post_seq + hist_auth_seq + cand_post_seq + cand_auth_seq, + user_end=user_end, + user_ip_end=user_end + ip_seq, + ) + geometry = PackedGeometry( + packed_history_len=int(hist["post_hashes"].shape[1]), + packed_candidate_len=int(cand["post_hashes"].shape[1]), + bs_per_device=int(batch_avals["user_hashes"].shape[1]), + merged_batch=int(batch_avals["user_hashes"].shape[0]), + ) + return slices, geometry + + +def _make_forward_fn( + export_cfg: Any, + embedding_slices: EmbeddingSlices, + mesh: jax.sharding.Mesh, + packed_geometry: PackedGeometry | None = None, +): + import haiku as hk + import jax.numpy as jnp + + model_config = export_cfg.model_config + + def _packed_recsys_embeddings(merged: jax.Array, sl: EmbeddingSlices) -> RecsysEmbeddings: + assert packed_geometry is not None + g = packed_geometry + + def _section(start: int, end: int, rows: int) -> jax.Array: + return merged[:, start:end, :].reshape(merged.shape[0], rows, -1, merged.shape[-1]) + + return RecsysEmbeddings( + history_post_embeddings=_section(0, sl.hist_post_end, g.packed_history_len), + history_author_embeddings=_section( + sl.hist_post_end, sl.hist_auth_end, g.packed_history_len + ), + candidate_post_embeddings=_section( + sl.hist_auth_end, sl.cand_post_end, g.packed_candidate_len + ), + candidate_author_embeddings=_section( + sl.cand_post_end, sl.cand_auth_end, g.packed_candidate_len + ), + user_embeddings=_section(sl.cand_auth_end, sl.user_end, g.bs_per_device), + user_ip_embeddings=( + _section(sl.user_end, sl.user_ip_end, g.bs_per_device) + if sl.user_ip_end > sl.user_end + else None + ), + ) + + @hk.transform + def forward_fn(batch: Any, merged_embeddings: jax.Array): + sl = embedding_slices + if packed_geometry is not None: + recsys_embeddings = _packed_recsys_embeddings(merged_embeddings, sl) + else: + recsys_embeddings = RecsysEmbeddings( + history_post_embeddings=merged_embeddings[:, : sl.hist_post_end, :], + history_author_embeddings=merged_embeddings[ + :, sl.hist_post_end : sl.hist_auth_end, : + ], + candidate_post_embeddings=merged_embeddings[ + :, sl.hist_auth_end : sl.cand_post_end, : + ], + candidate_author_embeddings=merged_embeddings[ + :, sl.cand_post_end : sl.cand_auth_end, : + ], + user_embeddings=merged_embeddings[:, sl.cand_auth_end : sl.user_end, :], + user_ip_embeddings=( + merged_embeddings[:, sl.user_end : sl.user_ip_end, :] + if sl.user_ip_end > sl.user_end + else None + ), + ) + model = model_config.make(sharding_context=make_legacy_sharding_context(mesh)) + logits, candidate_continuous_predictions = model.forward(batch, recsys_embeddings) + log_probs = jax.nn.log_sigmoid(logits).astype(jnp.bfloat16).astype(jnp.float32) + cont_preds = candidate_continuous_predictions.astype(jnp.bfloat16).astype(jnp.float32) + has_nan = jnp.any(jnp.isnan(log_probs), axis=tuple(range(1, log_probs.ndim))) + return log_probs, cont_preds, has_nan + + return forward_fn + + +def _scan_custom_call_targets(lowered_text: str) -> list[str]: + targets = set(re.findall(r"stablehlo\.custom_call\s*@([\w.$-]+)", lowered_text)) + targets |= set(re.findall(r'call_target_name\s*=\s*"([^"]+)"', lowered_text)) + return sorted(targets) + + +def _input_spec( + params_avals: Any, rng_aval: Any, batch_avals: Any, merged_aval: Any +) -> list[dict[str, Any]]: + from xai_checkpointing.tree_util import keystr + + spec: list[dict[str, Any]] = [] + + param_leaves = jax.tree.leaves(params_avals) + param_keys = list(tree_to_dict(unwrap_tree(params_avals), keep_none=False).keys()) + if len(param_keys) != len(param_leaves): + raise AssertionError( + f"params flatten mismatch: {len(param_keys)} keys vs {len(param_leaves)} leaves" + ) + for key, leaf in zip(param_keys, param_leaves): + spec.append({"kind": "weight", "key": f"params.{key}", **_aval_entry(leaf)}) + + spec.append({"kind": "rng", "key": "rng", **_aval_entry(rng_aval)}) + + for path, leaf in jax.tree_util.tree_flatten_with_path(batch_avals)[0]: + key = keystr(path) + kind = "packing_layout" if key.startswith("packing_layout.") else "batch" + spec.append({"kind": kind, "key": key, **_aval_entry(leaf)}) + + spec.append( + {"kind": "merged_embeddings", "key": "merged_embeddings", **_aval_entry(merged_aval)} + ) + return spec + + +def build_bundle(trainer: RecsysTrainer) -> list[BundleFile]: + import flatbuffers + from jax import export as jax_export + + from xrex.models.recsys_gen_recs_model import RecsysGenRecsModelConfig + from xrex.models.recsys_model import RecsysAggregatedModelConfig + + model_config = trainer.model_config + if not isinstance(model_config, RecsysAggregatedModelConfig) or isinstance( + model_config, RecsysGenRecsModelConfig + ): + raise NotImplementedError( + "StableHLO bundle export supports ranking (RecsysAggregatedModelConfig) only, " + f"got {type(model_config).__name__}" + ) + if trainer.using_seqpack and trainer.using_fa4: + raise NotImplementedError( + "StableHLO bundle export supports seqpack with pallas_ranker_varlen_attn only; " + "FA4 (cutedsl_ranker_varlen_attn) block-sparse layouts are not supported yet" + ) + if not trainer.checkpoint_config.copy_port: + raise ValueError("export_stablehlo_bundle requires checkpoint_config.copy_port") + + buckets = sorted( + {int(b) for b in str(trainer.export_bundle_bs_per_device).split(",") if b.strip()} + ) + if not buckets or buckets[0] < 1: + raise ValueError( + f"invalid export_bundle_bs_per_device: {trainer.export_bundle_bs_per_device!r}" + ) + + history_seq_len = trainer.export_bundle_history_seq_len or trainer.dataset.history_seq_len + candidate_seq_len = trainer.export_bundle_candidate_seq_len or trainer.dataset.candidate_seq_len + + if trainer.using_seqpack: + block = int(trainer._seqpack_block_size) + prefix = int(model_config.num_user_prefix_tokens) + total = prefix + candidate_seq_len + history_seq_len + if total % block: + raise ValueError( + f"seqpack export needs (num_user_prefix_tokens + candidate_seq_len + " + f"history_seq_len) to be a multiple of the attention block size " + f"{block}, got {prefix} + {candidate_seq_len} + {history_seq_len} = {total}; " + "set export_bundle_candidate_seq_len / export_bundle_history_seq_len " + "to block-aligned serving lengths" + ) + + ensure_parameter_serialization_registered() + if trainer.using_seqpack: + ensure_packing_layout_serialization_registered() + export_cfg = _make_export_config(trainer, history_seq_len, candidate_seq_len) + + axis_names = export_cfg.parallel_config.mesh_axis_names() + mesh_shape = export_cfg.parallel_config.mesh_shape() + if int(np.prod(mesh_shape)) != 1: + raise AssertionError(f"export mesh must be single-device, got {mesh_shape}") + device = jax.local_devices()[0] + mesh = jax.sharding.Mesh(np.array([device]).reshape(mesh_shape), axis_names) + + packed = bool(trainer.using_seqpack) + dense_slices = ( + None if packed else _embedding_slices(export_cfg, history_seq_len, candidate_seq_len) + ) + + params_avals = jax.tree.map( + lambda leaf: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype), trainer.state_shape.params + ) + rng_aval = jax.ShapeDtypeStruct((2,), np.uint32) + + files: list[BundleFile] = [] + programs: dict[str, Any] = {} + all_custom_call_targets: set[str] = set() + + for bs in buckets: + start = time.perf_counter() + batch_avals = _batch_avals(export_cfg, bs, packed=packed) + if packed: + embedding_slices, packed_geometry = _packed_embedding_slices(export_cfg, batch_avals) + merged_batch = packed_geometry.merged_batch + else: + assert dense_slices is not None + embedding_slices, packed_geometry = dense_slices, None + merged_batch = bs + forward_fn = _make_forward_fn(export_cfg, embedding_slices, mesh, packed_geometry) + merged_aval = jax.ShapeDtypeStruct( + (merged_batch, embedding_slices.user_ip_end, model_config.emb_table_width), + model_config.embedding_dtype, + ) + args = (params_avals, rng_aval, batch_avals, merged_aval) + + with mesh: + jitted = jax.jit(forward_fn.apply) + lowered_text = jitted.lower(*args).as_text(dialect="stablehlo") + custom_call_targets = _scan_custom_call_targets(lowered_text) + exported = jax_export.export( + jitted, + platforms=("cuda",), + disabled_checks=[ + jax_export.DisabledSafetyCheck.custom_call(t) for t in custom_call_targets + ], + )(*args) + + spec = _input_spec(params_avals, rng_aval, batch_avals, merged_aval) + if len(spec) != len(exported.in_avals): + raise AssertionError( + f"input spec mismatch for bs={bs}: {len(spec)} != {len(exported.in_avals)}" + ) + for entry, aval in zip(spec, exported.in_avals): + expected = _aval_entry(aval) + if entry["shape"] != expected["shape"] or entry["dtype"] != expected["dtype"]: + raise AssertionError( + f"input spec mismatch for bs={bs} {entry['key']}: {entry} vs {expected}" + ) + + mlir_name = f"{BUNDLE_DIR}/forward_bs{bs}.mlirbc" + jax_export_name = f"{BUNDLE_DIR}/forward_bs{bs}.jax_export" + mlir_bytes = bytes(exported.mlir_module_serialized) + jax_export_bytes = bytes(exported.serialize()) + files.append(BundleFile(mlir_name, mlir_bytes)) + files.append(BundleFile(jax_export_name, jax_export_bytes)) + all_custom_call_targets.update(custom_call_targets) + + outputs = [ + {"name": name, **_aval_entry(aval)} + for name, aval in zip(("log_probs", "cont_preds", "has_nan"), exported.out_avals) + ] + programs[str(bs)] = { + "batch_size": bs, + "mlir_module": {"file": mlir_name, "adler32": zlib.adler32(mlir_bytes)}, + "jax_export": {"file": jax_export_name, "adler32": zlib.adler32(jax_export_bytes)}, + "calling_convention_version": int(exported.calling_convention_version), + "module_kept_var_idx": [int(i) for i in exported.module_kept_var_idx], + "custom_call_targets": custom_call_targets, + "input_spec": spec, + "output_spec": outputs, + "seqpack": ( + {**packed_geometry._asdict(), "merged_slices": embedding_slices._asdict()} + if packed_geometry is not None + else None + ), + } + logger.info( + "Exported StableHLO ranking forward bs=%d in %.1fs (%d inputs, %d kept, " + "custom_calls=%s)", + bs, + time.perf_counter() - start, + len(spec), + len(exported.module_kept_var_idx), + custom_call_targets, + ) + + manifest = _build_manifest( + trainer, + export_cfg, + device, + programs, + dense_slices, + history_seq_len, + candidate_seq_len, + sorted(all_custom_call_targets), + packed=packed, + ) + files.insert(0, BundleFile(MANIFEST_NAME, json.dumps(manifest, indent=2).encode())) + return files + + +def _leaf_width(tree: Any, key: str, axis: int) -> int: + leaf = tree.get(key) + if leaf is None or len(leaf.shape) <= axis: + return 0 + return int(leaf.shape[axis]) + + +def _prep_spec( + export_cfg: Any, history_seq_len: int, candidate_seq_len: int, *, packed: bool +) -> dict[str, Any]: + dataset = export_cfg.dataset + model_config = export_cfg.model_config + ht = dataset.hash_table + hk = ht.hash_keys + batch = _batch_avals(export_cfg, 1, packed=False) + hist = batch["history_seq"] + cand = batch["candidate_seq"] + + return { + "user_id_table_size": int(ht.user_id_table_size), + "user_hash_scales": [int(x) for x in hk.user_hash_scales], + "user_biases": [int(x) for x in hk.user_biases], + "user_modulus": int(hk.user_modulus), + "item_id_table_size": int(ht.item_id_table_size), + "item_hash_vocab_size": int(getattr(hk, "item_hash_vocab_size", 0) or 0), + "item_hash_scales": [int(x) for x in hk.item_hash_scales], + "item_biases": [int(x) for x in hk.item_biases], + "item_modulus": int(hk.item_modulus), + "author_id_table_size": int(ht.author_id_table_size), + "author_hash_scales": [int(x) for x in hk.author_hash_scales], + "author_biases": [int(x) for x in hk.author_biases], + "author_modulus": int(hk.author_modulus), + "ip_id_table_size": int(ht.ip_id_table_size), + "ip_hash_scales": [int(x) for x in hk.ip_hash_scales], + "ip_biases": [int(x) for x in hk.ip_biases], + "ip_modulus": int(hk.ip_modulus), + "output_vocab_size": int(dataset.output_vocab_size), + "num_continuous_actions": int(model_config.num_continuous_actions), + "search_query_embedding_dim": _leaf_width(cand, "search_query_embeddings", 2), + "num_user_categorical_features": _leaf_width(batch, "user_categorical_features", 1), + "num_user_bool_features": _leaf_width(batch, "user_bool_features", 1), + "num_user_float_features": _leaf_width(batch, "user_float_features", 1), + "num_user_int64_features": _leaf_width(batch, "user_int64_features", 1), + "num_user_installed_apps": _leaf_width(batch, "user_installed_apps_multihot", 1), + "num_post_categorical_features": _leaf_width(hist, "categorical_features", 2), + "num_post_bool_features": _leaf_width(hist, "bool_features", 2), + "num_post_float_features": _leaf_width(hist, "float_features", 2), + "num_post_int64_features": _leaf_width(hist, "int64_features", 2), + "enable_stale_post": bool( + getattr(getattr(model_config, "feature_prep", None), "enable_stale_post", False) + ), + "history_seq_len": history_seq_len, + "candidate_seq_len": candidate_seq_len, + "sid_num_levels": _leaf_width(hist, "post_sids", 2), + "multimodal_embedding_dim": _leaf_width(cand, "embedding", 2), + "num_categorical_features": 0, + "use_ip": bool(model_config.use_ip_address), + "use_seqpack": packed, + "seqpack_block_size": int(export_cfg._seqpack_block_size) if packed else 0, + "num_user_prefix_tokens": (int(model_config.num_user_prefix_tokens) if packed else 0), + "transformer_candidate_seq_len": ( + (0 if cand.get("post_ids") is not None else candidate_seq_len) if packed else 0 + ), + } + + +def _build_manifest( + trainer: RecsysTrainer, + export_cfg: Any, + device: Any, + programs: dict[str, Any], + dense_slices: EmbeddingSlices | None, + history_seq_len: int, + candidate_seq_len: int, + custom_call_targets: list[str], + *, + packed: bool, +) -> dict[str, Any]: + import jaxlib + + model_config = export_cfg.model_config + dataset = export_cfg.dataset + + ep = int(trainer.mesh.shape["expert"]) + emb_rows = int(trainer.state_shape.emb_table.x.shape[0]) + emb_rows_padded = -(-emb_rows // ep) * ep + + return { + "bundle_schema_version": BUNDLE_SCHEMA_VERSION, + "kind": "recsys_ranking_forward", + "name": trainer.name, + "model_config_class": type(trainer.model_config).__name__, + "created_timestamp": time.time(), + "jax_version": jax.__version__, + "jaxlib_version": jaxlib.__version__, + "platforms": ["cuda"], + "device_kind": str(device.device_kind), + "compute_capability": str(getattr(device, "compute_capability", "")), + "custom_call_targets": custom_call_targets, + "history_seq_len": history_seq_len, + "candidate_seq_len": candidate_seq_len, + "output_vocab_size": int(dataset.output_vocab_size), + "num_continuous_actions": int(model_config.num_continuous_actions), + "use_seqpack": packed, + "prep_spec": _prep_spec(export_cfg, history_seq_len, candidate_seq_len, packed=packed), + "programs": programs, + "embedding": { + "table_key": "emb_table", + "rows": emb_rows, + "rows_padded": emb_rows_padded, + "num_shards": ep, + "width": int(model_config.emb_table_width), + "dtype": _dtype_name(model_config.embedding_dtype), + "merged_slices": dense_slices._asdict() if dense_slices is not None else None, + "hash_table": dataset.hash_table.to_dict(), + }, + } diff --git a/phoenix/xrex/train/trainer.py b/phoenix/xrex/train/trainer.py index fc9b7381..12245a50 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -244,6 +244,13 @@ class Trainer(Config): ) current_ckpt_index: int = field(init=False, repr=False, compare=False, default=0) + @property + def step(self) -> int: + state = self.state[0] if isinstance(self.state, list) else self.state + step = state.step + 0 + step = step.item() + return step + data_first_read_timeout = 60 data_read_timeout = 30 @@ -1149,6 +1156,19 @@ def _get_checkpoint_infos(self): mask = checkpointing_common.get_load_mask(axes, srcs, self.mesh) return axes, srcs, mask + def purge_opt_state_on_load(self, host_state): + rank_logger.info("Not loading optimizer state from checkpoint") + return host_state.purge_opt_state() + + def warm_start_staging_spec(self): + keep_fields = {"opt_state"} + if hasattr(self.state, "emb_table_state"): + keep_fields.add("emb_table_state") + return (lambda tree: tree.purge_opt_state()), keep_fields + + def _uses_tensorstore_save(self) -> bool: + return self.checkpoint_config.save_method == "tensorstore" + def maybe_load_checkpoint( self, ctx: TrainerContext, tag: str | None = None ) -> tuple[bool, int, int]: @@ -1156,28 +1176,74 @@ def maybe_load_checkpoint( rank_logger.info("Not loading checkpoint; starting from scratch") return False, 0, 0 + _src = Path(ctx.checkpoint.path) + self._restored_encrypted = any( + (p / "_DEK").exists() or checkpointing_load._is_encrypted_tree(p) + for p in (_src, _src / "orbax-ckpt") + ) + + do_not_load_opt_state = ( + self.checkpoint_config.no_opt_state or self.reinit_on_load + ) and ctx.checkpoint.is_manual_load() + + use_streamed_restore = ( + self.checkpoint_config.restore_streamed + and ctx.checkpoint.format == "orbax" + and not self._uses_tensorstore_save() + ) + if self.checkpoint_config.restore_streamed and not use_streamed_restore: + rank_logger.info( + "restore_streamed=True but falling back to whole-state staging " + "(format=%s, save_method=%s)", + ctx.checkpoint.format, + self.checkpoint_config.save_method, + ) + + if use_streamed_restore: + self.state = checkpointing_load.copy_aliased_arrays(self.state) + restore_kind = next( k for k in jax.tree.leaves(jax.tree.map(lambda s: s.memory_kind, self.host_sharding)) ) restore_staging_sharding = jax.tree.map( lambda s: s.with_memory_kind(restore_kind), self.state_sharding ) - self.host_state = jax.device_put(self.state, restore_staging_sharding) + warm_purge = None + warm_keep_fields: set[str] = set() + if do_not_load_opt_state and hasattr(self.state, "purge_opt_state"): + warm_purge, warm_keep_fields = self.warm_start_staging_spec() + rank_logger.info( + "Not loading optimizer state from checkpoint (params-only %s staging)", + "streamed" if use_streamed_restore else "pinned-host", + ) + staged_state = warm_purge(self.state) + staged_sharding = warm_purge(restore_staging_sharding) + else: + do_not_load_opt_state = False + staged_state = self.state + staged_sharding = restore_staging_sharding + + if use_streamed_restore: + self.host_state = None + else: + self.host_state = jax.device_put(staged_state, staged_sharding) rename = None loads: dict[str, dict[str, jax.Array]] = {} if ctx.checkpoint.format == "orbax": - host_state = unwrap_tree(self.host_state) - - do_not_load_opt_state = ( - self.checkpoint_config.no_opt_state or self.reinit_on_load - ) and ctx.checkpoint.is_manual_load() + if use_streamed_restore: + host_state = jax.tree.map( + lambda p: p.x if isinstance(p, Parameter) else p, + staged_state, + is_leaf=lambda x: isinstance(x, Parameter), + ) + else: + host_state = unwrap_tree(self.host_state) if do_not_load_opt_state: - rank_logger.info("Not loading optimizer state from checkpoint") - host_state = host_state.purge_opt_state() + host_state = self.purge_opt_state_on_load(host_state) host_state = tree_to_dict(host_state) domains = None @@ -1219,20 +1285,75 @@ def maybe_load_checkpoint( name = checkpointing_load.rename_tensor(name, rename_state_patterns) loads[checkpoint_path][name] = tensor + if use_streamed_restore: + del staged_state + + def _graft_replaced(pairs: list[tuple[jax.Array, jax.Array]]) -> None: + id_map = {id(old): new for old, new in pairs} + assert len(id_map) == len(pairs), ( + "restore_streamed does not support aliased state leaves: " + "multiple loaded tensors share one device array" + ) + grafted: set[int] = set() + + def _graft(x): + new = id_map.get(id(x)) + if new is None: + return x + grafted.add(id(x)) + return new + + self.state = jax.tree.map(_graft, self.state) + missing = len(id_map) - len(grafted) + assert not missing, ( + f"{missing} loaded tensors were not grafted back into the state tree" + ) + for checkpoint_path, partial_host_state in loads.items(): - checkpointing_load.load_checkpoint( - checkpoint_path, - partial_host_state, - load_mask=mask, - rename=rename, - domains=domains, - tag=tag, - timeout=self.checkpoint_config.timeout_secs, - ) + if use_streamed_restore: + checkpointing_load.load_checkpoint_streamed( + checkpoint_path, + partial_host_state, + load_mask=mask, + rename=rename, + domains=domains, + tag=tag, + timeout=self.checkpoint_config.timeout_secs, + window_gb=self.checkpoint_config.restore_window_gb, + window_cap_gb=( + self.checkpoint_config.restore_concurrent_gb + if self.checkpoint_config.restore_window_gb is None + else None + ), + on_replaced=_graft_replaced, + ) + else: + checkpointing_load.load_checkpoint( + checkpoint_path, + partial_host_state, + load_mask=mask, + rename=rename, + domains=domains, + tag=tag, + timeout=self.checkpoint_config.timeout_secs, + concurrent_gb=self.checkpoint_config.restore_concurrent_gb, + ) - self.state = None - self.state = jax.device_put(self.host_state, self.state_sharding) - self.host_state = jax.device_put(self.state, self.host_sharding) + if not use_streamed_restore: + if do_not_load_opt_state: + loaded = jax.device_put(self.host_state, warm_purge(self.state_sharding)) + self.state = self.state._replace( + **{ + field: getattr(loaded, field) + for field in self.state._fields + if field not in warm_keep_fields + } + ) + self.host_state = None + else: + self.state = None + self.state = jax.device_put(self.host_state, self.state_sharding) + self.host_state = jax.device_put(self.state, self.host_sharding) if mask: axes_sizes = {} @@ -1245,8 +1366,23 @@ def maybe_load_checkpoint( ) _, treedef = jax.tree.flatten(self.state_shape) - axes = jax.tree.unflatten(treedef, jax.tree.flatten(axes)[0]) - srcs = jax.tree.unflatten(treedef, jax.tree.flatten(srcs)[0]) + keep = [ + len(jax.tree.leaves(node)) > 0 + for node in jax.tree.leaves( + self.state_shape, is_leaf=lambda x: isinstance(x, Parameter) + ) + ] + + def _align_to_state(tree): + kept = [ + leaf + for leaf, keep_leaf in zip(jax.tree.leaves(tree), keep, strict=True) + if keep_leaf + ] + return jax.tree.unflatten(treedef, kept) + + axes = _align_to_state(axes) + srcs = _align_to_state(srcs) self.state = checkpointing_load.broadcast_replicated( self.state, axes, srcs, self.mesh ) @@ -1430,20 +1566,19 @@ def handle_metrics( mfu = self.model_config.compute_mfu(num_seq_per_sec_per_device) tflops = self.model_config.compute_tflops(num_seq_per_sec_per_device) + etas: list[float] = [] + fractions: list[float] = [] if self.max_steps is not None: - eta = (self.max_steps - self.prev_step) * step_time / 3600 - metrics["finished_percent"] = int(metrics["step"]) / self.max_steps - elif self.max_samples is not None: - eta = ( - (self.max_samples - self.elapsed_samples) - / float(metrics["examples_per_batch"]) - * step_time - / 3600 + etas.append((self.max_steps - self.prev_step) * step_time / 3600) + fractions.append(int(metrics["step"]) / self.max_steps) + if self.max_samples is not None: + remaining_steps = (self.max_samples - self.elapsed_samples) / float( + metrics["examples_per_batch"] ) - metrics["finished_percent"] = self.elapsed_samples / self.max_samples - else: - eta = -1 - metrics["finished_percent"] = -1 + etas.append(remaining_steps * step_time / 3600) + fractions.append(self.elapsed_samples / self.max_samples) + eta = min(etas) if etas else -1 + metrics["finished_percent"] = max(fractions) if fractions else -1 self.prev_step = int(metrics["step"]) metrics["soft_step"] = soft_step diff --git a/phoenix/xrex/train/trainer_recsys.py b/phoenix/xrex/train/trainer_recsys.py index 50c77db0..fd517792 100644 --- a/phoenix/xrex/train/trainer_recsys.py +++ b/phoenix/xrex/train/trainer_recsys.py @@ -9,6 +9,7 @@ import logging import math import os +import pathlib import shutil import signal import sys @@ -84,6 +85,7 @@ AsyncEmbOptimizer, ) from xrex.train.misc import ( + CheckpointConfig, PostEmbeddings, RecsysTrainingState, ) @@ -392,6 +394,11 @@ def _read_checkpoint_kafka_config(ctx) -> tuple[str | None, int | None]: return None, None +@configclass +class RecsysCheckpointConfig(CheckpointConfig): + keep_emb_opt_state: bool = False + + @configclass class RecsysTrainer(Trainer): offsets_to_commit: dict[int, int] = field(default_factory=dict) @@ -408,6 +415,11 @@ class RecsysTrainer(Trainer): checkpoint_storage_urls: str = "" + export_stablehlo_bundle: bool = False + export_bundle_bs_per_device: str = "1,2,4" + export_bundle_history_seq_len: int = 0 + export_bundle_candidate_seq_len: int = 0 + smoothing_windows: list[int] = field(default_factory=lambda: [1_048_576, 4_194_304]) reset_data_position: bool = False @@ -446,11 +458,14 @@ def __post_init__(self): assert isinstance(self.model_config, RecsysAggregatedModelConfig) hl = self.model_config.history_seq_len self.seqpack_distribution = FixedLengthDistribution(min_len=hl, max_len=hl, mean_len=hl) + if self.export_stablehlo_bundle and not self.checkpoint_config.copy_port: + raise ValueError("export_stablehlo_bundle requires checkpoint_config.copy_port") state: RecsysTrainingState = field(init=False, repr=False, compare=False) _pending_shmem_ckpt_write_s: float | None = field(default=None, init=False, repr=False) _pending_checksum_s: float | None = field(default=None, init=False, repr=False) _pending_gc_collect_s: float | None = field(default=None, init=False, repr=False) + _stablehlo_bundle_files: list | None = field(default=None, init=False, repr=False) _engine = None _shmem_write_pool = None @@ -596,6 +611,8 @@ def init(self, batch: RecsysFeaturesBatch, rng: jax.Array) -> RecsysTrainingStat dummy, dummy, enable_platform_metrics=self.model_config.enable_platform_metrics, + split_head_training_by_source=self.model_config.split_head_training_by_source, + metric_mask_keys=self.model_config.metric_mask_keys, ).keys() ) rce_ema = { @@ -1322,8 +1339,18 @@ def update( valid_step, grad_norm = self.is_valid_step(gradients) + segment_sum_result = self._segment_sum(emb_gradients, inverse_indices, num_unique) + + emb_gradients = replace( + emb_gradients.history_author_embeddings, + x=segment_sum_result, + ) + + emb_valid_step, emb_grad_norm = self.is_valid_step(emb_gradients) + keep_step = valid_step & emb_valid_step + def _update(updated, original): - return jnp.where(valid_step, updated, original) + return jnp.where(keep_step, updated, original) new_params = jax.tree.map(_update, new_params, state.params) new_opt_state = jax.tree.map(_update, new_opt_state, state.opt_state) @@ -1334,7 +1361,7 @@ def _update(updated, original): "step": state.step, "loss": loss, "examples_per_batch": examples, - "valid_step": valid_step, + "valid_step": keep_step, "global_grad_norm": grad_norm, "learning_rate": lr * new_opt_state.hyperparams["learning_rate"], "weight_decay": new_opt_state.hyperparams["weight_decay"], @@ -1348,15 +1375,6 @@ def _update(updated, original): new_calib_ema = stats.pop("_calib_ema", state.calib_ema) metrics.update(**stats) - segment_sum_result = self._segment_sum(emb_gradients, inverse_indices, num_unique) - - emb_gradients = replace( - emb_gradients.history_author_embeddings, - x=segment_sum_result, - ) - - emb_valid_step, emb_grad_norm = self.is_valid_step(emb_gradients) - new_emb_table, emb_new_opt_state, emb_optim_metrics = self._emb_optim.sparse_update( grads=emb_gradients, full_state=state.emb_table_state, @@ -1364,7 +1382,7 @@ def _update(updated, original): unique_tokens=unique_tokens, num_unique=num_unique, lr=lr, - valid_step=emb_valid_step, + valid_step=keep_step, carry=emb_carry, ) @@ -1393,7 +1411,7 @@ def _update(updated, original): metric_keys, metric_values = zip(*metrics.items()) metrics = { metric_keys: jnp.stack([jnp.float32(v) for v in metric_values], axis=0), - ("valid_step",): valid_step & emb_valid_step, + ("valid_step",): keep_step, } return new_state, metrics, {} @@ -1432,7 +1450,7 @@ def async_emb_step( self.batch_size, -1, flat_prefetched.shape[-1] ) - prev_update_pins, updating_table, updating_emb_state = ( + prev_update_pins, updating_table, updating_emb_state, emb_optim_metrics = ( self._emb_optim.gradient_update_start( self._async_emb_context, prev_step_grad_update, @@ -1503,8 +1521,12 @@ def loss_fn(params, embeddings): valid_step, grad_norm = self.is_valid_step(gradients) + emb_gradients = jax.tree.map(lambda x: x.astype(jnp.bfloat16), emb_gradients) + deferred_emb_valid_step, deferred_emb_grad_norm = self.is_valid_step(emb_gradients) + keep_step = valid_step & deferred_emb_valid_step + def _update(updated, original): - return jnp.where(valid_step, updated, original) + return jnp.where(keep_step, updated, original) new_params = jax.tree.map(_update, new_params, state.params) new_opt_state = jax.tree.map(_update, new_opt_state, state.opt_state) @@ -1513,7 +1535,7 @@ def _update(updated, original): "step": state.step, "loss": loss, "examples_per_batch": np.prod(data["user_hashes"].shape[:-1]), - "valid_step": valid_step, + "valid_step": keep_step, "global_grad_norm": grad_norm, "learning_rate": lr * new_opt_state.hyperparams["learning_rate"], "weight_decay": new_opt_state.hyperparams["weight_decay"], @@ -1521,7 +1543,10 @@ def _update(updated, original): "b2": new_opt_state.hyperparams["b2"], "emb_grad_norm": emb_grad_norm, "emb_valid_step": emb_valid_step, + "deferred_emb_grad_norm": deferred_emb_grad_norm, + "deferred_emb_valid_step": deferred_emb_valid_step, } + metrics.update(emb_optim_metrics) if self.track_norm_metrics: metrics.update(norm_metrics(new_params, new_opt_state, gradients, updates)) @@ -1534,10 +1559,9 @@ def _update(updated, original): metric_keys, metric_values = zip(*metrics.items()) metrics = { metric_keys: jnp.stack([jnp.float32(v) for v in metric_values], axis=0), - ("valid_step",): valid_step & emb_valid_step, + ("valid_step",): keep_step, } - emb_gradients = jax.tree.map(lambda x: x.astype(jnp.bfloat16), emb_gradients) next_step_grad_update = AsyncEmbGradientUpdate( unique_tokens=unique_tokens, grads=jax.lax.with_sharding_constraint( @@ -1545,7 +1569,7 @@ def _update(updated, original): P(self._async_emb_context.data_axis, None), ), segment_ids=segment_ids, - pending=jnp.asarray(True), + pending=keep_step, ) new_state = RecsysTrainingState( @@ -1663,6 +1687,8 @@ def add_block_sparse_layout(self, batch: RecsysFeaturesBatch) -> RecsysFeaturesB transformer_candidate_seq_len=transformer_candidate_seq_len, max_history_seq_len=(_mc.num_user_prefix_tokens + _mc.history_seq_len), packed_seq_len=int(layout.segment_ids.shape[1]), + padding_mask=layout.padding_mask, + num_user_prefix_tokens=_mc.num_user_prefix_tokens, ) return {**batch, "packing_layout": replace(layout, block_sparse=block_sparse)} @@ -1700,7 +1726,6 @@ def _create_async_emb_executables(self, init_data, lr_shape, compiler_options) - assert self.parallel_config.num_devices_per_process == 1 assert "expert" in data_axis assert all(self.mesh.shape[a] == 1 for a in data_axis if a != "expert") - assert self.grad_norm_keep_threshold is None assert isinstance(self._emb_optim, AsyncEmbOptimizer) tokens_per_example = jax.eval_shape(self.get_flattened_token_ids, init_data).shape[1] @@ -1754,7 +1779,7 @@ def empty_embedding_grad_update(): ) def apply_deferred_embedding_update(state, grad_update): - pins, updating_table, updating_emb_state = self._emb_optim.gradient_update_start( + pins, updating_table, updating_emb_state, _ = self._emb_optim.gradient_update_start( self._async_emb_context, grad_update, state.emb_table.x, @@ -2142,6 +2167,34 @@ def _is_copy_port_binder(self) -> bool: return self.ctx.rank == 0 return hostnames.index(hostnames[self.ctx.rank]) == self.ctx.rank + def _copy_port_channel(self, port: int) -> grpc.Channel: + cc = self.checkpoint_config + if not cc.copy_port_tls_cert: + return grpc.insecure_channel(f"127.0.0.1:{port}") + if not cc.copy_port_tls_ca or not cc.copy_port_tls_server_name: + raise ValueError( + "copy_port_tls_ca and copy_port_tls_server_name are required for the " + "loopback client when copy_port TLS is enabled" + ) + if cc.copy_port_tls_client_ca and not ( + cc.copy_port_tls_client_cert and cc.copy_port_tls_client_key + ): + raise ValueError( + "copy_port_tls_client_ca (mTLS) requires copy_port_tls_client_cert/key " + "for the loopback client" + ) + + def read(p: str) -> bytes | None: + return pathlib.Path(p).read_bytes() if p else None + + creds = grpc.ssl_channel_credentials( + root_certificates=read(cc.copy_port_tls_ca), + private_key=read(cc.copy_port_tls_client_key), + certificate_chain=read(cc.copy_port_tls_client_cert), + ) + options = (("grpc.ssl_target_name_override", cc.copy_port_tls_server_name),) + return grpc.secure_channel(f"127.0.0.1:{port}", creds, options=options) + def _free_ports(self): if port := self.checkpoint_config.copy_port: for conn in psutil.net_connections(kind="inet"): @@ -2294,6 +2347,8 @@ def _finalize_checkpoint_load( dummy, dummy, enable_platform_metrics=self.model_config.enable_platform_metrics, + split_head_training_by_source=self.model_config.split_head_training_by_source, + metric_mask_keys=self.model_config.metric_mask_keys, ).keys() ) loaded_rce = self.state.rce_ema @@ -2391,6 +2446,17 @@ def _has_own_local_checkpoint(self) -> bool: return False return any(d.startswith("elapsed_samples_") for d in os.listdir(own_dir)) + def purge_opt_state_on_load(self, host_state): + if getattr(self.checkpoint_config, "keep_emb_opt_state", False): + rank_logger.info("Not loading dense optimizer state (keeping emb_table_state)") + return host_state._replace(opt_state=None) + return super().purge_opt_state_on_load(host_state) + + def warm_start_staging_spec(self): + if getattr(self.checkpoint_config, "keep_emb_opt_state", False): + return (lambda tree: tree._replace(opt_state=None)), {"opt_state"} + return super().warm_start_staging_spec() + def maybe_load_checkpoint(self, ctx: TrainerContext, tag=None): assert isinstance( self.model_config, (RecsysAggregatedModelConfig, RecsysTwoTowerModelConfig) @@ -2766,6 +2832,51 @@ def _checkpoint_data_position(self) -> DataPosition | None: return self.dataset.get_data_position() return self._batch_pipeline.current.data_position + def _maybe_build_stablehlo_bundle(self) -> list | None: + if not self.export_stablehlo_bundle: + return None + if self._stablehlo_bundle_files is None: + from xrex.train.recsys_bundle_export import build_bundle + + try: + start = time.perf_counter() + self._stablehlo_bundle_files = build_bundle(self) + rank_logger.info( + "Built StableHLO bundle (%d files, %.1fs); it will be included in " + "every copy_port checkpoint publish", + len(self._stablehlo_bundle_files), + time.perf_counter() - start, + ) + except Exception: + rank_logger.exception( + "StableHLO bundle export failed; disabling for the rest of this run " + "(copy_port checkpoints continue without export/)" + ) + self._stablehlo_bundle_files = [] + if self._engine is None: + return None + return self._stablehlo_bundle_files or None + + def _write_stablehlo_bundle_files(self, prefix: str, bundle_files: list | None) -> None: + from xrex.train.recsys_bundle_export import MANIFEST_NAME, restamp_manifest + + try: + for bundle_file in bundle_files or (): + data = bundle_file.data + if bundle_file.name == MANIFEST_NAME: + data = restamp_manifest(data) + bundle_path = f"{OUT_PATH}/.{prefix}/{bundle_file.name}" + os.makedirs(os.path.dirname(bundle_path), exist_ok=True) + _write_all_bytes(bundle_path, memoryview(data)) + except OSError: + rank_logger.exception( + "StableHLO bundle write failed; disabling for the rest of this run " + "(this publish continues without export/)" + ) + self._stablehlo_bundle_files = [] + for subdir in {f.name.split("/", 1)[0] for f in bundle_files or ()}: + shutil.rmtree(f"{OUT_PATH}/.{prefix}/{subdir}", ignore_errors=True) + def _write_shmem_checkpoint( self, write_items: list[tuple[str, typing.Any, npt.NDArray]], @@ -2773,6 +2884,7 @@ def _write_shmem_checkpoint( data_pos: typing.Any, prefix: str, stub: typing.Any, + bundle_files: list | None = None, ) -> float: start = time.perf_counter() proc_idx = jax.process_index() @@ -2841,6 +2953,7 @@ def _write_shmem_checkpoint( }, f, ) + self._write_stablehlo_bundle_files(prefix, bundle_files) if data_pos is not None: with open(f"{OUT_PATH}/.{prefix}/{_DATA_POSITION_FILENAME}", "w") as f: json.dump(data_pos, f) @@ -2947,12 +3060,16 @@ def save_checkpoint(self, *args, **kwargs): if port: if self._engine is None and self._is_copy_port_binder(): + cc = self.checkpoint_config self._engine = xai_recsys_engine.RecsysPredictorServer( port, port + 1, 1, 0, - copy_max_entries=self.checkpoint_config.shm_max_entries, + copy_max_entries=cc.shm_max_entries, + tls_cert_path=cc.copy_port_tls_cert or None, + tls_key_path=cc.copy_port_tls_key or None, + tls_client_ca_path=cc.copy_port_tls_client_ca or None, ) multihost_utils.sync_global_devices("recsys-copy-port-bind") @@ -2964,7 +3081,7 @@ def save_checkpoint(self, *args, **kwargs): self._pending_shmem_ckpt_write_s = self._shmem_write_future.result() self._shmem_write_future = None multihost_utils.sync_global_devices("recsys-save-checkpoint1") - stub = copy_pb2_grpc.CopyStub(grpc.insecure_channel(f"127.0.0.1:{port}")) + stub = copy_pb2_grpc.CopyStub(self._copy_port_channel(port)) if not hasattr(self, "host_state") or self.host_state is None: self.host_state = jax.device_put(self.state, self.host_sharding) @@ -3111,6 +3228,8 @@ def save_checkpoint(self, *args, **kwargs): data_pos = self._checkpoint_data_position() if self._engine is not None else None + bundle_files = self._maybe_build_stablehlo_bundle() + self._shmem_write_future = self._shmem_write_pool.submit( self._write_shmem_checkpoint, write_items, @@ -3118,6 +3237,7 @@ def save_checkpoint(self, *args, **kwargs): data_pos, prefix, stub, + bundle_files, ) if store != Store.FS: @@ -3347,35 +3467,54 @@ def maybe_build_retrieval_post_embeddings(self): ) else: combined_hashes_shard = author_hashes_shard - combined_hashes_jax = jax.make_array_from_process_local_data( - self.data_sharding, np.asarray(combined_hashes_shard) - ) - post_author_embeddings = self._lookup(self.state.emb_table, combined_hashes_jax) - post_sids_raw_shard = post_sids_raw[start_idx:end_idx] if _use_post_sid: post_sids_u16_shard = (post_sids_raw_shard + 1).astype(np.uint16) else: post_sids_u16_shard = np.zeros_like(post_sids_raw_shard, dtype=np.uint16) - post_sids_jax = jax.make_array_from_process_local_data( - self.data_sharding, post_sids_u16_shard - ) - post_hashes_jax = jax.make_array_from_process_local_data( - self.data_sharding, np.asarray(post_hashes_shard) - ) rng, _new_rng = jax.random.split(self.state.rng) + combined_hashes_np = np.asarray(combined_hashes_shard) + post_hashes_np = np.asarray(post_hashes_shard) + _shard_rows = combined_hashes_np.shape[0] + if total_samples % self.data_world_size == 0: + _chunk_rows = min(_shard_rows, 65536) + else: + _chunk_rows = _shard_rows + + def _forward_chunked(head_index: int) -> jax.Array: + outs: list[np.ndarray] = [] + for _start in range(0, _shard_rows, _chunk_rows): + _sl = slice(_start, min(_start + _chunk_rows, _shard_rows)) + _hashes_jax = jax.make_array_from_process_local_data( + self.data_sharding, combined_hashes_np[_sl] + ) + _pae = self._lookup(self.state.emb_table, _hashes_jax) + _sids_jax = jax.make_array_from_process_local_data( + self.data_sharding, post_sids_u16_shard[_sl] + ) + _ph_jax = jax.make_array_from_process_local_data( + self.data_sharding, post_hashes_np[_sl] + ) + _out = self.candidate_tower_forward_jit( + self.state.params, + rng, + _pae.x, + _sids_jax, + _ph_jax, + head_index, + ) + _local_shards = sorted(_out.addressable_shards, key=lambda s: s.index[0].start or 0) + outs.append(np.concatenate([np.asarray(s.data) for s in _local_shards], axis=0)) + del _out, _pae + return jax.make_array_from_process_local_data( + self.data_sharding, np.concatenate(outs, axis=0) + ) + head_dataset_mapping = getattr(self.model_config, "head_dataset_mapping", None) num_heads = self.model_config.candidate_tower_config.num_candidate_heads - candidate_embeddings = self.candidate_tower_forward_jit( - self.state.params, - rng, - post_author_embeddings.x, - post_sids_jax, - post_hashes_jax, - 0, - ) + candidate_embeddings = _forward_chunked(0) if head_dataset_mapping is not None and num_heads > 1: dataset_to_head: dict[int, int] = {} for ds_name, head_idx in head_dataset_mapping.items(): @@ -3389,26 +3528,10 @@ def maybe_build_retrieval_post_embeddings(self): head_indices_shard, ) for h in range(1, num_heads): - emb_h = self.candidate_tower_forward_jit( - self.state.params, - rng, - post_author_embeddings.x, - post_sids_jax, - post_hashes_jax, - h, - ) + emb_h = _forward_chunked(h) mask_h = post_head_jax == h candidate_embeddings = jnp.where(mask_h, emb_h, candidate_embeddings) del emb_h - else: - candidate_embeddings = self.candidate_tower_forward_jit( - self.state.params, - rng, - post_author_embeddings.x, - post_sids_jax, - post_hashes_jax, - 0, - ) global_post_ids = multihost_utils.host_local_array_to_global_array( self.int64_to_two_int32(post_ids), self.mesh, P(None) diff --git a/phoenix/xrex/utils/checkpoint_cloud.py b/phoenix/xrex/utils/checkpoint_cloud.py index 8e90e7a6..1f5e06ef 100644 --- a/phoenix/xrex/utils/checkpoint_cloud.py +++ b/phoenix/xrex/utils/checkpoint_cloud.py @@ -7,7 +7,7 @@ logger = logging.getLogger(__name__) -_ORBAX_TMP_PREFIX = ".orbax-checkpoint-tmp-" +ORBAX_TMP_DIR_SUFFIX = ".orbax-checkpoint-tmp-" def ensure_orbax_dir_finalized(local_ckpt_dir: str, tag: str) -> None: @@ -15,7 +15,7 @@ def ensure_orbax_dir_finalized(local_ckpt_dir: str, tag: str) -> None: if os.path.isdir(orbax_dir): return for entry in os.listdir(local_ckpt_dir): - if entry.startswith(f"{tag}{_ORBAX_TMP_PREFIX}"): + if entry.startswith(f"{tag}{ORBAX_TMP_DIR_SUFFIX}"): tmp_dir = os.path.join(local_ckpt_dir, entry) try: os.rename(tmp_dir, orbax_dir) diff --git a/phoenix/xrex/utils/checkpointing.py b/phoenix/xrex/utils/checkpointing.py index 5047ae4e..953ab1a5 100644 --- a/phoenix/xrex/utils/checkpointing.py +++ b/phoenix/xrex/utils/checkpointing.py @@ -1,7 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import contextlib +import ctypes +import fcntl +import gc import logging import os +import sys +import time +from typing import Any import jax import jax.numpy as jnp @@ -16,16 +23,95 @@ rank_logger = logging.getLogger("rank") _CHECKPOINTER = None +_CHECKPOINTER_SAVE_CONCURRENT_GB: int | None = None +_SAVE_NODE_SERIALIZE_ENV = "XAI_SAVE_NODE_SERIALIZE" +_SAVE_NODE_LOCK_FILE_ENV = "XAI_SAVE_NODE_LOCK_FILE" -def get_checkpointer(timeout_secs=900): - global _CHECKPOINTER + +def _save_node_serialize_enabled() -> bool: + return os.getenv(_SAVE_NODE_SERIALIZE_ENV, "0").lower() not in ("0", "", "false") + + +def _save_node_lock_path() -> str: + path = os.getenv(_SAVE_NODE_LOCK_FILE_ENV) + if path: + return path + if os.path.isdir("/dev/shm"): + return "/dev/shm/xai_save_node_lock" + return "/tmp/xai_save_node_lock" + + +class _NodeBatchLock: + def __init__(self, path: str): + self._path = path + self._fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o666) + + def __enter__(self): + t0 = time.time() + fcntl.flock(self._fd, fcntl.LOCK_EX) + waited = time.time() - t0 + if waited > 1.0: + rank_logger.info( + "save node-serialize: waited %.1fs for node lock %s", waited, self._path + ) + return self + + def __exit__(self, *exc): + fcntl.flock(self._fd, fcntl.LOCK_UN) + + def close(self): + try: + os.close(self._fd) + except OSError: + pass + + +def _release_batch_memory(): + gc.collect() + try: + ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception: + pass + + +class _AsyncCheckpointer(ocp.AsyncCheckpointer): + def wait_until_finished(self): + super().wait_until_finished() + _orbax_encrypted = sys.modules.get("xai_checkpointing.orbax_encrypted") + if _orbax_encrypted is not None: + _orbax_encrypted.wait_until_finished() + + +def get_checkpointer(timeout_secs=900, save_concurrent_gb: int | None = None): + global _CHECKPOINTER, _CHECKPOINTER_SAVE_CONCURRENT_GB if _CHECKPOINTER is None: - _CHECKPOINTER = ocp.AsyncCheckpointer( - ocp.PyTreeCheckpointHandler(use_zarr3=True), timeout_secs + handler_kwargs: dict[str, Any] = {"use_zarr3": True, "restore_concurrent_gb": 1} + if save_concurrent_gb is not None: + handler_kwargs["save_concurrent_gb"] = save_concurrent_gb + handler_kwargs["restore_concurrent_gb"] = save_concurrent_gb + _CHECKPOINTER_SAVE_CONCURRENT_GB = save_concurrent_gb + rank_logger.info( + "Creating AsyncCheckpointer: PyTreeCheckpointHandler(use_zarr3=True, " + "save_concurrent_gb=%s, restore_concurrent_gb=%s) " + "[None => Orbax write-limiter default 96GB; D2H still all-at-once " + "unless save_checkpoint registers ThrottledD2HArrayHandler]", + save_concurrent_gb, + save_concurrent_gb, + ) + _CHECKPOINTER = _AsyncCheckpointer( + ocp.PyTreeCheckpointHandler(**handler_kwargs), timeout_secs ) if not hasattr(_CHECKPOINTER, "_post_finalization_callback"): raise RuntimeError("Orbax version is too old") + elif save_concurrent_gb is not None and _CHECKPOINTER_SAVE_CONCURRENT_GB != save_concurrent_gb: + rank_logger.warning( + "get_checkpointer(save_concurrent_gb=%s) ignored; checkpointer already " + "created with save_concurrent_gb=%s. D2H throttle (if enabled) still " + "uses the value passed to save_checkpoint.", + save_concurrent_gb, + _CHECKPOINTER_SAVE_CONCURRENT_GB, + ) return _CHECKPOINTER @@ -38,6 +124,101 @@ def _get_json_tspec_write(self, *args, **kwargs): return spec +class ThrottledD2HArrayHandler(ocp.type_handlers.ArrayHandler): + def __init__(self, concurrent_bytes: int, **kwargs): + if concurrent_bytes <= 0: + raise ValueError(f"concurrent_bytes must be > 0, got {concurrent_bytes}") + super().__init__(**kwargs) + self._concurrent_bytes = concurrent_bytes + + def _addressable_nbytes(self, arr: jax.Array) -> int: + total = 0 + for shard in arr.addressable_shards: + if self._replica_id is None or shard.replica_id == self._replica_id: + total += int(shard.data.nbytes) + return total + + async def serialize(self, values, infos, args=None): + args = args or [ocp.SaveArgs()] * len(values) + if not values: + return [] + + batches: list[tuple[list, list, list]] = [] + cur_v: list = [] + cur_i: list = [] + cur_a: list = [] + cur_b = 0 + max_leaf = 0 + total_b = 0 + for v, info, arg in zip(values, infos, args): + nb = self._addressable_nbytes(v) + max_leaf = max(max_leaf, nb) + total_b += nb + if cur_v and cur_b + nb > self._concurrent_bytes: + batches.append((cur_v, cur_i, cur_a)) + cur_v, cur_i, cur_a, cur_b = [], [], [], 0 + cur_v.append(v) + cur_i.append(info) + cur_a.append(arg) + cur_b += nb + if cur_v: + batches.append((cur_v, cur_i, cur_a)) + + rank_logger.info( + "ThrottledD2HArrayHandler: save_concurrent_bytes=%.2fGiB, " + "addressable_total=%.2fGiB, max_leaf=%.2fGiB, num_leaves=%d, " + "num_batches=%d (D2H+write one batch at a time)", + self._concurrent_bytes / (1 << 30), + total_b / (1 << 30), + max_leaf / (1 << 30), + len(values), + len(batches), + ) + if max_leaf > self._concurrent_bytes: + rank_logger.warning( + "ThrottledD2HArrayHandler: largest leaf %.2fGiB exceeds " + "save_concurrent_bytes %.2fGiB; that leaf still D2Hs in one shot. " + "Raise checkpoint_config.save_concurrent_gb if write limiter errors.", + max_leaf / (1 << 30), + self._concurrent_bytes / (1 << 30), + ) + + node_lock: _NodeBatchLock | None = None + if _save_node_serialize_enabled(): + node_lock = _NodeBatchLock(_save_node_lock_path()) + rank_logger.info( + "save node-serialize ACTIVE (flock per batch) lock=%s pid=%d num_batches=%d", + node_lock._path, + os.getpid(), + len(batches), + ) + + try: + for batch_idx, (bv, bi, ba) in enumerate(batches): + with node_lock if node_lock is not None else contextlib.nullcontext(): + batch_bytes = sum(self._addressable_nbytes(v) for v in bv) + rank_logger.info( + "ThrottledD2HArrayHandler: batch %d/%d leaves=%d bytes=%.2fGiB", + batch_idx + 1, + len(batches), + len(bv), + batch_bytes / (1 << 30), + ) + futs = await super().serialize(bv, bi, ba) + for fut in futs: + fut.result() + if node_lock is not None: + _release_batch_memory() + finally: + if node_lock is not None: + node_lock.close() + return [] + + +class ThrottledNoCompressionArrayHandler(ThrottledD2HArrayHandler, NoCompressionArrayHandler): + pass + + def save_checkpoint( state, path, @@ -49,6 +230,9 @@ def save_checkpoint( compressed=True, chunk_byte_size: int = 1024 * 1024 * 4, tracer: Tracer | None = None, + save_concurrent_gb: int | None = None, + kms_client: object | None = None, + encryption_chunk_size: int = 8 * 1024 * 1024, ): if not tracer: tracer = MockTracer() @@ -57,13 +241,27 @@ def save_checkpoint( tag = "orbax-ckpt" os.makedirs(path, exist_ok=True) - checkpointer = get_checkpointer(timeout_secs) + if kms_client is not None: + from xai_checkpointing import orbax_encrypted + + rank_logger.info("Saving ENCRYPTED orbax checkpoint (xai_encrypted driver) to %s", path) + orbax_encrypted._require_array_leaves(state) + if save_concurrent_gb is not None: + array_handler_cls = ( + ThrottledD2HArrayHandler if compressed else ThrottledNoCompressionArrayHandler + ) + elif not compressed: + array_handler_cls = NoCompressionArrayHandler + else: + array_handler_cls = ocp.type_handlers.ArrayHandler + checkpointer = orbax_encrypted.get_encrypted_checkpointer( + kms_client, encryption_chunk_size, timeout_secs, save_concurrent_gb, array_handler_cls + ) + else: + checkpointer = get_checkpointer(timeout_secs, save_concurrent_gb=save_concurrent_gb) with tracer.start_as_current_span("wait_for_previous_checkpoint"): - try: - checkpointer.wait_until_finished() - except Exception as e: - rank_logger.error("Error waiting for previous checkpoint: %s", e) + checkpointer.wait_until_finished() dest = os.path.join(path, tag) if os.path.exists(dest): @@ -86,10 +284,35 @@ def save_checkpoint( state = multihost_utils.host_local_array_to_global_array(state, mesh, pspecs) original_handler = ocp.type_handlers.get_type_handler(jax.Array) - if not compressed: - ocp.type_handlers.register_type_handler( - jax.Array, NoCompressionArrayHandler(), override=True - ) + if kms_client is None: + if save_concurrent_gb is not None: + concurrent_bytes = int(save_concurrent_gb) * 10**9 + if compressed: + handler = ThrottledD2HArrayHandler(concurrent_bytes) + else: + handler = ThrottledNoCompressionArrayHandler(concurrent_bytes) + ocp.type_handlers.register_type_handler(jax.Array, handler, override=True) + rank_logger.info( + "save_checkpoint: ThrottledD2HArrayHandler ACTIVE " + "save_concurrent_gb=%s compressed=%s path=%s", + save_concurrent_gb, + compressed, + path, + ) + elif not compressed: + ocp.type_handlers.register_type_handler( + jax.Array, NoCompressionArrayHandler(), override=True + ) + rank_logger.info( + "save_checkpoint: NoCompressionArrayHandler active (no D2H throttle) path=%s", + path, + ) + else: + rank_logger.info( + "save_checkpoint: stock Orbax ArrayHandler (no D2H throttle; " + "full addressable state staged to host at once) path=%s", + path, + ) def _callback(): ocp.type_handlers.register_type_handler(jax.Array, original_handler, override=True) @@ -105,22 +328,41 @@ def _callback(): lambda _: ocp.SaveArgs(chunk_byte_size=chunk_byte_size), state, ) - checkpointer.save( - dest, - args=ocp.args.PyTreeSave( - state, - save_args=save_args, - ), - ) + if kms_client is not None: + checkpointer.save(dest, args=ocp.args.PyTreeSave(state, save_args=save_args)) + else: + checkpointer.save( + dest, + args=ocp.args.PyTreeSave( + state, + save_args=save_args, + ), + ) - rank_logger.info("Started writing checkpoint to %s", path) + rank_logger.info( + "Started writing checkpoint to %s (save_concurrent_gb=%s, blocking=%s)", + path, + save_concurrent_gb, + blocking, + ) if blocking: + rank_logger.info("blocking wait until finished: waiting for Orbax write at %s", path) + t0 = time.time() checkpointer.wait_until_finished() + rank_logger.info( + "blocking wait until finished: Orbax write done in %.2fs at %s", + time.time() - t0, + path, + ) if callback and jax.process_index() == 0: callback() multihost_utils.sync_global_devices("blocking-checkpoint") def wait_until_finished(): - get_checkpointer().wait_until_finished() + if _CHECKPOINTER is not None: + _CHECKPOINTER.wait_until_finished() + _orbax_encrypted = sys.modules.get("xai_checkpointing.orbax_encrypted") + if _orbax_encrypted is not None: + _orbax_encrypted.wait_until_finished() diff --git a/phoenix/xrex/utils/log_timer.py b/phoenix/xrex/utils/log_timer.py index 47796b8a..84d66143 100644 --- a/phoenix/xrex/utils/log_timer.py +++ b/phoenix/xrex/utils/log_timer.py @@ -29,12 +29,13 @@ def duration(secs: float) -> str: class Timer: def __init__(self) -> None: self.start = time.time() + self._start_perf_counter = time.perf_counter() def __repr__(self) -> str: return duration(self.elapsed()) def elapsed(self) -> float: - return time.time() - self.start + return time.perf_counter() - self._start_perf_counter __float__ = elapsed diff --git a/phoenix/xrex/utils/log_util.py b/phoenix/xrex/utils/log_util.py index 0b7adc50..ffdb958e 100644 --- a/phoenix/xrex/utils/log_util.py +++ b/phoenix/xrex/utils/log_util.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import json import logging import os import socket @@ -40,6 +41,39 @@ def set_log_row(record_id: int | None) -> None: _LOG_ROW.set(str(record_id)) +_STANDARD_LOG_RECORD_ATTRS = set(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys()) +_SKIP_ATTRS = _STANDARD_LOG_RECORD_ATTRS | { + "message", + "asctime", + "taskName", + "rl_attr", +} + + +def _stringify(value: object) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (str, int, float)): + return str(value) + return json.dumps(value, default=str, separators=(",", ":")) + + +def _format_extra_kv(record: logging.LogRecord) -> str: + items: list[str] = [] + event: str | None = None + for key, value in record.__dict__.items(): + if key in _SKIP_ATTRS or value is None: + continue + rendered = _stringify(value) + if key == "event": + event = rendered + else: + items.append(f"{key}={rendered}") + if event is not None: + items.insert(0, f"event={event}") + return " ".join(items) + + def get_formatter( rank: int | None, prefix: str = "", worker_name: str | None = None ) -> logging.Formatter: @@ -82,7 +116,11 @@ def format(self, record): if row: attr += f"/row={row}" record.rl_attr = attr + extra_kv = _format_extra_kv(record) msg = super().format(record) + if extra_kv: + head, sep, tail = msg.partition("\n") + msg = f"{head} {extra_kv}{sep}{tail}" color = LEVEL_COLORS.get(record.levelno) if color: if msg.startswith(faint): diff --git a/phoenix/xrex/utils/metadata.py b/phoenix/xrex/utils/metadata.py index 721edd87..c4920ef9 100644 --- a/phoenix/xrex/utils/metadata.py +++ b/phoenix/xrex/utils/metadata.py @@ -22,6 +22,7 @@ from serde.json import from_dict, from_json, to_json from xrex import settings +from xrex.utils.checkpoint_cloud import ORBAX_TMP_DIR_SUFFIX from xrex.utils.launch_env import CHECKPOINT_DIR, XAI_USER logger = logging.getLogger(__name__) @@ -32,12 +33,14 @@ _ELAPSED_TOKENS = "elapsed_tokens" _CKPT_INDEX = "checkpoint_index" _CKPT_EXPIRY = "checkpoint_expiry" +_STEP = "step" COMPLETED_FILENAME = "completed" METADATA_FILENAME = "metadata.json" MetadataFromCheckpoint = namedtuple( - "MetadataFromCheckpoint", [_ELAPSED_SAMPLES, _CKPT_INDEX, _ELAPSED_TOKENS] + "MetadataFromCheckpoint", + [_ELAPSED_SAMPLES, _CKPT_INDEX, _ELAPSED_TOKENS, _STEP], ) @@ -56,7 +59,10 @@ def metadata_file_exists(checkpoint_path: Path): def read_checkpoint_metadata(metadata_path: Path) -> MetadataFromCheckpoint: record = json.loads(metadata_path.read_text()) return MetadataFromCheckpoint( - record[_ELAPSED_SAMPLES], record[_CKPT_INDEX], record[_ELAPSED_TOKENS] + record[_ELAPSED_SAMPLES], + record[_CKPT_INDEX], + record[_ELAPSED_TOKENS], + record.get(_STEP), ) @@ -66,6 +72,7 @@ def write_checkpoint_metadata( checkpoint_index: int, checkpoint_expiry: str, elapsed_tokens: int | None, + step: int, ) -> None: with metadata_path.open("x") as f: json.dump( @@ -74,6 +81,7 @@ def write_checkpoint_metadata( _CKPT_INDEX: checkpoint_index, _ELAPSED_TOKENS: elapsed_tokens, _CKPT_EXPIRY: checkpoint_expiry, + _STEP: step, }, f, ) @@ -92,7 +100,8 @@ def read_metadata_file(checkpoint_path: Path) -> MetadataFromCheckpoint | None: elapsed_samples = completion_path.read_text() checkpoint_index = None elapsed_tokens = None - return MetadataFromCheckpoint(elapsed_samples, checkpoint_index, elapsed_tokens) + step = None + return MetadataFromCheckpoint(elapsed_samples, checkpoint_index, elapsed_tokens, step) except Exception as e: logger.error(f"Unable to read checkpoint completion file: {str(e)}") return None @@ -104,6 +113,7 @@ def write_metadata_file( checkpoint_index: int, checkpoint_expiry: str, elapsed_tokens: int | None, + step: int, ): write_checkpoint_metadata( checkpoint_path / METADATA_FILENAME, @@ -111,6 +121,7 @@ def write_metadata_file( checkpoint_index, checkpoint_expiry, elapsed_tokens, + step, ) with (checkpoint_path / COMPLETED_FILENAME).open("x") as f: @@ -238,11 +249,40 @@ def guess_checkpoint_format(path): return "orbax" if (path / "ckpt-0" / "tensor00000_000").exists(): return "pickle" - if any(path.glob("orbax-ckpt*")): + if any(ORBAX_TMP_DIR_SUFFIX not in p.name for p in path.glob("orbax-ckpt*")): return "orbax" raise ValueError(f"Could not determine format of checkpoint at {path}") +def has_committed_payload(checkpoint_path: Path) -> bool: + if (checkpoint_path / "ckpt-0" / "tensor00000_000").exists(): + return True + final_names = set() + tmp_names = [] + for entry in checkpoint_path.glob("orbax-ckpt*"): + if ORBAX_TMP_DIR_SUFFIX in entry.name: + tmp_names.append(entry.name) + else: + final_names.add(entry.name) + if not final_names: + return False + return all(name.split(ORBAX_TMP_DIR_SUFFIX, 1)[0] in final_names for name in tmp_names) + + +def _is_loadable_checkpoint(checkpoint_path: Path) -> bool: + if not (checkpoint_path / COMPLETED_FILENAME).exists(): + return False + if has_committed_payload(checkpoint_path): + return True + logger.warning( + "Ignoring checkpoint at %s: it has a %r marker but its Orbax data was never" + " committed (crash between the marker write and the tmp-dir rename?)", + checkpoint_path, + COMPLETED_FILENAME, + ) + return False + + class MetadataProvider(ABC): @abstractmethod def discover_checkpoint( @@ -265,6 +305,7 @@ def record_checkpoint( checkpoint_index: int, checkpoint_ttl: int, elapsed_tokens: int, + step: int, ): ... @@ -290,13 +331,13 @@ def _search_for_latest_path_manual(search_dir: Path): if not m1: return None if m2 := PATH2_RE.search(str(search_dir)): - if (search_dir / COMPLETED_FILENAME).exists(): + if _is_loadable_checkpoint(search_dir): return m2["run_id"], search_dir, int(m1["samples"]) return None candidates: list[tuple[int, str, Path]] = [] for candidate, m2 in _matching_subdirs(search_dir, PATH2_RE): - if (candidate / COMPLETED_FILENAME).exists(): + if _is_loadable_checkpoint(candidate): candidates.append((int(m1["samples"]), m2["run_id"], candidate)) if candidates: @@ -315,14 +356,14 @@ def _search_for_latest_path(search_dir: Path) -> Optional[tuple[str, Path, int]] path = (search_dir / "latest").resolve() m2 = PATH2_RE.match(path.name) m1 = PATH1_RE.match(path.parent.name) - if m1 and m2 and (path / COMPLETED_FILENAME).exists(): + if m1 and m2 and _is_loadable_checkpoint(path): logger.info("Found via symlink %s", path) return m2["run_id"], path, int(m1["samples"]) candidates: list[tuple[int, str, Path]] = [] for path1, m1 in _matching_subdirs(search_dir, PATH1_RE): for candidate, m2 in _matching_subdirs(path1, PATH2_RE): - if (candidate / COMPLETED_FILENAME).exists(): + if _is_loadable_checkpoint(candidate): candidates.append((int(m1["samples"]), m2["run_id"], candidate)) if candidates: @@ -398,6 +439,7 @@ def record_checkpoint( checkpoint_index: int, checkpoint_ttl: int, elapsed_tokens: int, + step: int, ): base_dir = checkpoint_dir_or_default(base_dir) path = Path(run.checkpoint_path(base_dir, elapsed_samples)) @@ -418,7 +460,7 @@ def record_checkpoint( ).strftime("%Y-%m-%dT%H:%M:%SZ") write_metadata_file( - path, elapsed_samples, checkpoint_index, checkpoint_expiry, elapsed_tokens + path, elapsed_samples, checkpoint_index, checkpoint_expiry, elapsed_tokens, step ) rpath = os.path.join(base_dir, run.name) @@ -534,6 +576,7 @@ def record_completed_checkpoint( elapsed_samples: int, checkpoint_index: int, elapsed_tokens: int, + step: int, checkpoint_ttl: int = datetime.timedelta(weeks=2).total_seconds(), ): from xrex.utils.toolbox_notify import notify_checkpoint_async @@ -549,6 +592,7 @@ def record_completed_checkpoint( checkpoint_index=checkpoint_index, checkpoint_ttl=checkpoint_ttl, elapsed_tokens=elapsed_tokens, + step=step, ) path = run.checkpoint_path(base_dir, elapsed_samples) rank_logger.info(f"Recorded checkpoint {elapsed_samples=} {path=}") diff --git a/phoenix/xrex/utils/metrics.py b/phoenix/xrex/utils/metrics.py index 57fb72b3..58ee8172 100644 --- a/phoenix/xrex/utils/metrics.py +++ b/phoenix/xrex/utils/metrics.py @@ -98,6 +98,21 @@ def _find_moment_state(state): return None +def _moment_metrics(first_state) -> dict: + metrics = {} + for field, suffix in (("mu", "_mu"), ("nu", "_nu")): + tree = getattr(first_state, field, None) + if tree is None: + continue + flat = { + k: v + for k, v in flatten_dict(tree).items() + if isinstance(v, Parameter) and hasattr(v.x, "ndim") + } + metrics.update(split_metrics(flat, suffix)) + return metrics + + def norm_metrics(params, opt_state, gradients, updates): from xrex.optimizers.optim import InjectHyperparamsState @@ -105,12 +120,14 @@ def norm_metrics(params, opt_state, gradients, updates): metrics = {} if not isinstance(opt_state.inner_state, optax._src.combine.MultiTransformState): - _adam_state = _find_moment_state(opt_state.inner_state) - if _adam_state is not None: - mu_metric = split_metrics(flatten_dict(_adam_state.mu), "_mu") - metrics.update(mu_metric) - nu_metric = split_metrics(flatten_dict(_adam_state.nu), "_nu") - metrics.update(nu_metric) + optimizer_state = opt_state.inner_state[1] + if isinstance(optimizer_state, optax._src.combine.MultiTransformState): + for masked in optimizer_state.inner_states.values(): + metrics.update(_moment_metrics(masked.inner_state[0])) + else: + _adam_state = _find_moment_state(opt_state.inner_state) + if _adam_state is not None: + metrics.update(_moment_metrics(_adam_state)) weight_norm_metric = split_metrics(flatten_dict(params), "_norm") metrics.update(weight_norm_metric) diff --git a/scarecrow/legacy/DownstreamServices.scala b/scarecrow/legacy/DownstreamServices.scala index 44a28d3f..9a196240 100644 --- a/scarecrow/legacy/DownstreamServices.scala +++ b/scarecrow/legacy/DownstreamServices.scala @@ -75,4 +75,8 @@ object DownstreamServices { case object UrlToSlugStore extends DownstreamService { override val name = "URL_TO_SLUG_STORE" } + + case object XReviewIntake extends DownstreamService { + override val name = "XREVIEW_INTAKE" + } } diff --git a/scarecrow/legacy/XReviewIntakeRegistry.scala b/scarecrow/legacy/XReviewIntakeRegistry.scala new file mode 100644 index 00000000..7c08f0c5 --- /dev/null +++ b/scarecrow/legacy/XReviewIntakeRegistry.scala @@ -0,0 +1,179 @@ +package com.twitter.botmaker.app.scarecrow.legacy + +import com.google.inject.Exposed +import com.google.inject.Provides +import com.google.inject.Singleton +import com.twitter.botmaker.ASTNode +import com.twitter.botmaker.BotMakerFinatraModule.Migration +import com.twitter.botmaker.Context +import com.twitter.botmaker.DownstreamService +import com.twitter.botmaker.FunctionUnit5O2 +import com.twitter.botmaker.app.scarecrow.ScarecrowRuntime +import com.twitter.botmaker.compiler.ActionLevel +import com.twitter.botmaker.runtime.localMode +import com.twitter.botmaker.runtime.personalAccountMode +import com.twitter.finagle.mtls.authentication.ServiceIdentifier +import com.twitter.finagle.stats.StatsReceiver +import com.twitter.inject.Injector +import com.twitter.inject.TwitterPrivateModule +import com.twitter.inject.annotations.Flag +import com.twitter.useng.common.xreview.{ + XReviewIntakeClientModule => SharedXReviewIntakeClientModule +} +import com.twitter.util.Future +import java.lang.{Long => JLong} + +object XReviewIntakeRegistry extends TwitterPrivateModule { + + flag[String]( + "xreview.tls.ca-cert", + SharedXReviewIntakeClientModule.DefaultCaCertPath, + "CA cert path for XReview intake mTLS" + ) + + @Provides + @Singleton + @Exposed + def providesXReviewReportSubmitter( + serviceIdentifier: ServiceIdentifier, + @Migration statsReceiver: StatsReceiver, + @Flag("xreview.tls.ca-cert") caCertPath: String + ): XReviewReportSubmitter = { + if (localMode() || personalAccountMode()) { + XReviewReportSubmitter.noop(statsReceiver) + } else { + val prodClient = SharedXReviewIntakeClientModule.provideXReviewIntakeClient( + serviceIdentifier = serviceIdentifier, + statsReceiver = statsReceiver.scope("xreview"), + caCertPath = caCertPath + ) + val stagingClient = SharedXReviewIntakeClientModule.provideXReviewStagingIntakeClient( + serviceIdentifier = serviceIdentifier, + statsReceiver = statsReceiver.scope("xreview_staging"), + caCertPath = caCertPath + ) + XReviewReportSubmitter(prodClient, stagingClient, statsReceiver) + } + } + + override def singletonShutdown(injector: Injector): Unit = + injector.instance[XReviewReportSubmitter].close() +} + +object CreateXReviewReportProd + extends FunctionUnit5O2[ + ScarecrowRuntime, + String, + Long, + Long, + String, + Long, + String, + JLong, + Future[Unit] + ] { + + override def cacheLevel: ASTNode.CacheLevel = ASTNode.CacheLevel.Event + override def actionLevel: ActionLevel = ActionLevel.PROD_AGENT_WORKFLOW_ACTION + override def downstreams: Set[DownstreamService] = Set(DownstreamServices.XReviewIntake) + override def description: String = + "Submits a report to production XReview Intake over gRPC. " + + "entityType must be post or profile. reportType must be an XReview-allowlisted value. " + + "Impersonation bots should use bystander_impersonation (lane tags key on that type, " + + "not generic impersonation). Each evaluation includes detection_timestamp_ms so " + + "intake's content-hash report id is unique; evidence rollup then applies." + override def arguments: Seq[String] = Seq( + "reported entity type (post or profile)", + "reported entity id (tweet id or user id)", + "reported user id", + "XReview report_type (e.g. bystander_impersonation)", + "bot id (report-bag detection_bot_id; reporter_id is 0)", + "optional note", + "optional victim user id (victim_user_id)" + ) + override def examples: Seq[String] = Seq( + "CreateXReviewReportProd(\"profile\", :userId, :userId, \"bystander_impersonation\", :botId)", + "CreateXReviewReportProd(\"profile\", :userId, :userId, \"bystander_impersonation\", :botId, :note, :victimId)" + ) + + override def evaluate( + context: Context[ScarecrowRuntime], + entityType: String, + entityId: Long, + userId: Long, + reportType: String, + botId: Long, + note: Option[String], + victimId: Option[JLong] + ): Future[Unit] = { + context.getRuntime.fetcher20.xreviewReportSubmitter.submit( + entityType, + entityId, + userId, + reportType, + botId, + note, + victimId, + staging = false, + detectionTimestampMs = context.startMillis() + ) + } +} + +object CreateXReviewReportStaging + extends FunctionUnit5O2[ + ScarecrowRuntime, + String, + Long, + Long, + String, + Long, + String, + JLong, + Future[Unit] + ] { + + override def cacheLevel: ASTNode.CacheLevel = ASTNode.CacheLevel.Event + override def actionLevel: ActionLevel = ActionLevel.NO_ACTION + override def downstreams: Set[DownstreamService] = Set(DownstreamServices.XReviewIntake) + override def description: String = + "Submits a report to staging XReview Intake over gRPC. " + + "entityType must be post or profile. reportType must be an XReview-allowlisted value. " + + "Impersonation bots should use bystander_impersonation (lane tags key on that type, " + + "not generic impersonation)." + override def arguments: Seq[String] = Seq( + "reported entity type (post or profile)", + "reported entity id (tweet id or user id)", + "reported user id", + "XReview report_type (e.g. bystander_impersonation)", + "bot id (report-bag detection_bot_id; reporter_id is 0)", + "optional note", + "optional victim user id (victim_user_id)" + ) + override def examples: Seq[String] = Seq( + "CreateXReviewReportStaging(\"profile\", :userId, :userId, \"bystander_impersonation\", :botId)" + ) + + override def evaluate( + context: Context[ScarecrowRuntime], + entityType: String, + entityId: Long, + userId: Long, + reportType: String, + botId: Long, + note: Option[String], + victimId: Option[JLong] + ): Future[Unit] = { + context.getRuntime.fetcher20.xreviewReportSubmitter.submit( + entityType, + entityId, + userId, + reportType, + botId, + note, + victimId, + staging = true, + detectionTimestampMs = context.startMillis() + ) + } +} diff --git a/scarecrow/legacy/XReviewReportBuilder.scala b/scarecrow/legacy/XReviewReportBuilder.scala new file mode 100644 index 00000000..594a6a65 --- /dev/null +++ b/scarecrow/legacy/XReviewReportBuilder.scala @@ -0,0 +1,59 @@ +package com.twitter.botmaker.app.scarecrow.legacy + +import com.twitter.useng.common.xreview.XReviewReportKeys +import com.twitter.useng.common.xreview.XReviewReportKeys.pb +import intake_service.IntakeService.SubmitReportRequest +import java.lang.{Long => JLong} +import scala.jdk.CollectionConverters._ + +object XReviewReportBuilder { + + val ReportSurfaceValue = "x_app" + val ReportSourceValue = "proactive" + val Subject = "Botmaker report" + + val ReporterIdValue = "0" + + val DescriptionKey = "description" + + val PostEntityType = "post" + val ProfileEntityType = "profile" + val SupportedEntityTypes: Set[String] = Set(PostEntityType, ProfileEntityType) + + def normalizeEntityType(entityType: String): Option[String] = { + Option(entityType).map(_.trim.toLowerCase).filter(SupportedEntityTypes.contains) + } + + def toSubmitReportRequest( + entityType: String, + entityId: Long, + userId: Long, + reportType: String, + botId: Long, + note: Option[String], + victimId: Option[JLong], + detectionTimestampMs: Long + ): SubmitReportRequest = { + val reportFields = Seq( + pb(XReviewReportKeys.Subject, Subject), + pb(XReviewReportKeys.ReportSurface, ReportSurfaceValue), + pb(XReviewReportKeys.ReportType, reportType), + pb(XReviewReportKeys.ReportSource, ReportSourceValue), + pb(XReviewReportKeys.ReportedEntityType, entityType), + pb(XReviewReportKeys.ReportedEntityId, entityId.toString), + pb(XReviewReportKeys.ReportedUserId, userId.toString), + pb(XReviewReportKeys.DetectionBotId, botId.toString), + pb(XReviewReportKeys.DetectionTimestampMs, detectionTimestampMs.toString) + ) ++ + note.filter(_.nonEmpty).map(n => pb(DescriptionKey, n)).toSeq ++ + victimId.filter(_ != null).map(id => pb(XReviewReportKeys.VictimUserId, id.toString)).toSeq + + val reporterFields = Seq(pb(XReviewReportKeys.ReporterId, ReporterIdValue)) + + SubmitReportRequest + .newBuilder() + .addAllReport(reportFields.asJava) + .addAllReporter(reporterFields.asJava) + .build() + } +} diff --git a/scarecrow/legacy/XReviewReportSubmitter.scala b/scarecrow/legacy/XReviewReportSubmitter.scala new file mode 100644 index 00000000..cf142410 --- /dev/null +++ b/scarecrow/legacy/XReviewReportSubmitter.scala @@ -0,0 +1,91 @@ +package com.twitter.botmaker.app.scarecrow.legacy + +import com.twitter.finagle.stats.StatsReceiver +import com.twitter.useng.common.xreview.XReviewIntakeClient +import com.twitter.util.Future +import com.twitter.util.logging.Logging +import java.lang.{Long => JLong} +import scala.util.control.NonFatal + +class XReviewReportSubmitter( + prodClient: Option[XReviewIntakeClient], + stagingClient: Option[XReviewIntakeClient], + statsReceiver: StatsReceiver) + extends Logging { + + private val scoped = statsReceiver.scope("xreview", "botmaker") + private val attemptedCounter = scoped.counter("attempted") + private val submittedCounter = scoped.counter("submitted") + private val skippedCounter = scoped.counter("skipped") + private val unsupportedEntityCounter = scoped.counter("unsupported_entity") + private val rpcFailedCounter = scoped.counter("rpc_failed") + + def submit( + entityType: String, + entityId: Long, + userId: Long, + reportType: String, + botId: Long, + note: Option[String], + victimId: Option[JLong], + staging: Boolean, + detectionTimestampMs: Long + ): Future[Unit] = { + attemptedCounter.incr() + XReviewReportBuilder.normalizeEntityType(entityType) match { + case None => + unsupportedEntityCounter.incr() + Future.exception( + new IllegalArgumentException( + s"XReview entity type must be post or profile, got '$entityType'")) + case Some(normalizedType) => + clientFor(staging) match { + case None => + skippedCounter.incr() + Future.Unit + case Some(client) => + val request = XReviewReportBuilder.toSubmitReportRequest( + entityType = normalizedType, + entityId = entityId, + userId = userId, + reportType = reportType, + botId = botId, + note = note, + victimId = victimId, + detectionTimestampMs = detectionTimestampMs + ) + client + .submitReport(request) + .onSuccess(_ => submittedCounter.incr()) + .unit + .rescue { + case NonFatal(e) => + rpcFailedCounter.incr() + warn(s"[xreview] botmaker report submit failed: ${e.getMessage}", e) + Future.exception(e) + } + } + } + } + + def close(): Unit = { + prodClient.foreach(_.shutdown()) + stagingClient.foreach(_.shutdown()) + } + + private def clientFor(staging: Boolean): Option[XReviewIntakeClient] = + if (staging) stagingClient else prodClient +} + +object XReviewReportSubmitter { + + def apply( + prodClient: XReviewIntakeClient, + stagingClient: XReviewIntakeClient, + statsReceiver: StatsReceiver + ): XReviewReportSubmitter = + new XReviewReportSubmitter(Some(prodClient), Some(stagingClient), statsReceiver) + + def noop(statsReceiver: StatsReceiver): XReviewReportSubmitter = + new XReviewReportSubmitter(None, None, statsReceiver) +} diff --git a/thunder/args.rs b/thunder/args.rs index f328d40a..8e1c4963 100644 --- a/thunder/args.rs +++ b/thunder/args.rs @@ -69,6 +69,9 @@ pub struct Args { #[arg(long, default_value = "/s/kafka/phoenix-kafka-scram-bootstrap")] pub in_network_events_consumer_dest: String, + #[arg(long)] + pub in_network_events_consumer_mtls_zone: Option, + #[arg(long, default_value = "innetwork-posts")] pub o2_bucket: String, diff --git a/thunder/kafka_utils.rs b/thunder/kafka_utils.rs index 21c7c842..8ab020b2 100644 --- a/thunder/kafka_utils.rs +++ b/thunder/kafka_utils.rs @@ -1,7 +1,7 @@ use anyhow::{Context, Result}; use std::sync::Arc; use xai_kafka::config::SslConfig; -use xai_kafka::{KafkaConsumerBuilder, KafkaProducerBuilder}; +use xai_kafka::{KafkaConsumerBuilder, KafkaConsumerConfigBuilder, KafkaProducerBuilder}; use xai_wily::WilyConfig; use crate::{ @@ -15,9 +15,29 @@ use crate::{ const TWEET_EVENT_TOPIC: &str = "tweet_events"; const TWEET_EVENT_DEST: &str = "kafka.tweet-events.example.invalid"; +const IN_NETWORK_EVENTS_CLUSTER: &str = "phoenix"; const IN_NETWORK_EVENTS_DEST: &str = "kafka.phoenix-bootstrap.example.invalid"; const IN_NETWORK_EVENTS_TOPIC: &str = "innetwork_post"; +#[derive(Debug, PartialEq, Eq)] +enum InNetworkEventsConsumerAuth<'a> { + Scram, + Mtls { + cluster: &'static str, + zone: &'a str, + }, +} + +fn in_network_events_consumer_auth(args: &args::Args) -> InNetworkEventsConsumerAuth<'_> { + match args.in_network_events_consumer_mtls_zone.as_deref() { + Some(zone) => InNetworkEventsConsumerAuth::Mtls { + cluster: IN_NETWORK_EVENTS_CLUSTER, + zone, + }, + None => InNetworkEventsConsumerAuth::Scram, + } +} + pub async fn start_kafka( args: &args::Args, post_store: Arc, @@ -37,23 +57,39 @@ pub async fn start_kafka( if args.is_serving { let unique_id = uuid::Uuid::new_v4().to_string(); + let group_id = format!("{}-{}", args.kafka_group_id, unique_id); - let v2_tweet_events_consumer_config = KafkaConsumerBuilder::new( - args.in_network_events_consumer_dest.clone(), - IN_NETWORK_EVENTS_TOPIC.to_string(), - format!("{}-{}", args.kafka_group_id, unique_id), - ) - .with_wily_config(WilyConfig::default()) - .with_ssl(SslConfig { - security_protocol: args.security_protocol.clone(), - sasl_mechanism: Some(args.producer_sasl_mechanism.clone()), - sasl_username: Some(args.producer_sasl_username.clone()), - sasl_password: producer_sasl_password.clone(), - }) - .with_auto_offset_reset(args.auto_offset_reset.clone()) - .with_fetch_timeout_ms(args.fetch_timeout_ms) - .with_max_partition_fetch_bytes(1024 * 1024 * 100) - .with_skip_to_latest(args.skip_to_latest); + let consumer_builder = match in_network_events_consumer_auth(args) { + InNetworkEventsConsumerAuth::Mtls { cluster, zone } => { + KafkaConsumerConfigBuilder::for_cluster_mtls_auto( + cluster, + IN_NETWORK_EVENTS_TOPIC, + group_id.clone(), + Some(zone), + ) + .context("Failed to build Phoenix mTLS Kafka consumer config")? + } + InNetworkEventsConsumerAuth::Scram => KafkaConsumerConfigBuilder::new( + args.in_network_events_consumer_dest.clone(), + IN_NETWORK_EVENTS_TOPIC, + group_id, + ) + .with_wily_config(WilyConfig::default()) + .with_ssl(SslConfig { + security_protocol: args.security_protocol.clone(), + sasl_mechanism: Some(args.producer_sasl_mechanism.clone()), + sasl_username: Some(args.producer_sasl_username.clone()), + sasl_password: producer_sasl_password.clone(), + }), + }; + + let v2_tweet_events_consumer_config = consumer_builder + .with_auto_offset_reset(args.auto_offset_reset.clone()) + .with_enable_auto_offset_store(true) + .with_enable_auto_commit(true) + .with_fetch_timeout_ms(args.fetch_timeout_ms) + .with_max_partition_fetch_bytes(1024 * 1024 * 100) + .with_skip_to_latest(args.skip_to_latest); start_tweet_event_processing_v2( v2_tweet_events_consumer_config, @@ -100,3 +136,38 @@ pub async fn start_kafka( Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use clap::Parser; + + #[test] + fn capi_serving_consumer_uses_phoenix_mtls_in_explicit_zone() { + let args = args::Args::parse_from([ + "thunder", + "--kafka-group-id", + "thunder", + "--in-network-events-consumer-mtls-zone", + "atla", + ]); + + assert_eq!( + in_network_events_consumer_auth(&args), + InNetworkEventsConsumerAuth::Mtls { + cluster: "phoenix", + zone: "atla", + } + ); + } + + #[test] + fn legacy_serving_consumer_retains_scram_without_mtls_zone() { + let args = args::Args::parse_from(["thunder", "--kafka-group-id", "thunder"]); + + assert_eq!( + in_network_events_consumer_auth(&args), + InNetworkEventsConsumerAuth::Scram + ); + } +} diff --git a/visibility-filtering-client/tweet_safety_label.rs b/visibility-filtering-client/tweet_safety_label.rs index b6b89fd3..5227bf17 100644 --- a/visibility-filtering-client/tweet_safety_label.rs +++ b/visibility-filtering-client/tweet_safety_label.rs @@ -217,7 +217,7 @@ fn proto_to_safety_label(label: &vf_pb::SafetyLabel) -> SafetyLabel { } } -fn proto_to_safety_label_map(proto: &vf_pb::SafetyLabelMap) -> SafetyLabelMap { +pub(crate) fn proto_to_safety_label_map(proto: &vf_pb::SafetyLabelMap) -> SafetyLabelMap { proto .labels .iter() diff --git a/visibility-filtering-client/vf_client.rs b/visibility-filtering-client/vf_client.rs index f05e01d1..f9de7d25 100644 --- a/visibility-filtering-client/vf_client.rs +++ b/visibility-filtering-client/vf_client.rs @@ -1,5 +1,6 @@ use crate::discovery::{build_vf_channel, VfChannel, VfChannelError, VfChannelParams, VfDiscovery}; use crate::models::{Action, FilteredReason, SafetyResult}; +use crate::tweet_safety_label::{proto_to_safety_label_map, SafetyLabelFailure}; use anyhow::{anyhow, Result}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -12,6 +13,7 @@ use thrift::protocol::{ use tonic::async_trait; use tonic::codec::CompressionEncoding; use tonic::transport::Channel; +use xai_safety_label_store::types::SafetyLabelMap; use xai_stats_receiver::global_stats_receiver; use xai_strato::{ decode, encode, MValCodec, StratoGrpc, StratoGrpcConfig, StratoResult, StratoValue, @@ -103,6 +105,12 @@ impl MValCodec for SafetyLevel { } } +#[derive(Debug, Clone)] +pub struct TweetVisibility { + pub reason: Option, + pub safety_labels: Result, +} + #[async_trait] pub trait VfClient { async fn get_result( @@ -111,7 +119,7 @@ pub trait VfClient { safety_level: SafetyLevel, for_user_id: u64, context: Option, - ) -> HashMap>>; + ) -> HashMap>; } pub struct StratoVfClient { @@ -154,7 +162,7 @@ impl VfClient for StratoVfClient { safety_level: SafetyLevel, for_user_id: u64, context: Option, - ) -> HashMap>> { + ) -> HashMap> { let client = &self.grpc_client; let view = VisibilityFilteringLookupContext { safety_level, @@ -172,13 +180,16 @@ impl VfClient for StratoVfClient { }) .collect::>)>>(); let result_batch = client.batch_call(calls, context.as_ref()).await; - let mut result_map: HashMap>> = HashMap::new(); + let mut result_map: HashMap> = HashMap::new(); for (tweet_id, bytes_result) in tweet_ids.iter().zip(result_batch) { let item_result = match bytes_result { Ok(bytes) => { let decoded: StratoResult> = decode(&bytes); match decoded { - StratoResult::Ok(strato_value) => Ok(strato_value.v), + StratoResult::Ok(strato_value) => Ok(TweetVisibility { + reason: strato_value.v, + safety_labels: Err(SafetyLabelFailure::LookupFailed), + }), StratoResult::Err(err) => { Err(anyhow!("Strato error code {}: {}", err.code, err.message)) } @@ -315,17 +326,33 @@ fn result_to_reason(result: vf_pb::TweetVisibilityResult) -> Option TweetVisibility { + let safety_labels = result + .safety_labels + .take() + .map(|m| proto_to_safety_label_map(&m)) + .ok_or(SafetyLabelFailure::LookupFailed); + TweetVisibility { + reason: result_to_reason(result), + safety_labels, + } +} + fn results_to_map( requested_tweet_ids: &[u64], results: Vec, -) -> HashMap>> { - let mut map: HashMap>> = results +) -> HashMap> { + let mut map: HashMap> = results .into_iter() - .map(|r| (r.tweet_id, Ok(result_to_reason(r)))) + .map(|r| (r.tweet_id, Ok(result_to_visibility(r)))) .collect(); for &tweet_id in requested_tweet_ids { - map.entry(tweet_id) - .or_insert_with(|| Ok(Some(FilteredReason::UnspecifiedReason))); + map.entry(tweet_id).or_insert_with(|| { + Ok(TweetVisibility { + reason: Some(FilteredReason::UnspecifiedReason), + safety_labels: Err(SafetyLabelFailure::LookupFailed), + }) + }); } map } @@ -333,7 +360,7 @@ fn results_to_map( fn rpc_error_map( tweet_ids: &[u64], status: &tonic::Status, -) -> HashMap>> { +) -> HashMap> { tweet_ids .iter() .map(|&tweet_id| { @@ -397,7 +424,7 @@ impl VfClient for XaiVfClient { safety_level: SafetyLevel, for_user_id: u64, context: Option, - ) -> HashMap>> { + ) -> HashMap> { if tweet_ids.is_empty() { return HashMap::new(); } @@ -471,7 +498,7 @@ impl VfClient for MockVfClient { _safety_level: SafetyLevel, _for_user_id: u64, _context: Option, - ) -> HashMap>> { + ) -> HashMap> { HashMap::new() } } @@ -631,13 +658,22 @@ mod rust_vf_tests { let map = results_to_map(&[1, 2, 3, 4], results); - assert!(matches!(map.get(&1), Some(Ok(None)))); + assert!(matches!( + map.get(&1), + Some(Ok(TweetVisibility { reason: None, .. })) + )); assert!(matches!( map.get(&2), - Some(Ok(Some(FilteredReason::AuthorIsUnsafe))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::AuthorIsUnsafe), + .. + })) )); match map.get(&3) { - Some(Ok(Some(FilteredReason::SafetyResult(sr)))) => { + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::SafetyResult(sr)), + .. + })) => { assert!(matches!(sr.action, Action::Drop(_))); } other => panic!("unexpected mapping for id 3: {other:?}"), @@ -645,12 +681,75 @@ mod rust_vf_tests { assert!( matches!( map.get(&4), - Some(Ok(Some(FilteredReason::UnspecifiedReason))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::UnspecifiedReason), + .. + })) ), "missing response ids fail closed" ); } + #[test] + fn results_to_map_converts_labels_and_fails_closed() { + use xai_x_thrift::tweet_safety_label::SafetyLabelType; + + let results = vec![ + vf_pb::TweetVisibilityResult { + tweet_id: 1, + action: Some(vf_pb::Action { + kind: Some(vf_pb::action::Kind::Allow(true)), + }), + filtered_reason: None, + safety_labels: Some(vf_pb::SafetyLabelMap { + labels: HashMap::from([( + i32::from(SafetyLabelType::NSFW_HIGH_PRECISION), + vf_pb::SafetyLabel { + source: Some("some rule".to_string()), + ..Default::default() + }, + )]), + }), + }, + vf_pb::TweetVisibilityResult { + tweet_id: 2, + action: Some(vf_pb::Action { + kind: Some(vf_pb::action::Kind::Allow(true)), + }), + filtered_reason: None, + safety_labels: None, + }, + ]; + + let map = results_to_map(&[1, 2, 3], results); + + match map.get(&1) { + Some(Ok(r)) => { + assert_eq!(r.reason, None); + let labels = r.safety_labels.as_ref().expect("labels present"); + let label = labels + .get(&SafetyLabelType::NSFW_HIGH_PRECISION) + .expect("label converted"); + assert_eq!(label.source.as_deref(), Some("some rule")); + } + other => panic!("unexpected mapping for id 1: {other:?}"), + } + match map.get(&2) { + Some(Ok(r)) => { + assert_eq!(r.reason, None); + assert!(r.safety_labels.is_err(), "absent map stays unavailable"); + } + other => panic!("unexpected mapping for id 2: {other:?}"), + } + match map.get(&3) { + Some(Ok(r)) => { + assert_eq!(r.reason, Some(FilteredReason::UnspecifiedReason)); + assert!(r.safety_labels.is_err(), "missing ids fail closed"); + } + other => panic!("unexpected mapping for id 3: {other:?}"), + } + } + #[test] fn interstitial_wraps_as_safety_result() { let results = vec![vf_pb::TweetVisibilityResult { @@ -667,7 +766,10 @@ mod rust_vf_tests { let map = results_to_map(&[7], results); match map.get(&7) { - Some(Ok(Some(FilteredReason::SafetyResult(sr)))) => { + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::SafetyResult(sr)), + .. + })) => { assert!(matches!(sr.action, Action::Interstitial)); } other => panic!("interstitial should wrap as SafetyResult, got: {other:?}"), @@ -753,11 +855,17 @@ mod rust_vf_tests { assert_eq!(result.len(), 2); assert!(matches!( result.get(&10), - Some(Ok(Some(FilteredReason::AuthorIsUnsafe))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::AuthorIsUnsafe), + .. + })) )); assert!(matches!( result.get(&20), - Some(Ok(Some(FilteredReason::AuthorIsUnsafe))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::AuthorIsUnsafe), + .. + })) )); handle.abort(); } @@ -896,11 +1004,17 @@ mod rust_vf_tests { assert_eq!(result.len(), 2); assert!(matches!( result.get(&10), - Some(Ok(Some(FilteredReason::AuthorIsUnsafe))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::AuthorIsUnsafe), + .. + })) )); assert!(matches!( result.get(&20), - Some(Ok(Some(FilteredReason::UnspecifiedReason))) + Some(Ok(TweetVisibility { + reason: Some(FilteredReason::UnspecifiedReason), + .. + })) )); handle.abort(); } @@ -976,7 +1090,10 @@ mod rust_vf_tests { assert_eq!(result.len(), tweet_ids.len()); for id in 1..=XAI_VF_MAX_BATCH_SIZE as u64 { assert!( - matches!(result.get(&id), Some(Ok(None))), + matches!( + result.get(&id), + Some(Ok(TweetVisibility { reason: None, .. })) + ), "chunk-1 id {id} should Allow" ); } diff --git a/visibility-filtering/clients/gizmoduck_client.rs b/visibility-filtering/clients/gizmoduck_client.rs index ac7a2d52..e8fbb703 100644 --- a/visibility-filtering/clients/gizmoduck_client.rs +++ b/visibility-filtering/clients/gizmoduck_client.rs @@ -9,7 +9,7 @@ pub struct GizmoduckLookup { fn author_hydration_lookup_context() -> LookupContext { LookupContext { - for_user_id: 0, + for_user_id: None, include_deactivated: true, include_failed: true, include_erased: true, @@ -46,7 +46,7 @@ mod tests { #[test] fn author_lookup_context_sets_include_flags() { let ctx = author_hydration_lookup_context(); - assert_eq!(ctx.for_user_id, 0); + assert_eq!(ctx.for_user_id, None); assert!(ctx.include_deactivated); assert!(ctx.include_failed); assert!(ctx.include_erased); diff --git a/visibility-filtering/clients/socialgraph_client.rs b/visibility-filtering/clients/socialgraph_client.rs index f9da5ccf..cf559f09 100644 --- a/visibility-filtering/clients/socialgraph_client.rs +++ b/visibility-filtering/clients/socialgraph_client.rs @@ -29,58 +29,40 @@ pub trait SocialgraphClient: Send + Sync { ) -> HashMap; } -#[derive(Default)] -pub struct MockSocialgraphClient { - pub relationships: HashMap<(u64, u64), ViewerAuthorRelationship>, - pub super_follows: HashMap<(u64, u64), bool>, -} +#[cfg(test)] +pub struct FakeSocialgraphClient; +#[cfg(test)] #[async_trait] -impl SocialgraphClient for MockSocialgraphClient { +impl SocialgraphClient for FakeSocialgraphClient { async fn batch_check_relationships( &self, - viewer_id: u64, + _viewer_id: u64, author_ids: &[u64], ) -> HashMap { author_ids .iter() - .map(|&author_id| { - let rel = self - .relationships - .get(&(viewer_id, author_id)) - .cloned() - .unwrap_or_default(); - (author_id, rel) - }) + .map(|&author_id| (author_id, ViewerAuthorRelationship::default())) .collect() } async fn batch_check_super_follows( &self, - viewer_id: u64, + _viewer_id: u64, author_ids: &[u64], ) -> HashMap { author_ids .iter() - .map(|&author_id| { - let follows = self - .super_follows - .get(&(viewer_id, author_id)) - .copied() - .unwrap_or(false); - (author_id, follows) - }) + .map(|&author_id| (author_id, false)) .collect() } } fn decode_packed_ids(packed: &[u8]) -> HashSet { - packed - .chunks_exact(8) - .map(|chunk| { - let arr: [u8; 8] = chunk.try_into().unwrap(); - i64::from_le_bytes(arr) as u64 - }) + let (chunks, _remainder) = packed.as_chunks::<8>(); + chunks + .iter() + .map(|&arr| i64::from_le_bytes(arr) as u64) .collect() } diff --git a/visibility-filtering/config.rs b/visibility-filtering/config.rs index 8646f3ab..58fb1016 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -1,45 +1,51 @@ -pub const ENV_GRPC_MTLS_ENABLED: &str = "GRPC_MTLS_ENABLED"; -pub const ENV_GRPC_MTLS_SERVER_KEY_PATH: &str = "GRPC_MTLS_SERVER_KEY_PATH"; -pub const ENV_GRPC_MTLS_SERVER_CRT_PATH: &str = "GRPC_MTLS_SERVER_CRT_PATH"; -pub const ENV_GRPC_MTLS_SERVER_CHAIN_PATH: &str = "GRPC_MTLS_SERVER_CHAIN_PATH"; -pub const ENV_GRPC_MTLS_CLIENT_CA_PATH: &str = "GRPC_MTLS_CLIENT_CA_PATH"; -pub const ENV_FALLBACK_CACHE_SERVE_STALE_ENABLED: &str = "VF_FALLBACK_CACHE_SERVE_STALE_ENABLED"; -pub const ENV_FALLBACK_CACHE_POPULATE_ENABLED: &str = "VF_FALLBACK_CACHE_POPULATE_ENABLED"; -pub const ENV_MEDIA_FALLBACK_CACHE_SERVE_STALE_ENABLED: &str = - "VF_MEDIA_FALLBACK_CACHE_SERVE_STALE_ENABLED"; -pub const ENV_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED: &str = - "VF_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED"; +pub const ENV_DUAL_CALL_HARNESS_ENABLED: &str = "VF_DUAL_CALL_HARNESS_ENABLED"; +pub const ENV_FALLBACK_CACHE_ENABLED: &str = "VF_FALLBACK_CACHE_ENABLED"; +pub const ENV_CACHE_WARM_ENABLED: &str = "VF_CACHE_WARM_ENABLED"; +pub const ENV_APP_ENV: &str = "APP_ENV"; +pub const ENV_FS_PATH: &str = "VF_FS_PATH"; +pub const ENV_GIZMODUCK_CLIENT_ID: &str = "VF_GIZMODUCK_CLIENT_ID"; +pub const ENV_TWEMCACHE_CLIENT_NAME: &str = "VF_TWEMCACHE_CLIENT_NAME"; + +const DEFAULT_FS_PATH: &str = "/usr/local/config/features/visibility/main/rust_vf.yml"; + +pub fn fs_path() -> String { + std::env::var(ENV_FS_PATH).unwrap_or_else(|_| DEFAULT_FS_PATH.to_string()) +} -pub fn fallback_cache_serve_stale_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_FALLBACK_CACHE_SERVE_STALE_ENABLED) - .ok() - .as_deref(), +pub(crate) fn gizmoduck_client_id() -> String { + resolve_gizmoduck_client_id( + std::env::var(ENV_GIZMODUCK_CLIENT_ID).ok().as_deref(), + std::env::var(ENV_APP_ENV).ok().as_deref(), ) } -pub fn fallback_cache_populate_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_FALLBACK_CACHE_POPULATE_ENABLED) - .ok() - .as_deref(), - ) +pub(crate) fn twemcache_client_name() -> String { + resolve_twemcache_client_name(std::env::var(ENV_TWEMCACHE_CLIENT_NAME).ok().as_deref()) } -pub fn media_fallback_cache_serve_stale_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_MEDIA_FALLBACK_CACHE_SERVE_STALE_ENABLED) - .ok() - .as_deref(), - ) +pub fn resolve_gizmoduck_client_id(configured: Option<&str>, app_env: Option<&str>) -> String { + match configured { + Some(id) => id.to_string(), + None => format!("visibility-filtering-service.{}", app_env.unwrap_or("prod")), + } } -pub fn media_fallback_cache_populate_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED) - .ok() - .as_deref(), - ) +pub fn resolve_twemcache_client_name(configured: Option<&str>) -> String { + configured + .unwrap_or("visibility-filtering-service") + .to_string() +} + +pub(crate) fn dual_call_harness_enabled() -> bool { + parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) +} + +pub(crate) fn fallback_cache_enabled() -> bool { + parse_env_flag(std::env::var(ENV_FALLBACK_CACHE_ENABLED).ok().as_deref()) +} + +pub(crate) fn cache_warm_enabled() -> bool { + parse_env_flag(std::env::var(ENV_CACHE_WARM_ENABLED).ok().as_deref()) } fn parse_env_flag(value: Option<&str>) -> bool { @@ -51,74 +57,33 @@ fn parse_env_flag(value: Option<&str>) -> bool { }) } -#[derive(Debug, Clone)] -pub struct GrpcMtlsConfig { - pub server_key_path: String, - pub server_crt_path: String, - pub server_chain_path: Option, - pub client_ca_path: String, -} - -impl GrpcMtlsConfig { - pub fn from_env() -> anyhow::Result> { - let enabled = parse_env_flag(std::env::var(ENV_GRPC_MTLS_ENABLED).ok().as_deref()); - - if !enabled { - return Ok(None); - } - - let server_key_path = std::env::var(ENV_GRPC_MTLS_SERVER_KEY_PATH) - .ok() - .filter(|v| !v.is_empty()) - .ok_or_else(|| anyhow::anyhow!("{ENV_GRPC_MTLS_SERVER_KEY_PATH} must be set"))?; - - let server_crt_path = std::env::var(ENV_GRPC_MTLS_SERVER_CRT_PATH) - .ok() - .filter(|v| !v.is_empty()) - .ok_or_else(|| anyhow::anyhow!("{ENV_GRPC_MTLS_SERVER_CRT_PATH} must be set"))?; - - let server_chain_path = std::env::var(ENV_GRPC_MTLS_SERVER_CHAIN_PATH) - .ok() - .filter(|v| !v.is_empty()); - - let client_ca_path = std::env::var(ENV_GRPC_MTLS_CLIENT_CA_PATH) - .ok() - .filter(|v| !v.is_empty()) - .ok_or_else(|| anyhow::anyhow!("{ENV_GRPC_MTLS_CLIENT_CA_PATH} must be set"))?; - - Ok(Some(Self { - server_key_path, - server_crt_path, - server_chain_path, - client_ca_path, - })) - } - - pub fn server_tls_config(&self) -> anyhow::Result { - let mut cert_pem = std::fs::read(&self.server_crt_path)?; - let key_pem = std::fs::read(&self.server_key_path)?; - let client_ca_pem = std::fs::read(&self.client_ca_path)?; - - if let Some(chain_path) = self.server_chain_path.as_ref() { - let chain_pem = std::fs::read(chain_path)?; - if !cert_pem.ends_with(b"\n") { - cert_pem.push(b'\n'); - } - cert_pem.extend_from_slice(&chain_pem); - } - - let identity = tonic::transport::Identity::from_pem(cert_pem, key_pem); - let client_ca = tonic::transport::Certificate::from_pem(client_ca_pem); - - Ok(tonic::transport::ServerTlsConfig::new() - .identity(identity) - .client_ca_root(client_ca)) - } -} - #[cfg(test)] mod tests { - use super::parse_env_flag; + use super::{parse_env_flag, resolve_gizmoduck_client_id, resolve_twemcache_client_name}; + + #[test] + fn client_ids_default_to_historical_values_and_overrides_win() { + assert_eq!( + resolve_gizmoduck_client_id(None, Some("prod")), + "visibility-filtering-service.prod" + ); + assert_eq!( + resolve_gizmoduck_client_id(None, Some("staging")), + "visibility-filtering-service.staging" + ); + assert_eq!( + resolve_gizmoduck_client_id(None, None), + "visibility-filtering-service.prod" + ); + assert_eq!( + resolve_twemcache_client_name(None), + "visibility-filtering-service" + ); + assert_eq!( + resolve_gizmoduck_client_id(Some("xai-vf-service.staging"), Some("staging")), + "xai-vf-service.staging" + ); + } #[test] fn parses_enabled_environment_values() { diff --git a/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs new file mode 100644 index 00000000..e6562531 --- /dev/null +++ b/visibility-filtering/dark_traffic_setup.rs @@ -0,0 +1,342 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use envoy_types::pb::envoy::config::core::v3::Node; +use envoy_types::pb::envoy::config::listener::v3::Listener; +use envoy_types::pb::envoy::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient; +use envoy_types::pb::envoy::service::discovery::v3::DiscoveryRequest; +use prost::Message; +use tower::util::Either; +use tracing::info; + +use tonic::transport::Channel; +use xai_dark_traffic::{DarkTrafficLayer, ReloadableDarkTrafficConfigBuilder}; +use xai_x_rpc::dynamic_channel_manager::{ + ChannelFactory, DynamicChannelManager, EndpointDiscovery, EndpointInfo, +}; +use xai_x_rpc::grpc_client::TlsMode; +use xai_x_rpc::xds_channel_factory::XdsChannelFactory; + +const CONFIG_PATH: &str = "/config/dark-traffic/dark_traffic.yaml"; + +pub const MIRROR_NAMESPACE: &str = "visibility"; +pub const STAGING_APP_ENV: &str = "staging"; +pub const DEVEL_APP_ENV: &str = "devel"; +pub const MIRROR_PORT_ID: &str = "grpc"; +pub const MIRROR_WORKLOAD_PREFIX: &str = "xai-vf-service"; + +const LISTENER_TYPE_URL: &str = "type.googleapis.com/envoy.config.listener.v3.Listener"; +const LDS_MAX_DECODING_MESSAGE_SIZE: usize = 256 * 1024 * 1024; +const LDS_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30); +const CHANNEL_CREATE_TIMEOUT: Duration = Duration::from_secs(30); + +pub fn staging_tls_domain(dc: &str) -> String { + format!("visibility.visibility-filtering-service.staging.{dc}.s2s.twttr.net") +} + +pub fn devel_tls_domain(dc: &str) -> String { + format!("visibility.xai-vf-service.devel.{dc}.s2s.twttr.net") +} + +pub type DarkLayer = Either; + +pub fn parse_mirror_listener(listener_name: &str) -> Option { + let dest = listener_name.rsplit('/').next().unwrap_or(listener_name); + let dest = dest.split('?').next().unwrap_or(dest); + let (workload, app_env) = [STAGING_APP_ENV, DEVEL_APP_ENV] + .into_iter() + .find_map(|app_env| { + let suffix = format!(".{app_env}.{MIRROR_NAMESPACE}:{MIRROR_PORT_ID}"); + dest.strip_suffix(&suffix) + .map(|workload| (workload, app_env)) + })?; + if !workload.starts_with(MIRROR_WORKLOAD_PREFIX) || workload.contains('.') { + return None; + } + let name = if app_env == STAGING_APP_ENV { + workload.to_string() + } else { + format!("{workload}.{app_env}") + }; + Some(EndpointInfo { + name, + xds_dest: dest.to_string(), + }) +} + +struct XdsMirrorDiscovery { + server_uri: String, +} + +#[async_trait::async_trait] +impl EndpointDiscovery for XdsMirrorDiscovery { + async fn discover(&self) -> anyhow::Result> { + tokio::time::timeout(LDS_RESPONSE_TIMEOUT, self.fetch()) + .await + .context("wildcard LDS exchange timed out")? + } +} + +impl XdsMirrorDiscovery { + async fn fetch(&self) -> anyhow::Result> { + let channel = tonic::transport::Endpoint::from_shared(self.server_uri.clone()) + .context("invalid kube-discovery URI")? + .connect() + .await + .context("failed to connect to kube-discovery")?; + + let mut client = AggregatedDiscoveryServiceClient::new(channel) + .max_decoding_message_size(LDS_MAX_DECODING_MESSAGE_SIZE); + + let request = DiscoveryRequest { + type_url: LISTENER_TYPE_URL.to_string(), + resource_names: vec!["*".to_string()], + node: Some(Node { + id: format!("{MIRROR_WORKLOAD_PREFIX}-dark-traffic"), + cluster: MIRROR_NAMESPACE.to_string(), + ..Default::default() + }), + ..Default::default() + }; + + use futures::StreamExt; + let requests = futures::stream::iter([request]).chain(futures::stream::pending()); + let mut stream = client + .stream_aggregated_resources(requests) + .await + .context("wildcard LDS stream failed to open")? + .into_inner(); + + let response = loop { + let response = stream + .message() + .await + .context("wildcard LDS stream errored")? + .context("wildcard LDS stream ended without a listener-carrying response")?; + if !response.resources.is_empty() { + break response; + } + }; + + let names: Vec = response + .resources + .iter() + .filter_map(|any| Some(Listener::decode(any.value.as_ref()).ok()?.name)) + .collect(); + anyhow::ensure!( + !names.is_empty(), + "none of {} LDS resources decoded as Listener", + response.resources.len() + ); + let endpoints: Vec = names + .iter() + .filter_map(|name| parse_mirror_listener(name)) + .collect(); + + if endpoints.is_empty() { + tracing::warn!( + listeners = names.len(), + "dark_traffic: no staging or devel listeners matched" + ); + } else { + info!( + names = %endpoints.iter().map(|e| e.name.as_str()).collect::>().join(", "), + "dark_traffic: discovery complete" + ); + } + + Ok(endpoints) + } +} + +struct MirrorChannelFactory { + staging: XdsChannelFactory, + devel: XdsChannelFactory, +} + +#[async_trait::async_trait] +impl ChannelFactory for MirrorChannelFactory { + async fn create_channel(&self, ep: &EndpointInfo) -> anyhow::Result { + let factory = if ep.xds_dest.contains(&format!(".{DEVEL_APP_ENV}.")) { + &self.devel + } else { + &self.staging + }; + tokio::time::timeout(CHANNEL_CREATE_TIMEOUT, factory.create_channel(ep)) + .await + .with_context(|| format!("channel dial timed out for {}", ep.xds_dest))? + } +} + +pub fn resolve_layer() -> DarkLayer { + if !std::env::var("DARK_TRAFFIC_ENABLED") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false) + { + info!("dark_traffic: disabled"); + return Either::Right(tower::layer::util::Identity::new()); + } + + let workload = std::env::var("WORKLOAD_NAME").ok(); + let max_ordinal: Option = std::env::var("DARK_TRAFFIC_MAX_ORDINAL") + .ok() + .and_then(|s| s.parse().ok()); + let ordinal: Option = std::env::var("ORDINAL_NUMBER") + .ok() + .and_then(|s| s.parse().ok()); + if !should_mirror(workload.as_deref(), ordinal, max_ordinal) { + info!( + ?workload, + ?ordinal, + ?max_ordinal, + "dark_traffic: disabled (not a mirror host)" + ); + return Either::Right(tower::layer::util::Identity::new()); + } + + let dc = std::env::var("DATACENTER").unwrap_or_else(|_| "atla".to_string()); + let discovery = XdsMirrorDiscovery { + server_uri: format!("http://frontend.kube-discovery.prod.svc.{dc}.kube.int-x.ai:8082"), + }; + let staging_domain = staging_tls_domain(&dc); + let devel_domain = devel_tls_domain(&dc); + info!(staging_domain, devel_domain, "dark_traffic: enabled"); + + #[expect(clippy::expect_used, reason = "startup fail-fast: TLS is required")] + let tls = TlsMode::mtls_from_env().expect("S2S TLS config required"); + let factory = MirrorChannelFactory { + staging: XdsChannelFactory::new(tls.clone().with_domain_override(&staging_domain)), + devel: XdsChannelFactory::new(tls.with_domain_override(&devel_domain)), + }; + + let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(discovery)); + + let config = ReloadableDarkTrafficConfigBuilder::new(CONFIG_PATH) + .forwarders({ + let ch = Arc::clone(&channels); + move || ch.channels() + }) + .build(); + + Either::Left(DarkTrafficLayer::new(config)) +} + +fn should_mirror(workload: Option<&str>, ordinal: Option, max_ordinal: Option) -> bool { + let Some(workload) = workload else { + return false; + }; + if workload.ends_with("-canary") { + return false; + } + let Some(ord) = ordinal else { + return false; + }; + ord < max_ordinal.unwrap_or(1) +} + +#[cfg(test)] +mod tests { + use super::*; + + const PROD: Option<&str> = Some("xai-vf-service"); + const CANARY: Option<&str> = Some("xai-vf-service-canary"); + + #[test] + fn default_only_pod0() { + assert!(should_mirror(PROD, Some(0), None)); + assert!(!should_mirror(PROD, Some(1), None)); + assert!(!should_mirror(PROD, Some(99), None)); + } + + #[test] + fn no_ordinal_disables() { + assert!(!should_mirror(PROD, None, None)); + assert!(!should_mirror(PROD, None, Some(3))); + } + + #[test] + fn max_ordinal_threshold() { + assert!(should_mirror(PROD, Some(0), Some(3))); + assert!(should_mirror(PROD, Some(1), Some(3))); + assert!(should_mirror(PROD, Some(2), Some(3))); + assert!(!should_mirror(PROD, Some(3), Some(3))); + assert!(!should_mirror(PROD, Some(4), Some(3))); + } + + #[test] + fn max_ordinal_zero_disables_all() { + assert!(!should_mirror(PROD, Some(0), Some(0))); + } + + #[test] + fn canary_never_mirrors() { + assert!(!should_mirror(CANARY, Some(0), Some(11))); + assert!(!should_mirror(CANARY, Some(0), Some(u32::MAX))); + } + + #[test] + fn missing_workload_name_fails_closed() { + assert!(!should_mirror(None, Some(0), Some(11))); + } + + #[test] + fn parse_accepts_vf_staging_and_devel_listeners() { + for (listener, workload) in [ + ("xai-vf-service.staging.visibility:grpc", "xai-vf-service"), + ( + "xai-vf-service.devel.visibility:grpc", + "xai-vf-service.devel", + ), + ( + "xai-vf-service-pr-123.devel.visibility:grpc", + "xai-vf-service-pr-123.devel", + ), + ( + "xdstp://kube-discovery/envoy.config.listener.v3.Listener/xai-vf-service-pr-123.devel.visibility:grpc?key=val", + "xai-vf-service-pr-123.devel", + ), + ( + "xai-vf-service-user1-foo.staging.visibility:grpc", + "xai-vf-service-user1-foo", + ), + ( + "xdstp://kube-discovery/envoy.config.listener.v3.Listener/xai-vf-service.staging.visibility:grpc?key=val", + "xai-vf-service", + ), + ] { + let ep = parse_mirror_listener(listener).expect(listener); + assert_eq!(ep.name, workload); + assert_eq!( + ep.xds_dest, + listener + .rsplit("/") + .next() + .unwrap() + .split("?") + .next() + .unwrap() + ); + } + } + + #[test] + fn parse_rejects_out_of_scope_listeners() { + for name in [ + "xai-vf-service.prod.visibility:grpc", + "xai-vf-service.development.visibility:grpc", + "xai-vf-service.staging-devel.visibility:grpc", + "xai-vf-service.devel-staging.visibility:grpc", + "xai-vf-service.devel.visibility:grpc-extra", + "xai-vf-service.devel.visibility:grpc.other", + "other-svc.staging.visibility:grpc", + "xai-vf-service.staging.other:grpc", + "xai-vf-service.staging.visibility:metrics", + "evil.xai-vf-service.staging.visibility:grpc", + "xai-vf-service.staging.visibility", + "", + ] { + assert!(parse_mirror_listener(name).is_none(), "{name}"); + } + } +} diff --git a/visibility-filtering/filter.rs b/visibility-filtering/filter.rs index c2d6bb82..2b1b2f3b 100644 --- a/visibility-filtering/filter.rs +++ b/visibility-filtering/filter.rs @@ -1,9 +1,8 @@ use crate::hydration::{HydrationOutput, HydrationPipeline, HydrationRequest}; -use crate::models::{RawCandidate, TweetId, VfAction}; +use crate::models::{RawCandidate, TweetId}; use crate::rules::metrics as ft_metrics; -use crate::rules::{Policies, SafetyLevel, Verdict}; +use crate::rules::{RuleEngine, SafetyLevel, Verdict}; use std::collections::HashMap; -use tracing::{debug, info}; use xai_visibility_filtering_proto as vf_pb; pub struct FilterRequest { @@ -19,27 +18,20 @@ pub struct FilterOutcome { pub safety_labels: Option, } -pub struct FilterSummary { - pub tweet_count: usize, - pub drop_count: usize, - pub unresolved_author_count: usize, -} - pub struct FilterResponse { pub outcomes: Vec, - pub summary: FilterSummary, } pub struct FilterTweets { hydration_pipeline: HydrationPipeline, - policies: Policies, + rule_engine: RuleEngine, } impl FilterTweets { - pub(crate) fn new(hydration_pipeline: HydrationPipeline, policies: Policies) -> Self { + pub(crate) fn new(hydration_pipeline: HydrationPipeline, rule_engine: RuleEngine) -> Self { Self { hydration_pipeline, - policies, + rule_engine, } } @@ -58,20 +50,12 @@ impl FilterTweets { candidates: hydrated_candidates, safety_labels, } = hydration; - let unresolved_author_count = request.candidates.len() - hydrated_candidates.len(); - if unresolved_author_count > 0 { - info!( - count = unresolved_author_count, - "Tweets with unresolved author_id" - ); - } - let evaluated: HashMap = hydrated_candidates .iter() .map(|candidate| { ( TweetId(candidate.tweet_id), - self.policies + self.rule_engine .evaluate(request.safety_level, &viewer_features, candidate), ) }) @@ -85,18 +69,12 @@ impl FilterTweets { .get(&candidate.tweet_id) .cloned() .unwrap_or_else(Verdict::unresolved_author); - debug!( - viewer_id = request.viewer_id, - tweet_id = candidate.tweet_id.0, - action = ?verdict.action, - rule = ?verdict.decided_by, - safety_level = ?request.safety_level, - "VF verdict" - ); FilterOutcome { tweet_id: candidate.tweet_id, verdict, - safety_labels: safety_labels.get(&candidate.tweet_id).cloned(), + safety_labels: safety_labels + .get(&candidate.tweet_id) + .map(|labels| vf_pb::SafetyLabelMap::clone(labels)), } }) .collect(); @@ -106,24 +84,14 @@ impl FilterTweets { outcomes.iter().map(|outcome| &outcome.verdict), ); - FilterResponse { - summary: FilterSummary { - tweet_count: outcomes.len(), - drop_count: outcomes - .iter() - .filter(|outcome| matches!(outcome.verdict.action, VfAction::Drop(_))) - .count(), - unresolved_author_count, - }, - outcomes, - } + FilterResponse { outcomes } } } #[cfg(test)] -mod tests { +pub(crate) mod test_support { use super::*; - use crate::clients::socialgraph_client::MockSocialgraphClient; + use crate::clients::socialgraph_client::FakeSocialgraphClient; use crate::safety_label_source::lookup::{ManhattanLookup, RemoteSource, TwemcacheLookup}; use crate::safety_label_source::types::{ManhattanOutcome, TwemcacheOutcome}; use crate::safety_label_source::SafetyLabelSource; @@ -169,11 +137,15 @@ mod tests { } } - fn filter_tweets() -> FilterTweets { + pub(crate) fn filter_tweets() -> FilterTweets { + filter_tweets_with_gizmoduck(Arc::new(MockGizmoduckClient::default())) + } + + pub(crate) fn filter_tweets_with_gizmoduck( + gizmoduck: Arc, + ) -> FilterTweets { let tes: Arc = Arc::new(MockTESClient::default()); - let gizmoduck: Arc = - Arc::new(MockGizmoduckClient::default()); - let socialgraph = Arc::new(MockSocialgraphClient::default()); + let socialgraph = Arc::new(FakeSocialgraphClient); let twemcache = Arc::new(FakeTwemcache); let manhattan = Arc::new(FakeManhattan); let labels = Arc::new(SafetyLabelSource::new(Arc::new(RemoteSource::new( @@ -181,17 +153,17 @@ mod tests { )))); FilterTweets::new( - HydrationPipeline::new( - tes, - gizmoduck, - socialgraph, - labels, - crate::hydration::FallbackCacheMode::Disabled, - crate::hydration::FallbackCacheMode::Disabled, - ), - Policies::new(), + HydrationPipeline::new(tes, gizmoduck, socialgraph, labels, None), + RuleEngine::for_tests(), ) } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::test_support::filter_tweets; + use crate::models::VfAction; fn candidate(tweet_id: u64, author_id: Option) -> RawCandidate { RawCandidate { @@ -239,9 +211,6 @@ mod tests { response.outcomes[2].verdict.action, VfAction::Allow )); - assert_eq!(response.summary.tweet_count, 3); - assert_eq!(response.summary.drop_count, 1); - assert_eq!(response.summary.unresolved_author_count, 1); assert!(response .outcomes .iter() diff --git a/visibility-filtering/filter_tweets.rs b/visibility-filtering/filter_tweets.rs index aaa34654..fe653ff6 100644 --- a/visibility-filtering/filter_tweets.rs +++ b/visibility-filtering/filter_tweets.rs @@ -1,19 +1,26 @@ use crate::filter::{FilterOutcome, FilterRequest, FilterTweets}; use crate::models::{RawCandidate, TweetId, VfAction}; +use crate::reference_compare::{ReferenceCompareHarness, TweetVerdict}; use crate::rules::metrics::{self as ft_metrics, RequestMetricsGuard}; use crate::rules::SafetyLevel; -use std::time::Instant; +use std::sync::Arc; use tonic::{Request, Response, Status}; -use tracing::info; use xai_visibility_filtering_proto as vf_pb; pub struct FilterTweetsEndpoint { filter_tweets: FilterTweets, + reference_compare: Option>, } impl FilterTweetsEndpoint { - pub(crate) fn new(filter_tweets: FilterTweets) -> Self { - Self { filter_tweets } + pub(crate) fn new( + filter_tweets: FilterTweets, + reference_compare: Option>, + ) -> Self { + Self { + filter_tweets, + reference_compare, + } } pub async fn handle( @@ -21,9 +28,10 @@ impl FilterTweetsEndpoint { request: Request, ) -> Result, Status> { let request_metrics = RequestMetricsGuard::new(); - let start = Instant::now(); let req = request.into_inner(); ft_metrics::record_batch_size(req.tweets.len()); + let viewer_id = normalize_viewer_id(req.viewer_id); + ft_metrics::record_viewer_state(req.viewer_id, viewer_id); let safety_level = match req.safety_level() { vf_pb::SafetyLevel::FilterAll => SafetyLevel::FilterAll, @@ -32,14 +40,6 @@ impl FilterTweetsEndpoint { SafetyLevel::TimelineHomeRecommendations } }; - info!( - viewer_id = req.viewer_id, - tweet_count = req.tweets.len(), - country_code = ?req.country_code, - safety_level = ?safety_level, - "VF request" - ); - let candidates: Vec = req .tweets .iter() @@ -49,34 +49,53 @@ impl FilterTweetsEndpoint { }) .collect(); + let reference_compare = self.reference_compare.as_ref().and_then(|harness| { + harness.begin_compare( + req.viewer_id, + req.country_code.clone(), + safety_level, + req.tweets.iter().map(|t| t.tweet_id).collect(), + ) + }); + let response = self .filter_tweets .run(FilterRequest { - viewer_id: req.viewer_id, + viewer_id, country_code: req.country_code, safety_level, candidates, }) .await; + + if let Some(verdicts) = reference_compare { + verdicts.send( + response + .outcomes + .iter() + .map(|outcome| TweetVerdict { + tweet_id: outcome.tweet_id.0, + verdict: outcome.verdict.clone(), + }) + .collect(), + ); + } + let results = response .outcomes .into_iter() .map(to_visibility_result) .collect(); - info!( - tweet_count = response.summary.tweet_count, - drop_count = response.summary.drop_count, - not_found_count = response.summary.unresolved_author_count, - latency_ms = start.elapsed().as_millis(), - "VF response" - ); - request_metrics.mark_success(); Ok(Response::new(vf_pb::VisibilityFilterResponse { results })) } } +fn normalize_viewer_id(raw: Option) -> Option { + raw.filter(|&id| id as i64 > 0) +} + fn to_visibility_result(outcome: FilterOutcome) -> vf_pb::TweetVisibilityResult { let (kind, filtered_reason) = match outcome.verdict.action { VfAction::Allow => (vf_pb::action::Kind::Allow(true), None), @@ -113,52 +132,62 @@ mod tests { } } - #[test] - fn allow_maps_to_proto_result() { - let result = to_visibility_result(outcome(7, VfAction::Allow)); - - assert_eq!(result.tweet_id, 7); - assert!(matches!( - result.action.unwrap().kind, - Some(vf_pb::action::Kind::Allow(true)) - )); - assert!(result.filtered_reason.is_none()); + async fn gizmoduck_calls(viewer_id: Option) -> usize { + let gizmoduck = std::sync::Arc::new( + xai_core_entities::gizmoduck_client::MockGizmoduckClient::default(), + ); + let endpoint = FilterTweetsEndpoint::new( + crate::filter::test_support::filter_tweets_with_gizmoduck(gizmoduck.clone()), + None, + ); + let response = endpoint + .handle(Request::new(vf_pb::VisibilityFilterRequest { + safety_level: vf_pb::SafetyLevel::TimelineHome.into(), + tweets: vec![vf_pb::TweetInput { + tweet_id: 2, + author_id: Some(20), + }], + viewer_id, + country_code: None, + })) + .await + .unwrap(); + assert_eq!(response.into_inner().results.len(), 1); + gizmoduck.call_count() + } + + #[tokio::test] + async fn endpoint_treats_zero_viewer_id_as_logged_out() { + let logged_out = gizmoduck_calls(None).await; + assert_eq!(gizmoduck_calls(Some(0)).await, logged_out); + assert_eq!(gizmoduck_calls(Some(42)).await, logged_out + 1); } #[test] - fn drop_maps_to_proto_result_with_labels() { - let labels = vf_pb::SafetyLabelMap::default(); - let mut drop = outcome(8, VfAction::Drop(FilteredReason::ContainNsfwMedia)); - drop.safety_labels = Some(labels.clone()); - let result = to_visibility_result(drop); - - assert_eq!(result.tweet_id, 8); - assert!(matches!( - result.action.unwrap().kind, - Some(vf_pb::action::Kind::Drop(_)) - )); - assert!(matches!( - result.filtered_reason.unwrap().reason, - Some(vf_pb::filtered_reason::Reason::ContainNsfwMedia(true)) - )); - assert_eq!(result.safety_labels, Some(labels)); + fn normalize_viewer_id_cases() { + assert_eq!(normalize_viewer_id(Some(0)), None); + assert_eq!(normalize_viewer_id(Some(u64::MAX)), None); + assert_eq!(normalize_viewer_id(Some(42)), Some(42)); + assert_eq!(normalize_viewer_id(None), None); } #[test] - fn interstitial_maps_to_proto_result() { - let result = to_visibility_result(outcome( - 9, - VfAction::Interstitial(FilteredReason::ContainNsfwMedia), - )); - - assert_eq!(result.tweet_id, 9); - assert!(matches!( - result.action.unwrap().kind, - Some(vf_pb::action::Kind::Interstitial(true)) - )); - assert!(matches!( - result.filtered_reason.unwrap().reason, - Some(vf_pb::filtered_reason::Reason::ContainNsfwMedia(true)) - )); + fn actions_map_to_proto_kinds() { + let cases = [ + (VfAction::Allow, vf_pb::action::Kind::Allow(true)), + ( + VfAction::Drop(FilteredReason::ContainNsfwMedia), + vf_pb::action::Kind::Drop(vf_pb::DropReason {}), + ), + ( + VfAction::Interstitial(FilteredReason::ContainNsfwMedia), + vf_pb::action::Kind::Interstitial(true), + ), + ]; + + for (action, expected) in cases { + let result = to_visibility_result(outcome(1, action)); + assert_eq!(result.action.and_then(|action| action.kind), Some(expected)); + } } } diff --git a/visibility-filtering/get_safety_labels.rs b/visibility-filtering/get_safety_labels.rs index 84054fe4..6985aba5 100644 --- a/visibility-filtering/get_safety_labels.rs +++ b/visibility-filtering/get_safety_labels.rs @@ -58,18 +58,6 @@ impl GetSafetyLabelsEndpoint { metrics::record_lookup_failures(kind, count); } - tracing::info!( - requested_count = outcome.requested_count(), - success_count = outcome.success_count(), - failure_count = failed_count, - manhattan_fetch_failure_count = outcome.failures[FailureKind::ManhattanFetch], - manhattan_decode_failure_count = outcome.failures[FailureKind::ManhattanDecode], - other_failure_count = outcome.failures[FailureKind::Other], - is_partial = outcome.is_partial_failure(), - is_full_failure = outcome.is_full_failure(), - "GetSafetyLabels lookup complete" - ); - Ok(Response::new(vf_pb::GetSafetyLabelsResponse { results: outcome.results, failed_ids: outcome.failed_ids, @@ -87,7 +75,7 @@ pub(crate) struct GetSafetyLabelsOutcome { impl GetSafetyLabelsOutcome { pub(crate) fn try_from_resolved( - resolved: HashMap>, + resolved: HashMap, LookupError>>, requested_count: usize, ) -> Result { let mut results = HashMap::with_capacity(resolved.len()); @@ -97,7 +85,7 @@ impl GetSafetyLabelsOutcome { for (id, lookup_result) in resolved { match lookup_result { Ok(label_map) => { - results.insert(id, label_map); + results.insert(id, Arc::unwrap_or_clone(label_map)); } Err(e) => { failed_ids.push(id); @@ -126,10 +114,6 @@ impl GetSafetyLabelsOutcome { Ok(result) } - pub(crate) fn requested_count(&self) -> usize { - self.requested_count - } - pub(crate) fn success_count(&self) -> usize { self.results.len() } @@ -138,18 +122,6 @@ impl GetSafetyLabelsOutcome { self.failed_ids.len() } - pub(crate) fn has_failures(&self) -> bool { - !self.failed_ids.is_empty() - } - - pub(crate) fn is_partial_failure(&self) -> bool { - self.has_failures() && self.failure_count() < self.requested_count - } - - pub(crate) fn is_full_failure(&self) -> bool { - self.has_failures() && self.failure_count() == self.requested_count - } - fn accounted_count(&self) -> usize { self.results.len() + self.failed_ids.len() } @@ -182,28 +154,24 @@ mod tests { fn try_from_resolved_reports_all_success() { let outcome = GetSafetyLabelsOutcome::try_from_resolved( HashMap::from([ - (1, Ok(labels_with_entry(11))), - (2, Ok(labels_with_entry(22))), + (1, Ok(Arc::new(labels_with_entry(11)))), + (2, Ok(Arc::new(labels_with_entry(22)))), ]), 2, ) .unwrap(); - assert_eq!(outcome.requested_count(), 2); assert_eq!(outcome.success_count(), 2); assert_eq!(outcome.failure_count(), 0); assert!(outcome.failed_ids.is_empty()); assert!(outcome.failures.values().all(|&c| c == 0)); - assert!(!outcome.has_failures()); - assert!(!outcome.is_partial_failure()); - assert!(!outcome.is_full_failure()); } #[test] fn try_from_resolved_reports_partial_failure() { let outcome = GetSafetyLabelsOutcome::try_from_resolved( HashMap::from([ - (1, Ok(labels_with_entry(11))), + (1, Ok(Arc::new(labels_with_entry(11)))), ( 2, Err(LookupError::new(FailureKind::ManhattanDecode, "decode")), @@ -217,16 +185,12 @@ mod tests { ) .unwrap(); - assert_eq!(outcome.requested_count(), 3); assert_eq!(outcome.success_count(), 1); assert_eq!(outcome.results.len(), 1); assert_eq!(outcome.failed_ids, vec![2, 3]); assert_eq!(outcome.failures[FailureKind::ManhattanFetch], 1); assert_eq!(outcome.failures[FailureKind::ManhattanDecode], 1); - assert_eq!(outcome.failures[FailureKind::Other], 0); assert_eq!(outcome.failure_count(), 2); - assert!(outcome.is_partial_failure()); - assert!(!outcome.is_full_failure()); } #[test] @@ -246,22 +210,18 @@ mod tests { ) .unwrap(); - assert_eq!(outcome.requested_count(), 2); assert_eq!(outcome.success_count(), 0); assert!(outcome.results.is_empty()); assert_eq!(outcome.failed_ids, vec![4, 5]); assert_eq!(outcome.failures[FailureKind::ManhattanFetch], 2); assert_eq!(outcome.failures[FailureKind::ManhattanDecode], 0); - assert_eq!(outcome.failures[FailureKind::Other], 0); assert_eq!(outcome.failure_count(), 2); - assert!(!outcome.is_partial_failure()); - assert!(outcome.is_full_failure()); } #[test] fn try_from_resolved_rejects_missing_ids() { let status = GetSafetyLabelsOutcome::try_from_resolved( - HashMap::from([(7, Ok(labels_with_entry(22)))]), + HashMap::from([(7, Ok(Arc::new(labels_with_entry(22))))]), 2, ) .unwrap_err(); diff --git a/visibility-filtering/hydration/batch.rs b/visibility-filtering/hydration/batch.rs index a0712115..4f9dc9ac 100644 --- a/visibility-filtering/hydration/batch.rs +++ b/visibility-filtering/hydration/batch.rs @@ -25,10 +25,6 @@ impl Hydrated { Hydrated::NotFound | Hydrated::Failed(_) => None, } } - - pub(crate) fn is_failed(&self) -> bool { - matches!(self, Hydrated::Failed(_)) - } } impl From, E>> for Hydrated { @@ -106,14 +102,6 @@ impl HydrationBatch { } } - pub(crate) fn len(&self) -> usize { - self.results.len() - } - - pub(crate) fn failed_count(&self) -> usize { - self.results.values().filter(|r| r.is_failed()).count() - } - pub(crate) fn hydrated(&self, key: &K) -> Option<&Hydrated> { self.results.get(key) } @@ -214,7 +202,6 @@ mod tests { ); assert_eq!(batch.get_or_default(&2), 0); assert_eq!(batch.get_or_default(&3), 0); - assert_eq!(batch.failed_count(), 1); } #[test] @@ -226,7 +213,6 @@ mod tests { batch.hydrated(&2), Some(&Hydrated::Failed(HydrationError::MissingResponse)) ); - assert_eq!(batch.failed_count(), 1); } #[test] @@ -237,8 +223,6 @@ mod tests { batch.hydrated(&1), Some(&Hydrated::Failed(HydrationError::Timeout)) ); - assert_eq!(batch.failed_count(), 2); - assert_eq!(batch.len(), 2); } #[test] @@ -267,7 +251,6 @@ mod tests { by_tweet.hydrated(&TweetId(3)), Some(&Hydrated::Failed(_)) )); - assert_eq!(by_tweet.failed_count(), 1); } #[test] @@ -289,7 +272,6 @@ mod tests { assert_eq!(batch.get(&1), Some(&7)); assert_eq!(batch.get(&2), Some(&8)); - assert_eq!(batch.failed_count(), 0); } #[test] @@ -302,6 +284,5 @@ mod tests { batch.hydrated(&2), Some(&Hydrated::Failed(HydrationError::MissingResponse)) ); - assert_eq!(batch.failed_count(), 1); } } diff --git a/visibility-filtering/hydration/fallback_cache.rs b/visibility-filtering/hydration/fallback_cache.rs index 0c00a343..6e1d09e3 100644 --- a/visibility-filtering/hydration/fallback_cache.rs +++ b/visibility-filtering/hydration/fallback_cache.rs @@ -12,23 +12,6 @@ use crate::hydration::metrics::{record_fallback_cache_entries, record_fallback_c const CACHE_SHARDS: usize = 64; const OCCUPANCY_SAMPLE_INTERVAL: u64 = 1024; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum FallbackCacheMode { - Disabled, - Shadow, - ServeStale, -} - -impl FallbackCacheMode { - fn enabled(self) -> bool { - self != Self::Disabled - } - - fn serves_stale(self) -> bool { - self == Self::ServeStale - } -} - #[derive(Clone)] struct CacheEntry { generation: u64, @@ -39,8 +22,7 @@ type CacheShards = Vec>>>; pub(crate) struct FallbackCache { facet: &'static str, - mode: FallbackCacheMode, - shards: Option>, + shards: CacheShards, next_generation: AtomicU64, } @@ -49,24 +31,22 @@ where K: Eq + Hash + Clone, V: Clone, { - pub(crate) fn new(facet: &'static str, capacity: usize, mode: FallbackCacheMode) -> Self { - let shards = mode.enabled().then(|| { - let shard_count = CACHE_SHARDS.min(capacity.max(1)); - let shard_capacity = capacity.div_ceil(shard_count); - (0..shard_count) - .map(|_| Mutex::new(Cache::>::new(shard_capacity))) - .collect() - }); + pub(crate) fn new(facet: &'static str, capacity: usize) -> Self { + let shard_count = CACHE_SHARDS.min(capacity.max(1)); + let shard_capacity = capacity.div_ceil(shard_count); + let shards = (0..shard_count) + .map(|_| Mutex::new(Cache::>::new(shard_capacity))) + .collect(); Self { facet, - mode, shards, next_generation: AtomicU64::new(0), } } - pub(crate) fn enabled(&self) -> bool { - self.mode.enabled() + #[cfg(test)] + pub(crate) fn with_test_capacity(facet: &'static str) -> Self { + Self::new(facet, 8) } pub(crate) fn begin_request(&self) -> u64 { @@ -78,14 +58,9 @@ where generation: u64, batch: HydrationBatch, ) -> HydrationBatch { - if !self.mode.enabled() { - return batch; - } - let mut fresh = 0; let mut stale = 0; let mut stale_not_found = 0; - let mut shadow_hit = 0; let mut not_found = 0; let mut unavailable = 0; let resolved = batch @@ -106,18 +81,14 @@ where Hydrated::Failed(error) => match self.cached_entry(&key) { Some(CacheEntry { value: Some(value), .. - }) if self.mode.serves_stale() => { + }) => { stale += 1; Hydrated::Found(value) } - Some(CacheEntry { value: None, .. }) if self.mode.serves_stale() => { + Some(CacheEntry { value: None, .. }) => { stale_not_found += 1; Hydrated::NotFound } - Some(_) => { - shadow_hit += 1; - Hydrated::Failed(error) - } None => { unavailable += 1; Hydrated::Failed(error) @@ -133,7 +104,6 @@ where fresh, stale, stale_not_found, - shadow_hit, not_found, unavailable, ); @@ -144,7 +114,7 @@ where } fn entry_count(&self) -> usize { - self.shards() + self.shards .iter() .map(|shard| { shard @@ -156,7 +126,7 @@ where } fn write_entry(&self, generation: u64, key: &K, value: Option) { - let shard = cache_shard(self.shards(), key); + let shard = cache_shard(&self.shards, key); let mut cache = shard .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); @@ -170,20 +140,18 @@ where } fn cached_entry(&self, key: &K) -> Option> { - cache_shard(self.shards(), key) + cache_shard(&self.shards, key) .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) .get(key) .cloned() } - - fn shards(&self) -> &[Mutex>>] { - self.shards - .as_deref() - .expect("fallback cache shards unavailable while enabled") - } } +#[expect( + clippy::indexing_slicing, + reason = "index is modulo the non-empty shard list" +)] fn cache_shard<'a, K, V>(shards: &'a [Mutex>], key: &K) -> &'a Mutex> where K: Eq + Hash, @@ -199,8 +167,8 @@ mod tests { use super::*; use crate::hydration::batch::HydrationError; - fn cache(mode: FallbackCacheMode) -> FallbackCache { - FallbackCache::new("test", 8, mode) + fn cache() -> FallbackCache { + FallbackCache::with_test_capacity("test") } fn batch( @@ -215,7 +183,7 @@ mod tests { #[test] fn recovers_only_resident_failed_keys() { - let cache = cache(FallbackCacheMode::ServeStale); + let cache = cache(); cache.resolve_hydration_batch( cache.begin_request(), batch([(1, Hydrated::Found("cached".to_string()))]), @@ -237,7 +205,7 @@ mod tests { #[test] fn authoritative_not_found_invalidates_stale_value() { - let cache = cache(FallbackCacheMode::ServeStale); + let cache = cache(); cache.resolve_hydration_batch( cache.begin_request(), batch([(1, Hydrated::Found("cached".to_string()))]), @@ -249,39 +217,9 @@ mod tests { assert!(matches!(failed.hydrated(&1), Some(Hydrated::NotFound))); } - #[test] - fn shadow_mode_does_not_serve_stale() { - let cache = cache(FallbackCacheMode::Shadow); - cache.resolve_hydration_batch( - cache.begin_request(), - batch([(1, Hydrated::Found("cached".to_string()))]), - ); - - let failed = cache.resolve_hydration_batch(cache.begin_request(), batch([(1, failed())])); - - assert!(matches!(failed.hydrated(&1), Some(Hydrated::Failed(_)))); - } - - #[test] - fn shadow_mode_does_not_serve_cached_not_found() { - let cache = cache(FallbackCacheMode::Shadow); - cache.resolve_hydration_batch(cache.begin_request(), batch([(1, Hydrated::NotFound)])); - - let failed = cache.resolve_hydration_batch(cache.begin_request(), batch([(1, failed())])); - - assert!(matches!(failed.hydrated(&1), Some(Hydrated::Failed(_)))); - } - - #[test] - fn disabled_mode_does_not_allocate_cache_shards() { - let cache = cache(FallbackCacheMode::Disabled); - - assert!(cache.shards.is_none()); - } - #[test] fn late_older_value_does_not_resurrect_newer_not_found() { - let cache = cache(FallbackCacheMode::ServeStale); + let cache = cache(); let older = cache.begin_request(); let newer = cache.begin_request(); @@ -294,7 +232,7 @@ mod tests { #[test] fn late_older_not_found_does_not_suppress_newer_value() { - let cache = cache(FallbackCacheMode::ServeStale); + let cache = cache(); let older = cache.begin_request(); let newer = cache.begin_request(); diff --git a/visibility-filtering/hydration/gizmoduck_hydrator.rs b/visibility-filtering/hydration/gizmoduck_hydrator.rs index 175dc406..2f07001f 100644 --- a/visibility-filtering/hydration/gizmoduck_hydrator.rs +++ b/visibility-filtering/hydration/gizmoduck_hydrator.rs @@ -1,6 +1,6 @@ use crate::clients::gizmoduck_client::GizmoduckLookup; use crate::hydration::batch::{AuthorHydrationBatch, HydrationBatch, TweetHydrationBatch}; -use crate::hydration::fallback_cache::{FallbackCache, FallbackCacheMode}; +use crate::hydration::fallback_cache::FallbackCache; use crate::hydration::metrics::{record_batch_size, timed_results}; use crate::hydration::{keyed_by_author, tweets_per_author}; use crate::models::{AuthorFeatures, AuthorId, TweetCandidateInput, UserLabelSet}; @@ -16,26 +16,33 @@ const CACHE_CAPACITY: usize = 1_000_000; pub struct GizmoduckAuthorHydrator { pub gizmoduck_client: GizmoduckLookup, - fallback_cache: FallbackCache, + fallback_cache: Option>, } impl GizmoduckAuthorHydrator { - pub(crate) fn new(gizmoduck_client: GizmoduckLookup, cache_mode: FallbackCacheMode) -> Self { + pub(crate) fn new( + gizmoduck_client: GizmoduckLookup, + fallback_cache: Option>, + ) -> Self { Self { gizmoduck_client, - fallback_cache: FallbackCache::new("author", CACHE_CAPACITY, cache_mode), + fallback_cache, } } + pub(crate) fn fallback_cache() -> FallbackCache { + FallbackCache::new("author", CACHE_CAPACITY) + } + pub(crate) async fn hydrate( &self, candidates: &[TweetCandidateInput], safety_level: SafetyLevel, ) -> TweetHydrationBatch { - let generation = self + let cache_request = self .fallback_cache - .enabled() - .then(|| self.fallback_cache.begin_request()); + .as_ref() + .map(|cache| (cache, cache.begin_request())); let candidate_count_by_key = tweets_per_author(candidates); let author_ids: Vec = candidate_count_by_key.keys().map(|a| a.get()).collect(); @@ -61,9 +68,8 @@ impl GizmoduckAuthorHydrator { }; let author_features = user_results.map(author_features); - let author_features = if let Some(generation) = generation { - self.fallback_cache - .resolve_hydration_batch(generation, author_features) + let author_features = if let Some((cache, generation)) = cache_request { + cache.resolve_hydration_batch(generation, author_features) } else { author_features }; @@ -205,7 +211,7 @@ mod tests { let client = Arc::new(MockGizmoduckClient::default()); let hydrator = GizmoduckAuthorHydrator::new( GizmoduckLookup::new(client.clone()), - FallbackCacheMode::ServeStale, + Some(FallbackCache::with_test_capacity("author")), ); let candidates = vec![candidate(1, 10), candidate(2, 10)]; @@ -214,8 +220,6 @@ mod tests { .await; assert_eq!(client.call_count(), 1); - assert_eq!(features.len(), 2); - assert_eq!(features.failed_count(), 0); for tweet_id in [TweetId(1), TweetId(2)] { assert!(matches!( features.hydrated(&tweet_id), @@ -233,36 +237,23 @@ mod tests { } #[tokio::test] - async fn stale_recovery_respects_cache_mode() { + async fn stale_recovery_uses_resident_value() { let candidates = vec![candidate(1, 10)]; - for (mode, serves_stale) in [ - (FallbackCacheMode::ServeStale, true), - (FallbackCacheMode::Shadow, false), - (FallbackCacheMode::Disabled, false), - ] { - let hydrator = GizmoduckAuthorHydrator::new( - GizmoduckLookup::new(Arc::new(FailingAfterFirstClient { - calls: AtomicUsize::new(0), - })), - mode, - ); - let first = hydrator - .hydrate(&candidates, SafetyLevel::TimelineHome) - .await; - assert!(first.get_or_default(&TweetId(1)).is_suspended); + let hydrator = GizmoduckAuthorHydrator::new( + GizmoduckLookup::new(Arc::new(FailingAfterFirstClient { + calls: AtomicUsize::new(0), + })), + Some(FallbackCache::with_test_capacity("author")), + ); + let first = hydrator + .hydrate(&candidates, SafetyLevel::TimelineHome) + .await; + assert!(first.get_or_default(&TweetId(1)).is_suspended); - let second = hydrator - .hydrate(&candidates, SafetyLevel::TimelineHome) - .await; - if serves_stale { - assert!(second.get_or_default(&TweetId(1)).is_suspended); - } else { - assert!(matches!( - second.hydrated(&TweetId(1)), - Some(Hydrated::Failed(_)) - )); - } - } + let second = hydrator + .hydrate(&candidates, SafetyLevel::TimelineHome) + .await; + assert!(second.get_or_default(&TweetId(1)).is_suspended); } #[test] @@ -277,8 +268,11 @@ mod tests { .map(author_features) .project([(TweetId(1), a10), (TweetId(2), a20)]); - assert_eq!(by_tweet.failed_count(), 2); for tweet_id in [TweetId(1), TweetId(2)] { + assert!(matches!( + by_tweet.hydrated(&tweet_id), + Some(Hydrated::Failed(_)) + )); let feature = by_tweet.get_or_default(&tweet_id); assert!(!feature.is_suspended); assert!(!feature.is_deactivated); diff --git a/visibility-filtering/hydration/metrics.rs b/visibility-filtering/hydration/metrics.rs index 1cf2e3eb..3339673b 100644 --- a/visibility-filtering/hydration/metrics.rs +++ b/visibility-filtering/hydration/metrics.rs @@ -241,7 +241,6 @@ pub(crate) fn record_fallback_cache_keys( fresh: usize, stale: usize, stale_not_found: usize, - shadow_hit: usize, not_found: usize, unavailable: usize, ) { @@ -249,7 +248,6 @@ pub(crate) fn record_fallback_cache_keys( ("fresh", fresh), ("stale", stale), ("stale_not_found", stale_not_found), - ("shadow_hit", shadow_hit), ("not_found", not_found), ("unavailable", unavailable), ] { @@ -532,7 +530,6 @@ mod tests { ) .await; - assert_eq!(returned.failed_count(), 2); assert!(matches!( returned.hydrated(&1), Some(Hydrated::Failed(HydrationError::Timeout)) diff --git a/visibility-filtering/hydration/mod.rs b/visibility-filtering/hydration/mod.rs index ff66570c..0ea33a16 100644 --- a/visibility-filtering/hydration/mod.rs +++ b/visibility-filtering/hydration/mod.rs @@ -19,13 +19,12 @@ use crate::rules::SafetyLevel; use crate::safety_label_source::SafetyLabelSource; use batch::TweetHydrationBatch; use exclusive_content_hydrator::ExclusiveContentHydrator; -pub(crate) use fallback_cache::FallbackCacheMode; +use fallback_cache::FallbackCache; use gizmoduck_hydrator::GizmoduckAuthorHydrator; use safety_label_hydrator::{SafetyLabelHydration, SafetyLabelHydrator}; use socialgraph_hydrator::SocialgraphHydrator; use std::collections::HashMap; use std::sync::Arc; -use std::time::Instant; use tes_hydrator::TesHydrator; use viewer_hydrator::ViewerHydrator; use xai_core_entities::gizmoduck_client::GizmoduckClient; @@ -126,7 +125,7 @@ pub(crate) struct HydrationPipeline { pub(crate) struct HydrationOutput { pub(crate) viewer_features: ViewerFeatures, pub(crate) candidates: Vec, - pub(crate) safety_labels: HashMap, + pub(crate) safety_labels: HashMap>, } impl HydrationPipeline { @@ -135,17 +134,18 @@ impl HydrationPipeline { gizmoduck_client: Arc, socialgraph_client: Arc, safety_label_source: Arc, - fallback_cache_mode: FallbackCacheMode, - media_fallback_cache_mode: FallbackCacheMode, + fallback_cache: Option>, ) -> Self { Self { viewer_hydrator: ViewerHydrator { gizmoduck_client: gizmoduck_client.clone(), }, - tes_hydrator: TesHydrator::new(tes_client.clone(), media_fallback_cache_mode), + tes_hydrator: TesHydrator { + tes_client: tes_client.clone(), + }, gizmoduck_author_hydrator: GizmoduckAuthorHydrator::new( GizmoduckLookup::new(gizmoduck_client), - fallback_cache_mode, + fallback_cache, ), socialgraph_hydrator: SocialgraphHydrator { sg_client: socialgraph_client.clone(), @@ -172,100 +172,27 @@ impl HydrationPipeline { .viewer_hydrator .hydrate(viewer_id, country_code, safety_level); let candidate_hydration = async { - let start = Instant::now(); let tweet_ids: Vec = raw_candidates.iter().map(|c| c.tweet_id).collect(); let independent_group = async { tokio::join!( - async { - let hydrator_start = Instant::now(); - let result = self - .safety_label_hydrator - .hydrate(&tweet_ids, safety_level) - .await; - tracing::info!( - hydrator = "SafetyLabelHydrator", - result_count = result.label_types.len(), - latency_ms = hydrator_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); - result - }, - async { - let hydrator_start = Instant::now(); - let result = self - .tes_hydrator - .hydrate_tweets(&tweet_ids, safety_level) - .await; - tracing::info!( - hydrator = "TesHydrator.tweets", - result_count = result.media.len(), - failed_entries = result.failed_entries(), - latency_ms = hydrator_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); - result - }, - async { - let hydrator_start = Instant::now(); - let result = self - .exclusive_content_hydrator - .hydrate(&tweet_ids, viewer, safety_level) - .await; - tracing::info!( - hydrator = "ExclusiveContentHydrator", - result_count = result.len(), - latency_ms = hydrator_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); - result - }, + self.safety_label_hydrator.hydrate(&tweet_ids, safety_level), + self.tes_hydrator.hydrate_tweets(&tweet_ids, safety_level), + self.exclusive_content_hydrator + .hydrate(&tweet_ids, viewer, safety_level), ) }; let author_hop = async { - let pure_core_start = Instant::now(); let core_datas = self .tes_hydrator .fetch_pure_core(&tweet_ids, safety_level) .await; - tracing::info!( - hydrator = "TesHydrator.pure_core", - result_count = core_datas.len(), - latency_ms = pure_core_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); let candidates = resolve_candidates(raw_candidates, &core_datas); let (author_features, relationships) = tokio::join!( - async { - let hydrator_start = Instant::now(); - let result = self - .gizmoduck_author_hydrator - .hydrate(&candidates, safety_level) - .await; - tracing::info!( - hydrator = "GizmoduckAuthorHydrator", - result_count = result.len(), - failed_count = result.failed_count(), - latency_ms = hydrator_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); - result - }, - async { - let hydrator_start = Instant::now(); - let result = self - .socialgraph_hydrator - .hydrate(&candidates, viewer, safety_level) - .await; - tracing::info!( - hydrator = "SocialgraphHydrator", - result_count = result.len(), - failed_count = result.failed_count(), - latency_ms = hydrator_start.elapsed().as_millis() as u64, - "Hydrator completed" - ); - result - }, + self.gizmoduck_author_hydrator + .hydrate(&candidates, safety_level), + self.socialgraph_hydrator + .hydrate(&candidates, viewer, safety_level), ); (core_datas, candidates, author_features, relationships) }; @@ -295,11 +222,6 @@ impl HydrationPipeline { }; let hydrated_candidates = features.assemble(&candidates); - tracing::info!( - candidate_count = hydrated_candidates.len(), - total_latency_ms = start.elapsed().as_millis() as u64, - "All hydrators completed" - ); (hydrated_candidates, label_response) }; @@ -342,7 +264,7 @@ mod tests { TweetId(2), TweetFeatures { core: crate::models::CoreFeature { - text: "two".to_string(), + source_tweet_id: Some(2), ..Default::default() }, ..Default::default() @@ -382,7 +304,7 @@ mod tests { let c = &assembled[0]; assert_eq!(c.tweet_id, 2); assert_eq!(c.author_id, 200); - assert_eq!(c.tweet_features.core.text, "two"); + assert_eq!(c.tweet_features.core.source_tweet_id, Some(2)); assert!(c.author_features.is_suspended); assert!(c.relationship.viewer_follows_author); } diff --git a/visibility-filtering/hydration/safety_label_hydrator.rs b/visibility-filtering/hydration/safety_label_hydrator.rs index e7ec2644..5756bc66 100644 --- a/visibility-filtering/hydration/safety_label_hydrator.rs +++ b/visibility-filtering/hydration/safety_label_hydrator.rs @@ -15,7 +15,7 @@ pub struct SafetyLabelHydrator { pub struct SafetyLabelHydration { pub label_types: HashMap, - pub label_response: HashMap, + pub label_response: HashMap>, } impl SafetyLabelHydrator { @@ -48,7 +48,7 @@ impl SafetyLabelHydrator { Some(label_map) => { label_types .insert(*tweet_id, SafetyLabelMap::from_proto_label_types(label_map)); - label_response.insert(*tweet_id, label_map.clone()); + label_response.insert(*tweet_id, Arc::clone(label_map)); } None => { label_types.insert(*tweet_id, SafetyLabelMap::default()); diff --git a/visibility-filtering/hydration/socialgraph_hydrator.rs b/visibility-filtering/hydration/socialgraph_hydrator.rs index 7d833965..eddb7805 100644 --- a/visibility-filtering/hydration/socialgraph_hydrator.rs +++ b/visibility-filtering/hydration/socialgraph_hydrator.rs @@ -58,7 +58,7 @@ impl SocialgraphHydrator { #[cfg(test)] mod tests { use super::*; - use crate::clients::socialgraph_client::MockSocialgraphClient; + use crate::clients::socialgraph_client::FakeSocialgraphClient; use crate::hydration::batch::Hydrated; use crate::models::{resolve_candidate, RawCandidate, TweetId}; use std::collections::HashMap; @@ -97,7 +97,7 @@ mod tests { } #[tokio::test] - async fn rpc_error_empty_response_marks_every_key_failed() { + async fn rpc_empty_response_marks_every_key_failed() { let hydrator = SocialgraphHydrator { sg_client: Arc::new(EmptyResponseSocialgraphClient), }; @@ -107,13 +107,18 @@ mod tests { .hydrate(&candidates, Viewer::LoggedIn(99), SafetyLevel::TimelineHome) .await; - assert_eq!(relationships.failed_count(), 2); + for tweet_id in [TweetId(1), TweetId(2)] { + assert!(matches!( + relationships.hydrated(&tweet_id), + Some(Hydrated::Failed(_)) + )); + } } #[tokio::test] async fn logged_out_viewer_is_found_default_even_for_duplicate_tweets() { let hydrator = SocialgraphHydrator { - sg_client: Arc::new(MockSocialgraphClient::default()), + sg_client: Arc::new(FakeSocialgraphClient), }; let candidates = vec![candidate(1, 10), candidate(1, 10), candidate(2, 20)]; @@ -121,7 +126,6 @@ mod tests { .hydrate(&candidates, Viewer::LoggedOut, SafetyLevel::TimelineHome) .await; - assert_eq!(relationships.failed_count(), 0); for tweet_id in [TweetId(1), TweetId(2)] { assert!(matches!( relationships.hydrated(&tweet_id), diff --git a/visibility-filtering/hydration/tes_hydrator.rs b/visibility-filtering/hydration/tes_hydrator.rs index 220e2171..7b18980a 100644 --- a/visibility-filtering/hydration/tes_hydrator.rs +++ b/visibility-filtering/hydration/tes_hydrator.rs @@ -1,9 +1,7 @@ use crate::hydration::batch::TweetHydrationBatch; -use crate::hydration::fallback_cache::{FallbackCache, FallbackCacheMode}; use crate::hydration::metrics::{record_batch_size, timed_keyed_rpc, timed_results}; use crate::models::{ - CoreFeature, MediaFeature, NsfwFeature, TakedownFeature, TweetCandidateInput, TweetFeatures, - TweetId, + CoreFeature, MediaFeature, NsfwFeature, TweetCandidateInput, TweetFeatures, TweetId, }; use crate::rules::SafetyLevel; use std::collections::HashMap; @@ -14,11 +12,9 @@ use xai_core_entities::tweet_entity_service_client::TESClient; const CLIENT_TIMEOUT: Duration = Duration::from_millis(150); const CLIENT: &str = "tes"; -const CACHE_CAPACITY: usize = 1_000_000; pub struct TesHydrator { pub tes_client: Arc, - fallback_cache: FallbackCache, } #[derive(Default)] @@ -27,36 +23,12 @@ pub(crate) struct TweetHydration { pub(crate) community: TweetHydrationBatch, pub(crate) nsfw_user: TweetHydrationBatch, pub(crate) nsfw_admin: TweetHydrationBatch, - pub(crate) has_takedown: TweetHydrationBatch, pub(crate) takedown_reasons: TweetHydrationBatch>, pub(crate) edit_control: TweetHydrationBatch, pub(crate) media: TweetHydrationBatch, } -impl TweetHydration { - pub(crate) fn failed_entries(&self) -> usize { - self.nullcast.failed_count() - + self.community.failed_count() - + self.nsfw_user.failed_count() - + self.nsfw_admin.failed_count() - + self.has_takedown.failed_count() - + self.takedown_reasons.failed_count() - + self.edit_control.failed_count() - + self.media.failed_count() - } -} - impl TesHydrator { - pub(crate) fn new( - tes_client: Arc, - cache_mode: FallbackCacheMode, - ) -> Self { - Self { - tes_client, - fallback_cache: FallbackCache::new("media", CACHE_CAPACITY, cache_mode), - } - } - pub async fn fetch_pure_core( &self, tweet_ids: &[TweetId], @@ -88,10 +60,6 @@ impl TesHydrator { tweet_ids: &[TweetId], safety_level: SafetyLevel, ) -> TweetHydration { - let generation = self - .fallback_cache - .enabled() - .then(|| self.fallback_cache.begin_request()); let candidate_count_by_key = candidates_per_tweet(tweet_ids); let raw_ids: Vec = candidate_count_by_key.keys().copied().collect(); @@ -100,7 +68,6 @@ impl TesHydrator { community, nsfw_user, nsfw_admin, - has_takedown, takedown_reasons, edit_control, media_entities, @@ -137,14 +104,6 @@ impl TesHydrator { CLIENT_TIMEOUT, self.tes_client.get_nsfw_admin(raw_ids.clone()), ), - timed_results( - CLIENT, - "get_has_takedown", - safety_level, - &candidate_count_by_key, - CLIENT_TIMEOUT, - self.tes_client.get_has_takedown(raw_ids.clone()), - ), timed_results( CLIENT, "get_takedown_reasons", @@ -171,23 +130,14 @@ impl TesHydrator { ), ); - let media = media_entities.map_keys(TweetId).map(media_feature); - let media = if let Some(generation) = generation { - self.fallback_cache - .resolve_hydration_batch(generation, media) - } else { - media - }; - TweetHydration { nullcast: nullcast.map_keys(TweetId), community: community.map_keys(TweetId), nsfw_user: nsfw_user.map_keys(TweetId), nsfw_admin: nsfw_admin.map_keys(TweetId), - has_takedown: has_takedown.map_keys(TweetId), takedown_reasons: takedown_reasons.map_keys(TweetId), edit_control: edit_control.map_keys(TweetId), - media, + media: media_entities.map_keys(TweetId).map(media_feature), } } @@ -227,10 +177,7 @@ fn build_tweet_features( let media = tweet_keyed.media.get_or_default(&id); let is_nullcast = tweet_keyed.nullcast.get(&id).copied().unwrap_or(false); let is_community_tweet = tweet_keyed.community.get(&id).is_some(); - let takedown = TakedownFeature { - applied: tweet_keyed.has_takedown.get(&id).copied().unwrap_or(false), - reasons: tweet_keyed.takedown_reasons.get_or_default(&id), - }; + let takedown_reasons = tweet_keyed.takedown_reasons.get_or_default(&id); let nsfw = NsfwFeature { user: tweet_keyed.nsfw_user.get(&id).copied().unwrap_or(false), admin: tweet_keyed.nsfw_admin.get(&id).copied().unwrap_or(false), @@ -243,10 +190,9 @@ fn build_tweet_features( core: CoreFeature { text: core_data.text.clone(), source_tweet_id: core_data.source_tweet_id, - created_at_secs: core_data.created_at_secs, }, media, - takedown, + takedown_reasons, nsfw, is_nullcast, is_community_tweet, @@ -284,10 +230,7 @@ fn media_feature(entities: MediaEntities) -> MediaFeature { #[cfg(test)] mod tests { use super::*; - use crate::hydration::batch::Hydrated; use crate::models::{resolve_candidate, RawCandidate}; - use anyhow::Result; - use std::sync::atomic::{AtomicUsize, Ordering}; use xai_core_entities::entities::{MediaEntity, PureCoreData}; use xai_core_entities::tweet_entity_service_client::MockTESClient; use xai_x_thrift::media_information::{AdditionalMetadata, Restrictions}; @@ -318,10 +261,9 @@ mod tests { } fn hydrator() -> TesHydrator { - TesHydrator::new( - Arc::new(MockTESClient::default()), - FallbackCacheMode::Disabled, - ) + TesHydrator { + tes_client: Arc::new(MockTESClient::default()), + } } #[test] @@ -490,6 +432,27 @@ mod tests { assert!(!feature.has_dmca_media); } + #[test] + fn assemble_hydrates_text_from_core_data() { + let candidates = vec![candidate(10, 100)]; + let core_datas = HashMap::from([( + TweetId(10), + PureCoreData { + author_id: 100, + text: "muted words".to_string(), + ..Default::default() + }, + )]); + + let features = hydrator().assemble_tweet_features( + &candidates, + &core_datas, + &TweetHydration::default(), + ); + + assert_eq!(features[&TweetId(10)].core.text, "muted words"); + } + #[test] fn assemble_defaults_features_when_core_missing() { let candidates = vec![resolve_candidate( @@ -509,258 +472,6 @@ mod tests { let f = &features[&TweetId(10)]; assert!(f.core.text.is_empty()); - assert_eq!(f.core.created_at_secs, None); assert!(!f.media.has_media); } - - struct MediaFailingAfterFirstClient { - inner: MockTESClient, - media_calls: AtomicUsize, - } - - impl MediaFailingAfterFirstClient { - fn with_media(media_entities: HashMap>) -> Self { - Self { - inner: MockTESClient { - media_entities, - ..Default::default() - }, - media_calls: AtomicUsize::new(0), - } - } - } - - #[tonic::async_trait] - impl TESClient for MediaFailingAfterFirstClient { - async fn get_tweet_media_entities( - &self, - tweet_ids: Vec, - ) -> HashMap>> { - if self.media_calls.fetch_add(1, Ordering::SeqCst) == 0 { - self.inner.get_tweet_media_entities(tweet_ids).await - } else { - tweet_ids - .into_iter() - .map(|id| (id, Err(anyhow::anyhow!("tes unavailable")))) - .collect() - } - } - - async fn get_nullcast(&self, tweet_ids: Vec) -> HashMap>> { - self.inner.get_nullcast(tweet_ids).await - } - - async fn get_community(&self, tweet_ids: Vec) -> HashMap>> { - self.inner.get_community(tweet_ids).await - } - - async fn get_nsfw_user(&self, tweet_ids: Vec) -> HashMap>> { - self.inner.get_nsfw_user(tweet_ids).await - } - - async fn get_nsfw_admin(&self, tweet_ids: Vec) -> HashMap>> { - self.inner.get_nsfw_admin(tweet_ids).await - } - - async fn get_has_takedown( - &self, - tweet_ids: Vec, - ) -> HashMap>> { - self.inner.get_has_takedown(tweet_ids).await - } - - async fn get_takedown_reasons( - &self, - tweet_ids: Vec, - ) -> HashMap>>> { - self.inner.get_takedown_reasons(tweet_ids).await - } - - async fn get_edit_control( - &self, - tweet_ids: Vec, - ) -> HashMap>> { - self.inner.get_edit_control(tweet_ids).await - } - - async fn get_tweet_core_datas( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_subscription_author_ids( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_quoted_tweets( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_reaction_context( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_min_video_durations( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_media_count(&self, _tweet_ids: Vec) -> HashMap>> { - unreachable!() - } - - async fn get_takedown_country_codes( - &self, - _tweet_ids: Vec, - ) -> HashMap>>> { - unreachable!() - } - - async fn get_language_code( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_api_counts( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_is_article(&self, _tweet_ids: Vec) -> HashMap>> { - unreachable!() - } - - async fn get_is_premium(&self, _tweet_ids: Vec) -> HashMap>> { - unreachable!() - } - - async fn get_urls( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_exclusive_controls( - &self, - _tweet_ids: Vec, - ) -> HashMap>> - { - unreachable!() - } - - async fn get_grok_post_ids( - &self, - _tweet_ids: Vec, - ) -> HashMap>> { - unreachable!() - } - - async fn get_status_perspectives( - &self, - _tweet_ids: Vec, - _metadata: Option<&tonic::metadata::MetadataMap>, - ) -> HashMap>> { - unreachable!() - } - - async fn get_api_media_entities( - &self, - _tweet_ids: Vec, - _metadata: Option<&tonic::metadata::MetadataMap>, - ) -> HashMap>>> { - unreachable!() - } - - async fn get_escherbird_entity_annotations( - &self, - _tweet_ids: Vec, - ) -> HashMap< - u64, - Result>>, - > { - unreachable!() - } - } - - #[tokio::test] - async fn media_stale_recovery_respects_cache_mode() { - let tweet_ids = vec![TweetId(1)]; - for (mode, serves_stale) in [ - (FallbackCacheMode::ServeStale, true), - (FallbackCacheMode::Shadow, false), - (FallbackCacheMode::Disabled, false), - ] { - let hydrator = TesHydrator::new( - Arc::new(MediaFailingAfterFirstClient::with_media(HashMap::from([( - 1, - Some(vec![dmca_media_entity(true)]), - )]))), - mode, - ); - let first = hydrator - .hydrate_tweets(&tweet_ids, SafetyLevel::TimelineHome) - .await; - assert!(first.media.get_or_default(&TweetId(1)).has_dmca_media); - - let second = hydrator - .hydrate_tweets(&tweet_ids, SafetyLevel::TimelineHome) - .await; - if serves_stale { - let media = second.media.get_or_default(&TweetId(1)); - assert!(media.has_media); - assert!(media.has_dmca_media); - } else { - assert!(matches!( - second.media.hydrated(&TweetId(1)), - Some(Hydrated::Failed(_)) - )); - assert!(!second.media.get_or_default(&TweetId(1)).has_dmca_media); - } - } - } - - #[tokio::test] - async fn no_media_tweets_cache_default_entries_that_serve_stale() { - let tweet_ids = vec![TweetId(1)]; - let hydrator = TesHydrator::new( - Arc::new(MediaFailingAfterFirstClient::with_media(HashMap::from([( - 1, - Some(Vec::new()), - )]))), - FallbackCacheMode::ServeStale, - ); - - let first = hydrator - .hydrate_tweets(&tweet_ids, SafetyLevel::TimelineHome) - .await; - assert!(!first.media.get_or_default(&TweetId(1)).has_media); - - let second = hydrator - .hydrate_tweets(&tweet_ids, SafetyLevel::TimelineHome) - .await; - assert!(matches!( - second.media.hydrated(&TweetId(1)), - Some(Hydrated::Found(_)) - )); - assert!(!second.media.get_or_default(&TweetId(1)).has_media); - assert_eq!(second.media.failed_count(), 0); - } } diff --git a/visibility-filtering/hydration/viewer_hydrator.rs b/visibility-filtering/hydration/viewer_hydrator.rs index 51b9b3bd..313ebccf 100644 --- a/visibility-filtering/hydration/viewer_hydrator.rs +++ b/visibility-filtering/hydration/viewer_hydrator.rs @@ -61,12 +61,12 @@ impl ViewerHydrator { ) } Ok(Err(e)) => { - warn!(error = %e, "Gizmoduck viewer lookup failed; failing open"); - (false, ViewerAge::Unknown, None) + warn!(error = %e, "Gizmoduck viewer lookup failed; failing closed on age"); + (false, ViewerAge::LookupFailed, None) } Err(_) => { - warn!("Gizmoduck viewer lookup timed out; failing open"); - (false, ViewerAge::Unknown, None) + warn!("Gizmoduck viewer lookup timed out; failing closed on age"); + (false, ViewerAge::LookupFailed, None) } } } @@ -175,23 +175,27 @@ mod tests { } #[tokio::test] - async fn rpc_error_fails_open_to_logged_in_defaults() { + async fn rpc_error_fails_closed_on_age() { let viewer = hydrate_with_broken_client(ViewerLookup::Fails).await; assert_eq!(viewer.viewer, Viewer::LoggedIn(123)); assert!(!viewer.allows_sensitive_media); - assert_eq!(viewer.viewer_age, ViewerAge::Unknown); + assert_eq!(viewer.viewer_age, ViewerAge::LookupFailed); + assert!(viewer.viewer_age_lookup_failed()); + assert!(!viewer.viewer_is_underage()); + assert!(!viewer.viewer_has_no_stated_age()); assert_eq!(viewer.account_country_code, None); assert_eq!(viewer.country_code.as_deref(), Some("us")); } #[tokio::test] - async fn rpc_timeout_fails_open_to_logged_in_defaults() { + async fn rpc_timeout_fails_closed_on_age() { let viewer = hydrate_with_broken_client(ViewerLookup::Hangs).await; assert_eq!(viewer.viewer, Viewer::LoggedIn(123)); assert!(!viewer.allows_sensitive_media); - assert_eq!(viewer.viewer_age, ViewerAge::Unknown); + assert_eq!(viewer.viewer_age, ViewerAge::LookupFailed); + assert!(viewer.viewer_age_lookup_failed()); assert_eq!(viewer.account_country_code, None); assert_eq!(viewer.country_code.as_deref(), Some("us")); } diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index e83fc76d..bef575e7 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -1,12 +1,39 @@ -pub mod clients; +#![deny( + clippy::dbg_macro, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + clippy::print_stderr, + clippy::print_stdout, + clippy::todo, + clippy::unimplemented, + clippy::unwrap_used +)] +#![cfg_attr( + test, + allow( + clippy::dbg_macro, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + clippy::print_stderr, + clippy::print_stdout, + clippy::unwrap_used + ) +)] + +pub(crate) mod clients; pub mod config; +pub mod dark_traffic_setup; pub(crate) mod filter; pub(crate) mod filter_tweets; pub(crate) mod get_safety_labels; -pub mod hydration; -pub mod models; -pub mod rules; -pub mod safety_label_source; +pub(crate) mod hydration; +pub(crate) mod models; +pub mod params; +pub(crate) mod reference_compare; +pub(crate) mod rules; +pub(crate) mod safety_label_source; pub mod server; pub(crate) mod server_deps; -pub mod twemcache; +pub(crate) mod twemcache; diff --git a/visibility-filtering/main.rs b/visibility-filtering/main.rs index 0f47c471..19ab936c 100644 --- a/visibility-filtering/main.rs +++ b/visibility-filtering/main.rs @@ -1,6 +1,20 @@ +#![deny( + clippy::dbg_macro, + clippy::expect_used, + clippy::indexing_slicing, + clippy::panic, + clippy::print_stderr, + clippy::print_stdout, + clippy::todo, + clippy::unimplemented, + clippy::unwrap_used +)] + use clap::Parser; +use xai_dark_traffic::RejectDarkTrafficLayer; use xai_grpc_compression::GrpcZstdLayer; use xai_visibility_filtering_proto as vf_pb; +use xai_visibility_filtering_service::dark_traffic_setup; use xai_visibility_filtering_service::server::VFServer; use xai_x_rpc::grpc_client::TlsMode; use xai_x_service_builder::XServiceBuilder; @@ -23,6 +37,7 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); XServiceBuilder::new("visibility-filtering-service") + .with_featureswitches(xai_visibility_filtering_service::config::fs_path()) .grpc_port(args.grpc_port) .metrics_port(args.metrics_port) .datacenter(args.datacenter) @@ -30,6 +45,8 @@ async fn main() -> anyhow::Result<()> { .with_tls(TlsMode::server_mtls_from_env()?) .with_reflection(vf_pb::FILE_DESCRIPTOR_SET) .with_layer(GrpcZstdLayer) + .with_layer(dark_traffic_setup::resolve_layer()) + .with_layer(RejectDarkTrafficLayer::from_env()) .http_routes(xai_profiling::profiling_router()) .run::(()) .await diff --git a/visibility-filtering/models/mod.rs b/visibility-filtering/models/mod.rs index 6fbfc0bf..b00cd656 100644 --- a/visibility-filtering/models/mod.rs +++ b/visibility-filtering/models/mod.rs @@ -8,9 +8,9 @@ pub mod viewer; pub use author::{AuthorFeatures, UserLabelSet}; pub use exclusive_content::ExclusiveContentFeatures; pub use relationship::ViewerAuthorRelationship; -pub use safety_labels::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; -pub use tweet::{CoreFeature, MediaFeature, NsfwFeature, TakedownFeature, TweetFeatures}; -pub use viewer::{Viewer, ViewerAge, ViewerFeatures, ADULT_AGE_YEARS}; +pub use safety_labels::{SafetyLabelMap, SafetyLabelType}; +pub use tweet::{CoreFeature, MediaFeature, NsfwFeature, TweetFeatures}; +pub use viewer::{Viewer, ViewerAge, ViewerFeatures}; use std::collections::HashMap; use xai_core_entities::entities::PureCoreData; diff --git a/visibility-filtering/models/safety_labels.rs b/visibility-filtering/models/safety_labels.rs index cb7e0b7d..5389ea5d 100644 --- a/visibility-filtering/models/safety_labels.rs +++ b/visibility-filtering/models/safety_labels.rs @@ -1,34 +1,29 @@ -pub use xai_x_thrift::tweet_safety_label::{SafetyLabel, SafetyLabelType}; +pub use xai_x_thrift::tweet_safety_label::SafetyLabelType; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use xai_visibility_filtering_proto as vf_pb; #[derive(Clone, Debug, Default)] -pub struct SafetyLabelMap { - pub labels: HashMap, - label_types: HashSet, -} +pub struct SafetyLabelMap(HashSet); impl SafetyLabelMap { - pub fn new(labels: HashMap) -> Self { - let label_types = labels.keys().copied().collect(); - Self { - labels, - label_types, - } + #[cfg(test)] + pub fn new(label_types: HashSet) -> Self { + Self(label_types) } pub fn from_proto_label_types(proto: &vf_pb::SafetyLabelMap) -> Self { - let labels = proto - .labels - .keys() - .map(|label_type| (SafetyLabelType(*label_type), SafetyLabel::default())) - .collect(); - Self::new(labels) + Self( + proto + .labels + .keys() + .map(|label_type| SafetyLabelType(*label_type)) + .collect(), + ) } #[inline] pub fn has_label(&self, label_type: SafetyLabelType) -> bool { - self.label_types.contains(&label_type) + self.0.contains(&label_type) } } diff --git a/visibility-filtering/models/tweet.rs b/visibility-filtering/models/tweet.rs index 46c01db0..37098848 100644 --- a/visibility-filtering/models/tweet.rs +++ b/visibility-filtering/models/tweet.rs @@ -1,8 +1,11 @@ #[derive(Clone, Debug, Default)] pub struct CoreFeature { + #[cfg_attr( + not(test), + expect(dead_code, reason = "retained for the upcoming muted-keyword rule") + )] pub text: String, pub source_tweet_id: Option, - pub created_at_secs: Option, } #[derive(Clone, Debug, Default)] @@ -13,12 +16,6 @@ pub struct MediaFeature { pub geo_deny_list: Vec, } -#[derive(Clone, Debug, Default)] -pub struct TakedownFeature { - pub applied: bool, - pub reasons: Vec, -} - #[derive(Clone, Debug, Default)] pub struct NsfwFeature { pub user: bool, @@ -29,7 +26,7 @@ pub struct NsfwFeature { pub struct TweetFeatures { pub core: CoreFeature, pub media: MediaFeature, - pub takedown: TakedownFeature, + pub takedown_reasons: Vec, pub nsfw: NsfwFeature, pub is_nullcast: bool, pub is_community_tweet: bool, diff --git a/visibility-filtering/models/viewer.rs b/visibility-filtering/models/viewer.rs index 46a9b6b0..114f9d66 100644 --- a/visibility-filtering/models/viewer.rs +++ b/visibility-filtering/models/viewer.rs @@ -20,6 +20,9 @@ pub const ADULT_AGE_YEARS: i32 = 18; pub enum ViewerAge { Known(i32), NotStated, + /// Gizmoduck viewer RPC errored or timed out. Distinct from `Unknown` + /// (no such user) so age gates can fail closed instead of serving NSFW. + LookupFailed, #[default] Unknown, } @@ -41,6 +44,13 @@ impl ViewerFeatures { pub fn viewer_has_no_stated_age(&self) -> bool { matches!(self.viewer, Viewer::LoggedIn(_)) && self.viewer_age == ViewerAge::NotStated } + + /// Viewer age could not be confirmed because gizmoduck failed. Sensitive + /// media must hard-drop globally (same sink as underage), not fail open + /// as `Unknown` and not jurisdiction-scope as `NotStated`. + pub fn viewer_age_lookup_failed(&self) -> bool { + matches!(self.viewer, Viewer::LoggedIn(_)) && self.viewer_age == ViewerAge::LookupFailed + } } impl ViewerFeatures { diff --git a/visibility-filtering/params.rs b/visibility-filtering/params.rs new file mode 100644 index 00000000..15598882 --- /dev/null +++ b/visibility-filtering/params.rs @@ -0,0 +1,211 @@ +use arc_swap::ArcSwap; +use std::collections::BTreeSet; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use xai_feature_switches::{FeatureSwitches, RecipientBuilder, Value}; + +pub const NSFW_GATING_COUNTRIES_KEY: &str = "rust_vf_nsfw_gating_countries"; + +const SCALA_NSFW_GATING_FILE: &str = "country_specific_nsfw_content_gating.yml"; +const SCALA_NSFW_GATING_COUNTRIES_KEY: &str = "country_specific_nsfw_content_gating_countries"; + +const DRIFT_COUNTER: &str = "nsfw_gating_countries_drift"; + +pub fn default_nsfw_gating_countries() -> Vec { + [ + "ar", "au", "br", "ca", "de", "es", "fr", "gb", "id", "it", "kr", "mx", "nl", "ph", "pt", + "th", + ] + .map(str::to_string) + .to_vec() +} + +pub(crate) struct NsfwGatingCountries { + countries: ArcSwap>, +} + +impl NsfwGatingCountries { + pub fn starting_at_default() -> Self { + Self { + countries: ArcSwap::from_pointee(default_nsfw_gating_countries()), + } + } + + pub fn contains(&self, country_code: &str) -> bool { + self.countries.load().iter().any(|c| c == country_code) + } + + #[cfg(test)] + pub fn refresh_from(&self, feature_switches: &FeatureSwitches) { + let (_, resolved) = resolve_with_origin(feature_switches); + self.countries.store(Arc::new(resolved)); + } + + pub fn refresh_and_check_drift(&self, feature_switches: &FeatureSwitches, fs_path: &str) { + let (origin, resolved) = resolve_with_origin(feature_switches); + self.countries.store(Arc::new(resolved.clone())); + check_drift(origin, &resolved, fs_path); + } + + pub fn spawn_refresh( + self: &Arc, + feature_switches: Arc, + fs_path: String, + ) { + let cache = Arc::clone(self); + tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + loop { + interval.tick().await; + cache.refresh_and_check_drift(&feature_switches, &fs_path); + } + }); + } +} + +fn lowercased_codes(values: &[Value]) -> Option> { + values + .iter() + .map(|v| v.as_str().map(str::to_ascii_lowercase)) + .collect() +} + +fn resolve_with_origin(feature_switches: &FeatureSwitches) -> (&'static str, Vec) { + let configured = feature_switches + .match_recipient(&RecipientBuilder::new().build()) + .get_array_no_impression(NSFW_GATING_COUNTRIES_KEY) + .and_then(|values| lowercased_codes(values)); + match configured { + Some(list) => ("config", list), + None => ("default", default_nsfw_gating_countries()), + } +} + +fn check_drift(origin: &str, resolved: &[String], fs_path: &str) { + let default = default_nsfw_gating_countries(); + let matches_default = set_eq(resolved, &default); + let Some(scala_list) = scala_nsfw_gating_countries(fs_path) else { + tracing::debug!( + fs_path, + scala_file = SCALA_NSFW_GATING_FILE, + "scala gating file absent or unparsable; skipping drift check" + ); + return; + }; + let matches_scala = set_eq(resolved, &scala_list); + if matches_scala { + tracing::info!( + key = NSFW_GATING_COUNTRIES_KEY, + origin, + resolved = ?resolved, + scala_list = ?scala_list, + matches_scala, + matches_default, + "nsfw gating countries: drift check" + ); + } else { + tracing::warn!( + key = NSFW_GATING_COUNTRIES_KEY, + origin, + resolved = ?resolved, + scala_list = ?scala_list, + matches_scala, + matches_default, + "nsfw gating countries: drift check" + ); + if let Some(sr) = xai_stats_receiver::global_stats_receiver() { + sr.incr(DRIFT_COUNTER, &[], 1); + } + } +} + +fn scala_nsfw_gating_countries(fs_path: &str) -> Option> { + let path = Path::new(fs_path).parent()?.join(SCALA_NSFW_GATING_FILE); + let features = xai_feature_switches::load_yaml_file(&path).ok()?; + features + .iter() + .find_map(|f| f.parameters.get(SCALA_NSFW_GATING_COUNTRIES_KEY)) + .and_then(|param| param.default_value.as_array()) + .and_then(|values| lowercased_codes(values)) +} + +fn set_eq(a: &[String], b: &[String]) -> bool { + a.iter().collect::>() == b.iter().collect::>() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn engine(yaml: &str) -> FeatureSwitches { + FeatureSwitches::load_string(yaml).unwrap() + } + + #[test] + fn refresh_reads_key_and_fails_open() { + let cache = NsfwGatingCountries::starting_at_default(); + assert!(cache.contains("de")); + assert!(!cache.contains("xx")); + + cache.refresh_from(&engine( + r#" +rust_vf: + parameters: + rust_vf_nsfw_gating_countries: + type: array + default: + - "XX" +"#, + )); + assert!(cache.contains("xx")); + assert!(!cache.contains("de")); + + cache.refresh_from(&engine("other:\n parameters: {}\n")); + assert!(cache.contains("de")); + assert!(!cache.contains("xx")); + } + + #[test] + fn malformed_value_falls_back_whole_not_partial() { + let cache = NsfwGatingCountries::starting_at_default(); + cache.refresh_from(&engine( + r#" +rust_vf: + parameters: + rust_vf_nsfw_gating_countries: + type: array + default: + - "xx" + - 7 +"#, + )); + assert!(!cache.contains("xx")); + assert!(cache.contains("de")); + } + + #[test] + fn scala_list_parses_from_sibling_file_and_tolerates_absence() { + let dir = tempfile::tempdir().unwrap(); + let fs_path = dir.path().join("rust_vf.yml"); + assert_eq!(scala_nsfw_gating_countries(fs_path.to_str().unwrap()), None); + + std::fs::write( + dir.path().join(SCALA_NSFW_GATING_FILE), + r#" +country_specific_nsfw_content_gating: + parameters: + country_specific_nsfw_content_gating_countries: + type: array + default: + - "de" + - "FR" +"#, + ) + .unwrap(); + assert_eq!( + scala_nsfw_gating_countries(fs_path.to_str().unwrap()), + Some(vec!["de".to_string(), "fr".to_string()]) + ); + } +} diff --git a/visibility-filtering/reference_compare.rs b/visibility-filtering/reference_compare.rs new file mode 100644 index 00000000..ba878a47 --- /dev/null +++ b/visibility-filtering/reference_compare.rs @@ -0,0 +1,835 @@ +use crate::models::VfAction; +use crate::rules::{SafetyLevel, Verdict}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use tracing::info; +use xai_stats_receiver::global_stats_receiver; +use xai_twittercontext_proto::TwitterContextViewer; +use xai_visibility_filtering::models::{Action, FilteredReason}; +use xai_visibility_filtering::vf_client::{SafetyLevel as ReferenceSafetyLevel, VfClient}; + +const COMPARED: &str = "vf_reference_compared"; +const EXACT_MATCH: &str = "vf_reference_exact_match"; +const DIFFERED: &str = "vf_reference_differed"; +const ERROR: &str = "vf_reference_error"; +const SKIPPED: &str = "vf_reference_skipped"; +const ENABLED: &str = "vf_reference_enabled"; + +const HARNESS_LINE_MARKER: &str = "vf_reference_compare"; +const SCHEMA_VERSION: u32 = 1; +const LINE_BUDGET_BYTES: usize = 12 * 1024; + +const REFERENCE_TIMEOUT: Duration = Duration::from_millis(1500); + +pub(crate) fn should_build_harness( + flag_enabled: bool, + app_env: Option<&str>, +) -> Result { + match (flag_enabled, app_env) { + (false, _) => Ok(false), + (true, Some("prod")) => Err( + "VF_DUAL_CALL_HARNESS_ENABLED is set but APP_ENV=prod; the reference comparator is staging-only", + ), + (true, _) => Ok(true), + } +} + +fn service_pair(action: &VfAction) -> (&'static str, Option<&FilteredReason>) { + match action { + VfAction::Allow => ("allow", None), + VfAction::Drop(reason) => ("drop", Some(reason)), + VfAction::Interstitial(reason) => ("interstitial", Some(reason)), + } +} + +fn reference_action_label(reason: &Option) -> &'static str { + match reason { + None => "allow", + Some(FilteredReason::SafetyResult(safety_result)) => match safety_result.action { + Action::NotEvaluated => "not_evaluated", + Action::Allow => "allow", + Action::Drop(_) => "drop", + Action::Interstitial => "interstitial", + Action::Downrank => "downrank", + Action::Tombstone => "tombstone", + Action::Avoid => "avoid", + }, + Some(_) => "drop", + } +} + +fn reason_token(reason: &FilteredReason) -> String { + match reason { + FilteredReason::SafetyResult(safety_result) => match &safety_result.reason { + Some(inner) => format!("{inner:?}"), + None => "SafetyResult".to_string(), + }, + FilteredReason::TweetMatchesViewerMutedKeyword(_) => { + "TweetMatchesViewerMutedKeyword".to_string() + } + other => format!("{other:?}"), + } +} + +fn service_verdict_str(verdict: &Verdict) -> String { + let (action, reason) = service_pair(&verdict.action); + let mut out = action.to_string(); + if let Some(reason) = reason { + out.push(':'); + out.push_str(&reason_token(reason)); + } + if let Some(rule) = verdict.decided_by { + out.push('@'); + out.push_str(rule); + } + out +} + +fn reference_verdict_str(reference: &Option) -> String { + match reference { + None => "allow".to_string(), + Some(reason) => format!( + "{}:{}", + reference_action_label(reference), + reason_token(reason) + ), + } +} + +pub(crate) fn is_exact_match(service: &VfAction, reference: &Option) -> bool { + let (service_action, service_reason) = service_pair(service); + service_action == reference_action_label(reference) && service_reason == reference.as_ref() +} + +pub(crate) struct TweetVerdict { + pub tweet_id: u64, + pub verdict: Verdict, +} + +pub(crate) struct VerdictSender(tokio::sync::oneshot::Sender>); + +impl VerdictSender { + pub(crate) fn send(self, verdicts: Vec) { + let _ = self.0.send(verdicts); + } +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub(crate) struct CompareCounts { + pub compared: u64, + pub exact_match: u64, + pub differed: u64, + pub errors: HashMap<&'static str, u64>, +} + +pub(crate) struct CompareContext<'a> { + pub viewer_id: u64, + pub safety_level: SafetyLevel, + pub dc: &'a str, + pub build_sha: &'a str, +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct Diff { + pub tweet_id: u64, + pub service: String, + pub reference: String, +} + +pub(crate) fn compare_batch( + verdicts: &[TweetVerdict], + reference_results: &HashMap>>, +) -> (CompareCounts, Vec) { + let mut counts = CompareCounts::default(); + let mut diffs = Vec::new(); + for TweetVerdict { tweet_id, verdict } in verdicts { + let reference = match reference_results.get(tweet_id) { + None => { + *counts.errors.entry("missing_result").or_default() += 1; + continue; + } + Some(Err(_)) => { + *counts.errors.entry("reference_item").or_default() += 1; + continue; + } + Some(Ok(reason)) => reason, + }; + counts.compared += 1; + if is_exact_match(&verdict.action, reference) { + counts.exact_match += 1; + } else { + counts.differed += 1; + diffs.push(Diff { + tweet_id: *tweet_id, + service: service_verdict_str(verdict), + reference: reference_verdict_str(reference), + }); + } + } + (counts, diffs) +} + +fn batch_id() -> String { + use std::sync::atomic::{AtomicU64, Ordering}; + static SEQ: AtomicU64 = AtomicU64::new(0); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() as u64; + format!("{nanos:x}-{:x}", SEQ.fetch_add(1, Ordering::Relaxed)) +} + +struct Group<'a> { + service: &'a str, + reference: &'a str, + tweet_ids: Vec, +} + +fn group_diffs(diffs: &[Diff]) -> Vec> { + let mut index: HashMap<(&str, &str), usize> = HashMap::new(); + let mut groups: Vec> = Vec::new(); + for diff in diffs { + let at = *index + .entry((diff.service.as_str(), diff.reference.as_str())) + .or_insert_with(|| { + groups.push(Group { + service: &diff.service, + reference: &diff.reference, + tweet_ids: Vec::new(), + }); + groups.len() - 1 + }); + #[expect(clippy::indexing_slicing, reason = "at indexes a group already pushed")] + groups[at].tweet_ids.push(diff.tweet_id); + } + groups +} + +fn group_slices(group: &Group<'_>, budget: usize) -> Vec { + let whole = serde_json::json!([group.service, group.reference, group.tweet_ids]); + if whole.to_string().len() < budget { + return vec![whole]; + } + let fixed = serde_json::json!([group.service, group.reference, []]) + .to_string() + .len() + + 1; + let ids_per_slice = budget.saturating_sub(fixed).div_euclid(21).max(1); + group + .tweet_ids + .chunks(ids_per_slice) + .map(|ids| serde_json::json!([group.service, group.reference, ids])) + .collect() +} + +fn line_json( + context: &CompareContext<'_>, + batch: &str, + chunk: [usize; 2], + diffs: Vec, +) -> serde_json::Value { + serde_json::json!({ + "h": HARNESS_LINE_MARKER, + "v": SCHEMA_VERSION, + "batch": batch, + "chunk": chunk, + "build": context.build_sha, + "dc": context.dc, + "level": context.safety_level.as_str(), + "viewer": context.viewer_id, + "diffs": diffs, + }) +} + +pub(crate) fn chunk_lines( + context: &CompareContext<'_>, + batch: &str, + diffs: &[Diff], +) -> Vec { + if diffs.is_empty() { + return Vec::new(); + } + let header_len = line_json(context, batch, [1, 1], Vec::new()) + .to_string() + .len(); + let budget = LINE_BUDGET_BYTES.saturating_sub(header_len); + let mut pages: Vec> = Vec::new(); + let mut current: Vec = Vec::new(); + let mut used = 0; + for group in group_diffs(diffs) { + for slice in group_slices(&group, budget) { + let cost = slice.to_string().len() + 1; + if used + cost > budget && !current.is_empty() { + pages.push(std::mem::take(&mut current)); + used = 0; + } + used += cost; + current.push(slice); + } + } + pages.push(current); + let total = pages.len(); + pages + .into_iter() + .enumerate() + .map(|(index, page)| line_json(context, batch, [index + 1, total], page)) + .collect() +} + +pub(crate) fn comparable_request( + safety_level: SafetyLevel, + viewer_id: Option, +) -> Result<(ReferenceSafetyLevel, u64), &'static str> { + let level = match safety_level { + SafetyLevel::TimelineHome => ReferenceSafetyLevel::TimelineHome, + SafetyLevel::TimelineHomeRecommendations => { + ReferenceSafetyLevel::TimelineHomeRecommendations + } + SafetyLevel::FilterAll => return Err("level_unmapped"), + }; + match viewer_id { + Some(viewer_id) => Ok((level, viewer_id)), + None => Err("logged_out_viewer"), + } +} + +const BUILD_SHA_LEN: usize = 12; +const VF_IMAGE_ENV: &str = "VF_IMAGE"; + +fn resolve_build_sha(compiled: &str, image: Option<&str>) -> String { + if let Some(sha) = sha_prefix(compiled) { + return sha.to_owned(); + } + if let Some(image) = image + && let Some(tag) = image.rsplit(':').next() + && let Some(sha) = sha_prefix(tag) + { + return sha.to_owned(); + } + let mut fallback = compiled.to_owned(); + fallback.truncate(BUILD_SHA_LEN); + fallback +} + +fn sha_prefix(s: &str) -> Option<&str> { + let n = s.bytes().take_while(u8::is_ascii_hexdigit).count(); + (n >= BUILD_SHA_LEN).then(|| &s[..BUILD_SHA_LEN]) +} + +pub struct ReferenceCompareHarness { + reference: Arc, + dc: String, + build_sha: String, +} + +impl ReferenceCompareHarness { + pub(crate) fn new(reference: Arc, datacenter: &str) -> Self { + let compiled = xai_build_version::current_build_information().git_commit_sha; + let image = std::env::var(VF_IMAGE_ENV).ok(); + let build_sha = resolve_build_sha(&compiled, image.as_deref()); + let harness = Self { + reference, + dc: datacenter.to_string(), + build_sha, + }; + info!( + build_sha = %harness.build_sha, + "reference_compare: harness enabled" + ); + harness.incr(ENABLED, &[]); + harness + } + + pub(crate) fn begin_compare( + self: &Arc, + viewer_id: Option, + country_code: Option, + safety_level: SafetyLevel, + tweet_ids: Vec, + ) -> Option { + let (reference_level, viewer_id) = match comparable_request(safety_level, viewer_id) { + Ok(comparable) => comparable, + Err(reason) => { + self.incr(SKIPPED, &[("reason", reason)]); + return None; + } + }; + if tweet_ids.is_empty() { + return None; + } + let (tx, rx) = tokio::sync::oneshot::channel::>(); + let harness = Arc::clone(self); + tokio::spawn(async move { + let viewer = TwitterContextViewer { + user_id: viewer_id as i64, + request_country_code: country_code.unwrap_or_default(), + ..Default::default() + }; + let reference_fut = + harness + .reference + .get_result(tweet_ids, reference_level, viewer_id, Some(viewer)); + let (reference_outcome, verdicts) = + futures::future::join(tokio::time::timeout(REFERENCE_TIMEOUT, reference_fut), rx) + .await; + let reference_results: HashMap>> = + match reference_outcome { + Ok(results) => results + .into_iter() + .map(|(id, r)| (id, r.map(|t| t.reason))) + .collect(), + Err(_) => { + harness.incr(ERROR, &[("kind", "timeout")]); + return; + } + }; + let Ok(verdicts) = verdicts else { return }; + let context = CompareContext { + viewer_id, + safety_level, + dc: &harness.dc, + build_sha: &harness.build_sha, + }; + let (counts, diffs) = compare_batch(&verdicts, &reference_results); + harness.emit(safety_level, &counts); + if !diffs.is_empty() { + for line in chunk_lines(&context, &batch_id(), &diffs) { + #[expect(clippy::print_stdout, reason = "stdout is the diff sink")] + { + println!("{line}"); + } + } + } + }); + Some(VerdictSender(tx)) + } + + fn emit(&self, safety_level: SafetyLevel, counts: &CompareCounts) { + let level = safety_level.as_str(); + self.incr_nonzero(COMPARED, &[("safety_level", level)], counts.compared); + self.incr_nonzero(EXACT_MATCH, &[("safety_level", level)], counts.exact_match); + self.incr_nonzero(DIFFERED, &[("safety_level", level)], counts.differed); + for (kind, count) in &counts.errors { + self.incr_nonzero(ERROR, &[("kind", kind)], *count); + } + } + + fn incr(&self, metric: &str, labels: &[(&str, &str)]) { + self.incr_nonzero(metric, labels, 1); + } + + fn incr_nonzero(&self, metric: &str, labels: &[(&str, &str)], count: u64) { + if count == 0 { + return; + } + if let Some(sr) = global_stats_receiver() { + let mut stamped: Vec<(&str, &str)> = labels.to_vec(); + stamped.push(("build_sha", &self.build_sha)); + sr.incr(metric, &stamped, count); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use xai_visibility_filtering::models::{ + Action, DropReason, KeywordMatch, SafetyResult as ReferenceSafetyResult, + }; + use xai_visibility_filtering::tweet_safety_label::SafetyLabelFailure; + use xai_visibility_filtering::vf_client::TweetVisibility; + + fn reference_allow() -> Option { + None + } + + fn reference_bare_drop() -> Option { + Some(FilteredReason::AuthorIsSuspended) + } + + fn reference_safety_result(action: Action) -> Option { + Some(FilteredReason::SafetyResult(ReferenceSafetyResult { + reason: None, + action, + })) + } + + fn reference_muted_keyword() -> Option { + Some(FilteredReason::TweetMatchesViewerMutedKeyword( + KeywordMatch { + keyword: "spoilers".to_string(), + }, + )) + } + + fn service_allow() -> VfAction { + VfAction::Allow + } + + fn service_drop() -> VfAction { + VfAction::Drop(FilteredReason::AuthorIsSuspended) + } + + fn service_interstitial() -> VfAction { + VfAction::Interstitial(FilteredReason::ContainNsfwMedia) + } + + #[test] + fn strict_equality_no_normalization() { + assert!(is_exact_match(&service_allow(), &reference_allow())); + assert!(is_exact_match(&service_drop(), &reference_bare_drop())); + assert!(!is_exact_match( + &VfAction::Drop(FilteredReason::AuthorIsUnsafe), + &reference_bare_drop() + )); + assert!(!is_exact_match(&service_interstitial(), &reference_allow())); + assert!(!is_exact_match( + &service_allow(), + &reference_safety_result(Action::Avoid) + )); + assert!(!is_exact_match( + &service_drop(), + &reference_safety_result(Action::Drop(DropReason {})) + )); + assert!(!is_exact_match( + &service_allow(), + &reference_muted_keyword() + )); + assert!(!is_exact_match(&service_drop(), &reference_muted_keyword())); + } + + #[test] + fn comparable_request_maps_home_levels_and_skips_the_rest() { + assert_eq!( + comparable_request(SafetyLevel::TimelineHome, Some(7)), + Ok((ReferenceSafetyLevel::TimelineHome, 7)) + ); + assert_eq!( + comparable_request(SafetyLevel::TimelineHomeRecommendations, Some(7)), + Ok((ReferenceSafetyLevel::TimelineHomeRecommendations, 7)) + ); + assert_eq!( + comparable_request(SafetyLevel::FilterAll, Some(7)), + Err("level_unmapped") + ); + assert_eq!( + comparable_request(SafetyLevel::TimelineHome, None), + Err("logged_out_viewer") + ); + } + + #[test] + fn should_build_harness_requires_flag_and_rejects_prod() { + assert_eq!(should_build_harness(false, Some("prod")), Ok(false)); + assert_eq!(should_build_harness(false, Some("staging")), Ok(false)); + assert_eq!(should_build_harness(false, None), Ok(false)); + assert_eq!(should_build_harness(true, Some("staging")), Ok(true)); + assert_eq!(should_build_harness(true, None), Ok(true)); + assert!(should_build_harness(true, Some("prod")).is_err()); + } + + fn verdict(tweet_id: u64, action: VfAction, decided_by: Option<&'static str>) -> TweetVerdict { + TweetVerdict { + tweet_id, + verdict: Verdict { action, decided_by }, + } + } + + fn context() -> CompareContext<'static> { + CompareContext { + viewer_id: 99, + safety_level: SafetyLevel::TimelineHomeRecommendations, + dc: "atla", + build_sha: "abc123def456", + } + } + + #[test] + fn compare_batch_counts_policy_free_and_collects_differing_pairs_only() { + let verdicts = vec![ + verdict(1, service_allow(), None), + verdict(2, service_drop(), Some("DropSuspendedAuthorRule")), + verdict(3, service_allow(), None), + verdict(4, service_allow(), None), + verdict(5, service_allow(), None), + ]; + let reference_results: HashMap>> = + HashMap::from([ + (1, Ok(reference_allow())), + (2, Ok(reference_allow())), + (3, Ok(reference_bare_drop())), + (4, Err(anyhow::anyhow!("reference error"))), + ]); + + let (counts, diffs) = compare_batch(&verdicts, &reference_results); + + assert_eq!(counts.compared, 3); + assert_eq!(counts.exact_match, 1); + assert_eq!(counts.differed, 2); + assert_eq!( + counts.errors, + HashMap::from([("reference_item", 1), ("missing_result", 1)]) + ); + assert_eq!(diffs.len(), 2, "differing pairs only: {diffs:?}"); + } + + #[test] + fn verdict_grammar_encodes_action_reason_and_rule() { + let cases = [ + (verdict(0, service_allow(), None), "allow"), + ( + verdict(0, service_drop(), Some("drop_suspended_author")), + "drop:AuthorIsSuspended@drop_suspended_author", + ), + (verdict(0, service_drop(), None), "drop:AuthorIsSuspended"), + ( + verdict(0, service_interstitial(), Some("nsfw_media")), + "interstitial:ContainNsfwMedia@nsfw_media", + ), + ]; + for (v, expected) in &cases { + assert_eq!(service_verdict_str(&v.verdict), *expected); + } + + assert_eq!(reference_verdict_str(&reference_allow()), "allow"); + assert_eq!( + reference_verdict_str(&reference_bare_drop()), + "drop:AuthorIsSuspended" + ); + assert_eq!( + reference_verdict_str(&Some(FilteredReason::SafetyResult(ReferenceSafetyResult { + reason: Some( + xai_visibility_filtering::models::SafetyResultReason::NsfwHighPrecision + ), + action: Action::Avoid, + }))), + "avoid:NsfwHighPrecision" + ); + for (action, label) in [ + (Action::NotEvaluated, "not_evaluated"), + (Action::Allow, "allow"), + (Action::Drop(DropReason {}), "drop"), + (Action::Interstitial, "interstitial"), + (Action::Downrank, "downrank"), + (Action::Tombstone, "tombstone"), + (Action::Avoid, "avoid"), + ] { + assert_eq!( + reference_verdict_str(&reference_safety_result(action)), + format!("{label}:SafetyResult") + ); + } + } + + #[test] + fn muted_keyword_payload_never_reaches_the_line() { + let encoded = reference_verdict_str(&reference_muted_keyword()); + assert_eq!(encoded, "drop:TweetMatchesViewerMutedKeyword"); + assert!(!encoded.contains("spoilers"), "viewer content leaked"); + } + + fn diff(tweet_id: u64, service: &str, reference: &str) -> Diff { + Diff { + tweet_id, + service: service.to_string(), + reference: reference.to_string(), + } + } + + const ID: u64 = 1_000_000_000_000_000_000; + + #[test] + fn identical_pairs_group_into_one_diffs_entry() { + let diffs = vec![ + diff(1, "allow", "avoid:SafetyResult"), + diff(2, "drop:ContainNsfwMedia@nsfw_media", "allow"), + diff(3, "allow", "avoid:SafetyResult"), + diff(4, "allow", "avoid:SafetyResult"), + ]; + + let lines = chunk_lines(&context(), "b1", &diffs); + + assert_eq!(lines.len(), 1, "a full request fits one line"); + assert_eq!(lines[0]["chunk"], serde_json::json!([1, 1])); + assert_eq!( + lines[0]["diffs"], + serde_json::json!([ + ["allow", "avoid:SafetyResult", [1, 3, 4]], + ["drop:ContainNsfwMedia@nsfw_media", "allow", [2]], + ]) + ); + + let repeated: Vec = (0..150) + .map(|i| diff(ID + i, "allow", "avoid:SafetyResult")) + .collect(); + let lines = chunk_lines(&context(), "b2", &repeated); + assert_eq!(lines.len(), 1); + assert!(lines[0].to_string().len() <= LINE_BUDGET_BYTES); + assert_eq!(lines[0]["diffs"].as_array().unwrap().len(), 1); + } + + #[test] + fn lines_split_at_the_byte_budget_and_stay_self_contained() { + let diffs: Vec = (0..300) + .map(|i| { + diff( + ID + i, + &format!("interstitial:ContainNsfwMedia@NsfwAuthorInterstitialRule{i}"), + "avoid:SafetyResult", + ) + }) + .collect(); + + let lines = chunk_lines(&context(), "b3", &diffs); + + assert!(lines.len() > 1, "300 distinct pairs exceed one budget"); + let total = lines.len(); + let mut seen = 0; + for (i, line) in lines.iter().enumerate() { + assert!(line.to_string().len() <= LINE_BUDGET_BYTES); + assert_eq!(line["h"], "vf_reference_compare"); + assert_eq!(line["v"], 1); + assert_eq!(line["batch"], "b3"); + assert_eq!(line["chunk"], serde_json::json!([i + 1, total])); + assert_eq!(line["build"], "abc123def456"); + assert_eq!(line["dc"], "atla"); + assert_eq!(line["level"], "timeline_home_recommendations"); + assert_eq!(line["viewer"], 99); + seen += line["diffs"].as_array().unwrap().len(); + } + assert_eq!(seen, 300, "splitting is lossless"); + } + + #[test] + fn oversized_single_group_splits_its_id_list() { + let diffs: Vec = (0..1000) + .map(|i| diff(ID + i, "allow", "avoid:SafetyResult")) + .collect(); + + let lines = chunk_lines(&context(), "b4", &diffs); + + assert!(lines.len() > 1, "1000 ids exceed one budget"); + let ids: Vec = lines + .iter() + .flat_map(|line| line["diffs"].as_array().unwrap().iter()) + .flat_map(|group| group[2].as_array().unwrap().iter()) + .map(|id| id.as_u64().unwrap()) + .collect(); + assert_eq!(ids.len(), 1000, "splitting is lossless"); + assert_eq!(ids[0], ID); + assert_eq!(ids[999], ID + 999); + for line in &lines { + assert!(line.to_string().len() <= LINE_BUDGET_BYTES); + assert_eq!(line["diffs"][0][0], "allow"); + assert_eq!(line["diffs"][0][1], "avoid:SafetyResult"); + } + } + + type RecordedCall = (Vec, ReferenceSafetyLevel, u64, Option); + + struct FakeReference { + calls: std::sync::Mutex>, + } + + #[tonic::async_trait] + impl VfClient for FakeReference { + async fn get_result( + &self, + tweet_ids: Vec, + safety_level: ReferenceSafetyLevel, + for_user_id: u64, + context: Option, + ) -> HashMap> { + let results = tweet_ids + .iter() + .map(|&id| { + ( + id, + Ok(TweetVisibility { + reason: None, + safety_labels: Err(SafetyLabelFailure::LookupFailed), + }), + ) + }) + .collect(); + self.calls.lock().unwrap().push(( + tweet_ids, + safety_level, + for_user_id, + context.map(|c| c.request_country_code), + )); + results + } + } + + fn fake_harness() -> (Arc, Arc) { + let fake = Arc::new(FakeReference { + calls: std::sync::Mutex::new(Vec::new()), + }); + ( + Arc::new(ReferenceCompareHarness::new(fake.clone(), "atla")), + fake, + ) + } + + #[tokio::test] + async fn begin_compare_skips_unmappable_requests_without_calling_reference() { + let (harness, fake) = fake_harness(); + + assert!( + harness + .begin_compare(Some(7), None, SafetyLevel::FilterAll, vec![1]) + .is_none(), + "FilterAll has no reference level" + ); + assert!( + harness + .begin_compare(Some(7), None, SafetyLevel::TimelineHome, vec![]) + .is_none(), + "empty requests have nothing to compare" + ); + tokio::task::yield_now().await; + assert!(fake.calls.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn begin_compare_fetches_reference_concurrently_with_request_context() { + let (harness, fake) = fake_harness(); + + let sender = harness + .begin_compare( + Some(99), + Some("de".to_string()), + SafetyLevel::TimelineHomeRecommendations, + vec![1, 2], + ) + .expect("comparable request"); + sender.send(vec![ + verdict(1, service_allow(), None), + verdict(2, service_allow(), None), + ]); + + let deadline = std::time::Instant::now() + Duration::from_secs(2); + loop { + if !fake.calls.lock().unwrap().is_empty() { + break; + } + assert!( + std::time::Instant::now() < deadline, + "reference never called" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + let calls = fake.calls.lock().unwrap(); + assert_eq!( + *calls, + vec![( + vec![1, 2], + ReferenceSafetyLevel::TimelineHomeRecommendations, + 99, + Some("de".to_string()), + )] + ); + } +} diff --git a/visibility-filtering/rules/author_rules.rs b/visibility-filtering/rules/author_rules.rs new file mode 100644 index 00000000..1ebf83a6 --- /dev/null +++ b/visibility-filtering/rules/author_rules.rs @@ -0,0 +1,330 @@ +use crate::models::VfAction; +use crate::rules::rule_spec::RuleSpec; +use crate::rules::RuleContext; +use xai_visibility_filtering::models::FilteredReason; +use xai_x_thrift::user_labels::LabelValue; + +pub(super) const AUTHOR_STATE_DROPS: &[RuleSpec] = &[ + RuleSpec::Author { + name: "SuspendedAuthorRule", + when: |author| author.is_suspended(), + reason: FilteredReason::AuthorIsSuspended, + exempt_follower: false, + }, + RuleSpec::Author { + name: "DeactivatedAuthorRule", + when: |author| author.is_deactivated(), + reason: FilteredReason::AuthorIsDeactivated, + exempt_follower: false, + }, + RuleSpec::Author { + name: "ErasedAuthorRule", + when: |author| author.is_erased(), + reason: FilteredReason::AuthorAccountIsInactive, + exempt_follower: false, + }, + RuleSpec::Author { + name: "OffboardedAuthorRule", + when: |author| author.is_offboarded(), + reason: FilteredReason::AuthorAccountIsInactive, + exempt_follower: false, + }, + RuleSpec::Author { + name: "ProtectedAuthorDropRule", + when: |author| author.is_protected(), + reason: FilteredReason::AuthorIsProtected, + exempt_follower: true, + }, +]; + +pub(super) const OON_NSFW_AUTHOR_DROPS: &[RuleSpec] = &[ + RuleSpec::Author { + name: "DropNsfwUserAuthorRule", + when: |author| author.is_nsfw_user(), + reason: FilteredReason::ContainNsfwMedia, + exempt_follower: false, + }, + RuleSpec::Author { + name: "DropNsfwAdminAuthorRule", + when: |author| author.is_nsfw_admin(), + reason: FilteredReason::ContainNsfwMedia, + exempt_follower: false, + }, +]; + +pub(super) const OON_USER_LABEL_DROPS: &[RuleSpec] = &[ + RuleSpec::Author { + name: "NsfwHighRecallUserLabelRule", + when: |author| author.has_user_label(LabelValue::NSFW_HIGH_RECALL), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "NsfwHighPrecisionUserLabelRule", + when: |author| author.has_user_label(LabelValue::NSFW_HIGH_PRECISION), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "SpamHighRecallUserLabelRule", + when: |author| author.has_user_label(LabelValue::SPAM_HIGH_RECALL), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "CompromisedUserLabelRule", + when: |author| author.has_user_label(LabelValue::COMPROMISED), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "ReadOnlyUserLabelRule", + when: |author| author.has_user_label(LabelValue::READ_ONLY), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "ImpersonationHighPrecisionUserLabelRule", + when: |author| author.has_user_label(LabelValue::IMPERSONATION_HIGH_PRECISION), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "NsfwAvatarImageRule", + when: |author| author.has_user_label(LabelValue::NSFW_AVATAR_IMAGE), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "NsfwBannerImageRule", + when: |author| author.has_user_label(LabelValue::NSFW_BANNER_IMAGE), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "AbusiveHighRecallRule", + when: |author| author.has_user_label(LabelValue::ABUSIVE_HIGH_RECALL), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: true, + }, + RuleSpec::Author { + name: "NsfwNearPerfectAuthorRule", + when: |author| author.has_user_label(LabelValue::NSFW_NEAR_PERFECT), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: false, + }, + RuleSpec::Author { + name: "DoNotAmplifyNonFollowerRule", + when: |author| author.has_user_label(LabelValue::DO_NOT_AMPLIFY), + reason: FilteredReason::UnspecifiedReason, + exempt_follower: true, + }, +]; + +fn viewer_blocks_author(context: &RuleContext<'_>) -> VfAction { + if context.viewer().is_logged_out() { + return VfAction::Allow; + } + if context.viewer().blocks_author() { + return VfAction::Drop(FilteredReason::AuthorBlockViewer); + } + VfAction::Allow +} + +fn viewer_mutes_author(context: &RuleContext<'_>) -> VfAction { + if context.viewer().is_logged_out() { + return VfAction::Allow; + } + if context.viewer().mutes_author() { + return VfAction::Drop(FilteredReason::ViewerMutesAuthor); + } + VfAction::Allow +} + +fn muted_retweets(context: &RuleContext<'_>) -> VfAction { + if context.viewer().is_logged_out() { + return VfAction::Allow; + } + if context.tweet().is_retweet() && context.viewer().mutes_retweets_from_author() { + return VfAction::Drop(FilteredReason::UnspecifiedReason); + } + VfAction::Allow +} + +pub(super) const SOCIALGRAPH_DROPS: &[RuleSpec] = &[ + RuleSpec::Custom { + name: "ViewerBlocksAuthorRule", + evaluate: viewer_blocks_author, + }, + RuleSpec::Custom { + name: "ViewerMutesAuthorRule", + evaluate: viewer_mutes_author, + }, + RuleSpec::Custom { + name: "MutedRetweetsRule", + evaluate: muted_retweets, + }, +]; + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::{AuthorFeatures, ViewerAuthorRelationship}; + use crate::rules::fixtures::{ + assert_allows, assert_drops, author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID, + }; + + fn author_flag_features(name: &str) -> AuthorFeatures { + let mut features = AuthorFeatures::default(); + match name { + "SuspendedAuthorRule" => features.is_suspended = true, + "DeactivatedAuthorRule" => features.is_deactivated = true, + "ErasedAuthorRule" => features.is_erased = true, + "OffboardedAuthorRule" => features.is_offboarded = true, + "DropNsfwUserAuthorRule" => features.is_nsfw_user = true, + "DropNsfwAdminAuthorRule" => features.is_nsfw_admin = true, + "ProtectedAuthorDropRule" => features.is_protected = true, + _ => panic!("no trigger flags for rule {name}"), + } + features + } + + #[test] + fn author_flag_drop_axis() { + for spec in AUTHOR_STATE_DROPS.iter().chain(OON_NSFW_AUTHOR_DROPS) { + let RuleSpec::Author { + name, + reason, + exempt_follower, + .. + } = spec + else { + panic!("{} is not an author drop row", spec.name()); + }; + let firing = candidate() + .with_author_features(author_flag_features(name)) + .build(); + for v in [viewer(VIEWER_ID), logged_out_viewer()] { + assert_drops(spec, &v, &firing, reason); + } + let followed = candidate() + .with_author_features(author_flag_features(name)) + .followed() + .build(); + if *exempt_follower { + assert_allows(spec, &viewer(VIEWER_ID), &followed); + assert_drops(spec, &logged_out_viewer(), &followed, reason); + } else { + assert_drops(spec, &viewer(VIEWER_ID), &followed, reason); + } + let unflagged = candidate().build(); + assert_allows(spec, &viewer(VIEWER_ID), &unflagged); + assert_allows(spec, &author_viewer(), &firing); + } + } + + fn trigger_user_label(name: &str) -> LabelValue { + match name { + "NsfwHighRecallUserLabelRule" => LabelValue::NSFW_HIGH_RECALL, + "NsfwHighPrecisionUserLabelRule" => LabelValue::NSFW_HIGH_PRECISION, + "SpamHighRecallUserLabelRule" => LabelValue::SPAM_HIGH_RECALL, + "CompromisedUserLabelRule" => LabelValue::COMPROMISED, + "ReadOnlyUserLabelRule" => LabelValue::READ_ONLY, + "ImpersonationHighPrecisionUserLabelRule" => LabelValue::IMPERSONATION_HIGH_PRECISION, + "NsfwAvatarImageRule" => LabelValue::NSFW_AVATAR_IMAGE, + "NsfwBannerImageRule" => LabelValue::NSFW_BANNER_IMAGE, + "AbusiveHighRecallRule" => LabelValue::ABUSIVE_HIGH_RECALL, + "NsfwNearPerfectAuthorRule" => LabelValue::NSFW_NEAR_PERFECT, + "DoNotAmplifyNonFollowerRule" => LabelValue::DO_NOT_AMPLIFY, + _ => panic!("no trigger user label for rule {name}"), + } + } + + #[test] + fn user_label_drop_axis() { + for spec in OON_USER_LABEL_DROPS { + let RuleSpec::Author { + name, + reason, + exempt_follower, + .. + } = spec + else { + panic!("{} is not a user-label drop row", spec.name()); + }; + let firing = candidate() + .with_author_user_label(trigger_user_label(name)) + .build(); + for v in [viewer(VIEWER_ID), logged_out_viewer()] { + assert_drops(spec, &v, &firing, reason); + } + let followed = candidate() + .with_author_user_label(trigger_user_label(name)) + .followed() + .build(); + if *exempt_follower { + assert_allows(spec, &viewer(VIEWER_ID), &followed); + assert_drops(spec, &logged_out_viewer(), &followed, reason); + } else { + assert_drops(spec, &viewer(VIEWER_ID), &followed, reason); + } + let unrelated = candidate() + .with_author_user_label(LabelValue::LOW_QUALITY) + .build(); + assert_allows(spec, &viewer(VIEWER_ID), &unrelated); + assert_allows(spec, &author_viewer(), &firing); + } + } + + fn relationship_trigger(name: &str) -> (ViewerAuthorRelationship, bool, FilteredReason) { + match name { + "ViewerBlocksAuthorRule" => ( + ViewerAuthorRelationship { + viewer_blocks_author: true, + ..Default::default() + }, + false, + FilteredReason::AuthorBlockViewer, + ), + "ViewerMutesAuthorRule" => ( + ViewerAuthorRelationship { + viewer_mutes_author: true, + ..Default::default() + }, + false, + FilteredReason::ViewerMutesAuthor, + ), + "MutedRetweetsRule" => ( + ViewerAuthorRelationship { + viewer_mutes_retweets_from_author: true, + ..Default::default() + }, + true, + FilteredReason::UnspecifiedReason, + ), + _ => panic!("no relationship trigger for rule {name}"), + } + } + + #[test] + fn socialgraph_relationship_axis() { + for spec in SOCIALGRAPH_DROPS { + let RuleSpec::Custom { name, .. } = spec else { + panic!("{} is not a custom socialgraph row", spec.name()); + }; + let (rel, retweet, reason) = relationship_trigger(name); + let mut firing = candidate().with_relationship(rel.clone()); + if retweet { + firing = firing.retweet_of(99); + } + let firing = firing.build(); + assert_drops(spec, &viewer(VIEWER_ID), &firing, &reason); + assert_allows(spec, &logged_out_viewer(), &firing); + assert_allows(spec, &viewer(VIEWER_ID), &candidate().build()); + if *name == "MutedRetweetsRule" { + let non_retweet = candidate().with_relationship(rel).build(); + assert_allows(spec, &viewer(VIEWER_ID), &non_retweet); + } + } + } +} diff --git a/visibility-filtering/rules/context.rs b/visibility-filtering/rules/context.rs new file mode 100644 index 00000000..c371dee9 --- /dev/null +++ b/visibility-filtering/rules/context.rs @@ -0,0 +1,313 @@ +use crate::models::{HydratedTweetCandidate, SafetyLabelType, ViewerFeatures}; +use crate::params::NsfwGatingCountries; +use xai_core_entities::entities::TakedownReason; +use xai_x_thrift::user_labels::LabelValue; + +pub struct RuleContext<'a> { + viewer: &'a ViewerFeatures, + candidate: &'a HydratedTweetCandidate, + nsfw_gating_countries: &'a NsfwGatingCountries, +} + +impl<'a> RuleContext<'a> { + pub(super) fn new( + viewer: &'a ViewerFeatures, + candidate: &'a HydratedTweetCandidate, + nsfw_gating_countries: &'a NsfwGatingCountries, + ) -> Self { + Self { + viewer, + candidate, + nsfw_gating_countries, + } + } + + #[inline] + pub fn viewer(&self) -> ViewerPredicates<'_> { + ViewerPredicates { ctx: self } + } + + #[inline] + pub fn tweet(&self) -> TweetPredicates<'_> { + TweetPredicates { ctx: self } + } + + #[inline] + pub fn author(&self) -> AuthorPredicates<'_> { + AuthorPredicates { ctx: self } + } + + #[inline] + pub fn takedown(&self) -> TakedownPredicates<'_> { + TakedownPredicates { ctx: self } + } + + #[inline] + pub fn nsfw_gating_country(&self, country_code: &str) -> bool { + self.nsfw_gating_countries.contains(country_code) + } +} + +#[derive(Clone, Copy)] +pub struct ViewerPredicates<'a> { + ctx: &'a RuleContext<'a>, +} + +impl ViewerPredicates<'_> { + #[inline] + pub fn is_logged_out(&self) -> bool { + self.ctx.viewer.viewer_is_logged_out() + } + + #[inline] + pub fn is_underage(&self) -> bool { + self.ctx.viewer.viewer_is_underage() + } + + #[inline] + pub fn has_no_stated_age(&self) -> bool { + self.ctx.viewer.viewer_has_no_stated_age() + } + + #[inline] + pub fn age_lookup_failed(&self) -> bool { + self.ctx.viewer.viewer_age_lookup_failed() + } + + #[inline] + pub fn allows_sensitive_media(&self) -> bool { + self.ctx.viewer.allows_sensitive_media + } + + #[inline] + pub fn country(&self) -> Option<&str> { + self.ctx + .viewer + .account_country_code + .as_deref() + .or(self.ctx.viewer.country_code.as_deref()) + } + + #[inline] + pub fn is_author(&self) -> bool { + self.ctx.candidate.is_author_viewer(self.ctx.viewer.viewer) + } + + #[inline] + pub fn follows_author(&self) -> bool { + self.ctx.candidate.viewer_follows_author() + } + + #[inline] + pub fn blocks_author(&self) -> bool { + self.ctx.candidate.relationship.viewer_blocks_author + } + + #[inline] + pub fn mutes_author(&self) -> bool { + self.ctx.candidate.relationship.viewer_mutes_author + } + + #[inline] + pub fn mutes_retweets_from_author(&self) -> bool { + self.ctx + .candidate + .relationship + .viewer_mutes_retweets_from_author + } + + #[inline] + pub fn is_conversation_author(&self) -> bool { + match ( + &self.ctx.candidate.exclusive_content, + self.ctx.viewer.viewer_id(), + ) { + (Some(exclusive), Some(viewer_id)) => viewer_id == exclusive.conversation_author_id, + _ => false, + } + } + + #[inline] + pub fn super_follows_author(&self) -> bool { + self.ctx + .candidate + .exclusive_content + .as_ref() + .is_some_and(|exclusive| exclusive.viewer_super_follows_author) + } +} + +#[derive(Clone, Copy)] +pub struct TweetPredicates<'a> { + ctx: &'a RuleContext<'a>, +} + +impl TweetPredicates<'_> { + #[inline] + pub fn has_safety_label(&self, label: SafetyLabelType) -> bool { + self.ctx.candidate.has_safety_label(label) + } + + #[inline] + pub fn is_retweet(&self) -> bool { + self.ctx.candidate.is_retweet() + } + + #[inline] + pub fn is_stale(&self) -> bool { + self.ctx.candidate.is_stale_tweet() + } + + #[inline] + pub fn is_nullcast(&self) -> bool { + self.ctx.candidate.is_nullcast() + } + + #[inline] + pub fn is_community_tweet(&self) -> bool { + self.ctx.candidate.is_community_tweet() + } + + #[inline] + pub fn has_media(&self) -> bool { + self.ctx.candidate.has_media() + } + + #[inline] + pub fn has_dmca_media(&self) -> bool { + self.ctx.candidate.has_dmca_media() + } + + #[inline] + pub fn is_nsfw_flagged(&self) -> bool { + self.ctx.candidate.is_nsfw_flagged() + } + + #[inline] + pub fn has_nsfw_user_flag(&self) -> bool { + self.ctx.candidate.tweet_features.nsfw.user + } + + #[inline] + pub fn has_nsfw_admin_flag(&self) -> bool { + self.ctx.candidate.tweet_features.nsfw.admin + } + + #[inline] + pub fn is_exclusive(&self) -> bool { + self.ctx.candidate.exclusive_content.is_some() + } +} + +#[derive(Clone, Copy)] +pub struct AuthorPredicates<'a> { + ctx: &'a RuleContext<'a>, +} + +impl AuthorPredicates<'_> { + #[inline] + pub fn is_suspended(&self) -> bool { + self.ctx.candidate.author_features.is_suspended + } + + #[inline] + pub fn is_deactivated(&self) -> bool { + self.ctx.candidate.author_features.is_deactivated + } + + #[inline] + pub fn is_erased(&self) -> bool { + self.ctx.candidate.author_features.is_erased + } + + #[inline] + pub fn is_offboarded(&self) -> bool { + self.ctx.candidate.author_features.is_offboarded + } + + #[inline] + pub fn is_protected(&self) -> bool { + self.ctx.candidate.author_features.is_protected + } + + #[inline] + pub fn is_nsfw_user(&self) -> bool { + self.ctx.candidate.author_features.is_nsfw_user + } + + #[inline] + pub fn is_nsfw_admin(&self) -> bool { + self.ctx.candidate.author_features.is_nsfw_admin + } + + #[inline] + pub fn has_user_label(&self, label: LabelValue) -> bool { + self.ctx.candidate.author_has_user_label(label) + } +} + +#[derive(Clone, Copy)] +pub struct TakedownPredicates<'a> { + ctx: &'a RuleContext<'a>, +} + +impl TakedownPredicates<'_> { + #[inline] + pub fn legal_in_viewer_country(&self) -> bool { + self.in_viewer_country(legal_takedown_country) + } + + #[inline] + pub fn local_laws_in_viewer_country(&self) -> bool { + self.in_viewer_country(local_laws_takedown_country) + } + + #[inline] + fn in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool { + let viewer_country = self.ctx.viewer.country_code.as_deref(); + self.ctx + .candidate + .tweet_features + .takedown_reasons + .iter() + .filter_map(extractor) + .any(|c| { + c.eq_ignore_ascii_case(WORLDWIDE_COUNTRY_CODE) + || c.eq_ignore_ascii_case(WORLDWIDE_COPYRIGHT_COUNTRY_CODE) + || viewer_country.is_some_and(|v| c.eq_ignore_ascii_case(v)) + }) + } + + #[inline] + pub fn media_restricted_in_viewer_country(&self) -> bool { + let country = self + .ctx + .viewer + .country_code + .as_deref() + .unwrap_or(WORLDWIDE_COUNTRY_CODE); + let allow = &self.ctx.candidate.tweet_features.media.geo_allow_list; + let deny = &self.ctx.candidate.tweet_features.media.geo_deny_list; + (!allow.is_empty() && !allow.iter().any(|c| c.eq_ignore_ascii_case(country))) + || deny.iter().any(|c| c.eq_ignore_ascii_case(country)) + } +} + +const WORLDWIDE_COUNTRY_CODE: &str = "xx"; +const WORLDWIDE_COPYRIGHT_COUNTRY_CODE: &str = "xy"; + +fn legal_takedown_country(reason: &TakedownReason) -> Option<&str> { + match reason { + TakedownReason::LegalRequest { country_code } + | TakedownReason::UnspecifiedReason { country_code } => Some(country_code), + TakedownReason::Dmca => Some(WORLDWIDE_COPYRIGHT_COUNTRY_CODE), + _ => None, + } +} + +fn local_laws_takedown_country(reason: &TakedownReason) -> Option<&str> { + match reason { + TakedownReason::BystanderReport { country_code } => Some(country_code), + _ => None, + } +} diff --git a/visibility-filtering/rules/fixtures.rs b/visibility-filtering/rules/fixtures.rs new file mode 100644 index 00000000..614bb9ab --- /dev/null +++ b/visibility-filtering/rules/fixtures.rs @@ -0,0 +1,191 @@ +use crate::models::{ + AuthorFeatures, HydratedTweetCandidate, NsfwFeature, SafetyLabelMap, SafetyLabelType, + TweetFeatures, UserLabelSet, VfAction, Viewer, ViewerAuthorRelationship, ViewerFeatures, +}; +use crate::rules::rule_spec::RuleSpec; +use crate::rules::test_context; +use std::collections::HashSet; +use xai_visibility_filtering::models::FilteredReason; +use xai_x_thrift::user_labels::LabelValue; + +const TWEET_ID: u64 = 1; +const AUTHOR_ID: u64 = 100; +pub(crate) const VIEWER_ID: u64 = 999; + +pub(super) fn assert_drops( + spec: &RuleSpec, + viewer: &ViewerFeatures, + candidate: &HydratedTweetCandidate, + expected: &FilteredReason, +) { + let action = spec.evaluate(&test_context(viewer, candidate)); + assert!( + matches!(&action, VfAction::Drop(reason) if reason == expected), + "{} should drop with {expected:?}, got {action:?}", + spec.name() + ); +} + +pub(super) fn assert_allows( + spec: &RuleSpec, + viewer: &ViewerFeatures, + candidate: &HydratedTweetCandidate, +) { + let action = spec.evaluate(&test_context(viewer, candidate)); + assert!( + matches!(action, VfAction::Allow), + "{} should allow, got {action:?}", + spec.name() + ); +} + +pub(crate) fn viewer(id: u64) -> ViewerFeatures { + ViewerFeatures { + viewer: Viewer::LoggedIn(id), + ..Default::default() + } +} + +pub(crate) fn author_viewer() -> ViewerFeatures { + viewer(AUTHOR_ID) +} + +pub(crate) fn logged_out_viewer() -> ViewerFeatures { + ViewerFeatures { + viewer: Viewer::LoggedOut, + ..Default::default() + } +} + +pub(crate) fn sensitive_opt_in_viewer() -> ViewerFeatures { + ViewerFeatures { + allows_sensitive_media: true, + ..viewer(VIEWER_ID) + } +} + +pub(crate) fn candidate() -> CandidateBuilder { + CandidateBuilder { + candidate: HydratedTweetCandidate { + tweet_id: TWEET_ID, + author_id: AUTHOR_ID, + ..Default::default() + }, + labels: HashSet::new(), + user_labels: HashSet::new(), + } +} + +pub(crate) struct CandidateBuilder { + candidate: HydratedTweetCandidate, + labels: HashSet, + user_labels: HashSet, +} + +impl CandidateBuilder { + pub(crate) fn tweet_id(mut self, id: u64) -> Self { + self.candidate.tweet_id = id; + self + } + + pub(crate) fn author_id(mut self, id: u64) -> Self { + self.candidate.author_id = id; + self + } + + pub(crate) fn with_label(mut self, label: SafetyLabelType) -> Self { + self.labels.insert(label); + self + } + + pub(crate) fn with_author_user_label(mut self, label: LabelValue) -> Self { + self.user_labels.insert(label); + self + } + + pub(crate) fn with_tweet_features(mut self, features: TweetFeatures) -> Self { + self.candidate.tweet_features = features; + self + } + + pub(crate) fn with_author_features(mut self, features: AuthorFeatures) -> Self { + self.candidate.author_features = features; + self + } + + pub(crate) fn with_relationship(mut self, relationship: ViewerAuthorRelationship) -> Self { + self.candidate.relationship = relationship; + self + } + + pub(crate) fn followed(mut self) -> Self { + self.candidate.relationship.viewer_follows_author = true; + self + } + + pub(crate) fn with_media(mut self) -> Self { + self.candidate.tweet_features.media.has_media = true; + self + } + + pub(crate) fn retweet_of(mut self, source_tweet_id: u64) -> Self { + self.candidate.tweet_features.core.source_tweet_id = Some(source_tweet_id); + self + } + + pub(crate) fn build(self) -> HydratedTweetCandidate { + let mut candidate = self.candidate; + if !self.labels.is_empty() { + candidate.safety_labels = SafetyLabelMap::new(self.labels); + } + if !self.user_labels.is_empty() { + candidate.author_features.user_labels = UserLabelSet::new(self.user_labels); + } + candidate + } +} + +pub(super) fn nsfw_flag_media_candidates() -> [HydratedTweetCandidate; 4] { + let author_user = candidate() + .with_media() + .with_author_features(AuthorFeatures { + is_nsfw_user: true, + ..Default::default() + }) + .build(); + let tweet_user = candidate() + .with_tweet_features(TweetFeatures { + nsfw: NsfwFeature { + user: true, + admin: false, + }, + ..Default::default() + }) + .with_media() + .build(); + let tweet_admin = candidate() + .with_tweet_features(TweetFeatures { + nsfw: NsfwFeature { + user: false, + admin: true, + }, + ..Default::default() + }) + .with_media() + .build(); + let both = candidate() + .with_author_features(AuthorFeatures { + is_nsfw_admin: true, + ..Default::default() + }) + .with_tweet_features(TweetFeatures { + nsfw: NsfwFeature { + user: true, + admin: false, + }, + ..Default::default() + }) + .with_media() + .build(); + [author_user, tweet_user, tweet_admin, both] +} diff --git a/visibility-filtering/rules/golden_corpus.rs b/visibility-filtering/rules/golden_corpus.rs new file mode 100644 index 00000000..8a1598e6 --- /dev/null +++ b/visibility-filtering/rules/golden_corpus.rs @@ -0,0 +1,1144 @@ +use crate::models::{ + AuthorFeatures, ExclusiveContentFeatures, HydratedTweetCandidate, SafetyLabelType, + TweetFeatures, VfAction, ViewerAge, ViewerAuthorRelationship, ViewerFeatures, +}; +use crate::rules::fixtures::{ + author_viewer, candidate, logged_out_viewer, sensitive_opt_in_viewer, viewer, VIEWER_ID, +}; +use crate::rules::{RuleEngine, SafetyLevel}; +use std::collections::BTreeSet; +use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; +use xai_visibility_filtering::models::{ + Action, DropReason, FilteredReason, SafetyResult, SafetyResultReason, +}; +use xai_x_thrift::user_labels::LabelValue; +use SafetyLevel::{FilterAll, TimelineHome, TimelineHomeRecommendations}; +use VfAction::{Allow, Drop, Interstitial}; + +struct Case { + name: &'static str, + level: SafetyLevel, + viewer: ViewerFeatures, + candidate: HydratedTweetCandidate, + expected_action: VfAction, + expected_decided_by: Option<&'static str>, +} + +#[test] +fn golden_corpus_pins_policy_verdicts() { + let rule_engine = RuleEngine::for_tests(); + let mut failures = Vec::new(); + for case in cases() { + let verdict = rule_engine.evaluate(case.level, &case.viewer, &case.candidate); + if !action_eq(&verdict.action, &case.expected_action) + || verdict.decided_by != case.expected_decided_by + { + failures.push(format!( + "{} [{:?}]:\n expected {:?} decided_by {:?}\n got {:?} decided_by {:?}", + case.name, + case.level, + case.expected_action, + case.expected_decided_by, + verdict.action, + verdict.decided_by, + )); + } + } + assert!( + failures.is_empty(), + "{} corpus case(s) diverged:\n{}", + failures.len(), + failures.join("\n") + ); +} + +#[test] +fn every_wired_rule_decides_a_corpus_case() { + let rule_engine = RuleEngine::for_tests(); + let wired: BTreeSet<&'static str> = [FilterAll, TimelineHome, TimelineHomeRecommendations] + .into_iter() + .flat_map(|level| rule_engine.wired_rule_names(level)) + .collect(); + let deciders: BTreeSet<&'static str> = cases() + .iter() + .filter_map(|c| c.expected_decided_by) + .collect(); + let missing: Vec<&&'static str> = wired.difference(&deciders).collect(); + assert!( + missing.is_empty(), + "rules wired in RuleEngine but never the decider of any corpus case: {missing:?}" + ); +} + +#[test] +fn corpus_case_names_are_unique() { + let cases = cases(); + let names: BTreeSet<&'static str> = cases.iter().map(|c| c.name).collect(); + assert_eq!(names.len(), cases.len()); +} + +fn action_eq(left: &VfAction, right: &VfAction) -> bool { + match (left, right) { + (Allow, Allow) => true, + (Drop(l), Drop(r)) => l == r, + (Interstitial(l), Interstitial(r)) => l == r, + _ => false, + } +} + +fn cases() -> Vec { + let mut cases = filter_all_cases(); + cases.extend(baseline_cases()); + cases.extend(author_state_cases()); + cases.extend(relationship_cases()); + cases.extend(tweet_label_cases()); + cases.extend(tweet_shape_cases()); + cases.extend(age_gating_cases()); + cases.extend(exclusive_content_cases()); + cases.extend(interstitial_cases()); + cases.extend(oon_media_cases()); + cases.extend(oon_tweet_label_cases()); + cases.extend(oon_user_label_cases()); + cases.extend(interaction_cases()); + cases +} + +fn author_candidate(set: fn(&mut AuthorFeatures)) -> HydratedTweetCandidate { + let mut features = AuthorFeatures::default(); + set(&mut features); + candidate().with_author_features(features).build() +} + +fn tweet_candidate(set: fn(&mut TweetFeatures)) -> HydratedTweetCandidate { + let mut features = TweetFeatures::default(); + set(&mut features); + candidate().with_tweet_features(features).build() +} + +fn relationship_candidate(set: fn(&mut ViewerAuthorRelationship)) -> HydratedTweetCandidate { + let mut relationship = ViewerAuthorRelationship::default(); + set(&mut relationship); + candidate().with_relationship(relationship).build() +} + +fn labeled(label: SafetyLabelType) -> HydratedTweetCandidate { + candidate().with_label(label).build() +} + +fn labeled_media(label: SafetyLabelType) -> HydratedTweetCandidate { + candidate().with_label(label).with_media().build() +} + +fn user_labeled(label: LabelValue) -> HydratedTweetCandidate { + candidate().with_author_user_label(label).build() +} + +fn user_labeled_follower(label: LabelValue) -> HydratedTweetCandidate { + candidate().with_author_user_label(label).followed().build() +} + +fn stale_candidate() -> HydratedTweetCandidate { + tweet_candidate(|t| { + t.edit_control = Some(EditControl::Initial(EditControlInitial { + edit_tweet_ids: vec![1, 2], + ..Default::default() + })) + }) +} + +fn takedown_candidate(reason: TakedownReason) -> HydratedTweetCandidate { + candidate() + .with_tweet_features(TweetFeatures { + takedown_reasons: vec![reason], + ..Default::default() + }) + .build() +} + +fn exclusive_candidate(viewer_super_follows_author: bool) -> HydratedTweetCandidate { + let mut c = candidate().build(); + c.exclusive_content = Some(ExclusiveContentFeatures { + conversation_author_id: 42, + viewer_super_follows_author, + }); + c +} + +fn viewer_in_country(code: &str) -> ViewerFeatures { + ViewerFeatures { + country_code: Some(code.to_string()), + ..viewer(VIEWER_ID) + } +} + +fn viewer_with_age(age: ViewerAge) -> ViewerFeatures { + ViewerFeatures { + viewer_age: age, + ..viewer(VIEWER_ID) + } +} + +fn no_stated_age_viewer(account_country_code: &str) -> ViewerFeatures { + ViewerFeatures { + account_country_code: Some(account_country_code.to_string()), + ..viewer_with_age(ViewerAge::NotStated) + } +} + +fn nsfw_high_precision_reason() -> FilteredReason { + FilteredReason::SafetyResult(SafetyResult { + reason: Some(SafetyResultReason::NsfwHighPrecision), + action: Action::Drop(DropReason {}), + }) +} + +fn filter_all_cases() -> Vec { + vec![ + Case { + name: "filter_all_drops_pristine_candidate", + level: FilterAll, + viewer: viewer(VIEWER_ID), + candidate: candidate().build(), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("FilterAllRule"), + }, + Case { + name: "filter_all_drops_even_self_view", + level: FilterAll, + viewer: author_viewer(), + candidate: candidate().build(), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("FilterAllRule"), + }, + ] +} + +fn baseline_cases() -> Vec { + vec![ + Case { + name: "home_allows_pristine_candidate", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: candidate().build(), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "recommendations_allow_pristine_candidate", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: candidate().build(), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "home_allows_pristine_candidate_for_logged_out", + level: TimelineHome, + viewer: logged_out_viewer(), + candidate: candidate().build(), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "home_allows_egregious_nsfw_tweet_label", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::EGREGIOUS_NSFW), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "recommendations_allow_egregious_nsfw_tweet_label", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::EGREGIOUS_NSFW), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "home_allows_egregious_nsfw_user_label", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::EGREGIOUS_NSFW), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "recommendations_allow_egregious_nsfw_user_label", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::EGREGIOUS_NSFW), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "home_allows_recommendations_blacklist_user_label", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::RECOMMENDATIONS_BLACKLIST), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "recommendations_allow_recommendations_blacklist_user_label", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::RECOMMENDATIONS_BLACKLIST), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn author_state_cases() -> Vec { + vec![ + Case { + name: "suspended_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_suspended = true), + expected_action: Drop(FilteredReason::AuthorIsSuspended), + expected_decided_by: Some("SuspendedAuthorRule"), + }, + Case { + name: "suspended_author_allows_self_view", + level: TimelineHome, + viewer: author_viewer(), + candidate: author_candidate(|a| a.is_suspended = true), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "deactivated_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_deactivated = true), + expected_action: Drop(FilteredReason::AuthorIsDeactivated), + expected_decided_by: Some("DeactivatedAuthorRule"), + }, + Case { + name: "erased_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_erased = true), + expected_action: Drop(FilteredReason::AuthorAccountIsInactive), + expected_decided_by: Some("ErasedAuthorRule"), + }, + Case { + name: "offboarded_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_offboarded = true), + expected_action: Drop(FilteredReason::AuthorAccountIsInactive), + expected_decided_by: Some("OffboardedAuthorRule"), + }, + Case { + name: "protected_author_drops_non_follower", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_protected = true), + expected_action: Drop(FilteredReason::AuthorIsProtected), + expected_decided_by: Some("ProtectedAuthorDropRule"), + }, + Case { + name: "protected_author_allows_follower", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: { + let features = AuthorFeatures { + is_protected: true, + ..Default::default() + }; + candidate() + .with_author_features(features) + .followed() + .build() + }, + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "protected_author_drops_logged_out", + level: TimelineHome, + viewer: logged_out_viewer(), + candidate: author_candidate(|a| a.is_protected = true), + expected_action: Drop(FilteredReason::AuthorIsProtected), + expected_decided_by: Some("ProtectedAuthorDropRule"), + }, + ] +} + +fn relationship_cases() -> Vec { + vec![ + Case { + name: "viewer_blocking_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: relationship_candidate(|r| r.viewer_blocks_author = true), + expected_action: Drop(FilteredReason::AuthorBlockViewer), + expected_decided_by: Some("ViewerBlocksAuthorRule"), + }, + Case { + name: "viewer_muting_author_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: relationship_candidate(|r| r.viewer_mutes_author = true), + expected_action: Drop(FilteredReason::ViewerMutesAuthor), + expected_decided_by: Some("ViewerMutesAuthorRule"), + }, + Case { + name: "block_decides_before_mute", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: relationship_candidate(|r| { + r.viewer_blocks_author = true; + r.viewer_mutes_author = true; + }), + expected_action: Drop(FilteredReason::AuthorBlockViewer), + expected_decided_by: Some("ViewerBlocksAuthorRule"), + }, + Case { + name: "muted_retweets_drop_retweet", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: { + let relationship = ViewerAuthorRelationship { + viewer_mutes_retweets_from_author: true, + ..Default::default() + }; + candidate() + .with_relationship(relationship) + .retweet_of(2) + .build() + }, + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("MutedRetweetsRule"), + }, + Case { + name: "muted_retweets_allow_original_tweet", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: relationship_candidate(|r| r.viewer_mutes_retweets_from_author = true), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn tweet_label_cases() -> Vec { + vec![ + Case { + name: "pdna_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::PDNA), + expected_action: Drop(nsfw_high_precision_reason()), + expected_decided_by: Some("PdnaTweetLabelRule"), + }, + Case { + name: "pdna_label_allows_self_view", + level: TimelineHome, + viewer: author_viewer(), + candidate: labeled(SafetyLabelType::PDNA), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "bounce_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::BOUNCE), + expected_action: Drop(FilteredReason::TweetIsBounced), + expected_decided_by: Some("BounceTweetLabelRule"), + }, + Case { + name: "spam_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::SPAM), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("SpamTweetLabelRule"), + }, + Case { + name: "for_emergency_use_only_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOR_EMERGENCY_USE_ONLY), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("ForEmergencyUseOnlyDropRule"), + }, + Case { + name: "for_emergency_use_only_label_drops_even_self_view", + level: TimelineHome, + viewer: author_viewer(), + candidate: labeled(SafetyLabelType::FOR_EMERGENCY_USE_ONLY), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("ForEmergencyUseOnlyDropRule"), + }, + Case { + name: "fosnr_hateful_conduct_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_HATEFUL_CONDUCT), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("FosnrHatefulConductDropRule"), + }, + Case { + name: "fosnr_violent_speech_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_VIOLENT_SPEECH), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("FosnrViolentSpeechDropRule"), + }, + Case { + name: "fosnr_abuse_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_ABUSE), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("FosnrAbuseDropRule"), + }, + Case { + name: "fosnr_civic_integrity_label_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_CIVIC_INTEGRITY), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("FosnrCivicIntegrityDropRule"), + }, + ] +} + +fn tweet_shape_cases() -> Vec { + vec![ + Case { + name: "nullcast_tweet_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.is_nullcast = true), + expected_action: Drop(FilteredReason::TweetIsNullcast), + expected_decided_by: Some("NullcastedTweetDropRule"), + }, + Case { + name: "nullcast_retweet_allows", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: { + let features = TweetFeatures { + is_nullcast: true, + ..Default::default() + }; + candidate() + .with_tweet_features(features) + .retweet_of(2) + .build() + }, + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "stale_edit_tweet_drops", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: stale_candidate(), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropStaleTweetsRule"), + }, + Case { + name: "legal_takedown_drops_in_withheld_country", + level: TimelineHome, + viewer: viewer_in_country("us"), + candidate: takedown_candidate(TakedownReason::LegalRequest { + country_code: "us".to_string(), + }), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLegalTakendownPostRule"), + }, + Case { + name: "legal_takedown_allows_other_country", + level: TimelineHome, + viewer: viewer_in_country("fr"), + candidate: takedown_candidate(TakedownReason::LegalRequest { + country_code: "us".to_string(), + }), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "local_laws_takedown_drops_in_withheld_country", + level: TimelineHome, + viewer: viewer_in_country("de"), + candidate: takedown_candidate(TakedownReason::BystanderReport { + country_code: "de".to_string(), + }), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLocalLawsTakendownPostRule"), + }, + Case { + name: "legal_takedown_worldwide_drops_for_us_viewer", + level: TimelineHome, + viewer: viewer_in_country("us"), + candidate: takedown_candidate(TakedownReason::LegalRequest { + country_code: "xx".to_string(), + }), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLegalTakendownPostRule"), + }, + Case { + name: "legal_takedown_worldwide_drops_without_viewer_country", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: takedown_candidate(TakedownReason::LegalRequest { + country_code: "xx".to_string(), + }), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLegalTakendownPostRule"), + }, + Case { + name: "local_laws_takedown_worldwide_drops_for_any_viewer", + level: TimelineHome, + viewer: viewer_in_country("us"), + candidate: takedown_candidate(TakedownReason::BystanderReport { + country_code: "xx".to_string(), + }), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLocalLawsTakendownPostRule"), + }, + Case { + name: "dmca_takedown_drops_for_any_viewer", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: takedown_candidate(TakedownReason::Dmca), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropLegalTakendownPostRule"), + }, + Case { + name: "dmca_takedown_allows_author", + level: TimelineHome, + viewer: author_viewer(), + candidate: takedown_candidate(TakedownReason::Dmca), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn age_gating_cases() -> Vec { + vec![ + Case { + name: "logged_out_viewer_drops_sensitive_media", + level: TimelineHome, + viewer: logged_out_viewer(), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("SensitiveViewerLoggedOutDropRule"), + }, + Case { + name: "underage_viewer_drops_sensitive_media", + level: TimelineHome, + viewer: viewer_with_age(ViewerAge::Known(17)), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("SensitiveViewerUnderageDropRule"), + }, + Case { + name: "no_stated_age_in_gating_country_drops_sensitive_media", + level: TimelineHome, + viewer: no_stated_age_viewer("gb"), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("SensitiveViewerNoStatedAgeDropRule"), + }, + Case { + name: "no_stated_age_outside_gating_country_allows_sensitive_media", + level: TimelineHome, + viewer: no_stated_age_viewer("us"), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "known_adult_age_allows_sensitive_media_in_network", + level: TimelineHome, + viewer: viewer_with_age(ViewerAge::Known(30)), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "gizmoduck_age_lookup_failed_drops_sensitive_media", + level: TimelineHome, + viewer: viewer_with_age(ViewerAge::LookupFailed), + candidate: labeled_media(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("SensitiveViewerUnderageDropRule"), + }, + ] +} + +fn exclusive_content_cases() -> Vec { + vec![ + Case { + name: "exclusive_tweet_drops_non_subscriber", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: exclusive_candidate(false), + expected_action: Drop(FilteredReason::ExclusiveTweet), + expected_decided_by: Some("DropExclusiveTweetContentRule"), + }, + Case { + name: "exclusive_tweet_allows_super_follower", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: exclusive_candidate(true), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "exclusive_tweet_drops_logged_out", + level: TimelineHome, + viewer: logged_out_viewer(), + candidate: exclusive_candidate(false), + expected_action: Drop(FilteredReason::ExclusiveTweet), + expected_decided_by: Some("DropExclusiveTweetContentRule"), + }, + ] +} + +fn interstitial_cases() -> Vec { + vec![ + Case { + name: "nsfw_high_precision_label_interstitials_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_HIGH_PRECISION), + expected_action: Interstitial(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwHighPrecisionInterstitialRule"), + }, + Case { + name: "gore_and_violence_label_interstitials_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION), + expected_action: Interstitial(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("GoreAndViolenceInterstitialRule"), + }, + Case { + name: "nsfw_card_image_label_interstitials_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_CARD_IMAGE), + expected_action: Interstitial(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwCardImageInterstitialRule"), + }, + Case { + name: "nsfw_author_with_media_interstitials_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: { + let features = AuthorFeatures { + is_nsfw_user: true, + ..Default::default() + }; + candidate() + .with_author_features(features) + .with_media() + .build() + }, + expected_action: Interstitial(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwAuthorInterstitialRule"), + }, + Case { + name: "nsfw_author_without_media_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_nsfw_user = true), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_interstitial_exempts_self_view", + level: TimelineHome, + viewer: author_viewer(), + candidate: labeled(SafetyLabelType::NSFW_HIGH_PRECISION), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_interstitial_exempts_sensitive_opt_in_viewer", + level: TimelineHome, + viewer: sensitive_opt_in_viewer(), + candidate: labeled(SafetyLabelType::NSFW_HIGH_PRECISION), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn oon_media_cases() -> Vec { + vec![ + Case { + name: "dmca_media_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.media.has_dmca_media = true), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropTweetsWithDmcaMediaRule"), + }, + Case { + name: "dmca_media_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.media.has_dmca_media = true), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "geo_denied_media_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer_in_country("de"), + candidate: tweet_candidate(|t| t.media.geo_deny_list = vec!["de".to_string()]), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropTweetsWithGeoRestrictedMediaRule"), + }, + Case { + name: "geo_allow_listed_media_drops_unknown_country_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.media.geo_allow_list = vec!["us".to_string()]), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DropTweetsWithGeoRestrictedMediaRule"), + }, + Case { + name: "nsfw_user_author_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_nsfw_user = true), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("DropNsfwUserAuthorRule"), + }, + Case { + name: "nsfw_admin_author_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: author_candidate(|a| a.is_nsfw_admin = true), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("DropNsfwAdminAuthorRule"), + }, + Case { + name: "tweet_nsfw_user_flag_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.nsfw.user = true), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("TweetNsfwUserDropRule"), + }, + Case { + name: "tweet_nsfw_user_flag_drops_even_self_view_oon", + level: TimelineHomeRecommendations, + viewer: author_viewer(), + candidate: tweet_candidate(|t| t.nsfw.user = true), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("TweetNsfwUserDropRule"), + }, + Case { + name: "tweet_nsfw_admin_flag_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: tweet_candidate(|t| t.nsfw.admin = true), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("TweetNsfwAdminDropRule"), + }, + ] +} + +fn oon_tweet_label_cases() -> Vec { + vec![ + Case { + name: "nsfw_high_recall_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwHighRecallDropRule"), + }, + Case { + name: "nsfw_high_precision_label_drop_beats_interstitial_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_HIGH_PRECISION), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwHighPrecisionOonDropRule"), + }, + Case { + name: "gore_and_violence_label_drop_beats_interstitial_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("GoreAndViolenceOonDropRule"), + }, + Case { + name: "nsfw_card_image_label_drop_beats_interstitial_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_CARD_IMAGE), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwCardImageOonDropRule"), + }, + Case { + name: "sensitive_opt_in_does_not_save_oon_drop", + level: TimelineHomeRecommendations, + viewer: sensitive_opt_in_viewer(), + candidate: labeled(SafetyLabelType::NSFW_HIGH_PRECISION), + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("NsfwHighPrecisionOonDropRule"), + }, + Case { + name: "do_not_amplify_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::DO_NOT_AMPLIFY), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("DoNotAmplifyOonDropRule"), + }, + Case { + name: "malicious_url_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::MALICIOUS_URL), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("MaliciousUrlOonDropRule"), + }, + Case { + name: "malicious_url_label_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::MALICIOUS_URL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "malicious_url_label_allows_self_view_oon", + level: TimelineHomeRecommendations, + viewer: author_viewer(), + candidate: labeled(SafetyLabelType::MALICIOUS_URL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "spam_high_recall_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::SPAM_HIGH_RECALL), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("SpamHighRecallDropRule"), + }, + Case { + name: "spam_high_recall_label_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::SPAM_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_text_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::NSFW_TEXT), + expected_action: Drop(nsfw_high_precision_reason()), + expected_decided_by: Some("NsfwTextTweetLabelDropRule"), + }, + Case { + name: "fosnr_abuse_insults_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_ABUSE_INSULTS), + expected_action: Drop(FilteredReason::PossiblyUndesirable), + expected_decided_by: Some("FosnrAbuseInsultsOonDropRule"), + }, + Case { + name: "fosnr_abuse_insults_label_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: labeled(SafetyLabelType::FOSNR_ABUSE_INSULTS), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn oon_user_label_cases() -> Vec { + vec![ + Case { + name: "nsfw_high_recall_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_HIGH_RECALL), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("NsfwHighRecallUserLabelRule"), + }, + Case { + name: "nsfw_high_recall_user_label_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_high_recall_user_label_allows_self_view_oon", + level: TimelineHomeRecommendations, + viewer: author_viewer(), + candidate: user_labeled(LabelValue::NSFW_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_high_precision_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_HIGH_PRECISION), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("NsfwHighPrecisionUserLabelRule"), + }, + Case { + name: "spam_high_recall_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::SPAM_HIGH_RECALL), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("SpamHighRecallUserLabelRule"), + }, + Case { + name: "compromised_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::COMPROMISED), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("CompromisedUserLabelRule"), + }, + Case { + name: "read_only_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::READ_ONLY), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("ReadOnlyUserLabelRule"), + }, + Case { + name: "impersonation_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::IMPERSONATION_HIGH_PRECISION), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("ImpersonationHighPrecisionUserLabelRule"), + }, + Case { + name: "nsfw_avatar_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_AVATAR_IMAGE), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("NsfwAvatarImageRule"), + }, + Case { + name: "nsfw_banner_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_BANNER_IMAGE), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("NsfwBannerImageRule"), + }, + Case { + name: "abusive_high_recall_user_label_drops_non_follower_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::ABUSIVE_HIGH_RECALL), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("AbusiveHighRecallRule"), + }, + Case { + name: "abusive_high_recall_user_label_allows_follower_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled_follower(LabelValue::ABUSIVE_HIGH_RECALL), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "nsfw_near_perfect_user_label_drops_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_NEAR_PERFECT), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("NsfwNearPerfectAuthorRule"), + }, + Case { + name: "nsfw_near_perfect_user_label_allows_in_network", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::NSFW_NEAR_PERFECT), + expected_action: Allow, + expected_decided_by: None, + }, + Case { + name: "do_not_amplify_user_label_drops_non_follower_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled(LabelValue::DO_NOT_AMPLIFY), + expected_action: Drop(FilteredReason::UnspecifiedReason), + expected_decided_by: Some("DoNotAmplifyNonFollowerRule"), + }, + Case { + name: "do_not_amplify_user_label_allows_follower_oon", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: user_labeled_follower(LabelValue::DO_NOT_AMPLIFY), + expected_action: Allow, + expected_decided_by: None, + }, + ] +} + +fn interaction_cases() -> Vec { + vec![ + Case { + name: "drop_short_circuits_before_interstitial_attribution", + level: TimelineHome, + viewer: viewer(VIEWER_ID), + candidate: { + let features = AuthorFeatures { + is_suspended: true, + ..Default::default() + }; + candidate() + .with_author_features(features) + .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) + .with_media() + .build() + }, + expected_action: Drop(FilteredReason::AuthorIsSuspended), + expected_decided_by: Some("SuspendedAuthorRule"), + }, + Case { + name: "later_oon_drop_beats_earlier_nsfw_author_interstitial", + level: TimelineHomeRecommendations, + viewer: viewer(VIEWER_ID), + candidate: { + let features = AuthorFeatures { + is_nsfw_user: true, + ..Default::default() + }; + candidate() + .with_author_features(features) + .with_media() + .build() + }, + expected_action: Drop(FilteredReason::ContainNsfwMedia), + expected_decided_by: Some("DropNsfwUserAuthorRule"), + }, + ] +} diff --git a/visibility-filtering/rules/metrics.rs b/visibility-filtering/rules/metrics.rs index 32df2102..2fb790c6 100644 --- a/visibility-filtering/rules/metrics.rs +++ b/visibility-filtering/rules/metrics.rs @@ -12,6 +12,18 @@ const LATENCY_MS: &str = "filter_tweets_latency_ms"; const BATCH_SIZE: &str = "filter_tweets_batch_size"; const VERDICTS: &str = "filter_tweets_verdicts"; const VERDICTS_BY_RULE: &str = "filter_tweets_verdicts_by_rule"; +const LOGGED_OUT_VIEWER: &str = "filter_tweets_logged_out_viewer"; +const VIEWER_ID_NORMALIZED: &str = "filter_tweets_viewer_id_normalized"; + +pub(crate) fn record_viewer_state(raw: Option, normalized: Option) { + if normalized.is_some() { + return; + } + incr(LOGGED_OUT_VIEWER, &[], 1); + if raw.is_some() { + incr(VIEWER_ID_NORMALIZED, &[], 1); + } +} #[derive(Clone, Debug, Default, PartialEq, Eq)] pub(crate) struct AggregatedVerdicts { diff --git a/visibility-filtering/rules/mod.rs b/visibility-filtering/rules/mod.rs index 2750312a..83da0392 100644 --- a/visibility-filtering/rules/mod.rs +++ b/visibility-filtering/rules/mod.rs @@ -1,72 +1,19 @@ +mod author_rules; +pub mod context; +#[cfg(test)] +pub(crate) mod fixtures; +#[cfg(test)] +mod golden_corpus; pub mod metrics; -pub mod nsfw_age_gating; -pub mod nsfw_interstitial; -pub mod nullcast_rule; pub mod registry; -pub mod socialgraph_rules; -pub mod tes_rules; -pub mod tweet_flag_rules; -pub mod tweet_label_drops; -pub mod user_label_drops; -pub mod user_rules; +mod rule_spec; +mod tweet_rules; -use crate::models::{HydratedTweetCandidate, SafetyLabelType, VfAction, ViewerFeatures}; +use crate::models::VfAction; use xai_visibility_filtering::models::FilteredReason; -pub use registry::{Policies, SafetyLevel}; - -pub struct RuleContext<'a> { - safety_level: SafetyLevel, - viewer: &'a ViewerFeatures, - candidate: &'a HydratedTweetCandidate, -} - -impl<'a> RuleContext<'a> { - fn new( - safety_level: SafetyLevel, - viewer: &'a ViewerFeatures, - candidate: &'a HydratedTweetCandidate, - ) -> Self { - Self { - safety_level, - viewer, - candidate, - } - } - - pub fn safety_level(&self) -> SafetyLevel { - self.safety_level - } - - pub fn has_tweet_safety_label(&self, label: SafetyLabelType) -> bool { - self.candidate.has_safety_label(label) - } - - pub fn is_author_viewer(&self) -> bool { - self.candidate.is_author_viewer(self.viewer.viewer) - } - - pub fn viewer_follows_author(&self) -> bool { - self.candidate.viewer_follows_author() - } - - pub fn viewer_allows_sensitive_media(&self) -> bool { - self.viewer.allows_sensitive_media - } - - pub fn viewer(&self) -> &ViewerFeatures { - self.viewer - } - - pub fn candidate(&self) -> &HydratedTweetCandidate { - self.candidate - } -} - -pub trait Rule: Send + Sync { - fn name(&self) -> &'static str; - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction; -} +pub use context::RuleContext; +pub use registry::{RuleEngine, SafetyLevel}; #[derive(Clone, Debug)] pub struct Verdict { @@ -83,162 +30,14 @@ impl Verdict { } } -fn evaluate_rules(rules: &[Box], context: &RuleContext<'_>) -> Verdict { - let mut worst = VfAction::Allow; - let mut decided_by = None; - for rule in rules { - match rule.evaluate(context) { - VfAction::Drop(reason) => { - return Verdict { - action: VfAction::Drop(reason), - decided_by: Some(rule.name()), - }; - } - VfAction::Interstitial(reason) => { - if matches!(worst, VfAction::Allow) { - worst = VfAction::Interstitial(reason); - decided_by = Some(rule.name()); - } - } - VfAction::Allow => {} - } - } - Verdict { - action: worst, - decided_by, - } -} - #[cfg(test)] pub(crate) fn test_context<'a>( - viewer: &'a ViewerFeatures, - candidate: &'a HydratedTweetCandidate, + viewer: &'a crate::models::ViewerFeatures, + candidate: &'a crate::models::HydratedTweetCandidate, ) -> RuleContext<'a> { - RuleContext::new(SafetyLevel::TimelineHome, viewer, candidate) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::{Arc, Mutex}; - - struct FakeRule { - name: &'static str, - action: VfAction, - calls: Arc>>, - } - - impl Rule for FakeRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, _context: &RuleContext<'_>) -> VfAction { - self.calls.lock().unwrap().push(self.name); - self.action.clone() - } - } - - type CallLog = Arc>>; - - fn recording_rules(specs: &[(&'static str, VfAction)]) -> (Vec>, CallLog) { - let calls = CallLog::default(); - let rules = specs - .iter() - .map(|(name, action)| { - Box::new(FakeRule { - name, - action: action.clone(), - calls: calls.clone(), - }) as Box - }) - .collect(); - (rules, calls) - } - - fn context_inputs() -> (ViewerFeatures, HydratedTweetCandidate) { - (ViewerFeatures::default(), HydratedTweetCandidate::default()) - } + use std::sync::LazyLock; - #[test] - fn drop_short_circuits_later_rules() { - let (rules, calls) = recording_rules(&[ - ("allow", VfAction::Allow), - ("drop", VfAction::Drop(FilteredReason::AuthorIsSuspended)), - ("after_drop", VfAction::Allow), - ]); - let (viewer, candidate) = context_inputs(); - - let verdict = evaluate_rules(&rules, &test_context(&viewer, &candidate)); - - assert!(matches!( - verdict.action, - VfAction::Drop(FilteredReason::AuthorIsSuspended) - )); - assert_eq!(verdict.decided_by, Some("drop")); - assert_eq!(*calls.lock().unwrap(), vec!["allow", "drop"]); - } - - #[test] - fn first_interstitial_decided_by_sticks_without_short_circuit() { - let (rules, calls) = recording_rules(&[ - ( - "first_interstitial", - VfAction::Interstitial(FilteredReason::ContainNsfwMedia), - ), - ( - "second_interstitial", - VfAction::Interstitial(FilteredReason::UnspecifiedReason), - ), - ("after_interstitials", VfAction::Allow), - ]); - let (viewer, candidate) = context_inputs(); - - let verdict = evaluate_rules(&rules, &test_context(&viewer, &candidate)); - - assert!(matches!( - verdict.action, - VfAction::Interstitial(FilteredReason::ContainNsfwMedia) - )); - assert_eq!(verdict.decided_by, Some("first_interstitial")); - assert_eq!( - *calls.lock().unwrap(), - vec![ - "first_interstitial", - "second_interstitial", - "after_interstitials" - ] - ); - } - - #[test] - fn drop_after_interstitial_wins() { - let (rules, _) = recording_rules(&[ - ( - "interstitial", - VfAction::Interstitial(FilteredReason::ContainNsfwMedia), - ), - ("drop", VfAction::Drop(FilteredReason::AuthorIsSuspended)), - ]); - let (viewer, candidate) = context_inputs(); - - let verdict = evaluate_rules(&rules, &test_context(&viewer, &candidate)); - - assert!(matches!( - verdict.action, - VfAction::Drop(FilteredReason::AuthorIsSuspended) - )); - assert_eq!(verdict.decided_by, Some("drop")); - } - - #[test] - fn all_allows_is_allow_with_no_decider() { - let (rules, _) = recording_rules(&[("a", VfAction::Allow), ("b", VfAction::Allow)]); - let (viewer, candidate) = context_inputs(); - - let verdict = evaluate_rules(&rules, &test_context(&viewer, &candidate)); - - assert!(matches!(verdict.action, VfAction::Allow)); - assert_eq!(verdict.decided_by, None); - } + static NSFW_GATING_COUNTRIES: LazyLock = + LazyLock::new(crate::params::NsfwGatingCountries::starting_at_default); + RuleContext::new(viewer, candidate, &NSFW_GATING_COUNTRIES) } diff --git a/visibility-filtering/rules/nsfw_age_gating.rs b/visibility-filtering/rules/nsfw_age_gating.rs deleted file mode 100644 index a2a0e65e..00000000 --- a/visibility-filtering/rules/nsfw_age_gating.rs +++ /dev/null @@ -1,654 +0,0 @@ -use crate::models::{SafetyLabelType, VfAction, ViewerFeatures}; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -const NSFW_GATING_COUNTRIES: [&str; 16] = [ - "ar", "au", "br", "ca", "de", "es", "fr", "gb", "id", "it", "kr", "mx", "nl", "ph", "pt", "th", -]; - -fn nsfw_base_condition(context: &RuleContext<'_>) -> bool { - let candidate = context.candidate(); - !context.is_author_viewer() - && candidate.has_media() - && (context.has_tweet_safety_label(SafetyLabelType::NSFW_HIGH_PRECISION) - || context.has_tweet_safety_label(SafetyLabelType::NSFW_HIGH_RECALL) - || (candidate.is_nsfw_flagged() && !candidate.is_retweet())) -} - -fn graphic_base_condition(context: &RuleContext<'_>) -> bool { - let candidate = context.candidate(); - !context.is_author_viewer() - && candidate.has_media() - && context.has_tweet_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION) -} - -fn nsfw_no_media_label_condition(context: &RuleContext<'_>) -> bool { - !context.is_author_viewer() - && (context.has_tweet_safety_label(SafetyLabelType::NSFW_TEXT) - || context.has_tweet_safety_label(SafetyLabelType::NSFW_CARD_IMAGE)) -} - -fn sensitive_base_condition(context: &RuleContext<'_>) -> bool { - nsfw_base_condition(context) - || graphic_base_condition(context) - || nsfw_no_media_label_condition(context) -} - -fn viewer_in_gating_jurisdiction(viewer: &ViewerFeatures) -> bool { - viewer - .account_country_code - .as_deref() - .or(viewer.country_code.as_deref()) - .is_some_and(|c| NSFW_GATING_COUNTRIES.contains(&c)) -} - -pub struct SensitiveViewerLoggedOutDropRule; - -impl Rule for SensitiveViewerLoggedOutDropRule { - fn name(&self) -> &'static str { - "SensitiveViewerLoggedOutDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_is_logged_out() && sensitive_base_condition(context) { - VfAction::Drop(FilteredReason::ContainNsfwMedia) - } else { - VfAction::Allow - } - } -} - -pub struct SensitiveViewerUnderageDropRule; - -impl Rule for SensitiveViewerUnderageDropRule { - fn name(&self) -> &'static str { - "SensitiveViewerUnderageDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_is_underage() && sensitive_base_condition(context) { - VfAction::Drop(FilteredReason::ContainNsfwMedia) - } else { - VfAction::Allow - } - } -} - -pub struct SensitiveViewerNoStatedAgeDropRule; - -impl Rule for SensitiveViewerNoStatedAgeDropRule { - fn name(&self) -> &'static str { - "SensitiveViewerNoStatedAgeDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_has_no_stated_age() - && viewer_in_gating_jurisdiction(viewer) - && sensitive_base_condition(context) - { - VfAction::Drop(FilteredReason::ContainNsfwMedia) - } else { - VfAction::Allow - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, MediaFeature, NsfwFeature, SafetyLabel, - SafetyLabelMap, TweetFeatures, Viewer, ViewerAge, ViewerFeatures, - }; - use std::collections::HashMap; - - fn viewer(age: ViewerAge) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: false, - viewer_age: age, - country_code: Some("de".into()), - account_country_code: None, - } - } - - fn media_candidate_with_label(label: SafetyLabelType) -> HydratedTweetCandidate { - let mut labels = HashMap::new(); - labels.insert(label, SafetyLabel::default()); - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - tweet_features: TweetFeatures { - media: MediaFeature { - has_media: true, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - } - } - - fn nsfw_author_media_candidate() -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - media: MediaFeature { - has_media: true, - ..Default::default() - }, - ..Default::default() - }, - author_features: AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }, - ..Default::default() - } - } - - #[test] - fn underage_drops_nsfw_label_media() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_nsfw_high_recall_media() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_RECALL); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_gore_label_media() { - let c = media_candidate_with_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drop_even_when_opted_in() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - allows_sensitive_media: true, - ..viewer(ViewerAge::Known(15)) - }; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn adult_does_not_drop() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(18)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn unknown_age_fails_open() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerUnderageDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), - VfAction::Allow - )); - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), - VfAction::Allow - )); - } - - fn no_media_candidate_with_label(label: SafetyLabelType) -> HydratedTweetCandidate { - let mut candidate = media_candidate_with_label(label); - candidate.tweet_features.media.has_media = false; - candidate - } - - #[test] - fn underage_drops_nsfw_text_without_media() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_nsfw_card_image_without_media() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_CARD_IMAGE); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn no_stated_age_drops_nsfw_text_in_jurisdiction() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::NotStated), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn no_stated_age_allows_nsfw_text_outside_jurisdiction() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - let v = ViewerFeatures { - country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Allow - )); - } - - #[test] - fn logged_out_drops_nsfw_text_without_media() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - let v = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) - }; - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn logged_out_drops_nsfw_card_image_without_media() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_CARD_IMAGE); - let v = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) - }; - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn adult_allows_nsfw_text() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(18)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn unknown_age_allows_nsfw_text() { - let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - assert!(matches!( - SensitiveViewerUnderageDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), - VfAction::Allow - )); - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), - VfAction::Allow - )); - } - - #[test] - fn nsfw_text_self_view_is_exempt() { - let mut c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - c.author_id = 999; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn label_rule_requires_media() { - let mut c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.tweet_features.media.has_media = false; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn self_view_is_exempt() { - let mut c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.author_id = 999; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn no_stated_age_drops_in_jurisdiction() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::NotStated), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn no_stated_age_allows_outside_jurisdiction() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Allow - )); - } - - #[test] - fn no_stated_age_allows_missing_country() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - country_code: None, - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Allow - )); - } - - #[test] - fn no_stated_age_allows_non_gating_account_country_over_gating_request() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - country_code: Some("de".into()), - account_country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Allow - )); - } - - #[test] - fn no_stated_age_drops_gating_account_country() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - country_code: Some("us".into()), - account_country_code: Some("kr".into()), - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn no_stated_age_falls_back_to_request_country_when_account_absent() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - country_code: Some("de".into()), - account_country_code: None, - ..viewer(ViewerAge::NotStated) - }; - assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_nsfw_author_media() { - let c = nsfw_author_media_candidate(); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_nsfw_admin_author_media() { - let mut c = nsfw_author_media_candidate(); - c.author_features = AuthorFeatures { - is_nsfw_user: false, - is_nsfw_admin: true, - ..Default::default() - }; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - fn nsfw_tweet_flag_media_candidate() -> HydratedTweetCandidate { - let mut c = nsfw_author_media_candidate(); - c.author_features = AuthorFeatures::default(); - c.tweet_features.nsfw = NsfwFeature { - user: true, - admin: false, - }; - c - } - - #[test] - fn underage_drops_nsfw_tweet_flag_media() { - let c = nsfw_tweet_flag_media_candidate(); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_nsfw_admin_tweet_flag_media() { - let mut c = nsfw_tweet_flag_media_candidate(); - c.tweet_features.nsfw = NsfwFeature { - user: false, - admin: true, - }; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_drops_when_both_flag_sources_set() { - let mut c = nsfw_tweet_flag_media_candidate(); - c.author_features.is_nsfw_user = true; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Drop(_) - )); - } - - #[test] - fn underage_allows_no_flags_no_labels() { - let mut c = nsfw_tweet_flag_media_candidate(); - c.tweet_features.nsfw = NsfwFeature::default(); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn nsfw_tweet_flag_retweet_not_dropped() { - let mut c = nsfw_tweet_flag_media_candidate(); - c.tweet_features.core.source_tweet_id = Some(42); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn nsfw_tweet_flag_self_view_exempt() { - let mut c = nsfw_tweet_flag_media_candidate(); - c.author_id = 999; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn nsfw_author_retweet_not_dropped() { - let mut c = nsfw_author_media_candidate(); - c.tweet_features.core.source_tweet_id = Some(42); - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn nsfw_author_without_media_not_dropped() { - let mut c = nsfw_author_media_candidate(); - c.tweet_features.media.has_media = false; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } - - #[test] - fn logged_out_drops_nsfw_label_media() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - let v = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) - }; - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn logged_out_drops_nsfw_author_media() { - let c = nsfw_author_media_candidate(); - let v = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) - }; - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn logged_out_requires_media() { - let mut c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.tweet_features.media.has_media = false; - let v = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) - }; - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), - VfAction::Allow - )); - } - - #[test] - fn logged_in_not_handled_by_logged_out_rule() { - let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), - &c - )), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/nsfw_interstitial.rs b/visibility-filtering/rules/nsfw_interstitial.rs deleted file mode 100644 index c1d25ee3..00000000 --- a/visibility-filtering/rules/nsfw_interstitial.rs +++ /dev/null @@ -1,275 +0,0 @@ -use crate::models::{SafetyLabelType, VfAction}; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -#[derive(Clone, Copy)] -pub struct NsfwMediaInterstitialRule { - label: SafetyLabelType, - name: &'static str, -} - -impl NsfwMediaInterstitialRule { - pub const fn new(label: SafetyLabelType, name: &'static str) -> Self { - Self { label, name } - } -} - -impl Rule for NsfwMediaInterstitialRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.has_tweet_safety_label(self.label) - && !context.is_author_viewer() - && !context.viewer_allows_sensitive_media() - { - return VfAction::Interstitial(FilteredReason::ContainNsfwMedia); - } - VfAction::Allow - } -} - -pub static NSFW_HIGH_PRECISION_INTERSTITIAL: NsfwMediaInterstitialRule = - NsfwMediaInterstitialRule::new( - SafetyLabelType::NSFW_HIGH_PRECISION, - "NsfwHighPrecisionInterstitialRule", - ); - -pub static GORE_AND_VIOLENCE_INTERSTITIAL: NsfwMediaInterstitialRule = - NsfwMediaInterstitialRule::new( - SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, - "GoreAndViolenceInterstitialRule", - ); - -pub static NSFW_CARD_IMAGE_INTERSTITIAL: NsfwMediaInterstitialRule = NsfwMediaInterstitialRule::new( - SafetyLabelType::NSFW_CARD_IMAGE, - "NsfwCardImageInterstitialRule", -); - -pub struct NsfwAuthorInterstitialRule; - -impl Rule for NsfwAuthorInterstitialRule { - fn name(&self) -> &'static str { - "NsfwAuthorInterstitialRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.is_nsfw_flagged() - && candidate.has_media() - && !context.is_author_viewer() - && !context.viewer_allows_sensitive_media() - { - return VfAction::Interstitial(FilteredReason::ContainNsfwMedia); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, MediaFeature, NsfwFeature, SafetyLabel, - SafetyLabelMap, TweetFeatures, Viewer, ViewerFeatures, - }; - use std::collections::HashMap; - - fn viewer_default() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: false, - ..Default::default() - } - } - - fn viewer_sensitive_opt_in() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: true, - ..Default::default() - } - } - - fn candidate_with_label(label: SafetyLabelType) -> HydratedTweetCandidate { - let mut labels = HashMap::new(); - labels.insert(label, SafetyLabel::default()); - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - } - } - - #[test] - fn interstitial_blurs_non_opt_in() { - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Interstitial(_) - )); - } - - #[test] - fn interstitial_allows_opt_in() { - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_sensitive_opt_in(), &c)), - VfAction::Allow - )); - } - - #[test] - fn interstitial_allows_self_view() { - let mut c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.author_id = 999; - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Allow - )); - } - - #[test] - fn interstitial_allows_no_label() { - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Allow - )); - } - - fn nsfw_author_candidate() -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - media: MediaFeature { - has_media: true, - ..Default::default() - }, - ..Default::default() - }, - author_features: AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }, - ..Default::default() - } - } - - #[test] - fn author_interstitial_blurs_non_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_default(), - &nsfw_author_candidate() - )), - VfAction::Interstitial(_) - )); - } - - #[test] - fn author_interstitial_allows_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_sensitive_opt_in(), - &nsfw_author_candidate() - )), - VfAction::Allow - )); - } - - #[test] - fn author_interstitial_allows_when_no_media() { - let mut c = nsfw_author_candidate(); - c.tweet_features.media.has_media = false; - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Allow - )); - } - - fn nsfw_tweet_flag_candidate() -> HydratedTweetCandidate { - let mut c = nsfw_author_candidate(); - c.author_features = AuthorFeatures::default(); - c.tweet_features.nsfw = NsfwFeature { - user: true, - admin: false, - }; - c - } - - #[test] - fn tweet_flag_interstitial_blurs_non_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_default(), - &nsfw_tweet_flag_candidate() - )), - VfAction::Interstitial(_) - )); - } - - #[test] - fn tweet_admin_flag_interstitial_blurs_non_opt_in() { - let mut c = nsfw_tweet_flag_candidate(); - c.tweet_features.nsfw = NsfwFeature { - user: false, - admin: true, - }; - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Interstitial(_) - )); - } - - #[test] - fn both_flag_sources_interstitial_blurs_non_opt_in() { - let mut c = nsfw_tweet_flag_candidate(); - c.author_features.is_nsfw_admin = true; - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Interstitial(_) - )); - } - - #[test] - fn no_flags_allows() { - let mut c = nsfw_tweet_flag_candidate(); - c.tweet_features.nsfw = NsfwFeature::default(); - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Allow - )); - } - - #[test] - fn tweet_flag_interstitial_allows_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_sensitive_opt_in(), - &nsfw_tweet_flag_candidate() - )), - VfAction::Allow - )); - } - - #[test] - fn tweet_flag_interstitial_allows_self_view() { - let mut c = nsfw_tweet_flag_candidate(); - c.author_id = 999; - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/nullcast_rule.rs b/visibility-filtering/rules/nullcast_rule.rs deleted file mode 100644 index d6ce983f..00000000 --- a/visibility-filtering/rules/nullcast_rule.rs +++ /dev/null @@ -1,107 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -pub struct NullcastedTweetDropRule; - -impl Rule for NullcastedTweetDropRule { - fn name(&self) -> &'static str { - "NullcastedTweetDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.is_nullcast() && !candidate.is_retweet() && !candidate.is_community_tweet() { - return VfAction::Drop(FilteredReason::TweetIsNullcast); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - CoreFeature, HydratedTweetCandidate, TweetFeatures, Viewer, ViewerFeatures, - }; - - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } - } - - #[test] - fn nullcast_non_retweet_drops() { - let rule = NullcastedTweetDropRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - is_nullcast: true, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn nullcast_community_tweet_allows() { - let rule = NullcastedTweetDropRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - is_nullcast: true, - is_community_tweet: true, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn nullcast_retweet_allows() { - let rule = NullcastedTweetDropRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - is_nullcast: true, - core: CoreFeature { - source_tweet_id: Some(99), - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn non_nullcast_allows() { - let rule = NullcastedTweetDropRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - is_nullcast: false, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index 64ac414a..af79ac33 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -1,26 +1,9 @@ use crate::models::{HydratedTweetCandidate, VfAction, ViewerFeatures}; -use crate::rules::nsfw_age_gating::{ - SensitiveViewerLoggedOutDropRule, SensitiveViewerNoStatedAgeDropRule, - SensitiveViewerUnderageDropRule, -}; -use crate::rules::nsfw_interstitial::{ - NsfwAuthorInterstitialRule, GORE_AND_VIOLENCE_INTERSTITIAL, NSFW_CARD_IMAGE_INTERSTITIAL, - NSFW_HIGH_PRECISION_INTERSTITIAL, -}; -use crate::rules::nullcast_rule::NullcastedTweetDropRule; -use crate::rules::socialgraph_rules::{ - DropExclusiveTweetContentRule, MutedRetweetsRule, ViewerBlocksAuthorRule, ViewerMutesAuthorRule, -}; -use crate::rules::tes_rules::{ - DropLegalTakendownPostRule, DropLocalLawsTakendownPostRule, DropStaleTweetsRule, - DropTweetsWithDmcaMediaRule, DropTweetsWithGeoRestrictedMediaRule, -}; -use crate::rules::tweet_flag_rules as tweet_flag; -use crate::rules::tweet_label_drops as tweet_label; -use crate::rules::user_label_drops as user_label; -use crate::rules::user_rules::{self as author, ProtectedAuthorDropRule}; -use crate::rules::{evaluate_rules, Rule, RuleContext, Verdict}; -use xai_visibility_filtering::models::FilteredReason; +use crate::params::NsfwGatingCountries; +use crate::rules::rule_spec::RuleSpec; +use crate::rules::{author_rules, tweet_rules}; +use crate::rules::{RuleContext, Verdict}; +use std::sync::Arc; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum SafetyLevel { @@ -39,874 +22,414 @@ impl SafetyLevel { } } -pub struct Policies { - filter_all: Vec>, - timeline_home: Vec>, - timeline_home_recommendations: Vec>, +pub(super) struct Policy { + rules: &'static [&'static [RuleSpec]], + additional_rules: &'static [&'static [RuleSpec]], } -impl Policies { - pub fn new() -> Self { +impl Policy { + pub(super) const fn new(rules: &'static [&'static [RuleSpec]]) -> Self { Self { - filter_all: vec![Box::new(FilterAllRule)], - timeline_home: timeline_home_policy(), - timeline_home_recommendations: timeline_home_recommendations_policy(), - } - } - - fn select(&self, level: SafetyLevel) -> &[Box] { - match level { - SafetyLevel::FilterAll => &self.filter_all, - SafetyLevel::TimelineHome => &self.timeline_home, - SafetyLevel::TimelineHomeRecommendations => &self.timeline_home_recommendations, + rules, + additional_rules: &[], } } - pub fn evaluate( - &self, - level: SafetyLevel, - viewer: &ViewerFeatures, - candidate: &HydratedTweetCandidate, - ) -> Verdict { - let context = RuleContext::new(level, viewer, candidate); - evaluate_rules(self.select(level), &context) + fn rules(&self) -> impl Iterator + '_ { + self.rules + .iter() + .chain(self.additional_rules) + .copied() + .flatten() } - pub fn rule_counts(&self) -> (usize, usize) { - ( - self.timeline_home.len(), - self.timeline_home_recommendations.len(), - ) - } -} + pub(super) fn evaluate(&self, context: &RuleContext<'_>) -> Verdict { + let mut action = VfAction::Allow; + let mut decided_by = None; + + for rule in self.rules() { + match rule.evaluate(context) { + VfAction::Drop(reason) => { + return Verdict { + action: VfAction::Drop(reason), + decided_by: Some(rule.name()), + }; + } + VfAction::Interstitial(reason) if matches!(action, VfAction::Allow) => { + action = VfAction::Interstitial(reason); + decided_by = Some(rule.name()); + } + VfAction::Allow | VfAction::Interstitial(_) => {} + } + } -impl Default for Policies { - fn default() -> Self { - Self::new() + Verdict { action, decided_by } } -} -struct FilterAllRule; - -impl Rule for FilterAllRule { - fn name(&self) -> &'static str { - "FilterAllRule" + fn len(&self) -> usize { + self.rules().count() } - fn evaluate(&self, _context: &RuleContext<'_>) -> VfAction { - VfAction::Drop(FilteredReason::UnspecifiedReason) + #[cfg(test)] + fn rule_names(&self) -> impl Iterator + '_ { + self.rules().map(RuleSpec::name) } } -fn base_home_rules() -> Vec> { - vec![ - Box::new(author::SUSPENDED_AUTHOR_DROP), - Box::new(author::DEACTIVATED_AUTHOR_DROP), - Box::new(author::ERASED_AUTHOR_DROP), - Box::new(author::OFFBOARDED_AUTHOR_DROP), - Box::new(ProtectedAuthorDropRule), - Box::new(ViewerBlocksAuthorRule), - Box::new(ViewerMutesAuthorRule), - Box::new(MutedRetweetsRule), - Box::new(tweet_label::PDNA_DROP), - Box::new(tweet_label::BOUNCE_DROP), - Box::new(tweet_label::SPAM_DROP), - Box::new(tweet_label::FOR_EMERGENCY_USE_ONLY_DROP), - Box::new(tweet_label::FOSNR_HATEFUL_CONDUCT_DROP), - Box::new(tweet_label::FOSNR_VIOLENT_SPEECH_DROP), - Box::new(tweet_label::FOSNR_ABUSE_DROP), - Box::new(tweet_label::FOSNR_CIVIC_INTEGRITY_DROP), - Box::new(NullcastedTweetDropRule), - Box::new(DropStaleTweetsRule), - Box::new(DropLegalTakendownPostRule), - Box::new(DropLocalLawsTakendownPostRule), - Box::new(SensitiveViewerLoggedOutDropRule), - Box::new(SensitiveViewerUnderageDropRule), - Box::new(SensitiveViewerNoStatedAgeDropRule), - Box::new(DropExclusiveTweetContentRule), - Box::new(NSFW_HIGH_PRECISION_INTERSTITIAL), - Box::new(GORE_AND_VIOLENCE_INTERSTITIAL), - Box::new(NSFW_CARD_IMAGE_INTERSTITIAL), - Box::new(NsfwAuthorInterstitialRule), - ] -} - -fn timeline_home_policy() -> Vec> { - base_home_rules() -} +static FILTER_ALL_POLICY: Policy = Policy::new(&[tweet_rules::FILTER_ALL]); + +static TIMELINE_HOME_SHARED_RULES: [&[RuleSpec]; 9] = [ + author_rules::AUTHOR_STATE_DROPS, + author_rules::SOCIALGRAPH_DROPS, + tweet_rules::TWEET_LABEL_DROPS, + tweet_rules::NULLCAST_DROP, + tweet_rules::TES_HOME_DROPS, + tweet_rules::SENSITIVE_VIEWER_DROPS, + tweet_rules::EXCLUSIVE_TWEET_DROP, + tweet_rules::NSFW_MEDIA_INTERSTITIALS, + tweet_rules::NSFW_AUTHOR_INTERSTITIAL, +]; + +static TIMELINE_HOME_RECOMMENDATION_ONLY_RULES: [&[RuleSpec]; 5] = [ + tweet_rules::RECS_MEDIA_DROPS, + author_rules::OON_NSFW_AUTHOR_DROPS, + tweet_rules::OON_TWEET_FLAG_DROPS, + tweet_rules::OON_TWEET_LABEL_DROPS, + author_rules::OON_USER_LABEL_DROPS, +]; + +static TIMELINE_HOME_POLICY: Policy = Policy::new(&TIMELINE_HOME_SHARED_RULES); +static TIMELINE_HOME_RECOMMENDATIONS_POLICY: Policy = Policy { + rules: &TIMELINE_HOME_SHARED_RULES, + additional_rules: &TIMELINE_HOME_RECOMMENDATION_ONLY_RULES, +}; -fn timeline_home_recommendations_policy() -> Vec> { - let mut rules = base_home_rules(); - let oon_drops: Vec> = vec![ - Box::new(DropTweetsWithDmcaMediaRule), - Box::new(DropTweetsWithGeoRestrictedMediaRule), - Box::new(author::NSFW_USER_AUTHOR_DROP), - Box::new(author::NSFW_ADMIN_AUTHOR_DROP), - Box::new(tweet_flag::TWEET_NSFW_USER_DROP), - Box::new(tweet_flag::TWEET_NSFW_ADMIN_DROP), - Box::new(tweet_label::NSFW_HIGH_RECALL_DROP), - Box::new(tweet_label::NSFW_HIGH_PRECISION_DROP), - Box::new(tweet_label::GORE_AND_VIOLENCE_HIGH_PRECISION_DROP), - Box::new(tweet_label::NSFW_CARD_IMAGE_DROP), - Box::new(tweet_label::DO_NOT_AMPLIFY_DROP), - Box::new(tweet_label::MALICIOUS_URL_DROP), - Box::new(tweet_label::SPAM_HIGH_RECALL_DROP), - Box::new(tweet_label::NSFW_TEXT_DROP), - Box::new(tweet_label::FOSNR_ABUSE_INSULTS_OON_DROP), - Box::new(user_label::NSFW_HIGH_RECALL_USER_DROP), - Box::new(user_label::NSFW_HIGH_PRECISION_USER_DROP), - Box::new(user_label::SPAM_HIGH_RECALL_USER_DROP), - Box::new(user_label::COMPROMISED_USER_DROP), - Box::new(user_label::READ_ONLY_USER_DROP), - Box::new(user_label::IMPERSONATION_HIGH_PRECISION_USER_DROP), - Box::new(user_label::NSFW_AVATAR_IMAGE_USER_DROP), - Box::new(user_label::NSFW_BANNER_IMAGE_USER_DROP), - Box::new(user_label::ABUSIVE_HIGH_RECALL_USER_DROP), - Box::new(user_label::NSFW_NEAR_PERFECT_USER_DROP), - Box::new(user_label::DO_NOT_AMPLIFY_NON_FOLLOWER_USER_DROP), - ]; - rules.extend(oon_drops); - rules +pub struct RuleEngine { + nsfw_gating_countries: Arc, } -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - HydratedTweetCandidate, MediaFeature, TweetFeatures, Viewer, ViewerFeatures, - }; - use std::collections::HashMap; - - struct RecommendationsOnlyRule; - - impl Rule for RecommendationsOnlyRule { - fn name(&self) -> &'static str { - "RecommendationsOnlyRule" - } +impl RuleEngine { + #[cfg(test)] + pub(crate) fn for_tests() -> Self { + Self::with_nsfw_gating_countries(Arc::new(NsfwGatingCountries::starting_at_default())) + } - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - match context.safety_level() { - SafetyLevel::TimelineHomeRecommendations => { - VfAction::Drop(FilteredReason::UnspecifiedReason) - } - SafetyLevel::FilterAll | SafetyLevel::TimelineHome => VfAction::Allow, - } + pub fn with_nsfw_gating_countries(gating_countries: Arc) -> Self { + Self { + nsfw_gating_countries: gating_countries, } } - #[test] - fn policies_evaluate_uses_selected_safety_level_in_context() { - let policies = Policies { - filter_all: vec![Box::new(RecommendationsOnlyRule)], - timeline_home: vec![Box::new(RecommendationsOnlyRule)], - timeline_home_recommendations: vec![Box::new(RecommendationsOnlyRule)], - }; - let viewer = ViewerFeatures::default(); - let candidate = HydratedTweetCandidate::default(); - - assert!(matches!( - policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action, - VfAction::Allow - )); - assert!(matches!( - policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate - ) - .action, - VfAction::Drop(_) - )); + fn select(level: SafetyLevel) -> &'static Policy { + match level { + SafetyLevel::FilterAll => &FILTER_ALL_POLICY, + SafetyLevel::TimelineHome => &TIMELINE_HOME_POLICY, + SafetyLevel::TimelineHomeRecommendations => &TIMELINE_HOME_RECOMMENDATIONS_POLICY, + } } - #[test] - fn filter_all_rule_drops_even_self_view() { - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - }; - assert!(matches!( - FilterAllRule.evaluate(&crate::rules::test_context(&viewer, &candidate)), - VfAction::Drop(_) - )); + pub fn evaluate( + &self, + level: SafetyLevel, + viewer: &ViewerFeatures, + candidate: &HydratedTweetCandidate, + ) -> Verdict { + let context = RuleContext::new(viewer, candidate, &self.nsfw_gating_countries); + Self::select(level).evaluate(&context) } - #[test] - fn filter_all_policy_drops_pristine_candidate() { - let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; - let verdict = policies.evaluate( - SafetyLevel::FilterAll, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(verdict.action, VfAction::Drop(_))); - - let verdict = policies.evaluate( - SafetyLevel::TimelineHome, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(verdict.action, VfAction::Allow)); + #[cfg(test)] + pub(crate) fn wired_rule_names(&self, level: SafetyLevel) -> Vec<&'static str> { + Self::select(level).rule_names().collect() } - #[test] - fn dmca_media_drops_recommendations_only() { - let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - media: MediaFeature { - has_dmca_media: true, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - - let timeline_home = policies.evaluate( - SafetyLevel::TimelineHome, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(timeline_home.action, VfAction::Allow)); - - let recommendations = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(recommendations.action, VfAction::Drop(_))); + pub fn rule_counts(&self) -> (usize, usize) { + ( + TIMELINE_HOME_POLICY.len(), + TIMELINE_HOME_RECOMMENDATIONS_POLICY.len(), + ) } +} - #[test] - fn tweet_nsfw_flag_drops_recommendations_only() { - use crate::models::NsfwFeature; - let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - nsfw: NsfwFeature { - user: true, - admin: false, - }, - ..Default::default() - }, - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - let timeline_home = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(timeline_home, VfAction::Allow), - "in-network tweet nsfw_user flag should allow, got {timeline_home:?}" - ); - - let recommendations = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!(matches!(recommendations.action, VfAction::Drop(_))); - assert_eq!(recommendations.decided_by, Some("TweetNsfwUserDropRule")); - } +#[cfg(test)] +mod tests { + use super::*; + use crate::models::{VfAction, ViewerFeatures}; + use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; #[test] - fn nsfw_author_interstitials_in_network_but_drops_oon() { - use crate::models::{AuthorFeatures, TweetFeatures}; - let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }, - tweet_features: TweetFeatures { - media: MediaFeature { - has_media: true, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; + fn refreshed_config_country_reaches_the_wired_rule() { + let gating_countries = Arc::new(NsfwGatingCountries::starting_at_default()); + let rule_engine = RuleEngine::with_nsfw_gating_countries(Arc::clone(&gating_countries)); + let candidate = candidate() + .with_label(crate::models::SafetyLabelType::NSFW_HIGH_PRECISION) + .with_media() + .build(); let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: false, - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Interstitial(_)), - "in-network NSFW author should interstitial, got {in_network:?}" - ); - - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, + viewer_age: crate::models::ViewerAge::NotStated, + country_code: Some("us".into()), + ..viewer(VIEWER_ID) + }; + + let verdict = rule_engine.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); + assert!(!matches!(verdict.action, VfAction::Drop(_))); + + gating_countries.refresh_from( + &xai_feature_switches::FeatureSwitches::load_string( + r#" +rust_vf: + parameters: + rust_vf_nsfw_gating_countries: + type: array + default: + - "us" +"#, ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON NSFW author should drop, got {oon:?}" - ); - } - - #[test] - fn egregious_nsfw_does_not_drop() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - - let mut labels = HashMap::new(); - labels.insert(SafetyLabelType::EGREGIOUS_NSFW, SafetyLabel::default()); - let tweet_candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - }; - let user_candidate = candidate_with_author_user_label(LabelValue::EGREGIOUS_NSFW, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - for candidate in [&tweet_candidate, &user_candidate] { - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network EgregiousNsfw should allow after rule removal, got {in_network:?}" - ); - let oon = policies - .evaluate(SafetyLevel::TimelineHomeRecommendations, &viewer, candidate) - .action; - assert!( - matches!(oon, VfAction::Allow), - "OON EgregiousNsfw should allow after rule removal, got {oon:?}" - ); - } - } - - #[test] - fn spam_high_recall_drops_oon_but_allows_in_network() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; - let policies = Policies::new(); - let mut labels = HashMap::new(); - labels.insert(SafetyLabelType::SPAM_HIGH_RECALL, SafetyLabel::default()); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network SpamHighRecall should allow (Scala drops it OON only), got {in_network:?}" + .unwrap(), ); - - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON SpamHighRecall should drop, got {oon:?}" + let verdict = rule_engine.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); + assert!(matches!(verdict.action, VfAction::Drop(_))); + assert_eq!( + verdict.decided_by, + Some("SensitiveViewerNoStatedAgeDropRule") ); } #[test] - fn malicious_url_drops_oon_but_allows_in_network() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; - let policies = Policies::new(); - let mut labels = HashMap::new(); - labels.insert(SafetyLabelType::MALICIOUS_URL, SafetyLabel::default()); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network MaliciousUrl should allow (Scala drops it OON only), got {in_network:?}" - ); - - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON MaliciousUrl should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("MaliciousUrlOonDropRule")); - - let author = ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - }; - let oon_author = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &author, - &candidate, - ) - .action; - assert!( - matches!(oon_author, VfAction::Allow), - "OON MaliciousUrl should allow author, got {oon_author:?}" + fn wired_rule_order_matches_pre_migration_sequence() { + let rule_engine = RuleEngine::for_tests(); + assert_eq!( + rule_engine.wired_rule_names(SafetyLevel::FilterAll), + vec!["FilterAllRule"] + ); + let home = rule_engine.wired_rule_names(SafetyLevel::TimelineHome); + assert_eq!( + home, + vec![ + "SuspendedAuthorRule", + "DeactivatedAuthorRule", + "ErasedAuthorRule", + "OffboardedAuthorRule", + "ProtectedAuthorDropRule", + "ViewerBlocksAuthorRule", + "ViewerMutesAuthorRule", + "MutedRetweetsRule", + "PdnaTweetLabelRule", + "BounceTweetLabelRule", + "SpamTweetLabelRule", + "ForEmergencyUseOnlyDropRule", + "FosnrHatefulConductDropRule", + "FosnrViolentSpeechDropRule", + "FosnrAbuseDropRule", + "FosnrCivicIntegrityDropRule", + "NullcastedTweetDropRule", + "DropStaleTweetsRule", + "DropLegalTakendownPostRule", + "DropLocalLawsTakendownPostRule", + "SensitiveViewerLoggedOutDropRule", + "SensitiveViewerUnderageDropRule", + "SensitiveViewerNoStatedAgeDropRule", + "DropExclusiveTweetContentRule", + "NsfwHighPrecisionInterstitialRule", + "GoreAndViolenceInterstitialRule", + "NsfwCardImageInterstitialRule", + "NsfwAuthorInterstitialRule", + ] + ); + let mut recs = home.clone(); + recs.extend([ + "DropTweetsWithDmcaMediaRule", + "DropTweetsWithGeoRestrictedMediaRule", + "DropNsfwUserAuthorRule", + "DropNsfwAdminAuthorRule", + "TweetNsfwUserDropRule", + "TweetNsfwAdminDropRule", + "NsfwHighRecallDropRule", + "NsfwHighPrecisionOonDropRule", + "GoreAndViolenceOonDropRule", + "NsfwCardImageOonDropRule", + "DoNotAmplifyOonDropRule", + "MaliciousUrlOonDropRule", + "SpamHighRecallDropRule", + "NsfwTextTweetLabelDropRule", + "FosnrAbuseInsultsOonDropRule", + "NsfwHighRecallUserLabelRule", + "NsfwHighPrecisionUserLabelRule", + "SpamHighRecallUserLabelRule", + "CompromisedUserLabelRule", + "ReadOnlyUserLabelRule", + "ImpersonationHighPrecisionUserLabelRule", + "NsfwAvatarImageRule", + "NsfwBannerImageRule", + "AbusiveHighRecallRule", + "NsfwNearPerfectAuthorRule", + "DoNotAmplifyNonFollowerRule", + ]); + assert_eq!( + rule_engine.wired_rule_names(SafetyLevel::TimelineHomeRecommendations), + recs ); } - fn fosnr_candidate( - label: crate::models::SafetyLabelType, - follows: bool, - ) -> HydratedTweetCandidate { - use crate::models::{SafetyLabel, SafetyLabelMap}; - let mut labels = HashMap::new(); - labels.insert(label, SafetyLabel::default()); - let mut c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - }; - c.relationship.viewer_follows_author = follows; - c + #[derive(Debug, PartialEq, Eq)] + enum RowClass { + Drop, + Interstitial, } - #[test] - fn fosnr_labels_drop_non_author_non_follower_on_both_surfaces() { - use crate::models::SafetyLabelType; - let policies = Policies::new(); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - for label in [ - SafetyLabelType::FOSNR_HATEFUL_CONDUCT, - SafetyLabelType::FOSNR_VIOLENT_SPEECH, - SafetyLabelType::FOSNR_ABUSE, - SafetyLabelType::FOSNR_CIVIC_INTEGRITY, - ] { - let candidate = fosnr_candidate(label, false); - for level in [ - SafetyLevel::TimelineHome, - SafetyLevel::TimelineHomeRecommendations, - ] { - let action = policies.evaluate(level, &viewer, &candidate).action; - assert!( - matches!(action, VfAction::Drop(_)), - "{label:?} on {level:?} should drop non-follower, got {action:?}" - ); + fn row_class(spec: &RuleSpec) -> RowClass { + use crate::rules::rule_spec::RuleAction; + match spec { + RuleSpec::Tweet { + action: RuleAction::Drop(_), + .. } + | RuleSpec::Author { .. } + | RuleSpec::Custom { .. } => RowClass::Drop, + RuleSpec::Tweet { + action: RuleAction::SensitiveMediaInterstitial(_), + .. + } => RowClass::Interstitial, } } - #[test] - fn fosnr_never_drops_author() { - use crate::models::SafetyLabelType; - let policies = Policies::new(); - let author = ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - }; - for label in [ - SafetyLabelType::FOSNR_HATEFUL_CONDUCT, - SafetyLabelType::FOSNR_VIOLENT_SPEECH, - SafetyLabelType::FOSNR_ABUSE, - SafetyLabelType::FOSNR_CIVIC_INTEGRITY, - SafetyLabelType::FOSNR_ABUSE_INSULTS, - ] { - let candidate = fosnr_candidate(label, false); - for level in [ - SafetyLevel::TimelineHome, - SafetyLevel::TimelineHomeRecommendations, - ] { - let action = policies.evaluate(level, &author, &candidate).action; - assert!( - matches!(action, VfAction::Allow), - "{label:?} on {level:?} should allow author, got {action:?}" - ); + fn assert_drops_precede_interstitials(policy: &Policy, name: &str) { + let mut seen_interstitial = false; + for spec in policy.rules() { + match row_class(spec) { + RowClass::Interstitial => seen_interstitial = true, + RowClass::Drop if seen_interstitial => { + panic!("{name}: drop {} follows an interstitial row", spec.name()); + } + RowClass::Drop => {} } } } #[test] - fn fosnr_abuse_insults_drops_oon_but_allows_in_network() { - use crate::models::SafetyLabelType; - let policies = Policies::new(); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - let author = ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - }; - - for follows in [true, false] { - let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, follows); - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network FosnrAbuseInsults should allow (follows={follows}), got {in_network:?}" + fn drops_precede_interstitials_and_recommendation_only_rules_are_drops() { + assert_drops_precede_interstitials(&FILTER_ALL_POLICY, "FilterAll"); + assert_drops_precede_interstitials(&TIMELINE_HOME_POLICY, "TimelineHome"); + for spec in TIMELINE_HOME_RECOMMENDATION_ONLY_RULES + .iter() + .copied() + .flatten() + { + assert_eq!( + row_class(spec), + RowClass::Drop, + "recommendation-only row {} must be a drop", + spec.name() ); } - - let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, false); - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON FosnrAbuseInsults should drop non-author, got {oon:?}" - ); - - let oon_author = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &author, - &candidate, - ) - .action; - assert!( - matches!(oon_author, VfAction::Allow), - "OON FosnrAbuseInsults should allow author, got {oon_author:?}" - ); - } - - #[test] - fn geo_restricted_media_drops_oon_but_allows_in_network() { - use crate::models::TweetFeatures; - let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - media: MediaFeature { - geo_deny_list: vec!["de".to_string()], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - country_code: Some("de".to_string()), - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network geo-restricted media should allow (Scala wires the rule in THR only), got {in_network:?}" - ); - - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON geo-restricted media should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("DropTweetsWithGeoRestrictedMediaRule")); } - #[test] - fn nsfw_text_drops_oon_but_allows_in_network() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; - let policies = Policies::new(); - let mut labels = HashMap::new(); - labels.insert(SafetyLabelType::NSFW_TEXT, SafetyLabel::default()); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network NsfwText should allow (Scala drops it OON only), got {in_network:?}" - ); + mod engine { + use super::super::*; + use crate::models::{HydratedTweetCandidate, VfAction, ViewerFeatures}; + use crate::rules::test_context; + use xai_visibility_filtering::models::FilteredReason; + + const fn custom( + name: &'static str, + evaluate: fn(&RuleContext<'_>) -> VfAction, + ) -> RuleSpec { + RuleSpec::Custom { name, evaluate } + } - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON NsfwText should drop, got {oon:?}" - ); - } + fn allow(_: &RuleContext<'_>) -> VfAction { + VfAction::Allow + } - fn candidate_with_author_user_label( - label: xai_x_thrift::user_labels::LabelValue, - follows: bool, - ) -> HydratedTweetCandidate { - use crate::models::{AuthorFeatures, UserLabelSet}; - use std::collections::HashSet; - let mut c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { - user_labels: UserLabelSet::new(HashSet::from([label])), - ..Default::default() - }, - ..Default::default() - }; - c.relationship.viewer_follows_author = follows; - c - } + fn drop_suspended(_: &RuleContext<'_>) -> VfAction { + VfAction::Drop(FilteredReason::AuthorIsSuspended) + } - #[test] - fn nsfw_avatar_user_label_drops_oon_but_allows_in_network() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let candidate = candidate_with_author_user_label(LabelValue::NSFW_AVATAR_IMAGE, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + fn interstitial_nsfw(_: &RuleContext<'_>) -> VfAction { + VfAction::Interstitial(FilteredReason::ContainNsfwMedia) + } - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network NsfwAvatarImage should allow, got {in_network:?}" - ); + fn interstitial_unspecified(_: &RuleContext<'_>) -> VfAction { + VfAction::Interstitial(FilteredReason::UnspecifiedReason) + } - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON NsfwAvatarImage should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("NsfwAvatarImageRule")); - } + fn unreachable_after_drop(_: &RuleContext<'_>) -> VfAction { + panic!("a rule after a Drop must never be evaluated"); + } - #[test] - fn recommendations_blacklist_does_not_drop() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let candidate = - candidate_with_author_user_label(LabelValue::RECOMMENDATIONS_BLACKLIST, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + fn context_inputs() -> (ViewerFeatures, HydratedTweetCandidate) { + (ViewerFeatures::default(), HydratedTweetCandidate::default()) + } - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network RecommendationsBlacklist should allow, got {in_network:?}" - ); + static SHORT_CIRCUIT_ROWS: [RuleSpec; 3] = [ + custom("allow", allow), + custom("drop", drop_suspended), + custom("after_drop", unreachable_after_drop), + ]; + static SHORT_CIRCUIT: Policy = Policy::new(&[&SHORT_CIRCUIT_ROWS]); + + static INTERSTITIAL_ROWS: [RuleSpec; 2] = [ + custom("first_interstitial", interstitial_nsfw), + custom("second_interstitial", interstitial_unspecified), + ]; + static INTERSTITIALS: Policy = Policy::new(&[&INTERSTITIAL_ROWS]); + + static DROP_AFTER_INTERSTITIAL_ROWS: [RuleSpec; 2] = [ + custom("interstitial", interstitial_nsfw), + custom("drop", drop_suspended), + ]; + static DROP_AFTER_INTERSTITIAL: Policy = Policy::new(&[&DROP_AFTER_INTERSTITIAL_ROWS]); + + static ALL_ALLOWS_ROWS: [RuleSpec; 2] = [custom("a", allow), custom("b", allow)]; + static ALL_ALLOWS: Policy = Policy::new(&[&ALL_ALLOWS_ROWS]); + + #[test] + fn drop_short_circuits_later_rules() { + let (viewer, candidate) = context_inputs(); + + let verdict = SHORT_CIRCUIT.evaluate(&test_context(&viewer, &candidate)); + + assert!(matches!( + verdict.action, + VfAction::Drop(FilteredReason::AuthorIsSuspended) + )); + assert_eq!(verdict.decided_by, Some("drop")); + } - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Allow), - "OON RecommendationsBlacklist should allow after rule removal, got {oon:?}" - ); - } + #[test] + fn first_interstitial_sticks_without_short_circuit() { + let (viewer, candidate) = context_inputs(); - #[test] - fn abusive_high_recall_allows_follower_on_both_surfaces() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let candidate = candidate_with_author_user_label(LabelValue::ABUSIVE_HIGH_RECALL, true); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + let verdict = INTERSTITIALS.evaluate(&test_context(&viewer, &candidate)); - for level in [ - SafetyLevel::TimelineHome, - SafetyLevel::TimelineHomeRecommendations, - ] { - let action = policies.evaluate(level, &viewer, &candidate).action; - assert!( - matches!(action, VfAction::Allow), - "AbusiveHighRecall follower on {level:?} should allow, got {action:?}" - ); + assert!(matches!( + verdict.action, + VfAction::Interstitial(FilteredReason::ContainNsfwMedia) + )); + assert_eq!(verdict.decided_by, Some("first_interstitial")); } - } - #[test] - fn abusive_high_recall_drops_oon_non_follower_but_allows_in_network() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let candidate = candidate_with_author_user_label(LabelValue::ABUSIVE_HIGH_RECALL, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; - - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network AbusiveHighRecall should allow, got {in_network:?}" - ); + #[test] + fn drop_after_interstitial_wins() { + let (viewer, candidate) = context_inputs(); - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON AbusiveHighRecall non-follower should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("AbusiveHighRecallRule")); - } - - #[test] - fn nsfw_near_perfect_drops_oon_but_allows_in_network() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let candidate = candidate_with_author_user_label(LabelValue::NSFW_NEAR_PERFECT, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + let verdict = DROP_AFTER_INTERSTITIAL.evaluate(&test_context(&viewer, &candidate)); - let in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network NsfwNearPerfect should allow, got {in_network:?}" - ); + assert!(matches!( + verdict.action, + VfAction::Drop(FilteredReason::AuthorIsSuspended) + )); + assert_eq!(verdict.decided_by, Some("drop")); + } - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON NsfwNearPerfect should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("NsfwNearPerfectAuthorRule")); - } + #[test] + fn all_allows_is_allow_with_no_decider() { + let (viewer, candidate) = context_inputs(); - #[test] - fn do_not_amplify_drops_oon_non_follower_but_allows_follower() { - use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + let verdict = ALL_ALLOWS.evaluate(&test_context(&viewer, &candidate)); - let non_follower = candidate_with_author_user_label(LabelValue::DO_NOT_AMPLIFY, false); - let oon = policies.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &non_follower, - ); - assert!( - matches!(oon.action, VfAction::Drop(_)), - "OON DoNotAmplify non-follower should drop, got {:?}", - oon.action - ); - assert_eq!(oon.decided_by, Some("DoNotAmplifyNonFollowerRule")); - - let follower = candidate_with_author_user_label(LabelValue::DO_NOT_AMPLIFY, true); - let action = policies - .evaluate(SafetyLevel::TimelineHomeRecommendations, &viewer, &follower) - .action; - assert!( - matches!(action, VfAction::Allow), - "OON DoNotAmplify follower should allow, got {action:?}" - ); + assert!(matches!(verdict.action, VfAction::Allow)); + assert_eq!(verdict.decided_by, None); + } } } diff --git a/visibility-filtering/rules/rule_spec.rs b/visibility-filtering/rules/rule_spec.rs new file mode 100644 index 00000000..6bac699b --- /dev/null +++ b/visibility-filtering/rules/rule_spec.rs @@ -0,0 +1,87 @@ +use crate::models::VfAction; +use crate::rules::context::{AuthorPredicates, TweetPredicates}; +use crate::rules::RuleContext; +use xai_visibility_filtering::models::FilteredReason; + +pub(super) enum RuleSpec { + Tweet { + name: &'static str, + when: fn(TweetPredicates<'_>) -> bool, + action: RuleAction, + exempt_author: bool, + }, + Author { + name: &'static str, + when: fn(AuthorPredicates<'_>) -> bool, + reason: FilteredReason, + exempt_follower: bool, + }, + Custom { + name: &'static str, + evaluate: fn(&RuleContext<'_>) -> VfAction, + }, +} + +pub(super) enum RuleAction { + Drop(FilteredReason), + SensitiveMediaInterstitial(FilteredReason), +} + +impl RuleSpec { + pub(super) fn name(&self) -> &'static str { + match self { + RuleSpec::Tweet { name, .. } + | RuleSpec::Author { name, .. } + | RuleSpec::Custom { name, .. } => name, + } + } + + pub(super) fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { + match self { + RuleSpec::Tweet { + when, + action, + exempt_author, + .. + } => { + if !when(context.tweet()) { + return VfAction::Allow; + } + if *exempt_author && context.viewer().is_author() { + return VfAction::Allow; + } + match action { + RuleAction::Drop(reason) => VfAction::Drop(reason.clone()), + RuleAction::SensitiveMediaInterstitial(reason) => { + if context.viewer().allows_sensitive_media() { + VfAction::Allow + } else { + VfAction::Interstitial(reason.clone()) + } + } + } + } + RuleSpec::Author { + when, + reason, + exempt_follower, + .. + } => { + if !when(context.author()) { + return VfAction::Allow; + } + if context.viewer().is_author() { + return VfAction::Allow; + } + if *exempt_follower + && !context.viewer().is_logged_out() + && context.viewer().follows_author() + { + return VfAction::Allow; + } + VfAction::Drop(reason.clone()) + } + RuleSpec::Custom { evaluate, .. } => evaluate(context), + } + } +} diff --git a/visibility-filtering/rules/socialgraph_rules.rs b/visibility-filtering/rules/socialgraph_rules.rs deleted file mode 100644 index 691c13c7..00000000 --- a/visibility-filtering/rules/socialgraph_rules.rs +++ /dev/null @@ -1,337 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -pub struct ViewerBlocksAuthorRule; - -impl Rule for ViewerBlocksAuthorRule { - fn name(&self) -> &'static str { - "ViewerBlocksAuthorRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { - return VfAction::Allow; - } - if candidate.relationship.viewer_blocks_author { - return VfAction::Drop(FilteredReason::AuthorBlockViewer); - } - VfAction::Allow - } -} - -pub struct MutedRetweetsRule; - -impl Rule for MutedRetweetsRule { - fn name(&self) -> &'static str { - "MutedRetweetsRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { - return VfAction::Allow; - } - if candidate.is_retweet() && candidate.relationship.viewer_mutes_retweets_from_author { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -pub struct ViewerMutesAuthorRule; - -impl Rule for ViewerMutesAuthorRule { - fn name(&self) -> &'static str { - "ViewerMutesAuthorRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { - return VfAction::Allow; - } - if candidate.relationship.viewer_mutes_author { - return VfAction::Drop(FilteredReason::ViewerMutesAuthor); - } - VfAction::Allow - } -} - -pub struct DropExclusiveTweetContentRule; - -impl Rule for DropExclusiveTweetContentRule { - fn name(&self) -> &'static str { - "DropExclusiveTweetContentRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - let Some(exc) = &candidate.exclusive_content else { - return VfAction::Allow; - }; - - let Some(viewer_id) = viewer.viewer_id() else { - return VfAction::Drop(FilteredReason::ExclusiveTweet); - }; - - if viewer_id == exc.conversation_author_id { - return VfAction::Allow; - } - - if exc.viewer_super_follows_author { - return VfAction::Allow; - } - - if !candidate.is_retweet() && context.is_author_viewer() { - return VfAction::Allow; - } - - VfAction::Drop(FilteredReason::ExclusiveTweet) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - ExclusiveContentFeatures, HydratedTweetCandidate, TweetFeatures, Viewer, - ViewerAuthorRelationship, ViewerFeatures, - }; - - fn viewer(id: u64) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(id), - ..Default::default() - } - } - - fn logged_out_viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedOut, - ..Default::default() - } - } - - fn exclusive_candidate( - tweet_id: u64, - author_id: u64, - root_author_id: u64, - ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id, - author_id, - exclusive_content: Some(ExclusiveContentFeatures { - conversation_author_id: root_author_id, - viewer_super_follows_author: false, - }), - tweet_features: TweetFeatures::default(), - ..Default::default() - } - } - - fn non_exclusive_candidate() -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures::default(), - ..Default::default() - } - } - - fn candidate_with_relationship( - relationship: ViewerAuthorRelationship, - ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - relationship, - ..Default::default() - } - } - - #[test] - fn viewer_blocks_author_drops() { - let c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_blocks_author: true, - ..Default::default() - }); - assert!(matches!( - ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(FilteredReason::AuthorBlockViewer) - )); - } - - #[test] - fn viewer_does_not_block_author_allows() { - let c = candidate_with_relationship(ViewerAuthorRelationship::default()); - assert!(matches!( - ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn blocks_rule_allows_logged_out_viewer() { - let c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_blocks_author: true, - ..Default::default() - }); - assert!(matches!( - ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn viewer_mutes_author_drops() { - let c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_mutes_author: true, - ..Default::default() - }); - assert!(matches!( - ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(FilteredReason::ViewerMutesAuthor) - )); - } - - #[test] - fn viewer_does_not_mute_author_allows() { - let c = candidate_with_relationship(ViewerAuthorRelationship::default()); - assert!(matches!( - ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn mutes_rule_allows_logged_out_viewer() { - let c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_mutes_author: true, - ..Default::default() - }); - assert!(matches!( - ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn muted_retweets_drops_retweet_from_muting_viewer() { - let mut c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_mutes_retweets_from_author: true, - ..Default::default() - }); - c.tweet_features.core.source_tweet_id = Some(99); - assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - } - - #[test] - fn muted_retweets_allows_non_retweet() { - let c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_mutes_retweets_from_author: true, - ..Default::default() - }); - assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn muted_retweets_allows_logged_out_viewer() { - let mut c = candidate_with_relationship(ViewerAuthorRelationship { - viewer_mutes_retweets_from_author: true, - ..Default::default() - }); - c.tweet_features.core.source_tweet_id = Some(99); - assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn non_exclusive_tweet_is_allowed() { - let rule = DropExclusiveTweetContentRule; - let action = rule.evaluate(&crate::rules::test_context( - &viewer(999), - &non_exclusive_candidate(), - )); - assert!(matches!(action, VfAction::Allow)); - } - - #[test] - fn logged_out_viewer_drops_exclusive() { - let rule = DropExclusiveTweetContentRule; - let candidate = exclusive_candidate(1, 100, 100); - let action = rule.evaluate(&crate::rules::test_context( - &logged_out_viewer(), - &candidate, - )); - assert!(matches!( - action, - VfAction::Drop(FilteredReason::ExclusiveTweet) - )); - } - - #[test] - fn root_author_can_see_own_exclusive() { - let rule = DropExclusiveTweetContentRule; - let candidate = exclusive_candidate(1, 100, 100); - let action = rule.evaluate(&crate::rules::test_context(&viewer(100), &candidate)); - assert!(matches!(action, VfAction::Allow)); - } - - #[test] - fn super_follower_can_see_exclusive() { - let rule = DropExclusiveTweetContentRule; - let mut candidate = exclusive_candidate(1, 100, 100); - candidate - .exclusive_content - .as_mut() - .unwrap() - .viewer_super_follows_author = true; - let action = rule.evaluate(&crate::rules::test_context(&viewer(200), &candidate)); - assert!(matches!(action, VfAction::Allow)); - } - - #[test] - fn non_super_follower_drops_exclusive() { - let rule = DropExclusiveTweetContentRule; - let candidate = exclusive_candidate(1, 100, 100); - let action = rule.evaluate(&crate::rules::test_context(&viewer(200), &candidate)); - assert!(matches!( - action, - VfAction::Drop(FilteredReason::ExclusiveTweet) - )); - } - - #[test] - fn reply_author_can_see_own_reply_in_exclusive_convo() { - let rule = DropExclusiveTweetContentRule; - let candidate = exclusive_candidate(2, 200, 100); - let action = rule.evaluate(&crate::rules::test_context(&viewer(200), &candidate)); - assert!(matches!(action, VfAction::Allow)); - } - - #[test] - fn retweet_author_cannot_self_view_exclusive() { - let rule = DropExclusiveTweetContentRule; - let mut candidate = exclusive_candidate(2, 200, 100); - candidate.tweet_features.core.source_tweet_id = Some(99); - let action = rule.evaluate(&crate::rules::test_context(&viewer(200), &candidate)); - assert!(matches!( - action, - VfAction::Drop(FilteredReason::ExclusiveTweet) - )); - } -} diff --git a/visibility-filtering/rules/tes_rules.rs b/visibility-filtering/rules/tes_rules.rs deleted file mode 100644 index 6f90983a..00000000 --- a/visibility-filtering/rules/tes_rules.rs +++ /dev/null @@ -1,614 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_core_entities::entities::TakedownReason; -use xai_visibility_filtering::models::FilteredReason; - -pub struct DropStaleTweetsRule; - -impl Rule for DropStaleTweetsRule { - fn name(&self) -> &'static str { - "DropStaleTweetsRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if !candidate.is_stale_tweet() { - return VfAction::Allow; - } - if candidate.is_retweet() { - return VfAction::Allow; - } - if candidate.tweet_features.core.source_tweet_id.is_some() { - return VfAction::Allow; - } - VfAction::Drop(FilteredReason::UnspecifiedReason) - } -} - -pub struct DropLegalTakendownPostRule; - -impl Rule for DropLegalTakendownPostRule { - fn name(&self) -> &'static str { - "DropLegalTakendownPostRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if viewer_in_withheld_country(context, legal_takedown_country) { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -pub struct DropLocalLawsTakendownPostRule; - -impl Rule for DropLocalLawsTakendownPostRule { - fn name(&self) -> &'static str { - "DropLocalLawsTakendownPostRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if viewer_in_withheld_country(context, local_laws_takedown_country) { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -fn legal_takedown_country(reason: &TakedownReason) -> Option<&str> { - match reason { - TakedownReason::LegalRequest { country_code } - | TakedownReason::UnspecifiedReason { country_code } => Some(country_code), - _ => None, - } -} - -fn local_laws_takedown_country(reason: &TakedownReason) -> Option<&str> { - match reason { - TakedownReason::BystanderReport { country_code } => Some(country_code), - _ => None, - } -} - -fn viewer_in_withheld_country( - context: &RuleContext<'_>, - extractor: impl Fn(&TakedownReason) -> Option<&str>, -) -> bool { - let viewer = context.viewer(); - let candidate = context.candidate(); - if context.is_author_viewer() { - return false; - } - let Some(viewer_country) = &viewer.country_code else { - return false; - }; - candidate - .tweet_features - .takedown - .reasons - .iter() - .filter_map(extractor) - .any(|c| c.eq_ignore_ascii_case(viewer_country)) -} - -pub struct DropTweetsWithGeoRestrictedMediaRule; - -const WORLDWIDE_COUNTRY_CODE: &str = "xx"; - -impl Rule for DropTweetsWithGeoRestrictedMediaRule { - fn name(&self) -> &'static str { - "DropTweetsWithGeoRestrictedMediaRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - let country = viewer - .country_code - .as_deref() - .unwrap_or(WORLDWIDE_COUNTRY_CODE); - let allow = &candidate.tweet_features.media.geo_allow_list; - let deny = &candidate.tweet_features.media.geo_deny_list; - if (!allow.is_empty() && !allow.iter().any(|c| c.eq_ignore_ascii_case(country))) - || deny.iter().any(|c| c.eq_ignore_ascii_case(country)) - { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -pub struct DropTweetsWithDmcaMediaRule; - -impl Rule for DropTweetsWithDmcaMediaRule { - fn name(&self) -> &'static str { - "DropTweetsWithDmcaMediaRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.has_dmca_media() { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - CoreFeature, HydratedTweetCandidate, MediaFeature, TakedownFeature, TweetFeatures, Viewer, - ViewerFeatures, - }; - use xai_core_entities::entities::{EditControl, EditControlInitial}; - - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } - } - - fn stale_edit_control() -> Option { - Some(EditControl::Initial(EditControlInitial { - edit_tweet_ids: vec![1, 2], - ..Default::default() - })) - } - - #[test] - fn stale_edit_drops() { - let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - edit_control: stale_edit_control(), - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn non_stale_tweet_allows() { - let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn stale_retweet_allows() { - let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - edit_control: stale_edit_control(), - core: CoreFeature { - source_tweet_id: Some(99), - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - fn viewer_with_country(country: &str) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - country_code: Some(country.to_string()), - ..Default::default() - } - } - - #[test] - fn takedown_drops_in_matching_country() { - let rule = DropLegalTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![ - TakedownReason::LegalRequest { - country_code: "de".to_string(), - }, - TakedownReason::UnspecifiedReason { - country_code: "fr".to_string(), - }, - ], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn takedown_allows_in_non_matching_country() { - let rule = DropLegalTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), - VfAction::Allow - )); - } - - #[test] - fn takedown_allows_when_no_viewer_country() { - let rule = DropLegalTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn legal_rule_ignores_local_laws_countries() { - let rule = DropLegalTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - } - - #[test] - fn local_laws_drops_in_matching_country() { - let rule = DropLocalLawsTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![ - TakedownReason::BystanderReport { - country_code: "de".to_string(), - }, - TakedownReason::BystanderReport { - country_code: "fr".to_string(), - }, - ], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("fr"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn local_laws_allows_in_non_matching_country() { - let rule = DropLocalLawsTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), - VfAction::Allow - )); - } - - #[test] - fn local_laws_ignores_legal_countries() { - let rule = DropLocalLawsTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - } - - #[test] - fn legal_allows_author_viewing_own_withheld_post() { - let rule = DropLegalTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 999, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - } - - #[test] - fn local_laws_allows_author_viewing_own_withheld_post() { - let rule = DropLocalLawsTakendownPostRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 999, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - } - - #[test] - fn takedown_rules_ignore_non_country_reasons() { - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - takedown: TakedownFeature { - reasons: vec![ - TakedownReason::Dmca, - TakedownReason::HatefulImagery, - TakedownReason::Unknown, - ], - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - DropLegalTakendownPostRule - .evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - assert!(matches!( - DropLocalLawsTakendownPostRule - .evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Allow - )); - } - - fn geo_candidate(allow: &[&str], deny: &[&str]) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - media: MediaFeature { - geo_allow_list: allow.iter().map(|s| s.to_string()).collect(), - geo_deny_list: deny.iter().map(|s| s.to_string()).collect(), - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - } - } - - #[test] - fn geo_restricted_no_restrictions_allows() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&[], &[]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), - VfAction::Allow - )); - } - - #[test] - fn geo_restricted_denylist_drops_matching_country() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&[], &["de", "fr"]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_denylist_allows_other_country() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&[], &["de", "fr"]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), - VfAction::Allow - )); - } - - #[test] - fn geo_restricted_allowlist_allows_listed_country() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&["us", "gb"], &[]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), - VfAction::Allow - )); - } - - #[test] - fn geo_restricted_allowlist_drops_unlisted_country() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&["us", "gb"], &[]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_country_matching_is_case_insensitive() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - assert!(matches!( - rule.evaluate(&crate::rules::test_context( - &viewer_with_country("us"), - &geo_candidate(&["US"], &[]) - )), - VfAction::Allow - )); - assert!(matches!( - rule.evaluate(&crate::rules::test_context( - &viewer_with_country("de"), - &geo_candidate(&[], &["DE"]) - )), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_missing_country_fails_nonempty_allowlist() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&["us"], &[]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_missing_country_matches_xx_denylist() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&[], &["xx"]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_missing_country_not_in_denylist_allows() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let c = geo_candidate(&[], &["de"]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn geo_restricted_drops_even_for_author() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let mut c = geo_candidate(&[], &["de"]); - c.author_id = 999; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn geo_restricted_drops_retweets_too() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let mut c = geo_candidate(&[], &["de"]); - c.tweet_features.core.source_tweet_id = Some(99); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn dmca_drops() { - let rule = DropTweetsWithDmcaMediaRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - media: MediaFeature { - has_dmca_media: true, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(_) - )); - } -} diff --git a/visibility-filtering/rules/tweet_flag_rules.rs b/visibility-filtering/rules/tweet_flag_rules.rs deleted file mode 100644 index 0738c240..00000000 --- a/visibility-filtering/rules/tweet_flag_rules.rs +++ /dev/null @@ -1,125 +0,0 @@ -use crate::models::{TweetFeatures, VfAction}; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -#[derive(Clone)] -pub struct TweetFlagDropRule { - name: &'static str, - flag: fn(&TweetFeatures) -> bool, - reason: FilteredReason, -} - -impl TweetFlagDropRule { - pub const fn new( - name: &'static str, - flag: fn(&TweetFeatures) -> bool, - reason: FilteredReason, - ) -> Self { - Self { name, flag, reason } - } -} - -impl Rule for TweetFlagDropRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if (self.flag)(&context.candidate().tweet_features) { - return VfAction::Drop(self.reason.clone()); - } - VfAction::Allow - } -} - -pub const TWEET_NSFW_USER_DROP: TweetFlagDropRule = TweetFlagDropRule::new( - "TweetNsfwUserDropRule", - |t| t.nsfw.user, - FilteredReason::ContainNsfwMedia, -); -pub const TWEET_NSFW_ADMIN_DROP: TweetFlagDropRule = TweetFlagDropRule::new( - "TweetNsfwAdminDropRule", - |t| t.nsfw.admin, - FilteredReason::ContainNsfwMedia, -); - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - HydratedTweetCandidate, NsfwFeature, TweetFeatures, Viewer, ViewerFeatures, - }; - - fn candidate_with_nsfw_flags(user: bool, admin: bool) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { - nsfw: NsfwFeature { user, admin }, - ..Default::default() - }, - ..Default::default() - } - } - - fn viewer(id: u64) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(id), - ..Default::default() - } - } - - #[test] - fn tweet_nsfw_user_drops() { - let c = candidate_with_nsfw_flags(true, false); - assert!(matches!( - TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn tweet_nsfw_user_unset_allows() { - let c = candidate_with_nsfw_flags(false, false); - assert!(matches!( - TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn tweet_nsfw_user_drops_even_self_view() { - let c = candidate_with_nsfw_flags(true, false); - assert!(matches!( - TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn tweet_nsfw_admin_drops() { - let c = candidate_with_nsfw_flags(false, true); - assert!(matches!( - TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn tweet_nsfw_admin_unset_allows() { - let c = candidate_with_nsfw_flags(false, false); - assert!(matches!( - TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn tweet_nsfw_admin_drops_even_self_view() { - let c = candidate_with_nsfw_flags(false, true); - assert!(matches!( - TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } -} diff --git a/visibility-filtering/rules/tweet_label_drops.rs b/visibility-filtering/rules/tweet_label_drops.rs deleted file mode 100644 index d5c9b2ab..00000000 --- a/visibility-filtering/rules/tweet_label_drops.rs +++ /dev/null @@ -1,358 +0,0 @@ -use crate::models::{SafetyLabelType, VfAction}; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::{ - Action, DropReason, FilteredReason, SafetyResult, SafetyResultReason, -}; - -#[derive(Clone)] -pub struct SafetyLabelDropRule { - name: &'static str, - label: SafetyLabelType, - reason: FilteredReason, - exempt_author: bool, -} - -impl SafetyLabelDropRule { - pub const fn new( - name: &'static str, - label: SafetyLabelType, - reason: FilteredReason, - exempt_author: bool, - ) -> Self { - Self { - name, - label, - reason, - exempt_author, - } - } -} - -impl Rule for SafetyLabelDropRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.has_tweet_safety_label(self.label) { - return VfAction::Allow; - } - if self.exempt_author && context.is_author_viewer() { - return VfAction::Allow; - } - VfAction::Drop(self.reason.clone()) - } -} - -const NSFW_HIGH_PRECISION_REASON: FilteredReason = FilteredReason::SafetyResult(SafetyResult { - reason: Some(SafetyResultReason::NsfwHighPrecision), - action: Action::Drop(DropReason {}), -}); - -pub const PDNA_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "PdnaTweetLabelRule", - SafetyLabelType::PDNA, - NSFW_HIGH_PRECISION_REASON, - true, -); -pub const BOUNCE_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "BounceTweetLabelRule", - SafetyLabelType::BOUNCE, - FilteredReason::TweetIsBounced, - true, -); -pub const SPAM_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "SpamTweetLabelRule", - SafetyLabelType::SPAM, - FilteredReason::PossiblyUndesirable, - true, -); -pub const SPAM_HIGH_RECALL_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "SpamHighRecallDropRule", - SafetyLabelType::SPAM_HIGH_RECALL, - FilteredReason::PossiblyUndesirable, - true, -); -pub const NSFW_TEXT_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "NsfwTextTweetLabelDropRule", - SafetyLabelType::NSFW_TEXT, - NSFW_HIGH_PRECISION_REASON, - true, -); -pub const NSFW_HIGH_RECALL_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "NsfwHighRecallDropRule", - SafetyLabelType::NSFW_HIGH_RECALL, - FilteredReason::ContainNsfwMedia, - true, -); -pub const FOR_EMERGENCY_USE_ONLY_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "ForEmergencyUseOnlyDropRule", - SafetyLabelType::FOR_EMERGENCY_USE_ONLY, - FilteredReason::UnspecifiedReason, - false, -); - -pub const NSFW_HIGH_PRECISION_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "NsfwHighPrecisionOonDropRule", - SafetyLabelType::NSFW_HIGH_PRECISION, - FilteredReason::ContainNsfwMedia, - true, -); -pub const NSFW_CARD_IMAGE_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "NsfwCardImageOonDropRule", - SafetyLabelType::NSFW_CARD_IMAGE, - FilteredReason::ContainNsfwMedia, - true, -); -pub const GORE_AND_VIOLENCE_HIGH_PRECISION_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "GoreAndViolenceOonDropRule", - SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, - FilteredReason::ContainNsfwMedia, - true, -); -pub const DO_NOT_AMPLIFY_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "DoNotAmplifyOonDropRule", - SafetyLabelType::DO_NOT_AMPLIFY, - FilteredReason::PossiblyUndesirable, - true, -); -pub const MALICIOUS_URL_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "MaliciousUrlOonDropRule", - SafetyLabelType::MALICIOUS_URL, - FilteredReason::PossiblyUndesirable, - true, -); - -pub const FOSNR_HATEFUL_CONDUCT_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "FosnrHatefulConductDropRule", - SafetyLabelType::FOSNR_HATEFUL_CONDUCT, - FilteredReason::PossiblyUndesirable, - true, -); -pub const FOSNR_VIOLENT_SPEECH_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "FosnrViolentSpeechDropRule", - SafetyLabelType::FOSNR_VIOLENT_SPEECH, - FilteredReason::PossiblyUndesirable, - true, -); -pub const FOSNR_ABUSE_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "FosnrAbuseDropRule", - SafetyLabelType::FOSNR_ABUSE, - FilteredReason::PossiblyUndesirable, - true, -); -pub const FOSNR_CIVIC_INTEGRITY_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "FosnrCivicIntegrityDropRule", - SafetyLabelType::FOSNR_CIVIC_INTEGRITY, - FilteredReason::PossiblyUndesirable, - true, -); - -pub const FOSNR_ABUSE_INSULTS_OON_DROP: SafetyLabelDropRule = SafetyLabelDropRule::new( - "FosnrAbuseInsultsOonDropRule", - SafetyLabelType::FOSNR_ABUSE_INSULTS, - FilteredReason::PossiblyUndesirable, - true, -); - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - HydratedTweetCandidate, SafetyLabel, SafetyLabelMap, Viewer, ViewerFeatures, - }; - use std::collections::HashMap; - - fn candidate_with_label(label: SafetyLabelType) -> HydratedTweetCandidate { - let mut labels = HashMap::new(); - labels.insert(label, SafetyLabel::default()); - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - safety_labels: SafetyLabelMap::new(labels), - ..Default::default() - } - } - - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } - } - - fn author_viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - } - } - - #[test] - fn drops_non_author_with_mapped_reason() { - let c = candidate_with_label(SafetyLabelType::PDNA); - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::SafetyResult(_)) - )); - - let c = candidate_with_label(SafetyLabelType::BOUNCE); - assert!(matches!( - BOUNCE_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::TweetIsBounced) - )); - - let c = candidate_with_label(SafetyLabelType::SPAM); - assert!(matches!( - SPAM_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_RECALL); - assert!(matches!( - NSFW_HIGH_RECALL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn author_exempt_rules_allow_self_view() { - let c = candidate_with_label(SafetyLabelType::PDNA); - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - let c = candidate_with_label(SafetyLabelType::DO_NOT_AMPLIFY); - assert!(matches!( - DO_NOT_AMPLIFY_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn all_viewer_rules_drop_even_for_author() { - let c = candidate_with_label(SafetyLabelType::FOR_EMERGENCY_USE_ONLY); - assert!(matches!( - FOR_EMERGENCY_USE_ONLY_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - } - - #[test] - fn oon_media_drops_regardless_of_opt_in() { - let opted_in = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: true, - ..Default::default() - }; - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - assert!(matches!( - NSFW_HIGH_PRECISION_DROP.evaluate(&crate::rules::test_context(&opted_in, &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn no_label_allows() { - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - fn follower_viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } - } - - fn candidate_with_label_followed(label: SafetyLabelType) -> HydratedTweetCandidate { - let mut c = candidate_with_label(label); - c.relationship.viewer_follows_author = true; - c - } - - #[test] - fn fosnr_level3_drops_non_author_including_follower() { - for rule in [ - &FOSNR_HATEFUL_CONDUCT_DROP, - &FOSNR_VIOLENT_SPEECH_DROP, - &FOSNR_ABUSE_DROP, - &FOSNR_CIVIC_INTEGRITY_DROP, - ] { - let c = candidate_with_label(rule.label); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate_with_label_followed(rule.label); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&follower_viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - } - } - - #[test] - fn fosnr_level3_exempts_author() { - for rule in [ - &FOSNR_HATEFUL_CONDUCT_DROP, - &FOSNR_VIOLENT_SPEECH_DROP, - &FOSNR_ABUSE_DROP, - &FOSNR_CIVIC_INTEGRITY_DROP, - ] { - let c = candidate_with_label(rule.label); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } - } - - #[test] - fn malicious_url_drops_non_author_and_exempts_author() { - let c = candidate_with_label(SafetyLabelType::MALICIOUS_URL); - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn fosnr_abuse_insults_oon_drops_all_non_authors() { - let c = candidate_with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS); - assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate_with_label_followed(SafetyLabelType::FOSNR_ABUSE_INSULTS); - assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP - .evaluate(&crate::rules::test_context(&follower_viewer(), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate_with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS); - assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP - .evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/tweet_rules.rs b/visibility-filtering/rules/tweet_rules.rs new file mode 100644 index 00000000..ba214158 --- /dev/null +++ b/visibility-filtering/rules/tweet_rules.rs @@ -0,0 +1,1098 @@ +use crate::models::{SafetyLabelType, VfAction}; +use crate::rules::rule_spec::{RuleAction, RuleSpec}; +use crate::rules::RuleContext; +use xai_visibility_filtering::models::{ + Action, DropReason, FilteredReason, SafetyResult, SafetyResultReason, +}; + +const NSFW_HIGH_PRECISION_REASON: FilteredReason = FilteredReason::SafetyResult(SafetyResult { + reason: Some(SafetyResultReason::NsfwHighPrecision), + action: Action::Drop(DropReason {}), +}); + +pub(super) const TWEET_LABEL_DROPS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "PdnaTweetLabelRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::PDNA), + action: RuleAction::Drop(NSFW_HIGH_PRECISION_REASON), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "BounceTweetLabelRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::BOUNCE), + action: RuleAction::Drop(FilteredReason::TweetIsBounced), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "SpamTweetLabelRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::SPAM), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "ForEmergencyUseOnlyDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOR_EMERGENCY_USE_ONLY), + action: RuleAction::Drop(FilteredReason::UnspecifiedReason), + exempt_author: false, + }, + RuleSpec::Tweet { + name: "FosnrHatefulConductDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOSNR_HATEFUL_CONDUCT), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "FosnrViolentSpeechDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOSNR_VIOLENT_SPEECH), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "FosnrAbuseDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOSNR_ABUSE), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "FosnrCivicIntegrityDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOSNR_CIVIC_INTEGRITY), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, +]; + +pub(super) const NSFW_MEDIA_INTERSTITIALS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "NsfwHighPrecisionInterstitialRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_HIGH_PRECISION), + action: RuleAction::SensitiveMediaInterstitial(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "GoreAndViolenceInterstitialRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION), + action: RuleAction::SensitiveMediaInterstitial(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "NsfwCardImageInterstitialRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_CARD_IMAGE), + action: RuleAction::SensitiveMediaInterstitial(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, +]; + +pub(super) const OON_TWEET_FLAG_DROPS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "TweetNsfwUserDropRule", + when: |tweet| tweet.has_nsfw_user_flag(), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: false, + }, + RuleSpec::Tweet { + name: "TweetNsfwAdminDropRule", + when: |tweet| tweet.has_nsfw_admin_flag(), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: false, + }, +]; + +pub(super) const OON_TWEET_LABEL_DROPS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "NsfwHighRecallDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_HIGH_RECALL), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "NsfwHighPrecisionOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_HIGH_PRECISION), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "GoreAndViolenceOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "NsfwCardImageOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_CARD_IMAGE), + action: RuleAction::Drop(FilteredReason::ContainNsfwMedia), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "DoNotAmplifyOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::DO_NOT_AMPLIFY), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "MaliciousUrlOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::MALICIOUS_URL), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "SpamHighRecallDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::SPAM_HIGH_RECALL), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "NsfwTextTweetLabelDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::NSFW_TEXT), + action: RuleAction::Drop(NSFW_HIGH_PRECISION_REASON), + exempt_author: true, + }, + RuleSpec::Tweet { + name: "FosnrAbuseInsultsOonDropRule", + when: |tweet| tweet.has_safety_label(SafetyLabelType::FOSNR_ABUSE_INSULTS), + action: RuleAction::Drop(FilteredReason::PossiblyUndesirable), + exempt_author: true, + }, +]; + +fn drop_exclusive_tweet_content(context: &RuleContext<'_>) -> VfAction { + if !context.tweet().is_exclusive() { + return VfAction::Allow; + } + + if context.viewer().is_logged_out() { + return VfAction::Drop(FilteredReason::ExclusiveTweet); + } + + if context.viewer().is_conversation_author() { + return VfAction::Allow; + } + + if context.viewer().super_follows_author() { + return VfAction::Allow; + } + + if !context.tweet().is_retweet() && context.viewer().is_author() { + return VfAction::Allow; + } + + VfAction::Drop(FilteredReason::ExclusiveTweet) +} + +pub(super) const EXCLUSIVE_TWEET_DROP: &[RuleSpec] = &[RuleSpec::Custom { + name: "DropExclusiveTweetContentRule", + evaluate: drop_exclusive_tweet_content, +}]; + +pub(super) const NSFW_AUTHOR_INTERSTITIAL: &[RuleSpec] = &[RuleSpec::Tweet { + name: "NsfwAuthorInterstitialRule", + when: |tweet| tweet.is_nsfw_flagged() && tweet.has_media(), + action: RuleAction::SensitiveMediaInterstitial(FilteredReason::ContainNsfwMedia), + exempt_author: true, +}]; + +fn nsfw_base_condition(context: &RuleContext<'_>) -> bool { + let tweet = context.tweet(); + !context.viewer().is_author() + && tweet.has_media() + && (tweet.has_safety_label(SafetyLabelType::NSFW_HIGH_PRECISION) + || tweet.has_safety_label(SafetyLabelType::NSFW_HIGH_RECALL) + || (tweet.is_nsfw_flagged() && !tweet.is_retweet())) +} + +fn graphic_base_condition(context: &RuleContext<'_>) -> bool { + let tweet = context.tweet(); + !context.viewer().is_author() + && tweet.has_media() + && tweet.has_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION) +} + +fn nsfw_no_media_label_condition(context: &RuleContext<'_>) -> bool { + let tweet = context.tweet(); + !context.viewer().is_author() + && (tweet.has_safety_label(SafetyLabelType::NSFW_TEXT) + || tweet.has_safety_label(SafetyLabelType::NSFW_CARD_IMAGE)) +} + +fn sensitive_base_condition(context: &RuleContext<'_>) -> bool { + nsfw_base_condition(context) + || graphic_base_condition(context) + || nsfw_no_media_label_condition(context) +} + +fn sensitive_viewer_logged_out(context: &RuleContext<'_>) -> VfAction { + if context.viewer().is_logged_out() && sensitive_base_condition(context) { + VfAction::Drop(FilteredReason::ContainNsfwMedia) + } else { + VfAction::Allow + } +} + +fn sensitive_viewer_underage(context: &RuleContext<'_>) -> VfAction { + if (context.viewer().is_underage() || context.viewer().age_lookup_failed()) + && sensitive_base_condition(context) + { + VfAction::Drop(FilteredReason::ContainNsfwMedia) + } else { + VfAction::Allow + } +} + +fn sensitive_viewer_no_stated_age(context: &RuleContext<'_>) -> VfAction { + if context.viewer().has_no_stated_age() + && context + .viewer() + .country() + .is_some_and(|country| context.nsfw_gating_country(country)) + && sensitive_base_condition(context) + { + VfAction::Drop(FilteredReason::ContainNsfwMedia) + } else { + VfAction::Allow + } +} + +pub(super) const NULLCAST_DROP: &[RuleSpec] = &[RuleSpec::Tweet { + name: "NullcastedTweetDropRule", + when: |tweet| tweet.is_nullcast() && !tweet.is_retweet() && !tweet.is_community_tweet(), + action: RuleAction::Drop(FilteredReason::TweetIsNullcast), + exempt_author: false, +}]; + +fn drop_legal_takendown_post(context: &RuleContext<'_>) -> VfAction { + if !context.viewer().is_author() && context.takedown().legal_in_viewer_country() { + return VfAction::Drop(FilteredReason::UnspecifiedReason); + } + VfAction::Allow +} + +fn drop_local_laws_takendown_post(context: &RuleContext<'_>) -> VfAction { + if !context.viewer().is_author() && context.takedown().local_laws_in_viewer_country() { + return VfAction::Drop(FilteredReason::UnspecifiedReason); + } + VfAction::Allow +} + +fn drop_geo_restricted_media(context: &RuleContext<'_>) -> VfAction { + if context.takedown().media_restricted_in_viewer_country() { + VfAction::Drop(FilteredReason::UnspecifiedReason) + } else { + VfAction::Allow + } +} + +pub(super) const TES_HOME_DROPS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "DropStaleTweetsRule", + when: |tweet| tweet.is_stale() && !tweet.is_retweet(), + action: RuleAction::Drop(FilteredReason::UnspecifiedReason), + exempt_author: false, + }, + RuleSpec::Custom { + name: "DropLegalTakendownPostRule", + evaluate: drop_legal_takendown_post, + }, + RuleSpec::Custom { + name: "DropLocalLawsTakendownPostRule", + evaluate: drop_local_laws_takendown_post, + }, +]; + +pub(super) const FILTER_ALL: &[RuleSpec] = &[RuleSpec::Tweet { + name: "FilterAllRule", + when: |_| true, + action: RuleAction::Drop(FilteredReason::UnspecifiedReason), + exempt_author: false, +}]; + +pub(super) const RECS_MEDIA_DROPS: &[RuleSpec] = &[ + RuleSpec::Tweet { + name: "DropTweetsWithDmcaMediaRule", + when: |tweet| tweet.has_dmca_media(), + action: RuleAction::Drop(FilteredReason::UnspecifiedReason), + exempt_author: false, + }, + RuleSpec::Custom { + name: "DropTweetsWithGeoRestrictedMediaRule", + evaluate: drop_geo_restricted_media, + }, +]; + +pub(super) const SENSITIVE_VIEWER_DROPS: &[RuleSpec] = &[ + RuleSpec::Custom { + name: "SensitiveViewerLoggedOutDropRule", + evaluate: sensitive_viewer_logged_out, + }, + RuleSpec::Custom { + name: "SensitiveViewerUnderageDropRule", + evaluate: sensitive_viewer_underage, + }, + RuleSpec::Custom { + name: "SensitiveViewerNoStatedAgeDropRule", + evaluate: sensitive_viewer_no_stated_age, + }, +]; + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::{ + AuthorFeatures, ExclusiveContentFeatures, HydratedTweetCandidate, MediaFeature, + NsfwFeature, TweetFeatures, VfAction, Viewer, ViewerAge, ViewerFeatures, + }; + use crate::rules::fixtures::{ + assert_allows, assert_drops, author_viewer, candidate, logged_out_viewer, + nsfw_flag_media_candidates, sensitive_opt_in_viewer, viewer, VIEWER_ID, + }; + use crate::rules::test_context; + use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; + + fn trigger_label(name: &str) -> SafetyLabelType { + match name { + "PdnaTweetLabelRule" => SafetyLabelType::PDNA, + "BounceTweetLabelRule" => SafetyLabelType::BOUNCE, + "SpamTweetLabelRule" => SafetyLabelType::SPAM, + "ForEmergencyUseOnlyDropRule" => SafetyLabelType::FOR_EMERGENCY_USE_ONLY, + "FosnrHatefulConductDropRule" => SafetyLabelType::FOSNR_HATEFUL_CONDUCT, + "FosnrViolentSpeechDropRule" => SafetyLabelType::FOSNR_VIOLENT_SPEECH, + "FosnrAbuseDropRule" => SafetyLabelType::FOSNR_ABUSE, + "FosnrCivicIntegrityDropRule" => SafetyLabelType::FOSNR_CIVIC_INTEGRITY, + "NsfwHighRecallDropRule" => SafetyLabelType::NSFW_HIGH_RECALL, + "NsfwHighPrecisionOonDropRule" => SafetyLabelType::NSFW_HIGH_PRECISION, + "GoreAndViolenceOonDropRule" => SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, + "NsfwCardImageOonDropRule" => SafetyLabelType::NSFW_CARD_IMAGE, + "DoNotAmplifyOonDropRule" => SafetyLabelType::DO_NOT_AMPLIFY, + "MaliciousUrlOonDropRule" => SafetyLabelType::MALICIOUS_URL, + "SpamHighRecallDropRule" => SafetyLabelType::SPAM_HIGH_RECALL, + "NsfwTextTweetLabelDropRule" => SafetyLabelType::NSFW_TEXT, + "FosnrAbuseInsultsOonDropRule" => SafetyLabelType::FOSNR_ABUSE_INSULTS, + "NsfwHighPrecisionInterstitialRule" => SafetyLabelType::NSFW_HIGH_PRECISION, + "GoreAndViolenceInterstitialRule" => SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, + "NsfwCardImageInterstitialRule" => SafetyLabelType::NSFW_CARD_IMAGE, + _ => panic!("no trigger label for rule {name}"), + } + } + + const UNRELATED_LABEL: SafetyLabelType = SafetyLabelType::EGREGIOUS_NSFW; + + #[test] + fn tweet_label_drop_axis() { + for spec in TWEET_LABEL_DROPS.iter().chain(OON_TWEET_LABEL_DROPS) { + let RuleSpec::Tweet { + name, + action: RuleAction::Drop(reason), + exempt_author, + .. + } = spec + else { + panic!("{} is not a tweet-label drop row", spec.name()); + }; + let firing = candidate().with_label(trigger_label(name)).build(); + for v in [ + viewer(VIEWER_ID), + logged_out_viewer(), + sensitive_opt_in_viewer(), + ] { + assert_drops(spec, &v, &firing, reason); + } + let followed = candidate() + .with_label(trigger_label(name)) + .followed() + .build(); + assert_drops(spec, &viewer(VIEWER_ID), &followed, reason); + + let unrelated = candidate().with_label(UNRELATED_LABEL).build(); + assert_allows(spec, &viewer(VIEWER_ID), &unrelated); + + if *exempt_author { + assert_allows(spec, &author_viewer(), &firing); + } else { + assert_drops(spec, &author_viewer(), &firing, reason); + } + } + } + + #[test] + fn nsfw_media_interstitial_axis() { + for spec in NSFW_MEDIA_INTERSTITIALS { + let RuleSpec::Tweet { + name, + action: RuleAction::SensitiveMediaInterstitial(reason), + exempt_author: true, + .. + } = spec + else { + panic!( + "{} is not an author-exempt sensitive-media interstitial row", + spec.name() + ); + }; + let firing = candidate().with_label(trigger_label(name)).build(); + let action = spec.evaluate(&test_context(&viewer(VIEWER_ID), &firing)); + assert!( + matches!(&action, VfAction::Interstitial(r) if r == reason), + "{name} should interstitial, got {action:?}" + ); + assert_allows(spec, &sensitive_opt_in_viewer(), &firing); + assert_allows(spec, &author_viewer(), &firing); + let unrelated = candidate().with_label(UNRELATED_LABEL).build(); + assert_allows(spec, &viewer(VIEWER_ID), &unrelated); + } + } + + fn tweet_flag_features(name: &str) -> NsfwFeature { + match name { + "TweetNsfwUserDropRule" => NsfwFeature { + user: true, + admin: false, + }, + "TweetNsfwAdminDropRule" => NsfwFeature { + user: false, + admin: true, + }, + _ => panic!("no trigger flags for rule {name}"), + } + } + + #[test] + fn tweet_flag_drop_axis() { + for spec in OON_TWEET_FLAG_DROPS { + let RuleSpec::Tweet { + name, + action: RuleAction::Drop(reason), + exempt_author, + .. + } = spec + else { + panic!("{} is not a tweet-flag drop row", spec.name()); + }; + let nsfw = tweet_flag_features(name); + let firing = candidate() + .with_tweet_features(TweetFeatures { + nsfw, + ..Default::default() + }) + .build(); + assert_drops(spec, &viewer(VIEWER_ID), &firing, reason); + let unflagged = candidate().build(); + assert_allows(spec, &viewer(VIEWER_ID), &unflagged); + if *exempt_author { + assert_allows(spec, &author_viewer(), &firing); + } else { + assert_drops(spec, &author_viewer(), &firing, reason); + } + } + } + + fn exclusive_candidate( + tweet_id: u64, + author_id: u64, + root_author_id: u64, + ) -> HydratedTweetCandidate { + let mut c = candidate().tweet_id(tweet_id).author_id(author_id).build(); + c.exclusive_content = Some(ExclusiveContentFeatures { + conversation_author_id: root_author_id, + viewer_super_follows_author: false, + }); + c + } + + #[test] + fn nsfw_author_interstitial_axis() { + let spec = &NSFW_AUTHOR_INTERSTITIAL[0]; + let RuleSpec::Tweet { + name: "NsfwAuthorInterstitialRule", + action: RuleAction::SensitiveMediaInterstitial(reason), + exempt_author: true, + .. + } = spec + else { + panic!("{} is not the NSFW-author interstitial row", spec.name()); + }; + for firing in nsfw_flag_media_candidates() { + let action = spec.evaluate(&test_context(&viewer(VIEWER_ID), &firing)); + assert!( + matches!(&action, VfAction::Interstitial(r) if r == reason), + "{} should interstitial, got {action:?}", + spec.name() + ); + assert_allows(spec, &sensitive_opt_in_viewer(), &firing); + assert_allows(spec, &author_viewer(), &firing); + } + let [mut no_media, ..] = nsfw_flag_media_candidates(); + no_media.tweet_features.media.has_media = false; + assert_allows(spec, &viewer(VIEWER_ID), &no_media); + let [_, mut no_flags, ..] = nsfw_flag_media_candidates(); + no_flags.tweet_features.nsfw = NsfwFeature::default(); + assert_allows(spec, &viewer(VIEWER_ID), &no_flags); + } + + #[test] + fn exclusive_content_axis() { + let spec = &EXCLUSIVE_TWEET_DROP[0]; + assert_allows(spec, &viewer(VIEWER_ID), &candidate().build()); + + let exclusive = exclusive_candidate(1, 100, 100); + assert_drops( + spec, + &logged_out_viewer(), + &exclusive, + &FilteredReason::ExclusiveTweet, + ); + assert_allows(spec, &author_viewer(), &exclusive); + assert_drops( + spec, + &viewer(200), + &exclusive, + &FilteredReason::ExclusiveTweet, + ); + + let mut super_follow = exclusive_candidate(1, 100, 100); + super_follow + .exclusive_content + .as_mut() + .unwrap() + .viewer_super_follows_author = true; + assert_allows(spec, &viewer(200), &super_follow); + + let reply = exclusive_candidate(2, 200, 100); + assert_allows(spec, &viewer(200), &reply); + + let mut retweet = exclusive_candidate(2, 200, 100); + retweet.tweet_features.core.source_tweet_id = Some(99); + assert_drops( + spec, + &viewer(200), + &retweet, + &FilteredReason::ExclusiveTweet, + ); + } + + fn gating_viewer(age: ViewerAge) -> ViewerFeatures { + ViewerFeatures { + viewer_age: age, + country_code: Some("de".into()), + ..viewer(VIEWER_ID) + } + } + + fn media_label(label: SafetyLabelType) -> HydratedTweetCandidate { + candidate().with_label(label).with_media().build() + } + + fn no_media_label(label: SafetyLabelType) -> HydratedTweetCandidate { + let mut c = media_label(label); + c.tweet_features.media.has_media = false; + c + } + + fn nsfw_author_media() -> HydratedTweetCandidate { + candidate() + .with_media() + .with_author_features(AuthorFeatures { + is_nsfw_user: true, + ..Default::default() + }) + .build() + } + + fn nsfw_tweet_flag_media() -> HydratedTweetCandidate { + let mut c = nsfw_author_media(); + c.author_features = AuthorFeatures::default(); + c.tweet_features.nsfw = NsfwFeature { + user: true, + admin: false, + }; + c + } + + fn sensitive_spec(name: &str) -> &'static RuleSpec { + SENSITIVE_VIEWER_DROPS + .iter() + .find(|spec| spec.name() == name) + .unwrap_or_else(|| panic!("no sensitive-viewer row {name}")) + } + + fn sensitive_firing_candidates() -> Vec { + let mut admin_author = nsfw_author_media(); + admin_author.author_features = AuthorFeatures { + is_nsfw_admin: true, + ..Default::default() + }; + let mut admin_flag = nsfw_tweet_flag_media(); + admin_flag.tweet_features.nsfw = NsfwFeature { + user: false, + admin: true, + }; + let mut both_flags = nsfw_tweet_flag_media(); + both_flags.author_features.is_nsfw_user = true; + vec![ + media_label(SafetyLabelType::NSFW_HIGH_PRECISION), + media_label(SafetyLabelType::NSFW_HIGH_RECALL), + media_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION), + no_media_label(SafetyLabelType::NSFW_TEXT), + no_media_label(SafetyLabelType::NSFW_CARD_IMAGE), + nsfw_author_media(), + admin_author, + nsfw_tweet_flag_media(), + admin_flag, + both_flags, + ] + } + + #[test] + fn sensitive_viewer_content_axis() { + let underage = sensitive_spec("SensitiveViewerUnderageDropRule"); + let logged_out = sensitive_spec("SensitiveViewerLoggedOutDropRule"); + let no_age = sensitive_spec("SensitiveViewerNoStatedAgeDropRule"); + let reason = FilteredReason::ContainNsfwMedia; + let logged_out_viewer = ViewerFeatures { + viewer: Viewer::LoggedOut, + ..gating_viewer(ViewerAge::Unknown) + }; + for firing in sensitive_firing_candidates() { + assert_drops( + underage, + &gating_viewer(ViewerAge::Known(15)), + &firing, + &reason, + ); + assert_drops(logged_out, &logged_out_viewer, &firing, &reason); + assert_drops( + no_age, + &gating_viewer(ViewerAge::NotStated), + &firing, + &reason, + ); + } + } + + #[test] + fn sensitive_viewer_exemption_axis() { + let underage = sensitive_spec("SensitiveViewerUnderageDropRule"); + let logged_out = sensitive_spec("SensitiveViewerLoggedOutDropRule"); + let no_age = sensitive_spec("SensitiveViewerNoStatedAgeDropRule"); + let hp = media_label(SafetyLabelType::NSFW_HIGH_PRECISION); + let text = no_media_label(SafetyLabelType::NSFW_TEXT); + let reason = FilteredReason::ContainNsfwMedia; + + assert_allows(underage, &gating_viewer(ViewerAge::Known(18)), &hp); + assert_allows(underage, &gating_viewer(ViewerAge::Known(18)), &text); + assert_allows(underage, &gating_viewer(ViewerAge::Unknown), &hp); + assert_allows(no_age, &gating_viewer(ViewerAge::Unknown), &hp); + assert_allows(underage, &gating_viewer(ViewerAge::Unknown), &text); + assert_allows(no_age, &gating_viewer(ViewerAge::Unknown), &text); + + for firing in sensitive_firing_candidates() { + assert_drops( + underage, + &gating_viewer(ViewerAge::LookupFailed), + &firing, + &reason, + ); + } + assert_allows(no_age, &gating_viewer(ViewerAge::LookupFailed), &hp); + let us_failed = ViewerFeatures { + country_code: Some("us".into()), + ..gating_viewer(ViewerAge::LookupFailed) + }; + assert_drops(underage, &us_failed, &text, &reason); + let clean = candidate().build(); + assert_allows(underage, &gating_viewer(ViewerAge::LookupFailed), &clean); + + let opted_in = ViewerFeatures { + allows_sensitive_media: true, + ..gating_viewer(ViewerAge::Known(15)) + }; + assert_drops(underage, &opted_in, &hp, &reason); + + let mut self_hp = hp.clone(); + self_hp.author_id = VIEWER_ID; + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &self_hp); + let mut self_text = text.clone(); + self_text.author_id = VIEWER_ID; + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &self_text); + + let mut hp_no_media = hp.clone(); + hp_no_media.tweet_features.media.has_media = false; + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &hp_no_media); + let logged_out_viewer = ViewerFeatures { + viewer: Viewer::LoggedOut, + ..gating_viewer(ViewerAge::Unknown) + }; + assert_allows(logged_out, &logged_out_viewer, &hp_no_media); + assert_allows(logged_out, &gating_viewer(ViewerAge::Known(15)), &hp); + + let mut no_flags = nsfw_tweet_flag_media(); + no_flags.tweet_features.nsfw = NsfwFeature::default(); + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &no_flags); + + let mut flag_rt = nsfw_tweet_flag_media(); + flag_rt.tweet_features.core.source_tweet_id = Some(42); + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &flag_rt); + let mut flag_self = nsfw_tweet_flag_media(); + flag_self.author_id = VIEWER_ID; + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &flag_self); + + let mut author_rt = nsfw_author_media(); + author_rt.tweet_features.core.source_tweet_id = Some(42); + assert_allows(underage, &gating_viewer(ViewerAge::Known(15)), &author_rt); + let mut author_no_media = nsfw_author_media(); + author_no_media.tweet_features.media.has_media = false; + assert_allows( + underage, + &gating_viewer(ViewerAge::Known(15)), + &author_no_media, + ); + } + + fn tes_spec(name: &str) -> &'static RuleSpec { + TES_HOME_DROPS + .iter() + .chain(RECS_MEDIA_DROPS) + .find(|spec| spec.name() == name) + .unwrap_or_else(|| panic!("no TES row {name}")) + } + + fn stale_edit_control() -> Option { + Some(EditControl::Initial(EditControlInitial { + edit_tweet_ids: vec![1, 2], + ..Default::default() + })) + } + + fn takedown_candidate(reasons: Vec) -> HydratedTweetCandidate { + candidate() + .with_tweet_features(TweetFeatures { + takedown_reasons: reasons, + ..Default::default() + }) + .build() + } + + fn viewer_with_country(country: &str) -> ViewerFeatures { + ViewerFeatures { + country_code: Some(country.to_string()), + ..viewer(VIEWER_ID) + } + } + + fn geo_candidate(allow: &[&str], deny: &[&str]) -> HydratedTweetCandidate { + candidate() + .with_tweet_features(TweetFeatures { + media: MediaFeature { + geo_allow_list: allow.iter().map(|s| s.to_string()).collect(), + geo_deny_list: deny.iter().map(|s| s.to_string()).collect(), + ..Default::default() + }, + ..Default::default() + }) + .build() + } + + #[test] + fn filter_all_axis() { + let spec = &FILTER_ALL[0]; + let RuleSpec::Tweet { + name: "FilterAllRule", + action: RuleAction::Drop(reason), + exempt_author: false, + .. + } = spec + else { + panic!("{} is not the FilterAll row", spec.name()); + }; + let pristine = candidate().build(); + assert_drops(spec, &viewer(VIEWER_ID), &pristine, reason); + assert_drops(spec, &author_viewer(), &pristine, reason); + assert_drops(spec, &logged_out_viewer(), &pristine, reason); + } + + #[test] + fn stale_and_dmca_tweet_axis() { + let stale = tes_spec("DropStaleTweetsRule"); + let dmca = tes_spec("DropTweetsWithDmcaMediaRule"); + let reason = FilteredReason::UnspecifiedReason; + let stale_c = candidate() + .with_tweet_features(TweetFeatures { + edit_control: stale_edit_control(), + ..Default::default() + }) + .build(); + assert_drops(stale, &viewer(VIEWER_ID), &stale_c, &reason); + assert_allows(stale, &viewer(VIEWER_ID), &candidate().build()); + let stale_rt = candidate() + .with_tweet_features(TweetFeatures { + edit_control: stale_edit_control(), + ..Default::default() + }) + .retweet_of(99) + .build(); + assert_allows(stale, &viewer(VIEWER_ID), &stale_rt); + + let dmca_c = candidate() + .with_tweet_features(TweetFeatures { + media: MediaFeature { + has_dmca_media: true, + ..Default::default() + }, + ..Default::default() + }) + .build(); + assert_drops(dmca, &viewer(VIEWER_ID), &dmca_c, &reason); + assert_allows(dmca, &viewer(VIEWER_ID), &candidate().build()); + } + + #[test] + fn takedown_country_axis() { + let legal = tes_spec("DropLegalTakendownPostRule"); + let local = tes_spec("DropLocalLawsTakendownPostRule"); + let reason = FilteredReason::UnspecifiedReason; + let legal_c = takedown_candidate(vec![ + TakedownReason::LegalRequest { + country_code: "de".to_string(), + }, + TakedownReason::UnspecifiedReason { + country_code: "fr".to_string(), + }, + ]); + assert_drops(legal, &viewer_with_country("de"), &legal_c, &reason); + assert_allows(legal, &viewer_with_country("us"), &legal_c); + assert_allows(legal, &viewer(VIEWER_ID), &legal_c); + + let bystander = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "de".to_string(), + }]); + assert_allows(legal, &viewer_with_country("de"), &bystander); + assert_drops(local, &viewer_with_country("de"), &bystander, &reason); + assert_allows(local, &viewer_with_country("us"), &bystander); + + let legal_only = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); + assert_allows(local, &viewer_with_country("de"), &legal_only); + + let mut author_legal = legal_only.clone(); + author_legal.author_id = VIEWER_ID; + assert_allows(legal, &viewer_with_country("de"), &author_legal); + let mut author_local = bystander.clone(); + author_local.author_id = VIEWER_ID; + assert_allows(local, &viewer_with_country("de"), &author_local); + + let non_country = takedown_candidate(vec![ + TakedownReason::HatefulImagery, + TakedownReason::Unknown, + ]); + assert_allows(legal, &viewer_with_country("de"), &non_country); + assert_allows(local, &viewer_with_country("de"), &non_country); + } + + #[test] + fn takedown_worldwide_axis() { + let legal = tes_spec("DropLegalTakendownPostRule"); + let local = tes_spec("DropLocalLawsTakendownPostRule"); + let reason = FilteredReason::UnspecifiedReason; + + for code in ["xx", "xy", "XX"] { + let worldwide = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: code.to_string(), + }]); + assert_drops(legal, &viewer_with_country("us"), &worldwide, &reason); + assert_drops(legal, &viewer(VIEWER_ID), &worldwide, &reason); + } + + let bystander_worldwide = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "xx".to_string(), + }]); + assert_drops( + local, + &viewer_with_country("us"), + &bystander_worldwide, + &reason, + ); + assert_drops(local, &viewer(VIEWER_ID), &bystander_worldwide, &reason); + + let country_scoped = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); + assert_allows(legal, &viewer(VIEWER_ID), &country_scoped); + let bystander_scoped = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "de".to_string(), + }]); + assert_allows(local, &viewer(VIEWER_ID), &bystander_scoped); + + let dmca = takedown_candidate(vec![TakedownReason::Dmca]); + assert_drops(legal, &viewer_with_country("de"), &dmca, &reason); + assert_drops(legal, &viewer(VIEWER_ID), &dmca, &reason); + assert_allows(local, &viewer_with_country("de"), &dmca); + let mut author_dmca = dmca.clone(); + author_dmca.author_id = VIEWER_ID; + assert_allows(legal, &viewer(VIEWER_ID), &author_dmca); + } + + #[test] + fn geo_restricted_media_axis() { + let spec = tes_spec("DropTweetsWithGeoRestrictedMediaRule"); + let reason = FilteredReason::UnspecifiedReason; + assert_allows(spec, &viewer_with_country("us"), &geo_candidate(&[], &[])); + assert_drops( + spec, + &viewer_with_country("de"), + &geo_candidate(&[], &["de", "fr"]), + &reason, + ); + assert_allows( + spec, + &viewer_with_country("us"), + &geo_candidate(&[], &["de", "fr"]), + ); + assert_allows( + spec, + &viewer_with_country("us"), + &geo_candidate(&["us", "gb"], &[]), + ); + assert_drops( + spec, + &viewer_with_country("de"), + &geo_candidate(&["us", "gb"], &[]), + &reason, + ); + assert_allows( + spec, + &viewer_with_country("us"), + &geo_candidate(&["US"], &[]), + ); + assert_drops( + spec, + &viewer_with_country("de"), + &geo_candidate(&[], &["DE"]), + &reason, + ); + assert_drops( + spec, + &viewer(VIEWER_ID), + &geo_candidate(&["us"], &[]), + &reason, + ); + assert_drops( + spec, + &viewer(VIEWER_ID), + &geo_candidate(&[], &["xx"]), + &reason, + ); + assert_allows(spec, &viewer(VIEWER_ID), &geo_candidate(&[], &["de"])); + + let mut author = geo_candidate(&[], &["de"]); + author.author_id = VIEWER_ID; + assert_drops(spec, &viewer_with_country("de"), &author, &reason); + + let mut retweet = geo_candidate(&[], &["de"]); + retweet.tweet_features.core.source_tweet_id = Some(99); + assert_drops(spec, &viewer_with_country("de"), &retweet, &reason); + } + + #[test] + fn nullcast_drop_axis() { + let spec = &NULLCAST_DROP[0]; + let RuleSpec::Tweet { + name: "NullcastedTweetDropRule", + action: RuleAction::Drop(reason), + exempt_author: false, + .. + } = spec + else { + panic!("{} is not the nullcast drop row", spec.name()); + }; + let firing = candidate() + .with_tweet_features(TweetFeatures { + is_nullcast: true, + ..Default::default() + }) + .build(); + assert_drops(spec, &viewer(VIEWER_ID), &firing, reason); + assert_drops(spec, &author_viewer(), &firing, reason); + assert_allows(spec, &viewer(VIEWER_ID), &candidate().build()); + + let mut community = firing.clone(); + community.tweet_features.is_community_tweet = true; + assert_allows(spec, &viewer(VIEWER_ID), &community); + + let mut retweet = firing.clone(); + retweet.tweet_features.core.source_tweet_id = Some(99); + assert_allows(spec, &viewer(VIEWER_ID), &retweet); + } + + #[test] + fn no_stated_age_jurisdiction_axis() { + let no_age = sensitive_spec("SensitiveViewerNoStatedAgeDropRule"); + let hp = media_label(SafetyLabelType::NSFW_HIGH_PRECISION); + let text = no_media_label(SafetyLabelType::NSFW_TEXT); + let reason = FilteredReason::ContainNsfwMedia; + + let us = ViewerFeatures { + country_code: Some("us".into()), + ..gating_viewer(ViewerAge::NotStated) + }; + assert_allows(no_age, &us, &hp); + assert_allows(no_age, &us, &text); + + let missing = ViewerFeatures { + country_code: None, + ..gating_viewer(ViewerAge::NotStated) + }; + assert_allows(no_age, &missing, &hp); + + let account_overrides = ViewerFeatures { + country_code: Some("de".into()), + account_country_code: Some("us".into()), + ..gating_viewer(ViewerAge::NotStated) + }; + assert_allows(no_age, &account_overrides, &hp); + + let gating_account = ViewerFeatures { + country_code: Some("us".into()), + account_country_code: Some("kr".into()), + ..gating_viewer(ViewerAge::NotStated) + }; + assert_drops(no_age, &gating_account, &hp, &reason); + + let request_fallback = ViewerFeatures { + country_code: Some("de".into()), + account_country_code: None, + ..gating_viewer(ViewerAge::NotStated) + }; + assert_drops(no_age, &request_fallback, &hp, &reason); + } + + fn all_rule_slices() -> [&'static [RuleSpec]; 15] { + use crate::rules::author_rules::{ + AUTHOR_STATE_DROPS, OON_NSFW_AUTHOR_DROPS, OON_USER_LABEL_DROPS, SOCIALGRAPH_DROPS, + }; + [ + AUTHOR_STATE_DROPS, + TWEET_LABEL_DROPS, + NSFW_MEDIA_INTERSTITIALS, + OON_NSFW_AUTHOR_DROPS, + OON_TWEET_FLAG_DROPS, + OON_TWEET_LABEL_DROPS, + OON_USER_LABEL_DROPS, + SOCIALGRAPH_DROPS, + EXCLUSIVE_TWEET_DROP, + NSFW_AUTHOR_INTERSTITIAL, + NULLCAST_DROP, + TES_HOME_DROPS, + FILTER_ALL, + RECS_MEDIA_DROPS, + SENSITIVE_VIEWER_DROPS, + ] + } + + #[test] + fn wired_rule_names_are_unique_and_nonempty() { + let mut seen = std::collections::BTreeSet::new(); + for spec in all_rule_slices().into_iter().flatten() { + let name = spec.name(); + assert!(!name.is_empty(), "rule name must be non-empty"); + assert!(seen.insert(name), "duplicate wired rule name {name}"); + } + } +} diff --git a/visibility-filtering/rules/user_label_drops.rs b/visibility-filtering/rules/user_label_drops.rs deleted file mode 100644 index 54b87a8e..00000000 --- a/visibility-filtering/rules/user_label_drops.rs +++ /dev/null @@ -1,263 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; -use xai_x_thrift::user_labels::LabelValue; - -#[derive(Clone)] -pub struct UserSafetyLabelDropRule { - name: &'static str, - label: LabelValue, - reason: FilteredReason, - require_non_follower: bool, -} - -impl UserSafetyLabelDropRule { - pub const fn new( - name: &'static str, - label: LabelValue, - reason: FilteredReason, - require_non_follower: bool, - ) -> Self { - Self { - name, - label, - reason, - require_non_follower, - } - } -} - -impl Rule for UserSafetyLabelDropRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.is_author_viewer() { - return VfAction::Allow; - } - if !context.candidate().author_has_user_label(self.label) { - return VfAction::Allow; - } - if self.require_non_follower - && !context.viewer().viewer_is_logged_out() - && context.viewer_follows_author() - { - return VfAction::Allow; - } - VfAction::Drop(self.reason.clone()) - } -} - -pub const NSFW_HIGH_RECALL_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "NsfwHighRecallUserLabelRule", - LabelValue::NSFW_HIGH_RECALL, - FilteredReason::UnspecifiedReason, - false, -); -pub const NSFW_HIGH_PRECISION_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "NsfwHighPrecisionUserLabelRule", - LabelValue::NSFW_HIGH_PRECISION, - FilteredReason::UnspecifiedReason, - false, -); -pub const SPAM_HIGH_RECALL_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "SpamHighRecallUserLabelRule", - LabelValue::SPAM_HIGH_RECALL, - FilteredReason::UnspecifiedReason, - false, -); -pub const COMPROMISED_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "CompromisedUserLabelRule", - LabelValue::COMPROMISED, - FilteredReason::UnspecifiedReason, - false, -); -pub const READ_ONLY_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "ReadOnlyUserLabelRule", - LabelValue::READ_ONLY, - FilteredReason::UnspecifiedReason, - false, -); -pub const IMPERSONATION_HIGH_PRECISION_USER_DROP: UserSafetyLabelDropRule = - UserSafetyLabelDropRule::new( - "ImpersonationHighPrecisionUserLabelRule", - LabelValue::IMPERSONATION_HIGH_PRECISION, - FilteredReason::UnspecifiedReason, - false, - ); -pub const NSFW_AVATAR_IMAGE_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "NsfwAvatarImageRule", - LabelValue::NSFW_AVATAR_IMAGE, - FilteredReason::UnspecifiedReason, - false, -); -pub const NSFW_BANNER_IMAGE_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "NsfwBannerImageRule", - LabelValue::NSFW_BANNER_IMAGE, - FilteredReason::UnspecifiedReason, - false, -); -pub const ABUSIVE_HIGH_RECALL_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "AbusiveHighRecallRule", - LabelValue::ABUSIVE_HIGH_RECALL, - FilteredReason::UnspecifiedReason, - true, -); -pub const NSFW_NEAR_PERFECT_USER_DROP: UserSafetyLabelDropRule = UserSafetyLabelDropRule::new( - "NsfwNearPerfectAuthorRule", - LabelValue::NSFW_NEAR_PERFECT, - FilteredReason::UnspecifiedReason, - false, -); -pub const DO_NOT_AMPLIFY_NON_FOLLOWER_USER_DROP: UserSafetyLabelDropRule = - UserSafetyLabelDropRule::new( - "DoNotAmplifyNonFollowerRule", - LabelValue::DO_NOT_AMPLIFY, - FilteredReason::UnspecifiedReason, - true, - ); - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, UserLabelSet, Viewer, ViewerFeatures, - }; - use std::collections::HashSet; - - fn candidate_with_user_label(label: LabelValue) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { - user_labels: UserLabelSet::new(HashSet::from([label])), - ..Default::default() - }, - ..Default::default() - } - } - - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } - } - - fn author_viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - } - } - - #[test] - fn drops_when_author_has_label() { - let c = candidate_with_user_label(LabelValue::NSFW_HIGH_RECALL); - assert!(matches!( - NSFW_HIGH_RECALL_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - let c = candidate_with_user_label(LabelValue::COMPROMISED); - assert!(matches!( - COMPROMISED_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - let c = candidate_with_user_label(LabelValue::SPAM_HIGH_RECALL); - assert!(matches!( - SPAM_HIGH_RECALL_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - } - - #[test] - fn allows_author_self_view() { - let c = candidate_with_user_label(LabelValue::READ_ONLY); - assert!(matches!( - READ_ONLY_USER_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - - let c = candidate_with_user_label(LabelValue::NSFW_AVATAR_IMAGE); - assert!(matches!( - NSFW_AVATAR_IMAGE_USER_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - - let c = candidate_with_user_label(LabelValue::ABUSIVE_HIGH_RECALL); - assert!(matches!( - ABUSIVE_HIGH_RECALL_USER_DROP - .evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn allows_when_label_absent() { - let c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; - assert!(matches!( - IMPERSONATION_HIGH_PRECISION_USER_DROP - .evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn different_label_does_not_match() { - let c = candidate_with_user_label(LabelValue::LOW_QUALITY); - assert!(matches!( - NSFW_HIGH_PRECISION_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn avatar_banner_blacklist_drop_with_mapped_reason() { - let c = candidate_with_user_label(LabelValue::NSFW_AVATAR_IMAGE); - assert!(matches!( - NSFW_AVATAR_IMAGE_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - let c = candidate_with_user_label(LabelValue::NSFW_BANNER_IMAGE); - assert!(matches!( - NSFW_BANNER_IMAGE_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - } - - #[test] - fn abusive_high_recall_drops_non_followers_and_logged_out() { - let c = candidate_with_user_label(LabelValue::ABUSIVE_HIGH_RECALL); - assert!(matches!( - ABUSIVE_HIGH_RECALL_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - let logged_out = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..Default::default() - }; - assert!(matches!( - ABUSIVE_HIGH_RECALL_USER_DROP.evaluate(&crate::rules::test_context(&logged_out, &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - } - - #[test] - fn abusive_high_recall_allows_follower() { - let mut c = candidate_with_user_label(LabelValue::ABUSIVE_HIGH_RECALL); - c.relationship.viewer_follows_author = true; - assert!(matches!( - ABUSIVE_HIGH_RECALL_USER_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/user_rules.rs b/visibility-filtering/rules/user_rules.rs deleted file mode 100644 index 893ddd03..00000000 --- a/visibility-filtering/rules/user_rules.rs +++ /dev/null @@ -1,316 +0,0 @@ -use crate::models::{AuthorFeatures, VfAction}; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -#[derive(Clone)] -pub struct AuthorFlagDropRule { - name: &'static str, - flag: fn(&AuthorFeatures) -> bool, - reason: FilteredReason, -} - -impl AuthorFlagDropRule { - pub const fn new( - name: &'static str, - flag: fn(&AuthorFeatures) -> bool, - reason: FilteredReason, - ) -> Self { - Self { name, flag, reason } - } -} - -impl Rule for AuthorFlagDropRule { - fn name(&self) -> &'static str { - self.name - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if (self.flag)(&candidate.author_features) && !context.is_author_viewer() { - return VfAction::Drop(self.reason.clone()); - } - VfAction::Allow - } -} - -pub const SUSPENDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "SuspendedAuthorRule", - |a| a.is_suspended, - FilteredReason::AuthorIsSuspended, -); -pub const DEACTIVATED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DeactivatedAuthorRule", - |a| a.is_deactivated, - FilteredReason::AuthorIsDeactivated, -); -pub const ERASED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "ErasedAuthorRule", - |a| a.is_erased, - FilteredReason::AuthorAccountIsInactive, -); -pub const OFFBOARDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "OffboardedAuthorRule", - |a| a.is_offboarded, - FilteredReason::AuthorAccountIsInactive, -); -pub const NSFW_USER_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DropNsfwUserAuthorRule", - |a| a.is_nsfw_user, - FilteredReason::ContainNsfwMedia, -); -pub const NSFW_ADMIN_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DropNsfwAdminAuthorRule", - |a| a.is_nsfw_admin, - FilteredReason::ContainNsfwMedia, -); - -pub struct ProtectedAuthorDropRule; - -impl Rule for ProtectedAuthorDropRule { - fn name(&self) -> &'static str { - "ProtectedAuthorDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if candidate.author_features.is_protected - && !context.is_author_viewer() - && (viewer.viewer_is_logged_out() || !candidate.viewer_follows_author()) - { - return VfAction::Drop(FilteredReason::AuthorIsProtected); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, Viewer, ViewerAuthorRelationship, ViewerFeatures, - }; - - fn candidate_with_author( - suspended: bool, - deactivated: bool, - protected: bool, - ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { - is_suspended: suspended, - is_deactivated: deactivated, - is_protected: protected, - ..Default::default() - }, - ..Default::default() - } - } - - fn viewer(id: u64) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(id), - ..Default::default() - } - } - - #[test] - fn suspended_author_drops() { - let rule = SUSPENDED_AUTHOR_DROP; - let c = candidate_with_author(true, false, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn suspended_author_self_view_allows() { - let rule = SUSPENDED_AUTHOR_DROP; - let c = candidate_with_author(true, false, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Allow - )); - } - - #[test] - fn deactivated_author_drops() { - let rule = DEACTIVATED_AUTHOR_DROP; - let c = candidate_with_author(false, true, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn protected_author_drops_non_follower() { - let rule = ProtectedAuthorDropRule; - let c = candidate_with_author(false, false, true); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn protected_author_allows_follower() { - let rule = ProtectedAuthorDropRule; - let mut c = candidate_with_author(false, false, true); - c.relationship = ViewerAuthorRelationship { - viewer_follows_author: true, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - #[test] - fn protected_author_drops_logged_out_viewer() { - let rule = ProtectedAuthorDropRule; - let c = candidate_with_author(false, false, true); - let logged_out = ViewerFeatures { - viewer: Viewer::LoggedOut, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&logged_out, &c)), - VfAction::Drop(FilteredReason::AuthorIsProtected) - )); - } - - #[test] - fn protected_author_allows_self_view() { - let rule = ProtectedAuthorDropRule; - let c = candidate_with_author(false, false, true); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Allow - )); - } - - #[test] - fn erased_and_offboarded_drops_wire_flag_name_and_reason() { - let cases: [(&AuthorFlagDropRule, &str, AuthorFeatures); 2] = [ - ( - &ERASED_AUTHOR_DROP, - "ErasedAuthorRule", - AuthorFeatures { - is_erased: true, - ..Default::default() - }, - ), - ( - &OFFBOARDED_AUTHOR_DROP, - "OffboardedAuthorRule", - AuthorFeatures { - is_offboarded: true, - ..Default::default() - }, - ), - ]; - for (rule, name, author_features) in cases { - assert_eq!(rule.name(), name); - let flagged = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &flagged)), - VfAction::Drop(FilteredReason::AuthorAccountIsInactive) - )); - let unflagged = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &unflagged)), - VfAction::Allow - )); - } - } - - #[test] - fn normal_author_allows() { - let rule = SUSPENDED_AUTHOR_DROP; - let c = candidate_with_author(false, false, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } - - fn candidate_with_nsfw_author( - is_nsfw_user: bool, - is_nsfw_admin: bool, - ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { - is_nsfw_user, - is_nsfw_admin, - ..Default::default() - }, - ..Default::default() - } - } - - #[test] - fn nsfw_user_author_drops() { - let rule = NSFW_USER_AUTHOR_DROP; - let c = candidate_with_nsfw_author(true, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn nsfw_user_author_self_view_allows() { - let rule = NSFW_USER_AUTHOR_DROP; - let c = candidate_with_nsfw_author(true, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Allow - )); - } - - #[test] - fn nsfw_admin_author_drops() { - let rule = NSFW_ADMIN_AUTHOR_DROP; - let c = candidate_with_nsfw_author(false, true); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn nsfw_admin_author_self_view_allows() { - let rule = NSFW_ADMIN_AUTHOR_DROP; - let c = candidate_with_nsfw_author(false, true); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), - VfAction::Allow - )); - } - - #[test] - fn non_nsfw_author_allows() { - let rule = NSFW_USER_AUTHOR_DROP; - let c = candidate_with_nsfw_author(false, false); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/safety_label_source/codec.rs b/visibility-filtering/safety_label_source/codec.rs index bc52054c..36d23db0 100644 --- a/visibility-filtering/safety_label_source/codec.rs +++ b/visibility-filtering/safety_label_source/codec.rs @@ -1,33 +1,14 @@ -use xai_safety_label_store::error::SafetyLabelStoreError; -use xai_safety_label_store::types::{decode_lkey, SafetyLabelMap}; +use xai_safety_label_store::types::{decode_lkey_bytes, SafetyLabelMap}; use xai_visibility_filtering_proto as vf_pb; use xai_x_thrift::tweet_safety_label::SafetyLabel; use super::proto; use std::collections::HashMap; -use std::fmt; use std::panic::{catch_unwind, AssertUnwindSafe}; use tracing::debug; -const VERSION: u8 = 0x01; - -#[derive(Debug)] -pub enum CodecError { - TooManyLabels { count: usize }, -} - -impl fmt::Display for CodecError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::TooManyLabels { count } => { - write!(f, "too many labels to encode: {count} (max {})", u16::MAX) - } - } - } -} - #[derive(Debug, Clone)] pub struct LkeyBytes(pub [u8; 4]); @@ -40,76 +21,6 @@ pub struct RawSafetyLabel { pub mval: MvalBytes, } -pub struct EncodedSafetyLabelMap(Vec); - -impl EncodedSafetyLabelMap { - pub fn into_bytes(self) -> Vec { - self.0 - } -} - -pub fn encode(items: &[RawSafetyLabel]) -> Result { - let count: u16 = items - .len() - .try_into() - .map_err(|_| CodecError::TooManyLabels { count: items.len() })?; - let total: usize = 3 + items - .iter() - .map(|item| 8 + item.mval.0.len()) - .sum::(); - let mut buf = Vec::with_capacity(total); - buf.push(VERSION); - buf.extend_from_slice(&count.to_be_bytes()); - for item in items { - buf.extend_from_slice(&item.lkey.0); - buf.extend_from_slice(&(item.mval.0.len() as u32).to_be_bytes()); - buf.extend_from_slice(&item.mval.0); - } - Ok(EncodedSafetyLabelMap(buf)) -} - -pub fn decode(bytes: &[u8]) -> Option { - if bytes.len() < 3 { - return None; - } - if bytes[0] != VERSION { - return None; - } - let count = u16::from_be_bytes([bytes[1], bytes[2]]) as usize; - let mut cursor = 3; - let mut map = HashMap::with_capacity(count); - for _ in 0..count { - if cursor + 8 > bytes.len() { - return None; - } - let lkey = &bytes[cursor..cursor + 4]; - cursor += 4; - let len = u32::from_be_bytes(bytes[cursor..cursor + 4].try_into().ok()?) as usize; - cursor += 4; - if cursor + len > bytes.len() { - return None; - } - let mval = &bytes[cursor..cursor + len]; - cursor += len; - - let label_type = match decode_lkey(lkey) { - Ok(lt) => lt, - Err(SafetyLabelStoreError::InvalidLabelType(_)) => continue, - Err(_) => return None, - }; - let label = match xai_x_thrift::deserialize_mval(mval) { - Ok(l) => l, - Err(e) => { - debug!(label_type = ?label_type, error = %e, "skipping label value with corrupt mval from twemcache"); - map.insert(label_type, SafetyLabel::default()); - continue; - } - }; - map.insert(label_type, label); - } - Some(map) -} - pub fn decode_mval_payload(bytes: &[u8]) -> Option { xai_safety_label_store::mval_safety_label_map::decode_mval(bytes) .map(|labels| proto::label_map_to_proto(&labels)) @@ -117,28 +28,24 @@ pub fn decode_mval_payload(bytes: &[u8]) -> Option { pub(crate) enum DecodeAttempt { Success(vf_pb::SafetyLabelMap), - Failure(SafetyLabelStoreError), Panic, } pub(crate) fn decode_raw_labels(items: &[RawSafetyLabel]) -> DecodeAttempt { - match catch_unwind(AssertUnwindSafe(|| decode_raw_labels_inner(items))) { - Ok(Ok(map)) => DecodeAttempt::Success(proto::label_map_to_proto(&map)), - Ok(Err(e)) => DecodeAttempt::Failure(e), + contain_decode_panic(|| decode_raw_labels_inner(items)) +} + +fn contain_decode_panic(decode: impl FnOnce() -> SafetyLabelMap) -> DecodeAttempt { + match catch_unwind(AssertUnwindSafe(decode)) { + Ok(map) => DecodeAttempt::Success(proto::label_map_to_proto(&map)), Err(_) => DecodeAttempt::Panic, } } -fn decode_raw_labels_inner( - items: &[RawSafetyLabel], -) -> Result { +fn decode_raw_labels_inner(items: &[RawSafetyLabel]) -> SafetyLabelMap { let mut map = HashMap::with_capacity(items.len()); for item in items { - let label_type = match decode_lkey(&item.lkey.0) { - Ok(lt) => lt, - Err(SafetyLabelStoreError::InvalidLabelType(_)) => continue, - Err(e) => return Err(e), - }; + let label_type = decode_lkey_bytes(item.lkey.0); let label = match xai_x_thrift::deserialize_mval(&item.mval.0) { Ok(l) => l, Err(e) => { @@ -148,7 +55,7 @@ fn decode_raw_labels_inner( }; map.insert(label_type, label); } - Ok(map) + map } #[cfg(test)] @@ -177,102 +84,20 @@ mod tests { } } - #[test] - fn encode_decode_roundtrip_empty() { - let blob = encode(&[]).unwrap(); - assert!(blob.0.len() == 3); - let map = decode(&blob.0).unwrap(); - assert!(map.is_empty()); - } - - #[test] - fn encode_decode_roundtrip_one_label() { - let items = vec![raw_label(SafetyLabelType::SPAM, minimal_mval())]; - let blob = encode(&items).unwrap(); - let map = decode(&blob.0).unwrap(); - assert!(map.len() == 1); - assert!(map.contains_key(&SafetyLabelType::SPAM)); - } - - #[test] - fn encode_decode_roundtrip_multiple_labels() { - let items = vec![ - raw_label(SafetyLabelType::SPAM, minimal_mval()), - raw_label(SafetyLabelType::BOUNCE, mval_with_score()), - raw_label(SafetyLabelType::NSFW_HIGH_PRECISION, minimal_mval()), - ]; - let blob = encode(&items).unwrap(); - let map = decode(&blob.0).unwrap(); - assert!(map.len() == 3); - assert!(map.contains_key(&SafetyLabelType::SPAM)); - assert!(map.contains_key(&SafetyLabelType::BOUNCE)); - assert!(map.contains_key(&SafetyLabelType::NSFW_HIGH_PRECISION)); - assert_eq!( - map[&SafetyLabelType::BOUNCE].score.map(|f| f.into_inner()), - Some(0.9) - ); - } - - #[test] - fn decode_returns_none_on_empty() { - assert!(decode(&[]).is_none()); - } - - #[test] - fn decode_returns_none_on_garbage() { - assert!(decode(&[0xff; 16]).is_none()); - } - - #[test] - fn decode_returns_none_on_truncated() { - let items = vec![raw_label(SafetyLabelType::SPAM, minimal_mval())]; - let blob = encode(&items).unwrap(); - assert!(decode(&blob.0[..blob.0.len() - 2]).is_none()); - } - - #[test] - fn decode_returns_none_on_wrong_version() { - let mut blob = encode(&[]).unwrap().0; - blob[0] = 0x02; - assert!(decode(&blob).is_none()); - } - - #[test] - fn decode_includes_unknown_label_types() { - let unknown_lkey = (999i32 ^ i32::MIN).to_be_bytes(); - let items = vec![ - RawSafetyLabel { - lkey: LkeyBytes(unknown_lkey), - mval: MvalBytes(minimal_mval()), - }, - raw_label(SafetyLabelType::SPAM, minimal_mval()), - ]; - let blob = encode(&items).unwrap(); - let map = decode(&blob.0).unwrap(); - assert!(map.len() == 2); - assert!(map.contains_key(&SafetyLabelType::SPAM)); - assert!(map.contains_key(&SafetyLabelType::from(999))); - } - #[track_caller] fn decoded(items: &[RawSafetyLabel]) -> vf_pb::SafetyLabelMap { match decode_raw_labels(items) { DecodeAttempt::Success(map) => map, - DecodeAttempt::Failure(e) => panic!("expected Success, got Failure({e})"), DecodeAttempt::Panic => panic!("expected Success, got Panic"), } } #[test] - fn decode_raw_labels_matches_decode() { - let items = vec![ - raw_label(SafetyLabelType::SPAM, minimal_mval()), - raw_label(SafetyLabelType::BOUNCE, mval_with_score()), - ]; - let blob = encode(&items).unwrap(); - let from_decode = proto::label_map_to_proto(&decode(&blob.0).unwrap()); - let from_raw = decoded(&items); - assert!(from_decode == from_raw); + fn decode_raw_labels_contains_panics() { + assert!(matches!( + contain_decode_panic(|| panic!("test decode panic")), + DecodeAttempt::Panic + )); } #[test] @@ -294,20 +119,4 @@ mod tests { vf_pb::SafetyLabel::default() ); } - - #[test] - fn decode_preserves_label_type_on_corrupt_mval() { - let items = vec![ - raw_label(SafetyLabelType::SPAM, minimal_mval()), - raw_label(SafetyLabelType::BOUNCE, vec![0xDE, 0xAD]), - raw_label(SafetyLabelType::NSFW_HIGH_PRECISION, minimal_mval()), - ]; - let blob = encode(&items).unwrap(); - let map = decode(&blob.0).unwrap(); - assert_eq!(map.len(), 3); - assert!(map.contains_key(&SafetyLabelType::SPAM)); - assert!(map.contains_key(&SafetyLabelType::BOUNCE)); - assert!(map.contains_key(&SafetyLabelType::NSFW_HIGH_PRECISION)); - assert_eq!(map[&SafetyLabelType::BOUNCE], SafetyLabel::default()); - } } diff --git a/visibility-filtering/safety_label_source/lookup.rs b/visibility-filtering/safety_label_source/lookup.rs index 36e8fd04..fcabb52a 100644 --- a/visibility-filtering/safety_label_source/lookup.rs +++ b/visibility-filtering/safety_label_source/lookup.rs @@ -6,6 +6,7 @@ use xai_visibility_filtering_proto as vf_pb; use super::metrics::{self, BatchStage}; use super::types::{FailureKind, FallbackReason, LabelSource, ManhattanOutcome, TwemcacheOutcome}; +use super::warmer::Warmer; #[derive(Debug, Clone, thiserror::Error)] #[error("{kind:?}: {message}")] @@ -38,6 +39,7 @@ pub(crate) trait ManhattanLookup: Send + Sync { pub(crate) struct RemoteSource { twemcache: Arc, manhattan: Arc, + warmer: Option>, } impl RemoteSource { @@ -51,14 +53,21 @@ impl RemoteSource { Self { twemcache, manhattan, + warmer: None, } } + pub(crate) fn with_warmer(mut self, warmer: Arc) -> Self { + self.warmer = Some(warmer); + self + } + pub(crate) async fn get(&self, ids: &[u64]) -> LookupResults { let mut twemcache_results = self.twemcache.get(ids).await; let mut results = HashMap::with_capacity(ids.len()); let mut fallback_ids = Vec::new(); let mut fallback_counts: BTreeMap = BTreeMap::new(); + let mut warm_ids = Vec::new(); for &tweet_id in ids { match twemcache_results.remove(&tweet_id) { @@ -75,6 +84,9 @@ impl RemoteSource { } Some(TwemcacheOutcome::Miss) => { fallback_ids.push(tweet_id); + if self.warmer.is_some() { + warm_ids.push(tweet_id); + } } Some(TwemcacheOutcome::FallThrough(reason)) => { fallback_ids.push(tweet_id); @@ -93,6 +105,12 @@ impl RemoteSource { metrics::record_cache_fallback_keys(LabelSource::Twemcache, reason, count); } + if let Some(warmer) = &self.warmer + && !warm_ids.is_empty() + { + warmer.warm(warm_ids); + } + metrics::record_batch_size(BatchStage::ManhattanFallback, fallback_ids.len()); let mut manhattan_results = self.manhattan.get(&fallback_ids).await; for tweet_id in fallback_ids { @@ -182,6 +200,28 @@ mod tests { } } + struct FakeWarmer { + published: Mutex>>, + } + + impl FakeWarmer { + fn new() -> Arc { + Arc::new(Self { + published: Mutex::new(Vec::new()), + }) + } + + fn published(&self) -> Vec> { + self.published.lock().unwrap().clone() + } + } + + impl Warmer for FakeWarmer { + fn warm(&self, miss_ids: Vec) { + self.published.lock().unwrap().push(miss_ids); + } + } + fn empty_label_map() -> vf_pb::SafetyLabelMap { vf_pb::SafetyLabelMap { labels: HashMap::new(), @@ -284,6 +324,50 @@ mod tests { assert_eq!(manhattan.calls(), vec![vec![42]]); } + #[tokio::test] + async fn plain_miss_publishes_to_warmer() { + let twemcache = FakeTwemcache::new(HashMap::from([ + (1, TwemcacheOutcome::Hit(empty_label_map())), + (2, TwemcacheOutcome::NotFound), + (3, TwemcacheOutcome::Miss), + (4, TwemcacheOutcome::FallThrough(FallbackReason::Timeout)), + ])); + let manhattan = FakeManhattan::new(HashMap::from([ + (3, ManhattanOutcome::Resolved(empty_label_map())), + (4, ManhattanOutcome::Resolved(empty_label_map())), + (5, ManhattanOutcome::Resolved(empty_label_map())), + ])); + let warmer = FakeWarmer::new(); + let source = + RemoteSource::new(twemcache.clone(), manhattan.clone()).with_warmer(warmer.clone()); + + let results = source.get(&[1, 2, 3, 4, 5]).await; + + assert_eq!(results.len(), 5); + assert_eq!(warmer.published(), vec![vec![3]]); + } + + #[tokio::test] + async fn full_warm_channel_does_not_affect_fallback_result() { + use super::super::warmer::CacheWarmer; + + let (warmer, _rx) = CacheWarmer::without_drain_task(1); + let warmer = Arc::new(warmer); + warmer.warm(vec![0]); + + let twemcache = FakeTwemcache::new(HashMap::from([(42, TwemcacheOutcome::Miss)])); + let manhattan = FakeManhattan::new(HashMap::from([( + 42, + ManhattanOutcome::Resolved(empty_label_map()), + )])); + let source = RemoteSource::new(twemcache.clone(), manhattan.clone()).with_warmer(warmer); + + let results = source.get(&[42]).await; + + assert!(results.get(&42).unwrap().is_ok()); + assert_eq!(manhattan.calls(), vec![vec![42]]); + } + #[tokio::test] async fn get_mixed_results_merges_cache_and_manhattan() { let twemcache = FakeTwemcache::new(HashMap::from([ diff --git a/visibility-filtering/safety_label_source/manhattan.rs b/visibility-filtering/safety_label_source/manhattan.rs index 58f68a78..691de506 100644 --- a/visibility-filtering/safety_label_source/manhattan.rs +++ b/visibility-filtering/safety_label_source/manhattan.rs @@ -41,16 +41,6 @@ impl ManhattanLookup for ManhattanSource { successes += 1; results.insert(tweet_id, ManhattanOutcome::Resolved(label_map)); } - DecodeAttempt::Failure(e) => { - failures += 1; - results.insert( - tweet_id, - ManhattanOutcome::Failure(LookupError::new( - FailureKind::ManhattanDecode, - e.to_string(), - )), - ); - } DecodeAttempt::Panic => { failures += 1; results.insert( diff --git a/visibility-filtering/safety_label_source/metrics.rs b/visibility-filtering/safety_label_source/metrics.rs index 770eb6af..4e53d03c 100644 --- a/visibility-filtering/safety_label_source/metrics.rs +++ b/visibility-filtering/safety_label_source/metrics.rs @@ -15,6 +15,7 @@ const CACHE_KEYS: &str = "safety_labels_cache_keys"; const MANHATTAN_KEYS: &str = "safety_labels_manhattan_keys"; const CACHE_FALLBACK_KEYS: &str = "safety_labels_cache_fallback_keys"; const BATCH_SIZE: &str = "safety_labels_lookup_batch_size"; +const CACHE_WARM_KEYS: &str = "safety_labels_cache_warm_keys"; #[derive(Clone, Copy)] pub(crate) enum RequestOutcome { @@ -202,6 +203,35 @@ pub(crate) fn record_cache_fallback_keys( ); } +#[derive(Clone, Copy)] +pub(crate) enum WarmKeyResult { + EligibleMiss, + Enqueued, + DroppedChannelFull, + FetchIssued, + FetchFailed, +} + +impl WarmKeyResult { + fn as_str(self) -> &'static str { + match self { + Self::EligibleMiss => "eligible_miss", + Self::Enqueued => "enqueued", + Self::DroppedChannelFull => "dropped_channel_full", + Self::FetchIssued => "fetch_issued", + Self::FetchFailed => "fetch_failed", + } + } +} + +pub(crate) fn record_cache_warm_keys(result: WarmKeyResult, count: usize) { + incr_nonzero( + CACHE_WARM_KEYS, + &[("result", result.as_str())], + count as u64, + ); +} + pub(crate) fn record_batch_size(stage: BatchStage, size: usize) { observe( BATCH_SIZE, diff --git a/visibility-filtering/safety_label_source/mod.rs b/visibility-filtering/safety_label_source/mod.rs index 960bc379..6f72ac31 100644 --- a/visibility-filtering/safety_label_source/mod.rs +++ b/visibility-filtering/safety_label_source/mod.rs @@ -1,5 +1,5 @@ mod cached_value; -pub mod codec; +pub(crate) mod codec; mod expiring_cache; pub(crate) mod lookup; pub(crate) mod manhattan; @@ -9,6 +9,7 @@ mod proto; pub mod source; pub(crate) mod twemcache; pub(crate) mod types; +pub(crate) mod warmer; pub use lookup::LookupError; pub use mh_client::{ManhattanLabelFetcher, MhLabelClient}; diff --git a/visibility-filtering/safety_label_source/source.rs b/visibility-filtering/safety_label_source/source.rs index 92e9b263..ceee1472 100644 --- a/visibility-filtering/safety_label_source/source.rs +++ b/visibility-filtering/safety_label_source/source.rs @@ -33,7 +33,7 @@ fn ttl_for_tweet(tweet_id: u64, now: SystemTime) -> Option { } pub struct SafetyLabelSource { - cache: ExpiringCache, + cache: ExpiringCache>, remote: Arc, } @@ -59,17 +59,16 @@ impl SafetyLabelSource { pub async fn get( &self, ids: &[u64], - ) -> HashMap> { + ) -> HashMap, LookupError>> { let total = ids.len(); - let mut results: HashMap> = + let mut results: HashMap, LookupError>> = HashMap::with_capacity(total); let (local_misses, expired) = self.get_local(ids, &mut results); let batch_size = local_misses.len(); let remote_results = self.remote.get(&local_misses).await; - self.backfill_local(&remote_results); - results.extend(remote_results); + self.backfill_local(remote_results, &mut results); self.emit_stats(total - batch_size, batch_size, expired); @@ -79,7 +78,7 @@ impl SafetyLabelSource { fn get_local( &self, ids: &[u64], - results: &mut HashMap>, + results: &mut HashMap, LookupError>>, ) -> (Vec, usize) { let mut misses = Vec::with_capacity(ids.len()); let mut expired = 0; @@ -98,14 +97,20 @@ impl SafetyLabelSource { (misses, expired) } - fn backfill_local(&self, results: &HashMap>) { + fn backfill_local( + &self, + remote_results: HashMap>, + results: &mut HashMap, LookupError>>, + ) { let wall_now = SystemTime::now(); - for (&id, result) in results { - if let Ok(label_map) = result + for (id, result) in remote_results { + let result = result.map(Arc::new); + if let Ok(label_map) = &result && let Some(ttl) = ttl_for_tweet(id, wall_now) { - self.cache.insert(id, label_map.clone(), ttl); + self.cache.insert(id, Arc::clone(label_map), ttl); } + results.insert(id, result); } } @@ -257,12 +262,15 @@ mod tests { ); let results1 = source.get(&[42]).await; - let first = results1.get(&42).unwrap().as_ref().unwrap().clone(); + let first = Arc::clone(results1.get(&42).unwrap().as_ref().unwrap()); assert_eq!(remote_keys.load(Ordering::SeqCst), 1); assert!(source.cache.expiry_of(&42).is_some()); let results2 = source.get(&[42]).await; - assert_eq!(results2.get(&42).unwrap().as_ref().unwrap(), &first); + assert!(Arc::ptr_eq( + results2.get(&42).unwrap().as_ref().unwrap(), + &first + )); assert_eq!(remote_keys.load(Ordering::SeqCst), 1); } diff --git a/visibility-filtering/safety_label_source/twemcache.rs b/visibility-filtering/safety_label_source/twemcache.rs index 2813debe..fa5bb255 100644 --- a/visibility-filtering/safety_label_source/twemcache.rs +++ b/visibility-filtering/safety_label_source/twemcache.rs @@ -107,8 +107,8 @@ impl TwemcacheLookup for TwemcacheSource { let mut counts = ItemCounts::default(); let map = self.cache.multi_get(&keys).await; - for (i, &tweet_id) in ids.iter().enumerate() { - let outcome = match map.get(&keys[i]) { + for (&tweet_id, key) in ids.iter().zip(&keys) { + let outcome = match map.get(key) { Some(Ok(Some(bytes))) => match cached_value::decode(bytes) { CacheLookup::Hit(label_map) => { counts.hit += 1; diff --git a/visibility-filtering/safety_label_source/types.rs b/visibility-filtering/safety_label_source/types.rs index a62c348e..6d573d38 100644 --- a/visibility-filtering/safety_label_source/types.rs +++ b/visibility-filtering/safety_label_source/types.rs @@ -41,7 +41,6 @@ impl FallbackReason { pub(crate) enum FailureKind { ManhattanFetch, ManhattanDecode, - Other, } impl FailureKind { @@ -49,7 +48,6 @@ impl FailureKind { match self { Self::ManhattanFetch => "manhattan_fetch", Self::ManhattanDecode => "manhattan_decode", - Self::Other => "other", } } } diff --git a/visibility-filtering/safety_label_source/warmer.rs b/visibility-filtering/safety_label_source/warmer.rs new file mode 100644 index 00000000..c5047500 --- /dev/null +++ b/visibility-filtering/safety_label_source/warmer.rs @@ -0,0 +1,194 @@ +use std::sync::Arc; +use std::time::Duration; + +use tokio::sync::mpsc; +use tonic::async_trait; +use xai_strato::{encode, StratoGrpc}; + +use super::metrics::{self, WarmKeyResult}; + +const WARM_COLUMN_PATH: &str = "visibility/baseTweetSafetyLabelMap"; +const WARM_CHANNEL_CAPACITY: usize = 1024; +const WARM_LINGER: Duration = Duration::from_millis(500); +const WARM_FETCH_MAX_KEYS: usize = 50; + +pub(crate) trait Warmer: Send + Sync { + fn warm(&self, miss_ids: Vec); +} + +#[async_trait] +pub(crate) trait WarmFetcher: Send + Sync { + async fn fetch(&self, ids: &[u64]) -> usize; +} + +pub(crate) struct StratoWarmFetcher { + grpc: StratoGrpc, +} + +impl StratoWarmFetcher { + pub(crate) fn new(grpc: StratoGrpc) -> Self { + Self { grpc } + } +} + +#[async_trait] +impl WarmFetcher for StratoWarmFetcher { + async fn fetch(&self, ids: &[u64]) -> usize { + let calls = ids + .iter() + .map(|id| { + ( + WARM_COLUMN_PATH.to_string(), + "fetch".to_string(), + vec![encode(&(*id as i64, ()))], + ) + }) + .collect(); + self.grpc + .batch_call(calls, None) + .await + .iter() + .filter(|result| result.is_err()) + .count() + } +} + +pub(crate) struct CacheWarmer { + tx: mpsc::Sender>, +} + +impl CacheWarmer { + pub(crate) fn spawn(fetcher: Arc) -> Arc { + let (tx, rx) = mpsc::channel(WARM_CHANNEL_CAPACITY); + tokio::spawn(drain(rx, fetcher)); + Arc::new(Self { tx }) + } + + #[cfg(test)] + pub(crate) fn without_drain_task(capacity: usize) -> (Self, mpsc::Receiver>) { + let (tx, rx) = mpsc::channel(capacity); + (Self { tx }, rx) + } +} + +impl Warmer for CacheWarmer { + fn warm(&self, miss_ids: Vec) { + metrics::record_cache_warm_keys(WarmKeyResult::EligibleMiss, miss_ids.len()); + if miss_ids.is_empty() { + return; + } + let count = miss_ids.len(); + match self.tx.try_send(miss_ids) { + Ok(()) => metrics::record_cache_warm_keys(WarmKeyResult::Enqueued, count), + Err(_) => metrics::record_cache_warm_keys(WarmKeyResult::DroppedChannelFull, count), + } + } +} + +async fn drain(mut rx: mpsc::Receiver>, fetcher: Arc) { + while let Some(mut ids) = rx.recv().await { + let deadline = tokio::time::Instant::now() + WARM_LINGER; + while ids.len() < WARM_FETCH_MAX_KEYS { + match tokio::time::timeout_at(deadline, rx.recv()).await { + Ok(Some(mut more)) => ids.append(&mut more), + _ => break, + } + } + for chunk in ids.chunks(WARM_FETCH_MAX_KEYS) { + let failed = fetcher.fetch(chunk).await; + metrics::record_cache_warm_keys(WarmKeyResult::FetchIssued, chunk.len()); + metrics::record_cache_warm_keys(WarmKeyResult::FetchFailed, failed); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + struct FakeFetcher { + batches: Mutex>>, + failed_per_batch: usize, + } + + impl FakeFetcher { + fn new(failed_per_batch: usize) -> Arc { + Arc::new(Self { + batches: Mutex::new(Vec::new()), + failed_per_batch, + }) + } + + fn batches(&self) -> Vec> { + self.batches.lock().unwrap().clone() + } + } + + #[async_trait] + impl WarmFetcher for FakeFetcher { + async fn fetch(&self, ids: &[u64]) -> usize { + self.batches.lock().unwrap().push(ids.to_vec()); + self.failed_per_batch + } + } + + #[tokio::test(start_paused = true)] + async fn drain_lingers_then_flushes_in_chunks() { + let fetcher = FakeFetcher::new(0); + let warmer = CacheWarmer::spawn(fetcher.clone()); + + warmer.warm((0..30).collect()); + warmer.warm((30..60).collect()); + tokio::time::sleep(WARM_LINGER * 2).await; + warmer.warm(vec![100]); + tokio::time::sleep(WARM_LINGER * 2).await; + + let batches = fetcher.batches(); + assert_eq!(batches.len(), 3); + assert_eq!(batches[0], (0..50).collect::>()); + assert_eq!(batches[1], (50..60).collect::>()); + assert_eq!(batches[2], vec![100]); + } + + #[tokio::test(start_paused = true)] + async fn accumulation_is_capped_per_flush() { + let fetcher = FakeFetcher::new(0); + let warmer = CacheWarmer::spawn(fetcher.clone()); + + for start in (0..150).step_by(30) { + warmer.warm((start..start + 30).collect()); + } + tokio::time::sleep(WARM_LINGER * 2).await; + + let batches = fetcher.batches(); + assert!(batches + .iter() + .all(|batch| batch.len() <= WARM_FETCH_MAX_KEYS)); + assert_eq!(batches.concat(), (0..150).collect::>()); + } + + #[tokio::test(start_paused = true)] + async fn fetch_failure_does_not_stop_the_drain() { + let fetcher = FakeFetcher::new(1); + let warmer = CacheWarmer::spawn(fetcher.clone()); + + warmer.warm(vec![1]); + tokio::time::sleep(WARM_LINGER * 2).await; + warmer.warm(vec![2]); + tokio::time::sleep(WARM_LINGER * 2).await; + + assert_eq!(fetcher.batches(), vec![vec![1], vec![2]]); + } + + #[tokio::test] + async fn full_channel_drops_without_blocking() { + let (warmer, mut rx) = CacheWarmer::without_drain_task(1); + + warmer.warm(vec![1]); + warmer.warm(vec![2]); + + assert_eq!(rx.try_recv(), Ok(vec![1])); + assert!(rx.try_recv().is_err(), "the second publish was dropped"); + } +} diff --git a/visibility-filtering/server.rs b/visibility-filtering/server.rs index 64f760e3..1b58b4b1 100644 --- a/visibility-filtering/server.rs +++ b/visibility-filtering/server.rs @@ -15,7 +15,7 @@ impl xai_x_service_builder::XService for VFServer { type Config = (); async fn build(ctx: xai_x_service_builder::ServiceContext<()>) -> Self { - VFServer::new(&ctx.datacenter).await + VFServer::new(&ctx.datacenter, ctx.feature_switches).await } fn register(self: Arc, routes: &mut tonic::service::RoutesBuilder) { @@ -28,11 +28,14 @@ impl xai_x_service_builder::XService for VFServer { } impl VFServer { - pub async fn new(datacenter: &str) -> Self { - crate::server_deps::build_prod_server(datacenter).await + pub(crate) async fn new( + datacenter: &str, + feature_switches: Arc, + ) -> Self { + crate::server_deps::build_prod_server(datacenter, feature_switches).await } - pub fn from_endpoints( + pub(crate) fn from_endpoints( filter_tweets: FilterTweetsEndpoint, get_safety_labels: GetSafetyLabelsEndpoint, ) -> Self { diff --git a/visibility-filtering/server_deps.rs b/visibility-filtering/server_deps.rs index 6e41332a..12c8f579 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -2,12 +2,14 @@ use crate::clients::socialgraph_client::ProdSocialgraphClient; use crate::filter::{FilterRequest, FilterResponse, FilterTweets}; use crate::filter_tweets::FilterTweetsEndpoint; use crate::get_safety_labels::GetSafetyLabelsEndpoint; -use crate::hydration::{FallbackCacheMode, HydrationPipeline}; +use crate::hydration::HydrationPipeline; use crate::models::{RawCandidate, TweetId}; +use crate::reference_compare::ReferenceCompareHarness; use crate::rules::{SafetyLevel, Verdict}; use crate::safety_label_source::lookup::RemoteSource; use crate::safety_label_source::manhattan::ManhattanSource; use crate::safety_label_source::twemcache::TwemcacheSource; +use crate::safety_label_source::warmer::{CacheWarmer, StratoWarmFetcher, Warmer}; use crate::safety_label_source::{ManhattanLabelFetcher, MhLabelClient, SafetyLabelSource}; use crate::server::VFServer; use std::future::Future; @@ -18,6 +20,7 @@ use xai_core_entities::gizmoduck_client::{GizmoduckClientConfig, ProdGizmoduckCl use xai_core_entities::rpc_constants::{GizmoduckRpcConstants, RpcConstants, TESRpcConstants}; use xai_core_entities::s2s::{S2S_CHAIN_PATH, S2S_CRT_PATH, S2S_KEY_PATH}; use xai_core_entities::tweet_entity_service_client::{ProdTESClient, TESClientConfig}; +use xai_visibility_filtering::vf_client::{StratoVfClient, VfClient}; use xai_x_rpc::balanced_channel::LbPolicy; const CACHE_PATH: &str = "/s/cache/safety_label_store:twemcaches"; @@ -82,27 +85,22 @@ where } } -pub async fn build_prod_server(datacenter: &str) -> VFServer { +#[expect( + clippy::expect_used, + reason = "startup fail-fast: init failure is fatal" +)] +pub async fn build_prod_server( + datacenter: &str, + feature_switches: Arc, +) -> VFServer { info!("Initializing prod clients for datacenter={}", datacenter); let init_deadline = tokio::time::Instant::now() + CLIENT_INIT_RETRY_BUDGET; let deterministic_aperture = std::env::var("APP_ENV").as_deref() == Ok("prod"); - let fallback_cache_serve_stale = crate::config::fallback_cache_serve_stale_enabled(); - let fallback_cache_mode = if fallback_cache_serve_stale { - FallbackCacheMode::ServeStale - } else if crate::config::fallback_cache_populate_enabled() { - FallbackCacheMode::Shadow - } else { - FallbackCacheMode::Disabled - }; - let media_fallback_cache_mode = if crate::config::media_fallback_cache_serve_stale_enabled() { - FallbackCacheMode::ServeStale - } else if crate::config::media_fallback_cache_populate_enabled() { - FallbackCacheMode::Shadow - } else { - FallbackCacheMode::Disabled - }; + let fallback_cache_enabled = crate::config::fallback_cache_enabled(); + let fallback_cache = fallback_cache_enabled + .then(crate::hydration::gizmoduck_hydrator::GizmoduckAuthorHydrator::fallback_cache); let tes_client: Arc< dyn xai_core_entities::tweet_entity_service_client::TESClient + Send + Sync, @@ -119,10 +117,7 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { .expect("Failed to initialize TES client"), ); - let gizmoduck_client_id = format!( - "visibility-filtering-service.{}", - std::env::var("APP_ENV").unwrap_or_else(|_| "prod".to_string()) - ); + let gizmoduck_client_id = crate::config::gizmoduck_client_id(); let gizmoduck_client: Arc< dyn xai_core_entities::gizmoduck_client::GizmoduckClient + Send + Sync, > = Arc::new( @@ -170,11 +165,12 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { .expect("Failed to initialize MhLabelClient"), ); + let twemcache_client_name = crate::config::twemcache_client_name(); let twemcache = Arc::new( init_client_with_retry("twemcache", init_deadline, || { crate::twemcache::TwemcacheClient::new_with_tls_paths( CACHE_PATH, - "visibility-filtering-service", + twemcache_client_name.clone(), datacenter, &S2S_CHAIN_PATH, &S2S_CRT_PATH, @@ -190,12 +186,20 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { ); info!("Cache client connected to {CACHE_PATH}"); + let reference_compare = build_reference_compare_harness(datacenter, init_deadline).await; + warm_cache(&twemcache).await; warm_manhattan(mh_label_client.as_ref()).await; + let cache_warmer = build_cache_warmer(datacenter, init_deadline).await; + let twemcache_source = Arc::new(TwemcacheSource::new(twemcache)); let manhattan_source = Arc::new(ManhattanSource::new(mh_label_client)); - let remote = Arc::new(RemoteSource::new(twemcache_source, manhattan_source)); + let mut remote = RemoteSource::new(twemcache_source, manhattan_source); + if let Some(warmer) = cache_warmer { + remote = remote.with_warmer(warmer); + } + let remote = Arc::new(remote); let safety_label_source = Arc::new(SafetyLabelSource::new(remote)); let hydration_pipeline = HydrationPipeline::new( @@ -203,30 +207,118 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { gizmoduck_client, sg_client, safety_label_source.clone(), - fallback_cache_mode, - media_fallback_cache_mode, + fallback_cache, ); - let policies = crate::rules::Policies::new(); - let (home_rule_count, recommendations_rule_count) = policies.rule_counts(); - let filter_tweets = FilterTweets::new(hydration_pipeline, policies); + let gating_countries = Arc::new(crate::params::NsfwGatingCountries::starting_at_default()); + let fs_path = crate::config::fs_path(); + gating_countries.refresh_and_check_drift(&feature_switches, &fs_path); + gating_countries.spawn_refresh(feature_switches, fs_path); + let rule_engine = crate::rules::RuleEngine::with_nsfw_gating_countries(gating_countries); + let (home_rule_count, recommendations_rule_count) = rule_engine.rule_counts(); + let filter_tweets = FilterTweets::new(hydration_pipeline, rule_engine); warm_filter_tweets(&filter_tweets).await; info!( hydrator_count = 5, - ?fallback_cache_mode, - ?media_fallback_cache_mode, + fallback_cache_enabled, home_rule_count, recommendations_rule_count, "VFServer initialized with prod clients" ); VFServer::from_endpoints( - FilterTweetsEndpoint::new(filter_tweets), + FilterTweetsEndpoint::new(filter_tweets, reference_compare), GetSafetyLabelsEndpoint::new(safety_label_source), ) } +async fn build_reference_compare_harness( + datacenter: &str, + init_deadline: tokio::time::Instant, +) -> Option> { + #[expect(clippy::panic, reason = "startup fail-fast on misconfiguration")] + let should_build = crate::reference_compare::should_build_harness( + crate::config::dual_call_harness_enabled(), + std::env::var("APP_ENV").ok().as_deref(), + ) + .unwrap_or_else(|misconfiguration| panic!("{misconfiguration}")); + if !should_build { + return None; + } + + let client_id = format!( + "visibility-filtering-service.{}", + std::env::var("APP_ENV").unwrap_or_else(|_| "staging".to_string()) + ); + #[expect( + clippy::expect_used, + reason = "startup fail-fast: init failure is fatal" + )] + let strato: Arc = Arc::new( + init_client_with_retry("strato_vf", init_deadline, || { + let client_id = client_id.clone(); + async move { + StratoVfClient::new( + S2S_CHAIN_PATH.clone(), + S2S_CRT_PATH.clone(), + S2S_KEY_PATH.clone(), + client_id, + datacenter.to_string(), + ) + .await + .map_err(|e| e.to_string()) + } + }) + .await + .expect("Failed to initialize Strato VF client (reference comparator)"), + ); + Some(Arc::new(ReferenceCompareHarness::new(strato, datacenter))) +} + +const CACHE_WARM_REQUEST_TIMEOUT_MS: u64 = 500; + +#[expect( + clippy::expect_used, + reason = "startup fail-fast: init failure is fatal" +)] +async fn build_cache_warmer( + datacenter: &str, + init_deadline: tokio::time::Instant, +) -> Option> { + if !crate::config::cache_warm_enabled() { + return None; + } + + let client_id = format!( + "visibility-filtering-service.{}", + std::env::var("APP_ENV").unwrap_or_else(|_| "prod".to_string()) + ); + let grpc = init_client_with_retry("strato_cache_warm", init_deadline, || { + let config = xai_strato::StratoGrpcConfig { + ca_cert_path: S2S_CHAIN_PATH.clone(), + client_cert_path: S2S_CRT_PATH.clone(), + client_key_path: S2S_KEY_PATH.clone(), + num_endpoints: Some(12), + connect_timeout_ms: 400, + request_timeout_ms: CACHE_WARM_REQUEST_TIMEOUT_MS, + client_id: Some(client_id.clone()), + service_url: format!("stratostore.stratoserver.prod.{datacenter}.s2s.twttr.net"), + zone: datacenter.to_string(), + ..Default::default() + }; + async move { + xai_strato::StratoGrpc::new(config) + .await + .map_err(|e| e.to_string()) + } + }) + .await + .expect("Failed to initialize Strato cache-warm client"); + info!("L2 cache warmer enabled"); + Some(CacheWarmer::spawn(Arc::new(StratoWarmFetcher::new(grpc)))) +} + const TES_STRATO_REQUEST_TIMEOUT_MS: u64 = 100; fn tes_client_config(deterministic_aperture: bool) -> TESClientConfig { @@ -241,12 +333,15 @@ fn tes_client_config(deterministic_aperture: bool) -> TESClientConfig { } } +const GIZMODUCK_STRATO_REQUEST_TIMEOUT_MS: u64 = 80; + fn gizmoduck_client_config(deterministic_aperture: bool) -> GizmoduckClientConfig { GizmoduckClientConfig { aperture_size: Some(GizmoduckRpcConstants::num_endpoints()), deterministic_aperture, lb_policy: Some(LbPolicy::least_request()), readiness_probe_port: Some(GIZMODUCK_READINESS_PROBE_PORT), + request_timeout_ms: Some(GIZMODUCK_STRATO_REQUEST_TIMEOUT_MS), ..Default::default() } } @@ -354,7 +449,7 @@ async fn warm_manhattan(manhattan: &dyn ManhattanLabelFetcher) { #[cfg(test)] mod tests { use super::*; - use crate::filter::{FilterOutcome, FilterSummary}; + use crate::filter::FilterOutcome; use crate::models::VfAction; use std::cell::Cell; use xai_visibility_filtering::models::FilteredReason; @@ -363,6 +458,13 @@ mod tests { tokio::time::Instant::now() + budget } + #[tokio::test] + async fn reference_compare_harness_not_built_without_flag() { + let harness = + build_reference_compare_harness("atla", deadline_in(Duration::from_secs(1))).await; + assert!(harness.is_none()); + } + #[tokio::test(start_paused = true)] async fn init_retry_recovers_from_transient_failures() { let attempts = Cell::new(0u32); @@ -457,22 +559,7 @@ mod tests { safety_labels: None, }) .collect(); - FilterResponse { - summary: FilterSummary { - tweet_count: outcomes.len(), - drop_count: outcomes - .iter() - .filter(|outcome| matches!(outcome.verdict.action, VfAction::Drop(_))) - .count(), - unresolved_author_count: outcomes - .iter() - .filter(|outcome| { - outcome.verdict.decided_by == Verdict::unresolved_author().decided_by - }) - .count(), - }, - outcomes, - } + FilterResponse { outcomes } } #[test] diff --git a/visibility-filtering/twemcache/client.rs b/visibility-filtering/twemcache/client.rs index c55beef2..800eb32d 100644 --- a/visibility-filtering/twemcache/client.rs +++ b/visibility-filtering/twemcache/client.rs @@ -12,10 +12,6 @@ use super::key::{Key, Server, ServerSet, Value}; use super::metrics::Metrics; use super::ring::HashRing; -#[allow(dead_code)] -pub(crate) const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_millis(75); -#[allow(dead_code)] -pub(crate) const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_millis(200); const DEPTH_SAMPLE_INTERVAL: Duration = Duration::from_secs(5); pub(crate) trait PoolFactory: Send + Sync { diff --git a/visibility-filtering/twemcache/connection.rs b/visibility-filtering/twemcache/connection.rs index 150d5ce9..89495fb1 100644 --- a/visibility-filtering/twemcache/connection.rs +++ b/visibility-filtering/twemcache/connection.rs @@ -56,6 +56,7 @@ impl Shared { return; } self.metrics.record_connection_event(ConnEvent::TornDown); + #[expect(clippy::unwrap_used, reason = "propagate FIFO-lock poisoning")] let mut q = self.pending.lock().unwrap(); while let Some(p) = q.pop_front() { p.fail(err.clone()); @@ -176,6 +177,10 @@ impl PipelinedConnection { let len = 4 + key_bytes.len() + 2; let mut stack_buf = [0u8; 256 + 6]; let fallback; + #[expect( + clippy::indexing_slicing, + reason = "this arm requires len <= stack_buf.len()" + )] let cmd: &[u8] = if len <= stack_buf.len() { stack_buf[..4].copy_from_slice(b"get "); stack_buf[4..4 + key_bytes.len()].copy_from_slice(key_bytes); @@ -192,6 +197,7 @@ impl PipelinedConnection { return Err(TwemcacheError::Unavailable); } { + #[expect(clippy::unwrap_used, reason = "propagate FIFO-lock poisoning")] let mut q = self.shared.pending.lock().unwrap(); q.push_back(Pending::Get(tx)); self.shared.in_flight.fetch_add(1, Ordering::AcqRel); @@ -234,6 +240,7 @@ async fn dispatch_version( return Err(TwemcacheError::Unavailable); } { + #[expect(clippy::unwrap_used, reason = "propagate FIFO-lock poisoning")] let mut q = shared.pending.lock().unwrap(); q.push_back(Pending::Version(tx)); shared.in_flight.fetch_add(1, Ordering::AcqRel); @@ -308,6 +315,7 @@ async fn reader_loop( fn deliver(shared: &Shared, reply: Reply) -> std::result::Result<(), ()> { let pending = { + #[expect(clippy::unwrap_used, reason = "propagate FIFO-lock poisoning")] let mut q = shared.pending.lock().unwrap(); match q.pop_front() { Some(p) => { @@ -367,15 +375,18 @@ async fn read_one_reply(reader: &mut R) -> Result = header.split_whitespace().collect(); - if parts.len() < 4 || parts[0] != "VALUE" { - return Err(TwemcacheError::Io(format!( - "unexpected response: {header:?}" - ))); - } - let _flags: u32 = parts[2] + let (flags_part, byte_count_part) = match parts.as_slice() { + ["VALUE", _key, flags, byte_count, ..] => (*flags, *byte_count), + _ => { + return Err(TwemcacheError::Io(format!( + "unexpected response: {header:?}" + ))); + } + }; + let _flags: u32 = flags_part .parse() .map_err(|_| TwemcacheError::Io(format!("invalid flags: {header:?}")))?; - let byte_count: usize = parts[3] + let byte_count: usize = byte_count_part .parse() .map_err(|_| TwemcacheError::Io(format!("invalid byte count: {header:?}")))?; @@ -412,8 +423,9 @@ pub fn build_mtls_connector( use tokio_rustls::client::TlsConnector; use tokio_rustls::rustls::{client::ClientConfig, RootCertStore}; - let _ = - rustls::crypto::CryptoProvider::install_default(rustls::crypto::ring::default_provider()); + let _ = rustls::crypto::CryptoProvider::install_default( + rustls::crypto::aws_lc_rs::default_provider(), + ); let mut roots = RootCertStore::empty(); for cert in load_certs(ca_cert_path)? { @@ -483,7 +495,7 @@ impl ChainVerifyingSkipNameVerifier { fn new(roots: Arc) -> Result { let inner = rustls::client::WebPkiServerVerifier::builder_with_provider( roots, - Arc::new(rustls::crypto::ring::default_provider()), + Arc::new(rustls::crypto::aws_lc_rs::default_provider()), ) .build() .map_err(|e| TwemcacheError::Io(format!("build server cert verifier: {e}")))?; diff --git a/visibility-filtering/twemcache/host_pool.rs b/visibility-filtering/twemcache/host_pool.rs index d21ab4f5..66d45702 100644 --- a/visibility-filtering/twemcache/host_pool.rs +++ b/visibility-filtering/twemcache/host_pool.rs @@ -1,4 +1,4 @@ -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, MutexGuard}; use std::time::Duration; use quanta::{Clock, Instant}; @@ -78,9 +78,9 @@ type BestConn = Option<(usize, Arc, usize)>; pub struct HostPool { factory: Arc, slots: Vec>, - open_lock: TokioMutex<()>, + open_lock: TokioMutex<()>, depth_cap: usize, - probe_stale_after: Duration, + probe_stale_after: Duration, clock: Clock, metrics: Metrics, } @@ -155,6 +155,14 @@ impl HostPool { } } + #[expect( + clippy::indexing_slicing, + reason = "every idx comes from enumerate() over self.slots" + )] + fn slot_guard(&self, idx: usize) -> MutexGuard<'_, ConnSlot> { + lock_slot(&self.slots[idx]) + } + pub async fn get(&self, key: &Key) -> Result> { let now = self.clock.now(); let (conn, idx, is_probe) = match self.select(now).await? { @@ -165,7 +173,7 @@ impl HostPool { if !is_probe && result.as_ref().err().is_some_and(is_dead_socket_err) { { - let mut g = self.slots[idx].lock().unwrap(); + let mut g = self.slot_guard(idx); if g.conn.as_ref().is_some_and(|c| Arc::ptr_eq(c, &conn)) { g.conn = None; } @@ -187,7 +195,7 @@ impl HostPool { result } - async fn select(&self, now: Instant) -> Result { + async fn select(&self, now: Instant) -> Result { let (best, live_exists) = { let (best, live_exists, any_need) = self.scan(); if any_need { @@ -216,12 +224,12 @@ impl HostPool { } } - fn scan(&self) -> (BestConn, bool, bool) { + fn scan(&self) -> (BestConn, bool, bool) { let mut best: BestConn = None; let mut live_exists = false; let mut any_need = false; for (idx, slot) in self.slots.iter().enumerate() { - let g = slot.lock().unwrap(); + let g = lock_slot(slot); if g.health.is_live() && !g.has_living_conn() { any_need = true; } @@ -236,9 +244,9 @@ impl HostPool { (best, live_exists, any_need) } - async fn ensure_live_open(&self, now: Instant) { + async fn ensure_live_open(&self, now: Instant) { let any_need = self.slots.iter().any(|s| { - let g = s.lock().unwrap(); + let g = lock_slot(s); g.health.is_live() && !g.has_living_conn() }); if !any_need { @@ -248,7 +256,7 @@ impl HostPool { let _guard = self.open_lock.lock().await; for slot in &self.slots { let need = { - let g = slot.lock().unwrap(); + let g = lock_slot(slot); g.health.is_live() && !g.has_living_conn() }; if !need { @@ -256,14 +264,14 @@ impl HostPool { } match self.factory.open().await { Ok(conn) => { - let mut g = slot.lock().unwrap(); + let mut g = lock_slot(slot); if g.health.is_live() && !g.has_living_conn() { g.conn = Some(conn); } } Err(_) => { let tripped = { - let mut g = slot.lock().unwrap(); + let mut g = lock_slot(slot); if g.health.is_live() { g.health.record_connect_failure(now); true @@ -279,14 +287,14 @@ impl HostPool { } } - async fn try_open_probe(&self, now: Instant) -> Option<(usize, Arc)> { + async fn try_open_probe(&self, now: Instant) -> Option<(usize, Arc)> { let idx = { let mut claimed = None; for (i, slot) in self.slots.iter().enumerate() { - let mut g = slot.lock().unwrap(); + let mut g = lock_slot(slot); if g.health.is_probe_eligible(now, self.probe_stale_after) { g.health.begin_probe(now); - g.conn = None; + g.conn = None; claimed = Some(i); break; } @@ -297,12 +305,12 @@ impl HostPool { let _guard = self.open_lock.lock().await; match self.factory.open().await { Ok(conn) => { - let mut g = self.slots[idx].lock().unwrap(); + let mut g = self.slot_guard(idx); g.conn = Some(Arc::clone(&conn)); Some((idx, conn)) } Err(_) => { - let mut g = self.slots[idx].lock().unwrap(); + let mut g = self.slot_guard(idx); g.health.note_probe(Outcome::Failure, now); None } @@ -312,7 +320,7 @@ impl HostPool { fn record(&self, idx: usize, result: &Result>, is_probe: bool, now: Instant) { let outcome = classify(result); let event = { - let mut g = self.slots[idx].lock().unwrap(); + let mut g = self.slot_guard(idx); if is_probe { g.health.note_probe(outcome, now); match outcome { @@ -331,9 +339,9 @@ impl HostPool { } } - pub(crate) fn for_each_depth(&self, mut sink: impl FnMut(usize)) { + pub(crate) fn for_each_depth(&self, mut sink: impl FnMut(usize)) { for slot in &self.slots { - let g = slot.lock().unwrap(); + let g = lock_slot(slot); if let Some(c) = &g.conn { sink(c.in_flight()); } @@ -341,6 +349,11 @@ impl HostPool { } } +#[expect(clippy::unwrap_used, reason = "propagate slot-lock poisoning")] +fn lock_slot(slot: &Mutex) -> MutexGuard<'_, ConnSlot> { + slot.lock().unwrap() +} + #[cfg(test)] #[path = "host_pool_tests.rs"] mod tests; diff --git a/visibility-filtering/twemcache/key.rs b/visibility-filtering/twemcache/key.rs index db19ebaa..b42e915e 100644 --- a/visibility-filtering/twemcache/key.rs +++ b/visibility-filtering/twemcache/key.rs @@ -48,7 +48,7 @@ pub enum KeyError { } #[derive(Clone, Debug, Hash, Eq, PartialEq)] -pub struct Key(Vec); +pub struct Key(Box); impl Key { pub fn new(bytes: Vec) -> std::result::Result { @@ -58,29 +58,24 @@ impl Key { if bytes.iter().any(|&b| matches!(b, b' ' | b'\n' | b'\r')) { return Err(KeyError::Whitespace); } - std::str::from_utf8(&bytes).map_err(|_| KeyError::InvalidUtf8)?; - Ok(Key(bytes)) - } - - pub fn from_bytes(bytes: Vec) -> Self { - Key(bytes) + let s = String::from_utf8(bytes).map_err(|_| KeyError::InvalidUtf8)?; + Ok(Key(s.into_boxed_str())) } pub fn as_bytes(&self) -> &[u8] { - &self.0 + self.0.as_bytes() } } impl AsRef<[u8]> for Key { fn as_ref(&self) -> &[u8] { - &self.0 + self.0.as_bytes() } } impl std::fmt::Display for Key { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let s = std::str::from_utf8(self.as_bytes()).expect("Key::new validates UTF-8"); - f.write_str(s) + f.write_str(&self.0) } } diff --git a/visibility-filtering/twemcache/metrics.rs b/visibility-filtering/twemcache/metrics.rs index 3b7761f5..ceb0e241 100644 --- a/visibility-filtering/twemcache/metrics.rs +++ b/visibility-filtering/twemcache/metrics.rs @@ -41,7 +41,7 @@ enum Sink { } impl Metrics { - pub(crate) fn global() -> Self { + pub(crate) fn global() -> Self { Self(Sink::Global) } @@ -65,7 +65,7 @@ impl Metrics { } } - pub(crate) fn record_backpressure(&self, count: u64) { + pub(crate) fn record_backpressure(&self, count: u64) { if count > 0 && let Some(sr) = self.sink() { @@ -73,7 +73,7 @@ impl Metrics { } } - pub(crate) fn observe_pipeline_depth(&self, depth: usize) { + pub(crate) fn observe_pipeline_depth(&self, depth: usize) { if let Some(sr) = self.sink() { sr.observe( PIPELINE_DEPTH, diff --git a/visibility-filtering/twemcache/mod.rs b/visibility-filtering/twemcache/mod.rs index 4d426d6c..e17a3f70 100644 --- a/visibility-filtering/twemcache/mod.rs +++ b/visibility-filtering/twemcache/mod.rs @@ -8,6 +8,6 @@ pub(crate) mod key; mod metrics; pub(crate) mod ring; -pub use client::{TwemcacheClient, default_connections_per_host, default_depth_cap}; +pub use client::TwemcacheClient; pub use error::{Result, TwemcacheError}; -pub use key::{Key, KeyError, Value}; +pub use key::{Key, Value}; diff --git a/visibility-filtering/twemcache/ring.rs b/visibility-filtering/twemcache/ring.rs index 323f4d3f..026126e1 100644 --- a/visibility-filtering/twemcache/ring.rs +++ b/visibility-filtering/twemcache/ring.rs @@ -27,13 +27,9 @@ impl HashRing { md.update(i.to_string().as_bytes()); let hash = md.finalize_reset(); - for offset in [0, 4, 8, 12] { - let value = u32::from_le_bytes([ - hash[offset], - hash[offset + 1], - hash[offset + 2], - hash[offset + 3], - ]) as u64; + let (words, _remainder) = hash.as_slice().as_chunks::<4>(); + for &word in words { + let value = u32::from_le_bytes(word) as u64; ring.insert(value, server.clone()); } }